Replace Qwen3.5 FLA paths with Attention Gym - #1
Closed
drisspg wants to merge 1 commit into
Closed
Conversation
## 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 ```
Owner
Author
|
Closing this fork-only PR; the change is being resubmitted as a draft against pytorch/torchtitan on top of pytorch#4363. |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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 = 128contract while retaining two key heads for TP=2 coverage.
The remaining serving-side copy is the much smaller convolution history
[N, W - 1, C]duringvarlen 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:
Direct paging reduced this isolated subsystem time by 15.2%. This is not an end-to-end RL result.
Test Plan
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.