[Feat][1/N] CuTeDSL warmup infrastructure, FA4 MLA (#46182)
Signed-off-by: LopezCastroRoberto <rocastro@redhat.com> Signed-off-by: Roberto L. Castro <38211239+LopezCastroRoberto@users.noreply.github.com> Co-authored-by: Lucas Wilkinson <LucasWilkinson@users.noreply.github.com>
Roberto L. Castro committed
Jun 30, 2026 at 21:17 UTC
248d1fbb711b210784ab880593403d665a4731bd
9 files changed
+518
-2
cmake/external_projects/vllm_flash_attn.cmake
+1
-1
@@ -39,7 +39,7 @@ else()
39
FetchContent_Declare(
40
vllm-flash-attn
41
GIT_REPOSITORY https://github.com/vllm-project/flash-attention.git
42
- GIT_TAG b3964b1d8b95d8e8447435668ab169a2700bab65
42
+ GIT_TAG 2c839c33742309ec41e620bf837495ec9926c56e
43
GIT_PROGRESS TRUE
44
# Don't share the vllm-flash-attn build between build types
45
BINARY_DIR ${CMAKE_BINARY_DIR}/vllm-flash-attn
vllm/config/kernel.py
+9
-1
@@ -172,6 +172,9 @@ class KernelConfig:
172
enable_flashinfer_autotune: bool = None # type: ignore[assignment]
173
"""If True, run FlashInfer autotuning during kernel warmup."""
174
175
+ enable_cutedsl_warmup: bool = True
176
+ """If True, run CuTeDSL compile warmup during kernel warmup."""
177
+
178
moe_backend: MoEBackend = "auto"
179
"""Backend for MoE expert computation kernels. Available options:
180
@@ -237,6 +240,7 @@ class KernelConfig:
240
Any future fields that don't affect compilation should be excluded.
241
"""
242
ignored_factors = {
243
+ "enable_cutedsl_warmup",
244
"enable_flashinfer_autotune",
245
"ir_op_priority", # handled separately below
246
}
@@ -244,7 +248,11 @@ class KernelConfig:
248
factors["ir_op_priority"] = self.ir_op_priority.compute_hash()
249
return hash_factors(factors)
250
247
- @field_validator("enable_flashinfer_autotune", mode="wrap")
251
+ @field_validator(
252
+ "enable_flashinfer_autotune",
253
+ "enable_cutedsl_warmup",
254
+ mode="wrap",
255
+ )
256
@classmethod
257
def _skip_none_validation(cls, value: Any, handler: Callable) -> Any:
258
"""Skip validation if the value is `None` when initialization is delayed."""
vllm/model_executor/warmup/cutedsl_warmup.py
new
+113
@@ -0,0 +1,113 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Run registered CuTeDSL warmup compile units."""
4
+
5
+from __future__ import annotations
6
+
7
+import time
8
+import weakref
9
+from collections.abc import Callable, Hashable, Iterable
10
+from dataclasses import dataclass
11
+
12
+import torch
13
+
14
+from vllm.logger import init_logger
15
+from vllm.platforms import current_platform
16
+from vllm.tracing import instrument
17
+
18
+logger = init_logger(__name__)
19
+
20
+CuTeDSLCompileFn = Callable[[], None]
21
+
22
+
23
+@dataclass(frozen=True)
24
+class CuTeDSLCompileUnit:
25
+ name: str
26
+ key: Hashable
27
+ compile: CuTeDSLCompileFn
28
+
29
+
30
+_CUTEDSL_WARMUP_PROVIDERS: weakref.WeakSet[object] = weakref.WeakSet()
31
+
32
+
33
+def register_cutedsl_warmup_provider(provider: object) -> None:
34
+ """Register an object that can expose CuTeDSL warmup compile units."""
35
+ _CUTEDSL_WARMUP_PROVIDERS.add(provider)
36
+
37
+
38
+# Yield compile units from registered providers.
39
+def _iter_cutedsl_warmup_compile_units() -> Iterable[CuTeDSLCompileUnit]:
40
+ for provider in tuple(_CUTEDSL_WARMUP_PROVIDERS):
41
+ get_units = getattr(provider, "get_cutedsl_warmup_compile_units", None)
42
+ if not callable(get_units):
43
+ continue
44
+
45
+ compile_units = get_units()
46
+ if compile_units is None:
47
+ continue
48
+ for unit in compile_units:
49
+ if not isinstance(unit, CuTeDSLCompileUnit):
50
+ raise TypeError(
51
+ "get_cutedsl_warmup_compile_units must return "
52
+ "CuTeDSLCompileUnit objects"
53
+ )
54
+ yield unit
55
+
56
+
57
+# Drop duplicate compile units across providers.
58
+def _collect_unique_compile_units(
59
+ compile_units: Iterable[CuTeDSLCompileUnit],
60
+) -> list[CuTeDSLCompileUnit]:
61
+ seen: set[Hashable] = set()
62
+ unique_compile_units: list[CuTeDSLCompileUnit] = []
63
+
64
+ for unit in compile_units:
65
+ if unit.key in seen:
66
+ continue
67
+
68
+ seen.add(unit.key)
69
+ unique_compile_units.append(unit)
70
+
71
+ return unique_compile_units
72
+
73
+
74
+# Execute compile units under inference mode.
75
+def _compile_cutedsl_warmup_units(
76
+ compile_units: Iterable[CuTeDSLCompileUnit],
77
+) -> int:
78
+ compiled = 0
79
+ with torch.inference_mode():
80
+ for unit in compile_units:
81
+ unit.compile()
82
+ compiled += 1
83
+ torch.accelerator.synchronize()
84
+ return compiled
85
+
86
+
87
+# Run all CuTeDSL warmup before serving.
88
+@instrument(span_name="CuTeDSL warmup")
89
+def cutedsl_warmup() -> None:
90
+ """Run CuTeDSL compile providers before serving."""
91
+ if not current_platform.is_cuda():
92
+ logger.info("Skipping CuTeDSL warmup on non-CUDA platform.")
93
+ return
94
+
95
+ compile_units = _collect_unique_compile_units(_iter_cutedsl_warmup_compile_units())
96
+ if not compile_units:
97
+ logger.info("Skipping CuTeDSL warmup because no compile units were requested.")
98
+ return
99
+
100
+ unit_names = list(dict.fromkeys(unit.name for unit in compile_units))
101
+ logger.info(
102
+ "Warming up CuTeDSL compile_units=%d names=%s.",
103
+ len(compile_units),
104
+ unit_names,
105
+ )
106
+
107
+ start_time = time.perf_counter()
108
+ compiled_count = _compile_cutedsl_warmup_units(compile_units)
109
+ logger.info(
110
+ "CuTeDSL warmup compiled %d units in %.2f s.",
111
+ compiled_count,
112
+ time.perf_counter() - start_time,
113
+ )
vllm/model_executor/warmup/fa4_cutedsl_config.py
new
+204
@@ -0,0 +1,204 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""FA4 MLA prefill CuTeDSL compile warmup config."""
4
+
5
+from __future__ import annotations
6
+
7
+from collections.abc import Hashable, Iterator
8
+from dataclasses import dataclass
9
+from typing import TYPE_CHECKING, Literal
10
+
11
+import torch
12
+
13
+if TYPE_CHECKING:
14
+ from vllm.v1.attention.backends.fa_utils import (
15
+ FlashAttentionCuTeDSLCompileSpec,
16
+ )
17
+
18
+FA4ArchitectureFamily = Literal["sm90", "sm100f", "sm120"]
19
+
20
+FA4_STANDARD_DTYPES = (torch.bfloat16, torch.float16)
21
+
22
+# Current vLLM MLA prefill expands K/V to num_heads before FA4, so this plan
23
+# covers qhead_per_kvhead=1.
24
+# Batch is not a current FA4 MLA-prefill key field. Use b1 for compile-only
25
+# specs because it is the conservative case for Split-KV shape heuristics.
26
+# TODO(roberto): FA4 also has direct-GQA and qv/top-k absorbed-MLA paths, but vLLM
27
+# does not use them in this backend yet; they need a separate
28
+# num_kv_heads/qv/top-k-aware warmup plan if wired in later.
29
+FA4_MLA_PREFILL_COMPILE_BATCH_SIZE = 1
30
+FA4_MLA_PREFILL_Q_TILE = 128
31
+FA4_MLA_PREFILL_K_TILE = 128
32
+FA4_MLA_PREFILL_LONG_K_BLOCKS = 32
33
+FA4_MLA_PREFILL_VERY_LONG_K_BLOCKS = 64
34
+FA4_MLA_PREFILL_CAUSAL_OPTIONS = (False, True)
35
+FA4_MLA_PREFILL_LSE_OPTIONS = (False, True)
36
+
37
+
38
+@dataclass(frozen=True)
39
+class FA4MLAPrefillCompileContext:
40
+ dtype: torch.dtype
41
+ num_heads: int
42
+ qk_head_dim: int
43
+ v_head_dim: int
44
+ kv_nope_head_dim: int
45
+ requires_v_padding: bool
46
+ scale: float
47
+ num_splits: int
48
+ fa_version: int
49
+
50
+ # Return the V head dim FA4 sees.
51
+ @property
52
+ def effective_v_head_dim(self) -> int:
53
+ if self.requires_v_padding:
54
+ return self.qk_head_dim
55
+ return self.v_head_dim
56
+
57
+
58
+@dataclass(frozen=True)
59
+class FA4MLAPrefillCompileRequest:
60
+ """One compile-only FA4 MLA prefill request."""
61
+
62
+ key: Hashable
63
+ compile_spec: FlashAttentionCuTeDSLCompileSpec
64
+
65
+ # Compile this request.
66
+ def compile(self) -> None:
67
+ self.compile_spec.compile()
68
+
69
+
70
+# Yield deduped compile requests.
71
+def iter_fa4_mla_prefill_compile_requests(
72
+ ctx: FA4MLAPrefillCompileContext,
73
+) -> Iterator[FA4MLAPrefillCompileRequest]:
74
+ """Yield compile requests for this fixed MLA backend.
75
+
76
+ FA4 dedupes duplicate atomic kernel selections in its own JIT cache.
77
+ """
78
+ seen: set[Hashable] = set()
79
+ for compile_spec in iter_fa4_mla_prefill_compile_specs(ctx):
80
+ key = compile_spec.request_key()
81
+ if key in seen:
82
+ continue
83
+ seen.add(key)
84
+ yield FA4MLAPrefillCompileRequest(
85
+ key=key,
86
+ compile_spec=compile_spec,
87
+ )
88
+
89
+
90
+# Build compile specs for this setup.
91
+def iter_fa4_mla_prefill_compile_specs(
92
+ ctx: FA4MLAPrefillCompileContext,
93
+) -> Iterator[FlashAttentionCuTeDSLCompileSpec]:
94
+ """Yield compile-only FA4 MLA prefill requests for this fixed setup."""
95
+
96
+ arch_family = _fa4_architecture_family_from_compute_capability(
97
+ *torch.cuda.get_device_capability()
98
+ )
99
+ if not _supports_fa4_mla_prefill(ctx, arch_family):
100
+ return
101
+
102
+ from vllm.v1.attention.backends.fa_utils import (
103
+ FlashAttentionCuTeDSLCompileSpec,
104
+ )
105
+
106
+ batch_size = FA4_MLA_PREFILL_COMPILE_BATCH_SIZE
107
+ v_stride = None
108
+ if not ctx.requires_v_padding:
109
+ v_stride = (
110
+ ctx.num_heads * ctx.kv_nope_head_dim,
111
+ ctx.kv_nope_head_dim,
112
+ 1,
113
+ )
114
+
115
+ for _, max_seqlen_q, max_seqlen_k in _shape_probes_for_context(ctx, arch_family):
116
+ total_q_tokens = batch_size * max_seqlen_q
117
+ total_kv_tokens = batch_size * max_seqlen_k
118
+ for causal in FA4_MLA_PREFILL_CAUSAL_OPTIONS:
119
+ for return_lse in FA4_MLA_PREFILL_LSE_OPTIONS:
120
+ yield FlashAttentionCuTeDSLCompileSpec(
121
+ q_shape=(total_q_tokens, ctx.num_heads, ctx.qk_head_dim),
122
+ k_shape=(total_kv_tokens, ctx.num_heads, ctx.qk_head_dim),
123
+ v_shape=(
124
+ total_kv_tokens,
125
+ ctx.num_heads,
126
+ ctx.effective_v_head_dim,
127
+ ),
128
+ v_stride=v_stride,
129
+ q_dtype=ctx.dtype,
130
+ cu_seqlens_q_shape=(batch_size + 1,),
131
+ cu_seqlens_k_shape=(batch_size + 1,),
132
+ max_seqlen_q=max_seqlen_q,
133
+ max_seqlen_k=max_seqlen_k,
134
+ softmax_scale=ctx.scale,
135
+ causal=causal,
136
+ return_softmax_lse=return_lse,
137
+ num_splits=ctx.num_splits,
138
+ fa_version=ctx.fa_version,
139
+ )
140
+
141
+
142
+# Pick one q/k point per current FA4 MLA-prefill shape regime.
143
+def _shape_probes_for_context(
144
+ ctx: FA4MLAPrefillCompileContext,
145
+ arch_family: FA4ArchitectureFamily,
146
+) -> tuple[tuple[str, int, int], ...]:
147
+ q_stage1_q = 1
148
+ q_stage2_q = FA4_MLA_PREFILL_Q_TILE + 1
149
+ # FA4 never auto-splits when ceil(max_seqlen_k / tile_n) <= 4.
150
+ no_split_k = 4 * FA4_MLA_PREFILL_K_TILE
151
+ long_k = FA4_MLA_PREFILL_LONG_K_BLOCKS * FA4_MLA_PREFILL_K_TILE
152
+ # Diff-head-dim Blackwell Split-KV switches tile_n at 64 K blocks.
153
+ very_long_k = FA4_MLA_PREFILL_VERY_LONG_K_BLOCKS * FA4_MLA_PREFILL_K_TILE
154
+
155
+ base_probes = (
156
+ ("q_stage1", q_stage1_q, FA4_MLA_PREFILL_K_TILE),
157
+ ("q_stage2", q_stage2_q, no_split_k),
158
+ )
159
+ # SM120 currently rejects Split-KV in FA4; num_splits=1 also has no split
160
+ # shape regimes on any architecture.
161
+ if ctx.num_splits == 1 or arch_family == "sm120":
162
+ return base_probes
163
+
164
+ long_k_probes = (
165
+ ("q_stage1_long_k", q_stage1_q, long_k),
166
+ ("q_stage2_long_k", q_stage2_q, long_k),
167
+ )
168
+
169
+ # SM90 does not have the SM100 q_stage or diff-head-dim tile_n=64 branch.
170
+ # Same-dim SM100-family MLA also does not need the very-long-K probe.
171
+ if arch_family == "sm90" or ctx.qk_head_dim == ctx.effective_v_head_dim:
172
+ return (*base_probes, *long_k_probes)
173
+
174
+ very_long_k_probes = (
175
+ ("q_stage1_very_long_k", q_stage1_q, very_long_k),
176
+ ("q_stage2_very_long_k", q_stage2_q, very_long_k),
177
+ )
178
+ return (*base_probes, *long_k_probes, *very_long_k_probes)
179
+
180
+
181
+# Check whether this setup can use FA4 MLA prefill.
182
+def _supports_fa4_mla_prefill(
183
+ ctx: FA4MLAPrefillCompileContext,
184
+ arch_family: FA4ArchitectureFamily,
185
+) -> bool:
186
+ return (
187
+ ctx.dtype in FA4_STANDARD_DTYPES
188
+ and ctx.num_heads > 0
189
+ and (arch_family != "sm120" or ctx.num_splits == 1)
190
+ )
191
+
192
+
193
+# Map CUDA capability to the FA4 arch family used by warmup checks.
194
+def _fa4_architecture_family_from_compute_capability(
195
+ major: int,
196
+ minor: int,
197
+) -> FA4ArchitectureFamily:
198
+ if (major, minor) == (9, 0):
199
+ return "sm90"
200
+ if major == 10:
201
+ return "sm100f"
202
+ if (major, minor) == (12, 0):
203
+ return "sm120"
204
+ raise ValueError(f"FA4 warmup does not know CUDA capability {major}.{minor}")
vllm/model_executor/warmup/kernel_warmup.py
+4
@@ -12,6 +12,7 @@ import torch
12
13
import vllm.envs as envs
14
from vllm.logger import init_logger
15
+from vllm.model_executor.warmup.cutedsl_warmup import cutedsl_warmup
16
from vllm.model_executor.warmup.deep_gemm_warmup import deep_gemm_warmup
17
from vllm.model_executor.warmup.deepseek_v4_mhc_warmup import (
18
deepseek_v4_mhc_warmup,
@@ -125,6 +126,9 @@ def kernel_warmup(worker: "Worker"):
126
create_mixed_batch=True,
127
)
128
129
+ if worker.vllm_config.kernel_config.enable_cutedsl_warmup:
130
+ cutedsl_warmup()
131
+
132
133
def flashinfer_autotune(runner: "GPUModelRunner") -> None:
134
"""
vllm/v1/attention/backends/fa_utils.py
+76
@@ -1,8 +1,11 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
+from dataclasses import dataclass
5
from typing import Any
6
7
+import torch
8
+
9
import vllm.envs as envs
10
from vllm.logger import init_logger
11
from vllm.platforms import current_platform
@@ -18,6 +21,7 @@ _ROCM_FLASH_ATTN_AVAILABLE = False
21
if current_platform.is_cuda():
22
from vllm._custom_ops import reshape_and_cache_flash
23
from vllm.vllm_flash_attn import ( # type: ignore[attr-defined]
24
+ compile_flash_attn_varlen_func_from_specs,
25
flash_attn_varlen_func,
26
get_scheduler_metadata,
27
)
@@ -28,11 +32,14 @@ elif current_platform.is_xpu():
32
33
reshape_and_cache_flash = ops.reshape_and_cache_flash
34
flash_attn_varlen_func = xpu_ops.flash_attn_varlen_func # type: ignore[assignment]
35
+ compile_flash_attn_varlen_func_from_specs = None # type: ignore[assignment]
36
get_scheduler_metadata = xpu_ops.get_scheduler_metadata # type: ignore[assignment]
37
elif current_platform.is_rocm():
38
try:
39
from flash_attn import flash_attn_varlen_func # type: ignore[no-redef]
40
41
+ compile_flash_attn_varlen_func_from_specs = None # type: ignore[assignment]
42
+
43
# Mark that upstream flash-attn is available on ROCm
44
_ROCM_FLASH_ATTN_AVAILABLE = True
45
except ImportError:
@@ -43,6 +50,8 @@ elif current_platform.is_rocm():
50
"to be installed. Please install flash-attn first."
51
)
52
53
+ compile_flash_attn_varlen_func_from_specs = None # type: ignore[assignment]
54
+
55
# ROCm doesn't use scheduler metadata (FA3 feature), provide stub
56
def get_scheduler_metadata(*args: Any, **kwargs: Any) -> None: # type: ignore[misc]
57
return None
@@ -53,6 +62,73 @@ elif current_platform.is_rocm():
62
reshape_and_cache_flash = ops.reshape_and_cache_flash
63
64
65
+@dataclass(frozen=True)
66
+class FlashAttentionCuTeDSLCompileSpec:
67
+ """High-level FA4 compile-only request used by vLLM warmup.
68
+
69
+ This is not the CuTeDSL cache key. FA4 owns the selector that maps these
70
+ serving inputs to the actual compile-static fields: tile sizes, q_stage,
71
+ Split-KV, scheduler choice, layout-presence booleans, dtype/head dims,
72
+ arch, and related fields.
73
+ """
74
+
75
+ q_shape: tuple[int, ...]
76
+ k_shape: tuple[int, ...]
77
+ v_shape: tuple[int, ...]
78
+ q_dtype: torch.dtype
79
+ max_seqlen_q: int
80
+ max_seqlen_k: int
81
+ softmax_scale: float
82
+ causal: bool
83
+ fa_version: int
84
+ v_stride: tuple[int, ...] | None = None
85
+ cu_seqlens_q_shape: tuple[int, ...] | None = None
86
+ cu_seqlens_k_shape: tuple[int, ...] | None = None
87
+ window_size: tuple[int, int] | None = None
88
+ return_softmax_lse: bool = False
89
+ num_splits: int = 0
90
+
91
+ def compile(self) -> None:
92
+ assert compile_flash_attn_varlen_func_from_specs is not None
93
+ window_size = list(self.window_size) if self.window_size is not None else None
94
+ compile_flash_attn_varlen_func_from_specs(
95
+ q_shape=self.q_shape,
96
+ k_shape=self.k_shape,
97
+ v_shape=self.v_shape,
98
+ q_dtype=self.q_dtype,
99
+ v_stride=self.v_stride,
100
+ cu_seqlens_q_shape=self.cu_seqlens_q_shape,
101
+ cu_seqlens_k_shape=self.cu_seqlens_k_shape,
102
+ max_seqlen_q=self.max_seqlen_q,
103
+ max_seqlen_k=self.max_seqlen_k,
104
+ softmax_scale=self.softmax_scale,
105
+ causal=self.causal,
106
+ window_size=window_size,
107
+ return_softmax_lse=self.return_softmax_lse,
108
+ fa_version=self.fa_version,
109
+ num_splits=self.num_splits,
110
+ )
111
+
112
+ def request_key(self) -> tuple[object, ...]:
113
+ return (
114
+ self.q_shape,
115
+ self.k_shape,
116
+ self.v_shape,
117
+ self.q_dtype,
118
+ self.max_seqlen_q,
119
+ self.max_seqlen_k,
120
+ self.softmax_scale,
121
+ self.causal,
122
+ self.fa_version,
123
+ self.v_stride,
124
+ self.cu_seqlens_q_shape,
125
+ self.cu_seqlens_k_shape,
126
+ self.window_size,
127
+ self.return_softmax_lse,
128
+ self.num_splits,
129
+ )
130
+
131
+
132
def get_flash_attn_version(
133
requires_alibi: bool = False,
134
head_size: int | None = None,
vllm/v1/attention/backends/mla/prefill/flash_attn.py
+49
@@ -11,8 +11,17 @@ import vllm.envs as envs
11
from vllm.model_executor.layers.quantization.utils.quant_utils import (
12
kFp8StaticTensorSym,
13
)
14
+from vllm.model_executor.warmup.cutedsl_warmup import (
15
+ CuTeDSLCompileUnit,
16
+ register_cutedsl_warmup_provider,
17
+)
18
+from vllm.model_executor.warmup.fa4_cutedsl_config import (
19
+ FA4MLAPrefillCompileContext,
20
+ iter_fa4_mla_prefill_compile_requests,
21
+)
22
from vllm.platforms import current_platform
23
from vllm.v1.attention.backends.fa_utils import (
24
+ compile_flash_attn_varlen_func_from_specs,
25
get_flash_attn_version,
26
is_flash_attn_varlen_func_available,
27
)
@@ -90,6 +99,46 @@ class FlashAttnPrefillBackend(MLAPrefillBackend):
99
100
# Track whether we're using vllm's FA or upstream (for ROCm)
101
self._is_vllm_fa = current_platform.is_cuda() or current_platform.is_xpu()
102
+ if self.vllm_flash_attn_version == 4:
103
+ register_cutedsl_warmup_provider(self)
104
+
105
+ def get_cutedsl_warmup_compile_units(self) -> tuple[CuTeDSLCompileUnit, ...]:
106
+ if self.vllm_flash_attn_version != 4:
107
+ return ()
108
+ if compile_flash_attn_varlen_func_from_specs is None:
109
+ raise RuntimeError(
110
+ "FA4 compile-only API is unavailable; CuTeDSL warmup does not "
111
+ "fall back to synthetic forward passes."
112
+ )
113
+
114
+ dtype = self.vllm_config.model_config.dtype
115
+ if dtype not in self.supported_dtypes:
116
+ dtype = torch.bfloat16
117
+
118
+ qk_head_dim = self.qk_nope_head_dim + self.qk_rope_head_dim
119
+ ctx = FA4MLAPrefillCompileContext(
120
+ dtype=dtype,
121
+ num_heads=self.num_heads,
122
+ qk_head_dim=qk_head_dim,
123
+ v_head_dim=self.v_head_dim,
124
+ kv_nope_head_dim=self.qk_nope_head_dim + self.v_head_dim,
125
+ requires_v_padding=self.requires_v_padding,
126
+ scale=self.scale,
127
+ num_splits=1 if envs.VLLM_BATCH_INVARIANT else 0,
128
+ fa_version=self.vllm_flash_attn_version,
129
+ )
130
+ compile_requests = tuple(iter_fa4_mla_prefill_compile_requests(ctx))
131
+ if not compile_requests:
132
+ return ()
133
+
134
+ return tuple(
135
+ CuTeDSLCompileUnit(
136
+ name="fa4_mla_prefill",
137
+ key=request.key,
138
+ compile=request.compile,
139
+ )
140
+ for request in compile_requests
141
+ )
142
143
def supports_quant_output(self, quant_key: "QuantKey") -> bool:
144
device_capability = current_platform.get_device_capability()
vllm/vllm_flash_attn/__init__.py
+2
@@ -23,6 +23,7 @@ if os.path.islink(_cute_dir) and "flash_attn" not in sys.modules:
23
from vllm.vllm_flash_attn.flash_attn_interface import ( # noqa: E402
24
FA2_AVAILABLE,
25
FA3_AVAILABLE,
26
+ compile_flash_attn_varlen_func_from_specs,
27
fa_version_unsupported_reason,
28
flash_attn_varlen_func,
29
get_scheduler_metadata,
@@ -36,6 +37,7 @@ if not (FA2_AVAILABLE or FA3_AVAILABLE):
37
)
38
39
__all__ = [
40
+ "compile_flash_attn_varlen_func_from_specs",
41
"fa_version_unsupported_reason",
42
"flash_attn_varlen_func",
43
"get_scheduler_metadata",
vllm/vllm_flash_attn/flash_attn_interface.py
+60
@@ -417,6 +417,66 @@ def flash_attn_varlen_func(
417
return (out, softmax_lse) if return_softmax_lse else out
418
419
420
+def compile_flash_attn_varlen_func_from_specs(
421
+ *,
422
+ q_shape: tuple[int, ...],
423
+ k_shape: tuple[int, ...],
424
+ v_shape: tuple[int, ...],
425
+ q_dtype: torch.dtype,
426
+ v_stride: tuple[int, ...] | None = None,
427
+ cu_seqlens_q_shape: tuple[int, ...] | None = None,
428
+ cu_seqlens_k_shape: tuple[int, ...] | None = None,
429
+ max_seqlen_q: int | None = None,
430
+ max_seqlen_k: int | None = None,
431
+ dropout_p: float = 0.0,
432
+ softmax_scale=None,
433
+ causal=False,
434
+ window_size: list[int] | None = None,
435
+ deterministic=False,
436
+ return_softmax_lse=False,
437
+ num_splits: int = 0,
438
+ fa_version: int = DEFAULT_FA_VERSION,
439
+) -> None:
440
+ if fa_version != 4:
441
+ raise ValueError(
442
+ f"Compile-only FlashAttention is only supported for FA4, got FA{fa_version}"
443
+ )
444
+ if dropout_p != 0.0:
445
+ raise NotImplementedError("FA4 compile-only wrapper does not support dropout")
446
+ del deterministic
447
+
448
+ from vllm.vllm_flash_attn.cute.interface import (
449
+ compile_flash_attn_varlen_func_from_specs as _fa4_compile_flash_attn_varlen_func_from_specs,
450
+ )
451
+
452
+ real_window_size: tuple[int, int]
453
+ if window_size is None:
454
+ real_window_size = (-1, -1)
455
+ else:
456
+ assert len(window_size) == 2
457
+ real_window_size = (window_size[0], window_size[1])
458
+
459
+ if softmax_scale is None:
460
+ softmax_scale = q_shape[-1] ** (-0.5)
461
+
462
+ return _fa4_compile_flash_attn_varlen_func_from_specs(
463
+ q_shape=q_shape,
464
+ k_shape=k_shape,
465
+ v_shape=v_shape,
466
+ q_dtype=q_dtype,
467
+ v_stride=v_stride,
468
+ cu_seqlens_q_shape=cu_seqlens_q_shape,
469
+ cu_seqlens_k_shape=cu_seqlens_k_shape,
470
+ max_seqlen_q=max_seqlen_q,
471
+ max_seqlen_k=max_seqlen_k,
472
+ softmax_scale=softmax_scale,
473
+ causal=causal,
474
+ window_size=real_window_size,
475
+ num_splits=num_splits,
476
+ return_lse=return_softmax_lse,
477
+ )
478
+
479
+
480
def sparse_attn_func(
481
q,
482
k,