recipes/docker-images/mia-tp4/docs/chunked-prefill-memory.md

Running CHUNKED_PREFILL_SIZE above 1024 without the long-prompt memory problem

Analysis only (2026-09-11). Nothing here is implemented or tested unless it says so; the measurements come from the boots recorded in logs/profile-2026-09-10/REPORT.md §17-17b.

1. Why the chunk size is tied to memory at all

On the head (rank 0, ~4.5-5.3 GB free after boot) a prefill's memory peak is set by the DeepSeek-V4 low-ratio indexer, not by KV. For every prefill chunk, each of the ratio-1/2 source layers scores the chunk's T query rows against the whole compressed prefix L (deepseek_v4_backend._low_ratio_index_topk_dense → _dense_fp4_mqa_logits, output [T, max_seqlen_k] fp32, max_seqlen_k = the longest compressed length in the batch). Around that one buffer the same function keeps, at the same moment:

buffer	bytes per (row, key)	note
logits fp32	4	the kernel output
candidate-block copy (select_candidate_blocks, layer 20 only)	up to 4	"pads and pools a copy of its rows", capped at _TORCH_INDEXER_SCORE_BUDGET_BYTES = 1 GiB per call
candidate_masks bool	1	published by layer 20, consumed by every later low-ratio layer, lives for the whole chunk
j arange, compress_lens broadcast, top-k scratch	small

So the live peak of one chunk is roughly c × T × L with c ≈ 14 B measured (REPORT §17b: 2.97 GB at T=2048, L=100k; 2.6 GB at T=1024, L=208k; both ≈ 3.5 × 4 B × T × L, plus ~0.2-0.35 GB of touched KV pages), and the allocator's cache is emptied after each chunk by adapter/prefill_empty_cache.py, so the sum over chunks no longer accumulates.

With a 1.5 GB guard and ~0.5 GB of margin the usable budget is about 2.5-3 GB, which gives the rule that every measurement so far obeys:

T × L  ≲  2.0e8 token²            (T = chunk, L = prefix at that chunk)
2048 × 100k = 2.05e8  ok          2048 × 133k = 2.7e8  guard fired (boot 11, 128k)
1024 × 208k = 2.13e8  ok          1024 × 256k = 2.6e8  guard fired by 20 MB (boot 12)

Everything below is a way to raise T without letting T × L cross that line, or to move the line.

2. Options, ranked

2.1 Adaptive chunk: 2048 while the prefix is short, smaller as it grows (recommended)

The scheduler already asks a dynamic_chunk_sizer for the size of every continuation chunk (scheduler.py get_new_batch_prefill: if self.chunked_req is not None and self.dynamic_chunk_sizer is not None: chunked_prefill_size = sizer.predict(history_len), where history_len is the prefix in tokens). SGLang only installs that sizer for pipeline parallelism (maybe_init_dynamic_chunk_sizer, pp_size > 1), but the attribute is plain and the contract is one method: predict(history_len) -> Optional[int].

An adapter hook (sitecustomize on sglang.srt.managers.scheduler, ~30 lines) can install an object whose predict returns

chunk(L) = clamp( round_down_to_page( BUDGET / L ), CHUNK_MIN, CHUNK_MAX )
BUDGET = 2.0e8 token²  (env knob), CHUNK_MAX = 2048, CHUNK_MIN = 256, page = 256

which yields a schedule of

prefix so far	chunk	prefill rate (measured or extrapolated)
0 - 98k	2048	2.0-2.2k tok/s (measured)
98k - 195k	1024	1.5-1.7k tok/s (measured)
195k - 390k	512	~1.1-1.3k tok/s (extrapolated; boot 13 was stopped before measuring)
390k - 780k	256	~0.8k tok/s (extrapolated)

Consequences: every prompt under ~100k tokens prefills at the full 2048 speed (today all of them run at 1024, 15% slower); a 200k prompt takes about the same as today; 256k-500k prompts become possible at all (est. 500k ≈ 7 min, vs. impossible), and CONTEXT_LENGTH could go back to 512000 with the same memory guarantee. Decode is untouched (the chunk size never affects decode steps).

Edge to handle: the first chunk of a request is sized statically (the sizer is consulted only for chunked_req, i.e. continuation chunks). A brand-new request has L ≈ chunk, fine; a request that lands on a long cached prefix (radix hit of, say, 150k tokens) gets its first chunk at 2048 with L = 150k, over budget. Two fixes: clamp inside PrefillAdder.add_one_req (it knows req.extend_input_len and the prefix length), or set the static size to 1024 and let the sizer raise it to 2048 when L is small (the sizer's return value overrides the static size in both directions). The second needs no adder hook.

Cost: none on decode; on prefill only the chunks beyond 98k get smaller than today's 1024 beyond 195k. Validation: one boot, ramp 64k / 128k / 200k / 256k / 400k / 500k with scripts/verify/memguard.py, check prefill tok/s per band and that lowest MemAvailable stays above 2 GB at every size.

2.2 Lower the candidate-selection copy budget (cheap, moves the line)

_TORCH_INDEXER_SCORE_BUDGET_BYTES = 1 << 30 bounds the padded copy that layer 20 makes for block selection; at large T × L that copy is a full extra [T, L] fp32 (0.8 GB at 2048 × 100k). It is a module constant, so the adapter can set it to 256 MiB at import (module._TORCH_INDEXER_SCORE_BUDGET_BYTES = 256 << 20); the loop then processes fewer rows per step with identical results. Expected: c drops from ~14 B to ~11 B per (row, key), i.e. the budget rises from ~2.0e8 to ~2.6e8 token² and 2048-token chunks reach ~130k prefix instead of ~100k. Free to try; combine with 2.1 by raising BUDGET.

2.3 Block the indexer along the key axis and merge top-k (engine change, memory-flat)

The exact answer: score the T rows against key blocks of B_k (e.g. 32k), keep the per-block top-512 (values + indices), then take the top-512 of the concatenated candidates. topk(union) == topk(∪ per-block topk) holds exactly, so outputs do not change. Memory becomes c × T × B_k, independent of the prefix: 2048-token chunks (or 4096) at any context. The candidate-source layer (block pooling over 8-key blocks, top-2048 blocks) is decomposable the same way, and the consumer layers' masked_fill_ of non-candidates can be applied per block. Touch points: _low_ratio_index_topk_dense, _publish_or_consume_candidates, topk_transform_ragged_v2 (needs per-block offsets or a second pass), _mask_topk_scores. Effort: a day or two in a 180 KB file plus a correctness diff against the unblocked path on the 32k sweep. This is the upstream-quality fix; 2.1 is what to run in the meantime.

2.4 Context-parallel prefill over the three ranks (exists in the engine, untested here)

The engine carries a DSA prefill context-parallel mode (--enable-dsa-prefill-context-parallel, dsa_prefill_cp_mode round-robin-split, attn_cp_size; backend paths _forward_low_ratio_sources_cp, _late_layer_tail_cp_layout). It splits the query rows of a prefill across ranks, so each rank scores T / cp_size rows: at cp_size = 3 the per-rank buffer is a third, i.e. the same T × L line moves to 6e8 and 2048-token chunks reach ~300k. Unknowns on this stack: whether CP composes with TP3 padding (rank 2's all-padding attention shard), the late-layer tail, DSpark's per-chunk hidden-state injection, and whether the gathers cost more than the GEMMs save. Worth one boot with the flags and the 32k sweep for correctness before any ramp; not something to enable blind.

2.5 More headroom on the head (small, additive)

• The KV pool does not help: its pages are touched only when used, so shrinking MAX_TOTAL_TOKENS frees nothing at rest and costs concurrency.
• NCCL buffers are already trimmed (139 MB pinned).
• The tokenizer, detokenizer and HTTP server (~1.3 GB of rank-0-only RSS) cannot move to a worker in this SGLang build without a proxy layer; that is the one real reserve.
• The docker healthcheck process (~0.6 GB RSS each time it ran the engine's imports) is gone since the -S fix; it used to bite exactly during long prefills.
• Loosening the 1.5 GB guard is not a lever: the head survived at 1.0-1.3 GB in tests, but the failure past zero is a host wedge, and 0.5 GB of margin is thin at 1-second sampling.

2.6 What does not work

• expandable_segments:True (would coalesce the fragmentation) returns NaN logits on every >64-token prefill on this stack (REPORT §17). Banned.
• garbage_collection_threshold is a no-op: SGLang never calls set_per_process_memory_fraction, which that option requires.
• A static chunk of 2048 with the empty-cache hook alone: 100k prompts fit, 128k do not (boot 11).
• Emptying the cache more often than once per chunk: the peak is inside one layer's forward, not between layers.

3. Recommendation

Implement 2.1 with 2.2 folded in (two adapter hooks, no engine edit, no rebuild of anything but the image layer that copies adapter/), set CHUNKED_PREFILL_SIZE=2048 as the maximum and let the sizer descend, restore CONTEXT_LENGTH=512000, and verify with the ramp above. Expected result: 2048-speed prefill for all prompts up to ~100-130k, long prompts to 500k without touching the guard, decode unchanged. Keep 2.3 as the follow-up that removes the prefix dependence for good.