Repository: triton-lang/triton

Issue #1671: fused attention has issues with FP32 when allow_tf32 is enabled
Created: 2023-05-15T21:42:47Z
Hello,
I was testing the fused attention kernel in tutorials and think I may find a bug with allow_tf32.
I have attached my diff below. The kernel works correctly with fp16 and fp32 when allow_tf32 is disabled. However, when allow_tf32 is enabled, the fp32 results seem completely off.
$ python3 06-fused-attention.py
Testing dtype = torch.float16, allow_tf32 = True
trion = 0.03875732421875
torch = 0.03875732421875
diff = 0.0
Testing dtype = torch.float16, allow_tf32 = False
trion = 0.03875732421875
torch = 0.03875732421875
diff = 0.0
Testing dtype = torch.float32, allow_tf32 = True
trion = 0.2014191895723343
torch = 0.03874829784035683
diff = 0.16267089545726776
Testing dtype = torch.float32, allow_tf32 = False
trion = 0.03874829784035683
torch = 0.03874829784035683
diff = 0.0

diff --git a/python/tutorials/06-fused-attention.py b/python/tutorials/06-fused-attention.py
index c9875cf35..1e282677c 100644
--- a/python/tutorials/06-fused-attention.py
+++ b/python/tutorials/06-fused-attention.py
@@ -25,6 +25,7 @@ def _fwd_kernel(
     Z, H, N_CTX,
     BLOCK_M: tl.constexpr, BLOCK_DMODEL: tl.constexpr,
     BLOCK_N: tl.constexpr,
+    ALLOW_TF32: tl.constexpr,
 ):
     start_m = tl.program_id(0)
     off_hz = tl.program_id(1)
@@ -50,7 +51,7 @@ def _fwd_kernel(
         # -- compute qk ----
         k = tl.load(k_ptrs)
         qk = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32)
-        qk += tl.dot(q, k)
+        qk += tl.dot(q, k, allow_tf32=ALLOW_TF32)
         qk *= sm_scale
         qk = tl.where(offs_m[:, None] >= (start_n + offs_n[None, :]), qk, float("-inf"))
         # compute new m
@@ -67,7 +68,7 @@ def _fwd_kernel(
         # update acc
         p = p.to(Q.dtype.element_ty)
         v = tl.load(v_ptrs)
-        acc += tl.dot(p, v)
+        acc += tl.dot(p, v, allow_tf32=ALLOW_TF32)
         # update m_i and l_i
         l_prev = l_curr
         m_prev = m_curr
@@ -122,6 +123,7 @@ def _bwd_kernel(
     num_block,
     BLOCK_M: tl.constexpr, BLOCK_DMODEL: tl.constexpr,
     BLOCK_N: tl.constexpr,
+    ALLOW_TF32: tl.constexpr,
 ):
     off_hz = tl.program_id(0)
     off_z = off_hz // H
@@ -163,24 +165,24 @@ def _bwd_kernel(
             q = tl.load(q_ptrs)
             # recompute p = softmax(qk, dim=-1).T
             # NOTE: `do` is pre-divided by `l`; no normalization here
-            qk = tl.dot(q, tl.trans(k))
+            qk = tl.dot(q, tl.trans(k), allow_tf32=ALLOW_TF32)
             qk = tl.where(offs_m_curr[:, None] >= (offs_n[None, :]), qk, float("-inf"))
             m = tl.load(m_ptrs + offs_m_curr)
             p = tl.exp(qk * sm_scale - m[:, None])
             # compute dv
             do = tl.load(do_ptrs)
-            dv += tl.dot(tl.trans(p.to(Q.dtype.element_ty)), do)
+            dv += tl.dot(tl.trans(p.to(Q.dtype.element_ty)), do, allow_tf32=ALLOW_TF32)
             # compute dp = dot(v, do)
             Di = tl.load(D_ptrs + offs_m_curr)
             dp = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32) - Di[:, None]
-            dp += tl.dot(do, tl.trans(v))
+            dp += tl.dot(do, tl.trans(v), allow_tf32=ALLOW_TF32)
             # compute ds = p * (dp - delta[:, None])
             ds = p * dp * sm_scale
             # compute dk = dot(ds.T, q)
-            dk += tl.dot(tl.trans(ds.to(Q.dtype.element_ty)), q)
+            dk += tl.dot(tl.trans(ds.to(Q.dtype.element_ty)), q, allow_tf32=ALLOW_TF32)
             # compute dq
             dq = tl.load(dq_ptrs)
-            dq += tl.dot(ds.to(Q.dtype.element_ty), k)
+            dq += tl.dot(ds.to(Q.dtype.element_ty), k, allow_tf32=ALLOW_TF32)
             tl.store(dq_ptrs, dq)
             # increment pointers
             dq_ptrs += BLOCK_M * stride_qm
@@ -197,9 +199,8 @@ empty = torch.empty(128, device="cuda")
 
 
 class _attention(torch.autograd.Function):
-
     @staticmethod
-    def forward(ctx, q, k, v, sm_scale):
+    def forward(ctx, q, k, v, sm_scale, allow_tf32):
         BLOCK = 128
         # shape constraints
         Lq, Lk, Lv = q.shape[-1], k.shape[-1], v.shape[-1]
@@ -223,6 +224,7 @@ class _attention(torch.autograd.Function):
             BLOCK_M=BLOCK, BLOCK_N=BLOCK,
             BLOCK_DMODEL=Lk, num_warps=num_warps,
             num_stages=2,
+            ALLOW_TF32=allow_tf32,
         )
         # print(h.asm["ttgir"])
 
@@ -230,6 +232,7 @@ class _attention(torch.autograd.Function):
         ctx.grid = grid
         ctx.sm_scale = sm_scale
         ctx.BLOCK_DMODEL = Lk
+        ctx.allow_tf32 = allow_tf32
         return o
 
     @staticmethod
@@ -242,7 +245,7 @@ class _attention(torch.autograd.Function):
         dv = torch.empty_like(v)
         do_scaled = torch.empty_like(do)
         delta = torch.empty_like(l)
-        _bwd_preprocess[(ctx.grid[0] * ctx.grid[1], )](
+        _bwd_preprocess[(ctx.grid[0] * ctx.grid[1],)](
             o, do, l,
             do_scaled, delta,
             BLOCK_M=BLOCK, D_HEAD=ctx.BLOCK_DMODEL,
@@ -261,16 +264,17 @@ class _attention(torch.autograd.Function):
             BLOCK_M=BLOCK, BLOCK_N=BLOCK,
             BLOCK_DMODEL=ctx.BLOCK_DMODEL, num_warps=8,
             num_stages=1,
+            ALLOW_TF32=ctx.allow_tf32,
         )
         # print(h.asm["ttgir"])
-        return dq, dk, dv, None
+        return dq, dk, dv, None, None
 
 
 attention = _attention.apply
 
 
-@pytest.mark.parametrize('Z, H, N_CTX, D_HEAD', [(4, 48, 1024, 64)])
-def test_op(Z, H, N_CTX, D_HEAD, dtype=torch.float16):
+def test_op(Z, H, N_CTX, D_HEAD, dtype=torch.float16, allow_tf32=True):
+    print(f"Testing dtype = {dtype}, allow_tf32 = {allow_tf32}")
     torch.manual_seed(20)
     q = torch.empty((Z, H, N_CTX, D_HEAD), dtype=dtype, device="cuda").normal_(mean=0.1, std=0.2).requires_grad_()
     k = torch.empty((Z, H, N_CTX, D_HEAD), dtype=dtype, device="cuda").normal_(mean=0.4, std=0.2).requires_grad_()
@@ -283,79 +287,20 @@ def test_op(Z, H, N_CTX, D_HEAD, dtype=torch.float16):
     for z in range(Z):
         for h in range(H):
             p[:, :, M == 0] = float("-inf")
-    p = torch.softmax(p.float(), dim=-1).half()
-    # p = torch.exp(p)
+    p = torch.softmax(p.float(), dim=-1).to(dtype)
+    # ref implemenntation
     ref_out = torch.matmul(p, v)
-    ref_out.backward(dout)
-    ref_dv, v.grad = v.grad.clone(), None
-    ref_dk, k.grad = k.grad.clone(), None
-    ref_dq, q.grad = q.grad.clone(), None
     # # triton implementation
-    tri_out = attention(q, k, v, sm_scale)
-    # print(ref_out)
-    # print(tri_out)
-    tri_out.backward(dout)
-    tri_dv, v.grad = v.grad.clone(), None
-    tri_dk, k.grad = k.grad.clone(), None
-    tri_dq, q.grad = q.grad.clone(), None
-    # compare
-    assert torch.allclose(ref_out, tri_out, atol=1e-2, rtol=0)
-    assert torch.allclose(ref_dv, tri_dv, atol=1e-2, rtol=0)
-    assert torch.allclose(ref_dk, tri_dk, atol=1e-2, rtol=0)
-    assert torch.allclose(ref_dq, tri_dq, atol=1e-2, rtol=0)
-
-
-try:
-    from flash_attn.flash_attn_interface import flash_attn_func
-    HAS_FLASH = True
-except BaseException:
-    HAS_FLASH = False
+    tri_out = attention(q, k, v, sm_scale, allow_tf32)
+    print(f"trion = {tri_out[0][0][0][0].cpu().item()}")
+    print(f"torch = {ref_out[0][0][0][0].cpu().item()}")
+    print(f"diff = {(tri_out[0][0][0][0] - ref_out[0][0][0][0]).cpu().item()}")
 
-BATCH, N_HEADS, N_CTX, D_HEAD = 4, 48, 4096, 64
-# vary seq length for fixed head and batch=4
-configs = [triton.testing.Benchmark(
-    x_names=['N_CTX'],
-    x_vals=[2**i for i in range(10, 14)],
-    line_arg='provider',
-    line_vals=['triton'] + (['flash'] if HAS_FLASH else []),
-    line_names=['Triton'] + (['Flash'] if HAS_FLASH else []),
-    styles=[('red', '-'), ('blue', '-')],
-    ylabel='ms',
-    plot_name=f'fused-attention-batch{BATCH}-head{N_HEADS}-d{D_HEAD}-{mode}',
-    args={'H': N_HEADS, 'BATCH': BATCH, 'D_HEAD': D_HEAD, 'dtype': torch.float16, 'mode': mode}
-) for mode in ['fwd', 'bwd']]
 
+test_op(4, 48, 1024, 32, dtype=torch.float16, allow_tf32=True)
 
-@triton.testing.perf_report(configs)
-def bench_flash_attention(BATCH, H, N_CTX, D_HEAD, mode, provider, dtype=torch.float16, device="cuda"):
-    assert mode in ['fwd', 'bwd']
-    warmup = 25
-    rep = 100
-    if provider == "triton":
-        q = torch.randn((BATCH, H, N_CTX, D_HEAD), dtype=dtype, device="cuda", requires_grad=True)
-        k = torch.randn((BATCH, H, N_CTX, D_HEAD), dtype=dtype, device="cuda", requires_grad=True)
-        v = torch.randn((BATCH, H, N_CTX, D_HEAD), dtype=dtype, device="cuda", requires_grad=True)
-        sm_scale = 1.3
-        fn = lambda: attention(q, k, v, sm_scale)
-        if mode == 'bwd':
-            o = fn()
-            do = torch.randn_like(o)
-            fn = lambda: o.backward(do, retain_graph=True)
-        ms = triton.testing.do_bench(fn, warmup=warmup, rep=rep)
-        return ms
-    if provider == "flash":
-        lengths = torch.full((BATCH,), fill_value=N_CTX, device=device)
-        cu_seqlens = torch.zeros((BATCH + 1,), device=device, dtype=torch.int32)
-        cu_seqlens[1:] = lengths.cumsum(0)
-        qkv = torch.randn((BATCH * N_CTX, 3, H, D_HEAD), dtype=dtype, device=device, requires_grad=True)
-        fn = lambda: flash_attn_func(qkv, cu_seqlens, 0., N_CTX, causal=True)
-        if mode == 'bwd':
-            o = fn()
-            do = torch.randn_like(o)
-            fn = lambda: o.backward(do, retain_graph=True)
-        ms = triton.testing.do_bench(fn, warmup=warmup, rep=rep)
-        return ms
+test_op(4, 48, 1024, 32, dtype=torch.float16, allow_tf32=False)
 
+test_op(4, 48, 1024, 32, dtype=torch.float32, allow_tf32=True)
 
-# only works on post-Ampere GPUs right now
-bench_flash_attention.run(save_path='.', print_data=True)
+test_op(4, 48, 1024, 32, dtype=torch.float32, allow_tf32=False)

Duplicate issue #1840: Wrong results when using `tl.dot` with `tl.trans`
Created: 2023-06-25T10:25:49Z
I'm trying to implement a linear kernel i.e. C = A @ B.T in Triton. Since it's just a matmul with the second argument transposed, I start to modify based on the tutorial code. But I'm confused that it works on my PC (1650 Ti, sm_75 arch, Turing arch) but fails on A100 (Ampere).
The following is a reproduction:
import torch
import triton
import triton.language as tl


@triton.autotune(
    configs=[
        triton.Config(
            {
                "BLOCK_SIZE_M": 32,
                "BLOCK_SIZE_N": 64,
                "BLOCK_SIZE_K": 32,
                "GROUP_SIZE_M": 8,
            },
            num_stages=3,
            num_warps=8,
        ),
    ],
    key=["M", "N", "K"],
)
@triton.jit
def linear_kernel(
    a_ptr,
    b_ptr,
    c_ptr,
    # Matrix dimensions
    M,
    N,
    K,
    # Strides
    stride_am,
    stride_ak,
    stride_bn,
    stride_bk,
    stride_cm,
    stride_cn,
    # Meta-parameters
    BLOCK_SIZE_M: tl.constexpr,
    BLOCK_SIZE_N: tl.constexpr,
    BLOCK_SIZE_K: tl.constexpr,
    GROUP_SIZE_M: tl.constexpr,
):
    """Kernel for computing the linear C = A x B^T.
    A has shape (M, K), B has shape (N, K) and C has shape (M, N).
    Here M=batch_size, N=input_features, K=output_features.
    """
    # -----------------------------------------------------------
    # Map program ids `pid` to the block of C it should compute.
    # This is done in a grouped ordering to promote L2 data reuse
    # See above `L2 Cache Optimizations` section for details
    pid = tl.program_id(axis=0)
    num_pid_m = tl.cdiv(M, BLOCK_SIZE_M)
    num_pid_n = tl.cdiv(N, BLOCK_SIZE_N)
    num_pid_in_group = GROUP_SIZE_M * num_pid_n
    group_id = pid // num_pid_in_group
    first_pid_m = group_id * GROUP_SIZE_M
    group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M)
    pid_m = first_pid_m + (pid % group_size_m)
    pid_n = (pid % num_pid_in_group) // group_size_m

    # ----------------------------------------------------------
    # Create pointers for the first blocks of A and B.
    # We will advance this pointer as we move in the K direction
    # and accumulate
    # a_ptrs is a block of [BLOCK_SIZE_M, BLOCK_SIZE_K] pointers
    # b_ptrs is a block of [BLOCK_SIZE_N, BLOCK_SIZE_K] pointers
    # see above `Pointer Arithmetics` section for details
    offs_am = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
    offs_bn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
    offs_k = tl.arange(0, BLOCK_SIZE_K)
    a_ptrs = a_ptr + (offs_am[:, None] * stride_am + offs_k[None, :] * stride_ak)
    b_ptrs = b_ptr + (offs_bn[:, None] * stride_bn + offs_k[None, :] * stride_bk)

    # -----------------------------------------------------------
    # Iterate to compute a block of the C matrix
    # We accumulate into a `[BLOCK_SIZE_M, BLOCK_SIZE_N]` block
    # of fp32 values for higher accuracy.
    # `accumulator` will be converted back to fp16 after the loop
    accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
    for k in range(0, K, BLOCK_SIZE_K):
        # Note that for simplicity, we don't apply a mask here.
        # This means that if K is not a multiple of BLOCK_SIZE_K,
        # this will access out-of-bounds memory and produce an
        # error or (worse!) incorrect results.
        a = tl.load(a_ptrs)
        b = tl.load(b_ptrs)
        # We accumulate along the K dimension
        accumulator += tl.dot(a, tl.trans(b))
        # Advance the ptrs to the next K block
        a_ptrs += BLOCK_SIZE_K * stride_ak
        b_ptrs += BLOCK_SIZE_K * stride_bk

    # while the accumulator is still in FP32!
    c = accumulator.to(tl.float32)

    # -----------------------------------------------------------
    # Write back the block of the output matrix C
    offs_cm = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
    offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
    c_ptrs = c_ptr + stride_cm * offs_cm[:, None] + stride_cn * offs_cn[None, :]
    c_mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < N)
    tl.store(c_ptrs, c, mask=c_mask)


def linear(input, weight):
    # checks constraints
    assert input.shape[1] == weight.shape[1], "incompatible dimensions"
    assert input.is_contiguous(), "matrix A must be contiguous"
    assert weight.is_contiguous(), "matrix B must be contiguous"
    M, K = input.shape
    N, K = weight.shape
    assert (
        K % 32 == 0
    ), "We don't check memory-out-of-bounds with K so K must be divisible by BLOCK_SIZE_K"
    # allocates output
    c = torch.empty((M, N), device=a.device, dtype=a.dtype)
    # 1D launch kernel where each block gets its own program.
    grid = lambda META: (
        triton.cdiv(M, META["BLOCK_SIZE_M"]) * triton.cdiv(N, META["BLOCK_SIZE_N"]),
    )
    bin = linear_kernel[grid](
        a,
        b,
        c,
        M,
        N,
        K,
        a.stride(0),
        a.stride(1),
        b.stride(0),
        b.stride(1),
        c.stride(0),
        c.stride(1),
    )
    return c


# %%
# Unit Test
# -----------
#
# We can test our custom Linear against torch Linear

torch.manual_seed(0)
a = torch.randn((256, 128), device="cuda", dtype=torch.float32)
b = torch.randn((512, 128), device="cuda", dtype=torch.float32)
triton_output = linear(a, b)
torch_output = torch.nn.functional.linear(a, b)
print(f"triton_output={triton_output}")
print(f"torch_output={torch_output}")
if triton.testing.allclose(triton_output, torch_output):
    print("✅ Triton and Torch match")
else:
    print("❌ Triton and Torch differ")
Outputs on 1650Ti (Passed):
triton_output=tensor([[ 11.2542,   0.4675,  -8.1082,  ...,  13.9425,  -4.0245,  20.8765],
        [ 21.8792, -15.1901,  13.6344,  ...,  19.8364,  14.2783,  24.0153],
        [  1.5868, -11.9487,  -3.4485,  ...,  16.2666,  -4.1443, -11.2711],
        ...,
        [  0.1807, -13.5327,   2.1057,  ...,   5.6092,  16.9773,   6.4231],
        [ 10.8032,  -6.5762,   2.6077,  ...,   0.0420,   7.5391,   6.2915],
        [  7.9209,   8.9644, -14.0444,  ...,  -0.7635,   8.7893,   2.8634]],
       device='cuda:0')
torch_output=tensor([[ 11.2542,   0.4675,  -8.1082,  ...,  13.9425,  -4.0245,  20.8765],
        [ 21.8792, -15.1901,  13.6344,  ...,  19.8364,  14.2783,  24.0153],
        [  1.5868, -11.9487,  -3.4485,  ...,  16.2666,  -4.1443, -11.2711],
        ...,
        [  0.1807, -13.5327,   2.1057,  ...,   5.6092,  16.9773,   6.4231],
        [ 10.8032,  -6.5762,   2.6077,  ...,   0.0420,   7.5391,   6.2915],
        [  7.9209,   8.9644, -14.0444,  ...,  -0.7635,   8.7893,   2.8634]],
       device='cuda:0')
✅ Triton and Torch match

Outputs on A100 (Failed):
triton_output=tensor([[ 11.2513, -15.2858,  18.1877,  ..., -14.1791,   6.6081,   2.8859],
        [ 21.8570,  -4.1529,  19.5992,  ..., -14.1775, -11.1812,  18.0821],
        [  1.5818, -10.8298,   7.6850,  ...,   4.0299,   4.0591, -32.8007],
        ...,
        [ -0.0426,  19.5304, -11.0848,  ...,  -0.2982,  27.7024,  -8.0911],
        [-12.7682, -10.8570,  -7.1777,  ...,   0.0818,   1.2360,  -9.0717],
        [ -7.8679,   4.2744,  -4.0673,  ...,   5.8584,   9.2709,  10.0778]],
       device='cuda:0')
torch_output=tensor([[ 11.2542,   0.4675,  -8.1082,  ...,   4.9326,  -4.6819,  -2.3277],
        [ 21.8792, -15.1901,  13.6344,  ..., -11.2832,  -3.5730,   9.9869],
        [  1.5868, -11.9487,  -3.4485,  ...,  -4.7018,  -7.9433,  -7.3278],
        ...,
        [ -0.0400,  -2.2899,  10.6292,  ...,  -7.8206,   2.6337,   1.4667],
        [-12.7751,   0.1556,  11.1091,  ...,  19.4366,  -4.1139,   5.2731],
        [ -7.8825,  -6.2058,   4.9614,  ...,   9.5303,   5.1297,  -0.3788]],
       device='cuda:0')
❌ Triton and Torch differ

Closing comment (2023-06-26T17:00:31Z): What happens if you set allow_tf32=False?  Might be related to #1671 (just a random thought as it works on 1650 TI).