[Costa-SM]

Summary

Hi folks. Today I offer you an optimization to the GDN Layer.

This PR adds a guarded fast path in GatedDeltaNet.forward that runs the q/k/v short convolutions as a single causal_conv1d call. This drops from three conv kernel launches from the layer's convolution prologue to a single one + a cat/split without changing the math.

We concatenate the q/k/v projection outputs along the channel axis, concatenate the matching depthwise conv weights (and bias, if any), and run one causal_conv1d (same underlying kernel). Finally, we then split the result back into q/k/v.

Why

A short depthwise convolution is launch-bound at the default GatedDeltaNet layer shape: the three q/k/v conv calls each do little work, but pay full launch overhead. Since all three are depthwise (groups=channels) convolutions with the same backend, activation, and kernel size, there's no cross-channel mixing, and so, concatenating along the channel axis and running one wider depthwise convolution is mathematically identical to running the three separately while paying the launch cost only once.

This is a layer-prologue change and is independent of the GDN kernels, so it stacks cleanly with kernel optimizations.

Benchmark / Performance

This is a layer-prologue change, so it is most visible at the layer level. Full benchmark_training_throughput dilutes it to noise (~1.003x end-to-end), and the op runner benchmarks.ops.run --op chunk_gated_delta_rule does not see it since the GDN kernels are unchanged. The removal of two causal_conv1d launches shows up in a per-kernel breakdown (benchmark_training_throughput --profile), but does not provide speedup measurements. See PR#952 for a discussion around this. Therefore, we use the script below to get the speedup values:

import torch, triton
from unittest import mock
from fla.layers.gated_deltanet import GatedDeltaNet

# (B, T, hidden_size, head_dim, dtype): first row is the representative shape,
# the rest match the regression test.
SHAPES = [
    (2, 2048, 2048, 256, torch.bfloat16),
    (1, 128, 512, 64, torch.bfloat16),
    (2, 100, 512, 64, torch.bfloat16),
    (2, 256, 512, 64, torch.float16),
]

def bench(fn):
    for _ in range(8):
        fn()
    torch.cuda.synchronize()
    return triton.testing.do_bench(fn, warmup=50, rep=200, quantiles=[0.5, 0.2, 0.8])[0]

for B, T, H, hd, dt in SHAPES:
    torch.manual_seed(42)
    layer = GatedDeltaNet(hidden_size=H, head_dim=hd, num_heads=6, expand_v=2, mode='chunk').cuda().to(dt).train()
    x = torch.randn(B, T, H, device='cuda', dtype=dt)
    xg = x.detach().clone().requires_grad_(True)

    def fwd():
        return layer(x)[0]

    def fwdbwd():
        xg.grad = None
        layer(xg)[0].sum().backward()

    fwd_fused, fb_fused = bench(fwd), bench(fwdbwd)
    # Force the separate path once, around the whole measurement, so the per-call
    # patch/unpatch overhead never enters the timed region.
    with mock.patch.object(GatedDeltaNet, '_use_fused_qkv_conv', return_value=False):
        fwd_sep, fb_sep = bench(fwd), bench(fwdbwd)
    print(f"B{B}-T{T}-H{H}-hd{hd}-{dt}: fwd {fwd_sep/fwd_fused:.2f}x  fwd+bwd {fb_sep/fb_fused:.2f}x")

B200 (CC 10.0, CUDA 12.8, PyTorch 2.11)

Shape	Fwd sep (ms)	Fwd fused (ms)	Fwd	Fwd+bwd sep (ms)	Fwd+bwd fused (ms)	Fwd+bwd
B2-T2048-H2048-hd256 bf16 (representative)	1.739	1.540	~1.13x	4.779	4.084	~1.17x
B1-T128-H512-hd64 bf16	1.732	1.493	~1.16x	4.764	4.066	~1.17x
B2-T100-H512-hd64 bf16	1.707	1.486	~1.15x	4.786	4.088	~1.17x
B2-T256-H512-hd64 fp16	1.727	1.488	~1.16x	4.791	4.095	~1.17x

Test plan

Added a regression test at test_gated_deltanet.py that compares the fused path with the original one over a few shapes/dtypes. The forward output and the input gradient match exactly and every parameter gradient matches within 5e-3. Only the q/k/v conv weights differ at all (by ≲1e-3) due to accumulation noise from regrouping the concatenated conv weight's backward.

Breaking changes

None. Public API and layer outputs are unchanged.

Risk / scope

The fast path is intentionally narrow and is guarded to activate only when:

• use_short_conv is enabled
• there is no cache state (last_state is None, use_cache is falsy)
• there is no variable-length path (cu_seqlens is None)
• the q/k/v convolutions share the same backend, activation, and kernel size

If any are not true, we don't apply the optimization. No new Triton kernels are introduced and we use the existing causal_conv1d for the fused path.

[gemini-code-assist]

Code Review

This pull request introduces a fused QKV convolution optimization in GatedDeltaNet when there is no cache or variable-length state, and the Q, K, and V convolutions share the same configuration. This is achieved by concatenating the projections and weights to execute a single causal_conv1d instead of three separate convolutions. A comprehensive unit test is also added to verify correctness and ensure the fused path is correctly triggered. Feedback suggests replacing rearrange with native PyTorch .squeeze(1) to avoid einops overhead in the forward pass.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[zhiyuan1i]

Thank you for your contribution.
I’ll first merge this MR and apply similar optimizations to other similar models.
By the way, due to the sharp increase in maintenance costs, and given the limited potential for performance improvements, we don’t want to introduce too many specialized“No Grad” features. I think this is something that inference engines like VLLM should handle.