From a686c8e1b7e5870eadfc9f47447de2b8a26211fc Mon Sep 17 00:00:00 2001 From: Nick Riasanovsky Date: Mon, 12 Oct 2026 11:38:43 -0400 Subject: [PATCH] [Triton] Migrate Triton kernels off tl.make_block_ptr for triton-main compat (#60502) Signed-off-by: Nick Riasanovsky --- .../amd/ops/third_party/kda/kernels.py | 357 +++++----------- .../nvidia/ops/third_party/kda/kernels.py | 357 +++++----------- .../inkling/amd/ops/fa4_rel_attention.py | 40 +- .../nvidia/ops/third_party/kda/chunk.py | 260 ++++------- .../nvidia/ops/third_party/kda/chunk_intra.py | 270 +++++------- .../kda/chunk_intra_token_parallel.py | 16 +- .../minimax_m3/common/ops/index_topk.py | 18 +- .../ops/chunk_delta_h.py | 182 ++++---- .../flash_linear_attention/ops/chunk_o.py | 49 +-- .../ops/chunk_scaled_dot_kkt.py | 35 +- .../flash_linear_attention/ops/cumsum.py | 69 +-- .../ops/fused_norm_gate.py | 34 +- .../flash_linear_attention/ops/kda.py | 402 ++++++------------ .../flash_linear_attention/ops/l2norm.py | 15 +- .../flash_linear_attention/ops/solve_tril.py | 172 +++----- .../flash_linear_attention/ops/wy_fast.py | 75 ++-- 16 files changed, 836 insertions(+), 1515 deletions(-) diff --git a/vllm/models/glm5next/amd/ops/third_party/kda/kernels.py b/vllm/models/glm5next/amd/ops/third_party/kda/kernels.py index 9805b37471..0e32c94543 100644 --- a/vllm/models/glm5next/amd/ops/third_party/kda/kernels.py +++ b/vllm/models/glm5next/amd/ops/third_party/kda/kernels.py @@ -228,13 +228,13 @@ def chunk_kda_scaled_dot_kkt_fwd_kernel_intra_sub_inter( NC: tl.constexpr, IS_VARLEN: tl.constexpr, ): - i_t, i_c, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_t, i_c, i_bh = tl.program_id(0).to(tl.int64), tl.program_id(1), tl.program_id(2) i_b, i_h = i_bh // H, i_bh % H i_i, i_j = i_c // NC, i_c % NC if IS_VARLEN: i_n, i_t = ( tl.load(chunk_indices + i_t * 2).to(tl.int32), - tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), + tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64), ) bos, eos = ( tl.load(cu_seqlens + i_n).to(tl.int32), @@ -255,58 +255,51 @@ def chunk_kda_scaled_dot_kkt_fwd_kernel_intra_sub_inter( A += (bos * H + i_h) * BT Aqk += (bos * H + i_h) * BT - p_b = tl.make_block_ptr( - beta + bos * H + i_h, (T,), (H,), (i_t * BT + i_i * BC,), (BC,), (0,) - ) - b_b = tl.load(p_b, boundary_check=(0,)) + o_i = i_t * BT + i_i * BC + tl.arange(0, BC) + o_j = i_t * BT + i_j * BC + tl.arange(0, BC) + m_i = o_i < T + m_j = o_j < T + p_b = beta + bos * H + i_h + o_i * H + b_b = tl.load(p_b, mask=m_i, other=0.0) b_A = tl.zeros([BC, BC], dtype=tl.float32) b_Aqk = tl.zeros([BC, BC], 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_i * BC, i_k * BK), (BC, BK), (1, 0) - ) - p_k = tl.make_block_ptr( - k, (T, K), (H * K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0) - ) - p_g = tl.make_block_ptr( - g, (T, K), (H * K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0) - ) - b_kt = tl.make_block_ptr( - k, (K, T), (1, H * K), (i_k * BK, i_t * BT + i_j * BC), (BK, BC), (0, 1) - ) - p_gk = tl.make_block_ptr( - g, (K, T), (1, H * K), (i_k * BK, i_t * BT + i_j * BC), (BK, BC), (0, 1) - ) - o_k = i_k * BK + tl.arange(0, BK) m_k = o_k < K + m_qk = m_i[:, None] & m_k[None, :] + m_kj = m_k[:, None] & m_j[None, :] + + p_q = q + o_i[:, None] * (H * K) + o_k[None, :] + p_k = k + o_i[:, None] * (H * K) + o_k[None, :] + p_g = g + o_i[:, None] * (H * K) + o_k[None, :] + b_kt = k + o_k[:, None] + o_j[None, :] * (H * K) + p_gk = g + o_k[:, None] + o_j[None, :] * (H * K) + # [BK,] b_gn = tl.load(g + (i_t * BT + i_i * BC) * H * K + o_k, mask=m_k, other=0) # [BC, BK] - b_g = tl.load(p_g, boundary_check=(0, 1)) - b_k = tl.load(p_k, boundary_check=(0, 1)) * exp2(b_g - b_gn[None, :]) + b_g = tl.load(p_g, mask=m_qk, other=0.0) + b_k = tl.load(p_k, mask=m_qk, other=0.0) * exp2(b_g - b_gn[None, :]) # [BK, BC] - b_gk = tl.load(p_gk, boundary_check=(0, 1)) - b_kt = tl.load(b_kt, boundary_check=(0, 1)) + b_gk = tl.load(p_gk, mask=m_kj, other=0.0) + b_kt = tl.load(b_kt, mask=m_kj, other=0.0) # [BC, BC] b_ktg = b_kt * exp2(b_gn[:, None] - b_gk) b_A += tl.dot(b_k, b_ktg) - b_q = tl.load(p_q, boundary_check=(0, 1)) + b_q = tl.load(p_q, mask=m_qk, other=0.0) b_qg = b_q * exp2(b_g - b_gn[None, :]) * scale b_Aqk += tl.dot(b_qg, b_ktg) b_A *= b_b[:, None] - p_A = tl.make_block_ptr( - A, (T, BT), (H * BT, 1), (i_t * BT + i_i * BC, i_j * BC), (BC, BC), (1, 0) - ) - tl.store(p_A, b_A.to(A.dtype.element_ty), boundary_check=(0, 1)) - p_Aqk = tl.make_block_ptr( - Aqk, (T, BT), (H * BT, 1), (i_t * BT + i_i * BC, i_j * BC), (BC, BC), (1, 0) - ) - tl.store(p_Aqk, b_Aqk.to(Aqk.dtype.element_ty), boundary_check=(0, 1)) + o_jA = i_j * BC + tl.arange(0, BC) + m_A = m_i[:, None] & (o_jA[None, :] < BT) + p_A = A + o_i[:, None] * (H * BT) + o_jA[None, :] + tl.store(p_A, b_A.to(A.dtype.element_ty), mask=m_A) + p_Aqk = Aqk + o_i[:, None] * (H * BT) + o_jA[None, :] + tl.store(p_Aqk, b_Aqk.to(Aqk.dtype.element_ty), mask=m_A) @triton.heuristics({"IS_VARLEN": lambda args: args["cu_seqlens"] is not None}) @@ -333,12 +326,12 @@ def chunk_kda_scaled_dot_kkt_fwd_kernel_intra_sub_intra( BK: tl.constexpr, IS_VARLEN: tl.constexpr, ): - i_t, i_i, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_t, i_i, i_bh = tl.program_id(0).to(tl.int64), tl.program_id(1), tl.program_id(2) i_b, i_h = i_bh // H, i_bh % H if IS_VARLEN: i_n, i_t = ( tl.load(chunk_indices + i_t * 2).to(tl.int32), - tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), + tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64), ) bos, eos = ( tl.load(cu_seqlens + i_n).to(tl.int32), @@ -357,33 +350,14 @@ def chunk_kda_scaled_dot_kkt_fwd_kernel_intra_sub_intra( m_A = (i_t * BT + i_i * BC + o_i) < T o_A = (bos + i_t * BT + i_i * BC + o_i) * H * BT + i_h * BT + i_i * BC - p_q = tl.make_block_ptr( - q + (bos * H + i_h) * K, - (T, K), - (H * K, 1), - (i_t * BT + i_i * BC, 0), - (BC, BK), - (1, 0), - ) - p_k = tl.make_block_ptr( - k + (bos * H + i_h) * K, - (T, K), - (H * K, 1), - (i_t * BT + i_i * BC, 0), - (BC, BK), - (1, 0), - ) - p_g = tl.make_block_ptr( - g + (bos * H + i_h) * K, - (T, K), - (H * K, 1), - (i_t * BT + i_i * BC, 0), - (BC, BK), - (1, 0), - ) - b_q = tl.load(p_q, boundary_check=(0, 1)) - b_k = tl.load(p_k, boundary_check=(0, 1)) - b_g = tl.load(p_g, boundary_check=(0, 1)) + o_c = i_t * BT + i_i * BC + tl.arange(0, BC) + m_qk = m_A[:, None] & m_k[None, :] + p_q = q + (bos * H + i_h) * K + o_c[:, None] * (H * K) + o_k[None, :] + p_k = k + (bos * H + i_h) * K + o_c[:, None] * (H * K) + o_k[None, :] + p_g = g + (bos * H + i_h) * K + o_c[:, None] * (H * K) + o_k[None, :] + b_q = tl.load(p_q, mask=m_qk, other=0.0) + b_k = tl.load(p_k, mask=m_qk, other=0.0) + b_g = tl.load(p_g, mask=m_qk, other=0.0) p_b = beta + (bos + i_t * BT + i_i * BC + o_i) * H + i_h b_k = b_k * tl.load(p_b, mask=m_A, other=0)[:, None] @@ -536,12 +510,12 @@ def recompute_w_u_fwd_kernel( IS_VARLEN: tl.constexpr, DOT_PRECISION: tl.constexpr, ): - i_t, i_bh = tl.program_id(0), tl.program_id(1) + i_t, i_bh = tl.program_id(0).to(tl.int64), tl.program_id(1) i_b, i_h = i_bh // H, i_bh % H if IS_VARLEN: i_n, i_t = ( tl.load(chunk_indices + i_t * 2).to(tl.int32), - tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), + tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64), ) bos, eos = ( tl.load(cu_seqlens + i_n).to(tl.int32), @@ -550,108 +524,57 @@ 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,)) + o_t = i_t * BT + tl.arange(0, BT) + m_t = o_t < T + p_b = beta + bos * H + i_h + o_t * H + b_b = tl.load(p_b, mask=m_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) - ) - b_A = tl.load(p_A, boundary_check=(0, 1)) + o_A = tl.arange(0, BT) + m_A = m_t[:, None] & (o_A[None, :] < BT) + p_A = A + (bos * H + i_h) * BT + o_t[:, None] * (H * BT) + o_A[None, :] + b_A = tl.load(p_A, mask=m_A, other=0.0) 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_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)) + o_v = i_v * BV + tl.arange(0, BV) + m_v = m_t[:, None] & (o_v[None, :] < V) + p_v = v + (bos * H + i_h) * V + o_t[:, None] * (H * V) + o_v[None, :] + p_u = u + (bos * H + i_h) * V + o_t[:, None] * (H * V) + o_v[None, :] + b_v = tl.load(p_v, mask=m_v, other=0.0) b_vb = (b_v * b_b[:, None]).to(b_v.dtype) b_u = tl.dot(b_A, b_vb, input_precision=DOT_PRECISION) - tl.store(p_u, b_u.to(p_u.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_u, b_u.to(p_u.dtype.element_ty), mask=m_v) for i_k in range(tl.cdiv(K, BK)): - 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), - ) - 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), - ) - b_k = tl.load(p_k, boundary_check=(0, 1)) + o_k = i_k * BK + tl.arange(0, BK) + m_k = o_k < K + m_tk = m_t[:, None] & m_k[None, :] + p_w = w + (bos * H + i_h) * K + o_t[:, None] * (H * K) + o_k[None, :] + p_k = k + (bos * H + i_h) * K + o_t[:, None] * (H * K) + o_k[None, :] + b_k = tl.load(p_k, mask=m_tk, other=0.0) b_kb = b_k * b_b[:, None] - p_gk = tl.make_block_ptr( - gk + (bos * H + i_h) * K, - (T, K), - (H * K, 1), - (i_t * BT, i_k * BK), - (BT, BK), - (1, 0), - ) - b_gk = tl.load(p_gk, boundary_check=(0, 1)) + p_gk = gk + (bos * H + i_h) * K + o_t[:, None] * (H * K) + o_k[None, :] + b_gk = tl.load(p_gk, mask=m_tk, other=0.0) b_kb *= exp2(b_gk) if STORE_QG: - p_q = tl.make_block_ptr( - q + (bos * H + i_h) * K, - (T, K), - (H * K, 1), - (i_t * BT, i_k * BK), - (BT, BK), - (1, 0), - ) - p_qg = tl.make_block_ptr( - qg + (bos * H + i_h) * 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)) + p_q = q + (bos * H + i_h) * K + o_t[:, None] * (H * K) + o_k[None, :] + p_qg = qg + (bos * H + i_h) * K + o_t[:, None] * (H * K) + o_k[None, :] + b_q = tl.load(p_q, mask=m_tk, other=0.0) b_qg = b_q * exp2(b_gk) - tl.store(p_qg, b_qg.to(p_qg.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_qg, b_qg.to(p_qg.dtype.element_ty), mask=m_tk) if STORE_KG: last_idx = min(i_t * BT + BT, T) - 1 - o_k = i_k * BK + tl.arange(0, BK) - m_k = o_k < K b_gn = tl.load( gk + ((bos + last_idx) * H + i_h) * K + o_k, mask=m_k, other=0.0 ) b_kg = b_k * exp2(b_gn - b_gk) - p_kg = tl.make_block_ptr( - kg + (bos * H + i_h) * K, - (T, K), - (H * K, 1), - (i_t * BT, i_k * BK), - (BT, BK), - (1, 0), - ) - tl.store(p_kg, b_kg.to(p_kg.dtype.element_ty), boundary_check=(0, 1)) + p_kg = kg + (bos * H + i_h) * K + o_t[:, None] * (H * K) + o_k[None, :] + tl.store(p_kg, b_kg.to(p_kg.dtype.element_ty), mask=m_tk) 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)) + tl.store(p_w, b_w.to(p_w.dtype.element_ty), mask=m_tk) def recompute_w_u_fwd( @@ -732,13 +655,13 @@ def chunk_gla_fwd_kernel_o( BV: tl.constexpr, IS_VARLEN: tl.constexpr, ): - i_v, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_v, i_t, i_bh = tl.program_id(0), tl.program_id(1).to(tl.int64), tl.program_id(2) i_b, i_h = i_bh // H, i_bh % H if IS_VARLEN: i_tg = i_t i_n, i_t = ( tl.load(chunk_indices + i_t * 2).to(tl.int32), - tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), + tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64), ) bos, eos = ( tl.load(cu_seqlens + i_n).to(tl.int32), @@ -754,70 +677,44 @@ def chunk_gla_fwd_kernel_o( m_s = tl.arange(0, BT)[:, None] >= tl.arange(0, BT)[None, :] b_o = tl.zeros([BT, BV], dtype=tl.float32) + o_t = i_t * BT + tl.arange(0, BT) + o_v = i_v * BV + tl.arange(0, BV) + o_i = tl.arange(0, BT) + m_t = o_t < T + m_v = o_v < V + m_tv = m_t[:, None] & m_v[None, :] + m_A = m_t[:, None] & (o_i[None, :] < BT) for i_k in range(tl.cdiv(K, BK)): - p_q = tl.make_block_ptr( - q + (bos * H + i_h) * K, - (T, K), - (H * K, 1), - (i_t * BT, i_k * BK), - (BT, BK), - (1, 0), - ) - p_g = tl.make_block_ptr( - g + (bos * H + i_h) * K, - (T, K), - (H * K, 1), - (i_t * BT, i_k * BK), - (BT, BK), - (1, 0), - ) - p_h = tl.make_block_ptr( - h + (i_tg * H + i_h) * K * V, - (V, K), - (K, 1), - (i_v * BV, i_k * BK), - (BV, BK), - (1, 0), - ) + o_k = i_k * BK + tl.arange(0, BK) + m_k = o_k < K + m_qk = m_t[:, None] & m_k[None, :] + m_h = m_v[:, None] & m_k[None, :] + p_q = q + (bos * H + i_h) * K + o_t[:, None] * (H * K) + o_k[None, :] + p_g = g + (bos * H + i_h) * K + o_t[:, None] * (H * K) + o_k[None, :] + p_h = h + (i_tg * H + i_h) * K * V + o_v[:, None] * K + o_k[None, :] # [BT, BK] - b_q = tl.load(p_q, boundary_check=(0, 1)) + b_q = tl.load(p_q, mask=m_qk, other=0.0) b_q = (b_q * scale).to(b_q.dtype) # [BT, BK] - b_g = tl.load(p_g, boundary_check=(0, 1)) + b_g = tl.load(p_g, mask=m_qk, other=0.0) # [BT, BK] b_qg = (b_q * exp2(b_g)).to(b_q.dtype) # [BV, BK] - b_h = tl.load(p_h, boundary_check=(0, 1)) + b_h = tl.load(p_h, mask=m_h, other=0.0) # [BT, BV] if i_k >= 0: b_o += tl.dot(b_qg, tl.trans(b_h).to(b_qg.dtype)) - 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_o = tl.make_block_ptr( - o + (bos * H + i_h) * V, - (T, V), - (H * V, 1), - (i_t * BT, i_v * BV), - (BT, BV), - (1, 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_v = v + (bos * H + i_h) * V + o_t[:, None] * (H * V) + o_v[None, :] + p_o = o + (bos * H + i_h) * V + o_t[:, None] * (H * V) + o_v[None, :] + p_A = A + (bos * H + i_h) * BT + o_t[:, None] * (H * BT) + o_i[None, :] # [BT, BV] - b_v = tl.load(p_v, boundary_check=(0, 1)) + b_v = tl.load(p_v, mask=m_tv, other=0.0) # [BT, BT] - b_A = tl.load(p_A, boundary_check=(0, 1)) + b_A = tl.load(p_A, mask=m_A, other=0.0) b_A = tl.where(m_s, b_A, 0.0).to(b_v.dtype) b_o += tl.dot(b_A, b_v, allow_tf32=False) - tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_o, b_o.to(p_o.dtype.element_ty), mask=m_tv) def chunk_gla_fwd_o_gk( @@ -896,12 +793,12 @@ def kda_gate_cumsum_fwd_kernel( HAS_BIAS: tl.constexpr, IS_VARLEN: tl.constexpr, ): - i_d, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_d, i_t, i_bh = tl.program_id(0), tl.program_id(1).to(tl.int64), tl.program_id(2) i_b, i_h = i_bh // H, i_bh % H if IS_VARLEN: i_n, i_t = ( tl.load(chunk_indices + i_t * 2).to(tl.int32), - tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), + tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64), ) bos, eos = ( tl.load(cu_seqlens + i_n).to(tl.int32), @@ -911,26 +808,15 @@ def kda_gate_cumsum_fwd_kernel( else: bos = i_b * T - p_g = tl.make_block_ptr( - g + (bos * H + i_h) * D, - (T, D), - (H * D, 1), - (i_t * BT, i_d * BD), - (BT, BD), - (1, 0), - ) - p_y = tl.make_block_ptr( - y + (bos * H + i_h) * D, - (T, D), - (H * D, 1), - (i_t * BT, i_d * BD), - (BT, BD), - (1, 0), - ) + o_t = i_t * BT + tl.arange(0, BT) + o_d = i_d * BD + tl.arange(0, BD) + m_t = o_t < T + m_g = m_t[:, None] & (o_d[None, :] < D) + p_g = g + (bos * H + i_h) * D + o_t[:, None] * (H * D) + o_d[None, :] + p_y = y + (bos * H + i_h) * D + o_t[:, None] * (H * D) + o_d[None, :] - b_g = tl.load(p_g, boundary_check=(0, 1)).to(tl.float32) + b_g = tl.load(p_g, mask=m_g, other=0.0).to(tl.float32) if HAS_BIAS: - o_d = i_d * BD + tl.arange(0, BD) b_bias = tl.load(g_bias + i_h * D + o_d, mask=o_d < D, other=0.0).to(tl.float32) b_g = b_g + b_bias[None, :] @@ -951,12 +837,12 @@ def kda_gate_cumsum_fwd_kernel( # Out-of-bounds rows (load returns 0, but softplus/bias can still make # b_gate non-zero) participate in the dot product. They only contribute to - # out-of-bounds output rows, which are masked away by `boundary_check` on + # out-of-bounds output rows, which are masked away by `mask` on # the store, so visible output matches unfused gate + chunk-local cumsum. - o_t = tl.arange(0, BT) - m_cumsum = tl.where(o_t[:, None] >= o_t[None, :], 1.0, 0.0) + o_c = tl.arange(0, BT) + m_cumsum = tl.where(o_c[:, None] >= o_c[None, :], 1.0, 0.0) b_y = tl.dot(m_cumsum, b_gate, allow_tf32=False) * cumsum_scale - tl.store(p_y, b_y.to(p_y.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_y, b_y.to(p_y.dtype.element_ty), mask=m_g) def fused_kda_gate_chunk_cumsum( @@ -1271,7 +1157,7 @@ def kda_gate_fwd_kernel( BD: tl.constexpr, HAS_BIAS: tl.constexpr, ): - i_t, i_h = tl.program_id(0), tl.program_id(1) + i_t, i_h = tl.program_id(0).to(tl.int64), tl.program_id(1) n_t = i_t * BT b_a = tl.load(A + i_h).to(tl.float32) @@ -1280,25 +1166,14 @@ def kda_gate_fwd_kernel( stride_row = H * D stride_col = 1 - g_ptr = tl.make_block_ptr( - base=g + i_h * D, - shape=(T, D), - strides=(stride_row, stride_col), - offsets=(n_t, 0), - block_shape=(BT, BD), - order=(1, 0), - ) + o_t = n_t + tl.arange(0, BT) + o_d = tl.arange(0, BD) + m_t = o_t < T + m_g = m_t[:, None] & (o_d[None, :] < D) + g_ptr = g + i_h * D + o_t[:, None] * stride_row + o_d[None, :] * stride_col + y_ptr = y + i_h * D + o_t[:, None] * stride_row + o_d[None, :] * stride_col - y_ptr = tl.make_block_ptr( - base=y + i_h * D, - shape=(T, D), - strides=(stride_row, stride_col), - offsets=(n_t, 0), - block_shape=(BT, BD), - order=(1, 0), - ) - - b_g = tl.load(g_ptr, boundary_check=(0, 1)).to(tl.float32) + b_g = tl.load(g_ptr, mask=m_g, other=0.0).to(tl.float32) if HAS_BIAS: n_d = tl.arange(0, BD) @@ -1321,7 +1196,7 @@ def kda_gate_fwd_kernel( sp = tl.where(use_linear, b_g, (1.0 / beta) * log(1.0 + tl.exp(g_scaled))) b_y = b_a * sp - tl.store(y_ptr, b_y.to(y.dtype.element_ty), boundary_check=(0, 1)) + tl.store(y_ptr, b_y.to(y.dtype.element_ty), mask=m_g) def fused_kda_gate( diff --git a/vllm/models/glm5next/nvidia/ops/third_party/kda/kernels.py b/vllm/models/glm5next/nvidia/ops/third_party/kda/kernels.py index 8f662f68aa..da498b5855 100644 --- a/vllm/models/glm5next/nvidia/ops/third_party/kda/kernels.py +++ b/vllm/models/glm5next/nvidia/ops/third_party/kda/kernels.py @@ -230,13 +230,13 @@ def chunk_kda_scaled_dot_kkt_fwd_kernel_intra_sub_inter( NC: tl.constexpr, IS_VARLEN: tl.constexpr, ): - i_t, i_c, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_t, i_c, i_bh = tl.program_id(0).to(tl.int64), tl.program_id(1), tl.program_id(2) i_b, i_h = i_bh // H, i_bh % H i_i, i_j = i_c // NC, i_c % NC if IS_VARLEN: i_n, i_t = ( tl.load(chunk_indices + i_t * 2).to(tl.int32), - tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), + tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64), ) bos, eos = ( tl.load(cu_seqlens + i_n).to(tl.int32), @@ -257,58 +257,51 @@ def chunk_kda_scaled_dot_kkt_fwd_kernel_intra_sub_inter( A += (bos * H + i_h) * BT Aqk += (bos * H + i_h) * BT - p_b = tl.make_block_ptr( - beta + bos * H + i_h, (T,), (H,), (i_t * BT + i_i * BC,), (BC,), (0,) - ) - b_b = tl.load(p_b, boundary_check=(0,)) + o_i = i_t * BT + i_i * BC + tl.arange(0, BC) + o_j = i_t * BT + i_j * BC + tl.arange(0, BC) + m_i = o_i < T + m_j = o_j < T + p_b = beta + bos * H + i_h + o_i * H + b_b = tl.load(p_b, mask=m_i, other=0.0) b_A = tl.zeros([BC, BC], dtype=tl.float32) b_Aqk = tl.zeros([BC, BC], 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_i * BC, i_k * BK), (BC, BK), (1, 0) - ) - p_k = tl.make_block_ptr( - k, (T, K), (H * K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0) - ) - p_g = tl.make_block_ptr( - g, (T, K), (H * K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0) - ) - b_kt = tl.make_block_ptr( - k, (K, T), (1, H * K), (i_k * BK, i_t * BT + i_j * BC), (BK, BC), (0, 1) - ) - p_gk = tl.make_block_ptr( - g, (K, T), (1, H * K), (i_k * BK, i_t * BT + i_j * BC), (BK, BC), (0, 1) - ) - o_k = i_k * BK + tl.arange(0, BK) m_k = o_k < K + m_qk = m_i[:, None] & m_k[None, :] + m_kj = m_k[:, None] & m_j[None, :] + + p_q = q + o_i[:, None] * (H * K) + o_k[None, :] + p_k = k + o_i[:, None] * (H * K) + o_k[None, :] + p_g = g + o_i[:, None] * (H * K) + o_k[None, :] + b_kt = k + o_k[:, None] + o_j[None, :] * (H * K) + p_gk = g + o_k[:, None] + o_j[None, :] * (H * K) + # [BK,] b_gn = tl.load(g + (i_t * BT + i_i * BC) * H * K + o_k, mask=m_k, other=0) # [BC, BK] - b_g = tl.load(p_g, boundary_check=(0, 1)) - b_k = tl.load(p_k, boundary_check=(0, 1)) * exp2(b_g - b_gn[None, :]) + b_g = tl.load(p_g, mask=m_qk, other=0.0) + b_k = tl.load(p_k, mask=m_qk, other=0.0) * exp2(b_g - b_gn[None, :]) # [BK, BC] - b_gk = tl.load(p_gk, boundary_check=(0, 1)) - b_kt = tl.load(b_kt, boundary_check=(0, 1)) + b_gk = tl.load(p_gk, mask=m_kj, other=0.0) + b_kt = tl.load(b_kt, mask=m_kj, other=0.0) # [BC, BC] b_ktg = b_kt * exp2(b_gn[:, None] - b_gk) b_A += tl.dot(b_k, b_ktg) - b_q = tl.load(p_q, boundary_check=(0, 1)) + b_q = tl.load(p_q, mask=m_qk, other=0.0) b_qg = b_q * exp2(b_g - b_gn[None, :]) * scale b_Aqk += tl.dot(b_qg, b_ktg) b_A *= b_b[:, None] - p_A = tl.make_block_ptr( - A, (T, BT), (H * BT, 1), (i_t * BT + i_i * BC, i_j * BC), (BC, BC), (1, 0) - ) - tl.store(p_A, b_A.to(A.dtype.element_ty), boundary_check=(0, 1)) - p_Aqk = tl.make_block_ptr( - Aqk, (T, BT), (H * BT, 1), (i_t * BT + i_i * BC, i_j * BC), (BC, BC), (1, 0) - ) - tl.store(p_Aqk, b_Aqk.to(Aqk.dtype.element_ty), boundary_check=(0, 1)) + o_jA = i_j * BC + tl.arange(0, BC) + m_A = m_i[:, None] & (o_jA[None, :] < BT) + p_A = A + o_i[:, None] * (H * BT) + o_jA[None, :] + tl.store(p_A, b_A.to(A.dtype.element_ty), mask=m_A) + p_Aqk = Aqk + o_i[:, None] * (H * BT) + o_jA[None, :] + tl.store(p_Aqk, b_Aqk.to(Aqk.dtype.element_ty), mask=m_A) @triton.heuristics({"IS_VARLEN": lambda args: args["cu_seqlens"] is not None}) @@ -335,12 +328,12 @@ def chunk_kda_scaled_dot_kkt_fwd_kernel_intra_sub_intra( BK: tl.constexpr, IS_VARLEN: tl.constexpr, ): - i_t, i_i, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_t, i_i, i_bh = tl.program_id(0).to(tl.int64), tl.program_id(1), tl.program_id(2) i_b, i_h = i_bh // H, i_bh % H if IS_VARLEN: i_n, i_t = ( tl.load(chunk_indices + i_t * 2).to(tl.int32), - tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), + tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64), ) bos, eos = ( tl.load(cu_seqlens + i_n).to(tl.int32), @@ -359,33 +352,14 @@ def chunk_kda_scaled_dot_kkt_fwd_kernel_intra_sub_intra( m_A = (i_t * BT + i_i * BC + o_i) < T o_A = (bos + i_t * BT + i_i * BC + o_i) * H * BT + i_h * BT + i_i * BC - p_q = tl.make_block_ptr( - q + (bos * H + i_h) * K, - (T, K), - (H * K, 1), - (i_t * BT + i_i * BC, 0), - (BC, BK), - (1, 0), - ) - p_k = tl.make_block_ptr( - k + (bos * H + i_h) * K, - (T, K), - (H * K, 1), - (i_t * BT + i_i * BC, 0), - (BC, BK), - (1, 0), - ) - p_g = tl.make_block_ptr( - g + (bos * H + i_h) * K, - (T, K), - (H * K, 1), - (i_t * BT + i_i * BC, 0), - (BC, BK), - (1, 0), - ) - b_q = tl.load(p_q, boundary_check=(0, 1)) - b_k = tl.load(p_k, boundary_check=(0, 1)) - b_g = tl.load(p_g, boundary_check=(0, 1)) + o_c = i_t * BT + i_i * BC + tl.arange(0, BC) + m_qk = m_A[:, None] & m_k[None, :] + p_q = q + (bos * H + i_h) * K + o_c[:, None] * (H * K) + o_k[None, :] + p_k = k + (bos * H + i_h) * K + o_c[:, None] * (H * K) + o_k[None, :] + p_g = g + (bos * H + i_h) * K + o_c[:, None] * (H * K) + o_k[None, :] + b_q = tl.load(p_q, mask=m_qk, other=0.0) + b_k = tl.load(p_k, mask=m_qk, other=0.0) + b_g = tl.load(p_g, mask=m_qk, other=0.0) p_b = beta + (bos + i_t * BT + i_i * BC + o_i) * H + i_h b_k = b_k * tl.load(p_b, mask=m_A, other=0)[:, None] @@ -538,12 +512,12 @@ def recompute_w_u_fwd_kernel( IS_VARLEN: tl.constexpr, DOT_PRECISION: tl.constexpr, ): - i_t, i_bh = tl.program_id(0), tl.program_id(1) + i_t, i_bh = tl.program_id(0).to(tl.int64), tl.program_id(1) i_b, i_h = i_bh // H, i_bh % H if IS_VARLEN: i_n, i_t = ( tl.load(chunk_indices + i_t * 2).to(tl.int32), - tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), + tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64), ) bos, eos = ( tl.load(cu_seqlens + i_n).to(tl.int32), @@ -552,108 +526,57 @@ 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,)) + o_t = i_t * BT + tl.arange(0, BT) + m_t = o_t < T + p_b = beta + bos * H + i_h + o_t * H + b_b = tl.load(p_b, mask=m_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) - ) - b_A = tl.load(p_A, boundary_check=(0, 1)) + o_A = tl.arange(0, BT) + m_A = m_t[:, None] & (o_A[None, :] < BT) + p_A = A + (bos * H + i_h) * BT + o_t[:, None] * (H * BT) + o_A[None, :] + b_A = tl.load(p_A, mask=m_A, other=0.0) 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_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)) + o_v = i_v * BV + tl.arange(0, BV) + m_v = m_t[:, None] & (o_v[None, :] < V) + p_v = v + (bos * H + i_h) * V + o_t[:, None] * (H * V) + o_v[None, :] + p_u = u + (bos * H + i_h) * V + o_t[:, None] * (H * V) + o_v[None, :] + b_v = tl.load(p_v, mask=m_v, other=0.0) b_vb = (b_v * b_b[:, None]).to(b_v.dtype) b_u = tl.dot(b_A, b_vb, input_precision=DOT_PRECISION) - tl.store(p_u, b_u.to(p_u.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_u, b_u.to(p_u.dtype.element_ty), mask=m_v) for i_k in range(tl.cdiv(K, BK)): - 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), - ) - 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), - ) - b_k = tl.load(p_k, boundary_check=(0, 1)) + o_k = i_k * BK + tl.arange(0, BK) + m_k = o_k < K + m_tk = m_t[:, None] & m_k[None, :] + p_w = w + (bos * H + i_h) * K + o_t[:, None] * (H * K) + o_k[None, :] + p_k = k + (bos * H + i_h) * K + o_t[:, None] * (H * K) + o_k[None, :] + b_k = tl.load(p_k, mask=m_tk, other=0.0) b_kb = b_k * b_b[:, None] - p_gk = tl.make_block_ptr( - gk + (bos * H + i_h) * K, - (T, K), - (H * K, 1), - (i_t * BT, i_k * BK), - (BT, BK), - (1, 0), - ) - b_gk = tl.load(p_gk, boundary_check=(0, 1)) + p_gk = gk + (bos * H + i_h) * K + o_t[:, None] * (H * K) + o_k[None, :] + b_gk = tl.load(p_gk, mask=m_tk, other=0.0) b_kb *= exp2(b_gk) if STORE_QG: - p_q = tl.make_block_ptr( - q + (bos * H + i_h) * K, - (T, K), - (H * K, 1), - (i_t * BT, i_k * BK), - (BT, BK), - (1, 0), - ) - p_qg = tl.make_block_ptr( - qg + (bos * H + i_h) * 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)) + p_q = q + (bos * H + i_h) * K + o_t[:, None] * (H * K) + o_k[None, :] + p_qg = qg + (bos * H + i_h) * K + o_t[:, None] * (H * K) + o_k[None, :] + b_q = tl.load(p_q, mask=m_tk, other=0.0) b_qg = b_q * exp2(b_gk) - tl.store(p_qg, b_qg.to(p_qg.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_qg, b_qg.to(p_qg.dtype.element_ty), mask=m_tk) if STORE_KG: last_idx = min(i_t * BT + BT, T) - 1 - o_k = i_k * BK + tl.arange(0, BK) - m_k = o_k < K b_gn = tl.load( gk + ((bos + last_idx) * H + i_h) * K + o_k, mask=m_k, other=0.0 ) b_kg = b_k * exp2(b_gn - b_gk) - p_kg = tl.make_block_ptr( - kg + (bos * H + i_h) * K, - (T, K), - (H * K, 1), - (i_t * BT, i_k * BK), - (BT, BK), - (1, 0), - ) - tl.store(p_kg, b_kg.to(p_kg.dtype.element_ty), boundary_check=(0, 1)) + p_kg = kg + (bos * H + i_h) * K + o_t[:, None] * (H * K) + o_k[None, :] + tl.store(p_kg, b_kg.to(p_kg.dtype.element_ty), mask=m_tk) 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)) + tl.store(p_w, b_w.to(p_w.dtype.element_ty), mask=m_tk) def recompute_w_u_fwd( @@ -734,13 +657,13 @@ def chunk_gla_fwd_kernel_o( BV: tl.constexpr, IS_VARLEN: tl.constexpr, ): - i_v, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_v, i_t, i_bh = tl.program_id(0), tl.program_id(1).to(tl.int64), tl.program_id(2) i_b, i_h = i_bh // H, i_bh % H if IS_VARLEN: i_tg = i_t i_n, i_t = ( tl.load(chunk_indices + i_t * 2).to(tl.int32), - tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), + tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64), ) bos, eos = ( tl.load(cu_seqlens + i_n).to(tl.int32), @@ -756,70 +679,44 @@ def chunk_gla_fwd_kernel_o( m_s = tl.arange(0, BT)[:, None] >= tl.arange(0, BT)[None, :] b_o = tl.zeros([BT, BV], dtype=tl.float32) + o_t = i_t * BT + tl.arange(0, BT) + o_v = i_v * BV + tl.arange(0, BV) + o_i = tl.arange(0, BT) + m_t = o_t < T + m_v = o_v < V + m_tv = m_t[:, None] & m_v[None, :] + m_A = m_t[:, None] & (o_i[None, :] < BT) for i_k in range(tl.cdiv(K, BK)): - p_q = tl.make_block_ptr( - q + (bos * H + i_h) * K, - (T, K), - (H * K, 1), - (i_t * BT, i_k * BK), - (BT, BK), - (1, 0), - ) - p_g = tl.make_block_ptr( - g + (bos * H + i_h) * K, - (T, K), - (H * K, 1), - (i_t * BT, i_k * BK), - (BT, BK), - (1, 0), - ) - p_h = tl.make_block_ptr( - h + (i_tg * H + i_h) * K * V, - (V, K), - (K, 1), - (i_v * BV, i_k * BK), - (BV, BK), - (1, 0), - ) + o_k = i_k * BK + tl.arange(0, BK) + m_k = o_k < K + m_qk = m_t[:, None] & m_k[None, :] + m_h = m_v[:, None] & m_k[None, :] + p_q = q + (bos * H + i_h) * K + o_t[:, None] * (H * K) + o_k[None, :] + p_g = g + (bos * H + i_h) * K + o_t[:, None] * (H * K) + o_k[None, :] + p_h = h + (i_tg * H + i_h) * K * V + o_v[:, None] * K + o_k[None, :] # [BT, BK] - b_q = tl.load(p_q, boundary_check=(0, 1)) + b_q = tl.load(p_q, mask=m_qk, other=0.0) b_q = (b_q * scale).to(b_q.dtype) # [BT, BK] - b_g = tl.load(p_g, boundary_check=(0, 1)) + b_g = tl.load(p_g, mask=m_qk, other=0.0) # [BT, BK] b_qg = (b_q * exp2(b_g)).to(b_q.dtype) # [BV, BK] - b_h = tl.load(p_h, boundary_check=(0, 1)) + b_h = tl.load(p_h, mask=m_h, other=0.0) # [BT, BV] if i_k >= 0: b_o += tl.dot(b_qg, tl.trans(b_h).to(b_qg.dtype)) - 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_o = tl.make_block_ptr( - o + (bos * H + i_h) * V, - (T, V), - (H * V, 1), - (i_t * BT, i_v * BV), - (BT, BV), - (1, 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_v = v + (bos * H + i_h) * V + o_t[:, None] * (H * V) + o_v[None, :] + p_o = o + (bos * H + i_h) * V + o_t[:, None] * (H * V) + o_v[None, :] + p_A = A + (bos * H + i_h) * BT + o_t[:, None] * (H * BT) + o_i[None, :] # [BT, BV] - b_v = tl.load(p_v, boundary_check=(0, 1)) + b_v = tl.load(p_v, mask=m_tv, other=0.0) # [BT, BT] - b_A = tl.load(p_A, boundary_check=(0, 1)) + b_A = tl.load(p_A, mask=m_A, other=0.0) b_A = tl.where(m_s, b_A, 0.0).to(b_v.dtype) b_o += tl.dot(b_A, b_v, allow_tf32=False) - tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_o, b_o.to(p_o.dtype.element_ty), mask=m_tv) def chunk_gla_fwd_o_gk( @@ -898,12 +795,12 @@ def kda_gate_cumsum_fwd_kernel( HAS_BIAS: tl.constexpr, IS_VARLEN: tl.constexpr, ): - i_d, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_d, i_t, i_bh = tl.program_id(0), tl.program_id(1).to(tl.int64), tl.program_id(2) i_b, i_h = i_bh // H, i_bh % H if IS_VARLEN: i_n, i_t = ( tl.load(chunk_indices + i_t * 2).to(tl.int32), - tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), + tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64), ) bos, eos = ( tl.load(cu_seqlens + i_n).to(tl.int32), @@ -913,26 +810,15 @@ def kda_gate_cumsum_fwd_kernel( else: bos = i_b * T - p_g = tl.make_block_ptr( - g + (bos * H + i_h) * D, - (T, D), - (H * D, 1), - (i_t * BT, i_d * BD), - (BT, BD), - (1, 0), - ) - p_y = tl.make_block_ptr( - y + (bos * H + i_h) * D, - (T, D), - (H * D, 1), - (i_t * BT, i_d * BD), - (BT, BD), - (1, 0), - ) + o_t = i_t * BT + tl.arange(0, BT) + o_d = i_d * BD + tl.arange(0, BD) + m_t = o_t < T + m_g = m_t[:, None] & (o_d[None, :] < D) + p_g = g + (bos * H + i_h) * D + o_t[:, None] * (H * D) + o_d[None, :] + p_y = y + (bos * H + i_h) * D + o_t[:, None] * (H * D) + o_d[None, :] - b_g = tl.load(p_g, boundary_check=(0, 1)).to(tl.float32) + b_g = tl.load(p_g, mask=m_g, other=0.0).to(tl.float32) if HAS_BIAS: - o_d = i_d * BD + tl.arange(0, BD) b_bias = tl.load(g_bias + i_h * D + o_d, mask=o_d < D, other=0.0).to(tl.float32) b_g = b_g + b_bias[None, :] @@ -953,12 +839,12 @@ def kda_gate_cumsum_fwd_kernel( # Out-of-bounds rows (load returns 0, but softplus/bias can still make # b_gate non-zero) participate in the dot product. They only contribute to - # out-of-bounds output rows, which are masked away by `boundary_check` on + # out-of-bounds output rows, which are masked away by `mask` on # the store, so visible output matches unfused gate + chunk-local cumsum. - o_t = tl.arange(0, BT) - m_cumsum = tl.where(o_t[:, None] >= o_t[None, :], 1.0, 0.0) + o_c = tl.arange(0, BT) + m_cumsum = tl.where(o_c[:, None] >= o_c[None, :], 1.0, 0.0) b_y = tl.dot(m_cumsum, b_gate, allow_tf32=False) * cumsum_scale - tl.store(p_y, b_y.to(p_y.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_y, b_y.to(p_y.dtype.element_ty), mask=m_g) def fused_kda_gate_chunk_cumsum( @@ -1273,7 +1159,7 @@ def kda_gate_fwd_kernel( BD: tl.constexpr, HAS_BIAS: tl.constexpr, ): - i_t, i_h = tl.program_id(0), tl.program_id(1) + i_t, i_h = tl.program_id(0).to(tl.int64), tl.program_id(1) n_t = i_t * BT b_a = tl.load(A + i_h).to(tl.float32) @@ -1282,25 +1168,14 @@ def kda_gate_fwd_kernel( stride_row = H * D stride_col = 1 - g_ptr = tl.make_block_ptr( - base=g + i_h * D, - shape=(T, D), - strides=(stride_row, stride_col), - offsets=(n_t, 0), - block_shape=(BT, BD), - order=(1, 0), - ) + o_t = n_t + tl.arange(0, BT) + o_d = tl.arange(0, BD) + m_t = o_t < T + m_g = m_t[:, None] & (o_d[None, :] < D) + g_ptr = g + i_h * D + o_t[:, None] * stride_row + o_d[None, :] * stride_col + y_ptr = y + i_h * D + o_t[:, None] * stride_row + o_d[None, :] * stride_col - y_ptr = tl.make_block_ptr( - base=y + i_h * D, - shape=(T, D), - strides=(stride_row, stride_col), - offsets=(n_t, 0), - block_shape=(BT, BD), - order=(1, 0), - ) - - b_g = tl.load(g_ptr, boundary_check=(0, 1)).to(tl.float32) + b_g = tl.load(g_ptr, mask=m_g, other=0.0).to(tl.float32) if HAS_BIAS: n_d = tl.arange(0, BD) @@ -1323,7 +1198,7 @@ def kda_gate_fwd_kernel( sp = tl.where(use_linear, b_g, (1.0 / beta) * log(1.0 + tl.exp(g_scaled))) b_y = b_a * sp - tl.store(y_ptr, b_y.to(y.dtype.element_ty), boundary_check=(0, 1)) + tl.store(y_ptr, b_y.to(y.dtype.element_ty), mask=m_g) def fused_kda_gate( diff --git a/vllm/models/inkling/amd/ops/fa4_rel_attention.py b/vllm/models/inkling/amd/ops/fa4_rel_attention.py index d448ac8e2c..936b51eb6e 100644 --- a/vllm/models/inkling/amd/ops/fa4_rel_attention.py +++ b/vllm/models/inkling/amd/ops/fa4_rel_attention.py @@ -165,15 +165,23 @@ def _inkling_rel_attention_kernel( q_head_start = pid_kh * gqa_group_size bt_row = block_table_ptr + pid_b * stride_bt_b - q_block_ptr = tl.make_block_ptr( - base=q_ptr + q_start * stride_q_t + q_head_start * stride_q_h, - shape=(q_len, gqa_group_size, head_dim), - strides=(stride_q_t, stride_q_h, stride_q_d), - offsets=(q_block, 0, 0), - block_shape=(BLOCK_Q, BLOCK_H, BLOCK_D), - order=(2, 1, 0), + o0 = q_block + tl.arange(0, BLOCK_Q) + o1 = tl.arange(0, BLOCK_H) + o2 = tl.arange(0, BLOCK_D) + m_q = ( + (o0 < q_len)[:, None, None] + & (o1 < gqa_group_size)[None, :, None] + & (o2 < head_dim)[None, None, :] ) - q = tl.load(q_block_ptr, boundary_check=(0, 1, 2), padding_option="zero") + q_block_ptr = ( + q_ptr + + q_start * stride_q_t + + q_head_start * stride_q_h + + o0[:, None, None] * stride_q_t + + o1[None, :, None] * stride_q_h + + o2[None, None, :] * stride_q_d + ) + q = tl.load(q_block_ptr, mask=m_q, other=0.0) q = tl.reshape(q, (BLOCK_QH, BLOCK_D)) q_rows = q_block + tl.arange(0, BLOCK_Q) @@ -265,18 +273,18 @@ def _inkling_rel_attention_kernel( acc /= l_i[:, None] acc = tl.reshape(acc, (BLOCK_Q, BLOCK_H, BLOCK_D)) - out_block_ptr = tl.make_block_ptr( - base=out_ptr + q_start * stride_o_t + q_head_start * stride_o_h, - shape=(q_len, gqa_group_size, head_dim), - strides=(stride_o_t, stride_o_h, stride_o_d), - offsets=(q_block, 0, 0), - block_shape=(BLOCK_Q, BLOCK_H, BLOCK_D), - order=(2, 1, 0), + out_block_ptr = ( + out_ptr + + q_start * stride_o_t + + q_head_start * stride_o_h + + o0[:, None, None] * stride_o_t + + o1[None, :, None] * stride_o_h + + o2[None, None, :] * stride_o_d ) tl.store( out_block_ptr, acc.to(out_ptr.dtype.element_ty), - boundary_check=(0, 1, 2), + mask=m_q, ) diff --git a/vllm/models/kimi_k3/nvidia/ops/third_party/kda/chunk.py b/vllm/models/kimi_k3/nvidia/ops/third_party/kda/chunk.py index 3911d20d03..19fcd9a686 100644 --- a/vllm/models/kimi_k3/nvidia/ops/third_party/kda/chunk.py +++ b/vllm/models/kimi_k3/nvidia/ops/third_party/kda/chunk.py @@ -69,12 +69,12 @@ def recompute_w_u_fwd_kernel( IS_VARLEN: tl.constexpr, DOT_PRECISION: tl.constexpr, ): - i_t, i_bh = tl.program_id(0), tl.program_id(1) + i_t, i_bh = tl.program_id(0).to(tl.int64), tl.program_id(1) i_b, i_h = i_bh // H, i_bh % H if IS_VARLEN: i_n, i_t = ( tl.load(chunk_indices + i_t * 2).to(tl.int32), - tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), + tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64), ) bos, eos = ( tl.load(cu_seqlens + i_n).to(tl.int32), @@ -83,108 +83,57 @@ 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,)).to(tl.float32) + o_t = i_t * BT + tl.arange(0, BT) + m_t = o_t < T + p_b = beta + bos * H + i_h + o_t * H + b_b = tl.load(p_b, mask=m_t, other=0.0).to(tl.float32) - 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) - ) - b_A = tl.load(p_A, boundary_check=(0, 1)) + o_A = tl.arange(0, BT) + m_A = m_t[:, None] & (o_A[None, :] < BT) + p_A = A + (bos * H + i_h) * BT + o_t[:, None] * (H * BT) + o_A[None, :] + b_A = tl.load(p_A, mask=m_A, other=0.0) 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_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)) + o_v = i_v * BV + tl.arange(0, BV) + m_v = m_t[:, None] & (o_v[None, :] < V) + p_v = v + (bos * H + i_h) * V + o_t[:, None] * (H * V) + o_v[None, :] + p_u = u + (bos * H + i_h) * V + o_t[:, None] * (H * V) + o_v[None, :] + b_v = tl.load(p_v, mask=m_v, other=0.0) b_vb = (b_v * b_b[:, None]).to(b_v.dtype) b_u = tl.dot(b_A, b_vb, input_precision=DOT_PRECISION) - tl.store(p_u, b_u.to(p_u.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_u, b_u.to(p_u.dtype.element_ty), mask=m_v) for i_k in range(tl.cdiv(K, BK)): - 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), - ) - 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), - ) - b_k = tl.load(p_k, boundary_check=(0, 1)) + o_k = i_k * BK + tl.arange(0, BK) + m_k = o_k < K + m_tk = m_t[:, None] & m_k[None, :] + p_w = w + (bos * H + i_h) * K + o_t[:, None] * (H * K) + o_k[None, :] + p_k = k + (bos * H + i_h) * K + o_t[:, None] * (H * K) + o_k[None, :] + b_k = tl.load(p_k, mask=m_tk, other=0.0) b_kb = b_k * b_b[:, None] - p_gk = tl.make_block_ptr( - gk + (bos * H + i_h) * K, - (T, K), - (H * K, 1), - (i_t * BT, i_k * BK), - (BT, BK), - (1, 0), - ) - b_gk = tl.load(p_gk, boundary_check=(0, 1)) + p_gk = gk + (bos * H + i_h) * K + o_t[:, None] * (H * K) + o_k[None, :] + b_gk = tl.load(p_gk, mask=m_tk, other=0.0) b_kb *= exp2(b_gk) if STORE_QG: - p_q = tl.make_block_ptr( - q + (bos * H + i_h) * K, - (T, K), - (H * K, 1), - (i_t * BT, i_k * BK), - (BT, BK), - (1, 0), - ) - p_qg = tl.make_block_ptr( - qg + (bos * H + i_h) * 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)) + p_q = q + (bos * H + i_h) * K + o_t[:, None] * (H * K) + o_k[None, :] + p_qg = qg + (bos * H + i_h) * K + o_t[:, None] * (H * K) + o_k[None, :] + b_q = tl.load(p_q, mask=m_tk, other=0.0) b_qg = b_q * exp2(b_gk) - tl.store(p_qg, b_qg.to(p_qg.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_qg, b_qg.to(p_qg.dtype.element_ty), mask=m_tk) if STORE_KG: last_idx = min(i_t * BT + BT, T) - 1 - o_k = i_k * BK + tl.arange(0, BK) - m_k = o_k < K b_gn = tl.load( gk + ((bos + last_idx) * H + i_h) * K + o_k, mask=m_k, other=0.0 ) b_kg = b_k * exp2(b_gn - b_gk) - p_kg = tl.make_block_ptr( - kg + (bos * H + i_h) * K, - (T, K), - (H * K, 1), - (i_t * BT, i_k * BK), - (BT, BK), - (1, 0), - ) - tl.store(p_kg, b_kg.to(p_kg.dtype.element_ty), boundary_check=(0, 1)) + p_kg = kg + (bos * H + i_h) * K + o_t[:, None] * (H * K) + o_k[None, :] + tl.store(p_kg, b_kg.to(p_kg.dtype.element_ty), mask=m_tk) 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)) + tl.store(p_w, b_w.to(p_w.dtype.element_ty), mask=m_tk) def recompute_w_u_fwd( @@ -265,13 +214,13 @@ def chunk_gla_fwd_kernel_o( BV: tl.constexpr, IS_VARLEN: tl.constexpr, ): - i_v, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_v, i_t, i_bh = tl.program_id(0), tl.program_id(1).to(tl.int64), tl.program_id(2) i_b, i_h = i_bh // H, i_bh % H if IS_VARLEN: i_tg = i_t i_n, i_t = ( tl.load(chunk_indices + i_t * 2).to(tl.int32), - tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), + tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64), ) bos, eos = ( tl.load(cu_seqlens + i_n).to(tl.int32), @@ -287,70 +236,44 @@ def chunk_gla_fwd_kernel_o( m_s = tl.arange(0, BT)[:, None] >= tl.arange(0, BT)[None, :] b_o = tl.zeros([BT, BV], dtype=tl.float32) + o_t = i_t * BT + tl.arange(0, BT) + o_v = i_v * BV + tl.arange(0, BV) + o_i = tl.arange(0, BT) + m_t = o_t < T + m_v = o_v < V + m_tv = m_t[:, None] & m_v[None, :] + m_A = m_t[:, None] & (o_i[None, :] < BT) for i_k in range(tl.cdiv(K, BK)): - p_q = tl.make_block_ptr( - q + (bos * H + i_h) * K, - (T, K), - (H * K, 1), - (i_t * BT, i_k * BK), - (BT, BK), - (1, 0), - ) - p_g = tl.make_block_ptr( - g + (bos * H + i_h) * K, - (T, K), - (H * K, 1), - (i_t * BT, i_k * BK), - (BT, BK), - (1, 0), - ) - p_h = tl.make_block_ptr( - h + (i_tg * H + i_h) * K * V, - (V, K), - (K, 1), - (i_v * BV, i_k * BK), - (BV, BK), - (1, 0), - ) + o_k = i_k * BK + tl.arange(0, BK) + m_k = o_k < K + m_qk = m_t[:, None] & m_k[None, :] + p_q = q + (bos * H + i_h) * K + o_t[:, None] * (H * K) + o_k[None, :] + p_g = g + (bos * H + i_h) * K + o_t[:, None] * (H * K) + o_k[None, :] + p_h = h + (i_tg * H + i_h) * K * V + o_v[:, None] * K + o_k[None, :] + m_h = m_v[:, None] & m_k[None, :] # [BT, BK] - b_q = tl.load(p_q, boundary_check=(0, 1)) + b_q = tl.load(p_q, mask=m_qk, other=0.0) b_q = (b_q * scale).to(b_q.dtype) # [BT, BK] - b_g = tl.load(p_g, boundary_check=(0, 1)) + b_g = tl.load(p_g, mask=m_qk, other=0.0) # [BT, BK] b_qg = (b_q * exp2(b_g)).to(b_q.dtype) # [BV, BK] - b_h = tl.load(p_h, boundary_check=(0, 1)) + b_h = tl.load(p_h, mask=m_h, other=0.0) # [BT, BV] if i_k >= 0: b_o += tl.dot(b_qg, tl.trans(b_h).to(b_qg.dtype)) - 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_o = tl.make_block_ptr( - o + (bos * H + i_h) * V, - (T, V), - (H * V, 1), - (i_t * BT, i_v * BV), - (BT, BV), - (1, 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_v = v + (bos * H + i_h) * V + o_t[:, None] * (H * V) + o_v[None, :] + p_o = o + (bos * H + i_h) * V + o_t[:, None] * (H * V) + o_v[None, :] + p_A = A + (bos * H + i_h) * BT + o_t[:, None] * (H * BT) + o_i[None, :] # [BT, BV] - b_v = tl.load(p_v, boundary_check=(0, 1)) + b_v = tl.load(p_v, mask=m_tv, other=0.0) # [BT, BT] - b_A = tl.load(p_A, boundary_check=(0, 1)) + b_A = tl.load(p_A, mask=m_A, other=0.0) b_A = tl.where(m_s, b_A, 0.0).to(b_v.dtype) b_o += tl.dot(b_A, b_v, allow_tf32=False) - tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_o, b_o.to(p_o.dtype.element_ty), mask=m_tv) def chunk_gla_fwd_o_gk( @@ -434,12 +357,12 @@ def kda_gate_chunk_cumsum_vector_kernel( IS_VARLEN: tl.constexpr, USE_LOWER_BOUND: tl.constexpr, ): - i_s, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_s, i_t, i_bh = tl.program_id(0), tl.program_id(1).to(tl.int64), tl.program_id(2) i_b, i_h = i_bh // H, i_bh % H if IS_VARLEN: i_n, i_t = ( tl.load(chunk_indices + i_t * 2).to(tl.int32), - tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), + tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64), ) bos, eos = ( tl.load(cu_seqlens + i_n).to(tl.int32), @@ -472,34 +395,15 @@ def kda_gate_chunk_cumsum_vector_kernel( i_s -= 1 - 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), - ) + o_t = i_t * BT + tl.arange(0, BT) + o_s = i_s * BS + tl.arange(0, BS) + m_s = (o_t[:, None] < T) & (o_s[None, :] < S) + p_s = s + (bos * H + i_h) * S + o_t[:, None] * (H * S) + o_s[None, :] + p_o = o + (bos * H + i_h) * S + o_t[:, None] * (H * S) + o_s[None, :] - b_s = tl.load(p_s, boundary_check=(0, 1)).to(tl.float32) + b_s = tl.load(p_s, mask=m_s, other=0.0).to(tl.float32) if HAS_BIAS: - p_bias = tl.make_block_ptr( - g_bias + i_h * S, - (S,), - (1,), - (i_s * BS,), - (BS,), - (0,), - ) - b_bias = tl.load(p_bias, boundary_check=(0,)).to(tl.float32) + b_bias = tl.load(g_bias + i_h * S + o_s, mask=o_s < S, other=0.0).to(tl.float32) b_s += b_bias[None, :] b_a = tl.exp(tl.load(A_log + i_h).to(tl.float32)) @@ -517,7 +421,7 @@ def kda_gate_chunk_cumsum_vector_kernel( # Boundary loads return zero, but bias and gate activation can make padded # rows nonzero. Padding trails valid rows, so it only affects masked stores. b_o = tl.cumsum(b_gate, axis=0) * cumsum_scale - tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_o, b_o.to(p_o.dtype.element_ty), mask=m_s) def fused_kda_gate_chunk_cumsum( @@ -843,33 +747,19 @@ def kda_gate_fwd_kernel( HAS_BIAS: tl.constexpr, USE_LOWER_BOUND: tl.constexpr, ): - i_t, i_h = tl.program_id(0), tl.program_id(1) + i_t, i_h = tl.program_id(0).to(tl.int64), tl.program_id(1) n_t = i_t * BT b_a = tl.exp(tl.load(A + i_h).to(tl.float32)) - stride_row = H * D - stride_col = 1 + o_t = n_t + tl.arange(0, BT) + o_d = tl.arange(0, BD) + m_t = o_t < T + m_g = m_t[:, None] & (o_d[None, :] < D) + g_ptr = g + i_h * D + o_t[:, None] * (H * D) + o_d[None, :] + y_ptr = y + i_h * D + o_t[:, None] * (H * D) + o_d[None, :] - g_ptr = tl.make_block_ptr( - base=g + i_h * D, - shape=(T, D), - strides=(stride_row, stride_col), - offsets=(n_t, 0), - block_shape=(BT, BD), - order=(1, 0), - ) - - y_ptr = tl.make_block_ptr( - base=y + i_h * D, - shape=(T, D), - strides=(stride_row, stride_col), - offsets=(n_t, 0), - block_shape=(BT, BD), - order=(1, 0), - ) - - b_g = tl.load(g_ptr, boundary_check=(0, 1)).to(tl.float32) + b_g = tl.load(g_ptr, mask=m_g, other=0.0).to(tl.float32) if HAS_BIAS: n_d = tl.arange(0, BD) @@ -887,7 +777,7 @@ def kda_gate_fwd_kernel( sp = tl.where(use_linear, b_g, (1.0 / beta) * log(1.0 + tl.exp(g_scaled))) b_y = -b_a * sp - tl.store(y_ptr, b_y.to(y.dtype.element_ty), boundary_check=(0, 1)) + tl.store(y_ptr, b_y.to(y.dtype.element_ty), mask=m_g) def fused_kda_gate( diff --git a/vllm/models/kimi_k3/nvidia/ops/third_party/kda/chunk_intra.py b/vllm/models/kimi_k3/nvidia/ops/third_party/kda/chunk_intra.py index 07b4a38893..286d23c814 100644 --- a/vllm/models/kimi_k3/nvidia/ops/third_party/kda/chunk_intra.py +++ b/vllm/models/kimi_k3/nvidia/ops/third_party/kda/chunk_intra.py @@ -70,14 +70,14 @@ def chunk_kda_fwd_kernel_inter_solve_fused( 5. Computes merged Akk_inv 6. Writes Akk_inv to Akk """ - i_t, i_bh = tl.program_id(0), tl.program_id(1) + i_t, i_bh = tl.program_id(0).to(tl.int64), tl.program_id(1) i_b, i_hv = i_bh // HV, i_bh % HV i_h = i_hv // (HV // H) if IS_VARLEN: i_n, i_t = ( tl.load(chunk_indices + i_t * 2).to(tl.int32), - tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), + tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64), ) bos, eos = ( tl.load(cu_seqlens + i_n).to(tl.int32), @@ -106,6 +106,15 @@ def chunk_kda_fwd_kernel_inter_solve_fused( m_tc1 = (i_tc1 + o_i) < T m_tc2 = (i_tc2 + o_i) < T m_tc3 = (i_tc3 + o_i) < T + o_c0 = i_tc0 + o_i + o_c1 = i_tc1 + o_i + o_c2 = i_tc2 + o_i + o_c3 = i_tc3 + o_i + m_tc0 = o_c0 < T + m_A0 = m_tc0[:, None] & (o_i[None, :] < BT) + m_A1 = m_tc1[:, None] & (o_i[None, :] < BT) + m_A2 = m_tc2[:, None] & (o_i[None, :] < BT) + m_A3 = m_tc3[:, None] & (o_i[None, :] < BT) b_Aqk10 = tl.zeros([BC, BC], dtype=tl.float32) b_Akk10 = tl.zeros([BC, BC], dtype=tl.float32) @@ -129,29 +138,21 @@ def chunk_kda_fwd_kernel_inter_solve_fused( o_k = i_k * BK + tl.arange(0, BK) m_k = o_k < K - p_k0 = tl.make_block_ptr( - k, (T, K), (H * K, 1), (i_tc0, i_k * BK), (BC, BK), (1, 0) - ) - p_g0 = tl.make_block_ptr( - g, (T, K), (HV * K, 1), (i_tc0, i_k * BK), (BC, BK), (1, 0) - ) - b_k0 = tl.load(p_k0, boundary_check=(0, 1)).to(tl.float32) - b_g0 = tl.load(p_g0, boundary_check=(0, 1)).to(tl.float32) + m_ck0 = m_tc0[:, None] & m_k[None, :] + p_k0 = k + o_c0[:, None] * (H * K) + o_k[None, :] + p_g0 = g + o_c0[:, None] * (HV * K) + o_k[None, :] + b_k0 = tl.load(p_k0, mask=m_ck0, other=0.0).to(tl.float32) + b_g0 = tl.load(p_g0, mask=m_ck0, other=0.0).to(tl.float32) if i_tc1 < T: - p_q1 = tl.make_block_ptr( - q, (T, K), (H * K, 1), (i_tc1, i_k * BK), (BC, BK), (1, 0) - ) - p_k1 = tl.make_block_ptr( - k, (T, K), (H * K, 1), (i_tc1, i_k * BK), (BC, BK), (1, 0) - ) - p_g1 = tl.make_block_ptr( - g, (T, K), (HV * K, 1), (i_tc1, i_k * BK), (BC, BK), (1, 0) - ) + m_ck1 = m_tc1[:, None] & m_k[None, :] + p_q1 = q + o_c1[:, None] * (H * K) + o_k[None, :] + p_k1 = k + o_c1[:, None] * (H * K) + o_k[None, :] + p_g1 = g + o_c1[:, None] * (HV * K) + o_k[None, :] # [BC, BK] - b_q1 = tl.load(p_q1, boundary_check=(0, 1)).to(tl.float32) - b_k1 = tl.load(p_k1, boundary_check=(0, 1)).to(tl.float32) - b_g1 = tl.load(p_g1, boundary_check=(0, 1)).to(tl.float32) + b_q1 = tl.load(p_q1, mask=m_ck1, other=0.0).to(tl.float32) + b_k1 = tl.load(p_k1, mask=m_ck1, other=0.0).to(tl.float32) + b_g1 = tl.load(p_g1, mask=m_ck1, other=0.0).to(tl.float32) # [BK] b_gn1 = tl.load(g + i_tc1 * HV * K + o_k, mask=m_k, other=0).to(tl.float32) # [BC, BK] @@ -163,19 +164,14 @@ def chunk_kda_fwd_kernel_inter_solve_fused( b_Akk10 += tl.dot(b_k1 * b_gqn, b_kgt) if i_tc2 < T: - p_q2 = tl.make_block_ptr( - q, (T, K), (H * K, 1), (i_tc2, i_k * BK), (BC, BK), (1, 0) - ) - p_k2 = tl.make_block_ptr( - k, (T, K), (H * K, 1), (i_tc2, i_k * BK), (BC, BK), (1, 0) - ) - p_g2 = tl.make_block_ptr( - g, (T, K), (HV * K, 1), (i_tc2, i_k * BK), (BC, BK), (1, 0) - ) + m_ck2 = m_tc2[:, None] & m_k[None, :] + p_q2 = q + o_c2[:, None] * (H * K) + o_k[None, :] + p_k2 = k + o_c2[:, None] * (H * K) + o_k[None, :] + p_g2 = g + o_c2[:, None] * (HV * K) + o_k[None, :] # [BC, BK] - b_q2 = tl.load(p_q2, boundary_check=(0, 1)).to(tl.float32) - b_k2 = tl.load(p_k2, boundary_check=(0, 1)).to(tl.float32) - b_g2 = tl.load(p_g2, boundary_check=(0, 1)).to(tl.float32) + b_q2 = tl.load(p_q2, mask=m_ck2, other=0.0).to(tl.float32) + b_k2 = tl.load(p_k2, mask=m_ck2, other=0.0).to(tl.float32) + b_g2 = tl.load(p_g2, mask=m_ck2, other=0.0).to(tl.float32) # [BK] b_gn2 = tl.load(g + i_tc2 * HV * K + o_k, mask=m_k, other=0).to( tl.float32 @@ -195,19 +191,14 @@ def chunk_kda_fwd_kernel_inter_solve_fused( b_Akk21 += tl.dot(b_kg2, b_kgt) if i_tc3 < T: - p_q3 = tl.make_block_ptr( - q, (T, K), (H * K, 1), (i_tc3, i_k * BK), (BC, BK), (1, 0) - ) - p_k3 = tl.make_block_ptr( - k, (T, K), (H * K, 1), (i_tc3, i_k * BK), (BC, BK), (1, 0) - ) - p_g3 = tl.make_block_ptr( - g, (T, K), (HV * K, 1), (i_tc3, i_k * BK), (BC, BK), (1, 0) - ) + m_ck3 = m_tc3[:, None] & m_k[None, :] + p_q3 = q + o_c3[:, None] * (H * K) + o_k[None, :] + p_k3 = k + o_c3[:, None] * (H * K) + o_k[None, :] + p_g3 = g + o_c3[:, None] * (HV * K) + o_k[None, :] # [BC, BK] - b_q3 = tl.load(p_q3, boundary_check=(0, 1)).to(tl.float32) - b_k3 = tl.load(p_k3, boundary_check=(0, 1)).to(tl.float32) - b_g3 = tl.load(p_g3, boundary_check=(0, 1)).to(tl.float32) + b_q3 = tl.load(p_q3, mask=m_ck3, other=0.0).to(tl.float32) + b_k3 = tl.load(p_k3, mask=m_ck3, other=0.0).to(tl.float32) + b_g3 = tl.load(p_g3, mask=m_ck3, other=0.0).to(tl.float32) # [BK] b_gn3 = tl.load(g + i_tc3 * HV * K + o_k, mask=m_k, other=0).to( tl.float32 @@ -236,82 +227,44 @@ def chunk_kda_fwd_kernel_inter_solve_fused( # save off-diagonal Aqk blocks and prepare Akk ################################################################################ if i_tc1 < T: - p_Aqk10 = tl.make_block_ptr( - Aqk, (T, BT), (HV * BT, 1), (i_tc1, 0), (BC, BC), (1, 0) - ) - tl.store( - p_Aqk10, (b_Aqk10 * scale).to(Aqk.dtype.element_ty), boundary_check=(0, 1) - ) + p_Aqk10 = Aqk + o_c1[:, None] * (HV * BT) + o_i[None, :] + tl.store(p_Aqk10, (b_Aqk10 * scale).to(Aqk.dtype.element_ty), mask=m_A1) - p_b1 = tl.make_block_ptr( - beta + bos * HV + i_hv, (T,), (HV,), (i_tc1,), (BC,), (0,) - ) - b_b1 = tl.load(p_b1, boundary_check=(0,)).to(tl.float32) + p_b1 = beta + bos * HV + i_hv + o_c1 * HV + b_b1 = tl.load(p_b1, mask=m_tc1, other=0.0).to(tl.float32) b_Akk10 = b_Akk10 * b_b1[:, None] if i_tc2 < T: - p_Aqk20 = tl.make_block_ptr( - Aqk, (T, BT), (HV * BT, 1), (i_tc2, 0), (BC, BC), (1, 0) - ) - p_Aqk21 = tl.make_block_ptr( - Aqk, (T, BT), (HV * BT, 1), (i_tc2, BC), (BC, BC), (1, 0) - ) - tl.store( - p_Aqk20, (b_Aqk20 * scale).to(Aqk.dtype.element_ty), boundary_check=(0, 1) - ) - tl.store( - p_Aqk21, (b_Aqk21 * scale).to(Aqk.dtype.element_ty), boundary_check=(0, 1) - ) + p_Aqk20 = Aqk + o_c2[:, None] * (HV * BT) + o_i[None, :] + p_Aqk21 = Aqk + o_c2[:, None] * (HV * BT) + (o_i + BC)[None, :] + tl.store(p_Aqk20, (b_Aqk20 * scale).to(Aqk.dtype.element_ty), mask=m_A2) + tl.store(p_Aqk21, (b_Aqk21 * scale).to(Aqk.dtype.element_ty), mask=m_A2) - p_b2 = tl.make_block_ptr( - beta + bos * HV + i_hv, (T,), (HV,), (i_tc2,), (BC,), (0,) - ) - b_b2 = tl.load(p_b2, boundary_check=(0,)).to(tl.float32) + p_b2 = beta + bos * HV + i_hv + o_c2 * HV + b_b2 = tl.load(p_b2, mask=m_tc2, other=0.0).to(tl.float32) b_Akk20 = b_Akk20 * b_b2[:, None] b_Akk21 = b_Akk21 * b_b2[:, None] if i_tc3 < T: - p_Aqk30 = tl.make_block_ptr( - Aqk, (T, BT), (HV * BT, 1), (i_tc3, 0), (BC, BC), (1, 0) - ) - p_Aqk31 = tl.make_block_ptr( - Aqk, (T, BT), (HV * BT, 1), (i_tc3, BC), (BC, BC), (1, 0) - ) - p_Aqk32 = tl.make_block_ptr( - Aqk, (T, BT), (HV * BT, 1), (i_tc3, 2 * BC), (BC, BC), (1, 0) - ) - tl.store( - p_Aqk30, (b_Aqk30 * scale).to(Aqk.dtype.element_ty), boundary_check=(0, 1) - ) - tl.store( - p_Aqk31, (b_Aqk31 * scale).to(Aqk.dtype.element_ty), boundary_check=(0, 1) - ) - tl.store( - p_Aqk32, (b_Aqk32 * scale).to(Aqk.dtype.element_ty), boundary_check=(0, 1) - ) + p_Aqk30 = Aqk + o_c3[:, None] * (HV * BT) + o_i[None, :] + p_Aqk31 = Aqk + o_c3[:, None] * (HV * BT) + (o_i + BC)[None, :] + p_Aqk32 = Aqk + o_c3[:, None] * (HV * BT) + (o_i + 2 * BC)[None, :] + tl.store(p_Aqk30, (b_Aqk30 * scale).to(Aqk.dtype.element_ty), mask=m_A3) + tl.store(p_Aqk31, (b_Aqk31 * scale).to(Aqk.dtype.element_ty), mask=m_A3) + tl.store(p_Aqk32, (b_Aqk32 * scale).to(Aqk.dtype.element_ty), mask=m_A3) - p_b3 = tl.make_block_ptr( - beta + bos * HV + i_hv, (T,), (HV,), (i_tc3,), (BC,), (0,) - ) - b_b3 = tl.load(p_b3, boundary_check=(0,)).to(tl.float32) + p_b3 = beta + bos * HV + i_hv + o_c3 * HV + b_b3 = tl.load(p_b3, mask=m_tc3, other=0.0).to(tl.float32) b_Akk30 = b_Akk30 * b_b3[:, None] b_Akk31 = b_Akk31 * b_b3[:, None] b_Akk32 = b_Akk32 * b_b3[:, None] - p_Akk00 = tl.make_block_ptr( - Akkd, (T, BC), (HV * BC, 1), (i_tc0, 0), (BC, BC), (1, 0) - ) - p_Akk11 = tl.make_block_ptr( - Akkd, (T, BC), (HV * BC, 1), (i_tc1, 0), (BC, BC), (1, 0) - ) - p_Akk22 = tl.make_block_ptr( - Akkd, (T, BC), (HV * BC, 1), (i_tc2, 0), (BC, BC), (1, 0) - ) - p_Akk33 = tl.make_block_ptr( - Akkd, (T, BC), (HV * BC, 1), (i_tc3, 0), (BC, BC), (1, 0) - ) - b_Ai00 = tl.load(p_Akk00, boundary_check=(0, 1)).to(tl.float32) - b_Ai11 = tl.load(p_Akk11, boundary_check=(0, 1)).to(tl.float32) - b_Ai22 = tl.load(p_Akk22, boundary_check=(0, 1)).to(tl.float32) - b_Ai33 = tl.load(p_Akk33, boundary_check=(0, 1)).to(tl.float32) + p_Akk00 = Akkd + o_c0[:, None] * (HV * BC) + o_i[None, :] + p_Akk11 = Akkd + o_c1[:, None] * (HV * BC) + o_i[None, :] + p_Akk22 = Akkd + o_c2[:, None] * (HV * BC) + o_i[None, :] + p_Akk33 = Akkd + o_c3[:, None] * (HV * BC) + o_i[None, :] + b_Ai00 = tl.load(p_Akk00, mask=m_A0, other=0.0).to(tl.float32) + b_Ai11 = tl.load(p_Akk11, mask=m_A1, other=0.0).to(tl.float32) + b_Ai22 = tl.load(p_Akk22, mask=m_A2, other=0.0).to(tl.float32) + b_Ai33 = tl.load(p_Akk33, mask=m_A3, other=0.0).to(tl.float32) ################################################################################ # forward substitution on diagonals @@ -397,47 +350,27 @@ def chunk_kda_fwd_kernel_inter_solve_fused( # store full Akk_inv to Akk ################################################################################ - p_Akk00 = tl.make_block_ptr( - Akk, (T, BT), (HV * BT, 1), (i_tc0, 0), (BC, BC), (1, 0) - ) - p_Akk10 = tl.make_block_ptr( - Akk, (T, BT), (HV * BT, 1), (i_tc1, 0), (BC, BC), (1, 0) - ) - p_Akk11 = tl.make_block_ptr( - Akk, (T, BT), (HV * BT, 1), (i_tc1, BC), (BC, BC), (1, 0) - ) - p_Akk20 = tl.make_block_ptr( - Akk, (T, BT), (HV * BT, 1), (i_tc2, 0), (BC, BC), (1, 0) - ) - p_Akk21 = tl.make_block_ptr( - Akk, (T, BT), (HV * BT, 1), (i_tc2, BC), (BC, BC), (1, 0) - ) - p_Akk22 = tl.make_block_ptr( - Akk, (T, BT), (HV * BT, 1), (i_tc2, 2 * BC), (BC, BC), (1, 0) - ) - p_Akk30 = tl.make_block_ptr( - Akk, (T, BT), (HV * BT, 1), (i_tc3, 0), (BC, BC), (1, 0) - ) - p_Akk31 = tl.make_block_ptr( - Akk, (T, BT), (HV * BT, 1), (i_tc3, BC), (BC, BC), (1, 0) - ) - p_Akk32 = tl.make_block_ptr( - Akk, (T, BT), (HV * BT, 1), (i_tc3, 2 * BC), (BC, BC), (1, 0) - ) - p_Akk33 = tl.make_block_ptr( - Akk, (T, BT), (HV * BT, 1), (i_tc3, 3 * BC), (BC, BC), (1, 0) - ) + p_Akk00 = Akk + o_c0[:, None] * (HV * BT) + o_i[None, :] + p_Akk10 = Akk + o_c1[:, None] * (HV * BT) + o_i[None, :] + p_Akk11 = Akk + o_c1[:, None] * (HV * BT) + (o_i + BC)[None, :] + p_Akk20 = Akk + o_c2[:, None] * (HV * BT) + o_i[None, :] + p_Akk21 = Akk + o_c2[:, None] * (HV * BT) + (o_i + BC)[None, :] + p_Akk22 = Akk + o_c2[:, None] * (HV * BT) + (o_i + 2 * BC)[None, :] + p_Akk30 = Akk + o_c3[:, None] * (HV * BT) + o_i[None, :] + p_Akk31 = Akk + o_c3[:, None] * (HV * BT) + (o_i + BC)[None, :] + p_Akk32 = Akk + o_c3[:, None] * (HV * BT) + (o_i + 2 * BC)[None, :] + p_Akk33 = Akk + o_c3[:, None] * (HV * BT) + (o_i + 3 * BC)[None, :] - tl.store(p_Akk00, b_Ai00.to(Akk.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_Akk10, b_Ai10.to(Akk.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_Akk11, b_Ai11.to(Akk.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_Akk20, b_Ai20.to(Akk.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_Akk21, b_Ai21.to(Akk.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_Akk22, b_Ai22.to(Akk.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_Akk30, b_Ai30.to(Akk.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_Akk31, b_Ai31.to(Akk.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_Akk32, b_Ai32.to(Akk.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_Akk33, b_Ai33.to(Akk.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_Akk00, b_Ai00.to(Akk.dtype.element_ty), mask=m_A0) + tl.store(p_Akk10, b_Ai10.to(Akk.dtype.element_ty), mask=m_A1) + tl.store(p_Akk11, b_Ai11.to(Akk.dtype.element_ty), mask=m_A1) + tl.store(p_Akk20, b_Ai20.to(Akk.dtype.element_ty), mask=m_A2) + tl.store(p_Akk21, b_Ai21.to(Akk.dtype.element_ty), mask=m_A2) + tl.store(p_Akk22, b_Ai22.to(Akk.dtype.element_ty), mask=m_A2) + tl.store(p_Akk30, b_Ai30.to(Akk.dtype.element_ty), mask=m_A3) + tl.store(p_Akk31, b_Ai31.to(Akk.dtype.element_ty), mask=m_A3) + tl.store(p_Akk32, b_Ai32.to(Akk.dtype.element_ty), mask=m_A3) + tl.store(p_Akk33, b_Ai33.to(Akk.dtype.element_ty), mask=m_A3) @triton.heuristics( @@ -474,14 +407,14 @@ def chunk_kda_fwd_kernel_intra_sub_chunk( IS_VARLEN: tl.constexpr, USE_GATHER: tl.constexpr, ): - i_t, i_i, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_t, i_i, i_bh = tl.program_id(0).to(tl.int64), tl.program_id(1), tl.program_id(2) i_b, i_hv = i_bh // HV, i_bh % HV i_h = i_hv // (HV // H) if IS_VARLEN: i_n, i_t = ( tl.load(chunk_indices + i_t * 2).to(tl.int32), - tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), + tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64), ) bos, eos = ( tl.load(cu_seqlens + i_n).to(tl.int32), @@ -505,16 +438,19 @@ def chunk_kda_fwd_kernel_intra_sub_chunk( Aqk = Aqk + (bos * HV + i_hv) * BT Akk = Akk + (bos * HV + i_hv) * BC - p_q = tl.make_block_ptr(q, (T, K), (H * K, 1), (i_ti, 0), (BC, BK), (1, 0)) - p_k = tl.make_block_ptr(k, (T, K), (H * K, 1), (i_ti, 0), (BC, BK), (1, 0)) - p_g = tl.make_block_ptr(g, (T, K), (HV * K, 1), (i_ti, 0), (BC, BK), (1, 0)) + o_k = tl.arange(0, BK) + m_k = o_k < K + m_ck = m_c[:, None] & m_k[None, :] + p_q = q + o_c[:, None] * (H * K) + o_k[None, :] + p_k = k + o_c[:, None] * (H * K) + o_k[None, :] + p_g = g + o_c[:, None] * (HV * K) + o_k[None, :] - p_beta = tl.make_block_ptr(beta, (T,), (HV,), (i_ti,), (BC,), (0,)) + p_beta = beta + o_c * HV - b_q = tl.load(p_q, boundary_check=(0, 1)) - b_k = tl.load(p_k, boundary_check=(0, 1)) - b_g = tl.load(p_g, boundary_check=(0, 1)) - b_beta = tl.load(p_beta, boundary_check=(0,)).to(tl.float32) + b_q = tl.load(p_q, mask=m_ck, other=0.0) + b_k = tl.load(p_k, mask=m_ck, other=0.0) + b_g = tl.load(p_g, mask=m_ck, other=0.0) + b_beta = tl.load(p_beta, mask=m_c, other=0.0).to(tl.float32) if USE_GATHER: b_gn = gather( @@ -546,12 +482,12 @@ def chunk_kda_fwd_kernel_intra_sub_chunk( b_Aqk = tl.where(m_Aqk, b_Aqk, 0.0) b_Akk = tl.where(m_Akk, b_Akk, 0.0) - p_Aqk = tl.make_block_ptr( - Aqk, (T, BT), (HV * BT, 1), (i_ti, i_i * BC), (BC, BC), (1, 0) - ) - p_Akk = tl.make_block_ptr(Akk, (T, BC), (HV * BC, 1), (i_ti, 0), (BC, BC), (1, 0)) - tl.store(p_Aqk, b_Aqk.to(Aqk.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_Akk, b_Akk.to(Akk.dtype.element_ty), boundary_check=(0, 1)) + m_Aqk_st = m_c[:, None] & (o_i[None, :] < BT) + m_Akk_st = m_c[:, None] & (o_i[None, :] < BC) + p_Aqk = Aqk + o_c[:, None] * (HV * BT) + (i_i * BC + o_i)[None, :] + p_Akk = Akk + o_c[:, None] * (HV * BC) + o_i[None, :] + tl.store(p_Aqk, b_Aqk.to(Aqk.dtype.element_ty), mask=m_Aqk_st) + tl.store(p_Akk, b_Akk.to(Akk.dtype.element_ty), mask=m_Akk_st) tl.debug_barrier() @@ -566,7 +502,7 @@ def chunk_kda_fwd_kernel_intra_sub_chunk( b_a += tl.sum(b_a[:, None] * b_Ai, 0) b_Ai = tl.where((o_i == i)[:, None], b_a, b_Ai) b_Ai += m_I - tl.store(p_Akk, b_Ai.to(Akk.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_Akk, b_Ai.to(Akk.dtype.element_ty), mask=m_Akk_st) def chunk_kda_fwd_intra( diff --git a/vllm/models/kimi_k3/nvidia/ops/third_party/kda/chunk_intra_token_parallel.py b/vllm/models/kimi_k3/nvidia/ops/third_party/kda/chunk_intra_token_parallel.py index 86654f2644..7cb68b2fd4 100644 --- a/vllm/models/kimi_k3/nvidia/ops/third_party/kda/chunk_intra_token_parallel.py +++ b/vllm/models/kimi_k3/nvidia/ops/third_party/kda/chunk_intra_token_parallel.py @@ -108,20 +108,16 @@ def chunk_kda_fwd_kernel_intra_token_parallel( b_k = tl.load(k + i_t * H * K + p_qk, mask=m_hk, other=0).to(tl.float32) # g: [B, T, HV, K], beta: [B, T, HV] - p_g = tl.make_block_ptr( - g + i_t * HV * K, (HV, K), (K, 1), (i_hg * BH, 0), (BH, BK), (1, 0) - ) - p_beta = tl.make_block_ptr(beta + i_t * HV, (HV,), (1,), (i_hg * BH,), (BH,), (0,)) - b_g = tl.load(p_g, boundary_check=(0, 1)).to(tl.float32) - b_beta = tl.load(p_beta, boundary_check=(0,)).to(tl.float32) + p_g = g + i_t * HV * K + o_hv[:, None] * K + o_k[None, :] + p_beta = beta + i_t * HV + o_hv + b_g = tl.load(p_g, mask=m_hk, other=0.0).to(tl.float32) + b_beta = tl.load(p_beta, mask=m_hv, other=0.0).to(tl.float32) b_k *= b_beta[:, None] for j in range(i_ts, min(i_t + 1, min(T, i_ts + BC))): b_kj = tl.load(k + j * H * K + p_qk, mask=m_hk, other=0).to(tl.float32) - p_gj = tl.make_block_ptr( - g + j * HV * K, (HV, K), (K, 1), (i_hg * BH, 0), (BH, BK), (1, 0) - ) - b_gj = tl.load(p_gj, boundary_check=(0, 1)).to(tl.float32) + p_gj = g + j * HV * K + o_hv[:, None] * K + o_k[None, :] + b_gj = tl.load(p_gj, mask=m_hk, other=0.0).to(tl.float32) b_kgj = tl.where(m_k[None, :], b_kj * exp2(b_g - b_gj), 0.0) b_Aqk = tl.sum(b_q * b_kgj, axis=1) * scale diff --git a/vllm/models/minimax_m3/common/ops/index_topk.py b/vllm/models/minimax_m3/common/ops/index_topk.py index 025c3dd428..2d612c12f2 100644 --- a/vllm/models/minimax_m3/common/ops/index_topk.py +++ b/vllm/models/minimax_m3/common/ops/index_topk.py @@ -115,15 +115,17 @@ def _index_block_score_kernel( if BLOCK_SIZE_Q * pid_q >= q_len: return - q_ptrs = tl.make_block_ptr( - base=q_ptr + seq_start * stride_q_n + pid_h * stride_q_h, - shape=(q_len, head_dim), - strides=(stride_q_n, stride_q_d), - offsets=(pid_q * BLOCK_SIZE_Q, 0), - block_shape=(BLOCK_SIZE_Q, head_dim), - order=(1, 0), + o_q = pid_q * BLOCK_SIZE_Q + tl.arange(0, BLOCK_SIZE_Q) + o_d = tl.arange(0, head_dim) + m_q = o_q < q_len + q_ptrs = ( + q_ptr + + seq_start.to(tl.int64) * stride_q_n + + pid_h.to(tl.int64) * stride_q_h + + o_q[:, None].to(tl.int64) * stride_q_n + + o_d[None, :].to(tl.int64) * stride_q_d ) - q = tl.load(q_ptrs, boundary_check=(0,), padding_option="zero") + q = tl.load(q_ptrs, mask=m_q[:, None], other=0.0) # A matched fp8 pair (e.g. e4m3 x e4m3) lowers to a native FP8 MMA, so # upcast only when the operands differ: there is no mixed-dtype fp8 MMA, # including across fp8 flavours (fp8e4nv vs fp8e4b8). bf16 and fp32 loads diff --git a/vllm/third_party/flash_linear_attention/ops/chunk_delta_h.py b/vllm/third_party/flash_linear_attention/ops/chunk_delta_h.py index eb1c3af152..cc27437fac 100644 --- a/vllm/third_party/flash_linear_attention/ops/chunk_delta_h.py +++ b/vllm/third_party/flash_linear_attention/ops/chunk_delta_h.py @@ -110,109 +110,82 @@ def chunk_gated_delta_rule_fwd_kernel_h_blockdim64( ht = ht + i_nh * V * K # load initial state + o_v = i_v * BV + tl.arange(0, BV) + m_v = o_v < V + o_k1 = tl.arange(0, 64) + m_k1 = o_k1 < K + o_k2 = 64 + o_k1 + m_k2 = o_k2 < K + o_k3 = 128 + o_k1 + m_k3 = o_k3 < K + o_k4 = 192 + o_k1 + m_k4 = o_k4 < K if USE_INITIAL_STATE: - p_h0_1 = tl.make_block_ptr(h0, (V, K), (K, 1), (i_v * BV, 0), (BV, 64), (1, 0)) - b_h1 += tl.load(p_h0_1, boundary_check=(0, 1)).to(tl.float32) + p_h0_1 = h0 + o_v[:, None] * K + o_k1[None, :] + m_h0_1 = m_v[:, None] & m_k1[None, :] + b_h1 += tl.load(p_h0_1, mask=m_h0_1, other=0.0).to(tl.float32) if K > 64: - p_h0_2 = tl.make_block_ptr( - h0, (V, K), (K, 1), (i_v * BV, 64), (BV, 64), (1, 0) - ) - b_h2 += tl.load(p_h0_2, boundary_check=(0, 1)).to(tl.float32) + p_h0_2 = h0 + o_v[:, None] * K + o_k2[None, :] + m_h0_2 = m_v[:, None] & m_k2[None, :] + b_h2 += tl.load(p_h0_2, mask=m_h0_2, other=0.0).to(tl.float32) if K > 128: - p_h0_3 = tl.make_block_ptr( - h0, (V, K), (K, 1), (i_v * BV, 128), (BV, 64), (1, 0) - ) - b_h3 += tl.load(p_h0_3, boundary_check=(0, 1)).to(tl.float32) + p_h0_3 = h0 + o_v[:, None] * K + o_k3[None, :] + m_h0_3 = m_v[:, None] & m_k3[None, :] + b_h3 += tl.load(p_h0_3, mask=m_h0_3, other=0.0).to(tl.float32) if K > 192: - p_h0_4 = tl.make_block_ptr( - h0, (V, K), (K, 1), (i_v * BV, 192), (BV, 64), (1, 0) - ) - b_h4 += tl.load(p_h0_4, boundary_check=(0, 1)).to(tl.float32) + p_h0_4 = h0 + o_v[:, None] * K + o_k4[None, :] + m_h0_4 = m_v[:, None] & m_k4[None, :] + b_h4 += tl.load(p_h0_4, mask=m_h0_4, other=0.0).to(tl.float32) # main recurrence for i_t in range(NT): - p_h1 = tl.make_block_ptr( - h + i_t.to(tl.int64) * stride_h, - (V, K), - (K, 1), - (i_v * BV, 0), - (BV, 64), - (1, 0), - ) - tl.store(p_h1, b_h1.to(p_h1.dtype.element_ty), boundary_check=(0, 1)) + o_t = (i_t * BT + tl.arange(0, BT)).to(tl.int64) + m_t = o_t < T + p_h1 = h + i_t.to(tl.int64) * stride_h + o_v[:, None] * K + o_k1[None, :] + m_h1 = m_v[:, None] & m_k1[None, :] + tl.store(p_h1, b_h1.to(p_h1.dtype.element_ty), mask=m_h1) if K > 64: - p_h2 = tl.make_block_ptr( - h + i_t.to(tl.int64) * stride_h, - (V, K), - (K, 1), - (i_v * BV, 64), - (BV, 64), - (1, 0), - ) - tl.store(p_h2, b_h2.to(p_h2.dtype.element_ty), boundary_check=(0, 1)) + p_h2 = h + i_t.to(tl.int64) * stride_h + o_v[:, None] * K + o_k2[None, :] + m_h2 = m_v[:, None] & m_k2[None, :] + tl.store(p_h2, b_h2.to(p_h2.dtype.element_ty), mask=m_h2) if K > 128: - p_h3 = tl.make_block_ptr( - h + i_t.to(tl.int64) * stride_h, - (V, K), - (K, 1), - (i_v * BV, 128), - (BV, 64), - (1, 0), - ) - tl.store(p_h3, b_h3.to(p_h3.dtype.element_ty), boundary_check=(0, 1)) + p_h3 = h + i_t.to(tl.int64) * stride_h + o_v[:, None] * K + o_k3[None, :] + m_h3 = m_v[:, None] & m_k3[None, :] + tl.store(p_h3, b_h3.to(p_h3.dtype.element_ty), mask=m_h3) if K > 192: - p_h4 = tl.make_block_ptr( - h + i_t.to(tl.int64) * stride_h, - (V, K), - (K, 1), - (i_v * BV, 192), - (BV, 64), - (1, 0), - ) - tl.store(p_h4, b_h4.to(p_h4.dtype.element_ty), boundary_check=(0, 1)) + p_h4 = h + i_t.to(tl.int64) * stride_h + o_v[:, None] * K + o_k4[None, :] + m_h4 = m_v[:, None] & m_k4[None, :] + tl.store(p_h4, b_h4.to(p_h4.dtype.element_ty), mask=m_h4) - p_w = tl.make_block_ptr( - w, (T, K), (stride_w, 1), (i_t * BT, 0), (BT, 64), (1, 0) - ) - b_w = tl.load(p_w, boundary_check=(0, 1)) + p_w = w + o_t[:, None] * stride_w + o_k1[None, :] + b_w = tl.load(p_w, mask=m_t[:, None] & m_k1[None, :], other=0.0) b_v = tl.dot(b_w, tl.trans(b_h1).to(b_w.dtype)) if K > 64: - p_w = tl.make_block_ptr( - w, (T, K), (stride_w, 1), (i_t * BT, 64), (BT, 64), (1, 0) - ) - b_w = tl.load(p_w, boundary_check=(0, 1)) + p_w = w + o_t[:, None] * stride_w + o_k2[None, :] + b_w = tl.load(p_w, mask=m_t[:, None] & m_k2[None, :], other=0.0) b_v += tl.dot(b_w, tl.trans(b_h2).to(b_w.dtype)) if K > 128: - p_w = tl.make_block_ptr( - w, (T, K), (stride_w, 1), (i_t * BT, 128), (BT, 64), (1, 0) - ) - b_w = tl.load(p_w, boundary_check=(0, 1)) + p_w = w + o_t[:, None] * stride_w + o_k3[None, :] + b_w = tl.load(p_w, mask=m_t[:, None] & m_k3[None, :], other=0.0) b_v += tl.dot(b_w, tl.trans(b_h3).to(b_w.dtype)) if K > 192: - p_w = tl.make_block_ptr( - w, (T, K), (stride_w, 1), (i_t * BT, 192), (BT, 64), (1, 0) - ) - b_w = tl.load(p_w, boundary_check=(0, 1)) + p_w = w + o_t[:, None] * stride_w + o_k4[None, :] + b_w = tl.load(p_w, mask=m_t[:, None] & m_k4[None, :], other=0.0) b_v += tl.dot(b_w, tl.trans(b_h4).to(b_w.dtype)) - p_v = tl.make_block_ptr( - v, (T, V), (stride_v, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0) - ) - b_v = tl.load(p_v, boundary_check=(0, 1)) - b_v + p_v = v + o_t[:, None] * stride_v + o_v[None, :] + b_v = tl.load(p_v, mask=m_t[:, None] & m_v[None, :], other=0.0) - b_v if SAVE_NEW_VALUE: - p_v = tl.make_block_ptr( - v_new, (T, V), (stride_v, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0) + p_v = v_new + o_t[:, None] * stride_v + o_v[None, :] + tl.store( + p_v, b_v.to(p_v.dtype.element_ty), mask=m_t[:, None] & m_v[None, :] ) - tl.store(p_v, b_v.to(p_v.dtype.element_ty), boundary_check=(0, 1)) last_idx = min((i_t.to(tl.int64) + 1) * BT, T) - 1 if USE_G: - m_t = (i_t.to(tl.int64) * BT + tl.arange(0, BT)) < T b_g_last = tl.load(g + bos * H + last_idx * H + i_h) - 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 = g + bos * H + i_h + o_t * H + b_g = tl.load(p_g, mask=m_t, other=0.0) if USE_EXP2: b_v = b_v * tl.where(m_t, exp2(b_g_last - b_g), 0)[:, None] b_g_last = exp2(b_g_last) @@ -239,7 +212,6 @@ def chunk_gated_delta_rule_fwd_kernel_h_blockdim64( else: b_h1 *= exp(b_gk_last1)[None, :] if K > 64: - o_k2 = 64 + o_k1 b_gk_last2 = tl.load( gk + (bos + last_idx) * H * K + i_h * K + o_k2, mask=(o_k2 < K), @@ -250,7 +222,6 @@ def chunk_gated_delta_rule_fwd_kernel_h_blockdim64( else: b_h2 *= exp(b_gk_last2)[None, :] if K > 128: - o_k3 = 128 + o_k1 b_gk_last3 = tl.load( gk + (bos + last_idx) * H * K + i_h * K + o_k3, mask=(o_k3 < K), @@ -261,7 +232,6 @@ def chunk_gated_delta_rule_fwd_kernel_h_blockdim64( else: b_h3 *= exp(b_gk_last3)[None, :] if K > 192: - o_k4 = 192 + o_k1 b_gk_last4 = tl.load( gk + (bos + last_idx) * H * K + i_h * K + o_k4, mask=(o_k4 < K), @@ -273,48 +243,38 @@ def chunk_gated_delta_rule_fwd_kernel_h_blockdim64( b_h4 *= exp(b_gk_last4)[None, :] b_v = b_v.to(k.dtype.element_ty) - p_k = tl.make_block_ptr( - k, (K, T), (1, stride_k), (0, i_t * BT), (64, BT), (0, 1) - ) - b_k = tl.load(p_k, boundary_check=(0, 1)) + p_k = k + o_k1[:, None] + o_t[None, :] * stride_k + b_k = tl.load(p_k, mask=m_k1[:, None] & m_t[None, :], other=0.0) b_h1 += tl.trans(tl.dot(b_k, b_v)) if K > 64: - p_k = tl.make_block_ptr( - k, (K, T), (1, stride_k), (64, i_t * BT), (64, BT), (0, 1) - ) - b_k = tl.load(p_k, boundary_check=(0, 1)) + p_k = k + o_k2[:, None] + o_t[None, :] * stride_k + b_k = tl.load(p_k, mask=m_k2[:, None] & m_t[None, :], other=0.0) b_h2 += tl.trans(tl.dot(b_k, b_v)) if K > 128: - p_k = tl.make_block_ptr( - k, (K, T), (1, stride_k), (128, i_t * BT), (64, BT), (0, 1) - ) - b_k = tl.load(p_k, boundary_check=(0, 1)) + p_k = k + o_k3[:, None] + o_t[None, :] * stride_k + b_k = tl.load(p_k, mask=m_k3[:, None] & m_t[None, :], other=0.0) b_h3 += tl.trans(tl.dot(b_k, b_v)) if K > 192: - p_k = tl.make_block_ptr( - k, (K, T), (1, stride_k), (192, i_t * BT), (64, BT), (0, 1) - ) - b_k = tl.load(p_k, boundary_check=(0, 1)) + p_k = k + o_k4[:, None] + o_t[None, :] * stride_k + b_k = tl.load(p_k, mask=m_k4[:, None] & m_t[None, :], other=0.0) b_h4 += tl.trans(tl.dot(b_k, b_v)) # epilogue if STORE_FINAL_STATE: - p_ht = tl.make_block_ptr(ht, (V, K), (K, 1), (i_v * BV, 0), (BV, 64), (1, 0)) - tl.store(p_ht, b_h1.to(p_ht.dtype.element_ty), boundary_check=(0, 1)) + p_ht = ht + o_v[:, None] * K + o_k1[None, :] + m_ht = m_v[:, None] & m_k1[None, :] + tl.store(p_ht, b_h1.to(p_ht.dtype.element_ty), mask=m_ht) if K > 64: - p_ht = tl.make_block_ptr( - ht, (V, K), (K, 1), (i_v * BV, 64), (BV, 64), (1, 0) - ) - tl.store(p_ht, b_h2.to(p_ht.dtype.element_ty), boundary_check=(0, 1)) + p_ht = ht + o_v[:, None] * K + o_k2[None, :] + m_ht = m_v[:, None] & m_k2[None, :] + tl.store(p_ht, b_h2.to(p_ht.dtype.element_ty), mask=m_ht) if K > 128: - p_ht = tl.make_block_ptr( - ht, (V, K), (K, 1), (i_v * BV, 128), (BV, 64), (1, 0) - ) - tl.store(p_ht, b_h3.to(p_ht.dtype.element_ty), boundary_check=(0, 1)) + p_ht = ht + o_v[:, None] * K + o_k3[None, :] + m_ht = m_v[:, None] & m_k3[None, :] + tl.store(p_ht, b_h3.to(p_ht.dtype.element_ty), mask=m_ht) if K > 192: - p_ht = tl.make_block_ptr( - ht, (V, K), (K, 1), (i_v * BV, 192), (BV, 64), (1, 0) - ) - tl.store(p_ht, b_h4.to(p_ht.dtype.element_ty), boundary_check=(0, 1)) + p_ht = ht + o_v[:, None] * K + o_k4[None, :] + m_ht = m_v[:, None] & m_k4[None, :] + tl.store(p_ht, b_h4.to(p_ht.dtype.element_ty), mask=m_ht) def chunk_gated_delta_rule_fwd_h( diff --git a/vllm/third_party/flash_linear_attention/ops/chunk_o.py b/vllm/third_party/flash_linear_attention/ops/chunk_o.py index 0c323b8ce2..219c8ed43f 100644 --- a/vllm/third_party/flash_linear_attention/ops/chunk_o.py +++ b/vllm/third_party/flash_linear_attention/ops/chunk_o.py @@ -60,14 +60,18 @@ def chunk_fwd_kernel_o( USE_G: tl.constexpr, IS_VARLEN: tl.constexpr, ): - i_v, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_v, i_t, i_bh = ( + tl.program_id(0), + tl.program_id(1).to(tl.int64), + tl.program_id(2).to(tl.int64), + ) i_b, i_h = i_bh // H, i_bh % H if IS_VARLEN: i_tg = i_t i_n, i_t = ( tl.load(chunk_indices + i_t * 2).to(tl.int32), - tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), + tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64), ) bos, eos = ( tl.load(cu_seqlens + i_n).to(tl.int32), @@ -90,22 +94,21 @@ def chunk_fwd_kernel_o( b_o = tl.zeros([BT, BV], dtype=tl.float32) b_A = tl.zeros([BT, BT], dtype=tl.float32) + o_t = i_t * BT + tl.arange(0, BT) + m_t = o_t < T + o_v = i_v * BV + tl.arange(0, BV) 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) - ) + o_k = i_k * BK + tl.arange(0, BK) + m_k = o_k < K + p_q = q + o_t[:, None] * (Hg * K) + o_k[None, :] + p_k = k + o_k[:, None] + o_t[None, :] * (Hg * K) + p_h = h + o_v[:, None] * K + o_k[None, :] # [BT, BK] - b_q = tl.load(p_q, boundary_check=(0, 1)) + b_q = tl.load(p_q, mask=m_t[:, None] & m_k[None, :], other=0.0) # [BK, BT] - b_k = tl.load(p_k, boundary_check=(0, 1)) + b_k = tl.load(p_k, mask=m_k[:, None] & m_t[None, :], other=0.0) # [BV, BK] - b_h = tl.load(p_h, boundary_check=(0, 1)) + b_h = tl.load(p_h, mask=(o_v[:, None] < V) & m_k[None, :], other=0.0) # [BT, BK] @ [BK, BV] -> [BT, BV] b_o += tl.dot(b_q, tl.trans(b_h)) @@ -114,28 +117,22 @@ 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,)) + p_g = g + o_t * H + b_g = tl.load(p_g, mask=m_t, other=0.0) b_o = b_o * exp(b_g)[:, None] b_A = b_A * exp(b_g[:, None] - b_g[None, :]) - o_t = i_t * BT + tl.arange(0, BT) - m_t = o_t < T 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)) + p_v = v + o_t[:, None] * (H * V) + o_v[None, :] + p_o = o + o_t[:, None] * (H * V) + o_v[None, :] + b_v = tl.load(p_v, mask=m_t[:, None] & (o_v < V)[None, :], other=0.0) # 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)) + tl.store(p_o, b_o.to(p_o.dtype.element_ty), mask=m_t[:, None] & (o_v < V)[None, :]) def chunk_fwd_o( diff --git a/vllm/third_party/flash_linear_attention/ops/chunk_scaled_dot_kkt.py b/vllm/third_party/flash_linear_attention/ops/chunk_scaled_dot_kkt.py index 10ff970c80..f610fc7031 100644 --- a/vllm/third_party/flash_linear_attention/ops/chunk_scaled_dot_kkt.py +++ b/vllm/third_party/flash_linear_attention/ops/chunk_scaled_dot_kkt.py @@ -61,12 +61,12 @@ def chunk_scaled_dot_kkt_fwd_kernel( USE_G: tl.constexpr, CAST_DOT_TO_K_DTYPE: tl.constexpr, ): - i_t, i_bh = tl.program_id(0), tl.program_id(1) + i_t, i_bh = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64) i_b, i_h = i_bh // H, i_bh % H if IS_VARLEN: i_n, i_t = ( tl.load(chunk_indices + i_t * 2).to(tl.int32), - tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), + tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64), ) bos, eos = ( tl.load(cu_seqlens + i_n).to(tl.int32), @@ -78,22 +78,19 @@ def chunk_scaled_dot_kkt_fwd_kernel( o_t = i_t * BT + tl.arange(0, BT) m_t = o_t < T - p_beta = tl.make_block_ptr( - beta + bos * H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,) - ) - b_beta = tl.load(p_beta, boundary_check=(0,)) + p_beta = beta + bos * H + i_h + o_t * H + b_beta = tl.load(p_beta, mask=m_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 * 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) + p_k = ( + k + + (bos * Hg + i_h // (H // Hg)) * K + + o_t[:, None] * (Hg * K) + + o_k[None, :] ) - b_k = tl.load(p_k, boundary_check=(0, 1)) + b_k = tl.load(p_k, mask=m_t[:, None] & (o_k < K)[None, :], other=0.0) b_kb = b_k * b_beta[:, None] if CAST_DOT_TO_K_DTYPE: # RDNA: force operands to k's native dtype so WMMA is used. @@ -103,17 +100,15 @@ def chunk_scaled_dot_kkt_fwd_kernel( b_A += tl.dot(b_kb, tl.trans(b_k).to(b_kb.dtype)) 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 = g + bos * H + i_h + o_t * H + b_g = tl.load(p_g, mask=m_t, other=0.0) b_g_diff = b_g[:, None] - b_g[None, :] b_A = b_A * exp(b_g_diff) 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) - ) - tl.store(p_A, b_A.to(p_A.dtype.element_ty), boundary_check=(0, 1)) + p_A = A + (bos * H + i_h) * BT + o_t[:, None] * (BT * H) + tl.arange(0, BT)[None, :] + tl.store(p_A, b_A.to(p_A.dtype.element_ty), mask=m_t[:, None]) def chunk_scaled_dot_kkt_fwd( diff --git a/vllm/third_party/flash_linear_attention/ops/cumsum.py b/vllm/third_party/flash_linear_attention/ops/cumsum.py index b0820104b1..1fcab916ca 100644 --- a/vllm/third_party/flash_linear_attention/ops/cumsum.py +++ b/vllm/third_party/flash_linear_attention/ops/cumsum.py @@ -37,12 +37,12 @@ def chunk_local_cumsum_scalar_kernel( IS_VARLEN: tl.constexpr, HEAD_FIRST: tl.constexpr, ): - i_t, i_bh = tl.program_id(0), tl.program_id(1) + i_t, i_bh = tl.program_id(0).to(tl.int64), tl.program_id(1) i_b, i_h = i_bh // H, i_bh % H if IS_VARLEN: i_n, i_t = ( tl.load(chunk_indices + i_t * 2).to(tl.int32), - tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), + tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64), ) bos, eos = ( tl.load(cu_seqlens + i_n).to(tl.int32), @@ -52,23 +52,21 @@ def chunk_local_cumsum_scalar_kernel( else: bos, eos = i_b * T, i_b * T + T + 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,) - ) + p_s = s + bos * H + i_h * T + o_t + p_o = o + bos * H + i_h * T + o_t 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,)) + p_s = s + bos * H + i_h + o_t * H + p_o = o + bos * H + i_h + o_t * H # [BT] - b_s = tl.load(p_s, boundary_check=(0,)).to(tl.float32) + b_s = tl.load(p_s, 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 - tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0,)) + tl.store(p_o, b_o.to(p_o.dtype.element_ty), mask=m_t) @triton.heuristics({"IS_VARLEN": lambda args: args["cu_seqlens"] is not None}) @@ -96,12 +94,12 @@ def chunk_local_cumsum_vector_kernel( IS_VARLEN: tl.constexpr, HEAD_FIRST: tl.constexpr, ): - i_s, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_s, i_t, i_bh = tl.program_id(0), tl.program_id(1).to(tl.int64), tl.program_id(2) i_b, i_h = i_bh // H, i_bh % H if IS_VARLEN: i_n, i_t = ( tl.load(chunk_indices + i_t * 2).to(tl.int32), - tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), + tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64), ) bos, eos = ( tl.load(cu_seqlens + i_n).to(tl.int32), @@ -117,44 +115,19 @@ def chunk_local_cumsum_vector_kernel( else: m_s = tl.where(o_i[:, None] >= o_i[None, :], 1.0, 0.0) + o_t = i_t * BT + tl.arange(0, BT) + o_s = i_s * BS + tl.arange(0, BS) + m_ts = (o_t[:, None] < T) & (o_s[None, :] < S) 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), - ) + p_s = s + (bos * H + i_h * T) * S + o_t[:, None] * S + o_s[None, :] + p_o = o + (bos * H + i_h * T) * S + o_t[:, None] * S + o_s[None, :] 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), - ) + p_s = s + (bos * H + i_h) * S + o_t[:, None] * (H * S) + o_s[None, :] + p_o = o + (bos * H + i_h) * S + o_t[:, None] * (H * S) + o_s[None, :] # [BT, BS] - b_s = tl.load(p_s, boundary_check=(0, 1)).to(tl.float32) + b_s = tl.load(p_s, mask=m_ts, other=0.0).to(tl.float32) b_o = tl.dot(m_s, b_s, allow_tf32=False) - tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_o, b_o.to(p_o.dtype.element_ty), mask=m_ts) def chunk_local_cumsum_scalar( diff --git a/vllm/third_party/flash_linear_attention/ops/fused_norm_gate.py b/vllm/third_party/flash_linear_attention/ops/fused_norm_gate.py index 8f58ddcf8c..51ef432707 100644 --- a/vllm/third_party/flash_linear_attention/ops/fused_norm_gate.py +++ b/vllm/third_party/flash_linear_attention/ops/fused_norm_gate.py @@ -50,31 +50,30 @@ def layer_norm_gated_fwd_kernel( HAS_BIAS: tl.constexpr, launch_pdl: tl.constexpr, ): - i_t = tl.program_id(0) + i_t = tl.program_id(0).to(tl.int64) + o_t = i_t * BT + tl.arange(0, BT) o_d = tl.arange(0, BD) m_d = o_d < D + m_t = o_t < T + m_x = m_t[:, None] & m_d[None, :] if launch_pdl: tl.extra.cuda.gdc_wait() tl.extra.cuda.gdc_launch_dependents() - p_x = tl.make_block_ptr(x, (T, D), (D, 1), (i_t * BT, 0), (BT, BD), (1, 0)) - b_x = tl.load(p_x, boundary_check=(0, 1)).to(tl.float32) + p_x = x + o_t[:, None] * D + o_d[None, :] + b_x = tl.load(p_x, mask=m_x, other=0.0).to(tl.float32) if HAS_RESIDUAL: - p_res = tl.make_block_ptr( - residual, (T, D), (D, 1), (i_t * BT, 0), (BT, BD), (1, 0) - ) - b_x += tl.load(p_res, boundary_check=(0, 1)).to(tl.float32) + p_res = residual + o_t[:, None] * D + o_d[None, :] + b_x += tl.load(p_res, mask=m_x, other=0.0).to(tl.float32) if STORE_RESIDUAL_OUT: - p_res_out = tl.make_block_ptr( - residual_out, (T, D), (D, 1), (i_t * BT, 0), (BT, BD), (1, 0) - ) - tl.store(p_res_out, b_x.to(p_res_out.dtype.element_ty), boundary_check=(0, 1)) + p_res_out = residual_out + o_t[:, None] * D + o_d[None, :] + tl.store(p_res_out, b_x.to(p_res_out.dtype.element_ty), mask=m_x) if not IS_RMS_NORM: b_mean = tl.sum(b_x, axis=1) / D - p_mean = tl.make_block_ptr(mean, (T,), (1,), (i_t * BT,), (BT,), (0,)) - tl.store(p_mean, b_mean.to(p_mean.dtype.element_ty), boundary_check=(0,)) + p_mean = mean + o_t + tl.store(p_mean, b_mean.to(p_mean.dtype.element_ty), mask=m_t) b_xbar = tl.where(m_d[None, :], b_x - b_mean[:, None], 0.0) b_var = tl.sum(b_xbar * b_xbar, axis=1) / D else: @@ -82,8 +81,8 @@ def layer_norm_gated_fwd_kernel( b_var = tl.sum(b_xbar * b_xbar, axis=1) / D b_rstd = 1 / tl.sqrt(b_var + eps) - p_rstd = tl.make_block_ptr(rstd, (T,), (1,), (i_t * BT,), (BT,), (0,)) - tl.store(p_rstd, b_rstd.to(p_rstd.dtype.element_ty), boundary_check=(0,)) + p_rstd = rstd + o_t + tl.store(p_rstd, b_rstd.to(p_rstd.dtype.element_ty), mask=m_t) if HAS_WEIGHT: b_w = tl.load(w + o_d, mask=m_d).to(tl.float32) @@ -99,7 +98,6 @@ def layer_norm_gated_fwd_kernel( b_y = b_y + b_b[None, :] # swish/sigmoid output gate - o_t = i_t * BT + tl.arange(0, BT) o_g = (o_t // H) * g_stride_n + (o_t % H) * D b_g = tl.load( g + o_g[:, None] + o_d[None, :], @@ -112,8 +110,8 @@ def layer_norm_gated_fwd_kernel( b_y = b_y * tl.sigmoid(b_g) # Write output - p_y = tl.make_block_ptr(y, (T, D), (D, 1), (i_t * BT, 0), (BT, BD), (1, 0)) - tl.store(p_y, b_y.to(p_y.dtype.element_ty), boundary_check=(0, 1)) + p_y = y + o_t[:, None] * D + o_d[None, :] + tl.store(p_y, b_y.to(p_y.dtype.element_ty), mask=m_x) @triton.heuristics( diff --git a/vllm/third_party/flash_linear_attention/ops/kda.py b/vllm/third_party/flash_linear_attention/ops/kda.py index b5c3a92b65..9327b24ddd 100644 --- a/vllm/third_party/flash_linear_attention/ops/kda.py +++ b/vllm/third_party/flash_linear_attention/ops/kda.py @@ -178,27 +178,26 @@ def layer_norm_gated_fwd_kernel( HAS_WEIGHT: tl.constexpr, HAS_BIAS: tl.constexpr, ): - i_t = tl.program_id(0) + i_t = tl.program_id(0).to(tl.int64) + o_t = i_t * BT + tl.arange(0, BT) o_d = tl.arange(0, BD) m_d = o_d < D + m_t = o_t < T + m_x = m_t[:, None] & m_d[None, :] - p_x = tl.make_block_ptr(x, (T, D), (D, 1), (i_t * BT, 0), (BT, BD), (1, 0)) - b_x = tl.load(p_x, boundary_check=(0, 1)).to(tl.float32) + p_x = x + o_t[:, None] * D + o_d[None, :] + b_x = tl.load(p_x, mask=m_x, other=0.0).to(tl.float32) if HAS_RESIDUAL: - p_res = tl.make_block_ptr( - residual, (T, D), (D, 1), (i_t * BT, 0), (BT, BD), (1, 0) - ) - b_x += tl.load(p_res, boundary_check=(0, 1)).to(tl.float32) + p_res = residual + o_t[:, None] * D + o_d[None, :] + b_x += tl.load(p_res, mask=m_x, other=0.0).to(tl.float32) if STORE_RESIDUAL_OUT: - p_res_out = tl.make_block_ptr( - residual_out, (T, D), (D, 1), (i_t * BT, 0), (BT, BD), (1, 0) - ) - tl.store(p_res_out, b_x.to(p_res_out.dtype.element_ty), boundary_check=(0, 1)) + p_res_out = residual_out + o_t[:, None] * D + o_d[None, :] + tl.store(p_res_out, b_x.to(p_res_out.dtype.element_ty), mask=m_x) if not IS_RMS_NORM: b_mean = tl.sum(b_x, axis=1) / D - p_mean = tl.make_block_ptr(mean, (T,), (1,), (i_t * BT,), (BT,), (0,)) - tl.store(p_mean, b_mean.to(p_mean.dtype.element_ty), boundary_check=(0,)) + p_mean = mean + o_t + tl.store(p_mean, b_mean.to(p_mean.dtype.element_ty), mask=m_t) b_xbar = tl.where(m_d[None, :], b_x - b_mean[:, None], 0.0) b_var = tl.sum(b_xbar * b_xbar, axis=1) / D else: @@ -206,8 +205,8 @@ def layer_norm_gated_fwd_kernel( b_var = tl.sum(b_xbar * b_xbar, axis=1) / D b_rstd = 1 / tl.sqrt(b_var + eps) - p_rstd = tl.make_block_ptr(rstd, (T,), (1,), (i_t * BT,), (BT,), (0,)) - tl.store(p_rstd, b_rstd.to(p_rstd.dtype.element_ty), boundary_check=(0,)) + p_rstd = rstd + o_t + tl.store(p_rstd, b_rstd.to(p_rstd.dtype.element_ty), mask=m_t) if HAS_WEIGHT: b_w = tl.load(w + o_d, mask=m_d).to(tl.float32) @@ -223,7 +222,6 @@ def layer_norm_gated_fwd_kernel( b_y = b_y + b_b[None, :] # swish/sigmoid output gate - o_t = i_t * BT + tl.arange(0, BT) o_g = (o_t // H) * g_stride_n + (o_t % H) * D b_g = tl.load( g + o_g[:, None] + o_d[None, :], @@ -236,8 +234,8 @@ def layer_norm_gated_fwd_kernel( b_y = b_y * tl.sigmoid(b_g) # Write output - p_y = tl.make_block_ptr(y, (T, D), (D, 1), (i_t * BT, 0), (BT, BD), (1, 0)) - tl.store(p_y, b_y.to(p_y.dtype.element_ty), boundary_check=(0, 1)) + p_y = y + o_t[:, None] * D + o_d[None, :] + tl.store(p_y, b_y.to(p_y.dtype.element_ty), mask=m_x) @triton.heuristics( @@ -576,13 +574,13 @@ def chunk_kda_scaled_dot_kkt_fwd_kernel_intra_sub_inter( NC: tl.constexpr, IS_VARLEN: tl.constexpr, ): - i_t, i_c, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_t, i_c, i_bh = tl.program_id(0).to(tl.int64), tl.program_id(1), tl.program_id(2) i_b, i_h = i_bh // H, i_bh % H i_i, i_j = i_c // NC, i_c % NC if IS_VARLEN: i_n, i_t = ( tl.load(chunk_indices + i_t * 2).to(tl.int32), - tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), + tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64), ) bos, eos = ( tl.load(cu_seqlens + i_n).to(tl.int32), @@ -603,58 +601,51 @@ def chunk_kda_scaled_dot_kkt_fwd_kernel_intra_sub_inter( A += (bos * H + i_h) * BT Aqk += (bos * H + i_h) * BT - p_b = tl.make_block_ptr( - beta + bos * H + i_h, (T,), (H,), (i_t * BT + i_i * BC,), (BC,), (0,) - ) - b_b = tl.load(p_b, boundary_check=(0,)) + o_i = i_t * BT + i_i * BC + tl.arange(0, BC) + o_j = i_t * BT + i_j * BC + tl.arange(0, BC) + m_i = o_i < T + m_j = o_j < T + + p_b = beta + bos * H + i_h + o_i * H + b_b = tl.load(p_b, mask=m_i, other=0.0) b_A = tl.zeros([BC, BC], dtype=tl.float32) b_Aqk = tl.zeros([BC, BC], 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_i * BC, i_k * BK), (BC, BK), (1, 0) - ) - p_k = tl.make_block_ptr( - k, (T, K), (H * K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0) - ) - p_g = tl.make_block_ptr( - g, (T, K), (H * K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0) - ) - b_kt = tl.make_block_ptr( - k, (K, T), (1, H * K), (i_k * BK, i_t * BT + i_j * BC), (BK, BC), (0, 1) - ) - p_gk = tl.make_block_ptr( - g, (K, T), (1, H * K), (i_k * BK, i_t * BT + i_j * BC), (BK, BC), (0, 1) - ) - o_k = i_k * BK + tl.arange(0, BK) m_k = o_k < K + m_qk = m_i[:, None] & m_k[None, :] + m_kj = m_k[:, None] & m_j[None, :] + p_q = q + o_i[:, None] * (H * K) + o_k[None, :] + p_k = k + o_i[:, None] * (H * K) + o_k[None, :] + p_g = g + o_i[:, None] * (H * K) + o_k[None, :] + b_kt = k + o_k[:, None] + o_j[None, :] * (H * K) + p_gk = g + o_k[:, None] + o_j[None, :] * (H * K) + # [BK,] b_gn = tl.load(g + (i_t * BT + i_i * BC) * H * K + o_k, mask=m_k, other=0) # [BC, BK] - b_g = tl.load(p_g, boundary_check=(0, 1)) - b_k = tl.load(p_k, boundary_check=(0, 1)) * exp2(b_g - b_gn[None, :]) + b_g = tl.load(p_g, mask=m_qk, other=0.0) + b_k = tl.load(p_k, mask=m_qk, other=0.0) * exp2(b_g - b_gn[None, :]) # [BK, BC] - b_gk = tl.load(p_gk, boundary_check=(0, 1)) - b_kt = tl.load(b_kt, boundary_check=(0, 1)) + b_gk = tl.load(p_gk, mask=m_kj, other=0.0) + b_kt = tl.load(b_kt, mask=m_kj, other=0.0) # [BC, BC] b_ktg = b_kt * exp2(b_gn[:, None] - b_gk) b_A += tl.dot(b_k, b_ktg) - b_q = tl.load(p_q, boundary_check=(0, 1)) + b_q = tl.load(p_q, mask=m_qk, other=0.0) b_qg = b_q * exp2(b_g - b_gn[None, :]) * scale b_Aqk += tl.dot(b_qg, b_ktg) b_A *= b_b[:, None] - p_A = tl.make_block_ptr( - A, (T, BT), (H * BT, 1), (i_t * BT + i_i * BC, i_j * BC), (BC, BC), (1, 0) - ) - tl.store(p_A, b_A.to(A.dtype.element_ty), boundary_check=(0, 1)) - p_Aqk = tl.make_block_ptr( - Aqk, (T, BT), (H * BT, 1), (i_t * BT + i_i * BC, i_j * BC), (BC, BC), (1, 0) - ) - tl.store(p_Aqk, b_Aqk.to(Aqk.dtype.element_ty), boundary_check=(0, 1)) + o_jA = i_j * BC + tl.arange(0, BC) + m_A = m_i[:, None] & (o_jA[None, :] < BT) + p_A = A + o_i[:, None] * (H * BT) + o_jA[None, :] + tl.store(p_A, b_A.to(A.dtype.element_ty), mask=m_A) + p_Aqk = Aqk + o_i[:, None] * (H * BT) + o_jA[None, :] + tl.store(p_Aqk, b_Aqk.to(Aqk.dtype.element_ty), mask=m_A) @triton.heuristics({"IS_VARLEN": lambda args: args["cu_seqlens"] is not None}) @@ -681,12 +672,12 @@ def chunk_kda_scaled_dot_kkt_fwd_kernel_intra_sub_intra( BK: tl.constexpr, IS_VARLEN: tl.constexpr, ): - i_t, i_i, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_t, i_i, i_bh = tl.program_id(0).to(tl.int64), tl.program_id(1), tl.program_id(2) i_b, i_h = i_bh // H, i_bh % H if IS_VARLEN: i_n, i_t = ( tl.load(chunk_indices + i_t * 2).to(tl.int32), - tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), + tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64), ) bos, eos = ( tl.load(cu_seqlens + i_n).to(tl.int32), @@ -705,33 +696,14 @@ def chunk_kda_scaled_dot_kkt_fwd_kernel_intra_sub_intra( m_A = (i_t * BT + i_i * BC + o_i) < T o_A = (bos + i_t * BT + i_i * BC + o_i) * H * BT + i_h * BT + i_i * BC - p_q = tl.make_block_ptr( - q + (bos * H + i_h) * K, - (T, K), - (H * K, 1), - (i_t * BT + i_i * BC, 0), - (BC, BK), - (1, 0), - ) - p_k = tl.make_block_ptr( - k + (bos * H + i_h) * K, - (T, K), - (H * K, 1), - (i_t * BT + i_i * BC, 0), - (BC, BK), - (1, 0), - ) - p_g = tl.make_block_ptr( - g + (bos * H + i_h) * K, - (T, K), - (H * K, 1), - (i_t * BT + i_i * BC, 0), - (BC, BK), - (1, 0), - ) - b_q = tl.load(p_q, boundary_check=(0, 1)) - b_k = tl.load(p_k, boundary_check=(0, 1)) - b_g = tl.load(p_g, boundary_check=(0, 1)) + o_c = i_t * BT + i_i * BC + tl.arange(0, BC) + m_qk = m_A[:, None] & m_k[None, :] + p_q = q + (bos * H + i_h) * K + o_c[:, None] * (H * K) + o_k[None, :] + p_k = k + (bos * H + i_h) * K + o_c[:, None] * (H * K) + o_k[None, :] + p_g = g + (bos * H + i_h) * K + o_c[:, None] * (H * K) + o_k[None, :] + b_q = tl.load(p_q, mask=m_qk, other=0.0) + b_k = tl.load(p_k, mask=m_qk, other=0.0) + b_g = tl.load(p_g, mask=m_qk, other=0.0) p_b = beta + (bos + i_t * BT + i_i * BC + o_i) * H + i_h b_k = b_k * tl.load(p_b, mask=m_A, other=0)[:, None] @@ -878,12 +850,12 @@ def recompute_w_u_fwd_kernel( IS_VARLEN: tl.constexpr, DOT_PRECISION: tl.constexpr, ): - i_t, i_bh = tl.program_id(0), tl.program_id(1) + i_t, i_bh = tl.program_id(0).to(tl.int64), tl.program_id(1) i_b, i_h = i_bh // H, i_bh % H if IS_VARLEN: i_n, i_t = ( tl.load(chunk_indices + i_t * 2).to(tl.int32), - tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), + tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64), ) bos, eos = ( tl.load(cu_seqlens + i_n).to(tl.int32), @@ -892,108 +864,57 @@ 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,)) + o_t = i_t * BT + tl.arange(0, BT) + m_t = o_t < T + p_b = beta + bos * H + i_h + o_t * H + b_b = tl.load(p_b, mask=m_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) - ) - b_A = tl.load(p_A, boundary_check=(0, 1)) + o_A = tl.arange(0, BT) + m_A = m_t[:, None] & (o_A[None, :] < BT) + p_A = A + (bos * H + i_h) * BT + o_t[:, None] * (H * BT) + o_A[None, :] + b_A = tl.load(p_A, mask=m_A, other=0.0) 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_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)) + o_v = i_v * BV + tl.arange(0, BV) + m_v = m_t[:, None] & (o_v[None, :] < V) + p_v = v + (bos * H + i_h) * V + o_t[:, None] * (H * V) + o_v[None, :] + p_u = u + (bos * H + i_h) * V + o_t[:, None] * (H * V) + o_v[None, :] + b_v = tl.load(p_v, mask=m_v, other=0.0) b_vb = (b_v * b_b[:, None]).to(b_v.dtype) b_u = tl.dot(b_A, b_vb, input_precision=DOT_PRECISION) - tl.store(p_u, b_u.to(p_u.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_u, b_u.to(p_u.dtype.element_ty), mask=m_v) for i_k in range(tl.cdiv(K, BK)): - 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), - ) - 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), - ) - b_k = tl.load(p_k, boundary_check=(0, 1)) + o_k = i_k * BK + tl.arange(0, BK) + m_k = o_k < K + m_tk = m_t[:, None] & m_k[None, :] + p_w = w + (bos * H + i_h) * K + o_t[:, None] * (H * K) + o_k[None, :] + p_k = k + (bos * H + i_h) * K + o_t[:, None] * (H * K) + o_k[None, :] + b_k = tl.load(p_k, mask=m_tk, other=0.0) b_kb = b_k * b_b[:, None] - p_gk = tl.make_block_ptr( - gk + (bos * H + i_h) * K, - (T, K), - (H * K, 1), - (i_t * BT, i_k * BK), - (BT, BK), - (1, 0), - ) - b_gk = tl.load(p_gk, boundary_check=(0, 1)) + p_gk = gk + (bos * H + i_h) * K + o_t[:, None] * (H * K) + o_k[None, :] + b_gk = tl.load(p_gk, mask=m_tk, other=0.0) b_kb *= exp2(b_gk) if STORE_QG: - p_q = tl.make_block_ptr( - q + (bos * H + i_h) * K, - (T, K), - (H * K, 1), - (i_t * BT, i_k * BK), - (BT, BK), - (1, 0), - ) - p_qg = tl.make_block_ptr( - qg + (bos * H + i_h) * 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)) + p_q = q + (bos * H + i_h) * K + o_t[:, None] * (H * K) + o_k[None, :] + p_qg = qg + (bos * H + i_h) * K + o_t[:, None] * (H * K) + o_k[None, :] + b_q = tl.load(p_q, mask=m_tk, other=0.0) b_qg = b_q * exp2(b_gk) - tl.store(p_qg, b_qg.to(p_qg.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_qg, b_qg.to(p_qg.dtype.element_ty), mask=m_tk) if STORE_KG: last_idx = min(i_t * BT + BT, T) - 1 - o_k = i_k * BK + tl.arange(0, BK) - m_k = o_k < K b_gn = tl.load( gk + ((bos + last_idx) * H + i_h) * K + o_k, mask=m_k, other=0.0 ) b_kg = b_k * exp2(b_gn - b_gk) - p_kg = tl.make_block_ptr( - kg + (bos * H + i_h) * K, - (T, K), - (H * K, 1), - (i_t * BT, i_k * BK), - (BT, BK), - (1, 0), - ) - tl.store(p_kg, b_kg.to(p_kg.dtype.element_ty), boundary_check=(0, 1)) + p_kg = kg + (bos * H + i_h) * K + o_t[:, None] * (H * K) + o_k[None, :] + tl.store(p_kg, b_kg.to(p_kg.dtype.element_ty), mask=m_tk) 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)) + tl.store(p_w, b_w.to(p_w.dtype.element_ty), mask=m_tk) def recompute_w_u_fwd( @@ -1074,17 +995,17 @@ def chunk_gla_fwd_kernel_o( BV: tl.constexpr, IS_VARLEN: tl.constexpr, ): - i_v, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_v, i_t, i_bh = tl.program_id(0), tl.program_id(1).to(tl.int64), tl.program_id(2) i_b, i_h = i_bh // H, i_bh % H if IS_VARLEN: - i_tg = i_t + i_tg = i_t.to(tl.int64) i_n, i_t = ( tl.load(chunk_indices + i_t * 2).to(tl.int32), - tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), + tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64), ) bos, eos = ( - tl.load(cu_seqlens + i_n).to(tl.int32), - tl.load(cu_seqlens + i_n + 1).to(tl.int32), + tl.load(cu_seqlens + i_n).to(tl.int64), + tl.load(cu_seqlens + i_n + 1).to(tl.int64), ) T = eos - bos NT = tl.cdiv(T, BT) @@ -1096,70 +1017,44 @@ def chunk_gla_fwd_kernel_o( m_s = tl.arange(0, BT)[:, None] >= tl.arange(0, BT)[None, :] b_o = tl.zeros([BT, BV], dtype=tl.float32) + o_t = i_t * BT + tl.arange(0, BT) + o_v = i_v * BV + tl.arange(0, BV) + o_i = tl.arange(0, BT) + m_t = o_t < T + m_v = o_v < V + m_tv = m_t[:, None] & m_v[None, :] + m_A = m_t[:, None] & (o_i[None, :] < BT) for i_k in range(tl.cdiv(K, BK)): - p_q = tl.make_block_ptr( - q + (bos * H + i_h) * K, - (T, K), - (H * K, 1), - (i_t * BT, i_k * BK), - (BT, BK), - (1, 0), - ) - p_g = tl.make_block_ptr( - g + (bos * H + i_h) * K, - (T, K), - (H * K, 1), - (i_t * BT, i_k * BK), - (BT, BK), - (1, 0), - ) - p_h = tl.make_block_ptr( - h + (i_tg * H + i_h) * K * V, - (V, K), - (K, 1), - (i_v * BV, i_k * BK), - (BV, BK), - (1, 0), - ) + o_k = i_k * BK + tl.arange(0, BK) + m_k = o_k < K + m_qk = m_t[:, None] & m_k[None, :] + p_q = q + (bos * H + i_h) * K + o_t[:, None] * (H * K) + o_k[None, :] + p_g = g + (bos * H + i_h) * K + o_t[:, None] * (H * K) + o_k[None, :] + p_h = h + (i_tg * H + i_h) * K * V + o_v[:, None] * K + o_k[None, :] + m_h = m_v[:, None] & m_k[None, :] # [BT, BK] - b_q = tl.load(p_q, boundary_check=(0, 1)) + b_q = tl.load(p_q, mask=m_qk, other=0.0) b_q = (b_q * scale).to(b_q.dtype) # [BT, BK] - b_g = tl.load(p_g, boundary_check=(0, 1)) + b_g = tl.load(p_g, mask=m_qk, other=0.0) # [BT, BK] b_qg = (b_q * exp2(b_g)).to(b_q.dtype) # [BV, BK] - b_h = tl.load(p_h, boundary_check=(0, 1)) + b_h = tl.load(p_h, mask=m_h, other=0.0) # [BT, BV] if i_k >= 0: b_o += tl.dot(b_qg, tl.trans(b_h).to(b_qg.dtype)) - 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_o = tl.make_block_ptr( - o + (bos * H + i_h) * V, - (T, V), - (H * V, 1), - (i_t * BT, i_v * BV), - (BT, BV), - (1, 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_v = v + (bos * H + i_h) * V + o_t[:, None] * (H * V) + o_v[None, :] + p_o = o + (bos * H + i_h) * V + o_t[:, None] * (H * V) + o_v[None, :] + p_A = A + (bos * H + i_h) * BT + o_t[:, None] * (H * BT) + o_i[None, :] # [BT, BV] - b_v = tl.load(p_v, boundary_check=(0, 1)) + b_v = tl.load(p_v, mask=m_tv, other=0.0) # [BT, BT] - b_A = tl.load(p_A, boundary_check=(0, 1)) + b_A = tl.load(p_A, mask=m_A, other=0.0) b_A = tl.where(m_s, b_A, 0.0).to(b_v.dtype) b_o += tl.dot(b_A, b_v, allow_tf32=False) - tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_o, b_o.to(p_o.dtype.element_ty), mask=m_tv) def chunk_gla_fwd_o_gk( @@ -1238,12 +1133,12 @@ def kda_gate_cumsum_fwd_kernel( IS_VARLEN: tl.constexpr, USE_LOWER_BOUND: tl.constexpr, ): - i_d, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_d, i_t, i_bh = tl.program_id(0), tl.program_id(1).to(tl.int64), tl.program_id(2) i_b, i_h = i_bh // H, i_bh % H if IS_VARLEN: i_n, i_t = ( tl.load(chunk_indices + i_t * 2).to(tl.int32), - tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), + tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64), ) bos, eos = ( tl.load(cu_seqlens + i_n).to(tl.int32), @@ -1253,27 +1148,17 @@ def kda_gate_cumsum_fwd_kernel( else: bos = i_b * T - p_g = tl.make_block_ptr( - g + (bos * H + i_h) * D, - (T, D), - (H * D, 1), - (i_t * BT, i_d * BD), - (BT, BD), - (1, 0), - ) - p_y = tl.make_block_ptr( - y + (bos * H + i_h) * D, - (T, D), - (H * D, 1), - (i_t * BT, i_d * BD), - (BT, BD), - (1, 0), - ) + o_t = i_t * BT + tl.arange(0, BT) + o_d = i_d * BD + tl.arange(0, BD) + m_t = o_t < T + m_d = o_d < D + m_g = m_t[:, None] & m_d[None, :] + p_g = g + (bos * H + i_h) * D + o_t[:, None] * (H * D) + o_d[None, :] + p_y = y + (bos * H + i_h) * D + o_t[:, None] * (H * D) + o_d[None, :] - b_g = tl.load(p_g, boundary_check=(0, 1)).to(tl.float32) + b_g = tl.load(p_g, mask=m_g, other=0.0).to(tl.float32) if HAS_BIAS: - o_d = i_d * BD + tl.arange(0, BD) - b_bias = tl.load(g_bias + i_h * D + o_d, mask=o_d < D, other=0.0).to(tl.float32) + b_bias = tl.load(g_bias + i_h * D + o_d, mask=m_d, other=0.0).to(tl.float32) b_g = b_g + b_bias[None, :] b_a = tl.exp(tl.load(A + i_h).to(tl.float32)) @@ -1290,12 +1175,12 @@ def kda_gate_cumsum_fwd_kernel( # Out-of-bounds rows (load returns 0, but softplus/bias can still make # b_gate non-zero) participate in the dot product. They only contribute to - # out-of-bounds output rows, which are masked away by `boundary_check` on - # the store, so visible output matches unfused gate + chunk-local cumsum. - o_t = tl.arange(0, BT) - m_cumsum = tl.where(o_t[:, None] >= o_t[None, :], 1.0, 0.0) + # out-of-bounds output rows, which are masked away by the store mask, + # so visible output matches unfused gate + chunk-local cumsum. + o_c = tl.arange(0, BT) + m_cumsum = tl.where(o_c[:, None] >= o_c[None, :], 1.0, 0.0) b_y = tl.dot(m_cumsum, b_gate, allow_tf32=False) * cumsum_scale - tl.store(p_y, b_y.to(p_y.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_y, b_y.to(p_y.dtype.element_ty), mask=m_g) def fused_kda_gate_chunk_cumsum( @@ -1605,7 +1490,7 @@ def kda_gate_fwd_kernel( HAS_BIAS: tl.constexpr, USE_LOWER_BOUND: tl.constexpr, ): - i_t, i_h = tl.program_id(0), tl.program_id(1) + i_t, i_h = tl.program_id(0).to(tl.int64), tl.program_id(1) n_t = i_t * BT b_a = tl.exp(tl.load(A + i_h).to(tl.float32)) @@ -1613,25 +1498,14 @@ def kda_gate_fwd_kernel( stride_row = H * D stride_col = 1 - g_ptr = tl.make_block_ptr( - base=g + i_h * D, - shape=(T, D), - strides=(stride_row, stride_col), - offsets=(n_t, 0), - block_shape=(BT, BD), - order=(1, 0), - ) + o_t = n_t + tl.arange(0, BT) + o_d = tl.arange(0, BD) + m_t = o_t < T + m_g = m_t[:, None] & (o_d[None, :] < D) + g_ptr = g + i_h * D + o_t[:, None] * stride_row + o_d[None, :] * stride_col + y_ptr = y + i_h * D + o_t[:, None] * stride_row + o_d[None, :] * stride_col - y_ptr = tl.make_block_ptr( - base=y + i_h * D, - shape=(T, D), - strides=(stride_row, stride_col), - offsets=(n_t, 0), - block_shape=(BT, BD), - order=(1, 0), - ) - - b_g = tl.load(g_ptr, boundary_check=(0, 1)).to(tl.float32) + b_g = tl.load(g_ptr, mask=m_g, other=0.0).to(tl.float32) if HAS_BIAS: n_d = tl.arange(0, BD) @@ -1649,7 +1523,7 @@ def kda_gate_fwd_kernel( sp = tl.where(use_linear, b_g, (1.0 / beta) * log(1.0 + tl.exp(g_scaled))) b_y = -b_a * sp - tl.store(y_ptr, b_y.to(y.dtype.element_ty), boundary_check=(0, 1)) + tl.store(y_ptr, b_y.to(y.dtype.element_ty), mask=m_g) def fused_kda_gate( diff --git a/vllm/third_party/flash_linear_attention/ops/l2norm.py b/vllm/third_party/flash_linear_attention/ops/l2norm.py index 2eb137a242..4827df2d07 100644 --- a/vllm/third_party/flash_linear_attention/ops/l2norm.py +++ b/vllm/third_party/flash_linear_attention/ops/l2norm.py @@ -66,13 +66,18 @@ def l2norm_fwd_kernel( BT: tl.constexpr, BD: tl.constexpr, ): - i_t = tl.program_id(0) - p_x = tl.make_block_ptr(x, (T, D), (D, 1), (i_t * BT, 0), (BT, BD), (1, 0)) - b_x = tl.load(p_x, boundary_check=(0, 1)).to(tl.float32) + i_t = tl.program_id(0).to(tl.int64) + o_t = i_t * BT + tl.arange(0, BT) + o_d = tl.arange(0, BD) + m_t = o_t < T + m_x = m_t[:, None] & (o_d[None, :] < D) + p_x = x + o_t[:, None] * D + o_d[None, :] + p_y = y + o_t[:, None] * D + o_d[None, :] + + b_x = tl.load(p_x, mask=m_x, other=0.0).to(tl.float32) b_var = tl.sum(b_x * b_x, axis=1) b_y = b_x / tl.sqrt(b_var + eps)[:, None] - p_y = tl.make_block_ptr(y, (T, D), (D, 1), (i_t * BT, 0), (BT, BD), (1, 0)) - tl.store(p_y, b_y.to(p_y.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_y, b_y.to(p_y.dtype.element_ty), mask=m_x) @triton.jit diff --git a/vllm/third_party/flash_linear_attention/ops/solve_tril.py b/vllm/third_party/flash_linear_attention/ops/solve_tril.py index 8d3811ca4c..93ef770be6 100644 --- a/vllm/third_party/flash_linear_attention/ops/solve_tril.py +++ b/vllm/third_party/flash_linear_attention/ops/solve_tril.py @@ -69,12 +69,12 @@ def solve_tril_16x16_kernel( Ai = Ai + (bos * H + i_h) * 16 offset = (i_t * 16) % BT + o_t = (i_t * 16 + o_i).to(tl.int64) + m_t = o_t < T if not USE_TMA: - p_A = tl.make_block_ptr( - A, (T, BT), (H * BT, 1), (i_t * 16, offset), (16, 16), (1, 0) - ) + p_A = A + o_t[:, None] * (H * BT) + (o_i + offset)[None, :] # [16, 16] - b_A = tl.load(p_A, boundary_check=(0, 1)).to(tl.float32) + b_A = tl.load(p_A, mask=m_t[:, None], other=0.0).to(tl.float32) else: desc = make_tensor_descriptor(A, [T, BT], [H * BT, 1], [16, 16]) desc_o = make_tensor_descriptor(Ai, [T, 16], [H * 16, 1], [16, 16]) @@ -88,13 +88,11 @@ 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) - ) + p_Ai = Ai + o_t[:, None] * (H * 16) + o_i[None, :] tl.store( p_Ai, b_A.to(p_Ai.dtype.element_ty, fp_downcast_rounding="rtne"), - boundary_check=(0, 1), + mask=m_t[:, None], ) else: desc_o.store([i_t * 16, 0], b_A.to(desc_o.dtype, fp_downcast_rounding="rtne")) @@ -143,15 +141,14 @@ def merge_16x16_to_32x32_inverse_kernel( A += (bos * H + i_h) * BT Ai += (bos * H + i_h) * BT + o_t = (i_t * BT + o_i).to(tl.int64) 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_11 = A + o_t[:, None] * (H * BT) + o_i[None, :] + p_A_22 = A + (o_t[:, None] + 16) * (H * BT) + (o_i[None, :] + 16) + b_Ai_11 = tl.load(p_A_11, mask=(o_t[:, None] < T), other=0.0).to(tl.float32) + b_Ai_22 = tl.load(p_A_22, mask=(o_t[:, None] + 16 < T), other=0.0).to( + tl.float32 ) - 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) 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]) @@ -175,10 +172,8 @@ 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) + p_A_21 = A + (o_t[:, None] + 16) * (H * BT) + o_i[None, :] + b_A_21 = tl.load(p_A_21, mask=(o_t[:, None] + 16 < T), other=0.0).to(tl.float32) else: b_A_21 = desc.load([i_t * BT + 16, 0]).to(tl.float32) @@ -189,29 +184,23 @@ def merge_16x16_to_32x32_inverse_kernel( ) 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_11 = Ai + o_t[:, None] * (H * BT) + o_i[None, :] + p_Ai_21 = Ai + (o_t[:, None] + 16) * (H * BT) + o_i[None, :] + p_Ai_22 = Ai + (o_t[:, None] + 16) * (H * BT) + (o_i[None, :] + 16) tl.store( p_Ai_11, b_Ai_11.to(p_Ai_11.dtype.element_ty, fp_downcast_rounding="rtne"), - boundary_check=(0, 1), + mask=(o_t[:, None] < T), ) tl.store( p_Ai_22, b_Ai_22.to(p_Ai_22.dtype.element_ty, fp_downcast_rounding="rtne"), - boundary_check=(0, 1), + mask=(o_t[:, None] + 16 < T), ) tl.store( p_Ai_21, b_Ai_21.to(p_Ai_21.dtype.element_ty, fp_downcast_rounding="rtne"), - boundary_check=(0, 1), + mask=(o_t[:, None] + 16 < T), ) else: desc_o.store( @@ -268,23 +257,22 @@ def merge_16x16_to_64x64_inverse_kernel( A += (bos * H + i_h) * BT Ai += (bos * H + i_h) * BT + o_t = (i_t * BT + o_i).to(tl.int64) 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_11 = A + o_t[:, None] * (H * BT) + o_i[None, :] + p_A_22 = A + (o_t[:, None] + 16) * (H * BT) + (o_i[None, :] + 16) + p_A_33 = A + (o_t[:, None] + 32) * (H * BT) + (o_i[None, :] + 32) + p_A_44 = A + (o_t[:, None] + 48) * (H * BT) + (o_i[None, :] + 48) + b_Ai_11 = tl.load(p_A_11, mask=(o_t[:, None] < T), other=0.0).to(tl.float32) + b_Ai_22 = tl.load(p_A_22, mask=(o_t[:, None] + 16 < T), other=0.0).to( + tl.float32 ) - p_A_22 = tl.make_block_ptr( - A, (T, BT), (H * BT, 1), (i_t * BT + 16, 16), (16, 16), (1, 0) + b_Ai_33 = tl.load(p_A_33, mask=(o_t[:, None] + 32 < T), other=0.0).to( + tl.float32 ) - p_A_33 = tl.make_block_ptr( - A, (T, BT), (H * BT, 1), (i_t * BT + 32, 32), (16, 16), (1, 0) + b_Ai_44 = tl.load(p_A_44, mask=(o_t[:, None] + 48 < T), other=0.0).to( + tl.float32 ) - 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) 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]) @@ -321,30 +309,18 @@ 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) + p_A_21 = A + (o_t[:, None] + 16) * (H * BT) + o_i[None, :] + p_A_31 = A + (o_t[:, None] + 32) * (H * BT) + o_i[None, :] + p_A_32 = A + (o_t[:, None] + 32) * (H * BT) + (o_i[None, :] + 16) + p_A_41 = A + (o_t[:, None] + 48) * (H * BT) + o_i[None, :] + p_A_42 = A + (o_t[:, None] + 48) * (H * BT) + (o_i[None, :] + 16) + p_A_43 = A + (o_t[:, None] + 48) * (H * BT) + (o_i[None, :] + 32) + b_A_21 = tl.load(p_A_21, mask=(o_t[:, None] + 16 < T), other=0.0).to(tl.float32) + b_A_31 = tl.load(p_A_31, mask=(o_t[:, None] + 32 < T), other=0.0).to(tl.float32) + b_A_32 = tl.load(p_A_32, mask=(o_t[:, None] + 32 < T), other=0.0).to(tl.float32) + b_A_41 = tl.load(p_A_41, mask=(o_t[:, None] + 48 < T), other=0.0).to(tl.float32) + b_A_42 = tl.load(p_A_42, mask=(o_t[:, None] + 48 < T), other=0.0).to(tl.float32) + b_A_43 = tl.load(p_A_43, mask=(o_t[:, None] + 48 < T), other=0.0).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) @@ -390,85 +366,65 @@ def merge_16x16_to_64x64_inverse_kernel( ) 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) - ) + p_Ai_11 = Ai + o_t[:, None] * (H * BT) + o_i[None, :] + p_Ai_22 = Ai + (o_t[:, None] + 16) * (H * BT) + (o_i[None, :] + 16) + p_Ai_33 = Ai + (o_t[:, None] + 32) * (H * BT) + (o_i[None, :] + 32) + p_Ai_44 = Ai + (o_t[:, None] + 48) * (H * BT) + (o_i[None, :] + 48) + p_Ai_21 = Ai + (o_t[:, None] + 16) * (H * BT) + o_i[None, :] + p_Ai_31 = Ai + (o_t[:, None] + 32) * (H * BT) + o_i[None, :] + p_Ai_32 = Ai + (o_t[:, None] + 32) * (H * BT) + (o_i[None, :] + 16) + p_Ai_41 = Ai + (o_t[:, None] + 48) * (H * BT) + o_i[None, :] + p_Ai_42 = Ai + (o_t[:, None] + 48) * (H * BT) + (o_i[None, :] + 16) + p_Ai_43 = Ai + (o_t[:, None] + 48) * (H * BT) + (o_i[None, :] + 32) tl.store( p_Ai_11, b_Ai_11.to(p_Ai_11.dtype.element_ty, fp_downcast_rounding="rtne"), - boundary_check=(0, 1), + mask=(o_t[:, None] < T), ) tl.store( p_Ai_22, b_Ai_22.to(p_Ai_22.dtype.element_ty, fp_downcast_rounding="rtne"), - boundary_check=(0, 1), + mask=(o_t[:, None] + 16 < T), ) tl.store( p_Ai_33, b_Ai_33.to(p_Ai_33.dtype.element_ty, fp_downcast_rounding="rtne"), - boundary_check=(0, 1), + mask=(o_t[:, None] + 32 < T), ) tl.store( p_Ai_44, b_Ai_44.to(p_Ai_44.dtype.element_ty, fp_downcast_rounding="rtne"), - boundary_check=(0, 1), + mask=(o_t[:, None] + 48 < T), ) tl.store( p_Ai_21, b_Ai_21.to(p_Ai_21.dtype.element_ty, fp_downcast_rounding="rtne"), - boundary_check=(0, 1), + mask=(o_t[:, None] + 16 < T), ) tl.store( p_Ai_31, b_Ai_31.to(p_Ai_31.dtype.element_ty, fp_downcast_rounding="rtne"), - boundary_check=(0, 1), + mask=(o_t[:, None] + 32 < T), ) tl.store( p_Ai_32, b_Ai_32.to(p_Ai_32.dtype.element_ty, fp_downcast_rounding="rtne"), - boundary_check=(0, 1), + mask=(o_t[:, None] + 32 < T), ) tl.store( p_Ai_41, b_Ai_41.to(p_Ai_41.dtype.element_ty, fp_downcast_rounding="rtne"), - boundary_check=(0, 1), + mask=(o_t[:, None] + 48 < T), ) tl.store( p_Ai_42, b_Ai_42.to(p_Ai_42.dtype.element_ty, fp_downcast_rounding="rtne"), - boundary_check=(0, 1), + mask=(o_t[:, None] + 48 < T), ) tl.store( p_Ai_43, b_Ai_43.to(p_Ai_43.dtype.element_ty, fp_downcast_rounding="rtne"), - boundary_check=(0, 1), + mask=(o_t[:, None] + 48 < T), ) else: desc_o.store( diff --git a/vllm/third_party/flash_linear_attention/ops/wy_fast.py b/vllm/third_party/flash_linear_attention/ops/wy_fast.py index 52d2b28195..040c37ebfe 100644 --- a/vllm/third_party/flash_linear_attention/ops/wy_fast.py +++ b/vllm/third_party/flash_linear_attention/ops/wy_fast.py @@ -46,12 +46,12 @@ def recompute_w_u_fwd_kernel( BV: tl.constexpr, IS_VARLEN: tl.constexpr, ): - i_t, i_bh = tl.program_id(0), tl.program_id(1) + i_t, i_bh = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64) i_b, i_h = i_bh // H, i_bh % H if IS_VARLEN: i_n, i_t = ( tl.load(chunk_indices + i_t * 2).to(tl.int32), - tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), + tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64), ) bos, eos = ( tl.load(cu_seqlens + i_n).to(tl.int32), @@ -60,60 +60,41 @@ def recompute_w_u_fwd_kernel( T = eos - bos else: bos, eos = i_b * T, i_b * T + T - 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 + (bos * H + i_h), (T,), (H,), (i_t * BT,), (BT,), (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) - ) - b_beta = tl.load(p_beta, boundary_check=(0,)) - b_A = tl.load(p_A, boundary_check=(0, 1)) - b_g = tl.exp(tl.load(p_g, boundary_check=(0,))) + o_t = i_t * BT + tl.arange(0, BT) + o_A = tl.arange(0, BT) + m_t = o_t < T + m_A = m_t[:, None] & (o_A[None, :] < BT) + p_beta = beta + bos * H + i_h + o_t * H + p_g = g + (bos * H + i_h) + o_t * H + p_A = A + (bos * H + i_h) * BT + o_t[:, None] * (H * BT) + o_A[None, :] + b_beta = tl.load(p_beta, mask=m_t, other=0.0) + b_A = tl.load(p_A, mask=m_A, other=0.0) + b_g = tl.exp(tl.load(p_g, mask=m_t, other=0.0)) 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_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)) + o_v = i_v * BV + tl.arange(0, BV) + m_v = m_t[:, None] & (o_v[None, :] < V) + p_v = v + (bos * H + i_h) * V + o_t[:, None] * (H * V) + o_v[None, :] + p_u = u + (bos * H + i_h) * V + o_t[:, None] * (H * V) + o_v[None, :] + b_v = tl.load(p_v, mask=m_v, other=0.0) b_vb = (b_v * b_beta[:, 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)) + tl.store(p_u, b_u.to(p_u.dtype.element_ty), mask=m_v) 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 = m_t[:, None] & (o_k[None, :] < K) + p_k = ( + k + + (bos * Hg + i_h // (H // Hg)) * K + + o_t[:, None] * (Hg * K) + + o_k[None, :] ) - 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)) + p_w = w + (bos * H + i_h) * K + o_t[:, None] * (H * K) + o_k[None, :] + b_k = tl.load(p_k, mask=m_k, other=0.0) b_kb = (b_k * b_beta[:, None] * b_g[:, None]).to(b_k.dtype) b_w = tl.dot(b_A, b_kb) - tl.store(p_w, b_w.to(p_w.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_w, b_w.to(p_w.dtype.element_ty), mask=m_k) def recompute_w_u_fwd(