Remove redundant NVFP4 attention paths

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

Author: kaix-nv

examples/vllm_serve/sparse_attn_worker.py

@@ -16,7 +16,6 @@
 """Custom vLLM workers for checkpoint-driven sparse and fixed-NVFP4 attention."""
 
 import importlib
-import os
 from collections import Counter
 from functools import cache
 from types import SimpleNamespace
@@ -291,21 +290,12 @@ def _install_sparse_plans(plans) -> None:
 
 def _install_quant_plans(worker, plans) -> None:
     api = _quant_api()
-    quant_off = os.environ.get("MODELOPT_ATTN_QUANT_OFF") == "1"
     for plan in plans:
         module = plan.module
         module.device, module.dtype = plan.device, plan.dtype
         api.nn.QuantModuleRegistry.convert(module)
         module.p_bmm_quantizer = api.nn.TensorQuantizer()
         api.conversion.set_quantizer_by_cfg(module, _BMM_CFG)
-        if quant_off:
-            # Isolation knob: keep the ModelOpt kernel fixed while disabling all
-            # Q/K/P/V NVFP4 transforms, so quant-on minus quant-off isolates the
-            # fakequant loss without triggering the native dense fallback.
-            for name in ("q", "k", "p", "v"):
-                quantizer = getattr(module, f"{name}_bmm_quantizer", None)
-                if quantizer is not None:
-                    quantizer.disable()
         api.plugin._set_vllm_attention_kv_default_amax(module, plan.device)
         plan.new_impl.sparse_kw = plan.sparse_kw
         p_qdq, p_qdq_amax = _p_qdq_from_layer(module)
@@ -317,10 +307,8 @@ def _install_quant_plans(worker, plans) -> None:
             "v_qdq_amax": v_qdq_amax,
         }
         module.impl = plan.new_impl
-        module._query_quant_in_kernel = not quant_off
-        module._value_quant_in_kernel = not quant_off
-        if quant_off:
-            module._modelopt_force_kernel = True
+        module._query_quant_in_kernel = True
+        module._value_quant_in_kernel = True
     worker.model_runner.cascade_attn_enabled = False
     installed = dict(Counter(type(plan.new_impl).__name__ for plan in plans))
     print(f"[ModelOpt] Installed NVFP4 quant+sparse attention on {len(plans)} layers: {installed}")

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

@@ -71,8 +71,9 @@ def _load_qdq_helpers() -> None:
         _v_qdq_nvfp4 = _v_nvfp4
 
 
-# Maps public P/V QDQ options to kernel constexpr values.
-_QDQ_MODES = {None: 0, "fp8": 1, "nvfp4": 2}
+# Maps public QDQ options to kernel constexpr values.
+_P_QDQ_MODES = {None: 0, "fp8": 1, "nvfp4": 2}
+_V_QDQ_MODES = {None: 0, "nvfp4": 2}
 
 
 LOG2E: float = 1.44269504088896
@@ -289,7 +290,7 @@ def _attn_fwd(
     SKIP_THRESHOLD_LOG2: tl.constexpr = 0.0,  # log2(lambda) in the kernel's scaled log2 score space
     P_QDQ: tl.constexpr = 0,  # Fake quant-dequant of softmax P: 0=off, 1=FP8 E4M3, 2=NVFP4
     p_qdq_scale=1.0,  # Per-tensor scale for softmax qdq (runtime scalar; amax/448 or amax/(6*448))
-    V_QDQ: tl.constexpr = 0,  # Fake quant-dequant of V: 0=off, 1=FP8 E4M3, 2=NVFP4
+    V_QDQ: tl.constexpr = 0,  # Fake quant-dequant of V: 0=off, 2=NVFP4
     v_qdq_scale=1.0,
     V_CACHE_QUANTIZED: tl.constexpr = False,  # complete block-16 groups are already QDQ
     Sparsity_total=None,  # Optional int64 scalar for counting total tiles (atomic)
@@ -476,13 +477,10 @@ def _attn_fwd(
                     mask=((kv_start + kv_pos[:, None]) < seq_len_kv) & d_mask[None, :],
                     other=0.0,
                 )
-            if V_QDQ != 0 and (
+            if V_QDQ == 2 and (
                 (not V_CACHE_QUANTIZED) or (kv_start + BLOCK_N > v_quantized_boundary)
             ):
-                if V_QDQ == 1:
-                    v_qdq = _qdq_fp8(v.to(tl.float32), v_qdq_scale)
-                else:
-                    v_qdq = _v_qdq_nvfp4(v.to(tl.float32), v_qdq_scale, BLOCK_N, BLOCK_D)
+                v_qdq = _v_qdq_nvfp4(v.to(tl.float32), v_qdq_scale, BLOCK_N, BLOCK_D)
                 v_qdq = v_qdq.to(v.dtype)
                 if V_CACHE_QUANTIZED:
                     use_qdq = (kv_start + kv_pos) >= v_quantized_boundary
@@ -1298,12 +1296,11 @@ def attention(
             and the global scale ``amax / (6 * 448)`` for NVFP4. A runtime
             scalar — user-set or calibrated values do not recompile the
             kernel. Values above amax saturate.
-        v_qdq: Fake quant-dequant of V before ``P @ V``. ``"fp8"`` uses a
-            per-tensor E4M3 scale; ``"nvfp4"`` uses signed E2M1 values with
-            one E4M3 scale per 16 keys. ``None`` disables V QDQ.
+        v_qdq: Fake quant-dequant of V before ``P @ V``. ``"nvfp4"`` uses
+            signed E2M1 values with one E4M3 scale per 16 keys. ``None``
+            disables V QDQ.
         v_qdq_amax: Optional per-tensor V amax. ``None`` uses global scale 1;
-            otherwise converts to ``amax / 448`` (FP8) or ``amax / (6 * 448)``
-            (NVFP4).
+            otherwise converts to the NVFP4 global scale ``amax / (6 * 448)``.
         v_cache_quantized: Complete block-16 groups in the paged V cache are
             already QDQ; only the pristine partial group is QDQ on read.
         k_cache: Paged K cache [num_blocks, page_size, num_kv_heads, head_dim].
@@ -1330,11 +1327,11 @@ def attention(
     # silently reuse stale compiled kernels from the on-disk cache.
     _load_sparsity_helpers()
     _load_qdq_helpers()
-    if p_qdq not in _QDQ_MODES:
+    if p_qdq not in _P_QDQ_MODES:
         raise ValueError(
-            f"p_qdq must be one of {sorted(k for k in _QDQ_MODES if k)} or None, got {p_qdq!r}"
+            f"p_qdq must be one of {sorted(k for k in _P_QDQ_MODES if k)} or None, got {p_qdq!r}"
         )
-    p_qdq_mode = _QDQ_MODES[p_qdq]
+    p_qdq_mode = _P_QDQ_MODES[p_qdq]
     # Convert the per-tensor amax to the kernel's scale convention
     # (``q = cast(p / scale) * scale``): FP8 uses ``amax / 448``; NVFP4 uses the
     # global scale ``amax / (6 * 448)``. amax=1 (the default, the theoretical
@@ -1344,18 +1341,18 @@ def attention(
         if not (math.isfinite(p_qdq_amax) and p_qdq_amax > 0):
             raise ValueError(f"p_qdq_amax must be a finite positive value, got {p_qdq_amax}")
         p_qdq_scale = p_qdq_amax / 448.0 if p_qdq == "fp8" else p_qdq_amax / (6.0 * 448.0)
-    if v_qdq not in _QDQ_MODES:
+    if v_qdq not in _V_QDQ_MODES:
         raise ValueError(
-            f"v_qdq must be one of {sorted(k for k in _QDQ_MODES if k)} or None, got {v_qdq!r}"
+            f"v_qdq must be one of {sorted(k for k in _V_QDQ_MODES if k)} or None, got {v_qdq!r}"
         )
-    v_qdq_mode = _QDQ_MODES[v_qdq]
+    v_qdq_mode = _V_QDQ_MODES[v_qdq]
     if v_qdq_mode and any(t.requires_grad for t in (q, k, v)):
         raise NotImplementedError("v_qdq is inference-only and does not support autograd")
     v_qdq_scale = 1.0
     if v_qdq_mode and v_qdq_amax is not None:
         if not (math.isfinite(v_qdq_amax) and v_qdq_amax > 0):
             raise ValueError(f"v_qdq_amax must be a finite positive value, got {v_qdq_amax}")
-        v_qdq_scale = v_qdq_amax / 448.0 if v_qdq == "fp8" else v_qdq_amax / (6.0 * 448.0)
+        v_qdq_scale = v_qdq_amax / (6.0 * 448.0)
     if v_cache_quantized and v_qdq != "nvfp4":
         raise ValueError("v_cache_quantized requires v_qdq='nvfp4'")
     if v_cache_quantized and any(x is None for x in (k_cache, v_cache, block_table)):

modelopt/torch/kernels/quantization/attention/v_qdq.py

@@ -86,15 +86,14 @@ def fake_quant_v_onwrite(
     max_new_tokens: int,
     page_size: int = 16,
     v_qdq_scale: float = 1.0,
-    decode: bool = False,
 ) -> None:
     """NVFP4-finalize complete block-16 groups in ``[v_lo, v_hi)`` in place.
 
     ``max_new_tokens`` is host metadata used to size the masked launch grid.
-    Decode uses one fixed group per request. Eager prefill covers every group
-    that the largest query chunk can complete without reading device metadata.
-    ``v_lo`` and ``v_hi`` must describe aligned, completed block-16 boundaries;
-    their device values are not host-validated.
+    The grid covers every group that the largest query chunk can complete
+    without reading device metadata. ``v_lo`` and ``v_hi`` must describe
+    aligned, completed block-16 boundaries; their device values are not
+    host-validated.
     """
     if page_size != v_cache.shape[1]:
         raise ValueError(f"page_size {page_size} must match v_cache.shape[1] {v_cache.shape[1]}")
@@ -105,7 +104,7 @@ def fake_quant_v_onwrite(
 
     batch, max_blocks = block_table.shape
     num_kv_heads, head_dim = v_cache.shape[2:]
-    num_groups = 1 if decode else triton.cdiv(max_new_tokens, _BLOCK_N)
+    num_groups = triton.cdiv(max_new_tokens, _BLOCK_N)
 
     with torch.cuda.device(v_cache.device):
         _fake_quant_v_onwrite_kernel[(batch, num_kv_heads, num_groups)](

modelopt/torch/kernels/quantization/gemm/fp4_kernel_hopper.py

@@ -81,6 +81,7 @@ def fp4_fake_quant_kernel(
     block_max = tl.max(x_abs, axis=2, keep_dims=True)
 
     block_max_quant = fp8_quantize_scale(block_max, global_scale_safe)
+    block_max_quant = tl.where(block_max_quant >= 1e-5, block_max_quant, 1.0)
 
     block_max_quant_broadcast = tl.broadcast_to(
         block_max_quant, (TILE_M, NUM_FP4_BLOCKS, BLOCK_SIZE)

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

@@ -171,14 +171,9 @@ def _any_quant_active(layer, p_qdq, v_qdq) -> bool:
     )
 
 
-def _should_run_modelopt_kernel(sparse_kw, quant_active: bool, force_kernel: bool) -> bool:
+def _should_run_modelopt_kernel(sparse_kw, quant_active: bool) -> bool:
     """Return whether a launch has effective work for the ModelOpt kernel."""
-    return bool(sparse_kw or quant_active or force_kernel)
-
-
-def _mixed_batch_needs_dense_base(quant_active: bool, force_kernel: bool) -> bool:
-    """Return whether sparse-only mixed phases need one native output seed."""
-    return not quant_active and not force_kernel
+    return bool(sparse_kw or quant_active)
 
 
 @dataclass(frozen=True, slots=True)
@@ -188,7 +183,6 @@ class _ResolvedForward:
     v_qdq: str | None
     v_qdq_amax: float | None
     quant_active: bool
-    force_kernel: bool
 
 
 def _resolve_forward(
@@ -203,10 +197,7 @@ def _resolve_forward(
     """Resolve shared transform state or request the backend's native path."""
     p_qdq, p_qdq_amax, v_qdq, v_qdq_amax = _quant_kw_from_impl(impl, layer)
     quant_active = _any_quant_active(layer, p_qdq, v_qdq)
-    force_kernel = getattr(layer, "_modelopt_force_kernel", False)
-    transform_active = _should_run_modelopt_kernel(
-        getattr(impl, "sparse_kw", None), quant_active, force_kernel
-    )
+    transform_active = _should_run_modelopt_kernel(getattr(impl, "sparse_kw", None), quant_active)
 
     if getattr(attn_metadata, "use_cascade", False):
         if transform_active:
@@ -234,7 +225,6 @@ def _resolve_forward(
         v_qdq=v_qdq,
         v_qdq_amax=v_qdq_amax,
         quant_active=quant_active,
-        force_kernel=force_kernel,
     )
 
 
@@ -260,7 +250,6 @@ def _forward_modelopt(
     v_qdq: str | None,
     v_qdq_amax: float | None,
     quant_active: bool,
-    force_kernel: bool,
     dense_fallback,
     prepare_modelopt=None,
 ) -> torch.Tensor:
@@ -281,16 +270,9 @@ def _forward_modelopt(
         # N:M sparse softmax is prefill-only.
         for name in ("sparsity_n", "sparsity_m", "dense_sink_tokens", "dense_recent_tokens"):
             sparse_kw.pop(name, None)
-    if not _should_run_modelopt_kernel(
-        sparse_kw,
-        quant_active,
-        force_kernel,
-    ):
+    if not _should_run_modelopt_kernel(sparse_kw, quant_active):
         # Dynamic calibration can disable sparse work for a launch. Preserve the
         # backend's native dense path when no ModelOpt transform remains active.
-        # ``_modelopt_force_kernel`` (isolation knob) keeps the bf16 kernel path
-        # even with all quant disabled, so the on/off PPL diff holds the kernel
-        # constant.
         return dense_fallback()
     if prepare_modelopt is not None:
         prepare_modelopt()
@@ -309,7 +291,6 @@ def _forward_modelopt(
             max_new_tokens=max_query_len,
             page_size=page_size,
             v_qdq_scale=v_qdq_scale,
-            decode=is_decode_only,
         )
 
     q = query[:num_actual_tokens].contiguous()
@@ -459,7 +440,6 @@ def native_forward():
             v_qdq=resolved.v_qdq,
             v_qdq_amax=resolved.v_qdq_amax,
             quant_active=resolved.quant_active,
-            force_kernel=resolved.force_kernel,
             dense_fallback=native_forward,
         )
 
@@ -633,7 +613,6 @@ def prepare_modelopt():
         "v_qdq": resolved.v_qdq,
         "v_qdq_amax": resolved.v_qdq_amax,
         "quant_active": resolved.quant_active,
-        "force_kernel": resolved.force_kernel,
         "dense_fallback": dense_fallback,
         "prepare_modelopt": prepare_modelopt,
     }
@@ -662,10 +641,7 @@ def prepare_modelopt():
     # Sparse-only launches may have an inactive phase (for example N:M
     # is prefill-only). Compute the native result once, then overwrite
     # each active phase with its ModelOpt result.
-    if _mixed_batch_needs_dense_base(
-        resolved.quant_active,
-        resolved.force_kernel,
-    ):
+    if not resolved.quant_active:
         dense_fallback()
 
     block_table = attn_metadata._modelopt_block_table

tests/gpu/torch/kernels/common/attention/test_triton_fa_p_qdq.py

@@ -488,6 +488,7 @@ def test_invalid_mode_raises(self):
         ("kwargs", "match"),
         [
             ({"v_qdq": "int8"}, "v_qdq"),
+            ({"v_qdq": "fp8"}, "v_qdq"),
             ({"v_qdq": "nvfp4", "v_qdq_amax": 0.0}, "v_qdq_amax"),
             ({"v_qdq": "nvfp4", "v_cache_quantized": True}, "v_cache_quantized"),
         ],

tests/gpu/torch/kernels/common/attention/test_triton_fa_paged.py

@@ -279,60 +279,6 @@ def test_paged_different_page_sizes(self, page_size):
 
         torch.testing.assert_close(out_paged, out_contig, rtol=1e-2, atol=1e-2)
 
-    def test_paged_decode_ignores_sparse_nm(self):
-        """N:M sparse softmax is prefill-only; paged decode remains dense."""
-        batch = 2
-        seq_lens_k = [64, 128]
-        num_heads, num_kv_heads, head_dim = 4, 2, 64
-        page_size = 16
-        scale = 1.0 / (head_dim**0.5)
-        total_kv = sum(seq_lens_k)
-
-        torch.manual_seed(34)
-        q_flat = torch.randn(batch, num_heads, head_dim, device="cuda", dtype=torch.float16)
-        k_flat = torch.randn(total_kv, num_kv_heads, head_dim, device="cuda", dtype=torch.float16)
-        v_flat = torch.randn(total_kv, num_kv_heads, head_dim, device="cuda", dtype=torch.float16)
-
-        b_start_loc_q = torch.arange(batch, device="cuda", dtype=torch.int32)
-        b_seq_len_q = torch.ones(batch, device="cuda", dtype=torch.int32)
-        cumsum = [0]
-        for sl in seq_lens_k:
-            cumsum.append(cumsum[-1] + sl)
-        b_start_loc_k = torch.tensor(cumsum[:-1], device="cuda", dtype=torch.int32)
-        b_seq_len_k = torch.tensor(seq_lens_k, device="cuda", dtype=torch.int32)
-        k_cache, v_cache, block_table = _scatter_to_paged_cache(
-            k_flat, v_flat, b_start_loc_k, b_seq_len_k, num_kv_heads, head_dim, page_size
-        )
-
-        common_kwargs = {
-            "is_causal": False,
-            "softmax_scale": scale,
-            "b_start_loc_k": b_start_loc_k,
-            "b_seq_len_k": b_seq_len_k,
-            "max_input_len_k": max(seq_lens_k),
-            "k_cache": k_cache,
-            "v_cache": v_cache,
-            "block_table": block_table,
-            "page_size": page_size,
-        }
-        out_dense = attention(
-            q_flat, k_flat, v_flat, b_start_loc_q, b_seq_len_q, 1, **common_kwargs
-        )
-        out_sparse = attention(
-            q_flat,
-            k_flat,
-            v_flat,
-            b_start_loc_q,
-            b_seq_len_q,
-            1,
-            **common_kwargs,
-            sparsity_n=2,
-            sparsity_m=4,
-            dense_recent_tokens=64,
-        )
-
-        torch.testing.assert_close(out_sparse, out_dense, rtol=1e-2, atol=1e-2)
-
     def test_paged_mixed_prefill_decode_sparse_nm_keeps_decode_dense(self):
         """Decode rows stay dense even when batched with sparse prefill rows."""
         q_lens = [64, 1]
@@ -447,7 +393,6 @@ def test_v_cache_finalizes_complete_groups_once(self):
             torch.tensor([48], device="cuda", dtype=torch.int32),
             max_new_tokens=1,
             page_size=page_size,
-            decode=True,
         )
         torch.testing.assert_close(baked[:2], after_prefill[:2], rtol=0, atol=0)
         assert torch.all(baked[:3] == 0.015625)

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

@@ -134,29 +134,6 @@ def test_install_converts_only_attention_and_configures_fixed_recipe(monkeypatch
     assert state.model_runner.cascade_attn_enabled is False
 
 
-def test_quant_off_disables_recipe_but_keeps_modelopt_adapter(monkeypatch):
-    _patch_conversion(monkeypatch)
-    monkeypatch.setenv("MODELOPT_ATTN_QUANT_OFF", "1")
-    attention = _attention()
-    state = _worker(nn.ModuleDict({"attn": attention}))
-
-    worker_module._install_attention(state, quantize=True)
-
-    converted = state.model_runner.model["attn"]
-    for name in ("q_bmm_quantizer", "k_bmm_quantizer", "p_bmm_quantizer", "v_bmm_quantizer"):
-        assert not getattr(converted, name).is_enabled
-    assert converted._query_quant_in_kernel is False
-    assert converted._value_quant_in_kernel is False
-    assert converted._modelopt_force_kernel is True
-    assert type(converted.impl) is ModelOptSparseAttentionImpl
-    assert converted.impl.quant_kw == {
-        "p_qdq": None,
-        "p_qdq_amax": 1.0,
-        "v_qdq": None,
-        "v_qdq_amax": None,
-    }
-
-
 def test_validation_of_all_layouts_precedes_mutation(monkeypatch):
     _patch_conversion(monkeypatch)
     good, bad = _attention(), _attention(AttentionType.ENCODER)