Skip to content

docs(qwen35): TP2 implementation record + knowledge-eval tooling (split 4/4 of #946) - #1006

Open
Ma1oneZhang wants to merge 15 commits into
pegainfer-project:mainfrom
Ma1oneZhang:dev/qwen35-tp-eval-docs
Open

docs(qwen35): TP2 implementation record + knowledge-eval tooling (split 4/4 of #946)#1006
Ma1oneZhang wants to merge 15 commits into
pegainfer-project:mainfrom
Ma1oneZhang:dev/qwen35-tp-eval-docs

Conversation

@Ma1oneZhang

Copy link
Copy Markdown
Contributor

Description

Split 4/4 of #946, tracked in #1001. Independent of splits 1–3 (docs + bench tooling only; no runtime dependency).

  • Documents the Qwen3.5 TP implementation record and the 27B TP2 knowledge-eval setup.
  • Fixes scripts/eval_mc.py MMLU-Redux extraction (absorbs the codex review comment from feat(qwen35): shard GDR state per TP rank + batched eager/graph decode under TP (+27B benchmark validation) #946): answers are extracted via the shared sg_extract_labels(text, 'ABCD') marker logic instead of a first-capital-letter regex, matching the other benchmarks.
  • Flags the 27B TP2 MMLU-Redux 94.09 figure in the doc as produced by the pre-fix extractor: do not cite until rerun with the fixed extractor.

Type of Change

  • Documentation update
  • Bug fix (non-breaking change which fixes an issue)

Evidence

Checklist

  • My code follows the style guidelines of this project (see docs/conventions/coding-style.md).
  • I have performed a self-review of my own code.
  • I have formatted my commits according to Commitizen conventions.
  • I have run the local test suite and all tests pass (see CLAUDE.md).

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>
@Ma1oneZhang

Copy link
Copy Markdown
Contributor Author

@codex please review

@chatgpt-codex-connector

Copy link
Copy Markdown

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>
… 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>
Port of the batched eager TP decode step onto pegainfer-project#870's worker structure:
decode rows in one command now run as ONE batched forward per step on
every rank plus one batched rank-0 sampling pass, instead of a
per-request bs=1 loop.

- run_decode_batch resolves every decode row's worker state in command
  order, builds a step-scoped LinearStatePointerTables over the whole
  batch (from_recurrent_refs(..., bs, ...)), runs one
  batch_decode_eager_logits forward, then rank 0 snapshots all requested
  logprob rows before one batched select_batch over per-row params
- execute_decode_rows (used by both decode-only and unified steps) calls
  run_decode_batch once; per-row results fan out in command order
- TpRequestState.linear_pointer_tables (capacity-1, decode-only) removed;
  ensure_prefill_state no longer builds it. The step-scoped table is
  rebuilt every step, so swap_remove retirement can't stale it

Seeded rows keep per-row semantics: select_batch isolates each seeded
row into its own single-row philox call keyed on (request seed, step 0),
so seeded output stays independent of batch composition. Unseeded rows
decorrelate via the per-step command seed, same as the single-GPU
batched path.

Reference (27B TP2, 2x RTX 4090, eager): 16 concurrent 256-token
completions aggregate 24.9 -> 292.3 tok/s; single-request unchanged.

Signed-off-by: Ziyang Zhang <hafuhafu@qq.com>
P2c from docs/models/qwen35/tp-design.md. TP decode captures/replays
per-bucket CUDA Graphs when --cuda-graph is set AND the TP-local decode
GQA group has a compiled kernel (4B/9B TP2); 27B TP2 (group 6) keeps the
batched eager path byte-for-byte under the gate.

- Gate: drop the fail-closed TP+graph rejections in config.rs/lib.rs/
  tp_executor.rs; log once when graph was requested but the group gate
  keeps decode eager.
- State: scheduler owns dense decode slots (slot_idx on TP decode rows,
  slot_for_new_request at promote, compaction_after_retire on retire);
  workers hold a fixed-address BatchDecodeGraphState plus slot_map, D2D
  copy prefill state into the slot on the first decode row, and apply
  DropRequest compactions via move_slot_within with occupancy assertions
  (poison on mismatch).
- Capture/replay: startup pre-capture sweep ported from qwen3 (Warmup
  via new warmup_tp_collective, Capture/Launch per bucket, Finalize)
  behind a 600 s abort watchdog; batch_decode_graph gains DecodeGraphUse
  (Serve/CaptureOnly/Replay); TP serving is replay-only; graphs drop
  before the NCCL comm.
- Convenience executor API (model-local tests) keeps a slot tracker
  mirroring Qwen35Executor; scheduler flows pass explicit slots.

Validation (2x RTX 4090, venv NCCL): lib 105/105; 9B TP2 HF gates eager
+ graph (sequential replay, bucket-straddling, post-compaction) pass
within existing tolerances; 9B TP2 scheduler e2e eager+graph pass;
serving_tp2 now launches with graph on; 27B TP2 HF+e2e pass unchanged
(group-6 gate keeps eager, graph variant self-skips). Serving benchmark
9B TP2, 16 concurrent x 256 out: 767.15 tok/s graph vs 705.86 eager
(+8.7% steady output, TPOT 20.04 vs 21.99 ms).

Signed-off-by: Ziyang Zhang <hafuhafu@qq.com>
The graph-slot pre-reserve used bucket_for(requested_max_batch), but the
graph state is later allocated as bucket_for(effective_max_batch). On a
tight-memory rank the oversized reserve could starve the effective
recurrent capacity to zero and refuse startup.

Iterate the reserve bucket downward against the fitted capacity until it
stabilises (the bucket only shrinks, so it converges), and clamp the
fitted batch to the reserved bucket so the later bucket_for(effective)
allocation never exceeds the reserve.

Absorbs the codex review comment on Ma1oneZhang/pegainfer PR pegainfer-project#946
(tp_executor.rs graph_slot_reserve).

Signed-off-by: Ziyang Zhang <hafuhafu@qq.com>
…group predicate

The rank-local GQA group equals the global group under validated head
sharding, so the P2B review removed LocalGeometry::local_decode_group_is_compiled.
Route the P2c graph gate through Config35::decode_group_is_compiled instead
of re-adding the local duplicate; the gate decision was already identical on
every rank.

Evidence (2x RTX 4090, sm_89): cargo check/clippy --release --all-targets
-D warnings clean; qwen35 lib tests 107 passed / 0 failed; cargo fmt clean.

Signed-off-by: Ziyang Zhang <hafuhafu@qq.com>
Signed-off-by: Ziyang Zhang <hafuhafu@qq.com>
…eval

- tp-implementation.md: Phase 2a/2b + Step 3 batched-decode landing record
- benchmarks/qwen35-27b-tp2-knowledge-eval.md: MMLU-Redux 94.09 vs 93.2,
  C-Eval 88.11 vs 90.5 (full runs, truncation rerun-merged); in band
- scripts/eval_mc.py + eval_rerun_truncated.py: benchmark runner
  (OpenCompass/lm-eval recipe replicas over the chat completions API)
- developer-onboarding: NCCL libnccl.so dlopen note

Signed-off-by: Ziyang Zhang <hafuhafu@qq.com>
The MMLU-Redux extractor searched the whole completion for the first
[ABCD], so any answer prefixed with "Answer: B" scored as the "A" in
"Answer", silently marking correct B/C/D responses wrong (codex review
on pegainfer-project#946). Reuse the SuperGPQA marker-preferring patterns restricted to
A-D, and drop the now-dead MMLUREDUX_RE. Redo the MMLU-Redux numbers
before citing them.

Signed-off-by: Ziyang Zhang <hafuhafu@qq.com>
The 94.09 snapshot predates the answer-extraction fix; mark it as not
citable until rerun with the marker-preferring extractor.

Signed-off-by: Ziyang Zhang <hafuhafu@qq.com>
The second LocalGeometry guard was dropped in review: Config35 already
validates linear_num_value_heads % linear_num_key_heads == 0, so the
key-head % tp guard implies value-head divisibility. Align the P2/P2b
requirement wording with the single source of truth.

Signed-off-by: Ziyang Zhang <hafuhafu@qq.com>
@Ma1oneZhang
Ma1oneZhang force-pushed the dev/qwen35-tp-eval-docs branch from 6692b61 to 7a3f543 Compare September 5, 2026 10:16
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