Add vllm qdq benchmark plugin (#2523)

* add qdq

Signed-off-by: yiliu30 <yi4.liu@intel.com>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Signed-off-by: yiliu30 <yi4.liu@intel.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: yiliu30

benchmark/vllm-qdq-plugin/README.md

@@ -0,0 +1,114 @@
+# vllm-qdq-plugin
+
+Out-of-tree [vLLM](https://github.com/vllm-project/vllm) plugin that simulates activation quant-dequant (QDQ) before quantized GEMM kernels. Useful for studying the accuracy impact of "real" quantized compute vs weight-only dequant approaches.
+
+## How It Works
+
+The plugin registers as a `vllm.general_plugins` entry point, which vLLM loads automatically in **all processes** (main + workers). It monkey-patches the low-level op wrappers in `vllm._custom_ops` to inject QDQ on input activations before the actual kernel call. This means:
+
+- Zero vLLM source modifications
+- Works with both `LLM()` Python API and `vllm serve`
+- Covers all call sites automatically (dense linear, MoE gate+up, MoE down)
+
+## Installation
+
+```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/
+```
+
+## Usage
+
+```bash
+# Enable QDQ
+VLLM_QDQ=1 python my_script.py
+
+# With vllm serve
+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
+
+# 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
+```
+
+### Environment Variables
+
+| Variable | Default | Description |
+|---|---|---|
+| `VLLM_QDQ` | `0` | Set to `1` to enable QDQ |
+| `VLLM_QDQ_TRACE` | `0` | Set to `1` to print trace lines (up to 200) |
+| `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
+
+| Dtype | Op | Status | Notes |
+|---|---|---|---|
+| **MXFP4** (E2M1 + E8M0 scales) | `marlin_gemm` | ✅ Supported | Dense quantized linear (MXFP4 via Marlin) |
+| **MXFP4** (E2M1 + E8M0 scales) | `moe_wna16_marlin_gemm` | ✅ Supported | MoE quantized linear (MXFP4 via Marlin) |
+
+### How QDQ Works
+
+For MXFP4, the QDQ simulates:
+1. **Quantize**: Scale activations per group of 32 using E8M0 (power-of-2) scales, then round to nearest FP4 E2M1 value `{0, 0.5, 1, 1.5, 2, 3, 4, 6}`
+2. **Dequantize**: Multiply back by the scale to restore the original dtype
+
+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.
+
+## Adding New Dtypes
+
+1. Create a new QDQ implementation in `src/vllm_qdq_plugin/qdq/` (e.g., `fp8.py`)
+2. Add an `elif` branch in `patch.py` where the dtype check happens
+3. The QDQ function signature: `(x: Tensor, **config) -> Tensor` — same shape and dtype in/out
+
+## License
+
+Apache-2.0
+
+---
+
+## Sage3 Triton Attention Backend (vllm-omni)
+
+This plugin also provides an **out-of-tree diffusion attention backend** for [vllm-omni](https://github.com/vllm-project/vllm-omni), using the [SageAttention3](https://github.com/thu-ml/SageAttention) standalone Triton kernel.
+
+### How It Works
+
+Registers via the `vllm_omni.general_plugins` entry_point. When `VLLM_SAGE3_TRITON=1`, overrides the `SAGE_ATTN` diffusion attention backend with the sage3 Triton implementation. When disabled (default), the original in-tree backend is used unchanged.
+
+- Zero vllm-omni source modifications
+- Conditional activation — doesn't affect normal operation when off
+- Falls back to torch SDPA for cross-attention (different Q/K sequence lengths)
+
+### Usage
+
+```bash
+# Enable sage3 Triton attention for diffusion models
+VLLM_SAGE3_TRITON=1 \
+SAGE3_QUANT_FORMAT=mxfp4 \
+DIFFUSION_ATTENTION_BACKEND=SAGE_ATTN \
+python examples/offline_inference/text_to_image/text_to_image.py \
+  --model /path/to/model ...
+
+# Use original in-tree sage_attn (sageattention v2) — default
+DIFFUSION_ATTENTION_BACKEND=SAGE_ATTN python ...
+
+# Use torch SDPA (no sage at all)
+DIFFUSION_ATTENTION_BACKEND=TORCH_SDPA python ...
+```
+
+### Environment Variables
+
+| Variable | Default | Description |
+|---|---|---|
+| `VLLM_SAGE3_TRITON` | `0` | Set to `1` to enable sage3 Triton backend override |
+| `SAGE3_QUANT_FORMAT` | `mxfp4` | Quantization config for K/V (`mxfp4`, `nvfp4`, `mxfp8_s1`, `mxfp4_s1`) |
+| `SAGE3_ACC_DTYPE` | `fp32` | Accumulator dtype (`fp32`, `bf16_both_dot`, `bf16_pv_only`, etc.) |
+
+### Notes
+
+- **Shared memory requirement**: The sage3 fp32 kernel needs ~192KB shared memory per SM. On GPUs with less (e.g., RTX 6000D with 100KB), use `SAGE3_ACC_DTYPE=bf16_both_dot` or switch to TORCH_SDPA.
+- **Cross-attention**: sage3 requires Q and K to have the same sequence length. Cross-attention calls automatically fall back to torch SDPA.

benchmark/vllm-qdq-plugin/docs/plan/review_force_mxfp4_qdq_plan_20260507_064813.md

@@ -0,0 +1,17 @@
+## Goal
+
+Review the current local changes for forced MXFP4 QDQ support and the move from ad-hoc `os.getenv` reads to a shared `envs` module, then fix any correctness issues.
+
+## Finish Criteria
+
+- Forced MXFP4 mode cannot silently fail because of env parsing or case handling.
+- The package consistently reads public runtime switches through `vllm_qdq_plugin.envs`.
+- The updated behavior is documented where users would look for it.
+- Focused verification passes for import/syntax and env parsing behavior.
+
+## Plan
+
+1. Read the local diff and surrounding files to find correctness gaps.
+2. Patch code paths that mishandle or incompletely apply the new env-based controls.
+3. Add or update documentation for the new env knob if needed.
+4. Run targeted verification and only stop once the new behavior is confirmed.

benchmark/vllm-qdq-plugin/examples/README.md

@@ -0,0 +1,56 @@
+# sage3 attention route example
+
+`sage3_route.json` is an example route file for per-(layer, denoising-step)
+attention quant-config routing. Point the plugin at it with:
+
+```bash
+SAGE3_ROUTE_FILE=examples/sage3_route.json \
+SAGE3_ROUTE_DEBUG=1 \
+VLLM_SAGE3_TRITON=1 \
+DIFFUSION_ATTENTION_BACKEND=SAGE_ATTN \
+  python ...   # your Wan2.2 T2V run
+```
+
+## What this example does
+
+For a **40-step, 40-layer** Wan2.2 run:
+
+| Region | Config | Meaning |
+| --- | --- | --- |
+| steps `0-7` (all layers) | `sdpa` | first 20% of steps run full precision |
+| layer `0` (all steps) | `sdpa` | first layer stays full precision |
+| everything else | `mxfp4` | cheap quant attention |
+
+## Schema
+
+```json
+{
+  "default": "mxfp4",
+  "rules": [
+    {"steps": "0-7", "config": "sdpa"},
+    {"layers": "0",  "config": "sdpa"}
+  ]
+}
+```
+
+- Each rule has optional `layers` and/or `steps` (inclusive range strings like
+  `"0-7,9,12-15"`) and a required `config`.
+- `config` is any key in the sage3 `ATTENTION_CONFIGS` (e.g. `mxfp4`,
+  `mxfp8_s1`) or the literal `"sdpa"` (full precision).
+- Rules are **first-match-wins**. No matching rule falls back to `default`.
+  With no `default`, the impl's own config (`SAGE3_QUANT_FORMAT`) is left alone.
+- `SAGE3_ROUTE_DEBUG=1` logs each unique `(layer, step) -> config` once.
+
+## Indices are absolute — rescale for your step count
+
+`steps` and `layers` are absolute indices, not percentages. `"steps": "0-7"`
+equals "first 20%" **only at 40 steps**. If you change the step count, rescale
+the step rule to keep the same window:
+
+| Steps | First 20% |
+| --- | --- |
+| 20 | `"0-3"` |
+| 40 | `"0-7"` |
+| 50 | `"0-9"` |
+
+The layer rule never changes — Wan2.2 always has 40 layers.

benchmark/vllm-qdq-plugin/examples/sage3_route.json

@@ -0,0 +1,7 @@
+{
+  "default": "mxfp4",
+  "rules": [
+    {"steps": "0-7", "config": "sdpa"},
+    {"layers": "0", "config": "sdpa"}
+  ]
+}

benchmark/vllm-qdq-plugin/pyproject.toml

@@ -0,0 +1,16 @@
+[project]
+name = "vllm-qdq-plugin"
+version = "0.1.0"
+description = "Out-of-tree QDQ (quant-dequant) simulation plugin for vLLM"
+requires-python = ">=3.10"
+dependencies = ["torch"]
+
+[project.entry-points."vllm.general_plugins"]
+qdq = "vllm_qdq_plugin:register"
+
+[project.entry-points."vllm_omni.general_plugins"]
+sage3_attn = "vllm_qdq_plugin:register_omni"
+
+[build-system]
+requires = ["setuptools>=64"]
+build-backend = "setuptools.build_meta"

benchmark/vllm-qdq-plugin/run_wan_sparge_example.sh

@@ -0,0 +1,69 @@
+#!/usr/bin/env bash
+# Wan2.2 T2V using SpargeAttn (block-sparse + quantized attention) as the backend.
+#
+# SpargeAttn replaces self-attention via the plugin's SpargeAttnImpl. The plugin
+# overrides the SAGE_ATTN backend slot when VLLM_SPARGE_ATTN=1, so the host's
+# DIFFUSION_ATTENTION_BACKEND stays "SAGE_ATTN".
+#
+# SPARGE_TOPK=1.0 keeps ALL KV blocks (dense). This isolates SpargeAttn's
+# quantization error for a correctness-first run (no sparsity speedup yet).
+# To enable real block-sparsity later, lower SPARGE_TOPK (e.g. 0.8, 0.5) — but
+# note SpargeAttn is UNtuned on Wan2.2, so video quality at topk<1.0 is unverified.
+set -euo pipefail
+cd /home/yiliu7/workspace/vllm-omni
+
+source /home/yiliu7/workspace/venvs/omni/bin/activate
+export C_INCLUDE_PATH=$HOME/.local/include:$HOME/.local/include/python3.12${C_INCLUDE_PATH:+:$C_INCLUDE_PATH}
+
+# The omni venv has a stale vllm 0.20.0 wheel, but the editable vllm-omni (rebased
+# to v0.22.0) needs symbols only present in the newer vllm source checkout at
+# /home/yiliu7/workspace/vllm (0.21.1rc1, already compiled). Prepend it on
+# PYTHONPATH so `import vllm` resolves to that tree. Non-invasive: no venv mutation.
+# export PYTHONPATH=/home/yiliu7/workspace/vllm${PYTHONPATH:+:$PYTHONPATH}
+
+PY=/home/yiliu7/workspace/venvs/omni/bin/python
+MODEL=/storage/yiliu7/Wan-AI/Wan2.2-T2V-A14B-Diffusers
+
+# height=420
+# width=280
+# num_frames=81
+# num_inference_steps=40
+height=720
+width=1280
+num_frames=81
+num_inference_steps=40
+seed=42
+SPARGE_TOPK=0.5
+
+out_dir=/home/yiliu7/workspace/wan-res
+out_file=t2v_sparge_h${height}_w${width}_f${num_frames}_s${num_inference_steps}_topk${SPARGE_TOPK}_2
+mkdir -p "${out_dir}"
+
+# VLLM_SPARGE_ATTN=1 activates the plugin's SpargeAttn backend (overriding SAGE_ATTN).
+# SPARGE_ATTN_REPO is injected on sys.path so the prebuilt spas_sage_attn package
+# (built in the `sparge` venv) is importable from the `omni` venv — identical
+# torch/CUDA/python ABI, so no rebuild is needed.
+# NOTE: linear/MoE quantization is intentionally OFF (bf16). The host's
+# --quantization mxfp8 path is gated to NPU/XPU only (mxfp8_config.py:144) and
+# raises on CUDA — and it is orthogonal to attention. Leaving it off isolates the
+# SpargeAttn attention integration, which is what this demo exercises.
+VLLM_SPARGE_ATTN=1 \
+SPARGE_ATTN_REPO=/home/yiliu7/workspace/SpargeAttn \
+SPARGE_MODE=topk \
+SPARGE_TOPK=${SPARGE_TOPK} \
+DIFFUSION_ATTENTION_BACKEND=SAGE_ATTN \
+CUDA_VISIBLE_DEVICES=6,7 \
+PYTHONUNBUFFERED=1 \
+  ${PY} -u examples/offline_inference/text_to_video/text_to_video.py \
+    --model ${MODEL} \
+    --prompt "Two anthropomorphic cats in comfy boxing gear and bright gloves fight intensely on a spotlighted stage." \
+    --negative-prompt "色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走" \
+    --height ${height} \
+    --width ${width} \
+    --num-frames ${num_frames} \
+    --guidance-scale 4.0 \
+    --guidance-scale-high 3.0 \
+    --num-inference-steps ${num_inference_steps} \
+    --fps 16 \
+    --tensor-parallel-size 2 \
+    --output ${out_dir}/${out_file}.mp4 2>&1 | tee ${out_file}.log

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

@@ -0,0 +1,111 @@
+"""Benchmark: SAGE3 HW kernels vs BF16 SDPA on [1, 40, 75600, 128]."""
+
+import json
+import os
+import sys
+
+import torch
+import torch.nn.functional as F
+
+# Resolve paths relative to this script to avoid hardcoded user-specific paths
+REPO_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
+sys.path.insert(0, os.path.join(REPO_ROOT, "src"))
+# Bypass top-level vllm dependency by importing the sage3 subpackage directly
+sys.path.insert(0, os.path.join(REPO_ROOT, "src", "vllm_qdq_plugin", "sage3_attn"))
+from sage3.api import sageattn3_standalone
+
+# ── Config ──
+B, H, N, D = 1, 40, 75600, 128
+WARMUP = 3
+ITERS = 10
+DEVICE = "cuda"
+DTYPE = torch.bfloat16
+
+
+def bench(fn, warmup=WARMUP, iters=ITERS):
+    """Returns mean latency in ms."""
+    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(iters):
+        fn()
+    end.record()
+    torch.cuda.synchronize()
+    return start.elapsed_time(end) / iters
+
+
+def compute_attn_flops(B, H, N, D):
+    """FLOPs for standard attention: 2*B*H*N*N*D (QK) + 2*B*H*N*N*D (PV) = 4*B*H*N^2*D."""
+    return 4 * B * H * N * N * D
+
+
+def main():
+    print(f"Shape: [{B}, {H}, {N}, {D}] | dtype: {DTYPE} | warmup: {WARMUP} | iters: {ITERS}")
+    print(f"Device: {torch.cuda.get_device_name()}")
+    print("-" * 70)
+    total_flops = compute_attn_flops(B, H, N, D)
+
+    q = torch.randn(B, H, N, D, device=DEVICE, dtype=DTYPE)
+    k = torch.randn(B, H, N, D, device=DEVICE, dtype=DTYPE)
+    v = torch.randn(B, H, N, D, device=DEVICE, dtype=DTYPE)
+
+    configs = [
+        ("bf16_sdpa", lambda: F.scaled_dot_product_attention(q, k, v)),
+        ("mxfp8_hw", lambda: sageattn3_standalone(q, k, v, config="mxfp8_hw")),
+        ("mxfp4_hw", lambda: sageattn3_standalone(q, k, v, config="mxfp4_hw")),
+        ("mixed_mxfp8qk_mxfp4pv_hw", lambda: sageattn3_standalone(q, k, v, config="mixed_mxfp8qk_mxfp4pv_hw")),
+        ("mixed_mxfp4qk_mxfp8pv_hw", lambda: sageattn3_standalone(q, k, v, config="mixed_mxfp4qk_mxfp8pv_hw")),
+    ]
+
+    results = {}
+    baseline_ms = None
+
+    for name, fn in configs:
+        try:
+            # Verify it runs
+            out = fn()
+            assert out.shape == (B, H, N, D), f"Bad shape: {out.shape}"
+            del out
+            torch.cuda.empty_cache()
+
+            ms = bench(fn)
+            results[name] = {"ms": ms, "status": "ok"}
+            if baseline_ms is None:
+                baseline_ms = ms
+        except Exception as e:
+            results[name] = {"ms": None, "status": f"FAIL: {e}"}
+            print(f"  {name}: FAILED - {e}")
+
+    # Print table
+    print(f"{'Config':<30} {'Latency (ms)':>12} {'TFLOPS':>8} {'Speedup':>8}")
+    print("-" * 62)
+    for name, r in results.items():
+        if r["ms"] is not None:
+            speedup = baseline_ms / r["ms"] if baseline_ms else 0
+            tflops = total_flops / (r["ms"] * 1e-3) / 1e12
+            print(f"{name:<30} {r['ms']:>12.2f} {tflops:>7.1f} {speedup:>7.2f}x")
+            r["speedup"] = round(speedup, 3)
+            r["tflops"] = round(tflops, 2)
+        else:
+            print(f"{name:<30} {'FAIL':>12} {'N/A':>8}")
+
+    # Save raw results
+    out_path = os.path.join(REPO_ROOT, "bench_hw_vs_sdpa_results.json")
+    meta = {
+        "shape": [B, H, N, D],
+        "dtype": str(DTYPE),
+        "warmup": WARMUP,
+        "iters": ITERS,
+        "device": torch.cuda.get_device_name(),
+        "total_flops": total_flops,
+    }
+    with open(out_path, "w") as f:
+        json.dump({"meta": meta, "results": results}, f, indent=2)
+    print(f"\nRaw results saved to: {out_path}")
+
+
+if __name__ == "__main__":
+    main()