[Fix] Add Ascend UB management and fix NPU benchmark issues

Author: sunyi0505

benchmarks/modules/benchmark_conv.py

@@ -11,34 +11,28 @@
 
 from fla.modules.convolution import causal_conv1d
 from fla.ops.utils.index import prepare_sequence_ids
+from fla.utils import IS_NPU
 
 try:
     from causal_conv1d import causal_conv1d_fn
 except ImportError:
     causal_conv1d_fn = None
 
 
-@triton.testing.perf_report(
-    triton.testing.Benchmark(
-        # argument names to use as an x-axis for the plot
-        x_names=['T', 'D'],
-        # different possible values for `x_name`
-        x_vals=[(128 * 2 ** i, d) for d in [256, 512, 1024, 2048, 4096] for i in range(1, 10)],
-        # argument name whose value corresponds to a different line in the plot
-        line_arg='provider',
-        # possible values for `line_arg``
-        line_vals=['causal_conv1d_fwd', 'causal_conv1d_cuda_fwd', 'causal_conv1d_fwdbwd', 'causal_conv1d_cuda_fwdbwd'],
-        # label name for the lines
-        line_names=['causal_conv1d_fwd', 'causal_conv1d_cuda_fwd', 'causal_conv1d_fwdbwd', 'causal_conv1d_cuda_fwdbwd'],
-        # line styles
-        styles=[('green', '-'), ('blue', '--'), ('red', '-.'),
-                ('cyan', ':'), ('yellow', 'dotted'), ('cyan', '--'), ('cyan', '-'), ('black', ':')],
-        ylabel="Execution Time (ms)",  # label name for the y-axis
-        # name for the plot. Used also as a file name for saving the plot.
-        plot_name="Performance",
-        args={},
-    ),
-)
+def _conv_benchmark_providers():
+    providers = ['causal_conv1d_fwd', 'causal_conv1d_fwdbwd']
+    if not IS_NPU and causal_conv1d_fn is not None:
+        providers.extend(['causal_conv1d_cuda_fwd', 'causal_conv1d_cuda_fwdbwd'])
+    return providers
+
+
+_LINE_VALS = _conv_benchmark_providers()
+_STYLES = [
+    ('green', '-'), ('blue', '--'), ('red', '-.'), ('cyan', ':'),
+    ('yellow', 'dotted'), ('cyan', '--'), ('cyan', '-'), ('black', ':'),
+]
+
+
 def benchmark(T, D, provider):
     from fla.utils import device
     dtype = torch.bfloat16
@@ -100,5 +94,27 @@ def benchmark(T, D, provider):
     return results
 
 
+benchmark = triton.testing.perf_report(
+    triton.testing.Benchmark(
+        # argument names to use as an x-axis for the plot
+        x_names=['T', 'D'],
+        # different possible values for `x_name`
+        x_vals=[(128 * 2 ** i, d) for d in [256, 512, 1024, 2048, 4096] for i in range(1, 10)],
+        # argument name whose value corresponds to a different line in the plot
+        line_arg='provider',
+        # possible values for `line_arg``
+        line_vals=_LINE_VALS,
+        # label name for the lines
+        line_names=_LINE_VALS,
+        # line styles
+        styles=_STYLES[:len(_LINE_VALS)],
+        ylabel="Execution Time (ms)",  # label name for the y-axis
+        # name for the plot. Used also as a file name for saving the plot.
+        plot_name="Performance",
+        args={},
+    ),
+)(benchmark)
+
+
 if __name__ == '__main__':
     benchmark.run(print_data=True)

fla/modules/backends/triton_ascend/activations.py

@@ -14,20 +14,27 @@
 
 from fla.ops.utils.op import exp, log
 from fla.utils import autocast_custom_bwd, autocast_custom_fwd, input_guard
+from fla.utils.ascend_ub_manager import ASCEND_MAX_GRID_DIM, compute_activation_block_size
 
 # Ascend launch limits: grid dim and per-core vector width.
-_PREFERRED_BLOCK_SIZES = (2048, 1024, 512)
-_MAX_GRID_DIM = 65535
 _MAX_CORE_DIM = 65535
 
 
-def _activation_launch_config(T: int) -> tuple[tuple[int], int]:
-    """Pick block size under Ascend launch limits."""
-    for B in _PREFERRED_BLOCK_SIZES:
-        if triton.cdiv(T, B) <= _MAX_GRID_DIM and B * 8 <= _MAX_CORE_DIM:
-            return (triton.cdiv(T, B),), B
-    B = min(triton.cdiv(T, _MAX_GRID_DIM), _MAX_CORE_DIM // 8)
-    return (triton.cdiv(T, B),), max(B, 1)
+def _activation_launch_config(
+    T: int,
+    is_backward: bool = False,
+    *,
+    memory_multiplier: float | None = None,
+) -> tuple[tuple[int], int]:
+    """Pick block size under Ascend launch and UB limits."""
+    B = compute_activation_block_size(
+        T,
+        is_backward,
+        max_grid=ASCEND_MAX_GRID_DIM,
+        max_core_dim=_MAX_CORE_DIM,
+        memory_multiplier=memory_multiplier,
+    )
+    return (triton.cdiv(T, B),), B
 
 
 @triton.jit
@@ -470,7 +477,7 @@ def gelu_bwd_npu(g: torch.Tensor, x: torch.Tensor) -> torch.Tensor:
     g = _ensure_inner_contiguous(g)
     T, D = x.numel(), x.shape[-1]
     dx = _alloc_output(x)
-    grid, B = _activation_launch_config(T)
+    grid, B = _activation_launch_config(T, is_backward=True)
     gelu_bwd_kernel[grid](
         x, g, dx, T=T, D=D,
         stride_x_row=_get_stride(x),
@@ -502,7 +509,7 @@ def sqrelu_bwd_npu(g: torch.Tensor, x: torch.Tensor) -> torch.Tensor:
     g = _ensure_inner_contiguous(g)
     T, D = x.numel(), x.shape[-1]
     dx = _alloc_output(x)
-    grid, B = _activation_launch_config(T)
+    grid, B = _activation_launch_config(T, is_backward=True)
     sqrelu_bwd_kernel[grid](
         x, g, dx, T=T, D=D,
         stride_x_row=_get_stride(x),
@@ -534,7 +541,7 @@ def sigmoid_bwd_npu(x: torch.Tensor, dy: torch.Tensor, output_contiguous: bool =
     dy = _ensure_inner_contiguous(dy)
     T, D = x.numel(), x.shape[-1]
     dx = _alloc_output(x, output_contiguous)
-    grid, B = _activation_launch_config(T)
+    grid, B = _activation_launch_config(T, is_backward=True)
     sigmoid_bwd_kernel[grid](
         x, dy, dx, T=T, D=D,
         stride_x_row=_get_stride(x),
@@ -575,7 +582,7 @@ def logsigmoid_bwd_npu(
     dy = _ensure_inner_contiguous(dy)
     T, D = x.numel(), x.shape[-1]
     dx = _alloc_output(x, output_contiguous)
-    grid, B = _activation_launch_config(T)
+    grid, B = _activation_launch_config(T, is_backward=True)
     logsigmoid_bwd_kernel[grid](
         x=x,
         dx=dx,
@@ -612,7 +619,7 @@ def swish_bwd_npu(x: torch.Tensor, dy: torch.Tensor, output_contiguous: bool = F
     dy = _ensure_inner_contiguous(dy)
     T, D = x.numel(), x.shape[-1]
     dx = _alloc_output(x, output_contiguous)
-    grid, B = _activation_launch_config(T)
+    grid, B = _activation_launch_config(T, is_backward=True)
     swish_bwd_kernel[grid](
         x, dy, dx, T=T, D=D,
         stride_x_row=_get_stride(x),
@@ -660,7 +667,7 @@ def swiglu_fwdbwd_npu(
         z = _alloc_output(x, output_contiguous)
     else:
         z = None
-    grid, B = _activation_launch_config(T)
+    grid, B = _activation_launch_config(T, is_backward=True)
     swiglu_fwdbwd_kernel[grid](
         x, y, g, dx, dy, z, T=T, D=D,
         stride_x_row=_get_stride(x),
@@ -807,14 +814,19 @@ def powglu_fwdbwd_kernel(
         tl.store(z + z_off, b_z.to(z.dtype.element_ty), mask=mask)
 
 
+# Peak fp32 temporaries: sigmoid, sqrt, log, exp, pow, gate, output.
+_POWGLU_FWD_MEM_MULT = 8.0
+_POWGLU_BWD_MEM_MULT = 10.0
+
+
 @torch.compiler.disable
 def powglu_fwd_npu(x: torch.Tensor, y: torch.Tensor, power: float = 3.0, output_contiguous: bool = False) -> torch.Tensor:
     assert x.shape == y.shape, f"powglu_fwd: shape mismatch x={x.shape} y={y.shape}"
     x = _ensure_inner_contiguous(x)
     y = _ensure_inner_contiguous(y)
     T, D = x.numel(), x.shape[-1]
     z = _alloc_output(x, output_contiguous)
-    grid, B = _activation_launch_config(T)
+    grid, B = _activation_launch_config(T, memory_multiplier=_POWGLU_FWD_MEM_MULT)
     powglu_fwd_kernel[grid](
         x=x,
         y=y,
@@ -850,7 +862,7 @@ def powglu_fwdbwd_npu(
         z = _alloc_output(x, output_contiguous)
     else:
         z = None
-    grid, B = _activation_launch_config(T)
+    grid, B = _activation_launch_config(T, is_backward=True, memory_multiplier=_POWGLU_BWD_MEM_MULT)
     powglu_fwdbwd_kernel[grid](
         x=x,
         y=y,