Fix FP8 Quack GEMM call when inference_mode() and torch.compile() is used. (#5153)
Signed-off-by: Bartosz Stefaniak <bstefaniak@nvidia.com> Co-authored-by: Roger Wang <hey@rogerw.io>
This commit is contained in:
1 parent
ae949d129f
commit
07957e12f4
1 file changed
+1
-1
@@ -94,7 +94,7 @@ def quack_scaled_fp8_mm(
|
||||
return None
|
||||
out = torch.empty(a.shape[0], b.shape[1], device=a.device, dtype=out_dtype)
|
||||
alpha = scale_a.reshape(1).float() * scale_b.reshape(1).float()
|
||||
gemm.gemm_tuned(a, b, out, bias=bias, alpha=alpha)
|
||||
gemm.gemm(a, b, out=out, bias=bias, alpha=alpha, tuned=True)
|
||||
return out
|
||||
|
||||
|
||||
|
||||
Reference in new issue
Block a user