[MiniMax-M3] Retune decode indexer launch grid for occupancy (#130 )

Signed-off-by: Zijing Liu <liuzijing2014@gmail.com>
Add MiniMax-M3 modelopt nvfp4 support (#46380 )
2026-06-25 18:14:52 +00:00 · 2026-06-25 11:11:41 -07:00 · 2026-06-23 00:01:30 +00:00 · 2026-06-22 05:21:19 +00:00 · 2026-06-20 13:53:02 -07:00 · 2026-06-20 13:52:56 -07:00
25 changed files with 1297 additions and 131 deletions
@@ -17,7 +17,7 @@ else()
  FetchContent_Declare(
    fmha_sm100
    GIT_REPOSITORY https://github.com/vllm-project/MSA.git
-    GIT_TAG 544eee5e09ae2dfa774d5b06739013f9b7402c57
+    GIT_TAG fee783153f3efe57e3e933c5cb7e267a7cebcfb5
    GIT_PROGRESS TRUE
    CONFIGURE_COMMAND ""
    BUILD_COMMAND ""
@@ -36,13 +36,38 @@ set(FMHA_SM100_PY_ROOT "${fmha_sm100_SOURCE_DIR}/python/fmha_sm100")

 install(FILES
  "${FMHA_SM100_PY_ROOT}/__init__.py"
+  "${FMHA_SM100_PY_ROOT}/api.py"
+  "${FMHA_SM100_PY_ROOT}/bench_utils.py"
+  "${FMHA_SM100_PY_ROOT}/jit.py"
  "${FMHA_SM100_PY_ROOT}/sparse.py"
+  "${FMHA_SM100_PY_ROOT}/sparse_fmha_adapter.py"
  DESTINATION vllm/third_party/fmha_sm100
  COMPONENT fmha_sm100)

+install(DIRECTORY "${FMHA_SM100_PY_ROOT}/csrc/"
+  DESTINATION vllm/third_party/fmha_sm100/csrc
+  COMPONENT fmha_sm100
+  PATTERN "__pycache__" EXCLUDE
+  PATTERN "*.pyc" EXCLUDE
+  PATTERN ".git*" EXCLUDE)
+
 install(DIRECTORY "${FMHA_SM100_PY_ROOT}/cute/"
  DESTINATION vllm/third_party/fmha_sm100/cute
  COMPONENT fmha_sm100
  PATTERN "__pycache__" EXCLUDE
  PATTERN "*.pyc" EXCLUDE
  PATTERN ".git*" EXCLUDE)
+
+install(DIRECTORY "${FMHA_SM100_PY_ROOT}/cutlass/include/"
+  DESTINATION vllm/third_party/fmha_sm100/cutlass/include
+  COMPONENT fmha_sm100
+  PATTERN "__pycache__" EXCLUDE
+  PATTERN "*.pyc" EXCLUDE
+  PATTERN ".git*" EXCLUDE)
+
+install(DIRECTORY "${FMHA_SM100_PY_ROOT}/cutlass/tools/util/include/"
+  DESTINATION vllm/third_party/fmha_sm100/cutlass/tools/util/include
+  COMPONENT fmha_sm100
+  PATTERN "__pycache__" EXCLUDE
+  PATTERN "*.pyc" EXCLUDE
+  PATTERN ".git*" EXCLUDE)
@@ -67,6 +67,13 @@
  #include "../quantization/w8a8/fp8/nvidia/quant_utils.cuh"
 #endif

+// Direct float -> E4M3 FP8 conversion for the indexer Q / index-K outputs.
+#ifndef USE_ROCM
+  #include <cuda_fp8.h>
+#else
+  #include <hip/hip_fp8.h>
+#endif
+
 #ifndef FINAL_MASK
  #ifdef USE_ROCM
    #define FINAL_MASK 0xffffffffffffffffULL
@@ -75,6 +82,19 @@
  #endif
 #endif

+#ifdef USE_ROCM
+// ROCm-compatible direct float -> E4M3 FP8 conversion (mirrors the DeepSeek V4
+// fused kernel).
+__device__ __forceinline__ uint8_t rocm_cvt_float_to_fp8_e4m3(float val) {
+  #if defined(HIP_FP8_TYPE_OCP)
+  __hip_fp8_e4m3 fp8_val(val);
+  #else
+  __hip_fp8_e4m3_fnuz fp8_val(val);
+  #endif
+  return reinterpret_cast<uint8_t&>(fp8_val);
+}
+#endif
+
 namespace vllm {
 namespace minimax_m3_fused_ops {

@@ -193,6 +213,8 @@ __device__ __forceinline__ void storeElems(
  *reinterpret_cast<uint2*>(dst) = v;
 }

+// Main K/V cache store. kAuto = unquantized (cache_t == scalar_t); fp8 cache
+// dtypes use the scaled-convert path with identity scale.
 template <typename scalar_t, typename cache_t, Fp8KVCacheDataType kv_dt>
 __device__ __forceinline__ void storeCacheElems(
    cache_t* __restrict__ dst, float const (&elems)[kElemsPerLane]) {
@@ -208,6 +230,32 @@ __device__ __forceinline__ void storeCacheElems(
  }
 }

+// Store 4 fp32 registers -> 4 contiguous E4M3 FP8 bytes (direct cast,
+// saturating to ±448). Used for the fp8 indexer-Q / index-K outputs; no scale
+// (RMSNorm outputs are O(1) and the score path only needs relative block
+// ordering).
+__device__ __forceinline__ void storeElemsFp8(
+    uint8_t* __restrict__ dst, float const (&elems)[kElemsPerLane]) {
+  constexpr float kFp8Max = 448.0f;
+#ifndef USE_ROCM
+  __nv_fp8x2_storage_t out2[kElemsPerLane / 2];
+  #pragma unroll
+  for (int i = 0; i < kElemsPerLane / 2; i++) {
+    float2 vv = make_float2(elems[2 * i], elems[2 * i + 1]);
+    vv.x = fminf(fmaxf(vv.x, -kFp8Max), kFp8Max);
+    vv.y = fminf(fmaxf(vv.y, -kFp8Max), kFp8Max);
+    out2[i] = __nv_cvt_float2_to_fp8x2(vv, __NV_SATFINITE, __NV_E4M3);
+  }
+  *reinterpret_cast<uint32_t*>(dst) = *reinterpret_cast<uint32_t const*>(out2);
+#else
+  #pragma unroll
+  for (int i = 0; i < kElemsPerLane; i++) {
+    float vv = fminf(fmaxf(elems[i], -kFp8Max), kFp8Max);
+    dst[i] = rocm_cvt_float_to_fp8_e4m3(vv);
+  }
+#endif
+}
+
 // ────────────────────────────────────────────────────────────────────────────
 // Kernel
 // ────────────────────────────────────────────────────────────────────────────
@@ -224,12 +272,14 @@ __device__ __forceinline__ void storeCacheElems(
 //     V : nkv  only if kInsertKV        (V-cache insert; no warps in dense)
 //     IQ: niq  only if kIsSparse        (norm+RoPE)
 //     IK: 1    only if kIsSparse        (norm+RoPE; +index-cache insert)
+// cache_t/kv_dt: main attention KV-cache dtype (auto/fp8). out_idx_t/kFp8Idx:
+// indexer index-K cache + index-Q output dtype (scalar_t or e4m3 byte).
 template <typename scalar_t, typename cache_t, Fp8KVCacheDataType kv_dt,
-          bool kIsSparse, bool kInsertKV>
+          typename out_idx_t, bool kIsSparse, bool kInsertKV, bool kFp8Idx>
 __global__ void fusedMiniMaxM3QNormRopeKVInsertKernel(
    scalar_t* __restrict__ qkv,  // [N, qkv_row] in/out (packs index if sparse)
-    scalar_t* __restrict__ q_out,        // [N, nq*128] contiguous, or nullptr
-    scalar_t* __restrict__ index_q_out,  // [N, niq*128] contiguous, or nullptr
+    scalar_t* __restrict__ q_out,         // [N, nq*128] contiguous, or nullptr
+    out_idx_t* __restrict__ index_q_out,  // [N, niq*128]; scalar_t or e4m3 byte
    scalar_t const* __restrict__ q_norm_w,
    scalar_t const* __restrict__ k_norm_w,
    scalar_t const* __restrict__ iq_norm_w,
@@ -238,8 +288,8 @@ __global__ void fusedMiniMaxM3QNormRopeKVInsertKernel(
    int64_t const* __restrict__ positions,       // [N] i64
    int64_t const* __restrict__ slot_mapping,    // main K/V slots or nullptr
    int64_t const* __restrict__ index_slot_mapping,  // index K slots/nullptr
-    cache_t* __restrict__ kv_cache,      // [nb,2,bs,nkv,128] or nullptr
-    scalar_t* __restrict__ index_cache,  // [nb*bs, 128] or nullptr
+    cache_t* __restrict__ kv_cache,       // [nb,2,bs,nkv,128] or nullptr
+    out_idx_t* __restrict__ index_cache,  // [nb*bs, 128]; scalar_t or e4m3 byte
    float const eps, int const rotary_dim, int const num_tokens, int const nq,
    int const nkv, int const niq, int const block_size,
    // kv_cache strides (in elements) for logical shape [nb, 2, bs, nkv, 128].
@@ -334,9 +384,12 @@ __global__ void fusedMiniMaxM3QNormRopeKVInsertKernel(
      store_ptr = q_out + static_cast<int64_t>(tokenIdx) * nq * kHeadDim +
                  slot * kHeadDim;
    } else if (isIQ && index_q_out != nullptr) {
-      store_ptr = index_q_out +
-                  static_cast<int64_t>(tokenIdx) * niq * kHeadDim +
-                  (slot - iq_begin) * kHeadDim;
+      // bf16 index_q_out: gather here. fp8: written by the explicit fp8 store.
+      if constexpr (!kFp8Idx) {
+        store_ptr = index_q_out +
+                    static_cast<int64_t>(tokenIdx) * niq * kHeadDim +
+                    (slot - iq_begin) * kHeadDim;
+      }
    }

    // PDL: wait for the predecessor kernel (the qkv-projection GEMM that
@@ -356,7 +409,19 @@ __global__ void fusedMiniMaxM3QNormRopeKVInsertKernel(
      scalar_t const* cos_ptr = cos_sin_cache + pos * rotary_dim;
      normAndRope<scalar_t>(elems, laneId, eps, norm_w, do_rope, rotary_dim,
                            cos_ptr, /*apply_norm=*/norm_w != nullptr);
-      storeElems<scalar_t>(store_ptr + dim_base, elems);
+      if constexpr (kFp8Idx) {
+        // index_q is e4m3 bytes; Q/K (and in-place index_k) stay scalar_t.
+        if (isIQ && index_q_out != nullptr) {
+          storeElemsFp8(index_q_out +
+                            static_cast<int64_t>(tokenIdx) * niq * kHeadDim +
+                            (slot - iq_begin) * kHeadDim + dim_base,
+                        elems);
+        } else {
+          storeElems<scalar_t>(store_ptr + dim_base, elems);
+        }
+      } else {
+        storeElems<scalar_t>(store_ptr + dim_base, elems);
+      }
    }

    // ── Cache inserts (sparse serving only). ───────────────────────────────
@@ -367,8 +432,11 @@ __global__ void fusedMiniMaxM3QNormRopeKVInsertKernel(
                             : (isIK ? index_slot_mapping[tokenIdx] : -1);
      if (sm >= 0) {  // skip padded / unscheduled tokens
        if (isIK) {
-          scalar_t* dst = index_cache + sm * kHeadDim + dim_base;
-          storeElems<scalar_t>(dst, elems);
+          if constexpr (kFp8Idx) {
+            storeElemsFp8(index_cache + sm * kHeadDim + dim_base, elems);
+          } else {
+            storeElems<scalar_t>(index_cache + sm * kHeadDim + dim_base, elems);
+          }
        } else if (isK || isV) {
          // kv_cache logical shape [num_blocks, 2, block_size, nkv, head_dim].
          // Paging is logical (block = sm/block_size, token = sm%block_size);
@@ -398,19 +466,19 @@ __global__ void fusedMiniMaxM3QNormRopeKVInsertKernel(
 // Launch wrapper
 // ────────────────────────────────────────────────────────────────────────────
 template <typename scalar_t, typename cache_t, Fp8KVCacheDataType kv_dt>
-void launchFusedMiniMaxM3(scalar_t* qkv, scalar_t* q_out, scalar_t* index_q_out,
-                          scalar_t const* q_norm_w, scalar_t const* k_norm_w,
-                          scalar_t const* iq_norm_w, scalar_t const* ik_norm_w,
-                          scalar_t const* cos_sin_cache,
-                          int64_t const* positions, int64_t const* slot_mapping,
-                          int64_t const* index_slot_mapping, cache_t* kv_cache,
-                          scalar_t* index_cache, float const eps,
-                          int const rotary_dim, int const num_tokens,
-                          int const nq, int const nkv, int const niq,
-                          int const block_size, int64_t const kv_s_block,
-                          int64_t const kv_s_kv, int64_t const kv_s_token,
-                          int64_t const kv_s_head, bool const has_index,
-                          bool const insert_kv, cudaStream_t stream) {
+void launchFusedMiniMaxM3(
+    scalar_t* qkv, scalar_t* q_out, void* index_q_out, scalar_t const* q_norm_w,
+    scalar_t const* k_norm_w, scalar_t const* iq_norm_w,
+    scalar_t const* ik_norm_w, scalar_t const* cos_sin_cache,
+    int64_t const* positions, int64_t const* slot_mapping,
+    int64_t const* index_slot_mapping, cache_t* kv_cache, void* index_cache,
+    float const eps, int const rotary_dim, int const num_tokens, int const nq,
+    int const nkv, int const niq, int const block_size,
+    int64_t const kv_s_block, int64_t const kv_s_kv, int64_t const kv_s_token,
+    int64_t const kv_s_head, bool const has_index, bool const insert_kv,
+    bool const fp8_idx, cudaStream_t stream) {
+  // Index outputs are scalar_t (bf16) or e4m3 bytes (uint8_t); reinterpret the
+  // void* pointers per instantiation in the LAUNCH macro.
  // Slot count must match the kernel's compile-time gating.
  int const v_slots = insert_kv ? nkv : 0;
  int const idx_slots = has_index ? niq + 1 : 0;
@@ -440,25 +508,27 @@ void launchFusedMiniMaxM3(scalar_t* qkv, scalar_t* q_out, scalar_t* index_q_out,
  config.attrs = attrs;
  config.numAttrs = (sm_version >= 90) ? 1 : 0;

-  #define LAUNCH(IS_SPARSE, INSERT)                                           \
-    cudaLaunchKernelEx(                                                       \
-        &config,                                                              \
-        fusedMiniMaxM3QNormRopeKVInsertKernel<scalar_t, cache_t, kv_dt,       \
-                                              IS_SPARSE, INSERT>,             \
-        qkv, q_out, index_q_out, q_norm_w, k_norm_w, iq_norm_w, ik_norm_w,    \
-        cos_sin_cache, positions, slot_mapping, index_slot_mapping, kv_cache, \
-        index_cache, eps, rotary_dim, num_tokens, nq, nkv, niq, block_size,   \
-        kv_s_block, kv_s_kv, kv_s_token, kv_s_head)
+  #define LAUNCH(IS_SPARSE, INSERT, FP8, OUT_T)                                \
+    cudaLaunchKernelEx(                                                        \
+        &config,                                                               \
+        fusedMiniMaxM3QNormRopeKVInsertKernel<scalar_t, cache_t, kv_dt, OUT_T, \
+                                              IS_SPARSE, INSERT, FP8>,         \
+        qkv, q_out, reinterpret_cast<OUT_T*>(index_q_out), q_norm_w, k_norm_w, \
+        iq_norm_w, ik_norm_w, cos_sin_cache, positions, slot_mapping,          \
+        index_slot_mapping, kv_cache, reinterpret_cast<OUT_T*>(index_cache),   \
+        eps, rotary_dim, num_tokens, nq, nkv, niq, block_size, kv_s_block,     \
+        kv_s_kv, kv_s_token, kv_s_head)
 #else
  // ROCm: standard kernel launch syntax (no PDL/stream serialization).
  // clang-format off
-  #define LAUNCH(IS_SPARSE, INSERT)                                          \
-    fusedMiniMaxM3QNormRopeKVInsertKernel<scalar_t, cache_t, kv_dt,          \
-                                          IS_SPARSE, INSERT>                 \
+  #define LAUNCH(IS_SPARSE, INSERT, FP8, OUT_T)                              \
+    fusedMiniMaxM3QNormRopeKVInsertKernel<scalar_t, cache_t, kv_dt, OUT_T,   \
+                                          IS_SPARSE, INSERT, FP8>            \
        <<<grid, kBlockSize, 0, stream>>>(                                   \
-            qkv, q_out, index_q_out, q_norm_w, k_norm_w, iq_norm_w,          \
-            ik_norm_w, cos_sin_cache, positions, slot_mapping,               \
-            index_slot_mapping, kv_cache, index_cache, eps, rotary_dim,      \
+            qkv, q_out, reinterpret_cast<OUT_T*>(index_q_out), q_norm_w,     \
+            k_norm_w, iq_norm_w, ik_norm_w, cos_sin_cache, positions,        \
+            slot_mapping, index_slot_mapping, kv_cache,                      \
+            reinterpret_cast<OUT_T*>(index_cache), eps, rotary_dim,          \
            num_tokens, nq, nkv, niq, block_size, kv_s_block, kv_s_kv,       \
            kv_s_token, kv_s_head)
  // clang-format on
@@ -466,14 +536,22 @@ void launchFusedMiniMaxM3(scalar_t* qkv, scalar_t* q_out, scalar_t* index_q_out,

  if (has_index) {
    if (insert_kv) {
-      LAUNCH(true, true);  // sparse serving
+      if (fp8_idx) {
+        LAUNCH(true, true, true, uint8_t);  // sparse serving, fp8 index outputs
+      } else {
+        LAUNCH(true, true, false, scalar_t);  // sparse serving, bf16
+      }
    } else {
-      LAUNCH(true, false);  // sparse profiling
+      if (fp8_idx) {
+        LAUNCH(true, false, true, uint8_t);  // sparse profiling, fp8 index_q
+      } else {
+        LAUNCH(true, false, false, scalar_t);  // sparse profiling, bf16
+      }
    }
  } else {
    // Dense layer: never has an index branch and never inserts here (the
    // generic Attention layer owns the KV insert).
-    LAUNCH(false, false);
+    LAUNCH(false, false, false, scalar_t);
  }
 #undef LAUNCH
 }
@@ -485,8 +563,9 @@ void launchFusedMiniMaxM3(scalar_t* qkv, scalar_t* q_out, scalar_t* index_q_out,
  vllm::minimax_m3_fused_ops::launchFusedMiniMaxM3<st, CACHE_T, KV_DTYPE>(     \
      reinterpret_cast<st*>(qkv.data_ptr()),                                   \
      q_out.has_value() ? reinterpret_cast<st*>(q_out->data_ptr()) : nullptr,  \
-      index_q_out.has_value() ? reinterpret_cast<st*>(index_q_out->data_ptr()) \
-                              : nullptr,                                       \
+      index_q_out.has_value()                                                  \
+          ? reinterpret_cast<void*>(index_q_out->data_ptr())                   \
+          : nullptr,                                                           \
      reinterpret_cast<st const*>(q_norm_weight.data_ptr()),                   \
      reinterpret_cast<st const*>(k_norm_weight.data_ptr()),                   \
      has_index ? reinterpret_cast<st const*>(index_q_norm_weight->data_ptr()) \
@@ -502,11 +581,11 @@ void launchFusedMiniMaxM3(scalar_t* qkv, scalar_t* q_out, scalar_t* index_q_out,
                : nullptr,                                                     \
      insert_kv ? reinterpret_cast<CACHE_T*>(kv_cache->data_ptr()) : nullptr,  \
      (insert_kv && has_index)                                                 \
-          ? reinterpret_cast<st*>(index_cache->data_ptr())                     \
+          ? reinterpret_cast<void*>(index_cache->data_ptr())                   \
          : nullptr,                                                           \
      static_cast<float>(eps), static_cast<int>(rotary_dim), num_tokens, nq,   \
      nkv, niq, static_cast<int>(block_size), kv_s_block, kv_s_kv, kv_s_token, \
-      kv_s_head, has_index, insert_kv, stream)
+      kv_s_head, has_index, insert_kv, fp8_idx, stream)

 // ────────────────────────────────────────────────────────────────────────────
 // Torch op wrapper
@@ -612,6 +691,7 @@ void fused_minimax_m3_qknorm_rope_kv_insert(
                 torch::headeronly::ScalarType::Long &&
             index_slot_mapping->numel() == slot_mapping->numel()),
        "index_slot_mapping must be int64 CUDA with slot_mapping length");
+    // Main attention KV cache: auto matches qkv, fp8 uses uint8 storage.
    if (kv_dt == vllm::Fp8KVCacheDataType::kAuto) {
      STD_TORCH_CHECK(kv_cache->scalar_type() == qkv.scalar_type(),
                      "auto kv_cache dtype must match qkv");
@@ -620,9 +700,13 @@ void fused_minimax_m3_qknorm_rope_kv_insert(
          kv_cache->scalar_type() == torch::headeronly::ScalarType::Byte,
          "fp8 kv_cache must use uint8 storage");
    }
-    STD_TORCH_CHECK(index_cache.has_value() &&
-                        index_cache->scalar_type() == qkv.scalar_type(),
-                    "insert mode requires matching index_cache");
+    // Indexer index-K cache: independent dtype -- qkv dtype or fp8 e4m3.
+    STD_TORCH_CHECK(
+        index_cache.has_value() &&
+            (index_cache->scalar_type() == qkv.scalar_type() ||
+             index_cache->scalar_type() ==
+                 torch::headeronly::ScalarType::Float8_e4m3fn),
+        "insert mode requires index_cache matching qkv dtype or fp8 e4m3");
    STD_TORCH_CHECK(kv_cache->dim() == 5 && kv_cache->stride(4) == 1,
                    "kv_cache must be [nb,2,bs,nkv,head_dim] with contiguous "
                    "head_dim (stride(4)==1)");
@@ -652,14 +736,31 @@ void fused_minimax_m3_qknorm_rope_kv_insert(
        "index_q_out requires the index branch (num_index_heads > 0)");
    STD_TORCH_CHECK(
        index_q_out->is_cuda() && index_q_out->is_contiguous() &&
-            index_q_out->scalar_type() == qkv.scalar_type(),
-        "index_q_out must be a contiguous CUDA tensor matching qkv dtype");
+            (index_q_out->scalar_type() == qkv.scalar_type() ||
+             index_q_out->scalar_type() ==
+                 torch::headeronly::ScalarType::Float8_e4m3fn),
+        "index_q_out must be contiguous CUDA, qkv dtype or fp8 e4m3");
    STD_TORCH_CHECK(index_q_out->numel() ==
                        static_cast<int64_t>(num_tokens) * niq * kHeadDim,
                    "index_q_out must have num_tokens * num_index_heads * 128 "
                    "elements");
  }

+  // fp8 index path: the index-K cache and index-Q outputs are e4m3 bytes while
+  // q/k/v + q_out stay qkv dtype. Both index outputs must agree.
+  auto const kFp8 = torch::headeronly::ScalarType::Float8_e4m3fn;
+  bool const fp8_idx =
+      (index_cache.has_value() && index_cache->scalar_type() == kFp8) ||
+      (index_q_out.has_value() && index_q_out->scalar_type() == kFp8);
+  if (fp8_idx) {
+    STD_TORCH_CHECK(
+        !index_cache.has_value() || index_cache->scalar_type() == kFp8,
+        "fp8 index path: index_cache must be fp8 e4m3");
+    STD_TORCH_CHECK(
+        !index_q_out.has_value() || index_q_out->scalar_type() == kFp8,
+        "fp8 index path: index_q_out must be fp8 e4m3");
+  }
+
  const torch::stable::accelerator::DeviceGuard device_guard(
      qkv.get_device_index());
  auto stream = get_current_cuda_stream(qkv.get_device_index());
@@ -792,7 +792,7 @@ RUN --mount=type=cache,target=/opt/uv/cache \
 # Install FlashInfer JIT cache (requires CUDA-version-specific index URL)
 # https://docs.flashinfer.ai/installation.html
 # From versions.json: .flashinfer.version
-ARG FLASHINFER_VERSION=0.6.12
+ARG FLASHINFER_VERSION=0.6.13rc2
 RUN --mount=type=cache,target=/opt/uv/cache \
    uv pip install --system flashinfer-jit-cache==${FLASHINFER_VERSION} \
        --index-url https://flashinfer.ai/whl/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.')
@@ -68,7 +68,7 @@
      "default": "true"
    },
    "FLASHINFER_VERSION": {
-      "default": "0.6.12"
+      "default": "0.6.13rc2"
    },
    "GDRCOPY_CUDA_VERSION": {
      "default": "12.8"
@@ -9,8 +9,8 @@ torchaudio==2.11.0
 # These must be updated alongside torch
 torchvision==0.26.0 # Required for phi3v processor. See https://github.com/pytorch/vision?tab=readme-ov-file#installation for corresponding version
 # FlashInfer should be updated together with the Dockerfile
-flashinfer-python==0.6.12
-flashinfer-cubin==0.6.12
+flashinfer-python==0.6.13rc2
+flashinfer-cubin==0.6.13rc2
 apache-tvm-ffi==0.1.9
 tilelang==0.1.9
 nvidia-cudnn-frontend>=1.19.1
@@ -1171,7 +1171,15 @@ package_data = {
        "third_party/deep_gemm/include/**/*.h",
        "third_party/deep_gemm/include/**/*.hpp",
        # fmha_sm100 sparse CuTe-DSL helper kernels (vendored via cmake)
+        "third_party/fmha_sm100/csrc/**/*.cu",
+        "third_party/fmha_sm100/csrc/**/*.h",
+        "third_party/fmha_sm100/csrc/**/*.jinja",
+        "third_party/fmha_sm100/csrc/**/*.cu.jinja",
        "third_party/fmha_sm100/cute/**/*.cu",
+        "third_party/fmha_sm100/cutlass/include/**/*.h",
+        "third_party/fmha_sm100/cutlass/include/**/*.hpp",
+        "third_party/fmha_sm100/cutlass/tools/util/include/**/*.h",
+        "third_party/fmha_sm100/cutlass/tools/util/include/**/*.hpp",
    ]
 }

@@ -134,6 +134,7 @@ def _reference_index_topk(
    topk: int,
    init_blocks: int,
    local_blocks: int,
+    sm_scale: float = 1.0,
 ) -> torch.Tensor:
    total_q, num_idx_heads, _ = idx_q.shape
    out = torch.full(
@@ -149,7 +150,7 @@ def _reference_index_topk(
        num_blocks = (seq_len + BLOCK_SIZE - 1) // BLOCK_SIZE
        pages = block_table[req_id, :num_blocks]
        k = index_kv_cache[pages].reshape(num_blocks * BLOCK_SIZE, -1)
-        score = torch.einsum("qhd,kd->hqk", q.float(), k.float())
+        score = torch.einsum("qhd,kd->hqk", q.float(), k.float()) * sm_scale

        q_pos = prefix_len + torch.arange(q_len, device=idx_q.device)
        k_pos = torch.arange(k.shape[0], device=idx_q.device)
@@ -244,6 +245,270 @@ def test_prefill_index_topk_correctness():
    _assert_topk_indices_equal_unordered(actual, expected)


+# MSA indexer (SM100): fmha_sm100 OnlyScore for the per-block scores, then the
+# Triton minimax_m3_index_topk for selection (no sparse_topk_select). Uses a
+# deterministic construction (idx_q == 1, distinct e4m3-exact per-block values)
+# so scores are strictly monotonic in the block id -> exact top-k agreement.
+def _fmha_indexer_topk(
+    idx_q: torch.Tensor,  # [total_q, H, 128] bf16/e4m3
+    index_cache: torch.Tensor,  # [num_pages, 128, 128] bf16/e4m3
+    block_table: torch.Tensor,
+    q_lens: torch.Tensor,
+    seq_lens: torch.Tensor,
+    prefix_lens: torch.Tensor,
+    sm_scale: float,
+    topk: int,
+) -> torch.Tensor:
+    """Replicate MiniMaxM3IndexerMSAImpl's score path (single decode/prefill side)."""
+    from vllm.third_party.fmha_sm100.api import _fmha_sm100, _fmha_sm100_plan
+
+    num_idx_heads, head_dim = idx_q.shape[1], idx_q.shape[2]
+    nvp = [(s + 127) // 128 for s in seq_lens.tolist()]
+    kv_indices = torch.cat([block_table[r, : nvp[r]] for r in range(len(nvp))]).to(
+        torch.int32
+    )
+
+    qo = q_lens.cpu().to(torch.int32)
+    kv = seq_lens.cpu().to(torch.int32)
+    plan = _fmha_sm100_plan(
+        qo,
+        kv,
+        num_idx_heads,
+        num_kv_heads=1,
+        qo_offset=kv - qo,
+        page_size=128,
+        output_maxscore=True,
+        causal=True,
+        num_kv_splits=1,
+    )
+    k_pages = index_cache.view(index_cache.shape[0], 1, 128, head_dim)
+    _, max_score = _fmha_sm100(
+        idx_q,
+        k_pages,
+        k_pages,
+        plan,
+        kv_indices=kv_indices,
+        output_o=False,
+        output_maxscore=True,
+        sm_scale=sm_scale,
+    )
+
+    batch = q_lens.numel()
+    cu = torch.zeros(batch + 1, dtype=torch.int32, device=idx_q.device)
+    cu[1:] = q_lens.to(torch.int32).cumsum(0)
+    # max_score [H, k_tiles, total_q] -> transpose to [H, total_q, k_tiles].
+    return minimax_m3_index_topk(
+        max_score.transpose(1, 2),
+        cu,
+        prefix_lens.to(torch.int32),
+        int(q_lens.max()),
+        topk,
+        0,  # init_blocks
+        0,  # local_blocks
+    )
+
+
+# e4m3-exact, strictly-increasing per-block values: with idx_q == 1 (also exact)
+# the per-block scores are exact and distinct in BOTH bf16 and e4m3, so the fp8
+# score path selects the same top-k as the reference (no quantization ties).
+_E4M3_EXACT_VALUES = [
+    *range(1, 17),  # 1..16  (step 1)
+    *range(18, 33, 2),  # 18..32 (step 2)
+    *range(36, 65, 4),  # 36..64 (step 4)
+    *range(72, 129, 8),  # 72..128 (step 8)
+]
+
+
+@pytest.mark.skipif(
+    not current_platform.is_device_capability_family(100),
+    reason="fmha_sm100 indexer requires SM100 (Blackwell).",
+)
+@pytest.mark.parametrize("index_dtype", [torch.bfloat16, torch.float8_e4m3fn])
+@pytest.mark.parametrize(
+    ("q_lens", "prefix_lens"),
+    [
+        ((4, 3), (2048, 2560)),  # prefill: every token sees >= 16 causal blocks
+        ((1, 1, 1), (2048, 3000, 4096)),  # decode: one query token per request
+    ],
+)
+def test_fmha_sm100_indexer_matches_reference(q_lens, prefix_lens, index_dtype):
+    torch.manual_seed(0)
+    num_idx_heads, head_dim = 4, HEAD_DIM
+    device = "cuda"
+
+    q_lens_t = torch.tensor(q_lens, device=device, dtype=torch.int32)
+    prefix_lens_t = torch.tensor(prefix_lens, device=device, dtype=torch.int32)
+    seq_lens = prefix_lens_t + q_lens_t
+    batch = len(q_lens)
+    max_blocks = (int(seq_lens.max()) + BLOCK_SIZE - 1) // BLOCK_SIZE
+    assert max_blocks <= len(_E4M3_EXACT_VALUES)
+    num_pages = batch * max_blocks
+    block_table = torch.randperm(num_pages, device=device, dtype=torch.int32).reshape(
+        batch, max_blocks
+    )
+
+    idx_q = torch.ones(
+        int(q_lens_t.sum()), num_idx_heads, head_dim, device=device, dtype=index_dtype
+    )
+    index_cache = torch.empty(
+        num_pages, BLOCK_SIZE, head_dim, device=device, dtype=index_dtype
+    )
+    for r in range(batch):
+        for b in range(max_blocks):
+            index_cache[block_table[r, b]] = float(_E4M3_EXACT_VALUES[b])
+
+    sm_scale = head_dim**-0.5
+    actual = _fmha_indexer_topk(
+        idx_q,
+        index_cache,
+        block_table,
+        q_lens_t,
+        seq_lens,
+        prefix_lens_t,
+        sm_scale,
+        TOPK,
+    )
+    expected = _reference_index_topk(
+        idx_q,
+        index_cache,
+        block_table,
+        q_lens_t,
+        seq_lens,
+        prefix_lens_t,
+        TOPK,
+        init_blocks=0,
+        local_blocks=0,
+        sm_scale=sm_scale,
+    )
+    _assert_topk_indices_equal_unordered(actual, expected)
+
+
+# Full impl-level parity: drive both MiniMaxM3IndexerMSAImpl (fmha_sm100 score +
+# Triton top-k) and MiniMaxM3IndexerTritonImpl through their real metadata
+# builders on the SAME CommonAttentionMetadata + index cache, and assert the
+# selected blocks agree. This exercises all the metadata the impl/kernels consume
+# (decode/prefill split, cu_seqlens_q rebasing, prefix_lens, kv_indices gather,
+# decode_pages split) -- a metadata bug on either side shifts the causal window
+# or the block->page mapping and breaks the comparison.
+@pytest.mark.skipif(
+    not current_platform.is_device_capability_family(100),
+    reason="fmha_sm100 indexer requires SM100 (Blackwell).",
+)
+@pytest.mark.parametrize("topk", [8, 16])
+def test_msa_indexer_impl_matches_triton(topk, monkeypatch):
+    import vllm.models.minimax_m3.common.indexer as indexer_mod
+    from tests.v1.attention.utils import (
+        BatchSpec,
+        create_common_attn_metadata,
+        create_vllm_config,
+    )
+    from vllm.config import set_current_vllm_config
+    from vllm.forward_context import set_forward_context
+    from vllm.models.minimax_m3.common.indexer import (
+        MiniMaxM3IndexerTritonImpl,
+        MiniMaxM3IndexerTritonMetadataBuilder,
+    )
+    from vllm.models.minimax_m3.nvidia.indexer_msa import (
+        MiniMaxM3IndexerMSAImpl,
+        MiniMaxM3IndexerMSAMetadataBuilder,
+    )
+
+    torch.manual_seed(0)
+    device = torch.device("cuda")
+    num_idx_heads, head_dim = 4, HEAD_DIM
+    # TP=1: avoid requiring an initialized distributed group in a unit test.
+    monkeypatch.setattr(indexer_mod, "get_tensor_model_parallel_world_size", lambda: 1)
+
+    vllm_config = create_vllm_config(
+        block_size=BLOCK_SIZE, max_model_len=8192, max_num_batched_tokens=8192
+    )
+    vllm_config.model_config.hf_config.sparse_attention_config = {
+        "sparse_num_index_heads": num_idx_heads
+    }
+
+    # Decode-first mixed batch: 2 decode reqs (q_len 1) then 2 prefill reqs. Long
+    # prefixes so every token sees > TOPK causal blocks (non-trivial selection).
+    batch = BatchSpec(seq_lens=[2305, 2561, 2624, 2720], query_lens=[1, 1, 64, 96])
+    common = create_common_attn_metadata(
+        batch, BLOCK_SIZE, device, arange_block_indices=True
+    )
+    num_tokens = batch.compute_num_tokens()
+
+    # Deterministic index cache: distinct, monotonic per-logical-block values so
+    # the top-k is unambiguous (both kernels pick the same blocks, no fp ties).
+    block_table = common.block_table_tensor
+    num_pages = int(block_table.max().item()) + 1
+    index_cache = torch.zeros(
+        num_pages, BLOCK_SIZE, head_dim, device=device, dtype=DTYPE
+    )
+    for r, seq_len in enumerate(batch.seq_lens):
+        for b in range((seq_len + BLOCK_SIZE - 1) // BLOCK_SIZE):
+            index_cache[block_table[r, b]] = float(b + 1)
+    index_q = torch.ones(
+        num_tokens, num_idx_heads * head_dim, device=device, dtype=DTYPE
+    )
+
+    spec = MLAAttentionSpec(
+        block_size=BLOCK_SIZE, num_kv_heads=1, head_size=head_dim, dtype=DTYPE
+    )
+    impl_kwargs = dict(
+        num_kv_heads=num_idx_heads,
+        scale=head_dim**-0.5,
+        topk_blocks=topk,
+        sparse_block_size=BLOCK_SIZE,
+        num_index_heads=num_idx_heads,
+        index_head_dim=head_dim,
+        init_blocks=0,
+        local_blocks=0,
+    )
+
+    with set_current_vllm_config(vllm_config):
+        msa_impl = MiniMaxM3IndexerMSAImpl(prefix="idx_msa", **impl_kwargs)
+        triton_impl = MiniMaxM3IndexerTritonImpl(prefix="idx_triton", **impl_kwargs)
+        msa_builder = MiniMaxM3IndexerMSAMetadataBuilder(
+            spec, [msa_impl.index_cache.prefix], vllm_config, device
+        )
+        triton_builder = MiniMaxM3IndexerTritonMetadataBuilder(
+            spec, [triton_impl.index_cache.prefix], vllm_config, device
+        )
+
+    # Both impls score against the same index keys.
+    msa_impl.index_cache.kv_cache = index_cache
+    triton_impl.index_cache.kv_cache = index_cache
+
+    # Exercise the shared persistent top-k buffer for BOTH impls: each must write
+    # decode ([:, :nd]) and prefill ([:, nd:]) into its buffer and return views.
+    # Separate buffers so the two forwards don't clobber each other.
+    nd = sum(q for q in batch.query_lens if q <= 1)
+    msa_impl.topk_indices_buffer = torch.full(
+        (num_idx_heads, num_tokens, topk), -2, dtype=torch.int32, device=device
+    )
+    triton_impl.topk_indices_buffer = torch.full(
+        (num_idx_heads, num_tokens, topk), -2, dtype=torch.int32, device=device
+    )
+
+    attn_metadata = {
+        msa_impl.index_cache.prefix: msa_builder.build(0, common),
+        triton_impl.index_cache.prefix: triton_builder.build(0, common),
+    }
+    with set_forward_context(attn_metadata, vllm_config):
+        msa_decode, msa_prefill = msa_impl(index_q)
+        tri_decode, tri_prefill = triton_impl(index_q)
+
+    assert msa_decode is not None and tri_decode is not None
+    assert msa_prefill is not None and tri_prefill is not None
+    _assert_topk_indices_equal_unordered(msa_decode, tri_decode)
+    _assert_topk_indices_equal_unordered(msa_prefill, tri_prefill)
+    # decode/prefill outputs are views into each impl's persistent buffer.
+    for impl, dec, pre in (
+        (msa_impl, msa_decode, msa_prefill),
+        (triton_impl, tri_decode, tri_prefill),
+    ):
+        buf = impl.topk_indices_buffer
+        assert dec.data_ptr() == buf[:, :nd, :].data_ptr()
+        assert pre.data_ptr() == buf[:, nd:, :].data_ptr()
+
+
@pytest.mark.parametrize(
    ("decode_query_len", "max_decode_query_len"),
    [
@@ -317,6 +582,65 @@ def test_decode_index_topk_correctness(
    _assert_topk_indices_equal_unordered(actual, expected)


+@pytest.mark.skipif(
+    not current_platform.is_device_capability_family(100),
+    reason="fp8 e4m3 indexer cache is the SM100 (MSA) path.",
+)
+@pytest.mark.parametrize("num_idx_heads", [1, 4])
+def test_decode_index_topk_fp8(num_idx_heads: int):
+    """The fp8 (e4m3) indexer cache feeds the Triton decode kernel on the MSA
+    path. The kernel must score in fp32 (no scaling) so its top-k matches a
+    reference computed from the dequantized fp8 values."""
+    torch.manual_seed(0)
+    topk, init_blocks, local_blocks, head_dim = 8, 0, 1, 128
+    decode_query_len = 1
+    active_seq_lens = torch.tensor((129, 1025, 4097), device="cuda", dtype=torch.int32)
+    q_lens = torch.full_like(active_seq_lens, decode_query_len)
+    prefix_lens = active_seq_lens - decode_query_len
+    batch = active_seq_lens.numel()
+    max_seq_len = int(active_seq_lens.max())
+    max_blocks = (max_seq_len + BLOCK_SIZE - 1) // BLOCK_SIZE
+    num_pages = batch * max_blocks
+    block_table = torch.randperm(num_pages, device="cuda", dtype=torch.int32).reshape(
+        batch, max_blocks
+    )
+    idx_q = torch.randn(
+        batch * decode_query_len, num_idx_heads, head_dim, device="cuda"
+    ).to(torch.float8_e4m3fn)
+    index_kv_cache = torch.randn(num_pages, BLOCK_SIZE, head_dim, device="cuda").to(
+        torch.float8_e4m3fn
+    )
+
+    actual = minimax_m3_index_decode(
+        idx_q,
+        index_kv_cache,
+        block_table,
+        active_seq_lens,
+        max_seq_len=max_seq_len,
+        topk=topk,
+        init_blocks=init_blocks,
+        local_blocks=local_blocks,
+        num_kv_heads=num_idx_heads,
+        sm_scale=head_dim**-0.5,
+        decode_query_len=decode_query_len,
+    )
+    # Reference from the DEQUANTIZED fp8 values (the kernel computes the fp8 QK
+    # in fp32, so it must match an fp32 matmul of the same e4m3 values).
+    expected = _reference_index_topk(
+        idx_q.float(),
+        index_kv_cache.float(),
+        block_table,
+        q_lens,
+        active_seq_lens,
+        prefix_lens,
+        topk,
+        init_blocks,
+        local_blocks,
+        head_dim**-0.5,
+    )
+    _assert_topk_indices_equal_unordered(actual, expected)
+
+
 # Sparse attention kernels.
 def _reference_sparse_attn(
    q: torch.Tensor,
@@ -5,8 +5,9 @@ Tests for the FlashInfer TRTLLM NvFP4 MoE backend
 (`TrtLlmNvFp4ExpertsModular`).

 Covers the activations the wrapper claims to support — SiLU, RELU^2 (non-gated),
-and GELU — including a Gemma4-shaped case (128 experts, top-k 8,
-intermediate_size 704) that exercises the non-256-aligned padding path.
+GELU, and clamped SwiGLU-OAI (MiniMax-M3) — including a Gemma4-shaped case
+(128 experts, top-k 8, intermediate_size 704) that exercises the non-256-aligned
+padding path.
 """

 import pytest
@@ -80,6 +81,29 @@ if _CLAMP_OP_NAME not in op_registry:

 SILU_WITH_CLAMP = op_registry[_CLAMP_OP_NAME]

+# Clamped SwiGLU-OAI (MiniMax-M3): non-default alpha/beta so the kernel must
+# honor gemm1_alpha (raw) and gemm1_beta (folded by g1_alphas), not just clamp.
+_SWIGLU_ALPHA = 1.702
+_SWIGLU_BETA = 1.0
+_OAI_OP_NAME = "test_swigluoai_with_clamp"
+
+if _OAI_OP_NAME not in op_registry:
+
+    @CustomOp.register(_OAI_OP_NAME)
+    class _SwigluOAIWithClampTest(SiluAndMulWithClamp):
+        custom_op_name = _OAI_OP_NAME
+
+        def __init__(self, *, compile_native: bool = True) -> None:
+            super().__init__(
+                _SWIGLU_LIMIT,
+                alpha=_SWIGLU_ALPHA,
+                beta=_SWIGLU_BETA,
+                compile_native=compile_native,
+            )
+
+
+SWIGLUOAI_REF = op_registry[_OAI_OP_NAME]
+

 ACTIVATION_CASES = [
    pytest.param(MoEActivation.SILU, MoEActivation.SILU, None, id="silu"),
@@ -91,6 +115,12 @@ ACTIVATION_CASES = [
        id="relu2_no_mul",
    ),
    pytest.param(MoEActivation.GELU, MoEActivation.GELU, None, id="gelu"),
+    pytest.param(
+        MoEActivation.SWIGLUOAI_UNINTERLEAVE,
+        SWIGLUOAI_REF,
+        _SWIGLU_LIMIT,
+        id="swigluoai_uninterleave",
+    ),
 ]


@@ -148,6 +178,10 @@ def test_trtllm_fp4_moe_no_graph(
            is_scale_swizzled=False,
        )
        quant_config.gemm1_clamp_limit = swiglu_limit
+        is_oai = activation == MoEActivation.SWIGLUOAI_UNINTERLEAVE
+        if is_oai:
+            quant_config.gemm1_alpha = _SWIGLU_ALPHA
+            quant_config.gemm1_beta = _SWIGLU_BETA
        if swiglu_limit is not None:
            assert quant_config.g1_alphas is not None
            assert quant_config.a2_gscale is not None
@@ -192,6 +226,27 @@ def test_trtllm_fp4_moe_no_graph(
        fake_layer.w2_input_scale = torch.ones_like(quant_config.g2_alphas)
        trtllm_inner.process_weights_after_loading(fake_layer)

+        if is_oai:
+            # alpha stays raw; beta and clamp are folded by g1_alphas
+            # (== _LARGE_OUTPUT1_SCALE here), so the fold is load-bearing.
+            assert torch.allclose(
+                trtllm_inner.gemm1_alpha,
+                torch.full_like(trtllm_inner.gemm1_alpha, _SWIGLU_ALPHA),
+            )
+            assert torch.allclose(
+                trtllm_inner.gemm1_beta,
+                torch.full_like(
+                    trtllm_inner.gemm1_beta, _SWIGLU_BETA / _LARGE_OUTPUT1_SCALE
+                ),
+            )
+            assert torch.allclose(
+                trtllm_inner.gemm1_clamp_limit,
+                torch.full_like(
+                    trtllm_inner.gemm1_clamp_limit,
+                    _SWIGLU_LIMIT / _LARGE_OUTPUT1_SCALE,
+                ),
+            )
+
        trtllm_experts = mk.FusedMoEKernel(
            maybe_make_prepare_finalize(
                moe=moe_config,
@@ -278,3 +278,99 @@ def test_sparse_full(num_tokens, block_size, kv_cache_dtype):
    torch.testing.assert_close(
        index_cache.view(-1, HEAD_DIM), expected_index_cache, rtol=0, atol=0
    )
+
+
+# ── Test 3: fp8 (e4m3) index outputs ─────────────────────────────────────────
+# The fp8 score path stores index_q and the index-K cache as e4m3 while q/k/v +
+# q_out stay bf16. Asserts: (1) q/k/v/q_out are bit-identical to the bf16 run
+# (the index dtype must not perturb the main branch), and (2) the e4m3 index
+# outputs dequantize close to the bf16 reference.
+
+
+@pytest.mark.skipif(
+    not torch.cuda.is_available() or torch.cuda.get_device_capability() < (8, 9),
+    reason="e4m3 conversion requires CUDA SM89+.",
+)
+@pytest.mark.parametrize("num_tokens", [1, 7, 64, 513])
+@pytest.mark.parametrize("block_size", [16, 64])
+def test_sparse_full_fp8_index(num_tokens, block_size):
+    torch.manual_seed(1)
+    device, dtype, eps = "cuda", torch.bfloat16, 1e-6
+    base, max_pos = 5_000_000.0, 4096
+    num_heads, num_kv_heads, num_idx_heads = 16, 4, 4
+
+    q_w = torch.randn(HEAD_DIM, dtype=dtype, device=device) * 0.1
+    k_w = torch.randn(HEAD_DIM, dtype=dtype, device=device) * 0.1
+    iq_w = torch.randn(HEAD_DIM, dtype=dtype, device=device) * 0.1
+    ik_w = torch.randn(HEAD_DIM, dtype=dtype, device=device) * 0.1
+    cos_sin = make_cos_sin_cache(max_pos, ROTARY_DIM, base, dtype, device)
+    positions = torch.randint(
+        0, max_pos, (num_tokens,), dtype=torch.int64, device=device
+    )
+
+    qsz, kvsz = num_heads * HEAD_DIM, num_kv_heads * HEAD_DIM
+    iqsz, iksz = num_idx_heads * HEAD_DIM, HEAD_DIM
+    qkv0 = torch.randn(
+        num_tokens, qsz + 2 * kvsz + iqsz + iksz, dtype=dtype, device=device
+    )
+
+    num_blocks = (num_tokens + block_size - 1) // block_size + 1
+    slot_mapping = torch.randperm(
+        num_blocks * block_size, dtype=torch.int64, device=device
+    )[:num_tokens]
+    index_slot_mapping = torch.roll(slot_mapping, shifts=1)
+
+    def run(index_dtype):
+        qkv = qkv0.clone()
+        kv_cache = torch.zeros(
+            num_blocks,
+            2,
+            block_size,
+            num_kv_heads,
+            HEAD_DIM,
+            dtype=dtype,
+            device=device,
+        )
+        index_cache = torch.zeros(
+            num_blocks, block_size, HEAD_DIM, dtype=index_dtype, device=device
+        )
+        q_out = torch.empty(num_tokens, qsz, dtype=dtype, device=device)
+        index_q = torch.empty(num_tokens, iqsz, dtype=index_dtype, device=device)
+        ops.fused_minimax_m3_qknorm_rope_kv_insert(
+            qkv,
+            q_w,
+            k_w,
+            cos_sin,
+            positions,
+            num_heads,
+            num_kv_heads,
+            ROTARY_DIM,
+            eps,
+            iq_w,
+            ik_w,
+            num_idx_heads,
+            slot_mapping,
+            index_slot_mapping,
+            kv_cache,
+            index_cache,
+            block_size,
+            q_out,
+            index_q,
+        )
+        return qkv, kv_cache, index_cache, q_out, index_q
+
+    qkv_bf, kvc_bf, idxc_bf, qo_bf, iq_bf = run(torch.bfloat16)
+    qkv_fp, kvc_fp, idxc_fp, qo_fp, iq_fp = run(torch.float8_e4m3fn)
+
+    assert iq_fp.dtype == torch.float8_e4m3fn
+    assert idxc_fp.dtype == torch.float8_e4m3fn
+
+    # (1) The main branch (q/k/v in qkv, q_out, kv cache) must be bit-identical:
+    # the index output dtype must not perturb anything else.
+    torch.testing.assert_close(qo_fp, qo_bf, rtol=0, atol=0)
+    torch.testing.assert_close(qkv_fp, qkv_bf, rtol=0, atol=0)
+    torch.testing.assert_close(kvc_fp, kvc_bf, rtol=0, atol=0)
+
+    # (2) Dequantized e4m3 index outputs match the bf16 reference within fp8 ulp.
+    torch.testing.assert_close(iq_fp.float(), iq_bf.float(), rtol=0.13, atol=0.05)
+    torch.testing.assert_close(idxc_fp.float(), idxc_bf.float(), rtol=0.13, atol=0.05)
@@ -1579,10 +1579,9 @@ environment_variables: dict[str, Callable[[], Any]] = {
        os.getenv("VLLM_TOOL_PARSE_REGEX_TIMEOUT_SECONDS", "1")
    ),
    # Enforce function parameter schemas in structural-tag based tool calling.
-    "VLLM_ENFORCE_STRICT_TOOL_CALLING": lambda: os.getenv(
-        "VLLM_ENFORCE_STRICT_TOOL_CALLING", "True"
-    ).lower()
-    in ("true", "1"),
+    "VLLM_ENFORCE_STRICT_TOOL_CALLING": lambda: (
+        os.getenv("VLLM_ENFORCE_STRICT_TOOL_CALLING", "True").lower() in ("true", "1")
+    ),
    # Control the max chunk bytes (in MB) for the rpc message queue.
    # Object larger than this threshold will be broadcast to worker
    # processes via zmq.
@@ -56,6 +56,35 @@ class TrtLlmFp8ExpertsBase:
        self.moe_config = moe_config
        self.quant_config = quant_config

+        # Per-expert SwiGLU parameters from quant_config (MXFP8 + Swiglu only).
+        device = torch.accelerator.current_device_index()
+        if quant_config.gemm1_alpha is not None:
+            self.gemm1_alpha = torch.tensor(
+                [quant_config.gemm1_alpha] * self.local_num_experts,
+                dtype=torch.float32,
+                device=device,
+            )
+        else:
+            self.gemm1_alpha = None
+
+        if quant_config.gemm1_beta is not None:
+            self.gemm1_beta = torch.tensor(
+                [quant_config.gemm1_beta] * self.local_num_experts,
+                dtype=torch.float32,
+                device=device,
+            )
+        else:
+            self.gemm1_beta = None
+
+        if quant_config.gemm1_clamp_limit is not None:
+            self.gemm1_clamp_limit = torch.tensor(
+                [quant_config.gemm1_clamp_limit] * self.local_num_experts,
+                dtype=torch.float32,
+                device=device,
+            )
+        else:
+            self.gemm1_clamp_limit = None
+
    @staticmethod
    def activation_format() -> mk.FusedMoEActivationFormat:
        return mk.FusedMoEActivationFormat.Standard
@@ -77,8 +106,12 @@ class TrtLlmFp8ExpertsBase:

    @staticmethod
    def _supports_activation(activation: MoEActivation) -> bool:
-        """Supports only SiLU and RELU^2 non-gated activation."""
-        return activation in [MoEActivation.SILU, MoEActivation.RELU2_NO_MUL]
+        """Supports SiLU, SwiGLU-OAI (uninterleaved), and RELU^2 non-gated."""
+        return activation in [
+            MoEActivation.SILU,
+            MoEActivation.SWIGLUOAI_UNINTERLEAVE,
+            MoEActivation.RELU2_NO_MUL,
+        ]

    @staticmethod
    def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool:
@@ -198,6 +231,9 @@ class TrtLlmFp8ExpertsModular(TrtLlmFp8ExpertsBase, mk.FusedMoEExpertsModular):
            hidden_states_scale=hidden_states_scale,
            gemm1_weights=w1,
            gemm1_weights_scale=self.quant_config.w1_scale,
+            gemm1_alpha=self.gemm1_alpha,
+            gemm1_beta=self.gemm1_beta,
+            gemm1_clamp_limit=self.gemm1_clamp_limit,
            gemm2_weights=w2,
            gemm2_weights_scale=self.quant_config.w2_scale,
            num_experts=global_num_experts,
@@ -327,7 +363,11 @@ class TrtLlmFp8ExpertsMonolithic(TrtLlmFp8ExpertsBase, mk.FusedMoEExpertsMonolit
        from flashinfer.fused_moe import Fp8QuantizationType, WeightLayout

        assert not apply_router_weight_on_input
-        assert activation in [MoEActivation.SILU, MoEActivation.RELU2_NO_MUL]
+        assert activation in [
+            MoEActivation.SILU,
+            MoEActivation.SWIGLUOAI_UNINTERLEAVE,
+            MoEActivation.RELU2_NO_MUL,
+        ]
        activation_type = activation_to_flashinfer_int(activation)
        assert self.topk <= global_num_experts
        assert global_num_experts % 4 == 0
@@ -362,6 +402,9 @@ class TrtLlmFp8ExpertsMonolithic(TrtLlmFp8ExpertsBase, mk.FusedMoEExpertsMonolit
            hidden_states_scale=hidden_states_scale,
            gemm1_weights=w1,
            gemm1_weights_scale=self.quant_config.w1_scale,
+            gemm1_alpha=self.gemm1_alpha,
+            gemm1_beta=self.gemm1_beta,
+            gemm1_clamp_limit=self.gemm1_clamp_limit,
            gemm2_weights=w2,
            gemm2_weights_scale=self.quant_config.w2_scale,
            num_experts=global_num_experts,
@@ -66,16 +66,47 @@ class TrtLlmNvFp4ExpertsBase:
        else:
            self.g1_scale_c = self.quant_config.a2_gscale.clone()

-        if moe_config.is_act_and_mul and quant_config.gemm1_clamp_limit is not None:
-            device = torch.accelerator.current_device_index()
-            self.gemm1_clamp_limit = torch.full(
+        # Fall back to moe_config.swiglu_* when quant_config doesn't carry them
+        # (ModelOpt NVFP4 checkpoints store these on moe_config, not quant_config).
+        device = torch.accelerator.current_device_index()
+
+        def _per_expert(val: float | None) -> torch.Tensor | None:
+            if val is None:
+                return None
+            return torch.full(
                (self.local_num_experts,),
-                quant_config.gemm1_clamp_limit,
+                float(val),
                dtype=torch.float32,
                device=device,
            )
+
+        clamp = quant_config.gemm1_clamp_limit
+        if clamp is None:
+            clamp = getattr(moe_config, "swiglu_limit", None)
+        alpha = quant_config.gemm1_alpha
+        if alpha is None:
+            alpha = getattr(moe_config, "swiglu_alpha", None)
+        beta = quant_config.gemm1_beta
+        if beta is None:
+            beta = getattr(moe_config, "swiglu_beta", None)
+
+
+        if moe_config.is_act_and_mul:
+            self.gemm1_clamp_limit = _per_expert(clamp)
+            self.gemm1_alpha = _per_expert(alpha)
+            self.gemm1_beta = _per_expert(beta)
        else:
            self.gemm1_clamp_limit = None
+            self.gemm1_alpha = None
+            self.gemm1_beta = None
+
+        logger.info_once(
+            "activation=%s, gemm1_alpha=%s, gemm1_beta=%s, gemm1_clamp_limit=%s",
+            moe_config.activation,
+            alpha,
+            beta,
+            clamp,
+        )

    def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
        layer.w13_weight_scale_2.data.mul_(layer.w13_input_scale)
@@ -109,6 +140,25 @@ class TrtLlmNvFp4ExpertsBase:
            )
            self.gemm1_clamp_limit = layer.gemm1_clamp_limit

+        # beta shifts the raw GEMM1 accumulator, so fold by g1_alphas like the
+        # clamp limit. alpha is applied to the dequantized gate, so it stays
+        # raw. Register both on the layer so EPLB rearranges them with the
+        # other per-expert tensors.
+        if self.gemm1_beta is not None:
+            gemm1_beta = self.gemm1_beta / self.quant_config.g1_alphas
+            layer.register_parameter(
+                "gemm1_beta",
+                torch.nn.Parameter(gemm1_beta, requires_grad=False),
+            )
+            self.gemm1_beta = layer.gemm1_beta
+
+        if self.gemm1_alpha is not None:
+            layer.register_parameter(
+                "gemm1_alpha",
+                torch.nn.Parameter(self.gemm1_alpha, requires_grad=False),
+            )
+            self.gemm1_alpha = layer.gemm1_alpha
+
    @staticmethod
    def _supports_current_device() -> bool:
        """Supports only Blackwell-family GPUs."""
@@ -137,12 +187,14 @@ class TrtLlmNvFp4ExpertsBase:

    @staticmethod
    def _supports_activation(activation: MoEActivation) -> bool:
-        """Supports only SiLU, RELU^2 non-gated and GELU activation."""
+        """Supports SiLU, RELU^2 non-gated, GELU, and clamped SwiGLU-OAI."""
        return activation in [
            MoEActivation.SILU,
            MoEActivation.RELU2_NO_MUL,
            MoEActivation.GELU,
            MoEActivation.GELU_TANH,
+            MoEActivation.SWIGLUOAI,
+            MoEActivation.SWIGLUOAI_UNINTERLEAVE,
        ]

    @staticmethod
@@ -248,8 +300,8 @@ class TrtLlmNvFp4ExpertsModular(TrtLlmNvFp4ExpertsBase, mk.FusedMoEExpertsModula
            gemm1_weights=w1,
            gemm1_weights_scale=self.quant_config.w1_scale.view(torch.float8_e4m3fn),
            gemm1_bias=None,
-            gemm1_alpha=None,
-            gemm1_beta=None,
+            gemm1_alpha=self.gemm1_alpha,
+            gemm1_beta=self.gemm1_beta,
            gemm1_clamp_limit=self.gemm1_clamp_limit,
            gemm2_weights=w2,
            gemm2_weights_scale=self.quant_config.w2_scale.view(torch.float8_e4m3fn),
@@ -409,8 +461,8 @@ class TrtLlmNvFp4ExpertsMonolithic(
            gemm1_weights=w1,
            gemm1_weights_scale=self.quant_config.w1_scale.view(torch.float8_e4m3fn),
            gemm1_bias=None,
-            gemm1_alpha=None,
-            gemm1_beta=None,
+            gemm1_alpha=self.gemm1_alpha,
+            gemm1_beta=self.gemm1_beta,
            gemm1_clamp_limit=self.gemm1_clamp_limit,
            gemm2_weights=w2,
            gemm2_weights_scale=self.quant_config.w2_scale.view(torch.float8_e4m3fn),
@@ -140,6 +140,7 @@ def FusedMoE(
    apply_routed_scale_to_output: bool = False,
    zero_expert_type: str | None = None,
    hash_indices_table: torch.Tensor | None = None,
+    reduce_results: bool = True,
    runner_cls: type[MoERunner] | None = None,
    runner_args: dict[str, Any] | None = None,
    routed_experts_cls: type[RoutedExperts] | None = None,
@@ -198,6 +199,9 @@ def FusedMoE(
                                      output instead of topk_weights
        zero_expert_type: Type of zero expert handling
        hash_indices_table: Hash table for expert indices
+        reduce_results: Whether to all-reduce the final output across TP/EP
+            ranks. Set to False to defer the all-reduce (e.g. to fuse it into
+            a subsequent GemmaRMSNorm).
        runner_cls: Custom MoERunner class (None = use default MoERunner)
        runner_args: Additional arguments for runner constructor
        routed_experts_cls: Custom RoutedExperts class (None = use default)
@@ -385,6 +389,7 @@ def FusedMoE(
        routed_scaling_factor=routed_scaling_factor
        if apply_routed_scale_to_output
        else 1.0,
+        reduce_results=reduce_results,
        **runner_args if runner_args is not None else {},
    )

@@ -254,6 +254,7 @@ class MoERunner(MoERunnerInterface):
        routed_input_transform: torch.nn.Module | None = None,
        routed_output_transform: torch.nn.Module | None = None,
        routed_scaling_factor: float = 1.0,
+        reduce_results: bool = True,
    ):
        super().__init__()
        self.moe_config = moe_config
@@ -265,6 +266,7 @@ class MoERunner(MoERunnerInterface):
        self.shared_expert_gate = shared_expert_gate
        self.routed_experts = routed_experts
        self.enable_dbo = enable_dbo
+        self.reduce_results = reduce_results

        # When both gates are present and FSE is enabled, fuse their
        # weight matrices into [num_experts + num_shared, hidden] so one
@@ -420,6 +422,15 @@ class MoERunner(MoERunnerInterface):
        * If we have SP (TP=N, DP=M, EP), there is a separate AG step handled
          in the model.
        """
+        # A combine kernel that already reduces the fused output is
+        # incompatible with deferring the all-reduce (reduce_results=False,
+        # e.g. fusing it into a subsequent GemmaRMSNorm): the deferred
+        # all-reduce would double-reduce the fused output.
+        assert not (self._fused_output_is_reduced and not self.reduce_results), (
+            "reduce_results=False is incompatible with a combine kernel that "
+            "already reduces the fused output (e.g. DeepEP/Mori/NIXL/"
+            "FlashInfer-NVLink all2all backends)."
+        )
        if (
            shared_output is not None
            and not self.moe_config.is_sequence_parallel
@@ -447,6 +458,7 @@ class MoERunner(MoERunnerInterface):
            not self.moe_config.is_sequence_parallel
            and (self.moe_config.tp_size > 1 or self.moe_config.ep_size > 1)
            and not self._fused_output_is_reduced
+            and self.reduce_results
        ):
            states = tensor_model_parallel_all_reduce(states)

@@ -153,6 +153,8 @@ class CompressedTensorsW8A8Mxfp8MoEMethod(CompressedTensorsMoEMethod):
            a2_scale=layer.w2_input_scale,
            block_shape=self.weight_block_size,
            swiglu_limit=getattr(layer, "swiglu_limit", None),
+            gemm1_alpha=getattr(layer, "swiglu_alpha", None),
+            gemm1_beta=getattr(layer, "swiglu_beta", None),
        )

    def maybe_make_prepare_finalize(
@@ -2283,6 +2283,7 @@ class ModelOptMixedPrecisionConfig(ModelOptQuantConfigBase):
        fp8_config: ModelOptFp8Config,
        nvfp4_config: ModelOptNvFp4Config,
        w4a16_nvfp4_config: ModelOptNvFp4Config,
+        mxfp8_config: ModelOptMxFp8Config,
    ) -> None:
        super().__init__(exclude_modules)
        self.kv_cache_quant_method = kv_cache_quant_method
@@ -2290,6 +2291,7 @@ class ModelOptMixedPrecisionConfig(ModelOptQuantConfigBase):
        self.fp8_config = fp8_config
        self.nvfp4_config = nvfp4_config
        self.w4a16_nvfp4_config = w4a16_nvfp4_config
+        self.mxfp8_config = mxfp8_config

    def get_name(self) -> QuantizationMethods:
        return "modelopt_mixed"
@@ -2379,6 +2381,12 @@ class ModelOptMixedPrecisionConfig(ModelOptQuantConfigBase):
            group_size=group_size,
        )

+        mxfp8_config = ModelOptMxFp8Config(
+            is_checkpoint_mxfp8_serialized=True,
+            kv_cache_quant_algo=kv_cache_quant_method,
+            exclude_modules=[],
+        )
+
        return cls(
            kv_cache_quant_method=kv_cache_quant_method,
            exclude_modules=exclude_modules,
@@ -2386,6 +2394,7 @@ class ModelOptMixedPrecisionConfig(ModelOptQuantConfigBase):
            fp8_config=fp8_config,
            nvfp4_config=nvfp4_config,
            w4a16_nvfp4_config=w4a16_nvfp4_config,
+            mxfp8_config=mxfp8_config,
        )

    def _resolve_quant_algo(self, prefix: str) -> str | None:
@@ -2440,6 +2449,17 @@ class ModelOptMixedPrecisionConfig(ModelOptQuantConfigBase):
                if key.startswith(parent_dot):
                    return info["quant_algo"].upper()

+        # 4. Parent-prefix fallback for fused projections (qkv_proj, gate_up_proj).
+        for candidate in self._quantized_layer_prefix_candidates(prefix):
+            parent_dot = candidate.rsplit(".", 1)[0] + "."
+            algos = {
+                info["quant_algo"].upper()
+                for key, info in self.quantized_layers.items()
+                if key.startswith(parent_dot) and "." not in key[len(parent_dot):]
+            }
+            if len(algos) == 1:
+                return algos.pop()
+
        return None

    @staticmethod
@@ -2485,6 +2505,8 @@ class ModelOptMixedPrecisionConfig(ModelOptQuantConfigBase):
                return ModelOptNvFp4LinearMethod(self.nvfp4_config)
            if quant_algo == "W4A16_NVFP4":
                return ModelOptNvFp4W4A16LinearMethod(self.w4a16_nvfp4_config)
+            if quant_algo == "MXFP8":
+                return ModelOptMxFp8LinearMethod(self.mxfp8_config)
            # Layer not in quantized_layers — leave unquantized
            return UnquantizedLinearMethod()

@@ -2504,6 +2526,11 @@ class ModelOptMixedPrecisionConfig(ModelOptQuantConfigBase):
                    quant_config=self.w4a16_nvfp4_config,
                    moe_config=layer.moe_config,
                )
+            if quant_algo == "MXFP8":
+                return ModelOptMxFp8FusedMoE(
+                    quant_config=self.mxfp8_config,
+                    moe_config=layer.moe_config,
+                )
            return None

        return None
@@ -224,6 +224,8 @@ class Mxfp8OnlineMoEMethod(OnlineMoEMethodBase):
            w2_bias=getattr(layer, "w2_bias", None),
            block_shape=self.weight_block_size,
            swiglu_limit=getattr(layer, "swiglu_limit", None),
+            gemm1_alpha=getattr(layer, "swiglu_alpha", None),
+            gemm1_beta=getattr(layer, "swiglu_beta", None),
        )

    def process_weights_after_loading(self, layer: Module) -> None:
@@ -36,6 +36,13 @@ def activation_to_flashinfer_type(activation: MoEActivation) -> "ActivationType"
        MoEActivation.GELU: ActivationType.Geglu,
        MoEActivation.GELU_TANH: ActivationType.Geglu,
        MoEActivation.RELU2_NO_MUL: ActivationType.Relu2,
+        # Both OAI variants map to Swiglu: FlashInfer has no SwigluOAI enum;
+        # the clamped/biased behavior is driven by the per-expert gemm1_alpha/
+        # gemm1_beta/gemm1_clamp_limit tensors (see trtllm_nvfp4_moe.py).
+        # The interleaved-vs-contiguous row layout difference between the two
+        # is resolved in process_weights_after_loading, not here.
+        MoEActivation.SWIGLUOAI: ActivationType.Swiglu,
+        MoEActivation.SWIGLUOAI_UNINTERLEAVE: ActivationType.Swiglu,
    }
    return ACTIVATION_TO_FI_ACTIVATION[activation]

@@ -457,6 +457,7 @@ class MiniMaxM3SparseAttention(nn.Module, AttentionLayerBase):
        quant_config: QuantizationConfig | None = None,
        prefix: str = "",
        cache_config: CacheConfig | None = None,
+        topk_indices_buffer: torch.Tensor | None = None,
    ) -> None:
        super().__init__()
        self.hidden_size = config.hidden_size
@@ -535,6 +536,9 @@ class MiniMaxM3SparseAttention(nn.Module, AttentionLayerBase):
            self.kv_cache_dtype, vllm_config.model_config
        )

+        # Shared top-k buffer: the indexer writes the selected blocks into it and
+        # the attend impl reads them back (no Python value crosses the break).
+        self.topk_indices_buffer = topk_indices_buffer
        self.attn_backend = MiniMaxM3SparseBackend
        # Indexer and main attention are separate impls. On ROCm the SM100 gate
        # is always False, so both pick Triton and the index cache stays bf16.
@@ -565,6 +569,7 @@ class MiniMaxM3SparseAttention(nn.Module, AttentionLayerBase):
            local_blocks=sparse_cfg.get("sparse_local_block", 0),
            score_type=sparse_cfg.get("sparse_score_type", "max"),
            cache_config=cache_config,
+            topk_indices_buffer=topk_indices_buffer,
        )

        # Register the main K/V cache so the KV-cache manager allocates it.
@@ -657,9 +662,10 @@ class MiniMaxM3SparseAttention(nn.Module, AttentionLayerBase):
        output: torch.Tensor,
    ) -> torch.Tensor:
        # Single eager break around both: their split-K kernels read per-request
-        # metadata and can't be captured into a cudagraph.
-        topk_idx = self.indexer(index_query)
-        return self.impl.forward(self, query, self.kv_cache, topk_idx, output)
+        # metadata and can't be captured into a cudagraph. The indexer writes its
+        # top-k into the shared ``topk_indices_buffer``; the attend reads it back.
+        self.indexer(index_query)
+        return self.impl.forward(self, query, self.kv_cache, output)


 class MiniMaxM3DecoderLayer(nn.Module):
@@ -671,6 +677,7 @@ class MiniMaxM3DecoderLayer(nn.Module):
        quant_config: QuantizationConfig | None = None,
        force_sparse_attn: bool = False,
        force_moe: bool = False,
+        topk_indices_buffer: torch.Tensor | None = None,
    ) -> None:
        super().__init__()
        self.hidden_size = config.hidden_size
@@ -690,6 +697,7 @@ class MiniMaxM3DecoderLayer(nn.Module):
                quant_config=quant_config,
                prefix=f"{prefix}.self_attn",
                cache_config=cache_config,
+                topk_indices_buffer=topk_indices_buffer,
            )
        else:
            self.self_attn = MiniMaxM3Attention(
@@ -771,6 +779,22 @@ class MiniMaxM3Model(nn.Module, EagleModelMixin):
            prefix=f"{prefix}.embed_tokens",
        )

+        # Reserved top-k indices buffer shared by all sparse-attention indexer
+        # layers (mirrors DeepseekV4); the indexer writes its per-head decode/
+        # prefill block selection into it, the attend reads it back.
+        sparse_cfg = getattr(config, "sparse_attention_config", None)
+        if sparse_cfg is not None:
+            tp_size = get_tensor_model_parallel_world_size()
+            num_index_heads = max(1, sparse_cfg["sparse_num_index_heads"] // tp_size)
+            self.topk_indices_buffer = torch.empty(
+                num_index_heads,
+                vllm_config.scheduler_config.max_num_batched_tokens,
+                sparse_cfg["sparse_topk_blocks"],
+                dtype=torch.int32,
+            )
+        else:
+            self.topk_indices_buffer = None
+
        self.start_layer, self.end_layer, self.layers = make_layers(
            config.num_hidden_layers,
            lambda prefix: MiniMaxM3DecoderLayer(
@@ -778,6 +802,7 @@ class MiniMaxM3Model(nn.Module, EagleModelMixin):
                prefix,
                cache_config=cache_config,
                quant_config=quant_config,
+                topk_indices_buffer=self.topk_indices_buffer,
            ),
            prefix=f"{prefix}.layers",
        )
@@ -25,12 +25,14 @@ from vllm.config.attention import IndexerKVDType
 from vllm.config.cache import CacheDType
 from vllm.distributed import get_tensor_model_parallel_world_size
 from vllm.forward_context import get_forward_context
+from vllm.logger import init_logger
 from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
 from vllm.models.minimax_m3.common.ops.index_topk import (
    minimax_m3_index_decode,
    minimax_m3_index_score,
    minimax_m3_index_topk,
 )
+from vllm.platforms import current_platform
 from vllm.v1.attention.backend import (
    AttentionBackend,
    AttentionCGSupport,
@@ -46,6 +48,8 @@ from vllm.v1.kv_cache_interface import (
    MLAAttentionSpec,
 )

+logger = init_logger(__name__)
+

 class MiniMaxM3IndexerBackend(AttentionBackend):
    """Indexer side-cache backend (key-only)."""
@@ -120,16 +124,20 @@ class MiniMaxM3IndexerCache(nn.Module, AttentionLayerBase):
        backend_cls: type[AttentionBackend] = MiniMaxM3IndexerBackend,
    ) -> None:
        super().__init__()
-        if indexer_kv_dtype != "bf16":
+        if indexer_kv_dtype in ("fp8", "fp8_e4m3"):
+            cache_dtype = torch.float8_e4m3fn
+        elif indexer_kv_dtype == "bf16":
+            cache_dtype = torch.bfloat16
+        else:
            raise NotImplementedError(
-                f"indexer_kv_dtype={indexer_kv_dtype!r} is not supported yet "
-                "for the MiniMax M3 indexer cache (only 'bf16')."
+                f"indexer_kv_dtype={indexer_kv_dtype!r} is not supported by the "
+                "MiniMax M3 indexer cache (only 'bf16' or 'fp8'/'fp8_e4m3')."
            )
        self.kv_cache = torch.tensor([])
        self.head_dim = head_dim
        self.indexer_kv_dtype = indexer_kv_dtype
-        # Storage dtype for the side cache (bf16 today; quantized layouts later).
-        self.dtype = torch.bfloat16
+        # Side-cache storage dtype: bf16, or e4m3 for the fp8 score path.
+        self.dtype = cache_dtype
        self.prefix = prefix
        self.cache_config = cache_config
        # Impl-chosen backend -> each impl gets its own builder (get_attn_backend).
@@ -344,6 +352,7 @@ class MiniMaxM3IndexerImpl(nn.Module):
        score_type: str = "max",
        cache_config: CacheConfig | None = None,
        indexer_kv_dtype: IndexerKVDType = "bf16",
+        topk_indices_buffer: torch.Tensor | None = None,
    ) -> None:
        super().__init__()
        self.num_kv_heads = num_kv_heads
@@ -356,6 +365,9 @@ class MiniMaxM3IndexerImpl(nn.Module):
        self.num_index_heads = num_index_heads
        self.index_head_dim = index_head_dim
        self.indexer_kv_dtype = indexer_kv_dtype
+        # Shared, stable-address top-k output buffer (set by the model for the
+        # cudagraph-safe MSA impl); None -> impl allocates fresh (eager).
+        self.topk_indices_buffer = topk_indices_buffer
        # Owns the side cache (registers itself in the static forward context).
        self.index_cache = MiniMaxM3IndexerCache(
            head_dim=index_head_dim,
@@ -392,6 +404,10 @@ class MiniMaxM3IndexerTritonImpl(MiniMaxM3IndexerImpl):
        )
        kv = self.index_cache.kv_cache

+        # Both sides write into the single shared persistent topk_indices_buffer
+        # (decode at [:, :nd], prefill at [:, nd:]) and return views into it; the
+        # kernels' out= writes out[:, :total_q]. None -> allocate fresh.
+        buf = self.topk_indices_buffer
        decode_topk: torch.Tensor | None = None
        prefill_topk: torch.Tensor | None = None
        if index_md.num_decodes > 0:
@@ -409,6 +425,7 @@ class MiniMaxM3IndexerTritonImpl(MiniMaxM3IndexerImpl):
                self.num_kv_heads,
                d.decode_query_len,
                d.max_decode_query_len,
+                out=buf,
            )
        if index_md.num_prefills > 0:
            p = index_md.prefill
@@ -432,29 +449,61 @@ class MiniMaxM3IndexerTritonImpl(MiniMaxM3IndexerImpl):
                self.topk_blocks,
                self.init_blocks,
                self.local_blocks,
+                out=buf[:, nd:, :] if buf is not None else None,
            )
        return decode_topk, prefill_topk


 def select_indexer_impl_cls(
    *,
+    topk_blocks: int,
    indexer_kv_dtype: IndexerKVDType = "bf16",
 ) -> type[MiniMaxM3IndexerImpl]:
-    """Pick the indexer impl off the index-cache dtype.
+    """Pick the indexer impl off the platform, top-k count, and cache dtype.

-    The SM100 MSA indexer score path is disabled for now; use the local Triton
-    indexer. If re-enabled, add a NVIDIA-specific ``MiniMaxM3IndexerImpl`` here.
+    On Blackwell (SM100) with ``topk_blocks`` in ``(4, 8, 16, 32)`` (matching the
+    main MSA attend), the fmha_sm100 score path + Triton top-k is used for both
+    bf16 and fp8 index caches. Everything else falls back to the Triton indexer
+    (bf16 only).
    """
    if indexer_kv_dtype in ("mxfp4", "nvfp4"):
        raise NotImplementedError(
            f"indexer_kv_dtype={indexer_kv_dtype!r} needs the (not-yet-added) "
            "CuteDSL indexer impl."
        )
+    is_sm100 = (
+        current_platform.is_cuda() and current_platform.is_device_capability_family(100)
+    )
+    use_msa = (
+        is_sm100
+        and topk_blocks in (4, 8, 16, 32)
+        and indexer_kv_dtype in ("bf16", "fp8", "fp8_e4m3")
+    )
+    if use_msa:
+        # Lazy import so AMD / non-SM100 never import fmha_sm100.
+        from vllm.models.minimax_m3.nvidia.indexer_msa import (
+            MiniMaxM3IndexerMSAImpl,
+        )
+
+        logger.info_once(
+            "MiniMax M3 indexer: selected MSA (fmha_sm100 score + Triton top-k) "
+            "[topk_blocks=%d, indexer_kv_dtype=%s]",
+            topk_blocks,
+            indexer_kv_dtype,
+        )
+        return MiniMaxM3IndexerMSAImpl
    if indexer_kv_dtype != "bf16":
        raise NotImplementedError(
            f"indexer_kv_dtype={indexer_kv_dtype!r} is not supported by the "
            "Triton indexer impl."
        )
+    logger.info_once(
+        "MiniMax M3 indexer: selected Triton (no fmha_sm100) "
+        "[topk_blocks=%d, indexer_kv_dtype=%s, sm100=%s]",
+        topk_blocks,
+        indexer_kv_dtype,
+        is_sm100,
+    )
    return MiniMaxM3IndexerTritonImpl


@@ -480,9 +529,11 @@ class MiniMaxM3Indexer(nn.Module):
        score_type: str = "max",
        cache_config: CacheConfig | None = None,
        indexer_kv_dtype: IndexerKVDType = "bf16",
+        topk_indices_buffer: torch.Tensor | None = None,
    ) -> None:
        super().__init__()
        impl_cls = select_indexer_impl_cls(
+            topk_blocks=topk_blocks,
            indexer_kv_dtype=indexer_kv_dtype,
        )
        self.impl = impl_cls(
@@ -498,6 +549,7 @@ class MiniMaxM3Indexer(nn.Module):
            score_type=score_type,
            cache_config=cache_config,
            indexer_kv_dtype=indexer_kv_dtype,
+            topk_indices_buffer=topk_indices_buffer,
        )

    @property
@@ -373,7 +373,10 @@ def _decode_index_score_kernel(
            + off_k[:, None] * stride_ik_pos
            + off_d * stride_ik_d,
        )  # [N,D]
-        kq = tl.dot(k, q)  # [N,HQ]
+        # fp32 accumulation is required for the fp8 (e4m3) index cache: q/k are
+        # loaded in their stored dtype (bf16 or e4m3) and the MMA accumulates in
+        # fp32 so the per-block max score is exact for the fp8 indexer too.
+        kq = tl.dot(k, q, out_dtype=tl.float32)  # [N,HQ]
        kq = tl.where(pos_mask & q_mask[None, :], kq, float("-inf"))
        score = tl.max(kq, axis=0)  # [HQ]
        is_visible_block = blk < num_blocks_q
@@ -709,16 +712,25 @@ def minimax_m3_index_topk(
    topk: int,
    init_blocks: int,
    local_blocks: int,
+    out: torch.Tensor | None = None,
 ) -> torch.Tensor:
-    """Select index top-k from a precomputed score tensor."""
+    """Select index top-k from a precomputed score tensor.
+
+    When ``out`` is provided (a ``[num_idx_heads, >=total_q, topk]`` buffer), the
+    result is written into ``out[:, :total_q, :]`` instead of a fresh tensor --
+    used to keep the top-k output at a stable address for cudagraph capture.
+    """
    num_idx_heads = score.shape[0]
    batch = cu_seqlens_q.shape[0] - 1
    total_q = score.shape[1]
-    topk_idx = torch.empty(
-        (num_idx_heads, total_q, topk),
-        dtype=torch.int32,
-        device=score.device,
-    )
+    if out is not None:
+        topk_idx = out[:, :total_q, :]
+    else:
+        topk_idx = torch.empty(
+            (num_idx_heads, total_q, topk),
+            dtype=torch.int32,
+            device=score.device,
+        )
    # block_size_q == 1 -> query blocks coincide with query tokens.
    grid_topk = (max_query_len, batch, num_idx_heads)
    _topk_index_kernel[grid_topk](
@@ -757,10 +769,13 @@ def minimax_m3_index_decode(
    num_kv_heads: int,
    decode_query_len: int,
    max_decode_query_len: int,
+    out: torch.Tensor | None = None,
 ) -> torch.Tensor:
    """Decode index block-score + top-k, both split-K (cudagraph-safe).

    Returns topk_idx [num_kv_heads, total_q, topk] (0-indexed block ids, -1 pad).
+    When ``out`` ([num_kv_heads, >=total_q, topk]) is given, writes into
+    ``out[:, :total_q, :]`` (stable address for cudagraph) instead of allocating.
    """
    total_q, num_idx_heads, head_dim = idx_q.shape
    assert num_idx_heads == num_kv_heads, (
@@ -794,7 +809,7 @@ def minimax_m3_index_decode(
    )
    # split-K over seq blocks; chunk count depends only on shape constants so
    # the grid is fixed within a cuda graph.
-    TARGET_GRID = 512
+    TARGET_GRID = 4096
    MAX_NUM_KV_CHUNKS = 256
    # Use the configured max decode length to avoid Triton recompiles when
    # switching between qlen=1 and spec-decode verification batches.
@@ -834,14 +849,17 @@ def minimax_m3_index_decode(
        **score_kwargs,
    )

-    topk_idx = torch.empty(
-        (num_idx_heads, total_q, topk),
-        dtype=torch.int32,
-        device=idx_q.device,
-    )
+    if out is not None:
+        topk_idx = out[:, :total_q, :]
+    else:
+        topk_idx = torch.empty(
+            (num_idx_heads, total_q, topk),
+            dtype=torch.int32,
+            device=idx_q.device,
+        )
    # Chunk count is shape-constant (cudagraph-safe), capped so the merge sorts
    # pow2(num_topk_chunks * pow2(topk)) candidates.
-    TOPK_TARGET_GRID = 64
+    TOPK_TARGET_GRID = 512
    MAX_NUM_TOPK_CHUNKS = 16
    topk_target = max(
        1, min(MAX_NUM_TOPK_CHUNKS, TOPK_TARGET_GRID // max(1, batch * num_idx_heads))
@@ -2,10 +2,11 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 """Main block-sparse GQA attention for MiniMax M3 sparse layers.

-The lightning indexer (``indexer.py``) selects the top-k KV blocks; this module
-holds the main attention that attends only to those blocks: the paged K/V cache
-backend, its metadata + builder, and the impl that consumes the indexer's
-``topk_idx``. The Triton attend kernel lives here; the SM100 (MSA)
+The lightning indexer (``indexer.py``) selects the top-k KV blocks (written into
+the shared ``layer.topk_indices_buffer``); this module holds the main attention
+that attends only to those blocks: the paged K/V cache backend, its metadata +
+builder, and the impl that reads the indexer's top-k from that buffer. The Triton
+attend kernel lives here; the SM100 (MSA)
 ``build_k2q_csr`` + ``sparse_atten_func`` attend lives in
 ``nvidia/sparse_attention_msa.py``.

@@ -272,9 +273,10 @@ class MiniMaxM3SparseImpl(AttentionImplBase[MiniMaxM3SparseMetadata]):
    """Abstract base for block-sparse GQA over the indexer-selected blocks.

    Inherits ``AttentionImplBase`` for a custom forward signature (the layer
-    pre-inserts K/V and runs the indexer, so forward takes the queries +
-    ``topk_idx``). The Triton and MSA subclasses each own a full ``forward`` --
-    no shared forward code.
+    pre-inserts K/V and runs the indexer, which writes the selected blocks into
+    the shared ``layer.topk_indices_buffer``; the attend reads them back from
+    there). The Triton and MSA subclasses each own a full ``forward`` -- no
+    shared forward code.
    """

    def __init__(
@@ -311,10 +313,14 @@ class MiniMaxM3SparseImpl(AttentionImplBase[MiniMaxM3SparseMetadata]):
        layer: AttentionLayer,
        query: torch.Tensor,
        kv_cache: torch.Tensor,
-        topk_idx: tuple[torch.Tensor | None, torch.Tensor | None],
        output: torch.Tensor,
    ) -> torch.Tensor:
-        """Attend the queries to the indexer-selected blocks. Per kernel."""
+        """Attend the queries to the indexer-selected blocks. Per kernel.
+
+        The indexer has already written the top-k block ids into
+        ``layer.topk_indices_buffer`` (decode at ``[:, :nd]``, prefill at
+        ``[:, nd:num_tokens]``); the attend reads them from there.
+        """
        raise NotImplementedError


@@ -326,7 +332,6 @@ class MiniMaxM3SparseTritonImpl(MiniMaxM3SparseImpl):
        layer: AttentionLayer,
        query: torch.Tensor,
        kv_cache: torch.Tensor,
-        topk_idx: tuple[torch.Tensor | None, torch.Tensor | None],
        output: torch.Tensor,
    ) -> torch.Tensor:
        attn_metadata = get_forward_context().attn_metadata
@@ -334,10 +339,12 @@ class MiniMaxM3SparseTritonImpl(MiniMaxM3SparseImpl):
            return output  # profiling run; caches unbound
        main_md = attn_metadata[layer.layer_name]  # type: ignore[attr-defined]
        assert isinstance(main_md, MiniMaxM3SparseMetadata)
-        decode_topk, prefill_topk = topk_idx

        nd = main_md.num_decode_tokens
        num_tokens = main_md.num_actual_tokens
+        # Indexer top-k from the shared buffer: decode [:, :nd], prefill [:, nd:].
+        topk = layer.topk_indices_buffer  # type: ignore[attr-defined]
+        assert topk is not None
        hd = self.head_size
        q = query[:num_tokens].view(-1, self.num_heads, hd)
        out = output[:num_tokens].view(-1, self.num_heads, hd)
@@ -348,11 +355,11 @@ class MiniMaxM3SparseTritonImpl(MiniMaxM3SparseImpl):
        # Decode [:nd]: split-K over the selected blocks (request-major chunks).
        if main_md.num_decodes > 0:
            d = main_md.decode
-            assert d is not None and decode_topk is not None
+            assert d is not None
            minimax_m3_sparse_attn_decode(
                q[:nd],
                kv_cache,
-                decode_topk,
+                topk[:, :nd, :],
                d.block_table,
                d.seq_lens,
                self.num_kv_heads,
@@ -364,11 +371,11 @@ class MiniMaxM3SparseTritonImpl(MiniMaxM3SparseImpl):
        # Prefill [nd:]: cu_seqlens_q already rebased to 0.
        if main_md.num_prefills > 0:
            p = main_md.prefill
-            assert p is not None and prefill_topk is not None
+            assert p is not None
            minimax_m3_sparse_attn(
                q[nd:],
                kv_cache,
-                prefill_topk,
+                topk[:, nd:num_tokens, :],
                p.block_table,
                p.cu_seqlens_q,
                p.seq_lens,
@@ -0,0 +1,251 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""MSA (SM100/Blackwell) indexer impl for MiniMax M3.
+
+Prefill scores with ``fmha_sm100``'s score-only (``OnlyScore``) path then selects
+top-k blocks with the Triton ``minimax_m3_index_topk`` kernel -- fmha is much
+faster than Triton for the wide prefill score (benchmarked ~3-5x).
+
+Decode uses the Triton fused ``minimax_m3_index_decode`` (the same kernel the
+Triton indexer impl uses): for q_len==1 it is a purpose-built vector x matrix
+score (no wasted tensor-core tiles) with a 256-way split-K and a fused split-K
+top-k, which beats fmha's OnlyScore (wasted MMA on a single query, 64-split cap)
+by ~1.1-3.7x. It is cudagraph-safe by construction (shape-constant split grids)
+and writes the shared ``topk_indices_buffer`` via ``out=``.
+
+``fmha_sm100`` imports are function-local so this module is import-safe on
+AMD / non-SM100.
+"""
+
+from dataclasses import dataclass
+from typing import ClassVar
+
+import torch
+
+from vllm.forward_context import get_forward_context
+from vllm.models.minimax_m3.common.indexer import (
+    MiniMaxM3IndexerBackend,
+    MiniMaxM3IndexerDecodeMetadata,
+    MiniMaxM3IndexerImpl,
+    MiniMaxM3IndexerMetadata,
+    MiniMaxM3IndexerMetadataBuilder,
+)
+from vllm.models.minimax_m3.common.ops.index_topk import (
+    minimax_m3_index_decode,
+    minimax_m3_index_topk,
+)
+from vllm.v1.attention.backend import (
+    AttentionBackend,
+    AttentionCGSupport,
+    CommonAttentionMetadata,
+)
+from vllm.v1.attention.backends.utils import split_decodes_and_prefills
+
+# Page size == sparse block size == index-K block; fmha tile id == M3 block id.
+PAGE_SIZE = 128
+
+
+class MiniMaxM3IndexerMSABackend(MiniMaxM3IndexerBackend):
+    """Indexer side-cache backend selecting the MSA builder."""
+
+    @staticmethod
+    def get_builder_cls() -> type["MiniMaxM3IndexerMSAMetadataBuilder"]:
+        return MiniMaxM3IndexerMSAMetadataBuilder
+
+
+@dataclass
+class MiniMaxM3IndexerMSAPrefillMetadata:
+    """fmha score plan + Triton top-k inputs for the prefill side (eager)."""
+
+    plan: dict  # fmha_sm100 PlanInfo
+    cu_seqlens_q: torch.Tensor  # [num_prefills + 1] int32, rebased to 0
+    prefix_lens: torch.Tensor  # [num_prefills] int32, context tokens
+    max_query_len: int
+    page_table: torch.Tensor  # flat physical page indices for the prefill side
+
+
+@dataclass
+class MiniMaxM3IndexerMSAMetadata(MiniMaxM3IndexerMetadata):
+    """Decode reuses the inherited base ``decode`` field (the Triton decode
+    metadata); ``prefill_msa`` carries the fmha score plan for the prefill side
+    (the base ``prefill`` field is unused on this path)."""
+
+    prefill_msa: MiniMaxM3IndexerMSAPrefillMetadata | None = None
+
+
+class MiniMaxM3IndexerMSAMetadataBuilder(MiniMaxM3IndexerMetadataBuilder):
+    """Decode metadata is the cudagraph-safe Triton decode metadata; the prefill
+    fmha plan is built eagerly (prefill batches are not captured)."""
+
+    _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.UNIFORM_BATCH
+
+    def build(
+        self,
+        common_prefix_len: int,
+        common_attn_metadata: CommonAttentionMetadata,
+        fast_build: bool = False,
+    ) -> MiniMaxM3IndexerMSAMetadata:
+        num_reqs = common_attn_metadata.num_reqs
+        num_tokens = common_attn_metadata.num_actual_tokens
+        seq_lens = common_attn_metadata.seq_lens
+        block_table = common_attn_metadata.block_table_tensor
+        query_start_loc = common_attn_metadata.query_start_loc
+
+        num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens = (
+            split_decodes_and_prefills(
+                common_attn_metadata,
+                decode_threshold=self.reorder_batch_threshold,
+                require_uniform=True,
+            )
+        )
+        assert num_decodes + num_prefills == num_reqs
+        assert num_decode_tokens + num_prefill_tokens == num_tokens
+
+        # Context (prefix) lengths into the stable cudagraph buffer.
+        context_lens = self.context_len_buffer[:num_reqs]
+        context_lens.copy_(
+            common_attn_metadata.compute_num_computed_tokens(), non_blocking=True
+        )
+
+        decode: MiniMaxM3IndexerDecodeMetadata | None = None
+        if num_decodes > 0:
+            qsl_cpu = common_attn_metadata.query_start_loc_cpu
+            query_lens_cpu = qsl_cpu[1 : num_decodes + 1] - qsl_cpu[:num_decodes]
+            decode_query_len = int(query_lens_cpu[0].item())
+            assert decode_query_len > 0
+            assert torch.all(
+                (query_lens_cpu == decode_query_len) | (query_lens_cpu == 0)
+            )
+            decode = MiniMaxM3IndexerDecodeMetadata(
+                seq_lens=seq_lens[:num_decodes],
+                block_table=block_table[:num_decodes],
+                max_seq_len=common_attn_metadata.max_seq_len,
+                decode_query_len=decode_query_len,
+                max_decode_query_len=self.max_decode_query_len,
+            )
+
+        prefill: MiniMaxM3IndexerMSAPrefillMetadata | None = None
+        if num_prefills > 0:
+            # Prefill is eager (not captured); the host lengths it needs (and the
+            # _fmha_sm100_plan .tolist() inside) make the D->H sync acceptable.
+            from vllm.third_party.fmha_sm100.api import _fmha_sm100_plan
+
+            lo, hi = num_decodes, num_reqs
+            qsl_cpu = common_attn_metadata.query_start_loc_cpu[: num_reqs + 1]
+            qo_lens_cpu = (qsl_cpu[1:] - qsl_cpu[:-1]).to(torch.int32)
+            kv_lens_cpu = seq_lens[:num_reqs].cpu().to(torch.int32)
+            nvp = (kv_lens_cpu + PAGE_SIZE - 1) // PAGE_SIZE
+            side_qo = qo_lens_cpu[lo:hi]
+            side_kv = kv_lens_cpu[lo:hi]
+            plan = _fmha_sm100_plan(
+                side_qo,
+                side_kv,
+                self.num_index_heads,
+                num_kv_heads=1,
+                qo_offset=side_kv - side_qo,  # bottom-right causal
+                page_size=PAGE_SIZE,
+                output_maxscore=True,
+                causal=True,
+                num_kv_splits=1,
+            )
+            cols = torch.arange(block_table.shape[1], device=block_table.device)
+            valid = cols[None, :] < nvp[lo:hi].to(block_table.device)[:, None]
+            prefill = MiniMaxM3IndexerMSAPrefillMetadata(
+                plan=plan,
+                cu_seqlens_q=(query_start_loc[lo : hi + 1] - query_start_loc[lo]).to(
+                    torch.int32
+                ),
+                prefix_lens=context_lens[lo:hi],
+                max_query_len=int(side_qo.max()),
+                page_table=block_table[lo:hi][valid].to(torch.int32),
+            )
+
+        return MiniMaxM3IndexerMSAMetadata(
+            seq_lens=seq_lens,
+            max_seq_len=common_attn_metadata.max_seq_len,
+            slot_mapping=common_attn_metadata.slot_mapping,
+            num_actual_tokens=num_tokens,
+            num_decodes=num_decodes,
+            num_decode_tokens=num_decode_tokens,
+            num_prefills=num_prefills,
+            num_prefill_tokens=num_prefill_tokens,
+            decode=decode,
+            prefill_msa=prefill,
+        )
+
+
+class MiniMaxM3IndexerMSAImpl(MiniMaxM3IndexerImpl):
+    """Decode: Triton fused score+top-k. Prefill: fmha_sm100 OnlyScore + top-k."""
+
+    indexer_backend_cls: ClassVar[type[AttentionBackend]] = MiniMaxM3IndexerMSABackend
+
+    def forward(
+        self,
+        index_query: torch.Tensor,
+    ) -> tuple[torch.Tensor | None, torch.Tensor | None]:
+        attn_metadata = get_forward_context().attn_metadata
+        if not isinstance(attn_metadata, dict):
+            return None, None  # profiling run; caches unbound
+        md = attn_metadata[self.index_cache.prefix]
+        assert isinstance(md, MiniMaxM3IndexerMSAMetadata)
+
+        num_tokens = md.num_actual_tokens
+        nd = md.num_decode_tokens
+        index_q = index_query[:num_tokens].view(
+            -1, self.num_index_heads, self.index_head_dim
+        )
+        kv = self.index_cache.kv_cache
+        # Both sides write into the single shared persistent topk_indices_buffer:
+        # decode at [:, :nd], prefill at [:, nd:] (each kernel writes [:, :total_q]).
+        buf = self.topk_indices_buffer
+
+        decode_topk: torch.Tensor | None = None
+        if md.decode is not None:
+            d = md.decode
+            decode_topk = minimax_m3_index_decode(
+                index_q[:nd],
+                kv,
+                d.block_table,
+                d.seq_lens,
+                d.max_seq_len,
+                self.topk_blocks,
+                self.init_blocks,
+                self.local_blocks,
+                self.num_kv_heads,
+                d.decode_query_len,
+                d.max_decode_query_len,
+                out=buf,
+            )
+
+        prefill_topk: torch.Tensor | None = None
+        if md.prefill_msa is not None:
+            from vllm.third_party.fmha_sm100.api import _fmha_sm100
+
+            p = md.prefill_msa
+            # Index-K cache (num_blocks, 128, D) -> paged MQA (num_blocks,1,128,D).
+            k_pages = kv.view(kv.shape[0], 1, PAGE_SIZE, self.index_head_dim)
+            _, max_score = _fmha_sm100(
+                index_q[nd:],
+                k_pages,
+                k_pages,  # V placeholder; not read in OnlyScore
+                p.plan,
+                kv_indices=p.page_table,
+                output_o=False,
+                output_maxscore=True,
+                sm_scale=self.scale,
+            )
+            # Triton top-k wants [num_index_heads, num_tokens, max_block]; the
+            # transpose is a strided view (the kernel reads via strides).
+            out = buf[:, nd:, :] if buf is not None else None
+            prefill_topk = minimax_m3_index_topk(
+                max_score.transpose(1, 2),
+                p.cu_seqlens_q,
+                p.prefix_lens,
+                p.max_query_len,
+                self.topk_blocks,
+                self.init_blocks,
+                self.local_blocks,
+                out=out,
+            )
+
+        return decode_topk, prefill_topk
@@ -193,6 +193,7 @@ class MiniMaxM3MoE(nn.Module):
        layer_id: int,
        quant_config: QuantizationConfig | None = None,
        prefix: str = "",
+        reduce_results: bool = True,
    ) -> None:
        super().__init__()
        self.tp_size = get_tensor_model_parallel_world_size()
@@ -259,6 +260,7 @@ class MiniMaxM3MoE(nn.Module):
            shared_experts=self.shared_experts,
            quant_config=quant_config,
            prefix=f"{prefix}.experts",
+            reduce_results=reduce_results,
        )

    @staticmethod
@@ -402,6 +404,7 @@ class MiniMaxM3SparseAttention(nn.Module, AttentionLayerBase):
        quant_config: QuantizationConfig | None = None,
        prefix: str = "",
        cache_config: CacheConfig | None = None,
+        topk_indices_buffer: torch.Tensor | None = None,
    ) -> None:
        super().__init__()
        self.hidden_size = config.hidden_size
@@ -489,6 +492,10 @@ class MiniMaxM3SparseAttention(nn.Module, AttentionLayerBase):
        # cache (--attention-config '{"indexer_kv_dtype": ...}').
        self.indexer_kv_dtype = vllm_config.attention_config.indexer_kv_dtype

+        # Shared top-k buffer: the indexer writes the selected blocks into it and
+        # the attend impl reads them back (so nothing crosses the eager break as a
+        # Python value, which would freeze at capture).
+        self.topk_indices_buffer = topk_indices_buffer
        self.attn_backend = MiniMaxM3SparseBackend
        # Indexer (top-k selection) and main attention are separate impls, each
        # picking Triton vs MSA off its cache dtype. impl is AttentionImplBase
@@ -519,6 +526,7 @@ class MiniMaxM3SparseAttention(nn.Module, AttentionLayerBase):
            score_type=sparse_cfg.get("sparse_score_type", "max"),
            cache_config=cache_config,
            indexer_kv_dtype=self.indexer_kv_dtype,
+            topk_indices_buffer=topk_indices_buffer,
        )

        # Register the main K/V cache so the KV-cache manager allocates it.
@@ -576,7 +584,12 @@ class MiniMaxM3SparseAttention(nn.Module, AttentionLayerBase):
        main_slot_mapping = fwd_slot_mapping[self.layer_name]
        index_slot_mapping = fwd_slot_mapping[self.indexer.index_cache.prefix]
        q = qkv.new_empty((num_tokens, self.q_size))
-        index_q = qkv.new_empty((num_tokens, self.index_q_size))
+        # index_q matches the index-K cache dtype (e4m3 for the fp8 score path);
+        # the fused kernel emits fp8 directly when this buffer is e4m3.
+        index_q = qkv.new_empty(
+            (num_tokens, self.index_q_size),
+            dtype=self.indexer.index_cache.dtype,
+        )
        ops.fused_minimax_m3_qknorm_rope_kv_insert(
            qkv,
            self.q_norm.weight,
@@ -613,9 +626,10 @@ class MiniMaxM3SparseAttention(nn.Module, AttentionLayerBase):
        output: torch.Tensor,
    ) -> torch.Tensor:
        # Single eager break around both: their split-K kernels read per-request
-        # metadata and can't be captured into a cudagraph.
-        topk_idx = self.indexer(index_query)
-        return self.impl.forward(self, query, self.kv_cache, topk_idx, output)
+        # metadata and can't be captured into a cudagraph. The indexer writes its
+        # top-k into the shared ``topk_indices_buffer``; the attend reads it back.
+        self.indexer(index_query)
+        return self.impl.forward(self, query, self.kv_cache, output)


 class MiniMaxM3DecoderLayer(nn.Module):
@@ -627,6 +641,7 @@ class MiniMaxM3DecoderLayer(nn.Module):
        force_sparse_attn: bool = False,
        force_moe: bool = False,
        is_mtp_block: bool = False,
+        topk_indices_buffer: torch.Tensor | None = None,
    ) -> None:
        super().__init__()
        if is_mtp_block:
@@ -642,13 +657,12 @@ class MiniMaxM3DecoderLayer(nn.Module):
        layer_id = int(prefix.split(sep=".")[-1])
        self.layer_id = layer_id

-        # Complete the preceding dense MLP's deferred all-reduce
-        # (reduce_results=False), fused into this layer's input_layernorm.
-        # Disable this fusion when PP is set
+        # Complete the preceding FFN's deferred all-reduce (its down_proj / MoE
+        # combine ran with reduce_results=False), fused into this layer's
+        # input_layernorm. Both dense and MoE FFNs defer under PP==1, so every
+        # non-first layer fuses; disable when PP>1 (FFNs reduce themselves).
        self.fuse_input_allreduce = (
-            layer_id > 0
-            and not _is_moe_layer(config, layer_id - 1)
-            and vllm_config.parallel_config.pipeline_parallel_size == 1
+            layer_id > 0 and vllm_config.parallel_config.pipeline_parallel_size == 1
        )

        is_sparse_attention_layer = (
@@ -662,6 +676,7 @@ class MiniMaxM3DecoderLayer(nn.Module):
                quant_config=quant_config,
                prefix=f"{prefix}.self_attn",
                cache_config=cache_config,
+                topk_indices_buffer=topk_indices_buffer,
            )
        else:
            self.self_attn = MiniMaxM3Attention(
@@ -681,6 +696,12 @@ class MiniMaxM3DecoderLayer(nn.Module):
                layer_id=layer_id,
                quant_config=quant_config,
                prefix=f"{prefix}.block_sparse_moe",
+                # Defer the MoE all-reduce only when it can be fused into a
+                # following GemmaRMSNorm
+                reduce_results=(
+                    vllm_config.parallel_config.pipeline_parallel_size > 1
+                    or is_mtp_block
+                ),
            )
        else:
            self.mlp = MiniMaxM3MLP(
@@ -747,17 +768,43 @@ class MiniMaxM3Model(nn.Module, EagleModelMixin):
            prefix=f"{prefix}.embed_tokens",
        )

+        # Reserved top-k indices buffer shared by all sparse-attention indexer
+        # layers (mirrors DeepseekV4); kept at a stable address so the indexer's
+        # top-k output survives cudagraph capture/replay. Shape matches the
+        # per-head index top-k output [num_index_heads, total_q, topk].
+        sparse_cfg = getattr(config, "sparse_attention_config", None)
+        if sparse_cfg is not None:
+            tp_size = get_tensor_model_parallel_world_size()
+            num_index_heads = max(1, sparse_cfg["sparse_num_index_heads"] // tp_size)
+            max_num_batched_tokens = vllm_config.scheduler_config.max_num_batched_tokens
+            padded_num_tokens = (max_num_batched_tokens + 3) // 4 * 4
+            self.topk_indices_buffer = torch.empty(
+                num_index_heads,
+                padded_num_tokens,
+                sparse_cfg["sparse_topk_blocks"],
+                dtype=torch.int32,
+            )
+        else:
+            self.topk_indices_buffer = None
+
        self.start_layer, self.end_layer, self.layers = make_layers(
            config.num_hidden_layers,
            lambda prefix: MiniMaxM3DecoderLayer(
                vllm_config=vllm_config,
                prefix=prefix,
+                topk_indices_buffer=self.topk_indices_buffer,
            ),
            prefix=f"{prefix}.layers",
        )

        self.norm = MiniMAXGemmaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)

+        # The final decoder layer has no next layer, so its deferred all-reduce is
+        # completed here in the model norm.
+        self.fuse_final_allreduce = (
+            vllm_config.parallel_config.pipeline_parallel_size == 1
+        )
+
    def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
        return self.embed_tokens(input_ids)

@@ -781,7 +828,12 @@ class MiniMaxM3Model(nn.Module, EagleModelMixin):
                aux_hidden_states, idx + 1, hidden_states, residual
            )

-        hidden_states, _ = self.norm(hidden_states, residual)
+        if self.fuse_final_allreduce and residual is not None:
+            hidden_states, _ = fused_allreduce_gemma_rms_norm(
+                hidden_states, residual, self.norm
+            )
+        else:
+            hidden_states, _ = self.norm(hidden_states, residual)

        if len(aux_hidden_states) > 0:
            return hidden_states, aux_hidden_states
@@ -29,7 +29,6 @@ class MiniMaxM3SparseMSAImpl(MiniMaxM3SparseImpl):
        layer: AttentionLayer,
        query: torch.Tensor,
        kv_cache: torch.Tensor,
-        topk_idx: tuple[torch.Tensor | None, torch.Tensor | None],
        output: torch.Tensor,
    ) -> torch.Tensor:
        attn_metadata = get_forward_context().attn_metadata
@@ -37,10 +36,12 @@ class MiniMaxM3SparseMSAImpl(MiniMaxM3SparseImpl):
            return output  # profiling run; caches unbound
        main_md = attn_metadata[layer.layer_name]  # type: ignore[attr-defined]
        assert isinstance(main_md, MiniMaxM3SparseMetadata)
-        decode_topk, prefill_topk = topk_idx

        nd = main_md.num_decode_tokens
        num_tokens = main_md.num_actual_tokens
+        # Indexer top-k from the shared buffer: decode [:, :nd], prefill [:, nd:].
+        topk = layer.topk_indices_buffer  # type: ignore[attr-defined]
+        assert topk is not None
        hd = self.head_size
        q = query[:num_tokens].view(-1, self.num_heads, hd)
        out = output[:num_tokens].view(-1, self.num_heads, hd)
@@ -51,11 +52,11 @@ class MiniMaxM3SparseMSAImpl(MiniMaxM3SparseImpl):
        # Decode [:nd]: Triton split-K placeholder (no MSA decode yet).
        if main_md.num_decodes > 0:
            d = main_md.decode
-            assert d is not None and decode_topk is not None
+            assert d is not None
            minimax_m3_sparse_attn_decode(
                q[:nd],
                kv_cache,
-                decode_topk,
+                topk[:, :nd, :],
                d.block_table,
                d.seq_lens,
                self.num_kv_heads,
@@ -72,7 +73,9 @@ class MiniMaxM3SparseMSAImpl(MiniMaxM3SparseImpl):
            )

            p = main_md.prefill
-            assert p is not None and prefill_topk is not None
+            assert p is not None
+            # build_k2q_csr() doesn't support strided topk buffer
+            prefill_topk = topk[:, nd:num_tokens, :]
            qp = q[nd:]
            k_cache = kv_cache[:, 0].transpose(1, 2)
            v_cache = kv_cache[:, 1].transpose(1, 2)
Author	SHA1	Message	Date
Zijing LiuandYongye Zhu	8b00f41237	[MiniMax-M3] Retune decode indexer launch grid for occupancy (#130 ) Signed-off-by: Zijing Liu <liuzijing2014@gmail.com>	2026-06-25 18:14:52 +00:00
Xin LiandGitHub	0586a20184	Add MiniMax-M3 modelopt nvfp4 support (#46380 ) Signed-off-by: Xin Li <xinli@nvidia.com>	2026-06-25 11:11:41 -07:00
Yongye ZhuandClaude Opus 4.8	7a672233eb	Bump flashinfer-jit-cache to 0.6.13rc2 Align flashinfer-jit-cache with flashinfer-python/flashinfer-cubin, which are already pinned to 0.6.13rc2 in requirements/cuda.txt. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: Yongye Zhu <zyy1102000@gmail.com>	2026-06-23 00:01:30 +00:00
Yongye Zhu	ea890c8066	remove contiguous constraint for prefill_topk Signed-off-by: Yongye Zhu <zyy1102000@gmail.com>	2026-06-22 05:21:19 +00:00
Thien TranandRoger Wang	8717eccaf4	Fix non-contiguous topk buffer for MSA sparse GQA prefill (#105 ) (cherry picked from commit eff83fa5c1ecf47678a262ed85f5d751b5cf1b67)	2026-06-20 13:53:02 -07:00
Zijing LiuandRoger Wang	9ac6a8ed61	NVFP4 TRTLLM-Gen fused MoE: clamped SwiGLU-OAI support + NVFP4 (#97 ) Signed-off-by: Zijing Liu <liuzijing2014@gmail.com> (cherry picked from commit 20bafcab50f1f63de1e0a5f810d7b844888fc0ac)	2026-06-20 13:52:56 -07:00
Yongye ZhuandRoger Wang	ef762c26e5	[MoE] Plumb gemm1_alpha/beta/clamp_limit into TRT-LLM FP8 MoE The FlashInfer FP8 block-scale MoE kernels (trtllm_fp8_block_scale_moe and trtllm_fp8_block_scale_routed_moe) accept optional per-expert SwiGLU parameters gemm1_alpha, gemm1_beta, gemm1_clamp_limit that realize the OAI SwiGLU variant for MXFP8. The FP8 experts did not pass them. - Build per-expert alpha/beta/clamp tensors in TrtLlmFp8ExpertsBase and pass them to the block-scale (monolithic) and routed (modular) kernels, mirroring the existing MXFP4 experts. - Forward gemm1_alpha/gemm1_beta from layer.swiglu_alpha/beta in the MXFP8-capable FP8 quant-config builders (fp8, online mxfp8, compressed-tensors mxfp8). Previously only swiglu_limit was forwarded; only ModelOptMxfp8 forwarded all three, so other paths silently dropped alpha/beta. - Add swigluoai_uninterleave to the supported activations (maps to FlashInfer Swiglu; the OAI behavior comes from the gemm1_* params). Signed-off-by: Yongye Zhu <zyy1102000@gmail.com> Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> (cherry picked from commit `f4b8540570`)	2026-06-20 13:51:49 -07:00
Yongye ZhuandRoger Wang	985d00a751	[Test] M3 indexer: add sm_scale arg to _reference_index_topk Fixes mypy call-arg errors: two fp8 call sites pass sm_scale (by keyword and positionally) that the reference signature lacked. Add sm_scale (default 1.0) and apply it to the score; top-k selection is invariant to a positive scalar, so existing callers are unchanged. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: Yongye Zhu <zyy1102000@gmail.com> (cherry picked from commit `522f262265`)	2026-06-20 13:51:43 -07:00
Yongye ZhuandRoger Wang	70ae0bb9e5	run indexer eager Signed-off-by: Yongye Zhu <zyy1102000@gmail.com> (cherry picked from commit `e6dd784466`)	2026-06-20 13:51:43 -07:00
Thien TranandRoger Wang	18d9bd9399	[Build] fmha_sm100: install indexer csrc/cutlass + package_data globs PR #93 integrated the MSA indexer but did not update the cmake install rules / setup.py package_data to vendor the new files. Add install rules for csrc/, cutlass/include, cutlass/tools/util/include and the matching package_data globs. (cherry picked from commit ea59ac42ab8e4a68f3bc362034de5199b59246df) Signed-off-by: Yongye Zhu <zyy1102000@gmail.com> (cherry picked from commit `e6fe1ecd64`)	2026-06-20 13:51:43 -07:00
Yongye ZhuandRoger Wang	28ee7f57df	[Model] M3 MSA indexer: align decode path with rebased index API The MSA builder/impl were written against the pre-#45743 indexer decode API. After rebasing onto main, MiniMaxM3IndexerDecodeMetadata gained a required max_decode_query_len field and minimax_m3_index_decode dropped its sm_scale parameter, so the MSA path crashed at engine init during cudagraph profiling (missing max_decode_query_len). - Pass max_decode_query_len when building decode metadata. - Drop the stale sm_scale arg and pass max_decode_query_len in the decode kernel call, matching the Triton builder/impl. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: Yongye Zhu <zyy1102000@gmail.com> (cherry picked from commit `0b70ba38a1`)	2026-06-20 13:51:43 -07:00
Yongye ZhuandRoger Wang	c12ed89131	[Model] M3 MSA indexer: swap decode to Triton fused kernel, add fp8 Route the MSA (SM100) indexer decode through the Triton fused minimax_m3_index_decode kernel (the same kernel the Triton indexer impl uses) instead of fmha_sm100's OnlyScore path. For q_len==1 decode the Triton kernel is a purpose-built vector x matrix score with a 256-way split-K and fused split-K top-k, which beats fmha's OnlyScore (wasted MMA tiles on a single query, 64-split cap) by ~1.1-3.7x in benchmarks. It is cudagraph-safe by construction (shape-constant split grids) and writes the shared topk_indices_buffer via out=. Prefill keeps fmha OnlyScore + the single-pass Triton top-k, where fmha is ~3-5x faster for the wide score. This drops the persistent fmha decode plan buffers, the num_kv_splits estimate, and the now-unused VLLM_M3_INDEXER_CONTEXT_LEN env. Add fp8 (e4m3) index KV cache support to the Triton decode score kernel: make the QK MMA accumulate in fp32 (out_dtype=tl.float32) so the per-block max score is exact for the e4m3 cache. Top-k is invariant to positive scalar scaling, so fp8 needs no scale tuning. Tests: tests/kernels/attention/test_minimax_m3.py (44 passed), including a new test_decode_index_topk_fp8 validating fp8 decode top-k vs a dequantized-fp32 reference. mypy + ruff clean. AI assistance was used. Co-authored-by: Claude <noreply@anthropic.com> Signed-off-by: Yongye Zhu <yongye@inferact.ai> Signed-off-by: Yongye Zhu <zyy1102000@gmail.com> (cherry picked from commit `0a102417bb`)	2026-06-20 13:51:43 -07:00
Yongye ZhuandRoger Wang	703e492aab	[Model] M3: tune indexer decode split heuristic from fp8 sweep estimate_num_kv_splits was a static replica of fmha_sm100's auto cost heuristic, which over-splits the decode score: an fp8 e4m3 cudagraph-timed split sweep shows the fixed-split (cudagraph-safe) path hits a sharp split-KV cliff above ~16 splits with >1 work row -- up to ~9x slower (e.g. reqs=2 ctx=65536: s64=153us vs s16=15us). The previous formula picked 52-64 splits there. Replace it with a measured, static fit (no kernel calls, so it stays cudagraph-stable and sync-free): - bs==1: min(64, kv_tiles) -- no cliff, fill with many splits. - bs>=2: min(num_sms // work_rows, cap, kv_tiles), where the cliff-safe cap is 16 through the ~60-100k context target and 32 for longer context (the cliff ceiling rises with context). Across an fp8 cudagraph sweep this is within ~3% of the best fixed split for every (batch, context) measured. The planner's own auto path is faster for bs>=2 (adaptive split distribution) but is catastrophic for bs==1 (up to ~11x) and needs a device->host sync, so it is unsuitable for the captured decode path. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: Yongye Zhu <yongye@inferact.ai> Signed-off-by: Yongye Zhu <zyy1102000@gmail.com> (cherry picked from commit `6dd73260e8`)	2026-06-20 13:51:43 -07:00
Yongye ZhuandRoger Wang	5f0bafab12	[Model] M3: read sparse-attend top-k from the persistent buffer Fixes a cudagraph-replay AssertionError on bs=1 decode: with the indexer in the captured segment and only the attend eager-broken, the (decode_topk, prefill_topk) Python tuple handed across the eager break was frozen at capture and replayed stale against the current step's metadata. The attend now reads its top-k directly from the shared persistent ``topk_indices_buffer`` (decode at [:, :nd], prefill at [:, nd:num_tokens]), sliced by the current step's metadata -- nothing crosses the eager break as a Python value. The indexer (captured) writes the buffer; the attend (eager) reads it; the breakable-cudagraph segment ordering guarantees the write precedes the read on replay. - MiniMaxM3SparseImpl.forward (+ Triton and MSA subclasses): drop the topk_idx arg, read layer.topk_indices_buffer. - nvidia model: store the buffer on the sparse-attention layer; narrow the eager break so the indexer is captured and only the attend (_run_sparse_attn) is eager-broken. - amd model: store the buffer on the layer; indexer + attend stay in one eager break (no capture), but the attend reads the buffer too. 42/42 in test_minimax_m3.py pass. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: Yongye Zhu <yongye@inferact.ai> Signed-off-by: Yongye Zhu <zyy1102000@gmail.com> (cherry picked from commit `627f730369`)	2026-06-20 13:51:43 -07:00
Yongye ZhuandRoger Wang	545393a3ab	[Model] M3: log indexer kernel selection in select_indexer_impl_cls Emit an info_once line naming the chosen indexer impl (MSA fmha_sm100 vs Triton) and the deciding inputs (topk_blocks, indexer_kv_dtype, sm100), so the active kernel path is visible at startup. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: Yongye Zhu <yongye@inferact.ai> Signed-off-by: Yongye Zhu <zyy1102000@gmail.com> (cherry picked from commit `7fd090bd9c`)	2026-06-20 13:51:43 -07:00
Yongye ZhuandRoger Wang	d71445af45	[Model] M3: route Triton indexer top-k through shared topk_indices_buffer Unify the Triton indexer impl onto the same persistent top-k buffer as the MSA impl, and thread the buffer through the AMD model too. - minimax_m3_index_decode gains an out= param (writes out[:, :total_q]); the merge kernel already writes via strides, so a buffer view works. - MiniMaxM3IndexerTritonImpl.forward writes decode ([:, :nd]) and prefill ([:, nd:]) into the shared topk_indices_buffer and returns views into it (no fresh per-step top-k allocations), matching the MSA impl. - amd/model.py: allocate the model-level topk_indices_buffer and thread it model -> decoder layer -> sparse attention -> indexer (mirrors nvidia). AMD keeps its eager break, so this is purely allocation reuse there. test_msa_indexer_impl_matches_triton now gives each impl its own buffer and asserts both decode/prefill outputs are views into it. 42/42 pass. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: Yongye Zhu <yongye@inferact.ai> Signed-off-by: Yongye Zhu <zyy1102000@gmail.com> (cherry picked from commit `06324dd3df`)	2026-06-20 13:51:43 -07:00
Yongye ZhuandRoger Wang	eb8e264edd	[Model] M3 MSA indexer: unify top-k buffer, drop numpy, GPU-only decode plan Follow-ups on the cudagraph-capturable MSA indexer: - Top-k: both decode and prefill now write into the single shared, persistent topk_indices_buffer (decode at [:, :nd], prefill at [:, nd:]) and return views into it -- no fresh per-step top-k allocations. - Build the decode plan + flat page table entirely with torch on-GPU: drop numpy and CpuGpuBuffer; segment offsets/lengths are computed via torch.cumsum into the persistent int32 buffers, and the request-major page table is scattered into the buffer via the on-GPU page indptr (the run bounds reads by indptr, so the full buffer is passed and no host page count is needed). - No GPU->CPU sync on the decode path: scalars come from host ints (num_decode_tokens // num_decodes), and seq_lens.cpu() is confined to the eager prefill branch. The impl forward (fmha OnlyScore + Triton top-k) was already sync-free. test_msa_indexer_impl_matches_triton now also asserts both outputs are views into the persistent buffer. 42/42 in test_minimax_m3.py pass. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: Yongye Zhu <yongye@inferact.ai> Signed-off-by: Yongye Zhu <zyy1102000@gmail.com> (cherry picked from commit `d2fbaf73c1`)	2026-06-20 13:51:43 -07:00
Yongye ZhuandRoger Wang	8f5070c447	[Model] Make MiniMax M3 MSA indexer cudagraph-capturable The SM100 (MSA) lightning indexer declared AttentionCGSupport.NEVER and ran eager: its fmha_sm100 score plan allocated fresh buffers every build() and the run allocated a fresh max_score / top-k each call. Make the decode side cudagraph-replay-safe: - Reserve persistent plan buffers in the builder __init__ (sized over a scan of decode sizes [1, max_num_seqs]); build() fills them in place and calls the plan kernel directly with a fixed, batch-size-only num_kv_splits (estimate_num_kv_splits, a uniform-context replica of the planner's auto-split math; tunable via VLLM_M3_INDEXER_CONTEXT_LEN). A positive split count takes the deterministic plan path with no device->host sync. - workspace_o / workspace_lse / cute_workspace are builder-owned dedicated tensors (not the shared global _alloc_workspace_buf cache) so a larger fmha call elsewhere can't realloc and move an address a captured graph baked. - max_score is a 1-D-backed contiguous [H, max_k_tiles, nnz_qo] view (a sliced 3-D buffer is non-contiguous; the kernel assumes contiguous); max_k_tiles pinned for a stable shape. - Top-k output goes to a model-level topk_indices_buffer (DeepSeek-V3.2 pattern), threaded model -> decoder layer -> sparse attention -> indexer; index_topk gains an out= param. - Narrow _run_attention's eager break so the indexer runs in the captured segment and only the sparse attention is eager-broken. Builder reports UNIFORM_BATCH; prefill stays eager. Verified: tests/kernels/attention/test_minimax_m3.py 42/42 pass (incl. impl-level parity with num_kv_splits > 1 over short context); GPQA accuracy matches eager. AI assistance (Claude Code) was used for this change. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: Yongye Zhu <yongye@inferact.ai> Signed-off-by: Yongye Zhu <zyy1102000@gmail.com> (cherry picked from commit `54810663ff`)	2026-06-20 13:51:43 -07:00
Yongye ZhuandRoger Wang	934fa2b599	[Model] Add fmha_sm100 MSA indexer backend + fp8 index cache for MiniMax M3 Add an SM100/Blackwell lightning-indexer impl that computes the per-128-block QK max-scores with fmha_sm100's score-only (OnlyScore) path and selects the top-k blocks with the existing Triton minimax_m3_index_topk kernel, mirroring how the main MSA attention pairs the SM100 attend with Triton. Decode and prefill requests are split manually (decode-first batch) and each side gets its own _fmha_sm100_plan / _fmha_sm100 call. Auto-selected on SM100 when topk_blocks in (4, 8, 16, 32) for both bf16 and fp8 index caches; falls back to the Triton indexer otherwise. The builder declares AttentionCGSupport.NEVER (eager; the attention is broken out of the graph by _run_attention). Extend the fused qknorm+rope+kv-insert kernel to optionally emit fp8 (e4m3) for the index-K cache and index-Q via a direct cast with no scale tensors (RMSNorm outputs are O(1) and scalar scales do not change top-k ordering). Only the index outputs go fp8; q/k/v and q_out stay bf16 and bit-identical to the existing path. MiniMaxM3IndexerCache now accepts fp8 caches and the model allocates index_q in the cache dtype. Tests: test_fmha_sm100_indexer_matches_reference (bf16/fp8 x prefill/decode) and test_msa_indexer_impl_matches_triton (full impl parity vs the Triton indexer through the real metadata builders); fp8 fused-kernel parity is covered in test_fused_minimax_m3_qknorm_rope_kv_insert. AI assistance (Claude Code) was used for this change. Signed-off-by: Yongye Zhu <yongye@inferact.ai> Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: Yongye Zhu <zyy1102000@gmail.com> (cherry picked from commit `bb4844dba3`)	2026-06-20 13:51:43 -07:00
Jee Jee LiandRoger Wang	713fb6cdb7	Done (#87 ) Signed-off-by: Jee Jee Li <jeejeelee@inferact.ai> (cherry picked from commit 1d7ba0da2f0fedf056b0073fe7774f03fe4616f0)	2026-06-20 13:51:19 -07:00