feat(qwen35): shard linear-attention/GDR state per TP rank (split 1/4 of #946) - #1003
Conversation
|
@codex please review |
|
You have reached your Codex usage limits for code reviews. You can see your limits in the Codex usage dashboard. |
568e025 to
fcdeb5a
Compare
|
Thanks for splitting this out. Two items are still needed before this PR is independently mergeable:
|
|
Exact-head acceptance rerun on
Load/HBM: peak per-rank 35,988 / 36,822 MiB of 49,140 MiB -- 27B TP2 fits on 2x48GB with sharded GDR state; each test process fully freed HBM between runs. Host notes: no system NCCL on this box; used pip wheel The |
|
Thanks — the exact-code-head TP2 acceptance and the documentation update resolve the previous requests. One blocker remains in |
FeathBow
left a comment
FeathBow
left a comment
There was a problem hiding this comment.
Thanks! The core P2B design is coherent, some smaill changes for the acceptance boundary. The remaining notes are non-blocking cleanup suggestions intended to keep the patch focused.
| kv.advance(1); | ||
| recurrent_states[i].seq_len += 1; | ||
| positions.push(pos as i32); | ||
| start_positions.push(pos); |
There was a problem hiding this comment.
start_positions is allocated and filled before the GQA predicate is checked. It is needed for the 27B group-6 paged-prefill fallback, but for models whose group has a compiled decode kernel it is immediately discarded. Those compiled-group paths therefore perform one avoidable host allocation and fill per rank per token.
The local predicate also appears equivalent to the existing config predicate: both full-attention head counts are divided by the same TP width, so (q / tp) / (kv / tp) == q / kv for every validated geometry. As a result, LocalGeometry::local_decode_group_is_compiled() returns the same result as Config35::decode_group_is_compiled() for valid configurations.
Would it be possible to create or derive start_positions only inside the uncompiled branch and share the one-token paged-plan construction with the existing hybrid fallback? This would keep the compiled route simpler and avoid maintaining two similar builders. I have not measured a throughput regression here.
| world_size: tp.world_size(), | ||
| }); | ||
| } | ||
| if !config |
There was a problem hiding this comment.
LocalGeometry receives a validated Config35, and Config35::try_from already guarantees linear_num_value_heads % linear_num_key_heads == 0. Once the preceding guard proves linear_num_key_heads % world_size == 0, value-head divisibility follows algebraically; this branch cannot fail for a production config.
The new test reaches this branch by mutating linear_num_value_heads from 32 to 31 after validation, which constructs a state outside the Config35 invariant rather than exercising an input boundary. Could we keep the key-head TP guard, exercise it with an internally consistent case such as key/value = 17/34, and omit this derived branch, its test, and the corresponding documentation claim? That would leave one source of truth for the relationship between these head counts.
| } | ||
|
|
||
| #[test] | ||
| fn linear_attention_tp2_slices_match_synthetic_checkpoint() { |
There was a problem hiding this comment.
This test never calls load_tensor_2d_row_stitch, load_tensor_1d_stitch, either model loader wrapper, H2D, or a GDR kernel. It serializes/deserializes an F32 safetensors blob even though the production stitch loaders consume BF16, then manually slices CPU vectors with the same segments tuples it is asserting. The QKV expected value is built from those same tuples, while the conv and value-head assertions duplicate the two exact segment tests immediately above.
Could we remove safetensors_fixture_f32 and this test, and rely on the direct QKV/conv segment tests for the non-obvious local layout plus the real HF gate for the production loader/kernel path? The adjacent TP1 segment test also appears unnecessary because its own comment notes that the TP1 production branch does not call this helper. This would avoid implying production checkpoint coverage where the test currently verifies dependency parsing and repeated slice arithmetic.
| local_full_attn_gated_q_dim: usize, | ||
| local_linear_num_key_heads: usize, | ||
| local_linear_num_value_heads: usize, | ||
| local_linear_q_dim: usize, |
There was a problem hiding this comment.
No production code reads this field; it is stored only so the new arithmetic tests can inspect it. The constructor already needs a local q_dim to derive local_linear_qkv_dim, so do that calculation locally and omit the field. This avoids carrying test-only state in the runtime geometry contract.
| //! - Recurrent state: [local_value_heads, key_head_dim, value_head_dim] f32, V contiguous ([H,K,V]) | ||
| //! - Conv state: [local_qkv_dim × (conv_kernel_dim - 1)] bf16 | ||
| //! | ||
| //! Under TP both are rank-local: Phase 2b shards value heads (and fused qkv |
There was a problem hiding this comment.
The local shape and "never all-reduced" rule are useful context, but the same Phase-2b explanation now appears at the module, both fields, size helper, allocator callers, prefill, decode, loader sites, and worker capacity calculation. Could we keep one module/type-level ownership invariant plus the non-obvious fused-QKV stitching and post-out-proj reduction comments, while trimming phase labels and comments that restate the following expression or method name? That would preserve the durable design context with less repetition in production code.
FeathBow
left a comment
FeathBow
left a comment
There was a problem hiding this comment.
seems all resolved, please rebase origin main :)
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>
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>
… record Signed-off-by: Ziyang Zhang <hafuhafu@qq.com>
Signed-off-by: Ziyang Zhang <hafuhafu@qq.com>
…te, prune fixture tests Address the P2B review on the GDR state-sharding split: - config/tp.rs drops the linear_num_value_heads % world_size guard: Config35::try_from already validates value % key == 0, so the key-head guard implies value divisibility; one source of truth. The surviving test uses an internally consistent key/value = 17/34 fixture. - Drop the test-only local_linear_q_dim field; the constructor keeps q_dim as a local to derive local_linear_qkv_dim. Geometry tests re-derive q from qkv/z dims. - batch_decode_eager_logits derives start_positions only inside the paged-prefill fallback branch instead of allocating per rank per token on the compiled path, and routes on Config35::decode_group_is_compiled: head sharding leaves the q-per-kv group unchanged, so the rank-local predicate duplicated the config one. The one-token paged plan is now shared with the hybrid fallback via one_token_paged_plan(). - weights.rs deletes the F32 safetensors fixture test and the TP1 segment test: they re-asserted the segment tuples they built expectations from without exercising the production BF16 stitch loaders. Segment coverage stays on the direct QKV/conv tests plus the HF golden gates. - Trim the repeated Phase-2b ownership explanation across recurrent_state, prefill, prefill_buffers, batch_decode, tp_executor, and config/tp to one module-level invariant plus the non-obvious fused-QKV stitching and post-out_proj all-reduce notes. - tp-implementation.md P2B record updated for the dropped value-head guard, the deleted local predicate, and the removed local_linear_q_dim dim. Evidence (2x RTX 4090, sm_89): cargo check/clippy --release --all-targets -D warnings clean; qwen35 lib tests 104 passed / 0 failed; cargo fmt clean. TP2 gpu gates not rerun (no local weights). Signed-off-by: Ziyang Zhang <hafuhafu@qq.com>
6867111 to
fb8206a
Compare
FeathBow
left a comment
FeathBow
left a comment
There was a problem hiding this comment.
nothing below blocks merge. remaining non-blocking cleanup.
| self.config.head_dim | ||
| ) | ||
| })?; | ||
| "hybrid decode", |
There was a problem hiding this comment.
one_token_paged_plan now has two callers: the eager route passes geometry.local_num_attention_heads() / local_num_key_value_heads(), while the hybrid route passes self.config.num_attention_heads / num_key_value_heads. The hybrid route is only reachable from the graph path, and TP refuses graphs, so the values are equal today. I recommend using the geometry accessors in both call sites so the helper has one contract and the equality does not depend on a reader remembering the graph/TP exclusion.
| } | ||
|
|
||
| #[test] | ||
| fn computes_tp2_linear_attention_local_dimensions() { |
There was a problem hiding this comment.
None of them can fail without the formula changing, so they detect edits rather than protect a boundary. I recommend removing them and keeping requires_linear_attention_key_head_divisibility (a real input boundary) and linear_qkv_shard_segments_stitch_head_local_slices (the non-obvious per-segment tiling)
…ived-dim tests Review follow-ups on pegainfer-project#1003 (non-blocking): - one_token_paged_plan's hybrid caller now passes the geometry's local head counts like the eager caller, so the helper has one contract instead of relying on the graph/TP exclusion to keep the two equal. - The TP2/TP1 local-linear-dimension tests only restated the accessor formulas; removed. requires_linear_attention_key_head_divisibility stays as the real input boundary. Signed-off-by: Ziyang Zhang <hafuhafu@qq.com> Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
3 participants
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).