[Perf] Autotune batch invariance triton kernel in blackwell, 33.6% E2E latency reduction (#53649)
Signed-off-by: yewentao256 <zhyanwentao@126.com>
Wentao Ye committed
Aug 25, 2026 at 15:06 UTC
384d07969771a988eef0e2030ef6ac92b0e8b75e
2 files changed
+85
-2
tests/v1/determinism/test_matmul_batch_invariant.py
+6
-2
@@ -111,7 +111,8 @@ def test_matmul_batch_invariance(dtype):
111
112
@skip_unsupported
113
@pytest.mark.parametrize("m", [8, 32, 256, 2048])
114
-def test_matmul_batch_invariance_across_tuned_m_buckets(m):
114
+@pytest.mark.parametrize("transpose_b", [False, True], ids=["contiguous", "transposed"])
115
+def test_matmul_batch_invariance_across_tuned_m_buckets(m, transpose_b):
116
# Tuned M buckets must preserve each row's K-reduction order.
117
capability = (
118
current_platform.get_device_capability() if current_platform.is_cuda() else None
@@ -124,7 +125,10 @@ def test_matmul_batch_invariance_across_tuned_m_buckets(m):
125
n = k = 2048
126
torch.manual_seed(42)
127
a = torch.rand((m, k), dtype=torch.bfloat16, device=device)
127
- b = torch.rand((k, n), dtype=torch.bfloat16, device=device)
128
+ if transpose_b:
129
+ b = torch.rand((n, k), dtype=torch.bfloat16, device=device).t()
130
+ else:
131
+ b = torch.rand((k, n), dtype=torch.bfloat16, device=device)
132
133
single_output = matmul_batch_invariant(a[:1], b)
134
batch_output = matmul_batch_invariant(a, b)
vllm/model_executor/determinism/batch_invariant_configs.py
+79
@@ -177,6 +177,83 @@ _BATCH_INVARIANT_MATMUL_TUNED_CONFIGS: dict[
177
),
178
),
179
},
180
+ "blackwell": {
181
+ (12288, 2048): _MatmulShapeConfig(
182
+ block_k=64,
183
+ m_buckets=(
184
+ (1, _MatmulMConfig(16, 128, 8, 5)),
185
+ (4, _MatmulMConfig(16, 128, 4, 5)),
186
+ (8, _MatmulMConfig(16, 128, 4, 5)),
187
+ (16, _MatmulMConfig(16, 128, 4, 5)),
188
+ (32, _MatmulMConfig(32, 128, 4, 5)),
189
+ (64, _MatmulMConfig(64, 128, 8, 5)),
190
+ (256, _MatmulMConfig(128, 256, 4, 4)),
191
+ (512, _MatmulMConfig(128, 256, 4, 4)),
192
+ (1024, _MatmulMConfig(128, 256, 4, 4)),
193
+ (2048, _MatmulMConfig(128, 256, 4, 4)),
194
+ ),
195
+ ),
196
+ (2048, 6144): _MatmulShapeConfig(
197
+ block_k=128,
198
+ m_buckets=(
199
+ (1, _MatmulMConfig(16, 32, 4, 5)),
200
+ (4, _MatmulMConfig(16, 32, 4, 5)),
201
+ (8, _MatmulMConfig(16, 32, 4, 5)),
202
+ (16, _MatmulMConfig(16, 32, 4, 5)),
203
+ (32, _MatmulMConfig(16, 32, 4, 5)),
204
+ (64, _MatmulMConfig(32, 32, 4, 5)),
205
+ (256, _MatmulMConfig(64, 64, 4, 5)),
206
+ (512, _MatmulMConfig(128, 64, 4, 4)),
207
+ (1024, _MatmulMConfig(128, 128, 8, 3)),
208
+ (2048, _MatmulMConfig(128, 128, 8, 3)),
209
+ ),
210
+ ),
211
+ (4096, 2048): _MatmulShapeConfig(
212
+ block_k=128,
213
+ m_buckets=(
214
+ (1, _MatmulMConfig(16, 64, 4, 5)),
215
+ (4, _MatmulMConfig(16, 64, 4, 5)),
216
+ (8, _MatmulMConfig(16, 32, 4, 5)),
217
+ (16, _MatmulMConfig(16, 32, 4, 5)),
218
+ (32, _MatmulMConfig(32, 32, 4, 5)),
219
+ (64, _MatmulMConfig(32, 64, 4, 4)),
220
+ (256, _MatmulMConfig(128, 64, 4, 4)),
221
+ (512, _MatmulMConfig(128, 128, 8, 3)),
222
+ (1024, _MatmulMConfig(128, 128, 8, 3)),
223
+ (2048, _MatmulMConfig(128, 128, 8, 3)),
224
+ ),
225
+ ),
226
+ (151936, 2048): _MatmulShapeConfig(
227
+ block_k=64,
228
+ m_buckets=(
229
+ (1, _MatmulMConfig(16, 256, 4, 5)),
230
+ (4, _MatmulMConfig(16, 256, 4, 5)),
231
+ (8, _MatmulMConfig(16, 256, 4, 5)),
232
+ (16, _MatmulMConfig(16, 256, 4, 5)),
233
+ (32, _MatmulMConfig(64, 256, 4, 5)),
234
+ (64, _MatmulMConfig(64, 256, 4, 5)),
235
+ (256, _MatmulMConfig(128, 256, 4, 4)),
236
+ (512, _MatmulMConfig(128, 256, 4, 4)),
237
+ (1024, _MatmulMConfig(128, 256, 4, 4)),
238
+ (2048, _MatmulMConfig(128, 256, 4, 4)),
239
+ ),
240
+ ),
241
+ (2048, 2048): _MatmulShapeConfig(
242
+ block_k=128,
243
+ m_buckets=(
244
+ (1, _MatmulMConfig(16, 32, 4, 4)),
245
+ (4, _MatmulMConfig(16, 32, 4, 4)),
246
+ (8, _MatmulMConfig(16, 32, 4, 5)),
247
+ (16, _MatmulMConfig(16, 32, 8, 5)),
248
+ (32, _MatmulMConfig(16, 32, 4, 4)),
249
+ (64, _MatmulMConfig(32, 32, 4, 5)),
250
+ (256, _MatmulMConfig(64, 64, 4, 4)),
251
+ (512, _MatmulMConfig(128, 64, 4, 4)),
252
+ (1024, _MatmulMConfig(128, 128, 4, 3)),
253
+ (2048, _MatmulMConfig(128, 128, 8, 3)),
254
+ ),
255
+ ),
256
+ },
257
}
258
259
_TUNED_MATMUL_CONFIGS_FOR_DEVICE: dict[tuple[int, int], _MatmulShapeConfig] | None = (
@@ -188,6 +265,8 @@ _TUNED_MATMUL_CONFIGS_RESOLVED = False
265
def _get_tuned_matmul_arch_family(capability: DeviceCapability | None) -> str | None:
266
if capability is None:
267
return None
268
+ if capability.major == 10:
269
+ return "blackwell"
270
if capability.major == 9:
271
return "hopper"
272
if capability.major == 8 and capability.minor == 9: