[FlyDSL] [Feature] Tiered persistent radix-select decode Top-K + arch/shape dispatcher

[JH-Leon-KIM-AMD]

Tiered persistent radix-select decode Top-K + arch/shape dispatcher

Tickets: SILOTIGER-699 (gfx950 decode Top-K) ·
SILOTIGER-629 /
SILOTIGER-683 (the kernel this builds on) —
epic SILOTIGER-606 (GLM-5 MXFP4, MI355X perf uplift)

What this is

1. The FlyDSL decode Top-K kernel — context-length-independent tiered persistent radix-select
   for the sparse-attention indexer, which calls it ~78×/step. The HIP path it replaces
   (radix_topk_one_block<12,1024>) is an O(N_blocks) radix sort: 369 µs @ 8K → 1,165 µs @ 120K,
   ~15× slower than B200's persistent_topk_kernel. This one is O(topK·log topK).
2. A dispatcher inside aiter.top_k_per_row_decode, routing there only where it is measured to
   win. Without it the kernel ships unreachable: the public op still calls HIP.

Best cell vs HIP	1.42× gfx950, 2.19× gfx942
Worst in-window cell	0.61× (gfx942, k=512 rows=1 width=163840 seq=16384)
gfx942 shipped window	+1371 µs over 360 cells, median 0.976×, 169 won
Same grid, ungated	−2998 µs over 928 cells, median 0.861×
Correctness	set-equivalence with torch.topk, both arch sweeps

Ungated, this kernel loses on that grid; the gate is what turns it into a win — a routing result as
much as a kernel result.

Lineage

Sami's and Robin's commits ride along because that base is not in main yet. Base branch
RElbers/aiter@samremes/topk-per-row.

who	what	commits
@samremes	the initial FlyDSL kernel (SILOTIGER-683): K=512 graph-capture-friendly tiered persistent radix-select, runtime row_len tiers, row-barrier optimization, long-row cap	48f0a93 555569b 72c7745 5152d4d
@RElbers	made it the sole decode path for every K (256/512/1024/2048) by lifting top_k to a compile-time parameter; workspace API, batch-aware tier config, deadlock guard, gfx950 params, non-pow2 K, tier_mode refactor, AOT entry	b3a0d97 c0b6909 01741ab 755df3b c06b8db 135ebca c278797 f42bb80 4abc985 c6c30d8 ← base
@JH-Leon-KIM-AMD	gfx950 tuning (below), the dispatcher, the workspace cache, dispatch counters, the gfx942 window, the correctness harness, and the move onto aiter's buffer_ops/vector shims after #4501	below

The generic tiering serialized wide decode batches into sequential co-resident waves. The switches
that fix it are gfx950-only, since gfx942 runs the frozen configuration: short_max +
dead_block_trim + midbatch cap
(d2d2ad6 5383b16 d783300),
blocks_per_row floor 1 for the barrier-free short tier (2fc5ddc),
row_proportional_parts (a87146d),
batch_coresident_cap (197eda0),
early_stop (d839939), and a rows>63 / L≤65536
fold back to bpr=1 (e7dde59).

The gate

Five checks, cheapest first, all on host values. Per-row lengths live in seqLens on the device,
and reading them would sync every decode step.

	check	rejects to
1	AITER_DISABLE_FLYDSL_TOPK_DECODE	HIP
2	arch has a row in the gate table	HIP
3	k in ks, and numRows inside the row cap that logits.shape[1] earns	HIP
4	stride1 == 1, stride0 == logits.stride(0), next_n >= 1, logits.ndim == 2	HIP
5	flydsl importable	HIP

Step 4 is a contract check, not a filter: FlyDSL raises where HIP accepts.

# (min padded width, max rows) bands, widest first.
"gfx950": _DecodeGate(((163840, 15), (131072, 9)), ks={256, 512, 1024, 2048}),
"gfx942": _DecodeGate(((163840, 18), (131072, 11)), ks={256, 512, 1024, 2048}),

The first band whose width a call reaches sets its cap, and a call reaching none is refused. Each
cap sits two rows below its measured crossing.

graph-replay sign change	width 163840	width 131072
gfx950	rows 17 → cap 15	rows 11 → cap 9
gfx942	rows 20 → cap 18	rows 13 → cap 11

A wider buffer is more work per row, so cooperation keeps paying for more rows. One number would
refuse a wide band it wins, or admit a narrow band it loses.

Width asks "long-context model", not "long request". Callers size the buffer to the model's max
context — vLLM's indexer builds logits as (batch * next_n, max_model_len) — and under graph
capture the request length does not exist yet. Sweeping width at fixed real length moves the margin
under 2.2 µs, so the check works only because a request cannot outrun its buffer.

Thresholds come from graph replay only. Eager charges HIP 11.7 µs of host work per call against
FlyDSL's 2.6 µs, mostly an un-memoized topk_ob_workspace_size query; its kinder column (in-window
median 1.135×, 349/360 won) is that artifact.

env	effect
AITER_DISABLE_FLYDSL_TOPK_DECODE=1	every shape back to HIP
AITER_FLYDSL_TOPK_ARCHS	narrows which archs the table applies to
AITER_FLYDSL_TOPK_MAX_ROWS	replaces the cap in every band
AITER_FLYDSL_TOPK_MIN_WIDTH	also flattens the table to one band
AITER_FLYDSL_TOPK_COUNT=1	tallies flydsl vs hip per worker; an earlier PR-image A/B measured ~0% with HIP on both arms

Why the gate lives in AITER. The alternative duplicates arch detection, availability checks,
shape thresholds and the HIP fallback in every calling framework. gemm_op_a8w8.py does the same.

Host cost is the same order as the kernel win, so two things came along:

• workspace is a parameter, with a buffer cached on (device, stream, size rounded to pow2) —
  worth ~2.4 µs, and returning nothing under graph capture.
• The stream context manager around workspace.zero_() is skipped when no stream was asked for.

Measurements

	conditions
gfx950	rocprofv3 kernel-only min µs, k=2048, --boost-ms 500 --warmup 10 --iters 40
gfx942	928 cells (k × rows × width × seq), eager and graph replay, both arms through this public entry point so each pays the gate and the workspace sizing; dispatch counters checked per cell
regimes	A = uniform padded length, B = per-row random causal length, fixed seed

gfx950 vs HIP. hip/ours > 1.00 is faster. B tracks A, so the win is not a padding artifact.

rows	L	ours (µs)	hip (µs)	hip/ours (A)	hip/ours (B)
1	8K	10.0	10.1	1.01×	1.04×
4	8K	10.5	10.3	0.98×	1.04×
1	65K	21.1	30.1	1.42×	1.42×
4	65K	27.0	30.0	1.11×	1.10×
8	65K	30.0	30.6	1.02×	1.08×
1	256K	31.2	33.4	1.07×	1.05×
4	256K	42.1	55.6	1.32×	1.24×
8	256K	48.3	56.7	1.17×	1.16×

gfx950 through the dispatcher, width 163840, 200 back-to-back calls with one sync at the end.
hip / dispatch:

k	rows	seq 4,096	65,536	163,840
2048	1	0.97	1.16	1.05
2048	4	1.16	1.32	1.25
2048	8	1.19	1.22	1.17
2048	16	1.37	1.11	1.03

10 win, 1 tie, 1 marginal (0.97 is ~0.5 µs, inside run-to-run noise). The rows=16 line predates
the staircase; the cap here is 15, so it routes to HIP now.

gfx942 length curve, graph replay, median over k × rows, in the window that shipped when this
was measured (width ≥ 131072, rows ≤ 16). The advantage is a function of the real request length,
which the gate cannot see.

seq_len	4K	16K	32K	48K	64K	96K	128K	160K
HIP / FlyDSL (higher = we win)	0.80×	0.73×	0.86×	0.91×	1.09×	1.28×	1.54×	1.64×
cells won	0/48	0/48	3/48	6/48	40/48	48/48	48/48	24/24

• Profitable because the wins are larger than the losses (mean +13.3 against −4.6 µs), not because
  most cells win. Break-even needs long requests to be 25.6% of traffic.
• The staircase came later: it drops rows 12–16 at width 131072 and admits rows 17–18 at 163840,
  both measured as gains (+2.41 and +1.99 µs per cell). The shipped window scores no worse, but this
  table has not been re-run.
• gfx942 is an SR-IOV VF with no clock locking, so band noise is 2–6 µs against gfx950's 0.3 µs.
  Near-threshold cells were repeated in three independent runs.

vs Robin's gfx950 base, informational only — that base was not tuned for wide batch and this
path trails HIP there anyway: rows=1 1.15–1.19×, rows 4–8 parity, rows ≥ 16 2.0×–10.5×. Full sweep
in the SILOTIGER-699 results comment.

Correctness

Buffers are padded past the real length with the tail poisoned, so a kernel ignoring seqLens fails
loudly instead of passing by accident.

arch	result	grid
gfx950	504/504	A 252/252 + B 252/252, over K{512, 2048} × rows{1, 4, 16, 32, 64, 128, 256} × L{8K, 16K, 32K, 64K, 120K, 256K} × dist{random, ties, 10LSB}. No out-of-bounds reads
gfx942	all 928 cells, both arms, two independent runs	plus k{256, 1024} × rows{1, 4, 8} × L{131072, 163840} for the newly admitted k, and rows{12, 16} × k{256, 512, 1024, 2048} × L{4096, 70000, 131072} for the row band above the sweep's samples

Unit tests on gfx950: 195/195 (dispatcher + kernel, including the stable exact-order cases). Two of the 33 pin the staircase —
the bands descend with width, and one row count is refused at 131072 and admitted at 163840.

Run: pytest op_tests/test_topk_decode_dispatch.py,
pytest op_tests/flydsl_tests/test_flydsl_topk_per_row_decode.py,
python op_tests/topk_decode_correctness.py, and
python op_tests/benchmark_topk_per_row_decode.py under rocprofv3.

stable=True (ordered emit)

Previously stable=True forced HIP: FlyDSL returned an unordered set. This commit
adds an ordered emit to the tiered kernel and routes stable=True through
the same gate bands as the unordered path (ordered=stable forwarded to
FlyDSL).

Implementation highlights:

• One contiguous vec-block run per workgroup; run counts from pass histograms
  (s_own_hist) — no extra row read for a count pass.
• ordered=True disables the short-tier atomic scatter and early_stop
  (both assume unordered output); persistent path covers all tiers.
• Launcher cache key includes ordered (different kernel).

Correctness contract: exact index list equality, ascending with
smallest-index tie-break on the kth value — not set equivalence. Tie-dense
inputs in tests.

check	gfx950 result
pytest op_tests/flydsl_tests/test_flydsl_topk_per_row_decode.py + op_tests/test_topk_decode_dispatch.py	195/195
python op_tests/test_topk_per_row_stable.py	ALL PASS
Final bench (3 rep, all shapes)	exact failures 0

Performance (stable, k=2048, 3-rep median µs, vs #5011 ordered / HIP):

rows × L	ours	5011	ours/5011	HIP	ours/HIP
26 × 700K	126	121	1.04	933	0.14
26 × 1M	157	149	1.05	1335	0.12
28 × 700K	128	127	1.01	933	0.14
28 × 1M	159	160	0.99	1335	0.12
32 × 700K	134	128	1.05	936	0.14
32 × 1M	165	166	1.00	1336	0.12

Gate region (26–32 × 700K–1M): 5011 parity (0.99–1.05×), HIP 7–8×.
Stable overhead vs unordered: 1.12–1.18× in that region. Only regression:
rows=8 × 700K/1M ~1.11× vs 5011 (gate-outside; HIP ~92/118 µs there).

Raw: JHK_TASK/.../results_this_node/final/stable_fused_rep{1,2,3}.jsonl.

Follow-ups

	why it is left	blocker
Kernel self-reset	the barrier-counter zero_() is 3.7 µs of launch latency per call, the largest item left	HIP's multi-block kernel already self-resets; a side stream measured slower
gfx942 kernel tuning	the four gfx950 switches are frozen there; enabling them flipped rows=32 from 0.79× to 1.06× and widens the window	needs a full correctness pass
HIP multi-block for decode	the C++ decode entry is one-block only, though its own crossover says multi-block should win where FlyDSL does	graph capture freezes stride0, this gate's information problem

See also

• SILOTIGER-629 — topk algorithm optimization (Robin, Implemented)
• SILOTIGER-683 — K=2048 tiered persistent decode path (Sami/Robin, Implemented)
• SILOTIGER-636 — gfx942 indexer original (fuse 4 indexer kernels)
• vllm-project/vllm#50470 — the caller-side change that routes the sparse indexer through this op

[github-actions]

🏷️ CI Guide

Runs automatically on every PR:

• ✅ Pre-checks (submodule verification, code formatting)
• ✅ Aiter op tests (gfx942 + gfx950)
• ✅ Triton tests on MI35X (only when aiter/ops/triton/** or related paths are changed)

Extended tests (opt-in via labels):

Label	Tests
ci:triton-300x	Run an additional Triton test job on MI300X in PRs; main branch always runs both MI35X and MI300X
ci:sglang	SGLang integration tests: DeepSeek-R1-MXFP4 accuracy, Qwen 3.5 accuracy
ci:atom	ATOM benchmark: DeepSeek-R1-0528, GPT-OSS-120B
ci:atom_full	ATOM accuracy suite for PR and main models from ATOM models_accuracy.json
ci:vllm	vLLM benchmark: GPT-OSS-120B, DeepSeek-R1-0528, Kimi-K2.5
ci:all	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 4355 --add-label <label>

[Copilot]

Pull request overview

Copilot reviewed 10 out of 10 changed files in this pull request and generated no new comments.

Suppressed comments (1)

Previously missed (1) — in code that hasn't changed since the last review.

aiter/aot/flydsl/common.py:175

• In _compile_one_config_for, OpKind.GROUPED_MOE is handled twice: first by importing from .grouped_moe import compile_one_config, and then again by returning a stub lambda. The second branch is unreachable and the comment contradicts the actual behavior (grouped_moe AOT appears wired up). This should be de-duplicated (either remove the stub branch, or change it to the intended missing kind).

def _compile_one_config_for(kind: OpKind) -> Callable[..., dict[str, Any]]:
    if kind is OpKind.MOE:
        from .moe import compile_one_config
    elif kind is OpKind.MXFP4_MOE:
        from .mxfp4_moe import compile_one_config
    elif kind is OpKind.GEMM:
        from .gemm import compile_one_config
    elif kind is OpKind.GROUPED_MOE:
        from .grouped_moe import compile_one_config
    elif kind is OpKind.CHUNK_GDN_H:
        from .chunk_gdn_h import compile_one_config
    elif kind is OpKind.TOPK:
        from .topk import compile_one_config
    elif kind is OpKind.GROUPED_MOE:
        # grouped_moe AOT not wired up yet (no jobs are ever collected); keep a
        # trivial stub so the dispatch is total.
        return lambda **_kw: {}
    else:
        raise ValueError(f"unknown FlyDSL AOT kind: {kind!r}")

[Copilot]

Pull request overview

Copilot reviewed 10 out of 10 changed files in this pull request and generated 3 comments.

[Copilot]

Pull request overview

Copilot reviewed 10 out of 10 changed files in this pull request and generated no new comments.

Suppressed comments (2)

Previously missed (2) — in code that hasn't changed since the last review.

aiter/ops/topk.py:598

• _should_use_flydsl_decode() calls get_gfx_runtime() without guarding exceptions. If rocminfo probing fails or the live arch is not in the known map, get_gfx_runtime raises and the whole top_k_per_row_decode op will error out instead of safely falling back to the HIP kernel (this gate is an optimization, so it should be fail-open). Suggestion: wrap the get_gfx() call (and possibly the gate lookup) in try/except and return False on any exception so decode remains functional even when arch detection is unavailable.

    arch = get_gfx()
    if arch not in _FLYDSL_TOPK_DECODE_ARCHS:
        return False
    gate = _FLYDSL_TOPK_DECODE_GATES.get(arch)
    if gate is None:

aiter/ops/flydsl/kernels/topk_per_row_decode_tiered.py:3

• This new kernel module is missing the standard SPDX license header / copyright header used throughout aiter’s FlyDSL kernel sources (e.g., other files in aiter/ops/flydsl/kernels start with # SPDX-License-Identifier: MIT). Please add the SPDX header at the top of the file for consistency and license compliance automation.

"""FlyDSL decode TopK-per-row kernel (tiered persistent multi-block radix-select)

Computes an unordered Top-K index set per decode row, fusing a single-workgroup and

[Copilot]

Pull request overview

Copilot reviewed 10 out of 10 changed files in this pull request and generated 1 comment.

[Copilot]

🟡 Changes recommended

The public dispatch path in aiter/ops/topk.py has two verified fallback/stride-handling bugs that can cause unexpected exceptions or incorrect results instead of clean HIP fallback.

Once you've addressed the issues Copilot identified, you can request another Copilot review.

Review details

• Files reviewed: 10/11 changed files
• Comments generated: 2
• Review effort level: Lite

[samremes]

Agent review: AITER PR #4355 — tiered persistent decode Top-K

Reviewed: #4355
Title: [FlyDSL] [Feature] Tiered persistent radix-select decode Top-K + arch/shape dispatcher
Local worktree: /home/samremes/dev/aiter-pr4355-stageb-baseline
Base: origin/main @ f0321c0e8927d1d90a29385433f71e592b1c51f5
Merge base: 4ad99832823dde2315b361cbd3b54b1c5c12acd5
Head: f19d30a6113d848f32e894297ae0b8cca5bd6afe
Dirty files in this worktree: none
Review date: 2026-09-10
Skills applied: FlyDSL kernel authoring, FlyDSL kernel cleanup, FlyDSL layout algebra, AITER operation testing

This document is a read-only review. I inspected origin/main...HEAD, the pinned FlyDSL 0.3.2 package, the AOT path, the dispatcher, and both new test files. I ran Python bytecode compilation, Ruff, eight host-only dispatcher tests, and AOT job discovery. I did not run a GPU kernel test or benchmark.

Verdict

Do not merge this revision. The kernels can return indices outside the logical row when a device sequence length exceeds the logits width. The new compact kernel is also disconnected from runtime dispatch, AOT, and tests.

The public wrapper has useful dtype, device, stride, workspace, and stream checks. The host-only dispatcher boundary tests pass. The change still needs one correctness fix and material integration, layout, and test cleanup.

Review context

Item	Value
Change class	New kernel and operation; performance and tuning change; AITER dispatch and AOT integration
Targets	gfx942 and gfx950; wave64; 1024 threads per workgroup
Dtypes	fp32 logits; int32 sequence lengths, indices, histograms, and counters
Shape regime	Decode rows; k in 256, 512, 1024, or 2048 through the public gate; long padded widths from 131072
Input layout	Rank-2 logits with column stride 1; packed sequence lengths; packed (numRows, k) output; contiguous int32 workspace
Expected fallback	Existing HIP one-block decode path outside the architecture, width, row-count, k, availability, or stride gate
Declared FlyDSL pin	flydsl==0.3.2 in requirements.txt and pyproject.toml
Imported FlyDSL	0.3.2 from /opt/venv/lib/python3.14/site-packages/flydsl
Diff vs base	11 files, +7036 / -15

Changed files:

• aiter/aot/flydsl/common.py (AOT operation registration)
• aiter/aot/flydsl/topk.py (Top-K AOT job generation and compilation)
• aiter/configs/topk_decode_aot.csv (AOT shape and architecture matrix)
• aiter/ops/flydsl/__init__.py (FlyDSL exports)
• aiter/ops/flydsl/kernels/topk_per_row_decode.py (standalone one-workgroup kernel)
• aiter/ops/flydsl/kernels/topk_per_row_decode_tiered.py (runtime-selected tiered kernel)
• aiter/ops/flydsl/kernels/topk_per_row_decode_tiered_compact.py (new compact tiered kernel and tuning rules)
• aiter/ops/flydsl/topk_per_row_decode.py (configuration, validation, workspace, and launch wrapper)
• aiter/ops/topk.py (public architecture and shape dispatcher)
• op_tests/flydsl_tests/test_flydsl_topk_per_row_decode.py (direct FlyDSL pytest suite)
• op_tests/test_topk_decode_dispatch.py (public dispatcher pytest suite)

FlyDSL 0.3.2 provides the APIs used in the layout recommendations below. I verified make_layout, make_view, idx2crd, crd2idx, flat_divide, zipped_divide, slice, make_tiled_copy_tv, copy, SharedAllocator, rocdl.make_buffer_tensor, rocdl.BufferCopy128b, and UniversalAtomicAdd in the installed package.

Findings

Findings are ordered by severity. Each finding has a location, a trigger, a reason, and a proposed fix.

P1 — Sequence lengths beyond the logical width select data from outside the row

Where: aiter/ops/flydsl/kernels/topk_per_row_decode.py lines 182-191, load_row_vec (linear row load), and lines 273-299, row_len and vec_blocks_i32 (lower clamp only); aiter/ops/flydsl/kernels/topk_per_row_decode_tiered.py lines 319-348, descriptor and row geometry, and lines 533-539, load_row_vec; aiter/ops/flydsl/kernels/topk_per_row_decode_tiered_compact.py lines 955-984, descriptor and row geometry, and lines 1225-1231, load_row_vec; aiter/ops/flydsl/topk_per_row_decode.py lines 562-575, _validate_inputs (sequence storage validation without a value bound).

Trigger: A seqLens entry produces row_len > logits.shape[1]. For row zero of a multi-row contiguous tensor, the scan reads the next row as part of row zero. For the final row, the bounded descriptor returns zero for out-of-range loads, but col_i32 < row_len still marks those zero values as valid. Negative valid logits can therefore lose to synthetic zeros. The kernel can emit a column index that is greater than or equal to the logical width.

Reason: Each kernel clamps row_len only at zero. The kernel does not receive the logical width. The descriptor byte count is the physical row-stride footprint, not the logical per-row bound. A descriptor bound therefore cannot prevent an inner-row overrun into the next physical row.

Proposed fix: Pass max_model_len = logits.shape[1] as an explicit runtime kernel argument. Clamp row_len to [0, max_model_len] before tier selection and before vec_blocks_i32 is computed. Apply the same contract to every retained kernel. Add a regression with two rows, negative logits, and one sequence length larger than the logical width. Verify that no output index reaches the next row or exceeds the width.

P2 — The new compact kernel is not part of the operation

Where: aiter/ops/flydsl/kernels/topk_per_row_decode_tiered_compact.py lines 311-552, compact host policy, and lines 597-2810, create_topk_per_row_decode_compact_kernel; aiter/ops/flydsl/topk_per_row_decode.py lines 452-490, _build_launcher (always calls create_topk_per_row_decode_tiered_kernel); aiter/aot/flydsl/topk.py lines 50-106 and 110-183, AOT job creation and compilation (both call the non-compact configuration and launcher); op_tests/flydsl_tests/test_flydsl_topk_per_row_decode.py lines 1-710 and op_tests/test_topk_decode_dispatch.py lines 1-326 (no compact builder or compact policy call).

Trigger: Every public call, direct FlyDSL call, AOT build, and new test in this revision.

Reason: Repository-wide reference search finds the compact symbols only in their defining file. The latest commit adds 2810 lines of kernel and tuning code, but no caller can select or compile it. The standalone builder in aiter/ops/flydsl/kernels/topk_per_row_decode.py lines 47-461 is also unreferenced. The change therefore carries three complete implementations while only the non-compact tiered implementation is reachable. Ruff also reports 17 errors in the disconnected compact file, including unused loop results and self-comparisons.

Proposed fix: Decide which implementation defines the current operation. If the compact kernel is the replacement, connect its shape policy to the public wrapper, workspace sizing, cache key, AOT jobs, and canonical behavior test. Prove correctness and matched performance first. Then delete the superseded tiered and standalone implementations, their selectors, and stale configuration. If the compact kernel is not ready, remove it from this PR and land it only with its integration and verification.

P2 — The new kernels bypass FlyDSL layout ownership for their core mappings

Where: aiter/ops/flydsl/kernels/topk_per_row_decode_tiered.py lines 275-283, thread-to-wave mapping; lines 319-348 and 533-539, raw buffer resources and row offsets; lines 388-418, workspace and pointer offsets; lines 560-569, thread-to-bin-pair mapping; lines 620-633, histogram reload; lines 967-980 and 984-1027, block/thread/vector ownership; aiter/ops/flydsl/kernels/topk_per_row_decode_tiered_compact.py lines 910-918, 955-984, 1027-1067, 1225-1236, 1252-1264, 1440-1459, and 2058-2103, equivalent mappings plus compact ownership; aiter/ops/flydsl/kernels/topk_per_row_decode.py lines 103-123, 151-191, 230-233, and 273-320, equivalent standalone mappings.

Trigger: All radix scans, histogram scans, output writes, and compact candidate scans.

Reason: These new kernels use fx.SharedAllocator, but they express most global and ownership mappings through buffer_ops, raw element formulas, fx.Index, raw LLVM pointers, and repeated division and remainder. FlyDSL 0.3.2 has native representations for these regular mappings. The current form makes units and ownership depend on repeated formulas across three implementations. It also makes changes to vector width, block size, workspace regions, or output packing unsafe.

Proposed fix: Rewrite one retained kernel around these exact mappings:

• Use thread_layout = fx.make_layout((16, 64), (64, 1)) and wave, lane = fx.idx2crd(thread_x, thread_layout).unpack() for the 1024-thread wave64 decomposition.
• Use bin_owner = fx.make_layout((1024, 2), (2, 1)). Derive each histogram bin with fx.crd2idx((tid, pair), bin_owner) for pair in 0 and 1. This mapping replaces first_bin = tid * 2 and documents exact coverage of 2048 bins. Use a one-value mode for the 1024-bin case.
• Use vector_owner = fx.make_layout((blocks_per_row, 1024, 4), (4096, 4, 1)). Derive the first unordered scan column from fx.crd2idx((part, tid, value), vector_owner). Keep the persistent increment active_parts * 4096 explicit because active_parts is runtime state.
• For ordered and compact tiles, use a static tile layout (ORDERED_STAGES, 1024, 4) with stride (4096, 4, 1). Add the dynamic run_first_vblk * 4 base after the layout maps (stage, tid, value). Keep run_first_vblk and the persistent loop explicit.
• Build physical row-major views for logits with shape (grid_y, stride0) and stride (stride0, 1). Build logical output views with shape (grid_y, top_k) and stride (top_k, 1). Wrap them with fx.rocdl.make_buffer_tensor(..., max_size=False, num_records_bytes=...). Use BufferCopy128b plus fx.copy for vec4 row and histogram transfers. Use UniversalCopy32b or the matching buffer copy atom for scalar sequence and output transfers.
• Model workspace as a row-major (grid_y, row_workspace_slots) view. Create named views for counter groups, pass histograms, and compact (column, key) records. Keep acquire/release loads and exchanges in one low-level row-barrier helper because FlyDSL 0.3.2 does not expose those exact ordered operations. Use UniversalAtomicAdd(fx.Int32, fx.rocdl.SyncScope.Workgroup) for LDS additions and a buffer atomic atom for monotonic global histogram additions where their lowering matches the current instructions.
• Keep @fx.struct plus one fx.SharedAllocator. Change s_scan to a (2, 16):(16, 1) view for wave totals and exclusive offsets. Split s_meta and s_run into named fields or named views instead of numeric slot constants. Keep the histogram as a 1-D bin view because bucket selection is data-dependent.

Compile each mapping with FlyDSL 0.3.2. Compare frontend IR, post-layout IR, vector load/store ISA, LDS atomics, registers, LDS, scratch, occupancy, numerical results, and matched latency before removing the manual form.

P2 — Architecture selection is process-global instead of tensor-device-specific

Extra note: Using e.g. get_gfx_runtime is common practice in AITER, so parsing of the device name per-kernel should likely not be done as proposed in this comment.

Where: aiter/ops/topk.py lines 571-599, _should_use_flydsl_decode (calls the process-global get_gfx_runtime); aiter/ops/flydsl/topk_per_row_decode.py lines 440-448, _current_arch (one-entry process cache), and lines 626-633, flydsl_top_k_per_row_decode (reads that cache before entering the logits device context).

Trigger: A process has devices of different supported architectures, or the current device architecture differs from logits.device.

Reason: The selected architecture comes from process-global current-device detection. The wrapper then keys the launcher and target-specific options with that architecture. decode_cu_count correctly checks the tensor device, but it receives the already incorrect architecture and falls back to an architecture floor. The JIT launcher can then freeze a target that does not match the tensor device.

Proposed fix: Derive the architecture from torch.cuda.get_device_properties(logits.device).gcnArchName.split(":")[0]. Pass that value to the public gate, workspace sizing, and launcher cache. Cache by device identity and architecture. Enter the tensor device context before any target-sensitive FlyDSL call. Add a device-selection test where the current device differs from the logits device. Add a heterogeneous-architecture test when such CI hardware is available.

P2 — The tests do not follow the AITER operation-test contract

Note: Do verify first before reading this comment if extra op_tests/flydsl_tests/ are allowed in any documentation under AITER, or any AITER issue/RFC.

Where: op_tests/flydsl_tests/test_flydsl_topk_per_row_decode.py lines 1-18, pytest-only suite and module-level GPU skip; lines 116-143, local set comparison and torch.testing.assert_close; lines 166-529 and 637-710, large pytest matrices without the standard benchmark table; op_tests/test_topk_decode_dispatch.py lines 1-326, second pytest-only suite; lines 39-81 and 142-170, assertions coupled directly to the shipped tuning table; lines 40-45, external ticket and deployment-specific narrative.

Trigger: Normal operation-test maintenance, a tuning-table update, or the standard operation benchmark workflow.

Reason: The AITER test guidance requires one canonical op_tests/test_<op>.py suite with @benchmark, run_perftest, checkAllclose, per-candidate latency and throughput fields, an architecture allow-list in main(), an itertools.product sweep, and a final Markdown table. This PR adds two pytest-only suites. The direct kernel suite uses a local comparison instead of checkAllclose. The dispatcher suite pins mutable production thresholds rather than injecting a small synthetic gate. A threshold retune can therefore fail tests without changing the dispatch contract. The test narrative also includes an external identifier that does not describe a technical behavior.

Proposed fix: Create one canonical operation test at op_tests/test_topk_decode.py. Drive the public wrapper and add a direct launcher candidate only where implementation-specific coverage is required. Use an independent fp32 reference, checkAllclose, and the standard timing and summary flow. Keep focused pytest checks only for distinct failure behavior that does not fit the benchmark function. Test gate precedence and fallback with an injected synthetic gate. Remove the duplicate suite, mutable winner assertions, and external identifiers from test and production comments.

P3 — New-file metadata and comments do not meet repository guidance

Where: aiter/ops/flydsl/kernels/topk_per_row_decode.py lines 1-10; aiter/ops/flydsl/kernels/topk_per_row_decode_tiered.py lines 1-45; aiter/ops/flydsl/kernels/topk_per_row_decode_tiered_compact.py lines 1-45; aiter/ops/flydsl/topk_per_row_decode.py lines 1-22; op_tests/flydsl_tests/test_flydsl_topk_per_row_decode.py lines 1-7 (new files have SPDX text but omit the repository copyright line); aiter/ops/topk.py lines 480-483 (production comment refers to an external investigation identifier).

Trigger: Source distribution, generated notices, or reuse of these comments outside the original development context.

Reason: Nearby new AITER files use the current copyright convention. The cleanup guidance also requires comments to explain the technical property without external ticket or review-history references.

Proposed fix: Add the repository's current copyright line to retained new files. Replace external references with the measured width, row-count, architecture, and graph-replay facts that justify each threshold. Remove historical tuning narratives that do not define a maintained invariant.

Verification gaps

• I did not run GPU correctness on gfx942 or gfx950.
• I did not compile the kernel from an empty FlyDSL cache.
• I did not test AOT compile-only output or a packaged run-only cache hit.
• I did not test graph capture and replay, a caller-provided non-current stream, workspace reuse across streams, or heterogeneous devices.
• I did not inspect generated IR, ISA, VGPR, SGPR, LDS, scratch, or occupancy for this revision.
• I did not reproduce the claimed architecture and shape performance windows.
• Python bytecode compilation passed for the changed Python files.
• python3 -m pytest op_tests/test_topk_decode_dispatch.py -q -k 'shipped or staircase or routing_cases' passed 8 tests and deselected 53 tests.
• AOT job discovery completed and produced 224 jobs: 36 per k for the 304-CU rows and 20 per k for the 256-CU rows.
• git diff --check origin/main...HEAD passed.
• Ruff failed with 17 findings, all in topk_per_row_decode_tiered_compact.py.