[Perf] [NVIDIA] minimax-sparse: Run the decode indexer CP on Hopper at TP4 and TP8

python/sglang/srt/environ.py

@@ -2263,8 +2263,8 @@ class Envs:
     # Only honored on ROCm gfx95 without two-batch overlap; elsewhere the
     # backend pins 1.
     SGLANG_MINIMAX_M3_INDEX_TOPK_FREQ = EnvInt(1)
-    # Opt-in gfx950 TP4 decode indexer context partitioning. Keeps the index
-    # cache replicated; gathers Q and exchanges local top-k candidates.
+    # Opt-in decode indexer context partitioning (gfx950 TP4, SM90 TP4/TP8). Keeps
+    # the index cache replicated; gathers Q and exchanges local top-k candidates.
     SGLANG_MINIMAX_M3_INDEXER_CP = EnvBool(False)
     # gfx95: lightning-indexer K cache in fp8_e4m3fn (bf16 q x fp8 k in the scorers);
     # main attention K/V keep kv_cache_dtype.

python/sglang/srt/layers/attention/minimax_sparse_ops/indexer_cp.py

@@ -16,6 +16,7 @@
 def unsupported_reason(
     *,
     gfx950,
+    sm90,
     tp_size,
     attn_tp_size,
     attn_cp_size,
@@ -35,16 +36,25 @@ def unsupported_reason(
     dense_sparse_decode,
 ):
     """Run-level gates: identical on every rank, with no device data reads."""
-    if not gfx950:
-        return "requires AMD gfx950"
-    if (tp_size, attn_tp_size, attn_cp_size, attn_dp_size) != (4, 4, 1, 1):
-        return "requires TP4 with attention CP/DP size 1"
+    if gfx950:
+        if (tp_size, attn_tp_size, attn_cp_size, attn_dp_size) != (4, 4, 1, 1):
+            return "requires TP4 with attention CP/DP size 1"
+    elif sm90:
+        # At TP8 each rank still serves one index head; head pairs are replicated.
+        if attn_tp_size not in (4, 8) or (tp_size, attn_cp_size, attn_dp_size) != (
+            attn_tp_size,
+            1,
+            1,
+        ):
+            return "requires TP4 or TP8 with attention CP/DP size 1"
+    else:
+        return "requires AMD gfx950 or NVIDIA SM90"
     if (index_heads, kv_heads, head_dim, block_size, topk) != (4, 4, 128, 128, 16):
         return "requires four index/KV heads, dimension 128, block 128, top-k 16"
     if score_type != "max" or not radix_topk:
-        return "requires max scores and the ROCm radix top-k tie ordering"
+        return "requires max scores and the radix decode top-k"
     if not 0 < max_context_len <= 16384 * 128:
-        return "context length exceeds the ROCm radix selector contract"
+        return "context length exceeds the radix selector contract"
     if not draft_is_chain:
         return "verify rows must be chain drafts: no speculation, or EAGLE with top-k 1"
     if tbo or hisparse or fp8_query or dense_sparse_decode:
@@ -71,7 +81,7 @@ def make_indexer_cp(backend, runner, sparse_cfg):
         return None
     from sglang.srt.layers.moe.utils import is_tbo_enabled
     from sglang.srt.runtime_context import get_parallel, get_spec
-    from sglang.srt.utils import is_hip
+    from sglang.srt.utils import is_hip, is_sm90_supported
 
     parallel = get_parallel()
     spec = get_spec()
@@ -82,6 +92,7 @@ def make_indexer_cp(backend, runner, sparse_cfg):
     )
     reason = unsupported_reason(
         gfx950=arch.split(":")[0] == "gfx950",
+        sm90=not is_hip() and is_sm90_supported(),
         tp_size=parallel.tp_size,
         attn_tp_size=parallel.attn_tp_size,
         attn_cp_size=parallel.attn_cp_size,
@@ -108,23 +119,25 @@ def make_indexer_cp(backend, runner, sparse_cfg):
     cp = MiniMaxIndexerCP(parallel.attn_tp_group)
     cp.warmup()
     logger.info(
-        "MiniMax indexer CP enabled: TP4, decode-only, replicated index cache, "
-        "query gather + candidate gather, top-k 16; index-value layers use TP"
+        "MiniMax indexer CP enabled: TP%d, decode-only, replicated index cache, "
+        "query gather + candidate gather, top-k 16; index-value layers use TP",
+        cp.world,
     )
     return cp
 
 
 class MiniMaxIndexerCP:
     def __init__(self, group):
-        if group.world_size != 4:
-            raise ValueError("MiniMax indexer CP requires a four-rank group")
+        if group.world_size not in (4, 8):
+            raise ValueError("MiniMax indexer CP requires a four- or eight-rank group")
         self.group = group
+        self.world = group.world_size
         self.rank = group.rank_in_group
 
     def warmup(self):
         """Initialize both runtime and PyNCCL gather paths before graph capture."""
         probe = torch.zeros((1, 128), dtype=torch.bfloat16, device=self.group.device)
-        output = torch.empty((4, 128), dtype=probe.dtype, device=probe.device)
+        output = torch.empty((self.world, 128), dtype=probe.dtype, device=probe.device)
         self.group.all_gather_into_tensor(output, probe)
         comm = self.group.get_available_pynccl_comm()
         if comm is not None:
@@ -174,12 +187,13 @@ def __call__(
         if batch == 0:
             return torch.empty((1, 0, 16), dtype=torch.int32, device=q.device)
         local_q = q[:, 0].contiguous()
-        gathered_q = torch.empty((4 * batch, 128), dtype=q.dtype, device=q.device)
+        world = self.world
+        gathered_q = torch.empty((world * batch, 128), dtype=q.dtype, device=q.device)
         self.group.all_gather_into_tensor(gathered_q, local_q)
         sm_scale = 128**-0.5 if sm_scale is None else sm_scale
         sm_scale *= 1.0 if q_scale is None else q_scale
         scores = score_local_blocks(
-            gathered_q.view(4, batch, 128),
+            gathered_q.view(world, batch, 128),
             k_cache,
             req_to_token,
             slot_ids,
@@ -193,11 +207,13 @@ def __call__(
             packed_queries=packed_queries,
         )
         keys = select_local_candidates(scores, seq_lens, self.rank, max_seqlen)
-        received = torch.empty((4, 4, batch, 16), dtype=torch.int64, device=q.device)
+        received = torch.empty(
+            (world, world, batch, 16), dtype=torch.int64, device=q.device
+        )
         # Pure byte copies: numeric conversion would destroy the score/ID keys.
         # Float32 views also qualify for the runtime's AITER custom gather.
         self.group.all_gather_into_tensor(
-            received.view(torch.float32).view(16 * batch, 32),
-            keys.view(torch.float32).view(4 * batch, 32),
+            received.view(torch.float32).view(world * world * batch, 32),
+            keys.view(torch.float32).view(world * batch, 32),
         )
         return merge_candidates(received, self.rank)

test/registered/kernels/ops/attention/test_minimax_indexer_cp.py

@@ -0,0 +1,170 @@
+"""On SM90, indexer CP must select exactly the blocks the native page-128 selector does.
+
+Hopper serves the page-128 index cache through its own paged scorer, and at TP8 each
+index head is served by two ranks, so eight shards score four distinct queries.
+"""
+
+import pytest
+import torch
+
+from sglang.kernels.ops.attention.minimax_sparse.decode.flash_with_topk_idx import (
+    flash_decode_with_topk_idx,
+)
+from sglang.kernels.ops.attention.minimax_sparse.decode.indexer_cp import (
+    merge_candidates,
+    score_local_blocks,
+    select_local_candidates,
+)
+from sglang.test.ci.ci_register import register_cuda_ci
+
+register_cuda_ci(est_time=30, stage="base-b-kernel-unit", runner_config="1-gpu-large")
+
+BLOCK, TOPK, HEADS, DIM, PAGE = 128, 16, 4, 128, 128
+INIT_BLOCKS, LOCAL_BLOCKS = 0, 1
+
+pytestmark = pytest.mark.skipif(
+    not torch.cuda.is_available() or torch.cuda.get_device_capability() != (9, 0),
+    reason="the page-128 paged scorer is the SM90 decode path",
+)
+
+
+def _inputs(batch, max_len, dtype):
+    gen = torch.Generator(device="cuda").manual_seed(20261025)
+    # Spare rows so slots are never the batch index; pages shuffled so a rank's
+    # strided blocks never land on contiguous memory.
+    rows = 3 * batch
+    pages = torch.randperm(rows * max_len // PAGE, generator=gen, device="cuda")
+    table = (
+        (pages.to(torch.int32)[:, None] * PAGE + torch.arange(PAGE, device="cuda"))
+        .reshape(rows, max_len)
+        .to(torch.int32)
+    )
+    cache = torch.randn(
+        (rows * max_len, 1, DIM), generator=gen, device="cuda", dtype=torch.bfloat16
+    ).to(dtype)
+    q = torch.randn(
+        (batch, HEADS, DIM), generator=gen, device="cuda", dtype=torch.bfloat16
+    )
+    slots = (batch + 1 + 2 * torch.arange(batch, device="cuda")).flip(0)
+    # Off block boundaries, so the partial last block and the local block both count.
+    lengths = max_len - (37 * torch.arange(batch, device="cuda")) % 200
+    return q, cache, table, slots, lengths
+
+
+@pytest.mark.parametrize("world", [4, 8])
+@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float8_e4m3fn])
+def test_sharded_selection_matches_paged_native_selector(world, dtype):
+    batch, max_len = 4, 32768
+    q, cache, table, slots, lengths = _inputs(batch, max_len, dtype)
+    head_of_rank = [rank * HEADS // world for rank in range(world)]
+    gathered_q = q[:, head_of_rank].permute(1, 0, 2).contiguous()
+    keys = torch.stack(
+        [
+            select_local_candidates(
+                score_local_blocks(
+                    gathered_q,
+                    cache,
+                    table,
+                    slots,
+                    lengths,
+                    max_len,
+                    rank,
+                    INIT_BLOCKS,
+                    LOCAL_BLOCKS,
+                    DIM**-0.5,
+                    1.0,
+                ),
+                lengths,
+                rank,
+                max_len,
+                TOPK,
+            )
+            for rank in range(world)
+        ]
+    )
+    for rank, head in enumerate(head_of_rank):
+        native = flash_decode_with_topk_idx(
+            q=q[:, head : head + 1].contiguous(),
+            sink=None,
+            k_cache=cache,
+            v_cache=None,
+            req_to_token=table,
+            seq_lens=lengths,
+            max_seqlen=max_len,
+            slot_ids=slots,
+            block_size=BLOCK,
+            topk=TOPK,
+            init_blocks=INIT_BLOCKS,
+            local_blocks=LOCAL_BLOCKS,
+            disable_index_value=True,
+            page_size=PAGE,
+        )[1]
+        sharded = merge_candidates(keys, rank)
+        assert torch.equal(sharded, native), (
+            f"rank {rank}: {(sharded != native).any(-1).sum()} rows differ"
+        )
+
+
+@pytest.mark.parametrize("world", [4, 8])
+def test_chain_verify_rows_match_paged_native_selector(world):
+    """Verify rows straddle a block boundary; W=4 scores them packed, W=8 row by row."""
+    reqs, draft_tokens = 4, 4
+    prefix = 255 * BLOCK + 126  # rows end at 32767..32770, across block 256
+    max_len = (prefix + draft_tokens + PAGE - 1) // PAGE * PAGE
+    q, cache, table, slots, _ = _inputs(reqs * draft_tokens, max_len, torch.bfloat16)
+    slots = slots[:reqs].repeat_interleave(draft_tokens)
+    lengths = (prefix + torch.arange(1, draft_tokens + 1, device="cuda")).repeat(reqs)
+    head_of_rank = [rank * HEADS // world for rank in range(world)]
+    gathered_q = q[:, head_of_rank].permute(1, 0, 2).contiguous()
+    keys = torch.stack(
+        [
+            select_local_candidates(
+                score_local_blocks(
+                    gathered_q,
+                    cache,
+                    table,
+                    slots,
+                    lengths,
+                    max_len,
+                    rank,
+                    INIT_BLOCKS,
+                    LOCAL_BLOCKS,
+                    DIM**-0.5,
+                    1.0,
+                    packed_queries=draft_tokens,
+                ),
+                lengths,
+                rank,
+                max_len,
+                TOPK,
+            )
+            for rank in range(world)
+        ]
+    )
+    for rank, head in enumerate(head_of_rank):
+        native = flash_decode_with_topk_idx(
+            q=q[:, head : head + 1].contiguous(),
+            sink=None,
+            k_cache=cache,
+            v_cache=None,
+            req_to_token=table,
+            seq_lens=lengths,
+            max_seqlen=max_len,
+            slot_ids=slots,
+            block_size=BLOCK,
+            topk=TOPK,
+            init_blocks=INIT_BLOCKS,
+            local_blocks=LOCAL_BLOCKS,
+            disable_index_value=True,
+            page_size=PAGE,
+        )[1]
+        sharded = merge_candidates(keys, rank)
+        assert torch.equal(sharded, native), (
+            f"rank {rank}: {(sharded != native).any(-1).sum()} rows differ"
+        )
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main([__file__, "-v", "-s"]))

test/registered/unit/layers/attention/test_minimax_indexer_cp_gate.py

@@ -12,6 +12,7 @@
 # The run shape MiniMax-M3 CP was validated on; only the draft layout varies below.
 SUPPORTED = dict(
     gfx950=True,
+    sm90=False,
     tp_size=4,
     attn_tp_size=4,
     attn_cp_size=1,
@@ -56,6 +57,46 @@ def test_only_validated_draft_layouts_enable_cp(self):
                 )
                 self.assertEqual(reason is None, expected, reason)
 
+    def test_sm90_enables_only_validated_topologies(self):
+        """Each rank must serve exactly one index head of the four; any other split,
+        or another attention CP/DP layout, would merge candidates for the wrong query."""
+        enabled = {
+            (4, 4, 1, 1): True,
+            (8, 8, 1, 1): True,  # head pairs replicated across ranks
+            (2, 2, 1, 1): False,  # two index heads per rank
+            (8, 4, 1, 2): False,  # DP attention not validated
+            (8, 4, 2, 1): False,
+        }
+        for (tp, attn_tp, attn_cp, attn_dp), expected in enabled.items():
+            with self.subTest(tp=tp, attn_tp=attn_tp, attn_cp=attn_cp, attn_dp=attn_dp):
+                reason = unsupported_reason(
+                    **{
+                        **SUPPORTED,
+                        "gfx950": False,
+                        "sm90": True,
+                        "tp_size": tp,
+                        "attn_tp_size": attn_tp,
+                        "attn_cp_size": attn_cp,
+                        "attn_dp_size": attn_dp,
+                    }
+                )
+                self.assertEqual(reason is None, expected, reason)
+        for gfx950, sm90, attn_tp, expected in (
+            (True, False, 8, False),  # TP8 is validated on SM90 only
+            (False, False, 4, False),  # other CUDA architectures stay off
+        ):
+            with self.subTest(gfx950=gfx950, sm90=sm90, attn_tp=attn_tp):
+                reason = unsupported_reason(
+                    **{
+                        **SUPPORTED,
+                        "gfx950": gfx950,
+                        "sm90": sm90,
+                        "tp_size": attn_tp,
+                        "attn_tp_size": attn_tp,
+                    }
+                )
+                self.assertEqual(reason is None, expected, reason)
+
 
 if __name__ == "__main__":
     unittest.main()