iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[Triton] Migrate Triton kernels off tl.make_block_ptr for triton-main compat (#60502)

Browse Source 
Signed-off-by: Nick Riasanovsky <njriasan@meta.com>
This commit is contained in:
Nick Riasanovsky
2026-10-12 11:38:43 -04:00
committed by GitHub
parent ffec221276
commit a686c8e1b7
16 changed files with 836 additions and 1515 deletions
Download Patch File Download Diff File Expand all files Collapse all files
vllm/models/glm5next/amd/ops/third_party/kda/kernels.py
Vendored
+116 -241
View File
@@ -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(
vllm/models/glm5next/nvidia/ops/third_party/kda/kernels.py
Vendored
+116 -241
View File
@@ -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(
vllm/models/inkling/amd/ops/fa4_rel_attention.py
+24 -16
View File
@@ -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,
)
vllm/models/kimi_k3/nvidia/ops/third_party/kda/chunk.py
Vendored
+75 -185
View File
@@ -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(
vllm/models/kimi_k3/nvidia/ops/third_party/kda/chunk_intra.py
Vendored
+103 -167
View File
@@ -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(
vllm/models/kimi_k3/nvidia/ops/third_party/kda/chunk_intra_token_parallel.py
Vendored
+6 -10
View File
@@ -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
vllm/models/minimax_m3/common/ops/index_topk.py
+10 -8
View File
@@ -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
vllm/third_party/flash_linear_attention/ops/chunk_delta_h.py
Vendored
+71 -111
View File
@@ -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(
vllm/third_party/flash_linear_attention/ops/chunk_o.py
Vendored
+23 -26
View File
@@ -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(
vllm/third_party/flash_linear_attention/ops/chunk_scaled_dot_kkt.py
Vendored
+15 -20
View File
@@ -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(
vllm/third_party/flash_linear_attention/ops/cumsum.py
Vendored
+21 -48
View File
@@ -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(
vllm/third_party/flash_linear_attention/ops/fused_norm_gate.py
Vendored
+16 -18
View File
@@ -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(
vllm/third_party/flash_linear_attention/ops/kda.py
Vendored
+138 -264
View File
@@ -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(
vllm/third_party/flash_linear_attention/ops/l2norm.py
Vendored
+10 -5
View File
@@ -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
vllm/third_party/flash_linear_attention/ops/solve_tril.py
Vendored
+64 -108
View File
@@ -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(
vllm/third_party/flash_linear_attention/ops/wy_fast.py
Vendored
+28 -47
View File
@@ -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(