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 12, 13:30 UTC
[Perf] Optimize batch invariant BMM, 18.1% Throughput improvement, 10.7% TTFT improvement (#29345)
Signed-off-by: yewentao256 <zhyanwentao@126.com> Signed-off-by: Wentao Ye <44945378+yewentao256@users.noreply.github.com> Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
Wentao Ye committed
Nov 26, 2025 at 11:38 UTC
0b0aa874e85431a0f08a0d1fad95ae673e034392
3 files changed
+217
-16
tests/v1/determinism/test_batch_invariance.py
-3
@@ -159,7 +159,6 @@ def test_v1_generation_is_deterministic_across_batch_sizes_with_needle(
159
"backend",
160
BACKENDS,
161
)
162
-@pytest.mark.forked
162
def test_logprobs_bitwise_batch_invariance_bs1_vs_bsN(
163
backend, monkeypatch: pytest.MonkeyPatch
164
):
@@ -429,7 +428,6 @@ def test_simple_generation(backend, monkeypatch: pytest.MonkeyPatch):
428
"backend",
429
BACKENDS,
430
)
432
-@pytest.mark.forked
431
def test_logprobs_without_batch_invariance_should_fail(
432
backend, monkeypatch: pytest.MonkeyPatch
433
):
@@ -646,7 +644,6 @@ def test_logprobs_without_batch_invariance_should_fail(
644
645
@skip_unsupported
646
@pytest.mark.parametrize("backend", ["FLASH_ATTN"])
649
-@pytest.mark.forked
647
def test_decode_logprobs_match_prefill_logprobs(
648
backend, monkeypatch: pytest.MonkeyPatch
649
):
tests/v1/determinism/utils.py
+4
-1
@@ -8,6 +8,7 @@ import torch
8
9
from vllm.attention.utils.fa_utils import flash_attn_supports_mla
10
from vllm.platforms import current_platform
11
+from vllm.utils.flashinfer import has_flashinfer
12
13
skip_unsupported = pytest.mark.skipif(
14
not (current_platform.is_cuda() and current_platform.has_device_capability(90)),
@@ -16,9 +17,11 @@ skip_unsupported = pytest.mark.skipif(
17
18
BACKENDS: list[str] = [
19
"FLASH_ATTN",
19
- "FLASHINFER",
20
]
21
22
+if has_flashinfer():
23
+ BACKENDS.append("FLASHINFER")
24
+
25
if flash_attn_supports_mla():
26
BACKENDS.append("FLASH_ATTN_MLA")
27
vllm/model_executor/layers/batch_invariant.py
+213
-12
@@ -215,6 +215,139 @@ def matmul_persistent(
215
return c
216
217
218
+@triton.jit
219
+def bmm_kernel(
220
+ a_ptr, # (*, ) pointer to A, (B, M, K)
221
+ b_ptr, # (*, ) pointer to B, (B, K, N)
222
+ c_ptr, # (*, ) pointer to C, (B, M, N)
223
+ B, # int, batch size
224
+ M, # int, output rows
225
+ N, # int, output cols
226
+ K, # int, reduction dim
227
+ stride_ab,
228
+ stride_am,
229
+ stride_ak,
230
+ stride_bb,
231
+ stride_bk,
232
+ stride_bn,
233
+ stride_cb,
234
+ stride_cm,
235
+ stride_cn,
236
+ BLOCK_SIZE_M: tl.constexpr,
237
+ BLOCK_SIZE_N: tl.constexpr,
238
+ BLOCK_SIZE_K: tl.constexpr,
239
+ A_LARGE: tl.constexpr,
240
+ B_LARGE: tl.constexpr,
241
+ C_LARGE: tl.constexpr,
242
+):
243
+ """Batched GEMM: (B, M, K) x (B, K, N) -> (B, M, N)
244
+
245
+ Each program computes one (batch_idx, tile_m, tile_n) tile, accumulating
246
+ along K in a fixed order to preserve batch invariance.
247
+ """
248
+ pid_b = tl.program_id(0)
249
+ pid = tl.program_id(1)
250
+
251
+ if pid_b >= B:
252
+ return
253
+
254
+ # number of tiles along M / N
255
+ num_pid_m = tl.cdiv(M, BLOCK_SIZE_M)
256
+ num_pid_n = tl.cdiv(N, BLOCK_SIZE_N)
257
+
258
+ pid_m = pid // num_pid_n
259
+ pid_n = pid % num_pid_n
260
+
261
+ if pid_m >= num_pid_m or pid_n >= num_pid_n:
262
+ return
263
+
264
+ # offs_m / offs_n: raw global row/col indices for this tile
265
+ offs_m = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
266
+ offs_n = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
267
+ # masks for valid logical rows/cols within (M, N)
268
+ mask_m = offs_m < M # [BLOCK_SIZE_M]
269
+ mask_n = offs_n < N # [BLOCK_SIZE_N]
270
+
271
+ if A_LARGE or B_LARGE or C_LARGE:
272
+ offs_m = offs_m.to(tl.int64)
273
+ offs_n = offs_n.to(tl.int64)
274
+
275
+ offs_m = tl.where(mask_m, offs_m, 0)
276
+ offs_n = tl.where(mask_n, offs_n, 0)
277
+
278
+ # hint for triton contiguous memory
279
+ offs_m = tl.max_contiguous(tl.multiple_of(offs_m, BLOCK_SIZE_M), BLOCK_SIZE_M)
280
+ offs_n = tl.max_contiguous(tl.multiple_of(offs_n, BLOCK_SIZE_N), BLOCK_SIZE_N)
281
+
282
+ # base pointers for current batch, shape-wise:
283
+ # a_batch_ptr points to A[pid_b, 0, 0]
284
+ # b_batch_ptr points to B[pid_b, 0, 0]
285
+ # c_batch_ptr points to C[pid_b, 0, 0]
286
+ a_batch_ptr = a_ptr + pid_b * stride_ab
287
+ b_batch_ptr = b_ptr + pid_b * stride_bb
288
+ c_batch_ptr = c_ptr + pid_b * stride_cb
289
+
290
+ accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
291
+ # number of K-blocks this tile iterates over
292
+ k_tiles = tl.cdiv(K, BLOCK_SIZE_K)
293
+ offs_k_mask = tl.arange(0, BLOCK_SIZE_K)
294
+
295
+ for ki in range(k_tiles):
296
+ if A_LARGE or B_LARGE:
297
+ # offs_k: [BLOCK_SIZE_K], global K indices
298
+ offs_k = ki * BLOCK_SIZE_K + tl.arange(0, BLOCK_SIZE_K).to(tl.int64)
299
+ else:
300
+ offs_k = ki * BLOCK_SIZE_K + tl.arange(0, BLOCK_SIZE_K)
301
+
302
+ # a_ptrs: [BLOCK_SIZE_M, BLOCK_SIZE_K]
303
+ # element (i, j) points to A[pid_b, offs_m[i], offs_k[j]]
304
+ a_ptrs = a_batch_ptr + (
305
+ offs_m[:, None] * stride_am + offs_k[None, :] * stride_ak
306
+ )
307
+ # b_ptrs: [BLOCK_SIZE_K, BLOCK_SIZE_N]
308
+ # element (i, j) points to B[pid_b, offs_k[i], offs_n[j]]
309
+ b_ptrs = b_batch_ptr + (
310
+ offs_k[:, None] * stride_bk + offs_n[None, :] * stride_bn
311
+ )
312
+
313
+ # valid K lanes for this block
314
+ k_valid = offs_k_mask < (K - ki * BLOCK_SIZE_K)
315
+ # A mask within (M, K): [BLOCK_SIZE_M, BLOCK_SIZE_K]
316
+ a_mask = mask_m[:, None] & k_valid[None, :]
317
+ # B mask within (K, N): [BLOCK_SIZE_K, BLOCK_SIZE_N]
318
+ b_mask = k_valid[:, None] & mask_n[None, :]
319
+
320
+ # a: [BLOCK_SIZE_M, BLOCK_SIZE_K] from A[offs_m, offs_k]
321
+ a = tl.load(
322
+ a_ptrs,
323
+ mask=a_mask,
324
+ other=0.0,
325
+ )
326
+ # b: [BLOCK_SIZE_K, BLOCK_SIZE_N] from B[offs_k, offs_n]
327
+ b = tl.load(
328
+ b_ptrs,
329
+ mask=b_mask,
330
+ other=0.0,
331
+ )
332
+ accumulator = tl.dot(a, b, accumulator)
333
+
334
+ # c_m / c_n: [BLOCK_SIZE_M] / [BLOCK_SIZE_N], row/col indices for C
335
+ c_m = offs_m
336
+ c_n = offs_n
337
+ if C_LARGE:
338
+ c_m = c_m.to(tl.int64)
339
+ c_n = c_n.to(tl.int64)
340
+
341
+ # c_ptrs: [BLOCK_SIZE_M, BLOCK_SIZE_N]
342
+ # element (i, j) points to C[pid_b, c_m[i], c_n[j]]
343
+ c_ptrs = c_batch_ptr + stride_cm * c_m[:, None] + stride_cn * c_n[None, :]
344
+ # mask out elements that fall outside logical (M, N) range
345
+ c_mask = mask_m[:, None] & mask_n[None, :]
346
+ # cast FP32 accumulator back to original dtype of C
347
+ c = accumulator.to(c_ptr.dtype.element_ty)
348
+ tl.store(c_ptrs, c, mask=c_mask)
349
+
350
+
351
@triton.jit
352
def _log_softmax_kernel(
353
input_ptr,
@@ -526,23 +659,91 @@ def matmul_batch_invariant(a, b, *, out=None):
659
660
def bmm_batch_invariant(a, b, *, out=None):
661
# Batched matrix multiply: (B, M, K) x (B, K, N) -> (B, M, N)
529
- # Process each batch separately with our persistent kernel
530
- if a.ndim == 3 and b.ndim == 3:
531
- results = []
532
- for i in range(a.shape[0]):
533
- results.append(matmul_persistent(a[i], b[i]))
534
- result = torch.stack(results, dim=0)
535
-
536
- if out is not None:
537
- out.copy_(result)
538
- return out
539
- return result
540
- else:
662
+ if not (a.ndim == 3 and b.ndim == 3):
663
raise ValueError(
664
f"bmm_batch_invariant expects 3D tensors, "
665
f"got shapes {a.shape} and {b.shape}"
666
)
667
668
+ if a.shape[0] != b.shape[0]:
669
+ raise ValueError(
670
+ f"Batch dimensions of tensors must match, "
671
+ f"but got {a.shape[0]} and {b.shape[0]}."
672
+ )
673
+ if a.shape[2] != b.shape[1]:
674
+ raise ValueError(
675
+ f"Incompatible inner dimensions for matmul: got {a.shape} and {b.shape}."
676
+ )
677
+ if a.dtype != b.dtype:
678
+ raise ValueError(f"Incompatible dtypes: got {a.dtype} and {b.dtype}.")
679
+
680
+ B, M, K = a.shape
681
+ _, _, N = b.shape
682
+ dtype = a.dtype
683
+
684
+ if out is None:
685
+ c = torch.empty((B, M, N), device=a.device, dtype=dtype)
686
+ else:
687
+ assert out.shape == (B, M, N), "out tensor has incorrect shape"
688
+ assert out.dtype == dtype and out.device == a.device, "out tensor mismatch"
689
+ c = out
690
+
691
+ configs = {
692
+ torch.bfloat16: {
693
+ "BLOCK_SIZE_M": 128,
694
+ "BLOCK_SIZE_N": 128,
695
+ "BLOCK_SIZE_K": 64,
696
+ "num_stages": 3,
697
+ "num_warps": 8,
698
+ },
699
+ torch.float16: {
700
+ "BLOCK_SIZE_M": 128,
701
+ "BLOCK_SIZE_N": 256,
702
+ "BLOCK_SIZE_K": 64,
703
+ "num_stages": 3,
704
+ "num_warps": 8,
705
+ },
706
+ torch.float32: {
707
+ "BLOCK_SIZE_M": 128,
708
+ "BLOCK_SIZE_N": 128,
709
+ "BLOCK_SIZE_K": 32,
710
+ "num_stages": 3,
711
+ "num_warps": 8,
712
+ },
713
+ }
714
+
715
+ cfg = configs[dtype]
716
+ # grid = (B, num_tiles_per_matrix)
717
+ grid = (
718
+ B,
719
+ triton.cdiv(M, cfg["BLOCK_SIZE_M"]) * triton.cdiv(N, cfg["BLOCK_SIZE_N"]),
720
+ )
721
+
722
+ bmm_kernel[grid](
723
+ a,
724
+ b,
725
+ c,
726
+ B,
727
+ M,
728
+ N,
729
+ K,
730
+ a.stride(0),
731
+ a.stride(1),
732
+ a.stride(2),
733
+ b.stride(0),
734
+ b.stride(1),
735
+ b.stride(2),
736
+ c.stride(0),
737
+ c.stride(1),
738
+ c.stride(2),
739
+ A_LARGE=a.numel() > 2**31,
740
+ B_LARGE=b.numel() > 2**31,
741
+ C_LARGE=c.numel() > 2**31,
742
+ **cfg,
743
+ )
744
+
745
+ return c
746
+
747
748
def addmm_batch_invariant(bias, a, b):
749
return matmul_persistent(a, b, bias=bias)