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 4, 23:45 UTC
Added general ND x ND matmul and unit test for it (#39909)
Signed-off-by: Yusuf <yusufmohammad@live.com>
Yusuf Mohammad committed
Apr 18, 2026 at 15:05 UTC
b5f6c5f8343dd0ff5accb3f8a8f7d1f5c40f32aa
2 files changed
+129
-31
tests/v1/determinism/test_matmul_batch_invariant.py
new
+105
@@ -0,0 +1,105 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""
4
+Test batch-invariant matmul against torch.matmul for various shape combinations.
5
+
6
+Tests correctness (matches torch.matmul) and batch invariance (result for one
7
+item doesn't change based on other items in the batch).
8
+"""
9
+
10
+import pytest
11
+import torch
12
+from utils import skip_unsupported
13
+
14
+from vllm.model_executor.layers.batch_invariant import matmul_batch_invariant
15
+from vllm.platforms import current_platform
16
+
17
+DEVICE_TYPE = current_platform.device_type
18
+
19
+
20
+@skip_unsupported
21
+@pytest.mark.parametrize(
22
+ "a_shape,b_shape",
23
+ [
24
+ # 2D x 2D
25
+ ((32, 64), (64, 16)),
26
+ # 2D x 3D
27
+ ((64, 16), (4, 16, 32)),
28
+ # 3D x 2D
29
+ ((4, 32, 64), (64, 16)),
30
+ # 4D x 2D
31
+ ((1, 4, 32, 64), (64, 16)),
32
+ # 3D x 3D
33
+ ((4, 32, 64), (4, 64, 16)),
34
+ # 3D x 4D
35
+ ((2, 32, 64), (1, 2, 64, 16)),
36
+ # 4D x 3D (Gemma4 pattern)
37
+ ((1, 2, 32, 64), (2, 64, 16)),
38
+ # 4D x 4D
39
+ ((1, 2, 32, 64), (4, 2, 64, 16)),
40
+ # 2D x 4D
41
+ ((32, 64), (1, 2, 64, 16)),
42
+ # 2D x 5D
43
+ ((32, 64), (1, 2, 2, 64, 16)),
44
+ # 5D x 2D
45
+ ((1, 2, 2, 32, 64), (64, 16)),
46
+ # 5D x 5D
47
+ ((1, 2, 4, 32, 64), (1, 2, 4, 64, 16)),
48
+ ],
49
+)
50
+@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
51
+def test_matmul_correctness(a_shape, b_shape, dtype):
52
+ """
53
+ Compare matmul_batch_invariant against torch.matmul for various shapes.
54
+ """
55
+ device = torch.device(DEVICE_TYPE)
56
+
57
+ torch.manual_seed(42)
58
+ a = torch.rand(a_shape, dtype=dtype, device=device)
59
+ b = torch.rand(b_shape, dtype=dtype, device=device)
60
+
61
+ # Standard implementation (CUDA ops)
62
+ standard_output = torch.matmul(a, b)
63
+
64
+ # Batch-invariant implementation (Triton)
65
+ triton_output = matmul_batch_invariant(a, b)
66
+
67
+ # Compare outputs
68
+ # Use looser tolerance for bfloat16 due to its lower precision
69
+ if dtype == torch.bfloat16:
70
+ rtol, atol = 1e-1, 1e-1 # 10% relative tolerance for bfloat16
71
+ else:
72
+ rtol, atol = 1e-2, 1e-2 # 1% for float16/float32
73
+
74
+ torch.testing.assert_close(
75
+ triton_output,
76
+ standard_output,
77
+ rtol=rtol,
78
+ atol=atol,
79
+ msg=f"matmul mismatch for a ndim={a.ndim}, b ndim={b.ndim},",
80
+ )
81
+
82
+
83
+@skip_unsupported
84
+@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
85
+def test_matmul_batch_invariance(dtype):
86
+ """
87
+ Verify that the result for one item is bitwise identical regardless
88
+ of what other items are in the batch.
89
+ """
90
+
91
+ device = torch.device(DEVICE_TYPE)
92
+
93
+ torch.manual_seed(42)
94
+ a_single = torch.rand((1, 64, 32), dtype=dtype, device=device)
95
+ b = torch.rand((32, 128), dtype=dtype, device=device)
96
+
97
+ standard_output = matmul_batch_invariant(a_single, b)
98
+
99
+ a_batch = torch.rand((8, 64, 32), dtype=dtype, device=device)
100
+ a_batch[3] = a_single[0]
101
+
102
+ batch_output = matmul_batch_invariant(a_batch, b)
103
+ batch_output_a = batch_output[3]
104
+
105
+ assert torch.equal(standard_output[0], batch_output_a)
vllm/model_executor/layers/batch_invariant.py
+24
-31
@@ -1,5 +1,6 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import math
4
import os
5
from collections.abc import Callable
6
from typing import Any
@@ -611,51 +612,43 @@ def matmul_batch_invariant(a, b, *, out=None):
612
out.copy_(result)
613
return out
614
return result
614
- elif a.ndim == 3 and b.ndim == 3:
615
- # Handle batched case like bmm
616
- return bmm_batch_invariant(a, b, out=out)
617
- elif a.ndim == 3 and b.ndim == 2:
618
- # Handle 3D x 2D: common for linear layers
619
- # (batch, seq, hidden) @ (hidden, out) -> (batch, seq, out)
620
- # Reshape to 2D, do mm, reshape back
621
- batch, seq, hidden = a.shape
615
+ elif b.ndim == 2:
616
+ # Handle ND x 2D: Common for linear layers
617
+ # (..., batch, seq, hidden) @ (hidden, out) -> (..., batch, seq, out)
618
+ batch_dims = a.shape[:-1]
619
+ hidden = a.shape[-1]
620
+ out_dim = b.shape[-1]
621
a_2d = a.reshape(-1, hidden)
622
result_2d = matmul_persistent(a_2d, b)
624
- result = result_2d.reshape(batch, seq, -1)
623
+ result = result_2d.reshape(batch_dims + (out_dim,))
624
if out is not None:
625
out.copy_(result)
626
return out
627
return result
629
- elif a.ndim == 2 and b.ndim == 3:
630
- # Handle 2D x 3D: (M, K) @ (B, K, N) -> (B, M, N)
631
- # By broadcasting `a` to 3D, we can reuse the batched matrix
632
- # multiplication logic.
633
- a_expanded = a.unsqueeze(0).expand(b.shape[0], -1, -1)
634
- return bmm_batch_invariant(a_expanded, b, out=out)
635
- elif a.ndim == 4 and b.ndim == 4:
636
- # Handle 4D attention tensors: [batch, heads, seq, dim]
637
- # Reshape to 3D, process, reshape back
638
- batch, heads, seq_a, dim_a = a.shape
639
- _, _, dim_b, seq_b = b.shape
640
-
641
- # Reshape to [batch*heads, seq_a, dim_a]
642
- a_3d = a.reshape(batch * heads, seq_a, dim_a)
643
- b_3d = b.reshape(batch * heads, dim_b, seq_b)
644
-
628
+ elif a.ndim >= 2 and b.ndim >= 3:
629
+ # Generic handler for 2D x ND and ND x ND (except 1D)
630
+ # Broadcast dims to ensure both matrices have the same shape
631
+ # If 2D x ND, then unsqueeze to add a dim to a
632
+ if a.ndim == 2:
633
+ a = a.unsqueeze(0)
634
+ broadcast_shape = torch.broadcast_shapes(a.shape[:-2], b.shape[:-2])
635
+ a = a.expand(broadcast_shape + a.shape[-2:])
636
+ b = b.expand(broadcast_shape + b.shape[-2:])
637
+ batch_dim = math.prod(broadcast_shape)
638
+ # Reuse broadcast shape to get all dims except mm dims
639
+ a_3d = a.reshape(batch_dim, a.shape[-2], a.shape[-1])
640
+ b_3d = b.reshape(batch_dim, b.shape[-2], b.shape[-1])
641
# Do batched matmul
642
result_3d = bmm_batch_invariant(a_3d, b_3d)
647
-
648
- # Reshape back to [batch, heads, seq_a, seq_b]
649
- result = result_3d.reshape(batch, heads, seq_a, seq_b)
650
-
643
+ # Reshape back to [broadcast_shape, seq_a, seq_b]
644
+ result = result_3d.reshape(broadcast_shape + (a.shape[-2], b.shape[-1]))
645
if out is not None:
646
out.copy_(result)
647
return out
648
return result
649
else:
650
raise ValueError(
657
- f"matmul_batch_invariant currently only supports 2D x 2D, 3D x 3D, "
658
- f"3D x 2D, 2D x 3D, and 4D x 4D, "
651
+ f"matmul_batch_invariant requires both inputs be at least 2D "
652
f"got shapes {a.shape} and {b.shape}"
653
)
654