[minimax-m3] fp8 attention GEMMs on SM100 (fp8_e4m3 KV + trtllm_mha) (#30971)

Co-authored-by: qiuyue <qiuyue@minimaxi.com>
Co-authored-by: xuebi <xuebi@minimaxi.com>
Co-authored-by: Xiaoyu Zhang <1182563586@qq.com>

python/sglang/kernels/jit/csrc/minimax/minimax_decode_topk.cuh

@@ -9,7 +9,7 @@
 
 #include <cfloat>
 #include <cstdint>
-#if defined(__HIP_PLATFORM_AMD__)
+#ifdef USE_ROCM
 static constexpr unsigned long long kWarpSyncMask = 0xFFFFFFFFFFFFFFFFull;
 #else
 #include <math_constants.h>
@@ -30,7 +30,9 @@ namespace {
 // The trivial case num_blocks <= topk (every block selected) is handled by the
 // kernels below, outside the Trait.
 struct TopKTrait {
-  static constexpr uint32_t kMaxTopK = 32;
+  // Also sizes the kernels' smem staging for the ascending-order emit; the
+  // block-id path's test contract goes up to topk == 64.
+  static constexpr uint32_t kMaxTopK = 64;
   static constexpr uint32_t kCTASize = 512;
   static constexpr uint32_t kNumWarps = kCTASize / device::kWarpThreads;
   static constexpr uint32_t kMaxNumBlocks = 4096;  // block topk
@@ -67,8 +69,14 @@ struct TopKTrait {
     };
     constexpr auto warp_inclusive_sum = [](uint32_t lane_id, uint32_t val) {
 #pragma unroll
-      for (uint32_t offset = 1; offset < 32; offset *= 2) {
+      for (uint32_t offset = 1; offset < device::kWarpThreads; offset *= 2) {
-        uint32_t n = __shfl_up_sync(kWarpSyncMask, val, offset, 32);
+        // Width-32 up-shuffle. On wave64 HIP the un-suffixed __shfl_up takes the
+        // logical-warp width directly; CUDA needs the active mask.
+#ifdef USE_ROCM
+        uint32_t n = __shfl_up(val, offset, device::kWarpThreads);
+#else
+        uint32_t n = __shfl_up_sync(kWarpSyncMask, val, offset, device::kWarpThreads);
+#endif
         if (lane_id >= offset) val += n;
       }
       return val;
@@ -264,8 +272,10 @@ struct TopKTrait {
 // Trait; otherwise the Trait selects the top-k block ids.
 // -------------------------------------------------------------------------
 
-// Block-id output: topk_idx[h, b, 0:k_eff) = selected block ids (front-packed,
+// Block-id output: topk_idx[h, b, 0:k_eff) = selected block ids (sorted
-// unordered), [k_eff:topk) = -1.
+// ascending), [k_eff:topk) = -1. Ascending order is a hard requirement of the
+// MSA fmha_sm100 consumer (kv_block_indexes must be strictly ascending; its
+// sorted-order early-exit otherwise mis-masks the partial last block).
 template <typename SeqLenT, bool kUsePDL>
 __global__ void minimax_decode_topk_block_kernel(
     const float* __restrict__ score,
@@ -297,7 +307,43 @@ __global__ void minimax_decode_topk_block_kernel(
 
   const float* __restrict__ row = score + (static_cast<int64_t>(h) * batch + b) * max_seqblock;
   __shared__ TopKTrait::Smem smem;
-  TopKTrait::forward(row, static_cast<uint32_t>(num_blocks), out, static_cast<uint32_t>(topk), &smem);
+  __shared__ int32_t s_topk[TopKTrait::kMaxTopK];
+  TopKTrait::forward(row, static_cast<uint32_t>(num_blocks), s_topk, static_cast<uint32_t>(topk), &smem);
+  __syncthreads();  // s_topk fully written before the sort reads it
+
+  // Emit ascending: num_blocks > topk here, so all topk slots hold distinct
+  // ids and rank(v) = |{x : x < v}| is a permutation. deepseek_v4
+  // topk_impl.cuh warp sort (32x32 / 64x64 branches; topk <= kMaxTopK = 64):
+  // lanes hold the elements in registers (INT32_MAX sentinel past topk), warp
+  // w ranks targets {w, w + kNumWarps, ...} via ballot+popc, lane 0 emits.
+  static_assert(TopKTrait::kMaxTopK <= 2 * device::kWarpThreads);
+  const auto warp_id = tx / device::kWarpThreads;
+  const auto lane_id = tx % device::kWarpThreads;
+#ifdef USE_ROCM
+  // wave64: __ballot spans the full 64-lane wave (would count the sibling
+  // 32-lane logical warp too); use the file's width-32 shuffle reduction.
+  const auto count_lt = [](int32_t x, int32_t v) { return device::warp::reduce_sum(static_cast<int>(x < v)); };
+#else
+  const auto count_lt = [](int32_t x, int32_t v) { return __popc(__ballot_sync(kWarpSyncMask, x < v)); };
+#endif
+  if (topk <= static_cast<int>(device::kWarpThreads)) {  // 32 x 32
+    const int32_t tie = (lane_id < static_cast<uint32_t>(topk)) ? s_topk[lane_id] : INT32_MAX;
+    for (uint32_t t = warp_id; t < static_cast<uint32_t>(topk); t += TopKTrait::kNumWarps) {
+      const int32_t target = s_topk[t];
+      const auto rank = count_lt(tie, target);
+      if (lane_id == 0) out[rank] = target;
+    }
+  } else {  // 64 x 64: each lane takes 2 elements
+    const int32_t tie_0 = s_topk[lane_id];
+    const int32_t tie_1 = (lane_id + device::kWarpThreads < static_cast<uint32_t>(topk))
+                              ? s_topk[lane_id + device::kWarpThreads]
+                              : INT32_MAX;
+    for (uint32_t t = warp_id; t < static_cast<uint32_t>(topk); t += TopKTrait::kNumWarps) {
+      const int32_t target = s_topk[t];
+      const auto rank = count_lt(tie_0, target) + count_lt(tie_1, target);
+      if (lane_id == 0) out[rank] = target;
+    }
+  }
 }
 
 // Page-table output: for each (batch b, kv-head h) pseudo-request emit the
@@ -436,6 +482,7 @@ void minimax_decode_topk(
       topk_i,
       ")");
   RuntimeCheck(block_size > 0, "block_size must be > 0, got ", block_size);
+  RuntimeCheck(topk <= static_cast<int64_t>(TopKTrait::kMaxTopK), "topk exceeds kMaxTopK (ascending-sort smem buffer)");
   if (batch == 0 || num_heads == 0) return;
 
   const dim3 grid(static_cast<unsigned>(batch), static_cast<unsigned>(num_heads));

python/sglang/kernels/ops/attention/minimax_decode_topk.py

@@ -3,8 +3,10 @@
 Drop-in replacement for the 2-stage split-K Triton topk
 (``_topk_index_partial_kernel`` + ``_topk_index_merge_kernel``): given the
 decode score tensor ``[num_heads, batch, max_seqblock]`` it produces
-``topk_idx`` ``[num_heads, batch, topk]`` (0-indexed block ids, front-packed,
+``topk_idx`` ``[num_heads, batch, topk]`` (0-indexed block ids, sorted
-``-1`` padded), matching the consumer ``_gqa_share_sparse_decode_kernel``.
+ascending, ``-1`` padded at the tail). Ascending order is required by the MSA
+fmha_sm100 consumer; the Triton ``_gqa_share_sparse_decode_kernel`` is
+order-insensitive.
 
 ``minimax_decode_topk_page_table`` additionally fuses the page-table transform
 for the dense paged backend (trtllm_mha / fa3) and returns the page table plus
@@ -17,7 +19,12 @@ from typing import TYPE_CHECKING, Tuple
 
 import torch
 
-from sglang.kernels.jit.utils import cache_once, load_jit, make_cpp_args
+from sglang.kernels.jit.utils import (
+    cache_once,
+    is_arch_support_pdl,
+    load_jit,
+    make_cpp_args,
+)
 
 if TYPE_CHECKING:
     from tvm_ffi.module import Module
@@ -25,7 +32,7 @@ if TYPE_CHECKING:
 
 @cache_once
 def _jit_module(seq_dtype: torch.dtype) -> Module:
-    args = make_cpp_args(seq_dtype, True)  # SeqLenT, kUsePDL
+    args = make_cpp_args(seq_dtype, is_arch_support_pdl())  # SeqLenT, kUsePDL
     return load_jit(
         "minimax_decode_topk",
         *args,

python/sglang/kernels/ops/attention/minimax_sparse/common/utils.py

@@ -10,11 +10,14 @@ import triton.language as tl
 
 _tma_keep_alive_buf = deque(maxlen=200)
 
-# Q is always bf16/fp16. The paged main K/V cache may be fp8 (unit-scaled) under
+# The paged main K/V cache may be fp8 (unit-scaled) under --kv-cache-dtype
-# --kv-cache-dtype fp8_*; the kernel widens it to the Q dtype on load (IS_FP8
+# fp8_*; with a bf16/fp16 Q the kernel widens K/V to the Q dtype on load
-# branch). Accepted on both HIP and CUDA (the bf16->fp8 cache write is unit-scaled,
+# (IS_FP8 branch — the bf16->fp8 cache write is unit-scaled, so the widening
-# so the widening cast is the exact inverse dequant). The bf16/fp16-only MSA
+# cast is the exact inverse dequant). Under fp8 attn-GEMM mode Q itself is
-# (fmha_sm100) kernel is excluded for fp8 KV by the backend use_msa gate.
+# fp8_e4m3: the IS_FP8 casts become no-ops and tl.dot runs fp8x8 on tensor
+# cores (P is quantized to the V dtype for the PV MMA, same contract as the
+# fmha_sm100 fp8 kernel). Accepted on both HIP and CUDA. MSA (fmha_sm100)
+# accepts fp8 only in the uniform-e4m3 fp8 attn-GEMM mode (backend gate).
 SPARSE_KV_FP8_DTYPES = (
     torch.float8_e4m3fn,
     torch.float8_e5m2,
@@ -25,26 +28,56 @@ SPARSE_KV_FP8_DTYPES = (
 def check_sparse_kv_fp8(
     q: torch.Tensor,
     k_cache: torch.Tensor,
-    v_cache: torch.Tensor,
+    v_cache: Optional[torch.Tensor],
     *,
     label: str,
 ) -> bool:
-    """Validate the sparse-attention KV cache dtype contract.
+    """Validate the sparse-attention Q/KV dtype contract.
 
-    Returns True iff the K cache is fp8 (then widened to Q dtype in the kernel).
+    Returns True iff the K cache is fp8 (drives the kernels' IS_FP8 constexpr).
-    Raises AssertionError otherwise, mirroring the contract the decode and prefill
+    Two fp8 modes are allowed:
-    topk kernels both enforce. fp8 is accepted on both HIP and CUDA.
+      * widening (Q bf16/fp16, K/V any fp8): K/V widened to Q dtype on load;
+      * all-fp8 GEMM (fp8 attn-GEMM mode): Q/K/V all fp8_e4m3fn. e5m2 Q is
+        rejected — fmha_sm100's variant lookup silently mis-dispatches e5m2
+        to the e4m3 kernel, so uniform e4m3 is enforced on the sglang side.
     """
-    assert q.dtype in (torch.bfloat16, torch.float16)
     is_fp8 = k_cache.dtype in SPARSE_KV_FP8_DTYPES
+    if q.dtype == torch.float8_e4m3fn:
+        assert k_cache.dtype == torch.float8_e4m3fn, (
+            f"sparse {label} with fp8 Q requires an fp8_e4m3fn K cache, "
+            f"got {k_cache.dtype}"
+        )
+    else:
+        assert q.dtype in (
+            torch.bfloat16,
+            torch.float16,
+        ), f"sparse {label} expects Q dtype bf16/fp16/fp8_e4m3fn, got {q.dtype}"
         assert k_cache.dtype == q.dtype or is_fp8, (
             f"sparse {label} expects K cache dtype == Q dtype ({q.dtype}) "
             f"or fp8, got {k_cache.dtype}"
         )
+    if v_cache is not None:
         assert v_cache.dtype == k_cache.dtype
     return is_fp8
 
 
+def sparse_out_dtype(q: torch.Tensor) -> torch.dtype:
+    """Attention output dtype: bf16 for fp8 Q (fp8 accumulates to bf16 out,
+    matching fmha_sm100's fp8 variant), else the Q dtype."""
+    return torch.bfloat16 if q.dtype == torch.float8_e4m3fn else q.dtype
+
+
+def unit_scale(scale: Optional[float]) -> float:
+    """Normalize an optional per-tensor dequant scale: None means unit scale.
+
+    All sparse-op entry points take ``Optional[float] = None`` scales (matching
+    ``k_scale_float`` / ``v_scale_float`` on RadixAttention, which are None
+    unless a checkpoint provides them) and normalize here at the kernel-launch
+    boundary, where a concrete float is needed.
+    """
+    return 1.0 if scale is None else scale
+
+
 try:
     make_tensor_descriptor = tl.make_tensor_descriptor
 except Exception:
@@ -260,3 +293,26 @@ def _bitonic_merge(
     for i in tl.static_range(stage):
         x, ids = _compare_and_swap(x, ids, flip, i + (n_dims - stage), n_dims)
     return x, ids
+
+
+@triton.jit
+def _sort_ids_ascending(ids, valid_n, BLOCK_SIZE_T: tl.constexpr):
+    """Sort a top-k id row ascending, invalids packed as -1 at the tail.
+
+    ``ids``: int32 [BLOCK_SIZE_T] block ids (0-indexed, -1 = invalid); entries at
+    positions >= ``valid_n`` are also treated as invalid. Valid ids must be
+    distinct and < 2**30 (block ids always are). The MSA fmha_sm100 consumer
+    requires kv_block_indexes strictly ascending with -1 tail padding — its
+    sorted-order early-exit otherwise mis-masks the partial last block.
+
+    O(T^2) rank sort (T = BLOCK_SIZE_T <= 64): invalid entries get unique keys
+    above every valid id so ranks form a permutation.
+    """
+    off = tl.arange(0, BLOCK_SIZE_T)
+    invalid = (off >= valid_n) | (ids < 0)
+    key = tl.where(invalid, 0x40000000 + off, ids)
+    rank = tl.sum(tl.where(key[None, :] < key[:, None], 1, 0), axis=1)
+    sorted_key = tl.sum(
+        tl.where(rank[None, :] == off[:, None], key[None, :], 0), axis=1
+    )
+    return tl.where(sorted_key >= 0x40000000, -1, sorted_key)

python/sglang/kernels/ops/attention/minimax_sparse/decode/flash_with_topk_idx.py

@@ -8,7 +8,14 @@ import triton.language as tl
 
 from sglang.srt.environ import envs
 
-from ..common.utils import _bitonic_merge, robust_allocator
+from ..common.utils import (
+    _bitonic_merge,
+    _sort_ids_ascending,
+    check_sparse_kv_fp8,
+    robust_allocator,
+    sparse_out_dtype,
+    unit_scale,
+)
 
 
 @triton.heuristics(
@@ -53,6 +60,8 @@ def _decode_score_kernel(
     topk: tl.constexpr,
     # sm_scale
     sm_scale,
+    # per-tensor K dequant scale (1.0 when the cache is unit-scaled)
+    k_scale,
     # init and local blocks
     init_blocks,
     local_blocks,
@@ -75,6 +84,7 @@ def _decode_score_kernel(
     NUM_KV_CHUNKS: tl.constexpr,
     SCORE_TYPE: tl.constexpr,
     SKIP_TRIVIAL_TOPK_SCORE: tl.constexpr,
+    IS_FP8: tl.constexpr,
 ):
     tl.static_assert(SCORE_TYPE == "max" or SCORE_TYPE == "lse")
     sm_scale_log2e = sm_scale * 1.4426950409
@@ -160,11 +170,15 @@ def _decode_score_kernel(
             mask=dim_mask[:, None] & pos_mask[None, :],
             other=0.0,
         )
+        if IS_FP8:
+            # fp8 index K cache: widening cast with bf16/fp16 Q, no-op with fp8
+            # Q (fp8 attn-GEMM mode; tl.dot runs fp8x8). Compiled out for bf16.
+            k = k.to(q.dtype)
         # compute qk
         qk = tl.zeros((BLOCK_SIZE_H, BLOCK_SIZE_N), dtype=tl.float32)
         qk += tl.where(off_n[None, :] < chunk_end - i, 0, float("-inf"))
         # [H, D], [D, N] -> [H, N]
-        qk += tl.dot(q, k) * sm_scale_log2e
+        qk += tl.dot(q, k) * (sm_scale_log2e * k_scale)
         # save qk to score
         score = tl.reshape(
             qk,
@@ -237,6 +251,9 @@ def _decode_score_attn_kernel(
     topk: tl.constexpr,
     # sm_scale
     sm_scale,
+    # per-tensor KV dequant scales (1.0 when the cache is unit-scaled)
+    k_scale,
+    v_scale,
     # init and local blocks
     init_blocks,
     local_blocks,
@@ -271,6 +288,7 @@ def _decode_score_attn_kernel(
     HAS_SINK: tl.constexpr,
     SCORE_TYPE: tl.constexpr,
     SKIP_TRIVIAL_TOPK_SCORE: tl.constexpr,
+    IS_FP8: tl.constexpr,
 ):
     tl.static_assert(SCORE_TYPE == "max" or SCORE_TYPE == "lse")
     sm_scale_log2e = sm_scale * 1.4426950409
@@ -367,6 +385,10 @@ def _decode_score_attn_kernel(
             mask=dim_mask[:, None] & pos_mask[None, :],
             other=0.0,
         )
+        if IS_FP8:
+            # fp8 index K cache: widening cast with bf16/fp16 Q, no-op with fp8
+            # Q (fp8 attn-GEMM mode; tl.dot runs fp8x8). Compiled out for bf16.
+            k = k.to(q.dtype)
         # load V as (BLOCK_SIZE_N, head_dim) via indirect addressing
         v_off = (
             slots[:, None] * stride_v_s
@@ -378,11 +400,15 @@ def _decode_score_attn_kernel(
             mask=pos_mask[:, None] & dim_mask[None, :],
             other=0.0,
         )
+        if IS_FP8:
+            # Cast V to the compute dtype (widening for bf16/fp16 Q; no-op for
+            # fp8 Q where P is quantized to e4m3 for the fp8 PV MMA).
+            v = v.to(q.dtype)
         # compute qk
         qk = tl.zeros((BLOCK_SIZE_H, BLOCK_SIZE_N), dtype=tl.float32)
         qk += tl.where(off_n[None, :] < chunk_end - i, 0, float("-inf"))
         # [H, D], [D, N] -> [H, N]
-        qk += tl.dot(q, k) * sm_scale_log2e
+        qk += tl.dot(q, k) * (sm_scale_log2e * k_scale)
         # save qk to score
         score = tl.reshape(
             qk,
@@ -426,7 +452,7 @@ def _decode_score_attn_kernel(
         acc_o_scale = tl.exp2(m_i - m_ij)
         acc_o = acc_o * acc_o_scale[:, None]
         # [H, N], [N, D] -> [H, D]
-        acc_o += tl.dot(p.to(v.dtype), v)
+        acc_o += tl.dot(p.to(v.dtype), v) * v_scale
         m_i = m_ij
         l_i = l_i * acc_o_scale + l_ij
         # update ptrs
@@ -744,8 +770,14 @@ def _topk_index_merge_kernel(
         + pid_b * stride_tif_b
         + off_t * stride_tif_t
     )
-    topk_idx_final = tl.where(off_t < tl.minimum(topk, num_blocks), topk_idx_final, -1)
+    # Ascending by block id, -1 tail: the MSA fmha_sm100 consumer requires
-    tl.store(tif_ptrs, topk_idx_final.to(ti_final_ptr.dtype.element_ty))
+    # sorted kv_block_indexes (the bitonic pass above orders by score).
+    topk_idx_final = _sort_ids_ascending(
+        topk_idx_final, tl.minimum(topk, num_blocks), BLOCK_SIZE_T
+    )
+    tl.store(
+        tif_ptrs, topk_idx_final.to(ti_final_ptr.dtype.element_ty), mask=off_t < topk
+    )
 
 
 @torch.no_grad()
@@ -768,18 +800,21 @@ def flash_decode_with_topk_idx(
     disable_index_value: bool = False,
     use_dense_main_attn: bool = False,  # NOTE: need transform idx in this case
     page_size: int = 1,
+    q_scale: Optional[float] = None,
+    k_scale: Optional[float] = None,
+    v_scale: Optional[float] = None,
 ) -> torch.Tensor:
     assert score_type in (
         "max",
         "lse",
     ), f"score_type must be 'max' or 'lse', got {score_type!r}"
     triton.set_allocator(robust_allocator)
-    # dtype check
+    # dtype check (v_cache is None under disable_index_value)
-    assert (
+    is_fp8 = check_sparse_kv_fp8(
-        q.dtype == torch.bfloat16
+        q, k_cache, None if disable_index_value else v_cache, label="decode indexer"
-        or q.dtype == torch.float16
-        and k_cache.dtype == q.dtype
     )
+    k_scale = unit_scale(k_scale)
+    v_scale = unit_scale(v_scale)
     if not disable_index_value:
         assert v_cache is not None
     # shape
@@ -793,6 +828,10 @@ def flash_decode_with_topk_idx(
     # sm scale
     if sm_scale is None:
         sm_scale = head_dim**-0.5
+    # q_scale folds exactly into sm_scale: it multiplies every Q-side logit —
+    # the QK dot AND the sink logit — unlike k_scale, which must not touch the
+    # sink term and therefore stays a separate kernel argument.
+    sm_scale = sm_scale * unit_scale(q_scale)
     # NUM_KV_CHUNKS controls how many parallel chunks each (batch, kv_head) gets.
     # Total CTAs = batch_size * NUM_KV_CHUNKS * num_kv_heads.
     # TARGET_GRID is the desired total CTA count; NUM_KV_CHUNKS is derived by:
@@ -848,6 +887,7 @@ def flash_decode_with_topk_idx(
             block_size,
             topk,
             sm_scale,
+            k_scale,
             init_blocks,
             local_blocks,
             q.stride(0),
@@ -863,6 +903,7 @@ def flash_decode_with_topk_idx(
             NUM_KV_CHUNKS=NUM_KV_CHUNKS,
             SCORE_TYPE=score_type,
             SKIP_TRIVIAL_TOPK_SCORE=skip_trivial_topk_score,
+            IS_FP8=is_fp8,
         )
     else:
         assert v_cache is not None
@@ -871,7 +912,7 @@ def flash_decode_with_topk_idx(
             batch_size,
             num_q_heads,
             head_dim,
-            dtype=q.dtype,
+            dtype=sparse_out_dtype(q),
             device=q.device,
         )
         lse = torch.empty(
@@ -895,6 +936,8 @@ def flash_decode_with_topk_idx(
             block_size,
             topk,
             sm_scale,
+            k_scale,
+            v_scale,
             init_blocks,
             local_blocks,
             q.stride(0),
@@ -922,6 +965,7 @@ def flash_decode_with_topk_idx(
             NUM_KV_CHUNKS=NUM_KV_CHUNKS,
             SCORE_TYPE=score_type,
             SKIP_TRIVIAL_TOPK_SCORE=skip_trivial_topk_score,
+            IS_FP8=is_fp8,
         )
     # Fused top-k + page-table transform: emit the dense backend's page table
     # directly (page-size-aware) instead of block ids, skipping a separate gather.

python/sglang/kernels/ops/attention/minimax_sparse/decode/topk_sparse.py

@@ -6,7 +6,12 @@ import torch
 import triton
 import triton.language as tl
 
-from ..common.utils import check_sparse_kv_fp8, robust_allocator
+from ..common.utils import (
+    check_sparse_kv_fp8,
+    robust_allocator,
+    sparse_out_dtype,
+    unit_scale,
+)
 
 
 @triton.heuristics(
@@ -49,6 +54,9 @@ def _gqa_share_sparse_decode_kernel(
     max_kv_len,
     # sm_scale
     sm_scale,
+    # per-tensor KV dequant scales (1.0 when the cache is unit-scaled)
+    k_scale,
+    v_scale,
     # stride
     stride_q_b,
     stride_q_h,
@@ -178,10 +186,12 @@ def _gqa_share_sparse_decode_kernel(
             other=0.0,
         )
         if IS_FP8:
-            # fp8 KV cache is unit-scaled (set_kv_buffer casts bf16->fp8 with no
+            # fp8 KV cache: with bf16/fp16 Q this widens K to the compute dtype
-            # scale), so dequant is just a widening cast to the Q compute dtype
+            # (unit-scaled cache -> exact inverse dequant; k_scale covers
-            # before the tl.dot. Matches the bf16 path bit-for-bit when the cache
+            # calibrated caches). With fp8 Q (fp8 attn-GEMM mode) the cast is a
-            # is bf16 (IS_FP8 False -> this branch is compiled out).
+            # no-op and tl.dot below runs fp8x8 on tensor cores. Matches the
+            # bf16 path bit-for-bit when the cache is bf16 (IS_FP8 False ->
+            # this branch is compiled out).
             k = k.to(q.dtype)
         # load V as (BLOCK_SIZE_N, head_dim) via indirect addressing
         v_off = (
@@ -195,15 +205,16 @@ def _gqa_share_sparse_decode_kernel(
             other=0.0,
         )
         if IS_FP8:
-            # Widen V before the P@V dot. This also makes the `p.to(v.dtype)`
+            # Cast V to the compute dtype. With bf16/fp16 Q this widens (so the
-            # below cast P to the compute dtype (not to fp8, which would be
+            # `p.to(v.dtype)` below keeps P in the compute dtype); with fp8 Q it
-            # catastrophic precision loss on the attention weights).
+            # is a no-op and P is quantized to e4m3 for the fp8 PV MMA — the
+            # same accuracy contract as fmha_sm100's fp8 kernel.
             v = v.to(q.dtype)
         # compute qk
         qk = tl.zeros((BLOCK_SIZE_H, BLOCK_SIZE_N), dtype=tl.float32)
         qk += tl.where(off_n[None, :] < seq_len - c, 0, float("-inf"))
         # [H, D], [D, N] -> [H, N]
-        qk += tl.dot(q, k) * sm_scale
+        qk += tl.dot(q, k) * (sm_scale * k_scale)
         # compute m_ij and l_ij
         m_ij = tl.maximum(m_i, tl.max(qk, axis=1))
         p = tl.exp(qk - m_ij[:, None])
@@ -212,9 +223,8 @@ def _gqa_share_sparse_decode_kernel(
         acc_o_scale = tl.exp(m_i - m_ij)
         acc_o = acc_o * acc_o_scale[:, None]
         # load v and update acc_o
-        p = p.to(v.dtype)
         # [H, N], [N, D] -> [H, D]
-        acc_o += tl.dot(p.to(v.dtype), v)
+        acc_o += tl.dot(p.to(v.dtype), v) * v_scale
         # update statistics
         m_i = m_ij
         lse_i = m_ij + tl.log(tl.exp(lse_i - m_ij) + l_ij)
@@ -308,9 +318,14 @@ def flash_decode_with_gqa_share_sparse(
     topk_idx: torch.Tensor,  # [num_kv_heads, batch_size, topk]
     sm_scale: Optional[float] = None,
     use_tma: bool = True,
+    q_scale: Optional[float] = None,
+    k_scale: Optional[float] = None,
+    v_scale: Optional[float] = None,
 ) -> torch.Tensor:
     triton.set_allocator(robust_allocator)
     is_fp8 = check_sparse_kv_fp8(q, k_cache, v_cache, label="decode")
+    k_scale = unit_scale(k_scale)
+    v_scale = unit_scale(v_scale)
     # shape
     batch_size, num_q_heads, head_dim = q.shape
     max_slots, num_kv_heads, _ = k_cache.shape
@@ -328,6 +343,9 @@ def flash_decode_with_gqa_share_sparse(
     # sm scale
     if sm_scale is None:
         sm_scale = head_dim**-0.5
+    # q_scale multiplies every Q-side logit (QK dot and sink), so it folds into
+    # sm_scale; k_scale must not touch the sink term and stays a kernel arg.
+    sm_scale = sm_scale * unit_scale(q_scale)
     # Pick NUM_TOPK_CHUNKS so total grid ≈ TARGET_GRID. Same constraints as
     # flash_decode_with_topk_idx: must be power of 2 (Triton arange) and must
     # only depend on shape constants (so grid is fixed within a cuda graph).
@@ -345,7 +363,7 @@ def flash_decode_with_gqa_share_sparse(
         batch_size,
         num_q_heads,
         head_dim,
-        dtype=q.dtype,
+        dtype=sparse_out_dtype(q),
         device=q.device,
     )
     lse_partial = torch.empty(
@@ -375,6 +393,8 @@ def flash_decode_with_gqa_share_sparse(
         max_topk,
         max_kv_len,
         sm_scale,
+        k_scale,
+        v_scale,
         q.stride(0),
         q.stride(1),
         q.stride(2),

python/sglang/kernels/ops/attention/minimax_sparse/prefill/flash_with_topk_idx.py

@@ -6,7 +6,15 @@ import torch
 import triton
 import triton.language as tl
 
-from ..common.utils import _bitonic_merge, get_cu_seqblocks, robust_allocator
+from ..common.utils import (
+    _bitonic_merge,
+    _sort_ids_ascending,
+    check_sparse_kv_fp8,
+    get_cu_seqblocks,
+    robust_allocator,
+    sparse_out_dtype,
+    unit_scale,
+)
 
 
 @triton.heuristics(
@@ -81,6 +89,9 @@ def _flash_attn_fwd_with_block_score_kernel(
     block_size: tl.constexpr,
     # sm_scale
     sm_scale,
+    # per-tensor KV dequant scales (1.0 when the cache is unit-scaled)
+    k_scale,
+    v_scale,
     # gumbel topk
     use_gumbel_topk: tl.constexpr,
     gumbel_seed,
@@ -112,6 +123,7 @@ def _flash_attn_fwd_with_block_score_kernel(
     HAS_SINK: tl.constexpr,
     SCORE_TYPE: tl.constexpr,
     DISABLE_INDEX_VALUE: tl.constexpr,
+    IS_FP8: tl.constexpr,
 ):
     tl.static_assert(SCORE_TYPE == "max" or SCORE_TYPE == "lse")
     sm_scale_log2e = sm_scale * 1.4426950409
@@ -170,7 +182,10 @@ def _flash_attn_fwd_with_block_score_kernel(
     if HAS_SINK:
         m_i = tl.zeros((BLOCK_SIZE_Q,), dtype=tl.float32)
         lse_i = tl.zeros((BLOCK_SIZE_Q,), dtype=tl.float32)
-        qsink = tl.sum(q * sink[None, :], axis=1) * sm_scale_log2e  # (BLOCK_SIZE_Q,)
+        qsink = (
+            tl.sum(q.to(tl.float32) * sink[None, :].to(tl.float32), axis=1)
+            * sm_scale_log2e
+        )  # (BLOCK_SIZE_Q,)
         m_i += qsink
         lse_i += qsink
     else:
@@ -199,8 +214,12 @@ def _flash_attn_fwd_with_block_score_kernel(
             mask=kd_mask[:, None] & pos_mask[None, :],
             other=0.0,
         )
+        if IS_FP8:
+            # fp8 index K cache: widening cast with bf16/fp16 Q, no-op with fp8
+            # Q (fp8 attn-GEMM mode; tl.dot runs fp8x8). Compiled out for bf16.
+            k = k.to(q.dtype)
         # compute qk
-        qk = tl.dot(q, k) * sm_scale_log2e
+        qk = tl.dot(q, k) * (sm_scale_log2e * k_scale)
         if i >= diag_start:
             qk = tl.where(off_q[:, None] >= (i + off_k)[None, :], qk, float("-inf"))
         # K boundary mask: positions beyond seq_len contribute -inf
@@ -251,8 +270,12 @@ def _flash_attn_fwd_with_block_score_kernel(
                 mask=pos_mask[:, None] & vd_mask[None, :],
                 other=0.0,
             )
+            if IS_FP8:
+                # Cast V to the compute dtype (widening for bf16/fp16 Q; no-op
+                # for fp8 Q where P is quantized to e4m3 for the fp8 PV MMA).
+                v = v.to(q.dtype)
             p = p.to(v.dtype)
-            acc_o += tl.dot(p, v)
+            acc_o += tl.dot(p, v) * v_scale
             # update statistics
             m_i = m_ij
             lse_i = m_ij + tl.log2(tl.exp2(lse_i - m_ij) + l_ij)
@@ -402,6 +425,9 @@ def _topk_index_kernel(
         * tl.reshape(topk_idx - 1, [BLOCK_SIZE_K // BLOCK_SIZE_T, BLOCK_SIZE_T]),
         axis=0,
     )
+    # Ascending by block id, -1 tail: the MSA fmha_sm100 consumer requires
+    # sorted kv_block_indexes (the bitonic pass above orders by score).
+    topk_idx = _sort_ids_ascending(topk_idx, min(topk, valid_blocks), BLOCK_SIZE_T)
     # save topk
     ti_ptrs = (
         ti_ptr
@@ -438,15 +464,21 @@ def flash_prefill_with_topk_index(
     cu_seqblocks_q: Optional[torch.Tensor] = None,
     max_seqblock_q: Optional[int] = None,
     all_seqblock_q: Optional[int] = None,
+    q_scale: Optional[float] = None,
+    k_scale: Optional[float] = None,
+    v_scale: Optional[float] = None,
 ):
     assert score_type in (
         "max",
         "lse",
     ), f"score_type must be 'max' or 'lse', got {score_type!r}"
     triton.set_allocator(robust_allocator)
-    # dtype check
+    # dtype check (v_cache is None under disable_index_value)
-    assert q.dtype == torch.bfloat16 or q.dtype == torch.float16
+    is_fp8 = check_sparse_kv_fp8(
-    assert k_cache.dtype == q.dtype
+        q, k_cache, None if disable_index_value else v_cache, label="prefill indexer"
+    )
+    k_scale = unit_scale(k_scale)
+    v_scale = unit_scale(v_scale)
     assert cu_seqlens.dtype == torch.int32
     # shape
     total_q, num_heads, qk_head_dim = q.shape
@@ -455,7 +487,7 @@ def flash_prefill_with_topk_index(
         # placeholder for BLOCK_SIZE_VD; V is never loaded
         v_head_dim = qk_head_dim
     else:
-        assert v_cache is not None and v_cache.dtype == q.dtype
+        assert v_cache is not None
         assert v_cache.shape[1] == k_cache.shape[1]
         v_head_dim = v_cache.shape[-1]
     gqa_group_size = num_heads // num_kv_heads
@@ -468,6 +500,9 @@ def flash_prefill_with_topk_index(
     ), "init_blocks + local_blocks must be less than topk"
     if sm_scale is None:
         sm_scale = qk_head_dim**-0.5
+    # q_scale multiplies every Q-side logit (QK dot and sink), so it folds into
+    # sm_scale; k_scale must not touch the sink term and stays a kernel arg.
+    sm_scale = sm_scale * unit_scale(q_scale)
     if cu_seqblocks_q is None or max_seqblock_q is None or all_seqblock_q is None:
         cu_seqblocks_q, max_seqblock_q, all_seqblock_q, _, _, _ = get_cu_seqblocks(
             cu_seqlens, max_seqlen_q, block_size_q, block_size_k
@@ -476,7 +511,9 @@ def flash_prefill_with_topk_index(
     if disable_index_value:
         o = None
     else:
-        o = torch.empty(total_q, num_heads, v_head_dim, dtype=q.dtype, device=q.device)
+        o = torch.empty(
+            total_q, num_heads, v_head_dim, dtype=sparse_out_dtype(q), device=q.device
+        )
     score = torch.full(
         (num_heads, total_q, max_seqblock_k),
         float("-inf"),
@@ -507,6 +544,8 @@ def flash_prefill_with_topk_index(
         v_head_dim,
         block_size_k,
         sm_scale,
+        k_scale,
+        v_scale,
         False,
         1,
         q.stride(0),
@@ -529,6 +568,7 @@ def flash_prefill_with_topk_index(
         req_to_token.stride(0),
         SCORE_TYPE=score_type,
         DISABLE_INDEX_VALUE=disable_index_value,
+        IS_FP8=is_fp8,
     )
 
     # topk extraction kernel

python/sglang/kernels/ops/attention/minimax_sparse/prefill/topk_sparse.py

@@ -6,7 +6,13 @@ import torch
 import triton
 import triton.language as tl
 
-from ..common.utils import check_sparse_kv_fp8, get_cu_seqblocks, robust_allocator
+from ..common.utils import (
+    check_sparse_kv_fp8,
+    get_cu_seqblocks,
+    robust_allocator,
+    sparse_out_dtype,
+    unit_scale,
+)
 
 
 @triton.heuristics(
@@ -66,6 +72,9 @@ def _gqa_share_sparse_fwd_kernel(
     num_q_loop,
     # sm_scale
     sm_scale,
+    # per-tensor KV dequant scales (1.0 when the cache is unit-scaled)
+    k_scale,
+    v_scale,
     # stride
     stride_qn,
     stride_qh,
@@ -201,8 +210,10 @@ def _gqa_share_sparse_fwd_kernel(
                 other=0.0,
             )
             if IS_FP8:
-                # fp8 main K cache is unit-scaled; widen to the Q compute dtype
+                # fp8 main K cache: widening cast with bf16/fp16 Q (unit-scaled
-                # before the tl.dot (compiled out when the cache is bf16).
+                # cache -> exact inverse dequant; k_scale covers calibrated
+                # caches), no-op with fp8 Q (fp8 attn-GEMM mode) so tl.dot runs
+                # fp8x8. Compiled out when the cache is bf16.
                 k = k.to(q.dtype)
             # compute qk
             qk = tl.zeros((BLOCK_SIZE_Q, BLOCK_SIZE_H, BLOCK_SIZE_K), dtype=tl.float32)
@@ -211,7 +222,7 @@ def _gqa_share_sparse_fwd_kernel(
             qk = tl.reshape(qk, BLOCK_SIZE_QH, BLOCK_SIZE_K)
             # [BLOCK_SIZE_QH, qk_head_dim] @ [qk_head_dim, BLOCK_SIZE_K]
             #   -> [BLOCK_SIZE_QH, BLOCK_SIZE_K]
-            qk += tl.dot(q, k) * sm_scale_log2e
+            qk += tl.dot(q, k) * (sm_scale_log2e * k_scale)
             # K boundary mask: positions beyond seq_len contribute -inf
             qk += tl.where(pos_mask[None, :], 0, float("-inf"))
             # compute m_ij and l_ij
@@ -231,11 +242,13 @@ def _gqa_share_sparse_fwd_kernel(
                 other=0.0,
             )
             if IS_FP8:
-                # Widen V so `p.to(v.dtype)` casts P to the compute dtype rather
+                # Cast V to the compute dtype: widening with bf16/fp16 Q (so
-                # than to fp8 (which would wreck attention-weight precision).
+                # `p.to(v.dtype)` keeps P in the compute dtype), no-op with fp8
+                # Q where P is quantized to e4m3 for the fp8 PV MMA — the same
+                # accuracy contract as fmha_sm100's fp8 kernel.
                 v = v.to(q.dtype)
             p = p.to(v.dtype)
-            acc_o += tl.dot(p, v)
+            acc_o += tl.dot(p, v) * v_scale
             # update statistics
             m_i = m_ij
             lse_i = m_ij + tl.log2(tl.exp2(lse_i - m_ij) + l_ij)
@@ -273,9 +286,14 @@ def flash_prefill_with_gqa_share_sparse(
     use_tma: bool = True,
     cu_seqblocks_q: Optional[torch.Tensor] = None,
     max_seqblock_q: Optional[int] = None,
+    q_scale: Optional[float] = None,
+    k_scale: Optional[float] = None,
+    v_scale: Optional[float] = None,
 ) -> torch.Tensor:
     triton.set_allocator(robust_allocator)
     is_fp8 = check_sparse_kv_fp8(q, k_cache, v_cache, label="prefill")
+    k_scale = unit_scale(k_scale)
+    v_scale = unit_scale(v_scale)
     assert block_size_q in {1, 2, 4, 8, 16, 32, 64}
     assert block_size_k in {16, 32, 64, 128}
     # shape
@@ -292,12 +310,17 @@ def flash_prefill_with_gqa_share_sparse(
     assert gqa_group_size * block_size_q <= 128
     if sm_scale is None:
         sm_scale = qk_head_dim**-0.5
+    # q_scale multiplies every Q-side logit (QK dot and sink), so it folds into
+    # sm_scale; k_scale must not touch the sink term and stays a kernel arg.
+    sm_scale = sm_scale * unit_scale(q_scale)
     if cu_seqblocks_q is None or max_seqblock_q is None:
         cu_seqblocks_q, max_seqblock_q, _, _, _, _ = get_cu_seqblocks(
             cu_seqlens, max_seqlen_q, block_size_q, block_size_k
         )
     # output tensor
-    o = torch.empty(total_q, num_q_heads, v_head_dim, device=q.device, dtype=q.dtype)
+    o = torch.empty(
+        total_q, num_q_heads, v_head_dim, device=q.device, dtype=sparse_out_dtype(q)
+    )
     # launch kernel
     num_q_loop = (
         max_seqblock_q // 131072 + 1
@@ -330,6 +353,8 @@ def flash_prefill_with_gqa_share_sparse(
         topk,
         num_q_loop,
         sm_scale,
+        k_scale,
+        v_scale,
         q.stride(0),
         q.stride(1),
         q.stride(2),

python/sglang/srt/arg_groups/overrides.py

@@ -512,13 +512,26 @@ def _minimax_m3_overrides(server_args: Any, hf_config: Any) -> dict:
             overrides["disable_custom_all_reduce"] = True
     elif is_sm100_supported():
         if server_args.is_attention_backend_not_set():
-            overrides["attention_backend"] = "fa4"
-        page_resolved = server_args.page_size
             if (
-            page_resolved is None
+                server_args.kv_cache_dtype == "fp8_e4m3"
-            and overrides.get("attention_backend", server_args.attention_backend)
+                and not envs.SGLANG_DISABLE_M3_FP8_ATTN_GEMM.get()
-            == "fa4"
             ):
+                # fp8 attention GEMMs activate whenever possible
+                # (m3_fp8_attn_gemm_enabled); only trtllm_mha serves the dense
+                # fp8-q path, so prefer it over fa4 for fp8 KV. The
+                # SGLANG_DISABLE_M3_FP8_ATTN_GEMM kill switch keeps the fa4
+                # default (pre-fp8 behavior).
+                overrides["attention_backend"] = "trtllm_mha"
+            else:
+                overrides["attention_backend"] = "fa4"
+        backend_resolved = overrides.get(
+            "attention_backend", server_args.attention_backend
+        )
+        page_resolved = server_args.page_size
+        # fa4 (fmha_sm100) and trtllm_mha both allow the page_size == 128
+        # sparse block MSA needs (trtllm_mha via trtllm-gen's dynamic
+        # tokens-per-page kernels).
+        if page_resolved is None and backend_resolved in ("fa4", "trtllm_mha"):
             overrides["page_size"] = 128
             page_resolved = 128
         if server_args.moe_runner_backend == "auto" and quant_resolved == "mxfp8":
@@ -550,6 +563,41 @@ def _minimax_m3_overrides(server_args: Any, hf_config: Any) -> dict:
             "(MSA is SM100-only; sparse attention runs on the Triton path)."
         )
 
+    # fp8 attention GEMMs have no opt-in flag: m3_fp8_attn_gemm_enabled
+    # (server_args.py) derives the mode from kv_cache_dtype (fp8_e4m3) +
+    # attention_backend (trtllm_mha) + SM100 at runtime. Surface the
+    # resolution here: warn on fp8_e5m2 (fmha_sm100's variant lookup would
+    # silently dispatch the e4m3 kernel, so e5m2 stays on the widening Triton
+    # path), log when the fp8 GEMM mode is active, and log when the
+    # SGLANG_DISABLE_M3_FP8_ATTN_GEMM kill switch suppresses it.
+    if server_args.kv_cache_dtype == "fp8_e5m2":
+        logger.warning(
+            "MiniMax-M3 with kv_cache_dtype fp8_e5m2: fp8 attention GEMMs stay "
+            "DISABLED (fmha_sm100's variant lookup would silently dispatch the "
+            "e4m3 kernel for e5m2); sparse attention runs on the widening "
+            "Triton path. Use --kv-cache-dtype fp8_e4m3 for fp8 attention GEMMs."
+        )
+    elif (
+        server_args.kv_cache_dtype == "fp8_e4m3"
+        and overrides.get("attention_backend", server_args.attention_backend)
+        == "trtllm_mha"
+        and is_sm100_supported()
+    ):
+        if envs.SGLANG_DISABLE_M3_FP8_ATTN_GEMM.get():
+            logger.info(
+                "MiniMax-M3 fp8 attention GEMMs DISABLED by "
+                "SGLANG_DISABLE_M3_FP8_ATTN_GEMM: bf16 indexer + widening "
+                "Triton sparse path, bf16 q; dense layers keep trtllm_mha's "
+                "fp8 KV cache."
+            )
+        else:
+            logger.info(
+                "MiniMax-M3 fp8 attention GEMMs active (kv_cache_dtype fp8_e4m3 + "
+                "trtllm_mha on SM100): fp8 main/index KV, fp8-cast q, fp8 "
+                "sparse/MSA kernels. Set SGLANG_DISABLE_M3_FP8_ATTN_GEMM=1 to "
+                "force the pre-fp8 numerics."
+            )
+
     moe_runner_resolved = overrides.get(
         "moe_runner_backend", server_args.moe_runner_backend
     )
@@ -1794,9 +1842,12 @@ def _mla_backend_page_constraints(view: Any) -> dict:
         or view.decode_attention_backend == "trtllm_mha"
         or view.prefill_attention_backend == "trtllm_mha"
     ):
-        if page_size not in [16, 32, 64]:
+        # 128 runs on trtllm-gen's dynamic tokens-per-page kernels (flashinfer
+        # >= 0.6.12), which require GQA and equal QK/V head dims — validated at
+        # TRTLLMHAAttnBackend init where the model config is known.
+        if page_size not in [16, 32, 64, 128]:
             logger.warning(
-                f"TensorRT-LLM MHA only supports page_size of 16, 32 or 64, changing page_size from {page_size} to 64."
+                f"TensorRT-LLM MHA only supports page_size of 16, 32, 64 or 128, changing page_size from {page_size} to 64."
             )
             page_size = 64
     if (

python/sglang/srt/environ.py

@@ -1153,6 +1153,10 @@ class Envs:
     SGLANG_OPT_USE_MINIMAX_DENSE_SPARSE_DECODE = EnvBool(False)
     SGLANG_DISABLE_MSA = EnvBool(False)
     SGLANG_OPT_USE_MSA_DECODE_UNDER_GRAPH = EnvBool(False)
+    # Kill switch for the derived fp8 attention-GEMM mode (m3_fp8_attn_gemm_enabled):
+    # forces the pre-fp8 behavior (bf16 indexer + widening sparse path, bf16 q)
+    # even when kv_cache_dtype fp8_e4m3 + trtllm_mha + SM100 would activate it.
+    SGLANG_DISABLE_M3_FP8_ATTN_GEMM = EnvBool(False)
 
     # MiniMax-M3 sparse decode indexer: single JIT radix-select kernel replaces the 2-stage split-K Triton topk.
     SGLANG_OPT_USE_MINIMAX_DECODE_TOPK_RADIX = EnvBool(True)

python/sglang/srt/layers/attention/minimax_sparse_backend.py

@@ -18,6 +18,7 @@ from sglang.srt.layers.attention.minimax_sparse_ops.minimax_sparse import (
 )
 from sglang.srt.mem_cache.memory_pool import MiniMaxSparseKVPool
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch
+from sglang.srt.server_args import m3_fp8_attn_gemm_enabled
 
 if TYPE_CHECKING:
     from sglang.srt.model_executor.model_runner import ModelRunner
@@ -25,12 +26,26 @@ if TYPE_CHECKING:
 logger = logging.getLogger(__name__)
 
 
+def _quant_q_fp8(q: torch.Tensor, q_scale: Optional[float]) -> torch.Tensor:
+    # Same convention as the KV pools: the fp8 tensor stores value/scale and
+    # the attention kernels multiply the logits back by the scale (None = unit).
+    if q_scale is not None:
+        q = q / q_scale
+    return q.to(torch.float8_e4m3fn)
+
+
 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)
+        self.fp8_attn_gemm = m3_fp8_attn_gemm_enabled(runner.server_args)
+        if self.fp8_attn_gemm:
+            assert self.kv_pool.main_pool.dtype == torch.float8_e4m3fn, (
+                "fp8 attn-GEMM mode requires an fp8_e4m3fn main KV pool, got "
+                f"{self.kv_pool.main_pool.dtype}"
+            )
 
         hf_config = runner.model_config.hf_config
         sparse_cfg = get_minimax_sparse_attention_config(hf_config)
@@ -72,19 +87,26 @@ class MiniMaxSparseAttnBackend(AttentionBackend):
             msa_available,
         )
 
-        # MSA (fmha_sm100) is bf16/fp16-only; an fp8 main KV cache must stay on the
+        # MSA (fmha_sm100) runs bf16, or uniform fp8_e4m3 under fp8 attn-GEMM mode
-        # Triton sparse path (it dequants fp8 on load).
+        # (which also casts q to fp8). An fp8 main KV cache WITHOUT the flag
+        # would pair a bf16 q with fp8 K/V — unsupported by fmha_sm100's
+        # uniform-dtype kernels — so it stays on the Triton sparse path (which
+        # dequants fp8 on load). e5m2 is never allowed into MSA (fmha_sm100's
+        # variant lookup would silently dispatch the e4m3 kernel).
         _main_kv_is_fp8 = self.kv_pool.main_pool.dtype in (
             torch.float8_e4m3fn,
             torch.float8_e5m2,
         )
+        _msa_fp8_ok = (
+            self.fp8_attn_gemm and self.kv_pool.main_pool.dtype == torch.float8_e4m3fn
+        )
         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
+            and (not _main_kv_is_fp8 or _msa_fp8_ok)
         )
         if (
             not self.use_msa
@@ -96,8 +118,8 @@ class MiniMaxSparseAttnBackend(AttentionBackend):
             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 "
+                "e.g. fa4 or trtllm_mha) to enable the faster MSA kernel; falling "
-                "Triton sparse path.",
+                "back to the Triton sparse path.",
                 self.kv_pool.page_size,
                 self.block_size_k,
             )
@@ -118,11 +140,10 @@ class MiniMaxSparseAttnBackend(AttentionBackend):
         self.use_dense_sparse_decode = (
             envs.SGLANG_OPT_USE_MINIMAX_DENSE_SPARSE_DECODE.get()
             and self.block_size_k % self.page_size == 0
+            # _dense_sparse_main_decode calls trtllm decode with a bf16 q and
+            # unit bmm scales — no fp8 handling yet (follow-up).
+            and not self.fp8_attn_gemm
         )
-        # 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,
@@ -158,8 +179,18 @@ class MiniMaxSparseAttnBackend(AttentionBackend):
             f"[MiniMaxSparse] Backend initialized "
             f"(score_type={self.score_type!r}, "
             f"main_attn={'MSA' if self.use_msa else 'triton'}, "
+            f"msa_decode={self._use_msa_decode}, "
+            f"msa_owns_decode={self._msa_owns_decode}, "
+            f"decode_cuda_graph={_decode_cuda_graph}, "
+            f"fp8_attn_gemm={self.fp8_attn_gemm}, "
             f"disable_value_layers={sorted(self.disable_value_layer_ids)})"
         )
+        if self.fp8_attn_gemm and self.use_msa:
+            logger.info(
+                "[MiniMaxSparse] fp8 MSA active: the first forward may "
+                "JIT-compile fmha_sm100 fp8 kernel variants (cold cache can "
+                "take minutes; compiles serialize across TP ranks)."
+            )
 
     def init_forward_metadata_out_graph(
         self, forward_batch: ForwardBatch, in_capture: bool = False
@@ -201,6 +232,7 @@ class MiniMaxSparseAttnBackend(AttentionBackend):
                 self.topk_blocks,
                 bs,
                 device=device,
+                is_fp8=self.fp8_attn_gemm,
             )
             kv_indices_buf = torch.zeros(
                 bs * self._msa_nb_max, dtype=torch.int32, device=device
@@ -288,6 +320,10 @@ class MiniMaxSparseAttnBackend(AttentionBackend):
                 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:
@@ -321,6 +357,12 @@ class MiniMaxSparseAttnBackend(AttentionBackend):
             q = q[:actual_num_tokens]
             idx_q = idx_q[:actual_num_tokens]
 
+        # fp8 attention GEMMs: quantize q/idx_q AFTER the KV store (which reads
+        # the bf16 k/v) and the DP trim.
+        if self.fp8_attn_gemm:
+            q = _quant_q_fp8(q, layer.q_scale_float)
+            idx_q = _quant_q_fp8(idx_q, layer.idx_q_scale_float)
+
         idx_o, o = minimax_sparse_prefill(
             q,
             k_cache,
@@ -346,6 +388,12 @@ class MiniMaxSparseAttnBackend(AttentionBackend):
             disable_index_value=disable_value,
             use_msa=self.use_msa,
             seqlens_cpu=forward_batch.extend_seq_lens_cpu,
+            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,
         )
 
         if actual_num_tokens < original_num_tokens:
@@ -424,6 +472,10 @@ class MiniMaxSparseAttnBackend(AttentionBackend):
             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:
@@ -458,6 +510,12 @@ class MiniMaxSparseAttnBackend(AttentionBackend):
                     "did not prepare the plan for this forward (gate mismatch)."
                 )
 
+        # fp8 attention GEMMs: quantize q/idx_q AFTER the KV store (which reads
+        # the bf16 k/v).
+        if self.fp8_attn_gemm:
+            q = _quant_q_fp8(q, layer.q_scale_float)
+            idx_q = _quant_q_fp8(idx_q, layer.idx_q_scale_float)
+
         idx_o, o = minimax_sparse_decode(
             q,
             None,
@@ -483,6 +541,12 @@ class MiniMaxSparseAttnBackend(AttentionBackend):
             use_msa=self._use_msa_decode,
             msa_kv_indices=msa_kv_indices,
             msa_plan=msa_plan,
+            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,
         )
         return (
             None if idx_o is None else idx_o.reshape(q.shape[0], -1).contiguous(),

python/sglang/srt/layers/attention/minimax_sparse_ops/minimax_sparse.py

@@ -67,6 +67,12 @@ def minimax_sparse_prefill(
     max_seqblock_q: Optional[int] = None,
     all_seqblock_q: Optional[int] = None,
     seqlens_cpu: Optional[List[int]] = None,
+    q_scale: Optional[float] = None,
+    k_scale: Optional[float] = None,
+    v_scale: Optional[float] = None,
+    idx_q_scale: Optional[float] = None,
+    idx_k_scale: Optional[float] = None,
+    idx_v_scale: Optional[float] = None,
 ):
     """Run MiniMax-M3 sparse prefill.
 
@@ -106,6 +112,9 @@ def minimax_sparse_prefill(
         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]
@@ -134,6 +143,9 @@ def minimax_sparse_prefill(
                 prefix_lens=prefix_lens,
                 block_size_k=block_size_k,
                 sm_scale=sm_scale,
+                q_scale=q_scale,
+                k_scale=k_scale,
+                v_scale=v_scale,
             )
         except MSAUnavailableError as err:
             _warn_msa_fallback(err)
@@ -154,6 +166,9 @@ def minimax_sparse_prefill(
                 sm_scale=sm_scale,
                 cu_seqblocks_q=cu_seqblocks_q,
                 max_seqblock_q=max_seqblock_q,
+                q_scale=q_scale,
+                k_scale=k_scale,
+                v_scale=v_scale,
             )
     else:
         o = flash_prefill_with_gqa_share_sparse(
@@ -173,6 +188,9 @@ def minimax_sparse_prefill(
             sm_scale=sm_scale,
             cu_seqblocks_q=cu_seqblocks_q,
             max_seqblock_q=max_seqblock_q,
+            q_scale=q_scale,
+            k_scale=k_scale,
+            v_scale=v_scale,
         )
     return idx_o, o
 
@@ -208,6 +226,12 @@ def minimax_sparse_decode(
         torch.Tensor
     ] = None,  # per-forward MSA page table (cached)
     msa_plan=None,  # per-forward MSA fmha_sm100 plan (cached)
+    q_scale: Optional[float] = None,
+    k_scale: Optional[float] = None,
+    v_scale: Optional[float] = None,
+    idx_q_scale: Optional[float] = None,
+    idx_k_scale: Optional[float] = None,
+    idx_v_scale: Optional[float] = 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
@@ -230,6 +254,9 @@ def minimax_sparse_decode(
         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]
@@ -262,6 +289,9 @@ def minimax_sparse_decode(
                     sm_scale=sm_scale,
                     kv_indices=msa_kv_indices,
                     plan=msa_plan,
+                    q_scale=q_scale,
+                    k_scale=k_scale,
+                    v_scale=v_scale,
                 )
             except MSAUnavailableError as err:
                 _warn_msa_fallback(err)
@@ -276,6 +306,9 @@ def minimax_sparse_decode(
                     block_size=block_size_k,
                     topk_idx=topk_idx,
                     sm_scale=sm_scale,
+                    q_scale=q_scale,
+                    k_scale=k_scale,
+                    v_scale=v_scale,
                 )
         else:
             o = flash_decode_with_gqa_share_sparse(
@@ -289,5 +322,8 @@ def minimax_sparse_decode(
                 block_size=block_size_k,
                 topk_idx=topk_idx,
                 sm_scale=sm_scale,
+                q_scale=q_scale,
+                k_scale=k_scale,
+                v_scale=v_scale,
             )
     return idx_o, o

python/sglang/srt/layers/attention/minimax_sparse_ops/msa.py

@@ -3,6 +3,14 @@
 # Replaces only step 3 of MiniMax sparse prefill/decode. The lightning indexer
 # (steps 1-2) is unchanged and still produces `topk_idx`.
 # NVIDIA Blackwell (SM100/sm_103) only; callers gate on `msa_available()`.
+#
+# Dtypes: bf16 end-to-end, or uniform fp8_e4m3fn Q/K/V under fp8 attn-GEMM mode
+# (output bf16). fmha_sm100 selects its kernel variant from q.dtype alone and
+# casts k/v pointers to the same element type, so mixed bf16-q/fp8-KV is NOT
+# possible on the cutlass path and e5m2 would silently dispatch the e4m3
+# kernel — `_check_msa_dtypes` enforces uniformity here. The fp8 kernel
+# quantizes the unnormalized softmax P to e4m3 before the PV MMA (same
+# contract as the Triton fp8 path).
 
 from __future__ import annotations
 
@@ -11,11 +19,32 @@ from typing import Optional
 
 import torch
 
+from sglang.kernels.ops.attention.minimax_sparse.common.utils import unit_scale
+
 
 class MSAUnavailableError(RuntimeError):
     """Raised when fmha_sm100 cannot serve the MiniMax MSA path."""
 
 
+def _check_msa_dtypes(q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor):
+    # Uniform dtype required in BOTH modes: fmha_sm100 keys its kernel variant
+    # on q.dtype alone and casts the k/v pointers to the same element type, so
+    # a mismatched cache would be silently reinterpreted.
+    if q.dtype == torch.bfloat16:
+        assert (
+            k_cache.dtype == torch.bfloat16
+        ), f"MSA bf16 requires a bf16 K cache, got {k_cache.dtype}"
+    elif q.dtype == torch.float8_e4m3fn:
+        # e5m2 is rejected here too: fmha_sm100's variant lookup falls back to
+        # the e4m3 kernel for unknown dtype codes.
+        assert (
+            k_cache.dtype == torch.float8_e4m3fn
+        ), f"MSA fp8 requires an fp8_e4m3fn K cache, got {k_cache.dtype}"
+    else:
+        raise AssertionError(f"MSA supports bf16 or fp8_e4m3fn Q, got {q.dtype}")
+    assert v_cache.dtype == k_cache.dtype
+
+
 @functools.lru_cache(maxsize=1)
 def _load_fmha_sm100():
     try:
@@ -101,12 +130,23 @@ def msa_sparse_prefill_main(
     prefix_lens: torch.Tensor,  # [batch]
     block_size_k: int,  # == page_size == 128 for M3
     sm_scale: Optional[float] = None,
+    q_scale: Optional[float] = None,
+    k_scale: Optional[float] = None,
+    v_scale: Optional[float] = None,
 ) -> torch.Tensor:
     """Drop-in for flash_prefill_with_gqa_share_sparse using MSA fmha_sm100.
 
-    Returns o [total_q, num_q_heads, head_dim].
+    Returns o [total_q, num_q_heads, head_dim] (bf16 for fp8 inputs).
+
+    Scale semantics (per-tensor, None = unit): attention runs on Q*q_scale,
+    K*k_scale, V*v_scale. The long-q cute path honors only sm_scale, so
+    q_scale*k_scale is folded into sm_scale (exact for softmax) and v_scale is
+    applied on the output; the short-q cutlass path gets the same folded values.
     """
     fmha_sm100, _ = _load_fmha_sm100()
+    _check_msa_dtypes(q, k_cache, v_cache)
+    is_fp8 = q.dtype == torch.float8_e4m3fn
+    v_scale = unit_scale(v_scale)
 
     max_slots, num_kv_heads, head_dim = k_cache.shape
     num_q_heads = q.shape[1]
@@ -116,6 +156,7 @@ def msa_sparse_prefill_main(
         raise ValueError(f"max_slots={max_slots} not divisible by page_size={P}")
     if sm_scale is None:
         sm_scale = head_dim**-0.5
+    sm_scale = sm_scale * unit_scale(q_scale) * unit_scale(k_scale)
 
     # Whole pool as MSA paged KV: [num_phys_pages, num_kv_heads, P, head_dim].
     n_phys_pages = max_slots // P
@@ -138,6 +179,7 @@ def msa_sparse_prefill_main(
         kv_block_num=topk,
         causal=True,
         qo_offset=prefix_lens.to(torch.int32),
+        use_fp8_kvcache=is_fp8,
     )
     o, _ = fmha_sm100(
         q,
@@ -148,6 +190,10 @@ def msa_sparse_prefill_main(
         kv_indices=kv_indices,
         kv_block_indexes=kv_block_indexes,
     )
+    # The cute (long-q) sparse prefill backend honors sm_scale only; apply the
+    # V dequant scale on the output (exact: softmax normalization excludes V).
+    if v_scale != 1.0:
+        o = o * v_scale
     return o
 
 
@@ -159,6 +205,7 @@ def build_msa_decode_meta(
     num_q_heads: int,
     block_size_k: int,
     topk: int,
+    is_fp8: bool = False,
 ):
     """Per-forward MSA decode metadata (page table + fmha plan), shared across layers.
 
@@ -187,25 +234,28 @@ def build_msa_decode_meta(
         kv_block_num=topk,
         causal=False,
         qo_offset=seq_lens_i32 - 1,  # decode query sits at the last cached position
+        use_fp8_kvcache=is_fp8,
     )
     return kv_indices, plan
 
 
 # ---------------------------------------------------------------------------
-# Eager-only MSA decode plan (NOT used under CUDA graph)
+# MSA decode plan (persistent per batch size; used by eager decode AND under
+# CUDA graph)
 #
-# WARNING: the fmha_sm100 sparse decode kernel is NOT cuda-graph-safe — captured
+# History: MSA decode under CUDA graph was disabled after silently wrong
-# and replayed it returns silently wrong results that compound across replays
+# results (~14% GSM8K on B200). Root cause (2026-07): the topk producers
-# (~14% GSM8K loss on B200). The backend routes decode to the cuda-graph-safe
+# emitted block ids in score order, violating fmha_sm100's strictly-ascending
-# Triton sparse path whenever decode runs under a CUDA graph (see
+# kv_block_indexes contract — its sorted-order early-exit then mis-masked the
-# MiniMaxSparseAttnBackend._use_msa_decode); this plan is reachable ONLY in eager
+# partial last block for any row with seq_len > topk*block_size. The producers
-# decode (no decode CUDA graph), where there is no capture/replay. Do NOT wire it
+# now sort ascending (minimax_decode_topk.cuh, _topk_index_merge_kernel,
-# back into a captured graph — that reintroduces the ~14% regression.
+# prefill _topk_index_kernel), and capture/replay of the full pipeline is
+# bit-exact vs eager (see tests/repro_msa_decode_degenerate.py).
 #
-# The build-once / replay-update structure below (refreshing the four length
+# The build-once / update-in-place structure below refreshes the four length
-# tensors ``{kv_segment_lens, kv_segment_offsets, kv_page_indptr, qo_offset}`` and
+# tensors ``{kv_segment_lens, kv_segment_offsets, kv_page_indptr, qo_offset}``
-# the page table in place) is a leftover from the abandoned capture-once attempt;
+# and the page table each forward; the captured graph reads the same tensor
-# it is kept only because eager decode reuses one per-forward plan across layers.
+# addresses on replay.
 # ---------------------------------------------------------------------------
 
 _MSA_CG_LEN_KEYS = (
@@ -242,6 +292,7 @@ def build_msa_decode_cg_plan(
     topk: int,
     batch_size: int,
     device: Optional[torch.device] = None,
+    is_fp8: bool = False,
 ):
     """Persistent fmha_sm100 decode plan for one batch size (CUDA-graph stable).
 
@@ -264,6 +315,7 @@ def build_msa_decode_cg_plan(
         causal=False,
         qo_offset=kv - 1,
         device=device,
+        use_fp8_kvcache=is_fp8,
     )
     _check_cg_plan_layout(plan)
     return plan
@@ -280,37 +332,53 @@ def update_msa_decode_cg_meta(
     num_q_heads: int,
     num_kv_heads: int,
 ):
-    """Refresh the persistent decode plan's length-dependent tensors + page table IN PLACE.
+    """Refresh the persistent decode plan's length-dependent tensors + page table
+    IN PLACE, entirely with device-side ops (no host<->device sync).
 
-    Host-side (calls fmha_sm100_plan and one ``.item()``); MUST run outside CUDA-graph
+    Runs once per decode forward from ``init_forward_metadata_out_graph``; the
-    capture — i.e. only from ``init_forward_metadata_out_graph``. The captured graph then
+    captured graph then reads the same plan-tensor and ``kv_indices_buf``
-    reads the same plan-tensor and ``kv_indices_buf`` addresses on replay.
+    addresses on replay. A device sync here stalls the overlap scheduler, so the
+    previous implementation — a throwaway ``fmha_sm100_plan`` build per step
+    (``.tolist()``/``.item()`` syncs + a plan-kernel launch) just to copy four
+    length tensors — is replaced by computing their contents directly, matching
+    ``_fmha_sm100_plan``'s sparse-decode (qo_len==1, causal=False) layout:
+
+      kv_segment_lens    = seq_lens
+      kv_segment_offsets = [0, cumsum(seq_lens)]
+      kv_page_indptr     = [0, cumsum(ceil(seq_lens / P))]
+      qo_offset          = broadcast max(seq_lens)   (causal=False planner quirk)
+
+    The worklist tensors stay untouched: the plan schedules from the constant
+    ``topk * P`` per request, never the real lengths (see build_msa_decode_cg_plan).
     """
     P = block_size_k
     B = seq_lens.shape[0]
+    if B == 0:  # idle batch: serving guards this, but keep the helper total
+        return
+    pd = plan[3]
     seq_lens_i32 = seq_lens.to(torch.int32)
-    # Fresh plan for the real lengths; copy only its four length-dependent tensors into the
+    pd["kv_segment_lens"].copy_(seq_lens_i32)
-    # persistent plan (same shapes — they depend on batch size, not length). The fresh
+    kv_off = pd["kv_segment_offsets"]
-    # worklist is identical to the persistent one (topk*P based) and is discarded.
+    kv_off[0].zero_()
-    # qo_offset is clamped: graph replay pads the batch with seq_len==0 slots
+    torch.cumsum(seq_lens_i32, 0, out=kv_off[1:])
-    # (masked via kv_segment_lens==0, but seq_len-1 would be -1).
+    n_pages = torch.div(seq_lens_i32 + (P - 1), P, rounding_mode="floor")
-    fresh = _run_fmha_sm100_plan(
+    indptr = pd["kv_page_indptr"]
-        torch.ones(B, dtype=torch.int32),
+    indptr[0].zero_()
-        seq_lens_i32,
+    torch.cumsum(n_pages, 0, out=indptr[1:])
-        num_q_heads,
+    pd["qo_offset"].copy_(seq_lens_i32.max().expand(B))
-        num_kv_heads=num_kv_heads,
+
-        page_size=P,
+    # Page table, sync-free: fill the WHOLE persistent buffer (fixed size, no
-        kv_block_num=topk,
+    # host-side total-page count). Packed slot -> (request, logical page) via
-        causal=False,
+    # searchsorted; slots beyond the live page count land on clamped reads and
-        qo_offset=(seq_lens_i32 - 1).clamp_min(0),
+    # are never dereferenced (the kernel walks kv_page_indptr ranges only).
-        device=seq_lens.device,
+    n = kv_indices_buf.numel()
-    )
+    idx = torch.arange(n, device=kv_indices_buf.device)
-    _check_cg_plan_layout(fresh)
+    ends = torch.cumsum(n_pages.to(torch.int64), 0)
-    pd, fd = plan[3], fresh[3]
+    req = torch.searchsorted(ends, idx, right=True).clamp_max_(B - 1)
-    for k in _MSA_CG_LEN_KEYS:
+    starts = ends - n_pages
-        pd[k].copy_(fd[k])
+    logical_first = ((idx - starts[req]) * P).clamp_(0, req_to_token.shape[1] - 1)
-    table = _build_page_table(req_to_token, slot_ids, seq_lens, P)
+    rows = slot_ids[req].to(torch.int64)
-    kv_indices_buf[: table.numel()].copy_(table)
+    kv_indices_buf.copy_((req_to_token[rows, logical_first] // P).to(torch.int32))
 
 
 def msa_sparse_decode_main(
@@ -327,18 +395,27 @@ def msa_sparse_decode_main(
         torch.Tensor
     ] = None,  # precomputed page table (per-forward cache)
     plan=None,  # precomputed fmha_sm100 plan (per-forward cache)
+    q_scale: Optional[float] = None,
+    k_scale: Optional[float] = None,
+    v_scale: Optional[float] = None,
 ) -> torch.Tensor:
     """Drop-in for flash_decode_with_gqa_share_sparse using MSA fmha_sm100.
 
     Each request is one decode query at absolute position seq_len-1 attending to its
-    cached KV through the topk selected 128-blocks. Returns o [batch, num_q_heads, head_dim].
+    cached KV through the topk selected 128-blocks. Returns o [batch, num_q_heads,
+    head_dim] (bf16 for fp8 inputs).
 
     ``kv_indices`` / ``plan`` are shared across all sparse layers of a forward; the serving
     backend builds them once via ``build_msa_decode_cg_plan`` + ``update_msa_decode_cg_meta``
     (eager decode only) and passes them in. When omitted (only the standalone parity
     harnesses) they are built here via ``build_msa_decode_meta``.
+
+    Scales (None = unit) are passed natively: the cutlass decode path folds
+    q_scale*k_scale into the softmax scale and applies v_scale on the output
+    in-kernel.
     """
     fmha_sm100, _ = _load_fmha_sm100()
+    _check_msa_dtypes(q, k_cache, v_cache)
 
     max_slots, num_kv_heads, head_dim = k_cache.shape
     H = q.shape[1]
@@ -355,7 +432,14 @@ def msa_sparse_decode_main(
 
     if kv_indices is None or plan is None:
         kv_indices, plan = build_msa_decode_meta(
-            k_cache, req_to_token, slot_ids, seq_lens, H, P, topk
+            k_cache,
+            req_to_token,
+            slot_ids,
+            seq_lens,
+            H,
+            P,
+            topk,
+            is_fp8=q.dtype == torch.float8_e4m3fn,
         )
     kv_block_indexes = topk_idx.permute(1, 0, 2).contiguous().to(torch.int32)
 
@@ -365,6 +449,9 @@ def msa_sparse_decode_main(
         v_paged,
         plan,
         sm_scale=sm_scale,
+        q_scale=unit_scale(q_scale),
+        k_scale=unit_scale(k_scale),
+        v_scale=unit_scale(v_scale),
         kv_indices=kv_indices,
         kv_block_indexes=kv_block_indexes,
     )

python/sglang/srt/layers/attention/minimax_sparse_ops/tests/test_fp8_attn_gemm.py

@@ -0,0 +1,455 @@
+"""Unit tests for fp8 (fp8 attn-GEMM mode) support in the M3 sparse Triton kernels.
+
+Strategy: quantize random bf16 tensors to fp8_e4m3fn, then compare the fp8
+kernel run against the SAME kernel run in bf16 on the *dequantized* tensors.
+Both runs see numerically identical Q/K values, so the QK GEMMs match closely
+and top-k selection is stable; the only intended divergence is the fp8 PV MMA
+(P quantized to e4m3), which the tolerances cover. This isolates kernel
+arithmetic from quantization error.
+
+Covers: step-3 decode/prefill (gqa-share sparse), step-1 decode/prefill
+indexer, non-unit k_scale/v_scale semantics, and bf16-path regression.
+"""
+
+import pytest
+import torch
+
+from sglang.srt.layers.attention.minimax_sparse_ops.decode.flash_with_topk_idx import (
+    flash_decode_with_topk_idx,
+)
+from sglang.srt.layers.attention.minimax_sparse_ops.decode.topk_sparse import (
+    flash_decode_with_gqa_share_sparse,
+)
+from sglang.srt.layers.attention.minimax_sparse_ops.prefill.flash_with_topk_idx import (
+    flash_prefill_with_topk_index,
+)
+from sglang.srt.layers.attention.minimax_sparse_ops.prefill.topk_sparse import (
+    flash_prefill_with_gqa_share_sparse,
+)
+
+DEVICE = "cuda"
+FP8 = torch.float8_e4m3fn
+
+pytestmark = pytest.mark.skipif(
+    not torch.cuda.is_available(), reason="requires CUDA (Triton fp8 kernels)"
+)
+# fp8 PV (P quantized to e4m3, ~3-bit mantissa on [0,1] weights) dominates the
+# fp8-vs-dequantized-ref error; QK matches to fp32-accumulation noise.
+FP8_ATOL = 6e-2
+FP8_RTOL = 6e-2
+# widening mode (bf16 Q, fp8 KV) computes on exactly the dequantized values.
+WIDEN_ATOL = 1e-3
+WIDEN_RTOL = 1e-3
+
+
+def qdq(x: torch.Tensor):
+    """Quantize to e4m3 and return (fp8, dequantized-bf16) views of it."""
+    x8 = x.to(FP8)
+    return x8, x8.to(torch.bfloat16)
+
+
+def build_decode_inputs(
+    batch_size=4,
+    num_q_heads=8,
+    num_kv_heads=1,
+    head_dim=128,
+    block_size=128,
+    topk=8,
+    seq_lens_list=(513, 1023, 257, 769),
+):
+    seq_lens_list = list(seq_lens_list)[:batch_size]
+    max_kv_len = max(seq_lens_list)
+    max_slots = batch_size * max_kv_len
+    q = torch.randn(
+        batch_size, num_q_heads, head_dim, dtype=torch.bfloat16, device=DEVICE
+    )
+    k = torch.randn(
+        max_slots, num_kv_heads, head_dim, dtype=torch.bfloat16, device=DEVICE
+    )
+    v = torch.randn(
+        max_slots, num_kv_heads, head_dim, dtype=torch.bfloat16, device=DEVICE
+    )
+    req_to_token = torch.zeros(batch_size, max_kv_len, dtype=torch.int32, device=DEVICE)
+    slot_ids = torch.arange(batch_size, dtype=torch.int64, device=DEVICE)
+    seq_lens = torch.tensor(seq_lens_list, dtype=torch.int32, device=DEVICE)
+    for i in range(batch_size):
+        base = i * max_kv_len
+        req_to_token[i] = (torch.randperm(max_kv_len, device=DEVICE) + base).to(
+            torch.int32
+        )
+    topk_idx = torch.full(
+        (num_kv_heads, batch_size, topk), -1, dtype=torch.int32, device=DEVICE
+    )
+    for kh in range(num_kv_heads):
+        for b in range(batch_size):
+            nb = (seq_lens_list[b] + block_size - 1) // block_size
+            ak = min(topk, nb)
+            # sorted ascending, matching the production topk contract
+            sel = torch.randperm(nb, device=DEVICE)[:ak].sort().values
+            topk_idx[kh, b, :ak] = sel.to(torch.int32)
+    return q, k, v, req_to_token, seq_lens, slot_ids, topk_idx
+
+
+def build_prefill_inputs(
+    batch_size=2,
+    num_q_heads=8,
+    num_kv_heads=1,
+    head_dim=128,
+    seq_lens_list=(513, 769),
+    prefix_lens_list=(0, 257),
+):
+    seq_lens_list = list(seq_lens_list)[:batch_size]
+    prefix_lens_list = list(prefix_lens_list)[:batch_size]
+    q_lens = [s - p for s, p in zip(seq_lens_list, prefix_lens_list)]
+    total_q = sum(q_lens)
+    max_kv_len = max(seq_lens_list)
+    max_slots = batch_size * max_kv_len
+    q = torch.randn(total_q, num_q_heads, head_dim, dtype=torch.bfloat16, device=DEVICE)
+    k = torch.randn(
+        max_slots, num_kv_heads, head_dim, dtype=torch.bfloat16, device=DEVICE
+    )
+    v = torch.randn(
+        max_slots, num_kv_heads, head_dim, dtype=torch.bfloat16, device=DEVICE
+    )
+    req_to_token = torch.zeros(batch_size, max_kv_len, dtype=torch.int32, device=DEVICE)
+    slot_ids = torch.arange(batch_size, dtype=torch.int64, device=DEVICE)
+    for i in range(batch_size):
+        base = i * max_kv_len
+        req_to_token[i] = (torch.randperm(max_kv_len, device=DEVICE) + base).to(
+            torch.int32
+        )
+    cu_seqlens = torch.zeros(batch_size + 1, dtype=torch.int32, device=DEVICE)
+    cu_seqlens[1:] = torch.tensor(q_lens, device=DEVICE).cumsum(0)
+    seq_lens = torch.tensor(seq_lens_list, dtype=torch.int32, device=DEVICE)
+    prefix_lens = torch.tensor(prefix_lens_list, dtype=torch.int32, device=DEVICE)
+    return (
+        q,
+        k,
+        v,
+        req_to_token,
+        slot_ids,
+        cu_seqlens,
+        seq_lens,
+        prefix_lens,
+        max(q_lens),
+        max(seq_lens_list),
+    )
+
+
+def run_step3_decode(q, k, v, req_to_token, seq_lens, slot_ids, topk_idx, **kw):
+    return flash_decode_with_gqa_share_sparse(
+        q, None, k, v, req_to_token, seq_lens, slot_ids, 128, topk_idx, **kw
+    )
+
+
+# ---------------------------------------------------------------------------
+# step-3 decode (gqa-share sparse)
+# ---------------------------------------------------------------------------
+
+
+def test_step3_decode_all_fp8_vs_dequant_ref():
+    torch.manual_seed(0)
+    q, k, v, r2t, seq_lens, sids, tidx = build_decode_inputs()
+    q8, qr = qdq(q)
+    k8, kr = qdq(k)
+    v8, vr = qdq(v)
+    out8 = run_step3_decode(q8, k8, v8, r2t, seq_lens, sids, tidx)
+    ref = run_step3_decode(qr, kr, vr, r2t, seq_lens, sids, tidx)
+    assert out8.dtype == torch.bfloat16
+    torch.testing.assert_close(out8.float(), ref.float(), atol=FP8_ATOL, rtol=FP8_RTOL)
+
+
+def test_step3_decode_widening_mode():
+    torch.manual_seed(1)
+    q, k, v, r2t, seq_lens, sids, tidx = build_decode_inputs()
+    k8, kr = qdq(k)
+    v8, vr = qdq(v)
+    out = run_step3_decode(q, k8, v8, r2t, seq_lens, sids, tidx)
+    ref = run_step3_decode(q, kr, vr, r2t, seq_lens, sids, tidx)
+    torch.testing.assert_close(
+        out.float(), ref.float(), atol=WIDEN_ATOL, rtol=WIDEN_RTOL
+    )
+
+
+def test_step3_decode_scales():
+    torch.manual_seed(2)
+    q, k, v, r2t, seq_lens, sids, tidx = build_decode_inputs()
+    q8, qr = qdq(q)
+    k8, kr = qdq(k)
+    v8, vr = qdq(v)
+    q_scale, k_scale, v_scale = 1.5, 0.5, 2.0
+    out8 = run_step3_decode(
+        q8,
+        k8,
+        v8,
+        r2t,
+        seq_lens,
+        sids,
+        tidx,
+        q_scale=q_scale,
+        k_scale=k_scale,
+        v_scale=v_scale,
+    )
+    # reference: bf16 kernel on pre-scaled dequantized Q/K/V (scale semantics:
+    # the tensor stores value/scale; attention runs on value = stored * scale)
+    ref = run_step3_decode(
+        (qr * q_scale).to(torch.bfloat16),
+        (kr * k_scale).to(torch.bfloat16),
+        (vr * v_scale).to(torch.bfloat16),
+        r2t,
+        seq_lens,
+        sids,
+        tidx,
+    )
+    torch.testing.assert_close(
+        out8.float(), ref.float(), atol=FP8_ATOL * v_scale, rtol=FP8_RTOL * v_scale
+    )
+
+
+def test_step3_decode_bf16_regression():
+    torch.manual_seed(3)
+    q, k, v, r2t, seq_lens, sids, tidx = build_decode_inputs()
+    out = run_step3_decode(q, k, v, r2t, seq_lens, sids, tidx)
+    out_scaled = run_step3_decode(
+        q, k, v, r2t, seq_lens, sids, tidx, q_scale=1.0, k_scale=1.0, v_scale=1.0
+    )
+    assert out.dtype == q.dtype
+    torch.testing.assert_close(out, out_scaled, atol=0, rtol=0)
+
+
+# ---------------------------------------------------------------------------
+# step-3 prefill (gqa-share sparse)
+# ---------------------------------------------------------------------------
+
+
+def run_step3_prefill(q, k, v, r2t, sids, tidx, cu, seq_lens, prefix, max_q, **kw):
+    return flash_prefill_with_gqa_share_sparse(
+        q=q,
+        k_cache=k,
+        v_cache=v,
+        sink=None,
+        req_to_token=r2t,
+        slot_ids=sids,
+        topk_idx=tidx,
+        block_size_q=1,
+        block_size_k=128,
+        cu_seqlens=cu,
+        seq_lens=seq_lens,
+        prefix_lens=prefix,
+        max_seqlen_q=max_q,
+        **kw,
+    )
+
+
+def _prefill_topk_idx(cu_seqlens, seq_lens, prefix_lens, num_kv_heads, topk, block):
+    # per-token (block_size_q=1) causal topk: for query at absolute position p,
+    # any blocks with start <= p, sorted ascending, -1 padded.
+    total_q = int(cu_seqlens[-1].item())
+    tidx = torch.full(
+        (num_kv_heads, total_q, topk), -1, dtype=torch.int32, device=DEVICE
+    )
+    row = 0
+    for b in range(len(seq_lens)):
+        q_len = int(cu_seqlens[b + 1] - cu_seqlens[b])
+        prefix = int(prefix_lens[b])
+        for j in range(q_len):
+            nb = (prefix + j) // block + 1  # blocks fully/partially before pos
+            ak = min(topk, nb)
+            sel = torch.randperm(nb, device=DEVICE)[:ak].sort().values
+            for kh in range(num_kv_heads):
+                tidx[kh, row, :ak] = sel.to(torch.int32)
+            row += 1
+    return tidx
+
+
+def test_step3_prefill_all_fp8_vs_dequant_ref():
+    torch.manual_seed(4)
+    q, k, v, r2t, sids, cu, seq_lens, prefix, max_q, _ = build_prefill_inputs()
+    tidx = _prefill_topk_idx(cu.cpu(), seq_lens.cpu(), prefix.cpu(), 1, 8, 128)
+    q8, qr = qdq(q)
+    k8, kr = qdq(k)
+    v8, vr = qdq(v)
+    out8 = run_step3_prefill(q8, k8, v8, r2t, sids, tidx, cu, seq_lens, prefix, max_q)
+    ref = run_step3_prefill(qr, kr, vr, r2t, sids, tidx, cu, seq_lens, prefix, max_q)
+    assert out8.dtype == torch.bfloat16
+    torch.testing.assert_close(out8.float(), ref.float(), atol=FP8_ATOL, rtol=FP8_RTOL)
+
+
+def test_step3_prefill_scales():
+    torch.manual_seed(5)
+    q, k, v, r2t, sids, cu, seq_lens, prefix, max_q, _ = build_prefill_inputs()
+    tidx = _prefill_topk_idx(cu.cpu(), seq_lens.cpu(), prefix.cpu(), 1, 8, 128)
+    q8, qr = qdq(q)
+    k8, kr = qdq(k)
+    v8, vr = qdq(v)
+    q_scale, k_scale, v_scale = 1.5, 0.5, 2.0
+    out8 = run_step3_prefill(
+        q8,
+        k8,
+        v8,
+        r2t,
+        sids,
+        tidx,
+        cu,
+        seq_lens,
+        prefix,
+        max_q,
+        q_scale=q_scale,
+        k_scale=k_scale,
+        v_scale=v_scale,
+    )
+    ref = run_step3_prefill(
+        (qr * q_scale).to(torch.bfloat16),
+        (kr * k_scale).to(torch.bfloat16),
+        (vr * v_scale).to(torch.bfloat16),
+        r2t,
+        sids,
+        tidx,
+        cu,
+        seq_lens,
+        prefix,
+        max_q,
+    )
+    # v_scale multiplies the fp8 kernel's quantized-P PV output, amplifying
+    # the P-quantization error by the same factor vs the pre-scaled bf16 ref.
+    torch.testing.assert_close(
+        out8.float(), ref.float(), atol=FP8_ATOL * v_scale, rtol=FP8_RTOL * v_scale
+    )
+
+
+# ---------------------------------------------------------------------------
+# step-1 decode indexer
+# ---------------------------------------------------------------------------
+
+
+def run_indexer_decode(q, k, v, r2t, seq_lens, sids, **kw):
+    return flash_decode_with_topk_idx(
+        q=q,
+        sink=None,
+        k_cache=k,
+        v_cache=v,
+        req_to_token=r2t,
+        seq_lens=seq_lens,
+        slot_ids=sids,
+        max_seqlen=int(seq_lens.max().item()),
+        block_size=128,
+        topk=4,
+        init_blocks=1,
+        local_blocks=1,
+        **kw,
+    )
+
+
+def _topk_overlap(a: torch.Tensor, b: torch.Tensor) -> float:
+    """Mean per-row overlap of the selected (non-negative) block-id sets."""
+    total, hit = 0, 0
+    af, bf = a.reshape(-1, a.shape[-1]), b.reshape(-1, b.shape[-1])
+    for i in range(af.shape[0]):
+        sa = set(af[i][af[i] >= 0].tolist())
+        sb = set(bf[i][bf[i] >= 0].tolist())
+        if not sa and not sb:
+            continue
+        total += max(len(sa), len(sb))
+        hit += len(sa & sb)
+    return hit / max(total, 1)
+
+
+def test_indexer_decode_all_fp8_vs_dequant_ref():
+    torch.manual_seed(6)
+    q, k, v, r2t, seq_lens, sids, _ = build_decode_inputs(
+        num_q_heads=1, num_kv_heads=1, seq_lens_list=(1023, 769, 513, 257)
+    )
+    q8, qr = qdq(q)
+    k8, kr = qdq(k)
+    v8, vr = qdq(v)
+    o8, tidx8, _ = run_indexer_decode(q8, k8, v8, r2t, seq_lens, sids)
+    oref, tidxref, _ = run_indexer_decode(qr, kr, vr, r2t, seq_lens, sids)
+    assert o8.dtype == torch.bfloat16
+    # o is the full (non-sparse) indexer attention output — independent of the
+    # topk side-channel — so it must track the dequantized reference.
+    torch.testing.assert_close(o8.float(), oref.float(), atol=FP8_ATOL, rtol=FP8_RTOL)
+    # topk selection runs on QK scores that only differ by fp32-accumulation
+    # noise; require near-total agreement (ties may flip an occasional block).
+    assert _topk_overlap(tidx8, tidxref) >= 0.9
+
+
+def test_indexer_decode_score_only_fp8():
+    torch.manual_seed(7)
+    q, k, _, r2t, seq_lens, sids, _ = build_decode_inputs(
+        batch_size=2, num_q_heads=1, num_kv_heads=1, seq_lens_list=(1023, 769)
+    )
+    q8, qr = qdq(q)
+    k8, kr = qdq(k)
+    o8, tidx8, _ = run_indexer_decode(
+        q8, k8, None, r2t, seq_lens, sids, disable_index_value=True
+    )
+    oref, tidxref, _ = run_indexer_decode(
+        qr, kr, None, r2t, seq_lens, sids, disable_index_value=True
+    )
+    assert o8 is None and oref is None
+    assert _topk_overlap(tidx8, tidxref) >= 0.9
+
+
+# ---------------------------------------------------------------------------
+# step-1 prefill indexer
+# ---------------------------------------------------------------------------
+
+
+def run_indexer_prefill(q, k, v, r2t, sids, cu, seq_lens, prefix, max_q, max_k, **kw):
+    return flash_prefill_with_topk_index(
+        q=q,
+        k_cache=k,
+        v_cache=v,
+        sink=None,
+        req_to_token=r2t,
+        slot_ids=sids,
+        cu_seqlens=cu,
+        seq_lens=seq_lens,
+        prefix_lens=prefix,
+        max_seqlen_q=max_q,
+        max_seqlen_k=max_k,
+        block_size_q=1,
+        block_size_k=128,
+        topk=4,
+        init_blocks=1,
+        local_blocks=1,
+        **kw,
+    )
+
+
+def test_indexer_prefill_all_fp8_vs_dequant_ref():
+    torch.manual_seed(8)
+    q, k, v, r2t, sids, cu, seq_lens, prefix, max_q, max_k = build_prefill_inputs(
+        num_q_heads=1, num_kv_heads=1
+    )
+    q8, qr = qdq(q)
+    k8, kr = qdq(k)
+    v8, vr = qdq(v)
+    o8, tidx8 = run_indexer_prefill(
+        q8, k8, v8, r2t, sids, cu, seq_lens, prefix, max_q, max_k
+    )
+    oref, tidxref = run_indexer_prefill(
+        qr, kr, vr, r2t, sids, cu, seq_lens, prefix, max_q, max_k
+    )
+    assert o8.dtype == torch.bfloat16
+    torch.testing.assert_close(o8.float(), oref.float(), atol=FP8_ATOL, rtol=FP8_RTOL)
+    assert _topk_overlap(tidx8, tidxref) >= 0.9
+
+
+def test_dtype_contract_rejects_e5m2_q():
+    q = torch.randn(2, 8, 128, dtype=torch.bfloat16, device=DEVICE).to(
+        torch.float8_e5m2
+    )
+    k = torch.randn(256, 1, 128, dtype=torch.bfloat16, device=DEVICE).to(FP8)
+    from sglang.kernels.ops.attention.minimax_sparse.common.utils import (
+        check_sparse_kv_fp8,
+    )
+
+    with pytest.raises(AssertionError):
+        check_sparse_kv_fp8(q, k, None, label="test")
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main([__file__, "-v"]))

python/sglang/srt/layers/attention/minimax_sparse_ops/tests/test_msa_fp8_parity.py

@@ -0,0 +1,286 @@
+"""Parity tests for MSA (fmha_sm100) all-fp8 sparse attention (fp8 attn-GEMM mode).
+
+No upstream fp8 test exists for fmha_sm100's cutlass sparse-decode path, so
+this is the reference check: MSA fp8 vs the Triton fp8 sparse kernels on the
+same quantized tensors, plus fp8-vs-bf16 error bounds and CUDA-graph
+capture/replay bit-exactness of the fp8 decode.
+
+Both fp8 paths quantize the unnormalized softmax P to e4m3 before the PV MMA,
+but their QK/accumulation orders differ, so MSA-fp8 vs Triton-fp8 tolerances
+cover two independent P-quantization errors (~1e-1 worst-case elementwise).
+
+Requires SM100 + fmha_sm100 (first run JIT-compiles the fp8 variants).
+Run: pytest python/sglang/srt/layers/attention/minimax_sparse_ops/tests/test_msa_fp8_parity.py -v
+"""
+
+import pytest
+import torch
+
+from sglang.srt.layers.attention.minimax_sparse_ops.decode.topk_sparse import (
+    flash_decode_with_gqa_share_sparse,
+)
+from sglang.srt.layers.attention.minimax_sparse_ops.msa import (
+    build_msa_decode_cg_plan,
+    msa_available,
+    msa_sparse_decode_main,
+    msa_sparse_prefill_main,
+    update_msa_decode_cg_meta,
+)
+from sglang.srt.layers.attention.minimax_sparse_ops.prefill.topk_sparse import (
+    flash_prefill_with_gqa_share_sparse,
+)
+
+DEVICE = "cuda"
+FP8 = torch.float8_e4m3fn
+P = 128  # sparse block == page size
+# two independent e4m3-P quantizations (MSA + Triton)
+X_ATOL = 1e-1
+X_RTOL = 1e-1
+
+pytestmark = pytest.mark.skipif(
+    not torch.cuda.is_available() or not msa_available(),
+    reason="requires SM100 + fmha_sm100",
+)
+
+
+def qdq(x):
+    x8 = x.to(FP8)
+    return x8, x8.to(torch.bfloat16)
+
+
+def build_paged_inputs(seq_lens_list, num_q_heads=16, num_kv_heads=1, head_dim=128):
+    """Page-aligned pool: each logical 128-token page maps to one physical page
+    (contiguous 128 slots), as MSA's page-table builder requires."""
+    batch = len(seq_lens_list)
+    pages_per_req = [(s + P - 1) // P for s in seq_lens_list]
+    max_pages = max(pages_per_req)
+    total_pages = batch * max_pages
+    max_slots = total_pages * P
+    page_perm = torch.randperm(total_pages, device=DEVICE)
+    req_to_token = torch.zeros(batch, max_pages * P, dtype=torch.int32, device=DEVICE)
+    for b in range(batch):
+        for p in range(max_pages):
+            phys = page_perm[b * max_pages + p]
+            req_to_token[b, p * P : (p + 1) * P] = phys * P + torch.arange(
+                P, device=DEVICE
+            )
+    k = torch.randn(
+        max_slots, num_kv_heads, head_dim, dtype=torch.bfloat16, device=DEVICE
+    )
+    v = torch.randn(
+        max_slots, num_kv_heads, head_dim, dtype=torch.bfloat16, device=DEVICE
+    )
+    slot_ids = torch.arange(batch, dtype=torch.int64, device=DEVICE)
+    seq_lens = torch.tensor(seq_lens_list, dtype=torch.int32, device=DEVICE)
+    return k, v, req_to_token, slot_ids, seq_lens
+
+
+def make_topk_idx(seq_lens_list, num_kv_heads, rows_per_req, topk):
+    """Sorted-ascending causal-safe topk over each request's block count.
+    rows_per_req[b] = number of query rows for request b (1 for decode)."""
+    total_rows = sum(rows_per_req)
+    tidx = torch.full(
+        (num_kv_heads, total_rows, topk), -1, dtype=torch.int32, device=DEVICE
+    )
+    row = 0
+    for b, s in enumerate(seq_lens_list):
+        nb = (s + P - 1) // P
+        for _ in range(rows_per_req[b]):
+            ak = min(topk, nb)
+            sel = torch.randperm(nb, device=DEVICE)[:ak].sort().values
+            for kh in range(num_kv_heads):
+                tidx[kh, row, :ak] = sel.to(torch.int32)
+            row += 1
+    return tidx
+
+
+# ---------------------------------------------------------------------------
+# decode
+# ---------------------------------------------------------------------------
+
+
+def _decode_setup(seq_lens_list=(1023, 769, 513, 130), topk=4):
+    """topk*P = 512 < 1023/769: exercises the degenerate partial-last-block
+    masking (the historical unsorted-topk failure shape)."""
+    torch.manual_seed(0)
+    batch = len(seq_lens_list)
+    k, v, r2t, sids, seq_lens = build_paged_inputs(seq_lens_list)
+    q = torch.randn(batch, 16, 128, dtype=torch.bfloat16, device=DEVICE)
+    tidx = make_topk_idx(seq_lens_list, 1, [1] * batch, topk)
+    return q, k, v, r2t, sids, seq_lens, tidx
+
+
+def test_msa_fp8_decode_vs_triton_fp8():
+    q, k, v, r2t, sids, seq_lens, tidx = _decode_setup()
+    q8, qr = qdq(q)
+    k8, kr = qdq(k)
+    v8, vr = qdq(v)
+    o_msa = msa_sparse_decode_main(
+        q8, k8, v8, tidx, r2t, sids, seq_lens, block_size_k=P
+    )
+    o_triton = flash_decode_with_gqa_share_sparse(
+        q8, None, k8, v8, r2t, seq_lens, sids, P, tidx
+    )
+    assert o_msa.dtype == torch.bfloat16
+    torch.testing.assert_close(
+        o_msa.float(), o_triton.float(), atol=X_ATOL, rtol=X_RTOL
+    )
+    # bf16 MSA on the dequantized tensors bounds the pure fp8-kernel error
+    o_bf16 = msa_sparse_decode_main(
+        qr, kr, vr, tidx, r2t, sids, seq_lens, block_size_k=P
+    )
+    err = (o_msa.float() - o_bf16.float()).abs().mean()
+    ref = o_bf16.float().abs().mean()
+    assert err / ref < 0.06, f"mean rel err {err/ref:.4f} too high vs bf16 MSA"
+
+
+def test_msa_fp8_decode_scales():
+    q, k, v, r2t, sids, seq_lens, tidx = _decode_setup()
+    q8, qr = qdq(q)
+    k8, kr = qdq(k)
+    v8, vr = qdq(v)
+    q_scale, k_scale, v_scale = 1.5, 0.5, 2.0
+    o = msa_sparse_decode_main(
+        q8,
+        k8,
+        v8,
+        tidx,
+        r2t,
+        sids,
+        seq_lens,
+        block_size_k=P,
+        q_scale=q_scale,
+        k_scale=k_scale,
+        v_scale=v_scale,
+    )
+    # reference: bf16 MSA on pre-scaled dequantized Q/K/V
+    ref = msa_sparse_decode_main(
+        (qr * q_scale).to(torch.bfloat16),
+        (kr * k_scale).to(torch.bfloat16),
+        (vr * v_scale).to(torch.bfloat16),
+        tidx,
+        r2t,
+        sids,
+        seq_lens,
+        block_size_k=P,
+    )
+    torch.testing.assert_close(
+        o.float(), ref.float(), atol=X_ATOL * v_scale, rtol=X_RTOL * v_scale
+    )
+
+
+def test_msa_fp8_decode_capture_replay_bitexact():
+    q, k, v, r2t, sids, seq_lens, tidx = _decode_setup()
+    q8, _ = qdq(q)
+    k8, _ = qdq(k)
+    v8, _ = qdq(v)
+    bs = q8.shape[0]
+    nb_max = r2t.shape[1] // P
+    plan = build_msa_decode_cg_plan(
+        16, 1, P, tidx.shape[-1], bs, device=q8.device, is_fp8=True
+    )
+    kv_indices = torch.zeros(bs * nb_max, dtype=torch.int32, device=DEVICE)
+    update_msa_decode_cg_meta(
+        plan, kv_indices, r2t, sids, seq_lens, P, tidx.shape[-1], 16, 1
+    )
+
+    def run():
+        return msa_sparse_decode_main(
+            q8,
+            k8,
+            v8,
+            tidx,
+            r2t,
+            sids,
+            seq_lens,
+            block_size_k=P,
+            kv_indices=kv_indices,
+            plan=plan,
+        )
+
+    # eager warmups (also pays the fp8 JIT) on a side stream, as capture does
+    s = torch.cuda.Stream()
+    s.wait_stream(torch.cuda.current_stream())
+    with torch.cuda.stream(s):
+        for _ in range(2):
+            o_eager = run()
+    torch.cuda.current_stream().wait_stream(s)
+
+    g = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(g):
+        o_graph = run()
+    g.replay()
+    torch.cuda.synchronize()
+    assert torch.equal(o_graph, o_eager), "fp8 MSA decode capture/replay not bit-exact"
+
+
+# ---------------------------------------------------------------------------
+# prefill (cutlass short-q and cute long-q branches)
+# ---------------------------------------------------------------------------
+
+
+def _prefill_setup(seq_lens_list, prefix_lens_list, topk=4):
+    torch.manual_seed(1)
+    k, v, r2t, sids, seq_lens = build_paged_inputs(seq_lens_list)
+    q_lens = [s - p for s, p in zip(seq_lens_list, prefix_lens_list)]
+    total_q = sum(q_lens)
+    q = torch.randn(total_q, 16, 128, dtype=torch.bfloat16, device=DEVICE)
+    cu = torch.zeros(len(q_lens) + 1, dtype=torch.int32, device=DEVICE)
+    cu[1:] = torch.tensor(q_lens, device=DEVICE).cumsum(0)
+    prefix = torch.tensor(prefix_lens_list, dtype=torch.int32, device=DEVICE)
+    # causal-valid per-token topk (block_size_q == 1)
+    tidx = torch.full((1, total_q, topk), -1, dtype=torch.int32, device=DEVICE)
+    row = 0
+    for b, (s, p) in enumerate(zip(seq_lens_list, prefix_lens_list)):
+        for j in range(s - p):
+            nb = (p + j) // P + 1
+            ak = min(topk, nb)
+            sel = torch.randperm(nb, device=DEVICE)[:ak].sort().values
+            tidx[0, row, :ak] = sel.to(torch.int32)
+            row += 1
+    return q, k, v, r2t, sids, cu, seq_lens, prefix, tidx, max(q_lens)
+
+
+@pytest.mark.parametrize(
+    "seq_lens,prefix_lens,branch",
+    [
+        ((530, 700), (500, 680), "cutlass_short_q"),  # qo <= 32
+        ((513, 769), (0, 257), "cute_long_q"),  # qo > 32
+    ],
+    ids=["cutlass_short_q", "cute_long_q"],
+)
+def test_msa_fp8_prefill_vs_triton_fp8(seq_lens, prefix_lens, branch):
+    q, k, v, r2t, sids, cu, seq_lens_t, prefix, tidx, max_q = _prefill_setup(
+        list(seq_lens), list(prefix_lens)
+    )
+    q8, _ = qdq(q)
+    k8, _ = qdq(k)
+    v8, _ = qdq(v)
+    o_msa = msa_sparse_prefill_main(
+        q8, k8, v8, tidx, r2t, sids, cu, seq_lens_t, prefix, block_size_k=P
+    )
+    o_triton = flash_prefill_with_gqa_share_sparse(
+        q=q8,
+        k_cache=k8,
+        v_cache=v8,
+        sink=None,
+        req_to_token=r2t,
+        slot_ids=sids,
+        topk_idx=tidx,
+        block_size_q=1,
+        block_size_k=P,
+        cu_seqlens=cu,
+        seq_lens=seq_lens_t,
+        prefix_lens=prefix,
+        max_seqlen_q=max_q,
+    )
+    assert o_msa.dtype == torch.bfloat16
+    torch.testing.assert_close(
+        o_msa.float(), o_triton.float(), atol=X_ATOL, rtol=X_RTOL
+    )
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main([__file__, "-v"]))

python/sglang/srt/layers/attention/trtllm_mha_backend.py

@@ -204,6 +204,32 @@ class TRTLLMHAAttnBackend(FlashInferAttnBackend):
         #   KV fp8: q_type = fp8, out_type=model_runner.dtype
         self.is_xqa_impl = is_sm90_supported() or is_sm120_supported()
 
+        # trtllm-gen serves page_size >= 128 only through its dynamic
+        # tokens-per-page kernels, which exist solely for GQA with equal QK/V
+        # head dims (power-of-2 pages). Mirror that precondition here so an
+        # unsupported combo fails at construction instead of as a
+        # "Missing TRTLLM-GEN kernel" error during CUDA-graph capture.
+        # XQA (SM90/SM120 decode) has native page-128 kernels; no check needed.
+        if self.page_size >= 128 and not self.is_xqa_impl:
+            from sglang.srt.runtime_context import get_parallel
+
+            attn_tp_size = get_parallel().attn_tp_size
+            num_q_heads = config.num_attention_heads // attn_tp_size
+            num_kv_heads = config.get_num_kv_heads(attn_tp_size)
+            if (
+                num_q_heads // num_kv_heads <= 1
+                or config.head_dim != config.v_head_dim
+                or self.page_size & (self.page_size - 1) != 0
+            ):
+                raise ValueError(
+                    f"trtllm_mha with page_size={self.page_size} requires "
+                    f"trtllm-gen's dynamic tokens-per-page kernels, which only "
+                    f"support GQA (q heads per kv head > 1, got "
+                    f"{num_q_heads}/{num_kv_heads}) with equal QK/V head dims "
+                    f"(got {config.head_dim}/{config.v_head_dim}) and a "
+                    f"power-of-2 page size. Use --page-size 64 instead."
+                )
+
     def _check_decode_kv_access(self) -> None:
         supported_kinds = {
             KVCacheAttentionAccessKind.PLAIN,

python/sglang/srt/layers/radix_attention.py

@@ -128,6 +128,13 @@ class RadixAttention(nn.Module):
         self.v_scale = None
         self.k_scale_float = None
         self.v_scale_float = None
+        # MiniMax-M3 fp8 attention-GEMM scales (fp8 attn-GEMM mode): main q and
+        # lightning-indexer q/k/v. No checkpoint loader populates them yet;
+        # None means unit scale.
+        self.q_scale_float = None
+        self.idx_q_scale_float = None
+        self.idx_k_scale_float = None
+        self.idx_v_scale_float = None
         self.quant_method = None
 
         if quant_config is not None:

python/sglang/srt/mem_cache/hybrid_cache/hybrid_cache_controller.py

@@ -397,11 +397,18 @@ class HybridCacheController(BaseHiCacheController):
         if not self.write_queue:
             return
         op = CacheOperation.merge_ops(self.write_queue)
-        # Page-first write-back JIT kernels can keep destination host indices on CPU.
+        # Page-first staged write-back kernels need CPU destination host indices.
+        # A HostPoolGroup may mix staged and non-staged child pools, so let it
+        # normalize indices per child instead of moving the whole operation here.
         if (
             self.io_backend == "kernel"
             and self.mem_pool_host.layout == "page_first"
-            and getattr(self.mem_pool_host, "can_use_write_back_jit", False)
+            and (
+                getattr(self.mem_pool_host, "can_use_write_back_jit", False)
+                or getattr(
+                    self.mem_pool_host, "supports_per_pool_backup_indices", False
+                )
+            )
         ):
             host_indices = op.host_indices
             device_indices = op.device_indices

python/sglang/srt/mem_cache/kv_cache_configurator.py

@@ -1271,6 +1271,8 @@ class KVCacheConfigurator:
         return token_to_kv_pool
 
     def _build_minimax_sparse_kv_pool(self, *, max_total_num_tokens: int) -> KVCache:
+        from sglang.srt.server_args import m3_fp8_attn_gemm_enabled
+
         _hf_config = self.model_config.hf_config
         sparse_cfg = get_minimax_sparse_attention_config(_hf_config)
         dense_layer_ids, sparse_layer_ids = get_minimax_sparse_layer_ids(sparse_cfg)
@@ -1281,7 +1283,15 @@ class KVCacheConfigurator:
             size=max_total_num_tokens,
             page_size=get_schedule().page_size,
             dtype=self.kv_cache_dtype,
-            index_dtype=self.model_dtype,
+            # fp8 attn-GEMM mode opts the lightning-indexer cache into
+            # fp8 too (fp8 indexer GEMMs); fp8 KV without the mode
+            # (e5m2 or non-trtllm_mha backend) keeps the indexer bf16
+            # with the widening-dequant contract.
+            index_dtype=(
+                self.kv_cache_dtype
+                if m3_fp8_attn_gemm_enabled(self.server_args)
+                else self.model_dtype
+            ),
             head_num=self.model_config.get_num_kv_heads(get_parallel().attn_tp_size),
             head_dim=self.model_config.head_dim,
             idx_head_dim=sparse_cfg["sparse_index_dim"],

python/sglang/srt/mem_cache/memory_pool.py

@@ -4810,6 +4810,11 @@ class MiniMaxSparseKVPool(KVCache):
     ) -> None:
         self.layer_transfer_counter = layer_transfer_counter
 
+    def get_kv_cache_quant_method(self) -> Any:
+        # The base unwrap chain only knows full_kv_pool/swa_kv_pool; the dense
+        # KV (what attention backends quantize against) lives in main_pool here.
+        return self.main_pool.get_kv_cache_quant_method()
+
     def _wait_for_layer(self, layer_id: int) -> None:
         if self.layer_transfer_counter is not None:
             self.layer_transfer_counter.wait_until(layer_id - self.start_layer)
@@ -4858,10 +4863,14 @@ class MiniMaxSparseKVPool(KVCache):
         loc: torch.Tensor,
         cache_k: torch.Tensor,
         cache_v: torch.Tensor,
-        k_scale: float = 1.0,
+        k_scale: Optional[float] = None,
-        v_scale: float = 1.0,
+        v_scale: Optional[float] = None,
     ) -> None:
-        """Write main K/V at `loc`. Works for any layer (dense or sparse)."""
+        """Write main K/V at `loc`. Works for any layer (dense or sparse).
+
+        Scale semantics follow MHATokenToKVPool: None means unit scale;
+        a non-None scale is applied with an in-place div_ before the fp8 cast.
+        """
         self.main_pool.set_kv_buffer(
             layer,
             loc,
@@ -4877,8 +4886,8 @@ class MiniMaxSparseKVPool(KVCache):
         loc: torch.Tensor,
         cache_idx_k: torch.Tensor,
         cache_idx_v: torch.Tensor,
-        k_scale: float = 1.0,
+        k_scale: Optional[float] = None,
-        v_scale: float = 1.0,
+        v_scale: Optional[float] = None,
     ) -> None:
         mapped_id = self.index_kv_layer_id_mapping.get(layer.layer_id)
         if mapped_id is None:
@@ -4902,6 +4911,7 @@ class MiniMaxSparseKVPool(KVCache):
         layer: RadixAttention,
         loc: torch.Tensor,
         cache_idx_k: torch.Tensor,
+        k_scale: Optional[float] = None,
     ) -> None:
         mapped_id = self.index_k_layer_id_mapping.get(layer.layer_id)
         if mapped_id is None:
@@ -4912,6 +4922,8 @@ class MiniMaxSparseKVPool(KVCache):
             )
         sub_pool = self.index_k_pool
         if cache_idx_k.dtype != sub_pool.dtype:
+            if k_scale is not None:
+                cache_idx_k = cache_idx_k / k_scale
             cache_idx_k = cache_idx_k.to(sub_pool.dtype)
         if sub_pool.store_dtype != sub_pool.dtype:
             cache_idx_k = cache_idx_k.view(sub_pool.store_dtype)
@@ -4950,6 +4962,10 @@ class MiniMaxSparseKVPool(KVCache):
         cache_v: torch.Tensor,
         cache_idx_k: torch.Tensor,
         cache_idx_v: Optional[torch.Tensor],
+        k_scale: Optional[float] = None,
+        v_scale: Optional[float] = None,
+        idx_k_scale: Optional[float] = None,
+        idx_v_scale: Optional[float] = None,
     ) -> None:
         """Store main K/V + index K (+ optional index V) for a sparse layer in
         one fused JIT launch, falling back to separate stores when not applicable."""
@@ -4984,12 +5000,24 @@ class MiniMaxSparseKVPool(KVCache):
             )
             return
 
-        # Fallback: separate stores (identical semantics).
+        # Fallback: separate stores (identical semantics; quantizes for fp8
-        self.set_kv_buffer(layer, loc, cache_k, cache_v)
+        # pools — the fused raw-byte path is disqualified there by
+        # _can_fuse_kv_index_store's dtype-equality checks). Scales use the
+        # None-means-unit convention throughout: MHATokenToKVPool.set_kv_buffer
+        # applies any non-None scale with an IN-PLACE div_ (extra kernel +
+        # caller-tensor mutation), which must not fire for unit scale.
+        self.set_kv_buffer(layer, loc, cache_k, cache_v, k_scale, v_scale)
         if disable_value:
-            self.set_index_k_buffer(layer, loc, cache_idx_k)
+            self.set_index_k_buffer(layer, loc, cache_idx_k, idx_k_scale)
         else:
-            self.set_index_kv_buffer(layer, loc, cache_idx_k, cache_idx_v)
+            self.set_index_kv_buffer(
+                layer,
+                loc,
+                cache_idx_k,
+                cache_idx_v,
+                idx_k_scale,
+                idx_v_scale,
+            )
 
     def get_kv_size_bytes(self):
         sub_pools = [self.main_pool, self.index_kv_pool, self.index_k_pool]

python/sglang/srt/mem_cache/memory_pool_host.py

@@ -1528,10 +1528,12 @@ class HostPoolGroup:
         self.page_size = self.anchor_entry.host_pool.page_size
         self.device = self.anchor_entry.host_pool.device
         self.size = self.anchor_entry.host_pool.size
-        self.can_use_write_back_jit = all(
+        child_write_back_jit = [
             getattr(entry.host_pool, "can_use_write_back_jit", False)
             for entry in entries
-        )
+        ]
+        self.can_use_write_back_jit = all(child_write_back_jit)
+        self.supports_per_pool_backup_indices = any(child_write_back_jit)
 
     @property
     def kv_buffer(self):
@@ -1632,6 +1634,39 @@ class HostPoolGroup:
                 io_backend,
             )
 
+    def _backup_uses_cpu_host_indices(self, host_pool, io_backend) -> bool:
+        return (
+            io_backend == "kernel"
+            and getattr(host_pool, "layout", None) == "page_first"
+            and getattr(host_pool, "can_use_write_back_jit", False)
+        )
+
+    def _kernel_index_device(self, entry, device_indices):
+        if device_indices is not None and device_indices.is_cuda:
+            return device_indices.device
+        return getattr(entry.device_pool, "device", None)
+
+    def _normalize_backup_indices(
+        self, entry, host_indices, device_indices, io_backend
+    ):
+        if io_backend != "kernel":
+            return host_indices, device_indices
+
+        if self._backup_uses_cpu_host_indices(entry.host_pool, io_backend):
+            if host_indices.is_cuda:
+                host_indices = host_indices.cpu()
+            return host_indices, device_indices
+
+        if not host_indices.is_cuda:
+            target_device = self._kernel_index_device(entry, device_indices)
+            if target_device is not None:
+                host_indices = host_indices.to(target_device, non_blocking=True)
+                if host_indices.is_cuda:
+                    host_indices.record_stream(
+                        torch.cuda.current_stream(host_indices.device)
+                    )
+        return host_indices, device_indices
+
     def backup_from_device_all_layer(
         self,
         device_pool,
@@ -1641,10 +1676,13 @@ class HostPoolGroup:
         pool_transfers: Optional[list] = None,
     ) -> None:
         # 1. Anchor (KV) backup
+        anchor_host_indices, anchor_device_indices = self._normalize_backup_indices(
+            self.anchor_entry, host_indices, device_indices, io_backend
+        )
         self.anchor_entry.host_pool.backup_from_device_all_layer(
             self.anchor_entry.device_pool,
-            host_indices,
+            anchor_host_indices,
-            device_indices,
+            anchor_device_indices,
             io_backend,
         )
         # 2. Extra pool backup
@@ -1652,12 +1690,20 @@ class HostPoolGroup:
             entry = self.entry_map.get(transfer.name)
             if entry is None or transfer.host_indices is None:
                 continue
-            entry.host_pool.backup_from_device_all_layer(
+            transfer_host_indices, transfer_device_indices = (
-                entry.device_pool,
+                self._normalize_backup_indices(
+                    entry,
                     transfer.host_indices,
                     transfer.device_indices,
                     io_backend,
                 )
+            )
+            entry.host_pool.backup_from_device_all_layer(
+                entry.device_pool,
+                transfer_host_indices,
+                transfer_device_indices,
+                io_backend,
+            )
 
 
 class DSAIndexerPoolHost(HostKVCache):

python/sglang/srt/server_args.py

@@ -9025,6 +9025,28 @@ class ServerArgs:
         }
 
 
+def m3_fp8_attn_gemm_enabled(args) -> bool:
+    """Whether MiniMax-M3 attention GEMMs run in fp8 (no opt-in flag; active
+    whenever possible): fp8_e4m3 main + index KV caches, fp8-cast q, fp8
+    sparse/MSA kernels, with dense layers on trtllm_mha's fp8-q path. Needs
+    kv_cache_dtype fp8_e4m3 (e5m2 would silently mis-dispatch fmha_sm100's
+    e4m3 kernel), the trtllm_mha backend (the only dense backend with fp8-q
+    GEMMs), and SM100 (MSA fp8 variants and trtllm-gen fp8 dense kernels are
+    sm100-only). SGLANG_DISABLE_M3_FP8_ATTN_GEMM=1 is the kill switch:
+    it forces the pre-fp8 numerics (bf16 indexer + widening sparse path,
+    bf16 q) without having to move off trtllm_mha.
+    """
+    from sglang.srt.environ import envs
+    from sglang.srt.utils.common import is_sm100_supported
+
+    return (
+        args.kv_cache_dtype == "fp8_e4m3"
+        and args.attention_backend == "trtllm_mha"
+        and is_sm100_supported()
+        and not envs.SGLANG_DISABLE_M3_FP8_ATTN_GEMM.get()
+    )
+
+
 # NOTE: The process-wide ServerArgs is owned by the runtime context
 # (sglang.srt.runtime_context). The two functions below are LEGACY shims kept
 # for the existing call-sites; they publish/read the same live object by

test/registered/attention/unittests/dense/test_trtllm_mha.py

@@ -6,7 +6,11 @@ import torch
 
 from sglang.srt.model_executor.forward_batch_info import ForwardMode
 from sglang.srt.utils import is_flashinfer_available
-from sglang.srt.utils.common import is_sm90_supported, is_sm120_supported
+from sglang.srt.utils.common import (
+    is_sm90_supported,
+    is_sm100_supported,
+    is_sm120_supported,
+)
 from sglang.test.test_utils import CustomTestCase
 
 sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
@@ -14,6 +18,7 @@ sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
 from sglang.test.ci.ci_register import register_cuda_ci
 from sglang.test.kits.attention_unittest.attention_methods.dense_attention import (
     DenseAttentionCase,
+    build_dense_attention_fixture,
     run_dense_attention_case,
 )
 from sglang.test.kits.attention_unittest.runner_modes.cuda_graph_decode_runner import (
@@ -200,6 +205,148 @@ class TestTRTLLMMHADenseAttentionBackendCorrectness(CustomTestCase):
                     hidden_size=self.HIDDEN_SIZE,
                 )
 
+    # XQA has native page-128 kernels (any head layout). max_context_len must
+    # be a page multiple so the kit's per-request slot ranges stay page-aligned.
+    def test_page128_decode(self):
+        case = DenseAttentionCase(
+            name="trtllm_mha_xqa_decode_page128_boundary",
+            backend="trtllm_mha",
+            forward_mode=ForwardMode.DECODE,
+            num_heads=4,
+            num_kv_heads=2,
+            page_size=128,
+            prefix_lens=(127, 128, 200),
+        )
+        run_dense_attention_case(
+            self,
+            case,
+            head_dim=self.HEAD_DIM,
+            hidden_size=self.HIDDEN_SIZE,
+            max_context_len=512,
+        )
+
+
+@unittest.skipIf(
+    not torch.cuda.is_available()
+    or not is_flashinfer_available()
+    or not is_sm100_supported(),
+    "CUDA + FlashInfer TRT-LLM-GEN (SM100) are required",
+)
+class TestTRTLLMMHAPage128TrtllmGen(CustomTestCase):
+    """page_size=128 on trtllm-gen (SM100) via dynamic tokens-per-page kernels.
+
+    Those kernels only exist for GQA (q heads per kv head > 1) with equal QK/V
+    head dims, so every positive case here is GQA; the MHA layout must fail at
+    backend construction (see test_page128_mha_rejected_at_init). All cases
+    pass max_context_len=512: the kit's per-request slot ranges start at
+    ``page_size + req_idx * max_context_len``, so it must be a page multiple.
+    """
+
+    HEAD_DIM = 64
+    HIDDEN_SIZE = 256
+    MAX_CONTEXT_LEN = 512
+
+    DECODE_CASES = (
+        DenseAttentionCase(
+            name="trtllm_gen_gqa_decode_page128_boundary",
+            backend="trtllm_mha",
+            forward_mode=ForwardMode.DECODE,
+            num_heads=4,
+            num_kv_heads=2,
+            page_size=128,
+            prefix_lens=(127, 128, 200),
+        ),
+        DenseAttentionCase(
+            name="trtllm_gen_gqa4_decode_page128_bsz1",
+            backend="trtllm_mha",
+            forward_mode=ForwardMode.DECODE,
+            num_heads=8,
+            num_kv_heads=2,
+            page_size=128,
+            prefix_lens=(300,),
+        ),
+    )
+
+    EXTEND_CASES = (
+        DenseAttentionCase(
+            name="trtllm_gen_gqa_extend_page128",
+            backend="trtllm_mha",
+            forward_mode=ForwardMode.EXTEND,
+            num_heads=4,
+            num_kv_heads=2,
+            page_size=128,
+            prefix_lens=(0, 128),
+            extend_lens=(130, 5),
+        ),
+    )
+
+    CUDA_GRAPH_DECODE_CASES = (
+        DenseAttentionCase(
+            name="runner_cuda_graph_trtllm_gen_gqa_decode_page128",
+            backend="trtllm_mha",
+            forward_mode=ForwardMode.DECODE,
+            num_heads=4,
+            num_kv_heads=2,
+            page_size=128,
+            prefix_lens=(127, 128, 200),
+        ),
+    )
+
+    def test_page128_decode_cases(self):
+        for case in self.DECODE_CASES:
+            with self.subTest(case=case.name):
+                run_dense_attention_case(
+                    self,
+                    case,
+                    head_dim=self.HEAD_DIM,
+                    hidden_size=self.HIDDEN_SIZE,
+                    max_context_len=self.MAX_CONTEXT_LEN,
+                )
+
+    def test_page128_extend_cases(self):
+        for case in self.EXTEND_CASES:
+            with self.subTest(case=case.name):
+                run_dense_attention_case(
+                    self,
+                    case,
+                    head_dim=self.HEAD_DIM,
+                    hidden_size=self.HIDDEN_SIZE,
+                    max_context_len=self.MAX_CONTEXT_LEN,
+                )
+
+    def test_page128_cuda_graph_decode_cases(self):
+        for case in self.CUDA_GRAPH_DECODE_CASES:
+            with self.subTest(case=case.name):
+                run_dense_cuda_graph_decode_case(
+                    self,
+                    case,
+                    head_dim=self.HEAD_DIM,
+                    hidden_size=self.HIDDEN_SIZE,
+                    max_context_len=self.MAX_CONTEXT_LEN,
+                )
+
+    def test_page128_mha_rejected_at_init(self):
+        # heads_per_kv == 1 has no page-128 trtllm-gen kernel; the backend must
+        # refuse at construction (not fail mid-capture with a missing-kernel
+        # RuntimeError from flashinfer).
+        case = DenseAttentionCase(
+            name="trtllm_gen_mha_decode_page128_rejected",
+            backend="trtllm_mha",
+            forward_mode=ForwardMode.DECODE,
+            num_heads=4,
+            num_kv_heads=4,
+            page_size=128,
+            prefix_lens=(7,),
+        )
+        with self.assertRaisesRegex(ValueError, "dynamic tokens-per-page"):
+            build_dense_attention_fixture(
+                self,
+                case,
+                head_dim=self.HEAD_DIM,
+                hidden_size=self.HIDDEN_SIZE,
+                max_context_len=self.MAX_CONTEXT_LEN,
+            )
+
 
 if __name__ == "__main__":
     unittest.main()

test/registered/unit/test_model_overrides.py

@@ -22,6 +22,7 @@ from sglang.srt.arg_groups.overrides import (
     register_model_override,
     validate_declarations,
 )
+from sglang.srt.environ import envs
 from sglang.srt.runtime_context import (
     get_context,
     get_server_args,
@@ -1642,6 +1643,87 @@ class TestGoldenModelOverrides(_IsolatedPublish):
                 {},
             )
 
+    def test_m3_fp8_attn_gemm_resolution(self):
+        from sglang.srt.arg_groups.overrides import _minimax_m3_overrides
+        from sglang.srt.server_args import m3_fp8_attn_gemm_enabled
+
+        def _args(**kw):
+            defaults = dict(
+                attention_backend="trtllm_mha",
+                kv_cache_dtype="fp8_e4m3",
+            )
+            defaults.update(kw)
+            return SimpleNamespace(**defaults)
+
+        with patch("sglang.srt.utils.common.is_sm100_supported", return_value=True):
+            # e4m3 + trtllm_mha + SM100: mode active
+            self.assertTrue(m3_fp8_attn_gemm_enabled(_args()))
+            # fa4 dense backend: mode inactive (no fp8-q GEMM path)
+            self.assertFalse(m3_fp8_attn_gemm_enabled(_args(attention_backend="fa4")))
+            # bf16 KV: mode inactive
+            self.assertFalse(m3_fp8_attn_gemm_enabled(_args(kv_cache_dtype="auto")))
+            # e5m2: mode inactive (fmha_sm100's variant lookup would silently
+            # dispatch the e4m3 kernel)
+            self.assertFalse(m3_fp8_attn_gemm_enabled(_args(kv_cache_dtype="fp8_e5m2")))
+            # SGLANG_DISABLE_M3_FP8_ATTN_GEMM kill switch wins over an
+            # otherwise-active config
+            with envs.SGLANG_DISABLE_M3_FP8_ATTN_GEMM.override(True):
+                self.assertFalse(m3_fp8_attn_gemm_enabled(_args()))
+        with patch("sglang.srt.utils.common.is_sm100_supported", return_value=False):
+            # non-SM100: mode inactive
+            self.assertFalse(m3_fp8_attn_gemm_enabled(_args()))
+
+        def _m3_args(**kw):
+            defaults = dict(
+                quantization=None,
+                _quantization_explicitly_unset=True,
+                attention_backend=None,
+                prefill_attention_backend=None,
+                decode_attention_backend=None,
+                page_size=None,
+                moe_runner_backend="auto",
+                kv_cache_dtype="auto",
+            )
+            defaults.update(kw)
+            ns = SimpleNamespace(**defaults)
+            ns.is_attention_backend_not_set = lambda: (
+                ns.attention_backend is None
+                and ns.prefill_attention_backend is None
+                and ns.decode_attention_backend is None
+            )
+            return ns
+
+        hf = SimpleNamespace()
+        with patch.object(overrides_module, "is_hip", return_value=False), patch.object(
+            overrides_module, "is_sm100_supported", return_value=True
+        ), patch.object(overrides_module, "get_quantization_config", return_value=None):
+            # fp8_e4m3 KV: SM100 backend default flips to trtllm_mha (the only
+            # dense backend with the fp8-q GEMM path); page snaps to 128
+            ov = _minimax_m3_overrides(_m3_args(kv_cache_dtype="fp8_e4m3"), hf)
+            self.assertEqual(ov["attention_backend"], "trtllm_mha")
+            self.assertEqual(ov["page_size"], 128)
+            # auto KV: fa4 stays the SM100 default
+            ov = _minimax_m3_overrides(_m3_args(), hf)
+            self.assertEqual(ov["attention_backend"], "fa4")
+            self.assertEqual(ov["page_size"], 128)
+            # e5m2 KV: stays on fa4 + the widening Triton path, and warns
+            with self.assertLogs(
+                "sglang.srt.arg_groups.overrides", level="WARNING"
+            ) as logs:
+                ov = _minimax_m3_overrides(_m3_args(kv_cache_dtype="fp8_e5m2"), hf)
+            self.assertEqual(ov["attention_backend"], "fa4")
+            self.assertIn("fp8_e5m2", "\n".join(logs.output))
+            # explicit backend choice is never overridden
+            ov = _minimax_m3_overrides(
+                _m3_args(kv_cache_dtype="fp8_e4m3", attention_backend="fa4"), hf
+            )
+            self.assertNotIn("attention_backend", ov)
+            # kill switch also reverts the SM100 backend default to fa4
+            with envs.SGLANG_DISABLE_M3_FP8_ATTN_GEMM.override(True):
+                ov = _minimax_m3_overrides(_m3_args(kv_cache_dtype="fp8_e4m3"), hf)
+            self.assertEqual(ov["attention_backend"], "fa4")
+            self.assertEqual(ov["page_size"], 128)
+
     def test_page_constraint_passes_at_callable_level(self):
         from sglang.srt.arg_groups.overrides import (
             ResolvedView,
@@ -1682,6 +1764,30 @@ class TestGoldenModelOverrides(_IsolatedPublish):
             ),
             {"page_size": 64},
         )
+        # trtllm_mha accepts 128 (trtllm-gen dynamic tokens-per-page kernels)
+        self.assertEqual(
+            _mla_backend_page_constraints(
+                _view(attention_backend="trtllm_mha", page_size=128)
+            ),
+            {},
+        )
+        # trtllm_mha with an unsupported page still snaps to 64
+        self.assertEqual(
+            _mla_backend_page_constraints(
+                _view(attention_backend="trtllm_mha", page_size=256)
+            ),
+            {"page_size": 64},
+        )
+        # chained: cutlass_mla decode -> 128, then trtllm_mha prefill keeps 128
+        self.assertEqual(
+            _mla_backend_page_constraints(
+                _view(
+                    decode_attention_backend="cutlass_mla",
+                    prefill_attention_backend="trtllm_mha",
+                )
+            ),
+            {"page_size": 128},
+        )
         # no matching backend: nothing declared
         self.assertEqual(_mla_backend_page_constraints(_view()), {})