From 0b09b5dc6349556eab8f9a340b4ab201c7e6281f Mon Sep 17 00:00:00 2001 From: Zijie Chen <300606707+zijiecode@users.noreply.github.com> Date: Tue, 22 Sep 2026 17:39:09 -0700 Subject: [PATCH] [AMD] Fix MiniMax-M3 EAGLE linear verification on ROCm --- .../attention/minimax_sparse_backend.py | 63 ++++++++- .../amd/test_minimax_rocm_verify.py | 127 ++++++++++++++++++ 2 files changed, 186 insertions(+), 4 deletions(-) create mode 100644 test/registered/amd/test_minimax_rocm_verify.py diff --git a/python/sglang/srt/layers/attention/minimax_sparse_backend.py b/python/sglang/srt/layers/attention/minimax_sparse_backend.py index 11171a61f4ba..4f47fa5f7281 100644 --- a/python/sglang/srt/layers/attention/minimax_sparse_backend.py +++ b/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 ) @@ -309,8 +318,6 @@ def __init__(self, runner: ModelRunner): # that wide. Head split mirrors MiniMaxM3 sparse attention's. self._idx_group_size = 1 if self.index_cache_enabled: - from sglang.srt.runtime_context import get_parallel - _num_idx_heads = max( sparse_cfg["sparse_num_index_heads"] // get_parallel().attn_tp_size, 1 ) @@ -410,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( @@ -433,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. @@ -527,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 @@ -1414,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 diff --git a/test/registered/amd/test_minimax_rocm_verify.py b/test/registered/amd/test_minimax_rocm_verify.py new file mode 100644 index 000000000000..2131cc595a56 --- /dev/null +++ b/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()