Heterogeneous targets: loop-carried accumulator residency (declared + capacity-checked)
Nobody has claimed this yet.
- Dominant language
- Rust
- Stars
- 14
- Forks
- 2
- Avg merge
- 12h 42m
- Merged PRs (30d)
- 61
Description
Part of hiraditya/Vx.1#256. On the critical path for a real FlashAttention kernel; not covered by any existing P-item.
The gap
Vx can declare where a tensor is placed (transfer(x, Memory::SMEM)) but has no way to say a loop-carried value stays resident in a space across iterations. You can write
Memory RF { within: Memory::L2, capacity: 256 KB, managed: explicit, scope: thread }
and nothing can be pinned to it.
Why it matters
This is the difference between FlashAttention and merely-tiled attention. The online-softmax recurrence carries three values across the entire K/V stream:
m running max (scalar per row)
l running normalizer (scalar per row)
O_acc output accumulator (B_r x d tile) <- the expensive one
If O_acc round-trips to HBM on every KV iteration, the kernel does FlashAttention's arithmetic with worse-than-baseline memory traffic — the entire premise of the algorithm is defeated. On B200 the accumulator is supposed to live in TMEM (or registers) for the whole loop and touch memory once at the end.
Today tests/backend/pass/flash_attention_v4.vx threads m/l/o as ordinary mut locals. That is correct on the CPU path and says nothing about residency, so a GPU backend has no information to act on and LLVM would have to infer it.
What's needed
A way to state residency as a placement obligation rather than a hope, e.g.
spawn on(Topology::TensorCore5) {
resident<Memory::TMEM> let mut o_acc = Tensor<f32>([64, 128]);
resident<Memory::RF> let mut m = -inf;
for j in 0..n_kv {
// ... o_acc updated in place, never spilled
}
}
with the compiler:
- checking the residency set fits the space's
capacity— this is exactly the E6010 working-set check, applied to loop-carried values rather than transfers; - emitting the residency so the backend honours it (
gpu.funcprivate/workgroup attribution, or a TMEM allocation vianvvm.tcgen05.alloc); - rejecting — or at minimum warning — when a value marked resident is provably spilled (aliased, address-taken, or too large).
Point 1 is the part Vx is uniquely positioned to do: it already has the capacity algebra. Point 3 is the one that makes it a guarantee rather than a hint, and is what distinguishes this from register in C.
Relationships
- #247 (scope ladder) — residency targets a
scope:-qualified space, so the two want the same vocabulary. - #246 (layout) — the accumulator's byte size depends on layout/packing, so the capacity check needs it.
- #252 (execution model) — once warp roles exist, residency is per-role.
- #254 (
addressable_by:) — TMEM residency is only meaningful if TMEM is modelled as MMA-addressable.
Note on scope
This is deliberately narrower than general register allocation. The ask is a declared, checked residency for values the programmer names — not an optimiser that decides residency on its own. The optimiser version is out of scope; the declaration version is what the FA kernel needs and what Vx's capacity model can actually verify.
Contributor guide
No contributing guide indexed for this repository
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 with tests/backend/pass/flash_attention_v4.vx and review the residency-related requirements in this issue alongside relationships #247, #246, #252, and #254. Trace the existing capacity algebra and placement handling before deciding where a declared, scope-qualified residency obligation belongs. Done means loop-carried values can be declared resident, capacity-checked, emitted for the backend, and rejected or warned on when provably spilled.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- rust
- Domain
- compilers
- Issue type
- Feature
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Quiet
- Clarity
- Mostly clear
- Newbie friendliness
- 45/100