Skip to content

[Perf] [AMD] kimi-k3: Run the fused KDA decode with a BF16 state and Quark weights - #41794

Open
yuychang wants to merge 6 commits into
sgl-project:mainfrom
yuychang:k3-pr1-quark-weights
Open

yuychang wants to merge 6 commits into
sgl-project:mainfrom
yuychang:k3-pr1-quark-weights

Conversation

@yuychang

@yuychang yuychang commented Sep 30, 2026 •

Copy link
Copy Markdown
Contributor

Motivation

  • On ROCm, Kimi-K3 has a fused KDA decode kernel: f_b projection + convolution + recurrent update + gated RMSNorm in one gfx950 launch per KDA layer.
  • With the aiter backend it never turns on for the deployed Quark checkpoint (amd/Kimi-K3-Quark-MXFP4-AttnFP8), so decode silently runs the unfused chain and pays the extra launches every step:
Blocker Cause
Wrong AITER module SGLANG_K3_FLYDSL_SOURCE=aiter imports a module no AITER build ships; ROCm/aiter#6026 publishes the kernel under another name
BF16 recurrent state rejected the adapter accepts FP32 only, but the recommended --mamba-ssm-dtype bfloat16 pool is BF16
FP8 f_b the checkpoint stores f_b as per-channel FP8; the kernel takes dense BF16
  • The two fused-decode env vars were read with raw os.environ, outside Envs.

Modifications

  • Bind flydsl_kda_decode_with_f_b from aiter.ops.flydsl.linear_attention_kernels; if AITER lacks [FlyDSL] [CI] [Kernel] Add fused KDA decode kernels on gfx950 ROCm/aiter#6026, warn once and fall back.
  • Arm a fused kernel only if it takes the pool's state dtype (the AITER op advertises state_dtypes); otherwise warn once and keep the unfused path.
  • Dequantize the Quark per-channel FP8 f_b to a BF16 copy once, for the fused kernel only; the FP8 layer keeps serving the unfused path.
  • Mirror AITER's cache-slot and stride rules in covered(), so an unsupported layout falls back instead of raising.
  • Warm up only the pool's state dtype before graph capture; log the armed kernel and dtype.
  • Apply the deferred f_b on the ReplaySSM decode-ring path too (it has no fused kernel); before, that path received the raw f_a.
  • Scope: ROCm only. CUDA is unchanged.

Arming rules:

Backend State dtype Other conditions Result
gfx950 (default) FP32 gfx950 armed
gfx950 BF16 - not armed, one-time warning
aiter, source auto / sglang (vendored kernel) FP32 gfx950 armed
aiter, source auto / sglang BF16 - not armed, one-time warning
aiter, source aiter FP32 or BF16 AITER with ROCm/aiter#6026 armed
aiter, source aiter any AITER without #6026 not armed, one-time warning
any any Quark f_b already post-processed (ShardedStateLoader, RemoteModelLoader) not armed, one-time warning

Env vars (moved into Envs; the old names keep working with a DeprecationWarning):

Name Deprecated alias Default Effect
SGLANG_USE_HIP_K3_KDA_FUSED_BACKEND SGLANG_ROCM_K3_KDA_FUSED_BACKEND gfx950 gfx950: SGLang HIP kernel; aiter: FlyDSL boundary; other values: unfused decode
SGLANG_USE_HIP_K3_FLYDSL_SOURCE SGLANG_K3_FLYDSL_SOURCE auto FlyDSL module behind aiter: auto prefers the vendored copy, sglang / aiter pin one
  • Docs: both vars in the AMD & ROCm env-var table; the Kimi-K3 cookbook documents the FP8 dequant and the state-dtype rule (it keeps the old names for its pinned image).

Accuracy Tests

GSM8K, sgl-eval 0.1.3, --load-preset-from-model-id moonshotai/Kimi-K3 (temperature 1.0, top_p 0.95, thinking on), all 1319 questions, one run per side.

main This PR (fused KDA decode on)
97.27% 97.12%
  • The 0.15-point gap is 2 questions, within noise (standard error about +/-0.45 points).

Unit and kernel tests (one MI355X):

Test AITER Result
test_kimi_k3_kda_decode.py integration build with ROCm/aiter#6026 (6ae4bf0f2) 25 passed
test_kimi_k3_kda_decode.py pinned image 285579b62 (no #6026) 15 passed, 10 skipped
test_kda_fused_decode_gfx950.py pinned image 4 passed
test_kimi_k3_rocm_quant.py CPU 8 passed

Speed Tests and Profiling

  • Setup: 8x MI355X, TP8, random 8192 in / 1024 out, radix cache off, FP8 KV, AITER attention, --mamba-ssm-dtype bfloat16.
  • Base: main at dd9ea06c00. PR side adds SGLANG_ROCM_K3_KDA_FUSED_BACKEND=aiter SGLANG_K3_FLYDSL_SOURCE=aiter.
  • AITER be461934f (has [FlyDSL] [CI] [Kernel] Add fused KDA decode kernels on gfx950 ROCm/aiter#6026) on both sides. One run per side; a repeat of main matched its TPOT within 0.4%.
Conc TPOT base (ms) TPOT PR (ms) Delta
2 18.22 17.70 -2.9%
4 19.90 19.31 -2.9%
8 23.00 22.48 -2.3%
16 28.12 27.51 -2.2%
32 36.64 36.33 -0.8% (noise)
64 52.85 52.35 -0.9% (noise)
  • TPOT drops 2.2-2.9% at concurrency 2-16; at 32 and 64 the change is about twice the noise, so it is not claimed.
  • TTFT does not change (within +/-0.8%), as expected for a decode-only change.
  • These numbers come from an earlier revision; later changes touch load-time checks, warnings, warm-up and the ReplaySSM fallback, not the kernels.

Checklist


CI States

Latest PR Test (Base): ❌ Run #42094941736
Latest PR Test (Extra): ❌ Run #42094941906
Latest PR Test (AMD ROCm 10): ❌ Run #42094941657

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

@github-actions github-actions Bot added quant LLM Quantization amd jit-kernel labels Sep 30, 2026
@yuychang yuychang changed the title [AMD][Quark] [AMD][Quark] Lay out Quark weights to speed up ROCm decoding Sep 30, 2026
@yuychang
yuychang force-pushed the k3-pr1-quark-weights branch 2 times, most recently from 4ad0fa8 to 73885df Compare September 30, 2026 01:29
@github-actions github-actions Bot added the documentation Improvements or additions to documentation label Sep 30, 2026
@yuychang
yuychang marked this pull request as ready for review September 30, 2026 01:57
@sogalin sogalin added the run-ci CI: run the baseline test suite on this PR label Sep 30, 2026
@yuychang yuychang changed the title [AMD][Quark] Lay out Quark weights to speed up ROCm decoding [AMD][Quark] Kimi K3 Quark FP8/MXFP4 fusion Oct 1, 2026
@yuychang

yuychang commented Oct 2, 2026

Copy link
Copy Markdown
Contributor Author

/rerun-failed-ci

@yuychang
yuychang requested a review from BowenBao as a code owner October 5, 2026 08:25
@yuychang

yuychang commented Oct 6, 2026

Copy link
Copy Markdown
Contributor Author

/rerun-failed-ci

@kkHuang-amd

Copy link
Copy Markdown
Collaborator

@amd-bot ci-status

@yuychang yuychang changed the title [AMD][Kimi-K3] Run the Quark per-channel FP8 KDA in-projection and MLA absorb on FP8 kernels [Perf] [AMD] kimi-k3: Run Quark FP8 KDA in-proj and MLA absorb on FP8 kernels Oct 17, 2026
@yuychang
yuychang force-pushed the k3-pr1-quark-weights branch 2 times, most recently from 250b58b to ee938d5 Compare October 21, 2026 15:42
@yuychang yuychang changed the title [Perf] [AMD] kimi-k3: Run Quark FP8 KDA in-proj and MLA absorb on FP8 kernels [Perf] [AMD] kimi-k3: Run the fused KDA decode on AITER #6026 with a BF16 state and the Quark checkpoint Oct 21, 2026
@yuychang yuychang changed the title [Perf] [AMD] kimi-k3: Run the fused KDA decode on AITER #6026 with a BF16 state and the Quark checkpoint [Perf] [AMD] kimi-k3: Run the fused KDA decode on ROCm/aiter#6026 with a BF16 state and the Quark checkpoint Oct 21, 2026
@yuychang yuychang changed the title [Perf] [AMD] kimi-k3: Run the fused KDA decode on ROCm/aiter#6026 with a BF16 state and the Quark checkpoint [Perf] [AMD] kimi-k3: Run the fused KDA decode with a BF16 state and Quark weights Oct 21, 2026
@mintlify

mintlify Bot commented Oct 21, 2026 •

Copy link
Copy Markdown
Contributor

Preview approval required

This pull request comes from a fork, so Mintlify won't build a preview until someone with write access approves commit dea7326.

Approve preview

@yuychang

Copy link
Copy Markdown
Contributor Author

/rerun-failed-ci

1 similar comment
@chien-an-chen

Copy link
Copy Markdown
Contributor

/rerun-failed-ci

@yuychang
yuychang force-pushed the k3-pr1-quark-weights branch from 49e7c45 to 3fe31bf Compare October 29, 2026 02:50
@bingxche bingxche removed the run-ci-extra CI: also run the extra suite (requires run-ci) label Oct 30, 2026
@github-actions

This comment has been minimized.

…e pool's state dtype

- Load flydsl_kda_decode_with_f_b from aiter.ops.flydsl.linear_attention_kernels when the FlyDSL source is aiter.
- Fail closed with a one-time warning when AITER lacks that kernel or the FlyDSL import fails.
- Take the recurrent-state dtypes the bound kernel advertises; the vendored kernel stays FP32-only.
- available() and warmup() take the mamba pool's state dtype; a kernel that cannot take it is not armed and warns once.
- Arm the gfx950 kernel only for an FP32 pool, so a BF16 pool no longer arms it.
- Mirror AITER's cache-slot and stride checks in covered(), so an unsupported layout falls back instead of raising.
- Warm only the pool's state dtype before graph capture and log the armed kernel and dtype.
- Test the bindings, BF16 state, SGLang's cache pool layouts, warmup, arming, and vendored-vs-AITER agreement.
- End test_kimi_k3_kda_decode.py with the __main__ block test/README.md prescribes.
- Document the state-dtype arming rule, the backend default, and the aiter source in the Kimi-K3 cookbook.
- The fused KDA decode takes a dense BF16 f_b with no weight scale; Quark stores f_b as per-channel FP8.
- Dequantize it once to a BF16 copy for the fused kernel only; the FP8 layer keeps serving the unfused path.
- The helper lives in the ROCm-only kimi_k3_rocm_quant.py and runs only once a fused backend is available.
- ShardedStateLoader and RemoteModelLoader transpose the FP8 weight before post_load_weights.
- Skip the dequant for such a weight and warn once; the layer keeps the unfused path instead of failing the load.
- Add CPU tests for the dequant, the BF16 pass-through, and the post-processed weight.
- Note the BF16 dequant in the Kimi-K3 cookbook.
- Define SGLANG_USE_HIP_K3_KDA_FUSED_BACKEND and SGLANG_USE_HIP_K3_FLYDSL_SOURCE in Envs.
- Keep SGLANG_ROCM_K3_KDA_FUSED_BACKEND and SGLANG_K3_FLYDSL_SOURCE as deprecated aliases.
- Read both through envs in the two fused-decode adapters, the FlyDSL source loader and kimi_k3.py.
- Switch the tests to envs overrides and cover the two aliases.
- Document both in the AMD & ROCm table; the cookbook keeps the old names for its pinned image.
@yuychang
yuychang force-pushed the k3-pr1-quark-weights branch from c790c48 to 39f747c Compare November 2, 2026 06:01
# Conflicts:
#	python/sglang/kernels/ops/attention/kda_fused_decode_gfx950.py
…fused KDA decode checks

- Materialize the deferred f_b whenever it is deferred; the ReplaySSM ring path skipped it and got the raw f_a.
- Split the gfx950 arch check from available(), so covered() no longer logs a "not armed" warning per step.
- Restore the unset-default case in test_fused_kda_backend_is_opt_in.
- Clarify in the env-var docs that a BF16 state needs SGLANG_USE_HIP_K3_FLYDSL_SOURCE=aiter.
- Say in the cookbook which init failures log the one-time warning.
@yuychang
yuychang marked this pull request as ready for review November 5, 2026 17:40

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

…fused decode arming

- SGLANG_USE_HIP_K3_FLYDSL_SOURCE=auto still prefers the vendored kernel, but binds AITER's when only it takes the state
- Cache the kernel binding per (source, state dtype); covered() and run() no longer re-import or probe the device
- Warn separately when AITER lacks flydsl_kda_decode_with_f_b and when the op lacks the state_dtypes tag
- Reset the HIP fused decode arming before re-arming, so a re-run that fails cannot keep a stale f_b in use
- Take the merged BF16 f_b (block-FP8 and BF16 checkpoints) first; warn once when no BF16 f_b can be produced
- Make the decode adapters' available() keyword-only and fix the gfx950 docstring that described an unused device
- Re-wrap the new environ.py comments at clause boundaries and restore the decode test's -v main call
- Tests: auto binds AITER for a BF16 pool, both AITER failure messages, failed re-arm, merged block-FP8 f_b
- Tests: set the deprecated env names through temp_set_env(allow_sglang=True), as they have no descriptor
- Docs: describe the dtype-aware auto source and the BF16 f_b requirement in the cookbook and env var reference

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

amd documentation Improvements or additions to documentation jit-kernel quant LLM Quantization run-ci CI: run the baseline test suite on this PR

Projects

None yet

Development

Successfully merging this pull request may close these issues.

6 participants