Skip to content

fix(engine): account for attention workspace in cache-pool sizing #303

Description

@ralfkuh-lab

Summary

On a 6 GB card, ft serve with the auto-selected FlashInfer attention
backend can fail at startup with a CUDA OOM for FlashInfer's 256 MiB
float_workspace_buffer (attention/fi.py): the cache-budget planner
(engine/cache_budget.py) sizes the MoE/KV pools from --memory-ratio
without explicitly charging the attention backend's known fixed workspace,
so at higher ratios the workspace no longer fits into what is left. The
failure is clean (worker exits, API server shuts down), but the knob that
caused it is not obvious from the error.

Related: #110 discusses request-time OOM handling; this report is only
about the startup-time budget. At lower ratio settings that start
successfully, the same headroom must still cover warmup/graph-capture
allocations.

Environment

  • FreeToken 0.1.2, source install; upstream base 4b94bdc, run at local HEAD bd7482d (local patches touch the server layer only: an opt-in distributed-port override, unset here, and a thinking override, not passed here — neither is involved in engine memory planning). torch 2.11.0+cu130, CUDA toolkit 13.0.2, driver 595.84
  • RTX 3060 Laptop 6 GB (sm_86), Ryzen 9 5900HX, 30 GB RAM, Linux Mint 22.3
  • Checkpoint: nvidia/Qwen3.6-35B-A3B-NVFP4

Reproduce

ft serve --model /home/user/models/Qwen3.6-35B-A3B-NVFP4 --memory-ratio 0.95 --kv-reserve-tokens 4096
ERROR Backend supervisor: OutOfMemoryError: CUDA out of memory. Tried to allocate
256.00 MiB. GPU 0 has a total capacity of 5.67 GiB of which 238.12 MiB is free. …
  File ".../freetoken/attention/fi.py", line 114, in __init__
    self.float_workspace_buffer = torch.empty(

With --memory-ratio 0.92 and explicit pool sizing (the reference
configuration below) the same checkpoint starts and serves fine on this
card.

Full log, --memory-ratio 0.95 run
[2026-08-31|00:41:26] INFO     Parsed arguments:
ServerArgs(model_path='/home/user/models/Qwen3.6-35B-A3B-NVFP4', tp_info=DistributedInfo(rank=0, size=1), dtype=torch.bfloat16, max_running_req=4, attention_backend='auto', moe_backend='auto', nvfp4_backend='triton', expert_load='auto', moe_cache_size=0, moe_cache_rate=None, moe_cache_auto=False, kv_reserve_tokens=4096, moe_cache_policy='lru', moe_prefill_overlap=True, moe_prefill_hit_d2d=False, moe_collect_stats=False, moe_cpu_threads=0, moe_cpu_layers=None, moe_hybrid_max_fetch=-1, cuda_graph_bs=None, cuda_graph_max_bs=None, page_size=1, memory_ratio=0.95, linear_state_cache_ratio=2.0, swa_full_tokens_ratio=0.2, swa_num_pages_override=None, distributed_timeout=60.0, use_dummy_weight=False, use_pynccl=True, max_seq_len_override=None, num_page_override=None, num_token_override=None, max_extend_tokens=8192, cache_type='radix', offline_mode=False, decode_log_interval=40, special_token_ckpt=False, _unique_suffix='.pid=3065162', server_host='127.0.0.1', server_port=1919, num_tokenizer=0, silent_output=False, shell_mode=False, served_model_name='Qwen3.6-35B-A3B-NVFP4', tool_call_parser='qwen3_coder', reasoning_parser='qwen3', thinking='auto', sampling_defaults='model', max_output_tokens=None, enable_cache_report=False, cors_origins='tauri://localhost,http://tauri.localhost,http://localhost:1420', gpu=(), gpu_assigned=None)
[2026-08-31|00:41:26|FrontendAPI] INFO     Default sampling config (source=model): temperature=1.0, top_k=20, top_p=0.95
INFO:     Started server process [3065162]
INFO:     Waiting for application startup.
INFO:     Application startup complete.
INFO:     Uvicorn running on http://127.0.0.1:1919 (Press CTRL+C to quit)
/home/user/dev/local-ai/upstream/FreeToken/python/freetoken/engine/engine.py:1030: FutureWarning: torch.cuda._set_allocator_settings is deprecated. Use torch._C._accelerator_setAllocatorSettings instead.
  torch.cuda.memory._set_allocator_settings("expandable_segments:True")
[2026-08-31|00:41:29|core|rank=0] INFO     Enabled expandable_segments (override via PYTORCH_ALLOC_CONF)
[2026-08-31|00:41:29|core|rank=0] INFO     Auto-selected attention backend: fi
[2026-08-31|00:41:29|core|rank=0] INFO     benchbw profile recommends hybrid for 'nvfp4' experts on this GPU
[2026-08-31|00:41:29|core|rank=0] INFO     Auto-selected MoE backend: hybrid
[2026-08-31|00:41:29|core|rank=0] INFO     No MoE cache sizing flag given; defaulting to --moe-cache-auto for auto-selected backend 'hybrid'
[2026-08-31|00:41:29|core|rank=0] INFO     Resolved config: moe_backend='hybrid', attention_backend='fi', cache_type='hybrid_radix', page_size=1
[2026-08-31|00:41:29|core|rank=0] INFO     Free memory before loading model: 5.52 GiB

Loading mixed-fp8 weights:   0%|          | 0/3 [00:00<?, ?it/s]
Loading mixed-fp8 weights:  33%|███▎      | 1/3 [00:01<00:03,  1.67s/it]
Loading mixed-fp8 weights:  67%|██████▋   | 2/3 [00:02<00:00,  1.10it/s]
Loading mixed-fp8 weights: 100%|██████████| 3/3 [00:02<00:00,  1.26it/s]
Loading mixed-fp8 weights: 100%|██████████| 3/3 [00:02<00:00,  1.11it/s]
[2026-08-31|00:41:34|core|rank=0] WARNING  expert banks: low free RAM -> serial build (avoids parallel-reader OOM; override with --expert-load parallel)
[2026-08-31|00:41:34|core|rank=0] INFO     expert banks: slow path (serial build)

Loading Qwen3.5 NVFP4 experts:   0%|          | 0/3 [00:00<?, ?it/s]
Loading Qwen3.5 NVFP4 experts:  33%|███▎      | 1/3 [00:06<00:12,  6.25s/it]
Loading Qwen3.5 NVFP4 experts:  67%|██████▋   | 2/3 [00:17<00:08,  8.97s/it]
Loading Qwen3.5 NVFP4 experts: 100%|██████████| 3/3 [00:19<00:00,  5.74s/it]
Loading Qwen3.5 NVFP4 experts: 100%|██████████| 3/3 [00:19<00:00,  6.34s/it]
[2026-08-31|00:41:54|core|rank=0] INFO     --moe-cache-auto resolved moe_cache_size=610 num_pages=4163 (prefill_overlap=True)
[2026-08-31|00:41:55|core|rank=0] INFO     --moe-hybrid-max-fetch auto: fetching 20.1% of each decode step's expert misses over PCIe (benched PCIe/CPU bandwidth ratio), the rest on the CPU
[2026-08-31|00:41:55|core|rank=0] INFO     torch intra-op threads: 8 -> 1 (cores reserved for the pinned CPU MoE pool)
[2026-08-31|00:41:55|core|rank=0] INFO     CPU MoE executor ready: threads=7 (pinned to cores 0..12) isa=avx2 fmt=nvfp4 H=2048 I=512 experts=256 layers=40 top_k=8 act=silu max_tokens=4
[2026-08-31|00:41:55|core|rank=0] INFO     Allocating 4163 tokens for KV cache, K + V = 0.08 GiB
Process freetoken-TP0-scheduler:
[2026-08-31|00:41:55|FrontendAPI] ERROR    Backend supervisor: OutOfMemoryError: CUDA out of memory. Tried to allocate 256.00 MiB. GPU 0 has a total capacity of 5.67 GiB of which 238.12 MiB is free. Including non-PyTorch memory, this process has 5.38 GiB memory in use. Of the allocated memory 5.17 GiB is allocated by PyTorch, and 67.22 MiB is reserved by PyTorch but unallocated. If reserved but unallocated memory is large try setting PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True to avoid fragmentation.  See documentation for Memory Management  (https://docs.pytorch.org/docs/stable/notes/cuda.html#optimizing-memory-usage-with-pytorch-cuda-alloc-conf)
Traceback (most recent call last):
  File "/usr/lib/python3.12/multiprocessing/process.py", line 314, in _bootstrap
    self.run()
  File "/usr/lib/python3.12/multiprocessing/process.py", line 108, in run
    self._target(*self._args, **self._kwargs)
  File "/home/user/dev/local-ai/upstream/FreeToken/python/freetoken/server/launch.py", line 81, in _run_scheduler
    scheduler = Scheduler(args)
                ^^^^^^^^^^^^^^^
  File "/home/user/dev/local-ai/upstream/FreeToken/python/freetoken/scheduler/scheduler.py", line 65, in __init__
    self.engine = Engine(config)
                  ^^^^^^^^^^^^^^
  File "/home/user/dev/local-ai/upstream/FreeToken/python/freetoken/engine/engine.py", line 385, in __init__
    self.ctx.attn_backend = self.attn_backend = create_attention_backend(
                                                ^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/user/dev/local-ai/upstream/FreeToken/python/freetoken/attention/__init__.py", line 181, in create_attention_backend
    return SUPPORTED_ATTENTION_BACKENDS[backend](config)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/user/dev/local-ai/upstream/FreeToken/python/freetoken/attention/__init__.py", line 66, in create_fi_backend
    return FlashInferBackend(config)
           ^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/user/dev/local-ai/upstream/FreeToken/python/freetoken/attention/fi.py", line 114, in __init__
    self.float_workspace_buffer = torch.empty(
                                  ^^^^^^^^^^^^
torch.OutOfMemoryError: CUDA out of memory. Tried to allocate 256.00 MiB. GPU 0 has a total capacity of 5.67 GiB of which 238.12 MiB is free. Including non-PyTorch memory, this process has 5.38 GiB memory in use. Of the allocated memory 5.17 GiB is allocated by PyTorch, and 67.22 MiB is reserved by PyTorch but unallocated. If reserved but unallocated memory is large try setting PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True to avoid fragmentation.  See documentation for Memory Management  (https://docs.pytorch.org/docs/stable/notes/cuda.html#optimizing-memory-usage-with-pytorch-cuda-alloc-conf)
[2026-08-31|00:42:05|FrontendAPI] ERROR    Backend worker is gone and cannot be restarted; stopping the API server
INFO:     Shutting down
INFO:     Waiting for application shutdown.
INFO:     Application shutdown complete.
INFO:     Finished server process [3065162]
/usr/lib/python3.12/multiprocessing/resource_tracker.py:254: UserWarning: resource_tracker: There appear to be 3 leaked semaphore objects to clean up at shutdown
  warnings.warn('resource_tracker: There appear to be %d '

Suggestion

Please account for the attention backend's known fixed workspace explicitly
when sizing the cache pools, and validate startup warmup/graph capture
against the remaining runtime headroom — or fail the plan early with an
actionable message (for this case: lower memory_ratio or shrink the
pools; note the existing cache-budget assert suggests raising
memory_ratio, which would make this particular failure worse).

Stable configuration on this card, for reference:
--attn triton --memory-ratio 0.92 --max-prefill-length 2048, start with
--moe-cache-size 512, then grow via ft ctl cache --moe 940.

Transparency (home directory in paths/logs anonymised to /home/user): prepared with AI assistance; the failure and the working
configuration were reproduced and verified on the hardware above.

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions