[FlyDSL] Add FlyDSL Radix-Select TopK Path to the Existing Per-Row Decode Interface - #5011
[FlyDSL] Add FlyDSL Radix-Select TopK Path to the Existing Per-Row Decode Interface#5011lirui927 wants to merge 23 commits into
Conversation
Add a compile-time-specialized multi-CTA radix-select path with optional deterministic output for the existing per-row decode interface. Co-authored-by: Cursor <cursoragent@cursor.com>
🏷️ CI GuideRuns automatically on every PR:
Extended tests (opt-in via labels):
PR title tags: |
Reuse radix histograms for stable selection, add wide-count support for long contexts, and align launcher caching and validation with other FlyDSL operators. Co-authored-by: Cursor <cursoragent@cursor.com>
Vectorize stable writes with a compile-time width choice and remove redundant scan barriers to improve long-context latency without regressing short rows. Co-authored-by: Cursor <cursoragent@cursor.com>
|
FYI, there is also this PR with a similar kind of kernel: #4355 which seems faster than this in many cases in our quick comparisons. |
@samremes Thanks for your great job. GPU：MI355X ROCm：7.2.4 FlyDSL：0.3.1 HIP：7.2.53211
|
Initialize empty-row state and overwrite unused output slots so CUDA Graph replays cannot retain stale indices. Co-authored-by: Cursor <cursoragent@cursor.com>
Reduce long-row scan and dispatch overhead with adaptive three-pass radix, vectorized histogram reduction, fused stable prefixing, cached workspace, and short-row direct fill. Co-authored-by: Cursor <cursoragent@cursor.com>
Use a 1024-thread, 16-chunk data geometry with a 256-thread reduce to improve three-pass radix performance across stable and unstable batches. Co-authored-by: Cursor <cursoragent@cursor.com>
Remove four-pass and shape-specific routing so every decode path uses the tuned 11-bit radix geometry. Co-authored-by: Cursor <cursoragent@cursor.com>
Add single-launch stable paths and reuse dynamic-N multi-kernel binaries so model workloads avoid repeated launches and per-width compilation. Co-authored-by: Cursor <cursoragent@cursor.com>
Route long single-row calls through the safe multi-launch path because concurrent graph replays can deadlock non-cooperative barriers, while retaining the barrier-free short path. Co-authored-by: Cursor <cursoragent@cursor.com>
Route measured gfx950 shapes through dynamic FlyDSL kernels with HIP fallback, and remove unused persistent paths to keep graph-safe execution maintainable. Co-authored-by: Cursor <cursoragent@cursor.com>
lirui927
changed the title
Reuse the faster deterministic short-row kernel for unordered calls, whose output remains valid for the weaker ordering contract. Co-authored-by: Cursor <cursoragent@cursor.com>
Apply the formatter output required by the code-style check. Co-authored-by: Cursor <cursoragent@cursor.com>
Keep focused dispatch, end-to-end accuracy, and CUDA Graph coverage while removing redundant performance and implementation-detail cases. Co-authored-by: Cursor <cursoragent@cursor.com>
Enable only the measured short and long gfx942 ranges while retaining HIP fallback for unstable intermediate shapes. Co-authored-by: Cursor <cursoragent@cursor.com>
Reduce the short-context stable dispatch limit on gfx942 so FlyDSL is only selected for smaller batch sizes where it reliably wins. Co-authored-by: Cursor <cursoragent@cursor.com>
Update gate and e2e cases so short-context stable dispatch boundaries match the tightened gfx942 FlyDSL row gate. Co-authored-by: Cursor <cursoragent@cursor.com>
Replace inline fly_rocdl.update_dpp calls with the shared DPP helper. Co-authored-by: Cursor <cursoragent@cursor.com>
Add the blank line Ruff I001 expects between third-party and local imports. Co-authored-by: Cursor <cursoragent@cursor.com>
Specialize kernels at compile time to write selected values inline while preserving the no-values path and signed-zero ordering. Co-authored-by: Cursor <cursoragent@cursor.com>
Fall back to HIP for unsupported FlyDSL signatures and align edge-case ordering while covering ties, NaNs, and MTP row lengths. Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
4 participants
Add a FlyDSL Radix-Select TopK Path to the Existing Per-Row Decode Interface
Motivation
The latency of the existing HIP one-block implementation increases with the effective context length, making it less efficient for long-context decode workloads.
This PR aims to:
top_k_per_row_decodeinterface.next_nruntime parameters so different context lengths reuse the same compiled kernel.Technical Details
Multi-CTA Radix-Select
The long-sequence path splits each row into 16 chunks. Multiple CTAs build local histograms in parallel, and three 11/11/10-bit radix-selection passes progressively determine the TopK threshold.
stable=Falsegathers results directly from the threshold, whilestable=Trueadditionally computes chunk prefixes and writes results in original-index order for deterministic output. Launchers are cached by(rows, K, stable), while N,next_n, and stride are runtime parameters.Short-Sequence Path
For physical widths up to 20K, one workgroup completes all three radix-selection passes in a single kernel launch. The histogram, prefix scan, and deterministic writeback remain in LDS and require no cross-workgroup synchronization.
stable=Falsereuses the same result; deterministic ordering is stronger than required but remains valid for the unordered TopK interface.Public Dispatch
The FlyDSL path is currently enabled only on gfx950 for:
K ∈ {512, 1024, 2048, 4096}stable=Truestable=Falsestable=Truestable=FalseOther architectures, shapes, K values, or environments without FlyDSL fall back to HIP.
Set
AITER_DISABLE_FLYDSL_TOPK_DECODE=1to force the HIP path.Test Plan
Run on gfx942/gfx950 with FlyDSL available:
test_public_topk_decode_gate): arch/K/width/rows at each FlyDSL enable/disable edge, including unsupported K values.test_public_topk_decode_e2e): random logits throughtop_k_per_row_decodecompared against the Torch reference on both FlyDSL and HIP fallback paths; coversstable=True/Falseand slightly unaligned physical/effective lengths.test_public_topk_decode_cuda_graph): capture and replay on the long-row multi-kernel path.Representative shapes:
256, 512, 2048, 4096(plus unsupported K at boundaries)1, 4, 16, 17, 32)Test Result
Correctness:
The following results were measured with values=None. Enabling values output generally adds less than 2% performance overhead.
Kernel Performance at Dispatch Boundaries
gfx950 / MI355X Kernel Performance at Dispatch Boundaries
The following results use the final code on gfx950 with
K=2048and unpadded inputs. CUDA Graph replay measures the kernel path. Speedup isHIP / FlyDSL; values greater than 1 indicate that FlyDSL is faster.K=2048andstable=True, FlyDSL CUDA Graph performance was faster than HIP in all 77 cases, with speedups ranging from 1.02× to 4.02× and a median of 1.94×.The following cases use the maximum row count at each dispatch boundary, plus a 1M-length case.
gfx942 / MI308 Kernel Performance at Dispatch Boundaries
The following results use
lirui-vllm-latest-0623with FlyDSL 0.3.2 on AMD Instinct MI308X (gfx942). CUDA Graph replay measures the kernel path. Each value is the median of 101 graph replays after 10 warmup replays. Speedup isHIP / FlyDSL; values greater than 1 indicate that FlyDSL is faster. All listed cases matched the HIP reference.Based on these measurements, gfx942 dispatch is gated conservatively to:
stable=True:width <= 20K && rows <= 64, orwidth >= 128K && rows <= 16stable=False:width >= 512K && rows <= 32Submission Checklist