From 5eecc5f9c50dadad44c92dd4106a7e16e3562f31 Mon Sep 17 00:00:00 2001 From: laoyao0822 Date: Fri, 5 Jun 2026 07:45:29 +0800 Subject: [PATCH] Preserve CP ragged topk coordinates under batch planning Batch-planning metadata can now be attached even for single-request CP prefill, which routed FP8 flashmla_sparse through the batch cp_index path. That path used compact MQA-buffer row bases for score lookup but did not override the final ragged topk coordinate base consumed by attention, so topk indices could point at the wrong KV rows and produce low accept length or meaningless output.\n\nThis keeps ordinary long page tails out of compute padding, reserves compute padding for truly tiny suffixes, and makes cp_index RAGGED topk emit request-ragged offsets while preserving the compact buffer descriptors used for score materialization. The debug ledger records the rejected intermediate diagnoses and the confirmed coordinate-space failure.\n\nConstraint: CP shared-KV cache residency is page-granular, but attention/index compute must not consume synthetic long-tail rows.\nConstraint: FP8 CP prefill uses flashmla_sparse/RAGGED, where fused topk output is consumed directly as attention page_table_1.\nRejected: Disable current reuse or batch planning | would hide the regression and lose the intended bs>1 fast path.\nRejected: Treat all page tails as compute padding | regresses bs=1 semantics and can corrupt query/topk row contracts.\nConfidence: medium\nScope-risk: moderate\nDirective: Do not change cp_index RAGGED topk offset handling without verifying score-buffer row bases and final attention KV coordinate space independently.\nTested: python -m py_compile on touched Python/test files; git diff --check; remote targeted ragged cp_index offset regression test; remote test_nsa_cp_utils.py; remote test_cp_shared_kv_runtime.py; user-reported ETE output recovered after restart.\nNot-tested: Agent-driven full ETE traffic run; broad multi-request bs>1 production soak. --- ...accept_len1_garbage_output_debug_ledger.md | 114 ++++++++++++ .../srt/layers/attention/nsa/nsa_indexer.py | 22 +++ .../sglang/srt/layers/attention/nsa/utils.py | 33 ++-- .../unit/layers/test_nsa_cp_utils.py | 169 ++++++++++++++++-- .../mem_cache/test_cp_shared_kv_runtime.py | 10 +- 5 files changed, 316 insertions(+), 32 deletions(-) diff --git a/docs/advanced_features/nsa_prefill_cp_accept_len1_garbage_output_debug_ledger.md b/docs/advanced_features/nsa_prefill_cp_accept_len1_garbage_output_debug_ledger.md index 7254dbecf..2703d65f8 100644 --- a/docs/advanced_features/nsa_prefill_cp_accept_len1_garbage_output_debug_ledger.md +++ b/docs/advanced_features/nsa_prefill_cp_accept_len1_garbage_output_debug_ledger.md @@ -1382,3 +1382,117 @@ Verification status: - `test_nsa_cp_utils.py` 与 runtime tests 混跑时，远端 pytest collection 仍可能被当前环境的 `sgl_kernel` 安装不完整阻塞： `cannot import name 'gelu_tanh_and_mul' from 'sgl_kernel'`。 这是测试环境 import 问题，不作为本轮语义验证失败；关键合同用 targeted/direct 验证覆盖。 + +## 2026-06-05 regression focus update: baseline is pre-bs>1, but split-shape compute-padding fix is already in HEAD + +User-confirmed baseline: before the bs>1 CP shared-KV work, both BF16 KV cache and FP8 E4M3 KV cache produced sane output. Therefore the current garbage-output issue remains a bs>1-series regression until proven otherwise. + +Correction to avoid repeated re-checking: + +- The earlier concrete `40387/page-tail compute_padding_enabled` bug has already been fixed in HEAD: `compute_padding_enabled` is now based on valid-vs-compute split shape, so non-page-aligned long prompts enter the explicit compute-padding contract instead of mixing valid tensors with compute metadata. +- The just-committed current-slot fix also preserves MLA/index current reuse and TAI materialize; it does not disable fast paths. +- The latest remote logs available at `/mnt/beegfs/cjy/log/sglang_cp_hicache_20260604_221557.log` show the 40k bad request at `22:23:22`, which is before local commit `c8510938b` (`22:25:06 UTC-equivalent log time`). After that commit, the log only shows health/warmup 64-token traffic, not a new 40k business repro. + +Next investigation should not re-open the already-fixed split-shape bug unless a post-`c8510938b` business repro shows the same metadata inconsistency. The next code-level audit should focus on bs>1-series changes that still affect bs=1 despite the split-shape/current-slot fixes: EAGLE/draft side-channel gating, CP batch-plan vs scalar-path differences, and FP8 packed KV materialize/transfer byte-order/scale contracts. + +## 2026-06-05 user-confirmed regression boundary + +User explicitly confirmed: before the bs>1 CP shared-KV work, both FP8 E4M3 KV cache and BF16 KV cache were correct. Treat this as the regression boundary. + +Implication: + +- Do not spend more cycles treating FP8/BF16 support itself as the primary cause unless new evidence contradicts this boundary. +- The active suspect set is bs>1-era code that also affects bs=1: batch-plan metadata replacing scalar CP metadata, compute-padding/page-tail handling, current/partial reuse changes, and EAGLE/draft side-channel interaction with the new CP metadata. +- A valid fix should preserve the pre-bs>1 FP8/BF16 fast paths and repair the bs>1-induced semantic difference, not blanket-disable MLA/index cache reuse. + +## 2026-06-05 correction: bs>1-era compute padding over-applies to normal long page tails + +User confirmed the regression boundary: before the bs>1 work, both FP8 and BF16 paths were correct. That makes the current primary suspect the bs>1 batch-plan / compute-padding path changing the bs=1 prefill semantics, not FP8/BF16 storage itself. + +New RED tests on remote `g0034` show current production code treats `extend_len=40387,page_size=64,cp_size=8` as compute-padded: + +- valid local rows on CP rank0: `4995` +- compute local rows on CP rank0: `5056` +- `compute_padding_enabled=True` + +This is now considered an over-correction. A normal long request with only a final partial page should keep compute rows equal to valid rows; only tiny/zero-lane cases where padded page count is below CP size need synthetic compute rows. Physical cache page coverage remains page-granular, but attention/index compute must not ingest dummy tail rows for ordinary long page tails. + +Supersedes the earlier 2026-06-05 note that blamed `compute_padding_enabled=False` for the 40387 case. That older diagnosis matched one intermediate failure mode but no longer matches the current code or the user-confirmed regression boundary. + +## 2026-06-05 post long-tail compute-padding fix: still garbage output + +User reports the latest run still produces meaningless output after the long-tail compute-padding correction. + +Facts already checked and should not be re-opened blindly: + +- Remote `/mnt/beegfs/cjy/sglang-dev` hashes matched local for `utils.py`, `nsa_indexer.py`, `nsa_backend.py`, and `hiradix_cache.py` before the latest run. +- Latest log `/mnt/beegfs/cjy/log/sglang_cp_hicache_20260604_225727.log` is after the sync and contains the 40k request. +- For `extend_len` around 40386/40387/40391, `build_batch_page_aligned_in_seq_split_plan()` now reports `compute_padding_enabled=False` and local valid rows equal local compute rows on rank0. +- Latest bad request has `#cached-token: 0`, so HiCache L2->L1 load-back/prefix reuse is not the primary boundary for the first bad output. +- `Prefill batch #new-token=40448` is scheduler page-rounded budget/log accounting, not proof that model compute consumed 40448 rows. +- `HiCache-write owned_positions=5056` is page-owner backup/reservation coverage, not proof that attention/index compute consumed 5056 rows. + +Current narrowed suspect set: + +1. no-cache current-only CP shared-KV MLA materialize path; +2. no-cache current-only CP shared-KV index/MQA materialize path; +3. CP shared-KV direct-write local loc construction still returning physical/padded rows at a runtime callsite; +4. EAGLE/draft reads malformed prefill KV after target current-only materialize. + +Next action: inspect `nsa_backend.py` and `nsa_indexer.py` current-only callsites and add one-shot layer-0 diagnostics if static inspection cannot prove the row contracts. + + +## 2026-06-05 latest repro: still garbage after long-tail and current-slot fixes + +User reports output is still meaningless. Keep the confirmed regression boundary: before bs>1 CP shared-KV changes, BF16 and FP8 E4M3 were both correct. + +New facts from code audit to avoid repeated searches: + +- Because bs>1 work now attaches a CP shared-KV batch plan even for single-request prefill, several bs=1 callsites route through batch-plan logic. In particular, `NSAIndexer._get_topk_in_seq_cp_pair()` enters `_get_topk_in_seq_cp_pair_batch()` when `get_cp_shared_kv_batch_plan(forward_batch) is not None`, even if metadata `batch_size == 1`. +- FP8 prefill with CP forces NSA implementation to `flashmla_sparse` and topk transform method to `RAGGED`. With `SGLANG_NSA_FUSE_TOPK=1`, `forward_extend()` uses `topk_indices` directly as `page_table_1`; there is no later `topk_indices_offset` repair. Therefore any batch topk coordinate-space mismatch can directly corrupt attention. +- For no-cache current-only prefill, `forward_mla.forward_absorb_prepare()` may rebuild full current `k_nope/k_pe` after CP direct-write because `should_reuse_current_extend_kv(forward_batch)` is true. Attention then consumes the rebuilt full-current KV tensor, so the topk/index coordinate contract must match that rebuilt tensor layout exactly. +- Existing unit coverage checks batch descriptor shapes/lengths, but does not yet prove bs=1 scalar-path and batch-plan-path CP all-gather/rerange outputs are identical, nor that batch topk transform emits indices in the same coordinate space consumed by `flashmla_sparse`. + +Next focused audit: prove or disprove that `cp_all_gather_rerange_output()` batch-plan path is bitwise-equivalent to the old scalar CP rerange for bs=1. This is high-signal because both MLA current reuse and index current reuse depend on the rebuilt current tensor and it would explain BF16+FP8 garbage without cache hit. + + +## 2026-06-05 RED: batch cp_index RAGGED topk loses request-coordinate offset + +New focused unit test on remote `g0034` fails before the fix: + +```text +test_indexer_ragged_cp_index_batch_uses_request_ragged_offsets +AssertionError: topk_indices_offset_override is None +``` + +Root-cause evidence: + +- `fast_topk_transform_ragged_fused` returns topk indices as local indices plus `topk_indices_offset`; `row_starts` only selects the score window. +- The bs>1 cp_index path compacts each CP segment's K rows into a temporary MQA buffer and sets `row_starts/ks` to compact K bases. +- However attention consumes normal ragged KV coordinates (per-request flattened KV), not the temporary compact segment-coordinate space. +- The current cp_index path calls `metadata.topk_transform(..., cu_seqlens_q=actual_seq_q, batch_idx_list=...)` without `topk_indices_offset_override`, so RAGGED topk derives offsets from compact query cu-seqlens. That is neither request KV base nor compact K base. +- This path is reached for bs=1 after bs>1 metadata attaches `batch_plan`, so it can corrupt the first no-cache prefill and explains garbage output/low accept len without requiring a cache hit. + +Planned minimal fix: build a per-compact-query `topk_indices_offset_override` from cumulative `seq_lens_cpu` request KV bases and pass it to all cp_index RAGGED `topk_transform` calls. Do not disable current reuse or batch planning. + + +## 2026-06-05 GREEN: cp_index RAGGED topk offset fixed + +Implemented fix: + +- In `_get_topk_ragged_with_cp(cp_index=...)`, compute per-compact-query ragged topk offsets from cumulative `forward_batch.seq_lens_cpu` request KV bases. +- Pass `topk_indices_offset_override` to all cp_index `metadata.topk_transform()` calls: TAI shared batch, TAI current batch, and Python fallback. +- This keeps `row_starts/ks` in temporary compact-MQA-buffer coordinates while final topk indices stay in the normal ragged KV coordinate space consumed by `flashmla_sparse`. + +Verification: + +- RED test failed before fix on remote: `topk_indices_offset_override is None`. +- After fix, remote `g0034` passed: + - targeted RED test: `1 passed`; + - related indexer cp_index tests: `6 passed`; + - full `test_nsa_cp_utils.py`: `91 passed`; + - full `test_cp_shared_kv_runtime.py`: `111 passed, 2 subtests passed`. + +Remaining runtime validation: + +- Need user restart/run real traffic to confirm output text and EAGLE accept length recover. This fix addresses a concrete coordinate corruption in the bs>1-era path that also affected bs=1; it does not prove all true bs>1 cases are complete. diff --git a/python/sglang/srt/layers/attention/nsa/nsa_indexer.py b/python/sglang/srt/layers/attention/nsa/nsa_indexer.py index 3964a19d1..486efc92a 100644 --- a/python/sglang/srt/layers/attention/nsa/nsa_indexer.py +++ b/python/sglang/srt/layers/attention/nsa/nsa_indexer.py @@ -1416,6 +1416,12 @@ class Indexer(MultiPlatformOp): q_lens_list = [] k_bases_list = [] q_bases_list = [] + topk_offset_list = [] + request_kv_bases: List[int] = [] + request_kv_base = 0 + for seq_len in forward_batch.seq_lens_cpu.tolist(): + request_kv_bases.append(int(request_kv_base)) + request_kv_base += int(seq_len) k_cursor = 0 q_cursor = 0 for raw_batch_idx, start_seq_position, end_seq_position in cp_index: @@ -1453,9 +1459,22 @@ class Indexer(MultiPlatformOp): q_lens_list.append(extend_seq_len) k_bases_list.append(k_cursor) q_bases_list.append(q_cursor) + if batch_idx < 0 or batch_idx >= len(request_kv_bases): + raise RuntimeError( + "[CP_SHARED_KV_FAIL_FAST][index_topk] " + "reason=batch_cp_index_bad_batch_idx " + f"batch_idx={batch_idx} seq_lens={forward_batch.seq_lens_cpu.tolist()}" + ) + topk_offset_list.extend( + [request_kv_bases[batch_idx]] * extend_seq_len + ) k_cursor += kv_len_i q_cursor += extend_seq_len + topk_indices_offset_override = torch.tensor( + topk_offset_list, dtype=torch.int32, device=q_fp8.device + ) + if current_index_kv is None: assert index_buffer is not None assert block_tables is not None @@ -1512,6 +1531,7 @@ class Indexer(MultiPlatformOp): cu_seqlens_q=actual_seq_q, ke_offset=ke_offset, batch_idx_list=batch_idx_list, + topk_indices_offset_override=topk_indices_offset_override, ) return topk_result else: @@ -1576,6 +1596,7 @@ class Indexer(MultiPlatformOp): cu_seqlens_q=actual_seq_q, ke_offset=ke_offset, batch_idx_list=batch_idx_list, + topk_indices_offset_override=topk_indices_offset_override, ) return topk_result @@ -1670,6 +1691,7 @@ class Indexer(MultiPlatformOp): cu_seqlens_q=actual_seq_q, ke_offset=ke_offset, batch_idx_list=batch_idx_list, + topk_indices_offset_override=topk_indices_offset_override, ) else: seq_len = int(forward_batch.seq_lens_cpu[batch_idx].item()) diff --git a/python/sglang/srt/layers/attention/nsa/utils.py b/python/sglang/srt/layers/attention/nsa/utils.py index c073c7a0c..cf6eea7f0 100644 --- a/python/sglang/srt/layers/attention/nsa/utils.py +++ b/python/sglang/srt/layers/attention/nsa/utils.py @@ -610,15 +610,28 @@ def build_batch_page_aligned_in_seq_split_plan( f"extend_len={extend_len} page_size={page_size}" ) - compute_pages = max(split_info.extend_padded_pages, cp_size) - compute_split_list, compute_page_starts, compute_page_ends = ( - _build_full_page_unit_split( - page_units=compute_pages, - extend_prefix_len=prefix_len, - page_size=page_size, - cp_size=cp_size, + # Cache residency is page-granular, but attention/index compute should + # not consume synthetic rows for an ordinary final partial page. Only + # tiny suffixes with fewer physical pages than CP ranks need compute + # padding; otherwise the valid split is already sufficient and keeps + # the pre-bs>1 bs=1 semantics. + needs_compute_padding = split_info.extend_padded_pages < cp_size + if needs_compute_padding: + compute_pages = cp_size + compute_split_list, compute_page_starts, compute_page_ends = ( + _build_full_page_unit_split( + page_units=compute_pages, + extend_prefix_len=prefix_len, + page_size=page_size, + cp_size=cp_size, + ) ) - ) + else: + compute_pages = split_info.extend_padded_pages + compute_split_list = split_list + compute_page_starts = split_info.segment_page_starts + compute_page_ends = split_info.segment_page_ends + compute_tokens = sum(int(token_count) for token_count in compute_split_list) owner, local_offset = _get_in_seq_last_token_owner_and_offset( split_list=compute_split_list, cp_size=cp_size, @@ -671,8 +684,8 @@ def build_batch_page_aligned_in_seq_split_plan( request_compute_segment_page_starts.append(compute_page_starts) request_compute_segment_page_ends.append(compute_page_ends) request_compute_padded_pages.append(compute_pages) - request_compute_padded_tokens.append(compute_pages * page_size) - request_compute_padding_tokens.append(compute_pages * page_size - extend_len) + request_compute_padded_tokens.append(compute_tokens) + request_compute_padding_tokens.append(compute_tokens - extend_len) request_compute_rank_local_tokens.append(compute_rank_local_tokens) request_compute_actual_seq_q_prev.append(compute_split_list[cp_rank]) request_compute_actual_seq_q_next.append(compute_split_list[mirror_idx]) diff --git a/test/registered/unit/layers/test_nsa_cp_utils.py b/test/registered/unit/layers/test_nsa_cp_utils.py index 8ecc11b81..fe32030ed 100644 --- a/test/registered/unit/layers/test_nsa_cp_utils.py +++ b/test/registered/unit/layers/test_nsa_cp_utils.py @@ -586,18 +586,17 @@ class TestNSAInSeqCPUtils(unittest.TestCase): plan.request_compute_seq_q_prev[0] + plan.request_compute_seq_q_next[0] ) self.assertEqual(valid_local_rows, 4995) - self.assertEqual(compute_local_rows, 5056) - self.assertTrue(plan.compute_padding_enabled) + self.assertEqual(compute_local_rows, valid_local_rows) + self.assertFalse(plan.compute_padding_enabled) - local_compute = split_tensor_by_cp_batch_plan( + local = split_tensor_by_cp_batch_plan( torch.arange(40387, dtype=torch.int64), plan, mode="1d", ) - self.assertEqual(local_compute.numel(), compute_local_rows) - self.assertEqual(local_compute[:2560].tolist(), list(range(2560))) - self.assertEqual(local_compute[2560:4995].tolist(), list(range(37952, 40387))) - self.assertEqual(local_compute[4995:].tolist(), [0] * 61) + self.assertEqual(local.numel(), valid_local_rows) + self.assertEqual(local[:2560].tolist(), list(range(2560))) + self.assertEqual(local[2560:4995].tolist(), list(range(37952, 40387))) local_valid = select_cp_local_valid_rows_for_cache_write( SimpleNamespace( @@ -606,10 +605,10 @@ class TestNSAInSeqCPUtils(unittest.TestCase): batch_plan=plan, ) ), - local_compute, + local, ) self.assertEqual(local_valid.numel(), valid_local_rows) - self.assertEqual(local_valid.tolist(), local_compute[:valid_local_rows].tolist()) + self.assertEqual(local_valid.tolist(), local.tolist()) selected = _select_batch_topk_query_lengths( cp_metadata=NSAContextParallelMetadata(batch_size=1, batch_plan=plan), @@ -629,7 +628,7 @@ class TestNSAInSeqCPUtils(unittest.TestCase): selected.request_valid_seq_q_next, plan.request_valid_seq_q_next ) - selected_compute = _select_batch_topk_query_lengths( + selected_compute_alias = _select_batch_topk_query_lengths( cp_metadata=NSAContextParallelMetadata(batch_size=1, batch_plan=plan), batch_plan=plan, batch_size=1, @@ -637,20 +636,79 @@ class TestNSAInSeqCPUtils(unittest.TestCase): weights_tokens=compute_local_rows, ) - self.assertTrue(selected_compute.uses_compute_query_rows) + self.assertFalse(selected_compute_alias.uses_compute_query_rows) self.assertEqual( - selected_compute.request_seq_q_prev, plan.request_compute_seq_q_prev + selected_compute_alias.request_seq_q_prev, plan.request_compute_seq_q_prev ) self.assertEqual( - selected_compute.request_seq_q_next, plan.request_compute_seq_q_next + selected_compute_alias.request_seq_q_next, plan.request_compute_seq_q_next ) self.assertEqual( - selected_compute.request_valid_seq_q_prev, plan.request_valid_seq_q_prev + selected_compute_alias.request_valid_seq_q_prev, plan.request_valid_seq_q_prev ) self.assertEqual( - selected_compute.request_valid_seq_q_next, plan.request_valid_seq_q_next + selected_compute_alias.request_valid_seq_q_next, plan.request_valid_seq_q_next ) + def test_batch_plan_keeps_long_page_tail_out_of_compute_padding(self): + plan = build_batch_page_aligned_in_seq_split_plan( + extend_lens=[40387], + prefix_lens=[0], + page_size=64, + cp_size=8, + cp_rank=0, + ) + + self.assertFalse(plan.compute_padding_enabled) + self.assertEqual(plan.request_valid_padded_pages, [632]) + self.assertEqual(plan.request_valid_padded_tokens, [40448]) + self.assertEqual(plan.request_compute_padded_pages, [632]) + self.assertEqual(plan.request_compute_padded_tokens, [40387]) + self.assertEqual(plan.request_compute_padding_tokens, [0]) + self.assertEqual( + plan.request_compute_split_lists, + plan.request_valid_split_lists, + ) + + def test_batch_plan_compute_padding_only_pads_tiny_request_in_mixed_batch(self): + import torch + + plan = build_batch_page_aligned_in_seq_split_plan( + extend_lens=[65, 40387], + prefix_lens=[40320, 0], + page_size=64, + cp_size=8, + cp_rank=1, + ) + + self.assertTrue(plan.compute_padding_enabled) + self.assertEqual(plan.request_valid_padded_pages, [2, 632]) + self.assertEqual(plan.request_compute_padded_pages, [8, 632]) + self.assertEqual(plan.request_compute_padded_tokens, [512, 40387]) + self.assertEqual(plan.request_compute_padding_tokens, [447, 0]) + self.assertEqual( + plan.request_compute_split_lists[1], + plan.request_valid_split_lists[1], + ) + + local = split_tensor_by_cp_batch_plan( + torch.arange(65 + 40387, dtype=torch.int64), + plan, + mode="1d", + ) + self.assertEqual(local.numel(), sum(plan.request_compute_rank_local_tokens)) + + valid = select_cp_local_valid_rows_for_cache_write( + SimpleNamespace( + nsa_cp_metadata=NSAContextParallelMetadata( + batch_size=2, + batch_plan=plan, + ) + ), + local, + ) + self.assertEqual(valid.numel(), sum(plan.request_valid_rank_local_tokens)) + def test_batch_plan_compute_padding_is_per_request_not_batch_total(self): plan = build_batch_page_aligned_in_seq_split_plan( extend_lens=[65, 1024], @@ -687,14 +745,14 @@ class TestNSAInSeqCPUtils(unittest.TestCase): self.assertEqual(build_flat_page_owner_plan(plan), [0, 0, 1, 1]) local_1d = split_tensor_by_cp_batch_plan(torch.arange(13), plan, mode="1d") - self.assertEqual(local_1d.tolist(), [0, 0, 0, 0, 8, 9, 10, 11, 12, 0, 0, 0]) + self.assertEqual(local_1d.tolist(), [0, 0, 0, 0, 8, 9, 10, 11, 12]) local_data = split_tensor_by_cp_batch_plan( torch.arange(13 * 2).view(13, 2), plan, mode="data" ) self.assertEqual( local_data[:, 0].tolist(), - [0, 0, 0, 0, 16, 18, 20, 22, 24, 0, 0, 0], + [0, 0, 0, 0, 16, 18, 20, 22, 24], ) def test_collect_last_token_hidden_uses_batch_owner_metadata(self): @@ -3227,6 +3285,83 @@ class TestNSAInSeqCPUtils(unittest.TestCase): self.assertEqual(deep_gemm_calls[0]["ks"], [0, 0, 3, 6, 10, 10, 10]) self.assertEqual(deep_gemm_calls[0]["ke"], [2, 3, 6, 10, 12, 13, 14]) + def test_indexer_ragged_cp_index_batch_uses_request_ragged_offsets(self): + import contextlib + import torch + from types import SimpleNamespace + from unittest.mock import patch + + from sglang.srt.layers.attention.nsa import nsa_indexer + from sglang.srt.layers.attention.nsa.nsa_indexer import Indexer + + indexer = object.__new__(Indexer) + indexer.index_topk = 2 + indexer._with_real_sm_count = lambda: contextlib.nullcontext() + + topk_kwargs = [] + + def fake_prepare(**kwargs): + total_kv_len = int(kwargs["total_kv_len"]) + return ( + torch.zeros((total_kv_len, 1), dtype=torch.uint8), + torch.zeros((total_kv_len,), dtype=torch.float32), + torch.tensor([0, 0, 3, 6, 10, 10, 10], dtype=torch.int32), + torch.tensor([2, 3, 3, 4, 2, 3, 4], dtype=torch.int32), + ) + + def fake_logits(q_fp8, kv_fp8, weights, ks, ke, clean_logits=False): + return torch.zeros((int(q_fp8.shape[0]), 8), dtype=torch.float32) + + class Metadata: + def get_page_table_64(self): + raise AssertionError("current cp_index path must not materialize index pages") + + def topk_transform(self, logits, topk, **kwargs): + topk_kwargs.append(kwargs) + return torch.zeros((int(logits.shape[0]), topk), dtype=torch.int32) + + forward_batch = SimpleNamespace( + token_to_kv_pool=SimpleNamespace(page_size=64, index_head_dim=1), + seq_lens_cpu=torch.tensor([3, 4], dtype=torch.int64), + extend_seq_lens_cpu=[3, 4], + ) + current_index_kv = ( + torch.arange(7, dtype=torch.uint8).view(7, 1), + torch.arange(7, dtype=torch.float32).view(7, 1), + ) + + with patch.object( + nsa_indexer, + "try_tai_prepare_cp_mqa_current_index_batch", + side_effect=fake_prepare, + create=True, + ), patch.object( + nsa_indexer, + "deep_gemm", + SimpleNamespace(fp8_mqa_logits=fake_logits), + ): + Indexer._get_topk_ragged_with_cp( + indexer, + forward_batch, + layer_id=7, + q_fp8=torch.empty((7, 1), dtype=torch.float32), + weights=torch.empty((7, 1, 1), dtype=torch.float32), + metadata=Metadata(), + kv_len=0, + actual_seq_q=7, + cp_index=[(0, 1, 3), (0, 2, 3), (1, 3, 4), (1, 1, 4)], + current_index_kv=current_index_kv, + ) + + self.assertEqual(len(topk_kwargs), 1) + offset = topk_kwargs[0].get("topk_indices_offset_override") + self.assertIsNotNone(offset) + # Batch cp_index compacts each CP segment's K into a temporary buffer, + # but flashmla_sparse consumes the normal ragged KV layout. The fused + # ragged topk offset therefore has to stay in request/KV coordinates, + # not compact segment coordinates or compact-q cu-seqlens. + self.assertEqual(offset.tolist(), [0, 0, 0, 3, 3, 3, 3]) + def test_indexer_ragged_cp_index_shared_batch_uses_tai_prepare_once(self): import contextlib import torch diff --git a/test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py b/test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py index 58e6d46c0..f90231e41 100644 --- a/test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py +++ b/test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py @@ -1044,10 +1044,10 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase): extend_seq_lens_cpu=[40387], seq_lens_cpu=torch.tensor([40387], dtype=torch.int32), out_cache_loc=torch.arange(40392, dtype=torch.int64), - cp_local_out_cache_loc=torch.arange(5056, dtype=torch.int64), + cp_local_out_cache_loc=torch.arange(4995, dtype=torch.int64), ) - local_k = torch.empty((5056, 1), dtype=torch.float32) - local_rope = torch.empty((5056, 1), dtype=torch.float32) + local_k = torch.empty((4995, 1), dtype=torch.float32) + local_rope = torch.empty((4995, 1), dtype=torch.float32) with envs.SGLANG_CP_SHARED_KV_CURRENT_REUSE.override(True): self.assertTrue(runtime.should_reuse_current_extend_kv(forward_batch)) @@ -1057,12 +1057,12 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase): local_k, local_rope, ), - 5056, + 4995, ) self.assertIsNone( runtime.current_extend_kv_rows_for_reuse( forward_batch, - local_k[:5055], + local_k[:4994], local_rope, ) )