[Kernel][FA] Support paged FP8 Flash attention with asymmetric K/V head dim on gfx950 - #1066
[Kernel][FA] Support paged FP8 Flash attention with asymmetric K/V head dim on gfx950#1066sammysun0711 wants to merge 27 commits into
Conversation
There was a problem hiding this comment.
Pull request overview
Adds gfx950-only support for paged FP8 flash-attention with asymmetric K/V head dimensions (Q/K=192, V/O=128 or 192) through the existing flydsl_flash_attn_func paged-KV API, avoiding padding V to 192 and reducing cache/bandwidth waste.
Changes:
- Extend paged flash-attn dispatch/validation to allow FP8 packed-varlen vectorized-KV with independent
value_head_dim, producing native-width BF16 output. - Generalize gfx950 dualwave FP8 traits/context/utilities to handle paged+vectorized KV and asymmetric output strides.
- Add correctness and addressing-robustness tests for the new paged FP8 asymmetric-V path (including a >4GiB cache-offset case).
Reviewed changes
Copilot reviewed 4 out of 4 changed files in this pull request and generated 1 comment.
| File | Description |
|---|---|
| tests/kernels/test_flash_attn_fwd.py | Adds gfx950 paged-FP8 asymmetric-V correctness test and a large-shape >4GiB cache-offset validation. |
| kernels/attention/flash_attn_utils.py | Extends gfx950 dualwave FP8 trait/context/helpers for paged vectorized KV and asymmetric output width/stride. |
| kernels/attention/flash_attn_interface.py | Hooks paged FP8 path into the public API, validates shapes/scales, and allocates native-width BF16 output. |
| kernels/attention/flash_attn_fp8_paged_gfx950.py | New gfx950 paged FP8 packed-varlen kernel module specialized for Q/K=192 and V/O=128/192. |
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
| key_page_bytes = page_size * head_dim | ||
| value_page_bytes = page_size * value_head_dim | ||
| high_page = math.ceil(2**32 / min(key_page_bytes, value_page_bytes)) | ||
| num_pages = high_page + 1 |
There was a problem hiding this comment.
Fixed in 32dac38. The test_flash_attn_fwd.py now calculates the K/V cache allocation size, releases unused cached allocations, and checks available VRAM before allocating.
It skips with a clear message unless the approximately 10 GiB cache requirement plus 2 GiB headroom is available.
b22846f to
32dac38
Compare
jhinpan
left a comment
jhinpan
left a comment
There was a problem hiding this comment.
COMMENT — no blocking defect found in what I can verify offline, but I cannot run the device suite this change depends on, so I am not approving around that gap.
Prior review items
- Copilot,
tests/kernels/test_flash_attn_fwd.py:4413(≈10 GiB K/V cache can OOM underlarge_shape) — fixed. At this head,test_paged_fp8_d192_cache_offsets_above_4gibcomputescache_bytes = num_pages * (key_page_bytes + value_page_bytes), callstorch.cuda.empty_cache(), readstorch.cuda.mem_get_info()and skips unlesscache_bytes + 2 GiBis free. I checked the arithmetic:high_page = ceil(2**32 / 8192) = 524288,num_pages = 524289,cache_bytes = 524289 * 20480 ≈ 10.7 GB, so the guard threshold is ≈12.7 GB — consistent with the reply. The in-testassert high_page * key_page_bytes >= 2**32/assert high_page * value_page_bytes >= 2**32still hold exactly at that page count, so the >4 GiB property the test exists to prove is not weakened by the guard.
No other reviewer items are open.
What I verified at this commit
Environment: no build-fly/, no torch, and rocminfo reports no GPU agents — so pytest tests/kernels/test_flash_attn_fwd.py cannot run here at all.
What I did run / check:
python3 -m py_compileon all four changed files — clean.python3 scripts/check_repo.py—agent-docscheck OK (871 symbols, 16 skills). Thetyped-arithmeticcheck aborts withgit merge-base HEAD origin/mainreturning 128 because this is a single-commit shallow checkout; that is an environment limitation, not a finding. I substituted a grep of the added lines for the spellings it rejects (maximumf/minimumf/maxsi/minsi/ceildivsi/…) — none present; the new code usesfx.maxnumf.- No added line exceeds 120 columns, so the black gate should be clean (black/ruff are not installed here, so this is a proxy, not the real gate).
- "Dense FP8 D128 unchanged" claim — checked by hand rather than taken on trust.
store_final_onow callsglobal_idx_oinstead ofglobal_idx_q; that is equivalent on the dense path becauseDualwaveFp8KernelContext.__init__defaultsstride_o_ntostride_q_n(flash_attn_utils.py:4426),kernels/attention/flash_attn_fp8_gfx950.py:137does not passstride_o_n, andV_HEAD_DIMdefaults tohead_dimin_make_dualwave_swp_fp8_traits. Same foro_nrec_bytes.smem_d_rpt = ceil(head_dim / d_128b_size)is identical to the oldhead_dim // d_128b_sizeat D128,d_chunks/sdrpt_bfare unchanged whenvalue_head_dim is None, andfp8_pv/fp8_pv_directonly lose their value whenpaged=True. Every other new branch is behindconst_expr(PAGED)/const_expr(KV_VECTORIZED). The only dense-path side effect is a changedcache_tag, i.e. a one-time JIT cache miss. - LDS budget for the new config — with
varlen=True,bn128is False sonum_prefetch_k = 2:kv= 2 * (821040 + 821088) = 68096 B,vt= 2 * 8 *sdrpt_bf* 544 * 2 B = 34816 B (V128) / 52224 B (V192). Totals 102 KB / 118 KB, inside the gfx950 160 KB budget. - New LDS index arithmetic — walked the bounds of
_stage_vt_dequant_fp8_vectorizedand the vectorizedload_v: max written element is 13023 and max read is 13024 againstVT_BF16_ELEMS = 13056at V192 (8671/8672 vs 8704 at V128), and the vectorized K reads stay under 12288 B inside the 16640 B K region. No overrun that I can see. - Interface validation vs the existing fixtures — the new
expected_v_tail = (Hkv, page_size // kvs, value_head_dim, kvs)matches what_vectorize_paged_kv(tests/kernels/test_flash_attn_fwd.py:405-410) produces, and the newseqlen_k/block_table/outassertions are satisfied by the existing paged fixtures (seqlen_kis int32 withnumel() == B;o_t = torch.zeros_like(q_t)). So the in-repo BF16 paged tests should not regress on validation grounds. - New reference test math — the K permute
(0,3,1,2,4)and V permute(0,2,4,1,3)intest_paged_fp8_d192_asymmetric_value_matches_torchcorrectly invert the aiter 5D layouts, and the bottom-right causal mask matchescross_seqlen=True.
The gap (why COMMENT, not APPROVE)
Everything above is static. The actual correctness of the new 962-line gfx950 schedule — the sigma-permuted V dequantization, the folded logit scale in qk, _lazy_rescale_o_paged, and the page-rebased 64-bit descriptors — is only demonstrated by tests/kernels/test_flash_attn_fwd.py on real gfx950 hardware. I have no GPU here and no built flydsl, so I am relying entirely on the 103-passed run in the PR description. There are also zero check-runs on this head commit (gh api .../check-runs returns total_count: 0), so CI has not covered it either.
Also worth flagging as state, not as a code defect: this PR is mergeable_state: dirty (conflicts with main) and is stacked on the still-open #1065.
Non-blocking notes
Inline below. None of these are correctness bugs as far as I can tell; the first one is the one I would most like to see removed before merge because it reads as an unfinished edit.
|
|
||
| def floor_masked_max(self, row_max): | ||
| if const_expr(self.traits.PAGED): | ||
| return fx.maxnumf(row_max, self.c_neg_floor) |
There was a problem hiding this comment.
These three const_expr(self.traits.PAGED) branches are exact no-ops — the guarded return and the fall-through return are character-for-character identical:
floor_masked_max(5386-5388): bothreturn fx.maxnumf(row_max, self.c_neg_floor)reduce_sum(5402-5404): bothreturn l_row + self.tile_sum(v_p)apply_l_rescale(5445-5446): bothreturn l_row * rescale
The neighbouring sub_m / rescale_from_tile_max / lazy_rescale_o branches do differ (they drop the c_logit_scale factor, which is correct since the paged qk now folds it in), so these three read like an unfinished edit rather than an intentional pattern. I convinced myself they are harmless — c_neg_floor is a unit-agnostic -3.0e38 clamp, and both rescale forms are plain multiplies — but a future reader will burn time re-deriving that. Please just delete the three if blocks.
There was a problem hiding this comment.
Fixed in 2b2cbe9f. I removed the identical PAGED branches from floor_masked_max, reduce_sum, and apply_l_rescale; the behaviorally distinct paged softmax/rescale branches remain unchanged.
| src_div = self.make_page_view(self.k_base_iter, page_id) | ||
| # Each wave owns two 64-lane octet rows at D192. The physical K | ||
| # page has only 12 d-groups; the final four groups address beyond | ||
| # the page descriptor and the gfx950 buffer load supplies zeros. |
There was a problem hiding this comment.
This comment is only accurate for the last KV head. The physical K page is [Hkv, D/16, PageSize, 16], so a head's region is 12 * 64 * 16 = 12288 bytes, but NUM_DMA_K = SMEM_D_RPT = 2 makes each wave issue octets for dg in [0, 15]. For kv_head_idx < Hkv - 1, dg 12..15 lands inside the next head's data, not "beyond the page descriptor", so it does not read zeros.
As far as I can tell this is benign: _read_vec_key below only ever addresses dg <= 11 (ws runs over HEAD_DIM // 64 = 3, dg0 = ws*4 + lane_div_32*2, max dg0 + 1 = 11), so those LDS slots are written and never consumed, and on the final head the buffer descriptor clamps the read to zero without faulting. So this is a comment-accuracy issue plus ~25% wasted K DMA bandwidth, not a correctness bug. Worth rewording, since both new tests use num_kv_heads=1 and would not expose the difference. The same claim appears in the smem_d_rpt comment in _make_dualwave_swp_fp8_traits.
There was a problem hiding this comment.
Fixed in 2b2cbe9f. Both comments now describe the unused D256 LDS padding without claiming that every extra load is descriptor-clamped. The DMA schedule is unchanged.
An additional Hkv=2 gfx950 probe matches PyTorch with max absolute error 4.883e-4 and cosine 0.9999964.
There was a problem hiding this comment.
Fixed in 2b2cbe9f. The V staging permutation now calls _sigma_k_tile_n, making the required K/V permutation relationship explicit.
| ) | ||
| ni = oct_idx % traits.BLOCK_N | ||
| dg = oct_idx // traits.BLOCK_N | ||
| src_oct = dg * traits.BLOCK_N + _sigma_k_tile_n(ni) |
There was a problem hiding this comment.
This inlines exactly what _vec_k_src_elem (line 807) already computes — the oct_idx expression above matches _vec_k_dma_oct_idx (line 797) verbatim, and these two lines match lines 810-815. _vec_k_src_elem is already used at line 4067 for the bf16 vectorized path. Reusing it here would keep the two vectorized-K address derivations from drifting apart.
There was a problem hiding this comment.
Fixed in 2b2cbe9f. The paged FP8 K loader now calls _vec_k_src_elem instead of duplicating the vectorized address calculation.
| # sigma(n). Apply the same (involutive) permutation to V; | ||
| # otherwise selective attention pairs K[sigma(n)] with V[n] | ||
| # even though near-uniform low-amplitude tests appear valid. | ||
| bf = [bf[(i & 3) | ((i & 8) >> 1) | ((i & 4) << 1)] for i in range_constexpr(16)] |
There was a problem hiding this comment.
This bit-shuffle is _sigma_k_tile_n(i) (line 802) for i < 16 — the (ni & ~15) term is zero over that range. Calling the named helper would make the "same (involutive) permutation as K" claim in the comment above self-evident instead of something the reader has to verify by hand.
There was a problem hiding this comment.
Fixed in 2b2cbe9f. The V staging permutation now calls _sigma_k_tile_n, making the required K/V permutation relationship explicit.
| f"flydsl_flash_attn_func: block_table must be 2D on {q.device}, " | ||
| f"got shape={tuple(block_table.shape)} device={block_table.device}" | ||
| ) | ||
| if seqlen_k.dtype != torch.int32 or seqlen_k.device != q.device or seqlen_k.numel() != B: |
There was a problem hiding this comment.
This tightens the existing BF16/F16 paged ABI, not just the new FP8 one: seqlen_k must now be exactly int32 with numel() == B. Previously it was only used host-side (seqlen_k.reshape(-1).tolist(), and only when max_seqlen_kv is None), so an int64 tensor or an over-allocated buffer worked. The in-repo fixtures all pass int32 [B] so nothing here regresses, but external callers on the already-shipped paged path could start raising.
Same question for the block_table.dim() != 2 check just above. If the tightening is deliberate, it would be worth a line in the PR description — right now the description says the dense FP8 path is unchanged but says nothing about the paged BF16 ABI getting stricter.
There was a problem hiding this comment.
Addressed in 02ecfe12. The exact device-side int32 [B] requirement is now scoped to paged FP8. Existing BF16/F16 paged calls continue to accept host-only int64 or over-allocated seqlen_k;
the new regression covers both dtypes. I retained the 2D/same-device block-table validation because that is the documented layout consumed directly by the kernel.
| q_flat = q.contiguous() | ||
| out_dtype = torch.bfloat16 if paged_fp8 else q.dtype | ||
| out = torch.empty(expected_out_shape, dtype=out_dtype, device=q.device) | ||
| elif out.shape != expected_out_shape or not out.is_contiguous(): |
There was a problem hiding this comment.
Same class of change as the seqlen_k check: a caller-supplied out must now be contiguous with an exact shape match, on every paged path including the pre-existing BF16 one. This is strictly better than the old o_flat = out.contiguous(), which silently wrote results into a throwaway copy for a non-contiguous out, so I read it as a genuine fix rather than a regression — just noting it is a user-visible behaviour change beyond the PR's stated scope.
There was a problem hiding this comment.
Agreed that this is caller-visible. I retained the exact-shape and contiguous-output checks because the prior out.contiguous() behavior could write into a temporary and return the caller's untouched tensor.
Commit 68f0587d also validates that out is on the Q/K/V device.
jhinpan
left a comment
jhinpan
left a comment
There was a problem hiding this comment.
REQUEST_CHANGES after three independent passes (host/API boundaries, kernel/addressing/numerics, and tests/CI). The new paged-FP8 math itself held up on a real MI355X/gfx950: the full attention file passed 103/103, the >4 GiB page-offset test passed, and extra partial-page/GQA/fully-masked/peaked-logit/V128/V192 probes matched an independent PyTorch oracle (worst max abs error 7.8125e-3). Repo and Python-style checks also pass.
Two host-side correctness defects remain: flattened Q/O can overflow the signed-int32 C ABI, and internal contiguous/cast copies are not ordered before an explicit side-stream launch. Both are reproduced inline below. The PR is also still conflicted with current main, where #1020 added the analogous dense-FP8 overflow handling; rebasing alone will not cover this paged early-return path.
| # Keep serving-sized physical K/V caches rank-5 because flattening their | ||
| # dynamic memref shape can exceed signed int32. The FP8 schedule consumes | ||
| # Q/O as flat token-major buffers, matching its explicit runtime strides. | ||
| q_flat = q.contiguous().view(-1) if paged_fp8 else q.contiguous() |
There was a problem hiding this comment.
Blocking — internal copies race an explicit stream. q.contiguous() here (likewise K/V and the block_table.to(...).contiguous() above) executes on the ambient current stream, while exe(..., stream=launch_stream) immediately consumes the result on the caller-supplied stream. The new side-stream test only uses already-contiguous inputs, so every copy is a no-op and cannot catch this.
I reproduced this on MI355X/gfx950 with a stride-2 Q view. With no pending default-stream work, the side-stream result was bit-identical to the contiguous reference. After enqueueing torch.cuda._sleep(3_000_000_000) on the default stream, side.synchronize() returned in 0.043 s (so it did not wait for that copy) and the same call was no longer bit-identical (max_abs_err=0.0045166015625). Running the call inside with torch.cuda.stream(side): restored bit equality, confirming the missing ordering edge.
Please enqueue all wrapper-owned conversions/copies on launch_stream (including the block-table cast), or reject layouts that require a copy, and add a delayed non-current-stream regression with a non-contiguous input.
There was a problem hiding this comment.
Fixed in 1eacd9d0. All wrapper-owned casts, contiguous copies, output allocation, and workspace allocation now run inside torch.cuda.stream(launch_stream).
The regression uses non-contiguous Q/K/V, an int64 block table, and a blocked default stream; the side-stream result is bit-identical to the contiguous reference without waiting for the default stream.
| k_flat = k.contiguous() | ||
| v_flat = v.contiguous() | ||
| o_flat = out.contiguous() | ||
| o_flat = out.contiguous().view(-1) if paged_fp8 else out.contiguous() |
There was a problem hiding this comment.
Blocking — the flattened paged-FP8 ABI has no signed-int32 dimension guard. These 1-D Q/O views put numel() into the dynamic memref dimension, which the C ABI packs as int32. On the real gfx950 I used a valid packed Q shape [699051, 16, 192] (2,147,484,672 elements, just 1,024 above 2**31) with a one-page KV cache; the call reached <flydsl-cabi-fill> and failed with struct.error: 'i' format requires -2147483648 <= number <= 2147483647.
Current main's #1020 added _FP8_MAX_FLAT_ELEMS for the same dense-FP8 failure, but its guard explicitly excludes paged KV and the paged dispatch returns before the dense path, so rebasing does not fix this. Please reject every flattened Q/O dimension >= 2**31 before build/launch (or keep a non-flattened ABI) and add a lowered-limit regression so CI does not need multi-GiB allocations.
There was a problem hiding this comment.
Fixed in 2f03af8b. Paged FP8 now checks flattened Q and output sizes before build/launch and raises a descriptive NotImplementedError at 2**31.
A lowered-limit V128/V192 regression covers the guard, and the original [699051, 16, 192] reproducer now fails at the host boundary instead of inside C-ABI packing.
Signed-off-by: Xiake Sun <xiake.sun@amd.com>
- Calculate the paged K/V cache footprint and require 2 GiB of additional headroom before running the >4 GiB addressing regression. - Skip with a clear diagnostic when the selected GPU has insufficient free memory. Signed-off-by: Xiake Sun <xiake.sun@amd.com>
- Run wrapper-owned casts, contiguous copies, output allocation, and workspace allocation on the selected stream. - Reuse validated caller output directly without creating a throwaway contiguous copy. - Add delayed side-stream coverage for non-contiguous Q/K/V and block-table conversion. Signed-off-by: Xiake Sun <xiake.sun@amd.com>
- Reject flattened Q or output dimensions at the signed-int32 C ABI boundary. - Raise a descriptive error before kernel build or launch. - Add lowered-limit V128 and V192 regression coverage. Signed-off-by: Xiake Sun <xiake.sun@amd.com>
- Remove redundant paged conditionals. - Reuse shared K-address and sigma-permutation helpers. - Correct the D192 padding comments without changing the schedule. Signed-off-by: Xiake Sun <xiake.sun@amd.com>
- Apply the strict int32 [B] seqlen_k contract only to paged FP8. - Preserve existing BF16/F16 acceptance of host-only length metadata. - Retain common block-table and output-buffer safety validation. Signed-off-by: Xiake Sun <xiake.sun@amd.com>
32dac38 to
68f0587
Compare
- Reject caller-provided paged outputs that are not colocated with Q/K/V. - Prevent cross-device output buffers from reaching the kernel launch. - Add host-side regression coverage for a CPU output tensor. Signed-off-by: Xiake Sun <xiake.sun@amd.com>
jhinpan
left a comment
jhinpan
left a comment
There was a problem hiding this comment.
APPROVE. Both blocking items from my CHANGES_REQUESTED review at 32dac38 are fixed at 68f0587, the five non-blocking items are fixed as well, and the BF16 paged ABI regression I flagged as an unstated side effect has been deliberately narrowed. One verification gap is stated at the bottom.
Prior items, re-checked at this head
1. Blocking — internal copies race an explicit stream (comment) — fixed. Every wrapper-owned conversion now runs inside with torch.cuda.stream(launch_stream): (kernels/attention/flash_attn_interface.py:702-722): the block_table.to(torch.int32).contiguous().reshape(-1) cast, the torch.empty output allocation, and q/k/v.contiguous(). Nothing device-side is left on the ambient stream between the guards and exe(..., stream=launch_stream) — the only remaining pre-launch host interaction is seqlen_k.reshape(-1).tolist(), which is a blocking D2H and is order-safe by construction. The regression exists and is the right shape: test_paged_fp8_d192_asymmetric_value_matches_torch[side-copies] (tests/kernels/test_flash_attn_fwd.py) builds stride-2 non-contiguous Q/K/V plus an int64 block_table (so all four copies are real), enqueues torch.cuda._sleep(3_000_000_000) on the default stream, asserts not default_done.query() (the side stream did not wait for it), and then asserts rtol=0, atol=0 against a contiguous reference. Under the old code the copies would sit behind that sleep and the kernel would read uninitialized memory, so the test discriminates.
2. Blocking — no signed-int32 dimension guard on the flattened paged-FP8 ABI (comment) — fixed. flash_attn_interface.py:628-636 computes q_flat_elems = q.numel() and out_flat_elems = q_flat_elems // D * value_head_dim and raises NotImplementedError when either is >= _FP8_MAX_FLAT_ELEMS (2**31, line 47), before any build or launch. The shape I reproduced with, [699051, 16, 192] = 2,147,484,672 elements, is 1,024 above the bound and is rejected. Because paged FP8 requires packed varlen (q.dim() == 3, line 494) the numel() // D * value_head_dim derivation is exact for both the allocated and the caller-supplied out, and the guard sits before the out shape check so the derived value is the one that will actually be flattened. K/V stay rank-5 and the comment at lines 715-718 now says why. test_paged_fp8_d192_rejects_flattened_int32_overflow monkeypatches _FP8_MAX_FLAT_ELEMS down to the tiny fixture's size, so CI gets the coverage without a multi-GiB allocation.
3. Three no-op const_expr(PAGED) branches (comment) — fixed. floor_masked_max (flash_attn_utils.py:5389), reduce_sum (5419) and apply_l_rescale (5460) are now single unconditional returns. The branches that genuinely differ (sub_m 5392, rescale_from_tile_max 5453, lazy_rescale_o 5487) are untouched.
4. Inaccurate "reads zeros beyond the page descriptor" comment (comment) — fixed. Both sites now state the layout fact rather than a false zero-fill claim: flash_attn_utils.py:4984-4985 ("Each wave owns two 64-lane octet rows for the D256-shaped LDS tile. The QK reader ignores the four groups beyond physical D192.") and the smem_d_rpt comment at 1955-1957.
5. Inlined duplicate of _vec_k_src_elem (comment) — fixed. flash_attn_utils.py:4990-4996 calls _vec_k_src_elem(traits, d, self.wave_id_uni, self.lane_in_warp, self.kv_head_idx). I diffed the removed inline block against the helper body (lines 814-824): the oct_idx/ni/dg/src_oct/src_elem chain is term-for-term identical, so the refactor is behaviour-preserving.
6. Open-coded sigma bit-shuffle (comment) — fixed. flash_attn_utils.py:5112 is now bf = [bf[_sigma_k_tile_n(i)] for i in range_constexpr(16)]. The helper's extra (ni & ~15) term is zero over i < 16, so the substitution is exact.
7. seqlen_k / out ABI tightening on the pre-existing BF16 paged path (1, 2) — addressed. The strict int32 [B] requirement is now gated behind paged_fp8 (flash_attn_interface.py:615), so an int64 or over-allocated host-only seqlen_k still works on BF16/F16, with test_paged_legacy_dtype_accepts_host_only_seqlen_metadata pinning it. The out contiguity/shape/device checks do still apply to every paged path; as I said last round, that replaces a silent write-into-a-throwaway-copy, so I read it as a fix and am not blocking on it.
What I ran here
python3 scripts/check_repo.py—agent-docspasses.typed-arithmeticcannot resolveorigin/mainin this pinned single-commit checkout, so I reproduced it by hand: parsed the base…head patch forkernels/**, fed the 1,498 added lines tocheck_typed_arithmetic_usage.scan_source, zero violations.python3 -m py_compileon all four changed files — clean.- Structural diff of
flash_attn_utils.py/flash_attn_interface.py/flash_attn_fp8_paged_gfx950.pybetween32dac38and this head, to confirm the only semantic changes since my last review are the six items above plus the rebase onto main'srescale_threshold/_P_HEADROOM_LOG2work. The PR-owned part of that rebase isrescale_threshold=8.0for the paged FP8 module; paged PV is HIPREC (bf16 P), so the e4m3 headroom reasoning that motivates the dense 6.0/4.0 split does not apply there and 8.0 is consistent with the unscaledm_diffcomparison in_lazy_rescale_o_paged. - GitHub check runs at this commit:
Check Python Code StyleandCheck C++ Code Styleboth green.
Verification gap
This environment has no torch and no pytest, so I could not execute the new gfx950 tests, and unlike my last round I could not re-run the two GPU repros. Both fixes are closed above by code inspection plus a reading of the new regressions, not by execution. The test job at this commit is still in_progress; please confirm it goes green before merging — in particular test_paged_fp8_d192_asymmetric_value_matches_torch[side-copies-*], which is the one whose assert not default_done.query() could in principle be flaky if a runner is slow enough to drain a 3e9-cycle sleep during the side-stream launch.
|
@sammysun0711 seems large perf gap . 30-50% loss. have you validated this? not make sense to regression too much |
- Load V192 LDS fragments in two smaller groups. - Overlap tail V reads with softmax computation. - Enable post-RA scheduling and disable machine sinking. - Tune D192 setprio and long-context V192 staggering. - Preserve D192/V128 correctness and performance. Signed-off-by: Xiake Sun <xiake.sun@amd.com>
- Accept vectorized page-64 D128/V128 inputs on gfx950. - Use native FP8 P*V for D128 paged attention. - Apply V descale during final output normalization. - Convert vectorized V into the FP8 MFMA operand layout. - Cover ragged batches and explicit-stream execution. Signed-off-by: Xiake Sun <xiake.sun@amd.com>
- Process adjacent page-64 tiles through a BN128 softmax iteration. - Reuse paired 64-bit block-table loads for K and V. - Stage V through an 80-byte bank-padded D-major LDS layout. - Use DPP lane exchange instead of LDS cross-lane shuffles. - Validate direct BN128 launch assumptions. - Cover even-page fast dispatch, odd-page fallback, partial tails, and graph replay. Signed-off-by: Xiake Sun <xiake.sun@amd.com>
- Keep post-P*V accumulators in SSA for the paired-page path. - Reduce wait, MFMA, and packed-add stalls without increasing resources. - Improve D128 throughput across the 4K, 16K, and 64K targets. Signed-off-by: Xiake Sun <xiake.sun@amd.com>
- Load each 32-byte V fragment with two aligned 128-bit LDS operations. - Split vector loads into the four i64 FP8 MFMA K-steps. - Reduce V-fragment address and wait overhead without increasing resources. - Improve D128 throughput across the 4K, 16K, and 64K targets. Signed-off-by: Xiake Sun <xiake.sun@amd.com>
- Map each 16-lane row to one vectorized V token group. - Replace four DPP exchanges and selects with two permlane16 swaps. - Require a full EXEC mask through the static D128/V128 tile shape. - Improve D128 throughput across the 4K, 16K, and 64K targets. Signed-off-by: Xiake Sun <xiake.sun@amd.com>
- Store permlane16 swap results without a staging-register reshuffle. - Apply the matching fixed word permutation to packed P. - Preserve the FP8 contraction mapping and full-EXEC requirement. - Improve D128 throughput at the 16K and 64K targets. Signed-off-by: Xiake Sun <xiake.sun@amd.com>
- Split vectorized V loads from permutation and LDS staging. - Hide V4 latency behind K prefetch and V5 latency behind QK MFMAs. - Keep one prefetched V vector live to preserve register and scratch usage. - Improve random-input D128 latency across the 4K, 16K, and 64K targets. Signed-off-by: Xiake Sun <xiake.sun@amd.com>
- Keep paged BN128 QK accumulators unscaled through row-max reduction. - Fold the runtime logit scale into score subtraction and max correction. - Remove eager per-score multiplies and reduce VGPR use from 246 to 218. - Reach 92.57% of dense throughput at the 64K/128K target. Signed-off-by: Xiake Sun <xiake.sun@amd.com>
- Pair each future V load with an independent QK contraction. - Keep only one prefetched V vector live at a time. - Preserve 218-VGPR, zero-scratch kernel resource usage. - Improve the 16K and 64K random-input attention targets. Signed-off-by: Xiake Sun <xiake.sun@amd.com>
- Keep generic paged QK accumulators in their raw domain. - Fold logit scaling into score subtraction and max correction. - Compare lazy-rescale thresholds in scaled log2 units. - Improve D192 steady-state latency and reduce V192 scratch usage. Signed-off-by: Xiake Sun <xiake.sun@amd.com>
- Add an explicit bank-padded FP8 V-layout trait. - Reuse vectorized V staging for D192/V128 and generic D128 fallback. - Pack softmax probabilities in the matching MFMA word order. - Improve D192/V128 latency while preserving accuracy and zero scratch. Signed-off-by: Xiake Sun <xiake.sun@amd.com>
- Reuse the compiled paged launcher instead of resolving the FlyDSL JIT cache key on every invocation. - Cache per-device architecture detection and reuse it during paged dispatch. - Avoid redundant current-stream context setup while preserving explicit-stream ordering. - Reduce 4K host submission overhead to reach over 90% of D128 throughput compare with dense attention Signed-off-by: Xiake Sun <xiake.sun@amd.com>
Thanks for your feedback. I understand the original performance gap refer to
Here are the final performance evaluation results on MI350X (gfx950) with random per-tensor OCP E4M3 inputs, preallocated output, 10 warmups, 30 iters. Dense comparison is available only for D128/V128 because the upstream dense kernel does not support Q/K D192. D128
D192
|
- Load and consume each BF16 V192 P*V fragment immediately to reduce register pressure and scratch usage. - Enable the schedule for paged FP8 D192/V192 at KV lengths of 64K and above. - Share the long-context threshold with stagger selection and cover its dispatch boundary. - Validate random pages, partial tails, explicit streams, and graph replay. Signed-off-by: Xiake Sun <xiake.sun@amd.com>
|
|
||
| import flydsl.compiler as flyc | ||
| import flydsl.expr as fx | ||
| from flydsl.compiler.kernel_function import CompilationContext |
| from flydsl.expr import const_expr, range_constexpr, rocdl | ||
| from flydsl.expr.typing import T | ||
| from flydsl.expr.utils.arith import ArithValue | ||
| from flydsl.expr.utils.arith import _to_raw as _raw |
| inv_l_rcp = rocdl.rcp(T.f32, _raw(l_row)) | ||
| inv_l = ArithValue(fx.Float32(l_row) > ctx.c_zero_f).select(inv_l_rcp, ctx.c_zero_f) | ||
| if const_expr(traits.FP8_PV): | ||
| inv_l = ArithValue(inv_l) * ctx.vd_fp8 |
| result.append(self.anchor_scalar_f32(m_new)) | ||
| return result | ||
|
|
||
| result = scf_if_dispatch( |
|
much better perf now. still need more clean. use flydsl kernel code cleanup to do a deeper clean. |
- Factor value dimensions into generic H1/H2 segments and avoid padded D256 P*V work. - Reuse full-CTA H1 staging and stage H2 with complete active waves in the bank-padded FP8 layout. - Distribute FP8 P*V chunks across the existing pipeline and disable the BF16-specific wave stagger. - Preserve dense FP8 and other paged dimensions while covering factorization, dispatch, streams, tails, and graph replay. Signed-off-by: Xiake Sun <xiake.sun@amd.com>
- Carry two future V fragments across BN128 loop iterations. - Materialize prefetched fragments into bank-padded LDS after QK work. - Keep paired page-table indexing in the int32 kernel ABI. - Preserve D192 paths and the generic ragged/odd-page fallback. Signed-off-by: Xiake Sun <xiake.sun@amd.com>
- Decouple BN128 tile selection from packed-varlen metadata. - Load per-sequence Q/KV offsets and causal lengths for B>1. - Preserve active q-block-aware longest-first scheduling. - Maintain B=1-relative performance through batch size 32. - Cover ragged multi-q-block batches with distinct causal offsets. Signed-off-by: Xiake Sun <xiake.sun@amd.com>
- Interleave up to eight MiMo batches on the head-fast grid axis. - Issue active causal q-blocks longest-first within each batch group. - Preserve the cache-locality-favored V192 mapping above batch size 16. - Keep divisor-safe fallbacks for arbitrary batch counts. - Cover dispatch policy and ragged multi-q-block D192 accuracy. Signed-off-by: Xiake Sun <xiake.sun@amd.com>
- Reuse the shared safe reciprocal helper for output normalization. - Remove private arithmetic conversion imports from the paged kernel. - Express paged lazy rescaling with typed FlyDSL if/else state. - Preserve generated ISA, resources, numerics, and performance. Signed-off-by: Xiake Sun <xiake.sun@amd.com>
5130cbc to
840d793
Compare
- Delete split-K combine imports, kernel, and launch branches. - Remove dead workspace and partial-output plumbing. - Keep the existing num_kv_splits=1 validation contract. - Preserve generated ISA, resources, numerics, and performance. Signed-off-by: Xiake Sun <xiake.sun@amd.com>
|
I've further optimized D192/V192 and D128/V128 performance and extend B>1 support as follow:
Here are performance evaluation results on MI350X (gfx950)
B=1-32, Hq=16, Hkv=1, Page size=64, Q/KV=4K/8K
|
Motivation
Based on #1065.
FlyDSL FA cached prefill currently cannot consume MiMO-V2.5-Pro physical FP8 Q/K192-V128 pages through the public FlyDSL flash-attention API. Padding V to 192 wastes cache memory and bandwidth.
This PR aims to add gfx950 causal packed-varlen FP8 flash attention for page-64 vectorized-5D KV with Q/K D192 and independent V/output D128 or D192.
Technical Details
flydsl_flash_attn_funcusing the existing 2Dblock_tableABI.Initial scope is causal packed-varlen attention, OCP E4M3 Q/K/V, page size 64,
vectorized-5D KV, Q/K D192, V/O D128 or D192, one KV split, and scalar descales.
Performance (MI350X/gfx950)
Performance evaluation result with random per-tensor OCP E4M3 inputs, preallocated output, 10 warmups, 30 iters. Dense comparison is available only for D128/V128 because the upstream dense kernel does not support Q/K D192.
B=1, Hq=16, Hkv=1, Page size=64
B=1-32, Hq=16, Hkv=1, Page size=64, Q/KV=4K/8K
Test Plan
python -m pytest -q tests/kernels/test_flash_attn_fwd.pyTest Result
Accuracy verification with pytorch reference:
V128 error is the expected direct-FP8 P×V tradeoff, Dense D128 has the same
0.002319maximum error against the PyTorch reference, while D192/V192 retains BF16 P×V.Submission Checklist