from __future__ import annotations import logging from typing import TYPE_CHECKING, Optional import torch from sglang.srt.configs.model_config import ( get_minimax_sparse_attention_config, get_minimax_sparse_disable_value_layer_ids, get_minimax_sparse_layer_ids, get_minimax_sparse_score_type, ) from sglang.srt.layers.attention.base_attn_backend import AttentionBackend from sglang.srt.layers.attention.minimax_sparse_ops.minimax_sparse import ( minimax_sparse_decode, minimax_sparse_prefill, ) from sglang.srt.mem_cache.memory_pool import MiniMaxSparseKVPool from sglang.srt.model_executor.forward_batch_info import ForwardBatch if TYPE_CHECKING: from sglang.srt.model_executor.model_runner import ModelRunner logger = logging.getLogger(__name__) class MiniMaxSparseAttnBackend(AttentionBackend): def __init__(self, runner: ModelRunner): assert isinstance(runner.token_to_kv_pool, MiniMaxSparseKVPool) self.kv_pool = runner.token_to_kv_pool self.req_to_token = runner.req_to_token_pool.req_to_token self.max_context_len = int(runner.model_config.context_len) hf_config = runner.model_config.hf_config sparse_cfg = get_minimax_sparse_attention_config(hf_config) self.idx_head_dim = sparse_cfg["sparse_index_dim"] self.dense_layer_ids, self.sparse_layer_ids = get_minimax_sparse_layer_ids( sparse_cfg ) self.disable_value_layer_ids: set[int] = set( get_minimax_sparse_disable_value_layer_ids(sparse_cfg) ) self.score_type: str = get_minimax_sparse_score_type(sparse_cfg) # Plain Python int so it is safe inside CUDA graphs (no .item() at graph time). self._max_seqlen_q: int = 1 self._max_seqlen_k: int = 1 self.block_size_q = 1 self.block_size_k = sparse_cfg["sparse_block_size"] if "sparse_init_block" in sparse_cfg: self.init_blocks = sparse_cfg["sparse_init_block"] else: init_tokens = sparse_cfg["sparse_init_tokens"] self.init_blocks = ( init_tokens + self.block_size_k - 1 ) // self.block_size_k if "sparse_local_block" in sparse_cfg: self.local_blocks = sparse_cfg["sparse_local_block"] else: local_tokens = sparse_cfg["sparse_local_tokens"] self.local_blocks = ( local_tokens + self.block_size_k - 1 ) // self.block_size_k + 1 self.topk_blocks = sparse_cfg["sparse_topk_blocks"] # MSA (fmha_sm100) is SM100-only; fall back to the Triton sparse path when # the kernel is unavailable or its constraints don't hold. from sglang.srt.environ import envs from sglang.srt.layers.attention.minimax_sparse_ops.msa import ( msa_available, ) # MSA (fmha_sm100) is bf16/fp16-only; an fp8 main KV cache must stay on the # Triton sparse path (it dequants fp8 on load). _main_kv_is_fp8 = self.kv_pool.main_pool.dtype in ( torch.float8_e4m3fn, torch.float8_e5m2, ) self.use_msa = ( not envs.SGLANG_DISABLE_MSA.get() and msa_available() and self.block_size_k == 128 and self.kv_pool.page_size == self.block_size_k and self.topk_blocks in (4, 8, 16, 32) and not _main_kv_is_fp8 ) if ( not self.use_msa and not envs.SGLANG_DISABLE_MSA.get() and msa_available() and self.block_size_k == 128 and self.kv_pool.page_size != self.block_size_k ): logger.warning( "MiniMax-M3 MSA decode disabled: page_size=%d != sparse block size " "%d. Pass --page-size 128 (with an attention backend that allows it, " "e.g. fa4) to enable the faster MSA kernel; falling back to the " "Triton sparse path.", self.kv_pool.page_size, self.block_size_k, ) self._msa_dec_meta = None if self.use_msa: from sglang.srt.runtime_context import get_parallel self.num_q_heads = ( runner.model_config.num_attention_heads // get_parallel().attn_tp_size ) self.num_kv_heads = self.kv_pool.main_pool.head_num self._msa_nb_max = ( self.max_context_len + self.block_size_k - 1 ) // self.block_size_k self._msa_cg: dict[int, tuple] = {} self.page_size = self.kv_pool.page_size self.use_dense_sparse_decode = ( envs.SGLANG_OPT_USE_MINIMAX_DENSE_SPARSE_DECODE.get() and self.block_size_k % self.page_size == 0 ) # MSA fmha_sm100 decode is NOT cuda-graph-safe: captured/replayed it returns # wrong results (~14% GSM8K loss on B200). Gate capture via cuda_graph_config, # not legacy disable_* flags — they disagree under config-native flags and would # capture the unsafe MSA decode kernel. from sglang.srt.model_executor.cuda_graph_config import ( Backend, Phase, check_cuda_graph_backend, ) _sa = getattr(runner, "server_args", None) _decode_cuda_graph = not check_cuda_graph_backend( Phase.DECODE, Backend.DISABLED ) self._use_msa_decode = self.use_msa and ( 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. if ( self.use_msa and _decode_cuda_graph and getattr(_sa, "speculative_algorithm", None) is not None ): raise NotImplementedError( "MiniMax-M3 MSA attention does not support speculative decoding under " "CUDA graph. Use --disable-cuda-graph, set SGLANG_DISABLE_MSA=1, or " "disable speculative decoding." ) self._msa_owns_decode = self._use_msa_decode and not ( self.use_dense_sparse_decode and self.kv_pool.main_pool.head_num == 1 ) self.dense_backend: Optional[AttentionBackend] = None logger.info( f"[MiniMaxSparse] Backend initialized " f"(score_type={self.score_type!r}, " f"main_attn={'MSA' if self.use_msa else 'triton'}, " f"disable_value_layers={sorted(self.disable_value_layer_ids)})" ) def init_forward_metadata_out_graph( self, forward_batch: ForwardBatch, in_capture: bool = False ): # cuda-graph replay views are a SimpleNamespace without extend_seq_lens_cpu, # and TARGET_VERIFY sets it to None despite is_extend() — getattr covers both. self._msa_dec_meta = None extend_lens = getattr(forward_batch, "extend_seq_lens_cpu", None) if extend_lens is not None: self._max_seqlen_q = int(max(extend_lens)) else: self._max_seqlen_q = 1 if in_capture and forward_batch.forward_mode.is_decode_or_idle(): self._max_seqlen_k = self.max_context_len else: self._max_seqlen_k = int(forward_batch.seq_lens_cpu.max().item()) # Build plan + page table eager (outside capture) so captured forward_decode # runs only device-side ops; host-side code can't be captured. if self._msa_owns_decode and forward_batch.forward_mode.is_decode_or_idle(): self._prepare_msa_decode_meta(forward_batch) def _prepare_msa_decode_meta(self, forward_batch: ForwardBatch): from sglang.srt.layers.attention.minimax_sparse_ops.msa import ( build_msa_decode_cg_plan, update_msa_decode_cg_meta, ) bs = forward_batch.seq_lens.shape[0] if bs == 0: return entry = self._msa_cg.get(bs) if entry is None: device = forward_batch.seq_lens.device plan = build_msa_decode_cg_plan( self.num_q_heads, self.num_kv_heads, self.block_size_k, self.topk_blocks, bs, device=device, ) kv_indices_buf = torch.zeros( bs * self._msa_nb_max, dtype=torch.int32, device=device ) entry = (plan, kv_indices_buf) self._msa_cg[bs] = entry plan, kv_indices_buf = entry update_msa_decode_cg_meta( plan, kv_indices_buf, self.req_to_token, forward_batch.req_pool_indices, forward_batch.seq_lens, self.block_size_k, self.topk_blocks, self.num_q_heads, self.num_kv_heads, ) self._msa_dec_meta = (kv_indices_buf, plan) def init_forward_metadata_in_graph(self, forward_batch: ForwardBatch): pass def init_cuda_graph_state(self, max_bs: int, max_num_tokens: int): pass def get_cuda_graph_seq_len_fill_value(self): return 1 @staticmethod def _is_sparse_kv_cached_by_fusion( forward_batch: ForwardBatch, layer_id: int ) -> bool: layer_ids = forward_batch.minimax_m3_precached_sparse_layers return layer_ids is not None and layer_id in layer_ids def forward( self, q, k, v, layer, forward_batch: ForwardBatch, save_kv_cache: bool = True, **kwargs, ): if forward_batch.forward_mode.is_idle(): idx_q = kwargs.get("idx_q") num_idx_heads = idx_q.shape[1] disable_value = layer.layer_id in self.disable_value_layer_ids idx_out: Optional[torch.Tensor] = ( None if disable_value else q.new_zeros(q.shape[0], num_idx_heads * self.idx_head_dim) ) out = q.new_zeros(q.shape[0], layer.tp_q_head_num * layer.v_head_dim) return idx_out, out else: return super().forward( q, k, v, layer, forward_batch, save_kv_cache, **kwargs ) def forward_extend( self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, layer, forward_batch: ForwardBatch, save_kv_cache=True, *, idx_q: torch.Tensor, idx_k: torch.Tensor, idx_v: Optional[torch.Tensor], ): 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 ) if not kv_cached_by_fusion: self.kv_pool.set_fused_kv_index_buffer( layer, forward_batch.out_cache_loc, k, v, idx_k, None if disable_value else idx_v, ) k_cache, v_cache = self.kv_pool.get_kv_buffer(layer.layer_id) if disable_value: idx_k_cache = self.kv_pool.get_index_k_buffer(layer.layer_id) idx_v_cache = None else: idx_k_cache, idx_v_cache = self.kv_pool.get_index_kv_buffer(layer.layer_id) cu_seqlens = torch.cat( [ torch.zeros( 1, dtype=torch.int32, device=forward_batch.extend_seq_lens.device ), forward_batch.extend_seq_lens.to(torch.int32).cumsum(0).to(torch.int32), ] ) seq_lens = forward_batch.seq_lens.to(torch.int32) if forward_batch.extend_prefix_lens is not None: prefix_lens = forward_batch.extend_prefix_lens.to(torch.int32) else: prefix_lens = torch.zeros_like(seq_lens) # DP attention pads q beyond the real token count for collective alignment; # trim to actual tokens so the sparse kernel sees consistent shapes. if forward_batch.extend_seq_lens_cpu is not None: actual_num_tokens = int(sum(forward_batch.extend_seq_lens_cpu)) else: actual_num_tokens = int(cu_seqlens[-1].item()) original_num_tokens = q.shape[0] if actual_num_tokens < original_num_tokens: q = q[:actual_num_tokens] idx_q = idx_q[:actual_num_tokens] idx_o, o = minimax_sparse_prefill( q, k_cache, v_cache, None, idx_q, idx_k_cache, idx_v_cache, None, self.req_to_token, forward_batch.req_pool_indices, cu_seqlens, seq_lens, prefix_lens, self._max_seqlen_q, self._max_seqlen_k, self.block_size_q, self.block_size_k, self.topk_blocks, self.init_blocks, self.local_blocks, score_type=self.score_type, disable_index_value=disable_value, use_msa=self.use_msa, seqlens_cpu=forward_batch.extend_seq_lens_cpu, ) if actual_num_tokens < original_num_tokens: pad_len = original_num_tokens - actual_num_tokens o = torch.cat([o, o.new_zeros(pad_len, *o.shape[1:])], dim=0) if idx_o is not None: idx_o = torch.cat( [idx_o, idx_o.new_zeros(pad_len, *idx_o.shape[1:])], dim=0 ) return ( ( None if idx_o is None else idx_o.reshape(original_num_tokens, -1).contiguous() ), o.reshape(original_num_tokens, -1).contiguous(), ) def _dense_sparse_main_decode( self, q: torch.Tensor, page_table: torch.Tensor, real_seq_lens: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, layer, forward_batch: ForwardBatch, ) -> torch.Tensor: from sglang.srt.layers.attention.trtllm_mha_backend import TRTLLMHAAttnBackend if isinstance(self.dense_backend, TRTLLMHAAttnBackend): import flashinfer ps = self.page_size nkv = 1 head_dim = q.size(-1) # [max_slots, nkv, D] -> [num_pages, page_size, nkv, D] # -> [num_pages, nkv, page_size, D] (HND, trtllm default) kc = k_cache.view(-1, ps, nkv, head_dim).permute(0, 2, 1, 3) vc = v_cache.view(-1, ps, nkv, head_dim).permute(0, 2, 1, 3) return flashinfer.decode.trtllm_batch_decode_with_kv_cache( # type: ignore query=q.contiguous(), kv_cache=(kc, vc), workspace_buffer=self.dense_backend.workspace_buffer, block_tables=page_table, seq_lens=real_seq_lens, max_seq_len=self.topk_blocks * self.block_size_k, bmm1_scale=layer.scaling, bmm2_scale=1.0, ) raise NotImplementedError( "dense sparse decode currently supports trtllm_mha only (fa3 is TODO)" ) def forward_decode( self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, layer, forward_batch: ForwardBatch, save_kv_cache: bool = True, *, idx_q: torch.Tensor, idx_k: torch.Tensor, idx_v: Optional[torch.Tensor], **kwargs, ): assert len(kwargs) == 0 disable_value = layer.layer_id in self.disable_value_layer_ids self.kv_pool.set_fused_kv_index_buffer( layer, forward_batch.out_cache_loc, k, v, idx_k, None if disable_value else idx_v, ) k_cache, v_cache = self.kv_pool.get_kv_buffer(layer.layer_id) if disable_value: idx_k_cache = self.kv_pool.get_index_k_buffer(layer.layer_id) idx_v_cache = None else: idx_k_cache, idx_v_cache = self.kv_pool.get_index_kv_buffer(layer.layer_id) attn_fn = None if self.use_dense_sparse_decode and k_cache.shape[1] == 1: def attn_fn(main_q, page_table, real_seq_lens): return self._dense_sparse_main_decode( main_q, page_table, real_seq_lens, k_cache, v_cache, layer, forward_batch, ) msa_kv_indices = msa_plan = None if self._use_msa_decode 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: # Rebuilding the plan inline would run host-side code inside # CUDA-graph capture; fail loudly instead. raise RuntimeError( "MSA decode metadata missing: init_forward_metadata_out_graph " "did not prepare the plan for this forward (gate mismatch)." ) idx_o, o = minimax_sparse_decode( q, None, k_cache, v_cache, idx_q, None, idx_k_cache, idx_v_cache, self.req_to_token, forward_batch.req_pool_indices, forward_batch.seq_lens, self._max_seqlen_k, 1, self.block_size_k, self.topk_blocks, self.init_blocks, self.local_blocks, score_type=self.score_type, disable_index_value=disable_value, dense_main_attn_fn=attn_fn, page_size=self.page_size, use_msa=self._use_msa_decode, msa_kv_indices=msa_kv_indices, msa_plan=msa_plan, ) return ( None if idx_o is None else idx_o.reshape(q.shape[0], -1).contiguous(), o.reshape(q.shape[0], -1).contiguous(), ) class MiniMaxHybridAttnBackend(AttentionBackend): def __init__( self, dense_backend: AttentionBackend, sparse_backend: MiniMaxSparseAttnBackend, sparse_layer_ids: list[int], ): self.dense = dense_backend self.sparse = sparse_backend self.sparse_layer_ids = sparse_layer_ids self.sparse.dense_backend = dense_backend def init_forward_metadata(self, forward_batch: ForwardBatch): self.sparse.init_forward_metadata(forward_batch) self.dense.init_forward_metadata(forward_batch) def init_forward_metadata_out_graph( self, forward_batch: ForwardBatch, in_capture: bool = False ): self.sparse.init_forward_metadata_out_graph(forward_batch, in_capture) self.dense.init_forward_metadata_out_graph(forward_batch, in_capture) def init_forward_metadata_in_graph(self, forward_batch: ForwardBatch): self.sparse.init_forward_metadata_in_graph(forward_batch) self.dense.init_forward_metadata_in_graph(forward_batch) def init_cuda_graph_state(self, max_bs: int, max_num_tokens: int): self.dense.init_cuda_graph_state(max_bs, max_num_tokens) self.sparse.init_cuda_graph_state(max_bs, max_num_tokens) def get_cuda_graph_seq_len_fill_value(self): return self.sparse.get_cuda_graph_seq_len_fill_value() def forward( self, q, k, v, layer, forward_batch: ForwardBatch, save_kv_cache: bool = True, **kwargs, ): if layer.layer_id in self.sparse_layer_ids: return self.sparse.forward( q, k, v, layer, forward_batch, save_kv_cache, **kwargs ) # DP attention pads q to an even length but flashinfer builds qo_indptr from # extend_seq_lens, so padded q.shape[0] != qo_indptr[-1] and paged-prefill # raises. Trim q and re-pad output; k/v stay untrimmed so KV-cache writes # align with out_cache_loc. mode = forward_batch.forward_mode if mode.is_extend() and forward_batch.extend_seq_lens_cpu is not None: actual_num_tokens = int(sum(forward_batch.extend_seq_lens_cpu)) original_num_tokens = q.shape[0] if actual_num_tokens < original_num_tokens: o = self.dense.forward( q[:actual_num_tokens], k, v, layer, forward_batch, save_kv_cache, **kwargs, ) pad_len = original_num_tokens - actual_num_tokens return torch.cat([o, o.new_zeros(pad_len, *o.shape[1:])], dim=0) return self.dense.forward( q, k, v, layer, forward_batch, save_kv_cache, **kwargs ) def forward_extend( self, q, k, v, layer, forward_batch: ForwardBatch, save_kv_cache: bool = True, **kwargs, ): if layer.layer_id in self.sparse_layer_ids: return self.sparse.forward_extend( q, k, v, layer, forward_batch, save_kv_cache, **kwargs ) else: return self.dense.forward_extend( q, k, v, layer, forward_batch, save_kv_cache, **kwargs ) def forward_decode( self, q, k, v, layer, forward_batch: ForwardBatch, save_kv_cache: bool = True, **kwargs, ): if layer.layer_id in self.sparse_layer_ids: return self.sparse.forward_decode( q, k, v, layer, forward_batch, save_kv_cache, **kwargs ) else: return self.dense.forward_decode( q, k, v, layer, forward_batch, save_kv_cache, **kwargs )