Repository navigation
Conversation
…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.
# 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.
…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
…nel on ROCm - SGLANG_OPT_HIP_K3_KDA_FUSED_VERIFY (default off) arms flydsl_kda_verify_with_f_b with the aiter fused KDA decode - The model defers f_b and offers the o_norm gate in verify; one kernel replaces f_b GEMM, conv, recurrence and o_norm - Shapes the kernel does not cover (ragged, tree, request-indexed ring) project f_b and keep the unfused chain - Works with the BF16 state and ReplaySSM rings the AgentX recipes use; warms every wave-count build before capture - Declares the K3 handoff fields on RadixLinearAttention so the new path reads them without getattr
… verify - forward_extend sliced [1, tokens, H, K]; a deferred f_a is [tokens, K], so it cut columns and crashed graph capture - Slice the o_norm gate to the same verify rows; the parity test now goes through forward_extend with padded tokens
- Move the f_b materialization and the fused verify body into kda_fused_verify_aiter_hip; kda_backend keeps a call site - Name the token count num_draft_tokens, make available() keyword-only, and import the AITER verify op directly - Warm one batch-1 build per AITER wave count instead of max-batch snapshots; without the tag, warm batches 1 and 17 - Pass num_warps through run() and drop the max_batch plumbing and get_schedule import from kimi_k3.py - Log a rejected verify batch once per process, keeping the "K3 HIP fused KDA verify rejected a batch" prefix - Reset the fused verify arming together with the decode arming, so a failed re-arm cannot keep a stale f_b - Default the fused decode attributes on RadixLinearAttention and read them directly in kda_backend - Move the verify tests above the main block, import ForwardMode from forward_mode, and add warmup wave-count tests - Document SGLANG_OPT_HIP_K3_KDA_FUSED_VERIFY in the environment variable reference
yuychang
requested review from
BBuf,
DarkSharpness,
Edwardf0t1,
Fridge003,
HaiShaw,
Ying1123,
celve,
ch-wan,
hanming-lu,
hebiao064,
ispobock,
merrymercy,
yizhang2077 and
yuan-luo
as code owners
November 5, 2026 19:36
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
yuychang
requested review from
HydraQYH,
JustinTong0323,
sogalin,
wisclmy0611 and
zijiexia
as code owners
November 5, 2026 19:36
github-actions
Bot
added
the
documentation
Contributor
Preview approval requiredThis pull request comes from a fork, so Mintlify won't build a preview until someone with write access approves commit |
…rmed - Record the warmed draft-token count at arming; covered() rejects any other count (adaptive speculative decoding) - covered() also checks max_tokens, a BF16 conv_state, dense conv windows and the ReplaySSM ring dtypes and shapes - Skip arming and the warmup under tree speculative decoding (speculative_eagle_topk > 1) and log why once - Warm an untagged AITER op with num_warps 8 and 4 at batch 1; fall back to batches 1 and 17 only on TypeError - Log backend-side rejections (ragged layout, tree, request-indexed ring) through the once-per-process verify log - Reuse the head, dim, slot-stride constants and _dense_slots of kda_fused_decode_aiter_hip - Read _k3_deferred_f_b and _k3_onorm_gate directly in kda_backend now that RadixLinearAttention declares them - Make the warmup test pass with and without the num_warps tag; add covered() and tree-arming tests
yuychang
force-pushed
the
k3/kda-fused-verify-pr
branch
from
November 5, 2026 19:54
343e09b to
3a9117f
Compare
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_btiny GEMM,causal_conv1d_update, Triton target verify,rms_norm_gated.Modifications
environ.pySGLANG_OPT_HIP_K3_KDA_FUSED_VERIFY(defaultfalse); arms only when the aiter fused KDA decode is armed (SGLANG_USE_HIP_K3_KDA_FUSED_BACKEND=aiter)kda_fused_verify_aiter_hip.py(new)covered()contract (only the warmed draft-token count; BF16 conv cache; ring dtypes/shapes), run, warmup (one batch-1 build per wave count before graph capture)kimi_k3.pyf_b(passesf_a) and offers theo_normgate, as decode already doeskda_backend.pyf_aradix_linear_attention.py__init__Fallbacks:
K3 HIP fused KDA verify rejected a batchonceflydsl_kda_verify_with_f_bAccuracy Tests
Kimi-K3 (Quark MXFP4, attn FP8), 8x MI355X, c4 server (DSpark block 5), real acceptance, GSM8K via sgl-eval 0.1.3, 200 questions, temperature 1.0, top-p 0.95, reasoning effort low, seed 0.
Kimi-K3 fused KDA verify armed; no batch was rejected.c4f41fbce): ring off/on 0.980 / 0.980 (accept 4.11 / 4.04), snapshot off/on 0.985 / 0.980 (accept 4.05 / 4.05); every rank armed, no rejected batch.Speed Tests and Profiling
One KDA layer at TP8, MI355X, cold L2/MALL, per verify step:
InferenceX AgentX (
inferencex-agentx-mvp), Kimi-K3, 8x MI355X, pinned acceptance, 3600 s per arm, off then on, same build:Checklist
test_kimi_k3_kda_decode.py(39 passed, incl.test_aiter_fused_verify_matches_unfused_backend[ring/snapshot]against the unfused chain pluso_normwith MoE token padding) andtest_kda_fused_verify_backend.py(10 sm90/sm100 subtests fail on gfx950 with and without this change) on MI355X with ROCm/aiter#7529. pre-commit passes.🤖 Generated with Claude Code