diff --git a/tests/models/inkling/test_fa4_rel_attention.py b/tests/models/inkling/test_fa4_rel_attention.py index bdd1fbf944..089e7b5585 100644 --- a/tests/models/inkling/test_fa4_rel_attention.py +++ b/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 @@ 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, diff --git a/tests/v1/attention/test_mla_prefill_quant_output.py b/tests/v1/attention/test_mla_prefill_quant_output.py index d7659485aa..220f6a198d 100644 --- a/tests/v1/attention/test_mla_prefill_quant_output.py +++ b/tests/v1/attention/test_mla_prefill_quant_output.py @@ -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.""" diff --git a/vllm/model_executor/warmup/fa4_cutedsl_warmup.py b/vllm/model_executor/warmup/fa4_cutedsl_warmup.py deleted file mode 100644 index 5470da5034..0000000000 --- a/vllm/model_executor/warmup/fa4_cutedsl_warmup.py +++ /dev/null @@ -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) diff --git a/vllm/model_executor/warmup/kernel_warmup.py b/vllm/model_executor/warmup/kernel_warmup.py index e59400b570..d7a73ef2f6 100644 --- a/vllm/model_executor/warmup/kernel_warmup.py +++ b/vllm/model_executor/warmup/kernel_warmup.py @@ -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) diff --git a/vllm/models/inkling/nvidia/attention.py b/vllm/models/inkling/nvidia/attention.py index b18fa81f7e..cc78b6e9ba 100644 --- a/vllm/models/inkling/nvidia/attention.py +++ b/vllm/models/inkling/nvidia/attention.py @@ -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, diff --git a/vllm/models/inkling/nvidia/ops/__init__.py b/vllm/models/inkling/nvidia/ops/__init__.py index bb7f066ddb..fe6bc87e30 100644 --- a/vllm/models/inkling/nvidia/ops/__init__.py +++ b/vllm/models/inkling/nvidia/ops/__init__.py @@ -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, diff --git a/vllm/models/inkling/nvidia/ops/fa4_rel_attention.py b/vllm/models/inkling/nvidia/ops/fa4_rel_attention.py index 4385786135..87fcbdbee6 100644 --- a/vllm/models/inkling/nvidia/ops/fa4_rel_attention.py +++ b/vllm/models/inkling/nvidia/ops/fa4_rel_attention.py @@ -433,4 +433,4 @@ class InklingFA4RelAttentionKernel( ) -INKLING_FA4_REL_ATTENTION_KERNEL = InklingFA4RelAttentionKernel() +_INKLING_FA4_REL_ATTENTION_KERNEL = InklingFA4RelAttentionKernel() diff --git a/vllm/v1/attention/backends/mla/prefill/flash_attn.py b/vllm/v1/attention/backends/mla/prefill/flash_attn.py index eace6f5a10..3e0d544195 100644 --- a/vllm/v1/attention/backends/mla/prefill/flash_attn.py +++ b/vllm/v1/attention/backends/mla/prefill/flash_attn.py @@ -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()