feat(glm52): TP8 span row mapping — single-user speculative decode at 186 tok/s#610
Conversation
Generalizes the TP8 phase kernels' row-to-owner mapping behind
Glm52Tp8RowMap: Dp8 (row i = rank i's token, the bucket-1 c8 shape,
unchanged addressing) and Span { owner } (all 8 rows are the owner's —
1 committed token + 7 speculative drafts, so bs=1 turns pads into
drafts at unchanged MoE cost). Only where rows come from (AG push) and
where results go (RS) changes; union/GEMM/SiLU phases are mapping-
agnostic. The RS region grows a row dimension ([parity][row][src][h],
8x VA; dp8 uses its own-row stripe) so both mappings share one layout.
Span pad ranks zero-fill their 8-row mlp_out (deterministic, no NaN
leaks into the next layer through a pad row).
New gate (layer_moe_tp8_span_matches_dp8): the same 8 rows through both
mappings must be BIT-identical — same gathered xg, same union order,
same fixed-order reductions — plus the pad zero-fill contract.
First slice of the TP8+MTP campaign (span-8 speculative decode on the
TP8 topology).
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…owner The span owner id moves from a kernel launch parameter (baked at graph capture) to a device scalar read at kernel time, like the step epoch: the host stages it once per span step (stage_span_owner), so ONE captured span graph serves whichever rank holds the solo request — no owner-pinning in admission, no per-owner graph zoo. Every rank passes full 8-row buffers in span mode (any rank can be the staged owner); the kernels read inputs only on the owner. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
The tp8 scheduler serves exactly two shapes now: bucket-1 dp8 (any concurrency, one slot per rank) and the bucket-8 span shape, planned when the FLEET holds exactly one active request wanting more than one row. All 8 rows run on the owning rank through the span MoE mapping — prefill spans immediately ride it (8 prompt tokens/step instead of tp8's previous 1), and the speculative verify rounds land on the same shape next. Wiring: plan_step_shapes grows the solo-span rule (pure, unit-tested); Glm52StepFlags carries tp8_span_owner, staged to device memory in the step prologue before capture/replay so ONE captured bucket-8 graph serves any owner; the step-head epoch advance and the MoeTp8 arm handle the span batch; precapture warms bucket-8 x 2 tiers with rank-0-owner pad steps. Span steps never lease or consume launch-ahead (enforced). run_dp8_coordinator now takes the topology and derives its caps. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
… shape Lifts the tp8 x dspark launch ban: verify spans ride the bucket-8 span shape (K2), whose MoE cost is bucket-1-flat, so the EP8 span-4 verdict (bucket-4 round beats span-8's accepted tail) inverts and the tp8 coordinator feeds the drafter's full 7-draft proposal. The span cap is now a per-round argument derived from the topology; EP8 keeps its measured 3. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…decode Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…row bond Toxic-review findings: the span owner slice moves into the Glm52Tp8RowMap::Span variant so a span launch without a staged owner is unrepresentable (two runtime ensures deleted); a const assert bonds the scheduler's largest bucket to the kernel's row count (a bigger bucket would silently compute 8 of N rows); GLM52_DSPARK_EP8_SPAN_DRAFTS says what it now is; the RS-region KV-capacity cost (~900 MiB/rank, ~5.5k tokens) is recorded in the design doc. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 3720f799dc
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| let span_drafts = if tp8_solo_span { | ||
| crate::dspark::GLM52_DSPARK_DRAFTS | ||
| } else { |
There was a problem hiding this comment.
When --moe-topo tp8 has more than one active request, plan_step_shapes intentionally clamps the fleet back to bucket-1, so a request cannot consume any drafted rows. This unconditional TP8 cap still installs all 7 DSpark drafts after each step; on the next bucket-1 step advance_span feeds only the anchor (drafts_fed == 0) and clears those drafts, so concurrent TP8+DSpark spends a draft round every step without ever accepting a draft. Consider suppressing proposals or installing drafts only when the next step can actually be a solo span.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Right — the c≥2 draft round was pure waste under the bucket-1 clamp. Fixed in af0ec9a: proposals are suppressed while the fleet holds more than one active request (appends/resets still flow so the drafter's shadow KV stays fresh; drafting resumes one round after the fleet drains back to solo).
|
Greedy losslessness A/B (tp8+DSpark vs tp8 plain, same prompts, 400 tokens): counting output is byte-identical (688/688); code diverges at byte 599/1309 ( |
A concurrent tp8 fleet is clamped to bucket-1, where a verify span can never be fed — the drafter forward every step bought zero accepted tokens. Proposals are suppressed while more than one request is active; appends/resets still flow so the shadow KV stays fresh and drafting resumes one round after the fleet drains back to solo. Reported by Codex review on #610. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
1 participant
What
The TP8 span row mapping: at bs=1 the bucket-1 shape computes m=8 with 1 real row
and 7 pads — this branch turns those pads into work. All 8 rows of the TP8 MoE
kernel can now belong to ONE owner rank (1 committed token + 7 DSpark drafts, or 8
prefill positions), gathered from the owner instead of one-per-rank. The MoE cost is
unchanged: the compute phases (union/GEMM/SiLU) are mapping-agnostic, only the
all-gather source and reduce-scatter destination patterns differ. DSpark verify
spans and solo prefill ride the new bucket-8 span shape;
--moe-topo tp8+--dflash-draft-model-pathnow compose.Measured (8×H200, GLM5.2-FP8, solo, greedy, 500-token completions)
Rounds run at 24.6–25.6 ms p50 / ≤29 ms p99 regardless of workload; accept incl.
bonus: code 4.80 (19/104 rounds accept all 7 drafts), prose 2.64, counting 3.37 —
span-8 releases the accept tail that EP8's measured span-4 verdict had to cut
(bucket-4 vs bucket-8 round costs; under TP8 the 8-row round costs the same as 1).
Bonus from the same shape: solo prefill runs 8 prompt tokens/step — ~2400-token cold
TTFT 5.37 s vs ~36 s at tp8's previous 1 token/step.
How
glm52_moe_tp8.cu): row-to-owner mapping generalized. The owner id isread from device memory at kernel time (staged per step like the epoch), so ONE
captured bucket-8 span graph serves whichever rank holds the request — no
owner-pinning, no per-owner graph zoo. The RS LL region grows a row dimension
(
[parity][row][src][hidden], one layout for both mappings; ~900 MiB/rank for 75layers, deducted from KV capacity by the VRAM probe — recorded in the design doc).
Span pad ranks zero-fill their 8-row mlp_out. The spin iron rule (wait only on
cross-rank packets from previous kernels) and the two-step parity handshake hold
in mixed dp8/span step sequences.
bucket-8 span, planned only when the fleet holds exactly one active request
wanting >1 row (pure rule in
plan_step_shapes, unit-tested). Span steps neverlease/consume launch-ahead (structural:
feed_want > 1blocks the lease). Theverify-span draft budget is a per-round argument: EP8 keeps its measured 3, tp8
feeds the drafter's full 7.
Glm52Tp8RowMap::Spanvariant (a span launch without an owner is unrepresentable); a const assert bonds
the scheduler's largest bucket to the kernel's row count.
Gates
layer_moe_tp8_span_matches_dp8): the same 8rows through both mappings must match exactly (same gathered activations, same
expert-union order, same fixed-order reductions) + NaN-poisoned pad zero-fill
check. Green on the final SHA, alongside the existing TP8 layer gate (62/64
probes, 2 tie-flips within the 4 allowance).
shutdown ≤5 s, EP8+DSpark regression (span cap 3 intact — accept histogram shows
zero rounds beyond 3 drafts; determinism ×2).
Design record:
docs/models/glm52/moe-tp8-low-latency.md("TP8 + MTP" section,including the SSE-events≠tokens measurement trap). Next levers: M4 attention-TP
(cuts the ~16 ms non-MoE share of the 25 ms round), M3 multi-user bucket decision.
🤖 Generated with Claude Code