Repository navigation
Conversation
yuychang
force-pushed
the
k3-pr1-quark-weights
branch
2 times, most recently
from
September 30, 2026 01:29
4ad0fa8 to
73885df
Compare
yuychang
marked this pull request as ready for review
September 30, 2026 01:57
yuychang
requested review from
1am9trash,
BBuf,
DarkSharpness,
HaiShaw,
HydraQYH,
JustinTong0323,
celve,
hubertlu-tw,
kkHuang-amd,
sogalin,
wisclmy0611,
yichiche,
yuan-luo and
zijiexia
as code owners
September 30, 2026 01:57
This was referenced Sep 30, 2026
Contributor
Author
|
/rerun-failed-ci |
Contributor
Author
|
/rerun-failed-ci |
Collaborator
|
@amd-bot ci-status |
yuychang
force-pushed
the
k3-pr1-quark-weights
branch
2 times, most recently
from
October 21, 2026 15:42
250b58b to
ee938d5
Compare
Contributor
Preview approval requiredThis pull request comes from a fork, so Mintlify won't build a preview until someone with write access approves commit |
Contributor
Author
|
/rerun-failed-ci |
1 similar comment
Contributor
|
/rerun-failed-ci |
3 of 5 tasks
yuychang
force-pushed
the
k3-pr1-quark-weights
branch
from
October 29, 2026 02:50
49e7c45 to
3fe31bf
Compare
This comment has been minimized.
This comment has been minimized.
4 of 5 tasks
…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
force-pushed
the
k3-pr1-quark-weights
branch
from
November 2, 2026 06:01
c790c48 to
39f747c
Compare
# Conflicts: # python/sglang/kernels/ops/attention/kda_fused_decode_gfx950.py
yuychang
marked this pull request as draft
November 3, 2026 05:58
1 of 5 tasks
…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
marked this pull request as ready for review
November 5, 2026 17:40
yuychang
requested review from
hanming-lu,
hebiao064 and
yizhang2077
as code owners
November 5, 2026 17:40
…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
5 tasks done
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Motivation
f_bprojection + convolution + recurrent update + gated RMSNorm in one gfx950 launch per KDA layer.aiterbackend 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:SGLANG_K3_FLYDSL_SOURCE=aiterimports a module no AITER build ships; ROCm/aiter#6026 publishes the kernel under another name--mamba-ssm-dtype bfloat16pool is BF16f_bf_bas per-channel FP8; the kernel takes dense BF16os.environ, outsideEnvs.Modifications
flydsl_kda_decode_with_f_bfromaiter.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.state_dtypes); otherwise warn once and keep the unfused path.f_bto a BF16 copy once, for the fused kernel only; the FP8 layer keeps serving the unfused path.covered(), so an unsupported layout falls back instead of raising.f_bon the ReplaySSM decode-ring path too (it has no fused kernel); before, that path received the rawf_a.Arming rules:
gfx950(default)gfx950aiter, sourceauto/sglang(vendored kernel)aiter, sourceauto/sglangaiter, sourceaiteraiter, sourceaiterf_balready post-processed (ShardedStateLoader,RemoteModelLoader)Env vars (moved into
Envs; the old names keep working with aDeprecationWarning):SGLANG_USE_HIP_K3_KDA_FUSED_BACKENDSGLANG_ROCM_K3_KDA_FUSED_BACKENDgfx950gfx950: SGLang HIP kernel;aiter: FlyDSL boundary; other values: unfused decodeSGLANG_USE_HIP_K3_FLYDSL_SOURCESGLANG_K3_FLYDSL_SOURCEautoaiter:autoprefers the vendored copy,sglang/aiterpin oneAccuracy 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.mainUnit and kernel tests (one MI355X):
test_kimi_k3_kda_decode.py6ae4bf0f2)test_kimi_k3_kda_decode.py285579b62(no #6026)test_kda_fused_decode_gfx950.pytest_kimi_k3_rocm_quant.pyf9b5308b2(has [Fix] Reduce busy polling when scheduler is idle #6026).Speed Tests and Profiling
--mamba-ssm-dtype bfloat16.mainatdd9ea06c00. PR side addsSGLANG_ROCM_K3_KDA_FUSED_BACKEND=aiter SGLANG_K3_FLYDSL_SOURCE=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 ofmainmatched its TPOT within 0.4%.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-testto run the tests your change touches, or copy this prompt for an agent: