| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import argparse |
| 5 | |
| 6 | import torch |
| 7 | |
| 8 | from vllm import _custom_ops as ops |
| 9 | from vllm.triton_utils import triton |
| 10 | |
| 11 | # DeepSeek V3 dimensions |
| 12 | NOPE_DIM = 512 |
| 13 | ROPE_DIM = 64 |
| 14 | NUM_HEADS = 128 |
| 15 | |
| 16 | NUM_TOKENS = [8, 16, 32, 64, 128, 256, 512, 1024, 2048, 4096, 8192] |
| 17 | |
| 18 | |
| 19 | def get_configs(): |
| 20 | return NUM_TOKENS |
| 21 | |
| 22 | |
| 23 | def make_inputs(num_tokens, dtype): |
| 24 | """Create inputs matching the real code path. |
| 25 | |
| 26 | Args: |
| 27 | num_tokens: Number of tokens to generate inputs for. |
| 28 | dtype: dtype of the generated tensors. |
| 29 | |
| 30 | """ |
| 31 | # Simulate: bmm output [N, B, L].transpose(0, 1) -> [B, N, L] |
| 32 | raw = torch.randn(NUM_HEADS, num_tokens, NOPE_DIM, dtype=dtype, device="cuda") |
| 33 | ql_nope = raw.transpose(0, 1) |
| 34 | |
| 35 | q_pe = torch.randn(num_tokens, NUM_HEADS, ROPE_DIM, dtype=dtype, device="cuda") |
| 36 | return ql_nope, q_pe |
| 37 | |
| 38 | |
| 39 | # ---- Non-contiguous nope benchmark (real code path) ---- |
| 40 | @triton.testing.perf_report( |
| 41 | triton.testing.Benchmark( |
| 42 | x_names=["num_tokens"], |
| 43 | x_vals=get_configs(), |
| 44 | line_arg="provider", |
| 45 | line_vals=["torch_cat", "concat_mla_q"], |
| 46 | line_names=["torch.cat", "concat_mla_q (v8)"], |
| 47 | styles=[("blue", "--"), ("green", "-")], |
| 48 | ylabel="Latency (us)", |
| 49 | plot_name="concat_mla_q-transposed", |
| 50 | args={}, |
| 51 | ) |
| 52 | ) |
| 53 | def bench_transposed(num_tokens, provider): |
| 54 | dtype = torch.bfloat16 |
| 55 | ql_nope, q_pe = make_inputs(num_tokens, dtype) |
| 56 | |
| 57 | q_out = torch.empty( |
| 58 | num_tokens, NUM_HEADS, NOPE_DIM + ROPE_DIM, dtype=dtype, device="cuda" |
| 59 | ) |
| 60 | |
| 61 | quantiles = [0.5, 0.2, 0.8] |
| 62 | |
| 63 | if provider == "torch_cat": |
| 64 | ms, min_ms, max_ms = triton.testing.do_bench_cudagraph( |
| 65 | lambda: torch.cat((ql_nope, q_pe), dim=-1), quantiles=quantiles, rep=500 |
| 66 | ) |
| 67 | else: |
| 68 | ms, min_ms, max_ms = triton.testing.do_bench_cudagraph( |
| 69 | lambda: ops.concat_mla_q(ql_nope, q_pe, q_out), quantiles=quantiles, rep=500 |
| 70 | ) |
| 71 | |
| 72 | return ms * 1000, max_ms * 1000, min_ms * 1000 # us |
| 73 | |
| 74 | |
| 75 | if __name__ == "__main__": |
| 76 | parser = argparse.ArgumentParser(description="Benchmark concat_mla_q vs torch.cat") |
| 77 | parser.add_argument( |
| 78 | "--save-path", type=str, default=None, help="Path to save benchmark results" |
| 79 | ) |
| 80 | args = parser.parse_args() |
| 81 | |
| 82 | print("\n" + "=" * 70) |
| 83 | print("CONCAT MLA Q KERNEL BENCHMARKS") |
| 84 | print("=" * 70) |
| 85 | print(f"Dimensions: nope={NOPE_DIM}, rope={ROPE_DIM}, heads={NUM_HEADS}") |
| 86 | print( |
| 87 | f"Per-head output: {NOPE_DIM + ROPE_DIM} bf16 = " |
| 88 | f"{(NOPE_DIM + ROPE_DIM) * 2} bytes" |
| 89 | ) |
| 90 | print(f"num_tokens (decode=batch_size, prefill=chunk_size): {NUM_TOKENS}") |
| 91 | print("=" * 70) |
| 92 | |
| 93 | print("\n--- Non-contiguous nope inputs (transposed BMM output) ---") |
| 94 | bench_transposed.run(print_data=True, save_path=args.save_path) |
| 95 | |
| 96 | print("\n" + "=" * 70) |
| 97 | print("Benchmarking complete!") |
| 98 | print("=" * 70) |