Add nvfp4 attention support for vLLM serving

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>