iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[Kimi-K3][Perf] Prefetch ll_bf16 router weights for M=1 (#53524)

Browse Source 
Signed-off-by: ming <ming@inferact.ai>
Co-authored-by: OpenAI Codex <codex@openai.com>
This commit is contained in:
mingg26
2026-09-01 21:40:19 +00:00
committed by GitHub
co-authored by OpenAI Codex
parent a56e74afd6
commit 259a209bfa
3 changed files with 99 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/kernels/test_ll_bf16_gemm.py
+23
View File
@@ -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 = [
vllm/model_executor/kernels/linear/cute_dsl/_ll_bf16_dotprod.py
+61 -5
View File
@@ -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):
vllm/model_executor/kernels/linear/cute_dsl/ll_bf16.py
+15 -3
View File
@@ -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)