[3/N][warmup][DSv4] Migrate FA4 MLA and shared CuTeDSL kernels (#53565)
Signed-off-by: LopezCastroRoberto <rocastro@redhat.com> Co-authored-by: Codex <codex@openai.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Roberto L. Castro committed
Sep 7, 2026 at 15:36 UTC
42801b3a6b3bb92397a04117449f4699b09db1df
8 files changed
+52
-67
tests/models/inkling/test_fa4_rel_attention.py
+3
-3
@@ -2,7 +2,7 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
"""Correctness test for the Inkling FA4 relative-attention score-mod kernel.
4
5
-Checks ``INKLING_FA4_REL_ATTENTION_KERNEL`` against a pure-PyTorch reference
5
+Checks ``_INKLING_FA4_REL_ATTENTION_KERNEL`` against a pure-PyTorch reference
6
that implements the relative bias exactly as documented in the Inkling architecture
7
guide::
8
@@ -23,7 +23,7 @@ from vllm.models.inkling.nvidia.attention import (
23
compute_log_scaling_tau,
24
)
25
from vllm.models.inkling.nvidia.ops.fa4_rel_attention import (
26
- INKLING_FA4_REL_ATTENTION_KERNEL,
26
+ _INKLING_FA4_REL_ATTENTION_KERNEL,
27
_use_sheared_bias,
28
inkling_fa4_num_splits,
29
)
@@ -279,7 +279,7 @@ def _run_case(seq_lens, num_heads, num_kv_heads, rel_extent, window_left, seed=0
279
num_kv_heads=num_kv_heads,
280
max_kv_len=max(kv_lens),
281
)
282
- out = INKLING_FA4_REL_ATTENTION_KERNEL(
282
+ out = _INKLING_FA4_REL_ATTENTION_KERNEL(
283
q,
284
key_cache,
285
value_cache,
tests/v1/attention/test_mla_prefill_quant_output.py
+33
@@ -10,6 +10,7 @@ Covers two things:
10
+ standalone static-FP8-quant path it replaces (GPU-only, SM100/SM110).
11
"""
12
13
+from types import SimpleNamespace
14
from unittest.mock import patch
15
16
import pytest
@@ -99,6 +100,38 @@ def test_flash_attn_supports_quant_output_unknown_device():
100
assert backend.supports_quant_output(kFp8StaticTensorSym) is False
101
102
103
+@pytest.mark.parametrize(
104
+ ("version", "enable_jit_warmup", "expected_calls"),
105
+ [(4, True, 1), (4, False, 0), (3, True, 0)],
106
+)
107
+def test_flash_attn_registers_warmup_only_for_fa4(
108
+ version: int,
109
+ enable_jit_warmup: bool,
110
+ expected_calls: int,
111
+):
112
+ vllm_config = SimpleNamespace(
113
+ kernel_config=SimpleNamespace(enable_jit_warmup=enable_jit_warmup)
114
+ )
115
+ with (
116
+ patch(f"{_FA_MODULE}.flash_attn_varlen_func"),
117
+ patch(f"{_FA_MODULE}.get_flash_attn_version", return_value=version),
118
+ patch(f"{_FA_MODULE}._FA4_MLA_PREFILL_KERNEL.register_warmup") as register,
119
+ patch(f"{_FA_MODULE}.current_platform") as platform,
120
+ ):
121
+ platform.get_device_capability.return_value = DeviceCapability(10, 0)
122
+ FlashAttnPrefillBackend(
123
+ num_heads=16,
124
+ scale=1.0,
125
+ kv_lora_rank=512,
126
+ qk_nope_head_dim=128,
127
+ qk_rope_head_dim=64,
128
+ v_head_dim=128,
129
+ vllm_config=vllm_config,
130
+ )
131
+
132
+ assert register.call_count == expected_calls
133
+
134
+
135
def test_flash_attn_prefill_backend_signature_accepts_fused_kwargs():
136
"""run_prefill_new_tokens must accept out/output_scale so the direct
137
(non-**kwargs) call in forward_mha type- and runtime-checks."""
vllm/model_executor/warmup/fa4_cutedsl_warmup.py
deleted
-53
@@ -1,53 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-"""Warm up FA4 CuTeDSL kernels."""
4
-
5
-from __future__ import annotations
6
-
7
-from typing import TYPE_CHECKING
8
-
9
-from vllm.v1.attention.backends.mla.prefill import get_mla_prefill_backend
10
-
11
-if TYPE_CHECKING:
12
- from vllm.v1.worker.gpu_worker import Worker
13
-
14
-
15
-def _warm_fa4_mla_prefill(worker: Worker) -> None:
16
- runner = worker.model_runner
17
- if runner.is_pooling_model:
18
- return
19
-
20
- vllm_config = runner.vllm_config
21
- if not vllm_config.model_config.use_mla:
22
- return
23
-
24
- try:
25
- backend_cls = get_mla_prefill_backend(vllm_config)
26
- except ValueError:
27
- # fall back to top-k MQA prefill path.
28
- return
29
- if backend_cls.get_name() != "FLASH_ATTN":
30
- return
31
-
32
- from vllm.v1.attention.backends.mla.prefill import flash_attn
33
-
34
- flash_attn.FA4_MLA_PREFILL_KERNEL.warmup(vllm_config)
35
-
36
-
37
-def _warm_inkling_fa4_rel_attention(worker: Worker) -> None:
38
- from vllm.models.inkling.configs import InklingMMConfig, InklingModelConfig
39
- from vllm.models.inkling.nvidia.ops.fa4_rel_attention import (
40
- INKLING_FA4_REL_ATTENTION_KERNEL,
41
- )
42
-
43
- vllm_config = worker.vllm_config
44
- hf_config = vllm_config.model_config.hf_config
45
- if not isinstance(hf_config, (InklingMMConfig, InklingModelConfig)):
46
- return
47
-
48
- INKLING_FA4_REL_ATTENTION_KERNEL.warmup(vllm_config)
49
-
50
-
51
-def fa4_cutedsl_warmup(worker: Worker) -> None:
52
- _warm_fa4_mla_prefill(worker)
53
- _warm_inkling_fa4_rel_attention(worker)
vllm/model_executor/warmup/kernel_warmup.py
-4
@@ -20,9 +20,6 @@ from vllm.model_executor.warmup.deep_gemm_warmup import deep_gemm_warmup
20
from vllm.model_executor.warmup.deepseek_v4_mhc_warmup import (
21
deepseek_v4_mhc_warmup,
22
)
23
-from vllm.model_executor.warmup.fa4_cutedsl_warmup import (
24
- fa4_cutedsl_warmup,
25
-)
23
from vllm.model_executor.warmup.flashinfer_autotune_cache import (
24
resolve_flashinfer_autotune_file,
25
write_flashinfer_autotune_cache,
@@ -169,7 +166,6 @@ def kernel_warmup(worker: "Worker", *, process_local_only: bool = False):
166
# Run next so input-prep kernels JIT against pristine runner state.
167
if worker.vllm_config.kernel_config.enable_jit_warmup:
168
kimi_k3_triton_warmup(worker)
172
- fa4_cutedsl_warmup(worker)
169
spec_decode_rejection_warmup(worker)
170
qwen4_exp_qsa_triton_warmup(worker)
171
vllm/models/inkling/nvidia/attention.py
+5
-2
@@ -35,7 +35,7 @@ from vllm.v1.kv_cache_interface import (
35
from ..configs import InklingModelConfig
36
from .layernorm import InklingRMSNorm
37
from .ops.fa4_rel_attention import (
38
- INKLING_FA4_REL_ATTENTION_KERNEL,
38
+ _INKLING_FA4_REL_ATTENTION_KERNEL,
39
bucket_max_seqlen_q,
40
inkling_fa4_num_splits,
41
)
@@ -174,6 +174,9 @@ class InklingAttention(nn.Module, AttentionLayerBase):
174
compilation_config.static_forward_context[prefix] = self
175
self.kv_cache = torch.tensor([]) # replaced by bind_kv_cache
176
177
+ if vllm_config.kernel_config.enable_jit_warmup:
178
+ _INKLING_FA4_REL_ATTENTION_KERNEL.register_warmup()
179
+
180
def get_attn_backend(self) -> type[AttentionBackend]:
181
return FlashAttentionBackend
182
@@ -289,7 +292,7 @@ class InklingAttention(nn.Module, AttentionLayerBase):
292
num_kv_heads=self.num_kv_heads,
293
max_kv_len=self._max_kv_len,
294
)
292
- INKLING_FA4_REL_ATTENTION_KERNEL(
295
+ _INKLING_FA4_REL_ATTENTION_KERNEL(
296
q[:nt],
297
key_cache,
298
value_cache,
vllm/models/inkling/nvidia/ops/__init__.py
+2
-2
@@ -15,12 +15,12 @@ from .sconv import fused_sconv
15
_LAZY_EXPORTS = {
16
"silu_and_mul_triton": "silu_and_mul",
17
"sink_silu_mul_epilogue": "silu_and_mul",
18
- "INKLING_FA4_REL_ATTENTION_KERNEL": "fa4_rel_attention",
18
+ "_INKLING_FA4_REL_ATTENTION_KERNEL": "fa4_rel_attention",
19
}
20
21
if TYPE_CHECKING:
22
from .fa4_rel_attention import ( # noqa: F401
23
- INKLING_FA4_REL_ATTENTION_KERNEL,
23
+ _INKLING_FA4_REL_ATTENTION_KERNEL,
24
)
25
from .silu_and_mul import ( # noqa: F401
26
silu_and_mul_triton,
vllm/models/inkling/nvidia/ops/fa4_rel_attention.py
+1
-1
@@ -433,4 +433,4 @@ class InklingFA4RelAttentionKernel(
433
)
434
435
436
-INKLING_FA4_REL_ATTENTION_KERNEL = InklingFA4RelAttentionKernel()
436
+_INKLING_FA4_REL_ATTENTION_KERNEL = InklingFA4RelAttentionKernel()
vllm/v1/attention/backends/mla/prefill/flash_attn.py
+8
-2
@@ -365,6 +365,12 @@ class FlashAttnPrefillBackend(MLAPrefillBackend):
365
flash_attn_varlen_func, fa_version=self.vllm_flash_attn_version
366
)
367
368
+ if (
369
+ vllm_config.kernel_config.enable_jit_warmup
370
+ and self.vllm_flash_attn_version == 4
371
+ ):
372
+ _FA4_MLA_PREFILL_KERNEL.register_warmup()
373
+
374
# Determine if we need to pad V
375
# For MLA the v head dim is smaller than qk head dim so we pad out
376
# v with 0s to match the qk head dim for attention backends that do
@@ -422,7 +428,7 @@ class FlashAttnPrefillBackend(MLAPrefillBackend):
428
if envs.VLLM_BATCH_INVARIANT:
429
kwargs["num_splits"] = 1
430
425
- attn_out = FA4_MLA_PREFILL_KERNEL(
431
+ attn_out = _FA4_MLA_PREFILL_KERNEL(
432
q=q,
433
k=k,
434
v=maybe_padded_v,
@@ -494,4 +500,4 @@ class FlashAttnPrefillBackend(MLAPrefillBackend):
500
)
501
502
497
-FA4_MLA_PREFILL_KERNEL = FA4MLAPrefillKernel()
503
+_FA4_MLA_PREFILL_KERNEL = FA4MLAPrefillKernel()