feat(NPU): add UB Manager for auto tiling strategy management - #987
Conversation
88e9158 to
161eff9
Compare
Ascend NPU UB Manager Design DocumentOverviewThe UB Manager (Unified Buffer Manager) is a core component in Liger-Kernel responsible for managing the Unified Buffer (UB) capacity on Ascend NPUs. By automatically detecting UB capacity and providing best-practice-based tiling strategies, it helps Triton kernels avoid UB overflow errors while maintaining high performance. Design Goals
Architecture DesignCore ComponentsClass DiagramCore Functionality1. UB Capacity DetectionThe UB Manager detects UB capacity in the following priority order:
2. Strategy Registration SystemStrategies are registered via the Fixed StrategyReturns fixed tiling parameters directly, suitable for simple scenarios. Conditional StrategyThe strategy function dynamically computes tiling parameters based on input arguments: 3. Strategy Lookup FlowUsage ExamplesBasic UsageUsage Inside a KernelExtension GuideAdding a New Kernel Strategy
Register the strategy: Use the strategy in the kernel: |
70b8d01 to
56431aa
Compare
Test Result:In PR #986, the accuracy tolerance for the GEGLU operator was updated. On NPU, bfloat16 execution of the GEGLU operator involves low-precision accumulation in the underlying matmul, and the stricter tolerance may therefore lead to false negatives in correctness tests. For this reason, the GEGLU test continues to use the previous tolerance setting instead of the updated one. |
|
Hi @Tcc0403 @zheliuyu @TianHao324, Merry Christmas Eve! 🎄 |
Implement UB Manager to automatically handle UB overflow issues on Ascend NPU by providing dynamic tiling strategies based on UB capacity and operator parameters. - Add UBManager class with automatic UB capacity detection - Implement UB-aware GEGLU and ROPE operators with internal tiling - Support dynamic block size calculation based on UB constraints - Use 80% safety margin for conservative memory estimation This prevents MLIRCompilationError: ub overflow while maintaining performance. Co-authored-by: TianHao324 <854531745@qq.com>
d065ae6 to
fbb9680
Compare
Tcc0403
left a comment
There was a problem hiding this comment.
This is cool! Overall lgtm, just some questions about caching and parameter normalization.
Regarding tolerances in geglu test, I think #986 doesn't comletely fix the numerical issue, it still needs further investigation. It's fine to relax the tolerance as you need.
- Add LRU cache to cache strategy computation results (default 128 entries) - Simplify get_tiling_strategy by removing unnecessary parameter conversions - Move parameter handling logic into strategy functions for better separation of concerns
|
Hi @Tcc0403, Your suggestions were very helpful. I’ve updated the code accordingly, added caching, and improved the overall functionality. |
My pleasure! Thank you for the contribution during holidays as well. Merry Christmas! Just one more thing I forgot to mention, we should include your ub manager document for future contributors. You can put it under |
|
@Tcc0403, Thank you so much! I really appreciate it 😊 |
|
Fantastic! Thank you so much! |
Tcc0403
left a comment
There was a problem hiding this comment.
Waiting for refactoring
- Remove strategy registry and kernel_name-based lookup - Remove LRU cache system for simplicity - Implement unified compute_default_tiling_strategy function - Refactor parameter structure: split tiling_dims and unit_params - Extract dtype_size and memory_multiplier as separate parameters - Update all operator calls (GEGLU, ROPE) to use new interface - Update design documentation to reflect new architecture The new design uses a single unified strategy function for all kernels, making the code simpler and more maintainable. All kernels now directly call compute_default_tiling_strategy with explicit parameters instead of using kernel_name lookup.
|
Hi @Tcc0403 , requesting review for this UB Manager refactoring. Changes:
All tests pass for GEGLU and ROPE. See |
Tcc0403
left a comment
There was a problem hiding this comment.
Sorry for late response, it took a while to compile my thoughts.
Refactor tiling strategy API to use shapes + tiling_dims (indices) instead of tiling_dims (values) + unit_params. Fixed dimensions are now automatically extracted from shapes. - Change API: shapes + tiling_dims (indices) instead of tiling_dims (values) + unit_params - Return structure matches input shapes (list of lists) - Add _normalize_tiling_dims helper function - Update GEGLU and ROPE operators to use new API - Update design documentation
|
Hi @Tcc0403, Based on the considerations above, I’ve gone ahead and refactored the code following this approach, centering the interface around shapes and tiling_dims while keeping the existing parameters (such as safety_margin, dtype_size, and memory_multiplier) compatible. If you have some time, I’d really appreciate it if you could help review this refactor and share any feedback on the interface design or potential improvements. Thanks a lot. |
Change shapes parameter and return values from List[List[int]] to Tuple[Tuple[int, ...], ...] for better immutability and PyTorch integration.
|
Hi @Tcc0403, We've updated the code to use
The new interface is indeed cleaner and more consistent with PyTorch conventions. Thanks for catching this! |
|
That was fast! Overall LGTM, requesting reviews in case I overlook something. |
Background and Motivation
When developing Ascend NPU operators, we frequently encounter compilation failures caused by UB (Unified Buffer) overflow. During compilation, Triton kernels check UB usage, and if it exceeds the capacity, an error is raised:
MLIRCompilationError: ub overflow.To address this issue, developers usually have to:
This process is tedious, error-prone, and difficult to cover all scenarios.
Solution
We implemented a UB Manager that provides:
Core Features
Automatic Capacity Detection
Dynamic Strategy Computation
n_colsanddtype_sizeBLOCK_QandBLOCK_Kbased onpad_n_q_head,pad_n_kv_head, andpad_hdEasy Extensibility
Implementation Details
ub_manager.py: the core UB management classgeglu.pyandrope.py: integrated UB-aware implementationsTesting
Verified on Ascend NPU 910B4:
make testto ensure correctnessmake checkstyleto ensure code stylemake test-convergenceto ensure convergence