fix: resolve CUTLASS fmin compatibility for DeepSeek-V4 init

Signed-off-by: khluu <khluu000@gmail.com>

.buildkite/hardware_tests/cpu.yaml

@@ -54,20 +54,6 @@ steps:
       pytest -x -v -s tests/models/language/generation -m cpu_model
       pytest -x -v -s tests/models/language/pooling -m cpu_model"
 
-- label: CPU-ModelRunnerV2 Tests
-  depends_on: []
-  device: intel_cpu
-  no_plugin: true
-  soft_fail: true
-  source_file_dependencies:
-  - vllm/v1/worker/cpu/
-  - vllm/v1/worker/gpu/
-  commands:
-    - |
-      bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 30m "
-      uv pip install git+https://github.com/triton-lang/triton-cpu.git@270e696d
-      VLLM_USE_V2_MODEL_RUNNER=1 pytest -x -v -s tests/models/language/generation/test_granite.py -m cpu_model"
-
 - label: CPU-Quantization Model Tests
   depends_on: []
   device: intel_cpu

.buildkite/scripts/install-kv-connectors.sh

@@ -0,0 +1,39 @@
+#!/bin/bash
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+set -euo pipefail
+
+REQUIREMENTS_FILE="${KV_CONNECTORS_REQUIREMENTS:-/vllm-workspace/requirements/kv_connectors.txt}"
+
+uv pip install --system -r "${REQUIREMENTS_FILE}"
+
+NIXL_METADATA=$(python3 - <<'PY'
+import importlib.metadata as metadata
+
+import torch
+
+cuda_version = torch.version.cuda
+if cuda_version is None:
+    raise SystemExit("torch.version.cuda is not set")
+
+print(cuda_version.split(".", 1)[0], metadata.version("nixl"))
+PY
+)
+read -r CUDA_MAJOR NIXL_VERSION <<<"${NIXL_METADATA}"
+
+# nixl>=1.1.0 can install multiple CUDA wheel variants. Keep only the variant
+# matching this CI image so nixl_ep_cpp links against the available libcudart.
+uv pip uninstall --system nixl-cu12 nixl-cu13 2>/dev/null || true
+uv pip install --system --no-deps "nixl-cu${CUDA_MAJOR}==${NIXL_VERSION}"
+
+python3 - <<'PY'
+import importlib.metadata as metadata
+
+for package_name in ("nixl", "nixl-cu12", "nixl-cu13"):
+    try:
+        version = metadata.version(package_name)
+    except metadata.PackageNotFoundError:
+        version = "not installed"
+    print(f"{package_name}: {version}")
+PY

.buildkite/test_areas/disaggregated.yaml

@@ -11,7 +11,7 @@ steps:
     - vllm/distributed/kv_transfer/kv_connector/v1/nixl/
     - tests/v1/kv_connector/nixl_integration/
   commands:
-    - uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
+    - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
     - bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
 - label: Distributed FlashInfer NixlConnector PD accuracy (4 GPUs)
   key: distributed-flashinfer-nixlconnector-pd-accuracy-4-gpus
@@ -22,7 +22,7 @@ steps:
     - vllm/distributed/kv_transfer/kv_connector/v1/nixl/
     - tests/v1/kv_connector/nixl_integration/
   commands:
-    - uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
+    - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
     - FLASHINFER=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
 
 - label: DP EP Distributed NixlConnector PD accuracy tests (4 GPUs)
@@ -34,7 +34,7 @@ steps:
     - vllm/distributed/kv_transfer/kv_connector/v1/nixl/
     - tests/v1/kv_connector/nixl_integration/
   commands:
-    - uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
+    - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
     - DP_EP=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
 
 - label: CrossLayer KV layout Distributed NixlConnector PD accuracy tests (4 GPUs)
@@ -46,7 +46,7 @@ steps:
     - vllm/distributed/kv_transfer/kv_connector/v1/nixl/
     - tests/v1/kv_connector/nixl_integration/
   commands:
-    - uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
+    - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
     - CROSS_LAYERS_BLOCKS=True bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
 
 - label: Hybrid SSM NixlConnector PD accuracy tests (4 GPUs)
@@ -58,7 +58,7 @@ steps:
     - vllm/distributed/kv_transfer/kv_connector/v1/nixl/
     - tests/v1/kv_connector/nixl_integration/
   commands:
-    - uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
+    - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
     - HYBRID_SSM=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
 
 - label: MultiConnector (Nixl+Offloading) PD accuracy (2 GPUs)
@@ -73,7 +73,7 @@ steps:
     - vllm/distributed/kv_transfer/kv_connector/v1/offloading/
     - tests/v1/kv_connector/nixl_integration/
   commands:
-    - uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
+    - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
     - bash v1/kv_connector/nixl_integration/run_multi_connector_accuracy_test.sh
 
 - label: NixlConnector PD + Spec Decode acceptance (2 GPUs)
@@ -87,7 +87,7 @@ steps:
     - vllm/v1/worker/kv_connector_model_runner_mixin.py
     - tests/v1/kv_connector/nixl_integration/
   commands:
-    - uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
+    - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
     - bash v1/kv_connector/nixl_integration/config_sweep_spec_decode_test.sh
 
 - label: MultiConnector (Nixl+Offloading) PD edge cases (2 GPUs)
@@ -102,5 +102,5 @@ steps:
     - vllm/distributed/kv_transfer/kv_connector/v1/offloading/
     - tests/v1/kv_connector/nixl_integration/
   commands:
-    - uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
-    - bash v1/kv_connector/nixl_integration/run_multi_connector_edge_case_test.sh
+    - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
+    - bash v1/kv_connector/nixl_integration/run_multi_connector_edge_case_test.sh

.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

.buildkite/test_areas/misc.yaml

@@ -86,7 +86,7 @@ steps:
     - tests/v1/metrics
     - tests/entrypoints/openai/correctness/test_lmeval.py
   commands:
-    - uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
+    - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
     - export VLLM_WORKER_MULTIPROC_METHOD=spawn
     # split the test to avoid interference
     - pytest -v -s -m 'not cpu_test' v1/core

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 '.') && \
@@ -758,7 +760,7 @@ RUN --mount=type=cache,target=/opt/uv/cache \
 ARG FLASHINFER_VERSION=0.6.11.post2
 RUN --mount=type=cache,target=/opt/uv/cache \
     uv pip install --system flashinfer-jit-cache==${FLASHINFER_VERSION} \
-        --extra-index-url https://flashinfer.ai/whl/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.')
+        --index-url https://flashinfer.ai/whl/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.')
 
 # ============================================================
 # OPENAI API SERVER DEPENDENCIES

docker/Dockerfile.cpu

@@ -27,14 +27,11 @@ WORKDIR /workspace
 ARG PYTHON_VERSION=3.12
 ARG PIP_EXTRA_INDEX_URL="https://download.pytorch.org/whl/cpu"
 
-ARG max_jobs=32
-ENV MAX_JOBS=${max_jobs}
-
 # Install minimal dependencies and uv
 RUN --mount=type=cache,target=/var/cache/apt,sharing=locked \
     --mount=type=cache,target=/var/lib/apt,sharing=locked \
     apt-get update -y \
-    && apt-get install -y --no-install-recommends sudo ccache git curl wget ca-certificates zlib1g-dev \
+    && apt-get install -y --no-install-recommends sudo ccache git curl wget ca-certificates \
     gcc-12 g++-12 libtcmalloc-minimal4 libnuma-dev ffmpeg libsm6 libxext6 libgl1 jq lsof make xz-utils \
     && update-alternatives --install /usr/bin/gcc gcc /usr/bin/gcc-12 10 --slave /usr/bin/g++ g++ /usr/bin/g++-12 \
     && curl -LsSf https://astral.sh/uv/install.sh | sh
@@ -126,6 +123,9 @@ RUN --mount=type=cache,target=/root/.cargo/registry \
 ######################### BUILD IMAGE #########################
 FROM base AS vllm-build
 
+ARG max_jobs=32
+ENV MAX_JOBS=${max_jobs}
+
 ARG GIT_REPO_CHECK=0
 # Support for cross-compilation with x86 ISA including AVX2 and AVX512: docker build --build-arg VLLM_CPU_X86="true" ...
 ARG VLLM_CPU_X86=0
@@ -257,7 +257,8 @@ WORKDIR /vllm-workspace
 RUN --mount=type=cache,target=/root/.cache/uv \
     --mount=type=cache,target=/root/.cache/ccache \
     --mount=type=bind,from=vllm-build,src=/vllm-workspace/dist,target=dist \
-    uv pip install "$(realpath dist/*.whl)[audio,triton-cpu]"
+    uv pip install dist/*.whl && \
+    uv pip install "vllm[audio]"
 
 # Add labels to document build configuration
 LABEL org.opencontainers.image.title="vLLM CPU"

docs/features/reasoning_outputs.md

@@ -15,6 +15,7 @@ vLLM currently supports the following reasoning models:
 | ------------ | ----------- | ---------------- | ----------- |
 | [Cohere Command A Reasoning](https://huggingface.co/CohereLabs/command-a-reasoning-08-2025) | `cohere_command3` | `json`, `regex` | ✅ |
 | [DeepSeek R1 series](https://huggingface.co/collections/deepseek-ai/deepseek-r1-678e1e131c0169c0bc89728d) | `deepseek_r1` | `json`, `regex` | ❌ |
+| [Gemma 4 series](https://huggingface.co/google/gemma-4-26B-A4B-it) | `gemma4` | `json`, `regex` | ✅ |
 | [DeepSeek-V3.1](https://huggingface.co/collections/deepseek-ai/deepseek-v31-68a491bed32bd77e7fca048f) | `deepseek_v3` | `json`, `regex` | ❌ |
 | [ERNIE-4.5-VL series](https://huggingface.co/baidu/ERNIE-4.5-VL-28B-A3B-PT) | `ernie45` | `json`, `regex` | ❌ |
 | [ERNIE-4.5-21B-A3B-Thinking](https://huggingface.co/baidu/ERNIE-4.5-21B-A3B-Thinking) | `ernie45` | `json`, `regex` | ✅ |
@@ -29,6 +30,7 @@ vLLM currently supports the following reasoning models:
 !!! note
     IBM Granite 3.2 and DeepSeek-V3.1 reasoning is disabled by default; to enable it, you must also pass `thinking=True` in your `chat_template_kwargs`.
     The reasoning feature for the Qwen3 series is enabled by default. To disable it, you must pass `enable_thinking=False` in your `chat_template_kwargs`.
+    Gemma 4 reasoning is disabled by default; to enable it, pass `enable_thinking=True` in your `chat_template_kwargs` or set `reasoning_effort` (which enables it automatically).
     DeepSeek-V3.1 tool calling is supported in non-thinking mode.
     Holo2 reasoning is enabled by default. To disable it, you must also pass `thinking=False` in your `chat_template_kwargs`.
 
@@ -314,9 +316,44 @@ for output in outputs:
     print("text:", output.outputs[0].text)
 ```
 
+## Automatic `enable_thinking` Activation
+
+Some models (such as Gemma 4, DeepSeek-V4-Pro and IBM Granite 3.2) require `enable_thinking: true` in their chat template kwargs to activate thinking mode — without it, reasoning tokens are never generated regardless of other settings.
+
+When you set `reasoning_effort` in a Chat Completions request (or `reasoning.effort` in a Responses API request), vLLM automatically injects `enable_thinking` into the chat template kwargs:
+
+- `reasoning_effort` = `"low"`, `"medium"`, or `"high"` → `enable_thinking = true`
+- `reasoning_effort` = `"none"` → `enable_thinking = false`
+- `reasoning_effort` not set → `enable_thinking` is not injected (preserves existing behavior)
+
+This means you no longer need to manually pass `chat_template_kwargs: {"enable_thinking": true}` when using `reasoning_effort` — it is handled automatically.
+
+!!! note
+    If you explicitly set `enable_thinking` in `chat_template_kwargs`, your value takes priority over the automatic injection. This allows you to override the behavior if needed.
+
+    For models whose templates don't declare `enable_thinking` (e.g., DeepSeek R1), the injected kwarg is harmlessly filtered out by `resolve_chat_template_kwargs`.
+
+### Example
+
+```python
+from openai import OpenAI
+
+client = OpenAI(base_url="http://localhost:8000/v1", api_key="dummy")
+
+# reasoning_effort automatically enables thinking for models that need it
+response = client.chat.completions.create(
+    model="google/gemma-4-26B-A4B-it",
+    messages=[{"role": "user", "content": "What is 15 * 37?"}],
+    reasoning_effort="high",  # Automatically sets enable_thinking=true
+)
+
+print(response.choices[0].message.reasoning)
+print(response.choices[0].message.content)
+```
+
 ## Limitations
 
-- The reasoning content is only available for online serving's chat completion endpoint (`/v1/chat/completions`).
+- The reasoning content is only available for online serving's chat completion endpoint (`/v1/chat/completions`), Anthropic Messages API (`/v1/messages`) and the Responses API (`/v1/responses`).
 
 ## How to support a new reasoning model
 

docs/models/supported_models.md

@@ -438,6 +438,7 @@ th {
 | `LongcatFlashForCausalLM` | LongCat-Flash | `meituan-longcat/LongCat-Flash-Chat`, `meituan-longcat/LongCat-Flash-Chat-FP8` | ✅︎ | ✅︎ |
 | `MambaForCausalLM` | Mamba | `state-spaces/mamba-130m-hf`, `state-spaces/mamba-790m-hf`, `state-spaces/mamba-2.8b-hf`, etc. | | ✅︎ |
 | `Mamba2ForCausalLM` | Mamba2 | `mistralai/Mamba-Codestral-7B-v0.1`, etc. | | ✅︎ |
+| `MellumForCausalLM` | Mellum 2 | `JetBrains/Mellum2-12B-A2.5B-Base`, etc. | | ✅︎ |
 | `MiMoForCausalLM` | MiMo | `XiaomiMiMo/MiMo-7B-RL`, etc. | ✅︎ | ✅︎ |
 | `MiMoV2FlashForCausalLM` | MiMoV2Flash | `XiaomiMiMo/MiMo-V2-Flash`, etc. | | ✅︎ |
 | `MiMoV2ForCausalLM` | MiMoV2Pro | `XiaomiMiMo/MiMo-V2.5-Pro`, etc. | | ✅︎ |

requirements/build/rocm.txt

@@ -16,3 +16,4 @@ wheel
 jinja2>=3.1.6
 amdsmi==7.0.2
 timm>=1.0.17
+tilelang==0.1.10

requirements/cuda.txt

@@ -21,8 +21,11 @@ nvidia-cudnn-frontend>=1.13.0,<1.19.0
 fastsafetensors >= 0.2.2
 
 # QuACK and Cutlass DSL for FA4 (cute-DSL implementation)
-nvidia-cutlass-dsl[cu13]==4.5.0
+nvidia-cutlass-dsl[cu13]==4.5.2
 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

requirements/rocm.txt

@@ -22,3 +22,4 @@ timm>=1.0.17
 # amd-quark: required for Quark quantization on ROCm 
 # To be consistent with test_quark.py
 amd-quark>=0.8.99
+tilelang==0.1.10

requirements/test/rocm.in

@@ -43,6 +43,7 @@ schemathesis>=3.39.15 # Required for openai schema test
 # quantization
 bitsandbytes==0.49.2
 buildkite-test-collector==0.1.9
+tilelang==0.1.10
 
 genai_perf>=0.0.8
 tritonclient>=2.51.0

requirements/test/rocm.txt

@@ -43,7 +43,9 @@ anyio==4.13.0
     #   starlette
     #   watchfiles
 apache-tvm-ffi==0.1.10
-    # via xgrammar
+    # via
+    #   tilelang
+    #   xgrammar
 arctic-inference==0.1.1
     # via -r requirements/test/rocm.in
 argcomplete==3.6.3
@@ -129,7 +131,9 @@ click==8.3.1
     #   typer
     #   uvicorn
 cloudpickle==3.1.2
-    # via -r requirements/test/../common.txt
+    # via
+    #   -r requirements/test/../common.txt
+    #   tilelang
 colorama==0.4.6
     # via
     #   perceptron
@@ -511,6 +515,8 @@ mistral-common==1.11.2
     #   -c requirements/common.txt
     #   -r requirements/test/../common.txt
     #   -r requirements/test/rocm.in
+ml-dtypes==0.5.4
+    # via tilelang
 model-hosting-container-standards==0.1.14
     # via
     #   -c requirements/common.txt
@@ -587,6 +593,7 @@ numpy==2.2.6
     #   lm-eval
     #   matplotlib
     #   mistral-common
+    #   ml-dtypes
     #   mteb
     #   numba
     #   opencv-python-headless
@@ -610,6 +617,7 @@ numpy==2.2.6
     #   statsmodels
     #   tensorizer
     #   tifffile
+    #   tilelang
     #   torchvision
     #   transformers
     #   tritonclient
@@ -811,6 +819,7 @@ psutil==7.2.2
     #   accelerate
     #   peft
     #   tensorizer
+    #   tilelang
 py==1.11.0
     # via pytest-forked
 py-cpuinfo==9.0.0
@@ -1192,6 +1201,10 @@ tiktoken==0.12.0
     #   gpt-oss
     #   lm-eval
     #   mistral-common
+tilelang==0.1.10
+    # via
+    #   -c requirements/rocm.txt
+    #   -r requirements/test/rocm.in
 timm==1.0.17
     # via
     #   -c requirements/rocm.txt
@@ -1208,6 +1221,8 @@ tomli==2.4.0
     # via schemathesis
 tomli-w==1.2.0
     # via schemathesis
+torch-c-dlpack-ext==0.1.5
+    # via tilelang
 tqdm==4.67.3
     # via
     #   -r requirements/test/../common.txt
@@ -1225,6 +1240,7 @@ tqdm==4.67.3
     #   pqdm
     #   segmentation-models-pytorch
     #   sentence-transformers
+    #   tilelang
     #   transformers
 transformers==5.5.3
     # via
@@ -1293,6 +1309,7 @@ typing-extensions==4.15.0
     #   sentence-transformers
     #   sqlalchemy
     #   starlette
+    #   tilelang
     #   torch
     #   typeguard
     #   typing-inspection
@@ -1359,6 +1376,8 @@ yarl==1.23.0
     # via
     #   aiohttp
     #   schemathesis
+z3-solver==4.15.4.0
+    # via tilelang
 zipp==3.23.0
     # via importlib-metadata
 

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():
@@ -1163,9 +1165,7 @@ setup(
     install_requires=get_requirements(),
     extras_require={
         # AMD Zen CPU optimizations via zentorch
-        "zen": [
-            "zentorch-weekly==5.2.1.dev20260408"
-        ],  # Zentorch has weekly releases. This pulls the known-good version.
+        "zen": ["zentorch==2.11.0.0"],
         "bench": ["pandas", "matplotlib", "seaborn", "datasets", "scipy", "plotly"],
         "tensorizer": ["tensorizer==2.10.1"],
         "fastsafetensors": ["fastsafetensors >= 0.2.2"],
@@ -1193,11 +1193,6 @@ setup(
             "opentelemetry-exporter-otlp>=1.26.0",
             "opentelemetry-semantic-conventions-ai>=0.4.1",
         ],
-        "triton-cpu": [
-            "triton @ "
-            "git+https://github.com/triton-lang/triton-cpu.git@270e696d ; "
-            "platform_machine == 'x86_64'",
-        ],  # Remove after stable release
     },
     cmdclass=cmdclass,
     package_data=package_data,

tests/entrypoints/openai/test_reasoning_enable_thinking.py

@@ -0,0 +1,107 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Unit tests for reasoning_effort -> enable_thinking mapping.
+
+Models like Gemma4 require enable_thinking=True in chat_template_kwargs to
+activate thinking mode. This mapping ensures that when a user requests
+reasoning (via reasoning_effort or reasoning.effort), the template kwarg
+is injected automatically.
+"""
+
+import pytest
+from openai.types.shared import Reasoning
+
+from vllm.entrypoints.openai.chat_completion.protocol import (
+    ChatCompletionRequest,
+)
+from vllm.entrypoints.openai.responses.protocol import ResponsesRequest
+
+
+def _build_chat_request(**kwargs) -> ChatCompletionRequest:
+    defaults = dict(
+        model="test-model",
+        messages=[{"role": "user", "content": "Hello"}],
+    )
+    defaults.update(kwargs)
+    return ChatCompletionRequest(**defaults)
+
+
+def _build_responses_request(**kwargs) -> ResponsesRequest:
+    defaults = dict(
+        model="test-model",
+        input=[{"role": "user", "content": "Hello"}],
+    )
+    defaults.update(kwargs)
+    return ResponsesRequest(**defaults)
+
+
+class TestChatCompletionReasoningEffort:
+    """Chat Completions: reasoning_effort -> enable_thinking."""
+
+    @pytest.mark.parametrize("effort", ["low", "medium", "high"])
+    def test_non_none_effort_injects_enable_thinking_true(self, effort):
+        request = _build_chat_request(reasoning_effort=effort)
+        params = request.build_chat_params(None, "auto")
+        assert params.chat_template_kwargs["enable_thinking"] is True
+
+    def test_none_effort_injects_enable_thinking_false(self):
+        request = _build_chat_request(reasoning_effort="none")
+        params = request.build_chat_params(None, "auto")
+        assert params.chat_template_kwargs["enable_thinking"] is False
+
+    def test_no_effort_does_not_inject(self):
+        request = _build_chat_request()
+        params = request.build_chat_params(None, "auto")
+        assert "enable_thinking" not in params.chat_template_kwargs
+
+    def test_explicit_user_kwarg_not_overridden(self):
+        request = _build_chat_request(
+            reasoning_effort="high",
+            chat_template_kwargs={"enable_thinking": False},
+        )
+        params = request.build_chat_params(None, "auto")
+        assert params.chat_template_kwargs["enable_thinking"] is False
+
+    def test_reasoning_effort_still_in_kwargs(self):
+        request = _build_chat_request(reasoning_effort="high")
+        params = request.build_chat_params(None, "auto")
+        assert params.chat_template_kwargs["reasoning_effort"] == "high"
+
+
+class TestResponsesReasoningEffort:
+    """Responses API: reasoning.effort -> enable_thinking."""
+
+    @pytest.mark.parametrize("effort", ["low", "medium", "high"])
+    def test_non_none_effort_injects_enable_thinking_true(self, effort):
+        request = _build_responses_request(
+            reasoning=Reasoning(effort=effort),
+        )
+        params = request.build_chat_params(None, "auto")
+        assert params.chat_template_kwargs["enable_thinking"] is True
+
+    def test_none_effort_injects_enable_thinking_false(self):
+        request = _build_responses_request(
+            reasoning=Reasoning(effort="none"),
+        )
+        params = request.build_chat_params(None, "auto")
+        assert params.chat_template_kwargs["enable_thinking"] is False
+
+    def test_no_reasoning_does_not_inject(self):
+        request = _build_responses_request()
+        params = request.build_chat_params(None, "auto")
+        assert "enable_thinking" not in params.chat_template_kwargs
+
+    def test_explicit_user_kwarg_not_overridden(self):
+        request = _build_responses_request(
+            reasoning=Reasoning(effort="high"),
+            chat_template_kwargs={"enable_thinking": False},
+        )
+        params = request.build_chat_params(None, "auto")
+        assert params.chat_template_kwargs["enable_thinking"] is False
+
+    def test_reasoning_effort_still_in_kwargs(self):
+        request = _build_responses_request(
+            reasoning=Reasoning(effort="high"),
+        )
+        params = request.build_chat_params(None, "auto")
+        assert params.chat_template_kwargs["reasoning_effort"] == "high"

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/moe/test_moe.py

@@ -1656,7 +1656,7 @@ def test_unquantized_bf16_flashinfer_trtllm_backend(
         layer.routing_method_type = RoutingMethodType.Renormalize
         layer.expert_map = None
         layer.apply_router_weight_on_input = False
-        layer.routed_scaling_factor = None
+        layer.routed_scaling_factor = 2.446
         layer.shared_experts = None
         layer._expert_routing_tables = lambda: None
 
@@ -1678,7 +1678,10 @@ def test_unquantized_bf16_flashinfer_trtllm_backend(
         # Compute torch baseline
         w1_original = w1.clone()
         w2_original = w2.clone()
-        baseline_output = torch_moe(a, w1_original, w2_original, router_logits, topk)
+        baseline_output = (
+            torch_moe(a, w1_original, w2_original, router_logits, topk)
+            * layer.routed_scaling_factor
+        )
 
     close = torch.isclose(trtllm_output, baseline_output, atol=1e-1, rtol=0.85)
     assert close.float().mean() > 0.925

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/kernels/test_mhc_kernels.py

@@ -4,7 +4,12 @@ import pytest
 import torch
 
 import vllm.model_executor.kernels.mhc  # noqa: F401
+from vllm.model_executor.kernels.mhc.tilelang import (
+    _tilelang_hc_prenorm_gemm,
+    _torch_hc_prenorm_gemm,
+)
 from vllm.platforms import current_platform
+from vllm.utils.import_utils import has_tilelang
 from vllm.utils.torch_utils import set_random_seed
 
 DEVICE = current_platform.device_type
@@ -92,8 +97,128 @@ def hc_head_ref(
 
 
 @pytest.mark.skipif(
-    not current_platform.is_cuda(),
-    reason="CUDA required",
+    not (current_platform.is_cuda_alike() and has_tilelang()),
+    reason="CUDA or ROCm and tilelang required",
+)
+@pytest.mark.parametrize("num_tokens", [1, 4, 8, 128])
+@pytest.mark.parametrize("hidden_size", [4096, 7168])
+@pytest.mark.parametrize("hc_mult", [4])
+def test_mhc_pre_tilelang(num_tokens, hidden_size, hc_mult):
+    torch.set_default_device(DEVICE)
+    set_random_seed(0)
+
+    residual = torch.randn((num_tokens, hc_mult, hidden_size), dtype=torch.bfloat16)
+    hc_mult2 = hc_mult * hc_mult
+    hc_mult3 = 2 * hc_mult + hc_mult2
+    fn = (
+        torch.randn((hc_mult3, hc_mult, hidden_size), dtype=torch.float)
+        * 1e-4
+        * (1 + torch.arange(hc_mult).mul(0.01).view(1, -1, 1))
+    ).flatten(1, 2)
+    hc_scale = torch.randn((3,), dtype=torch.float) * 0.1
+    hc_base = torch.randn((hc_mult3,), dtype=torch.float) * 0.1
+
+    hc_sinkhorn_eps = hc_pre_eps = rms_eps = 1e-6
+    sinkhorn_repeat = 20
+    hc_post_alpha = 1.0
+
+    ref = mhc_pre_ref(
+        residual,
+        fn,
+        hc_scale,
+        hc_base,
+        rms_eps,
+        hc_pre_eps,
+        hc_sinkhorn_eps,
+        hc_post_alpha,
+        sinkhorn_repeat,
+    )
+    out = torch.ops.vllm.mhc_pre_tilelang(
+        residual,
+        fn,
+        hc_scale,
+        hc_base,
+        rms_eps,
+        hc_pre_eps,
+        hc_sinkhorn_eps,
+        hc_post_alpha,
+        sinkhorn_repeat,
+    )
+
+    for actual, expected in zip(out, ref, strict=True):
+        torch.testing.assert_close(actual, expected, atol=5e-2, rtol=1e-2)
+
+
+@pytest.mark.skipif(
+    not (current_platform.is_cuda_alike() and has_tilelang()),
+    reason="CUDA or ROCm and tilelang required",
+)
+@pytest.mark.parametrize(
+    ("num_tokens", "hidden_size"),
+    [
+        (1, 1280),
+        (512, 1280),
+        (2048, 1280),
+        (1, 4096),
+        (64, 4096),
+        (512, 4096),
+        (2048, 4096),
+        (1, 7168),
+        (64, 7168),
+        (512, 7168),
+        (2048, 7168),
+    ],
+)
+def test_hc_prenorm_gemm_tilelang(num_tokens, hidden_size):
+    torch.set_default_device(DEVICE)
+    set_random_seed(0)
+
+    hc_mult = 4
+    hc_mult3 = 2 * hc_mult + hc_mult * hc_mult
+    x = torch.randn((num_tokens, hc_mult * hidden_size), dtype=torch.bfloat16)
+    fn = torch.randn((hc_mult3, hc_mult * hidden_size), dtype=torch.float32) * 1e-4
+    out_ref = torch.empty((1, num_tokens, hc_mult3), dtype=torch.float32)
+    sqrsum_ref = torch.empty((1, num_tokens), dtype=torch.float32)
+    out = torch.empty_like(out_ref)
+    sqrsum = torch.empty_like(sqrsum_ref)
+
+    _torch_hc_prenorm_gemm(x, fn, out_ref, sqrsum_ref)
+    _tilelang_hc_prenorm_gemm(x, fn, out, sqrsum, hidden_size, hc_mult)
+
+    torch.testing.assert_close(out, out_ref, atol=1e-5, rtol=1e-4)
+    torch.testing.assert_close(sqrsum, sqrsum_ref, atol=8.0, rtol=5e-4)
+
+
+@pytest.mark.skipif(
+    not (current_platform.is_cuda_alike() and has_tilelang()),
+    reason="CUDA or ROCm and tilelang required",
+)
+@pytest.mark.parametrize("num_tokens", [1, 4, 8, 128])
+@pytest.mark.parametrize("hidden_size", [4096, 7168])
+@pytest.mark.parametrize("hc_mult", [4])
+def test_mhc_post_tilelang(num_tokens, hidden_size, hc_mult):
+    torch.set_default_device(DEVICE)
+    set_random_seed(0)
+
+    x = torch.randn((num_tokens, hidden_size), dtype=torch.bfloat16)
+    residual = torch.randn((num_tokens, hc_mult, hidden_size), dtype=torch.bfloat16)
+    post_layer_mix = torch.randn((num_tokens, hc_mult, 1), dtype=torch.float32)
+    comb_res_mix = torch.randn((num_tokens, hc_mult, hc_mult), dtype=torch.float32)
+
+    ref = mhc_post_ref(x, residual, post_layer_mix, comb_res_mix)
+    out = torch.ops.vllm.mhc_post_tilelang(
+        x,
+        residual,
+        post_layer_mix,
+        comb_res_mix,
+    )
+
+    torch.testing.assert_close(out, ref, atol=5e-2, rtol=1e-2)
+
+
+@pytest.mark.skipif(
+    not (current_platform.is_cuda_alike() and has_tilelang()),
+    reason="CUDA or ROCm and tilelang required",
 )
 @pytest.mark.parametrize("num_tokens", [1, 4, 8, 128])
 @pytest.mark.parametrize("hidden_size", [4096, 7168])
@@ -196,3 +321,42 @@ def test_hc_head_triton(num_tokens, hidden_size, hc_mult):
 
     out_ref = hc_head_ref(residual, fn, hc_scale, hc_base, rms_eps, hc_eps)
     torch.testing.assert_close(out, out_ref, atol=5e-2, rtol=1e-2)
+
+
+@pytest.mark.skipif(
+    not (current_platform.is_cuda_alike() and has_tilelang()),
+    reason="CUDA or ROCm and tilelang required",
+)
+@pytest.mark.parametrize("num_tokens", [1, 4, 8, 128])
+@pytest.mark.parametrize("hidden_size", [4096, 7168])
+@pytest.mark.parametrize("hc_mult", [4])
+def test_hc_head_tilelang(num_tokens, hidden_size, hc_mult):
+    torch.set_default_device(DEVICE)
+    set_random_seed(0)
+
+    residual = torch.randn((num_tokens, hc_mult, hidden_size), dtype=torch.bfloat16)
+    fn = torch.randn((hc_mult, hc_mult * hidden_size), dtype=torch.float32) * 1e-4
+    hc_scale = torch.randn((1,), dtype=torch.float32) * 0.1
+    hc_base = torch.randn((hc_mult,), dtype=torch.float32) * 0.1
+    rms_eps = hc_eps = 1e-6
+
+    out = torch.empty((num_tokens, hidden_size), dtype=torch.bfloat16)
+    out.fill_(float("nan"))
+
+    result = torch.ops.vllm.hc_head_fused_kernel_tilelang(
+        residual,
+        fn,
+        hc_scale,
+        hc_base,
+        out,
+        hidden_size,
+        rms_eps,
+        hc_eps,
+        hc_mult,
+    )
+
+    assert result is None
+    assert not torch.isnan(out).any()
+
+    out_ref = hc_head_ref(residual, fn, hc_scale, hc_base, rms_eps, hc_eps)
+    torch.testing.assert_close(out, out_ref, atol=5e-2, rtol=1e-2)

tests/model_executor/model_loader/runai_streamer_loader/test_weight_utils.py

@@ -6,6 +6,7 @@ import tempfile
 
 import huggingface_hub.constants
 import torch
+from safetensors.torch import save_file
 
 from vllm.model_executor.model_loader.weight_utils import (
     download_weights_from_hf,
@@ -14,6 +15,27 @@ from vllm.model_executor.model_loader.weight_utils import (
 )
 
 
+def test_runai_safetensors_weights_iterator_clones_reused_buffers(
+    tmp_path, monkeypatch
+):
+    monkeypatch.setenv("RUNAI_STREAMER_MEMORY_LIMIT", "0")
+    weights_file = tmp_path / "model.safetensors"
+    expected_tensors = {
+        "first": torch.tensor([1.0, 2.0]),
+        "second": torch.tensor([3.0, 4.0]),
+    }
+    save_file(expected_tensors, weights_file)
+
+    actual_tensors = dict(
+        runai_safetensors_weights_iterator([str(weights_file)], False)
+    )
+
+    assert actual_tensors.keys() == expected_tensors.keys()
+    assert actual_tensors["first"].data_ptr() != actual_tensors["second"].data_ptr()
+    for name, expected_tensor in expected_tensors.items():
+        assert torch.equal(actual_tensors[name], expected_tensor)
+
+
 def test_runai_model_loader():
     with tempfile.TemporaryDirectory() as tmpdir:
         huggingface_hub.constants.HF_HUB_OFFLINE = False

tests/models/registry.py

@@ -335,7 +335,7 @@ _TEXT_GENERATION_EXAMPLE_MODELS = {
     "HYV3ForCausalLM": _HfExamplesInfo("tencent/Hy3-preview", trust_remote_code=True),
     "HyperCLOVAXForCausalLM": _HfExamplesInfo(
         "naver-hyperclovax/HyperCLOVAX-SEED-Think-14B",
-        trust_remote_code=True,
+        min_transformers_version="5.9.0",
     ),
     "InternLMForCausalLM": _HfExamplesInfo(
         "internlm/internlm-chat-7b", trust_remote_code=True
@@ -523,6 +523,7 @@ _TEXT_GENERATION_EXAMPLE_MODELS = {
     "Qwen2MoeForCausalLM": _HfExamplesInfo("Qwen/Qwen1.5-MoE-A2.7B-Chat"),
     "Qwen3ForCausalLM": _HfExamplesInfo("Qwen/Qwen3-8B"),
     "Qwen3MoeForCausalLM": _HfExamplesInfo("Qwen/Qwen3-30B-A3B"),
+    "MellumForCausalLM": _HfExamplesInfo("JetBrains/Mellum2-12B-A2.5B-Base"),
     "Qwen3NextForCausalLM": _HfExamplesInfo(
         "Qwen/Qwen3-Next-80B-A3B-Instruct",
         extras={"tiny-random": "tiny-random/qwen3-next-moe"},

tests/v1/engine/test_core_engine_actor_manager.py

@@ -2,6 +2,8 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 import os
+import socket
+import time
 import uuid
 from pathlib import Path
 from types import SimpleNamespace
@@ -9,9 +11,17 @@ from typing import Any
 
 import pytest
 import ray
+import zmq
 
+from vllm.utils.network_utils import make_zmq_socket, split_zmq_path
 from vllm.v1.engine.core import EngineCoreActorMixin
-from vllm.v1.engine.utils import CoreEngineActorManager, EngineZmqAddresses
+from vllm.v1.engine.utils import (
+    CoreEngineActorManager,
+    EngineZmqAddresses,
+    get_engine_zmq_addresses,
+    launch_core_engines,
+)
+from vllm.v1.utils import APIServerProcessManager
 
 
 class _StubEngineCoreActor(EngineCoreActorMixin):
@@ -42,6 +52,48 @@ class _StubEngineCoreActor(EngineCoreActorMixin):
     def get_nixl_side_channel_host(self) -> str | None:
         return os.environ.get("VLLM_NIXL_SIDE_CHANNEL_HOST")
 
+    def get_addresses(self) -> tuple[list[str], list[str]]:
+        """Return the addresses snapshot the actor was constructed with.
+
+        Used by the Ray-DP regression test to assert that no ``tcp://host:0``
+        placeholders were pickled into the actor at ``.remote()`` time.
+        """
+        return list(self.addresses.inputs), list(self.addresses.outputs)
+
+
+# Module-level stub worker for the Ray-DP regression test. Must be importable
+# by ``multiprocessing.spawn`` (no closures, no nesting). Mirrors the worker
+# in ``tests/entrypoints/test_api_server_process_manager.py``.
+def _bind_and_report_worker(listen_address, sock, args, client_config):
+    """Bind ROUTER/PULL with a kernel-assigned port, report the actual
+    endpoints back via ``actual_address_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()
+
 
 class _DummyExecutor:
     pass
@@ -134,3 +186,172 @@ def test_driver_nixl_side_channel_host_does_not_leak_to_engine_core_actor(
         else:
             for pg in created_placement_groups:
                 ray.util.remove_placement_group(pg)
+
+
+@pytest.fixture
+def ray_context_dp2():
+    """Ray context sized for two stub actors (each PG needs ~1 CPU)."""
+    started_ray = False
+    if not ray.is_initialized():
+        project_root = str(Path(__file__).resolve().parents[3])
+        ray.init(
+            num_cpus=4,
+            runtime_env={"env_vars": {"PYTHONPATH": project_root}},
+            log_to_driver=False,
+        )
+        started_ray = True
+
+    yield
+
+    if started_ray:
+        ray.shutdown()
+
+
+def _make_vllm_config_ray_dp_multinode() -> SimpleNamespace:
+    """Minimal vllm_config that drives the Ray-DP multi-API-server path:
+    ``data_parallel_size != data_parallel_size_local`` forces TCP placeholders
+    (multi-node fan-out), and ``data_parallel_backend="ray"`` routes
+    ``launch_core_engines`` through the Ray branch.
+    """
+    return SimpleNamespace(
+        parallel_config=SimpleNamespace(
+            data_parallel_size=2,
+            data_parallel_size_local=1,
+            data_parallel_rank=0,
+            data_parallel_rank_local=None,
+            data_parallel_master_ip="127.0.0.1",
+            data_parallel_backend="ray",
+            data_parallel_rpc_port=29550,
+            local_engines_only=False,
+            enable_elastic_ep=False,
+            world_size=1,
+        ),
+        model_config=SimpleNamespace(multimodal_config=None, is_moe=False),
+        cache_config=SimpleNamespace(),
+        needs_dp_coordinator=False,
+        kv_transfer_config=None,
+        # ``_apply_dp_identity_suffix`` reads and rewrites this.
+        instance_id="vllm-ray-dp-regression-test",
+    )
+
+
+@pytest.mark.timeout(120)
+@pytest.mark.usefixtures("ray_context_dp2")
+def test_ray_dp_addresses_resolved_before_actor_creation(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    """Regression guard for the Ray-DP + multi-API-server hang from PR #42585.
+
+    ``launch_core_engines`` Ray branch pickles ``addresses`` into each engine
+    actor at ``.remote()`` time, and ``EngineCoreActorMixin._perform_handshakes``
+    is a no-op, so the actor uses that pickled snapshot for the rest of its
+    life. If ``run_multi_api_server`` allocates ``addresses`` as
+    ``tcp://host:0`` placeholders (its default), the actors hold placeholders
+    forever and DEALER-connect to port 0 — ZMQ ``connect`` is async and does
+    not raise, so the failure mode is a deterministic hang.
+
+    The Ray-DP carve-out in ``run_multi_api_server`` forces
+    ``defer_api_server_ports=False`` when ``data_parallel_backend == "ray"``
+    so addresses are pre-allocated in the driver and Ray pickles real ports
+    into each actor. This test mirrors that call-site logic and asserts the
+    actors hold real (non-placeholder) endpoints. If the carve-out is
+    removed without an alternative fix, the test fails.
+    """
+    created_placement_groups: list[Any] = []
+
+    def create_dp_placement_groups(vllm_config: Any):
+        pg1 = _make_cpu_placement_group()
+        pg2 = _make_cpu_placement_group()
+        created_placement_groups.extend([pg1, pg2])
+        return [pg1, pg2], [0, 0]
+
+    monkeypatch.setattr("vllm.v1.engine.core.EngineCoreActor", _StubEngineCoreActor)
+    monkeypatch.setattr(
+        CoreEngineActorManager,
+        "create_dp_placement_groups",
+        staticmethod(create_dp_placement_groups),
+    )
+
+    vllm_config = _make_vllm_config_ray_dp_multinode()
+
+    # Mirror run_multi_api_server's address-allocation logic. The Ray DP
+    # carve-out forces pre-allocation so the addresses pickled into engine
+    # actors at .remote() time are real, not ``tcp://host:0``.
+    is_ray_dp = vllm_config.parallel_config.data_parallel_backend == "ray"
+    addresses = get_engine_zmq_addresses(
+        vllm_config,
+        num_api_servers=2,
+        defer_api_server_ports=not is_ray_dp,
+    )
+
+    sock = socket.socket()
+    engine_manager: CoreEngineActorManager | None = None
+    actor_snapshots: list[tuple[list[str], list[str]]] = []
+    api_server_manager: APIServerProcessManager | None = None
+    try:
+        # Ray actors are spawned here, pickling ``addresses`` into each one.
+        with launch_core_engines(
+            vllm_config,
+            executor_class=_DummyExecutor,
+            log_stats=False,
+            addresses=addresses,
+            num_api_servers=2,
+        ) as (
+            engine_manager,
+            _coordinator,
+            _addresses_out,
+            _tensor_queue,
+        ):
+            assert isinstance(engine_manager, CoreEngineActorManager)
+
+            # API-server children bind to the pre-allocated ports.
+            api_server_manager = APIServerProcessManager(
+                listen_address="tcp://127.0.0.1:0",
+                sock=sock,
+                args="test_args",
+                num_servers=2,
+                input_addresses=addresses.inputs,
+                output_addresses=addresses.outputs,
+                target_server_fn=_bind_and_report_worker,
+            )
+
+            # run_multi_api_server skips ``gather_actual_addresses`` for
+            # Ray DP (addresses are already real). Mirror that.
+            if not is_ray_dp:
+                actual_inputs, actual_outputs = (
+                    api_server_manager.gather_actual_addresses(timeout=15.0)
+                )
+                addresses.inputs = actual_inputs
+                addresses.outputs = actual_outputs
+
+            # Snapshot what each Ray actor actually holds.
+            actors = (
+                engine_manager.local_engine_actors + engine_manager.remote_engine_actors
+            )
+            actor_snapshots = ray.get(
+                [actor.get_addresses.remote() for actor in actors]
+            )
+    finally:
+        if api_server_manager is not None:
+            api_server_manager.shutdown()
+            time.sleep(0.2)
+        sock.close()
+        if engine_manager is not None:
+            engine_manager.shutdown()
+        else:
+            for pg in created_placement_groups:
+                ray.util.remove_placement_group(pg)
+
+    # Every Ray actor must hold real, non-placeholder addresses.
+    assert actor_snapshots, "expected at least one Ray actor to be created"
+    for actor_inputs, actor_outputs in actor_snapshots:
+        for url in actor_inputs + actor_outputs:
+            scheme, _host, port = split_zmq_path(url)
+            assert scheme == "tcp", url
+            assert port and int(port) > 0, (
+                f"Ray actor was pickled with placeholder address {url!r}; "
+                "``run_multi_api_server`` must pre-allocate ports for the "
+                "Ray DP backend so the actors hold real endpoints by the "
+                "time they DEALER-connect. See PR #42585 / Ray-DP "
+                "multi-API-server regression."
+            )

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/_tilelang_ops.py

@@ -2,7 +2,7 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 import math
 from functools import cache
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Any
 
 import torch
 
@@ -10,8 +10,9 @@ from vllm.platforms import current_platform
 from vllm.utils.import_utils import has_tilelang
 from vllm.utils.math_utils import cdiv
 
-# tilelang is only available on CUDA platforms
-if TYPE_CHECKING or current_platform.is_cuda():
+# TileLang is used for MHC on CUDA and ROCm. Keep non-GPU imports cheap so
+# registering the Python wrapper modules does not require TileLang everywhere.
+if TYPE_CHECKING or current_platform.is_cuda_alike():
     if not has_tilelang():
         raise ImportError(
             "tilelang is required for mhc but is not installed. Install it with "
@@ -23,6 +24,8 @@ else:
     tilelang = None  # type: ignore[assignment]
     T = None  # type: ignore[assignment]
 
+ENABLE_PDL = current_platform.is_arch_support_pdl() and current_platform.is_cuda()
+
 
 @cache
 def compute_num_split(block_k: int, k: int | None, grid_size: int) -> int:
@@ -37,12 +40,17 @@ def compute_num_split(block_k: int, k: int | None, grid_size: int) -> int:
     return split_k
 
 
+pass_configs: dict[tilelang.PassConfigKey, Any] = {
+    tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True,
+    tilelang.PassConfigKey.TL_DISABLE_TMA_LOWER: True,
+}
+
+if current_platform.is_cuda():
+    pass_configs[tilelang.PassConfigKey.TL_PTXAS_REGISTER_USAGE_LEVEL] = 10
+
+
 @tilelang.jit(
-    pass_configs={
-        tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True,
-        tilelang.PassConfigKey.TL_DISABLE_TMA_LOWER: True,
-        tilelang.PassConfigKey.TL_PTXAS_REGISTER_USAGE_LEVEL: 10,
-    },
+    pass_configs=pass_configs,
 )
 def mhc_pre_big_fuse_tilelang(
     gemm_out_mul,
@@ -78,7 +86,8 @@ def mhc_pre_big_fuse_tilelang(
     layer_input: T.Tensor[[num_tokens, hidden_size], T.bfloat16]  # type: ignore[no-redef, valid-type]
 
     with T.Kernel(num_tokens, threads=96) as i:
-        T.pdl_sync()
+        if ENABLE_PDL:
+            T.pdl_sync()
         ##################################################################
         # _pre_norm_fn_fwd_norm
         rms = T.alloc_fragment(1, T.float32)
@@ -174,18 +183,16 @@ def mhc_pre_big_fuse_tilelang(
                         ol[i1_h] += pre * xl[i_hc, i1_h]
 
                 T.copy(ol, layer_input[i, i0_h * hidden_block])
-        T.pdl_trigger()
+
+        if ENABLE_PDL:
+            T.pdl_trigger()
 
 
 # Copied from https://github.com/sgl-project/sglang/blob/main/python/sglang/srt/layers/mhc.py#L478
 
 
 @tilelang.jit(
-    pass_configs={
-        tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True,
-        tilelang.PassConfigKey.TL_DISABLE_TMA_LOWER: True,
-        tilelang.PassConfigKey.TL_PTXAS_REGISTER_USAGE_LEVEL: 10,
-    },
+    pass_configs=pass_configs,
 )
 def mhc_pre_big_fuse_with_norm_tilelang(
     gemm_out_mul,
@@ -230,7 +237,8 @@ def mhc_pre_big_fuse_with_norm_tilelang(
         T.clear(mixes)
         rms[0] = 0
 
-        T.pdl_sync()
+        if ENABLE_PDL:
+            T.pdl_sync()
 
         for i_split in T.serial(n_splits):
             rms[0] += gemm_out_sqrsum[i_split, i]
@@ -341,15 +349,12 @@ def mhc_pre_big_fuse_with_norm_tilelang(
 
                 T.copy(ol, layer_input[i, i0_h * hidden_block])
 
-        T.pdl_trigger()
+        if ENABLE_PDL:
+            T.pdl_trigger()
 
 
 @tilelang.jit(
-    pass_configs={
-        tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True,
-        tilelang.PassConfigKey.TL_DISABLE_TMA_LOWER: True,
-        tilelang.PassConfigKey.TL_PTXAS_REGISTER_USAGE_LEVEL: 10,
-    },
+    pass_configs=pass_configs,
 )
 def mhc_fused_tilelang(
     comb_mix,
@@ -390,8 +395,8 @@ def mhc_fused_tilelang(
 
     with T.Kernel(m, n_tiles, split_k, threads=n_thr) as (i_n, i_nt, i_ks):
         tid = T.get_thread_binding()
-        warp_id = T.get_warp_idx()
-        lane = T.get_lane_idx()
+        warp_id = tid // 32
+        lane = tid % 32
 
         s_warp = T.alloc_shared((num_warps, tile_n + 1), T.float32)
         s_post = T.alloc_shared((hc,), T.float32)
@@ -407,7 +412,8 @@ def mhc_fused_tilelang(
         T.clear(sqr)
         h_split_start = i_ks * h_per_split
 
-        T.pdl_sync()
+        if ENABLE_PDL:
+            T.pdl_sync()
 
         T.copy(post_mix[i_n, 0], s_post)
         T.copy(comb_mix[i_n, 0, 0], s_comb)
@@ -466,15 +472,12 @@ def mhc_fused_tilelang(
                     v2 += s_warp[w, tile_n]
                 rp_out[i_ks, i_n] = v2
 
-        T.pdl_trigger()
+        if ENABLE_PDL:
+            T.pdl_trigger()
 
 
 @tilelang.jit(
-    pass_configs={
-        tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True,
-        tilelang.PassConfigKey.TL_DISABLE_TMA_LOWER: True,
-        tilelang.PassConfigKey.TL_PTXAS_REGISTER_USAGE_LEVEL: 10,
-    },
+    pass_configs=pass_configs,
 )
 def mhc_post_tilelang(
     a,
@@ -507,7 +510,8 @@ def mhc_post_tilelang(
 
         a_local = T.alloc_fragment((hc, hc), T.float32)
         c_local = T.alloc_fragment(hc, T.float32)
-        T.pdl_sync()
+        if ENABLE_PDL:
+            T.pdl_sync()
         T.copy(a[i_n, 0, 0], a_local)
         T.copy(c[i_n, 0], c_local)
 
@@ -523,15 +527,193 @@ def mhc_post_tilelang(
                     x_local[i_hco, i1_h] += a_local[i_hci, i_hco] * b_local[i_hci, i1_h]
 
             T.copy(x_local, x[i_n, 0, i0_h * h_blk])
-        T.pdl_trigger()
+        if ENABLE_PDL:
+            T.pdl_trigger()
 
 
 @tilelang.jit(
-    pass_configs={
-        tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True,
-        tilelang.PassConfigKey.TL_DISABLE_TMA_LOWER: True,
-        tilelang.PassConfigKey.TL_PTXAS_REGISTER_USAGE_LEVEL: 10,
-    },
+    pass_configs=pass_configs,
+)
+def hc_prenorm_gemm_tilelang(
+    x,
+    fn,
+    out,
+    sqrsum,
+    hidden_size: int,
+    hc_mult: int = 4,
+    n_out: int = 24,
+    n_thr: int = 512,
+    tile_n: int = 12,
+    n_splits: int = 1,
+) -> tilelang.JITKernel:
+    num_tokens = T.dynamic("num_tokens")
+    hc_hidden_size = hc_mult * hidden_size
+    k_per_split = hc_hidden_size // n_splits
+    k_iters = k_per_split // n_thr
+    n_tiles = T.ceildiv(n_out, tile_n)
+
+    x: T.Tensor((num_tokens, hc_hidden_size), T.bfloat16)  # type: ignore[no-redef, valid-type]
+    fn: T.Tensor((n_out, hc_hidden_size), T.float32)  # type: ignore[no-redef, valid-type]
+    out: T.Tensor((n_splits, num_tokens, n_out), T.float32)  # type: ignore[no-redef, valid-type]
+    sqrsum: T.Tensor((n_splits, num_tokens), T.float32)  # type: ignore[no-redef, valid-type]
+
+    with T.Kernel(num_tokens, n_tiles, n_splits, threads=n_thr) as (
+        i_n,
+        i_t,
+        i_s,
+    ):
+        tid = T.get_thread_binding()
+        acc = T.alloc_local((tile_n,), T.float32)
+        sqr = T.alloc_local((1,), T.float32)
+        T.clear(acc)
+        T.clear(sqr)
+
+        if ENABLE_PDL:
+            T.pdl_sync()
+
+        for it in T.serial(k_iters):
+            i_k = i_s * k_per_split + it * n_thr + tid
+            x_val = x[i_n, i_k]
+            for i_o in T.unroll(tile_n):
+                out_idx = i_t * tile_n + i_o
+                if out_idx < n_out:
+                    acc[i_o] += x_val * fn[out_idx, i_k]
+            if i_t == 0:
+                sqr[0] += x_val * x_val
+
+        for i_o in T.unroll(tile_n):
+            acc[i_o] = T.warp_reduce_sum(acc[i_o])
+        if i_t == 0:
+            sqr[0] = T.warp_reduce_sum(sqr[0])
+
+        lane = tid % 32
+        warp_id = tid // 32
+        num_warps = n_thr // 32
+        warp_acc = T.alloc_shared((num_warps, tile_n), T.float32)
+        warp_sqr = T.alloc_shared(num_warps, T.float32)
+
+        if lane == 0:
+            for i_o in T.unroll(tile_n):
+                warp_acc[warp_id, i_o] = acc[i_o]
+            if i_t == 0:
+                warp_sqr[warp_id] = sqr[0]
+        T.sync_threads()
+
+        if warp_id == 0:
+            if lane < tile_n:
+                reduced_acc = T.alloc_var(T.float32, init=0.0)
+                for i_w in T.unroll(num_warps):
+                    reduced_acc += warp_acc[i_w, lane]
+                out_idx = i_t * tile_n + lane
+                if out_idx < n_out:
+                    out[i_s, i_n, out_idx] = reduced_acc
+            if lane == 0 and i_t == 0:
+                reduced_sqr = T.alloc_var(T.float32, init=0.0)
+                for i_w in T.unroll(num_warps):
+                    reduced_sqr += warp_sqr[i_w]
+                sqrsum[i_s, i_n] = reduced_sqr
+
+        if ENABLE_PDL:
+            T.pdl_trigger()
+
+
+@tilelang.jit(
+    pass_configs=pass_configs,
+)
+def hc_prenorm_gemm_block_m_tilelang(
+    x,
+    fn,
+    out,
+    sqrsum,
+    hidden_size: int,
+    hc_mult: int = 4,
+    n_out: int = 24,
+    n_thr: int = 512,
+    tile_n: int = 12,
+    block_m: int = 2,
+) -> tilelang.JITKernel:
+    num_tokens = T.dynamic("num_tokens")
+    hc_hidden_size = hc_mult * hidden_size
+    k_iters = hc_hidden_size // n_thr
+    n_tiles = T.ceildiv(n_out, tile_n)
+    m_tiles = T.ceildiv(num_tokens, block_m)
+
+    x: T.Tensor((num_tokens, hc_hidden_size), T.bfloat16)  # type: ignore[no-redef, valid-type]
+    fn: T.Tensor((n_out, hc_hidden_size), T.float32)  # type: ignore[no-redef, valid-type]
+    out: T.Tensor((1, num_tokens, n_out), T.float32)  # type: ignore[no-redef, valid-type]
+    sqrsum: T.Tensor((1, num_tokens), T.float32)  # type: ignore[no-redef, valid-type]
+
+    with T.Kernel(m_tiles, n_tiles, threads=n_thr) as (i_mt, i_t):
+        tid = T.get_thread_binding()
+        acc = T.alloc_local((block_m, tile_n), T.float32)
+        sqr = T.alloc_local((block_m,), T.float32)
+        T.clear(acc)
+        T.clear(sqr)
+
+        if ENABLE_PDL:
+            T.pdl_sync()
+
+        for it in T.serial(k_iters):
+            i_k = it * n_thr + tid
+            fn_val = T.alloc_local((tile_n,), T.float32)
+            for i_o in T.unroll(tile_n):
+                out_idx = i_t * tile_n + i_o
+                if out_idx < n_out:
+                    fn_val[i_o] = fn[out_idx, i_k]
+                else:
+                    fn_val[i_o] = 0.0
+            for i_m in T.unroll(block_m):
+                token_idx = i_mt * block_m + i_m
+                if token_idx < num_tokens:
+                    x_val = x[token_idx, i_k]
+                    for i_o in T.unroll(tile_n):
+                        acc[i_m, i_o] += x_val * fn_val[i_o]
+                    if i_t == 0:
+                        sqr[i_m] += x_val * x_val
+
+        for i_m in T.unroll(block_m):
+            for i_o in T.unroll(tile_n):
+                acc[i_m, i_o] = T.warp_reduce_sum(acc[i_m, i_o])
+            if i_t == 0:
+                sqr[i_m] = T.warp_reduce_sum(sqr[i_m])
+
+        lane = tid % 32
+        warp_id = tid // 32
+        num_warps = n_thr // 32
+        warp_acc = T.alloc_shared((num_warps, block_m, tile_n), T.float32)
+        warp_sqr = T.alloc_shared((num_warps, block_m), T.float32)
+
+        if lane == 0:
+            for i_m in T.unroll(block_m):
+                for i_o in T.unroll(tile_n):
+                    warp_acc[warp_id, i_m, i_o] = acc[i_m, i_o]
+                if i_t == 0:
+                    warp_sqr[warp_id, i_m] = sqr[i_m]
+        T.sync_threads()
+
+        if warp_id == 0:
+            for i_m in T.unroll(block_m):
+                token_idx = i_mt * block_m + i_m
+                if token_idx < num_tokens:
+                    if lane < tile_n:
+                        reduced_acc = T.alloc_var(T.float32, init=0.0)
+                        for i_w in T.unroll(num_warps):
+                            reduced_acc += warp_acc[i_w, i_m, lane]
+                        out_idx = i_t * tile_n + lane
+                        if out_idx < n_out:
+                            out[0, token_idx, out_idx] = reduced_acc
+                    if lane == 0 and i_t == 0:
+                        reduced_sqr = T.alloc_var(T.float32, init=0.0)
+                        for i_w in T.unroll(num_warps):
+                            reduced_sqr += warp_sqr[i_w, i_m]
+                        sqrsum[0, token_idx] = reduced_sqr
+
+        if ENABLE_PDL:
+            T.pdl_trigger()
+
+
+@tilelang.jit(
+    pass_configs=pass_configs,
 )
 def hc_head_fuse_tilelang(
     residual,
@@ -566,7 +748,8 @@ def hc_head_fuse_tilelang(
     out: T.Tensor[[num_tokens, hidden_size], T.bfloat16]  # type: ignore[no-redef,valid-type]
 
     with T.Kernel(num_tokens, threads=n_thr) as i:
-        T.pdl_sync()
+        if ENABLE_PDL:
+            T.pdl_sync()
 
         # ------------------------------------------------------------------
         # Pass 1 – for each residual channel m_c and h_block:
@@ -624,4 +807,5 @@ def hc_head_fuse_tilelang(
 
             T.copy(ol, out[i, i0_h * h_block], disable_tma=True)
 
-        T.pdl_trigger()
+        if ENABLE_PDL:
+            T.pdl_trigger()

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,17 @@ 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)
+    # Defer port allocation to the child's bind() to avoid TOCTOU, except
+    # for Rust front-end and Ray DP, which can't see the post-bind rebind
+    # (CLI-arg subprocess / pickled-into-actor snapshot respectively) and
+    # so pre-allocate driver-side -- reintroducing the original race only
+    # there.
+    is_ray_dp = parallel_config.data_parallel_backend == "ray"
+    addresses = get_engine_zmq_addresses(
+        vllm_config,
+        num_api_servers,
+        defer_api_server_ports=not (rust_frontend_path or is_ray_dp),
+    )
 
     with launch_core_engines(
         vllm_config, executor_class, log_stats, addresses, num_api_servers
@@ -341,6 +351,16 @@ def run_multi_api_server(args: argparse.Namespace):
                 tensor_queue=tensor_queue,
             )
 
+            if not is_ray_dp:
+                # Forward each child's bound endpoints to the engine handshake
+                # (runs on ``with`` exit). Skipped for Ray DP, where addresses
+                # are pre-allocated above and Ray actors already hold them.
+                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/entrypoints/openai/chat_completion/protocol.py

@@ -472,17 +472,27 @@ class ChatCompletionRequest(OpenAIBaseModel):
         default_template: str | None,
         default_template_content_format: ChatTemplateContentFormatOption,
     ) -> ChatParams:
+        extra_kwargs: dict[str, Any] = dict(
+            add_generation_prompt=self.add_generation_prompt,
+            continue_final_message=self.continue_final_message,
+            documents=self.documents,
+            reasoning_effort=self.reasoning_effort,
+        )
+
+        # When reasoning is requested, activate thinking for models whose
+        # chat templates require explicit opt-in (e.g., Gemma4 defaults
+        # enable_thinking to false). For templates that don't declare the
+        # variable, resolve_chat_template_kwargs filters it out harmlessly.
+        user_kwargs = self.chat_template_kwargs or {}
+        if self.reasoning_effort is not None and "enable_thinking" not in user_kwargs:
+            extra_kwargs["enable_thinking"] = self.reasoning_effort != "none"
+
         return ChatParams(
             chat_template=self.chat_template or default_template,
             chat_template_content_format=default_template_content_format,
             chat_template_kwargs=merge_kwargs(
                 self.chat_template_kwargs,
-                dict(
-                    add_generation_prompt=self.add_generation_prompt,
-                    continue_final_message=self.continue_final_message,
-                    documents=self.documents,
-                    reasoning_effort=self.reasoning_effort,
-                ),
+                extra_kwargs,
             ),
             media_io_kwargs=self.media_io_kwargs,
         )

vllm/entrypoints/openai/responses/protocol.py

@@ -298,17 +298,28 @@ class ResponsesRequest(OpenAIBaseModel):
         continue_final = should_continue_final_message(self.input)
 
         reasoning = self.reasoning
+        reasoning_effort = None if reasoning is None else reasoning.effort
+
+        extra_kwargs: dict[str, Any] = dict(
+            add_generation_prompt=not continue_final,
+            continue_final_message=continue_final,
+            reasoning_effort=reasoning_effort,
+        )
+
+        # When reasoning is requested, activate thinking for models whose
+        # chat templates require explicit opt-in (e.g., Gemma4 defaults
+        # enable_thinking to false). For templates that don't declare the
+        # variable, resolve_chat_template_kwargs filters it out harmlessly.
+        user_kwargs = self.chat_template_kwargs or {}
+        if reasoning_effort is not None and "enable_thinking" not in user_kwargs:
+            extra_kwargs["enable_thinking"] = reasoning_effort != "none"
 
         return ChatParams(
             chat_template=default_template,
             chat_template_content_format=default_template_content_format,
-            chat_template_kwargs=merge_kwargs(  # To remove unset values
+            chat_template_kwargs=merge_kwargs(
                 self.chat_template_kwargs,
-                dict(
-                    add_generation_prompt=not continue_final,
-                    continue_final_message=continue_final,
-                    reasoning_effort=None if reasoning is None else reasoning.effort,
-                ),
+                extra_kwargs,
             ),
             media_io_kwargs=self.media_io_kwargs,
         )

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/kernels/linear/__init__.py

@@ -58,6 +58,9 @@ from vllm.model_executor.kernels.linear.mixed_precision.xpu import (
     XPUW4A8IntLinearKernel,
     XPUwNa16LinearKernel,
 )
+from vllm.model_executor.kernels.linear.mixed_precision.zentorch import (
+    ZentorchWNA16LinearKernel,
+)
 from vllm.model_executor.kernels.linear.mxfp4 import (
     MxFp4LinearKernel,
     MxFp4LinearLayerConfig,
@@ -157,6 +160,9 @@ from vllm.model_executor.kernels.linear.scaled_mm.triton import (
 from vllm.model_executor.kernels.linear.scaled_mm.xpu import (
     XPUFP8ScaledMMLinearKernel,
 )
+from vllm.model_executor.kernels.linear.scaled_mm.zentorch import (
+    ZentorchInt8ScaledMMLinearKernel,
+)
 from vllm.model_executor.layers.quantization.utils.quant_utils import QuantKey
 from vllm.platforms import PlatformEnum, current_platform
 
@@ -254,7 +260,7 @@ def _filter_kernels_by_backend(
 
 # in priority/performance order (when available)
 _POSSIBLE_INT8_KERNELS: dict[PlatformEnum, list[type[Int8ScaledMMLinearKernel]]] = {
-    PlatformEnum.CPU: [CPUInt8ScaledMMLinearKernel],
+    PlatformEnum.CPU: [ZentorchInt8ScaledMMLinearKernel, CPUInt8ScaledMMLinearKernel],
     PlatformEnum.CUDA: [
         CutlassInt8ScaledMMLinearKernel,
         TritonInt8ScaledMMLinearKernel,
@@ -348,6 +354,7 @@ _POSSIBLE_KERNELS: dict[PlatformEnum, list[type[MPLinearKernel]]] = {
     ],
     PlatformEnum.CPU: [
         Dynamic4bitLinearKernel,
+        ZentorchWNA16LinearKernel,
         CPUWNA16LinearKernel,
     ],
 }
@@ -1018,6 +1025,8 @@ __all__ = [
     "RowWiseTorchFP8ScaledMMLinearKernel",
     "ROCmFP8ScaledMMLinearKernel",
     "TritonInt8ScaledMMLinearKernel",
+    "ZentorchInt8ScaledMMLinearKernel",
+    "ZentorchWNA16LinearKernel",
     "MPLinearKernel",
     "MPLinearLayerConfig",
     "AllSparkLinearKernel",

vllm/model_executor/kernels/linear/mixed_precision/__init__.py

@@ -36,6 +36,9 @@ from vllm.model_executor.kernels.linear.mixed_precision.xpu import (
     XPUW4A8IntLinearKernel,
     XPUwNa16LinearKernel,
 )
+from vllm.model_executor.kernels.linear.mixed_precision.zentorch import (
+    ZentorchWNA16LinearKernel,
+)
 
 __all__ = [
     "MPLinearKernel",
@@ -51,4 +54,5 @@ __all__ = [
     "TritonW4A16LinearKernel",
     "XPUW4A8IntLinearKernel",
     "XPUwNa16LinearKernel",
+    "ZentorchWNA16LinearKernel",
 ]

vllm/model_executor/kernels/linear/mixed_precision/zentorch.py

@@ -0,0 +1,211 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Zentorch W4A16 GPTQ weight-only-quantized linear kernel for AMD Zen CPUs.
+
+Selected by ``choose_mp_linear_kernel`` ahead of the generic oneDNN-backed
+``CPUWNA16LinearKernel``. When ``can_implement`` rejects a layer, the selector
+falls through to the next kernel in ``_POSSIBLE_KERNELS[PlatformEnum.CPU]``.
+"""
+
+import torch
+
+from vllm.logger import init_logger
+from vllm.model_executor.kernels.linear.zentorch_utils import has_zentorch_op
+from vllm.platforms import current_platform
+from vllm.scalar_type import scalar_types
+
+from .cpu import CPUWNA16LinearKernel
+from .MPLinearKernel import MPLinearLayerConfig
+
+logger = init_logger(__name__)
+
+
+def _import_unpack_from_int32():
+    """Import compressed-tensors' ``unpack_from_int32`` across versions."""
+    try:
+        from compressed_tensors.compressors.pack_quantized.helpers import (
+            unpack_from_int32,
+        )
+    except ImportError:
+        from compressed_tensors.compressors.quantized_compressors.pack_quantized import (  # type: ignore[import-not-found]  # noqa: E501
+            unpack_from_int32,
+        )
+    return unpack_from_int32
+
+
+class ZentorchWNA16LinearKernel(CPUWNA16LinearKernel):
+    """W4A16 GPTQ kernel backed by ``torch.ops.zentorch.zentorch_woq_linear``."""
+
+    @classmethod
+    def can_implement(cls, c: MPLinearLayerConfig) -> tuple[bool, str | None]:
+        ok, reason = super().can_implement(c)
+        if not ok:
+            return ok, reason
+
+        if not current_platform.is_zen_cpu():
+            return False, "ZentorchWNA16 requires an AMD Zen CPU."
+
+        if not has_zentorch_op(["zentorch_woq_repack_weight", "zentorch_woq_linear"]):
+            return (
+                False,
+                "torch.ops.zentorch.{zentorch_woq_repack_weight, "
+                "zentorch_woq_linear} are not registered.",
+            )
+
+        if c.has_g_idx:
+            return False, "ZentorchWNA16 does not support activation re-ordering."
+        return True, None
+
+    def _zentorch_woq_eligible(self, layer: torch.nn.Module) -> bool:
+        """Eligibility predicate for the zentorch W4A16 GPTQ fast path.
+
+        Constraints (any failure -> ``cpu_gemm_wna16`` path via ``super()``
+        with ``layer`` untouched).
+        """
+        if (
+            self.w_gidx_name is not None
+            and getattr(layer, self.w_gidx_name, None) is not None
+        ) or (getattr(self.config, "has_g_idx", False)):
+            return False
+
+        weight_packed = getattr(layer, self.w_q_name, None)
+        weight_scale = getattr(layer, self.w_s_name, None)
+        if weight_packed is None or weight_scale is None:
+            return False
+
+        bits = self.config.weight_type.mantissa
+        pack_factor = torch.iinfo(weight_packed.dtype).bits // bits
+        # 4-bit -> 8 values per int32;
+        if pack_factor != 8:
+            return False
+
+        # GPTQ-only. AWQ packs along the output dim instead.
+        in_dim = getattr(weight_packed, "input_dim", None)
+        pk_dim = getattr(weight_packed, "packed_dim", None)
+        if in_dim is None or pk_dim is None or in_dim != pk_dim:
+            return False
+
+        is_ct_format = in_dim == pk_dim == 1
+        if not is_ct_format:
+            return False
+
+        if weight_packed.dim() != 2 or weight_scale.dim() != 2:
+            return False
+
+        # 4-bit -> 8 values per int32; in_features must be divisible by num_groups.
+        in_features = weight_packed.shape[1] * 8
+        num_groups = weight_scale.shape[1]
+        return num_groups > 0 and in_features % num_groups == 0
+
+    def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
+        """Repack CT GPTQ weights into the zentorch WOQ layout.
+
+        Falls back to ``CPUWNA16LinearKernel.process_weights_after_loading``
+        via ``super()`` when the layer doesn't satisfy
+        ``_zentorch_woq_eligible``.
+
+        On success, ``layer._zentorch_processed_weights`` is set to ``True``
+        """
+        if getattr(layer, "_zentorch_processed_weights", False):
+            return
+
+        if not self._zentorch_woq_eligible(layer):
+            logger.info_once(
+                "[zen_cpu] ZentorchWNA16 fast path not eligible for this "
+                "layer (AWQ pack layout, g_idx, or non-int32 storage); "
+                "falling back to CPUWNA16LinearKernel (cpu_gemm_wna16)."
+            )
+            super().process_weights_after_loading(layer)
+            return
+
+        if (not self.config.zero_points) and (self.w_zp_name is not None):
+            setattr(layer, self.w_zp_name, None)
+
+        if (not self.config.has_g_idx) and (self.w_gidx_name is not None):
+            setattr(layer, self.w_gidx_name, None)
+
+        weight_q = getattr(layer, self.w_q_name)
+        weight_s = getattr(layer, self.w_s_name)
+        weight_packed = weight_q.data if hasattr(weight_q, "data") else weight_q
+        weight_scale = weight_s.data if hasattr(weight_s, "data") else weight_s
+
+        bits = self.config.weight_type.mantissa
+        pack_factor = torch.iinfo(weight_packed.dtype).bits // bits
+        out_features, num_groups = weight_scale.shape[0], weight_scale.shape[1]
+        in_features = weight_packed.shape[1] * pack_factor
+        original_shape = torch.Size([out_features, in_features])
+        unpack_from_int32 = _import_unpack_from_int32()
+        repack_op = torch.ops.zentorch.zentorch_woq_repack_weight.default
+
+        weight_unpacked = unpack_from_int32(
+            weight_packed,
+            bits,
+            original_shape,
+            packed_dim=weight_q.packed_dim,
+        )
+
+        zp_param = (
+            getattr(layer, self.w_zp_name, None) if self.w_zp_name is not None else None
+        )
+        needs_unsigned_offset = self.config.weight_type == scalar_types.uint4
+
+        if needs_unsigned_offset:
+            weight_unpacked = (weight_unpacked.to(torch.int32) + 8).clamp(0, 15)
+        repacked = repack_op(weight_unpacked.to(torch.int8).contiguous())
+
+        if zp_param is None:
+            zp_tc = None
+        else:
+            zp_tensor = zp_param.data if hasattr(zp_param, "data") else zp_param
+            zp = unpack_from_int32(
+                zp_tensor,
+                bits,
+                (out_features, num_groups),
+                packed_dim=zp_param.packed_dim,
+            )
+            if needs_unsigned_offset:
+                zp = (zp.to(torch.int32) + 8).clamp(0, 15)
+            zp_tc = zp.to(torch.int8).t().contiguous()
+
+        layer._zentorch_woq_packed = repacked.t()
+        layer._zentorch_woq_scale = weight_scale.t().contiguous()
+        layer._zentorch_woq_zero_point = zp_tc
+
+        for param_name in (self.w_q_name, self.w_s_name, self.w_zp_name):
+            if param_name is None:
+                continue
+            param = getattr(layer, param_name, None)
+            if param is None:
+                continue
+            if hasattr(param, "data"):
+                param.data = torch.empty(0)
+            else:
+                setattr(layer, param_name, torch.empty(0))
+
+        layer._zentorch_kind = "compressed_tensors_w4a16_gptq"
+        layer._zentorch_processed_weights = True
+        logger.info_once(
+            "[zen_cpu] Using zentorch_woq_linear for W4A16 GPTQ "
+            "(weight_type=%s, has_zp=%s)",
+            self.config.weight_type,
+            zp_tc is not None,
+        )
+
+    def apply_weights(
+        self,
+        layer: torch.nn.Module,
+        x: torch.Tensor,
+        bias: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        if getattr(layer, "_zentorch_processed_weights", False):
+            return torch.ops.zentorch.zentorch_woq_linear.default(
+                x,
+                layer._zentorch_woq_packed,
+                layer._zentorch_woq_scale,
+                layer._zentorch_woq_zero_point,
+                bias,
+            )
+        return super().apply_weights(layer, x, bias)
+
+
+__all__ = ["ZentorchWNA16LinearKernel"]

vllm/model_executor/kernels/linear/scaled_mm/__init__.py

@@ -39,6 +39,9 @@ from vllm.model_executor.kernels.linear.scaled_mm.ScaledMMLinearKernel import (
 from vllm.model_executor.kernels.linear.scaled_mm.triton import (
     TritonInt8ScaledMMLinearKernel,
 )
+from vllm.model_executor.kernels.linear.scaled_mm.zentorch import (
+    ZentorchInt8ScaledMMLinearKernel,
+)
 
 __all__ = [
     "FP8ScaledMMLinearKernel",
@@ -58,6 +61,7 @@ __all__ = [
     "RowWiseTorchFP8ScaledMMLinearKernel",
     "ROCmFP8ScaledMMLinearKernel",
     "TritonInt8ScaledMMLinearKernel",
+    "ZentorchInt8ScaledMMLinearKernel",
     "Fp8BlockScaledMMLinearKernel",
     "CPUFp8BlockScaledMMKernel",
 ]

vllm/model_executor/kernels/linear/scaled_mm/zentorch.py

@@ -0,0 +1,98 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Zentorch dynamic-symmetric W8A8 int8 linear kernel for AMD Zen CPUs.
+
+Selected by ``choose_scaled_mm_linear_kernel`` ahead of the generic
+oneDNN-backed ``CPUInt8ScaledMMLinearKernel``. When ``is_supported`` or
+``can_implement`` rejects a layer, the selector falls through to the next
+kernel in ``_POSSIBLE_INT8_KERNELS[PlatformEnum.CPU]``.
+"""
+
+import torch
+
+from vllm.logger import init_logger
+from vllm.model_executor.kernels.linear.zentorch_utils import has_zentorch_op
+from vllm.model_executor.layers.quantization.utils import replace_parameter
+from vllm.platforms import current_platform
+
+from .ScaledMMLinearKernel import (
+    Int8ScaledMMLinearKernel,
+    Int8ScaledMMLinearLayerConfig,
+)
+
+logger = init_logger(__name__)
+
+
+class ZentorchInt8ScaledMMLinearKernel(Int8ScaledMMLinearKernel):
+    @classmethod
+    def is_supported(
+        cls, compute_capability: int | None = None
+    ) -> tuple[bool, str | None]:
+        if not current_platform.is_cpu():
+            return False, "requires CPU."
+        if not current_platform.is_zen_cpu():
+            return False, "requires AMD Zen CPU."
+        if not has_zentorch_op(["zentorch_dynamic_qlinear"]):
+            return (
+                False,
+                "torch.ops.zentorch.zentorch_dynamic_qlinear is not registered.",
+            )
+        return True, None
+
+    @classmethod
+    def can_implement(cls, c: Int8ScaledMMLinearLayerConfig) -> tuple[bool, str | None]:
+        if c.is_static_input_scheme:
+            return False, "requires dynamic activation quantization."
+        if not c.input_symmetric:
+            return False, "requires symmetric activation quantization."
+        if not c.is_channelwise:
+            return False, "requires per-channel weight quantization."
+        return True, None
+
+    def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
+        """Prepare weights for ``zentorch_dynamic_qlinear``.
+
+        Keeps weight in [N, K] layout (int8, contiguous) and converts the
+        per-channel weight scale to bf16 with shape ``(N,)``.
+        """
+        w_q_name, w_s_name, _, _, _ = self.layer_param_names
+        weight = getattr(layer, w_q_name)
+        n = weight.shape[0]
+        replace_parameter(
+            layer,
+            w_q_name,
+            torch.nn.Parameter(weight.data.contiguous(), requires_grad=False),
+        )
+
+        weight_scale = getattr(layer, w_s_name)
+        ws = weight_scale.data
+        if ws.dim() == 2 and ws.shape[-1] == 1:
+            ws = ws.squeeze(-1)
+        ws = ws.to(torch.bfloat16).contiguous()
+        assert ws.shape == (n,), (
+            f"[zen_cpu] expected weight scale shape ({n},), got {tuple(ws.shape)}"
+        )
+
+        replace_parameter(
+            layer,
+            w_s_name,
+            torch.nn.Parameter(ws, requires_grad=False),
+        )
+        logger.info_once(
+            "[zen_cpu] Using zentorch_dynamic_qlinear for W8A8 (dynamic-symmetric)"
+        )
+
+    def apply_weights(
+        self,
+        layer: torch.nn.Module,
+        x: torch.Tensor,
+        bias: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        w_q_name, w_s_name, _, _, _ = self.layer_param_names
+        return torch.ops.zentorch.zentorch_dynamic_qlinear(
+            x,
+            getattr(layer, w_q_name),
+            getattr(layer, w_s_name),
+            bias,
+            zentorch_op_name="zentorch::zentorch_dynamic_qlinear",
+        )

vllm/model_executor/kernels/linear/zentorch_utils.py

@@ -0,0 +1,23 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Gates zentorch CPU linear dispatch on platform/op availability."""
+
+from __future__ import annotations
+
+import torch
+
+from vllm.platforms import current_platform
+
+__all__ = ["has_zentorch_op"]
+
+
+def has_zentorch_op(op_names: list[str]) -> bool:
+    """Return ``True`` when running on Zen CPU with all named ops registered."""
+    if not op_names:
+        raise ValueError("has_zentorch_op requires at least one op name")
+    if not current_platform.is_zen_cpu():
+        return False
+    ns = getattr(torch.ops, "zentorch", None)
+    if ns is None:
+        return False
+    return all(hasattr(ns, op_name) for op_name in op_names)

vllm/model_executor/kernels/mhc/tilelang.py

@@ -5,6 +5,88 @@ import torch
 from vllm.utils.torch_utils import direct_register_custom_op
 
 
+def _torch_hc_prenorm_gemm(
+    x: torch.Tensor,
+    fn: torch.Tensor,
+    out: torch.Tensor,
+    sqrsum: torch.Tensor,
+) -> None:
+    assert out.shape[0] == 1
+    assert sqrsum.shape[0] == 1
+    x_float = x.float()
+    out[0].copy_(x_float @ fn.t())
+    sqrsum[0].copy_(x_float.square().sum(dim=-1))
+
+
+def _tilelang_hc_prenorm_gemm(
+    x: torch.Tensor,
+    fn: torch.Tensor,
+    out: torch.Tensor,
+    sqrsum: torch.Tensor,
+    hidden_size: int,
+    hc_mult: int,
+    tile_n: int = 12,
+    n_thr: int = 512,
+    n_splits: int = 1,
+) -> None:
+    from vllm._tilelang_ops import (
+        hc_prenorm_gemm_block_m_tilelang,
+        hc_prenorm_gemm_tilelang,
+    )
+
+    assert out.shape[0] == n_splits
+    assert sqrsum.shape[0] == n_splits
+    assert x.shape[1] == hc_mult * hidden_size
+    assert x.shape[1] % n_splits == 0
+    assert (x.shape[1] // n_splits) % n_thr == 0
+    use_default_config = tile_n == 12 and n_thr == 512
+    if n_splits == 1 and use_default_config and x.shape[0] >= 1024:
+        hc_prenorm_gemm_block_m_tilelang(
+            x,
+            fn,
+            out,
+            sqrsum,
+            hidden_size,
+            hc_mult,
+            fn.shape[0],
+            n_thr,
+            tile_n,
+            2,
+        )
+        return
+    if (
+        n_splits == 1
+        and use_default_config
+        and x.shape[0] < 128
+        and x.shape[1] % 1024 == 0
+    ):
+        hc_prenorm_gemm_tilelang(
+            x,
+            fn,
+            out,
+            sqrsum,
+            hidden_size,
+            hc_mult,
+            fn.shape[0],
+            1024,
+            4,
+            n_splits,
+        )
+        return
+    hc_prenorm_gemm_tilelang(
+        x,
+        fn,
+        out,
+        sqrsum,
+        hidden_size,
+        hc_mult,
+        fn.shape[0],
+        n_thr,
+        tile_n,
+        n_splits,
+    )
+
+
 def mhc_pre_tilelang(
     residual: torch.Tensor,
     fn: torch.Tensor,
@@ -80,10 +162,16 @@ def mhc_pre_tilelang(
     residual_flat = residual.view(-1, hc_mult, hidden_size)
     num_tokens = residual_flat.shape[0]
 
-    # these numbers are from deepgemm kernel impl
-    block_k = 64
-    block_m = 64
-    n_splits = compute_num_split(block_k, hc_hidden_size, cdiv(num_tokens, block_m))
+    from vllm.utils.deep_gemm import is_deep_gemm_supported
+
+    use_deep_gemm = is_deep_gemm_supported()
+    if use_deep_gemm:
+        # these numbers are from deepgemm kernel impl
+        block_k = 64
+        block_m = 64
+        n_splits = compute_num_split(block_k, hc_hidden_size, cdiv(num_tokens, block_m))
+    else:
+        n_splits = 1
 
     post_mix = torch.empty(
         num_tokens, hc_mult, dtype=torch.float32, device=residual.device
@@ -102,13 +190,24 @@ def mhc_pre_tilelang(
         n_splits, num_tokens, dtype=torch.float32, device=residual.device
     )
 
-    tf32_hc_prenorm_gemm(
-        residual_flat.view(num_tokens, hc_mult * hidden_size),
-        fn,
-        gemm_out_mul,
-        gemm_out_sqrsum,
-        n_splits,
-    )
+    residual_2d = residual_flat.view(num_tokens, hc_mult * hidden_size)
+    if use_deep_gemm:
+        tf32_hc_prenorm_gemm(
+            residual_2d,
+            fn,
+            gemm_out_mul,
+            gemm_out_sqrsum,
+            n_splits,
+        )
+    else:
+        _tilelang_hc_prenorm_gemm(
+            residual_2d,
+            fn,
+            gemm_out_mul,
+            gemm_out_sqrsum,
+            hidden_size,
+            hc_mult,
+        )
 
     if norm_weight is None:
         mhc_pre_big_fuse_tilelang(
@@ -304,16 +403,24 @@ def mhc_fused_post_pre_tilelang(
     post_layer_mix_flat = post_layer_mix.view(num_tokens, hc_mult)
     comb_res_mix_flat = comb_res_mix.view(num_tokens, hc_mult, hc_mult)
 
-    fma_token_threshold = 16
-    if num_tokens <= fma_token_threshold:
+    from vllm.utils.deep_gemm import is_deep_gemm_supported
+
+    use_deep_gemm = is_deep_gemm_supported()
+    use_small_fma = num_tokens <= 16
+    if use_small_fma:
         # TODO(gnovack): investigate autotuning these heuristics
         tile_n = 2 if num_tokens < 8 else 3
         n_splits = 8 if (num_tokens < 8 and hidden_size <= 4096) else 4
     else:
-        # these number are from deepgemm kernel impl
-        block_k = 64
-        block_m = 64
-        n_splits = compute_num_split(block_k, hc_hidden_size, cdiv(num_tokens, block_m))
+        if use_deep_gemm:
+            # these number are from deepgemm kernel impl
+            block_k = 64
+            block_m = 64
+            n_splits = compute_num_split(
+                block_k, hc_hidden_size, cdiv(num_tokens, block_m)
+            )
+        else:
+            n_splits = 1
 
     gemm_out_mul = torch.empty(
         n_splits,
@@ -348,7 +455,7 @@ def mhc_fused_post_pre_tilelang(
         device=residual.device,
     )
 
-    if num_tokens <= fma_token_threshold:
+    if use_small_fma:
         mhc_fused_tilelang(
             comb_res_mix_flat,
             residual_flat,
@@ -375,15 +482,26 @@ def mhc_fused_post_pre_tilelang(
             residual.shape[-1],
         )
 
-        from vllm.utils.deep_gemm import tf32_hc_prenorm_gemm
+        residual_cur_2d = residual_cur.view(num_tokens, hc_mult * hidden_size)
+        if use_deep_gemm:
+            from vllm.utils.deep_gemm import tf32_hc_prenorm_gemm
 
-        tf32_hc_prenorm_gemm(
-            residual_cur.view(num_tokens, hc_mult * hidden_size),
-            fn,
-            gemm_out_mul,
-            gemm_out_sqrsum,
-            n_splits,
-        )
+            tf32_hc_prenorm_gemm(
+                residual_cur_2d,
+                fn,
+                gemm_out_mul,
+                gemm_out_sqrsum,
+                n_splits,
+            )
+        else:
+            _tilelang_hc_prenorm_gemm(
+                residual_cur_2d,
+                fn,
+                gemm_out_mul,
+                gemm_out_sqrsum,
+                hidden_size,
+                hc_mult,
+            )
 
     if norm_weight is None:
         mhc_pre_big_fuse_tilelang(

vllm/model_executor/layers/fused_moe/experts/trtllm_bf16_moe.py

@@ -99,9 +99,6 @@ class TrtLlmBf16Experts(mk.FusedMoEExpertsMonolithic):
     ) -> bool:
         return True
 
-    def supports_chunking(self) -> bool:
-        return False
-
     def supports_expert_map(self) -> bool:
         return False
 
@@ -140,5 +137,6 @@ class TrtLlmBf16Experts(mk.FusedMoEExpertsMonolithic):
             intermediate_size=self.intermediate_size_per_partition,
             local_expert_offset=self.ep_rank * self.local_num_experts,
             local_num_experts=self.local_num_experts,
+            routed_scaling_factor=routed_scaling_factor,
             routing_method_type=self.routing_method_type,
         )

vllm/model_executor/layers/fused_moe/experts/trtllm_fp8_moe.py

@@ -88,9 +88,6 @@ class TrtLlmFp8ExpertsBase:
             or moe_parallel_config.use_ag_rs_all2all_kernels
         ) and not moe_parallel_config.enable_eplb
 
-    def supports_chunking(self) -> bool:
-        return False
-
     def supports_expert_map(self) -> bool:
         return False
 

vllm/model_executor/layers/fused_moe/experts/trtllm_mxfp4_moe.py

@@ -113,9 +113,6 @@ class TrtLlmMxfp4ExpertsBase:
     def activation_format() -> mk.FusedMoEActivationFormat:
         return mk.FusedMoEActivationFormat.Standard
 
-    def supports_chunking(self) -> bool:
-        return False
-
     def supports_expert_map(self) -> bool:
         return False
 

vllm/model_executor/layers/fused_moe/experts/trtllm_nvfp4_moe.py

@@ -157,8 +157,27 @@ class TrtLlmNvFp4ExpertsBase:
     def activation_format() -> mk.FusedMoEActivationFormat:
         return mk.FusedMoEActivationFormat.Standard
 
-    def supports_chunking(self) -> bool:
-        return False
+    def _get_chunk_size(self) -> int:
+        MAX_GRID_Y = 65535
+        MAX_TILE_TOKENS_DIM = 128
+
+        def _calc_max_supported_tokens(top_k: int, num_experts: int) -> int:
+            """Calculates the max number of supported tokens, so the CUDA grid.Y limit
+            won't be reached.
+            Based on getMaxNumCtasInBatchDim function in flashinfer's TRTLLM MoE runner:
+            https://github.com/flashinfer-ai/flashinfer/blob/719ee23fd82cb220d51ad118ca60198718f6c9d1/include/flashinfer/trtllm/fused_moe/runner.h#L97
+            Which given numTokens, topK, numExperts, tileTokensDim calculates maxNumCtas
+            which is used as the CUDA grid.Y dimension, which we want to
+            be <= MAX_GRID_Y. Solving for numTokens gives the formula below.
+            """
+            return (
+                num_experts + (MAX_GRID_Y - num_experts + 1) * MAX_TILE_TOKENS_DIM - 1
+            ) // top_k
+
+        # Using 305k or more causes IMA error in the kernel, so limit to 300k.
+        return min(
+            300000, _calc_max_supported_tokens(self.topk, self.moe_config.num_experts)
+        )
 
     def supports_expert_map(self) -> bool:
         return False
@@ -199,7 +218,7 @@ class TrtLlmNvFp4ExpertsModular(TrtLlmNvFp4ExpertsBase, mk.FusedMoEExpertsModula
     def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce:
         return TopKWeightAndReduceNoOP()
 
-    def apply(
+    def _invoke_kernel(
         self,
         output: torch.Tensor,
         hidden_states: torch.Tensor,
@@ -209,18 +228,10 @@ class TrtLlmNvFp4ExpertsModular(TrtLlmNvFp4ExpertsBase, mk.FusedMoEExpertsModula
         topk_ids: torch.Tensor,
         activation: MoEActivation,
         global_num_experts: int,
-        expert_map: torch.Tensor | None,
-        a1q_scale: torch.Tensor | None,
-        a2_scale: torch.Tensor | None,
-        workspace13: torch.Tensor,
-        workspace2: torch.Tensor,
-        expert_tokens_meta: mk.ExpertTokensMetadata | None,
-        apply_router_weight_on_input: bool,
+        a1q_scale: torch.Tensor,
     ):
         import flashinfer
 
-        assert self._supports_activation(activation)
-        assert a1q_scale is not None
         assert self.quant_config.w1_scale is not None
         assert self.quant_config.w2_scale is not None
 
@@ -262,6 +273,57 @@ class TrtLlmNvFp4ExpertsModular(TrtLlmNvFp4ExpertsBase, mk.FusedMoEExpertsModula
             output=output,
         )
 
+    def apply(
+        self,
+        output: torch.Tensor,
+        hidden_states: torch.Tensor,
+        w1: torch.Tensor,
+        w2: torch.Tensor,
+        topk_weights: torch.Tensor,
+        topk_ids: torch.Tensor,
+        activation: MoEActivation,
+        global_num_experts: int,
+        expert_map: torch.Tensor | None,
+        a1q_scale: torch.Tensor | None,
+        a2_scale: torch.Tensor | None,
+        workspace13: torch.Tensor,
+        workspace2: torch.Tensor,
+        expert_tokens_meta: mk.ExpertTokensMetadata | None,
+        apply_router_weight_on_input: bool,
+    ):
+        assert self._supports_activation(activation)
+        assert a1q_scale is not None
+
+        M = hidden_states.shape[0]
+        chunk_size = self._get_chunk_size()
+
+        if chunk_size >= M:
+            self._invoke_kernel(
+                output,
+                hidden_states,
+                w1,
+                w2,
+                topk_weights,
+                topk_ids,
+                activation,
+                global_num_experts,
+                a1q_scale,
+            )
+        else:
+            for start in range(0, M, chunk_size):
+                end = min(start + chunk_size, M)
+                self._invoke_kernel(
+                    output[start:end],
+                    hidden_states[start:end],
+                    w1,
+                    w2,
+                    topk_weights[start:end],
+                    topk_ids[start:end],
+                    activation,
+                    global_num_experts,
+                    a1q_scale[start:end],
+                )
+
 
 class TrtLlmNvFp4ExpertsMonolithic(
     TrtLlmNvFp4ExpertsBase, mk.FusedMoEExpertsMonolithic

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/mhc.py

@@ -3,8 +3,12 @@
 import torch
 
 # this import will also register the custom ops
+# import vllm.model_executor.kernels.mhc  # noqa: F401
 import vllm.model_executor.kernels.mhc as mhc_kernels
 from vllm.model_executor.custom_op import CustomOp
+from vllm.utils.import_utils import has_tilelang
+
+HAS_TILELANG = has_tilelang()
 
 
 # --8<-- [start:mhc_pre]
@@ -85,6 +89,52 @@ class MHCPreOp(CustomOp):
         #         sinkhorn_repeat,
         #     )
         # else:
+        if HAS_TILELANG:
+            return torch.ops.vllm.mhc_pre_tilelang(
+                residual,
+                fn,
+                hc_scale,
+                hc_base,
+                rms_eps,
+                hc_pre_eps,
+                hc_sinkhorn_eps,
+                hc_post_mult_value,
+                sinkhorn_repeat,
+                n_splits,
+                norm_weight,
+                norm_eps,
+            )
+        else:
+            return self.forward_native(
+                residual,
+                fn,
+                hc_scale,
+                hc_base,
+                rms_eps,
+                hc_pre_eps,
+                hc_sinkhorn_eps,
+                hc_post_mult_value,
+                sinkhorn_repeat,
+                n_splits,
+                norm_weight,
+                norm_eps,
+            )
+
+    def forward_native(
+        self,
+        residual: torch.Tensor,
+        fn: torch.Tensor,
+        hc_scale: torch.Tensor,
+        hc_base: torch.Tensor,
+        rms_eps: float,
+        hc_pre_eps: float,
+        hc_sinkhorn_eps: float,
+        hc_post_mult_value: float,
+        sinkhorn_repeat: int,
+        n_splits: int = 1,
+        norm_weight: torch.Tensor | None = None,
+        norm_eps: float = 0.0,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
         return mhc_kernels.mhc_pre_torch(
             residual,
             fn,
@@ -97,9 +147,6 @@ class MHCPreOp(CustomOp):
             sinkhorn_repeat,
         )
 
-    def forward_native(self, *args, **kwargs):
-        raise NotImplementedError("Native implementation of mhc_pre is not available")
-
 
 # --8<-- [start:mhc_post]
 @CustomOp.register("mhc_post")
@@ -147,6 +194,20 @@ class MHCPostOp(CustomOp):
         #         comb_res_mix,
         #     )
         # else:
+        if HAS_TILELANG:
+            return torch.ops.vllm.mhc_post_tilelang(
+                x, residual, post_layer_mix, comb_res_mix
+            )
+        else:
+            return self.forward_native(x, residual, post_layer_mix, comb_res_mix)
+
+    def forward_native(
+        self,
+        x: torch.Tensor,
+        residual: torch.Tensor,
+        post_layer_mix: torch.Tensor,
+        comb_res_mix: torch.Tensor,
+    ) -> torch.Tensor:
         return mhc_kernels.mhc_post_torch(
             x,
             residual,
@@ -154,9 +215,6 @@ class MHCPostOp(CustomOp):
             comb_res_mix,
         )
 
-    def forward_native(self, *args, **kwargs):
-        raise NotImplementedError("Native implementation of mhc_post is not available")
-
 
 # --8<-- [start:hc_head]
 @CustomOp.register("hc_head")
@@ -220,17 +278,32 @@ class HCHeadOp(CustomOp):
         out = torch.empty(
             num_tokens, hidden_size, dtype=torch.bfloat16, device=hidden_states.device
         )
-        torch.ops.vllm.hc_head_triton(
-            hs_flat,
-            hc_fn,
-            hc_scale,
-            hc_base,
-            out,
-            hidden_size,
-            rms_norm_eps,
-            hc_eps,
-            hc_mult,
-        )
+
+        if HAS_TILELANG:
+            torch.ops.vllm.hc_head_fused_kernel_tilelang(
+                hs_flat,
+                hc_fn,
+                hc_scale,
+                hc_base,
+                out,
+                hidden_size,
+                rms_norm_eps,
+                hc_eps,
+                hc_mult,
+            )
+        else:
+            torch.ops.vllm.hc_head_triton(
+                hs_flat,
+                hc_fn,
+                hc_scale,
+                hc_base,
+                out,
+                hidden_size,
+                rms_norm_eps,
+                hc_eps,
+                hc_mult,
+            )
+
         return out.view(*outer_shape, hidden_size)
 
     def forward_native(self, *args, **kwargs):
@@ -290,9 +363,42 @@ class MHCFusedPostPreOp(CustomOp):
             norm_eps,
         )
 
-    def forward_hip(self, *args, **kwargs):
-        raise NotImplementedError(
-            "Hip implementation of mhc_fused_post_pre is not available"
+    def forward_hip(
+        self,
+        x: torch.Tensor,
+        residual: torch.Tensor,
+        post_layer_mix: torch.Tensor,
+        comb_res_mix: torch.Tensor,
+        fn: torch.Tensor,
+        hc_scale: torch.Tensor,
+        hc_base: torch.Tensor,
+        rms_eps: float,
+        hc_pre_eps: float,
+        hc_sinkhorn_eps: float,
+        hc_post_mult_value: float,
+        sinkhorn_repeat: int,
+        n_splits: int = 1,
+        tile_n: int = 1,
+        norm_weight: torch.Tensor | None = None,
+        norm_eps: float = 0.0,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+        return torch.ops.vllm.mhc_fused_post_pre_tilelang(
+            x,
+            residual,
+            post_layer_mix,
+            comb_res_mix,
+            fn,
+            hc_scale,
+            hc_base,
+            rms_eps,
+            hc_pre_eps,
+            hc_sinkhorn_eps,
+            hc_post_mult_value,
+            sinkhorn_repeat,
+            n_splits,
+            tile_n,
+            norm_weight,
+            norm_eps,
         )
 
     def forward_native(self, *args, **kwargs):

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/model_executor/model_loader/weight_utils.py

@@ -1090,7 +1090,8 @@ def runai_safetensors_weights_iterator(
             mininterval=2,
         )
 
-        yield from tensor_iter
+        for name, tensor in tensor_iter:
+            yield name, tensor.clone()
 
 
 def _init_fastsafetensors_loader(

vllm/model_executor/models/mellum.py

@@ -0,0 +1,253 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from typing import Any
+
+from torch import nn
+
+from vllm.compilation.decorators import support_torch_compile
+from vllm.config import VllmConfig
+from vllm.distributed import get_tensor_model_parallel_world_size
+from vllm.model_executor.layers.attention import Attention
+from vllm.model_executor.layers.layernorm import RMSNorm
+from vllm.model_executor.layers.linear import QKVParallelLinear, RowParallelLinear
+from vllm.model_executor.layers.logits_processor import LogitsProcessor
+from vllm.model_executor.layers.rotary_embedding import get_rope
+from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead
+
+from .qwen3_moe import (
+    Qwen3MoeAttention,
+    Qwen3MoeDecoderLayer,
+    Qwen3MoeForCausalLM,
+    Qwen3MoeMLP,
+    Qwen3MoeModel,
+    Qwen3MoeSparseMoeBlock,
+)
+from .utils import PPMissingLayer, extract_layer_index, maybe_prefix
+
+
+class MellumAttention(Qwen3MoeAttention):
+    """
+    Differences from `Qwen3MoeAttention`:
+    - Supports `per_layer_sliding_window` for `Attention`.
+    """
+
+    def __init__(
+        self,
+        hidden_size: int,
+        num_heads: int,
+        num_kv_heads: int,
+        rope_parameters: dict[str, Any],
+        max_position_embeddings: int = 8192,
+        head_dim: int | None = None,
+        rms_norm_eps: float = 1e-06,
+        qkv_bias: bool = False,
+        cache_config: Any | None = None,
+        quant_config: Any | None = None,
+        prefix: str = "",
+        dual_chunk_attention_config: dict[str, Any] | None = None,
+        per_layer_sliding_window: int | None = None,
+    ) -> None:
+        nn.Module.__init__(self)
+
+        self.hidden_size = hidden_size
+        tp_size = get_tensor_model_parallel_world_size()
+        self.total_num_heads = num_heads
+        assert self.total_num_heads % tp_size == 0
+        self.num_heads = self.total_num_heads // tp_size
+        self.total_num_kv_heads = num_kv_heads
+        if self.total_num_kv_heads >= tp_size:
+            assert self.total_num_kv_heads % tp_size == 0
+        else:
+            assert tp_size % self.total_num_kv_heads == 0
+        self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size)
+        self.head_dim = head_dim or (hidden_size // self.total_num_heads)
+        self.q_size = self.num_heads * self.head_dim
+        self.kv_size = self.num_kv_heads * self.head_dim
+        self.scaling = self.head_dim**-0.5
+        self.max_position_embeddings = max_position_embeddings
+        self.dual_chunk_attention_config = dual_chunk_attention_config
+
+        self.qkv_proj = QKVParallelLinear(
+            hidden_size,
+            self.head_dim,
+            self.total_num_heads,
+            self.total_num_kv_heads,
+            bias=qkv_bias,
+            quant_config=quant_config,
+            prefix=f"{prefix}.qkv_proj",
+        )
+
+        self.o_proj = RowParallelLinear(
+            self.total_num_heads * self.head_dim,
+            hidden_size,
+            bias=False,
+            quant_config=quant_config,
+            prefix=f"{prefix}.o_proj",
+        )
+
+        self.rotary_emb = get_rope(
+            self.head_dim,
+            max_position=max_position_embeddings,
+            rope_parameters=rope_parameters,
+            dual_chunk_attention_config=dual_chunk_attention_config,
+        )
+        self.attn = Attention(
+            self.num_heads,
+            self.head_dim,
+            self.scaling,
+            num_kv_heads=self.num_kv_heads,
+            cache_config=cache_config,
+            quant_config=quant_config,
+            per_layer_sliding_window=per_layer_sliding_window,
+            prefix=f"{prefix}.attn",
+            **(
+                {
+                    "layer_idx": extract_layer_index(prefix),
+                    "dual_chunk_attention_config": dual_chunk_attention_config,
+                }
+                if dual_chunk_attention_config
+                else {}
+            ),
+        )
+
+        self.q_norm = RMSNorm(self.head_dim, eps=rms_norm_eps)
+        self.k_norm = RMSNorm(self.head_dim, eps=rms_norm_eps)
+
+
+class MellumDecoderLayer(Qwen3MoeDecoderLayer):
+    """
+    Differences from `Qwen3MoeDecoderLayer`:
+    - Supports interleaved SWA and per-layer RoPE scaling.
+    """
+
+    def __init__(self, vllm_config: VllmConfig, prefix: str = "") -> None:
+        nn.Module.__init__(self)
+
+        config = vllm_config.model_config.hf_text_config
+        cache_config = vllm_config.cache_config
+        quant_config = vllm_config.quant_config
+
+        self.hidden_size = config.hidden_size
+        max_position_embeddings = getattr(config, "max_position_embeddings", 8192)
+        dual_chunk_attention_config = getattr(
+            config, "dual_chunk_attention_config", None
+        )
+
+        layer_idx = extract_layer_index(prefix)
+        layer_type = config.layer_types[layer_idx]
+        if layer_type == "sliding_attention":
+            sliding_window = getattr(config, "sliding_window", None)
+        else:
+            sliding_window = None
+        rope_parameters = config.rope_parameters[layer_type]
+
+        self.self_attn = MellumAttention(
+            hidden_size=self.hidden_size,
+            num_heads=config.num_attention_heads,
+            num_kv_heads=config.num_key_value_heads,
+            rope_parameters=rope_parameters,
+            max_position_embeddings=max_position_embeddings,
+            rms_norm_eps=config.rms_norm_eps,
+            qkv_bias=getattr(config, "attention_bias", False),
+            head_dim=getattr(config, "head_dim", None),
+            cache_config=cache_config,
+            quant_config=quant_config,
+            prefix=f"{prefix}.self_attn",
+            dual_chunk_attention_config=dual_chunk_attention_config,
+            per_layer_sliding_window=sliding_window,
+        )
+
+        if config.mlp_layer_types[layer_idx] == "sparse":
+            self.mlp = Qwen3MoeSparseMoeBlock(
+                vllm_config=vllm_config, prefix=f"{prefix}.mlp"
+            )
+        else:
+            self.mlp = Qwen3MoeMLP(
+                hidden_size=config.hidden_size,
+                intermediate_size=config.intermediate_size,
+                hidden_act=config.hidden_act,
+                quant_config=quant_config,
+                prefix=f"{prefix}.mlp",
+            )
+
+        self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.post_attention_layernorm = RMSNorm(
+            config.hidden_size, eps=config.rms_norm_eps
+        )
+
+
+@support_torch_compile
+class MellumModel(Qwen3MoeModel):
+    """
+    Differences from `Qwen3MoeModel`:
+    - Uses `MellumDecoderLayer`.
+    """
+
+    def __init__(
+        self,
+        *,
+        vllm_config: VllmConfig,
+        prefix: str = "",
+    ):
+        super().__init__(
+            vllm_config=vllm_config,
+            prefix=prefix,
+            decoder_layer_type=MellumDecoderLayer,
+        )
+
+
+class MellumForCausalLM(Qwen3MoeForCausalLM):
+    """
+    Differences from `Qwen3MoeForCausalLM`:
+    - Uses `MellumModel`.
+    """
+
+    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
+        nn.Module.__init__(self)
+        config = vllm_config.model_config.hf_text_config
+        quant_config = vllm_config.quant_config
+        self.config = config
+        self.quant_config = quant_config
+        if "dense" in getattr(config, "mlp_layer_types", []):
+            self.packed_modules_mapping["gate_up_proj"] = ["gate_proj", "up_proj"]
+        self.model = MellumModel(
+            vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
+        )
+        self.lm_head = ParallelLMHead(
+            config.vocab_size,
+            config.hidden_size,
+            quant_config=quant_config,
+            prefix=maybe_prefix(prefix, "lm_head"),
+        )
+        if self.config.tie_word_embeddings:
+            self.lm_head.weight = self.model.embed_tokens.weight
+        self.logits_processor = LogitsProcessor(config.vocab_size)
+        self.make_empty_intermediate_tensors = (
+            self.model.make_empty_intermediate_tensors
+        )
+
+        self.expert_weights = []
+
+        self.moe_layers = []
+        example_layer = None
+        for layer in self.model.layers:
+            if isinstance(layer, PPMissingLayer):
+                continue
+
+            assert isinstance(layer, Qwen3MoeDecoderLayer)
+            if isinstance(layer.mlp, Qwen3MoeSparseMoeBlock):
+                example_layer = layer.mlp
+                self.moe_layers.append(layer.mlp.experts)
+
+        if example_layer is None:
+            raise RuntimeError("No MoE layer found in the model.layers.")
+
+        self.num_moe_layers = len(self.moe_layers)
+        self.num_expert_groups = 1
+        self.num_shared_experts = 0
+        self.num_logical_experts = example_layer.n_logical_experts
+        self.num_physical_experts = example_layer.n_physical_experts
+        self.num_local_physical_experts = example_layer.n_local_physical_experts
+        self.num_routed_experts = example_layer.n_routed_experts
+        self.num_redundant_experts = example_layer.n_redundant_experts

vllm/model_executor/models/registry.py

@@ -160,6 +160,7 @@ _TEXT_GENERATION_MODELS = {
     "LongcatFlashForCausalLM": ("longcat_flash", "LongcatFlashForCausalLM"),
     "MambaForCausalLM": ("mamba", "MambaForCausalLM"),
     "Mamba2ForCausalLM": ("mamba2", "Mamba2ForCausalLM"),
+    "MellumForCausalLM": ("mellum", "MellumForCausalLM"),
     "MiniCPMForCausalLM": ("minicpm", "MiniCPMForCausalLM"),
     "MiniCPM3ForCausalLM": ("minicpm3", "MiniCPM3ForCausalLM"),
     "MiniMaxForCausalLM": ("minimax_text_01", "MiniMaxText01ForCausalLM"),

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,14 @@ 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
+from vllm.utils.import_utils import has_tilelang
 
 
 class DeepseekV4MLP(nn.Module):
@@ -115,501 +108,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 +120,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 +130,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 +142,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 +173,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 +208,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 +226,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__(
@@ -1074,6 +474,7 @@ class DeepseekV4DecoderLayer(nn.Module):
         self.mhc_pre = MHCPreOp()
         self.mhc_post = MHCPostOp()
         self.mhc_fused_post_pre = MHCFusedPostPreOp()
+        self.has_tilelang = has_tilelang()
 
     def hc_pre(
         self,
@@ -1104,7 +505,7 @@ class DeepseekV4DecoderLayer(nn.Module):
     ):
         return self.mhc_post(x, residual, post, comb)
 
-    def _forward_cuda(
+    def _forward_fused_post_pre(
         self,
         x: torch.Tensor,
         positions: torch.Tensor,
@@ -1156,7 +557,7 @@ class DeepseekV4DecoderLayer(nn.Module):
         x = self.ffn(x, input_ids)
         return x, residual, post_mix, res_mix
 
-    def _forward_rocm(
+    def _forward_unfused_post_pre(
         self,
         x: torch.Tensor,
         positions: torch.Tensor,
@@ -1195,12 +596,13 @@ class DeepseekV4DecoderLayer(nn.Module):
     ) -> tuple[
         torch.Tensor, torch.Tensor | None, torch.Tensor | None, torch.Tensor | None
     ]:
-        if current_platform.is_rocm():
-            return self._forward_rocm(
+        if not self.has_tilelang:
+            return self._forward_unfused_post_pre(
                 x, positions, input_ids, post_mix, res_mix, residual
             )
-
-        return self._forward_cuda(x, positions, input_ids, post_mix, res_mix, residual)
+        return self._forward_fused_post_pre(
+            x, positions, input_ids, post_mix, res_mix, residual
+        )
 
 
 @support_torch_compile
@@ -1211,15 +613,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
@@ -1292,6 +685,7 @@ class DeepseekV4Model(nn.Module):
             requires_grad=False,
         )
         self.hc_head_op = HCHeadOp()
+        self.has_tilelang = has_tilelang()
         # Pre-hc_head residual stream buffer for the MTP draft. Stable
         # address (outside the cudagraph pool) so the copy_ in forward()
         # refreshes it correctly across captured shapes.
@@ -1348,9 +742,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(
@@ -1361,7 +752,7 @@ class DeepseekV4Model(nn.Module):
                 res_mix,
                 residual,
             )
-        if layer is not None and current_platform.is_cuda():
+        if layer is not None and self.has_tilelang:
             hidden_states = layer.hc_post(hidden_states, residual, post_mix, res_mix)
 
         if not get_pp_group().is_last_rank:
@@ -1482,9 +873,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 +883,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 +985,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

@@ -39,11 +39,9 @@ from vllm.model_executor.models.deepseek_v2 import get_spec_layer_idx_from_weigh
 from vllm.model_executor.models.utils import maybe_prefix
 from vllm.platforms import current_platform
 from vllm.sequence import IntermediateTensors
+from vllm.utils.import_utils import has_tilelang
 
-from .model import (
-    DeepseekV4DecoderLayer,
-    make_deepseek_v4_expert_params_mapping,
-)
+from .model import DeepseekV4DecoderLayer
 
 logger = init_logger(__name__)
 
@@ -121,6 +119,7 @@ class DeepSeekV4MultiTokenPredictorLayer(nn.Module):
         )
 
         self.hc_head_op = HCHeadOp()
+        self.has_tilelang = has_tilelang()
 
     def forward(
         self,
@@ -147,7 +146,7 @@ class DeepSeekV4MultiTokenPredictorLayer(nn.Module):
         hidden_states, residual, post_mix, res_mix = self.mtp_block(
             positions=positions, x=hidden_states, input_ids=None
         )
-        if current_platform.is_cuda():
+        if self.has_tilelang:
             hidden_states = self.mtp_block.hc_post(
                 hidden_states, residual, post_mix, res_mix
             )
@@ -330,19 +329,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 +458,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/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/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,46 @@ 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.sparse_attn_compress_cutedsl 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,7 +54,7 @@ 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,

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

@@ -5,4 +5,9 @@
 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.
+
+This ``__init__`` deliberately imports nothing: re-exporting the cutedsl
+modules here would eagerly ``import cutlass`` (initializing the CUDA driver) for
+anyone who imports ``vllm.models.deepseek_v4``, breaking forked subprocesses.
+Import the leaf modules directly under a ``has_cutedsl()``/``is_cuda()`` gate.
 """

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/platforms/cpu.py

@@ -14,7 +14,6 @@ from vllm.logger import init_logger
 from vllm.utils.cpu_resource_utils import (
     DEVICE_CONTROL_ENV_VAR,
     get_memory_node_info,
-    get_visible_memory_node,
 )
 from vllm.utils.mem_constants import GiB_bytes
 from vllm.v1.attention.backends.registry import AttentionBackendEnum
@@ -136,13 +135,9 @@ class CpuPlatform(Platform):
         scheduler_config.async_scheduling = False
 
         parallel_config = vllm_config.parallel_config
-        if (
-            os.environ.get("VLLM_ENABLE_V1_MULTIPROCESSING", "1") == "1"
-            and parallel_config.distributed_executor_backend == "uni"
-        ):
-            # OMP requires the MP executor to function correctly, UniProc
-            # is not supported as it is not possible to set the OMP
-            # environment correctly
+        # OMP requires the MP executor to function correctly, UniProc is not
+        # supported as it is not possible to set the OMP environment correctly
+        if parallel_config.distributed_executor_backend == "uni":
             parallel_config.distributed_executor_backend = "mp"
 
         if parallel_config.worker_cls == "auto":
@@ -486,15 +481,3 @@ class CpuPlatform(Platform):
             slot_mapping,
             isa,
         )
-
-    @classmethod
-    def get_current_memory_usage(
-        cls, device: torch.types.Device | None = None
-    ) -> float:
-        allowed_mem_node_list = get_visible_memory_node()
-        mem_status_list = [get_memory_node_info(i) for i in allowed_mem_node_list]
-        memory_usage = 0
-        for s in mem_status_list:
-            memory_usage += s.total_memory - s.available_memory
-
-        return memory_usage

vllm/platforms/cuda.py

@@ -592,6 +592,15 @@ class CudaPlatformBase(Platform):
             default, rms_norm=rms_norm, fused_add_rms_norm=rms_norm
         )
 
+    @classmethod
+    def is_arch_support_pdl(cls) -> bool:
+        try:
+            device = torch.cuda.current_device()
+            major, _ = torch.cuda.get_device_capability(device)
+        except Exception:
+            return False
+        return major >= 9
+
 
 # NVML utils
 # Note that NVML is not affected by `CUDA_VISIBLE_DEVICES`,

vllm/platforms/interface.py

@@ -1016,6 +1016,13 @@ class Platform:
         # Native always used by default. Platforms can override this behavior.
         return IrOpPriorityConfig.with_default(["native"])
 
+    @classmethod
+    def is_arch_support_pdl(cls) -> bool:
+        """
+        Does the current platform support PDL (Programmatic Dependent Launch)?
+        """
+        return False
+
 
 class UnspecifiedPlatform(Platform):
     _enum = PlatformEnum.UNSPECIFIED

vllm/transformers_utils/config.py

@@ -101,6 +101,7 @@ _CONFIG_REGISTRY: dict[str, type[PretrainedConfig]] = LazyConfigDict(
     fireredlid="FireRedLIDConfig",
     funaudiochat="FunAudioChatConfig",
     granite4_vision="Granite4VisionConfig",
+    hyperclovax="HyperCLOVAXConfig",
     hyperclovax_vlm="HCXVisionConfig",
     hunyuan_vl="HunYuanVLConfig",
     hy_v3="HYV3Config",
@@ -114,6 +115,7 @@ _CONFIG_REGISTRY: dict[str, type[PretrainedConfig]] = LazyConfigDict(
     jais="JAISConfig",
     mlp_speculator="MLPSpeculatorConfig",
     medusa="MedusaConfig",
+    mellum="MellumConfig",
     midashenglm="MiDashengLMConfig",
     moondream3="Moondream3Config",
     eagle="EAGLEConfig",
@@ -429,9 +431,9 @@ def patch_legacy_rope_type(rope_parameters: dict[str, Any] | None) -> None:
         if "rope_type" not in rope_parameters and "type" in rope_parameters:
             rope_parameters["rope_type"] = rope_parameters["type"]
             logger.info("Replacing legacy 'type' key with 'rope_type'")
-        # Case 3: No rope_type field at all - cannot determine RoPE type, raise error
+        # Case 3: No rope_type field present - nothing to patch
         if "rope_type" not in rope_parameters:
-            raise ValueError("rope_parameters should have a 'rope_type' key")
+            return
         # Patch legacy rope_type values with warning
         if rope_parameters["rope_type"] == "su":
             rope_parameters["rope_type"] = "longrope"

vllm/transformers_utils/configs/__init__.py

@@ -49,6 +49,7 @@ _CLASS_TO_MODULE: dict[str, str] = {
     "LagunaConfig": "vllm.transformers_utils.configs.laguna",
     "Lfm2MoeConfig": "vllm.transformers_utils.configs.lfm2_moe",
     "MedusaConfig": "vllm.transformers_utils.configs.medusa",
+    "MellumConfig": "vllm.transformers_utils.configs.mellum",
     "MiDashengLMConfig": "vllm.transformers_utils.configs.midashenglm",
     "MLPSpeculatorConfig": "vllm.transformers_utils.configs.mlp_speculator",
     "Moondream3Config": "vllm.transformers_utils.configs.moondream3",
@@ -117,6 +118,7 @@ __all__ = [
     "LagunaConfig",
     "Lfm2MoeConfig",
     "MedusaConfig",
+    "MellumConfig",
     "MiDashengLMConfig",
     "MLPSpeculatorConfig",
     "Moondream3Config",

vllm/transformers_utils/configs/mellum.py

@@ -0,0 +1,7 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from transformers import Qwen3MoeConfig
+
+
+class MellumConfig(Qwen3MoeConfig):
+    model_type = "mellum"

vllm/triton_utils/importing.py

@@ -3,7 +3,6 @@
 
 import os
 import types
-from importlib.metadata import version
 from importlib.util import find_spec
 
 from vllm.logger import init_logger
@@ -49,17 +48,6 @@ if HAS_TRITON:
                 len(active_drivers),
             )
             HAS_TRITON = False
-
-        # Check Triton CPU
-        if "cpu" in version("vllm"):
-            if "cpu" in backends:
-                HAS_TRITON = True
-            else:
-                logger.warning(
-                    "Triton is installed, but doesn't include CPU backend. "
-                    "Disabling Triton."
-                )
-                HAS_TRITON = False
     except ImportError:
         # This can occur if Triton is partially installed or triton.backends
         # is missing.

vllm/utils/import_utils.py

@@ -430,6 +430,7 @@ def has_triton_kernels() -> bool:
     return is_available
 
 
+@cache
 def has_tilelang() -> bool:
     """Whether the optional `tilelang` package is available."""
     return _has_module("tilelang")

vllm/utils/platform_utils.py

@@ -50,10 +50,8 @@ def is_pin_memory_available() -> bool:
 def is_uva_available() -> bool:
     """Check if Unified Virtual Addressing (UVA) is available."""
     # UVA requires pinned memory.
-    from vllm.platforms import current_platform
-
     # TODO: Add more requirements for UVA if needed.
-    return is_pin_memory_available() or current_platform.is_cpu()
+    return is_pin_memory_available()
 
 
 @cache

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 = envs.VLLM_ENGINE_READY_TIMEOUT_S,
+    ) -> 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/cpu/buffer_utils.py

@@ -1,16 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-from collections.abc import Sequence
-
-import torch
-
-from vllm.utils.platform_utils import is_uva_available
-
-
-class UvaBuffer:
-    def __init__(self, size: int | Sequence[int], dtype: torch.dtype):
-        if not is_uva_available():
-            raise RuntimeError("UVA is not available")
-        self.cpu = torch.zeros(size, dtype=dtype, device="cpu")
-        self.np = self.cpu.numpy()
-        self.uva = self.cpu

vllm/v1/worker/cpu/model_runner.py

@@ -1,16 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-from vllm.logger import init_logger
-from vllm.v1.worker.gpu.model_runner import GPUModelRunner
-
-logger = init_logger(__name__)
-
-
-class CPUModelRunner(GPUModelRunner):
-    # TBD: Whether need to move this to Worker?
-    def warming_up_model(self) -> None:
-        logger.info("Warming up model for the compilation...")
-        # Only generate graph for the generic shape
-        self.profile_run()
-        logger.info("Warming up done.")

vllm/v1/worker/cpu/shm.py

@@ -1,62 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-# isort: skip_file
-# ruff: noqa: E402
-# mypy: disable-error-code="misc, assignment"
-
-from typing import Any
-
-# Patch torch APIs
-import torch
-
-
-def noop(*args: Any, **kwargs: Any) -> None:
-    pass
-
-
-class _EventPlaceholder:
-    def __init__(self, *args, **kwargs) -> None:
-        self.record = noop
-        self.synchronize = noop
-
-
-class _StreamPlaceholder:
-    def __init__(self, *args, **kwargs) -> None:
-        self.wait_stream = noop
-
-    def __enter__(self, *args, **kwargs):
-        return self
-
-    def __exit__(self, exc_type, exc_val, exc_tb):
-        pass
-
-
-torch.Event = _EventPlaceholder
-torch.cuda.Event = _EventPlaceholder
-torch.cuda.Stream = _StreamPlaceholder
-torch.cuda.set_stream = noop
-torch.cuda.current_stream = lambda *args, **kwargs: _StreamPlaceholder()
-torch.accelerator.synchronize = noop
-torch.accelerator.empty_cache = noop
-
-# Patch vLLM torch utils
-import vllm.utils.torch_utils as torch_utils
-
-
-def async_tensor_h2d(
-    data: list,
-    dtype: torch.dtype,
-    device: str | torch.device,
-    pin_memory: bool = False,
-) -> torch.Tensor:
-    return torch.tensor(data, dtype=dtype, device="cpu")
-
-
-torch_utils.async_tensor_h2d = async_tensor_h2d
-
-# Patch model runner APIs
-import vllm.v1.worker.gpu.buffer_utils as gpu_buffer_utils
-import vllm.v1.worker.cpu.buffer_utils as cpu_buffer_utils
-
-gpu_buffer_utils.UvaBuffer = cpu_buffer_utils.UvaBuffer

vllm/v1/worker/cpu_worker.py

@@ -1,9 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-# Must be imported firstly
-import vllm.v1.worker.cpu.shm  # noqa # isort: skip
-
 import math
 import os
 import sys
@@ -105,8 +101,6 @@ class CPUWorker(Worker):
             )
 
     def init_device(self):
-        self.device = torch.device("cpu")
-
         # Check whether critical libraries are loaded
         def check_preloaded_libs(name: str):
             ld_preload_list = os.environ.get("LD_PRELOAD", "")
@@ -147,16 +141,9 @@ class CPUWorker(Worker):
         set_random_seed(self.model_config.seed)
 
         # Construct the model runner
-        if self.use_v2_model_runner:
-            from vllm.v1.worker.cpu.model_runner import (
-                CPUModelRunner as CPUModelRunnerV2,
-            )
-
-            self.model_runner: CPUModelRunner = CPUModelRunnerV2(  # type: ignore
-                self.vllm_config, self.device
-            )
-        else:
-            self.model_runner = CPUModelRunner(self.vllm_config, torch.device("cpu"))
+        self.model_runner: CPUModelRunner = CPUModelRunner(
+            self.vllm_config, torch.device("cpu")
+        )
 
     def sleep(self, level: int = 1) -> None:
         logger.warning("sleep mode is not supported on CPU, ignore it.")

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]