diff --git a/python/sglang/kernels/ops/attention/minimax_sparse/decode/flash_with_topk_idx.py b/python/sglang/kernels/ops/attention/minimax_sparse/decode/flash_with_topk_idx.py index 8f3a5e1f3..4e9b5b2a7 100644 --- a/python/sglang/kernels/ops/attention/minimax_sparse/decode/flash_with_topk_idx.py +++ b/python/sglang/kernels/ops/attention/minimax_sparse/decode/flash_with_topk_idx.py @@ -18,22 +18,53 @@ from ..common.utils import ( ) -@triton.heuristics( - { - "BLOCK_SIZE_H": lambda args: max( - 16, triton.next_power_of_2(args["gqa_group_size"]) - ), - "BLOCK_SIZE_D": lambda args: triton.next_power_of_2(args["head_dim"]), - "BATCH_SIZE_BUCKET": lambda args: triton.next_power_of_2(args["batch_size"]), - } -) -@triton.autotune( - configs=[ +def _prune_decode_configs(configs, named_args, **kwargs): + """Drop autotune configs whose token tile is smaller than a sparse block. + + BLOCKS_PER_K_BLOCK = BLOCK_SIZE_N // block_size is 0 for those, so they + cannot compile. Keep the full list if nothing survives (block_size larger + than every tile) and let triton report the real failure. + """ + block_size = named_args["block_size"] + kept = [c for c in configs if c.kwargs["BLOCK_SIZE_N"] >= block_size] + return kept or list(configs) + + +_ON_HIP = torch.version.hip is not None + + +def _decode_score_block_n(args): + # gfx950 pick: 512 for tiny batches (few CTAs), else 128. + return max(512 if args["batch_size"] <= 4 else 128, args["block_size"]) + + +# On ROCm the autotune sweep is replaced by a fixed config plus the +# BLOCK_SIZE_N heuristic above: runtime autotune can fire under CUDA-graph +# capture and mistune. +_DECODE_SCORE_CONFIGS = ( + [triton.Config({}, num_warps=4, num_stages=1)] + if _ON_HIP + else [ triton.Config({"BLOCK_SIZE_N": BN}, num_warps=nw, num_stages=ns) for BN in [64, 128, 256, 512] for nw in [4, 8, 16] for ns in [1, 2, 3] - ], + ] +) + +_DECODE_SCORE_HEURISTICS = { + "BLOCK_SIZE_H": lambda args: max( + 16, triton.next_power_of_2(args["gqa_group_size"]) + ), + "BLOCK_SIZE_D": lambda args: triton.next_power_of_2(args["head_dim"]), + "BATCH_SIZE_BUCKET": lambda args: triton.next_power_of_2(args["batch_size"]), + **({"BLOCK_SIZE_N": _decode_score_block_n} if _ON_HIP else {}), +} + + +@triton.heuristics(_DECODE_SCORE_HEURISTICS) +@triton.autotune( + configs=_DECODE_SCORE_CONFIGS, key=[ "BATCH_SIZE_BUCKET", "gqa_group_size", @@ -41,6 +72,9 @@ from ..common.utils import ( "block_size", "SCORE_TYPE", ], + prune_configs_by=( + None if _ON_HIP else {"early_config_prune": _prune_decode_configs} + ), ) @triton.jit def _decode_score_kernel( @@ -228,6 +262,7 @@ def _decode_score_kernel( "HAS_SINK", "SCORE_TYPE", ], + prune_configs_by={"early_config_prune": _prune_decode_configs}, ) @triton.jit def _decode_score_attn_kernel( diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index 6f64ca795..c0a2144dc 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -1579,6 +1579,14 @@ class Envs: # MiniMax-M3 MXFP8 MoE experimental fusion toggles (default off; A/B only). SGLANG_MINIMAX_M3_FUSED_SWIGLU_MXFP8 = EnvBool(False) SGLANG_MINIMAX_M3_FUSED_MOE_COMBINE = EnvBool(False) + + # MiniMax-M3 sparse-attention toggles for ROCm. + # Share one index top-k across every N sparse layers; 1 disables sharing. + # Changes which KV blocks the skip layers attend, so it applies on ROCm only + # (never under two-batch overlap); elsewhere the backend pins 1. + # 2 is the accuracy-safe default: higher values reuse staler selections + # in the skip layers. + SGLANG_MINIMAX_M3_INDEX_TOPK_FREQ = EnvInt(2) # MiniMax M3 NPU prefill MAIN-attention: route the sparse main attention through # the native Ascend FA op `torch.ops.npu.npu_fused_infer_attention_score` (FIA) # with a per-query CUSTOM block_table diff --git a/python/sglang/srt/layers/attention/minimax_sparse_backend.py b/python/sglang/srt/layers/attention/minimax_sparse_backend.py index 2a84669e6..dd3ee67d6 100644 --- a/python/sglang/srt/layers/attention/minimax_sparse_backend.py +++ b/python/sglang/srt/layers/attention/minimax_sparse_backend.py @@ -19,6 +19,7 @@ from sglang.srt.layers.attention.base_attn_backend import ( AttentionBackend, SharedReadEnds, ) +from sglang.srt.layers.moe.utils import is_tbo_enabled from sglang.srt.mem_cache.memory_pool import MiniMaxSparseKVPool from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode from sglang.srt.runtime_context import ( @@ -26,7 +27,7 @@ from sglang.srt.runtime_context import ( get_spec, ) from sglang.srt.server_args import m3_fp8_attn_gemm_enabled -from sglang.srt.utils import is_npu +from sglang.srt.utils import is_gfx95_supported, is_hip, is_npu if is_npu(): from sglang.kernels.ops.attention.minimax_sparse.common.index import ( @@ -147,6 +148,12 @@ class MiniMaxSparseAttnBackend(AttentionBackend): # NPU: per-forward cached metadata for the triton paths (rebuilt each forward). self._prefill_meta: Optional[SimpleNamespace] = None + # (owning ForwardBatch, cu_seqlens, seq_lens, prefix_lens, cu_seqblocks_q, + # max_seqblock_q, all_seqblock_q). The owner is part of the key because one + # metadata init can be followed by more than one ForwardBatch reaching the + # layers (two-batch overlap splits into two children with different + # extend_seq_lens); a hit requires the SAME object, not just a live cache. + self._prefill_seqblock_meta: Optional[tuple] = None self._extend_meta: Optional[SimpleNamespace] = None self._extend_meta_key: Optional[int] = None self._decode_seq_lens_i32_cg: dict[int, torch.Tensor] = {} @@ -276,10 +283,45 @@ class MiniMaxSparseAttnBackend(AttentionBackend): ) self.dense_backend: Optional[AttentionBackend] = None + self.index_topk_freq = ( + max(int(envs.SGLANG_MINIMAX_M3_INDEX_TOPK_FREQ.get()), 1) + if is_hip() and is_gfx95_supported() and not is_tbo_enabled() + else 1 + ) + self.index_cache_enabled = self.index_topk_freq > 1 + # topk_index_reduce widens the last dim to idx_group_size * topk_blocks + # (union of the group's selections), so the shared decode buffer must be + # 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 + ) + self._idx_group_size = max( + _num_idx_heads // self.kv_pool.main_pool.head_num, 1 + ) + # Persistent per-bs device buffer for decode top-k reuse. Allocated eagerly + # outside CUDA-graph capture; the captured graph only copies into and reads + # from a fixed address. + self._decode_topk_buf: dict = {} + self._topk_group_of_layer: dict[int, int] = {} + self._topk_is_source: dict[int, bool] = {} + for ordinal, lid in enumerate( + lid for lid in self.sparse_layer_ids if lid in self.disable_value_layer_ids + ): + group = ordinal // self.index_topk_freq + self._topk_group_of_layer[lid] = group + self._topk_is_source[lid] = (ordinal % self.index_topk_freq) == 0 + self._topk_cache: dict = {} + self._topk_cache_owner: Optional[ForwardBatch] = 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"index_topk_freq={self.index_topk_freq}, " f"msa_decode={self._use_msa_decode}, " f"msa_owns_decode={self._msa_owns_decode}, " f"decode_cuda_graph={_decode_cuda_graph}, " @@ -330,6 +372,23 @@ class MiniMaxSparseAttnBackend(AttentionBackend): ): # getattr covers replay views lacking extend_seq_lens_cpu and TARGET_VERIFY. self._msa_dec_meta = None + # New forward -> drop the per-forward index-cache top-k (prefill only). + if self.index_cache_enabled: + self._topk_cache = {} + 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(): + bs = forward_batch.seq_lens.shape[0] + 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( + (_nkv, bs, self.topk_blocks * self._idx_group_size), + dtype=torch.int32, + device=forward_batch.seq_lens.device, + ) + # Per-forward cache of the layer-invariant prefill seqblock trio. + self._prefill_seqblock_meta = None if self.is_npu: # Invalidate cached prefill/extend metadata; rebuilt on first sparse layer. self._prefill_meta = None @@ -1347,7 +1406,46 @@ class MiniMaxSparseAttnBackend(AttentionBackend): else: idx_k_cache, idx_v_cache = self.kv_pool.get_index_kv_buffer(layer.layer_id) - cu_seqlens, seq_lens, prefix_lens = self._resolve_extend_meta(forward_batch, q) + cached = self._prefill_seqblock_meta + if cached is None or cached[0] is not forward_batch: + cu_seqlens, seq_lens, prefix_lens = self._resolve_extend_meta( + forward_batch, q + ) + if self.is_npu: + cu_seqblocks_q = max_seqblock_q = all_seqblock_q = None + else: + from sglang.kernels.ops.attention.minimax_sparse.common.utils import ( + get_cu_seqblocks, + ) + + cu_seqblocks_q, max_seqblock_q, all_seqblock_q, _, _, _ = ( + get_cu_seqblocks( + cu_seqlens, + self._max_seqlen_q, + self.block_size_q, + self.block_size_k, + forward_batch.extend_seq_lens_cpu, + ) + ) + cached = ( + forward_batch, + cu_seqlens, + seq_lens, + prefix_lens, + cu_seqblocks_q, + max_seqblock_q, + all_seqblock_q, + ) + self._prefill_seqblock_meta = cached + ( + _, + cu_seqlens, + seq_lens, + prefix_lens, + cu_seqblocks_q, + max_seqblock_q, + all_seqblock_q, + ) = cached # DP attention pads q beyond real tokens; trim (CPU list avoids a sync). if forward_batch.extend_seq_lens_cpu is not None: @@ -1398,7 +1496,24 @@ class MiniMaxSparseAttnBackend(AttentionBackend): minimax_sparse_prefill, ) - idx_o, o = minimax_sparse_prefill( + # Index cache: only for disable_value layers (idx_o is None, + # so skipping the indexer has no output side effect). A group's source + # layer computes + stores the reduced top-k; the other layers reuse it. + use_index_cache = self.index_cache_enabled and disable_value + cached_topk_idx = None + want_topk = False + if use_index_cache: + if self._topk_cache_owner is not forward_batch: + self._topk_cache = {} + self._topk_cache_owner = forward_batch + group = self._topk_group_of_layer[layer.layer_id] + if self._topk_is_source[layer.layer_id]: + want_topk = True # compute and store for this group + else: + cached_topk_idx = self._topk_cache.get(group) + # Miss (e.g. source layer chunked differently) -> recompute safely. + + result = minimax_sparse_prefill( q, k_cache, v_cache, @@ -1423,13 +1538,23 @@ class MiniMaxSparseAttnBackend(AttentionBackend): disable_index_value=disable_value, use_msa=self.use_msa, seqlens_cpu=forward_batch.extend_seq_lens_cpu, + cu_seqblocks_q=cu_seqblocks_q, + max_seqblock_q=max_seqblock_q, + all_seqblock_q=all_seqblock_q, q_scale=layer.q_scale_float, k_scale=layer.k_scale_float, v_scale=layer.v_scale_float, idx_q_scale=layer.idx_q_scale_float, idx_k_scale=layer.idx_k_scale_float, idx_v_scale=layer.idx_v_scale_float, + cached_topk_idx=cached_topk_idx, + return_topk_idx=want_topk, ) + if want_topk: + idx_o, o, reduced_topk_idx = result + self._topk_cache[group] = reduced_topk_idx + else: + idx_o, o = result 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) @@ -1499,18 +1624,19 @@ class MiniMaxSparseAttnBackend(AttentionBackend): ): 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, - layer.k_scale_float, - layer.v_scale_float, - layer.idx_k_scale_float, - layer.idx_v_scale_float, - ) + if not self._is_sparse_kv_cached_by_fusion(forward_batch, layer.layer_id): + 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, + layer.k_scale_float, + layer.v_scale_float, + layer.idx_k_scale_float, + layer.idx_v_scale_float, + ) 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) @@ -1565,6 +1691,17 @@ class MiniMaxSparseAttnBackend(AttentionBackend): minimax_sparse_decode, ) + # Decode top-k reuse: group source layer computes+stores; skips reuse. + _use_reuse = self.index_cache_enabled and disable_value and attn_fn is None + _topk_buf = self._decode_topk_buf.get(q.shape[0]) if _use_reuse else None + _cached_topk = None + _want_topk = False + if _use_reuse and _topk_buf is not None: + if self._topk_is_source.get(layer.layer_id, True): + _want_topk = True + else: + _cached_topk = _topk_buf + idx_o, o = minimax_sparse_decode( q, None, @@ -1596,6 +1733,8 @@ class MiniMaxSparseAttnBackend(AttentionBackend): idx_q_scale=layer.idx_q_scale_float, idx_k_scale=layer.idx_k_scale_float, idx_v_scale=layer.idx_v_scale_float, + cached_topk_idx=_cached_topk, + topk_out=_topk_buf if _want_topk else None, ) return ( None if idx_o is None else idx_o.reshape(q.shape[0], -1).contiguous(), diff --git a/python/sglang/srt/layers/attention/minimax_sparse_ops/minimax_sparse.py b/python/sglang/srt/layers/attention/minimax_sparse_ops/minimax_sparse.py index 100b27e3a..02d995422 100644 --- a/python/sglang/srt/layers/attention/minimax_sparse_ops/minimax_sparse.py +++ b/python/sglang/srt/layers/attention/minimax_sparse_ops/minimax_sparse.py @@ -73,9 +73,18 @@ def minimax_sparse_prefill( idx_q_scale: Optional[float] = None, idx_k_scale: Optional[float] = None, idx_v_scale: Optional[float] = None, + cached_topk_idx: Optional[torch.Tensor] = None, + return_topk_idx: bool = False, ): """Run MiniMax-M3 sparse prefill. + Index cache: when ``cached_topk_idx`` is given, skip Step 1 + (the flash-index attention + top-k selection) and reuse the provided top-k + indices for Step 3's sparse attention. When ``return_topk_idx`` is True, the + reduced top-k tensor is returned as a third element so the caller can cache + it for later skip layers. Only valid for ``disable_index_value`` layers + (idx_o is None there, so skipping the indexer has no output side effect). + ``cu_seqblocks_q``, ``max_seqblock_q``, and ``all_seqblock_q`` are optional precomputed query-block metadata shared by the index and value sparse kernels. Supplying them avoids recomputing the same block layout twice. @@ -88,42 +97,52 @@ def minimax_sparse_prefill( ) # All seqlen is less than topk, use full attention - # Step 1: Flash attention with topk index (using index head) - idx_o, topk_idx = flash_prefill_with_topk_index( - q=idx_q, - k_cache=idx_k_cache, - v_cache=idx_v_cache, - sink=idx_sink, - req_to_token=req_to_token, - slot_ids=slot_ids, - cu_seqlens=cu_seqlens, - seq_lens=seq_lens, - prefix_lens=prefix_lens, - max_seqlen_q=max_seqlen_q, - max_seqlen_k=max_seqlen_k, - block_size_q=block_size_q, - block_size_k=block_size_k, - topk=topk, - init_blocks=init_blocks, - local_blocks=local_blocks, - sm_scale=idx_sm_scale, - score_type=score_type, - disable_index_value=disable_index_value, - cu_seqblocks_q=cu_seqblocks_q, - max_seqblock_q=max_seqblock_q, - all_seqblock_q=all_seqblock_q, - q_scale=idx_q_scale, - k_scale=idx_k_scale, - v_scale=idx_v_scale, - ) - # Step 2: Reduce topk idx if num_idx_heads > num_kv_heads - num_idx_heads = idx_q.shape[1] - num_kv_heads = k_cache.shape[1] - idx_group_size = num_idx_heads // num_kv_heads - if idx_group_size > 1: - topk_idx = topk_index_reduce( - topk_idx.view(num_kv_heads, idx_group_size, -1, topk), dim=1 + if cached_topk_idx is not None: + # Index cache hit: reuse a prior sparse layer's reduced + # top-k, skipping Step 1 (flash-index attention + top-k) and Step 2 + # (reduce). idx_o is unused downstream for disable_index_value layers. + idx_o = None + topk_idx = cached_topk_idx + else: + # Step 1: Flash attention with topk index (using index head) + idx_o, topk_idx = flash_prefill_with_topk_index( + q=idx_q, + k_cache=idx_k_cache, + v_cache=idx_v_cache, + sink=idx_sink, + req_to_token=req_to_token, + slot_ids=slot_ids, + cu_seqlens=cu_seqlens, + seq_lens=seq_lens, + prefix_lens=prefix_lens, + max_seqlen_q=max_seqlen_q, + max_seqlen_k=max_seqlen_k, + block_size_q=block_size_q, + block_size_k=block_size_k, + topk=topk, + init_blocks=init_blocks, + local_blocks=local_blocks, + sm_scale=idx_sm_scale, + score_type=score_type, + disable_index_value=disable_index_value, + cu_seqblocks_q=cu_seqblocks_q, + max_seqblock_q=max_seqblock_q, + all_seqblock_q=all_seqblock_q, + q_scale=idx_q_scale, + k_scale=idx_k_scale, + v_scale=idx_v_scale, ) + # Step 2: Reduce topk idx if num_idx_heads > num_kv_heads + num_idx_heads = idx_q.shape[1] + num_kv_heads = k_cache.shape[1] + idx_group_size = num_idx_heads // num_kv_heads + if idx_group_size > 1: + topk_idx = topk_index_reduce( + topk_idx.view(num_kv_heads, idx_group_size, -1, topk), dim=1 + ) + + # Reduced top-k cached by the caller for subsequent skip layers. + reduced_topk_idx = topk_idx # Step 3: Sparse attention using topk index (main head). The MSA path only # replaces this step; the indexer above is unchanged. MSA has no attn-sink # input, so keep the Triton path when sink is present. @@ -192,6 +211,8 @@ def minimax_sparse_prefill( k_scale=k_scale, v_scale=v_scale, ) + if return_topk_idx: + return idx_o, o, reduced_topk_idx return idx_o, o @@ -232,45 +253,72 @@ def minimax_sparse_decode( idx_q_scale: Optional[float] = None, idx_k_scale: Optional[float] = None, idx_v_scale: Optional[float] = None, + cached_topk_idx: Optional[torch.Tensor] = None, + topk_out: Optional[torch.Tensor] = None, ) -> Tuple[torch.Tensor, torch.Tensor]: - # Step 1: Flash decode with topk index (using index head). When the dense main - # attention is used, the indexer emits the page table directly (fused - # transform) instead of block ids, plus the per-query effective KV length. - idx_o, topk_idx, real_seq_lens = flash_decode_with_topk_idx( - q=idx_q, - sink=idx_sink, - k_cache=idx_k_cache, - v_cache=idx_v_cache, - req_to_token=req_to_token, - seq_lens=seq_lens, - max_seqlen=max_seqlen, - slot_ids=slot_ids, - block_size=block_size_k, - topk=topk, - init_blocks=init_blocks, - local_blocks=local_blocks, - sm_scale=idx_sm_scale, - score_type=score_type, - disable_index_value=disable_index_value, - use_dense_main_attn=dense_main_attn_fn is not None, - page_size=page_size, - q_scale=idx_q_scale, - k_scale=idx_k_scale, - v_scale=idx_v_scale, - ) + # Index top-k sharing for DECODE. A group's source layer passes ``topk_out`` + # (a persistent buffer) and publishes its reduced top-k there; the group's + # skip layers pass that buffer back as ``cached_topk_idx`` and skip Step 1 + # (flash-index decode + top-k) and Step 2 (reduce) entirely, never reading + # idx_k_cache. Only valid for disable_index_value layers on the + # non-dense-main path (idx_o is None there). All-device: CUDA-graph safe. + if cached_topk_idx is not None: + idx_o = None + real_seq_lens = None + topk_idx = cached_topk_idx + _skip_reduce = True + else: + _skip_reduce = False + # Step 1: Flash decode with topk index (using index head). When the dense main + # attention is used, the indexer emits the page table directly (fused + # transform) instead of block ids, plus the per-query effective KV length. + idx_o, topk_idx, real_seq_lens = flash_decode_with_topk_idx( + q=idx_q, + sink=idx_sink, + k_cache=idx_k_cache, + v_cache=idx_v_cache, + req_to_token=req_to_token, + seq_lens=seq_lens, + max_seqlen=max_seqlen, + slot_ids=slot_ids, + block_size=block_size_k, + topk=topk, + init_blocks=init_blocks, + local_blocks=local_blocks, + sm_scale=idx_sm_scale, + score_type=score_type, + disable_index_value=disable_index_value, + use_dense_main_attn=dense_main_attn_fn is not None, + page_size=page_size, + q_scale=idx_q_scale, + k_scale=idx_k_scale, + v_scale=idx_v_scale, + ) num_idx_heads = idx_q.shape[1] num_kv_heads = k_cache.shape[1] idx_group_size = num_idx_heads // num_kv_heads + assert dense_main_attn_fn is None or ( + cached_topk_idx is None and topk_out is None + ), "index top-k sharing is not available on the dense main-attention path" if dense_main_attn_fn is not None: # topk_idx is the page table; real_seq_lens is the per-query cache_seqlens assert idx_group_size == 1 o = dense_main_attn_fn(q, topk_idx, real_seq_lens) else: # Step 2: Reduce topk idx if num_idx_heads > num_kv_heads - if idx_group_size > 1: + if idx_group_size > 1 and not _skip_reduce: topk_idx = topk_index_reduce( topk_idx.view(num_kv_heads, idx_group_size, -1, topk), dim=1 ) + if topk_out is not None: + # Publish the group's selection to the persistent buffer the skip + # layers read (fixed address -> CUDA-graph safe). + if topk_out.shape != topk_idx.shape: + raise ValueError( + f"topk_out shape {tuple(topk_out.shape)} does not match " + f"reduced top-k shape {tuple(topk_idx.shape)}" + ) + topk_out.copy_(topk_idx) # Step 3: Sparse attention using topk index (main head). The MSA path # only replaces this step; keep the Triton path when sink is present. if use_msa and sink is None: