mirror of
https://github.com/tile-ai/tilelang.git
synced 2026-10-24 07:04:00 +08:00
cklxx
2a17fffd44
* [TileOP] Add SM70 GEMM FMA fallback * [TileOP] GemmFMA: drop redundant pre-staging sync and add docstrings Address PR #2279 review: - coderabbitai nit: the T.sync_threads() before the cooperative A/B stage loads is unnecessary because A_stage / B_stage are freshly allocated shared buffers — no other thread holds references to them, and the later sync before the GEMM compute is sufficient. - Docstring Coverage pre-merge check: add module-, helper-, method-, and prim_func-level docstrings so coverage clears the 80% threshold. No behavior change; the prim_func still stages through shared memory and accumulates via scalar FMA. SM70 pytest trio still expected to pass. * [TileOP] Harden SM70 GemmFMA fallback * [TileOP] Fix SM70 FMA staging dtypes * [TileOP] Tighten SM70 MMA dispatch for transposed A
126 lines
4.1 KiB
Python
126 lines
4.1 KiB
Python
import torch
|
|
|
|
import tilelang
|
|
import tilelang.language as T
|
|
import tilelang.testing
|
|
|
|
|
|
def _make_bf16_ss_gemm():
|
|
@T.prim_func
|
|
def main(
|
|
A: T.Tensor((64, 16), T.bfloat16),
|
|
B: T.Tensor((16, 16), T.bfloat16),
|
|
C: T.Tensor((64, 16), T.float32),
|
|
):
|
|
with T.Kernel(1, threads=128):
|
|
A_shared = T.alloc_shared((64, 16), T.bfloat16)
|
|
B_shared = T.alloc_shared((16, 16), T.bfloat16)
|
|
C_local = T.alloc_fragment((64, 16), T.float32)
|
|
|
|
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
|
|
|
|
|
|
def _make_bf16_ss_and_rs_gemm():
|
|
@T.prim_func
|
|
def main(
|
|
A: T.Tensor((64, 16), T.bfloat16),
|
|
B: T.Tensor((16, 16), T.bfloat16),
|
|
C: T.Tensor((64, 16), T.float32),
|
|
):
|
|
with T.Kernel(1, threads=128):
|
|
A_shared = T.alloc_shared((64, 16), T.bfloat16)
|
|
B_shared = T.alloc_shared((16, 16), T.bfloat16)
|
|
P = T.alloc_fragment((64, 16), T.float32)
|
|
P_bf16 = T.alloc_fragment((64, 16), T.bfloat16)
|
|
C_local = T.alloc_fragment((64, 16), T.float32)
|
|
|
|
T.copy(A, A_shared)
|
|
T.copy(B, B_shared)
|
|
|
|
T.clear(P)
|
|
T.gemm(A_shared, B_shared, P, policy=T.GemmWarpPolicy.FullRow)
|
|
T.copy(P, P_bf16)
|
|
|
|
T.clear(C_local)
|
|
T.gemm(P_bf16, B_shared, C_local, policy=T.GemmWarpPolicy.FullRow)
|
|
T.copy(C_local, C)
|
|
|
|
return main
|
|
|
|
|
|
def _make_fp16_transposed_a_ss_gemm():
|
|
@T.prim_func
|
|
def main(
|
|
A: T.Tensor((16, 64), T.float16),
|
|
B: T.Tensor((16, 16), T.float16),
|
|
C: T.Tensor((64, 16), T.float32),
|
|
):
|
|
with T.Kernel(1, threads=128):
|
|
A_shared = T.alloc_shared((16, 64), T.float16)
|
|
B_shared = T.alloc_shared((16, 16), T.float16)
|
|
C_local = T.alloc_fragment((64, 16), T.float32)
|
|
|
|
T.copy(A, A_shared)
|
|
T.copy(B, B_shared)
|
|
T.clear(C_local)
|
|
T.gemm(A_shared, B_shared, C_local, transpose_A=True, policy=T.GemmWarpPolicy.FullRow)
|
|
T.copy(C_local, C)
|
|
|
|
return main
|
|
|
|
|
|
@tilelang.testing.requires_cuda
|
|
@tilelang.testing.requires_cuda_compute_version_eq(7, 0)
|
|
def test_sm70_bf16_ss_gemm_uses_fma_fallback():
|
|
kernel = tilelang.compile(_make_bf16_ss_gemm(), target="cuda", out_idx=[2])
|
|
source = kernel.get_kernel_source()
|
|
assert "mma_sync_sm70<tl::DataType::kBFloat16" not in source
|
|
|
|
a = torch.randn((64, 16), device="cuda", dtype=torch.bfloat16)
|
|
b = torch.randn((16, 16), 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)
|
|
|
|
|
|
@tilelang.testing.requires_cuda
|
|
@tilelang.testing.requires_cuda_compute_version_eq(7, 0)
|
|
def test_sm70_bf16_rs_gemm_uses_fma_fallback():
|
|
kernel = tilelang.compile(_make_bf16_ss_and_rs_gemm(), target="cuda", out_idx=[2])
|
|
source = kernel.get_kernel_source()
|
|
assert "mma_sync_sm70<tl::DataType::kBFloat16" not in source
|
|
|
|
a = torch.randn((64, 16), device="cuda", dtype=torch.bfloat16) * 0.25
|
|
b = torch.randn((16, 16), device="cuda", dtype=torch.bfloat16) * 0.25
|
|
c = kernel(a, b)
|
|
p = (a.float() @ b.float()).to(torch.bfloat16)
|
|
ref = p.float() @ b.float()
|
|
|
|
tilelang.testing.torch_assert_close(c, ref, rtol=2e-2, atol=2e-2, max_mismatched_ratio=0.01)
|
|
|
|
|
|
@tilelang.testing.requires_cuda
|
|
@tilelang.testing.requires_cuda_compute_version_eq(7, 0)
|
|
def test_sm70_fp16_transposed_a_uses_fma_fallback():
|
|
kernel = tilelang.compile(_make_fp16_transposed_a_ss_gemm(), target="cuda", out_idx=[2])
|
|
source = kernel.get_kernel_source()
|
|
assert "mma_sync_sm70" not in source
|
|
|
|
a = torch.randn((16, 64), device="cuda", dtype=torch.float16) * 0.25
|
|
b = torch.randn((16, 16), device="cuda", dtype=torch.float16) * 0.25
|
|
c = kernel(a, b)
|
|
ref = a.t().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()
|