Skip to content

fix(npu): hard-cap UB tiles for A3/A5 - #1464

Draft
zheliuyu wants to merge 2 commits into
linkedin:mainfrom
zheliuyu:main
Draft

zheliuyu wants to merge 2 commits into
linkedin:mainfrom
zheliuyu:main

Conversation

@zheliuyu

@zheliuyu zheliuyu commented Sep 12, 2026

Copy link
Copy Markdown
Contributor

Summary

Hard-cap Ascend UB tiles by detected UB / (bits × 4/3). Triton-Ascend 3.2.2 defaults to auto-multi-buffer (~4/3 live tiles). The margin-only path can still pick an A5-sized tile (e.g. GLU fwd 16384) when reported UB is below 256KB. Cap uses detected UB only: A3=192KB, A5=256KB.

Expected GLU tiles (safety_margin=0.9, fp32): A3 8192/4096, A5 16384/8192.

Details

  • Do not use multiplier * 2; that would cut A5 GLU fwd 16384→8192.
  • On A3, 4/3 vs margin can straddle a power of two. Plain CE/FLCE fwd vocab tile goes 16384→8192. SwiGLU / RMSNorm / RoPE tiles stay the same. LCE tiling is follow-up.
  • mHC mix/atomic multibuffer=False launches are unchanged.

Testing Done

  • Hardware Type: Ascend A3 (192KB UB)
  • Hardware Type: Ascend A5 (256KB UB)
  • run make test to ensure correctness
  • run make checkstyle to ensure code style
  • run make test-convergence to ensure convergence
  • Same-machine A/B of ub_manager old vs new: SwiGLU / RMSNorm / RoPE tile+latency unchanged; CE V=32k/128k flat; FLCE BT=2048,H=4096,V=32000 fwd about +3% from the plain-CE fwd tile cut.

@zheliuyu zheliuyu changed the title fix(npu): hard-cap UB tiles for A3/A5 and accept JSD dispatch kwargs fix(npu): hard-cap UB tiles for A3/A5 Sep 15, 2026
@zheliuyu
zheliuyu force-pushed the main branch 3 times, most recently from 60a9049 to d1d4ee8 Compare September 16, 2026 07:21
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant