[CI/Build] Fix topk histogram build on SM75 (#46550)

Signed-off-by: Mohammad Miadh Angkad <176301910+mmangkad@users.noreply.github.com>
(cherry picked from commit 191826ec61)

.buildkite/test_areas/kernels.yaml

@@ -47,8 +47,10 @@ steps:
   - csrc/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu
   - vllm/models/deepseek_v4/common/ops/
   - tests/kernels/test_fused_deepseek_v4_qnorm_rope_kv_insert.py
+  - tests/kernels/test_top_k_per_row.py # it runs on Blackwell too - some kernels have arch-specific optimizations
   commands:
     - pytest -v -s kernels/test_fused_deepseek_v4_*.py
+    - pytest -v -s kernels/test_top_k_per_row.py
 
 - label: Deepseek V4 Kernel Test (B200)
   key: deepseek-v4-kernel-test-b200

CMakeLists.txt

@@ -382,6 +382,24 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
     "csrc/libtorch_stable/custom_all_reduce.cu"
     "csrc/libtorch_stable/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu")
 
+  if(VLLM_GPU_LANG STREQUAL "CUDA" AND
+     DEFINED CMAKE_CUDA_COMPILER_VERSION AND
+     CMAKE_CUDA_COMPILER_VERSION VERSION_GREATER_EQUAL 12.0)
+
+    if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
+      cuda_archs_loose_intersection(COOPERATIVE_TOPK_ARCHS
+        "9.0a;10.0f;10.1f;10.3f;11.0f;12.0f;12.1f" "${CUDA_ARCHS}")
+    else()
+      cuda_archs_loose_intersection(COOPERATIVE_TOPK_ARCHS
+        "9.0a;10.0a;10.1a;10.3a;12.0a;12.1a" "${CUDA_ARCHS}")
+    endif()
+
+    if(COOPERATIVE_TOPK_ARCHS)
+      list(APPEND VLLM_GPU_FLAGS "-DVLLM_ENABLE_COOPERATIVE_TOPK=1")
+
+    endif()
+  endif()
+
   if(VLLM_GPU_LANG STREQUAL "CUDA")
     SET(CUTLASS_ENABLE_HEADERS_ONLY ON CACHE BOOL "Enable only the header library")
 
@@ -498,6 +516,14 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
       SRCS "${VLLM_STABLE_EXT_SRC}"
       CUDA_ARCHS "${CUDA_ARCHS}")
 
+    if(COOPERATIVE_TOPK_ARCHS)
+      list(APPEND VLLM_STABLE_EXT_SRC
+        "csrc/libtorch_stable/cooperative_topk.cu")
+      set_gencode_flags_for_srcs(
+        SRCS "csrc/libtorch_stable/cooperative_topk.cu"
+        CUDA_ARCHS "${COOPERATIVE_TOPK_ARCHS}")
+    endif()
+
     # Only build Marlin kernels if we are building for at least some compatible archs.
     # Keep building Marlin for 9.0 as there are some group sizes and shapes that
     # are not supported by Machete yet.
@@ -1049,6 +1075,10 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
     target_compile_definitions(_C_stable_libtorch PRIVATE
       TORCH_TARGET_VERSION=0x020B000000000000ULL)
     target_compile_definitions(_C_stable_libtorch PRIVATE USE_CUDA)
+    if(COOPERATIVE_TOPK_ARCHS)
+      target_compile_definitions(_C_stable_libtorch PRIVATE
+        VLLM_ENABLE_COOPERATIVE_TOPK=1)
+    endif()
     # Needed by CUTLASS kernels
     target_compile_definitions(_C_stable_libtorch PRIVATE
       CUTLASS_ENABLE_DIRECT_CUDA_DRIVER_CALL=1)

cmake/external_projects/fmha_sm100.cmake

@@ -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)

csrc/libtorch_stable/cooperative_topk.cu

@@ -0,0 +1,146 @@
+// Cooperative cluster TopK for DeepSeek V3 sparse attention indexer.
+// See cooperative_topk.cuh for kernel implementation.
+
+#include <cuda_runtime.h>
+
+#include "torch_utils.h"
+
+#ifndef USE_ROCM
+  #include "cooperative_topk.cuh"
+namespace ct = vllm::cooperative;
+namespace hist4096 = vllm::topk_histogram_4096;
+#endif
+
+#ifndef USE_ROCM
+template <uint32_t TopK, uint32_t CS>
+void launch_cooperative_cluster(ct::CooperativeTopKParams<TopK>& params,
+                                size_t smem, cudaStream_t stream) {
+  auto kernel = []() {
+    if constexpr (CS == 16) {
+      return &ct::cooperative_topk_cs16<TopK>;
+    } else if constexpr (CS == 8) {
+      return &ct::cooperative_topk_cs8<TopK>;
+    } else {
+      static_assert(CS == 4, "unsupported cooperative_topk cluster size");
+      return &ct::cooperative_topk_cs4<TopK>;
+    }
+  }();
+  if constexpr (CS > 8) {
+    cudaFuncSetAttribute(kernel, cudaFuncAttributeNonPortableClusterSizeAllowed,
+                         1);
+  }
+  cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize,
+                       smem);
+
+  cudaLaunchConfig_t cfg = {};
+  cfg.gridDim = dim3(params.num_rows, CS);
+  cfg.blockDim = dim3(hist4096::kBlockSize);
+  cfg.dynamicSmemBytes = smem;
+  cfg.stream = stream;
+  cudaLaunchAttribute attrs[1];
+  attrs[0].id = cudaLaunchAttributeClusterDimension;
+  attrs[0].val.clusterDim = {1, CS, 1};
+  cfg.numAttrs = 1;
+  cfg.attrs = attrs;
+  cudaError_t err = cudaLaunchKernelEx(&cfg, kernel, params);
+  STD_TORCH_CHECK(err == cudaSuccess,
+                  "cooperative_topk launch failed: ", cudaGetErrorString(err));
+}
+
+template <uint32_t TopK>
+void launch_cooperative_topk_impl(const torch::stable::Tensor& logits,
+                                  const torch::stable::Tensor& lengths,
+                                  torch::stable::Tensor& output,
+                                  torch::stable::Tensor& workspace,
+                                  int64_t max_seq_len) {
+  (void)max_seq_len;  // Kept for signature parity with persistent_topk.
+  const int64_t num_rows = logits.size(0);
+  const cudaStream_t stream = get_current_cuda_stream();
+
+  const uint32_t stride = static_cast<uint32_t>(logits.stride(0));
+  // 32 = max clusters for CS=4 (32 x 4 = 128 CTAs = 66% of SMs, leaves
+  // headroom)
+  STD_TORCH_CHECK(
+      num_rows <= 32,
+      "cooperative_topk supports <=32 rows; use persistent_topk for "
+      "larger batches");
+
+  STD_TORCH_CHECK(stride % 4 == 0,
+                  "cooperative_topk: stride must be multiple of 4 for TMA "
+                  "alignment, got stride (max_model_len)=",
+                  stride);
+
+  STD_TORCH_CHECK(workspace.is_cuda(), "workspace must be CUDA tensor");
+  STD_TORCH_CHECK(
+      workspace.scalar_type() == torch::headeronly::ScalarType::Byte,
+      "workspace must be uint8");
+
+  ct::CooperativeTopKParams<TopK> params;
+  params.input = logits.const_data_ptr<float>();
+  params.output = output.mutable_data_ptr<int32_t>();
+  params.lengths = lengths.const_data_ptr<int32_t>();
+  params.num_rows = static_cast<uint32_t>(num_rows);
+  params.stride = stride;
+  params.tie_ws =
+      reinterpret_cast<hist4096::Tie*>(workspace.mutable_data_ptr<uint8_t>());
+
+  constexpr uint32_t kTieWsPerRow =
+      TopK <= hist4096::kBlockSize ? hist4096::kMaxTies : TopK;
+  STD_TORCH_CHECK(
+      workspace.size(0) >=
+          static_cast<int64_t>(num_rows * kTieWsPerRow * sizeof(hist4096::Tie)),
+      "workspace too small");
+
+  const bool supports_cluster16 = get_device_prop()->major >= 10;
+  if (num_rows <= 4 && supports_cluster16) {
+    launch_cooperative_cluster<TopK, 16>(params, ct::kSmemSize8, stream);
+  } else if (num_rows <= 8) {
+    launch_cooperative_cluster<TopK, 8>(params, ct::kSmemSize8, stream);
+  } else {
+    launch_cooperative_cluster<TopK, 4>(params, ct::kSmemSize4, stream);
+  }
+}
+#endif  // USE_ROCM
+
+void cooperative_topk(const torch::stable::Tensor& logits,
+                      const torch::stable::Tensor& lengths,
+                      torch::stable::Tensor& output,
+                      torch::stable::Tensor& workspace, int64_t k,
+                      int64_t max_seq_len) {
+#ifndef USE_ROCM
+  STD_TORCH_CHECK(logits.is_cuda(), "logits must be CUDA tensor");
+  STD_TORCH_CHECK(lengths.is_cuda(), "lengths must be CUDA tensor");
+  STD_TORCH_CHECK(output.is_cuda(), "output must be CUDA tensor");
+  STD_TORCH_CHECK(logits.scalar_type() == torch::headeronly::ScalarType::Float,
+                  "Only float32 supported");
+  STD_TORCH_CHECK(lengths.scalar_type() == torch::headeronly::ScalarType::Int,
+                  "lengths must be int32");
+  STD_TORCH_CHECK(output.scalar_type() == torch::headeronly::ScalarType::Int,
+                  "output must be int32");
+  STD_TORCH_CHECK(logits.dim() == 2, "logits must be 2D");
+  STD_TORCH_CHECK(lengths.dim() == 1 || lengths.dim() == 2,
+                  "lengths must be 1D or 2D");
+  STD_TORCH_CHECK(lengths.is_contiguous(), "lengths must be contiguous");
+  STD_TORCH_CHECK(output.dim() == 2, "output must be 2D");
+  const int64_t num_rows = logits.size(0);
+  STD_TORCH_CHECK(lengths.numel() == num_rows, "lengths size mismatch");
+  STD_TORCH_CHECK(output.size(0) == num_rows && output.size(1) == k,
+                  "output size mismatch");
+  STD_TORCH_CHECK(
+      k == 512 || k == 1024 || k == 2048,
+      "cooperative_topk supports k=512, k=1024, or k=2048, got k=", k);
+
+  if (k == 512) {
+    launch_cooperative_topk_impl<512>(logits, lengths, output, workspace,
+                                      max_seq_len);
+  } else if (k == 1024) {
+    launch_cooperative_topk_impl<1024>(logits, lengths, output, workspace,
+                                       max_seq_len);
+  } else {
+    launch_cooperative_topk_impl<2048>(logits, lengths, output, workspace,
+                                       max_seq_len);
+  }
+#else
+  STD_TORCH_CHECK(false, "cooperative_topk is not supported on ROCm");
+#endif
+}

csrc/libtorch_stable/cooperative_topk.cuh

@@ -0,0 +1,593 @@
+/*
+ * Cooperative TopK kernel for DSA Indexer
+ */
+
+#ifndef COOPERATIVE_TOPK_CUH_
+#define COOPERATIVE_TOPK_CUH_
+
+#include <cooperative_groups.h>
+#include <cuda.h>
+#include <cuda_fp16.h>
+#include <cuda_runtime.h>
+#include <cuda/ptx>
+#include <algorithm>
+#include <cstdint>
+
+#include "topk_histogram_4096.cuh"
+
+namespace vllm {
+namespace cooperative {
+
+namespace hist4096 = topk_histogram_4096;
+
+constexpr uint32_t kHistBits = 10;
+constexpr uint32_t kHistBins = 1 << kHistBits;
+constexpr uint32_t kMaxTopK = 2048;
+
+constexpr uint32_t kElemPerStage = 16;
+constexpr uint32_t kSizePerStage =
+    kElemPerStage * hist4096::kBlockSize;  // 16384
+
+// CS=4 two-pass path uses two TMA stages as a double buffer.
+constexpr uint32_t kStreamingStagesCS4 = 2;
+// CS=8/16 fused paths keep all loaded TMA stages resident in smem.
+constexpr uint32_t kFusedStagesCS8 = 2;
+constexpr uint32_t kFusedStagesCS16 = 2;
+
+// CS=4 single-pass path
+constexpr uint32_t kMaxSinglePassStages = 3;
+constexpr uint32_t kMaxSinglePassPerBlock =
+    kMaxSinglePassStages * kSizePerStage;  // 49152
+
+template <uint32_t TopK = 1024>
+struct CooperativeTopKParams {
+  const float* __restrict__ input;
+  int32_t* __restrict__ output;
+  const int32_t* __restrict__ lengths;
+  hist4096::Tie* __restrict__ tie_ws;  // per-row tie workspace, see
+                                       // kTieWsPerRow
+  uint32_t num_rows, stride;
+};
+
+// ============================================================================
+// Cooperative helpers
+// ============================================================================
+
+// only CS adjacent lanes participate (sub-warp reduce), in opposite to
+// warp_reduce_sum_full
+template <uint32_t N>
+__device__ __forceinline__ uint32_t warp_reduce_sum_subN(uint32_t v) {
+#pragma unroll
+  for (uint32_t m = N >> 1; m > 0; m >>= 1)
+    v += __shfl_xor_sync(0xFFFFFFFF, v, m, 32);
+  return v;
+}
+
+// ============================================================================
+// Helpers
+// ============================================================================
+
+__device__ __forceinline__ uint32_t extract_coarse_bin(float x) {
+  return hist4096::extract_coarse_bin_N<kHistBits>(x);
+}
+
+__device__ __forceinline__ void mbarrier_init(uint64_t* a, uint32_t n) {
+  cuda::ptx::mbarrier_init(a, n);
+}
+__device__ __forceinline__ void mbarrier_wait(uint64_t* a, uint32_t p) {
+  while (!cuda::ptx::mbarrier_try_wait_parity(cuda::ptx::sem_relaxed,
+                                              cuda::ptx::scope_cta, a, p));
+}
+__device__ __forceinline__ void mbarrier_arrive_expect_tx(uint64_t* a,
+                                                          uint32_t t) {
+  cuda::ptx::mbarrier_arrive_expect_tx(cuda::ptx::sem_relaxed,
+                                       cuda::ptx::scope_cta,
+                                       cuda::ptx::space_shared, a, t);
+}
+__device__ __forceinline__ void tma_load(void* d, const void* s, uint32_t n,
+                                         uint64_t* m) {
+  cuda::ptx::cp_async_bulk(cuda::ptx::space_shared, cuda::ptx::space_global, d,
+                           s, n, m);
+}
+
+// ============================================================================
+// DSMEM histogram reduce
+// ============================================================================
+
+template <uint32_t CS>
+__device__ __forceinline__ void dsmem_hist_reduce(uint32_t* histogram) {
+  static_assert(kHistBins <= hist4096::kBlockSize);
+  auto cluster = cooperative_groups::this_cluster();
+  cluster.sync();
+  const auto tx = threadIdx.x;
+  const auto rank = blockIdx.y;
+  constexpr auto kLocal = kHistBins / CS;
+  const auto off = kLocal * rank;
+  if (tx < kHistBins) {
+    const auto addr = &histogram[off + tx / CS];
+    const auto src = cluster.map_shared_rank(addr, tx % CS);
+    *src = warp_reduce_sum_subN<CS>(*src);
+  }
+  cluster.sync();
+}
+
+// ============================================================================
+// Find threshold from reduced histogram
+// ============================================================================
+
+// NOTE: caller must ensure a cluster.sync() or __syncthreads() happened
+// before calling this, so warp_sum writes are visible across warps.
+// The first internal __syncthreads() is still needed for the warp_sum exchange.
+template <uint32_t TopK>
+__device__ __forceinline__ void find_threshold(uint32_t* histogram,
+                                               uint32_t* warp_sum,
+                                               uint32_t* counter_gt,
+                                               uint32_t* counter_eq,
+                                               hist4096::MatchBin* match) {
+  const auto tx = threadIdx.x;
+  const auto li = tx % hist4096::kWarpSize, wi = tx / hist4096::kWarpSize;
+  const auto value = tx < kHistBins ? histogram[tx] : 0;
+  const auto winc = hist4096::warp_inclusive_sum(li, value);
+  if (li == hist4096::kWarpSize - 1) warp_sum[wi] = winc;
+  __syncthreads();
+  const auto tmp = warp_sum[li];
+  const auto total = hist4096::warp_reduce_sum_full(tmp);
+  auto pfx = hist4096::warp_reduce_sum_full(li < wi ? tmp : 0) + winc;
+  const auto above = total - pfx;
+  if (tx < kHistBins && above < TopK && above + value >= TopK) {
+    *counter_gt = *counter_eq = 0;
+    *match = {.bin = tx, .above_count = above, .equal_count = value};
+  }
+  __syncthreads();
+}
+
+// Streams data through shared memory in chunks, processing each chunk before
+// loading the next overwrites each buffer after processing it (the epilogue
+// prefetch loads the next chunk into the same slot)
+template <typename SmemType, uint32_t kStages, uint32_t kBinBits,
+          bool kIsScatter>
+__device__ void tma_stream_pass(const float* scores, uint32_t length,
+                                uint32_t thr_bin, int32_t* indices,
+                                uint32_t* phases, SmemType* smem) {
+  const auto tx = threadIdx.x;
+  const auto lane = tx % hist4096::kWarpSize;
+  const auto ni =
+      (length + kSizePerStage - 1) / kSizePerStage;  // total stages needed
+  const auto la =
+      (length + 3u) & ~3u;  // length rounded up to float4 (TMA alignment)
+  const auto pass =
+      kIsScatter ? 1 : 0;  // barrier dim: [0] for histogram, [1] for scatter
+
+  // Prologue: issue initial TMA loads - prefill the pipeline
+  if (tx == 0) {
+#pragma unroll
+    for (uint32_t i = 0; i < kStages; i++) {
+      if (i >= ni) {
+        break;
+      }
+      const auto o = i * kSizePerStage;
+      const auto sz = min(kSizePerStage, la - o) * sizeof(float);
+      tma_load(smem->score_buffer[i], scores + o, sz,
+               &smem->barrier[pass][i]);  // cp.async.bulk is non-blocking
+      mbarrier_arrive_expect_tx(&smem->barrier[pass][i], sz);
+    }
+  }
+
+  // Main loop: process stages
+  for (uint32_t it = 0; it < ni; it++) {
+    const auto b = it % kStages;  // which buffer slot (0 or 1)
+    const auto o = it * kSizePerStage;
+    const auto sz = min(kSizePerStage, length - o);
+
+    if (lane == 0) {
+      mbarrier_wait(&smem->barrier[pass][b],
+                    phases[b] & 1);  // wait for the data
+    }
+    phases[b]++;  // advances the phase for next time this slot is reused
+    __syncwarp();
+
+#pragma unroll
+    for (uint32_t i = 0; i < kElemPerStage; i++) {
+      const auto li = tx + i * hist4096::kBlockSize;
+      if (li >= sz) {
+        break;
+      }
+      const auto sc = smem->score_buffer[b][li];
+      const auto bn = hist4096::extract_coarse_bin_N<kBinBits>(sc);
+      if constexpr (kIsScatter) {  // compile-time branch
+        // Scatter pass: place above-threshold and collect ties
+        const auto gi = o + li;
+        if (bn > thr_bin) {
+          indices[atomicAdd(&smem->counter_gt, 1)] = gi;
+        } else if (bn == thr_bin) {
+          const auto p = atomicAdd(&smem->counter_eq, 1);
+          if (p < hist4096::kMaxTies) {
+            smem->tie_buffer[p] = {gi, sc};
+          }
+        }
+      } else {
+        // Histogram pass: just count
+        atomicAdd(&smem->histogram[bn], 1);
+      }
+    }
+    __syncthreads();  // ensures all threads finished processing their buffer
+                      // before next TMA load
+
+    // Epilogue: issue next TMA load
+    if (tx == 0 && it + kStages < ni) {
+      const auto no = (it + kStages) * kSizePerStage;
+      const auto nsz = min(kSizePerStage, la - no) * sizeof(float);
+      tma_load(smem->score_buffer[b], scores + no, nsz,
+               &smem->barrier[pass][b]);
+      mbarrier_arrive_expect_tx(&smem->barrier[pass][b], nsz);
+    }
+  }
+}
+
+// ============================================================================
+// Fused path: single TMA pass, rescan smem for scatter
+// ============================================================================
+
+// Fused shared memory layout for cluster cooperative paths.
+// kPasses=1 for single-pass (CS=8, CS=4 singlepass), kPasses=2 for two-pass
+// (CS=4).
+template <uint32_t kStages, uint32_t kPasses = 1>
+struct SmemFused {
+  uint64_t barrier[kPasses][kStages];
+  alignas(128) uint32_t counter_gt;
+  alignas(128) uint32_t counter_eq;
+  alignas(128) hist4096::MatchBin match;
+  uint32_t warp_sum[hist4096::kNumWarps];
+  union {
+    uint32_t histogram[kHistBins];
+    hist4096::Tie tie_buffer[kMaxTopK];
+  };
+  alignas(128) float score_buffer[kStages][kSizePerStage];
+};
+
+using Smem8 = SmemFused<kFusedStagesCS8>;
+using Smem16 = SmemFused<kFusedStagesCS16>;
+using Smem4 = SmemFused<kStreamingStagesCS4, 2>;
+using SmemSinglePass = SmemFused<kMaxSinglePassStages>;
+
+// Cluster-cooperative large path.
+// kFused=true: all TMA stages resident, single-pass histogram + scatter (rescan
+// from smem). kFused=false: TMA double-buffer streaming, two passes (histogram
+// then scatter).
+template <uint32_t TopK, uint32_t CS, typename SmemType, bool kFused>
+__device__ void large_topk(const float* __restrict__ row_input,
+                           int32_t* __restrict__ row_output, uint32_t seq_len,
+                           uint32_t* phases, hist4096::Tie* tie_ws) {
+  const auto rank = blockIdx.y;  // this block's position in cluster
+  const auto tx = threadIdx.x;
+  const auto lane = tx % hist4096::kWarpSize;
+
+  extern __shared__ uint8_t smem_raw[];
+  auto* smem = reinterpret_cast<SmemType*>(smem_raw);
+  int32_t* s_topk = reinterpret_cast<int32_t*>(smem_raw + sizeof(SmemType));
+
+  // Partition row across cluster ranks
+  constexpr uint32_t kAlign = 4;
+  const auto units =
+      (seq_len + kAlign - 1) / kAlign;  // float4-aligned element count
+  const auto base = units / CS, extra = units % CS;  // elements per block
+  const auto lu = base + (rank < extra ? 1u : 0u);   // remainder blocks
+  const auto ou =
+      rank * base + min(rank, extra);  // this block's count (load-balanced)
+  const auto my_start = ou * kAlign;   // global start offset
+  const auto my_len = min(my_start + lu * kAlign, seq_len) -
+                      my_start;  // actual length of this block
+  const auto num_iters =
+      (my_len + kSizePerStage - 1) / kSizePerStage;  // TMA stages needed
+  const auto len_aligned = (my_len + 3u) & ~3u;
+
+  if constexpr (kFused) {
+    // Fused init + TMA prologue
+    if (tx < kHistBins) {
+      smem->histogram[tx] = 0;  // all threads zero histogram
+    }
+    if (tx == 0) {  // thread 0 issues TMA - then all threads continue working
+                    // until mbarrier sync
+      smem->counter_gt = 0;
+      smem->counter_eq = 0;
+      for (uint32_t i = 0; i < num_iters; i++) {
+        const auto off = i * kSizePerStage;
+        const auto sz = min(kSizePerStage, len_aligned - off) * sizeof(float);
+        tma_load(smem->score_buffer[i], row_input + my_start + off, sz,
+                 &smem->barrier[0][i]);  // cp.async.bulk of size kSizePerStage
+                                         // × sizeof(float)
+        mbarrier_arrive_expect_tx(&smem->barrier[0][i], sz);
+      }
+    }
+    __syncthreads();
+
+    // Histogram build. ILP unroll-by-2, no inter-stage sync
+    for (uint32_t iter = 0; iter < num_iters; iter++) {
+      const auto off = iter * kSizePerStage;
+      const auto sz = min(kSizePerStage, my_len - off);
+      if (lane == 0) {
+        mbarrier_wait(&smem->barrier[0][iter],
+                      phases[iter] & 1);  // wait for TMA
+      }
+      phases[iter]++;
+      __syncwarp();
+#pragma unroll
+      for (uint32_t i = 0; i < kElemPerStage; i += 2) {
+        const auto li0 = tx + i * hist4096::kBlockSize;
+        const auto li1 = tx + (i + 1) * hist4096::kBlockSize;
+        if (li0 >= sz) {
+          break;
+        }
+        const auto b0 = extract_coarse_bin(smem->score_buffer[iter][li0]);
+        if (li1 < sz) {
+          const auto b1 = extract_coarse_bin(smem->score_buffer[iter][li1]);
+          atomicAdd(&smem->histogram[b0], 1);
+          atomicAdd(&smem->histogram[b1], 1);
+        } else {
+          atomicAdd(&smem->histogram[b0], 1);
+        }
+      }
+    }
+  } else {
+    // Twopass: init then stream histogram pass
+    if (tx < kHistBins) {
+      smem->histogram[tx] = 0;
+    }
+    if (tx == 0) {
+      smem->counter_gt = 0;
+      smem->counter_eq = 0;
+    }
+    __syncthreads();
+    tma_stream_pass<SmemType, kStreamingStagesCS4, kHistBits, false>(
+        row_input + my_start, my_len, 0, nullptr, phases, smem);
+  }
+
+  // DSMEM all-reduce + find threshold
+  dsmem_hist_reduce<CS>(
+      smem->histogram);  // each block histogram is summed across all CS blocks
+  find_threshold<TopK>(smem->histogram, smem->warp_sum, &smem->counter_gt,
+                       &smem->counter_eq, &smem->match);
+
+  const auto thr = smem->match.bin;
+
+  if constexpr (kFused) {
+    // Fused scatter: rescan score_buffer (still in smem)
+    for (uint32_t iter = 0; iter < num_iters; iter++) {
+      const auto off = iter * kSizePerStage;
+      const auto sz = min(kSizePerStage, my_len - off);
+#pragma unroll
+      for (uint32_t i = 0; i < kElemPerStage; i++) {
+        const auto li = tx + i * hist4096::kBlockSize;
+        if (li >= sz) {
+          break;
+        }
+        const auto score = smem->score_buffer[iter][li];  // still in smem
+        const auto bin = extract_coarse_bin(score);
+        const auto gidx = off + li;
+        if (bin > thr) {
+          s_topk[atomicAdd(&smem->counter_gt, 1)] = gidx;  // above -> s_topk
+        } else if (bin == thr) {
+          const auto p = atomicAdd(&smem->counter_eq,
+                                   1);  // equal -> ties (later refinement)
+          if (p < hist4096::kMaxTies) {
+            smem->tie_buffer[p] = {gidx, score};
+          }
+        }
+      }
+    }
+    __syncthreads();
+  } else {
+    // Twopass scatter: re-stream data via TMA
+    uint32_t scatter_phases[kStreamingStagesCS4] = {0, 0};
+    tma_stream_pass<SmemType, kStreamingStagesCS4, kHistBits, true>(
+        row_input + my_start, my_len, thr, s_topk, scatter_phases, smem);
+  }
+
+  // Output collection via DSMEM prefix sum
+  constexpr uint32_t kAboveBits = 16;
+  constexpr uint32_t kAboveMask = (1 << kAboveBits) - 1;
+  static_assert(kAboveMask >= TopK);
+  static_assert(kAboveMask >= kMaxSinglePassPerBlock,
+                "kAboveBits must cover max per-block element count");
+
+  const uint32_t la = smem->counter_gt;
+  const uint32_t le_full = smem->counter_eq;
+  const uint32_t le =
+      min(le_full, hist4096::kMaxTies);  // written smem tie_buffer entries
+
+  __shared__ uint32_t s_local_counts[CS];
+  __shared__ uint32_t s_prefix_packed;
+  __shared__ uint32_t s_total_above, s_total_equal;
+
+  auto cluster = cooperative_groups::this_cluster();
+  if (tx < CS) {
+    // Pack written tie counts into 32-bit: (equal << 16) | above.
+    // `le_full` may exceed the per-block tie buffer cap; using it here creates
+    // holes in tie_ws and can make TopK=2048 refine unwritten workspace slots.
+    const uint32_t packed = (le << kAboveBits) | la;
+    const auto dst = cluster.map_shared_rank(s_local_counts, tx);
+    dst[rank] = packed;  // write my count to every block's s_local_counts[rank]
+  }
+  cluster.sync();
+
+  // Thread 0 computes serial prefix sum
+  if (tx == 0) {
+    uint32_t prefix = 0, ta = 0, te = 0;
+    for (uint32_t i = 0; i < CS; i++) {
+      if (i == rank) {
+        s_prefix_packed = prefix;  // my prefix
+      }
+      ta += s_local_counts[i] & kAboveMask;   // total above
+      te += s_local_counts[i] >> kAboveBits;  // total equal
+      prefix += s_local_counts[i];
+    }
+    s_total_above = ta;
+    s_total_equal = te;
+  }
+  __syncthreads();
+
+  const uint32_t prefix_above = s_prefix_packed & kAboveMask;
+  const uint32_t prefix_equal = s_prefix_packed >> kAboveBits;
+
+  // Write to global output
+  for (uint32_t i = tx; i < la; i += hist4096::kBlockSize) {
+    // indices are placed contiguously starting at prefix_above
+    row_output[prefix_above + i] =
+        s_topk[i] + my_start;  // my_start: block-local -> row-global index
+  }
+  for (uint32_t i = tx; i < le; i += hist4096::kBlockSize) {
+    const auto t = smem->tie_buffer[i];
+    uint32_t p = s_total_above + prefix_equal + i;
+    if (p < TopK) {
+      row_output[p] = t.idx + my_start;
+    }
+    uint32_t tp = prefix_equal + i;
+    if (tp < (TopK <= hist4096::kBlockSize ? hist4096::kMaxTies : TopK)) {
+      tie_ws[tp] = hist4096::Tie{t.idx + my_start, t.score};
+    }
+  }
+
+  // Tie refinement
+  cooperative_groups::this_cluster().sync();
+  if (rank != 0) {  // only rank 0 does tie refinement
+    return;
+  }
+  if (s_total_above + s_total_equal <= TopK) {  // no ties to refine
+    return;
+  }
+
+  // Tie-breaking uses FP32 (4-round radix sort)
+  if constexpr (TopK <= hist4096::kBlockSize) {
+    // copy ties from tie_ws back to smem, then refine
+    const uint32_t num_ties = min(s_total_equal, hist4096::kMaxTies);
+    // TODO (roberto): could vectorize with uint2 (8 bytes = exactly one Tie)
+    for (uint32_t i = tx; i < num_ties; i += hist4096::kBlockSize) {
+      smem->tie_buffer[i] = hist4096::Tie{tie_ws[i].idx, tie_ws[i].score};
+    }
+    __syncthreads();
+    hist4096::tie_handle<TopK>(smem->tie_buffer, num_ties, s_total_above,
+                               row_output, smem);
+  } else {
+    // TopK=2048: process directly from tie_ws (GMEM)
+    const uint32_t num_ties = min(s_total_equal, static_cast<uint32_t>(TopK));
+    hist4096::tie_handle_large<TopK>(tie_ws, num_ties, s_total_above,
+                                     row_output, smem);
+  }
+}
+
+// ============================================================================
+// Adapted from https://github.com/sgl-project/sglang/pull/23600
+// sgl-project/sglang
+// (python/sglang/jit_kernel/include/sgl_kernel/deepseek_v4/topk/)
+// ============================================================================
+
+template <uint32_t TopK, uint32_t CS>
+__device__ void cooperative_topk_body(CooperativeTopKParams<TopK> params) {
+  const auto rank = blockIdx.y, row = blockIdx.x, tx = threadIdx.x;
+  const auto sl = params.lengths[row];
+  int32_t* out = params.output + row * TopK;
+  const float* in = params.input + row * params.stride;
+
+  // Trivial: seq_len <= TopK
+  if (sl <= static_cast<int32_t>(TopK)) {
+    if (rank == 0) {
+      for (uint32_t i = tx; i < TopK; i += hist4096::kBlockSize) {
+        out[i] = (i < static_cast<uint32_t>(sl)) ? static_cast<int32_t>(i) : -1;
+      }
+    }
+    return;
+  }
+
+  // Short-Medium path: histogram_4096_topk on rank 0 only - all data fits in RF
+  if (sl <= static_cast<int32_t>(hist4096::kHist4096MaxLen)) {
+    if (rank == 0) {
+      extern __shared__ uint8_t sr[];
+      hist4096::histogram_4096_topk<TopK, 12>(
+          in, out, sl, sr);  // 4096-bin (12-bit) histogram
+    }
+    return;
+  }
+
+  // Large path: init mbarriers + state, then dispatch fused or twopass
+  const uint32_t per_block =
+      (params.stride + CS - 1) / CS;  // how many elements per block
+  constexpr uint32_t kFusedMax = ((CS == 16)  ? kFusedStagesCS16
+                                  : (CS == 8) ? kFusedStagesCS8
+                                              : kMaxSinglePassStages) *
+                                 kSizePerStage;
+  const bool use_singlepass =
+      per_block <=
+      kFusedMax;  // single pass or TMA streaming: histogram+scatter
+
+  // Select smem type and stage count at compile time based on CS
+  constexpr uint32_t kFusedStages = (CS == 16)  ? kFusedStagesCS16
+                                    : (CS == 8) ? kFusedStagesCS8
+                                                : kMaxSinglePassStages;
+  using FusedSmem = SmemFused<kFusedStages>;
+
+  extern __shared__ uint8_t sr[];
+
+  constexpr uint32_t kTieWsPerRow =
+      TopK <= hist4096::kBlockSize ? hist4096::kMaxTies : TopK;
+  hist4096::Tie* row_tie_ws = params.tie_ws + row * kTieWsPerRow;
+
+  if (use_singlepass) {
+    auto* smem = reinterpret_cast<FusedSmem*>(sr);
+    const uint32_t sp_stages = (per_block + kSizePerStage - 1) / kSizePerStage;
+    if (tx < sp_stages) {
+      mbarrier_init(&smem->barrier[0][tx],
+                    1);  // init 1 barrier per TMA stage -
+                         // signal when async copies complete
+    }
+    __syncthreads();
+    uint32_t phases[kFusedStages] =
+        {};  // tracks the parity for mbarrier wait/arrive protocol
+    large_topk<TopK, CS, FusedSmem, true>(in, out, sl, phases, row_tie_ws);
+  } else {
+    // Two-pass: only CS=4 in practice (CS=8 always fits in singlepass)
+    auto* smem = reinterpret_cast<Smem4*>(sr);
+    if (tx < 2 * kStreamingStagesCS4) {
+      mbarrier_init(&smem->barrier[0][tx],
+                    1);  // init 2×2=4 barriers (2 passes × 2 stages)
+    }
+    __syncthreads();
+    uint32_t hp[kStreamingStagesCS4] = {0,
+                                        0};  // histogram+scatter pass counters
+    large_topk<TopK, CS, Smem4, false>(in, out, sl, hp, row_tie_ws);
+  }
+}
+
+template <uint32_t TopK>
+__global__ void __launch_bounds__(hist4096::kBlockSize, 1)
+    __cluster_dims__(1, 4, 1)
+        cooperative_topk_cs4(CooperativeTopKParams<TopK> params) {
+  cooperative_topk_body<TopK, 4>(params);
+}
+
+template <uint32_t TopK>
+__global__ void __launch_bounds__(hist4096::kBlockSize, 1)
+    __cluster_dims__(1, 8, 1)
+        cooperative_topk_cs8(CooperativeTopKParams<TopK> params) {
+  cooperative_topk_body<TopK, 8>(params);
+}
+
+template <uint32_t TopK>
+__global__ void __launch_bounds__(hist4096::kBlockSize, 1)
+    __cluster_dims__(1, 16, 1)
+        cooperative_topk_cs16(CooperativeTopKParams<TopK> params) {
+  cooperative_topk_body<TopK, 16>(params);
+}
+
+constexpr size_t kSmemSize4_base = sizeof(Smem4);
+constexpr size_t kSmemSize4_sp = sizeof(SmemSinglePass);
+constexpr size_t kSmemSize4 =
+    (kSmemSize4_base > kSmemSize4_sp ? kSmemSize4_base : kSmemSize4_sp) +
+    sizeof(int32_t) * 2048 + 128;
+constexpr size_t kSmemSize8 =
+    sizeof(SmemFused<kFusedStagesCS8>) + sizeof(int32_t) * 2048 + 128;
+
+}  // namespace cooperative
+
+}  // namespace vllm
+
+#endif  // COOPERATIVE_TOPK_CUH_

csrc/libtorch_stable/fused_minimax_m3_qknorm_rope_kv_insert_kernel.cu

@@ -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());

csrc/libtorch_stable/moe/dsv3_router_gemm_bf16_out.cu

@@ -286,3 +286,52 @@ template void invokeRouterGemmBf16Output<__nv_bfloat16, 15, 384, 7168>(
 
 template void invokeRouterGemmBf16Output<__nv_bfloat16, 16, 384, 7168>(
     __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
+
+// Template instantiations for GLM-5 (DEFAULT_NUM_EXPERTS, hidden_dim=6144)
+template void invokeRouterGemmBf16Output<__nv_bfloat16, 1, 256, 6144>(
+    __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
+
+template void invokeRouterGemmBf16Output<__nv_bfloat16, 2, 256, 6144>(
+    __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
+
+template void invokeRouterGemmBf16Output<__nv_bfloat16, 3, 256, 6144>(
+    __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
+
+template void invokeRouterGemmBf16Output<__nv_bfloat16, 4, 256, 6144>(
+    __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
+
+template void invokeRouterGemmBf16Output<__nv_bfloat16, 5, 256, 6144>(
+    __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
+
+template void invokeRouterGemmBf16Output<__nv_bfloat16, 6, 256, 6144>(
+    __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
+
+template void invokeRouterGemmBf16Output<__nv_bfloat16, 7, 256, 6144>(
+    __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
+
+template void invokeRouterGemmBf16Output<__nv_bfloat16, 8, 256, 6144>(
+    __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
+
+template void invokeRouterGemmBf16Output<__nv_bfloat16, 9, 256, 6144>(
+    __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
+
+template void invokeRouterGemmBf16Output<__nv_bfloat16, 10, 256, 6144>(
+    __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
+
+template void invokeRouterGemmBf16Output<__nv_bfloat16, 11, 256, 6144>(
+    __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
+
+template void invokeRouterGemmBf16Output<__nv_bfloat16, 12, 256, 6144>(
+    __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
+
+template void invokeRouterGemmBf16Output<__nv_bfloat16, 13, 256, 6144>(
+    __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
+
+template void invokeRouterGemmBf16Output<__nv_bfloat16, 14, 256, 6144>(
+    __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
+
+template void invokeRouterGemmBf16Output<__nv_bfloat16, 15, 256, 6144>(
+    __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
+
+template void invokeRouterGemmBf16Output<__nv_bfloat16, 16, 256, 6144>(
+    __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);

csrc/libtorch_stable/moe/dsv3_router_gemm_entry.cu

@@ -41,6 +41,7 @@ inline int getSMVersion() {
 static constexpr int DEFAULT_NUM_EXPERTS = 256;
 static constexpr int KIMI_K2_NUM_EXPERTS = 384;
 static constexpr int DEFAULT_HIDDEN_DIM = 7168;
+static constexpr int GLM_5_HIDDEN_DIM = 6144;
 
 template <typename T, int kNumTokens, int kNumExperts, int kHiddenDim>
 void invokeRouterGemmFloatOutput(float* output, T const* mat_a, T const* mat_b,
@@ -121,14 +122,21 @@ void dsv3_router_gemm(
 
   STD_TORCH_CHECK(mat_a.size(1) == mat_b.size(1),
                   "mat_a and mat_b must have the same hidden_dim");
-  STD_TORCH_CHECK(hidden_dim == DEFAULT_HIDDEN_DIM,
-                  "Expected hidden_dim=", DEFAULT_HIDDEN_DIM,
-                  ", but got hidden_dim=", hidden_dim);
+  STD_TORCH_CHECK(
+      hidden_dim == DEFAULT_HIDDEN_DIM || hidden_dim == GLM_5_HIDDEN_DIM,
+      "Expected hidden_dim=", DEFAULT_HIDDEN_DIM,
+      " or hidden_dim=", GLM_5_HIDDEN_DIM, ", but got hidden_dim=", hidden_dim);
   STD_TORCH_CHECK(
       num_experts == DEFAULT_NUM_EXPERTS || num_experts == KIMI_K2_NUM_EXPERTS,
       "Expected num_experts=", DEFAULT_NUM_EXPERTS,
       " or num_experts=", KIMI_K2_NUM_EXPERTS,
       ", but got num_experts=", num_experts);
+  // KIMI_K2_NUM_EXPERTS is only instantiated for the default hidden_dim.
+  STD_TORCH_CHECK(
+      hidden_dim == DEFAULT_HIDDEN_DIM || num_experts == DEFAULT_NUM_EXPERTS,
+      "hidden_dim=", GLM_5_HIDDEN_DIM,
+      " only supports num_experts=", DEFAULT_NUM_EXPERTS,
+      ", but got num_experts=", num_experts);
   STD_TORCH_CHECK(num_tokens >= 1 && num_tokens <= 16,
                   "currently num_tokens must be less than or equal to 16 for "
                   "router_gemm");
@@ -148,35 +156,49 @@ void dsv3_router_gemm(
 
   const cudaStream_t stream = get_current_cuda_stream(mat_a.get_device_index());
 
+  __nv_bfloat16 const* a_ptr =
+      reinterpret_cast<__nv_bfloat16 const*>(mat_a.data_ptr());
+  __nv_bfloat16 const* b_ptr =
+      reinterpret_cast<__nv_bfloat16 const*>(mat_b.data_ptr());
+
   if (output.scalar_type() == torch::headeronly::ScalarType::Float) {
-    if (num_experts == DEFAULT_NUM_EXPERTS) {
-      LoopUnroller<1, 16, DEFAULT_NUM_EXPERTS, DEFAULT_HIDDEN_DIM>::
-          unroll_float_output(
-              num_tokens, reinterpret_cast<float*>(output.mutable_data_ptr()),
-              reinterpret_cast<__nv_bfloat16 const*>(mat_a.data_ptr()),
-              reinterpret_cast<__nv_bfloat16 const*>(mat_b.data_ptr()), stream);
-    } else if (num_experts == KIMI_K2_NUM_EXPERTS) {
-      LoopUnroller<1, 16, KIMI_K2_NUM_EXPERTS, DEFAULT_HIDDEN_DIM>::
-          unroll_float_output(
-              num_tokens, reinterpret_cast<float*>(output.mutable_data_ptr()),
-              reinterpret_cast<__nv_bfloat16 const*>(mat_a.data_ptr()),
-              reinterpret_cast<__nv_bfloat16 const*>(mat_b.data_ptr()), stream);
+    float* out_ptr = reinterpret_cast<float*>(output.mutable_data_ptr());
+    if (hidden_dim == DEFAULT_HIDDEN_DIM) {
+      if (num_experts == DEFAULT_NUM_EXPERTS) {
+        LoopUnroller<1, 16, DEFAULT_NUM_EXPERTS,
+                     DEFAULT_HIDDEN_DIM>::unroll_float_output(num_tokens,
+                                                              out_ptr, a_ptr,
+                                                              b_ptr, stream);
+      } else {
+        LoopUnroller<1, 16, KIMI_K2_NUM_EXPERTS,
+                     DEFAULT_HIDDEN_DIM>::unroll_float_output(num_tokens,
+                                                              out_ptr, a_ptr,
+                                                              b_ptr, stream);
+      }
+    } else {  // GLM_5_HIDDEN_DIM
+      LoopUnroller<1, 16, DEFAULT_NUM_EXPERTS,
+                   GLM_5_HIDDEN_DIM>::unroll_float_output(num_tokens, out_ptr,
+                                                          a_ptr, b_ptr, stream);
     }
   } else if (output.scalar_type() == torch::headeronly::ScalarType::BFloat16) {
-    if (num_experts == DEFAULT_NUM_EXPERTS) {
-      LoopUnroller<1, 16, DEFAULT_NUM_EXPERTS, DEFAULT_HIDDEN_DIM>::
-          unroll_bf16_output(
-              num_tokens,
-              reinterpret_cast<__nv_bfloat16*>(output.mutable_data_ptr()),
-              reinterpret_cast<__nv_bfloat16 const*>(mat_a.data_ptr()),
-              reinterpret_cast<__nv_bfloat16 const*>(mat_b.data_ptr()), stream);
-    } else if (num_experts == KIMI_K2_NUM_EXPERTS) {
-      LoopUnroller<1, 16, KIMI_K2_NUM_EXPERTS, DEFAULT_HIDDEN_DIM>::
-          unroll_bf16_output(
-              num_tokens,
-              reinterpret_cast<__nv_bfloat16*>(output.mutable_data_ptr()),
-              reinterpret_cast<__nv_bfloat16 const*>(mat_a.data_ptr()),
-              reinterpret_cast<__nv_bfloat16 const*>(mat_b.data_ptr()), stream);
+    __nv_bfloat16* out_ptr =
+        reinterpret_cast<__nv_bfloat16*>(output.mutable_data_ptr());
+    if (hidden_dim == DEFAULT_HIDDEN_DIM) {
+      if (num_experts == DEFAULT_NUM_EXPERTS) {
+        LoopUnroller<1, 16, DEFAULT_NUM_EXPERTS,
+                     DEFAULT_HIDDEN_DIM>::unroll_bf16_output(num_tokens,
+                                                             out_ptr, a_ptr,
+                                                             b_ptr, stream);
+      } else {
+        LoopUnroller<1, 16, KIMI_K2_NUM_EXPERTS,
+                     DEFAULT_HIDDEN_DIM>::unroll_bf16_output(num_tokens,
+                                                             out_ptr, a_ptr,
+                                                             b_ptr, stream);
+      }
+    } else {  // GLM_5_HIDDEN_DIM
+      LoopUnroller<1, 16, DEFAULT_NUM_EXPERTS,
+                   GLM_5_HIDDEN_DIM>::unroll_bf16_output(num_tokens, out_ptr,
+                                                         a_ptr, b_ptr, stream);
     }
   }
 }

csrc/libtorch_stable/moe/dsv3_router_gemm_float_out.cu

@@ -286,3 +286,52 @@ template void invokeRouterGemmFloatOutput<__nv_bfloat16, 15, 384, 7168>(
 
 template void invokeRouterGemmFloatOutput<__nv_bfloat16, 16, 384, 7168>(
     float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
+
+// Template instantiations for GLM-5 (DEFAULT_NUM_EXPERTS, hidden_dim=6144)
+template void invokeRouterGemmFloatOutput<__nv_bfloat16, 1, 256, 6144>(
+    float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
+
+template void invokeRouterGemmFloatOutput<__nv_bfloat16, 2, 256, 6144>(
+    float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
+
+template void invokeRouterGemmFloatOutput<__nv_bfloat16, 3, 256, 6144>(
+    float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
+
+template void invokeRouterGemmFloatOutput<__nv_bfloat16, 4, 256, 6144>(
+    float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
+
+template void invokeRouterGemmFloatOutput<__nv_bfloat16, 5, 256, 6144>(
+    float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
+
+template void invokeRouterGemmFloatOutput<__nv_bfloat16, 6, 256, 6144>(
+    float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
+
+template void invokeRouterGemmFloatOutput<__nv_bfloat16, 7, 256, 6144>(
+    float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
+
+template void invokeRouterGemmFloatOutput<__nv_bfloat16, 8, 256, 6144>(
+    float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
+
+template void invokeRouterGemmFloatOutput<__nv_bfloat16, 9, 256, 6144>(
+    float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
+
+template void invokeRouterGemmFloatOutput<__nv_bfloat16, 10, 256, 6144>(
+    float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
+
+template void invokeRouterGemmFloatOutput<__nv_bfloat16, 11, 256, 6144>(
+    float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
+
+template void invokeRouterGemmFloatOutput<__nv_bfloat16, 12, 256, 6144>(
+    float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
+
+template void invokeRouterGemmFloatOutput<__nv_bfloat16, 13, 256, 6144>(
+    float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
+
+template void invokeRouterGemmFloatOutput<__nv_bfloat16, 14, 256, 6144>(
+    float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
+
+template void invokeRouterGemmFloatOutput<__nv_bfloat16, 15, 256, 6144>(
+    float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
+
+template void invokeRouterGemmFloatOutput<__nv_bfloat16, 16, 256, 6144>(
+    float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);

csrc/libtorch_stable/ops.h

@@ -343,6 +343,14 @@ void persistent_topk(const torch::stable::Tensor& logits,
                      torch::stable::Tensor& workspace, int64_t k,
                      int64_t max_seq_len);
 
+#ifdef VLLM_ENABLE_COOPERATIVE_TOPK
+void cooperative_topk(const torch::stable::Tensor& logits,
+                      const torch::stable::Tensor& lengths,
+                      torch::stable::Tensor& output,
+                      torch::stable::Tensor& workspace, int64_t k,
+                      int64_t max_seq_len);
+#endif
+
 void selective_scan_fwd(
     const torch::stable::Tensor& u, const torch::stable::Tensor& delta,
     const torch::stable::Tensor& A, const torch::stable::Tensor& B,

csrc/libtorch_stable/persistent_topk.cuh

@@ -11,6 +11,8 @@
 #include <cub/cub.cuh>
 #include <cstdint>
 
+#include "topk_histogram_4096.cuh"
+
 namespace vllm {
 namespace persistent {
 
@@ -935,8 +937,16 @@ __global__ void __launch_bounds__(kThreadsPerBlock, 2)
 }  // namespace persistent
 
 // ============================================================================
-// FlashInfer FilteredTopK (BS>32 dispatch) — float32 only.
-// Extracted from flashinfer_topk.cuh. Lives in namespace vllm (not persistent).
+// ============================================================================
+// Optimized FilteredTopK — single CTA per row for bs > 32.
+// Kept with persistent_topk so the portable fallback owns the non-cluster path.
+// ============================================================================
+namespace filtered_topk {
+
+namespace hist4096 = topk_histogram_4096;
+
+// ============================================================================
+// FilteredTopK — single CTA per row for bs > 32
 // Adapted from https://github.com/flashinfer-ai/flashinfer/pull/2215
 // ============================================================================
 
@@ -963,13 +973,6 @@ struct vec_t {
       data[i] = ptr[i];
     }
   }
-
-  FLASHINFER_INLINE void cast_store(T* ptr) const {
-#pragma unroll
-    for (size_t i = 0; i < N; ++i) {
-      ptr[i] = data[i];
-    }
-  }
 };
 #undef FLASHINFER_INLINE
 
@@ -1013,7 +1016,8 @@ constexpr size_t FILTERED_TOPK_SMEM_DYNAMIC =
  * \tparam IdType Index type (int32_t)
  * \tparam VEC_SIZE Vector size for input loads (1, 2, 4, or 8)
  */
-template <typename DType, typename IdType, int VEC_SIZE, uint32_t MAX_K = 2048>
+template <typename DType, typename IdType, int VEC_SIZE, uint32_t MAX_K = 2048,
+          bool UsePredicatedShortLoads = false>
 __global__ void __launch_bounds__(FILTERED_TOPK_BLOCK_THREADS)
     FilteredTopKUnifiedKernel(const DType* __restrict__ input,
                               IdType* __restrict__ output,
@@ -1042,6 +1046,19 @@ __global__ void __launch_bounds__(FILTERED_TOPK_BLOCK_THREADS)
     return;
   }
 
+  // Short path
+  if (length <= 32768) {
+    extern __shared__ uint8_t _smem_reg[];
+    if constexpr (UsePredicatedShortLoads) {
+      hist4096::histogram_4096_topk_predicated<MAX_K, 12, 8>(score, dst, length,
+                                                             _smem_reg);
+    } else {
+      hist4096::histogram_4096_topk<MAX_K, 12, 8>(score, dst, length,
+                                                  _smem_reg);
+    }
+    return;
+  }
+
   // Static shared memory
   alignas(128) __shared__ int s_histogram_buf[2][RADIX + 128];
   alignas(128) __shared__ int s_counter;
@@ -1285,14 +1302,15 @@ cudaError_t FilteredTopKRaggedTransform(const DType* input,
 
   const int vec_size = ComputeFilteredTopKVecSize<DType>(max_len);
 
-#define DISPATCH_VEC_SIZE(VS)                                               \
-  if (vec_size == VS) {                                                     \
-    auto kernel = FilteredTopKUnifiedKernel<DType, IdType, VS, MAX_K>;      \
-    FLASHINFER_CUDA_CALL(cudaFuncSetAttribute(                              \
-        kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size));   \
-    FLASHINFER_CUDA_CALL(cudaLaunchKernel((void*)kernel, grid, block, args, \
-                                          smem_size, stream));              \
-    return cudaSuccess;                                                     \
+#define DISPATCH_VEC_SIZE(VS)                                                 \
+  if (vec_size == VS) {                                                       \
+    auto kernel =                                                             \
+        FilteredTopKUnifiedKernel<DType, IdType, VS, MAX_K, (VS != MAX_VEC)>; \
+    FLASHINFER_CUDA_CALL(cudaFuncSetAttribute(                                \
+        kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size));     \
+    FLASHINFER_CUDA_CALL(cudaLaunchKernel((void*)kernel, grid, block, args,   \
+                                          smem_size, stream));                \
+    return cudaSuccess;                                                       \
   }
 
   DISPATCH_VEC_SIZE(1)
@@ -1306,6 +1324,19 @@ cudaError_t FilteredTopKRaggedTransform(const DType* input,
   return cudaSuccess;
 }
 
+}  // namespace filtered_topk
+
+template <typename DType, typename IdType, uint32_t MAX_K = 2048>
+cudaError_t FilteredTopKRaggedTransform(const DType* input,
+                                        IdType* output_indices,
+                                        const IdType* lengths,
+                                        uint32_t num_rows, uint32_t top_k_val,
+                                        uint32_t max_len,
+                                        cudaStream_t stream = 0) {
+  return filtered_topk::FilteredTopKRaggedTransform<DType, IdType, MAX_K>(
+      input, output_indices, lengths, num_rows, top_k_val, max_len, stream);
+}
+
 }  // namespace vllm
 
 #endif  // PERSISTENT_TOPK_CUH_

csrc/libtorch_stable/topk_histogram_4096.cuh

@@ -0,0 +1,563 @@
+/*
+ * Shared 4096-bin single-CTA TopK helpers.
+ */
+
+#ifndef TOPK_HISTOGRAM_4096_CUH_
+#define TOPK_HISTOGRAM_4096_CUH_
+
+#include <cuda_fp16.h>
+#include <cuda_runtime.h>
+#include <cstdint>
+
+namespace vllm {
+namespace topk_histogram_4096 {
+
+constexpr uint32_t kBlockSize = 1024;
+constexpr uint32_t RADIX = 256;
+constexpr uint32_t kMaxTies = 1024;
+static_assert(kMaxTies <= kBlockSize,
+              "tie_handle requires kMaxTies <= kBlockSize");
+constexpr uint32_t kWarpSize = 32;
+constexpr uint32_t kNumWarps = kBlockSize / kWarpSize;
+
+// Register path
+constexpr uint32_t kHist4096VecsPerThread = 4;
+constexpr uint32_t kHist4096MaxLen =
+    kHist4096VecsPerThread * 4 * kBlockSize;  // 16384
+
+struct alignas(16) MatchBin {
+  uint32_t bin, above_count, equal_count;
+};
+struct alignas(8) Tie {
+  uint32_t idx;
+  float score;
+};
+
+__device__ __forceinline__ void load_float4_predicated(const float* ptr,
+                                                       int base, int seq_len,
+                                                       float& v0, float& v1,
+                                                       float& v2, float& v3) {
+  uint32_t r0, r1, r2, r3;
+  const int p0 = (base < seq_len);
+  const int p1 = (base + 1 < seq_len);
+  const int p2 = (base + 2 < seq_len);
+  const int p3 = (base + 3 < seq_len);
+  asm volatile(
+      "{\n"
+      "  .reg .pred pr0, pr1, pr2, pr3;\n"
+      "  setp.ne.u32 pr0, %4, 0;\n"
+      "  setp.ne.u32 pr1, %5, 0;\n"
+      "  setp.ne.u32 pr2, %6, 0;\n"
+      "  setp.ne.u32 pr3, %7, 0;\n"
+      "  mov.u32 %0, 0xFF800000;\n"
+      "  mov.u32 %1, 0xFF800000;\n"
+      "  mov.u32 %2, 0xFF800000;\n"
+      "  mov.u32 %3, 0xFF800000;\n"
+      "  @pr0 ld.global.cg.u32 %0, [%8];\n"
+      "  @pr1 ld.global.cg.u32 %1, [%8+4];\n"
+      "  @pr2 ld.global.cg.u32 %2, [%8+8];\n"
+      "  @pr3 ld.global.cg.u32 %3, [%8+12];\n"
+      "}\n"
+      : "=r"(r0), "=r"(r1), "=r"(r2), "=r"(r3)
+      : "r"(p0), "r"(p1), "r"(p2), "r"(p3), "l"(ptr));
+  v0 = __uint_as_float(r0);
+  v1 = __uint_as_float(r1);
+  v2 = __uint_as_float(r2);
+  v3 = __uint_as_float(r3);
+}
+
+// converts the float32 score to a 32-bit ordered unsigned integer — the full
+// precision key for radix sorting
+__device__ __forceinline__ auto convert_to_uint32_v2(float x) -> uint32_t {
+  uint32_t bits = __float_as_uint(x);
+  return (bits & 0x80000000u) ? ~bits : (bits | 0x80000000u);
+}
+
+// Converts each score to a 12-bit bin (FP16 sign-magnitude -> top 12 bits ->
+// bin 0-4095)
+template <uint32_t kBits>
+__device__ __forceinline__ uint32_t extract_coarse_bin_N(float x) {
+  __half h = __float2half_rn(x);
+  uint16_t bits = __half_as_ushort(h);
+  uint16_t key = (bits & 0x8000) ? static_cast<uint16_t>(~bits)
+                                 : static_cast<uint16_t>(bits | 0x8000);
+  return key >> (16 - kBits);
+}
+
+// running sum within each warp — thread 0 gets its own value, thread 1 gets
+// thread 0 + thread 1, thread 2 gets threads 0+1+2, etc.
+__device__ __forceinline__ uint32_t warp_inclusive_sum(uint32_t lane,
+                                                       uint32_t v) {
+#pragma unroll
+  for (uint32_t o = 1; o < 32; o *= 2) {
+    uint32_t n = __shfl_up_sync(0xFFFFFFFF, v, o);
+    if (lane >= o) v += n;
+  }
+  return v;
+}
+
+// Returns the sum of a value across all 32 threads in the warp, and every
+// thread gets the same result. SM80+ uses redux.sync.add.u32, a single PTX
+// instruction for hardware warp-wide reduction. Older targets use the
+// __shfl_xor_sync butterfly tree, like warp::reduce_sum() (5 shuffles for 32
+// lanes).
+__device__ __forceinline__ uint32_t warp_reduce_sum_full(uint32_t v) {
+#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 800)
+  uint32_t r;
+  asm("redux.sync.add.u32 %0, %1, 0xFFFFFFFF;" : "=r"(r) : "r"(v));
+  return r;
+#else
+  #pragma unroll
+  for (uint32_t mask = kWarpSize >> 1; mask > 0; mask >>= 1) {
+    v += __shfl_xor_sync(0xFFFFFFFF, v, mask);
+  }
+  return v;
+#endif
+}
+
+// ============================================================================
+// Tie refinement (single CTA): 4-round radix-256 topK on the full FP32 ordered
+// key Each round narrows by 8 bits until ties are fully resolved
+// ============================================================================
+
+template <uint32_t TopK>
+__device__ void tie_handle(const Tie* ties, uint32_t num_ties,
+                           uint32_t num_above, int32_t* output, void* _smem) {
+  struct TS {
+    alignas(128) uint32_t counter;
+    alignas(128) MatchBin match;
+    uint32_t histogram[RADIX];
+    uint32_t warp_sum[kNumWarps];
+  };
+  auto* s = static_cast<TS*>(_smem);
+  const auto tx = threadIdx.x;
+  const auto li = tx % kWarpSize, wi = tx / kWarpSize;
+
+  // Each thread loads one tie element.
+  const bool has = tx < num_ties;
+  const auto tie = has ? ties[tx] : Tie{0, 0.0f};
+  const uint32_t key = convert_to_uint32_v2(tie.score);
+
+  bool active = has;  // tracks whether this thread's tie is still a candidate.
+  uint32_t remain =
+      TopK - num_above;  // decreases each round as ties are resolved.
+  uint32_t wpos = TopK;  // wpos will hold the final output position.
+  s->counter = 0;
+  __syncthreads();
+
+  // The 4-round radix loop - each round narrows by 8 bits until ties are fully
+  // resolved
+#pragma unroll
+  for (int r = 0; r < 4; r++) {
+    uint32_t sh = 24 - r * 8;  // round 0: bits 31-24, round 1: 23-16, etc.
+    uint32_t bin = (key >> sh) & 0xFF;  // this tie's 8-bit bin for this round
+
+    // Step 1: Build 256-bin histogram.
+    if (tx < RADIX) s->histogram[tx] = 0;
+    __syncthreads();
+    if (active) atomicAdd(&s->histogram[bin], 1);
+    __syncthreads();
+
+    // Step 2: Prefix scan to find threshold
+    uint32_t hv = 0, wi2 = 0;
+    if (tx < RADIX) {
+      hv = s->histogram[tx];
+      wi2 = warp_inclusive_sum(li, hv);
+      if (li == kWarpSize - 1) s->warp_sum[wi] = wi2;
+    }
+    __syncthreads();
+
+    if (tx < RADIX) {
+      auto tmp = (li < RADIX / kWarpSize) ? s->warp_sum[li] : 0;
+      auto tot = warp_reduce_sum_full(tmp);
+      auto inter = warp_reduce_sum_full(li < wi ? tmp : 0);
+      auto above = tot - (inter + wi2);
+      if (above < remain && above + hv >= remain) {
+        s->match = {tx, above, remain - above};
+      }
+    }
+    __syncthreads();
+
+    // Step 3: Scatter
+    auto [thr, na, _] = s->match;  // threshold bin, num above, unused
+    if (active) {
+      if (bin > thr) {
+        wpos = num_above +
+               atomicAdd(&s->counter, 1);  // above -> place in output directly
+        active = false;
+      } else if (bin < thr)
+        active = false;  // below -> discard
+      else if (r == 3)
+        wpos = TopK - atomicAdd(&s->match.equal_count,
+                                -1u);  // last round: place remaining
+    }
+    remain -= na;
+    if (!remain) break;  // all ties resolved early
+  }
+  // Final write
+  if (wpos < TopK) output[wpos] = tie.idx;
+}
+
+// Extended tie_handle for TopK > kBlockSize (e.g. TopK=2048).
+// tie_handle assumes 1 tie per thread (max 1024).
+// This version handles 2 ties per thread via kPerThread=2
+template <uint32_t TopK>
+__device__ void tie_handle_large(const Tie* ties, uint32_t num_ties,
+                                 uint32_t num_above, int32_t* output,
+                                 void* _smem) {
+  static_assert(TopK > kBlockSize);
+  struct TS {
+    alignas(128) uint32_t counter;
+    alignas(128) MatchBin match;
+    uint32_t histogram[RADIX];
+    uint32_t warp_sum[kNumWarps];
+  };
+  auto* s = static_cast<TS*>(_smem);
+  const auto tx = threadIdx.x;
+  const auto li = tx % kWarpSize;
+  const auto wi = tx / kWarpSize;
+
+  constexpr uint32_t kPerThread = (TopK + kBlockSize - 1) / kBlockSize;
+  Tie my_ties[kPerThread];
+  uint32_t keys[kPerThread];
+  bool active[kPerThread];
+
+  for (uint32_t e = 0; e < kPerThread; e++) {
+    uint32_t idx = e * kBlockSize + tx;
+    if (idx < num_ties) {
+      my_ties[e] = ties[idx];
+      keys[e] = convert_to_uint32_v2(ties[idx].score);
+      active[e] = true;
+    } else {
+      my_ties[e] = {0, 0.0f};
+      keys[e] = 0;
+      active[e] = false;
+    }
+  }
+
+  uint32_t remain = TopK - num_above;
+  s->counter = 0;
+  __syncthreads();
+
+  for (int r = 0; r < 4; r++) {
+    uint32_t sh = 24 - r * 8;
+    if (tx < RADIX) {
+      s->histogram[tx] = 0;
+    }
+    __syncthreads();
+
+    for (uint32_t e = 0; e < kPerThread; e++) {
+      if (active[e]) {
+        atomicAdd(&s->histogram[(keys[e] >> sh) & 0xFF], 1);
+      }
+    }
+    __syncthreads();
+
+    uint32_t hv = 0;
+    if (tx < RADIX) {
+      hv = s->histogram[tx];
+      auto wi2 = warp_inclusive_sum(li, hv);
+      if (li == kWarpSize - 1) {
+        s->warp_sum[wi] = wi2;
+      }
+    }
+    __syncthreads();
+    if (tx < RADIX) {
+      auto tmp2 = (li < RADIX / kWarpSize) ? s->warp_sum[li] : 0;
+      auto total = warp_reduce_sum_full(tmp2);
+      auto inter = warp_reduce_sum_full(li < wi ? tmp2 : 0);
+      auto wi2 = warp_inclusive_sum(li, hv);
+      auto above = total - (inter + wi2);
+      if (above < remain && above + hv >= remain) {
+        s->match = {
+            .bin = tx, .above_count = above, .equal_count = remain - above};
+      }
+    }
+    __syncthreads();
+
+    auto thr = s->match.bin;
+    auto na = s->match.above_count;
+
+    for (uint32_t e = 0; e < kPerThread; e++) {
+      if (!active[e]) {
+        continue;
+      }
+      uint32_t bin = (keys[e] >> sh) & 0xFF;
+      if (bin > thr) {
+        uint32_t wpos = num_above + atomicAdd(&s->counter, 1);
+        if (wpos < TopK) {
+          output[wpos] = my_ties[e].idx;
+        }
+        active[e] = false;
+      } else if (bin < thr) {
+        active[e] = false;
+      } else if (r == 3) {
+        uint32_t wpos = TopK - atomicAdd(&s->match.equal_count, -1u);
+        if (wpos < TopK) {
+          output[wpos] = my_ties[e].idx;
+        }
+      }
+    }
+
+    num_above += na;
+    remain -= na;
+    __syncthreads();
+    s->counter = 0;
+    __syncthreads();
+  }
+}
+
+// ============================================================================
+// Register-based single-CTA fast path for seq_len <= 16384
+// 4 float4 per thread × 1024 threads = 16384 elements max
+// Uses 4096-bin (12-bit) histogram for better precision
+// ============================================================================
+
+template <uint32_t TopK, uint32_t HIST_BITS>
+struct Histogram4096Smem {
+  static constexpr uint32_t HIST_BINS = 1 << HIST_BITS;
+  static constexpr uint32_t TIE_CAPACITY = TopK > kMaxTies ? TopK : kMaxTies;
+  alignas(128) uint32_t counter_gt;
+  alignas(128) uint32_t counter_eq;
+  MatchBin match;
+  uint32_t warp_sum[kNumWarps];
+  union {
+    uint32_t histogram[HIST_BINS];
+    Tie tie_buffer[TIE_CAPACITY];
+  };
+};
+
+template <uint32_t TopK, uint32_t HIST_BITS,
+          uint32_t VECS_PER_THREAD = kHist4096VecsPerThread,
+          bool UsePredicatedLoads = false>
+__device__ void histogram_4096_topk(const float* __restrict__ scores,
+                                    int32_t* __restrict__ output,
+                                    uint32_t length, void* _smem) {
+  constexpr uint32_t HIST_BINS = 1 << HIST_BITS;
+  constexpr uint32_t ITEMS_PER_THREAD = HIST_BINS / kBlockSize;
+  static_assert(HIST_BINS >= kBlockSize,
+                "HIST_BITS must give >= kBlockSize bins");
+
+  using Smem = Histogram4096Smem<TopK, HIST_BITS>;
+  auto* smem = static_cast<Smem*>(_smem);
+  const auto tx = threadIdx.x;
+  const auto lane_id = tx % kWarpSize;
+  const auto warp_id = tx / kWarpSize;
+
+  // Phase 1: Load all data into RF + build histogram
+  float4
+      vecs[VECS_PER_THREAD];  // 4 vectors x 4 floats = 16 elements per thread
+  if constexpr (ITEMS_PER_THREAD >= 4) {
+    // Zero the histogram (SMEM writes)
+    for (uint32_t i = 0; i < ITEMS_PER_THREAD / 4; i++)
+      reinterpret_cast<uint4*>(
+          smem->histogram)[tx * (ITEMS_PER_THREAD / 4) + i] =
+          make_uint4(0, 0, 0, 0);
+  } else {
+    if (tx < HIST_BINS) smem->histogram[tx] = 0;
+  }
+  if (tx == 0) {
+    smem->counter_gt = 0;
+    smem->counter_eq = 0;
+  }
+  if constexpr (UsePredicatedLoads) {
+    const bool row_aligned = (reinterpret_cast<uintptr_t>(scores) & 0xFu) == 0;
+#pragma unroll
+    for (uint32_t v = 0; v < VECS_PER_THREAD; v++) {
+      const uint32_t base = (tx + v * kBlockSize) * 4;
+      if (base < length) {
+        if (row_aligned && base + 3 < length) {
+          vecs[v] = *reinterpret_cast<const float4*>(scores + base);
+        } else {
+          load_float4_predicated(scores + base, static_cast<int>(base),
+                                 static_cast<int>(length), vecs[v].x, vecs[v].y,
+                                 vecs[v].z, vecs[v].w);
+        }
+      }
+    }
+  } else {
+#pragma unroll
+    for (uint32_t v = 0; v < VECS_PER_THREAD; v++) {
+      const uint32_t base = (tx + v * kBlockSize) * 4;
+      if (base < length) {
+        vecs[v] = *reinterpret_cast<const float4*>(scores + base);
+      }
+    }
+  }
+  __syncthreads();
+
+  // Build histogram from RF via atomic adds into the shared histogram
+  bool done = false;
+#pragma unroll
+  for (uint32_t v = 0; v < VECS_PER_THREAD && !done; v++) {
+    const float* elems = reinterpret_cast<const float*>(&vecs[v]);
+#pragma unroll
+    for (uint32_t e = 0; e < 4 && !done; e++) {
+      const uint32_t idx = (tx + v * kBlockSize) * 4 + e;
+      if (idx >= length) {
+        done = true;
+      } else {
+        atomicAdd(&smem->histogram[extract_coarse_bin_N<HIST_BITS>(elems[e])],
+                  1);
+      }
+    }
+  }
+  __syncthreads();
+
+  // Phase 2: Prefix scan to find threshold bin
+  // Multi-element scan (4096 bins: 4 per thread)
+  uint32_t orig[ITEMS_PER_THREAD];
+  uint32_t local_sum = 0;
+
+  // Step 1: Each thread sums its 4 bins
+#pragma unroll
+  for (uint32_t i = 0; i < ITEMS_PER_THREAD; i++) {
+    orig[i] = smem->histogram[tx * ITEMS_PER_THREAD + i];
+    local_sum += orig[i];
+  }
+
+  // Step 2: Warp-level inclusive prefix sum on local_sum
+  const auto warp_inc = warp_inclusive_sum(lane_id, local_sum);
+  if (lane_id == kWarpSize - 1) smem->warp_sum[warp_id] = warp_inc;
+  __syncthreads();
+
+  // Step 3: Inter-warp prefix across warp sums.
+  const auto tmp = smem->warp_sum[lane_id];
+  uint32_t prefix = warp_reduce_sum_full(
+      lane_id < warp_id ? tmp : 0);  // sum of all prior warps
+  prefix +=
+      warp_inc - local_sum;  // exclusive prefix within this thread's position
+
+  // Step 4: Find threshold - scan 4 bins, accumulate prefix
+#pragma unroll
+  for (uint32_t i = 0; i < ITEMS_PER_THREAD; i++) {
+    prefix += orig[i];
+    const auto above = length - prefix;  // elements in bins ABOVE this one
+    if (above < TopK && above + orig[i] >= TopK) {
+      smem->match = {.bin = tx * ITEMS_PER_THREAD + i,
+                     .above_count = above,
+                     .equal_count = orig[i]};
+    }
+  }
+
+  __syncthreads();
+
+  // Phase 3: Scatter from registers
+  const auto [thr_bin, num_above, num_equal] = smem->match;
+  const bool need_tie = (num_equal + num_above > TopK);
+
+  done = false;
+#pragma unroll
+  for (uint32_t v = 0; v < VECS_PER_THREAD && !done; v++) {
+    const float* elems = reinterpret_cast<const float*>(&vecs[v]);
+#pragma unroll
+    for (uint32_t e = 0; e < 4 && !done; e++) {
+      const uint32_t idx = (tx + v * kBlockSize) * 4 + e;
+      if (idx >= length) {
+        done = true;
+      } else {
+        const uint32_t bin = extract_coarse_bin_N<HIST_BITS>(elems[e]);
+        if (bin > thr_bin) {
+          output[atomicAdd(&smem->counter_gt, 1)] =
+              idx;  // above -> output directly
+        } else if (bin == thr_bin) {
+          const auto pos = atomicAdd(&smem->counter_eq, 1);
+          if (!need_tie) {
+            if (pos + num_above < TopK) {
+              output[pos + num_above] = idx;  // all fit
+            }
+          } else {
+            if (pos < TopK) {
+              smem->tie_buffer[pos] = {idx, elems[e]};  // store for refirement
+            }
+          }
+        }
+        // else: bin < thr_bin - discard (not in top-k)
+      }
+    }
+  }
+
+  // Phase 4: Tie-breaking
+  if (!need_tie) return;
+  __syncthreads();
+
+  // Fast warp-ballot tie-breaking for small tie counts
+  const uint32_t num_ties = min(num_equal, static_cast<uint32_t>(TopK));
+  const uint32_t topk_remain =
+      TopK - num_above;  // pick exactly remaining elements to fill topK
+
+  auto is_greater = [](const Tie& a, const Tie& b) {
+    return (a.score > b.score) || (a.score == b.score && a.idx < b.idx);
+  };
+
+  if (num_ties <= kWarpSize) {
+    // <=32 ties - Use warp ballot
+    // All-to-all comparison in one __ballot_sync. 32 ties x 32 warps = 1024
+    // comparisons in one instruction per warp. O(1) work.
+    const auto lane_id = tx % kWarpSize;
+    const auto warp_id = tx / kWarpSize;
+    if (lane_id >= num_ties || warp_id >= num_ties) return;
+    const uint32_t mask = (1ull << num_ties) - 1u;
+    const auto tie = smem->tie_buffer[lane_id];  // each lane holds one tie
+    const auto target =
+        smem->tie_buffer[warp_id];  // each warp evaluates one candidate
+    const bool pred =
+        is_greater(tie, target);  // compare all ties against target
+    const auto rank = static_cast<uint32_t>(
+        __popc(__ballot_sync(mask, pred)));  // count how many are greater
+    if (lane_id == 0 && rank < topk_remain) {
+      output[num_above + rank] = target.idx;  // place at correct position
+    }
+  } else if (num_ties <=
+             kWarpSize *
+                 2) {  // TODO (roberto): try to refactor this with <=32 case
+    //  Same idea but each thread handles 2 tie elements
+    const auto lane_id = tx % kWarpSize;
+    const auto warp_id = tx / kWarpSize;
+    const auto lane1 = lane_id + kWarpSize;
+    const auto warp1 = warp_id + kWarpSize;
+    const auto invalid = Tie{0xFFFFFFFF, -__FLT_MAX__};
+    const auto tie0 = smem->tie_buffer[lane_id];
+    const auto tie1 = lane1 < num_ties ? smem->tie_buffer[lane1] : invalid;
+    if (warp_id < num_ties) {
+      const auto target = smem->tie_buffer[warp_id];
+      const auto r0 =
+          __popc(__ballot_sync(0xFFFFFFFF, is_greater(tie0, target)));
+      const auto r1 =
+          __popc(__ballot_sync(0xFFFFFFFF, is_greater(tie1, target)));
+      if (lane_id == 0 && r0 + r1 < topk_remain)
+        output[num_above + r0 + r1] = target.idx;
+    }
+    if (warp1 < num_ties) {
+      const auto target = smem->tie_buffer[warp1];
+      const auto r0 =
+          __popc(__ballot_sync(0xFFFFFFFF, is_greater(tie0, target)));
+      const auto r1 =
+          __popc(__ballot_sync(0xFFFFFFFF, is_greater(tie1, target)));
+      if (lane_id == 0 && r0 + r1 < topk_remain)
+        output[num_above + r0 + r1] = target.idx;
+    }
+  } else {
+    // Large tie count: fall back to 4-round radix-256 sort
+    if constexpr (TopK <= kBlockSize) {
+      tie_handle<TopK>(smem->tie_buffer, num_ties, num_above, output, smem);
+    } else {
+      tie_handle_large<TopK>(smem->tie_buffer, num_ties, num_above, output,
+                             smem);
+    }
+  }
+}
+
+template <uint32_t TopK, uint32_t HIST_BITS,
+          uint32_t VECS_PER_THREAD = kHist4096VecsPerThread>
+__device__ __noinline__ void histogram_4096_topk_predicated(
+    const float* __restrict__ scores, int32_t* __restrict__ output,
+    uint32_t length, void* _smem) {
+  histogram_4096_topk<TopK, HIST_BITS, VECS_PER_THREAD, true>(scores, output,
+                                                              length, _smem);
+}
+
+}  // namespace topk_histogram_4096
+}  // namespace vllm
+
+#endif  // TOPK_HISTOGRAM_4096_CUH_

csrc/libtorch_stable/torch_bindings.cpp

@@ -493,6 +493,12 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
       "persistent_topk(Tensor logits, Tensor lengths, Tensor! output, "
       "Tensor workspace, int k, int max_seq_len) -> ()");
 
+#ifdef VLLM_ENABLE_COOPERATIVE_TOPK
+  ops.def(
+      "cooperative_topk(Tensor logits, Tensor lengths, Tensor! output, "
+      "Tensor workspace, int k, int max_seq_len) -> ()");
+#endif
+
   // Activation ops
   ops.def(
       "persistent_masked_m_silu_mul_quant(Tensor input, Tensor counts, Tensor! "
@@ -711,6 +717,9 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) {
   ops.impl("top_k_per_row_prefill", TORCH_BOX(&top_k_per_row_prefill));
   ops.impl("top_k_per_row_decode", TORCH_BOX(&top_k_per_row_decode));
   ops.impl("persistent_topk", TORCH_BOX(&persistent_topk));
+#ifdef VLLM_ENABLE_COOPERATIVE_TOPK
+  ops.impl("cooperative_topk", TORCH_BOX(&cooperative_topk));
+#endif
 
   // Activation kernels (shared CUDA/ROCm)
   ops.impl("persistent_masked_m_silu_mul_quant",

docs/models/supported_models.md

@@ -597,6 +597,7 @@ These models primarily accept the [`LLM.generate`](./generative_models.md#llmgen
 | `Mistral3ForConditionalGeneration` | Mistral3 (HF Transformers) | T + I<sup>+</sup> | `mistralai/Mistral-Small-3.1-24B-Instruct-2503`, etc. | ✅︎ | ✅︎ |
 | `MolmoForCausalLM` | Molmo | T + I<sup>+</sup> | `allenai/Molmo-7B-D-0924`, `allenai/Molmo-7B-O-0924`, etc. | ✅︎ | ✅︎ |
 | `Molmo2ForConditionalGeneration` | Molmo2 | T + I<sup>+</sup> / V | `allenai/Molmo2-4B`, `allenai/Molmo2-8B`, `allenai/Molmo2-O-7B`, `allenai/MolmoWeb-4B`<sup>^</sup>, `allenai/MolmoWeb-8B`<sup>^</sup> | ✅︎ | ✅︎ |
+| `MossAudioModel` | MOSS-Audio | T + A<sup>+</sup> | `OpenMOSS-Team/MOSS-Audio-4B-Instruct`, `OpenMOSS-Team/MOSS-Audio-4B-Thinking`, `OpenMOSS-Team/MOSS-Audio-8B-Instruct`, `OpenMOSS-Team/MOSS-Audio-8B-Thinking` | ✅︎ | ✅︎ |
 | `Moondream3ForCausalLM` | Moondream3 | T + I | `moondream/moondream3-preview` | | ✅︎ |
 | `MusicFlamingoForConditionalGeneration` | MusicFlamingo | T + A | `nvidia/music-flamingo-2601-hf`, `nvidia/music-flamingo-think-2601-hf` | ✅︎ | ✅︎ |
 | `NVLM_D_Model` | NVLM-D 1.0 | T + I<sup>+</sup> | `nvidia/NVLM-D-72B`, etc. | | ✅︎ |

setup.py

@@ -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",
     ]
 }
 

tests/benchmarks/test_throughput_cli.py

@@ -4,6 +4,13 @@ import subprocess
 
 import pytest
 
+from vllm.benchmarks.datasets import SampleRequest
+from vllm.benchmarks.throughput import (
+    _run_vllm_chat_requests,
+    add_cli_args,
+)
+from vllm.utils.argparse_utils import FlexibleArgumentParser
+
 MODEL_NAME = "meta-llama/Llama-3.2-1B-Instruct"
 
 
@@ -28,3 +35,70 @@ def test_bench_throughput():
     print(result.stderr)
 
     assert result.returncode == 0, f"Benchmark failed: {result.stderr}"
+
+
+def test_bench_throughput_accepts_custom_audio_args():
+    parser = FlexibleArgumentParser()
+    add_cli_args(parser)
+
+    args = parser.parse_args(
+        [
+            "--dataset-name",
+            "custom_audio",
+            "--dataset-path",
+            "audio.jsonl",
+            "--no-oversample",
+            "--custom-output-len",
+            "32",
+            "--enable-multimodal-chat",
+        ]
+    )
+
+    assert args.dataset_name == "custom_audio"
+    assert args.no_oversample
+    assert args.custom_output_len == 32
+    assert args.enable_multimodal_chat
+
+
+def test_vllm_chat_requests_include_multimodal_content():
+    class FakeLLM:
+        def __init__(self):
+            self.prompts = None
+
+        def chat(self, prompts, sampling_params, use_tqdm):
+            del sampling_params, use_tqdm
+            self.prompts = prompts
+            return []
+
+    llm = FakeLLM()
+    audio_content = {
+        "type": "input_audio",
+        "input_audio": {"data": "abc", "format": "wav"},
+    }
+    request = SampleRequest(
+        prompt="Transcribe this audio.",
+        prompt_len=1,
+        expected_output_len=8,
+        multi_modal_data=audio_content,
+    )
+
+    _run_vllm_chat_requests(
+        llm,
+        [request],
+        n=1,
+        disable_detokenize=False,
+        do_profile=False,
+        prequeue_requests=False,
+    )
+
+    assert llm.prompts == [
+        [
+            {
+                "role": "user",
+                "content": [
+                    {"type": "text", "text": "Transcribe this audio."},
+                    audio_content,
+                ],
+            }
+        ]
+    ]

tests/entrypoints/openai/chat_completion/test_chat.py

@@ -824,7 +824,10 @@ async def test_invocations(server: RemoteOpenAIServer, client: openai.AsyncOpenA
     chat_output = chat_response.json()
     invocation_output = invocation_response.json()
 
-    assert chat_output.keys() == invocation_output.keys()
+    extra_keys = invocation_output.keys() - chat_output.keys()
+    missing_keys = chat_output.keys() - invocation_output.keys()
+    assert missing_keys == set()
+    assert extra_keys <= {"moderation"}
     assert chat_output["choices"] == invocation_output["choices"]
 
 

tests/kernels/attention/test_minimax_m3.py

@@ -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,

tests/kernels/test_fused_minimax_m3_qknorm_rope_kv_insert.py

@@ -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)

tests/kernels/test_top_k_per_row.py

@@ -14,6 +14,70 @@ TOP_K_VALUES = [2048, 3000]
 BATCH_SIZE = [1, 2, 2048]
 NEXT_N = [1, 8]
 DATA_GENERATION = ["random", "10LSBits"]
+RADIX_TOPK_WORKSPACE_SIZE = 1024 * 1024
+
+
+def _has_device_capability(major: int) -> bool:
+    return current_platform.is_cuda() and current_platform.has_device_capability(major)
+
+
+COOPERATIVE_TOPK_BACKEND = pytest.param(
+    "cooperative_topk",
+    marks=pytest.mark.skipif(
+        not _has_device_capability(90),
+        reason="cooperative_topk requires SM90+",
+    ),
+)
+WORKSPACE_TOPK_BACKENDS = ["persistent_topk", COOPERATIVE_TOPK_BACKEND]
+TOPK_BACKENDS = ["top_k_per_row_decode", *WORKSPACE_TOPK_BACKENDS]
+
+
+def _run_topk_backend(
+    backend: str,
+    logits: torch.Tensor,
+    lengths: torch.Tensor,
+    indices: torch.Tensor,
+    top_k: int,
+    max_seq_len: int,
+    next_n: int = 1,
+) -> None:
+    if backend == "top_k_per_row_decode":
+        torch.ops._C.top_k_per_row_decode(
+            logits,
+            next_n,
+            lengths,
+            indices,
+            indices.shape[0],
+            logits.stride(0),
+            logits.stride(1),
+            top_k,
+        )
+    elif backend == "persistent_topk":
+        workspace = torch.empty(
+            RADIX_TOPK_WORKSPACE_SIZE, dtype=torch.uint8, device="cuda"
+        )
+        torch.ops._C.persistent_topk(
+            logits, lengths, indices, workspace, top_k, max_seq_len
+        )
+    elif backend == "cooperative_topk":
+        if indices.shape[0] > 32:
+            pytest.skip(
+                "cooperative_topk supports <=32 rows; "
+                "persistent_topk covers larger batches"
+            )
+        if logits.stride(0) % 4 != 0:
+            pytest.skip(
+                "cooperative_topk requires row stride divisible by 4; "
+                "persistent_topk covers unaligned strides"
+            )
+        workspace = torch.empty(
+            RADIX_TOPK_WORKSPACE_SIZE, dtype=torch.uint8, device="cuda"
+        )
+        torch.ops._C.cooperative_topk(
+            logits, lengths, indices, workspace, top_k, max_seq_len
+        )
+    else:
+        raise ValueError(f"Unknown top-k backend: {backend}")
 
 
 def create_random_logits(
@@ -322,16 +386,19 @@ def test_top_k_per_row_decode_large_vocab_size(clean_logits: bool) -> None:
 @pytest.mark.parametrize("clean_logits", [True, False])
 @pytest.mark.parametrize("top_k", [2048])
 @pytest.mark.parametrize("next_n", [1, 4])
+@pytest.mark.parametrize("backend", WORKSPACE_TOPK_BACKENDS)
 @torch.inference_mode()
-def test_deepseek_persistent_topk(
+def test_deepseek_workspace_topk(
     seq_len_range: tuple[int, int],
     test_id: str,
     clean_logits: bool,
     top_k: int,
     next_n: int,
+    backend: str,
 ) -> None:
     """
-    Test persistent_topk with varying sequence lengths and speculative decoding.
+    Test workspace top-k backends with varying sequence lengths and speculative
+    decoding.
     Supports speculative decoding with next_n > 1.
     """
     set_random_seed(42 if test_id == "short_sequences" else 43)
@@ -347,6 +414,7 @@ def test_deepseek_persistent_topk(
         dtype=torch.int32,
         device="cuda",
     )
+    seq_lens = (seq_lens + 3) & ~3  # align to 4 for TMA
 
     # Compute row boundaries for speculative decoding
     row_starts = torch.zeros(num_rows, dtype=torch.int32, device="cuda")
@@ -366,14 +434,11 @@ def test_deepseek_persistent_topk(
         offsets = torch.arange(next_n, device=logits.device, dtype=torch.int32)
         lengths = (seq_lens.unsqueeze(1) - next_n + 1 + offsets).flatten()
 
-    workspace = torch.empty(1024 * 1024, dtype=torch.uint8, device="cuda")
     max_seq_len = int(seq_lens.max().item())
-    torch.ops._C.persistent_topk(
-        logits, lengths, indices, workspace, top_k, max_seq_len
-    )
+    _run_topk_backend(backend, logits, lengths, indices, top_k, max_seq_len, next_n)
 
     validate_topk_against_reference(
-        logits, indices, row_starts, row_ends, top_k, f"persistent_topk ({test_id})"
+        logits, indices, row_starts, row_ends, top_k, f"{backend} ({test_id})"
     )
 
 
@@ -383,9 +448,10 @@ def run_large_context_topk_test(
     top_k: int,
     data_type: str = "random",
     seed: int = 42,
+    backend: str = "cooperative_topk",
 ) -> None:
     """
-    Helper to run persistent_topk kernel test with given parameters.
+    Helper to run a top-k backend test with given parameters.
 
     Args:
         batch_size: Number of rows/sequences
@@ -393,6 +459,7 @@ def run_large_context_topk_test(
         top_k: Number of top elements to select
         data_type: Type of test data to generate
         seed: Random seed for reproducibility
+        backend: Top-k backend to test
     """
     torch.set_default_device("cuda:0")
     set_random_seed(seed)
@@ -449,11 +516,8 @@ def run_large_context_topk_test(
     # Create output tensor
     indices = torch.empty((num_rows, top_k), dtype=torch.int32, device="cuda")
 
-    workspace = torch.empty(1024 * 1024, dtype=torch.uint8, device="cuda")
     max_seq_len = max(seq_lens)
-    torch.ops._C.persistent_topk(
-        logits, lengths, indices, workspace, top_k, max_seq_len
-    )
+    _run_topk_backend(backend, logits, lengths, indices, top_k, max_seq_len)
 
     torch.accelerator.synchronize()
 
@@ -605,8 +669,9 @@ def run_large_context_topk_test(
         ),
     ],
 )
+@pytest.mark.parametrize("backend", WORKSPACE_TOPK_BACKENDS)
 @torch.inference_mode()
-def test_persistent_topk_correctness(test_config: dict) -> None:
+def test_workspace_topk_correctness(test_config: dict, backend: str) -> None:
     """
     Comprehensive correctness tests covering:
     - Sequence length edge cases (trivial, boundary, varied)
@@ -620,6 +685,7 @@ def test_persistent_topk_correctness(test_config: dict) -> None:
         seq_lens=test_config["seq_lens"],
         top_k=test_config["top_k"],
         data_type=test_config.get("data_type", "random"),
+        backend=backend,
     )
 
 
@@ -668,8 +734,9 @@ def test_persistent_topk_correctness(test_config: dict) -> None:
         ),
     ],
 )
+@pytest.mark.parametrize("backend", WORKSPACE_TOPK_BACKENDS)
 @torch.inference_mode()
-def test_persistent_topk_algorithm_paths(test_config: dict) -> None:
+def test_workspace_topk_algorithm_paths(test_config: dict, backend: str) -> None:
     """
     Test different algorithm execution paths (capped at 163840 for DeepSeek V3.2):
     - Batch size scalability (1, 4, 32, 256)
@@ -680,12 +747,14 @@ def test_persistent_topk_algorithm_paths(test_config: dict) -> None:
         batch_size=test_config["batch_size"],
         seq_lens=[test_config["seq_len"]] * test_config["batch_size"],
         top_k=test_config["top_k"],
+        backend=backend,
     )
 
 
 @pytest.mark.skipif(not current_platform.is_cuda(), reason="This test requires CUDA")
+@pytest.mark.parametrize("backend", WORKSPACE_TOPK_BACKENDS)
 @torch.inference_mode()
-def test_persistent_topk_stress() -> None:
+def test_workspace_topk_stress(backend: str) -> None:
     """
     Stress test with random configurations to catch edge cases.
     Capped at 163840 (DeepSeek V3.2 max context) for realistic testing.
@@ -700,16 +769,73 @@ def test_persistent_topk_stress() -> None:
         batch_size = torch.randint(1, 32, (1,)).item()
 
         # Random sequence lengths capped at DeepSeek V3.2 max context
-        seq_lens = torch.randint(100, 163840, (batch_size,)).tolist()
+        seq_lens_tensor = torch.randint(100, 163840, (batch_size,))
+        if backend == "cooperative_topk":
+            seq_lens = ((seq_lens_tensor + 3) & ~3).tolist()
+        else:
+            seq_lens = seq_lens_tensor.tolist()
 
         run_large_context_topk_test(
             batch_size=batch_size,
             seq_lens=seq_lens,
             top_k=top_k,
             seed=seed,
+            backend=backend,
         )
 
 
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="This test requires CUDA")
+@pytest.mark.parametrize("backend", TOPK_BACKENDS)
+@pytest.mark.parametrize("top_k", [512, 1024, 2048])
+@torch.inference_mode()
+def test_deepseek_topk_backends_no_error_and_reference(
+    backend: str,
+    top_k: int,
+) -> None:
+    """Exercise every production top-k backend on the same inputs."""
+    run_large_context_topk_test(
+        batch_size=4,
+        seq_lens=[2049, 4097, 8191, 12000],
+        top_k=top_k,
+        data_type="random",
+        seed=123,
+        backend=backend,
+    )
+
+
+@pytest.mark.skipif(not _has_device_capability(90), reason="This test requires SM90+")
+@torch.inference_mode()
+def test_cooperative_topk_512_tie_workspace_is_per_row() -> None:
+    """Regression test for TopK=512 tie workspace row overlap."""
+    torch.set_default_device("cuda:0")
+
+    top_k = 512
+    num_rows = 2
+    stride = 65536
+    lengths = torch.tensor([40960, 65536], dtype=torch.int32, device="cuda")
+    logits = torch.full(
+        (num_rows, stride), float("-inf"), dtype=torch.float32, device="cuda"
+    )
+
+    # Row 0 must never select these low indices: many better row-0 ties exist.
+    logits[0, :2048] = -10.0
+    logits[0, 2048 : lengths[0]] = 1.0
+    # Row 1 has higher exact tie scores. With the old row * TopK tie_ws stride,
+    # these row-1 ties could overwrite row 0's TopK=512 refinement workspace.
+    logits[1, : lengths[1]] = 2.0
+
+    indices = torch.empty((num_rows, top_k), dtype=torch.int32, device="cuda")
+    workspace = torch.empty(RADIX_TOPK_WORKSPACE_SIZE, dtype=torch.uint8, device="cuda")
+    torch.ops._C.cooperative_topk(logits, lengths, indices, workspace, top_k, stride)
+    torch.accelerator.synchronize()
+
+    row0 = indices[0].cpu()
+    assert torch.all(row0 >= 2048), (
+        "cooperative_topk TopK=512 selected row-0 low-score indices, likely "
+        "from overlapping tie_ws rows"
+    )
+
+
 @pytest.mark.skipif(not current_platform.is_cuda(), reason="This test requires CUDA")
 @pytest.mark.parametrize(
     "test_config",
@@ -774,10 +900,11 @@ def test_persistent_topk_stress() -> None:
     ],
 )
 @pytest.mark.parametrize("top_k", [512, 2048])
+@pytest.mark.parametrize("backend", WORKSPACE_TOPK_BACKENDS)
 @torch.inference_mode()
-def test_persistent_topk(test_config: dict, top_k: int) -> None:
+def test_workspace_topk(test_config: dict, top_k: int, backend: str) -> None:
     """
-    Tests specific to the persistent_topk kernel:
+    Tests specific to workspace top-k backends:
     - Mixed medium/large rows in the same batch (dynamic per-row dispatch)
     - Boundary around LARGE_THRESHOLD (32K)
     - Trivial + medium + large rows in a single batch
@@ -787,15 +914,17 @@ def test_persistent_topk(test_config: dict, top_k: int) -> None:
         seq_lens=test_config["seq_lens"],
         top_k=top_k,
         data_type=test_config.get("data_type", "random"),
+        backend=backend,
     )
 
 
 @pytest.mark.skipif(not current_platform.is_cuda(), reason="This test requires CUDA")
 @pytest.mark.parametrize("top_k", [512, 2048])
+@pytest.mark.parametrize("backend", WORKSPACE_TOPK_BACKENDS)
 @torch.inference_mode()
-def test_persistent_topk_padded_stride(top_k: int) -> None:
+def test_workspace_topk_padded_stride(top_k: int, backend: str) -> None:
     """
-    Test persistent_topk with padded logits (large stride, small seq_len)
+    Test workspace top-k backends with padded logits (large stride, small seq_len)
     to simulate the e2e CUDAGraph scenario where fp8_paged_mqa_logits
     returns [B, max_model_len] with max_model_len=163840.
     """
@@ -818,11 +947,7 @@ def test_persistent_topk_padded_stride(top_k: int) -> None:
 
     lengths = torch.tensor(actual_seq_lens, dtype=torch.int32, device="cuda")
     indices = torch.empty((batch_size, top_k), dtype=torch.int32, device="cuda")
-    workspace = torch.empty(1024 * 1024, dtype=torch.uint8, device="cuda")
-
-    torch.ops._C.persistent_topk(
-        logits, lengths, indices, workspace, top_k, max(actual_seq_lens)
-    )
+    _run_topk_backend(backend, logits, lengths, indices, top_k, max(actual_seq_lens))
     torch.accelerator.synchronize()
 
     # Validate against torch.topk
@@ -840,6 +965,6 @@ def test_persistent_topk_padded_stride(top_k: int) -> None:
             expected_vals = logits[i, expected].cpu().sort(descending=True)[0]
             actual_vals = logits[i, actual].cpu().sort(descending=True)[0]
             assert torch.allclose(expected_vals, actual_vals, rtol=1e-4, atol=1e-4), (
-                f"Row {i}: persistent_topk with padded stride doesn't match. "
+                f"Row {i}: {backend} with padded stride doesn't match. "
                 f"seq_len={sl}, stride={padded_stride}"
             )

tests/models/multimodal/generation/test_moss_audio.py

@@ -0,0 +1,155 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import pytest
+
+from vllm.assets.audio import AudioAsset
+from vllm.model_executor.models.moss_audio import MOSS_AUDIO_PLACEHOLDER
+from vllm.platforms import current_platform
+
+from ...registry import HF_EXAMPLE_MODELS
+from ...utils import check_logprobs_close
+
+CORE_MODEL = pytest.param(
+    "OpenMOSS-Team/MOSS-Audio-4B-Instruct",
+    marks=pytest.mark.core_model,
+    id="4b-instruct",
+)
+
+EXTENDED_MODELS = [
+    "OpenMOSS-Team/MOSS-Audio-4B-Thinking",
+    "OpenMOSS-Team/MOSS-Audio-8B-Instruct",
+    "OpenMOSS-Team/MOSS-Audio-8B-Thinking",
+]
+
+ACCURACY_MODELS = [CORE_MODEL, *EXTENDED_MODELS]
+
+PARALLEL_SMOKE_CASES = [
+    pytest.param({"tensor_parallel_size": 2}, id="tp2"),
+    pytest.param({"pipeline_parallel_size": 2}, id="pp2"),
+    pytest.param(
+        {"tensor_parallel_size": 2, "pipeline_parallel_size": 2},
+        id="tp2_pp2",
+    ),
+]
+
+HF_ACCURACY_SKIP_REASON = (
+    "HF AutoModelForCausalLM cannot load remote MOSS-Audio configs; "
+    "vLLM generation coverage is provided by the smoke tests below."
+)
+
+
+@pytest.mark.core_model
+def test_moss_audio_generation_smoke(vllm_runner) -> None:
+    model = "OpenMOSS-Team/MOSS-Audio-4B-Instruct"
+    model_info = HF_EXAMPLE_MODELS.find_hf_info(model)
+    model_info.check_available_online(on_fail="skip")
+    model_info.check_transformers_version(on_fail="skip")
+
+    prompts = [f"{MOSS_AUDIO_PLACEHOLDER}\nBriefly describe this audio."]
+    audios = [[AudioAsset("mary_had_lamb").audio_and_sample_rate[0]]]
+
+    with vllm_runner(
+        model,
+        dtype="half",
+        enforce_eager=True,
+        max_model_len=1024,
+        limit_mm_per_prompt={"audio": 1},
+        trust_remote_code=True,
+    ) as vllm_model:
+        outputs = vllm_model.generate_greedy(
+            prompts,
+            max_tokens=4,
+            audios=audios,
+        )
+
+    assert len(outputs) == 1
+    assert len(outputs[0][1]) > 0
+
+
+@pytest.mark.skip(reason=HF_ACCURACY_SKIP_REASON)
+@pytest.mark.parametrize("model", ACCURACY_MODELS)
+@pytest.mark.parametrize("dtype", ["half"])
+@pytest.mark.parametrize("max_tokens", [8])
+@pytest.mark.parametrize("num_logprobs", [5])
+def test_moss_audio_hf_vllm_accuracy(
+    hf_runner,
+    vllm_runner,
+    model: str,
+    dtype: str,
+    max_tokens: int,
+    num_logprobs: int,
+) -> None:
+    model_info = HF_EXAMPLE_MODELS.find_hf_info(model)
+    model_info.check_available_online(on_fail="skip")
+    model_info.check_transformers_version(on_fail="skip")
+
+    prompts = [f"{MOSS_AUDIO_PLACEHOLDER}\nTranscribe this audio."]
+    audios = [[AudioAsset("mary_had_lamb").audio_and_sample_rate[0]]]
+
+    with vllm_runner(
+        model,
+        dtype=dtype,
+        enforce_eager=True,
+        max_model_len=1024,
+        limit_mm_per_prompt={"audio": 1},
+        trust_remote_code=True,
+    ) as vllm_model:
+        vllm_outputs = vllm_model.generate_greedy_logprobs(
+            prompts,
+            max_tokens,
+            num_logprobs=num_logprobs,
+            audios=audios,
+        )
+
+    with hf_runner(model, dtype=dtype, trust_remote_code=True) as hf_model:
+        hf_outputs = hf_model.generate_greedy_logprobs_limit(
+            prompts,
+            max_tokens,
+            num_logprobs=num_logprobs,
+            audios=audios,
+        )
+
+    check_logprobs_close(
+        outputs_0_lst=hf_outputs,
+        outputs_1_lst=vllm_outputs,
+        name_0="hf",
+        name_1="vllm",
+    )
+
+
+@pytest.mark.core_model
+@pytest.mark.parametrize("parallel_kwargs", PARALLEL_SMOKE_CASES)
+def test_moss_audio_parallel_smoke(vllm_runner, parallel_kwargs) -> None:
+    model = "OpenMOSS-Team/MOSS-Audio-4B-Instruct"
+    required_gpus = parallel_kwargs.get(
+        "tensor_parallel_size", 1
+    ) * parallel_kwargs.get("pipeline_parallel_size", 1)
+    if current_platform.device_count() < required_gpus:
+        # TP/PP integration smoke runs on local or multi-GPU CI only.
+        pytest.skip(f"Requires at least {required_gpus} GPUs")
+
+    model_info = HF_EXAMPLE_MODELS.find_hf_info(model)
+    model_info.check_available_online(on_fail="skip")
+    model_info.check_transformers_version(on_fail="skip")
+
+    prompts = [f"{MOSS_AUDIO_PLACEHOLDER}\nBriefly describe this audio."]
+    audios = [[AudioAsset("mary_had_lamb").audio_and_sample_rate[0]]]
+
+    with vllm_runner(
+        model,
+        dtype="half",
+        enforce_eager=True,
+        max_model_len=1024,
+        limit_mm_per_prompt={"audio": 1},
+        trust_remote_code=True,
+        **parallel_kwargs,
+    ) as vllm_model:
+        outputs = vllm_model.generate_greedy(
+            prompts,
+            max_tokens=4,
+            audios=audios,
+        )
+
+    assert len(outputs) == 1
+    assert len(outputs[0][1]) > 0

tests/models/multimodal/processing/test_common.py

@@ -432,6 +432,12 @@ def test_processing_correctness(
         )
     if model_id == "CohereLabs/cohere-transcribe-03-2026":
         pytest.skip("Fix later")
+    if model_id.startswith("OpenMOSS-Team/MOSS-Audio-"):
+        pytest.skip(
+            "MOSS-Audio uses a custom processor that dynamically expands "
+            "audio placeholders from processed audio lengths. Its vLLM "
+            "processor paths are covered by test_moss_audio.py."
+        )
 
     _test_processing_correctness(
         model_id,

tests/models/multimodal/processing/test_moss_audio.py

@@ -0,0 +1,694 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from types import SimpleNamespace
+
+import numpy as np
+import pytest
+import torch
+from transformers import Qwen3Config
+
+from vllm.model_executor.models.interfaces import SupportsLoRA, supports_lora
+from vllm.model_executor.models.moss_audio import (
+    MOSS_AUDIO_BOS_TOKEN,
+    MOSS_AUDIO_BOS_TOKEN_ID,
+    MOSS_AUDIO_EOS_TOKEN,
+    MOSS_AUDIO_EOS_TOKEN_ID,
+    MOSS_AUDIO_PLACEHOLDER,
+    MOSS_AUDIO_TOKEN,
+    MOSS_AUDIO_TOKEN_ID,
+    GatedMLP,
+    MossAudioConfig,
+    MossAudioDummyInputsBuilder,
+    MossAudioEncoder,
+    MossAudioEncoderConfig,
+    MossAudioModel,
+    MossAudioMultiModalProcessor,
+    MossAudioProcessingInfo,
+    MossAudioProcessor,
+    MossQwen3ForCausalLM,
+    MossQwen3Model,
+)
+from vllm.model_executor.models.utils import AutoWeightsLoader
+from vllm.multimodal.cache import MultiModalProcessorOnlyCache
+from vllm.multimodal.inputs import batched_tensors_equal
+from vllm.sequence import IntermediateTensors
+
+
+class _Tokenizer:
+    def encode(self, text, add_special_tokens=False):
+        del add_special_tokens
+        return [ord(char) for char in text]
+
+    def decode(self, token_ids, **kwargs):
+        del kwargs
+        return "".join(chr(token_id) for token_id in token_ids)
+
+    def batch_decode(self, batch_token_ids, **kwargs):
+        return [self.decode(token_ids, **kwargs) for token_ids in batch_token_ids]
+
+
+class _MMConfig:
+    enable_mm_embeds = False
+    mm_processor_cache_gb = 1
+
+    def merge_mm_processor_kwargs(self, kwargs):
+        return dict(kwargs)
+
+    def get_limit_per_prompt(self, modality):
+        del modality
+        return 3
+
+
+class _ModelConfig:
+    def __init__(self):
+        self.model = "OpenMOSS-Team/MOSS-Audio-4B-Instruct"
+        self.revision = None
+        self.max_model_len = 4096
+        self.encoder_config = {}
+        self.dtype = torch.float32
+        self.hf_config = MossAudioConfig(language_config=Qwen3Config())
+        self.multimodal_config = _MMConfig()
+
+    def get_multimodal_config(self):
+        return self.multimodal_config
+
+    def get_inputs_embeds_size(self):
+        return None
+
+
+class _ProcessingContext:
+    def __init__(self):
+        self.model_config = _ModelConfig()
+        self.tokenizer = _Tokenizer()
+
+    def get_tokenizer(self):
+        return self.tokenizer
+
+    def get_hf_config(self):
+        return self.model_config.hf_config
+
+    def get_mm_config(self):
+        return self.model_config.get_multimodal_config()
+
+    def get_merged_mm_kwargs(self, kwargs):
+        return self.get_mm_config().merge_mm_processor_kwargs(kwargs)
+
+    def call_hf_processor(self, hf_processor, data, kwargs):
+        merged_kwargs = self.get_merged_mm_kwargs(kwargs)
+        merged_kwargs.setdefault("return_tensors", "pt")
+        return hf_processor(**data, **merged_kwargs)
+
+
+class _TestMossAudioProcessingInfo(MossAudioProcessingInfo):
+    def _get_processor_config_defaults(self):
+        return {}
+
+
+def _vllm_config(tensor_parallel_size=1, pipeline_parallel_size=1, hf_config=None):
+    if hf_config is None:
+        hf_config = MossAudioConfig(language_config=Qwen3Config())
+    return SimpleNamespace(
+        model_config=SimpleNamespace(
+            hf_config=hf_config,
+            multimodal_config=None,
+        ),
+        quant_config=None,
+        parallel_config=SimpleNamespace(
+            tensor_parallel_size=tensor_parallel_size,
+            pipeline_parallel_size=pipeline_parallel_size,
+        ),
+    )
+
+
+class _FakeAudioEncoder:
+    dtype = torch.float32
+
+    def __init__(self, deepstack_layers=0):
+        self.deepstack_layers = deepstack_layers
+        self.output_deepstack_hidden_states = None
+        self.input_shape = None
+        self.feature_lens = None
+
+    def __call__(self, audio_data, *, feature_lens, output_deepstack_hidden_states):
+        self.input_shape = tuple(audio_data.shape)
+        self.feature_lens = feature_lens.detach().cpu().clone()
+        self.output_deepstack_hidden_states = output_deepstack_hidden_states
+        lengths = MossAudioEncoder._compute_downsampled_length(feature_lens)
+        hidden_states = torch.ones(1, int(lengths.sum().item()), 8)
+        if not output_deepstack_hidden_states:
+            return hidden_states, None
+        return hidden_states, [
+            hidden_states * scale for scale in range(2, 2 + self.deepstack_layers)
+        ]
+
+
+def _patch_tensor_parallel_for_linear_layers(monkeypatch, tp_size=1, tp_rank=0):
+    import vllm.model_executor.layers.linear as linear_layers
+    import vllm.model_executor.models.moss_audio as moss_audio_module
+    import vllm.model_executor.parameter as parameter_module
+
+    for module in (moss_audio_module, linear_layers, parameter_module):
+        monkeypatch.setattr(
+            module, "get_tensor_model_parallel_world_size", lambda: tp_size
+        )
+    monkeypatch.setattr(
+        linear_layers, "get_tensor_model_parallel_rank", lambda: tp_rank
+    )
+    monkeypatch.setattr(
+        parameter_module, "get_tensor_model_parallel_rank", lambda: tp_rank
+    )
+    monkeypatch.setattr(
+        linear_layers, "tensor_model_parallel_all_reduce", lambda tensor: tensor
+    )
+
+
+def _build_moss_audio_processor(cache=None):
+    ctx = _ProcessingContext()
+    info = _TestMossAudioProcessingInfo(ctx)
+    return (
+        MossAudioMultiModalProcessor(
+            info,
+            MossAudioDummyInputsBuilder(info),
+            cache=cache,
+        ),
+        ctx,
+    )
+
+
+def _assert_mm_inputs_equal(left, right):
+    assert left["prompt_token_ids"] == right["prompt_token_ids"]
+    assert left["mm_hashes"] == right["mm_hashes"]
+
+    left_placeholder = left["mm_placeholders"]["audio"][0]
+    right_placeholder = right["mm_placeholders"]["audio"][0]
+    assert left_placeholder.offset == right_placeholder.offset
+    assert left_placeholder.length == right_placeholder.length
+    assert left_placeholder.is_embed.tolist() == right_placeholder.is_embed.tolist()
+
+    assert batched_tensors_equal(
+        left["mm_kwargs"].get_data(),
+        right["mm_kwargs"].get_data(),
+    )
+
+
+@pytest.mark.parametrize(
+    ("prompt", "prefix"),
+    [
+        (
+            f"before {MOSS_AUDIO_PLACEHOLDER} after",
+            [*[ord(char) for char in "before "], MOSS_AUDIO_BOS_TOKEN_ID],
+        ),
+        (
+            f"before {MOSS_AUDIO_BOS_TOKEN}{MOSS_AUDIO_TOKEN}"
+            f"{MOSS_AUDIO_TOKEN}{MOSS_AUDIO_EOS_TOKEN} after",
+            [*[ord(char) for char in "before "], MOSS_AUDIO_BOS_TOKEN_ID],
+        ),
+        ("Describe this audio.", [MOSS_AUDIO_BOS_TOKEN_ID]),
+    ],
+)
+def test_moss_audio_processor_expands_audio_placeholders(prompt, prefix):
+    raw_mel_len = 17
+    processed = MossAudioProcessor(_Tokenizer())(
+        text=prompt, audio=[torch.zeros(160 * raw_mel_len)]
+    )
+    input_ids = processed["input_ids"][0].tolist()
+
+    assert input_ids[: len(prefix)] == prefix
+    assert input_ids.count(MOSS_AUDIO_BOS_TOKEN_ID) == 1
+    assert input_ids.count(MOSS_AUDIO_EOS_TOKEN_ID) == 1
+    assert input_ids.count(MOSS_AUDIO_TOKEN_ID) == (
+        MossAudioEncoder.compute_num_audio_tokens(raw_mel_len)
+    )
+    assert processed["audio_data"].shape == (1, 128, raw_mel_len)
+    assert processed["audio_data_seqlens"].tolist() == [raw_mel_len]
+
+
+def test_moss_audio_processor_preserves_placeholder_without_audio():
+    processed = MossAudioProcessor(_Tokenizer())(
+        text=f"before {MOSS_AUDIO_PLACEHOLDER} after"
+    )
+
+    assert processed["input_ids"][0].tolist() == [
+        *[ord(char) for char in "before "],
+        MOSS_AUDIO_BOS_TOKEN_ID,
+        MOSS_AUDIO_TOKEN_ID,
+        MOSS_AUDIO_EOS_TOKEN_ID,
+        *[ord(char) for char in " after"],
+    ]
+    assert "audio_data" not in processed
+    assert "audio_data_seqlens" not in processed
+
+
+def test_moss_audio_multimodal_processor_handles_token_and_cache_paths():
+    raw_mel_len = 17
+    audio = np.zeros(160 * raw_mel_len, dtype=np.float32)
+    prompt = f"{MOSS_AUDIO_PLACEHOLDER}\nTranscribe this audio."
+
+    baseline_processor, ctx = _build_moss_audio_processor()
+    mm_items = baseline_processor.info.parse_mm_data({"audio": [audio]})
+    token_prompt = ctx.get_tokenizer().encode(prompt, add_special_tokens=False)
+
+    baseline_text = baseline_processor(
+        prompt,
+        mm_items=mm_items,
+        hf_processor_mm_kwargs={},
+    )
+    baseline_token = baseline_processor(
+        token_prompt,
+        mm_items=mm_items,
+        hf_processor_mm_kwargs={},
+    )
+
+    cache = MultiModalProcessorOnlyCache(ctx.model_config)
+    cached_processor, _ = _build_moss_audio_processor(cache=cache)
+    cached_text_miss = cached_processor(
+        prompt,
+        mm_items=mm_items,
+        hf_processor_mm_kwargs={},
+    )
+    cached_text_hit = cached_processor(
+        prompt,
+        mm_items=mm_items,
+        hf_processor_mm_kwargs={},
+    )
+    cached_token_hit = cached_processor(
+        token_prompt,
+        mm_items=mm_items,
+        hf_processor_mm_kwargs={},
+    )
+
+    expected_audio_tokens = MossAudioEncoder.compute_num_audio_tokens(raw_mel_len)
+    prompt_token_ids = baseline_text["prompt_token_ids"]
+    assert prompt_token_ids.count(MOSS_AUDIO_TOKEN_ID) == expected_audio_tokens
+    assert baseline_text["mm_placeholders"]["audio"][0].length == (
+        expected_audio_tokens + 2
+    )
+
+    _assert_mm_inputs_equal(baseline_text, baseline_token)
+    _assert_mm_inputs_equal(baseline_text, cached_text_miss)
+    _assert_mm_inputs_equal(baseline_text, cached_text_hit)
+    _assert_mm_inputs_equal(baseline_text, cached_token_hit)
+
+
+def test_moss_audio_supports_language_model_lora_only():
+    assert supports_lora(MossAudioModel)
+
+    model = object.__new__(MossAudioModel)
+    assert isinstance(model, SupportsLoRA)
+
+    mapping = model.get_mm_mapping()
+    assert mapping.language_model == ["language_model."]
+    assert mapping.tower_model == []
+    assert mapping.connector == []
+
+
+def test_moss_audio_error_paths():
+    model = object.__new__(MossAudioModel)
+    with pytest.raises(ValueError, match="DeepStack audio token count mismatch"):
+        model._cache_deepstack_input_embeds(
+            inputs_embeds=torch.zeros(4, 8),
+            deepstack_embeddings=((torch.ones(1, 8),),),
+            is_multimodal=torch.tensor([False, True, True, False]),
+        )
+
+    with pytest.raises(ValueError, match="too short"):
+        MossAudioProcessor(_Tokenizer())(
+            text=MOSS_AUDIO_PLACEHOLDER, audio=[torch.empty(0)]
+        )
+    with pytest.raises(ValueError, match="too short"):
+        model._parse_and_validate_audio_input(
+            audio_data=torch.zeros(1, 128, 1),
+            audio_data_seqlens=torch.tensor([0], dtype=torch.long),
+        )
+
+
+def test_moss_audio_validates_tp_config():
+    vllm_config = _vllm_config(tensor_parallel_size=2)
+    vllm_config.model_config.hf_config.adapter_hidden_size = 7
+
+    with pytest.raises(ValueError, match="adapter_hidden_size"):
+        MossAudioModel(vllm_config=vllm_config)
+
+    vllm_config = _vllm_config(tensor_parallel_size=2)
+    vllm_config.model_config.hf_config.audio_config.d_model = 6
+    vllm_config.model_config.hf_config.audio_config.encoder_attention_heads = 3
+    with pytest.raises(ValueError, match="encoder_attention_heads"):
+        MossAudioModel(vllm_config=vllm_config)
+
+
+def test_moss_audio_rejects_audio_data_list_seqlen_count_mismatch():
+    model = object.__new__(MossAudioModel)
+
+    with pytest.raises(ValueError, match="audio_data batch size"):
+        model._parse_and_validate_audio_input(
+            audio_data=[torch.zeros(128, 8), torch.zeros(128, 11)],
+            audio_data_seqlens=torch.tensor([8], dtype=torch.long),
+        )
+
+
+@pytest.mark.parametrize("deepstack_scales", [(), (7, 11)])
+def test_moss_audio_embed_multimodal_packs_by_audio(deepstack_scales):
+    model = object.__new__(MossAudioModel)
+    model.audio_encoder = _FakeAudioEncoder(len(deepstack_scales))
+    model.audio_adapter = lambda hidden_states: hidden_states * 5
+    model.deepstack_audio_merger_list = [
+        lambda hidden_states, scale=scale: hidden_states * scale
+        for scale in deepstack_scales
+    ]
+    model.deepstack_input_embeds = None
+
+    embeddings = model.embed_multimodal(
+        audio_data=torch.zeros(2, 128, 9),
+        audio_data_seqlens=torch.tensor([8, 9], dtype=torch.long),
+    )
+
+    assert model.audio_encoder.output_deepstack_hidden_states is bool(deepstack_scales)
+    assert [embeds.shape for embeds in embeddings] == [
+        torch.Size([1, 8 * (1 + len(deepstack_scales))]),
+        torch.Size([2, 8 * (1 + len(deepstack_scales))]),
+    ]
+    if not deepstack_scales:
+        assert model.deepstack_input_embeds is None
+        return
+
+    main_embeddings, deepstack_embeddings = model._split_multimodal_embeddings(
+        embeddings, hidden_size=8
+    )
+    assert [embeds.shape for embeds in main_embeddings] == [
+        torch.Size([1, 8]),
+        torch.Size([2, 8]),
+    ]
+    assert [[e.shape for e in layer] for layer in deepstack_embeddings] == [
+        [torch.Size([1, 8]), torch.Size([2, 8])] for _ in deepstack_scales
+    ]
+    assert torch.equal(main_embeddings[0], torch.full((1, 8), 5.0))
+    for idx, scale in enumerate(deepstack_scales):
+        assert torch.equal(
+            deepstack_embeddings[idx][0],
+            torch.full((1, 8), float((idx + 2) * scale)),
+        )
+
+
+def test_moss_audio_embed_input_ids_caches_packed_deepstack():
+    class _FakeLanguageModel:
+        def embed_input_ids(self, input_ids):
+            return torch.zeros(input_ids.shape[0], 8)
+
+    model = object.__new__(MossAudioModel)
+    model.language_model = _FakeLanguageModel()
+    model.deepstack_audio_merger_list = [object(), object()]
+    model.deepstack_input_embeds = None
+    multimodal_embeddings = (
+        torch.cat([torch.full((1, 8), x) for x in (5.0, 14.0, 33.0)], dim=-1),
+        torch.cat([torch.full((2, 8), x) for x in (7.0, 22.0, 44.0)], dim=-1),
+    )
+    is_multimodal = torch.tensor([False, True, True, True, False])
+
+    inputs_embeds = model.embed_input_ids(
+        input_ids=torch.arange(5),
+        multimodal_embeddings=multimodal_embeddings,
+        is_multimodal=is_multimodal,
+    )
+
+    assert torch.equal(inputs_embeds[1], torch.full((8,), 5.0))
+    assert torch.equal(inputs_embeds[2], torch.full((8,), 7.0))
+    assert torch.equal(inputs_embeds[3], torch.full((8,), 7.0))
+    assert model.deepstack_input_embeds is not None
+    tensors = model.deepstack_input_embeds.tensors
+    assert set(tensors) == {"deepstack_input_embeds_0", "deepstack_input_embeds_1"}
+    for tensor in tensors.values():
+        assert tensor[is_multimodal].abs().sum() > 0
+        assert torch.equal(tensor[~is_multimodal], torch.zeros(2, 8))
+
+
+def _patch_pp_group(monkeypatch, *, first=True, last=True):
+    import vllm.model_executor.models.moss_audio as moss_audio_module
+
+    monkeypatch.setattr(
+        moss_audio_module,
+        "get_pp_group",
+        lambda: SimpleNamespace(is_first_rank=first, is_last_rank=last),
+    )
+
+
+def test_moss_audio_pp_forward_routes_deepstack(monkeypatch):
+    for first in (True, False):
+        calls: list[dict[str, object]] = []
+
+        def fake_lm(*args, _calls=calls, **kwargs):
+            del args
+            _calls.append(kwargs)
+            return torch.ones(1, 1)
+
+        _patch_pp_group(monkeypatch, first=first)
+        model = object.__new__(MossAudioModel)
+        torch.nn.Module.__init__(model)
+        model.language_model = fake_lm
+        cached = IntermediateTensors({"deepstack_input_embeds_0": torch.ones(3, 8)})
+        inter = IntermediateTensors(
+            {
+                "hidden_states": torch.ones(3, 8),
+                "residual": torch.zeros(3, 8),
+                "deepstack_input_embeds_0": torch.full((3, 8), 5.0),
+            }
+        )
+        inputs_embeds = torch.full((3, 8), 9.0)
+        model.deepstack_input_embeds = cached
+
+        model.forward(
+            input_ids=None,
+            positions=torch.arange(3),
+            intermediate_tensors=None if first else inter,
+            inputs_embeds=inputs_embeds if first else None,
+        )
+
+        kwargs = calls[0]
+        assert kwargs["inputs_embeds"] is (inputs_embeds if first else None)
+        assert kwargs["deepstack_input_embeds"] is (cached if first else inter)
+        assert model.deepstack_input_embeds is None
+
+    calls = []
+
+    def fake_lm_non_first_rank(*args, **kwargs):
+        del args
+        calls.append(kwargs)
+        return torch.ones(1, 1)
+
+    _patch_pp_group(monkeypatch, first=False)
+    model = object.__new__(MossAudioModel)
+    torch.nn.Module.__init__(model)
+    model.language_model = fake_lm_non_first_rank
+    model.deepstack_input_embeds = IntermediateTensors({})
+    inter = IntermediateTensors(
+        {
+            "hidden_states": torch.ones(3, 8),
+            "residual": torch.zeros(3, 8),
+        }
+    )
+
+    model.forward(
+        input_ids=None,
+        positions=torch.arange(3),
+        intermediate_tensors=inter,
+        inputs_embeds=torch.ones(3, 8),
+    )
+    assert calls[0]["inputs_embeds"] is None
+    assert calls[0]["deepstack_input_embeds"] is inter
+
+
+def test_moss_qwen3_deepstack_keys_for_pp(monkeypatch):
+    class AddOne(torch.nn.Module):
+        def forward(self, positions, hidden_states, residual):
+            del positions, residual
+            return hidden_states + 1, torch.zeros_like(hidden_states)
+
+    def make_model(num_layers, deepstack_layers=None):
+        model = object.__new__(MossQwen3Model)
+        torch.nn.Module.__init__(model)
+        model.start_layer, model.end_layer = 0, num_layers
+        model.layers = torch.nn.ModuleList([AddOne() for _ in range(num_layers)])
+        model.norm = lambda hidden_states, residual: (hidden_states, residual)
+        model._maybe_add_hidden_state = lambda aux, *args: aux
+        model.deepstack_inject_layer_indices = (
+            range(0) if deepstack_layers is None else deepstack_layers
+        )
+        return model
+
+    _patch_pp_group(monkeypatch, first=True, last=True)
+    output = make_model(3).forward(
+        input_ids=None,
+        positions=torch.arange(2),
+        inputs_embeds=torch.zeros(2, 4),
+        deepstack_input_embeds=IntermediateTensors(
+            {
+                "deepstack_input_embeds_2": torch.full((2, 4), 5.0),
+            }
+        ),
+    )
+    assert torch.equal(output, torch.full((2, 4), 8.0))
+
+    _patch_pp_group(monkeypatch, first=True, last=False)
+    deepstack = IntermediateTensors(
+        {
+            "deepstack_input_embeds_0": torch.full((2, 4), 7.0),
+            "deepstack_input_embeds_3": torch.full((2, 4), 11.0),
+        }
+    )
+    output = make_model(2, range(4)).forward(
+        input_ids=None,
+        positions=torch.arange(2),
+        inputs_embeds=torch.zeros(2, 4),
+        deepstack_input_embeds=deepstack,
+    )
+    assert isinstance(output, IntermediateTensors)
+    assert set(output.tensors) == {
+        "hidden_states",
+        "residual",
+        "deepstack_input_embeds_2",
+        "deepstack_input_embeds_3",
+    }
+    assert torch.equal(output["hidden_states"], torch.full((2, 4), 9.0))
+    assert torch.equal(output["deepstack_input_embeds_2"], torch.zeros(2, 4))
+    assert output["deepstack_input_embeds_3"] is deepstack["deepstack_input_embeds_3"]
+
+    inner_model = make_model(0, range(2))
+    inner_model.make_empty_intermediate_tensors = lambda batch, dtype, device: (
+        IntermediateTensors(
+            {
+                "hidden_states": torch.zeros(batch, 4, dtype=dtype, device=device),
+                "residual": torch.zeros(batch, 4, dtype=dtype, device=device),
+            }
+        )
+    )
+    language_model = object.__new__(MossQwen3ForCausalLM)
+    torch.nn.Module.__init__(language_model)
+    language_model.model = inner_model
+    language_model.config = SimpleNamespace(hidden_size=4)
+    language_model.deepstack_inject_layer_indices = range(2)
+
+    tensors = MossQwen3ForCausalLM.make_empty_intermediate_tensors(
+        language_model,
+        batch_size=3,
+        dtype=torch.float16,
+        device=torch.device("cpu"),
+    )
+
+    assert set(tensors.tensors) == {
+        "hidden_states",
+        "residual",
+        "deepstack_input_embeds_0",
+        "deepstack_input_embeds_1",
+    }
+    assert tensors["deepstack_input_embeds_0"].shape == (3, 4)
+    assert tensors["deepstack_input_embeds_0"].dtype == torch.float16
+
+    _patch_pp_group(monkeypatch, first=True, last=False)
+    forward_tensors = inner_model.forward(
+        input_ids=None,
+        positions=torch.arange(3),
+        inputs_embeds=torch.ones(3, 4, dtype=torch.float16),
+        deepstack_input_embeds=None,
+    )
+    assert isinstance(forward_tensors, IntermediateTensors)
+    assert set(forward_tensors.tensors) == set(tensors.tensors)
+
+
+@pytest.mark.parametrize("tp_size", [1, 2])
+def test_moss_audio_gated_mlp_tp_shapes_and_loading(monkeypatch, tp_size):
+    from vllm.config import VllmConfig, set_current_vllm_config
+    from vllm.config.device import DeviceConfig
+
+    _patch_tensor_parallel_for_linear_layers(monkeypatch, tp_size=tp_size)
+    with set_current_vllm_config(VllmConfig(device_config=DeviceConfig(device="cpu"))):
+        mlp = GatedMLP(input_size=4, hidden_size=8, output_size=6)
+
+    params = dict(mlp.named_parameters())
+    assert params["gate_up_proj.weight"].shape == torch.Size([16 // tp_size, 4])
+    assert params["down_proj.weight"].shape == torch.Size([6, 8 // tp_size])
+
+    gate_weight = torch.arange(32, dtype=torch.float32).reshape(8, 4)
+    up_weight = torch.arange(100, 132, dtype=torch.float32).reshape(8, 4)
+    down_weight = torch.arange(48, dtype=torch.float32).reshape(6, 8)
+    loaded = mlp.load_weights(
+        [
+            ("gate_proj.weight", gate_weight),
+            ("up_proj.weight", up_weight),
+            ("down_proj.weight", down_weight),
+        ]
+    )
+
+    assert loaded == {"gate_up_proj.weight", "down_proj.weight"}
+    shard = 8 // tp_size
+    assert torch.equal(params["gate_up_proj.weight"][:shard], gate_weight[:shard])
+    assert torch.equal(params["gate_up_proj.weight"][shard:], up_weight[:shard])
+    assert torch.equal(params["down_proj.weight"], down_weight[:, : 8 // tp_size])
+
+    with set_current_vllm_config(VllmConfig(device_config=DeviceConfig(device="cpu"))):
+        packed_mlp = GatedMLP(input_size=4, hidden_size=8, output_size=6)
+    packed_params = dict(packed_mlp.named_parameters())
+    loaded = packed_mlp.load_weights(
+        [("gate_up_proj.weight", torch.cat([gate_weight, up_weight], dim=0))]
+    )
+    assert loaded == {"gate_up_proj.weight"}
+    assert torch.equal(
+        packed_params["gate_up_proj.weight"][:shard],
+        gate_weight[:shard],
+    )
+    assert torch.equal(
+        packed_params["gate_up_proj.weight"][shard:],
+        up_weight[:shard],
+    )
+
+
+def test_moss_audio_encoder_loads_realistic_attention_weight_names(monkeypatch):
+    from vllm.config import VllmConfig, set_current_vllm_config
+    from vllm.config.device import DeviceConfig
+
+    _patch_tensor_parallel_for_linear_layers(monkeypatch, tp_size=2)
+    config = MossAudioEncoderConfig(
+        d_model=8,
+        output_dim=8,
+        num_mel_bins=8,
+        encoder_layers=1,
+        encoder_attention_heads=2,
+        encoder_ffn_dim=16,
+        downsample_hidden_size=2,
+        deepstack_encoder_layer_indexes=[],
+    )
+    with set_current_vllm_config(VllmConfig(device_config=DeviceConfig(device="cpu"))):
+        encoder = MossAudioEncoder(config)
+
+    attention = encoder.layers[0].self_attn
+    assert all(hasattr(attention, name) for name in ("q_proj", "k_proj", "v_proj"))
+    assert hasattr(attention, "out_proj")
+    assert not hasattr(attention, "qkv")
+    assert attention.k_proj.bias is None
+
+    weight_names = [
+        "layers.0.self_attn.q_proj.weight",
+        "layers.0.self_attn.q_proj.bias",
+        "layers.0.self_attn.k_proj.weight",
+        "layers.0.self_attn.v_proj.weight",
+        "layers.0.self_attn.v_proj.bias",
+        "layers.0.self_attn.out_proj.weight",
+        "layers.0.self_attn.out_proj.bias",
+        "conv1.weight",
+        "conv1.bias",
+    ]
+    params = dict(encoder.named_parameters(remove_duplicate=False))
+    assert "layers.0.self_attn.k_proj.bias" not in params
+    weights = {
+        name: torch.full_like(params[name], fill_value=float(i + 1))
+        for i, name in enumerate(weight_names)
+    }
+
+    loaded = AutoWeightsLoader(encoder).load_weights(weights.items())
+
+    assert "load_weights" not in MossAudioEncoder.__dict__
+    assert loaded == set(weight_names)
+    assert not any(".qkv." in name for name in loaded)
+    assert torch.equal(
+        params["layers.0.self_attn.q_proj.weight"],
+        weights["layers.0.self_attn.q_proj.weight"],
+    )

tests/models/registry.py

@@ -1127,6 +1127,15 @@ _MULTIMODAL_EXAMPLE_MODELS = {
         tokenizer="moondream/starmie-v1",
         trust_remote_code=True,
     ),
+    "MossAudioModel": _HfExamplesInfo(
+        "OpenMOSS-Team/MOSS-Audio-4B-Instruct",
+        extras={
+            "4b-thinking": "OpenMOSS-Team/MOSS-Audio-4B-Thinking",
+            "8b-instruct": "OpenMOSS-Team/MOSS-Audio-8B-Instruct",
+            "8b-thinking": "OpenMOSS-Team/MOSS-Audio-8B-Thinking",
+        },
+        trust_remote_code=True,
+    ),
     "HfMoondream": _HfExamplesInfo(
         "moondream/moondream3-preview",
         tokenizer="moondream/starmie-v1",

vllm/benchmarks/throughput.py

@@ -20,6 +20,7 @@ from vllm.benchmarks.datasets import (
     ASRDataset,
     BurstGPTDataset,
     ConversationDataset,
+    CustomAudioDataset,
     InstructCoderDataset,
     MultiModalConversationDataset,
     PrefixRepetitionRandomDataset,
@@ -240,9 +241,24 @@ def _run_vllm_chat_requests(
 ) -> tuple[float, list[RequestOutput]]:
     from vllm import SamplingParams
 
-    prompts = [request.prompt for request in requests]
+    prompts = []
     sampling_params: list[SamplingParams] = []
     for request in requests:
+        if isinstance(request.prompt, list):
+            prompts.append(request.prompt)
+        else:
+            content: list[dict[str, Any]] = [{"type": "text", "text": request.prompt}]
+            if request.multi_modal_data is not None:
+                if isinstance(request.multi_modal_data, list):
+                    content.extend(request.multi_modal_data)
+                elif isinstance(request.multi_modal_data, dict):
+                    content.append(request.multi_modal_data)
+                else:
+                    raise TypeError(
+                        "Could not process multimodal content of type: "
+                        f"{type(request.multi_modal_data)}"
+                    )
+            prompts.append([{"role": "user", "content": content}])
         sampling_params.append(
             SamplingParams(
                 n=n,
@@ -510,6 +526,7 @@ def get_requests(args, tokenizer):
         "max_loras": args.max_loras,
         "lora_assignment": getattr(args, "lora_assignment", "random"),
         "num_requests": args.num_prompts,
+        "no_oversample": getattr(args, "no_oversample", False),
     }
 
     if args.dataset_name == "random" or (
@@ -550,6 +567,16 @@ def get_requests(args, tokenizer):
             sample_kwargs["output_len"] = args.output_len
     elif args.dataset_name == "burstgpt":
         dataset_cls = BurstGPTDataset
+    elif args.dataset_name == "custom_audio":
+        dataset_cls = CustomAudioDataset
+        sample_kwargs["enable_multimodal_chat"] = getattr(
+            args, "enable_multimodal_chat", False
+        )
+        custom_output_len = getattr(args, "custom_output_len", None)
+        if custom_output_len is not None:
+            sample_kwargs["output_len"] = custom_output_len
+        elif args.output_len is not None:
+            sample_kwargs["output_len"] = args.output_len
     elif args.dataset_name == "hf":
         if args.output_len is not None:
             sample_kwargs["output_len"] = args.output_len
@@ -894,6 +921,7 @@ def add_cli_args(parser: argparse.ArgumentParser):
             "prefix_repetition",
             "random-mm",
             "random-rerank",
+            "custom_audio",
         ],
         help="Name of the dataset to benchmark on.",
         default="sharegpt",
@@ -910,6 +938,22 @@ def add_cli_args(parser: argparse.ArgumentParser):
     parser.add_argument(
         "--dataset-path", type=str, default=None, help="Path to the dataset"
     )
+    parser.add_argument(
+        "--no-oversample",
+        action="store_true",
+        help="Do not oversample if the dataset has fewer samples than num-prompts.",
+    )
+    parser.add_argument(
+        "--enable-multimodal-chat",
+        action="store_true",
+        help="Enable multimodal chat transformation for datasets that support it.",
+    )
+    parser.add_argument(
+        "--custom-output-len",
+        type=int,
+        default=None,
+        help="Number of output tokens per request for custom datasets.",
+    )
     parser.add_argument(
         "--input-len",
         type=int,

vllm/config/kernel.py

@@ -199,7 +199,7 @@ class KernelConfig:
     - "auto": Automatically select the best backend based on model and hardware
     - "cutlass": Use CUTLASS-based kernels
     - "flashinfer_cutlass": Use FlashInfer with CUTLASS kernels
-    - "flashinfer_cutedsl": Use FlashInfer with CuteDSL kernels
+    - "flashinfer_cutedsl": Use FlashInfer with CuTe-DSL kernels (NVFP4, MXFP8)
     - "flashinfer_trtllm": Use FlashInfer with TensorRT-LLM kernels
     - "flashinfer_cudnn": Use FlashInfer with cuDNN kernels
     - "flashinfer_b12x": Use FlashInfer b12x CuteDSL NVFP4 GEMM (SM120+)

vllm/envs.py

@@ -1580,10 +1580,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.

vllm/model_executor/kernels/linear/__init__.py

@@ -88,6 +88,7 @@ from vllm.model_executor.kernels.linear.mxfp8.emulation import (
     EmulationMxfp8LinearKernel,
 )
 from vllm.model_executor.kernels.linear.mxfp8.flashinfer import (
+    FlashInferCutedslMxfp8LinearKernel,
     FlashInferCutlassMxfp8LinearKernel,
 )
 from vllm.model_executor.kernels.linear.mxfp8.marlin import (
@@ -212,6 +213,7 @@ _LINEAR_BACKEND_KERNEL_MAP: dict[str, set[type]] = {
     },
     "flashinfer_cutedsl": {
         FlashInferCuteDslNvFp4LinearKernel,
+        FlashInferCutedslMxfp8LinearKernel,
     },
     "flashinfer_trtllm": {
         FlashInferTrtllmNvFp4LinearKernel,
@@ -385,6 +387,7 @@ _POSSIBLE_KERNELS: dict[PlatformEnum, list[type[MPLinearKernel]]] = {
 # in priority/performance order (when available)
 _POSSIBLE_MXFP8_KERNELS: dict[PlatformEnum, list[type[Mxfp8LinearKernel]]] = {
     PlatformEnum.CUDA: [
+        FlashInferCutedslMxfp8LinearKernel,
         FlashInferCutlassMxfp8LinearKernel,
         MarlinMxfp8LinearKernel,
         EmulationMxfp8LinearKernel,
@@ -1036,6 +1039,7 @@ __all__ = [
     "MxFp4LinearLayerConfig",
     "FlashInferMxFp4LinearKernel",
     "MarlinMxFp4LinearKernel",
+    "FlashInferCutedslMxfp8LinearKernel",
     "FlashInferCutlassMxfp8LinearKernel",
     "MarlinMxfp8LinearKernel",
     "XPUMxFp8LinearKernel",

vllm/model_executor/kernels/linear/mxfp8/flashinfer.py

@@ -11,6 +11,7 @@ from vllm.model_executor.layers.quantization.utils.mxfp8_utils import (
 )
 from vllm.platforms import current_platform
 from vllm.utils import flashinfer as vllm_flashinfer
+from vllm.utils.flashinfer import has_flashinfer_cutedsl
 
 from .Mxfp8LinearKernel import Mxfp8LinearKernel, Mxfp8LinearLayerConfig
 
@@ -91,3 +92,84 @@ class FlashInferCutlassMxfp8LinearKernel(Mxfp8LinearKernel):
 
         output_shape = (*input_shape[:-1], N)
         return output.view(output_shape)
+
+
+class FlashInferCutedslMxfp8LinearKernel(Mxfp8LinearKernel):
+    """MXFP8 W8A8 GEMM via FlashInfer CuTe-DSL (SM100/SM103)."""
+
+    @classmethod
+    def is_supported(
+        cls, compute_capability: int | None = None
+    ) -> tuple[bool, str | None]:
+        if not (
+            current_platform.is_cuda()
+            and current_platform.is_device_capability_family(100)
+        ):
+            return False, "requires sm_100/sm_103 (Blackwell)"
+        if not has_flashinfer_cutedsl():
+            return False, "requires FlashInfer CuTe-DSL module"
+        return True, None
+
+    @classmethod
+    def can_implement(cls, c: Mxfp8LinearLayerConfig) -> tuple[bool, str | None]:
+        return True, None
+
+    def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
+        weight = layer.weight.data  # [N, K]
+        N, K = weight.shape
+
+        scale_k = K // MXFP8_BLOCK_SIZE
+        weight_scale_2d = layer.weight_scale.data[:N, :scale_k].contiguous()
+        weight_scale_swizzled = swizzle_mxfp8_scale(weight_scale_2d, M=N, K=K)
+
+        # Store weight column-major [K, N] as mm_mxfp8 expects for operand B.
+        layer.weight = Parameter(weight.contiguous().t(), requires_grad=False)
+        layer.weight_scale = Parameter(
+            weight_scale_swizzled.contiguous(), requires_grad=False
+        )
+
+    def apply_weights(
+        self,
+        layer: torch.nn.Module,
+        x: torch.Tensor,
+        bias: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        weight = layer.weight  # [K, N], column-major
+        weight_scale = layer.weight_scale
+        out_dtype = x.dtype
+        K, N = weight.shape
+
+        input_shape = x.shape
+        input_2d = x.view(-1, K)
+        min_dim = 128
+
+        assert min_dim <= K, (
+            f"mm_mxfp8 requires K >= {min_dim}, got K={K}. "
+            f"in_features is too small for mm_mxfp8."
+        )
+        assert K % MXFP8_BLOCK_SIZE == 0, (
+            f"mm_mxfp8 requires K to be divisible by {MXFP8_BLOCK_SIZE}, got K={K}."
+        )
+        assert min_dim <= N, (
+            f"mm_mxfp8 requires N >= {min_dim}, got N={N}. "
+            f"out_features is too small for mm_mxfp8."
+        )
+
+        input_mxfp8, input_scale = mxfp8_e4m3_quantize(
+            input_2d, is_sf_swizzled_layout=True
+        )
+
+        output = vllm_flashinfer.mm_mxfp8(
+            input_mxfp8,
+            weight,
+            input_scale,
+            weight_scale,
+            out_dtype=out_dtype,
+            backend="cute-dsl",
+        )
+
+        if bias is not None:
+            output = output + bias
+
+        output_shape = (*input_shape[:-1], N)
+        return output.view(output_shape)

vllm/model_executor/layers/fused_moe/oracle/nvfp4.py

@@ -174,6 +174,8 @@ def select_nvfp4_moe_backend(
 
     NVFP4_BACKENDS_WITH_CLAMP = {
         NvFp4MoeBackend.FLASHINFER_TRTLLM,
+        NvFp4MoeBackend.FLASHINFER_CUTLASS,
+        NvFp4MoeBackend.MARLIN,
     }
 
     if config.swiglu_limit is not None:

vllm/model_executor/layers/fused_moe/router/gate_linear.py

@@ -14,7 +14,7 @@ from vllm.utils.torch_utils import direct_register_custom_op
 class GateLinear(ReplicatedLinear):
     """MoE gate linear layer with multi-tier GEMM dispatch:
 
-    1. DSV3 specialized kernel (SM90+, fp32 out, M<=16, H=7168, E=256/384)
+    1. DSV3 specialized kernel (SM90+, M<=16, H=7168 E=256/384, H=6144 E=256)
     2. fp32 specialized kernel  (SM90+, bf16/fp32 in, fp32 out,
        M<=32, H=3072, E=256)
     3. cuBLAS bf16×bf16→fp32 (SM90+ + bf16 weight + fp32 out_dtype)
@@ -25,9 +25,14 @@ class GateLinear(ReplicatedLinear):
     method which is only known later).
     """
 
-    # Dimensions supported by the DSV3 specialized kernel
+    # Dimensions supported by the DSV3 specialized kernel.
+    # Valid (hidden_size, num_experts) combinations:
+    #   (7168, 256) -> DeepSeek-V3,  (7168, 384) -> Kimi-K2,
+    #   (6144, 256) -> GLM-5
     DSV3_SUPPORTED_NUM_EXPERTS = [256, 384]
-    DSV3_SUPPORTED_HIDDEN_SIZES = [7168]
+    DSV3_SUPPORTED_HIDDEN_SIZES = [7168, 6144]
+    # num_experts=384 is only instantiated for hidden_size=7168.
+    DSV3_UNSUPPORTED_SHAPES = {(6144, 384)}
 
     # (hidden_size, num_experts) pairs with an instantiated fp32 kernel:
     #   (3072, 256) -> MiniMax-M2/M2.5,  (6144, 128) -> MiniMax-M3
@@ -71,6 +76,7 @@ class GateLinear(ReplicatedLinear):
             self.allow_specialized_router_gemm
             and output_size in self.DSV3_SUPPORTED_NUM_EXPERTS
             and input_size in self.DSV3_SUPPORTED_HIDDEN_SIZES
+            and (input_size, output_size) not in self.DSV3_UNSUPPORTED_SHAPES
         )
         # See https://github.com/vllm-project/vllm/pull/44217
         # for more details.

vllm/model_executor/layers/sparse_attn_indexer.py

@@ -334,7 +334,32 @@ def sparse_attn_indexer(
         num_rows = logits.shape[0]
         topk_indices = topk_indices_buffer[:num_padded_tokens, :topk_tokens]
 
-        if current_platform.is_cuda() and topk_tokens in (512, 1024, 2048):
+        use_cooperative_topk = (
+            current_platform.is_cuda()
+            and topk_tokens in (512, 1024, 2048)
+            and num_rows <= 32
+            and logits.stride(0) % 4 == 0  # TMA 16-byte alignment
+            and current_platform.has_device_capability(90)
+        )
+        use_persistent_topk = current_platform.is_cuda() and topk_tokens in (
+            512,
+            1024,
+            2048,
+        )
+        if use_cooperative_topk:
+            workspace_manager = current_workspace_manager()
+            (topk_workspace,) = workspace_manager.get_simultaneous(
+                ((RADIX_TOPK_WORKSPACE_SIZE,), torch.uint8),
+            )
+            torch.ops._C.cooperative_topk(
+                logits,
+                seq_lens,
+                topk_indices,
+                topk_workspace,
+                topk_tokens,
+                attn_metadata_narrowed.max_seq_len,
+            )
+        elif use_persistent_topk:
             workspace_manager = current_workspace_manager()
             (topk_workspace,) = workspace_manager.get_simultaneous(
                 ((RADIX_TOPK_WORKSPACE_SIZE,), torch.uint8),

vllm/model_executor/models/registry.py

@@ -500,6 +500,7 @@ _MULTIMODAL_MODELS = {
     "MolmoForCausalLM": ("molmo", "MolmoForCausalLM"),
     "Molmo2ForConditionalGeneration": ("molmo2", "Molmo2ForConditionalGeneration"),
     "Moondream3ForCausalLM": ("moondream3", "Moondream3ForCausalLM"),
+    "MossAudioModel": ("moss_audio", "MossAudioModel"),
     "HfMoondream": ("moondream3", "Moondream3ForCausalLM"),
     "NemotronH_Nano_VL_V2": ("nano_nemotron_vl", "NemotronH_Nano_VL_V2"),
     "NemotronH_Nano_Omni_Reasoning_V3": ("nano_nemotron_vl", "NemotronH_Nano_VL_V2"),

vllm/models/minimax_m3/amd/model.py

@@ -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",
         )

vllm/models/minimax_m3/common/indexer.py

@@ -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

vllm/models/minimax_m3/common/ops/index_topk.py

@@ -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, (
@@ -834,11 +849,14 @@ 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

vllm/models/minimax_m3/common/sparse_attention.py

@@ -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,

vllm/models/minimax_m3/nvidia/indexer_msa.py

@@ -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

vllm/models/minimax_m3/nvidia/model.py

@@ -402,6 +402,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 +490,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 +524,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 +582,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 +624,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 +639,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:
@@ -662,6 +675,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(
@@ -747,11 +761,33 @@ 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)
+            # Pad tokens to a multiple of 4 so the buffer head stride stays
+            # int4-aligned for build_k2q_csr's vectorised int4 loads.
+            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",
         )

vllm/models/minimax_m3/nvidia/sparse_attention_msa.py

@@ -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,8 @@ class MiniMaxM3SparseMSAImpl(MiniMaxM3SparseImpl):
             )
 
             p = main_md.prefill
-            assert p is not None and prefill_topk is not None
+            assert p is not None
+            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)

vllm/transformers_utils/model_arch_config_convertor.py

@@ -577,6 +577,50 @@ class Gemma4ModelArchConfigConvertor(ModelArchConfigConvertorBase):
         return max(head_dim, global_head_dim) or super().get_head_size()
 
 
+class MossAudioModelArchConfigConvertor(ModelArchConfigConvertorBase):
+    def _language_config(self) -> PretrainedConfig:
+        return self.hf_config.language_config
+
+    def get_num_hidden_layers(self) -> int:
+        return getattr(self._language_config(), "num_hidden_layers", 0)
+
+    def get_total_num_attention_heads(self) -> int:
+        return getattr(self._language_config(), "num_attention_heads", 0)
+
+    def get_vocab_size(self) -> int:
+        return getattr(self._language_config(), "vocab_size", 0)
+
+    def get_hidden_size(self) -> int:
+        return getattr(self._language_config(), "hidden_size", 0)
+
+    def get_head_size(self) -> int:
+        head_dim = getattr(self._language_config(), "head_dim", None)
+        if head_dim is not None:
+            return head_dim
+        total_num_attention_heads = self.get_total_num_attention_heads()
+        if total_num_attention_heads == 0:
+            return 0
+        return self.get_hidden_size() // total_num_attention_heads
+
+    def get_total_num_kv_heads(self) -> int:
+        return getattr(
+            self._language_config(),
+            "num_key_value_heads",
+            self.get_total_num_attention_heads(),
+        )
+
+    def derive_max_model_len_and_key(self) -> tuple[float, str | None]:
+        language_config = self._language_config()
+        max_position_embeddings = getattr(
+            language_config,
+            "max_position_embeddings",
+            None,
+        )
+        if max_position_embeddings is None:
+            return super().derive_max_model_len_and_key()
+        return max_position_embeddings, "language_config.max_position_embeddings"
+
+
 # hf_config.model_type -> convertor class
 MODEL_ARCH_CONFIG_CONVERTORS = {
     "cohere_asr": CohereAsrModelArchConfigConvertor,
@@ -601,6 +645,7 @@ MODEL_ARCH_CONFIG_CONVERTORS = {
     "mimo_v2_flash": MimoV2ModelArchConfigConvertor,
     "mimo_v2_mtp": MimoV2MTPModelArchConfigConvertor,
     "mimo_v2_omni_mtp": MimoV2MTPModelArchConfigConvertor,
+    "moss_audio": MossAudioModelArchConfigConvertor,
     "mpt": MPTModelArchConfigConvertor,
     "nemotron-nas": NemotronNasModelArchConfigConvertor,
     "pangu_ultra_moe_mtp": PanguUltraMoeMTPModelArchConfigConvertor,