diff --git a/aiter/ops/triton/_triton_kernels/gated_delta_rule/prefill/chunk_o.py b/aiter/ops/triton/_triton_kernels/gated_delta_rule/prefill/chunk_o.py index c7ceac4769..9d8095dfa4 100644 --- a/aiter/ops/triton/_triton_kernels/gated_delta_rule/prefill/chunk_o.py +++ b/aiter/ops/triton/_triton_kernels/gated_delta_rule/prefill/chunk_o.py @@ -29,6 +29,44 @@ NUM_WARPS = [2, 4] if IS_NVIDIA_HOPPER else [2, 4, 8] +# tl.make_block_ptr was removed in Triton 3.8 ("Block pointers have been removed +# in favor of the tensor descriptor API"), so the block loads/stores below are +# expressed as plain pointer arithmetic with explicit bounds masks. Both strides +# are passed so the transposed views (stride (1, H * K)) map over unchanged. +@triton.jit +def _bp_ld1d(base, N, stride, off, BL: tl.constexpr): + o = off + tl.arange(0, BL) + return tl.load(base + o * stride, mask=o < N, other=0.0) + + +@triton.jit +def _bp_st1d(base, N, stride, off, val, BL: tl.constexpr): + o = off + tl.arange(0, BL) + tl.store(base + o * stride, val, mask=o < N) + + +@triton.jit +def _bp_ld2d(base, R, C, rs, cs, r0, c0, BR: tl.constexpr, BC: tl.constexpr): + r = r0 + tl.arange(0, BR) + c = c0 + tl.arange(0, BC) + return tl.load( + base + r[:, None] * rs + c[None, :] * cs, + mask=(r < R)[:, None] & (c < C)[None, :], + other=0.0, + ) + + +@triton.jit +def _bp_st2d(base, R, C, rs, cs, r0, c0, val, BR: tl.constexpr, BC: tl.constexpr): + r = r0 + tl.arange(0, BR) + c = c0 + tl.arange(0, BC) + tl.store( + base + r[:, None] * rs + c[None, :] * cs, + val, + mask=(r < R)[:, None] & (c < C)[None, :], + ) + + @triton.heuristics( { "USE_G": lambda args: args["g"] is not None, @@ -99,21 +137,12 @@ def chunk_fwd_kernel_o( b_A = tl.zeros([BT, BT], dtype=tl.float32) for i_k in range(tl.cdiv(K, BK)): - p_q = tl.make_block_ptr( - q, (T, K), (H * K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0) - ) - p_k = tl.make_block_ptr( - k, (K, T), (1, H * K), (i_k * BK, i_t * BT), (BK, BT), (0, 1) - ) - p_h = tl.make_block_ptr( - h, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0) - ) # [BT, BK] - b_q = tl.load(p_q, boundary_check=(0, 1)) + b_q = _bp_ld2d(q, T, K, H * K, 1, i_t * BT, i_k * BK, BT, BK) # [BK, BT] - b_k = tl.load(p_k, boundary_check=(0, 1)) + b_k = _bp_ld2d(k, K, T, 1, H * K, i_k * BK, i_t * BT, BK, BT) # [BK, BV] - b_h = tl.load(p_h, boundary_check=(0, 1)) + b_h = _bp_ld2d(h, K, V, V, 1, i_k * BK, i_v * BV, BK, BV) # [BT, BK] @ [BK, BV] -> [BT, BV] b_o = tl.dot(b_q, b_h, acc=b_o) @@ -122,8 +151,7 @@ def chunk_fwd_kernel_o( if USE_G: g += bos * H + i_h - p_g = tl.make_block_ptr(g, (T,), (H,), (i_t * BT,), (BT,), (0,)) - b_g = tl.load(p_g, boundary_check=(0,)) + b_g = _bp_ld1d(g, T, H, i_t * BT, BT) b_o = b_o * exp(b_g)[:, None] b_A = b_A * exp(b_g[:, None] - b_g[None, :]) @@ -138,18 +166,22 @@ def chunk_fwd_kernel_o( m_A = (o_t[:, None] >= o_t[None, :]) & (m_t[:, None] & m_t) b_A = tl.where(m_A, b_A, 0) - p_v = tl.make_block_ptr( - v, (T, V), (H * V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0) - ) - p_o = tl.make_block_ptr( - o, (T, V), (H * V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0) - ) - - b_v = tl.load(p_v, boundary_check=(0, 1)) + b_v = _bp_ld2d(v, T, V, H * V, 1, i_t * BT, i_v * BV, BT, BV) # to fix mma -> mma layout conversion # already solved by triton v3.2 or higher b_o = b_o * scale + tl.dot(b_A.to(b_v.dtype), b_v) * scale - tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1)) + _bp_st2d( + o, + T, + V, + H * V, + 1, + i_t * BT, + i_v * BV, + b_o.to(o.dtype.element_ty), + BT, + BV, + ) @triton.heuristics( @@ -249,24 +281,12 @@ def chunk_bwd_kernel_dqkwg( b_dw = tl.zeros([BT, BK], dtype=tl.float32) if USE_DW else None for i_v in range(tl.cdiv(V, BV)): - p_v = tl.make_block_ptr( - v, (T, V), (H * V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0) - ) - p_do = tl.make_block_ptr( - do, (T, V), (H * V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0) - ) - p_h = tl.make_block_ptr( - h, (V, K), (1, V), (i_v * BV, i_k * BK), (BV, BK), (0, 1) - ) - p_dh = tl.make_block_ptr( - dh, (V, K), (1, V), (i_v * BV, i_k * BK), (BV, BK), (0, 1) - ) # [BT, BV] - b_v = tl.load(p_v, boundary_check=(0, 1)) - b_do = tl.load(p_do, boundary_check=(0, 1)) + b_v = _bp_ld2d(v, T, V, H * V, 1, i_t * BT, i_v * BV, BT, BV) + b_do = _bp_ld2d(do, T, V, H * V, 1, i_t * BT, i_v * BV, BT, BV) # [BV, BK] - b_h = tl.load(p_h, boundary_check=(0, 1)) - b_dh = tl.load(p_dh, boundary_check=(0, 1)) + b_h = _bp_ld2d(h, V, K, 1, V, i_v * BV, i_k * BK, BV, BK) + b_dh = _bp_ld2d(dh, V, K, 1, V, i_v * BV, i_k * BK, BV, BK) if USE_G: b_dg_last += tl.sum(b_h * b_dh) # [BT, BV] @ [BV, BT] -> [BT, BT] @@ -276,34 +296,26 @@ def chunk_bwd_kernel_dqkwg( # [BT, BV] @ [BV, BK] -> [BT, BK] b_dk = tl.dot(b_v, b_dh.to(b_v.dtype), acc=b_dk) if USE_DW: - p_dv = tl.make_block_ptr( - dv, (T, V), (H * V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0) - ) - b_dv = tl.load(p_dv, boundary_check=(0, 1)) + b_dv = _bp_ld2d(dv, T, V, H * V, 1, i_t * BT, i_v * BV, BT, BV) b_dw = tl.dot(b_dv.to(b_v.dtype), b_h.to(b_v.dtype), acc=b_dw) if USE_DW: - p_dw = tl.make_block_ptr( - dw, (T, K), (H * K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0) + _bp_st2d( + dw, + T, + K, + H * K, + 1, + i_t * BT, + i_k * BK, + -b_dw.to(dw.dtype.element_ty), + BT, + BK, ) - tl.store(p_dw, -b_dw.to(p_dw.dtype.element_ty), boundary_check=(0, 1)) tl.debug_barrier() - p_q = tl.make_block_ptr( - q, (T, K), (H * K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0) - ) - p_k = tl.make_block_ptr( - k, (T, K), (H * K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0) - ) - b_q = tl.load(p_q, boundary_check=(0, 1)) - b_k = tl.load(p_k, boundary_check=(0, 1)) - - p_dq = tl.make_block_ptr( - dq, (T, K), (H * K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0) - ) - p_dk = tl.make_block_ptr( - dk, (T, K), (H * K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0) - ) + b_q = _bp_ld2d(q, T, K, H * K, 1, i_t * BT, i_k * BK, BT, BK) + b_k = _bp_ld2d(k, T, K, H * K, 1, i_t * BT, i_k * BK, BT, BK) o_t = i_t * BT + tl.arange(0, BT) m_t = o_t < T @@ -312,8 +324,7 @@ def chunk_bwd_kernel_dqkwg( b_dg = tl.zeros([BT], dtype=tl.float32) g += bos * H + i_h dg += bos * H + i_h - p_g = tl.make_block_ptr(g, (T,), (H,), (i_t * BT,), (BT,), (0,)) - b_g = tl.load(p_g, boundary_check=(0,)) + b_g = _bp_ld1d(g, T, H, i_t * BT, BT) b_g_last = tl.load(g + (min(i_t * BT + BT, T) - 1) * H) b_dg_last *= exp(b_g_last) @@ -333,13 +344,16 @@ def chunk_bwd_kernel_dqkwg( # [BT, BK] b_dq = tl.dot(b_ds, b_k, acc=b_dq) b_dk = tl.dot(tl.trans(b_ds), b_q, acc=b_dk) - p_dg = tl.make_block_ptr(dg, (T,), (H,), (i_t * BT,), (BT,), (0,)) # (SY 09/21) revcumsum in a separate kernel due to strange triton compiler issue # b_dg = tl.dot(tl.where(o_t[:, None] <= o_t[None, :], 1., 0.), b_dg, allow_tf32=False) + b_dg_last) b_dg = tl.where(o_t < min(i_t * BT + BT, T) - 1, b_dg, b_dg + b_dg_last) - tl.store(p_dq, b_dq.to(p_dq.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_dg, b_dg.to(p_dg.dtype.element_ty), boundary_check=(0,)) + _bp_st2d( + dq, T, K, H * K, 1, i_t * BT, i_k * BK, b_dq.to(dq.dtype.element_ty), BT, BK + ) + _bp_st2d( + dk, T, K, H * K, 1, i_t * BT, i_k * BK, b_dk.to(dk.dtype.element_ty), BT, BK + ) + _bp_st1d(dg, T, H, i_t * BT, b_dg.to(dg.dtype.element_ty), BT) elif USE_G_GAMMA: b_dq = b_dq * exp(b_g)[:, None] * scale @@ -349,8 +363,12 @@ def chunk_bwd_kernel_dqkwg( # [BT, BK] b_dq = tl.dot(b_ds, b_k, acc=b_dq) b_dk = tl.dot(tl.trans(b_ds), b_q, acc=b_dk) - tl.store(p_dq, b_dq.to(p_dq.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1)) + _bp_st2d( + dq, T, K, H * K, 1, i_t * BT, i_k * BK, b_dq.to(dq.dtype.element_ty), BT, BK + ) + _bp_st2d( + dk, T, K, H * K, 1, i_t * BT, i_k * BK, b_dk.to(dk.dtype.element_ty), BT, BK + ) else: b_ds = tl.where(m_A, b_ds, 0) @@ -358,8 +376,12 @@ def chunk_bwd_kernel_dqkwg( b_dq = tl.dot(b_ds, b_k, acc=b_dq) b_dk += tl.dot(tl.trans(b_ds), b_q) * scale b_dq *= scale - tl.store(p_dq, b_dq.to(p_dq.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1)) + _bp_st2d( + dq, T, K, H * K, 1, i_t * BT, i_k * BK, b_dq.to(dq.dtype.element_ty), BT, BK + ) + _bp_st2d( + dk, T, K, H * K, 1, i_t * BT, i_k * BK, b_dk.to(dk.dtype.element_ty), BT, BK + ) @triton.heuristics( @@ -431,27 +453,17 @@ def chunk_bwd_kernel_dv( b_A = tl.zeros([BT, BT], dtype=tl.float32) for i_k in range(tl.cdiv(K, BK)): - p_k = tl.make_block_ptr( - k, (T, K), (H * K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0) - ) - p_q = tl.make_block_ptr( - q, (K, T), (1, H * K), (i_k * BK, i_t * BT), (BK, BT), (0, 1) - ) - b_q = tl.load(p_q, boundary_check=(0, 1)) - b_k = tl.load(p_k, boundary_check=(0, 1)) + b_q = _bp_ld2d(q, K, T, 1, H * K, i_k * BK, i_t * BT, BK, BT) + b_k = _bp_ld2d(k, T, K, H * K, 1, i_t * BT, i_k * BK, BT, BK) b_A = tl.dot(b_k, b_q, acc=b_A) - p_dh = tl.make_block_ptr( - dh, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0) - ) - b_dh = tl.load(p_dh, boundary_check=(0, 1)) + b_dh = _bp_ld2d(dh, K, V, V, 1, i_k * BK, i_v * BV, BK, BV) b_dv = tl.dot(b_k, b_dh.to(b_k.dtype), acc=b_dv) o_t = i_t * BT + tl.arange(0, BT) m_t = o_t < T if USE_G: g += bos * H + i_h - p_g = tl.make_block_ptr(g, (T,), (H,), (i_t * BT,), (BT,), (0,)) - b_g = tl.load(p_g, boundary_check=(0,)) + b_g = _bp_ld1d(g, T, H, i_t * BT, BT) b_g_last = tl.load(g + (min(i_t * BT + BT, T) - 1) * H) if USE_G_GAMMA: b_gamma = tl.load(g_gamma + i_h) @@ -466,15 +478,11 @@ def chunk_bwd_kernel_dv( b_dv *= tl.where(m_t, exp(-b_g + b_g_last), 0)[:, None] else: b_A = tl.where(m_A, b_A * scale, 0).to(do.dtype.element_ty) - p_do = tl.make_block_ptr( - do, (T, V), (H * V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0) - ) - p_dv = tl.make_block_ptr( - dv, (T, V), (H * V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0) - ) - b_do = tl.load(p_do, boundary_check=(0, 1)) + b_do = _bp_ld2d(do, T, V, H * V, 1, i_t * BT, i_v * BV, BT, BV) b_dv = tl.dot(b_A.to(b_do.dtype), b_do, acc=b_dv) - tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), boundary_check=(0, 1)) + _bp_st2d( + dv, T, V, H * V, 1, i_t * BT, i_v * BV, b_dv.to(dv.dtype.element_ty), BT, BV + ) @triton.heuristics( @@ -540,35 +548,19 @@ def chunk_bwd_kernel_dv_local( dv += (bos * H + i_h) * V if USE_A: - p_A = tl.make_block_ptr( - A + (bos * H + i_h) * BT, - (BT, T), - (1, H * BT), - (0, i_t * BT), - (BT, BT), - (0, 1), - ) - b_A = tl.load(p_A, boundary_check=(0, 1)) + b_A = _bp_ld2d(A + (bos * H + i_h) * BT, BT, T, 1, H * BT, 0, i_t * BT, BT, BT) else: if USE_G: g += bos * H + i_h - p_g = tl.make_block_ptr(g, (T,), (H,), (i_t * BT,), (BT,), (0,)) - b_g = tl.load(p_g, boundary_check=(0,)) + b_g = _bp_ld1d(g, T, H, i_t * BT, BT) if USE_G_GAMMA: b_gamma = tl.load(g_gamma + i_h) b_g = b_gamma * (tl.arange(0, BT) + 1) b_A = tl.zeros([BT, BT], dtype=tl.float32) for i_k in range(tl.cdiv(K, BK)): - p_k = tl.make_block_ptr( - k, (T, K), (H * K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0) - ) - p_q = tl.make_block_ptr( - q, (K, T), (1, H * K), (i_k * BK, i_t * BT), (BK, BT), (0, 1) - ) - - b_k = tl.load(p_k, boundary_check=(0, 1)) - b_q = tl.load(p_q, boundary_check=(0, 1)) + b_k = _bp_ld2d(k, T, K, H * K, 1, i_t * BT, i_k * BK, BT, BK) + b_q = _bp_ld2d(q, K, T, 1, H * K, i_k * BK, i_t * BT, BK, BT) b_A += tl.dot(b_k, b_q) * scale if USE_G or USE_G_GAMMA: b_A *= exp(b_g[None, :] - b_g[:, None]) @@ -579,15 +571,11 @@ def chunk_bwd_kernel_dv_local( b_A = tl.where(m_A, b_A, 0).to(do.dtype.element_ty) for i_v in range(tl.cdiv(V, BV)): - p_do = tl.make_block_ptr( - do, (T, V), (H * V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0) - ) - p_dv = tl.make_block_ptr( - dv, (T, V), (H * V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0) - ) - b_do = tl.load(p_do, boundary_check=(0, 1)) + b_do = _bp_ld2d(do, T, V, H * V, 1, i_t * BT, i_v * BV, BT, BV) b_dv = tl.dot(b_A.to(b_do.dtype), b_do) - tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), boundary_check=(0, 1)) + _bp_st2d( + dv, T, V, H * V, 1, i_t * BT, i_v * BV, b_dv.to(dv.dtype.element_ty), BT, BV + ) def chunk_fwd_o( @@ -713,26 +701,16 @@ def chunk_fwd_kernel_o_opt( b_A = tl.zeros([BT, BT], dtype=tl.float32) for i_k in range(tl.cdiv(K, BK)): - p_q = tl.make_block_ptr( - q, (T, K), (Hg * K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0) - ) - p_k = tl.make_block_ptr( - k, (K, T), (1, Hg * K), (i_k * BK, i_t * BT), (BK, BT), (0, 1) - ) - p_h = tl.make_block_ptr( - h, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0) - ) - b_q = tl.load(p_q, boundary_check=(0, 1)) - b_k = tl.load(p_k, boundary_check=(0, 1)) - b_h = tl.load(p_h, boundary_check=(0, 1)) + b_q = _bp_ld2d(q, T, K, Hg * K, 1, i_t * BT, i_k * BK, BT, BK) + b_k = _bp_ld2d(k, K, T, 1, Hg * K, i_k * BK, i_t * BT, BK, BT) + b_h = _bp_ld2d(h, K, V, V, 1, i_k * BK, i_v * BV, BK, BV) b_o = tl.dot(b_q, b_h, acc=b_o) b_A = tl.dot(b_q, b_k, acc=b_A) if USE_G: g += bos * H + i_h - p_g = tl.make_block_ptr(g, (T,), (H,), (i_t * BT,), (BT,), (0,)) - b_g = tl.load(p_g, boundary_check=(0,)) + b_g = _bp_ld1d(g, T, H, i_t * BT, BT) b_o = b_o * exp(b_g)[:, None] b_A = b_A * exp(b_g[:, None] - b_g[None, :]) @@ -741,14 +719,10 @@ def chunk_fwd_kernel_o_opt( m_A = (o_t[:, None] >= o_t[None, :]) & (m_t[:, None] & m_t) b_A = tl.where(m_A, b_A, 0) - p_v = tl.make_block_ptr(v, (T, V), (V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_o = tl.make_block_ptr( - o, (T, V), (H * V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0) - ) - b_v = tl.load(p_v, boundary_check=(0, 1)) + b_v = _bp_ld2d(v, T, V, V, 1, i_t * BT, i_v * BV, BT, BV) b_o = b_o * scale + tl.dot(b_A.to(b_v.dtype), b_v) * scale - tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1)) + _bp_st2d(o, T, V, H * V, 1, i_t * BT, i_v * BV, b_o.to(o.dtype.element_ty), BT, BV) def chunk_fwd_o_opt( @@ -909,18 +883,9 @@ def chunk_fwd_kernel_o_opt_vk( b_A = tl.zeros([BT, BT], dtype=tl.float32) for i_k in range(tl.cdiv(K, BK)): - p_q = tl.make_block_ptr( - q, (T, K), (Hg * K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0) - ) - p_k = tl.make_block_ptr( - k, (K, T), (1, Hg * K), (i_k * BK, i_t * BT), (BK, BT), (0, 1) - ) - p_h = tl.make_block_ptr( - h, (V, K), (K, 1), (i_v * BV, i_k * BK), (BV, BK), (1, 0) - ) - b_q = tl.load(p_q, boundary_check=(0, 1)) - b_k = tl.load(p_k, boundary_check=(0, 1)) - b_h = tl.load(p_h, boundary_check=(0, 1)) + b_q = _bp_ld2d(q, T, K, Hg * K, 1, i_t * BT, i_k * BK, BT, BK) + b_k = _bp_ld2d(k, K, T, 1, Hg * K, i_k * BK, i_t * BT, BK, BT) + b_h = _bp_ld2d(h, V, K, K, 1, i_v * BV, i_k * BK, BV, BK) # K6 requires matching tl.dot operand dtypes. This handles every # supported input/snapshot combination (FP16/BF16 input with @@ -930,8 +895,7 @@ def chunk_fwd_kernel_o_opt_vk( b_A = tl.dot(b_q, b_k, acc=b_A) if USE_G: - p_g = tl.make_block_ptr(g, (T,), (1,), (i_t * BT,), (BT,), (0,)) - b_g = tl.load(p_g, boundary_check=(0,)) + b_g = _bp_ld1d(g, T, 1, i_t * BT, BT) if USE_EXP2: b_o = b_o * tl.math.exp2(b_g)[:, None] b_A = b_A * tl.math.exp2(b_g[:, None] - b_g[None, :]) @@ -944,14 +908,10 @@ def chunk_fwd_kernel_o_opt_vk( m_A = (o_t[:, None] >= o_t[None, :]) & (m_t[:, None] & m_t) b_A = tl.where(m_A, b_A, 0) - p_v = tl.make_block_ptr(v, (T, V), (V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_o = tl.make_block_ptr( - o, (T, V), (H * V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0) - ) - b_v = tl.load(p_v, boundary_check=(0, 1)) + b_v = _bp_ld2d(v, T, V, V, 1, i_t * BT, i_v * BV, BT, BV) b_o = b_o * scale + tl.dot(b_A.to(b_v.dtype), b_v) * scale - tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1)) + _bp_st2d(o, T, V, H * V, 1, i_t * BT, i_v * BV, b_o.to(o.dtype.element_ty), BT, BV) def chunk_fwd_o_opt_vk( diff --git a/aiter/ops/triton/_triton_kernels/gated_delta_rule/prefill/fused_cumsum_kkt.py b/aiter/ops/triton/_triton_kernels/gated_delta_rule/prefill/fused_cumsum_kkt.py index 18805ef027..f7f3c1824e 100644 --- a/aiter/ops/triton/_triton_kernels/gated_delta_rule/prefill/fused_cumsum_kkt.py +++ b/aiter/ops/triton/_triton_kernels/gated_delta_rule/prefill/fused_cumsum_kkt.py @@ -54,45 +54,40 @@ def _fused_cumsum_kkt_kernel( o_t = tl.arange(0, BT) - p_g = tl.make_block_ptr( - g_ptr + bos * H + i_h, (T_seq,), (H,), (i_t * BT,), (BT,), (0,) - ) - b_g = tl.load(p_g, boundary_check=(0,)).to(tl.float32) + # Plain pointer arithmetic rather than tl.make_block_ptr, which is not + # available in every Triton build we target. + o_abs = i_t * BT + o_t + m_t = o_abs < T_seq + o_k = tl.arange(0, K) + + b_g = tl.load(g_ptr + bos * H + i_h + o_abs * H, mask=m_t, other=0.0).to(tl.float32) b_g_cumsum = tl.cumsum(b_g, axis=0) - p_g_out = tl.make_block_ptr( - g_cumsum_ptr + bos * H + i_h, (T_seq,), (H,), (i_t * BT,), (BT,), (0,) - ) - tl.store(p_g_out, b_g_cumsum.to(p_g_out.dtype.element_ty), boundary_check=(0,)) + g_out_ptrs = g_cumsum_ptr + bos * H + i_h + o_abs * H + tl.store(g_out_ptrs, b_g_cumsum.to(g_cumsum_ptr.dtype.element_ty), mask=m_t) - p_beta = tl.make_block_ptr( - beta_ptr + bos * H + i_h, (T_seq,), (H,), (i_t * BT,), (BT,), (0,) - ) - b_beta = tl.load(p_beta, boundary_check=(0,)).to(tl.float32) - - p_k = tl.make_block_ptr( - k_ptr + (bos * Hg + i_h // (H // Hg)) * K, - (T_seq, K), - (Hg * K, 1), - (i_t * BT, 0), - (BT, K), - (1, 0), + b_beta = tl.load(beta_ptr + bos * H + i_h + o_abs * H, mask=m_t, other=0.0).to( + tl.float32 ) - b_k = tl.load(p_k, boundary_check=(0, 1)).to(tl.float32) + + b_k = tl.load( + k_ptr + + (bos * Hg + i_h // (H // Hg)) * K + + o_abs[:, None] * (Hg * K) + + o_k[None, :], + mask=m_t[:, None], + other=0.0, + ).to(tl.float32) b_A = tl.dot(b_k, tl.trans(b_k)) b_g_diff = b_g_cumsum[:, None] - b_g_cumsum[None, :] b_A = b_A * safe_exp(b_g_diff) * b_beta[:, None] b_A = tl.where(o_t[:, None] > o_t[None, :], b_A, 0.0) - p_A = tl.make_block_ptr( - A_ptr + (bos * H + i_h) * BT, - (T_seq, BT), - (BT * H, 1), - (i_t * BT, 0), - (BT, BT), - (1, 0), + tl.store( + A_ptr + (bos * H + i_h) * BT + o_abs[:, None] * (BT * H) + o_t[None, :], + b_A.to(A_ptr.dtype.element_ty), + mask=m_t[:, None], ) - tl.store(p_A, b_A.to(A_ptr.dtype.element_ty), boundary_check=(0, 1)) def fused_cumsum_kkt( @@ -215,8 +210,9 @@ def fused_chunk_local_cumsum_scaled_dot_kkt_fwd_kernel( o_t = i_t * BT + tl.arange(0, BT) m_t = o_t < T - p_g = tl.make_block_ptr(g + bos * H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) - b_g = tl.load(p_g, boundary_check=(0,)).to(tl.float32) + # Plain pointer arithmetic rather than tl.make_block_ptr, which is not + # available in every Triton build we target. + b_g = tl.load(g + bos * H + i_h + o_t * H, mask=m_t, other=0.0).to(tl.float32) b_g_cumsum = tl.cumsum(b_g, axis=0) # Store g_cumsum in log2 space when downstream kernels consume it with exp2: # exp2(x * RCP_LN2) == exp(x), keeping results identical. The scale arrives @@ -233,25 +229,24 @@ def fused_chunk_local_cumsum_scaled_dot_kkt_fwd_kernel( g_out_base = g_cumsum_out + i_h * T_flat + bos else: g_out_base = g_cumsum_out + (i_b * H + i_h) * T_flat - p_go = tl.make_block_ptr(g_out_base, (T,), (1,), (i_t * BT,), (BT,), (0,)) - tl.store(p_go, b_g_cumsum.to(p_go.dtype.element_ty), boundary_check=(0,)) - - p_beta = tl.make_block_ptr( - beta + bos * H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,) + tl.store( + g_out_base + o_t, + b_g_cumsum.to(g_out_base.dtype.element_ty), + mask=m_t, ) - b_beta = tl.load(p_beta, boundary_check=(0,)) + + b_beta = tl.load(beta + bos * H + i_h + o_t * H, mask=m_t, other=0.0) b_A = tl.zeros([BT, BT], dtype=tl.float32) + k_base = k + (bos * Hg + i_h // (H // Hg)) * K for i_k in range(tl.cdiv(K, BK)): - p_k = tl.make_block_ptr( - k + (bos * Hg + i_h // (H // Hg)) * K, - (T, K), - (Hg * K, 1), - (i_t * BT, i_k * BK), - (BT, BK), - (1, 0), + o_k = i_k * BK + tl.arange(0, BK) + m_k = o_k < K + b_k = tl.load( + k_base + o_t[:, None] * (Hg * K) + o_k[None, :], + mask=m_t[:, None] & m_k[None, :], + other=0.0, ) - b_k = tl.load(p_k, boundary_check=(0, 1)) b_kb = b_k * b_beta[:, None] b_A = tl.dot(b_kb.to(b_k.dtype), tl.trans(b_k), acc=b_A) @@ -260,15 +255,12 @@ def fused_chunk_local_cumsum_scaled_dot_kkt_fwd_kernel( b_gate = tl.math.exp2(b_g_diff) if USE_EXP2 else exp(b_g_diff) b_A = tl.where(m_A, b_A * b_gate, 0.0) - p_A = tl.make_block_ptr( - A_out + (bos * H + i_h) * BT, - (T, BT), - (BT * H, 1), - (i_t * BT, 0), - (BT, BT), - (1, 0), + o_bt = tl.arange(0, BT) + tl.store( + A_out + (bos * H + i_h) * BT + o_t[:, None] * (BT * H) + o_bt[None, :], + b_A.to(A_out.dtype.element_ty), + mask=m_t[:, None], ) - tl.store(p_A, b_A.to(A_out.dtype.element_ty), boundary_check=(0, 1)) def fused_chunk_local_cumsum_scaled_dot_kkt_fwd( diff --git a/aiter/ops/triton/_triton_kernels/gated_delta_rule/prefill/fused_solve_tril_recompute.py b/aiter/ops/triton/_triton_kernels/gated_delta_rule/prefill/fused_solve_tril_recompute.py index 59be1068f6..b9cf3c47e0 100644 --- a/aiter/ops/triton/_triton_kernels/gated_delta_rule/prefill/fused_solve_tril_recompute.py +++ b/aiter/ops/triton/_triton_kernels/gated_delta_rule/prefill/fused_solve_tril_recompute.py @@ -84,6 +84,41 @@ def _should_use_split_path(execution_mode: str, nt: int, is_varlen: bool) -> boo ) +# --- Block-pointer compatibility helpers ------------------------------------- +# tl.make_block_ptr was removed in Triton 3.8 ("Block pointers have been removed +# in favor of the tensor descriptor API"). These helpers reproduce the exact +# semantics we relied on (masked load/store of a 1-D or 2-D tile) with plain +# pointer arithmetic, so the kernels build on every Triton version we target. +@triton.jit +def _bp_ld1d(base, T, stride, off, BL: tl.constexpr): + o = off + tl.arange(0, BL) + return tl.load(base + o * stride, mask=o < T, other=0.0) + + +@triton.jit +def _bp_ld2d(base, T, D, row_stride, row0, col0, BR: tl.constexpr, BD: tl.constexpr): + r = row0 + tl.arange(0, BR) + c = col0 + tl.arange(0, BD) + return tl.load( + base + r[:, None] * row_stride + c[None, :], + mask=(r < T)[:, None] & (c < D)[None, :], + other=0.0, + ) + + +@triton.jit +def _bp_st2d( + base, T, D, row_stride, row0, col0, val, BR: tl.constexpr, BD: tl.constexpr +): + r = row0 + tl.arange(0, BR) + c = col0 + tl.arange(0, BD) + tl.store( + base + r[:, None] * row_stride + c[None, :], + val, + mask=(r < T)[:, None] & (c < D)[None, :], + ) + + @triton.heuristics({"IS_VARLEN": lambda args: args["cu_seqlens"] is not None}) @triton.autotune( configs=gated_delta_rule_autotune_configs( @@ -155,22 +190,24 @@ def fused_solve_tril_recompute_w_u_kernel( m_id = o_i[:, None] == o_i[None, :] A_base = A_raw + (bos * H + i_h) * BT - p11 = tl.make_block_ptr( - A_base, (T, BT), (H * BT, 1), (i_t * BT, 0), (16, 16), (1, 0) + b11 = -tl.where( + m_lo, _bp_ld2d(A_base, T, BT, H * BT, i_t * BT, 0, 16, 16).to(tl.float32), 0 ) - p22 = tl.make_block_ptr( - A_base, (T, BT), (H * BT, 1), (i_t * BT + 16, 16), (16, 16), (1, 0) + b22 = -tl.where( + m_lo, + _bp_ld2d(A_base, T, BT, H * BT, i_t * BT + 16, 16, 16, 16).to(tl.float32), + 0, ) - p33 = tl.make_block_ptr( - A_base, (T, BT), (H * BT, 1), (i_t * BT + 32, 32), (16, 16), (1, 0) + b33 = -tl.where( + m_lo, + _bp_ld2d(A_base, T, BT, H * BT, i_t * BT + 32, 32, 16, 16).to(tl.float32), + 0, ) - p44 = tl.make_block_ptr( - A_base, (T, BT), (H * BT, 1), (i_t * BT + 48, 48), (16, 16), (1, 0) + b44 = -tl.where( + m_lo, + _bp_ld2d(A_base, T, BT, H * BT, i_t * BT + 48, 48, 16, 16).to(tl.float32), + 0, ) - b11 = -tl.where(m_lo, tl.load(p11, boundary_check=(0, 1)).to(tl.float32), 0) - b22 = -tl.where(m_lo, tl.load(p22, boundary_check=(0, 1)).to(tl.float32), 0) - b33 = -tl.where(m_lo, tl.load(p33, boundary_check=(0, 1)).to(tl.float32), 0) - b44 = -tl.where(m_lo, tl.load(p44, boundary_check=(0, 1)).to(tl.float32), 0) for i in range(2, min(16, T - i_t * BT)): r = -tl.load(A_base + (i_t * BT + i) * H * BT + o_i) @@ -193,42 +230,12 @@ def fused_solve_tril_recompute_w_u_kernel( b33 += m_id b44 += m_id - rA21 = tl.load( - tl.make_block_ptr( - A_base, (T, BT), (H * BT, 1), (i_t * BT + 16, 0), (16, 16), (1, 0) - ), - boundary_check=(0, 1), - ).to(tl.float32) - rA31 = tl.load( - tl.make_block_ptr( - A_base, (T, BT), (H * BT, 1), (i_t * BT + 32, 0), (16, 16), (1, 0) - ), - boundary_check=(0, 1), - ).to(tl.float32) - rA32 = tl.load( - tl.make_block_ptr( - A_base, (T, BT), (H * BT, 1), (i_t * BT + 32, 16), (16, 16), (1, 0) - ), - boundary_check=(0, 1), - ).to(tl.float32) - rA41 = tl.load( - tl.make_block_ptr( - A_base, (T, BT), (H * BT, 1), (i_t * BT + 48, 0), (16, 16), (1, 0) - ), - boundary_check=(0, 1), - ).to(tl.float32) - rA42 = tl.load( - tl.make_block_ptr( - A_base, (T, BT), (H * BT, 1), (i_t * BT + 48, 16), (16, 16), (1, 0) - ), - boundary_check=(0, 1), - ).to(tl.float32) - rA43 = tl.load( - tl.make_block_ptr( - A_base, (T, BT), (H * BT, 1), (i_t * BT + 48, 32), (16, 16), (1, 0) - ), - boundary_check=(0, 1), - ).to(tl.float32) + rA21 = _bp_ld2d(A_base, T, BT, H * BT, i_t * BT + 16, 0, 16, 16).to(tl.float32) + rA31 = _bp_ld2d(A_base, T, BT, H * BT, i_t * BT + 32, 0, 16, 16).to(tl.float32) + rA32 = _bp_ld2d(A_base, T, BT, H * BT, i_t * BT + 32, 16, 16, 16).to(tl.float32) + rA41 = _bp_ld2d(A_base, T, BT, H * BT, i_t * BT + 48, 0, 16, 16).to(tl.float32) + rA42 = _bp_ld2d(A_base, T, BT, H * BT, i_t * BT + 48, 16, 16, 16).to(tl.float32) + rA43 = _bp_ld2d(A_base, T, BT, H * BT, i_t * BT + 48, 32, 16, 16).to(tl.float32) b21 = -tl.dot( tl.dot(b22, rA21, input_precision=DOT_PRECISION), @@ -285,29 +292,21 @@ def fused_solve_tril_recompute_w_u_kernel( else: g_base = g + (i_b * H + i_h) * T_flat - p_b0 = tl.make_block_ptr(beta_base, (T,), (H,), (i_t * BT,), (16,), (0,)) - p_b1 = tl.make_block_ptr(beta_base, (T,), (H,), (i_t * BT + 16,), (16,), (0,)) - p_b2 = tl.make_block_ptr(beta_base, (T,), (H,), (i_t * BT + 32,), (16,), (0,)) - p_b3 = tl.make_block_ptr(beta_base, (T,), (H,), (i_t * BT + 48,), (16,), (0,)) - bb0 = tl.load(p_b0, boundary_check=(0,)) - bb1 = tl.load(p_b1, boundary_check=(0,)) - bb2 = tl.load(p_b2, boundary_check=(0,)) - bb3 = tl.load(p_b3, boundary_check=(0,)) - - p_g0 = tl.make_block_ptr(g_base, (T,), (1,), (i_t * BT,), (16,), (0,)) - p_g1 = tl.make_block_ptr(g_base, (T,), (1,), (i_t * BT + 16,), (16,), (0,)) - p_g2 = tl.make_block_ptr(g_base, (T,), (1,), (i_t * BT + 32,), (16,), (0,)) - p_g3 = tl.make_block_ptr(g_base, (T,), (1,), (i_t * BT + 48,), (16,), (0,)) + bb0 = _bp_ld1d(beta_base, T, H, i_t * BT, 16) + bb1 = _bp_ld1d(beta_base, T, H, i_t * BT + 16, 16) + bb2 = _bp_ld1d(beta_base, T, H, i_t * BT + 32, 16) + bb3 = _bp_ld1d(beta_base, T, H, i_t * BT + 48, 16) + if USE_EXP2: - eg0 = tl.math.exp2(tl.load(p_g0, boundary_check=(0,))) - eg1 = tl.math.exp2(tl.load(p_g1, boundary_check=(0,))) - eg2 = tl.math.exp2(tl.load(p_g2, boundary_check=(0,))) - eg3 = tl.math.exp2(tl.load(p_g3, boundary_check=(0,))) + eg0 = tl.math.exp2(_bp_ld1d(g_base, T, 1, i_t * BT, 16)) + eg1 = tl.math.exp2(_bp_ld1d(g_base, T, 1, i_t * BT + 16, 16)) + eg2 = tl.math.exp2(_bp_ld1d(g_base, T, 1, i_t * BT + 32, 16)) + eg3 = tl.math.exp2(_bp_ld1d(g_base, T, 1, i_t * BT + 48, 16)) else: - eg0 = exp(tl.load(p_g0, boundary_check=(0,))) - eg1 = exp(tl.load(p_g1, boundary_check=(0,))) - eg2 = exp(tl.load(p_g2, boundary_check=(0,))) - eg3 = exp(tl.load(p_g3, boundary_check=(0,))) + eg0 = exp(_bp_ld1d(g_base, T, 1, i_t * BT, 16)) + eg1 = exp(_bp_ld1d(g_base, T, 1, i_t * BT + 16, 16)) + eg2 = exp(_bp_ld1d(g_base, T, 1, i_t * BT + 32, 16)) + eg3 = exp(_bp_ld1d(g_base, T, 1, i_t * BT + 48, 16)) v_base = v + (bos * H + i_h) * V if IS_VARLEN: @@ -316,22 +315,21 @@ def fused_solve_tril_recompute_w_u_kernel( u_base = u + (((i_b * H + i_h) * T_flat) * V) for i_v in range(tl.cdiv(V, BV)): - pv0 = tl.make_block_ptr( - v_base, (T, V), (H * V, 1), (i_t * BT, i_v * BV), (16, BV), (1, 0) - ) - pv1 = tl.make_block_ptr( - v_base, (T, V), (H * V, 1), (i_t * BT + 16, i_v * BV), (16, BV), (1, 0) - ) - pv2 = tl.make_block_ptr( - v_base, (T, V), (H * V, 1), (i_t * BT + 32, i_v * BV), (16, BV), (1, 0) - ) - pv3 = tl.make_block_ptr( - v_base, (T, V), (H * V, 1), (i_t * BT + 48, i_v * BV), (16, BV), (1, 0) - ) - vb0 = (tl.load(pv0, boundary_check=(0, 1)) * bb0[:, None]).to(lowp_dtype) - vb1 = (tl.load(pv1, boundary_check=(0, 1)) * bb1[:, None]).to(lowp_dtype) - vb2 = (tl.load(pv2, boundary_check=(0, 1)) * bb2[:, None]).to(lowp_dtype) - vb3 = (tl.load(pv3, boundary_check=(0, 1)) * bb3[:, None]).to(lowp_dtype) + vb0 = ( + _bp_ld2d(v_base, T, V, H * V, i_t * BT, i_v * BV, 16, BV) * bb0[:, None] + ).to(lowp_dtype) + vb1 = ( + _bp_ld2d(v_base, T, V, H * V, i_t * BT + 16, i_v * BV, 16, BV) + * bb1[:, None] + ).to(lowp_dtype) + vb2 = ( + _bp_ld2d(v_base, T, V, H * V, i_t * BT + 32, i_v * BV, 16, BV) + * bb2[:, None] + ).to(lowp_dtype) + vb3 = ( + _bp_ld2d(v_base, T, V, H * V, i_t * BT + 48, i_v * BV, 16, BV) + * bb3[:, None] + ).to(lowp_dtype) u0 = tl.dot(h11, vb0, allow_tf32=False) u1 = tl.dot(h21, vb0, allow_tf32=False) + tl.dot(h22, vb1, allow_tf32=False) @@ -347,22 +345,42 @@ def fused_solve_tril_recompute_w_u_kernel( + tl.dot(h44, vb3, allow_tf32=False) ) - pu0 = tl.make_block_ptr( - u_base, (T, V), (V, 1), (i_t * BT, i_v * BV), (16, BV), (1, 0) - ) - pu1 = tl.make_block_ptr( - u_base, (T, V), (V, 1), (i_t * BT + 16, i_v * BV), (16, BV), (1, 0) + _bp_st2d( + u_base, T, V, V, i_t * BT, i_v * BV, u0.to(u_base.dtype.element_ty), 16, BV ) - pu2 = tl.make_block_ptr( - u_base, (T, V), (V, 1), (i_t * BT + 32, i_v * BV), (16, BV), (1, 0) - ) - pu3 = tl.make_block_ptr( - u_base, (T, V), (V, 1), (i_t * BT + 48, i_v * BV), (16, BV), (1, 0) + _bp_st2d( + u_base, + T, + V, + V, + i_t * BT + 16, + i_v * BV, + u1.to(u_base.dtype.element_ty), + 16, + BV, + ) + _bp_st2d( + u_base, + T, + V, + V, + i_t * BT + 32, + i_v * BV, + u2.to(u_base.dtype.element_ty), + 16, + BV, + ) + _bp_st2d( + u_base, + T, + V, + V, + i_t * BT + 48, + i_v * BV, + u3.to(u_base.dtype.element_ty), + 16, + BV, ) - tl.store(pu0, u0.to(pu0.dtype.element_ty), boundary_check=(0, 1)) - tl.store(pu1, u1.to(pu1.dtype.element_ty), boundary_check=(0, 1)) - tl.store(pu2, u2.to(pu2.dtype.element_ty), boundary_check=(0, 1)) - tl.store(pu3, u3.to(pu3.dtype.element_ty), boundary_check=(0, 1)) k_base = k + (bos * Hg + i_h // (H // Hg)) * K if IS_VARLEN: @@ -371,52 +389,68 @@ def fused_solve_tril_recompute_w_u_kernel( w_base = w + (((i_b * H + i_h) * T_flat) * K) for i_k in range(tl.cdiv(K, BK)): - pk0 = tl.make_block_ptr( - k_base, (T, K), (Hg * K, 1), (i_t * BT, i_k * BK), (16, BK), (1, 0) - ) - pk1 = tl.make_block_ptr( - k_base, (T, K), (Hg * K, 1), (i_t * BT + 16, i_k * BK), (16, BK), (1, 0) - ) - pk2 = tl.make_block_ptr( - k_base, (T, K), (Hg * K, 1), (i_t * BT + 32, i_k * BK), (16, BK), (1, 0) - ) - pk3 = tl.make_block_ptr( - k_base, (T, K), (Hg * K, 1), (i_t * BT + 48, i_k * BK), (16, BK), (1, 0) - ) - kb0 = (tl.load(pk0, boundary_check=(0, 1)) * bb0[:, None] * eg0[:, None]).to( - lowp_dtype - ) - kb1 = (tl.load(pk1, boundary_check=(0, 1)) * bb1[:, None] * eg1[:, None]).to( - lowp_dtype - ) - kb2 = (tl.load(pk2, boundary_check=(0, 1)) * bb2[:, None] * eg2[:, None]).to( - lowp_dtype - ) - kb3 = (tl.load(pk3, boundary_check=(0, 1)) * bb3[:, None] * eg3[:, None]).to( - lowp_dtype - ) + kb0 = ( + _bp_ld2d(k_base, T, K, Hg * K, i_t * BT, i_k * BK, 16, BK) + * bb0[:, None] + * eg0[:, None] + ).to(lowp_dtype) + kb1 = ( + _bp_ld2d(k_base, T, K, Hg * K, i_t * BT + 16, i_k * BK, 16, BK) + * bb1[:, None] + * eg1[:, None] + ).to(lowp_dtype) + kb2 = ( + _bp_ld2d(k_base, T, K, Hg * K, i_t * BT + 32, i_k * BK, 16, BK) + * bb2[:, None] + * eg2[:, None] + ).to(lowp_dtype) + kb3 = ( + _bp_ld2d(k_base, T, K, Hg * K, i_t * BT + 48, i_k * BK, 16, BK) + * bb3[:, None] + * eg3[:, None] + ).to(lowp_dtype) w0 = tl.dot(h11, kb0) w1 = tl.dot(h21, kb0) + tl.dot(h22, kb1) w2 = tl.dot(h31, kb0) + tl.dot(h32, kb1) + tl.dot(h33, kb2) w3 = tl.dot(h41, kb0) + tl.dot(h42, kb1) + tl.dot(h43, kb2) + tl.dot(h44, kb3) - pw0 = tl.make_block_ptr( - w_base, (T, K), (K, 1), (i_t * BT, i_k * BK), (16, BK), (1, 0) - ) - pw1 = tl.make_block_ptr( - w_base, (T, K), (K, 1), (i_t * BT + 16, i_k * BK), (16, BK), (1, 0) + _bp_st2d( + w_base, T, K, K, i_t * BT, i_k * BK, w0.to(w_base.dtype.element_ty), 16, BK ) - pw2 = tl.make_block_ptr( - w_base, (T, K), (K, 1), (i_t * BT + 32, i_k * BK), (16, BK), (1, 0) - ) - pw3 = tl.make_block_ptr( - w_base, (T, K), (K, 1), (i_t * BT + 48, i_k * BK), (16, BK), (1, 0) + _bp_st2d( + w_base, + T, + K, + K, + i_t * BT + 16, + i_k * BK, + w1.to(w_base.dtype.element_ty), + 16, + BK, + ) + _bp_st2d( + w_base, + T, + K, + K, + i_t * BT + 32, + i_k * BK, + w2.to(w_base.dtype.element_ty), + 16, + BK, + ) + _bp_st2d( + w_base, + T, + K, + K, + i_t * BT + 48, + i_k * BK, + w3.to(w_base.dtype.element_ty), + 16, + BK, ) - tl.store(pw0, w0.to(pw0.dtype.element_ty), boundary_check=(0, 1)) - tl.store(pw1, w1.to(pw1.dtype.element_ty), boundary_check=(0, 1)) - tl.store(pw2, w2.to(pw2.dtype.element_ty), boundary_check=(0, 1)) - tl.store(pw3, w3.to(pw3.dtype.element_ty), boundary_check=(0, 1)) # ============================================================================= @@ -511,21 +545,9 @@ def recompute_w_u_head_major_kernel( g_base = g + i_h * T_flat + bos else: g_base = g + (i_b * H + i_h) * T_flat - p_beta = tl.make_block_ptr( - beta + bos * H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,) - ) - p_g = tl.make_block_ptr(g_base, (T,), (1,), (i_t * BT,), (BT,), (0,)) - p_Ai = tl.make_block_ptr( - Ai + (bos * H + i_h) * BT, - (T, BT), - (H * BT, 1), - (i_t * BT, 0), - (BT, BT), - (1, 0), - ) - b_beta = tl.load(p_beta, boundary_check=(0,)) - b_Ai = tl.load(p_Ai, boundary_check=(0, 1)) - b_g_raw = tl.load(p_g, boundary_check=(0,)) + b_beta = _bp_ld1d(beta + bos * H + i_h, T, H, i_t * BT, BT) + b_Ai = _bp_ld2d(Ai + (bos * H + i_h) * BT, T, BT, H * BT, i_t * BT, 0, BT, BT) + b_g_raw = _bp_ld1d(g_base, T, 1, i_t * BT, BT) b_g = tl.math.exp2(b_g_raw) if USE_EXP2 else exp(b_g_raw) # ---- u = Ai @ (v * beta) -> head-major store ---- @@ -536,16 +558,12 @@ def recompute_w_u_head_major_kernel( u_base = u + ((i_b * H + i_h) * T_flat) * V for i_v in range(tl.cdiv(V, BV)): - p_v = tl.make_block_ptr( - v_base, (T, V), (H * V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0) - ) - p_u = tl.make_block_ptr( - u_base, (T, V), (V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0) - ) - b_v = tl.load(p_v, boundary_check=(0, 1)) + b_v = _bp_ld2d(v_base, T, V, H * V, i_t * BT, i_v * BV, BT, BV) b_vb = (b_v * b_beta[:, None]).to(b_v.dtype) b_u = tl.dot(b_Ai, b_vb, allow_tf32=False) - tl.store(p_u, b_u.to(p_u.dtype.element_ty), boundary_check=(0, 1)) + _bp_st2d( + u_base, T, V, V, i_t * BT, i_v * BV, b_u.to(u_base.dtype.element_ty), BT, BV + ) # ---- w = Ai @ (k * beta * exp(g)) -> head-major store ---- k_base = k + (bos * Hg + i_h // (H // Hg)) * K @@ -555,17 +573,13 @@ def recompute_w_u_head_major_kernel( w_base = w + ((i_b * H + i_h) * T_flat) * K for i_k in range(tl.cdiv(K, BK)): - p_k = tl.make_block_ptr( - k_base, (T, K), (Hg * K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0) - ) - p_w = tl.make_block_ptr( - w_base, (T, K), (K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0) - ) - b_k = tl.load(p_k, boundary_check=(0, 1)) + b_k = _bp_ld2d(k_base, T, K, Hg * K, i_t * BT, i_k * BK, BT, BK) # Single fused multiply-cast before the dot. b_kb = (b_k * b_beta[:, None] * b_g[:, None]).to(b_k.dtype) b_w = tl.dot(b_Ai, b_kb) - tl.store(p_w, b_w.to(p_w.dtype.element_ty), boundary_check=(0, 1)) + _bp_st2d( + w_base, T, K, K, i_t * BT, i_k * BK, b_w.to(w_base.dtype.element_ty), BT, BK + ) def _run_split_path( diff --git a/aiter/ops/triton/_triton_kernels/gated_delta_rule/utils/cumsum.py b/aiter/ops/triton/_triton_kernels/gated_delta_rule/utils/cumsum.py index b4f2f4a4ff..248c67a8db 100644 --- a/aiter/ops/triton/_triton_kernels/gated_delta_rule/utils/cumsum.py +++ b/aiter/ops/triton/_triton_kernels/gated_delta_rule/utils/cumsum.py @@ -63,25 +63,27 @@ def chunk_local_cumsum_scalar_kernel( else: bos, eos = i_b * T, i_b * T + T + # Plain pointer arithmetic rather than tl.make_block_ptr, which is not + # available in every Triton build we target. + o_t = i_t * BT + tl.arange(0, BT) + m_t = o_t < T if HEAD_FIRST: - p_s = tl.make_block_ptr( - s + bos * H + i_h * T, (T,), (1,), (i_t * BT,), (BT,), (0,) - ) - p_o = tl.make_block_ptr( - o + bos * H + i_h * T, (T,), (1,), (i_t * BT,), (BT,), (0,) - ) + s_base = s + bos * H + i_h * T + o_base = o + bos * H + i_h * T + stride_t = 1 else: - p_s = tl.make_block_ptr(s + bos * H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) - p_o = tl.make_block_ptr(o + bos * H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) + s_base = s + bos * H + i_h + o_base = o + bos * H + i_h + stride_t = H # [BT] - b_s = tl.load(p_s, boundary_check=(0,)).to(tl.float32) + b_s = tl.load(s_base + o_t * stride_t, mask=m_t, other=0.0).to(tl.float32) b_o = tl.cumsum(b_s, axis=0) if REVERSE: b_z = tl.sum(b_s, axis=0) b_o = -b_o + b_z[None] + b_s if HAS_SCALE: b_o *= scale - tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0,)) + tl.store(o_base + o_t * stride_t, b_o.to(o_base.dtype.element_ty), mask=m_t) @triton.heuristics( @@ -130,49 +132,29 @@ def chunk_local_cumsum_vector_kernel( else: bos, eos = i_b * T, i_b * T + T + # Plain pointer arithmetic rather than tl.make_block_ptr, which is not + # available in every Triton build we target. + o_t = i_t * BT + tl.arange(0, BT) + o_s = i_s * BS + tl.arange(0, BS) + msk = (o_t < T)[:, None] & (o_s < S)[None, :] if HEAD_FIRST: - p_s = tl.make_block_ptr( - s + (bos * H + i_h * T) * S, - (T, S), - (S, 1), - (i_t * BT, i_s * BS), - (BT, BS), - (1, 0), - ) - p_o = tl.make_block_ptr( - o + (bos * H + i_h * T) * S, - (T, S), - (S, 1), - (i_t * BT, i_s * BS), - (BT, BS), - (1, 0), - ) + s_base = s + (bos * H + i_h * T) * S + o_base = o + (bos * H + i_h * T) * S + stride_t = S else: - p_s = tl.make_block_ptr( - s + (bos * H + i_h) * S, - (T, S), - (H * S, 1), - (i_t * BT, i_s * BS), - (BT, BS), - (1, 0), - ) - p_o = tl.make_block_ptr( - o + (bos * H + i_h) * S, - (T, S), - (H * S, 1), - (i_t * BT, i_s * BS), - (BT, BS), - (1, 0), - ) + s_base = s + (bos * H + i_h) * S + o_base = o + (bos * H + i_h) * S + stride_t = H * S + offs = o_t[:, None] * stride_t + o_s[None, :] # [BT, BS] - b_s = tl.load(p_s, boundary_check=(0, 1)).to(tl.float32) + b_s = tl.load(s_base + offs, mask=msk, other=0.0).to(tl.float32) if REVERSE: b_o = tl.cumsum(b_s, axis=0, reverse=True) else: b_o = tl.cumsum(b_s, axis=0) if HAS_SCALE: b_o *= scale - tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1)) + tl.store(o_base + offs, b_o.to(o_base.dtype.element_ty), mask=msk) def chunk_local_cumsum_scalar( diff --git a/aiter/ops/triton/_triton_kernels/gated_delta_rule/utils/l2norm.py b/aiter/ops/triton/_triton_kernels/gated_delta_rule/utils/l2norm.py index 852a96d4ca..7cbac2f453 100644 --- a/aiter/ops/triton/_triton_kernels/gated_delta_rule/utils/l2norm.py +++ b/aiter/ops/triton/_triton_kernels/gated_delta_rule/utils/l2norm.py @@ -156,18 +156,35 @@ def l2norm_bwd_kernel( BT: tl.constexpr, ): i_t = tl.program_id(0) - p_y = tl.make_block_ptr(y, (T, D), (D, 1), (i_t * BT, 0), (BT, BD), (1, 0)) - p_rstd = tl.make_block_ptr(rstd, (T,), (1,), (i_t * BT,), (BT,), (0,)) - p_dy = tl.make_block_ptr(dy, (T, D), (D, 1), (i_t * BT, 0), (BT, BD), (1, 0)) - p_dx = tl.make_block_ptr(dx, (T, D), (D, 1), (i_t * BT, 0), (BT, BD), (1, 0)) - - b_y = tl.load(p_y, boundary_check=(0, 1)).to(tl.float32) - b_rstd = tl.load(p_rstd, boundary_check=(0,)).to(tl.float32) - b_dy = tl.load(p_dy, boundary_check=(0, 1)).to(tl.float32) + + _p_y_0 = (i_t * BT) + tl.arange(0, BT) + _p_y_1 = (0) + tl.arange(0, BD) + b_y = tl.load( + y + _p_y_0[:, None] * (D) + _p_y_1[None, :] * (1), + mask=(_p_y_0[:, None] < (T)) & (_p_y_1[None, :] < (D)), + other=0.0, + ).to(tl.float32) + _p_rstd_0 = (i_t * BT) + tl.arange(0, BT) + b_rstd = tl.load(rstd + _p_rstd_0 * (1), mask=(_p_rstd_0 < (T)), other=0.0).to( + tl.float32 + ) + _p_dy_0 = (i_t * BT) + tl.arange(0, BT) + _p_dy_1 = (0) + tl.arange(0, BD) + b_dy = tl.load( + dy + _p_dy_0[:, None] * (D) + _p_dy_1[None, :] * (1), + mask=(_p_dy_0[:, None] < (T)) & (_p_dy_1[None, :] < (D)), + other=0.0, + ).to(tl.float32) b_dx = ( b_dy * b_rstd[:, None] - tl.sum(b_dy * b_y, 1)[:, None] * b_y * b_rstd[:, None] ) - tl.store(p_dx, b_dx.to(p_dx.dtype.element_ty), boundary_check=(0, 1)) + _p_dx_0 = (i_t * BT) + tl.arange(0, BT) + _p_dx_1 = (0) + tl.arange(0, BD) + tl.store( + dx + _p_dx_0[:, None] * (D) + _p_dx_1[None, :] * (1), + b_dx.to(dx.dtype.element_ty), + mask=(_p_dx_0[:, None] < (T)) & (_p_dx_1[None, :] < (D)), + ) def l2norm_fwd( diff --git a/aiter/ops/triton/_triton_kernels/gated_delta_rule/utils/solve_tril.py b/aiter/ops/triton/_triton_kernels/gated_delta_rule/utils/solve_tril.py index 5dcace8bc3..3c06f2aa0a 100644 --- a/aiter/ops/triton/_triton_kernels/gated_delta_rule/utils/solve_tril.py +++ b/aiter/ops/triton/_triton_kernels/gated_delta_rule/utils/solve_tril.py @@ -35,6 +35,32 @@ ) +# tl.make_block_ptr was removed in Triton 3.8 ("Block pointers have been removed +# in favor of the tensor descriptor API"). Every block access in this file is a +# row-major (BR, BC) tile with unit column stride, so plain pointer arithmetic +# with an explicit bounds mask reproduces the boundary_check=(0, 1) behaviour. +@triton.jit +def _bp_ld2d(base, R, C, rs, r0, c0, BR: tl.constexpr, BC: tl.constexpr): + r = r0 + tl.arange(0, BR) + c = c0 + tl.arange(0, BC) + return tl.load( + base + r[:, None] * rs + c[None, :], + mask=(r < R)[:, None] & (c < C)[None, :], + other=0.0, + ) + + +@triton.jit +def _bp_st2d(base, R, C, rs, r0, c0, val, BR: tl.constexpr, BC: tl.constexpr): + r = r0 + tl.arange(0, BR) + c = c0 + tl.arange(0, BC) + tl.store( + base + r[:, None] * rs + c[None, :], + val, + mask=(r < R)[:, None] & (c < C)[None, :], + ) + + @triton.heuristics( { "IS_VARLEN": lambda args: args["cu_seqlens"] is not None, @@ -90,11 +116,8 @@ def solve_tril_16x16_kernel( offset = (i_t * 16) % BT if not USE_TMA: - p_A = tl.make_block_ptr( - A, (T, BT), (H * BT, 1), (i_t * 16, offset), (16, 16), (1, 0) - ) # [16, 16] - b_A = tl.load(p_A, boundary_check=(0, 1)).to(tl.float32) + b_A = _bp_ld2d(A, T, BT, H * BT, i_t * 16, offset, 16, 16).to(tl.float32) b_A = tl.where(m_A, b_A, 0) else: desc = make_tensor_descriptor(A, [T, BT], [H * BT, 1], [16, 16]) @@ -112,13 +135,16 @@ def solve_tril_16x16_kernel( b_A = tl.where((o_i == i)[:, None], b_a, b_A) b_A += m_I if not USE_TMA: - p_Ai = tl.make_block_ptr( - Ai, (T, 16), (H * 16, 1), (i_t * 16, 0), (16, 16), (1, 0) - ) - tl.store( - p_Ai, - b_A.to(p_Ai.dtype.element_ty, fp_downcast_rounding="rtne"), - boundary_check=(0, 1), + _bp_st2d( + Ai, + T, + 16, + H * 16, + i_t * 16, + 0, + b_A.to(Ai.dtype.element_ty, fp_downcast_rounding="rtne"), + 16, + 16, ) else: desc_o.store([i_t * 16, 0], b_A.to(desc_o.dtype, fp_downcast_rounding="rtne")) @@ -181,14 +207,8 @@ def merge_16x16_to_32x32_inverse_kernel( Ai += (bos * H + i_h) * BT if not USE_TMA: - p_A_11 = tl.make_block_ptr( - A, (T, BT), (H * BT, 1), (i_t * BT, 0), (16, 16), (1, 0) - ) - p_A_22 = tl.make_block_ptr( - A, (T, BT), (H * BT, 1), (i_t * BT + 16, 16), (16, 16), (1, 0) - ) - b_Ai_11 = tl.load(p_A_11, boundary_check=(0, 1)).to(tl.float32) - b_Ai_22 = tl.load(p_A_22, boundary_check=(0, 1)).to(tl.float32) + b_Ai_11 = _bp_ld2d(A, T, BT, H * BT, i_t * BT, 0, 16, 16).to(tl.float32) + b_Ai_22 = _bp_ld2d(A, T, BT, H * BT, i_t * BT + 16, 16, 16, 16).to(tl.float32) else: desc = make_tensor_descriptor(A, [T, BT], [H * BT, 1], [16, 16]) desc_o = make_tensor_descriptor(Ai, [T, BT], [H * BT, 1], [16, 16]) @@ -212,10 +232,7 @@ def merge_16x16_to_32x32_inverse_kernel( b_Ai_22 += m_I if not USE_TMA: - p_A_21 = tl.make_block_ptr( - A, (T, BT), (H * BT, 1), (i_t * BT + 16, 0), (16, 16), (1, 0) - ) - b_A_21 = tl.load(p_A_21, boundary_check=(0, 1)).to(tl.float32) + b_A_21 = _bp_ld2d(A, T, BT, H * BT, i_t * BT + 16, 0, 16, 16).to(tl.float32) else: b_A_21 = desc.load([i_t * BT + 16, 0]).to(tl.float32) @@ -232,38 +249,40 @@ def merge_16x16_to_32x32_inverse_kernel( z16 = tl.zeros([16, 16], dtype=b_Ai_11.dtype) if not USE_TMA: - p_Ai_11 = tl.make_block_ptr( - Ai, (T, BT), (H * BT, 1), (i_t * BT, 0), (16, 16), (1, 0) - ) - p_Ai_21 = tl.make_block_ptr( - Ai, (T, BT), (H * BT, 1), (i_t * BT + 16, 0), (16, 16), (1, 0) - ) - p_Ai_22 = tl.make_block_ptr( - Ai, (T, BT), (H * BT, 1), (i_t * BT + 16, 16), (16, 16), (1, 0) - ) - p_Ai_12 = tl.make_block_ptr( - Ai, (T, BT), (H * BT, 1), (i_t * BT, 16), (16, 16), (1, 0) - ) - tl.store( - p_Ai_11, - b_Ai_11.to(p_Ai_11.dtype.element_ty, fp_downcast_rounding="rtne"), - boundary_check=(0, 1), - ) - tl.store( - p_Ai_22, - b_Ai_22.to(p_Ai_22.dtype.element_ty, fp_downcast_rounding="rtne"), - boundary_check=(0, 1), - ) - tl.store( - p_Ai_21, - b_Ai_21.to(p_Ai_21.dtype.element_ty, fp_downcast_rounding="rtne"), - boundary_check=(0, 1), - ) - tl.store( - p_Ai_12, - z16.to(p_Ai_12.dtype.element_ty), - boundary_check=(0, 1), - ) + _bp_st2d( + Ai, + T, + BT, + H * BT, + i_t * BT, + 0, + b_Ai_11.to(Ai.dtype.element_ty, fp_downcast_rounding="rtne"), + 16, + 16, + ) + _bp_st2d( + Ai, + T, + BT, + H * BT, + i_t * BT + 16, + 16, + b_Ai_22.to(Ai.dtype.element_ty, fp_downcast_rounding="rtne"), + 16, + 16, + ) + _bp_st2d( + Ai, + T, + BT, + H * BT, + i_t * BT + 16, + 0, + b_Ai_21.to(Ai.dtype.element_ty, fp_downcast_rounding="rtne"), + 16, + 16, + ) + _bp_st2d(Ai, T, BT, H * BT, i_t * BT, 16, z16.to(Ai.dtype.element_ty), 16, 16) else: desc_o.store( [i_t * BT + 0, 0], b_Ai_11.to(desc_o.dtype, fp_downcast_rounding="rtne") @@ -334,22 +353,10 @@ def merge_16x16_to_64x64_inverse_kernel( Ai += (bos * H + i_h) * BT if not USE_TMA: - p_A_11 = tl.make_block_ptr( - A, (T, BT), (H * BT, 1), (i_t * BT, 0), (16, 16), (1, 0) - ) - p_A_22 = tl.make_block_ptr( - A, (T, BT), (H * BT, 1), (i_t * BT + 16, 16), (16, 16), (1, 0) - ) - p_A_33 = tl.make_block_ptr( - A, (T, BT), (H * BT, 1), (i_t * BT + 32, 32), (16, 16), (1, 0) - ) - p_A_44 = tl.make_block_ptr( - A, (T, BT), (H * BT, 1), (i_t * BT + 48, 48), (16, 16), (1, 0) - ) - b_Ai_11 = tl.load(p_A_11, boundary_check=(0, 1)).to(tl.float32) - b_Ai_22 = tl.load(p_A_22, boundary_check=(0, 1)).to(tl.float32) - b_Ai_33 = tl.load(p_A_33, boundary_check=(0, 1)).to(tl.float32) - b_Ai_44 = tl.load(p_A_44, boundary_check=(0, 1)).to(tl.float32) + b_Ai_11 = _bp_ld2d(A, T, BT, H * BT, i_t * BT, 0, 16, 16).to(tl.float32) + b_Ai_22 = _bp_ld2d(A, T, BT, H * BT, i_t * BT + 16, 16, 16, 16).to(tl.float32) + b_Ai_33 = _bp_ld2d(A, T, BT, H * BT, i_t * BT + 32, 32, 16, 16).to(tl.float32) + b_Ai_44 = _bp_ld2d(A, T, BT, H * BT, i_t * BT + 48, 48, 16, 16).to(tl.float32) else: desc = make_tensor_descriptor(A, [T, BT], [H * BT, 1], [16, 16]) desc_o = make_tensor_descriptor(Ai, [T, BT], [H * BT, 1], [16, 16]) @@ -390,30 +397,12 @@ def merge_16x16_to_64x64_inverse_kernel( b_Ai_44 += m_I if not USE_TMA: - p_A_21 = tl.make_block_ptr( - A, (T, BT), (H * BT, 1), (i_t * BT + 16, 0), (16, 16), (1, 0) - ) - p_A_31 = tl.make_block_ptr( - A, (T, BT), (H * BT, 1), (i_t * BT + 32, 0), (16, 16), (1, 0) - ) - p_A_32 = tl.make_block_ptr( - A, (T, BT), (H * BT, 1), (i_t * BT + 32, 16), (16, 16), (1, 0) - ) - p_A_41 = tl.make_block_ptr( - A, (T, BT), (H * BT, 1), (i_t * BT + 48, 0), (16, 16), (1, 0) - ) - p_A_42 = tl.make_block_ptr( - A, (T, BT), (H * BT, 1), (i_t * BT + 48, 16), (16, 16), (1, 0) - ) - p_A_43 = tl.make_block_ptr( - A, (T, BT), (H * BT, 1), (i_t * BT + 48, 32), (16, 16), (1, 0) - ) - b_A_21 = tl.load(p_A_21, boundary_check=(0, 1)).to(tl.float32) - b_A_31 = tl.load(p_A_31, boundary_check=(0, 1)).to(tl.float32) - b_A_32 = tl.load(p_A_32, boundary_check=(0, 1)).to(tl.float32) - b_A_41 = tl.load(p_A_41, boundary_check=(0, 1)).to(tl.float32) - b_A_42 = tl.load(p_A_42, boundary_check=(0, 1)).to(tl.float32) - b_A_43 = tl.load(p_A_43, boundary_check=(0, 1)).to(tl.float32) + b_A_21 = _bp_ld2d(A, T, BT, H * BT, i_t * BT + 16, 0, 16, 16).to(tl.float32) + b_A_31 = _bp_ld2d(A, T, BT, H * BT, i_t * BT + 32, 0, 16, 16).to(tl.float32) + b_A_32 = _bp_ld2d(A, T, BT, H * BT, i_t * BT + 32, 16, 16, 16).to(tl.float32) + b_A_41 = _bp_ld2d(A, T, BT, H * BT, i_t * BT + 48, 0, 16, 16).to(tl.float32) + b_A_42 = _bp_ld2d(A, T, BT, H * BT, i_t * BT + 48, 16, 16, 16).to(tl.float32) + b_A_43 = _bp_ld2d(A, T, BT, H * BT, i_t * BT + 48, 32, 16, 16).to(tl.float32) else: b_A_21 = desc.load([i_t * BT + 16, 0]).to(tl.float32) b_A_31 = desc.load([i_t * BT + 32, 0]).to(tl.float32) @@ -465,111 +454,129 @@ def merge_16x16_to_64x64_inverse_kernel( z16 = tl.zeros([16, 16], dtype=b_Ai_11.dtype) if not USE_TMA: - p_Ai_11 = tl.make_block_ptr( - Ai, (T, BT), (H * BT, 1), (i_t * BT, 0), (16, 16), (1, 0) - ) - p_Ai_22 = tl.make_block_ptr( - Ai, (T, BT), (H * BT, 1), (i_t * BT + 16, 16), (16, 16), (1, 0) - ) - p_Ai_33 = tl.make_block_ptr( - Ai, (T, BT), (H * BT, 1), (i_t * BT + 32, 32), (16, 16), (1, 0) - ) - p_Ai_44 = tl.make_block_ptr( - Ai, (T, BT), (H * BT, 1), (i_t * BT + 48, 48), (16, 16), (1, 0) - ) - p_Ai_21 = tl.make_block_ptr( - Ai, (T, BT), (H * BT, 1), (i_t * BT + 16, 0), (16, 16), (1, 0) - ) - p_Ai_31 = tl.make_block_ptr( - Ai, (T, BT), (H * BT, 1), (i_t * BT + 32, 0), (16, 16), (1, 0) - ) - p_Ai_32 = tl.make_block_ptr( - Ai, (T, BT), (H * BT, 1), (i_t * BT + 32, 16), (16, 16), (1, 0) - ) - p_Ai_41 = tl.make_block_ptr( - Ai, (T, BT), (H * BT, 1), (i_t * BT + 48, 0), (16, 16), (1, 0) - ) - p_Ai_42 = tl.make_block_ptr( - Ai, (T, BT), (H * BT, 1), (i_t * BT + 48, 16), (16, 16), (1, 0) - ) - p_Ai_43 = tl.make_block_ptr( - Ai, (T, BT), (H * BT, 1), (i_t * BT + 48, 32), (16, 16), (1, 0) - ) # 6 strict-upper sub-block zero ptrs - p_Ai_12 = tl.make_block_ptr( - Ai, (T, BT), (H * BT, 1), (i_t * BT, 16), (16, 16), (1, 0) - ) - p_Ai_13 = tl.make_block_ptr( - Ai, (T, BT), (H * BT, 1), (i_t * BT, 32), (16, 16), (1, 0) - ) - p_Ai_14 = tl.make_block_ptr( - Ai, (T, BT), (H * BT, 1), (i_t * BT, 48), (16, 16), (1, 0) - ) - p_Ai_23 = tl.make_block_ptr( - Ai, (T, BT), (H * BT, 1), (i_t * BT + 16, 32), (16, 16), (1, 0) - ) - p_Ai_24 = tl.make_block_ptr( - Ai, (T, BT), (H * BT, 1), (i_t * BT + 16, 48), (16, 16), (1, 0) - ) - p_Ai_34 = tl.make_block_ptr( - Ai, (T, BT), (H * BT, 1), (i_t * BT + 32, 48), (16, 16), (1, 0) - ) - tl.store( - p_Ai_11, - b_Ai_11.to(p_Ai_11.dtype.element_ty, fp_downcast_rounding="rtne"), - boundary_check=(0, 1), - ) - tl.store( - p_Ai_22, - b_Ai_22.to(p_Ai_22.dtype.element_ty, fp_downcast_rounding="rtne"), - boundary_check=(0, 1), - ) - tl.store( - p_Ai_33, - b_Ai_33.to(p_Ai_33.dtype.element_ty, fp_downcast_rounding="rtne"), - boundary_check=(0, 1), - ) - tl.store( - p_Ai_44, - b_Ai_44.to(p_Ai_44.dtype.element_ty, fp_downcast_rounding="rtne"), - boundary_check=(0, 1), - ) - tl.store( - p_Ai_21, - b_Ai_21.to(p_Ai_21.dtype.element_ty, fp_downcast_rounding="rtne"), - boundary_check=(0, 1), - ) - tl.store( - p_Ai_31, - b_Ai_31.to(p_Ai_31.dtype.element_ty, fp_downcast_rounding="rtne"), - boundary_check=(0, 1), - ) - tl.store( - p_Ai_32, - b_Ai_32.to(p_Ai_32.dtype.element_ty, fp_downcast_rounding="rtne"), - boundary_check=(0, 1), - ) - tl.store( - p_Ai_41, - b_Ai_41.to(p_Ai_41.dtype.element_ty, fp_downcast_rounding="rtne"), - boundary_check=(0, 1), - ) - tl.store( - p_Ai_42, - b_Ai_42.to(p_Ai_42.dtype.element_ty, fp_downcast_rounding="rtne"), - boundary_check=(0, 1), - ) - tl.store( - p_Ai_43, - b_Ai_43.to(p_Ai_43.dtype.element_ty, fp_downcast_rounding="rtne"), - boundary_check=(0, 1), + _bp_st2d( + Ai, + T, + BT, + H * BT, + i_t * BT, + 0, + b_Ai_11.to(Ai.dtype.element_ty, fp_downcast_rounding="rtne"), + 16, + 16, + ) + _bp_st2d( + Ai, + T, + BT, + H * BT, + i_t * BT + 16, + 16, + b_Ai_22.to(Ai.dtype.element_ty, fp_downcast_rounding="rtne"), + 16, + 16, + ) + _bp_st2d( + Ai, + T, + BT, + H * BT, + i_t * BT + 32, + 32, + b_Ai_33.to(Ai.dtype.element_ty, fp_downcast_rounding="rtne"), + 16, + 16, + ) + _bp_st2d( + Ai, + T, + BT, + H * BT, + i_t * BT + 48, + 48, + b_Ai_44.to(Ai.dtype.element_ty, fp_downcast_rounding="rtne"), + 16, + 16, + ) + _bp_st2d( + Ai, + T, + BT, + H * BT, + i_t * BT + 16, + 0, + b_Ai_21.to(Ai.dtype.element_ty, fp_downcast_rounding="rtne"), + 16, + 16, + ) + _bp_st2d( + Ai, + T, + BT, + H * BT, + i_t * BT + 32, + 0, + b_Ai_31.to(Ai.dtype.element_ty, fp_downcast_rounding="rtne"), + 16, + 16, + ) + _bp_st2d( + Ai, + T, + BT, + H * BT, + i_t * BT + 32, + 16, + b_Ai_32.to(Ai.dtype.element_ty, fp_downcast_rounding="rtne"), + 16, + 16, + ) + _bp_st2d( + Ai, + T, + BT, + H * BT, + i_t * BT + 48, + 0, + b_Ai_41.to(Ai.dtype.element_ty, fp_downcast_rounding="rtne"), + 16, + 16, + ) + _bp_st2d( + Ai, + T, + BT, + H * BT, + i_t * BT + 48, + 16, + b_Ai_42.to(Ai.dtype.element_ty, fp_downcast_rounding="rtne"), + 16, + 16, + ) + _bp_st2d( + Ai, + T, + BT, + H * BT, + i_t * BT + 48, + 32, + b_Ai_43.to(Ai.dtype.element_ty, fp_downcast_rounding="rtne"), + 16, + 16, + ) + _bp_st2d(Ai, T, BT, H * BT, i_t * BT, 16, z16.to(Ai.dtype.element_ty), 16, 16) + _bp_st2d(Ai, T, BT, H * BT, i_t * BT, 32, z16.to(Ai.dtype.element_ty), 16, 16) + _bp_st2d(Ai, T, BT, H * BT, i_t * BT, 48, z16.to(Ai.dtype.element_ty), 16, 16) + _bp_st2d( + Ai, T, BT, H * BT, i_t * BT + 16, 32, z16.to(Ai.dtype.element_ty), 16, 16 + ) + _bp_st2d( + Ai, T, BT, H * BT, i_t * BT + 16, 48, z16.to(Ai.dtype.element_ty), 16, 16 + ) + _bp_st2d( + Ai, T, BT, H * BT, i_t * BT + 32, 48, z16.to(Ai.dtype.element_ty), 16, 16 ) - tl.store(p_Ai_12, z16.to(p_Ai_12.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_Ai_13, z16.to(p_Ai_13.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_Ai_14, z16.to(p_Ai_14.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_Ai_23, z16.to(p_Ai_23.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_Ai_24, z16.to(p_Ai_24.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_Ai_34, z16.to(p_Ai_34.dtype.element_ty), boundary_check=(0, 1)) else: desc_o.store( [i_t * BT + 0, 0], b_Ai_11.to(desc_o.dtype, fp_downcast_rounding="rtne") diff --git a/aiter/ops/triton/_triton_kernels/gated_delta_rule/utils/wy_representation.py b/aiter/ops/triton/_triton_kernels/gated_delta_rule/utils/wy_representation.py index af4965a2f9..a49fec2ef4 100644 --- a/aiter/ops/triton/_triton_kernels/gated_delta_rule/utils/wy_representation.py +++ b/aiter/ops/triton/_triton_kernels/gated_delta_rule/utils/wy_representation.py @@ -68,35 +68,36 @@ def chunk_scaled_dot_kkt_fwd_kernel( o_t = i_t * BT + tl.arange(0, BT) m_t = o_t < T - p_b = tl.make_block_ptr(beta + bos * H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) - b_b = tl.load(p_b, boundary_check=(0,)) + _p_b_0 = (i_t * BT) + tl.arange(0, BT) + b_b = tl.load(beta + bos * H + i_h + _p_b_0 * (H), mask=(_p_b_0 < (T)), other=0.0) b_A = tl.zeros([BT, BT], dtype=tl.float32) for i_k in range(tl.cdiv(K, BK)): - p_k = tl.make_block_ptr( - k + (bos * H + i_h) * K, - (T, K), - (H * K, 1), - (i_t * BT, i_k * BK), - (BT, BK), - (1, 0), + _p_k_0 = (i_t * BT) + tl.arange(0, BT) + _p_k_1 = (i_k * BK) + tl.arange(0, BK) + b_k = tl.load( + k + (bos * H + i_h) * K + _p_k_0[:, None] * (H * K) + _p_k_1[None, :] * (1), + mask=(_p_k_0[:, None] < (T)) & (_p_k_1[None, :] < (K)), + other=0.0, ) - b_k = tl.load(p_k, boundary_check=(0, 1)) b_A = tl.dot(b_k, tl.trans(b_k), acc=b_A) if USE_G: - p_g = tl.make_block_ptr(g + bos * H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) - b_g = tl.load(p_g, boundary_check=(0,)) + _p_g_0 = (i_t * BT) + tl.arange(0, BT) + b_g = tl.load(g + bos * H + i_h + _p_g_0 * (H), mask=(_p_g_0 < (T)), other=0.0) b_g_diff = b_g[:, None] - b_g[None, :] b_A *= exp(b_g_diff) b_A *= b_b[:, None] m_A = (o_t[:, None] > o_t[None, :]) & (m_t[:, None] & m_t) b_A = tl.where(m_A, b_A, 0) - p_A = tl.make_block_ptr( - A + (bos * H + i_h) * BT, (T, BT), (BT * H, 1), (i_t * BT, 0), (BT, BT), (1, 0) + _p_A_0 = (i_t * BT) + tl.arange(0, BT) + _p_A_1 = (0) + tl.arange(0, BT) + tl.store( + A + (bos * H + i_h) * BT + _p_A_0[:, None] * (BT * H) + _p_A_1[None, :] * (1), + b_A.to(A.dtype.element_ty), + mask=(_p_A_0[:, None] < (T)) & (_p_A_1[None, :] < (BT)), ) - tl.store(p_A, b_A.to(p_A.dtype.element_ty), boundary_check=(0, 1)) def chunk_scaled_dot_kkt_fwd( @@ -200,65 +201,60 @@ def recompute_w_u_fwd_kernel( T = eos - bos else: bos, eos = i_b * T, i_b * T + T - p_b = tl.make_block_ptr(beta + bos * H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) - b_b = tl.load(p_b, boundary_check=(0,)) + _p_b_0 = (i_t * BT) + tl.arange(0, BT) + b_b = tl.load(beta + bos * H + i_h + _p_b_0 * (H), mask=(_p_b_0 < (T)), other=0.0) - p_A = tl.make_block_ptr( - A + (bos * H + i_h) * BT, (T, BT), (H * BT, 1), (i_t * BT, 0), (BT, BT), (1, 0) + _p_A_0 = (i_t * BT) + tl.arange(0, BT) + _p_A_1 = (0) + tl.arange(0, BT) + b_A = tl.load( + A + (bos * H + i_h) * BT + _p_A_0[:, None] * (H * BT) + _p_A_1[None, :] * (1), + mask=(_p_A_0[:, None] < (T)) & (_p_A_1[None, :] < (BT)), + other=0.0, ) - b_A = tl.load(p_A, boundary_check=(0, 1)) for i_v in range(tl.cdiv(V, BV)): - p_v = tl.make_block_ptr( - v + (bos * H + i_h) * V, - (T, V), - (H * V, 1), - (i_t * BT, i_v * BV), - (BT, BV), - (1, 0), + _p_v_0 = (i_t * BT) + tl.arange(0, BT) + _p_v_1 = (i_v * BV) + tl.arange(0, BV) + b_v = tl.load( + v + (bos * H + i_h) * V + _p_v_0[:, None] * (H * V) + _p_v_1[None, :] * (1), + mask=(_p_v_0[:, None] < (T)) & (_p_v_1[None, :] < (V)), + other=0.0, ) - p_u = tl.make_block_ptr( - u + (bos * H + i_h) * V, - (T, V), - (H * V, 1), - (i_t * BT, i_v * BV), - (BT, BV), - (1, 0), - ) - b_v = tl.load(p_v, boundary_check=(0, 1)) b_vb = (b_v * b_b[:, None]).to(b_v.dtype) b_u = tl.dot(b_A, b_vb, allow_tf32=False) - tl.store(p_u, b_u.to(p_u.dtype.element_ty), boundary_check=(0, 1)) + _p_u_0 = (i_t * BT) + tl.arange(0, BT) + _p_u_1 = (i_v * BV) + tl.arange(0, BV) + tl.store( + u + (bos * H + i_h) * V + _p_u_0[:, None] * (H * V) + _p_u_1[None, :] * (1), + b_u.to(u.dtype.element_ty), + mask=(_p_u_0[:, None] < (T)) & (_p_u_1[None, :] < (V)), + ) if USE_G: - p_g = tl.make_block_ptr( - g + (bos * H + i_h), (T,), (H,), (i_t * BT,), (BT,), (0,) + _p_g_0 = (i_t * BT) + tl.arange(0, BT) + b_g = exp( + tl.load(g + (bos * H + i_h) + _p_g_0 * (H), mask=(_p_g_0 < (T)), other=0.0) ) - b_g = exp(tl.load(p_g, boundary_check=(0,))) for i_k in range(tl.cdiv(K, BK)): - p_k = tl.make_block_ptr( - k + (bos * H + i_h) * K, - (T, K), - (H * K, 1), - (i_t * BT, i_k * BK), - (BT, BK), - (1, 0), + _p_k_0 = (i_t * BT) + tl.arange(0, BT) + _p_k_1 = (i_k * BK) + tl.arange(0, BK) + b_k = tl.load( + k + (bos * H + i_h) * K + _p_k_0[:, None] * (H * K) + _p_k_1[None, :] * (1), + mask=(_p_k_0[:, None] < (T)) & (_p_k_1[None, :] < (K)), + other=0.0, ) - p_w = tl.make_block_ptr( - w + (bos * H + i_h) * K, - (T, K), - (H * K, 1), - (i_t * BT, i_k * BK), - (BT, BK), - (1, 0), - ) - b_k = tl.load(p_k, boundary_check=(0, 1)) b_kb = b_k * b_b[:, None] if USE_G: b_kb *= b_g[:, None] b_w = tl.dot(b_A, b_kb.to(b_k.dtype)) - tl.store(p_w, b_w.to(p_w.dtype.element_ty), boundary_check=(0, 1)) + _p_w_0 = (i_t * BT) + tl.arange(0, BT) + _p_w_1 = (i_k * BK) + tl.arange(0, BK) + tl.store( + w + (bos * H + i_h) * K + _p_w_0[:, None] * (H * K) + _p_w_1[None, :] * (1), + b_w.to(w.dtype.element_ty), + mask=(_p_w_0[:, None] < (T)) & (_p_w_1[None, :] < (K)), + ) def recompute_w_u_fwd(