Repository navigation
5080deploy: opt-in on-demand experts for short prefills (FREETOKEN_MOE_SMALL_PREFILL_TOKENS, off) - #68
Open
gdevenyi wants to merge 3 commits into
Open
5080deploy: opt-in on-demand experts for short prefills (FREETOKEN_MOE_SMALL_PREFILL_TOKENS, off)#68gdevenyi wants to merge 3 commits into
gdevenyi wants to merge 3 commits into
Conversation
…_MOE_SMALL_PREFILL_TOKENS) A prefill streams every expert of every layer (the next layer's routing is unknown when its prefetch must start). For a short extension over a cached prefix -- an agent turn -- that is ~68 GB on Qwen3.8-Flash-Next and ~1.3 s before the first token, whatever the token count. Prefills of at most N tokens now take the decode path's on-demand LRU load (route, fetch the missing routed experts, slot GEMM); N = 0 (default) keeps the whole-layer path everywhere. Assisted-by: Claude Opus 5.5
…unked under Triton's limit The kernel's block is next_pow2(ids) x next_pow2(slots): a 156-token turn's 1560 routed ids against 1650 slots is a 2048 x 2048 block, past Triton's 1M-element limit, and the compile error took the engine down. The on-demand path now dedups the ids (at most num_experts), ensures them in chunks that fit, copies each chunk's misses before the next ensure, and maps the tokens' ids to slots; a layer needing more than half the cache takes the whole-layer path. The GPU tests now assert that the on-demand path is really taken. Assisted-by: Claude Opus 5.5
Chunking the unique ids only under Triton's element limit still handed the LRU kernel a 512 x 2048 block; on the serving GPU (~1 GiB free) its launch failed with CUDA out of memory on a 156-token turn and took the engine down. Each ensure now takes at most 32 ids, the width a decode step already gives it. The GPU test asserts the width. Assisted-by: Claude Opus 5.5
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.
Stacked on #67. This is an opt-in path that is off by default and left off in production. It is published with its measurements so the result isn't lost.
What:
FREETOKEN_MOE_SMALL_PREFILL_TOKENS=N. A prefill of at most N tokens makes only its routed experts resident through the LRU cache, instead of streaming every expert of every layer.Commits:
51a6321: the path and its dispatch tests.d47f8eb: handlru_ensurethe unique ids, not tokens × top-k. The kernel block is next_pow2(ids) × next_pow2(slots); a 156-token turn overflowed Triton's 1M-element limit and took the engine down.283cb2e: ensure in 32-id (decode-width) chunks. A 512 × 2048 block was still under the limit but failed to launch with CUDA OOM on the serving GPU.Each fix has a GPU test that fails before and passes after. The real-size test (200 tokens × top-10, 512 experts, 1100 slots) asserts the ensure width.
A/B on an RTX 5080, Qwen3.8-Flash-Next-NVFP4, 1650 slots, a ~28-40K cached conversation:
arith2flips from 1047576 to 1047572, the same fragile item a bf16->FP8 change flips. On-demand prefill goes through the decode GEMM and its rounding, so its outputs are not identical to the whole-layer path.Verdict: a ~0.5 s TTFT win on turns under ~200 new tokens. Past a few hundred tokens it is 2× slower: nearly all experts are needed, and chunked ensure + copy loses to the overlapped bulk stream. With the numerics change on top, production stays at 0 (off). A useful N would sit around 128-256 and would still need a quality sign-off.
🤖 Generated with Claude Code
https://claude.ai/code/session_01Bu6LgoxLR4wETqb7RPR2vt