Add FP8 BMM2 attention recipe with per-tensor FP8 P and V

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

Author: kaix-nv

examples/vllm_serve/sparse_attn_worker.py

@@ -22,7 +22,7 @@
     install_vllm_sparse_attention_from_checkpoint,
 )
 
-__all__ = ["SparseAttnWorker", "QuantSparseAttnWorker"]  # noqa: RUF022
+__all__ = ["SparseAttnWorker", "QuantSparseAttnWorker", "QuantSparseAttnFp8Bmm2Worker"]  # noqa: RUF022
 
 
 def _unwrapped_model(worker):
@@ -32,7 +32,7 @@ def _unwrapped_model(worker):
 
 def _print_install_report(policy, report) -> None:
     if report.installed_count:
-        if policy == "NVFP4 attention":
+        if "NVFP4" in policy:
             print(
                 "[ModelOpt] Installed NVFP4 quant+sparse attention on "
                 f"{report.installed_count} layers: {dict(report.backend_counts)}"
@@ -85,3 +85,15 @@ def determine_available_memory(self) -> int:
 
         with torch.inference_mode(), disable_compilation(_unwrapped_model(self)):
             return BaseWorker.determine_available_memory(self)
+
+
+class QuantSparseAttnFp8Bmm2Worker(QuantSparseAttnWorker):
+    """NVFP4 BMM1 (Q/K) with per-tensor FP8 E4M3 BMM2 (P/V) plus optional sparsity."""
+
+    def load_model(self, *args, **kwargs) -> None:
+        """Load the model, then install NVFP4-BMM1 + FP8-BMM2 attention."""
+        BaseWorker.load_model(self, *args, **kwargs)
+        report = install_vllm_nvfp4_attention(
+            self.model_runner, sparse_cfg="checkpoint", bmm2="fp8"
+        )
+        _print_install_report("NVFP4-BMM1/FP8-BMM2 attention", report)

modelopt/torch/kernels/common/attention/decode_attention.py

@@ -32,13 +32,15 @@
     _load_paged_v_tile,
 )
 from modelopt.torch.kernels.quantization.attention.bmm2_qdq import _p_qdq_nvfp4, _v_qdq_nvfp4
+from modelopt.torch.kernels.quantization.common.fp8_quant import fp8_scalar_qdq
 
 __all__ = ["attention_decode"]
 
 _BLOCK_N = 128
 _DEFAULT_KV_SPLITS = 32
 _MAX_KV_SPLITS = 32
-_QDQ_MODES = {None, "nvfp4"}
+_P_QDQ_MODES = {None: 0, "fp8": 1, "nvfp4": 2}
+_V_QDQ_MODES = {None, "nvfp4"}
 
 
 @triton.jit
@@ -137,7 +139,11 @@ def _decode_split_kernel(
         running_sum = running_sum * correction + tl.sum(p, axis=0)
         acc *= correction
 
-        if P_QDQ:
+        if P_QDQ == 1:
+            # FP8 E4M3 per-tensor QDQ of the split-local unnormalized P
+            # (elementwise -> split-count- and batch-shape-invariant).
+            p = fp8_scalar_qdq(p, p_qdq_scale)
+        elif P_QDQ == 2:
             p = tl.reshape(
                 _p_qdq_nvfp4(
                     tl.reshape(p.to(V_cache.dtype.element_ty).to(tl.float32), (1, BLOCK_N)),
@@ -238,15 +244,20 @@ def _decode_combine_kernel(
 
 
 def _qdq_scale(mode: str | None, amax: float | None, operand: str) -> float:
-    if mode not in _QDQ_MODES:
-        raise ValueError(f"{operand}_qdq must be 'nvfp4' or None, got {mode!r}")
+    allowed = _P_QDQ_MODES if operand == "p" else _V_QDQ_MODES
+    if mode not in allowed:
+        raise ValueError(
+            f"{operand}_qdq must be one of {sorted(m for m in allowed if m)} or None, got {mode!r}"
+        )
     if mode is None:
         return 1.0
     if amax is None:
         return 1.0
     if not (math.isfinite(amax) and amax > 0.0):
         raise ValueError(f"{operand}_qdq_amax must be finite and positive, got {amax}")
-    return amax / (6.0 * 448.0)
+    # FP8 uses the per-tensor convention ``amax / 448``; NVFP4 the two-level
+    # global scale ``amax / (6 * 448)`` (matches the prefill kernel).
+    return amax / 448.0 if mode == "fp8" else amax / (6.0 * 448.0)
 
 
 def attention_decode(
@@ -337,7 +348,7 @@ def attention_decode(
             PAGE_SIZE=page_size,
             max_blocks_per_seq=block_table.shape[1],
             NUM_KV_SPLITS=num_kv_splits,
-            P_QDQ=p_qdq == "nvfp4",
+            P_QDQ=_P_QDQ_MODES[p_qdq],
             V_QDQ=v_qdq == "nvfp4",
             V_CACHE_QUANTIZED=v_cache_quantized,
             num_warps=4,

modelopt/torch/kernels/common/attention/triton_fa.py

@@ -989,7 +989,9 @@ def forward(
             "IS_CAUSAL": is_causal,
             "HEAD_DIM": HEAD_DIM,
             "STORE_LSE": True,
-            "Q_IS_FP32": q.dtype == torch.float32 and (p_qdq_mode == 2 or v_qdq_mode == 2),
+            # An fp32 Q (the dynamic-quant carrier) always needs the fp32 BMM1 dot;
+            # gating on QDQ modes would miss recipes like fp8-BMM2 (P mode 1).
+            "Q_IS_FP32": q.dtype == torch.float32,
             "SPARSITY_N": sparsity_n,
             "SPARSITY_M": sparsity_m,
             "DENSE_SINK_TOKENS": dense_sink_tokens,

modelopt/torch/quantization/plugins/vllm.py

@@ -47,6 +47,27 @@
         for name in ("q", "k", "p", "v")
     ),
 ]
+_FP8_ATTENTION_QUANTIZER_CFG = {"num_bits": (4, 3)}  # per-tensor E4M3, static scale amax/448
+# Hybrid recipe: NVFP4 BMM1 (Q/K) + per-tensor FP8 E4M3 BMM2 (P/V).
+_VLLM_NVFP4_BMM1_FP8_BMM2_QUANT_CFG = [
+    {"quantizer_name": "*_bmm_quantizer", "enable": False},
+    *(
+        {
+            "quantizer_name": f"*{name}_bmm_quantizer",
+            "cfg": _NVFP4_ATTENTION_QUANTIZER_CFG,
+            "enable": True,
+        }
+        for name in ("q", "k")
+    ),
+    *(
+        {
+            "quantizer_name": f"*{name}_bmm_quantizer",
+            "cfg": _FP8_ATTENTION_QUANTIZER_CFG,
+            "enable": True,
+        }
+        for name in ("p", "v")
+    ),
+]
 
 
 def _import_attention_module():
@@ -224,11 +245,27 @@ def _set_vllm_attention_kv_default_amax(module, device: torch.device) -> None:
         quantizer.amax = torch.tensor(6.0 * 448.0, device=device, dtype=torch.float32)
 
 
+def _set_vllm_attention_fp8_bmm2_default_amax(module, device: torch.device) -> None:
+    """Set fixed amax on uncalibrated per-tensor FP8 BMM2 quantizers (P=1.0, V=448)."""
+    for name, default in (("p_bmm_quantizer", 1.0), ("v_bmm_quantizer", 448.0)):
+        quantizer = getattr(module, name, None)
+        if (
+            not isinstance(quantizer, TensorQuantizer)
+            or not quantizer.is_enabled
+            or getattr(quantizer, "num_bits", None) != (4, 3)
+            or getattr(quantizer, "block_sizes", None)
+            or hasattr(quantizer, "_amax")
+        ):
+            continue
+        quantizer.amax = torch.tensor(default, device=device, dtype=torch.float32)
+
+
 def configure_vllm_nvfp4_attention_quantizers(
     module: torch.nn.Module,
     *,
     device: torch.device | str,
     dtype: torch.dtype,
+    cfg: list | None = None,
 ) -> torch.nn.Module:
     """Configure one vLLM Attention module for fused NVFP4 fake quantization.
 
@@ -256,10 +293,11 @@ def configure_vllm_nvfp4_attention_quantizers(
     if not hasattr(module, "p_bmm_quantizer"):
         module.p_bmm_quantizer = TensorQuantizer()
 
-    set_quantizer_by_cfg(module, _VLLM_NVFP4_ATTENTION_QUANT_CFG)
+    set_quantizer_by_cfg(module, _VLLM_NVFP4_ATTENTION_QUANT_CFG if cfg is None else cfg)
     for name in ("q", "k", "p", "v"):
         getattr(module, f"{name}_bmm_quantizer").to(device=device)
     _set_vllm_attention_kv_default_amax(module, device)
+    _set_vllm_attention_fp8_bmm2_default_amax(module, device)
     return module
 
 

modelopt/torch/sparsity/attention_sparsity/plugins/vllm.py

@@ -121,23 +121,31 @@ def _build_sparse_kw(layer_cfg: dict) -> dict:
 
 
 def _bmm_qdq_from_layer(layer, attr: str, default_amax: float | None):
-    """Map an enabled BMM2 quantizer to the kernel's NVFP4 mode and scalar amax."""
+    """Map an enabled BMM2 quantizer to the kernel's QDQ mode and scalar amax."""
     quantizer = getattr(layer, attr, None)
     if quantizer is None or not getattr(quantizer, "is_enabled", False):
         return None, default_amax
     if (
-        not getattr(quantizer, "is_nvfp4_dynamic", False)
-        or (quantizer.block_sizes or {}).get(-1) != 16
+        getattr(quantizer, "is_nvfp4_dynamic", False)
+        and (quantizer.block_sizes or {}).get(-1) == 16
     ):
+        mode = "nvfp4"
+    elif getattr(quantizer, "num_bits", None) == (4, 3) and not getattr(
+        quantizer, "block_sizes", None
+    ):
+        # Per-tensor FP8 E4M3 (static scale amax/448)
+        mode = "fp8"
+    else:
         raise NotImplementedError(
-            f"{attr} is enabled with an unsupported format; only dynamic block-16 NVFP4 is supported"
+            f"{attr} is enabled with an unsupported format; only dynamic block-16 NVFP4 "
+            "or per-tensor FP8 E4M3 is supported"
         )
     amax = getattr(quantizer, "_amax", None)
     if amax is None:
-        return "nvfp4", default_amax
+        return mode, default_amax
     if getattr(amax, "numel", lambda: 1)() != 1:
         raise NotImplementedError(f"{attr} requires a scalar amax, got shape {tuple(amax.shape)}")
-    return "nvfp4", float(amax)
+    return mode, float(amax)
 
 
 def _p_qdq_from_layer(layer) -> tuple[str | None, float]:

modelopt/torch/sparsity/attention_sparsity/plugins/vllm_runtime.py

@@ -94,6 +94,7 @@ class _InstallPlan:
     layers: tuple[_AttentionPlan, ...]
     quantize: bool
     sparse_algorithm: str | None
+    bmm2: str = "nvfp4"
 
 
 def _unwrapped_model(model_runner):
@@ -315,7 +316,9 @@ def _raise_unsupported(errors: list[str], policy: str) -> None:
         )
 
 
-def _plan_vllm_attention(model_runner, *, quantize: bool, sparse_cfg) -> _InstallPlan:
+def _plan_vllm_attention(
+    model_runner, *, quantize: bool, sparse_cfg, bmm2: str = "nvfp4"
+) -> _InstallPlan:
     model = _unwrapped_model(model_runner)
     resolved_sparse_cfg, sparse_algorithm = _resolve_sparse_config(model_runner, sparse_cfg)
     candidates = []
@@ -329,7 +332,7 @@ def _plan_vllm_attention(model_runner, *, quantize: bool, sparse_cfg) -> _Instal
             candidates.append((name, module, sparse_kw))
 
     if not candidates and not quantize:
-        return _InstallPlan(model_runner, (), False, sparse_algorithm)
+        return _InstallPlan(model_runner, (), False, sparse_algorithm, bmm2)
 
     _require_supported_vllm()
     errors = _global_errors(model_runner) if quantize else []
@@ -373,7 +376,7 @@ def _plan_vllm_attention(model_runner, *, quantize: bool, sparse_cfg) -> _Instal
     if quantize and attention_count == 0:
         errors.append("no regular attention layers were found")
     _raise_unsupported(errors, "NVFP4 attention" if quantize else "sparse attention")
-    return _InstallPlan(model_runner, tuple(plans), quantize, sparse_algorithm)
+    return _InstallPlan(model_runner, tuple(plans), quantize, sparse_algorithm, bmm2)
 
 
 def _build_report(plan: _InstallPlan) -> VllmAttentionInstallReport:
@@ -405,15 +408,28 @@ def _apply_vllm_attention_plans(plan: _InstallPlan) -> VllmAttentionInstallRepor
     for layer in plan.layers:
         layer.new_impl.sparse_kw = layer.sparse_kw
         if plan.quantize:
+            # Pass cfg only for the non-default recipe: keeps the default call
+            # signature stable for callers/fakes that predate the cfg parameter.
+            _cfg_kwargs = (
+                {"cfg": quant_plugin._VLLM_NVFP4_BMM1_FP8_BMM2_QUANT_CFG}
+                if plan.bmm2 == "fp8"
+                else {}
+            )
             converted = quant_plugin.configure_vllm_nvfp4_attention_quantizers(
                 layer.module,
                 device=layer.device,
                 dtype=layer.dtype,
+                **_cfg_kwargs,
             )
             if converted is not None and converted is not layer.module:
                 raise RuntimeError("vLLM attention quantization must convert modules in place")
             p_qdq, p_qdq_amax = attention_plugin._p_qdq_from_layer(layer.module)
             v_qdq, v_qdq_amax = attention_plugin._v_qdq_from_layer(layer.module)
+            if plan.bmm2 == "fp8":
+                # Per-tensor FP8 V is quantized module-level BEFORE the cache
+                # write (each token is self-contained; no block geometry), so
+                # the kernel sees pre-QDQ'd V and needs no V transform at all.
+                v_qdq, v_qdq_amax = None, None
             layer.new_impl.quant_kw = {
                 "p_qdq": p_qdq,
                 "p_qdq_amax": p_qdq_amax,
@@ -426,7 +442,7 @@ def _apply_vllm_attention_plans(plan: _InstallPlan) -> VllmAttentionInstallRepor
         old_value_flag = getattr(layer.module, "_value_quant_in_kernel", missing)
         if plan.quantize:
             layer.module._query_quant_in_kernel = True
-            layer.module._value_quant_in_kernel = True
+            layer.module._value_quant_in_kernel = plan.bmm2 != "fp8"
         try:
             # Publish the adapter last so a native impl never runs with in-kernel
             # quantization flags that only the ModelOpt adapter understands.
@@ -463,6 +479,7 @@ def install_vllm_nvfp4_attention(
     model_runner,
     *,
     sparse_cfg="checkpoint",
+    bmm2: str = "nvfp4",
 ) -> VllmAttentionInstallReport:
     """Install fixed NVFP4 attention with optional checkpoint sparsity.
 
@@ -471,6 +488,8 @@ def install_vllm_nvfp4_attention(
         sparse_cfg: ``"checkpoint"`` to consume optional exported metadata, a
             resolved sparse config dict, or ``None`` for NVFP4-only attention.
     """
+    if bmm2 not in ("nvfp4", "fp8"):
+        raise ValueError(f"bmm2 must be 'nvfp4' or 'fp8', got {bmm2!r}")
     return _apply_vllm_attention_plans(
-        _plan_vllm_attention(model_runner, quantize=True, sparse_cfg=sparse_cfg)
+        _plan_vllm_attention(model_runner, quantize=True, sparse_cfg=sparse_cfg, bmm2=bmm2)
     )

tests/gpu_vllm/torch/quantization/test_vllm_dynamic_modules.py

@@ -48,6 +48,7 @@
 from modelopt.torch.quantization.plugins import vllm as vllm_plugin
 from modelopt.torch.quantization.plugins.vllm import (
     _ATTENTION_TYPES,
+    _VLLM_NVFP4_BMM1_FP8_BMM2_QUANT_CFG,
     VllmMLAAttention,
     _QuantFusedMoEBase,
     _QuantVLLMAttention,
@@ -447,3 +448,37 @@ def test_tiny_deepseek_mla_quantize(tiny_deepseek_llm):
     assert summary["moe_count"] >= 2, summary
 
     _assert_quantizer_amax_is_static(summary)
+
+
+def test_configure_vllm_attention_quantizers_fp8_bmm2(monkeypatch):
+    monkeypatch.setattr(
+        vllm_plugin,
+        "create_parallel_state",
+        lambda: vllm_plugin.ParallelState(data_parallel_group=None),
+    )
+    attention = object.__new__(vllm_plugin.vllm_attention.Attention)
+    torch.nn.Module.__init__(attention)
+
+    converted = configure_vllm_nvfp4_attention_quantizers(
+        attention, device="cpu", dtype=torch.bfloat16, cfg=_VLLM_NVFP4_BMM1_FP8_BMM2_QUANT_CFG
+    )
+
+    # BMM1 unchanged: Q/K dynamic block-16 NVFP4 (F1)
+    for name in ("q", "k"):
+        quantizer = getattr(converted, f"{name}_bmm_quantizer")
+        assert quantizer.is_enabled and quantizer.is_nvfp4_dynamic
+        assert quantizer.block_sizes[-1] == 16
+    assert converted.k_bmm_quantizer._amax == 6.0 * 448.0
+    # BMM2: P/V per-tensor FP8 E4M3 with fixed amax (P=1.0, V=448) (F3)
+    for name, amax in (("p", 1.0), ("v", 448.0)):
+        quantizer = getattr(converted, f"{name}_bmm_quantizer")
+        assert quantizer.is_enabled
+        assert quantizer.num_bits == (4, 3)
+        assert not quantizer.block_sizes
+        assert float(quantizer._amax) == amax
+    # idempotent: calibrated amax survives reconfiguration
+    converted.v_bmm_quantizer.amax = torch.tensor(96.0)
+    reconfigured = configure_vllm_nvfp4_attention_quantizers(
+        converted, device="cpu", dtype=torch.bfloat16, cfg=_VLLM_NVFP4_BMM1_FP8_BMM2_QUANT_CFG
+    )
+    assert float(reconfigured.v_bmm_quantizer._amax) == 96.0

tests/gpu_vllm/torch/sparsity/attention_sparsity/test_quant_sparse_attn_worker.py

@@ -80,3 +80,22 @@ def profile(actual_worker):
         ("exit", "compilation", model),
         ("exit", "inference", None),
     ]
+
+
+def test_fp8_bmm2_worker_delegates_with_fp8(monkeypatch):
+    worker_module = _load_worker_module()
+    calls = []
+    monkeypatch.setattr(BaseWorker, "load_model", lambda *_a, **_k: calls.append("base"))
+    monkeypatch.setattr(
+        worker_module,
+        "install_vllm_nvfp4_attention",
+        lambda runner, **kw: (
+            calls.append(("install", kw))
+            or SimpleNamespace(installed_count=0, sparse_algorithm=None, backend_counts={})
+        ),
+    )
+    instance = object.__new__(worker_module.QuantSparseAttnFp8Bmm2Worker)
+    instance.model_runner = SimpleNamespace()
+    instance.load_model()
+    assert calls[0] == "base"
+    assert calls[1][1] == {"sparse_cfg": "checkpoint", "bmm2": "fp8"}

tests/gpu_vllm/torch/sparsity/attention_sparsity/test_sparse_attn_worker.py

@@ -80,7 +80,11 @@ def guarded_import(name, *args, **kwargs):
     monkeypatch.setattr(builtins, "__import__", guarded_import)
     worker_module = _load_worker_module("sparse_attn_worker_import_test")
 
-    assert worker_module.__all__ == ["SparseAttnWorker", "QuantSparseAttnWorker"]
+    assert worker_module.__all__ == [
+        "SparseAttnWorker",
+        "QuantSparseAttnWorker",
+        "QuantSparseAttnFp8Bmm2Worker",
+    ]
 
 
 @pytest.mark.parametrize(