Skip to content

Commit 994bed0

Browse files
wuhulalalawuhulalala
authored
fix mix MatMul kernels fail with fp64 inputs due to accumulator type … (#2606)
1 parent acc295c commit 994bed0

6 files changed

Lines changed: 68 additions & 14 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
@@ -440,6 +456,7 @@ def alloc_fn(size: int, align: int, stream: Optional[int]):
440456
c.stride(0),
441457
c.stride(1),
442458
GROUP_M=8,
459+
IS_FP64=a.dtype == torch.float64,
443460
)
444461
return c
445462

@@ -476,6 +493,7 @@ def gemv_kernel(
476493
stride_bk,
477494
BLOCK_M: tl.constexpr,
478495
BLOCK_K: tl.constexpr,
496+
IS_FP64: tl.constexpr = False,
479497
):
480498
"""Optimized kernel for matrix-vector multiplication (N=1 case)"""
481499
pid = tl.program_id(0)
@@ -486,7 +504,10 @@ def gemv_kernel(
486504
row_mask = row_offset < M
487505

488506
# Accumulator for this block of rows
489-
acc = tl.zeros((BLOCK_M,), dtype=tl.float32)
507+
if IS_FP64:
508+
acc = tl.zeros((BLOCK_M,), dtype=tl.float64)
509+
else:
510+
acc = tl.zeros((BLOCK_M,), dtype=tl.float32)
490511

491512
# Iterate over K dimension
492513
for k_start in range(0, K, BLOCK_K):
@@ -530,6 +551,7 @@ def gemv_mm(a, b, c, M, K):
530551
a.stride(0),
531552
a.stride(1),
532553
b.stride(0),
554+
IS_FP64=a.dtype == torch.float64,
533555
)
534556
return c
535557

‎tests/test_mm.py‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@
2222
]
2323
FLOAT_DTYPES = utils.FLOAT_DTYPES
2424

25+
2526
MK_SHAPES = (
2627
[(1, 32)]
2728
if QUICK_MODE

0 commit comments

Comments
 (0)