[FIX] Fix atom_set_value signature mismatch and ROCDL compat - #381
Closed
coderfeli wants to merge 2 commits into
Closed
[FIX] Fix atom_set_value signature mismatch and ROCDL compat#381coderfeli wants to merge 2 commits into
coderfeli wants to merge 2 commits into
Conversation
Two fixes:
1. primitive.py: atom_set_value() was calling fly.atom_set_value(atom, field, value)
but the ODS-generated binding expects (result, atom, field, value). Also, the
field string must be wrapped in StringAttr for the MLIR op to accept it.
This broke any kernel using copy_atom.set_value("soffset", ...) — including
the fused RoPE+KV cache kernel introduced in #377.
2. rocdl/__init__.py: cluster_workgroup_id_{x,y,z} are not present in the
current LLVM ROCDL dialect, causing NameError on import. Use globals().get()
to make them optional, matching the pattern used for other optional ops.
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
coderfeli
added a commit
to ROCm/aiter
that referenced
this pull request
Apr 11, 2026
Add FlyDSL-based `fused_qk_rope_reshape_and_cache` as a drop-in replacement for the Triton version, with automatic Triton fallback for unsupported features (GPT-J rotation, offsets, KV scaling, zeros output, full-dim cos/sin). ## Performance **1.45x average speedup** over Triton across LLM model configs: - Llama-3.1-8B/70B/405B (TP1, TP8) decode + prefill - Qwen3-72B (TP8) decode ## Validation Tests aligned with Triton `test_fused_qk_rope_reshape_and_cache`: - Same input generation (`generate_rope_inputs`) and reference (`ref_rope_sbhd_fwd`) - Same tolerance (atol=0.1, rtol=0.1) - Same KV cache validation (slots + full tensor) - 66 test cases: Category A (direct kernel) + Category B (wrapper fallback) ## Changes (5 files, +1141 lines) - `aiter/ops/flydsl/rope_kernels.py` — wrapper with auto Triton fallback - `aiter/ops/flydsl/kernels/fused_rope_cache_kernel.py` — vendored FlyDSL kernel - `aiter/ops/flydsl/__init__.py` — export - `op_tests/flydsl_tests/test_flydsl_rope.py` — tests - `pyproject.toml` — flydsl >= 0.1.3 ## Dependencies - FlyDSL >= 0.1.3 with ROCm/FlyDSL#381 fix Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
3 tasks
coderfeli
force-pushed
the
fix/atom-set-value-and-rocdl-compat
branch
from
April 11, 2026 02:39
abd1ef0 to
efe49c2
Compare
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
1 participant
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.
Summary
Two fixes for issues discovered while integrating FlyDSL fused RoPE kernel into AITER:
1.
atom_set_valuesignature mismatch (primitive.py)Commit #377 changed
atom_set_valueto callfly.atom_set_value(atom, field, value), but the ODS-generated binding still expects(result, atom, field, value). Additionally, thefieldstring must be wrapped inir.StringAttrfor the MLIR op to accept it.This breaks any kernel using
copy_atom.set_value("soffset", ...), including the fused RoPE+KV cache kernel.2. Missing ROCDL ops (
rocdl/__init__.py)cluster_workgroup_id_{x,y,z}are referenced at module level but not present in the current LLVM ROCDL dialect, causingNameErroron import. Changed toglobals().get()to make them optional, matching the pattern used forwmma_scale_f32_*.Test
🤖 Generated with Claude Code