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
Open
aswinkumar1999 wants to merge 4 commits into
aswinkumar1999 wants to merge 4 commits into
Conversation
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.
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
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Add an automatic KV/MoE cache ladder
Summary
This adds
--enable-kv-ladder, an opt-in serving mode that grows the KV cache when anincoming 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.
With the default step and this model, the resulting ladder is:
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
could exactly fill the current capacity.
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:
RadixArk/Qwen3.8-Flash-Next-NVFP4Short-output results include proportionally more fixed launch and streaming overhead. The
4,096-token column is the more representative sustained-decode comparison.
Rebuild latency
approximately 0.67-0.99 seconds.
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
status=ok; no capacity rejection occurred.git diff --checkpassed.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.