[Kimi-K3][Perf] Prefetch ll_bf16 router weights for M=1 (#53524)
Signed-off-by: ming <ming@inferact.ai> Co-authored-by: OpenAI Codex <codex@openai.com>
mingg26 committed
Sep 1, 2026 at 14:40 UTC
259a209bfa8a561e87b72e9678fd3d3a35519ce6
3 files changed
+99
-8
tests/kernels/test_ll_bf16_gemm.py
+23
@@ -69,6 +69,29 @@ def _gemm(a, b):
69
return ll_bf16_gemm(a, b)
70
71
72
+def test_c1_pdl_kernel_is_selected_automatically(monkeypatch):
73
+ from vllm.model_executor.kernels.linear.cute_dsl import ll_bf16
74
+
75
+ calls = []
76
+
77
+ def default_kernel(hidden_states, router_weight, output_dtype):
78
+ calls.append(("default", hidden_states.shape[0]))
79
+ return hidden_states
80
+
81
+ def c1_pdl_kernel(hidden_states, router_weight, output_dtype):
82
+ calls.append(("c1_pdl", hidden_states.shape[0]))
83
+ return hidden_states
84
+
85
+ monkeypatch.setattr(ll_bf16, "ll_bf16_gemm_kernel", default_kernel)
86
+ monkeypatch.setattr(ll_bf16, "ll_bf16_gemm_c1_pdl_kernel", c1_pdl_kernel)
87
+ weight = torch.empty(4, 8, device="cuda", dtype=torch.bfloat16)
88
+
89
+ ll_bf16.ll_bf16_gemm(torch.empty(1, 8, device="cuda", dtype=torch.bfloat16), weight)
90
+ ll_bf16.ll_bf16_gemm(torch.empty(2, 8, device="cuda", dtype=torch.bfloat16), weight)
91
+
92
+ assert calls == [("c1_pdl", 1), ("default", 2)]
93
+
94
+
95
# ===== Shapes =====
96
97
SHAPES = [
vllm/model_executor/kernels/linear/cute_dsl/_ll_bf16_dotprod.py
+61
-5
@@ -48,6 +48,7 @@ class LLBf16Dotprod:
48
main_vec_width: int = 8,
49
tail_vec_width: int = 4,
50
use_pdl: bool = False,
51
+ prefetch_pdl_weights: bool = False,
52
):
53
"""Initialize the dot-product kernel configuration.
54
@@ -70,8 +71,10 @@ class LLBf16Dotprod:
71
self.main_vec_width = main_vec_width
72
self.tail_vec_width = tail_vec_width
73
self.use_pdl = use_pdl
74
+ self.prefetch_pdl_weights = prefetch_pdl_weights
75
self.num_warps = bs // cute.arch.WARP_SIZE
76
self._init_k_tiles(k)
77
+ self.main_prefetch_tiles = min(self.main_tiles, 8)
78
79
def _vectorized_elems(self, k_extent: int, vec_width: int) -> int:
80
vector_tile = vec_width * self.bs
@@ -117,6 +120,33 @@ class LLBf16Dotprod:
120
for v in cutlass.range_constexpr(vec_width):
121
acc[m] = acc[m] + ar[v].to(cutlass.Float32) * br_f32[v]
122
123
+ @cute.jit
124
+ def _vector_dotprod_prefetched(
125
+ self,
126
+ acc: cute.Tensor,
127
+ tA: cute.Tensor,
128
+ tB: cute.Tensor,
129
+ rB: cute.Tensor,
130
+ M: cutlass.Constexpr,
131
+ num_tiles: cutlass.Constexpr,
132
+ prefetch_tiles: cutlass.Constexpr,
133
+ ):
134
+ for tile in cutlass.range_constexpr(num_tiles):
135
+ if const_expr(tile < prefetch_tiles):
136
+ br_f32 = rB[tile, None].load().to(cutlass.Float32)
137
+ else:
138
+ bt = tB[None, tile]
139
+ br = cute.make_rmem_tensor_like(bt)
140
+ cute.autovec_copy(bt, br)
141
+ br_f32 = br.load().to(cutlass.Float32)
142
+ for m in cutlass.range_constexpr(M):
143
+ at = tA[m, None, tile]
144
+ ar = cute.make_rmem_tensor_like(at)
145
+ cute.autovec_copy(at, ar)
146
+ vec_width: cutlass.Constexpr = cute.size(ar)
147
+ for v in cutlass.range_constexpr(vec_width):
148
+ acc[m] = acc[m] + ar[v].to(cutlass.Float32) * br_f32[v]
149
+
150
def _make_thread_vector_slice(
151
self,
152
gA_vec: cute.Tensor,
@@ -240,17 +270,43 @@ class LLBf16Dotprod:
270
acc = cute.make_rmem_tensor((M,), cutlass.Float32)
271
acc.fill(0.0)
272
243
- if const_expr(self.use_pdl):
244
- cute.arch.griddepcontrol_wait()
245
-
246
- # 128-bit vectorized main loop
273
+ # In the K3 C=1 specialization, the static weight stripe is independent
274
+ # of the producer and can be resident while PDL waits for activation A.
275
if const_expr(k_main_elems > 0):
276
gA_main = self._make_k_slice(gA, 0, k_main_elems)
277
gB_main = self._make_k_slice(gB, 0, k_main_elems)
278
gA_vec = cute.logical_divide(gA_main, (None, main_vec_width))
279
gB_vec = cute.logical_divide(gB_main, (None, main_vec_width))
280
tA, tB = self._make_thread_vector_slice(gA_vec, gB_vec, tidx, n_idx, bs)
253
- self._vector_dotprod(acc, tA, tB, M, main_tiles, 16)
281
+ if const_expr(self.use_pdl and self.prefetch_pdl_weights):
282
+ prefetch_tiles: cutlass.Constexpr = self.main_prefetch_tiles
283
+ prefetched_b = cute.make_rmem_tensor(
284
+ cute.make_layout(
285
+ (prefetch_tiles, main_vec_width),
286
+ stride=(main_vec_width, 1),
287
+ ),
288
+ cutlass.BFloat16,
289
+ )
290
+ for tile in cutlass.range_constexpr(prefetch_tiles):
291
+ cute.autovec_copy(tB[None, tile], prefetched_b[tile, None])
292
+
293
+ if const_expr(self.use_pdl):
294
+ cute.arch.griddepcontrol_wait()
295
+
296
+ # 128-bit vectorized main loop
297
+ if const_expr(k_main_elems > 0):
298
+ if const_expr(self.use_pdl and self.prefetch_pdl_weights):
299
+ self._vector_dotprod_prefetched(
300
+ acc,
301
+ tA,
302
+ tB,
303
+ prefetched_b,
304
+ M,
305
+ main_tiles,
306
+ prefetch_tiles,
307
+ )
308
+ else:
309
+ self._vector_dotprod(acc, tA, tB, M, main_tiles, 16)
310
311
# 64-bit vectorized tail (K remainder after main loop)
312
if const_expr(k_tail_elems > 0):
vllm/model_executor/kernels/linear/cute_dsl/ll_bf16.py
+15
-3
@@ -133,7 +133,8 @@ class LLBf16Gemm:
133
split_k: int = 0
134
num_stages: int = 0
135
136
- def __init__(self) -> None:
136
+ def __init__(self, *, prefetch_pdl_weights: bool = False) -> None:
137
+ self._prefetch_pdl_weights = prefetch_pdl_weights
138
# Dot-prod: keyed on (M, K, bs), because M and K are Constexpr.
139
self._compiled_cache: dict[tuple[int, int, int], Any] = {}
140
# Split-K: keyed on (split_k, num_stages), fully shape-dynamic.
@@ -217,7 +218,12 @@ class LLBf16Gemm:
218
N=N,
219
divisibility=stride_divisibility,
220
)
220
- gemm = LLBf16Dotprod(k=compile_key.K, bs=compile_key.bs, use_pdl=_use_pdl())
221
+ gemm = LLBf16Dotprod(
222
+ k=compile_key.K,
223
+ bs=compile_key.bs,
224
+ use_pdl=_use_pdl(),
225
+ prefetch_pdl_weights=self._prefetch_pdl_weights,
226
+ )
227
compiled = cute.compile(
228
gemm,
229
hidden_states,
@@ -324,6 +330,7 @@ class LLBf16Gemm:
330
331
332
ll_bf16_gemm_kernel = LLBf16Gemm()
333
+ll_bf16_gemm_c1_pdl_kernel = LLBf16Gemm(prefetch_pdl_weights=True)
334
335
336
def ll_bf16_gemm(
@@ -331,4 +338,9 @@ def ll_bf16_gemm(
338
router_weight: torch.Tensor,
339
output_dtype: torch.dtype = torch.float32,
340
) -> torch.Tensor:
334
- return ll_bf16_gemm_kernel(hidden_states, router_weight, output_dtype)
341
+ kernel = (
342
+ ll_bf16_gemm_c1_pdl_kernel
343
+ if hidden_states.shape[0] == 1
344
+ else ll_bf16_gemm_kernel
345
+ )
346
+ return kernel(hidden_states, router_weight, output_dtype)