Add per-operand FP8 attention quant formats for Q/K/P/V

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

Author: kaix-nv

examples/vllm_serve/sparse_attn_worker.py

@@ -22,7 +22,12 @@
     install_vllm_sparse_attention_from_checkpoint,
 )
 
-__all__ = ["SparseAttnWorker", "QuantSparseAttnWorker"]  # noqa: RUF022
+__all__ = [  # noqa: RUF022
+    "SparseAttnWorker",
+    "QuantSparseAttnWorker",
+    "QuantSparseAttnFp8Bmm2Worker",
+    "QuantSparseAttnFp8QkpvWorker",
+]
 
 
 def _unwrapped_model(worker):
@@ -32,7 +37,7 @@ def _unwrapped_model(worker):
 
 def _print_install_report(policy, report) -> None:
     if report.installed_count:
-        if policy == "NVFP4 attention":
+        if policy != "Sparse attention":
             print(
                 "[ModelOpt] Installed NVFP4 quant+sparse attention on "
                 f"{report.installed_count} layers: {dict(report.backend_counts)}"
@@ -85,3 +90,32 @@ 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", p_format="fp8", v_format="fp8"
+        )
+        _print_install_report("NVFP4-BMM1/FP8-BMM2 attention", report)
+
+
+class QuantSparseAttnFp8QkpvWorker(QuantSparseAttnWorker):
+    """Per-tensor FP8 E4M3 on all four attention operands (Q/K/P/V) plus optional sparsity."""
+
+    def load_model(self, *args, **kwargs) -> None:
+        """Load the model, then install all-FP8 attention."""
+        BaseWorker.load_model(self, *args, **kwargs)
+        report = install_vllm_nvfp4_attention(
+            self.model_runner,
+            sparse_cfg="checkpoint",
+            q_format="fp8",
+            k_format="fp8",
+            p_format="fp8",
+            v_format="fp8",
+        )
+        _print_install_report("FP8-QKPV 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,39 @@
         for name in ("q", "k", "p", "v")
     ),
 ]
+_FP8_ATTENTION_QUANTIZER_CFG = {"num_bits": (4, 3)}  # per-tensor E4M3, static scale amax/448
+_BMM2_FORMAT_CFGS = {"nvfp4": _NVFP4_ATTENTION_QUANTIZER_CFG, "fp8": _FP8_ATTENTION_QUANTIZER_CFG}
+
+
+def build_vllm_attention_quant_cfg(
+    *,
+    q_format: str = "nvfp4",
+    k_format: str = "nvfp4",
+    p_format: str = "nvfp4",
+    v_format: str = "nvfp4",
+) -> list:
+    """Build the attention BMM quantizer config with per-operand formats.
+
+    Each of Q/K/P/V takes "nvfp4" (dynamic block-16, two-level scale) or
+    "fp8" (per-tensor E4M3, static scale amax/448).
+    """
+    formats = {"q": q_format, "k": k_format, "p": p_format, "v": v_format}
+    for name, fmt in formats.items():
+        if fmt not in _BMM2_FORMAT_CFGS:
+            raise ValueError(
+                f"{name}_format must be one of {sorted(_BMM2_FORMAT_CFGS)}, got {fmt!r}"
+            )
+    return [
+        {"quantizer_name": "*_bmm_quantizer", "enable": False},
+        *(
+            {
+                "quantizer_name": f"*{name}_bmm_quantizer",
+                "cfg": _BMM2_FORMAT_CFGS[fmt],
+                "enable": True,
+            }
+            for name, fmt in formats.items()
+        ),
+    ]
 
 
 def _import_attention_module():
@@ -224,11 +257,32 @@ 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 (
+        ("q_bmm_quantizer", 448.0),
+        ("k_bmm_quantizer", 448.0),
+        ("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 +310,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,10 @@ class _InstallPlan:
     layers: tuple[_AttentionPlan, ...]
     quantize: bool
     sparse_algorithm: str | None
+    q_format: str = "nvfp4"
+    k_format: str = "nvfp4"
+    p_format: str = "nvfp4"
+    v_format: str = "nvfp4"
 
 
 def _unwrapped_model(model_runner):
@@ -315,7 +319,16 @@ 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,
+    q_format: str = "nvfp4",
+    k_format: str = "nvfp4",
+    p_format: str = "nvfp4",
+    v_format: str = "nvfp4",
+) -> _InstallPlan:
     model = _unwrapped_model(model_runner)
     resolved_sparse_cfg, sparse_algorithm = _resolve_sparse_config(model_runner, sparse_cfg)
     candidates = []
@@ -329,7 +342,9 @@ 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, q_format, k_format, p_format, v_format
+        )
 
     _require_supported_vllm()
     errors = _global_errors(model_runner) if quantize else []
@@ -373,7 +388,16 @@ 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,
+        q_format,
+        k_format,
+        p_format,
+        v_format,
+    )
 
 
 def _build_report(plan: _InstallPlan) -> VllmAttentionInstallReport:
@@ -405,15 +429,36 @@ 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 non-default formats: keeps the default call
+            # signature stable for callers/fakes that predate the cfg parameter.
+            _cfg_kwargs = (
+                {
+                    "cfg": quant_plugin.build_vllm_attention_quant_cfg(
+                        q_format=plan.q_format,
+                        k_format=plan.k_format,
+                        p_format=plan.p_format,
+                        v_format=plan.v_format,
+                    )
+                }
+                if (plan.q_format, plan.k_format, plan.p_format, plan.v_format)
+                != ("nvfp4", "nvfp4", "nvfp4", "nvfp4")
+                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.v_format == "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,
@@ -425,8 +470,10 @@ def _apply_vllm_attention_plans(plan: _InstallPlan) -> VllmAttentionInstallRepor
         old_query_flag = getattr(layer.module, "_query_quant_in_kernel", missing)
         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
+            # fp8 Q is module-level (bf16 losslessly carries E4M3 QDQ values);
+            # the kernel then runs a plain bf16 BMM1 with no Q transform.
+            layer.module._query_quant_in_kernel = plan.q_format != "fp8"
+            layer.module._value_quant_in_kernel = plan.v_format != "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 +510,10 @@ def install_vllm_nvfp4_attention(
     model_runner,
     *,
     sparse_cfg="checkpoint",
+    q_format: str = "nvfp4",
+    k_format: str = "nvfp4",
+    p_format: str = "nvfp4",
+    v_format: str = "nvfp4",
 ) -> VllmAttentionInstallReport:
     """Install fixed NVFP4 attention with optional checkpoint sparsity.
 
@@ -471,6 +522,22 @@ def install_vllm_nvfp4_attention(
         sparse_cfg: ``"checkpoint"`` to consume optional exported metadata, a
             resolved sparse config dict, or ``None`` for NVFP4-only attention.
     """
+    for name, fmt in (
+        ("q_format", q_format),
+        ("k_format", k_format),
+        ("p_format", p_format),
+        ("v_format", v_format),
+    ):
+        if fmt not in ("nvfp4", "fp8"):
+            raise ValueError(f"{name} must be 'nvfp4' or 'fp8', got {fmt!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,
+            q_format=q_format,
+            k_format=k_format,
+            p_format=p_format,
+            v_format=v_format,
+        )
     )