Skip to content

Fixes MX formats build for blackwell #2214

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.

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

Open
wants to merge 3 commits into
base: main
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from all commits
Commits
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
31 changes: 30 additions & 1 deletion setup.py
Original file line number Diff line number Diff line change
Expand Up @@ -364,6 +364,7 @@ def get_extensions():

use_cutlass = False
cutlass_90a_sources = None
cutlass_100a_sources = None
if use_cuda and not IS_ROCM and not IS_WINDOWS:
use_cutlass = True
cutlass_dir = os.path.join(third_party_path, "cutlass")
Expand Down Expand Up @@ -395,6 +396,8 @@ def get_extensions():
cuda_arch_flags = _get_cuda_arch_flags()
build_for_sm90 = "-gencode=arch=compute_90,code=sm_90" in cuda_arch_flags
build_for_sm90a = "-gencode=arch=compute_90a,code=sm_90a" in cuda_arch_flags
build_for_sm100 = "-gencode=arch=compute_100,code=sm_100" in cuda_arch_flags
build_for_sm100a = "-gencode=arch=compute_100a,code=sm_100a" in cuda_arch_flags
if build_for_sm90 and not build_for_sm90a:
cutlass_90a_sources = [
os.path.join(
Expand All @@ -418,6 +421,17 @@ def get_extensions():
)
)
sources = [s for s in sources if s not in cutlass_90a_sources]

if build_for_sm100 and not build_for_sm100a:
cutlass_100a_sources = [
os.path.join(
extensions_cuda_dir,
"mx_kernels",
"mx_fp_cutlass_kernels.cu",
),
]
sources = [s for s in sources if s not in cutlass_100a_sources]

else:
# Remove CUTLASS-based kernels from the sources list. An
# assumption is that these files will have "cutlass" in its
Expand Down Expand Up @@ -448,14 +462,29 @@ def get_extensions():
)
ext_modules.append(
extension(
"torchao._C",
"torchao._C_cutlass_90a",
cutlass_90a_sources,
py_limited_api=True,
extra_compile_args=cutlass_90a_extra_compile_args,
extra_link_args=extra_link_args,
)
)

if cutlass_100a_sources is not None and len(cutlass_100a_sources) > 0:
cutlass_100a_extra_compile_args = copy.deepcopy(extra_compile_args)
cutlass_100a_extra_compile_args["nvcc"].extend(
cuda_arch_flags + ["-gencode=arch=compute_100a,code=sm_100a"]
)
ext_modules.append(
extension(
"torchao._C_cutlass_100a",
cutlass_100a_sources,
py_limited_api=True,
extra_compile_args=cutlass_100a_extra_compile_args,
extra_link_args=extra_link_args,
)
)

# Build CMakeLists from /torchao/experimental - additional options become available : TORCHAO_BUILD_CPU_AARCH64, TORCHAO_BUILD_KLEIDIAI, TORCHAO_BUILD_MPS_OPS, TORCHAO_PARALLEL_BACKEND
if build_macos_arm_auto or os.getenv("BUILD_TORCHAO_EXPERIMENTAL") == "1":
build_options = BuildOptions()
Expand Down
4 changes: 2 additions & 2 deletions torchao/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -25,8 +25,8 @@

so_files = list(Path(__file__).parent.glob("_C*.so"))
if len(so_files) > 0:
assert len(so_files) == 1, f"Expected one _C*.so file, found {len(so_files)}"
torch.ops.load_library(str(so_files[0]))
for file in so_files:
torch.ops.load_library(str(file))
from . import ops

# The following library contains CPU kernels from torchao/experimental
Expand Down
Loading