[AMD] Fix MiniMax-M3 EAGLE3 verification on ROCm

python/sglang/srt/layers/attention/minimax_sparse_backend.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import copy
 import logging
 import os
 from types import SimpleNamespace
@@ -115,6 +116,7 @@ class MiniMaxSparseAttnBackend(AttentionBackend):
     def __init__(self, runner: ModelRunner):
         assert isinstance(runner.token_to_kv_pool, MiniMaxSparseKVPool)
         self.is_npu = is_npu()
+        self.is_hip = is_hip()
         self.kv_pool = runner.token_to_kv_pool
         self.hisparse_coordinator = runner.hisparse_coordinator
         self.token_to_kv_pool = runner.token_to_kv_pool  # alias for TboAttnBackend
@@ -159,6 +161,7 @@ def __init__(self, runner: ModelRunner):
         self._extend_meta_key: Optional[int] = None
         self._decode_seq_lens_i32_cg: dict[int, torch.Tensor] = {}
         self._verify_meta_cg: dict[tuple, SimpleNamespace] = {}
+        self._linear_verify_meta: Optional[SimpleNamespace] = None
 
         self.block_size_q = 1
         self.block_size_k = sparse_cfg["sparse_block_size"]
@@ -274,6 +277,12 @@ def __init__(self, runner: ModelRunner):
 
         spec = get_spec()
         self.speculative_num_draft_tokens = spec.speculative_num_draft_tokens
+        if self.is_hip and spec.speculative_algorithm is not None:
+            if spec.speculative_eagle_topk != 1:
+                raise NotImplementedError(
+                    "MiniMax-M3 ROCm speculative attention requires a linear "
+                    "draft chain (--speculative-eagle-topk 1)."
+                )
         _decode_cuda_graph = not check_cuda_graph_backend(
             Phase.DECODE, Backend.DISABLED
         )
@@ -408,8 +417,13 @@ def init_forward_metadata_out_graph(
             self._topk_cache_owner = None
         # Decode top-k reuse: pre-allocate the per-bs persistent buffer so graph
         # capture never allocates. num_kv_heads == 1 at TP>=4 for M3.
-        if self.index_cache_enabled and forward_batch.forward_mode.is_decode_or_idle():
+        if self.index_cache_enabled and (
+            forward_batch.forward_mode.is_decode_or_idle()
+            or (self.is_hip and forward_batch.forward_mode.is_target_verify())
+        ):
             bs = forward_batch.seq_lens.shape[0]
+            if forward_batch.forward_mode.is_target_verify():
+                bs *= self.speculative_num_draft_tokens
             if bs > 0 and bs not in self._decode_topk_buf:
                 _nkv = self.kv_pool.main_pool.head_num
                 self._decode_topk_buf[bs] = torch.empty(
@@ -431,13 +445,18 @@ def init_forward_metadata_out_graph(
             self._max_seqlen_q = 1
         if in_capture and (
             forward_batch.forward_mode.is_decode_or_idle()
-            or (self.is_npu and forward_batch.forward_mode.is_target_verify())
+            or (
+                (self.is_npu or self.is_hip)
+                and forward_batch.forward_mode.is_target_verify()
+            )
         ):
             # Capture uses tiny dummy seq_lens; bound by full context so replay
             # (longer sequences) does not miss KV blocks.
             self._max_seqlen_k = self.max_context_len
         else:
             self._max_seqlen_k = int(forward_batch.seq_lens_cpu.max().item())
+            if self.is_hip and forward_batch.forward_mode.is_target_verify():
+                self._max_seqlen_k += self.speculative_num_draft_tokens
 
         # Build plan + page table eager (outside capture) so captured forward_decode
         # runs only device-side ops; host-side code can't be captured.
@@ -525,8 +544,25 @@ def _prepare_msa_decode_meta(self, forward_batch: ForwardBatch):
         )
         self._msa_dec_meta = (kv_indices_buf, plan)
 
+    def _init_rocm_linear_verify_metadata(self, forward_batch: ForwardBatch):
+        ndt = self.speculative_num_draft_tokens
+        # GPU seq_lens are the accepted prefix lengths. A linear EAGLE
+        # chain exposes one more KV token to each successive query.
+        offsets = torch.arange(
+            1,
+            ndt + 1,
+            device=forward_batch.seq_lens.device,
+            dtype=forward_batch.seq_lens.dtype,
+        )
+        self._linear_verify_meta = SimpleNamespace(
+            seq_lens=(forward_batch.seq_lens[:, None] + offsets[None, :]).reshape(-1),
+            req_pool_indices=forward_batch.req_pool_indices.repeat_interleave(ndt),
+        )
+
     def init_forward_metadata_in_graph(self, forward_batch: ForwardBatch):
         if not self.is_npu:
+            if self.is_hip and forward_batch.forward_mode.is_target_verify():
+                self._init_rocm_linear_verify_metadata(forward_batch)
             return
         # Layer-invariant decode/verify metadata as captured ops (re-read at replay).
         fm = forward_batch.forward_mode
@@ -1412,6 +1448,27 @@ def forward_extend(
         idx_k: torch.Tensor,
         idx_v: Optional[torch.Tensor],
     ):
+        if self.is_hip and forward_batch.forward_mode.is_target_verify():
+            meta = self._linear_verify_meta
+            if meta is None or meta.seq_lens.numel() != q.shape[0]:
+                raise RuntimeError("Missing MiniMax-M3 linear verify metadata")
+            # Keep the dense backend's original per-request metadata intact.
+            # Sparse decode accepts one causal query per row; cache stores
+            # still use the original flattened out_cache_loc exactly once.
+            verify_batch = copy.copy(forward_batch)
+            verify_batch.seq_lens = meta.seq_lens
+            verify_batch.req_pool_indices = meta.req_pool_indices
+            return self.forward_decode(
+                q,
+                k,
+                v,
+                layer,
+                verify_batch,
+                save_kv_cache,
+                idx_q=idx_q,
+                idx_k=idx_k,
+                idx_v=idx_v,
+            )
         disable_value = layer.layer_id in self.disable_value_layer_ids
         kv_cached_by_fusion = self._is_sparse_kv_cached_by_fusion(
             forward_batch, layer.layer_id

test/registered/amd/test_minimax_rocm_verify.py

@@ -0,0 +1,127 @@
+"""Causal verify metadata and routing with no ordinary prefill metadata."""
+
+import unittest
+from types import SimpleNamespace
+
+import torch
+
+from sglang.test.ci.ci_register import register_amd_ci
+
+register_amd_ci(est_time=15, suite="stage-b-test-1-gpu-small-amd")
+
+
+@unittest.skipUnless(torch.version.hip, "ROCm verify integration")
+class TestMiniMaxROCmVerify(unittest.TestCase):
+    def setUp(self):
+        from sglang.srt.layers.attention.minimax_sparse_backend import (
+            MiniMaxSparseAttnBackend,
+        )
+        from sglang.srt.model_executor.forward_batch_info import ForwardMode
+
+        self.backend = MiniMaxSparseAttnBackend.__new__(MiniMaxSparseAttnBackend)
+        self.backend.is_hip = True
+        self.backend.speculative_num_draft_tokens = 4
+        self.backend._linear_verify_meta = None
+        self.batch = SimpleNamespace(
+            forward_mode=ForwardMode.TARGET_VERIFY,
+            seq_lens=torch.tensor([127, 2047], dtype=torch.int32, device="cuda"),
+            req_pool_indices=torch.tensor([3, 1], dtype=torch.int64, device="cuda"),
+            out_cache_loc=torch.arange(8, dtype=torch.int64, device="cuda"),
+            extend_seq_lens=None,
+            extend_prefix_lens=None,
+        )
+        self.q = torch.zeros(8, 1, 128, dtype=torch.bfloat16, device="cuda")
+        self.seen = []
+
+        def decode(q, k, v, layer, batch, save_kv_cache, **kwargs):
+            self.seen.append(batch)
+            # Expose the metadata received by the unchanged sparse decoder.
+            return None, torch.stack((batch.req_pool_indices, batch.seq_lens), -1)
+
+        self.backend.forward_decode = decode
+
+    def run_verify(self):
+        self.backend._init_rocm_linear_verify_metadata(self.batch)
+        return self.backend.forward_extend(
+            self.q,
+            self.q,
+            self.q,
+            None,
+            self.batch,
+            idx_q=self.q,
+            idx_k=self.q,
+            idx_v=None,
+        )[1]
+
+    def assert_causal_rows(self, output, expected):
+        torch.testing.assert_close(
+            output,
+            torch.tensor(expected, dtype=torch.int64, device="cuda"),
+            rtol=0,
+            atol=0,
+        )
+        self.assertIsNot(self.seen[-1], self.batch)
+        self.assertIs(self.seen[-1].out_cache_loc, self.batch.out_cache_loc)
+        self.assertIsNone(self.batch.extend_seq_lens)
+        self.assertEqual(self.batch.seq_lens.numel(), 2)
+
+    def test_verify_keeps_original_batch_and_cache_locations(self):
+        output = self.run_verify()
+        self.assert_causal_rows(
+            output,
+            [
+                [3, 128],
+                [3, 129],
+                [3, 130],
+                [3, 131],
+                [1, 2048],
+                [1, 2049],
+                [1, 2050],
+                [1, 2051],
+            ],
+        )
+        self.assertEqual(self.batch.seq_lens.tolist(), [127, 2047])
+
+    def test_graph_replay_reads_new_prefixes_and_requests(self):
+        self.run_verify()
+        graph = torch.cuda.CUDAGraph()
+        with torch.cuda.graph(graph):
+            output = self.run_verify()
+        self.batch.seq_lens.copy_(
+            torch.tensor([4095, 255], dtype=torch.int32, device="cuda")
+        )
+        self.batch.req_pool_indices.copy_(
+            torch.tensor([2, 0], dtype=torch.int64, device="cuda")
+        )
+        graph.replay()
+        self.assert_causal_rows(
+            output,
+            [
+                [2, 4096],
+                [2, 4097],
+                [2, 4098],
+                [2, 4099],
+                [0, 256],
+                [0, 257],
+                [0, 258],
+                [0, 259],
+            ],
+        )
+        self.assertEqual(self.batch.seq_lens.tolist(), [4095, 255])
+
+    def test_missing_metadata_is_rejected(self):
+        with self.assertRaisesRegex(RuntimeError, "Missing MiniMax-M3 linear verify"):
+            self.backend.forward_extend(
+                self.q,
+                self.q,
+                self.q,
+                None,
+                self.batch,
+                idx_q=self.q,
+                idx_k=self.q,
+                idx_v=None,
+            )
+
+
+if __name__ == "__main__":
+    unittest.main()