Repository navigation
[Kernel][FA] Support paged FP8 Flash attention with asymmetric K/V head dim on gfx950 - #1066
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.
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.
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.
32dac38 to
68f0587
Compare
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 |
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
|
|
much better perf now. still need more clean. use flydsl kernel code cleanup to do a deeper clean. |
5130cbc to
840d793
Compare
|
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
|
8d51c4e to
a547839
Compare
621484f to
787efd7
Compare
|
I rebased on latest main to add with dense attention with D192/V128 case: 826efd0
Here are performance evaluation results on MI350X (gfx950)
B=1-32, Hq=16, Hkv=1, Page size=64, Q/KV=4K/8K
Additional odd and large batchesEach request uses Q/KV=4096/8192.
Accuracy verification
|
| for i in range_constexpr(16): | ||
| centered.append( | ||
| fx.Float32( | ||
| llvm.intr_experimental_constrained_fma( |
There was a problem hiding this comment.
could this inline asm replaced with fx.fma?
There was a problem hiding this comment.
Fixed, replace it with fx.fma in: e41c748
Signed-off-by: Xiake Sun <xiake.sun@amd.com>
- Compact K staging and XOR V addresses to fit six native FP8 K/V slots in LDS. - Stage future V directly into LDS and load current fragments after softmax. - Retain the public page-cache ABI and generic odd-page fallback. - Add ragged multi-KV-head accuracy coverage and pass the expanded attention suite. Signed-off-by: Xiake Sun <xiake.sun@amd.com>
- Derive native value staging from the supported head dimensions. - Remove obsolete BF16 V modes, unreferenced staging helpers, and copied trait state. - Drop redundant dimension guards and unused Q LDS scaffolding. - Preserve public attention behavior and native instruction schedules. Signed-off-by: Xiake Sun <xiake.sun@amd.com>
- Replace PR-introduced index wrappers with explicit 64-bit integer values. - Keep the ring state typed across loop boundaries and retain 32-bit ABI casts. - Use typed vector conversion and remove redundant scalar wrappers. Signed-off-by: Xiake Sun <xiake.sun@amd.com>
- Gather vectorized page caches in one independent reference helper. - Reuse the existing cross-length PyTorch attention oracle. - Preserve ragged, graph, long-context, and multi-KV-head coverage. Signed-off-by: Xiake Sun <xiake.sun@amd.com>
- Construct bounded page resources with 64-bit rebasing and make_buffer_tensor. - Use aligned 16-byte LDS tiles, typed vector loads, and fx.copy stores. - Replace manual paged control flow and pointer casts with DSL operations. - Share canonical value-width traits and the register-resident Q loader. Signed-off-by: Xiake Sun <xiake.sun@amd.com>
- Normalize launch arguments once for cached launches and explicit compilation. - Preserve the compiled launch fast path and existing argument validation. - Check explicit compilation against PyTorch for both paged schedules and all supported widths. Signed-off-by: Xiake Sun <xiake.sun@amd.com>
- Remove stale BF16 V, self-attention, and per-step schedule narration. - Document aligned LDS views and the remaining intrinsic boundaries. - Keep cache ABI, synchronization, and numerical rationale. Signed-off-by: Xiake Sun <xiake.sun@amd.com>
- Simplify the segmented V chunk address calculation and remove redundant segment offsets. - Keep lane, buffer, and chunk offsets in 16-byte tile units so LLVM shares LDS addresses. - Reduce generic V192 spilling while preserving wide loads and odd-page accuracy. Signed-off-by: Xiake Sun <xiake.sun@amd.com>
- Reduce both score tiles through a balanced max tree before one half-wave exchange. - Preserve causal masking, lazy rescaling, and the native FP8 cache layouts. Signed-off-by: Xiake Sun <xiake.sun@amd.com>
- Separate below-diagonal pairs from the masked tail to expose QK and max-reduction overlap. - Carry the same K/V ring state and preserve one cross-wave rendezvous per pair. - Keep D128 and V192 on their compact single-loop schedule. Signed-off-by: Xiake Sun <xiake.sun@amd.com>
- Use scalar FP32 FMA for single-sequence D128 and D192 softmax centering. - Keep packed centering for batched D128 and all dense and generic paths. - Localize instruction selection while preserving probability headroom scaling. Signed-off-by: Xiake Sun <xiake.sun@amd.com>
- Enable batch-fast paired-page ordering for the two-request D192/V128 configuration. - Preserve head-fast ordering in the generic fallback and validate packed-varlen group launches. - Cover ragged multiple-query-block batches with shuffled pages on a non-default stream. Signed-off-by: Xiake Sun <xiake.sun@amd.com>
- Bound future FP8 scores with the query L1 norm and the OCP E4M3 finite range. - Select a max-free loop only when every lane proves lazy rescaling unnecessary; retain the original fallback and cache ABI. - Initialize the first pair without infinite-seed fast math and keep proof comparisons strict. - Add ragged, multi-KV-head, mixed-wave, and rescaling regressions; validate 332 tests and 45 of 48 dense-parity comparisons. Signed-off-by: Xiake Sun <xiake.sun@amd.com>
- Use batch-fast grid-X ordering for paired D128 at B=2/3/5 while preserving other dispatch policies and the native cache ABI. - Expand ragged, multi-KV-head, bounded/fallback, mixed-wave, and side-stream-copy regression coverage. - Isolate the stream concurrency test from HIP hardware-queue sharing without changing benchmark priority or assertions. - Validate all 48 matched comparisons at or above dense throughput, with 359 tests passing and four existing skips. Signed-off-by: Xiake Sun <xiake.sun@amd.com>
- Share probability packing and active query-block ordering between paged schedules. - Remove no-op segmented V load-step wrappers while preserving all loads, waits, and rescaling order. - Retain scheduler-sensitive address expressions and vector construction; nine checked ISA files remain byte-identical. - Validate the complete paged subset with 124 passing tests and four existing skips. Signed-off-by: Xiake Sun <xiake.sun@amd.com>
- Centralize paired-batch interleave selection and reuse the validated page count. - Separate reusable numerical checks from pytest entry points and parameterize duplicate ragged cases. - Add 18 paired-policy cases without changing numerical assertions or stream concurrency checks. - Verify dispatch equivalence across 3840 shape, batch, and head combinations. Signed-off-by: Xiake Sun <xiake.sun@amd.com>
- Trim duplicated stream and phase-shift narration while retaining the barrier and cache-layout rationale. - Clarify asymmetric V shapes, packed probability order, and on-demand segmented V reads. - Document the intentionally separate tail address expression that avoids added spills. - Verify that the executable AST is unchanged; only comments and docstrings differ. Signed-off-by: Xiake Sun <xiake.sun@amd.com>
- Use generic_load and generic_store for the four typed, 16-byte-aligned paged FP8 LDS access sites. - Remove the temporary register tensor used solely to issue the V LDS store while preserving bounded global loads and all synchronization. - Retain the exact qualified prototype implementation: nine ISA files are byte-identical and all 48 comparisons remain above dense throughput. - Recheck the integrated attention suite with 377 passing tests and four existing skips. Signed-off-by: Xiake Sun <xiake.sun@amd.com>
- Replace the staggered odd-page fallback with one bounds-safe paired-page pipeline and remove obsolete helpers, scheduling traits, and dispatch options. - Use fx.gemm with MFMA_Scale, typed LLVM memory primitives, causal selects, and BF16 conversions; replace paged inline assembly with standard constrained-FMA and arithmetic-fence intrinsics. - Preserve the native cache ABI, stream behavior, and strict query-bound proof while adding odd-page, multi-KV-head, eager/lazy, and no-inline-assembly regressions.. Signed-off-by: Xiake Sun <xiake.sun@amd.com>
- Support native physical page sizes 1, 16, 64, and 1024 with independent QK/V widths and bounds-safe cache addressing. - Replace page-1 byte gathers with vector loads and an exact wave transpose; specialize page-16 K/V staging and reuse D128 page translations. - Use bounded cache descriptors for small pools while retaining the validated 64-bit fallback for larger caches. - Preserve page-64/page-1024 ISA and qualify accuracy, stream/graph behavior, large offsets, and the complete performance matrix. Signed-off-by: Xiake Sun <xiake.sun@amd.com>
- Use DPP byte transposes and width-specific LDS word placement for native page-1 attention. - Reuse small-page K/V translations and bounded cache resources while preserving the 64-bit fallback. - Shorten page-16 fragment lifetimes and stage V192 segments with uniform scalar K offsets. - Simplify unused launch plumbing and supported V-width traits while retaining official FlyDSL APIs. Signed-off-by: Xiake Sun <xiake.sun@amd.com>
- Honor packed cu-seqlens for every batch, including B=1, and validate launch metadata, head counts and supported public options. - Bound per-request page-table lookups and clear inactive V bytes, preserve full-packet fast path. - Normalize strided cumulative metadata and single-value descales, preserve graph replay and native page-1/16/64/1024 cache layouts. - Use public fx.fma with explicit nearest-even rounding and load the scalar V descale through fx.generic_load at the paged epilogue. - Remove unused paged debug counters and obsolete B=1/varlen plumbing; preserve lazy rescaling, non-FP8 debugging and dense split-K handling. Signed-off-by: Xiake Sun <xiake.sun@amd.com>
913775a to
08ac7cb
Compare
- Centralize KV layouts and simplify prefix/tail handling. - Stream V fragments for a common D192/V128 page schedule. - Deduplicate tests while preserving boundary and fallback coverage. Signed-off-by: Xiake Sun <xiake.sun@amd.com>
Signed-off-by: Xiake Sun <xiake.sun@amd.com>
Flydsl page attention kernel now supports page size 1 with linear kv layout, 16/64/1024 with vectorized 5d kv layout. For page size 1: I compared that both FlyDSL paged attention with AITER mha batch prefill with page size 1 vs 16, we can see 10+% latency. With page size 1. A 64-token KV tile spans 64 physical pages; with page size 16, it spans only four. Page 1 therefore needs more page-table lookups, address calculations, and data rearrangement to feed the matrix instructions.
But flydsl paged attention is still faster than aiter mha with same shape and page size setting.
|
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 16/64/1024 vectorized-5D KV and page 1 linear3d KV with Q/K D192 and independent V/output D128 or D192.
Technical Details
Performance (MI350X/gfx950)
Performance evaluation result with random per-tensor OCP E4M3 inputs, preallocated output, 10 warmups, 30 iters.
Here are performance evaluation results on MI350X (gfx950)
B=1, Hq=16, Hkv=1, Page size=64
B=1-32, Hq=16, Hkv=1, Page size=64, Q/KV=4K/8K
Hq=16, Hkv=1, causal Q/KV=1:2, FP8, linear3d KV with page size 1
Test Plan
python -m pytest -q tests/kernels/test_flash_attn_fwd.pyTest Result
Accuracy verification with pytorch reference:
Submission Checklist