use cute-dsl to speedup mxfp QDQ

benchmark/vllm-qdq-plugin/README.md

@@ -15,9 +15,11 @@ The plugin registers as a `vllm.general_plugins` entry point, which vLLM loads a
 ```bash
 pip install git+https://github.com/yiliu30/vllm-qdq-plugin.git
 
-# Or for development:
-git clone https://github.com/yiliu30/vllm-qdq-plugin.git
-pip install -e vllm-qdq-plugin/
+# Or install from the neural-compressor repository root for development:
+pip install -e benchmark/vllm-qdq-plugin/
+
+# When already inside benchmark/vllm-qdq-plugin:
+pip install -e .
 ```
 
 ## Usage
@@ -32,6 +34,9 @@ VLLM_QDQ=1 vllm serve /path/to/model --tensor-parallel-size 2
 # Enable trace logging (prints shape/dtype for each QDQ call)
 VLLM_QDQ=1 VLLM_QDQ_TRACE=1 vllm serve /path/to/model
 
+# Request the CuTe QDQ backend on SM80+ GPUs
+VLLM_QDQ=1 VLLM_QDQ_CUTE=1 vllm serve /path/to/model
+
 # Force MXFP4 QDQ on Marlin MoE when dtype-based detection is not enough
 VLLM_QDQ=1 VLLM_MARLIN_MOE_QDQ_MODE=FORCE_MXFP4 vllm serve /path/to/model
 ```
@@ -42,6 +47,7 @@ VLLM_QDQ=1 VLLM_MARLIN_MOE_QDQ_MODE=FORCE_MXFP4 vllm serve /path/to/model
 |---|---|---|
 | `VLLM_QDQ` | `0` | Set to `1` to enable QDQ |
 | `VLLM_QDQ_TRACE` | `0` | Set to `1` to print trace lines (up to 200) |
+| `VLLM_QDQ_CUTE` | `0` | Enable fused CuTe MXFP4/MXFP8 QDQ kernels. Requires CUDA, SM80+, NVIDIA CUTLASS DSL, contiguous input, group size 32, and `K` divisible by 32. Unsupported inputs fall back to the reference implementation. |
 | `VLLM_MARLIN_MOE_QDQ_MODE` | `0` | Set to `FORCE_MXFP4` to apply MXFP4 QDQ in `moe_wna16_marlin_gemm` when dtype-based routing is not sufficient. Matching is case-insensitive. |
 
 ## Support Status
@@ -59,6 +65,30 @@ For MXFP4, the QDQ simulates:
 
 This introduces the same quantization noise that a "real" MXFP4 GEMM would produce on the input side, while keeping the actual computation in bf16 via Marlin's weight-only dequant kernel.
 
+### CuTe QDQ Validation
+
+The optional CuTe backend checks CUDA capability at runtime and only accepts SM80 or newer GPUs. Verify the installed CuTe DSL first, then compare the CuTe and reference paths:
+
+```bash
+CUDA_VISIBLE_DEVICES=<idle-gpu> python scripts/verify_cute_dsl.py
+CUDA_VISIBLE_DEVICES=<idle-gpu> python scripts/bench_qdq_cute.py --shape 1024 4096
+```
+
+The benchmark reports exact output equality, maximum absolute error, latency, and speedup for MXFP4 and MXFP8. On an NVIDIA A100 with a `[1024, 4096]` bf16 input, 20 warmup iterations, and 100 measured iterations, the fused kernels produced exact reference output and measured:
+
+| Format | Reference | CuTe | Speedup |
+|---|---:|---:|---:|
+| MXFP4 | 1.437 ms | 0.088 ms | 16.32x |
+| MXFP8 | 0.859 ms | 0.087 ms | 9.83x |
+
+The first call includes CuTe JIT compilation. The table measures warmed-up steady-state execution.
+
+### CUDA Graphs
+
+The fused QDQ kernels support CUDA Graph capture and replay, including the graph path used by vLLM. Each `(format, dtype, shape, device)` specialization must execute once in eager mode before capture so CuTe JIT compilation stays outside the graph. vLLM's normal warmup satisfies this requirement. A cache miss during capture raises an actionable error instead of attempting an unsafe JIT compilation.
+
+Compiled specializations are cached per CUDA device. Kernel launches use PyTorch's current CUDA stream, so capture records the QDQ kernel in the same graph as the following Marlin operation.
+
 ## Adding New Dtypes
 
 1. Create a new QDQ implementation in `src/vllm_qdq_plugin/qdq/` (e.g., `fp8.py`)

benchmark/vllm-qdq-plugin/scripts/bench_qdq_cute.py

@@ -0,0 +1,74 @@
+"""Compare reference and CuTe QDQ accuracy and latency on SM80+ GPUs."""
+
+import argparse
+import os
+import sys
+from collections.abc import Callable
+
+import torch
+
+REPO_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
+sys.path.insert(0, os.path.join(REPO_ROOT, "src"))
+
+from vllm_qdq_plugin.qdq.cute import cute_qdq_status
+from vllm_qdq_plugin.qdq.mxfp4 import _mxfp4_qdq_reference, mxfp4_qdq
+from vllm_qdq_plugin.qdq.mxfp8 import _mxfp8_qdq_reference, mxfp8_qdq
+
+
+def benchmark(fn: Callable[[], torch.Tensor], warmup: int, iterations: int) -> float:
+    for _ in range(warmup):
+        fn()
+    torch.cuda.synchronize()
+    start = torch.cuda.Event(enable_timing=True)
+    end = torch.cuda.Event(enable_timing=True)
+    start.record()
+    for _ in range(iterations):
+        fn()
+    end.record()
+    torch.cuda.synchronize()
+    return start.elapsed_time(end) / iterations
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--shape", nargs=2, type=int, default=(1024, 4096), metavar=("M", "K"))
+    parser.add_argument("--dtype", choices=("float16", "bfloat16"), default="bfloat16")
+    parser.add_argument("--warmup", type=int, default=20)
+    parser.add_argument("--iterations", type=int, default=100)
+    args = parser.parse_args()
+
+    if not torch.cuda.is_available():
+        raise RuntimeError("CUDA is required")
+    dtype = getattr(torch, args.dtype)
+    x = torch.randn(*args.shape, device="cuda", dtype=dtype)
+    available, reason = cute_qdq_status(x)
+    print(f"device: {torch.cuda.get_device_name()} | CuTe status: {reason}")
+    if not available:
+        raise RuntimeError(reason)
+
+    previous_flag = os.environ.get("VLLM_QDQ_CUTE")
+    os.environ["VLLM_QDQ_CUTE"] = "1"
+    try:
+        for name, reference, cute in (
+            ("MXFP4", _mxfp4_qdq_reference, mxfp4_qdq),
+            ("MXFP8", _mxfp8_qdq_reference, mxfp8_qdq),
+        ):
+            reference_out = reference(x)
+            cute_out = cute(x)
+            max_abs_error = (reference_out.float() - cute_out.float()).abs().max().item()
+            equal = torch.equal(reference_out, cute_out)
+            reference_ms = benchmark(lambda: reference(x), args.warmup, args.iterations)
+            cute_ms = benchmark(lambda: cute(x), args.warmup, args.iterations)
+            print(
+                f"{name}: exact={equal} max_abs_error={max_abs_error:.6g} "
+                f"reference={reference_ms:.3f} ms cute={cute_ms:.3f} ms speedup={reference_ms / cute_ms:.2f}x"
+            )
+    finally:
+        if previous_flag is None:
+            os.environ.pop("VLLM_QDQ_CUTE", None)
+        else:
+            os.environ["VLLM_QDQ_CUTE"] = previous_flag
+
+
+if __name__ == "__main__":
+    main()

benchmark/vllm-qdq-plugin/scripts/verify_cute_dsl.py

@@ -0,0 +1,43 @@
+"""Verify that the installed CuTe DSL can compile and execute on the active GPU."""
+
+import torch
+from cutlass import cute
+from cutlass.cute.runtime import from_dlpack
+
+
+@cute.kernel
+def copy_kernel(source: cute.Tensor, destination: cute.Tensor, count: cute.Int32):
+    thread, _, _ = cute.arch.thread_idx()
+    block, _, _ = cute.arch.block_idx()
+    index = block * 256 + thread
+    if index < count:
+        destination[index] = source[index]
+
+
+@cute.jit
+def copy_host(source: cute.Tensor, destination: cute.Tensor, count: cute.Int32):
+    copy_kernel(source, destination, count).launch(
+        grid=[(count + 255) // 256, 1, 1],
+        block=[256, 1, 1],
+    )
+
+
+def main() -> None:
+    if not torch.cuda.is_available():
+        raise RuntimeError("CuTe DSL verification requires CUDA")
+    if torch.cuda.get_device_capability() < (8, 0):
+        raise RuntimeError("CuTe QDQ requires SM80 or newer")
+
+    source = torch.arange(1024, device="cuda", dtype=torch.float32)
+    destination = torch.empty_like(source)
+    cute_source = from_dlpack(source)
+    cute_destination = from_dlpack(destination)
+    compiled = cute.compile(copy_host, cute_source, cute_destination, source.numel())
+    compiled(cute_source, cute_destination, source.numel())
+    torch.cuda.synchronize()
+    torch.testing.assert_close(destination, source)
+    print(f"CuTe DSL SM80 smoke test passed on {torch.cuda.get_device_name()}")
+
+
+if __name__ == "__main__":
+    main()

benchmark/vllm-qdq-plugin/src/vllm_qdq_plugin/envs.py

@@ -72,6 +72,7 @@ def _get_validated_env() -> str | None:
 environment_variables: dict[str, Callable[[], Any]] = {
     "VLLM_QDQ_TRACE": lambda: _env_flag("VLLM_QDQ_TRACE"),
     "VLLM_QDQ": lambda: _env_flag("VLLM_QDQ"),
+    "VLLM_QDQ_CUTE": lambda: _env_flag("VLLM_QDQ_CUTE"),
     "VLLM_MARLIN_MOE_QDQ_MODE": env_with_choices(
         "VLLM_MARLIN_MOE_QDQ_MODE",
         default="0",

benchmark/vllm-qdq-plugin/src/vllm_qdq_plugin/qdq/cute.py

@@ -0,0 +1,73 @@
+# SPDX-License-Identifier: Apache-2.0
+"""CuTe QDQ backend selection and capability checks.
+
+The reference QDQ functions remain the correctness oracle. A fused CuTe kernel
+must be bitwise validated against them before it is enabled for inference.
+"""
+
+import importlib.util
+import warnings
+
+import torch
+
+_FALLBACK_WARNING_EMITTED = False
+
+
+def cute_qdq_status(x: torch.Tensor) -> tuple[bool, str]:
+    """Return whether this tensor can execute a CuTe QDQ kernel."""
+    if not x.is_cuda:
+        return False, "input is not a CUDA tensor"
+    if torch.cuda.get_device_capability(x.device) < (8, 0):
+        return False, "CuTe QDQ requires SM80 or newer"
+    if importlib.util.find_spec("cutlass") is None:
+        return False, "NVIDIA CUTLASS DSL is not installed"
+    return True, "CuTe DSL is available"
+
+
+def _reference_fallback(x: torch.Tensor, group_size: int, format_name: str, reason: str | None = None) -> torch.Tensor:
+    global _FALLBACK_WARNING_EMITTED
+    available, capability_reason = cute_qdq_status(x)
+    if not _FALLBACK_WARNING_EMITTED:
+        status = reason or ("unsupported input" if available else f"unavailable: {capability_reason}")
+        warnings.warn(
+            f"VLLM_QDQ_CUTE=1 requested for {format_name}; {status}. Falling back to the reference QDQ.",
+            RuntimeWarning,
+            stacklevel=2,
+        )
+        _FALLBACK_WARNING_EMITTED = True
+
+    if format_name == "MXFP4":
+        from .mxfp4 import _mxfp4_qdq_reference
+
+        return _mxfp4_qdq_reference(x, group_size)
+
+    from .mxfp8 import _mxfp8_qdq_reference
+
+    return _mxfp8_qdq_reference(x, group_size)
+
+
+def _run_cute_or_fallback(x: torch.Tensor, group_size: int, format_name: str) -> torch.Tensor:
+    available, capability_reason = cute_qdq_status(x)
+    if available and group_size == 32 and x.is_contiguous() and x.shape[-1] % group_size == 0:
+        from .cute_kernels import run_cute_qdq
+
+        return run_cute_qdq(x, format_name)
+    if not available:
+        reason = capability_reason
+    elif group_size != 32:
+        reason = f"group_size={group_size} is unsupported"
+    elif not x.is_contiguous():
+        reason = "input is not contiguous"
+    else:
+        reason = f"K={x.shape[-1]} is not divisible by 32"
+    return _reference_fallback(x, group_size, format_name, reason)
+
+
+def mxfp4_qdq_cute(x: torch.Tensor, group_size: int = 32) -> torch.Tensor:
+    """Run the MXFP4 CuTe backend, or the validated reference fallback."""
+    return _run_cute_or_fallback(x, group_size, "MXFP4")
+
+
+def mxfp8_qdq_cute(x: torch.Tensor, group_size: int = 32) -> torch.Tensor:
+    """Run the MXFP8 CuTe backend, or the validated reference fallback."""
+    return _run_cute_or_fallback(x, group_size, "MXFP8")