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, ) )