From 259a209bfa8a561e87b72e9678fd3d3a35519ce6 Mon Sep 17 00:00:00 2001 From: mingg26 Date: Tue, 1 Sep 2026 14:40:19 -0700 Subject: [PATCH] [Kimi-K3][Perf] Prefetch ll_bf16 router weights for M=1 (#53524) Signed-off-by: ming Co-authored-by: OpenAI Codex --- tests/kernels/test_ll_bf16_gemm.py | 23 +++++++ .../linear/cute_dsl/_ll_bf16_dotprod.py | 66 +++++++++++++++++-- .../kernels/linear/cute_dsl/ll_bf16.py | 18 ++++- 3 files changed, 99 insertions(+), 8 deletions(-) diff --git a/tests/kernels/test_ll_bf16_gemm.py b/tests/kernels/test_ll_bf16_gemm.py index e71c89aa0d..1c7a15d514 100644 --- a/tests/kernels/test_ll_bf16_gemm.py +++ b/tests/kernels/test_ll_bf16_gemm.py @@ -69,6 +69,29 @@ def _gemm(a, b): return ll_bf16_gemm(a, b) +def test_c1_pdl_kernel_is_selected_automatically(monkeypatch): + from vllm.model_executor.kernels.linear.cute_dsl import ll_bf16 + + calls = [] + + def default_kernel(hidden_states, router_weight, output_dtype): + calls.append(("default", hidden_states.shape[0])) + return hidden_states + + def c1_pdl_kernel(hidden_states, router_weight, output_dtype): + calls.append(("c1_pdl", hidden_states.shape[0])) + return hidden_states + + monkeypatch.setattr(ll_bf16, "ll_bf16_gemm_kernel", default_kernel) + monkeypatch.setattr(ll_bf16, "ll_bf16_gemm_c1_pdl_kernel", c1_pdl_kernel) + weight = torch.empty(4, 8, device="cuda", dtype=torch.bfloat16) + + ll_bf16.ll_bf16_gemm(torch.empty(1, 8, device="cuda", dtype=torch.bfloat16), weight) + ll_bf16.ll_bf16_gemm(torch.empty(2, 8, device="cuda", dtype=torch.bfloat16), weight) + + assert calls == [("c1_pdl", 1), ("default", 2)] + + # ===== Shapes ===== SHAPES = [ diff --git a/vllm/model_executor/kernels/linear/cute_dsl/_ll_bf16_dotprod.py b/vllm/model_executor/kernels/linear/cute_dsl/_ll_bf16_dotprod.py index 704fa7da4b..f06de6bd0c 100644 --- a/vllm/model_executor/kernels/linear/cute_dsl/_ll_bf16_dotprod.py +++ b/vllm/model_executor/kernels/linear/cute_dsl/_ll_bf16_dotprod.py @@ -48,6 +48,7 @@ class LLBf16Dotprod: main_vec_width: int = 8, tail_vec_width: int = 4, use_pdl: bool = False, + prefetch_pdl_weights: bool = False, ): """Initialize the dot-product kernel configuration. @@ -70,8 +71,10 @@ class LLBf16Dotprod: self.main_vec_width = main_vec_width self.tail_vec_width = tail_vec_width self.use_pdl = use_pdl + self.prefetch_pdl_weights = prefetch_pdl_weights self.num_warps = bs // cute.arch.WARP_SIZE self._init_k_tiles(k) + self.main_prefetch_tiles = min(self.main_tiles, 8) def _vectorized_elems(self, k_extent: int, vec_width: int) -> int: vector_tile = vec_width * self.bs @@ -117,6 +120,33 @@ class LLBf16Dotprod: for v in cutlass.range_constexpr(vec_width): acc[m] = acc[m] + ar[v].to(cutlass.Float32) * br_f32[v] + @cute.jit + def _vector_dotprod_prefetched( + self, + acc: cute.Tensor, + tA: cute.Tensor, + tB: cute.Tensor, + rB: cute.Tensor, + M: cutlass.Constexpr, + num_tiles: cutlass.Constexpr, + prefetch_tiles: cutlass.Constexpr, + ): + for tile in cutlass.range_constexpr(num_tiles): + if const_expr(tile < prefetch_tiles): + br_f32 = rB[tile, None].load().to(cutlass.Float32) + else: + bt = tB[None, tile] + br = cute.make_rmem_tensor_like(bt) + cute.autovec_copy(bt, br) + br_f32 = br.load().to(cutlass.Float32) + for m in cutlass.range_constexpr(M): + at = tA[m, None, tile] + ar = cute.make_rmem_tensor_like(at) + cute.autovec_copy(at, ar) + vec_width: cutlass.Constexpr = cute.size(ar) + for v in cutlass.range_constexpr(vec_width): + acc[m] = acc[m] + ar[v].to(cutlass.Float32) * br_f32[v] + def _make_thread_vector_slice( self, gA_vec: cute.Tensor, @@ -240,17 +270,43 @@ class LLBf16Dotprod: acc = cute.make_rmem_tensor((M,), cutlass.Float32) acc.fill(0.0) - if const_expr(self.use_pdl): - cute.arch.griddepcontrol_wait() - - # 128-bit vectorized main loop + # In the K3 C=1 specialization, the static weight stripe is independent + # of the producer and can be resident while PDL waits for activation A. if const_expr(k_main_elems > 0): gA_main = self._make_k_slice(gA, 0, k_main_elems) gB_main = self._make_k_slice(gB, 0, k_main_elems) gA_vec = cute.logical_divide(gA_main, (None, main_vec_width)) gB_vec = cute.logical_divide(gB_main, (None, main_vec_width)) tA, tB = self._make_thread_vector_slice(gA_vec, gB_vec, tidx, n_idx, bs) - self._vector_dotprod(acc, tA, tB, M, main_tiles, 16) + if const_expr(self.use_pdl and self.prefetch_pdl_weights): + prefetch_tiles: cutlass.Constexpr = self.main_prefetch_tiles + prefetched_b = cute.make_rmem_tensor( + cute.make_layout( + (prefetch_tiles, main_vec_width), + stride=(main_vec_width, 1), + ), + cutlass.BFloat16, + ) + for tile in cutlass.range_constexpr(prefetch_tiles): + cute.autovec_copy(tB[None, tile], prefetched_b[tile, None]) + + if const_expr(self.use_pdl): + cute.arch.griddepcontrol_wait() + + # 128-bit vectorized main loop + if const_expr(k_main_elems > 0): + if const_expr(self.use_pdl and self.prefetch_pdl_weights): + self._vector_dotprod_prefetched( + acc, + tA, + tB, + prefetched_b, + M, + main_tiles, + prefetch_tiles, + ) + else: + self._vector_dotprod(acc, tA, tB, M, main_tiles, 16) # 64-bit vectorized tail (K remainder after main loop) if const_expr(k_tail_elems > 0): diff --git a/vllm/model_executor/kernels/linear/cute_dsl/ll_bf16.py b/vllm/model_executor/kernels/linear/cute_dsl/ll_bf16.py index b0471b5657..564abbd84b 100644 --- a/vllm/model_executor/kernels/linear/cute_dsl/ll_bf16.py +++ b/vllm/model_executor/kernels/linear/cute_dsl/ll_bf16.py @@ -133,7 +133,8 @@ class LLBf16Gemm: split_k: int = 0 num_stages: int = 0 - def __init__(self) -> None: + def __init__(self, *, prefetch_pdl_weights: bool = False) -> None: + self._prefetch_pdl_weights = prefetch_pdl_weights # Dot-prod: keyed on (M, K, bs), because M and K are Constexpr. self._compiled_cache: dict[tuple[int, int, int], Any] = {} # Split-K: keyed on (split_k, num_stages), fully shape-dynamic. @@ -217,7 +218,12 @@ class LLBf16Gemm: N=N, divisibility=stride_divisibility, ) - gemm = LLBf16Dotprod(k=compile_key.K, bs=compile_key.bs, use_pdl=_use_pdl()) + gemm = LLBf16Dotprod( + k=compile_key.K, + bs=compile_key.bs, + use_pdl=_use_pdl(), + prefetch_pdl_weights=self._prefetch_pdl_weights, + ) compiled = cute.compile( gemm, hidden_states, @@ -324,6 +330,7 @@ class LLBf16Gemm: ll_bf16_gemm_kernel = LLBf16Gemm() +ll_bf16_gemm_c1_pdl_kernel = LLBf16Gemm(prefetch_pdl_weights=True) def ll_bf16_gemm( @@ -331,4 +338,9 @@ def ll_bf16_gemm( router_weight: torch.Tensor, output_dtype: torch.dtype = torch.float32, ) -> torch.Tensor: - return ll_bf16_gemm_kernel(hidden_states, router_weight, output_dtype) + kernel = ( + ll_bf16_gemm_c1_pdl_kernel + if hidden_states.shape[0] == 1 + else ll_bf16_gemm_kernel + ) + return kernel(hidden_states, router_weight, output_dtype)