feat(sm12x): default packed prefill on when SM120 FI backend active

wingcomm's external nc=8 + MTP validation (PR vllm-project#41834: GSM8K 0.977 at genuine
max_num_seqs=8, 0 errors, no OOM; the 0.90 that shelved packed does not
reproduce) clears packed prefill for promotion. Flip
VLLM_DEEPSEEK_V4_FLASHINFER_SM120_PREFILL default OFF -> ON so the packed
FlashInfer SM120 sparse-MLA prefill runner is used by default whenever its
backend is active, instead of requiring a second opt-in flag on top of
VLLM_DEEPSEEK_V4_FLASHINFER_SM120_DECODE.

Safety: the prefill-SWA index kernel is launched by the SHARED
DeepseekSparseSWAMetadataBuilder, which previously gated only on the (default-off)
PREFILL flag. With the flag now default-on that gate alone would fire the kernel
on the default FlashMLA/Triton path and fault (cudaErrorLaunchFailure). Add
is_dsv4_sm120_fi_prefill_active() (DECODE opted in + SM12x + FI SM120 kernel
present) and require it in both the builder gate and implicitly via the decode
backend's _forward_prefill (only reached when that backend is selected). Default
FlashMLA path, non-SM12x, kernel-absent builds, and explicit =0 are byte-for-byte
unchanged. Adds a unit test for the gate's AND logic.

Not yet runtime-confirmed on our boxes (GB10 offline; RTX GPU0 wedged +
flashinfer 0.6.11 lacks the kernel) -- staged on a feature branch, unpushed,
pending a box that can serve the SM120 FI backend.

Signed-off-by: jasl <jasl9187@hotmail.com>

Author: jasl

tests/utils_/test_dsv4_sm120_prefill_gate.py

@@ -0,0 +1,62 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Unit tests for the DSv4 SM120 packed-prefill backend-active gate.
+
+``is_dsv4_sm120_fi_prefill_active()`` is the signal the SHARED
+``DeepseekSparseSWAMetadataBuilder`` uses to decide whether to launch the
+prefill-SWA index kernel. Since ``VLLM_DEEPSEEK_V4_FLASHINFER_SM120_PREFILL`` now
+defaults ON, this gate is what keeps the default FlashMLA/Triton path from
+launching that kernel (which faults). It must be True ONLY when all three of
+{DECODE opted in, SM12x device, FI SM120 kernel present} hold.
+"""
+from unittest.mock import patch
+
+import pytest
+
+from vllm.utils import flashinfer
+
+
+@pytest.fixture(autouse=True)
+def _clear_cache():
+    flashinfer.is_dsv4_sm120_fi_prefill_active.cache_clear()
+    yield
+    flashinfer.is_dsv4_sm120_fi_prefill_active.cache_clear()
+
+
+def _run(monkeypatch, decode: bool, family120: bool, kernel: bool) -> bool:
+    # The DECODE env resolves through os.environ via the envs.py getter lambda.
+    monkeypatch.setenv(
+        "VLLM_DEEPSEEK_V4_FLASHINFER_SM120_DECODE", "1" if decode else "0"
+    )
+    with (
+        patch(
+            "vllm.platforms.current_platform.is_device_capability_family",
+            return_value=family120,
+        ),
+        patch.object(
+            flashinfer,
+            "has_flashinfer_trtllm_sparse_mla_dsv4",
+            return_value=kernel,
+        ),
+    ):
+        flashinfer.is_dsv4_sm120_fi_prefill_active.cache_clear()
+        return flashinfer.is_dsv4_sm120_fi_prefill_active()
+
+
+def test_all_three_true_is_active(monkeypatch):
+    assert _run(monkeypatch, decode=True, family120=True, kernel=True) is True
+
+
+@pytest.mark.parametrize(
+    "decode,family120,kernel",
+    [
+        (False, True, True),   # SM120 FI decode backend not opted in (the default)
+        (True, False, True),   # not SM12x (e.g. Hopper)
+        (True, True, False),   # FI SM120 kernel absent (flashinfer < 0.6.13)
+        (False, False, False),
+    ],
+)
+def test_any_condition_false_is_inactive(monkeypatch, decode, family120, kernel):
+    # Safety-critical: DECODE off (the default) -> inactive, so the shared SWA
+    # builder must NOT launch the prefill-SWA kernel on the default FlashMLA path.
+    assert _run(monkeypatch, decode=decode, family120=family120, kernel=kernel) is False

vllm/envs.py

@@ -185,7 +185,12 @@
     VLLM_DEEPSEEK_V4_INDEXED_D512_SPLIT_PREFILL_WARMUP: bool = True
     VLLM_DEEPSEEK_V4_INDEXED_D512_CHUNKED_PREFILL: bool = True
     VLLM_DEEPSEEK_V4_FLASHINFER_SM120_DECODE: bool = False
-    VLLM_DEEPSEEK_V4_FLASHINFER_SM120_PREFILL: bool = False
+    # Default ON: when the SM120 FlashInfer sparse-MLA backend is active (DECODE
+    # opted in + SM12x + kernel present), the packed prefill runner is used by
+    # default. Set =0 to force the FlashMLA indexed-D512 (Triton split/merge)
+    # prefill path. No effect unless that backend is active (see
+    # is_dsv4_sm120_fi_prefill_active).
+    VLLM_DEEPSEEK_V4_FLASHINFER_SM120_PREFILL: bool = True
     VLLM_TRITON_MLA_SPARSE: bool | None = None
     VLLM_TRITON_MLA_SPARSE_TOPK_CHUNK_SIZE: int = 512
     VLLM_TRITON_MLA_SPARSE_QUERY_CHUNK_SIZE: int = 256
@@ -1485,7 +1490,7 @@ def _resolve_rust_frontend_path() -> str | None:
         int(os.getenv("VLLM_DEEPSEEK_V4_FLASHINFER_SM120_DECODE", "0"))
     ),
     "VLLM_DEEPSEEK_V4_FLASHINFER_SM120_PREFILL": lambda: bool(
-        int(os.getenv("VLLM_DEEPSEEK_V4_FLASHINFER_SM120_PREFILL", "0"))
+        int(os.getenv("VLLM_DEEPSEEK_V4_FLASHINFER_SM120_PREFILL", "1"))
     ),
     # Experimental sparse MLA fallback controls.
     # ``VLLM_TRITON_MLA_SPARSE`` unset means auto-select where FlashMLA sparse

vllm/utils/flashinfer.py

@@ -228,6 +228,30 @@ def has_flashinfer_trtllm_sparse_mla_dsv4() -> bool:
     return True
 
 
+@functools.cache
+def is_dsv4_sm120_fi_prefill_active() -> bool:
+    """Return ``True`` iff the DeepSeek-V4 FlashInfer SM120 sparse-MLA backend
+    (which owns the packed-prefill path) is the selected attention backend.
+
+    Mirrors the ``_select_dsv4_attn_cls`` gate (deepseek_v4/nvidia/model.py): that
+    backend is chosen only when SM120 decode is opted in, the device is SM12x, and
+    the FI SM120 sparse-MLA kernel (PR3395, flashinfer >= 0.6.13) is importable.
+    The shared ``DeepseekSparseSWAMetadataBuilder`` reads this to gate the
+    prefill-SWA index kernel that ONLY the packed path consumes -- launching it on
+    the default FlashMLA/Triton path faults (``cudaErrorLaunchFailure``). So this
+    must reflect "the packed prefill backend is active", not merely "the kernel is
+    importable".
+    """
+    import vllm.envs as envs
+    from vllm.platforms import current_platform
+
+    return bool(
+        envs.VLLM_DEEPSEEK_V4_FLASHINFER_SM120_DECODE
+        and current_platform.is_device_capability_family(120)
+        and has_flashinfer_trtllm_sparse_mla_dsv4()
+    )
+
+
 def has_flashinfer_sparse_mla_sm120() -> bool:
     """Return ``True`` if FlashInfer sparse MLA decode support is available."""
     if not has_flashinfer():

vllm/v1/attention/backends/mla/sparse_swa.py

@@ -10,6 +10,7 @@
 from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
 from vllm.platforms import current_platform
 from vllm.triton_utils import tl, triton
+from vllm.utils.flashinfer import is_dsv4_sm120_fi_prefill_active
 from vllm.utils.math_utils import cdiv
 from vllm.v1.attention.backend import (
     AttentionBackend,
@@ -427,13 +428,18 @@ def build(
         # Prefill SWA indices (paged coords; `token_offset` lets the kernel read at
         # absolute prefill positions while writing from index 0) are consumed ONLY by
         # the FlashInfer SM120 sparse-MLA fork path. The stock FlashMLA/Triton prefill
-        # self-computes and never reads them, so gate the launch behind
-        # VLLM_DEEPSEEK_V4_FLASHINFER_SM120_PREFILL (default off). Running it
-        # unconditionally faulted `_compute_swa_indices_and_lens_kernel` over 32k
-        # prefill rows (unclamped block_table address arithmetic on masked-off lanes
-        # -> cudaErrorLaunchFailure under concurrent load).
+        # self-computes and never reads them, so gate the launch behind both the
+        # backend-active check AND VLLM_DEEPSEEK_V4_FLASHINFER_SM120_PREFILL. This
+        # builder is SHARED across DSv4 backends, so the env flag alone is not enough
+        # now that PREFILL defaults on: without the backend-active gate the default
+        # FlashMLA path would launch this kernel and fault. Running it unconditionally
+        # faulted `_compute_swa_indices_and_lens_kernel` over 32k prefill rows
+        # (unclamped block_table address arithmetic on masked-off lanes ->
+        # cudaErrorLaunchFailure under concurrent load).
         want_prefill_swa = (
-            num_prefill_tokens > 0 and envs.VLLM_DEEPSEEK_V4_FLASHINFER_SM120_PREFILL
+            num_prefill_tokens > 0
+            and is_dsv4_sm120_fi_prefill_active()
+            and envs.VLLM_DEEPSEEK_V4_FLASHINFER_SM120_PREFILL
         )
         if want_prefill_swa:
             prefill_swa_indices = self.prefill_swa_indices[:num_prefill_tokens]