fix(dsv4): honor an explicit --max-prefill-length instead of silently forcing single-pass prefill

[avlp12]

Symptom

Cold prompts past ~17K tokens kill the DSV4 backend mid-prefill:

RuntimeError: CUDA driver error: device not ready
ERROR Backend supervisor: backend worker freetoken-TP0-scheduler exited
ERROR Backend worker is gone and cannot be restarted; stopping the API server

On WSL2, dmesg logs misc dxg: dxgk: dxgkio_make_resident: Ioctl failed: -12 (ENOMEM) at the moment of every crash. Same failure signature as #87 (that report is the Qwen3.6 GDN path; this one is DSV4, and the mechanism below suggests they are the same class of bug).

Environment: RTX 5090 32GB, WSL2 Ubuntu 26.04 (kernel 6.18.33), driver 610.88, CUDA 13.3, FreeToken 0.1.2, DeepSeek-V4-Flash (ds_fp4), --memory-ratio 0.9 --moe-cache-size 512 --num-tokens 319872 --max-seq-len-override 319872.

What we ruled out first

Experiment	Result
~8K cold prompt	OK (repeatedly)
16K prompt with 7.9K prefix-cached	OK
20K / 30K cold prompt	crash, 7/7 across all configs below
--max-prefill-length 4096 + --disable-moe-prefill-overlap	crash (flag has no effect — see below)
--max-prefill-length 32768	crash
PYTORCH_ALLOC_CONF=expandable_segments:False	same failure as a permanent scheduler hang instead of a crash
--memory-ratio 0.85 / 0.9 / 0.95	crash at all values

Diagnosis

With CUDA_LAUNCH_BLOCKING=1 the failure stays at a plain allocation — sparse_attn_paged's o = torch.empty_like(q) — so it is a true out-of-memory, not an async kernel fault. Instrumentation immediately before that line, for a "20K" prompt:

[sparse_attn DIAG] q=(1, 22896, 64, 512) idx=(1, 22896, 128) free=0.00GiB

The entire 22.9K-token prompt is in one ragged forward, and free VRAM is zero. q alone is 1.5 GiB (~64KB/token at 64 heads x 512 head_dim); together with k/v/hidden/indexer transients, single-pass prefill sweeps away the post-init headroom (3.4 GiB on this box) somewhere past ~17K tokens. The WDDM residency failure (-12) and "device not ready" are downstream symptoms.

Why isn't it chunked? _adjust_dsv4_config unconditionally raises max_extend_tokens to max_seq_len, so the CLI flag documented as "Chunk Prefill maximum chunk size" is silently ignored on DSV4. The only remaining chunker is the window-pool cap (30,464 tokens here) — far above what the activation headroom can carry. Admission-side instrumentation confirmed:

[adder DIAG] chunk_size=22911 remain=22911 cached=0 budget=30464 chunked=False

Fix

Track whether --max-prefill-length was explicitly given (max_extend_tokens_explicit) and let _adjust_dsv4_config widen the chunk only for the default. Default behavior is unchanged (single-pass prefill, per the existing design comment); an explicit flag is now honored, which is the only lever a user has when prefill transients exceed free VRAM.

Validation (same box, fix applied, --max-prefill-length 8192)

• 22,925-token cold prompt → chunks 8192 / 8192 / 6541, completes in 15.5s, output correct:

  [adder DIAG] chunk_size=8192 remain=22925 cached=0 budget=8192 chunked=True
  [adder DIAG] chunk_size=8192 remain=14733 cached=8192 budget=8192 chunked=True
  [adder DIAG] chunk_size=6541 remain=6541 cached=16384 budget=8192 chunked=False
  

• 34,942-token cold prompt → chunks 8192 x4 + 2174, completes in 21.3s, output correct.
• No new dxgkio_make_resident failures in dmesg; decode throughput unchanged (~16 tok/s at a 320K-token KV).

Without the fix the identical requests crashed the backend 7/7 times.

A follow-up worth considering (out of scope here): deriving a default prefill chunk from actual free VRAM after init, so the out-of-the-box config cannot outgrow its own activation headroom — that would also close the default-config case and likely #87's path.

🤖 Generated with Claude Code