Repository navigation
Qwen3.8-Flash-Next: GDN prefill OOM on a long prompt kills the backend (prefill transient not budgeted) #3
Description
Environment: RTX 5080 16 GB (sm_120), AMD Ryzen 9 9950X, 123 GB RAM, Ubuntu 24.04 (kernel 7.1.9-xanmod), driver 610.57.04, CUDA 13.3, torch 2.11.0+cu130, triton 3.6.0, FreeToken 0.1.2 editable. Checkpoint: nvidia/Qwen3.8-Flash-Next-NVFP4.
Commit: main e0886cc (before any merges), bf16 KV.
Command:
ft serve --model ~/models/Qwen3.8-Flash-Next-NVFP4 --max-running-requests 1 --cuda-graph-max-bs 1 --moe-cache-size 512 --disable-moe-prefill-overlap --memory-ratio 0.95 [--max-prefill-length 4096]
Startup log: Allocating 97984 tokens for KV cache, K + V = 2.31 GiB, Free GPU memory after capturing CUDA graphs: 0.67 GiB.
What happens: a single ~94.5K-token chat request crashes the scheduler in the first prefill chunk:
File ".../models/qwen4_exp/gdn.py", line 185, in forward
result = gdn_prefill_chunk_fla(
File ".../models/qwen3_5_moe/gdn_kernels.py", line 33, in gdn_prefill_chunk_fla
torch.OutOfMemoryError: CUDA out of memory. Tried to allocate 32.00 MiB. GPU 0 has a total capacity of 15.51 GiB of which 41.06 MiB is free.
Backend supervisor: backend worker freetoken-TP0-scheduler exited
Backend worker is gone and cannot be restarted; stopping the API server
With the default --max-prefill-length 8192 it fails as above. With 4096 it fails the same way ("Tried to allocate 48.00 MiB ... 63.06 MiB is free").
Measured prefill transient (peak minus idle, nvidia-smi sampled every 200 ms, same flags at --memory-ratio 0.85):
--max-prefill-length |
peak over idle |
|---|---|
| 8192 | +1428 MiB |
| 2048 | +392 MiB |
Two problems:
- The cache budget (MoE + KV at
memory_ratio) doesn't reserve the prefill activation transient implied by--max-prefill-length.(1 - memory_ratio)is the only headroom, so a config that boots and captures graphs can still OOM on its first long prompt. - The OOM is fatal: one request takes down the backend permanently, and the API server exits instead of failing that request.
Workaround used: --max-prefill-length 2048 at bf16 (a 94.5K prompt then completes, peak +392 MiB), or nvfp4 KV (FlashML-org#408) with 4096 (peak +964 MiB with 1.09 GiB free).
Related
- Upstream: Unhandled CUDA OOM in prefill expert workspace kills the engine; server keeps accepting requests that never return (root cause behind #20-style headless state; likely also #72) FlashML-org/FreeToken#110 (unhandled CUDA OOM in prefill kills the engine), --moe-cache-auto leaves no room for prefill: boots and serves decode, then the first 8k prefill OOMs and kills the worker (NVFP4 dense) FlashML-org/FreeToken#401 (
--moe-cache-autoleaves no room for prefill and the first 8k prefill kills the worker), Backend worker is gone and cannot be restarted; stopping the API server FlashML-org/FreeToken#20 ("Backend worker is gone and cannot be restarted"). - Workaround used on this card: nvfp4 KV from feat(kvcache): add nvfp4 kv quantization FlashML-org/FreeToken#408 (fork PR 5080deploy 1/9: merge upstream #408, nvfp4 KV cache quantization #9), and later FP8 dense from feat(qwen4_exp): load-time per-tensor FP8 dense projections (W8A8 via _scaled_mm) FlashML-org/FreeToken#389 (fork PR 5080deploy 9/9: merge upstream #389 (+#385), load-time FP8 dense projections #17), which leave enough headroom for 4096- and then 8192-token chunks.
Activity
gdevenyi commentedon Sep 15, 2026
No PR for this yet: what a safe fix needs
I looked at making a prefill OOM fail only the offending request (problem 2) on
5080deploy@ 2d56708 and stopped short of a PR. Catchingtorch.OutOfMemoryErroraroundScheduler._forwardis the easy part. Releasing the request without corrupting later ones is not:OffloadMoeCache.ensure_experts(python/freetoken/moe/offload_cache.py) runs the device kernel that rewrites the slot map beforecopy_missingcopies the expert weights in. An OOM between the two leaves slots claiming experts whose weights never arrived. Later requests would get wrong outputs with no error. Recovery must callmoe_offload_cache.reset(), which cold-starts the cache and costs the next requests their hit rate._free_req_resourcescallscache_manager.cache_req(req, finished=True)._cache_req_hybriddonates the live GDN slot, and any pending ping-pong snapshot, into the radix tree (python/freetoken/scheduler/cache.py). After a mid-forward OOM the live state is advanced for the layers before the failure and not after, so the donated snapshot is garbage that a later prefix hit would COW-restore.[old_handle.cached_len, page_ceil(cached_len + extend))and every GDN slot, insert nothing.prefill_manager.pending_list: a chunked request is already re-queued viachunked_listwhen the batch is scheduled.ErrorReplyMsg._execute_pending_rebuildmakes for rebuild failures).That is a few hundred lines across the scheduler, all cache managers and the MoE cache, with a silent-corruption failure mode if any piece is missed. It needs a design agreed first, not a fix PR.
Smaller options that would help problem 1 (a config boots, then dies on its first long prompt), not implemented:
--max-prefill-length, as_warmup_prefilldoes (python/freetoken/engine/engine.py) but at the maximum chunk length, and turn an OOM there into a startup error that names--max-prefill-length/--memory-ratio. This costs a few seconds of startup (~7 s for 8192 tokens here) and catches the chunk-length-driven transient measured above. It won't catch transients that grow with context position.plan_cache_budget(--moe-cache-auto fails only after the expert-bank load, with a 1024-slot floor that ignores the byte budget #4 / fix(engine): let moe-cache-auto drop prefill overlap when its slot floor does not fit #20 touch the same planner). That needs per-model activation accounting.Upstream items for the same failure: FlashML-org#110, FlashML-org#401, FlashML-org#20.
The current deployment avoids it with FP8 dense + nvfp4 KV at 8192-token chunks (#17): prefill peak 15,124 of 16,303 MiB.
Merge pull request #3 from naerymdan/perf/kv-fp8-read-path