Qwen3.5 cp - #1
Open
alexsun07 wants to merge 3 commits into
Open
Conversation
Implements CP for Qwen3.5 hybrid models (GDN linear-attn + full-attn MoE).
Production-validated at cp ∈ {2, 4, 8} on AMD MI300X via verbatim-copy
semantic test (1072-token Chinese prompt).
Architecture: model layer sets a custom forward_batch._gdn_cp_metadata
field (NOT the standard attn_cp_metadata) to avoid triggering MoE
communicator and other downstream paths that assume CP-split hidden_states.
GDN backend slices full-T input internally and runs a per-rank zigzag
two-pass + (b,M) chain reduce with all_gather. Decode-after-CP works via
broadcasting final ssm_state from rank holding the last causal segment
(rank 0 in zigzag) so all CP ranks have correct local cache for the
standard non-CP decode path.
Patches:
- fla/chunk.py + fla/chunk_delta_h.py: add output_M_total option to
chunk_gated_delta_rule_fwd_h with PyTorch reference compute_M_total_pytorch_ref.
Default-off, zero-impact for non-CP callers.
- models/qwen3_5.py: in Qwen3_5MoeForCausalLM.forward, populate
forward_batch._gdn_cp_metadata when CP enabled (without calling
cp_split_and_rebuild_data so hidden_states stays full-T).
- layers/attention/linear/gdn_backend.py: GDNAttnBackend.forward_extend
adds _forward_extend_cp branch implementing zigzag two-pass +
fp32 chain reduce + all_gather output + ssm_state writeback for
decode-after-CP.
Known limitation: cp=2 + attn_tp=4 + Qwen3.5-397B-FP8 + AMD aiter triggers
an upstream sglang bug (V11 stock baseline also fails this combo).
…ents
Remove instrumentation that was added during V1->V27 development but is no
longer needed in production:
- DUMP_HS hidden_state dump in qwen3_5.py
- DUMP_SSM ssm_state dumps in gdn_backend.py (both non-CP and CP paths)
- V17 SSM DIAG block at top of _forward_extend_cp
- [V15 CP FIRE] one-shot print + _v15_cp_path_printed class attr
- Defensive try/except + import os/torch.distributed only used by dumps
Also rewrite comments in why-style instead of citing V-N debug sessions:
- Architectural rationale for the custom _gdn_cp_metadata field stays
(decoupling from MoE communicator's CP path is the key design point)
- Algorithm walkthrough on _forward_extend_cp now describes the math,
not the iteration history
- compute_M_total_pytorch_ref docstring trimmed; PoC v0/v1 wording
replaced with a fused-Triton-can-replace-this pointer
- 'NEW' / 'Option-A' annotations on M_out parameters removed
No algorithm changes; the cp=2 + Qwen3.5-397B-FP8 verbatim echo test
(see report.md \xc2\xa71.3) still passes with the same output fidelity as
the original V21 PASS run.
Net: -92 LOC across the 4 CP files (-213 / +121, mostly reflowed comments).
The previous comment said 'MoE / aiter full-attn layers are not CP-aware,
so all non-GDN layers must keep operating on full sequences', which lumped
two distinct reasons together and was misleading about MoE.
MoE compute itself is per-token and CP-orthogonal. What forces full-T on
the MoE path is sglang's standard CP wrapper chain that activates on
forward_batch.attn_cp_metadata is not None:
- cp_split_and_rebuild_data at the model entry (qwen2_moe.py:778)
- moe_cp_all_gather_into_tensor in the layer communicator
(communicator.py:1106)
- cp_all_gather_rerange_output before lm_head (qwen2_moe.py:831)
All three assume hidden_states are already split to local-T. By using a
private _gdn_cp_metadata field instead, we avoid triggering any of them
and keep the MoE path running unmodified on full-T.
aiter full-attn separately has no CP code at all and would always need
full-T regardless. Comment now states each reason precisely.
No code change.
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
1 participant
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.
Motivation
Modifications
Accuracy Tests
Benchmarking and Profiling
Checklist
Review Process
/tag-run-ci-label,/rerun-failed-ci,/tag-and-rerun-ci