Replace Qwen3.5 FLA paths with Attention Gym

[drisspg]

Human Note

Agent note

This incorporates Angel Li's Attention Gym migration from #4363 and extends it so the resulting
change can land as one self-contained PR against main. It uses Attention Gym for Qwen3.5 GDN
training, batch-invariant backward recomputation, recurrent execution, decode, and paged prefill.

The vLLM path advances both convolution and SSM cache slots in place. This removes the prior
prefill-time temporary state allocations, grouped-head expansion, gather/scatter copies, transposes,
and host-syncing continuation check while preserving the shared [slots, H, V, K] SSM layout.
The paged operations landed in Attention Gym as
meta-pytorch/attention-gym#431 and
meta-pytorch/attention-gym#432.

The now-unused FLA backend selection and direct dependency are removed. Q/K normalization is explicit
at the Attention Gym boundary, and debug configurations use the fused backend's K = V = 128
contract while retaining two key heads for TP=2 coverage.

Note

The required Attention Gym APIs are on its main branch at 6ee58c2a. PyPI currently stops at
attn-gym==0.0.6, so this draft still needs the next Attention Gym release and corresponding pin
update before it is ready to land. A clean wheel built from 6ee58c2a was installed into an isolated
target and successfully exported both paged APIs to this TorchTitan adapter.

Performance

End-to-end vLLM: before either PR vs combined candidate

GB200, Qwen3.5 debug model with production GDN dimensions K=V=128, prompt length 128, generation
length 64, TP=1, FULL_AND_PIECEWISE CUDA graphs, five measured rounds per engine launch. Three engine
launches were paired at the parent of #4363 (9aadc3616) and the combined candidate; the table reports
median throughput across launches and the median paired change.

Batch	Before either PR output tok/s	Combined output tok/s	Median paired change
1	449.3	482.9	+7.5%
2	869.2	936.0	+12.3%
4	1683.9	1786.8	+6.1%
8	3184.6	3389.6	+5.3%
16	5874.8	6251.2	+6.4%

Process-level throughput varied, but every paired before/after launch improved at every batch size.

GDN prefill ablation

GB200, fixed-pointer warm-cache, HK=8, H=16, K=V=128, BF16 QKV, FP32 state, five interleaved
rounds of 30 iterations:

Workload	Before: FLA + copies	Attention Gym + copies	Attention Gym direct paged	Direct vs before
N=1, T=1024	105.4 us	124.5 us	89.5 us	-15.1%
N=8, T=1024	109.2 us	110.2 us	86.9 us	-20.4%
N=8, T=4096	221.6 us	183.3 us	153.1 us	-30.9%
N=32, T=4096	314.3 us	269.0 us	203.3 us	-35.3%
N=64, T=8192	574.7 us	466.8 us	349.5 us	-39.2%

The middle column isolates the backend migration while retaining state copies; direct paging removes
another 16.5-28.1% from that path.

Convolution prefill ablation

At the Qwen3.5 TP=2-local shape (C=4096, W=4, BF16), direct paging reduced GPU time by
43.8-79.4% for fresh prompts and 46.2-80.4% for resumed prefixes across N=1-64 and 128-512
tokens per sequence.

Test Plan

# Attention Gym paged convolution, paged GDN, graph/fake-tensor, and namespace coverage
PYTHONPATH=/home/drisspg/meta/attention-gym .venv/bin/pytest -q -n 6 test/test_short_conv_cute.py test/linear/test_gdn_chunk_fused.py test/test_namespaces.py
# 151 passed, 5 skipped

# Focused TorchTitan Qwen3.5 coverage
PYTHONPATH=/home/drisspg/meta/attention-gym:/home/drisspg/meta/torchtitan pytest -q tests/unit_tests/gpu/test_qwen3_5_deltanet.py tests/unit_tests/test_qwen3_5_mrope_positions.py -x
# 10 passed

# Two-GPU trainer/vLLM parity with grouped-head TP=2 GDN
PYTHONPATH=/home/drisspg/meta/torchtitan:/home/drisspg/meta/attention-gym torchrun --nproc-per-node=2 -m pytest -q torchtitan/experiments/rl/tests/test_bitwise_parity.py::TestBitwiseParityQwen35DebugVarlen
# 4 passed; trainer/vLLM prefill and decode comparisons are bitwise identical

# Non-batch-invariant prefix-cache continuation and mixed decode+prefill
PYTHONPATH=/home/drisspg/meta/torchtitan:/home/drisspg/meta/attention-gym torchrun --nproc-per-node=1 artifacts/pr4389/test_vllm_paged_continuity.py
# Reused 528 cached tokens at a 1056-token boundary; cached/fresh generated IDs matched,
# max logprob delta=1.240e-3; forced mixed decode+prefill completed both requests

# Two-GPU non-batch-invariant RL integration smoke
python -m torchtitan.experiments.rl.train --module alphabet_sort --config rl_grpo_qwen3_5_debug_varlen --async-loop.num-training-steps 10 --async-loop.training-sample-builder.no-drop-zero-std-reward-groups ...
# 10/10 steps completed, including rollout, prefill/decode, forward/backward, and weight sync

The real Qwen3.5-9B vLLM benchmark was also attempted, but the installed SM100 FA4 path rejects
head_dim=256 with seqused_q/seqused_k before GDN executes. The matched debug-model benchmark and
TP=2 parity suite exercise the Qwen3.5 GDN paths without that unrelated attention-backend limitation.