diff --git a/benchmarks/kernels/benchmark_kimi_k3_kda_projection.py b/benchmarks/kernels/benchmark_kimi_k3_kda_projection.py new file mode 100644 index 0000000000..7a62caa941 --- /dev/null +++ b/benchmarks/kernels/benchmark_kimi_k3_kda_projection.py @@ -0,0 +1,221 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Compare Kimi-K3 TP8 KDA sequential and multistream projections on B300.""" + +from __future__ import annotations + +import argparse +import statistics +from collections.abc import Callable + +import torch + +from vllm.model_executor.kernels.linear.cute_dsl.skinny_gemm import ( + shape_dynamic_skinny_gemm, +) +from vllm.models.kimi_k3.nvidia.low_latency_gemm import ( + KDA_PROJECTION_OVERLAP_MAX_TOKENS, + run_kda_projection_overlap, + try_low_latency_gemm, +) + +ProjectionOutput = tuple[torch.Tensor, torch.Tensor, torch.Tensor] + + +def _capture( + fn: Callable[[], list[ProjectionOutput]], +) -> tuple[torch.cuda.CUDAGraph, list[ProjectionOutput]]: + capture_stream = torch.cuda.Stream() + capture_stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(capture_stream): + fn() + fn() + capture_stream.synchronize() + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=capture_stream): + outputs = fn() + torch.cuda.current_stream().wait_stream(capture_stream) + torch.accelerator.synchronize() + return graph, outputs + + +def _benchmark_graph( + graph: torch.cuda.CUDAGraph, + num_weight_sets: int, + repeats: int, +) -> float: + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + for _ in range(repeats): + graph.replay() + end.record() + end.synchronize() + return start.elapsed_time(end) * 1000 / repeats / num_weight_sets + + +def _validate( + output: ProjectionOutput, + hidden_states: torch.Tensor, + packed_weight: torch.Tensor, + f_b_weight: torch.Tensor, +) -> None: + expected_qkvg = torch.nn.functional.linear( + hidden_states.float(), packed_weight[:6144].float() + ) + expected_fab = torch.nn.functional.linear( + hidden_states.float(), packed_weight[6144:6284].float() + ) + expected_g1 = torch.nn.functional.linear( + expected_fab[:, :128].to(torch.bfloat16).float(), + f_b_weight.float(), + ) + for actual, expected in zip( + output, + (expected_qkvg, expected_g1, expected_fab[:, 128:]), + ): + cosine = torch.nn.functional.cosine_similarity( + actual.float().flatten(), expected.flatten(), dim=0 + ).item() + if cosine <= 0.999: + raise AssertionError(f"projection cosine similarity is {cosine}") + + +def _sequential_projection( + hidden_states: torch.Tensor, + packed_weight: torch.Tensor, + f_b_weight: torch.Tensor, +) -> ProjectionOutput: + projected = try_low_latency_gemm(hidden_states, packed_weight) + if projected is None: + projected = torch.mm(hidden_states, packed_weight.t()) + f_a = projected[:, 6144:6272] + g1 = torch.mm(f_a, f_b_weight.t()) + return projected[:, :6144], g1, projected[:, 6272:6284] + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--tokens", type=int, nargs="+", default=range(1, 17)) + parser.add_argument("--weight-sets", type=int, default=4) + parser.add_argument("--repeats", type=int, default=300) + parser.add_argument("--trials", type=int, default=5) + args = parser.parse_args() + + if torch.cuda.get_device_capability() != (10, 3): + raise SystemExit("This benchmark requires an SM103 GPU") + if not shape_dynamic_skinny_gemm.is_available(): + raise SystemExit("This benchmark requires CuTe DSL") + required_ops = ("dsv3_fused_a_gemm",) + if not all(hasattr(torch.ops._C, name) for name in required_ops): + raise SystemExit("This benchmark requires the KDA projection GEMMs") + if not all(1 <= num_tokens <= 16 for num_tokens in args.tokens): + raise SystemExit("--tokens must be in [1, 16]") + + torch.manual_seed(42) + packed_weights = [ + torch.randn(6288, 7168, dtype=torch.bfloat16, device="cuda") + for _ in range(args.weight_sets) + ] + f_b_weights = [ + torch.randn(1536, 128, dtype=torch.bfloat16, device="cuda") + for _ in range(args.weight_sets) + ] + + from flashinfer.autotuner import autotune + from flashinfer.gemm import mm_bf16 + + with torch.inference_mode(), autotune(tune_mode=True): + mm_bf16( + torch.empty( + (KDA_PROJECTION_OVERLAP_MAX_TOKENS, 7168), + dtype=torch.bfloat16, + device="cuda", + ), + packed_weights[0][:6144].t(), + pdl=True, + backend="cute-dsl", + ) + + print("M sequential (us) multistream (us) fastest") + for num_tokens in args.tokens: + hidden_states = [ + torch.randn( + num_tokens, + 7168, + dtype=torch.bfloat16, + device="cuda", + ) + for _ in range(args.weight_sets) + ] + + def sequential( + hidden_states: list[torch.Tensor] = hidden_states, + ) -> list[ProjectionOutput]: + return [ + _sequential_projection(x, weight, f_b_weight) + for x, weight, f_b_weight in zip( + hidden_states, packed_weights, f_b_weights + ) + ] + + candidates: list[tuple[str, Callable[[], list[ProjectionOutput]]]] = [ + ("sequential", sequential) + ] + if num_tokens <= KDA_PROJECTION_OVERLAP_MAX_TOKENS: + aux_stream = torch.cuda.Stream() + events = [ + (torch.cuda.Event(), torch.cuda.Event()) + for _ in range(args.weight_sets) + ] + + def overlap( + hidden_states: list[torch.Tensor] = hidden_states, + aux_stream: torch.cuda.Stream = aux_stream, + events: list[tuple[torch.cuda.Event, torch.cuda.Event]] = events, + ) -> list[ProjectionOutput]: + return [ + run_kda_projection_overlap(x, weight, f_b_weight, aux_stream, pair) + for x, weight, f_b_weight, pair in zip( + hidden_states, packed_weights, f_b_weights, events + ) + ] + + candidates.append(("multistream", overlap)) + + graphs = {} + for name, candidate in candidates: + graph, outputs = _capture(candidate) + graph.replay() + torch.accelerator.synchronize() + _validate(outputs[0], hidden_states[0], packed_weights[0], f_b_weights[0]) + graphs[name] = graph + + samples = {name: [] for name in graphs} + graph_items = list(graphs.items()) + for trial in range(args.trials): + ordered_graphs = graph_items if trial % 2 == 0 else graph_items[::-1] + for name, graph in ordered_graphs: + samples[name].append( + _benchmark_graph(graph, args.weight_sets, args.repeats) + ) + results = {name: statistics.median(values) for name, values in samples.items()} + + sequential_us = results["sequential"] + multistream_us = results.get("multistream") + multistream_text = ( + f"{multistream_us:18.3f}" + if multistream_us is not None + else " -" + ) + fastest = ( + "multistream" + if multistream_us is not None and multistream_us < sequential_us + else "sequential" + ) + print(f"{num_tokens:2d} {sequential_us:15.3f}{multistream_text} {fastest}") + + +if __name__ == "__main__": + main() diff --git a/tests/kernels/test_bf16_skinny_gemm.py b/tests/kernels/test_bf16_skinny_gemm.py index ca4e8c08e4..cbbaefdf14 100644 --- a/tests/kernels/test_bf16_skinny_gemm.py +++ b/tests/kernels/test_bf16_skinny_gemm.py @@ -253,6 +253,65 @@ def test_residual_cute_configs_match_measured_table() -> None: assert actual == EXPECTED_RESIDUAL_CUTE_CONFIGS +def test_kda_overlap_configs_match_measured_table() -> None: + assert k3_gemm.KDA_PROJECTION_OVERLAP_MAX_TOKENS == 14 + assert k3_gemm.KDA_SKINNY_N_MAX_TOKENS == 14 + assert k3_gemm.KDA_SKINNY_K_MAX_TOKENS == 14 + assert _config_tuple(k3_gemm.KDA_M1_QKVG_CONFIG) == (64, 4, 2, 8) + assert _config_tuple(k3_gemm.KDA_M1_FAB_CONFIG) == (224, 1, 2, 8) + assert { + num_tokens: _config_tuple(config) + for num_tokens, config in k3_gemm.KDA_QKVG_CONFIGS.items() + } == { + 1: (64, 4, 2, 8), + 2: (64, 3, 2, 8), + } + + +@pytest.mark.parametrize("tp_size", [1, 2, 4, 16]) +def test_kda_projection_overlap_is_tp8_only(tp_size: int) -> None: + from vllm.models.kimi_k3.nvidia.kda import KimiK3DeltaAttention + + kda = KimiK3DeltaAttention.__new__(KimiK3DeltaAttention) + nn.Module.__init__(kda) + kda.tp_size = tp_size + kda._projection_overlap_max_tokens = 0 + + assert not k3_gemm._enable_kda_projection_overlap(kda) + assert kda._projection_overlap_max_tokens == 0 + + +def test_kda_qkvg_autotune_enables_full_overlap( + monkeypatch: pytest.MonkeyPatch, +) -> None: + import flashinfer.gemm + + from vllm.models.kimi_k3.nvidia.kda import KimiK3DeltaAttention + + kda = KimiK3DeltaAttention.__new__(KimiK3DeltaAttention) + nn.Module.__init__(kda) + kda._projection_overlap_max_tokens = max(k3_gemm.KDA_QKVG_CONFIGS) + kda.in_proj_qkvgfab = nn.Module() + kda.in_proj_qkvgfab.weight = nn.Parameter( + torch.empty(6144, 8, dtype=torch.bfloat16), + requires_grad=False, + ) + calls = [] + + def fake_mm_bf16(a, b, *, pdl, backend): + calls.append((a.shape, b.shape, pdl, backend)) + return torch.empty(a.shape[0], b.shape[1], dtype=a.dtype) + + monkeypatch.setattr(flashinfer.gemm, "mm_bf16", fake_mm_bf16) + + k3_gemm.autotune_kda_qkvg(kda) + + assert calls == [(torch.Size([14, 8]), torch.Size([8, 6144]), True, "cute-dsl")] + assert ( + kda._projection_overlap_max_tokens == k3_gemm.KDA_PROJECTION_OVERLAP_MAX_TOKENS + ) + + def test_glm52_projection_plans_are_separate() -> None: qkv_a = glm52_gemm.GLM52_QKV_A_PROJECTION q_b = glm52_gemm.GLM52_Q_B_PROJECTION @@ -812,6 +871,65 @@ def test_cute_selected_shapes( assert cosine > 0.999 +@pytest.mark.parametrize("num_tokens", [1, 2, 3, 4, 8, 9, 10, 11, 12, 13, 14]) +def test_kda_projection_overlap_cuda_graph(num_tokens: int) -> None: + """The TP8 split must preserve projection results across both streams.""" + _require_sm103_and_cute() + _require_sm103_and_dsv3() + torch.manual_seed(43) + hidden_states = torch.randn(num_tokens, 7168, dtype=torch.bfloat16, device="cuda") + packed_weight = torch.randn(6288, 7168, dtype=torch.bfloat16, device="cuda") + f_b_weight = torch.randn(1536, 128, dtype=torch.bfloat16, device="cuda") + aux_stream = torch.cuda.Stream() + events = (torch.cuda.Event(), torch.cuda.Event()) + + k3_gemm.run_kda_projection_overlap( + hidden_states, + packed_weight, + f_b_weight, + aux_stream, + events, + ) + torch.accelerator.synchronize() + + capture_stream = torch.cuda.Stream() + capture_stream.wait_stream(torch.cuda.current_stream()) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=capture_stream): + projected_qkvg, g1, beta = k3_gemm.run_kda_projection_overlap( + hidden_states, + packed_weight, + f_b_weight, + aux_stream, + events, + ) + torch.cuda.current_stream().wait_stream(capture_stream) + graph.replay() + torch.accelerator.synchronize() + + assert beta.stride() == ((140 if num_tokens == 1 else 144), 1) + + expected_qkvg = torch.nn.functional.linear( + hidden_states.float(), packed_weight[:6144].float() + ) + expected_fab = torch.nn.functional.linear( + hidden_states.float(), packed_weight[6144:6284].float() + ) + expected_g1 = torch.nn.functional.linear( + expected_fab[:, :128].to(torch.bfloat16).float(), + f_b_weight.float(), + ) + for actual, expected in ( + (projected_qkvg, expected_qkvg), + (g1, expected_g1), + (beta, expected_fab[:, 128:]), + ): + cosine = torch.nn.functional.cosine_similarity( + actual.float().flatten(), expected.flatten(), dim=0 + ).item() + assert cosine > 0.999 + + def _dsv3_probe_tokens(tokens: frozenset[int]) -> set[int]: """Extremes, plus both sides of the kernel's num_tokens<=8 tile_n branch.""" if not tokens: diff --git a/vllm/model_executor/warmup/kernel_warmup.py b/vllm/model_executor/warmup/kernel_warmup.py index ad11976650..78d9dbb98e 100644 --- a/vllm/model_executor/warmup/kernel_warmup.py +++ b/vllm/model_executor/warmup/kernel_warmup.py @@ -114,6 +114,12 @@ def _warmup_kimi_k3_gemm_rs_ar() -> None: logger.info_once("Warmed up %d Kimi-K3 GEMM-RS/AR variants.", compiled) +def _autotune_kimi_k3_kda_qkvg(model: torch.nn.Module) -> None: + module = sys.modules.get("vllm.models.kimi_k3.nvidia.low_latency_gemm") + if module is not None: + module.autotune_kda_qkvg(model) + + def kernel_warmup(worker: "Worker", *, process_local_only: bool = False): from vllm.model_executor.warmup.minimax_m3_msa_warmup import ( minimax_m3_msa_warmup, @@ -343,6 +349,7 @@ def flashinfer_autotune(runner: "GPUModelRunner") -> None: ): _run_flashinfer_autotune_dummy_runs(runner) replayssm_autotune_warmup(runner) + _autotune_kimi_k3_kda_qkvg(runner.get_model()) finally: set_autotune_process_group(None) diff --git a/vllm/models/kimi_k3/nvidia/kda.py b/vllm/models/kimi_k3/nvidia/kda.py index 83460967a6..3b17445a39 100644 --- a/vllm/models/kimi_k3/nvidia/kda.py +++ b/vllm/models/kimi_k3/nvidia/kda.py @@ -399,6 +399,7 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention): config: KimiLinearConfig, vllm_config: VllmConfig, prefix: str = "", + aux_stream: torch.cuda.Stream | None = None, run_gemm_rs_ar: bool = False, ) -> None: super().__init__(config, vllm_config, prefix) @@ -421,6 +422,11 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention): assert kda_config.get("use_full_rank_gate", False), ( "KimiK3DeltaAttention requires a full-rank gate" ) + self._projection_aux_stream = aux_stream + self._projection_events = ( + (torch.cuda.Event(), torch.cuda.Event()) if aux_stream is not None else None + ) + self._projection_overlap_max_tokens = 0 # Keep f_a before the narrow beta shard, then pad each TP-local row # to select the aligned BF16 GEMM path. @@ -605,19 +611,44 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention): positions: torch.Tensor, ) -> torch.Tensor: num_tokens = hidden_states.size(0) - projected_qkvgfab = self.in_proj_qkvgfab(hidden_states)[0] - split_sizes = [ - 3 * self.local_projection_size, - self.local_projection_size, - self.head_dim, - self.local_num_heads, - ] - if self.in_proj_padding: - split_sizes.append(self.in_proj_padding) - projected = projected_qkvgfab.split(split_sizes, dim=-1) - mixed_qkv, g_proj_states, f_a, beta = projected[:4] + projection_events = self._projection_events + projection_aux_stream = self._projection_aux_stream + if ( + 0 < num_tokens <= self._projection_overlap_max_tokens + and hidden_states.stride() == (self.hidden_size, 1) + and projection_events is not None + and projection_aux_stream is not None + and torch.cuda.is_current_stream_capturing() + ): + from vllm.models.kimi_k3.nvidia.low_latency_gemm import ( + run_kda_projection_overlap, + ) + + projected_qkvg, g1, beta = run_kda_projection_overlap( + hidden_states, + self.in_proj_qkvgfab.weight, + self.f_b_proj.weight, + projection_aux_stream, + projection_events, + ) + mixed_qkv, g_proj_states = projected_qkvg.split( + [3 * self.local_projection_size, self.local_projection_size], + dim=-1, + ) + else: + projected_qkvgfab = self.in_proj_qkvgfab(hidden_states)[0] + split_sizes = [ + 3 * self.local_projection_size, + self.local_projection_size, + self.head_dim, + self.local_num_heads, + ] + if self.in_proj_padding: + split_sizes.append(self.in_proj_padding) + projected = projected_qkvgfab.split(split_sizes, dim=-1) + mixed_qkv, g_proj_states, f_a, beta = projected[:4] + g1 = self.f_b_proj(f_a)[0] - g1 = self.f_b_proj(f_a)[0] beta = beta.unsqueeze(0) g1 = rearrange(g1, "n (h d) -> 1 n h d", d=self.head_dim) g2 = rearrange(g_proj_states, "... (h d) -> ... h d", d=self.head_dim) diff --git a/vllm/models/kimi_k3/nvidia/low_latency_gemm.py b/vllm/models/kimi_k3/nvidia/low_latency_gemm.py index 52e6aa0475..2b23a8aec1 100644 --- a/vllm/models/kimi_k3/nvidia/low_latency_gemm.py +++ b/vllm/models/kimi_k3/nvidia/low_latency_gemm.py @@ -35,12 +35,32 @@ from vllm.model_executor.layers.vocab_parallel_embedding import ( UnquantizedEmbeddingMethod, ) from vllm.platforms import current_platform +from vllm.utils.multi_stream_utils import maybe_execute_in_parallel Backend = Literal["cute", "dsv3_fused_a"] # A resolved per-token-count call: the backend plus its CuTe config (None for # dsv3, which needs no config). ResolvedCall = tuple[Backend, SkinnyGemmConfig | None] +# TP8 KDA projection split, measured together under CUDA graph capture on B300. +# Q/K/V/G stay on the graph's main stream while F_A/beta and F_B run on the +# model's auxiliary stream. +KDA_M1_QKVG_CONFIG = SkinnyGemmConfig(1, 64, 4, 2, 8) +KDA_M1_FAB_CONFIG = SkinnyGemmConfig(1, 224, 1, 2, 8) +KDA_QKVG_CONFIGS = { + 1: KDA_M1_QKVG_CONFIG, + 2: SkinnyGemmConfig(2, 64, 3, 2, 8), +} +# The captured end-to-end projection sweep wins through M=14 and regresses at +# M=15 and M=16, where the original packed projection remains selected. +KDA_PROJECTION_OVERLAP_MAX_TOKENS = 14 +KDA_SKINNY_N_MAX_TOKENS = KDA_PROJECTION_OVERLAP_MAX_TOKENS +KDA_SKINNY_K_MAX_TOKENS = KDA_PROJECTION_OVERLAP_MAX_TOKENS +_KDA_QKVG_SIZE = 4 * 1536 +_KDA_FAB_SIZE = 128 + 12 +_KDA_PACKED_SIZE = 6288 +_KDA_TP_SIZE = 8 + @dataclass(frozen=True, slots=True) class ProjectionSpec: @@ -683,6 +703,137 @@ def _run_plan( return output +def run_kda_projection_overlap( + hidden_states: torch.Tensor, + packed_weight: torch.Tensor, + f_b_weight: torch.Tensor, + aux_stream: torch.cuda.Stream, + events: tuple[torch.cuda.Event, torch.cuda.Event], +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Run the TP8 KDA decode projection branches concurrently. + + Args: + hidden_states: Packed BF16 input with shape ``[M, 7168]``. + packed_weight: Existing Q/K/V/G/F_A/beta/pad weight with shape + ``[6288, 7168]``. + f_b_weight: F_B weight with shape ``[1536, 128]``. + aux_stream: Stream for the F_A/beta then F_B branch. + events: Start and completion events for the stream fork and join. + + Returns: + The Q/K/V/G projection, F_B output, and beta projection. + """ + num_tokens = hidden_states.shape[0] + qkvg_weight = packed_weight[:_KDA_QKVG_SIZE] + + def run_qkvg() -> torch.Tensor: + if config := KDA_QKVG_CONFIGS.get(num_tokens): + return shape_dynamic_skinny_gemm( + hidden_states, + qkvg_weight, + config, + None, + ) + if num_tokens <= KDA_PROJECTION_OVERLAP_MAX_TOKENS: + from flashinfer.gemm import mm_bf16 + + return mm_bf16( + hidden_states, + qkvg_weight.t(), + pdl=True, + backend="cute-dsl", + ) + return torch.mm(hidden_states, qkvg_weight.t()) + + def run_fab_fb() -> tuple[torch.Tensor, torch.Tensor]: + if num_tokens == 1: + projected_fab = shape_dynamic_skinny_gemm( + hidden_states, + packed_weight[_KDA_QKVG_SIZE : _KDA_QKVG_SIZE + _KDA_FAB_SIZE], + KDA_M1_FAB_CONFIG, + None, + ) + f_a, beta = projected_fab.split([128, 12], dim=-1) + f_a = f_a.as_strided((1, 128), (128, 1)) + g1 = torch.empty( + (1, 1536), dtype=hidden_states.dtype, device=hidden_states.device + ) + ops.dsv3_fused_a_gemm(g1, f_a, f_b_weight.t(), enable_pdl=True) + else: + from vllm.models.kimi_k3.nvidia.ops.cute_dsl.kda_skinny_gemm import ( + kda_skinny_gemm, + ) + + if num_tokens <= KDA_SKINNY_N_MAX_TOKENS: + projected_fab = kda_skinny_gemm.run_n( + hidden_states, + packed_weight[_KDA_QKVG_SIZE:], + ) + else: + projected_fab = torch.mm( + hidden_states, + packed_weight[_KDA_QKVG_SIZE:].t(), + ) + beta = projected_fab[:, 128:140] + if num_tokens <= KDA_SKINNY_K_MAX_TOKENS: + g1 = kda_skinny_gemm.run_k(projected_fab, f_b_weight) + else: + g1 = torch.mm(projected_fab[:, :128], f_b_weight.t()) + return g1, beta + + projected_qkvg, (g1, beta) = maybe_execute_in_parallel( + run_qkvg, + run_fab_fb, + events[0], + events[1], + aux_stream, + ) + return projected_qkvg, g1, beta + + +def autotune_kda_qkvg(model: nn.Module) -> None: + """Autotune the FlashInfer QKVG GEMM before CUDA graph capture.""" + from flashinfer.gemm import mm_bf16 + + from vllm.models.kimi_k3.nvidia.kda import KimiK3DeltaAttention + + children: list[KimiK3DeltaAttention] = [] + weights_by_shape: dict[ + tuple[int, int, int, torch.dtype, torch.device], torch.Tensor + ] = {} + for child in model.modules(): + if not isinstance(child, KimiK3DeltaAttention): + continue + if child._projection_overlap_max_tokens <= 0: + continue + children.append(child) + qkvg_weight = child.in_proj_qkvgfab.weight[:_KDA_QKVG_SIZE] + shape = ( + KDA_PROJECTION_OVERLAP_MAX_TOKENS, + qkvg_weight.shape[0], + qkvg_weight.shape[1], + qkvg_weight.dtype, + qkvg_weight.device, + ) + weights_by_shape.setdefault(shape, qkvg_weight) + + for shape, qkvg_weight in weights_by_shape.items(): + num_tokens = shape[0] + hidden_states = torch.empty( + (num_tokens, qkvg_weight.shape[1]), + dtype=qkvg_weight.dtype, + device=qkvg_weight.device, + ) + mm_bf16( + hidden_states, + qkvg_weight.t(), + pdl=True, + backend="cute-dsl", + ) + for child in children: + child._projection_overlap_max_tokens = KDA_PROJECTION_OVERLAP_MAX_TOKENS + + def _run_residual_plan( residual_plan: dict[int, SkinnyGemmConfig], x: torch.Tensor, @@ -777,6 +928,40 @@ class KimiK3LowLatencyEmbeddingMethod( pass +def _enable_kda_projection_overlap(module: nn.Module) -> bool: + from vllm.models.kimi_k3.nvidia.kda import KimiK3DeltaAttention + + if envs.VLLM_BATCH_INVARIANT: + return False + + enabled = False + for child in module.modules(): + if not isinstance(child, KimiK3DeltaAttention): + continue + if child.tp_size != _KDA_TP_SIZE: + continue + packed_weight = child.in_proj_qkvgfab.weight + f_b_weight = child.f_b_proj.weight + if ( + type(child.in_proj_qkvgfab.quant_method) is not KimiK3LowLatencyLinearMethod + or type(child.f_b_proj.quant_method) is not KimiK3LowLatencyLinearMethod + or child._projection_aux_stream is None + or child._projection_events is None + or packed_weight.shape != (_KDA_PACKED_SIZE, 7168) + or f_b_weight.shape != (1536, 128) + or packed_weight.dtype != torch.bfloat16 + or f_b_weight.dtype != torch.bfloat16 + or not packed_weight.is_cuda + or not f_b_weight.is_cuda + or not packed_weight.is_contiguous() + or not f_b_weight.is_contiguous() + ): + continue + child._projection_overlap_max_tokens = max(KDA_QKVG_CONFIGS) + enabled = True + return enabled + + def enable_kimi_k3_low_latency_gemm( module: nn.Module, dtype: torch.dtype, @@ -825,6 +1010,16 @@ def enable_kimi_k3_low_latency_gemm( # TP8 deployment does not compile TP4 configs and vice versa. warmup_configs.update(config for _, config in spec.cute_configs) residual_warmup_configs.update(config for _, config in spec.residual_configs) + if _enable_kda_projection_overlap(module): + from vllm.models.kimi_k3.nvidia.ops.cute_dsl.kda_skinny_gemm import ( + kda_skinny_gemm, + ) + + warmup_configs.update((*KDA_QKVG_CONFIGS.values(), KDA_M1_FAB_CONFIG)) + kda_skinny_gemm.request_warmup( + set(range(2, KDA_SKINNY_N_MAX_TOKENS + 1)), + set(range(2, KDA_SKINNY_K_MAX_TOKENS + 1)), + ) if shape_dynamic_skinny_gemm.is_available(): if warmup_configs: diff --git a/vllm/models/kimi_k3/nvidia/model.py b/vllm/models/kimi_k3/nvidia/model.py index dbf998fc27..2fc20ecf21 100644 --- a/vllm/models/kimi_k3/nvidia/model.py +++ b/vllm/models/kimi_k3/nvidia/model.py @@ -895,6 +895,7 @@ class KimiDecoderLayer(nn.Module): config, vllm_config, prefix=f"{prefix}.self_attn", + aux_stream=aux_stream, run_gemm_rs_ar=run_gemm_rs_ar, ) self._self_attn_writes_output = False diff --git a/vllm/models/kimi_k3/nvidia/ops/cute_dsl/kda_skinny_gemm.py b/vllm/models/kimi_k3/nvidia/ops/cute_dsl/kda_skinny_gemm.py new file mode 100644 index 0000000000..da69e2631b --- /dev/null +++ b/vllm/models/kimi_k3/nvidia/ops/cute_dsl/kda_skinny_gemm.py @@ -0,0 +1,507 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# Adapted from +# https://github.com/sgl-project/sglang/blob/main/python/sglang/kernels/jit/csrc/gemm/tiny_gemm.cuh +"""Kimi-K3 TP8 skinny GEMMs for the KDA F_A/beta and F_B projections.""" + +from __future__ import annotations + +from collections.abc import Callable +from functools import partial + +import cuda.bindings.driver as cuda +import cutlass +import cutlass.cute as cute +import torch +from cutlass import BFloat16, Float32 +from cutlass._mlir.dialects import llvm +from cutlass.cute.runtime import make_fake_stream, make_fake_tensor +from cutlass.cutlass_dsl import T, dsl_user_op + +_FAB_N = 144 +_FAB_K = 7168 +_FAB_VECTOR_WIDTH = 16 +_FAB_BLOCK_SIZE = _FAB_K // _FAB_VECTOR_WIDTH +_FAB_NUM_WARPS = _FAB_BLOCK_SIZE // cute.arch.WARP_SIZE + +_FB_N = 1536 +_FB_K = 128 +_FB_VECTOR_WIDTH = 8 +_FB_K_LANES = _FB_K // _FB_VECTOR_WIDTH + +SKINNY_N_SPLIT_N = { + 2: 1, + 3: 2, + 4: 3, + 5: 3, + 6: 3, + 7: 3, + 8: 3, + 9: 2, + 10: 2, + 11: 2, + 12: 2, + 13: 2, + 14: 1, + 15: 1, + 16: 1, +} + +SKINNY_K_SPLIT_N = { + 2: 8, + 3: 16, + 4: 16, + 5: 24, + 6: 24, + 7: 24, + 8: 24, + 9: 12, + 10: 16, + 11: 16, + 12: 12, + 13: 8, + 14: 8, + 15: 8, + 16: 8, +} + + +@dsl_user_op +def _fma_f32_bf16( + a: BFloat16, + b: BFloat16, + acc: Float32, + *, + loc=None, + ip=None, +) -> Float32: + a_bits = llvm.bitcast(T.i16(), a.ir_value(loc=loc, ip=ip), loc=loc, ip=ip) + b_bits = llvm.bitcast(T.i16(), b.ir_value(loc=loc, ip=ip), loc=loc, ip=ip) + result = llvm.inline_asm( + T.f32(), + [a_bits, b_bits, acc.ir_value(loc=loc, ip=ip)], + "fma.rn.f32.bf16 $0, $1, $2, $3;", + "=f,h,h,f", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + return Float32(result) + + +class _KdaSkinnyNGemm: + def __init__(self, num_rows: int, split_n: int) -> None: + if _FAB_N % split_n: + raise ValueError("split_n must divide the FAB output size") + if num_rows * split_n > _FAB_BLOCK_SIZE: + raise ValueError("the FAB output tile must fit in one block") + self.num_rows = num_rows + self.split_n = split_n + + @cute.jit + def __call__( + self, + a: cute.Tensor, + weight: cute.Tensor, + output: cute.Tensor, + stream: cuda.CUstream, + ) -> None: + copy_a = cute.make_copy_atom( + cute.nvgpu.CopyG2ROp(), + BFloat16, + num_bits_per_copy=_FAB_VECTOR_WIDTH * BFloat16.width, + load_cache_mode=cute.nvgpu.LoadCacheMode.ALWAYS, + ) + copy_weight = cute.make_copy_atom( + cute.nvgpu.CopyG2ROp(), + BFloat16, + num_bits_per_copy=_FAB_VECTOR_WIDTH * BFloat16.width, + load_cache_mode=cute.nvgpu.LoadCacheMode.STREAMING, + ) + self.kernel(a, weight, output, copy_a, copy_weight).launch( + grid=[_FAB_N // self.split_n, 1, 1], + block=[_FAB_BLOCK_SIZE, 1, 1], + smem=_FAB_NUM_WARPS * self.num_rows * self.split_n * 4, + stream=stream, + use_pdl=True, + min_blocks_per_mp=1, + ) + + @cute.kernel + def kernel( + self, + a: cute.Tensor, + weight: cute.Tensor, + output: cute.Tensor, + copy_a: cute.CopyAtom, + copy_weight: cute.CopyAtom, + ) -> None: + tidx, _, _ = cute.arch.thread_idx() + block_idx, _, _ = cute.arch.block_idx() + warp_idx = cute.arch.warp_idx() + num_rows: cutlass.Constexpr = self.num_rows + split_n: cutlass.Constexpr = self.split_n + + a_vectors = cute.logical_divide(a, (None, _FAB_VECTOR_WIDTH)) + weight_vectors = cute.logical_divide(weight, (None, _FAB_VECTOR_WIDTH)) + a_tiles = cute.logical_divide( + a_vectors, + (None, (None, _FAB_BLOCK_SIZE)), + ) + weight_tiles = cute.logical_divide( + weight_vectors, + (None, (None, _FAB_BLOCK_SIZE)), + ) + a_tile = a_tiles[None, (None, (tidx, None))] + n_base = block_idx * split_n + + a_regs = cute.make_rmem_tensor( + cute.make_layout( + (num_rows, _FAB_VECTOR_WIDTH), + stride=(_FAB_VECTOR_WIDTH, 1), + ), + BFloat16, + ) + weight_regs = cute.make_rmem_tensor( + cute.make_layout( + (split_n, _FAB_VECTOR_WIDTH), + stride=(_FAB_VECTOR_WIDTH, 1), + ), + BFloat16, + ) + acc = cute.make_rmem_tensor( + cute.make_layout((num_rows, split_n), stride=(split_n, 1)), + Float32, + ) + acc.fill(0.0) + + for ni in cutlass.range_constexpr(split_n): + weight_tile = weight_tiles[ + n_base + ni, + (None, (tidx, None)), + ] + cute.copy(copy_weight, weight_tile[None, 0], weight_regs[ni, None]) + + cute.arch.griddepcontrol_wait() + + for mi in cutlass.range_constexpr(num_rows): + cute.copy(copy_a, a_tile[mi, None, 0], a_regs[mi, None]) + for mi in cutlass.range_constexpr(num_rows): + for ni in cutlass.range_constexpr(split_n): + for vi in cutlass.range_constexpr(_FAB_VECTOR_WIDTH): + acc[mi, ni] = _fma_f32_bf16( + a_regs[mi, vi], + weight_regs[ni, vi], + acc[mi, ni], + ) + acc[mi, ni] = cute.arch.warp_reduction_sum(acc[mi, ni]) + + partials_layout = cute.make_layout( + (_FAB_NUM_WARPS, num_rows, split_n), + stride=(num_rows * split_n, split_n, 1), + ) + smem = cutlass.utils.SmemAllocator() + partials = smem.allocate_tensor( + Float32, + partials_layout, + byte_alignment=16, + ) + with cute.arch.elect_one(): + for mi in cutlass.range_constexpr(num_rows): + for ni in cutlass.range_constexpr(split_n): + partials[warp_idx, mi, ni] = acc[mi, ni] + + cute.arch.griddepcontrol_launch_dependents() + cute.arch.sync_threads() + + if tidx < num_rows * split_n: + mi = tidx // split_n + ni = tidx % split_n + total = ( + partials[None, mi, ni] + .load() + .reduce( + cute.ReductionOp.ADD, + init_val=Float32(0.0), + reduction_profile=0, + ) + ) + output[mi, n_base + ni] = Float32(total).to(BFloat16) + + +class _KdaSkinnyKGemm: + def __init__(self, num_rows: int, split_n: int) -> None: + if _FB_N % split_n: + raise ValueError("split_n must divide the FB output size") + if split_n * _FB_K_LANES > 1024: + raise ValueError("the FB output tile exceeds the block-size limit") + self.num_rows = num_rows + self.split_n = split_n + + @cute.jit + def __call__( + self, + a: cute.Tensor, + weight: cute.Tensor, + output: cute.Tensor, + stream: cuda.CUstream, + ) -> None: + copy_atom = cute.make_copy_atom( + cute.nvgpu.CopyG2ROp(), + BFloat16, + num_bits_per_copy=_FB_VECTOR_WIDTH * BFloat16.width, + load_cache_mode=cute.nvgpu.LoadCacheMode.ALWAYS, + ) + self.kernel(a, weight, output, copy_atom).launch( + grid=[_FB_N // self.split_n, 1, 1], + block=[self.split_n * _FB_K_LANES, 1, 1], + stream=stream, + use_pdl=True, + min_blocks_per_mp=1, + ) + + @cute.kernel + def kernel( + self, + a: cute.Tensor, + weight: cute.Tensor, + output: cute.Tensor, + copy_atom: cute.CopyAtom, + ) -> None: + tidx, _, _ = cute.arch.thread_idx() + block_idx, _, _ = cute.arch.block_idx() + num_rows: cutlass.Constexpr = self.num_rows + split_n: cutlass.Constexpr = self.split_n + k_lane = tidx % _FB_K_LANES + n_idx = block_idx * split_n + tidx // _FB_K_LANES + + a_vectors = cute.logical_divide(a, (None, _FB_VECTOR_WIDTH)) + weight_vectors = cute.logical_divide(weight, (None, _FB_VECTOR_WIDTH)) + a_tiles = cute.logical_divide( + a_vectors, + (None, (None, _FB_K_LANES)), + ) + weight_tiles = cute.logical_divide( + weight_vectors, + (None, (None, _FB_K_LANES)), + ) + a_tile = a_tiles[None, (None, (k_lane, None))] + weight_tile = weight_tiles[n_idx, (None, (k_lane, None))] + weight_regs = cute.make_rmem_tensor(_FB_VECTOR_WIDTH, BFloat16) + input_regs = cute.make_rmem_tensor( + cute.make_layout( + (num_rows, _FB_VECTOR_WIDTH), + stride=(_FB_VECTOR_WIDTH, 1), + ), + BFloat16, + ) + cute.copy( + copy_atom, + weight_tile[None, 0], + weight_regs, + ) + + cute.arch.griddepcontrol_wait() + + for mi in cutlass.range_constexpr(num_rows): + cute.copy( + copy_atom, + a_tile[mi, None, 0], + input_regs[mi, None], + ) + for mi in cutlass.range_constexpr(num_rows): + acc = Float32(0.0) + for vi in cutlass.range_constexpr(_FB_VECTOR_WIDTH): + acc = _fma_f32_bf16( + input_regs[mi, vi], + weight_regs[vi], + acc, + ) + for offset in (8, 4, 2, 1): + acc += cute.arch.shuffle_sync_bfly(acc, offset) + output[mi, n_idx] = acc.to(BFloat16) + + cute.arch.griddepcontrol_launch_dependents() + + +class KdaSkinnyGemm: + def __init__(self) -> None: + self._compiled_n: dict[tuple[int, int], Callable[..., None]] = {} + self._compiled_k: dict[tuple[int, int], Callable[..., None]] = {} + self._warmup_n: set[int] = set() + self._warmup_k: set[int] = set() + self._warmup_registered = False + + def _compile_n(self, num_rows: int) -> None: + device = torch.accelerator.current_device_index() + key = (device, num_rows) + if key in self._compiled_n: + return + split_n = SKINNY_N_SPLIT_N[num_rows] + self._compiled_n[key] = cute.compile( + _KdaSkinnyNGemm(num_rows, split_n), + make_fake_tensor( + BFloat16, + (num_rows, _FAB_K), + stride=(_FAB_K, 1), + assumed_align=32, + ), + make_fake_tensor( + BFloat16, + (_FAB_N, _FAB_K), + stride=(_FAB_K, 1), + assumed_align=32, + ), + make_fake_tensor( + BFloat16, + (num_rows, _FAB_N), + stride=(_FAB_N, 1), + assumed_align=32, + ), + make_fake_stream(use_tvm_ffi_env_stream=True), + options="--enable-tvm-ffi --ptxas-options -maxrregcount=255", + ) + + def _compile_k(self, num_rows: int) -> None: + device = torch.accelerator.current_device_index() + key = (device, num_rows) + if key in self._compiled_k: + return + split_n = SKINNY_K_SPLIT_N[num_rows] + self._compiled_k[key] = cute.compile( + _KdaSkinnyKGemm(num_rows, split_n), + make_fake_tensor( + BFloat16, + (num_rows, _FB_K), + stride=(_FAB_N, 1), + assumed_align=32, + ), + make_fake_tensor( + BFloat16, + (_FB_N, _FB_K), + stride=(_FB_K, 1), + assumed_align=32, + ), + make_fake_tensor( + BFloat16, + (num_rows, _FB_N), + stride=(_FB_N, 1), + assumed_align=32, + ), + make_fake_stream(use_tvm_ffi_env_stream=True), + options="--enable-tvm-ffi --ptxas-options -maxrregcount=255", + ) + + def request_warmup(self, n_tokens: set[int], k_tokens: set[int]) -> None: + if not n_tokens <= SKINNY_N_SPLIT_N.keys(): + raise ValueError("unsupported skinny-N warmup token count") + if not k_tokens <= SKINNY_K_SPLIT_N.keys(): + raise ValueError("unsupported skinny-K warmup token count") + self._warmup_n.update(n_tokens) + self._warmup_k.update(k_tokens) + if self._warmup_registered: + return + from vllm.model_executor.warmup.cutedsl_warmup import ( + register_cutedsl_warmup_provider, + ) + + register_cutedsl_warmup_provider(self) + self._warmup_registered = True + + def get_cutedsl_warmup_compile_units(self): + from vllm.model_executor.warmup.cutedsl_warmup import CuTeDSLCompileUnit + + n_units = tuple( + CuTeDSLCompileUnit( + name="Kimi-K3 KDA skinny-N GEMM", + key=("kimi-k3-kda-skinny-n", num_rows), + compile=partial(self._compile_n, num_rows), + ) + for num_rows in sorted(self._warmup_n) + ) + k_units = tuple( + CuTeDSLCompileUnit( + name="Kimi-K3 KDA skinny-K GEMM", + key=("kimi-k3-kda-skinny-k", num_rows), + compile=partial(self._compile_k, num_rows), + ) + for num_rows in sorted(self._warmup_k) + ) + return n_units + k_units + + @staticmethod + def _validate( + a: torch.Tensor, + weight: torch.Tensor, + *, + a_shape: tuple[int, int], + weight_shape: tuple[int, int], + ) -> None: + if a.shape != a_shape or weight.shape != weight_shape: + raise ValueError( + f"expected input {a_shape} and weight {weight_shape}, got " + f"{tuple(a.shape)} and {tuple(weight.shape)}" + ) + if a.dtype != torch.bfloat16 or weight.dtype != torch.bfloat16: + raise ValueError("KDA skinny GEMMs require BF16 input and weight") + if not a.is_cuda or not weight.is_cuda or a.device != weight.device: + raise ValueError("KDA skinny GEMM operands must share a CUDA device") + if not a.is_contiguous() or not weight.is_contiguous(): + raise ValueError("KDA skinny GEMM operands must be contiguous") + if a.data_ptr() % 32 or weight.data_ptr() % 32: + raise ValueError("KDA skinny GEMM operands must be 32-byte aligned") + + def run_n(self, a: torch.Tensor, weight: torch.Tensor) -> torch.Tensor: + num_rows = a.shape[0] + if num_rows not in SKINNY_N_SPLIT_N: + raise ValueError("skinny-N requires M in [2, 16]") + self._validate( + a, + weight, + a_shape=(num_rows, _FAB_K), + weight_shape=(_FAB_N, _FAB_K), + ) + output = torch.empty( + (num_rows, _FAB_N), + dtype=a.dtype, + device=a.device, + ) + with torch.accelerator.device_index(a.device.index): + self._compile_n(num_rows) + key = (a.device.index, num_rows) + self._compiled_n[key]( + a, + weight, + output, + ) + return output + + def run_k(self, packed_fab: torch.Tensor, weight: torch.Tensor) -> torch.Tensor: + num_rows = packed_fab.shape[0] + if num_rows not in SKINNY_K_SPLIT_N: + raise ValueError("skinny-K requires M in [2, 16]") + self._validate( + packed_fab, + weight, + a_shape=(num_rows, _FAB_N), + weight_shape=(_FB_N, _FB_K), + ) + output = torch.empty( + (num_rows, _FB_N), + dtype=packed_fab.dtype, + device=packed_fab.device, + ) + with torch.accelerator.device_index(packed_fab.device.index): + self._compile_k(num_rows) + key = (packed_fab.device.index, num_rows) + self._compiled_k[key]( + packed_fab[:, :_FB_K], + weight, + output, + ) + return output + + +kda_skinny_gemm = KdaSkinnyGemm()