Skip to content

perf: make the CHUNKED_PREFILL_SIZE=1024 default conditional on context > native (+15…23% prefill at 262144) - #15

Closed
sethforprivacy wants to merge 1 commit into
MiaAI-Lab:mainfrom
sethforprivacy:pr-chunk-native-default
Closed

sethforprivacy wants to merge 1 commit into
MiaAI-Lab:mainfrom
sethforprivacy:pr-chunk-native-default

Conversation

@sethforprivacy

Copy link
Copy Markdown
Contributor

The defect

Two separate mechanisms currently share one value:

CHUNKED_PREFILL_SIZE="${CHUNKED_PREFILL_SIZE:-1024}"   # unconditional default
…
if (( CONTEXT_LENGTH > YARN_NATIVE_CTX )); then
  if (( CHUNKED_PREFILL_SIZE > 1024 )); then
    warn "… froze GB10 at 300k history — clamping to 1024"

The clamp is right and this PR does not touch it. Above native context a
4096 chunk froze GB10 at 300k history; 1024 is the safe value and it should stay
enforced.

The unconditional default is the problem. CONTEXT_LENGTH=262144 is a
documented, supported configuration — it is the opt-out this repo tells you to
use (CHANGELOG: "Opt out with CONTEXT_LENGTH=262144"), and it is the only
option at all if you want bf16 or fp8 KV, since a bf16 pool at 1M now trips the
issue #8 pool guard and refuses to boot. Anyone taking that opt-out silently
inherits a chunk size chosen for a hazard that does not exist at native context,
and pays real prefill throughput for it.

start.sh's own recorded provenance already disagrees with its default —
the model card recipe it follows, quoted in the header:

#   * Model card serving recipe (validated TP=2 on GB300/B300):
#     --quantization modelopt_fp4, --page-size 64, mamba extra_buffer +
#     --mamba-track-interval 64, --chunked-prefill-size 4096,
#     --mem-fraction-static 0.80, --context-length 262144, --allow-auto-truncate.

4096 at --context-length 262144. The 1M default arrived and took the chunk
size down with it for every context.

What it costs

2× GB10 TP2, native 262,144 context, same pair, same harness, one variable
moved, medians of 2 measured reps (reps tight on both sides):

prompt chunk 1024 chunk 4096
8K 2,439 tok/s 2,799 tok/s +15%
32K 2,435 tok/s 2,823 tok/s +16%
128K 2,065 tok/s 2,548 tok/s +23%

The shape is what you would expect: the smaller chunk multiplies scheduler
steps, and the penalty grows with prompt length.

Second-order but worth knowing, since it is the same knob: a small chunk at a
long prompt is also what widens the /metrics handler's blocking window. On a
128K prefill we measured 62–77 s of scheduler churn at chunk 1024 versus ~52 s
before, which was long enough for a 10 s metrics scrape to time out and take the
backend briefly out of a router's rotation.

Why 4096 and not higher

An earlier sweep of ours at native context found 4096 / 8192 / 12288
statistically indistinguishable
— prefill within 1% above 8K, decode 250.0 vs
245.0 aggregate at c=16 against a ±8% run-to-run noise floor (six repeats,
config held fixed). So 4096 is the bottom of a measured-equivalent band, not
the maximum. Deliberately conservative: it restores the throughput, it matches
the model card, and it keeps the smallest activation footprint of the three (the
larger values also give back a little KV pool — 623,936 tokens at 8192 vs
602,176 at 12288 — which is a reason not to reach for the top of the band).

The change

  • Default CHUNKED_PREFILL_SIZE per context: 4096 when
    CONTEXT_LENGTH <= YARN_NATIVE_CTX, 1024 above it. The 1M default path is
    bit-for-bit unchanged.
  • An explicit CHUNKED_PREFILL_SIZE (shell or .env) still wins, and is still
    clamped above native context with the existing warning.
  • .env.example, the --help header block, and the CHANGELOG line that
    currently reads "default 1024" updated to match.

Validation

bash -n start.sh clean. Full truth table, driven through the real assignment
and the real clamp block out of the patched script:

CONTEXT_LENGTH=1048576  CHUNKED_PREFILL_SIZE=unset  -> chunk=1024          (unchanged)
CONTEXT_LENGTH=262144   CHUNKED_PREFILL_SIZE=unset  -> chunk=4096          (the fix)
CONTEXT_LENGTH=131072   CHUNKED_PREFILL_SIZE=unset  -> chunk=4096
CONTEXT_LENGTH=1048576  CHUNKED_PREFILL_SIZE=8192   -> [WARN] … clamping to 1024
                                                       chunk=1024          (clamp intact)
CONTEXT_LENGTH=262144   CHUNKED_PREFILL_SIZE=8192   -> chunk=8192          (override honoured)
CONTEXT_LENGTH=262144   CHUNKED_PREFILL_SIZE=1024   -> chunk=1024          (override honoured)
unset                   unset                       -> ctx=1048576 chunk=1024

Beyond the A/B: the value this PR proposes as the native-context default is what
we already run in production. One of our pairs serves continuously at
CONTEXT_LENGTH=262144 + CHUNKED_PREFILL_SIZE=4096, carried as a local .env
override against this recipe — /get_server_info on it right now reports
context_length=262144, chunked_prefill_size=4096,
max_total_num_tokens=1907264, kv_cache_dtype=nvfp4. It is stable there; the
override is the only reason it is not at 1024.

… above it

The >native clamp is right and is untouched: chunk 4096 at 300k history
froze GB10. The unconditional 1024 *default* is the problem. At native
262144 — the documented opt-out, and the only option if you want bf16/fp8
KV, since a bf16 pool at 1M now trips the issue MiaAI-Lab#8 guard — it costs real
prefill throughput for a hazard that is not present.

Measured on 2x GB10 TP2 at native context, one variable, medians of 2:

  prompt   chunk 1024   chunk 4096
    8K     2,439 t/s    2,799 t/s   +15%
   32K     2,435 t/s    2,823 t/s   +16%
  128K     2,065 t/s    2,548 t/s   +23%

4096 is also what the model card recipe quoted in this script's own header
uses at --context-length 262144, and it is the bottom of a band we
measured as indistinguishable (4096/8192/12288 within a +/-8% noise
floor), so this is not a reach for the maximum.

An explicit CHUNKED_PREFILL_SIZE still wins and is still clamped above
native context.
@sethforprivacy

Copy link
Copy Markdown
Contributor Author

Please do not merge this. I'm withdrawing the recommendation — my own follow-up measurement contradicts it.

The throughput numbers in the description are real and I stand behind them. What I failed to do before filing was sample output correctness at either setting. Having now done that, CHUNKED_PREFILL_SIZE=4096 materially raises the rate of the token-0 ! loop.

Same probe, same 12 seeds, same host, POST /flush_cache before every request, prompts at 16,384 tokens — an ordinary length, not a long-context edge case:

chunk ! loops rate
4096 6 / 12 50%
1024 (your default) 1 / 12 8%

Fisher's exact, one-tailed: p = 0.0343.

So the default you already ship is the safer one, and this PR would have raised the loop rate for every user of the recipe in exchange for prefill throughput. I'd rather say that plainly than leave it sitting in your queue looking reasonable.

Two things I want to be careful not to overclaim:

  • 1024 is not a fix. It's 1/12, not 0/12. Chunk amplifies this fault; it doesn't cause it. The underlying bug is still present at your default and still unexplained.
  • n=12 per arm, and this fault is bursty — we measured a run of 14 consecutive failures in 80 requests, with a 200,000-shuffle permutation test giving p < 5e-06 against that being chance. So a 12-request window can land badly. The direction here is significant; the magnitudes are not pinned, and a burst-aware A/B with several windows per arm would be needed to trust them.

There's also a knock-on for anyone reading our earlier reports: because I changed this knob midway through our !-loop investigation, chunk size is a confound across that whole dataset. Conclusions we drew about the loop being gated by prompt depth don't survive it — a later soak found 16,384 failing 28/40 and 65,536 failing 30/40, statistically identical. I'll re-derive those tables with chunk held fixed before offering you anything about the loop itself.

Happy to close this myself if you'd prefer — flagging rather than deleting so the measurement and the retraction stay on the record together. If any part of it is still useful to you, the honest residual is narrow: at native context max_prefill_tokens returns to 16384 while chunk stays at 1024, and that asymmetry might be worth documenting even though changing the chunk default is not.

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