From 86913cde387dcfce79576a63db0683898766b658 Mon Sep 17 00:00:00 2001 From: Hao Phan Date: Mon, 12 Oct 2026 15:13:39 +0700 Subject: [PATCH 1/4] [MiniMax-M3] Fix EAGLE3 verification on CUDA TARGET_VERIFY reaches the sparse layers' forward_extend, which builds prefill metadata from extend_seq_lens. Verify batches carry none, so EAGLE3 fails at the FlashInfer autotune warm-up on CUDA. ROCm already runs a linear draft chain through the sparse decode kernels, one causal row per draft token; run that path on every non-NPU GPU, for EAGLE chains only (draft_is_chain_layout). The CUDA verify graph also took its score-buffer bound from the 1-token capture dummies, so replay at a real context would write past the buffer. Bound verify capture by the context length on every platform and count the draft rows in the eager bound. Rows past the context length are never committed; clamp them so the bound, and with it the decode top-k register bucket, stays at the context length. Verify rows skip MSA decode, whose plan covers decode batches only. Co-Authored-By: Claude Opus 5.5 --- .../attention/minimax_sparse_backend.py | 55 ++-- .../amd/test_minimax_rocm_verify.py | 4 +- .../attention/test_minimax_linear_verify.py | 299 ++++++++++++++++++ 3 files changed, 335 insertions(+), 23 deletions(-) create mode 100644 test/registered/kernels/ops/attention/test_minimax_linear_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 b16e5039f4db..2f978d86884a 100644 --- a/python/sglang/srt/layers/attention/minimax_sparse_backend.py +++ b/python/sglang/srt/layers/attention/minimax_sparse_backend.py @@ -116,7 +116,6 @@ 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 @@ -269,6 +268,9 @@ def __init__(self, runner: ModelRunner): # unit bmm scales — no fp8 handling yet (follow-up). and not self.fp8_attn_gemm ) + from sglang.srt.layers.attention.minimax_sparse_ops.indexer_cp import ( + draft_is_chain_layout, + ) from sglang.srt.model_executor.cuda_graph_config import ( Backend, Phase, @@ -277,12 +279,15 @@ 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)." - ) + # GPU verify runs each draft token as one causal sparse-decode row. + self._linear_verify = not self.is_npu + if self._linear_verify and not draft_is_chain_layout( + spec.speculative_algorithm, spec.speculative_eagle_topk + ): + raise NotImplementedError( + "MiniMax-M3 speculative attention requires a linear EAGLE " + "draft chain (--speculative-eagle-topk 1)." + ) _decode_cuda_graph = not check_cuda_graph_backend( Phase.DECODE, Backend.DISABLED ) @@ -425,7 +430,7 @@ def init_forward_metadata_out_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() - or (self.is_hip and forward_batch.forward_mode.is_target_verify()) + or (self._linear_verify and forward_batch.forward_mode.is_target_verify()) ): bs = forward_batch.seq_lens.shape[0] if forward_batch.forward_mode.is_target_verify(): @@ -451,18 +456,19 @@ 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 or self.is_hip) - and forward_batch.forward_mode.is_target_verify() - ) + or 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 + if self._linear_verify and forward_batch.forward_mode.is_target_verify(): + # The last draft row reads ndt tokens past the accepted prefix. + self._max_seqlen_k = min( + self._max_seqlen_k + self.speculative_num_draft_tokens, + self.max_context_len, + ) # Build plan + page table eager (outside capture) so captured forward_decode # runs only device-side ops; host-side code can't be captured. @@ -550,7 +556,7 @@ 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): + def _init_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. @@ -560,15 +566,18 @@ def _init_rocm_linear_verify_metadata(self, forward_batch: ForwardBatch): device=forward_batch.seq_lens.device, dtype=forward_batch.seq_lens.dtype, ) + seq_lens = (forward_batch.seq_lens[:, None] + offsets[None, :]).reshape(-1) self._linear_verify_meta = SimpleNamespace( - seq_lens=(forward_batch.seq_lens[:, None] + offsets[None, :]).reshape(-1), + # Rows past the context length predict positions that are never committed; + # the clamp keeps them inside the context-length bound of the score buffer. + seq_lens=seq_lens.clamp(max=self.max_context_len), 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) + if forward_batch.forward_mode.is_target_verify(): + self._init_linear_verify_metadata(forward_batch) return # Layer-invariant decode/verify metadata as captured ops (re-read at replay). fm = forward_batch.forward_mode @@ -1454,7 +1463,7 @@ def forward_extend( idx_k: torch.Tensor, idx_v: Optional[torch.Tensor], ): - if self.is_hip and forward_batch.forward_mode.is_target_verify(): + if self._linear_verify 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") @@ -1755,7 +1764,11 @@ def attn_fn(main_q, page_table, real_seq_lens): ) msa_kv_indices = msa_plan = None - if self._use_msa_decode and attn_fn is None: + # The MSA decode plan covers decode batches only; verify rows stay on Triton. + use_msa = ( + self._use_msa_decode and not forward_batch.forward_mode.is_target_verify() + ) + if use_msa and attn_fn is None: if self._msa_dec_meta is not None: msa_kv_indices, msa_plan = self._msa_dec_meta elif q.shape[0] > 0: @@ -1830,7 +1843,7 @@ def hisparse_swap_in_fn(topk_idx): disable_index_value=disable_value, dense_main_attn_fn=attn_fn, page_size=self.page_size, - use_msa=self._use_msa_decode, + use_msa=use_msa, msa_kv_indices=msa_kv_indices, msa_plan=msa_plan, q_scale=layer.q_scale_float, diff --git a/test/registered/amd/test_minimax_rocm_verify.py b/test/registered/amd/test_minimax_rocm_verify.py index 2131cc595a56..94355dc59fb0 100644 --- a/test/registered/amd/test_minimax_rocm_verify.py +++ b/test/registered/amd/test_minimax_rocm_verify.py @@ -19,7 +19,7 @@ def setUp(self): from sglang.srt.model_executor.forward_batch_info import ForwardMode self.backend = MiniMaxSparseAttnBackend.__new__(MiniMaxSparseAttnBackend) - self.backend.is_hip = True + self.backend._linear_verify = True self.backend.speculative_num_draft_tokens = 4 self.backend._linear_verify_meta = None self.batch = SimpleNamespace( @@ -41,7 +41,7 @@ def decode(q, k, v, layer, batch, save_kv_cache, **kwargs): self.backend.forward_decode = decode def run_verify(self): - self.backend._init_rocm_linear_verify_metadata(self.batch) + self.backend._init_linear_verify_metadata(self.batch) return self.backend.forward_extend( self.q, self.q, diff --git a/test/registered/kernels/ops/attention/test_minimax_linear_verify.py b/test/registered/kernels/ops/attention/test_minimax_linear_verify.py new file mode 100644 index 000000000000..c23108ccb3aa --- /dev/null +++ b/test/registered/kernels/ops/attention/test_minimax_linear_verify.py @@ -0,0 +1,299 @@ +"""MiniMax-M3 EAGLE verify on CUDA: each draft row must see exactly the prefix up to its own token, +eagerly and when a graph captured on 1-token dummies replays at the end of the context. +""" + +import sys +from types import SimpleNamespace + +import pytest +import torch + +from sglang.srt.layers.attention.minimax_sparse_backend import ( + MiniMaxSparseAttnBackend, +) +from sglang.srt.layers.attention.minimax_sparse_ops.minimax_sparse import ( + minimax_sparse_decode, +) +from sglang.srt.model_executor.forward_batch_info import ForwardMode +from sglang.test.ci.ci_register import register_cuda_ci + +register_cuda_ci(est_time=60, stage="base-b-kernel-unit", runner_config="1-gpu-large") + +# One TP4 rank of MiniMax-M3: 16 query heads, 1 KV head, 1 index head. +NQH, NKH, NIH, HD = 16, 1, 1, 128 +BLK, TOPK, INIT_BLOCKS, LOCAL_BLOCKS, PAGE = 128, 16, 0, 1, 128 +NDT = 4 +MAX_CONTEXT = 64 * 1024 +# ReqToTokenPool width with speculative decoding: context + 4 + ndt. +WIDTH = MAX_CONTEXT + 4 + NDT +LAYER_ID = 3 + + +def _paged_req_to_token(num_reqs, first_page=0): + pages_per_req = (WIDTH + PAGE - 1) // PAGE + rows = [] + for r in range(num_reqs): + base = first_page + r * pages_per_req + pages = torch.randperm(pages_per_req, device="cuda") + base + slots = pages[:, None] * PAGE + torch.arange(PAGE, device="cuda") + rows.append(slots.reshape(-1)[:WIDTH]) + return torch.stack(rows).to(torch.int32), first_page + num_reqs * pages_per_req + + +def _caches(num_slots, with_value): + nan = float("nan") + k = torch.full((num_slots, NKH, HD), nan, dtype=torch.bfloat16, device="cuda") + v = torch.full_like(k, nan) + idx_k = torch.full((num_slots, 1, HD), nan, dtype=torch.bfloat16, device="cuda") + idx_v = torch.full_like(idx_k, nan) if with_value else None + return k, v, idx_k, idx_v + + +class _Pool: + """The slice of MiniMaxSparseKVPool that forward_decode touches.""" + + def __init__(self, k, v, idx_k, idx_v): + self.k, self.v, self.idx_k, self.idx_v = k, v, idx_k, idx_v + + def set_fused_kv_index_buffer(self, layer, loc, k, v, idx_k, idx_v, *scales): + self.k[loc] = k.view(-1, *self.k.shape[1:]) + self.v[loc] = v.view(-1, *self.v.shape[1:]) + self.idx_k[loc] = idx_k.view(-1, *self.idx_k.shape[1:]) + if idx_v is not None: + self.idx_v[loc] = idx_v.view(-1, *self.idx_v.shape[1:]) + + def get_kv_buffer(self, layer_id): + return self.k, self.v + + def get_index_k_buffer(self, layer_id): + return self.idx_k + + def get_index_kv_buffer(self, layer_id): + return self.idx_k, self.idx_v + + +def _backend(req_to_token, pool, with_value): + b = MiniMaxSparseAttnBackend.__new__(MiniMaxSparseAttnBackend) + b.is_npu = False + b._linear_verify = True + b.speculative_num_draft_tokens = NDT + b.max_context_len = MAX_CONTEXT + b.req_to_token = req_to_token + b.kv_pool = pool + b.page_size = PAGE + b.block_size_k, b.topk_blocks = BLK, TOPK + b.init_blocks, b.local_blocks = INIT_BLOCKS, LOCAL_BLOCKS + b.score_type = "max" + b.disable_value_layer_ids = set() if with_value else {LAYER_ID} + b.use_dense_sparse_decode = False + # MSA decode on (no plan is built for verify): verify rows must still run on Triton. + b._use_msa_decode, b._msa_owns_decode = True, False + b._msa_dec_meta = None + b.fp8_attn_gemm = False + b.index_cache_enabled = False + b._decode_topk_buf = {} + b._topk_is_source = {} + b.hisparse_coordinator = None + b.indexer_cp = None + b._linear_verify_meta = None + b._prefill_seqblock_meta = None + b._max_seqlen_q = b._max_seqlen_k = 1 + return b + + +_LAYER = SimpleNamespace( + layer_id=LAYER_ID, + q_scale_float=None, + k_scale_float=None, + v_scale_float=None, + idx_q_scale_float=None, + idx_k_scale_float=None, + idx_v_scale_float=None, +) + + +class _Case: + """bs requests with given prefixes, their prefix KV, and ndt draft tokens each.""" + + def __init__(self, prefixes, with_value, seed=0): + torch.manual_seed(seed) + self.with_value = with_value + self.bs = len(prefixes) + self.prefix = torch.tensor(prefixes, dtype=torch.int64, device="cuda") + self.req_to_token, num_pages = _paged_req_to_token(self.bs) + self.pool = _Pool(*_caches(num_pages * PAGE, with_value)) + rows = self.bs * NDT + self.q = torch.randn(rows, NQH, HD, dtype=torch.bfloat16, device="cuda") + self.k = torch.randn(rows, NKH, HD, dtype=torch.bfloat16, device="cuda") + self.v = torch.randn_like(self.k) + self.idx_q = torch.randn(rows, NIH, HD, dtype=torch.bfloat16, device="cuda") + self.idx_k = torch.randn(rows, 1, HD, dtype=torch.bfloat16, device="cuda") + self.idx_v = torch.randn_like(self.idx_k) if with_value else None + caches = [self.pool.k, self.pool.v, self.pool.idx_k, self.pool.idx_v] + for r, p in enumerate(prefixes): + slots = self.req_to_token[r, :p].long() + for cache in caches: + if cache is not None: + cache[slots] = torch.randn( + p, *cache.shape[1:], dtype=cache.dtype, device="cuda" + ) + draft_pos = self.prefix[:, None] + torch.arange(NDT, device="cuda") + self.out_cache_loc = ( + self.req_to_token.gather(1, draft_pos).reshape(-1).to(torch.int64) + ) + + def batch(self): + return SimpleNamespace( + forward_mode=ForwardMode.TARGET_VERIFY, + seq_lens=self.prefix.clone(), + seq_lens_cpu=self.prefix.cpu(), + req_pool_indices=torch.arange(self.bs, dtype=torch.int64, device="cuda"), + out_cache_loc=self.out_cache_loc.clone(), + extend_seq_lens=None, + extend_seq_lens_cpu=None, + extend_prefix_lens=None, + minimax_m3_precached_sparse_layers=None, + ) + + def reference(self): + """Decode every draft row alone, on a copy that ends at its own token.""" + rows = self.bs * NDT + ref_r2t, num_pages = _paged_req_to_token(rows) + k, v, idx_k, idx_v = _caches(num_pages * PAGE, self.with_value) + pairs = [ + (k, self.pool.k, self.k), + (v, self.pool.v, self.v), + (idx_k, self.pool.idx_k, self.idx_k), + (idx_v, self.pool.idx_v, self.idx_v), + ] + seq_lens = [] + for r in range(self.bs): + p = int(self.prefix[r]) + for j in range(NDT): + row = r * NDT + j + dst_prefix = ref_r2t[row, :p].long() + dst_draft = ref_r2t[row, p : p + j + 1].long() + src = self.req_to_token[r, :p].long() + for dst, pool_cache, draft in pairs: + if dst is None: + continue + dst[dst_prefix] = pool_cache[src] + dst[dst_draft] = draft[r * NDT : row + 1].view(-1, *dst.shape[1:]) + seq_lens.append(p + j + 1) + seq_lens = torch.tensor(seq_lens, dtype=torch.int64, device="cuda") + idx_o, o = minimax_sparse_decode( + self.q, + None, + k, + v, + self.idx_q, + None, + idx_k, + idx_v, + ref_r2t, + torch.arange(rows, dtype=torch.int64, device="cuda"), + seq_lens, + int(seq_lens.max()), + 1, + BLK, + TOPK, + INIT_BLOCKS, + LOCAL_BLOCKS, + score_type="max", + disable_index_value=not self.with_value, + page_size=PAGE, + ) + return ( + None if idx_o is None else idx_o.reshape(rows, -1), + o.reshape(rows, -1), + ) + + +def _verify(backend, batch, q, k, v, idx_q, idx_k, idx_v): + backend.init_forward_metadata_in_graph(batch) + return backend.forward_extend( + q, k, v, _LAYER, batch, idx_q=idx_q, idx_k=idx_k, idx_v=idx_v + ) + + +def _assert_same(got, want, exact, rows=None): + for g, w in zip(got, want): + if w is None: + assert g is None + continue + assert not torch.isnan(g).any() + if rows is not None: + g, w = g[rows], w[rows] + if exact: + assert torch.equal(g, w) + else: + torch.testing.assert_close(g, w, atol=2e-2, rtol=2e-2) + + +@pytest.mark.parametrize("with_value", [False, True]) +def test_eager_verify_matches_causal_rows(with_value): + """Rows of one request share its cache yet see only their own prefix.""" + # 5118 + ndt crosses into a new 128-token block, so the bound must count the drafts. + case = _Case([1000, 5118], with_value) + backend = _backend(case.req_to_token, case.pool, with_value) + batch = case.batch() + backend.init_forward_metadata_out_graph(batch) + assert backend._max_seqlen_k == 5118 + NDT + got = _verify( + backend, + batch, + case.q, + case.k, + case.v, + case.idx_q, + case.idx_k, + case.idx_v, + ) + _assert_same(got, case.reference(), exact=True) + # Verify leaves the batch's per-request metadata to the dense layers. + assert batch.seq_lens.tolist() == [1000, 5118] + assert batch.extend_seq_lens is None + + +@pytest.mark.parametrize("with_value", [False, True]) +def test_graph_replay_at_longer_context(with_value): + """A graph captured on 1-token dummies replays at the end of the context.""" + # Request 0's last two rows run past MAX_CONTEXT: they are clamped, never committed. + case = _Case([MAX_CONTEXT - 2, 33000], with_value) + backend = _backend(case.req_to_token, case.pool, with_value) + + static = case.batch() + static.seq_lens.fill_(1) # get_cuda_graph_seq_len_fill_value() + static.seq_lens_cpu = static.seq_lens.cpu() + # Capture-time KV writes land in the free tail of each row. + static.out_cache_loc.copy_( + case.req_to_token[:, WIDTH - NDT :].reshape(-1).to(torch.int64) + ) + inputs = [t.clone() for t in (case.q, case.k, case.v, case.idx_q, case.idx_k)] + inputs.append(None if case.idx_v is None else case.idx_v.clone()) + + backend.init_forward_metadata_out_graph(static, in_capture=True) + assert backend._max_seqlen_k == MAX_CONTEXT + _verify(backend, static, *inputs) # warm-up: autotune and JIT compile + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + out = _verify(backend, static, *inputs) + + static.seq_lens.copy_(case.prefix) + static.out_cache_loc.copy_(case.out_cache_loc) + graph.replay() + torch.cuda.synchronize() + # No row may read past the context-length bound the score buffer was sized for. + assert int(backend._linear_verify_meta.seq_lens.max()) == MAX_CONTEXT + # Compared to tolerance: the captured and the reference score buffers differ in width. + inside = [ + r + for r in range(case.bs * NDT) + if int(case.prefix[r // NDT]) + r % NDT + 1 <= MAX_CONTEXT + ] + _assert_same(out, case.reference(), exact=False, rows=inside) + + +if __name__ == "__main__": + sys.exit(pytest.main([__file__, "-v"])) From 84decace83b2364bf36aad31d2fd3d5c9f2a38e3 Mon Sep 17 00:00:00 2001 From: Hao Phan Date: Thu, 22 Oct 2026 10:13:31 +0700 Subject: [PATCH 2/4] Trim comments in the linear verify test --- .../kernels/ops/attention/test_minimax_linear_verify.py | 8 ++------ 1 file changed, 2 insertions(+), 6 deletions(-) diff --git a/test/registered/kernels/ops/attention/test_minimax_linear_verify.py b/test/registered/kernels/ops/attention/test_minimax_linear_verify.py index c23108ccb3aa..b044ad7d7ff3 100644 --- a/test/registered/kernels/ops/attention/test_minimax_linear_verify.py +++ b/test/registered/kernels/ops/attention/test_minimax_linear_verify.py @@ -1,6 +1,4 @@ -"""MiniMax-M3 EAGLE verify on CUDA: each draft row must see exactly the prefix up to its own token, -eagerly and when a graph captured on 1-token dummies replays at the end of the context. -""" +"""MiniMax-M3 EAGLE verify: each draft row sees exactly its causal prefix.""" import sys from types import SimpleNamespace @@ -113,8 +111,6 @@ def _backend(req_to_token, pool, with_value): class _Case: - """bs requests with given prefixes, their prefix KV, and ndt draft tokens each.""" - def __init__(self, prefixes, with_value, seed=0): torch.manual_seed(seed) self.with_value = with_value @@ -286,12 +282,12 @@ def test_graph_replay_at_longer_context(with_value): torch.cuda.synchronize() # No row may read past the context-length bound the score buffer was sized for. assert int(backend._linear_verify_meta.seq_lens.max()) == MAX_CONTEXT - # Compared to tolerance: the captured and the reference score buffers differ in width. inside = [ r for r in range(case.bs * NDT) if int(case.prefix[r // NDT]) + r % NDT + 1 <= MAX_CONTEXT ] + # The captured and reference score buffers differ in width; compare to tolerance. _assert_same(out, case.reference(), exact=False, rows=inside) From 69a54e109756e56dbd84ab3029cd1b469b623059 Mon Sep 17 00:00:00 2001 From: Hao Phan Date: Thu, 22 Oct 2026 15:26:33 +0700 Subject: [PATCH 3/4] Refuse HiSparse with speculative decoding; fix the MSA guard reason and the ROCm verify test bound --- .../layers/attention/minimax_sparse_backend.py | 15 +++++++++++++-- test/registered/amd/test_minimax_rocm_verify.py | 1 + 2 files changed, 14 insertions(+), 2 deletions(-) diff --git a/python/sglang/srt/layers/attention/minimax_sparse_backend.py b/python/sglang/srt/layers/attention/minimax_sparse_backend.py index 2f978d86884a..8f7f24e12451 100644 --- a/python/sglang/srt/layers/attention/minimax_sparse_backend.py +++ b/python/sglang/srt/layers/attention/minimax_sparse_backend.py @@ -288,6 +288,17 @@ def __init__(self, runner: ModelRunner): "MiniMax-M3 speculative attention requires a linear EAGLE " "draft chain (--speculative-eagle-topk 1)." ) + # HiSparse swap-in keeps one device buffer and one LRU slot set per request, + # while linear verify scores bs * ndt rows. + if ( + self._linear_verify + and self.hisparse_coordinator is not None + and spec.speculative_algorithm is not None + ): + raise NotImplementedError( + "MiniMax-M3 HiSparse does not support speculative decoding yet. " + "Disable --enable-hisparse or speculative decoding." + ) _decode_cuda_graph = not check_cuda_graph_backend( Phase.DECODE, Backend.DISABLED ) @@ -295,8 +306,8 @@ def __init__(self, runner: ModelRunner): not _decode_cuda_graph or envs.SGLANG_OPT_USE_MSA_DECODE_UNDER_GRAPH.get() ) - # MSA + spec decode + cuda graph crashes mid-capture: TARGET_VERIFY batches - # route to forward_extend, dereferencing absent extend metadata. Fail at startup. + # Verify rows take the Triton decode kernels, but MSA decode with speculative + # decoding under CUDA graph has never been validated; fail at startup. if ( self.use_msa and _decode_cuda_graph diff --git a/test/registered/amd/test_minimax_rocm_verify.py b/test/registered/amd/test_minimax_rocm_verify.py index 94355dc59fb0..bbf9f57c9271 100644 --- a/test/registered/amd/test_minimax_rocm_verify.py +++ b/test/registered/amd/test_minimax_rocm_verify.py @@ -20,6 +20,7 @@ def setUp(self): self.backend = MiniMaxSparseAttnBackend.__new__(MiniMaxSparseAttnBackend) self.backend._linear_verify = True + self.backend.max_context_len = 8192 self.backend.speculative_num_draft_tokens = 4 self.backend._linear_verify_meta = None self.batch = SimpleNamespace( From 70576c3b98cba6cb4c7e3e4c9d32cbc75f16256d Mon Sep 17 00:00:00 2001 From: Hao Phan Date: Mon, 26 Oct 2026 14:30:27 +0700 Subject: [PATCH 4/4] Keep ROCm's top-k 1 draft gate; run the linear verify test on B200 too --- .../srt/layers/attention/minimax_sparse_backend.py | 11 +++++++++-- .../ops/attention/test_minimax_linear_verify.py | 1 + 2 files changed, 10 insertions(+), 2 deletions(-) diff --git a/python/sglang/srt/layers/attention/minimax_sparse_backend.py b/python/sglang/srt/layers/attention/minimax_sparse_backend.py index 8f7f24e12451..88b1e80ab4e8 100644 --- a/python/sglang/srt/layers/attention/minimax_sparse_backend.py +++ b/python/sglang/srt/layers/attention/minimax_sparse_backend.py @@ -279,9 +279,16 @@ def __init__(self, runner: ModelRunner): spec = get_spec() self.speculative_num_draft_tokens = spec.speculative_num_draft_tokens - # GPU verify runs each draft token as one causal sparse-decode row. + # GPU verify runs each draft token as one causal sparse-decode row. ROCm + # accepts any draft with top-k 1; CUDA takes the chain-layout allowlist. self._linear_verify = not self.is_npu - if self._linear_verify and not draft_is_chain_layout( + if 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)." + ) + elif self._linear_verify and not draft_is_chain_layout( spec.speculative_algorithm, spec.speculative_eagle_topk ): raise NotImplementedError( diff --git a/test/registered/kernels/ops/attention/test_minimax_linear_verify.py b/test/registered/kernels/ops/attention/test_minimax_linear_verify.py index b044ad7d7ff3..a91d1e7fdf95 100644 --- a/test/registered/kernels/ops/attention/test_minimax_linear_verify.py +++ b/test/registered/kernels/ops/attention/test_minimax_linear_verify.py @@ -16,6 +16,7 @@ from sglang.test.ci.ci_register import register_cuda_ci register_cuda_ci(est_time=60, stage="base-b-kernel-unit", runner_config="1-gpu-large") +register_cuda_ci(est_time=60, stage="base-b-kernel-unit", runner_config="4-gpu-b200") # One TP4 rank of MiniMax-M3: 16 query heads, 1 KV head, 1 index head. NQH, NKH, NIH, HD = 16, 1, 1, 128