[Fix] fix benchamrk issues and add ub management for npu

[sunyi0505]

Summary

This PR introduces a unified Unified Buffer (UB) management module for the Triton-Ascend backend on Ascend NPUs. It refactors the block size and grid launch strategies of multiple kernels based on this module to resolve launch failures caused by UB overflow and the Ascend hardware grid dimension limit of 65535. In addition, it fixes runtime errors in benchmark_conv.py under the NPU environment, which were triggered by the unavailability of CUDA reference implementations.

Key Changes

New File: fla/utils/ascend_ub_manager.py (+568 lines)

• Automatically detects the NPU model and its available UB capacity; falls back to a 64 KiB UB limit with a warning if detection fails.
• Provides unified tiling strategy utilities: UBManager, compute_ub_block_size, compute_vocab_block_size, compute_activation_block_size, compute_row_tile_block_size, etc.
• Exposes grid chunking helpers: defines ASCEND_MAX_GRID_DIM = 65535, max_grid_axis_chunks, and iter_axis_launch_chunks to enable staged kernel launches on the host side.

Triton-Ascend Kernels Integrated with UB Management (8 files modified)

Module	Main Modifications
activations.py	Replaces hardcoded block sizes with compute_activation_block_size; enables backward-pass differentiation via the is_backward flag; configures a dedicated memory multiplier for PowGLU.
layernorm.py	Removes autotune; calculates block dimension BD subject to UB constraints; performs row-axis chunked launches when sequence length T exceeds 65535.
rotary.py	Adopts UB-aware BT tiling instead of empirical hardcoded rules such as R>=128 → 16; chunks along the NT dimension when grid limits are hit, and adds the NT_OFFSET argument to kernels.
fused_cross_entropy.py	Computes vocabulary block size via the UB manager; introduces ROW_OFFSET and chunked kernel launches for both forward and backward passes.
fused_linear_cross_entropy.py	Same UB-based block sizing as above; adds chunked execution for logsumexp in the forward pass; uses compute_elementwise_block_size for element-wise multiplication operations.
fused_kl_div.py	Routes both vocabulary and element-wise block size calculations through the unified UB manager.
grpo.py	Standardizes vocabulary block size configuration; adds ROW_OFFSET and chunked launches for forward and backward kernels.
cumsum.py	Fully integrates UB-aware (BT, BS) block size calculation for scalar/vector local and global cumulative sum operations; supports NT_OFFSET/BH_OFFSET and multi-dimensional grid chunking.

Benchmark Fix (benchmarks/modules/benchmark_conv.py)

• Skips unavailable causal_conv1d_cuda_* execution providers on NPU devices.
• Dynamically generates the provider list at runtime to prevent decorators from binding incorrect line_vals and styles during module import.

Change Scope

10 files modified, with +1220 new lines and -240 deleted lines.

Motivation

Ascend NPUs have limited UB capacity that varies across different chip models. Previously, all kernels relied on fixed block sizes calculated as 65536 // element_size or empirical heuristic values. Such hardcoded configurations frequently caused UB overflow or kernel launch failures when the grid dimension exceeded 65535, especially under large vocabulary sizes, high feature dimensions, or long sequence lengths. This PR consolidates these scattered magic numbers into a centralized UB manager and applies more conservative tiling policies by respecting the peak memory multiplier of forward and backward computation paths.

[zheliuyu]

👍

[sunyi0505]

@zhiyuan1i This pr is ready for review. Thanks!

[zhiyuan1i]

LGTM on the NPU changes. Verified the launch chunking math (iter_axis_launch_chunks and the 2D/3D cumsum chunking keep every grid <= 65535 with full coverage) and that the offset/slice patterns are applied correctly across activations / layernorm / rotary / cross_entropy / grpo / cumsum. NPU A2 CI passes.

The red H100 checks (test-models forgetting_transformer TVM tirx.Bind, test-ops test_nsa) are in CUDA code this PR does not touch — test-models is already failing on main (ac6c6483) with the same error, so these are pre-existing / unrelated, not regressions from this PR.

Non-blocking follow-ups for a separate PR: a few unit tests for the new ascend_ub_manager module would be valuable (it's pure Python, easy to test); compute_grid_limited_tile_size has a dead other_grid_product arg worth cleaning up; and the unchunked axes (cross_entropy n_splits, cumsum vector ns) could use a >65535 guard like #997.