| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Benchmark the Kimi-K3 latent MoE addmm against CuTe residual GEMM. |
| 4 | |
| 5 | The benchmark covers ``BF16[M, 3584] @ BF16[7168, 3584].T + BF16[M, 7168]`` |
| 6 | with FP32 accumulation and BF16 output. Both backends execute through CUDA |
| 7 | Graph replay. Weights and residuals rotate across buffers exceeding L2 so the |
| 8 | comparison models the full latent MoE projection-and-add path. |
| 9 | """ |
| 10 | |
| 11 | from __future__ import annotations |
| 12 | |
| 13 | import argparse |
| 14 | import dataclasses |
| 15 | import importlib.util |
| 16 | import json |
| 17 | import math |
| 18 | import statistics |
| 19 | from collections.abc import Callable, Sequence |
| 20 | from pathlib import Path |
| 21 | from typing import Any |
| 22 | |
| 23 | import cutlass |
| 24 | import cutlass.cute as cute |
| 25 | import torch |
| 26 | from cuda.bindings import driver as cuda |
| 27 | from cuda.bindings.driver import CUstream |
| 28 | from quack.compile_utils import make_fake_tensor |
| 29 | |
| 30 | N = 7168 |
| 31 | K = 3584 |
| 32 | |
| 33 | |
| 34 | @dataclasses.dataclass(frozen=True, slots=True) |
| 35 | class Config: |
| 36 | block_size: int |
| 37 | outputs_per_block: int |
| 38 | k_unroll: int |
| 39 | vector_width: int = 8 |
| 40 | |
| 41 | |
| 42 | def parse_config(value: str) -> Config: |
| 43 | try: |
| 44 | parts = [int(part) for part in value.split(",")] |
| 45 | except ValueError as error: |
| 46 | raise argparse.ArgumentTypeError( |
| 47 | "config must be BLOCK,OUTPUTS,K_UNROLL[,VECTOR_WIDTH]" |
| 48 | ) from error |
| 49 | if len(parts) == 3: |
| 50 | return Config(*parts) |
| 51 | if len(parts) == 4: |
| 52 | return Config(*parts) |
| 53 | raise argparse.ArgumentTypeError( |
| 54 | "config must be BLOCK,OUTPUTS,K_UNROLL[,VECTOR_WIDTH]" |
| 55 | ) |
| 56 | |
| 57 | |
| 58 | def production_residual_config(m: int) -> Config | None: |
| 59 | """The measured Latent-MoE residual config for M, from the K3 table.""" |
| 60 | from vllm.models.kimi_k3.nvidia.low_latency_gemm import KIMI_K3_PROJECTIONS |
| 61 | |
| 62 | spec = KIMI_K3_PROJECTIONS.get((N, K)) |
| 63 | config = spec.residual_config(m) if spec is not None else None |
| 64 | if config is None: |
| 65 | return None |
| 66 | return Config( |
| 67 | config.block_size, |
| 68 | config.outputs_per_block, |
| 69 | config.k_unroll, |
| 70 | config.vector_width, |
| 71 | ) |
| 72 | |
| 73 | |
| 74 | def candidate_configs(mode: str, selected: Config | None, m: int) -> list[Config]: |
| 75 | if mode == "selected": |
| 76 | if selected is not None: |
| 77 | return [selected] |
| 78 | # No explicit --config: fall back to the production table for this M. |
| 79 | config = production_residual_config(m) |
| 80 | return [config] if config is not None else [] |
| 81 | if mode == "baseline": |
| 82 | return [Config(224, 4, 2)] |
| 83 | return [ |
| 84 | Config(block_size, outputs_per_block, k_unroll, vector_width) |
| 85 | for vector_width in (4, 8) |
| 86 | for block_size in (32, 64, 128, 224, 448) |
| 87 | if block_size % 32 == 0 and K % (block_size * vector_width) == 0 |
| 88 | for outputs_per_block in (1, 2, 4, 7, 8) |
| 89 | if N % outputs_per_block == 0 |
| 90 | for k_unroll in (1, 2, 4) |
| 91 | ] |
| 92 | |
| 93 | |
| 94 | def load_kernel_class(path: Path): |
| 95 | spec = importlib.util.spec_from_file_location("cute_skinny_device", path) |
| 96 | if spec is None or spec.loader is None: |
| 97 | raise RuntimeError(f"cannot load CuTe kernel from {path}") |
| 98 | module = importlib.util.module_from_spec(spec) |
| 99 | spec.loader.exec_module(module) |
| 100 | return module.CuteSkinnyGemm |
| 101 | |
| 102 | |
| 103 | def stream() -> CUstream: |
| 104 | return CUstream(torch.cuda.current_stream().cuda_stream) |
| 105 | |
| 106 | |
| 107 | def compile_kernel(kernel_class, m: int, config: Config, max_registers: int): |
| 108 | element_type = cutlass.BFloat16 |
| 109 | n = cute.sym_int(divisibility=config.outputs_per_block) |
| 110 | k = cute.sym_int(divisibility=config.block_size * config.vector_width) |
| 111 | a = make_fake_tensor(element_type, (m, k), divisibility=config.vector_width) |
| 112 | b = make_fake_tensor(element_type, (n, k), divisibility=config.vector_width) |
| 113 | residual = make_fake_tensor(element_type, (m, n), divisibility=1) |
| 114 | c = make_fake_tensor(element_type, (m, n), divisibility=1) |
| 115 | kernel = kernel_class( |
| 116 | element_type=element_type, |
| 117 | num_rows=m, |
| 118 | block_size=config.block_size, |
| 119 | outputs_per_block=config.outputs_per_block, |
| 120 | vector_width=config.vector_width, |
| 121 | k_unroll=config.k_unroll, |
| 122 | has_residual=True, |
| 123 | use_pdl=True, |
| 124 | ) |
| 125 | return cute.compile( |
| 126 | kernel, |
| 127 | a, |
| 128 | b, |
| 129 | residual, |
| 130 | c, |
| 131 | stream(), |
| 132 | options=( |
| 133 | "--enable-tvm-ffi --keep-cubin " |
| 134 | f"--ptxas-options -maxrregcount={max_registers} " |
| 135 | "--ptxas-options -lineinfo" |
| 136 | ), |
| 137 | ) |
| 138 | |
| 139 | |
| 140 | def resource_usage(compiled) -> dict[str, Any]: |
| 141 | executor = getattr(compiled, "_default_executor", None) |
| 142 | context = getattr(executor, "exec_context", None) |
| 143 | functions = getattr(context, "kernel_functions", None) |
| 144 | if not functions: |
| 145 | return {"resource_metrics_available": False} |
| 146 | |
| 147 | def attribute(name, function) -> int: |
| 148 | error, value = cuda.cuFuncGetAttribute(name, function) |
| 149 | if error != cuda.CUresult.CUDA_SUCCESS: |
| 150 | raise RuntimeError(f"cuFuncGetAttribute failed with {error}") |
| 151 | return int(value) |
| 152 | |
| 153 | registers = [ |
| 154 | attribute(cuda.CUfunction_attribute.CU_FUNC_ATTRIBUTE_NUM_REGS, function) |
| 155 | for function in functions |
| 156 | ] |
| 157 | local_bytes = [ |
| 158 | attribute( |
| 159 | cuda.CUfunction_attribute.CU_FUNC_ATTRIBUTE_LOCAL_SIZE_BYTES, |
| 160 | function, |
| 161 | ) |
| 162 | for function in functions |
| 163 | ] |
| 164 | return { |
| 165 | "resource_metrics_available": True, |
| 166 | "registers_per_thread": max(registers, default=0), |
| 167 | "spill_bytes": max(local_bytes, default=0), |
| 168 | } |
| 169 | |
| 170 | |
| 171 | def rotating_buffer_count(m: int, multiplier: float, limit: int) -> int: |
| 172 | properties = torch.cuda.get_device_properties(0) |
| 173 | bytes_per_pair = (N * K + m * N) * 2 |
| 174 | target = math.ceil(multiplier * properties.L2_cache_size) |
| 175 | return max(2, min(limit, math.ceil(target / bytes_per_pair))) |
| 176 | |
| 177 | |
| 178 | def graph_samples( |
| 179 | launch: Callable[[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor], None], |
| 180 | activation: torch.Tensor, |
| 181 | weights: Sequence[torch.Tensor], |
| 182 | residuals: Sequence[torch.Tensor], |
| 183 | repeats: int, |
| 184 | replays: int, |
| 185 | ) -> tuple[list[float], list[torch.Tensor]]: |
| 186 | outputs = [torch.empty_like(residual) for residual in residuals] |
| 187 | for weight, residual, output in zip(weights, residuals, outputs): |
| 188 | launch(activation, weight, residual, output) |
| 189 | torch.accelerator.synchronize() |
| 190 | |
| 191 | graph = torch.cuda.CUDAGraph() |
| 192 | with torch.cuda.graph(graph): |
| 193 | for weight, residual, output in zip(weights, residuals, outputs): |
| 194 | launch(activation, weight, residual, output) |
| 195 | for _ in range(20): |
| 196 | graph.replay() |
| 197 | torch.accelerator.synchronize() |
| 198 | |
| 199 | samples = [] |
| 200 | for _ in range(repeats): |
| 201 | start = torch.cuda.Event(enable_timing=True) |
| 202 | end = torch.cuda.Event(enable_timing=True) |
| 203 | start.record() |
| 204 | for _ in range(replays): |
| 205 | graph.replay() |
| 206 | end.record() |
| 207 | end.synchronize() |
| 208 | samples.append(start.elapsed_time(end) * 1000.0 / (replays * len(weights))) |
| 209 | return samples, outputs |
| 210 | |
| 211 | |
| 212 | def summarize(samples: Sequence[float]) -> dict[str, Any]: |
| 213 | ordered = sorted(samples) |
| 214 | |
| 215 | def percentile(fraction: float) -> float: |
| 216 | position = fraction * (len(ordered) - 1) |
| 217 | lower = math.floor(position) |
| 218 | upper = math.ceil(position) |
| 219 | if lower == upper: |
| 220 | return ordered[lower] |
| 221 | weight = position - lower |
| 222 | return ordered[lower] * (1.0 - weight) + ordered[upper] * weight |
| 223 | |
| 224 | mean = statistics.mean(samples) |
| 225 | return { |
| 226 | "median_us": statistics.median(samples), |
| 227 | "p10_us": percentile(0.1), |
| 228 | "p90_us": percentile(0.9), |
| 229 | "mean_us": mean, |
| 230 | "cv_pct": statistics.pstdev(samples) / mean * 100.0, |
| 231 | "samples_us": list(samples), |
| 232 | } |
| 233 | |
| 234 | |
| 235 | def correctness( |
| 236 | output: torch.Tensor, |
| 237 | activation: torch.Tensor, |
| 238 | weight: torch.Tensor, |
| 239 | residual: torch.Tensor, |
| 240 | ) -> dict[str, Any]: |
| 241 | actual = output.float() |
| 242 | reference = activation.float() @ weight.float().t() + residual.float() |
| 243 | error = (actual - reference).abs() |
| 244 | scaled_error = error / (reference.abs() + 1.0) |
| 245 | cosine = torch.nn.functional.cosine_similarity( |
| 246 | actual.flatten(), reference.flatten(), dim=0 |
| 247 | ).item() |
| 248 | return { |
| 249 | "valid": cosine > 0.999, |
| 250 | "cosine": cosine, |
| 251 | "max_abs_error": error.max().item(), |
| 252 | "max_scaled_error": scaled_error.max().item(), |
| 253 | } |
| 254 | |
| 255 | |
| 256 | def main() -> None: |
| 257 | parser = argparse.ArgumentParser(description=__doc__) |
| 258 | parser.add_argument("--kernel", type=Path, required=True) |
| 259 | parser.add_argument("--output", type=Path, required=True) |
| 260 | parser.add_argument( |
| 261 | "--mode", choices=("baseline", "sweep", "selected"), default="baseline" |
| 262 | ) |
| 263 | parser.add_argument("--config", type=parse_config) |
| 264 | parser.add_argument("--m", type=int, action="append") |
| 265 | parser.add_argument("--config-shard", type=int, default=0) |
| 266 | parser.add_argument("--num-config-shards", type=int, default=1) |
| 267 | parser.add_argument("--repeats", type=int, default=21) |
| 268 | parser.add_argument("--replays", type=int, default=200) |
| 269 | parser.add_argument("--cache-multiplier", type=float, default=3.0) |
| 270 | parser.add_argument("--max-buffers", type=int, default=32) |
| 271 | parser.add_argument("--max-registers", type=int, default=64) |
| 272 | args = parser.parse_args() |
| 273 | |
| 274 | token_counts = args.m or list(range(1, 17)) |
| 275 | if any(not 1 <= m <= 16 for m in token_counts): |
| 276 | raise ValueError("expected 1 <= M <= 16") |
| 277 | if not 0 <= args.config_shard < args.num_config_shards: |
| 278 | raise ValueError("config shard must be in [0, num_config_shards)") |
| 279 | torch.accelerator.set_device_index(0) |
| 280 | if torch.cuda.get_device_capability() != (10, 3): |
| 281 | raise RuntimeError("this benchmark requires SM103") |
| 282 | |
| 283 | kernel_class = load_kernel_class(args.kernel) |
| 284 | properties = torch.cuda.get_device_properties(0) |
| 285 | metadata = { |
| 286 | "device": properties.name, |
| 287 | "compute_capability": list(torch.cuda.get_device_capability()), |
| 288 | "torch_version": torch.__version__, |
| 289 | "cuda_version": torch.version.cuda, |
| 290 | } |
| 291 | args.output.parent.mkdir(parents=True, exist_ok=True) |
| 292 | with args.output.open("w", encoding="utf-8") as output_file: |
| 293 | for m in token_counts: |
| 294 | configs = candidate_configs(args.mode, args.config, m) |
| 295 | torch.manual_seed(20260722 + m) |
| 296 | count = rotating_buffer_count(m, args.cache_multiplier, args.max_buffers) |
| 297 | activation = torch.randn((m, K), device="cuda", dtype=torch.bfloat16) |
| 298 | weights = [ |
| 299 | torch.randn((N, K), device="cuda", dtype=torch.bfloat16) |
| 300 | for _ in range(count) |
| 301 | ] |
| 302 | residuals = [ |
| 303 | torch.randn((m, N), device="cuda", dtype=torch.bfloat16) |
| 304 | for _ in range(count) |
| 305 | ] |
| 306 | candidates: list[tuple[str, Config | None]] = [("cublas_addmm", None)] |
| 307 | candidates.extend( |
| 308 | ("cute_residual", config) |
| 309 | for index, config in enumerate(configs) |
| 310 | if index % args.num_config_shards == args.config_shard |
| 311 | ) |
| 312 | for backend, config in candidates: |
| 313 | row: dict[str, Any] = { |
| 314 | "m": m, |
| 315 | "n": N, |
| 316 | "k": K, |
| 317 | "backend": backend, |
| 318 | "mode": args.mode, |
| 319 | "config": dataclasses.asdict(config) if config else {}, |
| 320 | "num_buffers": count, |
| 321 | "cache_multiplier": args.cache_multiplier, |
| 322 | **metadata, |
| 323 | } |
| 324 | try: |
| 325 | if backend == "cublas_addmm": |
| 326 | launch = lambda a, b, residual, c: torch.addmm( |
| 327 | residual, a, b.t(), out=c |
| 328 | ) |
| 329 | else: |
| 330 | if config is None: |
| 331 | raise AssertionError("missing CuTe config") |
| 332 | compiled = compile_kernel( |
| 333 | kernel_class, m, config, args.max_registers |
| 334 | ) |
| 335 | launch = lambda a, b, residual, c, fn=compiled: fn( |
| 336 | a, b, residual, c, stream() |
| 337 | ) |
| 338 | row.update(resource_usage(compiled)) |
| 339 | samples, outputs = graph_samples( |
| 340 | launch, |
| 341 | activation, |
| 342 | weights, |
| 343 | residuals, |
| 344 | args.repeats, |
| 345 | args.replays, |
| 346 | ) |
| 347 | row.update( |
| 348 | correctness(outputs[0], activation, weights[0], residuals[0]) |
| 349 | ) |
| 350 | row.update(summarize(samples)) |
| 351 | except Exception as error: # noqa: BLE001 |
| 352 | row.update( |
| 353 | { |
| 354 | "valid": False, |
| 355 | "error": f"{type(error).__name__}: {error}", |
| 356 | } |
| 357 | ) |
| 358 | output_file.write(json.dumps(row, sort_keys=True) + "\n") |
| 359 | output_file.flush() |
| 360 | print(json.dumps(row, sort_keys=True), flush=True) |
| 361 | |
| 362 | del activation, weights, residuals |
| 363 | torch.accelerator.empty_cache() |
| 364 | |
| 365 | |
| 366 | if __name__ == "__main__": |
| 367 | main() |