Skip to content

[KDA] fused bwd kernels inter and prepare wy #688

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
yzhangcs merged 16 commits into from
Dec 20, 2025
+141 −82
Merged

[KDA] fused bwd kernels inter and prepare wy #688

Show file tree
Hide file tree
Changes from 14 commits
Commits
Show all changes
16 commits
Select commit Hold shift + click to select a range
d1c5717
combined conv and chunk to save 3 act tensor
yzhangcs Dec 17, 2025
e8ed41f
fused kda inter and prepare wy kernel
yzhangcs Dec 18, 2025
fb5892b
reserv combined chunk
yzhangcs Dec 18, 2025
d798ea1
Merge branch 'main' into feature/fuse_bwd_kernels
yzhangcs Dec 18, 2025
bb03318
Fix lint
yzhangcs Dec 18, 2025
ecff380
Merge remote-tracking branch 'origin/main' into feature/fuse_bwd_kernels
yzhangcs Dec 18, 2025
d2f6591
Minor fix
yzhangcs Dec 19, 2025
5d6d0cc
Minor fix
yzhangcs Dec 19, 2025
afdcb39
Recover wy fast
yzhangcs Dec 19, 2025
c00bec3
Remove dup computation
yzhangcs Dec 19, 2025
ba0812d
Merge branch 'main' into feature/fuse_bwd_kernels
yzhangcs Dec 19, 2025
7b19585
work around for old triton
Dec 19, 2025
68de662
Adjust NUM_WARPS based on GPU architecture
yzhangcs Dec 20, 2025
845d665
Add debug barrier in chunk_bwd.py
yzhangcs Dec 20, 2025
4c4fe2a
Merge branch 'main' into feature/fuse_bwd_kernels
yzhangcs Dec 20, 2025
a26d3ba
Rename chunk_bwd_kernel_dAv to chunk_kda_bwd_kernel_dAv
yzhangcs Dec 20, 2025
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
29 changes: 10 additions & 19 deletions fla/ops/kda/chunk.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,10 +5,10 @@
from fla.modules.l2norm import l2norm_bwd, l2norm_fwd
from fla.ops.common.chunk_delta_h import chunk_gated_delta_rule_bwd_dhu, chunk_gated_delta_rule_fwd_h
from fla.ops.gla.chunk import chunk_gla_fwd_o_gk
from fla.ops.kda.chunk_bwd import chunk_kda_bwd_dAv, chunk_kda_bwd_dqkwg
from fla.ops.kda.chunk_bwd import chunk_kda_bwd_dAv, chunk_kda_bwd_wy_dqkg_fused
from fla.ops.kda.chunk_intra import chunk_kda_bwd_intra, chunk_kda_fwd_intra
from fla.ops.kda.gate import kda_gate_bwd, kda_gate_fwd
from fla.ops.kda.wy_fast import prepare_wy_repr_bwd, recompute_w_u_fwd
from fla.ops.kda.wy_fast import recompute_w_u_fwd
from fla.ops.utils import chunk_local_cumsum
from fla.ops.utils.constant import RCP_LN2
from fla.utils import autocast_custom_bwd, autocast_custom_fwd, input_guard
Expand Down Expand Up @@ -81,6 +81,8 @@ def chunk_kda_bwd(
chunk_indices: torch.LongTensor | None = None,
chunk_size: int = 64,
):
# w = Akk @ (k * beta)
# u = Akk @ (v * beta)
w, u, qg, kg = recompute_w_u_fwd(
q=q,
k=k,
Expand Down Expand Up @@ -130,34 +132,23 @@ def chunk_kda_bwd(
chunk_indices=chunk_indices,
use_exp2=True,
)
dq, dk, dw, dg = chunk_kda_bwd_dqkwg(
dq, dk, dv, db, dg, dAkk = chunk_kda_bwd_wy_dqkg_fused(
q=q,
k=k,
v=v_new,
w=w,
v=v,
v_new=v_new,
g=g,
beta=beta,
A=Akk,
h=h,
dv=dv,
do=do,
dh=dh,
dv=dv,
scale=scale,
cu_seqlens=cu_seqlens,
chunk_size=chunk_size,
chunk_indices=chunk_indices,
)
dk, dv, db, dg, dAkk = prepare_wy_repr_bwd(
k=k,
v=v,
beta=beta,
gk=g,
A=Akk,
dk=dk,
dw=dw,
du=dv,
dg=dg,
cu_seqlens=cu_seqlens,
chunk_indices=chunk_indices,
)
dq, dk, db, dg = chunk_kda_bwd_intra(
q=q,
k=k,
Expand Down
191 changes: 130 additions & 61 deletions fla/ops/kda/chunk_bwd.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -99,26 +99,31 @@ def chunk_bwd_kernel_dAv(
triton.Config({'BK': BK, 'BV': BV}, num_warps=num_warps, num_stages=num_stages)
for BK in BK_LIST
for BV in BV_LIST
for num_warps in [2, 4, 8]
for num_warps in NUM_WARPS
for num_stages in [2, 3, 4]
],
key=['BT'],
**autotune_cache_kwargs,
)
@triton.jit(do_not_specialize=['T'])
def chunk_kda_bwd_kernel_dqkwg(
def chunk_kda_bwd_kernel_wy_dqkg_fused(
q,
k,
v,
v_new,
g,
beta,
A,
h,
do,
dh,
dq,
dk,
dv,
dw,
dv2,
dg,
db,
dA,
cu_seqlens,
chunk_indices,
scale,
Expand All @@ -131,8 +136,9 @@ def chunk_kda_bwd_kernel_dqkwg(
BV: tl.constexpr,
IS_VARLEN: tl.constexpr,
):
i_k, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2)
i_t, i_bh = tl.program_id(0), tl.program_id(1)
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)
Expand All @@ -143,77 +149,129 @@ def chunk_kda_bwd_kernel_dqkwg(
NT = tl.cdiv(T, BT)
i_tg = i_b * NT + i_t
bos, eos = i_b * T, i_b * T + T
o_k = i_k * BK + tl.arange(0, BK)

o_t = i_t * BT + tl.arange(0, BT)
m_k = o_k < K
m_t = o_t < T
m_last = (o_t == min(T, i_t * BT + BT) - 1)

q += (bos * H + i_h) * K
k += (bos * H + i_h) * K
v += (bos * H + i_h) * V
v_new += (bos * H + i_h) * V
g += (bos * H + i_h) * K
beta += bos * H + i_h
A += (bos * H + i_h) * BT
h += (i_tg * H + i_h) * K*V
do += (bos * H + i_h) * V
dh += (i_tg * H + i_h) * K*V
dq += (bos * H + i_h) * K
dk += (bos * H + i_h) * K
dw += (bos * H + i_h) * K
dv += (bos * H + i_h) * V
dv2 += (bos * H + i_h) * V
dg += (bos * H + i_h) * K
db += bos * H + i_h
dA += (bos * H + i_h) * BT

p_g = tl.make_block_ptr(g, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0))
b_g = tl.load(p_g, boundary_check=(0, 1))
p_gn = g + (min(T, i_t * BT + BT) - 1) * H*K + o_k
b_gn = tl.load(p_gn, mask=m_k, other=0)
b_dq = tl.zeros([BT, BK], dtype=tl.float32)
b_dk = tl.zeros([BT, BK], dtype=tl.float32)
b_dw = tl.zeros([BT, BK], dtype=tl.float32)
b_dgk = tl.zeros([BK], dtype=tl.float32)
p_beta = tl.make_block_ptr(beta, (T,), (H,), (i_t * BT,), (BT,), (0,))
b_beta = tl.load(p_beta, boundary_check=(0,))

for i_v in range(tl.cdiv(V, BV)):
p_v = tl.make_block_ptr(v, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0))
p_do = tl.make_block_ptr(do, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0))
p_h = tl.make_block_ptr(h, (V, K), (1, V), (i_v * BV, i_k * BK), (BV, BK), (0, 1))
p_dh = tl.make_block_ptr(dh, (V, K), (1, V), (i_v * BV, i_k * BK), (BV, BK), (0, 1))
# [BT, BV]
b_v = tl.load(p_v, boundary_check=(0, 1))
b_do = tl.load(p_do, boundary_check=(0, 1))
# [BV, BK]
b_h = tl.load(p_h, boundary_check=(0, 1))
b_dh = tl.load(p_dh, boundary_check=(0, 1))
p_A = tl.make_block_ptr(A, (BT, T), (1, H * BT), (0, i_t * BT), (BT, BT), (0, 1))

@gemini-code-assist gemini-code-assist Bot Dec 18, 2025

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.

critical

The tl.make_block_ptr for matrix A is configured incorrectly. The shape and strides seem to be for a transposed view, and the start offset calculation is incorrect. This will lead to out-of-bounds memory access and incorrect results. The shape of A for a given head is (T, BT), so the pointer should be configured accordingly.

Suggested change
p_A = tl.make_block_ptr(A, (BT, T), (1, H * BT), (0, i_t * BT), (BT, BT), (0, 1))
p_A = tl.make_block_ptr(A, (T, BT), (H * BT, 1), (i_t * BT, 0), (BT, BT), (1, 0))

b_A = tl.load(p_A, boundary_check=(0, 1))
Comment on lines +178 to +179

@coderabbitai coderabbitai Bot Dec 19, 2025

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

🧩 Analysis chain

🏁 Script executed:

#!/bin/bash
# Search for A allocation and block pointer usage in chunk_kda_bwd.py and related files
echo "=== Looking for A allocation in chunk_kda_bwd.py and related forward kernels ==="

# Find the wrapper function for chunk_kda_bwd_wy_dqkg_fused
rg -n "def chunk_kda_bwd_wy_dqkg_fused" fla/ops/kda/ -A 50 | head -100

echo ""
echo "=== Looking for A creation in forward or other files ==="

# Search for A.new_empty or similar allocations related to KDA
fd . fla/ops/kda/ --type f -name "*.py" | head -20

echo ""
echo "=== Check what A is in the backward kernel context ==="
rg -n "p_A.*make_block_ptr" fla/ops/kda/ -B 5 -A 2

Repository: fla-org/flash-linear-attention

Length of output: 10739

Correct block pointer configuration to match forward kernel memory layout.

The forward kernel (chunk_intra.py) allocates and accesses A with shape (T, BT) and strides (H*BT, 1). The backward kernel must use the same configuration to access the same memory layout correctly.

Change line 178 from:

p_A = tl.make_block_ptr(A, (BT, T), (1, H * BT), (0, i_t * BT), (BT, BT), (0, 1))

To:

p_A = tl.make_block_ptr(A, (T, BT), (H * BT, 1), (i_t * BT, 0), (BT, BT), (1, 0))

The transposed configuration (BT, T) with strides (1, H*BT) is incompatible with how A is allocated in the forward pass and will cause incorrect memory access patterns.

🤖 Prompt for AI Agents 
In fla/ops/kda/chunk_bwd.py around lines 178 to 179, the block pointer for A is
configured with shape (BT, T) and strides (1, H * BT) which is transposed
relative to the forward kernel; replace that configuration so it matches the
forward kernel's allocation/access (shape (T, BT) and strides (H * BT, 1)),
i.e., change the make_block_ptr call to use (T, BT), (H * BT, 1), and adjust the
offset/transpose parameters accordingly so the pointer addresses the same memory
layout as the forward pass.


# [BK]
b_dgk += tl.sum(b_h * b_dh, axis=0)
# [BT, BK]
b_dq += tl.dot(b_do, b_h.to(b_do.dtype))
b_dk += tl.dot(b_v, b_dh.to(b_v.dtype))
b_dA = tl.zeros([BT, BT], dtype=tl.float32)
b_db = tl.zeros([BT], dtype=tl.float32)

p_dv = tl.make_block_ptr(dv, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0))
b_dv = tl.load(p_dv, boundary_check=(0, 1))
b_dw += tl.dot(b_dv.to(b_v.dtype), b_h.to(b_v.dtype))
for i_k in range(tl.cdiv(K, BK)):
o_k = i_k * BK + tl.arange(0, BK)
m_k = o_k < K

p_k = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0))
p_g = tl.make_block_ptr(g, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0))
b_k = tl.load(p_k, boundary_check=(0, 1))
b_g = tl.load(p_g, boundary_check=(0, 1))

p_gn = g + (min(T, i_t * BT + BT) - 1) * H*K + o_k
b_gn = tl.load(p_gn, mask=m_k, other=0)

b_dq = tl.zeros([BT, BK], dtype=tl.float32)
b_dk = tl.zeros([BT, BK], dtype=tl.float32)
b_dw = tl.zeros([BT, BK], dtype=tl.float32)
b_dgk = tl.zeros([BK], dtype=tl.float32)

p_dw = tl.make_block_ptr(dw, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0))
tl.store(p_dw, -b_dw.to(p_dw.dtype.element_ty), boundary_check=(0, 1))
for i_v in range(tl.cdiv(V, BV)):
p_v_new = tl.make_block_ptr(v_new, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0))
p_do = tl.make_block_ptr(do, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0))
p_h = tl.make_block_ptr(h, (V, K), (1, V), (i_v * BV, i_k * BK), (BV, BK), (0, 1))
p_dh = tl.make_block_ptr(dh, (V, K), (1, V), (i_v * BV, i_k * BK), (BV, BK), (0, 1))
p_dv = tl.make_block_ptr(dv, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0))
# [BT, BV]
b_v_new = tl.load(p_v_new, boundary_check=(0, 1))
b_do = tl.load(p_do, boundary_check=(0, 1))
# [BV, BK]
b_h = tl.load(p_h, boundary_check=(0, 1))
b_dh = tl.load(p_dh, boundary_check=(0, 1))
# [BT, BV]
b_dv = tl.load(p_dv, boundary_check=(0, 1))

b_dgk *= exp2(b_gn)
b_dq *= scale
b_dq = b_dq * exp2(b_g)
b_dk = b_dk * tl.where(m_t[:, None], exp2(b_gn[None, :] - b_g), 0)
b_dgk += tl.sum(b_h * b_dh, axis=0)
b_dq += tl.dot(b_do, b_h.to(b_do.dtype))
b_dk += tl.dot(b_v_new, b_dh.to(b_v_new.dtype))
b_dw += tl.dot(b_dv.to(b_v_new.dtype), b_h.to(b_v_new.dtype))

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_dq = tl.make_block_ptr(dq, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0))
p_dk = tl.make_block_ptr(dk, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0))
p_dg = tl.make_block_ptr(dg, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0))
b_q = tl.load(p_q, boundary_check=(0, 1))
b_k = tl.load(p_k, boundary_check=(0, 1))
b_dgk += tl.sum(b_dk * b_k, axis=0)
b_dg = b_q * b_dq - b_k * b_dk + m_last[:, None] * b_dgk
tl.debug_barrier()

@coderabbitai coderabbitai Bot Dec 20, 2025

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 | 🟡 Minor

Consider removing or conditionalizing the debug barrier.

The tl.debug_barrier() is useful for debugging synchronization issues but may impact performance in production. Consider removing it once the kernel is validated, or guard it with a compile-time flag.

🤖 Prompt for AI Agents 
In fla/ops/kda/chunk_bwd.py around line 221, tl.debug_barrier() is left in the
hot kernel path; remove it for production or guard it behind a debug flag so it
only runs during development. Replace the bare call with either deletion, or a
conditional such as checking a module-level DEBUG/ENABLE_DEBUG_BARRIER constant
or an environment variable (e.g. if DEBUG: tl.debug_barrier()), or use a
compile-time flag so the barrier is compiled out in release builds.

if i_k == 0:
p_v = tl.make_block_ptr(v, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0))
p_dv2 = tl.make_block_ptr(dv2, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0))

tl.store(p_dq, b_dq.to(p_dq.dtype.element_ty), boundary_check=(0, 1))
tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1))
tl.store(p_dg, b_dg.to(p_dg.dtype.element_ty), boundary_check=(0, 1))
b_v = tl.load(p_v, boundary_check=(0, 1))
b_vb = (b_v * b_beta[:, None]).to(b_v.dtype)

b_dA += tl.dot(b_dv, tl.trans(b_vb))

b_dvb = tl.dot(b_A, b_dv)

@coderabbitai coderabbitai Bot Dec 19, 2025

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

🧩 Analysis chain

🏁 Script executed:

# Find the forward kernel file
find . -name "*chunk_fwd*" -o -name "*kda*fwd*" | head -20

Repository: fla-org/flash-linear-attention

Length of output: 56

🏁 Script executed:

# List files in the kda directory to find forward pass
ls -la fla/ops/kda/

Repository: fla-org/flash-linear-attention

Length of output: 775

🏁 Script executed:

# Look at the forward chunk.py file to understand the forward pass
cat -n fla/ops/kda/chunk.py | head -100

Repository: fla-org/flash-linear-attention

Length of output: 3402

🏁 Script executed:

# Search for where b_dvb is computed and look at surrounding context
grep -n "b_dvb\|b_dkbg\|b_dA" fla/ops/kda/chunk_bwd.py | head -20

Repository: fla-org/flash-linear-attention

Length of output: 1046

🏁 Script executed:

# Look at the context around line 230 to understand the forward pass computation
sed -n '200,275p' fla/ops/kda/chunk_bwd.py

Repository: fla-org/flash-linear-attention

Length of output: 3681

🏁 Script executed:

# Look at the wy_fast.py file to see how the WY representation forward pass is computed
cat -n fla/ops/kda/wy_fast.py | grep -A 20 -B 5 "v_new\|vb"

Repository: fla-org/flash-linear-attention

Length of output: 3501

🏁 Script executed:

# Now let's verify the matrix inverse gradient computation at lines 268-269
sed -n '260,275p' fla/ops/kda/chunk_bwd.py

Repository: fla-org/flash-linear-attention

Length of output: 1022

🏁 Script executed:

# Check if there are any comments or documentation about the WY representation
grep -n "WY\|inverse\|transpose" fla/ops/kda/wy_fast.py | head -20

Repository: fla-org/flash-linear-attention

Length of output: 56

🏁 Script executed:

# Let me look at the complete recompute_w_u_fwd function to understand the forward pass better
sed -n '212,250p' fla/ops/kda/wy_fast.py

Repository: fla-org/flash-linear-attention

Length of output: 1134

🏁 Script executed:

# Let's examine the mathematical operation more carefully - check what the actual forward computation is
sed -n '40,75p' fla/ops/kda/wy_fast.py

Repository: fla-org/flash-linear-attention

Length of output: 1567

Missing transpose in backward gradient computation through WY representation.

Line 230 computes b_dvb = tl.dot(b_A, b_dv), but the forward pass applies b_u = tl.dot(b_A, b_vb) where b_vb = b_v * b_beta. For the gradient to flow correctly through a matrix multiplication, the backward pass must apply the transpose: b_dvb = tl.dot(tl.trans(b_A), b_dv).

The same issue affects line 247: b_dkbg = tl.dot(b_A, b_dw) should be b_dkbg = tl.dot(tl.trans(b_A), b_dw).

Additionally, lines 268-269 implement the gradient of matrix inverse incorrectly. The derivative of f(X^{-1}) requires dX = -X^{-T} @ dY @ X^{-T}, but the current code only computes A @ dA @ A without proper transposes.

🤖 Prompt for AI Agents 
In fla/ops/kda/chunk_bwd.py around lines 230, 247 and 268-269, the backward pass
applies matrix multiplications and inverse gradients without required
transposes: change line 230 b_dvb = tl.dot(b_A, b_dv) to b_dvb =
tl.dot(tl.trans(b_A), b_dv); change line 247 b_dkbg = tl.dot(b_A, b_dw) to
b_dkbg = tl.dot(tl.trans(b_A), b_dw); and replace the incorrect inverse-gradient
computation at lines 268-269 (which currently uses A @ dA @ A) with the correct
form using transposed inverses: dX = -A_T @ dY @ A_T where A_T is tl.trans(A)
(i.e., use A_T = tl.trans(A) and compute -tl.dot(tl.dot(A_T, dY), A_T)). Ensure
variable names match existing locals and use tl.trans consistently.

b_dv2 = b_dvb * b_beta[:, None]
b_db += tl.sum(b_dvb * b_v, 1)

tl.store(p_dv2, b_dv2.to(p_dv2.dtype.element_ty), boundary_check=(0, 1))

b_gk_exp = exp2(b_g)
b_dgk *= exp2(b_gn)
b_dq = b_dq * b_gk_exp * scale
b_dk = b_dk * tl.where(m_t[:, None], exp2(b_gn[None, :] - b_g), 0)

b_kbg = (b_k * b_beta[:, None] * b_gk_exp).to(b_A.dtype)

b_dw = -b_dw.to(b_A.dtype)
b_dA += tl.dot(b_dw, tl.trans(b_kbg))

b_dkbg = tl.dot(b_A, b_dw)
b_dkbgg = b_dkbg * b_gk_exp
b_db += tl.sum(b_dkbgg * b_k, 1)

p_q = tl.make_block_ptr(q, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0))
b_q = tl.load(p_q, boundary_check=(0, 1))
b_kdk = b_k * b_dk
b_dgk += tl.sum(b_kdk, axis=0)
b_dg = b_q * b_dq - b_kdk + m_last[:, None] * b_dgk + b_kbg * b_dkbg
b_dk = b_dk + b_dkbgg * b_beta[:, None]

p_dq = tl.make_block_ptr(dq, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0))
p_dk = tl.make_block_ptr(dk, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0))
p_dg = tl.make_block_ptr(dg, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0))
tl.store(p_dq, b_dq.to(p_dq.dtype.element_ty), boundary_check=(0, 1))
tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1))
tl.store(p_dg, b_dg.to(p_dg.dtype.element_ty), boundary_check=(0, 1))
Comment on lines +184 to +262

@coderabbitai coderabbitai Bot Dec 20, 2025

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.

🛠️ Refactor suggestion | 🟠 Major

Good fusion design, but correctness depends on fixing critical bugs.

The interleaved K-block loop design successfully fuses WY backward computation with gradient accumulation, keeping intermediate dw in registers and eliminating redundant global memory traffic. This achieves the PR objective of reducing bandwidth and kernel launches.

However, the kernel contains multiple critical mathematical errors (missing transposes at lines 230, 247 and incorrect inverse gradient at lines 264-268) that must be fixed before the fused kernel can be used.


m_A = (o_t[:, None] > o_t[None, :]) & (m_t[:, None] & m_t)
b_dA = tl.where(m_A, b_dA, 0)
b_dA = tl.dot(b_dA.to(b_A.dtype), b_A)
b_dA = tl.dot(b_A, b_dA.to(b_A.dtype))
Comment on lines +266 to +267

@gemini-code-assist gemini-code-assist Bot Dec 18, 2025

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.

critical

The gradient calculation for dAkk from dAkk_inv (represented by b_dA) is incorrect. The derivative of a matrix inverse A = X^{-1} is dX = -A^T dA A^T. The current implementation computes A @ (dA @ A), which is mathematically incorrect. The multiplication should be with the transpose of b_A.

Suggested change
b_dA = tl.dot(b_dA.to(b_A.dtype), b_A)
b_dA = tl.dot(b_A, b_dA.to(b_A.dtype))
b_dA_t = tl.dot(tl.trans(b_A), b_dA.to(b_A.dtype))
b_dA = tl.dot(b_dA_t, tl.trans(b_A))

b_dA = tl.where(m_A, -b_dA, 0)
Comment on lines +264 to +268

@coderabbitai coderabbitai Bot Dec 19, 2025

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

Critical: Incorrect gradient calculation for matrix inverse.

The gradient of a matrix inverse is computed incorrectly. For A = X^{-1}, the derivative is:

dX = -A^T @ dA @ A^T

However, the current implementation computes:

b_dA = tl.dot(b_dA.to(b_A.dtype), b_A)      # dA @ A
b_dA = tl.dot(b_A, b_dA.to(b_A.dtype))      # A @ (dA @ A)
b_dA = tl.where(m_A, -b_dA, 0)               # negation

This computes -(A @ dA @ A) instead of -(A^T @ dA @ A^T). The transposes are missing.

🔎 Proposed fix 
-    b_dA = tl.dot(b_dA.to(b_A.dtype), b_A)
-    b_dA = tl.dot(b_A, b_dA.to(b_A.dtype))
+    b_dA_t = tl.dot(tl.trans(b_A), b_dA.to(b_A.dtype))
+    b_dA = tl.dot(b_dA_t, tl.trans(b_A))
     b_dA = tl.where(m_A, -b_dA, 0)
🤖 Prompt for AI Agents 
In fla/ops/kda/chunk_bwd.py around lines 266 to 270, the gradient for the matrix
inverse is computed as -(A @ dA @ A) but must be -(A^T @ dA @ A^T); replace the
three dot calls with operations that compute b_dA = - (b_A.T @ b_dA @ b_A.T)
(respecting dtype casts as needed), then apply the existing mask m_A (i.e., b_dA
= tl.where(m_A, b_dA, 0)). Ensure you cast operands to matching dtypes before
each dot and perform the transposes on b_A (and on the intermediate if required)
so the final result matches -(A^T @ dA @ A^T).


p_dA = tl.make_block_ptr(dA, (T, BT), (H * BT, 1), (i_t * BT, 0), (BT, BT), (1, 0))
p_db = tl.make_block_ptr(db, (T,), (H,), (i_t * BT,), (BT,), (0,))
tl.store(p_dA, b_dA.to(p_dA.dtype.element_ty), boundary_check=(0, 1))
tl.store(p_db, b_db.to(p_db.dtype.element_ty), boundary_check=(0,))


def chunk_kda_bwd_dAv(
Expand Down Expand Up @@ -267,13 +325,15 @@ def chunk_kda_bwd_dAv(
return dA, dv


def chunk_kda_bwd_dqkwg(
def chunk_kda_bwd_wy_dqkg_fused(
q: torch.Tensor,
k: torch.Tensor,
w: torch.Tensor,
v: torch.Tensor,
h: torch.Tensor,
v_new: torch.Tensor,
g: torch.Tensor,
beta: torch.Tensor,
A: torch.Tensor,
h: torch.Tensor,
do: torch.Tensor,
dh: torch.Tensor,
dv: torch.Tensor,
Expand All @@ -291,22 +351,30 @@ def chunk_kda_bwd_dqkwg(

dq = torch.empty_like(q, dtype=torch.float)
dk = torch.empty_like(k, dtype=torch.float)
dw = torch.empty_like(w)
dg = torch.empty_like(g)
def grid(meta): return (triton.cdiv(K, meta['BK']), NT, B * H)
chunk_kda_bwd_kernel_dqkwg[grid](
dv2 = torch.empty_like(v)
dg = torch.empty_like(g, dtype=torch.float)
db = torch.empty_like(beta, dtype=torch.float)
dA = torch.empty_like(A, dtype=torch.float)

grid = (NT, B * H)
chunk_kda_bwd_kernel_wy_dqkg_fused[grid](
q=q,
k=k,
v=v,
v_new=v_new,
g=g,
beta=beta,
A=A,
h=h,
do=do,
dh=dh,
dq=dq,
dk=dk,
dv=dv,
dw=dw,
dv2=dv2,
dg=dg,
db=db,
dA=dA,
cu_seqlens=cu_seqlens,
chunk_indices=chunk_indices,
scale=scale,
Expand All @@ -316,4 +384,5 @@ def grid(meta): return (triton.cdiv(K, meta['BK']), NT, B * H)
V=V,
BT=BT,
)
return dq, dk, dw, dg
dv = dv2
return dq, dk, dv, db, dg, dA