A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 14, 05:30 UTC
[PERF] [Qwen3-next] Speed up gated RMSNorm (#26207)
Signed-off-by: Vadim Gimpelson <vadim.gimpelson@gmail.com> Signed-off-by: Vadim Gimpelson <156319763+vadiklyutiy@users.noreply.github.com> Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
Vadim Gimpelson committed
Oct 12, 2025 at 12:27 UTC
82e64c7a204671423c2c7914cbabdf06bc89c0c8
2 files changed
+475
-33
tests/kernels/test_fla_layernorm_guard.py
new
+388
@@ -0,0 +1,388 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import pytest
5
+import torch
6
+import torch.nn.functional as F
7
+
8
+from vllm.model_executor.layers.fla.ops.layernorm_guard import (
9
+ layer_norm_fwd,
10
+ layernorm_fn,
11
+ rms_norm_ref,
12
+)
13
+from vllm.platforms import current_platform
14
+
15
+
16
+def layer_norm_ref(
17
+ x,
18
+ weight,
19
+ bias,
20
+ z=None,
21
+ eps=1e-6,
22
+ group_size=None,
23
+ norm_before_gate=True,
24
+ is_rms_norm=False,
25
+):
26
+ """Reference implementation for both layer norm and RMS norm."""
27
+ if is_rms_norm:
28
+ # Use the imported rms_norm_ref for RMS norm cases
29
+ return rms_norm_ref(
30
+ x,
31
+ weight,
32
+ bias,
33
+ z=z,
34
+ eps=eps,
35
+ group_size=group_size,
36
+ norm_before_gate=norm_before_gate,
37
+ upcast=True,
38
+ )
39
+
40
+ # Layer norm implementation
41
+ dtype = x.dtype
42
+ x = x.float()
43
+ weight = weight.float()
44
+ bias = bias.float() if bias is not None else None
45
+ z = z.float() if z is not None else None
46
+
47
+ if z is not None and not norm_before_gate:
48
+ x = x * F.silu(z)
49
+
50
+ if group_size is None:
51
+ # Layer norm: subtract mean
52
+ mean = x.mean(dim=-1, keepdim=True)
53
+ var = ((x - mean).square()).mean(dim=-1, keepdim=True)
54
+ rstd = 1 / torch.sqrt(var + eps)
55
+ out = (x - mean) * rstd * weight
56
+ if bias is not None:
57
+ out = out + bias
58
+ else:
59
+ # Group norm
60
+ from einops import rearrange
61
+
62
+ x_group = rearrange(x, "... (g d) -> ... g d", d=group_size)
63
+ mean = x_group.mean(dim=-1, keepdim=True)
64
+ var = ((x_group - mean).square()).mean(dim=-1, keepdim=True)
65
+ rstd = 1 / torch.sqrt(var + eps)
66
+ x_group = (x_group - mean) * rstd
67
+ out = rearrange(x_group, "... g d -> ... (g d)") * weight
68
+ if bias is not None:
69
+ out = out + bias
70
+
71
+ if z is not None and norm_before_gate:
72
+ out *= F.silu(z)
73
+
74
+ return out.to(dtype)
75
+
76
+
77
+DTYPES = [torch.bfloat16, torch.float32]
78
+# Test various M sizes to ensure rows_per_block logic works correctly
79
+NUM_TOKENS = [
80
+ 1,
81
+ 7,
82
+ 16,
83
+ 63,
84
+ 128,
85
+ 256,
86
+ 512,
87
+ 1024,
88
+ 2048,
89
+ 4096,
90
+ 5789,
91
+ 8189,
92
+ 8191,
93
+ 16383,
94
+ 32767,
95
+]
96
+HIDDEN_SIZES = [64, 128, 256, 1024]
97
+GROUP_SIZES = [None, 64, 128] # None means full hidden size
98
+NORM_BEFORE_GATE = [True, False]
99
+IS_RMS_NORM = [True, False]
100
+SEEDS = [0, 42]
101
+
102
+
103
+@pytest.mark.parametrize("num_tokens", NUM_TOKENS)
104
+@pytest.mark.parametrize("hidden_size", HIDDEN_SIZES)
105
+@pytest.mark.parametrize("dtype", DTYPES)
106
+@pytest.mark.parametrize("seed", SEEDS)
107
+@pytest.mark.parametrize("is_rms_norm", IS_RMS_NORM)
108
+@torch.inference_mode()
109
+def test_layer_norm_fwd_basic(
110
+ num_tokens: int,
111
+ hidden_size: int,
112
+ dtype: torch.dtype,
113
+ seed: int,
114
+ is_rms_norm: bool,
115
+) -> None:
116
+ """Test basic layer norm forward pass without z (gate) tensor."""
117
+ current_platform.seed_everything(seed)
118
+ device = torch.device("cuda:0")
119
+
120
+ # Create inputs
121
+ x = torch.randn(num_tokens, hidden_size, dtype=dtype, device=device)
122
+ weight = torch.randn(hidden_size, dtype=dtype, device=device)
123
+ bias = None if is_rms_norm else torch.randn(hidden_size, dtype=dtype, device=device)
124
+ eps = 1e-6
125
+
126
+ # Run the triton kernel
127
+ out, mean, rstd = layer_norm_fwd(
128
+ x, weight, bias, eps, z=None, is_rms_norm=is_rms_norm
129
+ )
130
+
131
+ # Run reference implementation
132
+ ref_out = layer_norm_ref(x, weight, bias, z=None, eps=eps, is_rms_norm=is_rms_norm)
133
+
134
+ # Check outputs
135
+ assert out.shape == x.shape
136
+ assert out.dtype == x.dtype
137
+ torch.testing.assert_close(out, ref_out, atol=1e-2, rtol=1e-2)
138
+
139
+ # Check mean and rstd shapes
140
+ if not is_rms_norm:
141
+ assert mean.shape == (num_tokens,)
142
+ assert rstd.shape == (num_tokens,)
143
+
144
+
145
+@pytest.mark.parametrize("num_tokens", NUM_TOKENS)
146
+@pytest.mark.parametrize("hidden_size", [128, 256, 1024])
147
+@pytest.mark.parametrize("dtype", DTYPES)
148
+@pytest.mark.parametrize("norm_before_gate", NORM_BEFORE_GATE)
149
+@pytest.mark.parametrize("is_rms_norm", IS_RMS_NORM)
150
+@torch.inference_mode()
151
+def test_layer_norm_fwd_with_gate(
152
+ num_tokens: int,
153
+ hidden_size: int,
154
+ dtype: torch.dtype,
155
+ norm_before_gate: bool,
156
+ is_rms_norm: bool,
157
+) -> None:
158
+ """Test layer norm forward pass with z (gate) tensor."""
159
+ current_platform.seed_everything(42)
160
+ device = torch.device("cuda:0")
161
+
162
+ # Create inputs
163
+ x = torch.randn(num_tokens, hidden_size, dtype=dtype, device=device)
164
+ z = torch.randn(num_tokens, hidden_size, dtype=dtype, device=device)
165
+ weight = torch.randn(hidden_size, dtype=dtype, device=device)
166
+ bias = None if is_rms_norm else torch.randn(hidden_size, dtype=dtype, device=device)
167
+ eps = 1e-6
168
+
169
+ # Run the triton kernel
170
+ out, mean, rstd = layer_norm_fwd(
171
+ x,
172
+ weight,
173
+ bias,
174
+ eps,
175
+ z=z,
176
+ norm_before_gate=norm_before_gate,
177
+ is_rms_norm=is_rms_norm,
178
+ )
179
+
180
+ # Run reference implementation
181
+ ref_out = layer_norm_ref(
182
+ x,
183
+ weight,
184
+ bias,
185
+ z=z,
186
+ eps=eps,
187
+ norm_before_gate=norm_before_gate,
188
+ is_rms_norm=is_rms_norm,
189
+ )
190
+
191
+ # Check outputs
192
+ assert out.shape == x.shape
193
+ assert out.dtype == x.dtype
194
+ torch.testing.assert_close(out, ref_out, atol=1e-2, rtol=1e-2)
195
+
196
+
197
+@pytest.mark.parametrize("num_tokens", [128, 512])
198
+@pytest.mark.parametrize("hidden_size", [512, 1024])
199
+@pytest.mark.parametrize("group_size", [64, 128, 256])
200
+@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
201
+@pytest.mark.parametrize("is_rms_norm", IS_RMS_NORM)
202
+@torch.inference_mode()
203
+def test_layer_norm_fwd_with_groups(
204
+ num_tokens: int,
205
+ hidden_size: int,
206
+ group_size: int,
207
+ dtype: torch.dtype,
208
+ is_rms_norm: bool,
209
+) -> None:
210
+ """Test layer norm forward pass with group normalization."""
211
+ if hidden_size % group_size != 0:
212
+ pytest.skip(
213
+ f"hidden_size {hidden_size} not divisible by group_size {group_size}"
214
+ )
215
+
216
+ current_platform.seed_everything(42)
217
+ device = torch.device("cuda:0")
218
+
219
+ # Create inputs
220
+ x = torch.randn(num_tokens, hidden_size, dtype=dtype, device=device)
221
+ weight = torch.randn(hidden_size, dtype=dtype, device=device)
222
+ bias = None if is_rms_norm else torch.randn(hidden_size, dtype=dtype, device=device)
223
+ eps = 1e-6
224
+
225
+ ngroups = hidden_size // group_size
226
+
227
+ # Run the triton kernel
228
+ out, mean, rstd = layer_norm_fwd(
229
+ x, weight, bias, eps, z=None, group_size=group_size, is_rms_norm=is_rms_norm
230
+ )
231
+
232
+ # Run reference implementation
233
+ ref_out = layer_norm_ref(
234
+ x, weight, bias, z=None, eps=eps, group_size=group_size, is_rms_norm=is_rms_norm
235
+ )
236
+
237
+ # Check outputs
238
+ assert out.shape == x.shape
239
+ assert out.dtype == x.dtype
240
+ torch.testing.assert_close(out, ref_out, atol=1e-2, rtol=1e-2)
241
+
242
+ # Check mean and rstd shapes for groups
243
+ if not is_rms_norm:
244
+ assert mean.shape == (ngroups * num_tokens,)
245
+ assert rstd.shape == (ngroups * num_tokens,)
246
+
247
+
248
+@pytest.mark.parametrize("num_tokens", [7, 63, 128, 513, 1024, 2049])
249
+@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
250
+@torch.inference_mode()
251
+def test_layer_norm_rows_per_block(
252
+ num_tokens: int,
253
+ dtype: torch.dtype,
254
+) -> None:
255
+ """Test that rows_per_block logic works correctly for various M sizes."""
256
+ current_platform.seed_everything(42)
257
+ device = torch.device("cuda:0")
258
+ hidden_size = 1024
259
+
260
+ # Create inputs
261
+ x = torch.randn(num_tokens, hidden_size, dtype=dtype, device=device)
262
+ weight = torch.randn(hidden_size, dtype=dtype, device=device)
263
+ bias = torch.randn(hidden_size, dtype=dtype, device=device)
264
+ eps = 1e-6
265
+
266
+ # Run the triton kernel
267
+ out, mean, rstd = layer_norm_fwd(x, weight, bias, eps, z=None, is_rms_norm=False)
268
+
269
+ # Run reference implementation
270
+ ref_out = layer_norm_ref(x, weight, bias, z=None, eps=eps, is_rms_norm=False)
271
+
272
+ # Check outputs
273
+ torch.testing.assert_close(out, ref_out, atol=1e-2, rtol=1e-2)
274
+
275
+
276
+@pytest.mark.parametrize("dtype", [torch.bfloat16])
277
+@torch.inference_mode()
278
+def test_strided_input(dtype: torch.dtype) -> None:
279
+ """Test that the kernel handles non-contiguous (strided)
280
+ inputs correctly."""
281
+ current_platform.seed_everything(42)
282
+ device = torch.device("cuda:0")
283
+ num_tokens = 128
284
+ hidden_size = 1024
285
+
286
+ # Create a larger tensor and take a strided slice
287
+ x_large = torch.randn(num_tokens, hidden_size * 2, dtype=dtype, device=device)
288
+ x = x_large[:, :hidden_size]
289
+
290
+ # Make it contiguous for the kernel
291
+ x_contiguous = x.contiguous()
292
+
293
+ weight = torch.randn(hidden_size, dtype=dtype, device=device)
294
+ bias = torch.randn(hidden_size, dtype=dtype, device=device)
295
+ eps = 1e-6
296
+
297
+ # Run the triton kernel with contiguous input
298
+ out, mean, rstd = layer_norm_fwd(
299
+ x_contiguous, weight, bias, eps, z=None, is_rms_norm=False
300
+ )
301
+
302
+ # Run reference implementation
303
+ ref_out = layer_norm_ref(
304
+ x_contiguous, weight, bias, z=None, eps=eps, is_rms_norm=False
305
+ )
306
+
307
+ # Check outputs
308
+ torch.testing.assert_close(out, ref_out, atol=1e-2, rtol=1e-2)
309
+
310
+
311
+@pytest.mark.parametrize("num_tokens", [1, 128, 2048])
312
+@pytest.mark.parametrize("hidden_size", [768, 4096])
313
+@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
314
+@torch.inference_mode()
315
+def test_output_buffer_provided(
316
+ num_tokens: int,
317
+ hidden_size: int,
318
+ dtype: torch.dtype,
319
+) -> None:
320
+ """Test that the kernel works when an output buffer is provided."""
321
+ current_platform.seed_everything(42)
322
+ device = torch.device("cuda:0")
323
+
324
+ # Create inputs
325
+ x = torch.randn(num_tokens, hidden_size, dtype=dtype, device=device)
326
+ weight = torch.randn(hidden_size, dtype=dtype, device=device)
327
+ bias = torch.randn(hidden_size, dtype=dtype, device=device)
328
+ eps = 1e-6
329
+
330
+ # Pre-allocate output buffer
331
+ out_buffer = torch.empty_like(x)
332
+
333
+ # Run the triton kernel with provided output
334
+ out, mean, rstd = layer_norm_fwd(
335
+ x, weight, bias, eps, z=None, out=out_buffer, is_rms_norm=False
336
+ )
337
+
338
+ # Check that the provided buffer was used
339
+ assert out.data_ptr() == out_buffer.data_ptr()
340
+
341
+ # Run reference implementation
342
+ ref_out = layer_norm_ref(x, weight, bias, z=None, eps=eps, is_rms_norm=False)
343
+
344
+ # Check outputs
345
+ torch.testing.assert_close(out, ref_out, atol=1e-2, rtol=1e-2)
346
+
347
+
348
+@pytest.mark.parametrize(
349
+ "shape",
350
+ [
351
+ (4, 16, 1024), # 3D tensor
352
+ (2, 8, 512, 256), # 4D tensor
353
+ ],
354
+)
355
+@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
356
+@torch.inference_mode()
357
+def test_multidimensional_input(
358
+ shape: tuple,
359
+ dtype: torch.dtype,
360
+) -> None:
361
+ """Test that the autograd function handles multidimensional inputs."""
362
+ current_platform.seed_everything(42)
363
+ device = torch.device("cuda:0")
364
+ hidden_size = shape[-1]
365
+
366
+ # Create inputs
367
+ x = torch.randn(*shape, dtype=dtype, device=device)
368
+ weight = torch.randn(hidden_size, dtype=dtype, device=device)
369
+ bias = torch.randn(hidden_size, dtype=dtype, device=device)
370
+ eps = 1e-6
371
+
372
+ # Run through autograd function
373
+ out = layernorm_fn(x, weight, bias, z=None, eps=eps)
374
+
375
+ # Run reference implementation
376
+ ref_out = layer_norm_ref(x, weight, bias, z=None, eps=eps, is_rms_norm=False)
377
+
378
+ # Check outputs
379
+ assert out.shape == x.shape
380
+ torch.testing.assert_close(out, ref_out, atol=1e-2, rtol=1e-2)
381
+
382
+
383
+if __name__ == "__main__":
384
+ # Run a quick smoke test
385
+ test_layer_norm_fwd_basic(128, 1024, torch.float16, 42, False)
386
+ test_layer_norm_fwd_with_gate(128, 1024, torch.float16, True, False)
387
+ test_layer_norm_rows_per_block(513, torch.float16)
388
+ print("All smoke tests passed!")
vllm/model_executor/layers/fla/ops/layernorm_guard.py
+87
-33
@@ -13,6 +13,7 @@
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 functools import lru_cache
17
from typing import Optional
18
19
import torch
@@ -21,6 +22,7 @@ import torch.nn.functional as F
22
from einops import rearrange
23
24
from vllm.triton_utils import tl, triton
25
+from vllm.utils import cdiv, next_power_of_2
26
27
from .utils import input_guard
28
@@ -76,55 +78,103 @@ def layer_norm_fwd_kernel(
78
stride_y_row,
79
stride_z_row,
80
M, # number of rows in X
79
- N, # number of columns in X
81
+ N: tl.constexpr, # number of columns in X
82
eps, # epsilon to avoid division by zero
83
BLOCK_N: tl.constexpr,
84
+ ROWS_PER_BLOCK: tl.constexpr,
85
HAS_BIAS: tl.constexpr,
86
HAS_Z: tl.constexpr,
87
NORM_BEFORE_GATE: tl.constexpr,
88
IS_RMS_NORM: tl.constexpr,
89
):
87
- # Map the program id to the row of X and Y it should compute.
88
- row = tl.program_id(0)
90
+ # Map the program id to the starting row of X and Y it should compute.
91
+ row_start = tl.program_id(0) * ROWS_PER_BLOCK
92
group = tl.program_id(1)
90
- X += row * stride_x_row + group * N
91
- Y += row * stride_y_row + group * N
92
- if HAS_Z:
93
- Z += row * stride_z_row + group * N
94
- if not IS_RMS_NORM:
95
- Mean += group * M
96
- Rstd += group * M
97
- W += group * N
98
- if HAS_BIAS:
99
- B += group * N
100
- # Compute mean and variance
93
+
94
+ # Create 2D tile: [ROWS_PER_BLOCK, BLOCK_N]
95
+ rows = row_start + tl.arange(0, ROWS_PER_BLOCK)
96
cols = tl.arange(0, BLOCK_N)
102
- x = tl.load(X + cols, mask=cols < N, other=0.0).to(tl.float32)
97
+
98
+ # Compute offsets for 2D tile
99
+ row_offsets = rows[:, None] * stride_x_row
100
+ col_offsets = cols[None, :] + group * N
101
+
102
+ # Base pointers
103
+ X_base = X + row_offsets + col_offsets
104
+ Y_base = Y + rows[:, None] * stride_y_row + col_offsets
105
+
106
+ # Create mask for valid rows and columns
107
+ row_mask = rows[:, None] < M
108
+ col_mask = cols[None, :] < N
109
+ mask = row_mask & col_mask
110
+
111
+ # Load input data with 2D tile
112
+ x = tl.load(X_base, mask=mask, other=0.0).to(tl.float32)
113
+
114
if HAS_Z and not NORM_BEFORE_GATE:
104
- z = tl.load(Z + cols, mask=cols < N).to(tl.float32)
115
+ Z_base = Z + rows[:, None] * stride_z_row + col_offsets
116
+ z = tl.load(Z_base, mask=mask, other=0.0).to(tl.float32)
117
x *= z * tl.sigmoid(z)
118
+
119
+ # Compute mean and variance per row (reduce along axis 1)
120
if not IS_RMS_NORM:
107
- mean = tl.sum(x, axis=0) / N
108
- tl.store(Mean + row, mean)
109
- xbar = tl.where(cols < N, x - mean, 0.0)
110
- var = tl.sum(xbar * xbar, axis=0) / N
121
+ mean = tl.sum(x, axis=1) / N # Shape: [ROWS_PER_BLOCK]
122
+ # Store mean for each row
123
+ mean_offsets = group * M + rows
124
+ mean_mask = rows < M
125
+ tl.store(Mean + mean_offsets, mean, mask=mean_mask)
126
+ # Broadcast mean back to 2D for subtraction
127
+ xbar = tl.where(mask, x - mean[:, None], 0.0)
128
+ var = tl.sum(xbar * xbar, axis=1) / N # Shape: [ROWS_PER_BLOCK]
129
else:
112
- xbar = tl.where(cols < N, x, 0.0)
113
- var = tl.sum(xbar * xbar, axis=0) / N
114
- rstd = 1 / tl.sqrt(var + eps)
115
- tl.store(Rstd + row, rstd)
116
- # Normalize and apply linear transformation
117
- mask = cols < N
118
- w = tl.load(W + cols, mask=mask).to(tl.float32)
130
+ xbar = tl.where(mask, x, 0.0)
131
+ var = tl.sum(xbar * xbar, axis=1) / N # Shape: [ROWS_PER_BLOCK]
132
+ mean = 0.0 # Placeholder for RMS norm
133
+
134
+ rstd = tl.rsqrt(var + eps) # Shape: [ROWS_PER_BLOCK]
135
+
136
+ # Store rstd for each row
137
+ rstd_offsets = group * M + rows
138
+ rstd_mask = rows < M
139
+ tl.store(Rstd + rstd_offsets, rstd, mask=rstd_mask)
140
+
141
+ # Load weights and biases (broadcast across rows)
142
+ w_offsets = cols + group * N
143
+ w_mask = cols < N
144
+ w = tl.load(W + w_offsets, mask=w_mask, other=0.0).to(tl.float32)
145
+
146
if HAS_BIAS:
120
- b = tl.load(B + cols, mask=mask).to(tl.float32)
121
- x_hat = (x - mean) * rstd if not IS_RMS_NORM else x * rstd
122
- y = x_hat * w + b if HAS_BIAS else x_hat * w
147
+ b = tl.load(B + w_offsets, mask=w_mask, other=0.0).to(tl.float32)
148
+
149
+ # Normalize and apply linear transformation
150
+ if not IS_RMS_NORM:
151
+ x_hat = (x - mean[:, None]) * rstd[:, None]
152
+ else:
153
+ x_hat = x * rstd[:, None]
154
+
155
+ y = x_hat * w[None, :] + b[None, :] if HAS_BIAS else x_hat * w[None, :]
156
+
157
if HAS_Z and NORM_BEFORE_GATE:
124
- z = tl.load(Z + cols, mask=mask).to(tl.float32)
158
+ Z_base = Z + rows[:, None] * stride_z_row + col_offsets
159
+ z = tl.load(Z_base, mask=mask, other=0.0).to(tl.float32)
160
y *= z * tl.sigmoid(z)
161
+
162
# Write output
127
- tl.store(Y + cols, y, mask=mask)
163
+ tl.store(Y_base, y, mask=mask)
164
+
165
+
166
+@lru_cache
167
+def _get_sm_count(device: torch.device) -> int:
168
+ """Get and cache the SM count for a given device."""
169
+ props = torch.cuda.get_device_properties(device)
170
+ return props.multi_processor_count
171
+
172
+
173
+def calc_rows_per_block(M: int, device: torch.device) -> int:
174
+ sm_count = _get_sm_count(device)
175
+ rows_per_block = next_power_of_2(cdiv(M, 2 * sm_count))
176
+ rows_per_block = min(rows_per_block, 4)
177
+ return rows_per_block
178
179
180
def layer_norm_fwd(
@@ -171,7 +221,10 @@ def layer_norm_fwd(
221
raise RuntimeError("This layer norm doesn't support feature dim >= 64KB.")
222
# heuristics for number of warps
223
num_warps = min(max(BLOCK_N // 256, 1), 8)
174
- grid = (M, ngroups)
224
+ # Calculate rows per block based on SM count
225
+ rows_per_block = calc_rows_per_block(M, x.device)
226
+ # Update grid to use rows_per_block
227
+ grid = (cdiv(M, rows_per_block), ngroups)
228
layer_norm_fwd_kernel[grid](
229
x,
230
out,
@@ -187,6 +240,7 @@ def layer_norm_fwd(
240
group_size,
241
eps,
242
BLOCK_N=BLOCK_N,
243
+ ROWS_PER_BLOCK=rows_per_block,
244
NORM_BEFORE_GATE=norm_before_gate,
245
IS_RMS_NORM=is_rms_norm,
246
num_warps=num_warps,