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
{{ message }}
Repository navigation
[Triton/Gluon] [FlyDSL] [CI] Refactor PA decode, direct native autotune, and BF16 KV - #5809
Refactor FlyDSL PA decode into operation classes with explicit GPU work plans, add native BF16 KV, and separate tuning from decode execution in aiter/ops/flydsl/pa_decode_tuning.py.
Separate scheduling/traits, context, memory movement, GEMM, softmax, epilogue, and pipeline components. Retain packed work plans, query splitting, K/V prefetch scheduling, and compact/window-bounded reductions.
Require an explicit work_plan for decode and graph replay. Prepare or refresh plans before execution; max_context_length supplies a host scheduling bound independently of block-table padding. Decode does not import or invoke the tuning module.
Provide a standalone tuning CLI with batch/head dimensions, query and context lengths, page size, FP8/BF16 K/V, scale mode, V layout, sliding windows, sinks, GPU selection, and output-directory options. One invocation searches the requested shape and saves the native configuration while retaining results for other shapes.
Keep native FlyDSL candidate benchmarking, fastest-config selection, validation hooks, and JSON persistence. Deduplicate equivalent capacities and compare candidates with the 2*CU baseline. Candidate allocation, compilation, and graph capture happen before timed replay.
Export prepare_pa_decode_plan from aiter.ops.flydsl.pa_decode_tuning and aiter.ops.flydsl. Ordinary cache misses use 2*CU without preparing or timing candidates; tune=True explicitly searches without modifying environment variables. config_dir selects per-architecture result directories.
Remove tuning pytest cases and their environment/function monkeypatch helpers from op_tests/test_flydsl_pa_decode.py. Retain decode correctness, invalid-input, work-plan, and graph-replay coverage.
Results are saved to /tmp/pa-configs/<arch>/_pa_decode_autotuner.json. Without --output-dir, use aiter/ops/flydsl/configs/pa/<arch>/, or the native FLYDSL_AUTOTUNE_CACHE_DIR override. Python callers can use the same result root with config_dir="/tmp/pa-configs".
Example with actual BF16 query/KV tensors and page64/D192:
BF16 KV follows ROCm/FlyDSL#1065: require BF16 queries, unscaled vector8 K/V, and compute_type=torch.bfloat16. Support page sizes 16/64/128 and BF16 head dimensions in multiples of 64 through 1024, including D192. Use BF16 Q/K/P/V operands, FP32 softmax/accumulation, and query-position splitting to fit 64 KiB LDS. Allocate caller-owned graph scratch from plan.capacity; refresh plans when context lengths change.
Cache keys include attention geometry, execution modes, CU count, PA kernel/tuning source contents, and native toolchain/environment/device fingerprints. Tensor strides, allocated KV pages, block-table padding, scalar scale values, and max_context_length do not partition the tuning cache. Restart inference processes after updating cache files and use separate directories for parallel tuning processes.
Validation for the standalone-script extraction: Ruff, Black, Python syntax compilation, and diff checks pass. Tests and GPU tuning were not run for this follow-up, as requested.
The previously observed FP32-reference failure for varlen_b81_q4_c1027 occurred on both monolithic and refactored kernels (maximum absolute error 0.00712435 with atol=rtol=0.005). This extraction does not change its numerical behavior.
Use the unaligned MTP window union to bound reducer specialization,
while retaining the CU-count plan cap, scratch layout and query selector.
Cover partition-boundary crossings during refresh in the generic test.
Validated on dzm-mi355-gpu34-dev: 196 unit tests passed, with two
70-case native accuracy/performance sweeps. W1024/B200/QL4 improves
from 80.18 us to 29.14 us, excluding plan construction and refresh.
Signed-off-by: fsx950223 <fsx950223@outlook.com>
Include pending varlen KV buffer loads, prefetch scheduling and compact reduction. Bound dense planned work for native QL2/QL4 query splitting and cover the selector and address boundaries in the generic unit test.
Add offline model/batch budget sweeps with fixed-plan timing, strict FP32 validation and static recommendation keys. Keep context distributions in benchmark metadata only, and preserve failed numerical cases without recommendations.
Validate with MI355 native A/B and generic GPU tests, 22 CPU tuner tests, Black 26.5.1 and Ruff 0.16.0. The existing short-context FP8 numerical limitation is not changed.
Signed-off-by: fsx950223 <fsx950223@outlook.com>
…de-oop
Preserve planned query splitting, unsigned KV buffer loads and prefetch ordering in the PA operation classes. Track the full helper source closure in offline tuning keys and test persistent-cache separation for wide/buffer specializations.
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 5809 --add-label <label>
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 opt this PR out.
fsx950223
changed the title
[FlyDSL] Refactor PA decode and add offline budget tuning
[FlyDSL] Refactor PA decode and load CSV tuning results
Sep 24, 2026
Always normalize Q rows by absmax and scale P before FP8 conversion, compensating both scales in the QK and output stages. This prevents small queries and probability tails from underflowing on the per-tensor KV path.
Add four regression cases with quantized FP8 KV caches and non-unit scales. Validated on MI355: all 264 native PA decode tests pass.
Use one parameterized correctness test with shared input construction and an independent FP32 reference. Cover query dtypes, FP8 KV quantization, cache layouts, partitioning, plans, windows and sinks, including the small-Q and probability-tail regressions.
Remove the benchmark CLI, tuning CSV tests and internal scheduling assertions from this file. Run the unit test once through pytest in CI.
Validated on MI355: all 80 cases pass. The four original scaling regressions still fail against the old shortcut implementation.
github-actionsBot
changed the title
[FlyDSL] Refactor PA decode and load CSV tuning results
[FlyDSL] [CI] Refactor PA decode and load CSV tuning results
Sep 28, 2026
Remove static tile dispatch, direct output, uniform scratch layouts, and
the static reducer branches. All FlyDSL decode calls now use work plans,
including a partition cap of one, with sinks applied once by the reducer.
Configure partition caps on plan_pa_decode and make the optional decode
arguments keyword-only. Update the tuner and generic unit test, and adapt
Gluon's existing FlyDSL reducer fallback to provide reduction metadata.
Validate on MI355 with 80 generic cases, six Gluon reducer cases, tuner
smoke, and ten graph/eager replays. All replay outputs match the previous
planned implementation bit for bit, including its two-token FP8 error.
github-actionsBot
changed the title
[FlyDSL] [CI] Refactor PA decode and load CSV tuning results
[Triton/Gluon] [FlyDSL] [CI] Refactor PA decode and load CSV tuning results
Sep 28, 2026
Use the existing C++ PS reducer first and fall back directly to Triton.
Remove FlyDSL imports, dispatch checks, and reduction-plan metadata from
the Gluon wrapper.
This new module uses a relative import, contrary to the repository's absolute-import contract for FlyDSL code. Import the traits module through its absolute aiter.ops.flydsl... path.
Replace FlyDSL relative imports with absolute paths
These new imports use relative paths, contrary to the repository's absolute-import convention for FlyDSL code. Switch them to the corresponding fully qualified aiter.ops.flydsl... paths.
This new operation module uses a relative import, contrary to the repository's absolute-import contract for FlyDSL code. Import the traits module through its absolute aiter.ops.flydsl... path.
Replace FlyDSL relative imports with absolute paths
These new imports use relative paths, contrary to the repository's absolute-import convention for FlyDSL code. Switch them to the corresponding fully qualified aiter.ops.flydsl... paths.
Replace FlyDSL relative imports with absolute paths
These new imports use relative paths, contrary to the repository's absolute-import convention for FlyDSL code. Switch them to the corresponding fully qualified aiter.ops.flydsl... paths.
The new PA decode subpackage uses relative imports here, contrary to the repository's absolute-import convention for FlyDSL modules. Use the fully qualified aiter.ops.flydsl... paths so this refactor follows the established import boundary.
Use absolute imports in Triton/FlyDSL wrappers
aiter/ops/flydsl/pa_decode.py:30
These new modules introduce relative imports, but the repository's Triton/FlyDSL wrapper convention requires absolute aiter.ops... imports. Please switch this import block to absolute paths so the new package does not add another import style.
New optional query_start_loc ([num_seqs + 1], int32) for pa_decode_gluon,
PS path only. Sequence i owns packed query/output rows
[query_start_loc[i], query_start_loc[i + 1]); query_length must be >= every
per-sequence length; rows not owned by any sequence are left untouched.
Non-PS paths assert on it.
- Supported in the gluon head_1 and multi-KV attention kernels and in the
Triton PS reduce, each behind a compile-time IS_VARLEN flag; the HIP PS
reduce is used through its query_start_loc argument.
- Causal masks use the per-sequence query length, on top of the per-split
causal mask from ROCm#5915.
- Fix the default temporary_output dtype: the attention kernels store bf16
for fp8 compute (otherwise compute_type), not query.dtype. An fp8 query
with fp8 compute and no temporary_output failed to compile before.
- New test op_tests/triton_tests/test_pa_decode_gluon_varlen.py: each
sequence of a ragged call must match a batch-1 uniform call, and unowned
rows must keep a sentinel. Covers bf16, fp8 KV, fp8 Q+KV, sinks, sliding
window, block 16/64/128 with shuffled V, head_1 and multi-KV, one-shot and
split runs through the HIP and Triton reduces, and zero-length sequences.
Split out of ROCm#6099, rebased onto main. The FlyDSL reduce part of ROCm#6099 is
dropped: since ROCm#5809 the Gluon path no longer dispatches to the FlyDSL PS
reduce, and the no-work-plan mode it extended no longer exists.
Signed-off-by: Sergei Shliapnikov <sergei.shliapnikov@amd.com>
Signed-off-by: Andrei Kochin <andrei.kochin@amd.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
Transplant fp8-Q pa_decode onto the modular #5809 package (single
query-scale-aware launcher, output dtype in schedule key), migrate
fmha op_lds BT/V LDS access to recast_iter/add_offset/ptr_load/ptr_store,
and route the adapter decode path through a retained geometry-and-stream
keyed work plan with capture-safe refresh.
sshlyapn
added a commit
to andrei-kochin/aiter
that referenced
this pull request
Oct 15, 2026
New optional query_start_loc ([num_seqs + 1], int32) for pa_decode_gluon,
PS path only. Sequence i owns packed query/output rows
[query_start_loc[i], query_start_loc[i + 1]); query_length must be >= every
per-sequence length; rows not owned by any sequence are left untouched.
Non-PS paths assert on it.
- Supported in the gluon head_1 and multi-KV attention kernels and in the
Triton PS reduce, each behind a compile-time IS_VARLEN flag; the HIP PS
reduce is used through its query_start_loc argument.
- Causal masks use the per-sequence query length, on top of the per-split
causal mask from ROCm#5915.
- Fix the default temporary_output dtype: the attention kernels store bf16
for fp8 compute (otherwise compute_type), not query.dtype. An fp8 query
with fp8 compute and no temporary_output failed to compile before.
- New test op_tests/triton_tests/test_pa_decode_gluon_varlen.py: each
sequence of a ragged call must match a batch-1 uniform call, and unowned
rows must keep a sentinel. Covers bf16, fp8 KV, fp8 Q+KV, sinks, sliding
window, block 16/64/128 with shuffled V, head_1 and multi-KV, one-shot and
split runs through the HIP and Triton reduces, and zero-length sequences.
Split out of ROCm#6099, rebased onto main. The FlyDSL reduce part of ROCm#6099 is
dropped: since ROCm#5809 the Gluon path no longer dispatches to the FlyDSL PS
reduce, and the no-work-plan mode it extended no longer exists.
Signed-off-by: Sergei Shliapnikov <sergei.shliapnikov@amd.com>
Signed-off-by: Andrei Kochin <andrei.kochin@amd.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
sshlyapn
added a commit
to andrei-kochin/aiter
that referenced
this pull request
Oct 15, 2026
New optional query_start_loc ([num_seqs + 1], int32) for pa_decode_gluon,
PS path only. Sequence i owns packed query/output rows
[query_start_loc[i], query_start_loc[i + 1]); query_length must be >= every
per-sequence length; rows not owned by any sequence are left untouched.
Non-PS paths assert on it.
- Supported in the gluon head_1 and multi-KV attention kernels and in the
Triton PS reduce, each behind a compile-time IS_VARLEN flag; the HIP PS
reduce is used through its query_start_loc argument.
- Causal masks use the per-sequence query length, on top of the per-split
causal mask from ROCm#5915.
- Fix the default temporary_output dtype: the attention kernels store bf16
for fp8 compute (otherwise compute_type), not query.dtype. An fp8 query
with fp8 compute and no temporary_output failed to compile before.
- New test op_tests/triton_tests/test_pa_decode_gluon_varlen.py: each
sequence of a ragged call must match a batch-1 uniform call, and unowned
rows must keep a sentinel. Covers bf16, fp8 KV, fp8 Q+KV, sinks, sliding
window, block 16/64/128 with shuffled V, head_1 and multi-KV, one-shot and
split runs through the HIP and Triton reduces, and zero-length sequences.
Split out of ROCm#6099, rebased onto main. The FlyDSL reduce part of ROCm#6099 is
dropped: since ROCm#5809 the Gluon path no longer dispatches to the FlyDSL PS
reduce, and the no-work-plan mode it extended no longer exists.
Signed-off-by: Sergei Shliapnikov <sergei.shliapnikov@amd.com>
Signed-off-by: Andrei Kochin <andrei.kochin@amd.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
sshlyapn
added a commit
to andrei-kochin/aiter
that referenced
this pull request
Oct 16, 2026
New optional query_start_loc ([num_seqs + 1], int32) for pa_decode_gluon,
PS path only. Sequence i owns packed query/output rows
[query_start_loc[i], query_start_loc[i + 1]); query_length must be >= every
per-sequence length; rows not owned by any sequence are left untouched.
Non-PS paths assert on it.
- Supported in the gluon head_1 and multi-KV attention kernels and in the
Triton PS reduce, each behind a compile-time IS_VARLEN flag; the HIP PS
reduce is used through its query_start_loc argument.
- Causal masks use the per-sequence query length, on top of the per-split
causal mask from ROCm#5915.
- Fix the default temporary_output dtype: the attention kernels store bf16
for fp8 compute (otherwise compute_type), not query.dtype. An fp8 query
with fp8 compute and no temporary_output failed to compile before.
- New test op_tests/triton_tests/test_pa_decode_gluon_varlen.py: each
sequence of a ragged call must match a batch-1 uniform call, and unowned
rows must keep a sentinel. Covers bf16, fp8 KV, fp8 Q+KV, sinks, sliding
window, block 16/64/128 with shuffled V, head_1 and multi-KV, one-shot and
split runs through the HIP and Triton reduces, and zero-length sequences.
Split out of ROCm#6099, rebased onto main. The FlyDSL reduce part of ROCm#6099 is
dropped: since ROCm#5809 the Gluon path no longer dispatches to the FlyDSL PS
reduce, and the no-work-plan mode it extended no longer exists.
Signed-off-by: Sergei Shliapnikov <sergei.shliapnikov@amd.com>
Signed-off-by: Andrei Kochin <andrei.kochin@amd.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
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
Refactor FlyDSL PA decode into operation classes with explicit GPU work plans, add native BF16 KV, and separate tuning from decode execution in
aiter/ops/flydsl/pa_decode_tuning.py.work_planfor decode and graph replay. Prepare or refresh plans before execution;max_context_lengthsupplies a host scheduling bound independently of block-table padding. Decode does not import or invoke the tuning module.2*CUbaseline. Candidate allocation, compilation, and graph capture happen before timed replay.prepare_pa_decode_planfromaiter.ops.flydsl.pa_decode_tuningandaiter.ops.flydsl. Ordinary cache misses use2*CUwithout preparing or timing candidates;tune=Trueexplicitly searches without modifying environment variables.config_dirselects per-architecture result directories.op_tests/test_flydsl_pa_decode.py. Retain decode correctness, invalid-input, work-plan, and graph-replay coverage.Standalone tuning example:
python -m aiter.ops.flydsl.pa_decode_tuning \ --batch-size 32 --num-query-heads 16 --num-kv-heads 1 --head-dim 128 \ --query-length 4 --context-length 16384 --block-size 128 \ --kv-dtype fp8 --scale-mode per-token --trans-v 1 \ --output-dir /tmp/pa-configsResults are saved to
/tmp/pa-configs/<arch>/_pa_decode_autotuner.json. Without--output-dir, useaiter/ops/flydsl/configs/pa/<arch>/, or the nativeFLYDSL_AUTOTUNE_CACHE_DIRoverride. Python callers can use the same result root withconfig_dir="/tmp/pa-configs".Example with actual BF16 query/KV tensors and page64/D192:
BF16 KV follows ROCm/FlyDSL#1065: require BF16 queries, unscaled vector8 K/V, and
compute_type=torch.bfloat16. Support page sizes 16/64/128 and BF16 head dimensions in multiples of 64 through 1024, including D192. Use BF16 Q/K/P/V operands, FP32 softmax/accumulation, and query-position splitting to fit 64 KiB LDS. Allocate caller-owned graph scratch fromplan.capacity; refresh plans when context lengths change.Cache keys include attention geometry, execution modes, CU count, PA kernel/tuning source contents, and native toolchain/environment/device fingerprints. Tensor strides, allocated KV pages, block-table padding, scalar scale values, and
max_context_lengthdo not partition the tuning cache. Restart inference processes after updating cache files and use separate directories for parallel tuning processes.Validation for the standalone-script extraction: Ruff, Black, Python syntax compilation, and diff checks pass. Tests and GPU tuning were not run for this follow-up, as requested.
The previously observed FP32-reference failure for
varlen_b81_q4_c1027occurred on both monolithic and refactored kernels (maximum absolute error 0.00712435 withatol=rtol=0.005). This extraction does not change its numerical behavior.