tile-ai/tilelang

[Bug] Tilelang the Float32 data type looks like is not usable.

Open

#859 opened on Sep 22, 2025

View on GitHub
 (6 comments) (0 reactions) (0 assignees)Python (648 forks)auto 404
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.

Contributor guide