Skip to content

Replace Qwen3.5 FLA paths with Attention Gym - #1

Closed
drisspg wants to merge 1 commit into
gdn_attn_gymfrom
gdn-attn-gym-paged
Closed

Replace Qwen3.5 FLA paths with Attention Gym#1
drisspg wants to merge 1 commit into
gdn_attn_gymfrom
gdn-attn-gym-paged

Conversation

@drisspg

@drisspg drisspg commented Sep 1, 2026

Copy link
Copy Markdown
Owner

Human Note

Agent note

This is a draft follow-up to pytorch#4363 and depends on the paged chunk GDN addition
being prepared on top of meta-pytorch/attention-gym#425.

Use Attention Gym for Qwen3.5 GDN training, batch-invariant backward recomputation, recurrent
execution, decode, and paged chunk prefill. The paged prefill path advances vLLM SSM cache slots
in place, eliminating the previous gather/chunk/scatter copies and preserving the shared
[slots, H, V, K] layout across prefill and decode.

Remove the now-unused FLA backend selection and direct dependency, normalize Q/K explicitly at
the Attention Gym boundary, and update debug configurations to the fused backend K = V = 128
contract while retaining two key heads for TP=2 coverage.

The remaining serving-side copy is the much smaller convolution history [N, W - 1, C] during
varlen prefill. Attention Gym currently has paged single-token convolution decode, but not a paged
varlen convolution prefill API.

Performance

GB200, fixed-pointer warm-cache microbenchmark, T=4096, N=32, HK=8, H=16, K=V=128,
BF16 QKV, FP32 state:

Path Median GPU time
Caller gather + chunk + scatter 219.38 us
Direct paged chunk 185.95 us

Direct paging reduced this isolated subsystem time by 15.2%. This is not an end-to-end RL result.

Test Plan

PYTHONPATH=/home/drisspg/meta/attention-gym-pr425:/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
PYTHONPATH=/home/drisspg/meta/attention-gym-pr425:/home/drisspg/meta/torchtitan pytest -q tests/unit_tests/cpu/test_state_dict_keys.py tests/unit_tests/cpu/test_integration_test_definitions.py tests/unit_tests/cpu/test_train_spec.py -x

The focused suites passed locally on GB200. Full RL/vLLM validation was not run because the local
environments do not contain vLLM. All configured formatting and lint hooks passed; repository-wide
Pyrefly remains blocked by unrelated missing optional dependencies and existing baseline errors.

## Human Note

## Agent note

Use Attention Gym for Qwen3.5 GDN training, batch-invariant backward recomputation, recurrent
execution, decode, and paged chunk prefill. The paged prefill path advances vLLM SSM cache slots
in place, eliminating the previous gather/chunk/scatter copies and preserving the shared
[slots, H, V, K] layout across prefill and decode.

Remove the now-unused FLA backend selection and dependency, normalize Q/K explicitly at the
Attention Gym boundary, and update debug configurations to the fused backend K=V=128 contract
while retaining two key heads for TP=2 coverage.

## Performance

On GB200 with T=4096, N=32, HK=8, H=16, K=V=128, BF16 QKV, and FP32 state, five interleaved
30-iteration warm-cache profiler rounds measured 219.38 us for caller-side gather/chunk/scatter
and 185.95 us for direct paged prefill (15.2% lower GPU time).

## Test Plan

```bash
PYTHONPATH=/home/drisspg/meta/attention-gym-pr425:/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
PYTHONPATH=/home/drisspg/meta/attention-gym-pr425:/home/drisspg/meta/torchtitan pytest -q tests/unit_tests/cpu/test_state_dict_keys.py tests/unit_tests/cpu/test_integration_test_definitions.py tests/unit_tests/cpu/test_train_spec.py -x
```
@drisspg

drisspg commented Sep 1, 2026

Copy link
Copy Markdown
Owner Author

Closing this fork-only PR; the change is being resubmitted as a draft against pytorch/torchtitan on top of pytorch#4363.

@drisspg drisspg closed this Sep 1, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant