| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Microbenchmark for the fused Kimi-K3 KDA decode kernel on ROCm. |
| 4 | |
| 5 | Compares the single fused launch against the three-kernel Triton chain it |
| 6 | replaces (packed causal conv1d update -> recurrent delta-rule decode -> gated |
| 7 | output RMSNorm), at Kimi-K3 KDA shapes: 96 heads x 128, conv width 4, |
| 8 | gate_lower_bound -5.0. Per-step figures scale the per-layer number by the 69 |
| 9 | KDA layers in the model. |
| 10 | |
| 11 | The recurrent state dominates the traffic (num_heads x 128 x 128 fp32 read and |
| 12 | written per token), so the reported bandwidth is the useful metric: the fusion |
| 13 | removes launches and the intermediate QKV / core-output round trips, not the |
| 14 | state traffic itself. |
| 15 | |
| 16 | Example: |
| 17 | python benchmarks/kernels/benchmark_kimi_k3_kda_decode.py \ |
| 18 | --tokens 1 8 32 64 128 --heads 12 |
| 19 | |
| 20 | """ |
| 21 | |
| 22 | import argparse |
| 23 | import functools |
| 24 | |
| 25 | import torch |
| 26 | |
| 27 | from vllm.triton_utils import triton |
| 28 | |
| 29 | HEAD_DIM = 128 |
| 30 | CONV_WIDTH = 4 |
| 31 | GATE_LOWER_BOUND = -5.0 |
| 32 | NORM_EPS = 1e-5 |
| 33 | NUM_KDA_LAYERS = 69 |
| 34 | DTYPE = torch.bfloat16 |
| 35 | |
| 36 | |
| 37 | def _bench(fn) -> float: |
| 38 | return triton.testing.do_bench(fn, warmup=50, rep=300, return_mode="median") |
| 39 | |
| 40 | |
| 41 | def _bench_graph_layers(calls: list) -> float: |
| 42 | """Per-layer milliseconds for a graph holding one call per KDA layer.""" |
| 43 | stream = torch.cuda.Stream() |
| 44 | stream.wait_stream(torch.cuda.current_stream()) |
| 45 | with torch.cuda.stream(stream): |
| 46 | for call in calls[:3]: |
| 47 | call() |
| 48 | torch.cuda.current_stream().wait_stream(stream) |
| 49 | graph = torch.cuda.CUDAGraph() |
| 50 | with torch.cuda.graph(graph): |
| 51 | for call in calls: |
| 52 | call() |
| 53 | total = triton.testing.do_bench( |
| 54 | graph.replay, warmup=50, rep=300, return_mode="median" |
| 55 | ) |
| 56 | return total / len(calls) |
| 57 | |
| 58 | |
| 59 | def _bench_graph(fn, repeats: int = NUM_KDA_LAYERS) -> float: |
| 60 | """Per-call milliseconds under CUDA-graph replay, as decode actually runs. |
| 61 | |
| 62 | Eager timings credit a fusion for every Python dispatch it removes; inside a |
| 63 | captured graph those are gone, so this is the number that decides whether |
| 64 | the kernel is worth it in the server. The graph holds `repeats` calls (one |
| 65 | per KDA layer) because a one-call graph is swamped by the ~19 us HIP |
| 66 | graph-launch overhead, which a real 93-layer graph amortises away. |
| 67 | """ |
| 68 | stream = torch.cuda.Stream() |
| 69 | stream.wait_stream(torch.cuda.current_stream()) |
| 70 | with torch.cuda.stream(stream): |
| 71 | for _ in range(3): |
| 72 | fn() |
| 73 | torch.cuda.current_stream().wait_stream(stream) |
| 74 | graph = torch.cuda.CUDAGraph() |
| 75 | with torch.cuda.graph(graph): |
| 76 | for _ in range(repeats): |
| 77 | fn() |
| 78 | total = triton.testing.do_bench( |
| 79 | graph.replay, warmup=50, rep=300, return_mode="median" |
| 80 | ) |
| 81 | return total / repeats |
| 82 | |
| 83 | |
| 84 | class Inputs: |
| 85 | def __init__(self, num_tokens: int, num_heads: int, conv_layout: str) -> None: |
| 86 | torch.manual_seed(0) |
| 87 | device = "cuda" |
| 88 | dim = num_heads * HEAD_DIM |
| 89 | num_slots = num_tokens + 8 |
| 90 | self.num_heads = num_heads |
| 91 | self.num_tokens = num_tokens |
| 92 | self.mixed_qkv = torch.randn(num_tokens, 3 * dim, device=device, dtype=DTYPE) |
| 93 | self.conv_weights = torch.randn( |
| 94 | 3 * dim, CONV_WIDTH, device=device, dtype=torch.float32 |
| 95 | ) |
| 96 | self.decode_conv1d_weight = torch.stack( |
| 97 | [ |
| 98 | self.conv_weights[i * dim : (i + 1) * dim].transpose(0, 1).contiguous() |
| 99 | for i in range(3) |
| 100 | ] |
| 101 | ) |
| 102 | if conv_layout == "DS": |
| 103 | self.conv_state_t = torch.randn( |
| 104 | num_slots, 3 * dim, CONV_WIDTH - 1, device=device, dtype=DTYPE |
| 105 | ) |
| 106 | else: |
| 107 | conv_state = torch.randn( |
| 108 | num_slots, CONV_WIDTH - 1, 3 * dim, device=device, dtype=DTYPE |
| 109 | ) |
| 110 | self.conv_state_t = conv_state.transpose(-1, -2) |
| 111 | self.recurrent_state = torch.randn( |
| 112 | num_slots, num_heads, HEAD_DIM, HEAD_DIM, device=device, dtype=torch.float32 |
| 113 | ) |
| 114 | self.g1 = torch.randn( |
| 115 | 1, num_tokens, num_heads, HEAD_DIM, device=device, dtype=DTYPE |
| 116 | ) |
| 117 | self.g2 = torch.randn( |
| 118 | num_tokens, num_heads, HEAD_DIM, device=device, dtype=DTYPE |
| 119 | ) |
| 120 | self.beta = torch.randn(1, num_tokens, num_heads, device=device, dtype=DTYPE) |
| 121 | self.A_log = torch.randn(num_heads, device=device, dtype=torch.float32) |
| 122 | self.dt_bias = torch.randn(dim, device=device, dtype=torch.float32) |
| 123 | self.norm_weight_bf16 = torch.ones(HEAD_DIM, device=device, dtype=DTYPE) |
| 124 | self.decode_norm_weight = self.norm_weight_bf16.float() |
| 125 | # Slots start at 1: slot 0 is NULL_BLOCK_ID, which the fused kernel |
| 126 | # treats as a padded row and skips, so timing it measures nothing. |
| 127 | self.state_indices = torch.arange( |
| 128 | 1, num_tokens + 1, device=device, dtype=torch.int32 |
| 129 | ) |
| 130 | self.out = torch.empty( |
| 131 | 1, num_tokens, num_heads, HEAD_DIM, device=device, dtype=DTYPE |
| 132 | ) |
| 133 | self.conv_out = torch.empty_like(self.mixed_qkv) |
| 134 | |
| 135 | def state_bytes(self) -> int: |
| 136 | """Recurrent-state traffic, identical for both implementations.""" |
| 137 | return self.num_tokens * self.num_heads * HEAD_DIM * HEAD_DIM * 4 * 2 |
| 138 | |
| 139 | |
| 140 | def _gated_rmsnorm( |
| 141 | x: torch.Tensor, gate: torch.Tensor, weight: torch.Tensor, eps: float |
| 142 | ) -> torch.Tensor: |
| 143 | x_float = x.float() |
| 144 | variance = x_float.pow(2).mean(dim=-1, keepdim=True) |
| 145 | normed = x_float * torch.rsqrt(variance + eps) * weight.float() |
| 146 | return (normed * torch.sigmoid(gate.float())).to(x.dtype) |
| 147 | |
| 148 | |
| 149 | def _triton_gated_norm(inp: Inputs, core_attn_out: torch.Tensor) -> torch.Tensor: |
| 150 | from vllm.third_party.flash_linear_attention.ops.kda import rms_norm_gated |
| 151 | |
| 152 | return rms_norm_gated( |
| 153 | core_attn_out, |
| 154 | inp.g2, |
| 155 | inp.norm_weight_bf16, |
| 156 | None, |
| 157 | activation="sigmoid", |
| 158 | eps=NORM_EPS, |
| 159 | ) |
| 160 | |
| 161 | |
| 162 | def triton_chain(inp: Inputs, fused_norm: bool = False) -> None: |
| 163 | from vllm.model_executor.layers.mamba.ops.causal_conv1d import ( |
| 164 | causal_conv1d_update, |
| 165 | ) |
| 166 | from vllm.models.kimi_k3.amd.ops.third_party.kda import ( |
| 167 | fused_recurrent_kda_packed_decode, |
| 168 | ) |
| 169 | |
| 170 | causal_conv1d_update( |
| 171 | inp.mixed_qkv, |
| 172 | inp.conv_state_t, |
| 173 | inp.conv_weights, |
| 174 | None, |
| 175 | activation="silu", |
| 176 | conv_state_indices=inp.state_indices, |
| 177 | validate_data=False, |
| 178 | out=inp.conv_out, |
| 179 | ) |
| 180 | core_attn_out, _ = fused_recurrent_kda_packed_decode( |
| 181 | mixed_qkv=inp.conv_out, |
| 182 | raw_g=inp.g1, |
| 183 | raw_beta=inp.beta, |
| 184 | A_log=inp.A_log, |
| 185 | dt_bias=inp.dt_bias, |
| 186 | lower_bound=GATE_LOWER_BOUND, |
| 187 | initial_state=inp.recurrent_state, |
| 188 | state_indices=inp.state_indices, |
| 189 | ) |
| 190 | if fused_norm: |
| 191 | inp.out.copy_(_triton_gated_norm(inp, core_attn_out)) |
| 192 | else: |
| 193 | inp.out.copy_( |
| 194 | _gated_rmsnorm(core_attn_out, inp.g2, inp.norm_weight_bf16, NORM_EPS) |
| 195 | ) |
| 196 | |
| 197 | |
| 198 | def fused(inp: Inputs) -> None: |
| 199 | from vllm import _custom_ops as ops |
| 200 | |
| 201 | ops.fused_kda_decode( |
| 202 | x=inp.mixed_qkv, |
| 203 | weight=inp.decode_conv1d_weight, |
| 204 | bias=None, |
| 205 | conv_state=inp.conv_state_t, |
| 206 | raw_g=inp.g1, |
| 207 | raw_beta=inp.beta, |
| 208 | A_log=inp.A_log, |
| 209 | dt_bias=inp.dt_bias, |
| 210 | state_indices=inp.state_indices, |
| 211 | state=inp.recurrent_state, |
| 212 | out=inp.out, |
| 213 | lower_bound=GATE_LOWER_BOUND, |
| 214 | output_gate=inp.g2, |
| 215 | norm_weight=inp.decode_norm_weight, |
| 216 | norm_eps=NORM_EPS, |
| 217 | ) |
| 218 | |
| 219 | |
| 220 | def main() -> None: |
| 221 | parser = argparse.ArgumentParser(description=__doc__) |
| 222 | parser.add_argument("--tokens", type=int, nargs="+", default=[1, 8, 32, 64, 128]) |
| 223 | parser.add_argument( |
| 224 | "--graph", |
| 225 | action="store_true", |
| 226 | help="time under CUDA-graph replay instead of eager dispatch", |
| 227 | ) |
| 228 | parser.add_argument( |
| 229 | "--layers", |
| 230 | type=int, |
| 231 | default=1, |
| 232 | help=( |
| 233 | "with --graph, give each of N layers its own state buffers so the " |
| 234 | "recurrent state cannot stay resident in Infinity Cache. " |
| 235 | f"Use {NUM_KDA_LAYERS} for a production-shaped step (needs " |
| 236 | "~0.9 GB per layer at 128 tokens / 12 heads)." |
| 237 | ), |
| 238 | ) |
| 239 | parser.add_argument( |
| 240 | "--heads", |
| 241 | type=int, |
| 242 | nargs="+", |
| 243 | default=[12], |
| 244 | help="KDA heads per rank (96 total / TP size)", |
| 245 | ) |
| 246 | parser.add_argument( |
| 247 | "--layout", |
| 248 | choices=("SD", "DS"), |
| 249 | default="SD", |
| 250 | help="physical conv-state cache layout", |
| 251 | ) |
| 252 | args = parser.parse_args() |
| 253 | |
| 254 | if not hasattr(torch.ops._C, "fused_kda_decode"): |
| 255 | raise SystemExit("vLLM was built without the fused KDA decode kernel") |
| 256 | |
| 257 | props = torch.cuda.get_device_properties(0) |
| 258 | bench = _bench_graph if args.graph else _bench |
| 259 | mode = "cuda-graph replay" if args.graph else "eager dispatch" |
| 260 | print( |
| 261 | f"device: {props.name} ({props.gcnArchName}) timing: {mode} " |
| 262 | f"layout: {args.layout}" |
| 263 | ) |
| 264 | print( |
| 265 | f"{'heads':>6} {'tokens':>7} {'eager-norm':>11} {'triton-norm':>12} " |
| 266 | f"{'fused us':>9} {'speedup':>8} {'state TB/s':>11} {'saved ms/step':>14}" |
| 267 | ) |
| 268 | for num_heads in args.heads: |
| 269 | for num_tokens in args.tokens: |
| 270 | if args.graph and args.layers > 1: |
| 271 | layers = [ |
| 272 | Inputs(num_tokens, num_heads, args.layout) |
| 273 | for _ in range(args.layers) |
| 274 | ] |
| 275 | eager_ms = _bench_graph_layers( |
| 276 | [functools.partial(triton_chain, i, False) for i in layers] |
| 277 | ) |
| 278 | triton_ms = _bench_graph_layers( |
| 279 | [functools.partial(triton_chain, i, True) for i in layers] |
| 280 | ) |
| 281 | fused_ms = _bench_graph_layers( |
| 282 | [functools.partial(fused, i) for i in layers] |
| 283 | ) |
| 284 | inp = layers[0] |
| 285 | del layers |
| 286 | torch.accelerator.empty_cache() |
| 287 | else: |
| 288 | inp = Inputs(num_tokens, num_heads, args.layout) |
| 289 | eager_ms = bench(functools.partial(triton_chain, inp, False)) |
| 290 | triton_ms = bench(functools.partial(triton_chain, inp, True)) |
| 291 | fused_ms = bench(functools.partial(fused, inp)) |
| 292 | bandwidth = inp.state_bytes() * 1e3 / fused_ms / 1e12 |
| 293 | saved = (triton_ms - fused_ms) * NUM_KDA_LAYERS |
| 294 | print( |
| 295 | f"{num_heads:>6} {num_tokens:>7} {eager_ms * 1e3:>11.2f} " |
| 296 | f"{triton_ms * 1e3:>12.2f} {fused_ms * 1e3:>9.2f} " |
| 297 | f"{triton_ms / fused_ms:>7.2f}x {bandwidth:>10.2f} {saved:>14.3f}" |
| 298 | ) |
| 299 | print( |
| 300 | f"\n'eager-norm' is the chain as it runs today (FusedRMSNormGated falls " |
| 301 | f"back to ~10 eager ops when custom_ops are off);\n'triton-norm' uses the " |
| 302 | f"Triton rms_norm_gated kernel, which is the honest baseline for this " |
| 303 | f"fusion.\nspeedup and saved ms/step are against 'triton-norm', over " |
| 304 | f"{NUM_KDA_LAYERS} KDA layers per forward pass." |
| 305 | ) |
| 306 | |
| 307 | |
| 308 | if __name__ == "__main__": |
| 309 | main() |