You signed in with another tab or window. Reload to refresh your session.You signed out in another tab or window. Reload to refresh your session.You switched accounts on another tab or window. Reload to refresh your session.Dismiss alert
Adds opt-in fused FlyDSL KDA (Kimi Delta Attention) decode kernels for the Kimi-K3 TP8 layout on gfx950. They live in the existing gated-delta-rule (GDR) decode files.
New APIs, exported from aiter.ops.flydsl: flydsl_kda_decode, flydsl_kda_decode_with_f_b (adds the head-local f_b projection) and is_flydsl_kda_decode_supported.
The recurrent state may be FP32 or BF16; the recurrence always computes in FP32.
Motivation
Kimi-K3 KDA decode today runs as Triton fused_kda_decode plus a separate BF16 linear for the f_b projection.
Each launch reads and writes q/k/v, the gate and the recurrent output through HBM.
A single fused kernel keeps those in LDS and removes the extra launch. It is about 2.1-2.4x faster than the Triton path at batch 1-128 (table below).
Changes
Area
Change
kernels/gdr_decode.py
create_kda_decode_kernel: one launch fuses the width-4 causal conv, the FP32 state update with per-channel decay, and RMSNorm/sigmoid gating. fuse_f_b also computes f_b in FP32 and keeps it in LDS
linear_attention_kernels.py
Wrappers with full input validation (shapes, dtypes, devices, dense non-overlapping cache slots, int32 addressing)
__init__.py
Lazy exports of the three APIs
op_tests/test_flydsl_linear_attention.py
158 KDA cases plus a gfx950 benchmark table with a Triton baseline
Contract:
Input
Behavior
Negative or out-of-range state_indices
zero output; caches untouched
Slot 0
live, as in Triton fused_kda_decode
Caches > 4 GiB
supported (slot rebase in int64)
conv_state
channel-major, or SGLang's width-major view
Duplicate live slots
undefined; SGLang guarantees uniqueness
Env var
Default
Effect
AITER_KDA_DECODE_SLOT_CHECK
0
1 raises ValueError on duplicate live slots (host sync; skipped during graph capture)
Since the last revision:
The debug slot check now checks capture on the input's device, and its docs say it also flags SGLang's eager DP padding rows that share slot 0.
A conv_weight with no unit-stride axis is rejected before compiling, instead of failing inside FlyDSL.
The compiled-specialization cache is sized like the GDR caches (1024), so a captured CUDA graph never outlives its kernel module.
Three redundant shuffle_idx broadcasts are removed from the kernel. Results are unchanged and performance is within ±2%.
Performance (if applicable)
Kernel time, Triton fused_kda_decode (plus a BF16 linear for the projection) vs FlyDSL. Author run: 80a3b0b09 vs 301f09d83, median of 7 alternating graph-replay rounds.
Hardware
Workload/configuration
Before: Triton (us), decode / proj+decode
After: FlyDSL (us), decode / proj+decode
Change
MI355X
batch 1
13.69 / 17.04
6.40 / 7.28
2.14× / 2.34×
MI355X
batch 32
27.55 / 31.97
12.95 / 15.11
2.13× / 2.12×
MI355X
batch 128
90.67 / 94.56
37.43 / 43.85
2.42× / 2.16×
MI355X
batch 128 decode, FlyDSL FP32 vs BF16 state (in-tree run_perftest, channel-major)
52.59
34.13
about a third less
Re-run of the FlyDSL kernels at the new head (MI355X, HIP device 7, CUDA-graph replay, median of 3 runs, us). The Triton side was not re-run:
State
batch 1 decode / f_b
batch 32 decode / f_b
batch 128 decode / f_b
FP32
6.26 / 7.31
11.16 / 14.41
37.11 / 44.44
BF16
6.52 / 7.62
10.01 / 13.15
25.14 / 34.70
Testing
Unit tests added or updated
Performance benchmarks run
Tested on MI250X
Tested on MI300X
Documentation
Docstrings updated
User guide updated
Performance guide updated
No documentation change is needed, with the reason stated below
Dependencies
No new third-party dependencies added
New dependencies are listed and justified below
Breaking Changes
None. The KDA APIs are new and opt-in, and the GDR decode APIs are unchanged.
PR Self-Review
Reviewed revision
cc96f6ff5feb6df57751872f69a4fb27064a5cf7
validate-kernel-pr PASS at this head (209/209 repo tests, 225/225 independent batch grid incl. 3/33/256, perf ratio 1.0014); report summary in the PR comments. Since efa1d4a94: graph-replay test (b0301d375), main merge (849dfcae2), batch-independent mismatch allowance (cc96f6ff5).
Behavior and intent
Adds an opt-in fused FlyDSL KDA decode path, gated by is_flydsl_kda_decode_supported(device) (gfx950 only).
No existing op changes behavior.
Validation and unverified scope
At c401182ab, MI355X (gfx950), torch 2.11.0+rocm7.2, FlyDSL 0.4.0, HIP device 7:
validate-kernel-pr on that file (base 699d75091 + patch)
merge_sim, gpu_claim, runtime_compat, test_policy and correctness pass (197/197). Verdict INCONCLUSIVE only because no shape grid or route was supplied
Same test file on the merge base
ImportError: cannot import name 'flydsl_kda_decode' (author run)
3-way merge into current main (f60ee45dd)
clean
Model-level accuracy: not shown in this PR. No SGLang end-to-end run uses these kernels yet.
Unverified: SGLang end-to-end; slot uniqueness is unchecked by default; MI300X/MI250X are unsupported (tests skip); owner sign-off is pending.
PR-specific question
What behavior would each changed op test fail on before this PR, and why is that coverage not already provided by an existing canonical test?
Before this PR the file fails at collection: the KDA APIs do not exist on main.
The KDA cases check numerics against an FP32 reference, cache safety for padded and out-of-range slots, input rejection and the duplicate-slot check.
The canonical KDA tests (test_fused_kda_decode*.py) run only Triton. This file's existing tests cover only GDR/GDN.
Changed op tests
op_tests/test_flydsl_linear_attention.py: 158 KDA cases and a gfx950-only benchmark table, beside the unchanged GDR/GDN cases that share the kernel file.
All standard extended tests (excludes ci:atom_full)
Only add ci:atom_full for FlyDSL or Triton upgrades.
Add labels via the sidebar or gh pr edit 6026 --add-label <label>
One backend per PR:
A PR changes one kernel backend: [Triton/Gluon] (Triton and Gluon count as one), [HIP], [ASM], [CK], [OPUS] or [FlyDSL]. If the title ends up with two backend tags, split the PR -- as stacked pull requests when one part cannot merge without the other.
PR title tags & labels:
Component tags ([Triton/Gluon], [HIP], [CK], [ASM], ...) are added to the PR title and as PR labels automatically from the changed files and re-synced on every push — change-type tags like [fix]/[Perf], op tags like [MLA], and human labels (ci:*) are left untouched. Add the no-auto-title label to stop the title rewrites; labels stay in sync either way.
github-actionsBot
changed the title
feat(flydsl): add fused KDA decode kernels on gfx950
[FlyDSL] feat(flydsl): add fused KDA decode kernels on gfx950
Oct 1, 2026
…decode
- Match ROCm/aiter#6026, which upstreams the kernels as aiter.ops.flydsl.kda_decode
- Load the AITER source from aiter.ops.flydsl.kda_decode and its flydsl_kda_decode* entry points
- Rename the vendored module to kda_flydsl/kda_decode.py so both sources expose one API
- Rename the kernel builders to kda_decode_gfx950.py and kda_decode_fused_projection_gfx950.py
- Rename the public functions, builders, compiled kernel names and registry op ids to the AITER spelling
- Rename the FlyDSL KDA decode test, its test cases and benchmark to the operation-based names
- Keep the vendored batch-2 fused-projection tuning as optional builder keywords
This projection variant has the same unbounded specialization cache keyed by caller-provided floats. Distinct norm_eps/lower_bound pairs permanently retain JIT launchers and compiled modules, allowing memory to grow with workload diversity. Please make the scalars runtime arguments or cap this cache with a workload-appropriate LRU size.
Unbounded specialization cache leaks memory for varying float parameters
This unbounded cache is keyed by the public norm_eps and lower_bound floats, so every distinct pair retains another JIT launcher/compiled module for the process lifetime. A caller serving models with varying scalar values can therefore grow host/code-object memory without limit. Please either pass these scalars as runtime kernel arguments or use a deliberately bounded LRU cache.
The reason will be displayed to describe this comment to others. Learn more.
🔵 Needs a closer look
The new cache-mutating GPU kernels need final human review and confirmation that the reported gfx950 correctness and performance results match the current head.
The reason will be displayed to describe this comment to others. Learn more.
🔵 Needs a closer look
The precision-sensitive kernels and in-place cache updates warrant human review because the reported gfx950 correctness and performance results were not independently verified here.
The reason will be displayed to describe this comment to others. Learn more.
🔵 Needs a closer look
The new gfx950 kernels and BF16 state path require target-GPU correctness and paired performance validation, which this source-only environment cannot perform, and no independently verified head-matched report was supplied.
The reason will be displayed to describe this comment to others. Learn more.
🔵 Needs a closer look
The gfx950 kernel and BF16 recurrent-state changes require head-matched target-GPU evidence for numerical correctness, cache safety, and latency that this source-only review could not independently verify.
The reason will be displayed to describe this comment to others. Learn more.
🟡 Changes recommended
The KDA sweep gates on the live GPU rather than the configured test architecture, allowing gfx950 measurements to be mislabeled as gfx942 results when GPU_ARCHS=gfx942.
The reason will be displayed to describe this comment to others. Learn more.
🔵 Needs a closer look
The precision-sensitive gfx950 kernel changes require human review of exact-head GPU correctness and performance evidence that could not be independently validated in the dependency-free review environment.
- Add create_kda_decode_kernel to kernels/gdr_decode.py: KDA is the gated delta rule with a per-channel decay
- Fixed to the Kimi-K3 TP8 layout on gfx950: 12 local heads, 128-dim state, 256-thread blocks
- One launch fuses the width-4 Q/K/V causal conv, L2 norm, FP32 recurrent-state update and gated RMSNorm output
- fuse_f_b selects a variant that computes the head-local f_b projection in FP32 and keeps the BF16 gate in LDS
- The two variants share one kernel body; they differ in gate source, LDS layout and waves-per-EU (3 vs 2)
- state_dtype "f32" or "bf16": BF16 state is widened on load, updated in FP32 and rounded back on store
- State moves in 16-byte transactions (f32x4 or bf16x8); the FP32 kernels keep names without a dtype suffix
- Negative or out-of-range state indices zero the output and leave both caches untouched; slot 0 is live
- Cache slots are rebased with 64-bit byte offsets, so state and conv caches larger than 4 GiB work
- Specializations are cached per (norm_eps, lower_bound, fuse_f_b, state_dtype) and reuse the module's _LOG2E
…ps.flydsl
- Add flydsl_kda_decode and flydsl_kda_decode_with_f_b to linear_attention_kernels.py, next to flydsl_gdr_decode
- Add is_flydsl_kda_decode_supported(device): true for a visible device index when get_gfx_runtime() reports gfx950
- conv_state may be channel-major [cache, 4608, 3] or SGLang's width-major [cache, 3, 4608] seen via transpose(-1, -2)
- Validate shapes, dtypes, devices, dense non-overlapping cache slots, out row layout and int32 buffer addressing
- Cache slot strides are passed as int64; slots beyond min(state, conv_state) capacity are treated as padding
- Accept an FP32 or BF16 state and advertise it as state_dtypes; set zeroes_invalid_output = True on both APIs
- Accept a preallocated out buffer, as SGLang's KDA backend passes out=
- AITER_KDA_DECODE_SLOT_CHECK=1 rejects duplicate live state_indices; read once at import, skipped in graph capture
- Document that live slots must be unique and that the fused f_b projection pays off only at decode batch sizes
- Export the three APIs lazily from aiter.ops.flydsl
…n baseline
- Extend op_tests/test_flydsl_linear_attention.py next to the GDR tests; KDA tests skip unless the device is gfx950
- Check both APIs against an FP32 torch reference for both conv layouts at batch 1/8/16/64/128, FP32 and BF16 state
- Check BF16 state at batch 1/2/8/64/128 with dense or padded slots and over recurrent sequences; reject fp16/fp64
- Cover negative, out-of-range and slot-0 indices, the opt-in duplicate-slot check and >4 GiB cache slot offsets
- Cover input validation, unsafe and padded out layouts, int32 address overflow bounds and the public exports
- test_flydsl_kda_decode times flydsl, flydsl_f_b and Triton fused_kda_decode (width-major only) into a preallocated out
- Batches above 16 allow a 1e-3 mismatch ratio for one-ULP BF16 state rounding ties; conv state must match exactly
- main() adds --conv-layout and --state-dtype axes and a KDA summary table, run only when get_gfx() is gfx950
… broadcasts
- Check AITER_KDA_DECODE_SLOT_CHECK's capture state on the input's device, not the current one.
- Document that the debug slot check also flags SGLang's eager DP padding rows sharing scratch slot 0.
- Reject a conv_weight with no unit-stride axis before compiling; FlyDSL's tensor argument needs one.
- Size the KDA specialization cache like the GDR caches (1024), so captured graphs never outlive their module.
- Drop three shuffle_idx broadcasts: the xor butterfly already leaves the K-group sum in all eight lanes.
- MI355X graph replay, batch 1/32/128 with FP32 and BF16 state: within +-2% of the previous kernel.
- Add a test for the conv_weight stride check; 197 tests pass with FLYDSL_RUNTIME_ENABLE_CACHE=0.
…ion sweep
- aiter_test.sh ran op_tests/test_flydsl_linear_attention.py only as a script, which runs the benchmark sweep.
- Run pytest and the sweep for that file and fail the step if either fails.
The reason will be displayed to describe this comment to others. Learn more.
🔵 Needs a closer look
The gfx950 fused recurrent kernels require target-GPU numerical and graph-replay validation that this environment cannot run, and no exact-head validation report was supplied.
…it for bit
- Capture both APIs (plain and fused f_b) with FP32 and BF16 state, rewrite inputs in place, replay three steps.
- Output, recurrent state and conv state must equal an eager run exactly; a stale captured f_a fails the fused cases.
The reason will be displayed to describe this comment to others. Learn more.
🔵 Needs a closer look
The new gfx950 recurrent kernels require exact-head GPU validation of compilation, numerics, graph replay, and latency, which this environment cannot perform and earlier-revision results do not establish.
The reason will be displayed to describe this comment to others. Learn more.
🔵 Needs a closer look
The new gfx950 kernels and graph-replay paths lack GPU correctness and performance validation under pinned FlyDSL 0.4.1, since reported runs used 0.4.0 and this review environment cannot execute GPU tests.
… cover odd batches
- An independent grid failed batch 3: a BF16 rounding tie moved 23 of 983040 FP32 state elements by <= 4.9e-6.
- Ties occur at any batch (worst 6.9e-5 over batch 1-256 on MI355X), so use _TOL_ERR_RATIO for every batch.
- Parametrize batch 3 and 33; the pytest check skips the Triton baseline, which stays in the benchmark sweep.
Exact-head validation for the open review items (gfx950 numerics, graph replay, latency):
validate-kernel-pr: PASS at head cc96f6ff5feb6df57751872f69a4fb27064a5cf7 (base main @ d7fad35c8, patch applied on a clean base worktree, 1x MI355X verified idle).
aiter.ops.flydsl.linear_attention_kernels:flydsl_kda_decode observed at every grid batch
Perf, head vs base on one locked GPU
pass, median ratio 1.0014 over 56 rows
Index-width scan
info: 16 index/stride expressions without explicit 64-bit widening; Kimi-K3's state pool (~192 slots x 12 x 128 x 128 = 38M elements) stays far below 2^31, and the large-cache addressing regression covers the slot stride
Changes since the last review:
b0301d375: test_flydsl_kda_decode_graph_replay captures both APIs (plain and fused f_b, FP32 and BF16 state), rewrites inputs in place and replays 3 steps; output, recurrent state and conv state must equal eager bit for bit. A stale captured f_a makes the fused cases fail.
849dfcae2: merged main (conflict in .github/scripts/aiter_test.sh, both cases kept).
cc96f6ff5: an earlier grid run failed batch 3 by 23 of 983040 FP32 state elements (<= 4.9e-6, a BF16 rounding tie). A batch 1-256 sweep shows ties at any batch (worst 6.9e-5), so the mismatch allowance is now the same at every batch, and batch 3 and 33 are parametrized. The pytest check no longer asserts the Triton baseline, which stays in the benchmark sweep.
Note: GitHub has not synced this PR's head yet (it still shows 849dfcae2); the fork branch is at cc96f6ff5.
- is_flydsl_kda_decode_supported(device) range-checked device but returned the cached first-GPU arch.
- Read gcnArchName of the device asked about, and require FlyDSL's build target (GDR_GPU_ARCH) to be gfx950 too.
- Memoize the answer per device index, so eager calls skip the get_device_properties query.
- Test a mocked gfx942 device, a non-gfx950 build target, the per-device memo and an out-of-range index.
…s-outer
- The out check required stride(1) >= 11 * stride(2) + 128, so it rejected valid heads-outer [H, B, D] storage.
- Add _check_disjoint_rows: by increasing stride, each row dim must step past what the narrower dims span.
- Batch-outer layouts, padded rows and batch 1 keep their previous answer.
- Test padded heads-outer outputs against the reference, plus two overlapping heads-outer rejects.
…padding
- SGLang's eager DP-attention padding rows all map to its reserved mamba slot 0, so the opt-in check always fired there.
- AITER_KDA_DECODE_SLOT_CHECK_PADDING_SLOT=<slot> lets that slot repeat like a negative index; unset keeps it live.
- Mask the slot at the decode call site, so _check_unique_live_slots and its signature stay unchanged.
- The kernel still decodes the slot; only the debug check changes. A non-integer value fails at import.
- Test the default, the opt-in, a repeated live slot alongside it, and the env parsing.
- conv_weight, A_log and dt_bias must be FP32, norm_weight and activations BF16; the validation already enforces it.
- Note that the Triton fused_kda_decode also takes BF16 conv_weight and FP32 norm_weight, so callers cast once at load.
- _kda_reference skipped only negative slots; the kernel also pads slots >= min(state, conv_state) slot counts.
- Skip those too, so the reference states the documented contract instead of indexing past the cache.
- The out-of-range test now mixes in a live row and checks output and both caches against the reference.
- Seven KDA tests bound the generated KdaInputs to `seed`, next to helpers whose `seed` is the RNG integer.
- Rename them to `base_inputs`; no behaviour change.
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
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.
github-actions
Bot
changed the title
Summary
aiter.ops.flydsl:flydsl_kda_decode,flydsl_kda_decode_with_f_b(adds the head-localf_bprojection) andis_flydsl_kda_decode_supported.Motivation
fused_kda_decodeplus a separate BF16 linear for thef_bprojection.Changes
kernels/gdr_decode.pycreate_kda_decode_kernel: one launch fuses the width-4 causal conv, the FP32 state update with per-channel decay, and RMSNorm/sigmoid gating.fuse_f_balso computesf_bin FP32 and keeps it in LDSlinear_attention_kernels.py__init__.pyop_tests/test_flydsl_linear_attention.pyContract:
state_indicesfused_kda_decodeconv_stateAITER_KDA_DECODE_SLOT_CHECK01raisesValueErroron duplicate live slots (host sync; skipped during graph capture)Since the last revision:
conv_weightwith no unit-stride axis is rejected before compiling, instead of failing inside FlyDSL.shuffle_idxbroadcasts are removed from the kernel. Results are unchanged and performance is within ±2%.Performance (if applicable)
Kernel time, Triton
fused_kda_decode(plus a BF16 linear for the projection) vs FlyDSL. Author run:80a3b0b09vs301f09d83, median of 7 alternating graph-replay rounds.run_perftest, channel-major)Re-run of the FlyDSL kernels at the new head (MI355X, HIP device 7, CUDA-graph replay, median of 3 runs, us). The Triton side was not re-run:
Testing
Documentation
Dependencies
Breaking Changes
None. The KDA APIs are new and opt-in, and the GDR decode APIs are unchanged.
PR Self-Review
Reviewed revision
cc96f6ff5feb6df57751872f69a4fb27064a5cf7validate-kernel-prPASS at this head (209/209 repo tests, 225/225 independent batch grid incl. 3/33/256, perf ratio 1.0014); report summary in the PR comments. Sinceefa1d4a94: graph-replay test (b0301d375), main merge (849dfcae2), batch-independent mismatch allowance (cc96f6ff5).Behavior and intent
is_flydsl_kda_decode_supported(device)(gfx950 only).Validation and unverified scope
At
c401182ab, MI355X (gfx950), torch 2.11.0+rocm7.2, FlyDSL 0.4.0, HIP device 7:FLYDSL_RUNTIME_ENABLE_CACHE=0 python3 -m pytest -q op_tests/test_flydsl_linear_attention.pyvalidate-kernel-pron that file (base699d75091+ patch)ImportError: cannot import name 'flydsl_kda_decode'(author run)main(f60ee45dd)Model-level accuracy: not shown in this PR. No SGLang end-to-end run uses these kernels yet.
Unverified: SGLang end-to-end; slot uniqueness is unchecked by default; MI300X/MI250X are unsupported (tests skip); owner sign-off is pending.
PR-specific question
What behavior would each changed op test fail on before this PR, and why is that coverage not already provided by an existing canonical test?
main.test_fused_kda_decode*.py) run only Triton. This file's existing tests cover only GDR/GDN.Changed op tests
op_tests/test_flydsl_linear_attention.py: 158 KDA cases and a gfx950-only benchmark table, beside the unchanged GDR/GDN cases that share the kernel file.