mudler / mudler/vllm.cpp

feat(rocm/gemma4): integrate bc64 flash-attention prefill kernel (2.76x isolated) env-gated on gfx1201

Open
#697 0 comments 0 reactions 0 assignees View on GitHub

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_FA host 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 lab VT_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

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.