feat(pp) support dynamic chunked pipeline parallel

[wanzhenchn]

Dynamic Chunked Pipeline Parallelism

Background

feat(pp) support dynamic chunked pipeline parallel
Motivated by SGLang Dynamic Chunking and vLLM Ascend Dynamic CPP.

With fixed-size prefill chunks, attention cost grows with the cached prefix, so later chunks take longer than earlier ones. Pipeline stages fall out of sync and the pipeline develops bubbles. Dynamic chunking shrinks later chunks so that every chunk costs about the same.

Principle

Latency model

A chunk of x new tokens after a cached prefix of L tokens is modeled as:

$$t(L, x) = c + \gamma L + b x + a (2Lx + x^2)$$

• L is the cached-prefix length, x is the number of new tokens in the current chunk, and t(L, x) is the predicted chunk latency.
• c is the fixed per-forward overhead, paid even when L = 0.
• gamma * L captures prefix-dependent work such as rebuilding cached-prefix state in MLA models. It is kept separate from c rather than folded into one constant because the two scale differently with L: gamma * L vanishes at L = 0 and grows with the cached prefix, so collapsing them would erase exactly the prefix-dependent behavior this feature acts on.
• b * x captures work that grows linearly with the new-token count.
• a * (2 * L * x + x^2) models self-attention work: 2 * L * x for the new chunk against its cached prefix, x^2 for interactions within the new chunk.

a, b, c, and gamma are learned automatically from measured forward latency, so no offline coefficients or calibration flags are required.

Solving for the next chunk

Every chunk of a request should cost what the first one did. The first chunk runs at L = 0, so its cost is c + b*x0 + a*x0^2. Since c is paid by every forward it cancels from both sides, and the budget a later chunk has to match is

$$b x_0 + a x_0^2 - \gamma L$$

• gamma * L is subtracted because it is a floor the chunk owes before any of its own tokens are attended to. Equalizing only the attention-area terms would leave every chunk paying that floor on top of an already-equal budget, making later chunks both too large to be equal-latency and too numerous.
• The scheduler solves the resulting quadratic for the x that spends the remaining budget, then applies the serving constraints: the smoothing factor, the minimum chunk size, an upper bound of the initial chunk size and the batch-token budget, and alignment to max(KV block size, 64).
• When the prefix rebuild alone exhausts the budget, no chunk size can match the first chunk; ATOM stops shrinking and falls back to fixed-size chunking.

Equal-cost chunks keep pipeline stages in sync, which is what removes the bubbles and shortens TTFT for long requests.

Performance

Kimi-K2.5-MXFP4, prefill PP4×TP1 on MI355X, PD disaggregation, OSL=128.

• Input length 128k

• Input length 256k

Each cell is throughput change / mean TTFT improvement; positive means dynamic chunking wins.

Baseline: tuned fixed chunk size

Fixed and dynamic each take their best result across 8K/16K/32K:

Workload	conc=1	conc=2	conc=4
128K fixed-length	+2.9% / +3.3%	+0.6% / +0.5%	-9.6% / -10.7%
128K variable-length	+4.3% / +4.9%	-0.3% / -0.5%	+3.3% / +2.4%
256K fixed-length	-3.6% / -4.1%	-0.2% / -0.2%	-0.2% / -0.2%
256K variable-length	-3.0% / -3.4%	+0.8% / +0.4%	≈0% / ≈0%

Against a fully tuned fixed chunk size the gains are small and workload dependent; at 256K, fixed 8K is still the best configuration in most cases.

Baseline: default --max-num-batched-tokens=16384

Workload	conc=1	conc=2	conc=4
128K fixed-length	+10.7% / +11.3%	-0.3% / +1.0%	+0.5% / +0.6%
128K variable-length	+9.7% / +10.4%	-1.1% / -1.2%	-0.7% / -0.7%
256K fixed-length	-0.8% / -0.9%	+0.1% / +0.2%	+0.1% / +0.1%
256K variable-length	-3.3% / -3.7%	≈0% / ≈0%	-0.1% / -0.1%

At the default 16K budget, dynamic chunking mainly helps 128K at conc=1, worth roughly 10% on both throughput and TTFT. At 256K or higher concurrency it is essentially flat.

Baseline: --max-num-batched-tokens=32768 (common in agentic serving)

Workload	conc=1	conc=2	conc=4
128K fixed-length	+43.7% / +34.3%	-0.1% / -0.7%	-15.5% / -19.1%
128K variable-length	+39.4% / +32.0%	-4.6% / -6.6%	+3.3% / +2.4%
256K fixed-length	+19.7% / +17.7%	+3.8% / +3.9%	+3.1% / +3.2%
256K variable-length	+17.1% / +15.6%	-0.4% / -0.8%	+0.1% / +0.1%

The large 32K initial chunk benefits most, especially at long context and low concurrency. Gains at higher concurrency are not consistent, and 128K fixed-length at conc=4 still regresses noticeably.

Overall, dynamic chunking reduces how sensitive a large initial chunk is to tuning at long context and low concurrency, rather than beating a fully tuned fixed chunk size everywhere.

Correctness

GSM8K, 3-shot, 1319 samples, with a 512-token budget to force multi-chunk prefills:

• 1107 requests actually took the multi-chunk path.
• Both fixed and dynamic completed 1319/1319 with no timeouts, crashes, or KV handoff failures.
• Flexible exact match: 92.72% → 93.56%.
• Strict exact match: 92.57% → 93.48%.

The accuracy difference is within run-to-run noise, so dynamic chunking does not regress accuracy.

[github-actions]

🏷️ CI Guide

Runs automatically on every eligible PR before approval:

• ✅ Pre Checkin: Black, Ruff, catalog schema validation, non-GPU unit tests

Heavy model tests:

• ✅ Run after the PR is approved and Pre Checkin passes
• ✅ Run immediately when an approval review is submitted
• ✅ Can be requested before approval with labels

Label	Tests
ci:full	Run all heavy PR model tests: native ATOM, vLLM, and SGLang
ci:atom	Run native ATOM model accuracy tests
ci:vllm	Run ATOM vLLM OOT model accuracy tests
ci:sglang	Run ATOM SGLang model accuracy tests

Heavy jobs are skipped when the PR is not approved and no matching ci:* label is present.
Add labels via the sidebar or gh pr edit 2007 --add-label <label>

[yitingw1]

Maybe better to use try... except like line1367 fit_chunk_overhead function?

[wanzhenchn]

I agree it should fail gracefully like fit_chunk_overhead: I’ll log a warning and return an error so the server can continue with fixed chunking instead of failing startup.

[yitingw1]

This calibration sweep has no give-up mechanism: _calibrating only flips False inside install_chunk_latency_model(), which is only reached once a fit clears ChunkLatencyCalibrator.fit()'s quality gates. If those never pass on real hardware (timing jitter, MoE routing noise, EPLB rebalancing, background KV offload), the sweep never ends — every other request stays permanently chunked at max_num_batched_tokens // 4 for the life of the process, with no way to recover short of restarting without --enable-dynamic-chunking. Suggest capping the number of failed fit attempts and giving up to fixed chunking with a warning once exceeded.

[wanzhenchn]

Thanks for catching this. You’re right: repeated samples can keep triggering failed fits indefinitely while the scheduler remains in calibration mode. I’ll cap the number of failed fit attempts; once exceeded, calibration will stop, a warning will be logged, and the scheduler will fall back to fixed chunking.