[MRV2][BugFix] Fix KV connector handling in spec decode case (#43719)

Signed-off-by: Nick Hill <nickhill123@gmail.com>
Co-authored-by: Wentao Ye <44945378+yewentao256@users.noreply.github.com>
(cherry picked from commit 8c94938cfb)

.buildkite/test_areas/kernels.yaml

@@ -38,6 +38,28 @@ steps:
   commands:
     - pytest -v -s kernels/core/test_minimax_reduce_rms.py
 
+- label: Deepseek V4 Kernel Test (H100)
+  key: deepseek-v4-kernel-test-h100
+  timeout_in_minutes: 15
+  device: h100
+  source_file_dependencies:
+  - 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
+  commands:
+    - pytest -v -s kernels/test_fused_deepseek_v4_*.py
+
+- label: Deepseek V4 Kernel Test (B200)
+  key: deepseek-v4-kernel-test-b200
+  timeout_in_minutes: 15
+  device: b200-k8s
+  source_file_dependencies:
+  - 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
+  commands:
+    - pytest -v -s kernels/test_fused_deepseek_v4_*.py
+
 - label: Kernels Attention Test %N
   key: kernels-attention-test
   timeout_in_minutes: 35

csrc/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu

@@ -122,18 +122,45 @@ __device__ __forceinline__ float warpSum(float val) {
 // Per-slot inner pipeline
 // ────────────────────────────────────────────────────────────────────────────
 // Shared by both kernel variants: 1 CTA per (token, head) pair vs. 1 CTA per
-// token
-template <typename scalar_t_in>
+// token.  Templated on `kNumHeadsQPadded` so the KV-sentinel comparison and
+// q_out stride fold to compile-time constants.
+//
+// Slot layout (per token):
+//   slot < num_heads_q                          → live-Q   (RMSNorm + RoPE,
+//                                                           read q_in →
+//                                                           write q_out)
+//   num_heads_q <= slot < kNumHeadsQPadded      → pad-Q    (zero-fill q_out;
+//                                                           v0/v1 unused)
+//   slot == kNumHeadsQPadded                    → KV       (RoPE + UE8M0 quant
+//                                                           + paged-cache
+//                                                           insert)
+template <typename scalar_t_in, int kNumHeadsQPadded>
 __device__ __forceinline__ void processDeepseekV4Slot(
     uint4 v0, uint4 v1, int const tokenIdx, int const slotIdx,
     int const dim_base, int const laneId, int const num_heads_q,
-    float const eps, scalar_t_in* __restrict__ q_inout,
+    float const eps, scalar_t_in* __restrict__ q_out,
     uint8_t* __restrict__ k_cache, int64_t const* __restrict__ slot_mapping,
     int64_t const* __restrict__ position_ids,
     float const* __restrict__ cos_sin_cache, int const cache_block_size,
     int const kv_block_stride) {
   using Converter = vllm::_typeConvert<scalar_t_in>;
-  bool const isKV = (slotIdx == num_heads_q);
+  bool const isKV = (slotIdx == kNumHeadsQPadded);
+  bool const isPadQ = !isKV && (slotIdx >= num_heads_q);
+
+  // ── Pad-Q branch: write 32 B of zeros and exit. ─────────────────────────
+  // FlashMLA reads these slots; bf16 +0.0 is bit pattern 0x0000, so a uint4
+  // zero literal is correct.  Matches the live-Q branch's vectorized store.
+  if (isPadQ) {
+    scalar_t_in* dst =
+        q_out +
+        (static_cast<int64_t>(tokenIdx) * kNumHeadsQPadded + slotIdx) *
+            kHeadDim +
+        dim_base;
+    uint4 const zero4 = {0u, 0u, 0u, 0u};
+    *reinterpret_cast<uint4*>(dst) = zero4;
+    *reinterpret_cast<uint4*>(dst + 8) = zero4;
+    return;
+  }
 
   // ── Decode the bf16 → 16 fp32 registers ─────────────────────────────
   float elements[kElemsPerLane];
@@ -207,7 +234,7 @@ __device__ __forceinline__ void processDeepseekV4Slot(
   // triggering and per-iteration buffer rotation.
   // ═══════════════════════════════════════════════════════════════════
   if (!isKV) {
-    // ── Q: cast back to bf16 and store. ────────────────────────────
+    // ── Live-Q: cast back to bf16 and store into the padded q_out. ─────
     uint4 out0, out1;
     typename Converter::packed_hip_type* po0 =
         reinterpret_cast<typename Converter::packed_hip_type*>(&out0);
@@ -224,8 +251,9 @@ __device__ __forceinline__ void processDeepseekV4Slot(
           make_float2(elements[8 + 2 * i], elements[8 + 2 * i + 1]));
     }
     scalar_t_in* dst =
-        q_inout +
-        (static_cast<int64_t>(tokenIdx) * num_heads_q + slotIdx) * kHeadDim +
+        q_out +
+        (static_cast<int64_t>(tokenIdx) * kNumHeadsQPadded + slotIdx) *
+            kHeadDim +
         dim_base;
     *reinterpret_cast<uint4*>(dst) = out0;
     *reinterpret_cast<uint4*>(dst + 8) = out1;
@@ -313,20 +341,32 @@ __device__ __forceinline__ void processDeepseekV4Slot(
 // Kernel
 // ────────────────────────────────────────────────────────────────────────────
 //
-// Grid: 1D, gridDim.x = ceil(num_tokens_full * (num_heads_q + 1) /
+// Grid: 1D, gridDim.x = ceil(num_tokens_full * (kNumHeadsQPadded + 1) /
 // warps_per_block) Block: blockDim.x = 256 threads (8 warps per block) Each
-// warp handles one (token, head_slot) pair. head_slot < num_heads_q          →
-// Q branch (RMSNorm + RoPE, in place) head_slot == num_heads_q         → KV
-// branch (RoPE + UE8M0 quant + insert)
+// warp handles one (token, head_slot) pair.
+//   slot < num_heads_q                              → live-Q branch
+//                                                     (RMSNorm + RoPE,
+//                                                      read q_in → write q_out)
+//   num_heads_q <= slot < kNumHeadsQPadded          → pad-Q branch
+//                                                     (zero-fill q_out)
+//   slot == kNumHeadsQPadded                        → KV branch
+//                                                     (RoPE + UE8M0 quant +
+//                                                      paged-cache insert)
+//
+// `kNumHeadsQPadded` is a template parameter (compile-time constant) so the
+// divisions in the grid math and the KV-sentinel comparison fold to fast
+// constant operations.  The launch wrapper dispatches the runtime value to
+// the matching instantiation.
 //
 // With DP padding, q/kv/position_ids can have more rows than slot_mapping.
-// The Q branch covers all `num_tokens_full` rows (downstream attention uses
-// them).  The KV branch only inserts the first `num_tokens_insert` tokens
-// (= slot_mapping length) into the paged cache.
+// The live-Q and pad-Q branches cover all `num_tokens_full` rows (downstream
+// attention uses them).  The KV branch only inserts the first
+// `num_tokens_insert` tokens (= slot_mapping length) into the paged cache.
 //
-template <typename scalar_t_in>
+template <typename scalar_t_in, int kNumHeadsQPadded>
 __global__ void fusedDeepseekV4QNormRopeKVRopeQuantInsertKernel(
-    scalar_t_in* __restrict__ q_inout,         // [N, H, 512] bf16, in place
+    scalar_t_in const* __restrict__ q_in,      // [N, num_heads_q,      512]
+    scalar_t_in* __restrict__ q_out,           // [N, kNumHeadsQPadded, 512]
     scalar_t_in const* __restrict__ kv_in,     // [N, 512] bf16
     uint8_t* __restrict__ k_cache,             // [num_blocks, block_stride]
     int64_t const* __restrict__ slot_mapping,  // [num_tokens_insert] i64
@@ -335,7 +375,7 @@ __global__ void fusedDeepseekV4QNormRopeKVRopeQuantInsertKernel(
     float const eps,
     int const num_tokens_full,    // = q.size(0) = kv.size(0)
     int const num_tokens_insert,  // = slot_mapping.size(0), ≤ num_tokens_full
-    int const num_heads_q,        // H
+    int const num_heads_q,        // live Q heads (input layout)
     int const cache_block_size,   // tokens per paged-cache block
     int const kv_block_stride) {  // bytes per paged-cache block
 #if (!defined(__CUDA_ARCH__) || __CUDA_ARCH__ < 800) && !defined(USE_ROCM)
@@ -351,12 +391,13 @@ __global__ void fusedDeepseekV4QNormRopeKVRopeQuantInsertKernel(
     int const laneId = threadIdx.x % 32;
     int const globalWarpIdx = blockIdx.x * warpsPerBlock + warpId;
 
-    int const total_slots_per_token = num_heads_q + 1;
-    int const tokenIdx = globalWarpIdx / total_slots_per_token;
-    int const slotIdx = globalWarpIdx % total_slots_per_token;
+    constexpr int kTotalSlotsPerToken = kNumHeadsQPadded + 1;
+    int const tokenIdx = globalWarpIdx / kTotalSlotsPerToken;
+    int const slotIdx = globalWarpIdx % kTotalSlotsPerToken;
     if (tokenIdx >= num_tokens_full) return;
 
-    bool const isKV = (slotIdx == num_heads_q);
+    bool const isKV = (slotIdx == kNumHeadsQPadded);
+    bool const isPadQ = !isKV && (slotIdx >= num_heads_q);
     // KV branch: skip DP-padded tokens (no slot reserved for them).
     if (isKV && tokenIdx >= num_tokens_insert) return;
 
@@ -371,22 +412,26 @@ __global__ void fusedDeepseekV4QNormRopeKVRopeQuantInsertKernel(
     // Dim range this lane owns within the 512-wide head.
     int const dim_base = laneId * kElemsPerLane;  // in [0, 512) step 16
 
-    // Two 16-byte loads per thread (8 bf16 each).  Use uint4 as the vector
-    // type; the shared per-slot helper bitcasts to scalar_t_in packed pairs.
-    scalar_t_in const* src_ptr;
-    if (isKV) {
-      src_ptr = kv_in + static_cast<int64_t>(tokenIdx) * kHeadDim + dim_base;
-    } else {
-      int64_t const q_row_offset =
-          (static_cast<int64_t>(tokenIdx) * num_heads_q + slotIdx) * kHeadDim +
-          dim_base;
-      src_ptr = q_inout + q_row_offset;
+    // Load only for live-Q and KV slots; pad-Q skips the read (q_in beyond
+    // num_heads_q is out of bounds) and the helper zero-fills its output.
+    uint4 v0, v1;
+    if (!isPadQ) {
+      scalar_t_in const* src_ptr;
+      if (isKV) {
+        src_ptr = kv_in + static_cast<int64_t>(tokenIdx) * kHeadDim + dim_base;
+      } else {
+        int64_t const q_row_offset =
+            (static_cast<int64_t>(tokenIdx) * num_heads_q + slotIdx) *
+                kHeadDim +
+            dim_base;
+        src_ptr = q_in + q_row_offset;
+      }
+      v0 = *reinterpret_cast<uint4 const*>(src_ptr);
+      v1 = *reinterpret_cast<uint4 const*>(src_ptr + 8);
     }
-    uint4 const v0 = *reinterpret_cast<uint4 const*>(src_ptr);
-    uint4 const v1 = *reinterpret_cast<uint4 const*>(src_ptr + 8);
 
-    processDeepseekV4Slot<scalar_t_in>(
-        v0, v1, tokenIdx, slotIdx, dim_base, laneId, num_heads_q, eps, q_inout,
+    processDeepseekV4Slot<scalar_t_in, kNumHeadsQPadded>(
+        v0, v1, tokenIdx, slotIdx, dim_base, laneId, num_heads_q, eps, q_out,
         k_cache, slot_mapping, position_ids, cos_sin_cache, cache_block_size,
         kv_block_stride);
 
@@ -408,9 +453,9 @@ __global__ void fusedDeepseekV4QNormRopeKVRopeQuantInsertKernel(
 // Q branch (RMSNorm + RoPE, in place) head_slot == num_heads_q
 // KV branch (RoPE + UE8M0 quant + insert)
 //
-template <typename scalar_t_in>
+template <typename scalar_t_in, int kNumHeadsQPadded>
 __global__ void fusedDeepseekV4QNormRopeKVRopeQuantInsertKernelReducedGrid(
-    scalar_t_in* __restrict__ q_inout,  // [N, H, 512] bf16, in place
+    scalar_t_in const* __restrict__ q_in, scalar_t_in* __restrict__ q_out,
     scalar_t_in const* __restrict__ kv_in, uint8_t* __restrict__ k_cache,
     int64_t const* __restrict__ slot_mapping,
     int64_t const* __restrict__ position_ids,
@@ -435,25 +480,32 @@ __global__ void fusedDeepseekV4QNormRopeKVRopeQuantInsertKernelReducedGrid(
 #endif
 
     int const dim_base = laneId * kElemsPerLane;  // in [0, 512) step 16
-    int const slot_end =
-        (tokenIdx >= num_tokens_insert) ? num_heads_q : (num_heads_q + 1);
+    // Slot enumeration: live-Q + pad-Q + (KV if this token has a slot).
+    int const slot_end = (tokenIdx >= num_tokens_insert)
+                             ? kNumHeadsQPadded
+                             : (kNumHeadsQPadded + 1);
 
-    auto src_for_slot = [&](int s) -> scalar_t_in const* {
-      if (s == num_heads_q) {
-        return kv_in + static_cast<int64_t>(tokenIdx) * kHeadDim + dim_base;
+    auto load_slot = [&](int s, uint4& va, uint4& vb) {
+      // pad-Q slots skip the load — q_in beyond num_heads_q is OOB.
+      if (s >= num_heads_q && s < kNumHeadsQPadded) return;
+      scalar_t_in const* src;
+      if (s == kNumHeadsQPadded) {
+        src = kv_in + static_cast<int64_t>(tokenIdx) * kHeadDim + dim_base;
+      } else {
+        src = q_in +
+              (static_cast<int64_t>(tokenIdx) * num_heads_q +
+               static_cast<int64_t>(s)) *
+                  kHeadDim +
+              dim_base;
       }
-      return q_inout +
-             (static_cast<int64_t>(tokenIdx) * num_heads_q +
-              static_cast<int64_t>(s)) *
-                 kHeadDim +
-             dim_base;
+      va = *reinterpret_cast<uint4 const*>(src);
+      vb = *reinterpret_cast<uint4 const*>(src + 8);
     };
 
     if (warpId < slot_end) {
       int curr_slot = warpId;
-      scalar_t_in const* src_curr = src_for_slot(curr_slot);
-      uint4 v0_curr = *reinterpret_cast<uint4 const*>(src_curr);
-      uint4 v1_curr = *reinterpret_cast<uint4 const*>(src_curr + 8);
+      uint4 v0_curr, v1_curr;
+      load_slot(curr_slot, v0_curr, v1_curr);
 
       while (curr_slot < slot_end) {
         int const next_slot = curr_slot + warpsPerBlock;
@@ -462,14 +514,12 @@ __global__ void fusedDeepseekV4QNormRopeKVRopeQuantInsertKernelReducedGrid(
         // Prefetch src for the next slot
         uint4 v0_next, v1_next;
         if (has_next) {
-          scalar_t_in const* src_next = src_for_slot(next_slot);
-          v0_next = *reinterpret_cast<uint4 const*>(src_next);
-          v1_next = *reinterpret_cast<uint4 const*>(src_next + 8);
+          load_slot(next_slot, v0_next, v1_next);
         }
 
-        processDeepseekV4Slot<scalar_t_in>(
+        processDeepseekV4Slot<scalar_t_in, kNumHeadsQPadded>(
             v0_curr, v1_curr, tokenIdx, curr_slot, dim_base, laneId,
-            num_heads_q, eps, q_inout, k_cache, slot_mapping, position_ids,
+            num_heads_q, eps, q_out, k_cache, slot_mapping, position_ids,
             cos_sin_cache, cache_block_size, kv_block_stride);
 
         // ── Buffer rotation: hand the prefetched LDGs to the next iter.
@@ -490,10 +540,10 @@ __global__ void fusedDeepseekV4QNormRopeKVRopeQuantInsertKernelReducedGrid(
 // ────────────────────────────────────────────────────────────────────────────
 // Launch wrapper
 // ────────────────────────────────────────────────────────────────────────────
-template <typename scalar_t_in>
-void launchFusedDeepseekV4QNormRopeKVRopeQuantInsert(
-    scalar_t_in* q_inout, scalar_t_in const* kv_in, uint8_t* k_cache,
-    int64_t const* slot_mapping, int64_t const* position_ids,
+template <typename scalar_t_in, int kNumHeadsQPadded>
+static void launchFusedDeepseekV4Templated(
+    scalar_t_in const* q_in, scalar_t_in* q_out, scalar_t_in const* kv_in,
+    uint8_t* k_cache, int64_t const* slot_mapping, int64_t const* position_ids,
     float const* cos_sin_cache, float const eps, int const num_tokens_full,
     int const num_tokens_insert, int const num_heads_q,
     int const cache_block_size, int const kv_block_stride,
@@ -501,7 +551,7 @@ void launchFusedDeepseekV4QNormRopeKVRopeQuantInsert(
   constexpr int kBlockSize = 256;
   constexpr int kWarpsPerBlock = kBlockSize / 32;
   int64_t const total_warps =
-      static_cast<int64_t>(num_tokens_full) * (num_heads_q + 1);
+      static_cast<int64_t>(num_tokens_full) * (kNumHeadsQPadded + 1);
   int const grid =
       static_cast<int>((total_warps + kWarpsPerBlock - 1) / kWarpsPerBlock);
 
@@ -532,46 +582,85 @@ void launchFusedDeepseekV4QNormRopeKVRopeQuantInsert(
 
   if (num_tokens_full < NUM_TOKEN_CUTOFF) {
     cudaLaunchKernelEx(
-        &config, fusedDeepseekV4QNormRopeKVRopeQuantInsertKernel<scalar_t_in>,
-        q_inout, kv_in, k_cache, slot_mapping, position_ids, cos_sin_cache, eps,
-        num_tokens_full, num_tokens_insert, num_heads_q, cache_block_size,
+        &config,
+        fusedDeepseekV4QNormRopeKVRopeQuantInsertKernel<scalar_t_in,
+                                                        kNumHeadsQPadded>,
+        q_in, q_out, kv_in, k_cache, slot_mapping, position_ids, cos_sin_cache,
+        eps, num_tokens_full, num_tokens_insert, num_heads_q, cache_block_size,
         kv_block_stride);
   } else {
     config.gridDim = dim3(num_tokens_full);
     cudaLaunchKernelEx(
         &config,
-        fusedDeepseekV4QNormRopeKVRopeQuantInsertKernelReducedGrid<scalar_t_in>,
-        q_inout, kv_in, k_cache, slot_mapping, position_ids, cos_sin_cache, eps,
-        num_tokens_full, num_tokens_insert, num_heads_q, cache_block_size,
+        fusedDeepseekV4QNormRopeKVRopeQuantInsertKernelReducedGrid<
+            scalar_t_in, kNumHeadsQPadded>,
+        q_in, q_out, kv_in, k_cache, slot_mapping, position_ids, cos_sin_cache,
+        eps, num_tokens_full, num_tokens_insert, num_heads_q, cache_block_size,
         kv_block_stride);
   }
-
 #else
   // ROCm: use standard kernel launch syntax (no PDL/stream serialization)
   // clang-format off
-  fusedDeepseekV4QNormRopeKVRopeQuantInsertKernel<scalar_t_in>
+  fusedDeepseekV4QNormRopeKVRopeQuantInsertKernel<scalar_t_in, kNumHeadsQPadded>
       <<<grid, kBlockSize, 0, stream>>>(
-          q_inout, kv_in, k_cache, slot_mapping, position_ids, cos_sin_cache,
-          eps, num_tokens_full, num_tokens_insert, num_heads_q,
+          q_in, q_out, kv_in, k_cache, slot_mapping, position_ids,
+          cos_sin_cache, eps, num_tokens_full, num_tokens_insert, num_heads_q,
           cache_block_size, kv_block_stride);
 #endif
 }
 
+// Runtime dispatch into one of the precompiled `kNumHeadsQPadded`
+// instantiations.  Supported padded head counts: 8, 16, 32, 64, 128.
+template <typename scalar_t_in>
+void launchFusedDeepseekV4QNormRopeKVRopeQuantInsert(
+    scalar_t_in const* q_in, scalar_t_in* q_out, scalar_t_in const* kv_in,
+    uint8_t* k_cache, int64_t const* slot_mapping,
+    int64_t const* position_ids, float const* cos_sin_cache, float const eps,
+    int const num_tokens_full, int const num_tokens_insert,
+    int const num_heads_q, int const num_heads_q_padded,
+    int const cache_block_size, int const kv_block_stride,
+    cudaStream_t stream) {
+#define DISPATCH(N)                                                         \
+  case N:                                                                   \
+    launchFusedDeepseekV4Templated<scalar_t_in, N>(                         \
+        q_in, q_out, kv_in, k_cache, slot_mapping, position_ids,            \
+        cos_sin_cache, eps, num_tokens_full, num_tokens_insert, num_heads_q, \
+        cache_block_size, kv_block_stride, stream);                         \
+    return;
+
+  switch (num_heads_q_padded) {
+    DISPATCH(8)
+    DISPATCH(16)
+    DISPATCH(32)
+    DISPATCH(64)
+    DISPATCH(128)
+    default:
+      TORCH_CHECK(false,
+                  "fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert: "
+                  "unsupported num_heads_q_padded=",
+                  num_heads_q_padded,
+                  " (compiled instantiations: 8, 16, 32, 64, 128).");
+  }
+#undef DISPATCH
+}
+
 }  // namespace deepseek_v4_fused_ops
 }  // namespace vllm
 
 // ────────────────────────────────────────────────────────────────────────────
 // Torch op wrapper
 // ────────────────────────────────────────────────────────────────────────────
-void fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
-    torch::Tensor& q,                    // [N, H, 512] bf16, in place
+torch::Tensor fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
+    torch::Tensor const& q_in,           // [N, num_heads_q, 512] bf16
     torch::Tensor const& kv,             // [N, 512] bf16 (read-only)
     torch::Tensor& k_cache,              // [num_blocks, block_bytes] uint8
     torch::Tensor const& slot_mapping,   // [N] int64
     torch::Tensor const& position_ids,   // [N] int64
     torch::Tensor const& cos_sin_cache,  // [max_pos, rope_dim] bf16
+    int64_t q_head_padded,               // padded Q head count for output
     double eps, int64_t cache_block_size) {
-  TORCH_CHECK(q.is_cuda() && q.is_contiguous(), "q must be contiguous CUDA");
+  TORCH_CHECK(q_in.is_cuda() && q_in.is_contiguous(),
+              "q_in must be contiguous CUDA");
   TORCH_CHECK(kv.is_cuda() && kv.is_contiguous(), "kv must be contiguous CUDA");
   TORCH_CHECK(k_cache.is_cuda(), "k_cache must be CUDA");
   TORCH_CHECK(slot_mapping.is_cuda() && slot_mapping.dtype() == torch::kInt64,
@@ -579,9 +668,12 @@ void fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
   TORCH_CHECK(position_ids.is_cuda() && position_ids.dtype() == torch::kInt64,
               "position_ids must be int64 CUDA");
   TORCH_CHECK(cos_sin_cache.is_cuda(), "cos_sin_cache must be CUDA");
-  TORCH_CHECK(q.dim() == 3 && q.size(2) == 512, "q shape [N, H, 512]");
+  TORCH_CHECK(q_in.dim() == 3 && q_in.size(2) == 512,
+              "q_in shape [N, num_heads_q, 512]");
   TORCH_CHECK(kv.dim() == 2 && kv.size(1) == 512, "kv shape [N, 512]");
-  TORCH_CHECK(q.dtype() == kv.dtype(), "q and kv dtype must match");
+  TORCH_CHECK(q_in.dtype() == kv.dtype(), "q_in and kv dtype must match");
+  TORCH_CHECK(q_head_padded >= q_in.size(1),
+              "q_head_padded must be >= q_in.size(1) (num_heads_q)");
   TORCH_CHECK(k_cache.dtype() == torch::kUInt8, "k_cache must be uint8");
   TORCH_CHECK(cos_sin_cache.dim() == 2 && cos_sin_cache.size(1) == 64,
               "cos_sin_cache shape [max_pos, 64]");
@@ -591,32 +683,41 @@ void fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
   // With DP padding, slot_mapping can be shorter than q/kv/positions.
   // Q-norm+RoPE runs on all q.size(0) rows (downstream attention uses them);
   // KV quant+insert runs only on the first slot_mapping.size(0) rows.
-  int const num_tokens_full = static_cast<int>(q.size(0));
+  int const num_tokens_full = static_cast<int>(q_in.size(0));
   int const num_tokens_insert = static_cast<int>(slot_mapping.size(0));
   TORCH_CHECK(static_cast<int>(kv.size(0)) == num_tokens_full &&
                   static_cast<int>(position_ids.size(0)) == num_tokens_full,
               "q/kv/position_ids row counts must match");
   TORCH_CHECK(num_tokens_insert <= num_tokens_full,
               "slot_mapping must not exceed q row count");
-  int const num_heads_q = static_cast<int>(q.size(1));
+  int const num_heads_q = static_cast<int>(q_in.size(1));
+  int const num_heads_q_padded = static_cast<int>(q_head_padded);
   int const cache_block_size_i = static_cast<int>(cache_block_size);
   int const kv_block_stride = static_cast<int>(k_cache.stride(0));
 
-  at::cuda::OptionalCUDAGuard device_guard(device_of(q));
+  at::cuda::OptionalCUDAGuard device_guard(device_of(q_in));
   auto stream = at::cuda::getCurrentCUDAStream();
 
+  // Allocate the padded q output.  The kernel writes every element (live
+  // region gets RMSNorm+RoPE; pad region gets zeros), so `empty` is safe.
+  torch::Tensor q_out = torch::empty(
+      {q_in.size(0), q_head_padded, q_in.size(2)}, q_in.options());
+
   VLLM_DISPATCH_HALF_TYPES(
-      q.scalar_type(), "fused_deepseek_v4_qnorm_rope_kv_insert", [&] {
+      q_in.scalar_type(), "fused_deepseek_v4_qnorm_rope_kv_insert", [&] {
         using qkv_scalar_t = scalar_t;
         vllm::deepseek_v4_fused_ops::
             launchFusedDeepseekV4QNormRopeKVRopeQuantInsert<qkv_scalar_t>(
-                reinterpret_cast<qkv_scalar_t*>(q.data_ptr()),
+                reinterpret_cast<qkv_scalar_t const*>(q_in.data_ptr()),
+                reinterpret_cast<qkv_scalar_t*>(q_out.data_ptr()),
                 reinterpret_cast<qkv_scalar_t const*>(kv.data_ptr()),
                 reinterpret_cast<uint8_t*>(k_cache.data_ptr()),
                 reinterpret_cast<int64_t const*>(slot_mapping.data_ptr()),
                 reinterpret_cast<int64_t const*>(position_ids.data_ptr()),
                 cos_sin_cache.data_ptr<float>(), static_cast<float>(eps),
                 num_tokens_full, num_tokens_insert, num_heads_q,
-                cache_block_size_i, kv_block_stride, stream);
+                num_heads_q_padded, cache_block_size_i, kv_block_stride,
+                stream);
       });
-}
+  return q_out;
+}

csrc/ops.h

@@ -70,10 +70,11 @@ void rms_norm(torch::Tensor& out, torch::Tensor& input, torch::Tensor& weight,
 void fused_add_rms_norm(torch::Tensor& input, torch::Tensor& residual,
                         torch::Tensor& weight, double epsilon);
 
-void fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
-    torch::Tensor& q, torch::Tensor const& kv, torch::Tensor& k_cache,
+torch::Tensor fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
+    torch::Tensor const& q_in, torch::Tensor const& kv, torch::Tensor& k_cache,
     torch::Tensor const& slot_mapping, torch::Tensor const& position_ids,
-    torch::Tensor const& cos_sin_cache, double eps, int64_t cache_block_size);
+    torch::Tensor const& cos_sin_cache, int64_t q_head_padded, double eps,
+    int64_t cache_block_size);
 
 void apply_repetition_penalties_(torch::Tensor& logits,
                                  const torch::Tensor& prompt_mask,

csrc/torch_bindings.cpp

@@ -99,9 +99,9 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
   // kernel launch.
   ops.def(
       "fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert("
-      "Tensor! q, Tensor kv, Tensor! k_cache, "
+      "Tensor q_in, Tensor kv, Tensor! k_cache, "
       "Tensor slot_mapping, Tensor position_ids, Tensor cos_sin_cache, "
-      "float eps, int cache_block_size) -> ()");
+      "int q_head_padded, float eps, int cache_block_size) -> Tensor");
   ops.impl("fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert", torch::kCUDA,
            &fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert);
 

docker/Dockerfile

@@ -220,7 +220,8 @@ COPY use_existing_torch.py use_existing_torch.py
 COPY pyproject.toml pyproject.toml
 RUN --mount=type=cache,target=/opt/uv/cache \
     if [ "$(echo $CUDA_VERSION | cut -d. -f1)" = "12" ]; then \
-        sed -i 's/^nvidia-cutlass-dsl\[cu13\]>=/nvidia-cutlass-dsl>=/' requirements/cuda.txt; \
+        sed -i 's/^nvidia-cutlass-dsl\[cu13\]/nvidia-cutlass-dsl/' requirements/cuda.txt; \
+        sed -i 's/^humming-kernels\[cu13\]/humming-kernels[cu12]/' requirements/cuda.txt; \
     fi \
     && if [ "${PYTORCH_NIGHTLY}" = "1" ]; then \
         echo "Installing torch nightly..." \
@@ -746,7 +747,8 @@ COPY requirements/common.txt /tmp/common.txt
 COPY requirements/cuda.txt /tmp/requirements-cuda.txt
 RUN --mount=type=cache,target=/opt/uv/cache \
     if [ "$(echo $CUDA_VERSION | cut -d. -f1)" = "12" ]; then \
-        sed -i 's/^nvidia-cutlass-dsl\[cu13\]>=/nvidia-cutlass-dsl>=/' /tmp/requirements-cuda.txt; \
+        sed -i 's/^nvidia-cutlass-dsl\[cu13\]/nvidia-cutlass-dsl/' /tmp/requirements-cuda.txt; \
+        sed -i 's/^humming-kernels\[cu13\]/humming-kernels[cu12]/' /tmp/requirements-cuda.txt; \
     fi && \
     uv pip install --system -r /tmp/requirements-cuda.txt \
         --extra-index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.') && \

requirements/cuda.txt

@@ -25,4 +25,7 @@ nvidia-cutlass-dsl[cu13]==4.5.0
 quack-kernels>=0.3.3
 
 # Tokenspeed_MLA for faster mla with spec decode
-tokenspeed-mla==0.1.2
+tokenspeed-mla==0.1.2
+
+# Humming kernels for quantization gemm
+humming-kernels[cu13]==0.1.2

requirements/kv_connectors.txt

@@ -1,3 +1,6 @@
 lmcache >= 0.3.9
+# CuPy 14.1.0 imports pytest from cupy.testing._random. Use <14.1.0
+# until a fixed newer release is verified for runtime images.
+cupy-cuda13x < 14.1.0
 nixl >= 1.1.0 # Required for disaggregated prefill
 mooncake-transfer-engine >= 0.3.8

setup.py

@@ -1017,6 +1017,8 @@ def get_requirements() -> list[str]:
             if "nvidia-cutlass-dsl[cu13]" in req and cuda_major == "12":
                 # [cu13] extra is the default; strip it on CUDA 12 builds.
                 req = req.replace("nvidia-cutlass-dsl[cu13]", "nvidia-cutlass-dsl")
+            if "humming-kernels[cu13]" in req and cuda_major == "12":
+                req = req.replace("humming-kernels[cu13]", "humming-kernels[cu12]")
             modified_requirements.append(req)
         requirements = modified_requirements
     elif _is_hip():

tests/entrypoints/test_api_server_process_manager.py

@@ -8,8 +8,14 @@ import time
 from unittest.mock import patch
 
 import pytest
+import zmq
 
-from vllm.v1.utils import APIServerProcessManager, wait_for_completion_or_failure
+from vllm.utils.network_utils import make_zmq_socket, split_zmq_path
+from vllm.v1.utils import (
+    APIServerProcessManager,
+    get_engine_client_zmq_addr,
+    wait_for_completion_or_failure,
+)
 
 # Global variables to control worker behavior
 WORKER_RUNTIME_SECONDS = 0.5
@@ -23,6 +29,39 @@ def mock_run_api_server_worker(listen_address, sock, args, client_config=None):
     print("Mock worker completed successfully")
 
 
+# Module-level stub for the gather_actual_addresses test. Must be
+# importable by `multiprocessing.spawn` (no closures, no nesting).
+def defer_addresses_stub_worker(listen_address, sock, args, client_config):
+    """Bind ROUTER/PULL with a kernel-assigned port, report the actual
+    endpoints back via the pipe, then exit."""
+    ctx = zmq.Context()
+    try:
+        in_sock = make_zmq_socket(
+            ctx, client_config["input_address"], zmq.ROUTER, bind=True
+        )
+        out_sock = make_zmq_socket(
+            ctx, client_config["output_address"], zmq.PULL, bind=True
+        )
+        try:
+            pipe = client_config["actual_address_pipe"]
+            try:
+                pipe.send(
+                    {
+                        "input_address": in_sock.getsockopt(zmq.LAST_ENDPOINT).decode(),
+                        "output_address": out_sock.getsockopt(
+                            zmq.LAST_ENDPOINT
+                        ).decode(),
+                    }
+                )
+            finally:
+                pipe.close()
+        finally:
+            in_sock.close(linger=0)
+            out_sock.close(linger=0)
+    finally:
+        ctx.term()
+
+
 @pytest.fixture
 def api_server_args():
     """Fixture to provide arguments for APIServerProcessManager."""
@@ -268,3 +307,92 @@ def test_external_process_monitoring(api_server_args):
         manager.shutdown()
         mock_coordinator.shutdown()
         time.sleep(0.2)
+
+
+@pytest.mark.timeout(60)
+def test_gather_actual_addresses_end_to_end():
+    """Each child binds ROUTER/PULL with a kernel-picked port and reports
+    the bound endpoints back via its per-child pipe; the manager surfaces
+    them via :py:meth:`gather_actual_addresses`."""
+    host = "127.0.0.1"
+    num_servers = 4
+
+    placeholder_inputs = [
+        get_engine_client_zmq_addr(local_only=False, host=host)
+        for _ in range(num_servers)
+    ]
+    placeholder_outputs = [
+        get_engine_client_zmq_addr(local_only=False, host=host)
+        for _ in range(num_servers)
+    ]
+    for addr in placeholder_inputs + placeholder_outputs:
+        assert addr == f"tcp://{host}:0", addr
+
+    sock = socket.socket()
+    manager = APIServerProcessManager(
+        listen_address=f"tcp://{host}:0",
+        sock=sock,
+        args="test_args",
+        num_servers=num_servers,
+        input_addresses=placeholder_inputs,
+        output_addresses=placeholder_outputs,
+        target_server_fn=defer_addresses_stub_worker,
+    )
+
+    try:
+        assert len(manager.processes) == num_servers
+        actual_inputs, actual_outputs = manager.gather_actual_addresses(timeout=15.0)
+    finally:
+        manager.shutdown()
+        time.sleep(0.2)
+        sock.close()
+
+    assert len(actual_inputs) == num_servers
+    assert len(actual_outputs) == num_servers
+
+    for addr in actual_inputs + actual_outputs:
+        scheme, parsed_host, port = split_zmq_path(addr)
+        assert scheme == "tcp", addr
+        assert parsed_host == host, addr
+        assert port and int(port) > 0, addr
+
+    all_addrs = actual_inputs + actual_outputs
+    assert len(set(all_addrs)) == len(all_addrs), all_addrs
+
+
+@pytest.mark.timeout(30)
+def test_gather_actual_addresses_child_crash_before_report():
+    """A child that exits before sending its endpoints must surface a
+    clear ``RuntimeError`` rather than hang or return ``None`` slots."""
+    host = "127.0.0.1"
+    num_servers = 2
+    placeholder_inputs = [
+        get_engine_client_zmq_addr(local_only=False, host=host)
+        for _ in range(num_servers)
+    ]
+    placeholder_outputs = [
+        get_engine_client_zmq_addr(local_only=False, host=host)
+        for _ in range(num_servers)
+    ]
+
+    sock = socket.socket()
+    manager = APIServerProcessManager(
+        listen_address=f"tcp://{host}:0",
+        sock=sock,
+        args="test_args",
+        num_servers=num_servers,
+        input_addresses=placeholder_inputs,
+        output_addresses=placeholder_outputs,
+        # mock_run_api_server_worker exits without touching
+        # ``actual_address_pipe`` — simulates a child that dies before
+        # reporting its bound addresses.
+        target_server_fn=mock_run_api_server_worker,
+    )
+    try:
+        # Sentinel-first vs pipe-EOF-first both produce "reporting".
+        with pytest.raises(RuntimeError, match="reporting"):
+            manager.gather_actual_addresses(timeout=10.0)
+    finally:
+        manager.shutdown()
+        time.sleep(0.2)
+        sock.close()

tests/kernels/mamba/test_gdn_prefill_cutedsl.py

@@ -0,0 +1,199 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import math
+
+import pytest
+import torch
+import torch.nn.functional as F
+
+from vllm.platforms import current_platform
+
+if not (
+    current_platform.is_cuda() and current_platform.is_device_capability_family(100)
+):
+    pytest.skip(
+        reason="GDN CuteDSL prefill requires CUDA SM10x.",
+        allow_module_level=True,
+    )
+
+from vllm.model_executor.layers.fla.ops import (  # noqa: E402
+    chunk_gated_delta_rule,
+)
+from vllm.model_executor.layers.fla.ops.index import (  # noqa: E402
+    prepare_chunk_indices,
+    prepare_chunk_offsets,
+)
+from vllm.model_executor.layers.mamba.ops.gdn_chunk_cutedsl import (  # noqa: E402
+    chunk_gated_delta_rule_cutedsl,
+    prepare_metadata_cutedsl,
+)
+
+
+@pytest.mark.parametrize("num_seqs", [1, 5, 257])
+@pytest.mark.parametrize("state_dtype", [torch.bfloat16, torch.float32])
+def test_gdn_chunk_cutedsl_correctness(num_seqs: int, state_dtype: torch.dtype):
+    seq_lens = torch.randint(
+        1,
+        130,
+        (num_seqs,),
+        dtype=torch.int32,
+    )
+    cu_seqlens = torch.zeros(num_seqs + 1, device="cuda", dtype=torch.int32)
+    cu_seqlens[1:] = seq_lens.to(device="cuda").cumsum(0)
+    total_tokens = int(cu_seqlens[-1].item())
+
+    num_k_heads = 4
+    num_v_heads = 8
+    head_k_dim = 128
+    head_v_dim = 128
+    dtype = torch.bfloat16
+
+    q = torch.randn(
+        1,
+        total_tokens,
+        num_k_heads,
+        head_k_dim,
+        device="cuda",
+        dtype=dtype,
+    )
+    k = torch.randn_like(q)
+    v = torch.randn(
+        1,
+        total_tokens,
+        num_v_heads,
+        head_v_dim,
+        device="cuda",
+        dtype=dtype,
+    )
+    q = F.normalize(q.float(), p=2, dim=-1).to(dtype)
+    k = F.normalize(k.float(), p=2, dim=-1).to(dtype)
+    a = torch.randn(
+        1,
+        total_tokens,
+        num_v_heads,
+        device="cuda",
+        dtype=dtype,
+    )
+    b = torch.randn(
+        1,
+        total_tokens,
+        num_v_heads,
+        device="cuda",
+        dtype=dtype,
+    )
+    # Match upstream FLA GatedDeltaNet synthetic initialization:
+    # https://github.com/fla-org/flash-linear-attention/blob/main/fla/layers/gated_deltanet.py
+    A = torch.empty(num_v_heads, device="cuda", dtype=torch.float32).uniform_(0, 16)
+    A_log = torch.log(A)
+    dt = torch.exp(
+        torch.rand(num_v_heads, device="cuda", dtype=torch.float32)
+        * (math.log(0.1) - math.log(0.001))
+        + math.log(0.001)
+    )
+    dt = torch.clamp(dt, min=1e-4)
+    dt_bias = dt + torch.log(-torch.expm1(-dt))
+    g = -A_log.exp().view(1, 1, num_v_heads) * F.softplus(
+        a.float() + dt_bias.view(1, 1, num_v_heads)
+    )
+    beta = torch.sigmoid(b.float())
+    initial_state = (
+        torch.randn(
+            num_seqs,
+            num_v_heads,
+            head_v_dim,
+            head_k_dim,
+            device="cuda",
+            dtype=state_dtype,
+        )
+        * 0.05
+    )
+
+    # check metadata kernel
+    chunk_indices, chunk_offsets = prepare_metadata_cutedsl(cu_seqlens, total_tokens)
+    torch.accelerator.synchronize()
+
+    expected_indices = prepare_chunk_indices(cu_seqlens, 64)
+    expected_offsets = prepare_chunk_offsets(cu_seqlens, 64)
+    total_chunks = int(expected_offsets[-1].item())
+
+    torch.testing.assert_close(chunk_offsets, expected_offsets.to(torch.int32))
+    torch.testing.assert_close(
+        chunk_indices[:total_chunks],
+        expected_indices,
+    )
+
+    ref_o, ref_state = chunk_gated_delta_rule(
+        q=q,
+        k=k,
+        v=v,
+        g=g,
+        beta=beta,
+        initial_state=initial_state,
+        output_final_state=True,
+        cu_seqlens=cu_seqlens,
+        use_qk_l2norm_in_kernel=False,
+    )
+    actual_core_attn_out = torch.empty(
+        total_tokens,
+        num_v_heads,
+        head_v_dim,
+        device="cuda",
+        dtype=dtype,
+    )
+    actual_o, actual_state = chunk_gated_delta_rule_cutedsl(
+        q=q,
+        k=k,
+        v=v,
+        g=g,
+        beta=beta,
+        initial_state=initial_state,
+        cu_seqlens=cu_seqlens,
+        chunk_indices=chunk_indices,
+        chunk_offsets=chunk_offsets,
+        core_attn_out=actual_core_attn_out,
+    )
+    torch.accelerator.synchronize()
+
+    # check main kernel
+    o_error = (actual_o.float() - ref_o.float()).abs()
+    state_error = (
+        actual_state.float() - ref_state.to(actual_state.dtype).float()
+    ).abs()
+    assert o_error.max().item() < 2e-3
+    assert o_error.mean().item() < 6e-5
+    assert state_error.max().item() < 2e-2
+    assert state_error.mean().item() < 6e-4
+    core_attn_out_error = (
+        actual_core_attn_out.float() - actual_o.squeeze(0).float()
+    ).abs()
+    assert core_attn_out_error.max().item() == 0
+
+    # check main kernel when core_attn_out is not passed
+    no_buffer_o, no_buffer_state = chunk_gated_delta_rule_cutedsl(
+        q=q,
+        k=k,
+        v=v,
+        g=g,
+        beta=beta,
+        initial_state=initial_state,
+        cu_seqlens=cu_seqlens,
+        chunk_indices=chunk_indices,
+        chunk_offsets=chunk_offsets,
+    )
+    torch.accelerator.synchronize()
+
+    no_buffer_o_error = (no_buffer_o.float() - ref_o.float()).abs()
+    no_buffer_state_error = (
+        no_buffer_state.float() - ref_state.to(no_buffer_state.dtype).float()
+    ).abs()
+    buffer_o_error = (no_buffer_o.float() - actual_o.float()).abs()
+    buffer_state_error = (
+        no_buffer_state.float() - actual_state.to(no_buffer_state.dtype).float()
+    ).abs()
+    assert no_buffer_o_error.max().item() < 2e-3
+    assert no_buffer_o_error.mean().item() < 6e-5
+    assert no_buffer_state_error.max().item() < 2e-2
+    assert no_buffer_state_error.mean().item() < 6e-4
+    assert buffer_o_error.max().item() == 0
+    assert buffer_state_error.max().item() == 0

tests/kernels/test_fused_deepseek_v4_qnorm_rope_kv_insert.py

@@ -120,9 +120,19 @@ pytestmark = pytest.mark.skipif(
 )
 
 
-def _call_fused(q, kv, k_cache, slot_mapping, positions, cos_sin_cache, eps, bs):
-    torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
-        q, kv, k_cache, slot_mapping, positions, cos_sin_cache, eps, bs
+def _call_fused(
+    q_in, q_head_padded, kv, k_cache, slot_mapping, positions, cos_sin_cache, eps, bs
+):
+    return torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
+        q_in,
+        kv,
+        k_cache,
+        slot_mapping,
+        positions,
+        cos_sin_cache,
+        q_head_padded,
+        eps,
+        bs,
     )
 
 
@@ -130,8 +140,23 @@ def _call_fused(q, kv, k_cache, slot_mapping, positions, cos_sin_cache, eps, bs)
 
 
 @pytest.mark.parametrize("num_tokens", [1, 4, 17, 64, 2048])
-@pytest.mark.parametrize("n_heads", [8, 64])
-def test_q_path_matches_reference(num_tokens: int, n_heads: int):
+@pytest.mark.parametrize(
+    "n_heads,padded_heads",
+    [
+        # Each supported padded_heads instantiation: padded (n_heads <
+        # padded_heads) and unpadded (n_heads == padded_heads).
+        (1, 8),
+        (8, 8),
+        (8, 16),
+        (16, 16),
+        (16, 32),
+        (32, 32),
+        (8, 64),
+        (64, 64),
+        (64, 128),
+    ],
+)
+def test_q_path_matches_reference(num_tokens: int, n_heads: int, padded_heads: int):
     torch.manual_seed(0)
     device = "cuda"
     dtype = torch.bfloat16
@@ -156,10 +181,16 @@ def test_q_path_matches_reference(num_tokens: int, n_heads: int):
         num_blocks, bs, HEAD_BYTES, dtype=torch.uint8, device=device
     ).view(num_blocks, -1)
     slot_mapping = torch.full((num_tokens,), -1, dtype=torch.int64, device=device)
-    q_fused = q.clone()
-    _call_fused(q_fused, kv, k_cache, slot_mapping, positions, cos_sin_cache, eps, bs)
+    q_out = _call_fused(
+        q, padded_heads, kv, k_cache, slot_mapping, positions, cos_sin_cache, eps, bs
+    )
 
-    torch.testing.assert_close(q_fused, q_ref, rtol=1e-2, atol=1e-2)
+    torch.testing.assert_close(q_out[:, :n_heads], q_ref, rtol=1e-2, atol=1e-2)
+    if n_heads < padded_heads:
+        pad_region = q_out[:, n_heads:padded_heads]
+        assert pad_region.abs().max().item() == 0.0, (
+            "padded head slots must be exact zero"
+        )
 
 
 # ── Test 2: KV path round-trip byte/value parity ─────────────────────────────
@@ -201,11 +232,12 @@ def test_kv_path_matches_reference(num_tokens: int, block_size: int):
         kv_ref, k_cache_ref, slot_mapping, block_size=block_size
     )
 
-    # ── Fused path (dummy q, single head) ──────────────────────────────────
+    # ── Fused path (dummy q, padded to FlashMLA's min head count 64) ───────
     k_cache_fused = torch.zeros_like(k_cache_ref)
     q_dummy = torch.zeros(num_tokens, 1, HEAD_DIM, dtype=dtype, device=device)
-    _call_fused(
+    _ = _call_fused(
         q_dummy,
+        64,
         kv,
         k_cache_fused,
         slot_mapping,
@@ -298,8 +330,9 @@ def test_kv_path_with_dp_padding(num_tokens: int, pad: int, block_size: int):
     # Fused: pass full-sized q/kv/positions, shorter slot_mapping.
     q_dummy = torch.zeros(total, 1, HEAD_DIM, dtype=dtype, device=device)
     k_cache_fused = torch.zeros_like(k_cache_ref)
-    _call_fused(
+    _ = _call_fused(
         q_dummy,
+        64,
         kv,
         k_cache_fused,
         slot_mapping,
@@ -316,9 +349,26 @@ def test_kv_path_with_dp_padding(num_tokens: int, pad: int, block_size: int):
 
 
 @pytest.mark.parametrize("num_tokens", [1, 4, 17, 2048])
-@pytest.mark.parametrize("n_heads", [8, 64])
+@pytest.mark.parametrize(
+    "n_heads,padded_heads",
+    [
+        # Each supported padded_heads instantiation: padded (n_heads <
+        # padded_heads) and unpadded (n_heads == padded_heads).
+        (1, 8),
+        (8, 8),
+        (8, 16),
+        (16, 16),
+        (16, 32),
+        (32, 32),
+        (8, 64),
+        (64, 64),
+        (64, 128),
+    ],
+)
 @pytest.mark.parametrize("block_size", [16, 64])
-def test_combined_q_and_kv(num_tokens: int, n_heads: int, block_size: int):
+def test_combined_q_and_kv(
+    num_tokens: int, n_heads: int, padded_heads: int, block_size: int
+):
     torch.manual_seed(2)
     device = "cuda"
     dtype = torch.bfloat16
@@ -345,10 +395,10 @@ def test_combined_q_and_kv(num_tokens: int, n_heads: int, block_size: int):
     )
 
     # Fused single call.
-    q_fused = q.clone()
     k_cache_fused = torch.zeros_like(k_cache_ref)
-    _call_fused(
-        q_fused,
+    q_out = _call_fused(
+        q,
+        padded_heads,
         kv,
         k_cache_fused,
         slot_mapping,
@@ -358,5 +408,10 @@ def test_combined_q_and_kv(num_tokens: int, n_heads: int, block_size: int):
         block_size,
     )
 
-    torch.testing.assert_close(q_fused, q_ref, rtol=1e-2, atol=1e-2)
+    torch.testing.assert_close(q_out[:, :n_heads], q_ref, rtol=1e-2, atol=1e-2)
+    if n_heads < padded_heads:
+        pad_region = q_out[:, n_heads:padded_heads]
+        assert pad_region.abs().max().item() == 0.0, (
+            "padded head slots must be exact zero"
+        )
     torch.testing.assert_close(k_cache_fused, k_cache_ref, rtol=0, atol=0)

tests/kernels/test_fused_inv_rope_fp8_quant.py

@@ -8,7 +8,7 @@ the existing separate operations (inverse RoPE via rotate_neox + FP8 quant
 via per_token_group_quant_fp8).
 
 The reference faithfully reproduces the exact flow in
-deepseek_v4/nvidia/ops/attention.py:295-310:
+deepseek_v4/attention.py:295-310:
   1. Apply inverse RoPE (NeoX style, last rope_dim=64 dims of each head)
   2. Reshape [T, H, head_dim] -> [T, G, D]
   3. Transpose+flatten to [G*T, D], quantize, reshape back
@@ -668,7 +668,7 @@ def _unfused_inv_rope_fp8_quant(
     nope_dim: int = NOPE_DIM,
     rope_dim: int = ROPE_DIM,
 ) -> tuple[torch.Tensor, torch.Tensor]:
-    """Unfused path matching deepseek_v4/nvidia/ops/attention.py:295-310.
+    """Unfused path matching deepseek_v4/attention.py:295-310.
 
     Uses the production CUDA RoPE kernel + per_token_group_quant_fp8.
     """

tests/v1/worker/test_gpu_model_runner_v2_eplb.py

@@ -168,7 +168,7 @@ def test_v2_sample_tokens_runs_eplb_on_non_last_pp_rank(monkeypatch):
         slot_mappings_by_layer=None,
         hidden_states=None,
         aux_hidden_states=None,
-        kv_connector_output=None,
+        finished_req_ids=set(),
         num_tokens_across_dp=None,
     )
     runner.postprocess = lambda *args, **kwargs: events.append("postprocess")

vllm/cute_utils/__init__.py

@@ -0,0 +1,134 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from cutlass import BFloat16, Float32, Int64, Uint32, cute
+from cutlass._mlir import ir
+from cutlass._mlir.dialects import llvm, vector
+from cutlass.cute.nvgpu import cpasync
+from cutlass.cutlass_dsl import T, dsl_user_op
+
+# https://github.com/NVIDIA/cutlass/blob/v4.3.2/include/cute/arch/copy_sm90_desc.hpp#L193-L197
+EVICT_NORMAL = Int64(0x1000000000000000)
+EVICT_FIRST = Int64(0x12F0000000000000)
+EVICT_LAST = Int64(0x14F0000000000000)
+
+
+@dsl_user_op
+def recast_val(x, dtype, *, loc=None, ip=None):
+    return dtype(llvm.bitcast(dtype.mlir_type, x.ir_value(loc=loc, ip=ip)))
+
+
+def simple_tma_copy(atom, src, dst, mbar=None, cache_policy=None):
+    """A simple helper that wraps group_modes() and tma_partition()
+    NOTE: this should be called WITHOUT cute.elect_one()
+    """
+    if isinstance(atom.op, cpasync.CopyBulkTensorTileG2SOp):
+        gmem = src
+        smem = dst
+    elif isinstance(atom.op, cpasync.CopyBulkTensorTileS2GOp):
+        smem = src
+        gmem = dst
+    else:
+        raise ValueError
+
+    s_part, g_part = cpasync.tma_partition(
+        atom,
+        0,
+        cute.make_layout(1),
+        cute.group_modes(smem, 0),
+        cute.group_modes(gmem, 0),
+    )
+
+    if isinstance(atom.op, cpasync.CopyBulkTensorTileG2SOp):
+        cute.copy(atom, g_part, s_part, tma_bar_ptr=mbar, cache_policy=cache_policy)
+    elif isinstance(atom.op, cpasync.CopyBulkTensorTileS2GOp):
+        cute.copy(atom, s_part, g_part, cache_policy=cache_policy)
+    else:
+        raise ValueError
+
+
+# can't find the equivalent in nvvm
+@dsl_user_op
+def fence_before_tma_store(*, loc=None, ip=None):
+    llvm.inline_asm(
+        T.i32(),
+        [],
+        "mov.u32 $0, 0;\n\t"
+        "fence.proxy.async::generic.release.sync_restrict::shared::cta.cluster;",
+        "=r",
+        has_side_effects=True,
+        is_align_stack=False,
+        loc=loc,
+        ip=ip,
+    )
+
+
+@dsl_user_op
+def mma_bf16(
+    a: cute.TensorSSA, b: cute.TensorSSA, c: cute.TensorSSA, *, loc=None, ip=None
+):
+    if a.element_type == BFloat16:
+        a = cute.recast_tensor(a, Uint32)
+    if b.element_type == BFloat16:
+        b = cute.recast_tensor(b, Uint32)
+
+    mlir_ty = Float32.mlir_type
+    out = llvm.inline_asm(
+        llvm.StructType.get_literal([mlir_ty] * 4),
+        [a[i].ir_value(loc=loc, ip=ip) for i in range(4)]
+        + [b[i].ir_value(loc=loc, ip=ip) for i in range(2)]
+        + [c[i].ir_value(loc=loc, ip=ip) for i in range(4)],
+        "mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 "
+        "{$0, $1, $2, $3}, {$4, $5, $6, $7}, {$8, $9}, "
+        "{$10, $11, $12, $13};",
+        "=f,=f,=f,=f,r,r,r,r,r,r,f,f,f,f",
+        has_side_effects=False,
+        is_align_stack=False,
+        loc=loc,
+        ip=ip,
+    )
+    vec = vector.from_elements(
+        ir.VectorType.get([4], mlir_ty, loc=loc),
+        [llvm.extractvalue(mlir_ty, out, [i], loc=loc, ip=ip) for i in range(4)],
+        loc=loc,
+        ip=ip,
+    )
+    return cute.TensorSSA(vec, 4, Float32)
+
+
+@dsl_user_op
+def _bf16x2_abs(a: Uint32, *, loc=None, ip=None) -> Uint32:
+    out = llvm.inline_asm(
+        T.i32(),
+        [a.ir_value(loc=loc, ip=ip)],
+        "abs.bf16x2 $0, $1;",
+        "=r,r",
+        has_side_effects=False,
+        is_align_stack=False,
+    )
+    return Uint32(out)
+
+
+@dsl_user_op
+def _bf16x2_max(a: Uint32, b: Uint32, *, loc=None, ip=None) -> Uint32:
+    out = llvm.inline_asm(
+        T.i32(),
+        [a.ir_value(loc=loc, ip=ip), b.ir_value(loc=loc, ip=ip)],
+        "max.bf16x2 $0, $1, $2;",
+        "=r,r,r",
+        has_side_effects=False,
+        is_align_stack=False,
+    )
+    return Uint32(out)
+
+
+@dsl_user_op
+def _bf16x2_mul(a: Uint32, b: Uint32, *, loc=None, ip=None) -> Uint32:
+    out = llvm.inline_asm(
+        T.i32(),
+        [a.ir_value(loc=loc, ip=ip), b.ir_value(loc=loc, ip=ip)],
+        "mul.rn.bf16x2 $0, $1, $2;",
+        "=r,r,r",
+        has_side_effects=False,
+        is_align_stack=False,
+    )
+    return Uint32(out)

vllm/cute_utils/_tcgen05.py

@@ -0,0 +1,219 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# this module is named _tcgen05 to avoid name collision with cute.nvgpu.tcgen05
+
+import cutlass
+from cutlass import Boolean, Float32, Int32, Uint32, Uint64, cute
+from cutlass._mlir import ir
+from cutlass._mlir.dialects import llvm, nvvm, vector
+from cutlass.cutlass_dsl import dsl_user_op
+
+NVVM_CTA_GROUP_MAP = [
+    None,
+    nvvm.Tcgen05GroupKind.CTA_1,
+    nvvm.Tcgen05GroupKind.CTA_2,
+]
+LDST_MAP = {
+    "32x32b": (nvvm.Tcgen05LdStShape.SHAPE_32X32B, 1),
+    "16x128b": (nvvm.Tcgen05LdStShape.SHAPE_16X128B, 2),
+    "16x256b": (nvvm.Tcgen05LdStShape.SHAPE_16X256B, 4),
+}
+
+
+def _make_tmem_llvm_ptr(addr, *, loc=None, ip=None):
+    ptr_ty = llvm.PointerType.get(cute.AddressSpace.tmem.value)
+    val = Int32(addr).ir_value(loc=loc, ip=ip)
+    return llvm.inttoptr(ptr_ty, val, loc=loc, ip=ip)
+
+
+@dsl_user_op
+def alloc(
+    taddr: cute.Pointer,
+    cta_group: int = 1,
+    *,
+    loc=None,
+    ip=None,
+) -> None:
+    nvvm.tcgen05_alloc(
+        taddr.to_llvm_ptr(loc=loc, ip=ip),
+        Uint32(512).ir_value(loc=loc, ip=ip),
+        group=NVVM_CTA_GROUP_MAP[cta_group],
+        loc=loc,
+        ip=ip,
+    )
+
+
+@dsl_user_op
+def dealloc(cta_group: int = 1, *, loc=None, ip=None) -> None:
+    nvvm.tcgen05_dealloc(
+        _make_tmem_llvm_ptr(0, loc=loc, ip=ip),
+        Int32(512).ir_value(loc=loc, ip=ip),
+        group=NVVM_CTA_GROUP_MAP[cta_group],
+        loc=loc,
+        ip=ip,
+    )
+
+
+def make_bf16_idesc(
+    MMA_M: int,
+    MMA_N: int,
+    *,
+    negate_A: bool = False,
+    negate_B: bool = False,
+    transpose_A: bool = False,
+    transpose_B: bool = False,
+):
+    idesc = Uint32(
+        (1 << 4) | (1 << 7) | (1 << 10) | ((MMA_N >> 3) << 17) | ((MMA_M >> 4) << 24)
+    )
+    idesc |= Uint32(negate_A) << 13
+    idesc |= Uint32(negate_B) << 14
+    idesc |= Uint32(transpose_A) << 15
+    idesc |= Uint32(transpose_B) << 16
+    return idesc
+
+
+def make_sdesc_128B_swizzle(LBO: int):
+    SBO = 8 * 128
+    return Uint64((LBO >> 4 << 16) | (SBO >> 4 << 32) | (1 << 46) | (2 << 61))
+
+
+@dsl_user_op
+def mma_f16(
+    d_tmem,
+    a_desc,
+    b_desc,
+    idesc,
+    enable_input_d,
+    cta_group: int = 1,
+    *,
+    loc=None,
+    ip=None,
+) -> None:
+    nvvm.tcgen05_mma(
+        nvvm.Tcgen05MMAKind.F16,
+        NVVM_CTA_GROUP_MAP[cta_group],
+        _make_tmem_llvm_ptr(d_tmem, loc=loc, ip=ip),
+        Uint64(a_desc).ir_value(loc=loc, ip=ip),
+        Uint64(b_desc).ir_value(loc=loc, ip=ip),
+        Int32(idesc).ir_value(loc=loc, ip=ip),
+        Boolean(enable_input_d).ir_value(loc=loc, ip=ip),
+        loc=loc,
+        ip=ip,
+    )
+
+
+@dsl_user_op
+def mma_ts_f16(
+    d_tmem,
+    a_tmem,
+    b_desc,
+    idesc,
+    enable_input_d,
+    cta_group: int = 1,
+    *,
+    loc=None,
+    ip=None,
+) -> None:
+    nvvm.tcgen05_mma(
+        nvvm.Tcgen05MMAKind.F16,
+        NVVM_CTA_GROUP_MAP[cta_group],
+        _make_tmem_llvm_ptr(d_tmem, loc=loc, ip=ip),
+        _make_tmem_llvm_ptr(a_tmem, loc=loc, ip=ip),
+        Uint64(b_desc).ir_value(loc=loc, ip=ip),
+        Int32(idesc).ir_value(loc=loc, ip=ip),
+        Boolean(enable_input_d).ir_value(loc=loc, ip=ip),
+        loc=loc,
+        ip=ip,
+    )
+
+
+@dsl_user_op
+def commit(mbar, cta_mask=None, cta_group: int = 1, *, loc=None, ip=None):
+    mbar_llvm = mbar.to_llvm_ptr(loc=loc, ip=ip)
+    group = NVVM_CTA_GROUP_MAP[cta_group]
+    if cutlass.const_expr(cta_mask is not None):
+        nvvm.tcgen05_commit_arrive(
+            mbar_llvm,
+            multicast_mask=cta_mask.ir_value(loc=loc, ip=ip),
+            group=group,
+            loc=loc,
+            ip=ip,
+        )
+    else:
+        nvvm.tcgen05_commit_arrive(mbar_llvm, group=group, loc=loc, ip=ip)
+
+
+@dsl_user_op
+def ld(row, col, shape: str, num: int, *, loc=None, ip=None):
+    nvvm_shape, regs_per_num = LDST_MAP[shape]
+    num_regs = regs_per_num * num
+    tmem = (Int32(row) << Int32(16)) | Int32(col)
+    tmem_ptr = _make_tmem_llvm_ptr(tmem, loc=loc, ip=ip)
+
+    if num_regs == 1:
+        reg = nvvm.tcgen05_ld(Int32.mlir_type, nvvm_shape, tmem_ptr, loc=loc, ip=ip)
+        reg_f32 = llvm.bitcast(Float32.mlir_type, reg, loc=loc, ip=ip)
+        return Float32(reg_f32)
+
+    else:
+        vec_i32_ty = ir.VectorType.get([num_regs], Int32.mlir_type, loc=loc)
+        vec_f32_ty = ir.VectorType.get([num_regs], Float32.mlir_type, loc=loc)
+        regs = nvvm.tcgen05_ld(vec_i32_ty, nvvm_shape, tmem_ptr, loc=loc, ip=ip)
+        regs_f32 = llvm.bitcast(vec_f32_ty, regs, loc=loc, ip=ip)
+        return cute.TensorSSA(regs_f32, (num_regs,), Float32)
+
+
+@dsl_user_op
+def st(row, col, shape: str, num: int, vals, *, loc=None, ip=None) -> None:
+    # if input is TensorSSA, convert to Tensor so we can bitcast
+    if isinstance(vals, cute.TensorSSA):
+        vals_ = cute.make_rmem_tensor_like(vals)
+        vals_.store(vals)
+        vals = vals_
+
+    # bitcast to Int32
+    vals = cute.recast_tensor(vals, Int32)
+
+    nvvm_shape, regs_per_num = LDST_MAP[shape]
+    num_regs = regs_per_num * num
+    tmem = (Int32(row) << Int32(16)) | Int32(col)
+    tmem_ptr = _make_tmem_llvm_ptr(tmem, loc=loc, ip=ip)
+
+    if num_regs == 1:
+        nvvm.tcgen05_st(
+            nvvm_shape,
+            tmem_ptr,
+            vals[0].ir_value(loc=loc, ip=ip),
+            loc=loc,
+            ip=ip,
+        )
+    else:
+        vec_i32_ty = ir.VectorType.get([num_regs], Int32.mlir_type, loc=loc)
+        val_vec = vector.from_elements(
+            vec_i32_ty,
+            [vals[i].ir_value(loc=loc, ip=ip) for i in range(num_regs)],
+            loc=loc,
+            ip=ip,
+        )
+        nvvm.tcgen05_st(nvvm_shape, tmem_ptr, val_vec, loc=loc, ip=ip)
+
+
+@dsl_user_op
+def fence_after_thread_sync(*, loc=None, ip=None):
+    nvvm.tcgen05_fence(nvvm.Tcgen05FenceKind.AFTER_THREAD_SYNC, loc=loc, ip=ip)
+
+
+@dsl_user_op
+def fence_before_thread_sync(*, loc=None, ip=None):
+    nvvm.tcgen05_fence(nvvm.Tcgen05FenceKind.BEFORE_THREAD_SYNC, loc=loc, ip=ip)
+
+
+@dsl_user_op
+def wait_ld(*, loc=None, ip=None):
+    nvvm.tcgen05_wait(nvvm.Tcgen05WaitKind.LOAD, loc=loc, ip=ip)
+
+
+@dsl_user_op
+def wait_st(*, loc=None, ip=None):
+    nvvm.tcgen05_wait(nvvm.Tcgen05WaitKind.STORE, loc=loc, ip=ip)

vllm/cute_utils/cvt.py

@@ -1,21 +1,13 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-import cutlass
-import cutlass.cute as cute
-from cutlass import Float32, Uint32
+from cutlass import Constexpr, Float32, Uint32, cute
 from cutlass._mlir import ir
 from cutlass._mlir.dialects import llvm, vector
 from cutlass.cutlass_dsl import T, dsl_user_op
 
 
 @dsl_user_op
-def _recast_val(x, dtype, *, loc=None, ip=None):
-    return dtype(llvm.bitcast(dtype.mlir_type, x.ir_value(loc=loc, ip=ip)))
-
-
-@dsl_user_op
-def _fp32x2_to_bf16x2(a: Float32, b: Float32, *, loc=None, ip=None) -> Uint32:
+def fp32x2_to_bf16x2(a: Float32, b: Float32, *, loc=None, ip=None) -> Uint32:
     out = llvm.inline_asm(
         T.i32(),
         [a.ir_value(loc=loc, ip=ip), b.ir_value(loc=loc, ip=ip)],
@@ -28,62 +20,37 @@ def _fp32x2_to_bf16x2(a: Float32, b: Float32, *, loc=None, ip=None) -> Uint32:
 
 
 @dsl_user_op
-def _bf16x2_to_fp32(data: Uint32, *, loc=None, ip=None) -> tuple[Float32, Float32]:
-    out = llvm.inline_asm(
-        llvm.StructType.get_literal([T.f32(), T.f32()]),
-        [data.ir_value(loc=loc, ip=ip)],
-        "shl.b32 $0, $2, 16;\n\tand.b32 $1, $2, 0xFFFF0000;\n",
-        "=f,=f,r",
-        has_side_effects=False,
-        is_align_stack=False,
-    )
-    return (
-        Float32(llvm.extractvalue(T.f32(), out, [0], loc=loc, ip=ip)),
-        Float32(llvm.extractvalue(T.f32(), out, [1], loc=loc, ip=ip)),
-    )
+def bf16x2_to_fp32x2(data, *, loc=None, ip=None) -> tuple[Float32, Float32]:
+    if isinstance(data, Uint32):
+        out = llvm.inline_asm(
+            llvm.StructType.get_literal([T.f32(), T.f32()]),
+            [data.ir_value(loc=loc, ip=ip)],
+            "shl.b32 $0, $2, 16;\n\tand.b32 $1, $2, 0xFFFF0000;",
+            "=f,=f,r",
+            has_side_effects=False,
+            is_align_stack=False,
+            loc=loc,
+            ip=ip,
+        )
+        return (
+            Float32(llvm.extractvalue(T.f32(), out, [0], loc=loc, ip=ip)),
+            Float32(llvm.extractvalue(T.f32(), out, [1], loc=loc, ip=ip)),
+        )
+
+    elif isinstance(data, (cute.Tensor, cute.TensorSSA)):
+        # NOTE: the output is always 1D
+        size = cute.size(data.shape)
+        out = cute.make_rmem_tensor(size * 2, Float32)
+        for i in range(size):
+            out[i * 2], out[i * 2 + 1] = bf16x2_to_fp32x2(data[i])
+        return out
+
+    else:
+        raise ValueError(f"Unsupported type {type(data)}")
 
 
 @dsl_user_op
-def _bf16x2_abs(a: Uint32, *, loc=None, ip=None) -> Uint32:
-    out = llvm.inline_asm(
-        T.i32(),
-        [a.ir_value(loc=loc, ip=ip)],
-        "abs.bf16x2 $0, $1;",
-        "=r,r",
-        has_side_effects=False,
-        is_align_stack=False,
-    )
-    return Uint32(out)
-
-
-@dsl_user_op
-def _bf16x2_max(a: Uint32, b: Uint32, *, loc=None, ip=None) -> Uint32:
-    out = llvm.inline_asm(
-        T.i32(),
-        [a.ir_value(loc=loc, ip=ip), b.ir_value(loc=loc, ip=ip)],
-        "max.bf16x2 $0, $1, $2;",
-        "=r,r,r",
-        has_side_effects=False,
-        is_align_stack=False,
-    )
-    return Uint32(out)
-
-
-@dsl_user_op
-def _bf16x2_mul(a: Uint32, b: Uint32, *, loc=None, ip=None) -> Uint32:
-    out = llvm.inline_asm(
-        T.i32(),
-        [a.ir_value(loc=loc, ip=ip), b.ir_value(loc=loc, ip=ip)],
-        "mul.rn.bf16x2 $0, $1, $2;",
-        "=r,r,r",
-        has_side_effects=False,
-        is_align_stack=False,
-    )
-    return Uint32(out)
-
-
-@dsl_user_op
-def _fp8x4_to_bf16x4(x: Uint32, *, loc=None, ip=None) -> cute.TensorSSA:
+def fp8x4_to_bf16x4(x: Uint32, *, loc=None, ip=None) -> cute.TensorSSA:
     # there is only fp8->fp16 conversion, hence we need to go
     # round trip through fp16.
     out = llvm.inline_asm(
@@ -118,7 +85,7 @@ def _fp8x4_to_bf16x4(x: Uint32, *, loc=None, ip=None) -> cute.TensorSSA:
 
 
 @dsl_user_op
-def _fp32x4_to_fp8x4(
+def fp32x4_to_fp8x4(
     a0: Float32,
     a1: Float32,
     a2: Float32,
@@ -151,9 +118,9 @@ def _fp32x4_to_fp8x4(
 
 
 @dsl_user_op
-def _fp32x8_to_fp4x8(
+def fp32x8_to_fp4x8(
     vals: cute.Tensor,
-    offset: cutlass.Constexpr[int],
+    offset: Constexpr[int],
     *,
     loc=None,
     ip=None,

vllm/engine/arg_utils.py

@@ -712,7 +712,7 @@ class EngineArgs:
     )
 
     fail_on_environ_validation: bool = False
-    gdn_prefill_backend: Literal["flashinfer", "triton"] | None = None
+    gdn_prefill_backend: Literal["flashinfer", "triton", "cutedsl"] | None = None
 
     def __post_init__(self):
         # support `EngineArgs(compilation_config={...})`
@@ -1527,7 +1527,7 @@ class EngineArgs:
         parser.add_argument(
             "--gdn-prefill-backend",
             dest="gdn_prefill_backend",
-            choices=["flashinfer", "triton"],
+            choices=["flashinfer", "triton", "cutedsl"],
             default=None,
             help="Select GDN prefill backend.",
         )

vllm/entrypoints/cli/serve.py

@@ -308,7 +308,14 @@ def run_multi_api_server(args: argparse.Namespace):
 
     from vllm.v1.engine.utils import get_engine_zmq_addresses
 
-    addresses = get_engine_zmq_addresses(vllm_config, num_api_servers)
+    # Per-API-server ports are picked by the kernel at each child's bind()
+    # to avoid parent-probe vs child-bind TOCTOU; Rust front-end opts out
+    # because it has no port-report-back channel.
+    addresses = get_engine_zmq_addresses(
+        vllm_config,
+        num_api_servers,
+        defer_api_server_ports=not rust_frontend_path,
+    )
 
     with launch_core_engines(
         vllm_config, executor_class, log_stats, addresses, num_api_servers
@@ -341,6 +348,12 @@ def run_multi_api_server(args: argparse.Namespace):
                 tensor_queue=tensor_queue,
             )
 
+            # Forward each child's bound endpoints to the engine handshake
+            # (runs on ``with`` exit).
+            actual_inputs, actual_outputs = api_server_manager.gather_actual_addresses()
+            addresses.inputs = actual_inputs
+            addresses.outputs = actual_outputs
+
     # Wait for API servers.
     try:
         wait_for_completion_or_failure(

vllm/lora/ops/triton_ops/fused_moe_lora_op.py

@@ -7,6 +7,7 @@ from vllm.distributed import (
     tensor_model_parallel_all_gather,
     tensor_model_parallel_all_reduce,
 )
+from vllm.platforms import current_platform
 from vllm.triton_utils import tl, triton
 from vllm.triton_utils.allocation import set_triton_allocator
 from vllm.utils.torch_utils import direct_register_custom_op
@@ -406,7 +407,7 @@ def _run_fused_moe_lora_one_shot(
 
     # NPID_FACTOR heuristic: scale N-axis parallelism when base CTA count is
     # short of saturating the SM array. Cap by the cost of redundant shrink.
-    sm_count = torch.cuda.get_device_properties(device).multi_processor_count
+    sm_count = current_platform.num_compute_units(device.index)
     base_programs = max(M_blocks * num_slices * grid_lora_dim, 1)
     shrink_ratio = K / max(K + N_per_slice, 1)
     max_npid_by_budget = max(1, int(1.5 / max(shrink_ratio, 1e-3)) + 1)
@@ -786,7 +787,7 @@ def _run_fused_moe_lora_small_batch(
     N_tiles = triton.cdiv(N_per_slice, BLOCK_N)
     pair_slices = M_grid * num_slices
 
-    sm_count = torch.cuda.get_device_properties(device).multi_processor_count
+    sm_count = current_platform.num_compute_units(device.index)
     n_tiles_per_program = _pick_small_batch_chunk(pair_slices, N_tiles, sm_count)
     n_chunks = triton.cdiv(N_tiles, n_tiles_per_program)
     work_total = pair_slices * n_chunks

vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py

@@ -3,6 +3,7 @@
 """Inference-only Qwen3-Next/Qwen3.5 model."""
 
 import functools
+from typing import Literal
 
 import torch
 from einops import rearrange
@@ -83,7 +84,7 @@ logger = init_logger(__name__)
 
 
 # TODO(arpera): remove ``_is_libs_cu13_install_intact`` and its caller in
-# ``_should_use_flashinfer_gdn_prefill`` once the upstream packaging bug is
+# ``_resolve_gdn_prefill_backend`` once the upstream packaging bug is
 # fixed and the broken wheels are yanked / superseded on PyPI:
 #   https://github.com/NVIDIA/cutlass/issues/3170
 #   https://github.com/NVIDIA/cutlass/issues/3259
@@ -146,11 +147,12 @@ def _is_libs_cu13_install_intact() -> bool:
     return True
 
 
-def _should_use_flashinfer_gdn_prefill(backend: str, head_k_dim: int | None) -> bool:
-    """Whether to use FlashInfer's GDN prefill kernel instead of the
-    Triton/FLA fallback.
+def _resolve_gdn_prefill_backend(
+    vllm_config: VllmConfig,
+) -> tuple[str, Literal["triton", "flashinfer", "cutedsl"]]:
+    """Resolve GDN prefill backend.
 
-    Requirements:
+    FlashInfer's GDN prefill kernel is chosen when:
     * ``requested in ["flashinfer", "auto"]``;
     * ``platform == cuda``;
     * one of the following:
@@ -158,47 +160,78 @@ def _should_use_flashinfer_gdn_prefill(backend: str, head_k_dim: int | None) ->
       - Blackwell (SM10.x) with ``head_k_dim == 128``, ``cuda_runtime >= 13``,
         and an intact ``nvidia-cutlass-dsl-libs-cu13`` install on disk
         (see :func:`_is_libs_cu13_install_intact`).
+
+    In-tree CuteDSL GDN prefill kernel is chosen when:
+    * "cutedsl" is requested; (opt-in only)
+    * Blackwell (SM10.x) with ``head_k_dim == 128``;
     """
-    if backend not in ["flashinfer", "auto"]:
-        return False
+    additional_config = vllm_config.additional_config
+    backend_cfg = (
+        additional_config.get("gdn_prefill_backend", "auto")
+        if isinstance(additional_config, dict)
+        else "auto"
+    )
+    backend = str(backend_cfg).strip().lower()
+
     if not current_platform.is_cuda():
-        return False
+        return backend, "triton"
+
+    head_k_dim = getattr(
+        vllm_config.model_config.hf_config, "linear_key_head_dim", None
+    )
+
+    supports_flashinfer = False
+    supports_cutedsl = False
+
     if current_platform.is_device_capability(90):
-        return True  # Hopper — no further constraints.
-    if not current_platform.is_device_capability_family(100):
-        return False  # Neither Hopper nor Blackwell.
-    if head_k_dim != 128:
-        return False
-    if current_platform.get_cuda_runtime_major() < 13:
-        return False
-    if not _is_libs_cu13_install_intact():
-        logger.warning_once(
-            "FlashInfer Blackwell GDN requires an intact nvidia-cutlass-dsl"
-            "-libs-cu13 install, but some on-disk files do not match the "
-            "SHA-256 declared in its RECORD (install-order race in "
-            "nvidia-cutlass-dsl packaging — see "
-            "https://github.com/NVIDIA/cutlass/issues/3170 and "
-            "https://github.com/NVIDIA/cutlass/issues/3259). Falling back "
-            "to Triton/FLA. Repair with: pip install --force-reinstall "
-            "--no-deps nvidia-cutlass-dsl-libs-cu13"
-        )
-        return False
-    return True
+        supports_flashinfer = True
+    elif (
+        current_platform.is_device_capability_family(100)
+        and head_k_dim == 128
+        and current_platform.get_cuda_runtime_major() >= 13
+    ):
+        supports_flashinfer = _is_libs_cu13_install_intact()
+        supports_cutedsl = True
+        if not supports_flashinfer:
+            logger.warning_once(
+                "FlashInfer Blackwell GDN requires an intact nvidia-cutlass-dsl"
+                "-libs-cu13 install, but some on-disk files do not match the "
+                "SHA-256 declared in its RECORD (install-order race in "
+                "nvidia-cutlass-dsl packaging -- see "
+                "https://github.com/NVIDIA/cutlass/issues/3170 and "
+                "https://github.com/NVIDIA/cutlass/issues/3259). Falling back "
+                "to Triton/FLA. Repair with: pip install --force-reinstall "
+                "--no-deps nvidia-cutlass-dsl-libs-cu13"
+            )
+
+    if backend in ["flashinfer", "auto"] and supports_flashinfer:
+        return backend, "flashinfer"
+    if backend == "cutedsl" and supports_cutedsl:
+        return backend, "cutedsl"
+    return backend, "triton"
 
 
 def _log_gdn_backend_decision(
-    backend: str, head_k_dim: int | None, use_flashinfer: bool
+    vllm_config: VllmConfig,
+    requested_backend: str,
+    active_backend: str,
 ) -> None:
     """Log the GDN prefill backend choice in the attention-selector style."""
-    chosen = "FlashInfer" if use_flashinfer else "Triton/FLA"
+    head_k_dim = getattr(
+        vllm_config.model_config.hf_config, "linear_key_head_dim", None
+    )
+    chosen = {
+        "flashinfer": "FlashInfer",
+        "cutedsl": "CuteDSL",
+        "triton": "Triton/FLA",
+    }[active_backend]
     logger.info_once(
         "Using %s GDN prefill kernel (requested=%s, head_k_dim=%s).",
         chosen,
-        backend,
+        requested_backend,
         head_k_dim,
     )
-    # JIT-compiled cutlass path is only used on SM90 (Hopper).
-    if use_flashinfer and current_platform.is_device_capability(90):
+    if active_backend == "flashinfer" and current_platform.is_device_capability(90):
         logger.warning_once(
             "FlashInfer GDN prefill is JIT-compiled; first run may take a "
             "while. Set --gdn-prefill-backend triton to skip JIT.",
@@ -256,25 +289,26 @@ def fi_chunk_gated_delta_rule(
 
 @CustomOp.register("chunk_gated_delta_rule")
 class ChunkGatedDeltaRule(CustomOp):
-    def __init__(self, head_k_dim: int | None = None) -> None:
+    def __init__(self) -> None:
         super().__init__()
-        additional_config = get_current_vllm_config().additional_config
-        assert isinstance(additional_config, dict)
-        backend_cfg = additional_config.get("gdn_prefill_backend", "auto")
-        backend = str(backend_cfg).strip().lower()
+        vllm_config = get_current_vllm_config()
+        backend, active_backend = _resolve_gdn_prefill_backend(vllm_config)
+        self.gdn_prefill_backend = active_backend
 
-        use_flashinfer = _should_use_flashinfer_gdn_prefill(backend, head_k_dim)
-        if backend == "flashinfer" and not use_flashinfer:
+        if backend in ("flashinfer", "cutedsl") and active_backend != backend:
             logger.warning_once(
-                "GDN prefill backend 'flashinfer' is selected but "
-                "cannot use this kernel on the current platform. "
-                "Falling back to Triton/FLA."
+                "GDN prefill backend '%s' is selected but cannot use this "
+                "kernel on the current platform. Falling back to Triton/FLA.",
+                backend,
             )
-        _log_gdn_backend_decision(backend, head_k_dim, use_flashinfer)
+        _log_gdn_backend_decision(vllm_config, backend, active_backend)
 
-        self._forward_method = (
-            self.forward_cuda if use_flashinfer else self.forward_native
-        )
+        if active_backend == "flashinfer":
+            self._forward_method = self.forward_cuda
+        elif active_backend == "cutedsl":
+            self._forward_method = self.forward_cutedsl
+        else:
+            self._forward_method = self.forward_native
 
     def forward_cuda(
         self,
@@ -338,6 +372,49 @@ class ChunkGatedDeltaRule(CustomOp):
             core_attn_out=core_attn_out,
         )
 
+    def forward_cutedsl(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        g: torch.Tensor,
+        beta: torch.Tensor,
+        initial_state: torch.Tensor,
+        output_final_state: bool,
+        cu_seqlens: torch.Tensor | None = None,
+        chunk_indices: torch.Tensor | None = None,
+        chunk_offsets: torch.Tensor | None = None,
+        use_qk_l2norm_in_kernel: bool = True,
+        core_attn_out: torch.Tensor | None = None,
+    ):
+        from vllm.model_executor.layers.mamba.ops.gdn_chunk_cutedsl import (
+            chunk_gated_delta_rule_cutedsl,
+        )
+
+        if use_qk_l2norm_in_kernel:
+            q = l2norm_fwd(q)
+            k = l2norm_fwd(k)
+
+        assert cu_seqlens is not None
+        assert chunk_indices is not None
+        assert chunk_offsets is not None
+
+        o, final_state = chunk_gated_delta_rule_cutedsl(
+            q=q,
+            k=k,
+            v=v,
+            g=g,
+            beta=beta,
+            initial_state=initial_state,
+            cu_seqlens=cu_seqlens,
+            chunk_indices=chunk_indices,
+            chunk_offsets=chunk_offsets,
+            core_attn_out=core_attn_out,
+        )
+        if not output_final_state:
+            final_state = None
+        return o, final_state
+
 
 @PluggableLayer.register("qwen_gated_delta_net_attention")
 class QwenGatedDeltaNetAttention(GatedDeltaNetAttention):
@@ -474,7 +551,8 @@ class QwenGatedDeltaNetAttention(GatedDeltaNetAttention):
             prefix=f"{prefix}.out_proj",
         )
 
-        self.chunk_gated_delta_rule = ChunkGatedDeltaRule(head_k_dim=self.head_k_dim)
+        self.chunk_gated_delta_rule = ChunkGatedDeltaRule()
+        self.gdn_prefill_backend = self.chunk_gated_delta_rule.gdn_prefill_backend
         self._prefill_kernels_warmed_up = False
         self.enable_packed_recurrent_decode = (
             envs.VLLM_ENABLE_FLA_PACKED_RECURRENT_DECODE
@@ -1060,6 +1138,16 @@ class QwenGatedDeltaNetAttention(GatedDeltaNetAttention):
         )
         cu_seqlens = torch.tensor([0, T], device=device, dtype=torch.int32)
 
+        # CuteDSL kernels require metadata
+        chunk_indices = None
+        chunk_offsets = None
+        if self.gdn_prefill_backend == "cutedsl":
+            from vllm.model_executor.layers.mamba.ops.gdn_chunk_cutedsl import (
+                prepare_metadata_cutedsl,
+            )
+
+            chunk_indices, chunk_offsets = prepare_metadata_cutedsl(cu_seqlens, T)
+
         try:
             self.chunk_gated_delta_rule(
                 q=q,
@@ -1070,6 +1158,8 @@ class QwenGatedDeltaNetAttention(GatedDeltaNetAttention):
                 initial_state=state,
                 output_final_state=True,
                 cu_seqlens=cu_seqlens,
+                chunk_indices=chunk_indices,
+                chunk_offsets=chunk_offsets,
                 use_qk_l2norm_in_kernel=False,
             )
         except Exception:
@@ -1088,7 +1178,20 @@ class QwenGatedDeltaNetAttention(GatedDeltaNetAttention):
                 self.prefix,
             )
         finally:
-            del dummy_mixed_qkv, q, k, v, dummy_a, dummy_b, g, beta, state, cu_seqlens
+            del (
+                dummy_mixed_qkv,
+                q,
+                k,
+                v,
+                dummy_a,
+                dummy_b,
+                g,
+                beta,
+                state,
+                cu_seqlens,
+                chunk_indices,
+                chunk_offsets,
+            )
 
         torch.accelerator.empty_cache()
 

vllm/model_executor/layers/mamba/ops/gdn_chunk_cutedsl/__init__.py

@@ -0,0 +1,251 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from functools import cache
+
+import cutlass
+import torch
+from cuda.bindings.driver import CUstream
+from cutlass import Int32, cute
+from quack.compile_utils import make_fake_tensor
+
+from vllm.triton_utils import triton
+
+from .kernel_h import h_cutedsl
+from .kernel_kkt_inv_uw import kkt_inv_uw_cutedsl
+from .kernel_o import o_cutedsl
+
+
+class PrepMetaKernel:
+    def __init__(self, BT: int) -> None:
+        self.BT = BT
+        self.num_warps = 8
+
+    @cute.jit
+    def __call__(
+        self,
+        cu_seqlens: cute.Tensor,
+        chunk_indices: cute.Tensor,
+        chunk_offsets: cute.Tensor,
+        stream: CUstream,
+    ):
+        block = (self.num_warps * 32, 1, 1)
+        self.kernel(
+            cu_seqlens,
+            chunk_indices,
+            chunk_offsets,
+        ).launch(grid=(1, 1, 1), block=block, stream=stream)
+
+    @cute.kernel
+    def kernel(
+        self,
+        cu_seqlens: cute.Tensor,
+        chunk_indices: cute.Tensor,
+        chunk_offsets: cute.Tensor,
+    ):
+        tid, _, _ = cute.arch.thread_idx()
+        warp_id = cute.arch.make_warp_uniform(tid // 32)
+        lane_id = tid % 32
+
+        num_seqs = cu_seqlens.shape[0] - 1
+        num_warps = self.num_warps
+        tb_size = num_warps * 32
+
+        if tid == 0:
+            chunk_offsets[0] = 0
+
+        coarsen = cute.ceil_div(num_seqs, tb_size)
+        seq_start = tid * coarsen
+        num_iters = cutlass.min(seq_start + coarsen, num_seqs) - seq_start
+
+        # First pass: compute this thread's total chunk count.
+        thread_sum = Int32(0)
+        for i in range(num_iters):
+            seq_id = seq_start + i
+            seqlen = cu_seqlens[seq_id + 1] - cu_seqlens[seq_id]
+            thread_sum += cute.ceil_div(seqlen, self.BT)
+
+        # warp parallel scan
+        cu_num_chunks = thread_sum
+        for i in cutlass.range_constexpr(5):
+            offset = cutlass.const_expr(1 << i)
+            lower = cute.arch.shuffle_sync_up(
+                cu_num_chunks, offset=offset, mask_and_clamp=0
+            )
+            if lane_id >= offset:
+                cu_num_chunks += lower
+
+        # cross-warp cumsum (CTA-wide)
+        smem = cutlass.utils.SmemAllocator()
+        warp_num_chunks = smem.allocate_array(Int32, num_warps)
+        if lane_id == 31:
+            warp_num_chunks[warp_id] = cu_num_chunks
+        cute.arch.sync_threads()
+
+        for i in cutlass.range_constexpr(1, num_warps):
+            if warp_id >= i:
+                cu_num_chunks += warp_num_chunks[i - 1]
+
+        chunk_start = cu_num_chunks - thread_sum
+
+        # Second pass: recompute per-sequence chunk counts and write results.
+        for i in range(num_iters):
+            seq_id = seq_start + i
+            seqlen = cu_seqlens[seq_id + 1] - cu_seqlens[seq_id]
+            num_chunks = cute.ceil_div(seqlen, self.BT)
+            chunk_end = chunk_start + num_chunks
+            chunk_offsets[seq_id + 1] = chunk_end
+
+            for chunk_id in range(num_chunks):
+                chunk_indices[chunk_start + chunk_id, 0] = seq_id
+                chunk_indices[chunk_start + chunk_id, 1] = chunk_id
+
+            chunk_start = chunk_end
+
+    @cache
+    @staticmethod
+    def compile(BT: int):
+        cu_entries = cute.sym_int()
+        upper_bound_chunks = cute.sym_int()
+
+        cu_seqlens = make_fake_tensor(Int32, (cu_entries,), divisibility=1)
+        chunk_indices = make_fake_tensor(Int32, (upper_bound_chunks, 2), divisibility=2)
+        chunk_offsets = make_fake_tensor(Int32, (cu_entries,), divisibility=1)
+
+        kernel = PrepMetaKernel(BT)
+        stream = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True)
+        return cute.compile(
+            kernel,
+            cu_seqlens,
+            chunk_indices,
+            chunk_offsets,
+            stream,
+            options="--enable-tvm-ffi",
+        )
+
+
+def _upper_bound_chunks(num_seqs: int, total_tokens: int, chunk_size: int) -> int:
+    return (num_seqs - 1) + triton.cdiv(total_tokens - (num_seqs - 1), chunk_size)
+
+
+def prepare_metadata_cutedsl(
+    cu_seqlens: torch.Tensor,
+    total_tokens: int,
+    chunk_size: int = 64,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    num_seqs = cu_seqlens.numel() - 1
+    upper_bound_chunks = _upper_bound_chunks(num_seqs, total_tokens, chunk_size)
+    chunk_offsets = cu_seqlens.new_empty(num_seqs + 1, dtype=torch.int32)
+    chunk_indices = cu_seqlens.new_empty((upper_bound_chunks, 2), dtype=torch.int32)
+
+    PrepMetaKernel.compile(chunk_size)(cu_seqlens, chunk_indices, chunk_offsets)
+    return chunk_indices, chunk_offsets
+
+
+def chunk_gated_delta_rule_cutedsl(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    g: torch.Tensor,
+    beta: torch.Tensor,
+    initial_state: torch.Tensor,
+    cu_seqlens: torch.Tensor,
+    chunk_indices: torch.Tensor,
+    chunk_offsets: torch.Tensor,
+    core_attn_out: torch.Tensor | None = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Run the GDN chunk CuteDSL prefill kernels.
+
+    Args:
+        q: Query tensor with shape ``[1, T, H, K]``.
+        k: Key tensor with shape ``[1, T, H, K]``.
+        v: Value tensor with shape ``[1, T, Hv, V]``.
+        g: Log-space decay tensor with shape ``[1, T, Hv]``.
+        beta: Delta-rule beta tensor with shape ``[1, T, Hv]``.
+        initial_state: Recurrent state with shape ``[N, Hv, V, K]``.
+        cu_seqlens: Cumulative sequence lengths with shape ``[N + 1]``.
+        chunk_indices: Chunk index metadata with shape ``[NT, 2]``.
+        chunk_offsets: Cumulative chunk offsets with shape ``[N + 1]``.
+        core_attn_out: Optional output buffer with shape ``[T, Hv, V]``.
+
+    Returns:
+        A tuple ``(output, final_state)`` where ``output`` has shape
+        ``[1, T, Hv, V]`` and ``final_state`` has shape ``[N, Hv, V, K]``.
+        When ``core_attn_out`` is provided, ``output`` is an unsqueezed view of
+        that buffer.
+    """
+    q_3d = q.squeeze(0)
+    k_3d = k.squeeze(0)
+    v_3d = v.squeeze(0)
+    g_2d = g.squeeze(0)
+    beta_2d = beta.squeeze(0)
+
+    _, _, head_k_dim = k_3d.shape
+    _, num_v_heads, head_v_dim = v_3d.shape
+    chunk_size = 64
+    upper_bound_chunks = chunk_indices.shape[0]
+    pad_t = upper_bound_chunks * chunk_size
+    total_chunks_ptr = chunk_offsets[-1:]
+
+    g_cu = torch.empty_like(g_2d, dtype=torch.float32)
+    u = q_3d.new_empty(pad_t, num_v_heads, head_v_dim)
+    w = q_3d.new_empty(pad_t, num_v_heads, head_k_dim)
+
+    num_sms = torch.cuda.get_device_properties(q.device).multi_processor_count
+    kkt_inv_uw_cutedsl(
+        k_3d,
+        v_3d,
+        u,
+        w,
+        g_2d,
+        beta_2d,
+        g_cu,
+        cu_seqlens,
+        chunk_indices,
+        total_chunks_ptr,
+        num_sms=num_sms,
+    )
+
+    h = k_3d.new_empty(
+        upper_bound_chunks,
+        num_v_heads,
+        head_v_dim,
+        head_k_dim,
+    )
+    v_new = q_3d.new_empty(pad_t, num_v_heads, head_v_dim)
+    final_state = torch.empty_like(initial_state)
+    h_cutedsl(
+        k_3d,
+        u,
+        w,
+        v_new,
+        g_cu,
+        h,
+        initial_state,
+        final_state,
+        cu_seqlens,
+        chunk_offsets,
+    )
+
+    output = core_attn_out if core_attn_out is not None else torch.empty_like(v_3d)
+    scale = head_k_dim**-0.5
+    o_cutedsl(
+        q_3d,
+        k_3d,
+        v_new.view(upper_bound_chunks, chunk_size, num_v_heads, head_v_dim),
+        h,
+        g_cu,
+        output,
+        cu_seqlens,
+        chunk_indices,
+        total_chunks_ptr,
+        scale,
+        num_sms=num_sms,
+    )
+    return output.unsqueeze(0), final_state
+
+
+__all__ = [
+    "chunk_gated_delta_rule_cutedsl",
+    "prepare_metadata_cutedsl",
+]

vllm/model_executor/layers/mamba/ops/gdn_chunk_cutedsl/kernel_h.py

@@ -0,0 +1,753 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from functools import cache
+
+import cutlass
+import torch
+from cuda.bindings.driver import CUstream
+from cutlass import BFloat16, Float32, Int32, Int64, Uint32, cute
+from cutlass.cute.nvgpu import cpasync, warp
+from quack.compile_utils import make_fake_tensor
+
+from vllm.cute_utils import (
+    EVICT_FIRST,
+    _tcgen05,
+    cvt,
+    fence_before_tma_store,
+    simple_tma_copy,
+)
+
+
+class Sm100ChunkHKernel:
+    """For each sequence, compute the chunk recurrent update.
+
+    The input V tile is the U output from the KKT/UW kernel. For each chunk:
+        V_new = U - W @ H.T
+        (we actually do V_new.T = U.T - H @ W.T instead)
+
+        H_scaled = H * exp(g_last)
+        V_scaled = V_new * exp(g_last - g)
+        H_new = H_scaled + V_scaled.T @ K
+    """
+
+    def __init__(
+        self,
+        H: int,
+        Hv: int,
+        K_dim: int,
+        V_dim: int,
+        h_dtype: cutlass.Numeric = Float32,
+        BT: int = 64,
+        num_stages: int = 2,
+    ) -> None:
+        assert Hv % H == 0
+        assert K_dim == V_dim == 128
+        assert BT == 64
+        self.H = H
+        self.Hv = Hv
+        self.K_dim = K_dim
+        self.V_dim = V_dim
+        self.h_dtype = h_dtype
+        self.BT = BT
+        self.num_stages = num_stages
+        self.num_warps = 10
+
+    @cute.jit
+    def _make_bf16_tma_args(
+        self,
+        tensor: cute.Tensor,
+        dim: cutlass.Constexpr[int],
+        op: cpasync.TmaCopyOp,
+        stages: cutlass.Constexpr[int],
+    ):
+        swizzle_128B = cute.make_swizzle(3, 4, 3)
+        slayout = cute.make_layout(
+            (self.BT, 1, (64, dim // 64), stages),
+            stride=(64, 0, (1, self.BT * 64), self.BT * dim),
+        )
+        slayout = cute.make_composed_layout(swizzle_128B, 0, slayout)
+        atom, tma_tensor = cpasync.make_tiled_tma_atom(
+            op,
+            cute.logical_divide(tensor, (None, None, 64)),
+            slayout,
+            cta_tiler=(self.BT, 1, dim),
+        )
+        return atom, tma_tensor, slayout
+
+    @cute.jit
+    def _make_h_tma_args(self, tensor: cute.Tensor, op: cpasync.TmaCopyOp):
+        # number of elements to fill 128B
+        num_elems = 128 // (tensor.element_type.width // 8)
+        swizzle_128B = cute.make_swizzle(3, 4, 3)
+        slayout = cute.make_layout(
+            (1, 1, self.V_dim, (num_elems, self.K_dim // num_elems)),
+            stride=(0, 0, num_elems, (1, self.V_dim * num_elems)),
+        )
+        slayout = cute.make_composed_layout(swizzle_128B, 0, slayout)
+        atom, tma_tensor = cpasync.make_tiled_tma_atom(
+            op,
+            cute.logical_divide(tensor, (None, None, None, num_elems)),
+            slayout,
+            cta_tiler=(1, 1, self.V_dim, self.K_dim),
+        )
+        return atom, tma_tensor, slayout
+
+    @cute.jit
+    def __call__(
+        self,
+        K: cute.Tensor,
+        V: cute.Tensor,
+        W: cute.Tensor,
+        V_new: cute.Tensor,
+        g_cu: cute.Tensor,
+        h: cute.Tensor,
+        h0: cute.Tensor,
+        ht: cute.Tensor,
+        cu_seqlens: cute.Tensor,
+        chunk_offsets: cute.Tensor,
+        stream: CUstream,
+    ):
+        tma_g2s = cpasync.CopyBulkTensorTileG2SOp()
+        tma_s2g = cpasync.CopyBulkTensorTileS2GOp()
+
+        K_args = self._make_bf16_tma_args(K, self.K_dim, tma_g2s, self.num_stages)
+        V_args = self._make_bf16_tma_args(V, self.V_dim, tma_g2s, self.num_stages)
+        W_args = self._make_bf16_tma_args(W, self.K_dim, tma_g2s, self.num_stages)
+        V_new_args = self._make_bf16_tma_args(V_new, self.V_dim, tma_s2g, 1)
+        H0_args = self._make_h_tma_args(h0, tma_g2s)
+        HT_args = self._make_h_tma_args(ht, tma_s2g)
+        H_args = self._make_h_tma_args(h, tma_s2g)
+
+        grid = (self.Hv, h0.shape[0], 1)
+        block = (self.num_warps * 32, 1, 1)
+        self.kernel(
+            K_args,
+            V_args,
+            W_args,
+            V_new_args,
+            H0_args,
+            HT_args,
+            H_args,
+            g_cu,
+            cu_seqlens,
+            chunk_offsets,
+        ).launch(grid=grid, block=block, stream=stream)
+
+    @cute.kernel
+    def kernel(
+        self,
+        K_args: tuple[cute.CopyAtom, cute.Tensor, cute.ComposedLayout],
+        V_args: tuple[cute.CopyAtom, cute.Tensor, cute.ComposedLayout],
+        W_args: tuple[cute.CopyAtom, cute.Tensor, cute.ComposedLayout],
+        V_new_args: tuple[cute.CopyAtom, cute.Tensor, cute.ComposedLayout],
+        H0_args: tuple[cute.CopyAtom, cute.Tensor, cute.ComposedLayout],
+        HT_args: tuple[cute.CopyAtom, cute.Tensor, cute.ComposedLayout],
+        H_args: tuple[cute.CopyAtom, cute.Tensor, cute.ComposedLayout],
+        g_cu: cute.Tensor,
+        cu_seqlens: cute.Tensor,
+        chunk_offsets: cute.Tensor,
+    ):
+        tid, _, _ = cute.arch.thread_idx()
+        head_id, seq_id, _ = cute.arch.block_idx()
+        warp_id = cute.arch.make_warp_uniform(tid // 32)
+        lane_id = tid % 32
+
+        BT = self.BT
+        V_dim = self.V_dim
+        K_dim = self.K_dim
+        num_stages = self.num_stages
+        is_f32 = self.h_dtype == Float32
+
+        K_tma_atom, tmaK, sK_layout = K_args
+        V_tma_atom, tmaV, sV_layout = V_args
+        W_tma_atom, tmaW, sW_layout = W_args
+        V_new_tma_atom, tmaV_new, sV_new_layout = V_new_args
+        H0_tma_atom, tmaH0, sH0_layout = H0_args
+        HT_tma_atom, tmaHT, _ = HT_args
+        H_tma_atom, tmaH, sH_layout = H_args
+
+        def allocate_tensor(smem, dtype, layout):
+            return smem.allocate_tensor(
+                dtype, layout.outer, byte_alignment=128, swizzle=layout.inner
+            )
+
+        smem = cutlass.utils.SmemAllocator()
+
+        # remove size=1 modes
+        sW = allocate_tensor(smem, BFloat16, sW_layout)[None, 0, None, None]
+        sV = allocate_tensor(smem, BFloat16, sV_layout)[None, 0, None, None]
+        sK = allocate_tensor(smem, BFloat16, sK_layout)[None, 0, None, None]
+        sH0 = allocate_tensor(smem, self.h_dtype, sH0_layout)[0, 0, None, None]
+        sH = allocate_tensor(smem, BFloat16, sH_layout)[0, 0, None, None]
+        sV_new = allocate_tensor(smem, BFloat16, sV_new_layout)[None, 0, None, 0]
+
+        s_v_scale = smem.allocate_array(Float32, BT)
+        tma_mbar = smem.allocate_array(Int64, num_stages)
+        wh_in_mbar = smem.allocate_array(Int64, num_stages)
+        wh_done_mbar = smem.allocate_array(Int64, num_stages)
+        vk_in_mbar = smem.allocate_array(Int64, num_stages)
+        vk_done_mbar = smem.allocate_array(Int64, num_stages)
+        h0_mbar = smem.allocate_array(Int64, 1)
+        taddr = smem.allocate(Int32, 4)
+
+        wh_tmem = 0
+        vk_tmem = wh_tmem + BT
+        h_tmem_base = vk_tmem + K_dim
+        v_tmem_base = h_tmem_base + K_dim // 2
+
+        if warp_id == 0:
+            with cute.arch.elect_one():
+                for i in cutlass.range_constexpr(num_stages):
+                    cute.arch.mbarrier_init(tma_mbar + i, 1)
+                    cute.arch.mbarrier_init(wh_in_mbar + i, 256)
+                    cute.arch.mbarrier_init(wh_done_mbar + i, 1)
+                    cute.arch.mbarrier_init(vk_in_mbar + i, 256)
+                    cute.arch.mbarrier_init(vk_done_mbar + i, 1)
+                cute.arch.mbarrier_init(h0_mbar, 1)
+                cute.arch.mbarrier_init_fence()
+        elif warp_id == 1:
+            cpasync.prefetch_descriptor(H0_tma_atom)
+            cpasync.prefetch_descriptor(W_tma_atom)
+            cpasync.prefetch_descriptor(V_tma_atom)
+            cpasync.prefetch_descriptor(K_tma_atom)
+            cpasync.prefetch_descriptor(HT_tma_atom)
+            cpasync.prefetch_descriptor(H_tma_atom)
+            cpasync.prefetch_descriptor(V_new_tma_atom)
+        cute.arch.sync_threads()
+
+        bos = cu_seqlens[seq_id]
+        eos = cu_seqlens[seq_id + 1]
+        seqlen = eos - bos
+        num_chunks = cute.ceil_div(seqlen, BT)
+
+        if warp_id == 9:
+            # TMA warp
+            stage_id = 0
+            parity = 1
+
+            k_head_id = head_id // (self.Hv // self.H)
+            chunk_offset = chunk_offsets[seq_id]
+
+            # load H0
+            with cute.arch.elect_one():
+                H0_size = V_dim * K_dim * self.h_dtype.width // 8
+                cute.arch.mbarrier_arrive_and_expect_tx(h0_mbar, H0_size)
+            simple_tma_copy(
+                H0_tma_atom, tmaH0[seq_id, head_id, None, None], sH0, h0_mbar
+            )
+
+            # shape: ((BT, num_BT_tiles), (64, 2))
+            gW_tiles = cute.logical_divide(tmaW[None, head_id, None], (BT, None))
+            gV_tiles = cute.logical_divide(tmaV[None, head_id, None], (BT, None))
+            gK_tiles = cute.logical_divide(
+                cute.domain_offset((bos, 0), tmaK[None, k_head_id, None]),
+                (BT, None),
+            )
+
+            for chunk_id in range(num_chunks):
+                mbar = tma_mbar + stage_id
+                gW = gW_tiles[(None, chunk_offset + chunk_id), None]
+                gV = gV_tiles[(None, chunk_offset + chunk_id), None]
+                gK = gK_tiles[(None, chunk_id), None]
+
+                # wait for MMA to release the buffer
+                cute.arch.mbarrier_wait(vk_done_mbar + stage_id, parity)
+
+                # load W, V (i.e. U), and K
+                with cute.arch.elect_one():
+                    STAGE_SIZE = BT * (K_dim + V_dim + K_dim) * 2
+                    cute.arch.mbarrier_arrive_and_expect_tx(mbar, STAGE_SIZE)
+                simple_tma_copy(
+                    W_tma_atom, gW, sW[None, None, stage_id], mbar, EVICT_FIRST
+                )
+                simple_tma_copy(
+                    V_tma_atom, gV, sV[None, None, stage_id], mbar, EVICT_FIRST
+                )
+                simple_tma_copy(K_tma_atom, gK, sK[None, None, stage_id], mbar)
+
+                stage_id = (stage_id + 1) % num_stages
+                if stage_id == 0:
+                    parity ^= 1
+
+        elif warp_id == 8:
+            # MMA warp
+            _tcgen05.alloc(taddr)
+            stage_id = 0
+            parity = 0
+
+            wh_idesc = _tcgen05.make_bf16_idesc(V_dim, BT, negate_A=True)
+            vk_idesc = _tcgen05.make_bf16_idesc(V_dim, K_dim, transpose_B=True)
+
+            # LBO=BT*128 is ignored for K-major
+            sdesc_template = _tcgen05.make_sdesc_128B_swizzle(BT * 128)
+
+            # when using BF16 state, H is read from smem for the 1st iteration
+            # variable names in this conditional branch can't be the same as those
+            # in the mainloop below due to CuteDSL restrictions.
+            if cutlass.const_expr(not is_f32):
+                ##### 1st MMA: V_new.T = V.T - H @ W.T #####
+                Haddr0 = sH0[None, None].iterator.toint()
+                Waddr0 = sW[None, None, stage_id].iterator.toint()
+                hdesc0_base = sdesc_template | (Haddr0 >> 4)
+                wdesc0_base = sdesc_template | (Waddr0 >> 4)
+
+                cute.arch.mbarrier_wait(tma_mbar + stage_id, parity)
+                cute.arch.mbarrier_wait(wh_in_mbar + stage_id, parity)
+                _tcgen05.fence_after_thread_sync()
+
+                with cute.arch.elect_one():
+                    for i in cutlass.range_constexpr(K_dim // 64):
+                        for j in cutlass.range_constexpr(64 // 16):
+                            hdesc0 = hdesc0_base | ((i * V_dim * 128 + j * 32) >> 4)
+                            wdesc0 = wdesc0_base | ((i * BT * 128 + j * 32) >> 4)
+                            _tcgen05.mma_f16(wh_tmem, hdesc0, wdesc0, wh_idesc, True)
+                    _tcgen05.commit(wh_done_mbar + stage_id)
+
+                ##### 2nd MMA: H_new = H + V_new.T @ K #####
+                Kaddr0 = sK[None, None, stage_id].iterator.toint()
+                kdesc0_base = sdesc_template | (Kaddr0 >> 4)
+
+                cute.arch.mbarrier_wait(vk_in_mbar + stage_id, parity)
+                _tcgen05.fence_after_thread_sync()
+
+                with cute.arch.elect_one():
+                    for k in cutlass.range_constexpr(BT // 16):
+                        vtmem0 = v_tmem_base + k * 8
+                        kdesc0 = kdesc0_base | ((k * 16 * 128) >> 4)
+                        _tcgen05.mma_ts_f16(vk_tmem, vtmem0, kdesc0, vk_idesc, True)
+                    _tcgen05.commit(vk_done_mbar + stage_id)
+
+                stage_id = (stage_id + 1) % num_stages
+                if stage_id == 0:
+                    parity ^= 1
+
+            num_iters = num_chunks - int(not is_f32)
+            for _ in range(num_iters):
+                ##### 1st MMA: V_new.T = V.T - H @ W.T #####
+                Waddr = sW[None, None, stage_id].iterator.toint()
+                wdesc_base = sdesc_template | (Waddr >> 4)
+
+                cute.arch.mbarrier_wait(tma_mbar + stage_id, parity)
+                cute.arch.mbarrier_wait(wh_in_mbar + stage_id, parity)
+                _tcgen05.fence_after_thread_sync()
+
+                with cute.arch.elect_one():
+                    for i in cutlass.range_constexpr(K_dim // 64):
+                        for j in cutlass.range_constexpr(64 // 16):
+                            htmem = h_tmem_base + i * 32 + j * 8
+                            wdesc = wdesc_base | ((i * BT * 128 + j * 32) >> 4)
+                            _tcgen05.mma_ts_f16(wh_tmem, htmem, wdesc, wh_idesc, True)
+                    _tcgen05.commit(wh_done_mbar + stage_id)
+
+                ##### 2nd MMA: H_new = H + V_new.T @ K #####
+                Kaddr = sK[None, None, stage_id].iterator.toint()
+                kdesc_base = sdesc_template | (Kaddr >> 4)
+
+                cute.arch.mbarrier_wait(vk_in_mbar + stage_id, parity)
+                _tcgen05.fence_after_thread_sync()
+
+                with cute.arch.elect_one():
+                    for k in cutlass.range_constexpr(BT // 16):
+                        vtmem = v_tmem_base + k * 8
+                        kdesc = kdesc_base | ((k * 16 * 128) >> 4)
+                        _tcgen05.mma_ts_f16(vk_tmem, vtmem, kdesc, vk_idesc, True)
+                    _tcgen05.commit(vk_done_mbar + stage_id)
+
+                stage_id = (stage_id + 1) % num_stages
+                if stage_id == 0:
+                    parity ^= 1
+
+        elif warp_id >= 4:
+            # H warps
+            tid_ = tid % 128
+            warp_id_ = warp_id % 4
+            chunk_offset = chunk_offsets[seq_id]
+
+            stage_id = 0
+            vk_stage_id = 0
+            vk_parity = 0
+
+            op = cute.nvgpu.CopyUniversalOp()
+            cp_16B = cute.make_copy_atom(op, Float32, num_bits_per_copy=128)
+
+            ##### chunk_id = 0 #####
+            if True:
+                chunk_id = 0
+                end_t = min(bos + (chunk_id + 1) * BT, eos)
+                last_idx = end_t - 1
+                h_scale = cute.math.exp(g_cu[last_idx, head_id], fastmath=True)
+
+                # for 1st chunk, wait for H0 transfer from gmem
+                if warp_id_ == 0:
+                    cute.arch.mbarrier_wait(h0_mbar, 0)
+                cute.arch.barrier(barrier_id=1, number_of_threads=128)
+
+                # when H0 is FP32, we need to pack it to BF16
+                # also store to smem for TMA store later.
+                if cutlass.const_expr(is_f32):
+                    for i in cutlass.range_constexpr(K_dim // 32):
+                        # H0 smem layout: (V_dim, (32, K_dim/32))
+                        h_f32 = cute.make_rmem_tensor(32, Float32)
+                        cute.copy(cp_16B, sH0[tid_, (None, i)], h_f32)
+
+                        h_bf16 = cute.make_rmem_tensor(32, BFloat16)
+                        h_bf16.store(h_f32.load().to(BFloat16))
+                        _tcgen05.st(
+                            warp_id_ * 32, h_tmem_base + i * 16, "32x32b", 16, h_bf16
+                        )
+
+                        # H smem layout: (V_dim, (64, K_dim/64))
+                        dst = cute.local_tile(sH[tid_, None], (32,), (i,))
+                        cute.copy(cp_16B, h_bf16, dst)
+
+                _tcgen05.wait_st()
+                _tcgen05.fence_before_thread_sync()
+                cute.arch.mbarrier_arrive(wh_in_mbar + stage_id)
+
+                # scale H for 2nd MMA
+                for i in cutlass.range_constexpr(K_dim // 32):
+                    h_f32 = cute.make_rmem_tensor(32, Float32)
+
+                    if cutlass.const_expr(is_f32):
+                        cute.copy(cp_16B, sH0[tid_, (None, i)], h_f32)
+
+                    else:
+                        h_bf16 = cute.make_rmem_tensor(32, BFloat16)
+                        sH_src = cute.local_tile(sH0[tid_, None], (32,), (i,))
+                        cute.copy(cp_16B, sH_src, h_bf16)
+                        h_f32.store(
+                            cvt.bf16x2_to_fp32x2(
+                                cute.recast_tensor(h_bf16, Uint32)
+                            ).load()
+                        )
+
+                    for j in cutlass.range_constexpr(32):
+                        h_f32[j] *= h_scale
+                    _tcgen05.st(warp_id_ * 32, vk_tmem + i * 32, "32x32b", 32, h_f32)
+
+                _tcgen05.wait_st()
+                _tcgen05.fence_before_thread_sync()
+                cute.arch.mbarrier_arrive(vk_in_mbar + stage_id)
+
+                # for BF16 H0, we issue TMA store from H0 smem
+                # for FP32 H0, we issue TMA store from H smem (after packing)
+                cute.arch.barrier(barrier_id=1, number_of_threads=128)
+                fence_before_tma_store()
+                if warp_id_ == 3:
+                    h_src = sH if cutlass.const_expr(is_f32) else sH0
+                    h_dst = tmaH[chunk_offset + chunk_id, head_id, None, None]
+                    simple_tma_copy(H_tma_atom, h_src, h_dst)
+                    with cute.arch.elect_one():
+                        cute.arch.cp_async_bulk_commit_group()
+
+                        # When H0 is BF16, and there is only 1 chunk, storing
+                        # the final state to sH0 can race before this store
+                        # has finished. hence, we need to wait here.
+                        if cutlass.const_expr(not is_f32):
+                            cute.arch.cp_async_bulk_wait_group(0, read=True)
+
+                stage_id = (stage_id + 1) % num_stages
+
+            ##### subsequent chunks #####
+            for chunk_id in range(1, num_chunks):
+                end_t = min(bos + (chunk_id + 1) * BT, eos)
+                last_idx = end_t - 1
+                h_scale = cute.math.exp(g_cu[last_idx, head_id], fastmath=True)
+
+                # wait for H from previous vk MMA
+                if warp_id_ == 0:
+                    cute.arch.mbarrier_wait(vk_done_mbar + vk_stage_id, vk_parity)
+                    vk_stage_id = (vk_stage_id + 1) % num_stages
+                    if vk_stage_id == 0:
+                        vk_parity ^= 1
+                elif warp_id_ == 3:
+                    with cute.arch.elect_one():
+                        cute.arch.cp_async_bulk_wait_group(0, read=True)
+                cute.arch.barrier(barrier_id=1, number_of_threads=128)
+                _tcgen05.fence_after_thread_sync()
+
+                # load FP32 H from tmem, convert to BF16, store to tmem for 1st MMA,
+                # store to smem for TMA store later.
+                for i in cutlass.range_constexpr(K_dim // 32):
+                    h_f32 = _tcgen05.ld(warp_id_ * 32, vk_tmem + i * 32, "32x32b", 32)
+                    h_bf16 = cute.make_rmem_tensor(32, BFloat16)
+                    h_bf16.store(h_f32.to(BFloat16))
+                    _tcgen05.st(
+                        warp_id_ * 32, h_tmem_base + i * 16, "32x32b", 16, h_bf16
+                    )
+
+                    # H smem layout: (V_dim, (64, K_dim/64))
+                    dst = cute.local_tile(sH[tid_, None], (32,), (i,))
+                    cute.copy(cp_16B, h_bf16, dst)
+
+                _tcgen05.wait_st()
+                _tcgen05.fence_before_thread_sync()
+                cute.arch.mbarrier_arrive(wh_in_mbar + stage_id)
+
+                # scale H for 2nd MMA
+                for i in cutlass.range_constexpr(K_dim // 32):
+                    h_f32 = cute.make_rmem_tensor(32, Float32)
+                    h_f32.store(
+                        _tcgen05.ld(warp_id_ * 32, vk_tmem + i * 32, "32x32b", 32)
+                    )
+                    for j in cutlass.range_constexpr(32):
+                        h_f32[j] *= h_scale
+                    _tcgen05.st(warp_id_ * 32, vk_tmem + i * 32, "32x32b", 32, h_f32)
+                _tcgen05.wait_st()
+                _tcgen05.fence_before_thread_sync()
+                cute.arch.mbarrier_arrive(vk_in_mbar + stage_id)
+
+                # issue TMA store for O kernel
+                cute.arch.barrier(barrier_id=1, number_of_threads=128)
+                fence_before_tma_store()
+                if warp_id_ == 3:
+                    h_dst = tmaH[chunk_offset + chunk_id, head_id, None, None]
+                    simple_tma_copy(H_tma_atom, sH, h_dst)
+                    with cute.arch.elect_one():
+                        cute.arch.cp_async_bulk_commit_group()
+
+                stage_id = (stage_id + 1) % num_stages
+
+            # handle final state. reuse H0 smem.
+            if warp_id_ == 0:
+                cute.arch.mbarrier_wait(vk_done_mbar + vk_stage_id, vk_parity)
+            cute.arch.barrier(barrier_id=1, number_of_threads=128)
+            _tcgen05.fence_after_thread_sync()
+
+            for i in cutlass.range_constexpr(K_dim // 32):
+                h_f32 = cute.make_rmem_tensor(32, Float32)
+                h_f32.store(_tcgen05.ld(warp_id_ * 32, vk_tmem + i * 32, "32x32b", 32))
+
+                if cutlass.const_expr(is_f32):
+                    cute.copy(cp_16B, h_f32, sH0[tid_, (None, i)])
+
+                else:
+                    h_bf16 = cute.make_rmem_tensor(32, BFloat16)
+                    h_bf16.store(h_f32.load().to(BFloat16))
+                    sH0_dst = cute.local_tile(sH0[tid_, None], (32,), (i,))
+                    cute.copy(cp_16B, h_bf16, sH0_dst)
+
+            cute.arch.barrier(barrier_id=1, number_of_threads=128)
+
+            if warp_id_ == 0:
+                ht_dst = tmaHT[seq_id, head_id, None, None]
+                simple_tma_copy(HT_tma_atom, sH0, ht_dst)
+                with cute.arch.elect_one():
+                    cute.arch.cp_async_bulk_commit_group()
+            if warp_id_ == 1:
+                _tcgen05.dealloc()
+
+        else:
+            # V warps
+            stage_id = 0
+            parity = 0
+
+            chunk_offset = chunk_offsets[seq_id]
+
+            ldsm_trans_op = warp.LdMatrix8x8x16bOp(num_matrices=4, transpose=True)
+            stsm_trans_op = warp.StMatrix8x8x16bOp(num_matrices=4, transpose=True)
+            ldsm_trans_atom = cute.make_copy_atom(ldsm_trans_op, BFloat16)
+            stsm_trans_atom = cute.make_copy_atom(stsm_trans_op, BFloat16)
+
+            # ((BT, num_BT_tiles), V_dim)
+            gV_new_tiles = cute.logical_divide(
+                tmaV_new[None, head_id, None], (BT, None)
+            )
+
+            # sV shape: [BT, (64, V_dim/64), num_stages]
+            # sV_view shape: [BT, (8, (8,2)), num_stages]
+            sV_view = cute.logical_divide(sV, (None, 8, None))
+            sV_new_view = cute.logical_divide(sV_new, (None, 8))
+
+            # [BT, 8, num_stages]
+            s_col = warp_id * 4 + (lane_id // 8)
+            sV_view = sV_view[None, (None, s_col), None]
+            sV_new_view = sV_new_view[None, (None, s_col)]
+
+            for chunk_id in range(num_chunks):
+                # wait for V to arrive
+                if warp_id == 0:
+                    cute.arch.mbarrier_wait(tma_mbar + stage_id, parity)
+                cute.arch.barrier(barrier_id=2, number_of_threads=128)
+
+                # unpack V BF16->FP32, then store to tmem for 1st MMA
+                # V smem layout: [BT, (64, V_dim/64)] / [BT, V_dim]
+                # each iteration, CTA loads [8, V_dim] tile
+                # (warp loads [8, 32] tile)
+                for i in cutlass.range_constexpr(BT // 8):
+                    s_row = i * 8 + (lane_id % 8)
+                    v_bf16 = cute.make_rmem_tensor(8, BFloat16)
+                    cute.copy(ldsm_trans_atom, sV_view[s_row, None, stage_id], v_bf16)
+                    v_fp32 = cvt.bf16x2_to_fp32x2(cute.recast_tensor(v_bf16, Uint32))
+                    v_fp32 = cute.logical_divide(v_fp32, 4)  # (4, 2)
+
+                    tcol = wh_tmem + i * 8
+                    _tcgen05.st(warp_id * 32 + 0, tcol, "16x256b", 1, v_fp32[None, 0])
+                    _tcgen05.st(warp_id * 32 + 16, tcol, "16x256b", 1, v_fp32[None, 1])
+
+                _tcgen05.wait_st()
+                _tcgen05.fence_before_thread_sync()
+                cute.arch.mbarrier_arrive(wh_in_mbar + stage_id)
+
+                # load g_cu for scaling
+                if tid < BT:
+                    end_t = min(bos + (chunk_id + 1) * BT, eos)
+                    last_idx = end_t - 1
+                    t = bos + chunk_id * BT + tid
+                    val = Float32(0.0)
+                    if t < eos:
+                        val = cute.math.exp(
+                            g_cu[last_idx, head_id] - g_cu[t, head_id],
+                            fastmath=True,
+                        )
+                    s_v_scale[tid] = val
+
+                # wait for 1st MMA to finish
+                if warp_id == 2:
+                    cute.arch.mbarrier_wait(wh_done_mbar + stage_id, parity)
+                elif warp_id == 3:
+                    with cute.arch.elect_one():
+                        cute.arch.cp_async_bulk_wait_group(0, read=True)
+                cute.arch.barrier(barrier_id=2, number_of_threads=128)
+                _tcgen05.fence_after_thread_sync()
+
+                for i in cutlass.range_constexpr(BT // 8):
+                    v_new = cute.make_rmem_tensor((4, 2), Float32)
+                    tcol = wh_tmem + i * 8
+                    v_new[None, 0].store(
+                        _tcgen05.ld(warp_id * 32 + 0, tcol, "16x256b", 1)
+                    )
+                    v_new[None, 1].store(
+                        _tcgen05.ld(warp_id * 32 + 16, tcol, "16x256b", 1)
+                    )
+                    v_new_bf16 = cute.make_rmem_tensor(8, BFloat16)
+                    v_new_bf16.store(v_new.load().to(BFloat16))
+
+                    # scale V_new for 2nd MMA
+                    scale0 = s_v_scale[i * 8 + (lane_id % 4) * 2 + 0]
+                    scale1 = s_v_scale[i * 8 + (lane_id % 4) * 2 + 1]
+                    v_scaled = cute.make_rmem_tensor(8, Float32)
+                    for k in cutlass.range_constexpr(4):
+                        v_scaled[k * 2] = v_new[k * 2] * scale0
+                        v_scaled[k * 2 + 1] = v_new[k * 2 + 1] * scale1
+                    v_scaled_bf16 = v_scaled.load().to(BFloat16).reshape((4, 2))
+
+                    # store V_new BF16 for O kernel
+                    s_row = i * 8 + (lane_id % 8)
+                    cute.copy(stsm_trans_atom, v_new_bf16, sV_new_view[s_row, None])
+
+                    # store to tmem
+                    tcol = v_tmem_base + i * 4
+                    _tcgen05.st(
+                        warp_id * 32 + 0, tcol, "16x128b", 1, v_scaled_bf16[None, 0]
+                    )
+                    _tcgen05.st(
+                        warp_id * 32 + 16, tcol, "16x128b", 1, v_scaled_bf16[None, 1]
+                    )
+                _tcgen05.wait_st()
+                _tcgen05.fence_before_thread_sync()
+                cute.arch.mbarrier_arrive(vk_in_mbar + stage_id)
+
+                # issue TMA store for V_new
+                cute.arch.barrier(barrier_id=2, number_of_threads=128)
+                fence_before_tma_store()
+                if warp_id == 3:
+                    gV = gV_new_tiles[(None, chunk_offset + chunk_id), None]
+                    simple_tma_copy(V_new_tma_atom, sV_new, gV)
+                    with cute.arch.elect_one():
+                        cute.arch.cp_async_bulk_commit_group()
+
+                stage_id = (stage_id + 1) % num_stages
+                if stage_id == 0:
+                    parity ^= 1
+
+    @cache
+    @staticmethod
+    def compile(
+        H: int,
+        Hv: int,
+        K_dim: int,
+        V_dim: int,
+        h_dtype: cutlass.Numeric = Float32,
+        BT: int = 64,
+        num_stages: int = 2,
+    ):
+        total_t = cute.sym_int()
+        pad_t = cute.sym_int()
+        total_chunks_n = cute.sym_int()
+        num_sequences = cute.sym_int()
+        cu_entries = cute.sym_int()
+
+        K = make_fake_tensor(BFloat16, (total_t, H, K_dim), divisibility=16)
+        V = make_fake_tensor(BFloat16, (pad_t, Hv, V_dim), divisibility=16)
+        W = make_fake_tensor(BFloat16, (pad_t, Hv, K_dim), divisibility=16)
+        V_new = make_fake_tensor(BFloat16, (pad_t, Hv, V_dim), divisibility=16)
+        g_cu = make_fake_tensor(Float32, (total_t, Hv), divisibility=4)
+        h = make_fake_tensor(
+            BFloat16, (total_chunks_n, Hv, V_dim, K_dim), divisibility=16
+        )
+        h0 = make_fake_tensor(
+            h_dtype, (num_sequences, Hv, V_dim, K_dim), divisibility=16
+        )
+        ht = make_fake_tensor(
+            h_dtype, (num_sequences, Hv, V_dim, K_dim), divisibility=16
+        )
+        cu_seqlens = make_fake_tensor(Int32, (cu_entries,), divisibility=1)
+        chunk_offsets = make_fake_tensor(Int32, (cu_entries,), divisibility=1)
+
+        kernel = Sm100ChunkHKernel(H, Hv, K_dim, V_dim, h_dtype, BT, num_stages)
+        stream = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True)
+        return cute.compile(
+            kernel,
+            K,
+            V,
+            W,
+            V_new,
+            g_cu,
+            h,
+            h0,
+            ht,
+            cu_seqlens,
+            chunk_offsets,
+            stream,
+            options="--enable-tvm-ffi",
+        )
+
+
+def h_cutedsl(
+    K: torch.Tensor,
+    V: torch.Tensor,
+    W: torch.Tensor,
+    V_new: torch.Tensor,
+    g_cu: torch.Tensor,
+    h: torch.Tensor,
+    h0: torch.Tensor,
+    ht: torch.Tensor,
+    cu_seqlens: torch.Tensor,
+    chunk_offsets: torch.Tensor,
+    BT: int = 64,
+    num_stages: int = 2,
+) -> None:
+    """Compute H/V_new with the same argument order as the CUDA wrapper."""
+
+    _, H, K_dim = K.shape
+    _, Hv, V_dim = V.shape
+    h_dtype = {
+        torch.bfloat16: BFloat16,
+        torch.float32: Float32,
+    }[h0.dtype]
+    Sm100ChunkHKernel.compile(H, Hv, K_dim, V_dim, h_dtype, BT, num_stages)(
+        K,
+        V,
+        W,
+        V_new,
+        g_cu,
+        h,
+        h0,
+        ht,
+        cu_seqlens,
+        chunk_offsets,
+    )
+
+
+h_v2b_cutedsl = h_cutedsl

vllm/model_executor/layers/mamba/ops/gdn_chunk_cutedsl/kernel_kkt_inv_uw.py

@@ -0,0 +1,832 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from functools import cache
+
+import cutlass
+import torch
+from cuda.bindings.driver import CUstream
+from cutlass import BFloat16, Float32, Int32, Int64, Uint32, cute
+from cutlass.cute.nvgpu import cpasync, warp
+from quack.compile_utils import make_fake_tensor
+
+from vllm.cute_utils import (
+    EVICT_FIRST,
+    _tcgen05,
+    cvt,
+    fence_before_tma_store,
+    mma_bf16,
+    simple_tma_copy,
+)
+
+
+class Sm100ChunkUWKernel:
+    """Compute per-chunk KKT inverse preprocessing and U/W tiles.
+
+    Gamma[i,j] = exp(g_cu[i] - g_cu[j])
+    A = strictLower(beta * (K @ K.T) * Gamma)
+    Ai = inverse(I + A)
+    U = (Ai * beta) @ V
+    W = (Ai * beta * exp(g_cu)) @ K
+    """
+
+    def __init__(
+        self,
+        H: int,
+        Hv: int,
+        K_dim: int,
+        V_dim: int,
+        num_stages: int = 2,
+    ) -> None:
+        assert Hv % H == 0
+        assert K_dim == V_dim == 128
+        self.H = H
+        self.Hv = Hv
+        self.K_dim = K_dim
+        self.V_dim = V_dim
+        self.num_stages = num_stages
+
+        # hard-code
+        self.BT = 64
+        self.num_warps = 2 + 4 + 4
+
+    @cute.jit
+    def _make_tma_args(
+        self,
+        tensor: cute.Tensor,
+        dim: cutlass.Constexpr[int],
+        num_stages: int,
+        op: cpasync.TmaCopyOp,
+    ):
+        # logical layout: [BT, dim]
+        # permute for TMA: [dim/64, BT, 64] with swizzling
+        swizzle_128B = cute.make_swizzle(3, 4, 3)
+        slayout = cute.make_layout(
+            (self.BT, 1, (64, dim // 64), num_stages),
+            stride=(64, 0, (1, self.BT * 64), self.BT * dim),
+        )
+        slayout = cute.make_composed_layout(swizzle_128B, 0, slayout)
+
+        # we need to convert gmem layout to (T, H, (64, D/64)) for make_tiled_tma_atom()
+        # to emit a single 4D TMA. otherwise, it will emit (D/64)x 3D TMA.
+        atom, tma_tensor = cpasync.make_tiled_tma_atom(
+            op,
+            cute.logical_divide(tensor, (None, None, 64)),
+            slayout,
+            cta_tiler=(self.BT, 1, dim),
+        )
+        return atom, tma_tensor, slayout
+
+    @cute.jit
+    def __call__(
+        self,
+        K: cute.Tensor,
+        V: cute.Tensor,
+        U: cute.Tensor,
+        W: cute.Tensor,
+        g: cute.Tensor,
+        beta: cute.Tensor,
+        g_cu: cute.Tensor,
+        cu_seqlens: cute.Tensor,
+        chunk_indices: cute.Tensor,
+        total_chunks: cute.Tensor,
+        num_sms: Int32,
+        stream: CUstream,
+    ):
+        tma_g2s = cpasync.CopyBulkTensorTileG2SOp()
+        tma_s2g = cpasync.CopyBulkTensorTileS2GOp()
+
+        K_args = self._make_tma_args(K, self.K_dim, self.num_stages, tma_g2s)
+        V_args = self._make_tma_args(V, self.V_dim, self.num_stages, tma_g2s)
+        U_args = self._make_tma_args(U, self.V_dim, 1, tma_s2g)
+        W_args = self._make_tma_args(W, self.K_dim, 1, tma_s2g)
+
+        grid = (num_sms // self.Hv, self.Hv, 1)
+        block = (self.num_warps * 32, 1, 1)
+        self.kernel(
+            K_args,
+            V_args,
+            U_args,
+            W_args,
+            g,
+            beta,
+            g_cu,
+            cu_seqlens,
+            chunk_indices,
+            total_chunks,
+        ).launch(grid=grid, block=block, stream=stream)
+
+    @cute.kernel
+    def kernel(
+        self,
+        K_args: tuple[cute.CopyAtom, cute.Tensor, cute.ComposedLayout],
+        V_args: tuple[cute.CopyAtom, cute.Tensor, cute.ComposedLayout],
+        U_args: tuple[cute.CopyAtom, cute.Tensor, cute.ComposedLayout],
+        W_args: tuple[cute.CopyAtom, cute.Tensor, cute.ComposedLayout],
+        g: cute.Tensor,
+        beta: cute.Tensor,
+        g_cu: cute.Tensor,
+        cu_seqlens: cute.Tensor,
+        chunk_indices: cute.Tensor,
+        total_chunks: cute.Tensor,
+    ):
+        tid, _, _ = cute.arch.thread_idx()
+        bid, head_id, _ = cute.arch.block_idx()
+        grid_x, _, _ = cute.arch.grid_dim()
+
+        warp_id = cute.arch.make_warp_uniform(tid // 32)
+        lane_id = tid % 32
+        k_head_id = head_id // (self.Hv // self.H)
+
+        BT = self.BT
+        K_dim = self.K_dim
+        V_dim = self.V_dim
+        num_stages = self.num_stages
+
+        K_tma_atom, tmaK, sK_layout = K_args
+        V_tma_atom, tmaV, sV_layout = V_args
+        U_tma_atom, tmaU, sU_layout = U_args
+        W_tma_atom, tmaW, sW_layout = W_args
+
+        def allocate_tensor(smem, dtype, layout):
+            return smem.allocate_tensor(
+                dtype, layout.outer, byte_alignment=128, swizzle=layout.inner
+            )
+
+        smem = cutlass.utils.SmemAllocator()
+        sK = allocate_tensor(smem, BFloat16, sK_layout)[None, 0, None, None]
+        sV = allocate_tensor(smem, BFloat16, sV_layout)[None, 0, None, None]
+        sU = allocate_tensor(smem, BFloat16, sU_layout)[None, 0, None, 0]
+        sW = allocate_tensor(smem, BFloat16, sW_layout)[None, 0, None, 0]
+
+        swizzle_128B = cute.make_swizzle(3, 4, 3)
+        sA_layout = cute.make_layout((BT, (64, 1)), stride=(64, (1, BT * 64)))
+        sA_layout = cute.make_composed_layout(swizzle_128B, 0, sA_layout)
+        sA = allocate_tensor(smem, BFloat16, sA_layout)
+        sAi = allocate_tensor(smem, BFloat16, sA_layout)
+
+        s_beta = smem.allocate_array(Float32, BT)
+        s_g_cu_exp = smem.allocate_array(Float32, BT)
+        s_g_cu = smem.allocate_array(Float32, BT)
+
+        tma_mbar = smem.allocate_array(Int64, num_stages)
+        mma_kkt_mbar = smem.allocate_array(Int64, num_stages)
+        inv_mbar = smem.allocate_array(Int64, num_stages)
+        mma_u_mbar = smem.allocate_array(Int64, num_stages)
+        mma_w_mbar = smem.allocate_array(Int64, num_stages)
+        epi_mbar = smem.allocate_array(Int64, num_stages)
+        taddr = smem.allocate(Int32, 4)
+
+        kkt_tmem = 0
+        U_tmem_base = kkt_tmem + BT
+        Ab_tmem_base = U_tmem_base + V_dim * num_stages
+        assert Ab_tmem_base + (BT // 2) * num_stages <= 512
+
+        # prepare ldmatrix/stmatrix ops
+        ldsm_op = warp.LdMatrix8x8x16bOp(num_matrices=4)
+        stsm_op = warp.StMatrix8x8x16bOp(num_matrices=4)
+        ldsm_trans_op = warp.LdMatrix8x8x16bOp(num_matrices=4, transpose=True)
+        ldsm_atom = cute.make_copy_atom(ldsm_op, BFloat16)
+        stsm_atom = cute.make_copy_atom(stsm_op, BFloat16)
+        ldsm_trans_atom = cute.make_copy_atom(ldsm_trans_op, BFloat16)
+
+        if warp_id == 0:
+            with cute.arch.elect_one():
+                for i in cutlass.range_constexpr(num_stages):
+                    cute.arch.mbarrier_init(tma_mbar + i, 1)
+                    cute.arch.mbarrier_init(mma_kkt_mbar + i, 1)
+                    cute.arch.mbarrier_init(inv_mbar + i, 128)
+                    cute.arch.mbarrier_init(mma_u_mbar + i, 1)
+                    cute.arch.mbarrier_init(mma_w_mbar + i, 1)
+                    cute.arch.mbarrier_init(epi_mbar + i, 128)
+                cute.arch.mbarrier_init_fence()
+        elif warp_id == 1:
+            cpasync.prefetch_descriptor(K_tma_atom)
+            cpasync.prefetch_descriptor(V_tma_atom)
+            cpasync.prefetch_descriptor(U_tma_atom)
+            cpasync.prefetch_descriptor(W_tma_atom)
+        cute.arch.sync_threads()
+
+        num_global_chunks = total_chunks[0]
+        if warp_id == 9:
+            # TMA warp
+            stage_id = 0
+            parity = 1
+
+            for global_chunk_id in range(bid, num_global_chunks, grid_x):
+                seq_id = chunk_indices[global_chunk_id, 0]
+                chunk_id = chunk_indices[global_chunk_id, 1]
+                bos = cu_seqlens[seq_id]
+
+                # since off_t is not a multiple of BT, we need to use
+                # domain_offset() to shift the pointer first.
+                mbar = tma_mbar + stage_id
+                gK = cute.local_tile(
+                    cute.domain_offset((bos, 0), tmaK[None, k_head_id, None]),
+                    tiler=(BT, K_dim),
+                    coord=(chunk_id, 0),
+                )
+                gV = cute.local_tile(
+                    cute.domain_offset((bos, 0), tmaV[None, head_id, None]),
+                    tiler=(BT, V_dim),
+                    coord=(chunk_id, 0),
+                )
+
+                # when UW MMA is done, K and V TMA buffers are released
+                cute.arch.mbarrier_wait(mma_u_mbar + stage_id, parity)
+
+                with cute.arch.elect_one():
+                    STAGE_SIZE = BT * (K_dim + V_dim) * 2
+                    cute.arch.mbarrier_arrive_and_expect_tx(mbar, STAGE_SIZE)
+                simple_tma_copy(K_tma_atom, gK, sK[None, None, stage_id], mbar)
+                simple_tma_copy(
+                    V_tma_atom, gV, sV[None, None, stage_id], mbar, EVICT_FIRST
+                )
+
+                stage_id = (stage_id + 1) % num_stages
+                if stage_id == 0:
+                    parity ^= 1
+
+        elif warp_id == 8:
+            # MMA warp
+            _tcgen05.alloc(taddr)
+
+            stage_id = 0
+            parity = 0
+
+            kkt_idesc = _tcgen05.make_bf16_idesc(BT, BT)
+            u_idesc = _tcgen05.make_bf16_idesc(BT, V_dim, transpose_B=True)
+            w_idesc = _tcgen05.make_bf16_idesc(BT, K_dim, transpose_B=True)
+
+            # LBO=BT*128 is ignored for K-major
+            sdesc_template = _tcgen05.make_sdesc_128B_swizzle(BT * 128)
+
+            for global_chunk_id in range(bid, num_global_chunks, grid_x):
+                U_tmem = U_tmem_base + V_dim * stage_id
+                W_tmem = U_tmem | (16 << 16)
+                Ab_tmem = Ab_tmem_base + (BT // 2) * stage_id
+                Abg_tmem = Ab_tmem | (16 << 16)
+
+                ##### KKT MMA: KKT = K @ K.T #####
+                kaddr = sK[None, None, stage_id].iterator.toint()
+                kdesc_base = sdesc_template | (kaddr >> 4)
+
+                # wait for TMA data to arrive
+                # kkt tmem is guaranteed to be free as this is issued
+                # after the previous kkt's consumer (inv warps)
+                cute.arch.mbarrier_wait(tma_mbar + stage_id, parity)
+                _tcgen05.fence_after_thread_sync()
+
+                with cute.arch.elect_one():
+                    for i in cutlass.range_constexpr(K_dim // 64):
+                        for j in cutlass.range_constexpr(64 // 16):
+                            kdesc = kdesc_base | ((i * BT * 128 + j * 32) >> 4)
+                            _tcgen05.mma_f16(
+                                kkt_tmem,
+                                kdesc,
+                                kdesc,
+                                kkt_idesc,
+                                (i > 0) or (j > 0),
+                            )
+                    _tcgen05.commit(mma_kkt_mbar + stage_id)
+
+                ##### U/W MMA: U = Ab @ V, W = Abg @ K #####
+                vaddr = sV[None, None, stage_id].iterator.toint()
+                vdesc = sdesc_template | (vaddr >> 4)
+                kdesc = sdesc_template | (kaddr >> 4)
+
+                # wait for epilogue to release tmem buffer
+                cute.arch.mbarrier_wait(epi_mbar + stage_id, parity ^ 1)
+                cute.arch.mbarrier_wait(inv_mbar + stage_id, parity)
+                _tcgen05.fence_after_thread_sync()
+
+                with cute.arch.elect_one():
+                    for i in cutlass.range_constexpr(BT // 16):
+                        _tcgen05.mma_ts_f16(
+                            W_tmem, Abg_tmem + i * 8, kdesc, w_idesc, i > 0
+                        )
+                        kdesc += (16 * 128) >> 4
+                    _tcgen05.commit(mma_w_mbar + stage_id)
+
+                    for i in cutlass.range_constexpr(BT // 16):
+                        _tcgen05.mma_ts_f16(
+                            U_tmem, Ab_tmem + i * 8, vdesc, u_idesc, i > 0
+                        )
+                        vdesc += (16 * 128) >> 4
+                    _tcgen05.commit(mma_u_mbar + stage_id)
+
+                stage_id = (stage_id + 1) % num_stages
+                if stage_id == 0:
+                    parity ^= 1
+
+            cute.arch.mbarrier_wait(epi_mbar + stage_id, parity ^ 1)
+            _tcgen05.dealloc()
+
+        elif warp_id >= 4:
+            # inv warps
+            tid_ = tid % 128
+            warp_id_ = warp_id % 4
+
+            stage_id = 0
+            parity = 0
+
+            # view into (16,16) sub-tiles, then ldmatrix layout
+            sA_ldsm = cute.logical_divide(sA, (16, cute.make_layout((8, 2))))
+            sAi_ldsm = cute.logical_divide(sAi, (16, cute.make_layout((8, 2))))
+            sA_ldsm = sA_ldsm[(lane_id % 16, None), ((None, lane_id // 16), None)]
+            sAi_ldsm = sAi_ldsm[(lane_id % 16, None), ((None, lane_id // 16), None)]
+
+            # init Ai smem buffer with zeros (only the first 48 rows)
+            for i in cutlass.range_constexpr((BT // 4 * 3) * BT // 128):
+                idx = i * 128 + tid_
+                sAi[idx // BT, idx % BT] = BFloat16(0.0)
+
+            # indices for ldmatrix layout later
+            row_indices = cute.make_rmem_tensor((1, 2, 1), Int32)
+            row_indices[0, 0, 0] = warp_id_ * 16 + (lane_id // 4)
+            row_indices[0, 1, 0] = warp_id_ * 16 + (lane_id // 4) + 8
+            row_indices = row_indices.load()
+
+            col_indices = cute.make_rmem_tensor((2, 1, 2), Int32)
+            col_indices[0, 0, 0] = (lane_id % 4) * 2 + 0
+            col_indices[1, 0, 0] = (lane_id % 4) * 2 + 1
+            col_indices[0, 0, 1] = (lane_id % 4) * 2 + 8
+            col_indices[1, 0, 1] = (lane_id % 4) * 2 + 9
+            col_indices = col_indices.load()
+
+            for global_chunk_id in range(bid, num_global_chunks, grid_x):
+                seq_id = chunk_indices[global_chunk_id, 0]
+                chunk_id = chunk_indices[global_chunk_id, 1]
+                bos = cu_seqlens[seq_id]
+                eos = cu_seqlens[seq_id + 1]
+                off_t = bos + chunk_id * BT
+
+                t = off_t + tid_
+
+                ##### Phase 1: load g and beta #####
+                if tid_ < BT:
+                    in_bounds = t < eos
+                    beta_val = beta[t, head_id] if in_bounds else Float32(0.0)
+                    g_val = g[t, head_id] if in_bounds else Float32(0.0)
+
+                    s_beta[tid_] = beta_val
+
+                    # compute cumsum(g)
+                    # parallel scan within a warp
+                    for i in cutlass.range_constexpr(5):
+                        offset = cutlass.const_expr(1 << i)
+                        lower = cute.arch.shuffle_sync_up(
+                            g_val, offset, mask_and_clamp=0
+                        )
+                        if lane_id >= offset:
+                            g_val += lower
+
+                    # store warp sum
+                    if lane_id == 31:
+                        s_g_cu[warp_id_] = g_val
+                    cute.arch.barrier(barrier_id=3, number_of_threads=BT)
+
+                    # add warp sum from lower warps
+                    for i in cutlass.range_constexpr(1, BT // 32):
+                        if warp_id_ >= i:
+                            g_val += s_g_cu[i - 1]
+                    cute.arch.barrier(barrier_id=3, number_of_threads=BT)
+
+                    # store g_cu to gmem for H and O kernels
+                    if in_bounds:
+                        g_cu[t, head_id] = g_val
+
+                    # store g and g_cu to smem for later
+                    s_g_cu[tid_] = g_val
+                    s_g_cu_exp[tid_] = cute.math.exp(g_val) if in_bounds else 0.0
+
+                ##### Phase 2: A = strictLower(beta * kkt * Gamma) #####
+                if warp_id_ == 0:
+                    cute.arch.mbarrier_wait(mma_kkt_mbar + stage_id, parity)
+                cute.arch.barrier(barrier_id=1, number_of_threads=128)
+                _tcgen05.fence_after_thread_sync()
+
+                # tmem 16x256b layout / ldmatrix layout
+                # mode0 is 8 rows together
+                # mode1 is top and bottom 8 rows
+                # mode2 is groups of 16 rows
+                row_coord = (lane_id // 4, None, warp_id_)
+                s_beta_view = cute.make_tensor(s_beta, (8, 2, 4))
+                beta_row = s_beta_view[row_coord].load().reshape((1, 2, 1))
+
+                s_g_cu_view = cute.make_tensor(s_g_cu, (8, 2, 4))
+                g_cu_row = s_g_cu_view[row_coord].load().reshape((1, 2, 1))
+
+                # mode0 is 2 consecutive elems
+                # mode1 is top and bottom 8 rows
+                # mode2 is next 8 columns
+                # mode3 is repeating that 16x16 tile pattern
+                kkt = _tcgen05.ld(kkt_tmem, 0, "16x256b", BT // 8)
+                kkt = kkt.reshape((2, 2, 2, BT // 16))
+
+                for i in cutlass.range_constexpr(BT // 16):
+                    # mode0 is 2 elems next to each other
+                    # mode1 is 4 pairs of elems on 1 row
+                    # mode2 is top and bottom 8 rows
+                    # mode3 is next 16 columns
+                    col_coord = (None, lane_id % 4, None, i)
+                    s_g_cu_view = cute.make_tensor(s_g_cu, (2, 4, 2, BT // 16))
+                    g_cu_col = s_g_cu_view[col_coord].load().reshape((2, 1, 2))
+
+                    Gamma = cute.math.exp(g_cu_row - g_cu_col, fastmath=True)
+                    A = kkt[None, None, None, i] * beta_row * Gamma
+
+                    # strict lower mask
+                    # NOTE: for OOB t position, s_beta is filled with zeros.
+                    # hence, we don't need to apply bounds check for columns.
+                    A_masked = cute.where(row_indices > col_indices + i * 16, A, 0.0)
+
+                    # pack to BF16
+                    # CuteDSL doesn't generate cvt.bf16x2.f32 here for some reasons
+                    packed = cute.make_rmem_tensor(4, Uint32)
+                    packed[0] = cvt.fp32x2_to_bf16x2(
+                        A_masked[0, 0, 0], A_masked[1, 0, 0]
+                    )
+                    packed[1] = cvt.fp32x2_to_bf16x2(
+                        A_masked[0, 1, 0], A_masked[1, 1, 0]
+                    )
+                    packed[2] = cvt.fp32x2_to_bf16x2(
+                        A_masked[0, 0, 1], A_masked[1, 0, 1]
+                    )
+                    packed[3] = cvt.fp32x2_to_bf16x2(
+                        A_masked[0, 1, 1], A_masked[1, 1, 1]
+                    )
+
+                    # store to smem
+                    cute.copy(
+                        stsm_atom,
+                        cute.recast_tensor(packed, BFloat16),
+                        sA_ldsm[warp_id_, None, i],
+                    )
+
+                cute.arch.barrier(barrier_id=1, number_of_threads=128)
+
+                ##### Phase 3: matrix inverse #####
+                # we use Newton-Schulz iterations to compute the inverse
+                # of the four 16x16 diagonal blocks.
+                #   Ai_new = 2 Ai - Ai @ M @ Ai
+                #   where M = I + A
+                #
+                # we do this with 2 MMAs:
+                # 1. -AiM = Ai @ (-M)
+                # 2. Ai_new = 2 Ai + (-AiM) @ Ai
+                zeros_f32 = cute.make_rmem_tensor(4, Float32)
+                zeros_f32.fill(0.0)
+
+                def set_diagonal(A: cute.Tensor, lane_id: Int32):
+                    "Set the diagonal to 1s"
+                    if lane_id % 9 == 0:
+                        A[0] = (A[0] & Uint32(0xFFFF0000)) | Uint32(0x00003F80)
+                        A[3] = (A[3] & Uint32(0xFFFF0000)) | Uint32(0x00003F80)
+                    elif lane_id % 9 == 4:
+                        A[0] = (A[0] & Uint32(0x0000FFFF)) | Uint32(0x3F800000)
+                        A[3] = (A[3] & Uint32(0x0000FFFF)) | Uint32(0x3F800000)
+
+                Ai_bf16 = cute.make_rmem_tensor(8, BFloat16)
+                mma_B_bf16 = cute.make_rmem_tensor(8, BFloat16)
+                M_bf16 = cute.make_rmem_tensor(8, BFloat16)
+                acc = cute.make_rmem_tensor((4, 2), Float32)
+
+                # share the same storage
+                Ai = cute.recast_tensor(Ai_bf16, Uint32)
+                mma_B = cute.logical_divide(cute.recast_tensor(mma_B_bf16, Uint32), 2)
+                M = cute.logical_divide(cute.recast_tensor(M_bf16, Uint32), 2)
+
+                # initial guess: Ai = I-A
+                cute.copy(ldsm_atom, sA_ldsm[warp_id_, None, warp_id_], Ai_bf16)
+                for i in cutlass.range_constexpr(4):
+                    Ai[i] ^= Uint32(0x80008000)  # negate A
+                set_diagonal(Ai, lane_id)
+
+                # (4, 2)
+                Ai_f32 = cute.logical_divide(cvt.bf16x2_to_fp32x2(Ai), 4)
+
+                # M is holding -(I+A), stay constant throughout the iterations
+                cute.copy(ldsm_trans_atom, sA_ldsm[warp_id_, None, warp_id_], M_bf16)
+                set_diagonal(M, lane_id)
+                for i in cutlass.range_constexpr(4):
+                    M[i] ^= Uint32(0x80008000)
+
+                # 3 rounds of Newton-Schulz
+                for _ in cutlass.range_constexpr(3):
+                    # First MMA: -AiM = Ai @ (-M)
+                    cute.copy(stsm_atom, Ai_bf16, sA_ldsm[warp_id_, None, warp_id_])
+                    cute.arch.sync_warp()
+                    acc[None, 0] = mma_bf16(Ai, M[None, 0], zeros_f32)
+                    acc[None, 1] = mma_bf16(Ai, M[None, 1], zeros_f32)
+                    Ai_bf16.store(acc.load().to(BFloat16))
+
+                    # Second MMA: Ai_new = 2Ai + (-AiM) @ Ai
+                    for j in cutlass.range_constexpr(8):
+                        Ai_f32[j] *= 2.0
+                    cute.copy(
+                        ldsm_trans_atom,
+                        sA_ldsm[warp_id_, None, warp_id_],
+                        mma_B_bf16,
+                    )
+                    Ai_f32[None, 0] = mma_bf16(Ai, mma_B[None, 0], Ai_f32[None, 0])
+                    Ai_f32[None, 1] = mma_bf16(Ai, mma_B[None, 1], Ai_f32[None, 1])
+                    Ai_bf16.store(Ai_f32.load().to(BFloat16))
+
+                cute.copy(stsm_atom, Ai_bf16, sAi_ldsm[warp_id_, None, warp_id_])
+                cute.arch.barrier(barrier_id=1, number_of_threads=128)
+
+                # off-diagonal by 1
+                # given
+                # [ Ai00               ]
+                # [  A10 Ai11          ]
+                # [  A20  A21 Ai22     ]
+                # [  A30  A31  A32 Ai33]
+                # warp1: Ai10 = -Ai11 @ A10 @ Ai00
+                # warp2: Ai21 = -Ai22 @ A21 @ Ai11
+                # warp3: Ai32 = -Ai33 @ A32 @ Ai22
+                if warp_id_ > 0:
+                    neg_Ai = cute.make_rmem_tensor(4, Uint32)
+                    for i in cutlass.range_constexpr(4):
+                        neg_Ai[i] = Ai[i] ^ Uint32(0x80008000)
+
+                    cute.copy(
+                        ldsm_trans_atom,
+                        sA_ldsm[warp_id_, None, warp_id_ - 1],
+                        mma_B_bf16,
+                    )
+                    acc[None, 0] = mma_bf16(neg_Ai, mma_B[None, 0], zeros_f32)
+                    acc[None, 1] = mma_bf16(neg_Ai, mma_B[None, 1], zeros_f32)
+                    Ai_bf16.store(acc.load().to(BFloat16))
+
+                    cute.copy(
+                        ldsm_trans_atom,
+                        sAi_ldsm[warp_id_ - 1, None, warp_id_ - 1],
+                        mma_B_bf16,
+                    )
+                    acc[None, 0] = mma_bf16(Ai, mma_B[None, 0], zeros_f32)
+                    acc[None, 1] = mma_bf16(Ai, mma_B[None, 1], zeros_f32)
+                    Ai_bf16.store(acc.load().to(BFloat16))
+                    cute.copy(
+                        stsm_atom,
+                        Ai_bf16,
+                        sAi_ldsm[warp_id_, None, warp_id_ - 1],
+                    )
+                cute.arch.barrier(barrier_id=1, number_of_threads=128)
+
+                # off-diagonal by 2
+                # warp0: Ai20 = -Ai22 @ (A20 @ Ai00 + A21 @ Ai10)
+                # warp1: Ai31 = -Ai33 @ (A31 @ Ai11 + A32 @ Ai21)
+                if warp_id_ < 2:
+                    cute.copy(
+                        ldsm_atom,
+                        sA_ldsm[warp_id_ + 2, None, warp_id_],
+                        Ai_bf16,
+                    )
+                    cute.copy(
+                        ldsm_trans_atom,
+                        sAi_ldsm[warp_id_, None, warp_id_],
+                        mma_B_bf16,
+                    )
+                    acc[None, 0] = mma_bf16(Ai, mma_B[None, 0], zeros_f32)
+                    acc[None, 1] = mma_bf16(Ai, mma_B[None, 1], zeros_f32)
+
+                    cute.copy(
+                        ldsm_atom,
+                        sA_ldsm[warp_id_ + 2, None, warp_id_ + 1],
+                        Ai_bf16,
+                    )
+                    cute.copy(
+                        ldsm_trans_atom,
+                        sAi_ldsm[warp_id_ + 1, None, warp_id_],
+                        mma_B_bf16,
+                    )
+                    acc[None, 0] = mma_bf16(Ai, mma_B[None, 0], acc[None, 0])
+                    acc[None, 1] = mma_bf16(Ai, mma_B[None, 1], acc[None, 1])
+
+                    tmp = cute.make_rmem_tensor(8, BFloat16)
+                    tmp.store(acc.load().to(BFloat16))
+                    cute.copy(stsm_atom, tmp, sAi_ldsm[warp_id_ + 2, None, warp_id_])
+                    cute.arch.sync_warp()
+
+                    cute.copy(
+                        ldsm_atom, sAi_ldsm[warp_id_ + 2, None, warp_id_ + 2], Ai_bf16
+                    )
+                    for i in cutlass.range_constexpr(4):
+                        Ai[i] ^= Uint32(0x80008000)
+                    cute.copy(
+                        ldsm_trans_atom,
+                        sAi_ldsm[warp_id_ + 2, None, warp_id_],
+                        mma_B_bf16,
+                    )
+                    acc[None, 0] = mma_bf16(Ai, mma_B[None, 0], zeros_f32)
+                    acc[None, 1] = mma_bf16(Ai, mma_B[None, 1], zeros_f32)
+                    tmp.store(acc.load().to(BFloat16))
+                    cute.copy(stsm_atom, tmp, sAi_ldsm[warp_id_ + 2, None, warp_id_])
+                cute.arch.barrier(barrier_id=1, number_of_threads=128)
+
+                # off-diagonal by 3
+                # warp0: Ai30 = -Ai33 @ (A30 @ Ai00 + A31 @ Ai10 + A32 @ Ai20)
+                if warp_id_ == 0:
+                    cute.copy(ldsm_atom, sA_ldsm[3, None, 0], Ai_bf16)
+                    cute.copy(ldsm_trans_atom, sAi_ldsm[0, None, 0], mma_B_bf16)
+                    acc[None, 0] = mma_bf16(Ai, mma_B[None, 0], zeros_f32)
+                    acc[None, 1] = mma_bf16(Ai, mma_B[None, 1], zeros_f32)
+
+                    for i in cutlass.range_constexpr(1, 3):
+                        cute.copy(ldsm_atom, sA_ldsm[3, None, i], Ai_bf16)
+                        cute.copy(ldsm_trans_atom, sAi_ldsm[i, None, 0], mma_B_bf16)
+                        acc[None, 0] = mma_bf16(Ai, mma_B[None, 0], acc[None, 0])
+                        acc[None, 1] = mma_bf16(Ai, mma_B[None, 1], acc[None, 1])
+
+                    tmp = cute.make_rmem_tensor(8, BFloat16)
+                    tmp.store(acc.load().to(BFloat16))
+                    cute.copy(stsm_atom, tmp, sAi_ldsm[3, None, 0])
+                    cute.arch.sync_warp()
+
+                    cute.copy(ldsm_atom, sAi_ldsm[3, None, 3], Ai_bf16)
+                    for i in cutlass.range_constexpr(4):
+                        Ai[i] ^= Uint32(0x80008000)
+                    cute.copy(ldsm_trans_atom, sAi_ldsm[3, None, 0], mma_B_bf16)
+                    acc[None, 0] = mma_bf16(Ai, mma_B[None, 0], zeros_f32)
+                    acc[None, 1] = mma_bf16(Ai, mma_B[None, 1], zeros_f32)
+                    tmp.store(acc.load().to(BFloat16))
+                    cute.copy(stsm_atom, tmp, sAi_ldsm[3, None, 0])
+
+                ##### Phase 4: compute Ab, Abg #####
+                if warp_id_ == 3:
+                    cute.arch.mbarrier_wait(mma_u_mbar + stage_id, parity ^ 1)
+                cute.arch.barrier(barrier_id=1, number_of_threads=128)
+
+                for i in cutlass.range_constexpr(BT // 16):
+                    cute.copy(ldsm_atom, sAi_ldsm[warp_id_, None, i], Ai_bf16)
+
+                    col_coord = (None, lane_id % 4, None, i)
+                    s_beta_view = cute.make_tensor(s_beta, (2, 4, 2, BT // 16))
+                    beta_col = s_beta_view[col_coord].load().reshape((2, 1, 2))
+
+                    s_g_cu_view = cute.make_tensor(s_g_cu_exp, (2, 4, 2, BT // 16))
+                    g_cu_col = s_g_cu_view[col_coord].load().reshape((2, 1, 2))
+
+                    Ai_f32 = cvt.bf16x2_to_fp32x2(Ai).load().reshape((2, 2, 2))
+
+                    Ab_f32 = Ai_f32 * beta_col
+                    Ab = Ab_f32.to(BFloat16)
+                    Ab_tmem = Ab_tmem_base + (BT // 2) * stage_id + i * 8
+                    _tcgen05.st(warp_id_ * 32, Ab_tmem, "16x128b", 2, Ab)
+
+                    Abg_f32 = Ab_f32 * g_cu_col
+                    Abg = Abg_f32.to(BFloat16)
+                    _tcgen05.st(warp_id_ * 32 + 16, Ab_tmem, "16x128b", 2, Abg)
+
+                _tcgen05.wait_st()
+                _tcgen05.fence_before_thread_sync()
+                cute.arch.mbarrier_arrive(inv_mbar + stage_id)
+
+                stage_id = (stage_id + 1) % num_stages
+                if stage_id == 0:
+                    parity ^= 1
+
+        elif warp_id < 4:
+            # epi warps
+            stage_id = 0
+            parity = 0
+
+            # ((BT, num_global_chunks), V_dim)
+            gU_tiles = cute.logical_divide(tmaU[None, head_id, None], (BT, None))
+            gW_tiles = cute.logical_divide(tmaW[None, head_id, None], (BT, None))
+
+            # sW shape: [BT, (64, K_dim/64)]
+            # sW_view shape: [(8, 2), (4, K_dim/64)]
+            s_row = warp_id * 16 + lane_id % 16  # select the rows of [16,16] tile
+            sW_view = cute.zipped_divide(
+                sW[s_row, None],
+                tiler=cute.make_layout((8, 2)),
+            )
+            sU_view = cute.zipped_divide(
+                sU[s_row, None],
+                tiler=cute.make_layout((8, 2)),
+            )
+
+            # select the 8 columns within [16,16] tile
+            sW_view = sW_view[(None, lane_id // 16), None]
+            sU_view = sU_view[(None, lane_id // 16), None]
+
+            for global_chunk_id in range(bid, num_global_chunks, grid_x):
+                # wait for W MMA + previous TMA store to finish
+                U_tmem = U_tmem_base + V_dim * stage_id
+                if warp_id == 0:
+                    cute.arch.mbarrier_wait(mma_w_mbar + stage_id, parity)
+                elif warp_id == 1:
+                    with cute.arch.elect_one():
+                        cute.arch.cp_async_bulk_wait_group(0, read=True)
+                cute.arch.barrier(barrier_id=2, number_of_threads=128)
+                _tcgen05.fence_after_thread_sync()
+
+                w_f32 = _tcgen05.ld(warp_id * 32 + 16, U_tmem, "16x256b", K_dim // 8)
+                _tcgen05.wait_ld()
+                w_bf16 = cute.make_rmem_tensor((8, K_dim // 16), BFloat16)
+                w_bf16.store(w_f32.to(BFloat16))
+                cute.copy(stsm_atom, w_bf16, sW_view)
+
+                # wait for U MMA + issue W TMA store
+                cute.arch.barrier(barrier_id=2, number_of_threads=128)
+                fence_before_tma_store()
+                if warp_id == 0:
+                    cute.arch.mbarrier_wait(mma_u_mbar + stage_id, parity)
+                elif warp_id == 1:
+                    # don't need to commit
+                    simple_tma_copy(
+                        W_tma_atom, sW, gW_tiles[(None, global_chunk_id), None]
+                    )
+                cute.arch.barrier(barrier_id=2, number_of_threads=128)
+                _tcgen05.fence_after_thread_sync()
+
+                u_f32 = _tcgen05.ld(warp_id * 32, U_tmem, "16x256b", V_dim // 8)
+                _tcgen05.wait_ld()
+                _tcgen05.fence_before_thread_sync()
+                cute.arch.mbarrier_arrive(epi_mbar + stage_id)
+                u_bf16 = cute.make_rmem_tensor((8, V_dim // 16), BFloat16)
+                u_bf16.store(u_f32.to(BFloat16))
+                cute.copy(stsm_atom, u_bf16, sU_view)
+
+                cute.arch.barrier(barrier_id=2, number_of_threads=128)
+                fence_before_tma_store()
+                if warp_id == 1:
+                    simple_tma_copy(
+                        U_tma_atom, sU, gU_tiles[(None, global_chunk_id), None]
+                    )
+                    with cute.arch.elect_one():
+                        cute.arch.cp_async_bulk_commit_group()
+
+                stage_id = (stage_id + 1) % num_stages
+                if stage_id == 0:
+                    parity ^= 1
+
+    @cache
+    @staticmethod
+    def compile(H: int, Hv: int, K_dim: int, V_dim: int, num_stages: int = 2):
+        total_t = cute.sym_int()
+        pad_t = cute.sym_int()
+        total_chunks_n = cute.sym_int()
+        num_sequences = cute.sym_int()
+
+        K = make_fake_tensor(BFloat16, (total_t, H, K_dim), divisibility=16)
+        V = make_fake_tensor(BFloat16, (total_t, Hv, V_dim), divisibility=16)
+        U = make_fake_tensor(BFloat16, (pad_t, Hv, V_dim), divisibility=16)
+        W = make_fake_tensor(BFloat16, (pad_t, Hv, K_dim), divisibility=16)
+        g = make_fake_tensor(Float32, (total_t, Hv), divisibility=4)
+        beta = make_fake_tensor(Float32, (total_t, Hv), divisibility=4)
+        g_cu = make_fake_tensor(Float32, (total_t, Hv), divisibility=4)
+        cu_seqlens = make_fake_tensor(Int32, (num_sequences,), divisibility=1)
+        chunk_indices = make_fake_tensor(Int32, (total_chunks_n, 2), divisibility=2)
+        total_chunks = make_fake_tensor(Int32, (1,), divisibility=1)
+
+        kernel = Sm100ChunkUWKernel(H, Hv, K_dim, V_dim, num_stages)
+        stream = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True)
+        return cute.compile(
+            kernel,
+            K,
+            V,
+            U,
+            W,
+            g,
+            beta,
+            g_cu,
+            cu_seqlens,
+            chunk_indices,
+            total_chunks,
+            Int32(148),
+            stream,
+            options="--enable-tvm-ffi",
+        )
+
+
+def kkt_inv_uw_cutedsl(
+    K: torch.Tensor,
+    V: torch.Tensor,
+    U: torch.Tensor,
+    W: torch.Tensor,
+    g: torch.Tensor,
+    beta: torch.Tensor,
+    g_cu: torch.Tensor,
+    cu_seqlens: torch.Tensor,
+    chunk_indices: torch.Tensor,
+    total_chunks: torch.Tensor,
+    num_sms: int = 148,
+) -> None:
+    _, Hv, V_dim = V.shape
+    _, H, K_dim = K.shape
+
+    Sm100ChunkUWKernel.compile(H, Hv, K_dim, V_dim)(
+        K,
+        V,
+        U,
+        W,
+        g,
+        beta,
+        g_cu,
+        cu_seqlens,
+        chunk_indices,
+        total_chunks,
+        num_sms,
+    )

vllm/model_executor/layers/mamba/ops/gdn_chunk_cutedsl/kernel_o.py

@@ -0,0 +1,630 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from functools import cache
+
+import cutlass
+import torch
+from cuda.bindings.driver import CUstream
+from cutlass import BFloat16, Float32, Int32, Int64, Uint32, cute
+from cutlass.cute.nvgpu import cpasync, warp
+from quack.compile_utils import make_fake_tensor
+
+from vllm.cute_utils import (
+    EVICT_FIRST,
+    _tcgen05,
+    cvt,
+    fence_before_tma_store,
+    simple_tma_copy,
+)
+
+
+class Sm100ChunkOKernel:
+    """Compute per-token output from recurrent and intra-chunk terms.
+
+    Gamma[i,j] = exp(g_cu[i] - g_cu[j])
+    P = mask((Q @ K.T) * Gamma)
+    O = scale * (exp(g_cu) * (Q @ H.T) + P @ V)
+    """
+
+    def __init__(
+        self,
+        H: int,
+        Hv: int,
+        K_dim: int,
+        V_dim: int,
+        BT: int = 64,
+        num_stages: int = 2,
+    ) -> None:
+        assert Hv % H == 0
+        assert K_dim == 128
+        assert V_dim == 128
+        assert BT == 64
+        self.H = H
+        self.Hv = Hv
+        self.K_dim = K_dim
+        self.V_dim = V_dim
+        self.BT = BT
+        self.num_stages = num_stages
+        self.num_warps = 10
+
+    @cute.jit
+    def _make_bf16_tma_args(
+        self,
+        tensor: cute.Tensor,
+        dim: cutlass.Constexpr[int],
+        op: cpasync.TmaCopyOp,
+        stages: cutlass.Constexpr[int],
+    ):
+        swizzle_128B = cute.make_swizzle(3, 4, 3)
+        slayout = cute.make_layout(
+            (self.BT, 1, (64, dim // 64), stages),
+            stride=(64, 0, (1, self.BT * 64), self.BT * dim),
+        )
+        slayout = cute.make_composed_layout(swizzle_128B, 0, slayout)
+        atom, tma_tensor = cpasync.make_tiled_tma_atom(
+            op,
+            cute.logical_divide(tensor, (None, None, 64)),
+            slayout,
+            cta_tiler=(self.BT, 1, dim),
+        )
+        return atom, tma_tensor, slayout
+
+    @cute.jit
+    def _make_h_tma_args(
+        self,
+        tensor: cute.Tensor,
+        op: cpasync.TmaCopyOp,
+        stages: cutlass.Constexpr[int],
+    ):
+        num_elems = 128 // (tensor.element_type.width // 8)
+        swizzle_128B = cute.make_swizzle(3, 4, 3)
+        slayout = cute.make_layout(
+            (1, self.V_dim, (num_elems, self.K_dim // num_elems), stages),
+            stride=(0, num_elems, (1, self.V_dim * num_elems), self.V_dim * self.K_dim),
+        )
+        slayout = cute.make_composed_layout(swizzle_128B, 0, slayout)
+        atom, tma_tensor = cpasync.make_tiled_tma_atom(
+            op,
+            cute.logical_divide(tensor, (None, None, num_elems)),
+            slayout,
+            cta_tiler=(1, self.V_dim, self.K_dim),
+        )
+        return atom, tma_tensor, slayout
+
+    @cute.jit
+    def __call__(
+        self,
+        q: cute.Tensor,
+        k: cute.Tensor,
+        v_new_chunks: cute.Tensor,
+        h: cute.Tensor,
+        g_cu: cute.Tensor,
+        o: cute.Tensor,
+        cu_seqlens: cute.Tensor,
+        chunk_indices: cute.Tensor,
+        total_chunks: cute.Tensor,
+        scale: Float32,
+        num_sms: Int32,
+        stream: CUstream,
+    ):
+        grid = (num_sms // self.Hv, self.Hv, 1)
+        block = (self.num_warps * 32, 1, 1)
+        tma_g2s = cpasync.CopyBulkTensorTileG2SOp()
+        tma_s2g = cpasync.CopyBulkTensorTileS2GOp()
+        Q_args = self._make_bf16_tma_args(q, self.K_dim, tma_g2s, self.num_stages)
+        K_args = self._make_bf16_tma_args(k, self.K_dim, tma_g2s, self.num_stages)
+        V_args = self._make_bf16_tma_args(
+            v_new_chunks, self.V_dim, tma_g2s, self.num_stages
+        )
+        H_args = self._make_h_tma_args(h, tma_g2s, self.num_stages)
+        O_args = self._make_bf16_tma_args(o, self.V_dim, tma_s2g, 1)
+        self.kernel(
+            Q_args,
+            K_args,
+            V_args,
+            H_args,
+            O_args,
+            g_cu,
+            o,
+            cu_seqlens,
+            chunk_indices,
+            total_chunks,
+            scale,
+        ).launch(grid=grid, block=block, stream=stream)
+
+    @cute.kernel
+    def kernel(
+        self,
+        Q_args: tuple[cute.CopyAtom, cute.Tensor, cute.ComposedLayout],
+        K_args: tuple[cute.CopyAtom, cute.Tensor, cute.ComposedLayout],
+        V_args: tuple[cute.CopyAtom, cute.Tensor, cute.ComposedLayout],
+        H_args: tuple[cute.CopyAtom, cute.Tensor, cute.ComposedLayout],
+        O_args: tuple[cute.CopyAtom, cute.Tensor, cute.ComposedLayout],
+        g_cu: cute.Tensor,
+        o: cute.Tensor,
+        cu_seqlens: cute.Tensor,
+        chunk_indices: cute.Tensor,
+        total_chunks: cute.Tensor,
+        scale: Float32,
+    ):
+        tid, _, _ = cute.arch.thread_idx()
+        bid, v_head_id, _ = cute.arch.block_idx()
+        grid_x, _, _ = cute.arch.grid_dim()
+        warp_id = cute.arch.make_warp_uniform(tid // 32)
+        lane_id = tid % 32
+
+        BT = self.BT
+        K_dim = self.K_dim
+        V_dim = self.V_dim
+        num_stages = self.num_stages
+
+        heads_per_qk = self.Hv // self.H
+        k_head_id = v_head_id // heads_per_qk
+        num_global_chunks = total_chunks[0]
+
+        Q_tma_atom, tmaQ, sQ_layout = Q_args
+        K_tma_atom, tmaK, sK_layout = K_args
+        V_tma_atom, tmaV, sV_layout = V_args
+        H_tma_atom, tmaH, sH_layout = H_args
+        O_tma_atom, tmaO, sO_layout = O_args
+
+        def allocate_tensor(smem, dtype, layout):
+            return smem.allocate_tensor(
+                dtype, layout.outer, byte_alignment=128, swizzle=layout.inner
+            )
+
+        smem = cutlass.utils.SmemAllocator()
+        sQ = allocate_tensor(smem, BFloat16, sQ_layout)[None, 0, None, None]
+        sK = allocate_tensor(smem, BFloat16, sK_layout)[None, 0, None, None]
+        sV = allocate_tensor(smem, BFloat16, sV_layout)[None, 0, None, None]
+        sH = allocate_tensor(smem, BFloat16, sH_layout)[0, None, None, None]
+        sO = allocate_tensor(smem, BFloat16, sO_layout)[None, 0, None, 0]
+
+        s_g_cu = smem.allocate_array(Float32, BT)
+        qk_full_mbar = smem.allocate_array(Int64, num_stages)
+        hv_full_mbar = smem.allocate_array(Int64, num_stages)
+        qk_empty_mbar = smem.allocate_array(Int64, num_stages)
+        pv_mma_mbar = smem.allocate_array(Int64, num_stages)
+        qk_mbar = smem.allocate_array(Int64, 1)
+        mask_mbar = smem.allocate_array(Int64, 1)
+        epi_mbar = smem.allocate_array(Int64, 1)
+        taddr = smem.allocate(Int32, 4)
+
+        qk_tmem = 0
+        p_tmem = 64
+        out_tmem = 128
+        qh_tmem = 256
+
+        if warp_id == 0:
+            with cute.arch.elect_one():
+                for i in cutlass.range_constexpr(num_stages):
+                    cute.arch.mbarrier_init(qk_full_mbar + i, 1)
+                    cute.arch.mbarrier_init(qk_empty_mbar + i, 1)
+                    cute.arch.mbarrier_init(hv_full_mbar + i, 1)
+                    cute.arch.mbarrier_init(pv_mma_mbar + i, 1)
+                cute.arch.mbarrier_init(qk_mbar, 1)
+                cute.arch.mbarrier_init(mask_mbar, 128)
+                cute.arch.mbarrier_init(epi_mbar, 128)
+                cute.arch.mbarrier_init_fence()
+        elif warp_id == 9:
+            cpasync.prefetch_descriptor(Q_tma_atom)
+            cpasync.prefetch_descriptor(K_tma_atom)
+            cpasync.prefetch_descriptor(V_tma_atom)
+            cpasync.prefetch_descriptor(H_tma_atom)
+        cute.arch.sync_threads()
+
+        if warp_id == 9:
+            # TMA warp
+            stage_id = 0
+            parity = 1
+
+            for global_chunk_id in range(bid, num_global_chunks, grid_x):
+                seq_id = chunk_indices[global_chunk_id, 0]
+                chunk_id = chunk_indices[global_chunk_id, 1]
+                bos = cu_seqlens[seq_id]
+
+                # copy Q and K
+                q_tile = cute.local_tile(
+                    cute.domain_offset((bos, 0), tmaQ[None, k_head_id, None]),
+                    tiler=(BT, K_dim),
+                    coord=(chunk_id, 0),
+                )
+                k_tile = cute.local_tile(
+                    cute.domain_offset((bos, 0), tmaK[None, k_head_id, None]),
+                    tiler=(BT, K_dim),
+                    coord=(chunk_id, 0),
+                )
+                mbar = qk_full_mbar + stage_id
+
+                cute.arch.mbarrier_wait(qk_empty_mbar + stage_id, parity)
+
+                with cute.arch.elect_one():
+                    STAGE_SIZE = BT * (K_dim + K_dim) * 2
+                    cute.arch.mbarrier_arrive_and_expect_tx(mbar, STAGE_SIZE)
+                simple_tma_copy(Q_tma_atom, q_tile, sQ[None, None, stage_id], mbar)
+                simple_tma_copy(K_tma_atom, k_tile, sK[None, None, stage_id], mbar)
+
+                # copy H and V
+                gH = tmaH[global_chunk_id * self.Hv + v_head_id, None, None]
+                gV = cute.local_tile(
+                    tmaV[None, v_head_id, None],
+                    tiler=(BT, V_dim),
+                    coord=(global_chunk_id, 0),
+                )
+                mbar = hv_full_mbar + stage_id
+
+                cute.arch.mbarrier_wait(pv_mma_mbar + stage_id, parity)
+
+                with cute.arch.elect_one():
+                    H_STAGE_SIZE = V_dim * K_dim * 2
+                    V_STAGE_SIZE = BT * V_dim * 2
+                    cute.arch.mbarrier_arrive_and_expect_tx(
+                        mbar, H_STAGE_SIZE + V_STAGE_SIZE
+                    )
+                simple_tma_copy(
+                    H_tma_atom, gH, sH[None, None, stage_id], mbar, EVICT_FIRST
+                )
+                simple_tma_copy(
+                    V_tma_atom, gV, sV[None, None, stage_id], mbar, EVICT_FIRST
+                )
+
+                stage_id = (stage_id + 1) % num_stages
+                if stage_id == 0:
+                    parity ^= 1
+
+        elif warp_id == 8:
+            # MMA warp
+            _tcgen05.alloc(taddr)
+
+            # LBO=BT*128 is ignored for K-major
+            sdesc_template = _tcgen05.make_sdesc_128B_swizzle(BT * 128)
+            qk_idesc = _tcgen05.make_bf16_idesc(BT, BT)
+            qh_idesc = _tcgen05.make_bf16_idesc(BT, V_dim)
+            pv_idesc = _tcgen05.make_bf16_idesc(BT, V_dim, transpose_B=True)
+
+            stage_id = 0
+            tma_parity = 0
+            mask_parity = 0
+
+            for global_chunk_id in range(bid, num_global_chunks, grid_x):
+                qaddr = sQ[None, None, stage_id].iterator.toint()
+                kaddr = sK[None, None, stage_id].iterator.toint()
+                haddr = sH[None, None, stage_id].iterator.toint()
+                vaddr = sV[None, None, stage_id].iterator.toint()
+                qdesc_base = sdesc_template | (qaddr >> 4)
+                kdesc_base = sdesc_template | (kaddr >> 4)
+                hdesc_base = sdesc_template | (haddr >> 4)
+                vdesc_base = sdesc_template | (vaddr >> 4)
+
+                ##### 1st MMA: Q @ K.T #####
+                # do this first to unblock mask(QK)
+                cute.arch.mbarrier_wait(epi_mbar, mask_parity ^ 1)
+                cute.arch.mbarrier_wait(qk_full_mbar + stage_id, tma_parity)
+                _tcgen05.fence_after_thread_sync()
+
+                with cute.arch.elect_one():
+                    for i in cutlass.range_constexpr(K_dim // BT):
+                        for j in cutlass.range_constexpr(BT // 16):
+                            qdesc = qdesc_base | ((i * BT * 128 + j * 32) >> 4)
+                            kdesc = kdesc_base | ((i * BT * 128 + j * 32) >> 4)
+                            _tcgen05.mma_f16(
+                                qk_tmem, qdesc, kdesc, qk_idesc, (i > 0) or (j > 0)
+                            )
+                    _tcgen05.commit(qk_mbar)
+
+                ##### 2nd MMA: Q @ H.T #####
+                cute.arch.mbarrier_wait(hv_full_mbar + stage_id, tma_parity)
+                _tcgen05.fence_after_thread_sync()
+                with cute.arch.elect_one():
+                    for i in cutlass.range_constexpr(K_dim // BT):
+                        for j in cutlass.range_constexpr(BT // 16):
+                            qdesc = qdesc_base | ((i * BT * 128 + j * 32) >> 4)
+                            hdesc = hdesc_base | ((i * V_dim * 128 + j * 32) >> 4)
+                            _tcgen05.mma_f16(
+                                qh_tmem, qdesc, hdesc, qh_idesc, (i > 0) or (j > 0)
+                            )
+                    _tcgen05.commit(qk_empty_mbar + stage_id)
+
+                ##### 3rd MMA: P @ V #####
+                # stalled by mask(QK)
+                cute.arch.mbarrier_wait(mask_mbar, mask_parity)
+                _tcgen05.fence_after_thread_sync()
+                with cute.arch.elect_one():
+                    for i in cutlass.range_constexpr(BT // 16):
+                        vdesc = vdesc_base | ((i * 16 * 128) >> 4)
+                        _tcgen05.mma_ts_f16(
+                            out_tmem, p_tmem + i * 8, vdesc, pv_idesc, i > 0
+                        )
+                    _tcgen05.commit(pv_mma_mbar + stage_id)
+
+                stage_id = (stage_id + 1) % num_stages
+                if stage_id == 0:
+                    tma_parity ^= 1
+                mask_parity ^= 1
+
+            # wait for epilogue to finish for deallocation
+            cute.arch.mbarrier_wait(epi_mbar, mask_parity ^ 1)
+            _tcgen05.dealloc()
+
+        elif warp_id >= 4:
+            # masking warps
+            warp_id_ = warp_id % 4
+            tid_ = tid % 128
+            row0 = warp_id_ * 16 + lane_id // 4
+            row1 = row0 + 8
+
+            parity = 0
+
+            # for ldmatrix layout later
+            row_indices = cute.make_rmem_tensor(2, Int32)
+            row_indices[0] = warp_id_ * 16 + lane_id // 4
+            row_indices[1] = warp_id_ * 16 + lane_id // 4 + 8
+            row_indices = row_indices.load().reshape((1, 2))
+
+            col_indices = cute.make_rmem_tensor(2, Int32)
+            col_indices[0] = (lane_id % 4) * 2
+            col_indices[1] = (lane_id % 4) * 2 + 1
+            col_indices = col_indices.load().reshape((2, 1))
+
+            for global_chunk_id in range(bid, num_global_chunks, grid_x):
+                if tid_ < BT:
+                    seq_id = chunk_indices[global_chunk_id, 0]
+                    chunk_id = chunk_indices[global_chunk_id, 1]
+                    bos = cu_seqlens[seq_id]
+                    eos = cu_seqlens[seq_id + 1]
+
+                    t_ = bos + chunk_id * BT + tid_
+                    s_g_cu[tid_] = g_cu[t_, v_head_id] if t_ < eos else Float32(0.0)
+
+                # wait for QK MMA
+                if warp_id_ == 0:
+                    cute.arch.mbarrier_wait(qk_mbar, parity)
+                cute.arch.barrier(barrier_id=1, number_of_threads=128)
+                _tcgen05.fence_after_thread_sync()
+                qk = _tcgen05.ld(warp_id_ * 32, qk_tmem, "16x256b", BT // 8)
+                qk = qk.reshape((2, 2, BT // 8))
+                _tcgen05.wait_ld()
+
+                g_cu_rows = cute.make_rmem_tensor(2, Float32)
+                g_cu_rows[0] = s_g_cu[row0]
+                g_cu_rows[1] = s_g_cu[row1]
+                g_cu_rows = g_cu_rows.load().reshape((1, 2))
+
+                for i in cutlass.range_constexpr(BT // 8):
+                    col = i * 8 + (lane_id % 4) * 2
+                    g_cu_cols = cute.make_rmem_tensor(2, Float32)
+                    g_cu_cols[0] = s_g_cu[col]
+                    g_cu_cols[1] = s_g_cu[col + 1]
+                    g_cu_cols = g_cu_cols.load().reshape((2, 1))
+
+                    # apply gamma and causal mask
+                    Gamma = cute.math.exp(g_cu_rows - g_cu_cols, fastmath=True)
+                    tmp = qk[None, None, i] * Gamma
+                    tmp = cute.where(row_indices >= col_indices + i * 8, tmp, 0.0)
+
+                    # CuteDSL can't emit cvt.bf16x2.f32 here
+                    attn_lo = cute.make_rmem_tensor(2, Uint32)
+                    attn_lo[0] = cvt.fp32x2_to_bf16x2(tmp[0, 0], tmp[1, 0])
+                    attn_lo[1] = cvt.fp32x2_to_bf16x2(tmp[0, 1], tmp[1, 1])
+                    _tcgen05.st(warp_id_ * 32, p_tmem + i * 4, "16x128b", 1, attn_lo)
+
+                _tcgen05.wait_st()
+                _tcgen05.fence_before_thread_sync()
+                cute.arch.mbarrier_arrive(mask_mbar)
+
+                parity ^= 1
+
+        else:
+            # epilogue warps
+            # for ldmatrix layout later
+            row0 = warp_id * 16 + lane_id // 4
+            row1 = row0 + 8
+
+            stage_id = 0
+            mma_parity = 0
+
+            op = cute.nvgpu.CopyUniversalOp()
+            cp_4B = cute.make_copy_atom(op, BFloat16, num_bits_per_copy=32)
+            stsm_op = warp.StMatrix8x8x16bOp(num_matrices=4, transpose=False)
+            stsm_atom = cute.make_copy_atom(stsm_op, BFloat16)
+
+            # ldmatrix layout
+            # [total_seq_len, ((2, 4, WIDTH/8), V_DIM/WIDTH)]
+            WIDTH = 64
+            o_view = cute.logical_divide(
+                o[None, v_head_id, None],
+                (None, cute.make_layout((2, 4, WIDTH // 8))),
+            )
+            # select lane: [total_seq_len, 2, WIDTH/8, V_DIM/WIDTH]
+            o_view = o_view[None, ((None, lane_id % 4, None), None)]
+
+            for global_chunk_id in range(bid, num_global_chunks, grid_x):
+                seq_id = chunk_indices[global_chunk_id, 0]
+                chunk_id = chunk_indices[global_chunk_id, 1]
+                bos = cu_seqlens[seq_id]
+                eos = cu_seqlens[seq_id + 1]
+                chunk_start = bos + chunk_id * BT
+                full_chunk = chunk_start + BT <= eos
+
+                g_cu_rows = cute.make_rmem_tensor(2, Float32)
+                g_cu_rows.fill(0.0)
+
+                # load g_cu
+                if chunk_start + row0 < eos:
+                    g_cu_rows[0] = cute.math.exp(
+                        g_cu[chunk_start + row0, v_head_id], fastmath=True
+                    )
+                if chunk_start + row1 < eos:
+                    g_cu_rows[1] = cute.math.exp(
+                        g_cu[chunk_start + row1, v_head_id], fastmath=True
+                    )
+                g_cu_rows = g_cu_rows.load().reshape((1, 2, 1))
+
+                if warp_id == 0:
+                    cute.arch.mbarrier_wait(pv_mma_mbar + stage_id, mma_parity)
+                elif warp_id == 3 and full_chunk:
+                    cute.arch.cp_async_bulk_wait_group(0, read=True)
+                cute.arch.barrier(barrier_id=2, number_of_threads=128)
+                _tcgen05.fence_after_thread_sync()
+
+                if full_chunk:
+                    # use TMA store: tmem->rmem->smem->gmem
+                    for i in cutlass.range_constexpr(V_dim // WIDTH):
+                        qh = _tcgen05.ld(
+                            warp_id * 32, qh_tmem + i * WIDTH, "16x256b", WIDTH // 8
+                        )
+                        pv = _tcgen05.ld(
+                            warp_id * 32, out_tmem + i * WIDTH, "16x256b", WIDTH // 8
+                        )
+                        _tcgen05.wait_ld()
+                        if i == V_dim // WIDTH - 1:
+                            _tcgen05.fence_before_thread_sync()
+                            cute.arch.mbarrier_arrive(epi_mbar)
+
+                        qh = qh.reshape((2, 2, WIDTH // 8))
+                        pv = pv.reshape((2, 2, WIDTH // 8))
+
+                        out_f32 = scale * (g_cu_rows * qh + pv)
+                        out_bf16 = cute.make_rmem_tensor((8, WIDTH // 16), BFloat16)
+                        out_bf16.store(out_f32.to(BFloat16).reshape((8, WIDTH // 16)))
+
+                        # TODO: issue single cute.copy()
+                        for j in cutlass.range_constexpr(WIDTH // 16):
+                            s_row = warp_id * 16 + lane_id % 16
+                            s_col = i * (WIDTH // 8) + j * 2 + lane_id // 16
+                            sO_tile = cute.local_tile(sO[s_row, None], (8,), (s_col,))
+                            cute.copy(stsm_atom, out_bf16[None, j], sO_tile)
+
+                    cute.arch.barrier(barrier_id=2, number_of_threads=128)
+                    fence_before_tma_store()
+                    if warp_id == 3:
+                        gO = cute.local_tile(
+                            cute.domain_offset((bos, 0), tmaO[None, v_head_id, None]),
+                            tiler=(BT, V_dim),
+                            coord=(chunk_id, 0),
+                        )
+                        simple_tma_copy(O_tma_atom, sO, gO)
+                        with cute.arch.elect_one():
+                            cute.arch.cp_async_bulk_commit_group()
+
+                else:
+                    # direct gmem store
+                    # TODO: explore doing multiple 1D TMAs
+                    for i in cutlass.range_constexpr(V_dim // WIDTH):
+                        qh = _tcgen05.ld(
+                            warp_id * 32, qh_tmem + i * WIDTH, "16x256b", WIDTH // 8
+                        )
+                        pv = _tcgen05.ld(
+                            warp_id * 32, out_tmem + i * WIDTH, "16x256b", WIDTH // 8
+                        )
+                        _tcgen05.wait_ld()
+                        if i == V_dim // WIDTH - 1:
+                            _tcgen05.fence_before_thread_sync()
+                            cute.arch.mbarrier_arrive(epi_mbar)
+
+                        qh = qh.reshape((2, 2, WIDTH // 8))
+                        pv = pv.reshape((2, 2, WIDTH // 8))
+
+                        out_f32 = scale * (g_cu_rows * qh + pv)
+                        out_bf16 = cute.make_rmem_tensor((2, 2, WIDTH // 8), BFloat16)
+                        out_bf16.store(out_f32.to(BFloat16))
+
+                        if chunk_start + row0 < eos:
+                            cute.copy(
+                                cp_4B,
+                                out_bf16[None, 0, None],
+                                o_view[chunk_start + row0, None, None, i],
+                            )
+                        if chunk_start + row1 < eos:
+                            cute.copy(
+                                cp_4B,
+                                out_bf16[None, 1, None],
+                                o_view[chunk_start + row1, None, None, i],
+                            )
+
+                stage_id = (stage_id + 1) % num_stages
+                if stage_id == 0:
+                    mma_parity ^= 1
+
+    @cache
+    @staticmethod
+    def compile(
+        H: int,
+        Hv: int,
+        K_dim: int,
+        V_dim: int,
+        BT: int = 64,
+        num_stages: int = 2,
+    ):
+        total_t = cute.sym_int()
+        pad_t = cute.sym_int()
+        total_chunks_n = cute.sym_int()
+        h_outer_n = cute.sym_int()
+        cu_entries = cute.sym_int()
+
+        q = make_fake_tensor(BFloat16, (total_t, H, K_dim), divisibility=16)
+        k = make_fake_tensor(BFloat16, (total_t, H, K_dim), divisibility=16)
+        v_new = make_fake_tensor(BFloat16, (pad_t, Hv, V_dim), divisibility=16)
+        h_flat = make_fake_tensor(BFloat16, (h_outer_n, V_dim, K_dim), divisibility=16)
+        g_cu = make_fake_tensor(Float32, (total_t, Hv), divisibility=4)
+        o = make_fake_tensor(BFloat16, (total_t, Hv, V_dim), divisibility=16)
+        cu_seqlens = make_fake_tensor(Int32, (cu_entries,), divisibility=1)
+        chunk_indices = make_fake_tensor(Int32, (total_chunks_n, 2), divisibility=2)
+        total_chunks = make_fake_tensor(Int32, (1,), divisibility=1)
+
+        kernel = Sm100ChunkOKernel(
+            H,
+            Hv,
+            K_dim,
+            V_dim,
+            BT,
+            num_stages,
+        )
+        stream = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True)
+        return cute.compile(
+            kernel,
+            q,
+            k,
+            v_new,
+            h_flat,
+            g_cu,
+            o,
+            cu_seqlens,
+            chunk_indices,
+            total_chunks,
+            Float32(1.0),
+            Int32(148),
+            stream,
+            options="--enable-tvm-ffi",
+        )
+
+
+def o_cutedsl(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v_new_chunks: torch.Tensor,
+    h: torch.Tensor,
+    g_cu: torch.Tensor,
+    o: torch.Tensor,
+    cu_seqlens: torch.Tensor,
+    chunk_indices: torch.Tensor,
+    total_chunks: torch.Tensor,
+    scale: float,
+    num_sms: int = 148,
+) -> None:
+    _, H, K_dim = q.shape
+    _, Hv, V_dim = o.shape
+
+    Sm100ChunkOKernel.compile(H, Hv, K_dim, V_dim)(
+        q,
+        k,
+        v_new_chunks.view(-1, Hv, V_dim),
+        h.view(-1, V_dim, K_dim),
+        g_cu,
+        o,
+        cu_seqlens,
+        chunk_indices,
+        total_chunks,
+        float(scale),
+        num_sms,
+    )

vllm/model_executor/layers/quantization/humming.py

@@ -43,12 +43,9 @@ from vllm.model_executor.parameter import (
     RowvLLMParameter,
 )
 from vllm.model_executor.utils import set_weight_attrs
+from vllm.platforms import current_platform
 
-if TYPE_CHECKING:
-    from vllm.model_executor.models.utils import WeightsMapper
-
-
-try:
+if current_platform.is_cuda():
     from humming.dtypes import DataType
     from humming.layer import HummingMethod
     from humming.schema import (
@@ -65,16 +62,17 @@ try:
         HummingIndexedExperts,
         get_humming_moe_gemm_type,
     )
-except ModuleNotFoundError:
-    HummingMethod = None
 
-
-def assert_humming_available():
-    assert HummingMethod is not None, (
-        "humming is not available, please run "
-        "'pip install git+https://github.com/inclusionAI/humming' to install it."
+if TYPE_CHECKING:
+    from humming.schema import (
+        BaseInputSchema,
+        BaseWeightSchema,
+        HummingInputSchema,
+        HummingWeightSchema,
     )
 
+    from vllm.model_executor.models.utils import WeightsMapper
+
 
 def prepare_padded_shape(shape, x):
     padded_shape = math.ceil(shape / x) * x
@@ -186,7 +184,6 @@ class HummingConfig(QuantizationConfig):
     packed_modules_mapping: dict[str, list[str]] = {}
 
     def __init__(self, full_config: dict[str, Any] | None = None):
-        assert_humming_available()
         self.full_config: dict[str, Any] = full_config or {}
 
     @classmethod

vllm/models/deepseek_v4/amd/model.py

@@ -11,17 +11,12 @@ import torch.nn as nn
 from vllm.compilation.decorators import support_torch_compile
 from vllm.config import VllmConfig
 from vllm.distributed import (
-    get_ep_group,
     get_pp_group,
     get_tensor_model_parallel_rank,
     get_tensor_model_parallel_world_size,
 )
-from vllm.forward_context import get_forward_context
 from vllm.model_executor.layers.activation import SiluAndMul, SiluAndMulWithClamp
 from vllm.model_executor.layers.fused_moe import FusedMoE, GateLinear
-from vllm.model_executor.layers.fused_moe.router.fused_topk_bias_router import (
-    fused_topk_bias,
-)
 from vllm.model_executor.layers.layernorm import RMSNorm
 from vllm.model_executor.layers.linear import (
     ColumnParallelLinear,
@@ -52,16 +47,13 @@ from vllm.model_executor.models.utils import (
     make_layers,
     maybe_prefix,
 )
-from vllm.model_executor.utils import set_weight_attrs
-from vllm.models.deepseek_v4.nvidia.ops.attention import (
+from vllm.models.deepseek_v4.attention import (
     DeepseekV4Indexer,
     DeepseekV4MLAModules,
     DeepseekV4MultiHeadLatentAttentionWrapper,
 )
 from vllm.platforms import current_platform
 from vllm.sequence import IntermediateTensors
-from vllm.triton_utils import tl, triton
-from vllm.utils.torch_utils import direct_register_custom_op
 
 
 class DeepseekV4MLP(nn.Module):
@@ -115,501 +107,6 @@ class DeepseekV4MLP(nn.Module):
         return x
 
 
-@triton.jit
-def _deepseek_v4_stage_mega_moe_inputs_kernel(
-    hidden_states,
-    x_fp8,
-    x_sf,
-    topk_ids,
-    topk_weights,
-    topk_idx_out,
-    topk_weights_out,
-    hidden_stride_m: tl.constexpr,
-    hidden_stride_k: tl.constexpr,
-    x_stride_m: tl.constexpr,
-    x_stride_k: tl.constexpr,
-    x_sf_stride_m: tl.constexpr,
-    x_sf_stride_k: tl.constexpr,
-    topk_ids_stride_m: tl.constexpr,
-    topk_ids_stride_k: tl.constexpr,
-    topk_weights_stride_m: tl.constexpr,
-    topk_weights_stride_k: tl.constexpr,
-    topk_idx_stride_m: tl.constexpr,
-    topk_idx_stride_k: tl.constexpr,
-    topk_weights_out_stride_m: tl.constexpr,
-    topk_weights_out_stride_k: tl.constexpr,
-    hidden_size: tl.constexpr,
-    top_k: tl.constexpr,
-    BLOCK_K: tl.constexpr,
-    GROUP_K: tl.constexpr,
-    BLOCK_TOPK: tl.constexpr,
-) -> None:
-    token_id = tl.program_id(0)
-    k_block_id = tl.program_id(1)
-
-    k_offsets = k_block_id * BLOCK_K + tl.arange(0, BLOCK_K)
-    k_mask = k_offsets < hidden_size
-    hidden = tl.load(
-        hidden_states + token_id * hidden_stride_m + k_offsets * hidden_stride_k,
-        mask=k_mask,
-        other=0.0,
-    ).to(tl.float32)
-
-    num_groups: tl.constexpr = BLOCK_K // GROUP_K
-    hidden_groups = tl.reshape(tl.abs(hidden), [num_groups, GROUP_K])
-    amax = tl.max(hidden_groups, axis=1)
-    amax = tl.maximum(amax, 1.0e-4)
-
-    scale = amax / 448.0
-    scale_bits = scale.to(tl.uint32, bitcast=True)
-    scale_exp = ((scale_bits >> 23) & 0xFF) + ((scale_bits & 0x7FFFFF) != 0).to(
-        tl.uint32
-    )
-    scale_exp = tl.minimum(tl.maximum(scale_exp, 1), 254)
-    rounded_scale = (scale_exp << 23).to(tl.float32, bitcast=True)
-
-    hidden_groups = tl.reshape(hidden, [num_groups, GROUP_K])
-    scaled = hidden_groups * (1.0 / rounded_scale)[:, None]
-    scaled = tl.reshape(scaled, [BLOCK_K])
-    fp8 = scaled.to(tl.float8e4nv)
-    tl.store(
-        x_fp8 + token_id * x_stride_m + k_offsets * x_stride_k,
-        fp8,
-        mask=k_mask,
-    )
-
-    scale_offsets = tl.arange(0, num_groups)
-    packed_scale = tl.sum(scale_exp << (scale_offsets * 8), axis=0).to(tl.int32)
-    tl.store(
-        x_sf + token_id * x_sf_stride_m + k_block_id * x_sf_stride_k,
-        packed_scale,
-    )
-
-    if k_block_id == 0:
-        topk_offsets = tl.arange(0, BLOCK_TOPK)
-        topk_mask = topk_offsets < top_k
-
-        ids = tl.load(
-            topk_ids + token_id * topk_ids_stride_m + topk_offsets * topk_ids_stride_k,
-            mask=topk_mask,
-            other=0,
-        ).to(tl.int64)
-        tl.store(
-            topk_idx_out
-            + token_id * topk_idx_stride_m
-            + topk_offsets * topk_idx_stride_k,
-            ids,
-            mask=topk_mask,
-        )
-
-        weights = tl.load(
-            topk_weights
-            + token_id * topk_weights_stride_m
-            + topk_offsets * topk_weights_stride_k,
-            mask=topk_mask,
-            other=0.0,
-        )
-        tl.store(
-            topk_weights_out
-            + token_id * topk_weights_out_stride_m
-            + topk_offsets * topk_weights_out_stride_k,
-            weights,
-            mask=topk_mask,
-        )
-
-
-def _stage_deepseek_v4_mega_moe_inputs(
-    hidden_states: torch.Tensor,
-    topk_weights: torch.Tensor,
-    topk_ids: torch.Tensor,
-    x_fp8: torch.Tensor,
-    x_sf: torch.Tensor,
-    topk_idx_out: torch.Tensor,
-    topk_weights_out: torch.Tensor,
-) -> None:
-    num_tokens, hidden_size = hidden_states.shape
-    if num_tokens == 0:
-        return
-    if hidden_size % 128 != 0:
-        raise ValueError(
-            "DeepSeek V4 MegaMoE input staging requires hidden_size to be "
-            "a multiple of 128."
-        )
-    top_k = topk_ids.shape[1]
-    if topk_weights.shape != topk_ids.shape:
-        raise ValueError(
-            "DeepSeek V4 MegaMoE input staging requires topk_weights and "
-            "topk_ids to have the same shape."
-        )
-
-    block_k = 128
-    grid = (num_tokens, triton.cdiv(hidden_size, block_k))
-    block_topk = triton.next_power_of_2(top_k)
-    _deepseek_v4_stage_mega_moe_inputs_kernel[grid](
-        hidden_states,
-        x_fp8,
-        x_sf,
-        topk_ids,
-        topk_weights,
-        topk_idx_out,
-        topk_weights_out,
-        hidden_states.stride(0),
-        hidden_states.stride(1),
-        x_fp8.stride(0),
-        x_fp8.stride(1),
-        x_sf.stride(0),
-        x_sf.stride(1),
-        topk_ids.stride(0),
-        topk_ids.stride(1),
-        topk_weights.stride(0),
-        topk_weights.stride(1),
-        topk_idx_out.stride(0),
-        topk_idx_out.stride(1),
-        topk_weights_out.stride(0),
-        topk_weights_out.stride(1),
-        hidden_size,
-        top_k,
-        BLOCK_K=block_k,
-        GROUP_K=32,
-        BLOCK_TOPK=block_topk,
-        num_warps=4,
-    )
-
-
-def make_deepseek_v4_expert_params_mapping(
-    num_experts: int,
-) -> list[tuple[str, str, int, str]]:
-    return [
-        (
-            "experts.w13_" if shard_id in ("w1", "w3") else "experts.w2_",
-            f"experts.{expert_id}.{weight_name}.",
-            expert_id,
-            shard_id,
-        )
-        for expert_id in range(num_experts)
-        for shard_id, weight_name in [
-            ("w1", "w1"),
-            ("w2", "w2"),
-            ("w3", "w3"),
-        ]
-    ]
-
-
-class DeepseekV4MegaMoEExperts(nn.Module):
-    _symm_buffer_cache: dict[tuple[int, int, int, int, int, int, int], object] = {}
-
-    def __init__(
-        self,
-        vllm_config: VllmConfig,
-        *,
-        num_experts: int,
-        num_local_experts: int,
-        experts_start_idx: int,
-        top_k: int,
-        hidden_size: int,
-        intermediate_size: int,
-        prefix: str = "",
-    ):
-        super().__init__()
-        self.prefix = prefix
-        self.num_experts = num_experts
-        self.num_local_experts = num_local_experts
-        self.experts_start_idx = experts_start_idx
-        self.experts_end_idx = experts_start_idx + num_local_experts
-        self.top_k = top_k
-        self.hidden_size = hidden_size
-        self.intermediate_size = intermediate_size
-        self.max_num_tokens = vllm_config.scheduler_config.max_num_batched_tokens
-
-        weight_attrs = {"weight_loader": self.weight_loader}
-        self.w13_weight = nn.Parameter(
-            torch.zeros(
-                num_local_experts,
-                2 * intermediate_size,
-                hidden_size // 2,
-                dtype=torch.uint8,
-            ),
-            requires_grad=False,
-        )
-        set_weight_attrs(self.w13_weight, weight_attrs)
-
-        self.w13_weight_scale = nn.Parameter(
-            torch.zeros(
-                num_local_experts,
-                2 * intermediate_size,
-                hidden_size // 32,
-                dtype=torch.uint8,
-            ),
-            requires_grad=False,
-        )
-        set_weight_attrs(self.w13_weight_scale, weight_attrs)
-        self.w13_weight_scale.quant_method = "block"
-
-        self.w2_weight = nn.Parameter(
-            torch.zeros(
-                num_local_experts,
-                hidden_size,
-                intermediate_size // 2,
-                dtype=torch.uint8,
-            ),
-            requires_grad=False,
-        )
-        set_weight_attrs(self.w2_weight, weight_attrs)
-
-        self.w2_weight_scale = nn.Parameter(
-            torch.zeros(
-                num_local_experts,
-                hidden_size,
-                intermediate_size // 32,
-                dtype=torch.uint8,
-            ),
-            requires_grad=False,
-        )
-        set_weight_attrs(self.w2_weight_scale, weight_attrs)
-        self.w2_weight_scale.quant_method = "block"
-
-        self._transformed_l1_weights: tuple[torch.Tensor, torch.Tensor] | None = None
-        self._transformed_l2_weights: tuple[torch.Tensor, torch.Tensor] | None = None
-
-        # Register in the static forward context so the custom-op wrapper
-        # can look up this module by name from within a torch.compile graph.
-        compilation_config = vllm_config.compilation_config
-        if prefix in compilation_config.static_forward_context:
-            raise ValueError(f"Duplicate layer name: {prefix}")
-        compilation_config.static_forward_context[prefix] = self
-
-    def _map_global_expert_id(self, expert_id: int) -> int:
-        if expert_id < self.experts_start_idx or expert_id >= self.experts_end_idx:
-            return -1
-        return expert_id - self.experts_start_idx
-
-    def weight_loader(
-        self,
-        param: nn.Parameter,
-        loaded_weight: torch.Tensor,
-        weight_name: str,
-        shard_id: str,
-        expert_id: int,
-        return_success: bool = False,
-    ) -> bool | None:
-        local_expert_id = self._map_global_expert_id(expert_id)
-        if local_expert_id == -1:
-            return False if return_success else None
-
-        expert_data = param.data[local_expert_id]
-        if shard_id in ("w1", "w3"):
-            if "w13_" not in weight_name:
-                return False if return_success else None
-            shard_offset = 0 if shard_id == "w1" else self.intermediate_size
-            expert_data = expert_data.narrow(0, shard_offset, self.intermediate_size)
-        elif shard_id == "w2":
-            if "w2_" not in weight_name:
-                return False if return_success else None
-        else:
-            raise ValueError(f"Unsupported expert shard id: {shard_id}")
-
-        if expert_data.shape != loaded_weight.shape:
-            raise ValueError(
-                f"DeepSeek V4 MegaMoE expert weight shape mismatch for "
-                f"{weight_name}: parameter shard {tuple(expert_data.shape)} "
-                f"vs checkpoint {tuple(loaded_weight.shape)}"
-            )
-        expert_data.copy_(loaded_weight)
-        return True if return_success else None
-
-    @staticmethod
-    def _ue8m0_uint8_to_float(sf: torch.Tensor) -> torch.Tensor:
-        return (sf.to(torch.int32) << 23).view(torch.float32)
-
-    def _check_runtime_supported(self) -> None:
-        if not torch.cuda.is_available():
-            raise NotImplementedError("DeepSeek V4 MegaMoE requires CUDA.")
-        device = self.w13_weight.device
-        if device.type != "cuda":
-            raise NotImplementedError(
-                "DeepSeek V4 MegaMoE expert weights must be loaded on CUDA."
-            )
-        if torch.cuda.get_device_capability(device)[0] != 10:
-            raise NotImplementedError("DeepGEMM MegaMoE requires SM100 GPUs.")
-        if self.hidden_size % 128 != 0 or self.intermediate_size % 128 != 0:
-            raise ValueError(
-                "DeepGEMM MegaMoE requires hidden and intermediate sizes "
-                "to be multiples of 128."
-            )
-
-    def finalize_weights(self) -> None:
-        if self._transformed_l1_weights is not None:
-            return
-
-        self._check_runtime_supported()
-        import vllm.third_party.deep_gemm as deep_gemm
-
-        w13_scale = deep_gemm.transform_sf_into_required_layout(
-            self._ue8m0_uint8_to_float(self.w13_weight_scale.data).contiguous(),
-            2 * self.intermediate_size,
-            self.hidden_size,
-            (1, 32),
-            self.num_local_experts,
-        )
-        w2_scale = deep_gemm.transform_sf_into_required_layout(
-            self._ue8m0_uint8_to_float(self.w2_weight_scale.data).contiguous(),
-            self.hidden_size,
-            self.intermediate_size,
-            (1, 32),
-            self.num_local_experts,
-        )
-        self._transformed_l1_weights, self._transformed_l2_weights = (
-            deep_gemm.transform_weights_for_mega_moe(
-                (self.w13_weight.data.view(torch.int8).contiguous(), w13_scale),
-                (self.w2_weight.data.view(torch.int8).contiguous(), w2_scale),
-            )
-        )
-        # Drop the original loader-side parameters: the MegaMoE kernels only
-        # consume the transformed views above. transform_weights_for_mega_moe
-        # allocates a fresh tensor for the L1 weight (see _interleave_l1_weights)
-        # and fresh SF tensors for L1/L2; the L2 weight is the only tensor that
-        # aliases the original storage, and _transformed_l2_weights still holds
-        # it, so the storage stays live after we drop the Parameter.
-        self.w13_weight = None
-        self.w13_weight_scale = None
-        self.w2_weight = None
-        self.w2_weight_scale = None
-
-    def get_symm_buffer(self):
-        import vllm.third_party.deep_gemm as deep_gemm
-
-        group = get_ep_group().device_group
-        device = torch.accelerator.current_device_index()
-        key = (
-            id(group),
-            device,
-            self.num_experts,
-            self.max_num_tokens,
-            self.top_k,
-            self.hidden_size,
-            self.intermediate_size,
-        )
-        symm_buffer = self._symm_buffer_cache.get(key)
-        if symm_buffer is None:
-            symm_buffer = deep_gemm.get_symm_buffer_for_mega_moe(
-                group,
-                self.num_experts,
-                self.max_num_tokens,
-                self.top_k,
-                self.hidden_size,
-                self.intermediate_size,
-            )
-            self._symm_buffer_cache[key] = symm_buffer
-        return symm_buffer
-
-    def forward(
-        self,
-        hidden_states: torch.Tensor,
-        topk_weights: torch.Tensor,
-        topk_ids: torch.Tensor,
-        *,
-        activation_clamp: float | None,
-        fast_math: bool = True,
-    ) -> torch.Tensor:
-        if hidden_states.shape[0] > self.max_num_tokens:
-            raise ValueError(
-                f"DeepSeek V4 MegaMoE got {hidden_states.shape[0]} tokens, "
-                f"but the symmetric buffer was sized for {self.max_num_tokens}."
-            )
-        y = torch.empty_like(hidden_states, dtype=torch.bfloat16)
-        torch.ops.vllm.deepseek_v4_mega_moe_experts(
-            hidden_states,
-            topk_weights,
-            topk_ids,
-            y,
-            self.prefix,
-            activation_clamp,
-            fast_math,
-        )
-        return y
-
-    def _run_mega_moe(
-        self,
-        hidden_states: torch.Tensor,
-        topk_weights: torch.Tensor,
-        topk_ids: torch.Tensor,
-        y: torch.Tensor,
-        activation_clamp: float | None,
-        fast_math: bool,
-    ) -> None:
-        import vllm.third_party.deep_gemm as deep_gemm
-
-        symm_buffer = self.get_symm_buffer()
-        num_tokens = hidden_states.shape[0]
-        _stage_deepseek_v4_mega_moe_inputs(
-            hidden_states,
-            topk_weights,
-            topk_ids,
-            symm_buffer.x[:num_tokens],
-            symm_buffer.x_sf[:num_tokens],
-            symm_buffer.topk_idx[:num_tokens],
-            symm_buffer.topk_weights[:num_tokens],
-        )
-
-        # This method must have been already called during the weight loading phase.
-        # We call it again here to cover the dummy weight loading case.
-        self.finalize_weights()
-
-        assert self._transformed_l1_weights is not None
-        assert self._transformed_l2_weights is not None
-        deep_gemm.fp8_fp4_mega_moe(
-            y,
-            self._transformed_l1_weights,
-            self._transformed_l2_weights,
-            symm_buffer,
-            activation_clamp=activation_clamp,
-            fast_math=fast_math,
-        )
-
-
-DeepseekV4MegaMoEExperts.weight_loader.supports_moe_loading = True  # type: ignore[attr-defined]
-
-
-def _deepseek_v4_mega_moe_experts_op(
-    hidden_states: torch.Tensor,
-    topk_weights: torch.Tensor,
-    topk_ids: torch.Tensor,
-    out: torch.Tensor,
-    layer_name: str,
-    activation_clamp: float | None,
-    fast_math: bool,
-) -> None:
-    self = get_forward_context().no_compile_layers[layer_name]
-    self._run_mega_moe(
-        hidden_states,
-        topk_weights,
-        topk_ids,
-        out,
-        activation_clamp,
-        fast_math,
-    )
-
-
-def _deepseek_v4_mega_moe_experts_op_fake(
-    hidden_states: torch.Tensor,
-    topk_weights: torch.Tensor,
-    topk_ids: torch.Tensor,
-    out: torch.Tensor,
-    layer_name: str,
-    activation_clamp: float | None,
-    fast_math: bool,
-) -> None:
-    return None
-
-
-direct_register_custom_op(
-    op_name="deepseek_v4_mega_moe_experts",
-    op_func=_deepseek_v4_mega_moe_experts_op,
-    mutates_args=["out"],
-    fake_impl=_deepseek_v4_mega_moe_experts_op_fake,
-)
-
-
 class DeepseekV4MoE(nn.Module):
     def __init__(
         self,
@@ -622,15 +119,6 @@ class DeepseekV4MoE(nn.Module):
         config = vllm_config.model_config.hf_config
         quant_config = vllm_config.quant_config
         self.prefix = prefix
-        self.use_mega_moe = (
-            vllm_config.kernel_config.moe_backend == "deep_gemm_mega_moe"
-        )
-        if self.use_mega_moe and not vllm_config.parallel_config.enable_expert_parallel:
-            raise NotImplementedError(
-                "DeepSeek V4 MegaMoE currently requires expert parallel. "
-                "Enable it with --enable-expert-parallel, or pick a different "
-                "moe backend."
-            )
 
         self.routed_scaling_factor = getattr(config, "routed_scaling_factor", 1.0)
         self.hidden_size = config.hidden_size
@@ -641,16 +129,6 @@ class DeepseekV4MoE(nn.Module):
         self.swiglu_limit = config.swiglu_limit
         self.renormalize = config.norm_topk_prob
         self.scoring_func = getattr(config, "scoring_func", "sqrtsoftplus")
-        if self.use_mega_moe and self.scoring_func != "sqrtsoftplus":
-            raise NotImplementedError(
-                "DeepSeek V4 MegaMoE currently supports sqrtsoftplus routing only."
-            )
-        if self.use_mega_moe and getattr(config, "expert_dtype", "fp4") != "fp4":
-            raise NotImplementedError(
-                "DeepSeek V4 MegaMoE only supports fp4 experts; got expert_dtype="
-                f"{config.expert_dtype!r}. Drop --kernel-config moe_backend="
-                "deep_gemm_mega_moe for this checkpoint."
-            )
 
         self.gate = GateLinear(
             input_size=config.hidden_size,
@@ -663,7 +141,7 @@ class DeepseekV4MoE(nn.Module):
         self.gate.e_score_correction_bias = None
         self.gate.tid2eid = None
         is_hash_moe = extract_layer_index(prefix) < config.num_hash_layers
-        self.hash_indices_dtype = torch.int64 if self.use_mega_moe else torch.int32
+        self.hash_indices_dtype = torch.int32
         if is_hash_moe:
             # hash MoE doesn't use e_score_correction_bias
             # Use randint instead of empty to avoid garbage values causing
@@ -694,47 +172,10 @@ class DeepseekV4MoE(nn.Module):
                 hidden_act=config.hidden_act,
                 swiglu_limit=self.swiglu_limit,
                 quant_config=quant_config,
-                reduce_results=self.use_mega_moe,
+                reduce_results=False,
                 prefix=f"{prefix}.shared_experts",
             )
 
-        if self.use_mega_moe:
-            self._init_mega_moe_experts(vllm_config, config, prefix)
-        else:
-            self._init_fused_moe_experts(config, quant_config, prefix)
-
-    def _init_mega_moe_experts(
-        self,
-        vllm_config: VllmConfig,
-        config,
-        prefix: str,
-    ) -> None:
-        self.ep_group = get_ep_group()
-        self.ep_size = self.ep_group.world_size
-        self.ep_rank = self.ep_group.rank_in_group
-        assert config.n_routed_experts % self.ep_size == 0
-
-        self.n_local_experts = config.n_routed_experts // self.ep_size
-        self.experts_start_idx = self.ep_rank * self.n_local_experts
-        self.experts_end_idx = self.experts_start_idx + self.n_local_experts
-
-        self.experts = DeepseekV4MegaMoEExperts(
-            vllm_config,
-            num_experts=config.n_routed_experts,
-            num_local_experts=self.n_local_experts,
-            experts_start_idx=self.experts_start_idx,
-            top_k=config.num_experts_per_tok,
-            hidden_size=config.hidden_size,
-            intermediate_size=config.moe_intermediate_size,
-            prefix=f"{prefix}.experts",
-        )
-
-    def _init_fused_moe_experts(
-        self,
-        config,
-        quant_config,
-        prefix: str,
-    ) -> None:
         self.tp_rank = get_tensor_model_parallel_rank()
         assert config.n_routed_experts % self.tp_size == 0
 
@@ -766,44 +207,6 @@ class DeepseekV4MoE(nn.Module):
         if self.gate.tid2eid is not None and input_ids is None:
             raise ValueError("DeepSeek V4 hash MoE routing requires input_ids.")
 
-        if not self.use_mega_moe:
-            return self._forward_fused_moe(hidden_states, input_ids)
-
-        org_shape = hidden_states.shape
-        router_logits, _ = self.gate(hidden_states)
-        topk_weights, topk_ids = fused_topk_bias(
-            hidden_states=hidden_states,
-            gating_output=router_logits,
-            scoring_func=self.scoring_func,
-            e_score_correction_bias=self.gate.e_score_correction_bias.data
-            if self.gate.e_score_correction_bias is not None
-            else None,
-            topk=self.n_activated_experts,
-            renormalize=self.renormalize,
-            indices_type=self.hash_indices_dtype,
-            input_tokens=input_ids,
-            hash_indices_table=self.gate.tid2eid,
-            routed_scaling_factor=self.routed_scaling_factor,
-        )
-        activation_clamp = (
-            float(self.swiglu_limit) if self.swiglu_limit is not None else None
-        )
-        final_hidden_states = self.experts(
-            hidden_states,
-            topk_weights,
-            topk_ids,
-            activation_clamp=activation_clamp,
-        )
-
-        if self.shared_experts is not None:
-            shared_output = self.shared_experts(hidden_states)
-            final_hidden_states += shared_output
-
-        return final_hidden_states.view(org_shape)
-
-    def _forward_fused_moe(
-        self, hidden_states: torch.Tensor, input_ids: torch.Tensor | None = None
-    ) -> torch.Tensor:
         org_shape = hidden_states.shape
         if self.experts.is_internal_router:
             # In this case, the gate/router runs inside the FusedMoE class
@@ -822,10 +225,6 @@ class DeepseekV4MoE(nn.Module):
 
         return final_hidden_states.view(org_shape)
 
-    def finalize_mega_moe_weights(self) -> None:
-        if self.use_mega_moe:
-            self.experts.finalize_weights()
-
 
 class DeepseekV4Attention(nn.Module):
     def __init__(
@@ -1211,15 +610,6 @@ class DeepseekV4Model(nn.Module):
         config = vllm_config.model_config.hf_config
         quant_config = vllm_config.quant_config
         self.config = config
-        self.use_mega_moe = (
-            vllm_config.kernel_config.moe_backend == "deep_gemm_mega_moe"
-        )
-        if self.use_mega_moe and not vllm_config.parallel_config.enable_expert_parallel:
-            raise NotImplementedError(
-                "DeepSeek V4 MegaMoE currently requires expert parallel. "
-                "Enable it with --enable-expert-parallel, or pick a different "
-                "moe backend."
-            )
         self.vocab_size = config.vocab_size
         self.hc_eps = config.hc_eps
         self.hc_mult = config.hc_mult
@@ -1348,9 +738,6 @@ class DeepseekV4Model(nn.Module):
             assert intermediate_tensors is not None
             hidden_states = intermediate_tensors["hidden_states"]
 
-        if self.use_mega_moe:
-            input_ids = input_ids.to(torch.int64)
-
         residual, post_mix, res_mix = None, None, None
         for layer in islice(self.layers, self.start_layer, self.end_layer):
             hidden_states, residual, post_mix, res_mix = layer(
@@ -1482,9 +869,6 @@ class DeepseekV4Model(nn.Module):
         return loaded_params
 
     def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
-        first_layer = next(iter(islice(self.layers, self.start_layer, self.end_layer)))
-        if first_layer.ffn.use_mega_moe:
-            return make_deepseek_v4_expert_params_mapping(self.config.n_routed_experts)
         # Params for weights, fp8 weight scales, fp8 activation scales
         # (param_name, weight_name, expert_id, shard_id)
         return FusedMoE.make_expert_params_mapping(
@@ -1495,10 +879,6 @@ class DeepseekV4Model(nn.Module):
             num_experts=self.config.n_routed_experts,
         )
 
-    def finalize_mega_moe_weights(self) -> None:
-        for layer in islice(self.layers, self.start_layer, self.end_layer):
-            layer.ffn.finalize_mega_moe_weights()
-
 
 def _make_deepseek_v4_weights_mapper(expert_dtype: str) -> WeightsMapper:
     if expert_dtype == "fp4":
@@ -1601,7 +981,6 @@ class DeepseekV4ForCausalLM(nn.Module, SupportsPP):
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
         loader = AutoWeightsLoader(self, skip_substrs=["mtp."])
         loaded_params = loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
-        self.model.finalize_mega_moe_weights()
         return loaded_params
 
     def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:

vllm/models/deepseek_v4/amd/mtp.py

@@ -40,10 +40,7 @@ from vllm.model_executor.models.utils import maybe_prefix
 from vllm.platforms import current_platform
 from vllm.sequence import IntermediateTensors
 
-from .model import (
-    DeepseekV4DecoderLayer,
-    make_deepseek_v4_expert_params_mapping,
-)
+from .model import DeepseekV4DecoderLayer
 
 logger = init_logger(__name__)
 
@@ -330,19 +327,13 @@ class DeepSeekV4MTP(nn.Module):
         head_rank_end = n_local_head * (tp_rank + 1)
 
         # Pre-compute expert mapping ONCE.
-        first_layer = next(iter(self.model.layers.values()))
-        if first_layer.mtp_block.ffn.use_mega_moe:
-            expert_mapping = make_deepseek_v4_expert_params_mapping(
-                self.config.n_routed_experts
-            )
-        else:
-            expert_mapping = FusedMoE.make_expert_params_mapping(
-                self,
-                ckpt_gate_proj_name="w1",
-                ckpt_down_proj_name="w2",
-                ckpt_up_proj_name="w3",
-                num_experts=self.config.n_routed_experts,
-            )
+        expert_mapping = FusedMoE.make_expert_params_mapping(
+            self,
+            ckpt_gate_proj_name="w1",
+            ckpt_down_proj_name="w2",
+            ckpt_up_proj_name="w3",
+            num_experts=self.config.n_routed_experts,
+        )
 
         # FP8 experts register ``..._weight_scale_inv`` (block_quant) while
         # FP4/MXFP4 experts register ``..._weight_scale``. Choose the suffix
@@ -465,14 +456,9 @@ class DeepSeekV4MTP(nn.Module):
                     f"Use a checkpoint that includes MTP layer weights, "
                     f"or disable speculative decoding."
                 )
-        self.finalize_mega_moe_weights()
         logger.info_once("MTP draft model loaded: %d params", len(loaded_params))
         return loaded_params
 
-    def finalize_mega_moe_weights(self) -> None:
-        for layer in self.model.layers.values():
-            layer.mtp_block.ffn.finalize_mega_moe_weights()
-
     def _rewrite_spec_layer_name(self, spec_layer: int, name: str) -> str:
         """
         Rewrite the weight name to match the format of the original model.

vllm/models/deepseek_v4/amd/rocm.py

@@ -32,7 +32,7 @@ from vllm.v1.attention.ops.rocm_aiter_mla_sparse import (
 from vllm.v1.worker.workspace import current_workspace_manager
 
 if TYPE_CHECKING:
-    from vllm.models.deepseek_v4.nvidia.ops.attention import (
+    from vllm.models.deepseek_v4.attention import (
         DeepseekV4MLAAttention,
     )
 
@@ -592,6 +592,10 @@ class DeepseekV4ROCMAiterMLASparseImpl(DeepseekV4SparseMLAAttentionImpl):
 
     backend_cls = DeepseekV4ROCMAiterMLASparseBackend
 
+    @classmethod
+    def get_padded_num_q_heads(cls, num_heads: int) -> int:
+        return num_heads
+
     @classmethod
     def forward_mqa(  # type: ignore[override]
         cls,

vllm/models/deepseek_v4/attention.py

@@ -156,18 +156,6 @@ class DeepseekV4MultiHeadLatentAttentionWrapper(PluggableLayer):
         self.head_dim = head_dim
         self.scale = scale
 
-        # FlashMLA sparse kernel only supports 64 or 128 heads; pad up to the
-        # next supported size. Must match DeepseekV4MLAAttention.padded_heads.
-        if num_heads <= 64:
-            self.padded_heads = 64
-        elif num_heads <= 128:
-            self.padded_heads = 128
-        else:
-            raise ValueError(
-                f"DeepseekV4 attention does not support {num_heads} heads "
-                "(must be <= 128)."
-            )
-
         self.q_lora_rank = q_lora_rank
         self.kv_lora_rank = kv_lora_rank
         self.window_size = window_size
@@ -263,6 +251,9 @@ class DeepseekV4MultiHeadLatentAttentionWrapper(PluggableLayer):
             indexer=self.indexer,
             topk_indices_buffer=self.topk_indices_buffer,
         )
+        # Mirror the inner layer's padded head count (single source of truth).
+        self.padded_heads = self.mla_attn.padded_heads
+
         # Register this layer in the compilation config's static forward context
         # This allows the custom op to retrieve the layer during execution
         compilation_config = mla_modules.vllm_config.compilation_config
@@ -450,7 +441,7 @@ class DeepseekV4MultiHeadLatentAttentionWrapper(PluggableLayer):
 
             def wq_b_kv_insert() -> torch.Tensor:
                 q = self.wq_b(qr).view(-1, self.n_local_heads, self.head_dim)
-                self._fused_qnorm_rope_kv_insert(q, kv, positions, attn_metadata)
+                q = self._fused_qnorm_rope_kv_insert(q, kv, positions, attn_metadata)
                 return q
 
             # 3-way overlap (matches TRT-LLM PR #14142 Level 1): default runs
@@ -484,7 +475,7 @@ class DeepseekV4MultiHeadLatentAttentionWrapper(PluggableLayer):
 
             def wq_b_kv_insert() -> torch.Tensor:
                 q = self.wq_b(qr).view(-1, self.n_local_heads, self.head_dim)
-                self._fused_qnorm_rope_kv_insert(q, kv, positions, attn_metadata)
+                q = self._fused_qnorm_rope_kv_insert(q, kv, positions, attn_metadata)
                 return q
 
             q, _ = maybe_execute_in_parallel(
@@ -497,12 +488,7 @@ class DeepseekV4MultiHeadLatentAttentionWrapper(PluggableLayer):
         else:
             # SWA-only layer: no compressor, no overlap.
             q = self.wq_b(qr).view(-1, self.n_local_heads, self.head_dim)
-            self._fused_qnorm_rope_kv_insert(q, kv, positions, attn_metadata)
-
-        # Pad q to FlashMLA-required head count (64 or 128)
-        if self.n_local_heads < self.padded_heads:
-            pad_size = self.padded_heads - self.n_local_heads
-            q = F.pad(q, (0, 0, 0, pad_size), value=0.0)
+            q = self._fused_qnorm_rope_kv_insert(q, kv, positions, attn_metadata)
 
         # MLA attention writes into the pre-allocated `out` buffer
         # ([num_tokens, padded_heads, head_dim]).
@@ -516,9 +502,17 @@ class DeepseekV4MultiHeadLatentAttentionWrapper(PluggableLayer):
         attn_metadata: (
             dict[str, AttentionMetadata] | list[dict[str, AttentionMetadata]] | None
         ),
-    ) -> None:
+    ) -> torch.Tensor:
         if not isinstance(attn_metadata, dict):
-            return
+            # Profile run: kernel doesn't fire; produce a padded tensor so
+            # downstream FlashMLA gets the right shape.
+            if self.n_local_heads < self.padded_heads:
+                return F.pad(
+                    q,
+                    (0, 0, 0, self.padded_heads - self.n_local_heads),
+                    value=0.0,
+                )
+            return q
 
         swa_metadata = cast(
             "DeepseekSparseSWAMetadata | None",
@@ -530,16 +524,19 @@ class DeepseekV4MultiHeadLatentAttentionWrapper(PluggableLayer):
         swa_kv_cache_2d = swa_kv_cache.view(swa_kv_cache.shape[0], -1)
 
         # Horizontally fused:
-        #   Q side:  q_head_norm (per-head RMSNorm, no weight) + GPT-J RoPE
+        #   Q side:  q_head_norm (per-head RMSNorm, no weight) + GPT-J RoPE,
+        #            with zero-fill for the padding head slots.  The kernel
+        #            allocates and returns the padded q tensor.
         #   KV side: GPT-J RoPE + UE8M0 FP8 quant + paged cache insert
         # kv is unchanged; mla_attn reads kv solely via swa_kv_cache.
-        torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
+        return torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
             q,
             kv,
             swa_kv_cache_2d,
             swa_metadata.slot_mapping,
             positions.to(torch.int64),
             self.rotary_emb.cos_sin_cache,
+            self.padded_heads,
             self.eps,
             swa_metadata.block_size,
         )
@@ -607,9 +604,6 @@ direct_register_custom_op(
 
 
 class DeepseekV4MLAAttention(nn.Module, AttentionLayerBase):
-    # FlashMLA FP8 sparse only supports 64 or 128 heads
-    SUPPORTED_HEAD_COUNTS = (64, 128)
-
     def __init__(
         self,
         num_heads: int,
@@ -655,19 +649,8 @@ class DeepseekV4MLAAttention(nn.Module, AttentionLayerBase):
         self.aux_stream = aux_stream
         self.ln_events = [torch.cuda.Event(), torch.cuda.Event()]
 
-        # Determine padded head count for FlashMLA
-        if num_heads not in self.SUPPORTED_HEAD_COUNTS:
-            if num_heads < 64:
-                self.padded_heads = 64
-            elif num_heads < 128:
-                self.padded_heads = 128
-            else:
-                raise ValueError(
-                    f"DeepseekV4MLAAttention does not support {num_heads} heads. "
-                    f"Supported: <= 128 (will be padded to 64 or 128)"
-                )
-        else:
-            self.padded_heads = num_heads
+        # Padded Q head count is dictated by the selected impl.
+        self.padded_heads = self.impl_cls.get_padded_num_q_heads(num_heads)
 
         # Store attention sink
         assert attn_sink is not None

vllm/models/deepseek_v4/common/ops/__init__.py

@@ -10,6 +10,7 @@ from .cache_utils import (
 from .fused_indexer_q import MXFP4_BLOCK_SIZE, fused_indexer_q_rope_quant
 from .fused_inv_rope_fp8_quant import fused_inv_rope_fp8_quant
 from .fused_qk_rmsnorm import fused_q_kv_rmsnorm
+from .save_partial_states import save_partial_states
 
 __all__ = [
     "MXFP4_BLOCK_SIZE",
@@ -20,4 +21,5 @@ __all__ = [
     "fused_inv_rope_fp8_quant",
     "fused_q_kv_rmsnorm",
     "quantize_and_insert_k_cache",
+    "save_partial_states",
 ]

vllm/models/deepseek_v4/common/ops/cache_utils.py

@@ -366,7 +366,7 @@ def dequantize_and_gather_k_cache(
 ) -> None:
     if has_cutedsl():
         # lazily import, otherwise some tests fail due to CUDA driver init failure.
-        from vllm.models.deepseek_v4.nvidia.ops.dequant_gather_k_cutedsl import (
+        from vllm.models.deepseek_v4.nvidia.ops import (
             dequantize_and_gather_k_cache_cutedsl,
         )
 

vllm/models/deepseek_v4/common/ops/fused_compress_quant_cache.py

@@ -19,11 +19,93 @@ even/odd halves, producing (N_QUANT_BLOCKS, MXFP4_BLOCK/2) packed nibbles
 and N_QUANT_BLOCKS ue8m0 bytes.
 """
 
+from typing import Any
+
+import torch
+
 from vllm.triton_utils import tl, triton
 
 from .fused_indexer_q import _fp32x2_to_fp4x2
 
 
+def compress_norm_rope_store_triton(
+    state_cache: torch.Tensor,
+    num_actual: int,
+    token_to_req_indices: torch.Tensor,
+    positions: torch.Tensor,
+    slot_mapping: torch.Tensor,
+    block_table: torch.Tensor,
+    block_size: int,
+    state_width: int,
+    cos_sin_cache: torch.Tensor,
+    kv_cache: torch.Tensor,
+    k_cache_metadata: Any,
+    pdl_kwargs: dict,
+    head_dim: int,
+    rope_head_dim: int,
+    compress_ratio: int,
+    overlap: bool,
+    use_fp4_cache: bool,
+    rms_norm_weight: torch.Tensor,
+    rms_norm_eps: float,
+    quant_block: int,
+    token_stride: int,
+    scale_dim: int,
+) -> None:
+    """Shared triton launcher for the fused compress+norm+RoPE+insert path.
+
+    Picks one of the three kernels in this module based on ``head_dim`` and
+    ``use_fp4_cache``. Identical launch signature for all three.
+    """
+    if head_dim == 512:
+        kernel = _fused_kv_compress_norm_rope_insert_sparse_attn
+        num_warps = 4
+    elif use_fp4_cache:
+        kernel = _fused_kv_compress_norm_rope_insert_indexer_mxfp4_attn
+        num_warps = 1
+    else:
+        kernel = _fused_kv_compress_norm_rope_insert_indexer_attn
+        num_warps = 1
+
+    kernel[(num_actual,)](
+        # state cache
+        state_cache,
+        state_cache.stride(0),
+        state_cache.stride(1),
+        # metadata
+        token_to_req_indices,
+        positions,
+        slot_mapping,
+        block_table,
+        block_table.stride(0),
+        block_size,
+        # RMSNorm
+        rms_norm_weight,
+        rms_norm_eps,
+        # RoPE
+        cos_sin_cache,
+        cos_sin_cache.stride(0),
+        # KV cache
+        kv_cache,
+        k_cache_metadata.slot_mapping,
+        kv_cache.shape[1],  # paged KV cache block size (tokens per block)
+        # constexprs
+        HEAD_SIZE=head_dim,
+        TRITON_BLOCK_SIZE=triton.next_power_of_2(head_dim),
+        STATE_WIDTH=state_width,
+        COMPRESS_RATIO=compress_ratio,
+        OVERLAP=overlap,
+        ROPE_HEAD_DIM=rope_head_dim,
+        FP8_MAX=448.0,
+        QUANT_BLOCK=quant_block,
+        TOKEN_STRIDE=token_stride,
+        SCALE_DIM=scale_dim,
+        KV_BLOCK_STRIDE=kv_cache.stride(0),
+        num_warps=num_warps,
+        **pdl_kwargs,
+    )
+
+
 # =============================================================================
 # DeepseekV4 Attention path (head=512, nope=448 FP8 + rope=64 bf16)
 # =============================================================================

vllm/models/deepseek_v4/common/ops/fused_indexer_q.py

@@ -346,7 +346,7 @@ def fused_indexer_q_rope_quant(
         )
         if has_cutedsl():
             # lazily import, otherwise some tests fail due to CUDA driver init failure.
-            from vllm.models.deepseek_v4.nvidia.ops.fused_indexer_q_cutedsl import (
+            from vllm.models.deepseek_v4.nvidia.ops import (
                 fused_indexer_q_rope_quant_mxfp4_cutedsl,
             )
 
@@ -400,7 +400,7 @@ def fused_indexer_q_rope_quant(
     index_q_fp8 = torch.empty_like(index_q, dtype=torch.float8_e4m3fn)
     if has_cutedsl():
         # lazily import, otherwise some tests fail due to CUDA driver init failure.
-        from vllm.models.deepseek_v4.nvidia.ops.fused_indexer_q_cutedsl import (
+        from vllm.models.deepseek_v4.nvidia.ops import (
             fused_indexer_q_rope_quant_fp8_cutedsl,
         )
 

vllm/models/deepseek_v4/common/ops/save_partial_states.py

@@ -0,0 +1,101 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import torch
+
+from vllm.triton_utils import tl, triton
+
+
+def save_partial_states(
+    kv: torch.Tensor,
+    score: torch.Tensor,
+    ape: torch.Tensor,
+    positions: torch.Tensor,
+    state_cache: torch.Tensor,
+    slot_mapping: torch.Tensor,
+    block_size: int,
+    state_width: int,
+    compress_ratio: int,
+    pdl_kwargs: dict | None = None,
+) -> None:
+    """Write packed [kv, score+ape] partial states into the compressor cache.
+
+    One program per token; pads (slot_id == -1) are skipped.
+    """
+    num_actual = slot_mapping.shape[0]
+    head_size = kv.shape[-1]
+    _save_partial_states_kernel[(num_actual,)](
+        kv,
+        kv.stride(0),
+        score,
+        score.stride(0),
+        ape,
+        ape.stride(0),
+        positions,
+        state_cache,
+        state_cache.stride(0),
+        state_cache.stride(1),
+        slot_mapping,
+        block_size,
+        HEAD_SIZE=head_size,
+        TRITON_BLOCK_SIZE=triton.next_power_of_2(head_size),
+        STATE_WIDTH=state_width,
+        COMPRESS_RATIO=compress_ratio,
+        **(pdl_kwargs or {}),
+    )
+
+
+@triton.jit
+def _save_partial_states_kernel(
+    kv_ptr,
+    kv_stride,
+    score_ptr,
+    score_stride,
+    ape_ptr,
+    ape_stride,
+    positions_ptr,
+    state_cache_ptr,
+    state_cache_stride0,
+    state_cache_stride1,
+    slot_mapping_ptr,
+    block_size,
+    HEAD_SIZE: tl.constexpr,
+    TRITON_BLOCK_SIZE: tl.constexpr,
+    # state_cache last dim packs [kv_state, score_state], each STATE_WIDTH wide.
+    STATE_WIDTH: tl.constexpr,
+    COMPRESS_RATIO: tl.constexpr,
+):
+    token_idx = tl.program_id(0)
+    slot_id = tl.load(slot_mapping_ptr + token_idx)
+
+    # Skip padded / invalid tokens (slot_id == -1 is the PAD sentinel used
+    # by vLLM).  During CUDA graph replay the batch may contain padding
+    # tokens whose slot_mapping is -1; writing to kv_state[-1] would be an
+    # illegal memory access.
+    if slot_id < 0:
+        return
+
+    block_idx = slot_id // block_size
+    pos_in_block = slot_id % block_size
+    base_ptr = (
+        state_cache_ptr
+        + block_idx * state_cache_stride0
+        + pos_in_block * state_cache_stride1
+    )
+
+    block = tl.arange(0, TRITON_BLOCK_SIZE)
+    mask = block < HEAD_SIZE
+
+    kv = tl.load(kv_ptr + token_idx * kv_stride + block, mask=mask)
+    tl.store(base_ptr + block, kv, mask=mask)
+
+    # Fused: score += ape[position % compress_ratio]
+    position = tl.load(positions_ptr + token_idx)
+    ape_row = position % COMPRESS_RATIO
+    ape = tl.load(ape_ptr + ape_row * ape_stride + block, mask=mask)
+    score = tl.load(score_ptr + token_idx * score_stride + block, mask=mask)
+    tl.store(
+        base_ptr + STATE_WIDTH + block,
+        score + ape,
+        mask=mask,
+    )

vllm/models/deepseek_v4/compressor.py

@@ -13,13 +13,13 @@ from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
 from vllm.model_executor.layers.layernorm import RMSNorm
 from vllm.model_executor.layers.linear import MergedColumnParallelLinear
 from vllm.models.deepseek_v4.common.ops.fused_compress_quant_cache import (
-    _fused_kv_compress_norm_rope_insert_indexer_attn,
-    _fused_kv_compress_norm_rope_insert_indexer_mxfp4_attn,
-    _fused_kv_compress_norm_rope_insert_sparse_attn,
+    compress_norm_rope_store_triton,
 )
 from vllm.models.deepseek_v4.common.ops.fused_indexer_q import MXFP4_BLOCK_SIZE
+from vllm.models.deepseek_v4.common.ops.save_partial_states import (
+    save_partial_states,
+)
 from vllm.platforms import current_platform
-from vllm.triton_utils import tl, triton
 from vllm.v1.attention.backend import (
     AttentionBackend,
     AttentionCGSupport,
@@ -171,6 +171,16 @@ class CompressorStateCache(torch.nn.Module, AttentionLayerBase):
 
 
 class DeepseekCompressor(nn.Module):
+    """DeepSeek V4 KV/score compressor.
+
+    Owns the linear / norm / state-cache / ape state and the shared forward
+    prologue (kv/score split, save_partial_states launch). The
+    compress → norm → RoPE → store step is dispatched to a triton kernel
+    (``compress_norm_rope_store_triton``) by default, except for the NVIDIA
+    head_dim=128 indexer path which uses the cutedsl kernel
+    (``compress_norm_rope_store_cutedsl``) for better performance.
+    """
+
     def __init__(
         self,
         vllm_config: VllmConfig,
@@ -240,25 +250,18 @@ class DeepseekCompressor(nn.Module):
             assert not use_fp4_cache, (
                 "MXFP4 cache is only supported for indexer (head=128)"
             )
-            self._fused_kernel = _fused_kv_compress_norm_rope_insert_sparse_attn
             self._quant_block = 64
             self._token_stride = self.nope_head_dim + self.rope_head_dim * 2
             self._scale_dim = self.nope_head_dim // 64 + 1  # 7 real + 1 pad
-            self._num_warps = 4
         elif self.head_dim == 128:
             if use_fp4_cache:
-                self._fused_kernel = (
-                    _fused_kv_compress_norm_rope_insert_indexer_mxfp4_attn
-                )
                 self._quant_block = MXFP4_BLOCK_SIZE
                 self._token_stride = self.head_dim // 2
                 self._scale_dim = self.head_dim // MXFP4_BLOCK_SIZE
             else:
-                self._fused_kernel = _fused_kv_compress_norm_rope_insert_indexer_attn
                 self._quant_block = 128
                 self._token_stride = self.head_dim
                 self._scale_dim = 4  # single float32 scale
-            self._num_warps = 1
         else:
             raise ValueError(
                 f"Unsupported head_dim for fused quant+cache: {self.head_dim}"
@@ -303,29 +306,22 @@ class DeepseekCompressor(nn.Module):
         )
 
         # Store the KV and score (with fused APE addition) in the state.
-        # NOTE: PDL is disabled — both this kernel and _fused_kernel below
-        # depend on preceding kernel outputs (kv/score from the cublas GEMM;
-        # state_cache from this kernel) but neither emits/waits on PDL grid
-        # dependency primitives, so launch_pdl=True caused a read-after-write
-        # race and non-deterministic output.
-        _save_partial_states_kernel[(num_actual,)](
-            kv,
-            kv.stride(0),
-            score,
-            score.stride(0),
-            self.ape,
-            self.ape.stride(0),
-            positions,
-            state_cache,
-            state_cache.stride(0),
-            state_cache.stride(1),
-            slot_mapping,
-            block_size,
-            HEAD_SIZE=kv.shape[-1],
-            TRITON_BLOCK_SIZE=triton.next_power_of_2(kv.shape[-1]),
-            STATE_WIDTH=state_width,
-            COMPRESS_RATIO=self.compress_ratio,
-            **pdl_kwargs,
+        # NOTE: PDL is disabled — both this kernel and the compress kernels
+        # below depend on preceding kernel outputs (kv/score from the cublas
+        # GEMM; state_cache from this kernel) but neither emits/waits on PDL
+        # grid dependency primitives, so launch_pdl=True caused a
+        # read-after-write race and non-deterministic output.
+        save_partial_states(
+            kv=kv,
+            score=score,
+            ape=self.ape,
+            positions=positions,
+            state_cache=state_cache,
+            slot_mapping=slot_mapping,
+            block_size=block_size,
+            state_width=state_width,
+            compress_ratio=self.compress_ratio,
+            pdl_kwargs=pdl_kwargs,
         )
 
         # Fused: compress → RMSNorm → RoPE → FP8 quant → KV cache write.
@@ -339,96 +335,44 @@ class DeepseekCompressor(nn.Module):
         k_cache_metadata = cast(Any, attn_metadata[self.k_cache_prefix])
         kv_cache = self._static_forward_context[self.k_cache_prefix].kv_cache
 
-        self._fused_kernel[(num_actual,)](
-            # state cache
-            state_cache,
-            state_cache.stride(0),
-            state_cache.stride(1),
-            # metadata
-            token_to_req_indices,
-            positions,
-            slot_mapping,
-            block_table,
-            block_table.stride(0),
-            block_size,
-            # RMSNorm
-            self.norm.weight,
-            self.rms_norm_eps,
-            # RoPE
-            cos_sin_cache,
-            cos_sin_cache.stride(0),
-            # KV cache
-            kv_cache,
-            k_cache_metadata.slot_mapping,
-            kv_cache.shape[1],  # paged KV cache block size (tokens per block)
-            # constexprs
-            HEAD_SIZE=self.head_dim,
-            TRITON_BLOCK_SIZE=triton.next_power_of_2(self.head_dim),
-            STATE_WIDTH=state_width,
-            COMPRESS_RATIO=self.compress_ratio,
-            OVERLAP=self.overlap,
-            ROPE_HEAD_DIM=self.rope_head_dim,
-            FP8_MAX=448.0,
-            QUANT_BLOCK=self._quant_block,
-            TOKEN_STRIDE=self._token_stride,
-            SCALE_DIM=self._scale_dim,
-            KV_BLOCK_STRIDE=kv_cache.stride(0),
-            num_warps=self._num_warps,
-            **pdl_kwargs,
+        if current_platform.is_cuda():
+            # NVIDIA GPUs.
+            if self.head_dim == 512:
+                from .nvidia.ops import compress_norm_rope_store_cutedsl
+
+                # Main compressor path.
+                # Use a cutedsl kernel for better performance.
+                compress_norm_rope_store_fn = compress_norm_rope_store_cutedsl
+            else:
+                # Indexer path (head_dim == 128).
+                # Use a triton kernel.
+                compress_norm_rope_store_fn = compress_norm_rope_store_triton
+        else:
+            # AMD GPUs.
+            # Always use a triton kernel.
+            compress_norm_rope_store_fn = compress_norm_rope_store_triton
+
+        compress_norm_rope_store_fn(
+            state_cache=state_cache,
+            num_actual=num_actual,
+            token_to_req_indices=token_to_req_indices,
+            positions=positions,
+            slot_mapping=slot_mapping,
+            block_table=block_table,
+            block_size=block_size,
+            state_width=state_width,
+            cos_sin_cache=cos_sin_cache,
+            kv_cache=kv_cache,
+            k_cache_metadata=k_cache_metadata,
+            pdl_kwargs=pdl_kwargs,
+            head_dim=self.head_dim,
+            rope_head_dim=self.rope_head_dim,
+            compress_ratio=self.compress_ratio,
+            overlap=self.overlap,
+            use_fp4_cache=self.use_fp4_cache,
+            rms_norm_weight=self.norm.weight,
+            rms_norm_eps=self.rms_norm_eps,
+            quant_block=self._quant_block,
+            token_stride=self._token_stride,
+            scale_dim=self._scale_dim,
         )
-
-
-@triton.jit
-def _save_partial_states_kernel(
-    kv_ptr,
-    kv_stride,
-    score_ptr,
-    score_stride,
-    ape_ptr,
-    ape_stride,
-    positions_ptr,
-    state_cache_ptr,
-    state_cache_stride0,
-    state_cache_stride1,
-    slot_mapping_ptr,
-    block_size,
-    HEAD_SIZE: tl.constexpr,
-    TRITON_BLOCK_SIZE: tl.constexpr,
-    # state_cache last dim packs [kv_state, score_state], each STATE_WIDTH wide.
-    STATE_WIDTH: tl.constexpr,
-    COMPRESS_RATIO: tl.constexpr,
-):
-    token_idx = tl.program_id(0)
-    slot_id = tl.load(slot_mapping_ptr + token_idx)
-
-    # Skip padded / invalid tokens (slot_id == -1 is the PAD sentinel used
-    # by vLLM).  During CUDA graph replay the batch may contain padding
-    # tokens whose slot_mapping is -1; writing to kv_state[-1] would be an
-    # illegal memory access.
-    if slot_id < 0:
-        return
-
-    block_idx = slot_id // block_size
-    pos_in_block = slot_id % block_size
-    base_ptr = (
-        state_cache_ptr
-        + block_idx * state_cache_stride0
-        + pos_in_block * state_cache_stride1
-    )
-
-    block = tl.arange(0, TRITON_BLOCK_SIZE)
-    mask = block < HEAD_SIZE
-
-    kv = tl.load(kv_ptr + token_idx * kv_stride + block, mask=mask)
-    tl.store(base_ptr + block, kv, mask=mask)
-
-    # Fused: score += ape[position % compress_ratio]
-    position = tl.load(positions_ptr + token_idx)
-    ape_row = position % COMPRESS_RATIO
-    ape = tl.load(ape_ptr + ape_row * ape_stride + block, mask=mask)
-    score = tl.load(score_ptr + token_idx * score_stride + block, mask=mask)
-    tl.store(
-        base_ptr + STATE_WIDTH + block,
-        score + ape,
-        mask=mask,
-    )

vllm/models/deepseek_v4/nvidia/flashmla.py

@@ -28,7 +28,7 @@ from vllm.v1.attention.ops.flashmla import (
 from vllm.v1.worker.workspace import current_workspace_manager
 
 if TYPE_CHECKING:
-    from vllm.models.deepseek_v4.nvidia.ops.attention import (
+    from vllm.models.deepseek_v4.attention import (
         DeepseekV4MLAAttention,
     )
     from vllm.v1.attention.backends.mla.sparse_swa import DeepseekSparseSWAMetadata
@@ -63,6 +63,18 @@ class DeepseekV4SparseMLAAttentionImpl(SparseMLAAttentionImpl[FlashMLASparseMeta
     ) -> None:
         raise NotImplementedError
 
+    @classmethod
+    @abstractmethod
+    def get_padded_num_q_heads(cls, num_heads: int) -> int:
+        """Q head count the backend wants q allocated at.
+
+        The MLA wrapper allocates the q/output buffers at
+        ``[N, get_padded_num_q_heads(n_local_heads), head_dim]``. Must
+        satisfy ``result >= num_heads``. Backends with no padding constraint
+        return ``num_heads``.
+        """
+        raise NotImplementedError
+
 
 class DeepseekV4FlashMLASparseBackend(FlashMLASparseBackend):
     @staticmethod
@@ -104,6 +116,16 @@ class DeepseekV4FlashMLASparseImpl(DeepseekV4SparseMLAAttentionImpl):
 
     backend_cls = DeepseekV4FlashMLASparseBackend
 
+    @classmethod
+    def get_padded_num_q_heads(cls, num_heads: int) -> int:
+        # FP8 decode kernel only supports h_q = 64 or 128.
+        if num_heads > 128:
+            raise ValueError(
+                f"DeepseekV4 FlashMLA does not support {num_heads} heads "
+                "(FP8 decode kernel requires h_q in {64, 128})."
+            )
+        return 64 if num_heads <= 64 else 128
+
     @classmethod
     def forward_mqa(  # type: ignore[override]
         cls,

vllm/models/deepseek_v4/nvidia/model.py

@@ -54,12 +54,12 @@ from vllm.model_executor.models.utils import (
     maybe_prefix,
 )
 from vllm.model_executor.utils import set_weight_attrs
-from vllm.models.deepseek_v4.nvidia.ops.attention import (
+from vllm.models.deepseek_v4.attention import (
     DeepseekV4Indexer,
     DeepseekV4MLAModules,
     DeepseekV4MultiHeadLatentAttentionWrapper,
 )
-from vllm.models.deepseek_v4.nvidia.ops.prepare_megamoe import prepare_megamoe_inputs
+from vllm.models.deepseek_v4.nvidia.ops import prepare_megamoe_inputs
 from vllm.platforms import current_platform
 from vllm.sequence import IntermediateTensors
 from vllm.utils.torch_utils import direct_register_custom_op

vllm/models/deepseek_v4/nvidia/ops/__init__.py

@@ -6,3 +6,19 @@ These modules import ``cutlass``/``cutedsl`` at module top level, so they must
 not be imported on non-CUDA platforms. Callers should gate on
 ``vllm.utils.import_utils.has_cutedsl()`` before importing from here.
 """
+
+from .dequant_gather_k_cutedsl import dequantize_and_gather_k_cache_cutedsl
+from .fused_indexer_q_cutedsl import (
+    fused_indexer_q_rope_quant_fp8_cutedsl,
+    fused_indexer_q_rope_quant_mxfp4_cutedsl,
+)
+from .prepare_megamoe import prepare_megamoe_inputs
+from .sparse_attn_compress_cutedsl import compress_norm_rope_store_cutedsl
+
+__all__ = [
+    "compress_norm_rope_store_cutedsl",
+    "dequantize_and_gather_k_cache_cutedsl",
+    "fused_indexer_q_rope_quant_fp8_cutedsl",
+    "fused_indexer_q_rope_quant_mxfp4_cutedsl",
+    "prepare_megamoe_inputs",
+]

vllm/models/deepseek_v4/nvidia/ops/dequant_gather_k_cutedsl.py

@@ -11,10 +11,7 @@ from cutlass import BFloat16, Int32, Uint8, Uint32
 from cutlass.cute.nvgpu import cpasync
 from quack.compile_utils import make_fake_tensor
 
-from vllm.models.deepseek_v4.nvidia.ops.cutedsl_utils import (
-    _bf16x2_mul,
-    _fp8x4_to_bf16x4,
-)
+from vllm.cute_utils import _bf16x2_mul, cvt
 
 
 def dequantize_and_gather_k_cache_cutedsl(
@@ -268,8 +265,8 @@ class DequantGatherKCacheKernel:
             dequant0 = cute.make_rmem_tensor(4, Uint32)
             dequant1 = cute.make_rmem_tensor(4, Uint32)
             for j in cutlass.range_constexpr(2):
-                tmp0 = _fp8x4_to_bf16x4(data0[j])
-                tmp1 = _fp8x4_to_bf16x4(data1[j])
+                tmp0 = cvt.fp8x4_to_bf16x4(data0[j])
+                tmp1 = cvt.fp8x4_to_bf16x4(data1[j])
 
                 # BF16 multiply is safe because the scales are exact powers of 2.
                 dequant0[j * 2] = _bf16x2_mul(tmp0[0], scale0_bf16x2)

vllm/models/deepseek_v4/nvidia/ops/fused_indexer_q_cutedsl.py

@@ -9,14 +9,11 @@ from cuda.bindings.driver import CUstream
 from cutlass import BFloat16, Float32, Int64, Uint8, Uint32, const_expr
 from quack.compile_utils import make_fake_tensor
 
-from vllm.models.deepseek_v4.nvidia.ops.cutedsl_utils import (
+from vllm.cute_utils import (
     _bf16x2_abs,
     _bf16x2_max,
-    _bf16x2_to_fp32,
-    _fp32x2_to_bf16x2,
-    _fp32x4_to_fp8x4,
-    _fp32x8_to_fp4x8,
-    _recast_val,
+    cvt,
+    recast_val,
 )
 from vllm.vllm_flash_attn.cute import utils as cute_utils
 
@@ -225,8 +222,8 @@ class IndexerQRopeQuantKernel:
                 cute.copy(cp_u32x4, cute.recast_tensor(sin_src, Uint32), sin_bf16x2)
 
                 for i in cutlass.range_constexpr(4):
-                    cos0, cos1 = _bf16x2_to_fp32(cos_bf16x2[i])
-                    sin0, sin1 = _bf16x2_to_fp32(sin_bf16x2[i])
+                    cos0, cos1 = cvt.bf16x2_to_fp32x2(cos_bf16x2[i])
+                    sin0, sin1 = cvt.bf16x2_to_fp32x2(sin_bf16x2[i])
                     cos_vals[i * 2] = cos0
                     cos_vals[i * 2 + 1] = cos1
                     sin_vals[i * 2] = sin0
@@ -234,11 +231,11 @@ class IndexerQRopeQuantKernel:
 
             for i in cutlass.range_constexpr(self.coarsen):
                 for j in cutlass.range_constexpr(8):
-                    q0, q1 = _bf16x2_to_fp32(q_bf16x2[i, j])
+                    q0, q1 = cvt.bf16x2_to_fp32x2(q_bf16x2[i, j])
                     rot0 = q0 * cos_vals[j] - q1 * sin_vals[j]
                     rot1 = q0 * sin_vals[j] + q1 * cos_vals[j]
                     # convert back to BF16 to match numerics
-                    q_bf16x2[i, j] = _fp32x2_to_bf16x2(rot0, rot1)
+                    q_bf16x2[i, j] = cvt.fp32x2_to_bf16x2(rot0, rot1)
 
         return (
             q_bf16x2,
@@ -327,7 +324,7 @@ class IndexerQMxFp4Kernel(IndexerQRopeQuantKernel):
                 _bf16x2_max,
                 width=MXFP4_BLOCK_SIZE // 16,
             )
-            amax_pair = _bf16x2_to_fp32(amax_bf16x2)
+            amax_pair = cvt.bf16x2_to_fp32x2(amax_bf16x2)
             amax = cute_utils.fmax(amax_pair[0], amax_pair[1])
 
             if in_bounds:
@@ -336,7 +333,7 @@ class IndexerQMxFp4Kernel(IndexerQRopeQuantKernel):
                 # increments the exponent whenever fp4_scale is not exactly a power of 2
                 eps = cutlass.const_expr(float.fromhex("0x6p-126"))
                 fp4_scale = cute_utils.fmax(amax, eps) * Float32(1.0 / 6.0)
-                bits = _recast_val(fp4_scale, Uint32)
+                bits = recast_val(fp4_scale, Uint32)
                 ue8m0 = cute_utils.shr_u32(
                     bits + Uint32(0x7FFFFF), Uint32(23)
                 ) & Uint32(0xFF)
@@ -349,18 +346,18 @@ class IndexerQMxFp4Kernel(IndexerQRopeQuantKernel):
                 # If scale = 2^A and ue8m0 = A + 127, then inverse scale has exponent
                 # -A + 127 = 254 - ue8m0.
                 inv_scale_bits = (Uint32(254) - ue8m0) << Uint32(23)
-                inv_fp4_scale = _recast_val(inv_scale_bits, Float32)
+                inv_fp4_scale = recast_val(inv_scale_bits, Float32)
 
                 vals = cute.make_rmem_tensor(16, Float32)
                 for j in cutlass.range_constexpr(8):
-                    q0, q1 = _bf16x2_to_fp32(q_bf16x2[i, j])
+                    q0, q1 = cvt.bf16x2_to_fp32x2(q_bf16x2[i, j])
                     vals[j * 2] = q0 * inv_fp4_scale
                     vals[j * 2 + 1] = q1 * inv_fp4_scale
 
                 # pack to FP4
                 packed = cute.make_rmem_tensor((2,), Uint32)
-                packed[0] = _fp32x8_to_fp4x8(vals, 0)
-                packed[1] = _fp32x8_to_fp4x8(vals, 8)
+                packed[0] = cvt.fp32x8_to_fp4x8(vals, 0)
+                packed[1] = cvt.fp32x8_to_fp4x8(vals, 8)
 
                 dst = q_fp4_tile[i, None]
                 cp_u32x2 = cute.make_copy_atom(cp_op, Uint32, num_bits_per_copy=64)
@@ -519,24 +516,24 @@ class IndexerQFp8Kernel(IndexerQRopeQuantKernel):
                 _bf16x2_max,
                 width=self.subwarp_size,
             )
-            amax_pair = _bf16x2_to_fp32(amax_bf16x2)
+            amax_pair = cvt.bf16x2_to_fp32x2(amax_bf16x2)
             amax = cute_utils.fmax(amax_pair[0], amax_pair[1])
 
             # scale = max(amax, eps) / fp8_max, then rounded UP to the next
             # power of two. Adding the mantissa mask before shifting out the
             # mantissa bumps the exponent whenever s isn't a pure pow2.
             fp32_scale = cute_utils.fmax(amax, Float32(1e-4)) * Float32(1.0 / 448.0)
-            bits = _recast_val(fp32_scale, Uint32)
+            bits = recast_val(fp32_scale, Uint32)
             scale_exp = cute_utils.shr_u32(
                 bits + Uint32(0x7FFFFF), Uint32(23)
             ) & Uint32(0xFF)
 
             # rounded scale = 2^(scale_exp - 127); bit pattern is scale_exp << 23
             fp8_scale_bits = scale_exp << Uint32(23)
-            fp8_scale = _recast_val(fp8_scale_bits, Float32)
+            fp8_scale = recast_val(fp8_scale_bits, Float32)
             # inverse = 2^-(scale_exp - 127); bit pattern is (254 - scale_exp) << 23
             inv_scale_bits = (Uint32(254) - scale_exp) << Uint32(23)
-            inv_fp8_scale = _recast_val(inv_scale_bits, Float32)
+            inv_fp8_scale = recast_val(inv_scale_bits, Float32)
 
             # Weight fold: weights_out = weights * q_scale * scale_combined.
             # All threads in the subwarp share the same fp8_scale after the
@@ -553,9 +550,9 @@ class IndexerQFp8Kernel(IndexerQRopeQuantKernel):
                 # (one cp.async-shaped 128-bit store per row).
                 packed = cute.make_rmem_tensor((4,), Uint32)
                 for j in cutlass.range_constexpr(4):
-                    q0, q1 = _bf16x2_to_fp32(q_bf16x2[i, j * 2])
-                    q2, q3 = _bf16x2_to_fp32(q_bf16x2[i, j * 2 + 1])
-                    packed[j] = _fp32x4_to_fp8x4(
+                    q0, q1 = cvt.bf16x2_to_fp32x2(q_bf16x2[i, j * 2])
+                    q2, q3 = cvt.bf16x2_to_fp32x2(q_bf16x2[i, j * 2 + 1])
+                    packed[j] = cvt.fp32x4_to_fp8x4(
                         q0 * inv_fp8_scale,
                         q1 * inv_fp8_scale,
                         q2 * inv_fp8_scale,

vllm/v1/attention/backends/gdn_attn.py

@@ -3,6 +3,7 @@
 """Backend for GatedDeltaNet attention."""
 
 from dataclasses import dataclass
+from typing import Literal
 
 import torch
 
@@ -90,6 +91,12 @@ class GDNAttentionMetadataBuilder(AttentionMetadataBuilder[GDNAttentionMetadata]
         self.compilation_config = vllm_config.compilation_config
         self.speculative_config = vllm_config.speculative_config
         self.kv_cache_spec = kv_cache_spec
+        from vllm.model_executor.layers.mamba.gdn.qwen_gdn_linear_attn import (
+            _resolve_gdn_prefill_backend,
+        )
+
+        self.gdn_prefill_backend: Literal["triton", "flashinfer", "cutedsl"]
+        _, self.gdn_prefill_backend = _resolve_gdn_prefill_backend(vllm_config)
 
         if self.speculative_config:
             assert self.speculative_config.num_speculative_tokens is not None
@@ -316,22 +323,38 @@ class GDNAttentionMetadataBuilder(AttentionMetadataBuilder[GDNAttentionMetadata]
         chunk_indices: torch.Tensor | None = None
         chunk_offsets: torch.Tensor | None = None
         if num_prefills > 0:
-            # Only prefill batches use FLA chunk ops.
-            # Pre-compute on CPU and async-copy to GPU to avoid
-            # GPU→CPU sync (.tolist()) in prepare_chunk_indices.
-            from vllm.model_executor.layers.fla.ops.index import (
-                prepare_chunk_indices,
-                prepare_chunk_offsets,
-            )
             from vllm.model_executor.layers.fla.ops.utils import FLA_CHUNK_SIZE
 
-            gpu_device = query_start_loc.device
-            chunk_indices = prepare_chunk_indices(
-                non_spec_query_start_loc_cpu, FLA_CHUNK_SIZE
-            ).to(device=gpu_device, non_blocking=True)
-            chunk_offsets = prepare_chunk_offsets(
-                non_spec_query_start_loc_cpu, FLA_CHUNK_SIZE
-            ).to(device=gpu_device, non_blocking=True)
+            if self.gdn_prefill_backend == "cutedsl":
+                from vllm.model_executor.layers.mamba.ops.gdn_chunk_cutedsl import (
+                    prepare_metadata_cutedsl,
+                )
+
+                assert non_spec_query_start_loc is not None
+                assert non_spec_query_start_loc_cpu is not None
+                total_tokens = int(non_spec_query_start_loc_cpu[-1].item())
+                chunk_indices, chunk_offsets = prepare_metadata_cutedsl(
+                    non_spec_query_start_loc,
+                    total_tokens,
+                    FLA_CHUNK_SIZE,
+                )
+            else:
+                gpu_device = query_start_loc.device
+                # Only prefill batches use FLA chunk ops.
+                # Pre-compute on CPU and async-copy to GPU to avoid
+                # GPU→CPU sync (.tolist()) in prepare_chunk_indices.
+                from vllm.model_executor.layers.fla.ops.index import (
+                    prepare_chunk_indices,
+                    prepare_chunk_offsets,
+                )
+
+                assert non_spec_query_start_loc_cpu is not None
+                chunk_indices = prepare_chunk_indices(
+                    non_spec_query_start_loc_cpu, FLA_CHUNK_SIZE
+                ).to(device=gpu_device, non_blocking=True)
+                chunk_offsets = prepare_chunk_offsets(
+                    non_spec_query_start_loc_cpu, FLA_CHUNK_SIZE
+                ).to(device=gpu_device, non_blocking=True)
 
         if num_prefills > 0:
             has_initial_state = context_lens_tensor > 0

vllm/v1/engine/async_llm.py

@@ -81,7 +81,7 @@ class AsyncLLM(EngineClient):
         start_engine_loop: bool = True,
         stat_loggers: list[StatLoggerFactory] | None = None,
         aggregate_engine_logging: bool = False,
-        client_addresses: dict[str, str] | None = None,
+        client_addresses: dict[str, Any] | None = None,
         client_count: int = 1,
         client_index: int = 0,
     ) -> None:
@@ -209,7 +209,7 @@ class AsyncLLM(EngineClient):
         enable_log_requests: bool = False,
         aggregate_engine_logging: bool = False,
         disable_log_stats: bool = False,
-        client_addresses: dict[str, str] | None = None,
+        client_addresses: dict[str, Any] | None = None,
         client_count: int = 1,
         client_index: int = 0,
     ) -> "AsyncLLM":

vllm/v1/engine/coordinator.py

@@ -11,7 +11,7 @@ import zmq
 
 from vllm.config import ParallelConfig
 from vllm.logger import init_logger
-from vllm.utils.network_utils import get_tcp_uri, make_zmq_socket
+from vllm.utils.network_utils import make_zmq_socket
 from vllm.utils.system_utils import get_mp_context, set_process_title
 from vllm.v1.engine import EngineCoreOutputs, EngineCoreRequestType
 from vllm.v1.serial_utils import MsgpackDecoder
@@ -91,16 +91,9 @@ class DPCoordinator:
         if parallel_config.enable_elastic_ep:
             local_only_eng = False
 
-        def bind_address(local_only: bool) -> str:
-            return (
-                get_engine_client_zmq_addr(local_only=True, host=host)
-                if local_only
-                else get_tcp_uri(host, 0)
-            )
-
-        front_publish_address = bind_address(local_only)
-        back_publish_address = bind_address(local_only_eng)
-        back_output_address = bind_address(local_only_eng)
+        front_publish_address = get_engine_client_zmq_addr(local_only, host=host)
+        back_publish_address = get_engine_client_zmq_addr(local_only_eng, host=host)
+        back_output_address = get_engine_client_zmq_addr(local_only_eng, host=host)
 
         context = get_mp_context()
         parent_zmq_addr_pipe, child_zmq_addr_pipe = context.Pipe(duplex=False)

vllm/v1/engine/core_client.py

@@ -11,6 +11,7 @@ from collections import defaultdict, deque
 from collections.abc import Awaitable, Callable, Sequence
 from concurrent.futures import Future
 from dataclasses import dataclass
+from multiprocessing.connection import Connection
 from multiprocessing.queues import Queue
 from threading import Thread
 from typing import Any, TypeAlias, TypeVar
@@ -108,7 +109,7 @@ class EngineCoreClient(ABC):
         vllm_config: VllmConfig,
         executor_class: type[Executor],
         log_stats: bool,
-        client_addresses: dict[str, str] | None = None,
+        client_addresses: dict[str, Any] | None = None,
         client_count: int = 1,
         client_index: int = 0,
     ) -> "AsyncMPClient":
@@ -476,7 +477,7 @@ class MPClient(EngineCoreClient):
         vllm_config: VllmConfig,
         executor_class: type[Executor],
         log_stats: bool,
-        client_addresses: dict[str, str] | None = None,
+        client_addresses: dict[str, Any] | None = None,
     ):
         self.vllm_config = vllm_config
 
@@ -507,7 +508,7 @@ class MPClient(EngineCoreClient):
                 output_address = client_addresses["output_address"]
                 self.stats_update_address = client_addresses.get("stats_update_address")
                 # Tensor queues passed via client_addresses for multi-API-server case
-                tensor_queue = client_addresses.get("tensor_queue")  # type: ignore[assignment]
+                tensor_queue = client_addresses.get("tensor_queue")
                 self.input_socket = self.resources.input_socket = make_zmq_socket(
                     self.ctx,
                     input_address,
@@ -518,6 +519,28 @@ class MPClient(EngineCoreClient):
                 self.resources.output_socket = make_zmq_socket(
                     self.ctx, output_address, zmq.PULL
                 )
+
+                # Report bound endpoints back so the parent can forward
+                # them to engines (mirrors the DPCoordinator pattern).
+                actual_address_pipe: Connection | None = client_addresses.get(
+                    "actual_address_pipe"
+                )
+                if actual_address_pipe is not None:
+                    try:
+                        actual_input = self.input_socket.getsockopt(
+                            zmq.LAST_ENDPOINT
+                        ).decode()
+                        actual_output = self.resources.output_socket.getsockopt(
+                            zmq.LAST_ENDPOINT
+                        ).decode()
+                        actual_address_pipe.send(
+                            {
+                                "input_address": actual_input,
+                                "output_address": actual_output,
+                            }
+                        )
+                    finally:
+                        actual_address_pipe.close()
             else:
                 # Engines are managed by this client.
                 addresses = get_engine_zmq_addresses(vllm_config)
@@ -532,6 +555,15 @@ class MPClient(EngineCoreClient):
                     self.ctx, addresses.outputs[0], zmq.PULL
                 )
 
+                # Resolve ``tcp://host:0`` placeholders to bound endpoints
+                # before engines DEALER-connect. No-op for IPC.
+                addresses.inputs[0] = self.input_socket.getsockopt(
+                    zmq.LAST_ENDPOINT
+                ).decode()
+                addresses.outputs[0] = self.resources.output_socket.getsockopt(
+                    zmq.LAST_ENDPOINT
+                ).decode()
+
                 with launch_core_engines(
                     vllm_config, executor_class, log_stats, addresses
                 ) as (engine_manager, coordinator, addresses, tensor_queue):
@@ -893,7 +925,7 @@ class AsyncMPClient(MPClient):
         vllm_config: VllmConfig,
         executor_class: type[Executor],
         log_stats: bool,
-        client_addresses: dict[str, str] | None = None,
+        client_addresses: dict[str, Any] | None = None,
         client_count: int = 1,
         client_index: int = 0,
     ):
@@ -1143,7 +1175,7 @@ class DPAsyncMPClient(AsyncMPClient):
         vllm_config: VllmConfig,
         executor_class: type[Executor],
         log_stats: bool,
-        client_addresses: dict[str, str] | None = None,
+        client_addresses: dict[str, Any] | None = None,
         client_count: int = 1,
         client_index: int = 0,
     ):
@@ -1323,7 +1355,7 @@ class DPLBAsyncMPClient(DPAsyncMPClient):
         vllm_config: VllmConfig,
         executor_class: type[Executor],
         log_stats: bool,
-        client_addresses: dict[str, str] | None = None,
+        client_addresses: dict[str, Any] | None = None,
         client_count: int = 1,
         client_index: int = 0,
     ):

vllm/v1/engine/utils.py

@@ -23,7 +23,12 @@ from vllm.logger import init_logger
 from vllm.platforms import current_platform
 from vllm.ray.ray_env import get_env_vars_to_copy
 from vllm.utils import numa_utils
-from vllm.utils.network_utils import get_open_zmq_ipc_path, zmq_socket_ctx
+from vllm.utils.network_utils import (
+    get_open_port,
+    get_open_zmq_ipc_path,
+    get_tcp_uri,
+    zmq_socket_ctx,
+)
 from vllm.utils.system_utils import get_mp_context
 from vllm.v1.engine.coordinator import DPCoordinator
 from vllm.v1.executor import Executor
@@ -955,8 +960,19 @@ class CoreEngineActorManager:
 def get_engine_zmq_addresses(
     vllm_config: VllmConfig,
     num_api_servers: int = 1,
+    *,
+    defer_api_server_ports: bool = True,
 ) -> EngineZmqAddresses:
-    """Allocate ZMQ addresses for engine-client communication."""
+    """Allocate ZMQ addresses for engine-client communication.
+
+    By default each TCP address is a ``tcp://host:0`` placeholder; the
+    consumer (API-server child or single-process ``MPClient``) binds, then
+    recovers the kernel-assigned port via ``getsockopt(zmq.LAST_ENDPOINT)``
+    and writes it back into ``addresses`` before the engine handshake.
+
+    Set ``defer_api_server_ports=False`` only when the consumer cannot
+    report a bound port back (e.g. the Rust front-end). IPC paths are
+    unaffected."""
     parallel_config = vllm_config.parallel_config
     local_engine_count = parallel_config.data_parallel_size_local
     local_start_index = parallel_config.data_parallel_rank_local
@@ -978,15 +994,14 @@ def get_engine_zmq_addresses(
     if parallel_config.enable_elastic_ep:
         client_local_only = False
 
+    def _addr() -> str:
+        if client_local_only:
+            return get_open_zmq_ipc_path()
+        return get_tcp_uri(host, 0 if defer_api_server_ports else get_open_port())
+
     return EngineZmqAddresses(
-        inputs=[
-            get_engine_client_zmq_addr(client_local_only, host)
-            for _ in range(num_api_servers)
-        ],
-        outputs=[
-            get_engine_client_zmq_addr(client_local_only, host)
-            for _ in range(num_api_servers)
-        ],
+        inputs=[_addr() for _ in range(num_api_servers)],
+        outputs=[_addr() for _ in range(num_api_servers)],
     )
 
 
@@ -1095,9 +1110,11 @@ def launch_core_engines(
     if parallel_config.enable_elastic_ep:
         handshake_local_only = False
 
-    handshake_address = get_engine_client_zmq_addr(
-        handshake_local_only, host, parallel_config.data_parallel_rpc_port
-    )
+    # Preserve "port=0 means auto-pick" for the handshake address, which
+    # is consumed by engines spawned in this process and so cannot defer
+    # port resolution to bind time.
+    rpc_port = parallel_config.data_parallel_rpc_port or get_open_port()
+    handshake_address = get_engine_client_zmq_addr(handshake_local_only, host, rpc_port)
 
     if local_engines_only and dp_rank > 0:
         assert not handshake_local_only

vllm/v1/utils.py

@@ -29,7 +29,7 @@ from torch.autograd.profiler import record_function
 import vllm.envs as envs
 from vllm.logger import init_logger
 from vllm.usage.usage_lib import UsageContext, is_usage_stats_enabled, usage_message
-from vllm.utils.network_utils import get_open_port, get_open_zmq_ipc_path, get_tcp_uri
+from vllm.utils.network_utils import get_open_zmq_ipc_path, get_tcp_uri
 from vllm.utils.system_utils import decorate_logs, kill_process_tree, set_process_title
 from vllm.v1.core.sched.output import SchedulerOutput
 
@@ -144,20 +144,18 @@ class CpuGpuBuffer:
         return self.cpu[:n].copy_(self.gpu[:n], non_blocking=True)
 
 
-def get_engine_client_zmq_addr(local_only: bool, host: str, port: int = 0) -> str:
-    """Assign a new ZMQ socket address.
+def get_engine_client_zmq_addr(
+    local_only: bool,
+    host: str,
+    port: int = 0,
+) -> str:
+    """Return an IPC path (``local_only=True``) or ``tcp://host:port``.
 
-    If local_only is True, participants are colocated and so a unique IPC
-    address will be returned.
-
-    Otherwise, the provided host and port will be used to construct a TCP
-    address (port == 0 means assign an available port)."""
-
-    return (
-        get_open_zmq_ipc_path()
-        if local_only
-        else (get_tcp_uri(host, port or get_open_port()))
-    )
+    ``port=0`` lets the kernel assign the port at ``bind()`` time; the
+    caller must recover it via ``getsockopt(zmq.LAST_ENDPOINT)``."""
+    if local_only:
+        return get_open_zmq_ipc_path()
+    return get_tcp_uri(host, port)
 
 
 class APIServerProcessManager:
@@ -181,6 +179,12 @@ class APIServerProcessManager:
     ):
         """Initialize and start API server worker processes.
 
+        ``input_addresses``/``output_addresses`` may contain
+        ``tcp://host:0`` placeholders; each child must report the actual
+        bound endpoint over its ``actual_address_pipe`` in ``client_config``
+        and the parent collects them via
+        :py:meth:`gather_actual_addresses`.
+
         Args:
             target_server_fn: Override function to call for each API server process
             listen_address: Address to listen for client connections
@@ -196,14 +200,14 @@ class APIServerProcessManager:
         self.sock = sock
         self.args = args
 
-        # Start API servers
         spawn_context = multiprocessing.get_context("spawn")
         self.processes: list[BaseProcess] = []
+        self._address_pipes: list[connection.Connection] = []
 
         for i, in_addr, out_addr in zip(
             range(num_servers), input_addresses, output_addresses
         ):
-            client_config = {
+            client_config: dict[str, Any] = {
                 "input_address": in_addr,
                 "output_address": out_addr,
                 "client_count": num_servers,
@@ -214,6 +218,10 @@ class APIServerProcessManager:
             if tensor_queue is not None:
                 client_config["tensor_queue"] = tensor_queue
 
+            parent_recv, child_send = spawn_context.Pipe(duplex=False)
+            self._address_pipes.append(parent_recv)
+            client_config["actual_address_pipe"] = child_send
+
             proc = spawn_context.Process(
                 target=target_server_fn or run_api_server_worker_proc,
                 name=f"ApiServer_{i}",
@@ -222,14 +230,89 @@ class APIServerProcessManager:
             self.processes.append(proc)
             proc.start()
 
+            # Drop parent's write end so reader sees EOF on child death.
+            child_send.close()
+
         logger.info("Started %d API server processes", len(self.processes))
 
         # Shutdown only the API server processes on garbage collection
         # The extra processes are managed by their owners
         self._finalizer = weakref.finalize(self, shutdown, self.processes)
 
+    def gather_actual_addresses(
+        self,
+        timeout: float = 60.0,
+    ) -> tuple[list[str], list[str]]:
+        """Return (inputs, outputs) reported by each child, indexed by
+        ``client_index``. Raises ``RuntimeError`` on timeout or premature
+        child exit."""
+        n = len(self._address_pipes)
+        inputs: list[str | None] = [None] * n
+        outputs: list[str | None] = [None] * n
+        pending: dict[connection.Connection, int] = {
+            pipe: i for i, pipe in enumerate(self._address_pipes)
+        }
+        sentinel_to_idx: dict[Any, int] = {
+            proc.sentinel: i for i, proc in enumerate(self.processes)
+        }
+
+        deadline = time.monotonic() + timeout
+        try:
+            while pending:
+                remaining = deadline - time.monotonic()
+                if remaining <= 0:
+                    missing = [self.processes[i].name for i in pending.values()]
+                    raise RuntimeError(
+                        f"Timed out after {timeout:.1f}s waiting for "
+                        f"API server(s) to report bound ZMQ addresses: "
+                        f"{missing}"
+                    )
+                waitables: list[Any] = list(pending.keys()) + list(
+                    sentinel_to_idx.keys()
+                )
+                ready = connection.wait(waitables, timeout=remaining)
+                # Drain pipes before checking sentinels: a child that sent
+                # its message and then exited can surface both events in
+                # the same poll, and we must record the success first.
+                for item in ready:
+                    if isinstance(item, connection.Connection) and item in pending:
+                        idx = pending.pop(item)
+                        try:
+                            msg: dict[str, str] = item.recv()
+                        except EOFError as e:
+                            raise RuntimeError(
+                                f"API server {self.processes[idx].name} "
+                                f"closed its address pipe without "
+                                f"reporting its bound ZMQ addresses"
+                            ) from e
+                        inputs[idx] = msg["input_address"]
+                        outputs[idx] = msg["output_address"]
+                        item.close()
+                for item in ready:
+                    if item in sentinel_to_idx:
+                        idx = sentinel_to_idx.pop(item)
+                        pipe = self._address_pipes[idx]
+                        if pipe in pending:
+                            proc = self.processes[idx]
+                            raise RuntimeError(
+                                f"API server process {proc.name} exited "
+                                f"(code={proc.exitcode}) before reporting "
+                                f"its bound ZMQ addresses"
+                            )
+        finally:
+            for pipe in pending:
+                with contextlib.suppress(Exception):
+                    pipe.close()
+
+        return inputs, outputs  # type: ignore[return-value]
+
     def shutdown(self, timeout: float | None = None) -> None:
         """Shutdown API server processes with configurable timeout"""
+        for pipe in self._address_pipes:
+            with contextlib.suppress(Exception):
+                pipe.close()
+        self._address_pipes = []
+
         if self._finalizer.detach() is not None:
             shutdown(self.processes, timeout=timeout)
 

vllm/v1/worker/gpu/kv_connector.py

@@ -34,7 +34,7 @@ class KVConnector:
         pass
 
     def post_forward(
-        self, scheduler_output: "SchedulerOutput", wait_for_save: bool = True
+        self, finished_req_ids: set[str], wait_for_save: bool = True
     ) -> KVConnectorOutput | None:
         return None
 
@@ -76,10 +76,7 @@ class ActiveKVConnector(KVConnector):
                 self.kv_connector.start_load_kv(get_forward_context())
 
     def post_forward(
-        self,
-        scheduler_output: "SchedulerOutput",
-        wait_for_save: bool = True,
-        clear_metadata: bool = True,
+        self, finished_req_ids: set[str], wait_for_save: bool = True
     ) -> KVConnectorOutput | None:
         if self._disabled:
             return None
@@ -88,7 +85,7 @@ class ActiveKVConnector(KVConnector):
         if wait_for_save:
             self.kv_connector.wait_for_save()
         output.finished_sending, output.finished_recving = (
-            self.kv_connector.get_finished(scheduler_output.finished_req_ids)
+            self.kv_connector.get_finished(finished_req_ids)
         )
         output.invalid_block_ids = self.kv_connector.get_block_ids_with_load_errors()
         output.kv_connector_stats = self.kv_connector.get_kv_connector_stats()
@@ -96,9 +93,7 @@ class ActiveKVConnector(KVConnector):
         output.kv_connector_worker_meta = (
             self.kv_connector.build_connector_worker_meta()
         )
-
-        if clear_metadata:
-            self.kv_connector.clear_connector_metadata()
+        self.kv_connector.clear_connector_metadata()
         return output
 
     def no_forward(self, scheduler_output: "SchedulerOutput") -> ModelRunnerOutput:
@@ -106,7 +101,8 @@ class ActiveKVConnector(KVConnector):
             return EMPTY_MODEL_RUNNER_OUTPUT
 
         self.pre_forward(scheduler_output)
-        kv_connector_output = self.post_forward(scheduler_output, wait_for_save=False)
+        finished_req_ids = scheduler_output.finished_req_ids
+        kv_connector_output = self.post_forward(finished_req_ids, wait_for_save=False)
         if kv_connector_output is None or kv_connector_output.is_empty():
             return EMPTY_MODEL_RUNNER_OUTPUT
         output = copy.copy(EMPTY_MODEL_RUNNER_OUTPUT)

vllm/v1/worker/gpu/model_runner.py

@@ -50,7 +50,7 @@ from vllm.utils.mem_utils import DeviceMemoryProfiler, format_gib
 from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE
 from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput
 from vllm.v1.kv_cache_interface import KVCacheConfig, MambaSpec
-from vllm.v1.outputs import DraftTokenIds, KVConnectorOutput, ModelRunnerOutput
+from vllm.v1.outputs import DraftTokenIds, ModelRunnerOutput
 from vllm.v1.worker.cp_utils import check_attention_cp_compatibility
 from vllm.v1.worker.gpu.async_utils import AsyncOutput, AsyncPoolingOutput
 from vllm.v1.worker.gpu.attn_utils import (
@@ -1207,19 +1207,20 @@ class GPUModelRunner(LoRAModelRunnerMixin):
             aux_hidden_states = None
             output_intermediate_tensors = model_output
 
-        kv_connector_output = self.kv_connector.post_forward(scheduler_output)
+        finished_req_ids = scheduler_output.finished_req_ids
         self.execute_model_state = ExecuteModelState(
             input_batch=input_batch,
             attn_metadata=attn_metadata,
             slot_mappings_by_layer=slot_mappings_by_layer,
             hidden_states=hidden_states,
             aux_hidden_states=aux_hidden_states,
-            kv_connector_output=kv_connector_output,
+            finished_req_ids=finished_req_ids,
         )
 
         if not self.is_last_pp_rank:
             # Non-last PP rank: return IntermediateTensors for sending.
             assert output_intermediate_tensors is not None
+            kv_connector_output = self.kv_connector.post_forward(finished_req_ids)
             output_intermediate_tensors.kv_connector_output = kv_connector_output
             return output_intermediate_tensors
         return None
@@ -1238,7 +1239,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
         slot_mappings_by_layer = self.execute_model_state.slot_mappings_by_layer
         hidden_states = self.execute_model_state.hidden_states
         aux_hidden_states = self.execute_model_state.aux_hidden_states
-        kv_connector_output = self.execute_model_state.kv_connector_output
+        finished_req_ids = self.execute_model_state.finished_req_ids
         self.execute_model_state = None
 
         if not self.is_last_pp_rank:
@@ -1280,8 +1281,8 @@ class GPUModelRunner(LoRAModelRunnerMixin):
             req_id_to_index={req_id: i for i, req_id in enumerate(input_batch.req_ids)},
             sampled_token_ids=None,  # type: ignore
             prompt_logprobs_dict=prompt_logprobs_dict,  # type: ignore[arg-type]
-            kv_connector_output=kv_connector_output,
         )
+        # Start async output copy here so that it can overlap with speculator proposal.
         async_output = AsyncOutput(
             model_runner_output=model_runner_output,
             sampler_output=sampler_output,
@@ -1344,6 +1345,10 @@ class GPUModelRunner(LoRAModelRunnerMixin):
             self.req_states.draft_tokens[input_batch.idx_mapping] = draft_tokens
             self.draft_tokens_handler.set_draft_tokens(input_batch, draft_tokens)
 
+        # Post-step KV connector related operations.
+        kv_connector_output = self.kv_connector.post_forward(finished_req_ids)
+        model_runner_output.kv_connector_output = kv_connector_output
+
         if self.use_async_scheduling:
             return async_output
         return async_output.get_output()
@@ -1360,7 +1365,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
 
         input_batch = self.execute_model_state.input_batch
         hidden_states = self.execute_model_state.hidden_states
-        kv_connector_output = self.execute_model_state.kv_connector_output
+        finished_req_ids = self.execute_model_state.finished_req_ids
         self.execute_model_state = None
 
         if not self.is_last_pp_rank:
@@ -1372,6 +1377,9 @@ class GPUModelRunner(LoRAModelRunnerMixin):
             hidden_states, input_batch, self.req_states
         )
 
+        # Post-step KV connector related operations.
+        kv_connector_output = self.kv_connector.post_forward(finished_req_ids)
+
         # Build the model runner output.
         model_runner_output = ModelRunnerOutput(
             req_ids=input_batch.req_ids,
@@ -1455,4 +1463,4 @@ class ExecuteModelState(NamedTuple):
     slot_mappings_by_layer: dict[str, torch.Tensor] | None
     hidden_states: torch.Tensor | None
     aux_hidden_states: list[torch.Tensor] | None
-    kv_connector_output: KVConnectorOutput | None
+    finished_req_ids: set[str]