| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Compare Kimi-K3 TP8 KDA sequential and multistream projections on B300.""" |
| 4 | |
| 5 | from __future__ import annotations |
| 6 | |
| 7 | import argparse |
| 8 | import statistics |
| 9 | from collections.abc import Callable |
| 10 | |
| 11 | import torch |
| 12 | |
| 13 | from vllm.model_executor.kernels.linear.cute_dsl.skinny_gemm import ( |
| 14 | shape_dynamic_skinny_gemm, |
| 15 | ) |
| 16 | from vllm.models.kimi_k3.nvidia.low_latency_gemm import ( |
| 17 | KDA_PROJECTION_OVERLAP_MAX_TOKENS, |
| 18 | run_kda_projection_overlap, |
| 19 | try_low_latency_gemm, |
| 20 | ) |
| 21 | |
| 22 | ProjectionOutput = tuple[torch.Tensor, torch.Tensor, torch.Tensor] |
| 23 | |
| 24 | |
| 25 | def _capture( |
| 26 | fn: Callable[[], list[ProjectionOutput]], |
| 27 | ) -> tuple[torch.cuda.CUDAGraph, list[ProjectionOutput]]: |
| 28 | capture_stream = torch.cuda.Stream() |
| 29 | capture_stream.wait_stream(torch.cuda.current_stream()) |
| 30 | with torch.cuda.stream(capture_stream): |
| 31 | fn() |
| 32 | fn() |
| 33 | capture_stream.synchronize() |
| 34 | |
| 35 | graph = torch.cuda.CUDAGraph() |
| 36 | with torch.cuda.graph(graph, stream=capture_stream): |
| 37 | outputs = fn() |
| 38 | torch.cuda.current_stream().wait_stream(capture_stream) |
| 39 | torch.accelerator.synchronize() |
| 40 | return graph, outputs |
| 41 | |
| 42 | |
| 43 | def _benchmark_graph( |
| 44 | graph: torch.cuda.CUDAGraph, |
| 45 | num_weight_sets: int, |
| 46 | repeats: int, |
| 47 | ) -> float: |
| 48 | start = torch.cuda.Event(enable_timing=True) |
| 49 | end = torch.cuda.Event(enable_timing=True) |
| 50 | start.record() |
| 51 | for _ in range(repeats): |
| 52 | graph.replay() |
| 53 | end.record() |
| 54 | end.synchronize() |
| 55 | return start.elapsed_time(end) * 1000 / repeats / num_weight_sets |
| 56 | |
| 57 | |
| 58 | def _validate( |
| 59 | output: ProjectionOutput, |
| 60 | hidden_states: torch.Tensor, |
| 61 | packed_weight: torch.Tensor, |
| 62 | f_b_weight: torch.Tensor, |
| 63 | ) -> None: |
| 64 | expected_qkvg = torch.nn.functional.linear( |
| 65 | hidden_states.float(), packed_weight[:6144].float() |
| 66 | ) |
| 67 | expected_fab = torch.nn.functional.linear( |
| 68 | hidden_states.float(), packed_weight[6144:6284].float() |
| 69 | ) |
| 70 | expected_g1 = torch.nn.functional.linear( |
| 71 | expected_fab[:, :128].to(torch.bfloat16).float(), |
| 72 | f_b_weight.float(), |
| 73 | ) |
| 74 | for actual, expected in zip( |
| 75 | output, |
| 76 | (expected_qkvg, expected_g1, expected_fab[:, 128:]), |
| 77 | ): |
| 78 | cosine = torch.nn.functional.cosine_similarity( |
| 79 | actual.float().flatten(), expected.flatten(), dim=0 |
| 80 | ).item() |
| 81 | if cosine <= 0.999: |
| 82 | raise AssertionError(f"projection cosine similarity is {cosine}") |
| 83 | |
| 84 | |
| 85 | def _sequential_projection( |
| 86 | hidden_states: torch.Tensor, |
| 87 | packed_weight: torch.Tensor, |
| 88 | f_b_weight: torch.Tensor, |
| 89 | ) -> ProjectionOutput: |
| 90 | projected = try_low_latency_gemm(hidden_states, packed_weight) |
| 91 | if projected is None: |
| 92 | projected = torch.mm(hidden_states, packed_weight.t()) |
| 93 | f_a = projected[:, 6144:6272] |
| 94 | g1 = torch.mm(f_a, f_b_weight.t()) |
| 95 | return projected[:, :6144], g1, projected[:, 6272:6284] |
| 96 | |
| 97 | |
| 98 | def main() -> None: |
| 99 | parser = argparse.ArgumentParser(description=__doc__) |
| 100 | parser.add_argument("--tokens", type=int, nargs="+", default=range(1, 17)) |
| 101 | parser.add_argument("--weight-sets", type=int, default=4) |
| 102 | parser.add_argument("--repeats", type=int, default=300) |
| 103 | parser.add_argument("--trials", type=int, default=5) |
| 104 | args = parser.parse_args() |
| 105 | |
| 106 | if torch.cuda.get_device_capability() != (10, 3): |
| 107 | raise SystemExit("This benchmark requires an SM103 GPU") |
| 108 | if not shape_dynamic_skinny_gemm.is_available(): |
| 109 | raise SystemExit("This benchmark requires CuTe DSL") |
| 110 | required_ops = ("dsv3_fused_a_gemm",) |
| 111 | if not all(hasattr(torch.ops._C, name) for name in required_ops): |
| 112 | raise SystemExit("This benchmark requires the KDA projection GEMMs") |
| 113 | if not all(1 <= num_tokens <= 16 for num_tokens in args.tokens): |
| 114 | raise SystemExit("--tokens must be in [1, 16]") |
| 115 | |
| 116 | torch.manual_seed(42) |
| 117 | packed_weights = [ |
| 118 | torch.randn(6288, 7168, dtype=torch.bfloat16, device="cuda") |
| 119 | for _ in range(args.weight_sets) |
| 120 | ] |
| 121 | f_b_weights = [ |
| 122 | torch.randn(1536, 128, dtype=torch.bfloat16, device="cuda") |
| 123 | for _ in range(args.weight_sets) |
| 124 | ] |
| 125 | |
| 126 | from flashinfer.autotuner import autotune |
| 127 | from flashinfer.gemm import mm_bf16 |
| 128 | |
| 129 | with torch.inference_mode(), autotune(tune_mode=True): |
| 130 | mm_bf16( |
| 131 | torch.empty( |
| 132 | (KDA_PROJECTION_OVERLAP_MAX_TOKENS, 7168), |
| 133 | dtype=torch.bfloat16, |
| 134 | device="cuda", |
| 135 | ), |
| 136 | packed_weights[0][:6144].t(), |
| 137 | pdl=True, |
| 138 | backend="cute-dsl", |
| 139 | ) |
| 140 | |
| 141 | print("M sequential (us) multistream (us) fastest") |
| 142 | for num_tokens in args.tokens: |
| 143 | hidden_states = [ |
| 144 | torch.randn( |
| 145 | num_tokens, |
| 146 | 7168, |
| 147 | dtype=torch.bfloat16, |
| 148 | device="cuda", |
| 149 | ) |
| 150 | for _ in range(args.weight_sets) |
| 151 | ] |
| 152 | |
| 153 | def sequential( |
| 154 | hidden_states: list[torch.Tensor] = hidden_states, |
| 155 | ) -> list[ProjectionOutput]: |
| 156 | return [ |
| 157 | _sequential_projection(x, weight, f_b_weight) |
| 158 | for x, weight, f_b_weight in zip( |
| 159 | hidden_states, packed_weights, f_b_weights |
| 160 | ) |
| 161 | ] |
| 162 | |
| 163 | candidates: list[tuple[str, Callable[[], list[ProjectionOutput]]]] = [ |
| 164 | ("sequential", sequential) |
| 165 | ] |
| 166 | if num_tokens <= KDA_PROJECTION_OVERLAP_MAX_TOKENS: |
| 167 | aux_stream = torch.cuda.Stream() |
| 168 | events = [ |
| 169 | (torch.cuda.Event(), torch.cuda.Event()) |
| 170 | for _ in range(args.weight_sets) |
| 171 | ] |
| 172 | |
| 173 | def overlap( |
| 174 | hidden_states: list[torch.Tensor] = hidden_states, |
| 175 | aux_stream: torch.cuda.Stream = aux_stream, |
| 176 | events: list[tuple[torch.cuda.Event, torch.cuda.Event]] = events, |
| 177 | ) -> list[ProjectionOutput]: |
| 178 | return [ |
| 179 | run_kda_projection_overlap(x, weight, f_b_weight, aux_stream, pair) |
| 180 | for x, weight, f_b_weight, pair in zip( |
| 181 | hidden_states, packed_weights, f_b_weights, events |
| 182 | ) |
| 183 | ] |
| 184 | |
| 185 | candidates.append(("multistream", overlap)) |
| 186 | |
| 187 | graphs = {} |
| 188 | for name, candidate in candidates: |
| 189 | graph, outputs = _capture(candidate) |
| 190 | graph.replay() |
| 191 | torch.accelerator.synchronize() |
| 192 | _validate(outputs[0], hidden_states[0], packed_weights[0], f_b_weights[0]) |
| 193 | graphs[name] = graph |
| 194 | |
| 195 | samples = {name: [] for name in graphs} |
| 196 | graph_items = list(graphs.items()) |
| 197 | for trial in range(args.trials): |
| 198 | ordered_graphs = graph_items if trial % 2 == 0 else graph_items[::-1] |
| 199 | for name, graph in ordered_graphs: |
| 200 | samples[name].append( |
| 201 | _benchmark_graph(graph, args.weight_sets, args.repeats) |
| 202 | ) |
| 203 | results = {name: statistics.median(values) for name, values in samples.items()} |
| 204 | |
| 205 | sequential_us = results["sequential"] |
| 206 | multistream_us = results.get("multistream") |
| 207 | multistream_text = ( |
| 208 | f"{multistream_us:18.3f}" |
| 209 | if multistream_us is not None |
| 210 | else " -" |
| 211 | ) |
| 212 | fastest = ( |
| 213 | "multistream" |
| 214 | if multistream_us is not None and multistream_us < sequential_us |
| 215 | else "sequential" |
| 216 | ) |
| 217 | print(f"{num_tokens:2d} {sequential_us:15.3f}{multistream_text} {fastest}") |
| 218 | |
| 219 | |
| 220 | if __name__ == "__main__": |
| 221 | main() |