mirror of
https://github.com/tile-ai/tilelang.git
synced 2026-10-24 07:04:00 +08:00
Chennes
7fd95363cc
On SM75, SelectInst routed GEMMs to the plain MMA path even when their dtype combination had no native Turing mma.sync atom. In particular, fp32, fp64, and bf16 selected SM80-only atoms: bf16 failed an nvcc static_assert, while fp32 and fp64 compiled into runtime traps that invalidated the CUDA context. Add AllowTuringMma alongside the Volta check introduced in #2339 and select the existing cuda.fma implementation when SM75 has no native atom. Keep the check dtype-only because the SM75 emitter supports the same operand scopes and transposes as the generic MMA path. Valid f16 combinations remain on native MMA rather than being demoted to the slower and less accurate serial FMA fallback. Add compile-only dispatch checks pinned to an explicit sm_75 target, so they run on any CUDA runner, verifying that fp32, fp64, and bf16 use the fallback while f16 remains on MMA. Add execution checks gated on compute capability 7.5.
103 lines
4.0 KiB
Python
103 lines
4.0 KiB
Python
import torch
|
|
|
|
import tilelang
|
|
import tilelang.language as T
|
|
import tilelang.testing
|
|
|
|
_SM75 = {"kind": "cuda", "arch": "sm_75"}
|
|
|
|
|
|
def _make_ss_gemm(M, N, K, in_dtype, accum_dtype, threads=128):
|
|
@T.prim_func
|
|
def main(
|
|
A: T.Tensor((M, K), in_dtype),
|
|
B: T.Tensor((K, N), in_dtype),
|
|
C: T.Tensor((M, N), accum_dtype),
|
|
):
|
|
with T.Kernel(1, threads=threads):
|
|
A_shared = T.alloc_shared((M, K), in_dtype)
|
|
B_shared = T.alloc_shared((K, N), in_dtype)
|
|
C_local = T.alloc_fragment((M, N), accum_dtype)
|
|
|
|
T.copy(A, A_shared)
|
|
T.copy(B, B_shared)
|
|
T.clear(C_local)
|
|
T.gemm(A_shared, B_shared, C_local, policy=T.GemmWarpPolicy.FullRow)
|
|
T.copy(C_local, C)
|
|
|
|
return main
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Compile-only dispatch checks pinned to an explicit sm_75 target, so they run
|
|
# on any CUDA runner regardless of the host GPU. Turing has no fp32/tf32,
|
|
# fp64, or bf16 mma.sync atoms; these dtypes must lower through the FMA
|
|
# fallback instead of tl::mma_sync (which fails nvcc's static_assert for bf16
|
|
# and compiles into a cutlass runtime trap for fp32/fp64).
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@tilelang.testing.requires_cuda
|
|
def test_sm75_fp32_gemm_compiles_to_fma_fallback():
|
|
kernel = tilelang.compile(_make_ss_gemm(64, 64, 32, T.float32, T.float32), target=_SM75, out_idx=[2])
|
|
assert "tl::mma_sync" not in kernel.get_kernel_source()
|
|
|
|
|
|
@tilelang.testing.requires_cuda
|
|
def test_sm75_fp64_gemm_compiles_to_fma_fallback():
|
|
kernel = tilelang.compile(_make_ss_gemm(32, 32, 32, T.float64, T.float64), target=_SM75, out_idx=[2])
|
|
assert "tl::mma_sync" not in kernel.get_kernel_source()
|
|
|
|
|
|
@tilelang.testing.requires_cuda
|
|
def test_sm75_bf16_gemm_compiles_to_fma_fallback():
|
|
kernel = tilelang.compile(_make_ss_gemm(64, 64, 32, T.bfloat16, T.float32), target=_SM75, out_idx=[2])
|
|
assert "tl::mma_sync" not in kernel.get_kernel_source()
|
|
|
|
|
|
@tilelang.testing.requires_cuda
|
|
def test_sm75_fp16_gemm_keeps_mma_path():
|
|
"""f16 has native Turing atoms; the fallback must not demote it."""
|
|
kernel = tilelang.compile(_make_ss_gemm(64, 64, 32, T.float16, T.float32), target=_SM75, out_idx=[2])
|
|
assert "tl::mma_sync" in kernel.get_kernel_source()
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Execution checks on real Turing hardware.
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@tilelang.testing.requires_cuda
|
|
@tilelang.testing.requires_cuda_compute_version_eq(7, 5)
|
|
def test_sm75_fp32_gemm_fma_matches_torch():
|
|
kernel = tilelang.compile(_make_ss_gemm(64, 64, 32, T.float32, T.float32), target="cuda", out_idx=[2])
|
|
a = torch.randn((64, 32), device="cuda", dtype=torch.float32)
|
|
b = torch.randn((32, 64), device="cuda", dtype=torch.float32)
|
|
c = kernel(a, b)
|
|
tilelang.testing.torch_assert_close(c, a @ b, rtol=1e-4, atol=1e-4)
|
|
|
|
|
|
@tilelang.testing.requires_cuda
|
|
@tilelang.testing.requires_cuda_compute_version_eq(7, 5)
|
|
def test_sm75_fp64_gemm_fma_matches_torch():
|
|
kernel = tilelang.compile(_make_ss_gemm(32, 32, 32, T.float64, T.float64), target="cuda", out_idx=[2])
|
|
a = torch.randn((32, 32), device="cuda", dtype=torch.float64)
|
|
b = torch.randn((32, 32), device="cuda", dtype=torch.float64)
|
|
c = kernel(a, b)
|
|
tilelang.testing.torch_assert_close(c, a @ b, rtol=1e-8, atol=1e-8)
|
|
|
|
|
|
@tilelang.testing.requires_cuda
|
|
@tilelang.testing.requires_cuda_compute_version_eq(7, 5)
|
|
def test_sm75_bf16_gemm_fma_matches_torch():
|
|
kernel = tilelang.compile(_make_ss_gemm(64, 64, 32, T.bfloat16, T.float32), target="cuda", out_idx=[2])
|
|
a = torch.randn((64, 32), device="cuda", dtype=torch.bfloat16)
|
|
b = torch.randn((32, 64), device="cuda", dtype=torch.bfloat16)
|
|
c = kernel(a, b)
|
|
ref = a.float() @ b.float()
|
|
tilelang.testing.torch_assert_close(c, ref, rtol=1e-2, atol=1e-2, max_mismatched_ratio=0.01)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
tilelang.testing.main()
|