[Warmup] Gemma 4 de-JITification (#55768)
Signed-off-by: LopezCastroRoberto <rocastro@redhat.com> Co-authored-by: Codex <noreply@openai.com>
Roberto L. Castro committed
Sep 13, 2026 at 22:47 UTC
186a1e62221c02cdf8ef0433eb4ad33d1b48c7f7
4 files changed
+287
-27
vllm/v1/attention/backends/flash_attn.py
+227
-10
@@ -5,12 +5,13 @@
5
import copy
6
import functools
7
from dataclasses import dataclass
8
-from typing import ClassVar
8
+from typing import Any, ClassVar
9
10
import numpy as np
11
import torch
12
13
from vllm.model_executor.layers.attention import Attention
14
+from vllm.model_executor.warmup.jit_warmup import VllmJitKernel, WarmupChoices
15
from vllm.platforms import current_platform
16
from vllm.utils.torch_utils import (
17
PIN_MEMORY,
@@ -46,6 +47,7 @@ from vllm.v1.worker.workspace import current_workspace_manager
47
48
if is_flash_attn_varlen_func_available():
49
from vllm.v1.attention.backends.fa_utils import (
50
+ compile_flash_attn_varlen_func_from_specs,
51
flash_attn_supports_sinks,
52
flash_attn_varlen_func,
53
get_scheduler_metadata,
@@ -67,7 +69,13 @@ from vllm.v1.attention.backend import (
69
AttentionMetadataBuilder,
70
CommonAttentionMetadata,
71
)
70
-from vllm.v1.kv_cache_interface import AttentionSpec, KVCacheSpec
72
+from vllm.v1.kv_cache_interface import (
73
+ AttentionSpec,
74
+ FullAttentionSpec,
75
+ KVCacheSpec,
76
+ KVQuantMode,
77
+ SlidingWindowSpec,
78
+)
79
from vllm.v1.worker.cp_utils import (
80
run_split_fa2_dcp_context_attention,
81
should_skip_dcp_context_attention,
@@ -77,6 +85,200 @@ from vllm.v1.worker.cp_utils import (
85
86
logger = init_logger(__name__)
87
88
+FA4_DENSE_FLOAT_DTYPES = (torch.bfloat16, torch.float16)
89
+FA4_DENSE_Q_TILE = 128
90
+FA4_DENSE_NUM_BLOCKS = 256
91
+FA4_DENSE_MAX_SEQLEN_K = 8192
92
+
93
+
94
+class FA4DenseAttentionKernel(VllmJitKernel["FA4DenseAttentionKernel.CompileKey"]):
95
+ """Warm paged, causal FA4 on Blackwell and Hopper."""
96
+
97
+ @dataclass(frozen=True)
98
+ class CompileKey:
99
+ q_stage: int
100
+ dtype: torch.dtype
101
+ qhead_per_kvhead: int
102
+ head_dim: int
103
+ page_size: int
104
+ has_window_left: bool
105
+ has_window_right: bool
106
+ softcap: float
107
+ is_single_batch: bool
108
+ num_splits: int
109
+
110
+ @staticmethod
111
+ def kernel(*args: Any, **kwargs: Any) -> Any:
112
+ assert flash_attn_varlen_func is not None
113
+ return flash_attn_varlen_func(*args, **kwargs)
114
+
115
+ def dispatch( # type: ignore[override]
116
+ self,
117
+ *,
118
+ q_stage: int,
119
+ dtype: torch.dtype,
120
+ num_qo_heads: int,
121
+ num_kv_heads: int,
122
+ head_dim: int,
123
+ page_size: int,
124
+ window_size: tuple[int, int],
125
+ softcap: float,
126
+ is_single_batch: bool = True,
127
+ num_splits: int = 1,
128
+ ) -> CompileKey:
129
+ return self.CompileKey(
130
+ q_stage=q_stage,
131
+ dtype=dtype,
132
+ qhead_per_kvhead=num_qo_heads // num_kv_heads,
133
+ head_dim=head_dim,
134
+ page_size=page_size,
135
+ has_window_left=window_size[0] >= 0,
136
+ has_window_right=window_size[1] >= 0,
137
+ softcap=softcap,
138
+ is_single_batch=is_single_batch,
139
+ num_splits=num_splits,
140
+ )
141
+
142
+ def get_warmup_keys(
143
+ self,
144
+ *,
145
+ vllm_config: VllmConfig,
146
+ kv_cache_spec: AttentionSpec,
147
+ num_qo_heads: int,
148
+ ) -> list[CompileKey]:
149
+ capability = current_platform.get_device_capability()
150
+ major = capability.major if capability is not None else None
151
+ if (
152
+ vllm_config.parallel_config.decode_context_parallel_size != 1
153
+ or vllm_config.model_config.rswa_window is not None
154
+ or kv_cache_spec.kv_quant_mode != KVQuantMode.NONE
155
+ or kv_cache_spec.dtype not in FA4_DENSE_FLOAT_DTYPES
156
+ or not isinstance(kv_cache_spec, (FullAttentionSpec, SlidingWindowSpec))
157
+ ):
158
+ return []
159
+ head_dim = kv_cache_spec.head_size
160
+ page_size = kv_cache_spec.block_size
161
+ sliding_window = kv_cache_spec.sliding_window
162
+ window_size = (
163
+ (sliding_window - 1, 0) if sliding_window is not None else (-1, -1)
164
+ )
165
+ softcap = float(
166
+ getattr(
167
+ vllm_config.model_config.hf_text_config, "attn_logit_softcapping", 0
168
+ )
169
+ or 0
170
+ )
171
+ if (
172
+ major in (10, 11)
173
+ and uses_fa4_hd256_kernel(head_dim)
174
+ and page_size == FA4_HD256_PAGE_SIZE
175
+ ):
176
+ q_stages = WarmupChoices(1, 2)
177
+ single_batches = WarmupChoices(True, False)
178
+ split_counts = WarmupChoices(1)
179
+ elif (
180
+ major == 9 and head_dim == 512 and window_size == (-1, -1) and softcap == 0
181
+ ):
182
+ # SM90 forward has split/non-split variants; its transitive
183
+ # combine specializes at 32/64/128/256 splits.
184
+ q_stages = WarmupChoices(1)
185
+ single_batches = WarmupChoices(False)
186
+ split_counts = WarmupChoices(1, 32, 64, 128, 256)
187
+ else:
188
+ return []
189
+
190
+ return self._trace_dispatch(self.dispatch)(
191
+ q_stage=q_stages,
192
+ is_single_batch=single_batches,
193
+ num_splits=split_counts,
194
+ dtype=vllm_config.model_config.dtype,
195
+ num_qo_heads=num_qo_heads,
196
+ num_kv_heads=kv_cache_spec.num_kv_heads,
197
+ head_dim=head_dim,
198
+ page_size=page_size,
199
+ window_size=WarmupChoices(window_size),
200
+ softcap=softcap,
201
+ )
202
+
203
+ def compile(self, compile_key: CompileKey) -> None:
204
+ assert compile_flash_attn_varlen_func_from_specs is not None
205
+ max_seqlen_q = FA4_DENSE_Q_TILE + 1 if compile_key.q_stage == 2 else 1
206
+ batch_size = 1 if compile_key.is_single_batch else 2
207
+ kv_shape = (
208
+ FA4_DENSE_NUM_BLOCKS,
209
+ compile_key.page_size,
210
+ 1,
211
+ compile_key.head_dim,
212
+ )
213
+ kv_stride = (
214
+ 2 * compile_key.page_size * compile_key.head_dim,
215
+ 2 * compile_key.head_dim,
216
+ 2 * compile_key.head_dim,
217
+ 1,
218
+ )
219
+ compile_flash_attn_varlen_func_from_specs(
220
+ q_shape=(
221
+ batch_size * max_seqlen_q,
222
+ compile_key.qhead_per_kvhead,
223
+ compile_key.head_dim,
224
+ ),
225
+ k_shape=kv_shape,
226
+ v_shape=kv_shape,
227
+ k_stride=kv_stride,
228
+ v_stride=kv_stride,
229
+ q_dtype=compile_key.dtype,
230
+ cu_seqlens_q_shape=(batch_size + 1,),
231
+ seqused_k_shape=(batch_size,),
232
+ page_table_shape=(
233
+ batch_size,
234
+ FA4_DENSE_MAX_SEQLEN_K // compile_key.page_size,
235
+ ),
236
+ max_seqlen_q=max_seqlen_q,
237
+ max_seqlen_k=FA4_DENSE_MAX_SEQLEN_K,
238
+ causal=True,
239
+ window_size=[
240
+ 1 if compile_key.has_window_left else -1,
241
+ 1 if compile_key.has_window_right else -1,
242
+ ],
243
+ softcap=compile_key.softcap,
244
+ num_splits=2 if compile_key.num_splits > 1 else 1,
245
+ fa_version=4,
246
+ return_softmax_lse=False,
247
+ )
248
+ if compile_key.head_dim == 512 and compile_key.num_splits > 1:
249
+ from torch._subclasses.fake_tensor import FakeTensorMode
250
+
251
+ from vllm.vllm_flash_attn.cute.interface import _flash_attn_fwd_combine
252
+
253
+ # FA4's forward compile-only entry returns before its transitive
254
+ # SplitKV combine. Call the native combine entry in FakeTensorMode:
255
+ # it populates the native cache without launching the kernel.
256
+ with FakeTensorMode():
257
+ _flash_attn_fwd_combine(
258
+ torch.empty(
259
+ (compile_key.num_splits, 2, 1, 512),
260
+ dtype=torch.float32,
261
+ device="cuda",
262
+ ),
263
+ torch.empty(
264
+ (compile_key.num_splits, 2, 1),
265
+ dtype=torch.float32,
266
+ device="cuda",
267
+ ),
268
+ torch.empty((2, 1, 512), dtype=compile_key.dtype, device="cuda"),
269
+ cu_seqlens=torch.empty((3,), dtype=torch.int32, device="cuda"),
270
+ )
271
+
272
+ def __call__(
273
+ self,
274
+ *,
275
+ q: torch.Tensor,
276
+ k: torch.Tensor,
277
+ v: torch.Tensor,
278
+ **kwargs: Any,
279
+ ) -> Any:
280
+ return self.kernel(q=q, k=k, v=v, **kwargs)
281
+
282
283
class FlashAttentionBackend(AttentionBackend):
284
supported_dtypes: ClassVar[list[torch.dtype]] = [torch.float16, torch.bfloat16]
@@ -447,15 +649,27 @@ class FlashAttentionMetadataBuilder(AttentionMetadataBuilder[FlashAttentionMetad
649
650
self.max_num_splits = 0 # No upper bound on the number of splits.
651
self.aot_schedule = get_flash_attn_version() == 3
652
+ head_size_v = getattr(kv_cache_spec, "head_size_v", None)
653
+ fa_version = get_flash_attn_version(
654
+ head_size=self.headdim,
655
+ head_size_v=head_size_v,
656
+ kv_cache_block_size=self.block_size,
657
+ supports_fa4_hd256=True,
658
+ )
659
+ self.fa4_hd256 = fa_version == 4 and uses_fa4_hd256_kernel(
660
+ self.headdim, head_size_v
661
+ )
662
451
- self.fa4_hd256 = uses_fa4_hd256_kernel(self.headdim) and (
452
- get_flash_attn_version(
453
- head_size=self.headdim,
454
- kv_cache_block_size=self.block_size,
455
- supports_fa4_hd256=True,
663
+ if (
664
+ vllm_config.kernel_config.enable_jit_warmup
665
+ and self.model_config.hf_config.model_type in ("gemma4", "gemma4_unified")
666
+ and fa_version == 4
667
+ ):
668
+ _FA4_DENSE_ATTENTION_KERNEL.register_warmup(
669
+ vllm_config=vllm_config,
670
+ kv_cache_spec=kv_cache_spec,
671
+ num_qo_heads=self.num_heads_q,
672
)
457
- == 4
458
- )
673
674
try:
675
from vllm.distributed.parallel_state import get_dcp_group
@@ -1164,7 +1378,7 @@ class FlashAttentionImpl(AttentionImpl):
1378
block_table = block_table[:, :num_pages]
1379
num_splits = 1
1380
1167
- flash_attn_varlen_func(
1381
+ _FA4_DENSE_ATTENTION_KERNEL(
1382
q=query[:num_actual_tokens],
1383
k=key_cache,
1384
v=value_cache,
@@ -1837,3 +2051,6 @@ def cascade_attention(
2051
2052
# Merge prefix and suffix outputs, and store the result in output.
2053
merge_attn_states(output, prefix_output, prefix_lse, suffix_output, suffix_lse)
2054
+
2055
+
2056
+_FA4_DENSE_ATTENTION_KERNEL = FA4DenseAttentionKernel()
vllm/v1/worker/gpu/attn_utils.py
+4
-1
@@ -8,7 +8,10 @@ from typing import TYPE_CHECKING, Any, cast
8
import numpy as np
9
import torch
10
11
-from vllm.config import VllmConfig, get_layers_from_vllm_config
11
+from vllm.config import (
12
+ VllmConfig,
13
+ get_layers_from_vllm_config,
14
+)
15
from vllm.config.compilation import CUDAGraphMode
16
from vllm.model_executor.layers.attention import Attention
17
from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
vllm/v1/worker/gpu/model_runner.py
+12
-6
@@ -632,12 +632,18 @@ class GPUModelRunner(LoRAModelRunnerMixin):
632
draft_attn_layer_names = None
633
if isinstance(self.speculator, DraftModelSpeculator):
634
draft_attn_layer_names = self.speculator.draft_attn_layer_names
635
- self.attn_groups, attn_cg_support, self.kernel_block_sizes = init_attn_backend(
636
- self.kv_cache_config,
637
- self.vllm_config,
638
- self.device,
639
- draft_layer_names=draft_attn_layer_names,
640
- )
635
+ # Metadata builders select attention kernels that need JIT warmup.
636
+ with self.jit_warmup_registry.activate():
637
+ (
638
+ self.attn_groups,
639
+ attn_cg_support,
640
+ self.kernel_block_sizes,
641
+ ) = init_attn_backend(
642
+ self.kv_cache_config,
643
+ self.vllm_config,
644
+ self.device,
645
+ draft_layer_names=draft_attn_layer_names,
646
+ )
647
additional_attn_cg_support = self.model_state.get_additional_cg_support()
648
attn_cg_support = attn_cg_support.narrow(*additional_attn_cg_support)
649
# The speculator clears the flag at load time when the checkpoint has
vllm/vllm_flash_attn/flash_attn_interface.py
+44
-10
@@ -445,6 +445,7 @@ def compile_flash_attn_varlen_func_from_specs(
445
v_shape: tuple[int, ...],
446
q_dtype: torch.dtype,
447
v_stride: tuple[int, ...] | None = None,
448
+ k_stride: tuple[int, ...] | None = None,
449
cu_seqlens_q_shape: tuple[int, ...] | None = None,
450
cu_seqlens_k_shape: tuple[int, ...] | None = None,
451
max_seqlen_q: int | None = None,
@@ -457,6 +458,9 @@ def compile_flash_attn_varlen_func_from_specs(
458
return_softmax_lse=False,
459
num_splits: int = 0,
460
fa_version: int = DEFAULT_FA_VERSION,
461
+ seqused_k_shape: tuple[int, ...] | None = None,
462
+ page_table_shape: tuple[int, ...] | None = None,
463
+ softcap: float | None = None,
464
) -> None:
465
if fa_version != 4:
466
raise ValueError(
@@ -467,7 +471,8 @@ def compile_flash_attn_varlen_func_from_specs(
471
del deterministic
472
473
from vllm.vllm_flash_attn.cute.interface import (
470
- compile_flash_attn_varlen_func_from_specs as _fa4_compile_flash_attn_varlen_func_from_specs,
474
+ _flash_attn_fwd,
475
+ _make_compile_only_tensor_spec,
476
)
477
478
real_window_size: tuple[int, int]
@@ -480,21 +485,50 @@ def compile_flash_attn_varlen_func_from_specs(
485
if softmax_scale is None:
486
softmax_scale = q_shape[-1] ** (-0.5)
487
483
- return _fa4_compile_flash_attn_varlen_func_from_specs(
484
- q_shape=q_shape,
485
- k_shape=k_shape,
486
- v_shape=v_shape,
487
- q_dtype=q_dtype,
488
- v_stride=v_stride,
489
- cu_seqlens_q_shape=cu_seqlens_q_shape,
490
- cu_seqlens_k_shape=cu_seqlens_k_shape,
488
+ q = _make_compile_only_tensor_spec(q_shape, q_dtype)
489
+ k = _make_compile_only_tensor_spec(k_shape, q_dtype, stride=k_stride)
490
+ v = _make_compile_only_tensor_spec(v_shape, q_dtype, stride=v_stride)
491
+ cu_seqlens_q = _make_compile_only_tensor_spec(cu_seqlens_q_shape, torch.int32, 4)
492
+ out = _make_compile_only_tensor_spec(
493
+ (*q_shape[:-1], v_shape[-1]),
494
+ q_dtype,
495
+ )
496
+ lse = None
497
+ if return_softmax_lse:
498
+ assert q is not None
499
+ if cu_seqlens_q_shape is None:
500
+ lse_shape = (*q.shape[:-3], q.shape[-2], q.shape[-3])
501
+ lse_stride = None
502
+ else:
503
+ lse_shape = (q.shape[-2], q.shape[0])
504
+ lse_stride = (None, 1)
505
+ lse = _make_compile_only_tensor_spec(
506
+ lse_shape,
507
+ torch.float32,
508
+ 4,
509
+ stride=lse_stride,
510
+ )
511
+
512
+ _flash_attn_fwd(
513
+ q=q,
514
+ k=k,
515
+ v=v,
516
+ cu_seqlens_q=cu_seqlens_q,
517
+ cu_seqlens_k=_make_compile_only_tensor_spec(cu_seqlens_k_shape, torch.int32, 4),
518
+ seqused_k=_make_compile_only_tensor_spec(seqused_k_shape, torch.int32, 4),
519
+ page_table=_make_compile_only_tensor_spec(page_table_shape, torch.int32, 4),
520
max_seqlen_q=max_seqlen_q,
521
max_seqlen_k=max_seqlen_k,
522
softmax_scale=softmax_scale,
523
causal=causal,
495
- window_size=real_window_size,
524
+ softcap=softcap,
525
+ window_size_left=real_window_size[0],
526
+ window_size_right=real_window_size[1],
527
num_splits=num_splits,
528
return_lse=return_softmax_lse,
529
+ out=out,
530
+ lse=lse,
531
+ compile_only=True,
532
)
533
534