Repository navigation

Add a fused full-attention path for head_dim 256 on NAX devices - #3842

Merged

zcbenz merged 6 commits into

ml-explore:mainfrom

wyanzhao:nax-256-dsplit

Aug 19, 2026

Merged

Add a fused full-attention path for head_dim 256 on NAX devices#3842
zcbenz merged 6 commits into
ml-explore:mainfrom
wyanzhao:nax-256-dsplit

Commits on Aug 19, 2026
- Add a fused full-attention path for head_dim 256 on NAX devices
  
  wyanzhao
  authored and
  zcbenz
  committed
  55ea7b6
- Split the NAX head-dim-split attention path into its own kernel
  
  wyanzhao
  authored and
  zcbenz
  committed
  a8be3b3
- Drop the head_dim 256 padding copies; read KV-cache slices past their end
  
  wyanzhao
  authored and
  zcbenz
  committed
  667bb61
- nit
  zcbenz
  committed
  7d9ce26
- nit
  zcbenz
  committed
  5193924
- Merge branch 'main' into nax-256-dsplit
  wyanzhao
  authored
  b86d328