Skip to content

[release/2.6] fix scaled matmul and test_float8_basics_cuda #2739

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
jeffdaily merged 1 commit into from
Oct 28, 2025
+9 −1
Merged

[release/2.6] fix scaled matmul and test_float8_basics_cuda #2739

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 6 additions & 0 deletions aten/src/ATen/cuda/CUDABlas.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1447,6 +1447,12 @@ void scaled_gemm(
#if defined(USE_ROCM)
#if defined(HIPBLASLT_OUTER_VEC)
// this case is handled later as hipified CUBLASLT_MATMUL_MATRIX_SCALE_OUTER_VEC_32F
if (use_rowwise) {
// swapped
computeDesc.setAttribute(CUBLASLT_MATMUL_DESC_A_SCALE_POINTER, mat2_scale_ptr);
computeDesc.setAttribute(CUBLASLT_MATMUL_DESC_B_SCALE_POINTER, mat1_scale_ptr);
}
else
#elif defined(HIPBLASLT_VEC_EXT)
if (use_rowwise) {
// swapped
Expand Down
4 changes: 3 additions & 1 deletion test/test_matmul_cuda.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
# Owner(s): ["module: linear algebra"]

from contextlib import nullcontext
import unittest
from itertools import product
from functools import partial
Expand Down Expand Up @@ -356,7 +357,8 @@ def test_float8_basics(self, device) -> None:
self._test_tautological_mm(device, e4m3_type, e5m2_type, size=32)
self._test_tautological_mm(device, e5m2_type, e4m3_type, size=48)
# According to https://docs.nvidia.com/cuda/cublas/#id99 8F_E5M2 MM is unsupported
with self.assertRaises(RuntimeError):
# supported on ROCm but fails on CUDA
with self.assertRaises(RuntimeError) if torch.version.hip is None else nullcontext():
self._test_tautological_mm(device, e5m2_type, e5m2_type)

self._test_tautological_mm(device, size=64, out_dtype=torch.float16)
Expand Down