Repository navigation
[Perf] kimi-k3: Seed a validated FlashInfer prefill MoE tactic - #47653
Conversation
BBuf
left a comment
There was a problem hiding this comment.
[BBuf Review Agent] LGTM. Please fix the CUDA cache-phase fixture and wait for related CI.
|
/tag-and-rerun-ci |
|
/tag-and-rerun-ci |
1 similar comment
|
/tag-and-rerun-ci |
|
@BBuf Fixed the CUDA cache-phase fixture in On B300/B30Z (SM103), PyTorch 2.14.1+cu130, FlashInfer 0.7.2, the original fixture reproduced the reported Both complete files passed with zero skips, including the actual CUDA cache-phase test. The PR description now includes completed common-base full-service validation of the production changes: MoE-only +1.44% (one screening pair), and combined KDA+MoE +3.27% over two A/B pairs. The combination's gain is not attributed to this PR alone; the test-only follow-up does not change the measured production code. Marked ready for review. Please re-review the fixture fix and trigger the related upstream tests; my account does not have targeted |
|
/rerun-failed-ci |
|
CI rerun follow-up: the baseline gate now passes, but CPU shard 16 hits an unrelated upstream fixture mismatch: The run checked out merge commit
The fixture's The upstream DCP test fixture likely needs |
|
[BBuf Review Agent] @Codoge Confirmed the base already has this fixture mismatch, and neither file is changed by this PR. CPU shard 16 is unrelated. |
…roject#47653) Co-authored-by: Codoge <12482210+Codoge@users.noreply.github.com>
Motivation
For the measured Kimi-K3 TP1/PP8 prefill workload, FlashInfer 0.7.1's 16K-token MoE candidate enumeration exposed token tiles 192 and 256. A supported tile128/config0 tactic, discovered through smaller-shape candidate enumeration and then executed at the real 16K shape, performed better on captured early/late-layer routing.
This ports the tile128 approach to the repository's pinned FlashInfer 0.7.2, using its existing cache-loading API. Re-enumeration on B30Z identified tile128/config4, and real captured-input replay verified that tactic and the new record. The older 0.7.1/config0 record is not installed or accepted. Tactic IDs are version-specific; no dependency downgrade or version-check bypass is included.
Modifications
load_configsmechanism; do not replaceAutoTuner.choose_one, add a process-wide monkey patch, or add a new configuration switch.The record follows FlashInfer's tuning-bucket semantics. It is not a universal tile128 policy for arbitrary sequence lengths, models, hardware, EP or decode workloads. Existing explicit cache choices can override it.
Accuracy Tests
AutoTuner, then replayed captured Kimi-K3 input[16384, 3584], original expert weights, scales and top-k routing. The real operation selected (128, 4) and output was bitwise equal to the captured service reference.git diff --checkpassed.The 0.7.2 Python wheel and its matching NVIDIA kernels were tested in an isolated overlay, retaining version checks; the installed 0.7.1 package was not modified. These verify the policy, public cache adapter and actual MoE operation.
Full-model serving validation now passes for these production changes on the common upstream
mainbase592754fc5e, with FlashInfer 0.7.2; the standalone and repeated combined results are reported below. The current six full-service runs each pass 12 fixed greedy output checks covering input lengths 1,024 / 16,383 / 16,384 / 16,385 / 32,768 / 65,536, each repeated twice. All 60 comparisons against A1 match output token IDs and input SHA256 hashes. An additional eight C4 concurrent fixed-input checks match between A2 and B2. These are finite output-regression checks, not logits-level equality or a full task-accuracy benchmark.Cache-phase fixture review follow-up (2026-11-01)
Commit
cc5c303a8ec5db2c0168414c00f0a3b058161775repairs the existing CUDA test fixture: initializeis_draft_worker=Falsefor the target phase and set it toTruebefore draft-cache loading. This commit changes only the test; the production code measured below is unchanged.On B300/B30Z (SM103), PyTorch 2.14.1+cu130 and FlashInfer 0.7.2:
TestAutotuneCachePhasesreproducesAttributeError: 'types.SimpleNamespace' object has no attribute 'is_draft_worker'.test/registered/unit/model_executor/runner/test_flashinfer_autotune_sync.py: 12 tests passed, 0 skipped, including the real CUDA cache-phase case.test/registered/unit/layers/moe/test_kimi_k3_prefill_tactics.py: 2 tests passed, 0 skipped.Speed Tests and Profiling
Full serving validation of the upstream integration (2026-11-01)
Tested production changes from #47651 (
7933473906e) and #47653 (34238f31a5d) on the same upstreammainbase592754fc5e, using FlashInfer 0.7.2 for every arm. Kimi-K3, one 8-GPU B300/B30Z, TP1/PP8/EP1, chunk 16,384, pipeline async depth 1, client C4, ISL 65,536 / OSL 1. CUDA Graphs and radix caching disabled; model precision and all serving settings held fixed. Each formal run contains 128 requests after a separate 8-request warmup. Formal measurements exclude profiler/diagnostic requests and prototype hooks.Each run uses an independent autotune cache cloned from the same baseline startup cache. Baseline/KDA caches contain no 16K MoE seed; MoE/combined caches contain the FlashInfer 0.7.2
(128, 4)record on all eight PP ranks. Separate post-benchmark Torch GPU traces of all eight PP ranks confirm the actual production paths: 1,104 KDA input-copy kernels become 276 pack kernels, and all 736 MoE GEMMs change token tile 192 to 128. These C1 diagnostic traces verify execution paths; their timings are not used as C4 serving measurements.MoE-only throughput changes by +1.44% against A1. This is one standalone screening pair. The repeated combined KDA+MoE result changes throughput by +3.27% (paired runs +3.24% / +3.31%); it must not be attributed to either PR alone. Pooled throughput is total input tokens / total formal duration; P50/P99 are recomputed from all formal request TTFTs. All 768 formal requests succeed with zero cached tokens and retractions. The combined pooled P50/P99 TTFT decrease by 3.28% / 3.17%; no formal request reaches 15 seconds. OSL is 1, so TPOT/ITL are not reported.
Earlier prototype evidence
Historical FlashInfer 0.7.1/config0 prototype measurements, not end-to-end measurements of the new 0.7.2/config4 startup-cache integration. Kimi-K3 / SGLang 0.5.22 / FlashInfer 0.7.1, 8x B300, TP1/PP8, ISL 65,536 / OSL 1, chunk16K, client C4, unchanged precision, no profiler attached to formal requests.
This single historical screening pair shows approximately +1.41% serving throughput; that historical pair is not a measurement of the 0.7.2 record and is not a repeated standalone result. The independently repeated combined MoE+KDA prototype achieved +3.12%/+3.32% and +3.14% in a longer 512-request-per-arm comparison. Those combined gains must not be attributed to this PR alone.
In matched-layer NSYS analysis of the combined prototype, MoE GEMM1+GEMM2 changed from 10.330 to 9.622 ms per invocation (about -6.85%). Smaller token tiles reduce estimated per-expert padding, but also change CTA counts and pipeline/resource use. Padding estimates and isolated NCU timing are not substitutes for serving measurements.
Checklist
CI States
Latest PR Test (Base): ❌ Run #41301201477
Latest PR Test (Extra): ❌ Run #41301201431
Latest PR Test (AMD ROCm 10): ✅ Run #41301201457
Tip
Base CI only covers the core cases, and its workload may not exercise your change. Use
/rerun-testto run the tests your change touches, or copy this prompt for an agent: