MiniMax-M3: share the sparse index top-k across layers and reuse the decode top-k buffer (#36527)

Co-authored-by: Kevin Mi <mikevin920@yahoo.com>

python/sglang/kernels/ops/attention/minimax_sparse/decode/flash_with_topk_idx.py

@@ -18,22 +18,53 @@ from ..common.utils import (
 )
 
 
-@triton.heuristics(
-    {
-        "BLOCK_SIZE_H": lambda args: max(
-            16, triton.next_power_of_2(args["gqa_group_size"])
-        ),
-        "BLOCK_SIZE_D": lambda args: triton.next_power_of_2(args["head_dim"]),
-        "BATCH_SIZE_BUCKET": lambda args: triton.next_power_of_2(args["batch_size"]),
-    }
-)
-@triton.autotune(
-    configs=[
+def _prune_decode_configs(configs, named_args, **kwargs):
+    """Drop autotune configs whose token tile is smaller than a sparse block.
+
+    BLOCKS_PER_K_BLOCK = BLOCK_SIZE_N // block_size is 0 for those, so they
+    cannot compile. Keep the full list if nothing survives (block_size larger
+    than every tile) and let triton report the real failure.
+    """
+    block_size = named_args["block_size"]
+    kept = [c for c in configs if c.kwargs["BLOCK_SIZE_N"] >= block_size]
+    return kept or list(configs)
+
+
+_ON_HIP = torch.version.hip is not None
+
+
+def _decode_score_block_n(args):
+    # gfx950 pick: 512 for tiny batches (few CTAs), else 128.
+    return max(512 if args["batch_size"] <= 4 else 128, args["block_size"])
+
+
+# On ROCm the autotune sweep is replaced by a fixed config plus the
+# BLOCK_SIZE_N heuristic above: runtime autotune can fire under CUDA-graph
+# capture and mistune.
+_DECODE_SCORE_CONFIGS = (
+    [triton.Config({}, num_warps=4, num_stages=1)]
+    if _ON_HIP
+    else [
         triton.Config({"BLOCK_SIZE_N": BN}, num_warps=nw, num_stages=ns)
         for BN in [64, 128, 256, 512]
         for nw in [4, 8, 16]
         for ns in [1, 2, 3]
-    ],
+    ]
+)
+
+_DECODE_SCORE_HEURISTICS = {
+    "BLOCK_SIZE_H": lambda args: max(
+        16, triton.next_power_of_2(args["gqa_group_size"])
+    ),
+    "BLOCK_SIZE_D": lambda args: triton.next_power_of_2(args["head_dim"]),
+    "BATCH_SIZE_BUCKET": lambda args: triton.next_power_of_2(args["batch_size"]),
+    **({"BLOCK_SIZE_N": _decode_score_block_n} if _ON_HIP else {}),
+}
+
+
+@triton.heuristics(_DECODE_SCORE_HEURISTICS)
+@triton.autotune(
+    configs=_DECODE_SCORE_CONFIGS,
     key=[
         "BATCH_SIZE_BUCKET",
         "gqa_group_size",
@@ -41,6 +72,9 @@ from ..common.utils import (
         "block_size",
         "SCORE_TYPE",
     ],
+    prune_configs_by=(
+        None if _ON_HIP else {"early_config_prune": _prune_decode_configs}
+    ),
 )
 @triton.jit
 def _decode_score_kernel(
@@ -228,6 +262,7 @@ def _decode_score_kernel(
         "HAS_SINK",
         "SCORE_TYPE",
     ],
+    prune_configs_by={"early_config_prune": _prune_decode_configs},
 )
 @triton.jit
 def _decode_score_attn_kernel(