Skip to content

Commit d812771

Browse files
wuhulalalawuhulalala
committed
fix mix MatMul kernels fail with fp64 inputs due to accumulator type mismatch
1 parent a8a4452 commit d812771

9 files changed

Lines changed: 88 additions & 19 deletions

File tree

‎src/flag_gems/ops/addmm.py‎

Lines changed: 10 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -47,6 +47,7 @@ def addmm_kernel(
4747
BLOCK_SIZE_M: tl.constexpr,
4848
BLOCK_SIZE_N: tl.constexpr,
4949
BLOCK_SIZE_K: tl.constexpr,
50+
IS_FP64: tl.constexpr = False,
5051
):
5152
pid_m = tle.program_id(0)
5253
pid_n = tle.program_id(1)
@@ -57,7 +58,10 @@ def addmm_kernel(
5758
a_ptrs = a_ptr + (offs_am[:, None] * stride_am + offs_k[None, :] * stride_ak)
5859
b_ptrs = b_ptr + (offs_k[:, None] * stride_bk + offs_bn[None, :] * stride_bn)
5960

60-
accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
61+
if IS_FP64:
62+
accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float64)
63+
else:
64+
accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
6165
for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)):
6266
a = tl.load(
6367
a_ptrs,
@@ -69,6 +73,9 @@ def addmm_kernel(
6973
mask=(offs_k[:, None] < K - k * BLOCK_SIZE_K) & (offs_bn[None, :] < N),
7074
other=0.0,
7175
)
76+
if IS_FP64:
77+
a = a.to(tl.float32)
78+
b = b.to(tl.float32)
7279
accumulator += tl.dot(a, b, allow_tf32=False)
7380
a_ptrs += BLOCK_SIZE_K * stride_ak
7481
b_ptrs += BLOCK_SIZE_K * stride_bk
@@ -131,6 +138,7 @@ def addmm(bias, mat1, mat2, *, beta=1, alpha=1):
131138
bias.stride(1),
132139
out.stride(0),
133140
out.stride(1),
141+
IS_FP64=mat1.dtype == torch.float64,
134142
)
135143
return out
136144

@@ -182,5 +190,6 @@ def addmm_out(bias, mat1, mat2, *, beta=1, alpha=1, out=None):
182190
bias.stride(1),
183191
out.stride(0),
184192
out.stride(1),
193+
IS_FP64=mat1.dtype == torch.float64,
185194
)
186195
return out

‎src/flag_gems/ops/baddbmm.py‎

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -49,6 +49,7 @@ def baddbmm_kernel(
4949
bias_batch_stride: tl.constexpr,
5050
bias_M_stride: tl.constexpr,
5151
bias_N_stride: tl.constexpr,
52+
IS_FP64: tl.constexpr = False,
5253
):
5354
# batch offsets
5455
pid_b = tle.program_id(2)
@@ -89,7 +90,10 @@ def baddbmm_kernel(
8990
o_ptrs = O + offs_m[:, None] * N + offs_n[None, :]
9091

9192
num_iters = tl.cdiv(K, TILE_K)
92-
accumulator = tl.zeros((TILE_M, TILE_N), dtype=tl.float32)
93+
if IS_FP64:
94+
accumulator = tl.zeros((TILE_M, TILE_N), dtype=tl.float64)
95+
else:
96+
accumulator = tl.zeros((TILE_M, TILE_N), dtype=tl.float32)
9397
for _ in range(num_iters):
9498
if DIVISIBLE_K:
9599
if DIVISIBLE_M:
@@ -173,6 +177,7 @@ def forward(ctx, bias, A, B, beta, alpha):
173177
bias_batch_stride=bias_batch_stride,
174178
bias_M_stride=bias_M_stride,
175179
bias_N_stride=bias_N_stride,
180+
IS_FP64=A.dtype == torch.float64,
176181
)
177182
return out
178183

‎src/flag_gems/ops/bmm.py‎

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -54,6 +54,7 @@ def bmm_kernel(
5454
DIVISIBLE_M: tl.constexpr,
5555
DIVISIBLE_N: tl.constexpr,
5656
DIVISIBLE_K: tl.constexpr,
57+
IS_FP64: tl.constexpr = False,
5758
):
5859
# batch offsets
5960
pid_b = tle.program_id(2)
@@ -96,7 +97,10 @@ def bmm_kernel(
9697
o_ptrs = O + offs_m[:, None] * stride_om + offs_n[None, :] * stride_on
9798

9899
num_iters = tl.cdiv(K, TILE_K)
99-
o = tl.zeros((TILE_M, TILE_N), dtype=tl.float32)
100+
if IS_FP64:
101+
o = tl.zeros((TILE_M, TILE_N), dtype=tl.float64)
102+
else:
103+
o = tl.zeros((TILE_M, TILE_N), dtype=tl.float32)
100104
for _ in range(num_iters):
101105
if DIVISIBLE_K:
102106
if DIVISIBLE_M:
@@ -168,6 +172,7 @@ def bmm(A, B):
168172
out.stride(0),
169173
out.stride(1),
170174
out.stride(2),
175+
IS_FP64=A.dtype == torch.float64,
171176
)
172177
return out
173178

@@ -201,5 +206,6 @@ def bmm_out(A, B, out):
201206
out.stride(0),
202207
out.stride(1),
203208
out.stride(2),
209+
IS_FP64=A.dtype == torch.float64,
204210
)
205211
return out

‎src/flag_gems/ops/mm.py‎

Lines changed: 15 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -49,6 +49,7 @@ def mm_kernel_general(
4949
BLOCK_N: tl.constexpr,
5050
BLOCK_K: tl.constexpr,
5151
GROUP_M: tl.constexpr,
52+
IS_FP64: tl.constexpr = False,
5253
):
5354
# matrix multiplication
5455
pid = tle.program_id(0)
@@ -69,15 +70,21 @@ def mm_kernel_general(
6970
rn = rn.to(tl.int64)
7071
prev_multiple = prev_multiple_of(K, BLOCK_K)
7172

72-
acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
73+
if IS_FP64:
74+
acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float64)
75+
else:
76+
acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
7377
for start_k in range(0, prev_multiple, BLOCK_K):
7478
rk = (start_k + tl.arange(0, BLOCK_K)).to(tl.int64)
7579
a = tl.load(A + (ram[:, None] * stride_am + rk[None, :] * stride_ak))
7680
b = tl.load(B + (rk[:, None] * stride_bk + rbn[None, :] * stride_bn))
7781
if a.dtype != b.dtype:
7882
a = a.to(C.dtype.element_ty)
7983
b = b.to(C.dtype.element_ty)
80-
acc += tl.dot(a, b, out_dtype=tl.float32, allow_tf32=False)
84+
if IS_FP64:
85+
acc += tl.dot(a, b, allow_tf32=False)
86+
else:
87+
acc += tl.dot(a, b, out_dtype=tl.float32, allow_tf32=False)
8188

8289
# loop peeling
8390
rk = (prev_multiple + tl.arange(0, BLOCK_K)).to(tl.int64)
@@ -95,7 +102,10 @@ def mm_kernel_general(
95102
if a.dtype != b.dtype:
96103
a = a.to(C.dtype.element_ty)
97104
b = b.to(C.dtype.element_ty)
98-
acc += tl.dot(a, b, out_dtype=tl.float32, allow_tf32=False)
105+
if IS_FP64:
106+
acc += tl.dot(a, b, allow_tf32=False)
107+
else:
108+
acc += tl.dot(a, b, out_dtype=tl.float32, allow_tf32=False)
99109

100110
acc = acc.to(C.dtype.element_ty)
101111
# rematerialize rm and rn to save registers
@@ -107,7 +117,7 @@ def mm_kernel_general(
107117
tl.store(C, acc, mask=mask)
108118

109119

110-
_ordered_datatypes = [torch.float16, torch.bfloat16, torch.float32]
120+
_ordered_datatypes = [torch.float16, torch.bfloat16, torch.float32, torch.float64]
111121

112122

113123
def get_higher_dtype(a, b):
@@ -152,6 +162,7 @@ def general_mm(a, b, c, M, N, K):
152162
c.stride(0),
153163
c.stride(1),
154164
GROUP_M=8,
165+
IS_FP64=a.dtype == torch.float64,
155166
)
156167
return c
157168

‎src/flag_gems/runtime/backend/_nvidia/hopper/ops/mm.py‎

Lines changed: 29 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -119,6 +119,7 @@ def mm_kernel_general(
119119
BLOCK_N: tl.constexpr,
120120
BLOCK_K: tl.constexpr,
121121
GROUP_M: tl.constexpr,
122+
IS_FP64: tl.constexpr = False,
122123
):
123124
# matrix multiplication
124125
pid = tle.program_id(0)
@@ -167,11 +168,17 @@ def mm_kernel_general(
167168
block_shape=[BLOCK_M, BLOCK_N],
168169
)
169170

170-
acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
171+
if IS_FP64:
172+
acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float64)
173+
else:
174+
acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
171175
for k in range(0, tl.cdiv(K, BLOCK_K)):
172176
a = a_desc.load([offset_am.to(tl.int32), offset_k.to(tl.int32)])
173177
b = b_desc.load([offset_k.to(tl.int32), offset_bn.to(tl.int32)])
174-
acc += tl.dot(a, b, out_dtype=tl.float32, allow_tf32=False)
178+
if IS_FP64:
179+
acc += tl.dot(a, b, allow_tf32=False)
180+
else:
181+
acc += tl.dot(a, b, out_dtype=tl.float32, allow_tf32=False)
175182
offset_k += BLOCK_K
176183

177184
acc = acc.to(a_desc.dtype)
@@ -187,15 +194,21 @@ def mm_kernel_general(
187194
rn = rn.to(tl.int64)
188195
prev_multiple = prev_multiple_of(K, BLOCK_K)
189196

190-
acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
197+
if IS_FP64:
198+
acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float64)
199+
else:
200+
acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
191201
for start_k in range(0, prev_multiple, BLOCK_K):
192202
rk = (start_k + tl.arange(0, BLOCK_K)).to(tl.int64)
193203
a = tl.load(A + (ram[:, None] * stride_am + rk[None, :] * stride_ak))
194204
b = tl.load(B + (rk[:, None] * stride_bk + rbn[None, :] * stride_bn))
195205
if a.dtype != b.dtype:
196206
a = a.to(C.dtype.element_ty)
197207
b = b.to(C.dtype.element_ty)
198-
acc += tl.dot(a, b, out_dtype=tl.float32, allow_tf32=False)
208+
if IS_FP64:
209+
acc += tl.dot(a, b, allow_tf32=False)
210+
else:
211+
acc += tl.dot(a, b, out_dtype=tl.float32, allow_tf32=False)
199212

200213
# loop peeling
201214
rk = (prev_multiple + tl.arange(0, BLOCK_K)).to(tl.int64)
@@ -213,7 +226,10 @@ def mm_kernel_general(
213226
if a.dtype != b.dtype:
214227
a = a.to(C.dtype.element_ty)
215228
b = b.to(C.dtype.element_ty)
216-
acc += tl.dot(a, b, out_dtype=tl.float32, allow_tf32=False)
229+
if IS_FP64:
230+
acc += tl.dot(a, b, allow_tf32=False)
231+
else:
232+
acc += tl.dot(a, b, out_dtype=tl.float32, allow_tf32=False)
217233

218234
acc = acc.to(C.dtype.element_ty)
219235
# rematerialize rm and rn to save registers
@@ -343,7 +359,7 @@ def mm_kernel_general_host_tma(
343359

344360

345361
def get_higher_dtype(a, b):
346-
_ordered_datatypes = [torch.float16, torch.bfloat16, torch.float32]
362+
_ordered_datatypes = [torch.float16, torch.bfloat16, torch.float32, torch.float64]
347363

348364
if a is b:
349365
return a
@@ -435,6 +451,7 @@ def alloc_fn(size: int, align: int, stream: Optional[int]):
435451
c.stride(0),
436452
c.stride(1),
437453
GROUP_M=8,
454+
IS_FP64=a.dtype == torch.float64,
438455
)
439456
return c
440457

@@ -471,6 +488,7 @@ def gemv_kernel(
471488
stride_bk,
472489
BLOCK_M: tl.constexpr,
473490
BLOCK_K: tl.constexpr,
491+
IS_FP64: tl.constexpr = False,
474492
):
475493
"""Optimized kernel for matrix-vector multiplication (N=1 case)"""
476494
pid = tl.program_id(0)
@@ -481,7 +499,10 @@ def gemv_kernel(
481499
row_mask = row_offset < M
482500

483501
# Accumulator for this block of rows
484-
acc = tl.zeros((BLOCK_M,), dtype=tl.float32)
502+
if IS_FP64:
503+
acc = tl.zeros((BLOCK_M,), dtype=tl.float64)
504+
else:
505+
acc = tl.zeros((BLOCK_M,), dtype=tl.float32)
485506

486507
# Iterate over K dimension
487508
for k_start in range(0, K, BLOCK_K):
@@ -525,6 +546,7 @@ def gemv_mm(a, b, c, M, K):
525546
a.stride(0),
526547
a.stride(1),
527548
b.stride(0),
549+
IS_FP64=a.dtype == torch.float64,
528550
)
529551
return c
530552

‎tests/test_addmm.py‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,7 @@
1717
(15, 160, 1024),
1818
(495, 5333, 71),
1919
]
20-
FLOAT_DTYPES = utils.FLOAT_DTYPES
20+
FLOAT_DTYPES = utils.ALL_FLOAT_DTYPES
2121

2222

2323
@pytest.mark.addmm
@@ -28,6 +28,8 @@
2828
def test_addmm(monkeypatch, M, N, K, scalar, dtype, b_column_major):
2929
if flag_gems.vendor_name == "tsingmicro" and dtype == torch.float32:
3030
pytest.skip("Skiping fp32 addmm test on tsingmicro platform")
31+
if dtype == torch.float64 and torch.cuda.get_device_capability()[0] < 9:
32+
pytest.skip("tl.dot does not support fp64 on compute capability < 9.0")
3133

3234
if flag_gems.vendor_name == "mthreads":
3335
monkeypatch.env("MUSA_ENABLE_SQMMA", "1")
@@ -67,6 +69,8 @@ def test_addmm(monkeypatch, M, N, K, scalar, dtype, b_column_major):
6769
def test_addmm_out(M, N, K, scalar, dtype):
6870
if flag_gems.vendor_name == "tsingmicro" and dtype == torch.float32:
6971
pytest.skip("Skiping fp32 addmm_out test on tsingmicro platform")
72+
if dtype == torch.float64 and torch.cuda.get_device_capability()[0] < 9:
73+
pytest.skip("tl.dot does not support fp64 on compute capability < 9.0")
7074

7175
mat1 = torch.randn((M, K), dtype=dtype, device=flag_gems.device)
7276
mat2 = torch.randn((K, N), dtype=dtype, device=flag_gems.device)

‎tests/test_baddbmm.py‎

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@
33

44
import flag_gems
55

6-
from .accuracy_utils import FLOAT_DTYPES as ORIG_FLOAT_DTYPES
6+
from .accuracy_utils import ALL_FLOAT_DTYPES as ORIG_ALL_FLOAT_DTYPES
77
from .accuracy_utils import SCALARS, gems_assert_close, to_reference
88
from .conftest import QUICK_MODE
99

@@ -18,7 +18,7 @@
1818
(15, 160, 1024),
1919
(495, 5333, 71),
2020
]
21-
FLOAT_DTYPES = ORIG_FLOAT_DTYPES
21+
FLOAT_DTYPES = ORIG_ALL_FLOAT_DTYPES
2222

2323
GNK_SHAPES = [(16, 512, 2048), (16, 2560, 2048), (64, 2048, 128)]
2424

@@ -37,6 +37,8 @@
3737
@pytest.mark.parametrize("scalar", SCALARS)
3838
@pytest.mark.parametrize("dtype", FLOAT_DTYPES)
3939
def test_baddbmm(monkeypatch, M, N, K, scalar, dtype):
40+
if dtype == torch.float64 and torch.cuda.get_device_capability()[0] < 9:
41+
pytest.skip("tl.dot does not support fp64 on compute capability < 9.0")
4042
if flag_gems.vendor_name == "mthreads" and dtype in [torch.float16, torch.bfloat16]:
4143
monkeypatch.setenv("MUSA_ENABLE_SQMMA", "1")
4244
batch = 4
@@ -60,6 +62,8 @@ def test_baddbmm(monkeypatch, M, N, K, scalar, dtype):
6062
@pytest.mark.parametrize("scalar", SCALARS)
6163
@pytest.mark.parametrize("dtype", FLOAT_DTYPES)
6264
def test_baddbmm_backward(M, N, K, scalar, dtype):
65+
if dtype == torch.float64 and torch.cuda.get_device_capability()[0] < 9:
66+
pytest.skip("tl.dot does not support fp64 on compute capability < 9.0")
6367
batch = 2
6468
mat1 = torch.randn(
6569
(batch, M, K), dtype=dtype, device=flag_gems.device, requires_grad=True

‎tests/test_bmm.py‎

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,7 @@
2121
(15, 160, 1024),
2222
(495, 5333, 71),
2323
]
24-
FLOAT_DTYPES = utils.FLOAT_DTYPES
24+
FLOAT_DTYPES = utils.ALL_FLOAT_DTYPES
2525

2626

2727
@pytest.mark.bmm
@@ -30,6 +30,8 @@
3030
def test_bmm(M, N, K, dtype):
3131
if flag_gems.vendor_name == "tsingmicro" and dtype == torch.float32:
3232
pytest.skip("Skiping fp32 bmm test on tsingmicro platform")
33+
if dtype == torch.float64 and torch.cuda.get_device_capability()[0] < 9:
34+
pytest.skip("tl.dot does not support fp64 on compute capability < 9.0")
3335

3436
if flag_gems.vendor_name == "mthreads":
3537
os.environ["MUSA_ENABLE_SQMMA"] = "1"
@@ -62,6 +64,8 @@ def test_bmm(M, N, K, dtype):
6264
def test_bmm_non_contiguous(M, N, K, dtype):
6365
if flag_gems.vendor_name == "tsingmicro" and dtype == torch.float32:
6466
pytest.skip("Skiping fp32 bmm test on tsingmicro platform")
67+
if dtype == torch.float64 and torch.cuda.get_device_capability()[0] < 9:
68+
pytest.skip("tl.dot does not support fp64 on compute capability < 9.0")
6569

6670
if flag_gems.vendor_name == "kunlunxin":
6771
torch.manual_seed(0)
@@ -94,6 +98,8 @@ def test_bmm_non_contiguous(M, N, K, dtype):
9498
def test_bmm_out(M, N, K, dtype):
9599
if flag_gems.vendor_name == "tsingmicro" and dtype == torch.float32:
96100
pytest.skip("Skiping fp32 bmm test on tsingmicro platform")
101+
if dtype == torch.float64 and torch.cuda.get_device_capability()[0] < 9:
102+
pytest.skip("tl.dot does not support fp64 on compute capability < 9.0")
97103

98104
if flag_gems.vendor_name == "kunlunxin":
99105
torch.manual_seed(0)

‎tests/test_mm.py‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -20,7 +20,7 @@
2020
(15, 160, 1024),
2121
(495, 5333, 71),
2222
]
23-
FLOAT_DTYPES = utils.FLOAT_DTYPES
23+
FLOAT_DTYPES = utils.ALL_FLOAT_DTYPES
2424

2525

2626
# TODO: failed at (1, 1, 2)
@@ -31,6 +31,8 @@
3131
def test_mm(M, N, K, dtype, b_column_major):
3232
if flag_gems.vendor_name == "tsingmicro" and dtype == torch.float32:
3333
pytest.skip("Skiping fp32 mm test on tsingmicro platform")
34+
if dtype == torch.float64 and torch.cuda.get_device_capability()[0] < 9:
35+
pytest.skip("tl.dot does not support fp64 on compute capability < 9.0")
3436

3537
torch.manual_seed(0)
3638
torch.cuda.manual_seed_all(0)

0 commit comments

Comments
 (0)