vx-lang / vx-lang/Vx

Heterogeneous targets: loop-carried accumulator residency (declared + capacity-checked)

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

Nobody has claimed this yet.

core-lang enhancement
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:

  1. 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;
  2. emitting the residency so the backend honours it (gpu.func private/workgroup attribution, or a TMEM allocation via nvvm.tcgen05.alloc);
  3. 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

  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 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.