diff --git a/python/sglang/srt/layers/attention/minimax_sparse_backend.py b/python/sglang/srt/layers/attention/minimax_sparse_backend.py index b16e5039f4db..88b1e80ab4e8 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,33 @@ 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: + # 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 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( + "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 ) @@ -290,8 +313,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 @@ -425,7 +448,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 +474,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 +574,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 +584,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 +1481,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 +1782,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 +1861,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..bbf9f57c9271 100644 --- a/test/registered/amd/test_minimax_rocm_verify.py +++ b/test/registered/amd/test_minimax_rocm_verify.py @@ -19,7 +19,8 @@ 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.max_context_len = 8192 self.backend.speculative_num_draft_tokens = 4 self.backend._linear_verify_meta = None self.batch = SimpleNamespace( @@ -41,7 +42,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..a91d1e7fdf95 --- /dev/null +++ b/test/registered/kernels/ops/attention/test_minimax_linear_verify.py @@ -0,0 +1,296 @@ +"""MiniMax-M3 EAGLE verify: each draft row sees exactly its causal prefix.""" + +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") +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 +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: + 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 + 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) + + +if __name__ == "__main__": + sys.exit(pytest.main([__file__, "-v"]))