Skip to content

Live cache rebuild grows MoE before shrinking KV/mamba, so a resize that passes the fit check can OOM #643

Open

Description

@gdevenyi

Engine.rebuild_runtime_cache (engine/engine.py:1006-1021) always resizes MoE, then KV, then mamba, and each pool frees only its own old tensors. A live resize that moves memory from KV or mamba into MoE therefore allocates the larger MoE cache while the old KV/mamba pools are still resident. The fit check (validate_rebuild, engine.py:972) prices only the target geometry, so on a nearly full card a resize that passes it can OOM and roll back. (A separate rollback failure after such an OOM is #526.)

Reproduce: serve an offloaded-MoE model near full, then ft ctl cache --moe <larger> --kv <smaller> with a target total at or below the current one. A CPU check that drives rebuild_runtime_cache with fake pools under a shared cap (current moe=40, kv=50; target moe=60, kv=30; cap 100) fails on the MoE allocation while kv=50 is still allocated.

Expected: apply the shrinking resizes before the growing ones, so any geometry that passes the fit check is reachable live.

Environment: main 07813248 (source build), RTX 5080 16 GB (sm_120), driver 610.57.04, Ubuntu 24.04, Python 3.14, torch 2.11.0+cu130, triton 3.6.0; checkpoint nvidia/Qwen3.8-Flash-Next-NVFP4.

Implementation (on my fork, with a test): 50caea2. Fork PR: gdevenyi#22. #650 now fixes this upstream.

Activity

  1. void-mckenzie commented on Oct 8, 2026

    @void-mckenzie
    Contributor

    Hi @gdevenyi Fix is up in #650: the rebuild now applies shrinking resizes before growing ones, so the transient peak stays within max(current, target) and any plan that passes the fit check completes. Reproduced your scenario on a 5090 with Qwen3.8 (offload MoE, 1.97 GiB free): on main the MoE grow OOMs with the old KV resident and the rollback latches failed; on the branch the same ft ctl cache rebuild completes and serving continues. Please let me know if this fixes your issue on your card!

  2. gdevenyi commented on Oct 13, 2026

    @gdevenyi
    ContributorAuthor

    Included in the 1080 Ti deployment branch, rebuilt today on main 9c635a1 (torch 2.14.1+cu126, Triton 3.8, GTX 1080 Ti sm_61, Qwen3.8-Flash-Next-NVFP4 at 262K); tests/scheduler/test_cache_rebuild.py passes on sm_61. The deployment never rebuilds its caches live, so this is unit-test coverage only.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

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