[Kernel][Perf] Per-architecture tuned configs for batch-invariant persistent matmul (~3x decode kernels on RTX 4090D/H20) (#53247)
Signed-off-by: LioEinaudi <zhao3024667639@gmail.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Lio Einaudi committed
Aug 24, 2026 at 23:28 UTC
7797b6022c129b862e45ae6aed08822e65d1bccb
3 files changed
+283
-1
tests/v1/determinism/test_matmul_batch_invariant.py
+27
@@ -12,6 +12,10 @@ import torch
12
from utils import skip_unsupported
13
14
from vllm.model_executor.layers.batch_invariant import matmul_batch_invariant
15
+from vllm.model_executor.layers.batch_invariant_configs import (
16
+ _BATCH_INVARIANT_MATMUL_TUNED_CONFIGS,
17
+ _get_tuned_matmul_arch_family,
18
+)
19
from vllm.platforms import current_platform
20
21
DEVICE_TYPE = current_platform.device_type
@@ -103,3 +107,26 @@ def test_matmul_batch_invariance(dtype):
107
batch_output_a = batch_output[3]
108
109
assert torch.equal(standard_output[0], batch_output_a)
110
+
111
+
112
+@skip_unsupported
113
+@pytest.mark.parametrize("m", [8, 32, 256, 2048])
114
+def test_matmul_batch_invariance_across_tuned_m_buckets(m):
115
+ # Tuned M buckets must preserve each row's K-reduction order.
116
+ capability = (
117
+ current_platform.get_device_capability() if current_platform.is_cuda() else None
118
+ )
119
+ arch_family = _get_tuned_matmul_arch_family(capability)
120
+ if arch_family not in _BATCH_INVARIANT_MATMUL_TUNED_CONFIGS:
121
+ pytest.skip("No tuned persistent matmul config for this architecture")
122
+
123
+ device = torch.device(DEVICE_TYPE)
124
+ n = k = 2048
125
+ torch.manual_seed(42)
126
+ a = torch.rand((m, k), dtype=torch.bfloat16, device=device)
127
+ b = torch.rand((k, n), dtype=torch.bfloat16, device=device)
128
+
129
+ single_output = matmul_batch_invariant(a[:1], b)
130
+ batch_output = matmul_batch_invariant(a, b)
131
+
132
+ assert torch.equal(single_output[0], batch_output[0])
vllm/model_executor/layers/batch_invariant.py
+6
-1
@@ -8,6 +8,10 @@ from typing import Any
8
import torch
9
10
import vllm.envs as envs
11
+from vllm.model_executor.layers.batch_invariant_configs import (
12
+ _get_matmul_config,
13
+ resolve_tuned_matmul_configs,
14
+)
15
from vllm.platforms import current_platform
16
from vllm.triton_utils import tl, triton
17
from vllm.utils.mem_utils import get_max_shared_memory_bytes
@@ -200,7 +204,7 @@ def matmul_persistent(
204
B_LARGE=b.numel() > 2**31,
205
C_LARGE=c.numel() > 2**31,
206
HAS_BIAS=bias is not None,
203
- **configs[dtype],
207
+ **_get_matmul_config(M, N, K, dtype, configs[dtype]),
208
)
209
return c
210
@@ -997,6 +1001,7 @@ def override_envs_for_invariance():
1001
def init_batch_invariance():
1002
# this will hit all the csrc overrides as well
1003
if envs.VLLM_BATCH_INVARIANT:
1004
+ resolve_tuned_matmul_configs()
1005
override_envs_for_invariance()
1006
enable_batch_invariant_mode()
1007
vllm/model_executor/layers/batch_invariant_configs.py
new
+250
@@ -0,0 +1,250 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+from typing import NamedTuple
4
+
5
+import torch
6
+
7
+from vllm.platforms import current_platform
8
+from vllm.platforms.interface import DeviceCapability
9
+
10
+
11
+class _MatmulMConfig(NamedTuple):
12
+ block_m: int
13
+ block_n: int
14
+ num_warps: int
15
+ num_stages: int
16
+
17
+
18
+class _MatmulShapeConfig(NamedTuple):
19
+ block_k: int
20
+ m_buckets: tuple[tuple[int, _MatmulMConfig], ...]
21
+
22
+
23
+_BATCH_INVARIANT_MATMUL_TUNED_CONFIGS: dict[
24
+ str, dict[tuple[int, int], _MatmulShapeConfig]
25
+] = {
26
+ "ada": {
27
+ (12288, 2048): _MatmulShapeConfig(
28
+ block_k=64,
29
+ m_buckets=(
30
+ (1, _MatmulMConfig(16, 128, 8, 2)),
31
+ (4, _MatmulMConfig(16, 64, 8, 5)),
32
+ (8, _MatmulMConfig(16, 128, 8, 4)),
33
+ (16, _MatmulMConfig(16, 128, 4, 5)),
34
+ (32, _MatmulMConfig(32, 128, 4, 5)),
35
+ (64, _MatmulMConfig(32, 256, 4, 3)),
36
+ (256, _MatmulMConfig(128, 128, 8, 3)),
37
+ (512, _MatmulMConfig(128, 64, 4, 4)),
38
+ (1024, _MatmulMConfig(128, 128, 8, 3)),
39
+ (2048, _MatmulMConfig(128, 128, 4, 3)),
40
+ ),
41
+ ),
42
+ (2048, 6144): _MatmulShapeConfig(
43
+ block_k=128,
44
+ m_buckets=(
45
+ (1, _MatmulMConfig(16, 64, 4, 4)),
46
+ (4, _MatmulMConfig(16, 64, 8, 4)),
47
+ (8, _MatmulMConfig(16, 32, 8, 4)),
48
+ (16, _MatmulMConfig(16, 32, 8, 4)),
49
+ (32, _MatmulMConfig(16, 64, 8, 3)),
50
+ (64, _MatmulMConfig(32, 64, 4, 4)),
51
+ (256, _MatmulMConfig(32, 64, 4, 4)),
52
+ (512, _MatmulMConfig(64, 64, 4, 3)),
53
+ (1024, _MatmulMConfig(64, 64, 4, 3)),
54
+ (2048, _MatmulMConfig(64, 64, 4, 3)),
55
+ ),
56
+ ),
57
+ (4096, 2048): _MatmulShapeConfig(
58
+ block_k=64,
59
+ m_buckets=(
60
+ (1, _MatmulMConfig(16, 256, 8, 3)),
61
+ (4, _MatmulMConfig(16, 64, 4, 3)),
62
+ (8, _MatmulMConfig(16, 128, 8, 2)),
63
+ (16, _MatmulMConfig(16, 32, 4, 5)),
64
+ (32, _MatmulMConfig(16, 64, 4, 4)),
65
+ (64, _MatmulMConfig(64, 64, 8, 3)),
66
+ (256, _MatmulMConfig(64, 64, 4, 5)),
67
+ (512, _MatmulMConfig(64, 64, 4, 5)),
68
+ (1024, _MatmulMConfig(128, 64, 4, 4)),
69
+ (2048, _MatmulMConfig(128, 64, 8, 4)),
70
+ ),
71
+ ),
72
+ (151936, 2048): _MatmulShapeConfig(
73
+ block_k=64,
74
+ m_buckets=(
75
+ (1, _MatmulMConfig(16, 32, 8, 5)),
76
+ (4, _MatmulMConfig(16, 32, 8, 5)),
77
+ (8, _MatmulMConfig(16, 32, 8, 5)),
78
+ (16, _MatmulMConfig(16, 32, 4, 4)),
79
+ (32, _MatmulMConfig(16, 64, 4, 4)),
80
+ (64, _MatmulMConfig(32, 64, 4, 4)),
81
+ (256, _MatmulMConfig(128, 128, 4, 3)),
82
+ (512, _MatmulMConfig(128, 128, 4, 3)),
83
+ (1024, _MatmulMConfig(128, 128, 4, 3)),
84
+ (2048, _MatmulMConfig(128, 128, 4, 3)),
85
+ ),
86
+ ),
87
+ (2048, 2048): _MatmulShapeConfig(
88
+ block_k=64,
89
+ m_buckets=(
90
+ (1, _MatmulMConfig(32, 128, 8, 5)),
91
+ (4, _MatmulMConfig(16, 256, 8, 3)),
92
+ (8, _MatmulMConfig(32, 128, 8, 3)),
93
+ (16, _MatmulMConfig(32, 128, 8, 3)),
94
+ (32, _MatmulMConfig(16, 128, 8, 5)),
95
+ (64, _MatmulMConfig(16, 256, 4, 3)),
96
+ (256, _MatmulMConfig(64, 128, 8, 4)),
97
+ (512, _MatmulMConfig(64, 64, 4, 5)),
98
+ (1024, _MatmulMConfig(64, 64, 4, 5)),
99
+ (2048, _MatmulMConfig(64, 64, 4, 5)),
100
+ ),
101
+ ),
102
+ },
103
+ "hopper": {
104
+ (12288, 2048): _MatmulShapeConfig(
105
+ block_k=128,
106
+ m_buckets=(
107
+ (1, _MatmulMConfig(16, 256, 4, 3)),
108
+ (4, _MatmulMConfig(16, 256, 4, 4)),
109
+ (8, _MatmulMConfig(16, 64, 4, 4)),
110
+ (16, _MatmulMConfig(16, 64, 4, 4)),
111
+ (32, _MatmulMConfig(32, 64, 8, 5)),
112
+ (64, _MatmulMConfig(64, 64, 4, 4)),
113
+ (256, _MatmulMConfig(64, 128, 4, 3)),
114
+ (512, _MatmulMConfig(128, 128, 8, 3)),
115
+ (1024, _MatmulMConfig(128, 256, 8, 2)),
116
+ (2048, _MatmulMConfig(128, 256, 8, 2)),
117
+ ),
118
+ ),
119
+ (2048, 6144): _MatmulShapeConfig(
120
+ block_k=128,
121
+ m_buckets=(
122
+ (1, _MatmulMConfig(16, 32, 4, 4)),
123
+ (4, _MatmulMConfig(16, 64, 4, 4)),
124
+ (8, _MatmulMConfig(16, 64, 4, 4)),
125
+ (16, _MatmulMConfig(16, 32, 4, 3)),
126
+ (32, _MatmulMConfig(32, 32, 8, 3)),
127
+ (64, _MatmulMConfig(64, 32, 4, 5)),
128
+ (256, _MatmulMConfig(64, 128, 4, 3)),
129
+ (512, _MatmulMConfig(128, 128, 4, 3)),
130
+ (1024, _MatmulMConfig(64, 64, 4, 5)),
131
+ (2048, _MatmulMConfig(64, 128, 4, 4)),
132
+ ),
133
+ ),
134
+ (4096, 2048): _MatmulShapeConfig(
135
+ block_k=128,
136
+ m_buckets=(
137
+ (1, _MatmulMConfig(16, 64, 8, 5)),
138
+ (4, _MatmulMConfig(16, 64, 4, 5)),
139
+ (8, _MatmulMConfig(16, 128, 4, 3)),
140
+ (16, _MatmulMConfig(16, 64, 4, 5)),
141
+ (32, _MatmulMConfig(16, 64, 4, 4)),
142
+ (64, _MatmulMConfig(64, 32, 4, 3)),
143
+ (256, _MatmulMConfig(64, 256, 4, 2)),
144
+ (512, _MatmulMConfig(64, 64, 4, 3)),
145
+ (1024, _MatmulMConfig(64, 128, 4, 3)),
146
+ (2048, _MatmulMConfig(64, 64, 4, 4)),
147
+ ),
148
+ ),
149
+ (151936, 2048): _MatmulShapeConfig(
150
+ block_k=128,
151
+ m_buckets=(
152
+ (1, _MatmulMConfig(16, 256, 4, 4)),
153
+ (4, _MatmulMConfig(16, 256, 4, 4)),
154
+ (8, _MatmulMConfig(16, 256, 4, 3)),
155
+ (16, _MatmulMConfig(16, 256, 4, 3)),
156
+ (32, _MatmulMConfig(32, 256, 8, 3)),
157
+ (64, _MatmulMConfig(64, 128, 4, 3)),
158
+ (256, _MatmulMConfig(64, 128, 4, 3)),
159
+ (512, _MatmulMConfig(128, 128, 8, 3)),
160
+ (1024, _MatmulMConfig(128, 256, 8, 2)),
161
+ (2048, _MatmulMConfig(128, 256, 8, 2)),
162
+ ),
163
+ ),
164
+ (2048, 2048): _MatmulShapeConfig(
165
+ block_k=128,
166
+ m_buckets=(
167
+ (1, _MatmulMConfig(16, 64, 4, 5)),
168
+ (4, _MatmulMConfig(16, 128, 8, 3)),
169
+ (8, _MatmulMConfig(16, 32, 4, 3)),
170
+ (16, _MatmulMConfig(32, 64, 8, 5)),
171
+ (32, _MatmulMConfig(16, 128, 4, 5)),
172
+ (64, _MatmulMConfig(32, 64, 4, 2)),
173
+ (256, _MatmulMConfig(128, 64, 4, 4)),
174
+ (512, _MatmulMConfig(64, 256, 4, 2)),
175
+ (1024, _MatmulMConfig(64, 64, 4, 3)),
176
+ (2048, _MatmulMConfig(64, 128, 4, 3)),
177
+ ),
178
+ ),
179
+ },
180
+}
181
+
182
+_TUNED_MATMUL_CONFIGS_FOR_DEVICE: dict[tuple[int, int], _MatmulShapeConfig] | None = (
183
+ None
184
+)
185
+_TUNED_MATMUL_CONFIGS_RESOLVED = False
186
+
187
+
188
+def _get_tuned_matmul_arch_family(capability: DeviceCapability | None) -> str | None:
189
+ if capability is None:
190
+ return None
191
+ if capability.major == 9:
192
+ return "hopper"
193
+ if capability.major == 8 and capability.minor == 9:
194
+ return "ada"
195
+ return None
196
+
197
+
198
+def resolve_tuned_matmul_configs() -> None:
199
+ global _TUNED_MATMUL_CONFIGS_FOR_DEVICE
200
+ global _TUNED_MATMUL_CONFIGS_RESOLVED
201
+
202
+ if _TUNED_MATMUL_CONFIGS_RESOLVED:
203
+ return
204
+
205
+ # Normal vLLM initialization resolves this before torch.compile tracing.
206
+ capability = (
207
+ current_platform.get_device_capability() if current_platform.is_cuda() else None
208
+ )
209
+ arch_family = _get_tuned_matmul_arch_family(capability)
210
+ if arch_family is None:
211
+ _TUNED_MATMUL_CONFIGS_FOR_DEVICE = None
212
+ else:
213
+ _TUNED_MATMUL_CONFIGS_FOR_DEVICE = _BATCH_INVARIANT_MATMUL_TUNED_CONFIGS.get(
214
+ arch_family
215
+ )
216
+ _TUNED_MATMUL_CONFIGS_RESOLVED = True
217
+
218
+
219
+def _get_matmul_config(
220
+ M: int,
221
+ N: int,
222
+ K: int,
223
+ dtype: torch.dtype,
224
+ default: dict[str, int],
225
+) -> dict[str, int]:
226
+ if not _TUNED_MATMUL_CONFIGS_RESOLVED:
227
+ resolve_tuned_matmul_configs()
228
+ if dtype != torch.bfloat16:
229
+ return default
230
+ device_configs = _TUNED_MATMUL_CONFIGS_FOR_DEVICE
231
+ if device_configs is None:
232
+ return default
233
+ shape_config = device_configs.get((N, K))
234
+ if shape_config is None:
235
+ return default
236
+ # Values above the tuned range reuse the largest bucket;
237
+ # shape-wide BLOCK_K keeps this batch-invariant.
238
+ m_config = shape_config.m_buckets[-1][1]
239
+ for max_m, bucket_config in shape_config.m_buckets:
240
+ if max_m >= M:
241
+ m_config = bucket_config
242
+ break
243
+ return {
244
+ "BLOCK_SIZE_M": m_config.block_m,
245
+ "BLOCK_SIZE_N": m_config.block_n,
246
+ "BLOCK_SIZE_K": shape_config.block_k,
247
+ "GROUP_SIZE_M": 8,
248
+ "num_warps": m_config.num_warps,
249
+ "num_stages": m_config.num_stages,
250
+ }