A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 11, 16:15 UTC
[perf] Use direct copy (broadcast) instead of cat for k_nope/k_pe in MLA prefill (#29710)
Signed-off-by: Ming Yang <minos.future@gmail.com>
Ming Yang committed
Dec 11, 2025 at 00:20 UTC
fba89069302e9b4d0457bc8eeddeeec76f27f0b1
2 files changed
+180
-3
benchmarks/kernels/benchmark_mla_k_concat.py
new
+150
@@ -0,0 +1,150 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""
4
+Benchmark script comparing torch.cat vs direct copy for k_nope/k_pe concatenation
5
+in MLA (Multi-head Latent Attention) prefill.
6
+
7
+This validates that the optimization from commit 8d4142bd is beneficial across
8
+various batch sizes, not just the originally tested batch size of 32768.
9
+"""
10
+
11
+import time
12
+from collections.abc import Callable
13
+
14
+import torch
15
+
16
+# DeepSeek-V3 MLA dimensions
17
+NUM_HEADS = 128
18
+QK_NOPE_HEAD_DIM = 128
19
+PE_DIM = 64
20
+
21
+
22
+def cat_method(k_nope: torch.Tensor, k_pe: torch.Tensor) -> torch.Tensor:
23
+ """Original torch.cat approach with expand."""
24
+ return torch.cat((k_nope, k_pe.expand((*k_nope.shape[:-1], -1))), dim=-1)
25
+
26
+
27
+def direct_copy_method(k_nope: torch.Tensor, k_pe: torch.Tensor) -> torch.Tensor:
28
+ """Optimized direct copy approach (avoids expand + cat overhead)."""
29
+ k = torch.empty(
30
+ (*k_nope.shape[:-1], k_nope.shape[-1] + k_pe.shape[-1]),
31
+ dtype=k_nope.dtype,
32
+ device=k_nope.device,
33
+ )
34
+ k[..., : k_nope.shape[-1]] = k_nope
35
+ k[..., k_nope.shape[-1] :] = k_pe
36
+ return k
37
+
38
+
39
+def benchmark_method(
40
+ method: Callable,
41
+ k_nope: torch.Tensor,
42
+ k_pe: torch.Tensor,
43
+ num_warmup: int = 10,
44
+ num_iters: int = 100,
45
+) -> float:
46
+ """Benchmark a concatenation method and return mean latency in ms."""
47
+ # Warmup
48
+ for _ in range(num_warmup):
49
+ _ = method(k_nope, k_pe)
50
+ torch.cuda.synchronize()
51
+
52
+ # Benchmark
53
+ start = time.perf_counter()
54
+ for _ in range(num_iters):
55
+ _ = method(k_nope, k_pe)
56
+ torch.cuda.synchronize()
57
+ end = time.perf_counter()
58
+
59
+ return (end - start) / num_iters * 1000 # Convert to ms
60
+
61
+
62
+@torch.inference_mode()
63
+def run_benchmark(dtype: torch.dtype, dtype_name: str):
64
+ """Run benchmark for a specific dtype."""
65
+ torch.set_default_device("cuda")
66
+
67
+ # Batch sizes to test (powers of 2 from 32 to 65536)
68
+ batch_sizes = [32, 64, 128, 256, 512, 1024, 2048, 4096, 8192, 16384, 32768, 65536]
69
+
70
+ print("=" * 80)
71
+ print("Benchmark: torch.cat vs direct copy for MLA k_nope/k_pe concatenation")
72
+ print("=" * 80)
73
+ print(
74
+ f"Tensor shapes: k_nope=[B, {NUM_HEADS}, {QK_NOPE_HEAD_DIM}], "
75
+ f"k_pe=[B, 1, {PE_DIM}]"
76
+ )
77
+ print(f"dtype: {dtype_name}")
78
+ print()
79
+ print(
80
+ f"{'Batch Size':>12} | {'cat (ms)':>10} | {'direct (ms)':>12} | "
81
+ f"{'Speedup':>8} | {'Reduction':>10}"
82
+ )
83
+ print("-" * 70)
84
+
85
+ results = []
86
+ for batch_size in batch_sizes:
87
+ # Create input tensors (generate in float32 then convert for FP8 compatibility)
88
+ k_nope = torch.randn(
89
+ batch_size, NUM_HEADS, QK_NOPE_HEAD_DIM, dtype=torch.float32, device="cuda"
90
+ ).to(dtype)
91
+ k_pe = torch.randn(
92
+ batch_size, 1, PE_DIM, dtype=torch.float32, device="cuda"
93
+ ).to(dtype)
94
+
95
+ # Benchmark both methods
96
+ cat_time = benchmark_method(cat_method, k_nope, k_pe)
97
+ direct_time = benchmark_method(direct_copy_method, k_nope, k_pe)
98
+
99
+ speedup = cat_time / direct_time
100
+ reduction = (1 - direct_time / cat_time) * 100
101
+
102
+ results.append((batch_size, cat_time, direct_time, speedup, reduction))
103
+
104
+ print(
105
+ f"{batch_size:>12} | {cat_time:>10.3f} | {direct_time:>12.3f} | "
106
+ f"{speedup:>7.2f}x | {reduction:>9.1f}%"
107
+ )
108
+
109
+ print("=" * 80)
110
+
111
+ # Summary statistics
112
+ speedups = [r[3] for r in results]
113
+ print("\nSpeedup summary:")
114
+ print(f" Min: {min(speedups):.2f}x")
115
+ print(f" Max: {max(speedups):.2f}x")
116
+ print(f" Mean: {sum(speedups) / len(speedups):.2f}x")
117
+
118
+ # Find crossover point
119
+ crossover_batch = None
120
+ for batch_size, _, _, speedup, _ in results:
121
+ if speedup >= 1.0:
122
+ crossover_batch = batch_size
123
+ break
124
+
125
+ print("\nConclusion:")
126
+ if crossover_batch:
127
+ print(f" - Direct copy becomes beneficial at batch size >= {crossover_batch}")
128
+ # Filter for large batches (>= 512 which is typical for prefill)
129
+ large_batch_speedups = [r[3] for r in results if r[0] >= 512]
130
+ if large_batch_speedups:
131
+ avg_large = sum(large_batch_speedups) / len(large_batch_speedups)
132
+ print(f" - For batch sizes >= 512: avg speedup = {avg_large:.2f}x")
133
+ print(" - MLA prefill typically uses large batches, so optimization is effective")
134
+
135
+ return results
136
+
137
+
138
+@torch.inference_mode()
139
+def main():
140
+ # Test bfloat16
141
+ print("\n")
142
+ run_benchmark(torch.bfloat16, "bfloat16")
143
+
144
+ # Test float8_e4m3fn
145
+ print("\n")
146
+ run_benchmark(torch.float8_e4m3fn, "float8_e4m3fn")
147
+
148
+
149
+if __name__ == "__main__":
150
+ main()
vllm/v1/attention/backends/mla/common.py
+30
-3
@@ -1654,6 +1654,33 @@ class MLACommonImpl(MLACommonBaseImpl[M], Generic[M]):
1654
# Convert from (L, N, P) to (N, P, L)
1655
self.W_UK_T = W_UK.permute(1, 2, 0)
1656
1657
+ def _concat_k_nope_k_pe(
1658
+ self, k_nope: torch.Tensor, k_pe: torch.Tensor
1659
+ ) -> torch.Tensor:
1660
+ """
1661
+ Efficiently concatenate k_nope and k_pe tensors along the last dimension.
1662
+
1663
+ This function avoids the performance penalty of torch.cat with expanded
1664
+ non-contiguous tensors by pre-allocating the output and using direct copies.
1665
+
1666
+ Args:
1667
+ k_nope: Tensor of shape [..., nope_dim]
1668
+ k_pe: Tensor to broadcast and concatenate, typically shape [..., 1, pe_dim]
1669
+ or [..., pe_dim]
1670
+
1671
+ Returns:
1672
+ Tensor of shape [..., nope_dim + pe_dim]
1673
+ """
1674
+ k = torch.empty(
1675
+ (*k_nope.shape[:-1], k_nope.shape[-1] + k_pe.shape[-1]),
1676
+ dtype=k_nope.dtype,
1677
+ device=k_nope.device,
1678
+ )
1679
+ # Direct copies with efficient broadcasting
1680
+ k[..., : k_nope.shape[-1]] = k_nope
1681
+ k[..., k_nope.shape[-1] :] = k_pe
1682
+ return k
1683
+
1684
def _compute_prefill_context(
1685
self,
1686
q: torch.Tensor,
@@ -1690,7 +1717,7 @@ class MLACommonImpl(MLACommonBaseImpl[M], Generic[M]):
1717
)
1718
k_nope, v = kv_nope.split([self.qk_nope_head_dim, self.v_head_dim], dim=-1)
1719
1693
- k = torch.cat((k_nope, k_pe.expand((*k_nope.shape[:-1], -1))), dim=-1)
1720
+ k = self._concat_k_nope_k_pe(k_nope, k_pe)
1721
1722
attn_output, attn_softmax_lse = self._run_prefill_context_chunk(
1723
prefill=prefill_metadata,
@@ -1794,7 +1821,7 @@ class MLACommonImpl(MLACommonBaseImpl[M], Generic[M]):
1821
-1, self.num_heads, self.qk_nope_head_dim + self.v_head_dim
1822
)
1823
k_nope, v = kv_nope.split([self.qk_nope_head_dim, self.v_head_dim], dim=-1)
1797
- k = torch.cat((k_nope, k_pe.expand((*k_nope.shape[:-1], -1))), dim=-1)
1824
+ k = self._concat_k_nope_k_pe(k_nope, k_pe)
1825
1826
attn_output, attn_softmax_lse = self._run_prefill_context_chunk(
1827
prefill=prefill_metadata,
@@ -1843,7 +1870,7 @@ class MLACommonImpl(MLACommonBaseImpl[M], Generic[M]):
1870
)
1871
k_nope, v = kv_nope.split([self.qk_nope_head_dim, self.v_head_dim], dim=-1)
1872
1846
- k = torch.cat((k_nope, k_pe.expand((*k_nope.shape[:-1], -1))), dim=-1)
1873
+ k = self._concat_k_nope_k_pe(k_nope, k_pe)
1874
1875
output_prefill = self._run_prefill_new_tokens(
1876
prefill=attn_metadata.prefill,