feat(qwen35): shard linear-attention/GDR state per TP rank (split 1/4 of #946) - #1003
Open
Ma1oneZhang wants to merge 3 commits into
Open
feat(qwen35): shard linear-attention/GDR state per TP rank (split 1/4 of #946)#1003Ma1oneZhang wants to merge 3 commits into
Ma1oneZhang wants to merge 3 commits into
Conversation
Phase 2b ported onto pegainfer-project#870: recurrent/conv state, GDR scratch, and the linear-attention weight surface are allocated and addressed per rank instead of replicated, which is what makes 27B TP2 fit on 2x48 GB cards. - weight_loader: additive stitch/shard loaders (2D row stitch for the fused qkv [q|k|v] segments, 1D element stitch for conv1d channels, bf16/f32 1D shards for dt_bias/A_log) - config: local_linear_* accessors mirroring the full-attn TP style; linear head divisibility fails closed in TensorParallelConfig - weights: per-rank stitched shards for in_proj_qkv/conv1d (head-local slices per segment), row shards for z/b/a, col shard for out_proj, dt_bias/A_log sliced, norm_weight kept replicated (head-shared); loader reserve uses TP-aware estimates - recurrent_state/decode_buffers/prefill_buffers: state and GDR scratch at local value-head/qkv sizes; capacity math derives from local allocation_bytes - batch_decode/prefill: local head counts into the GDR decode/conv/ Triton-AOT prefill chains, gated RMSNorm at local v heads, all-reduce after linear out_proj; batch_decode_full_attention_via_prefill is now TP-local so eager decode routes 27B TP2 group-6 full attention through prefill (was FlashInfer Unsupported group_size: 6) - tp_executor: worker capacity math and per-request state use the rank-local sizes; decode rows still run as a per-request bs=1 loop (batched in a follow-up) Recurrent/conv state is never all-reduced. Signed-off-by: Ziyang Zhang <hafuhafu@qq.com>
This was referenced Aug 30, 2026
Contributor
Author
|
@codex please review |
|
You have reached your Codex usage limits for code reviews. You can see your limits in the Codex usage dashboard. |
The TP rank-sliced 1D f32 loader casts the safetensors byte span to f32 exactly like the whole-tensor loader next to it (which already carries the allow); the missed attribute trips clippy::cast-ptr-alignment under the workspace's -D warnings gates. Signed-off-by: Ziyang Zhang <hafuhafu@qq.com>
…f32s via as_chunks if_not_else and chunks_exact_to_as_chunks (pedantic/default) break the Qwen3.5 clippy gate under -D warnings. Signed-off-by: Ziyang Zhang <hafuhafu@qq.com>
Ma1oneZhang
force-pushed
the
feat/qwen35-tp-shard-state
branch
from
August 30, 2026 14:33
568e025 to
fcdeb5a
Compare
Collaborator
|
Thanks for splitting this out. Two items are still needed before this PR is independently mergeable:
|
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
2 participants
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.
Description
Split 1/4 of #946, tracked in #1001. Followed by #1004 → #1005; merge order is this PR first.
Shards the Qwen3.5 linear-attention (GDR) recurrent state and per-request TP state per rank:
LocalGeometrygains rank-local linear-attention head/dim accessors (with indivisibility validation onlinear_num_key_heads/linear_num_value_heads), andRecurrentState/GdrChunkwiseScratch35size themselves from the rank-local geometry instead of global dims.Type of Change
Evidence
cargo check -p pegainfer-qwen35 --features qwen35 --all-targets: clean (sm_89, 2×RTX-4090 build host).cargo test -p pegainfer-qwen35 --features qwen35 --lib: 106/106 pass (7 GPU-gated tests skipped: no model weights on this host).cargo fmt --all -- --check: clean.GPU-gated acceptance from the original #946 (27B TP2 HF/E2E + throughput A/B, measured 2026-08-20 on the pre-rebase branch) still needs a rerun on this rebased stack and is tracked in #1001; no numbers are re-claimed here.
Checklist
docs/conventions/coding-style.md).CLAUDE.md).