@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 236 lines 6.97 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """Benchmark vLLM Kimi-K3 AttnRes kernels on NVIDIA Blackwell.
4
5 The benchmark uses a zero delta so repeated in-place launches remain stable.
6
7 Example:
8
9 .. code-block:: console
10
11 .venv/bin/python benchmarks/kernels/benchmark_kimi_k3_attn_res.py \
12 --output /tmp/attn_res.json
13
14 """
15
16 from __future__ import annotations
17
18 import argparse
19 import json
20 from collections.abc import Callable
21 from dataclasses import asdict, dataclass
22 from functools import partial
23 from pathlib import Path
24
25 import torch
26
27 import vllm._custom_ops # noqa: F401 - loads torch.ops._C
28 from vllm.models.kimi_k3.nvidia.ops.attn_res import _attn_res_kernel
29 from vllm.platforms import current_platform
30 from vllm.triton_utils import triton
31
32 HIDDEN_SIZE = 7168
33 MAX_BLOCKS = 8
34 EPS = 1e-5
35
36
37 @dataclass
38 class Inputs:
39 prefix: torch.Tensor
40 delta: torch.Tensor
41 blocks: torch.Tensor
42 norm_weight: torch.Tensor
43 qk_weight: torch.Tensor
44 output_norm_weight: torch.Tensor
45
46
47 @dataclass
48 class Result:
49 num_tokens: int
50 num_blocks: int
51 backend: str
52 median_us: float
53 p20_us: float
54 p80_us: float
55
56
57 def parse_args() -> argparse.Namespace:
58 parser = argparse.ArgumentParser(description=__doc__)
59 parser.add_argument(
60 "--num-tokens",
61 type=int,
62 nargs="+",
63 default=[1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 1024, 2048, 4096, 8192],
64 )
65 parser.add_argument("--num-blocks", type=int, nargs="+", default=[1, 4, 8])
66 parser.add_argument("--warmup-ms", type=int, default=50)
67 parser.add_argument("--repeat-ms", type=int, default=200)
68 parser.add_argument("--output", type=Path)
69 return parser.parse_args()
70
71
72 def make_inputs(num_tokens: int) -> Inputs:
73 prefix = torch.randn(num_tokens, HIDDEN_SIZE, device="cuda", dtype=torch.bfloat16)
74 blocks = torch.randn(
75 num_tokens,
76 MAX_BLOCKS,
77 HIDDEN_SIZE,
78 device="cuda",
79 dtype=torch.bfloat16,
80 )
81 norm_weight = torch.ones(HIDDEN_SIZE, device="cuda", dtype=torch.bfloat16)
82 qk_weight = (
83 torch.randn(HIDDEN_SIZE, device="cuda", dtype=torch.bfloat16) / HIDDEN_SIZE**0.5
84 )
85 output_norm_weight = 1 + 0.1 * torch.randn(
86 HIDDEN_SIZE, device="cuda", dtype=torch.bfloat16
87 )
88 return Inputs(
89 prefix=prefix,
90 delta=torch.zeros_like(prefix),
91 blocks=blocks,
92 norm_weight=norm_weight,
93 qk_weight=qk_weight,
94 output_norm_weight=output_norm_weight,
95 )
96
97
98 def launch_triton(inputs: Inputs, output: torch.Tensor, num_blocks: int) -> None:
99 num_tokens = inputs.prefix.shape[0]
100 if num_tokens >= 256 or num_blocks <= 1:
101 block_l, num_warps = 1, 4
102 else:
103 block_l, num_warps = 4, 8
104 _attn_res_kernel[(num_tokens,)](
105 inputs.prefix,
106 inputs.delta,
107 inputs.blocks,
108 inputs.norm_weight,
109 inputs.qk_weight,
110 inputs.output_norm_weight,
111 output,
112 inputs.prefix.stride(0),
113 inputs.delta.stride(0),
114 inputs.blocks.stride(0),
115 inputs.blocks.stride(1),
116 output.stride(0),
117 num_blocks,
118 HIDDEN_SIZE,
119 -1,
120 EPS,
121 EPS,
122 HAS_DELTA=True,
123 WRITE_BLOCK=False,
124 APPLY_OUTPUT_NORM=True,
125 BLOCK_L=block_l,
126 BLOCK_D=triton.next_power_of_2(HIDDEN_SIZE),
127 num_warps=num_warps,
128 num_stages=2,
129 launch_pdl=current_platform.is_arch_support_pdl(),
130 )
131
132
133 def launch_cuda(inputs: Inputs, output: torch.Tensor, num_blocks: int) -> None:
134 torch.ops._C.kimi_k3_attn_res(
135 inputs.prefix,
136 inputs.delta,
137 inputs.blocks,
138 inputs.norm_weight,
139 inputs.qk_weight,
140 inputs.output_norm_weight,
141 output,
142 num_blocks,
143 -1,
144 EPS,
145 EPS,
146 )
147
148
149 def benchmark(
150 function: Callable[[], None], warmup_ms: int, repeat_ms: int
151 ) -> tuple[float, float, float]:
152 median_ms, p20_ms, p80_ms = triton.testing.do_bench(
153 function,
154 warmup=warmup_ms,
155 rep=repeat_ms,
156 quantiles=[0.5, 0.2, 0.8],
157 )
158 return median_ms * 1000, p20_ms * 1000, p80_ms * 1000
159
160
161 def validate(inputs: Inputs, num_blocks: int) -> None:
162 outputs = [torch.empty_like(inputs.prefix) for _ in range(2)]
163 launch_triton(inputs, outputs[0], num_blocks)
164 launch_cuda(inputs, outputs[1], num_blocks)
165 torch.accelerator.synchronize()
166 torch.testing.assert_close(outputs[1], outputs[0], atol=8e-2, rtol=3e-2)
167
168
169 def main() -> None:
170 args = parse_args()
171 if torch.cuda.get_device_capability()[0] != 10:
172 raise SystemExit("The vLLM CUDA AttnRes kernel requires the SM100 family")
173 if not hasattr(torch.ops._C, "kimi_k3_attn_res"):
174 raise SystemExit("This vLLM build does not contain kimi_k3_attn_res")
175 if any(not 0 <= blocks <= MAX_BLOCKS for blocks in args.num_blocks):
176 raise SystemExit(f"--num-blocks values must be in [0, {MAX_BLOCKS}]")
177
178 torch.manual_seed(0)
179 check_inputs = make_inputs(5)
180 for num_blocks in args.num_blocks:
181 validate(check_inputs, num_blocks)
182 print("correctness: all selected block counts passed", flush=True)
183
184 results: list[Result] = []
185 for num_blocks in args.num_blocks:
186 for num_tokens in args.num_tokens:
187 inputs = make_inputs(num_tokens)
188 output = torch.empty_like(inputs.prefix)
189 candidates = {
190 "vllm_triton": partial(launch_triton, inputs, output, num_blocks),
191 "vllm_cuda": partial(launch_cuda, inputs, output, num_blocks),
192 }
193 row: dict[str, float] = {}
194 for backend, candidate in candidates.items():
195 median_us, p20_us, p80_us = benchmark(
196 candidate, args.warmup_ms, args.repeat_ms
197 )
198 results.append(
199 Result(
200 num_tokens=num_tokens,
201 num_blocks=num_blocks,
202 backend=backend,
203 median_us=median_us,
204 p20_us=p20_us,
205 p80_us=p80_us,
206 )
207 )
208 row[backend] = median_us
209 print(
210 f"blocks={num_blocks} tokens={num_tokens:5d} "
211 f"triton={row['vllm_triton']:9.2f} us "
212 f"cuda={row['vllm_cuda']:9.2f} us",
213 flush=True,
214 )
215
216 if args.output is not None:
217 metadata = {
218 "device": torch.cuda.get_device_name(),
219 "capability": torch.cuda.get_device_capability(),
220 "torch": torch.__version__,
221 "cuda": torch.version.cuda,
222 "warmup_ms": args.warmup_ms,
223 "repeat_ms": args.repeat_ms,
224 }
225 args.output.write_text(
226 json.dumps(
227 {"metadata": metadata, "results": [asdict(item) for item in results]},
228 indent=2,
229 )
230 + "\n"
231 )
232 print(f"wrote {args.output}")
233
234
235 if __name__ == "__main__":
236 main()