[CI] Filter import-only files using function-level coverage

Skip files where only module-level code ran (imports, class defs)
but no named functions were actually called. Uses the
functions_called field from stripped coverage JSON.

Reduces false-positive mappings by ~78% — e.g. ompmultiprocessing.py
drops from 73 steps to 0 (only used on CPU but imported everywhere).

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

.buildkite/hardware_tests/cpu.yaml

@@ -54,6 +54,20 @@ 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/coverage/aggregate-coverage.py

@@ -0,0 +1,208 @@
+#!/usr/bin/env python3
+"""Aggregate per-step coverage JSON files into a test-selection mapping.
+
+Downloads all coverage_*.json artifacts from the current Buildkite build,
+then produces two output files:
+
+1. coverage_map.json — inverted index: {source_file: [step_keys]}
+   Used by the pipeline generator to determine which steps to trigger.
+
+2. step_coverage.json — forward index: {step_key: [source_files]}
+   Useful for debugging and understanding test coverage.
+
+Usage:
+    # Run as a Buildkite step at the end of nightly CI
+    python3 .buildkite/scripts/coverage/aggregate-coverage.py
+
+    # Or locally with downloaded artifacts
+    python3 .buildkite/scripts/coverage/aggregate-coverage.py --local-dir ./artifacts/
+"""
+
+import argparse
+import json
+import os
+import subprocess
+import sys
+import tempfile
+from collections import defaultdict
+from pathlib import Path
+
+
+def download_artifacts(dest_dir: str) -> list[str]:
+    """Download all coverage_*.json artifacts from the current build."""
+    try:
+        subprocess.run(
+            ["buildkite-agent", "artifact", "download", "coverage_*.json", dest_dir],
+            check=True,
+            capture_output=True,
+            text=True,
+        )
+    except FileNotFoundError:
+        print("buildkite-agent not found, skipping download", file=sys.stderr)
+        return []
+    except subprocess.CalledProcessError as e:
+        print(f"Artifact download failed: {e.stderr}", file=sys.stderr)
+        return []
+
+    return list(Path(dest_dir).glob("coverage_*.json"))
+
+
+def load_coverage_files(files: list[Path]) -> dict[str, list[str]]:
+    """Load coverage JSON files and extract source files per step.
+
+    Returns: {step_key: [source_files]}
+    """
+    step_coverage = {}
+
+    for filepath in files:
+        filename = filepath.name
+        # coverage_<step_key>.json -> step_key
+        step_key = filename.removeprefix("coverage_").removesuffix(".json")
+
+        try:
+            with open(filepath) as f:
+                data = json.load(f)
+        except (json.JSONDecodeError, OSError) as e:
+            print(f"Warning: skipping {filename}: {e}", file=sys.stderr)
+            continue
+
+        source_files = []
+        for fpath, fdata in data.get("files", {}).items():
+            # Skip files with zero executed lines — coverage.py reports
+            # all files in the source tree, not just those actually run.
+            # Supports both full format (summary.covered_lines) and
+            # stripped format (covered_lines directly).
+            covered = fdata.get("covered_lines") or fdata.get("summary", {}).get("covered_lines", 0)
+            if covered == 0:
+                continue
+            # If function-level data is available, skip import-only files
+            # (files where only module-level code ran but no named functions
+            # were actually called).
+            funcs_called = fdata.get("functions_called")
+            if funcs_called is not None and funcs_called == 0:
+                continue
+            # Normalize paths to be relative to the vllm package root.
+            # coverage.py may report absolute paths or paths relative to
+            # the installed package location. We only care about files
+            # under the vllm/ directory.
+            normalized = _normalize_path(fpath)
+            if normalized:
+                source_files.append(normalized)
+
+        if source_files:
+            step_coverage[step_key] = sorted(set(source_files))
+            print(f"  {step_key}: {len(source_files)} source files")
+
+    return step_coverage
+
+
+def _normalize_path(path: str) -> str | None:
+    """Normalize a coverage path to a vllm-relative path.
+
+    Returns None for paths outside the vllm package (tests, third-party, etc).
+    """
+    # Strip common prefixes from installed package paths
+    markers = ["/site-packages/", "/dist-packages/", "/vllm-workspace/src/"]
+    for marker in markers:
+        idx = path.find(marker)
+        if idx != -1:
+            path = path[idx + len(marker):]
+            break
+
+    # Also handle paths that are already relative
+    if path.startswith("vllm/"):
+        return path
+
+    # Handle absolute paths that contain /vllm/
+    idx = path.find("/vllm/")
+    if idx != -1:
+        return path[idx + 1:]
+
+    return None
+
+
+def build_inverted_index(
+    step_coverage: dict[str, list[str]],
+) -> dict[str, list[str]]:
+    """Build {source_file: [step_keys]} from {step_key: [source_files]}."""
+    inverted = defaultdict(list)
+    for step_key, source_files in step_coverage.items():
+        for src_file in source_files:
+            inverted[src_file].append(step_key)
+
+    # Sort step lists for deterministic output
+    return {k: sorted(v) for k, v in sorted(inverted.items())}
+
+
+def main():
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument(
+        "--local-dir",
+        help="Directory containing coverage_*.json files (skip artifact download)",
+    )
+    parser.add_argument(
+        "--output-dir",
+        default=".",
+        help="Directory to write output files (default: cwd)",
+    )
+    args = parser.parse_args()
+
+    if args.local_dir:
+        artifact_dir = args.local_dir
+        files = list(Path(artifact_dir).glob("coverage_*.json"))
+    else:
+        artifact_dir = tempfile.mkdtemp(prefix="coverage_artifacts_")
+        files = download_artifacts(artifact_dir)
+
+    if not files:
+        print("No coverage files found. Nothing to aggregate.")
+        sys.exit(0)
+
+    print(f"Found {len(files)} coverage files:")
+
+    # Build the forward index: step -> source files
+    step_coverage = load_coverage_files(files)
+
+    if not step_coverage:
+        print("No valid coverage data found.")
+        sys.exit(0)
+
+    # Build the inverted index: source file -> steps
+    coverage_map = build_inverted_index(step_coverage)
+
+    # Write outputs
+    output_dir = Path(args.output_dir)
+    output_dir.mkdir(parents=True, exist_ok=True)
+
+    step_coverage_path = output_dir / "step_coverage.json"
+    with open(step_coverage_path, "w") as f:
+        json.dump(step_coverage, f, indent=2)
+    print(f"\nWrote {step_coverage_path} ({len(step_coverage)} steps)")
+
+    coverage_map_path = output_dir / "coverage_map.json"
+    with open(coverage_map_path, "w") as f:
+        json.dump(coverage_map, f, indent=2)
+    print(f"Wrote {coverage_map_path} ({len(coverage_map)} source files)")
+
+    # Summary stats
+    total_files = len(coverage_map)
+    total_mappings = sum(len(v) for v in coverage_map.values())
+    print(f"\nSummary: {total_files} source files mapped to "
+          f"{len(step_coverage)} steps ({total_mappings} total mappings)")
+
+    # Upload aggregated files as artifacts
+    for output_file in [step_coverage_path, coverage_map_path]:
+        try:
+            subprocess.run(
+                ["buildkite-agent", "artifact", "upload", str(output_file)],
+                check=True,
+                capture_output=True,
+                text=True,
+            )
+            print(f"Uploaded {output_file}")
+        except (FileNotFoundError, subprocess.CalledProcessError):
+            pass  # Not in Buildkite or upload failed — that's fine for local runs
+
+
+if __name__ == "__main__":
+    main()

.buildkite/scripts/coverage/upload-step-coverage.sh

@@ -0,0 +1,42 @@
+#!/bin/bash
+# Upload coverage data for the current Buildkite step.
+# Called automatically at the end of each step when COLLECT_COVERAGE=1.
+#
+# Expects:
+#   - .coverage.${BUILDKITE_STEP_KEY} data file from coverage run --append
+#   - BUILDKITE_STEP_KEY, BUILDKITE_BUILD_NUMBER env vars
+#
+# Produces:
+#   - coverage_${BUILDKITE_STEP_KEY}.json uploaded as a Buildkite artifact
+
+set -euo pipefail
+
+STEP_KEY="${BUILDKITE_STEP_KEY:-unknown}"
+DATA_FILE=".coverage.${STEP_KEY}"
+OUTPUT_JSON="coverage_${STEP_KEY}.json"
+
+if [ ! -f "$DATA_FILE" ]; then
+    echo "~~~ No coverage data file found ($DATA_FILE), skipping upload"
+    exit 0
+fi
+
+echo "~~~ :bar_chart: Exporting coverage data for step: ${STEP_KEY}"
+
+coverage json \
+    --data-file="$DATA_FILE" \
+    -o "$OUTPUT_JSON" \
+    --omit='*/tests/*,*/test_*,*/__pycache__/*' \
+    2>&1 || {
+        echo "Warning: coverage json export failed, skipping"
+        exit 0
+    }
+
+FILE_COUNT=$(python3 -c "import json; d=json.load(open('$OUTPUT_JSON')); print(len(d.get('files', {})))" 2>/dev/null || echo "?")
+echo "Coverage captured ${FILE_COUNT} source files for step ${STEP_KEY}"
+
+buildkite-agent artifact upload "$OUTPUT_JSON" 2>&1 || {
+    echo "Warning: artifact upload failed"
+    exit 0
+}
+
+echo "Uploaded $OUTPUT_JSON"

.buildkite/test_areas/kernels.yaml

@@ -38,28 +38,6 @@ steps:
   commands:
     - pytest -v -s kernels/core/test_minimax_reduce_rms.py
 
-- label: Deepseek V4 Kernel Test (H100)
-  key: deepseek-v4-kernel-test-h100
-  timeout_in_minutes: 15
-  device: h100
-  source_file_dependencies:
-  - csrc/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu
-  - vllm/models/deepseek_v4/common/ops/
-  - tests/kernels/test_fused_deepseek_v4_qnorm_rope_kv_insert.py
-  commands:
-    - pytest -v -s kernels/test_fused_deepseek_v4_*.py
-
-- label: Deepseek V4 Kernel Test (B200)
-  key: deepseek-v4-kernel-test-b200
-  timeout_in_minutes: 15
-  device: b200-k8s
-  source_file_dependencies:
-  - csrc/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu
-  - vllm/models/deepseek_v4/common/ops/
-  - tests/kernels/test_fused_deepseek_v4_qnorm_rope_kv_insert.py
-  commands:
-    - pytest -v -s kernels/test_fused_deepseek_v4_*.py
-
 - label: Kernels Attention Test %N
   key: kernels-attention-test
   timeout_in_minutes: 35

csrc/cpu/cpu_attn_impl.hpp

@@ -408,9 +408,19 @@ class AttentionScheduler {
     const int64_t cache_size = cpu_utils::get_available_l2_size();
     const int32_t max_num_q_per_iter = input.max_num_q_per_iter;
     const int32_t kv_len_alignment = input.kv_block_alignment;
+    bool has_decode_request = false;
+    bool decode_only_batch = true;
+    for (int32_t req_id = 0; req_id < input.num_reqs; ++req_id) {
+      const int32_t q_token_num =
+          input.query_start_loc[req_id + 1] - input.query_start_loc[req_id];
+      has_decode_request = has_decode_request || (q_token_num == 1);
+      decode_only_batch = decode_only_batch && (q_token_num == 1);
+    }
     int32_t q_head_per_kv = input.num_heads_q / input.num_heads_kv;
-    const bool use_gqa = (max_num_q_per_iter % q_head_per_kv == 0);
-    if (!use_gqa) {
+    const bool supports_gqa = q_head_per_kv <= max_num_q_per_iter;
+    const bool use_gqa_fast_path = supports_gqa && decode_only_batch;
+    const bool use_gqa_scratchpad = supports_gqa && has_decode_request;
+    if (!use_gqa_scratchpad) {
       q_head_per_kv = 1;  // fallback to MHA
     }
     const int32_t min_split_kv_len =
@@ -680,7 +690,7 @@ class AttentionScheduler {
         metadata_ptr->attention_scratchpad_size_per_thread *
             metadata_ptr->thread_num +
         metadata_ptr->reduction_scratchpad_size_per_kv_head *
-            (use_gqa ? input.num_heads_kv : input.num_heads_q);
+            (use_gqa_fast_path ? input.num_heads_kv : input.num_heads_q);
     cpu_utils::ScratchPadManager::get_scratchpad_manager()->realloc(
         scratchpad_size);
 
@@ -1409,13 +1419,24 @@ class AttentionMainLoop {
       const int32_t q_head_num = input->num_heads;
       const int32_t kv_head_num = input->num_kv_heads;
       const int32_t q_heads_per_kv = q_head_num / kv_head_num;
-      const bool use_gqa =
-          (max_q_head_num_per_iter % q_heads_per_kv == 0) ? true : false;
-      const int32_t actual_kv_head_num = use_gqa ? kv_head_num : q_head_num;
-      const int32_t actual_q_heads_per_kv = use_gqa ? q_heads_per_kv : 1;
+      AttentionWorkItemGroup* const workitem_groups =
+          metadata.workitem_groups_ptr;
+      const int32_t* cu_workitem_num_per_thread =
+          metadata.cu_workitem_num_per_thread;
+      ReductionWorkItemGroup* const reduction_items =
+          metadata.reduction_items_ptr;
+      const bool supports_gqa = q_heads_per_kv <= max_q_head_num_per_iter;
+      bool decode_only_batch = true;
+      for (int32_t i = 0; i < metadata.workitem_group_num; ++i) {
+        decode_only_batch =
+            decode_only_batch && (workitem_groups[i].q_token_num == 1);
+      }
+      const bool use_gqa_fast_path = supports_gqa && decode_only_batch;
+      const int32_t actual_kv_head_num =
+          use_gqa_fast_path ? kv_head_num : q_head_num;
+      const int32_t actual_q_heads_per_kv =
+          use_gqa_fast_path ? q_heads_per_kv : 1;
       TORCH_CHECK_LE(actual_q_heads_per_kv, max_q_head_num_per_iter);
-      const int32_t max_q_token_num_per_iter =
-          max_q_head_num_per_iter / actual_q_heads_per_kv;
       const int64_t q_token_num_stride = input->query_num_tokens_stride;
       const int64_t q_head_num_stride = input->query_num_heads_stride;
       const int64_t kv_cache_head_num_stride = input->cache_num_kv_heads_stride;
@@ -1461,15 +1482,6 @@ class AttentionMainLoop {
               sizeof(q_buffer_t), sizeof(logits_buffer_t),
               sizeof(partial_output_buffer_t), max_q_head_num_per_iter,
               max_q_head_num_per_iter);
-      const int32_t default_q_tile_token_num =
-          default_tile_size / actual_q_heads_per_kv;
-
-      AttentionWorkItemGroup* const workitem_groups =
-          metadata.workitem_groups_ptr;
-      const int32_t* cu_workitem_num_per_thread =
-          metadata.cu_workitem_num_per_thread;
-      ReductionWorkItemGroup* const reduction_items =
-          metadata.reduction_items_ptr;
 
       const int32_t effective_thread_num = metadata.effective_thread_num;
       const int32_t reduction_item_num = metadata.reduction_item_num;
@@ -1513,8 +1525,6 @@ class AttentionMainLoop {
               cu_workitem_num_per_thread[thread_offset + 1] -
               cu_workitem_num_per_thread[thread_offset];
 
-          const int32_t q_head_start_idx = kv_head_idx * actual_q_heads_per_kv;
-
           for (int32_t workitem_group_idx = 0;
                workitem_group_idx < curr_workitem_groups_num;
                ++workitem_group_idx) {
@@ -1529,6 +1539,21 @@ class AttentionMainLoop {
             const int32_t q_token_id_start =
                 current_workitem_group->q_token_id_start;
             const int32_t q_token_num = current_workitem_group->q_token_num;
+            const bool curr_use_gqa =
+                use_gqa_fast_path || (supports_gqa && q_token_num == 1);
+            if (!use_gqa_fast_path && curr_use_gqa &&
+                kv_head_idx % q_heads_per_kv != 0) {
+              continue;
+            }
+            const int32_t curr_q_heads_per_kv =
+                curr_use_gqa ? q_heads_per_kv : 1;
+            const int32_t curr_max_q_token_num_per_iter =
+                max_q_head_num_per_iter / curr_q_heads_per_kv;
+            const int32_t curr_default_q_tile_token_num =
+                default_tile_size / curr_q_heads_per_kv;
+            const int32_t q_head_start_idx =
+                use_gqa_fast_path ? (kv_head_idx * q_heads_per_kv)
+                                  : kv_head_idx;
 
             // taskgroup general information
             const int32_t q_end = input->query_start_loc[current_group_idx + 1];
@@ -1542,7 +1567,7 @@ class AttentionMainLoop {
                              current_workitem_group->local_split_id == 0);
 
             for (int32_t q_token_offset = 0; q_token_offset < q_token_num;
-                 q_token_offset += default_q_tile_token_num) {
+                 q_token_offset += curr_default_q_tile_token_num) {
               bool first_iter_flag[AttentionScheduler::MaxQTileIterNum];
               for (int32_t i = 0; i < AttentionScheduler::MaxQTileIterNum;
                    ++i) {
@@ -1552,9 +1577,9 @@ class AttentionMainLoop {
               const int32_t q_token_start_idx =
                   q_start + q_token_offset + q_token_id_start;
               const int32_t actual_q_token_num = std::min(
-                  default_q_tile_token_num, q_token_num - q_token_offset);
+                  curr_default_q_tile_token_num, q_token_num - q_token_offset);
               const int32_t q_head_tile_size =
-                  actual_q_token_num * actual_q_heads_per_kv;
+                  actual_q_token_num * curr_q_heads_per_kv;
               const int32_t rounded_q_head_tile_size =
                   ((q_head_tile_size + max_q_head_num_per_iter - 1) /
                    max_q_head_num_per_iter) *
@@ -1591,10 +1616,9 @@ class AttentionMainLoop {
                   AttentionScheduler::align_kv_tile_pos(
                       kv_tile_start_pos, kv_tile_end_pos, blocksize_alignment);
 
-              int32_t curr_kv_head_idx =
-                  use_gqa ? kv_head_idx
-                          : (kv_head_idx /
-                             q_heads_per_kv);  // for GQA disabled case
+              const int32_t curr_kv_head_idx =
+                  use_gqa_fast_path ? kv_head_idx
+                                    : (kv_head_idx / q_heads_per_kv);
 
               // std::printf("thread_id: %d, req_id: %d, q_token_start: %d,
               // q_token_end: %d, q_head_start: %d, q_head_end: %d, kv_head_idx:
@@ -1629,12 +1653,12 @@ class AttentionMainLoop {
                   (s_aux != nullptr ? s_aux + q_head_start_idx : nullptr);
 
               // copy the Q tile to q_buffer, the logical layout of q_buffer is
-              // [actual_q_token_num, actual_q_heads_per_kv, head_dim]
+              // [actual_q_token_num, curr_q_heads_per_kv, head_dim]
               {
                 attn_impl.copy_q_heads_tile(
                     q_tile_ptr, q_buffer, actual_q_token_num,
-                    actual_q_heads_per_kv, q_token_num_stride,
-                    q_head_num_stride, scale);
+                    curr_q_heads_per_kv, q_token_num_stride, q_head_num_stride,
+                    scale);
               }
 
               if (use_sink) {
@@ -1648,29 +1672,29 @@ class AttentionMainLoop {
                 float* __restrict__ curr_max_buffer = max_buffer;
                 for (int32_t token_idx = 0; token_idx < actual_q_token_num;
                      ++token_idx) {
-                  for (int32_t head_idx = 0; head_idx < actual_q_heads_per_kv;
+                  for (int32_t head_idx = 0; head_idx < curr_q_heads_per_kv;
                        ++head_idx) {
                     curr_sum_buffer[head_idx] = 1.0f;
                     curr_max_buffer[head_idx] = s_aux_fp32[head_idx];
                   }
 
-                  curr_sum_buffer += actual_q_heads_per_kv;
-                  curr_max_buffer += actual_q_heads_per_kv;
+                  curr_sum_buffer += curr_q_heads_per_kv;
+                  curr_max_buffer += curr_q_heads_per_kv;
                 }
               } else {
                 float* __restrict__ curr_sum_buffer = sum_buffer;
                 float* __restrict__ curr_max_buffer = max_buffer;
                 for (int32_t token_idx = 0; token_idx < actual_q_token_num;
                      ++token_idx) {
-                  for (int32_t head_idx = 0; head_idx < actual_q_heads_per_kv;
+                  for (int32_t head_idx = 0; head_idx < curr_q_heads_per_kv;
                        ++head_idx) {
                     curr_sum_buffer[head_idx] = 0.0f;
                     curr_max_buffer[head_idx] =
                         std::numeric_limits<float>::lowest();
                   }
 
-                  curr_sum_buffer += actual_q_heads_per_kv;
-                  curr_max_buffer += actual_q_heads_per_kv;
+                  curr_sum_buffer += curr_q_heads_per_kv;
+                  curr_max_buffer += curr_q_heads_per_kv;
                 }
               }
 
@@ -1683,16 +1707,17 @@ class AttentionMainLoop {
                     kv_tile_pos_left + kv_tile_size, rounded_kv_tile_end_pos);
                 for (int32_t q_head_tile_token_offset = 0;
                      q_head_tile_token_offset < actual_q_token_num;
-                     q_head_tile_token_offset += max_q_token_num_per_iter) {
+                     q_head_tile_token_offset +=
+                     curr_max_q_token_num_per_iter) {
                   const int32_t q_tile_pos_left =
                       q_tile_start_pos + q_head_tile_token_offset;
                   const int32_t q_tile_token_num =
-                      std::min(max_q_token_num_per_iter,
+                      std::min(curr_max_q_token_num_per_iter,
                                actual_q_token_num - q_head_tile_token_offset);
                   const int32_t q_tile_head_offset =
-                      q_head_tile_token_offset * actual_q_heads_per_kv;
+                      q_head_tile_token_offset * curr_q_heads_per_kv;
                   const int32_t q_tile_head_num =
-                      q_tile_token_num * actual_q_heads_per_kv;
+                      q_tile_token_num * curr_q_heads_per_kv;
                   const int32_t q_tile_pos_right =
                       q_tile_pos_left + q_tile_token_num;
                   const auto [actual_kv_tile_pos_left,
@@ -1702,7 +1727,7 @@ class AttentionMainLoop {
                           q_tile_pos_right, sliding_window_left,
                           sliding_window_right);
                   const int32_t q_iter_idx =
-                      q_head_tile_token_offset / max_q_token_num_per_iter;
+                      q_head_tile_token_offset / curr_max_q_token_num_per_iter;
 
                   if (actual_kv_tile_pos_right <= actual_kv_tile_pos_left) {
                     continue;
@@ -1768,7 +1793,7 @@ class AttentionMainLoop {
                       aligned_actual_kv_tile_pos_left,
                       aligned_actual_kv_tile_pos_right, actual_kv_token_num,
                       kv_cache_block_num_stride, q_tile_head_num,
-                      q_tile_token_num, q_tile_pos_left, actual_q_heads_per_kv,
+                      q_tile_token_num, q_tile_pos_left, curr_q_heads_per_kv,
                       block_size, sliding_window_left, sliding_window_right,
                       scale, softcap_scale, curr_alibi_slopes,
                       first_iter_flag[q_iter_idx], use_sink, debug_info);
@@ -1782,11 +1807,11 @@ class AttentionMainLoop {
                   final_output(partial_q_buffer,
                                reinterpret_cast<query_t*>(input->output) +
                                    output_buffer_offset,
-                               sum_buffer, actual_q_heads_per_kv,
+                               sum_buffer, curr_q_heads_per_kv,
                                actual_q_token_num, q_head_num, output_v_scale);
                 } else {
                   const int32_t stride =
-                      actual_q_heads_per_kv * split_kv_q_token_num_threshold;
+                      curr_q_heads_per_kv * split_kv_q_token_num_threshold;
                   buffer_manager.update(kv_head_idx, total_reduction_split_num,
                                         head_dim, stride, sizeof(float));
                   volatile bool* split_flag_buffer =
@@ -1822,18 +1847,26 @@ class AttentionMainLoop {
           const int32_t curr_split_id = curr_workitem_groups->split_start_id;
           const int32_t curr_split_num = curr_workitem_groups->split_num;
           const int32_t current_group_idx = curr_workitem_groups->req_id;
+          const bool curr_use_gqa =
+              use_gqa_fast_path || (supports_gqa && curr_output_token_num == 1);
+          if (!use_gqa_fast_path && curr_use_gqa &&
+              kv_head_idx % q_heads_per_kv != 0) {
+            continue;
+          }
+          const int32_t curr_q_heads_per_kv = curr_use_gqa ? q_heads_per_kv : 1;
           const int32_t curr_output_head_num =
-              curr_output_token_num * actual_q_heads_per_kv;
+              curr_output_token_num * curr_q_heads_per_kv;
 
           const int32_t q_start = input->query_start_loc[current_group_idx];
           const int32_t q_token_start_idx = q_start + curr_output_token_idx;
-          const int32_t q_head_start_idx = kv_head_idx * actual_q_heads_per_kv;
+          const int32_t q_head_start_idx =
+              use_gqa_fast_path ? (kv_head_idx * q_heads_per_kv) : kv_head_idx;
           size_t output_buffer_offset =
               q_token_start_idx * q_head_num * head_dim +
               q_head_start_idx * head_dim;
 
           const int32_t stride =
-              actual_q_heads_per_kv * split_kv_q_token_num_threshold;
+              curr_q_heads_per_kv * split_kv_q_token_num_threshold;
           buffer_manager.update(kv_head_idx, total_reduction_split_num,
                                 head_dim, stride, sizeof(float));
           volatile bool* split_flag_buffer =
@@ -1852,7 +1885,7 @@ class AttentionMainLoop {
           final_output(
               split_output_buffer,
               reinterpret_cast<query_t*>(input->output) + output_buffer_offset,
-              split_sum_buffer, actual_q_heads_per_kv, curr_output_token_num,
+              split_sum_buffer, curr_q_heads_per_kv, curr_output_token_num,
               q_head_num, output_v_scale);
         }
       }

csrc/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu

@@ -122,45 +122,18 @@ __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.  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>
+// token
+template <typename scalar_t_in>
 __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_out,
+    float const eps, scalar_t_in* __restrict__ q_inout,
     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 == 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;
-  }
+  bool const isKV = (slotIdx == num_heads_q);
 
   // ── Decode the bf16 → 16 fp32 registers ─────────────────────────────
   float elements[kElemsPerLane];
@@ -234,7 +207,7 @@ __device__ __forceinline__ void processDeepseekV4Slot(
   // triggering and per-iteration buffer rotation.
   // ═══════════════════════════════════════════════════════════════════
   if (!isKV) {
-    // ── Live-Q: cast back to bf16 and store into the padded q_out. ─────
+    // ── Q: cast back to bf16 and store. ────────────────────────────
     uint4 out0, out1;
     typename Converter::packed_hip_type* po0 =
         reinterpret_cast<typename Converter::packed_hip_type*>(&out0);
@@ -251,9 +224,8 @@ __device__ __forceinline__ void processDeepseekV4Slot(
           make_float2(elements[8 + 2 * i], elements[8 + 2 * i + 1]));
     }
     scalar_t_in* dst =
-        q_out +
-        (static_cast<int64_t>(tokenIdx) * kNumHeadsQPadded + slotIdx) *
-            kHeadDim +
+        q_inout +
+        (static_cast<int64_t>(tokenIdx) * num_heads_q + slotIdx) * kHeadDim +
         dim_base;
     *reinterpret_cast<uint4*>(dst) = out0;
     *reinterpret_cast<uint4*>(dst + 8) = out1;
@@ -341,32 +313,20 @@ __device__ __forceinline__ void processDeepseekV4Slot(
 // Kernel
 // ────────────────────────────────────────────────────────────────────────────
 //
-// Grid: 1D, gridDim.x = ceil(num_tokens_full * (kNumHeadsQPadded + 1) /
+// Grid: 1D, gridDim.x = ceil(num_tokens_full * (num_heads_q + 1) /
 // warps_per_block) Block: blockDim.x = 256 threads (8 warps per block) Each
-// 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.
+// 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)
 //
 // With DP padding, q/kv/position_ids can have more rows than slot_mapping.
-// 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.
+// 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.
 //
-template <typename scalar_t_in, int kNumHeadsQPadded>
+template <typename scalar_t_in>
 __global__ void fusedDeepseekV4QNormRopeKVRopeQuantInsertKernel(
-    scalar_t_in const* __restrict__ q_in,      // [N, num_heads_q,      512]
-    scalar_t_in* __restrict__ q_out,           // [N, kNumHeadsQPadded, 512]
+    scalar_t_in* __restrict__ q_inout,         // [N, H, 512] bf16, in place
     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
@@ -375,7 +335,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,        // live Q heads (input layout)
+    int const num_heads_q,        // H
     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)
@@ -391,13 +351,12 @@ __global__ void fusedDeepseekV4QNormRopeKVRopeQuantInsertKernel(
     int const laneId = threadIdx.x % 32;
     int const globalWarpIdx = blockIdx.x * warpsPerBlock + warpId;
 
-    constexpr int kTotalSlotsPerToken = kNumHeadsQPadded + 1;
-    int const tokenIdx = globalWarpIdx / kTotalSlotsPerToken;
-    int const slotIdx = globalWarpIdx % kTotalSlotsPerToken;
+    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;
     if (tokenIdx >= num_tokens_full) return;
 
-    bool const isKV = (slotIdx == kNumHeadsQPadded);
-    bool const isPadQ = !isKV && (slotIdx >= num_heads_q);
+    bool const isKV = (slotIdx == num_heads_q);
     // KV branch: skip DP-padded tokens (no slot reserved for them).
     if (isKV && tokenIdx >= num_tokens_insert) return;
 
@@ -412,26 +371,22 @@ __global__ void fusedDeepseekV4QNormRopeKVRopeQuantInsertKernel(
     // Dim range this lane owns within the 512-wide head.
     int const dim_base = laneId * kElemsPerLane;  // in [0, 512) step 16
 
-    // 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);
+    // 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;
     }
+    uint4 const v0 = *reinterpret_cast<uint4 const*>(src_ptr);
+    uint4 const v1 = *reinterpret_cast<uint4 const*>(src_ptr + 8);
 
-    processDeepseekV4Slot<scalar_t_in, kNumHeadsQPadded>(
-        v0, v1, tokenIdx, slotIdx, dim_base, laneId, num_heads_q, eps, q_out,
+    processDeepseekV4Slot<scalar_t_in>(
+        v0, v1, tokenIdx, slotIdx, dim_base, laneId, num_heads_q, eps, q_inout,
         k_cache, slot_mapping, position_ids, cos_sin_cache, cache_block_size,
         kv_block_stride);
 
@@ -453,9 +408,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, int kNumHeadsQPadded>
+template <typename scalar_t_in>
 __global__ void fusedDeepseekV4QNormRopeKVRopeQuantInsertKernelReducedGrid(
-    scalar_t_in const* __restrict__ q_in, scalar_t_in* __restrict__ q_out,
+    scalar_t_in* __restrict__ q_inout,  // [N, H, 512] bf16, in place
     scalar_t_in const* __restrict__ kv_in, uint8_t* __restrict__ k_cache,
     int64_t const* __restrict__ slot_mapping,
     int64_t const* __restrict__ position_ids,
@@ -480,32 +435,25 @@ __global__ void fusedDeepseekV4QNormRopeKVRopeQuantInsertKernelReducedGrid(
 #endif
 
     int const dim_base = laneId * kElemsPerLane;  // in [0, 512) step 16
-    // Slot enumeration: live-Q + pad-Q + (KV if this token has a slot).
-    int const slot_end = (tokenIdx >= num_tokens_insert)
-                             ? kNumHeadsQPadded
-                             : (kNumHeadsQPadded + 1);
+    int const slot_end =
+        (tokenIdx >= num_tokens_insert) ? num_heads_q : (num_heads_q + 1);
 
-    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;
+    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;
       }
-      va = *reinterpret_cast<uint4 const*>(src);
-      vb = *reinterpret_cast<uint4 const*>(src + 8);
+      return q_inout +
+             (static_cast<int64_t>(tokenIdx) * num_heads_q +
+              static_cast<int64_t>(s)) *
+                 kHeadDim +
+             dim_base;
     };
 
     if (warpId < slot_end) {
       int curr_slot = warpId;
-      uint4 v0_curr, v1_curr;
-      load_slot(curr_slot, v0_curr, v1_curr);
+      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);
 
       while (curr_slot < slot_end) {
         int const next_slot = curr_slot + warpsPerBlock;
@@ -514,12 +462,14 @@ __global__ void fusedDeepseekV4QNormRopeKVRopeQuantInsertKernelReducedGrid(
         // Prefetch src for the next slot
         uint4 v0_next, v1_next;
         if (has_next) {
-          load_slot(next_slot, v0_next, v1_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);
         }
 
-        processDeepseekV4Slot<scalar_t_in, kNumHeadsQPadded>(
+        processDeepseekV4Slot<scalar_t_in>(
             v0_curr, v1_curr, tokenIdx, curr_slot, dim_base, laneId,
-            num_heads_q, eps, q_out, k_cache, slot_mapping, position_ids,
+            num_heads_q, eps, q_inout, 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.
@@ -540,10 +490,10 @@ __global__ void fusedDeepseekV4QNormRopeKVRopeQuantInsertKernelReducedGrid(
 // ────────────────────────────────────────────────────────────────────────────
 // Launch wrapper
 // ────────────────────────────────────────────────────────────────────────────
-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,
+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,
     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,
@@ -551,7 +501,7 @@ static void launchFusedDeepseekV4Templated(
   constexpr int kBlockSize = 256;
   constexpr int kWarpsPerBlock = kBlockSize / 32;
   int64_t const total_warps =
-      static_cast<int64_t>(num_tokens_full) * (kNumHeadsQPadded + 1);
+      static_cast<int64_t>(num_tokens_full) * (num_heads_q + 1);
   int const grid =
       static_cast<int>((total_warps + kWarpsPerBlock - 1) / kWarpsPerBlock);
 
@@ -582,85 +532,46 @@ static void launchFusedDeepseekV4Templated(
 
   if (num_tokens_full < NUM_TOKEN_CUTOFF) {
     cudaLaunchKernelEx(
-        &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,
+        &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,
         kv_block_stride);
   } else {
     config.gridDim = dim3(num_tokens_full);
     cudaLaunchKernelEx(
         &config,
-        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,
+        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,
         kv_block_stride);
   }
+
 #else
   // ROCm: use standard kernel launch syntax (no PDL/stream serialization)
   // clang-format off
-  fusedDeepseekV4QNormRopeKVRopeQuantInsertKernel<scalar_t_in, kNumHeadsQPadded>
+  fusedDeepseekV4QNormRopeKVRopeQuantInsertKernel<scalar_t_in>
       <<<grid, kBlockSize, 0, stream>>>(
-          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,
+          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, 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
 // ────────────────────────────────────────────────────────────────────────────
-torch::Tensor fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
-    torch::Tensor const& q_in,           // [N, num_heads_q, 512] bf16
+void fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
+    torch::Tensor& q,                    // [N, H, 512] bf16, in place
     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_in.is_cuda() && q_in.is_contiguous(),
-              "q_in must be contiguous CUDA");
+  TORCH_CHECK(q.is_cuda() && q.is_contiguous(), "q 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,
@@ -668,12 +579,9 @@ torch::Tensor 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_in.dim() == 3 && q_in.size(2) == 512,
-              "q_in shape [N, num_heads_q, 512]");
+  TORCH_CHECK(q.dim() == 3 && q.size(2) == 512, "q shape [N, H, 512]");
   TORCH_CHECK(kv.dim() == 2 && kv.size(1) == 512, "kv shape [N, 512]");
-  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(q.dtype() == kv.dtype(), "q and kv dtype must match");
   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]");
@@ -683,41 +591,32 @@ torch::Tensor 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_in.size(0));
+  int const num_tokens_full = static_cast<int>(q.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_in.size(1));
-  int const num_heads_q_padded = static_cast<int>(q_head_padded);
+  int const num_heads_q = static_cast<int>(q.size(1));
   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_in));
+  at::cuda::OptionalCUDAGuard device_guard(device_of(q));
   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_in.scalar_type(), "fused_deepseek_v4_qnorm_rope_kv_insert", [&] {
+      q.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 const*>(q_in.data_ptr()),
-                reinterpret_cast<qkv_scalar_t*>(q_out.data_ptr()),
+                reinterpret_cast<qkv_scalar_t*>(q.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,
-                num_heads_q_padded, cache_block_size_i, kv_block_stride,
-                stream);
+                cache_block_size_i, kv_block_stride, stream);
       });
-  return q_out;
-}
+}

csrc/ops.h

@@ -70,11 +70,10 @@ 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);
 
-torch::Tensor fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
-    torch::Tensor const& q_in, torch::Tensor const& kv, torch::Tensor& k_cache,
+void fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
+    torch::Tensor& q, torch::Tensor const& kv, torch::Tensor& k_cache,
     torch::Tensor const& slot_mapping, torch::Tensor const& position_ids,
-    torch::Tensor const& cos_sin_cache, int64_t q_head_padded, double eps,
-    int64_t cache_block_size);
+    torch::Tensor const& cos_sin_cache, 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_in, Tensor kv, Tensor! k_cache, "
+      "Tensor! q, Tensor kv, Tensor! k_cache, "
       "Tensor slot_mapping, Tensor position_ids, Tensor cos_sin_cache, "
-      "int q_head_padded, float eps, int cache_block_size) -> Tensor");
+      "float eps, int cache_block_size) -> ()");
   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,8 +220,7 @@ 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/^humming-kernels\[cu13\]/humming-kernels[cu12]/' requirements/cuda.txt; \
+        sed -i 's/^nvidia-cutlass-dsl\[cu13\]>=/nvidia-cutlass-dsl>=/' requirements/cuda.txt; \
     fi \
     && if [ "${PYTORCH_NIGHTLY}" = "1" ]; then \
         echo "Installing torch nightly..." \
@@ -747,8 +746,7 @@ 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/^humming-kernels\[cu13\]/humming-kernels[cu12]/' /tmp/requirements-cuda.txt; \
+        sed -i 's/^nvidia-cutlass-dsl\[cu13\]>=/nvidia-cutlass-dsl>=/' /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 '.') && \

docker/Dockerfile.cpu

@@ -27,11 +27,14 @@ 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 \
+    && apt-get install -y --no-install-recommends sudo ccache git curl wget ca-certificates zlib1g-dev \
     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
@@ -123,9 +126,6 @@ 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,8 +257,7 @@ 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 dist/*.whl && \
-    uv pip install "vllm[audio]"
+    uv pip install "$(realpath dist/*.whl)[audio,triton-cpu]"
 
 # Add labels to document build configuration
 LABEL org.opencontainers.image.title="vLLM CPU"

docs/features/reasoning_outputs.md

@@ -15,7 +15,6 @@ 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` | ✅ |
@@ -30,7 +29,6 @@ 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`.
 
@@ -316,44 +314,9 @@ 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`), Anthropic Messages API (`/v1/messages`) and the Responses API (`/v1/responses`).
+- The reasoning content is only available for online serving's chat completion endpoint (`/v1/chat/completions`).
 
 ## How to support a new reasoning model
 

requirements/build/rocm.txt

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

requirements/cuda.txt

@@ -21,11 +21,8 @@ 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.2
+nvidia-cutlass-dsl[cu13]==4.5.0
 quack-kernels>=0.3.3
 
 # Tokenspeed_MLA for faster mla with spec decode
-tokenspeed-mla==0.1.2
-
-# Humming kernels for quantization gemm
-humming-kernels[cu13]==0.1.2
+tokenspeed-mla==0.1.2

requirements/kv_connectors.txt

@@ -1,6 +1,3 @@
 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,4 +22,3 @@ 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,7 +43,6 @@ 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,9 +43,7 @@ anyio==4.13.0
     #   starlette
     #   watchfiles
 apache-tvm-ffi==0.1.10
-    # via
-    #   tilelang
-    #   xgrammar
+    # via xgrammar
 arctic-inference==0.1.1
     # via -r requirements/test/rocm.in
 argcomplete==3.6.3
@@ -131,9 +129,7 @@ click==8.3.1
     #   typer
     #   uvicorn
 cloudpickle==3.1.2
-    # via
-    #   -r requirements/test/../common.txt
-    #   tilelang
+    # via -r requirements/test/../common.txt
 colorama==0.4.6
     # via
     #   perceptron
@@ -515,8 +511,6 @@ 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
@@ -593,7 +587,6 @@ numpy==2.2.6
     #   lm-eval
     #   matplotlib
     #   mistral-common
-    #   ml-dtypes
     #   mteb
     #   numba
     #   opencv-python-headless
@@ -617,7 +610,6 @@ numpy==2.2.6
     #   statsmodels
     #   tensorizer
     #   tifffile
-    #   tilelang
     #   torchvision
     #   transformers
     #   tritonclient
@@ -819,7 +811,6 @@ psutil==7.2.2
     #   accelerate
     #   peft
     #   tensorizer
-    #   tilelang
 py==1.11.0
     # via pytest-forked
 py-cpuinfo==9.0.0
@@ -1201,10 +1192,6 @@ 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
@@ -1221,8 +1208,6 @@ 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
@@ -1240,7 +1225,6 @@ tqdm==4.67.3
     #   pqdm
     #   segmentation-models-pytorch
     #   sentence-transformers
-    #   tilelang
     #   transformers
 transformers==5.5.3
     # via
@@ -1309,7 +1293,6 @@ typing-extensions==4.15.0
     #   sentence-transformers
     #   sqlalchemy
     #   starlette
-    #   tilelang
     #   torch
     #   typeguard
     #   typing-inspection
@@ -1376,8 +1359,6 @@ yarl==1.23.0
     # via
     #   aiohttp
     #   schemathesis
-z3-solver==4.15.4.0
-    # via tilelang
 zipp==3.23.0
     # via importlib-metadata
 

requirements/tpu.txt

@@ -12,4 +12,4 @@ ray[data]
 setuptools==78.1.0
 setuptools-rust>=1.9.0
 nixl==0.3.0
-tpu-inference==0.19.0
+tpu-inference==0.20.0

setup.py

@@ -1017,8 +1017,6 @@ 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():
@@ -1195,6 +1193,11 @@ setup(
             "opentelemetry-exporter-otlp>=1.26.0",
             "opentelemetry-semantic-conventions-ai>=0.4.1",
         ],
+        "triton-cpu": [
+            "triton @ "
+            "git+https://github.com/triton-lang/triton-cpu.git@270e696d ; "
+            "platform_machine == 'x86_64'",
+        ],  # Remove after stable release
     },
     cmdclass=cmdclass,
     package_data=package_data,

tests/entrypoints/openai/test_reasoning_enable_thinking.py

@@ -1,107 +0,0 @@
-# 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,14 +8,8 @@ import time
 from unittest.mock import patch
 
 import pytest
-import zmq
 
-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,
-)
+from vllm.v1.utils import APIServerProcessManager, wait_for_completion_or_failure
 
 # Global variables to control worker behavior
 WORKER_RUNTIME_SECONDS = 0.5
@@ -29,39 +23,6 @@ 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."""
@@ -307,92 +268,3 @@ 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/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 = 2.446
+        layer.routed_scaling_factor = None
         layer.shared_experts = None
         layer._expert_routing_tables = lambda: None
 
@@ -1678,10 +1678,7 @@ 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)
-            * layer.routed_scaling_factor
-        )
+        baseline_output = torch_moe(a, w1_original, w2_original, router_logits, topk)
 
     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,19 +120,9 @@ pytestmark = pytest.mark.skipif(
 )
 
 
-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,
+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
     )
 
 
@@ -140,23 +130,8 @@ def _call_fused(
 
 
 @pytest.mark.parametrize("num_tokens", [1, 4, 17, 64, 2048])
-@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):
+@pytest.mark.parametrize("n_heads", [8, 64])
+def test_q_path_matches_reference(num_tokens: int, n_heads: int):
     torch.manual_seed(0)
     device = "cuda"
     dtype = torch.bfloat16
@@ -181,16 +156,10 @@ def test_q_path_matches_reference(num_tokens: int, n_heads: int, padded_heads: i
         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_out = _call_fused(
-        q, padded_heads, kv, k_cache, slot_mapping, positions, cos_sin_cache, eps, bs
-    )
+    q_fused = q.clone()
+    _call_fused(q_fused, kv, k_cache, slot_mapping, positions, cos_sin_cache, eps, bs)
 
-    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(q_fused, q_ref, rtol=1e-2, atol=1e-2)
 
 
 # ── Test 2: KV path round-trip byte/value parity ─────────────────────────────
@@ -232,12 +201,11 @@ 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, padded to FlashMLA's min head count 64) ───────
+    # ── Fused path (dummy q, single head) ──────────────────────────────────
     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,
@@ -330,9 +298,8 @@ 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,
@@ -349,26 +316,9 @@ 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,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("n_heads", [8, 64])
 @pytest.mark.parametrize("block_size", [16, 64])
-def test_combined_q_and_kv(
-    num_tokens: int, n_heads: int, padded_heads: int, block_size: int
-):
+def test_combined_q_and_kv(num_tokens: int, n_heads: int, block_size: int):
     torch.manual_seed(2)
     device = "cuda"
     dtype = torch.bfloat16
@@ -395,10 +345,10 @@ def test_combined_q_and_kv(
     )
 
     # Fused single call.
+    q_fused = q.clone()
     k_cache_fused = torch.zeros_like(k_cache_ref)
-    q_out = _call_fused(
-        q,
-        padded_heads,
+    _call_fused(
+        q_fused,
         kv,
         k_cache_fused,
         slot_mapping,
@@ -408,10 +358,5 @@ def test_combined_q_and_kv(
         block_size,
     )
 
-    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(q_fused, q_ref, rtol=1e-2, atol=1e-2)
     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/attention.py:295-310:
+deepseek_v4/nvidia/ops/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/attention.py:295-310.
+    """Unfused path matching deepseek_v4/nvidia/ops/attention.py:295-310.
 
     Uses the production CUDA RoPE kernel + per_token_group_quant_fp8.
     """

tests/kernels/test_mhc_kernels.py

@@ -4,12 +4,7 @@ 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
@@ -97,128 +92,8 @@ def hc_head_ref(
 
 
 @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_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",
+    not current_platform.is_cuda(),
+    reason="CUDA required",
 )
 @pytest.mark.parametrize("num_tokens", [1, 4, 8, 128])
 @pytest.mark.parametrize("hidden_size", [4096, 7168])
@@ -321,42 +196,3 @@ def test_hc_head_triton(num_tokens, hidden_size, hc_mult):
 
     out_ref = hc_head_ref(residual, fn, hc_scale, hc_base, rms_eps, hc_eps)
     torch.testing.assert_close(out, out_ref, atol=5e-2, rtol=1e-2)
-
-
-@pytest.mark.skipif(
-    not (current_platform.is_cuda_alike() and has_tilelang()),
-    reason="CUDA or ROCm and tilelang required",
-)
-@pytest.mark.parametrize("num_tokens", [1, 4, 8, 128])
-@pytest.mark.parametrize("hidden_size", [4096, 7168])
-@pytest.mark.parametrize("hc_mult", [4])
-def test_hc_head_tilelang(num_tokens, hidden_size, hc_mult):
-    torch.set_default_device(DEVICE)
-    set_random_seed(0)
-
-    residual = torch.randn((num_tokens, hc_mult, hidden_size), dtype=torch.bfloat16)
-    fn = torch.randn((hc_mult, hc_mult * hidden_size), dtype=torch.float32) * 1e-4
-    hc_scale = torch.randn((1,), dtype=torch.float32) * 0.1
-    hc_base = torch.randn((hc_mult,), dtype=torch.float32) * 0.1
-    rms_eps = hc_eps = 1e-6
-
-    out = torch.empty((num_tokens, hidden_size), dtype=torch.bfloat16)
-    out.fill_(float("nan"))
-
-    result = torch.ops.vllm.hc_head_fused_kernel_tilelang(
-        residual,
-        fn,
-        hc_scale,
-        hc_base,
-        out,
-        hidden_size,
-        rms_eps,
-        hc_eps,
-        hc_mult,
-    )
-
-    assert result is None
-    assert not torch.isnan(out).any()
-
-    out_ref = hc_head_ref(residual, fn, hc_scale, hc_base, rms_eps, hc_eps)
-    torch.testing.assert_close(out, out_ref, atol=5e-2, rtol=1e-2)

tests/model_executor/model_loader/runai_streamer_loader/test_weight_utils.py

@@ -6,7 +6,6 @@ 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,
@@ -15,27 +14,6 @@ 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/quantization/test_register_quantization_config.py

@@ -108,14 +108,14 @@ def test_register_quantization_config(caplog_vllm):
     assert get_quantization_config("custom_quant") == CustomQuantConfig
 
     # The quantization method `custom_quant` is already exists,
-    # should raise a warning when re-registering it.
-    with caplog_vllm.at_level(logging.WARNING):
+    # should raise a debug message when re-registering it.
+    with caplog_vllm.at_level(logging.DEBUG, logger="vllm"):
         register_quantization_config("custom_quant")(CustomQuantConfig)
 
     assert any(
         "The quantization method 'custom_quant' already exists" in message
         for message in caplog_vllm.messages
-    ), "Expected a warning when re-registering custom_quant"
+    ), "Expected a debug message when re-registering custom_quant"
 
 
 @pytest.mark.parametrize(

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,
-        finished_req_ids=set(),
+        kv_connector_output=None,
         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, Any
+from typing import TYPE_CHECKING
 
 import torch
 
@@ -10,9 +10,8 @@ from vllm.platforms import current_platform
 from vllm.utils.import_utils import has_tilelang
 from vllm.utils.math_utils import cdiv
 
-# 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():
+# tilelang is only available on CUDA platforms
+if TYPE_CHECKING or current_platform.is_cuda():
     if not has_tilelang():
         raise ImportError(
             "tilelang is required for mhc but is not installed. Install it with "
@@ -24,8 +23,6 @@ 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:
@@ -40,17 +37,12 @@ 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=pass_configs,
+    pass_configs={
+        tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True,
+        tilelang.PassConfigKey.TL_DISABLE_TMA_LOWER: True,
+        tilelang.PassConfigKey.TL_PTXAS_REGISTER_USAGE_LEVEL: 10,
+    },
 )
 def mhc_pre_big_fuse_tilelang(
     gemm_out_mul,
@@ -86,8 +78,7 @@ 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:
-        if ENABLE_PDL:
-            T.pdl_sync()
+        T.pdl_sync()
         ##################################################################
         # _pre_norm_fn_fwd_norm
         rms = T.alloc_fragment(1, T.float32)
@@ -183,16 +174,18 @@ 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])
-
-        if ENABLE_PDL:
-            T.pdl_trigger()
+        T.pdl_trigger()
 
 
 # Copied from https://github.com/sgl-project/sglang/blob/main/python/sglang/srt/layers/mhc.py#L478
 
 
 @tilelang.jit(
-    pass_configs=pass_configs,
+    pass_configs={
+        tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True,
+        tilelang.PassConfigKey.TL_DISABLE_TMA_LOWER: True,
+        tilelang.PassConfigKey.TL_PTXAS_REGISTER_USAGE_LEVEL: 10,
+    },
 )
 def mhc_pre_big_fuse_with_norm_tilelang(
     gemm_out_mul,
@@ -237,8 +230,7 @@ def mhc_pre_big_fuse_with_norm_tilelang(
         T.clear(mixes)
         rms[0] = 0
 
-        if ENABLE_PDL:
-            T.pdl_sync()
+        T.pdl_sync()
 
         for i_split in T.serial(n_splits):
             rms[0] += gemm_out_sqrsum[i_split, i]
@@ -349,12 +341,15 @@ def mhc_pre_big_fuse_with_norm_tilelang(
 
                 T.copy(ol, layer_input[i, i0_h * hidden_block])
 
-        if ENABLE_PDL:
-            T.pdl_trigger()
+        T.pdl_trigger()
 
 
 @tilelang.jit(
-    pass_configs=pass_configs,
+    pass_configs={
+        tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True,
+        tilelang.PassConfigKey.TL_DISABLE_TMA_LOWER: True,
+        tilelang.PassConfigKey.TL_PTXAS_REGISTER_USAGE_LEVEL: 10,
+    },
 )
 def mhc_fused_tilelang(
     comb_mix,
@@ -395,8 +390,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 = tid // 32
-        lane = tid % 32
+        warp_id = T.get_warp_idx()
+        lane = T.get_lane_idx()
 
         s_warp = T.alloc_shared((num_warps, tile_n + 1), T.float32)
         s_post = T.alloc_shared((hc,), T.float32)
@@ -412,8 +407,7 @@ def mhc_fused_tilelang(
         T.clear(sqr)
         h_split_start = i_ks * h_per_split
 
-        if ENABLE_PDL:
-            T.pdl_sync()
+        T.pdl_sync()
 
         T.copy(post_mix[i_n, 0], s_post)
         T.copy(comb_mix[i_n, 0, 0], s_comb)
@@ -472,12 +466,15 @@ def mhc_fused_tilelang(
                     v2 += s_warp[w, tile_n]
                 rp_out[i_ks, i_n] = v2
 
-        if ENABLE_PDL:
-            T.pdl_trigger()
+        T.pdl_trigger()
 
 
 @tilelang.jit(
-    pass_configs=pass_configs,
+    pass_configs={
+        tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True,
+        tilelang.PassConfigKey.TL_DISABLE_TMA_LOWER: True,
+        tilelang.PassConfigKey.TL_PTXAS_REGISTER_USAGE_LEVEL: 10,
+    },
 )
 def mhc_post_tilelang(
     a,
@@ -510,8 +507,7 @@ def mhc_post_tilelang(
 
         a_local = T.alloc_fragment((hc, hc), T.float32)
         c_local = T.alloc_fragment(hc, T.float32)
-        if ENABLE_PDL:
-            T.pdl_sync()
+        T.pdl_sync()
         T.copy(a[i_n, 0, 0], a_local)
         T.copy(c[i_n, 0], c_local)
 
@@ -527,193 +523,15 @@ 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])
-        if ENABLE_PDL:
-            T.pdl_trigger()
+        T.pdl_trigger()
 
 
 @tilelang.jit(
-    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,
+    pass_configs={
+        tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True,
+        tilelang.PassConfigKey.TL_DISABLE_TMA_LOWER: True,
+        tilelang.PassConfigKey.TL_PTXAS_REGISTER_USAGE_LEVEL: 10,
+    },
 )
 def hc_head_fuse_tilelang(
     residual,
@@ -748,8 +566,7 @@ 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:
-        if ENABLE_PDL:
-            T.pdl_sync()
+        T.pdl_sync()
 
         # ------------------------------------------------------------------
         # Pass 1 – for each residual channel m_c and h_block:
@@ -807,5 +624,4 @@ def hc_head_fuse_tilelang(
 
             T.copy(ol, out[i, i0_h * h_block], disable_tma=True)
 
-        if ENABLE_PDL:
-            T.pdl_trigger()
+        T.pdl_trigger()

vllm/entrypoints/cli/serve.py

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

vllm/entrypoints/openai/chat_completion/protocol.py

@@ -472,27 +472,17 @@ 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,
-                extra_kwargs,
+                dict(
+                    add_generation_prompt=self.add_generation_prompt,
+                    continue_final_message=self.continue_final_message,
+                    documents=self.documents,
+                    reasoning_effort=self.reasoning_effort,
+                ),
             ),
             media_io_kwargs=self.media_io_kwargs,
         )

vllm/entrypoints/openai/responses/protocol.py

@@ -298,28 +298,17 @@ 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(
+            chat_template_kwargs=merge_kwargs(  # To remove unset values
                 self.chat_template_kwargs,
-                extra_kwargs,
+                dict(
+                    add_generation_prompt=not continue_final,
+                    continue_final_message=continue_final,
+                    reasoning_effort=None if reasoning is None else reasoning.effort,
+                ),
             ),
             media_io_kwargs=self.media_io_kwargs,
         )

vllm/lora/ops/triton_ops/fused_moe_lora_op.py

@@ -7,7 +7,6 @@ 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
@@ -407,7 +406,7 @@ def _run_fused_moe_lora_one_shot(
 
     # NPID_FACTOR heuristic: scale N-axis parallelism when base CTA count is
     # short of saturating the SM array. Cap by the cost of redundant shrink.
-    sm_count = current_platform.num_compute_units(device.index)
+    sm_count = torch.cuda.get_device_properties(device).multi_processor_count
     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)
@@ -787,7 +786,7 @@ def _run_fused_moe_lora_small_batch(
     N_tiles = triton.cdiv(N_per_slice, BLOCK_N)
     pair_slices = M_grid * num_slices
 
-    sm_count = current_platform.num_compute_units(device.index)
+    sm_count = torch.cuda.get_device_properties(device).multi_processor_count
     n_tiles_per_program = _pick_small_batch_chunk(pair_slices, N_tiles, sm_count)
     n_chunks = triton.cdiv(N_tiles, n_tiles_per_program)
     work_total = pair_slices * n_chunks

vllm/model_executor/kernels/mhc/tilelang.py

@@ -5,88 +5,6 @@ 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,
@@ -162,16 +80,10 @@ def mhc_pre_tilelang(
     residual_flat = residual.view(-1, hc_mult, hidden_size)
     num_tokens = residual_flat.shape[0]
 
-    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
+    # 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))
 
     post_mix = torch.empty(
         num_tokens, hc_mult, dtype=torch.float32, device=residual.device
@@ -190,24 +102,13 @@ def mhc_pre_tilelang(
         n_splits, num_tokens, dtype=torch.float32, device=residual.device
     )
 
-    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,
-        )
+    tf32_hc_prenorm_gemm(
+        residual_flat.view(num_tokens, hc_mult * hidden_size),
+        fn,
+        gemm_out_mul,
+        gemm_out_sqrsum,
+        n_splits,
+    )
 
     if norm_weight is None:
         mhc_pre_big_fuse_tilelang(
@@ -403,24 +304,16 @@ 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)
 
-    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:
+    fma_token_threshold = 16
+    if num_tokens <= fma_token_threshold:
         # 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:
-        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
+        # 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))
 
     gemm_out_mul = torch.empty(
         n_splits,
@@ -455,7 +348,7 @@ def mhc_fused_post_pre_tilelang(
         device=residual.device,
     )
 
-    if use_small_fma:
+    if num_tokens <= fma_token_threshold:
         mhc_fused_tilelang(
             comb_res_mix_flat,
             residual_flat,
@@ -482,26 +375,15 @@ def mhc_fused_post_pre_tilelang(
             residual.shape[-1],
         )
 
-        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
+        from vllm.utils.deep_gemm import tf32_hc_prenorm_gemm
 
-            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,
-            )
+        tf32_hc_prenorm_gemm(
+            residual_cur.view(num_tokens, hc_mult * hidden_size),
+            fn,
+            gemm_out_mul,
+            gemm_out_sqrsum,
+            n_splits,
+        )
 
     if norm_weight is None:
         mhc_pre_big_fuse_tilelang(

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

@@ -99,6 +99,9 @@ class TrtLlmBf16Experts(mk.FusedMoEExpertsMonolithic):
     ) -> bool:
         return True
 
+    def supports_chunking(self) -> bool:
+        return False
+
     def supports_expert_map(self) -> bool:
         return False
 
@@ -137,6 +140,5 @@ 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,6 +88,9 @@ 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,6 +113,9 @@ 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,27 +157,8 @@ class TrtLlmNvFp4ExpertsBase:
     def activation_format() -> mk.FusedMoEActivationFormat:
         return mk.FusedMoEActivationFormat.Standard
 
-    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_chunking(self) -> bool:
+        return False
 
     def supports_expert_map(self) -> bool:
         return False
@@ -218,7 +199,7 @@ class TrtLlmNvFp4ExpertsModular(TrtLlmNvFp4ExpertsBase, mk.FusedMoEExpertsModula
     def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce:
         return TopKWeightAndReduceNoOP()
 
-    def _invoke_kernel(
+    def apply(
         self,
         output: torch.Tensor,
         hidden_states: torch.Tensor,
@@ -228,10 +209,18 @@ class TrtLlmNvFp4ExpertsModular(TrtLlmNvFp4ExpertsBase, mk.FusedMoEExpertsModula
         topk_ids: torch.Tensor,
         activation: MoEActivation,
         global_num_experts: int,
-        a1q_scale: torch.Tensor,
+        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,
     ):
         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
 
@@ -273,57 +262,6 @@ 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/mhc.py

@@ -3,12 +3,8 @@
 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]
@@ -89,52 +85,6 @@ 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,
@@ -147,6 +97,9 @@ 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")
@@ -194,20 +147,6 @@ 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,
@@ -215,6 +154,9 @@ 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")
@@ -278,32 +220,17 @@ class HCHeadOp(CustomOp):
         out = torch.empty(
             num_tokens, hidden_size, dtype=torch.bfloat16, device=hidden_states.device
         )
-
-        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,
-            )
-
+        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):
@@ -363,42 +290,9 @@ class MHCFusedPostPreOp(CustomOp):
             norm_eps,
         )
 
-    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_hip(self, *args, **kwargs):
+        raise NotImplementedError(
+            "Hip implementation of mhc_fused_post_pre is not available"
         )
 
     def forward_native(self, *args, **kwargs):

vllm/model_executor/layers/quantization/__init__.py

@@ -84,7 +84,7 @@ def register_quantization_config(quantization: str):
 
     def _wrapper(quant_config_cls):
         if quantization in QUANTIZATION_METHODS:
-            logger.warning(
+            logger.debug(
                 "The quantization method '%s' already exists and will be "
                 "overwritten by the quantization config %s.",
                 quantization,

vllm/model_executor/layers/quantization/humming.py

@@ -43,9 +43,12 @@ from vllm.model_executor.parameter import (
     RowvLLMParameter,
 )
 from vllm.model_executor.utils import set_weight_attrs
-from vllm.platforms import current_platform
 
-if current_platform.is_cuda():
+if TYPE_CHECKING:
+    from vllm.model_executor.models.utils import WeightsMapper
+
+
+try:
     from humming.dtypes import DataType
     from humming.layer import HummingMethod
     from humming.schema import (
@@ -62,17 +65,16 @@ if current_platform.is_cuda():
         HummingIndexedExperts,
         get_humming_moe_gemm_type,
     )
+except ModuleNotFoundError:
+    HummingMethod = None
 
-if TYPE_CHECKING:
-    from humming.schema import (
-        BaseInputSchema,
-        BaseWeightSchema,
-        HummingInputSchema,
-        HummingWeightSchema,
+
+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."
     )
 
-    from vllm.model_executor.models.utils import WeightsMapper
-
 
 def prepare_padded_shape(shape, x):
     padded_shape = math.ceil(shape / x) * x
@@ -184,6 +186,7 @@ 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,8 +1090,7 @@ def runai_safetensors_weights_iterator(
             mininterval=2,
         )
 
-        for name, tensor in tensor_iter:
-            yield name, tensor.clone()
+        yield from tensor_iter
 
 
 def _init_fastsafetensors_loader(

vllm/model_executor/models/deepseek_eagle3.py

@@ -407,6 +407,12 @@ class Eagle3DeepseekV2ForCausalLM(DeepseekV2ForCausalLM):
         hidden_states: torch.Tensor,
     ) -> torch.Tensor:
         # Combine multiple auxiliary hidden states returned by Eagle3
+        if self.model.fc_norm is not None:
+            chunks = hidden_states.chunk(self.model.num_aux_hidden_states, dim=-1)
+            hidden_states = torch.cat(
+                [norm(chunk) for norm, chunk in zip(self.model.fc_norm, chunks)],
+                dim=-1,
+            )
         return self.model.fc(hidden_states)
 
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):

vllm/model_executor/models/registry.py

@@ -980,7 +980,7 @@ class _ModelRegistry:
             raise TypeError(msg)
 
         if model_arch in self.models:
-            logger.warning(
+            logger.debug(
                 "Model architecture %s is already registered, and will be "
                 "overwritten by the new model class %s.",
                 model_arch,

vllm/models/deepseek_v4/amd/model.py

@@ -11,12 +11,17 @@ 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,
@@ -47,14 +52,16 @@ from vllm.model_executor.models.utils import (
     make_layers,
     maybe_prefix,
 )
-from vllm.models.deepseek_v4.attention import (
+from vllm.model_executor.utils import set_weight_attrs
+from vllm.models.deepseek_v4.nvidia.ops.attention import (
     DeepseekV4Indexer,
     DeepseekV4MLAModules,
     DeepseekV4MultiHeadLatentAttentionWrapper,
 )
 from vllm.platforms import current_platform
 from vllm.sequence import IntermediateTensors
-from vllm.utils.import_utils import has_tilelang
+from vllm.triton_utils import tl, triton
+from vllm.utils.torch_utils import direct_register_custom_op
 
 
 class DeepseekV4MLP(nn.Module):
@@ -108,6 +115,501 @@ 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,
@@ -120,6 +622,15 @@ 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
@@ -130,6 +641,16 @@ 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,
@@ -142,7 +663,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.int32
+        self.hash_indices_dtype = torch.int64 if self.use_mega_moe else 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
@@ -173,10 +694,47 @@ class DeepseekV4MoE(nn.Module):
                 hidden_act=config.hidden_act,
                 swiglu_limit=self.swiglu_limit,
                 quant_config=quant_config,
-                reduce_results=False,
+                reduce_results=self.use_mega_moe,
                 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
 
@@ -208,6 +766,44 @@ 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
@@ -226,6 +822,10 @@ 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__(
@@ -474,7 +1074,6 @@ 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,
@@ -505,7 +1104,7 @@ class DeepseekV4DecoderLayer(nn.Module):
     ):
         return self.mhc_post(x, residual, post, comb)
 
-    def _forward_fused_post_pre(
+    def _forward_cuda(
         self,
         x: torch.Tensor,
         positions: torch.Tensor,
@@ -557,7 +1156,7 @@ class DeepseekV4DecoderLayer(nn.Module):
         x = self.ffn(x, input_ids)
         return x, residual, post_mix, res_mix
 
-    def _forward_unfused_post_pre(
+    def _forward_rocm(
         self,
         x: torch.Tensor,
         positions: torch.Tensor,
@@ -596,13 +1195,12 @@ class DeepseekV4DecoderLayer(nn.Module):
     ) -> tuple[
         torch.Tensor, torch.Tensor | None, torch.Tensor | None, torch.Tensor | None
     ]:
-        if not self.has_tilelang:
-            return self._forward_unfused_post_pre(
+        if current_platform.is_rocm():
+            return self._forward_rocm(
                 x, positions, input_ids, post_mix, res_mix, residual
             )
-        return self._forward_fused_post_pre(
-            x, positions, input_ids, post_mix, res_mix, residual
-        )
+
+        return self._forward_cuda(x, positions, input_ids, post_mix, res_mix, residual)
 
 
 @support_torch_compile
@@ -613,6 +1211,15 @@ 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
@@ -685,7 +1292,6 @@ 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.
@@ -742,6 +1348,9 @@ 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(
@@ -752,7 +1361,7 @@ class DeepseekV4Model(nn.Module):
                 res_mix,
                 residual,
             )
-        if layer is not None and self.has_tilelang:
+        if layer is not None and current_platform.is_cuda():
             hidden_states = layer.hc_post(hidden_states, residual, post_mix, res_mix)
 
         if not get_pp_group().is_last_rank:
@@ -873,6 +1482,9 @@ 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(
@@ -883,6 +1495,10 @@ 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":
@@ -985,6 +1601,7 @@ 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,9 +39,11 @@ 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
+from .model import (
+    DeepseekV4DecoderLayer,
+    make_deepseek_v4_expert_params_mapping,
+)
 
 logger = init_logger(__name__)
 
@@ -119,7 +121,6 @@ class DeepSeekV4MultiTokenPredictorLayer(nn.Module):
         )
 
         self.hc_head_op = HCHeadOp()
-        self.has_tilelang = has_tilelang()
 
     def forward(
         self,
@@ -146,7 +147,7 @@ class DeepSeekV4MultiTokenPredictorLayer(nn.Module):
         hidden_states, residual, post_mix, res_mix = self.mtp_block(
             positions=positions, x=hidden_states, input_ids=None
         )
-        if self.has_tilelang:
+        if current_platform.is_cuda():
             hidden_states = self.mtp_block.hc_post(
                 hidden_states, residual, post_mix, res_mix
             )
@@ -329,13 +330,19 @@ class DeepSeekV4MTP(nn.Module):
         head_rank_end = n_local_head * (tp_rank + 1)
 
         # Pre-compute expert mapping ONCE.
-        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,
-        )
+        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,
+            )
 
         # FP8 experts register ``..._weight_scale_inv`` (block_quant) while
         # FP4/MXFP4 experts register ``..._weight_scale``. Choose the suffix
@@ -458,9 +465,14 @@ 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.attention import (
+    from vllm.models.deepseek_v4.nvidia.ops.attention import (
         DeepseekV4MLAAttention,
     )
 
@@ -592,10 +592,6 @@ 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/common/ops/__init__.py

@@ -10,7 +10,6 @@ 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",
@@ -21,5 +20,4 @@ __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

@@ -11,6 +11,12 @@ Three specialized kernels:
   - _fused_kv_compress_norm_rope_insert_indexer_mxfp4_attn:
         head=128, MXFP4 (block=32), 4 ue8m0 bytes
 
+Additional cutedsl kernels:
+  - _compress_kv_sparse_attn_cutedsl / _norm_rope_insert_sparse_attn_cutedsl:
+        CuTe DSL split kernels for C128
+  - _fused_kv_compress_norm_rope_insert_sparse_attn_cutedsl:
+        CuTe DSL fused kernels for C4
+
 RoPE is register-based via tl.reshape -> tl.split -> tl.interleave (or the
 even/odd halves are consumed directly for MXFP4, no interleave needed).
 FP8 UE8M0 quant uses tl.reshape to tile [N_QUANT_BLOCKS, QUANT_BLOCK] for
@@ -19,92 +25,42 @@ 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 functools import cache
 
 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,
+@cache
+def _get_sparse_attn_cutedsl_impls():
+    from .sparse_attn_compress_cutedsl import (
+        _compress_kv_sparse_attn_cutedsl,
+        _fused_kv_compress_norm_rope_insert_sparse_attn_cutedsl,
+        _norm_rope_insert_sparse_attn_cutedsl,
     )
 
+    return (
+        _compress_kv_sparse_attn_cutedsl,
+        _norm_rope_insert_sparse_attn_cutedsl,
+        _fused_kv_compress_norm_rope_insert_sparse_attn_cutedsl,
+    )
+
+
+def _compress_kv_sparse_attn_cutedsl(*args, **kwargs):
+    """CuTe DSL sparse-attention compress wrapper."""
+    return _get_sparse_attn_cutedsl_impls()[0](*args, **kwargs)
+
+
+def _norm_rope_insert_sparse_attn_cutedsl(*args, **kwargs):
+    """CuTe DSL RMSNorm/RoPE/FP8-store wrapper."""
+    return _get_sparse_attn_cutedsl_impls()[1](*args, **kwargs)
+
+
+def _fused_kv_compress_norm_rope_insert_sparse_attn_cutedsl(*args, **kwargs):
+    """CuTe DSL fused C4 sparse-attention compressor wrapper."""
+    return _get_sparse_attn_cutedsl_impls()[2](*args, **kwargs)
+
 
 # =============================================================================
 # DeepseekV4 Attention path (head=512, nope=448 FP8 + rope=64 bf16)

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

@@ -1,101 +0,0 @@
-# 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/common/ops/sparse_attn_compress_cutedsl.py

@@ -8,7 +8,6 @@ The public wrappers provide the C4 fused and C128 split kernels.
 from __future__ import annotations
 
 from functools import cache
-from typing import Any
 
 import cutlass
 import cutlass.cute as cute
@@ -1087,7 +1086,7 @@ class SparseAttnNormRopeStoreKernel:
         )
 
 
-def compress_kv_sparse_attn_cutedsl(
+def _compress_kv_sparse_attn_cutedsl(
     state_cache: torch.Tensor,
     token_to_req_indices: torch.Tensor,
     positions: torch.Tensor,
@@ -1119,7 +1118,7 @@ def compress_kv_sparse_attn_cutedsl(
     )
 
 
-def norm_rope_insert_sparse_attn_cutedsl(
+def _norm_rope_insert_sparse_attn_cutedsl(
     compressed_kv: torch.Tensor,
     positions: torch.Tensor,
     slot_mapping: torch.Tensor,
@@ -1175,7 +1174,7 @@ def norm_rope_insert_sparse_attn_cutedsl(
     )
 
 
-def fused_kv_compress_norm_rope_insert_sparse_attn_cutedsl(
+def _fused_kv_compress_norm_rope_insert_sparse_attn_cutedsl(
     state_cache: torch.Tensor,
     token_to_req_indices: torch.Tensor,
     positions: torch.Tensor,
@@ -1239,94 +1238,3 @@ def fused_kv_compress_norm_rope_insert_sparse_attn_cutedsl(
         kv_slot_mapping,
         kv_cache_block_size,
     )
-
-
-def compress_norm_rope_store_cutedsl(
-    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:
-    if compress_ratio == 4:
-        # For C4A, the single fused kernel is faster than the two-kernel version.
-        fused_kv_compress_norm_rope_insert_sparse_attn_cutedsl(
-            state_cache,
-            token_to_req_indices,
-            positions,
-            slot_mapping,
-            block_table,
-            block_size,
-            rms_norm_weight,
-            rms_norm_eps,
-            cos_sin_cache,
-            kv_cache,
-            k_cache_metadata.slot_mapping,
-            kv_cache.shape[1],  # paged KV cache block size
-            kv_cache.stride(0),
-            head_size=head_dim,
-            state_width=state_width,
-            rope_head_dim=rope_head_dim,
-            fp8_max=448.0,
-            quant_block=quant_block,
-            token_stride=token_stride,
-            scale_dim=scale_dim,
-            compress_ratio=compress_ratio,
-            overlap=overlap,
-        )
-    else:
-        # For C128, the two-kernel version is faster than the single fused kernel.
-        compressed_kv = torch.empty(
-            (num_actual, head_dim),
-            dtype=torch.float32,
-            device=state_cache.device,
-        )
-        compress_kv_sparse_attn_cutedsl(
-            state_cache,
-            token_to_req_indices,
-            positions,
-            slot_mapping,
-            block_table,
-            block_size,
-            compressed_kv,
-            head_size=head_dim,
-            state_width=state_width,
-            compress_ratio=compress_ratio,
-            overlap=overlap,
-        )
-        norm_rope_insert_sparse_attn_cutedsl(
-            compressed_kv,
-            positions,
-            slot_mapping,
-            rms_norm_weight,
-            rms_norm_eps,
-            cos_sin_cache,
-            kv_cache,
-            k_cache_metadata.slot_mapping,
-            kv_cache.shape[1],  # paged KV cache block size
-            kv_cache.stride(0),
-            head_size=head_dim,
-            rope_head_dim=rope_head_dim,
-            fp8_max=448.0,
-            quant_block=quant_block,
-            token_stride=token_stride,
-            scale_dim=scale_dim,
-            compress_ratio=compress_ratio,
-        )

vllm/models/deepseek_v4/compressor.py

@@ -13,13 +13,15 @@ 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 (
-    compress_norm_rope_store_triton,
+    _compress_kv_sparse_attn_cutedsl,
+    _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_cutedsl,
+    _norm_rope_insert_sparse_attn_cutedsl,
 )
 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,15 +173,32 @@ class CompressorStateCache(torch.nn.Module, AttentionLayerBase):
 
 
 class DeepseekCompressor(nn.Module):
-    """DeepSeek V4 KV/score compressor.
+    _compressed_kv_buffers: ClassVar[dict[tuple[str, int, int], torch.Tensor]] = {}
 
-    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.
-    """
+    @classmethod
+    def _get_compressed_kv_buffer(
+        cls,
+        device: str,
+        max_num_tokens: int,
+        head_dim: int,
+    ) -> torch.Tensor:
+        if device == "cuda" and torch.accelerator.is_available():
+            device_key = f"cuda:{torch.accelerator.current_device_index()}"
+            alloc_device = torch.device(device_key)
+        else:
+            device_key = str(device)
+            alloc_device = torch.device(device)
+
+        key = (device_key, max_num_tokens, head_dim)
+        buffer = cls._compressed_kv_buffers.get(key)
+        if buffer is None:
+            buffer = torch.empty(
+                (max_num_tokens, head_dim),
+                dtype=torch.float32,
+                device=alloc_device,
+            )
+            cls._compressed_kv_buffers[key] = buffer
+        return buffer
 
     def __init__(
         self,
@@ -250,18 +269,37 @@ class DeepseekCompressor(nn.Module):
             assert not use_fp4_cache, (
                 "MXFP4 cache is only supported for indexer (head=128)"
             )
+            self._use_cutedsl_sparse_compressor = True
+            self._use_cutedsl_fused_sparse_compressor = self.compress_ratio == 4
+            self._compress_kernel = _compress_kv_sparse_attn_cutedsl
+            self._norm_rope_store_kernel = _norm_rope_insert_sparse_attn_cutedsl
+            self._fused_sparse_kernel = (
+                _fused_kv_compress_norm_rope_insert_sparse_attn_cutedsl
+            )
+            self._compressed_kv_buffer = self._get_compressed_kv_buffer(
+                self.device,
+                vllm_config.scheduler_config.max_num_batched_tokens,
+                self.head_dim,
+            )
             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:
+            self._use_cutedsl_sparse_compressor = False
             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}"
@@ -306,22 +344,29 @@ class DeepseekCompressor(nn.Module):
         )
 
         # Store the KV and score (with fused APE addition) in the state.
-        # 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,
+        # 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,
         )
 
         # Fused: compress → RMSNorm → RoPE → FP8 quant → KV cache write.
@@ -335,46 +380,157 @@ 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
 
-        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,
+        if self._use_cutedsl_sparse_compressor:
+            if self._use_cutedsl_fused_sparse_compressor:
+                self._fused_sparse_kernel(
+                    state_cache,
+                    token_to_req_indices,
+                    positions,
+                    slot_mapping,
+                    block_table,
+                    block_size,
+                    self.norm.weight,
+                    self.rms_norm_eps,
+                    cos_sin_cache,
+                    kv_cache,
+                    k_cache_metadata.slot_mapping,
+                    kv_cache.shape[1],  # paged KV cache block size
+                    kv_cache.stride(0),
+                    head_size=self.head_dim,
+                    state_width=state_width,
+                    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,
+                    compress_ratio=self.compress_ratio,
+                    overlap=self.overlap,
                 )
-
-                # 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
+                compressed_kv = self._compressed_kv_buffer[:num_actual]
+                self._compress_kernel(
+                    state_cache,
+                    token_to_req_indices,
+                    positions,
+                    slot_mapping,
+                    block_table,
+                    block_size,
+                    compressed_kv,
+                    head_size=self.head_dim,
+                    state_width=state_width,
+                    compress_ratio=self.compress_ratio,
+                    overlap=self.overlap,
+                )
+                self._norm_rope_store_kernel(
+                    compressed_kv,
+                    positions,
+                    slot_mapping,
+                    self.norm.weight,
+                    self.rms_norm_eps,
+                    cos_sin_cache,
+                    kv_cache,
+                    k_cache_metadata.slot_mapping,
+                    kv_cache.shape[1],  # paged KV cache block size
+                    kv_cache.stride(0),
+                    head_size=self.head_dim,
+                    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,
+                    compress_ratio=self.compress_ratio,
+                )
         else:
-            # AMD GPUs.
-            # Always use a triton kernel.
-            compress_norm_rope_store_fn = compress_norm_rope_store_triton
+            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,
+            )
 
-        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.attention import (
+    from vllm.models.deepseek_v4.nvidia.ops.attention import (
         DeepseekV4MLAAttention,
     )
     from vllm.v1.attention.backends.mla.sparse_swa import DeepseekSparseSWAMetadata
@@ -63,18 +63,6 @@ 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
@@ -116,16 +104,6 @@ class DeepseekV4FlashMLASparseImpl(DeepseekV4SparseMLAAttentionImpl):
 
     backend_cls = DeepseekV4FlashMLASparseBackend
 
-    @classmethod
-    def get_padded_num_q_heads(cls, num_heads: int) -> int:
-        # FP8 decode kernel only supports h_q = 64 or 128.
-        if num_heads > 128:
-            raise ValueError(
-                f"DeepseekV4 FlashMLA does not support {num_heads} heads "
-                "(FP8 decode kernel requires h_q in {64, 128})."
-            )
-        return 64 if num_heads <= 64 else 128
-
     @classmethod
     def forward_mqa(  # type: ignore[override]
         cls,

vllm/models/deepseek_v4/nvidia/model.py

@@ -54,7 +54,7 @@ from vllm.model_executor.models.utils import (
     maybe_prefix,
 )
 from vllm.model_executor.utils import set_weight_attrs
-from vllm.models.deepseek_v4.attention import (
+from vllm.models.deepseek_v4.nvidia.ops.attention import (
     DeepseekV4Indexer,
     DeepseekV4MLAModules,
     DeepseekV4MultiHeadLatentAttentionWrapper,

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

@@ -5,9 +5,4 @@
 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/attention.py

@@ -156,6 +156,18 @@ 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
@@ -251,9 +263,6 @@ 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
@@ -441,7 +450,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)
-                q = self._fused_qnorm_rope_kv_insert(q, kv, positions, attn_metadata)
+                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
@@ -475,7 +484,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)
-                q = self._fused_qnorm_rope_kv_insert(q, kv, positions, attn_metadata)
+                self._fused_qnorm_rope_kv_insert(q, kv, positions, attn_metadata)
                 return q
 
             q, _ = maybe_execute_in_parallel(
@@ -488,7 +497,12 @@ 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)
-            q = self._fused_qnorm_rope_kv_insert(q, kv, positions, attn_metadata)
+            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)
 
         # MLA attention writes into the pre-allocated `out` buffer
         # ([num_tokens, padded_heads, head_dim]).
@@ -502,17 +516,9 @@ class DeepseekV4MultiHeadLatentAttentionWrapper(PluggableLayer):
         attn_metadata: (
             dict[str, AttentionMetadata] | list[dict[str, AttentionMetadata]] | None
         ),
-    ) -> torch.Tensor:
+    ) -> None:
         if not isinstance(attn_metadata, dict):
-            # 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
+            return
 
         swa_metadata = cast(
             "DeepseekSparseSWAMetadata | None",
@@ -524,19 +530,16 @@ 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,
-        #            with zero-fill for the padding head slots.  The kernel
-        #            allocates and returns the padded q tensor.
+        #   Q side:  q_head_norm (per-head RMSNorm, no weight) + GPT-J RoPE
         #   KV side: GPT-J RoPE + UE8M0 FP8 quant + paged cache insert
         # kv is unchanged; mla_attn reads kv solely via swa_kv_cache.
-        return torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
+        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,
         )
@@ -604,6 +607,9 @@ 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,
@@ -649,8 +655,19 @@ class DeepseekV4MLAAttention(nn.Module, AttentionLayerBase):
         self.aux_stream = aux_stream
         self.ln_events = [torch.cuda.Event(), torch.cuda.Event()]
 
-        # Padded Q head count is dictated by the selected impl.
-        self.padded_heads = self.impl_cls.get_padded_num_q_heads(num_heads)
+        # 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
 
         # Store attention sink
         assert attn_sink is not None

vllm/platforms/cpu.py

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

vllm/platforms/cuda.py

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

vllm/platforms/interface.py

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

vllm/triton_utils/importing.py

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

vllm/utils/import_utils.py

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

vllm/utils/platform_utils.py

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

vllm/v1/engine/async_llm.py

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

vllm/v1/engine/coordinator.py

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

vllm/v1/engine/core_client.py

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

vllm/v1/engine/utils.py

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

vllm/v1/utils.py

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

vllm/v1/worker/cpu/buffer_utils.py

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

vllm/v1/worker/cpu/model_runner.py

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

vllm/v1/worker/cpu/shm.py

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

vllm/v1/worker/cpu_worker.py

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

vllm/v1/worker/gpu/kv_connector.py

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

vllm/v1/worker/gpu/model_runner.py

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