Skip to content

[GDN] Add exp2 support across chunk kernels for improved performance #791

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 5 commits into from
Mar 25, 2026
Merged

[GDN] Add exp2 support across chunk kernels for improved performance #791

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
8b287d5
[GDN] Add exp2 support across chunk kernels for improved performance
yzhangcs Mar 24, 2026
0a745db
Remove extra blank lines between copyright header and imports
yzhangcs Mar 24, 2026
b16630c
Update comba
yzhangcs Mar 24, 2026
d03a419
Fix bugs
yzhangcs Mar 24, 2026
28abbab
Merge branch 'main' into exp2
yzhangcs Mar 24, 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
73 changes: 57 additions & 16 deletions fla/ops/common/chunk_o.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,7 @@
import triton.language as tl

from fla.ops.utils import prepare_chunk_indices
from fla.ops.utils.op import exp
from fla.ops.utils.op import exp, exp2
from fla.utils import IS_NVIDIA_HOPPER, autotune_cache_kwargs, check_shared_mem

BKV_LIST = [64, 128] if check_shared_mem() else ([32, 64] if check_shared_mem('ada') else [32])
Expand Down Expand Up @@ -47,6 +47,7 @@ def chunk_fwd_kernel_o(
BV: tl.constexpr,
USE_G: tl.constexpr,
USE_G_GAMMA: tl.constexpr,
USE_EXP2: tl.constexpr,
TRANSPOSE_STATE: tl.constexpr,
IS_VARLEN: tl.constexpr,
):
Expand Down Expand Up @@ -99,14 +100,22 @@ def chunk_fwd_kernel_o(
g += bos * H + i_h
p_g = tl.make_block_ptr(g, (T,), (H,), (i_t * BT,), (BT,), (0,))
b_g = tl.load(p_g, boundary_check=(0,))
b_o = b_o * exp(b_g)[:, None]
b_A = b_A * exp(b_g[:, None] - b_g[None, :])
if USE_EXP2:
b_o = b_o * exp2(b_g)[:, None]
b_A = b_A * exp2(b_g[:, None] - b_g[None, :])
else:
b_o = b_o * exp(b_g)[:, None]
b_A = b_A * exp(b_g[:, None] - b_g[None, :])

if USE_G_GAMMA:
b_gamma = tl.load(g_gamma + i_h)
b_g = b_gamma * (tl.arange(0, BT) + 1)
b_o = b_o * exp(b_g)[:, None]
b_A = b_A * exp(b_g[:, None] - b_g[None, :])
if USE_EXP2:
b_o = b_o * exp2(b_g)[:, None]
b_A = b_A * exp2(b_g[:, None] - b_g[None, :])
else:
b_o = b_o * exp(b_g)[:, None]
b_A = b_A * exp(b_g[:, None] - b_g[None, :])
Comment on lines 110 to +118

@coderabbitai coderabbitai Bot Mar 24, 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 | 🟠 Major

Scale g_gamma before the exp2 branches.

The USE_G path can stay equivalent because fla/ops/gated_delta_rule/chunk.py now pre-scales cumulative g by RCP_LN2 before it reaches this file. g_gamma is generated locally here, though, so the new exp2(...) branches are still consuming raw natural-log decays. With use_exp2=True, every USE_G_GAMMA path here changes the operator in both forward and backward.

💡 Suggested fix 
+from fla.ops.utils.constant import RCP_LN2
 from fla.ops.utils import prepare_chunk_indices
 from fla.ops.utils.op import exp, exp2
 if USE_G_GAMMA:
     b_gamma = tl.load(g_gamma + i_h)
     b_g = b_gamma * (tl.arange(0, BT) + 1)
+    if USE_EXP2:
+        b_g = b_g * RCP_LN2
 if USE_G_GAMMA:
     b_gamma = tl.load(g_gamma + i_h)
     b_g = b_gamma * (tl.arange(0, BT) + 1)
     b_g_last = b_gamma * min(BT, T - i_t * BT)
+    if USE_EXP2:
+        b_g = b_g * RCP_LN2
+        b_g_last = b_g_last * RCP_LN2

Apply the same normalization in each kernel here that materializes b_g / b_g_last from g_gamma.

Also applies to: 311-319, 421-427, 512-516

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

In `@fla/ops/common/chunk_o.py` around lines 110 - 118, The code uses
locally-built g_gamma (symbol g_gamma) to compute b_gamma and b_g, but when
USE_EXP2 is true b_g must be pre-scaled by RCP_LN2 (the same normalization used
upstream) before calling exp2; update every kernel that materializes b_g (and
b_g_last) from b_gamma—i.e., where b_gamma = tl.load(g_gamma + i_h) and b_g =
b_gamma * (tl.arange(...) + 1)—to multiply b_gamma (or b_g) by RCP_LN2 (the
reciprocal of ln(2)) prior to using exp2, and keep the USE_EXP2/else branches
otherwise unchanged so exp2 sees base-2-scaled exponents consistently (apply the
same change in all occurrences mentioned: around lines 110-118, 311-319,
421-427, 512-516).


o_t = i_t * BT + tl.arange(0, BT)
m_t = o_t < T
Expand Down Expand Up @@ -166,6 +175,7 @@ def chunk_bwd_kernel_dqkwg(
BV: tl.constexpr,
USE_G: tl.constexpr,
USE_G_GAMMA: tl.constexpr,
USE_EXP2: tl.constexpr,
USE_DW: tl.constexpr,
TRANSPOSE_STATE: tl.constexpr,
IS_VARLEN: tl.constexpr,
Expand Down Expand Up @@ -263,16 +273,25 @@ def chunk_bwd_kernel_dqkwg(
p_g = tl.make_block_ptr(g, (T,), (H,), (i_t * BT,), (BT,), (0,))
b_g = tl.load(p_g, boundary_check=(0,))
b_g_last = tl.load(g + (min(i_t * BT + BT, T) - 1) * H)
b_dg_last *= exp(b_g_last)

b_dq = b_dq * exp(b_g)[:, None] * scale
if USE_EXP2:
b_dg_last *= exp2(b_g_last)
b_dq = b_dq * exp2(b_g)[:, None] * scale
else:
b_dg_last *= exp(b_g_last)
b_dq = b_dq * exp(b_g)[:, None] * scale
b_dg += tl.sum(b_dq * b_q, axis=1)

b_dk = b_dk * tl.where(m_t, exp(-b_g + b_g_last), 0)[:, None]
if USE_EXP2:
b_dk = b_dk * tl.where(m_t, exp2(-b_g + b_g_last), 0)[:, None]
else:
b_dk = b_dk * tl.where(m_t, exp(-b_g + b_g_last), 0)[:, None]
b_dg -= tl.sum(b_k * b_dk, axis=1)
b_dg_last += tl.sum(b_dk * b_k)

b_ds = tl.where(m_A, b_ds * exp(b_g[:, None] - b_g[None, :]), 0) * scale
if USE_EXP2:
b_ds = tl.where(m_A, b_ds * exp2(b_g[:, None] - b_g[None, :]), 0) * scale
else:
b_ds = tl.where(m_A, b_ds * exp(b_g[:, None] - b_g[None, :]), 0) * scale
b_ds2 = b_ds * tl.dot(b_q, tl.trans(b_k))
b_dg += tl.sum(b_ds2, axis=1)
b_dg -= tl.sum(b_ds2, axis=0)
Expand All @@ -290,9 +309,14 @@ def chunk_bwd_kernel_dqkwg(
tl.store(p_dg, b_dg.to(p_dg.dtype.element_ty), boundary_check=(0,))

elif USE_G_GAMMA:
b_dq = b_dq * exp(b_g)[:, None] * scale
b_dk = b_dk * tl.where(m_t, exp(-b_g + b_g_last), 0)[:, None]
b_ds = tl.where(m_A, b_ds * exp(b_g[:, None] - b_g[None, :]), 0) * scale
if USE_EXP2:
b_dq = b_dq * exp2(b_g)[:, None] * scale
b_dk = b_dk * tl.where(m_t, exp2(-b_g + b_g_last), 0)[:, None]
b_ds = tl.where(m_A, b_ds * exp2(b_g[:, None] - b_g[None, :]), 0) * scale
else:
b_dq = b_dq * exp(b_g)[:, None] * scale
b_dk = b_dk * tl.where(m_t, exp(-b_g + b_g_last), 0)[:, None]
b_ds = tl.where(m_A, b_ds * exp(b_g[:, None] - b_g[None, :]), 0) * scale
b_ds = b_ds.to(b_k.dtype)
# [BT, BK]
b_dq += tl.dot(b_ds, b_k)
Expand Down Expand Up @@ -345,6 +369,7 @@ def chunk_bwd_kernel_dv(
BV: tl.constexpr,
USE_G: tl.constexpr,
USE_G_GAMMA: tl.constexpr,
USE_EXP2: tl.constexpr,
IS_VARLEN: tl.constexpr,
):
i_v, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2)
Expand Down Expand Up @@ -394,8 +419,12 @@ def chunk_bwd_kernel_dv(

m_A = (o_t[:, None] <= o_t[None, :]) & (m_t[:, None] & m_t)
if USE_G or USE_G_GAMMA:
b_A = tl.where(m_A, b_A * exp(b_g[None, :] - b_g[:, None]) * scale, 0).to(do.dtype.element_ty)
b_dv *= tl.where(m_t, exp(-b_g + b_g_last), 0)[:, None]
if USE_EXP2:
b_A = tl.where(m_A, b_A * exp2(b_g[None, :] - b_g[:, None]) * scale, 0).to(do.dtype.element_ty)
b_dv *= tl.where(m_t, exp2(-b_g + b_g_last), 0)[:, None]
else:
b_A = tl.where(m_A, b_A * exp(b_g[None, :] - b_g[:, None]) * scale, 0).to(do.dtype.element_ty)
b_dv *= tl.where(m_t, exp(-b_g + b_g_last), 0)[:, None]
else:
b_A = tl.where(m_A, b_A * scale, 0).to(do.dtype.element_ty)
p_do = tl.make_block_ptr(do, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0))
Expand Down Expand Up @@ -441,6 +470,7 @@ def chunk_bwd_kernel_dv_local(
BV: tl.constexpr,
USE_G: tl.constexpr,
USE_G_GAMMA: tl.constexpr,
USE_EXP2: tl.constexpr,
USE_A: tl.constexpr,
IS_VARLEN: tl.constexpr,
):
Expand Down Expand Up @@ -480,7 +510,10 @@ def chunk_bwd_kernel_dv_local(
b_q = tl.load(p_q, boundary_check=(0, 1))
b_A += tl.dot(b_k, b_q) * scale
if USE_G or USE_G_GAMMA:
b_A *= exp(b_g[None, :] - b_g[:, None])
if USE_EXP2:
b_A *= exp2(b_g[None, :] - b_g[:, None])
else:
b_A *= exp(b_g[None, :] - b_g[:, None])

o_t = i_t * BT + tl.arange(0, BT)
m_t = o_t < T
Expand All @@ -506,6 +539,7 @@ def chunk_fwd_o(
cu_seqlens: torch.LongTensor | None = None,
chunk_size: int = 64,
chunk_indices: torch.LongTensor | None = None,
use_exp2: bool = False,
transpose_state_layout: bool = False,
) -> torch.Tensor:
B, T, H, K, V = *q.shape, v.shape[-1]
Expand Down Expand Up @@ -534,6 +568,7 @@ def grid(meta): return (triton.cdiv(V, meta['BV']), NT, B * H)
K=K,
V=V,
BT=BT,
USE_EXP2=use_exp2,
TRANSPOSE_STATE=transpose_state_layout,
)
return o
Expand All @@ -550,6 +585,7 @@ def chunk_bwd_dv(
cu_seqlens: torch.LongTensor | None = None,
chunk_size: int = 64,
chunk_indices: torch.LongTensor | None = None,
use_exp2: bool = False,
) -> torch.Tensor:
B, T, H, K, V = *k.shape, do.shape[-1]
BT = chunk_size
Expand Down Expand Up @@ -589,6 +625,7 @@ def chunk_bwd_dv(
BT=BT,
BK=BK,
BV=BV,
USE_EXP2=use_exp2,
)
return dv

Expand All @@ -604,6 +641,7 @@ def chunk_bwd_dv_local(
cu_seqlens: torch.LongTensor | None = None,
chunk_size: int = 64,
chunk_indices: torch.LongTensor | None = None,
use_exp2: bool = False,
) -> torch.Tensor:
B, T, H, K, V = *k.shape, do.shape[-1]
BT = chunk_size
Expand Down Expand Up @@ -640,6 +678,7 @@ def chunk_bwd_dv_local(
BT=BT,
BK=BK,
BV=BV,
USE_EXP2=use_exp2,
)
return dv

Expand All @@ -659,6 +698,7 @@ def chunk_bwd_dqkwg(
cu_seqlens: torch.LongTensor | None = None,
chunk_size: int = 64,
chunk_indices: torch.LongTensor | None = None,
use_exp2: bool = False,
transpose_state_layout: bool = False,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:

Expand Down Expand Up @@ -708,6 +748,7 @@ def chunk_bwd_dqkwg(
BT=BT,
BK=BK,
BV=BV,
USE_EXP2=use_exp2,
TRANSPOSE_STATE=transpose_state_layout,
)

Expand Down
1 change: 1 addition & 0 deletions fla/ops/gated_delta_product/chunk.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -188,6 +188,7 @@ def backward(
dht=dht,
cu_seqlens=cu_seqlens * ctx.num_householder if cu_seqlens is not None else None,
chunk_indices=chunk_indices_dp,
use_exp2=False,
)
dg = rearrange(dg, 'b (l n) h -> b l n h ', n=ctx.num_householder)[:, :, 0].contiguous().to(g)
else:
Expand Down
22 changes: 21 additions & 1 deletion fla/ops/gated_delta_rule/chunk.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,7 @@
from fla.ops.gated_delta_rule.chunk_fwd import chunk_gated_delta_rule_fwd_intra
from fla.ops.gated_delta_rule.wy_fast import prepare_wy_repr_bwd, recompute_w_u_fwd
from fla.ops.utils import chunk_local_cumsum
from fla.ops.utils.constant import RCP_LN2
from fla.ops.utils.index import prepare_chunk_indices
from fla.utils import autocast_custom_bwd, autocast_custom_fwd, input_guard

Expand All @@ -33,9 +34,16 @@ def chunk_gated_delta_rule_fwd(
cu_seqlens: torch.LongTensor | None = None,
cp_context: FLACPContext | None = None,
chunk_indices: torch.LongTensor | None = None,
use_exp2: bool = True,

@gemini-code-assist gemini-code-assist Bot Mar 24, 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.

high

The use_exp2 flag is introduced here with a default value of True, but it's not configurable from the public-facing chunk_gated_delta_rule function. This means exp2 will always be used for this operation. To make this configurable, you should consider adding use_exp2 to the signature of chunk_gated_delta_rule and ChunkGatedDeltaRuleFunction, and then pass it down through the function calls. This would align its behavior with other functions in this PR like fused_recurrent_gated_delta_rule.

transpose_state_layout: bool = False,
):
g = chunk_local_cumsum(g, chunk_size=64, cu_seqlens=cu_seqlens, chunk_indices=chunk_indices)
g = chunk_local_cumsum(
g,
chunk_size=64,
scale=RCP_LN2 if use_exp2 else None,
cu_seqlens=cu_seqlens,
chunk_indices=chunk_indices,
)
# obtain WY representation. u is actually the new v.
# fused kkt + solve_tril + recompute_w_u
w, u, A = chunk_gated_delta_rule_fwd_intra(
Expand All @@ -45,6 +53,7 @@ def chunk_gated_delta_rule_fwd(
beta=beta,
cu_seqlens=cu_seqlens,
chunk_indices=chunk_indices,
use_exp2=use_exp2,
)

if cp_context is not None:
Expand All @@ -56,6 +65,7 @@ def chunk_gated_delta_rule_fwd(
cu_seqlens=cu_seqlens,
initial_state=initial_state,
context=cp_context,
use_exp2=use_exp2,
transpose_state_layout=transpose_state_layout,
)

Expand All @@ -68,6 +78,7 @@ def chunk_gated_delta_rule_fwd(
output_final_state=output_final_state,
cu_seqlens=cu_seqlens,
chunk_indices=chunk_indices,
use_exp2=use_exp2,
transpose_state_layout=transpose_state_layout,
)

Expand All @@ -83,6 +94,7 @@ def chunk_gated_delta_rule_fwd(
scale=scale,
cu_seqlens=cu_seqlens,
chunk_indices=chunk_indices,
use_exp2=use_exp2,
transpose_state_layout=transpose_state_layout,
)
return g, o, A, final_state, initial_state
Expand All @@ -102,6 +114,7 @@ def chunk_gated_delta_rule_bwd(
cu_seqlens: torch.LongTensor | None = None,
cp_context: FLACPContext | None = None,
chunk_indices: torch.LongTensor | None = None,
use_exp2: bool = True,
transpose_state_layout: bool = False,
):
w, u = recompute_w_u_fwd(
Expand All @@ -112,6 +125,7 @@ def chunk_gated_delta_rule_bwd(
g=g,
cu_seqlens=cu_seqlens,
chunk_indices=chunk_indices,
use_exp2=use_exp2,
)

if cp_context is not None:
Expand All @@ -126,6 +140,7 @@ def chunk_gated_delta_rule_bwd(
output_final_state=False,
cu_seqlens=cu_seqlens,
chunk_indices=chunk_indices,
use_exp2=use_exp2,
transpose_state_layout=transpose_state_layout,
)
dv = chunk_bwd_dv_local(
Expand All @@ -136,6 +151,7 @@ def chunk_gated_delta_rule_bwd(
scale=scale,
cu_seqlens=cu_seqlens,
chunk_indices=chunk_indices,
use_exp2=use_exp2,
)

if cp_context is not None:
Expand All @@ -153,6 +169,7 @@ def chunk_gated_delta_rule_bwd(
dht=dht,
initial_state=initial_state,
context=cp_context,
use_exp2=use_exp2,
transpose_state_layout=transpose_state_layout,
)

Expand All @@ -168,6 +185,7 @@ def chunk_gated_delta_rule_bwd(
scale=scale,
cu_seqlens=cu_seqlens,
chunk_indices=chunk_indices,
use_exp2=use_exp2,
transpose_state_layout=transpose_state_layout,
)
dq, dk, dw, dg = chunk_bwd_dqkwg(
Expand All @@ -183,6 +201,7 @@ def chunk_gated_delta_rule_bwd(
scale=scale,
cu_seqlens=cu_seqlens,
chunk_indices=chunk_indices,
use_exp2=use_exp2,
transpose_state_layout=transpose_state_layout,
)
dk2, dv, db, dg2 = prepare_wy_repr_bwd(
Expand All @@ -195,6 +214,7 @@ def chunk_gated_delta_rule_bwd(
du=dv,
cu_seqlens=cu_seqlens,
chunk_indices=chunk_indices,
use_exp2=use_exp2,
)
dk.add_(dk2)
dg.add_(dg2)
Expand Down
Loading
Loading