import unittest from typing import List, Optional, Tuple from unittest.mock import MagicMock, patch import torch from sglang.srt.environ import envs from sglang.srt.runtime_context import get_parallel from sglang.test.ci.ci_register import register_cuda_ci _parallel_override = get_parallel().override(attn_tp_size=1) _parallel_override.__enter__() from sglang.srt.configs.model_config import AttentionArch from sglang.srt.layers.attention.dsa.dsa_indexer import ( BaseIndexerMetadata, Indexer, rotate_activation, ) from sglang.srt.layers.attention.dsa.dsa_topk_backend import ( DSATopKBackend, TopkTransformMethod, ) from sglang.srt.layers.attention.dsa_backend import ( DeepseekSparseAttnBackend, DSAIndexerMetadata, DSAMetadata, ) from sglang.srt.layers.layernorm import LayerNorm from sglang.srt.layers.linear import LinearBase from sglang.srt.mem_cache.memory_pool import DSATokenToKVPool from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode from sglang.srt.server_args import ServerArgs, set_global_server_args_for_scheduler from sglang.test.test_utils import CustomTestCase register_cuda_ci(est_time=18, stage="base-b", runner_config="1-gpu-large") # Global configuration for all indexer tests DEFAULT_CONFIG = { "device": "cuda", "dtype": torch.bfloat16, "kv_cache_dtype": torch.float8_e4m3fn, "context_len": 2048, "max_bs": 64, "hidden_size": 5120, "index_n_heads": 32, "index_head_dim": 128, "rope_head_dim": 64, "index_topk": 64, "q_lora_rank": 1536, "kv_lora_rank": 512, "qk_rope_head_dim": 64, "qk_nope_head_dim": 128, "max_position_embeddings": 163840, "rope_theta": 10000.0, "layer_id": 0, "page_size": 64, } class MockIndexerMetadata(BaseIndexerMetadata): """Mock implementation of BaseIndexerMetadata for testing.""" def __init__(self, batch_size, seq_lens, page_table=None): self.batch_size = batch_size self.seq_lens = seq_lens self.page_table = page_table self.device = "cuda" def get_seqlens_int32(self) -> torch.Tensor: """Return: (batch_size,) int32 tensor""" return torch.tensor(self.seq_lens, dtype=torch.int32, device=self.device) def get_page_table_64(self) -> torch.Tensor: """Return: (batch_size, num_blocks) int32, page table with page size 64.""" if self.page_table is not None: return self.page_table # Create a simple page table for testing max_seq_len = max(self.seq_lens) num_blocks = (max_seq_len + 63) // 64 # Round up to page size 64 page_table = torch.zeros( (self.batch_size, num_blocks), dtype=torch.int32, device=self.device ) for i in range(self.batch_size): # Simple linear mapping: block i maps to page i num_blocks_needed = (self.seq_lens[i] + 63) // 64 page_table[i, :num_blocks_needed] = torch.arange( num_blocks_needed, device=self.device ) return page_table def get_page_table_1(self) -> torch.Tensor: """Return: (batch_size, num_blocks) int32, page table with page size 1.""" # Create a simple page table for testing with page size 1 max_seq_len = max(self.seq_lens) num_blocks = max_seq_len # Page size 1 means num_blocks == max_seq_len page_table = torch.zeros( (self.batch_size, num_blocks), dtype=torch.int32, device=self.device ) for i in range(self.batch_size): # Simple linear mapping: block i maps to page i num_blocks_needed = self.seq_lens[i] page_table[i, :num_blocks_needed] = torch.arange( num_blocks_needed, device=self.device ) return page_table def get_seqlens_expanded(self) -> torch.Tensor: """Return: (sum_extend_seq_len,) int32 tensor""" # For extend mode, each new token attends to progressively more tokens # For a sequence being extended from position 0 to seq_len, token i attends to i+1 tokens result = [] for seq_len in self.seq_lens: result.extend(range(1, seq_len + 1)) return torch.tensor(result, dtype=torch.int32, device=self.device) def get_indexer_kvcache_range(self) -> Tuple[torch.Tensor, torch.Tensor]: """ Return: (tokens, ), (tokens, ) int32, k_start and k_end in kv cache for each token. For extend mode, token i attends to tokens [0, i]. """ ks_list = [] ke_list = [] k_offset = 0 for seq_len in self.seq_lens: # For a sequence being extended from position 0 to seq_len # Token i attends to [k_offset, k_offset + i + 1) ks = torch.full((seq_len,), k_offset, dtype=torch.int32, device=self.device) ke = torch.arange( k_offset + 1, k_offset + seq_len + 1, dtype=torch.int32, device=self.device, ) ks_list.append(ks) ke_list.append(ke) k_offset += seq_len return torch.cat(ks_list, dim=0), torch.cat(ke_list, dim=0) def get_indexer_seq_len_cpu(self) -> torch.Tensor: """Return: seq lens for each batch.""" return torch.tensor(self.seq_lens, dtype=torch.int32, device="cpu") def get_indexer_seq_len(self) -> torch.Tensor: """Return: seq lens for each batch.""" return torch.tensor(self.seq_lens, dtype=torch.int32, device=self.device) def get_dsa_extend_len_cpu(self) -> List[int]: """ Return: extend seq lens for each batch. """ return list(self.seq_lens) def get_token_to_batch_idx(self) -> torch.Tensor: """Return: batch idx for each token.""" result = [] for batch_idx, seq_len in enumerate(self.seq_lens): result.extend([batch_idx] * seq_len) return torch.tensor(result, dtype=torch.int32, device=self.device) def topk_transform( self, logits: torch.Tensor, topk: int, ks: Optional[torch.Tensor] = None, cu_seqlens_q: Optional[torch.Tensor] = None, ke_offset: Optional[torch.Tensor] = None, batch_idx_list: Optional[torch.Tensor] = None, topk_indices_offset_override: Optional[torch.Tensor] = None, ) -> torch.Tensor: """ Perform topk selection on the logits. For testing, just return the topk indices. """ return torch.topk(logits, k=topk, dim=-1).indices class MockModelRunner: def __init__(self, config=None): self.device = "cuda" self.config = {**DEFAULT_CONFIG, **(config or {})} self.dtype = self.config["dtype"] self.kv_cache_dtype = self.config["kv_cache_dtype"] self.is_hybrid_swa = False # Model configuration attention_arch = AttentionArch.MLA max_context_len = self.config["context_len"] max_batch_size = self.config["max_bs"] # Create mock hf_config for DSA - instantiate it as an object, not a type hf_config = type( "HfConfig", (), { "architectures": ["DeepseekV3ForCausalLM"], "index_topk": self.config["index_topk"], "index_head_dim": self.config["index_head_dim"], "index_n_heads": self.config["index_n_heads"], }, )() self.model_config = type( "ModelConfig", (), { "context_len": max_context_len, "is_multimodal": False, "attention_arch": attention_arch, "num_attention_heads": 128, "kv_lora_rank": self.config["kv_lora_rank"], "qk_rope_head_dim": self.config["qk_rope_head_dim"], "qk_nope_head_dim": self.config["qk_nope_head_dim"], "hf_config": hf_config, }, )() self.sliding_window_size = None self.page_size = self.config["page_size"] # Create req_to_token_pool self.req_to_token_pool = type( "TokenPool", (), { "size": max_batch_size, "req_to_token": torch.zeros( max_batch_size, max_context_len, dtype=torch.int32, device=self.device, ), }, )() # Create DSATokenToKVPool max_total_num_tokens = max_batch_size * max_context_len self.token_to_kv_pool = DSATokenToKVPool( size=max_total_num_tokens, page_size=self.config["page_size"], dtype=self.config["kv_cache_dtype"], kv_lora_rank=self.config["kv_lora_rank"], qk_rope_head_dim=self.config["qk_rope_head_dim"], layer_num=1, device=self.device, index_head_dim=self.config["index_head_dim"], enable_memory_saver=False, kv_cache_dim=self.config["kv_lora_rank"] + self.config["qk_rope_head_dim"], ) # Required by backend with DSA-specific attributes self.server_args = type( "ServerArgs", (), { "kv_cache_dtype": "auto", "speculative_eagle_topk": None, "speculative_num_draft_tokens": 0, "enable_deterministic_inference": False, "dsa_prefill_backend": "flashmla_sparse", "dsa_decode_backend": "fa3", "dsa_topk_backend": "sgl-kernel", "dsa_paged_mqa_logits_backend": "auto", }, )() self.hisparse_coordinator = None @unittest.skipIf(not torch.cuda.is_available(), "Test requires CUDA") class TestDSAIndexer(CustomTestCase): @classmethod def setUpClass(cls): """Set up global server args for testing.""" server_args = ServerArgs(model_path="dummy") server_args.enable_dp_attention = False server_args.dsa_prefill_backend = "flashmla_sparse" server_args.dsa_decode_backend = "flashmla_sparse" set_global_server_args_for_scheduler(server_args) # Check GPU capability for FP8 if torch.cuda.is_available(): compute_capability = torch.cuda.get_device_capability() cls.supports_fp8 = compute_capability[0] >= 9 # Hopper or newer @classmethod def tearDownClass(cls): """Clean up after all tests.""" pass def setUp(self): # Test parameters self.batch_size = 2 self.seq_len = 128 self.config = DEFAULT_CONFIG.copy() self.device = "cuda" self.dtype = torch.bfloat16 def _init_model_runner(self, config_override=None): """Initialize model runner with optional config override.""" config = self.config.copy() if config_override: config.update(config_override) self.model_runner = MockModelRunner(config) self.backend = DeepseekSparseAttnBackend(self.model_runner) def _create_indexer(self, **kwargs): """Create an Indexer instance with default parameters.""" params = { "hidden_size": self.config["hidden_size"], "index_n_heads": self.config["index_n_heads"], "index_head_dim": self.config["index_head_dim"], "rope_head_dim": self.config["rope_head_dim"], "index_topk": self.config["index_topk"], "q_lora_rank": self.config["q_lora_rank"], "max_position_embeddings": self.config["max_position_embeddings"], "rope_theta": self.config["rope_theta"], "layer_id": self.config["layer_id"], "scale_fmt": "ue8m0", "block_size": 128, "quant_config": None, # No quantization for testing } params.update(kwargs) torch.set_default_dtype(self.dtype) with torch.device(self.device): indexer = Indexer(**params) indexer = indexer.to(device=self.device) # Convert linear layer weights to bfloat16 (but preserve LayerNorm's float32 # and weights_proj's float32 - it uses params_dtype=torch.float32 in production) # Need to recursively convert LinearBase submodules (like ReplicatedLinear) for name, module in indexer.named_modules(): # Check for LinearBase (parent of ReplicatedLinear) but exclude LayerNorm # Also exclude weights_proj which uses float32 params in production if isinstance(module, LinearBase) and not isinstance(module, LayerNorm): if "weights_proj" not in name: module.to(dtype=self.dtype) return indexer def _create_forward_batch( self, mode, batch_size=None, seq_len=None, extend_len=None ): """Create a forward batch for testing.""" batch_size = batch_size or self.batch_size seq_len = seq_len or self.seq_len if mode == ForwardMode.EXTEND: q_len = extend_len or seq_len total_len = seq_len forward_batch = ForwardBatch( batch_size=batch_size, input_ids=torch.randint( 0, 100, (batch_size, q_len), device=self.device ), out_cache_loc=torch.arange( batch_size * (total_len - q_len), batch_size * total_len, device=self.device, ), seq_lens_sum=batch_size * total_len, forward_mode=mode, req_pool_indices=torch.arange(batch_size, device=self.device), seq_lens=torch.tensor([total_len] * batch_size, device=self.device), seq_lens_cpu=torch.tensor([total_len] * batch_size, device="cpu"), extend_prefix_lens=torch.tensor( [total_len - q_len] * batch_size, device=self.device ), extend_prefix_lens_cpu=[total_len - q_len] * batch_size, extend_seq_lens=torch.tensor([q_len] * batch_size, device=self.device), extend_seq_lens_cpu=[q_len] * batch_size, ) else: # ForwardMode.DECODE decode_len = 1 total_len = seq_len + decode_len forward_batch = ForwardBatch( batch_size=batch_size, input_ids=torch.randint( 0, 100, (batch_size, decode_len), device=self.device ), out_cache_loc=torch.arange( batch_size * seq_len, batch_size * total_len, device=self.device ), seq_lens_sum=batch_size * total_len, forward_mode=mode, req_pool_indices=torch.arange(batch_size, device=self.device), seq_lens=torch.tensor([total_len] * batch_size, device=self.device), seq_lens_cpu=torch.tensor([total_len] * batch_size, device="cpu"), ) # Pool refs + attn_backend are now resolved via the ForwardContext; # publish ``self.backend`` for the duration of this fixture call so # ``get_attn_backend()`` / ``get_token_to_kv_pool()`` / # ``get_req_to_token_pool()`` resolve correctly. from sglang.srt.model_executor.forward_context import ( ForwardContext, set_forward_context, ) set_forward_context(ForwardContext(attn_backend=self.backend)) # Mock write to req_to_token_pool page_size = self.model_runner.page_size for i in range(batch_size): seq_length = total_len for j in range(seq_length): self.model_runner.req_to_token_pool.req_to_token[i, j] = ( i * seq_length + j + page_size ) return forward_batch def _verify_topk_output(self, topk_indices, batch_size, q_len, topk): """Verify the topk indices output shape and basic properties.""" self.assertIsNotNone(topk_indices) self.assertEqual(topk_indices.device.type, "cuda") # Check shape - should be (total_q_len, topk_padded) # where topk_padded is aligned to 2048 self.assertEqual(len(topk_indices.shape), 2) self.assertEqual(topk_indices.shape[0], batch_size * q_len) # Check that topk is padded to at least topk self.assertGreaterEqual(topk_indices.shape[1], topk) # Check for padding values (-1) has_padding = (topk_indices == -1).any() self.assertTrue( has_padding or topk_indices.shape[1] == topk, "Output should have padding or exact topk size", ) def _make_tie_free_logits( self, batch_size: int, max_score_len: int ) -> torch.Tensor: perm = torch.argsort( torch.randn( batch_size, max_score_len, dtype=torch.float32, device=self.device ), dim=-1, ) return torch.gather( torch.arange(max_score_len, device=self.device, dtype=torch.float32) .unsqueeze(0) .expand(batch_size, -1), dim=1, index=perm, ) def _run_unfused_topk_backend_validity_test( self, batch_size: int, max_score_len: int, topk: int, topk_backend: DSATopKBackend, with_row_starts: bool, ): logits = self._make_tie_free_logits(batch_size, max_score_len) if with_row_starts: row_starts = torch.randint( 0, max_score_len - 1, (batch_size,), dtype=torch.int32, device=self.device, ) max_lengths = max_score_len - row_starts random_lengths = torch.randint( 0, max_score_len - 1, (batch_size,), dtype=torch.int32, device=self.device, ) seq_lens_expanded = torch.minimum(max_lengths, random_lengths) else: row_starts = None seq_lens_expanded = torch.randint( 0, max_score_len - 1, (batch_size,), dtype=torch.int32, device=self.device, ) seq_lens_expanded = seq_lens_expanded.to(dtype=torch.int32, device=self.device) max_seq_len_k = int(seq_lens_expanded.max().item()) cu_seqlens_q = torch.arange( batch_size + 1, dtype=torch.int32, device=self.device ) dsa_cu_seqlens_k = torch.zeros( batch_size + 1, dtype=torch.int32, device=self.device ) dsa_cu_seqlens_k[1:] = torch.cumsum(seq_lens_expanded, dim=0) page_table_1 = ( torch.arange(max_seq_len_k, dtype=torch.int32, device=self.device) .unsqueeze(0) .expand(batch_size, -1) .contiguous() ) metadata = DSAIndexerMetadata( attn_metadata=DSAMetadata( page_size=1, cache_seqlens_int32=seq_lens_expanded.clone(), max_seq_len_q=1, max_seq_len_k=max_seq_len_k, cu_seqlens_q=cu_seqlens_q, cu_seqlens_k=cu_seqlens_q.clone(), page_table_1=page_table_1, real_page_table=page_table_1, dsa_cache_seqlens_int32=seq_lens_expanded.clone(), dsa_cu_seqlens_q=cu_seqlens_q.clone(), dsa_cu_seqlens_k=dsa_cu_seqlens_k, dsa_extend_seq_lens_list=seq_lens_expanded.cpu().tolist(), dsa_seqlens_expanded=seq_lens_expanded, ), topk_transform_method=TopkTransformMethod.PAGED, topk_backend=topk_backend, ) with envs.SGLANG_DSA_FUSE_TOPK.override(False): topk_test = metadata.topk_transform(logits, topk, ks=row_starts) self.assertEqual(topk_test.shape, (batch_size, topk)) self.assertEqual(topk_test.dtype, torch.int32) expected_valid = torch.minimum( seq_lens_expanded, torch.full_like(seq_lens_expanded, topk), ) actual_valid = (topk_test >= 0).sum(dim=-1).to(torch.int32) self.assertTrue(torch.equal(actual_valid, expected_valid)) starts = ( row_starts.to(torch.int32) if row_starts is not None else torch.zeros( (topk_test.shape[0],), dtype=torch.int32, device=topk_test.device ) ) for row in range(topk_test.shape[0]): test_row = topk_test[row] valid_test = test_row[test_row >= 0] expected_k = int(expected_valid[row].item()) self.assertEqual(valid_test.numel(), expected_k) if expected_k == 0: continue start = int(starts[row].item()) row_len = int(seq_lens_expanded[row].item()) self.assertTrue(torch.all((valid_test >= 0) & (valid_test < row_len))) self.assertEqual(torch.unique(valid_test).numel(), valid_test.numel()) row_scores = logits[row, start : start + row_len] ref_topk = torch.topk(row_scores, expected_k, dim=-1, sorted=False).indices self.assertTrue( torch.equal( torch.sort(valid_test.to(torch.int32)).values, torch.sort(ref_topk.to(torch.int32)).values, ) ) def _run_fused_topk_backend_equivalence_test( self, batch_size: int, max_score_len: int, topk: int, topk_transform_method: TopkTransformMethod, with_row_starts: bool, query_lens: Optional[List[int]] = None, ): num_rows = sum(query_lens) if query_lens is not None else batch_size logits = self._make_tie_free_logits(num_rows, max_score_len) if with_row_starts: row_starts = torch.randint( 0, max_score_len - 1, (num_rows,), dtype=torch.int32, device=self.device, ) max_lengths = max_score_len - row_starts random_lengths = torch.randint( 1, max_score_len, (num_rows,), dtype=torch.int32, device=self.device, ) seq_lens_expanded = torch.minimum(max_lengths, random_lengths) else: row_starts = None seq_lens_expanded = torch.randint( 1, max_score_len, (num_rows,), dtype=torch.int32, device=self.device, ) topk_indices_offset = ( torch.arange(num_rows, dtype=torch.int32, device=self.device) * max_score_len ) if query_lens is None: cu_seqlens_q = torch.arange( batch_size + 1, dtype=torch.int32, device=self.device ) q_lens = None batch_idx_list = None else: q_lens = torch.tensor(query_lens, dtype=torch.int32, device=self.device) cu_seqlens_q = torch.zeros( batch_size + 1, dtype=torch.int32, device=self.device ) cu_seqlens_q[1:] = torch.cumsum(q_lens, dim=0) batch_idx_list = list(range(batch_size)) cu_seqlens_k = torch.zeros( batch_size + 1, dtype=torch.int32, device=self.device ) dsa_cu_seqlens_k = torch.zeros( num_rows + 1, dtype=torch.int32, device=self.device ) dsa_cu_seqlens_k[1:] = torch.cumsum(seq_lens_expanded, dim=0) page_table_1 = ( ( torch.arange(max_score_len, dtype=torch.int32, device=self.device) .unsqueeze(0) .expand(batch_size, -1) ) + ( torch.arange( batch_size, dtype=torch.int32, device=self.device ).unsqueeze(1) * max_score_len ) ).contiguous() # The fused v2 PAGED dispatch requires the per-forward plan to be # preprocessed alongside the metadata (it asserts rather than silently # recomputing it) -- mirror what init_forward_metadata / # _build_forward_metadata_cuda_graph do. from sglang.jit_kernel.dsv4.topk import plan_topk_v2 attn_metadata = DSAMetadata( page_size=1, cache_seqlens_int32=seq_lens_expanded.clone(), max_seq_len_q=1, max_seq_len_k=max_score_len, cu_seqlens_q=cu_seqlens_q, cu_seqlens_k=cu_seqlens_k, page_table_1=page_table_1, real_page_table=page_table_1, dsa_cache_seqlens_int32=seq_lens_expanded.clone(), dsa_cu_seqlens_q=cu_seqlens_q.clone(), dsa_cu_seqlens_k=dsa_cu_seqlens_k, dsa_extend_seq_lens_list=seq_lens_expanded.cpu().tolist(), dsa_seqlens_expanded=seq_lens_expanded, topk_v2_plan=plan_topk_v2(seq_lens_expanded), topk_indices_offset=( topk_indices_offset if topk_transform_method == TopkTransformMethod.RAGGED else None ), ) metadata_sgl = DSAIndexerMetadata( attn_metadata=attn_metadata, topk_transform_method=topk_transform_method, topk_backend=DSATopKBackend.SGL_KERNEL, ) metadata_flashinfer = DSAIndexerMetadata( attn_metadata=attn_metadata, topk_transform_method=topk_transform_method, topk_backend=DSATopKBackend.FLASHINFER, ) with envs.SGLANG_DSA_FUSE_TOPK.override(True): out_sgl = metadata_sgl.topk_transform( logits, topk, ks=row_starts, cu_seqlens_q=q_lens, batch_idx_list=batch_idx_list, ) out_flashinfer = metadata_flashinfer.topk_transform( logits, topk, ks=row_starts, cu_seqlens_q=q_lens, batch_idx_list=batch_idx_list, ) self.assertEqual(out_sgl.shape, out_flashinfer.shape) self.assertEqual(out_sgl.dtype, out_flashinfer.dtype) self.assertEqual(out_sgl.dtype, torch.int32) self.assertTrue( torch.equal( torch.sort(out_sgl, dim=-1).values, torch.sort(out_flashinfer, dim=-1).values, ) ) @patch("sglang.srt.layers.attention.dsa.dsa_indexer.deep_gemm") def test_indexer_basic_creation(self, mock_deep_gemm): """Test basic indexer creation and initialization.""" mock_deep_gemm.get_num_sms.return_value = 132 indexer = self._create_indexer() self.assertEqual(indexer.hidden_size, self.config["hidden_size"]) self.assertEqual(indexer.n_heads, self.config["index_n_heads"]) self.assertEqual(indexer.head_dim, self.config["index_head_dim"]) self.assertEqual(indexer.rope_head_dim, self.config["rope_head_dim"]) self.assertEqual(indexer.index_topk, self.config["index_topk"]) self.assertEqual(indexer.layer_id, self.config["layer_id"]) @patch("sglang.srt.layers.attention.dsa.dsa_indexer.deep_gemm") @patch("sglang.srt.layers.attention.dsa.triton_kernel.act_quant") def test_forward_extend_mode(self, mock_act_quant, mock_deep_gemm): """Test indexer forward pass in extend mode.""" if not self.supports_fp8: self.skipTest("FP8 requires Hopper GPU or newer") # Setup mocks mock_deep_gemm.get_num_sms.return_value = 132 mock_deep_gemm.get_paged_mqa_logits_metadata.return_value = MagicMock() def mock_quant(x, *args, **kwargs): # Return FP8 tensor and scale return x.to(torch.float8_e4m3fn), torch.ones( x.shape[0], dtype=torch.float32, device=x.device ) mock_act_quant.side_effect = mock_quant # Mock deep_gemm.fp8_mqa_logits to return logits (ragged path) def mock_mqa_logits(q, kv, weights, ks, ke, *args, **kwargs): # q shape: (sum_extend_seq_len, ...), return logits for each query token num_queries = q.shape[0] # kv is a tuple (k_fp8, k_scale), get total number of keys from k_fp8 k_fp8, k_scale = kv max_kv_len = k_fp8.shape[0] # Total keys across all batches (k_offset) return torch.randn( num_queries, max_kv_len, dtype=torch.float32, device="cuda" ) mock_deep_gemm.fp8_mqa_logits.side_effect = mock_mqa_logits # Also mock the paged version for completeness def mock_paged_mqa_logits(q, kv, weights, *args, **kwargs): batch_size = q.shape[0] seq_len = 128 return torch.randn(batch_size, seq_len, dtype=torch.float32, device="cuda") mock_deep_gemm.fp8_paged_mqa_logits.side_effect = mock_paged_mqa_logits self._init_model_runner() indexer = self._create_indexer() forward_batch = self._create_forward_batch(ForwardMode.EXTEND) # Create input tensors total_tokens = self.batch_size * self.seq_len hidden_states = torch.randn( total_tokens, self.config["hidden_size"], dtype=self.dtype, device=self.device, ) q_lora = torch.randn( total_tokens, self.config["q_lora_rank"], dtype=self.dtype, device=self.device, ) positions = torch.arange(total_tokens, device=self.device) # Run forward pass with patch.object( self.backend, "get_indexer_metadata", return_value=MockIndexerMetadata( self.batch_size, [self.seq_len] * self.batch_size ), ): topk_indices = indexer( x=hidden_states, q_lora=q_lora, positions=positions, forward_batch=forward_batch, layer_id=self.config["layer_id"], ) # Verify output self._verify_topk_output( topk_indices, self.batch_size, self.seq_len, self.config["index_topk"] ) @patch("sglang.srt.layers.attention.dsa.dsa_indexer.deep_gemm") @patch("sglang.srt.layers.attention.dsa.triton_kernel.act_quant") def test_forward_decode_mode(self, mock_act_quant, mock_deep_gemm): """Test indexer forward pass in decode mode.""" if not self.supports_fp8: self.skipTest("FP8 requires Hopper GPU or newer") # Setup mocks mock_deep_gemm.get_num_sms.return_value = 132 mock_deep_gemm.get_paged_mqa_logits_metadata.return_value = MagicMock() def mock_quant(x, *args, **kwargs): return x.to(torch.float8_e4m3fn), torch.ones( x.shape[0], dtype=torch.float32, device=x.device ) mock_act_quant.side_effect = mock_quant def mock_paged_mqa_logits(q, kv, weights, *args, **kwargs): batch_size = q.shape[0] seq_len = 128 return torch.randn(batch_size, seq_len, dtype=torch.float32, device="cuda") mock_deep_gemm.fp8_paged_mqa_logits.side_effect = mock_paged_mqa_logits self._init_model_runner() indexer = self._create_indexer() forward_batch = self._create_forward_batch(ForwardMode.DECODE) # Create input tensors for decode (batch_size tokens only) hidden_states = torch.randn( self.batch_size, self.config["hidden_size"], dtype=self.dtype, device=self.device, ) q_lora = torch.randn( self.batch_size, self.config["q_lora_rank"], dtype=self.dtype, device=self.device, ) positions = torch.arange(self.batch_size, device=self.device) # Run forward pass with patch.object( self.backend, "get_indexer_metadata", return_value=MockIndexerMetadata( self.batch_size, [self.seq_len + 1] * self.batch_size ), ): topk_indices = indexer( x=hidden_states, q_lora=q_lora, positions=positions, forward_batch=forward_batch, layer_id=self.config["layer_id"], ) # Verify output - decode mode has q_len=1 self._verify_topk_output( topk_indices, self.batch_size, 1, self.config["index_topk"] ) def test_rotate_activation(self): """Test the Hadamard transform (rotate_activation) function.""" # Test with power-of-2 hidden size hidden_size = 128 x = torch.randn(16, hidden_size, dtype=torch.bfloat16, device=self.device) try: output = rotate_activation(x) self.assertEqual(output.shape, x.shape) self.assertEqual(output.dtype, torch.bfloat16) except Exception: self.skipTest("hadamard JIT kernel not available") def test_rotate_activation_invalid_size(self): """Test that rotate_activation fails with non-power-of-2 size.""" # Test with non-power-of-2 hidden size hidden_size = 129 # Not a power of 2 x = torch.randn(16, hidden_size, dtype=torch.bfloat16, device=self.device) with self.assertRaises(AssertionError): rotate_activation(x) def test_indexer_metadata_interface(self): """Test the BaseIndexerMetadata interface implementation.""" batch_size = 4 seq_lens = [64, 128, 96, 112] metadata = MockIndexerMetadata(batch_size, seq_lens) # Test get_seqlens_int32 seqlens = metadata.get_seqlens_int32() self.assertEqual(seqlens.shape, (batch_size,)) self.assertEqual(seqlens.dtype, torch.int32) self.assertTrue(torch.all(seqlens == torch.tensor(seq_lens, device="cuda"))) # Test get_page_table_64 page_table = metadata.get_page_table_64() self.assertEqual(len(page_table.shape), 2) self.assertEqual(page_table.shape[0], batch_size) self.assertEqual(page_table.dtype, torch.int32) # Test topk_transform logits = torch.randn(batch_size, 128, device="cuda") topk = 64 topk_indices = metadata.topk_transform(logits, topk) self.assertEqual(topk_indices.shape, (batch_size, topk)) def test_topk_unfused_backends_valid_selection(self): batch_size = 8 max_score_len = 16 * 1024 topk = 2048 for topk_backend in [ DSATopKBackend.SGL_KERNEL, DSATopKBackend.TORCH, DSATopKBackend.FLASHINFER, ]: tie_break_values = ( [None, "small", "large"] if topk_backend == DSATopKBackend.FLASHINFER else [None] ) for tie_break in tie_break_values: for with_row_starts in [False, True]: with self.subTest( topk_backend=topk_backend.value, tie_break=tie_break, with_row_starts=with_row_starts, ): with envs.SGLANG_DSA_TOPK_FLASHINFER_TIE_BREAK.override( tie_break ): self._run_unfused_topk_backend_validity_test( batch_size, max_score_len, topk, topk_backend=topk_backend, with_row_starts=with_row_starts, ) def test_topk_fused_backends_equivalence(self): batch_size = 8 max_score_len = 16 * 1024 topk = 2048 for tie_break in [None, "small", "large"]: for topk_transform_method in [ TopkTransformMethod.PAGED, TopkTransformMethod.RAGGED, ]: for with_row_starts in [False, True]: if ( topk_transform_method == TopkTransformMethod.PAGED and with_row_starts ): # The synthetic paged fixture uses the decode-like row mapping. # Ragged fused and unfused cases cover shifted row windows. continue with self.subTest( tie_break=tie_break, topk_transform_method=topk_transform_method.name, with_row_starts=with_row_starts, ): with envs.SGLANG_DSA_TOPK_FLASHINFER_TIE_BREAK.override( tie_break ): self._run_fused_topk_backend_equivalence_test( batch_size=batch_size, max_score_len=max_score_len, topk=topk, topk_transform_method=topk_transform_method, with_row_starts=with_row_starts, ) with self.subTest( tie_break=tie_break, topk_transform_method=TopkTransformMethod.PAGED.name, with_row_starts=False, query_lens="multi", ): with envs.SGLANG_DSA_TOPK_FLASHINFER_TIE_BREAK.override(tie_break): self._run_fused_topk_backend_equivalence_test( batch_size=batch_size, max_score_len=max_score_len, topk=topk, topk_transform_method=TopkTransformMethod.PAGED, with_row_starts=False, query_lens=[1, 2, 3, 1, 2, 1, 3, 2], ) # TODO: enable this test after indexer accuracy aligned # @patch("sglang.srt.layers.attention.dsa.dsa_indexer.deep_gemm") # def test_indexer_with_different_topk(self, mock_deep_gemm): # """Test indexer with different topk values.""" # mock_deep_gemm.get_num_sms.return_value = 132 # for topk in [32, 64, 128]: # with self.subTest(topk=topk): # indexer = self._create_indexer(index_topk=topk) # self.assertEqual(indexer.index_topk, topk) @patch("sglang.srt.layers.attention.dsa.dsa_indexer.deep_gemm") def test_indexer_with_fused_wk(self, mock_deep_gemm): """Test indexer creation with fused wk and weights projection.""" mock_deep_gemm.get_num_sms.return_value = 132 # Note: fuse_wk_and_weights_proj feature is not currently implemented # This test verifies basic indexer creation still works indexer = self._create_indexer() self.assertIsNotNone(indexer) @patch("sglang.srt.layers.attention.dsa.dsa_indexer.deep_gemm") def test_indexer_with_alt_stream(self, mock_deep_gemm): """Test indexer creation with alternative CUDA stream.""" mock_deep_gemm.get_num_sms.return_value = 132 alt_stream = torch.cuda.Stream() indexer = self._create_indexer(alt_stream=alt_stream) self.assertEqual(indexer.alt_stream, alt_stream) def test_shape_sanity_checks(self): """Test various shape combinations for consistency.""" test_configs = [ {"batch_size": 1, "seq_len": 64}, {"batch_size": 4, "seq_len": 128}, {"batch_size": 8, "seq_len": 256}, ] for config in test_configs: with self.subTest(**config): batch_size = config["batch_size"] seq_len = config["seq_len"] # Test metadata shapes metadata = MockIndexerMetadata(batch_size, [seq_len] * batch_size) seqlens = metadata.get_seqlens_int32() self.assertEqual(seqlens.shape, (batch_size,)) page_table = metadata.get_page_table_64() expected_blocks = (seq_len + 63) // 64 self.assertEqual(page_table.shape[0], batch_size) self.assertGreaterEqual(page_table.shape[1], expected_blocks) if __name__ == "__main__": unittest.main()