/test/registered/kernels/test_dsa_indexer.py

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