Add nvfp4 attention support for vLLM serving (#1898)

### What does this PR do?

Type of change: New feature

This PR adds an NVFP4+sparse attention serving path for vLLM and
consolidates it with the existing checkpoint-driven sparse-attention
integration.

#### NVFP4 attention

- Applies dynamic block-16 NVFP4 fake quantization to Q/K/P/V.
- Quantizes K before the KV-cache write. V remains pristine until a
complete 16-token group can be finalized once; an incomplete tail is QDQ
on read.
- Adds paged/chunked prefill support and a dedicated split-K decode
kernel. Decode uses a fixed 32-split, 128-key-tile schedule because
split-local P quantization is part of the numerical contract.

#### vLLM integration

- Provides one consolidated worker module:
  - `SparseAttnWorker` for checkpoint-driven sparse attention.
- `QuantSparseAttnWorker` for fixed NVFP4 Q/K/P/V plus optional
checkpoint sparsity.
- Supports the native vLLM FlashAttention and FlashInfer backends. The
worker replaces only the selected implementation and delegates inactive
launches back to that backend.
- Preserves FlashInfer NHD/HND cache layouts and separates mixed
decode/prefill launches so each phase uses the correct kernel contract.
- Restores checkpoint-driven N:M sparse-softmax and skip-softmax
metadata. N:M remains a prefill transform; calibrated skip-softmax
decode uses the shared paged kernel.
- Validates the complete attention plan before modifying the model and
reports unsupported configurations without leaving a partially converted
model.

Supported configurations are regular decoder self-attention with vLLM >=
0.15.0, FlashInfer or FlashAttention, fp16/bf16 model and KV cache,
equal Q/K/V head dimensions divisible by 16, and DCP 1. The README
documents unsupported features and CUDA-graph constraints.

### Usage

Let vLLM select the backend for the model and platform. For example,
Nemotron-H on Blackwell selects FlashInfer automatically.

```bash
cd examples/vllm_serve

python vllm_serve_sparse_attn.py <MODEL_PATH> -tp 8 \
  --no-enable-prefix-caching \
  --worker-cls sparse_attn_worker.QuantSparseAttnWorker
```

If `<MODEL_PATH>/config.json` contains `sparse_attention_config`, the
same worker also applies its N:M or skip-softmax settings. Otherwise, it
runs NVFP4 attention only.

### Testing

Focused attention kernels:

```bash
PYTEST_VERSION=1 PYTHONPATH=$PWD python -m pytest -q \
  tests/gpu/torch/kernels/common/attention/test_triton_fa_p_qdq.py \
  tests/gpu/torch/kernels/common/attention/test_decode_attention.py \
  tests/gpu/torch/kernels/common/attention/test_triton_fa_paged.py
```

- B200 NVFP4 run before final test deduplication: 47 passed.
- Current RTX A6000 run: 23 passed, 21 skipped. The skips require native
E4M3 support (compute capability >= 8.9).

Focused vLLM worker and adapter tests:

```bash
PYTHONPATH=$PWD python -m pytest -q \
  tests/gpu_vllm/torch/sparsity/attention_sparsity/test_sparse_attn_worker.py \
  tests/gpu_vllm/torch/sparsity/attention_sparsity/test_quant_sparse_attn_worker.py
```

Result: 43 passed.

CPU import and launch-contract tests:

```bash
PYTHONPATH=$PWD python -m pytest -q \
  tests/unit/torch/kernels/common/attention/test_triton_fa.py
```

Result: 9 passed.

GitHub Linux, Windows, multi-version, code-quality, documentation, and
required unit checks pass.

### Before your PR is "*Ready for review*"

Make sure you read and follow [Contributor
guidelines](https://github.com/NVIDIA/Model-Optimizer/blob/main/CONTRIBUTING.md)
and your commits are signed (`git commit -s -S`).

Make sure you read and follow the [Security Best
Practices](https://github.com/NVIDIA/Model-Optimizer/blob/main/SECURITY.md#security-coding-practices-for-contributors)
(e.g. avoiding hardcoded `trust_remote_code=True`, `torch.load(...,
weights_only=False)`, `pickle`, etc.).

- Is this change backward compatible?: ✅ — both serving policies are
opt-in; sparse-only remains the launcher default.
- If you copied code from any other sources or added a new PIP
dependency, did you follow guidance in `CONTRIBUTING.md`: N/A — no
copied code or new dependency.
- Did you write any new necessary tests?: ✅
- Did you update
[Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?:
N/A — the example README documents this opt-in serving path; no stable
public API changed.
- Did you get Claude approval on this PR?: N/A

### Additional Information

The fixed attention recipe intentionally does not expose the integration
branch's environment-variable matrix. Backend selection is automatic
unless an explicit vLLM backend override is needed.


<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit

* **New Features**
* Added quant+sparse vLLM serving worker support and clarified compact
NVFP4 attention worker behavior.
* Introduced paged split-K decode and NVFP4 QDQ, including on-write
V-cache quantization and new decode attention API.
* Expanded vLLM runtime installation to support NVFP4 quantization and
sparse attention with FlashInfer metadata compatibility.
* **Bug Fixes**
* Improved NVFP4 degenerate/underflow scale handling to reliably zero
out blocks.
* Strengthened validation for paged-cache and NVFP4 quantization
contracts.
* **Documentation**
* Updated the vLLM sparse-attention example docs with tested versions
and explicit limitations.
* **Tests**
* Expanded GPU correctness and integration coverage for split-K decode,
NVFP4 QDQ, vLLM workers, and runtime installation paths.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Signed-off-by: Kai Xu <kaix@nvidia.com>

Author: kaix-nv

examples/vllm_serve/README.md

@@ -4,7 +4,8 @@ This is a simple example to demonstrate calibrating and serving ModelOpt fakequa
 
 Compared with realquant, fakequant is 2-5x slower, but doesn't require dedicated kernel support and facilitates research.
 
-This example is tested with vllm 0.9.0 and 0.19.1
+The general fakequant example is tested with vLLM 0.9.0 and 0.19.1. The compact
+NVFP4 attention worker documented below requires vLLM 0.15.0 or newer.
 
 ## Prepare environment
 
@@ -101,9 +102,9 @@ QUANT_CFG=<quant_cfg> QUANT_FILE_PATH=<quantizer_state.pth> python vllm_serve_fa
 
 ## Serve a model with sparse attention in vLLM
 
-Apply ModelOpt sparse attention at serve time. The launcher replaces vLLM's `FlashAttentionImpl` with `ModelOptSparseAttentionImpl` (Triton kernel with paged KV cache support) on every attention layer right after model load.
+Apply ModelOpt sparse attention at serve time. Right after model load, the launcher replaces each native attention implementation with its matching ModelOpt adapter: `ModelOptSparseAttentionImpl` for FlashAttention or `ModelOptSparseFlashInferImpl` for FlashInfer. Both adapters use the same Triton kernel with paged KV cache support.
 
-The configuration is read from the checkpoint's `config.json` `sparse_attention_config` block, written by ModelOpt's HF export. The launcher restores calibrated skip-softmax metadata and N:M sparse-softmax metadata (`sparsity_n`, `sparsity_m`, `dense_sink_tokens`, `dense_recent_tokens`). Checkpoints exported with both metadata entries use ModelOpt Triton for sparse prefill launches; decode-only launches and launches without active sparse work delegate back to vLLM FlashAttention.
+The configuration is read from the checkpoint's `config.json` `sparse_attention_config` block, written by ModelOpt's HF export. The launcher restores calibrated skip-softmax metadata and N:M sparse-softmax metadata (`sparsity_n`, `sparsity_m`, `dense_sink_tokens`, `dense_recent_tokens`). Checkpoints exported with both metadata entries use ModelOpt Triton for sparse prefill launches; launches without active sparse work delegate back to the native backend selected by vLLM.
 
 Workflow:
 
@@ -114,12 +115,50 @@ Workflow:
    python vllm_serve_sparse_attn.py <EXPORT_DIR> --enforce-eager -tp 8 --host 0.0.0.0 --port 8000
    ```
 
-If the checkpoint has no `sparse_attention_config`, the worker logs a message and passes through — vLLM runs unchanged. Quant-only flows are handled by `vllm_serve_fakequant.py`; combined sparse + quant will land in a follow-up PR.
+If the checkpoint has no `sparse_attention_config`, the sparse-only installer passes through and vLLM runs unchanged. Whole-model fakequant flows remain handled by `vllm_serve_fakequant.py`; the compact attention-only path is below.
+
+The reusable serving policies live in `modelopt/torch/sparsity/attention_sparsity/plugins/vllm_runtime.py`. `install_vllm_sparse_attention_from_checkpoint` installs checkpoint-driven sparse-only attention, while `install_vllm_nvfp4_attention` installs fixed NVFP4 Q/K/P/V with optional checkpoint sparsity. Both validate every selected layer before publishing any replacement implementation and return a `VllmAttentionInstallReport` with the installed layer names and backend counts.
+
+`sparse_attn_worker.py` only invokes these APIs after vLLM loads the model. It retains `SparseAttnWorker` as the launcher's default and provides `QuantSparseAttnWorker` for the compact NVFP4 policy. Other vLLM integrations can invoke the same library APIs directly:
+
+```python
+from modelopt.torch.sparsity.attention_sparsity.plugins.vllm_runtime import (
+    install_vllm_nvfp4_attention,
+)
+
+report = install_vllm_nvfp4_attention(model_runner, sparse_cfg="checkpoint")
+```
 
 Limitations:
 
 - vLLM V1 chunked prefill and prefix-cache suffix attention are supported by offsetting query positions into the longer KV span.
-- CUDA graph capture is not validated yet — use `--enforce-eager`.
+- `SparseAttnWorker` CUDA graph capture is not validated yet — use `--enforce-eager`.
+
+### Compact NVFP4 attention worker
+
+vLLM 0.15.0 or newer is required when either worker activates a ModelOpt attention transform. Importing `SparseAttnWorker`, or using it with no checkpoint sparse metadata, does not resolve quant-only APIs.
+
+Use the same launcher with the compact worker. By default, vLLM selects the backend for the model and platform; NemotronH on Blackwell selects FlashInfer:
+
+```bash
+python vllm_serve_sparse_attn.py <MODEL_PATH> -tp 8 \
+  --no-enable-prefix-caching \
+  --worker-cls sparse_attn_worker.QuantSparseAttnWorker
+```
+
+The installer supports both FlashInfer and FlashAttention, and the worker prints the installed adapter counts. Pass `--attention-backend FLASHINFER` or `--attention-backend FLASH_ATTN` only when an explicit override is needed.
+
+This attention-only path applies a fixed dynamic block-16 NVFP4 fakequant format to Q/K/P/V. Q is dynamic; missing K/V scales default to global scale 1.0, and P defaults to amax 1.0. Existing scalar attention amax values are preserved, but this path does not calibrate or restore them itself. It does not re-quantize realquant Linear or MoE weights. An optional checkpoint `sparse_attention_config` is still honored.
+
+Decode uses a fixed 32-split, 128-key-tile schedule. P QDQ consumes split-local,
+unnormalized online-softmax probabilities, so changing that schedule can change
+quantized results; split count is part of the numerical contract.
+
+K is QDQ before its cache write, while V is written pristine. Complete 16-token V groups are finalized once in cache; an incomplete tail remains pristine and is QDQ on read. P@V therefore sees uniform fakequant values without re-quantizing the tail.
+
+Supported configurations are regular decoder self-attention with FlashInfer or FlashAttention, fp16/bf16 model and KV cache, equal Q/K/V head dimensions that are multiples of 16, and DCP 1. The FlashInfer adapter preserves both NHD and HND cache strides and separates mixed decode/prefill launches so each phase keeps its own kernel contract. The default `FULL_AND_PIECEWISE` mode remains enabled for fixed N:M and attention-only NVFP4; checkpoints with calibrated decode `threshold_scale_factor` must use a non-`FULL` decode graph mode such as `--enforce-eager` because the live sequence length is not replayed as a Python scalar.
+
+Unsupported features are sliding window, ALiBi, softcap, sinks, FP8 KV cache, cross/encoder/MLA attention, KV sharing or transfer, prefix caching, speculative decoding, DBO/ubatching, and `FULL` mixed/prefill CUDA graphs.
 
 ## Known Problems
 

examples/vllm_serve/sparse_attn_worker.py

@@ -13,108 +13,96 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""Custom vLLM worker for sparse attention.
-
-``SparseAttnWorker``: Replaces ``FlashAttentionImpl`` with
-``ModelOptSparseAttentionImpl`` on each Attention module after model loading.
-The sparse impl uses the ModelOpt Triton kernel for sparse prefill launches.
-Decode-only launches and launches without active sparse work delegate back to
-vLLM FlashAttention.
-
-Configuration flows exclusively through the loaded checkpoint's
-``sparse_attention_config`` block (written by ModelOpt's HF export). If the
-checkpoint has no such block, the worker logs a message and passes through
-unchanged.
-
-Quantization combined with sparse attention is not handled by this worker
-and will land in a follow-up PR once the combined path is tested.
-
-Usage:
-    python vllm_serve_sparse_attn.py <path/to/modelopt-exported-ckpt>
-"""
-
-import importlib
-
-try:
-    _has_legacy_attention_layer = importlib.util.find_spec("vllm.attention.layer") is not None
-except (ModuleNotFoundError, ValueError):
-    _has_legacy_attention_layer = False
-
-if _has_legacy_attention_layer:
-    from vllm.attention.layer import Attention as VLLMAttention
-else:
-    from vllm.model_executor.layers.attention import Attention as VLLMAttention
+"""vLLM worker lifecycle wiring for ModelOpt attention transforms."""
 
 from vllm.v1.worker.gpu_worker import Worker as BaseWorker
 
-from modelopt.torch.sparsity.attention_sparsity.plugins.sparse_attn_config import (
-    load_from_checkpoint_metadata,
-    match_sparse_config,
-)
-from modelopt.torch.sparsity.attention_sparsity.plugins.vllm import (
-    _build_sparse_kw,
-    _clone_sparse_impl,
+from modelopt.torch.sparsity.attention_sparsity.plugins.vllm_runtime import (
+    install_vllm_nvfp4_attention,
+    install_vllm_sparse_attention_from_checkpoint,
 )
 
+__all__ = ["SparseAttnWorker", "QuantSparseAttnWorker"]  # noqa: RUF022
 
-def _replace_attention_impl(worker):
-    """Replace FlashAttentionImpl with ModelOptSparseAttentionImpl on all Attention layers.
+_QUANT_FORMAT_KEYS = ("q_format", "k_format", "p_format", "v_format")
 
-    The sole configuration source is the checkpoint's ``sparse_attention_config``
-    metadata. No-op if the checkpoint has no such block.
-    """
-    hf_config = getattr(worker.model_runner.model_config, "hf_config", None)
-    detected = load_from_checkpoint_metadata(hf_config)
-    if detected is None:
+
+def _unwrapped_model(worker):
+    model = worker.model_runner.model
+    return model.unwrap() if hasattr(model, "unwrap") else model
+
+
+def _print_install_report(policy, report) -> None:
+    if report.installed_count:
+        if policy != "Sparse attention":
+            print(
+                f"[ModelOpt] Installed {policy} (quant+sparse) on "
+                f"{report.installed_count} layers: {dict(report.backend_counts)}"
+            )
+        else:
+            if report.sparse_algorithm:
+                print(f"[ModelOpt] Sparse attention config: algo -> {report.sparse_algorithm}")
+            print(
+                f"[ModelOpt] Sparse attention: replaced impl on {report.installed_count} "
+                f"attention layers: {dict(report.backend_counts)}"
+            )
+    elif report.sparse_algorithm:
+        print(
+            f"[ModelOpt] Sparse attention config {report.sparse_algorithm} matched no active "
+            "attention layers; vLLM remains unchanged"
+        )
+    else:
         print(
             "[ModelOpt] No sparse_attention_config found in the checkpoint; "
-            "skipping sparse attention. Run examples/llm_sparsity/"
-            "attention_sparsity/hf_sa.py to calibrate and export a checkpoint "
-            "with the config embedded."
+            "skipping sparse attention. Run examples/llm_sparsity/attention_sparsity/"
+            "hf_sa.py to calibrate and export a checkpoint with the config embedded."
         )
-        return
-    cfg, preset_name = detected
-    print(f"[ModelOpt] Sparse attention config: algo -> {preset_name}")
-
-    model = worker.model_runner.model
-    if hasattr(model, "unwrap"):
-        model = model.unwrap()
 
-    patched = 0
-    for name, module in model.named_modules():
-        if not isinstance(module, VLLMAttention):
-            continue
 
-        layer_cfg = match_sparse_config(name, cfg)
-        if layer_cfg is None or not layer_cfg.get("enable", True):
-            continue
-
-        sparse_kw = _build_sparse_kw(layer_cfg)
-        if not sparse_kw:
-            # Keep vLLM's original impl when the exported layer config does not
-            # enable any sparse feature.
-            continue
-        new_impl = _clone_sparse_impl(module.impl)
-        new_impl.sparse_kw = sparse_kw
-        module.impl = new_impl
-        patched += 1
-    print(f"[ModelOpt] Sparse attention: replaced impl on {patched} attention layers")
+class SparseAttnWorker(BaseWorker):
+    """Install checkpoint-driven sparse attention after model loading."""
 
+    def load_model(self, *args, **kwargs) -> None:
+        """Load the model, then install checkpoint-configured attention."""
+        super().load_model(*args, **kwargs)
+        report = install_vllm_sparse_attention_from_checkpoint(self.model_runner)
+        _print_install_report("Sparse attention", report)
 
-# ---------------------------------------------------------------------------
-# Workers
-# ---------------------------------------------------------------------------
 
+class QuantSparseAttnWorker(BaseWorker):
+    """Install quantized attention plus optional checkpoint sparsity.
 
-class SparseAttnWorker(BaseWorker):
-    """vLLM worker that uses the ModelOpt sparse attention backend.
+    Per-operand formats come from vLLM's ``--additional-config``; absent keys
+    default to NVFP4 on all four operands (Q/K/P/V)::
 
-    Replaces FlashAttentionImpl with ModelOptSparseAttentionImpl on each
-    Attention module right after model loading — before any forward pass
-    (including determine_available_memory profiling).
+        --additional-config '{"modelopt_attn_quant": {"p_format": "fp8", "v_format": "fp8"}}'
     """
 
+    def _quant_formats(self) -> dict[str, str]:
+        additional = getattr(self.vllm_config, "additional_config", None) or {}
+        formats = additional.get("modelopt_attn_quant", {})
+        unknown = set(formats) - set(_QUANT_FORMAT_KEYS)
+        if unknown:
+            raise ValueError(
+                f"unknown modelopt_attn_quant keys {sorted(unknown)}; "
+                f"allowed: {list(_QUANT_FORMAT_KEYS)}"
+            )
+        return dict(formats)
+
     def load_model(self, *args, **kwargs) -> None:
-        """Load model, then replace attention impl with sparse variant."""
+        """Load the model, then install the configured attention quant recipe."""
         super().load_model(*args, **kwargs)
-        _replace_attention_impl(self)
+        formats = self._quant_formats()
+        report = install_vllm_nvfp4_attention(self.model_runner, sparse_cfg="checkpoint", **formats)
+        policy = "NVFP4 attention" if not formats else f"Quant attention ({formats})"
+        _print_install_report(policy, report)
+
+    def determine_available_memory(self) -> int:
+        """Profile memory without compiling the dynamically converted modules."""
+        # Sparse-only imports must remain independent of quantization-specific APIs.
+        import torch
+
+        from modelopt.torch.quantization.plugins.vllm import disable_compilation
+
+        with torch.inference_mode(), disable_compilation(_unwrapped_model(self)):
+            return BaseWorker.determine_available_memory(self)

examples/vllm_serve/vllm_serve_sparse_attn.py

@@ -15,14 +15,14 @@
 
 """Launch vLLM with sparse attention.
 
-Configuration is read exclusively from ``<ckpt>/config.json``'s
-``sparse_attention_config`` block, written during calibration by
+The default ``SparseAttnWorker`` reads configuration exclusively from
+``<ckpt>/config.json``'s ``sparse_attention_config`` block, written by
 ``examples/llm_sparsity/attention_sparsity/hf_sa.py``. If the checkpoint has
-no such block, the worker logs a message and the server runs as standard
+no such block, that worker logs a message and the server runs as standard
 vLLM.
 
-Combined sparse attention + quantization is not handled by this launcher; it
-will be added in a follow-up PR once the combined path is tested.
+The launcher defaults to ``sparse_attn_worker.SparseAttnWorker``. Pass
+``--worker-cls sparse_attn_worker.QuantSparseAttnWorker`` for quant+sparse.
 
 Usage:
     python vllm_serve_sparse_attn.py <path/to/modelopt-exported-ckpt>