[DFlash] Fuse precompute kv per-layer rmsnorms (#46761)
Signed-off-by: Giancarlo Delfin <gdelfin@inferact.ai> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Giancarlo Delfin committed
Jun 25, 2026 at 21:32 UTC
02a1f23711c5bdbff81eb8a610dde39e1141d036
3 files changed
+108
-18
csrc/libtorch_stable/layernorm_kernels.cu
+25
-8
@@ -20,20 +20,27 @@ __global__ void rms_norm_kernel(
20
const int64_t input_stride_d4, // input.stride(-4)
21
const int64_t input_shape_d2, // input.size(-2)
22
const int64_t input_shape_d3, // input.size(-3)
23
- const scalar_t* __restrict__ weight, // [hidden_size], null if !HasWeight
23
+ const scalar_t* __restrict__ weight, // [hidden_size] or
24
+ // [num_groups, hidden_size];
25
+ // null if !HasWeight
26
+ const int64_t weight_stride, // 0 or weight.stride(0)
27
const float epsilon, const int num_tokens, const int hidden_size) {
28
__shared__ float s_variance;
29
float variance = 0.0f;
30
const scalar_t* input_row;
31
+ const scalar_t* weight_row;
32
+ int64_t weight_row_off = 0;
33
if constexpr (NUM_DIMS == 2) {
34
// 2D for layernorm normal case [batch_size, hidden]
35
input_row = input + blockIdx.x * input_stride_d2;
36
+ weight_row = weight + blockIdx.x * weight_stride;
37
} else if constexpr (NUM_DIMS == 3) {
38
// 3D for q/k norm [batch_size, num_heads, head_size]
39
int batch_idx = blockIdx.x / input_shape_d2;
40
int head_idx = blockIdx.x % input_shape_d2;
41
input_row =
42
input + batch_idx * input_stride_d3 + head_idx * input_stride_d2;
43
+ weight_row = weight + batch_idx * weight_stride;
44
} else if constexpr (NUM_DIMS == 4) {
45
// 4D for transformers model_impl qk norm [batch, seq, head, head_dim]
46
int batch_idx = blockIdx.x / (input_shape_d3 * input_shape_d2);
@@ -42,6 +49,7 @@ __global__ void rms_norm_kernel(
49
int head_idx = remaining % input_shape_d2;
50
input_row = input + batch_idx * input_stride_d4 +
51
seq_idx * input_stride_d3 + head_idx * input_stride_d2;
52
+ weight_row = weight + batch_idx * weight_stride;
53
}
54
55
auto vec_op = [&variance](const vec_n_t<scalar_t, VEC_SIZE>& vec) {
@@ -69,7 +77,7 @@ __global__ void rms_norm_kernel(
77
78
scalar_t* out_row = out + blockIdx.x * hidden_size;
79
auto* v_in = reinterpret_cast<const vec_n_t<scalar_t, VEC_SIZE>*>(input_row);
72
- auto* v_w = reinterpret_cast<const vec_n_t<scalar_t, VEC_SIZE>*>(weight);
80
+ auto* v_w = reinterpret_cast<const vec_n_t<scalar_t, VEC_SIZE>*>(weight_row);
81
auto* v_out = reinterpret_cast<vec_n_t<scalar_t, VEC_SIZE>*>(out_row);
82
for (int i = threadIdx.x; i < hidden_size / VEC_SIZE; i += blockDim.x) {
83
vec_n_t<scalar_t, VEC_SIZE> dst;
@@ -211,15 +219,24 @@ fused_add_rms_norm_kernel(
219
220
void rms_norm(torch::stable::Tensor& out, // [..., hidden_size]
221
torch::stable::Tensor& input, // [..., hidden_size]
214
- std::optional<torch::stable::Tensor> weight, // [hidden_size]
215
- double epsilon) {
222
+ std::optional<torch::stable::Tensor> weight, double epsilon) {
223
STD_TORCH_CHECK(out.is_contiguous());
224
if (input.stride(-1) != 1) {
225
input = torch::stable::contiguous(input);
226
}
227
STD_TORCH_CHECK(input.stride(-1) == 1);
228
+ int64_t weight_stride = 0;
229
if (weight.has_value()) {
230
STD_TORCH_CHECK(weight->is_contiguous());
231
+ if (weight->dim() == 1) {
232
+ STD_TORCH_CHECK(weight->size(0) == input.size(-1));
233
+ } else if (weight->dim() == 2) {
234
+ STD_TORCH_CHECK(weight->size(0) == input.size(0));
235
+ STD_TORCH_CHECK(weight->size(-1) == input.size(-1));
236
+ weight_stride = weight->stride(0);
237
+ } else {
238
+ STD_TORCH_CHECK(false, "rms_norm weight must be 1D or 2D");
239
+ }
240
}
241
242
int hidden_size = input.size(-1);
@@ -256,16 +273,16 @@ void rms_norm(torch::stable::Tensor& out, // [..., hidden_size]
273
out.mutable_data_ptr<scalar_t>(),
274
input.const_data_ptr<scalar_t>(), input_stride_d2,
275
input_stride_d3, input_stride_d4, input_shape_d2,
259
- input_shape_d3, weight_ptr, epsilon, num_tokens,
260
- hidden_size);
276
+ input_shape_d3, weight_ptr, weight_stride, epsilon,
277
+ num_tokens, hidden_size);
278
} else {
279
vllm::rms_norm_kernel<scalar_t, vec_size, tensor_rank, false>
280
<<<grid, block, 0, stream>>>(
281
out.mutable_data_ptr<scalar_t>(),
282
input.const_data_ptr<scalar_t>(), input_stride_d2,
283
input_stride_d3, input_stride_d4, input_shape_d2,
267
- input_shape_d3, weight_ptr, epsilon, num_tokens,
268
- hidden_size);
284
+ input_shape_d3, weight_ptr, /*weight_stride=*/0, epsilon,
285
+ num_tokens, hidden_size);
286
}
287
});
288
});
tests/kernels/core/test_batched_weight_rms_norm.py
new
+70
@@ -0,0 +1,70 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Tests for the batched-weight RMS norm kernel (vllm._custom_ops.rms_norm).
4
+
5
+``rms_norm`` can use the outermost input batch index to select the corresponding
6
+weight row. The result must match that of looping ``rms_norm`` over that dimension.
7
+"""
8
+
9
+import pytest
10
+import torch
11
+
12
+from vllm import _custom_ops as ops
13
+from vllm.platforms import current_platform
14
+from vllm.utils.torch_utils import set_random_seed
15
+
16
+pytestmark = pytest.mark.skipif(
17
+ not current_platform.is_cuda_alike(),
18
+ reason="rms_norm requires a CUDA/ROCm device",
19
+)
20
+
21
+
22
+@pytest.mark.parametrize(
23
+ "shape",
24
+ [
25
+ (28, 17, 128), # 3D: [num_rows, tokens, hidden]
26
+ (1, 5, 2, 128), # 4D: single row (edge case)
27
+ (28, 13, 8, 128), # 4D: [L, num_ctx, nkv, hd] (DFlash K-norm)
28
+ (6, 3, 4, 769), # 4D: non-power-of-two hidden size
29
+ ],
30
+)
31
+@pytest.mark.parametrize("dtype", [torch.half, torch.bfloat16, torch.float])
32
+@pytest.mark.parametrize("seed", [42])
33
+@torch.inference_mode()
34
+def test_rms_norm_matches_loop(
35
+ shape: tuple[int, ...], dtype: torch.dtype, seed: int
36
+) -> None:
37
+ set_random_seed(seed)
38
+ torch.set_default_device("cuda")
39
+
40
+ num_rows, hidden = shape[0], shape[-1]
41
+ eps = 1e-6
42
+
43
+ x = torch.randn(*shape, dtype=dtype) * 0.1
44
+ # Distinct weight per row so that a wrong row index would be caught.
45
+ weight = torch.randn(num_rows, hidden, dtype=dtype) * 0.1 + 1.0
46
+
47
+ # Reference batched-weight rms norm.
48
+ out_ref = torch.empty_like(x)
49
+ for i in range(x.shape[0]):
50
+ ops.rms_norm(out_ref[i], x[i], weight[i], eps)
51
+
52
+ out = torch.empty_like(x)
53
+ ops.rms_norm(out, x, weight, eps)
54
+
55
+ # Expect bitwise-identical results.
56
+ torch.testing.assert_close(out, out_ref, atol=0, rtol=0)
57
+
58
+
59
+@torch.inference_mode()
60
+def test_rms_norm_validates_shapes() -> None:
61
+ torch.set_default_device("cuda")
62
+
63
+ x = torch.randn(4, 8, 128, dtype=torch.float)
64
+ out = torch.empty_like(x)
65
+ # Expect num rows mismatch.
66
+ with pytest.raises(RuntimeError):
67
+ ops.rms_norm(out, x, torch.randn(3, 128), 1e-6)
68
+ # Expect hidden size mismatch.
69
+ with pytest.raises(RuntimeError):
70
+ ops.rms_norm(out, x, torch.randn(4, 64), 1e-6)
vllm/model_executor/models/qwen3_dflash.py
+13
-10
@@ -309,8 +309,11 @@ class DFlashQwen3Model(nn.Module):
309
else:
310
self._fused_kv_bias = None
311
312
- # K-norm weights: list of [head_dim] tensors, one per layer.
313
- self._k_norm_weights = [a.k_norm.weight.data for a in layers_attn]
312
+ # K-norm weights stacked into one contiguous [num_layers, head_dim]
313
+ # tensor so the per-layer K-norm runs as a single grouped kernel.
314
+ self._k_norm_weights = torch.stack(
315
+ [a.k_norm.weight.data for a in layers_attn], dim=0
316
+ ).contiguous()
317
318
# RoPE parameters
319
self._rope_head_size = attn0.rotary_emb.head_size
@@ -392,15 +395,15 @@ class DFlashQwen3Model(nn.Module):
395
all_k = all_kv[0] # [L, num_ctx, nkv, hd], contiguous
396
all_v = all_kv[1] # [L, num_ctx, nkv, hd], contiguous
397
395
- # --- Per-layer RMSNorm K (3D: [num_ctx, nkv, hd] per layer) ---
398
+ # --- Grouped RMSNorm K across all layers ([L, num_ctx, nkv, hd]) ---
399
+ # The weight is selected per layer by the outermost (layer) index.
400
all_k_normed = torch.empty_like(all_k)
397
- for i in range(L):
398
- ops.rms_norm(
399
- all_k_normed[i],
400
- all_k[i],
401
- self._k_norm_weights[i],
402
- self._rms_norm_eps,
403
- )
401
+ ops.rms_norm(
402
+ all_k_normed,
403
+ all_k,
404
+ self._k_norm_weights,
405
+ self._rms_norm_eps,
406
+ )
407
408
# --- Fused RoPE across all layers ---
409
# View as [L * num_ctx, kv] so RoPE sees one big batch (no copy).