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

benchmarks/kernels/benchmark_fused_moe_lora_one_shot.py

@@ -0,0 +1,465 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Benchmark the fused MoE-LoRA fast path (one-shot) vs two-kernel baseline.
+
+The "one_shot" provider goes through `vllm.lora.ops.triton_ops.fused_moe_lora`
+which dispatches to the single-kernel one-shot implementation when
+fully_sharded=False (the prefill default).
+
+The "two_kernel" provider drives `fused_moe_lora_shrink` + `fused_moe_lora_expand`
+directly, bypassing the dispatch and matching the legacy two-kernel path's
+work distribution. This isolates the win from kernel fusion.
+
+Run:
+    .venv/bin/python -m benchmarks.kernels.benchmark_fused_moe_lora_one_shot
+    .venv/bin/python -m benchmarks.kernels.benchmark_fused_moe_lora_one_shot \\
+        --model qwen3moe
+"""
+
+from __future__ import annotations
+
+import argparse
+import os
+import random
+
+import torch
+
+from vllm import _custom_ops as ops
+from vllm.lora.ops.triton_ops import (
+    fused_moe_lora,
+    fused_moe_lora_expand,
+    fused_moe_lora_shrink,
+)
+from vllm.triton_utils import triton
+
+DTYPE = torch.bfloat16
+DEVICE = "cuda"
+
+
+# ----- input fabrication -----------------------------------------------------
+
+
+def _round_up(x: int, base: int) -> int:
+    return ((x + base - 1) // base) * base
+
+
+def _ceildiv(x: int, y: int) -> int:
+    return (x + y - 1) // y
+
+
+def _assign_loras(num_tokens: int, num_sequences: int, max_loras: int) -> torch.Tensor:
+    tokens_per_seq = num_tokens // num_sequences
+    rem = num_tokens % num_sequences
+    out = torch.empty(num_tokens, dtype=torch.int32)
+    start = 0
+    for i in range(num_sequences):
+        end = start + tokens_per_seq + (1 if i < rem else 0)
+        out[start:end] = random.randint(0, max_loras - 1)
+        start = end
+    return out
+
+
+def _assign_experts(num_tokens: int, num_experts: int, top_k: int):
+    expert_indices = torch.empty((num_tokens, top_k), dtype=torch.int32)
+    for i in range(num_tokens):
+        expert_indices[i] = torch.randperm(num_experts)[:top_k]
+    weights = torch.rand((num_tokens, top_k), dtype=torch.float32)
+    weights = weights / weights.sum(dim=1, keepdim=True)
+    return expert_indices, weights
+
+
+def _make_inputs(
+    M: int,
+    K: int,
+    N_per_slice: int,
+    rank: int,
+    num_experts: int,
+    top_k: int,
+    max_loras: int,
+    num_slices: int,
+    block_size_m: int,
+):
+    """Mirrors the production caller's tensor layout."""
+    torch.manual_seed(0)
+    random.seed(0)
+
+    num_sequences = max(1, min(M, 8))
+    topk_ids_cpu, topk_weights_cpu = _assign_experts(M, num_experts, top_k)
+    token_lora_cpu = _assign_loras(M, num_sequences, max_loras)
+    lora_ids_cpu = torch.full((max_loras + 1,), -1, dtype=torch.int32)
+    uniq = torch.unique(token_lora_cpu, sorted=True)
+    lora_ids_cpu[: uniq.size(0)].copy_(uniq)
+
+    topk_ids = topk_ids_cpu.to(DEVICE)
+    topk_weights = topk_weights_cpu.to(device=DEVICE, dtype=DTYPE)
+    token_lora_mapping = token_lora_cpu.to(DEVICE)
+    lora_ids = lora_ids_cpu.to(DEVICE)
+    adapter_enabled = torch.ones(max_loras + 1, dtype=torch.int32, device=DEVICE)
+
+    lora_a = [
+        torch.randn((max_loras, num_experts, rank, K), dtype=DTYPE, device=DEVICE)
+        / max(K, 1) ** 0.5
+        for _ in range(num_slices)
+    ]
+    lora_b = [
+        torch.randn(
+            (max_loras, num_experts, N_per_slice, rank),
+            dtype=DTYPE,
+            device=DEVICE,
+        )
+        / max(rank, 1) ** 0.5
+        for _ in range(num_slices)
+    ]
+    hidden = torch.randn((M, K), dtype=DTYPE, device=DEVICE)
+    out_template = torch.zeros(
+        (M, top_k, num_slices * N_per_slice), dtype=DTYPE, device=DEVICE
+    )
+
+    # Sorted-path metadata (the prefill default).
+    max_pad = topk_ids.numel() + num_experts * (block_size_m - 1)
+    max_pad = _round_up(max_pad, block_size_m)
+    max_blocks = _ceildiv(max_pad, block_size_m)
+    sorted_token_ids = torch.empty(
+        (max_loras * max_pad,), dtype=torch.int32, device=DEVICE
+    )
+    expert_ids = torch.empty(
+        (max_loras * max_blocks,), dtype=torch.int32, device=DEVICE
+    )
+    num_post = torch.empty((max_loras,), dtype=torch.int32, device=DEVICE)
+    ops.moe_lora_align_block_size(
+        topk_ids,
+        token_lora_mapping,
+        num_experts,
+        block_size_m,
+        max_loras,
+        max_pad,
+        max_blocks,
+        sorted_token_ids,
+        expert_ids,
+        num_post,
+        adapter_enabled,
+        lora_ids,
+    )
+    expert_ids = expert_ids.view(max_loras, -1).contiguous()
+    sorted_token_ids = sorted_token_ids.view(max_loras, -1).contiguous()
+    num_active = torch.tensor([max_loras + 1], dtype=torch.int32, device="cpu")
+
+    return dict(
+        hidden=hidden,
+        lora_a=lora_a,
+        lora_b=lora_b,
+        topk_weights=topk_weights,
+        sorted_token_ids=sorted_token_ids,
+        expert_ids=expert_ids,
+        num_post=num_post,
+        token_lora_mapping=token_lora_mapping,
+        lora_ids=lora_ids,
+        num_active=num_active,
+        adapter_enabled=adapter_enabled,
+        out_template=out_template,
+        # bookkeeping
+        M=M,
+        K=K,
+        N_per_slice=N_per_slice,
+        rank=rank,
+        num_experts=num_experts,
+        top_k=top_k,
+        max_loras=max_loras,
+        num_slices=num_slices,
+        block_size_m=block_size_m,
+    )
+
+
+# ----- providers -------------------------------------------------------------
+
+
+def _run_one_shot(inp: dict):
+    """Drive `fused_moe_lora` with fully_sharded=False -> one-shot fast path."""
+    out = inp["out_template"].clone()
+    fused_moe_lora(
+        out,
+        inp["hidden"],
+        inp["lora_a"],
+        inp["lora_b"],
+        inp["topk_weights"],
+        inp["sorted_token_ids"],
+        inp["expert_ids"],
+        inp["num_post"],
+        inp["token_lora_mapping"],
+        inp["rank"],
+        inp["top_k"],
+        inp["lora_ids"],
+        inp["num_active"],
+        inp["adapter_enabled"],
+        inp["block_size_m"],
+        64,
+        32,
+        8,
+        4,
+        3,
+        1,
+        inp["block_size_m"],
+        64,
+        32,
+        8,
+        4,
+        3,
+        1,
+        False,
+        False,
+        0,
+    )
+    return out
+
+
+def _run_two_kernel(inp: dict):
+    """Drive `fused_moe_lora_shrink` + `fused_moe_lora_expand` directly,
+    bypassing the dispatch. Matches the legacy two-kernel work distribution.
+    """
+    M = inp["M"]
+    top_k = inp["top_k"]
+    rank = inp["rank"]
+    num_slices = inp["num_slices"]
+    N_per_slice = inp["N_per_slice"]
+    K = inp["K"]
+    num_experts = inp["num_experts"]
+    block_m = inp["block_size_m"]
+
+    intermediate = torch.zeros((num_slices, M, top_k, rank), dtype=DTYPE, device=DEVICE)
+    out = inp["out_template"].clone()
+    EM = inp["sorted_token_ids"].shape[1]
+    num_tokens = M * top_k
+
+    fused_moe_lora_shrink(
+        intermediate,
+        inp["hidden"],
+        inp["lora_a"],
+        inp["topk_weights"],
+        inp["sorted_token_ids"],
+        inp["expert_ids"],
+        inp["num_post"],
+        inp["token_lora_mapping"],
+        top_k,
+        inp["lora_ids"],
+        inp["adapter_enabled"],
+        torch.device(DEVICE),
+        rank,
+        M,
+        EM,
+        K,
+        num_tokens,
+        num_experts,
+        num_slices,
+        block_m,
+        64,
+        32,
+        8,
+        4,
+        3,
+        1,
+        inp["num_active"],
+        False,
+    )
+    fused_moe_lora_expand(
+        out,
+        intermediate,
+        inp["lora_b"],
+        inp["topk_weights"],
+        inp["sorted_token_ids"],
+        inp["expert_ids"],
+        inp["num_post"],
+        inp["token_lora_mapping"],
+        top_k,
+        inp["lora_ids"],
+        inp["adapter_enabled"],
+        torch.device(DEVICE),
+        rank,
+        M,
+        EM,
+        K,
+        num_tokens,
+        num_experts,
+        num_slices,
+        rank,
+        N_per_slice,
+        block_m,
+        64,
+        32,
+        8,
+        4,
+        3,
+        1,
+        inp["num_active"],
+        False,
+        0,
+    )
+    return out
+
+
+PROVIDER_FNS = {
+    "one_shot": _run_one_shot,
+    "two_kernel": _run_two_kernel,
+}
+
+
+# ----- model presets ---------------------------------------------------------
+
+
+MODEL_PRESETS: dict[str, dict] = {
+    # Mixtral-8x7B style: E=8, top_k=2, hidden=4096, intermediate=14336
+    "mixtral": dict(
+        K=4096,
+        N_per_slice=7168,
+        num_experts=8,
+        top_k=2,
+        max_loras=4,
+        num_slices=2,
+        block_size_m=64,
+    ),
+    # Qwen3-MoE / DeepSeek-V2 style: E=64, top_k=8, hidden=2048, inter=1408
+    "qwen3moe": dict(
+        K=2048,
+        N_per_slice=1408,
+        num_experts=64,
+        top_k=8,
+        max_loras=4,
+        num_slices=2,
+        block_size_m=64,
+    ),
+    # GLM-5.1 (zai-org/GLM-5.1-FP8): E=256, top_k=8, hidden=6144,
+    # moe_intermediate=2048
+    "glm5_1": dict(
+        K=6144,
+        N_per_slice=2048,
+        num_experts=256,
+        top_k=8,
+        max_loras=4,
+        num_slices=2,
+        block_size_m=64,
+    ),
+}
+
+
+M_RANGE = [16, 64, 256, 1024, 4096, 16384]
+RANK_RANGE = [8, 16, 32, 64]
+
+
+def get_benchmark(model: str, max_loras: int | None = None):
+    preset = dict(MODEL_PRESETS[model])
+    if max_loras is not None:
+        preset["max_loras"] = max_loras
+
+    @triton.testing.perf_report(
+        triton.testing.Benchmark(
+            x_names=["M", "rank"],
+            x_vals=[(M, R) for M in M_RANGE for R in RANK_RANGE],
+            line_arg="provider",
+            line_vals=list(PROVIDER_FNS.keys()),
+            line_names=["one_shot (fused)", "two_kernel (legacy)"],
+            styles=[("red", "-"), ("blue", "-")],
+            ylabel="ms",
+            plot_name=f"fused_moe_lora-{model}-loras{preset['max_loras']}",
+            args={"preset": preset},
+        )
+    )
+    def benchmark(M, rank, provider, preset):
+        inp = _make_inputs(
+            M=M,
+            K=preset["K"],
+            N_per_slice=preset["N_per_slice"],
+            rank=rank,
+            num_experts=preset["num_experts"],
+            top_k=preset["top_k"],
+            max_loras=preset["max_loras"],
+            num_slices=preset["num_slices"],
+            block_size_m=preset["block_size_m"],
+        )
+        fn = PROVIDER_FNS[provider]
+        quantiles = [0.5, 0.2, 0.8]
+        ms, min_ms, max_ms = triton.testing.do_bench(
+            lambda: fn(inp), quantiles=quantiles
+        )
+        return ms, max_ms, min_ms
+
+    return benchmark
+
+
+# ----- correctness sanity ---------------------------------------------------
+
+
+def calculate_diff(model: str, M: int, rank: int, max_loras: int | None = None):
+    preset = dict(MODEL_PRESETS[model])
+    if max_loras is not None:
+        preset["max_loras"] = max_loras
+    inp = _make_inputs(
+        M=M,
+        K=preset["K"],
+        N_per_slice=preset["N_per_slice"],
+        rank=rank,
+        num_experts=preset["num_experts"],
+        top_k=preset["top_k"],
+        max_loras=preset["max_loras"],
+        num_slices=preset["num_slices"],
+        block_size_m=preset["block_size_m"],
+    )
+    out_one = _run_one_shot(inp)
+    out_two = _run_two_kernel(inp)
+    max_abs = (out_one.float() - out_two.float()).abs().max().item()
+    print(
+        f"  model={model:<9} M={M:<6} rank={rank:<3}  "
+        f"max|one_shot - two_kernel|={max_abs:.4g}  "
+        f"ref|max|={out_two.float().abs().max().item():.3g}"
+    )
+    if max_abs <= 5e-2:
+        print("  ✅ outputs match within bf16 tolerance")
+    else:
+        print("  ❌ outputs differ beyond expected bf16 noise")
+
+
+# ----- main ------------------------------------------------------------------
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--model",
+        type=str,
+        default="mixtral",
+        choices=list(MODEL_PRESETS.keys()),
+        help="Model preset to sweep",
+    )
+    parser.add_argument(
+        "--save-path",
+        type=str,
+        default="./configs/fused_moe_lora_one_shot/",
+        help="Directory to save benchmark results",
+    )
+    parser.add_argument(
+        "--check-only",
+        action="store_true",
+        help="Run correctness sanity check only, no perf sweep",
+    )
+    parser.add_argument(
+        "--max-loras",
+        type=int,
+        default=None,
+        help="Override max_loras in the model preset (number of LoRA adapters "
+        "active in the batch). Defaults to the preset's value.",
+    )
+    args = parser.parse_args()
+
+    print(f"Correctness check ({args.model}):")
+    calculate_diff(args.model, M=256, rank=32, max_loras=args.max_loras)
+    if args.check_only:
+        raise SystemExit(0)
+
+    effective_max_loras = (
+        args.max_loras
+        if args.max_loras is not None
+        else MODEL_PRESETS[args.model]["max_loras"]
+    )
+    print(f"\nGPU: {torch.cuda.get_device_name()}")
+    print(f"Model preset: {args.model}  max_loras={effective_max_loras}\n")
+    benchmark = get_benchmark(args.model, max_loras=args.max_loras)
+    os.makedirs(args.save_path, exist_ok=True)
+    benchmark.run(print_data=True, save_path=args.save_path)

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/llm/test_mm_processor_kwargs.py

@@ -192,7 +192,7 @@ def test_preprocess_cmpl_applies_mm_processor_kwargs_to_renderer(
     llm.renderer = renderer
 
     monkeypatch.setattr(
-        "vllm.entrypoints.llm.parse_model_prompt",
+        "vllm.entrypoints.offline_utils.parse_model_prompt",
         lambda _model_config, parsed_prompt: parsed_prompt,
     )
 
@@ -226,7 +226,7 @@ def test_preprocess_cmpl_keeps_prompt_mm_processor_kwargs_when_no_override(
     llm.renderer = renderer
 
     monkeypatch.setattr(
-        "vllm.entrypoints.llm.parse_model_prompt",
+        "vllm.entrypoints.offline_utils.parse_model_prompt",
         lambda _model_config, parsed_prompt: parsed_prompt,
     )
 

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/lora/test_fused_moe_lora_kernel.py

@@ -141,6 +141,7 @@ def use_fused_moe_lora_kernel(
     block_size,
     fully_sharded=False,
     offset=0,
+    add_inputs=True,
 ):
     max_num_tokens_padded = topk_ids.numel() + num_experts * (block_size - 1)
     max_num_tokens_padded = round_up(max_num_tokens_padded, block_size)
@@ -222,6 +223,7 @@ def use_fused_moe_lora_kernel(
         mul_routed_weight,
         fully_sharded=fully_sharded,
         offset=offset,
+        add_inputs=add_inputs,
     )
 
 
@@ -371,6 +373,7 @@ def use_fused_moe_lora_kernel_naive(
     block_size,
     fully_sharded=False,
     offset=0,
+    add_inputs=True,
 ):
     """
     Test helper for naive_block_assignment path.
@@ -435,6 +438,7 @@ def use_fused_moe_lora_kernel_naive(
         mul_routed_weight=mul_routed_weight,
         fully_sharded=fully_sharded,
         offset=offset,
+        add_inputs=add_inputs,
     )
 
 
@@ -745,3 +749,494 @@ def use_fused_moe_lora_kernel_tensor_parallel(
         output = tensor_model_parallel_all_reduce(output)
 
     torch.testing.assert_close(output, ref_output, atol=1e-2, rtol=1e-2)
+
+
+# -- one-shot fast-path coverage --------------------------------------------
+# The fused shrink+expand one-shot kernel pads `BLOCK_R` to next_pow2(rank),
+# with a floor of 16 (tensor-core minimum). Small ranks (4, 8) exercise the
+# rank-dim masking and are not covered by the original tests, which start at
+# rank=16. The legacy two-kernel path additionally fails on rank=4 in TMA
+# mode because the rank-dim stride (rank * elem_size) is not 16-byte
+# aligned; the one-shot fast path takes precedence whenever fully_sharded
+# is False so this regression is hidden in normal use, but the test still
+# ensures the one-shot logic is correct against the pytorch reference.
+
+
+@pytest.mark.parametrize("num_tokens", [16, 100])
+@pytest.mark.parametrize("top_k_num", [2])
+@pytest.mark.parametrize("num_experts", [8, 64])
+@pytest.mark.parametrize("max_loras", [4])
+@pytest.mark.parametrize("N", [1408])
+@pytest.mark.parametrize("K", [2048])
+@pytest.mark.parametrize("max_lora_rank", [4, 8])
+@pytest.mark.parametrize("block_size", [16, 64])
+@pytest.mark.parametrize("num_slices", [1, 2])
+@pytest.mark.parametrize("dtype", DTYPES)
+@pytest.mark.parametrize("device", DEVICES)
+@pytest.mark.parametrize("seed", SEED)
+def test_fused_moe_lora_kernel_small_rank(
+    num_tokens,
+    top_k_num,
+    num_experts,
+    max_loras,
+    N,
+    K,
+    max_lora_rank,
+    block_size,
+    num_slices,
+    dtype,
+    device,
+    seed,
+):
+    """One-shot fast path covering rank<16 (padded to BLOCK_R=16 inside kernel)."""
+    torch.set_default_device(device)
+    set_random_seed(seed)
+    num_sequences = max(1, min(num_tokens, 8))
+    topk_ids, topk_weights, token_lora_mapping, lora_ids = sample_data(
+        num_tokens, num_sequences, max_loras, num_experts, top_k_num
+    )
+
+    lora_a_stacked = [
+        torch.rand(
+            (max_loras, num_experts, max_lora_rank, K),
+            dtype=dtype,
+        )
+        for _ in range(num_slices)
+    ]
+    lora_b_stacked = [
+        torch.rand(
+            (max_loras, num_experts, N // num_slices, max_lora_rank),
+            dtype=dtype,
+        )
+        for _ in range(num_slices)
+    ]
+    hidden_states = torch.rand((num_tokens, K), dtype=dtype)
+
+    output = torch.zeros((num_tokens, top_k_num, N), dtype=dtype)
+    use_fused_moe_lora_kernel(
+        topk_ids,
+        topk_weights,
+        token_lora_mapping,
+        max_lora_rank,
+        top_k_num,
+        lora_ids,
+        lora_a_stacked,
+        lora_b_stacked,
+        hidden_states,
+        output,
+        max_loras,
+        num_experts,
+        block_size,
+    )
+    output_ref = use_torch(
+        hidden_states,
+        token_lora_mapping,
+        topk_ids,
+        lora_a_stacked,
+        lora_b_stacked,
+        top_k_num,
+        num_slices,
+    )
+    torch.testing.assert_close(output, output_ref, atol=1e-2, rtol=1e-2)
+
+
+@pytest.mark.parametrize("num_tokens", [16, 64])
+@pytest.mark.parametrize("top_k_num", [2])
+@pytest.mark.parametrize("num_experts", [8])
+@pytest.mark.parametrize("max_loras", [4])
+@pytest.mark.parametrize("N", [2048])
+@pytest.mark.parametrize("K", [4096])
+@pytest.mark.parametrize("max_lora_rank", [8, 16, 32, 64])
+@pytest.mark.parametrize("block_size", [64])
+@pytest.mark.parametrize("num_slices", [2])
+@pytest.mark.parametrize("dtype", [torch.bfloat16])
+@pytest.mark.parametrize("device", DEVICES)
+@pytest.mark.parametrize("seed", SEED)
+def test_fused_moe_lora_kernel_npid_path(
+    num_tokens,
+    top_k_num,
+    num_experts,
+    max_loras,
+    N,
+    K,
+    max_lora_rank,
+    block_size,
+    num_slices,
+    dtype,
+    device,
+    seed,
+):
+    """Exercise the small-batch / NPID > 1 branch of the one-shot fast path.
+
+    With these sizes the one-shot wrapper computes NPID_FACTOR > 1 (base CTA count
+    < SM count), so each program covers only an outer chunk of N. The
+    cross-outer-block write mask is the correctness-critical bit.
+    """
+    torch.set_default_device(device)
+    set_random_seed(seed)
+    num_sequences = max(1, min(num_tokens, 4))
+    topk_ids, topk_weights, token_lora_mapping, lora_ids = sample_data(
+        num_tokens, num_sequences, max_loras, num_experts, top_k_num
+    )
+
+    lora_a_stacked = [
+        torch.rand(
+            (max_loras, num_experts, max_lora_rank, K),
+            dtype=dtype,
+        )
+        for _ in range(num_slices)
+    ]
+    lora_b_stacked = [
+        torch.rand(
+            (max_loras, num_experts, N // num_slices, max_lora_rank),
+            dtype=dtype,
+        )
+        for _ in range(num_slices)
+    ]
+    hidden_states = torch.rand((num_tokens, K), dtype=dtype)
+
+    output = torch.zeros((num_tokens, top_k_num, N), dtype=dtype)
+    use_fused_moe_lora_kernel(
+        topk_ids,
+        topk_weights,
+        token_lora_mapping,
+        max_lora_rank,
+        top_k_num,
+        lora_ids,
+        lora_a_stacked,
+        lora_b_stacked,
+        hidden_states,
+        output,
+        max_loras,
+        num_experts,
+        block_size,
+    )
+    output_ref = use_torch(
+        hidden_states,
+        token_lora_mapping,
+        topk_ids,
+        lora_a_stacked,
+        lora_b_stacked,
+        top_k_num,
+        num_slices,
+    )
+    torch.testing.assert_close(output, output_ref, atol=2e-2, rtol=2e-2)
+
+
+# -- one-shot corner-case coverage ------------------------------------------
+# Each of the following exercises a path where the kernel is launched but
+# every program early-exits, leaving the output unchanged. The contract is
+# additive (`output += contribution`), so an empty contribution must leave
+# the input residual untouched.
+
+
+def _build_one_shot_inputs(
+    num_tokens,
+    top_k_num,
+    num_experts,
+    max_loras,
+    max_lora_rank,
+    K,
+    N,
+    num_slices,
+    block_size,
+    dtype,
+):
+    """Common scaffolding for the corner-case tests below."""
+    num_sequences = max(1, min(num_tokens, 4)) if num_tokens > 0 else 1
+    if num_tokens > 0:
+        topk_ids, topk_weights, token_lora_mapping, lora_ids = sample_data(
+            num_tokens, num_sequences, max_loras, num_experts, top_k_num
+        )
+    else:
+        # M=0 path: caller may still hand us empty tensors with the right shape.
+        topk_ids = torch.empty((0, top_k_num), dtype=torch.int32)
+        topk_weights = torch.empty((0, top_k_num), dtype=torch.float32)
+        token_lora_mapping = torch.empty((0,), dtype=torch.int32)
+        lora_ids = torch.full((max_loras + 1,), -1, dtype=torch.int32)
+
+    lora_a_stacked = [
+        torch.rand((max_loras, num_experts, max_lora_rank, K), dtype=dtype)
+        for _ in range(num_slices)
+    ]
+    lora_b_stacked = [
+        torch.rand(
+            (max_loras, num_experts, N // num_slices, max_lora_rank), dtype=dtype
+        )
+        for _ in range(num_slices)
+    ]
+    hidden_states = torch.rand((max(num_tokens, 0), K), dtype=dtype)
+    return (
+        topk_ids,
+        topk_weights.to(dtype),
+        token_lora_mapping,
+        lora_ids,
+        lora_a_stacked,
+        lora_b_stacked,
+        hidden_states,
+    )
+
+
+def _call_one_shot(
+    output,
+    hidden_states,
+    lora_a_stacked,
+    lora_b_stacked,
+    topk_weights,
+    sorted_token_ids,
+    expert_ids,
+    num_tokens_post_padded,
+    token_lora_mapping,
+    max_lora_rank,
+    top_k_num,
+    lora_ids,
+    num_active_loras,
+    adapter_enabled,
+    block_size,
+    add_inputs=True,
+):
+    """Direct call into fused_moe_lora with one-shot-routed defaults."""
+    from vllm.lora.ops.triton_ops import fused_moe_lora as _op
+
+    _op(
+        output,
+        hidden_states,
+        lora_a_stacked,
+        lora_b_stacked,
+        topk_weights,
+        sorted_token_ids,
+        expert_ids,
+        num_tokens_post_padded,
+        token_lora_mapping,
+        max_lora_rank,
+        top_k_num,
+        lora_ids,
+        num_active_loras,
+        adapter_enabled,
+        block_size,
+        32,
+        64,
+        1,
+        4,
+        3,
+        1,
+        block_size,
+        32,
+        64,
+        1,
+        4,
+        3,
+        1,
+        False,
+        False,
+        0,
+        add_inputs,
+    )
+
+
+@pytest.mark.parametrize(
+    "trigger",
+    ["sorted_lora_ids_neg", "naive_mapping_neg", "naive_all_disabled"],
+)
+@pytest.mark.parametrize("device", DEVICES)
+def test_fused_moe_lora_kernel_one_shot_early_exit(trigger, device):
+    """one-shot must leave the residual byte-identical when every program
+    must early-exit. Three trigger conditions are covered:
+
+      - "sorted_lora_ids_neg":  sorted path, lora_ids all -1 (lora_id<0 check)
+      - "naive_mapping_neg":    naive path, token_lora_mapping all -1
+      - "naive_all_disabled":   naive path, adapter_enabled all 0
+    """
+    torch.set_default_device(device)
+    set_random_seed(0)
+
+    # Per-trigger shapes: naive_mapping_neg needs the naive dispatch gate
+    # `num_tokens*top_k*8 <= num_experts*max_loras` to hold, hence the
+    # larger E/max_loras and smaller num_tokens.
+    if trigger == "naive_mapping_neg":
+        num_tokens, top_k, E, max_loras, R = 8, 2, 64, 8, 16
+    elif trigger == "naive_all_disabled":
+        num_tokens, top_k, E, max_loras, R = 32, 2, 8, 4, 32
+    else:  # sorted_lora_ids_neg
+        num_tokens, top_k, E, max_loras, R = 32, 2, 8, 4, 16
+    K, N = 1024, 1024
+    block_size, num_slices, dtype = 16, 2, torch.bfloat16
+
+    (
+        topk_ids,
+        topk_weights,
+        token_lora_mapping,
+        lora_ids,
+        lora_a_stacked,
+        lora_b_stacked,
+        hidden_states,
+    ) = _build_one_shot_inputs(
+        num_tokens, top_k, E, max_loras, R, K, N, num_slices, block_size, dtype
+    )
+
+    adapter_enabled = torch.ones(max_loras + 1, dtype=torch.int32)
+    num_active_loras = torch.tensor([max_loras + 1], dtype=torch.int32, device="cpu")
+
+    if trigger == "sorted_lora_ids_neg":
+        lora_ids = torch.full((max_loras + 1,), -1, dtype=torch.int32)
+        max_pad = topk_ids.numel() + E * (block_size - 1)
+        max_pad = round_up(max_pad, block_size)
+        max_blocks = CEILDIV(max_pad, block_size)
+        sorted_token_ids = torch.zeros((max_loras, max_pad), dtype=torch.int32)
+        expert_ids = torch.full((max_loras, max_blocks), -1, dtype=torch.int32)
+        num_post = torch.zeros((max_loras,), dtype=torch.int32)
+    else:
+        sorted_token_ids = None
+        expert_ids = topk_ids.reshape(-1).contiguous()
+        num_post = None
+        if trigger == "naive_mapping_neg":
+            token_lora_mapping = torch.full((num_tokens,), -1, dtype=torch.int32)
+            lora_ids = torch.full((max_loras + 1,), -1, dtype=torch.int32)
+        else:  # naive_all_disabled
+            adapter_enabled = torch.zeros(max_loras + 1, dtype=torch.int32)
+
+    residual = torch.randn((num_tokens, top_k, N), dtype=dtype) * 0.1
+    output = residual.clone()
+
+    _call_one_shot(
+        output,
+        hidden_states,
+        lora_a_stacked,
+        lora_b_stacked,
+        topk_weights,
+        sorted_token_ids,
+        expert_ids,
+        num_post,
+        token_lora_mapping,
+        R,
+        top_k,
+        lora_ids,
+        num_active_loras,
+        adapter_enabled,
+        block_size,
+    )
+    torch.testing.assert_close(output, residual, atol=0, rtol=0)
+
+
+@pytest.mark.parametrize("device", DEVICES)
+def test_fused_moe_lora_kernel_zero_grid_no_crash(device):
+    """num_active_loras=0 (or num_slices=0) would otherwise launch a grid
+    with a zero dimension. one-shot wrapper must short-circuit before launch."""
+    torch.set_default_device(device)
+    set_random_seed(0)
+    num_tokens, top_k, E, max_loras, R, K, N = 8, 2, 8, 4, 16, 1024, 1024
+    block_size, num_slices, dtype = 16, 2, torch.bfloat16
+
+    (
+        topk_ids,
+        topk_weights,
+        token_lora_mapping,
+        lora_ids,
+        lora_a_stacked,
+        lora_b_stacked,
+        hidden_states,
+    ) = _build_one_shot_inputs(
+        num_tokens,
+        top_k,
+        E,
+        max_loras,
+        R,
+        K,
+        N,
+        num_slices,
+        block_size,
+        dtype,
+    )
+    adapter_enabled = torch.ones(max_loras + 1, dtype=torch.int32)
+    num_active_loras = torch.tensor([0], dtype=torch.int32, device="cpu")
+    residual = torch.randn((num_tokens, top_k, N), dtype=dtype) * 0.1
+    output = residual.clone()
+
+    # sorted path is the one that uses num_active_loras for grid axis 2
+    max_pad = topk_ids.numel() + E * (block_size - 1)
+    max_pad = round_up(max_pad, block_size)
+    max_blocks = CEILDIV(max_pad, block_size)
+    sorted_token_ids = torch.zeros((max_loras, max_pad), dtype=torch.int32)
+    expert_ids = torch.full((max_loras, max_blocks), -1, dtype=torch.int32)
+    num_post = torch.zeros((max_loras,), dtype=torch.int32)
+    _call_one_shot(
+        output,
+        hidden_states,
+        lora_a_stacked,
+        lora_b_stacked,
+        topk_weights,
+        sorted_token_ids,
+        expert_ids,
+        num_post,
+        token_lora_mapping,
+        R,
+        top_k,
+        lora_ids,
+        num_active_loras,
+        adapter_enabled,
+        block_size,
+    )
+    torch.testing.assert_close(output, residual, atol=0, rtol=0)
+
+
+@pytest.mark.parametrize("device", DEVICES)
+def test_fused_moe_lora_kernel_rejects_bad_block_size_m(device):
+    """one-shot must surface a clear assertion when shrink_block_size_m is not
+    a power of 2 / less than 16, instead of the cryptic Triton compile
+    failure (`arange's range must be a power of 2`)."""
+    torch.set_default_device(device)
+    set_random_seed(0)
+    num_tokens, top_k, E, max_loras, R, K, N = 32, 2, 8, 4, 16, 1024, 1024
+    num_slices, dtype = 2, torch.bfloat16
+    block_size = 24  # NOT a power of 2
+
+    (
+        topk_ids,
+        topk_weights,
+        token_lora_mapping,
+        lora_ids,
+        lora_a_stacked,
+        lora_b_stacked,
+        hidden_states,
+    ) = _build_one_shot_inputs(
+        num_tokens,
+        top_k,
+        E,
+        max_loras,
+        R,
+        K,
+        N,
+        num_slices,
+        16,
+        dtype,
+    )
+    # Build sorted-mode metadata at block_size=16 so shapes are sane,
+    # but pass block_size=24 to the op (the buggy combination).
+    max_pad = topk_ids.numel() + E * (16 - 1)
+    max_pad = round_up(max_pad, 16)
+    max_blocks = CEILDIV(max_pad, 16)
+    sorted_token_ids = torch.zeros((max_loras, max_pad), dtype=torch.int32)
+    expert_ids = torch.full((max_loras, max_blocks), -1, dtype=torch.int32)
+    num_post = torch.zeros((max_loras,), dtype=torch.int32)
+    adapter_enabled = torch.ones(max_loras + 1, dtype=torch.int32)
+    num_active_loras = torch.tensor([max_loras + 1], dtype=torch.int32, device="cpu")
+    output = torch.zeros((num_tokens, top_k, N), dtype=dtype)
+
+    with pytest.raises(AssertionError, match="shrink_block_size_m"):
+        _call_one_shot(
+            output,
+            hidden_states,
+            lora_a_stacked,
+            lora_b_stacked,
+            topk_weights,
+            sorted_token_ids,
+            expert_ids,
+            num_post,
+            token_lora_mapping,
+            R,
+            top_k,
+            lora_ids,
+            num_active_loras,
+            adapter_enabled,
+            block_size,
+        )

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/models/test_deepseek_v4_mega_moe.py

@@ -8,9 +8,9 @@ import torch
 
 from vllm.models.deepseek_v4.nvidia.model import (
     DeepseekV4MegaMoEExperts,
-    _stage_deepseek_v4_mega_moe_inputs,
     make_deepseek_v4_expert_params_mapping,
 )
+from vllm.models.deepseek_v4.nvidia.ops.prepare_megamoe import prepare_megamoe_inputs
 from vllm.platforms import current_platform
 
 pytestmark = pytest.mark.skipif(
@@ -164,7 +164,7 @@ def test_deepseek_v4_mega_moe_fused_input_staging_is_bitwise_exact():
     fused_topk_idx = torch.empty_like(ref_topk_idx)
     fused_topk_weights = torch.empty_like(ref_topk_weights)
 
-    _stage_deepseek_v4_mega_moe_inputs(
+    prepare_megamoe_inputs(
         hidden_states,
         topk_weights,
         topk_ids,

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/kv_connector/unit/test_mooncake_store_connector.py

@@ -103,14 +103,18 @@ def test_worker_methods_delegate_to_store_worker():
 
     worker = mock_worker_cls.return_value
     worker.get_finished.return_value = ({"req-1"}, {"req-2"})
+    worker.get_block_ids_with_load_errors.return_value = {3, 4}
     connector.bind_connector_metadata(metadata)
 
     connector.register_kv_caches(kv_caches)
     result = connector.get_finished(finished_req_ids)
+    invalid_block_ids = connector.get_block_ids_with_load_errors()
 
     worker.register_kv_caches.assert_called_once_with(kv_caches)
     worker.get_finished.assert_called_once_with(finished_req_ids, metadata)
     assert result == ({"req-1"}, {"req-2"})
+    worker.get_block_ids_with_load_errors.assert_called_once_with()
+    assert invalid_block_ids == {3, 4}
 
 
 def test_get_kv_connector_kv_cache_events_returns_none_when_empty():

tests/v1/kv_connector/unit/test_mooncake_store_coordinator.py

@@ -300,3 +300,85 @@ def test_store_mask_fast_path_single_attention_group():
     assert len(coord.attention_groups) == 1
     masks = coord.store_mask(64)
     assert masks == ([True] * 4, [True] * 4)
+
+
+# ----- Eagle / MTP interaction with load_mask -----
+
+
+def test_lookup_with_eagle_pops_last_full_attention_block():
+    """Sanity: with use_eagle, find_longest_cache_hit drops the last block.
+    Pairs with the load_mask test below to lock the round-trip contract."""
+    groups = [KVCacheGroupSpec(["L0"], _full(16))]
+    coord = _make_coord(groups, hash_block_size=16, use_eagle=True)
+    hs = _hashes(4)
+    cmap = ExternalCachedBlockPool({(0, bytes(h)) for h in hs})
+    _masks, hit = coord.find_longest_cache_hit(
+        hs, max_length=64, cached_block_pool=cmap
+    )
+    # 4 blocks present, eagle pops 1 → 3 blocks = 48 tokens.
+    assert hit == 48
+
+
+def test_load_mask_with_eagle_does_not_double_prune_full_attention():
+    """Regression for silent KV corruption with MTP/EAGLE-3.
+
+    The recv side calls ``load_mask(block_hashes, token_len)`` where
+    ``token_len`` is already the eagle-pruned hit length from ``lookup``.
+    A second eagle pop here used to shorten the mask by one extra block;
+    ``process_tokens`` then yielded a chunk past the mask, which the worker
+    silently skipped — leaving the trailing block of the loaded prefix
+    uninitialized in local KV.
+    """
+    groups = [KVCacheGroupSpec(["L0"], _full(16))]
+    coord = _make_coord(groups, hash_block_size=16, use_eagle=True)
+    hs = _hashes(4)
+    cmap = ExternalCachedBlockPool({(0, bytes(h)) for h in hs})
+    _masks, hit = coord.find_longest_cache_hit(
+        hs, max_length=64, cached_block_pool=cmap
+    )
+    assert hit == 48  # eagle popped 1 block
+
+    masks = coord.load_mask(hs, token_len=hit)
+    # Every chunk that process_tokens(token_len=48, ...) would yield must
+    # have a corresponding mask slot. process_tokens emits chunk_id 0..2
+    # (start=0, 16, 32), so the mask must be length 3, all True.
+    assert masks[0] == [True, True, True]
+
+
+def test_load_mask_with_eagle_hybrid_full_plus_swa():
+    """Hybrid (FullAttn + SWA) with eagle: load_mask must cover every chunk
+    in [0, token_len) for the FullAttn group; SWA group keeps its
+    tail-window mask."""
+    groups = [
+        KVCacheGroupSpec(["L0"], _full(16)),
+        KVCacheGroupSpec(["L1"], _swa(16, 32)),
+    ]
+    coord = _make_coord(groups, hash_block_size=16, use_eagle=True)
+    hs = _hashes(4)
+    exists = {(g, bytes(h)) for g in (0, 1) for h in hs}
+    cmap = ExternalCachedBlockPool(exists)
+    _masks, hit = coord.find_longest_cache_hit(
+        hs, max_length=64, cached_block_pool=cmap
+    )
+    # FullAttn dictates the convergence; eagle pops one block off it.
+    assert hit == 48
+
+    masks = coord.load_mask(hs, token_len=hit)
+    # FullAttn: all chunks populated locally.
+    assert masks[0] == [True, True, True]
+    # SWA: tail-window only (ceil((32-1)/16) = 2 trailing blocks).
+    assert masks[1][-2:] == [True, True]
+
+
+def test_load_mask_without_eagle_unchanged():
+    """Sanity: when eagle is off, load_mask is identical to the pre-fix path."""
+    groups = [KVCacheGroupSpec(["L0"], _full(16))]
+    coord = _make_coord(groups, hash_block_size=16, use_eagle=False)
+    hs = _hashes(4)
+    cmap = ExternalCachedBlockPool({(0, bytes(h)) for h in hs})
+    _masks, hit = coord.find_longest_cache_hit(
+        hs, max_length=64, cached_block_pool=cmap
+    )
+    assert hit == 64
+    masks = coord.load_mask(hs, token_len=hit)
+    assert masks[0] == [True, True, True, True]

tests/v1/kv_connector/unit/test_mooncake_store_worker.py

@@ -77,6 +77,7 @@ def _make_store_sending_thread(
 def _make_store_recving_thread(
     store: MagicMock,
     *,
+    tp_rank: int = 0,
     disk_offload_buffer_budget_bytes: int | None = None,
 ) -> mooncake_store_worker.KVCacheStoreRecvingThread:
     from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheGroupSpec
@@ -96,7 +97,7 @@ def _make_store_recving_thread(
         store=store,
         token_databases=[token_database],
         block_size=16,
-        tp_rank=0,
+        tp_rank=tp_rank,
         ready_event=threading.Event(),
         coord=coord,
         disk_offload_buffer_budget_bytes=disk_offload_buffer_budget_bytes,
@@ -460,6 +461,67 @@ def test_store_sending_thread_only_skips_on_no_available_handle():
     assert store.batch_put_from_multi_buffers.call_count == 2
 
 
+def test_store_recving_thread_reports_failed_block_ids():
+    store = MagicMock()
+    store.batch_get_into_multi_buffers.return_value = [256, -5, -7]
+    thread = _make_store_recving_thread(store)
+
+    thread._handle_request(
+        _make_load_req(
+            "req-a",
+            [b"a0", b"a1", b"a2"],
+            token_len=48,
+        )
+    )
+
+    assert thread.get_and_clear_finished_requests() == {"req-a"}
+    assert thread.get_and_clear_block_ids_with_load_errors() == {1, 2}
+    assert thread.get_and_clear_block_ids_with_load_errors() == set()
+
+
+def test_store_recving_thread_reports_failed_block_ids_after_rotation():
+    store = MagicMock()
+    store.batch_get_into_multi_buffers.return_value = [256, -5, 256]
+    thread = _make_store_recving_thread(store, tp_rank=1)
+
+    thread._handle_request(
+        _make_load_req(
+            "req-a",
+            [b"a0", b"a1", b"a2"],
+            token_len=48,
+        )
+    )
+
+    assert thread.get_and_clear_block_ids_with_load_errors() == {2}
+
+
+def test_store_recving_thread_reports_all_attempted_blocks_on_exception():
+    store = MagicMock()
+    store.batch_get_into_multi_buffers.side_effect = RuntimeError("boom")
+    thread = _make_store_recving_thread(store)
+
+    thread._handle_request(
+        _make_load_req(
+            "req-a",
+            [b"a0", b"a1", b"a2"],
+            token_len=48,
+        )
+    )
+
+    assert thread.get_and_clear_finished_requests() == {"req-a"}
+    assert thread.get_and_clear_block_ids_with_load_errors() == {0, 1, 2}
+
+
+def test_store_worker_get_block_ids_with_load_errors_delegates_to_recv_thread():
+    recv_thread = MagicMock()
+    recv_thread.get_and_clear_block_ids_with_load_errors.return_value = {3, 4}
+    w = _make_bare_worker()
+    w.kv_recv_thread = recv_thread
+
+    assert w.get_block_ids_with_load_errors() == {3, 4}
+    recv_thread.get_and_clear_block_ids_with_load_errors.assert_called_once_with()
+
+
 def test_store_sending_thread_passes_replicate_config_when_preferred_segment_set():
     store = MagicMock()
     store.batch_is_exist.side_effect = lambda keys: [0] * len(keys)
@@ -682,6 +744,7 @@ def test_recv_thread_stops_after_first_failing_disk_offload_sub_batch():
     thread._handle_request(req)
 
     assert store.batch_get_into_multi_buffers.call_count == 1
+    assert thread.get_and_clear_block_ids_with_load_errors() == {0, 1}
 
 
 def test_recv_thread_skips_split_when_budget_holds_all_keys():
@@ -711,21 +774,26 @@ def test_recv_thread_skips_split_when_budget_holds_all_keys():
 
 
 def test_recv_thread_reports_unsplittable_key_larger_than_budget():
+    # Non-zero tp_rank exercises the rotation: the oversized key may not sit
+    # at the original request's first block. Every block must still be marked
+    # invalid, since none are loaded.
     store = MagicMock()
     thread = _make_store_recving_thread(
         store,
+        tp_rank=2,
         disk_offload_buffer_budget_bytes=_DISK_OFFLOAD_BUDGET_TOO_SMALL,
     )
 
     req = _make_load_req(
         "req-a",
-        [b"a0"],
-        token_len=16,
+        [b"a0", b"a1", b"a2"],
+        token_len=48,
     )
 
     thread._handle_request(req)
 
     assert store.batch_get_into_multi_buffers.call_count == 0
+    assert thread.get_and_clear_block_ids_with_load_errors() == {0, 1, 2}
 
 
 def test_requester_worker_init_uses_positional_setup(tmp_path, monkeypatch):

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/distributed/kv_transfer/kv_connector/v1/mooncake/store/connector.py

@@ -270,6 +270,10 @@ class MooncakeStoreConnector(KVConnectorBase_V1, SupportsHMA):
         assert isinstance(metadata, MooncakeStoreConnectorMetadata)
         return self.connector_worker.get_finished(finished_req_ids, metadata)
 
+    def get_block_ids_with_load_errors(self) -> set[int]:
+        assert self.connector_worker is not None
+        return self.connector_worker.get_block_ids_with_load_errors()
+
     def get_kv_connector_kv_cache_events(
         self,
     ) -> MooncakeStoreKVEvents | None:

vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/coordinator.py

@@ -115,12 +115,21 @@ class MooncakeStoreCoordinator:
         block_hashes: list[BlockHash],
         max_length: int,
         cached_block_pool: ExternalCachedBlockPool,
+        *,
+        apply_eagle: bool = True,
     ) -> tuple[tuple[list[bool], ...], int]:
         """Returns ``(load_mask_per_group, hit_length)``. ``mask[g][i]`` is True iff
         group ``g`` populates chunk ``i`` locally (e.g. SWA and Mamba tail-only);
-        recv-side callers skip False slots."""
+        recv-side callers skip False slots.
+
+        ``apply_eagle`` controls whether the per-spec ``use_eagle`` last-block
+        pop is applied. Lookup callers want it (the drafter requires recomputing
+        the last block); per-chunk mask callers must not, because ``token_len``
+        already reflects the eagle-pruned hit length and a second pop would
+        leave the trailing block unloaded.
+        """
         blocks_per_group, hit_length = self._find_hit_blocks(
-            block_hashes, max_length, cached_block_pool
+            block_hashes, max_length, cached_block_pool, apply_eagle=apply_eagle
         )
         masks = tuple(
             [blk is not cached_block_pool.null_block for blk in blocks]
@@ -137,8 +146,17 @@ class MooncakeStoreCoordinator:
         spec would populate chunk ``i`` locally at length ``token_len``
         (e.g. SWA / Mamba tail-only).
         """
+        # ``apply_eagle=False`` because ``token_len`` is already the
+        # eagle-pruned hit length returned by ``client.lookup``. Re-applying
+        # the pop here would shorten the mask by one extra block; the recv
+        # thread would then silently skip the trailing chunk yielded by
+        # ``db.process_tokens`` and leave that block uninitialized in the
+        # local KV pool.
         masks, _ = self.find_longest_cache_hit(
-            block_hashes, token_len, ExternalCachedBlockPool()
+            block_hashes,
+            token_len,
+            ExternalCachedBlockPool(),
+            apply_eagle=False,
         )
         return masks
 
@@ -195,10 +213,17 @@ class MooncakeStoreCoordinator:
         block_hashes: list[BlockHash],
         max_length: int,
         cached_block_pool: ExternalCachedBlockPool,
+        *,
+        apply_eagle: bool = True,
     ) -> tuple[tuple[list[KVCacheBlock], ...], int]:
         """Mirrors HybridKVCacheCoordinator.find_longest_cache_hit but
         dispatches via spec_manager_map (we don't allocate managers).
+
+        When ``apply_eagle`` is False, ignore ``eagle_attn_group_indices`` —
+        used by ``load_mask`` to avoid popping a second block on top of the
+        one already removed by the lookup.
         """
+        eagle_indices = self.eagle_attn_group_indices if apply_eagle else set()
         if len(self.attention_groups) == 1:
             spec, group_ids, manager_cls = self.attention_groups[0]
             hashes = self.block_hashes_for_spec(block_hashes, spec)
@@ -208,7 +233,7 @@ class MooncakeStoreCoordinator:
                 kv_cache_group_ids=group_ids,
                 block_pool=cast(BlockPool, cached_block_pool),
                 kv_cache_spec=spec,
-                use_eagle=(0 in self.eagle_attn_group_indices),
+                use_eagle=(0 in eagle_indices),
                 alignment_tokens=spec.block_size,
             )
             num_groups = len(self.kv_cache_groups)
@@ -237,9 +262,7 @@ class MooncakeStoreCoordinator:
                     )
                     continue
 
-                use_eagle = (
-                    idx in self.eagle_attn_group_indices and idx not in eagle_verified
-                )
+                use_eagle = idx in eagle_indices and idx not in eagle_verified
                 _max_length = curr_hit_length
                 if use_eagle:
                     _max_length = min(curr_hit_length + spec.block_size, max_length)

vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py

@@ -20,7 +20,7 @@ import time
 from collections import defaultdict
 from collections.abc import Callable
 from dataclasses import dataclass
-from typing import Any, Literal
+from typing import Any, Literal, TypeVar
 
 import regex as re
 import torch
@@ -68,6 +68,12 @@ DEFAULT_GLOBAL_SEGMENT_SIZE = 4 * 1024 * 1024 * 1024  # 4 GiB
 DEFAULT_LOCAL_BUFFER_SIZE = 4 * 1024 * 1024 * 1024  # 4 GiB
 
 MOONCAKE_NO_AVAILABLE_HANDLE = -200
+_T = TypeVar("_T")
+
+
+def _rotate_list(values: list[_T], offset: int) -> list[_T]:
+    return values[offset:] + values[:offset]
+
 
 # Mirrors FileStorageConfig::local_buffer_size in Mooncake C++.
 DEFAULT_MOONCAKE_DISK_STAGING_BUFFER_BYTES = 1280 * 1024 * 1024
@@ -711,6 +717,9 @@ class KVCacheStoreRecvingThread(KVTransferThread):
             name="KVCacheStoreRecvingThread",
             record_operation=record_operation,
         )
+        # _invalid_block_ids can be access by both the Worker and RecvingThread
+        self._invalid_block_ids_lock = threading.Lock()
+        self._invalid_block_ids: set[int] = set()
         self.disk_offload_buffer_budget_bytes = disk_offload_buffer_budget_bytes
         self.usable_disk_offload_buffer_budget_bytes = (
             None
@@ -721,6 +730,16 @@ class KVCacheStoreRecvingThread(KVTransferThread):
         )
         self.coord = coord
 
+    def _add_load_error_block_ids(self, block_ids: list[int]) -> None:
+        with self._invalid_block_ids_lock:
+            self._invalid_block_ids.update(block_ids)
+
+    def get_and_clear_block_ids_with_load_errors(self) -> set[int]:
+        with self._invalid_block_ids_lock:
+            invalid_block_ids = self._invalid_block_ids.copy()
+            self._invalid_block_ids.clear()
+        return invalid_block_ids
+
     def _handle_request(self, req_meta: ReqMeta):
         token_len = req_meta.load_spec.token_len  # type: ignore[union-attr]
         req_id = req_meta.req_id
@@ -737,6 +756,7 @@ class KVCacheStoreRecvingThread(KVTransferThread):
         addr_list: list[list[int]] = []
         size_list: list[list[int]] = []
         key_list: list[str] = []
+        block_id_list: list[int] = []
         for g_idx, db in enumerate(self.token_databases):
             mask = load_mask_per_group[g_idx]
             for start, end, key in db.process_tokens(
@@ -745,25 +765,29 @@ class KVCacheStoreRecvingThread(KVTransferThread):
                 chunk_idx = start // db.block_size
                 if chunk_idx >= len(mask) or not mask[chunk_idx]:
                     continue
-                addr, size, _ = db.prepare_value(start, end, req_meta.block_ids[g_idx])
+                addr, size, block_id = db.prepare_value(
+                    start, end, req_meta.block_ids[g_idx]
+                )
                 key_list.append(key.to_string())
                 addr_list.append(addr)
                 size_list.append(size)
+                block_id_list.append(block_id)
 
-        # Rotate lists by tp_rank for load balancing
+        # Rotate aligned lists by tp_rank for load balancing.
         rotation = self.tp_rank % len(key_list)
-        key_list_c = key_list[rotation:] + key_list[:rotation]
-        addr_list_c = addr_list[rotation:] + addr_list[:rotation]
-        size_list_c = size_list[rotation:] + size_list[:rotation]
+        key_list_c = _rotate_list(key_list, rotation)
+        addr_list_c = _rotate_list(addr_list, rotation)
+        size_list_c = _rotate_list(size_list, rotation)
+        block_id_list_c = _rotate_list(block_id_list, rotation)
 
-        load_batches = [(key_list_c, addr_list_c, size_list_c)]
+        load_batches = [(key_list_c, addr_list_c, size_list_c, block_id_list_c)]
         if self.usable_disk_offload_buffer_budget_bytes is not None:
             total_staging_bytes = sum(
                 _estimate_disk_offload_staging_bytes(size) for size in size_list_c
             )
             if total_staging_bytes > self.usable_disk_offload_buffer_budget_bytes:
                 assert self.disk_offload_buffer_budget_bytes is not None
-                load_batches, oversized_key = _split_disk_offload_load_batches(
+                split_batches, oversized_key = _split_disk_offload_load_batches(
                     key_list_c,
                     addr_list_c,
                     size_list_c,
@@ -772,6 +796,10 @@ class KVCacheStoreRecvingThread(KVTransferThread):
                 )
                 if oversized_key is not None:
                     oversized_key_index = key_list_c.index(oversized_key)
+                    # Mark every block: we skip the whole request, and the
+                    # tp_rank rotation means oversized_key isn't necessarily
+                    # the first block in the request's original order.
+                    self._add_load_error_block_ids(block_id_list_c)
                     oversized_key_bytes = _estimate_disk_offload_staging_bytes(
                         size_list_c[oversized_key_index]
                     )
@@ -786,12 +814,25 @@ class KVCacheStoreRecvingThread(KVTransferThread):
                     self.set_finished_request(req_id)
                     self.request_queue.task_done()
                     return
+                load_batches = []
+                block_id_offset = 0
+                for batch_keys, batch_addrs, batch_sizes in split_batches:
+                    next_block_id_offset = block_id_offset + len(batch_keys)
+                    batch_block_ids = block_id_list_c[
+                        block_id_offset:next_block_id_offset
+                    ]
+                    load_batches.append(
+                        (batch_keys, batch_addrs, batch_sizes, batch_block_ids)
+                    )
+                    block_id_offset = next_block_id_offset
 
         current_batch_keys: list[str] = key_list_c
+        current_batch_block_ids: list[int] = block_id_list_c
         batch_bytes = 0
         try:
-            for batch_keys, batch_addrs, batch_sizes in load_batches:
+            for batch_keys, batch_addrs, batch_sizes, batch_block_ids in load_batches:
                 current_batch_keys = batch_keys
+                current_batch_block_ids = batch_block_ids
                 batch_bytes = _sum_batch_bytes(batch_sizes)
                 tiers_by_key: dict[str, str] | None = None
                 if envs.VLLM_MOONCAKE_STORE_TIER_LOG:
@@ -806,8 +847,10 @@ class KVCacheStoreRecvingThread(KVTransferThread):
                         req_id, batch_keys, res, tiers_by_key
                     )
                 failed = [
-                    (key, value)
-                    for key, value in zip(batch_keys, res, strict=True)
+                    (key, value, block_id)
+                    for key, value, block_id in zip(
+                        batch_keys, res, batch_block_ids, strict=True
+                    )
                     if value < 0
                 ]
                 self._record_operation(
@@ -819,15 +862,19 @@ class KVCacheStoreRecvingThread(KVTransferThread):
                     num_failed_keys=len(failed),
                 )
                 if failed:
+                    self._add_load_error_block_ids(
+                        [block_id for _, _, block_id in failed]
+                    )
                     logger.warning(
                         "Failed to get %d Mooncake keys from sub-batch "
                         "(batch_keys=%d, first_failures=%s)",
                         len(failed),
                         len(batch_keys),
-                        failed[:3],
+                        [(key, value) for key, value, _ in failed[:3]],
                     )
                     break
         except Exception as e:
+            self._add_load_error_block_ids(current_batch_block_ids)
             self._record_operation(
                 "load_get",
                 load_get_start,
@@ -1239,6 +1286,11 @@ class MooncakeStoreWorker:
         )
         return done_sending, done_recving
 
+    def get_block_ids_with_load_errors(self) -> set[int]:
+        if self.kv_recv_thread is None:
+            return set()
+        return self.kv_recv_thread.get_and_clear_block_ids_with_load_errors()
+
     def _record_kv_connector_operation(
         self,
         operation: str,

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/generate/beam_search/offline.py

@@ -2,17 +2,13 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 import itertools
-from abc import ABC, abstractmethod
-from collections.abc import Callable, Iterable, Sequence
-from typing import Any
 
 from tqdm import tqdm
 
-from vllm import PromptType, RequestOutput, TextPrompt, TokensPrompt
-from vllm.inputs import EngineInput
+from vllm import RequestOutput, TextPrompt, TokensPrompt
+from vllm.entrypoints.offline_utils import OfflineInferenceMixin
 from vllm.logger import init_logger
 from vllm.lora.request import LoRARequest
-from vllm.renderers import BaseRenderer
 from vllm.sampling_params import BeamSearchParams, SamplingParams
 
 from .utils import (
@@ -25,11 +21,9 @@ from .utils import (
 logger = init_logger(__name__)
 
 
-class BeamSearchOfflineMixin(ABC):
+class BeamSearchOfflineMixin(OfflineInferenceMixin):
     """Offline inference for beam search"""
 
-    renderer: BaseRenderer
-
     def beam_search(
         self,
         prompts: list[TokensPrompt | TextPrompt],
@@ -186,41 +180,3 @@ class BeamSearchOfflineMixin(ABC):
             outputs.append(BeamSearchOutput(sequences=best_beams))
 
         return outputs
-
-    @abstractmethod
-    def _preprocess_cmpl(
-        self,
-        prompts: Sequence[PromptType],
-        tokenization_kwargs: dict[str, Any] | None = None,
-        mm_processor_kwargs: dict[str, Any] | None = None,
-    ) -> Sequence[EngineInput]:
-        raise NotImplementedError
-
-    @abstractmethod
-    def _lora_request_to_seq(
-        self,
-        lora_request: LoRARequest | None | Sequence[LoRARequest | None],
-        num_requests: int,
-    ) -> Sequence[LoRARequest | None]:
-        raise NotImplementedError
-
-    @abstractmethod
-    def _params_to_seq(
-        self,
-        params: SamplingParams,
-        num_requests: int,
-    ) -> Sequence[SamplingParams]:
-        raise NotImplementedError
-
-    @abstractmethod
-    def _render_and_run_requests(
-        self,
-        prompts: Iterable[EngineInput],
-        params: Sequence[SamplingParams],
-        output_type: type[RequestOutput],
-        *,
-        lora_requests: Sequence[LoRARequest | None] | None = None,
-        priorities: Sequence[int] | None = None,
-        use_tqdm: bool | Callable[..., tqdm] = True,
-    ):
-        raise NotImplementedError

vllm/entrypoints/llm.py

@@ -1,7 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
-from collections.abc import Callable, Iterable, Sequence
+from collections.abc import Callable, Sequence
 from pathlib import Path
 from typing import TYPE_CHECKING, Any
 
@@ -9,7 +9,7 @@ import cloudpickle
 import torch.nn as nn
 from pydantic import ValidationError
 from tqdm.auto import tqdm
-from typing_extensions import TypeVar, overload
+from typing_extensions import overload
 
 from vllm.config import (
     AttentionConfig,
@@ -41,47 +41,29 @@ from vllm.entrypoints.chat_utils import (
 from vllm.entrypoints.generate.beam_search.offline import BeamSearchOfflineMixin
 from vllm.entrypoints.pooling.offline import PoolingOfflineMixin
 from vllm.entrypoints.utils import log_non_default_args
-from vllm.inputs import (
-    EngineInput,
-    PromptType,
-)
+from vllm.inputs import PromptType
 from vllm.logger import init_logger
 from vllm.lora.request import LoRARequest
 from vllm.model_executor.layers.quantization import QuantizationMethods
 from vllm.outputs import PoolingRequestOutput, RequestOutput
 from vllm.platforms import current_platform
-from vllm.pooling_params import PoolingParams
-from vllm.renderers import ChatParams, merge_kwargs
-from vllm.renderers.inputs.preprocess import (
-    conversation_to_seq,
-    parse_model_prompt,
-    prompt_to_seq,
-)
-from vllm.sampling_params import RequestOutputKind, SamplingParams
+from vllm.sampling_params import SamplingParams
 from vllm.tokenizers import TokenizerLike
 from vllm.usage.usage_lib import UsageContext
 from vllm.utils.counter import Counter
-from vllm.utils.mistral import is_mistral_tokenizer
-from vllm.utils.tqdm_utils import maybe_tqdm
 from vllm.v1.engine import PauseMode
 from vllm.v1.engine.llm_engine import LLMEngine
 from vllm.v1.sample.logits_processor import LogitsProcessor
 
+from .offline_utils import _O, _R, OfflineInferenceMixin
+
 if TYPE_CHECKING:
     from vllm.v1.metrics.reader import Metric
 
 logger = init_logger(__name__)
 
-_O = TypeVar(
-    "_O",
-    bound=RequestOutput | PoolingRequestOutput,
-    default=RequestOutput | PoolingRequestOutput,
-)
-_P = TypeVar("_P", bound=SamplingParams | PoolingParams | None)
-_R = TypeVar("_R", default=Any)
 
-
-class LLM(BeamSearchOfflineMixin, PoolingOfflineMixin):
+class LLM(BeamSearchOfflineMixin, PoolingOfflineMixin, OfflineInferenceMixin):
     """An LLM for generating texts from given prompts and sampling parameters.
 
     This class includes a tokenizer, a language model (possibly distributed
@@ -585,58 +567,6 @@ class LLM(BeamSearchOfflineMixin, PoolingOfflineMixin):
 
         return self._run_engine(output_type, use_tqdm=use_tqdm)
 
-    def _resolve_mm_lora(
-        self,
-        prompt: EngineInput,
-        lora_request: LoRARequest | None,
-    ) -> LoRARequest | None:
-        if prompt["type"] != "multimodal":
-            return lora_request
-
-        lora_config = self.llm_engine.vllm_config.lora_config
-        default_mm_loras = None if lora_config is None else lora_config.default_mm_loras
-        if not default_mm_loras:
-            return lora_request
-
-        prompt_modalities = prompt["mm_placeholders"].keys()
-        intersection = set(prompt_modalities).intersection(default_mm_loras.keys())
-        if not intersection:
-            return lora_request
-
-        if len(intersection) > 1:
-            # TODO: Would be nice to be able to have multiple loras per prompt
-            logger.warning(
-                "Multiple modality specific loras were registered and would be "
-                "used by a single prompt consuming several modalities; "
-                "currently we only support one lora per request; as such, "
-                "lora(s) registered with modalities: %s will be skipped",
-                intersection,
-            )
-            return lora_request
-
-        # Build the LoRA request; the ID of the default mm lora is the
-        # index of the modality name sorted alphabetically + 1.
-        modality_name = intersection.pop()
-        modality_lora_path = default_mm_loras[modality_name]
-        modality_lora_id = sorted(default_mm_loras).index(modality_name) + 1
-
-        # If we have a collision, warn if there is a collision,
-        # but always send the explicitly provided request.
-        if lora_request:
-            if lora_request.lora_int_id != modality_lora_id:
-                logger.warning(
-                    "A modality with a registered lora and a lora_request "
-                    "with a different ID were provided; falling back to the "
-                    "lora_request as we only apply one LoRARequest per prompt"
-                )
-            return lora_request
-
-        return LoRARequest(
-            modality_name,
-            modality_lora_id,
-            modality_lora_path,
-        )
-
     def collective_rpc(
         self,
         method: str | Callable[..., _R],
@@ -682,139 +612,6 @@ class LLM(BeamSearchOfflineMixin, PoolingOfflineMixin):
         """
         return self.llm_engine.apply_model(func)
 
-    def _preprocess_cmpl(
-        self,
-        prompts: Sequence[PromptType],
-        tokenization_kwargs: dict[str, Any] | None = None,
-        mm_processor_kwargs: dict[str, Any] | None = None,
-    ) -> Sequence[EngineInput]:
-        """
-        Convert prompt inputs from LLM APIs (other than [LLM.chat][]) into
-        a format that can be passed to `_add_request`.
-
-        Refer to [LLM.generate][] for a complete description of the arguments.
-
-        Returns:
-            A list of `EngineInput` objects ready to be passed into LLMEngine.
-        """
-        renderer = self.renderer
-        model_config = self.model_config
-
-        parsed_prompts = [
-            parse_model_prompt(model_config, prompt) for prompt in prompts
-        ]
-        tok_params = renderer.default_cmpl_tok_params.with_kwargs(
-            **(tokenization_kwargs or {})
-        )
-        prompt_extras = (
-            None
-            if mm_processor_kwargs is None
-            else {"mm_processor_kwargs": mm_processor_kwargs}
-        )
-
-        return renderer.render_cmpl(
-            parsed_prompts,
-            tok_params,
-            prompt_extras=prompt_extras,
-        )
-
-    def _preprocess_cmpl_one(
-        self,
-        prompt: PromptType,
-        tokenization_kwargs: dict[str, Any] | None = None,
-        mm_processor_kwargs: dict[str, Any] | None = None,
-    ) -> EngineInput:
-        (engine_input,) = self._preprocess_cmpl(
-            [prompt],
-            tokenization_kwargs,
-            mm_processor_kwargs=mm_processor_kwargs,
-        )
-        return engine_input
-
-    def _preprocess_chat(
-        self,
-        conversations: Sequence[list[ChatCompletionMessageParam]],
-        chat_template: str | None = None,
-        chat_template_content_format: ChatTemplateContentFormatOption = "auto",
-        chat_template_kwargs: dict[str, Any] | None = None,
-        add_generation_prompt: bool = True,
-        continue_final_message: bool = False,
-        tools: list[dict[str, Any]] | None = None,
-        tokenization_kwargs: dict[str, Any] | None = None,
-        mm_processor_kwargs: dict[str, Any] | None = None,
-    ) -> Sequence[EngineInput]:
-        """
-        Convert a list of conversations into prompts so that they can then
-        be used as input for other LLM APIs.
-
-        Refer to [LLM.chat][] for a complete description of the arguments.
-
-        Returns:
-            A list of `EngineInput` objects ready to be passed into LLMEngine.
-        """
-        renderer = self.renderer
-
-        chat_params = ChatParams(
-            chat_template=chat_template,
-            chat_template_content_format=chat_template_content_format,
-            chat_template_kwargs=merge_kwargs(
-                chat_template_kwargs,
-                dict(
-                    add_generation_prompt=add_generation_prompt,
-                    continue_final_message=continue_final_message,
-                    tools=tools,
-                    tokenize=(
-                        is_mistral_tokenizer(renderer.tokenizer)
-                        or self.model_config.enable_prompt_embeds
-                    ),
-                ),
-            ),
-            mm_processor_kwargs=mm_processor_kwargs,
-        )
-        tok_params = renderer.default_chat_tok_params.with_kwargs(
-            **(tokenization_kwargs or {})
-        )
-        prompt_extras = (
-            None
-            if mm_processor_kwargs is None
-            else {"mm_processor_kwargs": mm_processor_kwargs}
-        )
-
-        _, engine_inputs = renderer.render_chat(
-            conversations,
-            chat_params,
-            tok_params,
-            prompt_extras=prompt_extras,
-        )
-
-        return engine_inputs
-
-    def _preprocess_chat_one(
-        self,
-        conversation: list[ChatCompletionMessageParam],
-        chat_template: str | None = None,
-        chat_template_content_format: ChatTemplateContentFormatOption = "auto",
-        chat_template_kwargs: dict[str, Any] | None = None,
-        add_generation_prompt: bool = True,
-        continue_final_message: bool = False,
-        tools: list[dict[str, Any]] | None = None,
-        tokenization_kwargs: dict[str, Any] | None = None,
-        mm_processor_kwargs: dict[str, Any] | None = None,
-    ) -> EngineInput:
-        (engine_input,) = self._preprocess_chat(
-            [conversation],
-            chat_template=chat_template,
-            chat_template_content_format=chat_template_content_format,
-            chat_template_kwargs=chat_template_kwargs,
-            add_generation_prompt=add_generation_prompt,
-            continue_final_message=continue_final_message,
-            tools=tools,
-            tokenization_kwargs=tokenization_kwargs,
-            mm_processor_kwargs=mm_processor_kwargs,
-        )
-
-        return engine_input
-
     def chat(
         self,
         messages: list[ChatCompletionMessageParam]
@@ -1058,392 +855,6 @@ class LLM(BeamSearchOfflineMixin, PoolingOfflineMixin):
         """
         return self.llm_engine.get_metrics()
 
-    def _params_to_seq(
-        self,
-        params: _P | Sequence[_P],
-        num_requests: int,
-    ) -> Sequence[_P]:
-        if isinstance(params, Sequence):
-            if len(params) != num_requests:
-                raise ValueError(
-                    f"The lengths of prompts ({num_requests}) "
-                    f"and params ({len(params)}) must be the same."
-                )
-
-            return params
-
-        return [params] * num_requests
-
-    def _lora_request_to_seq(
-        self,
-        lora_request: LoRARequest | None | Sequence[LoRARequest | None],
-        num_requests: int,
-    ) -> Sequence[LoRARequest | None]:
-        if isinstance(lora_request, Sequence):
-            if len(lora_request) != num_requests:
-                raise ValueError(
-                    f"The lengths of prompts ({num_requests}) "
-                    f"and lora_request ({len(lora_request)}) must be the same."
-                )
-
-            return lora_request
-
-        return [lora_request] * num_requests
-
-    def _priority_to_seq(
-        self,
-        priority: list[int] | None,
-        num_requests: int,
-    ) -> Sequence[int]:
-        if priority is not None:
-            if len(priority) != num_requests:
-                raise ValueError(
-                    f"The lengths of prompts ({num_requests}) "
-                    f"and priority ({len(priority)}) must be the same."
-                )
-
-            return priority
-
-        return [0] * num_requests
-
-    def _add_completion_requests(
-        self,
-        prompts: PromptType | Sequence[PromptType],
-        params: SamplingParams
-        | PoolingParams
-        | Sequence[SamplingParams | PoolingParams],
-        *,
-        use_tqdm: bool | Callable[..., tqdm] = True,
-        lora_request: Sequence[LoRARequest] | LoRARequest | None = None,
-        priority: list[int] | None = None,
-        tokenization_kwargs: dict[str, Any] | None = None,
-        mm_processor_kwargs: dict[str, Any] | None = None,
-    ) -> list[str]:
-        seq_prompts = prompt_to_seq(prompts)
-        seq_params = self._params_to_seq(params, len(seq_prompts))
-        seq_lora_requests = self._lora_request_to_seq(lora_request, len(seq_prompts))
-        seq_priority = self._priority_to_seq(priority, len(seq_prompts))
-
-        return self._render_and_add_requests(
-            prompts=(
-                self._preprocess_cmpl_one(
-                    prompt,
-                    tokenization_kwargs,
-                    mm_processor_kwargs=mm_processor_kwargs,
-                )
-                for prompt in maybe_tqdm(
-                    seq_prompts,
-                    use_tqdm=use_tqdm,
-                    desc="Rendering prompts",
-                )
-            ),
-            params=seq_params,
-            lora_requests=seq_lora_requests,
-            priorities=seq_priority,
-        )
-
-    def _run_completion(
-        self,
-        prompts: PromptType | Sequence[PromptType],
-        params: SamplingParams
-        | PoolingParams
-        | Sequence[SamplingParams | PoolingParams],
-        output_type: type[_O],
-        *,
-        use_tqdm: bool | Callable[..., tqdm] = True,
-        lora_request: Sequence[LoRARequest] | LoRARequest | None = None,
-        priority: list[int] | None = None,
-        tokenization_kwargs: dict[str, Any] | None = None,
-        mm_processor_kwargs: dict[str, Any] | None = None,
-    ):
-        self._add_completion_requests(
-            prompts=prompts,
-            params=params,
-            use_tqdm=use_tqdm,
-            lora_request=lora_request,
-            priority=priority,
-            tokenization_kwargs=tokenization_kwargs,
-            mm_processor_kwargs=mm_processor_kwargs,
-        )
-        return self._run_engine(use_tqdm=use_tqdm, output_type=output_type)
-
-    def _run_chat(
-        self,
-        messages: list[ChatCompletionMessageParam]
-        | Sequence[list[ChatCompletionMessageParam]],
-        params: SamplingParams
-        | PoolingParams
-        | Sequence[SamplingParams | PoolingParams],
-        output_type: type[_O],
-        *,
-        use_tqdm: bool | Callable[..., tqdm] = True,
-        lora_request: Sequence[LoRARequest] | LoRARequest | None = None,
-        chat_template: str | None = None,
-        chat_template_content_format: ChatTemplateContentFormatOption = "auto",
-        add_generation_prompt: bool = True,
-        continue_final_message: bool = False,
-        tools: list[dict[str, Any]] | None = None,
-        chat_template_kwargs: dict[str, Any] | None = None,
-        tokenization_kwargs: dict[str, Any] | None = None,
-        mm_processor_kwargs: dict[str, Any] | None = None,
-    ):
-        self._add_chat_requests(
-            messages=messages,
-            params=params,
-            use_tqdm=use_tqdm,
-            lora_request=lora_request,
-            chat_template=chat_template,
-            chat_template_content_format=chat_template_content_format,
-            chat_template_kwargs=chat_template_kwargs,
-            add_generation_prompt=add_generation_prompt,
-            continue_final_message=continue_final_message,
-            tools=tools,
-            tokenization_kwargs=tokenization_kwargs,
-            mm_processor_kwargs=mm_processor_kwargs,
-        )
-        return self._run_engine(output_type=output_type, use_tqdm=use_tqdm)
-
-    def _add_chat_requests(
-        self,
-        messages: list[ChatCompletionMessageParam]
-        | Sequence[list[ChatCompletionMessageParam]],
-        params: SamplingParams
-        | PoolingParams
-        | Sequence[SamplingParams | PoolingParams],
-        *,
-        use_tqdm: bool | Callable[..., tqdm] = True,
-        lora_request: Sequence[LoRARequest] | LoRARequest | None = None,
-        priority: list[int] | None = None,
-        chat_template: str | None = None,
-        chat_template_content_format: ChatTemplateContentFormatOption = "auto",
-        add_generation_prompt: bool = True,
-        continue_final_message: bool = False,
-        tools: list[dict[str, Any]] | None = None,
-        chat_template_kwargs: dict[str, Any] | None = None,
-        tokenization_kwargs: dict[str, Any] | None = None,
-        mm_processor_kwargs: dict[str, Any] | None = None,
-    ) -> list[str]:
-        seq_convs = conversation_to_seq(messages)
-        seq_params = self._params_to_seq(params, len(seq_convs))
-        seq_lora_requests = self._lora_request_to_seq(lora_request, len(seq_convs))
-        seq_priority = self._priority_to_seq(priority, len(seq_convs))
-
-        # When thinking is enabled or tools are provided, and the model
-        # uses special tokens for structured output (e.g. Gemma4's
-        # <|channel>, <|tool_call>, <|"|>), automatically set
-        # skip_special_tokens=False so these tokens are preserved in
-        # output.text for downstream parsing.
-        needs_parsing = (
-            chat_template_kwargs and chat_template_kwargs.get("enable_thinking")
-        ) or tools
-        if needs_parsing:
-            self._adjust_params_for_parsing(seq_params)
-
-        return self._render_and_add_requests(
-            prompts=(
-                self._preprocess_chat_one(
-                    conversation,
-                    chat_template=chat_template,
-                    chat_template_content_format=chat_template_content_format,
-                    chat_template_kwargs=chat_template_kwargs,
-                    add_generation_prompt=add_generation_prompt,
-                    continue_final_message=continue_final_message,
-                    tools=tools,
-                    tokenization_kwargs=tokenization_kwargs,
-                    mm_processor_kwargs=mm_processor_kwargs,
-                )
-                for conversation in maybe_tqdm(
-                    seq_convs,
-                    use_tqdm=use_tqdm,
-                    desc="Rendering conversations",
-                )
-            ),
-            params=seq_params,
-            lora_requests=seq_lora_requests,
-            priorities=seq_priority,
-        )
-
-    def _adjust_params_for_parsing(
-        self, params: Sequence[SamplingParams | PoolingParams]
-    ) -> None:
-        """Set ``skip_special_tokens=False`` when the model encodes
-        structured output syntax as special tokens.
-
-        Models like Gemma4 register thinking delimiters
-        (``<|channel>``/``<channel|>``) and tool call tokens
-        (``<|tool_call>``/``<tool_call|>``/``<|"|>``) as special tokens.
-        The default ``skip_special_tokens=True`` strips them from
-        ``output.text``, breaking parsing of both reasoning blocks and
-        tool calls.
-
-        This is a no-op for models whose structured tokens are regular
-        text tokens (e.g. DeepSeek's ``<think>``/``</think>``).
-        """
-        # The offline API currently lacks a unified rendering pipeline.
-        # Until the planned Renderer refactor is complete, we hardcode
-        # this token preservation logic specifically for Gemma4 models
-        # to avoid regressions on other models.
-        hf_config = getattr(self.model_config, "hf_config", None)
-        architectures = getattr(hf_config, "architectures", [])
-
-        if any("Gemma4" in arch for arch in architectures):
-            tokenizer = self.renderer.get_tokenizer()
-            vocab = tokenizer.get_vocab()
-            special_ids = set(getattr(tokenizer, "all_special_ids", []))
-
-            # Tokens used for thinking delimiters and tool call syntax
-            # that some models (Gemma4) register as special tokens.
-            structured_tokens = (
-                "<|channel>",
-                "<channel|>",  # thinking delimiters
-                "<|tool_call>",
-                "<tool_call|>",  # tool call delimiters
-                '<|"|>',  # string quoting in tool args
-            )
-            needs_special = any(
-                vocab.get(tok) in special_ids
-                for tok in structured_tokens
-                if tok in vocab
-            )
-            if needs_special:
-                for sp in params:
-                    if isinstance(sp, SamplingParams) and sp.skip_special_tokens:
-                        sp.skip_special_tokens = False
-
-    def _render_and_run_requests(
-        self,
-        prompts: Iterable[EngineInput],
-        params: Sequence[SamplingParams | PoolingParams],
-        output_type: type[_O],
-        *,
-        lora_requests: Sequence[LoRARequest | None] | None = None,
-        priorities: Sequence[int] | None = None,
-        use_tqdm: bool | Callable[..., tqdm] = True,
-    ):
-        if isinstance(prompts, (list, tuple)):
-            logger.warning_once(
-                "Rendering all prompts before adding them to the engine "
-                "is less efficient than performing both on the same prompt "
-                "before processing the next prompt. You should instead pass "
-                "a generator that renders one prompt per iteration, as that allows "
-                "engine execution to begin for the first prompt while processing "
-                "the next prompt."
-            )
-
-        self._render_and_add_requests(
-            prompts=prompts,
-            params=params,
-            lora_requests=lora_requests,
-            priorities=priorities,
-        )
-
-        return self._run_engine(output_type, use_tqdm=use_tqdm)
-
-    def _render_and_add_requests(
-        self,
-        prompts: Iterable[EngineInput],
-        params: Sequence[SamplingParams | PoolingParams],
-        *,
-        lora_requests: Sequence[LoRARequest | None] | None = None,
-        priorities: Sequence[int] | None = None,
-    ) -> list[str]:
-        added_request_ids: list[str] = []
-
-        try:
-            for i, prompt in enumerate(prompts):
-                request_id = self._add_request(
-                    prompt,
-                    params[i],
-                    lora_request=self._resolve_mm_lora(
-                        prompt,
-                        None if lora_requests is None else lora_requests[i],
-                    ),
-                    priority=0 if priorities is None else priorities[i],
-                )
-                added_request_ids.append(request_id)
-        except Exception as e:
-            if added_request_ids:
-                self.llm_engine.abort_request(added_request_ids, internal=True)
-            raise e
-
-        return added_request_ids
-
-    def _add_request(
-        self,
-        prompt: EngineInput,
-        params: SamplingParams | PoolingParams,
-        lora_request: LoRARequest | None = None,
-        priority: int = 0,
-    ) -> str:
-        if isinstance(params, SamplingParams):
-            # We only care about the final output
-            params.output_kind = RequestOutputKind.FINAL_ONLY
-
-        request_id = str(next(self.request_counter))
-
-        return self.llm_engine.add_request(
-            request_id,
-            prompt,
-            params,
-            lora_request=lora_request,
-            priority=priority,
-        )
-
-    def _run_engine(
-        self,
-        output_type: type[_O] | tuple[type[_O], ...],
-        *,
-        use_tqdm: bool | Callable[..., tqdm] = True,
-    ) -> list[_O]:
-        # Initialize tqdm.
-        if use_tqdm:
-            num_requests = self.llm_engine.get_num_unfinished_requests()
-            tqdm_func = use_tqdm if callable(use_tqdm) else tqdm
-            pbar = tqdm_func(
-                total=num_requests,
-                desc="Processed prompts",
-                dynamic_ncols=True,
-                postfix=(f"est. speed input: {0:.2f} toks/s, output: {0:.2f} toks/s"),
-            )
-
-        # Run the engine.
-        outputs: list[_O] = []
-        total_in_toks = 0
-        total_out_toks = 0
-        while self.llm_engine.has_unfinished_requests():
-            step_outputs = self.llm_engine.step()
-            for output in step_outputs:
-                assert isinstance(output, output_type)
-                if output.finished:
-                    outputs.append(output)  # type: ignore[arg-type]
-                    if use_tqdm:
-                        if isinstance(output, RequestOutput):
-                            # Calculate tokens only for RequestOutput
-                            n = len(output.outputs)
-                            assert output.prompt_token_ids is not None
-                            total_in_toks += len(output.prompt_token_ids) * n
-                            in_spd = total_in_toks / pbar.format_dict["elapsed"]
-                            total_out_toks += sum(
-                                len(stp.token_ids) for stp in output.outputs
-                            )
-                            out_spd = total_out_toks / pbar.format_dict["elapsed"]
-                            pbar.postfix = (
-                                f"est. speed input: {in_spd:.2f} toks/s, "
-                                f"output: {out_spd:.2f} toks/s"
-                            )
-                            pbar.update(n)
-                        else:
-                            pbar.update(1)
-                        if pbar.n == num_requests:
-                            pbar.refresh()
-
-        if use_tqdm:
-            pbar.close()
-        # Sort the outputs by request ID.
-        # This is necessary because some requests may be finished earlier than
-        # its previous requests.
-        return sorted(outputs, key=lambda x: int(x.request_id))
-
     def init_weight_transfer_engine(
         self, request: WeightTransferInitRequest | dict
     ) -> None:

vllm/entrypoints/offline_utils.py

@@ -0,0 +1,626 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from collections.abc import Callable, Iterable, Sequence
+from typing import Any
+
+from tqdm import tqdm
+from typing_extensions import TypeVar
+
+from vllm import (
+    PoolingParams,
+    PoolingRequestOutput,
+    PromptType,
+    RequestOutput,
+    SamplingParams,
+)
+from vllm.config import ModelConfig
+from vllm.entrypoints.chat_utils import (
+    ChatCompletionMessageParam,
+    ChatTemplateContentFormatOption,
+)
+from vllm.inputs import EngineInput
+from vllm.logger import init_logger
+from vllm.lora.request import LoRARequest
+from vllm.renderers import BaseRenderer, ChatParams, merge_kwargs
+from vllm.renderers.inputs.preprocess import (
+    conversation_to_seq,
+    parse_model_prompt,
+    prompt_to_seq,
+)
+from vllm.sampling_params import RequestOutputKind
+from vllm.utils.counter import Counter
+from vllm.utils.mistral import is_mistral_tokenizer
+from vllm.utils.tqdm_utils import maybe_tqdm
+from vllm.v1.engine.llm_engine import LLMEngine
+
+logger = init_logger(__name__)
+
+
+_P = TypeVar("_P", bound=SamplingParams | PoolingParams | None)
+_O = TypeVar(
+    "_O",
+    bound=RequestOutput | PoolingRequestOutput,
+    default=RequestOutput | PoolingRequestOutput,
+)
+_R = TypeVar("_R", default=Any)
+
+
+class OfflineInferenceMixin:
+    """Offline inference utils"""
+
+    request_counter: Counter
+    renderer: BaseRenderer
+    llm_engine: "LLMEngine"
+    model_config: ModelConfig
+
+    def _resolve_mm_lora(
+        self,
+        prompt: EngineInput,
+        lora_request: LoRARequest | None,
+    ) -> LoRARequest | None:
+        if prompt["type"] != "multimodal":
+            return lora_request
+
+        lora_config = self.llm_engine.vllm_config.lora_config
+        default_mm_loras = None if lora_config is None else lora_config.default_mm_loras
+        if not default_mm_loras:
+            return lora_request
+
+        prompt_modalities = prompt["mm_placeholders"].keys()
+        intersection = set(prompt_modalities).intersection(default_mm_loras.keys())
+        if not intersection:
+            return lora_request
+
+        if len(intersection) > 1:
+            # TODO: Would be nice to be able to have multiple loras per prompt
+            logger.warning(
+                "Multiple modality specific loras were registered and would be "
+                "used by a single prompt consuming several modalities; "
+                "currently we only support one lora per request; as such, "
+                "lora(s) registered with modalities: %s will be skipped",
+                intersection,
+            )
+            return lora_request
+
+        # Build the LoRA request; the ID of the default mm lora is the
+        # index of the modality name sorted alphabetically + 1.
+        modality_name = intersection.pop()
+        modality_lora_path = default_mm_loras[modality_name]
+        modality_lora_id = sorted(default_mm_loras).index(modality_name) + 1
+
+        # If we have a collision, warn if there is a collision,
+        # but always send the explicitly provided request.
+        if lora_request:
+            if lora_request.lora_int_id != modality_lora_id:
+                logger.warning(
+                    "A modality with a registered lora and a lora_request "
+                    "with a different ID were provided; falling back to the "
+                    "lora_request as we only apply one LoRARequest per prompt"
+                )
+            return lora_request
+
+        return LoRARequest(
+            modality_name,
+            modality_lora_id,
+            modality_lora_path,
+        )
+
+    def _preprocess_cmpl(
+        self,
+        prompts: Sequence[PromptType],
+        tokenization_kwargs: dict[str, Any] | None = None,
+        mm_processor_kwargs: dict[str, Any] | None = None,
+    ) -> Sequence[EngineInput]:
+        """
+        Convert prompt inputs from LLM APIs (other than [LLM.chat][]) into
+        a format that can be passed to `_add_request`.
+
+        Refer to [LLM.generate][] for a complete description of the arguments.
+
+        Returns:
+            A list of `EngineInput` objects ready to be passed into LLMEngine.
+        """
+        renderer = self.renderer
+        model_config = self.model_config
+
+        parsed_prompts = [
+            parse_model_prompt(model_config, prompt) for prompt in prompts
+        ]
+        tok_params = renderer.default_cmpl_tok_params.with_kwargs(
+            **(tokenization_kwargs or {})
+        )
+        prompt_extras = (
+            None
+            if mm_processor_kwargs is None
+            else {"mm_processor_kwargs": mm_processor_kwargs}
+        )
+
+        return renderer.render_cmpl(
+            parsed_prompts,
+            tok_params,
+            prompt_extras=prompt_extras,
+        )
+
+    def _preprocess_cmpl_one(
+        self,
+        prompt: PromptType,
+        tokenization_kwargs: dict[str, Any] | None = None,
+        mm_processor_kwargs: dict[str, Any] | None = None,
+    ) -> EngineInput:
+        (engine_input,) = self._preprocess_cmpl(
+            [prompt],
+            tokenization_kwargs,
+            mm_processor_kwargs=mm_processor_kwargs,
+        )
+        return engine_input
+
+    def _preprocess_chat(
+        self,
+        conversations: Sequence[list[ChatCompletionMessageParam]],
+        chat_template: str | None = None,
+        chat_template_content_format: ChatTemplateContentFormatOption = "auto",
+        chat_template_kwargs: dict[str, Any] | None = None,
+        add_generation_prompt: bool = True,
+        continue_final_message: bool = False,
+        tools: list[dict[str, Any]] | None = None,
+        tokenization_kwargs: dict[str, Any] | None = None,
+        mm_processor_kwargs: dict[str, Any] | None = None,
+    ) -> Sequence[EngineInput]:
+        """
+        Convert a list of conversations into prompts so that they can then
+        be used as input for other LLM APIs.
+
+        Refer to [LLM.chat][] for a complete description of the arguments.
+
+        Returns:
+            A list of `EngineInput` objects ready to be passed into LLMEngine.
+        """
+        renderer = self.renderer
+
+        chat_params = ChatParams(
+            chat_template=chat_template,
+            chat_template_content_format=chat_template_content_format,
+            chat_template_kwargs=merge_kwargs(
+                chat_template_kwargs,
+                dict(
+                    add_generation_prompt=add_generation_prompt,
+                    continue_final_message=continue_final_message,
+                    tools=tools,
+                    tokenize=(
+                        is_mistral_tokenizer(renderer.tokenizer)
+                        or self.model_config.enable_prompt_embeds
+                    ),
+                ),
+            ),
+            mm_processor_kwargs=mm_processor_kwargs,
+        )
+        tok_params = renderer.default_chat_tok_params.with_kwargs(
+            **(tokenization_kwargs or {})
+        )
+        prompt_extras = (
+            None
+            if mm_processor_kwargs is None
+            else {"mm_processor_kwargs": mm_processor_kwargs}
+        )
+
+        _, engine_inputs = renderer.render_chat(
+            conversations,
+            chat_params,
+            tok_params,
+            prompt_extras=prompt_extras,
+        )
+
+        return engine_inputs
+
+    def _preprocess_chat_one(
+        self,
+        conversation: list[ChatCompletionMessageParam],
+        chat_template: str | None = None,
+        chat_template_content_format: ChatTemplateContentFormatOption = "auto",
+        chat_template_kwargs: dict[str, Any] | None = None,
+        add_generation_prompt: bool = True,
+        continue_final_message: bool = False,
+        tools: list[dict[str, Any]] | None = None,
+        tokenization_kwargs: dict[str, Any] | None = None,
+        mm_processor_kwargs: dict[str, Any] | None = None,
+    ) -> EngineInput:
+        (engine_input,) = self._preprocess_chat(
+            [conversation],
+            chat_template=chat_template,
+            chat_template_content_format=chat_template_content_format,
+            chat_template_kwargs=chat_template_kwargs,
+            add_generation_prompt=add_generation_prompt,
+            continue_final_message=continue_final_message,
+            tools=tools,
+            tokenization_kwargs=tokenization_kwargs,
+            mm_processor_kwargs=mm_processor_kwargs,
+        )
+
+        return engine_input
+
+    def _params_to_seq(
+        self,
+        params: _P | Sequence[_P],
+        num_requests: int,
+    ) -> Sequence[_P]:
+        if isinstance(params, Sequence):
+            if len(params) != num_requests:
+                raise ValueError(
+                    f"The lengths of prompts ({num_requests}) "
+                    f"and params ({len(params)}) must be the same."
+                )
+
+            return params
+
+        return [params] * num_requests
+
+    def _lora_request_to_seq(
+        self,
+        lora_request: LoRARequest | None | Sequence[LoRARequest | None],
+        num_requests: int,
+    ) -> Sequence[LoRARequest | None]:
+        if isinstance(lora_request, Sequence):
+            if len(lora_request) != num_requests:
+                raise ValueError(
+                    f"The lengths of prompts ({num_requests}) "
+                    f"and lora_request ({len(lora_request)}) must be the same."
+                )
+
+            return lora_request
+
+        return [lora_request] * num_requests
+
+    def _priority_to_seq(
+        self,
+        priority: list[int] | None,
+        num_requests: int,
+    ) -> Sequence[int]:
+        if priority is not None:
+            if len(priority) != num_requests:
+                raise ValueError(
+                    f"The lengths of prompts ({num_requests}) "
+                    f"and priority ({len(priority)}) must be the same."
+                )
+
+            return priority
+
+        return [0] * num_requests
+
+    def _add_completion_requests(
+        self,
+        prompts: PromptType | Sequence[PromptType],
+        params: SamplingParams
+        | PoolingParams
+        | Sequence[SamplingParams | PoolingParams],
+        *,
+        use_tqdm: bool | Callable[..., tqdm] = True,
+        lora_request: Sequence[LoRARequest] | LoRARequest | None = None,
+        priority: list[int] | None = None,
+        tokenization_kwargs: dict[str, Any] | None = None,
+        mm_processor_kwargs: dict[str, Any] | None = None,
+    ) -> list[str]:
+        seq_prompts = prompt_to_seq(prompts)
+        seq_params = self._params_to_seq(params, len(seq_prompts))
+        seq_lora_requests = self._lora_request_to_seq(lora_request, len(seq_prompts))
+        seq_priority = self._priority_to_seq(priority, len(seq_prompts))
+
+        return self._render_and_add_requests(
+            prompts=(
+                self._preprocess_cmpl_one(
+                    prompt,
+                    tokenization_kwargs,
+                    mm_processor_kwargs=mm_processor_kwargs,
+                )
+                for prompt in maybe_tqdm(
+                    seq_prompts,
+                    use_tqdm=use_tqdm,
+                    desc="Rendering prompts",
+                )
+            ),
+            params=seq_params,
+            lora_requests=seq_lora_requests,
+            priorities=seq_priority,
+        )
+
+    def _run_completion(
+        self,
+        prompts: PromptType | Sequence[PromptType],
+        params: SamplingParams
+        | PoolingParams
+        | Sequence[SamplingParams | PoolingParams],
+        output_type: type[_O],
+        *,
+        use_tqdm: bool | Callable[..., tqdm] = True,
+        lora_request: Sequence[LoRARequest] | LoRARequest | None = None,
+        priority: list[int] | None = None,
+        tokenization_kwargs: dict[str, Any] | None = None,
+        mm_processor_kwargs: dict[str, Any] | None = None,
+    ):
+        self._add_completion_requests(
+            prompts=prompts,
+            params=params,
+            use_tqdm=use_tqdm,
+            lora_request=lora_request,
+            priority=priority,
+            tokenization_kwargs=tokenization_kwargs,
+            mm_processor_kwargs=mm_processor_kwargs,
+        )
+        return self._run_engine(use_tqdm=use_tqdm, output_type=output_type)
+
+    def _run_chat(
+        self,
+        messages: list[ChatCompletionMessageParam]
+        | Sequence[list[ChatCompletionMessageParam]],
+        params: SamplingParams
+        | PoolingParams
+        | Sequence[SamplingParams | PoolingParams],
+        output_type: type[_O],
+        *,
+        use_tqdm: bool | Callable[..., tqdm] = True,
+        lora_request: Sequence[LoRARequest] | LoRARequest | None = None,
+        chat_template: str | None = None,
+        chat_template_content_format: ChatTemplateContentFormatOption = "auto",
+        add_generation_prompt: bool = True,
+        continue_final_message: bool = False,
+        tools: list[dict[str, Any]] | None = None,
+        chat_template_kwargs: dict[str, Any] | None = None,
+        tokenization_kwargs: dict[str, Any] | None = None,
+        mm_processor_kwargs: dict[str, Any] | None = None,
+    ):
+        self._add_chat_requests(
+            messages=messages,
+            params=params,
+            use_tqdm=use_tqdm,
+            lora_request=lora_request,
+            chat_template=chat_template,
+            chat_template_content_format=chat_template_content_format,
+            chat_template_kwargs=chat_template_kwargs,
+            add_generation_prompt=add_generation_prompt,
+            continue_final_message=continue_final_message,
+            tools=tools,
+            tokenization_kwargs=tokenization_kwargs,
+            mm_processor_kwargs=mm_processor_kwargs,
+        )
+        return self._run_engine(output_type=output_type, use_tqdm=use_tqdm)
+
+    def _add_chat_requests(
+        self,
+        messages: list[ChatCompletionMessageParam]
+        | Sequence[list[ChatCompletionMessageParam]],
+        params: SamplingParams
+        | PoolingParams
+        | Sequence[SamplingParams | PoolingParams],
+        *,
+        use_tqdm: bool | Callable[..., tqdm] = True,
+        lora_request: Sequence[LoRARequest] | LoRARequest | None = None,
+        priority: list[int] | None = None,
+        chat_template: str | None = None,
+        chat_template_content_format: ChatTemplateContentFormatOption = "auto",
+        add_generation_prompt: bool = True,
+        continue_final_message: bool = False,
+        tools: list[dict[str, Any]] | None = None,
+        chat_template_kwargs: dict[str, Any] | None = None,
+        tokenization_kwargs: dict[str, Any] | None = None,
+        mm_processor_kwargs: dict[str, Any] | None = None,
+    ) -> list[str]:
+        seq_convs = conversation_to_seq(messages)
+        seq_params = self._params_to_seq(params, len(seq_convs))
+        seq_lora_requests = self._lora_request_to_seq(lora_request, len(seq_convs))
+        seq_priority = self._priority_to_seq(priority, len(seq_convs))
+
+        # When thinking is enabled or tools are provided, and the model
+        # uses special tokens for structured output (e.g. Gemma4's
+        # <|channel>, <|tool_call>, <|"|>), automatically set
+        # skip_special_tokens=False so these tokens are preserved in
+        # output.text for downstream parsing.
+        needs_parsing = (
+            chat_template_kwargs and chat_template_kwargs.get("enable_thinking")
+        ) or tools
+        if needs_parsing:
+            self._adjust_params_for_parsing(seq_params)
+
+        return self._render_and_add_requests(
+            prompts=(
+                self._preprocess_chat_one(
+                    conversation,
+                    chat_template=chat_template,
+                    chat_template_content_format=chat_template_content_format,
+                    chat_template_kwargs=chat_template_kwargs,
+                    add_generation_prompt=add_generation_prompt,
+                    continue_final_message=continue_final_message,
+                    tools=tools,
+                    tokenization_kwargs=tokenization_kwargs,
+                    mm_processor_kwargs=mm_processor_kwargs,
+                )
+                for conversation in maybe_tqdm(
+                    seq_convs,
+                    use_tqdm=use_tqdm,
+                    desc="Rendering conversations",
+                )
+            ),
+            params=seq_params,
+            lora_requests=seq_lora_requests,
+            priorities=seq_priority,
+        )
+
+    def _adjust_params_for_parsing(
+        self, params: Sequence[SamplingParams | PoolingParams]
+    ) -> None:
+        """Set ``skip_special_tokens=False`` when the model encodes
+        structured output syntax as special tokens.
+
+        Models like Gemma4 register thinking delimiters
+        (``<|channel>``/``<channel|>``) and tool call tokens
+        (``<|tool_call>``/``<tool_call|>``/``<|"|>``) as special tokens.
+        The default ``skip_special_tokens=True`` strips them from
+        ``output.text``, breaking parsing of both reasoning blocks and
+        tool calls.
+
+        This is a no-op for models whose structured tokens are regular
+        text tokens (e.g. DeepSeek's ``<think>``/``</think>``).
+        """
+        # The offline API currently lacks a unified rendering pipeline.
+        # Until the planned Renderer refactor is complete, we hardcode
+        # this token preservation logic specifically for Gemma4 models
+        # to avoid regressions on other models.
+        hf_config = getattr(self.model_config, "hf_config", None)
+        architectures = getattr(hf_config, "architectures", [])
+
+        if any("Gemma4" in arch for arch in architectures):
+            tokenizer = self.renderer.get_tokenizer()
+            vocab = tokenizer.get_vocab()
+            special_ids = set(getattr(tokenizer, "all_special_ids", []))
+
+            # Tokens used for thinking delimiters and tool call syntax
+            # that some models (Gemma4) register as special tokens.
+            structured_tokens = (
+                "<|channel>",
+                "<channel|>",  # thinking delimiters
+                "<|tool_call>",
+                "<tool_call|>",  # tool call delimiters
+                '<|"|>',  # string quoting in tool args
+            )
+            needs_special = any(
+                vocab.get(tok) in special_ids
+                for tok in structured_tokens
+                if tok in vocab
+            )
+            if needs_special:
+                for sp in params:
+                    if isinstance(sp, SamplingParams) and sp.skip_special_tokens:
+                        sp.skip_special_tokens = False
+
+    def _render_and_run_requests(
+        self,
+        prompts: Iterable[EngineInput],
+        params: Sequence[SamplingParams | PoolingParams],
+        output_type: type[_O],
+        *,
+        lora_requests: Sequence[LoRARequest | None] | None = None,
+        priorities: Sequence[int] | None = None,
+        use_tqdm: bool | Callable[..., tqdm] = True,
+    ):
+        if isinstance(prompts, (list, tuple)):
+            logger.warning_once(
+                "Rendering all prompts before adding them to the engine "
+                "is less efficient than performing both on the same prompt "
+                "before processing the next prompt. You should instead pass "
+                "a generator that renders one prompt per iteration, as that allows "
+                "engine execution to begin for the first prompt while processing "
+                "the next prompt."
+            )
+
+        self._render_and_add_requests(
+            prompts=prompts,
+            params=params,
+            lora_requests=lora_requests,
+            priorities=priorities,
+        )
+
+        return self._run_engine(output_type, use_tqdm=use_tqdm)
+
+    def _render_and_add_requests(
+        self,
+        prompts: Iterable[EngineInput],
+        params: Sequence[SamplingParams | PoolingParams],
+        *,
+        lora_requests: Sequence[LoRARequest | None] | None = None,
+        priorities: Sequence[int] | None = None,
+    ) -> list[str]:
+        added_request_ids: list[str] = []
+
+        try:
+            for i, prompt in enumerate(prompts):
+                request_id = self._add_request(
+                    prompt,
+                    params[i],
+                    lora_request=self._resolve_mm_lora(
+                        prompt,
+                        None if lora_requests is None else lora_requests[i],
+                    ),
+                    priority=0 if priorities is None else priorities[i],
+                )
+                added_request_ids.append(request_id)
+        except Exception as e:
+            if added_request_ids:
+                self.llm_engine.abort_request(added_request_ids, internal=True)
+            raise e
+
+        return added_request_ids
+
+    def _add_request(
+        self,
+        prompt: EngineInput,
+        params: SamplingParams | PoolingParams,
+        lora_request: LoRARequest | None = None,
+        priority: int = 0,
+    ) -> str:
+        if isinstance(params, SamplingParams):
+            # We only care about the final output
+            params.output_kind = RequestOutputKind.FINAL_ONLY
+
+        request_id = str(next(self.request_counter))
+
+        return self.llm_engine.add_request(
+            request_id,
+            prompt,
+            params,
+            lora_request=lora_request,
+            priority=priority,
+        )
+
+    def _run_engine(
+        self,
+        output_type: type[_O] | tuple[type[_O], ...],
+        *,
+        use_tqdm: bool | Callable[..., tqdm] = True,
+    ) -> list[_O]:
+        # Initialize tqdm.
+        if use_tqdm:
+            num_requests = self.llm_engine.get_num_unfinished_requests()
+            tqdm_func = use_tqdm if callable(use_tqdm) else tqdm
+            pbar = tqdm_func(
+                total=num_requests,
+                desc="Processed prompts",
+                dynamic_ncols=True,
+                postfix=(f"est. speed input: {0:.2f} toks/s, output: {0:.2f} toks/s"),
+            )
+
+        # Run the engine.
+        outputs: list[_O] = []
+        total_in_toks = 0
+        total_out_toks = 0
+        while self.llm_engine.has_unfinished_requests():
+            step_outputs = self.llm_engine.step()
+            for output in step_outputs:
+                assert isinstance(output, output_type)
+                if output.finished:
+                    outputs.append(output)  # type: ignore[arg-type]
+                    if use_tqdm:
+                        if isinstance(output, RequestOutput):
+                            # Calculate tokens only for RequestOutput
+                            n = len(output.outputs)
+                            assert output.prompt_token_ids is not None
+                            total_in_toks += len(output.prompt_token_ids) * n
+                            in_spd = total_in_toks / pbar.format_dict["elapsed"]
+                            total_out_toks += sum(
+                                len(stp.token_ids) for stp in output.outputs
+                            )
+                            out_spd = total_out_toks / pbar.format_dict["elapsed"]
+                            pbar.postfix = (
+                                f"est. speed input: {in_spd:.2f} toks/s, "
+                                f"output: {out_spd:.2f} toks/s"
+                            )
+                            pbar.update(n)
+                        else:
+                            pbar.update(1)
+                        if pbar.n == num_requests:
+                            pbar.refresh()
+
+        if use_tqdm:
+            pbar.close()
+        # Sort the outputs by request ID.
+        # This is necessary because some requests may be finished earlier than
+        # its previous requests.
+        return sorted(outputs, key=lambda x: int(x.request_id))

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/entrypoints/pooling/offline.py

@@ -1,34 +1,24 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
-from abc import ABC, abstractmethod
-from collections.abc import Callable, Iterable, Sequence
+from collections.abc import Callable, Sequence
 from typing import Any
 
 from tqdm.auto import tqdm
-from typing_extensions import TypeVar
 
-from vllm.config import ModelConfig
 from vllm.entrypoints.chat_utils import ChatTemplateConfig
-from vllm.inputs import (
-    DataPrompt,
-    EngineInput,
-    PromptType,
-)
+from vllm.entrypoints.offline_utils import OfflineInferenceMixin
+from vllm.inputs import DataPrompt, PromptType
 from vllm.logger import init_logger
 from vllm.lora.request import LoRARequest
 from vllm.outputs import (
     ClassificationRequestOutput,
     EmbeddingRequestOutput,
     PoolingRequestOutput,
-    RequestOutput,
     ScoringRequestOutput,
 )
 from vllm.pooling_params import PoolingParams
-from vllm.renderers import BaseRenderer
-from vllm.sampling_params import SamplingParams
 from vllm.tasks import SCORE_TYPE_MAP, PoolingTask, SupportedTask
-from vllm.v1.engine.llm_engine import LLMEngine
 
 from .factories import init_pooling_io_processors
 from .scoring.io_processor import ScoringIOProcessor
@@ -37,20 +27,10 @@ from .typing import OfflineInputsContext, OfflineOutputsContext
 
 logger = init_logger(__name__)
 
-_P = TypeVar("_P", bound=SamplingParams | PoolingParams | None)
-_O = TypeVar(
-    "_O",
-    bound=RequestOutput | PoolingRequestOutput,
-    default=RequestOutput | PoolingRequestOutput,
-)
 
-
-class PoolingOfflineMixin(ABC):
+class PoolingOfflineMixin(OfflineInferenceMixin):
     """Offline inference for pooling models"""
 
-    renderer: BaseRenderer
-    llm_engine: "LLMEngine"
-    model_config: ModelConfig
     runner_type: str
     chat_template: str | None
     supported_tasks: tuple[SupportedTask, ...]
@@ -464,47 +444,3 @@ class PoolingOfflineMixin(ABC):
         )
 
         return [ScoringRequestOutput.from_base(item) for item in outputs]
-
-    @abstractmethod
-    def _params_to_seq(
-        self,
-        params: _P | Sequence[_P],
-        num_requests: int,
-    ) -> Sequence[_P]:
-        raise NotImplementedError
-
-    @abstractmethod
-    def _lora_request_to_seq(
-        self,
-        lora_request: LoRARequest | None | Sequence[LoRARequest | None],
-        num_requests: int,
-    ) -> Sequence[LoRARequest | None]:
-        raise NotImplementedError
-
-    @abstractmethod
-    def _priority_to_seq(
-        self,
-        priority: list[int] | None,
-        num_requests: int,
-    ) -> Sequence[int]:
-        raise NotImplementedError
-
-    @abstractmethod
-    def _render_and_add_requests(
-        self,
-        prompts: Iterable[EngineInput],
-        params: Sequence[SamplingParams | PoolingParams],
-        *,
-        lora_requests: Sequence[LoRARequest | None] | None = None,
-        priorities: Sequence[int] | None = None,
-    ) -> list[str]:
-        raise NotImplementedError
-
-    @abstractmethod
-    def _run_engine(
-        self,
-        output_type: type[_O] | tuple[type[_O], ...],
-        *,
-        use_tqdm: bool | Callable[..., tqdm] = True,
-    ) -> list[_O]:
-        raise NotImplementedError

vllm/envs.py

@@ -1971,6 +1971,8 @@ environment_variables: dict[str, Callable[[], Any]] = {
         int(os.getenv("VLLM_USE_SIMPLE_KV_OFFLOAD", "0"))
     ),
     # Whether to enable dual cuda streams for LoRA computation
+    # (used by both BaseLinearLayerWithLoRA and FusedMoEWithLoRA to
+    # overlap the base layer compute with the LoRA fast path).
     "VLLM_LORA_ENABLE_DUAL_STREAM": lambda: bool(
         int(os.getenv("VLLM_LORA_ENABLE_DUAL_STREAM", "0"))
     ),

vllm/lora/layers/base_linear.py

@@ -25,16 +25,9 @@ from vllm.utils.multi_stream_utils import maybe_execute_in_parallel
 from vllm.utils.torch_utils import direct_register_custom_op
 
 from .base import BaseLayerWithLoRA
-from .utils import _get_lora_device
+from .utils import _get_lora_aux_cuda_stream, _get_lora_device
 
 if envs.VLLM_LORA_ENABLE_DUAL_STREAM:
-    _lora_aux_cuda_stream: torch.cuda.Stream | None = None
-
-    def _get_lora_aux_cuda_stream() -> torch.cuda.Stream | None:
-        global _lora_aux_cuda_stream
-        if _lora_aux_cuda_stream is None and current_platform.is_cuda_alike():
-            _lora_aux_cuda_stream = torch.cuda.Stream()
-        return _lora_aux_cuda_stream
 
     def lora_linear_async(
         layer_name: str,

vllm/lora/layers/fused_moe.py

@@ -19,8 +19,9 @@ from vllm.model_executor.layers.fused_moe.modular_kernel import FusedMoEKernel
 from vllm.model_executor.layers.fused_moe.prepare_finalize import (
     MoEPrepareAndFinalizeNoDPEPModular,
 )
+from vllm.platforms import current_platform
 
-from .utils import _get_lora_device
+from .utils import _get_lora_aux_cuda_stream, _get_lora_device
 
 
 class FusedMoEWithLoRA(BaseLayerWithLoRA):
@@ -34,6 +35,9 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA):
         self.tp_size = self.base_layer.tp_size
         self.tp_rank = self.base_layer.tp_rank
         self.device = _get_lora_device(base_layer)
+
+        self._enable_aux_cuda_stream = envs.VLLM_LORA_ENABLE_DUAL_STREAM
+        self._init_lora_stream_context()
         # For non-gated MoE (is_act_and_mul=False), only 1 slice is needed
         # since there's only up_proj (w1), not gate_proj + up_proj (w1 + w3)
         self._w13_slices = 2 if base_layer.moe_config.is_act_and_mul else 1
@@ -65,7 +69,25 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA):
             FusedMoEModularMethod(self.base_layer.quant_method, moe_kernel)
         )
 
+    def _init_lora_stream_context(self) -> None:
+        self._lora_stream: torch.cuda.Stream | None = None
+        self._events: tuple[torch.cuda.Event, ...] | None = None
+        if not self._enable_aux_cuda_stream:
+            return
+        if not current_platform.is_cuda_alike():
+            return
+        self._lora_stream = _get_lora_aux_cuda_stream()
+        # 4 events: 2 per (base GEMM, LoRA) pair so w13 and w2 don't reuse
+        # the same event objects; reuse-within-a-pair is fine because the
+        # second pair starts only after intermediate_cache1.add_() has joined.
+        self._events = tuple(torch.cuda.Event() for _ in range(4))
+
     def _build_lora_context(self):
+        use_dual_stream = (
+            self._enable_aux_cuda_stream
+            and not self.fully_sharded
+            and self._lora_stream is not None
+        )
         return MoELoRAContext(
             w13_lora_a_stacked=self.w13_lora_a_stacked,
             w13_lora_b_stacked=self.w13_lora_b_stacked,
@@ -81,6 +103,8 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA):
             local_num_experts=self.base_layer.local_num_experts,
             punica_wrapper=self.punica_wrapper,
             use_tuned_config=bool(envs.VLLM_TUNED_CONFIG_FOLDER),
+            aux_stream=self._lora_stream if use_dual_stream else None,
+            events=self._events if use_dual_stream else None,
         )
 
     def _create_lora_a_weights(

vllm/lora/layers/utils.py

@@ -7,9 +7,22 @@ from enum import Enum
 import torch
 import torch.nn as nn
 
+from vllm import envs
 from vllm.model_executor.layers.fused_moe.fused_moe import try_get_optimal_moe_config
+from vllm.platforms import current_platform
 from vllm.utils.math_utils import next_power_of_2
 
+_lora_aux_cuda_stream: torch.cuda.Stream | None = None
+
+
+def _get_lora_aux_cuda_stream() -> torch.cuda.Stream | None:
+    if not envs.VLLM_LORA_ENABLE_DUAL_STREAM:
+        return None
+    global _lora_aux_cuda_stream
+    if _lora_aux_cuda_stream is None and current_platform.is_cuda_alike():
+        _lora_aux_cuda_stream = torch.cuda.Stream()
+    return _lora_aux_cuda_stream
+
 
 class LoRAMappingType(Enum):
     LANGUAGE = 1

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
@@ -105,6 +106,742 @@ def _get_c_ptrs(
 _LORA_PTR_DICT: dict[tuple[int, ...], torch.tensor] = {}
 
 
+# ---------------------------------------------------------------------------
+# Fully-fused MoE-LoRA kernel (one-shot): shrink + expand combined into a single
+# launch with the rank-dim intermediate kept in registers. Used by the fast
+# path of `_fused_moe_lora` for `fully_sharded=False`. The legacy two-kernel
+# path (`_fused_moe_lora_kernel` above) is retained for `fully_sharded=True`
+# because that path needs to materialise the intermediate cache for an
+# all_reduce / all_gather between shrink and expand.
+# ---------------------------------------------------------------------------
+
+
+@triton.heuristics({"EVEN_K": lambda args: args["K"] % args["BLOCK_K"] == 0})
+@triton.jit
+def _fused_moe_lora_one_shot_kernel(
+    # ---- pointers ----
+    x_ptr,
+    A_ptrs,
+    B_ptrs,
+    out_ptr,
+    topk_weights_ptr,
+    sorted_token_ids_ptr,
+    expert_ids_ptr,
+    num_tokens_post_padded_ptr,
+    token_lora_mapping_ptr,
+    lora_ids_ptr,
+    adapter_enabled_ptr,
+    # ---- dims ----
+    N,
+    K,
+    num_valid_tokens,
+    top_k_num,
+    max_loras,
+    # ---- strides ----
+    stride_xm,
+    stride_xk,
+    stride_A_lora,
+    stride_A_expert,
+    stride_A_r,
+    stride_A_k,
+    stride_B_lora,
+    stride_B_expert,
+    stride_B_n,
+    stride_B_r,
+    stride_om,
+    stride_on,
+    stride_tl_,
+    stride_el,
+    # ---- scalar ----
+    slice_n_offset,
+    # ---- constexpr (set per call) ----
+    token_mapping_factor: tl.constexpr,
+    naive_block_assignment: tl.constexpr,
+    MUL_ROUTED_WEIGHT: tl.constexpr,
+    BLOCK_M: tl.constexpr,
+    BLOCK_R: tl.constexpr,
+    actual_rank: tl.constexpr,
+    NPID_FACTOR: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    BLOCK_K: tl.constexpr,
+    EVEN_K: tl.constexpr,
+    ADD_INPUTS: tl.constexpr,
+):
+    pid_full = tl.program_id(axis=0)
+    pid_m = pid_full // NPID_FACTOR
+    pid_n_outer = pid_full % NPID_FACTOR
+    slice_id = tl.program_id(axis=1)
+    lora_idx = tl.program_id(axis=2)
+
+    # Resolve lora_id.
+    if naive_block_assignment:
+        token_idx_for_lora = pid_m // top_k_num
+        lora_id = tl.load(token_lora_mapping_ptr + token_idx_for_lora)
+    else:
+        lora_id = tl.load(lora_ids_ptr + lora_idx)
+    if lora_id < 0:
+        return
+    if lora_id >= max_loras:
+        return
+    enabled = tl.load(adapter_enabled_ptr + lora_id)
+    if enabled == 0:
+        return
+
+    if not naive_block_assignment:
+        ntpp = tl.load(num_tokens_post_padded_ptr + lora_id)
+        if pid_m * BLOCK_M >= ntpp:
+            return
+
+    # Resolve expert_id.
+    if naive_block_assignment:
+        expert_id = tl.load(expert_ids_ptr + pid_m)
+    else:
+        ind = lora_id * stride_el + pid_m
+        expert_id = tl.load(
+            expert_ids_ptr + ind, mask=ind < max_loras * stride_el, other=-1
+        )
+    if expert_id < 0:
+        return
+
+    # Compute offs_token (flat token ids).
+    offs = tl.arange(0, BLOCK_M).to(tl.int64)
+    if naive_block_assignment:
+        offs_token = tl.where(offs == 0, pid_m, num_valid_tokens)
+    else:
+        offs_token_id = pid_m * BLOCK_M + offs
+        token_ind = stride_tl_ * lora_id + offs_token_id
+        offs_token = tl.load(
+            sorted_token_ids_ptr + token_ind,
+            mask=token_ind < max_loras * stride_tl_,
+            other=num_valid_tokens,
+        )
+    token_mask = offs_token < num_valid_tokens
+
+    # N range owned by this program. Splitting [0, N) into NPID_FACTOR
+    # contiguous outer blocks lets us scale parallelism for small batches.
+    n_per_outer = tl.cdiv(N, NPID_FACTOR)
+    n_lo = pid_n_outer * n_per_outer
+    n_hi = tl.minimum((pid_n_outer + 1) * n_per_outer, N)
+    if n_lo >= N:
+        return
+
+    # Slice pointers.
+    cur_A_ptr = tl.load(A_ptrs + slice_id).to(tl.pointer_type(out_ptr.dtype.element_ty))
+    cur_B_ptr = tl.load(B_ptrs + slice_id).to(tl.pointer_type(out_ptr.dtype.element_ty))
+
+    A_base = cur_A_ptr + lora_id * stride_A_lora + expert_id * stride_A_expert
+    B_base = cur_B_ptr + lora_id * stride_B_lora + expert_id * stride_B_expert
+
+    # SHRINK: tmp[BLOCK_M, BLOCK_R] = x @ A^T, accumulated in fp32 registers.
+    offs_r = tl.arange(0, BLOCK_R)
+    rank_mask = offs_r < actual_rank
+    # Clamp rank offsets so OOB rows of A / B map to address 0; the mask
+    # zeros the loaded values. Required when BLOCK_R > actual_rank
+    # (e.g. rank=4 padded to 16) -- without clamping, tl.load would address
+    # the next expert's memory.
+    safe_offs_r = tl.where(rank_mask, offs_r, 0)
+    offs_k = tl.arange(0, BLOCK_K)
+
+    offs_x_row = offs_token // token_mapping_factor
+    x_ptrs = x_ptr + offs_x_row[:, None] * stride_xm + offs_k[None, :] * stride_xk
+    a_ptrs = A_base + offs_k[:, None] * stride_A_k + safe_offs_r[None, :] * stride_A_r
+
+    tmp = tl.zeros((BLOCK_M, BLOCK_R), dtype=tl.float32)
+    if EVEN_K:
+        for _ in range(0, K, BLOCK_K):
+            x = tl.load(x_ptrs, mask=token_mask[:, None], other=0.0)
+            a = tl.load(a_ptrs, mask=rank_mask[None, :], other=0.0)
+            tmp += tl.dot(x, a)
+            x_ptrs += BLOCK_K * stride_xk
+            a_ptrs += BLOCK_K * stride_A_k
+    else:
+        for kb in range(0, K, BLOCK_K):
+            k_remain = K - kb
+            k_mask = offs_k < k_remain
+            x = tl.load(x_ptrs, mask=token_mask[:, None] & k_mask[None, :], other=0.0)
+            a = tl.load(a_ptrs, mask=k_mask[:, None] & rank_mask[None, :], other=0.0)
+            tmp += tl.dot(x, a)
+            x_ptrs += BLOCK_K * stride_xk
+            a_ptrs += BLOCK_K * stride_A_k
+
+    tmp_typed = tmp.to(out_ptr.dtype.element_ty)
+
+    # EXPAND: out[tokens, n] += tmp @ B^T, looped over BLOCK_N tiles within
+    # this program's [n_lo, n_hi). The (offs_n < n_hi) mask is required
+    # whenever BLOCK_N > n_per_outer to keep adjacent outer blocks from
+    # writing into each other's columns.
+    if MUL_ROUTED_WEIGHT:
+        moe_w = tl.load(topk_weights_ptr + offs_token, mask=token_mask, other=0.0).to(
+            tl.float32
+        )
+
+    out_slice_base = out_ptr + slice_id * slice_n_offset
+
+    for n_start in range(n_lo, n_hi, BLOCK_N):
+        offs_n = n_start + tl.arange(0, BLOCK_N)
+        n_mask = (offs_n < N) & (offs_n < n_hi)
+
+        b_ptrs = (
+            B_base + safe_offs_r[:, None] * stride_B_r + offs_n[None, :] * stride_B_n
+        )
+        b = tl.load(b_ptrs, mask=rank_mask[:, None] & n_mask[None, :], other=0.0)
+
+        acc = tl.dot(tmp_typed, b)  # (BLOCK_M, BLOCK_N) fp32
+        if MUL_ROUTED_WEIGHT:
+            acc = acc * moe_w[:, None]
+
+        out_ptrs = (
+            out_slice_base
+            + offs_token[:, None] * stride_om
+            + offs_n[None, :] * stride_on
+        )
+        out_mask = token_mask[:, None] & n_mask[None, :]
+        if ADD_INPUTS:
+            prev = tl.load(out_ptrs, mask=out_mask, other=0.0)
+            tl.store(out_ptrs, prev + acc.to(out_ptr.dtype.element_ty), mask=out_mask)
+        else:
+            tl.store(out_ptrs, acc.to(out_ptr.dtype.element_ty), mask=out_mask)
+
+
+def _run_fused_moe_lora_one_shot(
+    output: torch.Tensor,
+    qcurr_hidden_states: torch.Tensor,
+    lora_a_stacked: list[torch.Tensor],
+    lora_b_stacked: list[torch.Tensor],
+    topk_weights: torch.Tensor,
+    sorted_token_ids: torch.Tensor | None,
+    expert_ids: torch.Tensor,
+    num_tokens_post_padded: torch.Tensor | None,
+    token_lora_mapping: torch.Tensor,
+    max_lora_rank: int,
+    top_k_num: int,
+    lora_ids: torch.Tensor,
+    num_active_loras: torch.Tensor,
+    adapter_enabled: torch.Tensor,
+    mul_routed_weight: bool,
+    block_size_m: int,
+    add_inputs: bool = True,
+) -> None:
+    """Fast-path wrapper: launches one fused shrink+expand kernel.
+
+    The shape contract matches `_fused_moe_lora`. `output` has shape
+    `(num_tokens, top_k_num, num_slices * N_per_slice)`. When
+    `add_inputs=True` (default) the kernel reads-modifies-writes `output`
+    in place; when `add_inputs=False` the kernel overwrites `output` with
+    the LoRA delta only. The latter is used by the dual-stream path that
+    sums LoRA into the base output on a separate stream.
+    """
+    num_slices = len(lora_a_stacked)
+    device = qcurr_hidden_states.device
+
+    A0 = lora_a_stacked[0]
+    B0 = lora_b_stacked[0]
+    max_loras_w = A0.shape[0]
+    rank = A0.shape[2]
+    K = A0.shape[3]
+    N_per_slice = B0.shape[2]
+
+    # rank padding is to next pow2 with a floor of 16 (tensor-core minimum
+    # K-dim). Beyond 128 the (BLOCK_M, BLOCK_R) accumulator outgrows the
+    # register file; rank tiling would be needed but is out of scope for
+    # this kernel. Tried floor=32 to double MMA density per K-step but it
+    # regressed across all M (+8 to +40%): the (64,32) fp32 accumulator +
+    # widened B tile pushed register count past spill threshold, lowering
+    # occupancy by more than the MMA gain saved.
+    assert rank <= 128, (
+        f"fused_moe_lora_one_shot supports max_lora_rank<=128; got rank={rank}"
+    )
+    BLOCK_R = max(triton.next_power_of_2(rank), 16)
+
+    num_experts = A0.shape[1]
+    naive = sorted_token_ids is None
+    if sorted_token_ids is None:
+        EM_grid = topk_weights.numel()
+        BLOCK_M = 16
+        stride_tl_ = 0
+        stride_el = 0
+        grid_lora_dim = 1
+    else:
+        EM_grid = sorted_token_ids.shape[1]
+        # BLOCK_M must equal moe_lora_align_block_size's block_size. The
+        # caller passes that explicitly; deriving it from tensor shapes is
+        # unsafe because sorted_token_ids.shape[1] is the raw padded length
+        # (not necessarily a multiple of block_size — e.g. OLMoE prefill
+        # produces sorted=139200 with expert_ids=1088 and block_size=128).
+        # tl.arange and tl.dot need block_size_m to be a power of 2 and at
+        # least 16. The Python-side assertion gives a clearer error than
+        # the cryptic Triton compile failure.
+        assert block_size_m >= 16 and (block_size_m & (block_size_m - 1)) == 0, (
+            f"shrink_block_size_m must be a power of 2 and >=16; got {block_size_m}"
+        )
+        BLOCK_M = block_size_m
+        stride_tl_ = sorted_token_ids.stride(0)
+        stride_el = expert_ids.stride(0)
+        grid_lora_dim = int(num_active_loras.item())
+
+    # Empty-work guards: the grid would otherwise have a zero dimension,
+    # which Triton rejects. None of these is a hot path in production -- a
+    # batch with zero tokens, an EM_grid of zero, or zero active LoRAs all
+    # mean there's nothing to add to `output`.
+    if EM_grid == 0 or grid_lora_dim == 0 or num_slices == 0:
+        return
+
+    token_mapping_factor = 1 if mul_routed_weight else top_k_num
+
+    A_ptrs = _get_ptr(lora_a_stacked, device)
+    B_ptrs = _get_ptr(lora_b_stacked, device)
+
+    # Flatten (num_tokens, top_k) → flat_token axis. The kernel addresses
+    # output via offs_token * stride_om, which is correct iff the dim-0 /
+    # dim-1 strides collapse cleanly: stride(0) == top_k * stride(1). All
+    # production callers pass contiguous output, so this always holds; the
+    # explicit check guards against future regressions where a non-trivial
+    # view (e.g. permute) would silently break in-place accumulation.
+    assert output.dim() == 3, f"output must be 3-D, got {output.shape}"
+    assert output.stride(0) == output.shape[1] * output.stride(1), (
+        "fused_moe_lora_one_shot requires output.stride(0) == top_k*stride(1); "
+        f"got shape={output.shape} strides={output.stride()}"
+    )
+    out_view = output.view(-1, output.shape[-1])
+    M_blocks = triton.cdiv(EM_grid, BLOCK_M) if not naive else EM_grid
+
+    # 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 = 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)
+    target = 2 * sm_count
+    if base_programs >= int(1.5 * sm_count):
+        npid = 1
+    else:
+        npid_occ = max(1, min(16, (target + base_programs - 1) // base_programs))
+        npid = min(npid_occ, max_npid_by_budget)
+    npid = max(1, min(npid, max(1, N_per_slice // 128)))
+
+    # Robust defaults across the prefill regime (H100/H200/B200, bf16/fp16).
+    # NPID > 1 is the small-M / under-saturated path -- more warps help
+    # amortise the inner-N expand loop. ns=3 instead of 4: GB200 ncu showed
+    # the 4-stage pipeline pushed register count to 168/thread and capped
+    # achieved occupancy at ~17% (3 blocks/SM, register-bound); ns=3 frees
+    # ~30 regs/thread which keeps a 4th block resident on small grids.
+    # Tried BLOCK_N=64 for w13 (N=192) to avoid the half-wasted second
+    # tile: regressed 11-29% because the "waste" was just masked stores
+    # (cheap) and the extra iteration added load + index overhead.
+    if npid > 1:
+        block_n, nw, ns = 128, 8, 3
+    else:
+        block_n, nw, ns = 128, 4, 3
+    # BLOCK_K choice: for hidden-sized K (≥256, i.e. the K=hidden_size
+    # shrink input on w13) force BLOCK_K=128 -- the wider tile halves the
+    # K-loop trip count and removes the scoreboard stalls that dominated
+    # M=16-64 on GB200 (kernel time -13% to -37% vs the work_per_expert
+    # heuristic which picked 64 for low-tokens-per-expert ratios). For
+    # small-K shapes (e.g. w2 with K=192 where the down-proj reads the
+    # MoE intermediate) keep the work_per_expert heuristic: BLOCK_K=128
+    # would force the EVEN_K=False masked path and add no K-loop savings
+    # (K/64=3 vs K/128=2 masked) while inflating per-program startup.
+    if K >= 256:
+        block_k = 128
+    else:
+        work_per_expert = topk_weights.numel() / max(num_experts, 1)
+        block_k = 128 if work_per_expert >= 16 else 64
+
+    grid = (M_blocks * npid, num_slices, grid_lora_dim)
+
+    _fused_moe_lora_one_shot_kernel[grid](
+        qcurr_hidden_states,
+        A_ptrs,
+        B_ptrs,
+        out_view,
+        topk_weights,
+        sorted_token_ids,
+        expert_ids,
+        num_tokens_post_padded,
+        token_lora_mapping,
+        lora_ids,
+        adapter_enabled,
+        N_per_slice,
+        K,
+        topk_weights.numel(),
+        top_k_num,
+        max_loras_w,
+        qcurr_hidden_states.stride(0),
+        qcurr_hidden_states.stride(1),
+        A0.stride(0),
+        A0.stride(1),
+        A0.stride(2),
+        A0.stride(3),
+        B0.stride(0),
+        B0.stride(1),
+        B0.stride(2),
+        B0.stride(3),
+        out_view.stride(0),
+        out_view.stride(1),
+        stride_tl_,
+        stride_el,
+        N_per_slice,
+        token_mapping_factor=token_mapping_factor,
+        naive_block_assignment=naive,
+        MUL_ROUTED_WEIGHT=mul_routed_weight,
+        BLOCK_M=BLOCK_M,
+        BLOCK_R=BLOCK_R,
+        actual_rank=rank,
+        NPID_FACTOR=npid,
+        BLOCK_N=block_n,
+        BLOCK_K=block_k,
+        ADD_INPUTS=add_inputs,
+        num_warps=nw,
+        num_stages=ns,
+    )
+
+
+# ---------------------------------------------------------------------------
+# Small-batch (decode-style) fused MoE-LoRA kernel — sub-path of the
+# one_shot fast path.
+# ---------------------------------------------------------------------------
+
+
+@triton.heuristics({"EVEN_K": lambda args: args["K"] % args["BLOCK_K"] == 0})
+@triton.jit
+def _fused_moe_lora_small_batch_kernel(
+    # ---- pointers ----
+    x_ptr,
+    A_ptrs,
+    B_ptrs,
+    out_ptr,
+    topk_weights_ptr,
+    expert_ids_ptr,  # (num_tokens * top_k_num,)
+    token_lora_mapping_ptr,  # (num_tokens,)
+    adapter_enabled_ptr,
+    # ---- dims ----
+    N,
+    K,
+    top_k_num,
+    max_loras,
+    work_total,  # = pair_slices * n_chunks_per_pair_slice
+    pair_slices,  # = num_tokens * top_k_num * NUM_SLICES
+    # ---- strides ----
+    stride_xm,
+    stride_xk,
+    stride_A_lora,
+    stride_A_expert,
+    stride_A_r,
+    stride_A_k,
+    stride_B_lora,
+    stride_B_expert,
+    stride_B_n,
+    stride_B_r,
+    stride_om,
+    stride_on,
+    # ---- scalar (runtime ints, NOT constexpr) ----
+    # n_tiles_per_program / n_chunks_per_pair_slice are deliberately
+    # runtime: each distinct value would otherwise trigger a fresh Triton
+    # compile -> fresh kernel binary -> fresh CUDA graph instance per
+    # batch size. Production traces showed that variant explosion adding
+    # ~5.9k graph instantiations on top of legacy. Runtime args mean one
+    # shared binary across all chunk sizes.
+    slice_n_offset,
+    n_tiles_per_program,
+    n_chunks_per_pair_slice,
+    # ---- constexpr ----
+    token_mapping_factor: tl.constexpr,
+    MUL_ROUTED_WEIGHT: tl.constexpr,
+    ADD_INPUTS: tl.constexpr,
+    BLOCK_R: tl.constexpr,
+    actual_rank: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    BLOCK_K: tl.constexpr,
+    NUM_SLICES: tl.constexpr,
+    EVEN_K: tl.constexpr,
+):
+    """Persistent fused MoE-LoRA kernel for naive_block_assignment inputs.
+
+    Each program owns one (pair × slice × n_chunk) work item. A "chunk"
+    covers `n_tiles_per_program` consecutive output-N tiles, all of which
+    share a single shrink — so the rank-vector is computed once per
+    program and the A weights for that (lora, expert, slice) are loaded
+    once instead of n_tiles_per_program times.
+
+    The wrapper picks `n_tiles_per_program` to keep the grid close to
+    2*SM_count: at very small batch (work_total ≤ SM_count) the chunk
+    size collapses to 1 and behaviour matches a per-tile GEMV; as batch
+    grows the chunk grows so we trade some N-axis parallelism for shrink
+    reuse. When `work_total` exceeds the launched grid, the outer stride
+    loop drains the leftover work units serially.
+    """
+    pid = tl.program_id(axis=0)
+    num_programs = tl.num_programs(axis=0)
+
+    offs_r = tl.arange(0, BLOCK_R)
+    rank_mask = offs_r < actual_rank
+    # Clamp OOB rank lanes so they address row 0 of A/B; the mask zeros
+    # the loaded values. Required when BLOCK_R > actual_rank (e.g. rank=4
+    # padded to 16) -- without clamping, tl.load would address the next
+    # expert's memory.
+    safe_offs_r = tl.where(rank_mask, offs_r, 0)
+    offs_k = tl.arange(0, BLOCK_K)
+
+    # Persistent stride loop: when grid < work_total each program walks
+    # multiple work items. When grid == work_total the loop runs exactly
+    # once and the kernel degenerates to the per-tile GEMV.
+    for work_id in range(pid, work_total, num_programs):
+        n_chunk_idx = work_id % n_chunks_per_pair_slice
+        pair_slice_idx = work_id // n_chunks_per_pair_slice
+        # NUM_SLICES is constexpr (typ. 1 or 2) so divmod folds.
+        pair_idx = pair_slice_idx // NUM_SLICES
+        slice_id = pair_slice_idx % NUM_SLICES
+
+        # Resolve lora_id / expert_id; skip the body for inactive lanes.
+        # Using a single `valid` flag instead of early `return` keeps the
+        # outer stride loop alive — `return` would exit the whole program
+        # and skip later work items assigned to this SM.
+        token_idx = pair_idx // top_k_num
+        lora_id = tl.load(token_lora_mapping_ptr + token_idx)
+        valid = (lora_id >= 0) & (lora_id < max_loras)
+        enabled = tl.load(adapter_enabled_ptr + tl.where(valid, lora_id, 0))
+        valid = valid & (enabled != 0)
+        expert_id = tl.load(expert_ids_ptr + pair_idx)
+        valid = valid & (expert_id >= 0)
+
+        if valid:
+            cur_A_ptr = tl.load(A_ptrs + slice_id).to(
+                tl.pointer_type(out_ptr.dtype.element_ty)
+            )
+            cur_B_ptr = tl.load(B_ptrs + slice_id).to(
+                tl.pointer_type(out_ptr.dtype.element_ty)
+            )
+            A_base = cur_A_ptr + lora_id * stride_A_lora + expert_id * stride_A_expert
+            B_base = cur_B_ptr + lora_id * stride_B_lora + expert_id * stride_B_expert
+
+            x_row = pair_idx // token_mapping_factor
+            x_row_ptr = x_ptr + x_row * stride_xm
+
+            # SHRINK GEMV (once per program; reused across n_tiles_per_program
+            # expand tiles below). Sum-reduction over BLOCK_K with fp32
+            # accumulator — same precision path as the one_shot kernel.
+            rank_vec = tl.zeros((BLOCK_R,), dtype=tl.float32)
+            if EVEN_K:
+                for kb in range(0, K, BLOCK_K):
+                    cur_k = kb + offs_k
+                    x_tile = tl.load(x_row_ptr + cur_k * stride_xk).to(tl.float32)
+                    a_tile = tl.load(
+                        A_base
+                        + safe_offs_r[:, None] * stride_A_r
+                        + cur_k[None, :] * stride_A_k,
+                        mask=rank_mask[:, None],
+                        other=0.0,
+                    ).to(tl.float32)
+                    rank_vec += tl.sum(a_tile * x_tile[None, :], axis=1)
+            else:
+                for kb in range(0, K, BLOCK_K):
+                    cur_k = kb + offs_k
+                    k_mask = cur_k < K
+                    x_tile = tl.load(
+                        x_row_ptr + cur_k * stride_xk, mask=k_mask, other=0.0
+                    ).to(tl.float32)
+                    a_tile = tl.load(
+                        A_base
+                        + safe_offs_r[:, None] * stride_A_r
+                        + cur_k[None, :] * stride_A_k,
+                        mask=rank_mask[:, None] & k_mask[None, :],
+                        other=0.0,
+                    ).to(tl.float32)
+                    rank_vec += tl.sum(a_tile * x_tile[None, :], axis=1)
+
+            # EXPAND: walk n_tiles_per_program consecutive output-N tiles
+            # using the same rank_vec. The loop is a runtime range (not
+            # tl.static_range) so a single compiled kernel handles every
+            # chunk size — see the note on the kernel signature.
+            n_tile_start = n_chunk_idx * n_tiles_per_program
+            out_row_ptr = out_ptr + slice_id * slice_n_offset + pair_idx * stride_om
+
+            if MUL_ROUTED_WEIGHT:
+                moe_w = tl.load(topk_weights_ptr + pair_idx).to(tl.float32)
+
+            for nt in range(n_tiles_per_program):
+                n_lo = (n_tile_start + nt) * BLOCK_N
+                if n_lo < N:
+                    offs_n = n_lo + tl.arange(0, BLOCK_N)
+                    n_mask = offs_n < N
+                    b_tile = tl.load(
+                        B_base
+                        + offs_n[:, None] * stride_B_n
+                        + safe_offs_r[None, :] * stride_B_r,
+                        mask=n_mask[:, None] & rank_mask[None, :],
+                        other=0.0,
+                    ).to(tl.float32)
+                    out_tile = tl.sum(b_tile * rank_vec[None, :], axis=1)
+
+                    if MUL_ROUTED_WEIGHT:
+                        out_tile = out_tile * moe_w
+
+                    out_ptrs = out_row_ptr + offs_n * stride_on
+                    if ADD_INPUTS:
+                        prev = tl.load(out_ptrs, mask=n_mask, other=0.0).to(tl.float32)
+                        tl.store(
+                            out_ptrs,
+                            (prev + out_tile).to(out_ptr.dtype.element_ty),
+                            mask=n_mask,
+                        )
+                    else:
+                        tl.store(
+                            out_ptrs,
+                            out_tile.to(out_ptr.dtype.element_ty),
+                            mask=n_mask,
+                        )
+
+
+def _pick_small_batch_chunk(pair_slices: int, N_tiles: int, sm_count: int) -> int:
+    """Pick `n_tiles_per_program` so the launched grid stays near
+    2*SM_count.
+
+    Sizes for occupancy first (more programs in flight → better latency
+    hiding for the K-loop A/x loads). Once the per-tile grid already
+    exceeds 2*SM_count we increase the chunk size to amortise the shrink
+    cost — at that point the GPU is saturated by per-program work and
+    packing more tiles per program lets the rank_vec be reused.
+    """
+    target_grid = max(1, 2 * sm_count)
+    total_work = pair_slices * N_tiles
+    if total_work <= target_grid:
+        return 1
+    ntpp = (total_work + target_grid - 1) // target_grid
+    return min(ntpp, N_tiles)
+
+
+def _run_fused_moe_lora_small_batch(
+    output: torch.Tensor,
+    qcurr_hidden_states: torch.Tensor,
+    lora_a_stacked: list[torch.Tensor],
+    lora_b_stacked: list[torch.Tensor],
+    topk_weights: torch.Tensor,
+    expert_ids_flat: torch.Tensor,  # (num_tokens * top_k_num,)
+    token_lora_mapping: torch.Tensor,
+    top_k_num: int,
+    adapter_enabled: torch.Tensor,
+    mul_routed_weight: bool,
+    add_inputs: bool = True,
+) -> None:
+    """Small-batch GEMV-style wrapper. Naive-block-assignment inputs only.
+
+    Shape contract matches `_run_fused_moe_lora_one_shot`: `output` is
+    `(num_tokens, top_k_num, num_slices * N_per_slice)` with
+    contiguous-style strides, `expert_ids_flat` is the flattened
+    `topk_ids` of shape `(num_tokens * top_k_num,)`, and the
+    rank-padded LoRA weights live in `lora_a_stacked` /
+    `lora_b_stacked`.
+
+    The kernel is persistent over (pair × slice × n_chunk) work items —
+    each program does one shrink and reuses the rank vector across
+    `n_tiles_per_program` expand tiles. The chunk size scales with the
+    pair-slice count so very small batches keep per-tile parallelism
+    while medium batches cut redundant shrinks.
+    """
+    num_slices = len(lora_a_stacked)
+    device = qcurr_hidden_states.device
+
+    A0 = lora_a_stacked[0]
+    B0 = lora_b_stacked[0]
+    max_loras_w = A0.shape[0]
+    rank = A0.shape[2]
+    K = A0.shape[3]
+    N_per_slice = B0.shape[2]
+
+    # Rank padding: floor 16 (tensor-core min K), ceil to next pow2. The
+    # ≤64 cap is set conservatively for the prototype: at rank 64 the
+    # per-program register footprint is rank_vec(64 fp32) + b_tile(BLOCK_N
+    # × 64 fp32) = e.g. 128*64*4 = 32 KiB, comfortably within the 64 KiB
+    # register file even with num_warps=8. Doubling to 128 would push us
+    # against the limit and require shared-memory staging.
+    assert rank <= 64, f"fused_moe_lora_small_batch supports rank<=64; got rank={rank}"
+    BLOCK_R = max(triton.next_power_of_2(rank), 16)
+
+    num_tokens = topk_weights.shape[0]
+    M_grid = num_tokens * top_k_num
+    if M_grid == 0 or num_slices == 0:
+        return
+
+    token_mapping_factor = 1 if mul_routed_weight else top_k_num
+
+    A_ptrs = _get_ptr(lora_a_stacked, device)
+    B_ptrs = _get_ptr(lora_b_stacked, device)
+
+    assert output.dim() == 3, f"output must be 3-D, got {output.shape}"
+    assert output.stride(0) == output.shape[1] * output.stride(1), (
+        "fused_moe_lora_small_batch requires output.stride(0) == "
+        f"top_k*stride(1); got shape={output.shape} strides={output.stride()}"
+    )
+    out_view = output.view(-1, output.shape[-1])
+
+    # Block sizes. BLOCK_N=128 matches the one_shot's expand tile and gives
+    # 6-24 N tiles for typical N ∈ [768, 3072], enough to saturate the SM
+    # array once M_grid * num_slices reaches ~SM_count. BLOCK_K=128 halves
+    # the K-loop trip count vs 64 and pays for itself once K ≥ 1024 (the
+    # only regime we care about — hidden sizes are always large here).
+    BLOCK_N = 128
+    BLOCK_K = 128
+    nw = 4
+    ns = 3
+
+    N_tiles = triton.cdiv(N_per_slice, BLOCK_N)
+    pair_slices = M_grid * num_slices
+
+    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
+
+    # Grid sizing: keep parallelism uncapped when work_total is small (so
+    # very small batches still spread across SMs); cap at 2*SM_count once
+    # we have plenty of work, letting the in-kernel stride loop drain the
+    # remainder.
+    grid_size = min(work_total, max(1, 2 * sm_count))
+    grid = (grid_size,)
+
+    _fused_moe_lora_small_batch_kernel[grid](
+        qcurr_hidden_states,
+        A_ptrs,
+        B_ptrs,
+        out_view,
+        topk_weights,
+        expert_ids_flat,
+        token_lora_mapping,
+        adapter_enabled,
+        N_per_slice,
+        K,
+        top_k_num,
+        max_loras_w,
+        work_total,
+        pair_slices,
+        qcurr_hidden_states.stride(0),
+        qcurr_hidden_states.stride(1),
+        A0.stride(0),
+        A0.stride(1),
+        A0.stride(2),
+        A0.stride(3),
+        B0.stride(0),
+        B0.stride(1),
+        B0.stride(2),
+        B0.stride(3),
+        out_view.stride(0),
+        out_view.stride(1),
+        N_per_slice,
+        n_tiles_per_program,
+        n_chunks,
+        token_mapping_factor=token_mapping_factor,
+        MUL_ROUTED_WEIGHT=mul_routed_weight,
+        ADD_INPUTS=add_inputs,
+        BLOCK_R=BLOCK_R,
+        actual_rank=rank,
+        BLOCK_N=BLOCK_N,
+        BLOCK_K=BLOCK_K,
+        NUM_SLICES=num_slices,
+        num_warps=nw,
+        num_stages=ns,
+    )
+
+
 def _get_ptr(lora_weights: list[torch.Tensor], device: torch.device):
     """
     `_LORA_PTR_DICT` collects the required information during `profile_run`,
@@ -706,6 +1443,7 @@ def _fused_moe_lora(
     mul_routed_weight: bool = False,
     fully_sharded: bool = False,
     offset: int = 0,
+    add_inputs: bool = True,
 ) -> None:
     assert len(lora_a_stacked) == len(lora_b_stacked) > 0
     assert topk_weights.dim() == qcurr_hidden_states.dim() == 2
@@ -728,6 +1466,59 @@ def _fused_moe_lora(
         )
     assert output.shape[0] == topk_weights.shape[0]
     assert top_k_num == topk_weights.shape[1]
+
+    # Fast path: single fused kernel
+    if not fully_sharded:
+        M_pairs = topk_weights.numel()
+        if (
+            sorted_token_ids is None
+            and max_lora_rank <= 64
+            and M_pairs * max_lora_rank <= 1024
+        ):
+            _run_fused_moe_lora_small_batch(
+                output,
+                qcurr_hidden_states,
+                lora_a_stacked,
+                lora_b_stacked,
+                topk_weights,
+                expert_ids,
+                token_lora_mapping,
+                top_k_num,
+                adapter_enabled,
+                mul_routed_weight,
+                add_inputs=add_inputs,
+            )
+            return
+        # shrink/expand BLOCK_SIZE_M must match the block_size that
+        # moe_lora_align_block_size used; both shrink and expand pass the
+        # same value (asserted by `shrink_block_size_m == expand_block_size_m`
+        # below).
+        _run_fused_moe_lora_one_shot(
+            output,
+            qcurr_hidden_states,
+            lora_a_stacked,
+            lora_b_stacked,
+            topk_weights,
+            sorted_token_ids,
+            expert_ids,
+            num_tokens_post_padded,
+            token_lora_mapping,
+            max_lora_rank,
+            top_k_num,
+            lora_ids,
+            num_active_loras,
+            adapter_enabled,
+            mul_routed_weight,
+            shrink_block_size_m,
+            add_inputs=add_inputs,
+        )
+        return
+
+    assert add_inputs, (
+        "fused_moe_lora(add_inputs=False) is only supported on the "
+        "fully_sharded=False fast path"
+    )
+
     device = qcurr_hidden_states.device
     num_slices = len(lora_a_stacked)
     w1_lora_b_stacked = lora_b_stacked[0]
@@ -894,6 +1685,7 @@ def _fused_moe_lora_fake(
     mul_routed_weight: bool = False,
     fully_sharded: bool = False,
     offset: int = 0,
+    add_inputs: bool = True,
 ) -> None:
     return
 

vllm/lora/punica_wrapper/punica_base.py

@@ -514,6 +514,7 @@ class PunicaWrapperBase(PunicaWrapperABC):
         num_slices: int,
         fully_sharded: bool,
         use_tuned_config: bool,
+        add_inputs: bool = True,
         token_lora_mapping: torch.Tensor | None = None,
     ) -> tuple[
         torch.Tensor | None,
@@ -554,6 +555,7 @@ class PunicaWrapperBase(PunicaWrapperABC):
         fully_sharded: bool,
         tp_rank: int,
         use_tuned_config: bool,
+        add_inputs: bool = True,
     ) -> None:
         """Apply w2 LoRA to y (intermediate_cache3) in-place before moe_sum.
 

vllm/lora/punica_wrapper/punica_gpu.py

@@ -435,6 +435,7 @@ class PunicaWrapperGPU(PunicaWrapperBase):
         fully_sharded: bool = False,
         offset: int = 0,
         token_lora_mapping: torch.Tensor | None = None,
+        add_inputs: bool = True,
     ):
         """
         Performs a fused forward computation for LoRA of Mixture-of-Experts (MoE) layer.
@@ -484,6 +485,7 @@ class PunicaWrapperGPU(PunicaWrapperBase):
             mul_routed_weight,
             fully_sharded,
             offset,
+            add_inputs,
         )
 
     def add_lora_w13(
@@ -506,6 +508,7 @@ class PunicaWrapperGPU(PunicaWrapperBase):
         num_slices: int,
         fully_sharded: bool,
         use_tuned_config: bool,
+        add_inputs: bool = True,
         token_lora_mapping: torch.Tensor | None = None,
     ) -> tuple[
         torch.Tensor | None,
@@ -610,6 +613,7 @@ class PunicaWrapperGPU(PunicaWrapperBase):
             adapter_enabled,
             fully_sharded=fully_sharded,
             token_lora_mapping=token_lora_mapping,
+            add_inputs=add_inputs,
         )
 
         return (
@@ -640,6 +644,7 @@ class PunicaWrapperGPU(PunicaWrapperBase):
         fully_sharded: bool,
         tp_rank: int,
         use_tuned_config: bool,
+        add_inputs: bool = True,
     ) -> None:
         import functools
 
@@ -722,4 +727,5 @@ class PunicaWrapperGPU(PunicaWrapperBase):
             fully_sharded=fully_sharded,
             offset=offset,
             token_lora_mapping=token_lora_mapping,
+            add_inputs=add_inputs,
         )

vllm/lora/punica_wrapper/punica_xpu.py

@@ -390,6 +390,7 @@ class PunicaWrapperXPU(PunicaWrapperBase):
         fully_sharded: bool = False,
         offset: int = 0,
         token_lora_mapping: torch.Tensor | None = None,
+        add_inputs: bool = True,
     ):
         """
         Performs a fused forward computation for LoRA of Mixture-of-Experts (MoE) layer.
@@ -439,6 +440,7 @@ class PunicaWrapperXPU(PunicaWrapperBase):
             mul_routed_weight,
             fully_sharded,
             offset,
+            add_inputs,
         )
 
     def add_lora_w13(
@@ -461,6 +463,7 @@ class PunicaWrapperXPU(PunicaWrapperBase):
         num_slices: int,
         fully_sharded: bool,
         use_tuned_config: bool,
+        add_inputs: bool = True,
         token_lora_mapping: torch.Tensor | None = None,
     ) -> tuple[
         torch.Tensor | None,
@@ -594,6 +597,7 @@ class PunicaWrapperXPU(PunicaWrapperBase):
         fully_sharded: bool,
         tp_rank: int,
         use_tuned_config: bool,
+        add_inputs: bool = True,
     ) -> None:
         import functools
 

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

@@ -43,6 +43,16 @@ class MoELoRAContext:
     # try_get_optimal_moe_lora_config for Triton kernel tile configs.
     use_tuned_config: bool
 
+    # Optional dual-stream support for overlapping each (base GEMM, LoRA)
+    # pair. When aux_stream is None, the experts.apply() path runs the
+    # original sequential schedule. When set, base GEMM runs on the default
+    # stream and the LoRA fast-path writes the delta into a fresh buffer on
+    # aux_stream, which the default stream sums in afterwards.
+    # Events are paired one-per-overlap-pair: events[0,1] for w13,
+    # events[2,3] for w2, so the two pairs do not race on the same event.
+    aux_stream: torch.cuda.Stream | None = None
+    events: tuple[torch.cuda.Event, ...] | None = None
+
     # Per-rank token→LoRA mapping after EP dispatch. Set by
     # FusedMoEPrepareAndFinalizeModular.prepare() when EP+LoRA is active, read
     # by LoRAExpertsMixin helpers in place of punica_wrapper's global mapping.

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

@@ -45,6 +45,7 @@ class LoRAExpertsMixin:
         w2: torch.Tensor,
         num_tokens: int,
         top_k_num: int,
+        add_inputs: bool = True,
     ) -> tuple[
         torch.Tensor | None,
         torch.Tensor | None,
@@ -70,6 +71,7 @@ class LoRAExpertsMixin:
             lora_context.w13_num_slices,
             lora_context.fully_sharded,
             lora_context.use_tuned_config,
+            add_inputs=add_inputs,
             token_lora_mapping=lora_context.local_token_lora_mapping,
         )
 
@@ -88,6 +90,7 @@ class LoRAExpertsMixin:
         w1: torch.Tensor,
         w2: torch.Tensor,
         top_k_num: int,
+        add_inputs: bool = True,
     ) -> None:
         lora_context.punica_wrapper.add_lora_w2(
             y,
@@ -109,4 +112,5 @@ class LoRAExpertsMixin:
             lora_context.fully_sharded,
             lora_context.tp_rank,
             lora_context.use_tuned_config,
+            add_inputs=add_inputs,
         )

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

@@ -48,6 +48,7 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
 )
 from vllm.platforms import current_platform
 from vllm.triton_utils import tl
+from vllm.utils.multi_stream_utils import maybe_execute_in_parallel
 
 
 class TritonExperts(LoRAExpertsMixin, mk.FusedMoEExpertsModular):
@@ -247,55 +248,99 @@ class TritonExperts(LoRAExpertsMixin, mk.FusedMoEExpertsModular):
             )
         )
 
-        invoke_fused_moe_triton_kernel(
-            hidden_states,
-            w1,
-            intermediate_cache1,
-            a1q_scale if a1q_scale is not None else self.a1_scale,
-            self.w1_scale,
-            None,  # topk_weights
-            sorted_token_ids,
-            expert_ids,
-            num_tokens_post_padded,
-            False,  # mul_routed_weights
-            top_k_num,
-            config,
-            compute_type=compute_type,
-            use_fp8_w8a8=self.quant_config.use_fp8_w8a8,
-            use_int8_w8a8=self.quant_config.use_int8_w8a8,
-            use_int8_w8a16=self.quant_config.use_int8_w8a16,
-            use_int4_w4a16=self.quant_config.use_int4_w4a16,
-            per_channel_quant=self.per_act_token_quant,
-            block_shape=self.block_shape,
-            B_bias=self.w1_bias,
-        )
+        # LoRA w13: applied to intermediate_cache1 before activation. When
+        # the LoRA layer requested a dual-stream schedule, we run base w13
+        # GEMM on the default stream and the LoRA fast-path on aux_stream;
+        # the LoRA writes its delta into a fresh zero buffer (add_inputs=
+        # False) and we sum it into intermediate_cache1 after both finish.
 
-        # LoRA w13: applied to intermediate_cache1 before activation, using
-        # hidden_states as the lora_a input.  moe_lora_align_block_size is
-        # called once here and results reused for the w2 LoRA below.
         sorted_token_ids_lora = None
         expert_ids_lora = None
         num_tokens_post_padded_lora = None
         token_lora_mapping = None
         lora_context = self._lora_context
-        if lora_context is not None:
+
+        def _base_w13_fn():
+            invoke_fused_moe_triton_kernel(
+                hidden_states,
+                w1,
+                intermediate_cache1,
+                a1q_scale if a1q_scale is not None else self.a1_scale,
+                self.w1_scale,
+                None,  # topk_weights
+                sorted_token_ids,
+                expert_ids,
+                num_tokens_post_padded,
+                False,  # mul_routed_weights
+                top_k_num,
+                config,
+                compute_type=compute_type,
+                use_fp8_w8a8=self.quant_config.use_fp8_w8a8,
+                use_int8_w8a8=self.quant_config.use_int8_w8a8,
+                use_int8_w8a16=self.quant_config.use_int8_w8a16,
+                use_int4_w4a16=self.quant_config.use_int4_w4a16,
+                per_channel_quant=self.per_act_token_quant,
+                block_shape=self.block_shape,
+                B_bias=self.w1_bias,
+            )
+
+        if lora_context is not None and lora_context.aux_stream is not None:
+            # add_inputs=False: kernel overwrites lora_delta_w13. zeros (not
+            # empty) so untouched rows -- e.g. blocks where every program
+            # early-exits because lora_id<0 -- stay at zero and the trailing
+            # add_() is a no-op there.
+            lora_delta_w13 = torch.zeros_like(intermediate_cache1)
+
+            def _lora_w13_fn():
+                return self.apply_w13_lora(
+                    lora_context,
+                    y=lora_delta_w13,
+                    x=hidden_states,
+                    topk_ids=topk_ids,
+                    topk_weights=topk_weights,
+                    expert_map=expert_map,
+                    w1=w1,
+                    w2=w2,
+                    num_tokens=num_tokens,
+                    top_k_num=top_k_num,
+                    add_inputs=False,
+                )
+
+            assert lora_context.events is not None
+            _, lora_meta = maybe_execute_in_parallel(
+                _base_w13_fn,
+                _lora_w13_fn,
+                lora_context.events[0],
+                lora_context.events[1],
+                lora_context.aux_stream,
+            )
             (
                 sorted_token_ids_lora,
                 expert_ids_lora,
                 num_tokens_post_padded_lora,
                 token_lora_mapping,
-            ) = self.apply_w13_lora(
-                lora_context,
-                y=intermediate_cache1,
-                x=hidden_states,
-                topk_ids=topk_ids,
-                topk_weights=topk_weights,
-                expert_map=expert_map,
-                w1=w1,
-                w2=w2,
-                num_tokens=num_tokens,
-                top_k_num=top_k_num,
-            )
+            ) = lora_meta
+            intermediate_cache1.add_(lora_delta_w13)
+        else:
+            _base_w13_fn()
+            if lora_context is not None:
+                (
+                    sorted_token_ids_lora,
+                    expert_ids_lora,
+                    num_tokens_post_padded_lora,
+                    token_lora_mapping,
+                ) = self.apply_w13_lora(
+                    lora_context,
+                    y=intermediate_cache1,
+                    x=hidden_states,
+                    topk_ids=topk_ids,
+                    topk_weights=topk_weights,
+                    expert_map=expert_map,
+                    w1=w1,
+                    w2=w2,
+                    num_tokens=num_tokens,
+                    top_k_num=top_k_num,
+                )
 
         a2q_scale: torch.Tensor | None = None
 
@@ -328,48 +373,82 @@ class TritonExperts(LoRAExpertsMixin, mk.FusedMoEExpertsModular):
                 quantization_emulation=self.quantization_emulation,
             )
 
-        invoke_fused_moe_triton_kernel(
-            qintermediate_cache2,
-            w2,
-            intermediate_cache3,
-            a2q_scale,
-            self.w2_scale,
-            topk_weights,
-            sorted_token_ids,
-            expert_ids,
-            num_tokens_post_padded,
-            not apply_router_weight_on_input,
-            1,
-            config,
-            compute_type=compute_type,
-            use_fp8_w8a8=self.quant_config.use_fp8_w8a8,
-            use_int8_w8a8=self.quant_config.use_int8_w8a8,
-            use_int8_w8a16=self.quant_config.use_int8_w8a16,
-            use_int4_w4a16=self.quant_config.use_int4_w4a16,
-            per_channel_quant=self.per_act_token_quant,
-            block_shape=self.block_shape,
-            B_bias=self.w2_bias,
-        )
-
         # LoRA w2: applied to intermediate_cache3 before moe_sum, using the
         # unquantized intermediate_cache2 as the lora_a input.  Reuses the
-        # sorted_token_ids_lora computed above.
-        if lora_context is not None:
-            self.apply_w2_lora(
-                lora_context,
-                y=intermediate_cache3,
-                x=intermediate_cache2,
-                topk_weights=topk_weights,
-                sorted_token_ids_lora=sorted_token_ids_lora,
-                expert_ids_lora=expert_ids_lora,
-                num_tokens_post_padded_lora=num_tokens_post_padded_lora,
-                token_lora_mapping=token_lora_mapping,
-                num_tokens=num_tokens,
-                w1=w1,
-                w2=w2,
-                top_k_num=top_k_num,
+        # sorted_token_ids_lora computed above. Same dual-stream pattern as
+        # the w13 pair: base GEMM on default stream, LoRA delta on aux,
+        # join via .add_() into intermediate_cache3.
+        def _base_w2_fn():
+            invoke_fused_moe_triton_kernel(
+                qintermediate_cache2,
+                w2,
+                intermediate_cache3,
+                a2q_scale,
+                self.w2_scale,
+                topk_weights,
+                sorted_token_ids,
+                expert_ids,
+                num_tokens_post_padded,
+                not apply_router_weight_on_input,
+                1,
+                config,
+                compute_type=compute_type,
+                use_fp8_w8a8=self.quant_config.use_fp8_w8a8,
+                use_int8_w8a8=self.quant_config.use_int8_w8a8,
+                use_int8_w8a16=self.quant_config.use_int8_w8a16,
+                use_int4_w4a16=self.quant_config.use_int4_w4a16,
+                per_channel_quant=self.per_act_token_quant,
+                block_shape=self.block_shape,
+                B_bias=self.w2_bias,
             )
 
+        if lora_context is not None and lora_context.aux_stream is not None:
+            lora_delta_w2 = torch.zeros_like(intermediate_cache3)
+
+            def _lora_w2_fn():
+                self.apply_w2_lora(
+                    lora_context,
+                    y=lora_delta_w2,
+                    x=intermediate_cache2,
+                    topk_weights=topk_weights,
+                    sorted_token_ids_lora=sorted_token_ids_lora,
+                    expert_ids_lora=expert_ids_lora,
+                    num_tokens_post_padded_lora=num_tokens_post_padded_lora,
+                    token_lora_mapping=token_lora_mapping,
+                    num_tokens=num_tokens,
+                    w1=w1,
+                    w2=w2,
+                    top_k_num=top_k_num,
+                    add_inputs=False,
+                )
+
+            assert lora_context.events is not None
+            maybe_execute_in_parallel(
+                _base_w2_fn,
+                _lora_w2_fn,
+                lora_context.events[2],
+                lora_context.events[3],
+                lora_context.aux_stream,
+            )
+            intermediate_cache3.add_(lora_delta_w2)
+        else:
+            _base_w2_fn()
+            if lora_context is not None:
+                self.apply_w2_lora(
+                    lora_context,
+                    y=intermediate_cache3,
+                    x=intermediate_cache2,
+                    topk_weights=topk_weights,
+                    sorted_token_ids_lora=sorted_token_ids_lora,
+                    expert_ids_lora=expert_ids_lora,
+                    num_tokens_post_padded_lora=num_tokens_post_padded_lora,
+                    token_lora_mapping=token_lora_mapping,
+                    num_tokens=num_tokens,
+                    w1=w1,
+                    w2=w2,
+                    top_k_num=top_k_num,
+                )
+
         # separate function is required for MoE + LoRA
         self.moe_sum(intermediate_cache3, output)
 

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/fused_moe/runner/moe_runner.py

@@ -164,6 +164,8 @@ def _moe_forward_shared_fake(
     return shared_out, fused_out
 
 
+# NOTE: `moe_forward` and `moe_forward_shared` being opaque custom ops is a
+# load-bearing assumption for the MoE-LoRA dual-stream path.
 direct_register_custom_op(
     op_name="moe_forward",
     op_func=_moe_forward,

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

@@ -339,7 +339,7 @@ class GPT2ForSequenceClassification(nn.Module, SupportsCrossEncoding):
         super().__init__()
         config = vllm_config.model_config.hf_config
         self.transformer = GPT2Model(
-            vllm_config=vllm_config, prefix=maybe_prefix(prefix, "gpt2")
+            vllm_config=vllm_config, prefix=maybe_prefix(prefix, "transformer")
         )
         self.score = nn.Linear(
             config.n_embd,
@@ -358,6 +358,7 @@ class GPT2ForSequenceClassification(nn.Module, SupportsCrossEncoding):
 
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
         loader = AutoWeightsLoader(self)
+        weights = _add_transformer_prefix(weights)
         return loader.load_weights(weights)
 
     def forward(
@@ -380,6 +381,10 @@ def _add_transformer_prefix(
     weights: Iterable[tuple[str, torch.Tensor]],
 ) -> Iterable[tuple[str, torch.Tensor]]:
     for name, tensor in weights:
-        if not name.startswith("transformer.") and not name.startswith("lm_head"):
+        if (
+            not name.startswith("transformer.")
+            and not name.startswith("lm_head.")
+            and not name.startswith("score.")
+        ):
             name = "transformer." + name
         yield name, tensor

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

@@ -378,7 +378,9 @@ class Resampler4_5(Resampler2_5):
                     )  # D
 
             pos_embed_2d.append(
-                self.pos_embed[:tgt_h, :tgt_w, :].reshape((tgt_h * tgt_w, -1)).to(dtype)
+                self.pos_embed[:tgt_h, :tgt_w, :]
+                .reshape((tgt_h * tgt_w, -1))
+                .to(device=device, dtype=dtype)
             )  # patches * D
             key_padding_mask[i, patch_len[i] :] = True
 

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,14 +54,14 @@ from vllm.model_executor.models.utils import (
     maybe_prefix,
 )
 from vllm.model_executor.utils import set_weight_attrs
-from vllm.models.deepseek_v4.nvidia.ops.attention import (
+from vllm.models.deepseek_v4.attention import (
     DeepseekV4Indexer,
     DeepseekV4MLAModules,
     DeepseekV4MultiHeadLatentAttentionWrapper,
 )
+from vllm.models.deepseek_v4.nvidia.ops.prepare_megamoe import prepare_megamoe_inputs
 from vllm.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
 
 
@@ -116,167 +116,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]]:
@@ -542,7 +381,7 @@ class DeepseekV4MegaMoEExperts(nn.Module):
 
         symm_buffer = self.get_symm_buffer()
         num_tokens = hidden_states.shape[0]
-        _stage_deepseek_v4_mega_moe_inputs(
+        prepare_megamoe_inputs(
             hidden_states,
             topk_weights,
             topk_ids,

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,