Skip to content

[Perf] kimi-k3: Seed a validated FlashInfer prefill MoE tactic - #47653

Merged
BBuf merged 3 commits into
sgl-project:mainfrom
Codoge:perf/k3-moe-prefill-tile128
Nov 2, 2026
Merged

BBuf merged 3 commits into
sgl-project:mainfrom
Codoge:perf/k3-moe-prefill-tile128

Conversation

@Codoge

@Codoge Codoge commented Oct 31, 2026 •

Copy link
Copy Markdown
Contributor

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

  • Seed one serving-derived 16K MoE tuning bucket during startup, before the user's persistent cache is loaded. Explicit user-cache entries retain precedence; disabling cache reuse keeps fresh tuning behavior.
  • Restrict loading to Kimi-K3, TP1/PP8/EP1, chunk16K, BF16 model dtype, SM103 and FlashInfer 0.7.2. FlashInfer's full environment metadata check further restricts the bundled record to the measured B30Z/CUDA/library combination.
  • The cache key includes both activation/weight formats, local/global expert counts, top-k, hidden/intermediate dimensions, activation and weight layout, plus the unpacked routing representation.
  • Fix the existing target/draft CUDA cache-phase fixture to supply the runner field read by the startup guard.
  • Use the existing load_configs mechanism; do not replace AutoTuner.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

  • 2 unit tests passed, with subcases rejecting other model types, parallelism, chunk sizes, dtypes, hardware and unvalidated FlashInfer 0.7.1/0.7.3.
  • On B300/B30Z, loaded this branch's record with the real FlashInfer 0.7.2 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.
  • Loaded an explicit user record afterwards: the same real operation selected (192, 4), confirming user-cache precedence; output again matched the reference bitwise.
  • Ruff, isort, AST/JSON parsing, registered-test lint, and git diff --check passed.

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 main base 592754fc5e, 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 cc5c303a8ec5db2c0168414c00f0a3b058161775 repairs the existing CUDA test fixture: initialize is_draft_worker=False for the target phase and set it to True before 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:

  • The original TestAutotuneCachePhases reproduces AttributeError: 'types.SimpleNamespace' object has no attribute 'is_draft_worker'.
  • After the fix, the complete 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 upstream main base 592754fc5e, 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.

Run Requests Input tok/s P50 TTFT s P99 TTFT s
baseline A1 128 42,203.78 6.128 7.046
KDA only 128 42,899.91 6.022 6.937
MoE only 128 42,813.05 6.034 6.926
combined B1 128 43,569.94 5.925 6.835
baseline A2 128 42,130.58 6.137 7.050
combined B2 128 43,524.13 5.940 6.823
baseline pooled 256 42,167.15 6.132 7.359
combined pooled 256 43,547.02 5.930 7.125

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.

Configuration Requests Input tok/s P50 TTFT s P99 TTFT s
same-condition baseline 128 41,910.90 6.164 7.128
MoE tile128/config0 only 128 42,500.14 6.080 7.015

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

  • Preserve routing, weights, precision, dimensions and operation implementation.
  • Avoid process-wide method replacement and new environment switches.
  • Verify the real cached tactic, captured-tensor output and user-cache precedence.
  • Keep strict library/device metadata instead of assuming cross-version tactic compatibility.
  • Validate and add a new record for the pinned FlashInfer 0.7.2; reject other versions.
  • Complete common-base full-model serving validation, including repeated baseline/combined A/B and fixed output checks.

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-test to run the tests your change touches, or copy this prompt for an agent:

Find the registered tests this PR affects and run them with /rerun-test, following https://github.com/sgl-project/sglang/blob/main/.agents/skills/rerun-test/SKILL.md

@BBuf BBuf left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[BBuf Review Agent] LGTM. Please fix the CUDA cache-phase fixture and wait for related CI.

@BBuf

BBuf commented Nov 1, 2026

Copy link
Copy Markdown
Collaborator

/tag-and-rerun-ci

@github-actions github-actions Bot added the run-ci CI: run the baseline test suite on this PR label Nov 1, 2026
@BBuf

BBuf commented Nov 1, 2026

Copy link
Copy Markdown
Collaborator

/tag-and-rerun-ci

1 similar comment
@BBuf

BBuf commented Nov 1, 2026

Copy link
Copy Markdown
Collaborator

/tag-and-rerun-ci

@Codoge
Codoge marked this pull request as ready for review November 1, 2026 11:19
@Codoge

Codoge commented Nov 1, 2026

Copy link
Copy Markdown
Contributor Author

@BBuf Fixed the CUDA cache-phase fixture in cc5c303a8ec5db2c0168414c00f0a3b058161775: the target runner starts with is_draft_worker=False, then switches to True for draft-cache loading. Only the existing test fixture changed.

On B300/B30Z (SM103), PyTorch 2.14.1+cu130, FlashInfer 0.7.2, the original fixture reproduced the reported AttributeError. With the fix:

python3 test/registered/unit/model_executor/runner/test_flashinfer_autotune_sync.py -v
Ran 12 tests in 36.790s
OK

python3 test/registered/unit/layers/moe/test_kimi_k3_prefill_tactics.py -v
Ran 2 tests in 0.004s
OK

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-test permission:

/rerun-test test/registered/unit/model_executor/runner/test_flashinfer_autotune_sync.py test/registered/unit/layers/moe/test_kimi_k3_prefill_tactics.py

@claude claude Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

@Codoge

Codoge commented Nov 1, 2026

Copy link
Copy Markdown
Contributor Author

/rerun-failed-ci

@Codoge

Codoge commented Nov 1, 2026

Copy link
Copy Markdown
Contributor Author

CI rerun follow-up: the baseline gate now passes, but CPU shard 16 hits an unrelated upstream fixture mismatch:

ERROR: TestDcpReplayMetadata.test_replay_reuses_the_captured_gather_buffer
AttributeError: 'types.SimpleNamespace' object has no attribute
'skips_prefill_metadata_on_idle_dp_rank'

The run checked out merge commit 5aaa1d39890c4c5402f34553b2e3f2c31ffa93b3 (this PR merged into base c757c4bb7c55d1eb69fdb2bdc6e6fb8507acb7a0). I compared both affected files against that base and current main; all three versions have identical blob IDs:

  • test/registered/unit/model_executor/test_prefill_cuda_graph_dcp.py: ef5beb976949ec8fb56be440812c166c5e2b2177.
  • python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py: 803428ce82564b55938605844b52b3e7142d3f4f.

The fixture's attn_backend = SimpleNamespace(...) omits the field that the replay path reads directly at line 1371. Neither file is changed by this PR. This is separate from the FlashInfer target/draft fixture fixed in cc5c303a8ec5db2c0168414c00f0a3b058161775, whose full test file passed on B300.

The upstream DCP test fixture likely needs skips_prefill_metadata_on_idle_dp_rank=False for this non-idle replay case. Reporting the exact blocker here so the new CI red is not confused with the fixed FlashInfer test; the other jobs are still running.

@BBuf

BBuf commented Nov 1, 2026

Copy link
Copy Markdown
Collaborator

[BBuf Review Agent] @Codoge The target/draft fixture fix looks correct. Confirmed cc5c303 only changes that test, so the production measurements still apply.

@BBuf

BBuf commented Nov 1, 2026

Copy link
Copy Markdown
Collaborator

[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.

@BBuf BBuf left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[BBuf Review Agent] LGTM.

@BBuf
BBuf merged commit 8319feb into sgl-project:main Nov 2, 2026
243 of 273 checks passed
Yuxingwang-intel pushed a commit to Yuxingwang-intel/sglang that referenced this pull request Nov 6, 2026
…roject#47653)

Co-authored-by: Codoge <12482210+Codoge@users.noreply.github.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

run-ci CI: run the baseline test suite on this PR

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants