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
49 lines
1.5 KiB
Python
49 lines
1.5 KiB
Python
import torch
|
|
|
|
import tilelang
|
|
import tilelang.language as T
|
|
import tilelang.testing
|
|
|
|
|
|
def _make_fragment_cast_into_rs_gemm():
|
|
@T.prim_func
|
|
def main(
|
|
A: T.Tensor((64, 16), 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((64, 16), T.float16)
|
|
B_shared = T.alloc_shared((16, 16), T.float16)
|
|
P = T.alloc_fragment((64, 16), T.float32)
|
|
P_half = T.alloc_fragment((64, 16), T.float16)
|
|
C_local = T.alloc_fragment((64, 16), T.float32)
|
|
|
|
T.copy(A, A_shared)
|
|
T.copy(B, B_shared)
|
|
|
|
for i, j in T.Parallel(64, 16):
|
|
P[i, j] = T.cast(A_shared[i, j], T.float32)
|
|
|
|
T.copy(P, P_half)
|
|
T.clear(C_local)
|
|
T.gemm(P_half, B_shared, C_local, 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_fragment_cast_copy_feeds_rs_gemm():
|
|
kernel = tilelang.compile(_make_fragment_cast_into_rs_gemm(), target="cuda", out_idx=[2])
|
|
source = kernel.get_kernel_source()
|
|
assert "tl_frag_copy" not in source
|
|
|
|
a = torch.randn((64, 16), device="cuda", dtype=torch.float16)
|
|
b = torch.randn((16, 16), device="cuda", dtype=torch.float16)
|
|
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)
|