kda in torchtitan (training) by liangel-02 · Pull Request #4164 · pytorch/torchtitan

liangel-02 · 2026-08-15T20:46:51Z

this PR adds attention module for KDA (relying on attention gym KDA kernels and user facing APIs)

correctness

tested with c4 dataset:

torchtitan attention gym vs hugging face (natural routing)

Full multimodal last-token logit parity (TorchTitan vs HuggingFace)
KL=7.0921e-02 cos=0.927007 max_diff=1.6885e+00 top1=N top5=20%

^i compared hugigng face and torchtitan layer by layer using identical weights and inputs, first divergence appears at block 2 where expert choice agreement was 94.5% and max router score diff was 0.0065. with nearly tied experts, small numerical differences led to different experts being in the top k which led to model executing different expert weights.

if we force expert routing to be the same, we see more alignment. i did this by first running hugging face and recording the top k experts for every token at every layer and then replacing torchtitan with those choices.

torchtitan attention gym vs hugging face (forced expert routing)

Full multimodal last-token logit parity (TorchTitan vs HuggingFace)
KL=2.1806e-04 cos=0.999769 max_diff=9.6069e-02 top1=Y top5=80%

torchtitan fla vs hugging face

Full multimodal last-token logit parity (TorchTitan vs HuggingFace)
KL=2.3385e-04 cos=0.999754 max_diff=9.7656e-02 top1=Y top5=80%

performance

^attention gym is 1.38x and uses 26% less memory

attention gym trace:

fla trace:

Stack from ghstack (oldest at bottom):

[ghstack-poisoned]

drisspg · 2026-08-15T22:23:22Z

+    from attn_gym.linear.kda.fwd.recurrent import recurrent_kda
+    from attn_gym.linear.kda.fwd.triton.gate_fwd import bounded_gate_cumsum
+    from attn_gym.linear.kda.fwd.triton.l2norm_fwd import l2norm
+    from attn_gym.linear.kda.short_conv import cute_causal_conv1d_silu


Let me peep this behind a better public interface

## Human Note ## Agent note This is the missing public entrypoint that torchtitan's KDA layer (pytorch/torchtitan#4164) expects for its "recurrent" backend. `recurrent_kda(q, k, v, gate, beta, initial_state=None, *, cu_seqlens=None, output_final_state=False)` mirrors `chunk_kda`'s signature and packed contract: device-resident int32 offsets, repeated offsets as empty padding slots whose state passes through bitwise, a terminal offset below physical capacity leaving the tail outside the contract, and FP32 recurrent states with one leading entry per logical sequence. The distinction from the chunked core is encoded in the argument name: `gate` is the per-token log2 decay (`bounded_gate_cumsum` with `chunk_size=1`), not the chunk-local cumulative gate. Queries scale by 1/sqrt(K) internally, matching `chunk_kda` and the `naive_recurrent_kda` default. The Triton kernel scans tokens sequentially per (sequence, head, value-block) program with the [K, BV] FP32 state held in registers, computing in FP32 regardless of input dtype. It is inference-only: gradient-requiring inputs are rejected with a pointer at `chunk_kda` rather than silently detaching. The op uses the repo's define/impl + fake registration with a fixed-arity (output, state) schema, and joins the lazy `KDA_OPS` surface in `attn_gym.linear`. ## Test Plan ```bash gpu-run auto -- ~/.venvs/ag-linear/bin/python -m pytest test/test_kda_recurrent.py -q # 14 passed on B200: dense/packed vs naive_recurrent_kda (fp32 exact-tolerance, bf16), # non-power-of-two K/V masking, empty-slot state passthrough, capacity tails, # cross-check vs chunk_kda via bounded_gate_cumsum(chunk_size=1 vs 64), gradient # rejection, opcheck, fullgraph compile, and CUDA-graph replay with mutated # boundaries and history. ``` stack-info: PR: #317, branch: drisspg/stack/79

This is the missing public entrypoint that torchtitan's KDA layer (pytorch/torchtitan#4164) expects for its "recurrent" backend. `recurrent_kda(q, k, v, gate, beta, initial_state=None, *, cu_seqlens=None, output_final_state=False)` mirrors `chunk_kda`'s signature and packed contract: device-resident int32 offsets, repeated offsets as empty padding slots whose state passes through bitwise, a terminal offset below physical capacity leaving the tail outside the contract, and FP32 recurrent states with one leading entry per logical sequence. The distinction from the chunked core is encoded in the argument name: `gate` is the per-token log2 decay (`bounded_gate_cumsum` with `chunk_size=1`), not the chunk-local cumulative gate. Queries scale by 1/sqrt(K) internally, matching `chunk_kda` and the `naive_recurrent_kda` default. The Triton kernel scans tokens sequentially per (sequence, head, value-block) program with the [K, BV] FP32 state held in registers, computing in FP32 regardless of input dtype. It is inference-only: gradient-requiring inputs are rejected with a pointer at `chunk_kda` rather than silently detaching. The op uses the repo's define/impl + fake registration with a fixed-arity (output, state) schema, and joins the lazy `KDA_OPS` surface in `attn_gym.linear`. ```bash gpu-run auto -- ~/.venvs/ag-linear/bin/python -m pytest test/test_kda_recurrent.py -q ``` stack-info: PR: #317, branch: drisspg/stack/79

This is the missing public entrypoint that torchtitan's KDA layer (pytorch/torchtitan#4164) expects for its "recurrent" backend. `recurrent_kda(q, k, v, gate, beta, initial_state=None, *, cu_seqlens=None, output_final_state=False)` mirrors `chunk_kda`'s signature and packed contract: device-resident int32 offsets, repeated offsets as empty padding slots whose state passes through bitwise, a terminal offset below physical capacity leaving the tail outside the contract, and FP32 recurrent states with one leading entry per logical sequence. The distinction from the chunked core is encoded in the argument name: `gate` is the per-token log2 decay (`bounded_gate_cumsum` with `chunk_size=1`), not the chunk-local cumulative gate. Queries scale by 1/sqrt(K) internally, matching `chunk_kda` and the `naive_recurrent_kda` default. The Triton kernel scans tokens sequentially per (sequence, head, value-block) program with the [K, BV] FP32 state held in registers, computing in FP32 regardless of input dtype. It is inference-only: gradient-requiring inputs are rejected with a pointer at `chunk_kda` rather than silently detaching. The op uses the repo's define/impl + fake registration with a fixed-arity (output, state) schema, and joins the lazy `KDA_OPS` surface in `attn_gym.linear`. ```bash gpu-run auto -- ~/.venvs/ag-linear/bin/python -m pytest test/test_kda_recurrent.py -q ```

[ghstack-poisoned]

shuhuayu

#4025 has merged, and we may replace fla kda with kda in this pr now.

[ghstack-poisoned]

tianyu-l

I will stamp after @drisspg and @shuhuayu are happy

[ghstack-poisoned]

drisspg

Just smaller comments and perf tweaks but otherwise LGTM

drisspg

mostly nits and some perf stuff otherwise LGTM

[ghstack-poisoned]

shuhuayu

lgtm, thanks! can you check the tests to ensure they pass?

[ghstack-poisoned]

(remaking bc i accideltny merged #4164 into my ghstack instead of main) this PR adds attention module for KDA (relying on attention gym KDA kernels and user facing APIs) **correctness** tested with c4 dataset: <img width="940" height="496" alt="Screenshot 2026-08-24 at 3 04 05 PM" src="https://github.com/user-attachments/assets/c07f5e37-688d-47a6-9e96-e8ea8493e72f" /> torchtitan attention gym vs hugging face (natural routing) ``` Full multimodal last-token logit parity (TorchTitan vs HuggingFace) KL=7.0921e-02 cos=0.927007 max_diff=1.6885e+00 top1=N top5=20% ``` ^i compared hugigng face and torchtitan layer by layer using identical weights and inputs, first divergence appears at block 2 where expert choice agreement was 94.5% and max router score diff was 0.0065. with nearly tied experts, small numerical differences led to different experts being in the top k which led to model executing different expert weights. if we force expert routing to be the same, we see more alignment. i did this by first running hugging face and recording the top k experts for every token at every layer and then replacing torchtitan with those choices. torchtitan attention gym vs hugging face (forced expert routing) ``` Full multimodal last-token logit parity (TorchTitan vs HuggingFace) KL=2.1806e-04 cos=0.999769 max_diff=9.6069e-02 top1=Y top5=80% ``` torchtitan fla vs hugging face ``` Full multimodal last-token logit parity (TorchTitan vs HuggingFace) KL=2.3385e-04 cos=0.999754 max_diff=9.7656e-02 top1=Y top5=80% ``` **performance** <img width="530" height="144" alt="Screenshot 2026-08-24 at 6 35 10 PM" src="https://github.com/user-attachments/assets/a880b257-1a21-4818-b148-e302125ff86c" /> ^attention gym is 1.38x and uses 26% less memory attention gym trace: <img width="1225" height="416" alt="Screenshot 2026-08-24 at 6 41 17 PM" src="https://github.com/user-attachments/assets/11950616-8467-4c85-9353-ebfaafdaf31b" /> fla trace: <img width="1246" height="323" alt="Screenshot 2026-08-24 at 6 41 49 PM" src="https://github.com/user-attachments/assets/6168e16c-a6f4-4e62-a5a9-848c7a02055d" /> Stack from [ghstack](https://github.com/ezyang/ghstack/tree/0.15.0) (oldest at bottom): * #4352 * __->__ #4351

acisseJZhong · 2026-08-27T23:41:03Z

+from torchtitan.models.common.nn_modules import Conv1d
 from torchtitan.protocols.module import Module

-# Shape suffixes:


can we recover this? helps me better understand the math below.

agreed. let's add it back in later prs. one caveat in this file, kda internally still has the B dim iiuc.

Update

7c59ed5

[ghstack-poisoned]

liangel-02 requested review from fegin, tianyu-l, wconstab and wwwjn as code owners August 15, 2026 20:46

pytorch-bot Bot added the ciflow/8gpu label Aug 15, 2026

meta-cla Bot added the CLA Signed This label is managed by the Meta Open Source bot. label Aug 15, 2026

This was referenced Aug 15, 2026

[WIP] kda in torchtitan (vllm integration) #4165

Open

KDA in torchtitan #4115

Closed

liangel-02 requested review from drisspg and shuhuayu August 15, 2026 20:49

liangel-02 added 3 commits August 15, 2026 13:50

Update

d14f950

[ghstack-poisoned]

Update

c123534

[ghstack-poisoned]

Update

c6270e6

[ghstack-poisoned]

drisspg reviewed Aug 15, 2026

View reviewed changes

drisspg mentioned this pull request Aug 16, 2026

Add fused recurrent_kda for decode and inference prefill meta-pytorch/attention-gym#317

Merged

tianyu-l reviewed Aug 17, 2026

View reviewed changes

Comment thread torchtitan/models/common/kda.py Outdated

Update

d374707

[ghstack-poisoned]

shuhuayu requested changes Aug 24, 2026

View reviewed changes

Comment thread torchtitan/models/common/kda.py Outdated

Comment thread torchtitan/models/common/decoder_sharding.py Outdated