Harden vLLM attention worker per review: split mixed-batch FA decode/prefill, share sparse layer validation, disable sparse cascade, require vLLM>=0.15, reject sm<8.9

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

Author: kaix-nv

examples/vllm_serve/sparse_attn_worker.py

@@ -88,8 +88,13 @@ def _load_quant_api(vllm_version: str):
     import torch
     from packaging import version
 
-    if version.parse(vllm_version) < version.parse("0.14.0"):
-        raise RuntimeError("The compact NVFP4 attention worker requires vLLM >= 0.14.0")
+    # vLLM >= 0.15 is required: the ModelOpt kernel bypasses the native
+    # FlashAttentionImpl.forward() and relies on the external do_kv_cache_update
+    # contract to write the current step's K/V before reading the paged cache.
+    # That contract only exists from 0.15; on 0.14 reshape_and_cache_flash ran
+    # inside the native forward we skip, so the cache could be read stale.
+    if version.parse(vllm_version) < version.parse("0.15.0"):
+        raise RuntimeError("The compact NVFP4 attention worker requires vLLM >= 0.15.0")
 
     from vllm.config import compilation
     from vllm.v1.attention import backend
@@ -146,12 +151,40 @@ def _global_errors(worker, api) -> list[str]:
     return errors
 
 
-def _quant_layer_errors(module, api) -> list[str]:
+def _device_capability_error(device) -> str | None:
+    """Reject GPUs below sm_89: the NVFP4 P/V helpers cast block scales through
+    ``tl.float8e4nv``, which requires CUDA compute capability >= 8.9 (Ada/Hopper+).
+    Validate here so unsupported devices fail before any module mutation rather
+    than late during Triton compilation."""
+    import torch
+
+    if device is None or getattr(device, "type", None) != "cuda":
+        return None
+    major, minor = torch.cuda.get_device_capability(device)
+    if (major, minor) < (8, 9):
+        return (
+            f"NVFP4 attention requires CUDA compute capability >= 8.9 (Ada/Hopper+); "
+            f"got sm_{major}{minor}"
+        )
+    return None
+
+
+def _layer_errors(module) -> list[str]:
+    """Per-layer attention semantics the ModelOpt kernel does not implement.
+
+    Shared by the quant and sparse plans: the sparse Triton path makes the same
+    assumptions (regular decoder self-attention, no sliding window / ALiBi /
+    softcap / sinks / fp8 KV / cross-layer KV sharing / mismatched head dims), so
+    accepting any of them would silently change model output.
+    """
     impl = getattr(module, "impl", None)
     errors = []
-    if type(module) is not api.plugin.vllm_attention.Attention:
+    if type(module) is not VLLMAttention:
         errors.append(f"layout {type(module).__name__} is not regular decoder self-attention")
-    if getattr(module, "attn_type", None) != api.backend.AttentionType.DECODER:
+    # vLLM's AttentionType.DECODER is the string constant "decoder"; compare to it
+    # directly so the shared sparse-only path stays independent of quant-only vLLM
+    # imports (vllm.v1.attention.backend).
+    if getattr(module, "attn_type", None) != "decoder":
         errors.append("attn_type must be DECODER")
     head_size = getattr(module, "head_size", None)
     if not isinstance(head_size, int) or head_size % 16:
@@ -229,9 +262,15 @@ def _sparse_plans(worker):
         sparse_kw = _sparse_kwargs(name, sparse_cfg)
         if not sparse_kw:
             continue
+        # Validate the full attention semantics before installing a sparse-only
+        # adapter: the sparse kernel shares the quant path's assumptions, so an
+        # unsupported layer must fail here rather than silently change output.
+        reasons = _layer_errors(module)
         new_impl, error = _select_new_impl(module)
         if error:
-            errors.append(f"{name or '<root>'}: {error}")
+            reasons.append(error)
+        if reasons:
+            errors.extend(f"{name or '<root>'}: {reason}" for reason in reasons)
         else:
             plans.append(_AttentionPlan(module, new_impl, sparse_kw, None, None))
     _raise_unsupported(errors, "sparse attention")
@@ -252,7 +291,7 @@ def _quant_plans(worker):
         if not isinstance(module, api.plugin._ATTENTION_TYPES):
             continue
         attention_count += 1
-        reasons = _quant_layer_errors(module, api)
+        reasons = _layer_errors(module)
         # Prefer the model compute dtype (fp16/bf16); _get_device_dtype's buffer scan
         # can otherwise report fp32 from the attention module's scale buffers.
         device, dtype = api.plugin._get_device_dtype(module)
@@ -262,6 +301,8 @@ def _quant_plans(worker):
             reasons.append("device/dtype could not be resolved")
         elif dtype not in (api.torch.float16, api.torch.bfloat16):
             reasons.append(f"resolved dtype {dtype} must be fp16 or bf16")
+        if capability_error := _device_capability_error(device):
+            reasons.append(capability_error)
         sparse_kw = _sparse_kwargs(name, sparse_cfg)
         if graph_error := _sparse_graph_error(sparse_kw, mode, api):
             reasons.append(graph_error)
@@ -278,10 +319,14 @@ def _quant_plans(worker):
     return tuple(plans)
 
 
-def _install_sparse_plans(plans) -> None:
+def _install_sparse_plans(worker, plans) -> None:
     for plan in plans:
         plan.new_impl.sparse_kw = plan.sparse_kw
         plan.module.impl = plan.new_impl
+    # Disable cascade attention: vLLM may auto-select it on shared-prefix batches,
+    # but the ModelOpt sparse kernel does not implement the cascade path. Mirror
+    # the quant install so an installed sparse layer never hits it at runtime.
+    worker.model_runner.cascade_attn_enabled = False
     installed = dict(Counter(type(plan.new_impl).__name__ for plan in plans))
     print(
         f"[ModelOpt] Sparse attention: replaced impl on {len(plans)} attention layers: {installed}"
@@ -320,7 +365,7 @@ def _install_attention(worker, *, quantize: bool) -> None:
     else:
         plans = _sparse_plans(worker)
         if plans is not None:
-            _install_sparse_plans(plans)
+            _install_sparse_plans(worker, plans)
 
 
 class _ModelOptAttentionWorker(BaseWorker):

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

@@ -200,9 +200,12 @@ def _resolve_forward(
     transform_active = _should_run_modelopt_kernel(getattr(impl, "sparse_kw", None), quant_active)
 
     if getattr(attn_metadata, "use_cascade", False):
-        if transform_active:
+        # Cascade is unimplemented by the ModelOpt kernel. Quantization must not be
+        # silently dropped (it would change numerics), so reject it; a sparse-only
+        # transform is numerically safe to delegate to the native dense path.
+        if transform_active and quant_active:
             raise NotImplementedError(
-                "vLLM cascade attention is incompatible with an active ModelOpt attention transform"
+                "vLLM cascade attention is incompatible with active ModelOpt attention quantization"
             )
         return None
 
@@ -350,6 +353,83 @@ def _forward_modelopt(
     return output
 
 
+def _dispatch_modelopt(
+    impl,
+    *,
+    query: torch.Tensor,
+    block_table: torch.Tensor,
+    seq_lens: torch.Tensor,
+    cu_seqlens_q: torch.Tensor,
+    num_actual_tokens: int,
+    max_query_len: int,
+    output: torch.Tensor,
+    num_decodes: int,
+    num_prefills: int,
+    num_decode_tokens: int,
+    num_prefill_tokens: int,
+    **common_kw,
+) -> torch.Tensor:
+    """Run the ModelOpt path, splitting mixed decode+prefill batches by phase.
+
+    NVFP4 P-QDQ is schedule-sensitive by design, so a decode result must not
+    depend on whether a prefill request is co-scheduled. When a batch mixes
+    ``q_len==1`` decode rows with ``q_len>1`` (chunked-)prefill rows,
+    ``max_query_len > 1`` and the whole batch would otherwise take the prefill
+    skip-softmax path. Split so each phase runs its own schedule -- decode rows
+    always take the fixed decode path. Both the FlashAttention and FlashInfer
+    adapters share this dispatch.
+    """
+    if not (num_decodes and num_prefills):
+        return _forward_modelopt(
+            impl,
+            query=query,
+            block_table=block_table,
+            seq_lens=seq_lens,
+            cu_seqlens_q=cu_seqlens_q,
+            num_actual_tokens=num_actual_tokens,
+            max_query_len=max_query_len,
+            output=output,
+            **common_kw,
+        )
+
+    if num_decode_tokens % num_decodes:
+        raise NotImplementedError("Non-uniform mixed decode is unsupported")
+    if num_decode_tokens + num_prefill_tokens != num_actual_tokens:
+        raise ValueError("Mixed-batch token counts do not match common metadata")
+
+    # Sparse-only launches may have an inactive phase (for example N:M sparsity
+    # is prefill-only). Compute the native result once, then overwrite each
+    # active phase with its ModelOpt result.
+    if not common_kw.get("quant_active", False):
+        common_kw["dense_fallback"]()
+
+    _forward_modelopt(
+        impl,
+        query=query[:num_decode_tokens],
+        block_table=block_table[:num_decodes],
+        seq_lens=seq_lens[:num_decodes],
+        cu_seqlens_q=cu_seqlens_q[: num_decodes + 1],
+        num_actual_tokens=num_decode_tokens,
+        max_query_len=num_decode_tokens // num_decodes,
+        output=output[:num_decode_tokens],
+        **common_kw,
+    )
+    prefill_start = num_decode_tokens
+    prefill_cu_seqlens_q = cu_seqlens_q[num_decodes:] - cu_seqlens_q[num_decodes]
+    _forward_modelopt(
+        impl,
+        query=query[prefill_start : prefill_start + num_prefill_tokens],
+        block_table=block_table[num_decodes : num_decodes + num_prefills],
+        seq_lens=seq_lens[num_decodes : num_decodes + num_prefills],
+        cu_seqlens_q=prefill_cu_seqlens_q,
+        num_actual_tokens=num_prefill_tokens,
+        max_query_len=max_query_len,
+        output=output[prefill_start : prefill_start + num_prefill_tokens],
+        **common_kw,
+    )
+    return output
+
+
 class ModelOptSparseAttentionImpl(FlashAttentionImpl):
     """FlashAttention adapter for the compact ModelOpt Triton path."""
 
@@ -396,18 +476,25 @@ def forward(
         if attn_metadata is None:
             return output.fill_(0)
 
+        native_result = None
+
         def native_forward():
-            return self._forward_vllm_flash_attn(
-                layer,
-                query,
-                key,
-                value,
-                kv_cache,
-                attn_metadata,
-                output,
-                output_scale,
-                output_block_scale,
-            )
+            # Memoized: a split mixed batch may request the native dense result
+            # for an inactive phase after it was already computed for the batch.
+            nonlocal native_result
+            if native_result is None:
+                native_result = self._forward_vllm_flash_attn(
+                    layer,
+                    query,
+                    key,
+                    value,
+                    kv_cache,
+                    attn_metadata,
+                    output,
+                    output_scale,
+                    output_block_scale,
+                )
+            return native_result
 
         resolved = _resolve_forward(
             self,
@@ -421,26 +508,35 @@ def native_forward():
 
         key_cache, value_cache = kv_cache.unbind(0)
         is_decode_only = attn_metadata.max_query_len <= 1
-        return _forward_modelopt(
+        common_kw = {
+            "layer": layer,
+            "key_cache": key_cache,
+            "value_cache": value_cache,
+            "max_seq_len": attn_metadata.max_seq_len,
+            "is_causal": getattr(attn_metadata, "causal", not is_decode_only),
+            "p_qdq": resolved.p_qdq,
+            "p_qdq_amax": resolved.p_qdq_amax,
+            "v_qdq": resolved.v_qdq,
+            "v_qdq_amax": resolved.v_qdq_amax,
+            "quant_active": resolved.quant_active,
+            "dense_fallback": native_forward,
+        }
+        # Split mixed decode+prefill batches so decode rows never fall into the
+        # schedule-sensitive prefill skip-softmax path (see _dispatch_modelopt).
+        return _dispatch_modelopt(
             self,
-            layer=layer,
             query=query,
-            key_cache=key_cache,
-            value_cache=value_cache,
             block_table=attn_metadata.block_table,
             seq_lens=attn_metadata.seq_lens,
             cu_seqlens_q=attn_metadata.query_start_loc,
             num_actual_tokens=attn_metadata.num_actual_tokens,
             max_query_len=attn_metadata.max_query_len,
-            max_seq_len=attn_metadata.max_seq_len,
-            is_causal=getattr(attn_metadata, "causal", not is_decode_only),
             output=output,
-            p_qdq=resolved.p_qdq,
-            p_qdq_amax=resolved.p_qdq_amax,
-            v_qdq=resolved.v_qdq,
-            v_qdq_amax=resolved.v_qdq_amax,
-            quant_active=resolved.quant_active,
-            dense_fallback=native_forward,
+            num_decodes=getattr(attn_metadata, "num_decodes", 0),
+            num_prefills=getattr(attn_metadata, "num_prefills", 0),
+            num_decode_tokens=getattr(attn_metadata, "num_decode_tokens", 0),
+            num_prefill_tokens=getattr(attn_metadata, "num_prefill_tokens", 0),
+            **common_kw,
         )
 
 
@@ -616,62 +712,21 @@ def prepare_modelopt():
         "dense_fallback": dense_fallback,
         "prepare_modelopt": prepare_modelopt,
     }
-    num_decodes = getattr(attn_metadata, "num_decodes", 0)
-    num_prefills = getattr(attn_metadata, "num_prefills", 0)
-    if not (num_decodes and num_prefills):
-        return _forward_modelopt(
-            impl,
-            query=query,
-            block_table=attn_metadata._modelopt_block_table,
-            seq_lens=attn_metadata._modelopt_seq_lens,
-            cu_seqlens_q=attn_metadata._modelopt_query_start_loc,
-            num_actual_tokens=attn_metadata._modelopt_num_actual_tokens,
-            max_query_len=max_query_len,
-            output=output,
-            **common_kw,
-        )
-
-    num_decode_tokens = attn_metadata.num_decode_tokens
-    num_prefill_tokens = attn_metadata.num_prefill_tokens
-    if num_decode_tokens % num_decodes:
-        raise NotImplementedError("Non-uniform mixed FlashInfer decode is unsupported")
-    if num_decode_tokens + num_prefill_tokens != attn_metadata._modelopt_num_actual_tokens:
-        raise ValueError("FlashInfer mixed-batch token counts do not match common metadata")
-
-    # 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 not resolved.quant_active:
-        dense_fallback()
-
-    block_table = attn_metadata._modelopt_block_table
-    seq_lens = attn_metadata._modelopt_seq_lens
-    cu_seqlens_q = attn_metadata._modelopt_query_start_loc
-    _forward_modelopt(
+    return _dispatch_modelopt(
         impl,
-        query=query[:num_decode_tokens],
-        block_table=block_table[:num_decodes],
-        seq_lens=seq_lens[:num_decodes],
-        cu_seqlens_q=cu_seqlens_q[: num_decodes + 1],
-        num_actual_tokens=num_decode_tokens,
-        max_query_len=num_decode_tokens // num_decodes,
-        output=output[:num_decode_tokens],
-        **common_kw,
-    )
-    prefill_start = num_decode_tokens
-    prefill_cu_seqlens_q = cu_seqlens_q[num_decodes:] - cu_seqlens_q[num_decodes]
-    _forward_modelopt(
-        impl,
-        query=query[prefill_start : prefill_start + num_prefill_tokens],
-        block_table=block_table[num_decodes : num_decodes + num_prefills],
-        seq_lens=seq_lens[num_decodes : num_decodes + num_prefills],
-        cu_seqlens_q=prefill_cu_seqlens_q,
-        num_actual_tokens=num_prefill_tokens,
+        query=query,
+        block_table=attn_metadata._modelopt_block_table,
+        seq_lens=attn_metadata._modelopt_seq_lens,
+        cu_seqlens_q=attn_metadata._modelopt_query_start_loc,
+        num_actual_tokens=attn_metadata._modelopt_num_actual_tokens,
         max_query_len=max_query_len,
-        output=output[prefill_start : prefill_start + num_prefill_tokens],
+        output=output,
+        num_decodes=getattr(attn_metadata, "num_decodes", 0),
+        num_prefills=getattr(attn_metadata, "num_prefills", 0),
+        num_decode_tokens=getattr(attn_metadata, "num_decode_tokens", 0),
+        num_prefill_tokens=getattr(attn_metadata, "num_prefill_tokens", 0),
         **common_kw,
     )
-    return output
 
 
 def get_flashinfer_sparse_impl_cls() -> type:

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

@@ -56,7 +56,7 @@ def test_quant_policy_rejects_old_vllm(monkeypatch):
     worker_module._load_quant_api.cache_clear()
 
     try:
-        with pytest.raises(RuntimeError, match=r"vLLM >= 0\.14\.0"):
+        with pytest.raises(RuntimeError, match=r"vLLM >= 0\.15\.0"):
             worker_module._install_attention(SimpleNamespace(), quantize=True)
     finally:
         worker_module._load_quant_api.cache_clear()