import ast import os from pathlib import Path import unittest import sys from types import SimpleNamespace from unittest.mock import patch from sglang.srt.environ import envs from sglang.srt.layers.attention.nsa import utils as nsa_utils from sglang.srt.layers.attention.nsa.utils import ( NSAContextParallelMetadata, PageAlignedCacheExtent, build_flat_page_owner_plan, build_batch_page_aligned_in_seq_split_plan, build_page_aligned_cache_extent, _get_in_seq_last_token_owner_and_offset, _build_batch_metadata_from_plan, build_page_aligned_in_seq_split_list, build_token_balanced_in_seq_split_list, can_cp_split, cp_all_gather_rerange_output, cp_collect_last_token_hidden, cp_split_and_rebuild_1d, cp_split_and_rebuild_data, cp_split_and_rebuild_position, _torch_batch_in_seq_all_gather_rerange, get_cp_shared_kv_batch_plan, get_cp_shared_kv_local_out_cache_loc, get_cp_shared_kv_local_physical_out_cache_loc, get_cp_local_embedding_padded_token_count, nsa_use_prefill_cp, pad_cp_local_input_ids_for_embedding, prepare_input_dp_with_cp_dsa, select_cp_current_valid_rows_for_reuse, select_cp_local_valid_rows_for_cache_write, split_tensor_by_cp_batch_plan, split_in_seq_cp_local_pair, ) from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode from sglang.srt.models.deepseek_nextn import DeepseekModelNextN from sglang.srt.speculative.eagle_info import EagleDraftInput from sglang.test.ci.ci_register import register_cpu_ci register_cpu_ci(est_time=1, suite="stage-a-test-cpu") class TestPageAlignedCacheExtent(unittest.TestCase): def test_extent_uses_page_boundary_not_cp_size(self): extent = build_page_aligned_cache_extent(valid_tokens=100, page_size=64) self.assertEqual(extent.valid_tokens, 100) self.assertEqual(extent.padded_pages, 2) self.assertEqual(extent.padded_tokens, 128) self.assertEqual(extent.padding_tokens, 28) def test_extent_handles_empty_and_aligned_lengths(self): self.assertEqual( build_page_aligned_cache_extent(valid_tokens=0, page_size=64), PageAlignedCacheExtent( valid_tokens=0, padded_pages=0, padded_tokens=0, padding_tokens=0, ), ) self.assertEqual( build_page_aligned_cache_extent(valid_tokens=128, page_size=64), PageAlignedCacheExtent( valid_tokens=128, padded_pages=2, padded_tokens=128, padding_tokens=0, ), ) class TestNSAInSeqCPUtils(unittest.TestCase): def test_contiguous_valid_cp_query_count(self): from sglang.srt.layers.attention.nsa.nsa_indexer import ( _compute_contiguous_valid_cp_query_count, ) self.assertEqual( _compute_contiguous_valid_cp_query_count( cp_kv_end=1024, actual_seq_q=128, logical_kv_limit=1024, ), 128, ) self.assertEqual( _compute_contiguous_valid_cp_query_count( cp_kv_end=1100, actual_seq_q=128, logical_kv_limit=1024, ), 52, ) self.assertEqual( _compute_contiguous_valid_cp_query_count( cp_kv_end=1100, actual_seq_q=64, logical_kv_limit=1000, ), 0, ) self.assertEqual( _compute_contiguous_valid_cp_query_count( cp_kv_end=100, actual_seq_q=0, logical_kv_limit=100, ), 0, ) def assert_page_aligned_boundaries( self, split_list, *, extend_prefix_len, extend_len, page_size ): cursor = 0 for segment_len in split_list[:-1]: cursor += segment_len if cursor < extend_len: self.assertEqual((extend_prefix_len + cursor) % page_size, 0) def test_page_aligned_split_keeps_boundaries_on_pages(self): split_list, info = build_page_aligned_in_seq_split_list( total_len=32768, extend_len=32768, extend_prefix_len=0, page_size=64, cp_size=8, ) self.assertTrue(info.page_aligned) self.assertEqual(sum(split_list), 32768) self.assertEqual(len(split_list), 16) self.assertTrue(all(segment_len > 0 for segment_len in split_list)) self.assert_page_aligned_boundaries( split_list, extend_prefix_len=0, extend_len=32768, page_size=64 ) def test_page_aligned_split_uses_prefix_for_boundary_alignment(self): split_list, info = build_page_aligned_in_seq_split_list( total_len=1024, extend_len=1024, extend_prefix_len=128, page_size=64, cp_size=8, ) self.assertTrue(info.page_aligned) self.assertEqual(sum(split_list), 1024) self.assert_page_aligned_boundaries( split_list, extend_prefix_len=128, extend_len=1024, page_size=64 ) def test_page_aligned_split_keeps_tail_partial_page_unsplit(self): split_list, info = build_page_aligned_in_seq_split_list( total_len=1100, extend_len=1100, extend_prefix_len=0, page_size=64, cp_size=8, ) self.assertTrue(info.page_aligned) self.assertEqual(sum(split_list), 1100) self.assertEqual(split_list[-1], 12) self.assert_page_aligned_boundaries( split_list, extend_prefix_len=0, extend_len=1100, page_size=64 ) def test_page_aligned_split_exposes_padded_extent_without_padding_split_list(self): split_list, info = build_page_aligned_in_seq_split_list( total_len=100, extend_len=100, extend_prefix_len=0, page_size=64, cp_size=8, ) self.assertTrue(info.page_aligned) self.assertEqual(sum(split_list), 100) self.assertEqual(split_list[:2], [64, 36]) self.assertEqual(split_list[2:], [0] * 14) self.assertEqual(info.extend_valid_tokens, 100) self.assertEqual(info.extend_padded_pages, 2) self.assertEqual(info.extend_padded_tokens, 128) self.assertEqual(info.extend_padding_tokens, 28) def test_page_aligned_split_falls_back_when_prefix_is_not_page_aligned(self): split_list, info = build_page_aligned_in_seq_split_list( total_len=1024, extend_len=1024, extend_prefix_len=1, page_size=64, cp_size=8, ) self.assertFalse(info.page_aligned) self.assertEqual(split_list, build_token_balanced_in_seq_split_list(1024, 8)) def test_page_aligned_split_pads_zero_segments_when_page_units_are_short(self): split_list, info = build_page_aligned_in_seq_split_list( total_len=512, extend_len=512, extend_prefix_len=0, page_size=64, cp_size=8, ) self.assertTrue(info.page_aligned) self.assertEqual(sum(split_list), 512) self.assertEqual(split_list[:8], [64] * 8) self.assertEqual(split_list[8:], [0] * 8) self.assert_page_aligned_boundaries( split_list, extend_prefix_len=0, extend_len=512, page_size=64 ) def test_page_aligned_split_allows_radix_hit_suffix_with_one_page_per_rank(self): split_list, info = build_page_aligned_in_seq_split_list( total_len=512, extend_len=512, extend_prefix_len=54464, page_size=64, cp_size=8, ) self.assertTrue(info.page_aligned) self.assertEqual(sum(split_list), 512) self.assertEqual(split_list[:8], [64] * 8) self.assertEqual(split_list[8:], [0] * 8) self.assert_page_aligned_boundaries( split_list, extend_prefix_len=54464, extend_len=512, page_size=64 ) def test_page_aligned_split_keeps_short_radix_hit_suffix_page_aligned(self): split_list, info = build_page_aligned_in_seq_split_list( total_len=256, extend_len=256, extend_prefix_len=54464, page_size=64, cp_size=8, ) self.assertTrue(info.page_aligned) self.assertEqual(sum(split_list), 256) self.assertEqual(split_list[:4], [64] * 4) self.assertEqual(split_list[4:], [0] * 12) self.assert_page_aligned_boundaries( split_list, extend_prefix_len=54464, extend_len=256, page_size=64 ) def test_can_cp_split_uses_compute_padding_for_short_radix_hit_suffix(self): class Mode: def is_context_parallel_extend(self): return True forward_batch = SimpleNamespace( uses_cp_shared_kv=True, extend_seq_lens_cpu=[65], extend_prefix_lens_cpu=[54464], token_to_kv_pool=SimpleNamespace(page_size=64), forward_mode=Mode(), ) with ( patch( "sglang.srt.layers.attention.nsa.utils.is_nsa_prefill_cp_round_robin_split", return_value=False, ), patch( "sglang.srt.layers.attention.nsa.utils.is_nsa_enable_prefill_cp", return_value=True, ), ): self.assertTrue(can_cp_split(128, 8, True, forward_batch)) def test_can_cp_split_uses_compute_padding_when_page_units_do_not_cover_all_lanes( self, ): class Mode: def is_context_parallel_extend(self): return True forward_batch = SimpleNamespace( uses_cp_shared_kv=True, extend_seq_lens_cpu=[256], extend_prefix_lens_cpu=[54464], token_to_kv_pool=SimpleNamespace(page_size=64), forward_mode=Mode(), ) with ( patch( "sglang.srt.layers.attention.nsa.utils.is_nsa_prefill_cp_round_robin_split", return_value=False, ), patch( "sglang.srt.layers.attention.nsa.utils.is_nsa_enable_prefill_cp", return_value=True, ), ): self.assertTrue(can_cp_split(256, 8, True, forward_batch)) def test_can_cp_split_uses_compute_padding_for_current_only_one_page_suffix( self, ): class Mode: def is_context_parallel_extend(self): return True forward_batch = SimpleNamespace( uses_cp_shared_kv=True, extend_seq_lens_cpu=[64], extend_prefix_lens_cpu=[0], token_to_kv_pool=SimpleNamespace(page_size=64), forward_mode=Mode(), ) with ( patch( "sglang.srt.layers.attention.nsa.utils.is_nsa_prefill_cp_round_robin_split", return_value=False, ), patch( "sglang.srt.layers.attention.nsa.utils.is_nsa_enable_prefill_cp", return_value=True, ), ): self.assertTrue(can_cp_split(64, 8, True, forward_batch)) def test_can_cp_split_enables_cp_draft_shared_kv_draft_extend(self): class DraftMode: def is_context_parallel_extend(self, include_draft_extend_v2=False): return False def is_draft_extend(self, include_v2=False): return True forward_batch = SimpleNamespace( uses_cp_shared_kv=True, extend_seq_lens_cpu=[4], extend_prefix_lens_cpu=[0], token_to_kv_pool=SimpleNamespace(page_size=64), forward_mode=DraftMode(), ) with ( patch.dict(os.environ, {"SGLANG_CP_DRAFT_SHARED_KV": "1"}), patch( "sglang.srt.layers.attention.nsa.utils.is_nsa_prefill_cp_round_robin_split", return_value=False, ), patch( "sglang.srt.layers.attention.nsa.utils.is_nsa_enable_prefill_cp", return_value=True, ), ): self.assertTrue(can_cp_split(8, 8, True, forward_batch)) def test_nsa_use_prefill_cp_enables_cp_draft_shared_kv_draft_extend(self): forward_batch = SimpleNamespace( uses_cp_shared_kv=True, forward_mode=ForwardMode.DRAFT_EXTEND, nsa_cp_metadata=NSAContextParallelMetadata(batch_size=1), ) with patch.dict(os.environ, {"SGLANG_CP_DRAFT_SHARED_KV": "1"}): self.assertTrue(nsa_use_prefill_cp(forward_batch, True)) def test_can_cp_split_uses_compute_padding_per_request_for_batched_tiny_suffix( self, ): class Mode: def is_context_parallel_extend(self): return True forward_batch = SimpleNamespace( uses_cp_shared_kv=True, extend_seq_lens_cpu=[65, 64], extend_prefix_lens_cpu=[54464, 8192], token_to_kv_pool=SimpleNamespace(page_size=64), forward_mode=Mode(), ) with ( patch( "sglang.srt.layers.attention.nsa.utils.is_nsa_prefill_cp_round_robin_split", return_value=False, ), patch( "sglang.srt.layers.attention.nsa.utils.is_nsa_enable_prefill_cp", return_value=True, ), ): self.assertTrue(can_cp_split(129, 8, True, forward_batch)) def test_can_cp_split_fails_on_non_page_aligned_cp_shared_prefix(self): class Mode: def is_context_parallel_extend(self): return True forward_batch = SimpleNamespace( uses_cp_shared_kv=True, extend_seq_lens_cpu=[1024], extend_prefix_lens_cpu=[65], token_to_kv_pool=SimpleNamespace(page_size=64), forward_mode=Mode(), ) with ( patch( "sglang.srt.layers.attention.nsa.utils.is_nsa_prefill_cp_round_robin_split", return_value=False, ), patch( "sglang.srt.layers.attention.nsa.utils.is_nsa_enable_prefill_cp", return_value=True, ), self.assertRaisesRegex( RuntimeError, r"\[CP_SHARED_KV_FAIL_FAST\]\[cp_split_non_page_aligned_prefix\]", ), ): can_cp_split(1089, 8, True, forward_batch) def test_can_cp_split_skips_page_plan_validator_for_target_verify(self): forward_batch = SimpleNamespace( uses_cp_shared_kv=True, extend_seq_lens_cpu=[64], extend_prefix_lens_cpu=None, token_to_kv_pool=SimpleNamespace(page_size=64), forward_mode=ForwardMode.TARGET_VERIFY, ) with ( patch( "sglang.srt.layers.attention.nsa.utils.is_nsa_prefill_cp_round_robin_split", return_value=False, ), patch( "sglang.srt.layers.attention.nsa.utils.is_nsa_enable_prefill_cp", return_value=True, ), ): self.assertFalse(can_cp_split(64, 8, True, forward_batch)) def test_can_cp_split_keeps_cp_for_radix_hit_suffix_with_one_page_per_rank(self): class Mode: def is_context_parallel_extend(self): return True forward_batch = SimpleNamespace( uses_cp_shared_kv=True, extend_seq_lens_cpu=[512], extend_prefix_lens_cpu=[54464], token_to_kv_pool=SimpleNamespace(page_size=64), forward_mode=Mode(), ) with ( patch( "sglang.srt.layers.attention.nsa.utils.is_nsa_prefill_cp_round_robin_split", return_value=False, ), patch( "sglang.srt.layers.attention.nsa.utils.is_nsa_enable_prefill_cp", return_value=True, ), ): self.assertTrue(can_cp_split(512, 8, True, forward_batch)) def test_page_aligned_split_adds_padding_tokens_to_last_segment(self): split_list, info = build_page_aligned_in_seq_split_list( total_len=1040, extend_len=1024, extend_prefix_len=0, page_size=64, cp_size=8, ) self.assertTrue(info.page_aligned) self.assertEqual(sum(split_list), 1040) self.assertEqual(split_list[-1], 80) self.assert_page_aligned_boundaries( split_list, extend_prefix_len=0, extend_len=1024, page_size=64 ) def test_last_token_owner_uses_actual_token_count_when_batch_is_padded(self): # Padded prefill can have 64 model tokens while the real prompt has only # 11 tokens. In in-seq split with cp_size=8, the real last token is in # segment 2, not in rank 0's trailing padded segment. split_list = [4] * 16 owner, local_offset = _get_in_seq_last_token_owner_and_offset( split_list=split_list, cp_size=8, actual_token_count=11, ) self.assertEqual(owner, 2) self.assertEqual(local_offset, 2) def test_last_token_owner_keeps_existing_unpadded_fast_path_location(self): split_list = [4] * 16 owner, local_offset = _get_in_seq_last_token_owner_and_offset( split_list=split_list, cp_size=8, actual_token_count=64, ) self.assertEqual(owner, 0) self.assertEqual(local_offset, 7) def test_batch_page_aligned_plan_keeps_request_boundaries_and_last_token_owners( self, ): plan = build_batch_page_aligned_in_seq_split_plan( extend_lens=[4, 9], prefix_lens=[0, 8], page_size=4, cp_size=2, cp_rank=0, ) self.assertEqual(plan.batch_size, 2) self.assertEqual(plan.request_split_lists, [[4, 0, 0, 0], [4, 4, 1, 0]]) self.assertEqual(plan.request_padded_pages, [1, 3]) self.assertEqual(plan.request_padded_tokens, [4, 12]) self.assertEqual(plan.request_token_offsets, [0, 4]) self.assertEqual(plan.request_padded_token_offsets, [0, 4]) self.assertEqual(plan.request_page_offsets, [0, 1]) self.assertEqual(plan.request_last_token_owner, [0, 1]) self.assertEqual(plan.request_last_token_local_offset, [3, 4]) self.assertEqual(plan.request_rank_local_tokens, [4, 4]) self.assertEqual(plan.request_rank_local_offsets, [0, 4]) self.assertEqual(plan.request_kv_len_prev, [4, 4]) self.assertEqual(plan.request_kv_len_next, [4, 9]) self.assertEqual(plan.request_actual_seq_q_prev, [4, 4]) self.assertEqual(plan.request_actual_seq_q_next, [0, 0]) self.assertEqual(plan.flat_segment_request_ids, [0, 0, 0, 0, 1, 1, 1, 1]) self.assertEqual(plan.flat_segment_offsets, [0, 4, 4, 4, 0, 4, 8, 9]) def test_batch_plan_exposes_compute_padding_without_inflating_valid_cache_extent( self, ): plan = build_batch_page_aligned_in_seq_split_plan( extend_lens=[65], prefix_lens=[40320], page_size=64, cp_size=8, cp_rank=1, ) self.assertTrue(plan.compute_padding_enabled) self.assertEqual( plan.request_valid_split_lists, [[64, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]], ) self.assertEqual( plan.request_compute_split_lists, [[64, 64, 64, 64, 64, 64, 64, 64, 0, 0, 0, 0, 0, 0, 0, 0]], ) self.assertEqual(plan.request_valid_padded_pages, [2]) self.assertEqual(plan.request_valid_padded_tokens, [128]) self.assertEqual(plan.request_compute_padded_pages, [8]) self.assertEqual(plan.request_compute_padded_tokens, [512]) self.assertEqual(plan.request_compute_padding_tokens, [447]) self.assertEqual(plan.request_compute_rank_local_tokens, [64]) self.assertEqual(plan.request_compute_rank_local_offsets, [0]) self.assertEqual(plan.request_valid_rank_local_tokens, [1]) self.assertEqual(plan.request_valid_rank_local_offsets, [0]) self.assertEqual(plan.request_last_token_owner, [1]) self.assertEqual(plan.request_last_token_local_offset, [0]) # Compatibility aliases for cache/page accounting stay valid-token # based. Query-length metadata exposes both valid and compute rows: # consumers must choose the view that matches their actual q layout. self.assertEqual(plan.request_split_lists, plan.request_valid_split_lists) self.assertEqual(plan.request_padded_pages, plan.request_valid_padded_pages) self.assertEqual(plan.request_actual_seq_q_prev, [64]) self.assertEqual(plan.request_actual_seq_q_next, [0]) self.assertEqual(plan.request_valid_seq_q_prev, [1]) self.assertEqual(plan.request_valid_seq_q_next, [0]) self.assertEqual(plan.request_compute_seq_q_prev, [64]) self.assertEqual(plan.request_compute_seq_q_next, [0]) def test_bs_gt1_debug_log_is_env_gated_and_limited(self): nsa_utils._CP_SHARED_KV_BS_GT1_DEBUG_COUNTS.clear() with envs.SGLANG_CP_SHARED_KV_BS_GT1_DEBUG.override(False): with patch.object(nsa_utils.logger, "info") as info: nsa_utils.log_cp_shared_kv_bs_gt1_debug( "unit_test", "bs=%s", 2, ) self.assertFalse(info.called) with envs.SGLANG_CP_SHARED_KV_BS_GT1_DEBUG.override(True): with envs.SGLANG_CP_SHARED_KV_BS_GT1_DEBUG_LIMIT.override(1): with patch.object(nsa_utils.logger, "info") as info: nsa_utils.log_cp_shared_kv_bs_gt1_debug( "unit_test", "bs=%s", 2, ) nsa_utils.log_cp_shared_kv_bs_gt1_debug( "unit_test", "bs=%s", 3, ) self.assertEqual(info.call_count, 1) self.assertIn( "[CP_SHARED_KV_BS_GT1_DEBUG]", info.call_args.args[0], ) def test_index_topk_batch_lengths_follow_actual_q_rows_not_compute_alias(self): import torch from sglang.srt.layers.attention.nsa.nsa_indexer import ( _select_batch_topk_query_lengths, ) plan = build_batch_page_aligned_in_seq_split_plan( extend_lens=[40387], prefix_lens=[0], page_size=64, cp_size=8, cp_rank=0, ) valid_local_rows = ( plan.request_valid_seq_q_prev[0] + plan.request_valid_seq_q_next[0] ) compute_local_rows = ( 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, valid_local_rows) self.assertFalse(plan.compute_padding_enabled) local = split_tensor_by_cp_batch_plan( torch.arange(40387, dtype=torch.int64), plan, mode="1d", ) 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( nsa_cp_metadata=NSAContextParallelMetadata( batch_size=1, batch_plan=plan, ) ), local, ) self.assertEqual(local_valid.numel(), valid_local_rows) self.assertEqual(local_valid.tolist(), local.tolist()) selected = _select_batch_topk_query_lengths( cp_metadata=NSAContextParallelMetadata(batch_size=1, batch_plan=plan), batch_plan=plan, batch_size=1, q_tokens=valid_local_rows, weights_tokens=valid_local_rows, ) self.assertFalse(selected.uses_compute_query_rows) self.assertEqual(selected.request_seq_q_prev, plan.request_valid_seq_q_prev) self.assertEqual(selected.request_seq_q_next, plan.request_valid_seq_q_next) self.assertEqual( selected.request_valid_seq_q_prev, plan.request_valid_seq_q_prev ) self.assertEqual( selected.request_valid_seq_q_next, plan.request_valid_seq_q_next ) selected_compute_alias = _select_batch_topk_query_lengths( cp_metadata=NSAContextParallelMetadata(batch_size=1, batch_plan=plan), batch_plan=plan, batch_size=1, q_tokens=compute_local_rows, weights_tokens=compute_local_rows, ) self.assertFalse(selected_compute_alias.uses_compute_query_rows) self.assertEqual( selected_compute_alias.request_seq_q_prev, plan.request_compute_seq_q_prev ) self.assertEqual( selected_compute_alias.request_seq_q_next, plan.request_compute_seq_q_next ) self.assertEqual( selected_compute_alias.request_valid_seq_q_prev, plan.request_valid_seq_q_prev ) self.assertEqual( 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_compute_padding_tiny_batch_valid_rows_are_not_suffix_maskable_for_moe( self, ): # Regression invariant for the warm-cache GSM8K bs>1 failure: # compute padding gives every tiny request a fixed local page slot, but # the valid rows inside those slots are per-request prefixes, not one # contiguous tensor prefix. DeepEP/MoE top-k's scalar # num_token_non_padded can only mask a suffix, so it cannot represent # this CP-local row layout. import torch plan = build_batch_page_aligned_in_seq_split_plan( extend_lens=[29, 9, 60, 9, 17], prefix_lens=[704, 704, 640, 704, 704], page_size=64, cp_size=8, cp_rank=0, ) compute_rows = sum(plan.request_compute_rank_local_tokens) valid_rows = sum(plan.request_valid_rank_local_tokens) self.assertTrue(plan.compute_padding_enabled) self.assertEqual(plan.request_compute_rank_local_tokens, [64] * 5) self.assertEqual(plan.request_valid_rank_local_tokens, [29, 9, 60, 9, 17]) self.assertEqual(compute_rows, 320) self.assertEqual(valid_rows, 124) valid_mask = torch.zeros(compute_rows, dtype=torch.bool) for req_offset, spans in zip( plan.request_compute_rank_local_offsets, plan.request_valid_query_row_spans, ): for start, end in spans: if end > start: valid_mask[req_offset + start : req_offset + end] = True # A suffix-padding scalar would keep rows [0, valid_rows) and mask the # rest. The CP layout has dummy page-tail rows inside that prefix and # later valid rows after it. scalar_prefix_mask = torch.arange(compute_rows) < valid_rows self.assertFalse(torch.equal(valid_mask, scalar_prefix_mask)) self.assertGreater(int((scalar_prefix_mask & ~valid_mask).sum().item()), 0) self.assertGreater(int((valid_mask & ~scalar_prefix_mask).sum().item()), 0) def test_compute_padding_non_owner_rank_scalar_non_padded_would_unmask_dummy_rows( self, ): # Same regression shape on a non-owner CP rank: the local tensor has # compute rows but no valid rows. Reusing the global input-token count # as num_token_non_padded would route dummy rows through MoE. extend_lens = [29, 9, 60, 9, 17] plan = build_batch_page_aligned_in_seq_split_plan( extend_lens=extend_lens, prefix_lens=[704, 704, 640, 704, 704], page_size=64, cp_size=8, cp_rank=1, ) compute_rows = sum(plan.request_compute_rank_local_tokens) valid_rows = sum(plan.request_valid_rank_local_tokens) global_num_token_non_padded = sum(extend_lens) self.assertTrue(plan.compute_padding_enabled) self.assertEqual(plan.request_compute_rank_local_tokens, [64] * 5) self.assertEqual(plan.request_valid_rank_local_tokens, [0] * 5) self.assertEqual(compute_rows, 320) self.assertEqual(valid_rows, 0) self.assertGreater(global_num_token_non_padded, valid_rows) 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], prefix_lens=[40320, 8192], page_size=64, cp_size=8, cp_rank=0, ) self.assertTrue(plan.compute_padding_enabled) self.assertEqual(plan.request_valid_padded_pages, [2, 16]) self.assertEqual(plan.request_compute_padded_pages, [8, 16]) self.assertEqual(plan.request_compute_padded_tokens, [512, 1024]) self.assertEqual(plan.request_compute_padding_tokens, [447, 0]) self.assertEqual(plan.request_compute_rank_local_tokens, [64, 128]) self.assertEqual(plan.request_compute_rank_local_offsets, [0, 64]) self.assertEqual(plan.request_valid_rank_local_tokens, [64, 128]) self.assertEqual(plan.request_valid_rank_local_offsets, [0, 64]) self.assertEqual(plan.request_last_token_owner, [1, 0]) def test_batch_plan_stable_helpers_split_and_build_page_owner_plan(self): import torch plan = build_batch_page_aligned_in_seq_split_plan( extend_lens=[4, 9], prefix_lens=[0, 8], page_size=4, cp_size=2, cp_rank=1, ) forward_batch = SimpleNamespace(nsa_cp_metadata=SimpleNamespace(batch_plan=plan)) self.assertIs(get_cp_shared_kv_batch_plan(forward_batch), plan) 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]) 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], ) def test_collect_last_token_hidden_uses_batch_owner_metadata(self): import torch hidden_states = torch.tensor( [[10.0], [11.0], [12.0], [13.0], [20.0], [21.0], [22.0], [23.0]] ) forward_batch = SimpleNamespace( extend_seq_lens_cpu=[4, 9], nsa_cp_metadata=NSAContextParallelMetadata( batch_size=2, request_last_token_owner=[0, 1], request_last_token_local_offset=[3, 4], request_rank_local_offsets=[0, 4], ), ) def fake_all_gather(output, local_last): self.assertEqual(local_last.tolist(), [[13.0], [0.0]]) output.copy_(torch.tensor([[13.0], [0.0], [0.0], [99.0]])) with ( patch( "sglang.srt.layers.attention.nsa.utils.is_nsa_prefill_cp_round_robin_split", return_value=False, ), patch( "sglang.srt.layers.attention.nsa.utils.get_attention_cp_rank", return_value=0, ), patch( "sglang.srt.layers.attention.nsa.utils.attn_cp_all_gather_into_tensor", side_effect=fake_all_gather, ), ): collected = cp_collect_last_token_hidden(hidden_states, forward_batch, 2) self.assertEqual(collected.tolist(), [[13.0], [99.0]]) def test_collect_last_token_hidden_uses_compute_padding_for_single_request(self): import torch plan = build_batch_page_aligned_in_seq_split_plan( extend_lens=[65], prefix_lens=[40320], page_size=64, cp_size=8, cp_rank=1, ) hidden_states = torch.zeros((64, 1), dtype=torch.float32) hidden_states[0] = 123.0 hidden_states[1] = 999.0 forward_batch = SimpleNamespace( extend_seq_lens_cpu=[65], nsa_cp_metadata=NSAContextParallelMetadata( batch_size=1, batch_plan=plan, ), ) def fake_all_gather(output, local_last): self.assertEqual(local_last.tolist(), [[123.0]]) output.zero_() output[1] = local_last[0] with ( patch( "sglang.srt.layers.attention.nsa.utils.is_nsa_prefill_cp_round_robin_split", return_value=False, ), patch( "sglang.srt.layers.attention.nsa.utils.get_attention_cp_rank", return_value=1, ), patch( "sglang.srt.layers.attention.nsa.utils.attn_cp_all_gather_into_tensor", side_effect=fake_all_gather, ), ): collected = cp_collect_last_token_hidden(hidden_states, forward_batch, 8) self.assertEqual(collected.tolist(), [[123.0]]) def test_collect_last_token_hidden_uses_compute_rank_offsets_for_batch(self): import torch hidden_states = torch.zeros((8, 1), dtype=torch.float32) hidden_states[0] = 10.0 hidden_states[1] = 99.0 hidden_states[4] = 20.0 forward_batch = SimpleNamespace( extend_seq_lens_cpu=[5, 5], nsa_cp_metadata=NSAContextParallelMetadata( batch_size=2, request_last_token_owner=[1, 1], request_last_token_local_offset=[0, 0], request_rank_local_offsets=[0, 1], request_compute_rank_local_offsets=[0, 4], compute_padding_enabled=True, ), ) def fake_all_gather(output, local_last): self.assertEqual(local_last.tolist(), [[10.0], [20.0]]) output.copy_(torch.tensor([[0.0], [0.0], [10.0], [20.0]])) with ( patch( "sglang.srt.layers.attention.nsa.utils.is_nsa_prefill_cp_round_robin_split", return_value=False, ), patch( "sglang.srt.layers.attention.nsa.utils.get_attention_cp_rank", return_value=1, ), patch( "sglang.srt.layers.attention.nsa.utils.attn_cp_all_gather_into_tensor", side_effect=fake_all_gather, ), ): collected = cp_collect_last_token_hidden(hidden_states, forward_batch, 2) self.assertEqual(collected.tolist(), [[10.0], [20.0]]) def test_collect_last_token_hidden_matches_parallel20_tiny_extend_layout(self): import torch # Regression shape from the parallel=20 GSM8K warm-cache probe: # tiny extend requests are compute-padded to one page per request. # All real current tokens and all last tokens are owned by CP0, while # the local hidden rows remain laid out as fixed 64-row request slots. plan = build_batch_page_aligned_in_seq_split_plan( extend_lens=[19, 21, 61, 33], prefix_lens=[704, 704, 640, 704], page_size=64, cp_size=8, cp_rank=0, ) self.assertEqual(plan.request_last_token_owner, [0, 0, 0, 0]) self.assertEqual(plan.request_last_token_local_offset, [18, 20, 60, 32]) self.assertEqual(plan.request_compute_rank_local_offsets, [0, 64, 128, 192]) hidden_states = torch.zeros((256, 1), dtype=torch.float32) expected = [[100.0], [200.0], [300.0], [400.0]] for value, rank_offset, local_offset in zip( [100.0, 200.0, 300.0, 400.0], plan.request_compute_rank_local_offsets, plan.request_last_token_local_offset, ): hidden_states[rank_offset + local_offset] = value forward_batch = SimpleNamespace( extend_seq_lens_cpu=[19, 21, 61, 33], nsa_cp_metadata=NSAContextParallelMetadata( batch_size=4, batch_plan=plan, ), ) def fake_all_gather(output, local_last): self.assertEqual(local_last.tolist(), expected) output.zero_() output[:4] = local_last with ( patch( "sglang.srt.layers.attention.nsa.utils.is_nsa_prefill_cp_round_robin_split", return_value=False, ), patch( "sglang.srt.layers.attention.nsa.utils.get_attention_cp_rank", return_value=0, ), patch( "sglang.srt.layers.attention.nsa.utils.attn_cp_all_gather_into_tensor", side_effect=fake_all_gather, ), ): collected = cp_collect_last_token_hidden(hidden_states, forward_batch, 8) self.assertEqual(collected.tolist(), expected) def test_collect_last_token_hidden_fails_fast_without_batch_owner_metadata(self): import torch forward_batch = SimpleNamespace( extend_seq_lens_cpu=[4, 9], nsa_cp_metadata=NSAContextParallelMetadata(batch_size=2), ) with ( patch( "sglang.srt.layers.attention.nsa.utils.is_nsa_prefill_cp_round_robin_split", return_value=False, ), patch( "sglang.srt.layers.attention.nsa.utils.get_attention_cp_rank", return_value=0, ), self.assertRaisesRegex( RuntimeError, r"\[CP_SHARED_KV_FAIL_FAST\]\[batch_gt1_missing_last_token_metadata\]", ), ): cp_collect_last_token_hidden(torch.zeros((8, 1)), forward_batch, 2) def test_deepseek_nextn_cp_draft_bs_gt1_fails_fast_on_hidden_shape_fallback( self, ): import torch model = DeepseekModelNextN.__new__(DeepseekModelNextN) model._debug_cp_draft_shared_kv = lambda _message: None forward_batch = SimpleNamespace( uses_cp_shared_kv=True, batch_size=2, extend_seq_lens_cpu=[4, 9], nsa_cp_metadata=NSAContextParallelMetadata(batch_size=2), ) with ( patch.dict(os.environ, {"SGLANG_CP_DRAFT_SHARED_KV": "1"}), patch( "sglang.srt.models.deepseek_nextn.get_attention_cp_rank", return_value=0, ), patch( "sglang.srt.models.deepseek_nextn.get_attention_cp_size", return_value=8, ), self.assertRaisesRegex( RuntimeError, r"\[CP_SHARED_KV_FAIL_FAST\]\[draft_batch_gt1_spec_hidden_shape_mismatch\]", ), ): model._get_cp_local_spec_hidden_states( forward_batch, torch.zeros((3, 2)), full_num_tokens=13, local_num_tokens=8, ) def test_cp_draft_padding_keeps_local_hidden_when_static_tokens_are_shorter(self): import torch class FakeAttnBackend: def get_cuda_graph_seq_len_fill_value(self): return 0 model_runner = SimpleNamespace(attn_backend=FakeAttnBackend()) spec_info = EagleDraftInput( hidden_states=torch.ones((64, 2), dtype=torch.float32), verified_id=torch.tensor([1], dtype=torch.int64), num_tokens_per_req=1, num_tokens_for_logprob_per_req=1, cp_local_hidden_states=True, ) forward_batch = ForwardBatch( forward_mode=ForwardMode.DRAFT_EXTEND, batch_size=1, input_ids=torch.arange(8, dtype=torch.int64), req_pool_indices=torch.tensor([0], dtype=torch.int64), seq_lens=torch.tensor([4], dtype=torch.int32), out_cache_loc=torch.arange(8, dtype=torch.int64), seq_lens_sum=4, positions=torch.arange(8, dtype=torch.int64), lora_ids=[None], spec_info=spec_info, uses_cp_shared_kv=True, ) with patch.dict(os.environ, {"SGLANG_CP_DRAFT_SHARED_KV": "1"}): forward_batch._pad_inputs_to_size(model_runner, num_tokens=8, bs=1) self.assertEqual(tuple(forward_batch.spec_info.hidden_states.shape), (64, 2)) self.assertTrue( torch.equal( forward_batch.hidden_states_backup, torch.ones((64, 2), dtype=torch.float32), ) ) def test_cp_draft_padding_keeps_marked_cp_local_hidden_before_cp_flags_are_visible( self, ): import torch class FakeAttnBackend: def get_cuda_graph_seq_len_fill_value(self): return 0 model_runner = SimpleNamespace(attn_backend=FakeAttnBackend()) spec_info = EagleDraftInput( hidden_states=torch.ones((64, 2), dtype=torch.float32), verified_id=torch.tensor([1], dtype=torch.int64), num_tokens_per_req=1, num_tokens_for_logprob_per_req=1, cp_local_hidden_states=True, ) forward_batch = ForwardBatch( forward_mode=ForwardMode.DRAFT_EXTEND, batch_size=1, input_ids=torch.arange(8, dtype=torch.int64), req_pool_indices=torch.tensor([0], dtype=torch.int64), seq_lens=torch.tensor([4], dtype=torch.int32), out_cache_loc=torch.arange(8, dtype=torch.int64), seq_lens_sum=4, positions=torch.arange(8, dtype=torch.int64), lora_ids=[None], spec_info=spec_info, uses_cp_shared_kv=False, ) forward_batch._pad_inputs_to_size(model_runner, num_tokens=8, bs=1) self.assertEqual(tuple(forward_batch.spec_info.hidden_states.shape), (64, 2)) self.assertTrue( torch.equal( forward_batch.hidden_states_backup, torch.ones((64, 2), dtype=torch.float32), ) ) def test_cp_draft_padding_keeps_marked_cp_local_hidden_after_forward_mode_rewrite( self, ): import torch class FakeAttnBackend: def get_cuda_graph_seq_len_fill_value(self): return 0 model_runner = SimpleNamespace(attn_backend=FakeAttnBackend()) spec_info = EagleDraftInput( hidden_states=torch.ones((64, 2), dtype=torch.float32), verified_id=torch.tensor([1], dtype=torch.int64), num_tokens_per_req=1, num_tokens_for_logprob_per_req=1, cp_local_hidden_states=True, ) forward_batch = ForwardBatch( # prepare_mlp_sync_batch can temporarily rewrite draft extend to # EXTEND while static DP padding is being prepared. The draft # side-channel contract must therefore be carried by spec_info, not # inferred from forward_mode. forward_mode=ForwardMode.EXTEND, batch_size=1, input_ids=torch.arange(8, dtype=torch.int64), req_pool_indices=torch.tensor([0], dtype=torch.int64), seq_lens=torch.tensor([4], dtype=torch.int32), out_cache_loc=torch.arange(8, dtype=torch.int64), seq_lens_sum=4, positions=torch.arange(8, dtype=torch.int64), lora_ids=[None], spec_info=spec_info, uses_cp_shared_kv=True, ) forward_batch._pad_inputs_to_size(model_runner, num_tokens=8, bs=1) self.assertEqual(tuple(forward_batch.spec_info.hidden_states.shape), (64, 2)) self.assertTrue( torch.equal( forward_batch.hidden_states_backup, torch.ones((64, 2), dtype=torch.float32), ) ) def test_cp_draft_padding_rejects_unmarked_oversized_hidden(self): import torch class FakeAttnBackend: def get_cuda_graph_seq_len_fill_value(self): return 0 model_runner = SimpleNamespace(attn_backend=FakeAttnBackend()) spec_info = EagleDraftInput( hidden_states=torch.ones((64, 2), dtype=torch.float32), verified_id=torch.tensor([1], dtype=torch.int64), num_tokens_per_req=1, num_tokens_for_logprob_per_req=1, ) forward_batch = ForwardBatch( forward_mode=ForwardMode.DRAFT_EXTEND, batch_size=1, input_ids=torch.arange(8, dtype=torch.int64), req_pool_indices=torch.tensor([0], dtype=torch.int64), seq_lens=torch.tensor([4], dtype=torch.int32), out_cache_loc=torch.arange(8, dtype=torch.int64), seq_lens_sum=4, positions=torch.arange(8, dtype=torch.int64), lora_ids=[None], spec_info=spec_info, uses_cp_shared_kv=False, ) with self.assertRaisesRegex( RuntimeError, r"\[CP_SHARED_KV_FAIL_FAST\]\[draft_hidden_static_padding_mismatch\]", ): forward_batch._pad_inputs_to_size(model_runner, num_tokens=8, bs=1) def test_full_rerange_fails_fast_for_batch_metadata(self): import torch forward_batch = SimpleNamespace( nsa_cp_metadata=NSAContextParallelMetadata(batch_size=2) ) with self.assertRaisesRegex( RuntimeError, r"\[CP_SHARED_KV_FAIL_FAST\]\[batch_gt1_full_rerange_unsupported\]", ): cp_all_gather_rerange_output( torch.zeros((8, 1)), 2, forward_batch, stream=None ) def test_cp_shared_kv_prepare_rejects_round_robin_mode(self): class Mode: def is_context_parallel_extend(self): return True forward_batch = SimpleNamespace( uses_cp_shared_kv=True, extend_seq_lens_cpu=[128], extend_prefix_lens_cpu=[0], token_to_kv_pool=SimpleNamespace(page_size=64), forward_mode=Mode(), ) with ( patch( "sglang.srt.layers.attention.nsa.utils.is_nsa_prefill_cp_round_robin_split", return_value=True, ), self.assertRaisesRegex( RuntimeError, r"\[CP_SHARED_KV_FAIL_FAST\]\[round_robin_unsupported\]", ), ): prepare_input_dp_with_cp_dsa( 128, cp_rank=0, cp_size=2, seqs_len=[128], forward_batch=forward_batch, page_size=64, ) def test_cp_shared_kv_prepare_uses_batch_plan_for_bs1_compute_padding(self): class Mode: def is_context_parallel_extend(self): return True forward_batch = SimpleNamespace( uses_cp_shared_kv=True, extend_seq_lens_cpu=[65], extend_prefix_lens_cpu=[0], token_to_kv_pool=SimpleNamespace(page_size=64), forward_mode=Mode(), ) with patch( "sglang.srt.layers.attention.nsa.utils.is_nsa_prefill_cp_round_robin_split", return_value=False, ): metadata = prepare_input_dp_with_cp_dsa( 65, cp_rank=1, cp_size=8, seqs_len=[65], forward_batch=forward_batch, page_size=64, ) self.assertIsNotNone(metadata.batch_plan) self.assertEqual(metadata.batch_size, 1) self.assertTrue(metadata.compute_padding_enabled) self.assertEqual( metadata.request_valid_split_lists, [[64, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]], ) self.assertEqual( metadata.request_compute_split_lists, [[64, 64, 64, 64, 64, 64, 64, 64, 0, 0, 0, 0, 0, 0, 0, 0]], ) self.assertEqual(metadata.split_list, metadata.request_compute_split_lists[0]) self.assertEqual(metadata.max_rank_len, [64] * 8) self.assertEqual(metadata.per_rank_actual_token, [64] * 8) self.assertEqual(metadata.actual_seq_q_prev, 64) self.assertEqual(metadata.actual_seq_q_next, 0) self.assertEqual(metadata.request_valid_seq_q_prev, [1]) self.assertEqual(metadata.request_valid_seq_q_next, [0]) def test_cp_shared_kv_all_gather_rejects_round_robin_mode(self): import torch forward_batch = SimpleNamespace( uses_cp_shared_kv=True, nsa_cp_metadata=NSAContextParallelMetadata( split_list=[2, 2, 2, 2], zigzag_index=[0, 3], reverse_split_len=[2, 2, 2, 2], cp_reverse_index=[0, 2, 3, 1], total_seq_lens=torch.tensor(8), ), ) with ( patch( "sglang.srt.layers.attention.nsa.utils.is_nsa_prefill_cp_round_robin_split", return_value=True, ), self.assertRaisesRegex( RuntimeError, r"\[CP_SHARED_KV_FAIL_FAST\]\[round_robin_unsupported\]", ), ): cp_all_gather_rerange_output( torch.zeros((4, 1)), 2, forward_batch, stream=None ) def test_batch_in_seq_all_gather_rerange_restores_request_order_bf16(self): import torch cp_size = 2 request_split_lists = [ [2, 1, 3, 0], [1, 2, 0, 1], ] input_tensor_all, expected = self._build_batch_rerange_case( cp_size=cp_size, request_split_lists=request_split_lists, row_width=3, dtype=torch.bfloat16, ) forward_batch = SimpleNamespace( nsa_cp_metadata=NSAContextParallelMetadata( batch_size=2, request_split_lists=request_split_lists, max_rank_len=[6, 6], ) ) actual = _torch_batch_in_seq_all_gather_rerange( input_tensor_all, forward_batch, cp_size=cp_size, ) self.assertEqual(actual.dtype, torch.bfloat16) self.assertTrue(torch.equal(actual, expected)) def test_batch_in_seq_all_gather_rerange_treats_fp8_payload_as_opaque_rows(self): import torch cp_size = 2 request_split_lists = [ [1, 2, 1, 0], [2, 0, 1, 1], ] input_tensor_all, expected = self._build_batch_rerange_case( cp_size=cp_size, request_split_lists=request_split_lists, row_width=5, dtype=torch.uint8, ) forward_batch = SimpleNamespace( nsa_cp_metadata=NSAContextParallelMetadata( batch_size=2, request_split_lists=request_split_lists, max_rank_len=[4, 4], ) ) actual = _torch_batch_in_seq_all_gather_rerange( input_tensor_all, forward_batch, cp_size=cp_size, ) self.assertEqual(actual.dtype, torch.uint8) self.assertTrue(torch.equal(actual, expected)) def test_batch_in_seq_all_gather_rerange_uses_compute_offsets_for_padded_source(self): import torch cp_size = 2 # Request 0 is a tiny suffix: valid output only has segment 0, but the # rank-major source payload contains synthetic compute-padding rows in # the rank-local mirror segment. Request 1 follows it on the same rank. # Source offsets must therefore be computed from compute splits, while # output rows must still be restored from valid splits only. valid_split_lists = [ [1, 0, 0, 0], [2, 0, 1, 0], ] compute_split_lists = [ [1, 1, 1, 1], [2, 0, 1, 0], ] row_width = 2 max_rank_token = 4 input_tensor_all = torch.zeros((max_rank_token * cp_size, row_width)) # Build source rank-major payload by compute split. Values 900+ are # synthetic padding rows and must never appear in the restored output. req0_seg0 = torch.tensor([[10.0, 11.0]]) req0_seg1_pad = torch.tensor([[900.0, 901.0]]) req0_seg2_pad = torch.tensor([[902.0, 903.0]]) req0_seg3_pad = torch.tensor([[904.0, 905.0]]) req1_seg0 = torch.tensor([[20.0, 21.0], [22.0, 23.0]]) req1_seg2 = torch.tensor([[24.0, 25.0]]) # rank0 owns segment 0 then mirror segment 3 for each request. input_tensor_all[0:1] = req0_seg0 input_tensor_all[1:2] = req0_seg3_pad input_tensor_all[2:4] = req1_seg0 # rank1 owns segment 1 then mirror segment 2 for each request. rank1 = max_rank_token input_tensor_all[rank1 : rank1 + 1] = req0_seg1_pad input_tensor_all[rank1 + 1 : rank1 + 2] = req0_seg2_pad input_tensor_all[rank1 + 2 : rank1 + 3] = req1_seg2 expected = torch.cat([req0_seg0, req1_seg0, req1_seg2], dim=0) forward_batch = SimpleNamespace( nsa_cp_metadata=NSAContextParallelMetadata( batch_size=2, request_split_lists=valid_split_lists, request_compute_split_lists=compute_split_lists, compute_padding_enabled=True, max_rank_len=[max_rank_token, max_rank_token], ) ) actual = _torch_batch_in_seq_all_gather_rerange( input_tensor_all, forward_batch, cp_size=cp_size, ) self.assertTrue(torch.equal(actual, expected)) def test_batch_in_seq_all_gather_rerange_matches_parallel20_tiny_extend_layout(self): import torch cp_size = 8 extend_lens = [19, 21, 61, 33] prefix_lens = [704, 704, 640, 704] plans = [ build_batch_page_aligned_in_seq_split_plan( extend_lens=extend_lens, prefix_lens=prefix_lens, page_size=64, cp_size=cp_size, cp_rank=rank, ) for rank in range(cp_size) ] valid_split_lists = plans[0].request_split_lists compute_split_lists = plans[0].request_compute_split_lists self.assertEqual(plans[0].request_valid_rank_local_tokens, extend_lens) for rank in range(1, cp_size): self.assertEqual(plans[rank].request_valid_rank_local_tokens, [0, 0, 0, 0]) max_rank_token = max( sum( split[rank] + split[cp_size * 2 - rank - 1] for split in compute_split_lists ) for rank in range(cp_size) ) self.assertEqual(max_rank_token, 64 * len(extend_lens)) input_tensor_all = torch.zeros((max_rank_token * cp_size, 1), dtype=torch.float32) expected_rows = [] rank0_cursor = 0 for req_id, extend_len in enumerate(extend_lens): valid_rows = torch.arange( req_id * 1000, req_id * 1000 + extend_len, dtype=torch.float32, ).view(-1, 1) input_tensor_all[rank0_cursor : rank0_cursor + extend_len] = valid_rows # Fill the rest of the 64-row compute slot with poison values that # must not appear after valid-output rerange. pad_len = 64 - extend_len if pad_len: input_tensor_all[ rank0_cursor + extend_len : rank0_cursor + 64 ] = 900000.0 + req_id expected_rows.append(valid_rows) rank0_cursor += 64 # Non-owner ranks only have compute padding in this tiny-extend case. input_tensor_all[max_rank_token:] = -777.0 forward_batch = SimpleNamespace( nsa_cp_metadata=NSAContextParallelMetadata( batch_size=len(extend_lens), request_split_lists=valid_split_lists, request_compute_split_lists=compute_split_lists, compute_padding_enabled=True, max_rank_len=[max_rank_token] * cp_size, ) ) actual = _torch_batch_in_seq_all_gather_rerange( input_tensor_all, forward_batch, cp_size=cp_size, ) expected = torch.cat(expected_rows, dim=0) self.assertTrue(torch.equal(actual, expected)) def _build_batch_rerange_case( self, *, cp_size, request_split_lists, row_width, dtype, ): import torch rank_tokens = [] for rank in range(cp_size): mirror = cp_size * 2 - rank - 1 rank_tokens.append( sum(split[rank] + split[mirror] for split in request_split_lists) ) max_rank_token = max(rank_tokens) total_tokens = sum(sum(split) for split in request_split_lists) input_tensor_all = torch.zeros( (max_rank_token * cp_size, row_width), dtype=dtype, ) expected = torch.empty((total_tokens, row_width), dtype=dtype) next_value = 1 request_segments = [] for split in request_split_lists: segments = [] for segment_len in split: if dtype == torch.uint8: rows = ( torch.arange( next_value, next_value + segment_len * row_width, dtype=torch.int64, ) .remainder(251) .to(torch.uint8) .view(segment_len, row_width) ) else: rows = ( torch.arange( next_value, next_value + segment_len * row_width, dtype=torch.float32, ) .view(segment_len, row_width) .to(dtype) ) next_value += segment_len * row_width segments.append(rows) request_segments.append(segments) output_cursor = 0 for segments in request_segments: for rows in segments: expected[output_cursor : output_cursor + rows.shape[0]] = rows output_cursor += rows.shape[0] for rank in range(cp_size): mirror = cp_size * 2 - rank - 1 rank_cursor = rank * max_rank_token for segments in request_segments: for segment_id in (rank, mirror): rows = segments[segment_id] input_tensor_all[ rank_cursor : rank_cursor + rows.shape[0] ] = rows rank_cursor += rows.shape[0] return input_tensor_all, expected def test_local_pair_split_uses_metadata_lengths_not_half_split(self): import torch tensor = torch.arange(9) prev, next_ = split_in_seq_cp_local_pair(tensor, 6, 3) self.assertEqual(prev.tolist(), [0, 1, 2, 3, 4, 5]) self.assertEqual(next_.tolist(), [6, 7, 8]) def test_local_pair_split_rejects_stale_metadata(self): import torch with self.assertRaisesRegex(RuntimeError, "local in-seq CP length mismatch"): split_in_seq_cp_local_pair(torch.arange(9), 5, 5, name="q_fp8") def test_cp_split_and_rebuild_1d_matches_in_seq_zigzag_order(self): import torch from types import SimpleNamespace forward_batch = SimpleNamespace( nsa_cp_metadata=NSAContextParallelMetadata( split_list=[2, 2, 2, 2, 2, 2, 2, 2], zigzag_index=[1, 6], ) ) local_locs = cp_split_and_rebuild_1d(forward_batch, torch.arange(16)) self.assertEqual(local_locs.tolist(), [2, 3, 12, 13]) def test_cp_split_and_rebuild_data_keeps_batch_request_boundaries(self): import torch forward_batch = SimpleNamespace( nsa_cp_metadata=NSAContextParallelMetadata( batch_size=2, request_extend_lens=[4, 9], request_split_lists=[[4, 0, 0, 0], [4, 4, 1, 0]], request_zigzag_indices=[[0, 3], [0, 3]], ) ) tensor = torch.arange(13 * 2).view(13, 2) with patch( "sglang.srt.layers.attention.nsa.utils.is_nsa_prefill_cp_round_robin_split", return_value=False, ): local = cp_split_and_rebuild_data(forward_batch, tensor) self.assertEqual(local[:, 0].tolist(), list(range(0, 16, 2))) def test_cp_split_and_rebuild_data_preserves_non_token_dimensions(self): import torch forward_batch = SimpleNamespace( nsa_cp_metadata=NSAContextParallelMetadata( batch_size=2, request_extend_lens=[4, 9], request_split_lists=[[4, 0, 0, 0], [4, 4, 1, 0]], request_zigzag_indices=[[0, 3], [0, 3]], ) ) tensor = torch.arange(13 * 2 * 3, dtype=torch.float32).view(13, 2, 3) with patch( "sglang.srt.layers.attention.nsa.utils.is_nsa_prefill_cp_round_robin_split", return_value=False, ): local = cp_split_and_rebuild_data(forward_batch, tensor) self.assertEqual(tuple(local.shape), (8, 2, 3)) self.assertTrue(torch.equal(local[0], tensor[0])) def test_cp_split_and_rebuild_data_uses_compute_padding_rows(self): import torch plan = build_batch_page_aligned_in_seq_split_plan( extend_lens=[65], prefix_lens=[40320], page_size=64, cp_size=8, cp_rank=1, ) forward_batch = SimpleNamespace( nsa_cp_metadata=NSAContextParallelMetadata( batch_size=1, batch_plan=plan, ) ) tensor = torch.arange(65 * 2, dtype=torch.float32).view(65, 2) with patch( "sglang.srt.layers.attention.nsa.utils.is_nsa_prefill_cp_round_robin_split", return_value=False, ): local = cp_split_and_rebuild_data(forward_batch, tensor) self.assertEqual(local.shape, (64, 2)) self.assertEqual(local[0].tolist(), [128.0, 129.0]) self.assertTrue(torch.equal(local[1:], torch.zeros((63, 2)))) def test_cp_split_and_rebuild_data_ignores_trailing_static_padding_rows(self): import torch plan = build_batch_page_aligned_in_seq_split_plan( extend_lens=[4], prefix_lens=[0], page_size=4, cp_size=2, cp_rank=1, ) forward_batch = SimpleNamespace( nsa_cp_metadata=NSAContextParallelMetadata( batch_size=1, batch_plan=plan, ) ) tensor = torch.arange(8 * 2, dtype=torch.float32).view(8, 2) with patch( "sglang.srt.layers.attention.nsa.utils.is_nsa_prefill_cp_round_robin_split", return_value=False, ): local = cp_split_and_rebuild_data(forward_batch, tensor) self.assertEqual(local.shape, (4, 2)) self.assertTrue(torch.equal(local, torch.zeros((4, 2)))) def test_cp_split_and_rebuild_data_ignores_mlp_sync_static_padding_without_compute_padding( self, ): import torch plan = build_batch_page_aligned_in_seq_split_plan( extend_lens=[8], prefix_lens=[0], page_size=4, cp_size=2, cp_rank=1, ) self.assertFalse(plan.compute_padding_enabled) forward_batch = SimpleNamespace( extend_num_tokens=9, nsa_cp_metadata=NSAContextParallelMetadata( batch_size=1, batch_plan=plan, ), ) tensor = torch.arange(9 * 2, dtype=torch.float32).view(9, 2) expected = split_tensor_by_cp_batch_plan( tensor[:8], plan, mode="data", ) with patch( "sglang.srt.layers.attention.nsa.utils.is_nsa_prefill_cp_round_robin_split", return_value=False, ): local = cp_split_and_rebuild_data(forward_batch, tensor) self.assertTrue(torch.equal(local, expected)) def test_cp_split_valid_kind_rejects_trailing_padding_rows(self): import torch plan = build_batch_page_aligned_in_seq_split_plan( extend_lens=[4], prefix_lens=[0], page_size=4, cp_size=2, cp_rank=1, ) with self.assertRaisesRegex( RuntimeError, "batch_gt1_split_input_len_mismatch", ): split_tensor_by_cp_batch_plan( torch.arange(8), plan, mode="1d", split_kind="valid", ) def test_cp_split_and_rebuild_1d_keeps_batch_request_boundaries(self): import torch forward_batch = SimpleNamespace( nsa_cp_metadata=NSAContextParallelMetadata( batch_size=2, request_extend_lens=[4, 9], request_split_lists=[[4, 0, 0, 0], [4, 4, 1, 0]], request_zigzag_indices=[[1, 2], [1, 2]], ) ) with patch( "sglang.srt.layers.attention.nsa.utils.is_nsa_prefill_cp_round_robin_split", return_value=False, ): local = cp_split_and_rebuild_1d(forward_batch, torch.arange(13)) self.assertEqual(local.tolist(), [8, 9, 10, 11, 12]) def test_cp_split_and_rebuild_1d_uses_zero_compute_padding_rows(self): import torch plan = build_batch_page_aligned_in_seq_split_plan( extend_lens=[65], prefix_lens=[40320], page_size=64, cp_size=8, cp_rank=1, ) forward_batch = SimpleNamespace( nsa_cp_metadata=NSAContextParallelMetadata( batch_size=1, batch_plan=plan, ) ) with patch( "sglang.srt.layers.attention.nsa.utils.is_nsa_prefill_cp_round_robin_split", return_value=False, ): local = cp_split_and_rebuild_1d(forward_batch, torch.arange(65)) self.assertEqual(local.shape, (64,)) self.assertEqual(local[0].item(), 64) self.assertEqual(local[1:].tolist(), [0] * 63) def test_select_cp_local_valid_rows_filters_compute_padding_rows(self): import torch plan = build_batch_page_aligned_in_seq_split_plan( extend_lens=[65], prefix_lens=[0], page_size=64, cp_size=8, cp_rank=1, ) forward_batch = SimpleNamespace( nsa_cp_metadata=NSAContextParallelMetadata( batch_size=1, batch_plan=plan, ) ) local_compute_rows = torch.full((64, 2), -1.0) local_compute_rows[0] = torch.tensor([50.0, 51.0]) selected = select_cp_local_valid_rows_for_cache_write( forward_batch, local_compute_rows ) self.assertEqual(selected.tolist(), [[50.0, 51.0]]) def test_restore_cp_local_valid_rows_for_moe_keeps_dummy_rows_zero(self): import torch from sglang.srt.layers.attention.nsa.utils import ( restore_cp_local_valid_rows_for_moe, ) plan = build_batch_page_aligned_in_seq_split_plan( extend_lens=[29, 9, 60, 9, 17], prefix_lens=[704, 704, 640, 704, 704], page_size=64, cp_size=8, cp_rank=0, ) forward_batch = SimpleNamespace( nsa_cp_metadata=NSAContextParallelMetadata( batch_size=5, batch_plan=plan, ) ) local_compute_rows = torch.zeros( (sum(plan.request_compute_rank_local_tokens), 2), dtype=torch.float32 ) compact_valid = torch.arange( sum(plan.request_valid_rank_local_tokens) * 2, dtype=torch.float32 ).view(-1, 2) restored = restore_cp_local_valid_rows_for_moe( forward_batch, compact_valid, local_compute_rows, ) self.assertEqual(tuple(restored.shape), tuple(local_compute_rows.shape)) valid_rows = select_cp_local_valid_rows_for_cache_write( forward_batch, restored, ) self.assertTrue(torch.equal(valid_rows, compact_valid)) valid_mask = torch.zeros(restored.shape[0], dtype=torch.bool) cursor = 0 for compute_len, valid_len in zip( plan.request_compute_rank_local_tokens, plan.request_valid_rank_local_tokens, ): valid_mask[cursor : cursor + valid_len] = True cursor += compute_len self.assertTrue(torch.equal(restored[~valid_mask], torch.zeros_like(restored[~valid_mask]))) def test_select_cp_current_valid_rows_accepts_global_rows_under_compute_padding( self, ): import torch plan = build_batch_page_aligned_in_seq_split_plan( extend_lens=[65], prefix_lens=[0], page_size=64, cp_size=8, cp_rank=1, ) self.assertTrue(plan.compute_padding_enabled) global_current = torch.arange(65 * 2, dtype=torch.float32).view(65, 2) expected = split_tensor_by_cp_batch_plan( global_current, plan, mode="data", split_kind="valid", ) forward_batch = SimpleNamespace( extend_seq_lens_cpu=[65], cp_local_out_cache_loc=torch.arange(expected.shape[0], dtype=torch.long), nsa_cp_metadata=NSAContextParallelMetadata( batch_size=1, batch_plan=plan, ), ) selected = select_cp_current_valid_rows_for_reuse( forward_batch, global_current, ) self.assertIsNotNone(selected) self.assertTrue(torch.equal(selected, expected)) def test_cp_split_and_rebuild_position_is_batch_aware_and_compute_padded(self): import torch plan = build_batch_page_aligned_in_seq_split_plan( extend_lens=[65], prefix_lens=[40320], page_size=64, cp_size=8, cp_rank=1, ) forward_batch = SimpleNamespace( nsa_cp_metadata=NSAContextParallelMetadata( batch_size=1, batch_plan=plan, ) ) positions = torch.arange(40320, 40385, dtype=torch.int32) with patch( "sglang.srt.layers.attention.nsa.utils.is_nsa_prefill_cp_round_robin_split", return_value=False, ): local = cp_split_and_rebuild_position(forward_batch, positions) self.assertEqual(local.shape, (64,)) self.assertEqual(local.tolist(), list(range(40384, 40448))) def test_cp_split_and_rebuild_position_ignores_mlp_sync_static_padding_without_compute_padding( self, ): import torch plan = build_batch_page_aligned_in_seq_split_plan( extend_lens=[8], prefix_lens=[0], page_size=4, cp_size=2, cp_rank=1, ) self.assertFalse(plan.compute_padding_enabled) forward_batch = SimpleNamespace( extend_num_tokens=9, nsa_cp_metadata=NSAContextParallelMetadata( batch_size=1, batch_plan=plan, ), ) positions = torch.arange(9, dtype=torch.int32) expected = split_tensor_by_cp_batch_plan( positions[:8], plan, mode="position", ) with patch( "sglang.srt.layers.attention.nsa.utils.is_nsa_prefill_cp_round_robin_split", return_value=False, ): local = cp_split_and_rebuild_position(forward_batch, positions) self.assertTrue(torch.equal(local, expected)) def test_cp_local_embedding_pad_len_uses_metadata_max_rank_len(self): from types import SimpleNamespace import torch forward_batch = SimpleNamespace( nsa_cp_metadata=NSAContextParallelMetadata(max_rank_len=[4096] * 8) ) self.assertEqual( get_cp_local_embedding_padded_token_count(forward_batch, 4040), 4096 ) self.assertEqual( get_cp_local_embedding_padded_token_count(forward_batch, 4096), 4096 ) self.assertEqual( pad_cp_local_input_ids_for_embedding( SimpleNamespace( nsa_cp_metadata=NSAContextParallelMetadata(max_rank_len=[6] * 8) ), torch.tensor([11, 12, 13, 14]), ).tolist(), [11, 12, 13, 14, 0, 0], ) self.assertEqual( pad_cp_local_input_ids_for_embedding( SimpleNamespace( nsa_cp_metadata=NSAContextParallelMetadata(max_rank_len=[4] * 8) ), torch.tensor([11, 12, 13, 14]), ).tolist(), [11, 12, 13, 14], ) missing_metadata = SimpleNamespace(nsa_cp_metadata=None) self.assertIsNone( get_cp_local_embedding_padded_token_count(missing_metadata, 4040) ) self.assertIsNone( pad_cp_local_input_ids_for_embedding( missing_metadata, torch.tensor([11, 12, 13, 14]) ) ) stale_metadata = SimpleNamespace( nsa_cp_metadata=NSAContextParallelMetadata(max_rank_len=[4039] * 8) ) self.assertIsNone( get_cp_local_embedding_padded_token_count(stale_metadata, 4040) ) def test_local_out_cache_loc_requires_compute_owner_pages(self): import torch from types import SimpleNamespace page_size = 4 # Segment order for cp_size=4, cp_rank=1 is segment 1 then 6. # The logical page ids below deliberately encode the same owners through # (logical_page - 1) % cp_size: # segment 1 -> page 2 owner 1 # segment 6 -> page 6 owner 1 segment_pages = [1, 2, 3, 4, 8, 7, 6, 5] out_cache_loc = torch.cat( [ torch.arange(page * page_size, (page + 1) * page_size) for page in segment_pages ] ) forward_batch = SimpleNamespace( uses_cp_shared_kv=True, cp_shared_kv_layout=CpSharedKVLayout( page_size=page_size, cp_size=4, cp_rank=1, ), nsa_cp_metadata=NSAContextParallelMetadata( split_list=[page_size] * 8, zigzag_index=[1, 6], page_aligned=True, page_size=page_size, extend_prefix_len=0, ), out_cache_loc=out_cache_loc, ) local_locs = get_cp_shared_kv_local_out_cache_loc(forward_batch) self.assertIsNotNone(local_locs) self.assertEqual( local_locs.tolist(), list(range(2 * page_size, 3 * page_size)) + list(range(6 * page_size, 7 * page_size)), ) def test_batch_local_out_cache_loc_keeps_request_boundaries(self): import torch from types import SimpleNamespace page_size = 4 # cp_size=2/cp_rank=1 selects segment 1 and 2 for each request. # req0 has no rank-1 local rows. req1 contributes segment 1 (page 2) # and segment 2 (tail page 4). The synthetic logical page ids encode # the owner-lane invariant through (page_id - 1) % cp_size. out_cache_loc = torch.cat( [ torch.arange(1 * page_size, 2 * page_size), # req0 seg0 owner 0 torch.arange(3 * page_size, 4 * page_size), # req1 seg0 owner 0 torch.arange(2 * page_size, 3 * page_size), # req1 seg1 owner 1 torch.tensor([4 * page_size]), # req1 seg2 owner 1 tail ] ) forward_batch = SimpleNamespace( uses_cp_shared_kv=True, extend_num_tokens=8, cp_shared_kv_layout=CpSharedKVLayout( page_size=page_size, cp_size=2, cp_rank=1, ), nsa_cp_metadata=NSAContextParallelMetadata( batch_size=2, split_list=[4, 0, 0, 0], zigzag_index=[1, 2], page_aligned=True, page_size=page_size, extend_prefix_len=0, request_extend_lens=[4, 9], request_split_lists=[[4, 0, 0, 0], [4, 4, 1, 0]], request_zigzag_indices=[[1, 2], [1, 2]], ), out_cache_loc=out_cache_loc, ) local_locs = get_cp_shared_kv_local_out_cache_loc(forward_batch) self.assertIsNotNone(local_locs) self.assertEqual( local_locs.tolist(), list(range(2 * page_size, 3 * page_size)) + [4 * page_size], ) def test_local_out_cache_loc_uses_valid_rows_under_compute_padding(self): import torch page_size = 64 plan = build_batch_page_aligned_in_seq_split_plan( extend_lens=[65], prefix_lens=[0], page_size=page_size, cp_size=8, cp_rank=1, ) out_cache_loc = torch.cat( [ torch.arange(page_size, 2 * page_size), torch.tensor([2 * page_size]), ] ).to(torch.int64) forward_batch = SimpleNamespace( uses_cp_shared_kv=True, cp_shared_kv_layout=CpSharedKVLayout( page_size=page_size, cp_size=8, cp_rank=1, ), nsa_cp_metadata=NSAContextParallelMetadata( batch_size=1, batch_plan=plan, page_aligned=True, page_size=page_size, extend_prefix_len=0, ), out_cache_loc=out_cache_loc, ) with patch( "sglang.srt.layers.attention.nsa.utils.is_nsa_prefill_cp_round_robin_split", return_value=False, ): local_locs = get_cp_shared_kv_local_out_cache_loc(forward_batch) self.assertEqual(local_locs.tolist(), [2 * page_size]) def test_local_out_cache_loc_ignores_trailing_static_padding_locs(self): import torch page_size = 4 plan = build_batch_page_aligned_in_seq_split_plan( extend_lens=[4], prefix_lens=[0], page_size=page_size, cp_size=2, cp_rank=0, ) valid_locs = torch.arange(1 * page_size, 2 * page_size, dtype=torch.int64) static_padding_locs = torch.arange( 99 * page_size, 100 * page_size, dtype=torch.int64 ) forward_batch = SimpleNamespace( uses_cp_shared_kv=True, extend_num_tokens=8, cp_shared_kv_layout=CpSharedKVLayout( page_size=page_size, cp_size=2, cp_rank=0, ), nsa_cp_metadata=NSAContextParallelMetadata( batch_size=1, batch_plan=plan, page_aligned=True, page_size=page_size, extend_prefix_len=0, ), out_cache_loc=torch.cat((valid_locs, static_padding_locs)), ) with patch( "sglang.srt.layers.attention.nsa.utils.is_nsa_prefill_cp_round_robin_split", return_value=False, ): local_locs = get_cp_shared_kv_local_out_cache_loc(forward_batch) self.assertEqual(local_locs.tolist(), valid_locs.tolist()) def test_local_out_cache_loc_ignores_mlp_sync_static_padding_without_compute_padding( self, ): import torch page_size = 4 plan = build_batch_page_aligned_in_seq_split_plan( extend_lens=[8], prefix_lens=[0], page_size=page_size, cp_size=2, cp_rank=1, ) self.assertFalse(plan.compute_padding_enabled) static_padding_locs = torch.tensor([99 * page_size], dtype=torch.int64) forward_batch = SimpleNamespace( uses_cp_shared_kv=True, extend_num_tokens=9, cp_shared_kv_layout=CpSharedKVLayout( page_size=page_size, cp_size=2, cp_rank=1, ), nsa_cp_metadata=NSAContextParallelMetadata( batch_size=1, batch_plan=plan, page_aligned=True, page_size=page_size, extend_prefix_len=0, ), out_cache_loc=torch.cat( ( torch.arange(1 * page_size, 3 * page_size, dtype=torch.int64), static_padding_locs, ) ), ) with patch( "sglang.srt.layers.attention.nsa.utils.is_nsa_prefill_cp_round_robin_split", return_value=False, ): local_locs = get_cp_shared_kv_local_out_cache_loc(forward_batch) self.assertEqual( local_locs.tolist(), [ 2 * page_size, 2 * page_size + 1, 2 * page_size + 2, 2 * page_size + 3, ], ) def test_local_out_cache_loc_rejects_unproven_trailing_padding_even_with_compute_padding( self, ): import torch page_size = 4 plan = build_batch_page_aligned_in_seq_split_plan( extend_lens=[4], prefix_lens=[0], page_size=page_size, cp_size=2, cp_rank=0, ) self.assertTrue(plan.compute_padding_enabled) forward_batch = SimpleNamespace( uses_cp_shared_kv=True, cp_shared_kv_layout=CpSharedKVLayout( page_size=page_size, cp_size=2, cp_rank=0, ), nsa_cp_metadata=NSAContextParallelMetadata( batch_size=1, batch_plan=plan, page_aligned=True, page_size=page_size, extend_prefix_len=0, ), out_cache_loc=torch.cat( ( torch.arange(1 * page_size, 2 * page_size, dtype=torch.int64), torch.arange(99 * page_size, 100 * page_size, dtype=torch.int64), ) ), ) with self.assertRaisesRegex(RuntimeError, "static_padded=None"): get_cp_shared_kv_local_out_cache_loc(forward_batch) def test_batch_local_physical_out_cache_loc_reuses_layer_invariant_plan(self): import torch from types import SimpleNamespace page_size = 4 out_cache_loc = torch.cat( [ torch.arange(1 * page_size, 2 * page_size), torch.arange(3 * page_size, 4 * page_size), torch.arange(2 * page_size, 3 * page_size), torch.tensor([4 * page_size]), ] ) forward_batch = SimpleNamespace( uses_cp_shared_kv=True, cp_shared_kv_layout=CpSharedKVLayout( page_size=page_size, cp_size=2, cp_rank=1, ), nsa_cp_metadata=NSAContextParallelMetadata( batch_size=2, split_list=[4, 0, 0, 0], zigzag_index=[1, 2], page_aligned=True, page_size=page_size, extend_prefix_len=0, request_extend_lens=[4, 9], request_split_lists=[[4, 0, 0, 0], [4, 4, 1, 0]], request_zigzag_indices=[[1, 2], [1, 2]], ), out_cache_loc=out_cache_loc, ) physical_locs = get_cp_shared_kv_local_physical_out_cache_loc(forward_batch) second_read = get_cp_shared_kv_local_physical_out_cache_loc(forward_batch) self.assertIs(physical_locs, second_read) self.assertEqual( physical_locs.tolist(), list(range(1 * page_size, 2 * page_size)) + [2 * page_size], ) def test_local_physical_out_cache_loc_is_cached(self): import torch from types import SimpleNamespace page_size = 4 segment_pages = [1, 2, 3, 4, 8, 7, 6, 5] out_cache_loc = torch.cat( [ torch.arange(page * page_size, (page + 1) * page_size) for page in segment_pages ] ) forward_batch = SimpleNamespace( uses_cp_shared_kv=True, cp_shared_kv_layout=CpSharedKVLayout( page_size=page_size, cp_size=4, cp_rank=1, ), nsa_cp_metadata=NSAContextParallelMetadata( split_list=[page_size] * 8, zigzag_index=[1, 6], page_aligned=True, page_size=page_size, extend_prefix_len=0, ), out_cache_loc=out_cache_loc, ) physical_locs = get_cp_shared_kv_local_physical_out_cache_loc(forward_batch) second_read = get_cp_shared_kv_local_physical_out_cache_loc(forward_batch) self.assertIs(physical_locs, second_read) self.assertEqual( physical_locs.tolist(), list(range(1 * page_size, 2 * page_size)) + list(range(2 * page_size, 3 * page_size)), ) def test_local_out_cache_loc_fails_fast_when_owner_mismatch(self): import torch from types import SimpleNamespace page_size = 4 out_cache_loc = torch.arange(page_size * 8, page_size * 16) forward_batch = SimpleNamespace( uses_cp_shared_kv=True, cp_shared_kv_layout=CpSharedKVLayout( page_size=page_size, cp_size=4, cp_rank=1, ), nsa_cp_metadata=NSAContextParallelMetadata( split_list=[page_size] * 8, zigzag_index=[1, 6], page_aligned=True, page_size=page_size, extend_prefix_len=0, ), out_cache_loc=out_cache_loc, ) with self.assertRaisesRegex( RuntimeError, "CP_SHARED_KV_FAIL_FAST.*local_loc_owner_mismatch", ): get_cp_shared_kv_local_out_cache_loc(forward_batch) def test_local_out_cache_loc_fails_fast_every_invalid_event(self): import torch from types import SimpleNamespace page_size = 4 forward_batch = SimpleNamespace( uses_cp_shared_kv=True, cp_shared_kv_layout=CpSharedKVLayout( page_size=page_size, cp_size=4, cp_rank=1, ), nsa_cp_metadata=NSAContextParallelMetadata( split_list=[page_size] * 8, zigzag_index=[1, 6], page_aligned=False, page_size=page_size, extend_prefix_len=0, ), out_cache_loc=torch.arange(page_size * 8, page_size * 16), ) with self.assertLogs( "sglang.srt.layers.attention.nsa.utils", level="ERROR" ) as cm: with self.assertRaisesRegex( RuntimeError, "CP_SHARED_KV_FAIL_FAST.*not_page_aligned", ): get_cp_shared_kv_local_out_cache_loc(forward_batch) with self.assertRaisesRegex( RuntimeError, "CP_SHARED_KV_FAIL_FAST.*not_page_aligned", ): get_cp_shared_kv_local_out_cache_loc(forward_batch) self.assertEqual(len(cm.output), 2) self.assertIn("[CP_SHARED_KV_FAIL_FAST][direct_write]", cm.output[0]) self.assertIn("metadata is not page-aligned", cm.output[0]) self.assertIn("[CP_SHARED_KV_FAIL_FAST][direct_write]", cm.output[1]) self.assertIn("metadata is not page-aligned", cm.output[1]) def test_indexer_direct_write_fails_fast_on_local_shape_mismatch(self): import torch from sglang.srt.layers.attention.nsa import nsa_indexer from sglang.srt.layers.attention.nsa.nsa_indexer import Indexer indexer = object.__new__(Indexer) indexer.nsa_enable_prefill_cp = True forward_batch = SimpleNamespace(nsa_cp_metadata=object()) with ( patch.object(nsa_indexer, "nsa_use_prefill_cp", return_value=True), patch.object( nsa_indexer, "get_cp_shared_kv_local_out_cache_loc", return_value=torch.tensor([1, 2], dtype=torch.int64), ), ): with self.assertRaisesRegex( RuntimeError, "CP_SHARED_KV_FAIL_FAST.*index_local_shape_mismatch", ): Indexer._store_cp_shared_local_index_k_cache( indexer, forward_batch, layer_id=0, local_key=torch.empty((1, 8)), act_quant=None, ) def test_indexer_direct_write_filters_compute_padding_rows(self): import torch from sglang.srt.layers.attention.nsa import nsa_indexer from sglang.srt.layers.attention.nsa.nsa_indexer import Indexer page_size = 64 plan = build_batch_page_aligned_in_seq_split_plan( extend_lens=[65], prefix_lens=[0], page_size=page_size, cp_size=8, cp_rank=1, ) out_cache_loc = torch.cat( [ torch.arange(page_size, 2 * page_size), torch.tensor([2 * page_size]), ] ).to(torch.int64) forward_batch = SimpleNamespace( uses_cp_shared_kv=True, cp_shared_kv_layout=CpSharedKVLayout( page_size=page_size, cp_size=8, cp_rank=1, ), nsa_cp_metadata=NSAContextParallelMetadata( batch_size=1, batch_plan=plan, page_aligned=True, page_size=page_size, extend_prefix_len=0, ), out_cache_loc=out_cache_loc, token_to_kv_pool=SimpleNamespace(page_size=page_size), ) indexer = object.__new__(Indexer) indexer.nsa_enable_prefill_cp = True calls = [] def fake_store_index_k_cache(**kwargs): calls.append(kwargs) indexer._store_index_k_cache = fake_store_index_k_cache local_key = torch.full((64, 2), -1.0) local_key[0] = torch.tensor([9.0, 10.0]) with patch.object(nsa_indexer, "nsa_use_prefill_cp", return_value=True): stored = Indexer._store_cp_shared_local_index_k_cache( indexer, forward_batch, layer_id=0, local_key=local_key, act_quant=None, ) self.assertTrue(stored) self.assertEqual(len(calls), 1) self.assertEqual(calls[0]["key"].tolist(), [[9.0, 10.0]]) def test_mla_direct_write_fails_fast_on_local_shape_mismatch(self): import torch from sglang.srt.models.deepseek_common.attention_forward_methods import ( forward_mla, ) mla = SimpleNamespace(attn_mqa=SimpleNamespace(layer_id=3)) forward_batch = SimpleNamespace() with patch.object( forward_mla, "get_cp_shared_kv_local_out_cache_loc", return_value=torch.tensor([1, 2], dtype=torch.int64), ): with self.assertRaisesRegex( RuntimeError, "CP_SHARED_KV_FAIL_FAST.*mla_local_shape_mismatch", ): forward_mla.DeepseekMLAForwardMixin._maybe_write_cp_shared_local_mla_kv( mla, forward_batch, k_nope=torch.empty((1, 8)), k_pe=torch.empty((2, 8)), ) def test_mla_direct_write_filters_compute_padding_rows(self): import torch from sglang.srt.models.deepseek_common.attention_forward_methods import ( forward_mla, ) page_size = 64 plan = build_batch_page_aligned_in_seq_split_plan( extend_lens=[65], prefix_lens=[0], page_size=page_size, cp_size=8, cp_rank=1, ) out_cache_loc = torch.cat( [ torch.arange(page_size, 2 * page_size), torch.tensor([2 * page_size]), ] ).to(torch.int64) forward_batch = SimpleNamespace( uses_cp_shared_kv=True, cp_shared_kv_layout=CpSharedKVLayout( page_size=page_size, cp_size=8, cp_rank=1, ), nsa_cp_metadata=NSAContextParallelMetadata( batch_size=1, batch_plan=plan, page_aligned=True, page_size=page_size, extend_prefix_len=0, ), out_cache_loc=out_cache_loc, token_to_kv_pool=SimpleNamespace(page_size=page_size), ) mla = SimpleNamespace(attn_mqa=SimpleNamespace(layer_id=0)) calls = [] def fake_tai_store(**kwargs): calls.append(kwargs) return True k_nope = torch.full((64, 2), -1.0) k_nope[0] = torch.tensor([1.0, 2.0]) k_pe = torch.full((64, 2), -1.0) k_pe[0] = torch.tensor([3.0, 4.0]) with patch.object(forward_mla, "try_tai_fused_mla_store", fake_tai_store): stored = ( forward_mla.DeepseekMLAForwardMixin._maybe_write_cp_shared_local_mla_kv( mla, forward_batch, k_nope=k_nope, k_pe=k_pe, ) ) self.assertTrue(stored) self.assertEqual(len(calls), 1) self.assertEqual(calls[0]["k_nope"].tolist(), [[1.0, 2.0]]) self.assertEqual(calls[0]["k_rope"].tolist(), [[3.0, 4.0]]) self.assertEqual(calls[0]["logical_locs"].tolist(), [2 * page_size]) def test_index_partial_current_compose_accepts_local_valid_compute_padding_rows( self, ): import torch from sglang.srt.layers.attention.nsa import nsa_indexer from sglang.srt.layers.attention.nsa.nsa_indexer import Indexer page_size = 64 plan = build_batch_page_aligned_in_seq_split_plan( extend_lens=[65], prefix_lens=[0], page_size=page_size, cp_size=8, cp_rank=1, ) class FakePool: page_size = 64 index_head_dim = 2 def get_index_k_with_scale_buffer(self, layer_id): return torch.zeros((4, 3), dtype=torch.float32) forward_batch = SimpleNamespace( uses_cp_shared_kv=True, cp_shared_kv_layout=CpSharedKVLayout( page_size=page_size, cp_size=8, cp_rank=1, ), cp_shared_kv_index_prefetcher=None, token_to_kv_pool=FakePool(), nsa_cp_metadata=NSAContextParallelMetadata( batch_size=1, batch_plan=plan, page_aligned=True, page_size=page_size, extend_prefix_len=0, ), out_cache_loc=torch.cat( [ torch.arange(page_size, 2 * page_size), torch.tensor([2 * page_size]), ] ).to(torch.int64), extend_prefix_lens_cpu=[page_size], extend_seq_lens_cpu=[65], ) logical_page_table = torch.tensor([[1, 2, 3]], dtype=torch.int32) current_index_kv = ( torch.tensor([[7.0, 8.0]], dtype=torch.float32), torch.tensor([[0.5]], dtype=torch.float32), ) materialize_calls = [] expected_buffer = torch.ones((3, 3), dtype=torch.float32) expected_pages = torch.tensor([[0, 1, 2]], dtype=torch.int32) indexer = object.__new__(Indexer) def fake_materialize(**kwargs): materialize_calls.append(kwargs) return expected_buffer, expected_pages with patch.object( nsa_indexer, "get_or_build_shared_paged_buffer_slot_remap", return_value=torch.tensor([0, 1, 2], dtype=torch.int64), ), patch.object( nsa_indexer, "materialize_prefix_and_reuse_current_index_page_slots", side_effect=fake_materialize, ): dense_buffer, dense_pages = indexer._maybe_materialize_shared_index_buffer( forward_batch, layer_id=0, logical_page_table=logical_page_table, current_index_kv=current_index_kv, ) self.assertIs(dense_buffer, expected_buffer) self.assertIs(dense_pages, expected_pages) self.assertEqual(len(materialize_calls), 1) self.assertIs(materialize_calls[0]["current_index_k"], current_index_kv[0]) self.assertEqual(materialize_calls[0]["current_locs"].tolist(), [2 * page_size]) def test_indexer_current_reuse_compute_padding_selects_local_key_not_gathered_key( self, ): import torch import types from sglang.srt.layers.attention.nsa import nsa_indexer from sglang.srt.layers.attention.nsa.nsa_indexer import Indexer page_size = 64 plan = build_batch_page_aligned_in_seq_split_plan( extend_lens=[65], prefix_lens=[0], page_size=page_size, cp_size=8, cp_rank=1, ) metadata_obj = _build_batch_metadata_from_plan(plan) class Mode: def is_extend_without_speculative(self): return True def is_decode_or_idle(self): return False def is_target_verify(self): return False def is_draft_extend(self, include_v2=False): return False def is_context_parallel_extend(self): return True class AttnBackend: def get_indexer_metadata(self, layer_id, forward_batch): return object() forward_batch = SimpleNamespace( uses_cp_shared_kv=True, cp_shared_kv_layout=CpSharedKVLayout( page_size=page_size, cp_size=8, cp_rank=1, ), token_to_kv_pool=SimpleNamespace(page_size=page_size), nsa_cp_metadata=metadata_obj, out_cache_loc=torch.cat( [ torch.arange(page_size, 2 * page_size), torch.tensor([2 * page_size]), ] ).to(torch.int64), extend_prefix_lens_cpu=[0], extend_seq_lens_cpu=[65], seq_lens_cpu=torch.tensor([65], dtype=torch.int64), forward_mode=Mode(), attn_backend=AttnBackend(), hisparse_coordinator=None, ) local_key = torch.full((64, 2), -1.0, dtype=torch.float32) local_key[0] = torch.tensor([11.0, 12.0]) gathered_key = torch.full((64, 2), 99.0, dtype=torch.float32) gathered_key[0] = torch.tensor([101.0, 102.0]) query = torch.zeros((64, 2), dtype=torch.float32) act_quant_inputs = [] topk_current_index_kv = [] def fake_act_quant(tensor, block_size, scale_fmt): act_quant_inputs.append(tensor.detach().clone()) return tensor.detach().clone(), torch.ones( (int(tensor.shape[0]), 1), dtype=torch.float32 ) fake_triton_kernel = types.ModuleType( "sglang.srt.layers.attention.nsa.triton_kernel" ) fake_triton_kernel.act_quant = fake_act_quant indexer = object.__new__(Indexer) indexer.alt_stream = None indexer.nsa_enable_prefill_cp = True indexer.index_topk = 2 indexer.block_size = 64 indexer.scale_fmt = None indexer._get_q_k_bf16 = ( lambda *args, **kwargs: (query, gathered_key, local_key) ) indexer._store_cp_shared_local_index_k_cache = lambda **kwargs: True indexer._can_reuse_current_index_kv = lambda forward_batch: True indexer._get_logits_head_gate = ( lambda x_for_gate, q_scale: torch.zeros((64, 1), dtype=torch.float32) ) def fake_topk(*args, **kwargs): topk_current_index_kv.append(kwargs["current_index_kv"]) return torch.zeros((64, 2), dtype=torch.int32) indexer._get_topk_in_seq_cp_pair = fake_topk with ( patch.dict( sys.modules, { "sglang.srt.layers.attention.nsa.triton_kernel": fake_triton_kernel }, ), patch.object(nsa_indexer, "_is_cuda", True), patch.object(nsa_indexer, "_is_hip", False), patch.object(nsa_indexer, "_is_npu", False), patch.object( nsa_indexer, "is_nsa_prefill_cp_in_seq_split", return_value=True, ), ): result = Indexer.forward_cuda( indexer, x=torch.zeros((64, 2), dtype=torch.float32), q_lora=torch.zeros((64, 2), dtype=torch.float32), positions=torch.arange(64, dtype=torch.int64), forward_batch=forward_batch, layer_id=0, return_indices=True, ) self.assertEqual(result.shape, (64, 2)) self.assertEqual(len(act_quant_inputs), 2) self.assertEqual(act_quant_inputs[1].tolist(), [[11.0, 12.0]]) self.assertNotEqual(act_quant_inputs[1].tolist(), [[101.0, 102.0]]) self.assertEqual(len(topk_current_index_kv), 1) self.assertEqual(topk_current_index_kv[0][0].tolist(), [[11.0, 12.0]]) def test_indexer_direct_write_does_not_log_missing_metadata_for_non_cp_batch(self): import torch from types import SimpleNamespace from sglang.srt.layers.attention.nsa.nsa_indexer import Indexer indexer = object.__new__(Indexer) indexer.nsa_enable_prefill_cp = True forward_batch = SimpleNamespace( uses_cp_shared_kv=True, nsa_cp_metadata=None, ) with self.assertNoLogs( "sglang.srt.layers.attention.nsa.utils", level="INFO" ): stored = Indexer._store_cp_shared_local_index_k_cache( indexer, forward_batch, layer_id=0, local_key=torch.empty(0), act_quant=None, ) self.assertFalse(stored) def test_indexer_in_seq_cp_pair_materializes_index_once_for_prev_next(self): import torch from sglang.srt.layers.attention.nsa.nsa_indexer import Indexer indexer = object.__new__(Indexer) logical_pages = torch.tensor([[1, 2, 3, 4]], dtype=torch.int32) materialized_index = torch.tensor([11], dtype=torch.int32) dense_pages = torch.tensor([[1, 2, 3, 4]], dtype=torch.int32) materialize_calls = [] topk_calls = [] class Metadata: def get_page_table_64(self): return logical_pages def get_page_table_1(self): return torch.empty((1, 1000), dtype=torch.int32) def fake_materialize(forward_batch, layer_id, logical_page_table): materialize_calls.append((layer_id, logical_page_table)) return materialized_index, dense_pages def fake_get_topk( forward_batch, layer_id, q_fp8, weights, metadata, kv_len, actual_seq_q, cp_index=None, current_index_kv=None, shared_index_buffer=None, shared_block_tables=None, actual_seq_q_tensor=None, actual_seq_q_cu_tensor=None, ): topk_calls.append( { "kv_len": kv_len, "actual_seq_q": actual_seq_q, "actual_seq_q_tensor": actual_seq_q_tensor, "actual_seq_q_cu_tensor": actual_seq_q_cu_tensor, "shared_index_buffer": shared_index_buffer, "shared_block_tables": shared_block_tables, "current_index_kv": current_index_kv, } ) return torch.full((actual_seq_q, 2), len(topk_calls), dtype=torch.int32) indexer._maybe_materialize_shared_index_buffer = fake_materialize indexer._get_topk_ragged_with_cp = fake_get_topk class Mode: def is_extend_without_speculative(self): return True forward_batch = type( "ForwardBatchStub", (), { "forward_mode": Mode(), "extend_prefix_lens_cpu": [0], "extend_seq_lens_cpu": [5], "seq_lens_cpu": torch.tensor([5], dtype=torch.int64), "nsa_cp_metadata": NSAContextParallelMetadata( kv_len_prev=5, kv_len_next=9, actual_seq_q_prev=3, actual_seq_q_next=2, actual_seq_q_prev_cu_tensor=torch.tensor([0, 3], dtype=torch.int32), actual_seq_q_next_cu_tensor=torch.tensor([0, 2], dtype=torch.int32), ) }, )() q_fp8 = torch.arange(5 * 4, dtype=torch.float32).view(5, 4) weights = torch.arange(5 * 2, dtype=torch.float32).view(5, 2) result = Indexer._get_topk_in_seq_cp_pair( indexer, forward_batch, layer_id=7, q_fp8=q_fp8, weights=weights, metadata=Metadata(), current_index_kv=None, ) self.assertEqual(len(materialize_calls), 1) self.assertIs(materialize_calls[0][1], logical_pages) self.assertEqual(len(topk_calls), 2) self.assertIs(topk_calls[0]["shared_index_buffer"], materialized_index) self.assertIs(topk_calls[1]["shared_index_buffer"], materialized_index) self.assertIs(topk_calls[0]["shared_block_tables"], dense_pages) self.assertIs(topk_calls[1]["shared_block_tables"], dense_pages) self.assertIsNone(topk_calls[0]["current_index_kv"]) self.assertEqual(topk_calls[0]["kv_len"], 5) self.assertEqual(topk_calls[1]["kv_len"], 9) self.assertEqual(topk_calls[0]["actual_seq_q_cu_tensor"].tolist(), [0, 3]) self.assertEqual(topk_calls[1]["actual_seq_q_cu_tensor"].tolist(), [0, 2]) self.assertEqual(result.tolist(), [[1, 1], [1, 1], [1, 1], [2, 2], [2, 2]]) def test_indexer_in_seq_cp_pair_batch_preserves_request_segment_order(self): import torch from sglang.srt.layers.attention.nsa.nsa_indexer import Indexer indexer = object.__new__(Indexer) indexer.index_topk = 2 logical_pages = torch.tensor([[1, 2, 3], [4, 5, 6]], dtype=torch.int32) materialized_index = torch.tensor([11], dtype=torch.int32) dense_pages = torch.tensor([[1, 2, 3], [4, 5, 6]], dtype=torch.int32) materialize_calls = [] topk_calls = [] class Metadata: def get_page_table_64(self): return logical_pages def get_page_table_1(self): return torch.empty((2, 1000), dtype=torch.int32) def fake_materialize(forward_batch, layer_id, logical_page_table): materialize_calls.append((layer_id, logical_page_table)) return materialized_index, dense_pages def fake_get_topk( forward_batch, layer_id, q_fp8, weights, metadata, kv_len, actual_seq_q, cp_index=None, current_index_kv=None, shared_index_buffer=None, shared_block_tables=None, actual_seq_q_tensor=None, actual_seq_q_cu_tensor=None, batch_idx=0, ): topk_calls.append( { "batch_idx": batch_idx, "kv_len": kv_len, "actual_seq_q": actual_seq_q, "cp_index": cp_index, "q": q_fp8.flatten().tolist(), "weights": weights.flatten().tolist(), "actual_seq_q_tensor": actual_seq_q_tensor, "actual_seq_q_cu_tensor": actual_seq_q_cu_tensor, "shared_index_buffer": shared_index_buffer, "shared_block_tables": shared_block_tables, "current_index_kv": current_index_kv, } ) rows = int(q_fp8.shape[0]) return torch.arange(1, rows + 1, dtype=torch.int32).view(rows, 1).repeat(1, 2) indexer._maybe_materialize_shared_index_buffer = fake_materialize indexer._get_topk_ragged_with_cp = fake_get_topk forward_batch = SimpleNamespace( batch_size=2, forward_mode=SimpleNamespace( is_extend_without_speculative=lambda: True, ), extend_prefix_lens_cpu=[0, 0], extend_seq_lens_cpu=[1000, 1000], seq_lens_cpu=torch.tensor([1000, 1000], dtype=torch.int64), nsa_cp_metadata=NSAContextParallelMetadata( batch_size=2, kv_len_prev=100, kv_len_next=200, actual_seq_q_prev=2, actual_seq_q_next=1, actual_seq_q_prev_cu_tensor=torch.tensor([0, 2], dtype=torch.int32), actual_seq_q_next_cu_tensor=torch.tensor([0, 1], dtype=torch.int32), request_kv_len_prev=[2, 1], request_kv_len_next=[3, 4], request_actual_seq_q_prev=[2, 1], request_actual_seq_q_next=[1, 3], ), ) q_fp8 = torch.arange(7, dtype=torch.float32).view(7, 1) weights = (torch.arange(7, dtype=torch.float32) + 100).view(7, 1) result = Indexer._get_topk_in_seq_cp_pair( indexer, forward_batch, layer_id=7, q_fp8=q_fp8, weights=weights, metadata=Metadata(), current_index_kv=None, ) self.assertEqual(len(materialize_calls), 1) self.assertIs(materialize_calls[0][1], logical_pages) self.assertEqual(len(topk_calls), 1) self.assertEqual(topk_calls[0]["batch_idx"], 0) self.assertEqual(topk_calls[0]["actual_seq_q"], 7) self.assertEqual( topk_calls[0]["cp_index"], [(0, 0, 2), (0, 2, 3), (1, 0, 1), (1, 1, 4)], ) self.assertEqual(topk_calls[0]["q"], [0.0, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0]) self.assertEqual( topk_calls[0]["weights"], [100.0, 101.0, 102.0, 103.0, 104.0, 105.0, 106.0], ) self.assertIsNone(topk_calls[0]["actual_seq_q_tensor"]) self.assertIsNone(topk_calls[0]["actual_seq_q_cu_tensor"]) self.assertTrue(all(call["shared_index_buffer"] is materialized_index for call in topk_calls)) self.assertTrue(all(call["shared_block_tables"] is dense_pages for call in topk_calls)) self.assertTrue(all(call["current_index_kv"] is None for call in topk_calls)) self.assertEqual( result.tolist(), [[1, 1], [2, 2], [3, 3], [4, 4], [5, 5], [6, 6], [7, 7]], ) def test_indexer_in_seq_cp_pair_compute_padding_outputs_dummy_safe_rows(self): import torch from sglang.srt.layers.attention.nsa.nsa_indexer import Indexer page_size = 64 plan = build_batch_page_aligned_in_seq_split_plan( extend_lens=[65], prefix_lens=[0], page_size=page_size, cp_size=8, cp_rank=1, ) metadata_obj = _build_batch_metadata_from_plan(plan) indexer = object.__new__(Indexer) indexer.index_topk = 2 logical_pages = torch.tensor([[1, 2]], dtype=torch.int32) materialized_index = torch.tensor([11], dtype=torch.int32) dense_pages = torch.tensor([[1, 2]], dtype=torch.int32) materialize_calls = [] topk_calls = [] class Metadata: def get_page_table_64(self): return logical_pages def get_page_table_1(self): return torch.empty((1, 65), dtype=torch.int32) def fake_materialize(forward_batch, layer_id, logical_page_table): materialize_calls.append((layer_id, logical_page_table)) return materialized_index, dense_pages def fake_get_topk( forward_batch, layer_id, q_fp8, weights, metadata, kv_len, actual_seq_q, cp_index=None, current_index_kv=None, shared_index_buffer=None, shared_block_tables=None, actual_seq_q_tensor=None, actual_seq_q_cu_tensor=None, batch_idx=0, ): topk_calls.append( { "actual_seq_q": actual_seq_q, "cp_index": cp_index, "q": q_fp8.flatten().tolist(), "weights": weights.flatten().tolist(), "shared_index_buffer": shared_index_buffer, "shared_block_tables": shared_block_tables, } ) rows = int(q_fp8.shape[0]) return ( torch.arange(1, rows + 1, dtype=torch.int32) .view(rows, 1) .repeat(1, 2) ) indexer._maybe_materialize_shared_index_buffer = fake_materialize indexer._get_topk_ragged_with_cp = fake_get_topk forward_batch = SimpleNamespace( batch_size=1, forward_mode=SimpleNamespace( is_extend_without_speculative=lambda: True, ), extend_prefix_lens_cpu=[0], extend_seq_lens_cpu=[65], seq_lens_cpu=torch.tensor([65], dtype=torch.int64), nsa_cp_metadata=metadata_obj, ) q_fp8 = torch.arange(64, dtype=torch.float32).view(64, 1) weights = (torch.arange(64, dtype=torch.float32) + 100).view(64, 1) result = Indexer._get_topk_in_seq_cp_pair( indexer, forward_batch, layer_id=7, q_fp8=q_fp8, weights=weights, metadata=Metadata(), current_index_kv=None, ) self.assertEqual(len(materialize_calls), 1) self.assertIs(materialize_calls[0][1], logical_pages) self.assertEqual(len(topk_calls), 1) self.assertEqual(topk_calls[0]["actual_seq_q"], 1) self.assertEqual(topk_calls[0]["cp_index"], [(0, 64, 65)]) self.assertEqual(topk_calls[0]["q"], [0.0]) self.assertEqual(topk_calls[0]["weights"], [100.0]) self.assertIs(topk_calls[0]["shared_index_buffer"], materialized_index) self.assertIs(topk_calls[0]["shared_block_tables"], dense_pages) self.assertEqual(result.shape, (64, 2)) self.assertEqual(result[0].tolist(), [1, 1]) self.assertTrue( torch.equal(result[1:], torch.full((63, 2), -1, dtype=torch.int32)) ) def test_indexer_in_seq_cp_pair_batch_materializes_partial_current_index_reuse_once(self): import torch from sglang.srt.layers.attention.nsa.nsa_indexer import Indexer indexer = object.__new__(Indexer) indexer.index_topk = 2 current_index_kv = (torch.arange(7), torch.arange(7)) logical_pages = torch.tensor([[1, 2, 3], [4, 5, 6]], dtype=torch.int32) materialized_index = torch.tensor([11], dtype=torch.int32) dense_pages = torch.tensor([[1, 2, 3], [4, 5, 6]], dtype=torch.int32) materialize_calls = [] topk_calls = [] class Metadata: def get_page_table_64(self): return logical_pages def get_page_table_1(self): return torch.empty((2, 512), dtype=torch.int32) def fake_materialize( forward_batch, layer_id, logical_page_table, current_index_kv=None, ): materialize_calls.append( { "layer_id": layer_id, "logical_page_table": logical_page_table, "current_index_kv": current_index_kv, } ) return materialized_index, dense_pages def fake_get_topk( forward_batch, layer_id, q_fp8, weights, metadata, kv_len, actual_seq_q, cp_index=None, current_index_kv=None, shared_index_buffer=None, shared_block_tables=None, actual_seq_q_tensor=None, actual_seq_q_cu_tensor=None, batch_idx=0, ): topk_calls.append( { "batch_idx": batch_idx, "actual_seq_q": actual_seq_q, "cp_index": cp_index, "q_rows": int(q_fp8.shape[0]), "current_index_kv": current_index_kv, "shared_index_buffer": shared_index_buffer, "shared_block_tables": shared_block_tables, "actual_seq_q_cu_tensor": actual_seq_q_cu_tensor, } ) rows = int(q_fp8.shape[0]) return torch.arange(1, rows + 1, dtype=torch.int32).view(rows, 1).repeat(1, 2) indexer._maybe_materialize_shared_index_buffer = fake_materialize indexer._get_topk_ragged_with_cp = fake_get_topk forward_batch = SimpleNamespace( batch_size=2, forward_mode=SimpleNamespace( is_extend_without_speculative=lambda: True, ), extend_prefix_lens_cpu=[64, 64], extend_seq_lens_cpu=[3, 4], seq_lens_cpu=torch.tensor([67, 68], dtype=torch.int64), nsa_cp_metadata=NSAContextParallelMetadata( batch_size=2, request_kv_len_prev=[2, 1], request_kv_len_next=[3, 4], request_actual_seq_q_prev=[2, 1], request_actual_seq_q_next=[1, 3], ), ) result = Indexer._get_topk_in_seq_cp_pair( indexer, forward_batch, layer_id=7, q_fp8=torch.empty(7, 1), weights=torch.empty(7, 1), metadata=Metadata(), current_index_kv=current_index_kv, ) self.assertEqual(len(materialize_calls), 1) self.assertIs(materialize_calls[0]["logical_page_table"], logical_pages) self.assertIs(materialize_calls[0]["current_index_kv"], current_index_kv) self.assertEqual(len(topk_calls), 1) self.assertEqual(topk_calls[0]["actual_seq_q"], 7) self.assertEqual( topk_calls[0]["cp_index"], [(0, 0, 2), (0, 2, 3), (1, 0, 1), (1, 1, 4)], ) self.assertEqual(topk_calls[0]["q_rows"], 7) self.assertTrue( all(call["current_index_kv"] is None for call in topk_calls) ) self.assertTrue( all(call["shared_index_buffer"] is materialized_index for call in topk_calls) ) self.assertTrue( all(call["shared_block_tables"] is dense_pages for call in topk_calls) ) self.assertEqual([call["batch_idx"] for call in topk_calls], [0]) self.assertIsNone(topk_calls[0]["actual_seq_q_cu_tensor"]) self.assertEqual( result.tolist(), [[1, 1], [2, 2], [3, 3], [4, 4], [5, 5], [6, 6], [7, 7]], ) def test_indexer_in_seq_cp_pair_batch_composes_current_only_index_reuse(self): import torch from sglang.srt.layers.attention.nsa.nsa_indexer import Indexer indexer = object.__new__(Indexer) indexer.index_topk = 2 current_index_kv = (torch.arange(7), torch.arange(7)) logical_pages = torch.tensor([[1, 2], [3, 4]], dtype=torch.int32) materialized_index = torch.tensor([17], dtype=torch.int32) dense_pages = torch.tensor([[10, 11], [12, 13]], dtype=torch.int32) materialize_calls = [] topk_calls = [] class Mode: def is_extend_without_speculative(self): return True class Metadata: def get_page_table_64(self): return logical_pages def get_page_table_1(self): return torch.empty((2, 512), dtype=torch.int32) def fake_materialize( forward_batch, layer_id, logical_page_table, current_index_kv=None, ): materialize_calls.append( { "layer_id": layer_id, "logical_page_table": logical_page_table, "current_index_kv": current_index_kv, } ) return materialized_index, dense_pages def fake_get_topk( forward_batch, layer_id, q_fp8, weights, metadata, kv_len, actual_seq_q, cp_index=None, current_index_kv=None, shared_index_buffer=None, shared_block_tables=None, actual_seq_q_tensor=None, actual_seq_q_cu_tensor=None, batch_idx=0, ): topk_calls.append( { "batch_idx": batch_idx, "actual_seq_q": actual_seq_q, "cp_index": cp_index, "q_rows": int(q_fp8.shape[0]), "current_index_kv": current_index_kv, "shared_index_buffer": shared_index_buffer, "shared_block_tables": shared_block_tables, "actual_seq_q_cu_tensor": actual_seq_q_cu_tensor, } ) rows = int(q_fp8.shape[0]) return torch.arange(1, rows + 1, dtype=torch.int32).view(rows, 1).repeat(1, 2) indexer._maybe_materialize_shared_index_buffer = fake_materialize indexer._get_topk_ragged_with_cp = fake_get_topk forward_batch = SimpleNamespace( batch_size=2, forward_mode=Mode(), extend_prefix_lens_cpu=[0, 0], extend_seq_lens_cpu=[3, 4], seq_lens_cpu=torch.tensor([3, 4], dtype=torch.int64), nsa_cp_metadata=NSAContextParallelMetadata( batch_size=2, request_kv_len_prev=[3, 4], request_kv_len_next=[3, 4], request_actual_seq_q_prev=[2, 1], request_actual_seq_q_next=[1, 3], ), ) result = Indexer._get_topk_in_seq_cp_pair( indexer, forward_batch, layer_id=7, q_fp8=torch.empty(7, 1), weights=torch.empty(7, 1), metadata=Metadata(), current_index_kv=current_index_kv, ) self.assertEqual(len(materialize_calls), 1) self.assertIs(materialize_calls[0]["logical_page_table"], logical_pages) self.assertIs(materialize_calls[0]["current_index_kv"], current_index_kv) self.assertEqual(len(topk_calls), 1) self.assertTrue(all(call["current_index_kv"] is None for call in topk_calls)) self.assertTrue( all(call["shared_index_buffer"] is materialized_index for call in topk_calls) ) self.assertTrue( all(call["shared_block_tables"] is dense_pages for call in topk_calls) ) self.assertEqual(topk_calls[0]["batch_idx"], 0) self.assertEqual(topk_calls[0]["actual_seq_q"], 7) self.assertEqual( topk_calls[0]["cp_index"], [(0, 1, 3), (0, 2, 3), (1, 3, 4), (1, 1, 4)], ) self.assertEqual(topk_calls[0]["q_rows"], 7) self.assertIsNone(topk_calls[0]["actual_seq_q_cu_tensor"]) self.assertEqual( result.tolist(), [[1, 1], [2, 2], [3, 3], [4, 4], [5, 5], [6, 6], [7, 7]], ) def test_indexer_shared_index_materialize_accepts_current_only_compose(self): import torch from sglang.srt.layers.attention.nsa import nsa_indexer from sglang.srt.layers.attention.nsa.nsa_indexer import Indexer indexer = object.__new__(Indexer) page_buffer = torch.zeros((8, 32), dtype=torch.uint8) logical_pages = torch.tensor([[1, 2]], dtype=torch.int64) current_index_kv = ( torch.zeros((2, 4), dtype=torch.uint8), torch.zeros((2, 1), dtype=torch.float32), ) compose_calls = [] class Pool: page_size = 4 index_head_dim = 4 def get_index_k_with_scale_buffer(self, layer_id): return page_buffer class Mode: def is_extend_without_speculative(self): return True forward_batch = SimpleNamespace( uses_cp_shared_kv=True, cp_shared_kv_layout=CpSharedKVLayout(page_size=4, cp_size=2, cp_rank=0), token_to_kv_pool=Pool(), cp_shared_kv_index_prefetcher=None, forward_mode=Mode(), extend_prefix_lens_cpu=[0], extend_seq_lens_cpu=[5], seq_lens_cpu=torch.tensor([5], dtype=torch.int64), cp_local_out_cache_loc=torch.tensor([4, 5], dtype=torch.int64), ) def fake_compose(**kwargs): compose_calls.append(kwargs) return torch.empty((3, 32), dtype=torch.uint8), torch.tensor([[1, 2]]) with patch.object( nsa_indexer, "materialize_prefix_and_reuse_current_index_page_slots", side_effect=fake_compose, ): materialized, dense_pages = Indexer._maybe_materialize_shared_index_buffer( indexer, forward_batch, layer_id=7, logical_page_table=logical_pages, current_index_kv=current_index_kv, ) self.assertEqual(len(compose_calls), 1) self.assertEqual(compose_calls[0]["prefix_pages"], 0) self.assertEqual(compose_calls[0]["current_slot_spans"], [(0, 2)]) self.assertIs(compose_calls[0]["current_index_k"], current_index_kv[0]) self.assertIs(compose_calls[0]["current_index_scale"], current_index_kv[1]) self.assertEqual(list(materialized.shape), [3, 32]) self.assertEqual(dense_pages.tolist(), [[1, 2]]) def test_indexer_ragged_cp_index_current_batch_does_not_materialize(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() deep_gemm_calls = [] def fake_logits(q_fp8, kv_fp8, weights, ks, ke, clean_logits=False): deep_gemm_calls.append( { "q_rows": int(q_fp8.shape[0]), "kv_rows": int(kv_fp8[0].shape[0]), "weights_rows": int(weights.shape[0]), "ks": ks.tolist(), "ke": ke.tolist(), } ) 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): return ( torch.arange(1, int(logits.shape[0]) + 1, dtype=torch.int32) .view(-1, 1) .repeat(1, topk) ) forward_batch = SimpleNamespace( token_to_kv_pool=SimpleNamespace(page_size=64), seq_lens_cpu=torch.tensor([3, 4], dtype=torch.int64), extend_seq_lens_cpu=[3, 4], ) q_fp8 = torch.empty((7, 1), dtype=torch.float32) weights = torch.empty((7, 1, 1), dtype=torch.float32) 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, "deep_gemm", SimpleNamespace(fp8_mqa_logits=fake_logits), ): result = Indexer._get_topk_ragged_with_cp( indexer, forward_batch, layer_id=7, q_fp8=q_fp8, weights=weights, 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(result.tolist(), [[1, 1], [2, 2], [3, 3], [4, 4], [5, 5], [6, 6], [7, 7]]) self.assertEqual(len(deep_gemm_calls), 1) self.assertEqual(deep_gemm_calls[0]["q_rows"], 7) self.assertEqual(deep_gemm_calls[0]["weights_rows"], 7) self.assertEqual(deep_gemm_calls[0]["kv_rows"], 14) 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 from types import SimpleNamespace from unittest.mock import patch from sglang.srt.layers.attention.nsa import index_buf_accessor, 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() prepare_calls = [] deep_gemm_calls = [] def fake_prepare(**kwargs): prepare_calls.append(kwargs) total_kv_len = int(kwargs["total_kv_len"]) total_q_count = int(kwargs["total_q_count"]) 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): deep_gemm_calls.append( { "q_rows": int(q_fp8.shape[0]), "kv_rows": int(kv_fp8[0].shape[0]), "weights_rows": int(weights.shape[0]), "ks": ks.tolist(), "ke": ke.tolist(), } ) return torch.zeros((int(q_fp8.shape[0]), 8), dtype=torch.float32) class Metadata: def topk_transform(self, logits, topk, **kwargs): return ( torch.arange(1, int(logits.shape[0]) + 1, dtype=torch.int32) .view(-1, 1) .repeat(1, topk) ) 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], ) q_fp8 = torch.empty((7, 1), dtype=torch.float32) weights = torch.empty((7, 1, 1), dtype=torch.float32) shared_index_buffer = torch.zeros((8, 264), dtype=torch.uint8) shared_block_tables = torch.arange(8, dtype=torch.int64).view(2, 4) with patch.object( nsa_indexer, "try_tai_prepare_cp_mqa_index_batch", side_effect=fake_prepare, create=True, ), patch.object( index_buf_accessor.GetK, "execute", side_effect=AssertionError("batched path must not call per-segment GetK"), ), patch.object( index_buf_accessor.GetS, "execute", side_effect=AssertionError("batched path must not call per-segment GetS"), ), patch.object( nsa_indexer, "deep_gemm", SimpleNamespace(fp8_mqa_logits=fake_logits), ): result = Indexer._get_topk_ragged_with_cp( indexer, forward_batch, layer_id=7, q_fp8=q_fp8, weights=weights, metadata=Metadata(), kv_len=0, actual_seq_q=7, cp_index=[(0, 1, 3), (0, 2, 3), (1, 3, 4), (1, 1, 4)], shared_index_buffer=shared_index_buffer, shared_block_tables=shared_block_tables, ) self.assertEqual(result.tolist(), [[1, 1], [2, 2], [3, 3], [4, 4], [5, 5], [6, 6], [7, 7]]) self.assertEqual(len(prepare_calls), 1) call = prepare_calls[0] self.assertIs(call["index_buffer"], shared_index_buffer) self.assertIs(call["block_tables"], shared_block_tables) self.assertEqual(call["batch_indices"].dtype, torch.int32) self.assertEqual(call["batch_indices"].tolist(), [0, 0, 1, 1]) self.assertEqual(call["kv_lens"].dtype, torch.int32) self.assertEqual(call["kv_lens"].tolist(), [3, 3, 4, 4]) self.assertEqual(call["q_starts"].tolist(), [1, 2, 3, 1]) self.assertEqual(call["q_lens"].tolist(), [2, 1, 1, 3]) self.assertEqual(call["k_bases"].tolist(), [0, 3, 6, 10]) self.assertEqual(call["q_bases"].tolist(), [0, 2, 3, 4]) self.assertEqual(call["total_kv_len"], 14) self.assertEqual(call["total_q_count"], 7) self.assertEqual(call["max_kv_len"], 4) self.assertEqual(call["max_q_len"], 3) self.assertEqual(len(deep_gemm_calls), 1) self.assertEqual(deep_gemm_calls[0]["q_rows"], 7) self.assertEqual(deep_gemm_calls[0]["weights_rows"], 7) self.assertEqual(deep_gemm_calls[0]["kv_rows"], 14) 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_current_batch_uses_tai_compact_once(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() prepare_calls = [] deep_gemm_calls = [] def fake_prepare(**kwargs): prepare_calls.append(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): deep_gemm_calls.append( { "q_rows": int(q_fp8.shape[0]), "kv_rows": int(kv_fp8[0].shape[0]), "ks": ks.tolist(), "ke": ke.tolist(), } ) 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): return ( torch.arange(1, int(logits.shape[0]) + 1, dtype=torch.int32) .view(-1, 1) .repeat(1, topk) ) 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), ): result = 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(result.tolist(), [[1, 1], [2, 2], [3, 3], [4, 4], [5, 5], [6, 6], [7, 7]]) self.assertEqual(len(prepare_calls), 1) call = prepare_calls[0] self.assertIs(call["current_index_k"], current_index_kv[0]) self.assertIs(call["current_index_scale"], current_index_kv[1]) self.assertEqual(call["current_bases"].tolist(), [0, 0, 3, 3]) self.assertEqual(call["kv_lens"].tolist(), [3, 3, 4, 4]) self.assertEqual(call["q_starts"].tolist(), [1, 2, 3, 1]) self.assertEqual(call["q_lens"].tolist(), [2, 1, 1, 3]) self.assertEqual(call["k_bases"].tolist(), [0, 3, 6, 10]) self.assertEqual(call["q_bases"].tolist(), [0, 2, 3, 4]) self.assertEqual(call["total_kv_len"], 14) self.assertEqual(call["total_q_count"], 7) self.assertEqual(len(deep_gemm_calls), 1) self.assertEqual(deep_gemm_calls[0]["kv_rows"], 14) 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_current_batch_uses_cp_local_bases_and_uint8_k( 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() prepare_calls = [] def fake_prepare(**kwargs): prepare_calls.append(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, 3, 3], dtype=torch.int32), torch.tensor([2, 3, 2, 3, 4], dtype=torch.int32), ) 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): return ( torch.arange(1, int(logits.shape[0]) + 1, dtype=torch.int32) .view(-1, 1) .repeat(1, topk) ) forward_batch = SimpleNamespace( token_to_kv_pool=SimpleNamespace(page_size=64, index_head_dim=1), seq_lens_cpu=torch.tensor([5, 7], dtype=torch.int64), extend_seq_lens_cpu=[5, 7], nsa_cp_metadata=NSAContextParallelMetadata( batch_size=2, batch_plan=SimpleNamespace( request_rank_local_offsets=[0, 2], request_valid_rank_local_offsets=[0, 2], ), ), ) current_index_k = ( torch.arange(6, dtype=torch.float32) .to(torch.float8_e4m3fn) .view(6, 1) ) current_index_scale = torch.arange(6, dtype=torch.float32).view(6, 1) current_index_kv = (current_index_k, current_index_scale) 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=lambda q_fp8, kv_fp8, weights, ks, ke, clean_logits=False: torch.zeros( (int(q_fp8.shape[0]), 8), dtype=torch.float32 ) ), ): result = Indexer._get_topk_ragged_with_cp( indexer, forward_batch, layer_id=7, q_fp8=torch.empty((5, 1), dtype=torch.float32), weights=torch.empty((5, 1, 1), dtype=torch.float32), metadata=Metadata(), kv_len=0, actual_seq_q=5, cp_index=[(0, 1, 3), (1, 1, 4)], current_index_kv=current_index_kv, ) self.assertEqual(result.tolist(), [[1, 1], [2, 2], [3, 3], [4, 4], [5, 5]]) self.assertEqual(len(prepare_calls), 1) call = prepare_calls[0] self.assertEqual(call["current_index_k"].dtype, torch.uint8) self.assertEqual(call["current_bases"].tolist(), [0, 2]) self.assertEqual(call["kv_lens"].tolist(), [3, 4]) self.assertEqual(call["q_lens"].tolist(), [2, 3]) def test_eagle_capture_for_decode_clears_cp_local_hidden_marker(self): import torch from sglang.srt.layers.logits_processor import LogitsProcessorOutput from sglang.srt.speculative.eagle_worker import EAGLEWorker worker = object.__new__(EAGLEWorker) worker.topk = 1 draft_input = EagleDraftInput( hidden_states=torch.full((4, 2), -1.0), cp_local_hidden_states=True, ) draft_output_hidden = torch.tensor([[1.0, 2.0], [3.0, 4.0]]) logits_output = LogitsProcessorOutput( next_token_logits=torch.tensor([[0.1, 0.9], [0.7, 0.3]]), hidden_states=draft_output_hidden, ) worker.capture_for_decode(logits_output, draft_input) self.assertIs(draft_input.hidden_states, draft_output_hidden) self.assertFalse(draft_input.cp_local_hidden_states) def test_eagle_v2_draft_extend_for_prefill_preserves_cp_local_hidden_marker( self, ): import torch from sglang.srt.layers.logits_processor import LogitsProcessorOutput from sglang.srt.speculative import eagle_worker_v2 worker = object.__new__(eagle_worker_v2.EagleDraftWorker) worker.topk = 1 forward_batch = SimpleNamespace(mm_input_embeds=None) class DraftRunner: def forward(self, _forward_batch): return SimpleNamespace( logits_output=LogitsProcessorOutput( next_token_logits=torch.tensor([[0.1, 0.9]]), hidden_states=torch.tensor([[3.0, 4.0]]), ) ) worker.draft_runner = DraftRunner() batch = SimpleNamespace( forward_mode=ForwardMode.EXTEND, extend_seq_lens=[2], input_ids=torch.tensor([10, 11]), seq_lens=[2], spec_info=None, ) target_hidden_states = torch.tensor([[1.0, 2.0]]) next_token_ids = torch.tensor([99]) init_new_checks = [] def init_new_side_effect(_batch, _draft_runner): init_new_checks.append(True) self.assertTrue(_batch.spec_info.cp_local_hidden_states) self.assertIs(_batch.spec_info.hidden_states, target_hidden_states) return forward_batch with patch.object( eagle_worker_v2.ForwardBatch, "init_new", side_effect=init_new_side_effect, ): next_draft_input = ( eagle_worker_v2.EagleDraftWorker._draft_extend_for_prefill( worker, batch, target_hidden_states, next_token_ids, cp_local_hidden_states=True, ) ) self.assertEqual(init_new_checks, [True]) self.assertTrue(next_draft_input.cp_local_hidden_states) self.assertEqual(next_draft_input.hidden_states.tolist(), [[3.0, 4.0]]) def test_indexer_in_seq_cp_pair_composes_current_only_index_reuse(self): import torch from sglang.srt.layers.attention.nsa.nsa_indexer import Indexer indexer = object.__new__(Indexer) current_index_kv = (torch.tensor([1]), torch.tensor([2])) logical_pages = torch.tensor([[1, 2]], dtype=torch.int32) materialized_index = torch.tensor([7], dtype=torch.int32) dense_pages = torch.tensor([[3, 4]], dtype=torch.int32) materialize_calls = [] topk_calls = [] class Metadata: def get_page_table_64(self): return logical_pages def get_page_table_1(self): return torch.empty((1, 512), dtype=torch.int32) def fake_materialize( forward_batch, layer_id, logical_page_table, current_index_kv=None, ): materialize_calls.append( { "layer_id": layer_id, "logical_page_table": logical_page_table, "current_index_kv": current_index_kv, } ) return materialized_index, dense_pages def fake_get_topk( forward_batch, layer_id, q_fp8, weights, metadata, kv_len, actual_seq_q, cp_index=None, current_index_kv=None, shared_index_buffer=None, shared_block_tables=None, actual_seq_q_tensor=None, actual_seq_q_cu_tensor=None, ): topk_calls.append( { "current_index_kv": current_index_kv, "actual_seq_q_tensor": actual_seq_q_tensor, "actual_seq_q_cu_tensor": actual_seq_q_cu_tensor, "shared_index_buffer": shared_index_buffer, "shared_block_tables": shared_block_tables, } ) return torch.full((actual_seq_q, 2), len(topk_calls), dtype=torch.int32) indexer._maybe_materialize_shared_index_buffer = fake_materialize indexer._get_topk_ragged_with_cp = fake_get_topk class Mode: def is_extend_without_speculative(self): return True forward_batch = type( "ForwardBatchStub", (), { "forward_mode": Mode(), "extend_prefix_lens_cpu": [0], "extend_seq_lens_cpu": [5], "seq_lens_cpu": torch.tensor([5], dtype=torch.int64), "nsa_cp_metadata": NSAContextParallelMetadata( kv_len_prev=5, kv_len_next=9, actual_seq_q_prev=3, actual_seq_q_next=2, actual_seq_q_prev_cu_tensor=torch.tensor([0, 3], dtype=torch.int32), actual_seq_q_next_cu_tensor=torch.tensor([0, 2], dtype=torch.int32), ) }, )() result = Indexer._get_topk_in_seq_cp_pair( indexer, forward_batch, layer_id=7, q_fp8=torch.empty(5, 4), weights=torch.empty(5, 2), metadata=Metadata(), current_index_kv=current_index_kv, ) self.assertEqual(len(materialize_calls), 1) self.assertIs(materialize_calls[0]["logical_page_table"], logical_pages) self.assertIs(materialize_calls[0]["current_index_kv"], current_index_kv) self.assertEqual(len(topk_calls), 2) self.assertIsNone(topk_calls[0]["current_index_kv"]) self.assertIsNone(topk_calls[1]["current_index_kv"]) self.assertIs(topk_calls[0]["shared_index_buffer"], materialized_index) self.assertIs(topk_calls[1]["shared_block_tables"], dense_pages) self.assertEqual(topk_calls[0]["actual_seq_q_cu_tensor"].tolist(), [0, 3]) self.assertEqual(topk_calls[1]["actual_seq_q_cu_tensor"].tolist(), [0, 2]) self.assertEqual(result.tolist(), [[1, 1], [1, 1], [1, 1], [2, 2], [2, 2]]) def test_paged_topk_transform_uses_cu_override_without_scan_metadata_ops(self): import torch from sglang.srt.layers.attention.nsa_backend import ( NSAMetadata, NSAIndexerMetadata, TopkTransformMethod, ) cu_override = torch.tensor([0, 4], dtype=torch.int32) attn_metadata = NSAMetadata( page_size=64, cache_seqlens_int32=torch.tensor([4], dtype=torch.int32), max_seq_len_q=4, max_seq_len_k=8, cu_seqlens_q=torch.tensor([0, 4], dtype=torch.int32), cu_seqlens_k=torch.tensor([0, 8], dtype=torch.int32), page_table_1=torch.arange(8, dtype=torch.int32).view(1, 8), real_page_table=torch.arange(8, dtype=torch.int32).view(1, 8), nsa_cache_seqlens_int32=torch.tensor([4], dtype=torch.int32), nsa_cu_seqlens_q=torch.arange(2, dtype=torch.int32), nsa_cu_seqlens_k=torch.tensor([0, 4], dtype=torch.int32), nsa_extend_seq_lens_list=[4], nsa_seqlens_expanded=torch.arange(1, 5, dtype=torch.int32), topk_indices_offset=torch.zeros(4, dtype=torch.int32), ) metadata = NSAIndexerMetadata( attn_metadata=attn_metadata, topk_transform_method=TopkTransformMethod.PAGED, ) logits = torch.zeros(4, 8) lengths = torch.arange(1, 5, dtype=torch.int32) expected = torch.full((4, 2), 7, dtype=torch.int32) def fake_fused(**kwargs): self.assertIs(kwargs["cu_seqlens_q"], cu_override) self.assertIs(kwargs["lengths"], lengths) return expected fake_sgl_kernel = SimpleNamespace( fast_topk_transform_fused=fake_fused, fast_topk_transform_ragged_fused=lambda **_: (_ for _ in ()).throw( AssertionError("ragged path should not run") ), fast_topk_v2=lambda *_, **__: (_ for _ in ()).throw( AssertionError("unfused path should not run") ), ) with ( patch.dict(sys.modules, {"sgl_kernel": fake_sgl_kernel}), patch( "sglang.srt.layers.attention.nsa_backend.envs.SGLANG_NSA_FUSE_TOPK.get", return_value=True, ), patch( "sglang.srt.layers.attention.nsa_backend.compute_cu_seqlens", side_effect=AssertionError("paged override should skip cumsum"), ), patch( "torch.repeat_interleave", side_effect=AssertionError("paged topk should not build ragged offsets"), ), ): actual = metadata.topk_transform( logits, topk=2, cu_seqlens_q=torch.tensor([4], dtype=torch.int32), ke_offset=lengths, cu_seqlens_q_topk_override=cu_override, ) self.assertIs(actual, expected) def test_mla_partial_current_path_fails_fast_instead_of_compact_fallback(self): backend_path = ( Path(__file__).resolve().parents[4] / "python/sglang/srt/layers/attention/nsa_backend.py" ) source = backend_path.read_text() tree = ast.parse(source) compact_merge_calls = [ node for node in ast.walk(tree) if isinstance(node, ast.Call) and isinstance(node.func, ast.Name) and node.func.id == "merge_materialized_and_current_kv" ] self.assertEqual( compact_merge_calls, [], "MLA partial-current reuse must not fall back to compact " "materialize/current merge when page-slot prefetch compose is unavailable.", ) self.assertIn( "[CP_SHARED_KV_FAIL_FAST][mla_partial_current_sync]", source, ) if __name__ == "__main__": unittest.main()