[3/N][warmup][DSv4] Migrate FA4 MLA and shared CuTeDSL kernels

tests/models/inkling/test_fa4_rel_attention.py

@@ -2,7 +2,7 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 """Correctness test for the Inkling FA4 relative-attention score-mod kernel.
 
-Checks ``INKLING_FA4_REL_ATTENTION_KERNEL`` against a pure-PyTorch reference
+Checks ``_INKLING_FA4_REL_ATTENTION_KERNEL`` against a pure-PyTorch reference
 that implements the relative bias exactly as documented in the Inkling architecture
 guide::
 
@@ -23,7 +23,7 @@
     compute_log_scaling_tau,
 )
 from vllm.models.inkling.nvidia.ops.fa4_rel_attention import (
-    INKLING_FA4_REL_ATTENTION_KERNEL,
+    _INKLING_FA4_REL_ATTENTION_KERNEL,
     _use_sheared_bias,
     inkling_fa4_num_splits,
 )
@@ -279,7 +279,7 @@ def _run_case(seq_lens, num_heads, num_kv_heads, rel_extent, window_left, seed=0
         num_kv_heads=num_kv_heads,
         max_kv_len=max(kv_lens),
     )
-    out = INKLING_FA4_REL_ATTENTION_KERNEL(
+    out = _INKLING_FA4_REL_ATTENTION_KERNEL(
         q,
         key_cache,
         value_cache,

tests/v1/attention/test_mla_prefill_quant_output.py

@@ -10,6 +10,7 @@
     + standalone static-FP8-quant path it replaces (GPU-only, SM100/SM110).
 """
 
+from types import SimpleNamespace
 from unittest.mock import patch
 
 import pytest
@@ -99,6 +100,38 @@ def test_flash_attn_supports_quant_output_unknown_device():
         assert backend.supports_quant_output(kFp8StaticTensorSym) is False
 
 
+@pytest.mark.parametrize(
+    ("version", "enable_jit_warmup", "expected_calls"),
+    [(4, True, 1), (4, False, 0), (3, True, 0)],
+)
+def test_flash_attn_registers_warmup_only_for_fa4(
+    version: int,
+    enable_jit_warmup: bool,
+    expected_calls: int,
+):
+    vllm_config = SimpleNamespace(
+        kernel_config=SimpleNamespace(enable_jit_warmup=enable_jit_warmup)
+    )
+    with (
+        patch(f"{_FA_MODULE}.flash_attn_varlen_func"),
+        patch(f"{_FA_MODULE}.get_flash_attn_version", return_value=version),
+        patch(f"{_FA_MODULE}._FA4_MLA_PREFILL_KERNEL.register_warmup") as register,
+        patch(f"{_FA_MODULE}.current_platform") as platform,
+    ):
+        platform.get_device_capability.return_value = DeviceCapability(10, 0)
+        FlashAttnPrefillBackend(
+            num_heads=16,
+            scale=1.0,
+            kv_lora_rank=512,
+            qk_nope_head_dim=128,
+            qk_rope_head_dim=64,
+            v_head_dim=128,
+            vllm_config=vllm_config,
+        )
+
+    assert register.call_count == expected_calls
+
+
 def test_flash_attn_prefill_backend_signature_accepts_fused_kwargs():
     """run_prefill_new_tokens must accept out/output_scale so the direct
     (non-**kwargs) call in forward_mha type- and runtime-checks."""

vllm/model_executor/warmup/kernel_warmup.py

@@ -20,9 +20,6 @@
 from vllm.model_executor.warmup.deepseek_v4_mhc_warmup import (
     deepseek_v4_mhc_warmup,
 )
-from vllm.model_executor.warmup.fa4_cutedsl_warmup import (
-    fa4_cutedsl_warmup,
-)
 from vllm.model_executor.warmup.flashinfer_autotune_cache import (
     resolve_flashinfer_autotune_file,
     write_flashinfer_autotune_cache,
@@ -169,7 +166,6 @@ def kernel_warmup(worker: "Worker", *, process_local_only: bool = False):
     # Run next so input-prep kernels JIT against pristine runner state.
     if worker.vllm_config.kernel_config.enable_jit_warmup:
         kimi_k3_triton_warmup(worker)
-        fa4_cutedsl_warmup(worker)
         spec_decode_rejection_warmup(worker)
         qwen4_exp_qsa_triton_warmup(worker)
 

vllm/models/inkling/nvidia/attention.py

@@ -35,7 +35,7 @@
 from ..configs import InklingModelConfig
 from .layernorm import InklingRMSNorm
 from .ops.fa4_rel_attention import (
-    INKLING_FA4_REL_ATTENTION_KERNEL,
+    _INKLING_FA4_REL_ATTENTION_KERNEL,
     bucket_max_seqlen_q,
     inkling_fa4_num_splits,
 )
@@ -174,6 +174,9 @@ def __init__(
         compilation_config.static_forward_context[prefix] = self
         self.kv_cache = torch.tensor([])  # replaced by bind_kv_cache
 
+        if vllm_config.kernel_config.enable_jit_warmup:
+            _INKLING_FA4_REL_ATTENTION_KERNEL.register_warmup()
+
     def get_attn_backend(self) -> type[AttentionBackend]:
         return FlashAttentionBackend
 
@@ -289,7 +292,7 @@ def _attention(
             num_kv_heads=self.num_kv_heads,
             max_kv_len=self._max_kv_len,
         )
-        INKLING_FA4_REL_ATTENTION_KERNEL(
+        _INKLING_FA4_REL_ATTENTION_KERNEL(
             q[:nt],
             key_cache,
             value_cache,

vllm/models/inkling/nvidia/ops/__init__.py

@@ -15,12 +15,12 @@
 _LAZY_EXPORTS = {
     "silu_and_mul_triton": "silu_and_mul",
     "sink_silu_mul_epilogue": "silu_and_mul",
-    "INKLING_FA4_REL_ATTENTION_KERNEL": "fa4_rel_attention",
+    "_INKLING_FA4_REL_ATTENTION_KERNEL": "fa4_rel_attention",
 }
 
 if TYPE_CHECKING:
     from .fa4_rel_attention import (  # noqa: F401
-        INKLING_FA4_REL_ATTENTION_KERNEL,
+        _INKLING_FA4_REL_ATTENTION_KERNEL,
     )
     from .silu_and_mul import (  # noqa: F401
         silu_and_mul_triton,

vllm/models/inkling/nvidia/ops/fa4_rel_attention.py

@@ -433,4 +433,4 @@ def __call__(
         )
 
 
-INKLING_FA4_REL_ATTENTION_KERNEL = InklingFA4RelAttentionKernel()
+_INKLING_FA4_REL_ATTENTION_KERNEL = InklingFA4RelAttentionKernel()

vllm/v1/attention/backends/mla/prefill/flash_attn.py

@@ -365,6 +365,12 @@ def __init__(
                 flash_attn_varlen_func, fa_version=self.vllm_flash_attn_version
             )
 
+        if (
+            vllm_config.kernel_config.enable_jit_warmup
+            and self.vllm_flash_attn_version == 4
+        ):
+            _FA4_MLA_PREFILL_KERNEL.register_warmup()
+
         # Determine if we need to pad V
         # For MLA the v head dim is smaller than qk head dim so we pad out
         # v with 0s to match the qk head dim for attention backends that do
@@ -422,7 +428,7 @@ def _flash_attn_varlen_diff_headdims(
         if envs.VLLM_BATCH_INVARIANT:
             kwargs["num_splits"] = 1
 
-        attn_out = FA4_MLA_PREFILL_KERNEL(
+        attn_out = _FA4_MLA_PREFILL_KERNEL(
             q=q,
             k=k,
             v=maybe_padded_v,
@@ -494,4 +500,4 @@ def run_prefill_context_chunk(
         )
 
 
-FA4_MLA_PREFILL_KERNEL = FA4MLAPrefillKernel()
+_FA4_MLA_PREFILL_KERNEL = FA4MLAPrefillKernel()