Skip to content

feat(NPU): add UB Manager for auto tiling strategy management - #987

Merged
Tcc0403 merged 8 commits into
linkedin:mainfrom
noemotiovon:ub_manager
Dec 29, 2025
Merged

feat(NPU): add UB Manager for auto tiling strategy management#987
Tcc0403 merged 8 commits into
linkedin:mainfrom
noemotiovon:ub_manager

Conversation

@noemotiovon

@noemotiovon noemotiovon commented Dec 24, 2025

Copy link
Copy Markdown
Contributor

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:

  • Manually adjust block sizes
  • Tune separately for different NPU models and input sizes
  • Repeatedly trade off between performance and UB safety

This process is tedious, error-prone, and difficult to cover all scenarios.


Solution

We implemented a UB Manager that provides:

  • Automatic UB capacity detection: retrieved from device properties or environment variables
  • Best-practice-based tiling strategies: dynamically compute safe block sizes based on UB capacity and operator parameters
  • Unified strategy registration system: supports both fixed and conditional strategies, making it easy to extend
  • Integration with GEGLU and ROPE: automatically handles UB constraints and prevents overflow

Core Features

Automatic Capacity Detection

  • Supports three sources: environment variables, device properties, and model default values
  • Supports multiple Ascend models such as Ascend 910B1 / 910B4

Dynamic Strategy Computation

  • GEGLU: computes a safe block size based on n_cols and dtype_size
  • ROPE: computes BLOCK_Q and BLOCK_K based on pad_n_q_head, pad_n_kv_head, and pad_hd
  • Uses an 80% safety margin to balance performance and safety

Easy Extensibility

  • Simple interface: adding a new kernel strategy only requires registering a function
  • Supports parameterized strategies to adapt to different input sizes

Implementation Details

  • Added ub_manager.py: the core UB management class
  • Updated geglu.py and rope.py: integrated UB-aware implementations

Testing

Verified on Ascend NPU 910B4:

  • GEGLU forward and backward pass tests
  • ROPE forward and backward pass tests
  • Works correctly across different input sizes and data types
  • Hardware Type:
  • run make test to ensure correctness
  • run make checkstyle to ensure code style
  • run make test-convergence to ensure convergence

@noemotiovon

noemotiovon commented Dec 24, 2025

Copy link
Copy Markdown
Contributor Author

Ascend NPU UB Manager Design Document

Overview

The 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

  1. Automated UB Management: Automatically detect device UB capacity without manual configuration
  2. Best-Practice-Based: Use proven tiling strategies to avoid UB overflow
  3. Flexible Strategy System: Support both fixed strategies and conditional strategies to adapt to different scenarios
  4. Easy to Extend: Simple interfaces for adding new kernel strategies
  5. Performance Optimization: Maximize performance while ensuring UB safety

Architecture Design

Core Components

┌─────────────────────────────────────────────────────────┐
│                    UB Manager System                    │
├─────────────────────────────────────────────────────────┤
│                                                         │
│  ┌──────────────┐         ┌──────────────────┐          │
│  │  UBManager   │────────▶│ Strategy Registry│          │
│  │   (Singleton)│         │  (Best Practices)│          │
│  └──────────────┘         └──────────────────┘          │
│         │                            │                  │
│         │                            │                  │
│         ▼                            ▼                  │
│  ┌──────────────┐         ┌──────────────────┐          │
│  │   Capacity   │         │  Strategy        │          │
│  │  Detection   │         │  Functions       │          │
│  └──────────────┘         └──────────────────┘          │
│                                                         │
└─────────────────────────────────────────────────────────┘
         │                            │
         │                            │
         ▼                            ▼
┌──────────────┐         ┌──────────────────┐
│   GEGLU      │         │      ROPE        │
│   Kernel     │         │     Kernel       │
└──────────────┘         └──────────────────┘

Class Diagram

┌─────────────────────────────────────┐
│          UBManager                  │
├─────────────────────────────────────┤
│ - _npu_model: str                   │
│ - _ub_capacity_bits: int            │
├─────────────────────────────────────┤
│ + ub_capacity_bits: int             │
│ + ub_capacity_bytes: int            │
│ + npu_model: str                    │
│ + get_tiling_strategy()             │
│ - _detect_npu_model()               │
│ - _detect_ub_capacity()             │
└─────────────────────────────────────┘

Core Functionality

1. UB Capacity Detection

The UB Manager detects UB capacity in the following priority order:

  1. Environment Variable: ASCEND_UB_CAPACITY_BITS
  2. Device Properties: Retrieved from torch.npu.get_device_properties(0).ub_capacity_bits
  3. Model Defaults: Use predefined values based on the detected NPU model
# Default UB capacity configuration
_DEFAULT_UB_CAPACITIES = {
    "Ascend910B1": 2097152,  # ~256 KB
    "Ascend910B4": 1572864,  # ~192 KB
    "default": 2097152,       # ~256 KB
}

2. Strategy Registration System

Strategies are registered via the _TILING_STRATEGY_BEST_PRACTICES dictionary and support two formats:

Fixed Strategy

("kernel_name", ub_capacity_bits): (block_size, ...)

Returns fixed tiling parameters directly, suitable for simple scenarios.

Conditional Strategy

("kernel_name", ub_capacity_bits): strategy_function

The strategy function dynamically computes tiling parameters based on input arguments:

def strategy_function(key_params: Optional[Union[Tuple, Dict]]) -> Optional[Tuple]:
    # Compute the tiling strategy based on key_params and UB capacity
    # Return (block_size, ...) or None

3. Strategy Lookup Flow

User calls get_tiling_strategy()
         │
         ▼
Build lookup key: (kernel_name, ub_capacity_bits)
         │
         ▼
Look up in _TILING_STRATEGY_BEST_PRACTICES
         │
         ├─── Not found ────▶ Return None
         │
         ▼
      Strategy found
         │
         ├─── Fixed tuple ────▶ Return directly
         │
         ▼
     Callable function
         │
         ▼
  Normalize key_params
         │
         ▼
  Call strategy function
         │
         ▼
   Return strategy result

Usage Examples

Basic Usage

from liger_kernel.ops.backends._ascend.ub_manager import get_tiling_strategy

# GEGLU forward
strategy = get_tiling_strategy("geglu_forward", (4096, 2))
if strategy:
    block_size = strategy[0]
    # Call kernel with block_size

# ROPE forward
strategy = get_tiling_strategy("rope_forward", (32, 32, 128, 4))
if strategy:
    BLOCK_Q, BLOCK_K = strategy
    # Call kernel with BLOCK_Q and BLOCK_K

Usage Inside a Kernel

# GEGLU example
def geglu_forward(a, b):
    n_cols = a.shape[-1]
    dtype_size = 4 if a.dtype == torch.float32 else 2
    
    # Get strategy
    strategy = get_tiling_strategy("geglu_forward", (n_cols, dtype_size))
    
    if strategy is not None:
        block_size = strategy[0]
    else:
        block_size = triton.next_power_of_2(n_cols)  # Fallback
    
    # Call kernel
    kernel[(n_rows,)](a, b, c, BLOCK_SIZE=block_size)

Extension Guide

Adding a New Kernel Strategy

  1. Define the strategy function:
def _my_kernel_strategy(key_params: Optional[Union[Tuple, Dict]]) -> Optional[Tuple]:
    """
    My kernel tiling strategy.
    
    Args:
        key_params: (param1, param2, dtype_size) or dict
    
    Returns:
        (block_size1, block_size2, ...) or None
    """
    if key_params is None:
        return None
    
    # Get UB capacity
    ub_manager = get_ub_manager()
    ub_capacity_bits = ub_manager.ub_capacity_bits
    
    # Extract parameters
    if isinstance(key_params, dict):
        param1 = key_params["param1"]
        param2 = key_params["param2"]
        dtype_size = key_params.get("dtype_size", 4)
    else:
        param1 = key_params[0]
        param2 = key_params[1]
        dtype_size = key_params[2] if len(key_params) > 2 else 4
    
    # Compute strategy
    SAFE_UB_CAPACITY = int(ub_capacity_bits * 0.80)
    # ... compute block_size based on memory estimation ...
    
    return (block_size1, block_size2)

Register the strategy:

_TILING_STRATEGY_BEST_PRACTICES = {
    # ... existing strategies ...
    
    # New strategies
    ("my_kernel_forward", 1572864): _my_kernel_strategy,
    ("my_kernel_backward", 1572864): _my_kernel_strategy,
}

Use the strategy in the kernel:

def my_kernel_forward(input):
    # Prepare parameters
    param1 = input.shape[0]
    param2 = input.shape[1]
    dtype_size = 4 if input.dtype == torch.float32 else 2
    
    # Get strategy
    strategy = get_tiling_strategy("my_kernel_forward", (param1, param2, dtype_size))
    
    if strategy:
        block_size1, block_size2 = strategy
    else:
        # Fallback logic
        block_size1, block_size2 = default_sizes()
    
    # Call kernel
    kernel[(grid_size,)](
        input,
        BLOCK_SIZE1=block_size1,
        BLOCK_SIZE2=block_size2,
    )

@noemotiovon
noemotiovon force-pushed the ub_manager branch 2 times, most recently from 70b8d01 to 56431aa Compare December 24, 2025 02:14
@noemotiovon

Copy link
Copy Markdown
Contributor Author

Test Result:

In PR #986, the accuracy tolerance for the GEGLU operator was updated.
The previous tolerance was chosen based on the numerical properties of bfloat16 execution in the GEGLU operator, while the new (atol=1e-2, rtol=1e-2) implicitly assumes FP32-level precision.

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.

TOTAL                                                                8228   7118   2342     34    11%
Coverage HTML written to dir htmlcov
================================================================== slowest durations ===================================================================
11.90s call     test/transformers/test_geglu.py::test_correctness[dtype0-1.0-2e-06-2-2048-2048-4096]
2.11s call     test/transformers/test_geglu.py::test_correctness[dtype0-1.0-2e-06-9-41-341-4231]
2.10s call     test/transformers/test_geglu.py::test_correctness[dtype1-10000.0-0.006-9-41-341-4231]
2.05s call     test/transformers/test_geglu.py::test_correctness[dtype1-10000.0-0.006-2-2048-2048-4096]
0.15s call     test/transformers/test_rope.py::test_correctness[True-dtype0-1e-05-1e-05-1-128-32-32-64]
0.09s call     test/transformers/test_rope.py::test_correctness[True-dtype1-0.1-1e-05-1-128-32-32-64]
0.03s call     test/transformers/test_rope.py::test_correctness[True-dtype0-1e-05-1e-05-3-423-73-213-92]
0.03s call     test/transformers/test_rope.py::test_correctness[False-dtype0-1e-05-1e-05-3-423-73-213-92]
0.02s call     test/transformers/test_rope.py::test_functional_correctness[True-dtype0-1e-05-1e-05-1-2-2-2-8]
0.02s call     test/transformers/test_rope.py::test_correctness[False-dtype0-1e-05-1e-05-3-423-73-155-92]
0.02s call     test/transformers/test_rope.py::test_correctness[True-dtype0-1e-05-1e-05-3-423-73-155-92]
0.02s call     test/transformers/test_rope.py::test_correctness[False-dtype1-0.1-1e-05-3-423-73-213-92]
0.02s call     test/transformers/test_rope.py::test_correctness[True-dtype1-0.1-1e-05-3-423-73-213-92]
0.02s call     test/transformers/test_geglu.py::test_correctness_functional[dtype0-1.0-2e-06-2-2-8]
0.02s call     test/transformers/test_rope.py::test_correctness[True-dtype0-1e-05-1e-05-2-128-32-32-64]
0.02s call     test/transformers/test_rope.py::test_functional_correctness[True-dtype0-1e-05-1e-05-9-7-41-41-41]
0.02s call     test/transformers/test_rope.py::test_correctness[True-dtype1-0.1-1e-05-3-423-73-155-92]
0.02s call     test/transformers/test_rope.py::test_functional_correctness[True-dtype1-0.1-1e-05-1-2-2-2-8]
0.02s call     test/transformers/test_rope.py::test_functional_correctness[True-dtype1-0.1-1e-05-9-7-41-41-41]
0.02s call     test/transformers/test_rope.py::test_functional_correctness[False-dtype0-1e-05-1e-05-9-7-41-41-41]
0.02s call     test/transformers/test_rope.py::test_functional_correctness[False-dtype1-0.1-1e-05-9-7-41-41-41]
0.02s call     test/transformers/test_rope.py::test_functional_correctness[True-dtype1-0.1-1e-05-1-2-1-2-8]
0.02s call     test/transformers/test_rope.py::test_functional_correctness[True-dtype0-1e-05-1e-05-1-2-1-2-8]
0.02s call     test/transformers/test_rope.py::test_correctness[False-dtype1-0.1-1e-05-3-423-73-155-92]
0.01s call     test/transformers/test_geglu.py::test_correctness_functional[dtype1-10000.0-0.006-2-2-8]
0.01s call     test/transformers/test_geglu.py::test_correctness_functional[dtype0-1.0-2e-06-9-7-41]
0.01s call     test/transformers/test_geglu.py::test_correctness_functional[dtype1-10000.0-0.006-9-7-41]
0.01s call     test/transformers/test_rope.py::test_correctness[True-dtype1-0.1-1e-05-2-128-32-32-64]
0.01s call     test/transformers/test_rope.py::test_correctness[True-dtype1-0.1-1e-05-2-128-32-8-64]
0.01s call     test/transformers/test_rope.py::test_correctness[True-dtype0-1e-05-1e-05-2-128-32-8-64]
0.01s call     test/transformers/test_rope.py::test_correctness[True-dtype1-0.1-1e-05-1-128-32-8-64]
0.01s call     test/transformers/test_rope.py::test_correctness[True-dtype0-1e-05-1e-05-1-128-32-8-64]
0.01s call     test/transformers/test_rope.py::test_correctness[False-dtype1-0.1-1e-05-2-128-32-8-64]
0.01s call     test/transformers/test_rope.py::test_correctness[False-dtype1-0.1-1e-05-2-128-32-32-64]
0.01s call     test/transformers/test_rope.py::test_correctness[False-dtype0-1e-05-1e-05-2-128-32-32-64]
0.01s call     test/transformers/test_rope.py::test_correctness[False-dtype0-1e-05-1e-05-2-128-32-8-64]
0.01s teardown test/transformers/test_geglu.py::test_correctness[dtype0-1.0-2e-06-2-2048-2048-4096]
0.01s teardown test/transformers/test_rope.py::test_correctness[True-dtype0-1e-05-1e-05-3-423-73-213-92]
0.01s teardown test/transformers/test_rope.py::test_correctness[False-dtype0-1e-05-1e-05-3-423-73-213-92]
0.01s teardown test/transformers/test_rope.py::test_correctness[True-dtype0-1e-05-1e-05-3-423-73-155-92]
0.01s teardown test/transformers/test_rope.py::test_correctness[False-dtype0-1e-05-1e-05-3-423-73-155-92]
0.01s call     test/transformers/test_rope.py::test_correctness[False-dtype1-0.1-1e-05-1-128-32-32-64]
0.01s call     test/transformers/test_rope.py::test_correctness[False-dtype1-0.1-1e-05-1-128-32-8-64]
0.01s call     test/transformers/test_rope.py::test_functional_correctness[False-dtype1-0.1-1e-05-1-2-2-2-8]
0.01s call     test/transformers/test_rope.py::test_functional_correctness[False-dtype0-1e-05-1e-05-1-2-2-2-8]
0.01s call     test/transformers/test_rope.py::test_functional_correctness[False-dtype1-0.1-1e-05-1-2-1-2-8]
0.01s call     test/transformers/test_rope.py::test_correctness[False-dtype0-1e-05-1e-05-1-128-32-8-64]
0.01s call     test/transformers/test_rope.py::test_correctness[False-dtype0-1e-05-1e-05-1-128-32-32-64]
0.01s call     test/transformers/test_rope.py::test_functional_correctness[False-dtype0-1e-05-1e-05-1-2-1-2-8]
0.01s teardown test/transformers/test_geglu.py::test_correctness[dtype1-10000.0-0.006-2-2048-2048-4096]

(82 durations < 0.005s hidden.  Use -vv to show these durations.)
================================================================= 44 passed in 29.44s ==================================================================

@noemotiovon
noemotiovon marked this pull request as ready for review December 24, 2025 08:01
@noemotiovon

noemotiovon commented Dec 24, 2025

Copy link
Copy Markdown
Contributor Author

Hi @Tcc0403 @zheliuyu @TianHao324, Merry Christmas Eve! 🎄
When you have a moment, could you help take a look at this code? Thanks!

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>
@noemotiovon noemotiovon changed the title feat(ascend): add ub manager Add UB Manager for automatic tiling strategy management Dec 24, 2025
@noemotiovon noemotiovon changed the title Add UB Manager for automatic tiling strategy management feat(NPU): add UB Manager for auto tiling strategy management Dec 24, 2025

@Tcc0403 Tcc0403 left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Comment thread src/liger_kernel/ops/backends/_ascend/ops/geglu.py Outdated
Comment thread src/liger_kernel/ops/backends/_ascend/ops/rope.py Outdated
Comment thread src/liger_kernel/ops/backends/_ascend/ub_manager.py Outdated
Comment thread src/liger_kernel/ops/backends/_ascend/ub_manager.py Outdated
Comment thread src/liger_kernel/ops/backends/_ascend/ub_manager.py Outdated
- 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
@noemotiovon

Copy link
Copy Markdown
Contributor Author

Hi @Tcc0403,
Thank you so much for taking the time to review this, especially during the Christmas holidays — really appreciate it. Merry Christmas! 🎄

Your suggestions were very helpful. I’ve updated the code accordingly, added caching, and improved the overall functionality.

@Tcc0403

Tcc0403 commented Dec 25, 2025

Copy link
Copy Markdown
Collaborator

Thank you so much for taking the time to review this, especially during the Christmas holidays — really appreciate it. Merry Christmas! 🎄

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 src/liger_kernel/ops/backends/_ascend directly without having to change anything since it's well-written already. After adding the document, I think we are good to merge this PR!

@noemotiovon

Copy link
Copy Markdown
Contributor Author

@Tcc0403, Thank you so much! I really appreciate it 😊
I actually wanted to add the document earlier as well, but wasn’t sure where the most appropriate place would be, haha. Thanks a lot for the guidance.
I’ll put it under src/liger_kernel/ops/backends/_ascend as suggested, and I’ll also expand the document based on the existing content to include the cache-related details, so it can be as comprehensive as possible for future contributors.
Thanks again for the helpful feedback!

@Tcc0403

Tcc0403 commented Dec 25, 2025

Copy link
Copy Markdown
Collaborator

Fantastic! Thank you so much!

Comment thread src/liger_kernel/ops/backends/_ascend/ub_manager.py Outdated

@Tcc0403 Tcc0403 left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.
@noemotiovon

Copy link
Copy Markdown
Contributor Author

Hi @Tcc0403 , requesting review for this UB Manager refactoring.

Changes:

  • Removed strategy registry, unified to use _default_strategy (sufficient for current scenarios, extensible for future needs)
  • Clarified parameter structure: safety_margin, dtype_size, memory_multiplier, tiling_dims, unit_params
  • Removed caching (current computation is simple; can reconsider if complex strategies are added later)

All tests pass for GEGLU and ROPE. See ascend-ub-manager-design.md for details.
Please review, thanks!

@noemotiovon
noemotiovon requested a review from Tcc0403 December 26, 2025 04:33

@Tcc0403 Tcc0403 left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Sorry for late response, it took a while to compile my thoughts.

Comment thread src/liger_kernel/ops/backends/_ascend/ascend-ub-manager-design.md Outdated
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
@noemotiovon

noemotiovon commented Dec 27, 2025

Copy link
Copy Markdown
Contributor Author

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.

Comment thread src/liger_kernel/ops/backends/_ascend/ascend-ub-manager-design.md Outdated
Change shapes parameter and return values from List[List[int]] to
Tuple[Tuple[int, ...], ...] for better immutability and PyTorch integration.
@noemotiovon

Copy link
Copy Markdown
Contributor Author

Hi @Tcc0403, We've updated the code to use Tuple[Tuple[int, ...], ...] throughout, including:

  • Function signatures and return types
  • All operator calls (ROPE and GEGLU)
  • Documentation and examples

The new interface is indeed cleaner and more consistent with PyTorch conventions. Thanks for catching this!

Comment thread src/liger_kernel/ops/backends/_ascend/ascend-ub-manager-design.md
@Tcc0403
Tcc0403 requested a review from hipudding December 27, 2025 09:49
@Tcc0403

Tcc0403 commented Dec 27, 2025

Copy link
Copy Markdown
Collaborator

That was fast! Overall LGTM, requesting reviews in case I overlook something.

Tcc0403
Tcc0403 previously approved these changes Dec 27, 2025
Comment thread src/liger_kernel/ops/backends/_ascend/ub_manager.py Outdated
Comment thread src/liger_kernel/ops/backends/_ascend/ops/geglu.py
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.

3 participants