Skip to content

perf(moe): add Kimi-K3 B1 prepared A8W4 pipeline - #2

Closed
JohnQinAMD wants to merge 1 commit into
perf/kimi-k3-preroute-bf16from
perf/kimi-k3-persistent-moe-pr
Closed

JohnQinAMD wants to merge 1 commit into
perf/kimi-k3-preroute-bf16from
perf/kimi-k3-persistent-moe-pr

Conversation

@JohnQinAMD

Copy link
Copy Markdown
Owner

perf(moe): add Kimi-K3 B1 prepared A8W4 pipeline

Summary

Add a narrow gfx950 Kimi-K3 batch-one MoE pipeline that owns route metadata,
MXFP8 activation quantization, tuned A8W4 expert GEMMs, SiTU, and weighted BF16
reduction behind one typed request.

The fast path is fail-closed. It requires the exact native TP8 checkpoint
layout, validated kernel registry entry, routing semantics, device, dtypes, and
runtime features. Unsupported inputs return to the existing path.

Design

  • Generate stable-priority top-16 route metadata for 896 experts while
    preserving exact incumbent tie behavior.
  • Overlap the one-row MXFP8 preparation with route selection.
  • Register the independently validated 14-wave stage-1 schedule as one explicit
    production variant rather than widening the generic tuning combinations.
  • Reuse the tuned stage-1 A8W4 SiTU/quant pipeline and explicitly select the
    native-384 stage-2 atomic weighted reduction.
  • Keep eligibility in one named public predicate, split internally into native
    tensor layout, model configuration, and runtime-feature contracts.
  • Expose immutable request and metadata dataclasses so the consumer owns
    prepared routing exactly once.
  • Leave generic registry selection and non-Kimi paths unchanged.

Performance

MI355X, one GPU microbenchmark, rotating through all 896 experts:

Seed Candidate Native incumbent Speedup
111 18.8658 us/layer 27.5647 us/layer 1.4611x
223 18.8490 us/layer 27.5779 us/layer 1.4631x

On the full TP8 Kimi-K3 endpoint (B1, 8,192 input tokens, 1,024 output tokens,
no speculative decoding), two measured runs were 69.7324 and 69.8257 tok/s.
The 69.7791 tok/s mean is 8.56% above the hash-pinned 64.2744 tok/s parent and
reduces mean TPOT by 1.2275 ms.

The final clean signed commit was rebuilt as a separately hash-pinned image
and confirmed at 69.4669 and 69.8105 tok/s (69.6387 mean), within 0.21% of
the accuracy-tested image.

Correctness and tests

  • Native-layout kernel correctness: 20/20 cases, 1,000 graph replays.
  • Minimum cosine similarity: 0.9999848008.
  • Maximum relative L2: 0.0055199412.
  • NaN/Inf: zero.
  • Focused AITER contract and GPU tests: 31 passed.
  • Ruff format/check and git diff --check: passed.
  • Frozen GSM8K completion contract: 1,256/1,319 (0.952237), one invalid
    response. The hash-pinned parent scored 1,244/1,319 (0.943139), with two
    invalid responses; paired outcomes were 32 candidate wins and 20 losses.

Scope and dependencies

This is stacked on JohnQinAMD/aiter-amd:perf/kimi-k3-preroute-bf16.
It intentionally excludes the rejected direct-BF16 stage-1 experiment and does
not change the generic Kimi-K3 tuning CSV.

The matching vLLM integration is a separate PR against
JohnQinAMD/vllm-amd.

AI assistance

OpenAI Codex assisted with implementation, profiling analysis, test
construction, and PR preparation. The human submitter reviewed the changed
lines and is responsible for the design, measurements, and submission.

Add a fail-closed gfx950 batch-one pipeline that prepares stable route metadata and MXFP8 activation once, then consumes it through tuned MXFP4 expert GEMMs and weighted BF16 reduction.

The specialization requires the exact native TP8 checkpoint layout and routing semantics. Unsupported configurations retain the existing path.

Assisted-by: OpenAI Codex
Signed-off-by: Yanyuan Qin <yanyuan.qin@amd.com>
@JohnQinAMD

Copy link
Copy Markdown
Owner Author

Superseded by DCO-clean, patch-equivalent replacement #22. The replacement adds the required Signed-off-by trailers without force-pushing this reviewed branch; performance, correctness, and accuracy evidence is preserved.

@JohnQinAMD JohnQinAMD closed this Jul 31, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant