[Kernel] Warm up Qwen GDN gated RMSNorm (#54251)
Signed-off-by: Zupeng Wang <71580390+zupengwang@users.noreply.github.com> Co-authored-by: OpenAI <support@openai.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Zupeng Wang committed
Sep 3, 2026 at 18:19 UTC
c21751c90b1b1d4cfd8fd6e951feaeb539171cfc
3 files changed
+361
-26
tests/kernels/test_fla_layernorm_guard.py
+56
@@ -7,10 +7,13 @@ import torch.nn.functional as F
7
8
from vllm.platforms import current_platform
9
from vllm.third_party.flash_linear_attention.ops.layernorm_guard import (
10
+ LayerNormFwdKernel,
11
+ calc_rows_per_block,
12
layer_norm_fwd,
13
layernorm_fn,
14
rms_norm_ref,
15
)
16
+from vllm.triton_utils import triton
17
from vllm.utils.torch_utils import set_random_seed
18
19
DEVICE = "xpu" if current_platform.is_xpu() else "cuda"
@@ -103,6 +106,59 @@ IS_RMS_NORM = [True, False]
106
SEEDS = [0, 42]
107
108
109
+@pytest.mark.parametrize("rows_per_token", [1, 2, 4, 8, 16])
110
+def test_layer_norm_fwd_warmup_keys_cover_qwen_gdn(
111
+ rows_per_token: int,
112
+) -> None:
113
+ device = torch.device(DEVICE)
114
+ group_size = 128
115
+ max_num_tokens = 512
116
+ kernel = LayerNormFwdKernel()
117
+ warmup_keys = set(
118
+ kernel.get_warmup_keys(
119
+ max_num_tokens=max_num_tokens,
120
+ rows_per_token=rows_per_token,
121
+ group_size=group_size,
122
+ x_dtype=torch.bfloat16,
123
+ weight_dtype=torch.bfloat16,
124
+ device=device,
125
+ norm_before_gate=True,
126
+ is_rms_norm=True,
127
+ activation="silu",
128
+ )
129
+ )
130
+
131
+ for num_tokens in range(1, max_num_tokens + 1):
132
+ num_rows = num_tokens * rows_per_token
133
+ runtime_key = kernel.dispatch(
134
+ x_dtype=torch.bfloat16,
135
+ y_dtype=torch.bfloat16,
136
+ weight_dtype=torch.bfloat16,
137
+ bias_dtype=None,
138
+ z_dtype=torch.bfloat16,
139
+ mean_dtype=None,
140
+ rstd_dtype=torch.float32,
141
+ x_aligned=True,
142
+ y_aligned=True,
143
+ weight_aligned=True,
144
+ bias_aligned=True,
145
+ z_aligned=True,
146
+ mean_aligned=True,
147
+ rstd_aligned=True,
148
+ stride_x_row=group_size,
149
+ stride_y_row=group_size,
150
+ stride_z_row=group_size,
151
+ M=num_rows,
152
+ N=group_size,
153
+ BLOCK_N=triton.next_power_of_2(group_size),
154
+ ROWS_PER_BLOCK=calc_rows_per_block(num_rows, device),
155
+ norm_before_gate=True,
156
+ is_rms_norm=True,
157
+ activation="silu",
158
+ )
159
+ assert runtime_key in warmup_keys
160
+
161
+
162
@pytest.mark.parametrize("num_tokens", NUM_TOKENS)
163
@pytest.mark.parametrize("hidden_size", HIDDEN_SIZES)
164
@pytest.mark.parametrize("dtype", DTYPES)
vllm/model_executor/warmup/qwen_triton_warmup.py
+35
-1
@@ -10,6 +10,7 @@ import torch
10
from vllm.logger import init_logger
11
12
if TYPE_CHECKING:
13
+ from vllm.config import ModelConfig
14
from vllm.v1.worker.gpu_model_runner import GPUModelRunner
15
16
logger = init_logger(__name__)
@@ -37,6 +38,9 @@ class _QwenGDNWarmupConfig:
38
conv_kernel_size: int
39
conv_state: torch.Tensor
40
conv_dtype: torch.dtype
41
+ norm_weight_dtype: torch.dtype
42
+ norm_before_gate: bool
43
+ norm_activation: str
44
a_log: torch.Tensor
45
dt_bias: torch.Tensor
46
state_stride_token: int
@@ -62,6 +66,7 @@ def _is_qwen_gdn_layer(module: object) -> bool:
66
"conv_kernel_size",
67
"tp_size",
68
"kv_cache",
69
+ "norm",
70
"A_log",
71
"dt_bias",
72
)
@@ -112,6 +117,7 @@ def _qwen_gdn_warmup_config(
117
tp_size = int(layer.tp_size)
118
h = int(layer.num_k_heads) // tp_size
119
hv = int(layer.num_v_heads) // tp_size
120
+ norm = layer.norm
121
122
return _QwenGDNWarmupConfig(
123
h=h,
@@ -121,6 +127,9 @@ def _qwen_gdn_warmup_config(
127
conv_kernel_size=int(layer.conv_kernel_size),
128
conv_state=conv_state,
129
conv_dtype=conv_state.dtype,
130
+ norm_weight_dtype=norm.weight.dtype,
131
+ norm_before_gate=bool(norm.norm_before_gate),
132
+ norm_activation=str(norm.activation),
133
a_log=layer.A_log,
134
dt_bias=layer.dt_bias,
135
state_stride_token=int(ssm_state.stride(0)),
@@ -134,6 +143,29 @@ def _qwen_gdn_warmup_config(
143
return None
144
145
146
+def _warm_gated_rms_norm_kernel(
147
+ device: torch.device,
148
+ config: _QwenGDNWarmupConfig,
149
+ max_num_tokens: int,
150
+ x_dtype: torch.dtype,
151
+) -> None:
152
+ from vllm.third_party.flash_linear_attention.ops.layernorm_guard import (
153
+ warmup_layer_norm_fwd,
154
+ )
155
+
156
+ warmup_layer_norm_fwd(
157
+ max_num_tokens=max_num_tokens,
158
+ rows_per_token=config.hv,
159
+ group_size=config.v,
160
+ x_dtype=x_dtype,
161
+ weight_dtype=config.norm_weight_dtype,
162
+ device=device,
163
+ norm_before_gate=config.norm_before_gate,
164
+ is_rms_norm=True,
165
+ activation=config.norm_activation,
166
+ )
167
+
168
+
169
def _warm_causal_conv1d_fwd_kernel(
170
device: torch.device, config: _QwenGDNWarmupConfig
171
) -> None:
@@ -249,7 +281,7 @@ def _synchronize_device(device: torch.device) -> None:
281
@torch.inference_mode()
282
def qwen_triton_warmup(
283
runner: "GPUModelRunner",
252
- model_config: object,
284
+ model_config: "ModelConfig",
285
) -> None:
286
"""Warm Qwen Triton kernels reported by the JIT monitor."""
287
if runner.is_pooling_model:
@@ -275,6 +307,8 @@ def qwen_triton_warmup(
307
if gdn_config is None:
308
return
309
310
+ max_num_tokens = max(1, int(getattr(runner, "max_num_tokens", 1)))
311
+ _warm_gated_rms_norm_kernel(device, gdn_config, max_num_tokens, model_config.dtype)
312
_warm_causal_conv1d_fwd_kernel(device, gdn_config)
313
_warm_fused_post_conv_kernel(device, gdn_config)
314
_warm_fused_sigmoid_gating_delta_rule_update_kernel(device, gdn_config)
vllm/third_party/flash_linear_attention/ops/layernorm_guard.py
+270
-25
@@ -13,11 +13,19 @@
13
# This backward pass is faster for dimensions up to 8k, but after that it's much slower due to register spilling.
14
# The models we train have hidden dim up to 8k anyway (e.g. Llama 70B), so this is fine.
15
16
+from dataclasses import dataclass
17
+
18
import torch
19
import torch.nn as nn
20
import torch.nn.functional as F
21
from einops import rearrange
22
23
+from vllm.model_executor.warmup.jit_warmup import VllmJitKernel
24
+from vllm.model_executor.warmup.jit_warmup_triton_helper import (
25
+ TritonPointerInputVariant,
26
+ TritonWarmupTensor,
27
+ triton_scalar_specialization_rep,
28
+)
29
from vllm.triton_utils import tl, triton
30
from vllm.utils.math_utils import cdiv, next_power_of_2
31
from vllm.utils.platform_utils import num_compute_units
@@ -175,6 +183,263 @@ def calc_rows_per_block(M: int, device: torch.device) -> int:
183
return rows_per_block
184
185
186
+class LayerNormFwdKernel(VllmJitKernel["LayerNormFwdKernel.CompileKey"]):
187
+ """Own the runtime and compile-only dispatch for FLA layer norm."""
188
+
189
+ @dataclass(frozen=True)
190
+ class CompileKey:
191
+ x_dtype: torch.dtype
192
+ y_dtype: torch.dtype
193
+ weight_dtype: torch.dtype
194
+ bias_dtype: torch.dtype | None
195
+ z_dtype: torch.dtype | None
196
+ mean_dtype: torch.dtype | None
197
+ rstd_dtype: torch.dtype
198
+ input_variant: TritonPointerInputVariant
199
+ stride_x_row: int
200
+ stride_y_row: int
201
+ stride_z_row: int
202
+ M: int
203
+ N: int
204
+ BLOCK_N: int
205
+ ROWS_PER_BLOCK: int
206
+ NORM_BEFORE_GATE: bool
207
+ IS_RMS_NORM: bool
208
+ ACTIVATION: str
209
+
210
+ kernel = staticmethod(layer_norm_fwd_kernel)
211
+
212
+ def dispatch( # type: ignore[override]
213
+ self,
214
+ *,
215
+ x_dtype: torch.dtype,
216
+ y_dtype: torch.dtype,
217
+ weight_dtype: torch.dtype,
218
+ bias_dtype: torch.dtype | None,
219
+ z_dtype: torch.dtype | None,
220
+ mean_dtype: torch.dtype | None,
221
+ rstd_dtype: torch.dtype,
222
+ x_aligned: bool,
223
+ y_aligned: bool,
224
+ weight_aligned: bool,
225
+ bias_aligned: bool,
226
+ z_aligned: bool,
227
+ mean_aligned: bool,
228
+ rstd_aligned: bool,
229
+ stride_x_row: int,
230
+ stride_y_row: int,
231
+ stride_z_row: int,
232
+ M: int,
233
+ N: int,
234
+ BLOCK_N: int,
235
+ ROWS_PER_BLOCK: int,
236
+ norm_before_gate: bool,
237
+ is_rms_norm: bool,
238
+ activation: str,
239
+ ) -> CompileKey:
240
+ input_variant = TritonPointerInputVariant.from_alignment(
241
+ x=x_aligned,
242
+ y=y_aligned,
243
+ weight=weight_aligned,
244
+ bias=bias_aligned,
245
+ z=z_aligned,
246
+ mean=mean_aligned,
247
+ rstd=rstd_aligned,
248
+ )
249
+ return self.CompileKey(
250
+ x_dtype=x_dtype,
251
+ y_dtype=y_dtype,
252
+ weight_dtype=weight_dtype,
253
+ bias_dtype=bias_dtype,
254
+ z_dtype=z_dtype,
255
+ mean_dtype=mean_dtype,
256
+ rstd_dtype=rstd_dtype,
257
+ input_variant=input_variant,
258
+ stride_x_row=triton_scalar_specialization_rep(stride_x_row),
259
+ stride_y_row=triton_scalar_specialization_rep(stride_y_row),
260
+ stride_z_row=triton_scalar_specialization_rep(stride_z_row),
261
+ M=triton_scalar_specialization_rep(M),
262
+ N=N,
263
+ BLOCK_N=BLOCK_N,
264
+ ROWS_PER_BLOCK=ROWS_PER_BLOCK,
265
+ NORM_BEFORE_GATE=norm_before_gate,
266
+ IS_RMS_NORM=is_rms_norm,
267
+ ACTIVATION=activation,
268
+ )
269
+
270
+ def get_warmup_keys(
271
+ self,
272
+ *,
273
+ max_num_tokens: int,
274
+ rows_per_token: int,
275
+ group_size: int,
276
+ x_dtype: torch.dtype,
277
+ weight_dtype: torch.dtype,
278
+ device: torch.device,
279
+ norm_before_gate: bool,
280
+ is_rms_norm: bool,
281
+ activation: str,
282
+ ) -> list[CompileKey]:
283
+ if max_num_tokens < 1 or rows_per_token < 1:
284
+ return []
285
+
286
+ max_fused_size = 65536 // torch.empty((), dtype=x_dtype).element_size()
287
+ block_n = min(max_fused_size, triton.next_power_of_2(group_size))
288
+ if group_size > block_n:
289
+ return []
290
+
291
+ # Triton's runtime key only distinguishes three integer specialization
292
+ # classes. Keep one real M for each (M specialization, rows/block) pair.
293
+ m_values: dict[tuple[int, int], int] = {}
294
+ for num_tokens in range(1, max_num_tokens + 1):
295
+ M = num_tokens * rows_per_token
296
+ signature = (
297
+ triton_scalar_specialization_rep(M),
298
+ calc_rows_per_block(M, device),
299
+ )
300
+ m_values.setdefault(signature, M)
301
+
302
+ keys: dict[LayerNormFwdKernel.CompileKey, None] = {}
303
+ for M in m_values.values():
304
+ key = self.dispatch(
305
+ x_dtype=x_dtype,
306
+ y_dtype=x_dtype,
307
+ weight_dtype=weight_dtype,
308
+ bias_dtype=None,
309
+ z_dtype=x_dtype,
310
+ mean_dtype=None if is_rms_norm else torch.float32,
311
+ rstd_dtype=torch.float32,
312
+ x_aligned=True,
313
+ y_aligned=True,
314
+ weight_aligned=True,
315
+ bias_aligned=True,
316
+ z_aligned=True,
317
+ mean_aligned=True,
318
+ rstd_aligned=True,
319
+ stride_x_row=group_size,
320
+ stride_y_row=group_size,
321
+ stride_z_row=group_size,
322
+ M=M,
323
+ N=group_size,
324
+ BLOCK_N=block_n,
325
+ ROWS_PER_BLOCK=calc_rows_per_block(M, device),
326
+ norm_before_gate=norm_before_gate,
327
+ is_rms_norm=is_rms_norm,
328
+ activation=activation,
329
+ )
330
+ keys[key] = None
331
+ return list(keys)
332
+
333
+ def compile(self, compile_key: CompileKey) -> None:
334
+ warmup = getattr(self.kernel, "warmup", None)
335
+ assert warmup is not None
336
+ variant = compile_key.input_variant
337
+
338
+ def pointer(
339
+ name: str, dtype: torch.dtype | None
340
+ ) -> TritonWarmupTensor | None:
341
+ return None if dtype is None else variant.pointer(name, dtype)
342
+
343
+ warmup(
344
+ pointer("x", compile_key.x_dtype),
345
+ pointer("y", compile_key.y_dtype),
346
+ pointer("weight", compile_key.weight_dtype),
347
+ pointer("bias", compile_key.bias_dtype),
348
+ pointer("z", compile_key.z_dtype),
349
+ pointer("mean", compile_key.mean_dtype),
350
+ pointer("rstd", compile_key.rstd_dtype),
351
+ compile_key.stride_x_row,
352
+ compile_key.stride_y_row,
353
+ compile_key.stride_z_row,
354
+ compile_key.M,
355
+ compile_key.N,
356
+ 1e-5,
357
+ BLOCK_N=compile_key.BLOCK_N,
358
+ ROWS_PER_BLOCK=compile_key.ROWS_PER_BLOCK,
359
+ HAS_BIAS=compile_key.bias_dtype is not None,
360
+ HAS_Z=compile_key.z_dtype is not None,
361
+ NORM_BEFORE_GATE=compile_key.NORM_BEFORE_GATE,
362
+ IS_RMS_NORM=compile_key.IS_RMS_NORM,
363
+ ACTIVATION=compile_key.ACTIVATION,
364
+ num_warps=min(max(compile_key.BLOCK_N // 256, 1), 8),
365
+ grid=(1, 1),
366
+ )
367
+
368
+ def __call__(
369
+ self,
370
+ x: torch.Tensor,
371
+ out: torch.Tensor,
372
+ weight: torch.Tensor,
373
+ bias: torch.Tensor | None,
374
+ z: torch.Tensor | None,
375
+ mean: torch.Tensor | None,
376
+ rstd: torch.Tensor,
377
+ eps: float,
378
+ group_size: int,
379
+ norm_before_gate: bool,
380
+ is_rms_norm: bool,
381
+ activation: str,
382
+ ) -> None:
383
+ M = x.shape[0]
384
+ max_fused_size = 65536 // x.element_size()
385
+ block_n = min(max_fused_size, triton.next_power_of_2(group_size))
386
+ if group_size > block_n:
387
+ raise RuntimeError("This layer norm doesn't support feature dim >= 64KB.")
388
+ rows_per_block = calc_rows_per_block(M, x.device)
389
+ grid = (cdiv(M, rows_per_block), x.shape[1] // group_size)
390
+ self.kernel[grid](
391
+ x,
392
+ out,
393
+ weight,
394
+ bias,
395
+ z,
396
+ mean,
397
+ rstd,
398
+ x.stride(0),
399
+ out.stride(0),
400
+ z.stride(0) if z is not None else 0,
401
+ M,
402
+ group_size,
403
+ eps,
404
+ BLOCK_N=block_n,
405
+ ROWS_PER_BLOCK=rows_per_block,
406
+ HAS_BIAS=bias is not None,
407
+ HAS_Z=z is not None,
408
+ NORM_BEFORE_GATE=norm_before_gate,
409
+ IS_RMS_NORM=is_rms_norm,
410
+ ACTIVATION=activation,
411
+ num_warps=min(max(block_n // 256, 1), 8),
412
+ )
413
+
414
+
415
+_LAYER_NORM_FWD_KERNEL = LayerNormFwdKernel()
416
+
417
+
418
+def warmup_layer_norm_fwd(
419
+ *,
420
+ max_num_tokens: int,
421
+ rows_per_token: int,
422
+ group_size: int,
423
+ x_dtype: torch.dtype,
424
+ weight_dtype: torch.dtype,
425
+ device: torch.device,
426
+ norm_before_gate: bool,
427
+ is_rms_norm: bool,
428
+ activation: str,
429
+) -> None:
430
+ _LAYER_NORM_FWD_KERNEL.warmup(
431
+ max_num_tokens=max_num_tokens,
432
+ rows_per_token=rows_per_token,
433
+ group_size=group_size,
434
+ x_dtype=x_dtype,
435
+ weight_dtype=weight_dtype,
436
+ device=device,
437
+ norm_before_gate=norm_before_gate,
438
+ is_rms_norm=is_rms_norm,
439
+ activation=activation,
440
+ )
441
+
442
+
443
def layer_norm_fwd(
444
x: torch.Tensor,
445
weight: torch.Tensor,
@@ -213,18 +478,7 @@ def layer_norm_fwd(
478
else None
479
)
480
rstd = torch.empty((ngroups * M,), dtype=torch.float32, device=x.device)
216
- # Less than 64KB per feature: enqueue fused kernel
217
- MAX_FUSED_SIZE = 65536 // x.element_size()
218
- BLOCK_N = min(MAX_FUSED_SIZE, triton.next_power_of_2(group_size))
219
- if group_size > BLOCK_N:
220
- raise RuntimeError("This layer norm doesn't support feature dim >= 64KB.")
221
- # heuristics for number of warps
222
- num_warps = min(max(BLOCK_N // 256, 1), 8)
223
- # Calculate rows per block based on SM count
224
- rows_per_block = calc_rows_per_block(M, x.device)
225
- # Update grid to use rows_per_block
226
- grid = (cdiv(M, rows_per_block), ngroups)
227
- layer_norm_fwd_kernel[grid](
481
+ _LAYER_NORM_FWD_KERNEL(
482
x,
483
out,
484
weight,
@@ -232,20 +486,11 @@ def layer_norm_fwd(
486
z,
487
mean,
488
rstd,
235
- x.stride(0),
236
- out.stride(0),
237
- z.stride(0) if z is not None else 0,
238
- M,
239
- group_size,
489
eps,
241
- BLOCK_N=BLOCK_N,
242
- ROWS_PER_BLOCK=rows_per_block,
243
- HAS_BIAS=bias is not None,
244
- HAS_Z=z is not None,
245
- NORM_BEFORE_GATE=norm_before_gate,
246
- IS_RMS_NORM=is_rms_norm,
247
- num_warps=num_warps,
248
- ACTIVATION=activation,
490
+ group_size,
491
+ norm_before_gate,
492
+ is_rms_norm,
493
+ activation,
494
)
495
return out, mean, rstd
496