Repository navigation

[Triton/Gluon] Add 3D neighborhood flash attention (na3d_flash) - #5230

Merged

Boss2002n merged 56 commits into

ROCm:mainfrom

jjuvonen-amd:feat/3d-neighborhood-flash-attn

Oct 2, 2026

Merged

[Triton/Gluon] Add 3D neighborhood flash attention (na3d_flash)#5230
Boss2002n merged 56 commits into
ROCm:mainfrom
jjuvonen-amd:feat/3d-neighborhood-flash-attn

Commits on Sep 3, 2026
- add 3d neighborhood flash attention
  jjuvonen-amd
  committed
  9b301cb
- na3d_flash kernel: input type and bounds validation
  
  jjuvonen-amd
  and
  Copilot
  authored
  19f8b88
- use repr= in jit to include autotuned configs in kernel names
  jjuvonen-amd
  committed
  85671cc
- fix cost heuristics in _na3d_sdpa_exact, correct docstring for rightmost block
  jjuvonen-amd
  committed
  90a4b4e
- avoid importing from kernels module in test
  jjuvonen-amd
  committed
  508338d
- black format
  jjuvonen-amd
  committed
  c526551
- fix cross-row bug when W % BLOCK_Q != 0 and add test shape to validate fix
  jjuvonen-amd
  committed
  b448a1b
- remove unnecessary correct_rightmost_block parameter and related tests
  jjuvonen-amd
  committed
  6dda9d3
- Prune autotune configs where BLOCK_KV is too small for KW
  
  jjuvonen-amd
  and
  Copilot
  authored
  2617cc8
- fix global state leak in pytest, fix config prune conditions
  jjuvonen-amd
  committed
  3cf7288
- use actually highest value W in min(), instead of arbitary high value
  jjuvonen-amd
  committed
  510e65a
- Add a fast-fail validation to attn call
  jjuvonen-amd
  committed
  e00b5d3
- remove duplicate check
  
  jjuvonen-amd
  and
  Copilot
  authored
  0846c23
- cast indices to int64 to avoid potential OOB reads/writes
  
  jjuvonen-amd
  and
  Copilot
  authored
  af1eab5
- cast int multiplications to int64 to avoid potential OOB reads/writes
  
  jjuvonen-amd
  and
  Copilot
  authored
  028f9ca
- cast stride_bnh to int64 to avoid exceeding int32 range
  jjuvonen-amd
  committed
  bdd1f75
Commits on Sep 4, 2026
- drop lru_cache on _na3d_sdpa_mask in test file
  jjuvonen-amd
  committed
  82cfbbe
- replace f-string logging with lazy %-style logging
  jjuvonen-amd
  committed
  400e965
- add arch guard to public na3d_flash_attn launcher
  jjuvonen-amd
  committed
  1a3857a
- split test_na3d_flash.py into separate bench and test files
  jjuvonen-amd
  committed
  44e8fe4
Commits on Sep 6, 2026
Commits on Sep 7, 2026
- renamed autotune configs for clarity
  jjuvonen-amd
  committed
  aa8e3fd
- assert lower bounds of KT, KH and KW
  jjuvonen-amd
  committed
  98607a7
- assert HD >= 16 in na3d_flash
  jjuvonen-amd
  committed
  17c39a7
- verify that q/k/v are CUDA tensors
  jjuvonen-amd
  committed
  9035e3b
- fix num_stages forwarding bug in autotune configs
  jjuvonen-amd
  committed
  ce22c69
- use runtime variant of get_gfx
  
  jjuvonen-amd
  and
  Copilot
  authored
  d9b9afd
- fix checkAllClose asserts in tests
  jjuvonen-amd
  committed
  2c18894
- tighten test boundary check from 5% mismatch to 1%
  jjuvonen-amd
  committed
  085f542
Commits on Sep 8, 2026
- correct HBM formula in bench to report effective rate
  jjuvonen-amd
  committed
  c81c79a
- add test shapes for KW > 17 edge case
  jjuvonen-amd
  committed
  c6e7afe
- document na3d_flash_attn is forward-only
  jjuvonen-amd
  committed
  a97d6c6
- use runtime gfx check in bench and tests
  jjuvonen-amd
  committed
  24f3120
Commits on Sep 9, 2026
- fix W.to() crash when W=1, add test shape with W=1
  jjuvonen-amd
  committed
  44c90fd
Commits on Sep 14, 2026
- set default configs and keep opt-in autotune
  jjuvonen-amd
  committed
  e9eff68
- pruner to read named_args and kwargs; derive triton fallback configs from generator
  jjuvonen-amd
  committed
  c307a74
Commits on Sep 17, 2026
Commits on Sep 25, 2026
- Merge branch 'main' into feat/3d-neighborhood-flash-attn
  vgokhale
  authored
  2886f58
Commits on Sep 26, 2026
- tune(na3d_flash): use measured gfx942 tile for the large-KW config (num_stages=4 is 5.6% faster)
  jjuvonen-amd
  committed
  23b9b43
- docs(na3d_flash): describe reference masks as deduplicated, not cached
  jjuvonen-amd
  committed
  a597430
Commits on Sep 30, 2026
- Merge branch 'main' into feat/3d-neighborhood-flash-attn
  jjuvonen-amd
  authored
  1ac0278
- Merge branch 'main' into feat/3d-neighborhood-flash-attn
  Boss2002n
  authored
  1f16800
- Update comments to reflect edge-case coverage changes
  
  jjuvonen-amd
  and
  Copilot
  authored
  26481fb
- Merge branch 'main' into feat/3d-neighborhood-flash-attn
  Boss2002n
  authored
  d239b26
Commits on Oct 1, 2026

Commits

Commits on Sep 3, 2026

Commits on Sep 4, 2026

Commits on Sep 6, 2026

Commits on Sep 7, 2026

Commits on Sep 8, 2026

Commits on Sep 9, 2026

Commits on Sep 14, 2026

Commits on Sep 17, 2026

Commits on Sep 25, 2026

Commits on Sep 26, 2026

Commits on Sep 30, 2026

Commits on Oct 1, 2026