[Perf] Prefetch the weight before the PDL wait in fused_q_kv_rmsnorm (#55020)
Signed-off-by: Yongye Zhu <zyy1102000@gmail.com> Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
Yongye Zhu committed
Sep 3, 2026 at 15:40 UTC
bc2ee480738d7dcc558262a0c6d81956b515b050
1 file changed
+12
-5
vllm/models/common/ops/fused_qk_rmsnorm.py
+12
-5
@@ -43,19 +43,23 @@ def _fused_q_kv_rmsnorm_kernel(
43
weight_ptr = kv_weight_ptr
44
row_out = kv_out_ptr + token_idx * kv_out_stride
45
46
- if launch_pdl:
47
- tl.extra.cuda.gdc_wait()
48
- tl.extra.cuda.gdc_launch_dependents()
49
-
46
# RMSNorm in fp32 throughout — matches csrc/layernorm_kernels.cu's
47
# `(scalar_t)(x * s_variance * w)` and DeepseekV4's compressor kernel, which
48
# keep x, rrms, and w all in fp32 and perform a single cast at store.
49
block = tl.arange(0, BLOCK_SIZE)
50
mask = block < SIZE
51
+ # The weight load does not depend on the producer's output, so issue it
52
+ # before the PDL wait: gamma streams in while the producer finishes,
53
+ # leaving a single dependent global round trip (x) after the wait.
54
+ w = tl.load(weight_ptr + block, mask=mask, other=0.0).to(tl.float32)
55
+
56
+ if launch_pdl:
57
+ tl.extra.cuda.gdc_wait()
58
+ tl.extra.cuda.gdc_launch_dependents()
59
+
60
x = tl.load(row_in + block, mask=mask, other=0.0).to(tl.float32)
61
variance = tl.sum(x * x, axis=0) / SIZE
62
rrms = tl.rsqrt(variance + eps)
58
- w = tl.load(weight_ptr + block, mask=mask, other=0.0).to(tl.float32)
63
y = x * rrms * w
64
tl.store(row_out + block, y.to(row_out.dtype.element_ty), mask=mask)
65
@@ -106,5 +110,8 @@ def fused_q_kv_rmsnorm(
110
KV_SIZE=kv_size,
111
BLOCK_SIZE=block_size,
112
launch_pdl=current_platform.is_arch_support_pdl(),
113
+ # One vectorized 128-bit access per thread for K3/DSv4-sized rows,
114
+ # instead of a serial per-thread chain at the default 4 warps.
115
+ num_warps=8 if block_size >= 2048 else 4,
116
)
117
return qr_out, kv_out