Pull requests / #3545
#3545 [Engine] verify: STRATA_VERIFY_ECAP - cap the verify window's expert set (relaxed MTP, opt-in)
open · @vlupilin · 0 commentaires · Sur GitHub
BenchmarksAMD / HIPNVIDIA / CUDAModels & quantsLinux
Description
Problem and result: In a spec-decode verify window the engine reads every unique expert the window's T x K routed picks name, from VRAM. The picks beyond position 0 belong to draft positions that are often rejected - on this card their expert reads are ~30% of the window's GPU time, spent on tokens that get thrown away. With STRATA_VERIFY_ECAP=N (opt-in) the window keeps token 0's full K picks plus the top (N-K) tail entries by routing weight, and the expert reads drop accordingly. Measured on gfx1151 (OrcaRouter 256-expert Qwen3.8-Flash-Next, 512K context, --spec 4, interleaved A/B, fresh engine per arm, identical prompt per round): decode +10.5% at 32K, +24.0% at 131K, +13.0% at 262K, +17.6% at 524K context - all 8 paired runs positive; prefill unchanged; the window's expert phase 15.72 -> 10.95 ms (-30%); draft acceptance held (78% vs 72%). Default off: the uncapped path's bits are unchanged.
Related discussion: none in this tracker. The concept is vLLM PR #60689's moe_expert_cap (Gemma4 MTP); this is an independent implementation for Strata's resident-plan verify path, found while searching for decode ideas.
Approach: One small kernel (ecap_redirect_kernel) runs between the router and resident_plan, inside the captured window (zero-doorbell graph compatible). Token 0 (the guaranteed position) keeps all K picks; the speculative tail keeps the top (cap-K) entries by (weight desc, index asc); every dropped entry is redirected round-robin onto a KEPT expert with weight 0 - the plan reads no extra expert for it and its contribution is exactly zero (every combine path multiplies by the routing weight). The round-robin matters: folding all drops into one group was measured and rejected - a single M=45 group serialized the grouped GEMV and the expert phase got SLOWER. The keep-list is built serially by thread 0 in shared memory (n <= 80 entries, no atomics - an atomicAdd build crashed on divergent branches). Five files, +180 lines (63 of them the kernel and its insertion, the rest a contract test); no kernel signature changes: resident_plan and the combines consume the plan exactly as before.
Validation:
Passed: HIP build of this exact revision (gfx1151, ROCm 10.1, Linux 7.0).
ecap_parity unit test (this tree, this revision): a cap >= the entry count is a bitwise no-op; token 0's entries are never touched; kept entries keep id and weight; every dropped entry's weight is exactly 0.0 and its id is one of the kept set's; the same input twice gives the same output - T 2/3/4/6 x cap 11/16/20/26 on random router-shaped data (256 experts, top-10).
Engine window profile (STRATA_VERIFY_PROFILE=1 STRATA_DECODE_TIMING=1, 32K prompt, thinking on): "VRAM hits" 15.72 -> 10.95 ms/window, the whole window 51.41 -> 46.50 ms, no other phase moved beyond noise.
End-to-end A/B, 2 rounds interleaved, fresh engine per arm, 512 output tokens, identical prompt per round, no other load, prod sampling (temp 1.0, thinking on): decode 32K 47.0/46.8 vs 41.6/43.3, 131K 48.6/46.5 vs 37.6/39.1, 262K 44.5/42.3 vs 38.4/38.4, 524K 48.0/36.2 vs 36.2/35.4 tok/s (capped/uncapped); prefill at every depth within 1.4% (parity, as expected - the cap is decode-only).
Draft acceptance: 140 of 180 (78%) capped vs 143 of 198 (72%) uncapped.
Quality, capped arm vs uncapped control the same day: a 6-prompt suite (Rust O(n log n) codegen with tests, a modified river-crossing trap, a quadratic Diophantine, a forbidden-word constraint, a Grover lower-bound attribution, an HTTP semantics table) run 4x per arm - indistinguishable pass/fail profiles between arms; needle-in-haystack retrieval 5/5 keywords at ~500K tokens with the cap on.
Failed: two intermediate variants, kept out of the final diff and worth recording: all-drops-into-one-group (the expert phase got SLOWER at cap 16: the M=45 group serialized the grouped GEMV - group balance is the whole game), and an atomicAdd-based keep-list build in divergent branches (engine error "the all-resident plan met an expert that is not in VRAM"; replaced by the serial thread-0 build).
Not run: CUDA build (no NVIDIA toolchain on this machine; the kernel is plain CUDA C with no HIP-only intrinsics); other models (only the 256-expert Qwen3.8-Flash-Next family measured); gfx12 and pre-gfx11 cards.
Limits: Relaxed-MTP semantics: the acceptance checks for draft positions run on approximately-capped logits, so the output can differ from the uncapped model where a dropped low-weight expert would have flipped a draft acceptance (the suite and the needle test above did not catch such a flip; the knob is one env away from off). The gain grows with window length, so deep contexts benefit most. Covers the all-resident (ar_on) verify path; the host-plan path is untouched.
Sur le site
Liens install, modèles, releases.