iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[3/N][warmup][DSv4] Migrate FA4 MLA and shared CuTeDSL kernels (#53565)

Browse Source 
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>
This commit is contained in:
Roberto L. Castro
2026-09-07 09:36:35 -04:00
committed by GitHub
co-authored by Codex mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
parent 1713b9866a
commit 42801b3a6b
8 changed files with 52 additions and 67 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/models/inkling/test_fa4_rel_attention.py
+3 -3
View File
@@ -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 @@ from vllm.models.inkling.nvidia.attention import (
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
+33
View File
@@ -10,6 +10,7 @@ Covers two things:
+ 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/fa4_cutedsl_warmup.py
-53
View File
@@ -1,53 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Warm up FA4 CuTeDSL kernels."""
from __future__ import annotations
from typing import TYPE_CHECKING
from vllm.v1.attention.backends.mla.prefill import get_mla_prefill_backend
if TYPE_CHECKING:
from vllm.v1.worker.gpu_worker import Worker
def _warm_fa4_mla_prefill(worker: Worker) -> None:
runner = worker.model_runner
if runner.is_pooling_model:
return
vllm_config = runner.vllm_config
if not vllm_config.model_config.use_mla:
return
try:
backend_cls = get_mla_prefill_backend(vllm_config)
except ValueError:
# fall back to top-k MQA prefill path.
return
if backend_cls.get_name() != "FLASH_ATTN":
return
from vllm.v1.attention.backends.mla.prefill import flash_attn
flash_attn.FA4_MLA_PREFILL_KERNEL.warmup(vllm_config)
def _warm_inkling_fa4_rel_attention(worker: Worker) -> None:
from vllm.models.inkling.configs import InklingMMConfig, InklingModelConfig
from vllm.models.inkling.nvidia.ops.fa4_rel_attention import (
INKLING_FA4_REL_ATTENTION_KERNEL,
)
vllm_config = worker.vllm_config
hf_config = vllm_config.model_config.hf_config
if not isinstance(hf_config, (InklingMMConfig, InklingModelConfig)):
return
INKLING_FA4_REL_ATTENTION_KERNEL.warmup(vllm_config)
def fa4_cutedsl_warmup(worker: Worker) -> None:
_warm_fa4_mla_prefill(worker)
_warm_inkling_fa4_rel_attention(worker)
vllm/model_executor/warmup/kernel_warmup.py
-4
View File
@@ -20,9 +20,6 @@ from vllm.model_executor.warmup.deep_gemm_warmup import deep_gemm_warmup
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
+5 -2
View File
@@ -35,7 +35,7 @@ from vllm.v1.kv_cache_interface import (
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 @@ class InklingAttention(nn.Module, AttentionLayerBase):
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 @@ class InklingAttention(nn.Module, AttentionLayerBase):
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
+2 -2
View File
@@ -15,12 +15,12 @@ from .sconv import fused_sconv
_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
+1 -1
View File
@@ -433,4 +433,4 @@ class InklingFA4RelAttentionKernel(
)
INKLING_FA4_REL_ATTENTION_KERNEL = InklingFA4RelAttentionKernel()
_INKLING_FA4_REL_ATTENTION_KERNEL = InklingFA4RelAttentionKernel()
vllm/v1/attention/backends/mla/prefill/flash_attn.py
+8 -2
View File
@@ -365,6 +365,12 @@ class FlashAttnPrefillBackend(MLAPrefillBackend):
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 @@ class FlashAttnPrefillBackend(MLAPrefillBackend):
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 @@ class FlashAttnPrefillBackend(MLAPrefillBackend):
)
FA4_MLA_PREFILL_KERNEL = FA4MLAPrefillKernel()
_FA4_MLA_PREFILL_KERNEL = FA4MLAPrefillKernel()