Skip to content

[GDN] Add GVA support #799

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
zhiyuan1i merged 7 commits into from
Mar 31, 2026
Merged

[GDN] Add GVA support #799

Show file tree
Hide file tree
Changes from 3 commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
44563eb
[GDN] Add GQA support (Hq != H)
zhiyuan1i Mar 29, 2026
5ce7930
remove unused kernels
zhiyuan1i Mar 29, 2026
2b2b0e8
[GDN] Address review comments for GQA support
zhiyuan1i Mar 29, 2026
012c473
enhance ci
zhiyuan1i Mar 30, 2026
c7c7a5f
Merge branch 'main' into lzy/gdn-different-kvheads
yzhangcs Mar 30, 2026
78434d1
enhance ci
zhiyuan1i Mar 31, 2026
d412c84
Merge branch 'main' into lzy/gdn-different-kvheads
zhiyuan1i Mar 31, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
467 changes: 0 additions & 467 deletions benchmarks/cp/benchmark_chunk_delta_h_kernels.py
View file
Open in desktop

This file was deleted.

42 changes: 25 additions & 17 deletions fla/ops/common/chunk_delta_h.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -46,6 +46,7 @@ def chunk_gated_delta_rule_fwd_kernel_h_blockdim64(
chunk_offsets,
T,
H: tl.constexpr,
Hq: tl.constexpr,
K: tl.constexpr,
V: tl.constexpr,
BT: tl.constexpr,
Expand Down Expand Up @@ -91,7 +92,7 @@ def chunk_gated_delta_rule_fwd_kernel_h_blockdim64(
# calculate offset
h += (boh * H + i_h).to(tl.int64) * K*V
v += (bos * H + i_h).to(tl.int64) * V
k += (bos * H + i_h).to(tl.int64) * K
k += (bos * Hq + i_h // (H // Hq)).to(tl.int64) * K
w += (bos * H + i_h).to(tl.int64) * K
if SAVE_NEW_VALUE:
v_new += (bos * H + i_h).to(tl.int64) * V
Expand Down Expand Up @@ -263,28 +264,28 @@ def chunk_gated_delta_rule_fwd_kernel_h_blockdim64(

b_v = b_v.to(k.dtype.element_ty)

p_k = tl.make_block_ptr(k, (K, T), (1, H*K), (0, i_t * BT), (64, BT), (0, 1))
p_k = tl.make_block_ptr(k, (K, T), (1, Hq*K), (0, i_t * BT), (64, BT), (0, 1))
b_k = tl.load(p_k, boundary_check=(0, 1))
if TRANSPOSE_STATE:
b_h1 += tl.trans(tl.dot(b_k, b_v))
else:
b_h1 += tl.dot(b_k, b_v)
if K > 64:
p_k = tl.make_block_ptr(k, (K, T), (1, H*K), (64, i_t * BT), (64, BT), (0, 1))
p_k = tl.make_block_ptr(k, (K, T), (1, Hq*K), (64, i_t * BT), (64, BT), (0, 1))
b_k = tl.load(p_k, boundary_check=(0, 1))
if TRANSPOSE_STATE:
b_h2 += tl.trans(tl.dot(b_k, b_v))
else:
b_h2 += tl.dot(b_k, b_v)
if K > 128:
p_k = tl.make_block_ptr(k, (K, T), (1, H*K), (128, i_t * BT), (64, BT), (0, 1))
p_k = tl.make_block_ptr(k, (K, T), (1, Hq*K), (128, i_t * BT), (64, BT), (0, 1))
b_k = tl.load(p_k, boundary_check=(0, 1))
if TRANSPOSE_STATE:
b_h3 += tl.trans(tl.dot(b_k, b_v))
else:
b_h3 += tl.dot(b_k, b_v)
if K > 192:
p_k = tl.make_block_ptr(k, (K, T), (1, H*K), (192, i_t * BT), (64, BT), (0, 1))
p_k = tl.make_block_ptr(k, (K, T), (1, Hq*K), (192, i_t * BT), (64, BT), (0, 1))
b_k = tl.load(p_k, boundary_check=(0, 1))
if TRANSPOSE_STATE:
b_h4 += tl.trans(tl.dot(b_k, b_v))
Expand Down Expand Up @@ -353,6 +354,7 @@ def chunk_gated_delta_rule_bwd_kernel_dhu_blockdim64(
scale,
T,
H: tl.constexpr,
Hq: tl.constexpr,
K: tl.constexpr,
V: tl.constexpr,
BT: tl.constexpr,
Expand Down Expand Up @@ -395,8 +397,8 @@ def chunk_gated_delta_rule_bwd_kernel_dhu_blockdim64(
b_dh4 = tl.zeros([64, BV], dtype=tl.float32)

# calculate offset
q += (bos * H + i_h).to(tl.int64) * K
k += (bos * H + i_h).to(tl.int64) * K
q += (bos * Hq + i_h // (H // Hq)).to(tl.int64) * K
k += (bos * Hq + i_h // (H // Hq)).to(tl.int64) * K
w += (bos * H + i_h).to(tl.int64) * K
do += (bos * H + i_h).to(tl.int64) * V
dv += (bos * H + i_h).to(tl.int64) * V
Expand Down Expand Up @@ -480,7 +482,7 @@ def chunk_gated_delta_rule_bwd_kernel_dhu_blockdim64(
b_do = tl.load(p_do, boundary_check=(0, 1))

# Update dv
p_k = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_t * BT, 0), (BT, 64), (1, 0))
p_k = tl.make_block_ptr(k, (T, K), (Hq*K, 1), (i_t * BT, 0), (BT, 64), (1, 0))
b_k = tl.load(p_k, boundary_check=(0, 1))
if USE_GK:
o_k1 = tl.arange(0, 64)
Expand All @@ -491,7 +493,7 @@ def chunk_gated_delta_rule_bwd_kernel_dhu_blockdim64(
b_dv = tl.dot(b_k, b_dh1.to(b_k.dtype))

if K > 64:
p_k = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_t * BT, 64), (BT, 64), (1, 0))
p_k = tl.make_block_ptr(k, (T, K), (Hq*K, 1), (i_t * BT, 64), (BT, 64), (1, 0))
b_k = tl.load(p_k, boundary_check=(0, 1))
if USE_GK:
o_k2 = 64 + o_k1
Expand All @@ -502,7 +504,7 @@ def chunk_gated_delta_rule_bwd_kernel_dhu_blockdim64(
b_dv += tl.dot(b_k, b_dh2.to(b_k.dtype))

if K > 128:
p_k = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_t * BT, 128), (BT, 64), (1, 0))
p_k = tl.make_block_ptr(k, (T, K), (Hq*K, 1), (i_t * BT, 128), (BT, 64), (1, 0))
b_k = tl.load(p_k, boundary_check=(0, 1))
if USE_GK:
o_k3 = 128 + o_k1
Expand All @@ -513,7 +515,7 @@ def chunk_gated_delta_rule_bwd_kernel_dhu_blockdim64(
b_dv += tl.dot(b_k, b_dh3.to(b_k.dtype))

if K > 192:
p_k = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_t * BT, 192), (BT, 64), (1, 0))
p_k = tl.make_block_ptr(k, (T, K), (Hq*K, 1), (i_t * BT, 192), (BT, 64), (1, 0))
b_k = tl.load(p_k, boundary_check=(0, 1))
if USE_GK:
o_k4 = 192 + o_k1
Expand All @@ -534,7 +536,7 @@ def chunk_gated_delta_rule_bwd_kernel_dhu_blockdim64(
tl.store(p_dv2, b_dv.to(p_dv.dtype.element_ty), boundary_check=(0, 1))
# Update dh
p_w = tl.make_block_ptr(w, (K, T), (1, H*K), (0, i_t * BT), (64, BT), (0, 1))
p_q = tl.make_block_ptr(q, (K, T), (1, H*K), (0, i_t * BT), (64, BT), (0, 1))
p_q = tl.make_block_ptr(q, (K, T), (1, Hq*K), (0, i_t * BT), (64, BT), (0, 1))
b_w = tl.load(p_w, boundary_check=(0, 1))
b_q = tl.load(p_q, boundary_check=(0, 1))
if USE_G:
Expand All @@ -556,7 +558,7 @@ def chunk_gated_delta_rule_bwd_kernel_dhu_blockdim64(
else:
b_dh1 += tl.dot(b_q.to(b_q.dtype), b_do.to(b_q.dtype)) * scale - tl.dot(b_w, b_dv.to(b_w.dtype))
if K > 64:
p_q = tl.make_block_ptr(q, (K, T), (1, H*K), (64, i_t * BT), (64, BT), (0, 1))
p_q = tl.make_block_ptr(q, (K, T), (1, Hq*K), (64, i_t * BT), (64, BT), (0, 1))
p_w = tl.make_block_ptr(w, (K, T), (1, H*K), (64, i_t * BT), (64, BT), (0, 1))
b_q = tl.load(p_q, boundary_check=(0, 1))
b_w = tl.load(p_w, boundary_check=(0, 1))
Expand All @@ -579,7 +581,7 @@ def chunk_gated_delta_rule_bwd_kernel_dhu_blockdim64(
else:
b_dh2 += tl.dot(b_q.to(b_q.dtype), b_do.to(b_q.dtype)) * scale - tl.dot(b_w, b_dv.to(b_w.dtype))
if K > 128:
p_q = tl.make_block_ptr(q, (K, T), (1, H*K), (128, i_t * BT), (64, BT), (0, 1))
p_q = tl.make_block_ptr(q, (K, T), (1, Hq*K), (128, i_t * BT), (64, BT), (0, 1))
p_w = tl.make_block_ptr(w, (K, T), (1, H*K), (128, i_t * BT), (64, BT), (0, 1))
b_q = tl.load(p_q, boundary_check=(0, 1))
b_w = tl.load(p_w, boundary_check=(0, 1))
Expand All @@ -602,7 +604,7 @@ def chunk_gated_delta_rule_bwd_kernel_dhu_blockdim64(
else:
b_dh3 += tl.dot(b_q.to(b_q.dtype), b_do.to(b_q.dtype)) * scale - tl.dot(b_w, b_dv.to(b_w.dtype))
if K > 192:
p_q = tl.make_block_ptr(q, (K, T), (1, H*K), (192, i_t * BT), (64, BT), (0, 1))
p_q = tl.make_block_ptr(q, (K, T), (1, Hq*K), (192, i_t * BT), (64, BT), (0, 1))
p_w = tl.make_block_ptr(w, (K, T), (1, H*K), (192, i_t * BT), (64, BT), (0, 1))
b_q = tl.load(p_q, boundary_check=(0, 1))
b_w = tl.load(p_w, boundary_check=(0, 1))
Expand Down Expand Up @@ -668,7 +670,9 @@ def chunk_gated_delta_rule_fwd_h(
use_exp2: bool = False,
transpose_state_layout: bool = False,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]:
B, T, H, K, V = *k.shape, u.shape[-1]
B, T, Hq, K = k.shape
V = u.shape[-1]
H = u.shape[2]
Comment on lines +673 to +675

@coderabbitai coderabbitai Bot Mar 31, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🔴 Critical

Validate H/Hq before launching the remapped-head kernels.

Line 95 and Lines 400-401 use i_h // (H // Hq), but these wrappers never enforce the GQA invariant. If Hq > H, the divisor becomes zero; if H % Hq != 0, the mapped head can reach Hq, so the q/k base pointers walk past the (B, T, Hq, K) buffers. fla/ops/gated_delta_rule/chunk.py:390-399 already guards this, so these lower-level entry points should do the same.

🛡️ Proposed fix 
 def chunk_gated_delta_rule_fwd_h(
     k: torch.Tensor,
     w: torch.Tensor,
     u: torch.Tensor,
@@
 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]:
     B, T, Hq, K = k.shape
     V = u.shape[-1]
     H = u.shape[2]
+    if H % Hq != 0:
+        raise RuntimeError(
+            f"H (num_heads={H}) must be divisible by "
+            f"num_kv_heads (Hq={Hq}), but got H % Hq = {H % Hq}"
+        )
     BT = chunk_size
@@
 def chunk_gated_delta_rule_bwd_dhu(
     q: torch.Tensor,
     k: torch.Tensor,
     w: torch.Tensor,
@@
 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
     B, T, Hq, K = q.shape
     V = do.shape[-1]
     H = do.shape[2]
+    if H % Hq != 0:
+        raise RuntimeError(
+            f"H (num_heads={H}) must be divisible by "
+            f"num_kv_heads (Hq={Hq}), but got H % Hq = {H % Hq}"
+        )
     # N: the actual number of sequences in the batch with either equal or variable lengths

Also applies to: 737-739

🤖 Prompt for AI Agents 
Verify each finding against the current code and only fix it if needed.

In `@fla/ops/common/chunk_delta_h.py` around lines 673 - 675, Before launching the
remapped-head kernels, validate the head counts: ensure Hq <= H and H % Hq == 0
(the same invariant enforced in fla/ops/gated_delta_rule/chunk.py:390-399).
Specifically, where B, T, Hq, K = k.shape and H = u.shape[2] and where the
mapping uses i_h // (H // Hq), add a guard that raises/returns a clear error (or
assert) if Hq > H or H % Hq != 0 to prevent division by zero and out-of-bounds
base pointer walks; apply the same check in both entry points that set H/Hq (the
occurrences around the B,T,Hq,K and V,u/H assignments).

BT = chunk_size

if chunk_indices is None and cu_seqlens is not None:
Expand Down Expand Up @@ -703,6 +707,7 @@ def grid(meta): return (triton.cdiv(V, meta['BV']), N*H)
chunk_offsets=chunk_offsets,
T=T,
H=H,
Hq=Hq,
K=K,
V=V,
BT=BT,
Expand All @@ -729,7 +734,9 @@ def chunk_gated_delta_rule_bwd_dhu(
use_exp2: bool = False,
transpose_state_layout: bool = False,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
B, T, H, K, V = *q.shape, do.shape[-1]
B, T, Hq, K = q.shape
V = do.shape[-1]
H = do.shape[2]
# N: the actual number of sequences in the batch with either equal or variable lengths
BT = 64
assert K <= 256, "current kernel does not support head dimension being larger than 256."
Expand Down Expand Up @@ -766,6 +773,7 @@ def grid(meta): return (triton.cdiv(V, meta['BV']), N*H)
scale=scale,
T=T,
H=H,
Hq=Hq,
K=K,
V=V,
BT=BT,
Expand Down
63 changes: 41 additions & 22 deletions fla/ops/common/chunk_o.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -40,6 +40,7 @@ def chunk_fwd_kernel_o(
scale,
T,
H: tl.constexpr,
Hq: tl.constexpr,
K: tl.constexpr,
V: tl.constexpr,
BT: tl.constexpr,
Expand All @@ -66,8 +67,8 @@ def chunk_fwd_kernel_o(
bos, eos = i_b * T, i_b * T + T

# offset calculation
q += (bos * H + i_h) * K
k += (bos * H + i_h) * K
q += (bos * Hq + i_h // (H // Hq)) * K
k += (bos * Hq + i_h // (H // Hq)) * K
v += (bos * H + i_h) * V
o += (bos * H + i_h) * V
h += (i_tg * H + i_h).to(tl.int64) * K*V
Expand All @@ -76,8 +77,8 @@ def chunk_fwd_kernel_o(
b_A = tl.zeros([BT, BT], dtype=tl.float32)

for i_k in range(tl.cdiv(K, BK)):
p_q = tl.make_block_ptr(q, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0))
p_k = tl.make_block_ptr(k, (K, T), (1, H*K), (i_k * BK, i_t * BT), (BK, BT), (0, 1))
p_q = tl.make_block_ptr(q, (T, K), (Hq*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0))
p_k = tl.make_block_ptr(k, (K, T), (1, Hq*K), (i_k * BK, i_t * BT), (BK, BT), (0, 1))
if TRANSPOSE_STATE:
p_h = tl.make_block_ptr(h, (V, K), (K, 1), (i_v * BV, i_k * BK), (BV, BK), (1, 0))
else:
Expand Down Expand Up @@ -168,6 +169,7 @@ def chunk_bwd_kernel_dqkwg(
B: tl.constexpr,
T,
H: tl.constexpr,
Hq: tl.constexpr,
K: tl.constexpr,
V: tl.constexpr,
BT: tl.constexpr,
Expand Down Expand Up @@ -200,8 +202,8 @@ def chunk_bwd_kernel_dqkwg(
do += (bos * H + i_h) * V
h += (i_tg * H + i_h).to(tl.int64) * K*V
dh += (i_tg * H + i_h).to(tl.int64) * K*V
q += (bos * H + i_h) * K
k += (bos * H + i_h) * K
q += (bos * Hq + i_h // (H // Hq)) * K
k += (bos * Hq + i_h // (H // Hq)) * K
dq += (bos * H + i_h) * K
dk += (bos * H + i_h) * K

Expand Down Expand Up @@ -255,8 +257,8 @@ def chunk_bwd_kernel_dqkwg(
tl.store(p_dw, -b_dw.to(p_dw.dtype.element_ty), boundary_check=(0, 1))

tl.debug_barrier()
p_q = tl.make_block_ptr(q, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0))
p_k = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0))
p_q = tl.make_block_ptr(q, (T, K), (Hq*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0))
p_k = tl.make_block_ptr(k, (T, K), (Hq*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0))
b_q = tl.load(p_q, boundary_check=(0, 1))
b_k = tl.load(p_k, boundary_check=(0, 1))

Expand Down Expand Up @@ -362,6 +364,7 @@ def chunk_bwd_kernel_dv(
scale,
T,
H: tl.constexpr,
Hq: tl.constexpr,
K: tl.constexpr,
V: tl.constexpr,
BT: tl.constexpr,
Expand All @@ -388,16 +391,16 @@ def chunk_bwd_kernel_dv(
b_dv = tl.zeros([BT, BV], dtype=tl.float32)

# offset calculation
q += (bos * H + i_h) * K
k += (bos * H + i_h) * K
q += (bos * Hq + i_h // (H // Hq)) * K
k += (bos * Hq + i_h // (H // Hq)) * K
do += (bos * H + i_h) * V
dv += (bos * H + i_h) * V
dh += (i_tg * H + i_h).to(tl.int64) * K*V

b_A = tl.zeros([BT, BT], dtype=tl.float32)
for i_k in range(tl.cdiv(K, BK)):
p_k = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0))
p_q = tl.make_block_ptr(q, (K, T), (1, H*K), (i_k * BK, i_t * BT), (BK, BT), (0, 1))
p_k = tl.make_block_ptr(k, (T, K), (Hq*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0))
p_q = tl.make_block_ptr(q, (K, T), (1, Hq*K), (i_k * BK, i_t * BT), (BK, BT), (0, 1))
b_q = tl.load(p_q, boundary_check=(0, 1))
b_k = tl.load(p_k, boundary_check=(0, 1))
b_A += tl.dot(b_k, b_q)
Expand Down Expand Up @@ -463,6 +466,7 @@ def chunk_bwd_kernel_dv_local(
scale,
T,
H: tl.constexpr,
Hq: tl.constexpr,
K: tl.constexpr,
V: tl.constexpr,
BT: tl.constexpr,
Expand All @@ -484,8 +488,8 @@ def chunk_bwd_kernel_dv_local(
bos, eos = i_b * T, i_b * T + T

# offset calculation
q += (bos * H + i_h) * K
k += (bos * H + i_h) * K
q += (bos * Hq + i_h // (H // Hq)) * K
k += (bos * Hq + i_h // (H // Hq)) * K
do += (bos * H + i_h) * V
dv += (bos * H + i_h) * V

Expand All @@ -503,8 +507,8 @@ def chunk_bwd_kernel_dv_local(

b_A = tl.zeros([BT, BT], dtype=tl.float32)
for i_k in range(tl.cdiv(K, BK)):
p_k = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0))
p_q = tl.make_block_ptr(q, (K, T), (1, H*K), (i_k * BK, i_t * BT), (BK, BT), (0, 1))
p_k = tl.make_block_ptr(k, (T, K), (Hq*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0))
p_q = tl.make_block_ptr(q, (K, T), (1, Hq*K), (i_k * BK, i_t * BT), (BK, BT), (0, 1))

b_k = tl.load(p_k, boundary_check=(0, 1))
b_q = tl.load(p_q, boundary_check=(0, 1))
Expand Down Expand Up @@ -542,7 +546,9 @@ def chunk_fwd_o(
use_exp2: bool = False,
transpose_state_layout: bool = False,
) -> torch.Tensor:
B, T, H, K, V = *q.shape, v.shape[-1]
B, T, Hq, K = q.shape
V = v.shape[-1]
H = v.shape[2]
BT = chunk_size
if chunk_indices is None and cu_seqlens is not None:
chunk_indices = prepare_chunk_indices(cu_seqlens, BT)
Expand All @@ -565,6 +571,7 @@ def grid(meta): return (triton.cdiv(V, meta['BV']), NT, B * H)
scale=scale,
T=T,
H=H,
Hq=Hq,
K=K,
V=V,
BT=BT,
Expand All @@ -587,7 +594,9 @@ def chunk_bwd_dv(
chunk_indices: torch.LongTensor | None = None,
use_exp2: bool = False,
) -> torch.Tensor:
B, T, H, K, V = *k.shape, do.shape[-1]
B, T, Hq, K = k.shape
V = do.shape[-1]
H = do.shape[2]
BT = chunk_size
if chunk_indices is None and cu_seqlens is not None:
chunk_indices = prepare_chunk_indices(cu_seqlens, BT)
Expand Down Expand Up @@ -620,6 +629,7 @@ def chunk_bwd_dv(
scale=scale,
T=T,
H=H,
Hq=Hq,
K=K,
V=V,
BT=BT,
Expand All @@ -643,7 +653,9 @@ def chunk_bwd_dv_local(
chunk_indices: torch.LongTensor | None = None,
use_exp2: bool = False,
) -> torch.Tensor:
B, T, H, K, V = *k.shape, do.shape[-1]
B, T, Hq, K = k.shape
V = do.shape[-1]
H = do.shape[2]
BT = chunk_size
if chunk_indices is None and cu_seqlens is not None:
chunk_indices = prepare_chunk_indices(cu_seqlens, BT)
Expand Down Expand Up @@ -673,6 +685,7 @@ def chunk_bwd_dv_local(
scale=scale,
T=T,
H=H,
Hq=Hq,
K=K,
V=V,
BT=BT,
Expand Down Expand Up @@ -702,7 +715,9 @@ def chunk_bwd_dqkwg(
transpose_state_layout: bool = False,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:

B, T, H, K, V = *k.shape, v.shape[-1]
B, T, Hq, K = k.shape
V = v.shape[-1]
H = v.shape[2]
BT = chunk_size
if chunk_indices is None and cu_seqlens is not None:
chunk_indices = prepare_chunk_indices(cu_seqlens, BT)
Expand All @@ -717,8 +732,8 @@ def chunk_bwd_dqkwg(
BK = min(max(triton.next_power_of_2(K), 16), CONST_TILING)
BV = min(max(triton.next_power_of_2(V), 16), CONST_TILING)
NK = triton.cdiv(K, BK)
dq = torch.empty_like(q)
dk = torch.empty_like(k)
dq = q.new_empty(B, T, H, K)
dk = k.new_empty(B, T, H, K)
dg = torch.empty(NK, *g.shape, dtype=torch.float32, device=g.device) if g is not None else None
dw = torch.empty_like(w) if w is not None else None

Expand All @@ -743,6 +758,7 @@ def chunk_bwd_dqkwg(
B=B,
T=T,
H=H,
Hq=Hq,
K=K,
V=V,
BT=BT,
Expand All @@ -752,6 +768,9 @@ def chunk_bwd_dqkwg(
TRANSPOSE_STATE=transpose_state_layout,
)

if Hq != H:
dq = dq.view(B, T, Hq, H // Hq, K).sum(3)
dk = dk.view(B, T, Hq, H // Hq, K).sum(3)
Comment thread
zhiyuan1i marked this conversation as resolved.
if dg is not None:
dg = dg.sum(0)
return dq, dk, dw, dg
Loading
Loading