tile-ai/tilelang
View on GitHub[Bug] Tilelang the Float32 data type looks like is not usable.
Open
#859 opened on Sep 22, 2025
enhancementgood first issue
Repository metrics
- Stars
- (6,817 stars)
- PR merge metrics
- (PR metrics pending)
Description
import torch
import torch.nn.functional as F
import tilelang
import tilelang.language as T
@tilelang.jit(out_idx=[-1])
def single_gemm_kernel(
M, N, K,
block_M, block_N, block_K,
dtype="float32",
accum_dtype="float"
):
"""
TileLang kernel for a single GEMM operation in float32.
Computes: X @ W
"""
@T.prim_func
def single_gemm(
X: T.Tensor((M, K), dtype),
W: T.Tensor((K, N), dtype),
Out: T.Tensor((M, N), dtype),
):
with T.Kernel(T.ceildiv(N, block_N), T.ceildiv(M, block_M), threads=128) as (bx, by):
X_shared = T.alloc_shared((block_M, block_K), dtype)
W_shared = T.alloc_shared((block_K, block_N), dtype)
output_local = T.alloc_fragment((block_M, block_N), accum_dtype)
T.clear(output_local)
for k in T.Pipelined(T.ceildiv(K, block_K), num_stages=3):
T.copy(X[by * block_M, k * block_K], X_shared)
T.copy(W[k * block_K, bx * block_N], W_shared)
# Perform only ONE GEMM operation
T.gemm(X_shared, W_shared, output_local)
T.copy(output_local, Out[by * block_M, bx * block_N])
return single_gemm
def main():
# Define problem size
M, K, N = 4096, 2048, 1024
dtype = torch.float32
device = "cuda"
# Create random input tensors
X = torch.randn(M, K, dtype=dtype, device=device)
W = torch.randn(K, N, dtype=dtype, device=device)
# Define tiling parameters
block_M = 64
block_N = 64
block_K = 32
# Instantiate and compile the simplified kernel
kernel = single_gemm_kernel(M, N, K, block_M, block_N, block_K, dtype="float32")
# Run the TileLang kernel
output_tilelang = kernel(X, W)
# --- Verification using PyTorch ---
output_ref = X @ W
print("Verifying correctness for a single float32 GEMM operation...")
try:
torch.testing.assert_close(output_tilelang, output_ref, rtol=1e-5, atol=1e-5)
print("✅ Correctness check passed for single GEMM!")
except AssertionError as e:
print("❌ Correctness check FAILED for single GEMM!")
print(e)
if __name__ == "__main__":
main()
(venv) ➜ test git:(master) ✗ python tilelang_swiglu.py
Verifying correctness for a single float32 GEMM operation...
❌ Correctness check FAILED for single GEMM!
Tensor-likes are not close!
Mismatched elements: 4179519 / 4194304 (99.6%)
Greatest absolute difference: 0.2025909423828125 at index (599, 166) (up to 1e-05 allowed)
Greatest relative difference: 2067.269287109375 at index (2778, 409) (up to 1e-05 allowed)
When data type been switch to float16, 100% pass. The problem should be reproducible on your Tilelang environment.