[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/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"

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);
         }
       }

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

tests/kernels/mamba/test_gdn_prefill_cutedsl.py

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

tests/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(

vllm/cute_utils/__init__.py

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

vllm/cute_utils/_tcgen05.py

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

vllm/cute_utils/cvt.py

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

vllm/engine/arg_utils.py

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

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

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

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

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

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

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

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

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

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

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

vllm/model_executor/layers/quantization/__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/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/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,11 +25,43 @@ even/odd halves, producing (N_QUANT_BLOCKS, MXFP4_BLOCK/2) packed nibbles
 and N_QUANT_BLOCKS ue8m0 bytes.
 """
 
+from functools import cache
+
 from vllm.triton_utils import tl, triton
 
 from .fused_indexer_q import _fp32x2_to_fp4x2
 
 
+@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/compressor.py

@@ -13,9 +13,11 @@ 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_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,
+    _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.platforms import current_platform
@@ -171,6 +173,33 @@ class CompressorStateCache(torch.nn.Module, AttentionLayerBase):
 
 
 class DeepseekCompressor(nn.Module):
+    _compressed_kv_buffers: ClassVar[dict[tuple[str, int, int], torch.Tensor]] = {}
+
+    @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,
         vllm_config: VllmConfig,
@@ -240,12 +269,24 @@ class DeepseekCompressor(nn.Module):
             assert not use_fp4_cache, (
                 "MXFP4 cache is only supported for indexer (head=128)"
             )
-            self._fused_kernel = _fused_kv_compress_norm_rope_insert_sparse_attn
+            self._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
@@ -339,43 +380,104 @@ class DeepseekCompressor(nn.Module):
         k_cache_metadata = cast(Any, attn_metadata[self.k_cache_prefix])
         kv_cache = self._static_forward_context[self.k_cache_prefix].kv_cache
 
-        self._fused_kernel[(num_actual,)](
-            # state cache
-            state_cache,
-            state_cache.stride(0),
-            state_cache.stride(1),
-            # metadata
-            token_to_req_indices,
-            positions,
-            slot_mapping,
-            block_table,
-            block_table.stride(0),
-            block_size,
-            # RMSNorm
-            self.norm.weight,
-            self.rms_norm_eps,
-            # RoPE
-            cos_sin_cache,
-            cos_sin_cache.stride(0),
-            # KV cache
-            kv_cache,
-            k_cache_metadata.slot_mapping,
-            kv_cache.shape[1],  # paged KV cache block size (tokens per block)
-            # constexprs
-            HEAD_SIZE=self.head_dim,
-            TRITON_BLOCK_SIZE=triton.next_power_of_2(self.head_dim),
-            STATE_WIDTH=state_width,
-            COMPRESS_RATIO=self.compress_ratio,
-            OVERLAP=self.overlap,
-            ROPE_HEAD_DIM=self.rope_head_dim,
-            FP8_MAX=448.0,
-            QUANT_BLOCK=self._quant_block,
-            TOKEN_STRIDE=self._token_stride,
-            SCALE_DIM=self._scale_dim,
-            KV_BLOCK_STRIDE=kv_cache.stride(0),
-            num_warps=self._num_warps,
-            **pdl_kwargs,
-        )
+        if 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,
+                )
+            else:
+                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:
+            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,
+            )
 
 
 @triton.jit

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

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

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

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

vllm/v1/attention/backends/gdn_attn.py

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