Repository navigation
Conversation
🏷️ CI GuideRuns automatically on every PR:
Extended tests (opt-in via labels):
One backend per PR: PR title tags & labels: |
ea761a4 to
ed5cf9f
Compare
akii96
changed the title
|
I tested on MI325X, and the PR seems a solid one: all tests pass, and I measured similar speedups. Couple things though:
Also a suggestion: LDS seems to be capping the occupancy now to 3 workgroups per CU. Scales are what tips it over, 4 x KV tiles itself would fit the 64kB. Could the scales be loaded directly to registers? |
ed5cf9f to
2a3c08f
Compare
NOT READY FOR HUMAN REVIEWPR Self-ReviewRevision Author self-review accepted. Readiness report: Author answers
Show the author's self-review answersBehavior and intentValidation and unverified scopePR-specific questionChanged op testsPR state
Copilot findings
Code and op-test hygiene
Required CI
Shadow mode reports what would block but keeps the status successful. Set |
…ss waves
On gfx942 the direct-load kernel is limited by memory traffic, not MFMA
throughput: each wave loads its own slice of every KV tile and holds the Q
and weights of all of the block's rows, which keeps the VMEM queue full and
caps occupancy at 2 waves/SIMD.
Add an LDS-staged variant, mfma_bkv64_r{1,2,4}_w4_lds2
(kernels/mqa_logits/fp8_mqa_logits_gfx942_lds.py), and auto-select it on
gfx942. The block's 4 waves DMA one 64-column KV tile into LDS and all read
it back, so each tile is loaded once per block, and each wave keeps only its
own rows in registers (waves_per_eu=3). The variant handles a short last
block itself, so the host does not pad rows for it.
Selection: rows per wave start from a num_heads register cap and drop for
short or packed windows and for few rows against a long KV range; calls
with fewer than 16 rows keep the direct-load variant. The new variant has
its own _SplitPolicy entry. variant= and FLYDSL_FP8_MQA_LOGITS_VARIANT still
force any variant.
Rows with an empty window no longer widen a block's KV range, in both the
new and the direct-load kernel. gfx950 selection is unchanged.
op_tests: --packed-shapes, a small packed subset in the default lane, and
pins for the gfx942 variant selection.
Signed-off-by: Aakif Nawaz <aakif.nawaz@amd.com>
2a3c08f to
910d863
Compare
|
Thanks for testing and the feedback @juuso-oskari , I had forgotten about this but your review got me to pick this back up 😄 The newest pushes fixes some of the issues you raised with one todo:
(opening this now coz I am happy with how it looks for now but happy to iterate further) |
zhiding512
left a comment
zhiding512
left a comment
There was a problem hiding this comment.
Request changes: host-side row padding regresses gfx950
Thanks for the LDS-staged gfx942 work. On a gfx950 box I verified that the new mfma_bkv64_r*_w4_lds2 path itself is not selected, but the host launcher change still lands on every gfx950 call and regresses latency whenever seq_len is not a multiple of the rows-per-block of the auto-selected variant.
Root cause
After rebase onto main, this PR reintroduces the host pad of Q / weights / cu_starts / cu_ends for every variant that is not in _LDS942_VARIANTS:
aiter/aiter/ops/flydsl/fp8_mqa_logits_kernels.py
Lines 705 to 722 in 7580cd5
Main already dropped that padding in #6001: every builder (including all gfx950 LDS variants) masks a short last block in-kernel, and the module docstring still says "Query rows are not padded". The new gfx942 LDS kernels correctly skip the pad; everything else, including all gfx950 variants, gets the old host copies back.
On an unaligned seq_len one call therefore goes from 1 GPU kernel to 9 (zeros + cat for the four tensors, then the real kernel). Outputs stay bitwise identical to main.
Reproduce on gfx950 (MI355X)
# worktrees at the PR merge-base and at this PR tip
git fetch origin main pull/6164/head:pr/6164
MB=$(git merge-base origin/main pr/6164) # a97e99922...
git worktree add --detach /tmp/aiter-base $MB
git worktree add --detach /tmp/aiter-pr pr/6164
# flydsl matching aiter's pin
pip install --target /tmp/flydsl042 'flydsl==0.4.2'
# copy the PR test harness onto the base worktree so both sides run the same cases
cp /tmp/aiter-pr/op_tests/test_flydsl_fp8_mqa_logits.py \
/tmp/aiter-base/op_tests/test_flydsl_fp8_mqa_logits_pr6164harness.pyCorrectness (both sides pass):
# default + packed-shapes verify
(cd /tmp/aiter-pr && HIP_VISIBLE_DEVICES=0 PYTHONPATH=/tmp/flydsl042 \
FLYDSL_RUNTIME_CACHE_DIR=/tmp/flydsl_cache_pr \
python3 op_tests/test_flydsl_fp8_mqa_logits.py --scenario verify)
# -> all 73 case runs passed
(cd /tmp/aiter-base && HIP_VISIBLE_DEVICES=0 PYTHONPATH=/tmp/flydsl042 \
FLYDSL_RUNTIME_CACHE_DIR=/tmp/flydsl_cache_base \
python3 op_tests/test_flydsl_fp8_mqa_logits_pr6164harness.py --scenario verify)
# -> all 73 case runs passed
# same for --packed-shapes (36/36 both sides)Latency A/B with the PR's own graph-replay bench (aligned control vs unaligned / packed):
# model shapes are all multiples of RPB -> geomean PR/base ≈ 1.00 (control)
(cd /tmp/aiter-base && HIP_VISIBLE_DEVICES=0 PYTHONPATH=/tmp/flydsl042 \
FLYDSL_RUNTIME_CACHE_DIR=/tmp/flydsl_cache_base \
python3 op_tests/test_flydsl_fp8_mqa_logits_pr6164harness.py --scenario bench --model-shapes)
(cd /tmp/aiter-pr && HIP_VISIBLE_DEVICES=0 PYTHONPATH=/tmp/flydsl042 \
FLYDSL_RUNTIME_CACHE_DIR=/tmp/flydsl_cache_pr \
python3 op_tests/test_flydsl_fp8_mqa_logits.py --scenario bench --model-shapes)
# packed multi-seq chunks (most s_q not a multiple of RPB) -> geomean ~1.20× slower
... --scenario bench --packed-shapes
# deliberately odd rows -> geomean ~1.31× slower
... --scenario bench \
-s 1023,1023 4095,4095 8191,8191 16383,16383 4095,36863 8191,40959 \
-w causal --head-dim 128 --clean-logits 0 1 --num-heads 64Spot measurements I got on one MI355X (steady-state CUDA-graph replay, no LLC flush; same seed / inputs):
| shape | base µs | PR µs | note |
|---|---|---|---|
| 8192×8192 H64 | 326 | 326 | aligned control |
| 8191×8191 H64 | 330 | 431 | +30%, 1→9 kernels |
| 8191×8191 H32 | 193 | 260 | +35% |
| 16383×16383 H128 | 2203 | 2508 | +14% |
| 1023×1023 H64 | 18 | 47 | 2.5× |
| 7×20007 H64 | 12 | 29 | 2.4× |
A profiler on the unaligned shapes shows the PR launching zeros / vectorized_elementwise cats around the FlyDSL kernel; base launches only the FlyDSL kernel.
Requested fix
Please drop the host pad again and always pass the real seq_len, the same way main does after #6001. Keep _LDS942_SPLIT_POLICY for the new gfx942 LDS variants only. That is a ~27-line deletion in flydsl_fp8_mqa_logits; after that change on this box, the unaligned shapes above matched main again (73/73 and 36/36 still green).
Sketch of the resulting launcher body:
# No torch.full even when clean_logits: the kernel writes -inf itself.
aligned_size = 256
seq_len_kv_aligned = (seq_len_kv + aligned_size - 1) // aligned_size * aligned_size
logits = torch.empty(
(seq_len, seq_len_kv_aligned),
dtype=torch.float32,
device=Q.device,
)[:, :seq_len_kv]
...
scalars = (int(seq_len), int(seq_len_kv), int(logits.stride(0)))
...
split_pol = _LDS942_SPLIT_POLICY if variant in _LDS942_VARIANTS else None
num_splits = _auto_num_splits(
seq_len, seq_len_kv, _ROWS_PER_BLOCK, _BKV, Q.device.index,
split_pol, blocks_per_cu,
)Happy to re-check on gfx950 after the push.
…tical to main The merge with main reintroduced host-side padding of Q/weights/cu_starts/ cu_ends to a multiple of the block's rows for every variant outside the gfx942 LDS set. main dropped that in ROCm#6001 (every builder masks a short last block in-kernel), so on gfx950 an unaligned seq_len went from one launch to nine and regressed (reported on MI355X). Pass the real seq_len to every variant again, as main does; only the split policy stays variant-specific (_LDS942_SPLIT_POLICY for the gfx942 LDS variants). Also hard-gate the PR to gfx942: revert the shared direct-load builder (_build_kernel_mfma_r_w, used by gfx950's direct-load variants) to main's version. The empty-row union-window change only existed to absorb padded rows, and the hoisted _fn_to_fnuz_i64 now lives in the gfx942 LDS module, which is its only other user. fp8_mqa_logits.py is now byte-identical to main. Add a launch-free launcher-contract check to the op test (real seq_len, no host cat/new_zeros) that runs on every arch. Signed-off-by: Aakif Nawaz <aakif.nawaz@amd.com>
|
Thanks @zhiding512, good catch, and thanks for the detailed repro. You were right: the merge with main brought the host-side row padding back for every non-LDS variant. Fixed in cd113ea:
On an MI355X, your unaligned shapes now match main (e.g. 1023² H32 26.86 → 26.80 µs, 4095×36863 H64 1638 → 1635 µs), and cross-compiled IR/ISA is identical to main for all 28 variant × clean_logits combos I checked. A rerun on your box would be great whenever you get a chance. |
… policy Rework the gfx942 LDS-staged kernel's inner loop and launch policy, keeping its algorithm and variants: - Software-pipelined MFMA chains: B fragments read one n-tile ahead, each chain's ReLU x weight epilogue interleaved under the next chain's MFMAs, and the next tile's DMA issued after the tile's last LDS read and spread over the last MFMA chains, so the loop-top vmcnt is exact. - Rows pre-permuted per lane group, so the head reduce-scatter needs no selects; it is staged across MFMA chains, and each tile's last reduce and store is carried past the next barrier. - One DMA descriptor per KV tile, all window loads issued up front, and paired non-temporal RPW=2 stores and a dwordx4 -inf fill. - grid.y slices of decreasing size, so the last blocks to start are the smallest. - Compile the LDS variants with the SLP vectorizer off, so the epilogue stays scalar v_fmac that issues under the MFMAs. 128 heads with an FN operand compile at waves_per_eu=2, which removes its VGPR spills. - Split grid.y further for long cached-prefix windows, sized from the variant's waves_per_eu, so graph capture and eager calls pick the same split; cap the split count so the slice bounds fit in i32. - RPW=2 from 4096 rows without a cached prefix, for every head count. gfx950 and the direct-load kernel are unchanged. Signed-off-by: Aakif Nawaz <aakif.nawaz@amd.com>
|
Pushed a new commit from Amir (@aghamari) that makes the gfx942 LDS kernel faster (thanks for the great work!). Variant names and gfx950 are unchanged.
Results on MI325X vs the previous commit:
will add E2E validation numbers soon! |
Summary
Adds an LDS-staged
fp8_mqa_logitsvariant for gfx942 and makes it the default for ≥16 query rows. gfx950 is unchanged.Motivation
On gfx942 the direct-load kernel is bound by memory traffic: every wave loads its own slice of each KV tile and holds all of the block's rows in registers, capping occupancy at 2 waves/SIMD.
Changes
mfma_bkv64_r{1,2,4}_w4_lds2(fp8_mqa_logits_gfx942_lds.py): the block's 4 waves share each KV tile through LDS, and each wave keeps only its own rows, giving 3 waves/SIMD.fp8_mqa_logits.py) is byte-identical to main; all new variants, hints and selection are behind_ARCH == "gfx942".head_size ≤ 256(D=512 would overflow 64 KB LDS). Falls back to direct-load for larger heads.vmcntis exact.dwordx4-inf fill.v_fmacunder the MFMAs); 128 heads with an FN operand compile atwaves_per_eu=2, removing its VGPR spills.Rollback: set
FLYDSL_FP8_MQA_LOGITS_VARIANT=mfma_r4_w2(or passvariant=) to pin the previous direct-load kernel for every shape.Performance
MI325X (gfx942), vLLM 0.34.0rc1.dev206 (
nightly-f5d93296c8ec), flydsl 0.4.2.Kernel benchmarks (graph-replay timing, L2 flushed per replay)
--full/--packed-shapes/ default / FN Q+KV)--model-shapes, 24 shapes)--model-shapes): main → PRLong-context chunks are the shapes the DSv4 indexer issues at 128k+ prompts under the default 512 MB logits budget (e.g. 1024 rows × 128k KV, 2048 × 64k):
End-to-end serving (prefix caching off)
DeepSeek-V4-Flash, 4× MI325X, TP=4, random dataset,
--ignore-eos, OSL 256. Previous revision of this PR (cd113ea, 2 runs) → this revision (2–4 runs), medians. TPOT is unchanged within noise.GSM8K (8-shot, 1319 questions)
No accuracy regression.
gfx950 (MI355X): unchanged vs main
Testing
Documentation
Dependencies
Breaking Changes
None; the PR only modifies FlyDSL kernel files and their test.
PR Self-Review
Reviewed revision
107aaeb
Behavior and intent
Adds a new LDS-staged kernel variant for gfx942 without changing gfx950 or other archs; no variant gets host-side row padding, matching main after #6001. The direct-load variant stays and is auto-selected for short calls (<16 rows) or large heads (D>256). The LDS variant lets 4 waves share one KV tile through LDS instead of each loading its own copy, bumping occupancy from 2→3 waves/SIMD. Also adds a head_size guard so D=512 doesn't overflow the 64 KB LDS budget — it falls back to the existing direct-load path. The latest commit reworks the LDS kernel's schedule and its launch policy (grid.y split sizing, slice shape, compile hints) and moves the RPW=2 threshold for causal chunks to 4096 rows for every head count.
Validation and unverified scope
On MI325X at the latest commit:
test_flydsl_fp8_mqa_logits.pydefault 133/133,--packed-shapes94/94,--full1618/1618, FN Q+KV 133/133, plus 78/78 extra H16/H128 and head_dim 256 cases. All 54 LDS kernel builds compile. Kernel bench: 1.73× vs main on the 24 model shapes and 1.77× on 23 long-context cached-prefix shapes (32k–128k KV), 0 regressions; the final launch-policy cleanup measured 1.00× against the revision those numbers come from. E2E with DeepSeek-V4-Flash on 4× MI325X (TP=4, prefix caching off) at 128k–248k ISL: +1.3–1.9% total throughput and 1.5–2.7% lower TTFT vs the previous revision of this PR. GSM8K 8-shot: 0.9234→0.9287 strict, no regression. black/ruff pass.gfx950: cross-compiled IR/ISA identical to main for 28/28 variant × clean_logits combos, and on MI355X FlyDSL timings on unaligned shapes match main.
Unverified: on gfx950,
--packed-shapesfails at H64 on both main and this PR because the Gluon reference it compares against is wrong there (single-warp race, fixed separately in #7823); FlyDSL had no errors. GSM8K and the gfx950 checks are from the previous revision (gfx950 code is unchanged since).PR-specific question
The LDS variants register under
mfma_bkv64_r{1,2,4}_w4_lds2in_VARIANT_BUILDERS._resolve_variantauto-selects them only when arch=gfx942, num_heads∈{16,32,64,128}, seq_len≥16, and head_size≤256. They are also reachable throughvariant=, theFLYDSL_FP8_MQA_LOGITS_VARIANTenv var, and theaiter.mqa_logitstuner, which enumerates every registered variant. Direct-load variants stay registered and reachable the same way.Changed op tests
test_flydsl_fp8_mqa_logits.py: added--packed-shapesfor multi-sequence chunked-prefill shapes and pins for the LDS variant choice, including the 4096-row RPW boundary at 16/32/64 heads.--packed-shapescovers short per-sequence windows and short last blocks at the model head counts. Addedverify_launcher_contract, a launch-free check on every arch that fails on any host-side row padding and, for the gfx942 LDS variants, on a grid.y split count whose slice bounds overflow i32 (1.5M- and 1M-column cases). The default lane also runs one 16-row × 1.5M-column case end to end. No existing test removed.