Add Blackwell MXAttention scaled-MMA backend

README.md

@@ -112,3 +112,39 @@ DIFFUSION_ATTENTION_BACKEND=TORCH_SDPA python ...
 
 - **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.
+
+## MXAttention (UOS + PNQ MXFP4)
+
+The plugin also provides an experimental MXAttention backend for NVIDIA
+Blackwell GPUs. It uses MXAttention's UOS scaling (`Qmax=7.25`) and
+Pre-Normalization Quantization (PNQ) for the online probability update. On the
+current Triton 3.6/3.7 Blackwell compiler, the MXFP4 `tl.dot_scaled` lowering
+fails inside the online-softmax loop. MXAttention installs a process-local
+Python workaround for the faulty accumulator-init pass; Triton 3.7.1 is
+required for the split-PV TMEM path. The source-level Triton patch is still
+preferable for production deployment. The MXAttention PNQ path uses native
+`tl.dot_scaled`; the generic non-PNQ compatibility path retains the regular-dot
+fallback.
+
+It is opt-in and overrides the existing `SAGE_ATTN` slot:
+
+```bash
+VLLM_MXATTENTION=1 \
+DIFFUSION_ATTENTION_BACKEND=SAGE_ATTN \
+python examples/offline_inference/image_to_video/image_to_video.py \
+  --model /path/to/Wan2.2-TI2V-5B-Diffusers \
+  --image /path/to/input.jpg \
+  --prompt "smooth natural motion" \
+  --height 320 --width 576 --num-frames 17 --num-inference-steps 2
+```
+
+| Variable | Default | Description |
+|---|---|---|
+| `VLLM_MXATTENTION` | `0` | Enable the MXAttention `SAGE_ATTN` override |
+| `MXATTENTION_MODE` | `mxattention_full` | Ablation mode |
+| `MXATTENTION_QMAX` | `7.25` | UOS maximum |
+| `MXATTENTION_USE_HADAMARD` | `1` | Enable normalized D=128 FWHT rotation |
+
+The initial optimized path supports CUDA Blackwell, head dimension 128,
+equal Q/KV heads, and self-attention. Unsupported masks, cross-attention,
+GQA/MQA, and other shapes fall back to PyTorch SDPA.

docs/MXATTENTION_HANDOFF.md

@@ -0,0 +1,139 @@
+# MXAttention Handoff
+
+## Scope
+
+This branch adds an experimental MXAttention backend for NVIDIA Blackwell
+GPUs. It implements UOS MXFP4 quantization and PNQ probability updates for
+Wan/vLLM-Omni self-attention.
+
+The optimized PNQ path uses native `tl.dot_scaled` for both QK and PV. PV is
+split into four 32-column accumulators because Blackwell TMEM layout
+legalization rejects one `[128, 128]` scaled-MMA accumulator.
+
+## Runtime requirements
+
+- NVIDIA Blackwell / SM100
+- Triton 3.7.1 or newer
+- Torch 2.11.0+cu130 was used for validation
+- Head dimension 128
+- Equal Q/K/V head and sequence counts for the optimized path
+
+The plugin installs a process-local Python workaround for Triton's
+`OptimizeAccumulatorInit` pass when `VLLM_MXATTENTION=1`. This bypasses the
+Blackwell pass that creates an immutable uninitialized TMEM accumulator for
+loop-carried `tl.dot_scaled`. A source-level Triton compiler fix remains the
+preferred production solution.
+
+## Enablement
+
+```bash
+VLLM_MXATTENTION=1 \
+DIFFUSION_ATTENTION_BACKEND=SAGE_ATTN
+```
+
+Defaults:
+
+```text
+MXATTENTION_MODE=mxattention_full
+MXATTENTION_QMAX=7.25
+MXATTENTION_USE_HADAMARD=1
+```
+
+Unsupported masks, cross-attention, GQA/MQA, and unsupported head sizes use
+the SDPA fallback.
+
+## Wan model-card E2E validation
+
+Validated on GPU3 with the local
+`Wan2.2-TI2V-5B-Diffusers` checkpoint using:
+
+- Resolution: 1280x704
+- Frames: 121
+- Inference steps: 50
+- Guidance scale: 5.0
+- FPS: 24
+- Flow shift: 5.0
+- Solver: UniPC
+
+The run completed all 50 steps in approximately 260.5 seconds. The output
+was validated as an MP4 with 121 frames at 1280x704 and 24 FPS:
+
+```text
+/software/hshen/yiliu7/docker_tmp/wan_ti2v_mxattention_modelcard_1280x704_121f_50steps.mp4
+```
+
+Example command:
+
+```bash
+CUDA_VISIBLE_DEVICES=3 \
+PYTHONPATH=/workspace/vllm-qdq-plugin/src \
+VLLM_MXATTENTION=1 \
+DIFFUSION_ATTENTION_BACKEND=SAGE_ATTN \
+/workspace/vllm-omni/.venv/bin/python \
+  /workspace/vllm-omni/examples/offline_inference/image_to_video/image_to_video.py \
+  --model /dev/shm/.hf.tmp/Wan-AI/Wan2.2-TI2V-5B-Diffusers \
+  --image /dev/shm/.hf.tmp/Wan-AI/Wan2.2-TI2V-5B-Diffusers/examples/i2v_input.JPG \
+  --prompt "Two anthropomorphic cats in comfy boxing gear and bright gloves fight intensely on a spotlighted stage." \
+  --height 704 --width 1280 --num-frames 121 \
+  --num-inference-steps 50 --guidance-scale 5.0 \
+  --flow-shift 5.0 --fps 24 --seed 42 \
+  --output wan_ti2v_mxattention_modelcard_1280x704_121f_50steps.mp4
+```
+
+## Wan T2V A14B model-card E2E validation
+
+Validated on GPU3 with the local
+`Wan2.2-T2V-A14B-Diffusers` checkpoint using the model-card settings:
+
+- Resolution: 1280x720
+- Frames: 81
+- Inference steps: 40
+- Guidance scale: 4.0
+- High-noise guidance scale: 3.0
+- FPS: 16
+- Flow shift: 5.0
+- Seed: 42
+
+The run completed all 40 steps in 2684.1 seconds. The output was validated
+as an MP4 with 81 frames at 1280x720 and 16 FPS. CPU offload was enabled to
+fit the A14B checkpoint:
+
+```text
+/software/hshen/yiliu7/docker_tmp/wan_t2v_a14b_mxattention_modelcard_1280x720_81f_40steps.mp4
+```
+
+Example command:
+
+```bash
+CUDA_VISIBLE_DEVICES=3 \
+PYTHONPATH=/workspace/vllm-qdq-plugin/src \
+VLLM_MXATTENTION=1 \
+DIFFUSION_ATTENTION_BACKEND=SAGE_ATTN \
+/workspace/vllm-omni/.venv/bin/python \
+  /workspace/vllm-omni/examples/offline_inference/text_to_video/text_to_video.py \
+  --model /dev/shm/.hf.tmp/Wan-AI/Wan2.2-T2V-A14B-Diffusers \
+  --prompt "Two anthropomorphic cats in comfy boxing gear and bright gloves fight intensely on a spotlighted stage." \
+  --height 720 --width 1280 --num-frames 81 \
+  --num-inference-steps 40 --guidance-scale 4.0 \
+  --guidance-scale-high 3.0 --flow-shift 5.0 --fps 16 --seed 42 \
+  --enable-cpu-offload \
+  --output wan_t2v_a14b_mxattention_modelcard_1280x720_81f_40steps.mp4
+```
+
+## Validation completed
+
+- Native QK/PV `tl.dot_scaled` probes on GPU3
+- Native MXAttention forward, causal and noncausal
+- Padded and multi-block sequence cases
+- Public MXAttention API with `allow_fallback=False`
+- Wan E2E generation with explicit `SAGE_ATTN`
+- MP4 frame-count and resolution validation
+- Python compilation checks and `git diff --check`
+
+## Follow-up
+
+1. Replace the runtime Python workaround with the upstream Triton Blackwell
+   accumulator-init fix when it is available.
+2. Add a full GPU CI job for Triton 3.7.1+ and SM100.
+3. Benchmark the native path against the platform attention backend at the
+   model-card resolution and sequence length.

scripts/run_wan_t2v_a14b_mxattention.sh

@@ -0,0 +1,34 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+MODEL=/dev/shm/.hf.tmp/Wan-AI/Wan2.2-T2V-A14B-Diffusers
+OUTPUT=/workspace/vllm-omni/wan_t2v_a14b_mxattention_modelcard_1280x720_81f_40steps.mp4
+RESULTS_DIR=/software/hshen/yiliu7/docker_tmp
+
+cd /workspace/vllm-omni
+
+CUDA_VISIBLE_DEVICES=7 \
+PYTHONPATH=/workspace/vllm-qdq-plugin/src \
+VLLM_MXATTENTION=1 \
+DIFFUSION_ATTENTION_BACKEND=SAGE_ATTN \
+/workspace/vllm-omni/.venv/bin/python \
+  /workspace/vllm-omni/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 720 \
+  --width 1280 \
+  --num-frames 81 \
+  --num-inference-steps 40 \
+  --guidance-scale 4.0 \
+  --guidance-scale-high 3.0 \
+  --flow-shift 5.0 \
+  --fps 16 \
+  --seed 42 \
+  --enable-cpu-offload \
+  --output "$OUTPUT"
+
+test -s "$OUTPUT"
+cp "$OUTPUT" "$RESULTS_DIR/"
+echo "Saved: $OUTPUT"
+echo "Copied: $RESULTS_DIR/$(basename "$OUTPUT")"

src/vllm_qdq_plugin/__init__.py

@@ -149,6 +149,41 @@ def register_omni_sla_attn():
         )
 
 
+def register_omni_mxattention():
+    try:
+        from .mxattention.triton_workaround import install as install_triton_workaround
+
+        if install_triton_workaround():
+            logger.warning(
+                "vllm-qdq-plugin: enabled SM100 Triton accumulator-init workaround "
+                "for MXAttention"
+            )
+        else:
+            logger.warning(
+                "vllm-qdq-plugin: Triton accumulator-init workaround unavailable"
+            )
+
+        from vllm_omni.diffusion.attention.backends.registry import (
+            DiffusionAttentionBackendEnum,
+            register_diffusion_backend,
+        )
+
+        register_diffusion_backend(
+            DiffusionAttentionBackendEnum.SAGE_ATTN,
+            "vllm_qdq_plugin.mxattention.backend.MXAttentionBackend",
+        )
+        logger.warning(
+            "vllm-qdq-plugin: registered MXAttention backend as SAGE_ATTN "
+            "(VLLM_MXATTENTION=1)"
+        )
+    except ImportError as e:
+        logger.warning(
+            "vllm-qdq-plugin: cannot register MXAttention backend — "
+            "vllm_omni not available (%s)",
+            e,
+        )
+
+
 def _maybe_install_route(route_file: str):
     """Install per-(layer, step) attention routing if a route file is set.
 
@@ -180,15 +215,23 @@ def register_omni():
     sage3_requested = envs.VLLM_SAGE3_TRITON or envs.VLLM_SAGE3_CUTE
     overrides_requested = sum(
         bool(flag)
-        for flag in (envs.VLLM_SLA_ATTN, envs.VLLM_SPARGE_ATTN, sage3_requested)
+        for flag in (
+            envs.VLLM_SLA_ATTN,
+            envs.VLLM_SPARGE_ATTN,
+            envs.VLLM_MXATTENTION,
+            sage3_requested,
+        )
     )
     if overrides_requested > 1:
         raise RuntimeError(
-            "vllm-qdq-plugin: VLLM_SLA_ATTN, VLLM_SPARGE_ATTN, and "
-            "VLLM_SAGE3_{TRITON,CUTE} are mutually exclusive"
+            "vllm-qdq-plugin: VLLM_SLA_ATTN, VLLM_SPARGE_ATTN, "
+            "VLLM_MXATTENTION, and VLLM_SAGE3_{TRITON,CUTE} are mutually exclusive"
         )
 
-    if envs.VLLM_SLA_ATTN:
+    if envs.VLLM_MXATTENTION:
+        register_omni_mxattention()
+        logger.warning_once("vllm-qdq-plugin: registered MXAttention backend for vllm-omni")
+    elif envs.VLLM_SLA_ATTN:
         register_omni_sla_attn()
         logger.warning_once("vllm-qdq-plugin: registered SLA backend for vllm-omni")
     elif envs.VLLM_SPARGE_ATTN:
@@ -211,5 +254,6 @@ def register_omni():
         logger.warning_once(
             "vllm-qdq-plugin: no custom attention backend registered for "
             "vllm-omni — set VLLM_SLA_ATTN=1, VLLM_SPARGE_ATTN=1, "
+            "VLLM_MXATTENTION=1, "
             "VLLM_SAGE3_TRITON=1, or VLLM_SAGE3_CUTE=1 to enable"
         )

src/vllm_qdq_plugin/envs.py

@@ -110,6 +110,11 @@ def _get_validated_env() -> str | None:
     "SLA_COMPARE_MAX_CALLS": lambda: os.getenv("SLA_COMPARE_MAX_CALLS", "0"),
     # SpargeAttn block-sparse attention backend (mutually exclusive with sage3).
     "VLLM_SPARGE_ATTN": lambda: _env_flag("VLLM_SPARGE_ATTN"),
+    # MXAttention UOS/PNQ backend (overrides the SAGE_ATTN slot).
+    "VLLM_MXATTENTION": lambda: _env_flag("VLLM_MXATTENTION"),
+    "MXATTENTION_MODE": lambda: os.getenv("MXATTENTION_MODE", "mxattention_full").lower(),
+    "MXATTENTION_QMAX": lambda: os.getenv("MXATTENTION_QMAX", "7.25"),
+    "MXATTENTION_USE_HADAMARD": lambda: _env_flag("MXATTENTION_USE_HADAMARD", "1"),
     # Path to the SpargeAttn repo; injected on sys.path at registration when the
     # spas_sage_attn package is not already importable. Empty string = rely on the
     # package already being installed/importable.

src/vllm_qdq_plugin/mxattention/__init__.py

@@ -0,0 +1,15 @@
+"""MXAttention: UOS/PNQ MXFP4 attention for Blackwell GPUs."""
+
+from .api import mxattention_forward
+from .reference import (
+    dequantize_mxfp4_uos,
+    materialized_pnq_attention,
+    quantize_mxfp4_uos,
+)
+
+__all__ = [
+    "mxattention_forward",
+    "quantize_mxfp4_uos",
+    "dequantize_mxfp4_uos",
+    "materialized_pnq_attention",
+]