feat(rocm/gemma4): integrate bc64 flash-attention prefill kernel (2.76x isolated) env-gated on gfx1201
Nobody has claimed this yet.
- Dominant language
- C++
- Stars
- 423
- Forks
- 53
- Avg merge
- 20h 26m
- Merged PRs (30d)
- 310
Description
Goal
Integrate the validated bc64 flash-attention prefill kernel (striped online softmax + WG256 packaging + PV 8-wave, isolated best 8744us = 2.76x over the current SharedK path at T=2048/d512, identity-verified) into the production ROCm gfx1201 prefill attention path, env-gated, default OFF.
Context
Isolated kernel development (thread on the bus) took the gfx1201 prefill attention from SharedK-class to 2.76x via three evidence-backed fixes: eliminating a 21x-redundant lane-replicated softmax, WG256 active-wave packaging, and PV 8-wave distribution. The remaining ~1.94x to Vulkan/ACO is dynamic-scheduling/codegen and was ruled a ceiling. This issue tracks landing the 2.76x construction into the serving path (projected ~1.63x whole-prefill once integrated).
Scope
- Add the bc64 kernel adjacent to
PagedAttnPrefillSharedKWmma(src/vt/rocm/rocm_paged_attn.hip:1307), adapted to the paged-KV interface (block_table / seq_lens / query_start_loc), not the isolated contiguous form. - Add a default-OFF
VT_ATTN_PREFILL_BC64_FAhost branch inside the eligible BF16/GQA block, before the SharedK dispatch (~:1694). - Supported shapes only: d512 qg8 full window; d256 qg2 sliding window. Unsupported shape or env-absent/0 → byte-identical fallthrough.
Out of scope
- Replacing
PagedAttnPrefillSharedKWmma. Importing the dirty labVT_ATTN_PREFILL_*experiment matrix. Decode path. Non-BF16.
Gates
- off/default path byte-identical to current KEEP.
- on: shape reachability for d512-full + d256-sliding via real block_table/seq_lens/query_start_loc.
- token-exact vs the current SharedK path on real serving prompts BEFORE any perf number.
- same-binary A/B (flag on vs off) product p42k, fair PC=0, reverse-order medians.
Owner
Coordinator + hermes-vllm implement (don-agent offline, developer-directed); research fresh-review + p42k gate rerun.
Kind: feature
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Start in src/vt/rocm/rocm_paged_attn.hip at PagedAttnPrefillSharedKWmma around line 1307, then inspect the host dispatch near line 1694 and the existing paged-KV inputs. Completion requires the default-off path to remain byte-identical, supported shapes to reach the bc64 kernel with real block_table, seq_lens, and query_start_loc data, and token-exact plus same-binary A/B validation to pass before performance comparison.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- cpp
- Domain
- backend, performance
- Issue type
- Feature
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Quiet
- Clarity
- Clearly specified
- Newbie friendliness
- 35/100