【benchmark】新增TritonNPUBench，提供vllm/sglang/fbgemm仓triton算子描述

benchmarks/TritonNPUBench/fbgemm/_bwd_preprocess_do_o_dot/_bwd_preprocess_do_o_dot.py

@@ -0,0 +1,49 @@
+# Extracted Triton Kernels
+# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/fbgemm_operator_cases/newtest_cases/test__bwd_preprocess_do_o_dot_npu.py
+# Main kernel: _bwd_preprocess_do_o_dot
+# PT file: _bwd_preprocess_do_o_dot_v2.pt
+
+import triton
+import triton.language as tl
+
+
+# === _bwd_preprocess_do_o_dot ===
+@triton.jit
+def _bwd_preprocess_do_o_dot(
+    o_ptr,
+    do_ptr,
+    delta_ptr,
+    T,
+    stride_ob,
+    stride_ot,
+    stride_od,
+    stride_do_b,
+    stride_do_t,
+    stride_do_d,
+    BLOCK_T: tl.constexpr,
+    BLOCK_D: tl.constexpr,
+):
+    start_t = tl.program_id(0)
+    offs_t = start_t * BLOCK_T + tl.arange(0, BLOCK_T)
+    pid_b = tl.program_id(1)
+    offs_d = tl.arange(0, BLOCK_D)
+
+    o_ptrs = (
+        o_ptr
+        + pid_b * stride_ob
+        + offs_t[:, None] * stride_ot
+        + offs_d[None, :] * stride_od
+    )
+    do_ptrs = (
+        do_ptr
+        + pid_b * stride_do_b
+        + offs_t[:, None] * stride_do_t
+        + offs_d[None, :] * stride_do_d
+    )
+    o = tl.load(o_ptrs, mask=(offs_t[:, None] < T), other=0.0)
+    do = tl.load(do_ptrs, mask=(offs_t[:, None] < T), other=0.0)
+    delta = tl.sum(o * do, axis=1)
+
+    delta_ptrs = delta_ptr + pid_b * T + offs_t
+    tl.store(delta_ptrs, delta, mask=(offs_t < T))
+

benchmarks/TritonNPUBench/fbgemm/_fbgemm_gather_scale_dense_tokens/_fbgemm_gather_scale_dense_tokens.py

@@ -0,0 +1,67 @@
+# Extracted Triton Kernels
+# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/fbgemm_operator_cases/newtest_cases/test__fbgemm_gather_scale_dense_tokens_npu.py
+# Main kernel: _fbgemm_gather_scale_dense_tokens
+# PT file: test__fbgemm_gather_scale_dense_tokens_v2.pt
+
+import triton
+import triton.language as tl
+
+
+# === _fbgemm_gather_scale_dense_tokens ===
+@triton.jit
+def _fbgemm_gather_scale_dense_tokens(
+    out,
+    x,
+    token_indices,
+    expert_indices,
+    scores,
+    stride_t,
+    stride_e,
+    valid_token_count,
+    D: tl.constexpr,
+    BLOCK_D_OUTER: tl.constexpr,
+    BLOCK_D_INNER: tl.constexpr,
+):
+    output_token_index = tl.program_id(0)
+    feature_offset = tl.program_id(1) * BLOCK_D_OUTER
+
+    if valid_token_count is not None:
+        valid_token_count = tl.load(
+            valid_token_count, None, eviction_policy="evict_last"
+        )
+        if output_token_index >= valid_token_count:
+            return
+
+    input_token_index = tl.load(
+        token_indices + output_token_index, None, eviction_policy="evict_last"
+    )
+    input_expert_index = tl.load(
+        expert_indices + output_token_index, None, eviction_policy="evict_last"
+    )
+
+    input_score = tl.load(
+        scores + input_token_index * stride_t + input_expert_index * stride_e,
+        None,
+        eviction_policy="evict_last",
+    ).to(tl.float32)
+
+    for _ in range(0, BLOCK_D_OUTER // BLOCK_D_INNER):
+        input_token_value = tl.load(
+            x
+            + input_token_index.to(tl.int64) * D
+            + feature_offset
+            + tl.arange(0, BLOCK_D_INNER)[:],
+            None,
+        ).to(tl.float32)
+        output_token_value = input_token_value * input_score
+
+        tl.store(
+            out
+            + output_token_index.to(tl.int64) * D
+            + feature_offset
+            + tl.arange(0, BLOCK_D_INNER)[:],
+            output_token_value,
+            None,
+        )
+        feature_offset += BLOCK_D_INNER
+

benchmarks/TritonNPUBench/fbgemm/_fbgemm_gather_scale_fp8_rowwise_quant_dense_tokens/_fbgemm_gather_scale_fp8_rowwise_quant_dense_tokens.py

@@ -0,0 +1,110 @@
+# Extracted Triton Kernels
+# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/fbgemm_operator_cases/newtest_cases/test__fbgemm_gather_scale_fp8_rowwise_quant_dense_tokens_npu.py
+# Main kernel: _fbgemm_gather_scale_fp8_rowwise_quant_dense_tokens
+# PT file: test__fbgemm_gather_scale_fp8_rowwise_quant_dense_tokens_v2.pt
+
+import triton
+import triton.language as tl
+
+
+# === _fbgemm_gather_scale_fp8_rowwise_quant_dense_tokens ===
+@triton.jit
+def _fbgemm_gather_scale_fp8_rowwise_quant_dense_tokens(
+    output_ptr,
+    output_scale_ptr,
+    input_ptr,
+    token_indices_ptr,
+    expert_indices_ptr,
+    scores_ptr,
+    scale_ub_ptr,
+    stride_t,
+    stride_e,
+    valid_token_count,
+    D: tl.constexpr,
+    TL_FP8_DTYPE: tl.constexpr,
+    MAX_FP8: tl.constexpr,
+    EPS: tl.constexpr,
+    CLAMP_MAX: tl.constexpr,
+    BLOCK_D: tl.constexpr,
+):
+    tl.static_assert(D % BLOCK_D == 0, "D must be a multiple of BLOCK_D")
+
+    output_token_index = tl.program_id(0)
+
+    if valid_token_count is not None:
+        valid_token_count = tl.load(
+            valid_token_count, None, eviction_policy="evict_last"
+        )
+        if output_token_index >= valid_token_count:
+            return
+
+    input_token_index = tl.load(
+        token_indices_ptr + output_token_index, None, eviction_policy="evict_first"
+    )
+    input_expert_index = tl.load(
+        expert_indices_ptr + output_token_index, None, eviction_policy="evict_first"
+    )
+    input_score = tl.load(
+        scores_ptr + input_token_index * stride_t + input_expert_index * stride_e,
+        None,
+        eviction_policy="evict_first",
+    ).to(tl.float32)
+
+    row_max = 0.0
+    in_2d_ptr = (
+        input_ptr + input_token_index.to(tl.int64) * D + tl.arange(0, BLOCK_D)[:]
+    )
+    for _ in range(0, D, BLOCK_D):
+        input_token_value = tl.load(
+            in_2d_ptr,
+            None,
+            eviction_policy="evict_last",
+        ).to(tl.float32)
+        output_token_value = input_token_value * input_score
+
+        tile_max = tl.max(tl.abs(output_token_value))
+        row_max = tl.maximum(tile_max, row_max)
+        in_2d_ptr += BLOCK_D
+
+    # Clamp max value appropriately.
+    if CLAMP_MAX:
+        ub = tl.load(scale_ub_ptr, eviction_policy="evict_last")
+        row_max = tl.clamp(row_max, EPS, ub)
+    else:
+        row_max = tl.maximum(row_max, EPS)
+
+    # Scale and quantize.
+    output_scale = MAX_FP8 / row_max
+    tl.store(output_scale_ptr + output_token_index, 1.0 / output_scale)
+
+    in_2d_ptr = (
+        input_ptr + input_token_index.to(tl.int64) * D + tl.arange(0, BLOCK_D)[:]
+    )
+    out_2d_ptr = (
+        output_ptr + output_token_index.to(tl.int64) * D + tl.arange(0, BLOCK_D)[:]
+    )
+    for _ in range(0, D, BLOCK_D):
+        # Load from L2
+        input_token_value = tl.load(
+            in_2d_ptr,
+            None,
+            eviction_policy="evict_first",
+        ).to(tl.float32)
+        # Rematerilize
+        output_token_value_fp8 = (input_token_value * input_score) * output_scale
+
+        # Clamp A to fp8 range to make sure there's no overflow.
+        # This is required for AMD. Nvidia's default saturation
+        # handles it, but it's nice to have anyway.
+        output_token_value_fp8 = tl.clamp(output_token_value_fp8, -MAX_FP8, MAX_FP8).to(
+            TL_FP8_DTYPE
+        )
+        tl.store(
+            out_2d_ptr,
+            output_token_value_fp8,
+            None,
+            cache_modifier=".cg",
+        )
+        in_2d_ptr += BLOCK_D
+        out_2d_ptr += BLOCK_D
+

benchmarks/TritonNPUBench/fbgemm/_fbgemm_scatter_add_dense_tokens/_fbgemm_scatter_add_dense_tokens.py

@@ -0,0 +1,49 @@
+# Extracted Triton Kernels
+# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/fbgemm_operator_cases/newtest_cases/test__fbgemm_scatter_add_dense_tokens_npu.py
+# Main kernel: _fbgemm_scatter_add_dense_tokens
+# PT file: test__fbgemm_scatter_add_dense_tokens_v2.pt
+
+import triton
+import triton.language as tl
+
+
+# === _fbgemm_scatter_add_dense_tokens ===
+@triton.jit
+def _fbgemm_scatter_add_dense_tokens(
+    out_tokens,
+    in_tokens,
+    token_indices,
+    valid_token_count,
+    D: tl.constexpr,
+    BLOCK_D_OUTER: tl.constexpr,
+    BLOCK_D_INNER: tl.constexpr,
+):
+    input_token_index = tl.program_id(0).to(tl.int64)
+    feature_offset = tl.program_id(1) * BLOCK_D_OUTER + tl.arange(0, BLOCK_D_INNER)[:]
+
+    if valid_token_count is not None:
+        valid_token_count = tl.load(
+            valid_token_count, None, eviction_policy="evict_last"
+        )
+        if input_token_index >= valid_token_count:
+            return
+
+    output_token_index = tl.load(
+        token_indices + input_token_index, None, eviction_policy="evict_last"
+    ).to(tl.int64)
+
+    for _ in range(0, BLOCK_D_OUTER // BLOCK_D_INNER):
+        input_token_value = tl.load(
+            in_tokens + input_token_index * D + feature_offset,
+            None,
+            eviction_policy="evict_first",
+        )
+
+        tl.atomic_add(
+            out_tokens + output_token_index * D + feature_offset,
+            input_token_value,
+            None,
+            sem="relaxed",
+        )
+        feature_offset += BLOCK_D_INNER
+

benchmarks/TritonNPUBench/fbgemm/_fbgemm_scatter_add_padded_tokens/_fbgemm_scatter_add_padded_tokens.py

@@ -0,0 +1,74 @@
+# Extracted Triton Kernels
+# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/fbgemm_operator_cases/newtest_cases/test__fbgemm_scatter_add_padded_tokens_npu.py
+# Main kernel: _fbgemm_scatter_add_padded_tokens
+# PT file: test__fbgemm_scatter_add_padded_tokens_v2.pt
+
+import triton
+import triton.language as tl
+
+
+# === _fbgemm_scatter_add_padded_tokens ===
+@triton.jit
+def _fbgemm_scatter_add_padded_tokens(
+    in_tokens_ptr,
+    token_counts_ptr,
+    token_indices_ptr,
+    out_tokens_ptr,
+    EP: tl.constexpr,
+    E: tl.constexpr,
+    T_BUCKET,
+    T_K,
+    D: tl.constexpr,
+    BLOCK_E: tl.constexpr,
+    SPLIT_T: tl.constexpr,
+    BLOCK_D: tl.constexpr,
+):
+    """
+    in_tokens: [EP, T_K, D]
+    token_counts: [E]
+    out_tokens: [T, D]
+    """
+    expert = tl.program_id(0)
+    t_tile = tl.program_id(1)
+
+    tl.static_assert(D % BLOCK_D == 0)
+    NUM_D_BLOCKS: tl.constexpr = D // BLOCK_D
+
+    num_tokens = tl.load(token_counts_ptr + expert)
+    if num_tokens == 0:
+        return
+
+    num_tokens_per_cta = tl.cdiv(num_tokens, SPLIT_T)
+    start_token = t_tile * num_tokens_per_cta
+    end_token = min(start_token + num_tokens_per_cta, num_tokens)
+
+    tl.static_assert(E % EP == 0)
+    EXPERT_PER_RANK: tl.constexpr = E // EP
+    rank = expert // EXPERT_PER_RANK
+
+    offs_e = tl.arange(0, BLOCK_E)
+    token_counts = tl.load(token_counts_ptr + offs_e, mask=(offs_e < E), other=0)
+    input_local_offset = (
+        tl.sum(tl.where(offs_e < expert, token_counts, 0)) + start_token
+    ).to(tl.int64)
+
+    for _t in range(start_token, end_token):
+        output_local_offset = tl.load(token_indices_ptr + input_local_offset).to(
+            tl.int64
+        )
+        output_global_offset = output_local_offset * D
+
+        d_ptr = tl.arange(0, BLOCK_D)
+        input_global_ptr = (
+            in_tokens_ptr + rank * T_K * D + input_local_offset * D + d_ptr
+        )
+        output_global_ptr = out_tokens_ptr + output_global_offset + d_ptr
+
+        for _d in range(NUM_D_BLOCKS):
+            vec = tl.load(input_global_ptr)
+            tl.atomic_add(output_global_ptr, vec, sem="relaxed")
+            input_global_ptr += BLOCK_D
+            output_global_ptr += BLOCK_D
+
+        input_local_offset += 1
+