Skip to content

feat: Automatic KV/MoE Laddering for decode speed vs context-length trade off ( upto 33% faster decode ) - #300

Open
aswinkumar1999 wants to merge 4 commits into
FlashML-org:mainfrom
aswinkumar1999:main
Open

aswinkumar1999 wants to merge 4 commits into
FlashML-org:mainfrom
aswinkumar1999:main

Conversation

@aswinkumar1999

@aswinkumar1999 aswinkumar1999 commented Aug 30, 2026 •

Copy link
Copy Markdown

Add an automatic KV/MoE cache ladder

Summary

This adds --enable-kv-ladder, an opt-in serving mode that grows the KV cache when an
incoming request could reach its current capacity. Each growth trades GPU-resident MoE cache
slots for KV pages while staying inside the engine's existing measured cache budget.

It also adds --ladder-step-size, which defaults to 32,768 tokens. With the ladder enabled,
startup KV is at least twice the configured step size and grows by that step until reaching the
model's context limit.

ft serve --model RadixArk/Qwen3.8-Flash-Next-NVFP4 \
  --moe-backend offload --moe-cache-auto \
  --max-running-requests 1 --memory-ratio 0.93 \
  --enable-kv-ladder

With the default step and this model, the resulting ladder is:

65,536 -> 98,304 -> 131,072 -> 163,840 -> 196,608 -> 229,376 -> 262,144 ( or Model Maximum )

Motivation

Preallocating KV for the full model context reduces the number of MoE expert slots that can
remain resident on the GPU, even when most requests use much shorter contexts. Reserving too
little KV preserves decode throughput but prevents long sessions from reaching the model's full
context.

The ladder keeps more experts resident for common shorter requests and spends that memory on KV
only when a request actually needs the next context rung.

Implementation

  • Compute the possible request length as input tokens plus the request's maximum output tokens.
  • If that length reaches the current KV bound, hold the request until the scheduler is idle.
  • Select the smallest configured rung that fits, always moving at least one rung when a request
    could exactly fill the current capacity.
  • Derive the replacement MoE slot count from measured KV-page and expert-slot costs.
  • Retain at least the model's minimum expert floor.
  • Reuse the existing rollback-safe runtime cache-rebuild path.
  • Admit the held request after a successful rebuild while preserving queued request order.
  • Keep serving with the old cache geometry if a rebuild is rejected or fails.
  • Cap growth at the checkpoint's real context limit.

The initial version is intentionally restricted to TP=1, --max-running-requests 1,
--moe-cache-auto, and an offload-family MoE backend.

Measured results

Environment:

  • NVIDIA RTX 5090, 32 GiB
  • RadixArk/Qwen3.8-Flash-Next-NVFP4
  • offloaded MoE with automatic LRU cache sizing
  • one running request
  • memory ratio 0.93
  • one warm-up request at each geometry
KV capacity MoE slots 128-token decode 1,024-token decode 4,096-token decode Long-decode gain vs fixed 256K
65,536 6,425 69.3 tok/s 81.8 tok/s 83.8 tok/s +33.1%
98,304 6,125 69.0 tok/s 76.7 tok/s 80.9 tok/s +28.5%
131,072 5,826 68.3 tok/s 73.5 tok/s 77.9 tok/s +23.7%
163,840 5,526 67.8 tok/s 69.1 tok/s 74.4 tok/s +18.1%
196,608 5,226 66.1 tok/s 65.0 tok/s 70.7 tok/s +12.3%
229,376 4,927 64.2 tok/s 61.7 tok/s 67.1 tok/s +6.5%
262,144 4,627 61.1 tok/s 57.1 tok/s 63.0 tok/s baseline

Short-output results include proportionally more fixed launch and streaming overhead. The
4,096-token column is the more representative sustained-decode comparison.

Rebuild latency

  • Warm/manual rebuilds between previously exercised geometries completed end to end in
    approximately 0.67-0.99 seconds.
  • A first automatic visit to a new geometry required approximately 3-4 seconds, primarily
    for CUDA graph capture.

Long-context request TTFT should not be interpreted as rebuild latency. For example, the final
196,608-token cold prompt reached first token in 125.8 seconds, but almost all of that time was
prompt prefill rather than the cache rebuild.

Validation

  • Exercised every automatic rung from 65,536 through 262,144 tokens on the target model.
  • Every automatic rebuild returned status=ok; no capacity rejection occurred.
  • Verified short, medium, and long forced decode lengths at every geometry.
  • Focused scheduler/cache/parser test run: 67 passed locally.
  • Python byte-compilation and git diff --check passed.

Known limitation

This implementation is a grow-only high-watermark ladder for single-decode use cases. A new, unrelated prompt with no
prefix-cache reuse does not yet shrink KV and restore MoE slots automatically. The server sees
stateless prompts and has no direct session-ended signal; shrinking immediately on idle would
also erase useful KV between consecutive turns of the same conversation.

A follow-up can inspect the next tokenized request's reusable-prefix hit and, when reuse is
negligible, rebuild to the smallest rung that fits that request. That should include hysteresis
to prevent rebuild thrashing.

Added arguments for enabling KV ladder and setting ladder step size.
Add KV ladder policy implementation and management
This module implements a policy for managing key-value (KV) storage growth at request boundaries by optimizing the use of MoE (Mixture of Experts) cache slots. It includes classes for capacity errors and planning growth strategies based on current and target token requirements.
Added new CLI options for KV ladder and MoE configurations.
@aswinkumar1999 aswinkumar1999 changed the title feat: Automatic KV/MoE Laddering for decode speed vs context-length trade off! feat: Automatic KV/MoE Laddering for decode speed vs context-length trade off Aug 30, 2026
@aswinkumar1999 aswinkumar1999 changed the title feat: Automatic KV/MoE Laddering for decode speed vs context-length trade off feat: Automatic KV/MoE Laddering for decode speed vs context-length trade off ( upto 33% faster decode ) Aug 30, 2026
jomcgi added a commit to jomcgi/FreeToken that referenced this pull request Sep 3, 2026
Port of upstream FreeToken FlashML-org#300 with the FlashML-org#340 dummy-page floor fix.
With --moe-cache-auto the KV pool starts at a floor (at least two
growth steps, clamped to the cap) and grows in 32,768-token steps at
request boundaries by taking expert slots back through the runtime
cache rebuild, which drains hot adaptation and preserves protected
slots where it can; an explicit --num-pages becomes the growth cap.
One eligibility decision (flag, cache-auto, max-running 1, TP 1,
runtime rebuild support, no DSV4, live MoE cache) is shared by the
engine and the scheduler; when ineligible the pool is sized exactly as
before. Planning uses max_tokens clamped to the cap, parked requests
keep priority order with a starvation bound, are visible in the queue
stats, and growth is one-way for the process lifetime. The dummy page
is charged exactly once. On node-4 the startup pool is 65,536 tokens
at 3,753 expert slots, a 2k prompt with the default 32k output budget
needs no growth, and a long context grows to the 100,352 cap at about
3,590 slots.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01A88MCbnLtwsFSHmqwuJezY
MT-z added a commit to MT-z/FreeToken that referenced this pull request Sep 16, 2026
…easured on this box)

gdevenyi's FlashML-org#391, open since 2026-09-05. OpenAI's `developer` role is the current spelling
of `system`; a template that does not spell it raises, and the request dies at
`could not encode request: Unexpected message role.` The PR maps it to `system` unless the
template handles the role itself.

Reproduced before taking it, on this branch, Ornith-1.5-35B-A3B-NVFP4:

    role=developer   400  could not encode request: Unexpected message role.
    role=system      200

and after the merge both answer 200 with the same text ('Paris is the capital of France.').
`/v1/responses` already had its own mapping for codex (`responses_api.py:258`); this is the
chat-completions path, which did not.

Conflicted twice in `tokenizer/tokenize.py`, both "two things added at the same spot" --
the vision bundle on this side, `_map_developer_role` on the PR's -- so both are kept. The
Anthropic route still answers 200.

Full suite: 11 failed, 1861 passed, 60 skipped -- the same 11, plus the test the PR adds.

Not applicable on this box, checked and set aside in the same pass:

  FlashML-org#275  the GDN in_proj per-tensor FP8 shape it fixes IS Ornith's (config group_0 lists
        linear_attn.in_proj_qkv / in_proj_z as FP8), but the PR is the qwen3_5 *dense*
        loader and the MoE path already handles it -- Ornith loads and serves today
  FlashML-org#294  needs experts NVFP4 with the shared expert per-tensor FP8; Ornith's group_1 puts
        the shared expert in NVFP4 alongside the routed experts, so the shape does not match
  FlashML-org#296  compressed-tensors (llm-compressor); Ornith is quant_method=modelopt
  FlashML-org#300  the KV ladder is `--max-running-requests 1` only, and a Claude Code request's
        input + max_output_tokens lands on a high rung on the first call. What it automates
        IS real here though -- see FINDINGS: KV 98,304 buys +17% decode at 87k of context,
        because slots/layer crosses experts_for_90pct

Assisted-by: Claude Opus 5
MT-z added a commit to MT-z/FreeToken that referenced this pull request Sep 16, 2026
… tests still drive it

FlashML-org#454 reads five config attributes directly. Every one of them exists on a real
EngineConfig/ServerArgs, so FlashML-org#454's own suite never notices -- but this branch drives
the same code from tests that build the config as a SimpleNamespace, and those stubs
predate the fields. Five failures, all AttributeError, none of them about multimodal:

  kvcache/__init__.py    model_config.model_is_mrope   test_qsa_pool_fp8 (factory)
  kvcache/qsa_pool.py    same, inside kv_cost          test_qsa_pool_fp8 (kv_cost)
  server/stats.py        config.served_modalities      test_effective_context_length
  server/generation.py   state.config.mm.max_pixels    test_logprobs_api, test_openai_extras
  server/generation.py   the second call site

This is the shape FlashML-org#300 flagged for the scheduler hooks and that this branch already
carries a fix for: a hook reached from a stub must not require the whole config
object. derive_model_card is additionally a metadata route, which must never raise.

Only these five. The direct reads in engine.py (nine) and mm/media.py (three) are on
paths that see a real config and nothing else -- turning those defensive would hide a
genuinely missing field rather than tolerate a deliberately partial stub.

Also threads FlashML-org#411's default_max_tokens from _resolve_sampling into resolve_sampling.
The merge kept this branch's richer _resolve_sampling body, which did not carry the
argument, so /v1/chat/completions and /v1/completions both fell back to the 32k
built-in and ignored --max-output-tokens. FlashML-org#411's own test caught it.

Full suite against try/all before the vision swap: 8 failed / 2004 passed vs
8 failed / 1996 passed, same failure set -- the three fp8 ones are this branch's
standing failures from the FlashML-org#354 take, not new.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
MT-z added a commit to MT-z/FreeToken that referenced this pull request Sep 16, 2026
merge origin/main が出した新規失敗 3 件(standing 11 件とは別)。

1. attention/triton.py: 上流の backend 試験は fp8 scale を持たない KV cache の fake を
   渡すので、k_scale/v_scale を無条件に読むと落ちる。getattr 越しにした
   (FlashML-org#300 が scheduler hook で指摘したのと同じ stub の形)。
   accessor が無い = fp8 pool ではない = どのみち scale 無しの経路。

2-3. models/test_qwen3_5_moe_weight.py の 2 件。**機構が置き換わっていた。**
   こちらは W4A16 で欠けた input_scale を 1.0 で埋めていた。上流の scheme reader
   (FlashML-org#438) は源で解決していて、W4A16 module は input_scale role を持たないので
   **キーがそもそも出ない**。例外は出ておらず、埋めた値の断定だけが落ちていた。

   **Ornith-1.5-35B-A3B-NVFP4 を実際に読ませて確かめた**: ロード 12.7s、
   24 token 生成(枠ちょうど)。上流の方式で serve できている。
   よって機構を戻さず、試験を上流の実挙動へ合わせた:
   - キーが出ないこと+**重みは届いていること**を見る(「例外が出ない」だけにしない)
   - 本当に欠けた重みは今も error。文言が変わったので、module 名が
     メッセージに入ることを別途 assert する

全スイート: 11 failed(**standing のみ、新規ゼロ**)/ 2,108 passed。

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

This branch has not been deployed

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

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant