From 66b21643c57c6b88e7de6a5259f054076885a13f Mon Sep 17 00:00:00 2001 From: zhumingming Date: Fri, 17 Apr 2026 21:33:26 +0800 Subject: [PATCH 1/3] =?UTF-8?q?=E3=80=90benchmark=E3=80=91=E6=96=B0?= =?UTF-8?q?=E5=A2=9Evllm=E4=BB=93116=E4=B8=AA=E7=AE=97=E5=AD=90=E6=8F=8F?= =?UTF-8?q?=E8=BF=B0?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../_apply_grammar_bitmask_kernel.py | 42 ++ .../_apply_write_kernel.py | 36 ++ .../vllm/_bias_kernel/_bias_kernel.py | 100 ++++ .../vllm/_bincount_kernel/_bincount_kernel.py | 37 ++ .../_bmm_chunk_fwd_kernel.py | 122 +++++ .../_causal_conv1d_fwd_kernel.py | 462 ++++++++++++++++++ .../_causal_conv1d_update_kernel.py | 330 +++++++++++++ .../_chunk_cumsum_fwd_kernel.py | 104 ++++ .../_chunk_scan_fwd_kernel.py | 281 +++++++++++ .../_chunk_state_fwd_kernel.py | 118 +++++ .../_chunk_state_varlen_kernel.py | 192 ++++++++ ...combine_sampled_and_draft_tokens_kernel.py | 66 +++ .../_compute_pid_wrapper.py | 38 ++ .../_compute_slot_mappings_kernel.py | 62 +++ ...onvert_req_index_to_global_index_kernel.py | 76 +++ .../_copy_page_indices_kernel.py | 33 ++ .../_correct_attn_cp_out_kernel.py | 97 ++++ .../_count_expert_num_tokens.py | 41 ++ .../_expand_idx_mapping_kernel.py | 27 + .../vllm/_fwd_diag_kernel/_fwd_diag_kernel.py | 138 ++++++ .../_fwd_grouped_kernel_stage1.py | 176 +++++++ .../vllm/_fwd_kernel/_fwd_kernel.py | 329 +++++++++++++ .../_fwd_kernel_alibi/_fwd_kernel_alibi.py | 289 +++++++++++ .../_fwd_kernel_ep_gather.py | 75 +++ .../_fwd_kernel_ep_scatter_1.py | 56 +++ .../_fwd_kernel_ep_scatter_2.py | 84 ++++ .../_fwd_kernel_stage1/_fwd_kernel_stage1.py | 144 ++++++ .../_fwd_kernel_stage2/_fwd_kernel_stage2.py | 72 +++ .../vllm/_fwd_kv_parallel/_fwd_kv_parallel.py | 118 +++++ .../vllm/_fwd_kv_reduce/_fwd_kv_reduce.py | 74 +++ .../_fwd_none_diag_kernel.py | 92 ++++ .../_gather_block_tables_kernel.py | 47 ++ .../_get_num_sampled_and_rejected_kernel.py | 38 ++ .../_gumbel_sample_kernel.py | 66 +++ .../_layer_norm_fwd_1pass_kernel.py | 75 +++ .../_linear_attn_decode_kernel.py | 91 ++++ .../_log_softmax_kernel.py | 72 +++ .../_lora_expand_kernel.py | 345 +++++++++++++ .../_lora_shrink_kernel.py | 320 ++++++++++++ .../vllm/_min_p_kernel/_min_p_kernel.py | 45 ++ .../vllm/_num_nans_kernel/_num_nans_kernel.py | 32 ++ .../vllm/_pack_seq_kernel/_pack_seq_kernel.py | 57 +++ .../_penalties_and_temperature_kernel.py | 85 ++++ .../_per_token_group_quant_fp8.py | 63 +++ .../_per_token_group_quant_fp8_colmajor.py | 73 +++ .../_per_token_group_quant_int8.py | 51 ++ .../_per_token_quant_int8.py | 41 ++ .../_post_update_kernel.py | 51 ++ .../_prepare_eagle_docode_kernel.py | 78 +++ .../_prepare_eagle_inputs_kernel.py | 59 +++ .../_prepare_mrope_positions_kernel.py | 60 +++ .../_prepare_pos_seq_lens_kernel.py | 46 ++ .../_prepare_prefill_inputs_kernel.py | 46 ++ .../vllm/_ranks_kernel/_ranks_kernel.py | 32 ++ .../_rejection_sample_kernel.py | 42 ++ .../vllm/_rms_norm_kernel/_rms_norm_kernel.py | 60 +++ .../_selective_scan_update_kernel.py | 243 +++++++++ .../_silu_mul_fp8_quant_deep_gemm.py | 83 ++++ ..._mul_per_token_group_quant_fp8_colmajor.py | 80 +++ .../_state_passing_fwd_kernel.py | 92 ++++ .../_topk_log_softmax_kernel.py | 51 ++ .../_triton_mrope_forward.py | 128 +++++ .../_trtllm_prefill_attn_kvfp8_dequant.py | 56 +++ .../_unpack_seq_triton_kernel.py | 52 ++ .../_update_eagle_inputs_kernel.py | 55 +++ .../_w8a8_block_int8_matmul.py | 94 ++++ .../_w8a8_triton_block_scaled_mm.py | 94 ++++ .../awq_dequantize_kernel.py | 106 ++++ .../vllm/awq_gemm_kernel/awq_gemm_kernel.py | 128 +++++ .../batched_triton_kernel.py | 351 +++++++++++++ .../vllm/bmm_kernel/bmm_kernel.py | 142 ++++++ .../chunk_fwd_kernel_o/chunk_fwd_kernel_o.py | 109 +++++ ...ated_delta_rule_fwd_kernel_h_blockdim64.py | 248 ++++++++++ .../chunk_gla_fwd_kernel_o.py | 118 +++++ ...aled_dot_kkt_fwd_kernel_intra_sub_inter.py | 110 +++++ ...aled_dot_kkt_fwd_kernel_intra_sub_intra.py | 100 ++++ .../chunk_local_cumsum_scalar_kernel.py | 57 +++ .../chunk_local_cumsum_vector_kernel.py | 86 ++++ .../chunk_scaled_dot_kkt_fwd_kernel.py | 76 +++ .../compute_identity_kernel.py | 48 ++ .../cp_mha_gather_cache_kernel.py | 76 +++ .../eagle_prepare_inputs_padded_kernel.py | 52 ++ .../eagle_prepare_next_token_padded_kernel.py | 77 +++ .../vllm/expand_kernel/expand_kernel.py | 32 ++ .../fused_gdn_gating_kernel.py | 45 ++ .../vllm/fused_moe_kernel/fused_moe_kernel.py | 291 +++++++++++ .../fused_moe_kernel_gptq_awq.py | 263 ++++++++++ ...d_recurrent_gated_delta_rule_fwd_kernel.py | 158 ++++++ .../kda_gate_fwd_kernel.py | 71 +++ .../kernel_paged_attention_2d.py | 198 ++++++++ .../kernel_unified_attention_2d.py | 397 +++++++++++++++ .../kernel_unified_attention_3d.py | 395 +++++++++++++++ .../l2norm_fwd_kernel/l2norm_fwd_kernel.py | 29 ++ .../l2norm_fwd_kernel1/l2norm_fwd_kernel1.py | 32 ++ .../l2norm_fwd_kernel2/l2norm_fwd_kernel2.py | 22 + .../layer_norm_fwd_kernel.py | 107 ++++ .../layer_norm_gated_fwd_kernel.py | 89 ++++ .../layer_norm_gated_fwd_kernel1.py | 77 +++ .../matmul_kernel_persistent.py | 143 ++++++ .../vllm/mean_kernel/mean_kernel.py | 59 +++ .../merge_16x16_to_32x32_inverse_kernel.py | 125 +++++ .../merge_16x16_to_64x64_inverse_kernel.py | 278 +++++++++++ .../merge_attn_states_kernel.py | 84 ++++ .../vllm/pack_bitmatrix/pack_bitmatrix.py | 53 ++ .../recompute_w_u_fwd_kernel.py | 100 ++++ .../vllm/reduce_segments/reduce_segments.py | 130 +++++ .../rejection_greedy_sample_kernel.py | 53 ++ .../rejection_random_sample_kernel.py | 70 +++ .../reshape_and_cache_kernel_flash.py | 113 +++++ .../reshape_and_cache_kernel_flash_diffkv.py | 88 ++++ .../sample_recovered_tokens_kernel.py | 63 +++ .../vllm/scaled_mm_kernel/scaled_mm_kernel.py | 130 +++++ .../vllm/softplus_kernel/softplus_kernel.py | 26 + .../solve_tril_16x16_kernel.py | 75 +++ .../test_round_int8_kernel.py | 26 + .../triton_scale_swizzle.py | 73 +++ .../TritonNPUBench/vllm/vllm_gpu_perf.csv | 113 +++++ 117 files changed, 12806 insertions(+) create mode 100644 benchmarks/TritonNPUBench/vllm/_apply_grammar_bitmask_kernel/_apply_grammar_bitmask_kernel.py create mode 100644 benchmarks/TritonNPUBench/vllm/_apply_write_kernel/_apply_write_kernel.py create mode 100644 benchmarks/TritonNPUBench/vllm/_bias_kernel/_bias_kernel.py create mode 100644 benchmarks/TritonNPUBench/vllm/_bincount_kernel/_bincount_kernel.py create mode 100644 benchmarks/TritonNPUBench/vllm/_bmm_chunk_fwd_kernel/_bmm_chunk_fwd_kernel.py create mode 100644 benchmarks/TritonNPUBench/vllm/_causal_conv1d_fwd_kernel/_causal_conv1d_fwd_kernel.py create mode 100644 benchmarks/TritonNPUBench/vllm/_causal_conv1d_update_kernel/_causal_conv1d_update_kernel.py create mode 100644 benchmarks/TritonNPUBench/vllm/_chunk_cumsum_fwd_kernel/_chunk_cumsum_fwd_kernel.py create mode 100644 benchmarks/TritonNPUBench/vllm/_chunk_scan_fwd_kernel/_chunk_scan_fwd_kernel.py create mode 100644 benchmarks/TritonNPUBench/vllm/_chunk_state_fwd_kernel/_chunk_state_fwd_kernel.py create mode 100644 benchmarks/TritonNPUBench/vllm/_chunk_state_varlen_kernel/_chunk_state_varlen_kernel.py create mode 100644 benchmarks/TritonNPUBench/vllm/_combine_sampled_and_draft_tokens_kernel/_combine_sampled_and_draft_tokens_kernel.py create mode 100644 benchmarks/TritonNPUBench/vllm/_compute_pid_wrapper/_compute_pid_wrapper.py create mode 100644 benchmarks/TritonNPUBench/vllm/_compute_slot_mappings_kernel/_compute_slot_mappings_kernel.py create mode 100644 benchmarks/TritonNPUBench/vllm/_convert_req_index_to_global_index_kernel/_convert_req_index_to_global_index_kernel.py create mode 100644 benchmarks/TritonNPUBench/vllm/_copy_page_indices_kernel/_copy_page_indices_kernel.py create mode 100644 benchmarks/TritonNPUBench/vllm/_correct_attn_cp_out_kernel/_correct_attn_cp_out_kernel.py create mode 100644 benchmarks/TritonNPUBench/vllm/_count_expert_num_tokens/_count_expert_num_tokens.py create mode 100644 benchmarks/TritonNPUBench/vllm/_expand_idx_mapping_kernel/_expand_idx_mapping_kernel.py create mode 100644 benchmarks/TritonNPUBench/vllm/_fwd_diag_kernel/_fwd_diag_kernel.py create mode 100644 benchmarks/TritonNPUBench/vllm/_fwd_grouped_kernel_stage1/_fwd_grouped_kernel_stage1.py create mode 100644 benchmarks/TritonNPUBench/vllm/_fwd_kernel/_fwd_kernel.py create mode 100644 benchmarks/TritonNPUBench/vllm/_fwd_kernel_alibi/_fwd_kernel_alibi.py create mode 100644 benchmarks/TritonNPUBench/vllm/_fwd_kernel_ep_gather/_fwd_kernel_ep_gather.py create mode 100644 benchmarks/TritonNPUBench/vllm/_fwd_kernel_ep_scatter_1/_fwd_kernel_ep_scatter_1.py create mode 100644 benchmarks/TritonNPUBench/vllm/_fwd_kernel_ep_scatter_2/_fwd_kernel_ep_scatter_2.py create mode 100644 benchmarks/TritonNPUBench/vllm/_fwd_kernel_stage1/_fwd_kernel_stage1.py create mode 100644 benchmarks/TritonNPUBench/vllm/_fwd_kernel_stage2/_fwd_kernel_stage2.py create mode 100644 benchmarks/TritonNPUBench/vllm/_fwd_kv_parallel/_fwd_kv_parallel.py create mode 100644 benchmarks/TritonNPUBench/vllm/_fwd_kv_reduce/_fwd_kv_reduce.py create mode 100644 benchmarks/TritonNPUBench/vllm/_fwd_none_diag_kernel/_fwd_none_diag_kernel.py create mode 100644 benchmarks/TritonNPUBench/vllm/_gather_block_tables_kernel/_gather_block_tables_kernel.py create mode 100644 benchmarks/TritonNPUBench/vllm/_get_num_sampled_and_rejected_kernel/_get_num_sampled_and_rejected_kernel.py create mode 100644 benchmarks/TritonNPUBench/vllm/_gumbel_sample_kernel/_gumbel_sample_kernel.py create mode 100644 benchmarks/TritonNPUBench/vllm/_layer_norm_fwd_1pass_kernel/_layer_norm_fwd_1pass_kernel.py create mode 100644 benchmarks/TritonNPUBench/vllm/_linear_attn_decode_kernel/_linear_attn_decode_kernel.py create mode 100644 benchmarks/TritonNPUBench/vllm/_log_softmax_kernel/_log_softmax_kernel.py create mode 100644 benchmarks/TritonNPUBench/vllm/_lora_expand_kernel/_lora_expand_kernel.py create mode 100644 benchmarks/TritonNPUBench/vllm/_lora_shrink_kernel/_lora_shrink_kernel.py create mode 100644 benchmarks/TritonNPUBench/vllm/_min_p_kernel/_min_p_kernel.py create mode 100644 benchmarks/TritonNPUBench/vllm/_num_nans_kernel/_num_nans_kernel.py create mode 100644 benchmarks/TritonNPUBench/vllm/_pack_seq_kernel/_pack_seq_kernel.py create mode 100644 benchmarks/TritonNPUBench/vllm/_penalties_and_temperature_kernel/_penalties_and_temperature_kernel.py create mode 100644 benchmarks/TritonNPUBench/vllm/_per_token_group_quant_fp8/_per_token_group_quant_fp8.py create mode 100644 benchmarks/TritonNPUBench/vllm/_per_token_group_quant_fp8_colmajor/_per_token_group_quant_fp8_colmajor.py create mode 100644 benchmarks/TritonNPUBench/vllm/_per_token_group_quant_int8/_per_token_group_quant_int8.py create mode 100644 benchmarks/TritonNPUBench/vllm/_per_token_quant_int8/_per_token_quant_int8.py create mode 100644 benchmarks/TritonNPUBench/vllm/_post_update_kernel/_post_update_kernel.py create mode 100644 benchmarks/TritonNPUBench/vllm/_prepare_eagle_docode_kernel/_prepare_eagle_docode_kernel.py create mode 100644 benchmarks/TritonNPUBench/vllm/_prepare_eagle_inputs_kernel/_prepare_eagle_inputs_kernel.py create mode 100644 benchmarks/TritonNPUBench/vllm/_prepare_mrope_positions_kernel/_prepare_mrope_positions_kernel.py create mode 100644 benchmarks/TritonNPUBench/vllm/_prepare_pos_seq_lens_kernel/_prepare_pos_seq_lens_kernel.py create mode 100644 benchmarks/TritonNPUBench/vllm/_prepare_prefill_inputs_kernel/_prepare_prefill_inputs_kernel.py create mode 100644 benchmarks/TritonNPUBench/vllm/_ranks_kernel/_ranks_kernel.py create mode 100644 benchmarks/TritonNPUBench/vllm/_rejection_sample_kernel/_rejection_sample_kernel.py create mode 100644 benchmarks/TritonNPUBench/vllm/_rms_norm_kernel/_rms_norm_kernel.py create mode 100644 benchmarks/TritonNPUBench/vllm/_selective_scan_update_kernel/_selective_scan_update_kernel.py create mode 100644 benchmarks/TritonNPUBench/vllm/_silu_mul_fp8_quant_deep_gemm/_silu_mul_fp8_quant_deep_gemm.py create mode 100644 benchmarks/TritonNPUBench/vllm/_silu_mul_per_token_group_quant_fp8_colmajor/_silu_mul_per_token_group_quant_fp8_colmajor.py create mode 100644 benchmarks/TritonNPUBench/vllm/_state_passing_fwd_kernel/_state_passing_fwd_kernel.py create mode 100644 benchmarks/TritonNPUBench/vllm/_topk_log_softmax_kernel/_topk_log_softmax_kernel.py create mode 100644 benchmarks/TritonNPUBench/vllm/_triton_mrope_forward/_triton_mrope_forward.py create mode 100644 benchmarks/TritonNPUBench/vllm/_trtllm_prefill_attn_kvfp8_dequant/_trtllm_prefill_attn_kvfp8_dequant.py create mode 100644 benchmarks/TritonNPUBench/vllm/_unpack_seq_triton_kernel/_unpack_seq_triton_kernel.py create mode 100644 benchmarks/TritonNPUBench/vllm/_update_eagle_inputs_kernel/_update_eagle_inputs_kernel.py create mode 100644 benchmarks/TritonNPUBench/vllm/_w8a8_block_int8_matmul/_w8a8_block_int8_matmul.py create mode 100644 benchmarks/TritonNPUBench/vllm/_w8a8_triton_block_scaled_mm/_w8a8_triton_block_scaled_mm.py create mode 100644 benchmarks/TritonNPUBench/vllm/awq_dequantize_kernel/awq_dequantize_kernel.py create mode 100644 benchmarks/TritonNPUBench/vllm/awq_gemm_kernel/awq_gemm_kernel.py create mode 100644 benchmarks/TritonNPUBench/vllm/batched_triton_kernel/batched_triton_kernel.py create mode 100644 benchmarks/TritonNPUBench/vllm/bmm_kernel/bmm_kernel.py create mode 100644 benchmarks/TritonNPUBench/vllm/chunk_fwd_kernel_o/chunk_fwd_kernel_o.py create mode 100644 benchmarks/TritonNPUBench/vllm/chunk_gated_delta_rule_fwd_kernel_h_blockdim64/chunk_gated_delta_rule_fwd_kernel_h_blockdim64.py create mode 100644 benchmarks/TritonNPUBench/vllm/chunk_gla_fwd_kernel_o/chunk_gla_fwd_kernel_o.py create mode 100644 benchmarks/TritonNPUBench/vllm/chunk_kda_scaled_dot_kkt_fwd_kernel_intra_sub_inter/chunk_kda_scaled_dot_kkt_fwd_kernel_intra_sub_inter.py create mode 100644 benchmarks/TritonNPUBench/vllm/chunk_kda_scaled_dot_kkt_fwd_kernel_intra_sub_intra/chunk_kda_scaled_dot_kkt_fwd_kernel_intra_sub_intra.py create mode 100644 benchmarks/TritonNPUBench/vllm/chunk_local_cumsum_scalar_kernel/chunk_local_cumsum_scalar_kernel.py create mode 100644 benchmarks/TritonNPUBench/vllm/chunk_local_cumsum_vector_kernel/chunk_local_cumsum_vector_kernel.py create mode 100644 benchmarks/TritonNPUBench/vllm/chunk_scaled_dot_kkt_fwd_kernel/chunk_scaled_dot_kkt_fwd_kernel.py create mode 100644 benchmarks/TritonNPUBench/vllm/compute_identity_kernel/compute_identity_kernel.py create mode 100644 benchmarks/TritonNPUBench/vllm/cp_mha_gather_cache_kernel/cp_mha_gather_cache_kernel.py create mode 100644 benchmarks/TritonNPUBench/vllm/eagle_prepare_inputs_padded_kernel/eagle_prepare_inputs_padded_kernel.py create mode 100644 benchmarks/TritonNPUBench/vllm/eagle_prepare_next_token_padded_kernel/eagle_prepare_next_token_padded_kernel.py create mode 100644 benchmarks/TritonNPUBench/vllm/expand_kernel/expand_kernel.py create mode 100644 benchmarks/TritonNPUBench/vllm/fused_gdn_gating_kernel/fused_gdn_gating_kernel.py create mode 100644 benchmarks/TritonNPUBench/vllm/fused_moe_kernel/fused_moe_kernel.py create mode 100644 benchmarks/TritonNPUBench/vllm/fused_moe_kernel_gptq_awq/fused_moe_kernel_gptq_awq.py create mode 100644 benchmarks/TritonNPUBench/vllm/fused_recurrent_gated_delta_rule_fwd_kernel/fused_recurrent_gated_delta_rule_fwd_kernel.py create mode 100644 benchmarks/TritonNPUBench/vllm/kda_gate_fwd_kernel/kda_gate_fwd_kernel.py create mode 100644 benchmarks/TritonNPUBench/vllm/kernel_paged_attention_2d/kernel_paged_attention_2d.py create mode 100644 benchmarks/TritonNPUBench/vllm/kernel_unified_attention_2d/kernel_unified_attention_2d.py create mode 100644 benchmarks/TritonNPUBench/vllm/kernel_unified_attention_3d/kernel_unified_attention_3d.py create mode 100644 benchmarks/TritonNPUBench/vllm/l2norm_fwd_kernel/l2norm_fwd_kernel.py create mode 100644 benchmarks/TritonNPUBench/vllm/l2norm_fwd_kernel1/l2norm_fwd_kernel1.py create mode 100644 benchmarks/TritonNPUBench/vllm/l2norm_fwd_kernel2/l2norm_fwd_kernel2.py create mode 100644 benchmarks/TritonNPUBench/vllm/layer_norm_fwd_kernel/layer_norm_fwd_kernel.py create mode 100644 benchmarks/TritonNPUBench/vllm/layer_norm_gated_fwd_kernel/layer_norm_gated_fwd_kernel.py create mode 100644 benchmarks/TritonNPUBench/vllm/layer_norm_gated_fwd_kernel1/layer_norm_gated_fwd_kernel1.py create mode 100644 benchmarks/TritonNPUBench/vllm/matmul_kernel_persistent/matmul_kernel_persistent.py create mode 100644 benchmarks/TritonNPUBench/vllm/mean_kernel/mean_kernel.py create mode 100644 benchmarks/TritonNPUBench/vllm/merge_16x16_to_32x32_inverse_kernel/merge_16x16_to_32x32_inverse_kernel.py create mode 100644 benchmarks/TritonNPUBench/vllm/merge_16x16_to_64x64_inverse_kernel/merge_16x16_to_64x64_inverse_kernel.py create mode 100644 benchmarks/TritonNPUBench/vllm/merge_attn_states_kernel/merge_attn_states_kernel.py create mode 100644 benchmarks/TritonNPUBench/vllm/pack_bitmatrix/pack_bitmatrix.py create mode 100644 benchmarks/TritonNPUBench/vllm/recompute_w_u_fwd_kernel/recompute_w_u_fwd_kernel.py create mode 100644 benchmarks/TritonNPUBench/vllm/reduce_segments/reduce_segments.py create mode 100644 benchmarks/TritonNPUBench/vllm/rejection_greedy_sample_kernel/rejection_greedy_sample_kernel.py create mode 100644 benchmarks/TritonNPUBench/vllm/rejection_random_sample_kernel/rejection_random_sample_kernel.py create mode 100644 benchmarks/TritonNPUBench/vllm/reshape_and_cache_kernel_flash/reshape_and_cache_kernel_flash.py create mode 100644 benchmarks/TritonNPUBench/vllm/reshape_and_cache_kernel_flash_diffkv/reshape_and_cache_kernel_flash_diffkv.py create mode 100644 benchmarks/TritonNPUBench/vllm/sample_recovered_tokens_kernel/sample_recovered_tokens_kernel.py create mode 100644 benchmarks/TritonNPUBench/vllm/scaled_mm_kernel/scaled_mm_kernel.py create mode 100644 benchmarks/TritonNPUBench/vllm/softplus_kernel/softplus_kernel.py create mode 100644 benchmarks/TritonNPUBench/vllm/solve_tril_16x16_kernel/solve_tril_16x16_kernel.py create mode 100644 benchmarks/TritonNPUBench/vllm/test_round_int8_kernel/test_round_int8_kernel.py create mode 100644 benchmarks/TritonNPUBench/vllm/triton_scale_swizzle/triton_scale_swizzle.py create mode 100644 benchmarks/TritonNPUBench/vllm/vllm_gpu_perf.csv diff --git a/benchmarks/TritonNPUBench/vllm/_apply_grammar_bitmask_kernel/_apply_grammar_bitmask_kernel.py b/benchmarks/TritonNPUBench/vllm/_apply_grammar_bitmask_kernel/_apply_grammar_bitmask_kernel.py new file mode 100644 index 00000000..bdb41482 --- /dev/null +++ b/benchmarks/TritonNPUBench/vllm/_apply_grammar_bitmask_kernel/_apply_grammar_bitmask_kernel.py @@ -0,0 +1,42 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/vllm_operator_cases/newtest_cases/test_apply_grammar_bitmask_kernel.py +# Main kernel: _apply_grammar_bitmask_kernel +# PT file: apply_grammar_bitmask_kernel_test_data_v3.pt + +import triton +import triton.language as tl + + +# === _apply_grammar_bitmask_kernel === +@triton.jit +def _apply_grammar_bitmask_kernel( + logits_ptr, + logits_stride, + logits_indices_ptr, + bitmask_ptr, + bitmask_stride, + vocab_size, + BLOCK_SIZE: tl.constexpr, +): + bitmask_idx = tl.program_id(0) + logits_idx = tl.load(logits_indices_ptr + bitmask_idx) + + # Load the bitmask. + block_id = tl.program_id(1) + bitmask_offset = (block_id * BLOCK_SIZE) // 32 + tl.arange(0, BLOCK_SIZE // 32) + packed_bitmask = tl.load( + bitmask_ptr + bitmask_idx * bitmask_stride + bitmask_offset, + mask=bitmask_offset < bitmask_stride, + ) + # Unpack the bitmask. + bitmask = ((packed_bitmask[:, None] >> (tl.arange(0, 32)[None, :])) & 1) == 0 + bitmask = bitmask.reshape(BLOCK_SIZE) + + # Apply the bitmask to the logits. + block_offset = block_id * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + tl.store( + logits_ptr + logits_idx * logits_stride + block_offset, + -float("inf"), + mask=bitmask & (block_offset < vocab_size), + ) + diff --git a/benchmarks/TritonNPUBench/vllm/_apply_write_kernel/_apply_write_kernel.py b/benchmarks/TritonNPUBench/vllm/_apply_write_kernel/_apply_write_kernel.py new file mode 100644 index 00000000..3cf6a16b --- /dev/null +++ b/benchmarks/TritonNPUBench/vllm/_apply_write_kernel/_apply_write_kernel.py @@ -0,0 +1,36 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/vllm_operator_cases/newtest_cases/test__apply_write_kernel.py +# Main kernel: _apply_write_kernel +# PT file: apply_write_kernel_test_data.pt + +import triton +import triton.language as tl + + +# === _apply_write_kernel === +@triton.jit +def _apply_write_kernel( + output_ptr, + output_stride, + write_indices_ptr, + write_starts_ptr, + write_contents_ptr, + write_cu_lens_ptr, + BLOCK_SIZE: tl.constexpr, +): + pid = tl.program_id(0) + row_idx = tl.load(write_indices_ptr + pid) + start_idx = tl.load(write_starts_ptr + pid) + + cu_start = tl.load(write_cu_lens_ptr + pid - 1) if pid > 0 else 0 + cu_end = tl.load(write_cu_lens_ptr + pid) + content_len = cu_end - cu_start + + for i in range(0, content_len, BLOCK_SIZE): + block = i + tl.arange(0, BLOCK_SIZE) + mask = block < content_len + content = tl.load(write_contents_ptr + cu_start + block, mask=mask) + tl.store( + output_ptr + row_idx * output_stride + start_idx + block, content, mask=mask + ) + diff --git a/benchmarks/TritonNPUBench/vllm/_bias_kernel/_bias_kernel.py b/benchmarks/TritonNPUBench/vllm/_bias_kernel/_bias_kernel.py new file mode 100644 index 00000000..7a7c7cd1 --- /dev/null +++ b/benchmarks/TritonNPUBench/vllm/_bias_kernel/_bias_kernel.py @@ -0,0 +1,100 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/vllm_operator_cases/newtest_cases/test_bias_kernel.py +# Main kernel: _bias_kernel +# PT file: bias_kernel_test_v3.pt + +import triton +import triton.language as tl + + +# === _bias_kernel === +@triton.jit +def _bias_kernel( + logits_ptr, + logits_stride, + vocab_size, + idx_mapping_ptr, + # Allowed token IDs. + num_allowed_token_ids_ptr, + allowed_token_ids_ptr, + allowed_token_ids_stride, + # Logit bias. + num_logit_bias_ptr, + bias_token_ids_ptr, + bias_token_ids_stride, + bias_ptr, + bias_stride, + # Min tokens. + pos_ptr, + min_lens_ptr, + num_stop_token_ids_ptr, + stop_token_ids_ptr, + stop_token_ids_stride, + BLOCK_SIZE: tl.constexpr, + LOGITS_BLOCK_SIZE: tl.constexpr, +): + batch_idx = tl.program_id(0) + req_state_idx = tl.load(idx_mapping_ptr + batch_idx) + + block = tl.arange(0, BLOCK_SIZE) + + # Allowed token IDs. + num_allowed_token_ids = tl.load(num_allowed_token_ids_ptr + req_state_idx) + if num_allowed_token_ids > 0: + block = tl.arange(0, BLOCK_SIZE) + mask = block < num_allowed_token_ids + + # Save logits for allowed token IDs. + allowed_token_ids = tl.load( + allowed_token_ids_ptr + req_state_idx * allowed_token_ids_stride + block, + mask=mask, + ) + logits = tl.load( + logits_ptr + batch_idx * logits_stride + allowed_token_ids, mask=mask + ) + + # Set logits to -inf for all tokens. + for i in range(0, vocab_size, LOGITS_BLOCK_SIZE): + offset = i + tl.arange(0, LOGITS_BLOCK_SIZE) + tl.store( + logits_ptr + batch_idx * logits_stride + offset, + -float("inf"), + mask=offset < vocab_size, + ) + + # Restore logits for allowed token IDs. + tl.store( + logits_ptr + batch_idx * logits_stride + allowed_token_ids, + logits, + mask=mask, + ) + + # Logit bias. + num_logit_bias = tl.load(num_logit_bias_ptr + req_state_idx) + if num_logit_bias > 0: + mask = block < num_logit_bias + token_ids = tl.load( + bias_token_ids_ptr + req_state_idx * bias_token_ids_stride + block, + mask=mask, + ) + bias = tl.load(bias_ptr + req_state_idx * bias_stride + block, mask=mask) + logits = tl.load(logits_ptr + batch_idx * logits_stride + token_ids, mask=mask) + logits += bias + tl.store(logits_ptr + batch_idx * logits_stride + token_ids, logits, mask=mask) + + # Apply min tokens. + num_stop_token_ids = tl.load(num_stop_token_ids_ptr + req_state_idx) + pos = tl.load(pos_ptr + batch_idx) + min_len = tl.load(min_lens_ptr + req_state_idx) + if num_stop_token_ids > 0 and pos < min_len: + mask = block < num_stop_token_ids + stop_token_ids = tl.load( + stop_token_ids_ptr + req_state_idx * stop_token_ids_stride + block, + mask=mask, + ) + tl.store( + logits_ptr + batch_idx * logits_stride + stop_token_ids, + -float("inf"), + mask=mask, + ) + diff --git a/benchmarks/TritonNPUBench/vllm/_bincount_kernel/_bincount_kernel.py b/benchmarks/TritonNPUBench/vllm/_bincount_kernel/_bincount_kernel.py new file mode 100644 index 00000000..16b36961 --- /dev/null +++ b/benchmarks/TritonNPUBench/vllm/_bincount_kernel/_bincount_kernel.py @@ -0,0 +1,37 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/vllm_operator_cases/newtest_cases/test__bincount_kernel.py +# Main kernel: _bincount_kernel +# PT file: bincount_kernel_test_v2.pt + +import triton +import triton.language as tl + + +# === _bincount_kernel === +@triton.jit(do_not_specialize=["prefill_len", "prompt_len"]) +def _bincount_kernel( + prefill_token_ids_ptr, + prefill_len, + prompt_len, + prompt_bin_mask_ptr, + output_bin_counts_ptr, + BLOCK_SIZE: tl.constexpr, +): + block_idx = tl.program_id(0) + if block_idx * BLOCK_SIZE >= prefill_len: + return + + block = block_idx * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + if block_idx * BLOCK_SIZE < prompt_len: + mask = block < prompt_len + prefill_tokens = tl.load(prefill_token_ids_ptr + block, mask=mask) + idx = prefill_tokens // 32 + bit_idx = prefill_tokens % 32 + bit = tl.full((BLOCK_SIZE,), 1, tl.int32) << bit_idx + tl.atomic_or(prompt_bin_mask_ptr + idx, bit, mask=mask) + if (block_idx + 1) * BLOCK_SIZE >= prompt_len: + mask = block < prefill_len + mask &= block >= prompt_len + prefill_tokens = tl.load(prefill_token_ids_ptr + block, mask=mask) + tl.atomic_add(output_bin_counts_ptr + prefill_tokens, 1, mask=mask) + diff --git a/benchmarks/TritonNPUBench/vllm/_bmm_chunk_fwd_kernel/_bmm_chunk_fwd_kernel.py b/benchmarks/TritonNPUBench/vllm/_bmm_chunk_fwd_kernel/_bmm_chunk_fwd_kernel.py new file mode 100644 index 00000000..e7752500 --- /dev/null +++ b/benchmarks/TritonNPUBench/vllm/_bmm_chunk_fwd_kernel/_bmm_chunk_fwd_kernel.py @@ -0,0 +1,122 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/vllm_operator_cases/newtest_cases/test__bmm_chunk_fwd_kernel.py +# Main kernel: _bmm_chunk_fwd_kernel +# PT file: test__bmm_chunk_fwd_kernel_v2.pt + +import triton +import triton.language as tl + + +# === _bmm_chunk_fwd_kernel === +@triton.jit +def _bmm_chunk_fwd_kernel( + # Pointers to matrices + a_ptr, + b_ptr, + out_ptr, + seq_idx_ptr, + # Matrix dimensions + seqlen, + chunk_size, + K, + ngroups, + stride_a_batch, + stride_a_seqlen, + stride_a_head, + stride_ak, + stride_b_batch, + stride_b_seqlen, + stride_b_head, + stride_bk, + stride_out_batch, + stride_out_chunk, + stride_out_head, + stride_outm, + stride_outn, + stride_seq_idx_batch, + stride_seq_idx_seqlen, + # Meta-parameters + IS_CAUSAL: tl.constexpr, + dot_dtype: tl.constexpr, + HAS_SEQ_IDX: tl.constexpr, + BLOCK_SIZE_M: tl.constexpr = 16, + BLOCK_SIZE_N: tl.constexpr = 16, + BLOCK_SIZE_K: tl.constexpr = 16, +): + pid_b = tl.program_id(axis=1) + pid_ch = tl.program_id(axis=2).to(tl.int64) + pid_c = pid_ch // ngroups + pid_h = pid_ch - pid_c * ngroups + num_pid_n = tl.cdiv(chunk_size, BLOCK_SIZE_N) + pid_m = tl.program_id(axis=0) // num_pid_n + pid_n = tl.program_id(axis=0) % num_pid_n + if IS_CAUSAL: + if pid_n * BLOCK_SIZE_N >= (pid_m + 1) * BLOCK_SIZE_M: + return + a_ptr += ( + pid_b * stride_a_batch + + pid_c * chunk_size * stride_a_seqlen + + pid_h * stride_a_head + ) + b_ptr += ( + pid_b * stride_b_batch + + pid_c * chunk_size * stride_b_seqlen + + pid_h * stride_b_head + ) + if HAS_SEQ_IDX: + seq_idx_ptr += ( + pid_b * stride_seq_idx_batch + pid_c * chunk_size * stride_seq_idx_seqlen + ) + + offs_m = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) + offs_n = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) + offs_k = tl.arange(0, BLOCK_SIZE_K) + a_ptrs = a_ptr + (offs_m[:, None] * stride_a_seqlen + offs_k[None, :] * stride_ak) + b_ptrs = b_ptr + (offs_k[:, None] * stride_bk + offs_n[None, :] * stride_b_seqlen) + chunk_size_limit = min(chunk_size, seqlen - pid_c * chunk_size) + + acc = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32) + for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)): + a = tl.load( + a_ptrs, + mask=(offs_m[:, None] < chunk_size_limit) + & (offs_k[None, :] < K - k * BLOCK_SIZE_K), + other=0.0, + ).to(dot_dtype) + b = tl.load( + b_ptrs, + mask=(offs_k[:, None] < K - k * BLOCK_SIZE_K) + & (offs_n[None, :] < chunk_size_limit), + other=0.0, + ).to(dot_dtype) + acc += tl.dot(a, b) + a_ptrs += BLOCK_SIZE_K * stride_ak + b_ptrs += BLOCK_SIZE_K * stride_bk + + offs_m = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) + offs_n = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) + if HAS_SEQ_IDX: + chunk_size_limit = min(chunk_size, seqlen - pid_c * chunk_size) + seq_idx_m = tl.load( + seq_idx_ptr + offs_m * stride_seq_idx_seqlen, + mask=offs_m < chunk_size_limit, + other=-1, + ) + seq_idx_n = tl.load( + seq_idx_ptr + offs_n * stride_seq_idx_seqlen, + mask=offs_n < chunk_size_limit, + other=-2, + ) + acc = tl.where(seq_idx_m[:, None] == seq_idx_n[None, :], acc, 0.0) + out = acc.to(out_ptr.dtype.element_ty) + + out_ptr += ( + pid_b * stride_out_batch + pid_c * stride_out_chunk + pid_h * stride_out_head + ) + out_ptrs = out_ptr + (stride_outm * offs_m[:, None] + offs_n[None, :] * stride_outn) + tl.store( + out_ptrs, + out, + mask=(offs_m[:, None] < chunk_size) & (offs_n[None, :] < chunk_size), + ) + diff --git a/benchmarks/TritonNPUBench/vllm/_causal_conv1d_fwd_kernel/_causal_conv1d_fwd_kernel.py b/benchmarks/TritonNPUBench/vllm/_causal_conv1d_fwd_kernel/_causal_conv1d_fwd_kernel.py new file mode 100644 index 00000000..14e50ca8 --- /dev/null +++ b/benchmarks/TritonNPUBench/vllm/_causal_conv1d_fwd_kernel/_causal_conv1d_fwd_kernel.py @@ -0,0 +1,462 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/vllm_operator_cases/newtest_cases/test__causal_conv1d_fwd_kernel.py +# Main kernel: _causal_conv1d_fwd_kernel +# PT file: _causal_conv1d_fwd_kernel_v2.pt + +import triton +import triton.language as tl + + +# === _causal_conv1d_fwd_kernel === +@triton.jit() +def _causal_conv1d_fwd_kernel( # continuous batching + # Pointers to matrices + x_ptr, # (dim, cu_seqlen) holding `batch` of actual sequences + padded sequences + w_ptr, # (dim, width) + bias_ptr, + initial_states_ptr, # conv_states_ptr + cache_indices_ptr, # (batch, n_blocks + padding) The second dimension contains + # the block indices relevant for each sequence + # plus potential 0-padding at the beginning and at the end + has_initial_states_ptr, + query_start_loc_ptr, + batch_ptr, + token_chunk_offset_ptr, + block_idx_first_scheduled_token, # (batch,) + block_idx_last_scheduled_token, # (batch,) + initial_state_idx, # (batch,) + num_computed_tokens, # (batch,) + o_ptr, # (dim, seqlen) - actually pointing to x_ptr + # Matrix dimensions + dim: tl.constexpr, + seqlen: tl.int32, # cu_seqlen + num_cache_lines: tl.constexpr, # added to support vLLM larger cache lines + # Strides + stride_x_dim: tl.constexpr, # stride to get to next feature-value, + stride_x_token: tl.constexpr, # stride to get to next token (same feature-index, same sequence-index) + stride_w_dim: tl.constexpr, # stride to get to next dim-axis value + stride_w_width: tl.constexpr, # stride to get to next width-axis value + stride_istate_seq: tl.constexpr, + stride_istate_dim: tl.constexpr, + stride_istate_token: tl.constexpr, + stride_cache_indices: tl.constexpr, + stride_o_dim: tl.constexpr, + stride_o_token: tl.constexpr, + stride_block_m: tl.constexpr, # Stride block to align divided by BLOCK_M + # others + pad_slot_id: tl.constexpr, + # Meta-parameters + HAS_BIAS: tl.constexpr, + KERNEL_WIDTH: tl.constexpr, + SILU_ACTIVATION: tl.constexpr, + IS_APC_ENABLED: tl.constexpr, + USE_PAD_SLOT: tl.constexpr, + NP2_STATELEN: tl.constexpr, + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, +): + conv_states_ptr = initial_states_ptr + conv_state_indices_ptr = cache_indices_ptr + stride_conv_state_seq = stride_istate_seq + stride_conv_state_dim = stride_istate_dim + stride_conv_state_tok = stride_istate_token + state_len = ( + KERNEL_WIDTH - 1 + ) # can be passed via argument if it's not the same as this value + + # one program handles one chunk in a single sequence + # rather than mixing sequences - to make updating initial_states across sequences efficiently + + # single-sequence id + idx_seq = tl.load(batch_ptr + tl.program_id(0)).to(tl.int64) + chunk_offset = tl.load(token_chunk_offset_ptr + tl.program_id(0)) + + # BLOCK_N elements along the feature-dimension (channel) + idx_feats = tl.program_id(1) * BLOCK_N + tl.arange(0, BLOCK_N) + + if idx_seq == pad_slot_id: + return + + sequence_start_index = tl.load(query_start_loc_ptr + idx_seq) + sequence_end_index = tl.load(query_start_loc_ptr + idx_seq + 1) + # find the actual sequence length + seqlen = sequence_end_index - sequence_start_index + + B_size: tl.constexpr = stride_block_m * BLOCK_M + + if IS_APC_ENABLED: + # Handle the case if prefix caching is enabled. + # In particular, if prefix caching is enabled, the program write additional cache states to "cache_indices_ptr" + + # Get the length of the completed sequence so far and compute the offset. + current_first_index = tl.load(block_idx_first_scheduled_token + idx_seq) + current_last_index = tl.load(block_idx_last_scheduled_token + idx_seq) + sequence_completed_index = tl.load(num_computed_tokens + idx_seq) + + # Compute the offset where the first stride_block_m-aligned first full block is + # Value in "token-space" + sequence_completed_offset_token = sequence_completed_index % B_size + seq_completed_offset = B_size - sequence_completed_offset_token + seq_end_offset = (seqlen - seq_completed_offset) % B_size + last_full_block_token_index = sequence_end_index - seq_end_offset + # If the sequence without the sequence_offset_index is stride_cache_chunk-aligned, then the last full chunk is the second-to-last one + if seq_end_offset == 0: + last_full_block_token_index = last_full_block_token_index - B_size + + # Get the number of blocks to be filled for the current sequence + # If n_block_to_fill = 0, then only the state at the sequence end is stored + n_block_to_fill = current_last_index - current_first_index + + # Get the index of the init block + conv_state_init_index = tl.load(initial_state_idx + idx_seq) + else: + n_block_to_fill = 0 + current_last_index = 0 + conv_state_init_index = 0 + current_first_index = 0 + last_full_block_token_index = 0 + + token_offset = BLOCK_M * chunk_offset + segment_len = min(BLOCK_M, seqlen - token_offset) + + # base of the sequence + x_base = ( + x_ptr + sequence_start_index * stride_x_token + idx_feats * stride_x_dim + ) # [BLOCK_N,] + + # cache_idx + conv_states_input_coord = tl.load( + conv_state_indices_ptr + idx_seq * stride_cache_indices + conv_state_init_index + ).to(tl.int64) + + if USE_PAD_SLOT: # noqa + if conv_states_input_coord == pad_slot_id: + # not processing as this is not the actual sequence + return + conv_states_base = ( + conv_states_ptr + + (conv_states_input_coord * stride_conv_state_seq) + + (idx_feats * stride_conv_state_dim) + ) # [BLOCK_N,] + + w_base = w_ptr + (idx_feats * stride_w_dim) # [BLOCK_N,] + + # Does 2 things: + # 1. READ prior-block init-state data - [done by every Triton programs] + # 2. update conv_state with new data [only by the Triton program handles chunk_offset=0] + if chunk_offset == 0: + # read from conv_states + load_init_state = tl.load(has_initial_states_ptr + idx_seq).to(tl.int1) + if load_init_state: + # load from conv_states + prior_tokens = conv_states_base + (state_len - 1) * stride_conv_state_tok + mask_w = idx_feats < dim + if KERNEL_WIDTH == 2: + conv_states_ptrs = prior_tokens # [BLOCK_N] + col0 = tl.load(conv_states_ptrs, mask_w, 0.0) + if KERNEL_WIDTH == 3: + conv_states_ptrs = prior_tokens # [BLOCK_N] + col1 = tl.load(conv_states_ptrs, mask_w, 0.0) + conv_states_ptrs = prior_tokens - 1 * stride_conv_state_tok # [BLOCK_N] + col0 = tl.load(conv_states_ptrs, mask_w, 0.0) + if KERNEL_WIDTH == 4: + conv_states_ptrs = prior_tokens # [BLOCK_N] + col2 = tl.load(conv_states_ptrs, mask_w, 0.0) + conv_states_ptrs = prior_tokens - 1 * stride_conv_state_tok # [BLOCK_N] + col1 = tl.load(conv_states_ptrs, mask_w, 0.0) + conv_states_ptrs = prior_tokens - 2 * stride_conv_state_tok # [BLOCK_N] + col0 = tl.load(conv_states_ptrs, mask_w, 0.0) + if KERNEL_WIDTH == 5: + conv_states_ptrs = prior_tokens # [BLOCK_N] + col3 = tl.load(conv_states_ptrs, mask_w, 0.0) + conv_states_ptrs = prior_tokens - 1 * stride_conv_state_tok # [BLOCK_N] + col2 = tl.load(conv_states_ptrs, mask_w, 0.0) + conv_states_ptrs = prior_tokens - 2 * stride_conv_state_tok # [BLOCK_N] + col1 = tl.load(conv_states_ptrs, mask_w, 0.0) + conv_states_ptrs = prior_tokens - 3 * stride_conv_state_tok # [BLOCK_N] + col0 = tl.load(conv_states_ptrs, mask_w, 0.0) + else: + # prior-tokens are zeros + if KERNEL_WIDTH >= 2: # STRATEGY1 + # first chunk and does not have prior-token, so just set to 0 + col0 = tl.zeros((BLOCK_N,), dtype=x_ptr.dtype.element_ty) + if KERNEL_WIDTH >= 3: # STRATEGY1 + col1 = tl.zeros((BLOCK_N,), dtype=x_ptr.dtype.element_ty) + if KERNEL_WIDTH >= 4: # STRATEGY1 + col2 = tl.zeros((BLOCK_N,), dtype=x_ptr.dtype.element_ty) + if KERNEL_WIDTH >= 5: # STRATEGY1 + col3 = tl.zeros((BLOCK_N,), dtype=x_ptr.dtype.element_ty) + + # STEP 2: + # here prepare data for updating conv_state + if ( + state_len <= seqlen + ): # SMALL_CACHE=True (only move part of 'x' into conv_state cache) + # just read from 'x' + # copy 'x' data to conv_state + # load only 'x' data (and set 0 before 'x' if seqlen < state_len) + idx_tokens_last = (seqlen - state_len) + tl.arange( + 0, NP2_STATELEN + ) # [BLOCK_M] + x_ptrs = ( + x_ptr + + ((sequence_start_index + idx_tokens_last) * stride_x_token)[:, None] + + (idx_feats * stride_x_dim)[None, :] + ) # [BLOCK_M,BLOCK_N,] + mask_x = ( + (idx_tokens_last >= 0)[:, None] + & (idx_tokens_last < seqlen)[:, None] + & (idx_feats < dim)[None, :] + ) # token-index # token-index # feature-index + loaded_x = tl.load(x_ptrs, mask_x, 0.0) + idx_tokens_conv = tl.arange(0, NP2_STATELEN) # [BLOCK_M] + + # Compute the offset where the last block should be written in the conv_states + conv_states_output_coord = tl.load( + conv_state_indices_ptr + + idx_seq * stride_cache_indices + + current_last_index + ).to(tl.int64) + + conv_states_ptrs_target = ( + conv_states_ptr + + (conv_states_output_coord * stride_conv_state_seq) # Offset from seq + + (idx_feats * stride_conv_state_dim) + )[None, :] + ( # [BLOCK_N,] + idx_tokens_conv * stride_conv_state_tok + )[:, None] + + mask = (idx_tokens_conv < state_len)[:, None] & (idx_feats < dim)[None, :] + tl.debug_barrier() # NOTE: use this due to bug in Triton compiler + tl.store(conv_states_ptrs_target, loaded_x, mask) + + else: + if load_init_state: + # update conv_state by shifting left, i.e. take last few cols from conv_state + cols from 'x' + idx_tokens_conv = tl.arange(0, NP2_STATELEN) # [BLOCK_M] + + conv_states_ptrs_source = ( + conv_states_ptr + + (conv_states_input_coord * stride_conv_state_seq) + + (idx_feats * stride_conv_state_dim)[None, :] + + ((idx_tokens_conv + seqlen) * stride_conv_state_tok)[:, None] + ) # [BLOCK_M, BLOCK_N] + mask = ( + (conv_states_input_coord < num_cache_lines) + & ((idx_tokens_conv + seqlen) < state_len)[:, None] + & (idx_feats < dim)[None, :] + ) + conv_state = tl.load(conv_states_ptrs_source, mask, other=0.0) + + VAL = state_len - seqlen + + x_ptrs = ( + x_base[None, :] + + ((idx_tokens_conv - VAL) * stride_x_token)[:, None] + ) # [BLOCK_M, BLOCK_N] + + mask_x = ( + (idx_tokens_conv - VAL >= 0)[:, None] + & (idx_tokens_conv - VAL < seqlen)[:, None] + & (idx_feats < dim)[None, :] + ) # token-index # token-index # feature-index + loaded_x = tl.load(x_ptrs, mask_x, 0.0) + + tl.debug_barrier() # need this due to the bug in tl.where not enforcing this when data is the result of another tl.load + new_conv_state = tl.where( + mask, conv_state, loaded_x + ) # BUG in 'tl.where' which requires a barrier before this + conv_states_ptrs_target = ( + conv_states_base + + (idx_tokens_conv * stride_conv_state_tok)[:, None] + ) # [BLOCK_M, BLOCK_N] + mask = (idx_tokens_conv < state_len)[:, None] & (idx_feats < dim)[ + None, : + ] + tl.store(conv_states_ptrs_target, new_conv_state, mask) + else: # load_init_state == False + # update conv_state by shifting left, BUT + # set cols prior to 'x' as zeros + cols from 'x' + idx_tokens_conv = tl.arange(0, NP2_STATELEN) # [BLOCK_M] + + VAL = state_len - seqlen + + x_ptrs = ( + x_base[None, :] + + ((idx_tokens_conv - VAL) * stride_x_token)[:, None] + ) # [BLOCK_M, BLOCK_N] + + mask_x = ( + (idx_tokens_conv - VAL >= 0)[:, None] + & (idx_tokens_conv - VAL < seqlen)[:, None] + & (idx_feats < dim)[None, :] + ) # token-index # token-index # feature-index + new_conv_state = tl.load(x_ptrs, mask_x, 0.0) + + conv_states_ptrs_target = ( + conv_states_base + + (idx_tokens_conv * stride_conv_state_tok)[:, None] + ) # [BLOCK_M, BLOCK_N] + mask = (idx_tokens_conv < state_len)[:, None] & (idx_feats < dim)[ + None, : + ] + tl.store(conv_states_ptrs_target, new_conv_state, mask) + + else: # chunk_offset > 0 + # read prior-token data from `x` + load_init_state = True + prior_tokens = x_base + (token_offset - 1) * stride_x_token + mask_w = idx_feats < dim + if KERNEL_WIDTH == 2: + conv_states_ptrs = prior_tokens # [BLOCK_N] + col0 = tl.load(conv_states_ptrs, mask_w, 0.0, cache_modifier=".ca") + if KERNEL_WIDTH == 3: + conv_states_ptrs = prior_tokens # [BLOCK_N] + col1 = tl.load(conv_states_ptrs, mask_w, 0.0, cache_modifier=".ca") + conv_states_ptrs = prior_tokens - 1 * stride_x_token # [BLOCK_N] + col0 = tl.load(conv_states_ptrs, mask_w, 0.0, cache_modifier=".ca") + if KERNEL_WIDTH == 4: + conv_states_ptrs = prior_tokens # [BLOCK_N] + col2 = tl.load(conv_states_ptrs, mask_w, 0.0, cache_modifier=".ca") + conv_states_ptrs = prior_tokens - 1 * stride_x_token # [BLOCK_N] + col1 = tl.load(conv_states_ptrs, mask_w, 0.0, cache_modifier=".ca") + conv_states_ptrs = prior_tokens - 2 * stride_x_token # [BLOCK_N] + col0 = tl.load(conv_states_ptrs, mask_w, 0.0, cache_modifier=".ca") + if KERNEL_WIDTH == 5: + # ruff: noqa: F841 + conv_states_ptrs = prior_tokens # [BLOCK_N] + col3 = tl.load(conv_states_ptrs, mask_w, 0.0, cache_modifier=".ca") + conv_states_ptrs = prior_tokens - 1 * stride_x_token # [BLOCK_N] + col2 = tl.load(conv_states_ptrs, mask_w, 0.0, cache_modifier=".ca") + conv_states_ptrs = prior_tokens - 2 * stride_x_token # [BLOCK_N] + col1 = tl.load(conv_states_ptrs, mask_w, 0.0, cache_modifier=".ca") + conv_states_ptrs = prior_tokens - 3 * stride_x_token # [BLOCK_N] + col0 = tl.load(conv_states_ptrs, mask_w, 0.0, cache_modifier=".ca") + + # Store intermediate states aligned with stride_block_m + # The additional states are cached starting from the last stride_block_m. + # For example: + # If n_block_to_fill = 0, then only the state at the sequence end is cached and the process below is not involved. + # If n_block_to_fill > 0, then the states at the sequence end and at the n_block_to_fill-last + # stride_block_m are cached. + # For example chunk_offset = n_block_to_fill stores the state at last_full_block + if (chunk_offset - 1) < n_block_to_fill: + # Store the states at the chunk boundaries from the start of the sequence + idx_tokens_last = ( + last_full_block_token_index + - (n_block_to_fill - chunk_offset) * B_size + - state_len + ) + tl.arange(0, NP2_STATELEN) # [BLOCK_M] + x_ptrs = ( + x_ptr + + (idx_tokens_last * stride_x_token)[:, None] + + (idx_feats * stride_x_dim)[None, :] + ) # [BLOCK_M,BLOCK_N,] + + mask_x = (idx_tokens_last >= 0)[:, None] & (idx_feats < dim)[ + None, : + ] # token-index # token-index # feature-index + loaded_x = tl.load(x_ptrs, mask_x, 0.0) + idx_tokens_conv = tl.arange(0, NP2_STATELEN) # [BLOCK_M] + + # cache_idx + conv_states_output_coord = tl.load( + conv_state_indices_ptr + + idx_seq * stride_cache_indices + + current_first_index + + (chunk_offset - 1) + ).to(tl.int64) + + conv_states_ptrs_target = ( + conv_states_ptr + + (conv_states_output_coord * stride_conv_state_seq) # Offset from seq + + (idx_feats * stride_conv_state_dim) + )[None, :] + ( # [BLOCK_N,] + idx_tokens_conv * stride_conv_state_tok + )[:, None] + + mask = (idx_tokens_conv < state_len)[:, None] & (idx_feats < dim)[None, :] + tl.debug_barrier() # NOTE: use this due to bug in Triton compiler + tl.store(conv_states_ptrs_target, loaded_x, mask) + + if HAS_BIAS: + bias = bias_ptr + idx_feats + mask_bias = idx_feats < dim + acc_preload = tl.load(bias, mask=mask_bias, other=0.0).to( + tl.float32 + ) # [BLOCK_N] + else: + acc_preload = tl.zeros((BLOCK_N,), dtype=tl.float32) + + x_base_1d = x_base + token_offset * stride_x_token # starting of chunk + + # PRE-LOAD WEIGHTS + mask_w = idx_feats < dim + if KERNEL_WIDTH >= 2: + w_ptrs = w_base + (0 * stride_w_width) # [BLOCK_N] tensor + w_col0 = tl.load(w_ptrs, mask_w, other=0.0) + w_ptrs = w_base + (1 * stride_w_width) # [BLOCK_N] tensor + w_col1 = tl.load(w_ptrs, mask_w, other=0.0) + if KERNEL_WIDTH >= 3: + w_ptrs = w_base + (2 * stride_w_width) # [BLOCK_N] tensor + w_col2 = tl.load(w_ptrs, mask_w, other=0.0) + if KERNEL_WIDTH >= 4: + w_ptrs = w_base + (3 * stride_w_width) # [BLOCK_N] tensor + w_col3 = tl.load(w_ptrs, mask_w, other=0.0) + mask_x_1d = idx_feats < dim + for idx_token in range(segment_len): + acc = acc_preload + + matrix_w = w_col0 + matrix_x = col0 + for j in tl.static_range(KERNEL_WIDTH): + if KERNEL_WIDTH == 2: + if j == 1: # KERNEL_WIDTH-1: + matrix_w = w_col1 + x_ptrs_1d = x_base_1d + idx_token * stride_x_token # [BLOCK_N] + matrix_x = tl.load(x_ptrs_1d, mask=mask_x_1d) + elif KERNEL_WIDTH == 3: + if j == 1: + matrix_w = w_col1 + matrix_x = col1 + elif j == 2: + matrix_w = w_col2 + x_ptrs_1d = x_base_1d + idx_token * stride_x_token # [BLOCK_N] + matrix_x = tl.load(x_ptrs_1d, mask=mask_x_1d) + elif KERNEL_WIDTH == 4: + if j == 1: + matrix_w = w_col1 + matrix_x = col1 + elif j == 2: + matrix_w = w_col2 + matrix_x = col2 + elif j == 3: + matrix_w = w_col3 + x_ptrs_1d = x_base_1d + idx_token * stride_x_token # [BLOCK_N] + matrix_x = tl.load(x_ptrs_1d, mask=mask_x_1d) + + acc += matrix_x * matrix_w # [BLOCK_N] + + if KERNEL_WIDTH == 2: + col0 = matrix_x + elif KERNEL_WIDTH == 3: + col0 = col1 + col1 = matrix_x + elif KERNEL_WIDTH == 4: + col0 = col1 + col1 = col2 + col2 = matrix_x + + if SILU_ACTIVATION: + acc = acc / (1 + tl.exp(-acc)) + mask_1d = (idx_token < segment_len) & ( + idx_feats < dim + ) # token-index # feature-index + o_ptrs = ( + o_ptr + + (sequence_start_index + token_offset + idx_token) * stride_o_token + + (idx_feats * stride_o_dim) + ) + + tl.store(o_ptrs, acc, mask=mask_1d) + diff --git a/benchmarks/TritonNPUBench/vllm/_causal_conv1d_update_kernel/_causal_conv1d_update_kernel.py b/benchmarks/TritonNPUBench/vllm/_causal_conv1d_update_kernel/_causal_conv1d_update_kernel.py new file mode 100644 index 00000000..f34cc312 --- /dev/null +++ b/benchmarks/TritonNPUBench/vllm/_causal_conv1d_update_kernel/_causal_conv1d_update_kernel.py @@ -0,0 +1,330 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/vllm_operator_cases/newtest_cases/test__causal_conv1d_update_kernel.py +# Main kernel: _causal_conv1d_update_kernel +# PT file: _causal_conv1d_update_kernel_v2.pt + +import triton +import triton.language as tl + + +# === _causal_conv1d_update_kernel === +@triton.jit() +def _causal_conv1d_update_kernel( + # Pointers to matrices + x_ptr, # (batch, dim, seqlen) + w_ptr, # (dim, width) + bias_ptr, + conv_state_ptr, + conv_state_indices_ptr, + num_accepted_tokens_ptr, + query_start_loc_ptr, # (batch + 1) + block_idx_last_scheduled_token, # (batch,) + initial_state_idx, # (batch,) + o_ptr, # (batch, dim, seqlen) + # Matrix dimensions + batch: int, + dim: tl.constexpr, + seqlen: tl.constexpr, + state_len: tl.constexpr, + num_cache_lines: tl.constexpr, # added to support vLLM larger cache lines + # Strides + stride_x_seq: tl.constexpr, + stride_x_dim: tl.constexpr, + stride_x_token: tl.constexpr, + stride_w_dim: tl.constexpr, + stride_w_width: tl.constexpr, + stride_conv_state_seq: tl.constexpr, + stride_conv_state_dim: tl.constexpr, + stride_conv_state_tok: tl.constexpr, + stride_state_indices: tl.constexpr, + stride_o_seq: tl.constexpr, + stride_o_dim: tl.constexpr, + stride_o_token: tl.constexpr, + # others + pad_slot_id: tl.constexpr, + # Meta-parameters + HAS_BIAS: tl.constexpr, + KERNEL_WIDTH: tl.constexpr, + SILU_ACTIVATION: tl.constexpr, + IS_VARLEN: tl.constexpr, + IS_APC_ENABLED: tl.constexpr, + IS_SPEC_DECODING: tl.constexpr, + NP2_STATELEN: tl.constexpr, + USE_PAD_SLOT: tl.constexpr, + BLOCK_N: tl.constexpr, +): + # ruff: noqa: E501 + idx_seq = tl.program_id(0) + if idx_seq >= batch: + return + + # [BLOCK_N,] elements along the feature-dimension (channel) + idx_feats = tl.program_id(1) * BLOCK_N + tl.arange(0, BLOCK_N) + + if IS_APC_ENABLED: + # Get the state from the initial_state_idx + conv_state_init = tl.load(initial_state_idx + idx_seq) + current_last_index = tl.load(block_idx_last_scheduled_token + idx_seq) + else: + conv_state_init = 0 + current_last_index = 0 + + # cache_idx + conv_states_input_coord = tl.load( + conv_state_indices_ptr + idx_seq * stride_state_indices + conv_state_init + ).to(tl.int64) + + if USE_PAD_SLOT: # noqa + if conv_states_input_coord == pad_slot_id: + # not processing as this is not the actual sequence + return + + if IS_VARLEN: + query_start_index = tl.load(query_start_loc_ptr + idx_seq).to(tl.int64) + query_end_index = tl.load(query_start_loc_ptr + (idx_seq + 1)).to(tl.int64) + # revise state_len and seqlen + state_len = state_len - (seqlen - (query_end_index - query_start_index)) + seqlen = query_end_index - query_start_index + x_offset = query_start_index * stride_x_token + o_offset = query_start_index * stride_o_token + else: + query_start_index = idx_seq * seqlen + query_end_index = query_start_index + seqlen + x_offset = idx_seq * stride_x_seq + o_offset = idx_seq * stride_o_seq + + if query_start_index == query_end_index: + return + + if IS_SPEC_DECODING: + # The rolling of conv state: + # + # Before forward, the conv_state is: + # [history1, history2, ..., historyM]. + # + # After forward, the conv_state becomes: + # [history2, ..., historyM, draft1, draft2, ..., draftN]. + # + # After acceptance, it becomes: + # + # - accept 1 tokens: [history2, ..., historyM, draft1] + # - accept 2 tokens: [history3, ..., historyM, draft1, draft2] + # - and so on. + conv_state_token_offset = ( + tl.load(num_accepted_tokens_ptr + idx_seq).to(tl.int64) - 1 + ) + else: + conv_state_token_offset = 0 + + # STEP 1: READ init_state data + conv_states_base = ( + conv_state_ptr + + (conv_states_input_coord * stride_conv_state_seq) + + (idx_feats * stride_conv_state_dim) + ) + mask_w = idx_feats < dim + + prior_tokens = conv_states_base + conv_state_token_offset * stride_conv_state_tok + if KERNEL_WIDTH >= 2: + conv_states_ptrs = prior_tokens # [BLOCK_N] + col0 = tl.load(conv_states_ptrs, mask_w, 0.0) + if KERNEL_WIDTH >= 3: + conv_states_ptrs = prior_tokens + 1 * stride_conv_state_tok # [BLOCK_N] + col1 = tl.load(conv_states_ptrs, mask_w, 0.0) + if KERNEL_WIDTH >= 4: + conv_states_ptrs = prior_tokens + 2 * stride_conv_state_tok # [BLOCK_N] + col2 = tl.load(conv_states_ptrs, mask_w, 0.0) + if KERNEL_WIDTH >= 5: + conv_states_ptrs = prior_tokens + 3 * stride_conv_state_tok # [BLOCK_N] + col3 = tl.load(conv_states_ptrs, mask_w, 0.0) + if KERNEL_WIDTH >= 6: + conv_states_ptrs = prior_tokens + 4 * stride_conv_state_tok # [BLOCK_N] + col4 = tl.load(conv_states_ptrs, mask_w, 0.0) + + # STEP 2: assume state_len > seqlen + idx_tokens = tl.arange(0, NP2_STATELEN) # [BLOCK_M] + + # With speculative decoding, the conv_state updates works in a sliding + # window manner, at each forward pass, the tokens are shift by 1, so we + # load since idx_tokens + 1. + conv_state_ptrs_source = ( + conv_state_ptr + + (conv_states_input_coord * stride_conv_state_seq) + + conv_state_token_offset * stride_conv_state_tok + + (idx_feats * stride_conv_state_dim)[None, :] + + ((idx_tokens + (1 if IS_SPEC_DECODING else seqlen)) * stride_conv_state_tok)[ + :, None + ] + ) # [BLOCK_M, BLOCK_N] + mask = ( + (conv_states_input_coord < num_cache_lines) + & ((idx_tokens + seqlen) < state_len)[:, None] + & (idx_feats < dim)[None, :] + ) + conv_state = tl.load(conv_state_ptrs_source, mask, other=0.0) + + VAL = state_len - seqlen + x_base = x_ptr + x_offset + (idx_feats * stride_x_dim) # [BLOCK_N] + + x_ptrs = ( + x_base[None, :] + ((idx_tokens - VAL) * stride_x_token)[:, None] + ) # [BLOCK_M, BLOCK_N] + + mask_x = ( + (idx_tokens - VAL >= 0)[:, None] + & (idx_tokens - VAL < seqlen)[:, None] + & (idx_feats < dim)[None, :] + ) # token-index # token-index # feature-index + loaded_x = tl.load(x_ptrs, mask_x, 0.0) + tl.debug_barrier() + + new_conv_state = tl.where(mask, conv_state, loaded_x) + + # Get the state from the initial_state_idx + # cache_idx + conv_states_offset = tl.load( + conv_state_indices_ptr + idx_seq * stride_state_indices + current_last_index + ).to(tl.int64) + conv_state_ptrs_target = ( + conv_state_ptr + + (conv_states_offset * stride_conv_state_seq) # Offset from seq + + (idx_feats * stride_conv_state_dim) + )[None, :] + ( # [BLOCK_N,] + idx_tokens * stride_conv_state_tok + )[:, None] + mask = (idx_tokens < state_len)[:, None] & (idx_feats < dim)[None, :] + tl.store(conv_state_ptrs_target, new_conv_state, mask) + + # STEP 3: init accumulator + if HAS_BIAS: + bias = bias_ptr + idx_feats + mask_bias = idx_feats < dim + acc_preload = tl.load(bias, mask=mask_bias, other=0.0).to( + tl.float32 + ) # [BLOCK_N] + else: + acc_preload = tl.zeros((BLOCK_N,), dtype=tl.float32) + + # STEP 4: + # PRE-LOAD WEIGHTS + # first kernel column, configured for weights to handle BLOCK_N features in range + w_base = w_ptr + (idx_feats * stride_w_dim) # [BLOCK_N,] + mask_w = idx_feats < dim + if KERNEL_WIDTH >= 2: + w_ptrs = w_base + (0 * stride_w_width) # [BLOCK_N] tensor + w_col0 = tl.load(w_ptrs, mask_w, other=0.0) + w_ptrs = w_base + (1 * stride_w_width) # [BLOCK_N] tensor + w_col1 = tl.load(w_ptrs, mask_w, other=0.0) + if KERNEL_WIDTH >= 3: + w_ptrs = w_base + (2 * stride_w_width) # [BLOCK_N] tensor + w_col2 = tl.load(w_ptrs, mask_w, other=0.0) + if KERNEL_WIDTH >= 4: + w_ptrs = w_base + (3 * stride_w_width) # [BLOCK_N] tensor + w_col3 = tl.load(w_ptrs, mask_w, other=0.0) + if KERNEL_WIDTH >= 5: + w_ptrs = w_base + (4 * stride_w_width) # [BLOCK_N] tensor + w_col4 = tl.load(w_ptrs, mask_w, other=0.0) + if KERNEL_WIDTH >= 6: + w_ptrs = w_base + (5 * stride_w_width) # [BLOCK_N] tensor + w_col5 = tl.load(w_ptrs, mask_w, other=0.0) + + x_base_1d = x_base # starting of chunk [BLOCK_N] + mask_x_1d = idx_feats < dim + + # STEP 5: compute each token + for idx_token in tl.range(seqlen): + acc = acc_preload + + matrix_w = w_col0 + matrix_x = col0 + for j in tl.static_range(KERNEL_WIDTH): + if KERNEL_WIDTH == 2: + if j == 1: # KERNEL_WIDTH-1: + matrix_w = w_col1 + x_ptrs_1d = x_base_1d + idx_token * stride_x_token # [BLOCK_N] + matrix_x = tl.load(x_ptrs_1d, mask=mask_x_1d) + elif KERNEL_WIDTH == 3: + if j == 1: + matrix_w = w_col1 + matrix_x = col1 + elif j == 2: + matrix_w = w_col2 + x_ptrs_1d = x_base_1d + idx_token * stride_x_token # [BLOCK_N] + matrix_x = tl.load(x_ptrs_1d, mask=mask_x_1d) + elif KERNEL_WIDTH == 4: + if j == 1: + matrix_w = w_col1 + matrix_x = col1 + elif j == 2: + matrix_w = w_col2 + matrix_x = col2 + elif j == 3: + matrix_w = w_col3 + x_ptrs_1d = x_base_1d + idx_token * stride_x_token # [BLOCK_N] + matrix_x = tl.load(x_ptrs_1d, mask=mask_x_1d) + elif KERNEL_WIDTH == 5: + if j == 1: + matrix_w = w_col1 + matrix_x = col1 + elif j == 2: + matrix_w = w_col2 + matrix_x = col2 + elif j == 3: + matrix_w = w_col3 + matrix_x = col3 + elif j == 4: + matrix_w = w_col4 + x_ptrs_1d = x_base_1d + idx_token * stride_x_token # [BLOCK_N] + matrix_x = tl.load(x_ptrs_1d, mask=mask_x_1d) + elif KERNEL_WIDTH == 6: + if j == 1: + matrix_w = w_col1 + matrix_x = col1 + elif j == 2: + matrix_w = w_col2 + matrix_x = col2 + elif j == 3: + matrix_w = w_col3 + matrix_x = col3 + elif j == 4: + matrix_w = w_col4 + matrix_x = col4 + elif j == 5: + matrix_w = w_col5 + x_ptrs_1d = x_base_1d + idx_token * stride_x_token # [BLOCK_N] + matrix_x = tl.load(x_ptrs_1d, mask=mask_x_1d) + + acc += matrix_x * matrix_w # [BLOCK_N] + + if KERNEL_WIDTH == 2: + col0 = matrix_x + elif KERNEL_WIDTH == 3: + col0 = col1 + col1 = matrix_x + elif KERNEL_WIDTH == 4: + col0 = col1 + col1 = col2 + col2 = matrix_x + elif KERNEL_WIDTH == 5: + col0 = col1 + col1 = col2 + col2 = col3 + col3 = matrix_x + elif KERNEL_WIDTH == 6: + col0 = col1 + col1 = col2 + col2 = col3 + col3 = col4 + col4 = matrix_x + + if SILU_ACTIVATION: + acc = acc / (1 + tl.exp(-acc)) + mask_1d = (idx_token < seqlen) & ( + idx_feats < dim + ) # token-index # feature-index + o_ptrs = ( + o_ptr + o_offset + idx_token * stride_o_token + (idx_feats * stride_o_dim) + ) + + tl.store(o_ptrs, acc, mask=mask_1d) + diff --git a/benchmarks/TritonNPUBench/vllm/_chunk_cumsum_fwd_kernel/_chunk_cumsum_fwd_kernel.py b/benchmarks/TritonNPUBench/vllm/_chunk_cumsum_fwd_kernel/_chunk_cumsum_fwd_kernel.py new file mode 100644 index 00000000..d6c55434 --- /dev/null +++ b/benchmarks/TritonNPUBench/vllm/_chunk_cumsum_fwd_kernel/_chunk_cumsum_fwd_kernel.py @@ -0,0 +1,104 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/vllm_operator_cases/newtest_cases/test__chunk_cumsum_fwd_kernel.py +# Main kernel: _chunk_cumsum_fwd_kernel +# PT file: _chunk_cumsum_fwd_kernel_v2.pt + +import triton +import triton.language as tl + +# === softplus === +@triton.jit +def softplus(dt): + dt = tl.where(dt <= 20.0, tl.math.log(tl.math.exp(dt) + 1), dt) + return dt + +# === _chunk_cumsum_fwd_kernel === +@triton.jit +def _chunk_cumsum_fwd_kernel( + # Pointers to matrices + dt_ptr, + A_ptr, + dt_bias_ptr, + dt_out_ptr, + dA_cumsum_ptr, + cu_chunk_seqlens_ptr, + # Matrix dimension + seqlen, + nheads: tl.constexpr, + chunk_size: tl.constexpr, + dt_min: tl.constexpr, + dt_max: tl.constexpr, + # Strides + stride_dt_seqlen: tl.int64, + stride_dt_head: tl.constexpr, + stride_A_head: tl.constexpr, + stride_dt_bias_head: tl.constexpr, + stride_dt_out_head: tl.int64, + stride_dt_out_chunk: tl.int64, + stride_dt_out_csize: tl.constexpr, + stride_dA_cs_head: tl.int64, + stride_dA_cs_chunk: tl.int64, + stride_dA_cs_csize: tl.constexpr, + # Meta-parameters + DT_SOFTPLUS: tl.constexpr, + HAS_DT_BIAS: tl.constexpr, + BLOCK_SIZE_H: tl.constexpr, + BLOCK_SIZE_CHUNK: tl.constexpr, +): + # if dt is long, may cause problems, so use 64 bit + # https://github.com/triton-lang/triton/issues/1058 + pid_c = tl.program_id(axis=0).to(tl.int64) + pid_h = tl.program_id(axis=1) + + chunk_seqlen_start = tl.load(cu_chunk_seqlens_ptr + pid_c) + chunk_seqlen_end = tl.load(cu_chunk_seqlens_ptr + pid_c + 1) + + dt_ptr += chunk_seqlen_start * stride_dt_seqlen + dt_out_ptr += pid_c * stride_dt_out_chunk + dA_cumsum_ptr += pid_c * stride_dA_cs_chunk + + offs_h = pid_h * BLOCK_SIZE_H + tl.arange(0, BLOCK_SIZE_H) + offs_c = tl.arange(0, BLOCK_SIZE_CHUNK) + dt_ptrs = dt_ptr + ( + offs_h[:, None] * stride_dt_head + offs_c[None, :] * stride_dt_seqlen + ) + A_ptrs = A_ptr + offs_h * stride_A_head + dt_out_ptrs = dt_out_ptr + ( + offs_h[:, None] * stride_dt_out_head + offs_c[None, :] * stride_dt_out_csize + ) + dA_cs_ptrs = dA_cumsum_ptr + ( + offs_h[:, None] * stride_dA_cs_head + offs_c[None, :] * stride_dA_cs_csize + ) + chunk_size_limit = chunk_seqlen_end - chunk_seqlen_start + + dt = tl.load( + dt_ptrs, + mask=(offs_h[:, None] < nheads) & (offs_c[None, :] < chunk_size_limit), + other=0.0, + ).to(tl.float32) + if HAS_DT_BIAS: + dt_bias = tl.load( + dt_bias_ptr + offs_h * stride_dt_bias_head, mask=offs_h < nheads, other=0.0 + ).to(tl.float32) + dt += dt_bias[:, None] + if DT_SOFTPLUS: + dt = tl.where(dt <= 20.0, softplus(dt), dt) + + dt = tl.clamp(dt, dt_min, dt_max) + dt = tl.where( + (offs_h[:, None] < nheads) & (offs_c[None, :] < chunk_size_limit), dt, 0.0 + ) + tl.store( + dt_out_ptrs, + dt, + mask=(offs_h[:, None] < nheads) & (offs_c[None, :] < chunk_size), + ) + A = tl.load(A_ptrs, mask=offs_h < nheads, other=0.0).to(tl.float32) + dA = dt * A[:, None] + dA_cs = tl.cumsum(dA, axis=1) + tl.store( + dA_cs_ptrs, + dA_cs, + mask=(offs_h[:, None] < nheads) & (offs_c[None, :] < chunk_size), + ) + diff --git a/benchmarks/TritonNPUBench/vllm/_chunk_scan_fwd_kernel/_chunk_scan_fwd_kernel.py b/benchmarks/TritonNPUBench/vllm/_chunk_scan_fwd_kernel/_chunk_scan_fwd_kernel.py new file mode 100644 index 00000000..652972ff --- /dev/null +++ b/benchmarks/TritonNPUBench/vllm/_chunk_scan_fwd_kernel/_chunk_scan_fwd_kernel.py @@ -0,0 +1,281 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/vllm_operator_cases/newtest_cases/test__chunk_scan_fwd_kernel.py +# Main kernel: _chunk_scan_fwd_kernel +# PT file: _chunk_scan_fwd_kernel_unit_v2.pt + +import triton +import triton.language as tl + + +# === _chunk_scan_fwd_kernel === +@triton.jit +def _chunk_scan_fwd_kernel( + # Pointers to matrices + cb_ptr, + x_ptr, + z_ptr, + out_ptr, + dt_ptr, + dA_cumsum_ptr, + seq_idx_ptr, + C_ptr, + states_ptr, + D_ptr, + initstates_ptr, + cu_chunk_seqlens_ptr, + # Matrix dimensions + chunk_size: tl.constexpr, + hdim: tl.constexpr, + dstate: tl.constexpr, + seqlen, + nheads_ngroups_ratio: tl.constexpr, + # Strides + stride_cb_chunk: tl.int64, + stride_cb_head: tl.int64, + stride_cb_csize_m: tl.int64, + stride_cb_csize_k: tl.constexpr, + stride_x_seqlen: tl.int64, + stride_x_head: tl.int64, + stride_x_hdim: tl.constexpr, + stride_z_seqlen: tl.int64, + stride_z_head: tl.int64, + stride_z_hdim: tl.constexpr, + stride_out_seqlen: tl.int64, + stride_out_head: tl.int64, + stride_out_hdim: tl.constexpr, + stride_dt_chunk: tl.int64, + stride_dt_head: tl.int64, + stride_dt_csize: tl.constexpr, + stride_dA_cs_chunk: tl.int64, + stride_dA_cs_head: tl.int64, + stride_dA_cs_csize: tl.constexpr, + stride_seq_idx_chunk: tl.constexpr, + stride_C_seqlen: tl.int64, + stride_C_head: tl.int64, + stride_C_dstate: tl.constexpr, + stride_states_chunk: tl.int64, + stride_states_head: tl.int64, + stride_states_hdim: tl.int64, + stride_states_dstate: tl.constexpr, + stride_init_states_batch: tl.int64, + stride_init_states_head: tl.int64, + stride_init_states_hdim: tl.int64, + stride_init_states_dstate: tl.constexpr, + stride_D_head: tl.constexpr, + # Meta-parameters + IS_CAUSAL: tl.constexpr, + HAS_D: tl.constexpr, + D_HAS_HDIM: tl.constexpr, + HAS_Z: tl.constexpr, + BLOCK_SIZE_M: tl.constexpr, + BLOCK_SIZE_N: tl.constexpr, + BLOCK_SIZE_K: tl.constexpr, + BLOCK_SIZE_DSTATE: tl.constexpr, + IS_TRITON_22: tl.constexpr, + HAS_INITSTATES: tl.constexpr, +): + pid_c = tl.program_id(axis=1).to(tl.int64) + pid_h = tl.program_id(axis=2) + num_pid_n = tl.cdiv(hdim, BLOCK_SIZE_N) + pid_m = tl.program_id(axis=0) // num_pid_n + pid_n = tl.program_id(axis=0) % num_pid_n + cb_ptr += pid_c * stride_cb_chunk + (pid_h // nheads_ngroups_ratio) * stride_cb_head + chunk_seqlen_start = tl.load(cu_chunk_seqlens_ptr + pid_c) + chunk_seqlen_end = tl.load(cu_chunk_seqlens_ptr + pid_c + 1) + x_ptr += chunk_seqlen_start * stride_x_seqlen + pid_h * stride_x_head + dt_ptr += pid_c * stride_dt_chunk + pid_h * stride_dt_head + dA_cumsum_ptr += pid_c * stride_dA_cs_chunk + pid_h * stride_dA_cs_head + C_ptr += ( + chunk_seqlen_start * stride_C_seqlen + + (pid_h // nheads_ngroups_ratio) * stride_C_head + ) + + # M-block offsets and prev states + # - logic in next block may override these if there is an active offset + offs_m = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) + + seq_idx_ptr += pid_c * stride_seq_idx_chunk + seq_idx = tl.load(seq_idx_ptr) + seq_idx_prev = tl.load( + seq_idx_ptr - stride_seq_idx_chunk, mask=pid_c >= 1, other=-1 + ) + + if HAS_INITSTATES and (seq_idx != seq_idx_prev): + prev_states_ptr = ( + initstates_ptr + + seq_idx * stride_init_states_batch + + pid_h * stride_init_states_head + ) + prev_states_hdim = stride_init_states_hdim + prev_states_dstate = stride_init_states_dstate + else: + prev_states_ptr = ( + states_ptr + (pid_c - 1) * stride_states_chunk + pid_h * stride_states_head + ) + prev_states_hdim = stride_states_hdim + prev_states_dstate = stride_states_dstate + + chunk_size_limit = chunk_seqlen_end - chunk_seqlen_start + + offs_n = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) + dA_cs_m = tl.load( + dA_cumsum_ptr + offs_m * stride_dA_cs_csize, mask=offs_m < chunk_size, other=0.0 + ).to(tl.float32) + + acc = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32) + + offs_out_m = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) + offs_out_n = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) + + # Faster to just do 1 iteration with larger BLOCK_SIZE_K, up to block size 128 + offs_k_dstate = tl.arange( + 0, BLOCK_SIZE_DSTATE if BLOCK_SIZE_DSTATE <= 128 else BLOCK_SIZE_K + ) + C_ptrs = C_ptr + ( + offs_m[:, None] * stride_C_seqlen + offs_k_dstate[None, :] * stride_C_dstate + ) + + scale_m = tl.exp(dA_cs_m) + if BLOCK_SIZE_DSTATE <= 128: + C = tl.load( + C_ptrs, + mask=(offs_m[:, None] < chunk_size_limit) + & (offs_k_dstate[None, :] < dstate), + other=0.0, + ) + + if not HAS_INITSTATES and (seq_idx != seq_idx_prev): + # if no init states AND starting a new sequence, we need zeros + prev_states = tl.zeros( + (BLOCK_SIZE_DSTATE, BLOCK_SIZE_N), dtype=C_ptr.dtype.element_ty + ) + else: + # otherwise read the previous state + prev_states_ptrs = ( + prev_states_ptr + + offs_n[None, :] * prev_states_hdim + + offs_k_dstate[:, None] * prev_states_dstate + ) + prev_states = tl.load( + prev_states_ptrs, + mask=(offs_k_dstate[:, None] < dstate) & (offs_n[None, :] < hdim), + other=0.0, + ) + prev_states = prev_states.to(C_ptr.dtype.element_ty) + + acc = tl.dot(C, prev_states) * scale_m[:, None] + + else: + prev_states_ptrs = ( + prev_states_ptr + + offs_n[None, :] * prev_states_hdim + + offs_k_dstate[:, None] * prev_states_dstate + ) + for k in range(0, dstate, BLOCK_SIZE_K): + C = tl.load( + C_ptrs, + mask=(offs_m[:, None] < chunk_size_limit) + & (offs_k_dstate[None, :] < dstate - k), + other=0.0, + ) + if not HAS_INITSTATES and (seq_idx != seq_idx_prev): + prev_states = tl.zeros( + (BLOCK_SIZE_K, BLOCK_SIZE_N), dtype=C_ptr.dtype.element_ty + ) + else: + prev_states = tl.load( + prev_states_ptrs, + mask=(offs_k_dstate[:, None] < dstate - k) + & (offs_n[None, :] < hdim), + other=0.0, + ) + prev_states = prev_states.to(C_ptr.dtype.element_ty) + acc += tl.dot(C, prev_states) + C_ptrs += BLOCK_SIZE_K + prev_states_ptrs += BLOCK_SIZE_K + acc *= scale_m[:, None] + + offs_k = tl.arange(0, BLOCK_SIZE_K) + cb_ptrs = cb_ptr + ( + offs_m[:, None] * stride_cb_csize_m + offs_k[None, :] * stride_cb_csize_k + ) + x_ptrs = x_ptr + ( + offs_k[:, None] * stride_x_seqlen + offs_n[None, :] * stride_x_hdim + ) + dt_ptrs = dt_ptr + offs_k * stride_dt_csize + dA_cumsum_ptrs = dA_cumsum_ptr + offs_k * stride_dA_cs_csize + K_MAX = ( + chunk_size_limit + if not IS_CAUSAL + else min((pid_m + 1) * BLOCK_SIZE_M, chunk_size_limit) + ) + for k in range(0, K_MAX, BLOCK_SIZE_K): + cb = tl.load( + cb_ptrs, + mask=(offs_m[:, None] < chunk_size) & (offs_k[None, :] < chunk_size - k), + other=0.0, + ).to(tl.float32) + dA_cs_k = tl.load(dA_cumsum_ptrs, mask=offs_k < chunk_size - k, other=0.0).to( + tl.float32 + ) + # If there's seq_idx, we already set cb[i, j] = 0 for seq_idx[i] != seq_idx[j]. + # So we don't need masking wrt seq_idx here. + cb *= tl.exp(dA_cs_m[:, None] - dA_cs_k[None, :]) + dt_k = tl.load(dt_ptrs, mask=offs_k < chunk_size - k, other=0.0).to(tl.float32) + cb *= dt_k + if IS_CAUSAL: + mask = offs_m[:, None] >= k + offs_k[None, :] + cb = tl.where(mask, cb, 0.0) + cb = cb.to(x_ptr.dtype.element_ty) + x = tl.load( + x_ptrs, + mask=(offs_k[:, None] < chunk_size_limit - k) & (offs_n[None, :] < hdim), + other=0.0, + ) + acc += tl.dot(cb, x) + cb_ptrs += BLOCK_SIZE_K * stride_cb_csize_k + x_ptrs += BLOCK_SIZE_K * stride_x_seqlen + dt_ptrs += BLOCK_SIZE_K * stride_dt_csize + dA_cumsum_ptrs += BLOCK_SIZE_K * stride_dA_cs_csize + + offs_out_m = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) + offs_out_n = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) + + if HAS_D: + if D_HAS_HDIM: + D = tl.load( + D_ptr + pid_h * stride_D_head + offs_n, mask=offs_n < hdim, other=0.0 + ).to(tl.float32) + else: + D = tl.load(D_ptr + pid_h * stride_D_head).to(tl.float32) + x_residual = tl.load( + x_ptr + + (offs_m[:, None] * stride_x_seqlen + offs_n[None, :] * stride_x_hdim), + mask=(offs_m[:, None] < chunk_size_limit) & (offs_n[None, :] < hdim), + other=0.0, + ).to(tl.float32) + acc += x_residual * D + + if HAS_Z: + z_ptr += chunk_seqlen_start * stride_z_seqlen + pid_h * stride_z_head + z_ptrs = z_ptr + ( + stride_z_seqlen * offs_out_m[:, None] + stride_z_hdim * offs_out_n[None, :] + ) + z = tl.load( + z_ptrs, + mask=(offs_out_m[:, None] < chunk_size_limit) + & (offs_out_n[None, :] < hdim), + other=0.0, + ).to(tl.float32) + acc *= z * tl.sigmoid(z) + + out_ptr += chunk_seqlen_start * stride_out_seqlen + pid_h * stride_out_head + out_ptrs = out_ptr + ( + stride_out_seqlen * offs_out_m[:, None] + offs_out_n[None, :] * stride_out_hdim + ) + tl.store( + out_ptrs, + acc, + mask=(offs_out_m[:, None] < chunk_size_limit) & (offs_out_n[None, :] < hdim), + ) + diff --git a/benchmarks/TritonNPUBench/vllm/_chunk_state_fwd_kernel/_chunk_state_fwd_kernel.py b/benchmarks/TritonNPUBench/vllm/_chunk_state_fwd_kernel/_chunk_state_fwd_kernel.py new file mode 100644 index 00000000..a76d203c --- /dev/null +++ b/benchmarks/TritonNPUBench/vllm/_chunk_state_fwd_kernel/_chunk_state_fwd_kernel.py @@ -0,0 +1,118 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/vllm_operator_cases/newtest_cases/test__chunk_state_fwd_kernel.py +# Main kernel: _chunk_state_fwd_kernel +# PT file: _chunk_state_fwd_kernel_v3.pt + +import triton +import triton.language as tl + + +# === _chunk_state_fwd_kernel === +@triton.jit +def _chunk_state_fwd_kernel( + # Pointers to matrices + x_ptr, + b_ptr, + states_ptr, + dt_ptr, + dA_cumsum_ptr, + cu_chunk_seqlens_ptr, + # Matrix dimensions + hdim: tl.constexpr, + dstate: tl.constexpr, + chunk_size: tl.constexpr, + seqlen, + nheads_ngroups_ratio: tl.constexpr, + # Strides + stride_x_seqlen: tl.int64, + stride_x_head: tl.int64, + stride_x_hdim: tl.constexpr, + stride_b_seqlen: tl.int64, + stride_b_head: tl.int64, + stride_b_dstate: tl.constexpr, + stride_states_chunk: tl.int64, + stride_states_head: tl.int64, + stride_states_hdim: tl.int64, + stride_states_dstate: tl.constexpr, + stride_dt_head: tl.int64, + stride_dt_chunk: tl.int64, + stride_dt_csize: tl.constexpr, + stride_dA_cs_head: tl.int64, + stride_dA_cs_chunk: tl.int64, + stride_dA_cs_csize: tl.constexpr, + # Meta-parameters + BLOCK_SIZE_M: tl.constexpr, + BLOCK_SIZE_N: tl.constexpr, + BLOCK_SIZE_K: tl.constexpr, +): + pid_c = tl.program_id(axis=1).to(tl.int64) + pid_h = tl.program_id(axis=2) + num_pid_n = tl.cdiv(dstate, BLOCK_SIZE_N) + pid_m = tl.program_id(axis=0) // num_pid_n + pid_n = tl.program_id(axis=0) % num_pid_n + chunk_seqlen_start = tl.load(cu_chunk_seqlens_ptr + pid_c) + chunk_seqlen_end = tl.load(cu_chunk_seqlens_ptr + pid_c + 1) + b_ptr += ( + chunk_seqlen_start * stride_b_seqlen + + (pid_h // nheads_ngroups_ratio) * stride_b_head + ) + x_ptr += chunk_seqlen_start * stride_x_seqlen + pid_h * stride_x_head + dt_ptr += pid_c * stride_dt_chunk + pid_h * stride_dt_head + dA_cumsum_ptr += pid_c * stride_dA_cs_chunk + pid_h * stride_dA_cs_head + + offs_m = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) + offs_n = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) + offs_k = tl.arange(0, BLOCK_SIZE_K) + x_ptrs = x_ptr + ( + offs_m[:, None] * stride_x_hdim + offs_k[None, :] * stride_x_seqlen + ) + b_ptrs = b_ptr + ( + offs_n[None, :] * stride_b_dstate + offs_k[:, None] * stride_b_seqlen + ) + dt_ptrs = dt_ptr + offs_k * stride_dt_csize + dA_cs_last = tl.load(dA_cumsum_ptr + (chunk_size - 1) * stride_dA_cs_csize).to( + tl.float32 + ) + dA_cumsum_ptrs = dA_cumsum_ptr + offs_k * stride_dA_cs_csize + + chunk_size_limit = chunk_seqlen_end - chunk_seqlen_start + + acc = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32) + for k in range(0, chunk_size_limit, BLOCK_SIZE_K): + x = tl.load( + x_ptrs, + mask=(offs_m[:, None] < hdim) & (offs_k[None, :] < chunk_size_limit - k), + other=0.0, + ) + b = tl.load( + b_ptrs, + mask=(offs_k[:, None] < chunk_size_limit - k) & (offs_n[None, :] < dstate), + other=0.0, + ).to(tl.float32) + dA_cs_k = tl.load( + dA_cumsum_ptrs, mask=offs_k < chunk_size_limit - k, other=0.0 + ).to(tl.float32) + dt_k = tl.load(dt_ptrs, mask=offs_k < chunk_size_limit - k, other=0.0).to( + tl.float32 + ) + scale = tl.exp(dA_cs_last - dA_cs_k) * dt_k + b *= scale[:, None] + b = b.to(x_ptr.dtype.element_ty) + acc += tl.dot(x, b) + + x_ptrs += BLOCK_SIZE_K * stride_x_seqlen + b_ptrs += BLOCK_SIZE_K * stride_b_seqlen + dt_ptrs += BLOCK_SIZE_K * stride_dt_csize + dA_cumsum_ptrs += BLOCK_SIZE_K * stride_dA_cs_csize + + states = acc.to(states_ptr.dtype.element_ty) + + states_ptr += pid_c * stride_states_chunk + pid_h * stride_states_head + offs_m = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) + offs_n = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) + states_ptrs = states_ptr + ( + offs_m[:, None] * stride_states_hdim + offs_n[None, :] * stride_states_dstate + ) + c_mask = (offs_m[:, None] < hdim) & (offs_n[None, :] < dstate) + tl.store(states_ptrs, states, mask=c_mask) + diff --git a/benchmarks/TritonNPUBench/vllm/_chunk_state_varlen_kernel/_chunk_state_varlen_kernel.py b/benchmarks/TritonNPUBench/vllm/_chunk_state_varlen_kernel/_chunk_state_varlen_kernel.py new file mode 100644 index 00000000..a0147b24 --- /dev/null +++ b/benchmarks/TritonNPUBench/vllm/_chunk_state_varlen_kernel/_chunk_state_varlen_kernel.py @@ -0,0 +1,192 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/vllm_operator_cases/newtest_cases/test__chunk_state_varlen_kernel.py +# Main kernel: _chunk_state_varlen_kernel +# PT file: _chunk_state_varlen_kernel_v3.pt + +import triton +import triton.language as tl + + +# === _chunk_state_varlen_kernel === +@triton.jit +def _chunk_state_varlen_kernel( + # Pointers to matrices + x_ptr, + b_ptr, + dt_ptr, + dA_cumsum_ptr, + chunk_states_ptr, + cu_seqlens_ptr, + states_ptr, + initstates_ptr, + # Matrix dimensions + hdim: tl.constexpr, + dstate: tl.constexpr, + chunk_size: tl.constexpr, + nheads_ngroups_ratio: tl.constexpr, + # Strides + stride_x_seqlen: tl.int64, + stride_x_head: tl.int64, + stride_x_hdim: tl.constexpr, + stride_b_seqlen: tl.int64, + stride_b_head: tl.int64, + stride_b_dstate: tl.constexpr, + stride_dt_head: tl.int64, + stride_dt_chunk: tl.int64, + stride_dt_csize: tl.constexpr, + stride_dA_cs_head: tl.int64, + stride_dA_cs_chunk: tl.int64, + stride_dA_cs_csize: tl.constexpr, + stride_chunk_states_chunk: tl.int64, + stride_chunk_states_head: tl.int64, + stride_chunk_states_hdim: tl.int64, + stride_chunk_states_dstate: tl.constexpr, + stride_states_batch: tl.int64, + stride_states_head: tl.int64, + stride_states_hdim: tl.int64, + stride_states_dstate: tl.constexpr, + stride_init_states_batch: tl.int64, + stride_init_states_head: tl.int64, + stride_init_states_hdim: tl.int64, + stride_init_states_dstate: tl.constexpr, + # Meta-parameters + BLOCK_SIZE_M: tl.constexpr, + BLOCK_SIZE_N: tl.constexpr, + BLOCK_SIZE_K: tl.constexpr, + HAS_INITSTATES: tl.constexpr, +): + pid_b = tl.program_id(axis=1) + pid_h = tl.program_id(axis=2) + num_pid_n = tl.cdiv(dstate, BLOCK_SIZE_N) + pid_m = tl.program_id(axis=0) // num_pid_n + pid_n = tl.program_id(axis=0) % num_pid_n + end_idx = tl.load(cu_seqlens_ptr + pid_b + 1) + pid_c = (end_idx - 1) // chunk_size + b_ptr += ( + pid_c * chunk_size * stride_b_seqlen + + (pid_h // nheads_ngroups_ratio) * stride_b_head + ) + x_ptr += pid_c * chunk_size * stride_x_seqlen + pid_h * stride_x_head + dt_ptr += pid_c * stride_dt_chunk + pid_h * stride_dt_head + dA_cumsum_ptr += pid_c * stride_dA_cs_chunk + pid_h * stride_dA_cs_head + chunk_states_ptr += ( + pid_c * stride_chunk_states_chunk + pid_h * stride_chunk_states_head + ) + + if HAS_INITSTATES: + # if there are init states provided, we differentiate between states (which + # are boundary conditions at a chunk boundary) and initstates (which are boundary + # conditions when a new example in a cont batch starts) + initstates_ptr += pid_h * stride_init_states_head + + offs_m = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) + offs_n = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) + offs_k = tl.arange(0, BLOCK_SIZE_K) + x_ptrs = x_ptr + ( + offs_m[:, None] * stride_x_hdim + offs_k[None, :] * stride_x_seqlen + ) + b_ptrs = b_ptr + ( + offs_n[None, :] * stride_b_dstate + offs_k[:, None] * stride_b_seqlen + ) + dt_ptrs = dt_ptr + offs_k * stride_dt_csize + dA_cs_last = tl.load( + dA_cumsum_ptr + (end_idx - pid_c * chunk_size - 1) * stride_dA_cs_csize + ).to(tl.float32) + dA_cumsum_ptrs = dA_cumsum_ptr + offs_k * stride_dA_cs_csize + + chunk_size_limit = end_idx - pid_c * chunk_size + start_idx = tl.load(cu_seqlens_ptr + pid_b) + start_idx_cur = tl.maximum(start_idx - pid_c * chunk_size, 0) + + acc = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32) + for k in range(0, chunk_size_limit, BLOCK_SIZE_K): + x = tl.load( + x_ptrs, + mask=(offs_m[:, None] < hdim) + & (offs_k[None, :] < chunk_size_limit - k) + & (offs_k[None, :] >= start_idx_cur - k), + other=0.0, + ) + b = tl.load( + b_ptrs, + mask=(offs_k[:, None] < chunk_size_limit - k) + & (offs_n[None, :] < dstate) + & (offs_k[:, None] >= start_idx_cur - k), + other=0.0, + ).to(tl.float32) + dA_cs_k = tl.load( + dA_cumsum_ptrs, mask=offs_k < chunk_size_limit - k, other=0.0 + ).to(tl.float32) + dt_k = tl.load(dt_ptrs, mask=offs_k < chunk_size_limit - k, other=0.0).to( + tl.float32 + ) + scale = tl.where( + (offs_k >= start_idx_cur - k) & (offs_k < chunk_size_limit - k), + tl.exp(dA_cs_last - dA_cs_k) * dt_k, + 0.0, + ) + b *= scale[:, None] + b = b.to(x_ptr.dtype.element_ty) + acc += tl.dot(x, b) + x_ptrs += BLOCK_SIZE_K * stride_x_seqlen + b_ptrs += BLOCK_SIZE_K * stride_b_seqlen + dt_ptrs += BLOCK_SIZE_K * stride_dt_csize + dA_cumsum_ptrs += BLOCK_SIZE_K * stride_dA_cs_csize + + # If the sequence starts after the last chunk idx, we don't need to add the contribution from the last chunk + # If HAS_INITSTATES==True need to consider two possibilities + # - if start_idx < pid_c * chunk_size, then we need to take the past_states_ptrs + # - if state_idx >= pid * chunk_size, then we need to insert initstates + if ( + (start_idx < pid_c * chunk_size) # first chunk + or (HAS_INITSTATES) + ): + dA_cs_boundary = 0.0 # default + + ptrs_chunk = chunk_states_ptr + ( + offs_m[:, None] * stride_chunk_states_hdim + + offs_n[None, :] * stride_chunk_states_dstate + ) + + ptrs_init = initstates_ptr + ( + pid_b * stride_init_states_batch + + offs_m[:, None] * stride_init_states_hdim + + offs_n[None, :] * stride_init_states_dstate + ) + + if HAS_INITSTATES: + c_mask_chunk = (offs_m[:, None] < hdim) & (offs_n[None, :] < dstate) + c_mask_init = (offs_m[:, None] < hdim) & (offs_n[None, :] < dstate) + + state_chunk = tl.load(ptrs_chunk, mask=c_mask_chunk, other=0.0).to(tl.float32) + state_init = tl.load(ptrs_init, mask=c_mask_init, other=0.0).to(tl.float32) + + if start_idx > pid_c * chunk_size: + dA_cs_boundary = tl.load( + dA_cumsum_ptr + + (start_idx - pid_c * chunk_size - 1) * stride_dA_cs_csize + ).to(tl.float32) + + past_states = tl.where( + start_idx[:, None] < pid_c * chunk_size, + state_chunk, + state_init + ) + else: + c_mask = (offs_m[:, None] < hdim) & (offs_n[None, :] < dstate) + past_states = tl.load(ptrs_chunk, mask=c_mask, other=0.0) + + scale = tl.exp(dA_cs_last - dA_cs_boundary) + acc += past_states * scale + + states = acc.to(states_ptr.dtype.element_ty) + + states_ptr += pid_b * stride_states_batch + pid_h * stride_states_head + offs_m = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) + offs_n = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) + states_ptrs = states_ptr + ( + offs_m[:, None] * stride_states_hdim + offs_n[None, :] * stride_states_dstate + ) + c_mask = (offs_m[:, None] < hdim) & (offs_n[None, :] < dstate) + tl.store(states_ptrs, states, mask=c_mask) + diff --git a/benchmarks/TritonNPUBench/vllm/_combine_sampled_and_draft_tokens_kernel/_combine_sampled_and_draft_tokens_kernel.py b/benchmarks/TritonNPUBench/vllm/_combine_sampled_and_draft_tokens_kernel/_combine_sampled_and_draft_tokens_kernel.py new file mode 100644 index 00000000..0915d763 --- /dev/null +++ b/benchmarks/TritonNPUBench/vllm/_combine_sampled_and_draft_tokens_kernel/_combine_sampled_and_draft_tokens_kernel.py @@ -0,0 +1,66 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/vllm_operator_cases/newtest_cases/test__combine_sampled_and_draft_tokens_kernel.py +# Main kernel: _combine_sampled_and_draft_tokens_kernel +# PT file: combine_sampled_and_draft_tokens_test_v2.pt + +import triton +import triton.language as tl + + +# === _combine_sampled_and_draft_tokens_kernel === +@triton.jit +def _combine_sampled_and_draft_tokens_kernel( + input_ids_ptr, + idx_mapping_ptr, + last_sampled_tokens_ptr, + query_start_loc_ptr, + seq_lens_ptr, + prefill_len_ptr, + draft_tokens_ptr, + draft_tokens_stride, + cu_num_logits_ptr, + logits_indices_ptr, + BLOCK_SIZE: tl.constexpr, +): + batch_idx = tl.program_id(0) + req_state_idx = tl.load(idx_mapping_ptr + batch_idx) + + # Get the number of logits and draft tokens. + cu_num_logits_start = tl.load(cu_num_logits_ptr + batch_idx) + cu_num_logits_end = tl.load(cu_num_logits_ptr + batch_idx + 1) + num_logits = cu_num_logits_end - cu_num_logits_start + num_draft_tokens = num_logits - 1 + + # Compute the logits indices. + block = tl.arange(0, BLOCK_SIZE) + query_end = tl.load(query_start_loc_ptr + batch_idx + 1) + logits_start = query_end - num_logits + tl.store( + logits_indices_ptr + cu_num_logits_start + block, + logits_start + block, + mask=block < num_logits, + ) + + seq_len = tl.load(seq_lens_ptr + batch_idx) + prefill_len = tl.load(prefill_len_ptr + req_state_idx) + if seq_len <= prefill_len: + # Handling prefill tokens. No sampled or draft tokens. + return + + # Write the last sampled token ID to input_ids. + last_token_id = tl.load(last_sampled_tokens_ptr + req_state_idx) + tl.store(input_ids_ptr + query_end - num_logits, last_token_id) + + # Write the draft tokens (if any) to input_ids. + if num_draft_tokens > 0: + mask = block < num_draft_tokens + draft_tokens = tl.load( + draft_tokens_ptr + req_state_idx * draft_tokens_stride + block, + mask=mask, + ) + tl.store( + input_ids_ptr + query_end - num_draft_tokens + block, + draft_tokens, + mask=mask, + ) + diff --git a/benchmarks/TritonNPUBench/vllm/_compute_pid_wrapper/_compute_pid_wrapper.py b/benchmarks/TritonNPUBench/vllm/_compute_pid_wrapper/_compute_pid_wrapper.py new file mode 100644 index 00000000..34d9c295 --- /dev/null +++ b/benchmarks/TritonNPUBench/vllm/_compute_pid_wrapper/_compute_pid_wrapper.py @@ -0,0 +1,38 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/vllm_operator_cases/newtest_cases/test__compute_pid.py +# Main kernel: _compute_pid_wrapper +# PT file: _compute_pid_wrapper_v2.pt + +import triton +import triton.language as tl + +@triton.jit +def _compute_pid(tile_id, num_pid_in_group, num_pid_m, GROUP_SIZE_M, NUM_SMS): + group_id = tile_id // num_pid_in_group + first_pid_m = group_id * GROUP_SIZE_M + group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M) + pid_m = first_pid_m + (tile_id % group_size_m) + pid_n = (tile_id % num_pid_in_group) // group_size_m + return pid_m, pid_n + +# === _compute_pid_wrapper === +@triton.jit +def _compute_pid_wrapper( + tile_id_ptr, + num_pid_in_group_ptr, + num_pid_m_ptr, + GROUP_SIZE_M_ptr, + NUM_SMS_ptr, + output_ptr, +): + if tl.program_id(0) == 0: + tile_id = tl.load(tile_id_ptr) + num_pid_in_group = tl.load(num_pid_in_group_ptr) + num_pid_m = tl.load(num_pid_m_ptr) + GROUP_SIZE_M = tl.load(GROUP_SIZE_M_ptr) + NUM_SMS = tl.load(NUM_SMS_ptr) + + pid_m, pid_n = _compute_pid(tile_id, num_pid_in_group, num_pid_m, GROUP_SIZE_M, NUM_SMS) + tl.store(output_ptr + 0, pid_m) + tl.store(output_ptr + 1, pid_n) + diff --git a/benchmarks/TritonNPUBench/vllm/_compute_slot_mappings_kernel/_compute_slot_mappings_kernel.py b/benchmarks/TritonNPUBench/vllm/_compute_slot_mappings_kernel/_compute_slot_mappings_kernel.py new file mode 100644 index 00000000..48ff6a77 --- /dev/null +++ b/benchmarks/TritonNPUBench/vllm/_compute_slot_mappings_kernel/_compute_slot_mappings_kernel.py @@ -0,0 +1,62 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/vllm_operator_cases/newtest_cases/test__compute_slot_mappings_kernel.py +# Main kernel: _compute_slot_mappings_kernel +# PT file: _compute_slot_mappings_kernel_v2.pt + +import triton +import triton.language as tl + + +# === _compute_slot_mappings_kernel === +@triton.jit +def _compute_slot_mappings_kernel( + num_tokens, + max_num_tokens, + idx_mapping, # [num_reqs] + query_start_loc, # [num_reqs + 1] + pos, # [num_tokens] + block_table_ptrs, # [num_kv_cache_groups] + block_table_strides, # [num_kv_cache_groups] + block_sizes, # [num_kv_cache_groups] + slot_mappings_ptr, # [num_kv_cache_groups, max_num_tokens] + slot_mappings_stride, + PAD_ID: tl.constexpr, + TRITON_BLOCK_SIZE: tl.constexpr, +): + # kv cache group id + group_id = tl.program_id(0) + batch_idx = tl.program_id(1) + slot_mapping_ptr = slot_mappings_ptr + group_id * slot_mappings_stride + + if batch_idx == tl.num_programs(1) - 1: + # Pad remaining slots to -1. This is needed for CUDA graphs. + for i in range(num_tokens, max_num_tokens, TRITON_BLOCK_SIZE): + offset = i + tl.arange(0, TRITON_BLOCK_SIZE) + tl.store(slot_mapping_ptr + offset, PAD_ID, mask=offset < max_num_tokens) + return + + block_table_ptr = _load_ptr(block_table_ptrs + group_id, tl.int32) + block_table_stride = tl.load(block_table_strides + group_id) + block_size = tl.load(block_sizes + group_id) + + req_state_idx = tl.load(idx_mapping + batch_idx) + start_idx = tl.load(query_start_loc + batch_idx) + end_idx = tl.load(query_start_loc + batch_idx + 1) + for i in range(start_idx, end_idx, TRITON_BLOCK_SIZE): + offset = i + tl.arange(0, TRITON_BLOCK_SIZE) + positions = tl.load(pos + offset, mask=offset < end_idx, other=0) + block_indices = positions // block_size + block_numbers = tl.load( + block_table_ptr + req_state_idx * block_table_stride + block_indices + ) + slot_ids = block_numbers * block_size + positions % block_size + tl.store(slot_mapping_ptr + offset, slot_ids, mask=offset < end_idx) + + +# === _load_ptr === +@triton.jit +def _load_ptr(ptr_to_ptr, elem_dtype): + ptr = tl.load(ptr_to_ptr) + ptr = tl.cast(ptr, tl.pointer_type(elem_dtype)) + return tl.multiple_of(ptr, 16) + diff --git a/benchmarks/TritonNPUBench/vllm/_convert_req_index_to_global_index_kernel/_convert_req_index_to_global_index_kernel.py b/benchmarks/TritonNPUBench/vllm/_convert_req_index_to_global_index_kernel/_convert_req_index_to_global_index_kernel.py new file mode 100644 index 00000000..fcab2633 --- /dev/null +++ b/benchmarks/TritonNPUBench/vllm/_convert_req_index_to_global_index_kernel/_convert_req_index_to_global_index_kernel.py @@ -0,0 +1,76 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/vllm_operator_cases/newtest_cases/test__convert_req_index_to_global_index_kernel.py +# Main kernel: _convert_req_index_to_global_index_kernel +# PT file: _convert_req_index_to_global_index_kernel_v2.pt + +import triton +import triton.language as tl + + +# === _convert_req_index_to_global_index_kernel === +@triton.jit +def _convert_req_index_to_global_index_kernel( + req_id_ptr, # int32 [num_tokens] + block_table_ptr, # int32 [num_requests, max_num_blocks_per_req] + token_indices_ptr, # int32 [num_tokens, NUM_TOPK_TOKENS] + out_ptr, # int32 [num_tokens, NUM_TOPK_TOKENS] + prefill_request_id_ptr, # int32 [num_tokens], -1 for decode, >=0 for prefill + workspace_starts_ptr, # int32 [num_prefill_reqs+1] or nullptr + # shapes (compile-time where possible) + max_num_blocks_per_req: tl.constexpr, + BLOCK_SIZE: tl.constexpr, + BLOCK_N: tl.constexpr, # tile width along columns + HAS_PREFILL: tl.constexpr, + # strides (in elements) + bt_stride0, + bt_stride1, + ti_stride0, + ti_stride1, + out_stride0, + out_stride1, +): + # program_id(0) -> token_id (row) + # program_id(1) -> tile index along columns + token_id = tl.program_id(0) + tile_id = tl.program_id(1) + + # Each program covers BLOCK_N consecutive columns + indice_id = tile_id * BLOCK_N + tl.arange(0, BLOCK_N) + + # Load request id for this token (no mask: grid is exact) + req = tl.load(req_id_ptr + token_id) + + # Load token indices for this tile + ti_ptr = token_indices_ptr + token_id * ti_stride0 + indice_id * ti_stride1 + tok = tl.load(ti_ptr) # int32 + + # Only token == -1 should propagate as -1 + is_invalid_tok = tok < 0 + is_prefill = False + if HAS_PREFILL: + prefill_req_id = tl.load(prefill_request_id_ptr + token_id) + is_prefill = prefill_req_id >= 0 + # Compute block id and in-block offset + block_id = tok // BLOCK_SIZE + inblock_off = tok % BLOCK_SIZE + + # Guard block_table access + valid_block = (block_id < max_num_blocks_per_req) & (block_id >= 0) + bt_ptr = block_table_ptr + req * bt_stride0 + block_id * bt_stride1 + is_invalid_tok |= ~valid_block + base = tl.load(bt_ptr, mask=valid_block & ~is_prefill, other=0) + out_val = base * BLOCK_SIZE + inblock_off + + # Override with prefill output if prefill is enabled + if HAS_PREFILL: + workspace_start = tl.load( + workspace_starts_ptr + prefill_req_id, mask=is_prefill, other=0 + ) + prefill_out = workspace_start + tok + out_val = tl.where(is_prefill, prefill_out, out_val) + out_val = tl.where(is_invalid_tok, -1, out_val) + + # Store results + out_ptr_ij = out_ptr + token_id * out_stride0 + indice_id * out_stride1 + tl.store(out_ptr_ij, out_val) + diff --git a/benchmarks/TritonNPUBench/vllm/_copy_page_indices_kernel/_copy_page_indices_kernel.py b/benchmarks/TritonNPUBench/vllm/_copy_page_indices_kernel/_copy_page_indices_kernel.py new file mode 100644 index 00000000..5bc6932c --- /dev/null +++ b/benchmarks/TritonNPUBench/vllm/_copy_page_indices_kernel/_copy_page_indices_kernel.py @@ -0,0 +1,33 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/vllm_operator_cases/newtest_cases/test__copy_page_indices_kernel.py +# Main kernel: _copy_page_indices_kernel +# PT file: _copy_page_indices_kernel_v2.pt + +import triton +import triton.language as tl + + +# === _copy_page_indices_kernel === +@triton.jit +def _copy_page_indices_kernel( + page_indices, + block_table, + block_table_stride, + cu_num_blocks, + BLOCK_SIZE: tl.constexpr, +): + req_idx = tl.program_id(0) + row_ptr = block_table + req_idx * block_table_stride + start_idx = tl.load(cu_num_blocks + req_idx) + end_idx = tl.load(cu_num_blocks + req_idx + 1) + num_blocks = end_idx - start_idx + + offset = tl.arange(0, BLOCK_SIZE) + for i in tl.range(0, num_blocks, BLOCK_SIZE): + block_ids = tl.load(row_ptr + i + offset, mask=i + offset < num_blocks) + tl.store( + page_indices + start_idx + i + offset, + block_ids, + mask=i + offset < num_blocks, + ) + diff --git a/benchmarks/TritonNPUBench/vllm/_correct_attn_cp_out_kernel/_correct_attn_cp_out_kernel.py b/benchmarks/TritonNPUBench/vllm/_correct_attn_cp_out_kernel/_correct_attn_cp_out_kernel.py new file mode 100644 index 00000000..2357a756 --- /dev/null +++ b/benchmarks/TritonNPUBench/vllm/_correct_attn_cp_out_kernel/_correct_attn_cp_out_kernel.py @@ -0,0 +1,97 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/vllm_operator_cases/newtest_cases/test__correct_attn_cp_out_kernel.py +# Main kernel: _correct_attn_cp_out_kernel +# PT file: correct_attn_cp_out_test_v2.pt + +import triton +import triton.language as tl + + +# === _correct_attn_cp_out_kernel === +@triton.jit +def _correct_attn_cp_out_kernel( + outputs_ptr, + new_output_ptr, + lses_ptr, + vlse_ptr, + outputs_stride_B, + outputs_stride_H, + outputs_stride_D, + lses_stride_N, + lses_stride_B, + lses_stride_H, + lse_idx, + HEAD_DIM: tl.constexpr, + N_ROUNDED: tl.constexpr, + IS_BASE_E: tl.constexpr, +): + """ + Apply the all-gathered lses to correct each local rank's attention + output. we still need perform a cross-rank reduction to obtain the + final attention output. + + Args: + outputs_ptr (triton.PointerType): + Pointer to input tensor of shape [ B, H, D ] + lses_ptr (triton.PointerType): + Pointer to input tensor of shape [ N, B, H ] + new_output_ptr (triton.PointerType): + Pointer to output tensor of shape [ B, H, D ] + vlse_ptr (triton.PointerType): + Pointer to output tensor of shape [ B, H ] + """ + batch_idx = tl.program_id(axis=0).to(tl.int64) + head_idx = tl.program_id(axis=1).to(tl.int64) + d_offsets = tl.arange(0, HEAD_DIM) + num_n_offsets = tl.arange(0, N_ROUNDED) + + # shape = [N] + lse_offsets = ( + num_n_offsets * lses_stride_N + + batch_idx * lses_stride_B + + head_idx * lses_stride_H + ) + + # calc final lse + lse = tl.load(lses_ptr + lse_offsets) + lse = tl.where((lse != lse) | (lse == float("inf")), -float("inf"), lse) + lse_max = tl.max(lse, axis=0) + lse_max = tl.where(lse_max == -float("inf"), 0, lse_max) + lse -= lse_max + if IS_BASE_E: + lse_exp = tl.exp(lse) + lse_acc = tl.sum(lse_exp, axis=0) + lse = tl.log(lse_acc) + else: + lse_exp = tl.exp2(lse) + lse_acc = tl.sum(lse_exp, axis=0) + lse = tl.log2(lse_acc) + lse += lse_max + + lse_offsets = batch_idx * lses_stride_B + head_idx * lses_stride_H + tl.store(vlse_ptr + lse_offsets, lse) + + # shape = [D] + output_offsets = ( + batch_idx * outputs_stride_B + + head_idx * outputs_stride_H + + d_offsets * outputs_stride_D + ) + + # correct output + lse_offset = ( + lse_idx * lses_stride_N + batch_idx * lses_stride_B + head_idx * lses_stride_H + ) + lse_tmp = tl.load(lses_ptr + lse_offset) + lse_finally = lse_tmp - lse + lse_finally = tl.where( + (lse_finally != lse_finally) | (lse_finally == float("inf")), + -float("inf"), + lse_finally, + ) + factor = tl.exp(lse_finally) if IS_BASE_E else tl.exp2(lse_finally) + output = tl.load(outputs_ptr + output_offsets) + output = output * factor + + tl.store(new_output_ptr + output_offsets, output) + diff --git a/benchmarks/TritonNPUBench/vllm/_count_expert_num_tokens/_count_expert_num_tokens.py b/benchmarks/TritonNPUBench/vllm/_count_expert_num_tokens/_count_expert_num_tokens.py new file mode 100644 index 00000000..75d54fc6 --- /dev/null +++ b/benchmarks/TritonNPUBench/vllm/_count_expert_num_tokens/_count_expert_num_tokens.py @@ -0,0 +1,41 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/vllm_operator_cases/newtest_cases/test__count_expert_num_tokens.py +# Main kernel: _count_expert_num_tokens +# PT file: _count_expert_num_tokens_v2.pt + +import triton +import triton.language as tl + + +# === _count_expert_num_tokens === +@triton.jit +def _count_expert_num_tokens( + topk_ids_ptr, + expert_num_tokens_ptr, + num_experts, + topk_numel, + expert_map, + HAS_EXPERT_MAP: tl.constexpr, + BLOCK_SIZE: tl.constexpr, +): + curr_expert = tl.program_id(0) + + offsets = tl.arange(0, BLOCK_SIZE) + topk_ids_ptrs = topk_ids_ptr + offsets + + acc = tl.zeros((BLOCK_SIZE,), dtype=tl.int32) + for x in range(tl.cdiv(topk_numel, BLOCK_SIZE)): + mask = offsets < (topk_numel - x * BLOCK_SIZE) + expert_ids = tl.load(topk_ids_ptrs, mask=mask, other=-1) + if HAS_EXPERT_MAP: + expert_map_ptrs = expert_map + expert_ids + expert_map_mask = expert_ids >= 0 + expert_ids = tl.load(expert_map_ptrs, mask=expert_map_mask, other=-1) + + has_curr_expert = tl.where(expert_ids == curr_expert, 1, 0) + acc = acc + has_curr_expert + topk_ids_ptrs += BLOCK_SIZE + + if curr_expert < num_experts: + tl.store(expert_num_tokens_ptr + curr_expert, tl.sum(acc)) + diff --git a/benchmarks/TritonNPUBench/vllm/_expand_idx_mapping_kernel/_expand_idx_mapping_kernel.py b/benchmarks/TritonNPUBench/vllm/_expand_idx_mapping_kernel/_expand_idx_mapping_kernel.py new file mode 100644 index 00000000..2b50b8ed --- /dev/null +++ b/benchmarks/TritonNPUBench/vllm/_expand_idx_mapping_kernel/_expand_idx_mapping_kernel.py @@ -0,0 +1,27 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/vllm_operator_cases/newtest_cases/test_expand_idx_mapping_kernel.py +# Main kernel: _expand_idx_mapping_kernel +# PT file: expand_idx_mapping_test_v2.pt + +import triton +import triton.language as tl + + +# === _expand_idx_mapping_kernel === +@triton.jit +def _expand_idx_mapping_kernel( + idx_mapping_ptr, + expanded_idx_mapping_ptr, + cu_num_logits_ptr, + BLOCK_SIZE: tl.constexpr, +): + req_idx = tl.program_id(0) + start_idx = tl.load(cu_num_logits_ptr + req_idx) + end_idx = tl.load(cu_num_logits_ptr + req_idx + 1) + num_tokens = end_idx - start_idx + + block = tl.arange(0, BLOCK_SIZE) + mask = block < num_tokens + req_state_idx = tl.load(idx_mapping_ptr + req_idx) + tl.store(expanded_idx_mapping_ptr + start_idx + block, req_state_idx, mask=mask) + diff --git a/benchmarks/TritonNPUBench/vllm/_fwd_diag_kernel/_fwd_diag_kernel.py b/benchmarks/TritonNPUBench/vllm/_fwd_diag_kernel/_fwd_diag_kernel.py new file mode 100644 index 00000000..d2414470 --- /dev/null +++ b/benchmarks/TritonNPUBench/vllm/_fwd_diag_kernel/_fwd_diag_kernel.py @@ -0,0 +1,138 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/vllm_operator_cases/newtest_cases/test__fwd_diag_kernel.py +# Main kernel: _fwd_diag_kernel +# PT file: _fwd_diag_kernel_v2.pt + +import triton +import triton.language as tl + + +# === _fwd_diag_kernel === +@triton.jit +def _fwd_diag_kernel( + Q, + K, + V, + Out, + S, + b: tl.constexpr, + h: tl.constexpr, + n, + d: tl.constexpr, + e: tl.constexpr, + BLOCK: tl.constexpr, + NUM_BLOCK, + CBLOCK: tl.constexpr, +): + # This kernel computes the diagonal blocks of the attention matrix + # Each diagonal block represents attention + # where queries attend to keys in the same block + off = tl.program_id(0) + off_bh = off // NUM_BLOCK # batch-head index + off_block = off % NUM_BLOCK # block index within the sequence + off_cblock = tl.program_id(1) # sub-block index within a block + + off_h = off_bh % h # head index + + # Calculate base offsets for the current batch and head + qk_offset = off_bh * n * d + v_offset = off_bh * n * e + o_offset = off_bh * n * e + + # Calculate offsets for the current block + block_offset = off_block * BLOCK + qk_block_offset = block_offset * d + v_block_offset = block_offset * e + o_block_offset = block_offset * e + + # Calculate offsets for the current sub-block + cblock_offset = off_cblock * CBLOCK + q_cblock_offset = cblock_offset * d + o_cblock_offset = cblock_offset * e + + # Calculate pointers to the query, key, value, and output tensors + Q_block_ptr = ( + Q + + qk_offset + + qk_block_offset + + q_cblock_offset + + tl.arange(0, CBLOCK)[:, None] * d + + tl.arange(0, d)[None, :] + ) + K_trans_block_ptr = ( + K + + qk_offset + + qk_block_offset + + tl.arange(0, CBLOCK)[None, :] * d + + tl.arange(0, d)[:, None] + ) + V_block_ptr = ( + V + + v_offset + + v_block_offset + + tl.arange(0, CBLOCK)[:, None] * e + + tl.arange(0, e)[None, :] + ) + O_block_ptr = ( + Out + + o_offset + + o_block_offset + + o_cblock_offset + + tl.arange(0, CBLOCK)[:, None] * e + + tl.arange(0, e)[None, :] + ) + + # Load the decay rate for the current head + S_block_ptr = S + off_h + s = tl.load(S_block_ptr) + + i = off_cblock + q_index = tl.arange(0, CBLOCK) + i * CBLOCK + + # Load query values + q = tl.load(Q_block_ptr, mask=block_offset + q_index[:, None] < n, other=0.0).to( + tl.float32 + ) + + # Initialize output accumulator + qkv = tl.zeros([CBLOCK, e], dtype=tl.float32) + + # Process all sub-blocks up to and + # including the current one (causal attention) + for j in range(i + 1): + kv_index = tl.arange(0, CBLOCK) + j * CBLOCK + diff = q_index[:, None] - kv_index[None, :] + s_index = s * diff + # Apply causal mask: only attend to positions before the current one + s_index = tl.where(diff >= 0, -s_index, float("-inf")) + decay = tl.exp(s_index) + + # Load key and value + k_trans = tl.load( + K_trans_block_ptr, + mask=block_offset + kv_index[None, :] < n, + other=0.0, + ).to(tl.float32) + v = tl.load( + V_block_ptr, + mask=block_offset + kv_index[:, None] < n, + other=0.0, + ).to(tl.float32) + + # Compute attention scores and apply decay + qk = tl.dot(q, k_trans) * decay + + # Compute weighted values and accumulate + qkv += tl.dot(qk, v) + + # Move to the next sub-block + K_trans_block_ptr += CBLOCK * d + V_block_ptr += CBLOCK * e + + # Store the result + tl.store( + O_block_ptr, + qkv.to(O_block_ptr.dtype.element_ty), + mask=block_offset + q_index[:, None] < n, + ) + diff --git a/benchmarks/TritonNPUBench/vllm/_fwd_grouped_kernel_stage1/_fwd_grouped_kernel_stage1.py b/benchmarks/TritonNPUBench/vllm/_fwd_grouped_kernel_stage1/_fwd_grouped_kernel_stage1.py new file mode 100644 index 00000000..b6df76f8 --- /dev/null +++ b/benchmarks/TritonNPUBench/vllm/_fwd_grouped_kernel_stage1/_fwd_grouped_kernel_stage1.py @@ -0,0 +1,176 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/vllm_operator_cases/newtest_cases/test__fwd_grouped_kernel_stage1.py +# Main kernel: _fwd_grouped_kernel_stage1 +# PT file: _fwd_grouped_kernel_stage1_v3.pt + +import triton +import triton.language as tl + + +# === _fwd_grouped_kernel_stage1 === +@triton.jit +def _fwd_grouped_kernel_stage1( + Q, + K_Buffer, + V_Buffer, + sm_scale, + Req_to_tokens, + B_Seqlen, + Att_Out, + stride_req_to_tokens_b, + stride_qbs, + stride_qh, + stride_buf_kbs, + stride_buf_kh, + stride_buf_vbs, + stride_buf_vh, + stride_mid_ob, + stride_mid_oh, + stride_mid_os, + kv_group_num: tl.constexpr, + q_head_num: tl.constexpr, + BLOCK_DMODEL: tl.constexpr, + BLOCK_DPE: tl.constexpr, + BLOCK_DV: tl.constexpr, + BLOCK_N: tl.constexpr, + BLOCK_H: tl.constexpr, + NUM_KV_SPLITS: tl.constexpr, + PAGE_SIZE: tl.constexpr, + logit_cap: tl.constexpr, + Lk: tl.constexpr, + Lv: tl.constexpr, +): + cur_batch = tl.program_id(0) + cur_head_id = tl.program_id(1) + cur_kv_head = cur_head_id // tl.cdiv(kv_group_num, BLOCK_H) + split_kv_id = tl.program_id(2) + + VALID_BLOCK_H: tl.constexpr = BLOCK_H if kv_group_num > BLOCK_H else kv_group_num + cur_head = cur_head_id * VALID_BLOCK_H + tl.arange(0, BLOCK_H) + mask_h = cur_head < (cur_head_id + 1) * VALID_BLOCK_H + mask_h = mask_h & (cur_head < q_head_num) + + offs_d = tl.arange(0, BLOCK_DMODEL) + offs_dv = tl.arange(0, BLOCK_DV) + mask_d = offs_d < Lk + mask_dv = offs_dv < Lv + cur_batch_seq_len = tl.load(B_Seqlen + cur_batch) + cur_batch_req_idx = cur_batch + + offs_q = cur_batch * stride_qbs + cur_head[:, None] * stride_qh + offs_d[None, :] + q = tl.load(Q + offs_q, mask=(mask_h[:, None]) & (mask_d[None, :]), other=0.0) + + if BLOCK_DPE > 0: + offs_dpe = BLOCK_DMODEL + tl.arange(0, BLOCK_DPE) + mask_dpe = offs_dpe < Lk + off_qpe = ( + cur_batch * stride_qbs + cur_head[:, None] * stride_qh + offs_dpe[None, :] + ) + qpe = tl.load( + Q + off_qpe, mask=(mask_h[:, None]) & (mask_dpe[None, :]), other=0.0 + ) + + kv_len_per_split = tl.cdiv(cur_batch_seq_len, NUM_KV_SPLITS) + split_kv_start = kv_len_per_split * split_kv_id + split_kv_end = tl.minimum(split_kv_start + kv_len_per_split, cur_batch_seq_len) + + e_max = tl.zeros([BLOCK_H], dtype=tl.float32) - float("inf") + e_sum = tl.zeros([BLOCK_H], dtype=tl.float32) + acc = tl.zeros([BLOCK_H, BLOCK_DV], dtype=tl.float32) + + if split_kv_end > split_kv_start: + for start_n in range(split_kv_start, split_kv_end, BLOCK_N): + offs_n = start_n + tl.arange(0, BLOCK_N) + kv_page_number = tl.load( + Req_to_tokens + + stride_req_to_tokens_b * cur_batch_req_idx + + offs_n // PAGE_SIZE, + mask=offs_n < split_kv_end, + other=0, + ) + kv_loc = kv_page_number * PAGE_SIZE + offs_n % PAGE_SIZE + offs_buf_k = ( + kv_loc[None, :] * stride_buf_kbs + + cur_kv_head * stride_buf_kh + + offs_d[:, None] + ) + k = tl.load( + K_Buffer + offs_buf_k, + mask=(offs_n[None, :] < split_kv_end) & (mask_d[:, None]), + other=0.0, + ) + qk = tl.dot(q, k.to(q.dtype)) + if BLOCK_DPE > 0: + offs_buf_kpe = ( + kv_loc[None, :] * stride_buf_kbs + + cur_kv_head * stride_buf_kh + + offs_dpe[:, None] + ) + kpe = tl.load( + K_Buffer + offs_buf_kpe, + mask=(offs_n[None, :] < split_kv_end) & (mask_dpe[:, None]), + other=0.0, + ) + qk += tl.dot(qpe, kpe.to(qpe.dtype)) + qk *= sm_scale + + if logit_cap > 0: + qk = logit_cap * tanh(qk / logit_cap) + + qk = tl.where( + mask_h[:, None] & (offs_n[None, :] < split_kv_end), qk, float("-inf") + ) + + offs_buf_v = ( + kv_loc[:, None] * stride_buf_vbs + + cur_kv_head * stride_buf_vh + + offs_dv[None, :] + ) + v = tl.load( + V_Buffer + offs_buf_v, + mask=(offs_n[:, None] < split_kv_end) & (mask_dv[None, :]), + other=0.0, + ) + + n_e_max = tl.maximum(tl.max(qk, 1), e_max) + re_scale = tl.exp(e_max - n_e_max) + p = tl.exp(qk - n_e_max[:, None]) + acc *= re_scale[:, None] + acc += tl.dot(p.to(v.dtype), v) + + e_sum = e_sum * re_scale + tl.sum(p, 1) + e_max = n_e_max + + offs_mid_o = ( + cur_batch * stride_mid_ob + + cur_head[:, None] * stride_mid_oh + + split_kv_id * stride_mid_os + + offs_dv[None, :] + ) + + tl.store( + Att_Out + offs_mid_o, + acc / e_sum[:, None], + mask=(mask_h[:, None]) & (mask_dv[None, :]), + ) + + offs_mid_o_1 = ( + cur_batch * stride_mid_ob + + cur_head * stride_mid_oh + + split_kv_id * stride_mid_os + + Lv + ) + + tl.store( + Att_Out + offs_mid_o_1, + e_max + tl.log(e_sum), + mask=mask_h, + ) + + +# === tanh === +@triton.jit +def tanh(x): + # Tanh is just a scaled sigmoid + return 2 * tl.sigmoid(2 * x) - 1 + diff --git a/benchmarks/TritonNPUBench/vllm/_fwd_kernel/_fwd_kernel.py b/benchmarks/TritonNPUBench/vllm/_fwd_kernel/_fwd_kernel.py new file mode 100644 index 00000000..beab7ec7 --- /dev/null +++ b/benchmarks/TritonNPUBench/vllm/_fwd_kernel/_fwd_kernel.py @@ -0,0 +1,329 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/vllm_operator_cases/newtest_cases/test__fwd_kernel.py +# Main kernel: _fwd_kernel +# PT file: _fwd_kernel_v2.pt + +import triton +import triton.language as tl +import torch + +float8_info = torch.finfo(torch.float16) + +# === _fwd_kernel === +@triton.jit +def _fwd_kernel( + Q, + K, + V, + K_cache, + V_cache, + sink_ptr, + B_Loc, + sm_scale, + k_scale, + v_scale, + out_scale_inv, + B_Start_Loc, + B_Seqlen, + x: tl.constexpr, + Out, + stride_b_loc_b, + stride_b_loc_s, + stride_qbs, + stride_qh, + stride_qd, + stride_kbs, + stride_kh, + stride_kd, + stride_vbs, + stride_vh, + stride_vd, + stride_obs, + stride_oh, + stride_od, + stride_k_cache_bs, + stride_k_cache_h, + stride_k_cache_d, + stride_k_cache_bl: tl.constexpr, + stride_k_cache_x, + stride_v_cache_bs, + stride_v_cache_h, + stride_v_cache_d, + stride_v_cache_bl, + num_queries_per_kv: tl.constexpr, + IN_PRECISION: tl.constexpr, + BLOCK_M: tl.constexpr, + BLOCK_DMODEL: tl.constexpr, + BLOCK_DMODEL_PADDED: tl.constexpr, + BLOCK_SIZE: tl.constexpr, + PHYSICAL_BLOCK_SIZE: tl.constexpr, + BLOCK_N: tl.constexpr, + SLIDING_WINDOW: tl.constexpr, + num_unroll_cache: tl.constexpr, + num_unroll_request: tl.constexpr, + SKIP_DECODE: tl.constexpr, + USE_SINKS: tl.constexpr, + USE_FP8: tl.constexpr, + MAX_Q_LEN: tl.constexpr = 0, + MAX_CTX_LEN: tl.constexpr = 0, + FP8_MIN: tl.constexpr = float8_info.min, + FP8_MAX: tl.constexpr = float8_info.max, +): + cur_batch = tl.program_id(0) + cur_head = tl.program_id(1) + start_m = tl.program_id(2) + + cur_kv_head = cur_head // num_queries_per_kv + + cur_batch_seq_len = tl.load(B_Seqlen + cur_batch) + cur_batch_in_all_start_index = tl.load(B_Start_Loc + cur_batch) + cur_batch_in_all_stop_index = tl.load(B_Start_Loc + cur_batch + 1) + cur_batch_query_len = cur_batch_in_all_stop_index - cur_batch_in_all_start_index + cur_batch_ctx_len = cur_batch_seq_len - cur_batch_query_len + + if SKIP_DECODE and cur_batch_query_len == 1: + return + + # start position inside of the query + # generally, N goes over kv, while M goes over query_len + block_start_loc = BLOCK_M * start_m + + # initialize offsets + # [BLOCK_SIZE]; starts at 0 + offs_bs_n = tl.arange(0, BLOCK_SIZE) + # [N]; starts at 0 + offs_n = tl.arange(0, BLOCK_N) + # [D]; starts at 0 + offs_d = tl.arange(0, BLOCK_DMODEL_PADDED) + # [M]; starts at current position in query + offs_m = start_m * BLOCK_M + tl.arange(0, BLOCK_M) + # [M,D] + off_q = ( + (cur_batch_in_all_start_index + offs_m[:, None]) * stride_qbs + + cur_head * stride_qh + + offs_d[None, :] * stride_qd + ) + + dim_mask = tl.where(tl.arange(0, BLOCK_DMODEL_PADDED) < BLOCK_DMODEL, 1, 0).to( + tl.int1 + ) # [D] + + q = tl.load( + Q + off_q, + mask=dim_mask[None, :] & (offs_m[:, None] < cur_batch_query_len), + other=0.0, + ) # [M,D] + + # initialize pointer to m and l + if not USE_SINKS: + m_i = tl.full([BLOCK_M], float("-inf"), dtype=tl.float32) + l_i = tl.zeros([BLOCK_M], dtype=tl.float32) + else: + m_i = tl.load( + sink_ptr + tl.full([BLOCK_M], cur_head, dtype=tl.int64), + mask=(offs_m < cur_batch_query_len), + other=float("-inf"), + ).to(dtype=tl.float32) + l_i = tl.where(m_i > float("-inf"), 1.0, 0.0) + + acc = tl.zeros([BLOCK_M, BLOCK_DMODEL_PADDED], dtype=tl.float32) # [M,D] + + # compute query against context (no causal mask here) + for start_n in tl.range( + 0, cur_batch_ctx_len, BLOCK_SIZE, loop_unroll_factor=num_unroll_cache + ): + # Under a block size of 544 (Qwen/Qwen3-Next-80B-A3B-Thinking), + # replace one physical block every 17 32-Tile blocks + # Calculate the logical block index of each of the 32 tokens + # in the current Tile (handling cross-block cases). + token_indices = start_n + offs_bs_n + bn_logical_indices = token_indices // PHYSICAL_BLOCK_SIZE + + # 2. Vectorized loading of physical block IDs from B_Loc + bn = tl.load( + B_Loc + cur_batch * stride_b_loc_b + bn_logical_indices * stride_b_loc_s + ).to(tl.int64) + + # 3. Calculate the exact offset of + # each token within its physical block. + internal_offsets = token_indices % PHYSICAL_BLOCK_SIZE + + # Addressing of K (5D) + off_k = ( + bn[None, :] * stride_k_cache_bs + + cur_kv_head * stride_k_cache_h + + (offs_d[:, None] // x) * stride_k_cache_d + + internal_offsets[None, :] * stride_k_cache_bl + + (offs_d[:, None] % x) * stride_k_cache_x + ) + + # Addressing of V (4D) + off_v = ( + bn[:, None] * stride_v_cache_bs + + cur_kv_head * stride_v_cache_h + + offs_d[None, :] * stride_v_cache_d + + internal_offsets[:, None] * stride_v_cache_bl + ) + + if ( + start_n + BLOCK_SIZE > cur_batch_ctx_len + or BLOCK_DMODEL != BLOCK_DMODEL_PADDED + ): + k_load = tl.load( + K_cache + off_k, + mask=dim_mask[:, None] + & ((start_n + offs_bs_n[None, :]) < cur_batch_ctx_len), + other=0.0, + ) # [D,N] + else: + k_load = tl.load(K_cache + off_k) + + if k_load.dtype.is_fp8(): + k = (k_load.to(tl.float32) * tl.load(k_scale)).to(q.dtype) + else: + k = k_load + + # qk = tl.zeros([BLOCK_M, BLOCK_SIZE], dtype=tl.float32) # [M,N] + qk = sm_scale * tl.dot(q, k, input_precision=IN_PRECISION) + qk = tl.where( + (start_n + offs_bs_n[None, :]) < cur_batch_ctx_len, qk, float("-inf") + ) + # qk *= sm_scale + if SLIDING_WINDOW > 0: + # (cur_batch_ctx_len + offs_m[:, None]) are the positions of + # Q entries in sequence + # (start_n + offs_bs_n[None, :]) are the positions of + # KV entries in sequence + # So the condition makes sure each entry in Q only attends + # to KV entries not more than SLIDING_WINDOW away. + # + # We can't use -inf here, because the + # sliding window may lead to the entire row being masked. + # This then makes m_ij contain -inf, which causes NaNs in + # exp(). + qk = tl.where( + (cur_batch_ctx_len + offs_m[:, None]) - (start_n + offs_bs_n[None, :]) + < SLIDING_WINDOW, + qk, + float("-inf"), + ) + + # compute running maximum + m_ij = tl.maximum(m_i, tl.max(qk, axis=1)) + p = tl.exp(qk - m_ij[:, None]) + p = tl.where(m_ij[:, None] == float("-inf"), 0.0, p) + l_ij = tl.sum(p, axis=1) + alpha = tl.exp(m_i - m_ij) + alpha = tl.where(m_i == float("-inf"), 0.0, alpha) + acc = acc * alpha[:, None] + + # update acc + if ( + start_n + BLOCK_SIZE > cur_batch_ctx_len + or BLOCK_DMODEL != BLOCK_DMODEL_PADDED + ): + v_load = tl.load( + V_cache + off_v, + mask=dim_mask[None, :] + & ((start_n + offs_bs_n[:, None]) < cur_batch_ctx_len), + other=0.0, + ) # [N,D] + else: + v_load = tl.load(V_cache + off_v) + + if v_load.dtype.is_fp8(): + v = (v_load.to(tl.float32) * tl.load(v_scale)).to(q.dtype) + else: + v = v_load + p = p.to(v.dtype) + + acc = tl.dot(p, v, acc=acc, input_precision=IN_PRECISION) + # # update m_i and l_i + l_i = l_i * alpha + l_ij + m_i = m_ij + + off_k = ( + offs_n[None, :] * stride_kbs + + cur_kv_head * stride_kh + + offs_d[:, None] * stride_kd + ) + off_v = ( + offs_n[:, None] * stride_vbs + + cur_kv_head * stride_vh + + offs_d[None, :] * stride_vd + ) + k_ptrs = K + off_k + v_ptrs = V + off_v + + # block_mask is 0 when we're already past the current query length + block_mask = tl.where(block_start_loc < cur_batch_query_len, 1, 0) + + # compute query against itself (with causal mask) + for start_n in tl.range( + 0, + block_mask * (start_m + 1) * BLOCK_M, + BLOCK_N, + loop_unroll_factor=num_unroll_request, + ): + start_n = tl.multiple_of(start_n, BLOCK_N) + # -- compute qk ---- + k = tl.load( + k_ptrs + (cur_batch_in_all_start_index + start_n) * stride_kbs, + mask=dim_mask[:, None] + & ((start_n + offs_n[None, :]) < cur_batch_query_len), + other=0.0, + ) + + qk = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32) + qk = tl.dot(q, k, acc=qk, input_precision=IN_PRECISION) + qk *= sm_scale + # apply causal mask + qk = tl.where(offs_m[:, None] >= (start_n + offs_n[None, :]), qk, float("-inf")) + if SLIDING_WINDOW > 0: + qk = tl.where( + offs_m[:, None] - (start_n + offs_n[None, :]) < SLIDING_WINDOW, + qk, + float("-inf"), + ) + + # compute running maximum + m_ij = tl.maximum(m_i, tl.max(qk, axis=1)) + p = tl.exp(qk - m_ij[:, None]) + p = tl.where(m_ij[:, None] == float("-inf"), 0.0, p) + l_ij = tl.sum(p, axis=1) + alpha = tl.exp(m_i - m_ij) + # To prevent NaN from appearing in the first round + alpha = tl.where(m_i == float("-inf"), 0.0, alpha) + acc = acc * alpha[:, None] + + # update acc + v = tl.load( + v_ptrs + (cur_batch_in_all_start_index + start_n) * stride_vbs, + mask=dim_mask[None, :] + & ((start_n + offs_n[:, None]) < cur_batch_query_len), + other=0.0, + ) + p = p.to(v.dtype) + + acc = tl.dot(p, v, acc=acc, input_precision=IN_PRECISION) + # update m_i and l_i + l_i = l_i * alpha + l_ij + m_i = m_ij + + acc = acc / (l_i[:, None] + 1e-10) + + # initialize pointers to output + off_o = ( + (cur_batch_in_all_start_index + offs_m[:, None]) * stride_obs + + cur_head * stride_oh + + offs_d[None, :] * stride_od + ) + out_ptrs = Out + off_o + if USE_FP8: + acc = acc * tl.load(out_scale_inv) + acc = tl.clamp(acc, FP8_MIN, FP8_MAX) + tl.store( + out_ptrs, acc, mask=dim_mask[None, :] & (offs_m[:, None] < cur_batch_query_len) + ) + return + diff --git a/benchmarks/TritonNPUBench/vllm/_fwd_kernel_alibi/_fwd_kernel_alibi.py b/benchmarks/TritonNPUBench/vllm/_fwd_kernel_alibi/_fwd_kernel_alibi.py new file mode 100644 index 00000000..540fdb71 --- /dev/null +++ b/benchmarks/TritonNPUBench/vllm/_fwd_kernel_alibi/_fwd_kernel_alibi.py @@ -0,0 +1,289 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/vllm_operator_cases/newtest_cases/test__fwd_kernel_alibi.py +# Main kernel: _fwd_kernel_alibi +# PT file: _fwd_kernel_alibi_v2.pt + +import triton +import triton.language as tl + + +# === _fwd_kernel_alibi === +@triton.jit +def _fwd_kernel_alibi( + Q, + K, + V, + K_cache, + V_cache, + B_Loc, + sm_scale, + k_scale, + v_scale, + B_Start_Loc, + B_Seqlen, + Alibi_slopes, + block_size, + x, + Out, + stride_b_loc_b, + stride_b_loc_s, + stride_qbs, + stride_qh, + stride_qd, + stride_kbs, + stride_kh, + stride_kd, + stride_vbs, + stride_vh, + stride_vd, + stride_obs, + stride_oh, + stride_od, + stride_k_cache_bs, + stride_k_cache_h, + stride_k_cache_d, + stride_k_cache_bl, + stride_k_cache_x, + stride_v_cache_bs, + stride_v_cache_h, + stride_v_cache_d, + stride_v_cache_bl, + num_queries_per_kv: int, + IN_PRECISION: tl.constexpr, + BLOCK_M: tl.constexpr, + BLOCK_DMODEL: tl.constexpr, # head size + BLOCK_DMODEL_PADDED: tl.constexpr, # head size padded to a power of 2 + BLOCK_N: tl.constexpr, + SKIP_DECODE: tl.constexpr, +): + # attn_bias[] + cur_batch = tl.program_id(0) + cur_head = tl.program_id(1) + start_m = tl.program_id(2) + + cur_kv_head = cur_head // num_queries_per_kv + + # cur_batch_seq_len: the length of prompts + # cur_batch_ctx_len: the length of prefix + # cur_batch_in_all_start_index: the start id of the dim=0 + cur_batch_seq_len = tl.load(B_Seqlen + cur_batch) + cur_batch_in_all_start_index = tl.load(B_Start_Loc + cur_batch) + cur_batch_in_all_stop_index = tl.load(B_Start_Loc + cur_batch + 1) + cur_batch_query_len = cur_batch_in_all_stop_index - cur_batch_in_all_start_index + cur_batch_ctx_len = cur_batch_seq_len - cur_batch_query_len + + if SKIP_DECODE and cur_batch_query_len == 1: + return + + block_start_loc = BLOCK_M * start_m + + # initialize offsets + offs_n = tl.arange(0, BLOCK_N) + offs_d = tl.arange(0, BLOCK_DMODEL_PADDED) + offs_m = start_m * BLOCK_M + tl.arange(0, BLOCK_M) + off_q = ( + (cur_batch_in_all_start_index + offs_m[:, None]) * stride_qbs + + cur_head * stride_qh + + offs_d[None, :] * stride_qd + ) + + dim_mask = tl.where(tl.arange(0, BLOCK_DMODEL_PADDED) < BLOCK_DMODEL, 1, 0).to( + tl.int1 + ) + + q = tl.load( + Q + off_q, + mask=dim_mask[None, :] + & (offs_m[:, None] < cur_batch_seq_len - cur_batch_ctx_len), + other=0.0, + ) + + # # initialize pointer to m and l + m_i = tl.zeros([BLOCK_M], dtype=tl.float32) - float("inf") + l_i = tl.zeros([BLOCK_M], dtype=tl.float32) + acc = tl.zeros([BLOCK_M, BLOCK_DMODEL_PADDED], dtype=tl.float32) + + alibi_slope = tl.load(Alibi_slopes + cur_head) + alibi_start_q = tl.arange(0, BLOCK_M) + block_start_loc + cur_batch_ctx_len + alibi_start_k = 0 + for start_n in range(0, cur_batch_ctx_len, BLOCK_N): + start_n = tl.multiple_of(start_n, BLOCK_N) + # -- compute qk ---- + bn = tl.load( + B_Loc + + cur_batch * stride_b_loc_b + + ((start_n + offs_n) // block_size) * stride_b_loc_s, + mask=(start_n + offs_n) < cur_batch_ctx_len, + other=0, + ).to(tl.int64) + off_k = ( + bn[None, :] * stride_k_cache_bs + + cur_kv_head * stride_k_cache_h + + (offs_d[:, None] // x) * stride_k_cache_d + + ((start_n + offs_n[None, :]) % block_size) * stride_k_cache_bl + + (offs_d[:, None] % x) * stride_k_cache_x + ) + off_v = ( + bn[:, None] * stride_v_cache_bs + + cur_kv_head * stride_v_cache_h + + offs_d[None, :] * stride_v_cache_d + + (start_n + offs_n[:, None]) % block_size * stride_v_cache_bl + ) + k_load = tl.load( + K_cache + off_k, + mask=dim_mask[:, None] & ((start_n + offs_n[None, :]) < cur_batch_ctx_len), + other=0.0, + ) # [D,N] + + if k_load.dtype.is_fp8(): + k = (k_load.to(tl.float32) * tl.load(k_scale)).to(q.dtype) + else: + k = k_load + + qk = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32) + qk = tl.dot(q, k, acc=qk, input_precision=IN_PRECISION) + qk = tl.where( + (start_n + offs_n[None, :]) < cur_batch_ctx_len, qk, float("-inf") + ) + qk *= sm_scale + + # load alibi + alibi = ( + tl.arange(0, BLOCK_N)[None, :] + alibi_start_k - alibi_start_q[:, None] + ) * alibi_slope + alibi = tl.where( + (alibi <= 0) & (alibi_start_q[:, None] < cur_batch_seq_len), + alibi, + float("-inf"), + ) + qk += alibi + alibi_start_k += BLOCK_N + + # -- compute m_ij, p, l_ij + m_ij = tl.max(qk, 1) + m_i_new = tl.maximum(m_i, m_ij) + p = tl.math.exp(qk - m_i_new[:, None]) + l_ij = tl.sum(p, 1) + # -- update m_i and l_i + + alpha = tl.math.exp(m_i - m_i_new) + l_i_new = alpha * l_i + l_ij + # -- update output accumulator -- + # scale p + # scale acc + acc_scale = alpha + # acc_scale = l_i / l_i_new * alpha + acc = acc * acc_scale[:, None] + # update acc + v_load = tl.load( + V_cache + off_v, + mask=dim_mask[None, :] & ((start_n + offs_n[:, None]) < cur_batch_ctx_len), + other=0.0, + ) + if v_load.dtype.is_fp8(): + v = (v_load.to(tl.float32) * tl.load(v_scale)).to(q.dtype) + else: + v = v_load + p = p.to(v.dtype) + + acc = tl.dot(p, v, acc=acc, input_precision="ieee") + # update m_i and l_i + l_i = l_i_new + m_i = m_i_new + + off_k = ( + offs_n[None, :] * stride_kbs + + cur_kv_head * stride_kh + + offs_d[:, None] * stride_kd + ) + off_v = ( + offs_n[:, None] * stride_vbs + + cur_kv_head * stride_vh + + offs_d[None, :] * stride_vd + ) + k_ptrs = K + off_k + v_ptrs = V + off_v + + block_mask = tl.where(block_start_loc < cur_batch_seq_len - cur_batch_ctx_len, 1, 0) + + # init alibi + alibi_slope = tl.load(Alibi_slopes + cur_head) + alibi_start_q = tl.arange(0, BLOCK_M) + block_start_loc + cur_batch_ctx_len + alibi_start_k = cur_batch_ctx_len + # # init debugger + # offset_db_q = tl.arange(0, BLOCK_M) + block_start_loc + # offset_db_k = tl.arange(0, BLOCK_N) + # calc q[BLOCK_M, BLOCK_MODEL] mul k[prefix_len: , BLOCK_DMODEL] + for start_n in range(0, block_mask * (start_m + 1) * BLOCK_M, BLOCK_N): + start_n = tl.multiple_of(start_n, BLOCK_N) + # -- compute qk ---- + k = tl.load( + k_ptrs + (cur_batch_in_all_start_index + start_n) * stride_kbs, + mask=dim_mask[:, None] + & ((start_n + offs_n[None, :]) < cur_batch_seq_len - cur_batch_ctx_len), + other=0.0, + ) + + qk = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32) + qk = tl.dot(q, k, acc=qk, input_precision="ieee") + qk *= sm_scale + qk = tl.where(offs_m[:, None] >= (start_n + offs_n[None, :]), qk, float("-inf")) + + # load alibi + alibi = ( + tl.arange(0, BLOCK_N)[None, :] + alibi_start_k - alibi_start_q[:, None] + ) * alibi_slope + alibi = tl.where( + (alibi <= 0) & (alibi_start_q[:, None] < cur_batch_seq_len), + alibi, + float("-inf"), + ) + qk += alibi + alibi_start_k += BLOCK_N + + # -- compute m_ij, p, l_ij + m_ij = tl.max(qk, 1) + m_i_new = tl.maximum(m_i, m_ij) + p = tl.math.exp(qk - m_i_new[:, None]) + l_ij = tl.sum(p, 1) + # -- update m_i and l_i + + alpha = tl.math.exp(m_i - m_i_new) + l_i_new = alpha * l_i + l_ij + # -- update output accumulator -- + # scale p + # scale acc + acc_scale = alpha + # acc_scale = l_i / l_i_new * alpha + acc = acc * acc_scale[:, None] + # update acc + v = tl.load( + v_ptrs + (cur_batch_in_all_start_index + start_n) * stride_vbs, + mask=dim_mask[None, :] + & ((start_n + offs_n[:, None]) < cur_batch_seq_len - cur_batch_ctx_len), + other=0.0, + ) + p = p.to(v.dtype) + + acc = tl.dot(p, v, acc=acc, input_precision="ieee") + # update m_i and l_i + l_i = l_i_new + m_i = m_i_new + + acc = acc / l_i[:, None] + + # initialize pointers to output + off_o = ( + (cur_batch_in_all_start_index + offs_m[:, None]) * stride_obs + + cur_head * stride_oh + + offs_d[None, :] * stride_od + ) + out_ptrs = Out + off_o + tl.store( + out_ptrs, + acc, + mask=dim_mask[None, :] + & (offs_m[:, None] < cur_batch_seq_len - cur_batch_ctx_len), + ) + return + diff --git a/benchmarks/TritonNPUBench/vllm/_fwd_kernel_ep_gather/_fwd_kernel_ep_gather.py b/benchmarks/TritonNPUBench/vllm/_fwd_kernel_ep_gather/_fwd_kernel_ep_gather.py new file mode 100644 index 00000000..dc0076a6 --- /dev/null +++ b/benchmarks/TritonNPUBench/vllm/_fwd_kernel_ep_gather/_fwd_kernel_ep_gather.py @@ -0,0 +1,75 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/vllm_operator_cases/newtest_cases/test__fwd_kernel_ep_gather.py +# Main kernel: _fwd_kernel_ep_gather +# PT file: test__fwd_kernel_ep_gather_v2.pt + +import triton +import triton.language as tl + + +# === _fwd_kernel_ep_gather === +@triton.jit +def _fwd_kernel_ep_gather( + total_token_num, + input_tensor, + input_tensor_stride0, + input_tensor_stride1, + recv_topk_ids, + recv_topk_ids_stride0, + recv_topk_ids_stride1, + recv_topk_weight, + recv_topk_weight_stride0, + recv_topk_weight_stride1, + input_index, + input_index_stride0, + input_index_stride1, + output_tensor, + output_tensor_stride0, + output_tensor_stride1, + topk_num: tl.constexpr, + BLOCK_D: tl.constexpr, +): + cur_block_int32 = tl.program_id(0) + cur_block = cur_block_int32.to(tl.int64) + + start_cur_token_int32 = tl.program_id(1) + + grid_num = tl.num_programs(1) + + for cur_token_int32 in range(start_cur_token_int32, total_token_num, grid_num): + cur_token = cur_token_int32.to(tl.int64) + + off_d = tl.arange(0, BLOCK_D) + accumulator = tl.zeros([BLOCK_D], dtype=tl.float32) + + for topk_index_int32 in range(0, topk_num): + topk_index = topk_index_int32.to(tl.int64) + + expert_id = tl.load( + recv_topk_ids + cur_token * recv_topk_ids_stride0 + topk_index + ) + if expert_id >= 0: + source_token_index_int32 = tl.load( + input_index + cur_token * input_index_stride0 + topk_index + ) + source_token_index = source_token_index_int32.to(tl.int64) + + acc_weight = tl.load( + recv_topk_weight + cur_token * recv_topk_weight_stride0 + topk_index + ) + tmp = tl.load( + input_tensor + + source_token_index * input_tensor_stride0 + + cur_block * BLOCK_D + + off_d + ) + accumulator += tmp.to(tl.float32) * acc_weight + + tl.store( + output_tensor + + cur_token * output_tensor_stride0 + + cur_block * BLOCK_D + + off_d, + accumulator.to(output_tensor.dtype.element_ty), + ) + diff --git a/benchmarks/TritonNPUBench/vllm/_fwd_kernel_ep_scatter_1/_fwd_kernel_ep_scatter_1.py b/benchmarks/TritonNPUBench/vllm/_fwd_kernel_ep_scatter_1/_fwd_kernel_ep_scatter_1.py new file mode 100644 index 00000000..0b515ed1 --- /dev/null +++ b/benchmarks/TritonNPUBench/vllm/_fwd_kernel_ep_scatter_1/_fwd_kernel_ep_scatter_1.py @@ -0,0 +1,56 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/vllm_operator_cases/newtest_cases/test__fwd_kernel_ep_scatter_1.py +# Main kernel: _fwd_kernel_ep_scatter_1 +# PT file: _fwd_kernel_ep_scatter_1_v2.pt + +import triton +import triton.language as tl + + +# === _fwd_kernel_ep_scatter_1 === +@triton.jit +def _fwd_kernel_ep_scatter_1( + num_recv_tokens_per_expert, + expert_start_loc, + m_indices, + num_experts: tl.constexpr, + BLOCK_E: tl.constexpr, + BLOCK_EXPERT_NUM: tl.constexpr, +): + cur_expert = tl.program_id(0) + + offset_cumsum = tl.arange(0, BLOCK_EXPERT_NUM) + tokens_per_expert = tl.load( + num_recv_tokens_per_expert + offset_cumsum, + mask=offset_cumsum < num_experts, + other=0, + ) + tokens_per_expert = round_up_128(tokens_per_expert) + cumsum = tl.cumsum(tokens_per_expert) - tokens_per_expert + tl.store(expert_start_loc + offset_cumsum, cumsum, mask=offset_cumsum < num_experts) + + cur_expert_start = tl.load(expert_start_loc + cur_expert) + cur_expert_token_num = tl.load(num_recv_tokens_per_expert + cur_expert) + + m_indices_start_ptr = m_indices + cur_expert_start + off_expert = tl.arange(0, BLOCK_E) + + # any rows in the per-expert aligned region that do not correspond to + # real tokens are left untouched here and should remain initialized to + # -1 so DeepGEMM can skip them + for start_m in tl.range(0, cur_expert_token_num, BLOCK_E, num_stages=4): + offs = start_m + off_expert + mask = offs < cur_expert_token_num + tl.store( + m_indices_start_ptr + offs, + cur_expert, + mask=mask, + ) + + +# === round_up_128 === +@triton.jit +def round_up_128(x: int) -> int: + y = 128 + return ((x + y - 1) // y) * y + diff --git a/benchmarks/TritonNPUBench/vllm/_fwd_kernel_ep_scatter_2/_fwd_kernel_ep_scatter_2.py b/benchmarks/TritonNPUBench/vllm/_fwd_kernel_ep_scatter_2/_fwd_kernel_ep_scatter_2.py new file mode 100644 index 00000000..2a1dc947 --- /dev/null +++ b/benchmarks/TritonNPUBench/vllm/_fwd_kernel_ep_scatter_2/_fwd_kernel_ep_scatter_2.py @@ -0,0 +1,84 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/vllm_operator_cases/newtest_cases/test__fwd_kernel_ep_scatter_2.py +# Main kernel: _fwd_kernel_ep_scatter_2 +# PT file: None + +import triton +import triton.language as tl + + +# === _fwd_kernel_ep_scatter_2 === +@triton.jit +def _fwd_kernel_ep_scatter_2( + total_token_num, + expert_start_loc, + recv_x, + recv_x_stride0, + recv_x_stride1, + recv_x_scale, + recv_x_scale_stride0, + recv_x_scale_stride1, + recv_topk, + recv_topk_stride0, + recv_topk_stride1, + output_tensor, + output_tensor_stride0, + output_tensor_stride1, + output_tensor_scale, + output_tensor_scale_stride0, + output_tensor_scale_stride1, + output_index, + output_index_stride0, + output_index_stride1, + topk_num: tl.constexpr, + expert_map, + HAS_EXPERT_MAP: tl.constexpr, + HIDDEN_SIZE: tl.constexpr, + HIDDEN_SIZE_PAD: tl.constexpr, + SCALE_HIDDEN_SIZE: tl.constexpr, + SCALE_HIDDEN_SIZE_PAD: tl.constexpr, +): + start_token_id = tl.program_id(0) + grid_num = tl.num_programs(0) + + offset_in = tl.arange(0, HIDDEN_SIZE_PAD) + mask = offset_in < HIDDEN_SIZE + + offset_in_s = tl.arange(0, SCALE_HIDDEN_SIZE_PAD) + mask_s = offset_in_s < SCALE_HIDDEN_SIZE + + for token_id in range(start_token_id, total_token_num, grid_num): + to_copy = tl.load(recv_x + token_id * recv_x_stride0 + offset_in, mask=mask) + to_copy_s = tl.load( + recv_x_scale + token_id * recv_x_scale_stride0 + offset_in_s, mask=mask_s + ) + + for topk_index in tl.range(0, topk_num, 1, num_stages=4): + expert_id = tl.load(recv_topk + token_id * recv_topk_stride0 + topk_index) + + if HAS_EXPERT_MAP: + expert_id = apply_expert_map(expert_id, expert_map) + + if expert_id >= 0: + dest_token_index = tl.atomic_add(expert_start_loc + expert_id, 1) + tl.store( + output_index + token_id * output_index_stride0 + topk_index, + dest_token_index, + ) + output_tensor_ptr = ( + output_tensor + dest_token_index * output_tensor_stride0 + ) + output_tensor_scale_ptr = ( + output_tensor_scale + dest_token_index * output_tensor_scale_stride0 + ) + tl.store(output_tensor_ptr + offset_in, to_copy, mask=mask) + tl.store(output_tensor_scale_ptr + offset_in_s, to_copy_s, mask=mask_s) + + +# === apply_expert_map === +@triton.jit +def apply_expert_map(expert_id, expert_map): + if expert_id != -1: + expert_id = tl.load(expert_map + expert_id).to(expert_id.dtype) + return expert_id + diff --git a/benchmarks/TritonNPUBench/vllm/_fwd_kernel_stage1/_fwd_kernel_stage1.py b/benchmarks/TritonNPUBench/vllm/_fwd_kernel_stage1/_fwd_kernel_stage1.py new file mode 100644 index 00000000..6da7ae05 --- /dev/null +++ b/benchmarks/TritonNPUBench/vllm/_fwd_kernel_stage1/_fwd_kernel_stage1.py @@ -0,0 +1,144 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/vllm_operator_cases/newtest_cases/test__fwd_kernel_stage1.py +# Main kernel: _fwd_kernel_stage1 +# PT file: test__fwd_kernel_stage1_v3.pt + +import triton +import triton.language as tl + + +# === _fwd_kernel_stage1 === +@triton.jit +def _fwd_kernel_stage1( + Q, + K_Buffer, + V_Buffer, + sm_scale, + Req_to_tokens, + B_Seqlen, + Att_Out, + stride_req_to_tokens_b, + stride_qbs, + stride_qh, + stride_buf_kbs, + stride_buf_kh, + stride_buf_vbs, + stride_buf_vh, + stride_mid_ob, + stride_mid_oh, + stride_mid_os, + kv_group_num: tl.constexpr, + BLOCK_DMODEL: tl.constexpr, + BLOCK_DV: tl.constexpr, + BLOCK_N: tl.constexpr, + NUM_KV_SPLITS: tl.constexpr, + PAGE_SIZE: tl.constexpr, + logit_cap: tl.constexpr, + Lk: tl.constexpr, + Lv: tl.constexpr, +): + cur_batch = tl.program_id(0) + cur_head = tl.program_id(1) + split_kv_id = tl.program_id(2) + + cur_kv_head = cur_head // kv_group_num + + offs_d = tl.arange(0, BLOCK_DMODEL) + offs_dv = tl.arange(0, BLOCK_DV) + mask_d = offs_d < Lk + mask_dv = offs_dv < Lv + cur_batch_seq_len = tl.load(B_Seqlen + cur_batch) + cur_batch_req_idx = cur_batch + + off_q = cur_batch * stride_qbs + cur_head * stride_qh + offs_d + q = tl.load(Q + off_q, mask=mask_d, other=0.0) + + kv_len_per_split = tl.cdiv(cur_batch_seq_len, NUM_KV_SPLITS) + split_kv_start = kv_len_per_split * split_kv_id + split_kv_end = tl.minimum(split_kv_start + kv_len_per_split, cur_batch_seq_len) + + e_max = -float("inf") + e_sum = 0.0 + acc = tl.zeros([BLOCK_DV], dtype=tl.float32) + + if split_kv_end > split_kv_start: + for start_n in range(split_kv_start, split_kv_end, BLOCK_N): + offs_n = start_n + tl.arange(0, BLOCK_N) + kv_page_number = tl.load( + Req_to_tokens + + stride_req_to_tokens_b * cur_batch_req_idx + + offs_n // PAGE_SIZE, + mask=offs_n < split_kv_end, + other=0, + ) + kv_loc = kv_page_number * PAGE_SIZE + offs_n % PAGE_SIZE + offs_buf_k = ( + kv_loc[:, None] * stride_buf_kbs + + cur_kv_head * stride_buf_kh + + offs_d[None, :] + ) + k = tl.load( + K_Buffer + offs_buf_k, + mask=(offs_n[:, None] < split_kv_end) & (mask_d[None, :]), + other=0.0, + ) + qk = tl.sum(q[None, :] * k, 1) + qk *= sm_scale + + if logit_cap > 0: + qk = logit_cap * tanh(qk / logit_cap) + + qk = tl.where(offs_n < split_kv_end, qk, float("-inf")) + + offs_buf_v = ( + kv_loc[:, None] * stride_buf_vbs + + cur_kv_head * stride_buf_vh + + offs_dv[None, :] + ) + v = tl.load( + V_Buffer + offs_buf_v, + mask=(offs_n[:, None] < split_kv_end) & (mask_dv[None, :]), + other=0.0, + ) + + n_e_max = tl.maximum(tl.max(qk, 0), e_max) + re_scale = tl.exp(e_max - n_e_max) + p = tl.exp(qk - n_e_max) + acc *= re_scale + acc += tl.sum(p[:, None] * v, 0) + + e_sum = e_sum * re_scale + tl.sum(p, 0) + e_max = n_e_max + + offs_mid_o = ( + cur_batch * stride_mid_ob + + cur_head * stride_mid_oh + + split_kv_id * stride_mid_os + + offs_dv + ) + + tl.store( + Att_Out + offs_mid_o, + acc / e_sum, + mask=(mask_dv), + ) + + offs_mid_o_1 = ( + cur_batch * stride_mid_ob + + cur_head * stride_mid_oh + + split_kv_id * stride_mid_os + + Lv + ) + + tl.store( + Att_Out + offs_mid_o_1, + e_max + tl.log(e_sum), + ) + + +# === tanh === +@triton.jit +def tanh(x): + # Tanh is just a scaled sigmoid + return 2 * tl.sigmoid(2 * x) - 1 + diff --git a/benchmarks/TritonNPUBench/vllm/_fwd_kernel_stage2/_fwd_kernel_stage2.py b/benchmarks/TritonNPUBench/vllm/_fwd_kernel_stage2/_fwd_kernel_stage2.py new file mode 100644 index 00000000..b78aa8f5 --- /dev/null +++ b/benchmarks/TritonNPUBench/vllm/_fwd_kernel_stage2/_fwd_kernel_stage2.py @@ -0,0 +1,72 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/vllm_operator_cases/newtest_cases/test__fwd_kernel_stage2.py +# Main kernel: _fwd_kernel_stage2 +# PT file: test__fwd_kernel_stage2_v2.pt + +import triton +import triton.language as tl + + +# === _fwd_kernel_stage2 === +@triton.jit +def _fwd_kernel_stage2( + Mid_O, + o, + lse, + B_Seqlen, + stride_mid_ob, + stride_mid_oh, + stride_mid_os, + stride_obs, + stride_oh, + stride_lse_bs, + NUM_KV_SPLITS: tl.constexpr, + BLOCK_DV: tl.constexpr, + Lv: tl.constexpr, +): + cur_batch = tl.program_id(0) + cur_head = tl.program_id(1) + + cur_batch_seq_len = tl.load(B_Seqlen + cur_batch) + + offs_d = tl.arange(0, BLOCK_DV) + mask_d = offs_d < Lv + + e_sum = 0.0 + e_max = -float("inf") + acc = tl.zeros([BLOCK_DV], dtype=tl.float32) + + offs_v = cur_batch * stride_mid_ob + cur_head * stride_mid_oh + offs_d + offs_logic = cur_batch * stride_mid_ob + cur_head * stride_mid_oh + Lv + + for split_kv_id in range(0, NUM_KV_SPLITS): + kv_len_per_split = tl.cdiv(cur_batch_seq_len, NUM_KV_SPLITS) + split_kv_start = kv_len_per_split * split_kv_id + split_kv_end = tl.minimum(split_kv_start + kv_len_per_split, cur_batch_seq_len) + + if split_kv_end > split_kv_start: + tv = tl.load( + Mid_O + offs_v + split_kv_id * stride_mid_os, mask=mask_d, other=0.0 + ) + tlogic = tl.load(Mid_O + offs_logic + split_kv_id * stride_mid_os) + n_e_max = tl.maximum(tlogic, e_max) + + old_scale = tl.exp(e_max - n_e_max) + acc *= old_scale + exp_logic = tl.exp(tlogic - n_e_max) + acc += exp_logic * tv + + e_sum = e_sum * old_scale + exp_logic + e_max = n_e_max + + tl.store( + o + cur_batch * stride_obs + cur_head * stride_oh + offs_d, + acc / e_sum, + mask=mask_d, + ) + lse_val = e_max + tl.log(e_sum) + tl.store( + lse + cur_batch * stride_lse_bs + cur_head, + lse_val, + ) + diff --git a/benchmarks/TritonNPUBench/vllm/_fwd_kv_parallel/_fwd_kv_parallel.py b/benchmarks/TritonNPUBench/vllm/_fwd_kv_parallel/_fwd_kv_parallel.py new file mode 100644 index 00000000..0c122ed6 --- /dev/null +++ b/benchmarks/TritonNPUBench/vllm/_fwd_kv_parallel/_fwd_kv_parallel.py @@ -0,0 +1,118 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/vllm_operator_cases/newtest_cases/test__fwd_kv_parallel.py +# Main kernel: _fwd_kv_parallel +# PT file: _fwd_kv_parallel_v2.pt + +import triton +import triton.language as tl + + +# === _fwd_kv_parallel === +@triton.jit +def _fwd_kv_parallel( + K, + V, + K_decay, + KV, + b: tl.constexpr, + h: tl.constexpr, + n, + d: tl.constexpr, + e: tl.constexpr, + BLOCK: tl.constexpr, + NUM_BLOCK, + D_FBLOCK: tl.constexpr, + E_FBLOCK: tl.constexpr, + NUM_FBLOCK: tl.constexpr, + CBLOCK: tl.constexpr, + NUM_CBLOCK: tl.constexpr, +): + # This kernel computes the key-value outer + # products for each block in parallel + off_bh = tl.program_id(0) # batch-head index + off_block = tl.program_id(1) # block index + + off_h = off_bh % h # head index + + block_offset = off_block * BLOCK + + # Calculate offsets for the current block + k_block_offset = block_offset * d + v_block_offset = block_offset * e + kv_block_offset = off_block * d * e + + # Calculate base offsets for the current batch and head + k_offset = off_bh * n * d + v_offset = off_bh * n * e + kv_offset = off_bh * NUM_BLOCK * d * e + + # Calculate pointers to the key, value, and key-value tensors + K_trans_block_ptr = ( + K + + k_offset + + k_block_offset + + tl.arange(0, CBLOCK)[None, :] * d + + tl.arange(0, D_FBLOCK)[:, None] + ) + V_block_ptr = ( + V + + v_offset + + v_block_offset + + tl.arange(0, CBLOCK)[:, None] * e + + tl.arange(0, E_FBLOCK)[None, :] + ) + KV_block_ptr = ( + KV + + kv_offset + + kv_block_offset + + tl.arange(0, D_FBLOCK)[:, None] * e + + tl.arange(0, E_FBLOCK)[None, :] + ) + + # Load the decay factors for the current head and block + k_decay_ptr = K_decay + off_h * BLOCK + tl.arange(0, CBLOCK) + + kv_index = tl.arange(0, CBLOCK) + + # Initialize the key-value outer product accumulator + kv = tl.zeros([D_FBLOCK, E_FBLOCK], dtype=tl.float32) + + # Handle the last block which might be smaller than BLOCK + split_n = n - (NUM_BLOCK - 1) * BLOCK if off_block == NUM_BLOCK - 1 else BLOCK + left_shift = tl.cdiv(split_n, CBLOCK) * CBLOCK - split_n + num_blocks = min(tl.cdiv(split_n, CBLOCK), NUM_CBLOCK) + k_decay_ptr += (NUM_CBLOCK - num_blocks) * CBLOCK + + # Process all sub-blocks in the current block + for j in range(num_blocks): + left_bound = (1 - j) * left_shift + # Load key and value, handling boundary conditions + k_trans = tl.load( + K_trans_block_ptr - left_shift * d, + mask=kv_index[None, :] >= left_bound, + other=0.0, + ) + v = tl.load( + V_block_ptr - left_shift * e, + mask=kv_index[:, None] >= left_bound, + other=0.0, + ) + + # Load decay factor and compute weighted key-value outer product + k_decay = tl.load(k_decay_ptr) + + # NOTE: Need to add the extra dim here due to AMD MLIR lowering error. + # Please don't move it back until issue is resolved. + # Issue: https://github.com/ROCm/triton/issues/907 + k_decay = k_decay[None, :] + + kv += tl.dot(k_trans * k_decay, v) + + # Move to the next sub-block + K_trans_block_ptr += CBLOCK * d + V_block_ptr += CBLOCK * e + k_decay_ptr += CBLOCK + + # Store the result + tl.store(KV_block_ptr, kv.to(KV_block_ptr.dtype.element_ty)) + diff --git a/benchmarks/TritonNPUBench/vllm/_fwd_kv_reduce/_fwd_kv_reduce.py b/benchmarks/TritonNPUBench/vllm/_fwd_kv_reduce/_fwd_kv_reduce.py new file mode 100644 index 00000000..0f18ce0e --- /dev/null +++ b/benchmarks/TritonNPUBench/vllm/_fwd_kv_reduce/_fwd_kv_reduce.py @@ -0,0 +1,74 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/vllm_operator_cases/newtest_cases/test__fwd_kv_reduce.py +# Main kernel: _fwd_kv_reduce +# PT file: _fwd_kv_reduce_v2.pt + +import triton +import triton.language as tl + + +# === _fwd_kv_reduce === +@triton.jit +def _fwd_kv_reduce( + S, + KV, + KV_HISTORY, + b: tl.constexpr, + h: tl.constexpr, + n, + d: tl.constexpr, + e: tl.constexpr, + BLOCK: tl.constexpr, + NUM_BLOCK, + D_FBLOCK: tl.constexpr, + E_FBLOCK: tl.constexpr, +): + # This kernel reduces the key-value outer products + # across blocks and updates the KV history + off_bh = tl.program_id(0) # batch-head index + off_h = off_bh % h # head index + + kv_offset = off_bh * NUM_BLOCK * d * e + + # Calculate pointer to the key-value tensor + KV_block_ptr = ( + KV + + kv_offset + + tl.arange(0, D_FBLOCK)[:, None] * e + + tl.arange(0, E_FBLOCK)[None, :] + ) + + # Load the decay rate for the current head + s_ptrs = S + off_h + s = tl.load(s_ptrs) + + # Calculate pointer to the key-value history tensor + kv_history_offset = off_bh * d * e + KV_HISTORY_block_ptr = ( + KV_HISTORY + + kv_history_offset + + tl.arange(0, D_FBLOCK)[:, None] * e + + tl.arange(0, E_FBLOCK)[None, :] + ) + + # Load the previous key-value history + kv_pre = tl.load(KV_HISTORY_block_ptr).to(tl.float32) + + # Process all blocks in reverse order to compute the prefix sum + for i in range(NUM_BLOCK): + block_size = min(n - i * BLOCK, BLOCK) + # Compute decay factor for the current block + block_decay = tl.exp(-s.to(tl.float32) * block_size) + + # Load the current key-value outer product + kv_cur = tl.load(KV_block_ptr).to(tl.float32) + # Store the previous key-value history to the current block + tl.store(KV_block_ptr, kv_pre.to(KV_block_ptr.dtype.element_ty)) + + # Update the key-value history with the current block + kv_pre = block_decay * kv_pre + kv_cur + KV_block_ptr += d * e + + # Store the updated key-value history + tl.store(KV_HISTORY_block_ptr, kv_pre) + diff --git a/benchmarks/TritonNPUBench/vllm/_fwd_none_diag_kernel/_fwd_none_diag_kernel.py b/benchmarks/TritonNPUBench/vllm/_fwd_none_diag_kernel/_fwd_none_diag_kernel.py new file mode 100644 index 00000000..8a36fe68 --- /dev/null +++ b/benchmarks/TritonNPUBench/vllm/_fwd_none_diag_kernel/_fwd_none_diag_kernel.py @@ -0,0 +1,92 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/vllm_operator_cases/newtest_cases/test__fwd_none_diag_kernel.py +# Main kernel: _fwd_none_diag_kernel +# PT file: _fwd_none_diag_kernel_new.pt + +import triton +import triton.language as tl + + +# === _fwd_none_diag_kernel === +@triton.jit +def _fwd_none_diag_kernel( + Q, + Out, + S, + KV, + b: tl.constexpr, + h: tl.constexpr, + n, + d: tl.constexpr, + e: tl.constexpr, + BLOCK: tl.constexpr, + NUM_BLOCK, + E_FBLOCK: tl.constexpr, + CBLOCK: tl.constexpr, + NUM_CBLOCK: tl.constexpr, +): + # This kernel computes the non-diagonal blocks of the attention matrix + # Each non-diagonal block represents attention + # where queries attend to keys in different blocks + off_bh = tl.program_id(0) # batch-head index + off_h = off_bh % h # head index + + off_nc = tl.program_id(1) + off_n = off_nc // NUM_CBLOCK # block index + off_c = off_nc % NUM_CBLOCK # sub-block index + off_e = tl.program_id(2) # output feature block index + + n_offset = off_n * BLOCK + c_offset = off_c * CBLOCK + e_offset = off_e * E_FBLOCK + block_offset = n_offset + c_offset + + # Calculate offsets for the current batch, head, and block + q_offset = off_bh * n * d + (n_offset + c_offset) * d + o_offset = off_bh * n * e + (n_offset + c_offset) * e + e_offset + kv_offset = off_bh * NUM_BLOCK * d * e + off_n * d * e + e_offset + + # Calculate pointers to the query, output, and key-value tensors + Q_block_ptr = ( + Q + q_offset + tl.arange(0, CBLOCK)[:, None] * d + tl.arange(0, d)[None, :] + ) + O_block_ptr = ( + Out + + o_offset + + tl.arange(0, CBLOCK)[:, None] * e + + tl.arange(0, E_FBLOCK)[None, :] + ) + KV_block_ptr = ( + KV + kv_offset + tl.arange(0, d)[:, None] * e + tl.arange(0, E_FBLOCK)[None, :] + ) + + # Load the decay rate for the current head + S_block_ptr = S + off_h + s = tl.load(S_block_ptr) + + c_array = tl.arange(0, CBLOCK) + + # Load the key-value outer product for the current block + kv = tl.load(KV_block_ptr).to(tl.float32) + q_index = block_offset + tl.arange(0, CBLOCK) + + # Load query values + q = tl.load(Q_block_ptr, mask=q_index[:, None] < n, other=0.0).to(tl.float32) + + # Compute decay factors for the current sub-block + q_decay = tl.exp(-s.to(tl.float32) * (off_c * CBLOCK + c_array[:, None])) + + # Compute non-diagonal attention output + qkv_none_diag = tl.dot(q, kv) * q_decay + + # Load diagonal attention output (computed by _fwd_diag_kernel) + qkv_diag = tl.load(O_block_ptr, mask=q_index[:, None] < n, other=0.0).to(tl.float32) + + # Combine diagonal and non-diagonal attention outputs + qkv = qkv_diag + qkv_none_diag + + # Store the result + tl.store( + O_block_ptr, qkv.to(O_block_ptr.dtype.element_ty), mask=q_index[:, None] < n + ) + diff --git a/benchmarks/TritonNPUBench/vllm/_gather_block_tables_kernel/_gather_block_tables_kernel.py b/benchmarks/TritonNPUBench/vllm/_gather_block_tables_kernel/_gather_block_tables_kernel.py new file mode 100644 index 00000000..723d4138 --- /dev/null +++ b/benchmarks/TritonNPUBench/vllm/_gather_block_tables_kernel/_gather_block_tables_kernel.py @@ -0,0 +1,47 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/vllm_operator_cases/newtest_cases/test__gather_block_tables_kernel.py +# Main kernel: _gather_block_tables_kernel +# PT file: test__gather_block_tables_kernel_v2.pt + +import triton +import triton.language as tl + + +# === _gather_block_tables_kernel === +@triton.jit +def _gather_block_tables_kernel( + batch_idx_to_req_idx, # [batch_size] + src_block_table_ptrs, # [num_kv_cache_groups] + dst_block_table_ptrs, # [num_kv_cache_groups] + block_table_strides, # [num_kv_cache_groups] + num_blocks_ptr, # [num_kv_cache_groups, max_num_reqs] + num_blocks_stride, + BLOCK_SIZE: tl.constexpr, +): + # kv cache group id + group_id = tl.program_id(0) + batch_idx = tl.program_id(1) + req_idx = tl.load(batch_idx_to_req_idx + batch_idx) + + group_num_blocks_ptr = num_blocks_ptr + group_id * num_blocks_stride + num_blocks = tl.load(group_num_blocks_ptr + req_idx) + + stride = tl.load(block_table_strides + group_id) + src_block_table_ptr = _load_ptr(src_block_table_ptrs + group_id, tl.int32) + src_row_ptr = src_block_table_ptr + req_idx * stride + dst_block_table_ptr = _load_ptr(dst_block_table_ptrs + group_id, tl.int32) + dst_row_ptr = dst_block_table_ptr + batch_idx * stride + + for i in tl.range(0, num_blocks, BLOCK_SIZE): + offset = i + tl.arange(0, BLOCK_SIZE) + block_ids = tl.load(src_row_ptr + offset, mask=offset < num_blocks) + tl.store(dst_row_ptr + offset, block_ids, mask=offset < num_blocks) + + +# === _load_ptr === +@triton.jit +def _load_ptr(ptr_to_ptr, elem_dtype): + ptr = tl.load(ptr_to_ptr) + ptr = tl.cast(ptr, tl.pointer_type(elem_dtype)) + return tl.multiple_of(ptr, 16) + diff --git a/benchmarks/TritonNPUBench/vllm/_get_num_sampled_and_rejected_kernel/_get_num_sampled_and_rejected_kernel.py b/benchmarks/TritonNPUBench/vllm/_get_num_sampled_and_rejected_kernel/_get_num_sampled_and_rejected_kernel.py new file mode 100644 index 00000000..7f2b0613 --- /dev/null +++ b/benchmarks/TritonNPUBench/vllm/_get_num_sampled_and_rejected_kernel/_get_num_sampled_and_rejected_kernel.py @@ -0,0 +1,38 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/vllm_operator_cases/newtest_cases/test__get_num_sampled_and_rejected_kernel.py +# Main kernel: _get_num_sampled_and_rejected_kernel +# PT file: get_num_sampled_and_rejected_test_data_v2.pt + +import triton +import triton.language as tl + + +# === _get_num_sampled_and_rejected_kernel === +@triton.jit +def _get_num_sampled_and_rejected_kernel( + num_sampled_ptr, + num_rejected_ptr, + seq_lens_ptr, + cu_num_logits_ptr, + idx_mapping_ptr, + prefill_len_ptr, +): + batch_idx = tl.program_id(0) + req_state_idx = tl.load(idx_mapping_ptr + batch_idx) + + seq_len = tl.load(seq_lens_ptr + batch_idx) + prefill_len = tl.load(prefill_len_ptr + req_state_idx) + is_chunked_prefilling = seq_len < prefill_len + + num_sampled = tl.load(num_sampled_ptr + batch_idx) + num_sampled = tl.where(is_chunked_prefilling, 0, num_sampled) + tl.store(num_sampled_ptr + batch_idx, num_sampled) + + logits_start = tl.load(cu_num_logits_ptr + batch_idx) + logits_end = tl.load(cu_num_logits_ptr + batch_idx + 1) + num_logits = logits_end - logits_start + + num_rejected = num_logits - num_sampled + num_rejected = tl.where(is_chunked_prefilling, 0, num_rejected) + tl.store(num_rejected_ptr + batch_idx, num_rejected) + diff --git a/benchmarks/TritonNPUBench/vllm/_gumbel_sample_kernel/_gumbel_sample_kernel.py b/benchmarks/TritonNPUBench/vllm/_gumbel_sample_kernel/_gumbel_sample_kernel.py new file mode 100644 index 00000000..1db4cb55 --- /dev/null +++ b/benchmarks/TritonNPUBench/vllm/_gumbel_sample_kernel/_gumbel_sample_kernel.py @@ -0,0 +1,66 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/vllm_operator_cases/newtest_cases/test_gumbel_sample_kernel.py +# Main kernel: _gumbel_sample_kernel +# PT file: gumbel_sample_kernel_test_data_v2.pt + +import triton +import triton.language as tl + + +# === _gumbel_sample_kernel === +@triton.jit +def _gumbel_sample_kernel( + local_argmax_ptr, + local_argmax_stride, + local_max_ptr, + local_max_stride, + logits_ptr, + logits_stride, + idx_mapping_ptr, + seeds_ptr, + pos_ptr, + temp_ptr, + vocab_size, + BLOCK_SIZE: tl.constexpr, + APPLY_TEMPERATURE: tl.constexpr, +): + batch_idx = tl.program_id(0) + req_state_idx = tl.load(idx_mapping_ptr + batch_idx) + + block_idx = tl.program_id(1) + block = block_idx * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + mask = block < vocab_size + logits = tl.load( + logits_ptr + batch_idx * logits_stride + block, + mask=mask, + other=float("-inf"), + ) + logits = logits.to(tl.float32) + + temp = tl.load(temp_ptr + req_state_idx).to(tl.float32) + if temp != 0.0: + # Calculate the seed for gumbel noise. + seed = tl.load(seeds_ptr + req_state_idx) + pos = tl.load(pos_ptr + batch_idx) + gumbel_seed = tl.randint(seed, pos) + + # Generate gumbel noise. + r = tl.full([BLOCK_SIZE], 0.1, dtype=tl.float32) + gumbel_noise = -tl.log(-tl.log(r + 1e-20) + 1e-20) + gumbel_noise = gumbel_noise.to(tl.float32) + + # Apply temperature. + if APPLY_TEMPERATURE: + # NOTE(woosuk): Match the behavior of _penalties_and_temperature_kernel. + # E.g., if the kernel uses tl.div_rn, we should use tl.div_rn here too. + logits = logits / temp + + # Apply gumbel noise. + logits = tl.where(mask, logits + gumbel_noise, float("-inf")) + + idx = tl.argmax(logits, axis=0) + token_id = block_idx * BLOCK_SIZE + idx + value = tl.max(logits, axis=0) + tl.store(local_argmax_ptr + batch_idx * local_argmax_stride + block_idx, token_id) + tl.store(local_max_ptr + batch_idx * local_max_stride + block_idx, value) + diff --git a/benchmarks/TritonNPUBench/vllm/_layer_norm_fwd_1pass_kernel/_layer_norm_fwd_1pass_kernel.py b/benchmarks/TritonNPUBench/vllm/_layer_norm_fwd_1pass_kernel/_layer_norm_fwd_1pass_kernel.py new file mode 100644 index 00000000..b4bc398c --- /dev/null +++ b/benchmarks/TritonNPUBench/vllm/_layer_norm_fwd_1pass_kernel/_layer_norm_fwd_1pass_kernel.py @@ -0,0 +1,75 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/vllm_operator_cases/newtest_cases/test__layer_norm_fwd_1pass_kernel.py +# Main kernel: _layer_norm_fwd_1pass_kernel +# PT file: _layer_norm_fwd_1pass_kernel_v2.pt + +import triton +import triton.language as tl + + +# === _layer_norm_fwd_1pass_kernel === +@triton.heuristics({"HAS_BIAS": lambda args: args["B"] is not None}) +@triton.heuristics({"HAS_Z": lambda args: args["Z"] is not None}) +@triton.jit +def _layer_norm_fwd_1pass_kernel( + X, # pointer to the input + Y, # pointer to the output + W, # pointer to the weights + B, # pointer to the biases + Z, # pointer to the other branch + Mean, # pointer to the mean + Rstd, # pointer to the 1/std + stride_x_row: tl.int64, + stride_y_row: tl.int64, + stride_z_row: tl.int64, + M: tl.int64, # number of rows in X + N: tl.int64, # number of columns in X + eps, # epsilon to avoid division by zero + BLOCK_N: tl.constexpr, + HAS_BIAS: tl.constexpr, + HAS_Z: tl.constexpr, + NORM_BEFORE_GATE: tl.constexpr, + IS_RMS_NORM: tl.constexpr, +): + # Map the program id to the row of X and Y it should compute. + row = tl.program_id(0) + group = tl.program_id(1) + X += row * stride_x_row + group * N + Y += row * stride_y_row + group * N + if HAS_Z: + Z += row * stride_z_row + group * N + if not IS_RMS_NORM: + Mean += group * M + Rstd += group * M + W += group * N + if HAS_BIAS: + B += group * N + # Compute mean and variance + cols = tl.arange(0, BLOCK_N) + x = tl.load(X + cols, mask=cols < N, other=0.0).to(tl.float32) + if HAS_Z and not NORM_BEFORE_GATE: + z = tl.load(Z + cols, mask=cols < N).to(tl.float32) + x *= z * tl.sigmoid(z) + if not IS_RMS_NORM: + mean = tl.sum(x, axis=0) / N + tl.store(Mean + row, mean) + xbar = tl.where(cols < N, x - mean, 0.0) + var = tl.sum(xbar * xbar, axis=0) / N + else: + xbar = tl.where(cols < N, x, 0.0) + var = tl.sum(xbar * xbar, axis=0) / N + rstd = 1 / tl.sqrt(var + eps) + tl.store(Rstd + row, rstd) + # Normalize and apply linear transformation + mask = cols < N + w = tl.load(W + cols, mask=mask).to(tl.float32) + if HAS_BIAS: + b = tl.load(B + cols, mask=mask).to(tl.float32) + x_hat = (x - mean) * rstd if not IS_RMS_NORM else x * rstd + y = x_hat * w + b if HAS_BIAS else x_hat * w + if HAS_Z and NORM_BEFORE_GATE: + z = tl.load(Z + cols, mask=mask).to(tl.float32) + y *= z * tl.sigmoid(z) + # Write output + tl.store(Y + cols, y, mask=mask) + diff --git a/benchmarks/TritonNPUBench/vllm/_linear_attn_decode_kernel/_linear_attn_decode_kernel.py b/benchmarks/TritonNPUBench/vllm/_linear_attn_decode_kernel/_linear_attn_decode_kernel.py new file mode 100644 index 00000000..f186df6d --- /dev/null +++ b/benchmarks/TritonNPUBench/vllm/_linear_attn_decode_kernel/_linear_attn_decode_kernel.py @@ -0,0 +1,91 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/vllm_operator_cases/newtest_cases/test__linear_attn_decode_kernel.py +# Main kernel: _linear_attn_decode_kernel +# PT file: _linear_attn_decode_kernel.pt + +import triton +import triton.language as tl + + +# === _linear_attn_decode_kernel === +@triton.jit +def _linear_attn_decode_kernel( + q_ptr, + k_ptr, + v_ptr, + kv_cache_ptr, + slope_rate, + slot_idx, + output_ptr, + D: tl.constexpr, + qkv_b_stride, + qkv_h_stride, + cache_b_stride, + cache_h_stride, + cache_d0_stride, + cache_d1_stride, + BLOCK_SIZE: tl.constexpr, +): + """ + Kernel for linear attention decoding with KV cache. + + This kernel computes attention for a single token using the KV cache. + """ + pid_b = tl.program_id(0) # batch index + pid_h = tl.program_id(1) # head index + pid_d = tl.program_id(2) # dimension block index + + # Load slot index for the current batch + slot_id = tl.load(slot_idx + pid_b).to(tl.int64) + + # Skip if slot_id is -1 (padding) + if slot_id == -1: + return + + batch_id = pid_b + head_id = pid_h + + # Load decay rate for the current head + ratio = tl.load(slope_rate + pid_h) + + # Calculate offsets for dimensions + qk_d_offsets = tl.arange(0, D) + v_d_offsets = tl.arange(0, BLOCK_SIZE) + pid_d * BLOCK_SIZE + cache_d_offsets = ( + qk_d_offsets[:, None] * cache_d0_stride + v_d_offsets[None, :] * cache_d1_stride + ) + + # Calculate offsets for the current batch and head + q_offset = batch_id * qkv_b_stride + head_id * qkv_h_stride + k_offset = batch_id * qkv_b_stride + head_id * qkv_h_stride + v_offset = batch_id * qkv_b_stride + head_id * qkv_h_stride + + cache_offset = slot_id * cache_b_stride + head_id * cache_h_stride + + # Create masks for loading tensors + qk_mask = qk_d_offsets < D + v_mask = v_d_offsets < D + + # Load query, key, and value tensors + q = tl.load(q_ptr + q_offset + qk_d_offsets, mask=qk_mask, other=0.0) + k = tl.load(k_ptr + k_offset + qk_d_offsets, mask=qk_mask, other=0.0) + v = tl.load(v_ptr + v_offset + v_d_offsets, mask=v_mask, other=0.0) + + # Compute key-value outer product + kv_outer = k[:, None] * v[None, :] + kv_mask = qk_mask[:, None] & v_mask[None, :] + + # Apply decay to previous KV cache + ratio = tl.exp(-ratio) + kv_ptr = kv_cache_ptr + cache_offset + cache_d_offsets + kv_cache_old = tl.load(kv_ptr, mask=kv_mask, other=0.0) + kv_outer = kv_outer + ratio * kv_cache_old + + # Compute attention output + output = q[:, None].to(tl.float32) * kv_outer + output = tl.sum(output, axis=0) + + # Update KV cache and store output + tl.store(kv_ptr, kv_outer, mask=kv_mask) + tl.store(output_ptr + q_offset + v_d_offsets, output, mask=v_mask) + diff --git a/benchmarks/TritonNPUBench/vllm/_log_softmax_kernel/_log_softmax_kernel.py b/benchmarks/TritonNPUBench/vllm/_log_softmax_kernel/_log_softmax_kernel.py new file mode 100644 index 00000000..a4fe56ef --- /dev/null +++ b/benchmarks/TritonNPUBench/vllm/_log_softmax_kernel/_log_softmax_kernel.py @@ -0,0 +1,72 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/vllm_operator_cases/newtest_cases/test__log_softmax_kernel.py +# Main kernel: _log_softmax_kernel +# PT file: _log_softmax_kernel_v2.pt + +import triton +import triton.language as tl + + +# === _log_softmax_kernel === +@triton.jit +def _log_softmax_kernel( + input_ptr, + output_ptr, + input_row_stride, + output_row_stride, + n_cols, + BLOCK_SIZE: tl.constexpr, +): + """ + Compute log_softmax along the last dimension of a 2D tensor. + Each block handles one row of the input tensor. + """ + # Get the row index for this block + row_idx = tl.program_id(0).to(tl.int64) + + # Compute base pointers for input and output rows + row_start_ptr = input_ptr + row_idx * input_row_stride + output_row_start_ptr = output_ptr + row_idx * output_row_stride + + # Step 1: Find maximum value in the row for numerical stability + max_val = -float("inf") + for col_offset in range(0, n_cols, BLOCK_SIZE): + col_idx = col_offset + tl.arange(0, BLOCK_SIZE) + mask = col_idx < n_cols + + # Load values + vals = tl.load(row_start_ptr + col_idx, mask=mask, other=-float("inf")) + + # Update maximum + max_val = tl.max(tl.maximum(vals, max_val)) + + # Step 2: Compute sum of exp(x - max_val) + sum_exp = 0.0 + for col_offset in range(0, n_cols, BLOCK_SIZE): + col_idx = col_offset + tl.arange(0, BLOCK_SIZE) + mask = col_idx < n_cols + + # Load values + vals = tl.load(row_start_ptr + col_idx, mask=mask, other=0.0) + + # Compute exp(x - max_val) and accumulate + exp_vals = tl.exp(vals - max_val) + sum_exp += tl.sum(tl.where(mask, exp_vals, 0.0)) + + # Compute log(sum_exp) + log_sum_exp = tl.log(sum_exp) + + # Step 3: Compute final log_softmax values: x - max_val - log_sum_exp + for col_offset in range(0, n_cols, BLOCK_SIZE): + col_idx = col_offset + tl.arange(0, BLOCK_SIZE) + mask = col_idx < n_cols + + # Load values + vals = tl.load(row_start_ptr + col_idx, mask=mask) + + # Compute log_softmax + output = vals - max_val - log_sum_exp + + # Store results + tl.store(output_row_start_ptr + col_idx, output, mask=mask) + diff --git a/benchmarks/TritonNPUBench/vllm/_lora_expand_kernel/_lora_expand_kernel.py b/benchmarks/TritonNPUBench/vllm/_lora_expand_kernel/_lora_expand_kernel.py new file mode 100644 index 00000000..5c5248f4 --- /dev/null +++ b/benchmarks/TritonNPUBench/vllm/_lora_expand_kernel/_lora_expand_kernel.py @@ -0,0 +1,345 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/vllm_operator_cases/newtest_cases/test__lora_expand_kernel.py +# Main kernel: _lora_expand_kernel +# PT file: _lora_expand_kernel.pt + +import triton +import triton.language as tl + +@triton.jit +def mm_k( + a_ptr, + b_ptr, + ak_stride, + bk_stride, + offset_k, + K: tl.constexpr, + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, + BLOCK_K: tl.constexpr, + EVEN_K: tl.constexpr, + SPLIT_K: tl.constexpr, + CAST_TYPE: tl.constexpr, + b_dtype: tl.constexpr, + USE_GDC: tl.constexpr, + base_k, +): + """ + Given a_ptr and b_ptr, that identify the rows of A (m x k) and columns of + B (k x n), iterate, through the K dimension to compute the partial/complete + matrix block product. + If SPLIT_K == 1, the output m x n product is complete. + If SPLIT_K > 1, the thread block computes partial outputs. The partial + outputs are then atomically summed in the caller code. + Args: + a_ptr: Array of pointers, identifying rows of A + b_ptr: Array of pointers, identifying columns of B + ak_stride: K dimension stride of the A matrix + bk_stride: K dimension stride of the B matrix + K: Length of the K dimension + BLOCK_M: M dimension of the output block m x n + BLOCK_N: N dimension of the output block m x n + BLOCK_K: K dimension atom + EVEN_K: True if the blocks of A and B can be loaded without any + masking. + SPLIT_K: Parameter signifying parallelism in the K dimension. + CAST_TYPE: if True, cast the values from the A matrix to the B + matrix dtype. + b_dtype: datatype of the B matrix + USE_GDC: Whether to use PDL. True indicates use. + base_k: Base offset along K dimension for current SPLIT_K group + """ + accumulator = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32) + + # Step size along K for each iteration + STEP_K = BLOCK_K * SPLIT_K + + # Total number of iterations (compile-time constant) + num_iters = tl.cdiv(K, STEP_K) + + for k in range(num_iters): + # Current iteration's global K offset + iter_k = k * STEP_K + base_k + + # Check if this iteration is completely valid (no masking needed) + block_end = iter_k + BLOCK_K + + if EVEN_K: + # K is divisible by BLOCK_K, no masking ever needed + # pre-fetch lora weight + tiled_b = tl.load(b_ptr) + # if USE_GDC: + # tl.extra.cuda.gdc_wait() + tiled_a = tl.load(a_ptr) + if CAST_TYPE: + tiled_a = tiled_a.to(b_dtype) + accumulator += tl.dot(tiled_a, tiled_b) + else: + # Check if we need element-wise masking + if iter_k >= K: + # Entire block out of range, skip + pass + elif block_end <= K: + # Entire block in range, no masking needed (fast path) + tiled_b = tl.load(b_ptr) + # if USE_GDC: + # tl.extra.cuda.gdc_wait() + tiled_a = tl.load(a_ptr) + if CAST_TYPE: + tiled_a = tiled_a.to(b_dtype) + accumulator += tl.dot(tiled_a, tiled_b) + else: + # Partial block, need masking (only last iteration) + k_offsets = tl.arange(0, BLOCK_K) + mask = iter_k + k_offsets < K + tiled_b = tl.load(b_ptr, mask=mask[:, None], other=0.0) + # if USE_GDC: + # tl.extra.cuda.gdc_wait() + tiled_a = tl.load(a_ptr, mask=mask[None, :], other=0.0) + if CAST_TYPE: + tiled_a = tiled_a.to(b_dtype) + accumulator += tl.dot(tiled_a, tiled_b) + + a_ptr += STEP_K * ak_stride + b_ptr += STEP_K * bk_stride + + return accumulator + +@triton.jit +def do_expand_kernel( + pid_n, + lora_index, + slice_id, + input_ptr, + lora_ptr, + out_ptr, + N, + K, + M_LEN, + ram, # array identifying the rows of Input ptr to operate on + slice_start_loc, + # input ptr strides + input_d0_stride, + input_d1_stride, + input_d2_stride, + # lora ptr strides + ls_d0_ptr, + ls_d1_ptr, + ls_d2_ptr, + # out ptr strides + output_d0_stride, + output_d1_stride, + # constants + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, + BLOCK_K: tl.constexpr, + SAME_STRIDE: tl.constexpr, + SLICE_NUM: tl.constexpr, + EVEN_K: tl.constexpr, + CAST_TYPE: tl.constexpr, + ADD_INPUTS: tl.constexpr, + USE_GDC: tl.constexpr, +): + """ + Given an array of integers that identifies the rows of A, ram, + a lora index that identifies which LoRA to use from lora_ptr, lora_index, + a slice_id that identifies the input/output slice, + compute the matrix product and store in the appropriate output location. + Given that this is an expand kernel, we don't perform any split-K reduction + as the K dimension is assumed to be small. + """ + + # ls_d*_ptr can be either an integer or a pointer + if SAME_STRIDE: + # integer + cur_lora_d0_stride = ls_d0_ptr + cur_lora_d1_stride = ls_d1_ptr + cur_lora_d2_stride = ls_d2_ptr + else: + # pointer + cur_lora_d0_stride = tl.load(ls_d0_ptr + slice_id) + cur_lora_d1_stride = tl.load(ls_d1_ptr + slice_id) + cur_lora_d2_stride = tl.load(ls_d2_ptr + slice_id) + + # Identify the input_ptr and lora_ptr from slice_id. + if SLICE_NUM == 1: + cur_input_ptr = input_ptr + cur_lora_ptr = lora_ptr + else: + cur_input_ptr = input_ptr + slice_id * input_d0_stride + cur_lora_ptr = tl.load(lora_ptr + slice_id).to( + tl.pointer_type(out_ptr.dtype.element_ty) + ) + + # Identify the column indices of B to process. + offset_n = tl.arange(0, BLOCK_N) + pid_n * BLOCK_N + rbn = tl.max_contiguous(tl.multiple_of(offset_n % N, BLOCK_N), BLOCK_N) + + # Identify A and B block pointers + offset_k = tl.arange(0, BLOCK_K) + a_ptr = ( + cur_input_ptr + + ram[:, None] * input_d1_stride + + offset_k[None, :] * input_d2_stride + ) + b_ptr = ( + cur_lora_ptr + + cur_lora_d0_stride * lora_index + + offset_k[:, None] * cur_lora_d2_stride + + rbn[None, :] * cur_lora_d1_stride + ) + + # Compute the block matrix product. + SPLIT_K = 1 + + accumulator = mm_k( + a_ptr, + b_ptr, + input_d2_stride, + cur_lora_d2_stride, + offset_k, + K, + BLOCK_M, + BLOCK_N, + BLOCK_K, + EVEN_K, + SPLIT_K, + CAST_TYPE, + cur_lora_ptr.dtype.element_ty, + USE_GDC, + base_k=0, + ) + + tiled_c = accumulator.to(cur_lora_ptr.dtype.element_ty) + if SLICE_NUM == 1: + cur_slice_start = slice_start_loc + else: + cur_slice_start = tl.load(slice_start_loc + slice_id) + + # Identify the C output pointers to store the results of the accumulator. + offset_cn = tl.arange(0, BLOCK_N) + pid_n * BLOCK_N + cur_slice_start + offset_cm = tl.arange(0, BLOCK_M) + c_ptr = ( + out_ptr + + ram[:, None] * output_d0_stride + + offset_cn[None, :] * output_d1_stride + ) + c_mask = (offset_cm[:, None] < M_LEN) & (offset_cn[None, :] < (cur_slice_start + N)) + + if ADD_INPUTS: + tiled_out = tl.load(c_ptr, mask=c_mask) + tiled_c += tiled_out + tl.store(c_ptr, tiled_c, mask=c_mask) + +# === _lora_expand_kernel === +@triton.jit +def _lora_expand_kernel( + input_ptr, + lora_ptr, + out_ptr, + M, + N, + K, + token_indices_sorted_by_lora_ids, + num_tokens_per_lora, + lora_token_start_loc, + lora_ids, + slice_start_loc, + input_d0_stride, + input_d1_stride, + input_d2_stride, # 1 + ls_d0_ptr, + ls_d1_ptr, + ls_d2_ptr, # 1 + output_d0_stride, + output_d1_stride, # 1 + output_hs_ptr, + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, + BLOCK_K: tl.constexpr, + EVEN_K: tl.constexpr, + ADD_INPUTS: tl.constexpr, + CAST_TYPE: tl.constexpr, + SLICE_NUM: tl.constexpr, + SAME_STRIDE: tl.constexpr, + USE_GDC: tl.constexpr, + launch_pdl: tl.constexpr, +): + cta_n_num = tl.cdiv(N, BLOCK_N) + cta_m_num = tl.cdiv(M, BLOCK_M) + + pid_mn = tl.program_id(axis=0) + pid_m = pid_mn % cta_m_num + pid_n = (pid_mn // cta_m_num) % cta_n_num + + slice_id = tl.program_id(axis=1) + lora_idx = tl.program_id(axis=2) + + lora_id = tl.load(lora_ids + lora_idx) + if lora_id == -1: + # Early exit for the no-lora case. + return + + lora_m_size = tl.load(num_tokens_per_lora + lora_idx) + + cta_m_offset = pid_m * BLOCK_M + if cta_m_offset >= lora_m_size: + # Early exit CTA. + return + + # When the output dimensions of each slice are the same,cur_n=N, otherwise + # cur_n=tl.load(output_hs_ptr + slice_id), this situation exists in GQA's + # qkv linear. + curr_N = N if SAME_STRIDE else tl.load(output_hs_ptr + slice_id) + if pid_n * BLOCK_N >= curr_N: + # Early exit CTA. + return + + # num rows this CTA should process. + cta_m_len = min(BLOCK_M, lora_m_size - cta_m_offset) + + # Identify all rows that this CTA should process. + lora_m_indices_start = tl.load(lora_token_start_loc + lora_idx) + cta_lora_seq_indices = ( + token_indices_sorted_by_lora_ids + lora_m_indices_start + cta_m_offset + ) + + # Load all relevant row indices. + offset_m = tl.arange(0, BLOCK_M) % cta_m_len + ram = tl.load(cta_lora_seq_indices + offset_m) + + do_expand_kernel( + pid_n, + lora_id, + slice_id, + input_ptr, + lora_ptr, + out_ptr, + curr_N, + K, + cta_m_len, + ram, # array identifying the rows of Input ptr to operate on + slice_start_loc, + # input ptr strides + input_d0_stride, + input_d1_stride, + input_d2_stride, + # lora ptr strides + ls_d0_ptr, + ls_d1_ptr, + ls_d2_ptr, + # out ptr strides + output_d0_stride, + output_d1_stride, + # constants + BLOCK_M, + BLOCK_N, + BLOCK_K, + SAME_STRIDE, + SLICE_NUM, + EVEN_K, + CAST_TYPE, + ADD_INPUTS, + USE_GDC, + ) + diff --git a/benchmarks/TritonNPUBench/vllm/_lora_shrink_kernel/_lora_shrink_kernel.py b/benchmarks/TritonNPUBench/vllm/_lora_shrink_kernel/_lora_shrink_kernel.py new file mode 100644 index 00000000..d80a821b --- /dev/null +++ b/benchmarks/TritonNPUBench/vllm/_lora_shrink_kernel/_lora_shrink_kernel.py @@ -0,0 +1,320 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/vllm_operator_cases/newtest_cases/test__lora_shrink_kernel.py +# Main kernel: _lora_shrink_kernel +# PT file: _lora_shrink_kernel_v2.pt + +import triton +import triton.language as tl + +@triton.jit +def mm_k( + a_ptr, + b_ptr, + ak_stride, + bk_stride, + offset_k, + K: tl.constexpr, + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, + BLOCK_K: tl.constexpr, + EVEN_K: tl.constexpr, + SPLIT_K: tl.constexpr, + CAST_TYPE: tl.constexpr, + b_dtype: tl.constexpr, + USE_GDC: tl.constexpr, + base_k, +): + """ + Given a_ptr and b_ptr, that identify the rows of A (m x k) and columns of + B (k x n), iterate, through the K dimension to compute the partial/complete + matrix block product. + If SPLIT_K == 1, the output m x n product is complete. + If SPLIT_K > 1, the thread block computes partial outputs. The partial + outputs are then atomically summed in the caller code. + Args: + a_ptr: Array of pointers, identifying rows of A + b_ptr: Array of pointers, identifying columns of B + ak_stride: K dimension stride of the A matrix + bk_stride: K dimension stride of the B matrix + K: Length of the K dimension + BLOCK_M: M dimension of the output block m x n + BLOCK_N: N dimension of the output block m x n + BLOCK_K: K dimension atom + EVEN_K: True if the blocks of A and B can be loaded without any + masking. + SPLIT_K: Parameter signifying parallelism in the K dimension. + CAST_TYPE: if True, cast the values from the A matrix to the B + matrix dtype. + b_dtype: datatype of the B matrix + USE_GDC: Whether to use PDL. True indicates use. + base_k: Base offset along K dimension for current SPLIT_K group + """ + accumulator = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32) + + # Step size along K for each iteration + STEP_K = BLOCK_K * SPLIT_K + + # Total number of iterations (compile-time constant) + num_iters = tl.cdiv(K, STEP_K) + + for k in range(num_iters): + # Current iteration's global K offset + iter_k = k * STEP_K + base_k + + # Check if this iteration is completely valid (no masking needed) + block_end = iter_k + BLOCK_K + + if EVEN_K: + # K is divisible by BLOCK_K, no masking ever needed + # pre-fetch lora weight + tiled_b = tl.load(b_ptr) + # if USE_GDC: + # tl.extra.cuda.gdc_wait() + tiled_a = tl.load(a_ptr) + if CAST_TYPE: + tiled_a = tiled_a.to(b_dtype) + accumulator += tl.dot(tiled_a, tiled_b) + else: + # Check if we need element-wise masking + if iter_k >= K: + # Entire block out of range, skip + pass + elif block_end <= K: + # Entire block in range, no masking needed (fast path) + tiled_b = tl.load(b_ptr) + # if USE_GDC: + # tl.extra.cuda.gdc_wait() + tiled_a = tl.load(a_ptr) + if CAST_TYPE: + tiled_a = tiled_a.to(b_dtype) + accumulator += tl.dot(tiled_a, tiled_b) + else: + # Partial block, need masking (only last iteration) + k_offsets = tl.arange(0, BLOCK_K) + mask = iter_k + k_offsets < K + tiled_b = tl.load(b_ptr, mask=mask[:, None], other=0.0) + # if USE_GDC: + # tl.extra.cuda.gdc_wait() + tiled_a = tl.load(a_ptr, mask=mask[None, :], other=0.0) + if CAST_TYPE: + tiled_a = tiled_a.to(b_dtype) + accumulator += tl.dot(tiled_a, tiled_b) + + a_ptr += STEP_K * ak_stride + b_ptr += STEP_K * bk_stride + + return accumulator + +@triton.jit +def do_shrink_kernel( + pid_n, + pid_sk, + slice_id, + lora_index, + input_ptr, + lora_ptr, + out_ptr, + N, + K, + M_LEN, + ram, + # input strides + input_d0_stride, + input_d1_stride, + # lora strides + lora_d0_stride, + lora_d1_stride, + lora_d2_stride, + # output strides + output_d0_stride, + output_d1_stride, + output_d2_stride, + scaling, + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, + BLOCK_K: tl.constexpr, + EVEN_K: tl.constexpr, + SPLIT_K: tl.constexpr, + SLICE_NUM: tl.constexpr, + USE_GDC: tl.constexpr, +): + """ + Given an array of integers that identifies the rows of A, ram, + a lora index that identifies which LoRA to use from lora_ptr, lora_index, + a slice_id that identifies the input/output slice, compute the + matrix product and store in the appropriate output location. + """ + + # Identify the lora_ptr from slice_id. + if SLICE_NUM == 1: + # current lora ptr + cur_lora_ptr = lora_ptr + else: + # current lora ptr + cur_lora_ptr = tl.load(lora_ptr + slice_id).to( + tl.pointer_type(input_ptr.dtype.element_ty) + ) + + # Identify the column indices of B to process. + offset_n = tl.arange(0, BLOCK_N) + pid_n * BLOCK_N + rbn = tl.max_contiguous(tl.multiple_of(offset_n % N, BLOCK_N), BLOCK_N) + + # Identify A and B block pointers + offset_k = pid_sk * BLOCK_K + tl.arange(0, BLOCK_K) + a_ptr = ( + input_ptr + ram[:, None] * input_d0_stride + offset_k[None, :] * input_d1_stride + ) + b_ptr = ( + cur_lora_ptr + + lora_d0_stride * lora_index + + rbn[None, :] * lora_d1_stride + + offset_k[:, None] * lora_d2_stride + ) + + # Compute partial/complete block matrix product. + accumulator = mm_k( + a_ptr, + b_ptr, + input_d1_stride, + lora_d2_stride, + offset_k, + K, + BLOCK_M, + BLOCK_N, + BLOCK_K, + EVEN_K, + SPLIT_K, + False, + cur_lora_ptr.dtype.element_ty, + False, # USE_GDC is always False in shrink kernel + base_k=pid_sk * BLOCK_K, + ) + # GDC launch dependents hints the runtime system to launch dependent kernels. + # if USE_GDC: + # tl.extra.cuda.gdc_launch_dependents() + # Identify the C output pointers to store the results of the accumulator. + offset_cn = tl.arange(0, BLOCK_N) + pid_n * BLOCK_N + offset_cm = tl.arange(0, BLOCK_M) + cur_out_ptr = out_ptr if SLICE_NUM == 1 else out_ptr + slice_id * output_d0_stride + c_ptr = ( + cur_out_ptr + + ram[:, None] * output_d1_stride + + offset_cn[None, :] * output_d2_stride + ) + c_mask = (offset_cm[:, None] < M_LEN) & (offset_cn[None, :] < N) + accumulator *= scaling + + # handles write-back with reduction-splitting + if SPLIT_K == 1: + tl.store(c_ptr, accumulator, mask=c_mask) + else: + tl.atomic_add(c_ptr, accumulator, mask=c_mask, sem="relaxed") + +# === _lora_shrink_kernel === +@triton.jit +def _lora_shrink_kernel( + input_ptr, + lora_ptr, + out_ptr, + M, + N, + K, + token_indices_sorted_by_lora_ids, + num_tokens_per_lora, + lora_token_start_loc, + lora_ids, + scaling, + input_d0_stride, + input_d1_stride, + lora_d0_stride, + lora_d1_stride, + lora_d2_stride, + output_d0_stride, + output_d1_stride, + output_d2_stride, + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, + BLOCK_K: tl.constexpr, + EVEN_K: tl.constexpr, + SPLIT_K: tl.constexpr, + GROUP_SIZE_M: tl.constexpr, + SLICE_NUM: tl.constexpr, + USE_GDC: tl.constexpr, + launch_pdl: tl.constexpr, +): + cta_n_num = tl.cdiv(N, BLOCK_N) + cta_m_num = tl.cdiv(M, BLOCK_M) + + pid_sk_m_n = tl.program_id(axis=0) + pid_sk = pid_sk_m_n % SPLIT_K + + pid_m_n = pid_sk_m_n // SPLIT_K + num_pid_in_group = GROUP_SIZE_M * cta_n_num + group_id = pid_m_n // num_pid_in_group + first_pid_m = group_id * GROUP_SIZE_M + group_size_m = min(cta_m_num - first_pid_m, GROUP_SIZE_M) + + # Column-major ordering within groups for better cache reuse + pid_m = first_pid_m + ((pid_m_n % num_pid_in_group) % group_size_m) + pid_n = (pid_m_n % num_pid_in_group) // group_size_m + + slice_id = tl.program_id(axis=1) + lora_idx = tl.program_id(axis=2) + + lora_id = tl.load(lora_ids + lora_idx) + if lora_id == -1: + # Early exit for the no-lora case. + return + + lora_m_size = tl.load(num_tokens_per_lora + lora_idx) + + cta_m_offset = pid_m * BLOCK_M + if cta_m_offset >= lora_m_size: + # Early exit CTA. + return + + # num rows this CTA should process. + cta_m_len = min(BLOCK_M, lora_m_size - cta_m_offset) + + # Identify all rows that this CTA should process. + lora_m_indices_start = tl.load(lora_token_start_loc + lora_idx) + cta_lora_seq_indices = ( + token_indices_sorted_by_lora_ids + lora_m_indices_start + cta_m_offset + ) + # Load all relevant row indices. + offset_m = tl.arange(0, BLOCK_M) % cta_m_len + ram = tl.load(cta_lora_seq_indices + offset_m) + + do_shrink_kernel( + pid_n, + pid_sk, + slice_id, + lora_id, + input_ptr, + lora_ptr, + out_ptr, + N, + K, + cta_m_len, + ram, # array identifying the rows of Input ptr to operate on + # input strides + input_d0_stride, + input_d1_stride, + # lora strides + lora_d0_stride, + lora_d1_stride, + lora_d2_stride, + # output strides + output_d0_stride, + output_d1_stride, + output_d2_stride, + scaling, + BLOCK_M, + BLOCK_N, + BLOCK_K, + EVEN_K, + SPLIT_K, + SLICE_NUM, + USE_GDC, + ) + diff --git a/benchmarks/TritonNPUBench/vllm/_min_p_kernel/_min_p_kernel.py b/benchmarks/TritonNPUBench/vllm/_min_p_kernel/_min_p_kernel.py new file mode 100644 index 00000000..ba491c8b --- /dev/null +++ b/benchmarks/TritonNPUBench/vllm/_min_p_kernel/_min_p_kernel.py @@ -0,0 +1,45 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/vllm_operator_cases/newtest_cases/test_min_p_kernel.py +# Main kernel: _min_p_kernel +# PT file: min_p_kernel_test_v2.pt + +import triton +import triton.language as tl + + +# === _min_p_kernel === +@triton.jit +def _min_p_kernel( + logits_ptr, + logits_stride, + idx_mapping_ptr, + min_p_ptr, + vocab_size, + BLOCK_SIZE: tl.constexpr, +): + req_idx = tl.program_id(0) + req_state_idx = tl.load(idx_mapping_ptr + req_idx) + min_p = tl.load(min_p_ptr + req_state_idx).to(tl.float32) + if min_p == 0.0: + return + + max_val = float("-inf") + for i in range(0, vocab_size, BLOCK_SIZE): + block = i + tl.arange(0, BLOCK_SIZE) + mask = block < vocab_size + logits = tl.load( + logits_ptr + req_idx * logits_stride + block, mask=mask, other=float("-inf") + ) + max_val = tl.max(tl.maximum(logits, max_val)) + max_val = max_val.to(tl.float32) # type: ignore + + threshold = max_val + tl.log(min_p) + for i in range(0, vocab_size, BLOCK_SIZE): + block = i + tl.arange(0, BLOCK_SIZE) + mask = block < vocab_size + logits = tl.load( + logits_ptr + req_idx * logits_stride + block, mask=mask, other=float("-inf") + ) + logits = tl.where(logits < threshold, float("-inf"), logits) + tl.store(logits_ptr + req_idx * logits_stride + block, logits, mask=mask) + diff --git a/benchmarks/TritonNPUBench/vllm/_num_nans_kernel/_num_nans_kernel.py b/benchmarks/TritonNPUBench/vllm/_num_nans_kernel/_num_nans_kernel.py new file mode 100644 index 00000000..2815ca86 --- /dev/null +++ b/benchmarks/TritonNPUBench/vllm/_num_nans_kernel/_num_nans_kernel.py @@ -0,0 +1,32 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/vllm_operator_cases/newtest_cases/test__num_nans_kernel.py +# Main kernel: _num_nans_kernel +# PT file: None + +import triton +import triton.language as tl +from triton.language.extra.cann import libdevice + +# === _num_nans_kernel === +@triton.jit +def _num_nans_kernel( + logits_ptr, + logits_stride, + num_nans_ptr, + vocab_size, + BLOCK_SIZE: tl.constexpr, +): + req_idx = tl.program_id(0) + num_nans = 0 + for i in range(0, vocab_size, BLOCK_SIZE): + block = i + tl.arange(0, BLOCK_SIZE) + mask = block < vocab_size + logits = tl.load( + logits_ptr + req_idx * logits_stride + block, mask=mask, other=0 + ) + logits = logits.to(tl.float32) + is_nan = libdevice.isnan(logits).to(tl.int32) + num_nans += tl.sum(is_nan) + tl.store(num_nans_ptr + req_idx, num_nans) + tl.device_print('num_nans = ', num_nans) + diff --git a/benchmarks/TritonNPUBench/vllm/_pack_seq_kernel/_pack_seq_kernel.py b/benchmarks/TritonNPUBench/vllm/_pack_seq_kernel/_pack_seq_kernel.py new file mode 100644 index 00000000..9562dd5f --- /dev/null +++ b/benchmarks/TritonNPUBench/vllm/_pack_seq_kernel/_pack_seq_kernel.py @@ -0,0 +1,57 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/vllm_operator_cases/newtest_cases/test__pack_seq_kernel.py +# Main kernel: _pack_seq_kernel +# PT file: pack_seq_test_v2.pt + +import triton +import triton.language as tl + + +# === _pack_seq_kernel === +@triton.jit +def _pack_seq_kernel( + x_ptr, # [N, D] + out_ptr, # [B, Lmax, D] + lengths_ptr, # *i32, [B] + N: tl.constexpr, + D: tl.constexpr, + Lmax: tl.constexpr, + PAD_VALUE: tl.constexpr, + BLOCK_T: tl.constexpr, # timesteps per program + BLOCK_D: tl.constexpr, # features per program +): + pid_b = tl.program_id(0) # batch id + pid_t = tl.program_id(1) # block over time dimension + pid_d = tl.program_id(2) # block over feature dimension + off_t = pid_t * BLOCK_T + tl.arange(0, BLOCK_T) # [BLOCK_T] + off_d = pid_d * BLOCK_D + tl.arange(0, BLOCK_D) # [BLOCK_D] + + # Compute start index and sequence length from cumulative lengths + in_start = 0 + for i in range(pid_b): + in_start += tl.load(lengths_ptr + i) + seq_len = tl.load(lengths_ptr + pid_b) + + # valid time positions for this block + t_mask = off_t < Lmax + + # compute input row indices for valid (b, t) + in_row = in_start + off_t + valid_row = (off_t < seq_len) & t_mask + + # Pointers + # x_ptr: row-major [N, D] + x_row_ptr = x_ptr + in_row[:, None] * D + off_d[None, :] + + # out_ptr: row-major [B, Lmax, D] + out_row_ptr = out_ptr + (pid_b * Lmax + off_t)[:, None] * D + off_d[None, :] + + # Initialize with PAD (cast will occur as needed based on out_ptr dtype) + d_mask = off_d[None, :] < D + pad_vals = tl.full([BLOCK_T, BLOCK_D], PAD_VALUE, tl.float32) + tl.store(out_row_ptr, pad_vals, mask=t_mask[:, None] & d_mask) + + # Load & write only where within seq_len + x_vals = tl.load(x_row_ptr, mask=valid_row[:, None] & d_mask) + tl.store(out_row_ptr, x_vals, mask=valid_row[:, None] & d_mask) + diff --git a/benchmarks/TritonNPUBench/vllm/_penalties_and_temperature_kernel/_penalties_and_temperature_kernel.py b/benchmarks/TritonNPUBench/vllm/_penalties_and_temperature_kernel/_penalties_and_temperature_kernel.py new file mode 100644 index 00000000..047e9625 --- /dev/null +++ b/benchmarks/TritonNPUBench/vllm/_penalties_and_temperature_kernel/_penalties_and_temperature_kernel.py @@ -0,0 +1,85 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/vllm_operator_cases/newtest_cases/test__penalties_and_temperature_kernel.py +# Main kernel: _penalties_and_temperature_kernel +# PT file: penalties_and_temperature_kernel_test_data_v2.pt + +import triton +import triton.language as tl + + +# === _penalties_and_temperature_kernel === +@triton.jit +def _penalties_and_temperature_kernel( + logits_ptr, + logits_stride, + idx_mapping_ptr, + repetition_penalty_ptr, + frequency_penalty_ptr, + presence_penalty_ptr, + temperature_ptr, + prompt_bin_mask_ptr, + prompt_bin_mask_stride, + output_bin_counts_ptr, + output_bin_counts_stride, + vocab_size, + BLOCK_SIZE: tl.constexpr, +): + batch_idx = tl.program_id(0) + req_state_idx = tl.load(idx_mapping_ptr + batch_idx) + rep_penalty = tl.load(repetition_penalty_ptr + req_state_idx) + freq_penalty = tl.load(frequency_penalty_ptr + req_state_idx) + pres_penalty = tl.load(presence_penalty_ptr + req_state_idx) + temperature = tl.load(temperature_ptr + req_state_idx) + temperature = tl.where(temperature == 0.0, 1.0, temperature) + + use_rep_penalty = rep_penalty != 1.0 + use_freq_penalty = freq_penalty != 0.0 + use_pres_penalty = pres_penalty != 0.0 + use_penalty = use_rep_penalty or (use_freq_penalty or use_pres_penalty) + use_temperature = temperature != 1.0 + if not (use_penalty or use_temperature): + # Early return to avoid loading logits. + return + + block_idx = tl.program_id(1) + block = block_idx * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + mask = block < vocab_size + logits = tl.load(logits_ptr + batch_idx * logits_stride + block, mask=mask) + logits = logits.to(tl.float32) + + if use_penalty: + output_bin_counts = tl.load( + output_bin_counts_ptr + req_state_idx * output_bin_counts_stride + block, + mask=mask, + ) + output_bin_mask = output_bin_counts > 0 + + # Apply repetition penalties. + if use_rep_penalty: + packed_block = block_idx * BLOCK_SIZE // 32 + tl.arange(0, BLOCK_SIZE // 32) + packed_mask = tl.load( + prompt_bin_mask_ptr + + req_state_idx * prompt_bin_mask_stride + + packed_block, + mask=packed_block < tl.cdiv(vocab_size, 32), + ) + prompt_bin_mask = (packed_mask[:, None] >> (tl.arange(0, 32)[None, :])) & 1 + prompt_bin_mask = prompt_bin_mask.to(tl.int1) + prompt_bin_mask = prompt_bin_mask.reshape(BLOCK_SIZE) + + # If token appears in prompt or output, apply, otherwise use 1.0 for no-op. + scale = tl.where(prompt_bin_mask | output_bin_mask, rep_penalty, 1.0) + # If logits are positive, divide by penalty, otherwise multiply by penalty. + logits *= tl.where(logits > 0, 1.0 / scale, scale) + + # Apply frequency penalties. + logits -= freq_penalty * output_bin_counts + # Apply presence penalties. + logits -= pres_penalty * output_bin_mask + + # Apply temperature. + logits = logits / temperature + + # Store back to logits. + tl.store(logits_ptr + batch_idx * logits_stride + block, logits, mask=mask) + diff --git a/benchmarks/TritonNPUBench/vllm/_per_token_group_quant_fp8/_per_token_group_quant_fp8.py b/benchmarks/TritonNPUBench/vllm/_per_token_group_quant_fp8/_per_token_group_quant_fp8.py new file mode 100644 index 00000000..86109d7c --- /dev/null +++ b/benchmarks/TritonNPUBench/vllm/_per_token_group_quant_fp8/_per_token_group_quant_fp8.py @@ -0,0 +1,63 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/vllm_operator_cases/newtest_cases/test__per_token_group_quant_fp8.py +# Main kernel: _per_token_group_quant_fp8 +# PT file: _per_token_group_quant_fp8_v2.pt + +import triton +import triton.language as tl + + +# === _per_token_group_quant_fp8 === +@triton.jit +def _per_token_group_quant_fp8( + # Pointers to inputs and output + y_ptr, + y_q_ptr, + y_s_ptr, + group_size, + # Num columns of y + y_num_columns, + y_row_stride, + # Avoid to divide zero + eps, + # Information for float8 + fp8_min, + fp8_max, + use_ue8m0: tl.constexpr, + # Meta-parameters + BLOCK: tl.constexpr, +): + """A Triton-accelerated function to perform per-token-group + quantization on a tensor. + This function converts the tensor values into float8 values. + """ + groups_per_row = y_num_columns // group_size + + # Map the program id to the row of X and Y it should compute. + g_id = tl.program_id(0) + row = g_id // groups_per_row + row_g_id = g_id % groups_per_row + + # Ensure offset calculations use int64 to prevent overflow + y_ptr_offset = (row.to(tl.int64) * y_row_stride) + ( + row_g_id.to(tl.int64) * group_size + ) + y_ptr += y_ptr_offset + + y_q_ptr_offset = g_id.to(tl.int64) * group_size + y_q_ptr += y_q_ptr_offset + y_s_ptr += g_id + + cols = tl.arange(0, BLOCK) # N <= BLOCK + mask = cols < group_size + + y = tl.load(y_ptr + cols, mask=mask, other=0.0).to(tl.float32) + # Quant + _absmax = tl.maximum(tl.max(tl.abs(y)), eps) + scale_raw = _absmax / fp8_max + y_s = tl.math.exp2(tl.ceil(tl.log2(scale_raw))) if use_ue8m0 else scale_raw + y_q = tl.clamp(y / y_s, fp8_min, fp8_max).to(y_q_ptr.dtype.element_ty) + + tl.store(y_q_ptr + cols, y_q, mask=mask) + tl.store(y_s_ptr, y_s) + diff --git a/benchmarks/TritonNPUBench/vllm/_per_token_group_quant_fp8_colmajor/_per_token_group_quant_fp8_colmajor.py b/benchmarks/TritonNPUBench/vllm/_per_token_group_quant_fp8_colmajor/_per_token_group_quant_fp8_colmajor.py new file mode 100644 index 00000000..f76fd890 --- /dev/null +++ b/benchmarks/TritonNPUBench/vllm/_per_token_group_quant_fp8_colmajor/_per_token_group_quant_fp8_colmajor.py @@ -0,0 +1,73 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/vllm_operator_cases/newtest_cases/test__per_token_group_quant_fp8_colmajor.py +# Main kernel: _per_token_group_quant_fp8_colmajor +# PT file: _per_token_group_quant_fp8_colmajor_v2.pt + +import triton +import triton.language as tl + + +# === _per_token_group_quant_fp8_colmajor === +@triton.jit +def _per_token_group_quant_fp8_colmajor( + # Pointers to inputs and output + y_ptr, + y_q_ptr, + y_s_ptr, + group_size, + # Num columns of y + y_num_columns, + y_row_stride, + # Stride from one column to the next of y_s + y_s_col_stride, + # Avoid to divide zero + eps, + # Information for float8 + fp8_min, + fp8_max, + use_ue8m0: tl.constexpr, + # Meta-parameters + BLOCK: tl.constexpr, +): + """A Triton-accelerated function to perform per-token-group + quantization on a tensor. + This function converts the tensor values into float8 values. + """ + groups_per_row = y_num_columns // group_size + + # Map the program id to the row of X and Y it should compute. + g_id = tl.program_id(0) + row = g_id // groups_per_row + row_g_id = g_id % groups_per_row + + # Ensure offset calculations use int64 to prevent overflow + y_ptr_offset = (row.to(tl.int64) * y_row_stride) + ( + row_g_id.to(tl.int64) * group_size + ) + y_ptr += y_ptr_offset + + y_q_ptr_offset = g_id.to(tl.int64) * group_size + y_q_ptr += y_q_ptr_offset + + # Convert g_id the flattened block coordinate to 2D so we can index + # into the output y_scales matrix + blocks_per_row = y_num_columns // group_size + scale_col = g_id % blocks_per_row + scale_row = g_id // blocks_per_row + # Ensure offset calculation uses int64 for y_s_ptr + y_s_ptr_offset = (scale_col.to(tl.int64) * y_s_col_stride) + scale_row.to(tl.int64) + y_s_ptr += y_s_ptr_offset + + cols = tl.arange(0, BLOCK) # group_size <= BLOCK + mask = cols < group_size + + y = tl.load(y_ptr + cols, mask=mask, other=0.0).to(tl.float32) + # Quant + _absmax = tl.maximum(tl.max(tl.abs(y)), eps) + scale_raw = _absmax / fp8_max + y_s = tl.math.exp2(tl.ceil(tl.log2(scale_raw))) if use_ue8m0 else scale_raw + y_q = tl.clamp(y / y_s, fp8_min, fp8_max).to(y_q_ptr.dtype.element_ty) + + tl.store(y_q_ptr + cols, y_q, mask=mask) + tl.store(y_s_ptr, y_s) + diff --git a/benchmarks/TritonNPUBench/vllm/_per_token_group_quant_int8/_per_token_group_quant_int8.py b/benchmarks/TritonNPUBench/vllm/_per_token_group_quant_int8/_per_token_group_quant_int8.py new file mode 100644 index 00000000..08763ec2 --- /dev/null +++ b/benchmarks/TritonNPUBench/vllm/_per_token_group_quant_int8/_per_token_group_quant_int8.py @@ -0,0 +1,51 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/vllm_operator_cases/newtest_cases/test__per_token_group_quant_int8.py +# Main kernel: _per_token_group_quant_int8 +# PT file: _per_token_group_quant_int8_v2.pt + +import triton +import triton.language as tl + + +# === _per_token_group_quant_int8 === +@triton.jit +def _per_token_group_quant_int8( + # Pointers to inputs and output + y_ptr, + y_q_ptr, + y_s_ptr, + # Stride of input + y_stride, + # Columns of input + N, + # Avoid to divide zero + eps, + # Information for int8 + int8_min, + int8_max, + # Meta-parameters + BLOCK: tl.constexpr, +): + """A Triton-accelerated function to perform per-token-group + quantization on a tensor. + + This function converts the tensor values into int8 values. + """ + # Map the program id to the row of X and Y it should compute. + g_id = tl.program_id(0) + y_ptr += g_id * y_stride + y_q_ptr += g_id * y_stride + y_s_ptr += g_id + + cols = tl.arange(0, BLOCK) # N <= BLOCK + mask = cols < N + + y = tl.load(y_ptr + cols, mask=mask, other=0.0).to(tl.float32) + # Quant + _absmax = tl.maximum(tl.max(tl.abs(y)), eps) + y_s = _absmax / int8_max + y_q = tl.clamp(y / y_s, int8_min, int8_max).to(y_q_ptr.dtype.element_ty) + + tl.store(y_q_ptr + cols, y_q, mask=mask) + tl.store(y_s_ptr, y_s) + diff --git a/benchmarks/TritonNPUBench/vllm/_per_token_quant_int8/_per_token_quant_int8.py b/benchmarks/TritonNPUBench/vllm/_per_token_quant_int8/_per_token_quant_int8.py new file mode 100644 index 00000000..ba227150 --- /dev/null +++ b/benchmarks/TritonNPUBench/vllm/_per_token_quant_int8/_per_token_quant_int8.py @@ -0,0 +1,41 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/vllm_operator_cases/newtest_cases/test__per_token_quant_int8.py +# Main kernel: _per_token_quant_int8 +# PT file: _per_token_quant_int8_v2.pt + +import triton +import triton.language as tl + + +# === _per_token_quant_int8 === +@triton.jit +def _per_token_quant_int8( + x_ptr, + xq_ptr, + scale_ptr, + stride_x, + stride_xq, + N, + BLOCK: tl.constexpr, +): + # Adapted from https://github.com/InternLM/lmdeploy/blob/086481ed84b59bee3b8e4274e5fc69620040c048/lmdeploy/pytorch/kernels/cuda/w8a8_triton_kernels.py#L282 + row_id = tl.program_id(0) + + cols = tl.arange(0, BLOCK) + mask = cols < N + + x = tl.load(x_ptr + row_id * stride_x + cols, mask=mask, other=0.0).to(tl.float32) + absmax = tl.maximum(tl.max(tl.abs(x)), 1e-10) + scale_x = absmax / 127 + x_q = x * (127 / absmax) + x_q = round_int8(x_q) + + tl.store(xq_ptr + row_id * stride_xq + cols, x_q, mask=mask) + tl.store(scale_ptr + row_id, scale_x) + + +# === round_int8 === +@triton.jit +def round_int8(x): + return tl.extra.cann.libdevice.round(x).to(tl.int8) + diff --git a/benchmarks/TritonNPUBench/vllm/_post_update_kernel/_post_update_kernel.py b/benchmarks/TritonNPUBench/vllm/_post_update_kernel/_post_update_kernel.py new file mode 100644 index 00000000..98c5fdab --- /dev/null +++ b/benchmarks/TritonNPUBench/vllm/_post_update_kernel/_post_update_kernel.py @@ -0,0 +1,51 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/vllm_operator_cases/newtest_cases/test__post_update_kernel.py +# Main kernel: _post_update_kernel +# PT file: post_update_kernel_test_data.pt + +import triton +import triton.language as tl + + +# === _post_update_kernel === +@triton.jit +def _post_update_kernel( + idx_mapping_ptr, + num_computed_tokens_ptr, + last_sampled_tokens_ptr, + output_bin_counts_ptr, + output_bin_counts_stride, + sampled_tokens_ptr, + sampled_tokens_stride, + num_sampled_ptr, + num_rejected_ptr, + query_start_loc_ptr, +): + req_id = tl.program_id(0) + req_state_idx = tl.load(idx_mapping_ptr + req_id) + + num_sampled = tl.load(num_sampled_ptr + req_id) + if num_sampled > 0: + token_id = tl.load( + sampled_tokens_ptr + req_id * sampled_tokens_stride + num_sampled - 1 + ) + tl.store(last_sampled_tokens_ptr + req_state_idx, token_id) + + for i in range(num_sampled): + token_id = tl.load(sampled_tokens_ptr + req_id * sampled_tokens_stride + i) + token_ptr = ( + output_bin_counts_ptr + req_state_idx * output_bin_counts_stride + token_id + ) + count = tl.load(token_ptr) + count += 1 + tl.store(token_ptr, count) + + query_start = tl.load(query_start_loc_ptr + req_id) + query_end = tl.load(query_start_loc_ptr + req_id + 1) + query_len = query_end - query_start + num_rejected = tl.load(num_rejected_ptr + req_id) + + num_computed = tl.load(num_computed_tokens_ptr + req_state_idx) + num_computed += query_len - num_rejected + tl.store(num_computed_tokens_ptr + req_state_idx, num_computed) + diff --git a/benchmarks/TritonNPUBench/vllm/_prepare_eagle_docode_kernel/_prepare_eagle_docode_kernel.py b/benchmarks/TritonNPUBench/vllm/_prepare_eagle_docode_kernel/_prepare_eagle_docode_kernel.py new file mode 100644 index 00000000..cbe4bdd9 --- /dev/null +++ b/benchmarks/TritonNPUBench/vllm/_prepare_eagle_docode_kernel/_prepare_eagle_docode_kernel.py @@ -0,0 +1,78 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/vllm_operator_cases/newtest_cases/test__prepare_eagle_docode_kernel.py +# Main kernel: _prepare_eagle_docode_kernel +# PT file: prepare_eagle_docode_kernel_test_v2.pt + +import triton +import triton.language as tl + + +# === _prepare_eagle_docode_kernel === +@triton.jit +def _prepare_eagle_docode_kernel( + draft_tokens_ptr, + output_hidden_states_ptr, + output_hidden_states_stride, + last_token_indices_ptr, + target_seq_lens_ptr, + num_rejected_ptr, + input_ids_ptr, + positions_ptr, + input_hidden_states_ptr, + input_hidden_states_stride, + query_start_loc_ptr, + seq_lens_ptr, + hidden_size, + max_model_len, + max_num_reqs, + BLOCK_SIZE: tl.constexpr, +): + req_idx = tl.program_id(0) + num_reqs = tl.num_programs(0) - 1 + if req_idx == num_reqs: + # Compute query_start_loc. Pad it with the last query_start_loc + # for CUDA graphs. + for i in range(0, max_num_reqs + 1, BLOCK_SIZE): + block = i + tl.arange(0, BLOCK_SIZE) + q = tl.where(block < num_reqs, block, num_reqs) + mask = block < max_num_reqs + 1 + tl.store(query_start_loc_ptr + block, q, mask=mask) + # Pad seq_lens for CUDA graphs. + for i in range(req_idx, max_num_reqs, BLOCK_SIZE): + block = i + tl.arange(0, BLOCK_SIZE) + mask = block < max_num_reqs + tl.store(seq_lens_ptr + block, 0, mask=mask) + return + + # draft token -> input id. + draft_token = tl.load(draft_tokens_ptr + req_idx) + tl.store(input_ids_ptr + req_idx, draft_token) + + # output hidden states -> input hidden states. + src_idx = tl.load(last_token_indices_ptr + req_idx) + for i in range(0, hidden_size, BLOCK_SIZE): + block = i + tl.arange(0, BLOCK_SIZE) + mask = block < hidden_size + output_hidden_states = tl.load( + output_hidden_states_ptr + src_idx * output_hidden_states_stride + block, + mask=mask, + ) + tl.store( + input_hidden_states_ptr + req_idx * input_hidden_states_stride + block, + output_hidden_states, + mask=mask, + ) + + # Compute position and seq_lens. + # NOTE(woosuk): To prevent out-of-range access, we clamp these values + # if they reach the max model length. + position = tl.load(positions_ptr + req_idx) + position = tl.minimum(position + 1, max_model_len - 1) + tl.store(positions_ptr + req_idx, position) + + target_seq_len = tl.load(target_seq_lens_ptr + req_idx) + num_rejected = tl.load(num_rejected_ptr + req_idx) + seq_len = target_seq_len - num_rejected + seq_len = tl.minimum(seq_len + 1, max_model_len) + tl.store(seq_lens_ptr + req_idx, seq_len) + diff --git a/benchmarks/TritonNPUBench/vllm/_prepare_eagle_inputs_kernel/_prepare_eagle_inputs_kernel.py b/benchmarks/TritonNPUBench/vllm/_prepare_eagle_inputs_kernel/_prepare_eagle_inputs_kernel.py new file mode 100644 index 00000000..84f67e1e --- /dev/null +++ b/benchmarks/TritonNPUBench/vllm/_prepare_eagle_inputs_kernel/_prepare_eagle_inputs_kernel.py @@ -0,0 +1,59 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/vllm_operator_cases/newtest_cases/test__prepare_eagle_inputs_kernel.py +# Main kernel: _prepare_eagle_inputs_kernel +# PT file: prepare_eagle_inputs_kernel_test_data_v2.pt + +import triton +import triton.language as tl + + +# === _prepare_eagle_inputs_kernel === +@triton.jit +def _prepare_eagle_inputs_kernel( + last_token_indices_ptr, + eagle_input_ids_ptr, + eagle_positions_ptr, + target_input_ids_ptr, + target_positions_ptr, + last_sampled_ptr, + next_prefill_tokens_ptr, + num_sampled_ptr, + num_rejected_ptr, + query_start_loc_ptr, + BLOCK_SIZE: tl.constexpr, +): + batch_idx = tl.program_id(0) + query_start = tl.load(query_start_loc_ptr + batch_idx) + query_end = tl.load(query_start_loc_ptr + batch_idx + 1) + query_len = query_end - query_start + + # Get the true query length and next token after accounting for rejected tokens. + num_rejected = tl.load(num_rejected_ptr + batch_idx) + query_len -= num_rejected + + num_sampled = tl.load(num_sampled_ptr + batch_idx) + if num_sampled > 0: + next_token = tl.load(last_sampled_ptr + batch_idx).to(tl.int32) + else: + # Chunked prefilling. + # Get the next prefill token. + next_token = tl.load(next_prefill_tokens_ptr + batch_idx) + + # Shift target_input_ids by one. + for i in range(1, query_len, BLOCK_SIZE): + block = i + tl.arange(0, BLOCK_SIZE) + mask = block < query_len + input_ids = tl.load(target_input_ids_ptr + query_start + block, mask=mask) + tl.store(eagle_input_ids_ptr + query_start + block - 1, input_ids, mask=mask) + + last_token_index = query_start + query_len - 1 + tl.store(last_token_indices_ptr + batch_idx, last_token_index) + tl.store(eagle_input_ids_ptr + last_token_index, next_token) + + # Copy positions. + for i in range(0, query_len, BLOCK_SIZE): + block = i + tl.arange(0, BLOCK_SIZE) + mask = block < query_len + target_pos = tl.load(target_positions_ptr + query_start + block, mask=mask) + tl.store(eagle_positions_ptr + query_start + block, target_pos, mask=mask) + diff --git a/benchmarks/TritonNPUBench/vllm/_prepare_mrope_positions_kernel/_prepare_mrope_positions_kernel.py b/benchmarks/TritonNPUBench/vllm/_prepare_mrope_positions_kernel/_prepare_mrope_positions_kernel.py new file mode 100644 index 00000000..4ed91abc --- /dev/null +++ b/benchmarks/TritonNPUBench/vllm/_prepare_mrope_positions_kernel/_prepare_mrope_positions_kernel.py @@ -0,0 +1,60 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/vllm_operator_cases/newtest_cases/test__prepare_mrope_positions_kernel.py +# Main kernel: _prepare_mrope_positions_kernel +# PT file: prepare_mrope_positions_test_v2.pt + +import triton +import triton.language as tl + + +# === _prepare_mrope_positions_kernel === +@triton.jit +def _prepare_mrope_positions_kernel( + mrope_positions_ptr, + mrope_positions_stride, + prefill_mrope_positions_ptr, + prefill_mrope_positions_stride0, + prefill_mrope_positions_stride1, + prefill_mrope_delta_ptr, + idx_mapping_ptr, + query_start_loc_ptr, + prefill_lens_ptr, + num_computed_tokens_ptr, + BLOCK_SIZE: tl.constexpr, +): + batch_idx = tl.program_id(0) + req_state_idx = tl.load(idx_mapping_ptr + batch_idx) + + prefill_len = tl.load(prefill_lens_ptr + req_state_idx) + num_computed = tl.load(num_computed_tokens_ptr + req_state_idx) + is_prefill = num_computed < prefill_len + + query_start = tl.load(query_start_loc_ptr + batch_idx) + query_end = tl.load(query_start_loc_ptr + batch_idx + 1) + query_len = query_end - query_start + + mrope_delta = tl.load(prefill_mrope_delta_ptr + req_state_idx) + for i in range(0, query_len, BLOCK_SIZE): + block = i + tl.arange(0, BLOCK_SIZE) + mask = block < query_len + orig_pos = num_computed + block + + for j in tl.static_range(3): + if is_prefill: + # Read from pre-computed M-RoPE positions. + pos = tl.load( + prefill_mrope_positions_ptr + + req_state_idx * prefill_mrope_positions_stride0 + + j * prefill_mrope_positions_stride1 + + orig_pos, + mask=mask, + ) + else: + # Apply M-RoPE delta. + pos = orig_pos + mrope_delta + tl.store( + mrope_positions_ptr + j * mrope_positions_stride + query_start + block, + pos, + mask=mask, + ) + diff --git a/benchmarks/TritonNPUBench/vllm/_prepare_pos_seq_lens_kernel/_prepare_pos_seq_lens_kernel.py b/benchmarks/TritonNPUBench/vllm/_prepare_pos_seq_lens_kernel/_prepare_pos_seq_lens_kernel.py new file mode 100644 index 00000000..f67e1d80 --- /dev/null +++ b/benchmarks/TritonNPUBench/vllm/_prepare_pos_seq_lens_kernel/_prepare_pos_seq_lens_kernel.py @@ -0,0 +1,46 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/vllm_operator_cases/newtest_cases/test__prepare_pos_seq_lens_kernel.py +# Main kernel: _prepare_pos_seq_lens_kernel +# PT file: prepare_pos_seq_lens_test_data_v2.pt + +import triton +import triton.language as tl + + +# === _prepare_pos_seq_lens_kernel === +@triton.jit +def _prepare_pos_seq_lens_kernel( + pos_ptr, + seq_lens_ptr, + idx_mapping_ptr, + query_start_loc_ptr, + num_computed_tokens_ptr, + max_num_reqs, + BLOCK_SIZE: tl.constexpr, +): + req_id = tl.program_id(0) + num_reqs = tl.num_programs(0) - 1 + if req_id == num_reqs: + # Pad unused seq_lens as 0 for full CUDA graphs. + for i in tl.range(num_reqs, max_num_reqs, BLOCK_SIZE): + block = i + tl.arange(0, BLOCK_SIZE) + mask = block < max_num_reqs + tl.store(seq_lens_ptr + block, 0, mask=mask) + return + + req_state_idx = tl.load(idx_mapping_ptr + req_id) + num_computed_tokens = tl.load(num_computed_tokens_ptr + req_state_idx) + + start = tl.load(query_start_loc_ptr + req_id) + end = tl.load(query_start_loc_ptr + req_id + 1) + query_len = end - start + + seq_len = num_computed_tokens + query_len + tl.store(seq_lens_ptr + req_id, seq_len) + + for i in tl.range(0, query_len, BLOCK_SIZE): + block = i + tl.arange(0, BLOCK_SIZE) + mask = block < query_len + pos = num_computed_tokens + block + tl.store(pos_ptr + start + block, pos, mask=mask) + diff --git a/benchmarks/TritonNPUBench/vllm/_prepare_prefill_inputs_kernel/_prepare_prefill_inputs_kernel.py b/benchmarks/TritonNPUBench/vllm/_prepare_prefill_inputs_kernel/_prepare_prefill_inputs_kernel.py new file mode 100644 index 00000000..ce163893 --- /dev/null +++ b/benchmarks/TritonNPUBench/vllm/_prepare_prefill_inputs_kernel/_prepare_prefill_inputs_kernel.py @@ -0,0 +1,46 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/vllm_operator_cases/newtest_cases/test__prepare_prefill_inputs_kernel.py +# Main kernel: _prepare_prefill_inputs_kernel +# PT file: prepare_prefill_inputs_test_v2.pt + +import triton +import triton.language as tl + + +# === _prepare_prefill_inputs_kernel === +@triton.jit +def _prepare_prefill_inputs_kernel( + input_ids_ptr, + next_prefill_tokens_ptr, + idx_mapping_ptr, + query_start_loc_ptr, + prefill_token_ids_ptr, + prefill_token_ids_stride, + prefill_lens_ptr, + num_computed_tokens_ptr, + BLOCK_SIZE: tl.constexpr, +): + batch_idx = tl.program_id(0) + req_state_idx = tl.load(idx_mapping_ptr + batch_idx) + prefill_len = tl.load(prefill_lens_ptr + req_state_idx) + num_computed = tl.load(num_computed_tokens_ptr + req_state_idx) + if num_computed >= prefill_len: + # Not prefill. + return + + query_start = tl.load(query_start_loc_ptr + batch_idx) + query_end = tl.load(query_start_loc_ptr + batch_idx + 1) + query_len = query_end - query_start + + prefill_ptr = prefill_token_ids_ptr + req_state_idx * prefill_token_ids_stride + for i in range(0, query_len, BLOCK_SIZE): + block = i + tl.arange(0, BLOCK_SIZE) + mask = block < query_len + tokens = tl.load(prefill_ptr + num_computed + block, mask=mask) + tl.store(input_ids_ptr + query_start + block, tokens, mask=mask) + + next_pos = num_computed + query_len + if next_pos < prefill_len: + next_token = tl.load(prefill_ptr + next_pos) + tl.store(next_prefill_tokens_ptr + req_state_idx, next_token) + diff --git a/benchmarks/TritonNPUBench/vllm/_ranks_kernel/_ranks_kernel.py b/benchmarks/TritonNPUBench/vllm/_ranks_kernel/_ranks_kernel.py new file mode 100644 index 00000000..f01360a4 --- /dev/null +++ b/benchmarks/TritonNPUBench/vllm/_ranks_kernel/_ranks_kernel.py @@ -0,0 +1,32 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/vllm_operator_cases/newtest_cases/test__ranks_kernel.py +# Main kernel: _ranks_kernel +# PT file: ranks_kernel_test_v2.pt + +import triton +import triton.language as tl + + +# === _ranks_kernel === +@triton.jit +def _ranks_kernel( + output_ptr, + logits_ptr, + logits_stride, + token_ids_ptr, + vocab_size, + BLOCK_SIZE: tl.constexpr, +): + req_idx = tl.program_id(0) + row_ptr = logits_ptr + req_idx * logits_stride + + token_id = tl.load(token_ids_ptr + req_idx) + x = tl.load(row_ptr + token_id) + + n = 0 + for i in range(0, vocab_size, BLOCK_SIZE): + block = i + tl.arange(0, BLOCK_SIZE) + logits = tl.load(row_ptr + block, mask=block < vocab_size, other=float("-inf")) + n += tl.sum((logits > x).to(tl.int32)) + tl.store(output_ptr + req_idx, n) + diff --git a/benchmarks/TritonNPUBench/vllm/_rejection_sample_kernel/_rejection_sample_kernel.py b/benchmarks/TritonNPUBench/vllm/_rejection_sample_kernel/_rejection_sample_kernel.py new file mode 100644 index 00000000..90b0fc74 --- /dev/null +++ b/benchmarks/TritonNPUBench/vllm/_rejection_sample_kernel/_rejection_sample_kernel.py @@ -0,0 +1,42 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/vllm_operator_cases/newtest_cases/test__rejection_sample_kernel.py +# Main kernel: _rejection_sample_kernel +# PT file: rejection_sample_kernel_test_v2.pt + +import triton +import triton.language as tl + + +# === _rejection_sample_kernel === +@triton.jit +def _rejection_sample_kernel( + sampled_ptr, # [num_reqs, num_speculative_steps + 1] + sampled_stride, + num_sampled_ptr, # [num_reqs] + target_sampled_ptr, # [num_draft_tokens + num_reqs] + input_ids_ptr, # [num_draft_tokens + num_reqs] + cu_num_logits_ptr, # [num_reqs + 1] +): + req_idx = tl.program_id(0) + start_idx = tl.load(cu_num_logits_ptr + req_idx) + end_idx = tl.load(cu_num_logits_ptr + req_idx + 1) + num_tokens = end_idx - start_idx + + num_sampled = 0 + rejected = False + for i in range(num_tokens - 1): + if not rejected: + target_sampled = tl.load(target_sampled_ptr + start_idx + i) + draft_sampled = tl.load(input_ids_ptr + start_idx + i + 1) + tl.store(sampled_ptr + req_idx * sampled_stride + i, target_sampled) + num_sampled += 1 + if target_sampled != draft_sampled: + rejected = True + if not rejected: + target_sampled = tl.load(target_sampled_ptr + start_idx + num_tokens - 1) + tl.store( + sampled_ptr + req_idx * sampled_stride + num_tokens - 1, target_sampled + ) + num_sampled += 1 + tl.store(num_sampled_ptr + req_idx, num_sampled) + diff --git a/benchmarks/TritonNPUBench/vllm/_rms_norm_kernel/_rms_norm_kernel.py b/benchmarks/TritonNPUBench/vllm/_rms_norm_kernel/_rms_norm_kernel.py new file mode 100644 index 00000000..8899ea69 --- /dev/null +++ b/benchmarks/TritonNPUBench/vllm/_rms_norm_kernel/_rms_norm_kernel.py @@ -0,0 +1,60 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/vllm_operator_cases/newtest_cases/test__rms_norm_kernel.py +# Main kernel: _rms_norm_kernel +# PT file: _rms_norm_kernel_v2.pt + +import triton +import triton.language as tl + + +# === _rms_norm_kernel === +@triton.jit +def _rms_norm_kernel( + input_ptr, + weight_ptr, + output_ptr, + input_row_stride, + output_row_stride, + n_cols, + eps, + BLOCK_SIZE: tl.constexpr, +): + """ + Compute RMS normalization along the last dimension of a 2D tensor. + RMS Norm: y = x / sqrt(mean(x^2) + eps) * weight + Each block handles one row of the input tensor. + """ + row_idx = tl.program_id(0).to(tl.int64) + row_start_ptr = input_ptr + row_idx * input_row_stride + output_row_start_ptr = output_ptr + row_idx * output_row_stride + + # Step 1: Compute sum of squares in float32 to avoid overflow + sum_sq = tl.zeros([1], dtype=tl.float32) + for col_offset in range(0, n_cols, BLOCK_SIZE): + col_idx = col_offset + tl.arange(0, BLOCK_SIZE) + mask = col_idx < n_cols + + vals = tl.load(row_start_ptr + col_idx, mask=mask, other=0.0) + # Convert to float32 for accumulation to prevent overflow + vals_f32 = vals.to(tl.float32) + sq_vals = vals_f32 * vals_f32 + sum_sq += tl.sum(tl.where(mask, sq_vals, 0.0)) + + # Step 2: Compute RMS (root mean square) in float32 + mean_sq = sum_sq / n_cols + rms = tl.sqrt(mean_sq + eps) + inv_rms = 1.0 / rms + + # Step 3: Normalize and apply weight + for col_offset in range(0, n_cols, BLOCK_SIZE): + col_idx = col_offset + tl.arange(0, BLOCK_SIZE) + mask = col_idx < n_cols + vals = tl.load(row_start_ptr + col_idx, mask=mask, other=0.0) + weight = tl.load(weight_ptr + col_idx, mask=mask, other=1.0) + # Compute in float32 then convert back to input dtype + vals_f32 = vals.to(tl.float32) + weight_f32 = weight.to(tl.float32) + output_f32 = vals_f32 * inv_rms * weight_f32 + output = output_f32.to(vals.dtype) + tl.store(output_row_start_ptr + col_idx, output, mask=mask) + diff --git a/benchmarks/TritonNPUBench/vllm/_selective_scan_update_kernel/_selective_scan_update_kernel.py b/benchmarks/TritonNPUBench/vllm/_selective_scan_update_kernel/_selective_scan_update_kernel.py new file mode 100644 index 00000000..f9c88ced --- /dev/null +++ b/benchmarks/TritonNPUBench/vllm/_selective_scan_update_kernel/_selective_scan_update_kernel.py @@ -0,0 +1,243 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/vllm_operator_cases/newtest_cases/test__selective_scan_update_kernel.py +# Main kernel: _selective_scan_update_kernel +# PT file: _selective_scan_update_kernel_v2.pt + +import triton +import triton.language as tl + +@triton.jit +def softplus(dt): + dt = tl.where(dt <= 20.0, tl.math.log(tl.math.exp(dt) + 1), dt) + return dt + + +# === _selective_scan_update_kernel === +@triton.jit(do_not_specialize=["N"]) +def _selective_scan_update_kernel( + # Pointers to matrices + state_ptr, + x_ptr, + dt_ptr, + dt_bias_ptr, + A_ptr, + B_ptr, + C_ptr, + D_ptr, + z_ptr, + out_ptr, + state_batch_indices_ptr, + dst_state_batch_indices_ptr, + pad_slot_id, + num_accepted_tokens_ptr, + cu_seqlens_ptr, + # Matrix dimensions + N, + nheads, + dim, + dstate, + nheads_ngroups_ratio, + # Strides + stride_state_batch, + stride_state_head, + stride_state_dim, + stride_state_dstate, + stride_x_batch, + stride_x_head, + stride_x_dim, + stride_dt_batch, + stride_dt_head, + stride_dt_dim, + stride_dt_bias_head, + stride_dt_bias_dim, + stride_A_head, + stride_A_dim, + stride_A_dstate, + stride_B_batch, + stride_B_group, + stride_B_dstate, + stride_C_batch, + stride_C_group, + stride_C_dstate, + stride_D_head, + stride_D_dim, + stride_z_batch, + stride_z_head, + stride_z_dim, + stride_out_batch, + stride_out_head, + stride_out_dim, + stride_state_indices_batch, + stride_state_indices_T, + stride_dst_state_indices_batch, + stride_dst_state_indices_T, + # Meta-parameters + DT_SOFTPLUS: tl.constexpr, + TIE_HDIM: tl.constexpr, + BLOCK_SIZE_M: tl.constexpr, + HAS_DT_BIAS: tl.constexpr, + HAS_D: tl.constexpr, + HAS_Z: tl.constexpr, + HAS_STATE_BATCH_INDICES: tl.constexpr, + IS_SPEC_DECODING: tl.constexpr, + IS_VARLEN: tl.constexpr, + BLOCK_SIZE_DSTATE: tl.constexpr, +): + pid_m = tl.program_id(axis=0) + pid_b = tl.program_id(axis=1) + pid_h = tl.program_id(axis=2) + + if IS_VARLEN: + bos = tl.load(cu_seqlens_ptr + pid_b).to(tl.int64) + eos = tl.load(cu_seqlens_ptr + pid_b + 1).to(tl.int64) + seq_len = eos - bos + + if seq_len == 0: + return + else: + bos = pid_b + seq_len = 1 + + state_ptr_base = state_ptr + + # If HAS_STATE_BATCH_INDICES is true, then the ssm state's batch coordinate + # is taken from the state_batch_indices_ptr Otherwise, the state coordinate + # is the same as the batch id. + if HAS_STATE_BATCH_INDICES: + if IS_SPEC_DECODING: + num_accepted = tl.load(num_accepted_tokens_ptr + pid_b).to(tl.int64) + init_token_idx = tl.maximum(num_accepted - 1, 0) + else: + init_token_idx = 0 + + dst_state_batch_indices_ptr += pid_b * stride_dst_state_indices_batch + if not IS_SPEC_DECODING: + dst_state_batch_idx = tl.load( + dst_state_batch_indices_ptr + + init_token_idx * stride_dst_state_indices_T + ).to(tl.int64) + dst_state_ptr = state_ptr + ( + dst_state_batch_idx * stride_state_batch + pid_h * stride_state_head + ) + + state_batch_indices_ptr += ( + pid_b * stride_state_indices_batch + init_token_idx * stride_state_indices_T + ) + state_batch_idx = tl.load(state_batch_indices_ptr).to(tl.int64) + state_ptr += state_batch_idx * stride_state_batch + pid_h * stride_state_head + else: + dst_state_ptr = ( + state_ptr + pid_b * stride_state_batch + pid_h * stride_state_head + ) + state_ptr += pid_b * stride_state_batch + pid_h * stride_state_head + + x_ptr += bos * stride_x_batch + pid_h * stride_x_head + dt_ptr += bos * stride_dt_batch + pid_h * stride_dt_head + if HAS_DT_BIAS: + dt_bias_ptr += pid_h * stride_dt_bias_head + A_ptr += pid_h * stride_A_head + B_ptr += bos * stride_B_batch + (pid_h // nheads_ngroups_ratio) * stride_B_group + C_ptr += bos * stride_C_batch + (pid_h // nheads_ngroups_ratio) * stride_C_group + if HAS_Z: + z_ptr += bos * stride_z_batch + pid_h * stride_z_head + out_ptr += bos * stride_out_batch + pid_h * stride_out_head + + offs_m = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) + offs_n = tl.arange(0, BLOCK_SIZE_DSTATE) + state_ptrs = state_ptr + ( + offs_m[:, None] * stride_state_dim + offs_n[None, :] * stride_state_dstate + ) + if not IS_SPEC_DECODING: + dst_state_ptrs = dst_state_ptr + ( + offs_m[:, None] * stride_state_dim + offs_n[None, :] * stride_state_dstate + ) + + mask = (offs_m[:, None] < dim) & (offs_n[None, :] < dstate) + if HAS_STATE_BATCH_INDICES: + mask &= state_batch_idx != pad_slot_id + state = tl.load(state_ptrs, mask=mask, other=0.0).to(tl.float32) + + if HAS_DT_BIAS: + dt_bias_ptrs = dt_bias_ptr + offs_m * stride_dt_bias_dim + if HAS_D: + D_ptr += pid_h * stride_D_head + D_ptrs = D_ptr + offs_m * stride_D_dim + A_ptrs = A_ptr + offs_m[:, None] * stride_A_dim + offs_n[None, :] * stride_A_dstate + + for i_t in range(seq_len): + x_ptrs = x_ptr + offs_m * stride_x_dim + dt_ptrs = dt_ptr + offs_m * stride_dt_dim + B_ptrs = B_ptr + offs_n * stride_B_dstate + C_ptrs = C_ptr + offs_n * stride_C_dstate + if HAS_Z: + z_ptrs = z_ptr + offs_m * stride_z_dim + out_ptrs = out_ptr + offs_m * stride_out_dim + + x = tl.load(x_ptrs, mask=offs_m < dim, other=0.0).to(tl.float32) + if not TIE_HDIM: + dt = tl.load(dt_ptrs, mask=offs_m < dim, other=0.0).to(tl.float32) + if HAS_DT_BIAS: + dt += tl.load(dt_bias_ptrs, mask=offs_m < dim, other=0.0).to(tl.float32) + if DT_SOFTPLUS: + dt = softplus(dt) + A = tl.load( + A_ptrs, + mask=(offs_m[:, None] < dim) & (offs_n[None, :] < dstate), + other=0.0, + ).to(tl.float32) + dA = tl.exp(A * dt[:, None]) + else: + dt = tl.load(dt_ptr).to(tl.float32) + if HAS_DT_BIAS: + dt += tl.load(dt_bias_ptr).to(tl.float32) + if DT_SOFTPLUS: + dt = softplus(dt) + A = tl.load(A_ptr).to(tl.float32) + dA = tl.exp(A * dt) # scalar, not a matrix + + B = tl.load(B_ptrs, mask=offs_n < dstate, other=0.0).to(tl.float32) + C = tl.load(C_ptrs, mask=offs_n < dstate, other=0.0).to(tl.float32) + if HAS_D: + D = tl.load(D_ptrs, mask=offs_m < dim, other=0.0).to(tl.float32) + if HAS_Z: + z = tl.load(z_ptrs, mask=offs_m < dim, other=0.0).to(tl.float32) + + dB = B[None, :] * dt[:, None] if not TIE_HDIM else B * dt + state = state * dA + dB * x[:, None] + + if IS_SPEC_DECODING: + dst_idx_ptr = dst_state_batch_indices_ptr + i_t * stride_dst_state_indices_T + token_dst_idx = tl.load(dst_idx_ptr).to(tl.int64) + if token_dst_idx != pad_slot_id: + token_dst_ptrs = ( + state_ptr_base + + token_dst_idx * stride_state_batch + + pid_h * stride_state_head + + offs_m[:, None] * stride_state_dim + + offs_n[None, :] * stride_state_dstate + ) + tl.store( + token_dst_ptrs, state.to(token_dst_ptrs.dtype.element_ty), mask=mask + ) + + out = tl.sum(state * C[None, :], axis=1) + if HAS_D: + out += x * D + if HAS_Z: + out *= z * tl.sigmoid(z) + tl.store(out_ptrs, out, mask=offs_m < dim) + + x_ptr += stride_x_batch + dt_ptr += stride_dt_batch + B_ptr += stride_B_batch + C_ptr += stride_C_batch + out_ptr += stride_out_batch + if HAS_Z: + z_ptr += stride_z_batch + + if not IS_SPEC_DECODING: + tl.store(dst_state_ptrs, state.to(dst_state_ptrs.dtype.element_ty), mask=mask) + + + + diff --git a/benchmarks/TritonNPUBench/vllm/_silu_mul_fp8_quant_deep_gemm/_silu_mul_fp8_quant_deep_gemm.py b/benchmarks/TritonNPUBench/vllm/_silu_mul_fp8_quant_deep_gemm/_silu_mul_fp8_quant_deep_gemm.py new file mode 100644 index 00000000..4bb1d2b3 --- /dev/null +++ b/benchmarks/TritonNPUBench/vllm/_silu_mul_fp8_quant_deep_gemm/_silu_mul_fp8_quant_deep_gemm.py @@ -0,0 +1,83 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/vllm_operator_cases/newtest_cases/test__silu_mul_fp8_quant_deep_gemm.py +# Main kernel: _silu_mul_fp8_quant_deep_gemm +# PT file: _silu_mul_fp8_quant_deep_gemm_v2.pt + +import triton +import triton.language as tl + + +# === _silu_mul_fp8_quant_deep_gemm === +@triton.jit +def _silu_mul_fp8_quant_deep_gemm( + # Pointers ------------------------------------------------------------ + input_ptr, # 16-bit activations (E, T, 2*H) + y_q_ptr, # fp8 quantized activations (E, T, H) + y_s_ptr, # 16-bit scales (E, T, G) + counts_ptr, # int32 num tokens per expert (E) + # Sizes --------------------------------------------------------------- + H: tl.constexpr, # hidden dimension (per output) + GROUP_SIZE: tl.constexpr, # elements per group (usually 128) + # Strides for input (elements) --------------------------------------- + stride_i_e, + stride_i_t, + stride_i_h, + # Strides for y_q (elements) ----------------------------------------- + stride_yq_e, + stride_yq_t, + stride_yq_h, + # Strides for y_s (elements) ----------------------------------------- + stride_ys_e, + stride_ys_t, + stride_ys_g, + # Stride for counts (elements) + stride_counts_e, + # Numeric params ------------------------------------------------------ + eps: tl.constexpr, + fp8_min: tl.constexpr, + fp8_max: tl.constexpr, + use_ue8m0: tl.constexpr, + # Meta --------------------------------------------------------------- + BLOCK: tl.constexpr, + NUM_STAGES: tl.constexpr, +): + G = H // GROUP_SIZE + + # map program id -> (e, g) + pid = tl.program_id(0) + e = pid // G + g = pid % G + + e = e.to(tl.int64) + g = g.to(tl.int64) + + # number of valid tokens for this expert + n_tokens = tl.load(counts_ptr + e * stride_counts_e).to(tl.int64) + + cols = tl.arange(0, BLOCK).to(tl.int64) + mask = cols < BLOCK + + base_input_offset = e * stride_i_e + g * GROUP_SIZE * stride_i_h + base_gate_offset = base_input_offset + cols * stride_i_h + base_up_offset = base_input_offset + H * stride_i_h + cols * stride_i_h + base_yq_offset = e * stride_yq_e + g * GROUP_SIZE * stride_yq_h + cols * stride_yq_h + base_ys_offset = e * stride_ys_e + g * stride_ys_g + + for t in tl.range(0, n_tokens, num_stages=NUM_STAGES): + gate = tl.load( + input_ptr + base_gate_offset + t * stride_i_t, mask=mask, other=0.0 + ).to(tl.float32) + up = tl.load(input_ptr + base_up_offset + t * stride_i_t, mask=mask, other=0.0) + + gate = gate * (1.0 / (1.0 + tl.exp(-gate))) + y = gate * up + + y_s = tl.maximum(tl.max(tl.abs(y)), eps) / fp8_max + if use_ue8m0: + y_s = tl.exp2(tl.ceil(tl.log2(y_s))) + + y_q = tl.clamp(y / y_s, fp8_min, fp8_max).to(y_q_ptr.dtype.element_ty) + + tl.store(y_q_ptr + base_yq_offset + t * stride_yq_t, y_q, mask=mask) + tl.store(y_s_ptr + base_ys_offset + t * stride_ys_t, y_s) + diff --git a/benchmarks/TritonNPUBench/vllm/_silu_mul_per_token_group_quant_fp8_colmajor/_silu_mul_per_token_group_quant_fp8_colmajor.py b/benchmarks/TritonNPUBench/vllm/_silu_mul_per_token_group_quant_fp8_colmajor/_silu_mul_per_token_group_quant_fp8_colmajor.py new file mode 100644 index 00000000..72cc2321 --- /dev/null +++ b/benchmarks/TritonNPUBench/vllm/_silu_mul_per_token_group_quant_fp8_colmajor/_silu_mul_per_token_group_quant_fp8_colmajor.py @@ -0,0 +1,80 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/vllm_operator_cases/newtest_cases/test__silu_mul_per_token_group_quant_fp8_colmajor.py +# Main kernel: _silu_mul_per_token_group_quant_fp8_colmajor +# PT file: _silu_mul_per_token_group_quant_fp8_colmajor_v2.pt + +import triton +import triton.language as tl + + +# === _silu_mul_per_token_group_quant_fp8_colmajor === +@triton.jit +def _silu_mul_per_token_group_quant_fp8_colmajor( + y_ptr, # [M, N] + y_q_ptr, # [M, N // 2] + y_s_ptr, # [M, (N // 2) // GROUP_SIZE] + M, # num tokens + N, # intermediate size + # Stride + y_s_col_stride: tl.int64, + # Information for float8 + eps, + fp8_min, + fp8_max, + use_ue8m0: tl.constexpr, + # Meta-parameters + GROUP_SIZE: tl.constexpr, + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, +): + # TODO(varun) : Add expert_ids so we may early-exit no-op thread blocks. + """ + Each thread block (BLOCK_N) computes [BLOCK_M, GROUP_SIZE] act-mul outputs. Then + the thread block quantizes the [BLOCK_M, GROUP_SIZE] block of values and fills + the outputs tensors at the right positions. + """ + + pid_m = tl.program_id(0) + pid_n = tl.program_id(1) + N_2 = N // 2 + + m_offset = pid_m * BLOCK_M + n_offset = pid_n * BLOCK_N + if m_offset >= M: + return + + offs_n = tl.arange(0, BLOCK_N).to(tl.int64) + offs_m = tl.arange(0, BLOCK_M).to(tl.int64) + + base_y_ptr = y_ptr + m_offset * N + n_offset + + act_in_ptrs = base_y_ptr + offs_m[:, None] * N + offs_n[None, :] + + act_in = tl.load(act_in_ptrs) + mul_in = tl.load(act_in_ptrs + N_2) + + # silu & mul + act_in = act_in.to(tl.float32) + one_f32 = tl.cast(1, tl.float32) + silu_out = (act_in / (one_f32 + tl.exp(-act_in))).to(y_ptr.dtype.element_ty) + y = (silu_out * mul_in).to(tl.float32) + + # quant + _absmax = tl.maximum(tl.max(tl.abs(y), axis=1), eps) + scale_raw = _absmax / fp8_max + y_s = tl.math.exp2(tl.ceil(tl.log2(scale_raw))) if use_ue8m0 else scale_raw + y_s = tl.reshape(y_s, (BLOCK_M, 1)) + y_q = tl.clamp(y / y_s, fp8_min, fp8_max).to(y_q_ptr.dtype.element_ty) + + # store y_q + base_y_q_ptr = y_q_ptr + m_offset * N_2 + n_offset + y_q_ptrs = base_y_q_ptr + offs_m[:, None] * N_2 + offs_n[None, :] + tl.store(y_q_ptrs, y_q) + + # store y_s + group_id = n_offset // GROUP_SIZE + base_y_s_ptr = y_s_ptr + group_id * y_s_col_stride + m_offset + y_s_ptrs = base_y_s_ptr + offs_m + y_s = tl.reshape(y_s, (BLOCK_M,)) + tl.store(y_s_ptrs, y_s) + diff --git a/benchmarks/TritonNPUBench/vllm/_state_passing_fwd_kernel/_state_passing_fwd_kernel.py b/benchmarks/TritonNPUBench/vllm/_state_passing_fwd_kernel/_state_passing_fwd_kernel.py new file mode 100644 index 00000000..8adcd6e9 --- /dev/null +++ b/benchmarks/TritonNPUBench/vllm/_state_passing_fwd_kernel/_state_passing_fwd_kernel.py @@ -0,0 +1,92 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/vllm_operator_cases/newtest_cases/test__state_passing_fwd_kernel.py +# Main kernel: _state_passing_fwd_kernel +# PT file: _state_passing_fwd_kernel_v2.pt + +import triton +import triton.language as tl + + +# === _state_passing_fwd_kernel === +@triton.jit +def _state_passing_fwd_kernel( + # Pointers to matrices + states_ptr, + out_ptr, + dA_cs_ptr, + initstates_ptr, + seq_idx_ptr, + cu_chunk_seqlens_ptr, + # Matrix dimensions + dim: tl.constexpr, + nchunks, + seqlen, + chunk_size: tl.constexpr, + # Strides + stride_states_chunk: tl.int64, + stride_states_head: tl.int64, + stride_states_dim: tl.constexpr, + stride_out_chunk: tl.int64, + stride_out_head: tl.int64, + stride_out_dim: tl.constexpr, + stride_dA_cs_head: tl.int64, + stride_dA_cs_chunk: tl.int64, + stride_dA_cs_csize: tl.constexpr, + stride_initstates_batch: tl.int64, + stride_initstates_head: tl.int64, + stride_initstates_dim: tl.constexpr, + stride_seq_idx_chunk: tl.constexpr, + # Meta-parameters + HAS_INITSTATES: tl.constexpr, + BLOCK_SIZE: tl.constexpr, +): + pid_h = tl.program_id(axis=1) + pid_m = tl.program_id(axis=0) + + states_ptr += pid_h * stride_states_head + dA_cs_ptr += pid_h * stride_dA_cs_head + (chunk_size - 1) * stride_dA_cs_csize + out_ptr += pid_h * stride_out_head + + offs_m = pid_m * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + states_ptrs = states_ptr + offs_m * stride_states_dim + out_ptrs = out_ptr + offs_m * stride_out_dim + + if HAS_INITSTATES: + initstates_ptrs = ( + initstates_ptr + + pid_h * stride_initstates_head + + offs_m * stride_initstates_dim + ) + + states = tl.load(initstates_ptrs, mask=offs_m < dim, other=0.0).to(tl.float32) + else: + states = tl.zeros((BLOCK_SIZE,), dtype=tl.float32) + + prev_seq_idx = 0 + for c in range(nchunks): + new_states = tl.load(states_ptrs, mask=offs_m < dim, other=0.0).to(tl.float32) + dA_cs = tl.load(dA_cs_ptr).to(tl.float32) + seq_idx = tl.load(seq_idx_ptr + c * stride_seq_idx_chunk) + # we have started a new sequence + if prev_seq_idx != seq_idx: + if HAS_INITSTATES: + initstates_ptrs = ( + initstates_ptr + + seq_idx * stride_initstates_batch + + pid_h * stride_initstates_head + + offs_m * stride_initstates_dim + ) + states = tl.load(initstates_ptrs, mask=offs_m < dim, other=0.0).to( + tl.float32 + ) + else: + states = tl.zeros((BLOCK_SIZE,), dtype=tl.float32) + + prev_seq_idx = seq_idx + states = tl.exp(dA_cs) * states + new_states + tl.store(out_ptrs, states, mask=offs_m < dim) + + states_ptrs += stride_states_chunk + dA_cs_ptr += stride_dA_cs_chunk + out_ptrs += stride_out_chunk + diff --git a/benchmarks/TritonNPUBench/vllm/_topk_log_softmax_kernel/_topk_log_softmax_kernel.py b/benchmarks/TritonNPUBench/vllm/_topk_log_softmax_kernel/_topk_log_softmax_kernel.py new file mode 100644 index 00000000..fc2c0e0c --- /dev/null +++ b/benchmarks/TritonNPUBench/vllm/_topk_log_softmax_kernel/_topk_log_softmax_kernel.py @@ -0,0 +1,51 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/vllm_operator_cases/newtest_cases/test__topk_log_softmax_kernel.py +# Main kernel: _topk_log_softmax_kernel +# PT file: topk_log_softmax_kernel_test_data_v2.pt + +import triton +import triton.language as tl + + +# === _topk_log_softmax_kernel === +@triton.jit +def _topk_log_softmax_kernel( + output_ptr, + logits_ptr, + logits_stride, + topk_ids_ptr, + topk, + vocab_size, + BLOCK_SIZE: tl.constexpr, + PADDED_TOPK: tl.constexpr, +): + req_idx = tl.program_id(0) + row_ptr = logits_ptr + req_idx * logits_stride + + max_val = float("-inf") + for i in range(0, vocab_size, BLOCK_SIZE): + block = i + tl.arange(0, BLOCK_SIZE) + logits = tl.load(row_ptr + block, mask=block < vocab_size, other=float("-inf")) + max_val = tl.max(tl.maximum(logits, max_val)) + max_val = max_val.to(tl.float32) # type: ignore + + se = 0.0 + for i in range(0, vocab_size, BLOCK_SIZE): + block = i + tl.arange(0, BLOCK_SIZE) + logits = tl.load(row_ptr + block, mask=block < vocab_size, other=0.0) + # NOTE(woosuk): Make sure that logits and all following operations use FP32. + logits = logits.to(tl.float32) + e = tl.exp(logits - max_val) + e = tl.where(block < vocab_size, e, 0.0) + se += tl.sum(e) + lse = tl.log(se) + + k_offset = tl.arange(0, PADDED_TOPK) + k_mask = k_offset < topk + topk_ids = tl.load(topk_ids_ptr + req_idx * topk + k_offset, mask=k_mask, other=0) + + logits = tl.load(row_ptr + topk_ids, mask=k_mask) + logits = logits.to(tl.float32) + o = logits - max_val - lse + tl.store(output_ptr + req_idx * topk + k_offset, o, mask=k_mask) + diff --git a/benchmarks/TritonNPUBench/vllm/_triton_mrope_forward/_triton_mrope_forward.py b/benchmarks/TritonNPUBench/vllm/_triton_mrope_forward/_triton_mrope_forward.py new file mode 100644 index 00000000..a842141e --- /dev/null +++ b/benchmarks/TritonNPUBench/vllm/_triton_mrope_forward/_triton_mrope_forward.py @@ -0,0 +1,128 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/vllm_operator_cases/newtest_cases/test__triton_mrope_forward.py +# Main kernel: _triton_mrope_forward +# PT file: _triton_mrope_forward_v2.pt + +import triton +import triton.language as tl + + +# === _triton_mrope_forward === +@triton.jit +def _triton_mrope_forward( + q_ptr, + k_ptr, + cos, + sin, + num_tokens, + n_qh: tl.constexpr, + n_kh: tl.constexpr, + hd: tl.constexpr, + rd: tl.constexpr, + pad_n_qh: tl.constexpr, + pad_n_kh: tl.constexpr, + pad_hd: tl.constexpr, + mrope_section_t: tl.constexpr, + mrope_section_h: tl.constexpr, + mrope_section_w: tl.constexpr, + is_interleaved: tl.constexpr, +): + # Adapted from + # https://github.com/linkedin/Liger-Kernel/blob/main/src/liger_kernel/ops/qwen2vl_mrope.py + # This version supports flatten input tensors from vllm + # and supports cos and sin cache with shape (3, num_tokens, head_dim // 2) + # instead of (3, bsz, seq_len, head_dim), also supports interleaved rotary + pid = tl.program_id(0) + # locate start address + q_ptr = q_ptr + pid * (n_qh * hd) + k_ptr = k_ptr + pid * (n_kh * hd) + + # #################################################################### + # get the cos(mθ_{i...d/2}) and sin(mθ_{i...d/2}) for token position + # m of this program instance + # #################################################################### + # Note: cos and sin now have shape (3, num_tokens, head_dim // 2) + + # Updated stride calculation for half head_dim + half_rd = rd // 2 + t_cos = cos + pid * half_rd + h_cos = t_cos + num_tokens * half_rd + w_cos = h_cos + num_tokens * half_rd + t_sin = sin + pid * half_rd + h_sin = t_sin + num_tokens * half_rd + w_sin = h_sin + num_tokens * half_rd + + # Updated offsets for half head_dim + cos_offsets = tl.arange(0, pad_hd // 2) + if is_interleaved: + h_mask = ((cos_offsets % 3) == 1) & (cos_offsets <= 3 * mrope_section_h) + w_mask = ((cos_offsets % 3) == 2) & (cos_offsets <= 3 * mrope_section_w) + t_mask = ~(h_mask | w_mask) + else: + t_end = mrope_section_t + h_end = t_end + mrope_section_h + t_mask = cos_offsets < mrope_section_t + h_mask = (t_end <= cos_offsets) & (cos_offsets < h_end) + w_mask = (h_end <= cos_offsets) & (cos_offsets < half_rd) + + t_cos_row = tl.load(t_cos + cos_offsets, mask=t_mask, other=0) + h_cos_row = tl.load(h_cos + cos_offsets, mask=h_mask, other=0) + w_cos_row = tl.load(w_cos + cos_offsets, mask=w_mask, other=0) + t_sin_row = tl.load(t_sin + cos_offsets, mask=t_mask, other=0) + h_sin_row = tl.load(h_sin + cos_offsets, mask=h_mask, other=0) + w_sin_row = tl.load(w_sin + cos_offsets, mask=w_mask, other=0) + + cos_row = t_cos_row + h_cos_row + w_cos_row + sin_row = t_sin_row + h_sin_row + w_sin_row + + # #################################################################### + # Load the left and right half of q and k for the current + # program instance (i.e. for the current token) separately + # #################################################################### + # left half of the head + first_half_q_offsets = ( + tl.arange(0, pad_n_qh)[:, None] * hd + tl.arange(0, pad_hd // 2)[None, :] + ) + first_half_k_offsets = ( + tl.arange(0, pad_n_kh)[:, None] * hd + tl.arange(0, pad_hd // 2)[None, :] + ) + first_q_mask = (tl.arange(0, pad_n_qh)[:, None] < n_qh) & ( + tl.arange(0, pad_hd // 2)[None, :] < rd // 2 + ) + first_k_mask = (tl.arange(0, pad_n_kh)[:, None] < n_kh) & ( + tl.arange(0, pad_hd // 2)[None, :] < rd // 2 + ) + + q_tile_1 = tl.load(q_ptr + first_half_q_offsets, mask=first_q_mask, other=0).to( + sin_row.dtype + ) + k_tile_1 = tl.load(k_ptr + first_half_k_offsets, mask=first_k_mask, other=0).to( + sin_row.dtype + ) + + # right half of the head + second_half_q_offsets = first_half_q_offsets + (rd // 2) + second_half_k_offsets = first_half_k_offsets + (rd // 2) + second_q_mask = first_q_mask + second_k_mask = first_k_mask + + q_tile_2 = tl.load(q_ptr + second_half_q_offsets, mask=second_q_mask, other=0).to( + sin_row.dtype + ) + k_tile_2 = tl.load(k_ptr + second_half_k_offsets, mask=second_k_mask, other=0).to( + sin_row.dtype + ) + + # y = [x1, x2] * [cos, cos] + [-x2, x1] * [sin, sin] + # Since cos and sin are now half-size, + # we use the same cos_row and sin_row for both halves + new_q_tile_1 = q_tile_1 * cos_row - q_tile_2 * sin_row + tl.store(q_ptr + first_half_q_offsets, new_q_tile_1, mask=first_q_mask) + new_q_tile_2 = q_tile_2 * cos_row + q_tile_1 * sin_row + tl.store(q_ptr + second_half_q_offsets, new_q_tile_2, mask=second_q_mask) + + new_k_tile_1 = k_tile_1 * cos_row - k_tile_2 * sin_row + tl.store(k_ptr + first_half_k_offsets, new_k_tile_1, mask=first_k_mask) + new_k_tile_2 = k_tile_2 * cos_row + k_tile_1 * sin_row + tl.store(k_ptr + second_half_k_offsets, new_k_tile_2, mask=second_k_mask) + diff --git a/benchmarks/TritonNPUBench/vllm/_trtllm_prefill_attn_kvfp8_dequant/_trtllm_prefill_attn_kvfp8_dequant.py b/benchmarks/TritonNPUBench/vllm/_trtllm_prefill_attn_kvfp8_dequant/_trtllm_prefill_attn_kvfp8_dequant.py new file mode 100644 index 00000000..d2398f75 --- /dev/null +++ b/benchmarks/TritonNPUBench/vllm/_trtllm_prefill_attn_kvfp8_dequant/_trtllm_prefill_attn_kvfp8_dequant.py @@ -0,0 +1,56 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/vllm_operator_cases/newtest_cases/test__trtllm_prefill_attn_kvfp8_dequant.py +# Main kernel: _trtllm_prefill_attn_kvfp8_dequant +# PT file: _trtllm_prefill_attn_kvfp8_dequant_v2.pt + +import triton +import triton.language as tl + + +# === _trtllm_prefill_attn_kvfp8_dequant === +@triton.jit +def _trtllm_prefill_attn_kvfp8_dequant( + kv_cache_ptr, + block_tables_prefill_ptr, + block_table_stride, + mock_kv_cache_ptr, + k_scale_ptr, + v_scale_ptr, + K_CACHE_STRIDE: tl.constexpr, + KV_CACHE_STRIDE: tl.constexpr, +): + batch_idx = tl.program_id(0).to(tl.int64) + mock_block_table_idx = tl.program_id(1).to(tl.int64) + orig_page_num = tl.load( + block_tables_prefill_ptr + batch_idx * block_table_stride + mock_block_table_idx + ).to(tl.int64) + if orig_page_num <= 0: + return + dequant_dtype = mock_kv_cache_ptr.dtype.element_ty + + # Dequantize K + k_scale_val = tl.load(k_scale_ptr) + offset = orig_page_num * KV_CACHE_STRIDE + tl.arange(0, K_CACHE_STRIDE) + fp8_vals = tl.load(kv_cache_ptr + offset) + dequantized_vals = fp8_vals.to(tl.float32) * k_scale_val + mock_cache_offset = ( + batch_idx * block_table_stride + mock_block_table_idx + 1 + ) * KV_CACHE_STRIDE + tl.arange(0, K_CACHE_STRIDE) + dequantized_vals = dequantized_vals.to(dequant_dtype) + tl.store(mock_kv_cache_ptr + mock_cache_offset, dequantized_vals) + + # Dequantize V + v_scale_val = tl.load(v_scale_ptr) + offset = ( + orig_page_num * KV_CACHE_STRIDE + K_CACHE_STRIDE + tl.arange(0, K_CACHE_STRIDE) + ) + fp8_vals = tl.load(kv_cache_ptr + offset) + dequantized_vals = fp8_vals.to(tl.float32) * v_scale_val + mock_cache_offset = ( + (batch_idx * block_table_stride + mock_block_table_idx + 1) * KV_CACHE_STRIDE + + K_CACHE_STRIDE + + tl.arange(0, K_CACHE_STRIDE) + ) + dequantized_vals = dequantized_vals.to(dequant_dtype) + tl.store(mock_kv_cache_ptr + mock_cache_offset, dequantized_vals) + diff --git a/benchmarks/TritonNPUBench/vllm/_unpack_seq_triton_kernel/_unpack_seq_triton_kernel.py b/benchmarks/TritonNPUBench/vllm/_unpack_seq_triton_kernel/_unpack_seq_triton_kernel.py new file mode 100644 index 00000000..6b33a7fd --- /dev/null +++ b/benchmarks/TritonNPUBench/vllm/_unpack_seq_triton_kernel/_unpack_seq_triton_kernel.py @@ -0,0 +1,52 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/vllm_operator_cases/newtest_cases/test__unpack_seq_triton_kernel.py +# Main kernel: _unpack_seq_triton_kernel +# PT file: unpack_seq_test_v2.pt + +import triton +import triton.language as tl + + +# === _unpack_seq_triton_kernel === +@triton.jit +def _unpack_seq_triton_kernel( + packed_ptr, # [B, Lmax, D] + out_ptr, # [N, D] + lengths_ptr, # *i32, [B] + B: tl.constexpr, + Lmax: tl.constexpr, + D: tl.constexpr, + BLOCK_T: tl.constexpr, # timesteps per program + BLOCK_D: tl.constexpr, # features per program +): + pid_b = tl.program_id(0) # batch id + pid_t = tl.program_id(1) # block over time dimension + pid_d = tl.program_id(2) # block over feature dimension + off_t = pid_t * BLOCK_T + tl.arange(0, BLOCK_T) # [BLOCK_T] + off_d = pid_d * BLOCK_D + tl.arange(0, BLOCK_D) # [BLOCK_D] + + # bounds: compute start from cumulative lengths + in_start = 0 + for i in range(pid_b): + in_start += tl.load(lengths_ptr + i) + seq_len = tl.load(lengths_ptr + pid_b) + + # valid time positions for this block + t_mask = off_t < Lmax + valid_row = (off_t < seq_len) & t_mask + + # compute output row indices for valid (b, t) + out_row = in_start + off_t + + # Pointers + # packed_ptr: row-major [B, Lmax, D] + packed_row_ptr = packed_ptr + (pid_b * Lmax + off_t)[:, None] * D + off_d[None, :] + + # out_ptr: row-major [N, D] + out_row_ptr = out_ptr + out_row[:, None] * D + off_d[None, :] + + # Load from packed tensor and store to output + d_mask = off_d[None, :] < D + packed_vals = tl.load(packed_row_ptr, mask=valid_row[:, None] & d_mask) + tl.store(out_row_ptr, packed_vals, mask=valid_row[:, None] & d_mask) + diff --git a/benchmarks/TritonNPUBench/vllm/_update_eagle_inputs_kernel/_update_eagle_inputs_kernel.py b/benchmarks/TritonNPUBench/vllm/_update_eagle_inputs_kernel/_update_eagle_inputs_kernel.py new file mode 100644 index 00000000..91a4f017 --- /dev/null +++ b/benchmarks/TritonNPUBench/vllm/_update_eagle_inputs_kernel/_update_eagle_inputs_kernel.py @@ -0,0 +1,55 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/vllm_operator_cases/newtest_cases/test__update_eagle_inputs_kernel.py +# Main kernel: _update_eagle_inputs_kernel +# PT file: _update_eagle_inputs_kernel_v2.pt + +import triton +import triton.language as tl + + +# === _update_eagle_inputs_kernel === +@triton.jit +def _update_eagle_inputs_kernel( + input_ids_ptr, + positions_ptr, + input_hidden_states_ptr, + input_hidden_states_stride, + seq_lens_ptr, + max_model_len, + draft_tokens_ptr, + output_hidden_states_ptr, + output_hidden_states_stride, + hidden_size, + BLOCK_SIZE: tl.constexpr, +): + req_idx = tl.program_id(0) + + # Draft token -> Input ID. + draft_token = tl.load(draft_tokens_ptr + req_idx) + tl.store(input_ids_ptr + req_idx, draft_token) + + # Output hidden states -> Input hidden states. + for i in range(0, hidden_size, BLOCK_SIZE): + block = i + tl.arange(0, BLOCK_SIZE) + mask = block < hidden_size + output_hidden_states = tl.load( + output_hidden_states_ptr + req_idx * output_hidden_states_stride + block, + mask=mask, + ) + tl.store( + input_hidden_states_ptr + req_idx * input_hidden_states_stride + block, + output_hidden_states, + mask=mask, + ) + + # Increment position and seq_lens. + # NOTE(woosuk): To prevent out-of-range access, we clamp these values + # if they reach the max model length. + position = tl.load(positions_ptr + req_idx) + position = tl.minimum(position + 1, max_model_len - 1) + tl.store(positions_ptr + req_idx, position) + + seq_len = tl.load(seq_lens_ptr + req_idx) + seq_len = tl.minimum(seq_len + 1, max_model_len) + tl.store(seq_lens_ptr + req_idx, seq_len) + diff --git a/benchmarks/TritonNPUBench/vllm/_w8a8_block_int8_matmul/_w8a8_block_int8_matmul.py b/benchmarks/TritonNPUBench/vllm/_w8a8_block_int8_matmul/_w8a8_block_int8_matmul.py new file mode 100644 index 00000000..b16dd486 --- /dev/null +++ b/benchmarks/TritonNPUBench/vllm/_w8a8_block_int8_matmul/_w8a8_block_int8_matmul.py @@ -0,0 +1,94 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/vllm_operator_cases/newtest_cases/test__w8a8_block_int8_matmul.py +# Main kernel: _w8a8_block_int8_matmul +# PT file: _w8a8_block_int8_matmul_v2.pt + +import triton +import triton.language as tl + + +# === _w8a8_block_int8_matmul === +@triton.jit +def _w8a8_block_int8_matmul( + # Pointers to inputs and output + A, + B, + C, + As, + Bs, + # Shape for matmul + M, + N, + K, + # Block size for block-wise quantization + group_n, + group_k, + # Stride for inputs and output + stride_am, + stride_ak, + stride_bk, + stride_bn, + stride_cm, + stride_cn, + stride_As_m, + stride_As_k, + stride_Bs_k, + stride_Bs_n, + # Meta-parameters + BLOCK_SIZE_M: tl.constexpr, + BLOCK_SIZE_N: tl.constexpr, + BLOCK_SIZE_K: tl.constexpr, + GROUP_SIZE_M: tl.constexpr, +): + """Triton-accelerated function used to perform linear operations (dot + product) on input tensors `A` and `B` with block-wise quantization, and + store the result in output tensor `C`. + """ + + pid = tl.program_id(axis=0) + num_pid_m = tl.cdiv(M, BLOCK_SIZE_M) + num_pid_n = tl.cdiv(N, BLOCK_SIZE_N) + num_pid_in_group = GROUP_SIZE_M * num_pid_n + group_id = pid // num_pid_in_group + first_pid_m = group_id * GROUP_SIZE_M + group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M) + pid_m = first_pid_m + (pid % group_size_m) + pid_n = (pid % num_pid_in_group) // group_size_m + + offs_am = (pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)) % M + offs_bn = (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)) % N + offs_k = tl.arange(0, BLOCK_SIZE_K) + a_ptrs = A + (offs_am[:, None] * stride_am + offs_k[None, :] * stride_ak) + b_ptrs = B + (offs_k[:, None] * stride_bk + offs_bn[None, :] * stride_bn) + + As_ptrs = As + offs_am * stride_As_m + offs_bsn = offs_bn // group_n + Bs_ptrs = Bs + offs_bsn * stride_Bs_n + + accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32) + for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)): + a = tl.load(a_ptrs, mask=offs_k[None, :] < K - k * BLOCK_SIZE_K, other=0.0) + b = tl.load(b_ptrs, mask=offs_k[:, None] < K - k * BLOCK_SIZE_K, other=0.0) + + k_start = k * BLOCK_SIZE_K + offs_ks = k_start // group_k + a_s = tl.load(As_ptrs + offs_ks * stride_As_k) + b_s = tl.load(Bs_ptrs + offs_ks * stride_Bs_k) + + accumulator += tl.dot(a, b).to(tl.float32) * a_s[:, None] * b_s[None, :] + a_ptrs += BLOCK_SIZE_K * stride_ak + b_ptrs += BLOCK_SIZE_K * stride_bk + + if C.dtype.element_ty == tl.bfloat16: + c = accumulator.to(tl.bfloat16) + elif C.dtype.element_ty == tl.float16: + c = accumulator.to(tl.float16) + else: + c = accumulator.to(tl.float32) + + offs_cm = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) + offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) + c_ptrs = C + stride_cm * offs_cm[:, None] + stride_cn * offs_cn[None, :] + c_mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < N) + tl.store(c_ptrs, c, mask=c_mask) + diff --git a/benchmarks/TritonNPUBench/vllm/_w8a8_triton_block_scaled_mm/_w8a8_triton_block_scaled_mm.py b/benchmarks/TritonNPUBench/vllm/_w8a8_triton_block_scaled_mm/_w8a8_triton_block_scaled_mm.py new file mode 100644 index 00000000..600395e3 --- /dev/null +++ b/benchmarks/TritonNPUBench/vllm/_w8a8_triton_block_scaled_mm/_w8a8_triton_block_scaled_mm.py @@ -0,0 +1,94 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/vllm_operator_cases/newtest_cases/test__w8a8_triton_block_scaled_mm.py +# Main kernel: _w8a8_triton_block_scaled_mm +# PT file: _w8a8_triton_block_scaled_mm_v2.pt + +import triton +import triton.language as tl + + +# === _w8a8_triton_block_scaled_mm === +@triton.jit +def _w8a8_triton_block_scaled_mm( + # Pointers to inputs and output + A, + B, + C, + As, + Bs, + # Shape for matmul + M, + N, + K, + # Block size for block-wise quantization + group_n, + group_k, + # Stride for inputs and output + stride_am, + stride_ak, + stride_bk, + stride_bn, + stride_cm, + stride_cn, + stride_As_m, + stride_As_k, + stride_Bs_k, + stride_Bs_n, + # Meta-parameters + BLOCK_SIZE_M: tl.constexpr, + BLOCK_SIZE_N: tl.constexpr, + BLOCK_SIZE_K: tl.constexpr, + GROUP_SIZE_M: tl.constexpr, +): + """Triton-accelerated function used to perform linear operations (dot + product) on input tensors `A` and `B` with block-wise quantization, and + store the result in output tensor `C`. + """ + + pid = tl.program_id(axis=0) + num_pid_m = tl.cdiv(M, BLOCK_SIZE_M) + num_pid_n = tl.cdiv(N, BLOCK_SIZE_N) + num_pid_in_group = GROUP_SIZE_M * num_pid_n + group_id = pid // num_pid_in_group + first_pid_m = group_id * GROUP_SIZE_M + group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M) + pid_m = first_pid_m + (pid % group_size_m) + pid_n = (pid % num_pid_in_group) // group_size_m + + offs_am = (pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)) % M + offs_bn = (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)) % N + offs_k = tl.arange(0, BLOCK_SIZE_K) + a_ptrs = A + (offs_am[:, None] * stride_am + offs_k[None, :] * stride_ak) + b_ptrs = B + (offs_k[:, None] * stride_bk + offs_bn[None, :] * stride_bn) + + As_ptrs = As + offs_am * stride_As_m + offs_bsn = offs_bn // group_n + Bs_ptrs = Bs + offs_bsn * stride_Bs_n + + accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32) + for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)): + a = tl.load(a_ptrs, mask=offs_k[None, :] < K - k * BLOCK_SIZE_K, other=0.0) + b = tl.load(b_ptrs, mask=offs_k[:, None] < K - k * BLOCK_SIZE_K, other=0.0) + + k_start = k * BLOCK_SIZE_K + offs_ks = k_start // group_k + a_s = tl.load(As_ptrs + offs_ks * stride_As_k) + b_s = tl.load(Bs_ptrs + offs_ks * stride_Bs_k) + + accumulator += tl.dot(a, b) * a_s[:, None] * b_s[None, :] + a_ptrs += BLOCK_SIZE_K * stride_ak + b_ptrs += BLOCK_SIZE_K * stride_bk + + if C.dtype.element_ty == tl.bfloat16: + c = accumulator.to(tl.bfloat16) + elif C.dtype.element_ty == tl.float16: + c = accumulator.to(tl.float16) + else: + c = accumulator.to(tl.float32) + + offs_cm = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) + offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) + c_ptrs = C + stride_cm * offs_cm[:, None] + stride_cn * offs_cn[None, :] + c_mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < N) + tl.store(c_ptrs, c, mask=c_mask) + diff --git a/benchmarks/TritonNPUBench/vllm/awq_dequantize_kernel/awq_dequantize_kernel.py b/benchmarks/TritonNPUBench/vllm/awq_dequantize_kernel/awq_dequantize_kernel.py new file mode 100644 index 00000000..af17035b --- /dev/null +++ b/benchmarks/TritonNPUBench/vllm/awq_dequantize_kernel/awq_dequantize_kernel.py @@ -0,0 +1,106 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/vllm_operator_cases/newtest_cases/test_awq_dequantize_kernel.py +# Main kernel: awq_dequantize_kernel +# PT file: awq_dequantize_kernel_v2.pt + +import triton +import triton.language as tl + + +# === awq_dequantize_kernel === +@triton.jit +def awq_dequantize_kernel( + qweight_ptr, # quantized matrix + scales_ptr, # scales, per group + zeros_ptr, # zeros, per group + group_size, # Should always be one of the supported group sizes + result_ptr, # Output matrix + num_cols, # input num cols in qweight + num_rows, # input num rows in qweight + BLOCK_SIZE_X: tl.constexpr, + BLOCK_SIZE_Y: tl.constexpr, +): + # Set up the pids. + pid_x = tl.program_id(axis=0) + pid_y = tl.program_id(axis=1) + + # Compute offsets and masks for qweight_ptr. + offsets_y = pid_y * BLOCK_SIZE_Y + tl.arange(0, BLOCK_SIZE_Y) + offsets_x = pid_x * BLOCK_SIZE_X + tl.arange(0, BLOCK_SIZE_X) + offsets = num_cols * offsets_y[:, None] + offsets_x[None, :] + + masks_y = offsets_y < num_rows + masks_x = offsets_x < num_cols + + masks = masks_y[:, None] & masks_x[None, :] + + # Compute offsets and masks for result output ptr. + result_offsets_y = pid_y * BLOCK_SIZE_Y + tl.arange(0, BLOCK_SIZE_Y) + result_offsets_x = pid_x * BLOCK_SIZE_X * 8 + tl.arange(0, BLOCK_SIZE_X * 8) + result_offsets = ( + 8 * num_cols * result_offsets_y[:, None] + result_offsets_x[None, :] + ) + + result_masks_y = result_offsets_y < num_rows + result_masks_x = result_offsets_x < num_cols * 8 + result_masks = result_masks_y[:, None] & result_masks_x[None, :] + + # Load the weights. + iweights = tl.load(qweight_ptr + offsets, masks, 0.0) + iweights = tl.interleave(iweights, iweights) + iweights = tl.interleave(iweights, iweights) + iweights = tl.interleave(iweights, iweights) + + # Create reverse AWQ order as tensor: [0, 4, 1, 5, 2, 6, 3, 7] + # that will map given indices to the correct order. + reverse_awq_order_tensor = ( + (tl.arange(0, 2) * 4)[None, :] + tl.arange(0, 4)[:, None] + ).reshape(8) + + # Use this to compute a set of shifts that can be used to unpack and + # reorder the values in iweights and zeros. + shifts = reverse_awq_order_tensor * 4 + shifts = tl.broadcast_to(shifts[None, :], (BLOCK_SIZE_Y * BLOCK_SIZE_X, 8)) + shifts = tl.reshape(shifts, (BLOCK_SIZE_Y, BLOCK_SIZE_X * 8)) + + # Unpack and reorder: shift out the correct 4-bit value and mask. + iweights = (iweights >> shifts) & 0xF + + # Compute zero offsets and masks. + zero_offsets_y = pid_y * BLOCK_SIZE_Y // group_size + tl.arange(0, 1) + zero_offsets_x = pid_x * BLOCK_SIZE_X + tl.arange(0, BLOCK_SIZE_X) + zero_offsets = num_cols * zero_offsets_y[:, None] + zero_offsets_x[None, :] + + zero_masks_y = zero_offsets_y < num_rows // group_size + zero_masks_x = zero_offsets_x < num_cols + zero_masks = zero_masks_y[:, None] & zero_masks_x[None, :] + + # Load the zeros. + zeros = tl.load(zeros_ptr + zero_offsets, zero_masks, 0.0) + zeros = tl.interleave(zeros, zeros) + zeros = tl.interleave(zeros, zeros) + zeros = tl.interleave(zeros, zeros) + zeros = tl.broadcast_to(zeros, (BLOCK_SIZE_Y, BLOCK_SIZE_X * 8)) + + # Unpack and reorder: shift out the correct 4-bit value and mask. + zeros = (zeros >> shifts) & 0xF + + # Compute scale offsets and masks. + scale_offsets_y = pid_y * BLOCK_SIZE_Y // group_size + tl.arange(0, 1) + scale_offsets_x = pid_x * BLOCK_SIZE_X * 8 + tl.arange(0, BLOCK_SIZE_X * 8) + scale_offsets = num_cols * 8 * scale_offsets_y[:, None] + scale_offsets_x[None, :] + scale_masks_y = scale_offsets_y < num_rows // group_size + scale_masks_x = scale_offsets_x < num_cols * 8 + scale_masks = scale_masks_y[:, None] & scale_masks_x[None, :] + + # Load the scales. + scales = tl.load(scales_ptr + scale_offsets, scale_masks, 0.0) + scales = tl.broadcast_to(scales, (BLOCK_SIZE_Y, BLOCK_SIZE_X * 8)) + + # Dequantize. + iweights = (iweights - zeros) * scales + iweights = iweights.to(result_ptr.type.element_ty) + + # Finally, store. + tl.store(result_ptr + result_offsets, iweights, result_masks) + diff --git a/benchmarks/TritonNPUBench/vllm/awq_gemm_kernel/awq_gemm_kernel.py b/benchmarks/TritonNPUBench/vllm/awq_gemm_kernel/awq_gemm_kernel.py new file mode 100644 index 00000000..be218b49 --- /dev/null +++ b/benchmarks/TritonNPUBench/vllm/awq_gemm_kernel/awq_gemm_kernel.py @@ -0,0 +1,128 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/vllm_operator_cases/newtest_cases/test_awq_gemm_kernel.py +# Main kernel: awq_gemm_kernel +# PT file: awq_gemm_kernel_v3.pt + +import triton +import triton.language as tl + + +# === awq_gemm_kernel === +@triton.jit +def awq_gemm_kernel( + a_ptr, + b_ptr, + c_ptr, + zeros_ptr, + scales_ptr, + M, + N, + K, + group_size, + BLOCK_SIZE_M: tl.constexpr, + BLOCK_SIZE_N: tl.constexpr, + BLOCK_SIZE_K: tl.constexpr, + SPLIT_K: tl.constexpr, +): + pid = tl.program_id(axis=0) + pid_z = tl.program_id(1) + + # NOTE: This doesn't work in TRITON_INTERPRET=1 mode. Use below instead. + # num_pid_n = (N + BLOCK_SIZE_N - 1) // BLOCK_SIZE_N + num_pid_n = tl.cdiv(N, BLOCK_SIZE_N) + + pid_m = pid // num_pid_n + pid_n = pid % num_pid_n + + accumulator_dtype = tl.float32 + + # NOTE: This doesn't work in TRITON_INTERPRET=1 mode. Use below instead. + # accumulator = tl.arange(0, BLOCK_SIZE_N) + # accumulator = tl.broadcast_to(accumulator[None, :], + # (BLOCK_SIZE_M, BLOCK_SIZE_N)) + # accumulator = accumulator & 0x0 + # accumulator = accumulator.to(accumulator_dtype) + accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=accumulator_dtype) + + # Create reverse AWQ order as tensor: [0, 4, 1, 5, 2, 6, 3, 7] + # that will map given indices to the correct order. + reverse_awq_order_tensor = ( + (tl.arange(0, 2) * 4)[None, :] + tl.arange(0, 4)[:, None] + ).reshape(8) + + # Create the necessary shifts to use to unpack. + shifts = reverse_awq_order_tensor * 4 + shifts = tl.broadcast_to(shifts[None, :], (BLOCK_SIZE_K * (BLOCK_SIZE_N // 8), 8)) + shifts = tl.reshape(shifts, (BLOCK_SIZE_K, BLOCK_SIZE_N)) + + # Offsets and masks. + offsets_am = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) + masks_am = offsets_am < M + + offsets_bn = pid_n * (BLOCK_SIZE_N // 8) + tl.arange(0, BLOCK_SIZE_N // 8) + masks_bn = offsets_bn < N // 8 + + offsets_zn = pid_n * (BLOCK_SIZE_N // 8) + tl.arange(0, BLOCK_SIZE_N // 8) + masks_zn = offsets_zn < N // 8 + + offsets_sn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) + masks_sn = offsets_sn < N + + offsets_k = pid_z * BLOCK_SIZE_K + tl.arange(0, BLOCK_SIZE_K) + offsets_a = K * offsets_am[:, None] + offsets_k[None, :] + offsets_b = (N // 8) * offsets_k[:, None] + offsets_bn[None, :] + + a_ptrs = a_ptr + offsets_a + b_ptrs = b_ptr + offsets_b + + for k in range(0, tl.cdiv(K, BLOCK_SIZE_K * SPLIT_K)): + masks_k = offsets_k < K + masks_a = masks_am[:, None] & masks_k[None, :] + a = tl.load(a_ptrs, mask=masks_a, other=0.0) + + masks_b = masks_k[:, None] & masks_bn[None, :] + b = tl.load(b_ptrs, mask=masks_b, other=0.0) + b = tl.interleave(b, b) + b = tl.interleave(b, b) + b = tl.interleave(b, b) + + # Dequantize b. + offsets_szk = ( + BLOCK_SIZE_K * SPLIT_K * k + pid_z * BLOCK_SIZE_K + ) // group_size + tl.arange(0, 1) + offsets_z = (N // 8) * offsets_szk[:, None] + offsets_zn[None, :] + masks_zk = offsets_szk < K // group_size + masks_z = masks_zk[:, None] & masks_zn[None, :] + zeros_ptrs = zeros_ptr + offsets_z + zeros = tl.load(zeros_ptrs, mask=masks_z, other=0.0) + zeros = tl.interleave(zeros, zeros) + zeros = tl.interleave(zeros, zeros) + zeros = tl.interleave(zeros, zeros) + zeros = tl.broadcast_to(zeros, (BLOCK_SIZE_K, BLOCK_SIZE_N)) + + offsets_s = N * offsets_szk[:, None] + offsets_sn[None, :] + masks_sk = offsets_szk < K // group_size + masks_s = masks_sk[:, None] & masks_sn[None, :] + scales_ptrs = scales_ptr + offsets_s + scales = tl.load(scales_ptrs, mask=masks_s, other=0.0) + scales = tl.broadcast_to(scales, (BLOCK_SIZE_K, BLOCK_SIZE_N)) + + b = (b >> shifts) & 0xF + zeros = (zeros >> shifts) & 0xF + b = (b - zeros) * scales + b = b.to(c_ptr.type.element_ty) + + # Accumulate results. + accumulator = tl.dot(a, b, accumulator, out_dtype=accumulator_dtype) + + offsets_k += BLOCK_SIZE_K * SPLIT_K + a_ptrs += BLOCK_SIZE_K * SPLIT_K + b_ptrs += BLOCK_SIZE_K * SPLIT_K * (N // 8) + + c = accumulator.to(c_ptr.type.element_ty) + offs_cm = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) + offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) + c_ptrs = c_ptr + pid_z * N * M + N * offs_cm[:, None] + offs_cn[None, :] + c_mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < N) + tl.store(c_ptrs, c, mask=c_mask) + diff --git a/benchmarks/TritonNPUBench/vllm/batched_triton_kernel/batched_triton_kernel.py b/benchmarks/TritonNPUBench/vllm/batched_triton_kernel/batched_triton_kernel.py new file mode 100644 index 00000000..0fc9b547 --- /dev/null +++ b/benchmarks/TritonNPUBench/vllm/batched_triton_kernel/batched_triton_kernel.py @@ -0,0 +1,351 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/vllm_operator_cases/newtest_cases/test_batched_triton_kernel.py +# Main kernel: batched_triton_kernel +# PT file: test_batched_triton_kernel_v2.pt + +import triton +import triton.language as tl + + +# === batched_triton_kernel === +@triton.jit +def batched_triton_kernel( + a_ptr, # [E, max_num_tokens, K] + b_ptr, # [E, K, N] + c_ptr, # [E, max_num_tokens, N] + expert_num_tokens, # [E] + compute_type: tl.constexpr, + # Dimensions + max_num_tokens, + K, + N, + # Quantization data + a_scale_ptr, + b_scale_ptr, + b_zp_ptr, + # The stride variables represent how much to increase the ptr by when + # moving by 1 element in a particular dimension. E.g. `stride_am` is + # how much to increase `a_ptr` by to get the element one row down + # (A has M rows). + stride_ae: tl.int64, + stride_am: tl.int64, + stride_ak: tl.int64, + stride_be: tl.int64, + stride_bk: tl.int64, + stride_bn: tl.int64, + stride_ce: tl.int64, + stride_cm: tl.int64, + stride_cn: tl.int64, + stride_ase: tl.int64, + stride_asm: tl.int64, + stride_ask: tl.int64, + stride_bse: tl.int64, + stride_bsk: tl.int64, + stride_bsn: tl.int64, + # Blockwise quantization data + group_n: tl.constexpr, + group_k: tl.constexpr, + # Quantization schemes + use_fp8_w8a8: tl.constexpr, + use_int8_w8a16: tl.constexpr, + per_act_token_quant: tl.constexpr, + # Kernel config + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, + BLOCK_K: tl.constexpr, +): + expert_id = tl.program_id(axis=0) + e_num_tokens = tl.load(expert_num_tokens + expert_id) + if e_num_tokens == 0: + # Early exit + return + + # axis 1 is M_blocks * N_blocks + pid_mn = tl.program_id(axis=1) + # num_pid_m = tl.cdiv(max_num_tokens, BLOCK_M) + num_pid_n = tl.cdiv(N, BLOCK_N) + pid_m = pid_mn // num_pid_n + pid_n = pid_mn % num_pid_n + + cta_m_start = pid_m * BLOCK_M + cta_n_start = pid_n * BLOCK_N + if cta_m_start >= e_num_tokens: + # Early exit + return + + cta_m_size = min(BLOCK_M, e_num_tokens - cta_m_start) + cta_n_size = min(BLOCK_N, N - cta_n_start) + + a_ptr = a_ptr + expert_id * stride_ae + cta_m_start * stride_am + b_ptr = b_ptr + expert_id * stride_be + cta_n_start * stride_bn + c_ptr = ( + c_ptr + + expert_id * stride_ce + + cta_m_start * stride_cm + + cta_n_start * stride_cn + ) + + offs_bn = (pid_n * BLOCK_N + tl.arange(0, BLOCK_N).to(tl.int64)) % N + + if use_fp8_w8a8: + a_scale_ptr = a_scale_ptr + expert_id * stride_ase + b_scale_ptr = b_scale_ptr + expert_id * stride_bse + + # block-wise + if group_k > 0 and group_n > 0 or per_act_token_quant: + a_scale_ptr = a_scale_ptr + cta_m_start * stride_asm + + expert_triton_kernel( + a_ptr, + b_ptr, + c_ptr, + expert_id, + compute_type, + cta_m_size, # M + cta_n_size, # N + K, # K + a_scale_ptr, + b_scale_ptr, + b_zp_ptr, + # Strides + stride_am, + stride_ak, + stride_bk, + stride_bn, + stride_cm, + stride_cn, + stride_ase, + stride_asm, + stride_ask, + stride_bse, + stride_bsk, + stride_bsn, + # offsets + offs_bn, + # Blockwise quantization data + group_n, + group_k, + # Quantization schemes + use_fp8_w8a8, + use_int8_w8a16, + per_act_token_quant, + # Kernel config + BLOCK_M, + BLOCK_N, + BLOCK_K, + ) + + +# === expert_triton_kernel === +@triton.jit +def expert_triton_kernel( + a_ptr, # [max_tokens, K] + b_ptr, # [K, N] + c_ptr, # [max_tokens, N] + expert_id, + compute_type: tl.constexpr, + # Dimensions + M, + N, + K, + # Quantization data + a_scale_ptr, + b_scale_ptr, + b_zp_ptr, + # strides + stride_am: tl.int64, + stride_ak: tl.int64, + stride_bk: tl.int64, + stride_bn: tl.int64, + stride_cm: tl.int64, + stride_cn: tl.int64, + stride_ase: tl.int64, + stride_asm: tl.int64, + stride_ask: tl.int64, + stride_bse: tl.int64, + stride_bsk: tl.int64, + stride_bsn: tl.int64, + # offsets + offs_bn, + # Blockwise quantization data + group_n, + group_k, + # Quantization schemes + use_fp8_w8a8: tl.constexpr, + use_int8_w8a16: tl.constexpr, + per_act_token_quant: tl.constexpr, + # Kernel config + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, + BLOCK_K: tl.constexpr, +): + offs_m = tl.arange(0, BLOCK_M) + offs_n = tl.arange(0, BLOCK_N) % N + offs_k = tl.arange(0, BLOCK_K) + mask_m = offs_m < M + + # Make grids of a + b pointers + a_ptrs = a_ptr + offs_m[:, None] * stride_am + offs_k[None, :] * stride_ak + b_ptrs = b_ptr + offs_k[:, None] * stride_bk + offs_n[None, :] * stride_bn + + accumulator = moe_mmk( + a_ptrs, + b_ptrs, + K, + expert_id, + a_scale_ptr, + b_scale_ptr, + # The stride variables represent how much to increase the ptr by when + # moving by 1 element in a particular dimension. E.g. `stride_am` is + # how much to increase `a_ptr` by to get the element one row down + # (A has M rows). + stride_ak, + stride_bk, + stride_ase, + stride_asm, + stride_ask, + stride_bse, + stride_bsk, + stride_bsn, + # Offsets and masks + offs_m, + offs_n, + offs_bn, + mask_m, + # Block size for block-wise quantization + group_n, + group_k, + # Meta-parameters + BLOCK_M, + BLOCK_N, + BLOCK_K, + compute_type, + use_fp8_w8a8, + use_int8_w8a16, + per_act_token_quant, + ) + + # store in C + offs_cn = tl.arange(0, BLOCK_N) + c_ptrs = c_ptr + offs_m[:, None] * stride_cm + offs_cn[None, :] * stride_cn + c_mask = mask_m[:, None] & (offs_cn[None, :] < N) + tl.store(c_ptrs, accumulator, mask=c_mask) + + +# === moe_mmk === +@triton.jit +def moe_mmk( + a_ptrs, + b_ptrs, + K, + expert_id, + a_scale_ptr, + b_scale_ptr, + # The stride variables represent how much to increase the ptr by when + # moving by 1 element in a particular dimension. E.g. `stride_am` is + # how much to increase `a_ptr` by to get the element one row down + # (A has M rows). + stride_ak: tl.int64, + stride_bk: tl.int64, + stride_ase: tl.int64, + stride_asm: tl.int64, + stride_ask: tl.int64, + stride_bse: tl.int64, + stride_bsk: tl.int64, + stride_bsn: tl.int64, + # Offsets and masks + offs_m, + offs_n, + offs_bn, + mask_m, + # Block size for block-wise quantization + group_n: tl.constexpr, + group_k: tl.constexpr, + # Meta-parameters + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, + BLOCK_K: tl.constexpr, + compute_type: tl.constexpr, + use_w8a8: tl.constexpr, + use_w8a16: tl.constexpr, + per_act_token_quant: tl.constexpr, +): + offs_k = tl.arange(0, BLOCK_K) + + if use_w8a16: + b_scale_ptrs = ( + b_scale_ptr + expert_id * stride_bse + offs_n[None, :] * stride_bsn + ) + b_scale = tl.load(b_scale_ptrs) + + if use_w8a8: + # block-wise + if group_k > 0 and group_n > 0: + a_scale_ptrs = a_scale_ptr + offs_m * stride_asm + offs_bsn = offs_bn // group_n + b_scale_ptrs = b_scale_ptr + offs_bsn * stride_bsn + + # per act token + elif per_act_token_quant: + # Load per-token scale for activations + a_scale_ptrs = a_scale_ptr + offs_m * stride_asm + a_scale = tl.load(a_scale_ptrs, mask=mask_m, other=0.0)[:, None] + + b_scale_ptrs = b_scale_ptr + offs_bn[None, :] * stride_bsn + b_scale = tl.load(b_scale_ptrs) + + # tensor-wise + else: + a_scale = tl.load(a_scale_ptr) + b_scale = tl.load(b_scale_ptr) + + # ----------------------------------------------------------- + # Iterate to compute a block of the C matrix. + # We accumulate into a `[BLOCK_SIZE_M, BLOCK_SIZE_N]` block + # of fp32 values for higher accuracy. + # `accumulator` will be converted back to fp16 after the loop. + accumulator = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32) + for k in range(0, tl.cdiv(K, BLOCK_K)): + # Load the next block of A and B, generate a mask by checking the + # K dimension. + a = tl.load( + a_ptrs, + mask=mask_m[:, None] & (offs_k[None, :] < K - k * BLOCK_K), + other=0.0, + ) + b = tl.load(b_ptrs, mask=offs_k[:, None] < K - k * BLOCK_K, other=0.0) + # We accumulate along the K dimension. + if use_w8a16: + accumulator = tl.dot(a, b.to(compute_type), acc=accumulator) + elif use_w8a8: + if group_k > 0 and group_n > 0: + k_start = k * BLOCK_K + offs_ks = k_start // group_k + a_scale = tl.load( + a_scale_ptrs + offs_ks * stride_ask, mask=mask_m, other=0.0 + ) + b_scale = tl.load(b_scale_ptrs + offs_ks * stride_bsk) + + accumulator += tl.dot(a, b) * a_scale[:, None] * b_scale[None, :] + else: + # acc used to enable fp8_fast_accum + accumulator = tl.dot(a, b, acc=accumulator) + else: + accumulator += tl.dot(a, b) + + # Advance the ptrs to the next K block. + a_ptrs += BLOCK_K * stride_ak + b_ptrs += BLOCK_K * stride_bk + + if use_w8a16: + accumulator = (accumulator * b_scale).to(compute_type) + elif use_w8a8: + if group_k > 0 and group_n > 0: + accumulator = accumulator.to(compute_type) + else: + accumulator = (accumulator * a_scale * b_scale).to(compute_type) + else: + accumulator = accumulator.to(compute_type) + + return accumulator + diff --git a/benchmarks/TritonNPUBench/vllm/bmm_kernel/bmm_kernel.py b/benchmarks/TritonNPUBench/vllm/bmm_kernel/bmm_kernel.py new file mode 100644 index 00000000..6ad4795e --- /dev/null +++ b/benchmarks/TritonNPUBench/vllm/bmm_kernel/bmm_kernel.py @@ -0,0 +1,142 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/vllm_operator_cases/newtest_cases/test_bmm_kernel.py +# Main kernel: bmm_kernel +# PT file: bmm_kernel_v3.pt + +import triton +import triton.language as tl + + +# === bmm_kernel === +@triton.jit +def bmm_kernel( + a_ptr, # (*, ) pointer to A, (B, M, K) + b_ptr, # (*, ) pointer to B, (B, K, N) + c_ptr, # (*, ) pointer to C, (B, M, N) + B, # int, batch size + M, # int, output rows + N, # int, output cols + K, # int, reduction dim + stride_ab, + stride_am, + stride_ak, + stride_bb, + stride_bk, + stride_bn, + stride_cb, + stride_cm, + stride_cn, + BLOCK_SIZE_M: tl.constexpr, + BLOCK_SIZE_N: tl.constexpr, + BLOCK_SIZE_K: tl.constexpr, + A_LARGE: tl.constexpr, + B_LARGE: tl.constexpr, + C_LARGE: tl.constexpr, +): + """Batched GEMM: (B, M, K) x (B, K, N) -> (B, M, N) + + Each program computes one (batch_idx, tile_m, tile_n) tile, accumulating + along K in a fixed order to preserve batch invariance. + """ + pid_b = tl.program_id(0) + pid = tl.program_id(1) + + if pid_b >= B: + return + + # number of tiles along M / N + num_pid_m = tl.cdiv(M, BLOCK_SIZE_M) + num_pid_n = tl.cdiv(N, BLOCK_SIZE_N) + + pid_m = pid // num_pid_n + pid_n = pid % num_pid_n + + if pid_m >= num_pid_m or pid_n >= num_pid_n: + return + + # offs_m / offs_n: raw global row/col indices for this tile + offs_m = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) + offs_n = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) + # masks for valid logical rows/cols within (M, N) + mask_m = offs_m < M # [BLOCK_SIZE_M] + mask_n = offs_n < N # [BLOCK_SIZE_N] + + if A_LARGE or (B_LARGE or C_LARGE): + offs_m = offs_m.to(tl.int64) + offs_n = offs_n.to(tl.int64) + + offs_m = tl.where(mask_m, offs_m, 0) + offs_n = tl.where(mask_n, offs_n, 0) + + # hint for triton contiguous memory + offs_m = tl.max_contiguous(tl.multiple_of(offs_m, BLOCK_SIZE_M), BLOCK_SIZE_M) + offs_n = tl.max_contiguous(tl.multiple_of(offs_n, BLOCK_SIZE_N), BLOCK_SIZE_N) + + # base pointers for current batch, shape-wise: + # a_batch_ptr points to A[pid_b, 0, 0] + # b_batch_ptr points to B[pid_b, 0, 0] + # c_batch_ptr points to C[pid_b, 0, 0] + a_batch_ptr = a_ptr + pid_b * stride_ab + b_batch_ptr = b_ptr + pid_b * stride_bb + c_batch_ptr = c_ptr + pid_b * stride_cb + + accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32) + # number of K-blocks this tile iterates over + k_tiles = tl.cdiv(K, BLOCK_SIZE_K) + offs_k_mask = tl.arange(0, BLOCK_SIZE_K) + + for ki in range(k_tiles): + if A_LARGE or B_LARGE: + # offs_k: [BLOCK_SIZE_K], global K indices + offs_k = ki * BLOCK_SIZE_K + tl.arange(0, BLOCK_SIZE_K).to(tl.int64) + else: + offs_k = ki * BLOCK_SIZE_K + tl.arange(0, BLOCK_SIZE_K) + + # a_ptrs: [BLOCK_SIZE_M, BLOCK_SIZE_K] + # element (i, j) points to A[pid_b, offs_m[i], offs_k[j]] + a_ptrs = a_batch_ptr + ( + offs_m[:, None] * stride_am + offs_k[None, :] * stride_ak + ) + # b_ptrs: [BLOCK_SIZE_K, BLOCK_SIZE_N] + # element (i, j) points to B[pid_b, offs_k[i], offs_n[j]] + b_ptrs = b_batch_ptr + ( + offs_k[:, None] * stride_bk + offs_n[None, :] * stride_bn + ) + + # valid K lanes for this block + k_valid = offs_k_mask < (K - ki * BLOCK_SIZE_K) + # A mask within (M, K): [BLOCK_SIZE_M, BLOCK_SIZE_K] + a_mask = mask_m[:, None] & k_valid[None, :] + # B mask within (K, N): [BLOCK_SIZE_K, BLOCK_SIZE_N] + b_mask = k_valid[:, None] & mask_n[None, :] + + # a: [BLOCK_SIZE_M, BLOCK_SIZE_K] from A[offs_m, offs_k] + a = tl.load( + a_ptrs, + mask=a_mask, + other=0.0, + ) + # b: [BLOCK_SIZE_K, BLOCK_SIZE_N] from B[offs_k, offs_n] + b = tl.load( + b_ptrs, + mask=b_mask, + other=0.0, + ) + accumulator = tl.dot(a, b, accumulator) + + # c_m / c_n: [BLOCK_SIZE_M] / [BLOCK_SIZE_N], row/col indices for C + c_m = offs_m + c_n = offs_n + if C_LARGE: + c_m = c_m.to(tl.int64) + c_n = c_n.to(tl.int64) + + # c_ptrs: [BLOCK_SIZE_M, BLOCK_SIZE_N] + # element (i, j) points to C[pid_b, c_m[i], c_n[j]] + c_ptrs = c_batch_ptr + stride_cm * c_m[:, None] + stride_cn * c_n[None, :] + # mask out elements that fall outside logical (M, N) range + c_mask = mask_m[:, None] & mask_n[None, :] + # cast FP32 accumulator back to original dtype of C + c = accumulator.to(c_ptr.dtype.element_ty) + tl.store(c_ptrs, c, mask=c_mask) + diff --git a/benchmarks/TritonNPUBench/vllm/chunk_fwd_kernel_o/chunk_fwd_kernel_o.py b/benchmarks/TritonNPUBench/vllm/chunk_fwd_kernel_o/chunk_fwd_kernel_o.py new file mode 100644 index 00000000..acde3ef2 --- /dev/null +++ b/benchmarks/TritonNPUBench/vllm/chunk_fwd_kernel_o/chunk_fwd_kernel_o.py @@ -0,0 +1,109 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/vllm_operator_cases/newtest_cases/test_chunk_fwd_kernel_o.py +# Main kernel: chunk_fwd_kernel_o +# PT file: chunk_fwd_kernel_o_v3.pt + +import triton +import triton.language as tl + + +# === chunk_fwd_kernel_o === +@triton.jit(do_not_specialize=["T"]) +def chunk_fwd_kernel_o( + q, + k, + v, + h, + g, + o, + cu_seqlens, + chunk_indices, + scale, + T, + H: tl.constexpr, + Hg: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + USE_G: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_v, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_b, i_h = i_bh // H, i_bh % H + + if IS_VARLEN: + i_tg = i_t + i_n, i_t = ( + tl.load(chunk_indices + i_t * 2).to(tl.int32), + tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), + ) + bos, eos = ( + tl.load(cu_seqlens + i_n).to(tl.int32), + tl.load(cu_seqlens + i_n + 1).to(tl.int32), + ) + T = eos - bos + NT = tl.cdiv(T, BT) + else: + NT = tl.cdiv(T, BT) + i_tg = i_b * NT + i_t + bos, eos = i_b * T, i_b * T + T + + # offset calculation + q += (bos * Hg + i_h // (H // Hg)) * K + k += (bos * Hg + i_h // (H // Hg)) * K + v += (bos * H + i_h) * V + o += (bos * H + i_h) * V + h += (i_tg * H + i_h).to(tl.int64) * K * V + + b_o = tl.zeros([BT, BV], dtype=tl.float32) + b_A = tl.zeros([BT, BT], dtype=tl.float32) + + for i_k in range(tl.cdiv(K, BK)): + p_q = tl.make_block_ptr( + q, (T, K), (Hg * K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0) + ) + p_k = tl.make_block_ptr( + k, (K, T), (1, Hg * K), (i_k * BK, i_t * BT), (BK, BT), (0, 1) + ) + p_h = tl.make_block_ptr( + h, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0) + ) + # [BT, BK] + b_q = tl.load(p_q, boundary_check=(0, 1)) + # [BK, BT] + b_k = tl.load(p_k, boundary_check=(0, 1)) + # [BK, BV] + b_h = tl.load(p_h, boundary_check=(0, 1)) + + # [BT, BK] @ [BK, BV] -> [BT, BV] + b_o += tl.dot(b_q, b_h) + # [BT, BK] @ [BK, BT] -> [BT, BT] + b_A += tl.dot(b_q, b_k) + + if USE_G: + g += bos * H + i_h + p_g = tl.make_block_ptr(g, (T,), (H,), (i_t * BT,), (BT,), (0,)) + b_g = tl.load(p_g, boundary_check=(0,)) + b_o = b_o * tl.exp(b_g)[:, None] + b_A = b_A * tl.exp(b_g[:, None] - b_g[None, :]) + + o_t = i_t * BT + tl.arange(0, BT) + m_t = o_t < T + m_A = (o_t[:, None] >= o_t[None, :]) & (m_t[:, None] & m_t) + b_A = tl.where(m_A, b_A, 0) + + p_v = tl.make_block_ptr( + v, (T, V), (H * V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0) + ) + p_o = tl.make_block_ptr( + o, (T, V), (H * V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0) + ) + b_v = tl.load(p_v, boundary_check=(0, 1)) + + # to fix mma -> mma layout conversion + # already solved by triton v3.2 or higher + b_o = b_o * scale + tl.dot(b_A.to(b_v.dtype), b_v) * scale + tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1)) + diff --git a/benchmarks/TritonNPUBench/vllm/chunk_gated_delta_rule_fwd_kernel_h_blockdim64/chunk_gated_delta_rule_fwd_kernel_h_blockdim64.py b/benchmarks/TritonNPUBench/vllm/chunk_gated_delta_rule_fwd_kernel_h_blockdim64/chunk_gated_delta_rule_fwd_kernel_h_blockdim64.py new file mode 100644 index 00000000..27d522a6 --- /dev/null +++ b/benchmarks/TritonNPUBench/vllm/chunk_gated_delta_rule_fwd_kernel_h_blockdim64/chunk_gated_delta_rule_fwd_kernel_h_blockdim64.py @@ -0,0 +1,248 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/vllm_operator_cases/newtest_cases/test_chunk_gated_delta_rule_fwd_kernel_h_blockdim64.py +# Main kernel: chunk_gated_delta_rule_fwd_kernel_h_blockdim64 +# PT file: chunk_gated_delta_rule_fwd_kernel_h_blockdim64_v2.pt + +import triton +import triton.language as tl + + +# === chunk_gated_delta_rule_fwd_kernel_h_blockdim64 === +@triton.jit(do_not_specialize=["T"]) +def chunk_gated_delta_rule_fwd_kernel_h_blockdim64( + k, + v, + w, + v_new, + g, + gk, + h, + h0, + ht, + cu_seqlens, + chunk_offsets, + T, + H: tl.constexpr, + Hg: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BV: tl.constexpr, + USE_G: tl.constexpr, + USE_GK: tl.constexpr, + USE_INITIAL_STATE: tl.constexpr, + STORE_FINAL_STATE: tl.constexpr, + SAVE_NEW_VALUE: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_v, i_nh = tl.program_id(0), tl.program_id(1) + i_n, i_h = i_nh // H, i_nh % H + if IS_VARLEN: + bos, eos = ( + tl.load(cu_seqlens + i_n).to(tl.int32), + tl.load(cu_seqlens + i_n + 1).to(tl.int32), + ) + T = eos - bos + NT = tl.cdiv(T, BT) + boh = tl.load(chunk_offsets + i_n).to(tl.int32) + else: + bos, eos = i_n * T, i_n * T + T + NT = (T + BT - 1) // BT + boh = i_n * NT + + # [BK, BV] + b_h1 = tl.zeros([64, BV], dtype=tl.float32) + if K > 64: + b_h2 = tl.zeros([64, BV], dtype=tl.float32) + if K > 128: + b_h3 = tl.zeros([64, BV], dtype=tl.float32) + if K > 192: + b_h4 = tl.zeros([64, BV], dtype=tl.float32) + + # calculate offset + h += ((boh * H + i_h) * K * V).to(tl.int64) + v += ((bos * H + i_h) * V).to(tl.int64) + k += ((bos * Hg + i_h // (H // Hg)) * K).to(tl.int64) + w += ((bos * H + i_h) * K).to(tl.int64) + if SAVE_NEW_VALUE: + v_new += ((bos * H + i_h) * V).to(tl.int64) + stride_v = H * V + stride_h = H * K * V + stride_k = Hg * K + stride_w = H * K + if USE_INITIAL_STATE: + h0 = h0 + i_nh * K * V + if STORE_FINAL_STATE: + ht = ht + i_nh * K * V + + # load initial state + if USE_INITIAL_STATE: + p_h0_1 = tl.make_block_ptr(h0, (K, V), (V, 1), (0, i_v * BV), (64, BV), (1, 0)) + b_h1 += tl.load(p_h0_1, boundary_check=(0, 1), padding_option='zero').to(tl.float32) + if K > 64: + p_h0_2 = tl.make_block_ptr( + h0, (K, V), (V, 1), (64, i_v * BV), (64, BV), (1, 0) + ) + b_h2 += tl.load(p_h0_2, boundary_check=(0, 1), padding_option='zero').to(tl.float32) + if K > 128: + p_h0_3 = tl.make_block_ptr( + h0, (K, V), (V, 1), (128, i_v * BV), (64, BV), (1, 0) + ) + b_h3 += tl.load(p_h0_3, boundary_check=(0, 1), padding_option='zero').to(tl.float32) + if K > 192: + p_h0_4 = tl.make_block_ptr( + h0, (K, V), (V, 1), (192, i_v * BV), (64, BV), (1, 0) + ) + b_h4 += tl.load(p_h0_4, boundary_check=(0, 1), padding_option='zero').to(tl.float32) + + # main recurrence + for i_t in range(NT): + p_h1 = tl.make_block_ptr( + h + i_t * stride_h, (K, V), (V, 1), (0, i_v * BV), (64, BV), (1, 0) + ) + tl.store(p_h1, b_h1.to(p_h1.dtype.element_ty), boundary_check=(0, 1)) + if K > 64: + p_h2 = tl.make_block_ptr( + h + i_t * stride_h, (K, V), (V, 1), (64, i_v * BV), (64, BV), (1, 0) + ) + tl.store(p_h2, b_h2.to(p_h2.dtype.element_ty), boundary_check=(0, 1)) + if K > 128: + p_h3 = tl.make_block_ptr( + h + i_t * stride_h, (K, V), (V, 1), (128, i_v * BV), (64, BV), (1, 0) + ) + tl.store(p_h3, b_h3.to(p_h3.dtype.element_ty), boundary_check=(0, 1)) + if K > 192: + p_h4 = tl.make_block_ptr( + h + i_t * stride_h, (K, V), (V, 1), (192, i_v * BV), (64, BV), (1, 0) + ) + tl.store(p_h4, b_h4.to(p_h4.dtype.element_ty), boundary_check=(0, 1)) + + p_w = tl.make_block_ptr( + w, (T, K), (stride_w, 1), (i_t * BT, 0), (BT, 64), (1, 0) + ) + b_w = tl.load(p_w, boundary_check=(0, 1), padding_option='zero') + b_v = tl.dot(b_w, b_h1.to(b_w.dtype)) + if K > 64: + p_w = tl.make_block_ptr( + w, (T, K), (stride_w, 1), (i_t * BT, 64), (BT, 64), (1, 0) + ) + b_w = tl.load(p_w, boundary_check=(0, 1), padding_option='zero') + b_v += tl.dot(b_w, b_h2.to(b_w.dtype)) + if K > 128: + p_w = tl.make_block_ptr( + w, (T, K), (stride_w, 1), (i_t * BT, 128), (BT, 64), (1, 0) + ) + b_w = tl.load(p_w, boundary_check=(0, 1), padding_option='zero') + b_v += tl.dot(b_w, b_h3.to(b_w.dtype)) + if K > 192: + p_w = tl.make_block_ptr( + w, (T, K), (stride_w, 1), (i_t * BT, 192), (BT, 64), (1, 0) + ) + b_w = tl.load(p_w, boundary_check=(0, 1), padding_option='zero') + b_v += tl.dot(b_w, b_h4.to(b_w.dtype)) + p_v = tl.make_block_ptr( + v, (T, V), (stride_v, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0) + ) + b_v = tl.load(p_v, boundary_check=(0, 1), padding_option='zero') - b_v + + if SAVE_NEW_VALUE: + p_v = tl.make_block_ptr( + v_new, (T, V), (stride_v, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0) + ) + tl.store(p_v, b_v.to(p_v.dtype.element_ty), boundary_check=(0, 1)) + + last_idx = min((i_t + 1) * BT, T) - 1 + if USE_G: + m_t = (i_t * BT + tl.arange(0, BT)) < T + b_g_last = tl.load(g + bos * H + last_idx * H + i_h) + p_g = tl.make_block_ptr( + g + bos * H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,) + ) + b_g = tl.load(p_g, boundary_check=(0,), padding_option='zero') + b_v = b_v * tl.where(m_t, tl.exp(b_g_last - b_g), 0)[:, None] + b_g_last = tl.exp(b_g_last) + b_h1 *= b_g_last + if K > 64: + b_h2 *= b_g_last + if K > 128: + b_h3 *= b_g_last + if K > 192: + b_h4 *= b_g_last + + if USE_GK: + o_k1 = tl.arange(0, 64) + b_gk_last1 = tl.load( + gk + (bos + last_idx) * H * K + i_h * K + o_k1, + mask=(o_k1 < K), + other=0.0, + ) + b_h1 *= tl.exp(b_gk_last1)[:, None] + if K > 64: + o_k2 = 64 + o_k1 + b_gk_last2 = tl.load( + gk + (bos + last_idx) * H * K + i_h * K + o_k2, + mask=(o_k2 < K), + other=0.0, + ) + b_h2 *= tl.exp(b_gk_last2)[:, None] + if K > 128: + o_k3 = 128 + o_k1 + b_gk_last3 = tl.load( + gk + (bos + last_idx) * H * K + i_h * K + o_k3, + mask=(o_k3 < K), + other=0.0, + ) + b_h3 *= tl.exp(b_gk_last3)[:, None] + if K > 192: + o_k4 = 192 + o_k1 + b_gk_last4 = tl.load( + gk + (bos + last_idx) * H * K + i_h * K + o_k4, + mask=(o_k4 < K), + other=0.0, + ) + b_h4 *= tl.exp(b_gk_last4)[:, None] + b_v = b_v.to(k.dtype.element_ty) + + p_k = tl.make_block_ptr( + k, (K, T), (1, stride_k), (0, i_t * BT), (64, BT), (0, 1) + ) + b_k = tl.load(p_k, boundary_check=(0, 1), padding_option='zero') + b_h1 += tl.dot(b_k, b_v) + if K > 64: + p_k = tl.make_block_ptr( + k, (K, T), (1, stride_k), (64, i_t * BT), (64, BT), (0, 1) + ) + b_k = tl.load(p_k, boundary_check=(0, 1), padding_option='zero') + b_h2 += tl.dot(b_k, b_v) + if K > 128: + p_k = tl.make_block_ptr( + k, (K, T), (1, stride_k), (128, i_t * BT), (64, BT), (0, 1) + ) + b_k = tl.load(p_k, boundary_check=(0, 1), padding_option='zero') + b_h3 += tl.dot(b_k, b_v) + if K > 192: + p_k = tl.make_block_ptr( + k, (K, T), (1, stride_k), (192, i_t * BT), (64, BT), (0, 1) + ) + b_k = tl.load(p_k, boundary_check=(0, 1), padding_option='zero') + b_h4 += tl.dot(b_k, b_v) + # epilogue + if STORE_FINAL_STATE: + p_ht = tl.make_block_ptr(ht, (K, V), (V, 1), (0, i_v * BV), (64, BV), (1, 0)) + tl.store(p_ht, b_h1.to(p_ht.dtype.element_ty), boundary_check=(0, 1)) + if K > 64: + p_ht = tl.make_block_ptr( + ht, (K, V), (V, 1), (64, i_v * BV), (64, BV), (1, 0) + ) + tl.store(p_ht, b_h2.to(p_ht.dtype.element_ty), boundary_check=(0, 1)) + if K > 128: + p_ht = tl.make_block_ptr( + ht, (K, V), (V, 1), (128, i_v * BV), (64, BV), (1, 0) + ) + tl.store(p_ht, b_h3.to(p_ht.dtype.element_ty), boundary_check=(0, 1)) + if K > 192: + p_ht = tl.make_block_ptr( + ht, (K, V), (V, 1), (192, i_v * BV), (64, BV), (1, 0) + ) + tl.store(p_ht, b_h4.to(p_ht.dtype.element_ty), boundary_check=(0, 1)) + diff --git a/benchmarks/TritonNPUBench/vllm/chunk_gla_fwd_kernel_o/chunk_gla_fwd_kernel_o.py b/benchmarks/TritonNPUBench/vllm/chunk_gla_fwd_kernel_o/chunk_gla_fwd_kernel_o.py new file mode 100644 index 00000000..5b57b8bd --- /dev/null +++ b/benchmarks/TritonNPUBench/vllm/chunk_gla_fwd_kernel_o/chunk_gla_fwd_kernel_o.py @@ -0,0 +1,118 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/vllm_operator_cases/newtest_cases/test_chunk_gla_fwd_kernel_o.py +# Main kernel: chunk_gla_fwd_kernel_o +# PT file: chunk_gla_fwd_kernel_o_v3.pt + +import triton +import triton.language as tl + + +# === chunk_gla_fwd_kernel_o === +@triton.jit(do_not_specialize=["T"]) +def chunk_gla_fwd_kernel_o( + q, + v, + g, + h, + o, + A, + cu_seqlens, + chunk_indices, + scale, + T, + H: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_v, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_b, i_h = i_bh // H, i_bh % H + if IS_VARLEN: + i_tg = i_t + i_n, i_t = ( + tl.load(chunk_indices + i_t * 2).to(tl.int32), + tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), + ) + bos, eos = ( + tl.load(cu_seqlens + i_n).to(tl.int32), + tl.load(cu_seqlens + i_n + 1).to(tl.int32), + ) + T = eos - bos + NT = tl.cdiv(T, BT) + else: + NT = tl.cdiv(T, BT) + i_tg = i_b * NT + i_t + bos, eos = i_b * T, i_b * T + T + + m_s = tl.arange(0, BT)[:, None] >= tl.arange(0, BT)[None, :] + + b_o = tl.zeros([BT, BV], dtype=tl.float32) + for i_k in range(tl.cdiv(K, BK)): + p_q = tl.make_block_ptr( + q + (bos * H + i_h) * K, + (T, K), + (H * K, 1), + (i_t * BT, i_k * BK), + (BT, BK), + (1, 0), + ) + p_g = tl.make_block_ptr( + g + (bos * H + i_h) * K, + (T, K), + (H * K, 1), + (i_t * BT, i_k * BK), + (BT, BK), + (1, 0), + ) + p_h = tl.make_block_ptr( + h + (i_tg * H + i_h) * K * V, + (K, V), + (V, 1), + (i_k * BK, i_v * BV), + (BK, BV), + (1, 0), + ) + + # [BT, BK] + b_q = tl.load(p_q, boundary_check=(0, 1)) + b_q = (b_q * scale).to(b_q.dtype) + # [BT, BK] + b_g = tl.load(p_g, boundary_check=(0, 1)) + # [BT, BK] + b_qg = (b_q * exp(b_g)).to(b_q.dtype) + # [BK, BV] + b_h = tl.load(p_h, boundary_check=(0, 1)) + # works but dkw, owing to divine benevolence + # [BT, BV] + if i_k >= 0: + b_o += tl.dot(b_qg, b_h.to(b_qg.dtype)) + p_v = tl.make_block_ptr( + v + (bos * H + i_h) * V, + (T, V), + (H * V, 1), + (i_t * BT, i_v * BV), + (BT, BV), + (1, 0), + ) + p_o = tl.make_block_ptr( + o + (bos * H + i_h) * V, + (T, V), + (H * V, 1), + (i_t * BT, i_v * BV), + (BT, BV), + (1, 0), + ) + p_A = tl.make_block_ptr( + A + (bos * H + i_h) * BT, (T, BT), (H * BT, 1), (i_t * BT, 0), (BT, BT), (1, 0) + ) + # [BT, BV] + b_v = tl.load(p_v, boundary_check=(0, 1)) + # [BT, BT] + b_A = tl.load(p_A, boundary_check=(0, 1)) + b_A = tl.where(m_s, b_A, 0.0).to(b_v.dtype) + b_o += tl.dot(b_A, b_v, allow_tf32=False) + tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1)) + diff --git a/benchmarks/TritonNPUBench/vllm/chunk_kda_scaled_dot_kkt_fwd_kernel_intra_sub_inter/chunk_kda_scaled_dot_kkt_fwd_kernel_intra_sub_inter.py b/benchmarks/TritonNPUBench/vllm/chunk_kda_scaled_dot_kkt_fwd_kernel_intra_sub_inter/chunk_kda_scaled_dot_kkt_fwd_kernel_intra_sub_inter.py new file mode 100644 index 00000000..4d919d02 --- /dev/null +++ b/benchmarks/TritonNPUBench/vllm/chunk_kda_scaled_dot_kkt_fwd_kernel_intra_sub_inter/chunk_kda_scaled_dot_kkt_fwd_kernel_intra_sub_inter.py @@ -0,0 +1,110 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/vllm_operator_cases/newtest_cases/test_chunk_kda_scaled_dot_kkt_fwd_kernel_intra_sub_inter.py +# Main kernel: chunk_kda_scaled_dot_kkt_fwd_kernel_intra_sub_inter +# PT file: chunk_kda_scaled_dot_kkt_fwd_kernel_intra_sub_intra_v3.pt + +import triton +import triton.language as tl + + +# === chunk_kda_scaled_dot_kkt_fwd_kernel_intra_sub_inter === +@triton.jit(do_not_specialize=["T"]) +def chunk_kda_scaled_dot_kkt_fwd_kernel_intra_sub_inter( + q, + k, + g, + beta, + A, + Aqk, + scale, + cu_seqlens, + chunk_indices, + T, + H: tl.constexpr, + K: tl.constexpr, + BT: tl.constexpr, + BC: tl.constexpr, + BK: tl.constexpr, + NC: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_t, i_c, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_b, i_h = i_bh // H, i_bh % H + i_i, i_j = i_c // NC, i_c % NC + if IS_VARLEN: + i_n, i_t = ( + tl.load(chunk_indices + i_t * 2).to(tl.int32), + tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), + ) + bos, eos = ( + tl.load(cu_seqlens + i_n).to(tl.int32), + tl.load(cu_seqlens + i_n + 1).to(tl.int32), + ) + T = eos - bos + else: + bos, eos = i_b * T, i_b * T + T + + if i_t * BT + i_i * BC >= T: + return + if i_i <= i_j: + return + + q += (bos * H + i_h) * K + k += (bos * H + i_h) * K + g += (bos * H + i_h) * K + A += (bos * H + i_h) * BT + Aqk += (bos * H + i_h) * BT + + p_b = tl.make_block_ptr( + beta + bos * H + i_h, (T,), (H,), (i_t * BT + i_i * BC,), (BC,), (0,) + ) + b_b = tl.load(p_b, boundary_check=(0,)) + + b_A = tl.zeros([BC, BC], dtype=tl.float32) + b_Aqk = tl.zeros([BC, BC], dtype=tl.float32) + for i_k in range(tl.cdiv(K, BK)): + p_q = tl.make_block_ptr( + q, (T, K), (H * K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0) + ) + p_k = tl.make_block_ptr( + k, (T, K), (H * K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0) + ) + p_g = tl.make_block_ptr( + g, (T, K), (H * K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0) + ) + b_kt = tl.make_block_ptr( + k, (K, T), (1, H * K), (i_k * BK, i_t * BT + i_j * BC), (BK, BC), (0, 1) + ) + p_gk = tl.make_block_ptr( + g, (K, T), (1, H * K), (i_k * BK, i_t * BT + i_j * BC), (BK, BC), (0, 1) + ) + + o_k = i_k * BK + tl.arange(0, BK) + m_k = o_k < K + # [BK,] + b_gn = tl.load(g + (i_t * BT + i_i * BC) * H * K + o_k, mask=m_k, other=0) + # [BC, BK] + b_g = tl.load(p_g, boundary_check=(0, 1)) + b_k = tl.load(p_k, boundary_check=(0, 1)) * exp(b_g - b_gn[None, :]) + # [BK, BC] + b_gk = tl.load(p_gk, boundary_check=(0, 1)) + b_kt = tl.load(b_kt, boundary_check=(0, 1)) + # [BC, BC] + b_ktg = b_kt * exp(b_gn[:, None] - b_gk) + b_A += tl.dot(b_k, b_ktg) + + b_q = tl.load(p_q, boundary_check=(0, 1)) + b_qg = b_q * exp(b_g - b_gn[None, :]) * scale + b_Aqk += tl.dot(b_qg, b_ktg) + + b_A *= b_b[:, None] + + p_A = tl.make_block_ptr( + A, (T, BT), (H * BT, 1), (i_t * BT + i_i * BC, i_j * BC), (BC, BC), (1, 0) + ) + tl.store(p_A, b_A.to(A.dtype.element_ty), boundary_check=(0, 1)) + p_Aqk = tl.make_block_ptr( + Aqk, (T, BT), (H * BT, 1), (i_t * BT + i_i * BC, i_j * BC), (BC, BC), (1, 0) + ) + tl.store(p_Aqk, b_Aqk.to(Aqk.dtype.element_ty), boundary_check=(0, 1)) + diff --git a/benchmarks/TritonNPUBench/vllm/chunk_kda_scaled_dot_kkt_fwd_kernel_intra_sub_intra/chunk_kda_scaled_dot_kkt_fwd_kernel_intra_sub_intra.py b/benchmarks/TritonNPUBench/vllm/chunk_kda_scaled_dot_kkt_fwd_kernel_intra_sub_intra/chunk_kda_scaled_dot_kkt_fwd_kernel_intra_sub_intra.py new file mode 100644 index 00000000..cb02b31e --- /dev/null +++ b/benchmarks/TritonNPUBench/vllm/chunk_kda_scaled_dot_kkt_fwd_kernel_intra_sub_intra/chunk_kda_scaled_dot_kkt_fwd_kernel_intra_sub_intra.py @@ -0,0 +1,100 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/vllm_operator_cases/newtest_cases/test_chunk_kda_scaled_dot_kkt_fwd_kernel_intra_sub_intra.py +# Main kernel: chunk_kda_scaled_dot_kkt_fwd_kernel_intra_sub_intra +# PT file: chunk_kda_scaled_dot_kkt_fwd_kernel_intra_sub_intra_v3.pt + +import triton +import triton.language as tl + + +# === chunk_kda_scaled_dot_kkt_fwd_kernel_intra_sub_intra === +@triton.jit(do_not_specialize=["T"]) +def chunk_kda_scaled_dot_kkt_fwd_kernel_intra_sub_intra( + q, + k, + g, + beta, + A, + Aqk, + scale, + cu_seqlens, + chunk_indices, + T, + H: tl.constexpr, + K: tl.constexpr, + BT: tl.constexpr, + BC: tl.constexpr, + BK: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_t, i_i, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_b, i_h = i_bh // H, i_bh % H + if IS_VARLEN: + i_n, i_t = ( + tl.load(chunk_indices + i_t * 2).to(tl.int32), + tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), + ) + bos, eos = ( + tl.load(cu_seqlens + i_n).to(tl.int32), + tl.load(cu_seqlens + i_n + 1).to(tl.int32), + ) + T = eos - bos + else: + bos, eos = i_b * T, i_b * T + T + + if i_t * BT + i_i * BC >= T: + return + + o_i = tl.arange(0, BC) + o_k = tl.arange(0, BK) + m_k = o_k < K + m_A = (i_t * BT + i_i * BC + o_i) < T + o_A = (bos + i_t * BT + i_i * BC + o_i) * H * BT + i_h * BT + i_i * BC + + p_q = tl.make_block_ptr( + q + (bos * H + i_h) * K, + (T, K), + (H * K, 1), + (i_t * BT + i_i * BC, 0), + (BC, BK), + (1, 0), + ) + p_k = tl.make_block_ptr( + k + (bos * H + i_h) * K, + (T, K), + (H * K, 1), + (i_t * BT + i_i * BC, 0), + (BC, BK), + (1, 0), + ) + p_g = tl.make_block_ptr( + g + (bos * H + i_h) * K, + (T, K), + (H * K, 1), + (i_t * BT + i_i * BC, 0), + (BC, BK), + (1, 0), + ) + b_q = tl.load(p_q, boundary_check=(0, 1)) + b_k = tl.load(p_k, boundary_check=(0, 1)) + b_g = tl.load(p_g, boundary_check=(0, 1)) + + p_b = beta + (bos + i_t * BT + i_i * BC + o_i) * H + i_h + b_k = b_k * tl.load(p_b, mask=m_A, other=0)[:, None] + + p_kt = k + (bos + i_t * BT + i_i * BC) * H * K + i_h * K + o_k + p_gk = g + (bos + i_t * BT + i_i * BC) * H * K + i_h * K + o_k + + for j in range(0, min(BC, T - i_t * BT - i_i * BC)): + b_kt = tl.load(p_kt, mask=m_k, other=0).to(tl.float32) + b_gk = tl.load(p_gk, mask=m_k, other=0).to(tl.float32) + b_ktg = b_kt[None, :] * exp(b_g - b_gk[None, :]) + b_A = tl.sum(b_k * b_ktg, 1) + b_A = tl.where(o_i > j, b_A, 0.0) + b_Aqk = tl.sum(b_q * b_ktg, 1) + b_Aqk = tl.where(o_i >= j, b_Aqk * scale, 0.0) + tl.store(A + o_A + j, b_A, mask=m_A) + tl.store(Aqk + o_A + j, b_Aqk, mask=m_A) + p_kt += H * K + p_gk += H * K + diff --git a/benchmarks/TritonNPUBench/vllm/chunk_local_cumsum_scalar_kernel/chunk_local_cumsum_scalar_kernel.py b/benchmarks/TritonNPUBench/vllm/chunk_local_cumsum_scalar_kernel/chunk_local_cumsum_scalar_kernel.py new file mode 100644 index 00000000..c77b3fb8 --- /dev/null +++ b/benchmarks/TritonNPUBench/vllm/chunk_local_cumsum_scalar_kernel/chunk_local_cumsum_scalar_kernel.py @@ -0,0 +1,57 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/vllm_operator_cases/newtest_cases/test_chunk_local_cumsum_scalar_kernel.py +# Main kernel: chunk_local_cumsum_scalar_kernel +# PT file: chunk_local_cumsum_scalar_kernel_gpu2.pt + +import triton +import triton.language as tl + + +# === chunk_local_cumsum_scalar_kernel === +@triton.jit(do_not_specialize=["T"]) +def chunk_local_cumsum_scalar_kernel( + s, + o, + cu_seqlens, + chunk_indices, + T, + B: tl.constexpr, + H: tl.constexpr, + BT: tl.constexpr, + REVERSE: tl.constexpr, + IS_VARLEN: tl.constexpr, + HEAD_FIRST: tl.constexpr, +): + i_t, i_bh = tl.program_id(0), tl.program_id(1) + i_b, i_h = i_bh // H, i_bh % H + if IS_VARLEN: + i_n, i_t = ( + tl.load(chunk_indices + i_t * 2).to(tl.int32), + tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), + ) + bos, eos = ( + tl.load(cu_seqlens + i_n).to(tl.int32), + tl.load(cu_seqlens + i_n + 1).to(tl.int32), + ) + T = eos - bos + else: + bos, eos = i_b * T, i_b * T + T + + if HEAD_FIRST: + p_s = tl.make_block_ptr( + s + bos * H + i_h * T, (T,), (1,), (i_t * BT,), (BT,), (0,) + ) + p_o = tl.make_block_ptr( + o + bos * H + i_h * T, (T,), (1,), (i_t * BT,), (BT,), (0,) + ) + else: + p_s = tl.make_block_ptr(s + bos * H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) + p_o = tl.make_block_ptr(o + bos * H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) + # [BT] + b_s = tl.load(p_s, boundary_check=(0,)).to(tl.float32) + b_o = tl.cumsum(b_s, axis=0) + if REVERSE: + b_z = tl.sum(b_s, axis=0) + b_o = -b_o + b_z[None] + b_s + tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0,)) + diff --git a/benchmarks/TritonNPUBench/vllm/chunk_local_cumsum_vector_kernel/chunk_local_cumsum_vector_kernel.py b/benchmarks/TritonNPUBench/vllm/chunk_local_cumsum_vector_kernel/chunk_local_cumsum_vector_kernel.py new file mode 100644 index 00000000..dcb7038c --- /dev/null +++ b/benchmarks/TritonNPUBench/vllm/chunk_local_cumsum_vector_kernel/chunk_local_cumsum_vector_kernel.py @@ -0,0 +1,86 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/vllm_operator_cases/newtest_cases/test_chunk_local_cumsum_vector_kernel.py +# Main kernel: chunk_local_cumsum_vector_kernel +# PT file: chunk_local_cumsum_vector_kernel_v2.pt + +import triton +import triton.language as tl + + +# === chunk_local_cumsum_vector_kernel === +@triton.jit(do_not_specialize=["T"]) +def chunk_local_cumsum_vector_kernel( + s, + o, + cu_seqlens, + chunk_indices, + T, + B: tl.constexpr, + H: tl.constexpr, + S: tl.constexpr, + BT: tl.constexpr, + BS: tl.constexpr, + REVERSE: tl.constexpr, + IS_VARLEN: tl.constexpr, + HEAD_FIRST: tl.constexpr, +): + i_s, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_b, i_h = i_bh // H, i_bh % H + if IS_VARLEN: + i_n, i_t = ( + tl.load(chunk_indices + i_t * 2).to(tl.int32), + tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), + ) + bos, eos = ( + tl.load(cu_seqlens + i_n).to(tl.int32), + tl.load(cu_seqlens + i_n + 1).to(tl.int32), + ) + T = eos - bos + else: + bos, eos = i_b * T, i_b * T + T + + o_i = tl.arange(0, BT) + if REVERSE: + m_s = tl.where(o_i[:, None] <= o_i[None, :], 1.0, 0.0) + else: + m_s = tl.where(o_i[:, None] >= o_i[None, :], 1.0, 0.0) + + if HEAD_FIRST: + p_s = tl.make_block_ptr( + s + (bos * H + i_h * T) * S, + (T, S), + (S, 1), + (i_t * BT, i_s * BS), + (BT, BS), + (1, 0), + ) + p_o = tl.make_block_ptr( + o + (bos * H + i_h * T) * S, + (T, S), + (S, 1), + (i_t * BT, i_s * BS), + (BT, BS), + (1, 0), + ) + else: + p_s = tl.make_block_ptr( + s + (bos * H + i_h) * S, + (T, S), + (H * S, 1), + (i_t * BT, i_s * BS), + (BT, BS), + (1, 0), + ) + p_o = tl.make_block_ptr( + o + (bos * H + i_h) * S, + (T, S), + (H * S, 1), + (i_t * BT, i_s * BS), + (BT, BS), + (1, 0), + ) + # [BT, BS] + b_s = tl.load(p_s, boundary_check=(0, 1)).to(tl.float32) + b_o = tl.dot(m_s, b_s, allow_tf32=False) + tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1)) + diff --git a/benchmarks/TritonNPUBench/vllm/chunk_scaled_dot_kkt_fwd_kernel/chunk_scaled_dot_kkt_fwd_kernel.py b/benchmarks/TritonNPUBench/vllm/chunk_scaled_dot_kkt_fwd_kernel/chunk_scaled_dot_kkt_fwd_kernel.py new file mode 100644 index 00000000..acf71df4 --- /dev/null +++ b/benchmarks/TritonNPUBench/vllm/chunk_scaled_dot_kkt_fwd_kernel/chunk_scaled_dot_kkt_fwd_kernel.py @@ -0,0 +1,76 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/vllm_operator_cases/newtest_cases/test_chunk_scaled_dot_kkt_fwd_kernel.py +# Main kernel: chunk_scaled_dot_kkt_fwd_kernel +# PT file: chunk_scaled_dot_kkt_fwd_kernel_v2.pt + +import triton +import triton.language as tl + + +# === chunk_scaled_dot_kkt_fwd_kernel === +@triton.jit(do_not_specialize=["T"]) +def chunk_scaled_dot_kkt_fwd_kernel( + k, + beta, + g, + A, + cu_seqlens, + chunk_indices, + T, + H: tl.constexpr, + Hg: tl.constexpr, + K: tl.constexpr, + BT: tl.constexpr, + BK: tl.constexpr, + IS_VARLEN: tl.constexpr, + USE_G: tl.constexpr, +): + i_t, i_bh = tl.program_id(0), tl.program_id(1) + i_b, i_h = i_bh // H, i_bh % H + if IS_VARLEN: + i_n, i_t = ( + tl.load(chunk_indices + i_t * 2).to(tl.int32), + tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), + ) + bos, eos = ( + tl.load(cu_seqlens + i_n).to(tl.int32), + tl.load(cu_seqlens + i_n + 1).to(tl.int32), + ) + T = eos - bos + else: + bos, eos = i_b * T, i_b * T + T + o_t = i_t * BT + tl.arange(0, BT) + m_t = o_t < T + + p_beta = tl.make_block_ptr( + beta + bos * H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,) + ) + b_beta = tl.load(p_beta, boundary_check=(0,)) + + b_A = tl.zeros([BT, BT], dtype=tl.float32) + for i_k in range(tl.cdiv(K, BK)): + p_k = tl.make_block_ptr( + k + (bos * Hg + i_h // (H // Hg)) * K, + (T, K), + (Hg * K, 1), + (i_t * BT, i_k * BK), + (BT, BK), + (1, 0), + ) + b_k = tl.load(p_k, boundary_check=(0, 1)) + b_kb = b_k * b_beta[:, None] + b_A += tl.dot(b_kb.to(b_k.dtype), tl.trans(b_k)) + + if USE_G: + p_g = tl.make_block_ptr(g + bos * H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) + b_g = tl.load(p_g, boundary_check=(0,)) + b_g_diff = b_g[:, None] - b_g[None, :] + b_A = b_A * tl.exp(b_g_diff) + + m_A = (o_t[:, None] > o_t[None, :]) & (m_t[:, None] & m_t) + b_A = tl.where(m_A, b_A, 0) + p_A = tl.make_block_ptr( + A + (bos * H + i_h) * BT, (T, BT), (BT * H, 1), (i_t * BT, 0), (BT, BT), (1, 0) + ) + tl.store(p_A, b_A.to(p_A.dtype.element_ty), boundary_check=(0, 1)) + diff --git a/benchmarks/TritonNPUBench/vllm/compute_identity_kernel/compute_identity_kernel.py b/benchmarks/TritonNPUBench/vllm/compute_identity_kernel/compute_identity_kernel.py new file mode 100644 index 00000000..3a720fa0 --- /dev/null +++ b/benchmarks/TritonNPUBench/vllm/compute_identity_kernel/compute_identity_kernel.py @@ -0,0 +1,48 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/vllm_operator_cases/newtest_cases/test_compute_identity_kernel.py +# Main kernel: compute_identity_kernel +# PT file: compute_identity_kernel_v2.pt + +import triton +import triton.language as tl + + +# === compute_identity_kernel === +@triton.jit +def compute_identity_kernel( + top_k: int, + hidden_states_ptr: tl.tensor, + expert_scales_ptr: tl.tensor, + num_tokens: int, + output_ptr: tl.tensor, + hidden_dim: int, + scales_stride: int, + BLOCK_SIZE: tl.constexpr, +) -> None: + pid = tl.program_id(0) + + batch_id = pid // (hidden_dim // BLOCK_SIZE) + dim_offset = pid % (hidden_dim // BLOCK_SIZE) * BLOCK_SIZE + + if batch_id >= num_tokens or dim_offset >= hidden_dim: + return + + h = tl.load( + hidden_states_ptr + + batch_id * hidden_dim + + dim_offset + + tl.arange(0, BLOCK_SIZE), + mask=(dim_offset + tl.arange(0, BLOCK_SIZE)) < hidden_dim, + ) + + result = tl.zeros([BLOCK_SIZE], dtype=tl.float32) + for i in range(top_k): + scale = tl.load(expert_scales_ptr + batch_id * scales_stride + i) + result += h * scale + + tl.store( + output_ptr + batch_id * hidden_dim + dim_offset + tl.arange(0, BLOCK_SIZE), + result, + mask=(dim_offset + tl.arange(0, BLOCK_SIZE)) < hidden_dim, + ) + diff --git a/benchmarks/TritonNPUBench/vllm/cp_mha_gather_cache_kernel/cp_mha_gather_cache_kernel.py b/benchmarks/TritonNPUBench/vllm/cp_mha_gather_cache_kernel/cp_mha_gather_cache_kernel.py new file mode 100644 index 00000000..1a2e6e1c --- /dev/null +++ b/benchmarks/TritonNPUBench/vllm/cp_mha_gather_cache_kernel/cp_mha_gather_cache_kernel.py @@ -0,0 +1,76 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/vllm_operator_cases/newtest_cases/test_cp_mha_gather_cache_kernel.py +# Main kernel: cp_mha_gather_cache_kernel +# PT file: cp_mha_gather_cache_kernel_v2.pt + +import triton +import triton.language as tl + + +# === cp_mha_gather_cache_kernel === +@triton.jit +def cp_mha_gather_cache_kernel( + key_cache_ptr, # [num_blocks, page_size, num_head, head_size] + value_cache_ptr, # [num_blocks, page_size, num_head, head_size] + key_ptr, # [num_tokens, num_heads, head_size] + value_ptr, # [num_tokens, num_heads, head_size] + block_table_ptr, # [num_batches, max_block_num] + cu_seqlens_kv_ptr, # [num_batches + 1] + token_to_batch_ptr, # [max_cum_tokens] + seq_start_ptr, # [num_batches] + k_scale_ptr, + v_scale_ptr, + num_heads, + head_size, + x, + max_block_num, + DEQUANT: tl.constexpr, + PAGE_SIZE: tl.constexpr, + CACHE_FORMAT: tl.constexpr, + BLOCK_SIZE: tl.constexpr, +): + token_id = tl.program_id(0) + col_offsets = tl.arange(0, BLOCK_SIZE) + if DEQUANT: + k_scale = tl.load(k_scale_ptr) + v_scale = tl.load(v_scale_ptr) + + key_ptr_offset = key_ptr + token_id * head_size * num_heads + value_ptr_offset = value_ptr + token_id * head_size * num_heads + batch_idx = tl.load(token_to_batch_ptr + token_id) + batch_start = tl.load(seq_start_ptr + batch_idx) + token_start = tl.load(cu_seqlens_kv_ptr + batch_idx) + batch_offset = token_id - token_start + batch_start + block_offset = batch_offset // PAGE_SIZE + block_id = tl.load( + block_table_ptr + max_block_num * batch_idx + block_offset + ).to(tl.int64) + slot_id = batch_offset % PAGE_SIZE + + if CACHE_FORMAT == "NHD": + # for kv cache layout as + # K: [num_blocks, page_size, num_head, head_dim] + # V: [num_blocks, page_size, num_head, head_dim] + key_cache_ptr_offset = ( + key_cache_ptr + + block_id * num_heads * head_size * PAGE_SIZE + + slot_id * num_heads * head_size + ) + value_cache_ptr_offset = ( + value_cache_ptr + + block_id * num_heads * head_size * PAGE_SIZE + + slot_id * num_heads * head_size + ) + + for i in tl.range(0, head_size * num_heads, BLOCK_SIZE): + mask = (col_offsets + i) < head_size * num_heads + k_reg = tl.load(key_cache_ptr_offset + col_offsets + i, mask=mask) + v_reg = tl.load(value_cache_ptr_offset + col_offsets + i, mask=mask) + if DEQUANT: + k_dtype = k_reg.dtype + v_dtype = v_reg.dtype + k_reg = (k_reg.to(tl.float32) * k_scale).to(k_dtype) + v_reg = (v_reg.to(tl.float32) * v_scale).to(v_dtype) + tl.store(key_ptr_offset + col_offsets + i, k_reg, mask=mask) + tl.store(value_ptr_offset + col_offsets + i, v_reg, mask=mask) + diff --git a/benchmarks/TritonNPUBench/vllm/eagle_prepare_inputs_padded_kernel/eagle_prepare_inputs_padded_kernel.py b/benchmarks/TritonNPUBench/vllm/eagle_prepare_inputs_padded_kernel/eagle_prepare_inputs_padded_kernel.py new file mode 100644 index 00000000..78e17359 --- /dev/null +++ b/benchmarks/TritonNPUBench/vllm/eagle_prepare_inputs_padded_kernel/eagle_prepare_inputs_padded_kernel.py @@ -0,0 +1,52 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/vllm_operator_cases/newtest_cases/test_eagle_prepare_inputs_padded_kernel.py +# Main kernel: eagle_prepare_inputs_padded_kernel +# PT file: eagle_prepare_inputs_padded_kernel_v2.pt + +import triton +import triton.language as tl + + +# === eagle_prepare_inputs_padded_kernel === +@triton.jit +def eagle_prepare_inputs_padded_kernel( + cu_num_draft_tokens_ptr, # [num_reqs] + valid_sampled_tokens_count_ptr, # [num_reqs] + query_start_loc_gpu_ptr, # [num_reqs + 1] + token_indices_to_sample_ptr, # [num_reqs] (output) + num_rejected_tokens_gpu_ptr, # [num_reqs] (output) + num_reqs, # tl.int32 +): + """ + Fused kernel for Eagle prepare_input_padded. This kernel computes the + token index to sample for each request, taking into account the number + of draft tokens and the number of valid sampled tokens (which is one more than + the number of accepted tokens). + """ + req_idx = tl.program_id(axis=0) + if req_idx >= num_reqs: + return + + # Calculate num_draft_tokens from cu_num_draft_tokens, which is an inclusive + # cumulative sum (first entry is the first value, not zero). + cu_draft_curr = tl.load(cu_num_draft_tokens_ptr + req_idx) + + num_draft_tokens = 0 + if req_idx == 0: + num_draft_tokens = cu_draft_curr + else: + cu_draft_prev = tl.load(cu_num_draft_tokens_ptr + req_idx - 1) + num_draft_tokens = cu_draft_curr - cu_draft_prev + + valid_count = tl.load(valid_sampled_tokens_count_ptr + req_idx) + num_rejected_tokens = num_draft_tokens + 1 - valid_count + num_rejected_tokens = tl.where(num_draft_tokens > 0, num_rejected_tokens, 0) + + # query_start_loc[req_idx + 1] is the start position of the next request, + # which is one past the last token of this request. + q_last_tok_idx = tl.load(query_start_loc_gpu_ptr + req_idx + 1) - 1 + + index_to_sample = q_last_tok_idx - num_rejected_tokens + tl.store(token_indices_to_sample_ptr + req_idx, index_to_sample) + tl.store(num_rejected_tokens_gpu_ptr + req_idx, num_rejected_tokens) + diff --git a/benchmarks/TritonNPUBench/vllm/eagle_prepare_next_token_padded_kernel/eagle_prepare_next_token_padded_kernel.py b/benchmarks/TritonNPUBench/vllm/eagle_prepare_next_token_padded_kernel/eagle_prepare_next_token_padded_kernel.py new file mode 100644 index 00000000..e650c816 --- /dev/null +++ b/benchmarks/TritonNPUBench/vllm/eagle_prepare_next_token_padded_kernel/eagle_prepare_next_token_padded_kernel.py @@ -0,0 +1,77 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/vllm_operator_cases/newtest_cases/test_eagle_prepare_next_token_padded_kernel.py +# Main kernel: eagle_prepare_next_token_padded_kernel +# PT file: eagle_prepare_next_token_padded_test_v2.pt + +import triton +import triton.language as tl + + +# === eagle_prepare_next_token_padded_kernel === +@triton.jit +def eagle_prepare_next_token_padded_kernel( + sampled_token_ids_ptr, # [num_reqs, num_sampled_tokens_per_req] + discard_request_mask_ptr, # [num_reqs] + backup_next_token_ids_ptr, # [num_reqs] + next_token_ids_ptr, # [num_reqs] (output) + valid_sampled_tokens_count_ptr, # [num_reqs] (output) + vocab_size, # tl.int32 + num_sampled_tokens_per_req, # tl.int32 (num_spec_tokens + 1) + num_reqs, # tl.int32 + stride_sampled_token_ids, # tl.int32 (stride for dim 0) + BLOCK_SIZE_TOKENS: tl.constexpr, # Power-of-2 >= num_sampled_tokens_per_req +): + """ + Fused kernel for Eagle prepare_next_token_ids_padded. This kernel computes the + number of valid (1 + accepted) tokens for each request, and the corresponding + "next" token id to sample from during speculative decoding. This is the + "last accepted token" from the sampled tokens, or the backup token if no + tokens were accepted or if the request is marked as discarded. + """ + req_idx = tl.program_id(axis=0) + if req_idx >= num_reqs: + return + + # Check if this request is discarded. + is_discarded = tl.load(discard_request_mask_ptr + req_idx) + + if is_discarded: + backup_token = tl.load(backup_next_token_ids_ptr + req_idx) + valid_count = tl.full((), 0, dtype=tl.uint32) + tl.store(next_token_ids_ptr + req_idx, backup_token) + tl.store(valid_sampled_tokens_count_ptr + req_idx, valid_count) + else: + # Count the number of valid tokens among the sampled tokens. + token_offs = tl.arange(0, BLOCK_SIZE_TOKENS) + token_mask = token_offs < num_sampled_tokens_per_req + + row_ptr = sampled_token_ids_ptr + req_idx * stride_sampled_token_ids + token_ids = tl.load(row_ptr + token_offs, mask=token_mask, other=-1) + + # Rejected tokens are -1, valid tokens are in [0, vocab_size) + is_valid_mask = (token_ids != -1) & (token_ids < vocab_size) & token_mask + ##################################### + # origin + # valid_count = tl.sum(is_valid_mask) + # modified + valid_count = tl.sum(is_valid_mask.to(tl.uint32)) + ##################################### + + if valid_count > 0: + # Guaranteed to be well-defined since + # valid_count > 0 implies is_valid_mask is not empty + last_valid_index = tl.max(tl.where(is_valid_mask, token_offs, -1)) + + # Select the token at that index, using a sum trick since + # we don't want to load again to access token_ids[last_valid_index]. + last_valid_token = tl.sum( + tl.where(token_offs == last_valid_index, token_ids, 0) + ) + tl.store(next_token_ids_ptr + req_idx, last_valid_token) + else: + # No valid tokens found, use backup token + backup_token = tl.load(backup_next_token_ids_ptr + req_idx) + tl.store(next_token_ids_ptr + req_idx, backup_token) + + tl.store(valid_sampled_tokens_count_ptr + req_idx, valid_count) + diff --git a/benchmarks/TritonNPUBench/vllm/expand_kernel/expand_kernel.py b/benchmarks/TritonNPUBench/vllm/expand_kernel/expand_kernel.py new file mode 100644 index 00000000..ba85fa5b --- /dev/null +++ b/benchmarks/TritonNPUBench/vllm/expand_kernel/expand_kernel.py @@ -0,0 +1,32 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/vllm_operator_cases/newtest_cases/test_expand_kernel.py +# Main kernel: expand_kernel +# PT file: test_expand_kernel_v2.pt + +import triton +import triton.language as tl + + +# === expand_kernel === +@triton.jit(do_not_specialize=["replace_from", "replace_to"]) +def expand_kernel( + output_ptr, # [num_tokens] + input_ptr, # [batch_size] + cu_num_tokens_ptr, # [batch_size] + replace_from, + replace_to, + MAX_NUM_TOKENS: tl.constexpr, +): + req_idx = tl.program_id(0) + if req_idx == 0: # noqa: SIM108 + start_idx = 0 + else: + start_idx = tl.load(cu_num_tokens_ptr + req_idx - 1) + end_idx = tl.load(cu_num_tokens_ptr + req_idx) + num_tokens = end_idx - start_idx + + src_val = tl.load(input_ptr + req_idx) + src_val = tl.where(src_val == replace_from, replace_to, src_val) + offset = tl.arange(0, MAX_NUM_TOKENS) + tl.store(output_ptr + start_idx + offset, src_val, mask=offset < num_tokens) + diff --git a/benchmarks/TritonNPUBench/vllm/fused_gdn_gating_kernel/fused_gdn_gating_kernel.py b/benchmarks/TritonNPUBench/vllm/fused_gdn_gating_kernel/fused_gdn_gating_kernel.py new file mode 100644 index 00000000..821bcdb3 --- /dev/null +++ b/benchmarks/TritonNPUBench/vllm/fused_gdn_gating_kernel/fused_gdn_gating_kernel.py @@ -0,0 +1,45 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/vllm_operator_cases/newtest_cases/test_fused_gdn_gating_kernel.py +# Main kernel: fused_gdn_gating_kernel +# PT file: fused_gdn_gating_kernel_v2.pt + +import triton +import triton.language as tl + + +# === fused_gdn_gating_kernel === +@triton.jit +def fused_gdn_gating_kernel( + g, + beta_output, + A_log, + a, + b, + dt_bias, + seq_len, + NUM_HEADS: tl.constexpr, + beta: tl.constexpr, + threshold: tl.constexpr, + BLK_HEADS: tl.constexpr, +): + i_b, i_s, i_d = tl.program_id(0), tl.program_id(1), tl.program_id(2) + head_off = i_d * BLK_HEADS + tl.arange(0, BLK_HEADS) + off = i_b * seq_len * NUM_HEADS + i_s * NUM_HEADS + head_off + mask = head_off < NUM_HEADS + blk_A_log = tl.load(A_log + head_off, mask=mask) + blk_a = tl.load(a + off, mask=mask) + blk_b = tl.load(b + off, mask=mask) + blk_bias = tl.load(dt_bias + head_off, mask=mask) + # If the model is loaded in fp16, without the .float() here, A might be -inf + x = blk_a.to(tl.float32) + blk_bias.to(tl.float32) + softplus_x = tl.where( + beta * x <= threshold, (1 / beta) * tl.log(1 + tl.exp(beta * x)), x + ) + blk_g = -tl.exp(blk_A_log.to(tl.float32)) * softplus_x + tl.store(g + off, blk_g.to(g.dtype.element_ty), mask=mask) + # compute beta_output = sigmoid(b) + blk_beta_output = tl.sigmoid(blk_b.to(tl.float32)) + tl.store( + beta_output + off, blk_beta_output.to(beta_output.dtype.element_ty), mask=mask + ) + diff --git a/benchmarks/TritonNPUBench/vllm/fused_moe_kernel/fused_moe_kernel.py b/benchmarks/TritonNPUBench/vllm/fused_moe_kernel/fused_moe_kernel.py new file mode 100644 index 00000000..30325414 --- /dev/null +++ b/benchmarks/TritonNPUBench/vllm/fused_moe_kernel/fused_moe_kernel.py @@ -0,0 +1,291 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/vllm_operator_cases/newtest_cases/test_fused_moe_kernel.py +# Main kernel: fused_moe_kernel +# PT file: fused_moe_kernel_v2.pt + +import triton +import triton.language as tl + + +# === fused_moe_kernel === +@triton.jit +def fused_moe_kernel( + # Pointers to matrices + a_ptr, + b_ptr, + c_ptr, + b_bias_ptr, + a_scale_ptr, + b_scale_ptr, + topk_weights_ptr, + sorted_token_ids_ptr, + expert_ids_ptr, + num_tokens_post_padded_ptr, + # Matrix dimensions + N, + K, + EM, + num_valid_tokens, + # The stride variables represent how much to increase the ptr by when + # moving by 1 element in a particular dimension. E.g. `stride_am` is + # how much to increase `a_ptr` by to get the element one row down + # (A has M rows). + stride_am, + stride_ak, + stride_be, + stride_bk, + stride_bn, + stride_cm, + stride_cn, + stride_asm, + stride_ask, + stride_bse, + stride_bsk, + stride_bsn, + stride_bbe, # bias expert stride + stride_bbn, # bias N stride + # Block size for block-wise quantization + group_n: tl.constexpr, + group_k: tl.constexpr, + naive_block_assignment: tl.constexpr, + # Meta-parameters + BLOCK_SIZE_M: tl.constexpr, + BLOCK_SIZE_N: tl.constexpr, + BLOCK_SIZE_K: tl.constexpr, + GROUP_SIZE_M: tl.constexpr, + SPLIT_K: tl.constexpr, + MUL_ROUTED_WEIGHT: tl.constexpr, + top_k: tl.constexpr, + compute_type: tl.constexpr, + use_fp8_w8a8: tl.constexpr, + use_int8_w8a8: tl.constexpr, + use_int8_w8a16: tl.constexpr, + per_channel_quant: tl.constexpr, + HAS_BIAS: tl.constexpr, +): + """ + Implements the fused computation for a Mixture of Experts (MOE) using + token and expert matrices. + + Key Parameters: + - A: The input tensor representing tokens with shape (*, K), where '*' can + be any shape representing batches and K is the feature dimension of + each token. + - B: The stacked MOE weight tensor with shape (E, N, K), where E is + the number of experts, K is the input feature dimension, and N is + the output feature dimension. + - C: The output cache tensor with shape (M, topk, N), where M is the + total number of tokens post padding, topk is the number of times + each token is repeated, and N is the output feature dimension. + - sorted_token_ids: A tensor containing the sorted indices of tokens, + repeated topk times and arranged by the expert index they are + assigned to. + - expert_ids: A tensor containing the indices of the expert for each + block. It determines which expert matrix from B should be used for + each block in A. + - naive_block_assignment: A boolean flag indicating whether to use naive + token wise block assignment. If True, each block corresponds to a + single token. + This kernel performs the multiplication of a token by its corresponding + expert matrix as determined by `expert_ids`. The sorting of + `sorted_token_ids` by expert index and padding ensures divisibility by + BLOCK_SIZE_M, which is necessary to maintain consistency in block matrix + multiplication across different blocks processed by the same expert. + """ + # ----------------------------------------------------------- + # Map program ids `pid` to the block of C it should compute. + # This is done in a grouped ordering to promote L2 data reuse. + pid = tl.program_id(axis=0) + num_pid_m = tl.cdiv(EM, BLOCK_SIZE_M) + num_pid_n = tl.cdiv(N, BLOCK_SIZE_N) + num_pid_in_group = GROUP_SIZE_M * num_pid_n + group_id = pid // num_pid_in_group + first_pid_m = group_id * GROUP_SIZE_M + group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M) + pid_m = first_pid_m + ((pid % num_pid_in_group) % group_size_m) + pid_n = (pid % num_pid_in_group) // group_size_m + + # ---------------------------------------------------------- + # Create pointers for the first blocks of A and B. + # We will advance this pointer as we move in the K direction + # and accumulate + # `a_ptrs` is a block of [BLOCK_SIZE_M, BLOCK_SIZE_K] pointers + # `b_ptrs` is a block of [BLOCK_SIZE_K, BLOCK_SIZE_N] pointers + offs = tl.arange(0, BLOCK_SIZE_M).to(tl.int64) + num_tokens_post_padded = tl.load(num_tokens_post_padded_ptr) + if pid_m * BLOCK_SIZE_M >= num_tokens_post_padded: + return + if not naive_block_assignment: + offs_token_id = pid_m * BLOCK_SIZE_M + offs + offs_token = tl.load(sorted_token_ids_ptr + offs_token_id) + else: + offs_token = tl.where( + offs == 0, + pid_m, # first element = pid_m + num_valid_tokens, # remaining elements = constant + ) + + token_mask = offs_token < num_valid_tokens + + off_experts = tl.load(expert_ids_ptr + pid_m).to(tl.int64) + if off_experts == -1: + # ----------------------------------------------------------- + # Write back zeros to the output when the expert is not + # in the current expert parallel rank. + write_zeros_to_output( + c_ptr, + stride_cm, + stride_cn, + pid_n, + N, + offs_token, + token_mask, + BLOCK_SIZE_M, + BLOCK_SIZE_N, + compute_type, + ) + return + + offs_bn = (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N).to(tl.int64)) % N + offs_k = tl.arange(0, BLOCK_SIZE_K) + a_ptrs = a_ptr + ( + offs_token[:, None] // top_k * stride_am + offs_k[None, :] * stride_ak + ) + + b_ptrs = ( + b_ptr + + off_experts * stride_be + + (offs_k[:, None] * stride_bk + offs_bn[None, :] * stride_bn) + ) + if use_int8_w8a16: + b_scale_ptrs = ( + b_scale_ptr + off_experts * stride_bse + offs_bn[None, :] * stride_bsn + ) + b_scale = tl.load(b_scale_ptrs) + + if use_fp8_w8a8 or use_int8_w8a8: + # block-wise + if group_k > 0 and group_n > 0: + a_scale_ptrs = a_scale_ptr + (offs_token // top_k) * stride_asm + offs_bsn = offs_bn // group_n + b_scale_ptrs = ( + b_scale_ptr + off_experts * stride_bse + offs_bsn * stride_bsn + ) + # channel-wise + elif per_channel_quant: + b_scale_ptrs = ( + b_scale_ptr + off_experts * stride_bse + offs_bn[None, :] * stride_bsn + ) + b_scale = tl.load(b_scale_ptrs) + # Load per-token scale for activations + a_scale_ptrs = a_scale_ptr + (offs_token // top_k) * stride_asm + a_scale = tl.load(a_scale_ptrs, mask=token_mask, other=0.0)[:, None] + # tensor-wise + else: + a_scale = tl.load(a_scale_ptr) + b_scale = tl.load(b_scale_ptr + off_experts) + if HAS_BIAS: + # bias shape: [num_experts, N] + bias_ptrs = b_bias_ptr + off_experts * stride_bbe + offs_bn * stride_bbn + bias = tl.load(bias_ptrs, mask=(offs_bn < N), other=0.0) + # ----------------------------------------------------------- + # Iterate to compute a block of the C matrix. + # We accumulate into a `[BLOCK_SIZE_M, BLOCK_SIZE_N]` block + # of fp32 values for higher accuracy. + # `accumulator` will be converted back to fp16 after the loop. + accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32) + for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)): + # Load the next block of A and B, generate a mask by checking the + # K dimension. + a = tl.load( + a_ptrs, + mask=token_mask[:, None] & (offs_k[None, :] < K - k * BLOCK_SIZE_K), + other=0.0, + ) + b = tl.load(b_ptrs, mask=offs_k[:, None] < K - k * BLOCK_SIZE_K, other=0.0) + # We accumulate along the K dimension. + if use_int8_w8a16: + accumulator = tl.dot(a, b.to(compute_type), acc=accumulator) + elif use_fp8_w8a8 or use_int8_w8a8: + if group_k > 0 and group_n > 0: + k_start = k * BLOCK_SIZE_K + offs_ks = k_start // group_k + a_scale = tl.load( + a_scale_ptrs + offs_ks * stride_ask, mask=token_mask, other=0.0 + ) + b_scale = tl.load(b_scale_ptrs + offs_ks * stride_bsk) + + accumulator += tl.dot(a, b) * a_scale[:, None] * b_scale[None, :] + else: + if use_fp8_w8a8: + # acc used to enable fp8_fast_accum + accumulator = tl.dot(a, b, acc=accumulator) + else: + accumulator += tl.dot(a, b) + else: + accumulator += tl.dot(a, b) + # Advance the ptrs to the next K block. + a_ptrs += BLOCK_SIZE_K * stride_ak + b_ptrs += BLOCK_SIZE_K * stride_bk + + # Dequantization for supported quantization schemes: + # - int8_w8a16 + # - fp8_w8a8 + # - int8_w8a8 + # Accumulator and scalings are in float32 to preserve numerical accuracy. + if use_int8_w8a16: + accumulator = accumulator * b_scale + elif (use_fp8_w8a8 or use_int8_w8a8) and not (group_k > 0 and group_n > 0): + accumulator = accumulator * a_scale * b_scale + + # Bias addition: + # Bias must be applied after dequantization: + # - Since bias is typically not quantized + # - Bias should not be scaled by quantization factors + if HAS_BIAS: + accumulator += bias[None, :] + + # Router (MoE) weight multiplication: + # This multiplication MUST be performed in float32 before any precision + # conversion to ensure numerical stability, which is especially critical + # on ROCm platforms. + if MUL_ROUTED_WEIGHT: + moe_weight = tl.load( + topk_weights_ptr + offs_token, + mask=token_mask, + other=0, + ) + accumulator *= moe_weight[:, None] + + # Final precision conversion: + # Cast once at the end to the desired compute/output dtype. + accumulator = accumulator.to(compute_type) + + # ----------------------------------------------------------- + # Write back the block of the output + offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) + c_ptrs = c_ptr + stride_cm * offs_token[:, None] + stride_cn * offs_cn[None, :] + c_mask = token_mask[:, None] & (offs_cn[None, :] < N) + tl.store(c_ptrs, accumulator, mask=c_mask) + + +# === write_zeros_to_output === +@triton.jit +def write_zeros_to_output( + c_ptr, + stride_cm, + stride_cn, + pid_n, + N, + offs_token, + token_mask, + BLOCK_SIZE_M, + BLOCK_SIZE_N, + compute_type, +): + accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=compute_type) + offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) + c_ptrs = c_ptr + stride_cm * offs_token[:, None] + stride_cn * offs_cn[None, :] + c_mask = token_mask[:, None] & (offs_cn[None, :] < N) + tl.store(c_ptrs, accumulator, mask=c_mask) + diff --git a/benchmarks/TritonNPUBench/vllm/fused_moe_kernel_gptq_awq/fused_moe_kernel_gptq_awq.py b/benchmarks/TritonNPUBench/vllm/fused_moe_kernel_gptq_awq/fused_moe_kernel_gptq_awq.py new file mode 100644 index 00000000..f44c461a --- /dev/null +++ b/benchmarks/TritonNPUBench/vllm/fused_moe_kernel_gptq_awq/fused_moe_kernel_gptq_awq.py @@ -0,0 +1,263 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/vllm_operator_cases/newtest_cases/test_fused_moe_kernel_gptq_awq.py +# Main kernel: fused_moe_kernel_gptq_awq +# PT file: fused_moe_kernel_gptq_awq_v2.pt + +import triton +import triton.language as tl + + +# === fused_moe_kernel_gptq_awq === +@triton.jit +def fused_moe_kernel_gptq_awq( + # Pointers to matrices + a_ptr, + b_ptr, + c_ptr, + b_scale_ptr, + b_zp_ptr, + topk_weights_ptr, + sorted_token_ids_ptr, + expert_ids_ptr, + num_tokens_post_padded_ptr, + # Matrix dimensions + N: tl.constexpr, + K: tl.constexpr, + EM, + num_valid_tokens, + # The stride variables represent how much to increase the ptr by when + # moving by 1 element in a particular dimension. E.g. `stride_am` is + # how much to increase `a_ptr` by to get the element one row down + # (A has M rows). + stride_am, + stride_ak, + stride_be, + stride_bk, + stride_bn, + stride_cm, + stride_cn, + stride_bse, + stride_bsk, + stride_bsn, + stride_bze, + stride_bzk, + stride_bzn, + block_k_diviable: tl.constexpr, + group_size: tl.constexpr, + # Meta-parameters + BLOCK_SIZE_M: tl.constexpr, + BLOCK_SIZE_N: tl.constexpr, + BLOCK_SIZE_K: tl.constexpr, + GROUP_SIZE_M: tl.constexpr, + SPLIT_K: tl.constexpr, + MUL_ROUTED_WEIGHT: tl.constexpr, + top_k: tl.constexpr, + compute_type: tl.constexpr, + has_zp: tl.constexpr, + use_int4_w4a16: tl.constexpr, + use_int8_w8a16: tl.constexpr, +): + """ + Implements the fused computation for a Mixture of Experts (MOE) using + token and expert matrices. + + Key Parameters: + - A: The input tensor representing tokens with shape (*, K), where '*' can + be any shape representing batches and K is the feature dimension of + each token. + - B: The stacked MOE weight tensor with shape (E, N, K), where E is + the number of experts, K is the input feature dimension, and N is + the output feature dimension. + - C: The output cache tensor with shape (M, topk, N), where M is the + total number of tokens post padding, topk is the number of times + each token is repeated, and N is the output feature dimension. + - sorted_token_ids: A tensor containing the sorted indices of tokens, + repeated topk times and arranged by the expert index they are + assigned to. + - expert_ids: A tensor containing the indices of the expert for each + block. It determines which expert matrix from B should be used for + each block in A. + This kernel performs the multiplication of a token by its corresponding + expert matrix as determined by `expert_ids`. The sorting of + `sorted_token_ids` by expert index and padding ensures divisibility by + BLOCK_SIZE_M, which is necessary to maintain consistency in block matrix + multiplication across different blocks processed by the same expert. + """ + # ----------------------------------------------------------- + # Map program ids `pid` to the block of C it should compute. + # This is done in a grouped ordering to promote L2 data reuse. + pid = tl.program_id(axis=0) + num_pid_m = tl.cdiv(EM, BLOCK_SIZE_M) + num_pid_n = tl.cdiv(N, BLOCK_SIZE_N) + num_pid_in_group = GROUP_SIZE_M * num_pid_n + group_id = pid // num_pid_in_group + first_pid_m = group_id * GROUP_SIZE_M + group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M) + pid_m = first_pid_m + ((pid % num_pid_in_group) % group_size_m) + pid_n = (pid % num_pid_in_group) // group_size_m + + # ---------------------------------------------------------- + # Create pointers for the first blocks of A and B. + # We will advance this pointer as we move in the K direction + # and accumulate + # `a_ptrs` is a block of [BLOCK_SIZE_M, BLOCK_SIZE_K] pointers + # `b_ptrs` is a block of [BLOCK_SIZE_K, BLOCK_SIZE_N] pointers + num_tokens_post_padded = tl.load(num_tokens_post_padded_ptr) + if pid_m * BLOCK_SIZE_M >= num_tokens_post_padded: + return + offs_token_id = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M).to(tl.int64) + offs_token = tl.load(sorted_token_ids_ptr + offs_token_id) + token_mask = offs_token < num_valid_tokens + + off_experts = tl.load(expert_ids_ptr + pid_m).to(tl.int64) + if off_experts == -1: + # ----------------------------------------------------------- + # Write back zeros to the output when the expert is not + # in the current expert parallel rank. + write_zeros_to_output( + c_ptr, + stride_cm, + stride_cn, + pid_n, + N, + offs_token, + token_mask, + BLOCK_SIZE_M, + BLOCK_SIZE_N, + compute_type, + ) + return + + offs_bn = (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N).to(tl.int64)) % N + offs_k = tl.arange(0, BLOCK_SIZE_K) + a_ptrs = a_ptr + ( + offs_token[:, None] // top_k * stride_am + offs_k[None, :] * stride_ak + ) + + if use_int4_w4a16: + b_ptrs = ( + b_ptr + + off_experts * stride_be + + (offs_k[:, None] // 2) * stride_bk + + offs_bn[None, :] * stride_bn + ) + b_shifter = (offs_k[:, None] % 2) * 4 + elif use_int8_w8a16: + b_ptrs = ( + b_ptr + + off_experts * stride_be + + offs_k[:, None] * stride_bk + + offs_bn[None, :] * stride_bn + ) + + if not has_zp and use_int4_w4a16: + b_zp_num = 8 + if not has_zp and use_int8_w8a16: + b_zp_num = 128 + elif has_zp and use_int4_w4a16: + b_zp_shifter = (offs_bn[None, :] % 2) * 4 + + # ----------------------------------------------------------- + # Iterate to compute a block of the C matrix. + # We accumulate into a `[BLOCK_SIZE_M, BLOCK_SIZE_N]` block + # of fp32 values for higher accuracy. + # `accumulator` will be converted back to fp16 after the loop. + accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32) + for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)): + # Load the next block of A and B, generate a mask by checking the + # K dimension. + + if not block_k_diviable: + k_mask = offs_k[:, None] < K - k * BLOCK_SIZE_K + k_other = 0.0 + else: + k_mask = None + k_other = None + + a = tl.load( + a_ptrs, + mask=token_mask[:, None] & (offs_k[None, :] < K - k * BLOCK_SIZE_K), + other=0.0, + ) + b = tl.load(b_ptrs) + if use_int4_w4a16: + b = (b >> b_shifter) & 0xF + + b_scale_ptrs = ( + b_scale_ptr + + off_experts * stride_bse + + offs_bn[None, :] * stride_bsn + + ((offs_k[:, None] + BLOCK_SIZE_K * k) // group_size) * stride_bsk + ) + b_scale = tl.load(b_scale_ptrs, mask=k_mask, other=k_other) + b_scale = b_scale.to(tl.float32) + + if has_zp and use_int4_w4a16: + offs_k_true = (offs_k[:, None] + BLOCK_SIZE_K * k) // group_size + b_zp_ptrs = ( + b_zp_ptr + + off_experts * stride_bze + + (offs_bn[None, :] // 2) * stride_bzn + + offs_k_true * stride_bzk + ) + b_zp = tl.load(b_zp_ptrs, mask=k_mask, other=k_other) + b_zp = (b_zp >> b_zp_shifter) & 0xF + b_zp = b_zp.to(tl.float32) + elif has_zp and use_int8_w8a16: + offs_k_true = (offs_k[:, None] + BLOCK_SIZE_K * k) // group_size + b_zp_ptrs = ( + b_zp_ptr + + off_experts * stride_bze + + offs_bn[None, :] * stride_bzn + + offs_k_true * stride_bzk + ) + b_zp = tl.load(b_zp_ptrs, mask=k_mask, other=k_other) + b_zp = b_zp.to(tl.float32) + + # We accumulate along the K dimension. + if has_zp: + b = ((b.to(tl.float32) - b_zp) * b_scale).to(compute_type) + else: + b = ((b.to(tl.float32) - b_zp_num) * b_scale).to(compute_type) + accumulator = tl.dot(a, b, acc=accumulator) + + # Advance the ptrs to the next K block. + a_ptrs += BLOCK_SIZE_K * stride_ak + if use_int4_w4a16: + b_ptrs += (BLOCK_SIZE_K // 2) * stride_bk + else: + b_ptrs += BLOCK_SIZE_K * stride_bk + + if MUL_ROUTED_WEIGHT: + moe_weight = tl.load(topk_weights_ptr + offs_token, mask=token_mask, other=0) + accumulator = accumulator * moe_weight[:, None] + + accumulator = accumulator.to(compute_type) + # ----------------------------------------------------------- + # Write back the block of the output + offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) + c_ptrs = c_ptr + stride_cm * offs_token[:, None] + stride_cn * offs_cn[None, :] + c_mask = token_mask[:, None] & (offs_cn[None, :] < N) + tl.store(c_ptrs, accumulator, mask=c_mask) + + +# === write_zeros_to_output === +@triton.jit +def write_zeros_to_output( + c_ptr, + stride_cm, + stride_cn, + pid_n, + N, + offs_token, + token_mask, + BLOCK_SIZE_M, + BLOCK_SIZE_N, + compute_type, +): + accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=compute_type) + offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) + c_ptrs = c_ptr + stride_cm * offs_token[:, None] + stride_cn * offs_cn[None, :] + c_mask = token_mask[:, None] & (offs_cn[None, :] < N) + tl.store(c_ptrs, accumulator, mask=c_mask) + diff --git a/benchmarks/TritonNPUBench/vllm/fused_recurrent_gated_delta_rule_fwd_kernel/fused_recurrent_gated_delta_rule_fwd_kernel.py b/benchmarks/TritonNPUBench/vllm/fused_recurrent_gated_delta_rule_fwd_kernel/fused_recurrent_gated_delta_rule_fwd_kernel.py new file mode 100644 index 00000000..40787ba9 --- /dev/null +++ b/benchmarks/TritonNPUBench/vllm/fused_recurrent_gated_delta_rule_fwd_kernel/fused_recurrent_gated_delta_rule_fwd_kernel.py @@ -0,0 +1,158 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/vllm_operator_cases/newtest_cases/test_fused_recurrent_gated_delta_rule_fwd_kernel.py +# Main kernel: fused_recurrent_gated_delta_rule_fwd_kernel +# PT file: fused_recurrent_gated_delta_rule_fwd_kernel_v2.pt + +import triton +import triton.language as tl + + +# === fused_recurrent_gated_delta_rule_fwd_kernel === +@triton.jit(do_not_specialize=["N", "T"]) +def fused_recurrent_gated_delta_rule_fwd_kernel( + q, + k, + v, + g, + beta, + o, + h0, + ht, + cu_seqlens, + ssm_state_indices, + num_accepted_tokens, + scale, + N: tl.int64, # num of sequences + T: tl.int64, # num of tokens + B: tl.constexpr, + H: tl.constexpr, + HV: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + stride_init_state_token: tl.constexpr, + stride_final_state_token: tl.constexpr, + stride_indices_seq: tl.constexpr, + stride_indices_tok: tl.constexpr, + USE_INITIAL_STATE: tl.constexpr, # whether to use initial state + INPLACE_FINAL_STATE: tl.constexpr, # whether to store final state inplace + IS_BETA_HEADWISE: tl.constexpr, # whether beta is headwise vector or scalar, + USE_QK_L2NORM_IN_KERNEL: tl.constexpr, + IS_VARLEN: tl.constexpr, + IS_CONTINUOUS_BATCHING: tl.constexpr, + IS_SPEC_DECODING: tl.constexpr, + IS_KDA: tl.constexpr, +): + i_k, i_v, i_nh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_n, i_hv = i_nh // HV, i_nh % HV + i_h = i_hv // (HV // H) + if IS_VARLEN: + bos, eos = ( + tl.load(cu_seqlens + i_n).to(tl.int64), + tl.load(cu_seqlens + i_n + 1).to(tl.int64), + ) + all = T + T = eos - bos + else: + bos, eos = i_n * T, i_n * T + T + all = B * T + + if T == 0: + # no tokens to process for this sequence + return + + o_k = i_k * BK + tl.arange(0, BK) + o_v = i_v * BV + tl.arange(0, BV) + + p_q = q + (bos * H + i_h) * K + o_k + p_k = k + (bos * H + i_h) * K + o_k + p_v = v + (bos * HV + i_hv) * V + o_v + if IS_BETA_HEADWISE: + p_beta = beta + (bos * HV + i_hv) * V + o_v + else: + p_beta = beta + bos * HV + i_hv + + if not IS_KDA: + p_g = g + bos * HV + i_hv + else: + p_gk = g + (bos * HV + i_hv) * K + o_k + + p_o = o + ((i_k * all + bos) * HV + i_hv) * V + o_v + + mask_k = o_k < K + mask_v = o_v < V + mask_h = mask_k[:, None] & mask_v[None, :] + + b_h = tl.zeros([BK, BV], dtype=tl.float32) + if USE_INITIAL_STATE: + if IS_CONTINUOUS_BATCHING: + if IS_SPEC_DECODING: + i_t = tl.load(num_accepted_tokens + i_n).to(tl.int64) - 1 + else: + i_t = 0 + p_h0 = ( + h0 + + tl.load(ssm_state_indices + i_n * stride_indices_seq + i_t).to( + tl.int64 + ) + * stride_init_state_token + ) + else: + p_h0 = h0 + bos * HV * K * V + p_h0 = p_h0 + i_hv * K * V + o_k[:, None] * V + o_v[None, :] + b_h += tl.load(p_h0, mask=mask_h, other=0).to(tl.float32) + + for i_t in range(0, T): + b_q = tl.load(p_q, mask=mask_k, other=0).to(tl.float32) + b_k = tl.load(p_k, mask=mask_k, other=0).to(tl.float32) + b_v = tl.load(p_v, mask=mask_v, other=0).to(tl.float32) + + if USE_QK_L2NORM_IN_KERNEL: + b_q = b_q / tl.sqrt(tl.sum(b_q * b_q) + 1e-6) + b_k = b_k / tl.sqrt(tl.sum(b_k * b_k) + 1e-6) + b_q = b_q * scale + # [BK, BV] + if not IS_KDA: + b_g = tl.load(p_g).to(tl.float32) + b_h *= tl.exp(b_g) + else: + b_gk = tl.load(p_gk).to(tl.float32) + b_h *= tl.exp(b_gk[:, None]) + # [BV] + b_v -= tl.sum(b_h * b_k[:, None], 0) + if IS_BETA_HEADWISE: + b_beta = tl.load(p_beta, mask=mask_v, other=0).to(tl.float32) + else: + b_beta = tl.load(p_beta).to(tl.float32) + b_v *= b_beta + # [BK, BV] + b_h += b_k[:, None] * b_v[None, :] + # [BV] + b_o = tl.sum(b_h * b_q[:, None], 0) + tl.store(p_o, b_o.to(p_o.dtype.element_ty), mask=mask_v) + + # keep the states for multi-query tokens + if INPLACE_FINAL_STATE: + p_ht = ( + ht + + tl.load(ssm_state_indices + i_n * stride_indices_seq + i_t).to( + tl.int64 + ) + * stride_final_state_token + ) + else: + p_ht = ht + (bos + i_t) * stride_final_state_token + p_ht = p_ht + i_hv * K * V + o_k[:, None] * V + o_v[None, :] + tl.store(p_ht, b_h.to(p_ht.dtype.element_ty), mask=mask_h) + + p_q += H * K + p_k += H * K + p_o += HV * V + p_v += HV * V + if not IS_KDA: + p_g += HV + else: + p_gk += HV * K + p_beta += HV * (V if IS_BETA_HEADWISE else 1) + diff --git a/benchmarks/TritonNPUBench/vllm/kda_gate_fwd_kernel/kda_gate_fwd_kernel.py b/benchmarks/TritonNPUBench/vllm/kda_gate_fwd_kernel/kda_gate_fwd_kernel.py new file mode 100644 index 00000000..ad52531d --- /dev/null +++ b/benchmarks/TritonNPUBench/vllm/kda_gate_fwd_kernel/kda_gate_fwd_kernel.py @@ -0,0 +1,71 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/vllm_operator_cases/newtest_cases/test_kda_gate_fwd_kernel.py +# Main kernel: kda_gate_fwd_kernel +# PT file: test_kda_gate_fwd_kernel_v2.pt + +import triton +import triton.language as tl + +# === kda_gate_fwd_kernel === +@triton.jit +def kda_gate_fwd_kernel( + g, + A, + y, + g_bias, + beta: tl.constexpr, + threshold: tl.constexpr, + T, + H, + D: tl.constexpr, + BT: tl.constexpr, + BD: tl.constexpr, + HAS_BIAS: tl.constexpr, +): + i_t, i_h = tl.program_id(0), tl.program_id(1) + n_t = i_t * BT + + b_a = tl.load(A + i_h).to(tl.float32) + b_a = -tl.exp(b_a) + + stride_row = H * D + stride_col = 1 + + g_ptr = tl.make_block_ptr( + base=g + i_h * D, + shape=(T, D), + strides=(stride_row, stride_col), + offsets=(n_t, 0), + block_shape=(BT, BD), + order=(1, 0), + ) + + y_ptr = tl.make_block_ptr( + base=y + i_h * D, + shape=(T, D), + strides=(stride_row, stride_col), + offsets=(n_t, 0), + block_shape=(BT, BD), + order=(1, 0), + ) + + b_g = tl.load(g_ptr, boundary_check=(0, 1)).to(tl.float32) + + if HAS_BIAS: + n_d = tl.arange(0, BD) + bias_mask = n_d < D + b_bias = tl.load(g_bias + i_h * D + n_d, mask=bias_mask, other=0.0).to( + tl.float32 + ) + b_g = b_g + b_bias[None, :] + + # softplus(x, beta) = (1/beta) * log(1 + exp(beta * x)) + # When beta * x > threshold, use linear approximation x + # Use threshold to switch to linear when beta*x > threshold + g_scaled = b_g * beta + use_linear = g_scaled > threshold + sp = tl.where(use_linear, b_g, (1.0 / beta) * tl.log(1.0 + tl.exp(g_scaled))) + b_y = b_a * sp + + tl.store(y_ptr, b_y.to(y.dtype.element_ty), boundary_check=(0, 1)) + diff --git a/benchmarks/TritonNPUBench/vllm/kernel_paged_attention_2d/kernel_paged_attention_2d.py b/benchmarks/TritonNPUBench/vllm/kernel_paged_attention_2d/kernel_paged_attention_2d.py new file mode 100644 index 00000000..c0c746ab --- /dev/null +++ b/benchmarks/TritonNPUBench/vllm/kernel_paged_attention_2d/kernel_paged_attention_2d.py @@ -0,0 +1,198 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/vllm_operator_cases/newtest_cases/test_kernel_paged_attention_2d.py +# Main kernel: kernel_paged_attention_2d +# PT file: kernel_paged_attention_2d_v2.pt + +import triton +import triton.language as tl + + +# === kernel_paged_attention_2d === +@triton.jit +def kernel_paged_attention_2d( + output_ptr, # [num_tokens, num_query_heads, head_size] + query_ptr, # [num_tokens, num_query_heads, head_size] + key_cache_ptr, # [num_blks, num_kv_heads, head_size // x, blk_size, x] + value_cache_ptr, # [num_blks, num_kv_heads, head_size, blk_size] + block_tables_ptr, # [num_seqs, max_num_blocks_per_seq] + seq_lens_ptr, # [num_seqs] + alibi_slopes_ptr, # [num_query_heads] + scale, # float32 + k_scale, # float32 + v_scale, # float32 + num_query_heads: tl.constexpr, # int + num_queries_per_kv: tl.constexpr, # int + num_queries_per_kv_padded: tl.constexpr, # int + block_table_stride: tl.int64, # int + query_stride_0: tl.int64, # int + query_stride_1: tl.int64, # int, should be equal to head_size + output_stride_0: tl.int64, # int + output_stride_1: tl.int64, # int, should be equal to head_size + BLOCK_SIZE: tl.constexpr, # int + HEAD_SIZE: tl.constexpr, # int + HEAD_SIZE_PADDED: tl.constexpr, # int, must be power of 2 + USE_ALIBI_SLOPES: tl.constexpr, # bool + SLIDING_WINDOW: tl.constexpr, # int + x: tl.constexpr, # int + stride_k_cache_0: tl.int64, # int + stride_k_cache_1: tl.int64, # int + stride_k_cache_2: tl.int64, # int + stride_k_cache_3: tl.int64, # int + stride_k_cache_4: tl.int64, # int + stride_v_cache_0: tl.int64, # int + stride_v_cache_1: tl.int64, # int + stride_v_cache_2: tl.int64, # int + stride_v_cache_3: tl.int64, # int + filter_by_query_len: tl.constexpr, # bool + query_start_len_ptr, # [num_seqs+1] +): + seq_idx = tl.program_id(0) + kv_head_idx = tl.program_id(1) + + if filter_by_query_len: + cur_batch_in_all_start_index = tl.load(query_start_len_ptr + seq_idx) + cur_batch_in_all_stop_index = tl.load(query_start_len_ptr + seq_idx + + 1) + cur_batch_query_len = cur_batch_in_all_stop_index \ + - cur_batch_in_all_start_index + if cur_batch_query_len > 1: + return + else: + cur_batch_in_all_start_index = seq_idx + + query_head_idx = kv_head_idx * num_queries_per_kv + tl.arange( + 0, num_queries_per_kv_padded) + + query_offset = (cur_batch_in_all_start_index * query_stride_0 + + query_head_idx[:, None] * query_stride_1) + + head_mask = query_head_idx < (kv_head_idx + 1) * num_queries_per_kv + head_mask = head_mask & (query_head_idx < num_query_heads) + + dim_mask = tl.where(tl.arange(0, HEAD_SIZE_PADDED) < HEAD_SIZE, 1, + 0).to(tl.int1) + + # Q : (num_queries_per_kv, HEAD_SIZE,) + Q = tl.load( + query_ptr + query_offset + tl.arange(0, HEAD_SIZE_PADDED)[None, :], + mask=dim_mask[None, :] & head_mask[:, None], + other=0.0, + ) + + block_table_offset = seq_idx * block_table_stride + + M = tl.full([num_queries_per_kv_padded], float("-inf"), dtype=tl.float32) + L = tl.full([num_queries_per_kv_padded], 1.0, dtype=tl.float32) + acc = tl.zeros([num_queries_per_kv_padded, HEAD_SIZE_PADDED], + dtype=tl.float32) + + # sequence len for this particular sequence + seq_len = tl.load(seq_lens_ptr + seq_idx) + + # alibi slope for this head + if USE_ALIBI_SLOPES: + alibi_slope = tl.load(alibi_slopes_ptr + query_head_idx, + mask=head_mask, + other=0.0) + + num_blocks = cdiv_fn(seq_len, BLOCK_SIZE) + + # iterate through tiles + for j in range(0, num_blocks): + + physical_block_idx = tl.load(block_tables_ptr + block_table_offset + j) + + offs_n = tl.arange(0, BLOCK_SIZE) + offs_d = tl.arange(0, HEAD_SIZE_PADDED) + + v_offset = (physical_block_idx * stride_v_cache_0 + + kv_head_idx * stride_v_cache_1 + + offs_d[None, :] * stride_v_cache_2 + + offs_n[:, None] * stride_v_cache_3) + + k_offset = (physical_block_idx * stride_k_cache_0 + + kv_head_idx * stride_k_cache_1 + + (offs_d[:, None] // x) * stride_k_cache_2 + + offs_n[None, :] * stride_k_cache_3 + + (offs_d[:, None] % x) * stride_k_cache_4) + + # K : (HEAD_SIZE, BLOCK_SIZE) + K_load = tl.load(key_cache_ptr + k_offset, + mask=dim_mask[:, None], + other=0.0) + + if K_load.dtype.is_fp8(): + K = (K_load.to(tl.float32) * tl.load(k_scale)).to(Q.dtype) + else: + K = K_load + + # V : (BLOCK_SIZE, HEAD_SIZE) + V_load = tl.load(value_cache_ptr + v_offset, + mask=dim_mask[None, :], + other=0.0) + + if V_load.dtype.is_fp8(): + V = (V_load.to(tl.float32) * tl.load(v_scale)).to(Q.dtype) + else: + V = V_load + + seq_offset = j * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + boundary = tl.full([BLOCK_SIZE], seq_len, dtype=tl.int32) + seq_mask = seq_offset[None, :] < boundary + + # S : (num_queries_per_kv, BLOCK_SIZE,) + S = tl.where(head_mask[:, None] & seq_mask, 0.0, + float("-inf")).to(tl.float32) + S += scale * tl.dot(Q, K) + + context_len = seq_len - 1 + + if SLIDING_WINDOW > 0: + S = tl.where((context_len - seq_offset) < SLIDING_WINDOW, S, + -10000) + + if USE_ALIBI_SLOPES: + S += alibi_slope[:, None] * (seq_offset - context_len) + + # compute running maximum + # m_j : (num_queries_per_kv,) + m_j = tl.maximum(M, tl.max(S, axis=1)) + + # P : (num_queries_per_kv, BLOCK_SIZE,) + P = tl.exp(S - m_j[:, None]) + + # l_j : (num_queries_per_kv,) + l_j = tl.sum(P, axis=1) + + # alpha : (num_queries_per_kv, ) + alpha = tl.exp(M - m_j) + + # acc : (num_queries_per_kv, BLOCK_SIZE,) + acc = acc * alpha[:, None] + + # update constants + L = L * alpha + l_j + M = m_j + + # acc : (num_queries_per_kv, BLOCK_SIZE,) + acc += tl.dot(P.to(V.dtype), V) + + # epilogue + acc = acc / L[:, None] + + output_offset = (cur_batch_in_all_start_index * output_stride_0 + + query_head_idx * output_stride_1) + + tl.store( + output_ptr + output_offset[:, None] + + tl.arange(0, HEAD_SIZE_PADDED)[None, :], + acc, + mask=dim_mask[None, :] & head_mask[:, None], + ) + + +# === cdiv_fn === +@triton.jit +def cdiv_fn(x, y): + return (x + y - 1) // y + diff --git a/benchmarks/TritonNPUBench/vllm/kernel_unified_attention_2d/kernel_unified_attention_2d.py b/benchmarks/TritonNPUBench/vllm/kernel_unified_attention_2d/kernel_unified_attention_2d.py new file mode 100644 index 00000000..61646c8e --- /dev/null +++ b/benchmarks/TritonNPUBench/vllm/kernel_unified_attention_2d/kernel_unified_attention_2d.py @@ -0,0 +1,397 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/vllm_operator_cases/newtest_cases/test_kernel_unified_attention_2d.py +# Main kernel: kernel_unified_attention_2d +# PT file: test_kernel_unified_attention_2d_v2.pt + +import triton +import triton.language as tl +import torch +float8_info = torch.finfo(torch.float16) + +# === kernel_unified_attention_2d === +@triton.jit +def kernel_unified_attention_2d( + output_ptr, # [num_tokens, num_query_heads, head_size] + query_ptr, # [num_tokens, num_query_heads, head_size] + key_cache_ptr, # [num_blks, blk_size, num_kv_heads, head_size] + value_cache_ptr, # [num_blks, blk_size, num_kv_heads, head_size] + sink_ptr, # [num_query_heads] + block_tables_ptr, # [num_seqs, max_num_blocks_per_seq] + seq_lens_ptr, # [num_seqs] + alibi_slopes_ptr, # [num_query_heads] + qq_bias_ptr, # [num_query_tokens, num_query_tokens] + scale, # float32 + k_scale, # float32 + v_scale, # float32 + out_scale, # float32 + softcap, # float32 + num_query_heads: tl.constexpr, # int + num_queries_per_kv: tl.constexpr, # int + block_table_stride: tl.int64, # int + query_stride_0: tl.int64, # int + query_stride_1: tl.int64, # int, should be equal to head_size + output_stride_0: tl.int64, # int + output_stride_1: tl.int64, # int, should be equal to head_size + qq_bias_stride_0: tl.int64, # int + BLOCK_SIZE: tl.constexpr, # int + TILE_SIZE: tl.constexpr, # int must be power of 2 + HEAD_SIZE: tl.constexpr, # int + HEAD_SIZE_PADDED: tl.constexpr, # int, must be power of 2 + USE_ALIBI_SLOPES: tl.constexpr, # bool + USE_ALIBI_SQRT: tl.constexpr, # bool + USE_QQ_BIAS: tl.constexpr, # bool + USE_SOFTCAP: tl.constexpr, # bool + USE_SINKS: tl.constexpr, # bool + SLIDING_WINDOW: tl.constexpr, # int + USE_MM_PREFIX: tl.constexpr, # bool + MAX_MM_RANGES: tl.constexpr, # int + mm_prefix_range_ptr, # [num_seqs] - prefix length for each sequence + stride_k_cache_0: tl.int64, # int + stride_k_cache_1: tl.int64, # int + stride_k_cache_2: tl.int64, # int + stride_k_cache_3: tl.constexpr, # int + stride_v_cache_0: tl.int64, # int + stride_v_cache_1: tl.int64, # int + stride_v_cache_2: tl.int64, # int + stride_v_cache_3: tl.constexpr, # int + query_start_len_ptr, # [num_seqs+1] + BLOCK_Q: tl.constexpr, # int + num_seqs: tl.int32, + BLOCK_M: tl.constexpr, # int + USE_FP8: tl.constexpr, # bool + FP8_MIN: tl.constexpr = float8_info.min, + FP8_MAX: tl.constexpr = float8_info.max, +): + q_block_global_idx = tl.program_id(0) + kv_head_idx = tl.program_id(1) + + seq_idx = find_seq_idx( + query_start_len_ptr, q_block_global_idx, num_seqs, BLOCK_Q, True + ) + + q_block_start_idx = tl.load(query_start_len_ptr + seq_idx) // BLOCK_Q + seq_idx + + q_block_local_idx = q_block_global_idx - q_block_start_idx + + cur_batch_in_all_start_index = tl.load(query_start_len_ptr + seq_idx) + cur_batch_in_all_stop_index = tl.load(query_start_len_ptr + seq_idx + 1) + + cur_batch_query_len = cur_batch_in_all_stop_index - cur_batch_in_all_start_index + + if q_block_local_idx * BLOCK_Q >= cur_batch_query_len: + return + + offs_m = tl.arange(0, BLOCK_M) + offs_d = tl.arange(0, HEAD_SIZE_PADDED) + offs_t = tl.arange(0, TILE_SIZE) + query_pos = q_block_local_idx * BLOCK_Q + offs_m // num_queries_per_kv + + query_offset_0 = cur_batch_in_all_start_index + query_pos + query_offset_1 = kv_head_idx * num_queries_per_kv + offs_m % num_queries_per_kv + query_offset = ( + query_offset_0[:, None] * query_stride_0 + + query_offset_1[:, None] * query_stride_1 + + offs_d[None, :] + ) + + dim_mask = tl.where(offs_d < HEAD_SIZE, 1, 0).to(tl.int1) + query_mask_0 = tl.where(query_pos < cur_batch_query_len, 1, 0).to(tl.int1) + query_mask_1 = tl.where(query_offset_1 < num_query_heads, 1, 0).to(tl.int1) + + # Q : (BLOCK_M, HEAD_SIZE_PADDED) + Q = tl.load( + query_ptr + query_offset, + mask=dim_mask[None, :] & query_mask_0[:, None] & query_mask_1[:, None], + other=0.0, + ) + + block_table_offset = seq_idx * block_table_stride + + if not USE_SINKS: + M = tl.full([BLOCK_M], float("-inf"), dtype=tl.float32) + else: + M = tl.load( + sink_ptr + query_offset_1, + mask=query_mask_1, + other=float("-inf"), + ).to(dtype=tl.float32) + + L = tl.full([BLOCK_M], 1.0, dtype=tl.float32) + acc = tl.zeros([BLOCK_M, HEAD_SIZE_PADDED], dtype=tl.float32) + + # sequence len for this particular sequence + seq_len = tl.load(seq_lens_ptr + seq_idx) + + # context length for this particular sequences + context_len = seq_len - cur_batch_query_len + + # alibi slope for this head + if USE_ALIBI_SLOPES: + alibi_slope = tl.load( + alibi_slopes_ptr + query_offset_1, mask=query_mask_1, other=0.0 + ) + + # query-query attention bias + if USE_QQ_BIAS: + qq_bias_row_ptrs = ( + qq_bias_ptr + query_pos[:, None] * qq_bias_stride_0 + ) # shape: [BLOCK_M] + + # compute the length of the longest sequence prefix spanned by any + # query token in the current q_block (q_block_local_idx) + max_seq_prefix_len = ( + context_len + + q_block_local_idx * BLOCK_Q + + (BLOCK_M - 1) // num_queries_per_kv + + 1 + ) + + if USE_MM_PREFIX: + # image bidirectional attention ranges require a full range + # including q_block padding to make sure doc mask is correct + max_seq_prefix_len = tl.maximum(max_seq_prefix_len, seq_len) + else: + # adjust for potential padding in the last q_block by considering the + # actual sequence length + max_seq_prefix_len = tl.minimum(max_seq_prefix_len, seq_len) + + # calculate the number of tiles that need to be processed to + # cover the longest sequence prefix (due to causal masking, tiles beyond + # this prefix can be skipped) + num_tiles = cdiv_fn(max_seq_prefix_len, TILE_SIZE) + + # ---- Sliding-window tile pruning -------------------- + # Default: keep previous global behavior + tile_start = 0 + tile_end = num_tiles + # TODO(Isotr0py): sliding window pruning with image bidirectional mask + if SLIDING_WINDOW > 0 and not USE_MM_PREFIX: + # Query rows covered by this Q-block + qpos_lo = q_block_local_idx * BLOCK_Q + qpos_hi = tl.minimum( + qpos_lo + (BLOCK_M - 1) // num_queries_per_kv, + cur_batch_query_len - 1, + ) + # For sliding window, each query position q can only attend to + # keys in the range [q_abs - SLIDING_WINDOW + 1, q_abs] + # where q_abs = context_len + q + # The union of allowed key positions for this Q-block is: + # [context_len + qpos_lo - SLIDING_WINDOW + 1, context_len + qpos_hi] + first_allowed_key = context_len + qpos_lo - SLIDING_WINDOW + 1 + last_allowed_key = context_len + qpos_hi + # Convert to tile indices and clamp + tile_start = tl.maximum(0, first_allowed_key // TILE_SIZE) + tile_end = tl.minimum((last_allowed_key // TILE_SIZE) + 1, num_tiles) + + # iterate through tiles (now limited to the sliding window range) + for j in range(tile_start, tile_end): + seq_offset = j * TILE_SIZE + offs_t + tile_mask = seq_offset < max_seq_prefix_len + + physical_block_idx = tl.load( + block_tables_ptr + block_table_offset + seq_offset // BLOCK_SIZE + ).to(tl.int64) + + v_offset = ( + physical_block_idx[:, None] * stride_v_cache_0 + + kv_head_idx * stride_v_cache_2 + + offs_d[None, :] * stride_v_cache_3 + + (seq_offset % BLOCK_SIZE)[:, None] * stride_v_cache_1 + ) + + k_offset = ( + physical_block_idx[None, :] * stride_k_cache_0 + + kv_head_idx * stride_k_cache_2 + + offs_d[:, None] * stride_k_cache_3 + + (seq_offset % BLOCK_SIZE)[None, :] * stride_k_cache_1 + ) + + # K : (HEAD_SIZE, TILE_SIZE) + K_load = tl.load( + key_cache_ptr + k_offset, + mask=dim_mask[:, None] & tile_mask[None, :], + other=0.0, + ) + + if K_load.dtype.is_fp8(): + if Q.dtype.is_fp8(): + K = K_load + else: + K = (K_load.to(tl.float32) * tl.load(k_scale)).to(Q.dtype) + else: + K = K_load + + # V : (TILE_SIZE, HEAD_SIZE) + V_load = tl.load( + value_cache_ptr + v_offset, + mask=dim_mask[None, :] & tile_mask[:, None], + other=0.0, + ) + + if V_load.dtype.is_fp8(): + if Q.dtype.is_fp8(): + V = V_load + else: + V = (V_load.to(tl.float32) * tl.load(v_scale)).to(Q.dtype) + else: + V = V_load + + # Compute attention mask: causal by default (key <= query) + query_abs_pos = context_len + query_pos[:, None] + seq_mask = seq_offset[None, :] <= query_abs_pos + + # Apply sliding window to base mask BEFORE mm_prefix OR. + # Order must match FlexAttention: (causal AND sliding_window) OR mm_prefix + if SLIDING_WINDOW > 0: + seq_mask = seq_mask & ((query_abs_pos - seq_offset) < SLIDING_WINDOW) + + # PrefixLM: extend mask with bidirectional ranges for multimodal tokens. + # Applied AFTER sliding window so mm_prefix ranges override SW restriction. + if USE_MM_PREFIX: + for i in range(MAX_MM_RANGES): + range_start = tl.load( + mm_prefix_range_ptr + seq_idx * MAX_MM_RANGES * 2 + i * 2 + ) + range_end = tl.load( + mm_prefix_range_ptr + seq_idx * MAX_MM_RANGES * 2 + i * 2 + 1 + ) + + is_valid = range_start < range_end + q_in_range = ( + (query_abs_pos >= range_start) + & (query_abs_pos <= range_end) + & is_valid + ) + k_in_range = ( + (seq_offset[None, :] >= range_start) + & (seq_offset[None, :] <= range_end) + & is_valid + ) + seq_mask |= q_in_range & k_in_range + + # S : (BLOCK_M, TILE_SIZE) + S = tl.zeros(shape=(BLOCK_M, TILE_SIZE), dtype=tl.float32) + + S += scale * tl.dot(Q, K) + + if USE_SOFTCAP: + S = apply_softcap(S, softcap) + + S = tl.where( + query_mask_1[:, None] & query_mask_0[:, None] & seq_mask, S, float("-inf") + ) + + if USE_ALIBI_SLOPES: + if USE_ALIBI_SQRT: + relative_pos = seq_offset - (context_len + query_pos[:, None]) + alibi_offset = tl.where( + relative_pos <= 0, + -tl.sqrt((-relative_pos).to(tl.float32)), + 0.0, + ) + else: + alibi_offset = seq_offset - context_len + S += alibi_slope[:, None] * alibi_offset + + if USE_QQ_BIAS: + # compute key positions relative to query section + key_rel_pos = seq_offset - context_len # shape: [BLOCK_SIZE] + # load bias only for keys that correspond to queries + is_query_key = key_rel_pos >= 0 and key_rel_pos < qq_bias_stride_0 + qq_bias = tl.load( + qq_bias_row_ptrs + key_rel_pos[None, :], + mask=is_query_key[None, :], # avoid OOB for context keys + other=0.0, + ) + S += qq_bias + + # compute running maximum + # m_j : (BLOCK_M,) + m_j = tl.maximum(M, tl.max(S, axis=1)) + + # For sliding window there's a chance the max is -inf due to masking of + # the entire row. In this case we need to set m_j 0 to avoid NaN + m_j = tl.where(m_j > float("-inf"), m_j, 0.0) + + # P : (BLOCK_M, TILE_SIZE) + P = tl.exp(S - m_j[:, None]) + + # l_j : (BLOCK_M,) + l_j = tl.sum(P, axis=1) + + # alpha : (BLOCK_M, ) + alpha = tl.exp(M - m_j) + + # acc : (BLOCK_M, HEAD_SIZE_PADDED) + acc = acc * alpha[:, None] + + # update constants + L = L * alpha + l_j + M = m_j + + if SLIDING_WINDOW: + qpos_lo = q_block_local_idx * BLOCK_Q + V = tl.where( + (context_len + qpos_lo - seq_offset[:, None]) < SLIDING_WINDOW, V, 0.0 + ) + + # acc : (BLOCK_M, HEAD_SIZE_PADDED) + acc += tl.dot(P.to(V.dtype), V) + + # epilogue + acc = acc / L[:, None] + if USE_FP8: + acc = acc * tl.load(out_scale) + acc = tl.clamp(acc, FP8_MIN, FP8_MAX) + + output_offset = ( + query_offset_0[:, None] * output_stride_0 + + query_offset_1[:, None] * output_stride_1 + + offs_d[None, :] + ) + + tl.store( + output_ptr + output_offset, + acc, + mask=dim_mask[None, :] & query_mask_0[:, None] & query_mask_1[:, None], + ) + + +# === cdiv_fn === +@triton.jit +def cdiv_fn(x, y): + return (x + y - 1) // y + + +# === find_seq_idx === +@triton.jit +def find_seq_idx( + query_start_len_ptr, + target_idx, + num_seqs, + BLOCK_Q: tl.constexpr, + use_q_block_mode: tl.constexpr, +): + left: tl.int32 = 0 + right = num_seqs + while left < right: + mid = (left + right) // 2 + val = tl.load(query_start_len_ptr + mid) + mid_val = val // BLOCK_Q + mid if use_q_block_mode else val + + if mid_val <= target_idx: + left = mid + 1 + else: + right = mid + + return left - 1 + + +# === apply_softcap === +@triton.jit +def apply_softcap(S, x): + Sdiv = S / x + p1 = tl.exp(Sdiv) + p2 = tl.exp(-Sdiv) + return x * (p1 - p2) / (p1 + p2) + diff --git a/benchmarks/TritonNPUBench/vllm/kernel_unified_attention_3d/kernel_unified_attention_3d.py b/benchmarks/TritonNPUBench/vllm/kernel_unified_attention_3d/kernel_unified_attention_3d.py new file mode 100644 index 00000000..18f724c0 --- /dev/null +++ b/benchmarks/TritonNPUBench/vllm/kernel_unified_attention_3d/kernel_unified_attention_3d.py @@ -0,0 +1,395 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/vllm_operator_cases/newtest_cases/test_kernel_unified_attention_3d.py +# Main kernel: kernel_unified_attention_3d +# PT file: kernel_unified_attention_3d_test_data_fixed_v2.pt + +import triton +import triton.language as tl + + +# === kernel_unified_attention_3d === +@triton.jit +def kernel_unified_attention_3d( + segm_output_ptr, + # [num_tokens, num_query_heads, num_segments, head_size_padded] + segm_max_ptr, # [num_tokens, num_query_heads, num_segments] + segm_expsum_ptr, # [num_tokens, num_query_heads, num_segments] + query_ptr, # [num_tokens, num_query_heads, head_size] + key_cache_ptr, # [num_blks, num_kv_heads, head_size // x, blk_size, x] + value_cache_ptr, # [num_blks, num_kv_heads, head_size, blk_size] + sink_ptr, # [num_query_heads] + block_tables_ptr, # [num_seqs, max_num_blocks_per_seq] + seq_lens_ptr, # [num_seqs] + alibi_slopes_ptr, # [num_query_heads] + qq_bias_ptr, # [num_query_tokens, num_query_tokens] + scale, # float32 + k_scale, # float32 + v_scale, # float32 + softcap, # float32 + num_query_heads: tl.constexpr, # int + num_queries_per_kv: tl.constexpr, # int + block_table_stride: tl.int64, # int + query_stride_0: tl.int64, # int + query_stride_1: tl.int64, # int, should be equal to head_size + qq_bias_stride_0: tl.int64, # int + BLOCK_SIZE: tl.constexpr, # int + TILE_SIZE: tl.constexpr, # int, must be power of 2 + HEAD_SIZE: tl.constexpr, # int + HEAD_SIZE_PADDED: tl.constexpr, # int, must be power of 2 + USE_ALIBI_SLOPES: tl.constexpr, # bool + USE_QQ_BIAS: tl.constexpr, # bool + USE_SOFTCAP: tl.constexpr, # bool + USE_SINKS: tl.constexpr, # bool + SLIDING_WINDOW: tl.constexpr, # int + stride_k_cache_0: tl.int64, # int + stride_k_cache_1: tl.int64, # int + stride_k_cache_2: tl.int64, # int + stride_k_cache_3: tl.constexpr, # int + stride_v_cache_0: tl.int64, # int + stride_v_cache_1: tl.int64, # int + stride_v_cache_2: tl.int64, # int + stride_v_cache_3: tl.constexpr, # int + query_start_len_ptr, # [num_seqs+1] + BLOCK_Q: tl.constexpr, # int + num_seqs: tl.int32, + BLOCK_M: tl.constexpr, # int + NUM_SEGMENTS_PER_SEQ: tl.constexpr, # int + USE_MM_PREFIX: tl.constexpr, # bool + MAX_MM_RANGES: tl.constexpr, # int + mm_prefix_range_ptr, # [num_seqs] - prefix length for each sequence +): + q_block_global_idx = tl.program_id(0) + kv_head_idx = tl.program_id(1) + segm_idx = tl.program_id(2) + + seq_idx = find_seq_idx( + query_start_len_ptr, q_block_global_idx, num_seqs, BLOCK_Q, True + ) + + q_block_start_idx = tl.load(query_start_len_ptr + seq_idx) // BLOCK_Q + seq_idx + + q_block_local_idx = q_block_global_idx - q_block_start_idx + + cur_batch_in_all_start_index = tl.load(query_start_len_ptr + seq_idx) + cur_batch_in_all_stop_index = tl.load(query_start_len_ptr + seq_idx + 1) + + cur_batch_query_len = cur_batch_in_all_stop_index - cur_batch_in_all_start_index + + if q_block_local_idx * BLOCK_Q >= cur_batch_query_len: + return + + # sequence len for this particular sequence + seq_len = tl.load(seq_lens_ptr + seq_idx) + + # number of segments for this particular sequence + num_segments = NUM_SEGMENTS_PER_SEQ + tiles_per_segment = cdiv_fn(seq_len, num_segments * TILE_SIZE) + + if segm_idx * tiles_per_segment * TILE_SIZE >= seq_len: + return + + offs_m = tl.arange(0, BLOCK_M) + offs_d = tl.arange(0, HEAD_SIZE_PADDED) + offs_t = tl.arange(0, TILE_SIZE) + query_pos = q_block_local_idx * BLOCK_Q + offs_m // num_queries_per_kv + + query_offset_0 = cur_batch_in_all_start_index + query_pos + query_offset_1 = kv_head_idx * num_queries_per_kv + offs_m % num_queries_per_kv + query_offset = ( + query_offset_0[:, None] * query_stride_0 + + query_offset_1[:, None] * query_stride_1 + + offs_d[None, :] + ) + + dim_mask = tl.where(offs_d < HEAD_SIZE, 1, 0).to(tl.int1) + query_mask_0 = tl.where(query_pos < cur_batch_query_len, 1, 0).to(tl.int1) + query_mask_1 = tl.where(query_offset_1 < num_query_heads, 1, 0).to(tl.int1) + + # Q : (BLOCK_M, HEAD_SIZE_PADDED) + Q = tl.load( + query_ptr + query_offset.to(tl.int32), + mask=dim_mask[None, :] & query_mask_0[:, None] & query_mask_1[:, None], + other=0.0, + ) + + block_table_offset = seq_idx * block_table_stride + + if USE_SINKS: + if segm_idx == 0: + M = tl.load( + sink_ptr + query_offset_1, + mask=query_mask_1, + other=float("-inf"), + ).to(dtype=tl.float32) + else: + M = tl.full([BLOCK_M], float("-inf"), dtype=tl.float32) + else: + M = tl.full([BLOCK_M], float("-inf"), dtype=tl.float32) + + L = tl.full([BLOCK_M], 1.0, dtype=tl.float32) + acc = tl.zeros([BLOCK_M, HEAD_SIZE_PADDED], dtype=tl.float32) + + # context length for this particular sequences + context_len = seq_len - cur_batch_query_len + + # alibi slope for this head + if USE_ALIBI_SLOPES: + alibi_slope = tl.load( + alibi_slopes_ptr + query_offset_1, mask=query_mask_1, other=0.0 + ) + + # query-query attention bias + if USE_QQ_BIAS: + qq_bias_row_ptrs = ( + qq_bias_ptr + query_pos[:, None] * qq_bias_stride_0 + ) # shape: [BLOCK_M] + + # compute the length of the longest sequence prefix spanned by any + # query token in the current q_block (q_block_local_idx) + max_seq_prefix_len = ( + context_len + + q_block_local_idx * BLOCK_Q + + (BLOCK_M - 1) // num_queries_per_kv + + 1 + ) + + # adjust for potential padding in the last q_block by considering the + # actual sequence length + max_seq_prefix_len = tl.minimum(max_seq_prefix_len, seq_len) + + # calculate the number of tiles that need to be processed to + # cover the longest sequence prefix (due to causal masking, tiles beyond + # this prefix can be skipped) + num_tiles = cdiv_fn(max_seq_prefix_len, TILE_SIZE) + + # ---- Sliding-window tile pruning -------------------- + # Default: keep previous global behavior + tile_start = 0 + tile_end = num_tiles + # TODO(Isotr0py): sliding window pruning with image bidirectional mask + if SLIDING_WINDOW > 0 and not USE_MM_PREFIX: + # Query rows covered by this Q-block + qpos_lo = q_block_local_idx * BLOCK_Q + qpos_hi = tl.minimum( + qpos_lo + (BLOCK_M - 1) // num_queries_per_kv, + cur_batch_query_len - 1, + ) + # For sliding window, each query position q can only attend to + # keys in the range [q_abs - SLIDING_WINDOW + 1, q_abs] + # where q_abs = context_len + q + # The union of allowed key positions for this Q-block is: + # [context_len + qpos_lo - SLIDING_WINDOW + 1, context_len + qpos_hi] + first_allowed_key = context_len + qpos_lo - SLIDING_WINDOW + 1 + last_allowed_key = context_len + qpos_hi + # Convert to tile indices and clamp + tile_start = tl.maximum(0, first_allowed_key // TILE_SIZE) + tile_end = tl.minimum((last_allowed_key // TILE_SIZE) + 1, num_tiles) + + # iterate through tiles (now limited to the sliding window range) + for j in range( + max(segm_idx * tiles_per_segment, tile_start), + min((segm_idx + 1) * tiles_per_segment, tile_end), + ): + seq_offset = j * TILE_SIZE + offs_t + tile_mask = seq_offset < max_seq_prefix_len + + physical_block_idx = tl.load( + block_tables_ptr + block_table_offset + seq_offset // BLOCK_SIZE + ).to(tl.int64) + + v_offset = ( + physical_block_idx[:, None] * stride_v_cache_0 + + kv_head_idx * stride_v_cache_2 + + offs_d[None, :] * stride_v_cache_3 + + (seq_offset % BLOCK_SIZE)[:, None] * stride_v_cache_1 + ) + + k_offset = ( + physical_block_idx[None, :] * stride_k_cache_0 + + kv_head_idx * stride_k_cache_2 + + offs_d[:, None] * stride_k_cache_3 + + (seq_offset % BLOCK_SIZE)[None, :] * stride_k_cache_1 + ) + + # K : (HEAD_SIZE, TILE_SIZE) + K_load = tl.load( + key_cache_ptr + k_offset, + mask=dim_mask[:, None] & tile_mask[None, :], + other=0.0, + ) + + if K_load.dtype.is_fp8(): + if Q.dtype.is_fp8(): + K = K_load + else: + K = (K_load.to(tl.float32) * tl.load(k_scale)).to(Q.dtype) + else: + K = K_load + + # V : (TILE_SIZE, HEAD_SIZE) + V_load = tl.load( + value_cache_ptr + v_offset, + mask=dim_mask[None, :] & tile_mask[:, None], + other=0.0, + ) + + if V_load.dtype.is_fp8(): + if Q.dtype.is_fp8(): + V = V_load + else: + V = (V_load.to(tl.float32) * tl.load(v_scale)).to(Q.dtype) + else: + V = V_load + + # Compute attention mask: causal by default (key <= query) + query_abs_pos = context_len + query_pos[:, None] + seq_mask = seq_offset[None, :] <= query_abs_pos + + # Apply sliding window to base mask BEFORE mm_prefix OR. + # Order must match FlexAttention: (causal AND sliding_window) OR mm_prefix + if SLIDING_WINDOW > 0: + seq_mask = seq_mask & ((query_abs_pos - seq_offset) < SLIDING_WINDOW) + + # PrefixLM: extend mask with bidirectional ranges for multimodal tokens. + # Applied AFTER sliding window so mm_prefix ranges override SW restriction. + if USE_MM_PREFIX: + for i in range(MAX_MM_RANGES): + range_start = tl.load( + mm_prefix_range_ptr + seq_idx * MAX_MM_RANGES * 2 + i * 2 + ) + range_end = tl.load( + mm_prefix_range_ptr + seq_idx * MAX_MM_RANGES * 2 + i * 2 + 1 + ) + + is_valid = range_start < range_end + q_in_range = ( + (query_abs_pos >= range_start) + & (query_abs_pos <= range_end) + & is_valid + ) + k_in_range = ( + (seq_offset[None, :] >= range_start) + & (seq_offset[None, :] <= range_end) + & is_valid + ) + seq_mask |= q_in_range & k_in_range + + # S : (BLOCK_M, TILE_SIZE) + S = tl.zeros(shape=(BLOCK_M, TILE_SIZE), dtype=tl.float32) + S += scale * tl.dot(Q, K) + + if USE_SOFTCAP: + S = apply_softcap(S, softcap) + + S = tl.where( + query_mask_1[:, None] & query_mask_0[:, None] & seq_mask, S, float("-inf") + ) + + if USE_ALIBI_SLOPES: + S += alibi_slope[:, None] * (seq_offset - context_len) + + if USE_QQ_BIAS: + # compute key positions relative to query section + key_rel_pos = seq_offset - context_len # shape: [BLOCK_SIZE] + # load bias only for keys that correspond to queries + is_query_key = key_rel_pos >= 0 and key_rel_pos < qq_bias_stride_0 + qq_bias = tl.load( + qq_bias_row_ptrs + key_rel_pos[None, :], + mask=is_query_key[None, :], # avoid OOB for context keys + other=0.0, + ) + S += qq_bias + + # compute running maximum + # m_j : (BLOCK_M,) + m_j = tl.maximum(M, tl.max(S, axis=1)) + + # For sliding window there's a chance the max is -inf due to masking of + # the entire row. In this case we need to set m_j 0 to avoid NaN + m_j = tl.where(m_j > float("-inf"), m_j, 0.0) + + # P : (BLOCK_M, TILE_SIZE,) + P = tl.exp(S - m_j[:, None]) + + # l_j : (BLOCK_M,) + l_j = tl.sum(P, axis=1) + + # alpha : (BLOCK_M, ) + alpha = tl.exp(M - m_j) + + # acc : (BLOCK_M, HEAD_SIZE_PADDED) + acc = acc * alpha[:, None] + + # update constants + L = L * alpha + l_j + M = m_j + + if SLIDING_WINDOW: + qpos_lo = q_block_local_idx * BLOCK_Q + V = tl.where( + (context_len + qpos_lo - seq_offset[:, None]) < SLIDING_WINDOW, V, 0.0 + ) + + # acc : (BLOCK_M, HEAD_SIZE_PADDED) + acc += tl.dot(P.to(V.dtype), V) + + segm_output_offset = ( + query_offset_0[:, None].to(tl.int64) + * (num_query_heads * NUM_SEGMENTS_PER_SEQ * HEAD_SIZE_PADDED) + + query_offset_1[:, None] * (NUM_SEGMENTS_PER_SEQ * HEAD_SIZE_PADDED) + + segm_idx * HEAD_SIZE_PADDED + + tl.arange(0, HEAD_SIZE_PADDED)[None, :] + ) + tl.store( + segm_output_ptr + segm_output_offset.to(tl.int32), + acc, + mask=dim_mask[None, :] & query_mask_0[:, None] & query_mask_1[:, None], + ) + segm_offset = ( + query_offset_0.to(tl.int64) * (num_query_heads * NUM_SEGMENTS_PER_SEQ) + + query_offset_1 * NUM_SEGMENTS_PER_SEQ + + segm_idx + ) + tl.store(segm_max_ptr + segm_offset.to(tl.int32), M, mask=query_mask_0 & query_mask_1) + tl.store(segm_expsum_ptr + segm_offset.to(tl.int32), L, mask=query_mask_0 & query_mask_1) + + +# === cdiv_fn === +@triton.jit +def cdiv_fn(x, y): + return (x + y - 1) // y + + +# === find_seq_idx === +@triton.jit +def find_seq_idx( + query_start_len_ptr, + target_idx, + num_seqs, + BLOCK_Q: tl.constexpr, + use_q_block_mode: tl.constexpr, +): + #left: tl.int32 = 0 + left = 0 + right = num_seqs + while left < right: + mid = (left + right) // 2 + val = tl.load(query_start_len_ptr + mid) + mid_val = val // BLOCK_Q + mid if use_q_block_mode else val + + if mid_val <= target_idx: + left = mid + 1 + else: + right = mid + + return left - 1 + + +# === apply_softcap === +@triton.jit +def apply_softcap(S, x): + Sdiv = S / x + p1 = tl.exp(Sdiv) + p2 = tl.exp(-Sdiv) + return x * (p1 - p2) / (p1 + p2) + diff --git a/benchmarks/TritonNPUBench/vllm/l2norm_fwd_kernel/l2norm_fwd_kernel.py b/benchmarks/TritonNPUBench/vllm/l2norm_fwd_kernel/l2norm_fwd_kernel.py new file mode 100644 index 00000000..848a129c --- /dev/null +++ b/benchmarks/TritonNPUBench/vllm/l2norm_fwd_kernel/l2norm_fwd_kernel.py @@ -0,0 +1,29 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/vllm_operator_cases/newtest_cases/test_l2norm_fwd_kernel.py +# Main kernel: l2norm_fwd_kernel +# PT file: l2norm_fwd_kernel.pt + +import triton +import triton.language as tl + + +# === l2norm_fwd_kernel === +@triton.jit(do_not_specialize=["NB"]) +def l2norm_fwd_kernel( + x, + y, + eps, + NB, + T, + D: tl.constexpr, + BT: tl.constexpr, + BD: tl.constexpr, +): + i_t = tl.program_id(0) + p_x = tl.make_block_ptr(x, (T, D), (D, 1), (i_t * BT, 0), (BT, BD), (1, 0)) + b_x = tl.load(p_x, boundary_check=(0, 1)).to(tl.float32) + b_var = tl.sum(b_x * b_x, axis=1) + b_y = b_x / tl.sqrt(b_var + eps)[:, None] + p_y = tl.make_block_ptr(y, (T, D), (D, 1), (i_t * BT, 0), (BT, BD), (1, 0)) + tl.store(p_y, b_y.to(p_y.dtype.element_ty), boundary_check=(0, 1)) + diff --git a/benchmarks/TritonNPUBench/vllm/l2norm_fwd_kernel1/l2norm_fwd_kernel1.py b/benchmarks/TritonNPUBench/vllm/l2norm_fwd_kernel1/l2norm_fwd_kernel1.py new file mode 100644 index 00000000..8888d37f --- /dev/null +++ b/benchmarks/TritonNPUBench/vllm/l2norm_fwd_kernel1/l2norm_fwd_kernel1.py @@ -0,0 +1,32 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/vllm_operator_cases/newtest_cases/test_l2norm_fwd_kernel1.py +# Main kernel: l2norm_fwd_kernel1 +# PT file: l2norm_fwd_kernel1_v2.pt + +import triton +import triton.language as tl + + +# === l2norm_fwd_kernel1 === +@triton.jit +def l2norm_fwd_kernel1( + x, + y, + D, + BD: tl.constexpr, + eps, +): + i_t = tl.program_id(0) + x += i_t * D + y += i_t * D + # Compute mean and variance + cols = tl.arange(0, BD) + mask = cols < D + b_x = tl.load(x + cols, mask=mask, other=0.0).to(tl.float32) + b_var = tl.sum(b_x * b_x, axis=0) + b_rstd = 1 / tl.sqrt(b_var + eps) + # tl.store(Rstd + i_t, rstd) + # Normalize and apply linear transformation + b_y = b_x * b_rstd + tl.store(y + cols, b_y, mask=mask) + diff --git a/benchmarks/TritonNPUBench/vllm/l2norm_fwd_kernel2/l2norm_fwd_kernel2.py b/benchmarks/TritonNPUBench/vllm/l2norm_fwd_kernel2/l2norm_fwd_kernel2.py new file mode 100644 index 00000000..1b316081 --- /dev/null +++ b/benchmarks/TritonNPUBench/vllm/l2norm_fwd_kernel2/l2norm_fwd_kernel2.py @@ -0,0 +1,22 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/vllm_operator_cases/newtest_cases/test_l2norm_fwd_kernel2.py +# Main kernel: l2norm_fwd_kernel2 +# PT file: l2norm_fwd_kernel2_v2.pt + +import triton +import triton.language as tl + + +# === l2norm_fwd_kernel2 === +@triton.jit +def l2norm_fwd_kernel2(X, Y, eps, M, N: tl.constexpr, MBLOCK: tl.constexpr): + xoffset = tl.program_id(0) * MBLOCK + row_idx = xoffset + tl.arange(0, MBLOCK)[:, None] + xmask = row_idx < M + rindex = tl.arange(0, N)[None, :] + xs = tl.load(X + (rindex + N * row_idx), xmask).to(tl.float32) + square = tl.broadcast_to(xs * xs, [MBLOCK, N]) + square_sum = tl.sum(tl.where(xmask, square, 0), 1)[:, None] + rsqrt = tl.rsqrt(square_sum + eps) + tl.store(Y + (rindex + N * row_idx), xs * rsqrt, xmask) + diff --git a/benchmarks/TritonNPUBench/vllm/layer_norm_fwd_kernel/layer_norm_fwd_kernel.py b/benchmarks/TritonNPUBench/vllm/layer_norm_fwd_kernel/layer_norm_fwd_kernel.py new file mode 100644 index 00000000..bbcff0cc --- /dev/null +++ b/benchmarks/TritonNPUBench/vllm/layer_norm_fwd_kernel/layer_norm_fwd_kernel.py @@ -0,0 +1,107 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/vllm_operator_cases/newtest_cases/test_layer_norm_fwd_kernel.py +# Main kernel: layer_norm_fwd_kernel +# PT file: test_layer_norm_fwd_kernel.pt + +import triton +import triton.language as tl + + +# === layer_norm_fwd_kernel === +@triton.jit +def layer_norm_fwd_kernel( + X, # pointer to the input + Y, # pointer to the output + W, # pointer to the weights + B, # pointer to the biases + Z, # pointer to the other branch + Mean, # pointer to the mean + Rstd, # pointer to the 1/std + stride_x_row, # how much to increase the pointer when moving by 1 row + stride_y_row, + stride_z_row, + M, # number of rows in X + N: tl.constexpr, # number of columns in X + eps, # epsilon to avoid division by zero + BLOCK_N: tl.constexpr, + ROWS_PER_BLOCK: tl.constexpr, + HAS_BIAS: tl.constexpr, + HAS_Z: tl.constexpr, + NORM_BEFORE_GATE: tl.constexpr, + IS_RMS_NORM: tl.constexpr, +): + # Map the program id to the starting row of X and Y it should compute. + row_start = tl.program_id(0) * ROWS_PER_BLOCK + group = tl.program_id(1) + + # Create 2D tile: [ROWS_PER_BLOCK, BLOCK_N] + rows = row_start + tl.arange(0, ROWS_PER_BLOCK) + cols = tl.arange(0, BLOCK_N) + + # Compute offsets for 2D tile + row_offsets = rows[:, None] * stride_x_row + col_offsets = cols[None, :] + group * N + + # Base pointers + X_base = X + row_offsets + col_offsets + Y_base = Y + rows[:, None] * stride_y_row + col_offsets + + # Create mask for valid rows and columns + row_mask = rows[:, None] < M + col_mask = cols[None, :] < N + mask = row_mask & col_mask + + # Load input data with 2D tile + x = tl.load(X_base, mask=mask, other=0.0).to(tl.float32) + + if HAS_Z and not NORM_BEFORE_GATE: + Z_base = Z + rows[:, None] * stride_z_row + col_offsets + z = tl.load(Z_base, mask=mask, other=0.0).to(tl.float32) + x *= z * tl.sigmoid(z) + + # Compute mean and variance per row (reduce along axis 1) + if not IS_RMS_NORM: + mean = tl.sum(x, axis=1) / N # Shape: [ROWS_PER_BLOCK] + # Store mean for each row + mean_offsets = group * M + rows + mean_mask = rows < M + tl.store(Mean + mean_offsets, mean, mask=mean_mask) + # Broadcast mean back to 2D for subtraction + xbar = tl.where(mask, x - mean[:, None], 0.0) + var = tl.sum(xbar * xbar, axis=1) / N # Shape: [ROWS_PER_BLOCK] + else: + xbar = tl.where(mask, x, 0.0) + var = tl.sum(xbar * xbar, axis=1) / N # Shape: [ROWS_PER_BLOCK] + mean = 0.0 # Placeholder for RMS norm + + rstd = tl.rsqrt(var + eps) # Shape: [ROWS_PER_BLOCK] + + # Store rstd for each row + rstd_offsets = group * M + rows + rstd_mask = rows < M + tl.store(Rstd + rstd_offsets, rstd, mask=rstd_mask) + + # Load weights and biases (broadcast across rows) + w_offsets = cols + group * N + w_mask = cols < N + w = tl.load(W + w_offsets, mask=w_mask, other=0.0).to(tl.float32) + + if HAS_BIAS: + b = tl.load(B + w_offsets, mask=w_mask, other=0.0).to(tl.float32) + + # Normalize and apply linear transformation + if not IS_RMS_NORM: + x_hat = (x - mean[:, None]) * rstd[:, None] + else: + x_hat = x * rstd[:, None] + + y = x_hat * w[None, :] + b[None, :] if HAS_BIAS else x_hat * w[None, :] + + if HAS_Z and NORM_BEFORE_GATE: + Z_base = Z + rows[:, None] * stride_z_row + col_offsets + z = tl.load(Z_base, mask=mask, other=0.0).to(tl.float32) + y *= z * tl.sigmoid(z) + + # Write output + tl.store(Y_base, y, mask=mask) + diff --git a/benchmarks/TritonNPUBench/vllm/layer_norm_gated_fwd_kernel/layer_norm_gated_fwd_kernel.py b/benchmarks/TritonNPUBench/vllm/layer_norm_gated_fwd_kernel/layer_norm_gated_fwd_kernel.py new file mode 100644 index 00000000..71bb3f90 --- /dev/null +++ b/benchmarks/TritonNPUBench/vllm/layer_norm_gated_fwd_kernel/layer_norm_gated_fwd_kernel.py @@ -0,0 +1,89 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/vllm_operator_cases/newtest_cases/test_layer_norm_gated_fwd_kernel.py +# Main kernel: layer_norm_gated_fwd_kernel +# PT file: layer_norm_gated_fwd_kernel_v2.pt + +import triton +import triton.language as tl + + +# === layer_norm_gated_fwd_kernel === +@triton.jit +def layer_norm_gated_fwd_kernel( + x, # pointer to the input + g, # pointer to the gate + y, # pointer to the output + w, # pointer to the weights + b, # pointer to the biases + residual, # pointer to the residual + residual_out, # pointer to the residual + mean, # pointer to the mean + rstd, # pointer to the 1/std + eps, # epsilon to avoid division by zero + T, # number of rows in x + D: tl.constexpr, # number of columns in x + BT: tl.constexpr, + BD: tl.constexpr, + ACTIVATION: tl.constexpr, + IS_RMS_NORM: tl.constexpr, + STORE_RESIDUAL_OUT: tl.constexpr, + HAS_RESIDUAL: tl.constexpr, + HAS_WEIGHT: tl.constexpr, + HAS_BIAS: tl.constexpr, +): + i_t = tl.program_id(0) + + o_d = tl.arange(0, BD) + m_d = o_d < D + + p_x = tl.make_block_ptr(x, (T, D), (D, 1), (i_t * BT, 0), (BT, BD), (1, 0)) + b_x = tl.load(p_x, boundary_check=(0, 1)).to(tl.float32) + if HAS_RESIDUAL: + p_res = tl.make_block_ptr( + residual, (T, D), (D, 1), (i_t * BT, 0), (BT, BD), (1, 0) + ) + b_x += tl.load(p_res, boundary_check=(0, 1)).to(tl.float32) + if STORE_RESIDUAL_OUT: + p_res_out = tl.make_block_ptr( + residual_out, (T, D), (D, 1), (i_t * BT, 0), (BT, BD), (1, 0) + ) + tl.store(p_res_out, b_x.to(p_res_out.dtype.element_ty), boundary_check=(0, 1)) + if not IS_RMS_NORM: + b_mean = tl.sum(b_x, axis=1) / D + p_mean = tl.make_block_ptr(mean, (T,), (1,), (i_t * BT,), (BT,), (0,)) + tl.store(p_mean, b_mean.to(p_mean.dtype.element_ty), boundary_check=(0,)) + b_xbar = tl.where(m_d[None, :], b_x - b_mean[:, None], 0.0) + b_var = tl.sum(b_xbar * b_xbar, axis=1) / D + else: + b_xbar = tl.where(m_d[None, :], b_x, 0.0) + b_var = tl.sum(b_xbar * b_xbar, axis=1) / D + b_rstd = 1 / tl.sqrt(b_var + eps) + + p_rstd = tl.make_block_ptr(rstd, (T,), (1,), (i_t * BT,), (BT,), (0,)) + tl.store(p_rstd, b_rstd.to(p_rstd.dtype.element_ty), boundary_check=(0,)) + + if HAS_WEIGHT: + b_w = tl.load(w + o_d, mask=m_d).to(tl.float32) + if HAS_BIAS: + b_b = tl.load(b + o_d, mask=m_d).to(tl.float32) + b_x_hat = ( + (b_x - b_mean[:, None]) * b_rstd[:, None] + if not IS_RMS_NORM + else b_x * b_rstd[:, None] + ) + b_y = b_x_hat * b_w[None, :] if HAS_WEIGHT else b_x_hat + if HAS_BIAS: + b_y = b_y + b_b[None, :] + + # swish/sigmoid output gate + p_g = tl.make_block_ptr(g, (T, D), (D, 1), (i_t * BT, 0), (BT, BD), (1, 0)) + b_g = tl.load(p_g, boundary_check=(0, 1)).to(tl.float32) + if ACTIVATION == "swish" or ACTIVATION == "silu": + b_y = b_y * b_g * tl.sigmoid(b_g) + elif ACTIVATION == "sigmoid": + b_y = b_y * tl.sigmoid(b_g) + + # Write output + p_y = tl.make_block_ptr(y, (T, D), (D, 1), (i_t * BT, 0), (BT, BD), (1, 0)) + tl.store(p_y, b_y.to(p_y.dtype.element_ty), boundary_check=(0, 1)) + diff --git a/benchmarks/TritonNPUBench/vllm/layer_norm_gated_fwd_kernel1/layer_norm_gated_fwd_kernel1.py b/benchmarks/TritonNPUBench/vllm/layer_norm_gated_fwd_kernel1/layer_norm_gated_fwd_kernel1.py new file mode 100644 index 00000000..c04ad81f --- /dev/null +++ b/benchmarks/TritonNPUBench/vllm/layer_norm_gated_fwd_kernel1/layer_norm_gated_fwd_kernel1.py @@ -0,0 +1,77 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/vllm_operator_cases/newtest_cases/test_layer_norm_gated_fwd_kernel1.py +# Main kernel: layer_norm_gated_fwd_kernel1 +# PT file: layer_norm_gated_fwd_kernel1_v2.pt + +import triton +import triton.language as tl + + +# === layer_norm_gated_fwd_kernel1 === +@triton.jit +def layer_norm_gated_fwd_kernel1( + x, # pointer to the input + g, # pointer to the gate + y, # pointer to the output + w, # pointer to the weights + b, # pointer to the biases + residual, # pointer to the residual + residual_out, # pointer to the residual + mean, # pointer to the mean + rstd, # pointer to the 1/std + eps, # epsilon to avoid division by zero + D: tl.constexpr, # number of columns in x + BD: tl.constexpr, + ACTIVATION: tl.constexpr, + IS_RMS_NORM: tl.constexpr, + STORE_RESIDUAL_OUT: tl.constexpr, + HAS_RESIDUAL: tl.constexpr, + HAS_WEIGHT: tl.constexpr, + HAS_BIAS: tl.constexpr, +): + i_t = tl.program_id(0) + x += i_t * D + y += i_t * D + g += i_t * D + if HAS_RESIDUAL: + residual += i_t * D + if STORE_RESIDUAL_OUT: + residual_out += i_t * D + + o_d = tl.arange(0, BD) + m_d = o_d < D + b_x = tl.load(x + o_d, mask=m_d, other=0.0).to(tl.float32) + if HAS_RESIDUAL: + b_x += tl.load(residual + o_d, mask=m_d, other=0.0).to(tl.float32) + if STORE_RESIDUAL_OUT: + tl.store(residual_out + o_d, b_x, mask=m_d) + if not IS_RMS_NORM: + b_mean = tl.sum(b_x, axis=0) / D + tl.store(mean + i_t, b_mean) + b_xbar = tl.where(m_d, b_x - b_mean, 0.0) + b_var = tl.sum(b_xbar * b_xbar, axis=0) / D + else: + b_xbar = tl.where(m_d, b_x, 0.0) + b_var = tl.sum(b_xbar * b_xbar, axis=0) / D + b_rstd = 1 / tl.sqrt(b_var + eps) + tl.store(rstd + i_t, b_rstd) + + if HAS_WEIGHT: + b_w = tl.load(w + o_d, mask=m_d).to(tl.float32) + if HAS_BIAS: + b_b = tl.load(b + o_d, mask=m_d).to(tl.float32) + b_x_hat = (b_x - b_mean) * b_rstd if not IS_RMS_NORM else b_x * b_rstd + b_y = b_x_hat * b_w if HAS_WEIGHT else b_x_hat + if HAS_BIAS: + b_y = b_y + b_b + + # swish/sigmoid output gate + b_g = tl.load(g + o_d, mask=m_d, other=0.0).to(tl.float32) + if ACTIVATION == "swish" or ACTIVATION == "silu": + b_y = b_y * b_g * tl.sigmoid(b_g) + elif ACTIVATION == "sigmoid": + b_y = b_y * tl.sigmoid(b_g) + + # Write output + tl.store(y + o_d, b_y, mask=m_d) + diff --git a/benchmarks/TritonNPUBench/vllm/matmul_kernel_persistent/matmul_kernel_persistent.py b/benchmarks/TritonNPUBench/vllm/matmul_kernel_persistent/matmul_kernel_persistent.py new file mode 100644 index 00000000..a2e1d779 --- /dev/null +++ b/benchmarks/TritonNPUBench/vllm/matmul_kernel_persistent/matmul_kernel_persistent.py @@ -0,0 +1,143 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/vllm_operator_cases/newtest_cases/test_matmul_kernel_persistent.py +# Main kernel: matmul_kernel_persistent +# PT file: matmul_kernel_persistent_cpu_v2.pt + +import triton +import triton.language as tl +from collections.abc import Callable +from typing import Any + +def _matmul_launch_metadata( + grid: Callable[..., Any], kernel: Any, args: dict[str, Any] +) -> dict[str, Any]: + ret = {} + m, n, k = args["M"], args["N"], args["K"] + ret["name"] = f"{kernel.name} [M={m}, N={n}, K={k}]" + if "tiles_per_update" in args: + ret["name"] = ( + f"{kernel.name} [M={m}, N={n}, K={k}, " + f"tiles_per_update={args['tiles_per_update']:02}]" + ) + if "c_ptr" in args: + bytes_per_elem = args["c_ptr"].element_size() + else: + bytes_per_elem = 1 if args["FP8_OUTPUT"] else 2 + ret[f"flops{bytes_per_elem * 8}"] = 2.0 * m * n * k + ret["bytes"] = bytes_per_elem * (m * k + n * k + m * n) + return ret + + +@triton.jit +def _compute_pid(tile_id, num_pid_in_group, num_pid_m, GROUP_SIZE_M, NUM_SMS): + group_id = tile_id // num_pid_in_group + first_pid_m = group_id * GROUP_SIZE_M + group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M) + pid_m = first_pid_m + (tile_id % group_size_m) + pid_n = (tile_id % num_pid_in_group) // group_size_m + return pid_m, pid_n + +# === matmul_kernel_persistent === +@triton.jit(launch_metadata=_matmul_launch_metadata) +def matmul_kernel_persistent( + a_ptr, + b_ptr, + c_ptr, # + bias_ptr, + M, + N, + K, # + stride_am, + stride_ak, + stride_bk, + stride_bn, + stride_cm, + stride_cn, + BLOCK_SIZE_M: tl.constexpr, # + BLOCK_SIZE_N: tl.constexpr, # + BLOCK_SIZE_K: tl.constexpr, # + GROUP_SIZE_M: tl.constexpr, # + NUM_SMS: tl.constexpr, # + A_LARGE: tl.constexpr, + B_LARGE: tl.constexpr, + C_LARGE: tl.constexpr, + HAS_BIAS: tl.constexpr, +): + start_pid = tl.program_id(axis=0) + num_pid_m = tl.cdiv(M, BLOCK_SIZE_M) + num_pid_n = tl.cdiv(N, BLOCK_SIZE_N) + k_tiles = tl.cdiv(K, BLOCK_SIZE_K) + num_tiles = num_pid_m * num_pid_n + + tile_id_c = start_pid - NUM_SMS + + offs_k_for_mask = tl.arange(0, BLOCK_SIZE_K) + num_pid_in_group = GROUP_SIZE_M * num_pid_n + + for tile_id in tl.range(start_pid, num_tiles, NUM_SMS, flatten=True): + pid_m, pid_n = _compute_pid( + tile_id, num_pid_in_group, num_pid_m, GROUP_SIZE_M, NUM_SMS + ) + start_m = pid_m * BLOCK_SIZE_M + start_n = pid_n * BLOCK_SIZE_N + offs_am = start_m + tl.arange(0, BLOCK_SIZE_M) + offs_bn = start_n + tl.arange(0, BLOCK_SIZE_N) + if A_LARGE: + offs_am = offs_am.to(tl.int64) + if B_LARGE: + offs_bn = offs_bn.to(tl.int64) + offs_am = tl.where(offs_am < M, offs_am, 0) + offs_bn = tl.where(offs_bn < N, offs_bn, 0) + offs_am = tl.max_contiguous(tl.multiple_of(offs_am, BLOCK_SIZE_M), BLOCK_SIZE_M) + offs_bn = tl.max_contiguous(tl.multiple_of(offs_bn, BLOCK_SIZE_N), BLOCK_SIZE_N) + + accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32) + for ki in range(k_tiles): + if A_LARGE or B_LARGE: + offs_k = ki * BLOCK_SIZE_K + tl.arange(0, BLOCK_SIZE_K).to(tl.int64) + else: + offs_k = ki * BLOCK_SIZE_K + tl.arange(0, BLOCK_SIZE_K) + a_ptrs = a_ptr + ( + offs_am[:, None] * stride_am + offs_k[None, :] * stride_ak + ) + b_ptrs = b_ptr + ( + offs_k[:, None] * stride_bk + offs_bn[None, :] * stride_bn + ) + + a = tl.load( + a_ptrs, mask=offs_k_for_mask[None, :] < K - ki * BLOCK_SIZE_K, other=0.0 + ) + b = tl.load( + b_ptrs, mask=offs_k_for_mask[:, None] < K - ki * BLOCK_SIZE_K, other=0.0 + ) + accumulator = tl.dot(a, b, accumulator) + + tile_id_c += NUM_SMS + pid_m, pid_n = _compute_pid( + tile_id_c, num_pid_in_group, num_pid_m, GROUP_SIZE_M, NUM_SMS + ) + offs_cm = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) + offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) + if C_LARGE: + offs_cm = offs_cm.to(tl.int64) + offs_cn = offs_cn.to(tl.int64) + c_ptrs = c_ptr + stride_cm * offs_cm[:, None] + stride_cn * offs_cn[None, :] + c_mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < N) + if HAS_BIAS: + bias_ptrs = bias_ptr + offs_cn + bias = tl.load(bias_ptrs, mask=offs_cn < N, other=0.0).to(tl.float32) + accumulator += bias + c = accumulator.to(c_ptr.dtype.element_ty) + tl.store(c_ptrs, c, mask=c_mask) + + +# === _compute_pid === +@triton.jit +def _compute_pid(tile_id, num_pid_in_group, num_pid_m, GROUP_SIZE_M, NUM_SMS): + group_id = tile_id // num_pid_in_group + first_pid_m = group_id * GROUP_SIZE_M + group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M) + pid_m = first_pid_m + (tile_id % group_size_m) + pid_n = (tile_id % num_pid_in_group) // group_size_m + return pid_m, pid_n + diff --git a/benchmarks/TritonNPUBench/vllm/mean_kernel/mean_kernel.py b/benchmarks/TritonNPUBench/vllm/mean_kernel/mean_kernel.py new file mode 100644 index 00000000..61f22727 --- /dev/null +++ b/benchmarks/TritonNPUBench/vllm/mean_kernel/mean_kernel.py @@ -0,0 +1,59 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/vllm_operator_cases/newtest_cases/test_mean_kernel.py +# Main kernel: mean_kernel +# PT file: mean_kernel_v2.pt + +import triton +import triton.language as tl + + +# === mean_kernel === +@triton.jit +def mean_kernel( + input_ptr, + output_ptr, + input_stride0, + input_stride1, + input_stride2, + output_stride0, + output_stride1, + M, # size before reduction dim + N, # size of reduction dim + K, # size after reduction dim + BLOCK_SIZE: tl.constexpr, +): + """ + Kernel for computing mean along a single dimension. + Input is viewed as (M, N, K) where N is the dimension being reduced. + """ + # Program ID gives us which output element we're computing + pid = tl.program_id(0) + + # Compute output indices + m_idx = pid // K + k_idx = pid % K + + # Bounds check + if m_idx >= M or k_idx >= K: + return + + # Accumulate sum across reduction dimension + acc = 0.0 + for n_start in range(0, N, BLOCK_SIZE): + n_offsets = n_start + tl.arange(0, BLOCK_SIZE) + mask = n_offsets < N + + # Calculate input indices + input_idx = ( + m_idx * input_stride0 + n_offsets * input_stride1 + k_idx * input_stride2 + ) + + # Load and accumulate + vals = tl.load(input_ptr + input_idx, mask=mask, other=0.0) + acc += tl.sum(vals) + + # Compute mean and store + mean_val = acc / N + output_idx = m_idx * output_stride0 + k_idx * output_stride1 + tl.store(output_ptr + output_idx, mean_val) + diff --git a/benchmarks/TritonNPUBench/vllm/merge_16x16_to_32x32_inverse_kernel/merge_16x16_to_32x32_inverse_kernel.py b/benchmarks/TritonNPUBench/vllm/merge_16x16_to_32x32_inverse_kernel/merge_16x16_to_32x32_inverse_kernel.py new file mode 100644 index 00000000..2558b7bd --- /dev/null +++ b/benchmarks/TritonNPUBench/vllm/merge_16x16_to_32x32_inverse_kernel/merge_16x16_to_32x32_inverse_kernel.py @@ -0,0 +1,125 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/vllm_operator_cases/newtest_cases/test_merge_16x16_to_32x32_inverse_kernel.py +# Main kernel: merge_16x16_to_32x32_inverse_kernel +# PT file: test_merge_16x16_to_32x32_inverse_kernel_v3.pt + +import triton +import triton.language as tl + + +# === merge_16x16_to_32x32_inverse_kernel === +@triton.jit(do_not_specialize=["T"]) +def merge_16x16_to_32x32_inverse_kernel( + A, + Ai, + cu_seqlens, + chunk_indices, + T, + H: tl.constexpr, + BT: tl.constexpr, + USE_TMA: tl.constexpr, + IS_VARLEN: tl.constexpr, + DOT_PRECISION: tl.constexpr, +): + i_t, i_bh = tl.program_id(0), tl.program_id(1) + i_b, i_h = i_bh // H, i_bh % H + if IS_VARLEN: + i_n, i_t = ( + tl.load(chunk_indices + i_t * 2).to(tl.int32), + tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), + ) + bos, eos = ( + tl.load(cu_seqlens + i_n).to(tl.int32), + tl.load(cu_seqlens + i_n + 1).to(tl.int32), + ) + T = eos - bos + else: + bos, eos = i_b * T, i_b * T + T + + o_i = tl.arange(0, 16) + m_A = o_i[:, None] > o_i[None, :] + m_I = o_i[:, None] == o_i[None, :] + A += (bos * H + i_h) * BT + Ai += (bos * H + i_h) * BT + + if not USE_TMA: + p_A_11 = tl.make_block_ptr( + A, (T, BT), (H * BT, 1), (i_t * BT, 0), (16, 16), (1, 0) + ) + p_A_22 = tl.make_block_ptr( + A, (T, BT), (H * BT, 1), (i_t * BT + 16, 16), (16, 16), (1, 0) + ) + b_Ai_11 = tl.load(p_A_11, boundary_check=(0, 1)).to(tl.float32) + b_Ai_22 = tl.load(p_A_22, boundary_check=(0, 1)).to(tl.float32) + else: + desc = make_tensor_descriptor(A, [T, BT], [H * BT, 1], [16, 16]) + desc_o = make_tensor_descriptor(Ai, [T, BT], [H * BT, 1], [16, 16]) + b_Ai_11 = desc.load([i_t * BT + 0, 0]).to(tl.float32) + b_Ai_22 = desc.load([i_t * BT + 16, 16]).to(tl.float32) + + # [16, 16] + b_Ai_11 = -tl.where(m_A, b_Ai_11, 0) + b_Ai_22 = -tl.where(m_A, b_Ai_22, 0) + + for i in range(2, min(16, T - i_t * BT)): + b_a_11 = -tl.load(A + (i_t * BT + i) * H * BT + o_i) + b_a_11 += tl.sum(b_a_11[:, None] * b_Ai_11, 0) + b_Ai_11 = tl.where((o_i == i)[:, None], b_a_11, b_Ai_11) + for i in range(16 + 2, min(32, T - i_t * BT)): + b_a_22 = -tl.load(A + (i_t * BT + i) * H * BT + o_i + 16) + b_a_22 += tl.sum(b_a_22[:, None] * b_Ai_22, 0) + b_Ai_22 = tl.where((o_i == i - 16)[:, None], b_a_22, b_Ai_22) + + b_Ai_11 += m_I + b_Ai_22 += m_I + + if not USE_TMA: + p_A_21 = tl.make_block_ptr( + A, (T, BT), (H * BT, 1), (i_t * BT + 16, 0), (16, 16), (1, 0) + ) + b_A_21 = tl.load(p_A_21, boundary_check=(0, 1)).to(tl.float32) + else: + b_A_21 = desc.load([i_t * BT + 16, 0]).to(tl.float32) + + b_Ai_21 = -tl.dot( + tl.dot(b_Ai_22, b_A_21, input_precision=DOT_PRECISION), + b_Ai_11, + input_precision=DOT_PRECISION, + ) + + if not USE_TMA: + p_Ai_11 = tl.make_block_ptr( + Ai, (T, BT), (H * BT, 1), (i_t * BT, 0), (16, 16), (1, 0) + ) + p_Ai_21 = tl.make_block_ptr( + Ai, (T, BT), (H * BT, 1), (i_t * BT + 16, 0), (16, 16), (1, 0) + ) + p_Ai_22 = tl.make_block_ptr( + Ai, (T, BT), (H * BT, 1), (i_t * BT + 16, 16), (16, 16), (1, 0) + ) + tl.store( + p_Ai_11, + b_Ai_11.to(p_Ai_11.dtype.element_ty, fp_downcast_rounding="rtne"), + boundary_check=(0, 1), + ) + tl.store( + p_Ai_22, + b_Ai_22.to(p_Ai_22.dtype.element_ty, fp_downcast_rounding="rtne"), + boundary_check=(0, 1), + ) + tl.store( + p_Ai_21, + b_Ai_21.to(p_Ai_21.dtype.element_ty, fp_downcast_rounding="rtne"), + boundary_check=(0, 1), + ) + else: + desc_o.store( + [i_t * BT + 0, 0], b_Ai_11.to(desc_o.dtype, fp_downcast_rounding="rtne") + ) + desc_o.store( + [i_t * BT + 16, 0], b_Ai_21.to(desc_o.dtype, fp_downcast_rounding="rtne") + ) + desc_o.store( + [i_t * BT + 16, 16], b_Ai_22.to(desc_o.dtype, fp_downcast_rounding="rtne") + ) + diff --git a/benchmarks/TritonNPUBench/vllm/merge_16x16_to_64x64_inverse_kernel/merge_16x16_to_64x64_inverse_kernel.py b/benchmarks/TritonNPUBench/vllm/merge_16x16_to_64x64_inverse_kernel/merge_16x16_to_64x64_inverse_kernel.py new file mode 100644 index 00000000..f568f351 --- /dev/null +++ b/benchmarks/TritonNPUBench/vllm/merge_16x16_to_64x64_inverse_kernel/merge_16x16_to_64x64_inverse_kernel.py @@ -0,0 +1,278 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/vllm_operator_cases/newtest_cases/test_merge_16x16_to_64x64_inverse_kernel.py +# Main kernel: merge_16x16_to_64x64_inverse_kernel +# PT file: test_merge_16x16_to_64x64_inverse_kernel_v2.pt + +import triton +import triton.language as tl + + +# === merge_16x16_to_64x64_inverse_kernel === +@triton.jit(do_not_specialize=["T"]) +def merge_16x16_to_64x64_inverse_kernel( + A, + Ai, + cu_seqlens, + chunk_indices, + T, + H: tl.constexpr, + BT: tl.constexpr, + USE_TMA: tl.constexpr, + IS_VARLEN: tl.constexpr, + DOT_PRECISION: tl.constexpr, +): + i_t, i_bh = tl.program_id(0), tl.program_id(1) + i_b, i_h = i_bh // H, i_bh % H + if IS_VARLEN: + i_n, i_t = ( + tl.load(chunk_indices + i_t * 2).to(tl.int32), + tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), + ) + bos, eos = ( + tl.load(cu_seqlens + i_n).to(tl.int32), + tl.load(cu_seqlens + i_n + 1).to(tl.int32), + ) + T = eos - bos + else: + bos, eos = i_b * T, i_b * T + T + + o_i = tl.arange(0, 16) + m_A = o_i[:, None] > o_i[None, :] + m_I = o_i[:, None] == o_i[None, :] + A += (bos * H + i_h) * BT + Ai += (bos * H + i_h) * BT + + if not USE_TMA: + p_A_11 = tl.make_block_ptr( + A, (T, BT), (H * BT, 1), (i_t * BT, 0), (16, 16), (1, 0) + ) + p_A_22 = tl.make_block_ptr( + A, (T, BT), (H * BT, 1), (i_t * BT + 16, 16), (16, 16), (1, 0) + ) + p_A_33 = tl.make_block_ptr( + A, (T, BT), (H * BT, 1), (i_t * BT + 32, 32), (16, 16), (1, 0) + ) + p_A_44 = tl.make_block_ptr( + A, (T, BT), (H * BT, 1), (i_t * BT + 48, 48), (16, 16), (1, 0) + ) + b_Ai_11 = tl.load(p_A_11, boundary_check=(0, 1)).to(tl.float32) + b_Ai_22 = tl.load(p_A_22, boundary_check=(0, 1)).to(tl.float32) + b_Ai_33 = tl.load(p_A_33, boundary_check=(0, 1)).to(tl.float32) + b_Ai_44 = tl.load(p_A_44, boundary_check=(0, 1)).to(tl.float32) + else: + desc = make_tensor_descriptor(A, [T, BT], [H * BT, 1], [16, 16]) + desc_o = make_tensor_descriptor(Ai, [T, BT], [H * BT, 1], [16, 16]) + b_Ai_11 = desc.load([i_t * BT + 0, 0]).to(tl.float32) + b_Ai_22 = desc.load([i_t * BT + 16, 16]).to(tl.float32) + b_Ai_33 = desc.load([i_t * BT + 32, 32]).to(tl.float32) + b_Ai_44 = desc.load([i_t * BT + 48, 48]).to(tl.float32) + + # [16, 16] + b_Ai_11 = -tl.where(m_A, b_Ai_11, 0) + b_Ai_22 = -tl.where(m_A, b_Ai_22, 0) + b_Ai_33 = -tl.where(m_A, b_Ai_33, 0) + b_Ai_44 = -tl.where(m_A, b_Ai_44, 0) + + for i in range(2, min(16, T - i_t * BT)): + b_a_11 = -tl.load(A + (i_t * BT + i) * H * BT + o_i) + b_a_11 += tl.sum(b_a_11[:, None] * b_Ai_11, 0) + b_Ai_11 = tl.where((o_i == i)[:, None], b_a_11, b_Ai_11) + for i in range(16 + 2, min(32, T - i_t * BT)): + b_a_22 = -tl.load(A + (i_t * BT + i) * H * BT + o_i + 16) + b_a_22 += tl.sum(b_a_22[:, None] * b_Ai_22, 0) + b_Ai_22 = tl.where((o_i == i - 16)[:, None], b_a_22, b_Ai_22) + for i in range(32 + 2, min(48, T - i_t * BT)): + b_a_33 = -tl.load(A + (i_t * BT + i) * H * BT + o_i + 32) + b_a_33 += tl.sum(b_a_33[:, None] * b_Ai_33, 0) + b_Ai_33 = tl.where((o_i == i - 32)[:, None], b_a_33, b_Ai_33) + for i in range(48 + 2, min(64, T - i_t * BT)): + b_a_44 = -tl.load(A + (i_t * BT + i) * H * BT + o_i + 48) + b_a_44 += tl.sum(b_a_44[:, None] * b_Ai_44, 0) + b_Ai_44 = tl.where((o_i == i - 48)[:, None], b_a_44, b_Ai_44) + b_Ai_11 += m_I + b_Ai_22 += m_I + b_Ai_33 += m_I + b_Ai_44 += m_I + + if not USE_TMA: + p_A_21 = tl.make_block_ptr( + A, (T, BT), (H * BT, 1), (i_t * BT + 16, 0), (16, 16), (1, 0) + ) + p_A_31 = tl.make_block_ptr( + A, (T, BT), (H * BT, 1), (i_t * BT + 32, 0), (16, 16), (1, 0) + ) + p_A_32 = tl.make_block_ptr( + A, (T, BT), (H * BT, 1), (i_t * BT + 32, 16), (16, 16), (1, 0) + ) + p_A_41 = tl.make_block_ptr( + A, (T, BT), (H * BT, 1), (i_t * BT + 48, 0), (16, 16), (1, 0) + ) + p_A_42 = tl.make_block_ptr( + A, (T, BT), (H * BT, 1), (i_t * BT + 48, 16), (16, 16), (1, 0) + ) + p_A_43 = tl.make_block_ptr( + A, (T, BT), (H * BT, 1), (i_t * BT + 48, 32), (16, 16), (1, 0) + ) + b_A_21 = tl.load(p_A_21, boundary_check=(0, 1)).to(tl.float32) + b_A_31 = tl.load(p_A_31, boundary_check=(0, 1)).to(tl.float32) + b_A_32 = tl.load(p_A_32, boundary_check=(0, 1)).to(tl.float32) + b_A_41 = tl.load(p_A_41, boundary_check=(0, 1)).to(tl.float32) + b_A_42 = tl.load(p_A_42, boundary_check=(0, 1)).to(tl.float32) + b_A_43 = tl.load(p_A_43, boundary_check=(0, 1)).to(tl.float32) + else: + b_A_21 = desc.load([i_t * BT + 16, 0]).to(tl.float32) + b_A_31 = desc.load([i_t * BT + 32, 0]).to(tl.float32) + b_A_32 = desc.load([i_t * BT + 32, 16]).to(tl.float32) + b_A_41 = desc.load([i_t * BT + 48, 0]).to(tl.float32) + b_A_42 = desc.load([i_t * BT + 48, 16]).to(tl.float32) + b_A_43 = desc.load([i_t * BT + 48, 32]).to(tl.float32) + + b_Ai_21 = -tl.dot( + tl.dot(b_Ai_22, b_A_21, input_precision=DOT_PRECISION), + b_Ai_11, + input_precision=DOT_PRECISION, + ) + b_Ai_32 = -tl.dot( + tl.dot(b_Ai_33, b_A_32, input_precision=DOT_PRECISION), + b_Ai_22, + input_precision=DOT_PRECISION, + ) + b_Ai_43 = -tl.dot( + tl.dot(b_Ai_44, b_A_43, input_precision=DOT_PRECISION), + b_Ai_33, + input_precision=DOT_PRECISION, + ) + + b_Ai_31 = -tl.dot( + b_Ai_33, + tl.dot(b_A_31, b_Ai_11, input_precision=DOT_PRECISION) + + tl.dot(b_A_32, b_Ai_21, input_precision=DOT_PRECISION), + input_precision=DOT_PRECISION, + ) + b_Ai_42 = -tl.dot( + b_Ai_44, + tl.dot(b_A_42, b_Ai_22, input_precision=DOT_PRECISION) + + tl.dot(b_A_43, b_Ai_32, input_precision=DOT_PRECISION), + input_precision=DOT_PRECISION, + ) + b_Ai_41 = -tl.dot( + b_Ai_44, + tl.dot(b_A_41, b_Ai_11, input_precision=DOT_PRECISION) + + tl.dot(b_A_42, b_Ai_21, input_precision=DOT_PRECISION) + + tl.dot(b_A_43, b_Ai_31, input_precision=DOT_PRECISION), + input_precision=DOT_PRECISION, + ) + + if not USE_TMA: + p_Ai_11 = tl.make_block_ptr( + Ai, (T, BT), (H * BT, 1), (i_t * BT, 0), (16, 16), (1, 0) + ) + p_Ai_22 = tl.make_block_ptr( + Ai, (T, BT), (H * BT, 1), (i_t * BT + 16, 16), (16, 16), (1, 0) + ) + p_Ai_33 = tl.make_block_ptr( + Ai, (T, BT), (H * BT, 1), (i_t * BT + 32, 32), (16, 16), (1, 0) + ) + p_Ai_44 = tl.make_block_ptr( + Ai, (T, BT), (H * BT, 1), (i_t * BT + 48, 48), (16, 16), (1, 0) + ) + p_Ai_21 = tl.make_block_ptr( + Ai, (T, BT), (H * BT, 1), (i_t * BT + 16, 0), (16, 16), (1, 0) + ) + p_Ai_31 = tl.make_block_ptr( + Ai, (T, BT), (H * BT, 1), (i_t * BT + 32, 0), (16, 16), (1, 0) + ) + p_Ai_32 = tl.make_block_ptr( + Ai, (T, BT), (H * BT, 1), (i_t * BT + 32, 16), (16, 16), (1, 0) + ) + p_Ai_41 = tl.make_block_ptr( + Ai, (T, BT), (H * BT, 1), (i_t * BT + 48, 0), (16, 16), (1, 0) + ) + p_Ai_42 = tl.make_block_ptr( + Ai, (T, BT), (H * BT, 1), (i_t * BT + 48, 16), (16, 16), (1, 0) + ) + p_Ai_43 = tl.make_block_ptr( + Ai, (T, BT), (H * BT, 1), (i_t * BT + 48, 32), (16, 16), (1, 0) + ) + tl.store( + p_Ai_11, + b_Ai_11.to(p_Ai_11.dtype.element_ty, fp_downcast_rounding="rtne"), + boundary_check=(0, 1), + ) + tl.store( + p_Ai_22, + b_Ai_22.to(p_Ai_22.dtype.element_ty, fp_downcast_rounding="rtne"), + boundary_check=(0, 1), + ) + tl.store( + p_Ai_33, + b_Ai_33.to(p_Ai_33.dtype.element_ty, fp_downcast_rounding="rtne"), + boundary_check=(0, 1), + ) + tl.store( + p_Ai_44, + b_Ai_44.to(p_Ai_44.dtype.element_ty, fp_downcast_rounding="rtne"), + boundary_check=(0, 1), + ) + tl.store( + p_Ai_21, + b_Ai_21.to(p_Ai_21.dtype.element_ty, fp_downcast_rounding="rtne"), + boundary_check=(0, 1), + ) + tl.store( + p_Ai_31, + b_Ai_31.to(p_Ai_31.dtype.element_ty, fp_downcast_rounding="rtne"), + boundary_check=(0, 1), + ) + tl.store( + p_Ai_32, + b_Ai_32.to(p_Ai_32.dtype.element_ty, fp_downcast_rounding="rtne"), + boundary_check=(0, 1), + ) + tl.store( + p_Ai_41, + b_Ai_41.to(p_Ai_41.dtype.element_ty, fp_downcast_rounding="rtne"), + boundary_check=(0, 1), + ) + tl.store( + p_Ai_42, + b_Ai_42.to(p_Ai_42.dtype.element_ty, fp_downcast_rounding="rtne"), + boundary_check=(0, 1), + ) + tl.store( + p_Ai_43, + b_Ai_43.to(p_Ai_43.dtype.element_ty, fp_downcast_rounding="rtne"), + boundary_check=(0, 1), + ) + else: + desc_o.store( + [i_t * BT + 0, 0], b_Ai_11.to(desc_o.dtype, fp_downcast_rounding="rtne") + ) + desc_o.store( + [i_t * BT + 16, 16], b_Ai_22.to(desc_o.dtype, fp_downcast_rounding="rtne") + ) + desc_o.store( + [i_t * BT + 32, 32], b_Ai_33.to(desc_o.dtype, fp_downcast_rounding="rtne") + ) + desc_o.store( + [i_t * BT + 48, 48], b_Ai_44.to(desc_o.dtype, fp_downcast_rounding="rtne") + ) + desc_o.store( + [i_t * BT + 16, 0], b_Ai_21.to(desc_o.dtype, fp_downcast_rounding="rtne") + ) + desc_o.store( + [i_t * BT + 32, 0], b_Ai_31.to(desc_o.dtype, fp_downcast_rounding="rtne") + ) + desc_o.store( + [i_t * BT + 32, 16], b_Ai_32.to(desc_o.dtype, fp_downcast_rounding="rtne") + ) + desc_o.store( + [i_t * BT + 48, 0], b_Ai_41.to(desc_o.dtype, fp_downcast_rounding="rtne") + ) + desc_o.store( + [i_t * BT + 48, 16], b_Ai_42.to(desc_o.dtype, fp_downcast_rounding="rtne") + ) + desc_o.store( + [i_t * BT + 48, 32], b_Ai_43.to(desc_o.dtype, fp_downcast_rounding="rtne") + ) + diff --git a/benchmarks/TritonNPUBench/vllm/merge_attn_states_kernel/merge_attn_states_kernel.py b/benchmarks/TritonNPUBench/vllm/merge_attn_states_kernel/merge_attn_states_kernel.py new file mode 100644 index 00000000..d4dcb41d --- /dev/null +++ b/benchmarks/TritonNPUBench/vllm/merge_attn_states_kernel/merge_attn_states_kernel.py @@ -0,0 +1,84 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/vllm_operator_cases/newtest_cases/test_merge_attn_states_kernel.py +# Main kernel: merge_attn_states_kernel +# PT file: merge_attn_states_test_v2.pt + +import triton +import triton.language as tl + + +# === merge_attn_states_kernel === +@triton.jit +def merge_attn_states_kernel( + output, # [NUM_TOKENS, NUM_HEADS, HEAD_SIZE] + output_lse, # [NUM_HEADS, NUM_TOKENS] + prefix_output, # [NUM_TOKENS, NUM_HEADS, HEAD_SIZE] + prefix_lse, # [NUM_HEADS, NUM_TOKENS] + suffix_output, # [NUM_TOKENS, NUM_HEADS, HEAD_SIZE] + suffix_lse, # [NUM_HEADS, NUM_TOKENS] + prefix_head_stride, + output_head_stride, + HEAD_SIZE: tl.constexpr, + PADDED_HEAD_SIZE: tl.constexpr, + OUTPUT_LSE: tl.constexpr, +): + token_idx = tl.program_id(0) + num_tokens = tl.num_programs(0) + head_idx = tl.program_id(1) + num_heads = tl.num_programs(1) + + p_lse = tl.load(prefix_lse + head_idx * num_tokens + token_idx) + s_lse = tl.load(suffix_lse + head_idx * num_tokens + token_idx) + + # FA2 and FA3 have different behavior for when the sum-exp is 0, this namely + # arises with 0 len seqlens. FA3 returns -inf here while FA2 returns inf. + # If we see an inf assume FA2 and convert inf to -inf for consistency + # and correctness. Inf generally doesn't make sense in this context outside + # of undefined-behavior/FA2-case, so I think this a safe assumption. + p_lse = float("-inf") if p_lse == float("inf") else p_lse + s_lse = float("-inf") if s_lse == float("inf") else s_lse + + max_lse = tl.maximum(p_lse, s_lse) + p_lse = p_lse - max_lse + s_lse = s_lse - max_lse + # Will reuse precomputed Exp values for scale factor computation. + p_se = tl.exp(p_lse) + s_se = tl.exp(s_lse) + out_se = p_se + s_se + + if OUTPUT_LSE: + out_lse = tl.log(out_se) + max_lse + tl.store(output_lse + head_idx * num_tokens + token_idx, out_lse) + + head_arange = tl.arange(0, PADDED_HEAD_SIZE) + head_mask = head_arange < HEAD_SIZE + p_out = tl.load( + prefix_output + + token_idx * num_heads * prefix_head_stride + + head_idx * prefix_head_stride + + head_arange, + mask=head_mask, + ) + s_out = tl.load( + suffix_output + + token_idx * num_heads * prefix_head_stride + + head_idx * prefix_head_stride + + head_arange, + mask=head_mask, + ) + + # NOTE(woosuk): Be careful with the numerical stability. + # We should compute the scale first, and then multiply it with the output. + # Do not multiply the output with tl.exp(p_lse) or tl.exp(s_lse) directly. + p_scale = p_se / out_se + s_scale = s_se / out_se + out = p_out * p_scale + s_out * s_scale + tl.store( + output + + token_idx * num_heads * output_head_stride + + head_idx * output_head_stride + + head_arange, + out, + mask=head_mask, + ) + diff --git a/benchmarks/TritonNPUBench/vllm/pack_bitmatrix/pack_bitmatrix.py b/benchmarks/TritonNPUBench/vllm/pack_bitmatrix/pack_bitmatrix.py new file mode 100644 index 00000000..96cf9920 --- /dev/null +++ b/benchmarks/TritonNPUBench/vllm/pack_bitmatrix/pack_bitmatrix.py @@ -0,0 +1,53 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/vllm_operator_cases/newtest_cases/test_pack_bitmatrix.py +# Main kernel: pack_bitmatrix +# PT file: pack_bitmatrix_new.pt + +import triton +import triton.language as tl + +@triton.jit +def _or_combine(x, y): + return x | y + + +# === pack_bitmatrix === +@triton.jit +def pack_bitmatrix( + bitmatrix, + topk_ids, + n_rows, # n_rows in bitmatrix / topk_ids + bm_cols: tl.constexpr, # n int32_t bitpacks in bitmatrix + n_expts_act, # num_topk + BLOCK_SIZE_M: tl.constexpr, + BLOCK_SIZE_K: tl.constexpr, +): + """ + Packs topk_ids into a bitmatrix. + code reference: + https://github.com/triton-lang/triton/blob/dd1bbc52b34d202dfe5ffea1e04fb16166c5c04e/python/triton_kernels/bench/distributed.py#L264 + """ + pid_m = tl.program_id(0) + offsets_m = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) + offsets_k = tl.arange(0, BLOCK_SIZE_K) + offsets = offsets_m[:, None] * n_expts_act + offsets_k[None, :] + mask = (offsets_m < n_rows)[:, None] & (offsets_k < n_expts_act)[None, :] + indices = tl.load(topk_ids + offsets, mask=mask, other=-1) + div = indices // 32 + rem = indices % 32 + one = tl.cast(1, tl.uint32) + + # Iterate through all the relevant bitmatrix columns. + for i in range(bm_cols): + # When BLOCK_SIZE_K=32, offs is just the column index. + offs = tl.arange(0, BLOCK_SIZE_K // 32) + i * (BLOCK_SIZE_K // 32) + # All topks that need to go into this column has the correct bit set. + # Other bits are 0. x is a 2D tensor. + x = tl.where( + div[:, :, None] == offs[None, None, :], (one << rem)[:, :, None], 0 + ) + # Reduce x to get a single int32_t bitpack. + y = tl.reduce(x, axis=1, combine_fn=_or_combine) + bitmatrix_ptrs = bitmatrix + offsets_m[:, None] * bm_cols + offs[None, :] + tl.store(bitmatrix_ptrs, y, mask=offsets_m[:, None] < n_rows) + diff --git a/benchmarks/TritonNPUBench/vllm/recompute_w_u_fwd_kernel/recompute_w_u_fwd_kernel.py b/benchmarks/TritonNPUBench/vllm/recompute_w_u_fwd_kernel/recompute_w_u_fwd_kernel.py new file mode 100644 index 00000000..bfd92588 --- /dev/null +++ b/benchmarks/TritonNPUBench/vllm/recompute_w_u_fwd_kernel/recompute_w_u_fwd_kernel.py @@ -0,0 +1,100 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/vllm_operator_cases/newtest_cases/test_recompute_w_u_fwd_kernel.py +# Main kernel: recompute_w_u_fwd_kernel +# PT file: test_recompute_w_u_fwd_kernel_v2.pt + +import triton +import triton.language as tl + + +# === recompute_w_u_fwd_kernel === +@triton.jit(do_not_specialize=["T"]) +def recompute_w_u_fwd_kernel( + k, + v, + beta, + w, + u, + A, + g, + cu_seqlens, + chunk_indices, + T, + H: tl.constexpr, + Hg: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_t, i_bh = tl.program_id(0), tl.program_id(1) + i_b, i_h = i_bh // H, i_bh % H + if IS_VARLEN: + i_n, i_t = ( + tl.load(chunk_indices + i_t * 2).to(tl.int32), + tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), + ) + bos, eos = ( + tl.load(cu_seqlens + i_n).to(tl.int32), + tl.load(cu_seqlens + i_n + 1).to(tl.int32), + ) + T = eos - bos + else: + bos, eos = i_b * T, i_b * T + T + p_beta = tl.make_block_ptr( + beta + bos * H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,) + ) + p_g = tl.make_block_ptr(g + (bos * H + i_h), (T,), (H,), (i_t * BT,), (BT,), (0,)) + p_A = tl.make_block_ptr( + A + (bos * H + i_h) * BT, (T, BT), (H * BT, 1), (i_t * BT, 0), (BT, BT), (1, 0) + ) + b_beta = tl.load(p_beta, boundary_check=(0,)) + b_A = tl.load(p_A, boundary_check=(0, 1)) + b_g = tl.exp(tl.load(p_g, boundary_check=(0,))) + + for i_v in range(tl.cdiv(V, BV)): + p_v = tl.make_block_ptr( + v + (bos * H + i_h) * V, + (T, V), + (H * V, 1), + (i_t * BT, i_v * BV), + (BT, BV), + (1, 0), + ) + p_u = tl.make_block_ptr( + u + (bos * H + i_h) * V, + (T, V), + (H * V, 1), + (i_t * BT, i_v * BV), + (BT, BV), + (1, 0), + ) + b_v = tl.load(p_v, boundary_check=(0, 1)) + b_vb = (b_v * b_beta[:, None]).to(b_v.dtype) + b_u = tl.dot(b_A, b_vb, allow_tf32=False) + tl.store(p_u, b_u.to(p_u.dtype.element_ty), boundary_check=(0, 1)) + + for i_k in range(tl.cdiv(K, BK)): + p_k = tl.make_block_ptr( + k + (bos * Hg + i_h // (H // Hg)) * K, + (T, K), + (Hg * K, 1), + (i_t * BT, i_k * BK), + (BT, BK), + (1, 0), + ) + p_w = tl.make_block_ptr( + w + (bos * H + i_h) * K, + (T, K), + (H * K, 1), + (i_t * BT, i_k * BK), + (BT, BK), + (1, 0), + ) + b_k = tl.load(p_k, boundary_check=(0, 1)) + b_kb = (b_k * b_beta[:, None] * b_g[:, None]).to(b_k.dtype) + b_w = tl.dot(b_A, b_kb) + tl.store(p_w, b_w.to(p_w.dtype.element_ty), boundary_check=(0, 1)) + diff --git a/benchmarks/TritonNPUBench/vllm/reduce_segments/reduce_segments.py b/benchmarks/TritonNPUBench/vllm/reduce_segments/reduce_segments.py new file mode 100644 index 00000000..e602f4ab --- /dev/null +++ b/benchmarks/TritonNPUBench/vllm/reduce_segments/reduce_segments.py @@ -0,0 +1,130 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/vllm_operator_cases/newtest_cases/test_reduce_segments.py +# Main kernel: reduce_segments +# PT file: reduce_segments_test_v3.pt + +import triton +import triton.language as tl +import torch +float8_info = torch.finfo(torch.float16) + +# === reduce_segments === +@triton.jit +def reduce_segments( + output_ptr, # [num_tokens, num_query_heads, head_size] + segm_output_ptr, + # [num_tokens, num_query_heads, max_num_segments, head_size] + segm_max_ptr, # [num_tokens, num_query_heads, max_num_segments] + segm_expsum_ptr, # [num_tokens, num_query_heads, max_num_segments] + seq_lens_ptr, # [num_seqs] + num_seqs, # int + num_query_heads: tl.constexpr, # int + out_scale_inv, # float32 + output_stride_0: tl.int64, # int + output_stride_1: tl.int64, # int, should be equal to head_size + block_table_stride: tl.int64, # int + TILE_SIZE: tl.constexpr, # int + HEAD_SIZE: tl.constexpr, # int, must be power of 2 + HEAD_SIZE_PADDED: tl.constexpr, # int, must be power of 2 + query_start_len_ptr, # [num_seqs+1] + BLOCK_Q: tl.constexpr, # int + NUM_SEGMENTS_PER_SEQ: tl.constexpr, # int + USE_FP8: tl.constexpr, # bool + FP8_MIN: tl.constexpr = float8_info.min, + FP8_MAX: tl.constexpr = float8_info.max, +): + query_token_idx = tl.program_id(0) + query_head_idx = tl.program_id(1) + + seq_idx = find_seq_idx( + query_start_len_ptr, query_token_idx, num_seqs, BLOCK_Q, False + ) + + # sequence len for this particular sequence + seq_len = tl.load(seq_lens_ptr + seq_idx) + + # number of segments for this particular sequence + num_segments = NUM_SEGMENTS_PER_SEQ + tiles_per_segment = cdiv_fn(seq_len, num_segments * TILE_SIZE) + + # create masks for subsequent loads + act_num_segments = cdiv_fn(seq_len, tiles_per_segment * TILE_SIZE) + segm_mask = tl.arange(0, NUM_SEGMENTS_PER_SEQ) < tl.full( + [NUM_SEGMENTS_PER_SEQ], act_num_segments, dtype=tl.int32 + ) + dim_mask = tl.where(tl.arange(0, HEAD_SIZE_PADDED) < HEAD_SIZE, 1, 0).to(tl.int1) + + # load segment maxima + segm_offset = ( + query_token_idx.to(tl.int64) * (num_query_heads * NUM_SEGMENTS_PER_SEQ) + + query_head_idx * NUM_SEGMENTS_PER_SEQ + + tl.arange(0, NUM_SEGMENTS_PER_SEQ) + ) + segm_max = tl.load(segm_max_ptr + segm_offset, mask=segm_mask, other=float("-inf")) + overall_max = tl.max(segm_max) + + # load and rescale segment exp sums + segm_expsum = tl.load(segm_expsum_ptr + segm_offset, mask=segm_mask, other=0.0) + segm_expsum = segm_expsum * tl.exp(segm_max - overall_max) + overall_expsum = tl.sum(segm_expsum) + + # load, rescale, and add segment attention outputs + segm_output_offset = ( + query_token_idx.to(tl.int64) + * (num_query_heads * NUM_SEGMENTS_PER_SEQ * HEAD_SIZE_PADDED) + + query_head_idx * (NUM_SEGMENTS_PER_SEQ * HEAD_SIZE_PADDED) + + tl.arange(0, NUM_SEGMENTS_PER_SEQ)[:, None] * HEAD_SIZE_PADDED + + tl.arange(0, HEAD_SIZE_PADDED)[None, :] + ) + segm_output = tl.load( + segm_output_ptr + segm_output_offset, + mask=segm_mask[:, None] & dim_mask[None, :], + other=0.0, + ) + segm_output *= tl.exp(segm_max - overall_max)[:, None] + acc_sum = tl.sum(segm_output, axis=0) + # safely divide by overall_expsum, returning 0.0 if overall_expsum is 0 + acc = tl.where(overall_expsum == 0.0, 0.0, acc_sum / overall_expsum) + + if USE_FP8: + acc = acc * tl.load(out_scale_inv) + acc = tl.clamp(acc, FP8_MIN, FP8_MAX) + + # write result + output_offset = ( + query_token_idx * output_stride_0 + + query_head_idx * output_stride_1 + + tl.arange(0, HEAD_SIZE_PADDED) + ) + tl.store(output_ptr + output_offset, acc, mask=dim_mask) + + +# === cdiv_fn === +@triton.jit +def cdiv_fn(x, y): + return (x + y - 1) // y + + +# === find_seq_idx === +@triton.jit +def find_seq_idx( + query_start_len_ptr, + target_idx, + num_seqs, + BLOCK_Q: tl.constexpr, + use_q_block_mode: tl.constexpr, +): + left: tl.int32 = 0 + right = num_seqs + while left < right: + mid = (left + right) // 2 + val = tl.load(query_start_len_ptr + mid) + mid_val = val // BLOCK_Q + mid if use_q_block_mode else val + + if mid_val <= target_idx: + left = mid + 1 + else: + right = mid + + return left - 1 + diff --git a/benchmarks/TritonNPUBench/vllm/rejection_greedy_sample_kernel/rejection_greedy_sample_kernel.py b/benchmarks/TritonNPUBench/vllm/rejection_greedy_sample_kernel/rejection_greedy_sample_kernel.py new file mode 100644 index 00000000..80046083 --- /dev/null +++ b/benchmarks/TritonNPUBench/vllm/rejection_greedy_sample_kernel/rejection_greedy_sample_kernel.py @@ -0,0 +1,53 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/vllm_operator_cases/newtest_cases/test_rejection_greedy_sample_kernel.py +# Main kernel: rejection_greedy_sample_kernel +# PT file: rejection_greedy_sample_kernel_v2.pt + +import triton +import triton.language as tl + + +# === rejection_greedy_sample_kernel === +@triton.jit(do_not_specialize=["max_spec_len"]) +def rejection_greedy_sample_kernel( + output_token_ids_ptr, # [batch_size, max_spec_len + 1] + cu_num_draft_tokens_ptr, # [batch_size] + draft_token_ids_ptr, # [num_tokens] + target_argmax_ptr, # [num_tokens] + bonus_token_ids_ptr, # [batch_size] + is_greedy_ptr, # [batch_size] or None + max_spec_len, +): + req_idx = tl.program_id(0) + # FIXME(woosuk): Because is_greedy_ptr is not None at profiling run, + # re-compilation may happen during runtime when is_greedy_ptr is None. + is_greedy = True if is_greedy_ptr is None else tl.load(is_greedy_ptr + req_idx) + if not is_greedy: + # Early exit for non-greedy sampling requests. + return + + start_idx = 0 if req_idx == 0 else tl.load(cu_num_draft_tokens_ptr + req_idx - 1) + end_idx = tl.load(cu_num_draft_tokens_ptr + req_idx) + num_draft_tokens = end_idx - start_idx + + rejected = False + for pos in range(num_draft_tokens): + if not rejected: + draft_token_id = tl.load(draft_token_ids_ptr + start_idx + pos) + target_argmax_id = tl.load(target_argmax_ptr + start_idx + pos) + tl.store( + output_token_ids_ptr + req_idx * (max_spec_len + 1) + pos, + target_argmax_id, + ) + if draft_token_id != target_argmax_id: + # Reject. + rejected = True + + if not rejected: + # If all tokens are accepted, append the bonus token. + bonus_token_id = tl.load(bonus_token_ids_ptr + req_idx) + tl.store( + output_token_ids_ptr + req_idx * (max_spec_len + 1) + num_draft_tokens, + bonus_token_id, + ) + diff --git a/benchmarks/TritonNPUBench/vllm/rejection_random_sample_kernel/rejection_random_sample_kernel.py b/benchmarks/TritonNPUBench/vllm/rejection_random_sample_kernel/rejection_random_sample_kernel.py new file mode 100644 index 00000000..66009ca2 --- /dev/null +++ b/benchmarks/TritonNPUBench/vllm/rejection_random_sample_kernel/rejection_random_sample_kernel.py @@ -0,0 +1,70 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/vllm_operator_cases/newtest_cases/test_rejection_random_sample_kernel.py +# Main kernel: rejection_random_sample_kernel +# PT file: test_rejection_random_sample_kernel_v2.pt + +import triton +import triton.language as tl + + +# === rejection_random_sample_kernel === +@triton.jit(do_not_specialize=["max_spec_len"]) +def rejection_random_sample_kernel( + output_token_ids_ptr, # [batch_size, max_spec_len + 1] + cu_num_draft_tokens_ptr, # [batch_size] + draft_token_ids_ptr, # [num_tokens] + draft_probs_ptr, # [num_tokens, vocab_size] or None + target_probs_ptr, # [num_tokens, vocab_size] + bonus_token_ids_ptr, # [batch_size] + recovered_token_ids_ptr, # [num_tokens] + uniform_probs_ptr, # [num_tokens] + is_greedy_ptr, # [batch_size] + max_spec_len, + vocab_size, + NO_DRAFT_PROBS: tl.constexpr, +): + req_idx = tl.program_id(0) + is_greedy = tl.load(is_greedy_ptr + req_idx) + if is_greedy: + # Early exit for greedy sampling requests. + return + + start_idx = 0 if req_idx == 0 else tl.load(cu_num_draft_tokens_ptr + req_idx - 1) + end_idx = tl.load(cu_num_draft_tokens_ptr + req_idx) + num_draft_tokens = end_idx - start_idx + + rejected = False + for pos in range(num_draft_tokens): + if not rejected: + draft_token_id = tl.load(draft_token_ids_ptr + start_idx + pos) + if NO_DRAFT_PROBS: + draft_prob = 1 + else: + draft_prob = tl.load( + draft_probs_ptr + (start_idx + pos) * vocab_size + draft_token_id + ) + target_prob = tl.load( + target_probs_ptr + (start_idx + pos) * vocab_size + draft_token_id + ) + uniform_prob = tl.load(uniform_probs_ptr + start_idx + pos) + # NOTE(woosuk): While the draft probability should never be 0, + # we check it to avoid NaNs. If it happens to be 0, we reject. + if draft_prob > 0 and target_prob / draft_prob >= uniform_prob: + # Accept. + token_id = draft_token_id + else: + # Reject. Use recovered token. + rejected = True + token_id = tl.load(recovered_token_ids_ptr + start_idx + pos) + tl.store( + output_token_ids_ptr + req_idx * (max_spec_len + 1) + pos, token_id + ) + + if not rejected: + # If all tokens are accepted, append the bonus token. + bonus_token_id = tl.load(bonus_token_ids_ptr + req_idx) + tl.store( + output_token_ids_ptr + req_idx * (max_spec_len + 1) + num_draft_tokens, + bonus_token_id, + ) + diff --git a/benchmarks/TritonNPUBench/vllm/reshape_and_cache_kernel_flash/reshape_and_cache_kernel_flash.py b/benchmarks/TritonNPUBench/vllm/reshape_and_cache_kernel_flash/reshape_and_cache_kernel_flash.py new file mode 100644 index 00000000..c4362b95 --- /dev/null +++ b/benchmarks/TritonNPUBench/vllm/reshape_and_cache_kernel_flash/reshape_and_cache_kernel_flash.py @@ -0,0 +1,113 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/vllm_operator_cases/newtest_cases/test_reshape_and_cache_kernel_flash.py +# Main kernel: reshape_and_cache_kernel_flash +# PT file: reshape_and_cache_test.pt + +import triton +import triton.language as tl + + +# === reshape_and_cache_kernel_flash === +@triton.jit +def reshape_and_cache_kernel_flash( + key_ptr, # [num_tokens, num_heads, head_size] + value_ptr, # [num_tokens, num_heads, head_size] + key_cache_ptr, # [num_blocks, block_size, num_heads, head_size] + value_cache_ptr, # [num_blocks, block_size, num_heads, head_size] + slot_mapping_ptr, # [num_tokens] + k_scale, # float32 + v_scale, # float32 + # strides + key_stride: tl.int64, + value_stride: tl.int64, + block_stride: tl.int64, + head_stride: tl.int64, + dim_stride_k: tl.int64, + dim_stride_v: tl.int64, + page_stride: tl.int64, + num_heads: tl.constexpr, + head_size: tl.constexpr, + block_size: tl.constexpr, + x: tl.constexpr, + USE_HEAD_MAJOR_LAYOUT: tl.constexpr, + # FP8 flags + FP8_KV_CACHE: tl.constexpr, + # tune parameters + TILE_SIZE: tl.constexpr, +): + token_idx = tl.program_id(axis=0) + slot_idx = tl.load(slot_mapping_ptr + token_idx).to(tl.int64) + if slot_idx < 0: + # Padding token that should be ignored. + return + + block_idx = slot_idx // block_size + block_offset = slot_idx % block_size + + tile_i = tl.program_id(axis=1) + tile_offs = tl.arange(0, TILE_SIZE) + tile_pos = tile_i * TILE_SIZE + tile_offs + src_key_idx = token_idx * key_stride + src_value_idx = token_idx * value_stride + + if USE_HEAD_MAJOR_LAYOUT: + # Decompose the tile index back into head and dim coordinates. + cur_head = tile_pos // head_size + cur_dim = tile_pos % head_size + # Value addressing (4D): [Block, Head, Dim, Slot] + tgt_idx_v = ( + block_idx * block_stride + + cur_head * head_stride + + cur_dim * dim_stride_v + + block_offset * 1 + ) + # Key addressing (5D): [Block, Head, Dim//8, Slot, 8] + tgt_idx_k = ( + block_idx * block_stride + + cur_head * head_stride + + (cur_dim // x) * dim_stride_k + + block_offset * x + + (cur_dim % x) + ) + else: + tgt_base = block_idx * block_stride + block_offset * page_stride + tgt_idx_k = tgt_base + tile_pos + tgt_idx_v = tgt_base + tile_pos + + # [TILE_SIZE] + key_load = tl.load( + key_ptr + src_key_idx + tile_pos, mask=tile_pos < (num_heads * head_size) + ) + if FP8_KV_CACHE: + # tl.store will do the correct implicit cast to fp8, + # based on the key_cache_ptr.dtype.element_ty + key_tile = key_load if key_load.dtype.is_fp8() else key_load / tl.load(k_scale) + else: + key_tile = key_load + + # [TILE_SIZE] + value_load = tl.load( + value_ptr + src_value_idx + tile_pos, mask=tile_pos < (num_heads * head_size) + ) + if FP8_KV_CACHE: + if value_load.dtype.is_fp8(): + value_tile = value_load + else: + # tl.store will do the correct implicit cast to fp8, + # based on the value_cache_ptr.dtype.element_ty + value_tile = value_load / tl.load(v_scale) + else: + value_tile = value_load + + tl.store( + key_cache_ptr + tgt_idx_k, + key_tile, + mask=tile_pos < (num_heads * head_size), + ) + tl.store( + value_cache_ptr + tgt_idx_v, + value_tile, + mask=tile_pos < (num_heads * head_size), + ) + return + diff --git a/benchmarks/TritonNPUBench/vllm/reshape_and_cache_kernel_flash_diffkv/reshape_and_cache_kernel_flash_diffkv.py b/benchmarks/TritonNPUBench/vllm/reshape_and_cache_kernel_flash_diffkv/reshape_and_cache_kernel_flash_diffkv.py new file mode 100644 index 00000000..48462da3 --- /dev/null +++ b/benchmarks/TritonNPUBench/vllm/reshape_and_cache_kernel_flash_diffkv/reshape_and_cache_kernel_flash_diffkv.py @@ -0,0 +1,88 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/vllm_operator_cases/newtest_cases/test_reshape_and_cache_kernel_flash_diffkv.py +# Main kernel: reshape_and_cache_kernel_flash_diffkv +# PT file: reshape_and_cache_diffkv_test_v3.pt + +import triton +import triton.language as tl + + +# === reshape_and_cache_kernel_flash_diffkv === +@triton.jit +def reshape_and_cache_kernel_flash_diffkv( + key_ptr, # [num_tokens, num_heads, head_size] + value_ptr, # [num_tokens, num_heads, head_size_v] + kv_cache_ptr, # [num_blocks, block_size, num_heads, head_size + head_size_v] + slot_mapping_ptr, # [num_tokens] + k_scale, # float32 + v_scale, # float32 + # strides + key_stride: tl.int64, + value_stride: tl.int64, + block_stride: tl.int64, + page_stride: tl.int64, + num_heads: tl.constexpr, + head_size_k: tl.constexpr, + head_size_v: tl.constexpr, + block_size: tl.constexpr, + # FP8 flags + FP8_KV_CACHE: tl.constexpr, + # tune parameters + TILE_SIZE: tl.constexpr, +): + token_idx = tl.program_id(axis=0) + slot_idx = tl.load(slot_mapping_ptr + token_idx).to(tl.int64) + if slot_idx < 0: + # Padding token that should be ignored. + return + + tile_i = tl.program_id(axis=1) + tile_offs = tl.arange(0, TILE_SIZE) + + block_idx = slot_idx // block_size + block_offset = slot_idx % block_size + + src_key_idx = token_idx * key_stride + tile_i * head_size_k + src_value_idx = token_idx * value_stride + tile_i * head_size_v + + tgt_idx = ( + block_idx * block_stride + + block_offset * page_stride + + tile_i * (head_size_k + head_size_v) + ) + + # [TILE_SIZE] + key_load = tl.load(key_ptr + src_key_idx + tile_offs, mask=tile_offs < head_size_k) + if FP8_KV_CACHE: + # tl.store will do the correct implicit cast to fp8, + # based on the key_cache_ptr.dtype.element_ty + key_tile = key_load if key_load.dtype.is_fp8() else key_load / tl.load(k_scale) + else: + key_tile = key_load + + # [TILE_SIZE] + value_load = tl.load( + value_ptr + src_value_idx + tile_offs, mask=tile_offs < head_size_v + ) + if FP8_KV_CACHE: + if value_load.dtype.is_fp8(): + value_tile = value_load + else: + # tl.store will do the correct implicit cast to fp8, + # based on the value_cache_ptr.dtype.element_ty + value_tile = value_load / tl.load(v_scale) + else: + value_tile = value_load + + tl.store( + kv_cache_ptr + tgt_idx + tile_offs, + key_tile, + mask=tile_offs < head_size_k, + ) + tl.store( + kv_cache_ptr + tgt_idx + head_size_k + tile_offs, + value_tile, + mask=tile_offs < head_size_v, + ) + return + diff --git a/benchmarks/TritonNPUBench/vllm/sample_recovered_tokens_kernel/sample_recovered_tokens_kernel.py b/benchmarks/TritonNPUBench/vllm/sample_recovered_tokens_kernel/sample_recovered_tokens_kernel.py new file mode 100644 index 00000000..8d64eff8 --- /dev/null +++ b/benchmarks/TritonNPUBench/vllm/sample_recovered_tokens_kernel/sample_recovered_tokens_kernel.py @@ -0,0 +1,63 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/vllm_operator_cases/newtest_cases/test_sample_recovered_tokens_kernel.py +# Main kernel: sample_recovered_tokens_kernel +# PT file: sample_recovered_tokens_kernel_v2.pt + +import triton +import triton.language as tl + + +# === sample_recovered_tokens_kernel === +@triton.jit +def sample_recovered_tokens_kernel( + output_token_ids_ptr, # [num_tokens] + cu_num_draft_tokens_ptr, # [batch_size] + draft_token_ids_ptr, # [num_tokens] + draft_probs_ptr, # [num_tokens, vocab_size] or None + target_probs_ptr, # [num_tokens, vocab_size] + q_ptr, # [batch_size, vocab_size] + vocab_size, + PADDED_VOCAB_SIZE: tl.constexpr, + NO_DRAFT_PROBS: tl.constexpr, +): + req_idx = tl.program_id(0) + start_idx = 0 if req_idx == 0 else tl.load(cu_num_draft_tokens_ptr + req_idx - 1) + end_idx = tl.load(cu_num_draft_tokens_ptr + req_idx) + num_draft_tokens = end_idx - start_idx + + # Early exit for out-of-range positions. + pos = tl.program_id(1) + if pos >= num_draft_tokens: + return + + vocab_offset = tl.arange(0, PADDED_VOCAB_SIZE) + if NO_DRAFT_PROBS: + draft_token_id = tl.load(draft_token_ids_ptr + start_idx + pos) + prob = tl.load( + target_probs_ptr + (start_idx + pos) * vocab_size + vocab_offset, + mask=((vocab_offset < vocab_size) & (vocab_offset != draft_token_id)), + other=0, + ) + else: + draft_prob = tl.load( + draft_probs_ptr + (start_idx + pos) * vocab_size + vocab_offset, + mask=vocab_offset < vocab_size, + other=0, + ) + target_prob = tl.load( + target_probs_ptr + (start_idx + pos) * vocab_size + vocab_offset, + mask=vocab_offset < vocab_size, + other=0, + ) + prob = tl.maximum(target_prob - draft_prob, 0) + # NOTE(woosuk): We don't need `prob = prob / tl.sum(prob)` here because + # `tl.argmax` will select the maximum value. + + q = tl.load( + q_ptr + req_idx * vocab_size + vocab_offset, + mask=vocab_offset < vocab_size, + other=float("-inf"), + ) + recovered_id = tl.argmax(prob / q, axis=-1) + tl.store(output_token_ids_ptr + start_idx + pos, recovered_id) + diff --git a/benchmarks/TritonNPUBench/vllm/scaled_mm_kernel/scaled_mm_kernel.py b/benchmarks/TritonNPUBench/vllm/scaled_mm_kernel/scaled_mm_kernel.py new file mode 100644 index 00000000..3ee7d0b3 --- /dev/null +++ b/benchmarks/TritonNPUBench/vllm/scaled_mm_kernel/scaled_mm_kernel.py @@ -0,0 +1,130 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/vllm_operator_cases/newtest_cases/test_scaled_mm_kernel.py +# Main kernel: scaled_mm_kernel +# PT file: scaled_mm_kernel_v3.pt + +import triton +import triton.language as tl + + +# === scaled_mm_kernel === +@triton.jit +def scaled_mm_kernel( + a_ptr, + b_ptr, + scale_a_ptr, + scale_b_ptr, + c_ptr, + bias_ptr, + M, + N, + K, + stride_am, + stride_ak, + stride_bk, + stride_bn, + stride_cm, + stride_cn, + ACCUMULATOR_DTYPE: tl.constexpr, + BLOCK_SIZE_M: tl.constexpr, + BLOCK_SIZE_N: tl.constexpr, + BLOCK_SIZE_K: tl.constexpr, + BLOCK_SIZE_SCALE_A: tl.constexpr, + BLOCK_SIZE_SCALE_B: tl.constexpr, +): + pid = tl.program_id(axis=0) + + num_pid_n = tl.cdiv(N, BLOCK_SIZE_N) + + pid_m = pid // num_pid_n + pid_n = pid % num_pid_n + + accumulator_dtype = ACCUMULATOR_DTYPE + accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=accumulator_dtype) + + # NOTE: Some tensor inputs are so large, they will cause int32 overflow + # so it is necessary to use tl.int64 for all the offsets, else SEGV will + # eventually occur. + + # Offsets and masks. + offsets_am = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M).to(tl.int64) + masks_am = offsets_am < M + + offsets_bn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N).to(tl.int64) + masks_bn = offsets_bn < N + + offsets_k = tl.arange(0, BLOCK_SIZE_K).to(tl.int64) + offsets_a = stride_am * offsets_am[:, None] + stride_ak * offsets_k[None, :] + offsets_b = stride_bk * offsets_k[:, None] + stride_bn * offsets_bn[None, :] + + # NOTE: BLOCK_SIZE_SCALE_A could be 1 or BLOCK_SIZE_M, so need to create + # appropriate offsets and masks for each case. Same goes for + # BLOCK_SIZE_SCALE_B. + offsets_scale_am = ( + tl.arange(0, BLOCK_SIZE_SCALE_A) + + (BLOCK_SIZE_SCALE_A > 1) * pid_m * BLOCK_SIZE_M + ) + masks_scale_am = offsets_scale_am < M + + offsets_scale_bn = ( + tl.arange(0, BLOCK_SIZE_SCALE_B) + + (BLOCK_SIZE_SCALE_B > 1) * pid_n * BLOCK_SIZE_N + ) + masks_scale_bn = offsets_scale_bn < N + + a_ptrs = a_ptr + offsets_a + b_ptrs = b_ptr + offsets_b + + scale_a_ptrs = scale_a_ptr + offsets_scale_am + scale_b_ptrs = scale_b_ptr + offsets_scale_bn + + for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)): + masks_k = offsets_k < K + masks_a = masks_am[:, None] & masks_k[None, :] + a = tl.load(a_ptrs, mask=masks_a) + + masks_b = masks_k[:, None] & masks_bn[None, :] + b = tl.load(b_ptrs, mask=masks_b) + + # Accumulate results. + accumulator = tl.dot(a, b, accumulator, out_dtype=accumulator_dtype) + + offsets_k += BLOCK_SIZE_K + a_ptrs += BLOCK_SIZE_K * stride_ak + b_ptrs += BLOCK_SIZE_K * stride_bk + + # Apply scale at end. + masks_scale_a = masks_scale_am[:, None] & (tl.arange(0, 1) < 1)[:, None] + scale_a = tl.load(scale_a_ptrs[:, None], masks_scale_a) + # Need to broadcast to the appropriate size, if scale_a is already + # (BLOCK_SIZE_M, 1) then it will broadcast to its own shape. Same goes + # for scale_b below. + scale_a = scale_a.broadcast_to((BLOCK_SIZE_M, 1)) + accumulator = scale_a * accumulator.to(tl.float32) + + masks_scale_b = masks_scale_bn[:, None] & (tl.arange(0, 1) < 1)[None, :] + scale_b = tl.load(scale_b_ptrs[:, None], masks_scale_b) + scale_b = scale_b.broadcast_to((BLOCK_SIZE_N, 1)) + accumulator = scale_b.T * accumulator.to(tl.float32) + + # Convert to output format. + c = accumulator.to(c_ptr.type.element_ty) + + # Add bias, it's already in output format, so add it after conversion. + if bias_ptr: + offsets_bias = offsets_bn + bias_ptrs = bias_ptr + offsets_bias + bias_mask = offsets_bias < N + bias = tl.load(bias_ptrs, bias_mask) + c += bias + + # Save output + offs_cm = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M).to(tl.int64) + offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N).to(tl.int64) + offs_cm = offs_cm.to(tl.int64) + offs_cn = offs_cn.to(tl.int64) + c_ptrs = c_ptr + stride_cm * offs_cm[:, None] + stride_cn * offs_cn[None, :] + c_mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < N) + + tl.store(c_ptrs, c, mask=c_mask) + diff --git a/benchmarks/TritonNPUBench/vllm/softplus_kernel/softplus_kernel.py b/benchmarks/TritonNPUBench/vllm/softplus_kernel/softplus_kernel.py new file mode 100644 index 00000000..f752d441 --- /dev/null +++ b/benchmarks/TritonNPUBench/vllm/softplus_kernel/softplus_kernel.py @@ -0,0 +1,26 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/vllm_operator_cases/newtest_cases/test_softplus.py +# Main kernel: softplus_kernel +# PT file: softplus_v2.pt + +import triton +import triton.language as tl + +@triton.jit +def softplus(dt): + dt = tl.where(dt <= 20.0, tl.math.log(tl.math.exp(dt) + 1), dt) + return dt + + +# === softplus_kernel === +@triton.jit +def softplus_kernel(input_ptr, output_ptr, n_elements, BLOCK_SIZE: tl.constexpr): + pid = tl.program_id(axis=0) + block_start = pid * BLOCK_SIZE + offsets = block_start + tl.arange(0, BLOCK_SIZE) + mask = offsets < n_elements + + input_block = tl.load(input_ptr + offsets, mask=mask) + output_block = softplus(input_block) + tl.store(output_ptr + offsets, output_block, mask=mask) + diff --git a/benchmarks/TritonNPUBench/vllm/solve_tril_16x16_kernel/solve_tril_16x16_kernel.py b/benchmarks/TritonNPUBench/vllm/solve_tril_16x16_kernel/solve_tril_16x16_kernel.py new file mode 100644 index 00000000..2c4d800f --- /dev/null +++ b/benchmarks/TritonNPUBench/vllm/solve_tril_16x16_kernel/solve_tril_16x16_kernel.py @@ -0,0 +1,75 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/vllm_operator_cases/newtest_cases/test_solve_tril_16x16_kernel.py +# Main kernel: solve_tril_16x16_kernel +# PT file: test_solve_tril_16x16_kernel_v3.pt + +import triton +import triton.language as tl + + +# === solve_tril_16x16_kernel === +@triton.jit(do_not_specialize=["T"]) +def solve_tril_16x16_kernel( + A, + Ai, + cu_seqlens, + chunk_indices, + T, + H: tl.constexpr, + BT: tl.constexpr, + USE_TMA: tl.constexpr, + IS_VARLEN: tl.constexpr, + DOT_PRECISION: tl.constexpr, +): + i_t, i_bh = tl.program_id(0), tl.program_id(1) + i_b, i_h = i_bh // H, i_bh % H + if IS_VARLEN: + i_n, i_t = ( + tl.load(chunk_indices + i_t * 2).to(tl.int32), + tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), + ) + bos, eos = ( + tl.load(cu_seqlens + i_n).to(tl.int32), + tl.load(cu_seqlens + i_n + 1).to(tl.int32), + ) + T = eos - bos + else: + bos, eos = i_b * T, i_b * T + T + o_i = tl.arange(0, 16) + m_A = o_i[:, None] > o_i[None, :] + m_I = o_i[:, None] == o_i[None, :] + + A = A + (bos * H + i_h) * BT + Ai = Ai + (bos * H + i_h) * 16 + + offset = (i_t * 16) % BT + if not USE_TMA: + p_A = tl.make_block_ptr( + A, (T, BT), (H * BT, 1), (i_t * 16, offset), (16, 16), (1, 0) + ) + # [16, 16] + b_A = tl.load(p_A, boundary_check=(0, 1)).to(tl.float32) + else: + desc = make_tensor_descriptor(A, [T, BT], [H * BT, 1], [16, 16]) + desc_o = make_tensor_descriptor(Ai, [T, 16], [H * 16, 1], [16, 16]) + b_A = desc.load([i_t * 16, offset]).to(tl.float32) + b_A = -tl.where(m_A, b_A, 0) + + for i in range(2, min(16, T - i_t * 16)): + # [16] + b_a = -tl.load(A + (i_t * 16 + i) * H * BT + o_i + offset) + b_a = b_a + tl.sum(b_a[:, None] * b_A, 0) + b_A = tl.where((o_i == i)[:, None], b_a, b_A) + b_A += m_I + if not USE_TMA: + p_Ai = tl.make_block_ptr( + Ai, (T, 16), (H * 16, 1), (i_t * 16, 0), (16, 16), (1, 0) + ) + tl.store( + p_Ai, + b_A.to(p_Ai.dtype.element_ty, fp_downcast_rounding="rtne"), + boundary_check=(0, 1), + ) + else: + desc_o.store([i_t * 16, 0], b_A.to(desc_o.dtype, fp_downcast_rounding="rtne")) + diff --git a/benchmarks/TritonNPUBench/vllm/test_round_int8_kernel/test_round_int8_kernel.py b/benchmarks/TritonNPUBench/vllm/test_round_int8_kernel/test_round_int8_kernel.py new file mode 100644 index 00000000..f17c9290 --- /dev/null +++ b/benchmarks/TritonNPUBench/vllm/test_round_int8_kernel/test_round_int8_kernel.py @@ -0,0 +1,26 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/vllm_operator_cases/newtest_cases/test_round_int8.py +# Main kernel: test_round_int8_kernel +# PT file: round_int8.pt + +import triton +import triton.language as tl + +@triton.jit +def round_int8(x): + return tl.extra.cann.libdevice.round(x).to(tl.int8) + +@triton.jit +def test_round_int8_kernel( + x_ptr, # 输入数据指针 + output_ptr, # 输出数据指针 + N: tl.constexpr, + BLOCK: tl.constexpr, +): + pid = tl.program_id(0) + offsets = pid * BLOCK + tl.arange(0, BLOCK) + mask = offsets < N + x = tl.load(x_ptr + offsets, mask=mask, other=0.0) + result = round_int8(x) + tl.store(output_ptr + offsets, result, mask=mask) + diff --git a/benchmarks/TritonNPUBench/vllm/triton_scale_swizzle/triton_scale_swizzle.py b/benchmarks/TritonNPUBench/vllm/triton_scale_swizzle/triton_scale_swizzle.py new file mode 100644 index 00000000..970d9838 --- /dev/null +++ b/benchmarks/TritonNPUBench/vllm/triton_scale_swizzle/triton_scale_swizzle.py @@ -0,0 +1,73 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/vllm_operator_cases/newtest_cases/test_triton_scale_swizzle.py +# Main kernel: triton_scale_swizzle +# PT file: triton_scale_swizzle_v2.pt + +import triton +import triton.language as tl +import torch + +# === triton_scale_swizzle === +@triton.jit +def triton_scale_swizzle( + scale_ptr: torch.Tensor, + scale_rows: int, + scale_cols: int, + output_ptr: torch.Tensor, + input_row_stride: int, + output_block_stride: int, + BLOCK_ROWS: tl.constexpr, + BLOCK_COLS: tl.constexpr, +): + """ + Rearranges tensor data from row-major to block-scaled swizzle format. + + Args: + scale_ptr: Pointer to the input scale tensor + scale_rows: Number of rows in the scale tensor + scale_cols: Number of columns in the scale tensor + output_ptr: Pointer to the output tensor + input_row_stride: Stride between rows in the input tensor + output_block_stride: Stride between blocks in the output tensor + BLOCK_ROWS: Number of rows in a tile (compile-time constant) + BLOCK_COLS: Number of columns in a tile (compile-time constant) + """ + pid_row = tl.program_id(0) + pid_col = tl.program_id(1) + + rows = tl.arange(0, BLOCK_ROWS)[:, None] + cols = tl.arange(0, BLOCK_COLS)[None, :] + + # Calculate starting row and column for this tile + start_row = pid_row * BLOCK_ROWS + start_col = pid_col * BLOCK_COLS + global_rows = start_row + rows + global_cols = start_col + cols + + mask = (global_rows < scale_rows) & (global_cols < scale_cols) + + input_scales = tl.load( + scale_ptr + global_rows * input_row_stride + global_cols, + mask=mask, + other=0.0, + ) + + r_div_32 = rows // 32 + r_mod_32 = rows % 32 + + # 2) Rearrange to (32, 4, 4) then to final (32, 16) coordinates + dest_indices = r_mod_32 * 16 + r_div_32 * 4 + cols + + # Flatten + dest_indices_flat = tl.reshape(dest_indices, (BLOCK_ROWS * BLOCK_COLS)) + scales_flat = tl.reshape(input_scales, (BLOCK_ROWS * BLOCK_COLS)) + + # Calculate block offset using provided output block stride + LOCAL_NUMEL = BLOCK_ROWS * BLOCK_COLS + block_offset = pid_col * LOCAL_NUMEL + (pid_row * output_block_stride) + + tl.store( + output_ptr + block_offset + dest_indices_flat, + scales_flat, + ) + diff --git a/benchmarks/TritonNPUBench/vllm/vllm_gpu_perf.csv b/benchmarks/TritonNPUBench/vllm/vllm_gpu_perf.csv new file mode 100644 index 00000000..0aaa8a2f --- /dev/null +++ b/benchmarks/TritonNPUBench/vllm/vllm_gpu_perf.csv @@ -0,0 +1,113 @@ +Name,Type,Input Shapes,Input Data Types,Duration(us), +_apply_grammar_bitmask_kernel,VV,None,None,1.847, +_apply_write_kernel,VV,None,None,1.342, +_bias_kernel,VV,None,None,2.164, +_bincount_kernel,VV,None,None,5.447, +_bmm_chunk_fwd_kernel,CV,None,None,3.074, +_chunk_cumsum_fwd_kernel,VV,None,None,2.366, +_chunk_scan_fwd_kernel,CV,None,None,8.4, +_chunk_state_fwd_kernel,CV,None,None,2.886, +_chunk_state_varlen_kernel,CV,None,None,3.638, +_combine_sampled_and_draft_tokens_kernel,VV,None,None,2.338, +_compute_pid,VV,None,None,1.113, +_compute_slot_mappings_kernel,VV,None,None,3.516, +_convert_req_index_to_global_index_kernel,VV,None,None,2.479, +_copy_page_indices_kernel,VV,None,None,1.41, +_correct_attn_cp_out_kernel,VV,None,None,2.072, +_count_expert_num_tokens,VV,None,None,2.028, +_expand_idx_mapping_kernel,VV,None,None,1.126, +_fwd_diag_kernel,CV,None,None,9.549, +_fwd_grouped_kernel_stage1,CV,None,None,17.393, +_fwd_kernel,CV,None,None,14.454, +_fwd_kernel_alibi,CV,None,None,4.615, +_fwd_kernel_ep_gather,VV,None,None,5.32, +_fwd_kernel_ep_scatter_1,VV,None,None,1.145, +_fwd_kernel_ep_scatter_2,VV,None,None,4.942, +_fwd_kernel_stage1,VV,None,None,2.443, +_fwd_kernel_stage2,VV,None,None,1.555, +_fwd_kv_parallel,CV,None,None,6.597, +_fwd_kv_reduce,VV,None,None,3.186, +_fwd_none_diag_kernel,CV,None,None,7.17, +_gather_block_tables_kernel,VV,None,None,, +_get_num_sampled_and_rejected_kernel,VV,None,None,1.962, +_gumbel_sample_kernel,VV,None,None,2.95, +_layer_norm_fwd_1pass_kernel,VV,None,None,1.938, +_linear_attn_decode_kernel,VV,None,None,4.024, +_log_softmax_kernel,VV,None,None,4.498, +_lora_expand_kernel,CV,None,None,1.966, +_lora_shrink_kernel,CV,None,None,2.194, +_min_p_kernel,VV,None,None,1.599, +_num_nans_kernel,VV,None,None,1.413, +_pack_seq_kernel,VV,None,None,1.48, +_penalties_and_temperature_kernel,VV,None,None,1.78, +_per_token_group_quant_fp8,VV,None,None,2.045, +_per_token_group_quant_fp8_colmajor,VV,None,None,2.044, +_per_token_group_quant_int8,VV,None,None,2.033, +_per_token_quant_int8,VV,None,None,1.192, +_post_update_kernel,VV,None,None,1.24, +_prepare_eagle_docode_kernel,VV,None,None,4.519, +_prepare_eagle_inputs_kernel,VV,None,None,1.487, +_prepare_mrope_positions_kernel,VV,None,None,2.929, +_prepare_pos_seq_lens_kernel,VV,None,None,1.277, +_prepare_prefill_inputs_kernel,VV,None,None,1.503, +_ranks_kernel,VV,None,None,1.277, +_rejection_sample_kernel,VV,None,None,1.255, +_rms_norm_kernel,VV,None,None,3.336, +_selective_scan_update_kernel,VV,None,None,1.027, +_silu_mul_fp8_quant_deep_gemm,VV,None,None,5.562, +_silu_mul_per_token_group_quant_fp8_colmajor,VV,None,None,1.399, +_state_passing_fwd_kernel,VV,None,None,3.734, +_topk_log_softmax_kernel,VV,None,None,2.573, +_triton_mrope_forward,VV,None,None,2.121, +_trtllm_prefill_attn_kvfp8_dequant,VV,None,None,1.67, +_unpack_seq_triton_kernel,VV,None,None,3.671, +_update_eagle_inputs_kernel,VV,None,None,1.504, +_w8a8_block_int8_matmul,CV,None,None,2.812, +_w8a8_triton_block_scaled_mm,CV,None,None,4.028, +awq_dequantize_kernel,VV,None,None,1.183, +awq_gemm_kernel,CV,None,None,1.713, +batched_triton_kernel,CV,None,None,7.166, +bmm_kernel,CV,None,None,9.533, +chunk_fwd_kernel_o,CV,None,None,3.549, +chunk_gated_delta_rule_fwd_kernel_h_blockdim64,CV,None,None,, +chunk_gla_fwd_kernel_o,CV,None,None,6.268, +chunk_kda_scaled_dot_kkt_fwd_kernel_intra_sub_inter,CV,None,None,4.927, +chunk_kda_scaled_dot_kkt_fwd_kernel_intra_sub_intra,VV,None,None,10.149, +chunk_local_cumsum_scalar_kernel,VV,None,None,1.055, +chunk_local_cumsum_vector_kernel,CV,None,None,2.53, +chunk_scaled_dot_kkt_fwd_kernel,CV,None,None,1.679, +compute_identity_kernel,VV,None,None,1.449, +cp_mha_gather_cache_kernel,VV,None,None,4.347, +eagle_prepare_inputs_padded_kernel,VV,None,None,1.546, +eagle_prepare_next_token_padded_kernel,VV,None,None,1.341, +expand_kernel,VV,None,None,1.019, +fused_gdn_gating_kernel,VV,None,None,1.289, +fused_moe_kernel,CV,None,None,2.098, +fused_moe_kernel_gptq_awq,CV,None,None,15.196, +fused_recurrent_gated_delta_rule_fwd_kernel,VV,None,None,2.196, +kda_gate_fwd_kernel,VV,None,None,3.452, +kernel_paged_attention_2d,CV,None,None,6.763, +kernel_unified_attention_2d,CV,None,None,3.054, +kernel_unified_attention_3d,CV,None,None,2.583, +l2norm_fwd_kernel,VV,None,None,11.807, +l2norm_fwd_kernel1,VV,None,None,4.226, +l2norm_fwd_kernel2,VV,None,None,2.201, +layer_norm_fwd_kernel,VV,None,None,2.736, +layer_norm_gated_fwd_kernel,VV,None,None,8.181, +layer_norm_gated_fwd_kernel1,VV,None,None,2.061, +matmul_kernel_persistent,CV,None,None,7.463, +mean_kernel,VV,None,None,17.353, +merge_16x16_to_32x32_inverse_kernel,CV,None,None,7.178, +merge_16x16_to_64x64_inverse_kernel,CV,None,None,2.891, +merge_attn_states_kernel,VV,None,None,2.72, +pack_bitmatrix,VV,None,None,15.087, +recompute_w_u_fwd_kernel,CV,None,None,2.926, +reduce_segments,VV,None,None,10.498, +rejection_greedy_sample_kernel,VV,None,None,1.503, +rejection_random_sample_kernel,VV,None,None,1.797, +reshape_and_cache_kernel_flash,VV,None,None,2.381, +reshape_and_cache_kernel_flash_diffkv,VV,None,None,3.129, +sample_recovered_tokens_kernel,VV,None,None,2.309, +scaled_mm_kernel,CV,None,None,2.97, +solve_tril_16x16_kernel,VV,None,None,4.207, +triton_scale_swizzle,VV,None,None,1.173, From 7e8424279aacbd893b489262584b850fbcacdc13 Mon Sep 17 00:00:00 2001 From: zhumingming Date: Fri, 17 Apr 2026 21:35:45 +0800 Subject: [PATCH 2/3] =?UTF-8?q?=E3=80=90benchmark=E3=80=91=E6=96=B0?= =?UTF-8?q?=E5=A2=9Esglang=E4=BB=9395=E4=B8=AA=E7=AE=97=E5=AD=90=E6=8F=8F?= =?UTF-8?q?=E8=BF=B0?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../_act_quant_kernel/_act_quant_kernel.py | 86 ++++ .../_bmm_chunk_fwd_kernel.py | 122 ++++++ .../_causal_conv1d_fwd_kernel.py | 373 ++++++++++++++++ .../_chunk_cumsum_fwd_kernel.py | 105 +++++ .../_chunk_scan_fwd_kernel.py | 412 ++++++++++++++++++ .../_chunked_lora_expand_kernel.py | 142 ++++++ .../_chunked_lora_shrink_kernel.py | 118 +++++ .../_dequantize_k_cache_fast_kernel.py | 57 +++ .../_dequantize_k_cache_paged_kernel.py | 63 +++ .../_embedding_lora_a_kernel.py | 106 +++++ .../_fwd_grouped_kernel_stage1_rope.py | 281 ++++++++++++ .../_fwd_kernel_ep_gather.py | 75 ++++ .../_fwd_kernel_ep_scatter_1.py | 42 ++ .../_fwd_kernel_ep_scatter_2.py | 81 ++++ .../_gate_up_lora_b_kernel.py | 127 ++++++ .../_layer_norm_fwd_1pass_kernel.py | 107 +++++ .../_moe_sum_reduce_kernel.py | 60 +++ .../_qkv_lora_b_kernel/_qkv_lora_b_kernel.py | 128 ++++++ .../_quantize_k_cache_fast_kernel.py | 69 +++ .../_selective_scan_update_kernel.py | 235 ++++++++++ .../_sgemm_lora_a_kernel.py | 115 +++++ .../_sgemm_lora_b_kernel.py | 119 +++++ ...lu_and_mul_post_per_tensor_quant_kernel.py | 75 ++++ .../_silu_and_mul_post_quant_kernel.py | 87 ++++ .../_state_passing_fwd_kernel.py | 164 +++++++ .../_tma_align_input_scale_kernel.py | 39 ++ .../act_and_mul_kernel/act_and_mul_kernel.py | 83 ++++ .../add_tree_reduce_u64_kernel.py | 22 + .../chunk_fwd_kernel_o/chunk_fwd_kernel_o.py | 106 +++++ ...ated_delta_rule_fwd_kernel_h_blockdim64.py | 254 +++++++++++ ...aled_dot_kkt_fwd_kernel_intra_sub_inter.py | 110 +++++ ...aled_dot_kkt_fwd_kernel_intra_sub_intra.py | 100 +++++ .../chunk_local_cumsum_scalar_kernel.py | 59 +++ .../chunk_local_cumsum_vector_kernel.py | 88 ++++ .../chunk_scaled_dot_kkt_fwd_kernel.py | 80 ++++ .../compute_identity_kernel.py | 48 ++ .../compute_masked_m_triton_kernel.py | 17 + .../compute_problem_sizes_w4a8_kernel.py | 47 ++ .../compute_seg_indptr_triton_kernel.py | 26 ++ .../compute_src2dst_triton_kernel.py | 20 + .../concat_and_cast_mha_k_kernel.py | 48 ++ .../create_flashinfer_kv_indices_triton.py | 48 ++ .../create_flashmla_kv_indices_triton.py | 65 +++ .../deepep_compute_src2dst_triton_kernel.py | 21 + .../deepep_permute_triton_kernel.py | 40 ++ .../deepep_post_reorder_triton_kernel.py | 42 ++ .../deepgemm_compute_src2dst_triton_kernel.py | 30 ++ .../experts_combine_kernel.py | 40 ++ .../fill_gateup_input_triton_kernel.py | 50 +++ .../fn_triton_kernel/fn_triton_kernel.py | 60 +++ .../fused_dual_residual_rmsnorm_kernel.py | 55 +++ .../fused_gdn_gating_kernel.py | 41 ++ .../fused_moe_kernel/fused_moe_kernel.py | 291 +++++++++++++ .../fused_moe_kernel_gptq_awq.py | 262 +++++++++++ .../fused_moe_router_cudacore_kernel.py | 112 +++++ .../fused_moe_router_tensorcore_kernel.py | 136 ++++++ ...d_recurrent_gated_delta_rule_fwd_kernel.py | 158 +++++++ ...rent_gated_delta_rule_update_fwd_kernel.py | 188 ++++++++ .../fused_rmsnorm_kernel.py | 40 ++ ...sigmoid_gating_delta_rule_update_kernel.py | 158 +++++++ .../fused_softcap_kernel.py | 31 ++ .../gelu_and_mul_kernel.py | 47 ++ .../get_last_loc_kernel.py | 33 ++ .../get_num_kv_splits_triton.py | 61 +++ .../hash_tiles32_kernel_blocked.py | 92 ++++ .../l2norm_fwd_kernel/l2norm_fwd_kernel.py | 29 ++ .../l2norm_fwd_kernel1/l2norm_fwd_kernel1.py | 32 ++ .../layer_norm_gated_fwd_kernel.py | 89 ++++ .../layer_norm_gated_fwd_kernel1.py | 77 ++++ .../memcpy_triton_kernel.py | 35 ++ .../merge_16x16_to_32x32_inverse_kernel.py | 137 ++++++ .../merge_16x16_to_64x64_inverse_kernel.py | 243 +++++++++++ .../merge_state_kernel/merge_state_kernel.py | 62 +++ .../moe_align_block_size_stage1.py | 28 ++ .../moe_align_block_size_stage2.py | 22 + .../moe_align_block_size_stage3.py | 26 ++ .../moe_align_block_size_stage4.py | 39 ++ .../nsa_cp_round_robin_split_q_seqs_kernel.py | 31 ++ .../pad_draft_extend_query_kernel.py | 82 ++++ .../pad_sequence_with_mask_kernel.py | 63 +++ .../post_reorder_triton_kernel.py | 49 +++ ...t_reorder_triton_kernel_for_cutlass_moe.py | 59 +++ ...e_reorder_triton_kernel_for_cutlass_moe.py | 54 +++ .../recompute_w_u_fwd_kernel.py | 98 +++++ .../sglang/rotl32_kernel/rotl32_kernel.py | 29 ++ .../sglang/safe_exp_kernel/safe_exp_kernel.py | 25 ++ .../TritonNPUBench/sglang/sglang_gpu_perf.csv | 89 ++++ .../silu_and_mul_kernel.py | 47 ++ ...ise_quant_triton_kernel_for_cutlass_moe.py | 42 ++ .../solve_tril_16x16_kernel.py | 58 +++ .../sglang/tanh_kernel/tanh_kernel.py | 25 ++ .../track_mamba_state_if_needed_kernel.py | 69 +++ ...ransform_index_page_table_decode_kernel.py | 31 ++ .../unpad_draft_extend_output_kernel.py | 80 ++++ .../write_req_to_token_pool_triton.py | 46 ++ ...write_req_to_token_pool_triton_optimize.py | 52 +++ 96 files changed, 8445 insertions(+) create mode 100644 benchmarks/TritonNPUBench/sglang/_act_quant_kernel/_act_quant_kernel.py create mode 100644 benchmarks/TritonNPUBench/sglang/_bmm_chunk_fwd_kernel/_bmm_chunk_fwd_kernel.py create mode 100644 benchmarks/TritonNPUBench/sglang/_causal_conv1d_fwd_kernel/_causal_conv1d_fwd_kernel.py create mode 100644 benchmarks/TritonNPUBench/sglang/_chunk_cumsum_fwd_kernel/_chunk_cumsum_fwd_kernel.py create mode 100644 benchmarks/TritonNPUBench/sglang/_chunk_scan_fwd_kernel/_chunk_scan_fwd_kernel.py create mode 100644 benchmarks/TritonNPUBench/sglang/_chunked_lora_expand_kernel/_chunked_lora_expand_kernel.py create mode 100644 benchmarks/TritonNPUBench/sglang/_chunked_lora_shrink_kernel/_chunked_lora_shrink_kernel.py create mode 100644 benchmarks/TritonNPUBench/sglang/_dequantize_k_cache_fast_kernel/_dequantize_k_cache_fast_kernel.py create mode 100644 benchmarks/TritonNPUBench/sglang/_dequantize_k_cache_paged_kernel/_dequantize_k_cache_paged_kernel.py create mode 100644 benchmarks/TritonNPUBench/sglang/_embedding_lora_a_kernel/_embedding_lora_a_kernel.py create mode 100644 benchmarks/TritonNPUBench/sglang/_fwd_grouped_kernel_stage1_rope/_fwd_grouped_kernel_stage1_rope.py create mode 100644 benchmarks/TritonNPUBench/sglang/_fwd_kernel_ep_gather/_fwd_kernel_ep_gather.py create mode 100644 benchmarks/TritonNPUBench/sglang/_fwd_kernel_ep_scatter_1/_fwd_kernel_ep_scatter_1.py create mode 100644 benchmarks/TritonNPUBench/sglang/_fwd_kernel_ep_scatter_2/_fwd_kernel_ep_scatter_2.py create mode 100644 benchmarks/TritonNPUBench/sglang/_gate_up_lora_b_kernel/_gate_up_lora_b_kernel.py create mode 100644 benchmarks/TritonNPUBench/sglang/_layer_norm_fwd_1pass_kernel/_layer_norm_fwd_1pass_kernel.py create mode 100644 benchmarks/TritonNPUBench/sglang/_moe_sum_reduce_kernel/_moe_sum_reduce_kernel.py create mode 100644 benchmarks/TritonNPUBench/sglang/_qkv_lora_b_kernel/_qkv_lora_b_kernel.py create mode 100644 benchmarks/TritonNPUBench/sglang/_quantize_k_cache_fast_kernel/_quantize_k_cache_fast_kernel.py create mode 100644 benchmarks/TritonNPUBench/sglang/_selective_scan_update_kernel/_selective_scan_update_kernel.py create mode 100644 benchmarks/TritonNPUBench/sglang/_sgemm_lora_a_kernel/_sgemm_lora_a_kernel.py create mode 100644 benchmarks/TritonNPUBench/sglang/_sgemm_lora_b_kernel/_sgemm_lora_b_kernel.py create mode 100644 benchmarks/TritonNPUBench/sglang/_silu_and_mul_post_per_tensor_quant_kernel/_silu_and_mul_post_per_tensor_quant_kernel.py create mode 100644 benchmarks/TritonNPUBench/sglang/_silu_and_mul_post_quant_kernel/_silu_and_mul_post_quant_kernel.py create mode 100644 benchmarks/TritonNPUBench/sglang/_state_passing_fwd_kernel/_state_passing_fwd_kernel.py create mode 100644 benchmarks/TritonNPUBench/sglang/_tma_align_input_scale_kernel/_tma_align_input_scale_kernel.py create mode 100644 benchmarks/TritonNPUBench/sglang/act_and_mul_kernel/act_and_mul_kernel.py create mode 100644 benchmarks/TritonNPUBench/sglang/add_tree_reduce_u64_kernel/add_tree_reduce_u64_kernel.py create mode 100644 benchmarks/TritonNPUBench/sglang/chunk_fwd_kernel_o/chunk_fwd_kernel_o.py create mode 100644 benchmarks/TritonNPUBench/sglang/chunk_gated_delta_rule_fwd_kernel_h_blockdim64/chunk_gated_delta_rule_fwd_kernel_h_blockdim64.py create mode 100644 benchmarks/TritonNPUBench/sglang/chunk_kda_scaled_dot_kkt_fwd_kernel_intra_sub_inter/chunk_kda_scaled_dot_kkt_fwd_kernel_intra_sub_inter.py create mode 100644 benchmarks/TritonNPUBench/sglang/chunk_kda_scaled_dot_kkt_fwd_kernel_intra_sub_intra/chunk_kda_scaled_dot_kkt_fwd_kernel_intra_sub_intra.py create mode 100644 benchmarks/TritonNPUBench/sglang/chunk_local_cumsum_scalar_kernel/chunk_local_cumsum_scalar_kernel.py create mode 100644 benchmarks/TritonNPUBench/sglang/chunk_local_cumsum_vector_kernel/chunk_local_cumsum_vector_kernel.py create mode 100644 benchmarks/TritonNPUBench/sglang/chunk_scaled_dot_kkt_fwd_kernel/chunk_scaled_dot_kkt_fwd_kernel.py create mode 100644 benchmarks/TritonNPUBench/sglang/compute_identity_kernel/compute_identity_kernel.py create mode 100644 benchmarks/TritonNPUBench/sglang/compute_masked_m_triton_kernel/compute_masked_m_triton_kernel.py create mode 100644 benchmarks/TritonNPUBench/sglang/compute_problem_sizes_w4a8_kernel/compute_problem_sizes_w4a8_kernel.py create mode 100644 benchmarks/TritonNPUBench/sglang/compute_seg_indptr_triton_kernel/compute_seg_indptr_triton_kernel.py create mode 100644 benchmarks/TritonNPUBench/sglang/compute_src2dst_triton_kernel/compute_src2dst_triton_kernel.py create mode 100644 benchmarks/TritonNPUBench/sglang/concat_and_cast_mha_k_kernel/concat_and_cast_mha_k_kernel.py create mode 100644 benchmarks/TritonNPUBench/sglang/create_flashinfer_kv_indices_triton/create_flashinfer_kv_indices_triton.py create mode 100644 benchmarks/TritonNPUBench/sglang/create_flashmla_kv_indices_triton/create_flashmla_kv_indices_triton.py create mode 100644 benchmarks/TritonNPUBench/sglang/deepep_compute_src2dst_triton_kernel/deepep_compute_src2dst_triton_kernel.py create mode 100644 benchmarks/TritonNPUBench/sglang/deepep_permute_triton_kernel/deepep_permute_triton_kernel.py create mode 100644 benchmarks/TritonNPUBench/sglang/deepep_post_reorder_triton_kernel/deepep_post_reorder_triton_kernel.py create mode 100644 benchmarks/TritonNPUBench/sglang/deepgemm_compute_src2dst_triton_kernel/deepgemm_compute_src2dst_triton_kernel.py create mode 100644 benchmarks/TritonNPUBench/sglang/experts_combine_kernel/experts_combine_kernel.py create mode 100644 benchmarks/TritonNPUBench/sglang/fill_gateup_input_triton_kernel/fill_gateup_input_triton_kernel.py create mode 100644 benchmarks/TritonNPUBench/sglang/fn_triton_kernel/fn_triton_kernel.py create mode 100644 benchmarks/TritonNPUBench/sglang/fused_dual_residual_rmsnorm_kernel/fused_dual_residual_rmsnorm_kernel.py create mode 100644 benchmarks/TritonNPUBench/sglang/fused_gdn_gating_kernel/fused_gdn_gating_kernel.py create mode 100644 benchmarks/TritonNPUBench/sglang/fused_moe_kernel/fused_moe_kernel.py create mode 100644 benchmarks/TritonNPUBench/sglang/fused_moe_kernel_gptq_awq/fused_moe_kernel_gptq_awq.py create mode 100644 benchmarks/TritonNPUBench/sglang/fused_moe_router_cudacore_kernel/fused_moe_router_cudacore_kernel.py create mode 100644 benchmarks/TritonNPUBench/sglang/fused_moe_router_tensorcore_kernel/fused_moe_router_tensorcore_kernel.py create mode 100644 benchmarks/TritonNPUBench/sglang/fused_recurrent_gated_delta_rule_fwd_kernel/fused_recurrent_gated_delta_rule_fwd_kernel.py create mode 100644 benchmarks/TritonNPUBench/sglang/fused_recurrent_gated_delta_rule_update_fwd_kernel/fused_recurrent_gated_delta_rule_update_fwd_kernel.py create mode 100644 benchmarks/TritonNPUBench/sglang/fused_rmsnorm_kernel/fused_rmsnorm_kernel.py create mode 100644 benchmarks/TritonNPUBench/sglang/fused_sigmoid_gating_delta_rule_update_kernel/fused_sigmoid_gating_delta_rule_update_kernel.py create mode 100644 benchmarks/TritonNPUBench/sglang/fused_softcap_kernel/fused_softcap_kernel.py create mode 100644 benchmarks/TritonNPUBench/sglang/gelu_and_mul_kernel/gelu_and_mul_kernel.py create mode 100644 benchmarks/TritonNPUBench/sglang/get_last_loc_kernel/get_last_loc_kernel.py create mode 100644 benchmarks/TritonNPUBench/sglang/get_num_kv_splits_triton/get_num_kv_splits_triton.py create mode 100644 benchmarks/TritonNPUBench/sglang/hash_tiles32_kernel_blocked/hash_tiles32_kernel_blocked.py create mode 100644 benchmarks/TritonNPUBench/sglang/l2norm_fwd_kernel/l2norm_fwd_kernel.py create mode 100644 benchmarks/TritonNPUBench/sglang/l2norm_fwd_kernel1/l2norm_fwd_kernel1.py create mode 100644 benchmarks/TritonNPUBench/sglang/layer_norm_gated_fwd_kernel/layer_norm_gated_fwd_kernel.py create mode 100644 benchmarks/TritonNPUBench/sglang/layer_norm_gated_fwd_kernel1/layer_norm_gated_fwd_kernel1.py create mode 100644 benchmarks/TritonNPUBench/sglang/memcpy_triton_kernel/memcpy_triton_kernel.py create mode 100644 benchmarks/TritonNPUBench/sglang/merge_16x16_to_32x32_inverse_kernel/merge_16x16_to_32x32_inverse_kernel.py create mode 100644 benchmarks/TritonNPUBench/sglang/merge_16x16_to_64x64_inverse_kernel/merge_16x16_to_64x64_inverse_kernel.py create mode 100644 benchmarks/TritonNPUBench/sglang/merge_state_kernel/merge_state_kernel.py create mode 100644 benchmarks/TritonNPUBench/sglang/moe_align_block_size_stage1/moe_align_block_size_stage1.py create mode 100644 benchmarks/TritonNPUBench/sglang/moe_align_block_size_stage2/moe_align_block_size_stage2.py create mode 100644 benchmarks/TritonNPUBench/sglang/moe_align_block_size_stage3/moe_align_block_size_stage3.py create mode 100644 benchmarks/TritonNPUBench/sglang/moe_align_block_size_stage4/moe_align_block_size_stage4.py create mode 100644 benchmarks/TritonNPUBench/sglang/nsa_cp_round_robin_split_q_seqs_kernel/nsa_cp_round_robin_split_q_seqs_kernel.py create mode 100644 benchmarks/TritonNPUBench/sglang/pad_draft_extend_query_kernel/pad_draft_extend_query_kernel.py create mode 100644 benchmarks/TritonNPUBench/sglang/pad_sequence_with_mask_kernel/pad_sequence_with_mask_kernel.py create mode 100644 benchmarks/TritonNPUBench/sglang/post_reorder_triton_kernel/post_reorder_triton_kernel.py create mode 100644 benchmarks/TritonNPUBench/sglang/post_reorder_triton_kernel_for_cutlass_moe/post_reorder_triton_kernel_for_cutlass_moe.py create mode 100644 benchmarks/TritonNPUBench/sglang/pre_reorder_triton_kernel_for_cutlass_moe/pre_reorder_triton_kernel_for_cutlass_moe.py create mode 100644 benchmarks/TritonNPUBench/sglang/recompute_w_u_fwd_kernel/recompute_w_u_fwd_kernel.py create mode 100644 benchmarks/TritonNPUBench/sglang/rotl32_kernel/rotl32_kernel.py create mode 100644 benchmarks/TritonNPUBench/sglang/safe_exp_kernel/safe_exp_kernel.py create mode 100644 benchmarks/TritonNPUBench/sglang/sglang_gpu_perf.csv create mode 100644 benchmarks/TritonNPUBench/sglang/silu_and_mul_kernel/silu_and_mul_kernel.py create mode 100644 benchmarks/TritonNPUBench/sglang/silu_mul_static_tensorwise_quant_triton_kernel_for_cutlass_moe/silu_mul_static_tensorwise_quant_triton_kernel_for_cutlass_moe.py create mode 100644 benchmarks/TritonNPUBench/sglang/solve_tril_16x16_kernel/solve_tril_16x16_kernel.py create mode 100644 benchmarks/TritonNPUBench/sglang/tanh_kernel/tanh_kernel.py create mode 100644 benchmarks/TritonNPUBench/sglang/track_mamba_state_if_needed_kernel/track_mamba_state_if_needed_kernel.py create mode 100644 benchmarks/TritonNPUBench/sglang/transform_index_page_table_decode_kernel/transform_index_page_table_decode_kernel.py create mode 100644 benchmarks/TritonNPUBench/sglang/unpad_draft_extend_output_kernel/unpad_draft_extend_output_kernel.py create mode 100644 benchmarks/TritonNPUBench/sglang/write_req_to_token_pool_triton/write_req_to_token_pool_triton.py create mode 100644 benchmarks/TritonNPUBench/sglang/write_req_to_token_pool_triton_optimize/write_req_to_token_pool_triton_optimize.py diff --git a/benchmarks/TritonNPUBench/sglang/_act_quant_kernel/_act_quant_kernel.py b/benchmarks/TritonNPUBench/sglang/_act_quant_kernel/_act_quant_kernel.py new file mode 100644 index 00000000..551b1ef0 --- /dev/null +++ b/benchmarks/TritonNPUBench/sglang/_act_quant_kernel/_act_quant_kernel.py @@ -0,0 +1,86 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/sglang_operator_cases/newtest_cases/test__act_quant_kernel.py +# Main kernel: _act_quant_kernel +# PT file: test__act_quant_kernel_v2.pt + +import triton +import triton.language as tl + + +# === _act_quant_kernel === +@triton.jit +def _act_quant_kernel( + X_ptr, + Y_ptr, + S_ptr, + M, + N, + group_size: tl.constexpr, + round_scale: tl.constexpr, + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, +): + """ + Triton kernel for activation quantization. + + Each block processes BLOCK_M rows and group_size columns. + """ + # Get block IDs + pid_m = tl.program_id(0) + pid_n = tl.program_id(1) + + # FP8 constants + fp8_min = -448.0 + fp8_max = 448.0 + fp8_max_inv = 1.0 / fp8_max + + # Calculate row and column offsets + row_start = pid_m * BLOCK_M + col_start = pid_n * group_size + + # Create offset arrays + rows = row_start + tl.arange(0, BLOCK_M) + cols = col_start + tl.arange(0, BLOCK_N) + + # Mask for valid rows and columns + row_mask = rows < M + col_mask = cols < N + mask = row_mask[:, None] & col_mask[None, :] + + # Load input data + x_ptrs = X_ptr + rows[:, None] * N + cols[None, :] + x = tl.load(x_ptrs, mask=mask, other=0.0).to(tl.float32) + + # Compute absolute max along columns (group_size dimension) for each row + x_abs = tl.abs(x) + amax = tl.max(x_abs, axis=1) # Shape: (BLOCK_M,) + + # Clamp amax to avoid division by zero + amax = tl.maximum(amax, 1e-4) + + # Compute scale + if round_scale: + # Fast round scale using bit manipulation approximation + # This is a simplified version - the exact bit manipulation is harder in Triton + # Using log2 + ceil + pow2 as approximation + log_val = tl.log2(amax * fp8_max_inv) + log_ceil = tl.ceil(log_val) + scale = tl.exp2(log_ceil) + else: + scale = amax * fp8_max_inv + + # Quantize: y = clamp(x / scale, fp8_min, fp8_max) + scale_broadcast = scale[:, None] + y = x / scale_broadcast + y = tl.minimum(tl.maximum(y, fp8_min), fp8_max) + + # Store quantized output + y_ptrs = Y_ptr + rows[:, None] * N + cols[None, :] + tl.store(y_ptrs, y, mask=mask) + + # Store scales + s_cols = pid_n + s_ptrs = S_ptr + rows * (N // group_size) + s_cols + s_mask = row_mask + tl.store(s_ptrs, scale, mask=s_mask) + diff --git a/benchmarks/TritonNPUBench/sglang/_bmm_chunk_fwd_kernel/_bmm_chunk_fwd_kernel.py b/benchmarks/TritonNPUBench/sglang/_bmm_chunk_fwd_kernel/_bmm_chunk_fwd_kernel.py new file mode 100644 index 00000000..d911e0e8 --- /dev/null +++ b/benchmarks/TritonNPUBench/sglang/_bmm_chunk_fwd_kernel/_bmm_chunk_fwd_kernel.py @@ -0,0 +1,122 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/sglang_operator_cases/newtest_cases/test__bmm_chunk_fwd_kernel.py +# Main kernel: _bmm_chunk_fwd_kernel +# PT file: test__bmm_chunk_fwd_kernel_v2.pt + +import triton +import triton.language as tl + + +# === _bmm_chunk_fwd_kernel === +@triton.jit +def _bmm_chunk_fwd_kernel( + # Pointers to matrices + a_ptr, + b_ptr, + out_ptr, + seq_idx_ptr, + # Matrix dimensions + seqlen, + chunk_size, + K, + ngroups, + stride_a_batch, + stride_a_seqlen, + stride_a_head, + stride_ak, + stride_b_batch, + stride_b_seqlen, + stride_b_head, + stride_bk, + stride_out_batch, + stride_out_chunk, + stride_out_head, + stride_outm, + stride_outn, + stride_seq_idx_batch, + stride_seq_idx_seqlen, + # Meta-parameters + IS_CAUSAL: tl.constexpr, + dot_dtype: tl.constexpr, + HAS_SEQ_IDX: tl.constexpr, + BLOCK_SIZE_M: tl.constexpr = 16, + BLOCK_SIZE_N: tl.constexpr = 16, + BLOCK_SIZE_K: tl.constexpr = 16, +): + pid_b = tl.program_id(axis=1) + pid_ch = tl.program_id(axis=2).to(tl.int64) + pid_c = pid_ch // ngroups + pid_h = pid_ch - pid_c * ngroups + num_pid_n = tl.cdiv(chunk_size, BLOCK_SIZE_N) + pid_m = tl.program_id(axis=0) // num_pid_n + pid_n = tl.program_id(axis=0) % num_pid_n + if IS_CAUSAL: + if pid_n * BLOCK_SIZE_N >= (pid_m + 1) * BLOCK_SIZE_M: + return + a_ptr += ( + pid_b * stride_a_batch + + pid_c * chunk_size * stride_a_seqlen + + pid_h * stride_a_head + ) + b_ptr += ( + pid_b * stride_b_batch + + pid_c * chunk_size * stride_b_seqlen + + pid_h * stride_b_head + ) + if HAS_SEQ_IDX: + seq_idx_ptr += ( + pid_b * stride_seq_idx_batch + pid_c * chunk_size * stride_seq_idx_seqlen + ) + + offs_m = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) + offs_n = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) + offs_k = tl.arange(0, BLOCK_SIZE_K) + a_ptrs = a_ptr + (offs_m[:, None] * stride_a_seqlen + offs_k[None, :] * stride_ak) + b_ptrs = b_ptr + (offs_k[:, None] * stride_bk + offs_n[None, :] * stride_b_seqlen) + chunk_size_limit = min(chunk_size, seqlen - pid_c * chunk_size) + + acc = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32) + for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)): + a = tl.load( + a_ptrs, + mask=(offs_m[:, None] < chunk_size_limit) + & (offs_k[None, :] < K - k * BLOCK_SIZE_K), + other=0.0, + ).to(dot_dtype) + b = tl.load( + b_ptrs, + mask=(offs_k[:, None] < K - k * BLOCK_SIZE_K) + & (offs_n[None, :] < chunk_size_limit), + other=0.0, + ).to(dot_dtype) + acc += tl.dot(a, b) + a_ptrs += BLOCK_SIZE_K * stride_ak + b_ptrs += BLOCK_SIZE_K * stride_bk + + offs_m = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) + offs_n = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) + if HAS_SEQ_IDX: + chunk_size_limit = min(chunk_size, seqlen - pid_c * chunk_size) + seq_idx_m = tl.load( + seq_idx_ptr + offs_m * stride_seq_idx_seqlen, + mask=offs_m < chunk_size_limit, + other=-1, + ) + seq_idx_n = tl.load( + seq_idx_ptr + offs_n * stride_seq_idx_seqlen, + mask=offs_n < chunk_size_limit, + other=-2, + ) + acc = tl.where(seq_idx_m[:, None] == seq_idx_n[None, :], acc, 0.0) + out = acc.to(out_ptr.dtype.element_ty) + + out_ptr += ( + pid_b * stride_out_batch + pid_c * stride_out_chunk + pid_h * stride_out_head + ) + out_ptrs = out_ptr + (stride_outm * offs_m[:, None] + offs_n[None, :] * stride_outn) + tl.store( + out_ptrs, + out, + mask=(offs_m[:, None] < chunk_size) & (offs_n[None, :] < chunk_size), + ) + diff --git a/benchmarks/TritonNPUBench/sglang/_causal_conv1d_fwd_kernel/_causal_conv1d_fwd_kernel.py b/benchmarks/TritonNPUBench/sglang/_causal_conv1d_fwd_kernel/_causal_conv1d_fwd_kernel.py new file mode 100644 index 00000000..ca7c8470 --- /dev/null +++ b/benchmarks/TritonNPUBench/sglang/_causal_conv1d_fwd_kernel/_causal_conv1d_fwd_kernel.py @@ -0,0 +1,373 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/sglang_operator_cases/newtest_cases/test__causal_conv1d_fwd_kernel.py +# Main kernel: _causal_conv1d_fwd_kernel +# PT file: test__causal_conv1d_fwd_kernel_v2.pt + +import triton +import triton.language as tl + + +# === _causal_conv1d_fwd_kernel === +@triton.jit() +def _causal_conv1d_fwd_kernel( # continuous batching + # Pointers to matrices + x_ptr, # (dim, cu_seqlen) holding `batch` of actual sequences + padded sequences + w_ptr, # (dim, width) + bias_ptr, + initial_states_ptr, # conv_states_ptr + cache_indices_ptr, # conv_state_indices_ptr + has_initial_states_ptr, + query_start_loc_ptr, + o_ptr, # (dim, seqlen) - actually pointing to x_ptr + # Matrix dimensions + dim: tl.constexpr, + seqlen: tl.int32, # cu_seqlen + num_cache_lines: tl.constexpr, # added to support vLLM larger cache lines + # Strides + stride_x_seq: tl.constexpr, # stride to get to next sequence, + stride_x_dim: tl.constexpr, # stride to get to next feature-value, + stride_x_token: tl.constexpr, # stride to get to next token (same feature-index, same sequence-index) + stride_w_dim: tl.constexpr, # stride to get to next dim-axis value + stride_w_width: tl.constexpr, # stride to get to next width-axis value + stride_istate_seq: tl.constexpr, + stride_istate_dim: tl.constexpr, + stride_istate_token: tl.constexpr, + stride_o_seq: tl.constexpr, + stride_o_dim: tl.constexpr, + stride_o_token: tl.constexpr, + # others + pad_slot_id: tl.constexpr, + # Meta-parameters + HAS_BIAS: tl.constexpr, + KERNEL_WIDTH: tl.constexpr, + SILU_ACTIVATION: tl.constexpr, + HAS_INITIAL_STATES: tl.constexpr, + HAS_CACHE: tl.constexpr, + IS_CONTINUOUS_BATCHING: tl.constexpr, + USE_PAD_SLOT: tl.constexpr, + NP2_STATELEN: tl.constexpr, + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, +): + conv_states_ptr = initial_states_ptr + conv_state_indices_ptr = cache_indices_ptr + stride_conv_state_seq = stride_istate_seq + stride_conv_state_dim = stride_istate_dim + stride_conv_state_tok = stride_istate_token + state_len = ( + KERNEL_WIDTH - 1 + ) # can be passed via argument if it's not the same as this value + + # one program handles one chunk in a single sequence + # rather than mixing sequences - to make updating initial_states across sequences efficiently + + # single-sequence id + idx_seq = tl.program_id(0) + chunk_offset = tl.program_id(1) + + # BLOCK_N elements along the feature-dimension (channel) + idx_feats = tl.program_id(2) * BLOCK_N + tl.arange(0, BLOCK_N) + + if idx_seq == pad_slot_id: + return + + sequence_start_index = tl.load(query_start_loc_ptr + idx_seq) + sequence_end_index = tl.load(query_start_loc_ptr + idx_seq + 1) + # find the actual sequence length + seqlen = sequence_end_index - sequence_start_index + + token_offset = BLOCK_M * chunk_offset + segment_len = min(BLOCK_M, seqlen - token_offset) + + if segment_len <= 0: + return + + # base of the sequence + x_base = ( + x_ptr + sequence_start_index * stride_x_token + idx_feats * stride_x_dim + ) # [BLOCK_N,] + + if IS_CONTINUOUS_BATCHING: + # cache_idx + conv_state_batch_coord = tl.load(conv_state_indices_ptr + idx_seq).to(tl.int64) + else: + # cache_idx + conv_state_batch_coord = idx_seq + if USE_PAD_SLOT: # noqa + if conv_state_batch_coord == pad_slot_id: + # not processing as this is not the actual sequence + return + conv_states_base = ( + conv_states_ptr + + (conv_state_batch_coord * stride_conv_state_seq) + + (idx_feats * stride_conv_state_dim) + ) # [BLOCK_N,] + + w_base = w_ptr + (idx_feats * stride_w_dim) # [BLOCK_N,] + + # Does 2 things: + # 1. READ prior-block init-state data - [done by every Triton programs] + # 2. update conv_state with new data [only by the Triton program handles chunk_offset=0] + if chunk_offset == 0: + # read from conv_states + load_init_state = False + if HAS_INITIAL_STATES: # the new HAS_INITIAL_STATES + load_init_state = tl.load(has_initial_states_ptr + idx_seq).to(tl.int1) + if load_init_state: + # load from conv_states + prior_tokens = conv_states_base + (state_len - 1) * stride_conv_state_tok + mask_w = idx_feats < dim + if KERNEL_WIDTH == 2: + conv_states_ptrs = prior_tokens # [BLOCK_N] + col0 = tl.load(conv_states_ptrs, mask_w, 0.0) + if KERNEL_WIDTH == 3: + conv_states_ptrs = prior_tokens # [BLOCK_N] + col1 = tl.load(conv_states_ptrs, mask_w, 0.0) + conv_states_ptrs = prior_tokens - 1 * stride_conv_state_tok # [BLOCK_N] + col0 = tl.load(conv_states_ptrs, mask_w, 0.0) + if KERNEL_WIDTH == 4: + conv_states_ptrs = prior_tokens # [BLOCK_N] + col2 = tl.load(conv_states_ptrs, mask_w, 0.0) + conv_states_ptrs = prior_tokens - 1 * stride_conv_state_tok # [BLOCK_N] + col1 = tl.load(conv_states_ptrs, mask_w, 0.0) + conv_states_ptrs = prior_tokens - 2 * stride_conv_state_tok # [BLOCK_N] + col0 = tl.load(conv_states_ptrs, mask_w, 0.0) + if KERNEL_WIDTH == 5: + conv_states_ptrs = prior_tokens # [BLOCK_N] + col3 = tl.load(conv_states_ptrs, mask_w, 0.0) + conv_states_ptrs = prior_tokens - 1 * stride_conv_state_tok # [BLOCK_N] + col2 = tl.load(conv_states_ptrs, mask_w, 0.0) + conv_states_ptrs = prior_tokens - 2 * stride_conv_state_tok # [BLOCK_N] + col1 = tl.load(conv_states_ptrs, mask_w, 0.0) + conv_states_ptrs = prior_tokens - 3 * stride_conv_state_tok # [BLOCK_N] + col0 = tl.load(conv_states_ptrs, mask_w, 0.0) + else: + # prior-tokens are zeros + if KERNEL_WIDTH >= 2: # STRATEGY1 + # first chunk and does not have prior-token, so just set to 0 + col0 = tl.zeros((BLOCK_N,), dtype=x_ptr.dtype.element_ty) + if KERNEL_WIDTH >= 3: # STRATEGY1 + col1 = tl.zeros((BLOCK_N,), dtype=x_ptr.dtype.element_ty) + if KERNEL_WIDTH >= 4: # STRATEGY1 + col2 = tl.zeros((BLOCK_N,), dtype=x_ptr.dtype.element_ty) + if KERNEL_WIDTH >= 5: # STRATEGY1 + col3 = tl.zeros((BLOCK_N,), dtype=x_ptr.dtype.element_ty) + + # STEP 2: + # here prepare data for updating conv_state + if ( + state_len <= seqlen + ): # SMALL_CACHE=True (only move part of 'x' into conv_state cache) + # just read from 'x' + # copy 'x' data to conv_state + # load only 'x' data (and set 0 before 'x' if seqlen < state_len) + idx_tokens_last = (seqlen - state_len) + tl.arange( + 0, NP2_STATELEN + ) # [BLOCK_M] + x_ptrs = ( + x_ptr + + ((sequence_start_index + idx_tokens_last) * stride_x_token)[:, None] + + (idx_feats * stride_x_dim)[None, :] + ) # [BLOCK_M,BLOCK_N,] + mask_x = ( + (idx_tokens_last >= 0)[:, None] + & (idx_tokens_last < seqlen)[:, None] + & (idx_feats < dim)[None, :] + ) # token-index # token-index # feature-index + loaded_x = tl.load(x_ptrs, mask_x, 0.0) + new_conv_state = tl.load(x_ptrs, mask_x, 0.0) + idx_tokens_conv = tl.arange(0, NP2_STATELEN) # [BLOCK_M] + conv_states_ptrs_target = ( + conv_states_base[None, :] + + (idx_tokens_conv * stride_conv_state_tok)[:, None] + ) + + mask = (idx_tokens_conv < state_len)[:, None] & (idx_feats < dim)[None, :] + tl.debug_barrier() # NOTE: use this due to bug in Triton compiler + tl.store(conv_states_ptrs_target, new_conv_state, mask) + + else: + if load_init_state: + # update conv_state by shifting left, i.e. take last few cols from conv_state + cols from 'x' + idx_tokens_conv = tl.arange(0, NP2_STATELEN) # [BLOCK_M] + + conv_states_ptrs_source = ( + conv_states_ptr + + (conv_state_batch_coord * stride_conv_state_seq) + + (idx_feats * stride_conv_state_dim)[None, :] + + ((idx_tokens_conv + seqlen) * stride_conv_state_tok)[:, None] + ) # [BLOCK_M, BLOCK_N] + mask = ( + (conv_state_batch_coord < num_cache_lines) + & ((idx_tokens_conv + seqlen) < state_len)[:, None] + & (idx_feats < dim)[None, :] + ) + conv_state = tl.load(conv_states_ptrs_source, mask, other=0.0) + + VAL = state_len - seqlen + + x_ptrs = ( + x_base[None, :] + + ((idx_tokens_conv - VAL) * stride_x_token)[:, None] + ) # [BLOCK_M, BLOCK_N] + + mask_x = ( + (idx_tokens_conv - VAL >= 0)[:, None] + & (idx_tokens_conv - VAL < seqlen)[:, None] + & (idx_feats < dim)[None, :] + ) # token-index # token-index # feature-index + loaded_x = tl.load(x_ptrs, mask_x, 0.0) + + tl.debug_barrier() # need this due to the bug in tl.where not enforcing this when data is the result of another tl.load + new_conv_state = tl.where( + mask, conv_state, loaded_x + ) # BUG in 'tl.where' which requires a barrier before this + conv_states_ptrs_target = ( + conv_states_base + + (idx_tokens_conv * stride_conv_state_tok)[:, None] + ) # [BLOCK_M, BLOCK_N] + mask = (idx_tokens_conv < state_len)[:, None] & (idx_feats < dim)[ + None, : + ] + tl.store(conv_states_ptrs_target, new_conv_state, mask) + else: # load_init_state == False + # update conv_state by shifting left, BUT + # set cols prior to 'x' as zeros + cols from 'x' + idx_tokens_conv = tl.arange(0, NP2_STATELEN) # [BLOCK_M] + + VAL = state_len - seqlen + + x_ptrs = ( + x_base[None, :] + + ((idx_tokens_conv - VAL) * stride_x_token)[:, None] + ) # [BLOCK_M, BLOCK_N] + + mask_x = ( + (idx_tokens_conv - VAL >= 0)[:, None] + & (idx_tokens_conv - VAL < seqlen)[:, None] + & (idx_feats < dim)[None, :] + ) # token-index # token-index # feature-index + new_conv_state = tl.load(x_ptrs, mask_x, 0.0) + + conv_states_ptrs_target = ( + conv_states_base + + (idx_tokens_conv * stride_conv_state_tok)[:, None] + ) # [BLOCK_M, BLOCK_N] + mask = (idx_tokens_conv < state_len)[:, None] & (idx_feats < dim)[ + None, : + ] + tl.store(conv_states_ptrs_target, new_conv_state, mask) + + else: # chunk_offset > 0 + # read prior-token data from `x` + load_init_state = True + prior_tokens = x_base + (token_offset - 1) * stride_x_token + mask_w = idx_feats < dim + if KERNEL_WIDTH == 2: + conv_states_ptrs = prior_tokens # [BLOCK_N] + col0 = tl.load(conv_states_ptrs, mask_w, 0.0, cache_modifier=".ca") + if KERNEL_WIDTH == 3: + conv_states_ptrs = prior_tokens # [BLOCK_N] + col1 = tl.load(conv_states_ptrs, mask_w, 0.0, cache_modifier=".ca") + conv_states_ptrs = prior_tokens - 1 * stride_x_token # [BLOCK_N] + col0 = tl.load(conv_states_ptrs, mask_w, 0.0, cache_modifier=".ca") + if KERNEL_WIDTH == 4: + conv_states_ptrs = prior_tokens # [BLOCK_N] + col2 = tl.load(conv_states_ptrs, mask_w, 0.0, cache_modifier=".ca") + conv_states_ptrs = prior_tokens - 1 * stride_x_token # [BLOCK_N] + col1 = tl.load(conv_states_ptrs, mask_w, 0.0, cache_modifier=".ca") + conv_states_ptrs = prior_tokens - 2 * stride_x_token # [BLOCK_N] + col0 = tl.load(conv_states_ptrs, mask_w, 0.0, cache_modifier=".ca") + if KERNEL_WIDTH == 5: + # ruff: noqa: F841 + conv_states_ptrs = prior_tokens # [BLOCK_N] + col3 = tl.load(conv_states_ptrs, mask_w, 0.0, cache_modifier=".ca") + conv_states_ptrs = prior_tokens - 1 * stride_x_token # [BLOCK_N] + col2 = tl.load(conv_states_ptrs, mask_w, 0.0, cache_modifier=".ca") + conv_states_ptrs = prior_tokens - 2 * stride_x_token # [BLOCK_N] + col1 = tl.load(conv_states_ptrs, mask_w, 0.0, cache_modifier=".ca") + conv_states_ptrs = prior_tokens - 3 * stride_x_token # [BLOCK_N] + col0 = tl.load(conv_states_ptrs, mask_w, 0.0, cache_modifier=".ca") + + if HAS_BIAS: + bias = bias_ptr + idx_feats + mask_bias = idx_feats < dim + acc_preload = tl.load(bias, mask=mask_bias, other=0.0).to( + tl.float32 + ) # [BLOCK_N] + else: + acc_preload = tl.zeros((BLOCK_N,), dtype=tl.float32) + + x_base_1d = x_base + token_offset * stride_x_token # starting of chunk + + # PRE-LOAD WEIGHTS + mask_w = idx_feats < dim + if KERNEL_WIDTH >= 2: + w_ptrs = w_base + (0 * stride_w_width) # [BLOCK_N] tensor + w_col0 = tl.load(w_ptrs, mask_w, other=0.0) + w_ptrs = w_base + (1 * stride_w_width) # [BLOCK_N] tensor + w_col1 = tl.load(w_ptrs, mask_w, other=0.0) + if KERNEL_WIDTH >= 3: + w_ptrs = w_base + (2 * stride_w_width) # [BLOCK_N] tensor + w_col2 = tl.load(w_ptrs, mask_w, other=0.0) + if KERNEL_WIDTH >= 4: + w_ptrs = w_base + (3 * stride_w_width) # [BLOCK_N] tensor + w_col3 = tl.load(w_ptrs, mask_w, other=0.0) + mask_x_1d = idx_feats < dim + for idx_token in range(segment_len): + acc = acc_preload + + matrix_w = w_col0 + matrix_x = col0 + for j in tl.static_range(KERNEL_WIDTH): + + if KERNEL_WIDTH == 2: + if j == 1: # KERNEL_WIDTH-1: + matrix_w = w_col1 + x_ptrs_1d = x_base_1d + idx_token * stride_x_token # [BLOCK_N] + matrix_x = tl.load(x_ptrs_1d, mask=mask_x_1d) + elif KERNEL_WIDTH == 3: + if j == 1: + matrix_w = w_col1 + matrix_x = col1 + elif j == 2: + matrix_w = w_col2 + x_ptrs_1d = x_base_1d + idx_token * stride_x_token # [BLOCK_N] + matrix_x = tl.load(x_ptrs_1d, mask=mask_x_1d) + elif KERNEL_WIDTH == 4: + if j == 1: + matrix_w = w_col1 + matrix_x = col1 + elif j == 2: + matrix_w = w_col2 + matrix_x = col2 + elif j == 3: + matrix_w = w_col3 + x_ptrs_1d = x_base_1d + idx_token * stride_x_token # [BLOCK_N] + matrix_x = tl.load(x_ptrs_1d, mask=mask_x_1d) + + acc += matrix_x * matrix_w # [BLOCK_N] + + if KERNEL_WIDTH == 2: + col0 = matrix_x + elif KERNEL_WIDTH == 3: + col0 = col1 + col1 = matrix_x + elif KERNEL_WIDTH == 4: + col0 = col1 + col1 = col2 + col2 = matrix_x + + if SILU_ACTIVATION: + acc = acc / (1 + tl.exp(-acc)) + mask_1d = (idx_token < segment_len) & ( + idx_feats < dim + ) # token-index # feature-index + o_ptrs = ( + o_ptr + + (sequence_start_index + token_offset + idx_token) * stride_o_token + + (idx_feats * stride_o_dim) + ) + + tl.store(o_ptrs, acc, mask=mask_1d) + diff --git a/benchmarks/TritonNPUBench/sglang/_chunk_cumsum_fwd_kernel/_chunk_cumsum_fwd_kernel.py b/benchmarks/TritonNPUBench/sglang/_chunk_cumsum_fwd_kernel/_chunk_cumsum_fwd_kernel.py new file mode 100644 index 00000000..283145fb --- /dev/null +++ b/benchmarks/TritonNPUBench/sglang/_chunk_cumsum_fwd_kernel/_chunk_cumsum_fwd_kernel.py @@ -0,0 +1,105 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/sglang_operator_cases/newtest_cases/test__chunk_cumsum_fwd_kernel.py +# Main kernel: _chunk_cumsum_fwd_kernel +# PT file: test__chunk_cumsum_fwd_kernel_v2.pt + +import triton +import triton.language as tl + +@triton.jit +def softplus(dt): + dt = tl.where(dt <= 20.0, tl.math.log(tl.math.exp(dt) + 1), dt) + return dt + +# === _chunk_cumsum_fwd_kernel === +@triton.jit +def _chunk_cumsum_fwd_kernel( + # Pointers to matrices + dt_ptr, + A_ptr, + dt_bias_ptr, + dt_out_ptr, + dA_cumsum_ptr, + # Matrix dimension + batch, + seqlen, + nheads, + chunk_size, + dt_min, + dt_max, + # Strides + stride_dt_batch, + stride_dt_seqlen, + stride_dt_head, + stride_A_head, + stride_dt_bias_head, + stride_dt_out_batch, + stride_dt_out_chunk, + stride_dt_out_head, + stride_dt_out_csize, + stride_dA_cs_batch, + stride_dA_cs_chunk, + stride_dA_cs_head, + stride_dA_cs_csize, + # Meta-parameters + DT_SOFTPLUS: tl.constexpr, + HAS_DT_BIAS: tl.constexpr, + BLOCK_SIZE_CHUNK: tl.constexpr, + BLOCK_SIZE_H: tl.constexpr = 16, +): + pid_b = tl.program_id(axis=0) + + # if dt is long, may cause problems, so use 64 bit + # https://github.com/triton-lang/triton/issues/1058 + pid_c = tl.program_id(axis=1).to(tl.int64) + pid_h = tl.program_id(axis=2) + dt_ptr += pid_b * stride_dt_batch + pid_c * chunk_size * stride_dt_seqlen + dt_out_ptr += pid_b * stride_dt_out_batch + pid_c * stride_dt_out_chunk + dA_cumsum_ptr += pid_b * stride_dA_cs_batch + pid_c * stride_dA_cs_chunk + + offs_h = pid_h * BLOCK_SIZE_H + tl.arange(0, BLOCK_SIZE_H) + offs_c = tl.arange(0, BLOCK_SIZE_CHUNK) + dt_ptrs = dt_ptr + ( + offs_h[:, None] * stride_dt_head + offs_c[None, :] * stride_dt_seqlen + ) + A_ptrs = A_ptr + offs_h * stride_A_head + dt_out_ptrs = dt_out_ptr + ( + offs_h[:, None] * stride_dt_out_head + offs_c[None, :] * stride_dt_out_csize + ) + dA_cs_ptrs = dA_cumsum_ptr + ( + offs_h[:, None] * stride_dA_cs_head + offs_c[None, :] * stride_dA_cs_csize + ) + chunk_size_limit = min(chunk_size, seqlen - pid_c * chunk_size) + + dt = tl.load( + dt_ptrs, + mask=(offs_h[:, None] < nheads) & (offs_c[None, :] < chunk_size_limit), + other=0.0, + ).to(tl.float32) + if HAS_DT_BIAS: + dt_bias = tl.load( + dt_bias_ptr + offs_h * stride_dt_bias_head, mask=offs_h < nheads, other=0.0 + ).to(tl.float32) + dt += dt_bias[:, None] + if DT_SOFTPLUS: + dt = tl.where(dt <= 20.0, softplus(dt), dt) + # As of Triton 2.2.0, tl.clamp is not available yet + # dt = tl.clamp(dt, dt_min, dt_max) + dt = tl.minimum(tl.maximum(dt, dt_min), dt_max) + dt = tl.where( + (offs_h[:, None] < nheads) & (offs_c[None, :] < chunk_size_limit), dt, 0.0 + ) + tl.store( + dt_out_ptrs, + dt, + mask=(offs_h[:, None] < nheads) & (offs_c[None, :] < chunk_size), + ) + A = tl.load(A_ptrs, mask=offs_h < nheads, other=0.0).to(tl.float32) + dA = dt * A[:, None] + dA_cs = tl.cumsum(dA, axis=1) + tl.store( + dA_cs_ptrs, + dA_cs, + mask=(offs_h[:, None] < nheads) & (offs_c[None, :] < chunk_size), + ) + diff --git a/benchmarks/TritonNPUBench/sglang/_chunk_scan_fwd_kernel/_chunk_scan_fwd_kernel.py b/benchmarks/TritonNPUBench/sglang/_chunk_scan_fwd_kernel/_chunk_scan_fwd_kernel.py new file mode 100644 index 00000000..043ddb40 --- /dev/null +++ b/benchmarks/TritonNPUBench/sglang/_chunk_scan_fwd_kernel/_chunk_scan_fwd_kernel.py @@ -0,0 +1,412 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/sglang_operator_cases/newtest_cases/test__chunk_scan_fwd_kernel.py +# Main kernel: _chunk_scan_fwd_kernel +# PT file: test__chunk_scan_fwd_kernel_v2.pt + +import triton +import triton.language as tl + + +# === _chunk_scan_fwd_kernel === +@triton.jit +def _chunk_scan_fwd_kernel( + # Pointers to matrices + cb_ptr, + x_ptr, + z_ptr, + out_ptr, + out_x_ptr, + dt_ptr, + dA_cumsum_ptr, + seq_idx_ptr, + C_ptr, + states_ptr, + D_ptr, + initstates_ptr, + chunk_indices_ptr, + chunk_offsets_ptr, + chunk_meta_num, + # Matrix dimensions + chunk_size, + hdim, + dstate, + batch, + seqlen, + nheads_ngroups_ratio, + # Strides + stride_cb_batch, + stride_cb_chunk, + stride_cb_head, + stride_cb_csize_m, + stride_cb_csize_k, + stride_x_batch, + stride_x_seqlen, + stride_x_head, + stride_x_hdim, + stride_z_batch, + stride_z_seqlen, + stride_z_head, + stride_z_hdim, + stride_out_batch, + stride_out_seqlen, + stride_out_head, + stride_out_hdim, + stride_dt_batch, + stride_dt_chunk, + stride_dt_head, + stride_dt_csize, + stride_dA_cs_batch, + stride_dA_cs_chunk, + stride_dA_cs_head, + stride_dA_cs_csize, + stride_seq_idx_batch, + stride_seq_idx_seqlen, + stride_C_batch, + stride_C_seqlen, + stride_C_head, + stride_C_dstate, + stride_states_batch, + stride_states_chunk, + stride_states_head, + stride_states_hdim, + stride_states_dstate, + stride_init_states_batch, + stride_init_states_head, + stride_init_states_hdim, + stride_init_states_dstate, + stride_D_head, + # Meta-parameters + IS_CAUSAL: tl.constexpr, + HAS_D: tl.constexpr, + D_HAS_HDIM: tl.constexpr, + HAS_Z: tl.constexpr, + HAS_SEQ_IDX: tl.constexpr, + BLOCK_SIZE_DSTATE: tl.constexpr, + IS_TRITON_22: tl.constexpr, + HAS_INITSTATES: tl.constexpr, + BLOCK_SIZE_M: tl.constexpr = 16, + BLOCK_SIZE_N: tl.constexpr = 16, + BLOCK_SIZE_K: tl.constexpr = 16, +): + pid_bc = tl.program_id(axis=1).to(tl.int64) + pid_c = pid_bc // batch + pid_b = pid_bc - pid_c * batch + if not HAS_INITSTATES: + c_idx = pid_c + c_off = 0 + else: + c_idx = tl.load(chunk_indices_ptr + pid_c, mask=pid_c > -1, other=0) + c_off = tl.load(chunk_offsets_ptr + pid_c, mask=pid_c > -1, other=0) + + pid_h = tl.program_id(axis=2) + num_pid_n = tl.cdiv(hdim, BLOCK_SIZE_N) + pid_m = tl.program_id(axis=0) // num_pid_n + pid_n = tl.program_id(axis=0) % num_pid_n + cb_ptr += ( + pid_b * stride_cb_batch + + c_idx * stride_cb_chunk + + (pid_h // nheads_ngroups_ratio) * stride_cb_head + ) + x_ptr += ( + pid_b * stride_x_batch + + c_idx * chunk_size * stride_x_seqlen + + pid_h * stride_x_head + ) + dt_ptr += pid_b * stride_dt_batch + c_idx * stride_dt_chunk + pid_h * stride_dt_head + dA_cumsum_ptr += ( + pid_b * stride_dA_cs_batch + + c_idx * stride_dA_cs_chunk + + pid_h * stride_dA_cs_head + ) + C_ptr += ( + pid_b * stride_C_batch + + c_idx * chunk_size * stride_C_seqlen + + (pid_h // nheads_ngroups_ratio) * stride_C_head + ) + + # M-block offsets and prev states + # - logic in next block may override these if there is an active offset + offs_m = pid_m * BLOCK_SIZE_M + c_off + tl.arange(0, BLOCK_SIZE_M) + prev_states_ptr = ( + states_ptr + + pid_b * stride_states_batch + + c_idx * stride_states_chunk + + pid_h * stride_states_head + ) + prev_states_hdim = stride_states_hdim + prev_states_dstate = stride_states_dstate + + chunk_size_limit = min(chunk_size, seqlen - c_idx * chunk_size) + if HAS_SEQ_IDX: + seq_idx_ptr += ( + pid_b * stride_seq_idx_batch + c_idx * chunk_size * stride_seq_idx_seqlen + ) + + # - we only need seq_idx_prev to be aligned to chunk boundary + seq_idx_prev = tl.load( + seq_idx_ptr - stride_seq_idx_seqlen, mask=c_idx >= 1, other=0 + ) + + if HAS_INITSTATES: + # if there are init states, we only need seq_idx_m to point + # what is the current seq_idx + + # get current seq idx + if (pid_m * BLOCK_SIZE_M + c_off) < chunk_size_limit: + seq_idx_m = tl.load( + seq_idx_ptr + + (pid_m * BLOCK_SIZE_M + c_off) * stride_seq_idx_seqlen, + ) + + # - recall that in ssd_state_passing, for the case c_off == 0 + # i.e., the very first sequence, we made states_ptr hold its initial state + # so this edge case is taken care of + if ( + (c_off == 0) + and ( + seq_idx_prev != seq_idx_m + ) # if a seq is changed exactly on boundary + or (c_off > 0) # implies a new example (pseudo chunk) + ): + + # - replace prev_states_ptr with init_states + prev_states_ptr = ( + initstates_ptr + + seq_idx_m * stride_init_states_batch + + pid_h * stride_init_states_head + ) + prev_states_hdim = stride_init_states_hdim # override strides + prev_states_dstate = stride_init_states_dstate + + offs_n = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) + dA_cs_m = tl.load( + dA_cumsum_ptr + offs_m * stride_dA_cs_csize, mask=offs_m < chunk_size, other=0.0 + ).to(tl.float32) + + # - handle chunk state limit + if HAS_INITSTATES: + + # have to split this if otherwise compilation will have problems + dA_cs_m_boundary = 0.0 + + # get the c_idx for the next (logica) chunk + c_idx_n = tl.load( + chunk_indices_ptr + (pid_c + 1), + mask=pid_c > -1 and (pid_c + 1) < chunk_meta_num, + other=-1, # to trigger different chunk + ) + + # - there are things to consider + # A. if c_off > 0 then we need to move the dA_cs boundary to ensure correct + # contribution of past states + # B. if c_off_n < chunk_size_limit, then we need to adjust this so as not to + # encroach into the next sequence, where c_off_n is the offset of the next + # (logical) chunk. + # An equivalent check for B is c_idx == c_idx_n, where there is repetition in + # (logical) chunk indices. + + if (c_idx == c_idx_n) or c_off > 0: + + # get the next offset + c_off_n = tl.load( + chunk_offsets_ptr + (pid_c + 1), + mask=pid_c > -1 and (pid_c + 1) < chunk_meta_num, + other=chunk_size, + ) + + # in this case, adjust down the chunk_size_limit + if c_idx == c_idx_n: + chunk_size_limit = min(c_off_n, chunk_size_limit) + + # get the cs at the offset boundary + # - c_off == 0 is a passthrough + # - We need dA_cs at the boundary, defined by c_off - no need + # to increase pointer by pid_m (it is a constant offset, + # i.e. the same for all blocks) + dA_cs_m_boundary = tl.load( + dA_cumsum_ptr + (c_off - 1) * stride_dA_cs_csize, + mask=(((c_off - 1) > -1) and ((c_off) < chunk_size)), + other=0.0, + ).to(tl.float32) + + if HAS_SEQ_IDX: + # - handle seq idx when HAS_INITSTATES==False + if not HAS_INITSTATES: + seq_idx_m = tl.load( + seq_idx_ptr + offs_m * stride_seq_idx_seqlen, + mask=offs_m < chunk_size_limit, + other=-1, + ) + + acc = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32) + + # Without the if (pid_c > -1), with Triton 2.1.0, I get + # Assertion `!(srcMmaLayout && dstMmaLayout) && "Unexpected mma -> mm a layout conversion"' failed. + # With Triton 2.2.0, this works + if IS_TRITON_22 or c_idx > -1: + # Faster to just do 1 iteration with larger BLOCK_SIZE_K, up to block size 128 + offs_k_dstate = tl.arange( + 0, BLOCK_SIZE_DSTATE if BLOCK_SIZE_DSTATE <= 128 else BLOCK_SIZE_K + ) + C_ptrs = C_ptr + ( + offs_m[:, None] * stride_C_seqlen + offs_k_dstate[None, :] * stride_C_dstate + ) + + prev_states_ptrs = prev_states_ptr + ( + offs_n[None, :] * prev_states_hdim + + offs_k_dstate[:, None] * prev_states_dstate + ) + if HAS_SEQ_IDX: + + if not HAS_INITSTATES: + # - this is for continuous batching where there is no init states + scale_m = tl.where(seq_idx_m == seq_idx_prev, tl.exp(dA_cs_m), 0.0) + else: + # - if there is initstates, we will rely on prev_states, no zeroing + # required. + scale_m = tl.exp(dA_cs_m - dA_cs_m_boundary) + else: + scale_m = tl.exp(dA_cs_m) + if BLOCK_SIZE_DSTATE <= 128: + C = tl.load( + C_ptrs, + mask=(offs_m[:, None] < chunk_size_limit) + & (offs_k_dstate[None, :] < dstate), + other=0.0, + ) + + prev_states = tl.load( + prev_states_ptrs, + mask=(offs_k_dstate[:, None] < dstate) & (offs_n[None, :] < hdim), + other=0.0, + ) + prev_states = prev_states.to(C_ptr.dtype.element_ty) + acc = tl.dot(C, prev_states) * scale_m[:, None] + else: + for k in range(0, dstate, BLOCK_SIZE_K): + C = tl.load( + C_ptrs, + mask=(offs_m[:, None] < chunk_size_limit) + & (offs_k_dstate[None, :] < dstate - k), + other=0.0, + ) + # C = (C * scale_m[:, None]).to(C_ptr.dtype.element_ty) + prev_states = tl.load( + prev_states_ptrs, + mask=(offs_k_dstate[:, None] < dstate - k) + & (offs_n[None, :] < hdim), + other=0.0, + ) + prev_states = prev_states.to(C_ptr.dtype.element_ty) + acc += tl.dot(C, prev_states) + C_ptrs += BLOCK_SIZE_K + prev_states_ptrs += BLOCK_SIZE_K + acc *= scale_m[:, None] + + offs_k = tl.arange(0, BLOCK_SIZE_K) + c_off + cb_ptrs = cb_ptr + ( + offs_m[:, None] * stride_cb_csize_m + offs_k[None, :] * stride_cb_csize_k + ) + x_ptrs = x_ptr + ( + offs_k[:, None] * stride_x_seqlen + offs_n[None, :] * stride_x_hdim + ) + dt_ptrs = dt_ptr + offs_k * stride_dt_csize + dA_cumsum_ptrs = dA_cumsum_ptr + offs_k * stride_dA_cs_csize + K_MAX = ( + chunk_size_limit + if not IS_CAUSAL + else min((pid_m + 1) * BLOCK_SIZE_M, chunk_size_limit) + ) + for k in range(0, K_MAX, BLOCK_SIZE_K): + cb = tl.load( + cb_ptrs, + mask=(offs_m[:, None] < chunk_size) & (offs_k[None, :] < chunk_size - k), + other=0.0, + ).to(tl.float32) + dA_cs_k = tl.load(dA_cumsum_ptrs, mask=offs_k < chunk_size - k, other=0.0).to( + tl.float32 + ) + # If there's seq_idx, we already set cb[i, j] = 0 for seq_idx[i] != seq_idx[j]. + # So we don't need masking wrt seq_idx here. + cb *= tl.exp(dA_cs_m[:, None] - dA_cs_k[None, :]) + dt_k = tl.load(dt_ptrs, mask=offs_k < chunk_size - k, other=0.0).to(tl.float32) + cb *= dt_k + if IS_CAUSAL: + mask = offs_m[:, None] >= k + offs_k[None, :] + cb = tl.where(mask, cb, 0.0) + cb = cb.to(x_ptr.dtype.element_ty) + x = tl.load( + x_ptrs, + mask=(offs_k[:, None] < chunk_size_limit - k) & (offs_n[None, :] < hdim), + other=0.0, + ) + acc += tl.dot(cb, x) + cb_ptrs += BLOCK_SIZE_K * stride_cb_csize_k + x_ptrs += BLOCK_SIZE_K * stride_x_seqlen + dt_ptrs += BLOCK_SIZE_K * stride_dt_csize + dA_cumsum_ptrs += BLOCK_SIZE_K * stride_dA_cs_csize + + offs_out_m = pid_m * BLOCK_SIZE_M + c_off + tl.arange(0, BLOCK_SIZE_M) + offs_out_n = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) + + if HAS_D: + if D_HAS_HDIM: + D = tl.load( + D_ptr + pid_h * stride_D_head + offs_n, mask=offs_n < hdim, other=0.0 + ).to(tl.float32) + else: + D = tl.load(D_ptr + pid_h * stride_D_head).to(tl.float32) + x_residual = tl.load( + x_ptr + + (offs_m[:, None] * stride_x_seqlen + offs_n[None, :] * stride_x_hdim), + mask=(offs_m[:, None] < chunk_size_limit) & (offs_n[None, :] < hdim), + other=0.0, + ).to(tl.float32) + acc += x_residual * D + + if HAS_Z: + out_x_ptr += ( + pid_b * stride_out_batch + + c_idx * chunk_size * stride_out_seqlen + + pid_h * stride_out_head + ) + out_x_ptrs = out_x_ptr + ( + stride_out_seqlen * offs_out_m[:, None] + offs_out_n[None, :] + ) + tl.store( + out_x_ptrs, + acc, + mask=(offs_out_m[:, None] < chunk_size_limit) + & (offs_out_n[None, :] < hdim), + ) + + z_ptr += ( + pid_b * stride_z_batch + + c_idx * chunk_size * stride_z_seqlen + + pid_h * stride_z_head + ) + z_ptrs = z_ptr + ( + stride_z_seqlen * offs_out_m[:, None] + stride_z_hdim * offs_out_n[None, :] + ) + z = tl.load( + z_ptrs, + mask=(offs_out_m[:, None] < chunk_size_limit) + & (offs_out_n[None, :] < hdim), + other=0.0, + ).to(tl.float32) + acc *= z * tl.sigmoid(z) + + out_ptr += ( + pid_b * stride_out_batch + + c_idx * chunk_size * stride_out_seqlen + + pid_h * stride_out_head + ) + out_ptrs = out_ptr + ( + stride_out_seqlen * offs_out_m[:, None] + offs_out_n[None, :] * stride_out_hdim + ) + tl.store( + out_ptrs, + acc, + mask=(offs_out_m[:, None] < chunk_size_limit) & (offs_out_n[None, :] < hdim), + ) + diff --git a/benchmarks/TritonNPUBench/sglang/_chunked_lora_expand_kernel/_chunked_lora_expand_kernel.py b/benchmarks/TritonNPUBench/sglang/_chunked_lora_expand_kernel/_chunked_lora_expand_kernel.py new file mode 100644 index 00000000..80fff9b1 --- /dev/null +++ b/benchmarks/TritonNPUBench/sglang/_chunked_lora_expand_kernel/_chunked_lora_expand_kernel.py @@ -0,0 +1,142 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/sglang_operator_cases/newtest_cases/test__chunked_lora_expand_kernel.py +# Main kernel: _chunked_lora_expand_kernel +# PT file: test__chunked_lora_expand_kernel_uint2_v2.pt + +import triton +import triton.language as tl + + +# === _chunked_lora_expand_kernel === +@triton.jit(do_not_specialize=["num_segs"]) +def _chunked_lora_expand_kernel( + # Pointers to matrices + x, + weights, + output, + # Information on sequence lengths and weight id + seg_indptr, + weight_indices, + lora_ranks, + permutation, + num_segs, + # For fused output scaling + scalings, + # Offsets of q/k/v slice on output dimension + slice_offsets, + # Meta parameters + NUM_SLICES: tl.constexpr, + OUTPUT_DIM: tl.constexpr, + MAX_RANK: tl.constexpr, # K = R + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, + BLOCK_K: tl.constexpr, +): + """ + Computes a chunked SGMV for LoRA expand operations. + + When a sequence's rank is 0, the kernel is essentially a no-op, following + the convention in pytorch where the product of two matrices of shape (m, 0) + and (0, n) is an all-zero matrix of shape (m, n). + + Args: + x (Tensor): The input tensor, which is the result of the LoRA A projection. + Shape: (s, num_slices * K), where s is the sum of all sequence lengths in the + batch and K is the maximum LoRA rank. + weights (Tensor): The LoRA B weights for all adapters. + Shape: (num_lora, output_dim, K). + output (Tensor): The output tensor where the result is stored. + Shape: (s, output_dim). + """ + tl.static_assert(NUM_SLICES <= 3) + + x_stride_0: tl.constexpr = NUM_SLICES * MAX_RANK + x_stride_1: tl.constexpr = 1 + + w_stride_0: tl.constexpr = OUTPUT_DIM * MAX_RANK + w_stride_1: tl.constexpr = MAX_RANK + w_stride_2: tl.constexpr = 1 + + output_stride_0: tl.constexpr = OUTPUT_DIM + output_stride_1: tl.constexpr = 1 + + pid_s = tl.program_id(axis=2) + if pid_s >= num_segs: + return + + # Current block computes sequence with batch_id, + # which starts from row seg_start of x with length seg_len. + # qkv_id decides which of q,k,v to compute (0: q, 1: k, 2: v) + w_index = tl.load(weight_indices + pid_s) + cur_rank = tl.load(lora_ranks + w_index) + + # If rank is 0, this kernel is a no-op. + if cur_rank == 0: + return + + seg_start = tl.load(seg_indptr + pid_s) + seg_end = tl.load(seg_indptr + pid_s + 1) + + slice_id = tl.program_id(axis=1) + slice_start = tl.load(slice_offsets + slice_id) + slice_end = tl.load(slice_offsets + slice_id + 1) + + scaling = tl.load(scalings + w_index) + # Adjust K (rank) according to the specific LoRA adapter + cur_rank = tl.minimum(MAX_RANK, cur_rank) + + # Map logical sequence index to physical index + s_offset_logical = tl.arange(0, BLOCK_M) + seg_start + s_offset_physical = tl.load( + permutation + s_offset_logical, mask=s_offset_logical < seg_end + ) + + # Create pointers for the first block of x and weights[batch_id][n_start: n_end][:] + # The pointers will be advanced as we move in the K direction + # and accumulate + pid_n = tl.program_id(axis=0) + n_offset = tl.arange(0, BLOCK_N) + pid_n * BLOCK_N + slice_start + k_offset = tl.arange(0, BLOCK_K) + + x_ptrs = ( + x + + slice_id * cur_rank * x_stride_1 + + (s_offset_physical[:, None] * x_stride_0 + k_offset[None, :] * x_stride_1) + ) + w_ptrs = (weights + w_index * w_stride_0) + ( + k_offset[:, None] * w_stride_2 + n_offset[None, :] * w_stride_1 + ) + + # Iterate to compute the block in output matrix + partial_sum = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32) + for k in range(0, tl.cdiv(cur_rank, BLOCK_K)): + x_tile = tl.load( + x_ptrs, + mask=(s_offset_logical[:, None] < seg_end) + & (k_offset[None, :] < cur_rank - k * BLOCK_K), + other=0.0, + ) + w_tile = tl.load( + w_ptrs, + mask=(k_offset[:, None] < cur_rank - k * BLOCK_K) + & (n_offset[None, :] < slice_end), + other=0.0, + ) + partial_sum += tl.dot(x_tile, w_tile) + + x_ptrs += BLOCK_K * x_stride_1 + w_ptrs += BLOCK_K * w_stride_2 + + # Store result to output matrix + partial_sum *= scaling + partial_sum = partial_sum.to(x.dtype.element_ty) + output_ptr = output + ( + s_offset_physical[:, None] * output_stride_0 + + n_offset[None, :] * output_stride_1 + ) + output_mask = (s_offset_logical[:, None] < seg_end) & ( + n_offset[None, :] < slice_end + ) + partial_sum += tl.load(output_ptr, mask=output_mask, other=0.0) + tl.store(output_ptr, partial_sum, mask=output_mask) + diff --git a/benchmarks/TritonNPUBench/sglang/_chunked_lora_shrink_kernel/_chunked_lora_shrink_kernel.py b/benchmarks/TritonNPUBench/sglang/_chunked_lora_shrink_kernel/_chunked_lora_shrink_kernel.py new file mode 100644 index 00000000..3143d463 --- /dev/null +++ b/benchmarks/TritonNPUBench/sglang/_chunked_lora_shrink_kernel/_chunked_lora_shrink_kernel.py @@ -0,0 +1,118 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/sglang_operator_cases/newtest_cases/test__chunked_lora_shrink_kernel.py +# Main kernel: _chunked_lora_shrink_kernel +# PT file: test__chunked_lora_shrink_kernel_v2_unit2.pt + +import triton +import triton.language as tl + + +# === _chunked_lora_shrink_kernel === +@triton.jit(do_not_specialize=["num_segs"]) +def _chunked_lora_shrink_kernel( + # Pointers to matrices + x, + weights, + output, + # Information on sequence lengths,ranks and weight id + seg_indptr, + weight_indices, + lora_ranks, + permutation, + num_segs, + # Meta parameters + N: tl.constexpr, # num_slices * r + K: tl.constexpr, # input_dim + NUM_SLICES: tl.constexpr, + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, + BLOCK_K: tl.constexpr, + # index, +): + """ + Computes a chunked SGMV for LoRA shrink operations. + + The kernel ensures that output[seg_start:seg_start + seg_len, :rank * num_slices] + stores the product of the input `x` and the LoRA weights for the corresponding + sequence. This implies that when rank is 0, the kernel is essentially a no-op, + as output[seg_start:seg_start + seg_len, :0] is trivially correct (empty). + + Args: + x (torch.Tensor): The input activations tensor of shape `(s, K)`, where `s` + is the sum of all sequence lengths in the batch. + weights (torch.Tensor): The LoRA A weights for all available adapters, + with shape `(num_lora, N, K)` where N = num_slices * r. + output (torch.Tensor): The output tensor of shape `(s, N)`. + """ + x_stride_1: tl.constexpr = 1 + x_stride_0: tl.constexpr = K + + w_stride_0: tl.constexpr = N * K + w_stride_1: tl.constexpr = K + w_stride_2: tl.constexpr = 1 + + output_stride_0: tl.constexpr = N + output_stride_1: tl.constexpr = 1 + + pid_s = tl.program_id(1) + if pid_s >= num_segs: + return + + pid_n = tl.program_id(0) + + # Current block computes sequence with batch_id, + # which starts from row seg_start of x with length seg_len + w_index = tl.load(weight_indices + pid_s) + rank = tl.load(lora_ranks + w_index) + + # If rank is 0, this kernel becomes a no-op as the output is always trivially correct. + if rank == 0: + return + + seg_start = tl.load(seg_indptr + pid_s) + seg_end = tl.load(seg_indptr + pid_s + 1) + + # Adjust N dim according to the specific LoRA adapter + cur_n = tl.minimum(N, rank * NUM_SLICES) + + # Map logical sequence index to physical index + s_offset_logical = tl.arange(0, BLOCK_M) + seg_start + s_offset_physical = tl.load( + permutation + s_offset_logical, mask=s_offset_logical < seg_end + ) + + n_offset = tl.arange(0, BLOCK_N) + pid_n * BLOCK_N + k_offset = tl.arange(0, BLOCK_K) + x_ptrs = x + ( + s_offset_physical[:, None] * x_stride_0 + k_offset[None, :] * x_stride_1 + ) + w_ptrs = (weights + w_index * w_stride_0) + ( + k_offset[:, None] * w_stride_2 + n_offset[None, :] * w_stride_1 + ) + + # Iterate to compute the block in output matrix + partial_sum = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32) + for k in range(0, tl.cdiv(K, BLOCK_K)): + x_tile = tl.load( + x_ptrs, + mask=(s_offset_logical[:, None] < seg_end) + & (k_offset[None, :] < K - k * BLOCK_K), + other=0.0, + ) + w_tile = tl.load( + w_ptrs, + mask=(k_offset[:, None] < K - k * BLOCK_K) & (n_offset[None, :] < cur_n), + other=0.0, + ) + partial_sum += tl.dot(x_tile, w_tile) + + x_ptrs += BLOCK_K * x_stride_1 + w_ptrs += BLOCK_K * w_stride_2 + partial_sum = partial_sum.to(x.dtype.element_ty) + output_ptr = output + ( + s_offset_physical[:, None] * output_stride_0 + + n_offset[None, :] * output_stride_1 + ) + output_mask = (s_offset_logical[:, None] < seg_end) & (n_offset[None, :] < cur_n) + tl.store(output_ptr, partial_sum, mask=output_mask) + diff --git a/benchmarks/TritonNPUBench/sglang/_dequantize_k_cache_fast_kernel/_dequantize_k_cache_fast_kernel.py b/benchmarks/TritonNPUBench/sglang/_dequantize_k_cache_fast_kernel/_dequantize_k_cache_fast_kernel.py new file mode 100644 index 00000000..557b27a9 --- /dev/null +++ b/benchmarks/TritonNPUBench/sglang/_dequantize_k_cache_fast_kernel/_dequantize_k_cache_fast_kernel.py @@ -0,0 +1,57 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/sglang_operator_cases/newtest_cases/test__dequantize_k_cache_fast_kernel.py +# Main kernel: _dequantize_k_cache_fast_kernel +# PT file: test__dequantize_k_cache_fast_kernel_v2.pt + +import triton +import triton.language as tl + + +# === _dequantize_k_cache_fast_kernel === +@triton.jit +def _dequantize_k_cache_fast_kernel( + output_ptr, + input_nope_q_ptr, + input_nope_s_ptr, + input_rope_ptr, + output_stride_0: int, + input_nope_q_stride_0: int, + input_nope_s_stride_0: int, + input_rope_stride_0: int, + NUM_NOPE_BLOCKS: tl.constexpr, + GROUP_SIZE: tl.constexpr, + DIM_NOPE: tl.constexpr, + DIM_ROPE: tl.constexpr, +): + token_id = tl.program_id(0) + raw_block_id = tl.program_id(1) + + if raw_block_id < NUM_NOPE_BLOCKS: + # a. dequant nope + effective_block_id = raw_block_id + + offs_q = effective_block_id * GROUP_SIZE + tl.arange(0, GROUP_SIZE) + mask = offs_q < DIM_NOPE + ptr_q = input_nope_q_ptr + token_id * input_nope_q_stride_0 + offs_q + ptr_s = input_nope_s_ptr + token_id * input_nope_s_stride_0 + effective_block_id + + y_q = tl.load(ptr_q, mask=mask, other=0.0).to(tl.float32) + y_s = tl.load(ptr_s) + + y = (y_q * y_s).to(output_ptr.dtype.element_ty) + + dst_ptr = output_ptr + token_id * output_stride_0 + offs_q + tl.store(dst_ptr, y, mask=mask) + else: + # b. copy rope + effective_block_id = raw_block_id - NUM_NOPE_BLOCKS + + offs = effective_block_id * GROUP_SIZE + tl.arange(0, GROUP_SIZE) + mask = offs < DIM_ROPE + + src_ptr = input_rope_ptr + token_id * input_rope_stride_0 + offs + dst_ptr = output_ptr + token_id * output_stride_0 + DIM_NOPE + offs + + data = tl.load(src_ptr, mask=mask).to(tl.bfloat16) + tl.store(dst_ptr, data, mask=mask) + diff --git a/benchmarks/TritonNPUBench/sglang/_dequantize_k_cache_paged_kernel/_dequantize_k_cache_paged_kernel.py b/benchmarks/TritonNPUBench/sglang/_dequantize_k_cache_paged_kernel/_dequantize_k_cache_paged_kernel.py new file mode 100644 index 00000000..8c676c5e --- /dev/null +++ b/benchmarks/TritonNPUBench/sglang/_dequantize_k_cache_paged_kernel/_dequantize_k_cache_paged_kernel.py @@ -0,0 +1,63 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/sglang_operator_cases/newtest_cases/test__dequantize_k_cache_paged_kernel.py +# Main kernel: _dequantize_k_cache_paged_kernel +# PT file: test__dequantize_k_cache_paged_kernel_v2.pt + +import triton +import triton.language as tl + + +# === _dequantize_k_cache_paged_kernel === +@triton.jit +def _dequantize_k_cache_paged_kernel( + output_ptr, + input_nope_q_ptr, + input_nope_s_ptr, + input_rope_ptr, + page_table_1_ptr, + output_stride_0: int, + input_nope_q_stride_0: int, + input_nope_s_stride_0: int, + input_rope_stride_0: int, + NUM_NOPE_BLOCKS: tl.constexpr, + GROUP_SIZE: tl.constexpr, + DIM_NOPE: tl.constexpr, + DIM_ROPE: tl.constexpr, +): + token_id = tl.program_id(0) + token_id_paged = tl.load(page_table_1_ptr + token_id).to(tl.int32) + raw_block_id = tl.program_id(1) + + if raw_block_id < NUM_NOPE_BLOCKS: + # a. dequant nope + effective_block_id = raw_block_id + + offs_q = effective_block_id * GROUP_SIZE + tl.arange(0, GROUP_SIZE) + mask = offs_q < DIM_NOPE + ptr_q = input_nope_q_ptr + token_id_paged * input_nope_q_stride_0 + offs_q + ptr_s = ( + input_nope_s_ptr + + token_id_paged * input_nope_s_stride_0 + + effective_block_id + ) + + y_q = tl.load(ptr_q, mask=mask, other=0.0).to(tl.float32) + y_s = tl.load(ptr_s) + + y = (y_q * y_s).to(output_ptr.dtype.element_ty) + + dst_ptr = output_ptr + token_id * output_stride_0 + offs_q + tl.store(dst_ptr, y, mask=mask) + else: + # b. copy rope + effective_block_id = raw_block_id - NUM_NOPE_BLOCKS + + offs = effective_block_id * GROUP_SIZE + tl.arange(0, GROUP_SIZE) + mask = offs < DIM_ROPE + + src_ptr = input_rope_ptr + token_id_paged * input_rope_stride_0 + offs + dst_ptr = output_ptr + token_id * output_stride_0 + DIM_NOPE + offs + + data = tl.load(src_ptr, mask=mask).to(tl.bfloat16) + tl.store(dst_ptr, data, mask=mask) + diff --git a/benchmarks/TritonNPUBench/sglang/_embedding_lora_a_kernel/_embedding_lora_a_kernel.py b/benchmarks/TritonNPUBench/sglang/_embedding_lora_a_kernel/_embedding_lora_a_kernel.py new file mode 100644 index 00000000..3f56c0ad --- /dev/null +++ b/benchmarks/TritonNPUBench/sglang/_embedding_lora_a_kernel/_embedding_lora_a_kernel.py @@ -0,0 +1,106 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/sglang_operator_cases/newtest_cases/test__embedding_lora_a_kernel.py +# Main kernel: _embedding_lora_a_kernel +# PT file: test__embedding_lora_a_kernel_v2.pt + +import triton +import triton.language as tl + + +# === _embedding_lora_a_kernel === +@triton.jit +def _embedding_lora_a_kernel( + # Pointers to tensors + input_ids, + weights, + output, + extra_embeddings, + # Dimensions + vocab_size, + rank, + num_loras, + # Strides + w_stride_0, # stride for lora index + w_stride_1, # stride for rank + w_stride_2, # stride for vocab + output_stride_0, + output_stride_1, + extra_emb_stride_0, # stride for lora index + extra_emb_stride_1, # stride for token + extra_emb_stride_2, # stride for hidden dim (= rank for extra embeddings) + # Batch info + seg_lens, + seg_indptr, + weight_indices, + lora_ranks, + # Meta-parameters + BLOCK_RANK: tl.constexpr, + HAS_EXTRA_EMBEDDINGS: tl.constexpr, +): + """ + Embedding lookup for LoRA A weights with support for extra tokens. + + Each program handles one token across a block of rank dimensions. + Grid: (cdiv(max_len, 1), bs) - one program per token in each batch + """ + batch_id = tl.program_id(axis=1) + token_idx = tl.program_id(axis=0) + + w_index = tl.load(weight_indices + batch_id) + rank_val = tl.load(lora_ranks + w_index) + + # If rank is 0, skip + if rank_val == 0: + return + + seg_start = tl.load(seg_indptr + batch_id) + seg_len = tl.load(seg_lens + batch_id) + + # Check if this token is within the segment + if token_idx >= seg_len: + return + + # Load the token ID + token_id = tl.load(input_ids + seg_start + token_idx) + + # Process in chunks of BLOCK_RANK dimensions + num_blocks = tl.cdiv(rank_val, BLOCK_RANK) + + for block_id in range(num_blocks): + rank_offset = tl.arange(0, BLOCK_RANK) + block_id * BLOCK_RANK + rank_mask = rank_offset < rank_val + + # Check if this is an extra token + is_extra_token = token_id >= vocab_size + + if HAS_EXTRA_EMBEDDINGS and is_extra_token: + # Use extra embeddings + extra_token_id = token_id - vocab_size + extra_emb_ptr = ( + extra_embeddings + + w_index * extra_emb_stride_0 + + extra_token_id * extra_emb_stride_1 + + rank_offset * extra_emb_stride_2 + ) + emb_values = tl.load(extra_emb_ptr, mask=rank_mask, other=0.0) + else: + # Use regular LoRA A weights + # weights shape: (num_loras, rank, vocab_size) + # We need to load weights[w_index, rank_offset, token_id] + token_id_clamped = tl.minimum(token_id, vocab_size - 1) + weight_ptr = ( + weights + + w_index * w_stride_0 + + rank_offset * w_stride_1 + + token_id_clamped * w_stride_2 + ) + emb_values = tl.load(weight_ptr, mask=rank_mask, other=0.0) + + # Write to output + output_ptr = ( + output + + (seg_start + token_idx) * output_stride_0 + + rank_offset * output_stride_1 + ) + tl.store(output_ptr, emb_values, mask=rank_mask) + diff --git a/benchmarks/TritonNPUBench/sglang/_fwd_grouped_kernel_stage1_rope/_fwd_grouped_kernel_stage1_rope.py b/benchmarks/TritonNPUBench/sglang/_fwd_grouped_kernel_stage1_rope/_fwd_grouped_kernel_stage1_rope.py new file mode 100644 index 00000000..6189d96c --- /dev/null +++ b/benchmarks/TritonNPUBench/sglang/_fwd_grouped_kernel_stage1_rope/_fwd_grouped_kernel_stage1_rope.py @@ -0,0 +1,281 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/sglang_operator_cases/newtest_cases/test__fwd_grouped_kernel_stage1_rope.py +# Main kernel: _fwd_grouped_kernel_stage1_rope +# PT file: test__fwd_grouped_kernel_stage1_rope_v2.pt + +import triton +import triton.language as tl + + +# === _fwd_grouped_kernel_stage1_rope === +@triton.jit +def _fwd_grouped_kernel_stage1_rope( + Q, # Holds [Q_NOPE; Q_PE], b x h x (d+r) + K_Buffer, # Holds [KV; K_PE], b*s x (c+r) + V_buffer, # Holds [KV], b*s x (c) + cos_sin_cache, # max_seq_len x (rotary_dim * 2) + positions, # sequence positions + sm_scale, + kv_indptr, + kv_indices, + Att_Out, # b x h x NUM_KV_SPLITS x (kv_lora_rank + 1) + k_pe_t_out, + stride_qb, + stride_qh, + stride_buf_kbs, + stride_buf_vbs, + stride_mid_ob, + stride_mid_oh, + stride_mid_os, + stride_kpe_tokens_out_b, + stride_cos_sin_cache_s, + stride_positions_b, + rotary_dim: tl.constexpr, + kv_lora_rank: tl.constexpr, + qk_rope_head_dim: tl.constexpr, + kv_group_num: tl.constexpr, + q_head_num: tl.constexpr, + BLOCK_C: tl.constexpr, + BLOCK_R: tl.constexpr, + BLOCK_N: tl.constexpr, + BLOCK_H: tl.constexpr, + NUM_KV_SPLITS: tl.constexpr, + logit_cap: tl.constexpr, + USE_ROPE: tl.constexpr, + IS_NEOX_STYLE: tl.constexpr, +): + + cur_batch = tl.program_id(0) + cur_head_id = tl.program_id(1) + split_kv_id = tl.program_id(2) + + if BLOCK_H < kv_group_num: + VALID_BLOCK_H: tl.constexpr = BLOCK_H + else: + VALID_BLOCK_H: tl.constexpr = kv_group_num + cur_head = cur_head_id * VALID_BLOCK_H + tl.arange(0, BLOCK_H) + mask_h = cur_head < (cur_head_id + 1) * VALID_BLOCK_H + mask_h = mask_h & (cur_head < q_head_num) + + offs_c = tl.arange(0, BLOCK_C) + offs_qk_r = tl.arange(kv_lora_rank, kv_lora_rank + BLOCK_R) # to get the k_pe + + off_q_pe = ( + cur_batch * stride_qb + cur_head[:, None] * stride_qh + offs_qk_r[None, :] + ) + offs_q = cur_batch * stride_qb + cur_head[:, None] * stride_qh + offs_c[None, :] + + mask_c = offs_c < kv_lora_rank + mask_qk_r = offs_qk_r < (kv_lora_rank + qk_rope_head_dim) + + cur_batch_kv_start_idx = tl.load(kv_indptr + cur_batch) + cur_batch_seq_len = tl.load(kv_indptr + cur_batch + 1) - cur_batch_kv_start_idx + + q = tl.load(Q + offs_q, mask=(mask_h[:, None]) & (mask_c[None, :]), other=0.0) + q_pe = tl.load( + Q + off_q_pe, mask=(mask_h[:, None]) & (mask_qk_r[None, :]), other=0.0 + ) + + kv_len_per_split = tl.cdiv(cur_batch_seq_len, NUM_KV_SPLITS) + split_kv_start = kv_len_per_split * split_kv_id + split_kv_end = tl.minimum(split_kv_start + kv_len_per_split, cur_batch_seq_len) + + # apply rotary embedding for q_pe, and k_pe (last token per batch of K_PE) + LAST_SPLIT = split_kv_end == cur_batch_seq_len + k_pe_last_token = tl.zeros([BLOCK_R], dtype=q.dtype) + + if USE_ROPE: + if IS_NEOX_STYLE: + # [BLOCK_ROTARY // 2, BLOCK_ROTARY // 2 + 1, BLOCK_ROTARY // 2 + 2, ..., 0, 1, 2, ..., BLOCK_ROTARY // 2 - 1, pass:] + offs_qk_rot_r = kv_lora_rank + ( + (tl.arange(0, BLOCK_R) + (rotary_dim // 2)) % rotary_dim + ) + # Which elements to flip + mask_rotate = tl.arange(0, BLOCK_R) < (rotary_dim // 2) + # [0 , 1, 2, ..., rotary_dim // 2 - 1, 0 , 1, 2, ..., rotary_dim // 2 - 1] + offs_rotary = tl.arange(0, BLOCK_R) % (rotary_dim // 2) + else: + # [1, 0, 3, 2, 5, 4, ..., BLOCK_R, BLOCK_R - 1] + offs_qk_rot_r = ( + kv_lora_rank + + (((tl.arange(0, BLOCK_R) + 1) % 2) * 2) + - 1 + + tl.arange(0, BLOCK_R) + ) + mask_rotate = tl.arange(0, BLOCK_R) % 2 < 1 + # [0, 0, 1, 1, ..., rotary_dim // 2 - 1, rotary_dim // 2 - 1] + offs_rotary = tl.arange(0, BLOCK_R) // 2 + + if qk_rope_head_dim > rotary_dim: + offs_qk_rot_r = tl.where( + tl.arange(0, BLOCK_R) < rotary_dim, offs_qk_rot_r, tl.arange(0, BLOCK_R) + ) + offs_rotary = tl.where( + tl.arange(0, BLOCK_R) < rotary_dim, offs_rotary, tl.arange(0, BLOCK_R) + ) + + mask_rotary = tl.arange(0, BLOCK_R) < rotary_dim + + pos = tl.load(positions + cur_batch * stride_positions_b) + cos = tl.load( + cos_sin_cache + pos * stride_cos_sin_cache_s + offs_rotary, + mask=mask_rotary, + other=1.0, + ) + sin = tl.load( + cos_sin_cache + + pos * stride_cos_sin_cache_s + + offs_rotary + + rotary_dim // 2, + mask_rotary, + other=0.0, + ) + + off_q_pe_rot = ( + cur_batch * stride_qb + + cur_head[:, None] * stride_qh + + offs_qk_rot_r[None, :] + ) + mask_qk_rot_r = offs_qk_rot_r < (kv_lora_rank + qk_rope_head_dim) + + # 0, 2, 4,.... 1, 3, 5... + q_pe_rot = tl.load( + Q + off_q_pe_rot, + mask=(mask_h[:, None]) & (mask_qk_rot_r[None, :]), + other=0.0, + ) + q_pe_rot = tl.where(mask_rotate[None, :], -q_pe_rot, q_pe_rot) + + q_pe = q_pe * cos + q_pe_rot * sin + + # we only apply to the last token in the K_PE + if LAST_SPLIT: + # debug assert + # if (cur_batch == 0 and cur_head == 0) and split_kv_id < NUM_KV_SPLITS - 1: + # tl.device_assert(False, "Only last split should compute k_pe") + + kv_loc = tl.load( + kv_indices + cur_batch_kv_start_idx + cur_batch_seq_len - 1 + ) + offs_buf_k_pe_last_token = kv_loc * stride_buf_kbs + offs_qk_r + offs_buf_k_pe_rot_last_token = kv_loc * stride_buf_kbs + offs_qk_rot_r + k_pe_last_token = tl.load(K_Buffer + offs_buf_k_pe_last_token) + + k_pe_rot_last_token = tl.load(K_Buffer + offs_buf_k_pe_rot_last_token) + k_pe_rot_last_token = tl.where( + mask_rotate, -k_pe_rot_last_token, k_pe_rot_last_token + ) + + k_pe_last_token = k_pe_last_token * cos + k_pe_rot_last_token * sin + + e_max = tl.zeros([BLOCK_H], dtype=tl.float32) - float("inf") + e_sum = tl.zeros([BLOCK_H], dtype=tl.float32) + acc = tl.zeros([BLOCK_H, BLOCK_C], dtype=tl.float32) + + if split_kv_end > split_kv_start: + for start_n in range(split_kv_start, split_kv_end, BLOCK_N): + offs_n = start_n + tl.arange(0, BLOCK_N) + kv_loc = tl.load( + kv_indices + cur_batch_kv_start_idx + offs_n, + mask=offs_n < split_kv_end, + other=0, + ) + + offs_buf_kv = kv_loc[None, :] * stride_buf_kbs + offs_c[:, None] + offs_buf_k_pe = kv_loc[None, :] * stride_buf_kbs + offs_qk_r[:, None] + + k_pe = tl.load( + K_Buffer + offs_buf_k_pe, + mask=(offs_n[None, :] < split_kv_end) & (mask_qk_r[:, None]), + other=0.0, + ) # positional embedding part of keys + + if (USE_ROPE and LAST_SPLIT) and start_n >= cur_batch_seq_len - BLOCK_N: + k_pe = tl.where( + offs_n[None, :] != (split_kv_end - 1), + k_pe, + k_pe_last_token[:, None], + ) + + # (16, 64) x (64, 32) + # dot product of rope parts + qk = tl.dot(q_pe, k_pe.to(q_pe.dtype)) + + kv = tl.load( + K_Buffer + offs_buf_kv, + mask=(offs_n[None, :] < split_kv_end) & (mask_c[:, None]), + other=0.0, + ) # the shared latent tensor for keys and values + + # (16, 512) x (512, 32) + # dot product of nope parts + qk += tl.dot(q, kv) + + qk *= sm_scale + + if logit_cap > 0: + qk = logit_cap * tanh(qk / logit_cap) + + qk = tl.where( + mask_h[:, None] & (offs_n[None, :] < split_kv_end), qk, float("-inf") + ) + + offs_buf_v = kv_loc[:, None] * stride_buf_vbs + offs_c[None, :] + v = tl.load( + V_buffer + offs_buf_v, + mask=(offs_n[:, None] < split_kv_end) & (mask_c[None, :]), + other=0.0, + ) + + n_e_max = tl.maximum(tl.max(qk, 1), e_max) + re_scale = tl.exp(e_max - n_e_max) + p = tl.exp(qk - n_e_max[:, None]) + acc *= re_scale[:, None] + # (16, 32) x (32, 512) + acc += tl.dot(p.to(v.dtype), v) + + e_sum = e_sum * re_scale + tl.sum(p, 1) + e_max = n_e_max + + offs_mid_o = ( + cur_batch * stride_mid_ob + + cur_head[:, None] * stride_mid_oh + + split_kv_id * stride_mid_os + + offs_c[None, :] + ) + + if USE_ROPE: + if LAST_SPLIT: + k_pe_last_token_ptrs = ( + k_pe_t_out + + cur_batch * stride_kpe_tokens_out_b + + tl.arange(0, BLOCK_R) + ) + tl.store(k_pe_last_token_ptrs, k_pe_last_token, mask=mask_qk_r) + + tl.store( + Att_Out + offs_mid_o, + acc / e_sum[:, None], + mask=(mask_h[:, None]) & (mask_c[None, :]), + ) + + offs_mid_o_1 = ( + cur_batch * stride_mid_ob + + cur_head * stride_mid_oh + + split_kv_id * stride_mid_os + + kv_lora_rank + ) + + tl.store( + Att_Out + offs_mid_o_1, + e_max + tl.log(e_sum), + mask=mask_h, + ) + + +# === tanh === +@triton.jit +def tanh(x): + # Tanh is just a scaled sigmoid + return 2 * tl.sigmoid(2 * x) - 1 + diff --git a/benchmarks/TritonNPUBench/sglang/_fwd_kernel_ep_gather/_fwd_kernel_ep_gather.py b/benchmarks/TritonNPUBench/sglang/_fwd_kernel_ep_gather/_fwd_kernel_ep_gather.py new file mode 100644 index 00000000..ed7048c8 --- /dev/null +++ b/benchmarks/TritonNPUBench/sglang/_fwd_kernel_ep_gather/_fwd_kernel_ep_gather.py @@ -0,0 +1,75 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/sglang_operator_cases/newtest_cases/test__fwd_kernel_ep_gather.py +# Main kernel: _fwd_kernel_ep_gather +# PT file: test__fwd_kernel_ep_gather_v2.pt + +import triton +import triton.language as tl + + +# === _fwd_kernel_ep_gather === +@triton.jit +def _fwd_kernel_ep_gather( + total_token_num, + input_tensor, + input_tensor_stride0, + input_tensor_stride1, + recv_topk_ids, + recv_topk_ids_stride0, + recv_topk_ids_stride1, + recv_topk_weight, + recv_topk_weight_stride0, + recv_topk_weight_stride1, + input_index, + input_index_stride0, + input_index_stride1, + output_tensor, + output_tensor_stride0, + output_tensor_stride1, + topk_num: tl.constexpr, + BLOCK_D: tl.constexpr, +): + cur_block_int32 = tl.program_id(0) + cur_block = cur_block_int32.to(tl.int64) + + start_cur_token_int32 = tl.program_id(1) + + grid_num = tl.num_programs(1) + + for cur_token_int32 in range(start_cur_token_int32, total_token_num, grid_num): + cur_token = cur_token_int32.to(tl.int64) + + off_d = tl.arange(0, BLOCK_D) + accumulator = tl.zeros([BLOCK_D], dtype=tl.float32) + + for topk_index_int32 in range(0, topk_num): + topk_index = topk_index_int32.to(tl.int64) + + expert_id = tl.load( + recv_topk_ids + cur_token * recv_topk_ids_stride0 + topk_index + ) + if expert_id >= 0: + source_token_index_int32 = tl.load( + input_index + cur_token * input_index_stride0 + topk_index + ) + source_token_index = source_token_index_int32.to(tl.int64) + + acc_weight = tl.load( + recv_topk_weight + cur_token * recv_topk_weight_stride0 + topk_index + ) + tmp = tl.load( + input_tensor + + source_token_index * input_tensor_stride0 + + cur_block * BLOCK_D + + off_d + ) + accumulator += tmp.to(tl.float32) * acc_weight + + tl.store( + output_tensor + + cur_token * output_tensor_stride0 + + cur_block * BLOCK_D + + off_d, + accumulator.to(output_tensor.dtype.element_ty), + ) + diff --git a/benchmarks/TritonNPUBench/sglang/_fwd_kernel_ep_scatter_1/_fwd_kernel_ep_scatter_1.py b/benchmarks/TritonNPUBench/sglang/_fwd_kernel_ep_scatter_1/_fwd_kernel_ep_scatter_1.py new file mode 100644 index 00000000..ee976bd3 --- /dev/null +++ b/benchmarks/TritonNPUBench/sglang/_fwd_kernel_ep_scatter_1/_fwd_kernel_ep_scatter_1.py @@ -0,0 +1,42 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/sglang_operator_cases/newtest_cases/test__fwd_kernel_ep_scatter_1.py +# Main kernel: _fwd_kernel_ep_scatter_1 +# PT file: test__fwd_kernel_ep_scatter_1_v2.pt + +import triton +import triton.language as tl + + +# === _fwd_kernel_ep_scatter_1 === +@triton.jit +def _fwd_kernel_ep_scatter_1( + num_recv_tokens_per_expert, + expert_start_loc, + m_indices, + num_experts: tl.constexpr, + BLOCK_E: tl.constexpr, + BLOCK_EXPERT_NUM: tl.constexpr, +): + cur_expert = tl.program_id(0) + + offset_cumsum = tl.arange(0, BLOCK_EXPERT_NUM) + tokens_per_expert = tl.load( + num_recv_tokens_per_expert + offset_cumsum, + mask=offset_cumsum < num_experts, + other=0, + ) + cumsum = tl.cumsum(tokens_per_expert) - tokens_per_expert + tl.store(expert_start_loc + offset_cumsum, cumsum, mask=offset_cumsum < num_experts) + + cur_expert_start = tl.load(expert_start_loc + cur_expert) + cur_expert_token_num = tl.load(num_recv_tokens_per_expert + cur_expert) + + m_indices_start_ptr = m_indices + cur_expert_start + off_expert = tl.arange(0, BLOCK_E) + + for start_m in tl.range(0, cur_expert_token_num, BLOCK_E, num_stages=4): + tl.store( + m_indices_start_ptr + start_m + off_expert, + cur_expert, + ) + diff --git a/benchmarks/TritonNPUBench/sglang/_fwd_kernel_ep_scatter_2/_fwd_kernel_ep_scatter_2.py b/benchmarks/TritonNPUBench/sglang/_fwd_kernel_ep_scatter_2/_fwd_kernel_ep_scatter_2.py new file mode 100644 index 00000000..652cd1df --- /dev/null +++ b/benchmarks/TritonNPUBench/sglang/_fwd_kernel_ep_scatter_2/_fwd_kernel_ep_scatter_2.py @@ -0,0 +1,81 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/sglang_operator_cases/newtest_cases/test__fwd_kernel_ep_scatter_2.py +# Main kernel: _fwd_kernel_ep_scatter_2 +# PT file: test__fwd_kernel_ep_scatter_2_v2.pt + +import triton +import triton.language as tl + + +# === _fwd_kernel_ep_scatter_2 === +@triton.jit +def _fwd_kernel_ep_scatter_2( + total_token_num, + expert_start_loc, + recv_x, + recv_x_stride0, + recv_x_stride1, + recv_x_scale, + recv_x_scale_stride0, + recv_x_scale_stride1, + recv_topk, + recv_topk_stride0, + recv_topk_stride1, + output_tensor, + output_tensor_stride0, + output_tensor_stride1, + output_tensor_scale, + output_tensor_scale_stride0, + output_tensor_scale_stride1, + output_index, + output_index_stride0, + output_index_stride1, + topk_num: tl.constexpr, + HIDDEN_SIZE: tl.constexpr, + HIDDEN_SIZE_PAD: tl.constexpr, + SCALE_HIDDEN_SIZE: tl.constexpr, + SCALE_HIDDEN_SIZE_PAD: tl.constexpr, +): + start_token_id = tl.program_id(0) + grid_num = tl.num_programs(0) + + offset_in = tl.arange(0, HIDDEN_SIZE_PAD) + mask = offset_in < HIDDEN_SIZE + + index_in_s = tl.arange(0, SCALE_HIDDEN_SIZE_PAD) + mask_s = index_in_s < SCALE_HIDDEN_SIZE + + for token_id_int32 in range(start_token_id, total_token_num, grid_num): + token_id = token_id_int32.to(tl.int64) + to_copy = tl.load(recv_x + token_id * recv_x_stride0 + offset_in, mask=mask) + to_copy_s = tl.load( + recv_x_scale + + token_id * recv_x_scale_stride0 + + index_in_s * recv_x_scale_stride1, + mask=mask_s, + ) + + for topk_idx_int32 in tl.range(0, topk_num, 1, num_stages=4): + topk_index = topk_idx_int32.to(tl.int64) + expert_id = tl.load(recv_topk + token_id * recv_topk_stride0 + topk_index) + if expert_id >= 0: + dest_token_index_int32 = tl.atomic_add(expert_start_loc + expert_id, 1) + dest_token_index = dest_token_index_int32.to(tl.int64) + + tl.store( + output_index + token_id * output_index_stride0 + topk_index, + dest_token_index_int32, + ) + output_tensor_ptr = ( + output_tensor + dest_token_index * output_tensor_stride0 + ) + output_tensor_scale_ptr = ( + output_tensor_scale + dest_token_index * output_tensor_scale_stride0 + ) + tl.store(output_tensor_ptr + offset_in, to_copy, mask=mask) + tl.store( + output_tensor_scale_ptr + index_in_s * output_tensor_scale_stride1, + to_copy_s, + mask=mask_s, + ) + diff --git a/benchmarks/TritonNPUBench/sglang/_gate_up_lora_b_kernel/_gate_up_lora_b_kernel.py b/benchmarks/TritonNPUBench/sglang/_gate_up_lora_b_kernel/_gate_up_lora_b_kernel.py new file mode 100644 index 00000000..03b23e7d --- /dev/null +++ b/benchmarks/TritonNPUBench/sglang/_gate_up_lora_b_kernel/_gate_up_lora_b_kernel.py @@ -0,0 +1,127 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/sglang_operator_cases/newtest_cases/test__gate_up_lora_b_kernel.py +# Main kernel: _gate_up_lora_b_kernel +# PT file: test__gate_up_lora_b_kernel_v2.pt + +import triton +import triton.language as tl + + +# === _gate_up_lora_b_kernel === +@triton.jit +def _gate_up_lora_b_kernel( + # Pointers to matrices + x, + weights, + output, + # Parameters of size + K, # K = R + output_dim, + # Strides + x_stride_0, + x_stride_1, + w_stride_0, + w_stride_1, + w_stride_2, + output_stride_0, + output_stride_1, + # Information on sequence lengths,ranks and weight id + seg_lens, + seg_indptr, + weight_indices, + lora_ranks, + # Meta parameters + BLOCK_S: tl.constexpr, + BLOCK_N: tl.constexpr, + BLOCK_K: tl.constexpr, + # For fused output scaling + scalings, +): + """ + This kernel packs 2 sgemms (gate/up) into a single kernel. The multiplication + results are accumulated into the output tensor. + + When a sequence's rank is 0, the kernel is essentially a no-op, following + the convention in pytorch where the product of two matrices of shape (m, 0) + and (0, n) is an all-zero matrix of shape (m, n). + + Args: + x (Tensor): The input tensor, which is the result of the LoRA A projection. + Shape: (s, 2 * K), where s is the sum of all sequence lengths in the + batch and K is the maximum LoRA rank. + weights (Tensor): The LoRA B weights for all adapters. + Shape: (num_lora, 2 * output_dim, K). + output (Tensor): The output tensor where the result is stored. + Shape: (s, 2 * output_dim). + """ + # output_dim >> K + + # Current block computes sequence with batch_id, + # which starts from row seg_start of x with length seg_len. + # gate_up_id decides which of gate or up (0: gate, 1: up) + batch_id = tl.program_id(axis=2) + w_index = tl.load(weight_indices + batch_id) + rank = tl.load(lora_ranks + w_index) + + # If rank is 0, this kernel is a no-op. + if rank == 0: + return + + gate_up_id = tl.program_id(axis=1) + pid = tl.program_id(axis=0) + seg_len = tl.load(seg_lens + batch_id) + seg_start = tl.load(seg_indptr + batch_id) + n_start = gate_up_id * output_dim # offset on output dim + scaling = tl.load(scalings + w_index) + + # Adjust K (rank) according to the specific LoRA adapter + K = tl.minimum(K, rank) + + # The tile in output matrix will have (pid_s, pid_n) as id + num_pid_n = tl.cdiv(output_dim, BLOCK_N) + pid_s = pid // num_pid_n + pid_n = pid % num_pid_n + + # Create pointers for the first block of x and weights + # The pointers will be advanced as we move in the K direction + # and accumulate + s_offset = tl.arange(0, BLOCK_S) + pid_s * BLOCK_S + n_offset = tl.arange(0, BLOCK_N) + pid_n * BLOCK_N + k_offset = tl.arange(0, BLOCK_K) + + x_ptrs = (x + seg_start * x_stride_0 + (gate_up_id * K) * x_stride_1) + ( + s_offset[:, None] * x_stride_0 + k_offset[None, :] * x_stride_1 + ) + w_ptrs = (weights + w_index * w_stride_0 + n_start * w_stride_1) + ( + k_offset[:, None] * w_stride_2 + n_offset[None, :] * w_stride_1 + ) + + # Iterate to compute the block in output matrix + partial_sum = tl.zeros((BLOCK_S, BLOCK_N), dtype=tl.float32) + for k in range(0, tl.cdiv(K, BLOCK_K)): + x_tile = tl.load( + x_ptrs, + mask=(s_offset[:, None] < seg_len) & (k_offset[None, :] < K - k * BLOCK_K), + other=0.0, + ) + w_tile = tl.load( + w_ptrs, + mask=(k_offset[:, None] < K - k * BLOCK_K) + & (n_offset[None, :] < output_dim), + other=0.0, + ) + partial_sum += tl.dot(x_tile, w_tile) + + x_ptrs += BLOCK_K * x_stride_1 + w_ptrs += BLOCK_K * w_stride_2 + + # Store result to output matrix + partial_sum *= scaling + partial_sum = partial_sum.to(x.dtype.element_ty) + output_ptr = (output + seg_start * output_stride_0 + n_start * output_stride_1) + ( + s_offset[:, None] * output_stride_0 + n_offset[None, :] * output_stride_1 + ) + output_mask = (s_offset[:, None] < seg_len) & (n_offset[None, :] < output_dim) + partial_sum += tl.load(output_ptr, mask=output_mask) + tl.store(output_ptr, partial_sum, mask=output_mask) + diff --git a/benchmarks/TritonNPUBench/sglang/_layer_norm_fwd_1pass_kernel/_layer_norm_fwd_1pass_kernel.py b/benchmarks/TritonNPUBench/sglang/_layer_norm_fwd_1pass_kernel/_layer_norm_fwd_1pass_kernel.py new file mode 100644 index 00000000..d1a690f6 --- /dev/null +++ b/benchmarks/TritonNPUBench/sglang/_layer_norm_fwd_1pass_kernel/_layer_norm_fwd_1pass_kernel.py @@ -0,0 +1,107 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/sglang_operator_cases/newtest_cases/test__layer_norm_fwd_1pass_kernel.py +# Main kernel: _layer_norm_fwd_1pass_kernel +# PT file: test__layer_norm_fwd_1pass_kernel_v2.pt + +import triton +import triton.language as tl + + +# === _layer_norm_fwd_1pass_kernel === +@triton.jit +def _layer_norm_fwd_1pass_kernel( + X, # pointer to the input + Y, # pointer to the output + W, # pointer to the weights + B, # pointer to the biases + Z, # pointer to the other branch + Mean, # pointer to the mean + Rstd, # pointer to the 1/std + stride_x_row, # how much to increase the pointer when moving by 1 row + stride_y_row, + stride_z_row, + M, # number of rows in X + N: tl.constexpr, # number of columns in X + eps, # epsilon to avoid division by zero + BLOCK_N: tl.constexpr, + ROWS_PER_BLOCK: tl.constexpr, + HAS_BIAS: tl.constexpr, + HAS_Z: tl.constexpr, + NORM_BEFORE_GATE: tl.constexpr, + IS_RMS_NORM: tl.constexpr, +): + # Map the program id to the starting row of X and Y it should compute. + row_start = tl.program_id(0) * ROWS_PER_BLOCK + group = tl.program_id(1) + + # Create 2D tile: [ROWS_PER_BLOCK, BLOCK_N] + rows = row_start + tl.arange(0, ROWS_PER_BLOCK) + cols = tl.arange(0, BLOCK_N) + + # Compute offsets for 2D tile + row_offsets = rows[:, None] * stride_x_row + col_offsets = cols[None, :] + group * N + + # Base pointers + X_base = X + row_offsets + col_offsets + Y_base = Y + rows[:, None] * stride_y_row + col_offsets + + # Create mask for valid rows and columns + row_mask = rows[:, None] < M + col_mask = cols[None, :] < N + mask = row_mask & col_mask + + # Load input data with 2D tile + x = tl.load(X_base, mask=mask, other=0.0).to(tl.float32) + + if HAS_Z and not NORM_BEFORE_GATE: + Z_base = Z + rows[:, None] * stride_z_row + col_offsets + z = tl.load(Z_base, mask=mask, other=0.0).to(tl.float32) + x *= z * tl.sigmoid(z) + + # Compute mean and variance per row (reduce along axis 1) + if not IS_RMS_NORM: + mean = tl.sum(x, axis=1) / N # Shape: [ROWS_PER_BLOCK] + # Store mean for each row + mean_offsets = group * M + rows + mean_mask = rows < M + tl.store(Mean + mean_offsets, mean, mask=mean_mask) + # Broadcast mean back to 2D for subtraction + xbar = tl.where(mask, x - mean[:, None], 0.0) + var = tl.sum(xbar * xbar, axis=1) / N # Shape: [ROWS_PER_BLOCK] + else: + xbar = tl.where(mask, x, 0.0) + var = tl.sum(xbar * xbar, axis=1) / N # Shape: [ROWS_PER_BLOCK] + mean = 0.0 # Placeholder for RMS norm + + rstd = tl.rsqrt(var + eps) # Shape: [ROWS_PER_BLOCK] + + # Store rstd for each row + rstd_offsets = group * M + rows + rstd_mask = rows < M + tl.store(Rstd + rstd_offsets, rstd, mask=rstd_mask) + + # Load weights and biases (broadcast across rows) + w_offsets = cols + group * N + w_mask = cols < N + w = tl.load(W + w_offsets, mask=w_mask, other=0.0).to(tl.float32) + + if HAS_BIAS: + b = tl.load(B + w_offsets, mask=w_mask, other=0.0).to(tl.float32) + + # Normalize and apply linear transformation + if not IS_RMS_NORM: + x_hat = (x - mean[:, None]) * rstd[:, None] + else: + x_hat = x * rstd[:, None] + + y = x_hat * w[None, :] + b[None, :] if HAS_BIAS else x_hat * w[None, :] + + if HAS_Z and NORM_BEFORE_GATE: + Z_base = Z + rows[:, None] * stride_z_row + col_offsets + z = tl.load(Z_base, mask=mask, other=0.0).to(tl.float32) + y *= z * tl.sigmoid(z) + + # Write output + tl.store(Y_base, y, mask=mask) + diff --git a/benchmarks/TritonNPUBench/sglang/_moe_sum_reduce_kernel/_moe_sum_reduce_kernel.py b/benchmarks/TritonNPUBench/sglang/_moe_sum_reduce_kernel/_moe_sum_reduce_kernel.py new file mode 100644 index 00000000..9195f613 --- /dev/null +++ b/benchmarks/TritonNPUBench/sglang/_moe_sum_reduce_kernel/_moe_sum_reduce_kernel.py @@ -0,0 +1,60 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/sglang_operator_cases/newtest_cases/test__moe_sum_reduce_kernel.py +# Main kernel: _moe_sum_reduce_kernel +# PT file: test__moe_sum_reduce_kernel_v2.pt + +import triton +import triton.language as tl + + +# === _moe_sum_reduce_kernel === +@triton.jit +def _moe_sum_reduce_kernel( + input_ptr, + input_stride_0, + input_stride_1, + input_stride_2, + output_ptr, + output_stride_0, + output_stride_1, + token_num: int, + topk_num: int, + hidden_dim: int, + routed_scaling_factor: tl.constexpr, + BLOCK_M: tl.constexpr, + BLOCK_DIM: tl.constexpr, + NUM_STAGE: tl.constexpr, +): + input_stride_0 = tl.cast(input_stride_0, dtype=tl.int64) + input_stride_1 = tl.cast(input_stride_1, dtype=tl.int64) + output_stride_0 = tl.cast(output_stride_0, dtype=tl.int64) + + token_block_id = tl.program_id(0) + dim_block_id = tl.program_id(1) + + offs_token = token_block_id * BLOCK_M + tl.arange(0, BLOCK_M) + offs_dim = dim_block_id * BLOCK_DIM + tl.arange(0, BLOCK_DIM) + + mask_token = offs_token < token_num + mask_dim = offs_dim < hidden_dim + + base_ptrs = input_ptr + offs_token[:, None] * input_stride_0 + offs_dim[None, :] + + accumulator = tl.zeros((BLOCK_M, BLOCK_DIM), dtype=tl.float32) + for i in tl.range(0, topk_num, num_stages=NUM_STAGE): + tile = tl.load( + base_ptrs + i * input_stride_1, + mask=mask_token[:, None] & mask_dim[None, :], + other=0.0, + ) + accumulator += tile.to(tl.float32) + accumulator *= routed_scaling_factor + + # -------- Write back -------- + store_ptrs = output_ptr + offs_token[:, None] * output_stride_0 + offs_dim[None, :] + tl.store( + store_ptrs, + accumulator.to(input_ptr.dtype.element_ty), + mask=mask_token[:, None] & mask_dim[None, :], + ) + diff --git a/benchmarks/TritonNPUBench/sglang/_qkv_lora_b_kernel/_qkv_lora_b_kernel.py b/benchmarks/TritonNPUBench/sglang/_qkv_lora_b_kernel/_qkv_lora_b_kernel.py new file mode 100644 index 00000000..7ab87a0a --- /dev/null +++ b/benchmarks/TritonNPUBench/sglang/_qkv_lora_b_kernel/_qkv_lora_b_kernel.py @@ -0,0 +1,128 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/sglang_operator_cases/newtest_cases/test__qkv_lora_b_kernel.py +# Main kernel: _qkv_lora_b_kernel +# PT file: test__qkv_lora_b_kernel_v2.pt + +import triton +import triton.language as tl + + +# === _qkv_lora_b_kernel === +@triton.jit +def _qkv_lora_b_kernel( + # Pointers to matrices + x, + weights, + output, + # Parameters of size + K, # K = R + max_qkv_out_dim, # max(output_q_dim, output_kv_dim) + # Strides + x_stride_0, + x_stride_1, + w_stride_0, + w_stride_1, + w_stride_2, + output_stride_0, + output_stride_1, + # Information on sequence lengths and weight id + seg_lens, + seg_indptr, + weight_indices, + lora_ranks, + # Offsets of q/k/v slice on output dimension + n_offs, + # Meta parameters + BLOCK_S: tl.constexpr, + BLOCK_N: tl.constexpr, + BLOCK_K: tl.constexpr, + # For fused output scaling + scalings, +): + """ + This kernel packs 3 sgemms (q/k/v) into a single kernel. The multiplication + results are accumulated into the output tensor. + + When a sequence's rank is 0, the kernel is essentially a no-op, following + the convention in pytorch where the product of two matrices of shape (m, 0) + and (0, n) is an all-zero matrix of shape (m, n). + + Args: + x (Tensor): The input tensor, which is the result of the LoRA A projection. + Shape: (s, 3 * K), where s is the sum of all sequence lengths in the + batch and K is the maximum LoRA rank. The second dimension is partitioned + for Q, K, and V. + weights (Tensor): The LoRA B weights for all adapters. + Shape: (num_lora, N_Q + 2 * N_KV, K). + output (Tensor): The output tensor where the result is stored. + Shape: (s, N_Q + 2 * N_KV). + """ + + # Current block computes sequence with batch_id, + # which starts from row seg_start of x with length seg_len. + # qkv_id decides which of q,k,v to compute (0: q, 1: k, 2: v) + batch_id = tl.program_id(axis=2) + w_index = tl.load(weight_indices + batch_id) + rank = tl.load(lora_ranks + w_index) + + # If rank is 0, this kernel is a no-op. + if rank == 0: + return + + qkv_id = tl.program_id(axis=1) + pid = tl.program_id(axis=0) + seg_len = tl.load(seg_lens + batch_id) + seg_start = tl.load(seg_indptr + batch_id) + n_start = tl.load(n_offs + qkv_id) + n_size = tl.load(n_offs + qkv_id + 1) - n_start + scaling = tl.load(scalings + w_index) + # Adjust K (rank) according to the specific LoRA adapter + K = tl.minimum(K, rank) + + # The tile in output matrix will have (pid_s, pid_n) as id + num_pid_n = tl.cdiv(max_qkv_out_dim, BLOCK_N) + pid_s = pid // num_pid_n + pid_n = pid % num_pid_n + + # Create pointers for the first block of x and weights[batch_id][n_start: n_end][:] + # The pointers will be advanced as we move in the K direction + # and accumulate + s_offset = tl.arange(0, BLOCK_S) + pid_s * BLOCK_S + n_offset = tl.arange(0, BLOCK_N) + pid_n * BLOCK_N + k_offset = tl.arange(0, BLOCK_K) + + x_ptrs = (x + seg_start * x_stride_0 + (qkv_id * K) * x_stride_1) + ( + s_offset[:, None] * x_stride_0 + k_offset[None, :] * x_stride_1 + ) + w_ptrs = (weights + w_index * w_stride_0 + n_start * w_stride_1) + ( + k_offset[:, None] * w_stride_2 + n_offset[None, :] * w_stride_1 + ) + + # Iterate to compute the block in output matrix + partial_sum = tl.zeros((BLOCK_S, BLOCK_N), dtype=tl.float32) + for k in range(0, tl.cdiv(K, BLOCK_K)): + x_tile = tl.load( + x_ptrs, + mask=(s_offset[:, None] < seg_len) & (k_offset[None, :] < K - k * BLOCK_K), + other=0.0, + ) + w_tile = tl.load( + w_ptrs, + mask=(k_offset[:, None] < K - k * BLOCK_K) & (n_offset[None, :] < n_size), + other=0.0, + ) + partial_sum += tl.dot(x_tile, w_tile) + + x_ptrs += BLOCK_K * x_stride_1 + w_ptrs += BLOCK_K * w_stride_2 + + # Store result to output matrix + partial_sum *= scaling + partial_sum = partial_sum.to(x.dtype.element_ty) + output_ptr = (output + seg_start * output_stride_0 + n_start * output_stride_1) + ( + s_offset[:, None] * output_stride_0 + n_offset[None, :] * output_stride_1 + ) + output_mask = (s_offset[:, None] < seg_len) & (n_offset[None, :] < n_size) + partial_sum += tl.load(output_ptr, mask=output_mask) + tl.store(output_ptr, partial_sum, mask=output_mask) + diff --git a/benchmarks/TritonNPUBench/sglang/_quantize_k_cache_fast_kernel/_quantize_k_cache_fast_kernel.py b/benchmarks/TritonNPUBench/sglang/_quantize_k_cache_fast_kernel/_quantize_k_cache_fast_kernel.py new file mode 100644 index 00000000..dbf6f2dd --- /dev/null +++ b/benchmarks/TritonNPUBench/sglang/_quantize_k_cache_fast_kernel/_quantize_k_cache_fast_kernel.py @@ -0,0 +1,69 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/sglang_operator_cases/newtest_cases/test__quantize_k_cache_fast_kernel.py +# Main kernel: _quantize_k_cache_fast_kernel +# PT file: test__quantize_k_cache_fast_kernel_v2.pt + +import triton +import triton.language as tl + + +# === _quantize_k_cache_fast_kernel === +@triton.jit +def _quantize_k_cache_fast_kernel( + output_nope_q_ptr, + output_nope_s_ptr, + output_rope_ptr, + k_nope_ptr, + k_rope_ptr, + output_nope_q_stride_0: int, + output_nope_s_stride_0: int, + output_rope_stride_0: int, + k_nope_stride_0: int, + k_rope_stride_0: int, + NUM_NOPE_BLOCKS: tl.constexpr, + GROUP_SIZE: tl.constexpr, + DIM_NOPE: tl.constexpr, + DIM_ROPE: tl.constexpr, + FP8_MIN: tl.constexpr, + FP8_MAX: tl.constexpr, +): + token_id = tl.program_id(0) + raw_block_id = tl.program_id(1) + + if raw_block_id < NUM_NOPE_BLOCKS: + # a. quant nope + effective_block_id = raw_block_id + + offs = effective_block_id * GROUP_SIZE + tl.arange(0, GROUP_SIZE) + mask = offs < DIM_NOPE + ptr = k_nope_ptr + token_id * k_nope_stride_0 + offs + + y = tl.load(ptr, mask=mask, other=0.0).to(tl.float32) + + # the ref impl do not have a `tl.maximum(... eps)`, so we remove it here + y_s = tl.max(tl.abs(y)) / FP8_MAX + y_s_inv = 1.0 / y_s + y_q = tl.clamp(y * y_s_inv, FP8_MIN, FP8_MAX).to( + output_nope_q_ptr.dtype.element_ty + ) + + dst_q_ptr = output_nope_q_ptr + token_id * output_nope_q_stride_0 + offs + dst_s_ptr = ( + output_nope_s_ptr + token_id * output_nope_s_stride_0 + effective_block_id + ) + + tl.store(dst_q_ptr, y_q, mask=mask) + tl.store(dst_s_ptr, y_s) + else: + # b. copy rope + effective_block_id = raw_block_id - NUM_NOPE_BLOCKS + + offs = effective_block_id * GROUP_SIZE + tl.arange(0, GROUP_SIZE) + mask = offs < DIM_ROPE + + src_ptr = k_rope_ptr + token_id * k_rope_stride_0 + offs + dst_ptr = output_rope_ptr + token_id * output_rope_stride_0 + offs + + data = tl.load(src_ptr, mask=mask) + tl.store(dst_ptr, data, mask=mask) + diff --git a/benchmarks/TritonNPUBench/sglang/_selective_scan_update_kernel/_selective_scan_update_kernel.py b/benchmarks/TritonNPUBench/sglang/_selective_scan_update_kernel/_selective_scan_update_kernel.py new file mode 100644 index 00000000..8f4c9f04 --- /dev/null +++ b/benchmarks/TritonNPUBench/sglang/_selective_scan_update_kernel/_selective_scan_update_kernel.py @@ -0,0 +1,235 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/sglang_operator_cases/newtest_cases/test__selective_scan_update_kernel.py +# Main kernel: _selective_scan_update_kernel +# PT file: test__selective_scan_update_kernel_v2.pt + +import triton +import triton.language as tl + + +# === _selective_scan_update_kernel === +@triton.jit(do_not_specialize=["N"]) +def _selective_scan_update_kernel( + # Pointers to matrices + state_ptr, + x_ptr, + dt_ptr, + dt_bias_ptr, + A_ptr, + B_ptr, + C_ptr, + D_ptr, + z_ptr, + out_ptr, + state_batch_indices_ptr, + dst_state_batch_indices_ptr, + pad_slot_id, + num_accepted_tokens_ptr, + cu_seqlens_ptr, + # Matrix dimensions + N, + nheads, + dim, + dstate, + nheads_ngroups_ratio, + # Strides + stride_state_batch, + stride_state_head, + stride_state_dim, + stride_state_dstate, + stride_x_batch, + stride_x_head, + stride_x_dim, + stride_dt_batch, + stride_dt_head, + stride_dt_dim, + stride_dt_bias_head, + stride_dt_bias_dim, + stride_A_head, + stride_A_dim, + stride_A_dstate, + stride_B_batch, + stride_B_group, + stride_B_dstate, + stride_C_batch, + stride_C_group, + stride_C_dstate, + stride_D_head, + stride_D_dim, + stride_z_batch, + stride_z_head, + stride_z_dim, + stride_out_batch, + stride_out_head, + stride_out_dim, + stride_state_indices_batch, + stride_state_indices_T, + stride_dst_state_indices_batch, + stride_dst_state_indices_T, + # Meta-parameters + DT_SOFTPLUS: tl.constexpr, + TIE_HDIM: tl.constexpr, + BLOCK_SIZE_M: tl.constexpr, + HAS_DT_BIAS: tl.constexpr, + HAS_D: tl.constexpr, + HAS_Z: tl.constexpr, + HAS_STATE_BATCH_INDICES: tl.constexpr, + IS_SPEC_DECODING: tl.constexpr, + IS_VARLEN: tl.constexpr, + BLOCK_SIZE_DSTATE: tl.constexpr, +): + pid_m = tl.program_id(axis=0) + pid_b = tl.program_id(axis=1) + pid_h = tl.program_id(axis=2) + + if IS_VARLEN: + bos = tl.load(cu_seqlens_ptr + pid_b).to(tl.int64) + eos = tl.load(cu_seqlens_ptr + pid_b + 1).to(tl.int64) + seq_len = eos - bos + + if seq_len == 0: + return + else: + bos = pid_b + seq_len = 1 + + state_ptr_base = state_ptr + + # If HAS_STATE_BATCH_INDICES is true, then the ssm state's batch coordinate + # is taken from the state_batch_indices_ptr Otherwise, the state coordinate + # is the same as the batch id. + if HAS_STATE_BATCH_INDICES: + if IS_SPEC_DECODING: + num_accepted = tl.load(num_accepted_tokens_ptr + pid_b).to(tl.int64) + init_token_idx = tl.maximum(num_accepted - 1, 0) + else: + init_token_idx = 0 + + dst_state_batch_indices_ptr += pid_b * stride_dst_state_indices_batch + if not IS_SPEC_DECODING: + dst_state_batch_idx = tl.load( + dst_state_batch_indices_ptr + + init_token_idx * stride_dst_state_indices_T + ).to(tl.int64) + dst_state_ptr = state_ptr + ( + dst_state_batch_idx * stride_state_batch + pid_h * stride_state_head + ) + + state_batch_indices_ptr += ( + pid_b * stride_state_indices_batch + init_token_idx * stride_state_indices_T + ) + state_batch_idx = tl.load(state_batch_indices_ptr).to(tl.int64) + state_ptr += state_batch_idx * stride_state_batch + pid_h * stride_state_head + else: + dst_state_ptr = ( + state_ptr + pid_b * stride_state_batch + pid_h * stride_state_head + ) + state_ptr += pid_b * stride_state_batch + pid_h * stride_state_head + + x_ptr += bos * stride_x_batch + pid_h * stride_x_head + dt_ptr += bos * stride_dt_batch + pid_h * stride_dt_head + if HAS_DT_BIAS: + dt_bias_ptr += pid_h * stride_dt_bias_head + A_ptr += pid_h * stride_A_head + B_ptr += bos * stride_B_batch + (pid_h // nheads_ngroups_ratio) * stride_B_group + C_ptr += bos * stride_C_batch + (pid_h // nheads_ngroups_ratio) * stride_C_group + if HAS_Z: + z_ptr += bos * stride_z_batch + pid_h * stride_z_head + out_ptr += bos * stride_out_batch + pid_h * stride_out_head + + offs_m = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) + offs_n = tl.arange(0, BLOCK_SIZE_DSTATE) + state_ptrs = state_ptr + ( + offs_m[:, None] * stride_state_dim + offs_n[None, :] * stride_state_dstate + ) + if not IS_SPEC_DECODING: + dst_state_ptrs = dst_state_ptr + ( + offs_m[:, None] * stride_state_dim + offs_n[None, :] * stride_state_dstate + ) + + mask = (offs_m[:, None] < dim) & (offs_n[None, :] < dstate) + if HAS_STATE_BATCH_INDICES: + mask &= state_batch_idx != pad_slot_id + state = tl.load(state_ptrs, mask=mask, other=0.0).to(tl.float32) + + if HAS_DT_BIAS: + dt_bias_ptrs = dt_bias_ptr + offs_m * stride_dt_bias_dim + if HAS_D: + D_ptr += pid_h * stride_D_head + D_ptrs = D_ptr + offs_m * stride_D_dim + A_ptrs = A_ptr + offs_m[:, None] * stride_A_dim + offs_n[None, :] * stride_A_dstate + + for i_t in range(seq_len): + x_ptrs = x_ptr + offs_m * stride_x_dim + dt_ptrs = dt_ptr + offs_m * stride_dt_dim + B_ptrs = B_ptr + offs_n * stride_B_dstate + C_ptrs = C_ptr + offs_n * stride_C_dstate + if HAS_Z: + z_ptrs = z_ptr + offs_m * stride_z_dim + out_ptrs = out_ptr + offs_m * stride_out_dim + + x = tl.load(x_ptrs, mask=offs_m < dim, other=0.0).to(tl.float32) + if not TIE_HDIM: + dt = tl.load(dt_ptrs, mask=offs_m < dim, other=0.0).to(tl.float32) + if HAS_DT_BIAS: + dt += tl.load(dt_bias_ptrs, mask=offs_m < dim, other=0.0).to(tl.float32) + if DT_SOFTPLUS: + dt = softplus(dt) + A = tl.load( + A_ptrs, + mask=(offs_m[:, None] < dim) & (offs_n[None, :] < dstate), + other=0.0, + ).to(tl.float32) + dA = tl.exp(A * dt[:, None]) + else: + dt = tl.load(dt_ptr).to(tl.float32) + if HAS_DT_BIAS: + dt += tl.load(dt_bias_ptr).to(tl.float32) + if DT_SOFTPLUS: + dt = softplus(dt) + A = tl.load(A_ptr).to(tl.float32) + dA = tl.exp(A * dt) # scalar, not a matrix + + B = tl.load(B_ptrs, mask=offs_n < dstate, other=0.0).to(tl.float32) + C = tl.load(C_ptrs, mask=offs_n < dstate, other=0.0).to(tl.float32) + if HAS_D: + D = tl.load(D_ptrs, mask=offs_m < dim, other=0.0).to(tl.float32) + if HAS_Z: + z = tl.load(z_ptrs, mask=offs_m < dim, other=0.0).to(tl.float32) + + dB = B[None, :] * dt[:, None] if not TIE_HDIM else B * dt + state = state * dA + dB * x[:, None] + + if IS_SPEC_DECODING: + dst_idx_ptr = dst_state_batch_indices_ptr + i_t * stride_dst_state_indices_T + token_dst_idx = tl.load(dst_idx_ptr).to(tl.int64) + if token_dst_idx != pad_slot_id: + token_dst_ptrs = ( + state_ptr_base + + token_dst_idx * stride_state_batch + + pid_h * stride_state_head + + offs_m[:, None] * stride_state_dim + + offs_n[None, :] * stride_state_dstate + ) + tl.store( + token_dst_ptrs, state.to(token_dst_ptrs.dtype.element_ty), mask=mask + ) + + out = tl.sum(state * C[None, :], axis=1) + if HAS_D: + out += x * D + if HAS_Z: + out *= z * tl.sigmoid(z) + tl.store(out_ptrs, out, mask=offs_m < dim) + + x_ptr += stride_x_batch + dt_ptr += stride_dt_batch + B_ptr += stride_B_batch + C_ptr += stride_C_batch + out_ptr += stride_out_batch + if HAS_Z: + z_ptr += stride_z_batch + + if not IS_SPEC_DECODING: + tl.store(dst_state_ptrs, state.to(dst_state_ptrs.dtype.element_ty), mask=mask) + diff --git a/benchmarks/TritonNPUBench/sglang/_sgemm_lora_a_kernel/_sgemm_lora_a_kernel.py b/benchmarks/TritonNPUBench/sglang/_sgemm_lora_a_kernel/_sgemm_lora_a_kernel.py new file mode 100644 index 00000000..5d1adf0c --- /dev/null +++ b/benchmarks/TritonNPUBench/sglang/_sgemm_lora_a_kernel/_sgemm_lora_a_kernel.py @@ -0,0 +1,115 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/sglang_operator_cases/newtest_cases/test__sgemm_lora_a_kernel.py +# Main kernel: _sgemm_lora_a_kernel +# PT file: test__sgemm_lora_a_kernel_v2.pt + +import triton +import triton.language as tl + + +# === _sgemm_lora_a_kernel === +@triton.jit +def _sgemm_lora_a_kernel( + # Pointers to matrices + x, + weights, + output, + # Matrix dimensions + N, # stack_num * r + K, # input_dim + stack_num, + # Strides + x_stride_0, + x_stride_1, + w_stride_0, + w_stride_1, + w_stride_2, + output_stride_0, + output_stride_1, + # Information on sequence lengths,ranks and weight id + seg_lens, + seg_indptr, + weight_indices, + lora_ranks, + # Meta parameters + BLOCK_S: tl.constexpr, + BLOCK_N: tl.constexpr, + BLOCK_K: tl.constexpr, +): + """ + Computes a segmented batched matrix multiplication for the LoRA A matrix. + + The kernel ensures that output[seg_start:seg_start + seg_len, :rank * stack_num] + stores the product of the input `x` and the LoRA weights for the corresponding + sequence. This implies that when rank is 0, the kernel is essentially a no-op, + as output[seg_start:seg_start + seg_len, :0] is trivially correct (empty). + + Args: + x (torch.Tensor): The input activations tensor of shape `(s, K)`, where `s` + is the sum of all sequence lengths in the batch. + weights (torch.Tensor): The LoRA 'A' weights for all available adapters, + with shape `(num_lora, N, K)`. + output (torch.Tensor): The output tensor of shape `(s, N)`. + """ + + # Current block computes sequence with batch_id, + # which starts from row seg_start of x with length seg_len + batch_id = tl.program_id(axis=1) + w_index = tl.load(weight_indices + batch_id) + rank = tl.load(lora_ranks + w_index) + + # If rank is 0, this kernel becomes a no-op as the output is always trivially correct. + if rank == 0: + return + + pid = tl.program_id(axis=0) + seg_start = tl.load(seg_indptr + batch_id) + seg_len = tl.load(seg_lens + batch_id) + + # Adjust N (stack_num * max_rank) according to the specific LoRA adapter + N = tl.minimum(N, rank * stack_num) + + # The tile in output matrix will have (pid_s, pid_n) as id + num_pid_n = tl.cdiv(N, BLOCK_N) + pid_s = pid // num_pid_n + pid_n = pid % num_pid_n + + # Create pointers for the first block of x and weights[batch_id] + # The pointers will be advanced as we move in the K direction + # and accumulate + s_offset = tl.arange(0, BLOCK_S) + pid_s * BLOCK_S + n_offset = tl.arange(0, BLOCK_N) + pid_n * BLOCK_N + k_offset = tl.arange(0, BLOCK_K) + x_ptrs = (x + seg_start * x_stride_0) + ( + s_offset[:, None] * x_stride_0 + k_offset[None, :] * x_stride_1 + ) + w_ptrs = (weights + w_index * w_stride_0) + ( + k_offset[:, None] * w_stride_2 + n_offset[None, :] * w_stride_1 + ) + + # Iterate to compute the block in output matrix + partial_sum = tl.zeros((BLOCK_S, BLOCK_N), dtype=tl.float32) + for k in range(0, tl.cdiv(K, BLOCK_K)): + x_tile = tl.load( + x_ptrs, + mask=(s_offset[:, None] < seg_len) & (k_offset[None, :] < K - k * BLOCK_K), + other=0.0, + ) + w_tile = tl.load( + w_ptrs, + mask=(k_offset[:, None] < K - k * BLOCK_K) & (n_offset[None, :] < N), + other=0.0, + ) + partial_sum += tl.dot(x_tile, w_tile) + + x_ptrs += BLOCK_K * x_stride_1 + w_ptrs += BLOCK_K * w_stride_2 + + # Store result to output matrix + partial_sum = partial_sum.to(x.dtype.element_ty) + output_ptr = (output + seg_start * output_stride_0) + ( + s_offset[:, None] * output_stride_0 + n_offset[None, :] * output_stride_1 + ) + output_mask = (s_offset[:, None] < seg_len) & (n_offset[None, :] < N) + tl.store(output_ptr, partial_sum, mask=output_mask) + diff --git a/benchmarks/TritonNPUBench/sglang/_sgemm_lora_b_kernel/_sgemm_lora_b_kernel.py b/benchmarks/TritonNPUBench/sglang/_sgemm_lora_b_kernel/_sgemm_lora_b_kernel.py new file mode 100644 index 00000000..859c80f6 --- /dev/null +++ b/benchmarks/TritonNPUBench/sglang/_sgemm_lora_b_kernel/_sgemm_lora_b_kernel.py @@ -0,0 +1,119 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/sglang_operator_cases/newtest_cases/test__sgemm_lora_b_kernel.py +# Main kernel: _sgemm_lora_b_kernel +# PT file: test__sgemm_lora_b_kernel_v2.pt + +import triton +import triton.language as tl + + +# === _sgemm_lora_b_kernel === +@triton.jit +def _sgemm_lora_b_kernel( + # Pointers to matrices + x, + weights, + output, + # Matrix dimensions + N, # output_dim + K, # r + # Strides + x_stride_0, + x_stride_1, + w_stride_0, + w_stride_1, + w_stride_2, + output_stride_0, + output_stride_1, + # Information on sequence lengths and weight id + seg_lens, + seg_indptr, + weight_indices, + lora_ranks, + # Meta parameters + BLOCK_S: tl.constexpr, + BLOCK_N: tl.constexpr, + BLOCK_K: tl.constexpr, + # For fused output scaling + scalings, +): + """ + Computes a segmented batched matrix multiplication for the LoRA B matrix + and adds the result to the output in-place. + + When a sequence's rank is 0, the kernel is essentially a no-op, following + the convention in pytorch where the product of two matrices of shape (m, 0) + and (0, n) is an all-zero matrix of shape (m, n). + + Args: + x (torch.Tensor): The intermediate tensor from the LoRA 'A' multiplication, + of shape `(s, K)`, where `s` is the total number of tokens. + weights (torch.Tensor): The LoRA 'B' weights for all available adapters, + with shape `(num_lora, N, K)`. + output (torch.Tensor): The output tensor of shape `(s, N)`. This can be + the base model's output for a fused add operation. + """ + + # Current block computes sequence with batch_id, + # which starts from row seg_start of x with length seg_len + batch_id = tl.program_id(axis=1) + w_index = tl.load(weight_indices + batch_id) + rank = tl.load(lora_ranks + w_index) + + # If rank is 0, this kernel is a no-op. + if rank == 0: + return + + pid = tl.program_id(axis=0) + seg_len = tl.load(seg_lens + batch_id) + seg_start = tl.load(seg_indptr + batch_id) + scaling = tl.load(scalings + w_index) + # Adjust K (rank) according to the specific LoRA adapter + K = tl.minimum(K, rank) + + # The tile in output matrix will have (pid_s, pid_n) as id + num_pid_n = tl.cdiv(N, BLOCK_N) + pid_s = pid // num_pid_n + pid_n = pid % num_pid_n + + # Create pointers for the first block of x and weights[batch_id] + # The pointers will be advanced as we move in the K direction + # and accumulate + s_offset = tl.arange(0, BLOCK_S) + pid_s * BLOCK_S + n_offset = tl.arange(0, BLOCK_N) + pid_n * BLOCK_N + k_offset = tl.arange(0, BLOCK_K) + x_ptrs = (x + seg_start * x_stride_0) + ( + s_offset[:, None] * x_stride_0 + k_offset[None, :] * x_stride_1 + ) + w_ptrs = (weights + w_index * w_stride_0) + ( + k_offset[:, None] * w_stride_2 + n_offset[None, :] * w_stride_1 + ) + + # Iterate to compute the block in output matrix + partial_sum = tl.zeros((BLOCK_S, BLOCK_N), dtype=tl.float32) + for k in range(0, tl.cdiv(K, BLOCK_K)): + x_tile = tl.load( + x_ptrs, + mask=(s_offset[:, None] < seg_len) & (k_offset[None, :] < K - k * BLOCK_K), + other=0.0, + ) + w_tile = tl.load( + w_ptrs, + mask=(k_offset[:, None] < K - k * BLOCK_K), + other=0.0, + ) + partial_sum += tl.dot(x_tile, w_tile) + + x_ptrs += BLOCK_K * x_stride_1 + w_ptrs += BLOCK_K * w_stride_2 + + # Store result to output matrix + partial_sum *= scaling + partial_sum = partial_sum.to(x.dtype.element_ty) + output_ptr = (output + seg_start * output_stride_0) + ( + s_offset[:, None] * output_stride_0 + n_offset[None, :] * output_stride_1 + ) + output_mask = s_offset[:, None] < seg_len + partial_sum += tl.load(output_ptr, mask=output_mask) + tl.store(output_ptr, partial_sum, mask=output_mask) + diff --git a/benchmarks/TritonNPUBench/sglang/_silu_and_mul_post_per_tensor_quant_kernel/_silu_and_mul_post_per_tensor_quant_kernel.py b/benchmarks/TritonNPUBench/sglang/_silu_and_mul_post_per_tensor_quant_kernel/_silu_and_mul_post_per_tensor_quant_kernel.py new file mode 100644 index 00000000..cdd07256 --- /dev/null +++ b/benchmarks/TritonNPUBench/sglang/_silu_and_mul_post_per_tensor_quant_kernel/_silu_and_mul_post_per_tensor_quant_kernel.py @@ -0,0 +1,75 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/sglang_operator_cases/newtest_cases/test__silu_and_mul_post_per_tensor_quant_kernel.py +# Main kernel: _silu_and_mul_post_per_tensor_quant_kernel +# PT file: test__silu_and_mul_post_per_tensor_quant_kernel_v2.pt + +import triton +import triton.language as tl + + +# === _silu_and_mul_post_per_tensor_quant_kernel === +@triton.jit +def _silu_and_mul_post_per_tensor_quant_kernel( + input_ptr, + stride_input_expert, + stride_input_token, + stride_input_dim, + output_ptr, + stride_output_expert, + stride_output_token, + stride_output_dim, + scale_ptr, + masked_m_ptr, + inner_dim, + fp8_max, + fp8_min, + BLOCK_N: tl.constexpr, + NUM_STAGE: tl.constexpr, +): + """ + Triton kernel: fused SiLU(gate) * up + per-tensor FP8 quantization. + + Shape: + input: [E, T_padded, 2*D] -> gate: [:,:,D], up: [:,:,D] + output: [E, T_padded, D], dtype=float8_e4m3fn + """ + expert_id = tl.program_id(2) + block_id_token = tl.program_id(1) + block_id_dim = tl.program_id(0) + + num_token_blocks = tl.num_programs(1) + + token_num_cur_expert = tl.load(masked_m_ptr + expert_id) + + scale = 1.0 / tl.load(scale_ptr).to(tl.float32) + + stride_input_expert = tl.cast(stride_input_expert, tl.int32) + stride_output_expert = tl.cast(stride_output_expert, tl.int32) + stride_input_token = tl.cast(stride_input_token, tl.int32) + stride_output_token = tl.cast(stride_output_token, tl.int32) + + offset_d = block_id_dim * BLOCK_N + tl.arange(0, BLOCK_N) + mask_d = offset_d < inner_dim + + # base pointers for current expert and dim block + input_base_offs = input_ptr + expert_id * stride_input_expert + offset_d + output_base_offs = output_ptr + expert_id * stride_output_expert + offset_d + + for token_idx in tl.range( + block_id_token, token_num_cur_expert, num_token_blocks, num_stages=NUM_STAGE + ): + gate_ptr = input_base_offs + token_idx * stride_input_token + up_ptr = gate_ptr + inner_dim + gate = tl.load(gate_ptr, mask=mask_d, other=0.0).to(tl.float32) + up = tl.load(up_ptr, mask=mask_d, other=0.0).to(tl.float32) + + # SiLU: x * sigmoid(x) + gate = gate / (1 + tl.exp(-gate)) + gate = gate.to(input_ptr.dtype.element_ty) + gate_up = up * gate + + scaled = gate_up * scale + output_q = tl.clamp(scaled, fp8_min, fp8_max).to(output_ptr.dtype.element_ty) + out_ptr = output_base_offs + token_idx * stride_output_token + tl.store(out_ptr, output_q, mask=mask_d) + diff --git a/benchmarks/TritonNPUBench/sglang/_silu_and_mul_post_quant_kernel/_silu_and_mul_post_quant_kernel.py b/benchmarks/TritonNPUBench/sglang/_silu_and_mul_post_quant_kernel/_silu_and_mul_post_quant_kernel.py new file mode 100644 index 00000000..ba595683 --- /dev/null +++ b/benchmarks/TritonNPUBench/sglang/_silu_and_mul_post_quant_kernel/_silu_and_mul_post_quant_kernel.py @@ -0,0 +1,87 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/sglang_operator_cases/newtest_cases/test__silu_and_mul_post_quant_kernel.py +# Main kernel: _silu_and_mul_post_quant_kernel +# PT file: test__silu_and_mul_post_quant_kernel_v2.pt + +import triton +import triton.language as tl + + +# === _silu_and_mul_post_quant_kernel === +@triton.jit +def _silu_and_mul_post_quant_kernel( + input_ptr, + stride_input_0, + stride_input_1, + stride_input_2, + output_ptr, + stride_output_0, + stride_output_1, + stride_output_2, + output_scale_ptr, + stride_output_scale_0, + stride_output_scale_1, + stride_output_scale_2, + masked_m_ptr, + size_n, + fp8_max, + fp8_min, + BLOCK_N: tl.constexpr, + NUM_STAGE: tl.constexpr, + SCALE_UE8M0: tl.constexpr, +): + expert_id = tl.program_id(2) + token_id = tl.program_id(1) + hidden_dim_block_index = tl.program_id(0) + + block_num_per_expert = tl.num_programs(1) + + token_num_cur_expert = tl.load(masked_m_ptr + expert_id) + + stride_input_0 = tl.cast(stride_input_0, dtype=tl.int64) + stride_output_0 = tl.cast(stride_output_0, dtype=tl.int64) + stride_input_1 = tl.cast(stride_input_1, dtype=tl.int64) + stride_output_1 = tl.cast(stride_output_1, dtype=tl.int64) + + offs_in_d = hidden_dim_block_index * BLOCK_N + tl.arange(0, BLOCK_N) + input_ptr_offs = input_ptr + expert_id * stride_input_0 + offs_in_d + output_ptr_offs = output_ptr + expert_id * stride_output_0 + offs_in_d + output_scale_offs = ( + output_scale_ptr + + expert_id * stride_output_scale_0 + + hidden_dim_block_index * stride_output_scale_2 + ) + + for token_index in tl.range( + token_id, token_num_cur_expert, block_num_per_expert, num_stages=NUM_STAGE + ): + gate = tl.load( + input_ptr_offs + token_index * stride_input_1, + mask=offs_in_d < size_n, + other=0.0, + ).to(tl.float32) + up = tl.load( + input_ptr_offs + token_index * stride_input_1 + size_n, + mask=offs_in_d < size_n, + other=0.0, + ) + gate = gate / (1 + tl.exp(-gate)) + gate = gate.to(input_ptr.dtype.element_ty) + gate_up = up * gate + _absmax = tl.maximum(tl.max(tl.abs(gate_up)), 1e-10) + output_s = _absmax / fp8_max + if SCALE_UE8M0: + output_s = tl.exp2(tl.ceil(tl.log2(tl.abs(output_s)))) + output_q = tl.clamp(gate_up / output_s, fp8_min, fp8_max).to( + output_ptr.dtype.element_ty + ) + tl.store( + output_ptr_offs + token_index * stride_output_1, + output_q, + mask=offs_in_d < size_n, + ) + tl.store( + output_scale_offs + token_index * stride_output_scale_1, + output_s, + ) + diff --git a/benchmarks/TritonNPUBench/sglang/_state_passing_fwd_kernel/_state_passing_fwd_kernel.py b/benchmarks/TritonNPUBench/sglang/_state_passing_fwd_kernel/_state_passing_fwd_kernel.py new file mode 100644 index 00000000..d67c39a7 --- /dev/null +++ b/benchmarks/TritonNPUBench/sglang/_state_passing_fwd_kernel/_state_passing_fwd_kernel.py @@ -0,0 +1,164 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/sglang_operator_cases/newtest_cases/test__state_passing_fwd_kernel.py +# Main kernel: _state_passing_fwd_kernel +# PT file: test__state_passing_fwd_kernel_v2.pt + +import triton +import triton.language as tl + + +# === _state_passing_fwd_kernel === +@triton.jit +def _state_passing_fwd_kernel( + # Pointers to matrices + states_ptr, + out_ptr, + final_states_ptr, + dA_cs_ptr, + initstates_ptr, + seq_idx_ptr, + chunk_offsets_ptr, + chunk_meta_num, + # Matrix dimensions + dim, + nchunks, + seqlen, + chunk_size, + # Strides + stride_states_batch, + stride_states_chunk, + stride_states_head, + stride_states_dim, + stride_out_batch, + stride_out_chunk, + stride_out_head, + stride_out_dim, + stride_final_states_batch, + stride_final_states_head, + stride_final_states_dim, + stride_dA_cs_batch, + stride_dA_cs_chunk, + stride_dA_cs_head, + stride_dA_cs_csize, + stride_initstates_batch, + stride_initstates_head, + stride_initstates_dim, + stride_seq_idx_batch, + stride_seq_idx_seqlen, + # Meta-parameters + HAS_INITSTATES: tl.constexpr, + HAS_SEQ_IDX: tl.constexpr, + IS_CONT_BATCHED: tl.constexpr, + BLOCK_SIZE: tl.constexpr = 16, +): + pid_b = tl.program_id(axis=1) + pid_h = tl.program_id(axis=2) + pid_m = tl.program_id(axis=0) + states_ptr += pid_b * stride_states_batch + pid_h * stride_states_head + dA_cs_ptr += ( + pid_b * stride_dA_cs_batch + + pid_h * stride_dA_cs_head + + (chunk_size - 1) * stride_dA_cs_csize + ) + out_ptr += pid_b * stride_out_batch + pid_h * stride_out_head + final_states_ptr += ( + pid_b * stride_final_states_batch + pid_h * stride_final_states_head + ) + if HAS_INITSTATES: + initstates_ptr += pid_h * stride_initstates_head + if not IS_CONT_BATCHED: + initstates_ptr += pid_b * stride_initstates_batch + + if HAS_SEQ_IDX: + seq_idx_ptr += pid_b * stride_seq_idx_batch + + offs_m = pid_m * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + states_ptrs = states_ptr + offs_m * stride_states_dim + out_ptrs = out_ptr + offs_m * stride_out_dim + final_states_ptrs = final_states_ptr + offs_m * stride_final_states_dim + + # - states will be the past state of the sequence that continues on the current check + if not HAS_INITSTATES: + states = tl.zeros((BLOCK_SIZE,), dtype=tl.float32) + else: + initstates_ptr += offs_m * stride_initstates_dim + initstates_ptrs = initstates_ptr + # - for cont batches, for the first chunk mean it will be the first batch's + # init state + states = tl.load(initstates_ptrs, mask=offs_m < dim, other=0.0).to(tl.float32) + + tl.store(out_ptrs, states, mask=offs_m < dim) + out_ptrs += stride_out_chunk + prev_seq_idx_chunk_end = 0 + logical_chunk_idx = 0 + for c in range(nchunks): + new_states = tl.load(states_ptrs, mask=offs_m < dim, other=0.0).to(tl.float32) + dA_cs = tl.load(dA_cs_ptr).to(tl.float32) + scale_mask = True + if HAS_SEQ_IDX: + # - the seq to pass forward is the one that is flushed to the right + # boundary. + # - that is given by seq_idx_chunk_end below: the sequence index at the end of the chunk. + seq_idx_chunk_end = tl.load( + seq_idx_ptr + + (min((c + 1) * chunk_size, seqlen) - 1) * stride_seq_idx_seqlen + ) + if HAS_INITSTATES: + if IS_CONT_BATCHED and prev_seq_idx_chunk_end != seq_idx_chunk_end: + # this means in the current chunk the rightmost flushed seq + # has changed. + # - so we do not propagate the state from previous chunk + # - but rather we load that sequence's init state + initstates_ptrs = ( + initstates_ptr + seq_idx_chunk_end * stride_initstates_batch + ) + + # - update state with seq_idx_new's init state + states = tl.load(initstates_ptrs, mask=offs_m < dim, other=0.0).to( + tl.float32 + ) + + # - we need to consider the cumsum only of the last sequence in the chunk + # - find its starting position (given by c_off of the logical chunk index) + # - and subtract the cumsum just before that position from the total cumsum + # - first, update the logical chunk index (add the number of sequences in the current physical chunk): + # sequence index at the start of the current chunk + seq_idx_chunk_start = tl.load( + seq_idx_ptr + + min(c * chunk_size, seqlen) * stride_seq_idx_seqlen + ) + logical_chunk_idx += seq_idx_chunk_end - seq_idx_chunk_start + # - load the chunk offset: + c_off = tl.load( + chunk_offsets_ptr + logical_chunk_idx, + mask=logical_chunk_idx < chunk_meta_num, + other=0, + ) + # - if offset is 0, then the sequence starts at the beginning of the chunk, and we don't need to subtract anything + if c_off > 0: + # - dA_cs_ptr currently points to the cumsum at the end of the chunk - subtract the chunk size and add the offset + dA_cs_boundary = tl.load( + dA_cs_ptr + - (chunk_size - 1) * stride_dA_cs_csize + + (c_off - 1) * stride_dA_cs_csize, + mask=(c_off - 1) > -1 and c_off < chunk_size, + other=0.0, + ) + dA_cs -= dA_cs_boundary + + # - increment logical chunk index for every physical chunk + logical_chunk_idx += 1 + else: + scale_mask = seq_idx_chunk_end == prev_seq_idx_chunk_end + prev_seq_idx_chunk_end = seq_idx_chunk_end + + scale = tl.where(scale_mask, tl.exp(dA_cs), 0.0) + states = scale * states + new_states + if c < nchunks - 1: + tl.store(out_ptrs, states, mask=offs_m < dim) + else: + tl.store(final_states_ptrs, states, mask=offs_m < dim) + states_ptrs += stride_states_chunk + dA_cs_ptr += stride_dA_cs_chunk + out_ptrs += stride_out_chunk + diff --git a/benchmarks/TritonNPUBench/sglang/_tma_align_input_scale_kernel/_tma_align_input_scale_kernel.py b/benchmarks/TritonNPUBench/sglang/_tma_align_input_scale_kernel/_tma_align_input_scale_kernel.py new file mode 100644 index 00000000..7ad52474 --- /dev/null +++ b/benchmarks/TritonNPUBench/sglang/_tma_align_input_scale_kernel/_tma_align_input_scale_kernel.py @@ -0,0 +1,39 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/sglang_operator_cases/newtest_cases/test__tma_align_input_scale_kernel.py +# Main kernel: _tma_align_input_scale_kernel +# PT file: test__tma_align_input_scale_kernel_v2.pt + +import triton +import triton.language as tl + + +# === _tma_align_input_scale_kernel === +@triton.jit +def _tma_align_input_scale_kernel( + input_scale_ptr, + output_ptr, + m, + k_div_block_size, + input_scale_stride_m, + input_scale_stride_k, + output_stride_m, + output_stride_k, + BLOCK_SIZE_K: tl.constexpr, +): + pid_m = tl.program_id(axis=0) + grid_m = tl.num_programs(0) + k_offsets = tl.arange(0, BLOCK_SIZE_K) + + for m_base in range(pid_m, m, grid_m): + input_offset = ( + input_scale_ptr + + m_base * input_scale_stride_m + + k_offsets * input_scale_stride_k + ) + input_data = tl.load(input_offset, mask=k_offsets < k_div_block_size) + + output_offset = ( + output_ptr + k_offsets * output_stride_k + m_base * output_stride_m + ) + tl.store(output_offset, input_data, mask=k_offsets < k_div_block_size) + diff --git a/benchmarks/TritonNPUBench/sglang/act_and_mul_kernel/act_and_mul_kernel.py b/benchmarks/TritonNPUBench/sglang/act_and_mul_kernel/act_and_mul_kernel.py new file mode 100644 index 00000000..216f65f2 --- /dev/null +++ b/benchmarks/TritonNPUBench/sglang/act_and_mul_kernel/act_and_mul_kernel.py @@ -0,0 +1,83 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/sglang_operator_cases/newtest_cases/test_act_and_mul_kernel.py +# Main kernel: act_and_mul_kernel +# PT file: test_act_and_mul_kernel_v2.pt + +import triton +import triton.language as tl + + +# === act_and_mul_kernel === +@triton.jit +def act_and_mul_kernel( + gateup_output, + down_input, + hidden_size, + expert_ids_ptr, + expert_step: tl.constexpr, + BLOCK_SIZE: tl.constexpr, + ACTIVATION_TYPE: tl.constexpr, +): + """ + Unified activation and multiply kernel that handles both sorted and unsorted routing, + and both SiLU and GELU activations using compile-time constants. + """ + InDtype = gateup_output.dtype.element_ty + OutDtype = down_input.dtype.element_ty + + half_hidden_size = hidden_size // 2 + pid = tl.program_id(0) + + expert_id = tl.load(expert_ids_ptr + pid // expert_step) + + if expert_id == -1: + return + + gateup_output_ptr = gateup_output + pid * hidden_size + down_input_ptr = down_input + pid * half_hidden_size + gate_output_ptr = gateup_output_ptr + up_output_ptr = gateup_output_ptr + half_hidden_size + + for start_offset in tl.range(0, half_hidden_size, BLOCK_SIZE): + offset = start_offset + tl.arange(0, BLOCK_SIZE) + mask = offset < half_hidden_size + + gate_output = tl.load(gate_output_ptr + offset, mask=mask) + up_output = tl.load(up_output_ptr + offset, mask=mask) + + gate_output_activated = _apply_activation(gate_output, ACTIVATION_TYPE) + gate_output_activated = gate_output_activated.to(InDtype) + + act_mul_output = gate_output_activated * up_output + act_mul_output = act_mul_output.to(OutDtype) + tl.store(down_input_ptr + offset, act_mul_output, mask=mask) + + +# === _apply_activation === +@triton.jit +def _apply_activation(x, ACTIVATION_TYPE: tl.constexpr): + """ + Apply activation function based on compile-time constant. + + Args: + x: Input tensor (converted to float32 inside) + ACTIVATION_TYPE: Compile-time constant string ("silu" or "gelu") + + Returns: + Activated output in the same dtype as input + """ + x = x.to(tl.float32) + if ACTIVATION_TYPE == "silu": + return x * tl.sigmoid(x) + elif ACTIVATION_TYPE == "gelu": + kAlpha = 0.7978845608028654 + return 0.5 * x * (1 + tanh(kAlpha * (x + 0.044715 * x * x * x))) + else: + raise ValueError(f"Unsupported activation: {ACTIVATION_TYPE}") + + +# === tanh === +@triton.jit +def tanh(x): + return 2 * tl.sigmoid(2 * x) - 1 + diff --git a/benchmarks/TritonNPUBench/sglang/add_tree_reduce_u64_kernel/add_tree_reduce_u64_kernel.py b/benchmarks/TritonNPUBench/sglang/add_tree_reduce_u64_kernel/add_tree_reduce_u64_kernel.py new file mode 100644 index 00000000..20e73ae5 --- /dev/null +++ b/benchmarks/TritonNPUBench/sglang/add_tree_reduce_u64_kernel/add_tree_reduce_u64_kernel.py @@ -0,0 +1,22 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/sglang_operator_cases/newtest_cases/test_add_tree_reduce_u64_kernel.py +# Main kernel: add_tree_reduce_u64_kernel +# PT file: test_add_tree_reduce_u64_kernel_v2.pt + +import triton +import triton.language as tl + + +# === add_tree_reduce_u64_kernel === +@triton.jit +def add_tree_reduce_u64_kernel(in_ptr, out_ptr, n_elems, CHUNK: tl.constexpr): + pid = tl.program_id(axis=0) + start = pid * CHUNK + h = tl.zeros((), dtype=tl.int64) + for i in tl.static_range(0, CHUNK): + idx = start + i + m = idx < n_elems + v = tl.load(in_ptr + idx, mask=m, other=0).to(tl.int64) + h += v + tl.store(out_ptr + pid, h) + diff --git a/benchmarks/TritonNPUBench/sglang/chunk_fwd_kernel_o/chunk_fwd_kernel_o.py b/benchmarks/TritonNPUBench/sglang/chunk_fwd_kernel_o/chunk_fwd_kernel_o.py new file mode 100644 index 00000000..8c65ecc2 --- /dev/null +++ b/benchmarks/TritonNPUBench/sglang/chunk_fwd_kernel_o/chunk_fwd_kernel_o.py @@ -0,0 +1,106 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/sglang_operator_cases/newtest_cases/test_chunk_fwd_kernel_o.py +# Main kernel: chunk_fwd_kernel_o +# PT file: test_chunk_fwd_kernel_o_v2.pt + +import triton +import triton.language as tl + + +# === chunk_fwd_kernel_o === +@triton.jit(do_not_specialize=["T"]) +def chunk_fwd_kernel_o( + q, + k, + v, + h, + g, + o, + cu_seqlens, + chunk_indices, + scale, + T, + H: tl.constexpr, + Hg: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + USE_G: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_v, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_b, i_h = i_bh // H, i_bh % H + + if IS_VARLEN: + i_tg = i_t + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load( + chunk_indices + i_t * 2 + 1 + ).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load( + cu_seqlens + i_n + 1 + ).to(tl.int32) + T = eos - bos + NT = tl.cdiv(T, BT) + else: + NT = tl.cdiv(T, BT) + i_tg = i_b * NT + i_t + bos, eos = i_b * T, i_b * T + T + + # offset calculation + q += (bos * Hg + i_h // (H // Hg)) * K + k += (bos * Hg + i_h // (H // Hg)) * K + v += (bos * H + i_h) * V + o += (bos * H + i_h) * V + h += (i_tg * H + i_h).to(tl.int64) * K * V + + b_o = tl.zeros([BT, BV], dtype=tl.float32) + b_A = tl.zeros([BT, BT], dtype=tl.float32) + + for i_k in range(tl.cdiv(K, BK)): + p_q = tl.make_block_ptr( + q, (T, K), (Hg * K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0) + ) + p_k = tl.make_block_ptr( + k, (K, T), (1, Hg * K), (i_k * BK, i_t * BT), (BK, BT), (0, 1) + ) + p_h = tl.make_block_ptr( + h, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0) + ) + # [BT, BK] + b_q = tl.load(p_q, boundary_check=(0, 1)) + # [BK, BT] + b_k = tl.load(p_k, boundary_check=(0, 1)) + # [BK, BV] + b_h = tl.load(p_h, boundary_check=(0, 1)) + + # [BT, BK] @ [BK, BV] -> [BT, BV] + b_o += tl.dot(b_q, b_h) + # [BT, BK] @ [BK, BT] -> [BT, BT] + b_A += tl.dot(b_q, b_k) + + if USE_G: + g += bos * H + i_h + p_g = tl.make_block_ptr(g, (T,), (H,), (i_t * BT,), (BT,), (0,)) + b_g = tl.load(p_g, boundary_check=(0,)) + b_o = b_o * exp(b_g)[:, None] + b_A = b_A * safe_exp(b_g[:, None] - b_g[None, :]) + + o_i = tl.arange(0, BT) + m_A = o_i[:, None] >= o_i[None, :] + b_A = tl.where(m_A, b_A, 0) + + p_v = tl.make_block_ptr( + v, (T, V), (H * V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0) + ) + p_o = tl.make_block_ptr( + o, (T, V), (H * V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0) + ) + b_v = tl.load(p_v, boundary_check=(0, 1)) + + # to fix mma -> mma layout conversion + # already solved by triton v3.2 or higher + b_o = b_o * scale + tl.dot(b_A.to(b_v.dtype), b_v) * scale + tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1)) + diff --git a/benchmarks/TritonNPUBench/sglang/chunk_gated_delta_rule_fwd_kernel_h_blockdim64/chunk_gated_delta_rule_fwd_kernel_h_blockdim64.py b/benchmarks/TritonNPUBench/sglang/chunk_gated_delta_rule_fwd_kernel_h_blockdim64/chunk_gated_delta_rule_fwd_kernel_h_blockdim64.py new file mode 100644 index 00000000..b30cda94 --- /dev/null +++ b/benchmarks/TritonNPUBench/sglang/chunk_gated_delta_rule_fwd_kernel_h_blockdim64/chunk_gated_delta_rule_fwd_kernel_h_blockdim64.py @@ -0,0 +1,254 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/sglang_operator_cases/newtest_cases/test_chunk_gated_delta_rule_fwd_kernel_h_blockdim64.py +# Main kernel: chunk_gated_delta_rule_fwd_kernel_h_blockdim64 +# PT file: test_chunk_gated_delta_rule_fwd_kernel_h_blockdim64_v2.pt + +import triton +import triton.language as tl + +@triton.jit +def safe_exp(x): + return tl.exp(tl.where(x <= 0, x, float("-inf"))) + +# === chunk_gated_delta_rule_fwd_kernel_h_blockdim64 === +@triton.jit(do_not_specialize=["T"]) +def chunk_gated_delta_rule_fwd_kernel_h_blockdim64( + k, + v, + w, + v_new, + g, + gk, + h, + initial_state, + initial_state_indices, + cu_seqlens, + chunk_offsets, + T, + H: tl.constexpr, + Hg: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BV: tl.constexpr, + USE_G: tl.constexpr, + USE_GK: tl.constexpr, + USE_INITIAL_STATE: tl.constexpr, + INPLACE_UPDATE: tl.constexpr, + SAVE_NEW_VALUE: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_v, i_nh = tl.program_id(0), tl.program_id(1) + i_n, i_h = i_nh // H, i_nh % H + if IS_VARLEN: + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load( + cu_seqlens + i_n + 1 + ).to(tl.int32) + T = eos - bos + NT = tl.cdiv(T, BT) + boh = tl.load(chunk_offsets + i_n).to(tl.int32) + else: + bos, eos = i_n * T, i_n * T + T + NT = tl.cdiv(T, BT) + boh = i_n * NT + + # [BK, BV] + b_h1 = tl.zeros([64, BV], dtype=tl.float32) + if K > 64: + b_h2 = tl.zeros([64, BV], dtype=tl.float32) + if K > 128: + b_h3 = tl.zeros([64, BV], dtype=tl.float32) + if K > 192: + b_h4 = tl.zeros([64, BV], dtype=tl.float32) + + # calculate offset + h += ((boh * H + i_h) * K * V).to(tl.int64) + v += ((bos * H + i_h) * V).to(tl.int64) + k += ((bos * Hg + i_h // (H // Hg)) * K).to(tl.int64) + w += ((bos * H + i_h) * K).to(tl.int64) + if SAVE_NEW_VALUE: + v_new += ((bos * H + i_h) * V).to(tl.int64) + stride_v = H * V + stride_h = H * K * V + stride_k = Hg * K + stride_w = H * K + + index = tl.load(initial_state_indices + i_n).to(tl.int32) + h0 = initial_state + index * stride_h + ht = initial_state + index * stride_h + if USE_INITIAL_STATE: + h0 = h0 + i_h * K * V + if INPLACE_UPDATE: + ht = ht + i_h * K * V + + # load initial state + if USE_INITIAL_STATE: + p_h0_1 = tl.make_block_ptr(h0, (K, V), (V, 1), (0, i_v * BV), (64, BV), (1, 0)) + b_h1 += tl.load(p_h0_1, boundary_check=(0, 1)).to(tl.float32) + if K > 64: + p_h0_2 = tl.make_block_ptr( + h0, (K, V), (V, 1), (64, i_v * BV), (64, BV), (1, 0) + ) + b_h2 += tl.load(p_h0_2, boundary_check=(0, 1)).to(tl.float32) + if K > 128: + p_h0_3 = tl.make_block_ptr( + h0, (K, V), (V, 1), (128, i_v * BV), (64, BV), (1, 0) + ) + b_h3 += tl.load(p_h0_3, boundary_check=(0, 1)).to(tl.float32) + if K > 192: + p_h0_4 = tl.make_block_ptr( + h0, (K, V), (V, 1), (192, i_v * BV), (64, BV), (1, 0) + ) + b_h4 += tl.load(p_h0_4, boundary_check=(0, 1)).to(tl.float32) + + # main recurrence + for i_t in range(NT): + p_h1 = tl.make_block_ptr( + h + i_t * stride_h, (K, V), (V, 1), (0, i_v * BV), (64, BV), (1, 0) + ) + tl.store(p_h1, b_h1.to(p_h1.dtype.element_ty), boundary_check=(0, 1)) + if K > 64: + p_h2 = tl.make_block_ptr( + h + i_t * stride_h, (K, V), (V, 1), (64, i_v * BV), (64, BV), (1, 0) + ) + tl.store(p_h2, b_h2.to(p_h2.dtype.element_ty), boundary_check=(0, 1)) + if K > 128: + p_h3 = tl.make_block_ptr( + h + i_t * stride_h, (K, V), (V, 1), (128, i_v * BV), (64, BV), (1, 0) + ) + tl.store(p_h3, b_h3.to(p_h3.dtype.element_ty), boundary_check=(0, 1)) + if K > 192: + p_h4 = tl.make_block_ptr( + h + i_t * stride_h, (K, V), (V, 1), (192, i_v * BV), (64, BV), (1, 0) + ) + tl.store(p_h4, b_h4.to(p_h4.dtype.element_ty), boundary_check=(0, 1)) + + p_w = tl.make_block_ptr( + w, (T, K), (stride_w, 1), (i_t * BT, 0), (BT, 64), (1, 0) + ) + b_w = tl.load(p_w, boundary_check=(0, 1)) + b_v = tl.dot(b_w, b_h1.to(b_w.dtype)) + if K > 64: + p_w = tl.make_block_ptr( + w, (T, K), (stride_w, 1), (i_t * BT, 64), (BT, 64), (1, 0) + ) + b_w = tl.load(p_w, boundary_check=(0, 1)) + b_v += tl.dot(b_w, b_h2.to(b_w.dtype)) + if K > 128: + p_w = tl.make_block_ptr( + w, (T, K), (stride_w, 1), (i_t * BT, 128), (BT, 64), (1, 0) + ) + b_w = tl.load(p_w, boundary_check=(0, 1)) + b_v += tl.dot(b_w, b_h3.to(b_w.dtype)) + if K > 192: + p_w = tl.make_block_ptr( + w, (T, K), (stride_w, 1), (i_t * BT, 192), (BT, 64), (1, 0) + ) + b_w = tl.load(p_w, boundary_check=(0, 1)) + b_v += tl.dot(b_w, b_h4.to(b_w.dtype)) + p_v = tl.make_block_ptr( + v, (T, V), (stride_v, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0) + ) + b_v = tl.load(p_v, boundary_check=(0, 1)) - b_v + + if SAVE_NEW_VALUE: + p_v = tl.make_block_ptr( + v_new, (T, V), (stride_v, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0) + ) + tl.store(p_v, b_v.to(p_v.dtype.element_ty), boundary_check=(0, 1)) + + last_idx = min((i_t + 1) * BT, T) - 1 + if USE_G: + b_g_last = tl.load(g + bos * H + last_idx * H + i_h) + p_g = tl.make_block_ptr( + g + bos * H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,) + ) + b_g = tl.load(p_g, boundary_check=(0,)) + b_v = b_v * safe_exp(b_g_last - b_g)[:, None] + b_g_last = tl.exp(b_g_last) + b_h1 = b_h1 * b_g_last + if K > 64: + b_h2 = b_h2 * b_g_last + if K > 128: + b_h3 = b_h3 * b_g_last + if K > 192: + b_h4 = b_h4 * b_g_last + + if USE_GK: + o_k1 = tl.arange(0, 64) + b_gk_last1 = tl.load( + gk + (bos + last_idx) * H * K + i_h * K + o_k1, + mask=(o_k1 < K), + other=0.0, + ) + b_h1 *= tl.exp(b_gk_last1)[:, None] + if K > 64: + o_k2 = 64 + o_k1 + b_gk_last2 = tl.load( + gk + (bos + last_idx) * H * K + i_h * K + o_k2, + mask=(o_k2 < K), + other=0.0, + ) + b_h2 *= tl.exp(b_gk_last2)[:, None] + if K > 128: + o_k3 = 128 + o_k1 + b_gk_last3 = tl.load( + gk + (bos + last_idx) * H * K + i_h * K + o_k3, + mask=(o_k3 < K), + other=0.0, + ) + b_h3 *= tl.exp(b_gk_last3)[:, None] + if K > 192: + o_k4 = 192 + o_k1 + b_gk_last4 = tl.load( + gk + (bos + last_idx) * H * K + i_h * K + o_k4, + mask=(o_k4 < K), + other=0.0, + ) + b_h4 *= tl.exp(b_gk_last4)[:, None] + b_v = b_v.to(k.dtype.element_ty) + + p_k = tl.make_block_ptr( + k, (K, T), (1, stride_k), (0, i_t * BT), (64, BT), (0, 1) + ) + b_k = tl.load(p_k, boundary_check=(0, 1)) + b_h1 += tl.dot(b_k, b_v) + if K > 64: + p_k = tl.make_block_ptr( + k, (K, T), (1, stride_k), (64, i_t * BT), (64, BT), (0, 1) + ) + b_k = tl.load(p_k, boundary_check=(0, 1)) + b_h2 += tl.dot(b_k, b_v) + if K > 128: + p_k = tl.make_block_ptr( + k, (K, T), (1, stride_k), (128, i_t * BT), (64, BT), (0, 1) + ) + b_k = tl.load(p_k, boundary_check=(0, 1)) + b_h3 += tl.dot(b_k, b_v) + if K > 192: + p_k = tl.make_block_ptr( + k, (K, T), (1, stride_k), (192, i_t * BT), (64, BT), (0, 1) + ) + b_k = tl.load(p_k, boundary_check=(0, 1)) + b_h4 += tl.dot(b_k, b_v) + + # epilogue + if INPLACE_UPDATE: + p_ht = tl.make_block_ptr(ht, (K, V), (V, 1), (0, i_v * BV), (64, BV), (1, 0)) + tl.store(p_ht, b_h1.to(p_ht.dtype.element_ty), boundary_check=(0, 1)) + if K > 64: + p_ht = tl.make_block_ptr( + ht, (K, V), (V, 1), (64, i_v * BV), (64, BV), (1, 0) + ) + tl.store(p_ht, b_h2.to(p_ht.dtype.element_ty), boundary_check=(0, 1)) + if K > 128: + p_ht = tl.make_block_ptr( + ht, (K, V), (V, 1), (128, i_v * BV), (64, BV), (1, 0) + ) + tl.store(p_ht, b_h3.to(p_ht.dtype.element_ty), boundary_check=(0, 1)) + if K > 192: + p_ht = tl.make_block_ptr( + ht, (K, V), (V, 1), (192, i_v * BV), (64, BV), (1, 0) + ) + tl.store(p_ht, b_h4.to(p_ht.dtype.element_ty), boundary_check=(0, 1)) + diff --git a/benchmarks/TritonNPUBench/sglang/chunk_kda_scaled_dot_kkt_fwd_kernel_intra_sub_inter/chunk_kda_scaled_dot_kkt_fwd_kernel_intra_sub_inter.py b/benchmarks/TritonNPUBench/sglang/chunk_kda_scaled_dot_kkt_fwd_kernel_intra_sub_inter/chunk_kda_scaled_dot_kkt_fwd_kernel_intra_sub_inter.py new file mode 100644 index 00000000..1ed8451f --- /dev/null +++ b/benchmarks/TritonNPUBench/sglang/chunk_kda_scaled_dot_kkt_fwd_kernel_intra_sub_inter/chunk_kda_scaled_dot_kkt_fwd_kernel_intra_sub_inter.py @@ -0,0 +1,110 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/sglang_operator_cases/newtest_cases/test_chunk_kda_scaled_dot_kkt_fwd_kernel_intra_sub_inter.py +# Main kernel: chunk_kda_scaled_dot_kkt_fwd_kernel_intra_sub_inter +# PT file: test_chunk_kda_scaled_dot_kkt_fwd_unit_v2.pt + +import triton +import triton.language as tl + + +# === chunk_kda_scaled_dot_kkt_fwd_kernel_intra_sub_inter === +@triton.jit(do_not_specialize=["T"]) +def chunk_kda_scaled_dot_kkt_fwd_kernel_intra_sub_inter( + q, + k, + g, + beta, + A, + Aqk, + scale, + cu_seqlens, + chunk_indices, + T, + H: tl.constexpr, + K: tl.constexpr, + BT: tl.constexpr, + BC: tl.constexpr, + BK: tl.constexpr, + NC: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_t, i_c, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_b, i_h = i_bh // H, i_bh % H + i_i, i_j = i_c // NC, i_c % NC + if IS_VARLEN: + i_n, i_t = ( + tl.load(chunk_indices + i_t * 2).to(tl.int32), + tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), + ) + bos, eos = ( + tl.load(cu_seqlens + i_n).to(tl.int32), + tl.load(cu_seqlens + i_n + 1).to(tl.int32), + ) + T = eos - bos + else: + bos, eos = i_b * T, i_b * T + T + + if i_t * BT + i_i * BC >= T: + return + if i_i <= i_j: + return + + q += (bos * H + i_h) * K + k += (bos * H + i_h) * K + g += (bos * H + i_h) * K + A += (bos * H + i_h) * BT + Aqk += (bos * H + i_h) * BT + + p_b = tl.make_block_ptr( + beta + bos * H + i_h, (T,), (H,), (i_t * BT + i_i * BC,), (BC,), (0,) + ) + b_b = tl.load(p_b, boundary_check=(0,)) + + b_A = tl.zeros([BC, BC], dtype=tl.float32) + b_Aqk = tl.zeros([BC, BC], dtype=tl.float32) + for i_k in range(tl.cdiv(K, BK)): + p_q = tl.make_block_ptr( + q, (T, K), (H * K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0) + ) + p_k = tl.make_block_ptr( + k, (T, K), (H * K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0) + ) + p_g = tl.make_block_ptr( + g, (T, K), (H * K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0) + ) + b_kt = tl.make_block_ptr( + k, (K, T), (1, H * K), (i_k * BK, i_t * BT + i_j * BC), (BK, BC), (0, 1) + ) + p_gk = tl.make_block_ptr( + g, (K, T), (1, H * K), (i_k * BK, i_t * BT + i_j * BC), (BK, BC), (0, 1) + ) + + o_k = i_k * BK + tl.arange(0, BK) + m_k = o_k < K + # [BK,] + b_gn = tl.load(g + (i_t * BT + i_i * BC) * H * K + o_k, mask=m_k, other=0) + # [BC, BK] + b_g = tl.load(p_g, boundary_check=(0, 1)) + b_k = tl.load(p_k, boundary_check=(0, 1)) * exp(b_g - b_gn[None, :]) + # [BK, BC] + b_gk = tl.load(p_gk, boundary_check=(0, 1)) + b_kt = tl.load(b_kt, boundary_check=(0, 1)) + # [BC, BC] + b_ktg = b_kt * exp(b_gn[:, None] - b_gk) + b_A += tl.dot(b_k, b_ktg) + + b_q = tl.load(p_q, boundary_check=(0, 1)) + b_qg = b_q * exp(b_g - b_gn[None, :]) * scale + b_Aqk += tl.dot(b_qg, b_ktg) + + b_A *= b_b[:, None] + + p_A = tl.make_block_ptr( + A, (T, BT), (H * BT, 1), (i_t * BT + i_i * BC, i_j * BC), (BC, BC), (1, 0) + ) + tl.store(p_A, b_A.to(A.dtype.element_ty), boundary_check=(0, 1)) + p_Aqk = tl.make_block_ptr( + Aqk, (T, BT), (H * BT, 1), (i_t * BT + i_i * BC, i_j * BC), (BC, BC), (1, 0) + ) + tl.store(p_Aqk, b_Aqk.to(Aqk.dtype.element_ty), boundary_check=(0, 1)) + diff --git a/benchmarks/TritonNPUBench/sglang/chunk_kda_scaled_dot_kkt_fwd_kernel_intra_sub_intra/chunk_kda_scaled_dot_kkt_fwd_kernel_intra_sub_intra.py b/benchmarks/TritonNPUBench/sglang/chunk_kda_scaled_dot_kkt_fwd_kernel_intra_sub_intra/chunk_kda_scaled_dot_kkt_fwd_kernel_intra_sub_intra.py new file mode 100644 index 00000000..8c9470a1 --- /dev/null +++ b/benchmarks/TritonNPUBench/sglang/chunk_kda_scaled_dot_kkt_fwd_kernel_intra_sub_intra/chunk_kda_scaled_dot_kkt_fwd_kernel_intra_sub_intra.py @@ -0,0 +1,100 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/sglang_operator_cases/newtest_cases/test_chunk_kda_scaled_dot_kkt_fwd_kernel_intra_sub_intra.py +# Main kernel: chunk_kda_scaled_dot_kkt_fwd_kernel_intra_sub_intra +# PT file: chunk_kda_scaled_dot_kkt_fwd_kernel_intra_sub_intra_v2.pt + +import triton +import triton.language as tl + + +# === chunk_kda_scaled_dot_kkt_fwd_kernel_intra_sub_intra === +@triton.jit(do_not_specialize=["T"]) +def chunk_kda_scaled_dot_kkt_fwd_kernel_intra_sub_intra( + q, + k, + g, + beta, + A, + Aqk, + scale, + cu_seqlens, + chunk_indices, + T, + H: tl.constexpr, + K: tl.constexpr, + BT: tl.constexpr, + BC: tl.constexpr, + BK: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_t, i_i, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_b, i_h = i_bh // H, i_bh % H + if IS_VARLEN: + i_n, i_t = ( + tl.load(chunk_indices + i_t * 2).to(tl.int32), + tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), + ) + bos, eos = ( + tl.load(cu_seqlens + i_n).to(tl.int32), + tl.load(cu_seqlens + i_n + 1).to(tl.int32), + ) + T = eos - bos + else: + bos, eos = i_b * T, i_b * T + T + + if i_t * BT + i_i * BC >= T: + return + + o_i = tl.arange(0, BC) + o_k = tl.arange(0, BK) + m_k = o_k < K + m_A = (i_t * BT + i_i * BC + o_i) < T + o_A = (bos + i_t * BT + i_i * BC + o_i) * H * BT + i_h * BT + i_i * BC + + p_q = tl.make_block_ptr( + q + (bos * H + i_h) * K, + (T, K), + (H * K, 1), + (i_t * BT + i_i * BC, 0), + (BC, BK), + (1, 0), + ) + p_k = tl.make_block_ptr( + k + (bos * H + i_h) * K, + (T, K), + (H * K, 1), + (i_t * BT + i_i * BC, 0), + (BC, BK), + (1, 0), + ) + p_g = tl.make_block_ptr( + g + (bos * H + i_h) * K, + (T, K), + (H * K, 1), + (i_t * BT + i_i * BC, 0), + (BC, BK), + (1, 0), + ) + b_q = tl.load(p_q, boundary_check=(0, 1)) + b_k = tl.load(p_k, boundary_check=(0, 1)) + b_g = tl.load(p_g, boundary_check=(0, 1)) + + p_b = beta + (bos + i_t * BT + i_i * BC + o_i) * H + i_h + b_k = b_k * tl.load(p_b, mask=m_A, other=0)[:, None] + + p_kt = k + (bos + i_t * BT + i_i * BC) * H * K + i_h * K + o_k + p_gk = g + (bos + i_t * BT + i_i * BC) * H * K + i_h * K + o_k + + for j in range(0, min(BC, T - i_t * BT - i_i * BC)): + b_kt = tl.load(p_kt, mask=m_k, other=0).to(tl.float32) + b_gk = tl.load(p_gk, mask=m_k, other=0).to(tl.float32) + b_ktg = b_kt[None, :] * tl.exp(b_g - b_gk[None, :]) + b_A = tl.sum(b_k * b_ktg, 1) + b_A = tl.where(o_i > j, b_A, 0.0) + b_Aqk = tl.sum(b_q * b_ktg, 1) + b_Aqk = tl.where(o_i >= j, b_Aqk * scale, 0.0) + tl.store(A + o_A + j, b_A, mask=m_A) + tl.store(Aqk + o_A + j, b_Aqk, mask=m_A) + p_kt += H * K + p_gk += H * K + diff --git a/benchmarks/TritonNPUBench/sglang/chunk_local_cumsum_scalar_kernel/chunk_local_cumsum_scalar_kernel.py b/benchmarks/TritonNPUBench/sglang/chunk_local_cumsum_scalar_kernel/chunk_local_cumsum_scalar_kernel.py new file mode 100644 index 00000000..251747ec --- /dev/null +++ b/benchmarks/TritonNPUBench/sglang/chunk_local_cumsum_scalar_kernel/chunk_local_cumsum_scalar_kernel.py @@ -0,0 +1,59 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/sglang_operator_cases/newtest_cases/test_chunk_local_cumsum_scalar_kernel.py +# Main kernel: chunk_local_cumsum_scalar_kernel +# PT file: test_chunk_local_cumsum_scalar_kernel_v2.pt + +import triton +import triton.language as tl + + +# === chunk_local_cumsum_scalar_kernel === +@triton.jit(do_not_specialize=["T"]) +def chunk_local_cumsum_scalar_kernel( + s, + o, + scale, + cu_seqlens, + chunk_indices, + T, + B: tl.constexpr, + H: tl.constexpr, + BT: tl.constexpr, + REVERSE: tl.constexpr, + HAS_SCALE: tl.constexpr, + IS_VARLEN: tl.constexpr, + HEAD_FIRST: tl.constexpr, +): + i_t, i_bh = tl.program_id(0), tl.program_id(1) + i_b, i_h = i_bh // H, i_bh % H + if IS_VARLEN: + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load( + chunk_indices + i_t * 2 + 1 + ).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load( + cu_seqlens + i_n + 1 + ).to(tl.int32) + T = eos - bos + else: + bos, eos = i_b * T, i_b * T + T + + if HEAD_FIRST: + p_s = tl.make_block_ptr( + s + bos * H + i_h * T, (T,), (1,), (i_t * BT,), (BT,), (0,) + ) + p_o = tl.make_block_ptr( + o + bos * H + i_h * T, (T,), (1,), (i_t * BT,), (BT,), (0,) + ) + else: + p_s = tl.make_block_ptr(s + bos * H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) + p_o = tl.make_block_ptr(o + bos * H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) + # [BT] + b_s = tl.load(p_s, boundary_check=(0,)).to(tl.float32) + b_o = tl.cumsum(b_s, axis=0) + if REVERSE: + b_z = tl.sum(b_s, axis=0) + b_o = -b_o + b_z[None] + b_s + if HAS_SCALE: + b_o *= scale + tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0,)) + diff --git a/benchmarks/TritonNPUBench/sglang/chunk_local_cumsum_vector_kernel/chunk_local_cumsum_vector_kernel.py b/benchmarks/TritonNPUBench/sglang/chunk_local_cumsum_vector_kernel/chunk_local_cumsum_vector_kernel.py new file mode 100644 index 00000000..85b79ec3 --- /dev/null +++ b/benchmarks/TritonNPUBench/sglang/chunk_local_cumsum_vector_kernel/chunk_local_cumsum_vector_kernel.py @@ -0,0 +1,88 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/sglang_operator_cases/newtest_cases/test_chunk_local_cumsum_vector_kernel.py +# Main kernel: chunk_local_cumsum_vector_kernel +# PT file: test_chunk_local_cumsum_vector_kernel_v2.pt + +import triton +import triton.language as tl + + +# === chunk_local_cumsum_vector_kernel === +@triton.jit(do_not_specialize=["T"]) +def chunk_local_cumsum_vector_kernel( + s, + o, + scale, + cu_seqlens, + chunk_indices, + T, + B: tl.constexpr, + H: tl.constexpr, + S: tl.constexpr, + BT: tl.constexpr, + BS: tl.constexpr, + REVERSE: tl.constexpr, + HAS_SCALE: tl.constexpr, + IS_VARLEN: tl.constexpr, + HEAD_FIRST: tl.constexpr, +): + i_s, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_b, i_h = i_bh // H, i_bh % H + if IS_VARLEN: + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load( + chunk_indices + i_t * 2 + 1 + ).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load( + cu_seqlens + i_n + 1 + ).to(tl.int32) + T = eos - bos + else: + bos, eos = i_b * T, i_b * T + T + + o_i = tl.arange(0, BT) + if REVERSE: + m_s = tl.where(o_i[:, None] <= o_i[None, :], 1.0, 0.0) + else: + m_s = tl.where(o_i[:, None] >= o_i[None, :], 1.0, 0.0) + + if HEAD_FIRST: + p_s = tl.make_block_ptr( + s + (bos * H + i_h * T) * S, + (T, S), + (S, 1), + (i_t * BT, i_s * BS), + (BT, BS), + (1, 0), + ) + p_o = tl.make_block_ptr( + o + (bos * H + i_h * T) * S, + (T, S), + (S, 1), + (i_t * BT, i_s * BS), + (BT, BS), + (1, 0), + ) + else: + p_s = tl.make_block_ptr( + s + (bos * H + i_h) * S, + (T, S), + (H * S, 1), + (i_t * BT, i_s * BS), + (BT, BS), + (1, 0), + ) + p_o = tl.make_block_ptr( + o + (bos * H + i_h) * S, + (T, S), + (H * S, 1), + (i_t * BT, i_s * BS), + (BT, BS), + (1, 0), + ) + # [BT, BS] + b_s = tl.load(p_s, boundary_check=(0, 1)).to(tl.float32) + b_o = tl.dot(m_s, b_s, allow_tf32=False) + if HAS_SCALE: + b_o *= scale + tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1)) + diff --git a/benchmarks/TritonNPUBench/sglang/chunk_scaled_dot_kkt_fwd_kernel/chunk_scaled_dot_kkt_fwd_kernel.py b/benchmarks/TritonNPUBench/sglang/chunk_scaled_dot_kkt_fwd_kernel/chunk_scaled_dot_kkt_fwd_kernel.py new file mode 100644 index 00000000..71f9e05c --- /dev/null +++ b/benchmarks/TritonNPUBench/sglang/chunk_scaled_dot_kkt_fwd_kernel/chunk_scaled_dot_kkt_fwd_kernel.py @@ -0,0 +1,80 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/sglang_operator_cases/newtest_cases/test_chunk_scaled_dot_kkt_fwd_kernel.py +# Main kernel: chunk_scaled_dot_kkt_fwd_kernel +# PT file: test_chunk_scaled_dot_kkt_fwd_kernel__unit_v2.pt + +import triton +import triton.language as tl + + +# === chunk_scaled_dot_kkt_fwd_kernel === +@triton.jit(do_not_specialize=["T"]) +def chunk_scaled_dot_kkt_fwd_kernel( + k, + beta, + g_cumsum, + A, + cu_seqlens, + chunk_indices, + T, + H: tl.constexpr, + Hg: tl.constexpr, + K: tl.constexpr, + BT: tl.constexpr, + BK: tl.constexpr, + IS_VARLEN: tl.constexpr, + USE_G: tl.constexpr, +): + i_t, i_bh = tl.program_id(0), tl.program_id(1) + i_b, i_h = i_bh // H, i_bh % H + if IS_VARLEN: + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load( + chunk_indices + i_t * 2 + 1 + ).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load( + cu_seqlens + i_n + 1 + ).to(tl.int32) + T = eos - bos + else: + bos, eos = i_b * T, i_b * T + T + o_t = tl.arange(0, BT) + + p_beta = tl.make_block_ptr( + beta + bos * H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,) + ) + b_beta = tl.load(p_beta, boundary_check=(0,)) + + b_A = tl.zeros([BT, BT], dtype=tl.float32) + for i_k in range(tl.cdiv(K, BK)): + p_k = tl.make_block_ptr( + k + (bos * Hg + i_h // (H // Hg)) * K, + (T, K), + (Hg * K, 1), + (i_t * BT, i_k * BK), + (BT, BK), + (1, 0), + ) + b_k = tl.load(p_k, boundary_check=(0, 1)) + b_A += tl.dot(b_k, tl.trans(b_k)) + + if USE_G: + p_g = tl.make_block_ptr( + g_cumsum + bos * H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,) + ) + b_g = tl.load(p_g, boundary_check=(0,)) + b_g_diff = b_g[:, None] - b_g[None, :] + b_A = b_A * safe_exp(b_g_diff) + + b_A *= b_beta[:, None] + b_A = tl.where(o_t[:, None] > o_t[None, :], b_A, 0) + p_A = tl.make_block_ptr( + A + (bos * H + i_h) * BT, (T, BT), (BT * H, 1), (i_t * BT, 0), (BT, BT), (1, 0) + ) + tl.store(p_A, b_A.to(p_A.dtype.element_ty), boundary_check=(0, 1)) + + +# === safe_exp === +@triton.jit +def safe_exp(x): + return exp(tl.where(x <= 0, x, float("-inf"))) + diff --git a/benchmarks/TritonNPUBench/sglang/compute_identity_kernel/compute_identity_kernel.py b/benchmarks/TritonNPUBench/sglang/compute_identity_kernel/compute_identity_kernel.py new file mode 100644 index 00000000..ef83ef1b --- /dev/null +++ b/benchmarks/TritonNPUBench/sglang/compute_identity_kernel/compute_identity_kernel.py @@ -0,0 +1,48 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/sglang_operator_cases/newtest_cases/test_compute_identity_kernel.py +# Main kernel: compute_identity_kernel +# PT file: test_compute_identity_kernel_v2.pt + +import triton +import triton.language as tl + + +# === compute_identity_kernel === +@triton.jit +def compute_identity_kernel( + top_k, + hidden_states_ptr, + expert_scales_ptr, + num_tokens, + output_ptr, + hidden_dim, + scales_stride, + BLOCK_SIZE: tl.constexpr, +): + pid = tl.program_id(0) + + batch_id = pid // (hidden_dim // BLOCK_SIZE) + dim_offset = pid % (hidden_dim // BLOCK_SIZE) * BLOCK_SIZE + + if batch_id >= num_tokens or dim_offset >= hidden_dim: + return + + h = tl.load( + hidden_states_ptr + + batch_id * hidden_dim + + dim_offset + + tl.arange(0, BLOCK_SIZE), + mask=(dim_offset + tl.arange(0, BLOCK_SIZE)) < hidden_dim, + ) + + result = tl.zeros([BLOCK_SIZE], dtype=tl.float32) + for i in range(top_k): + scale = tl.load(expert_scales_ptr + batch_id * scales_stride + i) + result += h * scale + + tl.store( + output_ptr + batch_id * hidden_dim + dim_offset + tl.arange(0, BLOCK_SIZE), + result, + mask=(dim_offset + tl.arange(0, BLOCK_SIZE)) < hidden_dim, + ) + diff --git a/benchmarks/TritonNPUBench/sglang/compute_masked_m_triton_kernel/compute_masked_m_triton_kernel.py b/benchmarks/TritonNPUBench/sglang/compute_masked_m_triton_kernel/compute_masked_m_triton_kernel.py new file mode 100644 index 00000000..5bff411e --- /dev/null +++ b/benchmarks/TritonNPUBench/sglang/compute_masked_m_triton_kernel/compute_masked_m_triton_kernel.py @@ -0,0 +1,17 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/sglang_operator_cases/newtest_cases/test_compute_masked_m_triton_kernel.py +# Main kernel: compute_masked_m_triton_kernel +# PT file: test_compute_masked_m_triton_kernel_v2.pt + +import triton +import triton.language as tl + + +# === compute_masked_m_triton_kernel === +@triton.jit +def compute_masked_m_triton_kernel(seg_indptr, masked_m): + expert_id = tl.program_id(0) + start = tl.load(seg_indptr + expert_id) + end = tl.load(seg_indptr + expert_id + 1) + tl.store(masked_m + expert_id, (end - start)) + diff --git a/benchmarks/TritonNPUBench/sglang/compute_problem_sizes_w4a8_kernel/compute_problem_sizes_w4a8_kernel.py b/benchmarks/TritonNPUBench/sglang/compute_problem_sizes_w4a8_kernel/compute_problem_sizes_w4a8_kernel.py new file mode 100644 index 00000000..03c988ff --- /dev/null +++ b/benchmarks/TritonNPUBench/sglang/compute_problem_sizes_w4a8_kernel/compute_problem_sizes_w4a8_kernel.py @@ -0,0 +1,47 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/sglang_operator_cases/newtest_cases/test_compute_problem_sizes_w4a8_kernel.py +# Main kernel: compute_problem_sizes_w4a8_kernel +# PT file: test_compute_problem_sizes_w4a8_kernel_v2.pt + +import triton +import triton.language as tl + + +# === compute_problem_sizes_w4a8_kernel === +@triton.jit +def compute_problem_sizes_w4a8_kernel( + masked_m_ptr, + problem_sizes1_ptr, + problem_sizes2_ptr, + n, + k, + num_experts, + BLOCK_SIZE: tl.constexpr, +): + pid = tl.program_id(axis=0) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + mask = pid < num_experts + final_occurrences = tl.load(masked_m_ptr + pid, mask=mask, other=0) + + ps1_idx_0 = pid * 3 + ps1_idx_1 = ps1_idx_0 + 1 + ps1_idx_2 = ps1_idx_0 + 2 + + ps2_idx_0 = pid * 3 + ps2_idx_1 = ps2_idx_0 + 1 + ps2_idx_2 = ps2_idx_0 + 2 + + ps1_mask_0 = ps1_idx_0 < num_experts * 3 + ps1_mask_1 = ps1_idx_1 < num_experts * 3 + ps1_mask_2 = ps1_idx_2 < num_experts * 3 + ps2_mask_0 = ps2_idx_0 < num_experts * 3 + ps2_mask_1 = ps2_idx_1 < num_experts * 3 + ps2_mask_2 = ps2_idx_2 < num_experts * 3 + + tl.store(problem_sizes1_ptr + ps1_idx_0, 2 * n, mask=ps1_mask_0) + tl.store(problem_sizes1_ptr + ps1_idx_1, final_occurrences, mask=ps1_mask_1) + tl.store(problem_sizes1_ptr + ps1_idx_2, k, mask=ps1_mask_2) + + tl.store(problem_sizes2_ptr + ps2_idx_0, k, mask=ps2_mask_0) + tl.store(problem_sizes2_ptr + ps2_idx_1, final_occurrences, mask=ps2_mask_1) + tl.store(problem_sizes2_ptr + ps2_idx_2, n, mask=ps2_mask_2) + diff --git a/benchmarks/TritonNPUBench/sglang/compute_seg_indptr_triton_kernel/compute_seg_indptr_triton_kernel.py b/benchmarks/TritonNPUBench/sglang/compute_seg_indptr_triton_kernel/compute_seg_indptr_triton_kernel.py new file mode 100644 index 00000000..4a04e603 --- /dev/null +++ b/benchmarks/TritonNPUBench/sglang/compute_seg_indptr_triton_kernel/compute_seg_indptr_triton_kernel.py @@ -0,0 +1,26 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/sglang_operator_cases/newtest_cases/test_compute_seg_indptr_triton_kernel.py +# Main kernel: compute_seg_indptr_triton_kernel +# PT file: test_compute_seg_indptr_triton_kernel_v2.pt + +import triton +import triton.language as tl + + +# === compute_seg_indptr_triton_kernel === +@triton.jit +def compute_seg_indptr_triton_kernel(reorder_topk_ids, seg_indptr, num_toks): + expert_id_minus_1 = tl.program_id(0) - 1 + low = 0 + high = num_toks - 1 + target_location = -1 + while low <= high: + mid = (low + high) // 2 + + if tl.load(reorder_topk_ids + mid) > expert_id_minus_1: + high = mid - 1 + else: + low = mid + 1 + target_location = mid + tl.store(seg_indptr + expert_id_minus_1 + 1, target_location + 1) + diff --git a/benchmarks/TritonNPUBench/sglang/compute_src2dst_triton_kernel/compute_src2dst_triton_kernel.py b/benchmarks/TritonNPUBench/sglang/compute_src2dst_triton_kernel/compute_src2dst_triton_kernel.py new file mode 100644 index 00000000..576395bf --- /dev/null +++ b/benchmarks/TritonNPUBench/sglang/compute_src2dst_triton_kernel/compute_src2dst_triton_kernel.py @@ -0,0 +1,20 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/sglang_operator_cases/newtest_cases/test_compute_src2dst_triton_kernel.py +# Main kernel: compute_src2dst_triton_kernel +# PT file: test_compute_src2dst_triton_kernel_v2.pt + +import triton +import triton.language as tl + + +# === compute_src2dst_triton_kernel === +@triton.jit +def compute_src2dst_triton_kernel( + reorder_ids, src2dst, num_toks, BLOCK_SIZE: tl.constexpr +): + pid = tl.program_id(axis=0) + dst_id = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + mask = dst_id < num_toks + src_id = tl.load(reorder_ids + dst_id, mask=mask) + tl.store(src2dst + src_id, dst_id, mask=mask) + diff --git a/benchmarks/TritonNPUBench/sglang/concat_and_cast_mha_k_kernel/concat_and_cast_mha_k_kernel.py b/benchmarks/TritonNPUBench/sglang/concat_and_cast_mha_k_kernel/concat_and_cast_mha_k_kernel.py new file mode 100644 index 00000000..835044e6 --- /dev/null +++ b/benchmarks/TritonNPUBench/sglang/concat_and_cast_mha_k_kernel/concat_and_cast_mha_k_kernel.py @@ -0,0 +1,48 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/sglang_operator_cases/newtest_cases/test_concat_and_cast_mha_k_kernel.py +# Main kernel: concat_and_cast_mha_k_kernel +# PT file: test_concat_and_cast_mha_k_kernel_v2.pt + +import triton +import triton.language as tl + + +# === concat_and_cast_mha_k_kernel === +@triton.jit +def concat_and_cast_mha_k_kernel( + k_ptr, + k_nope_ptr, + k_rope_ptr, + head_cnt: tl.constexpr, + k_stride0: tl.constexpr, + k_stride1: tl.constexpr, + nope_stride0: tl.constexpr, + nope_stride1: tl.constexpr, + rope_stride0: tl.constexpr, + nope_dim: tl.constexpr, + rope_dim: tl.constexpr, +): + pid_loc = tl.program_id(0) + head_range = tl.arange(0, head_cnt) + + k_head_ptr = k_ptr + pid_loc * k_stride0 + head_range[:, None] * k_stride1 + + nope_offs = tl.arange(0, nope_dim) + + src_nope_ptr = ( + k_nope_ptr + + pid_loc * nope_stride0 + + head_range[:, None] * nope_stride1 + + nope_offs[None, :] + ) + dst_nope_ptr = k_head_ptr + nope_offs[None, :] + + src_nope = tl.load(src_nope_ptr) + tl.store(dst_nope_ptr, src_nope) + + rope_offs = tl.arange(0, rope_dim) + src_rope_ptr = k_rope_ptr + pid_loc * rope_stride0 + rope_offs[None, :] + dst_rope_ptr = k_head_ptr + nope_dim + rope_offs[None, :] + src_rope = tl.load(src_rope_ptr) + tl.store(dst_rope_ptr, src_rope) + diff --git a/benchmarks/TritonNPUBench/sglang/create_flashinfer_kv_indices_triton/create_flashinfer_kv_indices_triton.py b/benchmarks/TritonNPUBench/sglang/create_flashinfer_kv_indices_triton/create_flashinfer_kv_indices_triton.py new file mode 100644 index 00000000..b292c42a --- /dev/null +++ b/benchmarks/TritonNPUBench/sglang/create_flashinfer_kv_indices_triton/create_flashinfer_kv_indices_triton.py @@ -0,0 +1,48 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/sglang_operator_cases/newtest_cases/test_create_flashinfer_kv_indices_triton.py +# Main kernel: create_flashinfer_kv_indices_triton +# PT file: test_create_flashinfer_kv_indices_triton_v2.pt + +import triton +import triton.language as tl + + +# === create_flashinfer_kv_indices_triton === +@triton.jit +def create_flashinfer_kv_indices_triton( + req_to_token_ptr, # [max_batch, max_context_len] + req_pool_indices_ptr, + page_kernel_lens_ptr, + kv_indptr, + kv_start_idx, + kv_indices_ptr, + req_to_token_ptr_stride: tl.constexpr, +): + BLOCK_SIZE: tl.constexpr = 512 + pid = tl.program_id(axis=0) + + # find the req pool idx, this is for batch to token + req_pool_index = tl.load(req_pool_indices_ptr + pid) + kv_indices_offset = tl.load(kv_indptr + pid) + + kv_start = 0 + kv_end = 0 + if kv_start_idx: + kv_start = tl.load(kv_start_idx + pid).to(tl.int32) + kv_end = kv_start + kv_end += tl.load(page_kernel_lens_ptr + pid).to(tl.int32) + + num_loop = tl.cdiv(kv_end - kv_start, BLOCK_SIZE) + for i in range(num_loop): + # index into req_to_token_ptr needs to be int64 + offset = tl.arange(0, BLOCK_SIZE).to(tl.int64) + i * BLOCK_SIZE + mask = offset < kv_end - kv_start + data = tl.load( + req_to_token_ptr + + req_pool_index * req_to_token_ptr_stride + + kv_start + + offset, + mask=mask, + ) + tl.store(kv_indices_ptr + kv_indices_offset + offset, data, mask=mask) + diff --git a/benchmarks/TritonNPUBench/sglang/create_flashmla_kv_indices_triton/create_flashmla_kv_indices_triton.py b/benchmarks/TritonNPUBench/sglang/create_flashmla_kv_indices_triton/create_flashmla_kv_indices_triton.py new file mode 100644 index 00000000..57b79769 --- /dev/null +++ b/benchmarks/TritonNPUBench/sglang/create_flashmla_kv_indices_triton/create_flashmla_kv_indices_triton.py @@ -0,0 +1,65 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/sglang_operator_cases/newtest_cases/test_create_flashmla_kv_indices_triton.py +# Main kernel: create_flashmla_kv_indices_triton +# PT file: test_create_flashmla_kv_indices_triton_v2.pt + +import triton +import triton.language as tl + +_FLASHMLA_CREATE_KV_BLOCK_SIZE = 4096 +FLASHMLA_CREATE_KV_BLOCK_SIZE_TRITON = tl.constexpr(_FLASHMLA_CREATE_KV_BLOCK_SIZE) + +# === create_flashmla_kv_indices_triton === +@triton.jit +def create_flashmla_kv_indices_triton( + req_to_token_ptr, # [max_batch, max_context_len] + req_pool_indices_ptr, + page_kernel_lens_ptr, + kv_start_idx, + kv_indices_ptr, + req_to_token_ptr_stride: tl.constexpr, + kv_indices_ptr_stride: tl.constexpr, + PAGED_SIZE: tl.constexpr = 64, +): + NUM_PAGE_PER_BLOCK: tl.constexpr = ( + FLASHMLA_CREATE_KV_BLOCK_SIZE_TRITON // PAGED_SIZE + ) + pid = tl.program_id(axis=0) + + # find the req pool idx, this is for batch to token + req_pool_index = tl.load(req_pool_indices_ptr + pid) + + kv_start = 0 + kv_end = 0 + if kv_start_idx: + kv_start = tl.load(kv_start_idx + pid).to(tl.int32) + kv_end = kv_start + + kv_end += tl.load(page_kernel_lens_ptr + pid).to(tl.int32) + + num_paged = tl.cdiv(kv_end - kv_start, PAGED_SIZE) + num_pages_loop = tl.cdiv(kv_end - kv_start, FLASHMLA_CREATE_KV_BLOCK_SIZE_TRITON) + + for i in range(num_pages_loop): + # index into req_to_token_ptr needs to be int64 + paged_offset = ( + tl.arange(0, NUM_PAGE_PER_BLOCK).to(tl.int64) + i * NUM_PAGE_PER_BLOCK + ) * PAGED_SIZE + paged_offset_out = tl.arange(0, NUM_PAGE_PER_BLOCK) + i * NUM_PAGE_PER_BLOCK + + mask = paged_offset < num_paged * PAGED_SIZE + mask_out = paged_offset_out < num_paged + + data = tl.load( + req_to_token_ptr + + req_pool_index * req_to_token_ptr_stride + + kv_start + + paged_offset, + mask=mask, + ) + tl.store( + kv_indices_ptr + pid * kv_indices_ptr_stride + paged_offset_out, + data // PAGED_SIZE, + mask=mask_out, + ) + diff --git a/benchmarks/TritonNPUBench/sglang/deepep_compute_src2dst_triton_kernel/deepep_compute_src2dst_triton_kernel.py b/benchmarks/TritonNPUBench/sglang/deepep_compute_src2dst_triton_kernel/deepep_compute_src2dst_triton_kernel.py new file mode 100644 index 00000000..9652f543 --- /dev/null +++ b/benchmarks/TritonNPUBench/sglang/deepep_compute_src2dst_triton_kernel/deepep_compute_src2dst_triton_kernel.py @@ -0,0 +1,21 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/sglang_operator_cases/newtest_cases/test_deepep_compute_src2dst_triton_kernel.py +# Main kernel: deepep_compute_src2dst_triton_kernel +# PT file: test_deepep_compute_src2dst_triton_kernel_v2.pt + +import triton +import triton.language as tl + + +# === deepep_compute_src2dst_triton_kernel === +@triton.jit +def deepep_compute_src2dst_triton_kernel( + reorder_ids, src2dst, num_toks, num_minus_one, BLOCK_SIZE: tl.constexpr +): + pid = tl.program_id(axis=0) + dst_id = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + mask = dst_id < num_toks + src_id = tl.load(reorder_ids + dst_id, mask=mask) + num_invalid = tl.load(num_minus_one) + tl.store(src2dst + src_id, dst_id - num_invalid, mask=mask) + diff --git a/benchmarks/TritonNPUBench/sglang/deepep_permute_triton_kernel/deepep_permute_triton_kernel.py b/benchmarks/TritonNPUBench/sglang/deepep_permute_triton_kernel/deepep_permute_triton_kernel.py new file mode 100644 index 00000000..1cda25f6 --- /dev/null +++ b/benchmarks/TritonNPUBench/sglang/deepep_permute_triton_kernel/deepep_permute_triton_kernel.py @@ -0,0 +1,40 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/sglang_operator_cases/newtest_cases/test_deepep_permute_triton_kernel.py +# Main kernel: deepep_permute_triton_kernel +# PT file: test_deepep_permute_triton_kernel_v2.pt + +import triton +import triton.language as tl + + +# === deepep_permute_triton_kernel === +@triton.jit +def deepep_permute_triton_kernel( + input_ptr, + gateup_input_ptr, + src2dst_ptr, + topk_ids_ptr, + a1_scales_ptr, + topk, + hidden_size, + BLOCK_SIZE: tl.constexpr, +): + OutDtype = gateup_input_ptr.dtype.element_ty + + src_idx = tl.program_id(0) + src2dst_ptr = src2dst_ptr + src_idx * topk + topk_ids_ptr = topk_ids_ptr + src_idx * topk + + src_ptr = input_ptr + src_idx * hidden_size + + for start_offset in tl.range(0, hidden_size, BLOCK_SIZE): + offset = start_offset + tl.arange(0, BLOCK_SIZE) + mask = offset < hidden_size + in_data = tl.load(src_ptr + offset, mask=mask).to(OutDtype) + + for idx in range(topk): + dst_idx = tl.load(src2dst_ptr + idx) + if dst_idx >= 0: + dst_ptr = gateup_input_ptr + dst_idx * hidden_size + tl.store(dst_ptr + offset, in_data, mask=mask) + diff --git a/benchmarks/TritonNPUBench/sglang/deepep_post_reorder_triton_kernel/deepep_post_reorder_triton_kernel.py b/benchmarks/TritonNPUBench/sglang/deepep_post_reorder_triton_kernel/deepep_post_reorder_triton_kernel.py new file mode 100644 index 00000000..6ce57639 --- /dev/null +++ b/benchmarks/TritonNPUBench/sglang/deepep_post_reorder_triton_kernel/deepep_post_reorder_triton_kernel.py @@ -0,0 +1,42 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/sglang_operator_cases/newtest_cases/test_deepep_post_reorder_triton_kernel.py +# Main kernel: deepep_post_reorder_triton_kernel +# PT file: test_deepep_post_reorder_triton_kernel_v2.pt + +import triton +import triton.language as tl + + +# === deepep_post_reorder_triton_kernel === +@triton.jit +def deepep_post_reorder_triton_kernel( + down_output_ptr, + output_ptr, + src2dst_ptr, + topk_ids_ptr, + topk_weights_ptr, + topk, + hidden_size, + BLOCK_SIZE: tl.constexpr, +): + InDtype = down_output_ptr.dtype.element_ty + + src_idx = tl.program_id(0) + src2dst_ptr = src2dst_ptr + src_idx * topk + topk_ids_ptr = topk_ids_ptr + src_idx * topk + topk_weights_ptr = topk_weights_ptr + src_idx * topk + + store_ptr = output_ptr + src_idx * hidden_size + for start_offset in tl.range(0, hidden_size, BLOCK_SIZE): + offset = start_offset + tl.arange(0, BLOCK_SIZE) + mask = offset < hidden_size + sum_vec = tl.zeros([BLOCK_SIZE], dtype=InDtype) + for idx in range(topk): + dst_idx = tl.load(src2dst_ptr + idx) + if dst_idx >= 0: + weigh_scale = tl.load(topk_weights_ptr + idx).to(InDtype) + load_ptr = down_output_ptr + dst_idx * hidden_size + in_data = tl.load(load_ptr + offset, mask=mask) + sum_vec += in_data * weigh_scale + tl.store(store_ptr + offset, sum_vec, mask=mask) + diff --git a/benchmarks/TritonNPUBench/sglang/deepgemm_compute_src2dst_triton_kernel/deepgemm_compute_src2dst_triton_kernel.py b/benchmarks/TritonNPUBench/sglang/deepgemm_compute_src2dst_triton_kernel/deepgemm_compute_src2dst_triton_kernel.py new file mode 100644 index 00000000..3dd281be --- /dev/null +++ b/benchmarks/TritonNPUBench/sglang/deepgemm_compute_src2dst_triton_kernel/deepgemm_compute_src2dst_triton_kernel.py @@ -0,0 +1,30 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/sglang_operator_cases/newtest_cases/test_deepgemm_compute_src2dst_triton_kernel.py +# Main kernel: deepgemm_compute_src2dst_triton_kernel +# PT file: test_deepgemm_compute_src2dst_triton_kernel_v2.pt + +import triton +import triton.language as tl + + +# === deepgemm_compute_src2dst_triton_kernel === +@triton.jit +def deepgemm_compute_src2dst_triton_kernel( + topk_ids, + reorder_ids, + seg_indptr, + src2dst, + m_max, + num_toks, + BLOCK_SIZE: tl.constexpr, +): + pid = tl.program_id(axis=0) + dst_id = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + mask = dst_id < num_toks + src_id = tl.load(reorder_ids + dst_id, mask=mask) + expert_id = tl.load(topk_ids + src_id, mask=(src_id < num_toks)) + expert_dst_start = tl.load(seg_indptr + expert_id, mask=(expert_id >= 0)) + expert_dst_offset = dst_id - expert_dst_start + dst_id = expert_id * m_max + expert_dst_offset + tl.store(src2dst + src_id, dst_id, mask=mask) + diff --git a/benchmarks/TritonNPUBench/sglang/experts_combine_kernel/experts_combine_kernel.py b/benchmarks/TritonNPUBench/sglang/experts_combine_kernel/experts_combine_kernel.py new file mode 100644 index 00000000..1d890afb --- /dev/null +++ b/benchmarks/TritonNPUBench/sglang/experts_combine_kernel/experts_combine_kernel.py @@ -0,0 +1,40 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/sglang_operator_cases/newtest_cases/test_experts_combine_kernel.py +# Main kernel: experts_combine_kernel +# PT file: test_experts_combine_kernel_v2.pt + +import triton +import triton.language as tl + + +# === experts_combine_kernel === +@triton.jit +def experts_combine_kernel( + out_hidden_states, + moe_hidden_states, + mlp_hidden_states, + combine_k: tl.constexpr, + hidden_dim: tl.constexpr, + BLOCK_SIZE: tl.constexpr, +): + pid = tl.program_id(0) + start_index_mlp = pid * hidden_dim + start_index_rmoe = pid * hidden_dim * combine_k + offsets = tl.arange(0, BLOCK_SIZE) + mask = offsets < hidden_dim + combine_k_offsets = tl.arange(0, combine_k) + + moe_x = tl.load( + moe_hidden_states + + start_index_rmoe + + combine_k_offsets[:, None] * hidden_dim + + offsets[None, :], + mask=mask[None, :], + other=0.0, + ) + moe_x = tl.sum(moe_x, axis=0) + mlp_x = tl.load(mlp_hidden_states + start_index_mlp + offsets, mask=mask, other=0.0) + combined_x = (moe_x + mlp_x) / 1.4142135623730951 + + tl.store(out_hidden_states + start_index_mlp + offsets, combined_x, mask=mask) + diff --git a/benchmarks/TritonNPUBench/sglang/fill_gateup_input_triton_kernel/fill_gateup_input_triton_kernel.py b/benchmarks/TritonNPUBench/sglang/fill_gateup_input_triton_kernel/fill_gateup_input_triton_kernel.py new file mode 100644 index 00000000..57cb35b9 --- /dev/null +++ b/benchmarks/TritonNPUBench/sglang/fill_gateup_input_triton_kernel/fill_gateup_input_triton_kernel.py @@ -0,0 +1,50 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/sglang_operator_cases/newtest_cases/test_fill_gateup_input_triton_kernel.py +# Main kernel: fill_gateup_input_triton_kernel +# PT file: test_fill_gateup_input_triton_kernel_v2.pt + +import triton +import triton.language as tl + + +# === fill_gateup_input_triton_kernel === +@triton.jit +def fill_gateup_input_triton_kernel( + input_ptr, + scale_ptr, + gateup_input_ptr, + gateup_input_scale_ptr, + src2dst_ptr, + topk_ids_ptr, + topk, + hidden_size, + scale_size, + BLOCK_SIZE: tl.constexpr, +): + + src_idx_int32 = tl.program_id(0) + src_idx = src_idx_int32.to(tl.int64) + src2dst_ptr = src2dst_ptr + src_idx * topk + topk_ids_ptr = topk_ids_ptr + src_idx * topk + src_ptr = input_ptr + src_idx * hidden_size + scale_src_ptr = scale_ptr + src_idx * scale_size + + vec = tl.arange(0, BLOCK_SIZE) + for idx in range(topk): + expert_id = tl.load(topk_ids_ptr + idx) + if expert_id >= 0: + dst_idx_int32 = tl.load(src2dst_ptr + idx) + dst_idx = dst_idx_int32.to(tl.int64) + dst_ptr = gateup_input_ptr + dst_idx * hidden_size + for start_offset in tl.range(0, hidden_size, BLOCK_SIZE): + offset = start_offset + vec + mask = offset < hidden_size + in_data = tl.load(src_ptr + offset, mask=mask) + tl.store(dst_ptr + offset, in_data, mask=mask) + scale_dst_ptr = gateup_input_scale_ptr + dst_idx * scale_size + for start_offset in tl.range(0, scale_size, BLOCK_SIZE): + offset = start_offset + vec + mask = offset < scale_size + in_scale = tl.load(scale_src_ptr + offset, mask=mask) + tl.store(scale_dst_ptr + offset, in_scale, mask=mask) + diff --git a/benchmarks/TritonNPUBench/sglang/fn_triton_kernel/fn_triton_kernel.py b/benchmarks/TritonNPUBench/sglang/fn_triton_kernel/fn_triton_kernel.py new file mode 100644 index 00000000..5a71e938 --- /dev/null +++ b/benchmarks/TritonNPUBench/sglang/fn_triton_kernel/fn_triton_kernel.py @@ -0,0 +1,60 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/sglang_operator_cases/newtest_cases/test_fn_triton_kernel.py +# Main kernel: fn_triton_kernel +# PT file: test_fn_triton_kernel_v2.pt + +import triton +import triton.language as tl + + +# === fn_triton_kernel === +@triton.jit +def fn_triton_kernel( + k_ptr, + k_nope_ptr, + k_rope_ptr, + num_tokens, + QK_NOPE_HEAD_DIM: tl.constexpr, + QK_ROPE_HEAD_DIM: tl.constexpr, + NUM_LOCAL_HEADS: tl.constexpr, + K_NOPE_STRIDE_0: tl.constexpr, + K_NOPE_STRIDE_1: tl.constexpr, + K_STRIDE_0: tl.constexpr, + K_STRIDE_1: tl.constexpr, + K_ROPE_STRIDE_0: tl.constexpr, + BLOCK_ROWS: tl.constexpr, +): + pid = tl.program_id(axis=0) + + token_id = pid * BLOCK_ROWS + tl.arange(0, BLOCK_ROWS) + token_mask = token_id < num_tokens + + head_id = tl.arange(0, NUM_LOCAL_HEADS) + + # nope + nope_sub_id = tl.arange(0, QK_NOPE_HEAD_DIM) + offs_nope = ( + token_id[:, None, None] * K_NOPE_STRIDE_0 + + head_id[None, :, None] * K_NOPE_STRIDE_1 + + nope_sub_id[None, None, :] + ) + offs_k = ( + token_id[:, None, None] * K_STRIDE_0 + + head_id[None, :, None] * K_STRIDE_1 + + nope_sub_id[None, None, :] + ) + vals_nope = tl.load(k_nope_ptr + offs_nope, mask=token_mask[:, None, None]) + tl.store(k_ptr + offs_k, vals_nope, mask=token_mask[:, None, None]) + + # rope + rope_sub_id = tl.arange(0, QK_ROPE_HEAD_DIM) + offs_rope = token_id[:, None, None] * K_ROPE_STRIDE_0 + rope_sub_id[None, None, :] + offs_k = ( + token_id[:, None, None] * K_STRIDE_0 + + head_id[None, :, None] * K_STRIDE_1 + + rope_sub_id[None, None, :] + + QK_NOPE_HEAD_DIM + ) + vals_rope = tl.load(k_rope_ptr + offs_rope, mask=token_mask[:, None, None]) + tl.store(k_ptr + offs_k, vals_rope, mask=token_mask[:, None, None]) + diff --git a/benchmarks/TritonNPUBench/sglang/fused_dual_residual_rmsnorm_kernel/fused_dual_residual_rmsnorm_kernel.py b/benchmarks/TritonNPUBench/sglang/fused_dual_residual_rmsnorm_kernel/fused_dual_residual_rmsnorm_kernel.py new file mode 100644 index 00000000..295007b5 --- /dev/null +++ b/benchmarks/TritonNPUBench/sglang/fused_dual_residual_rmsnorm_kernel/fused_dual_residual_rmsnorm_kernel.py @@ -0,0 +1,55 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/sglang_operator_cases/newtest_cases/test_fused_dual_residual_rmsnorm_kernel.py +# Main kernel: fused_dual_residual_rmsnorm_kernel +# PT file: test_fused_dual_residual_rmsnorm_kernel_v2.pt + +import triton +import triton.language as tl + + +# === fused_dual_residual_rmsnorm_kernel === +@triton.jit +def fused_dual_residual_rmsnorm_kernel( + output_ptr, + mid_ptr, + activ_ptr, + residual_ptr, + weight1_ptr, + weight2_ptr, + eps: tl.constexpr, + hidden_dim: tl.constexpr, + BLOCK_SIZE: tl.constexpr, +): + pid = tl.program_id(axis=0) + input_start = pid * hidden_dim + + offsets = tl.arange(0, BLOCK_SIZE) + mask = offsets < hidden_dim + + a_ = tl.load(activ_ptr + input_start + offsets, mask=mask, other=0.0) + a = a_.to(tl.float32) + rms = tl.sqrt(tl.sum(a * a, axis=0) / hidden_dim + eps) + + r = tl.load(residual_ptr + input_start + offsets, mask=mask, other=0.0) + w1_ = tl.load(weight1_ptr + offsets, mask=mask, other=0.0) + w1 = w1_.to(tl.float32) + + a2r = r + (a / rms * w1).to(r.dtype) + tl.store( + mid_ptr + input_start + offsets, + a2r, + mask=mask, + ) + + a2r = a2r.to(tl.float32) + rms2 = tl.sqrt(tl.sum(a2r * a2r, axis=0) / hidden_dim + eps) + + w2_ = tl.load(weight2_ptr + offsets, mask=mask, other=0.0) + w2 = w2_.to(tl.float32) + + tl.store( + output_ptr + input_start + offsets, + a2r / rms2 * w2, # implicitly casts to output dtype here + mask=mask, + ) + diff --git a/benchmarks/TritonNPUBench/sglang/fused_gdn_gating_kernel/fused_gdn_gating_kernel.py b/benchmarks/TritonNPUBench/sglang/fused_gdn_gating_kernel/fused_gdn_gating_kernel.py new file mode 100644 index 00000000..a2c34d17 --- /dev/null +++ b/benchmarks/TritonNPUBench/sglang/fused_gdn_gating_kernel/fused_gdn_gating_kernel.py @@ -0,0 +1,41 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/sglang_operator_cases/newtest_cases/test_fused_gdn_gating_kernel.py +# Main kernel: fused_gdn_gating_kernel +# PT file: test_fused_gdn_gating_kernel_v2.pt + +import triton +import triton.language as tl + + +# === fused_gdn_gating_kernel === +@triton.jit +def fused_gdn_gating_kernel( + g, + beta_output, + A_log, + a, + b, + dt_bias, + seq_len, + NUM_HEADS: tl.constexpr, + beta: tl.constexpr, + threshold: tl.constexpr, + BLK_HEADS: tl.constexpr, +): + i_b, i_s, i_d = tl.program_id(0), tl.program_id(1), tl.program_id(2) + head_off = i_d * BLK_HEADS + tl.arange(0, BLK_HEADS) + off = i_b * seq_len * NUM_HEADS + i_s * NUM_HEADS + head_off + mask = head_off < NUM_HEADS + blk_A_log = tl.load(A_log + head_off, mask=mask) + blk_a = tl.load(a + off, mask=mask) + blk_b = tl.load(b + off, mask=mask) + blk_bias = tl.load(dt_bias + head_off, mask=mask) + x = blk_a.to(tl.float32) + blk_bias.to(tl.float32) + softplus_x = tl.where( + beta * x <= threshold, (1 / beta) * tl.log(1 + tl.exp(beta * x)), x + ) + blk_g = -tl.exp(blk_A_log.to(tl.float32)) * softplus_x + tl.store(g + off, blk_g.to(g.dtype.element_ty), mask=mask) + blk_beta_output = tl.sigmoid(blk_b.to(tl.float32)) + tl.store(beta_output + off, blk_beta_output.to(b.dtype.element_ty), mask=mask) + diff --git a/benchmarks/TritonNPUBench/sglang/fused_moe_kernel/fused_moe_kernel.py b/benchmarks/TritonNPUBench/sglang/fused_moe_kernel/fused_moe_kernel.py new file mode 100644 index 00000000..e44edca9 --- /dev/null +++ b/benchmarks/TritonNPUBench/sglang/fused_moe_kernel/fused_moe_kernel.py @@ -0,0 +1,291 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/sglang_operator_cases/newtest_cases/test_fused_moe_kernel.py +# Main kernel: fused_moe_kernel +# PT file: test_fused_moe_kernel_v2.pt + +import triton +import triton.language as tl + + +# === fused_moe_kernel === +@triton.jit +def fused_moe_kernel( + # Pointers to matrices + a_ptr, + b_ptr, + c_ptr, + b_bias_ptr, + a_scale_ptr, + b_scale_ptr, + topk_weights_ptr, + sorted_token_ids_ptr, + expert_ids_ptr, + num_tokens_post_padded_ptr, + # Matrix dimensions + N, + K, + EM, + num_valid_tokens, + # The stride variables represent how much to increase the ptr by when + # moving by 1 element in a particular dimension. E.g. `stride_am` is + # how much to increase `a_ptr` by to get the element one row down + # (A has M rows). + stride_am, + stride_ak, + stride_be, + stride_bk, + stride_bn, + stride_cm, + stride_cn, + stride_asm, + stride_ask, + stride_bse, + stride_bsk, + stride_bsn, + stride_bbe, # bias expert stride + stride_bbn, # bias N stride + # Block size for block-wise quantization + group_n: tl.constexpr, + group_k: tl.constexpr, + naive_block_assignment: tl.constexpr, + # Meta-parameters + BLOCK_SIZE_M: tl.constexpr, + BLOCK_SIZE_N: tl.constexpr, + BLOCK_SIZE_K: tl.constexpr, + GROUP_SIZE_M: tl.constexpr, + SPLIT_K: tl.constexpr, + MUL_ROUTED_WEIGHT: tl.constexpr, + top_k: tl.constexpr, + compute_type: tl.constexpr, + use_fp8_w8a8: tl.constexpr, + use_int8_w8a8: tl.constexpr, + use_int8_w8a16: tl.constexpr, + per_channel_quant: tl.constexpr, + HAS_BIAS: tl.constexpr, +): + """ + Implements the fused computation for a Mixture of Experts (MOE) using + token and expert matrices. + + Key Parameters: + - A: The input tensor representing tokens with shape (*, K), where '*' can + be any shape representing batches and K is the feature dimension of + each token. + - B: The stacked MOE weight tensor with shape (E, N, K), where E is + the number of experts, K is the input feature dimension, and N is + the output feature dimension. + - C: The output cache tensor with shape (M, topk, N), where M is the + total number of tokens post padding, topk is the number of times + each token is repeated, and N is the output feature dimension. + - sorted_token_ids: A tensor containing the sorted indices of tokens, + repeated topk times and arranged by the expert index they are + assigned to. + - expert_ids: A tensor containing the indices of the expert for each + block. It determines which expert matrix from B should be used for + each block in A. + - naive_block_assignment: A boolean flag indicating whether to use naive + token wise block assignment. If True, each block corresponds to a + single token. + This kernel performs the multiplication of a token by its corresponding + expert matrix as determined by `expert_ids`. The sorting of + `sorted_token_ids` by expert index and padding ensures divisibility by + BLOCK_SIZE_M, which is necessary to maintain consistency in block matrix + multiplication across different blocks processed by the same expert. + """ + # ----------------------------------------------------------- + # Map program ids `pid` to the block of C it should compute. + # This is done in a grouped ordering to promote L2 data reuse. + pid = tl.program_id(axis=0) + num_pid_m = tl.cdiv(EM, BLOCK_SIZE_M) + num_pid_n = tl.cdiv(N, BLOCK_SIZE_N) + num_pid_in_group = GROUP_SIZE_M * num_pid_n + group_id = pid // num_pid_in_group + first_pid_m = group_id * GROUP_SIZE_M + group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M) + pid_m = first_pid_m + ((pid % num_pid_in_group) % group_size_m) + pid_n = (pid % num_pid_in_group) // group_size_m + + # ---------------------------------------------------------- + # Create pointers for the first blocks of A and B. + # We will advance this pointer as we move in the K direction + # and accumulate + # `a_ptrs` is a block of [BLOCK_SIZE_M, BLOCK_SIZE_K] pointers + # `b_ptrs` is a block of [BLOCK_SIZE_K, BLOCK_SIZE_N] pointers + offs = tl.arange(0, BLOCK_SIZE_M).to(tl.int64) + num_tokens_post_padded = tl.load(num_tokens_post_padded_ptr) + if pid_m * BLOCK_SIZE_M >= num_tokens_post_padded: + return + if not naive_block_assignment: + offs_token_id = pid_m * BLOCK_SIZE_M + offs + offs_token = tl.load(sorted_token_ids_ptr + offs_token_id) + else: + offs_token = tl.where( + offs == 0, + pid_m, # first element = pid_m + num_valid_tokens, # remaining elements = constant + ) + + token_mask = offs_token < num_valid_tokens + + off_experts = tl.load(expert_ids_ptr + pid_m).to(tl.int64) + if off_experts == -1: + # ----------------------------------------------------------- + # Write back zeros to the output when the expert is not + # in the current expert parallel rank. + write_zeros_to_output( + c_ptr, + stride_cm, + stride_cn, + pid_n, + N, + offs_token, + token_mask, + BLOCK_SIZE_M, + BLOCK_SIZE_N, + compute_type, + ) + return + + offs_bn = (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N).to(tl.int64)) % N + offs_k = tl.arange(0, BLOCK_SIZE_K) + a_ptrs = a_ptr + ( + offs_token[:, None] // top_k * stride_am + offs_k[None, :] * stride_ak + ) + + b_ptrs = ( + b_ptr + + off_experts * stride_be + + (offs_k[:, None] * stride_bk + offs_bn[None, :] * stride_bn) + ) + if use_int8_w8a16: + b_scale_ptrs = ( + b_scale_ptr + off_experts * stride_bse + offs_bn[None, :] * stride_bsn + ) + b_scale = tl.load(b_scale_ptrs) + + if use_fp8_w8a8 or use_int8_w8a8: + # block-wise + if group_k > 0 and group_n > 0: + a_scale_ptrs = a_scale_ptr + (offs_token // top_k) * stride_asm + offs_bsn = offs_bn // group_n + b_scale_ptrs = ( + b_scale_ptr + off_experts * stride_bse + offs_bsn * stride_bsn + ) + # channel-wise + elif per_channel_quant: + b_scale_ptrs = ( + b_scale_ptr + off_experts * stride_bse + offs_bn[None, :] * stride_bsn + ) + b_scale = tl.load(b_scale_ptrs) + # Load per-token scale for activations + a_scale_ptrs = a_scale_ptr + (offs_token // top_k) * stride_asm + a_scale = tl.load(a_scale_ptrs, mask=token_mask, other=0.0)[:, None] + # tensor-wise + else: + a_scale = tl.load(a_scale_ptr) + b_scale = tl.load(b_scale_ptr + off_experts) + if HAS_BIAS: + # bias shape: [num_experts, N] + bias_ptrs = b_bias_ptr + off_experts * stride_bbe + offs_bn * stride_bbn + bias = tl.load(bias_ptrs, mask=(offs_bn < N), other=0.0) + # ----------------------------------------------------------- + # Iterate to compute a block of the C matrix. + # We accumulate into a `[BLOCK_SIZE_M, BLOCK_SIZE_N]` block + # of fp32 values for higher accuracy. + # `accumulator` will be converted back to fp16 after the loop. + accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32) + for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)): + # Load the next block of A and B, generate a mask by checking the + # K dimension. + a = tl.load( + a_ptrs, + mask=token_mask[:, None] & (offs_k[None, :] < K - k * BLOCK_SIZE_K), + other=0.0, + ) + b = tl.load(b_ptrs, mask=offs_k[:, None] < K - k * BLOCK_SIZE_K, other=0.0) + # We accumulate along the K dimension. + if use_int8_w8a16: + accumulator = tl.dot(a, b.to(compute_type), acc=accumulator) + elif use_fp8_w8a8 or use_int8_w8a8: + if group_k > 0 and group_n > 0: + k_start = k * BLOCK_SIZE_K + offs_ks = k_start // group_k + a_scale = tl.load( + a_scale_ptrs + offs_ks * stride_ask, mask=token_mask, other=0.0 + ) + b_scale = tl.load(b_scale_ptrs + offs_ks * stride_bsk) + + accumulator += tl.dot(a, b) * a_scale[:, None] * b_scale[None, :] + else: + if use_fp8_w8a8: + # acc used to enable fp8_fast_accum + accumulator = tl.dot(a, b, acc=accumulator) + else: + accumulator += tl.dot(a, b) + else: + accumulator += tl.dot(a, b) + # Advance the ptrs to the next K block. + a_ptrs += BLOCK_SIZE_K * stride_ak + b_ptrs += BLOCK_SIZE_K * stride_bk + + # Dequantization for supported quantization schemes: + # - int8_w8a16 + # - fp8_w8a8 + # - int8_w8a8 + # Accumulator and scalings are in float32 to preserve numerical accuracy. + if use_int8_w8a16: + accumulator = accumulator * b_scale + elif (use_fp8_w8a8 or use_int8_w8a8) and not (group_k > 0 and group_n > 0): + accumulator = accumulator * a_scale * b_scale + + # Bias addition: + # Bias must be applied after dequantization: + # - Since bias is typically not quantized + # - Bias should not be scaled by quantization factors + if HAS_BIAS: + accumulator += bias[None, :] + + # Router (MoE) weight multiplication: + # This multiplication MUST be performed in float32 before any precision + # conversion to ensure numerical stability, which is especially critical + # on ROCm platforms. + if MUL_ROUTED_WEIGHT: + moe_weight = tl.load( + topk_weights_ptr + offs_token, + mask=token_mask, + other=0, + ) + accumulator *= moe_weight[:, None] + + # Final precision conversion: + # Cast once at the end to the desired compute/output dtype. + accumulator = accumulator.to(compute_type) + + # ----------------------------------------------------------- + # Write back the block of the output + offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) + c_ptrs = c_ptr + stride_cm * offs_token[:, None] + stride_cn * offs_cn[None, :] + c_mask = token_mask[:, None] & (offs_cn[None, :] < N) + tl.store(c_ptrs, accumulator, mask=c_mask) + + +# === write_zeros_to_output === +@triton.jit +def write_zeros_to_output( + c_ptr, + stride_cm, + stride_cn, + pid_n, + N, + offs_token, + token_mask, + BLOCK_SIZE_M, + BLOCK_SIZE_N, + compute_type, +): + accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=compute_type) + offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) + c_ptrs = c_ptr + stride_cm * offs_token[:, None] + stride_cn * offs_cn[None, :] + c_mask = token_mask[:, None] & (offs_cn[None, :] < N) + tl.store(c_ptrs, accumulator, mask=c_mask) + diff --git a/benchmarks/TritonNPUBench/sglang/fused_moe_kernel_gptq_awq/fused_moe_kernel_gptq_awq.py b/benchmarks/TritonNPUBench/sglang/fused_moe_kernel_gptq_awq/fused_moe_kernel_gptq_awq.py new file mode 100644 index 00000000..f58da42c --- /dev/null +++ b/benchmarks/TritonNPUBench/sglang/fused_moe_kernel_gptq_awq/fused_moe_kernel_gptq_awq.py @@ -0,0 +1,262 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/sglang_operator_cases/newtest_cases/test_fused_moe_kernel_gptq_awq.py +# Main kernel: fused_moe_kernel_gptq_awq +# PT file: test_fused_moe_kernel_gptq_awq_v2.pt + +import triton +import triton.language as tl + + +# === fused_moe_kernel_gptq_awq === +@triton.jit +def fused_moe_kernel_gptq_awq( + # Pointers to matrices + a_ptr, + b_ptr, + c_ptr, + b_scale_ptr, + b_zp_ptr, + topk_weights_ptr, + sorted_token_ids_ptr, + expert_ids_ptr, + num_tokens_post_padded_ptr, + # Matrix dimensions + N: tl.constexpr, + K: tl.constexpr, + EM, + num_valid_tokens, + # The stride variables represent how much to increase the ptr by when + # moving by 1 element in a particular dimension. E.g. `stride_am` is + # how much to increase `a_ptr` by to get the element one row down + # (A has M rows). + stride_am, + stride_ak, + stride_be, + stride_bk, + stride_bn, + stride_cm, + stride_cn, + stride_bse, + stride_bsk, + stride_bsn, + stride_bze, + stride_bzk, + stride_bzn, + group_size: tl.constexpr, + # Meta-parameters + BLOCK_SIZE_M: tl.constexpr, + BLOCK_SIZE_N: tl.constexpr, + BLOCK_SIZE_K: tl.constexpr, + GROUP_SIZE_M: tl.constexpr, + MUL_ROUTED_WEIGHT: tl.constexpr, + top_k: tl.constexpr, + compute_type: tl.constexpr, + has_zp: tl.constexpr, + use_int4_w4a16: tl.constexpr, + use_int8_w8a16: tl.constexpr, + even_Ks: tl.constexpr, + filter_expert: tl.constexpr, +): + """ + Implements the fused computation for a Mixture of Experts (MOE) using + token and expert matrices. + Key Parameters: + - A: The input tensor representing tokens with shape (*, K), where '*' can + be any shape representing batches and K is the feature dimension of + each token. + - B: The stacked MOE weight tensor with shape (E, N, K), where E is + the number of experts, K is the input feature dimension, and N is + the output feature dimension. + - C: The output cache tensor with shape (M, topk, N), where M is the + total number of tokens post padding, topk is the number of times + each token is repeated, and N is the output feature dimension. + - sorted_token_ids: A tensor containing the sorted indices of tokens, + repeated topk times and arranged by the expert index they are + assigned to. + - expert_ids: A tensor containing the indices of the expert for each + block. It determines which expert matrix from B should be used for + each block in A. + This kernel performs the multiplication of a token by its corresponding + expert matrix as determined by `expert_ids`. The sorting of + `sorted_token_ids` by expert index and padding ensures divisibility by + BLOCK_SIZE_M, which is necessary to maintain consistency in block matrix + multiplication across different blocks processed by the same expert. + """ + # ----------------------------------------------------------- + # Map program ids `pid` to the block of C it should compute. + # This is done in a grouped ordering to promote L2 data reuse. + pid = tl.program_id(axis=0) + num_pid_m = tl.cdiv(EM, BLOCK_SIZE_M) + num_pid_n = tl.cdiv(N, BLOCK_SIZE_N) + num_pid_in_group = GROUP_SIZE_M * num_pid_n + group_id = pid // num_pid_in_group + first_pid_m = group_id * GROUP_SIZE_M + group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M) + pid_m = first_pid_m + ((pid % num_pid_in_group) % group_size_m) + pid_n = (pid % num_pid_in_group) // group_size_m + + # ---------------------------------------------------------- + # Create pointers for the first blocks of A and B. + # We will advance this pointer as we move in the K direction + # and accumulate + # `a_ptrs` is a block of [BLOCK_SIZE_M, BLOCK_SIZE_K] pointers + # `b_ptrs` is a block of [BLOCK_SIZE_K, BLOCK_SIZE_N] pointers + num_tokens_post_padded = tl.load(num_tokens_post_padded_ptr) + if pid_m * BLOCK_SIZE_M >= num_tokens_post_padded: + return + offs_token_id = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M).to(tl.int64) + offs_token = tl.load(sorted_token_ids_ptr + offs_token_id) + token_mask = offs_token < num_valid_tokens + + off_experts = tl.load(expert_ids_ptr + pid_m).to(tl.int64) + if filter_expert and off_experts == -1: + # ----------------------------------------------------------- + # Write back zeros to the output when the expert is not + # in the current expert parallel rank. + write_zeros_to_output( + c_ptr, + stride_cm, + stride_cn, + pid_n, + N, + offs_token, + token_mask, + BLOCK_SIZE_M, + BLOCK_SIZE_N, + compute_type, + ) + return + + offs_bn = (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N).to(tl.int64)) % N + offs_k = tl.arange(0, BLOCK_SIZE_K) + a_ptrs = a_ptr + ( + offs_token[:, None] // top_k * stride_am + offs_k[None, :] * stride_ak + ) + + if use_int4_w4a16: + b_ptrs = ( + b_ptr + + off_experts * stride_be + + (offs_k[:, None] // 2) * stride_bk + + offs_bn[None, :] * stride_bn + ) + b_shifter = (offs_k[:, None] % 2) * 4 + elif use_int8_w8a16: + b_ptrs = ( + b_ptr + + off_experts * stride_be + + offs_k[:, None] * stride_bk + + offs_bn[None, :] * stride_bn + ) + + if not has_zp and use_int4_w4a16: + b_zp_num = 8 + if not has_zp and use_int8_w8a16: + b_zp_num = 128 + elif has_zp and use_int4_w4a16: + b_zp_shifter = (offs_bn[None, :] % 2) * 4 + + # ----------------------------------------------------------- + # Iterate to compute a block of the C matrix. + # We accumulate into a `[BLOCK_SIZE_M, BLOCK_SIZE_N]` block + # of fp32 values for higher accuracy. + # `accumulator` will be converted back to fp16 after the loop. + accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32) + for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)): + # Load the next block of A and B, generate a mask by checking the + # K dimension. + + if not even_Ks: + k_mask = offs_k[:, None] < K - k * BLOCK_SIZE_K + k_other = 0.0 + else: + k_mask = None + k_other = None + + a = tl.load( + a_ptrs, + mask=token_mask[:, None] & (offs_k[None, :] < K - k * BLOCK_SIZE_K), + other=0.0, + ) + b = tl.load(b_ptrs) + if use_int4_w4a16: + b = (b >> b_shifter) & 0xF + + b_scale_ptrs = ( + b_scale_ptr + + off_experts * stride_bse + + offs_bn[None, :] * stride_bsn + + ((offs_k[:, None] + BLOCK_SIZE_K * k) // group_size) * stride_bsk + ) + b_scale = tl.load(b_scale_ptrs, mask=k_mask, other=k_other) + b_scale = b_scale.to(tl.float32) + + if has_zp and use_int4_w4a16: + offs_k_true = (offs_k[:, None] + BLOCK_SIZE_K * k) // group_size + b_zp_ptrs = ( + b_zp_ptr + + off_experts * stride_bze + + (offs_bn[None, :] // 2) * stride_bzn + + offs_k_true * stride_bzk + ) + b_zp = tl.load(b_zp_ptrs, mask=k_mask, other=k_other) + b_zp = (b_zp >> b_zp_shifter) & 0xF + b_zp = b_zp.to(tl.float32) + elif has_zp and use_int8_w8a16: + offs_k_true = (offs_k[:, None] + BLOCK_SIZE_K * k) // group_size + b_zp_ptrs = ( + b_zp_ptr + + off_experts * stride_bze + + offs_bn[None, :] * stride_bzn + + offs_k_true * stride_bzk + ) + b_zp = tl.load(b_zp_ptrs, mask=k_mask, other=k_other) + b_zp = b_zp.to(tl.float32) + + # We accumulate along the K dimension. + if has_zp: + b = ((b.to(tl.float32) - b_zp) * b_scale).to(compute_type) + else: + b = ((b.to(tl.float32) - b_zp_num) * b_scale).to(compute_type) + accumulator = tl.dot(a, b, acc=accumulator) + + # Advance the ptrs to the next K block. + a_ptrs += BLOCK_SIZE_K * stride_ak + if use_int4_w4a16: + b_ptrs += (BLOCK_SIZE_K // 2) * stride_bk + else: + b_ptrs += BLOCK_SIZE_K * stride_bk + + if MUL_ROUTED_WEIGHT: + moe_weight = tl.load(topk_weights_ptr + offs_token, mask=token_mask, other=0) + accumulator = accumulator * moe_weight[:, None] + + accumulator = accumulator.to(compute_type) + # ----------------------------------------------------------- + # Write back the block of the output + offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) + c_ptrs = c_ptr + stride_cm * offs_token[:, None] + stride_cn * offs_cn[None, :] + c_mask = token_mask[:, None] & (offs_cn[None, :] < N) + tl.store(c_ptrs, accumulator, mask=c_mask) + + +# === write_zeros_to_output === +@triton.jit +def write_zeros_to_output( + c_ptr, + stride_cm, + stride_cn, + pid_n, + N, + offs_token, + token_mask, + BLOCK_SIZE_M, + BLOCK_SIZE_N, + compute_type, +): + accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=compute_type) + offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) + c_ptrs = c_ptr + stride_cm * offs_token[:, None] + stride_cn * offs_cn[None, :] + c_mask = token_mask[:, None] & (offs_cn[None, :] < N) + tl.store(c_ptrs, accumulator, mask=c_mask) + diff --git a/benchmarks/TritonNPUBench/sglang/fused_moe_router_cudacore_kernel/fused_moe_router_cudacore_kernel.py b/benchmarks/TritonNPUBench/sglang/fused_moe_router_cudacore_kernel/fused_moe_router_cudacore_kernel.py new file mode 100644 index 00000000..566fe4a1 --- /dev/null +++ b/benchmarks/TritonNPUBench/sglang/fused_moe_router_cudacore_kernel/fused_moe_router_cudacore_kernel.py @@ -0,0 +1,112 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/sglang_operator_cases/newtest_cases/test_fused_moe_router_cudacore_kernel.py +# Main kernel: fused_moe_router_cudacore_kernel +# PT file: test_fused_moe_router_cudacore_kernel_v2.pt + +import triton +import triton.language as tl + + +# === fused_moe_router_cudacore_kernel === +@triton.jit +def fused_moe_router_cudacore_kernel( + input_ptr, # input (bs, hidden_dim) + moe_router_weight_ptr, # input (num_experts, hidden_dim) + topk_weights_ptr, # output (bs, topk) + topk_ids_ptr, # output (bs, topk) + correction_bias_ptr, + is_correction_bias: tl.constexpr, + num_experts: tl.constexpr, + topk: tl.constexpr, + moe_softcapping: tl.constexpr, + moe_renormalize: tl.constexpr, # not supported + hidden_dim: tl.constexpr, + BLOCK_SIZE: tl.constexpr, +): + pid = tl.program_id(axis=0) + + offsets = tl.arange(0, BLOCK_SIZE) + mask = offsets < hidden_dim + + # moe_router_weight is k major + expert_offsets = tl.arange(0, num_experts)[:, None] + router_mask = mask[None, :] + w_router = tl.load( + moe_router_weight_ptr + expert_offsets * hidden_dim + offsets[None, :], + mask=router_mask, + other=0.0, + ) + + x = tl.load(input_ptr + pid * hidden_dim + offsets, mask=mask, other=0.0) + + # todo: tl.dot? + logits = tl.sum((w_router.to(tl.float32) * x[None, :].to(tl.float32)), axis=-1) + + # logit softcap + if moe_softcapping == 0: + logits_softcapped = logits + else: + logits_scaled = logits / moe_softcapping + exped = tl.exp(2 * logits_scaled) + top = exped - 1 + bottom = exped + 1 + logits_softcapped = top / bottom * moe_softcapping + + # Add bias after softcapping + if is_correction_bias: + bias = tl.load(correction_bias_ptr + tl.arange(0, num_experts)) + logits_softcapped = logits_softcapped + bias + + # topk + # assert 1 <= topk <= num_experts + + # 5.38 us + + top1 = tl.argmax(logits_softcapped, axis=0) + tl.store(topk_ids_ptr + pid * topk + 0, top1) # 5.63 us + + top1_v = tl.max(logits_softcapped, axis=0) + invsumexp = 1.0 / tl.sum(tl.exp(logits_softcapped - top1_v), axis=0) + + tl.store( + topk_weights_ptr + pid * topk + 0, + invsumexp, + ) # 5.73 us + + if topk >= 2: + top2 = tl.argmax( + tl.where( + tl.arange(0, num_experts) != top1, logits_softcapped, float("-inf") + ), + axis=0, + ) + tl.store(topk_ids_ptr + pid * topk + 1, top2) + top2_v = tl.sum(logits_softcapped * (tl.arange(0, num_experts) == top2), axis=0) + tl.store( + topk_weights_ptr + pid * topk + 1, + tl.exp(top2_v - top1_v) * invsumexp, + ) # 5.95us + + # probably slow + if topk > 2: + topk_mask = tl.full(logits_softcapped.shape, 1.0, dtype=logits_softcapped.dtype) + topk_mask = tl.where( + tl.arange(0, num_experts) != top1, topk_mask, float("-inf") + ) + topk_mask = tl.where( + tl.arange(0, num_experts) != top2, topk_mask, float("-inf") + ) + for i in range(2, topk): + topi = tl.argmax(logits_softcapped + topk_mask, axis=0) + topk_mask = tl.where( + tl.arange(0, num_experts) != topi, topk_mask, float("-inf") + ) + tl.store(topk_ids_ptr + pid * topk + i, topi) + topi_v = tl.sum( + logits_softcapped * (tl.arange(0, num_experts) == topi), axis=0 + ) + tl.store( + topk_weights_ptr + pid * topk + i, + tl.exp(topi_v - top1_v) * invsumexp, + ) + diff --git a/benchmarks/TritonNPUBench/sglang/fused_moe_router_tensorcore_kernel/fused_moe_router_tensorcore_kernel.py b/benchmarks/TritonNPUBench/sglang/fused_moe_router_tensorcore_kernel/fused_moe_router_tensorcore_kernel.py new file mode 100644 index 00000000..88d38298 --- /dev/null +++ b/benchmarks/TritonNPUBench/sglang/fused_moe_router_tensorcore_kernel/fused_moe_router_tensorcore_kernel.py @@ -0,0 +1,136 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/sglang_operator_cases/newtest_cases/test_fused_moe_router_tensorcore_kernel.py +# Main kernel: fused_moe_router_tensorcore_kernel +# PT file: test_fused_moe_router_tensorcore_kernel_v2.pt + +import triton +import triton.language as tl + + +# === fused_moe_router_tensorcore_kernel === +@triton.jit +def fused_moe_router_tensorcore_kernel( + a_ptr, # input (bs, hidden_dim) + b_ptr, # input (num_experts, hidden_dim) + topk_weights_ptr, # output (bs, topk) + topk_ids_ptr, # output (bs, topk) + bs, + num_experts: tl.constexpr, + topk: tl.constexpr, # only support topk <= 2 + moe_softcapping: tl.constexpr, + moe_renormalize: tl.constexpr, # not supported + correction_bias_ptr, + is_correction_bias: tl.constexpr, + K: tl.constexpr, + BLOCK_SIZE_M: tl.constexpr, + BLOCK_SIZE_N: tl.constexpr, + BLOCK_SIZE_K: tl.constexpr, + stride_am: tl.constexpr, + stride_bn: tl.constexpr, + dp_attn_workaround_flag: tl.constexpr, +): + + # 1. get block id + pid = tl.program_id(axis=0) + + # 2. create pointers for the first block of A and B + # 2.1. setup a_ptrs with offsets in m and k + offs_m = pid * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)[:, None] + bs_mask = offs_m < bs + offs_k = tl.arange(0, BLOCK_SIZE_K)[None, :] + a_ptrs = a_ptr + (offs_m * stride_am + offs_k) + + # 2.2. setup b_ptrs with offsets in k and n. + # Note: b matrix is k-major. + offs_k = tl.arange(0, BLOCK_SIZE_K)[None, :] + offs_n = tl.arange(0, BLOCK_SIZE_N)[:, None] + expert_mask = offs_n < num_experts + b_ptrs = b_ptr + (offs_n * stride_bn + offs_k) + + # 3. Create an accumulator of float32 of size [BLOCK_SIZE_M, BLOCK_SIZE_N] + # 3.1. iterate in K dimension + # 3.2. transpose tile B + acc = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32) + for k in range(0, K // BLOCK_SIZE_K): # hidden_dim % BLOCK_SIZE_K == 0 + a = tl.load( + a_ptrs, + mask=bs_mask, + other=0.0, + ).to(tl.float32) + b = tl.load(b_ptrs, mask=expert_mask, other=0.0).to(tl.float32).T + acc += tl.dot(a, b) + + # Advance the ptrs to the next K block. + a_ptrs += BLOCK_SIZE_K + b_ptrs += BLOCK_SIZE_K + + # 4. logit softcap + if moe_softcapping == 0: + logits_softcapped = acc + else: + logits_scaled = acc / moe_softcapping + exped = tl.exp(2 * logits_scaled) + logits_softcapped = (exped - 1) / (exped + 1) * moe_softcapping + + # Add bias after softcapping + if is_correction_bias: + bias = tl.load( + correction_bias_ptr + tl.arange(0, BLOCK_SIZE_N)[None, :], + mask=expert_mask.T, + other=0.0, + ) + logits_softcapped = logits_softcapped + bias + + if dp_attn_workaround_flag: + logits_softcapped = tl.where( + logits_softcapped != logits_softcapped, -1e9, logits_softcapped + ) + + # 5. top1 + arange_block_size_n = tl.arange(0, BLOCK_SIZE_N)[None, :] + cond_top1 = arange_block_size_n < num_experts + top1 = tl.argmax(tl.where(cond_top1, logits_softcapped, float("-inf")), axis=1) + top1_v = tl.max( + tl.where(cond_top1, logits_softcapped, float("-inf")), axis=1, keep_dims=True + ) + top1_invsumexp = 1.0 / tl.sum( + tl.where(cond_top1, tl.exp(logits_softcapped - top1_v), 0.0), axis=1 + ) + + # 6. store top1 to output + offs_top1 = pid * topk * BLOCK_SIZE_M + topk * tl.arange(0, BLOCK_SIZE_M) + top1_mask = offs_top1 < bs * topk + tl.store(topk_ids_ptr + offs_top1, top1, mask=top1_mask) + tl.store( + topk_weights_ptr + offs_top1, + top1_invsumexp, + mask=top1_mask, + ) + + # 7. handle topk == 2 + if topk == 2: + cond_top2 = (arange_block_size_n < num_experts) & ( + arange_block_size_n != top1[:, None] + ) + top2 = tl.argmax( + tl.where(cond_top2, logits_softcapped, float("-inf")), + axis=1, + keep_dims=True, + ) + top2_v = tl.sum( + logits_softcapped * (arange_block_size_n == top2), axis=1, keep_dims=True + ) + top2_invsumexp = tl.exp(top2_v - top1_v) * top1_invsumexp[:, None] + + # store top2 + offs_top2 = ( + pid * topk * BLOCK_SIZE_M + topk * tl.arange(0, BLOCK_SIZE_M)[:, None] + 1 + ) + top2_mask = offs_top2 < bs * topk + tl.store(topk_ids_ptr + offs_top2, top2, mask=top2_mask) + tl.store( + topk_weights_ptr + offs_top2, + top2_invsumexp, + mask=top2_mask, + ) + diff --git a/benchmarks/TritonNPUBench/sglang/fused_recurrent_gated_delta_rule_fwd_kernel/fused_recurrent_gated_delta_rule_fwd_kernel.py b/benchmarks/TritonNPUBench/sglang/fused_recurrent_gated_delta_rule_fwd_kernel/fused_recurrent_gated_delta_rule_fwd_kernel.py new file mode 100644 index 00000000..2a9a5bba --- /dev/null +++ b/benchmarks/TritonNPUBench/sglang/fused_recurrent_gated_delta_rule_fwd_kernel/fused_recurrent_gated_delta_rule_fwd_kernel.py @@ -0,0 +1,158 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/sglang_operator_cases/newtest_cases/test_fused_recurrent_gated_delta_rule_fwd_kernel.py +# Main kernel: fused_recurrent_gated_delta_rule_fwd_kernel +# PT file: test_fused_recurrent_gated_delta_rule_fwd_kernel_v2.pt + +import triton +import triton.language as tl + + +# === fused_recurrent_gated_delta_rule_fwd_kernel === +@triton.jit(do_not_specialize=["N", "T"]) +def fused_recurrent_gated_delta_rule_fwd_kernel( + q, + k, + v, + g, + beta, + o, + h0, + ht, + cu_seqlens, + ssm_state_indices, + num_accepted_tokens, + scale, + N: tl.int64, # num of sequences + T: tl.int64, # num of tokens + B: tl.constexpr, + H: tl.constexpr, + HV: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + stride_init_state_token: tl.constexpr, + stride_final_state_token: tl.constexpr, + stride_indices_seq: tl.constexpr, + stride_indices_tok: tl.constexpr, + USE_INITIAL_STATE: tl.constexpr, # whether to use initial state + INPLACE_FINAL_STATE: tl.constexpr, # whether to store final state inplace + IS_BETA_HEADWISE: tl.constexpr, # whether beta is headwise vector or scalar, + USE_QK_L2NORM_IN_KERNEL: tl.constexpr, + IS_VARLEN: tl.constexpr, + IS_CONTINUOUS_BATCHING: tl.constexpr, + IS_SPEC_DECODING: tl.constexpr, + IS_KDA: tl.constexpr, +): + i_k, i_v, i_nh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_n, i_hv = i_nh // HV, i_nh % HV + i_h = i_hv // (HV // H) + if IS_VARLEN: + bos, eos = ( + tl.load(cu_seqlens + i_n).to(tl.int64), + tl.load(cu_seqlens + i_n + 1).to(tl.int64), + ) + all = T + T = eos - bos + else: + bos, eos = i_n * T, i_n * T + T + all = B * T + + if T == 0: + # no tokens to process for this sequence + return + + o_k = i_k * BK + tl.arange(0, BK) + o_v = i_v * BV + tl.arange(0, BV) + + p_q = q + (bos * H + i_h) * K + o_k + p_k = k + (bos * H + i_h) * K + o_k + p_v = v + (bos * HV + i_hv) * V + o_v + if IS_BETA_HEADWISE: + p_beta = beta + (bos * HV + i_hv) * V + o_v + else: + p_beta = beta + bos * HV + i_hv + + if not IS_KDA: + p_g = g + bos * HV + i_hv + else: + p_gk = g + (bos * HV + i_hv) * K + o_k + + p_o = o + ((i_k * all + bos) * HV + i_hv) * V + o_v + + mask_k = o_k < K + mask_v = o_v < V + mask_h = mask_k[:, None] & mask_v[None, :] + + b_h = tl.zeros([BK, BV], dtype=tl.float32) + if USE_INITIAL_STATE: + if IS_CONTINUOUS_BATCHING: + if IS_SPEC_DECODING: + i_t = tl.load(num_accepted_tokens + i_n).to(tl.int64) - 1 + else: + i_t = 0 + p_h0 = ( + h0 + + tl.load(ssm_state_indices + i_n * stride_indices_seq + i_t).to( + tl.int64 + ) + * stride_init_state_token + ) + else: + p_h0 = h0 + bos * HV * K * V + p_h0 = p_h0 + i_hv * K * V + o_k[:, None] * V + o_v[None, :] + b_h += tl.load(p_h0, mask=mask_h, other=0).to(tl.float32) + + for i_t in range(0, T): + b_q = tl.load(p_q, mask=mask_k, other=0).to(tl.float32) + b_k = tl.load(p_k, mask=mask_k, other=0).to(tl.float32) + b_v = tl.load(p_v, mask=mask_v, other=0).to(tl.float32) + + if USE_QK_L2NORM_IN_KERNEL: + b_q = b_q / tl.sqrt(tl.sum(b_q * b_q) + 1e-6) + b_k = b_k / tl.sqrt(tl.sum(b_k * b_k) + 1e-6) + b_q = b_q * scale + # [BK, BV] + if not IS_KDA: + b_g = tl.load(p_g).to(tl.float32) + b_h *= tl.exp(b_g) + else: + b_gk = tl.load(p_gk).to(tl.float32) + b_h *= tl.exp(b_gk[:, None]) + # [BV] + b_v -= tl.sum(b_h * b_k[:, None], 0) + if IS_BETA_HEADWISE: + b_beta = tl.load(p_beta, mask=mask_v, other=0).to(tl.float32) + else: + b_beta = tl.load(p_beta).to(tl.float32) + b_v *= b_beta + # [BK, BV] + b_h += b_k[:, None] * b_v[None, :] + # [BV] + b_o = tl.sum(b_h * b_q[:, None], 0) + tl.store(p_o, b_o.to(p_o.dtype.element_ty), mask=mask_v) + + # keep the states for multi-query tokens + if INPLACE_FINAL_STATE: + p_ht = ( + ht + + tl.load(ssm_state_indices + i_n * stride_indices_seq + i_t).to( + tl.int64 + ) + * stride_final_state_token + ) + else: + p_ht = ht + (bos + i_t) * stride_final_state_token + p_ht = p_ht + i_hv * K * V + o_k[:, None] * V + o_v[None, :] + tl.store(p_ht, b_h.to(p_ht.dtype.element_ty), mask=mask_h) + + p_q += H * K + p_k += H * K + p_o += HV * V + p_v += HV * V + if not IS_KDA: + p_g += HV + else: + p_gk += HV * K + p_beta += HV * (V if IS_BETA_HEADWISE else 1) + diff --git a/benchmarks/TritonNPUBench/sglang/fused_recurrent_gated_delta_rule_update_fwd_kernel/fused_recurrent_gated_delta_rule_update_fwd_kernel.py b/benchmarks/TritonNPUBench/sglang/fused_recurrent_gated_delta_rule_update_fwd_kernel/fused_recurrent_gated_delta_rule_update_fwd_kernel.py new file mode 100644 index 00000000..701436cb --- /dev/null +++ b/benchmarks/TritonNPUBench/sglang/fused_recurrent_gated_delta_rule_update_fwd_kernel/fused_recurrent_gated_delta_rule_update_fwd_kernel.py @@ -0,0 +1,188 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/sglang_operator_cases/newtest_cases/test_fused_recurrent_gated_delta_rule_update_fwd_kernel.py +# Main kernel: fused_recurrent_gated_delta_rule_update_fwd_kernel +# PT file: test_fused_recurrent_gated_delta_rule_update_fwd_kernel_v2_uint2.pt + +import triton +import triton.language as tl + + +# === fused_recurrent_gated_delta_rule_update_fwd_kernel === +@triton.jit(do_not_specialize=["T"]) +def fused_recurrent_gated_delta_rule_update_fwd_kernel( + q, + k, + v, + g, + beta, + o, + h0_source, + h0_indices, + cu_seqlens, + scale, + intermediate_states_buffer, + intermediate_state_indices, + cache_steps, + retrieve_parent_token_ptr, + stride_retrieve_parent_token_seq: tl.constexpr, + stride_retrieve_parent_token_token: tl.constexpr, + T, + NP2_T: tl.constexpr, + B: tl.constexpr, + H: tl.constexpr, + HV: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + USE_INITIAL_STATE: tl.constexpr, # whether to use initial state + IS_BETA_HEADWISE: tl.constexpr, # whether beta is headwise vector or scalar, + USE_QK_L2NORM_IN_KERNEL: tl.constexpr, + IS_VARLEN: tl.constexpr, + DISABLE_STATE_UPDATE: tl.constexpr, # whether to disable final state update + DISABLE_OUTPUT_CALCULATION: tl.constexpr, # whether to disable output calculation + CACHE_INTERMEDIATE_STATES: tl.constexpr, + HAS_EAGLE_TREE_CUSTOM_ATTN_MASK: tl.constexpr, +): + i_k, i_v, i_nh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_n, i_hv = i_nh // HV, i_nh % HV + i_h = i_hv // (HV // H) + if IS_VARLEN: + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load( + cu_seqlens + i_n + 1 + ).to(tl.int64) + all = T + T = eos - bos + else: + bos, eos = i_n * T, i_n * T + T + all = B * T + o_k = i_k * BK + tl.arange(0, BK) + o_v = i_v * BV + tl.arange(0, BV) + + p_q = q + (bos * H + i_h) * K + o_k + p_k = k + (bos * H + i_h) * K + o_k + p_v = v + (bos * HV + i_hv) * V + o_v + if IS_BETA_HEADWISE: + p_beta = beta + (bos * HV + i_hv) * V + o_v + else: + p_beta = beta + bos * HV + i_hv + p_g = g + bos * HV + i_hv + p_o = o + ((i_k * all + bos) * HV + i_hv) * V + o_v + + if HAS_EAGLE_TREE_CUSTOM_ATTN_MASK: + token_indices = tl.arange(0, NP2_T) + mask_retrieve = token_indices < T + retrieve_parent_token_base = ( + retrieve_parent_token_ptr + + (i_n * stride_retrieve_parent_token_seq) + + token_indices * stride_retrieve_parent_token_token + ) + parent_idx_tokens = tl.load(retrieve_parent_token_base, mask_retrieve) + + mask_k = o_k < K + mask_v = o_v < V + mask_h = mask_k[:, None] & mask_v[None, :] + + b_h = tl.zeros([BK, BV], dtype=tl.float32) + if USE_INITIAL_STATE: + idx = tl.load(h0_indices + i_n) + # Add bounds checking for idx + if idx >= 0: # Assuming negative indices are invalid + p_h0 = ( + h0_source + + idx * HV * K * V + + i_hv * K * V + + o_k[:, None] * V + + o_v[None, :] + ) + b_h += tl.load(p_h0, mask=mask_h, other=0).to(tl.float32) + + # Prepare intermediate state cache variables if enabled + cache_idx = -1 + if CACHE_INTERMEDIATE_STATES: + cache_idx = tl.load(intermediate_state_indices + i_n) + + step_idx = 0 + for _ in range(0, T): + if HAS_EAGLE_TREE_CUSTOM_ATTN_MASK: + # step_idx = 0 should use the b_h from USE_INITIAL_STATE + if step_idx != 0 and cache_idx >= 0: + # when calculating current step's attention, load the state from the parent token + parent_step_idx = tl.sum( + tl.where(token_indices == step_idx, parent_idx_tokens, 0) + ) + step_offset = parent_step_idx * HV * K * V + cache_ptr = ( + intermediate_states_buffer + + cache_idx * cache_steps * HV * K * V + + step_offset + + i_hv * K * V + + o_k[:, None] * V + + o_v[None, :] + ) + b_h = tl.load(cache_ptr, mask=mask_h, other=0).to(tl.float32) + + b_q = tl.load(p_q, mask=mask_k, other=0).to(tl.float32) + b_k = tl.load(p_k, mask=mask_k, other=0).to(tl.float32) + b_v = tl.load(p_v, mask=mask_v, other=0).to(tl.float32) + b_g = tl.load(p_g).to(tl.float32) + + if USE_QK_L2NORM_IN_KERNEL: + b_q = b_q / (tl.sqrt(tl.sum(b_q * b_q) + 1e-6)) + b_k = b_k / (tl.sqrt(tl.sum(b_k * b_k) + 1e-6)) + b_q = b_q * scale + # [BK, BV] + b_h *= exp(b_g) + # [BV] + b_v -= tl.sum(b_h * b_k[:, None], 0) + if IS_BETA_HEADWISE: + b_beta = tl.load(p_beta, mask=mask_v, other=0).to(tl.float32) + else: + b_beta = tl.load(p_beta).to(tl.float32) + b_v *= b_beta + # [BK, BV] + b_h += b_k[:, None] * b_v[None, :] + # [BV] + if not DISABLE_OUTPUT_CALCULATION: + b_o = tl.sum(b_h * b_q[:, None], 0) + # core attn output + tl.store(p_o, b_o.to(p_o.dtype.element_ty), mask=mask_v) + + # store intermediate states if enabled + if CACHE_INTERMEDIATE_STATES: + if cache_idx >= 0: + # Compute cache pointer for this step + step_offset = step_idx * HV * K * V + cache_ptr = ( + intermediate_states_buffer + + cache_idx * cache_steps * HV * K * V + + step_offset + + i_hv * K * V + + o_k[:, None] * V + + o_v[None, :] + ) + tl.store(cache_ptr, b_h.to(cache_ptr.dtype.element_ty), mask=mask_h) + + step_idx += 1 + + p_q += H * K + p_k += H * K + p_o += HV * V + p_v += HV * V + p_g += HV + p_beta += HV * (V if IS_BETA_HEADWISE else 1) + + # Store final state back to h0_source with bounds checking + # ssm states + if not DISABLE_STATE_UPDATE: + idx = tl.load(h0_indices + i_n) + if idx >= 0: # Add bounds checking + p_h0 = ( + h0_source + + idx * HV * K * V + + i_hv * K * V + + o_k[:, None] * V + + o_v[None, :] + ) + tl.store(p_h0, b_h.to(p_h0.dtype.element_ty), mask=mask_h) + diff --git a/benchmarks/TritonNPUBench/sglang/fused_rmsnorm_kernel/fused_rmsnorm_kernel.py b/benchmarks/TritonNPUBench/sglang/fused_rmsnorm_kernel/fused_rmsnorm_kernel.py new file mode 100644 index 00000000..263779af --- /dev/null +++ b/benchmarks/TritonNPUBench/sglang/fused_rmsnorm_kernel/fused_rmsnorm_kernel.py @@ -0,0 +1,40 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/sglang_operator_cases/newtest_cases/test_fused_rmsnorm_kernel.py +# Main kernel: fused_rmsnorm_kernel +# PT file: test_fused_rmsnorm_kernel_v2.pt + +import triton +import triton.language as tl + + +# === fused_rmsnorm_kernel === +@triton.jit +def fused_rmsnorm_kernel( + output_ptr, + activ_ptr, + weight_ptr, + eps: tl.constexpr, + hidden_dim: tl.constexpr, + BLOCK_SIZE: tl.constexpr, +): + pid = tl.program_id(axis=0).to(tl.int64) + input_start = pid * hidden_dim + + offsets = tl.arange(0, BLOCK_SIZE) + mask = offsets < hidden_dim + + a_ = tl.load(activ_ptr + input_start + offsets, mask=mask, other=0.0) + a = a_.to(tl.float32) + rms = tl.sqrt(tl.sum(a * a, axis=0) / hidden_dim + eps) + + w1_ = tl.load(weight_ptr + offsets, mask=mask, other=0.0) + w1 = w1_.to(tl.float32) + + a_rms = a / rms * w1 + + tl.store( + output_ptr + input_start + offsets, + a_rms, # implicitly casts to output dtype here + mask=mask, + ) + diff --git a/benchmarks/TritonNPUBench/sglang/fused_sigmoid_gating_delta_rule_update_kernel/fused_sigmoid_gating_delta_rule_update_kernel.py b/benchmarks/TritonNPUBench/sglang/fused_sigmoid_gating_delta_rule_update_kernel/fused_sigmoid_gating_delta_rule_update_kernel.py new file mode 100644 index 00000000..a150773c --- /dev/null +++ b/benchmarks/TritonNPUBench/sglang/fused_sigmoid_gating_delta_rule_update_kernel/fused_sigmoid_gating_delta_rule_update_kernel.py @@ -0,0 +1,158 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/sglang_operator_cases/newtest_cases/test_fused_sigmoid_gating_delta_rule_update_kernel.py +# Main kernel: fused_sigmoid_gating_delta_rule_update_kernel +# PT file: test_fused_sigmoid_gating_delta_rule_update_kernel_v2.pt + +import triton +import triton.language as tl + + +# === fused_sigmoid_gating_delta_rule_update_kernel === +@triton.jit(do_not_specialize=["T"]) +def fused_sigmoid_gating_delta_rule_update_kernel( + A_log, + a, + dt_bias, + softplus_beta, + softplus_threshold, + q, + k, + v, + b, + o, + h0_source, + h0_indices, + cu_seqlens, + scale, + T, + B: tl.constexpr, + H: tl.constexpr, + HV: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + USE_INITIAL_STATE: tl.constexpr, + USE_QK_L2NORM_IN_KERNEL: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + """ + Fused kernel that combines sigmoid gating computation with recurrent delta rule update. + """ + i_k, i_v, i_nh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_n, i_hv = i_nh // HV, i_nh % HV + i_h = i_hv // (HV // H) + + if IS_VARLEN: + bos, eos = ( + tl.load(cu_seqlens + i_n).to(tl.int64), + tl.load(cu_seqlens + i_n + 1).to(tl.int64), + ) + all = T + T = eos - bos + else: + bos, eos = i_n * T, i_n * T + T + all = B * T + + o_k = i_k * BK + tl.arange(0, BK) + o_v = i_v * BV + tl.arange(0, BV) + + p_q = q + (bos * H + i_h) * K + o_k + p_k = k + (bos * H + i_h) * K + o_k + p_v = v + (bos * HV + i_hv) * V + o_v + p_b = b + bos * HV + i_hv + p_o = o + ((i_k * all + bos) * HV + i_hv) * V + o_v + + # Gating computation pointers + p_A_log = A_log + i_hv + p_a = a + bos * HV + i_hv + p_dt_bias = dt_bias + i_hv + + mask_k = o_k < K + mask_v = o_v < V + mask_h = mask_k[:, None] & mask_v[None, :] + + b_h = tl.zeros([BK, BV], dtype=tl.float32) + if USE_INITIAL_STATE: + idx = tl.load(h0_indices + i_n) + if idx >= 0: + p_h0 = ( + h0_source + + idx * HV * K * V + + i_hv * K * V + + o_k[:, None] * V + + o_v[None, :] + ) + b_h += tl.load(p_h0, mask=mask_h, other=0).to(tl.float32) + + for _ in range(0, T): + # Load inputs + b_q = tl.load(p_q, mask=mask_k, other=0).to(tl.float32) + b_k = tl.load(p_k, mask=mask_k, other=0).to(tl.float32) + b_v = tl.load(p_v, mask=mask_v, other=0).to(tl.float32) + b_b = tl.load(p_b).to(tl.float32) + + # Compute sigmoid gating + # Load gating parameters + b_A_log = tl.load(p_A_log).to(tl.float32) + b_a = tl.load(p_a).to(tl.float32) + b_dt_bias = tl.load(p_dt_bias).to(tl.float32) + + # Compute g = -exp(A_log) * softplus(a + dt_bias) + x = b_a + b_dt_bias + beta_x = softplus_beta * x + # Apply softplus with numerical stability + softplus_x = tl.where( + beta_x <= softplus_threshold, + (1.0 / softplus_beta) * tl.log(1.0 + tl.exp(beta_x)), + x, + ) + b_g = -tl.exp(b_A_log) * softplus_x + + # Compute beta = sigmoid(b) + b_beta = 1.0 / (1.0 + tl.exp(-b_b)) + + # Apply L2 normalization if enabled + if USE_QK_L2NORM_IN_KERNEL: + b_q = b_q / (tl.sqrt(tl.sum(b_q * b_q) + 1e-6)) + b_k = b_k / (tl.sqrt(tl.sum(b_k * b_k) + 1e-6)) + + b_q = b_q * scale + + # Apply gating to hidden state: h *= exp(g) + b_h *= tl.exp(b_g) + + # Delta rule: v -= sum(h * k, dim=0) + b_v -= tl.sum(b_h * b_k[:, None], 0) + + # Apply beta gating: v *= beta + b_v *= b_beta + + # Update hidden state: h += k[:, None] * v[None, :] + b_h += b_k[:, None] * b_v[None, :] + + # Compute output: o = sum(h * q, dim=0) + b_o = tl.sum(b_h * b_q[:, None], 0) + tl.store(p_o, b_o.to(p_o.dtype.element_ty), mask=mask_v) + + # Update pointers for next timestep + p_q += H * K + p_k += H * K + p_o += HV * V + p_v += HV * V + p_b += HV + p_a += HV + + # Store final state back to h0_source with bounds checking + if USE_INITIAL_STATE: + idx = tl.load(h0_indices + i_n) + if idx >= 0: + p_h0 = ( + h0_source + + idx * HV * K * V + + i_hv * K * V + + o_k[:, None] * V + + o_v[None, :] + ) + tl.store(p_h0, b_h.to(p_h0.dtype.element_ty), mask=mask_h) + diff --git a/benchmarks/TritonNPUBench/sglang/fused_softcap_kernel/fused_softcap_kernel.py b/benchmarks/TritonNPUBench/sglang/fused_softcap_kernel/fused_softcap_kernel.py new file mode 100644 index 00000000..92cc2bb1 --- /dev/null +++ b/benchmarks/TritonNPUBench/sglang/fused_softcap_kernel/fused_softcap_kernel.py @@ -0,0 +1,31 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/sglang_operator_cases/newtest_cases/test_fused_softcap_kernel.py +# Main kernel: fused_softcap_kernel +# PT file: test_fused_softcap_kernel_v2.pt + +import triton +import triton.language as tl + + +# === fused_softcap_kernel === +@triton.jit +def fused_softcap_kernel( + output_ptr, + input_ptr, + n_ele, + softcap_const: tl.constexpr, + BLOCK_SIZE: tl.constexpr, +): + pid = tl.program_id(axis=0) + block_start = pid * BLOCK_SIZE + offsets = block_start + tl.arange(0, BLOCK_SIZE) + mask = offsets < n_ele + x = tl.load(input_ptr + offsets, mask=mask) + fx = x.to(tl.float32) + fxs = fx / softcap_const + exped = tl.exp(2 * fxs) + top = exped - 1 + bottom = exped + 1 + output = top / bottom * softcap_const + tl.store(output_ptr + offsets, output, mask=mask) + diff --git a/benchmarks/TritonNPUBench/sglang/gelu_and_mul_kernel/gelu_and_mul_kernel.py b/benchmarks/TritonNPUBench/sglang/gelu_and_mul_kernel/gelu_and_mul_kernel.py new file mode 100644 index 00000000..0b8af6d2 --- /dev/null +++ b/benchmarks/TritonNPUBench/sglang/gelu_and_mul_kernel/gelu_and_mul_kernel.py @@ -0,0 +1,47 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/sglang_operator_cases/newtest_cases/test_gelu_and_mul_kernel.py +# Main kernel: gelu_and_mul_kernel +# PT file: test_gelu_and_mul_kernel_v2.pt + +import triton +import triton.language as tl + + +# === gelu_and_mul_kernel === +@triton.jit +def gelu_and_mul_kernel( + out_hidden_states_ptr, # (bs, hidden_dim) + out_scales_ptr, # (bs,) + hidden_states_ptr, # (bs, hidden_dim * 2) + quant_max: tl.constexpr, + static_scale: tl.constexpr, + hidden_dim: tl.constexpr, # the output hidden_dim + BLOCK_SIZE: tl.constexpr, +): + pid = tl.program_id(axis=0) + + input_start = pid * hidden_dim * 2 + output_start = pid * hidden_dim + + input1_offs = tl.arange(0, BLOCK_SIZE) + mask = tl.arange(0, BLOCK_SIZE) < hidden_dim # shared for input1, input3, output + input3_offs = hidden_dim + tl.arange(0, BLOCK_SIZE) + output_offs = tl.arange(0, BLOCK_SIZE) + + x1 = tl.load( + hidden_states_ptr + input_start + input1_offs, mask=mask, other=0.0 + ).to(tl.float32) + x3 = tl.load( + hidden_states_ptr + input_start + input3_offs, mask=mask, other=0.0 + ).to(tl.float32) + + # gelu + # cast down before mul to better match training? + gelu_x1 = 0.5 * (1.0 + tl.erf(x1 * 0.7071067811865475)) * x1 + out = x3 * gelu_x1.to(hidden_states_ptr.dtype.element_ty) + + if quant_max is not None: + raise NotImplementedError() + + tl.store(out_hidden_states_ptr + output_start + output_offs, out, mask=mask) + diff --git a/benchmarks/TritonNPUBench/sglang/get_last_loc_kernel/get_last_loc_kernel.py b/benchmarks/TritonNPUBench/sglang/get_last_loc_kernel/get_last_loc_kernel.py new file mode 100644 index 00000000..f1431425 --- /dev/null +++ b/benchmarks/TritonNPUBench/sglang/get_last_loc_kernel/get_last_loc_kernel.py @@ -0,0 +1,33 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/sglang_operator_cases/newtest_cases/test_get_last_loc_kernel.py +# Main kernel: get_last_loc_kernel +# PT file: test_get_last_loc_kernel_v2.pt + +import triton +import triton.language as tl + + +# === get_last_loc_kernel === +@triton.jit +def get_last_loc_kernel( + req_to_token, + req_pool_indices_tensor, + prefix_lens_tensor, + result, + num_tokens, + req_to_token_stride, + BLOCK_SIZE: tl.constexpr, +): + pid = tl.program_id(0) + offset = tl.arange(0, BLOCK_SIZE) + pid * BLOCK_SIZE + mask = offset < num_tokens + + prefix_lens = tl.load(prefix_lens_tensor + offset, mask=mask, other=0) + req_pool_indices = tl.load(req_pool_indices_tensor + offset, mask=mask, other=0) + + token_mask = prefix_lens > 0 + token_index = req_pool_indices * req_to_token_stride + (prefix_lens - 1) + tokens = tl.load(req_to_token + token_index, mask=token_mask, other=-1) + + tl.store(result + offset, tokens, mask=mask) + diff --git a/benchmarks/TritonNPUBench/sglang/get_num_kv_splits_triton/get_num_kv_splits_triton.py b/benchmarks/TritonNPUBench/sglang/get_num_kv_splits_triton/get_num_kv_splits_triton.py new file mode 100644 index 00000000..ab3f32c5 --- /dev/null +++ b/benchmarks/TritonNPUBench/sglang/get_num_kv_splits_triton/get_num_kv_splits_triton.py @@ -0,0 +1,61 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/sglang_operator_cases/newtest_cases/test_get_num_kv_splits_triton.py +# Main kernel: get_num_kv_splits_triton +# PT file: test_get_num_kv_splits_triton_v2.pt + +import triton +import triton.language as tl + + +# === get_num_kv_splits_triton === +@triton.jit +def get_num_kv_splits_triton( + num_kv_splits_ptr, + seq_lens_ptr, + num_seq, + num_group, + num_head, + num_kv_head, + max_kv_splits, + device_core_count, + MAX_NUM_SEQ: tl.constexpr, +): + # TODO: this method is tunable, we need more online serving data to tune it + offs_seq = tl.arange(0, MAX_NUM_SEQ) + mask_seq = offs_seq < num_seq + + seq_lens = tl.load(seq_lens_ptr + offs_seq, mask=mask_seq, other=0) + max_seq_len = tl.max(seq_lens) + seq_lens = tl.load(seq_lens_ptr + offs_seq, mask=mask_seq, other=max_seq_len) + min_seq_len = tl.min(seq_lens) + if max_seq_len * 8 < min_seq_len * 10: + min_seq_len = max_seq_len + max_kv_splits_1 = tl.minimum(tl.cdiv(max_seq_len, min_seq_len), max_kv_splits) + kv_chunk_size_1 = tl.cdiv(max_seq_len, max_kv_splits_1) + + # NOTE: this is a hack to let num_kv_split grows up with seqlen gradually + ext_seq_len = tl.cast(max_seq_len, tl.float32) / 64.0 + ext_device_core_count = tl.cast( + device_core_count * tl.maximum(tl.log2(ext_seq_len), 1.0), tl.int32 + ) + block_h, num_kv_group = 16, num_head // num_kv_head + if num_kv_group == 1: + token_grid = num_seq * num_group * num_head + else: + # from triton_ops/decode_attention.py:_decode_grouped_att_m_fwd + block_h = tl.minimum(block_h, num_kv_group) + token_grid = num_seq * num_group * tl.cdiv(num_head, block_h) + max_kv_splits_2 = tl.minimum( + tl.cdiv(ext_device_core_count, token_grid), max_kv_splits + ) + kv_chunk_size_2 = tl.cdiv(max_seq_len, max_kv_splits_2) + + num_kv_splits = tl.maximum( + tl.cdiv(seq_lens, kv_chunk_size_1), tl.cdiv(seq_lens, kv_chunk_size_2) + ) + + offs_token = offs_seq * num_group + mask_token = offs_token < num_seq * num_group + for i in range(0, num_group): + tl.store(num_kv_splits_ptr + i + offs_token, num_kv_splits, mask=mask_token) + diff --git a/benchmarks/TritonNPUBench/sglang/hash_tiles32_kernel_blocked/hash_tiles32_kernel_blocked.py b/benchmarks/TritonNPUBench/sglang/hash_tiles32_kernel_blocked/hash_tiles32_kernel_blocked.py new file mode 100644 index 00000000..3f538d41 --- /dev/null +++ b/benchmarks/TritonNPUBench/sglang/hash_tiles32_kernel_blocked/hash_tiles32_kernel_blocked.py @@ -0,0 +1,92 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/sglang_operator_cases/newtest_cases/test_hash_tiles32_kernel_blocked.py +# Main kernel: hash_tiles32_kernel_blocked +# PT file: test_hash_tiles32_kernel_blocked_v2.pt + +import triton +import triton.language as tl + + +# === hash_tiles32_kernel_blocked === +@triton.jit +def hash_tiles32_kernel_blocked( + in_ptr, + out_ptr, + n_u32, + seed1, + seed2, + FM_C1: tl.constexpr, + FM_C2: tl.constexpr, + POS_A: tl.constexpr, + POS_B: tl.constexpr, + TILE: tl.constexpr, + BLOCK: tl.constexpr, + USE_CG: tl.constexpr, +): + pid = tl.program_id(axis=0) + base = pid * TILE + + s1 = tl.full((), seed1, tl.uint32) + s2 = tl.full((), seed2, tl.uint32) + posA = tl.full((), POS_A, tl.uint32) + posB = tl.full((), POS_B, tl.uint32) + + h1 = tl.zeros((), dtype=tl.uint32) + h2 = tl.zeros((), dtype=tl.uint32) + + for off in tl.static_range(0, TILE, BLOCK): + idx = base + off + tl.arange(0, BLOCK) + m = idx < n_u32 + + if USE_CG: + v = tl.load(in_ptr + idx, mask=m, other=0, cache_modifier=".cg") + else: + v = tl.load(in_ptr + idx, mask=m, other=0) + v = v.to(tl.uint32) + + iu = idx.to(tl.uint32) + p1 = (iu * posA + s1) ^ _rotl32(iu, 15) + p2 = (iu * posB + s2) ^ _rotl32(iu, 13) + + k1 = _fmix32(v ^ p1, C1=FM_C1, C2=FM_C2) + k2 = _fmix32(v ^ p2, C1=FM_C1, C2=FM_C2) + + zero32 = tl.zeros_like(k1) + k1 = tl.where(m, k1, zero32) + k2 = tl.where(m, k2, zero32) + + h1 += tl.sum(k1, axis=0).to(tl.uint32) + h2 += tl.sum(k2, axis=0).to(tl.uint32) + + nbytes = tl.full((), n_u32 * 4, tl.uint32) + h1 ^= nbytes + h2 ^= nbytes + h1 = _fmix32(h1, C1=FM_C1, C2=FM_C2) + h2 = ( + _fmix32(h2, C1=FMIX32_C1, C2=FMIX32_C2) + if False + else _fmix32(h2, C1=FM_C1, C2=FM_C2) + ) + + out = (h1.to(tl.uint64) << 32) | h2.to(tl.uint64) + tl.store(out_ptr + pid, out) + + +# === _rotl32 === +@triton.jit +def _rotl32(x, r: tl.constexpr): + return (x << r) | (x >> (32 - r)) + + +# === _fmix32 === +@triton.jit +def _fmix32(x, C1: tl.constexpr, C2: tl.constexpr): + c1 = tl.full((), C1, tl.uint32) + c2 = tl.full((), C2, tl.uint32) + x ^= x >> 16 + x = x * c1 + x ^= x >> 13 + x = x * c2 + x ^= x >> 16 + return x + diff --git a/benchmarks/TritonNPUBench/sglang/l2norm_fwd_kernel/l2norm_fwd_kernel.py b/benchmarks/TritonNPUBench/sglang/l2norm_fwd_kernel/l2norm_fwd_kernel.py new file mode 100644 index 00000000..4d8440b6 --- /dev/null +++ b/benchmarks/TritonNPUBench/sglang/l2norm_fwd_kernel/l2norm_fwd_kernel.py @@ -0,0 +1,29 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/sglang_operator_cases/newtest_cases/test_l2norm_fwd_kernel.py +# Main kernel: l2norm_fwd_kernel +# PT file: test_l2norm_fwd_kernel_v2.pt + +import triton +import triton.language as tl + + +# === l2norm_fwd_kernel === +@triton.jit +def l2norm_fwd_kernel( + x, + y, + eps, + NB: tl.constexpr, + T: tl.constexpr, + D: tl.constexpr, + BT: tl.constexpr, + BD: tl.constexpr, +): + i_t = tl.program_id(0) + p_x = tl.make_block_ptr(x, (T, D), (D, 1), (i_t * BT, 0), (BT, BD), (1, 0)) + b_x = tl.load(p_x, boundary_check=(0, 1)).to(tl.float32) + b_var = tl.sum(b_x * b_x, axis=1) + b_y = b_x / tl.sqrt(b_var + eps)[:, None] + p_y = tl.make_block_ptr(y, (T, D), (D, 1), (i_t * BT, 0), (BT, BD), (1, 0)) + tl.store(p_y, b_y.to(p_y.dtype.element_ty), boundary_check=(0, 1)) + diff --git a/benchmarks/TritonNPUBench/sglang/l2norm_fwd_kernel1/l2norm_fwd_kernel1.py b/benchmarks/TritonNPUBench/sglang/l2norm_fwd_kernel1/l2norm_fwd_kernel1.py new file mode 100644 index 00000000..5ade7943 --- /dev/null +++ b/benchmarks/TritonNPUBench/sglang/l2norm_fwd_kernel1/l2norm_fwd_kernel1.py @@ -0,0 +1,32 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/sglang_operator_cases/newtest_cases/test_l2norm_fwd_kernel1.py +# Main kernel: l2norm_fwd_kernel1 +# PT file: test_l2norm_fwd_kernel1_v2.pt + +import triton +import triton.language as tl + + +# === l2norm_fwd_kernel1 === +@triton.jit +def l2norm_fwd_kernel1( + x, + y, + D, + BD: tl.constexpr, + eps, +): + i_t = tl.program_id(0) + x += i_t * D + y += i_t * D + # Compute mean and variance + cols = tl.arange(0, BD) + mask = cols < D + b_x = tl.load(x + cols, mask=mask, other=0.0).to(tl.float32) + b_var = tl.sum(b_x * b_x, axis=0) + b_rstd = 1 / tl.sqrt(b_var + eps) + # tl.store(Rstd + i_t, rstd) + # Normalize and apply linear transformation + b_y = b_x * b_rstd + tl.store(y + cols, b_y, mask=mask) + diff --git a/benchmarks/TritonNPUBench/sglang/layer_norm_gated_fwd_kernel/layer_norm_gated_fwd_kernel.py b/benchmarks/TritonNPUBench/sglang/layer_norm_gated_fwd_kernel/layer_norm_gated_fwd_kernel.py new file mode 100644 index 00000000..5d4da689 --- /dev/null +++ b/benchmarks/TritonNPUBench/sglang/layer_norm_gated_fwd_kernel/layer_norm_gated_fwd_kernel.py @@ -0,0 +1,89 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/sglang_operator_cases/newtest_cases/test_layer_norm_gated_fwd_kernel.py +# Main kernel: layer_norm_gated_fwd_kernel +# PT file: test_layer_norm_gated_fwd_kernel_v2.pt + +import triton +import triton.language as tl + + +# === layer_norm_gated_fwd_kernel === +@triton.jit +def layer_norm_gated_fwd_kernel( + x, # pointer to the input + g, # pointer to the gate + y, # pointer to the output + w, # pointer to the weights + b, # pointer to the biases + residual, # pointer to the residual + residual_out, # pointer to the residual + mean, # pointer to the mean + rstd, # pointer to the 1/std + eps, # epsilon to avoid division by zero + T, # number of rows in x + D: tl.constexpr, # number of columns in x + BT: tl.constexpr, + BD: tl.constexpr, + ACTIVATION: tl.constexpr, + IS_RMS_NORM: tl.constexpr, + STORE_RESIDUAL_OUT: tl.constexpr, + HAS_RESIDUAL: tl.constexpr, + HAS_WEIGHT: tl.constexpr, + HAS_BIAS: tl.constexpr, +): + i_t = tl.program_id(0) + + o_d = tl.arange(0, BD) + m_d = o_d < D + + p_x = tl.make_block_ptr(x, (T, D), (D, 1), (i_t * BT, 0), (BT, BD), (1, 0)) + b_x = tl.load(p_x, boundary_check=(0, 1)).to(tl.float32) + if HAS_RESIDUAL: + p_res = tl.make_block_ptr( + residual, (T, D), (D, 1), (i_t * BT, 0), (BT, BD), (1, 0) + ) + b_x += tl.load(p_res, boundary_check=(0, 1)).to(tl.float32) + if STORE_RESIDUAL_OUT: + p_res_out = tl.make_block_ptr( + residual_out, (T, D), (D, 1), (i_t * BT, 0), (BT, BD), (1, 0) + ) + tl.store(p_res_out, b_x.to(p_res_out.dtype.element_ty), boundary_check=(0, 1)) + if not IS_RMS_NORM: + b_mean = tl.sum(b_x, axis=1) / D + p_mean = tl.make_block_ptr(mean, (T,), (1,), (i_t * BT,), (BT,), (0,)) + tl.store(p_mean, b_mean.to(p_mean.dtype.element_ty), boundary_check=(0,)) + b_xbar = tl.where(m_d[None, :], b_x - b_mean[:, None], 0.0) + b_var = tl.sum(b_xbar * b_xbar, axis=1) / D + else: + b_xbar = tl.where(m_d[None, :], b_x, 0.0) + b_var = tl.sum(b_xbar * b_xbar, axis=1) / D + b_rstd = 1 / tl.sqrt(b_var + eps) + + p_rstd = tl.make_block_ptr(rstd, (T,), (1,), (i_t * BT,), (BT,), (0,)) + tl.store(p_rstd, b_rstd.to(p_rstd.dtype.element_ty), boundary_check=(0,)) + + if HAS_WEIGHT: + b_w = tl.load(w + o_d, mask=m_d).to(tl.float32) + if HAS_BIAS: + b_b = tl.load(b + o_d, mask=m_d).to(tl.float32) + b_x_hat = ( + (b_x - b_mean[:, None]) * b_rstd[:, None] + if not IS_RMS_NORM + else b_x * b_rstd[:, None] + ) + b_y = b_x_hat * b_w[None, :] if HAS_WEIGHT else b_x_hat + if HAS_BIAS: + b_y = b_y + b_b[None, :] + + # swish/sigmoid output gate + p_g = tl.make_block_ptr(g, (T, D), (D, 1), (i_t * BT, 0), (BT, BD), (1, 0)) + b_g = tl.load(p_g, boundary_check=(0, 1)).to(tl.float32) + if ACTIVATION == "swish" or ACTIVATION == "silu": + b_y = b_y * b_g * tl.sigmoid(b_g) + elif ACTIVATION == "sigmoid": + b_y = b_y * tl.sigmoid(b_g) + + # Write output + p_y = tl.make_block_ptr(y, (T, D), (D, 1), (i_t * BT, 0), (BT, BD), (1, 0)) + tl.store(p_y, b_y.to(p_y.dtype.element_ty), boundary_check=(0, 1)) + diff --git a/benchmarks/TritonNPUBench/sglang/layer_norm_gated_fwd_kernel1/layer_norm_gated_fwd_kernel1.py b/benchmarks/TritonNPUBench/sglang/layer_norm_gated_fwd_kernel1/layer_norm_gated_fwd_kernel1.py new file mode 100644 index 00000000..7b0b3660 --- /dev/null +++ b/benchmarks/TritonNPUBench/sglang/layer_norm_gated_fwd_kernel1/layer_norm_gated_fwd_kernel1.py @@ -0,0 +1,77 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/sglang_operator_cases/newtest_cases/test_layer_norm_gated_fwd_kernel1.py +# Main kernel: layer_norm_gated_fwd_kernel1 +# PT file: test_layer_norm_gated_fwd_kernel1_v2.pt + +import triton +import triton.language as tl + + +# === layer_norm_gated_fwd_kernel1 === +@triton.jit +def layer_norm_gated_fwd_kernel1( + x, # pointer to the input + g, # pointer to the gate + y, # pointer to the output + w, # pointer to the weights + b, # pointer to the biases + residual, # pointer to the residual + residual_out, # pointer to the residual + mean, # pointer to the mean + rstd, # pointer to the 1/std + eps, # epsilon to avoid division by zero + D: tl.constexpr, # number of columns in x + BD: tl.constexpr, + ACTIVATION: tl.constexpr, + IS_RMS_NORM: tl.constexpr, + STORE_RESIDUAL_OUT: tl.constexpr, + HAS_RESIDUAL: tl.constexpr, + HAS_WEIGHT: tl.constexpr, + HAS_BIAS: tl.constexpr, +): + i_t = tl.program_id(0) + x += i_t * D + y += i_t * D + g += i_t * D + if HAS_RESIDUAL: + residual += i_t * D + if STORE_RESIDUAL_OUT: + residual_out += i_t * D + + o_d = tl.arange(0, BD) + m_d = o_d < D + b_x = tl.load(x + o_d, mask=m_d, other=0.0).to(tl.float32) + if HAS_RESIDUAL: + b_x += tl.load(residual + o_d, mask=m_d, other=0.0).to(tl.float32) + if STORE_RESIDUAL_OUT: + tl.store(residual_out + o_d, b_x, mask=m_d) + if not IS_RMS_NORM: + b_mean = tl.sum(b_x, axis=0) / D + tl.store(mean + i_t, b_mean) + b_xbar = tl.where(m_d, b_x - b_mean, 0.0) + b_var = tl.sum(b_xbar * b_xbar, axis=0) / D + else: + b_xbar = tl.where(m_d, b_x, 0.0) + b_var = tl.sum(b_xbar * b_xbar, axis=0) / D + b_rstd = 1 / tl.sqrt(b_var + eps) + tl.store(rstd + i_t, b_rstd) + + if HAS_WEIGHT: + b_w = tl.load(w + o_d, mask=m_d).to(tl.float32) + if HAS_BIAS: + b_b = tl.load(b + o_d, mask=m_d).to(tl.float32) + b_x_hat = (b_x - b_mean) * b_rstd if not IS_RMS_NORM else b_x * b_rstd + b_y = b_x_hat * b_w if HAS_WEIGHT else b_x_hat + if HAS_BIAS: + b_y = b_y + b_b + + # swish/sigmoid output gate + b_g = tl.load(g + o_d, mask=m_d, other=0.0).to(tl.float32) + if ACTIVATION == "swish" or ACTIVATION == "silu": + b_y = b_y * b_g * tl.sigmoid(b_g) + elif ACTIVATION == "sigmoid": + b_y = b_y * tl.sigmoid(b_g) + + # Write output + tl.store(y + o_d, b_y, mask=m_d) + diff --git a/benchmarks/TritonNPUBench/sglang/memcpy_triton_kernel/memcpy_triton_kernel.py b/benchmarks/TritonNPUBench/sglang/memcpy_triton_kernel/memcpy_triton_kernel.py new file mode 100644 index 00000000..1122c655 --- /dev/null +++ b/benchmarks/TritonNPUBench/sglang/memcpy_triton_kernel/memcpy_triton_kernel.py @@ -0,0 +1,35 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/sglang_operator_cases/newtest_cases/test_memcpy_triton_kernel.py +# Main kernel: memcpy_triton_kernel +# PT file: test_memcpy_triton_kernel_v2.pt + +import triton +import triton.language as tl + + +# === memcpy_triton_kernel === +@triton.jit +def memcpy_triton_kernel( + dst_ptr, + src_ptr, + offset_ptr, + sz_ptr, + offset_src: tl.constexpr, + chunk_size, # multiplied for offset and sz + BLOCK_SIZE: tl.constexpr, +): + pid = tl.program_id(axis=0).to(tl.int64) + offset = tl.load(offset_ptr).to(tl.int64) * chunk_size + sz = tl.load(sz_ptr).to(tl.int64) * chunk_size + + start_index = pid * BLOCK_SIZE + offs = tl.arange(0, BLOCK_SIZE) + mask = start_index + offs < sz + + if offset_src: + data = tl.load(src_ptr + offset + start_index + offs, mask=mask) + tl.store(dst_ptr + start_index + offs, data, mask=mask) + else: + data = tl.load(src_ptr + start_index + offs, mask=mask) + tl.store(dst_ptr + offset + start_index + offs, data, mask=mask) + diff --git a/benchmarks/TritonNPUBench/sglang/merge_16x16_to_32x32_inverse_kernel/merge_16x16_to_32x32_inverse_kernel.py b/benchmarks/TritonNPUBench/sglang/merge_16x16_to_32x32_inverse_kernel/merge_16x16_to_32x32_inverse_kernel.py new file mode 100644 index 00000000..8aaac01a --- /dev/null +++ b/benchmarks/TritonNPUBench/sglang/merge_16x16_to_32x32_inverse_kernel/merge_16x16_to_32x32_inverse_kernel.py @@ -0,0 +1,137 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/sglang_operator_cases/newtest_cases/test_merge_16x16_to_32x32_inverse_kernel.py +# Main kernel: merge_16x16_to_32x32_inverse_kernel +# PT file: merge_16x16_to_32x32_inverse_kernel.pt + +import triton +import triton.language as tl + + +# === merge_16x16_to_32x32_inverse_kernel === +@triton.jit(do_not_specialize=["T"]) +def merge_16x16_to_32x32_inverse_kernel( + A, + Ai, + cu_seqlens, + chunk_indices, + T, + H: tl.constexpr, + BT: tl.constexpr, + USE_TMA: tl.constexpr, + IS_VARLEN: tl.constexpr, + DOT_PRECISION: tl.constexpr, +): + i_t, i_bh = tl.program_id(0), tl.program_id(1) + i_b, i_h = i_bh // H, i_bh % H + if IS_VARLEN: + i_n, i_t = ( + tl.load(chunk_indices + i_t * 2).to(tl.int32), + tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), + ) + bos, eos = ( + tl.load(cu_seqlens + i_n).to(tl.int32), + tl.load(cu_seqlens + i_n + 1).to(tl.int32), + ) + T = eos - bos + else: + bos, eos = i_b * T, i_b * T + T + + o_i = tl.arange(0, 16) + m_A = o_i[:, None] > o_i[None, :] + m_I = o_i[:, None] == o_i[None, :] + A += (bos * H + i_h) * BT + Ai += (bos * H + i_h) * BT + + if not USE_TMA: + p_A_11 = tl.make_block_ptr( + A, (T, BT), (H * BT, 1), (i_t * BT, 0), (16, 16), (1, 0) + ) + p_A_22 = tl.make_block_ptr( + A, (T, BT), (H * BT, 1), (i_t * BT + 16, 16), (16, 16), (1, 0) + ) + b_Ai_11 = tl.load(p_A_11, boundary_check=(0, 1)).to(tl.float32) + b_Ai_22 = tl.load(p_A_22, boundary_check=(0, 1)).to(tl.float32) + else: + desc = make_tensor_descriptor(A, [T, BT], [H * BT, 1], [16, 16]) + desc_o = make_tensor_descriptor(Ai, [T, BT], [H * BT, 1], [16, 16]) + b_Ai_11 = desc.load([i_t * BT + 0, 0]).to(tl.float32) + b_Ai_22 = desc.load([i_t * BT + 16, 16]).to(tl.float32) + + # [16, 16] + b_Ai_11 = -tl.where(m_A, b_Ai_11, 0) + b_Ai_22 = -tl.where(m_A, b_Ai_22, 0) + + for i in range(2, min(16, T - i_t * BT)): + b_a_11 = -tl.load(A + (i_t * BT + i) * H * BT + o_i) + b_a_11 += tl.sum(b_a_11[:, None] * b_Ai_11, 0) + b_Ai_11 = tl.where((o_i == i)[:, None], b_a_11, b_Ai_11) + for i in range(16 + 2, min(32, T - i_t * BT)): + b_a_22 = -tl.load(A + (i_t * BT + i) * H * BT + o_i + 16) + b_a_22 += tl.sum(b_a_22[:, None] * b_Ai_22, 0) + b_Ai_22 = tl.where((o_i == i - 16)[:, None], b_a_22, b_Ai_22) + + b_Ai_11 += m_I + b_Ai_22 += m_I + + if not USE_TMA: + p_A_21 = tl.make_block_ptr( + A, (T, BT), (H * BT, 1), (i_t * BT + 16, 0), (16, 16), (1, 0) + ) + b_A_21 = tl.load(p_A_21, boundary_check=(0, 1)).to(tl.float32) + else: + b_A_21 = desc.load([i_t * BT + 16, 0]).to(tl.float32) + + b_Ai_21 = -tl.dot( + tl.dot(b_Ai_22, b_A_21, input_precision=DOT_PRECISION), + b_Ai_11, + input_precision=DOT_PRECISION, + ) + + if not USE_TMA: + p_Ai_11 = tl.make_block_ptr( + Ai, (T, BT), (H * BT, 1), (i_t * BT, 0), (16, 16), (1, 0) + ) + p_Ai_21 = tl.make_block_ptr( + Ai, (T, BT), (H * BT, 1), (i_t * BT + 16, 0), (16, 16), (1, 0) + ) + p_Ai_22 = tl.make_block_ptr( + Ai, (T, BT), (H * BT, 1), (i_t * BT + 16, 16), (16, 16), (1, 0) + ) + tl.store( + p_Ai_11, + b_Ai_11.to(p_Ai_11.dtype.element_ty, fp_downcast_rounding="rtne"), + boundary_check=(0, 1), + ) + tl.store( + p_Ai_22, + b_Ai_22.to(p_Ai_22.dtype.element_ty, fp_downcast_rounding="rtne"), + boundary_check=(0, 1), + ) + tl.store( + p_Ai_21, + b_Ai_21.to(p_Ai_21.dtype.element_ty, fp_downcast_rounding="rtne"), + boundary_check=(0, 1), + ) + else: + desc_o.store( + [i_t * BT + 0, 0], b_Ai_11.to(desc_o.dtype, fp_downcast_rounding="rtne") + ) + desc_o.store( + [i_t * BT + 16, 0], b_Ai_21.to(desc_o.dtype, fp_downcast_rounding="rtne") + ) + desc_o.store( + [i_t * BT + 16, 16], b_Ai_22.to(desc_o.dtype, fp_downcast_rounding="rtne") + ) + + +# === make_tensor_descriptor === + @triton.jit + def make_tensor_descriptor( + base, + shape, + strides, + block_shape, + _builder=None, + ): + return None + diff --git a/benchmarks/TritonNPUBench/sglang/merge_16x16_to_64x64_inverse_kernel/merge_16x16_to_64x64_inverse_kernel.py b/benchmarks/TritonNPUBench/sglang/merge_16x16_to_64x64_inverse_kernel/merge_16x16_to_64x64_inverse_kernel.py new file mode 100644 index 00000000..5beb6669 --- /dev/null +++ b/benchmarks/TritonNPUBench/sglang/merge_16x16_to_64x64_inverse_kernel/merge_16x16_to_64x64_inverse_kernel.py @@ -0,0 +1,243 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/sglang_operator_cases/newtest_cases/test_merge_16x16_to_64x64_inverse_kernel.py +# Main kernel: merge_16x16_to_64x64_inverse_kernel +# PT file: test_merge_16x16_to_64x64_inverse_kernel_v2_uint2.pt + +import triton +import triton.language as tl + + +# === merge_16x16_to_64x64_inverse_kernel === +@triton.jit(do_not_specialize=["T"]) +def merge_16x16_to_64x64_inverse_kernel( + A, + Ad, + Ai, + cu_seqlens, + chunk_indices, + T, + H: tl.constexpr, + BT: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_t, i_bh = tl.program_id(0), tl.program_id(1) + i_b, i_h = i_bh // H, i_bh % H + if IS_VARLEN: + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load( + chunk_indices + i_t * 2 + 1 + ).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load( + cu_seqlens + i_n + 1 + ).to(tl.int32) + T = eos - bos + else: + bos, eos = i_b * T, i_b * T + T + + A += (bos * H + i_h) * 64 + Ad += (bos * H + i_h) * 16 + Ai += (bos * H + i_h) * 64 + + p_A_21 = tl.make_block_ptr( + A, (T, 64), (H * 64, 1), (i_t * 64 + 16, 0), (16, 16), (1, 0) + ) + p_A_32 = tl.make_block_ptr( + A, (T, 64), (H * 64, 1), (i_t * 64 + 32, 16), (16, 16), (1, 0) + ) + p_A_31 = tl.make_block_ptr( + A, (T, 64), (H * 64, 1), (i_t * 64 + 32, 0), (16, 16), (1, 0) + ) + p_A_43 = tl.make_block_ptr( + A, (T, 64), (H * 64, 1), (i_t * 64 + 48, 32), (16, 16), (1, 0) + ) + p_A_42 = tl.make_block_ptr( + A, (T, 64), (H * 64, 1), (i_t * 64 + 48, 16), (16, 16), (1, 0) + ) + p_A_41 = tl.make_block_ptr( + A, (T, 64), (H * 64, 1), (i_t * 64 + 48, 0), (16, 16), (1, 0) + ) + p_Ad_11 = tl.make_block_ptr( + Ad, (T, 16), (H * 16, 1), (i_t * 64, 0), (16, 16), (1, 0) + ) + p_Ad_22 = tl.make_block_ptr( + Ad, (T, 16), (H * 16, 1), (i_t * 64 + 16, 0), (16, 16), (1, 0) + ) + p_Ad_33 = tl.make_block_ptr( + Ad, (T, 16), (H * 16, 1), (i_t * 64 + 32, 0), (16, 16), (1, 0) + ) + p_Ad_44 = tl.make_block_ptr( + Ad, (T, 16), (H * 16, 1), (i_t * 64 + 48, 0), (16, 16), (1, 0) + ) + + A_21 = tl.load(p_A_21, boundary_check=(0, 1)).to(tl.float32) + A_32 = tl.load(p_A_32, boundary_check=(0, 1)).to(tl.float32) + A_31 = tl.load(p_A_31, boundary_check=(0, 1)).to(tl.float32) + A_43 = tl.load(p_A_43, boundary_check=(0, 1)).to(tl.float32) + A_42 = tl.load(p_A_42, boundary_check=(0, 1)).to(tl.float32) + A_41 = tl.load(p_A_41, boundary_check=(0, 1)).to(tl.float32) + + Ai_11 = tl.load(p_Ad_11, boundary_check=(0, 1)).to(tl.float32) + Ai_22 = tl.load(p_Ad_22, boundary_check=(0, 1)).to(tl.float32) + Ai_33 = tl.load(p_Ad_33, boundary_check=(0, 1)).to(tl.float32) + Ai_44 = tl.load(p_Ad_44, boundary_check=(0, 1)).to(tl.float32) + + Ai_21 = -tl.dot( + tl.dot(Ai_22, A_21, input_precision="ieee"), Ai_11, input_precision="ieee" + ) + Ai_32 = -tl.dot( + tl.dot(Ai_33, A_32, input_precision="ieee"), Ai_22, input_precision="ieee" + ) + Ai_43 = -tl.dot( + tl.dot(Ai_44, A_43, input_precision="ieee"), Ai_33, input_precision="ieee" + ) + + Ai_31 = -tl.dot( + Ai_33, + tl.dot(A_31, Ai_11, input_precision="ieee") + + tl.dot(A_32, Ai_21, input_precision="ieee"), + input_precision="ieee", + ) + Ai_42 = -tl.dot( + Ai_44, + tl.dot(A_42, Ai_22, input_precision="ieee") + + tl.dot(A_43, Ai_32, input_precision="ieee"), + input_precision="ieee", + ) + Ai_41 = -tl.dot( + Ai_44, + tl.dot(A_41, Ai_11, input_precision="ieee") + + tl.dot(A_42, Ai_21, input_precision="ieee") + + tl.dot(A_43, Ai_31, input_precision="ieee"), + input_precision="ieee", + ) + + p_Ai_11 = tl.make_block_ptr( + Ai, (T, 64), (H * 64, 1), (i_t * 64, 0), (16, 16), (1, 0) + ) + p_Ai_22 = tl.make_block_ptr( + Ai, (T, 64), (H * 64, 1), (i_t * 64 + 16, 16), (16, 16), (1, 0) + ) + p_Ai_33 = tl.make_block_ptr( + Ai, (T, 64), (H * 64, 1), (i_t * 64 + 32, 32), (16, 16), (1, 0) + ) + p_Ai_44 = tl.make_block_ptr( + Ai, (T, 64), (H * 64, 1), (i_t * 64 + 48, 48), (16, 16), (1, 0) + ) + p_Ai_21 = tl.make_block_ptr( + Ai, (T, 64), (H * 64, 1), (i_t * 64 + 16, 0), (16, 16), (1, 0) + ) + p_Ai_31 = tl.make_block_ptr( + Ai, (T, 64), (H * 64, 1), (i_t * 64 + 32, 0), (16, 16), (1, 0) + ) + p_Ai_32 = tl.make_block_ptr( + Ai, (T, 64), (H * 64, 1), (i_t * 64 + 32, 16), (16, 16), (1, 0) + ) + p_Ai_41 = tl.make_block_ptr( + Ai, (T, 64), (H * 64, 1), (i_t * 64 + 48, 0), (16, 16), (1, 0) + ) + p_Ai_42 = tl.make_block_ptr( + Ai, (T, 64), (H * 64, 1), (i_t * 64 + 48, 16), (16, 16), (1, 0) + ) + p_Ai_43 = tl.make_block_ptr( + Ai, (T, 64), (H * 64, 1), (i_t * 64 + 48, 32), (16, 16), (1, 0) + ) + tl.store( + p_Ai_11, + Ai_11.to(p_Ai_11.dtype.element_ty, fp_downcast_rounding="rtne"), + boundary_check=(0, 1), + ) + tl.store( + p_Ai_22, + Ai_22.to(p_Ai_22.dtype.element_ty, fp_downcast_rounding="rtne"), + boundary_check=(0, 1), + ) + tl.store( + p_Ai_33, + Ai_33.to(p_Ai_33.dtype.element_ty, fp_downcast_rounding="rtne"), + boundary_check=(0, 1), + ) + tl.store( + p_Ai_44, + Ai_44.to(p_Ai_44.dtype.element_ty, fp_downcast_rounding="rtne"), + boundary_check=(0, 1), + ) + tl.store( + p_Ai_21, + Ai_21.to(p_Ai_21.dtype.element_ty, fp_downcast_rounding="rtne"), + boundary_check=(0, 1), + ) + tl.store( + p_Ai_31, + Ai_31.to(p_Ai_31.dtype.element_ty, fp_downcast_rounding="rtne"), + boundary_check=(0, 1), + ) + tl.store( + p_Ai_32, + Ai_32.to(p_Ai_32.dtype.element_ty, fp_downcast_rounding="rtne"), + boundary_check=(0, 1), + ) + tl.store( + p_Ai_41, + Ai_41.to(p_Ai_41.dtype.element_ty, fp_downcast_rounding="rtne"), + boundary_check=(0, 1), + ) + tl.store( + p_Ai_42, + Ai_42.to(p_Ai_42.dtype.element_ty, fp_downcast_rounding="rtne"), + boundary_check=(0, 1), + ) + tl.store( + p_Ai_43, + Ai_43.to(p_Ai_43.dtype.element_ty, fp_downcast_rounding="rtne"), + boundary_check=(0, 1), + ) + + fill_zeros = tl.zeros((16, 16), dtype=tl.float32) + p_Ai_12 = tl.make_block_ptr( + Ai, (T, 64), (H * 64, 1), (i_t * 64, 16), (16, 16), (1, 0) + ) + p_Ai_13 = tl.make_block_ptr( + Ai, (T, 64), (H * 64, 1), (i_t * 64, 32), (16, 16), (1, 0) + ) + p_Ai_14 = tl.make_block_ptr( + Ai, (T, 64), (H * 64, 1), (i_t * 64, 48), (16, 16), (1, 0) + ) + p_Ai_23 = tl.make_block_ptr( + Ai, (T, 64), (H * 64, 1), (i_t * 64 + 16, 32), (16, 16), (1, 0) + ) + p_Ai_24 = tl.make_block_ptr( + Ai, (T, 64), (H * 64, 1), (i_t * 64 + 16, 48), (16, 16), (1, 0) + ) + p_Ai_34 = tl.make_block_ptr( + Ai, (T, 64), (H * 64, 1), (i_t * 64 + 32, 48), (16, 16), (1, 0) + ) + tl.store( + p_Ai_12, + fill_zeros.to(p_Ai_12.dtype.element_ty, fp_downcast_rounding="rtne"), + boundary_check=(0, 1), + ) + tl.store( + p_Ai_13, + fill_zeros.to(p_Ai_13.dtype.element_ty, fp_downcast_rounding="rtne"), + boundary_check=(0, 1), + ) + tl.store( + p_Ai_14, + fill_zeros.to(p_Ai_14.dtype.element_ty, fp_downcast_rounding="rtne"), + boundary_check=(0, 1), + ) + tl.store( + p_Ai_23, + fill_zeros.to(p_Ai_23.dtype.element_ty, fp_downcast_rounding="rtne"), + boundary_check=(0, 1), + ) + tl.store( + p_Ai_24, + fill_zeros.to(p_Ai_24.dtype.element_ty, fp_downcast_rounding="rtne"), + boundary_check=(0, 1), + ) + tl.store( + p_Ai_34, + fill_zeros.to(p_Ai_34.dtype.element_ty, fp_downcast_rounding="rtne"), + boundary_check=(0, 1), + ) + diff --git a/benchmarks/TritonNPUBench/sglang/merge_state_kernel/merge_state_kernel.py b/benchmarks/TritonNPUBench/sglang/merge_state_kernel/merge_state_kernel.py new file mode 100644 index 00000000..730e5641 --- /dev/null +++ b/benchmarks/TritonNPUBench/sglang/merge_state_kernel/merge_state_kernel.py @@ -0,0 +1,62 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/sglang_operator_cases/newtest_cases/test_merge_state_kernel_case2.py +# Main kernel: merge_state_kernel +# PT file: test_merge_state_kernel_v2.pt + +import triton +import triton.language as tl + + +# === merge_state_kernel === +@triton.jit +def merge_state_kernel( + v_a_ptr, + s_a_ptr, + v_b_ptr, + s_b_ptr, + v_merged_ptr, + s_merged_ptr, + num_heads, + head_dim, + bdx: tl.constexpr, + bdy: tl.constexpr, +): + pos = tl.program_id(axis=0) + for tx in tl.range(bdx): + for head_idx in tl.range(bdy): + s_a_val = tl.load(s_a_ptr + pos * num_heads + head_idx) + s_b_val = tl.load(s_b_ptr + pos * num_heads + head_idx) + + offsets = (pos * num_heads + head_idx) * head_dim + tx + v_a = tl.load(v_a_ptr + offsets) + v_b = tl.load(v_b_ptr + offsets) + + v_merged, s_max, d = state_merge( + o=v_a, m=s_a_val, d=1, other_o=v_b, other_m=s_b_val, other_d=1 + ) + v_merged, s_max, d = state_normalize(v_merged, s_max, d) + v_merged_offset = (pos * num_heads + head_idx) * head_dim + tx + tl.store(v_merged_ptr + v_merged_offset, v_merged) + + if s_merged_ptr: + tl.store( + s_merged_ptr + pos * num_heads + head_idx, + tl.log2(d) + s_max, + ) + + +# === state_normalize === +@triton.jit +def state_normalize(o, m, d): + o = o / d + return o, m, d + + +# === state_merge === +@triton.jit +def state_merge(o, m, d, other_o, other_m, other_d): + m_max = tl.maximum(m, other_m) + d = d * tl.exp2(m - m_max) + other_d * tl.exp2(other_m - m_max) + o = o * tl.exp2(m - m_max) + other_o * tl.exp2(other_m - m_max) + return o, m_max, d + diff --git a/benchmarks/TritonNPUBench/sglang/moe_align_block_size_stage1/moe_align_block_size_stage1.py b/benchmarks/TritonNPUBench/sglang/moe_align_block_size_stage1/moe_align_block_size_stage1.py new file mode 100644 index 00000000..c1b96595 --- /dev/null +++ b/benchmarks/TritonNPUBench/sglang/moe_align_block_size_stage1/moe_align_block_size_stage1.py @@ -0,0 +1,28 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/sglang_operator_cases/newtest_cases/test_moe_align_block_size_stage1.py +# Main kernel: moe_align_block_size_stage1 +# PT file: test_moe_align_block_size_stage1_v2.pt + +import triton +import triton.language as tl + + +# === moe_align_block_size_stage1 === +@triton.jit +def moe_align_block_size_stage1( + topk_ids_ptr, + tokens_cnts_ptr, + num_experts: tl.constexpr, + numel: tl.constexpr, + tokens_per_thread: tl.constexpr, +): + pid = tl.program_id(0) + start_idx = pid * tokens_per_thread + off_c = (pid + 1) * num_experts + + for i in range(tokens_per_thread): + if start_idx + i < numel: + idx = tl.load(topk_ids_ptr + start_idx + i) + token_cnt = tl.load(tokens_cnts_ptr + off_c + idx) + tl.store(tokens_cnts_ptr + off_c + idx, token_cnt + 1) + diff --git a/benchmarks/TritonNPUBench/sglang/moe_align_block_size_stage2/moe_align_block_size_stage2.py b/benchmarks/TritonNPUBench/sglang/moe_align_block_size_stage2/moe_align_block_size_stage2.py new file mode 100644 index 00000000..d5945988 --- /dev/null +++ b/benchmarks/TritonNPUBench/sglang/moe_align_block_size_stage2/moe_align_block_size_stage2.py @@ -0,0 +1,22 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/sglang_operator_cases/newtest_cases/test_moe_align_block_size_stage2.py +# Main kernel: moe_align_block_size_stage2 +# PT file: test_moe_align_block_size_stage2_v2.pt + +import triton +import triton.language as tl + + +# === moe_align_block_size_stage2 === +@triton.jit +def moe_align_block_size_stage2( + tokens_cnts_ptr, + num_experts: tl.constexpr, +): + pid = tl.program_id(0) + last_cnt = 0 + for i in range(1, num_experts + 1): + token_cnt = tl.load(tokens_cnts_ptr + i * num_experts + pid) + last_cnt = last_cnt + token_cnt + tl.store(tokens_cnts_ptr + i * num_experts + pid, last_cnt) + diff --git a/benchmarks/TritonNPUBench/sglang/moe_align_block_size_stage3/moe_align_block_size_stage3.py b/benchmarks/TritonNPUBench/sglang/moe_align_block_size_stage3/moe_align_block_size_stage3.py new file mode 100644 index 00000000..dc48b413 --- /dev/null +++ b/benchmarks/TritonNPUBench/sglang/moe_align_block_size_stage3/moe_align_block_size_stage3.py @@ -0,0 +1,26 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/sglang_operator_cases/newtest_cases/test_moe_align_block_size_stage3.py +# Main kernel: moe_align_block_size_stage3 +# PT file: test_moe_align_block_size_stage3_v2.pt + +import triton +import triton.language as tl + + +# === moe_align_block_size_stage3 === +@triton.jit +def moe_align_block_size_stage3( + total_tokens_post_pad_ptr, + tokens_cnts_ptr, + cumsum_ptr, + num_experts: tl.constexpr, + block_size: tl.constexpr, +): + last_cumsum = 0 + off_cnt = num_experts * num_experts + for i in range(1, num_experts + 1): + token_cnt = tl.load(tokens_cnts_ptr + off_cnt + i - 1) + last_cumsum = last_cumsum + tl.cdiv(token_cnt, block_size) * block_size + tl.store(cumsum_ptr + i, last_cumsum) + tl.store(total_tokens_post_pad_ptr, last_cumsum) + diff --git a/benchmarks/TritonNPUBench/sglang/moe_align_block_size_stage4/moe_align_block_size_stage4.py b/benchmarks/TritonNPUBench/sglang/moe_align_block_size_stage4/moe_align_block_size_stage4.py new file mode 100644 index 00000000..67091d1e --- /dev/null +++ b/benchmarks/TritonNPUBench/sglang/moe_align_block_size_stage4/moe_align_block_size_stage4.py @@ -0,0 +1,39 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/sglang_operator_cases/newtest_cases/test_moe_align_block_size_stage4.py +# Main kernel: moe_align_block_size_stage4 +# PT file: test_moe_align_block_size_stage4_v2_new.pt + +import triton +import triton.language as tl + + +# === moe_align_block_size_stage4 === +@triton.jit +def moe_align_block_size_stage4( + topk_ids_ptr, + sorted_token_ids_ptr, + expert_ids_ptr, + tokens_cnts_ptr, + cumsum_ptr, + num_experts: tl.constexpr, + block_size: tl.constexpr, + numel: tl.constexpr, + tokens_per_thread: tl.constexpr, +): + pid = tl.program_id(0) + start_idx = tl.load(cumsum_ptr + pid) + end_idx = tl.load(cumsum_ptr + pid + 1) + + for i in range(start_idx, end_idx, block_size): + tl.store(expert_ids_ptr + i // block_size, pid) + + start_idx = pid * tokens_per_thread + off_t = pid * num_experts + + for i in range(start_idx, tl.minimum(start_idx + tokens_per_thread, numel)): + expert_id = tl.load(topk_ids_ptr + i) + token_cnt = tl.load(tokens_cnts_ptr + off_t + expert_id) + rank_post_pad = token_cnt + tl.load(cumsum_ptr + expert_id) + tl.store(sorted_token_ids_ptr + rank_post_pad, i) + tl.store(tokens_cnts_ptr + off_t + expert_id, token_cnt + 1) + diff --git a/benchmarks/TritonNPUBench/sglang/nsa_cp_round_robin_split_q_seqs_kernel/nsa_cp_round_robin_split_q_seqs_kernel.py b/benchmarks/TritonNPUBench/sglang/nsa_cp_round_robin_split_q_seqs_kernel/nsa_cp_round_robin_split_q_seqs_kernel.py new file mode 100644 index 00000000..6865a1fb --- /dev/null +++ b/benchmarks/TritonNPUBench/sglang/nsa_cp_round_robin_split_q_seqs_kernel/nsa_cp_round_robin_split_q_seqs_kernel.py @@ -0,0 +1,31 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/sglang_operator_cases/newtest_cases/test_nsa_cp_round_robin_split_q_seqs_kernel.py +# Main kernel: nsa_cp_round_robin_split_q_seqs_kernel +# PT file: test_nsa_cp_round_robin_split_q_seqs_kernel_v2.pt + +import triton +import triton.language as tl + + +# === nsa_cp_round_robin_split_q_seqs_kernel === +@triton.jit +def nsa_cp_round_robin_split_q_seqs_kernel( + in_seqs_ptr, + out_seqs_ptr, + bs_idx_ptr, + tokens: tl.constexpr, + cp_size: tl.constexpr, + cp_rank: tl.constexpr, +): + extra_seq = 0 + bs_idx = 0 + for bs in range(tokens): + cur_len = tl.load(in_seqs_ptr + bs) + cur_len += extra_seq + cur_seq = cur_len // cp_size + (cur_len % cp_size > cp_rank) + if cur_seq > 0: + tl.store(bs_idx_ptr + bs_idx, bs) + tl.store(out_seqs_ptr + bs_idx, cur_seq) + bs_idx += 1 + extra_seq = cur_len - cur_seq * cp_size + diff --git a/benchmarks/TritonNPUBench/sglang/pad_draft_extend_query_kernel/pad_draft_extend_query_kernel.py b/benchmarks/TritonNPUBench/sglang/pad_draft_extend_query_kernel/pad_draft_extend_query_kernel.py new file mode 100644 index 00000000..7e9a2a1c --- /dev/null +++ b/benchmarks/TritonNPUBench/sglang/pad_draft_extend_query_kernel/pad_draft_extend_query_kernel.py @@ -0,0 +1,82 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/sglang_operator_cases/newtest_cases/test_pad_draft_extend_query_kernel.py +# Main kernel: pad_draft_extend_query_kernel +# PT file: test_pad_draft_extend_query_kernel_v2.pt + +import triton +import triton.language as tl + + +# === pad_draft_extend_query_kernel === +@triton.jit +def pad_draft_extend_query_kernel( + q_ptr, # Input query tensor [total_seq_len, num_heads, head_dim] + padded_q_ptr, # Output padded query tensor [batch_size, max_seq_len, num_heads, head_dim] + seq_lens_q_ptr, # Sequence lengths for each sequence [batch_size] + cumsum_ptr, # Cumulative sum of accept lengths [batch_size + 1] + batch_size, + max_seq_len, + num_heads, + head_dim, + BLOCK_SIZE: tl.constexpr, +): + """Triton kernel for padding draft extended query tensor with parallelized head and dim processing.""" + # Use 3D program IDs: (batch_seq, head_block, dim_block) + batch_seq_pid = tl.program_id(0) + head_pid = tl.program_id(1) + dim_pid = tl.program_id(2) + + batch_id = batch_seq_pid // max_seq_len + seq_pos = batch_seq_pid % max_seq_len + + if batch_id >= batch_size: + return + + # Load accept length for this batch + seq_len = tl.load(seq_lens_q_ptr + batch_id) + + if seq_pos >= seq_len: + return + + # Load cumulative sum to get start position in input tensor + input_start = tl.load(cumsum_ptr + batch_id) + input_pos = input_start + seq_pos + + # Calculate head and dim block ranges + head_start = head_pid * BLOCK_SIZE + head_end = tl.minimum(head_start + BLOCK_SIZE, num_heads) + head_mask = tl.arange(0, BLOCK_SIZE) < (head_end - head_start) + + dim_start = dim_pid * BLOCK_SIZE + dim_end = tl.minimum(dim_start + BLOCK_SIZE, head_dim) + dim_mask = tl.arange(0, BLOCK_SIZE) < (dim_end - dim_start) + + # Calculate input offset + input_offset = ( + input_pos * num_heads * head_dim + + (head_start + tl.arange(0, BLOCK_SIZE))[:, None] * head_dim + + (dim_start + tl.arange(0, BLOCK_SIZE))[None, :] + ) + + # Load data + data = tl.load( + q_ptr + input_offset, + mask=head_mask[:, None] & dim_mask[None, :], + other=0.0, + ) + + # Calculate output offset + output_offset = ( + batch_id * max_seq_len * num_heads * head_dim + + seq_pos * num_heads * head_dim + + (head_start + tl.arange(0, BLOCK_SIZE))[:, None] * head_dim + + (dim_start + tl.arange(0, BLOCK_SIZE))[None, :] + ) + + # Store data + tl.store( + padded_q_ptr + output_offset, + data, + mask=head_mask[:, None] & dim_mask[None, :], + ) + diff --git a/benchmarks/TritonNPUBench/sglang/pad_sequence_with_mask_kernel/pad_sequence_with_mask_kernel.py b/benchmarks/TritonNPUBench/sglang/pad_sequence_with_mask_kernel/pad_sequence_with_mask_kernel.py new file mode 100644 index 00000000..d58301b1 --- /dev/null +++ b/benchmarks/TritonNPUBench/sglang/pad_sequence_with_mask_kernel/pad_sequence_with_mask_kernel.py @@ -0,0 +1,63 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/sglang_operator_cases/newtest_cases/test_pad_sequence_with_mask_kernel.py +# Main kernel: pad_sequence_with_mask_kernel +# PT file: test_pad_sequence_with_mask_kernel_v2.pt + +import triton +import triton.language as tl + + +# === pad_sequence_with_mask_kernel === +@triton.jit +def pad_sequence_with_mask_kernel( + input_ptr, # (total_tokens, hidden) + offsets_ptr, # (B,) + lengths_ptr, # (B,) + output_ptr, # (B, max_len, hidden) + mask_ptr, # (B, max_len) + max_len, + hidden_dim, + BLOCK_M: tl.constexpr, # seq block + BLOCK_D: tl.constexpr, # hidden block +): + b = tl.program_id(0) # batch index + m = tl.program_id(1) # seq block index + + offset = tl.load(offsets_ptr + b) + length = tl.load(lengths_ptr + b) + + seq_ids = m * BLOCK_M + tl.arange(0, BLOCK_M) + hid_ids = tl.arange(0, BLOCK_D) + + seq_mask = seq_ids < max_len + valid_token = seq_ids < length + + # input index + in_token = offset + seq_ids + in_ptr = input_ptr + in_token[:, None] * hidden_dim + hid_ids[None, :] + + # output index + out_ptr = ( + output_ptr + + b * max_len * hidden_dim + + seq_ids[:, None] * hidden_dim + + hid_ids[None, :] + ) + + values = tl.load( + in_ptr, + mask=valid_token[:, None] & (hid_ids[None, :] < hidden_dim), + other=0.0, + ) + + tl.store( + out_ptr, + values, + mask=seq_mask[:, None] & (hid_ids[None, :] < hidden_dim), + ) + + # attention mask + if tl.program_id(2) == 0: + mask_out_ptr = mask_ptr + b * max_len + seq_ids + tl.store(mask_out_ptr, valid_token, mask=seq_mask) + diff --git a/benchmarks/TritonNPUBench/sglang/post_reorder_triton_kernel/post_reorder_triton_kernel.py b/benchmarks/TritonNPUBench/sglang/post_reorder_triton_kernel/post_reorder_triton_kernel.py new file mode 100644 index 00000000..58990312 --- /dev/null +++ b/benchmarks/TritonNPUBench/sglang/post_reorder_triton_kernel/post_reorder_triton_kernel.py @@ -0,0 +1,49 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/sglang_operator_cases/newtest_cases/test_post_reorder_triton_kernel.py +# Main kernel: post_reorder_triton_kernel +# PT file: test_post_reorder_triton_kernel_v2_unit2.pt + +import triton +import triton.language as tl + + +# === post_reorder_triton_kernel === +@triton.jit +def post_reorder_triton_kernel( + down_output_ptr, + output_ptr, + src2dst_ptr, + topk_ids_ptr, + topk_weights_ptr, + topk, + hidden_size, + BLOCK_SIZE: tl.constexpr, +): + InDtype = down_output_ptr.dtype.element_ty + + src_idx_int32 = tl.program_id(0) + src_idx = src_idx_int32.to(tl.int64) + src2dst_ptr = src2dst_ptr + src_idx * topk + topk_ids_ptr = topk_ids_ptr + src_idx * topk + topk_weights_ptr = topk_weights_ptr + src_idx * topk + + store_ptr = output_ptr + src_idx * hidden_size + + vec = tl.arange(0, BLOCK_SIZE) + + for start_offset in tl.range(0, hidden_size, BLOCK_SIZE): + offset = start_offset + vec + mask = offset < hidden_size + + sum_vec = tl.zeros([BLOCK_SIZE], dtype=InDtype) + for idx in range(topk): + expert_id = tl.load(topk_ids_ptr + idx) + if expert_id > 0: + dst_idx_int32 = tl.load(src2dst_ptr + idx) + dst_idx = dst_idx_int32.to(tl.int64) + weigh_scale = tl.load(topk_weights_ptr + idx).to(InDtype) + load_ptr = down_output_ptr + dst_idx * hidden_size + in_data = tl.load(load_ptr + offset, mask=mask) + sum_vec += in_data * weigh_scale + tl.store(store_ptr + offset, sum_vec, mask=mask) + diff --git a/benchmarks/TritonNPUBench/sglang/post_reorder_triton_kernel_for_cutlass_moe/post_reorder_triton_kernel_for_cutlass_moe.py b/benchmarks/TritonNPUBench/sglang/post_reorder_triton_kernel_for_cutlass_moe/post_reorder_triton_kernel_for_cutlass_moe.py new file mode 100644 index 00000000..d18ed0fb --- /dev/null +++ b/benchmarks/TritonNPUBench/sglang/post_reorder_triton_kernel_for_cutlass_moe/post_reorder_triton_kernel_for_cutlass_moe.py @@ -0,0 +1,59 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/sglang_operator_cases/newtest_cases/test_post_reorder_triton_kernel_for_cutlass_moe.py +# Main kernel: post_reorder_triton_kernel_for_cutlass_moe +# PT file: test_post_reorder_triton_kernel_for_cutlass_moe_v2.pt + +import triton +import triton.language as tl + + +# === post_reorder_triton_kernel_for_cutlass_moe === +@triton.jit +def post_reorder_triton_kernel_for_cutlass_moe( + down_output_ptr, + output_ptr, + src2dst_ptr, + topk_ids_ptr, + topk_weights_ptr, + num_local_experts, + topk, + num_tokens, + hidden_size, + routed_scaling_factor: float, + BLOCK_SIZE: tl.constexpr, + NUM_STAGES: tl.constexpr, +): + OutDtype = output_ptr.dtype.element_ty + + offset = BLOCK_SIZE * tl.program_id(1) + tl.arange(0, BLOCK_SIZE) + mask = offset < hidden_size + + down_output_ptr_offs = down_output_ptr + offset + output_ptr_offs = output_ptr + offset + + start_src_idx = tl.program_id(0) + step = tl.num_programs(0) + + for src_idx_int32 in tl.range( + start_src_idx, num_tokens, step, num_stages=NUM_STAGES + ): + src_idx = src_idx_int32.to(tl.int64) + token_src2dst_ptr = src2dst_ptr + src_idx * topk + token_topk_ids_ptr = topk_ids_ptr + src_idx * topk + token_topk_weights_ptr = topk_weights_ptr + src_idx * topk + + sum_vec = tl.zeros([BLOCK_SIZE], dtype=tl.float32) + for idx in range(topk): + expert_id = tl.load(token_topk_ids_ptr + idx) + if expert_id != num_local_experts: + dst_idx_int32 = tl.load(token_src2dst_ptr + idx) + dst_idx = dst_idx_int32.to(tl.int64) + dst_idx = dst_idx + weight_scale = tl.load(token_topk_weights_ptr + idx).to(tl.float32) + load_ptr_offs = down_output_ptr_offs + dst_idx * hidden_size + in_data = tl.load(load_ptr_offs, mask=mask).to(tl.float32) + sum_vec += in_data * weight_scale + sum_vec *= routed_scaling_factor + store_ptr_offs = output_ptr_offs + src_idx * hidden_size + tl.store(store_ptr_offs, sum_vec.to(OutDtype), mask=mask) + diff --git a/benchmarks/TritonNPUBench/sglang/pre_reorder_triton_kernel_for_cutlass_moe/pre_reorder_triton_kernel_for_cutlass_moe.py b/benchmarks/TritonNPUBench/sglang/pre_reorder_triton_kernel_for_cutlass_moe/pre_reorder_triton_kernel_for_cutlass_moe.py new file mode 100644 index 00000000..4acdd337 --- /dev/null +++ b/benchmarks/TritonNPUBench/sglang/pre_reorder_triton_kernel_for_cutlass_moe/pre_reorder_triton_kernel_for_cutlass_moe.py @@ -0,0 +1,54 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/sglang_operator_cases/newtest_cases/test_pre_reorder_triton_kernel_for_cutlass_moe.py +# Main kernel: pre_reorder_triton_kernel_for_cutlass_moe +# PT file: test_pre_reorder_triton_kernel_for_cutlass_moe_v2.pt + +import triton +import triton.language as tl + + +# === pre_reorder_triton_kernel_for_cutlass_moe === +@triton.jit +def pre_reorder_triton_kernel_for_cutlass_moe( + input_ptr, + gateup_input_ptr, + src2dst_ptr, + topk_ids_ptr, + a1_scales_ptr, + num_local_experts, + topk, + num_tokens, + hidden_size, + BLOCK_SIZE: tl.constexpr, + NUM_STAGES: tl.constexpr, +): + OutDtype = gateup_input_ptr.dtype.element_ty + + if a1_scales_ptr is not None: + a1_scale = 1.0 / tl.load(a1_scales_ptr) + else: + a1_scale = 1.0 + + offset = BLOCK_SIZE * tl.program_id(1) + tl.arange(0, BLOCK_SIZE) + mask = offset < hidden_size + + start_src_idx = tl.program_id(0) + step = tl.num_programs(0) + + for src_idx_int32 in tl.range( + start_src_idx, num_tokens, step, num_stages=NUM_STAGES + ): + src_idx = src_idx_int32.to(tl.int64) + token_src2dst_ptr = src2dst_ptr + src_idx * topk + token_topk_ids_ptr = topk_ids_ptr + src_idx * topk + + src_ptr_offs = input_ptr + src_idx * hidden_size + offset + dst_ptr_offs = gateup_input_ptr + offset + in_data = tl.load(src_ptr_offs, mask=mask).to(tl.float32) + out_data = (in_data * a1_scale).to(OutDtype) + for idx in range(topk): + expert_id = tl.load(token_topk_ids_ptr + idx) + if expert_id != num_local_experts: + dst_idx = tl.load(token_src2dst_ptr + idx) + tl.store(dst_ptr_offs + dst_idx * hidden_size, out_data, mask=mask) + diff --git a/benchmarks/TritonNPUBench/sglang/recompute_w_u_fwd_kernel/recompute_w_u_fwd_kernel.py b/benchmarks/TritonNPUBench/sglang/recompute_w_u_fwd_kernel/recompute_w_u_fwd_kernel.py new file mode 100644 index 00000000..7da6e613 --- /dev/null +++ b/benchmarks/TritonNPUBench/sglang/recompute_w_u_fwd_kernel/recompute_w_u_fwd_kernel.py @@ -0,0 +1,98 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/sglang_operator_cases/newtest_cases/test_recompute_w_u_fwd_kernel.py +# Main kernel: recompute_w_u_fwd_kernel +# PT file: test_recompute_w_u_fwd_kernel_v2.pt + +import triton +import triton.language as tl + + +# === recompute_w_u_fwd_kernel === +@triton.jit(do_not_specialize=["T"]) +def recompute_w_u_fwd_kernel( + k, + v, + beta, + w, + u, + A, + g, + cu_seqlens, + chunk_indices, + T, + H: tl.constexpr, + Hg: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_t, i_bh = tl.program_id(0), tl.program_id(1) + i_b, i_h = i_bh // H, i_bh % H + if IS_VARLEN: + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load( + chunk_indices + i_t * 2 + 1 + ).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load( + cu_seqlens + i_n + 1 + ).to(tl.int32) + T = eos - bos + else: + bos, eos = i_b * T, i_b * T + T + p_beta = tl.make_block_ptr( + beta + bos * H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,) + ) + p_g = tl.make_block_ptr(g + (bos * H + i_h), (T,), (H,), (i_t * BT,), (BT,), (0,)) + p_A = tl.make_block_ptr( + A + (bos * H + i_h) * BT, (T, BT), (H * BT, 1), (i_t * BT, 0), (BT, BT), (1, 0) + ) + b_beta = tl.load(p_beta, boundary_check=(0,)) + b_A = tl.load(p_A, boundary_check=(0, 1)) + b_g = tl.exp(tl.load(p_g, boundary_check=(0,))) + + for i_v in range(tl.cdiv(V, BV)): + p_v = tl.make_block_ptr( + v + (bos * H + i_h) * V, + (T, V), + (H * V, 1), + (i_t * BT, i_v * BV), + (BT, BV), + (1, 0), + ) + p_u = tl.make_block_ptr( + u + (bos * H + i_h) * V, + (T, V), + (H * V, 1), + (i_t * BT, i_v * BV), + (BT, BV), + (1, 0), + ) + b_v = tl.load(p_v, boundary_check=(0, 1)) + b_vb = (b_v * b_beta[:, None]).to(b_v.dtype) + b_u = tl.dot(b_A, b_vb, allow_tf32=False) + tl.store(p_u, b_u.to(p_u.dtype.element_ty), boundary_check=(0, 1)) + + for i_k in range(tl.cdiv(K, BK)): + p_k = tl.make_block_ptr( + k + (bos * Hg + i_h // (H // Hg)) * K, + (T, K), + (Hg * K, 1), + (i_t * BT, i_k * BK), + (BT, BK), + (1, 0), + ) + p_w = tl.make_block_ptr( + w + (bos * H + i_h) * K, + (T, K), + (H * K, 1), + (i_t * BT, i_k * BK), + (BT, BK), + (1, 0), + ) + b_k = tl.load(p_k, boundary_check=(0, 1)) + b_kb = (b_k * b_beta[:, None] * b_g[:, None]).to(b_k.dtype) + b_w = tl.dot(b_A, b_kb) + tl.store(p_w, b_w.to(p_w.dtype.element_ty), boundary_check=(0, 1)) + diff --git a/benchmarks/TritonNPUBench/sglang/rotl32_kernel/rotl32_kernel.py b/benchmarks/TritonNPUBench/sglang/rotl32_kernel/rotl32_kernel.py new file mode 100644 index 00000000..0a0cee21 --- /dev/null +++ b/benchmarks/TritonNPUBench/sglang/rotl32_kernel/rotl32_kernel.py @@ -0,0 +1,29 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/sglang_operator_cases/newtest_cases/test__rotl32.py +# Main kernel: rotl32_kernel +# PT file: test_rotl32_kernel_v2.pt + +import triton +import triton.language as tl + +@triton.jit +def _rotl32(x, r: tl.constexpr): + return (x << r) | (x >> (32 - r)) + +# === rotl32_kernel === +@triton.jit +def rotl32_kernel( + x_ptr, + y_ptr, + n_elements, + r: tl.constexpr, + BLOCK_SIZE: tl.constexpr +): + pid = tl.program_id(axis=0) + block_start = pid * BLOCK_SIZE + offsets = block_start + tl.arange(0, BLOCK_SIZE) + mask = offsets < n_elements + x = tl.load(x_ptr + offsets, mask=mask, other=0) + y = _rotl32(x, r) + tl.store(y_ptr + offsets, y, mask=mask) + diff --git a/benchmarks/TritonNPUBench/sglang/safe_exp_kernel/safe_exp_kernel.py b/benchmarks/TritonNPUBench/sglang/safe_exp_kernel/safe_exp_kernel.py new file mode 100644 index 00000000..24fbadfd --- /dev/null +++ b/benchmarks/TritonNPUBench/sglang/safe_exp_kernel/safe_exp_kernel.py @@ -0,0 +1,25 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/sglang_operator_cases/newtest_cases/test_safe_exp.py +# Main kernel: safe_exp_kernel +# PT file: test_safe_exp_kernel_v2.pt + +import triton +import triton.language as tl + + +# === safe_exp_kernel === +@triton.jit +def safe_exp_kernel( + x_ptr, + y_ptr, + n_elements, + BLOCK_SIZE: tl.constexpr +): + pid = tl.program_id(axis=0) + block_start = pid * BLOCK_SIZE + offsets = block_start + tl.arange(0, BLOCK_SIZE) + mask = offsets < n_elements + x = tl.load(x_ptr + offsets, mask=mask, other=0.0) + y = tl.exp(tl.where(x <= 0, x, float("-inf"))) + tl.store(y_ptr + offsets, y, mask=mask) + diff --git a/benchmarks/TritonNPUBench/sglang/sglang_gpu_perf.csv b/benchmarks/TritonNPUBench/sglang/sglang_gpu_perf.csv new file mode 100644 index 00000000..7a859e28 --- /dev/null +++ b/benchmarks/TritonNPUBench/sglang/sglang_gpu_perf.csv @@ -0,0 +1,89 @@ +Name,Type,Input Shapes,Input Data Types,Duration(us) +_act_quant_kernel,VV,None,None,1.713 +_bmm_chunk_fwd_kernel,CV,None,None,1.619 +_chunk_cumsum_fwd_kernel,VV,None,None,3.44 +_chunk_scan_fwd_kernel,CV,None,None,2.233 +_chunked_lora_expand_kernel,CV,None,None,20.066 +_chunked_lora_shrink_kernel,CV,None,None, +_dequantize_k_cache_fast_kernel,VV,None,None,1.058 +_dequantize_k_cache_paged_kernel,VV,None,None,1.323 +_embedding_lora_a_kernel,VV,None,None,1.656 +_fwd_kernel_ep_gather,VV,None,None,1.318 +_fwd_kernel_ep_scatter_1,VV,None,None,1.111 +_fwd_kernel_ep_scatter_2,VV,None,None,955.982 +_gate_up_lora_b_kernel,CV,None,None,1.806 +_layer_norm_fwd_1pass_kernel,VV,None,None,1.971 +_moe_sum_reduce_kernel,VV,None,None,3.386 +_qkv_lora_b_kernel,CV,None,None,1.872 +_selective_scan_update_kernel,VV,None,None,3.328 +_sgemm_lora_a_kernel,CV,None,None,2.102 +_sgemm_lora_b_kernel,CV,None,None,1.794 +_silu_and_mul_post_per_tensor_quant_kernel,VV,None,None,19.337 +_silu_and_mul_post_quant_kernel,VV,None,None,25.125 +_state_passing_fwd_kernel,VV,None,None,10.245 +_tma_align_input_scale_kernel,VV,None,None,1.04 +act_and_mul_kernel,VV,None,None,1.291 +add_tree_reduce_u64_kernel,VV,None,None,53.616 +chunk_fwd_kernel_o,CV,None,None,6.71 +chunk_gated_delta_rule_fwd_kernel_h_blockdim64,CV,None,None,7.512 +chunk_kda_scaled_dot_kkt_fwd_kernel_intra_sub_inter,CV,None,None,2.747 +chunk_kda_scaled_dot_kkt_fwd_kernel_intra_sub_intra,VV,None,None,10.311 +chunk_local_cumsum_scalar_kernel,VV,None,None,1.322 +chunk_local_cumsum_vector_kernel,CV,None,None,1.608 +chunk_scaled_dot_kkt_fwd_kernel,CV,None,None,2.368 +compute_identity_kernel,VV,None,None,1.705 +compute_masked_m_triton_kernel,VV,None,None,1.007 +compute_problem_sizes_w4a8_kernel,VV,None,None,1.038 +compute_seg_indptr_triton_kernel,VV,None,None,1.657 +compute_src2dst_triton_kernel,VV,None,None,1.206 +concat_and_cast_mha_k_kernel,VV,None,None,1.44 +create_flashinfer_kv_indices_triton,VV,None,None,1.724 +create_flashmla_kv_indices_triton,VV,None,None,1.34 +deepep_compute_src2dst_triton_kernel,VV,None,None,1.023 +deepep_permute_triton_kernel,VV,None,None,1.116 +deepep_post_reorder_triton_kernel,VV,None,None,1.387 +deepgemm_compute_src2dst_triton_kernel,VV,None,None,1.209 +experts_combine_kernel,VV,None,None,2.092 +fill_gateup_input_triton_kernel,VV,None,None,1.702 +fn_triton_kernel,VV,None,None,5.961 +fused_dual_residual_rmsnorm_kernel,VV,None,None,1.582 +fused_gdn_gating_kernel,VV,None,None,1.232 +fused_moe_kernel,CV,None,None,2.187 +fused_moe_kernel_gptq_awq,CV,None,None,29.271 +fused_moe_router_cudacore_kernel,VV,None,None,2.277 +fused_moe_router_tensorcore_kernel,CV,None,None,2.379 +fused_recurrent_gated_delta_rule_fwd_kernel,VV,None,None,2.74 +fused_recurrent_gated_delta_rule_update_fwd_kernel,VV,None,None,33.569 +fused_rmsnorm_kernel,VV,None,None,1.384 +fused_sigmoid_gating_delta_rule_update_kernel,VV,None,None,5.662 +fused_softcap_kernel,VV,None,None,1.1 +gelu_and_mul_kernel,VV,None,None,1.052 +get_last_loc_kernel,VV,None,None,1.15 +get_num_kv_splits_triton,VV,None,None,1.537 +l2norm_fwd_kernel,VV,None,None,2.002 +l2norm_fwd_kernel1,VV,None,None,1.271 +layer_norm_gated_fwd_kernel,VV,None,None,4.007 +layer_norm_gated_fwd_kernel1,VV,None,None,1.873 +memcpy_triton_kernel,VV,None,None,1.155 +merge_16x16_to_32x32_inverse_kernel,CV,None,None,3.8 +merge_16x16_to_64x64_inverse_kernel,CV,None,None,2.906 +merge_state_kernel,VV,None,None,7.519 +moe_align_block_size_stage1,VV,None,None,1.301 +moe_align_block_size_stage2,VV,None,None,1.03 +moe_align_block_size_stage3,VV,None,None,1.115 +moe_align_block_size_stage4,VV,None,None,2.113 +nsa_cp_round_robin_split_q_seqs_kernel,VV,None,None,11.912 +pad_draft_extend_query_kernel,VV,None,None,149.019 +pad_sequence_with_mask_kernel,VV,None,None,4.989 +post_reorder_triton_kernel,VV,None,None,1.794 +post_reorder_triton_kernel_for_cutlass_moe,VV,None,None,4.128 +pre_reorder_triton_kernel_for_cutlass_moe,VV,None,None,1.519 +recompute_w_u_fwd_kernel,CV,None,None,531.835 +silu_and_mul_kernel,VV,None,None,1.135 +silu_mul_static_tensorwise_quant_triton_kernel_for_cutlass_moe,VV,None,None,1.343 +solve_tril_16x16_kernel,VV,None,None,3.137 +track_mamba_state_if_needed_kernel,VV,None,None,2.156 +transform_index_page_table_decode_kernel,VV,None,None,1.6 +unpad_draft_extend_output_kernel,VV,None,None,181.388 +write_req_to_token_pool_triton,VV,None,None,3.824 +write_req_to_token_pool_triton_optimize,VV,None,None,1.231 diff --git a/benchmarks/TritonNPUBench/sglang/silu_and_mul_kernel/silu_and_mul_kernel.py b/benchmarks/TritonNPUBench/sglang/silu_and_mul_kernel/silu_and_mul_kernel.py new file mode 100644 index 00000000..abab8cd1 --- /dev/null +++ b/benchmarks/TritonNPUBench/sglang/silu_and_mul_kernel/silu_and_mul_kernel.py @@ -0,0 +1,47 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/sglang_operator_cases/newtest_cases/test_silu_and_mul_kernel.py +# Main kernel: silu_and_mul_kernel +# PT file: test_silu_and_mul_kernel_v2.pt + +import triton +import triton.language as tl + + +# === silu_and_mul_kernel === +@triton.jit +def silu_and_mul_kernel( + out_hidden_states_ptr, # (bs, hidden_dim) + out_scales_ptr, # (bs,) + hidden_states_ptr, # (bs, hidden_dim * 2) + quant_max: tl.constexpr, + static_scale: tl.constexpr, + hidden_dim: tl.constexpr, # the output hidden_dim + BLOCK_SIZE: tl.constexpr, +): + pid = tl.program_id(axis=0) + + input_start = pid * hidden_dim * 2 + output_start = pid * hidden_dim + + input1_offs = tl.arange(0, BLOCK_SIZE) + mask = tl.arange(0, BLOCK_SIZE) < hidden_dim # shared for input1, input3, output + input3_offs = hidden_dim + tl.arange(0, BLOCK_SIZE) + output_offs = tl.arange(0, BLOCK_SIZE) + + x1 = tl.load( + hidden_states_ptr + input_start + input1_offs, mask=mask, other=0.0 + ).to(tl.float32) + x3 = tl.load( + hidden_states_ptr + input_start + input3_offs, mask=mask, other=0.0 + ).to(tl.float32) + + # silu + # cast down before mul to better match training? + silu_x1 = x1 * tl.sigmoid(x1) + out = x3 * silu_x1.to(hidden_states_ptr.dtype.element_ty) + + if quant_max is not None: + raise NotImplementedError() + + tl.store(out_hidden_states_ptr + output_start + output_offs, out, mask=mask) + diff --git a/benchmarks/TritonNPUBench/sglang/silu_mul_static_tensorwise_quant_triton_kernel_for_cutlass_moe/silu_mul_static_tensorwise_quant_triton_kernel_for_cutlass_moe.py b/benchmarks/TritonNPUBench/sglang/silu_mul_static_tensorwise_quant_triton_kernel_for_cutlass_moe/silu_mul_static_tensorwise_quant_triton_kernel_for_cutlass_moe.py new file mode 100644 index 00000000..3a6a225a --- /dev/null +++ b/benchmarks/TritonNPUBench/sglang/silu_mul_static_tensorwise_quant_triton_kernel_for_cutlass_moe/silu_mul_static_tensorwise_quant_triton_kernel_for_cutlass_moe.py @@ -0,0 +1,42 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/sglang_operator_cases/newtest_cases/test_silu_mul_static_tensorwise_quant_triton_kernel_for_cutlass_moe.py +# Main kernel: silu_mul_static_tensorwise_quant_triton_kernel_for_cutlass_moe +# PT file: test_silu_mul_static_tensorwise_quant_triton_kernel_for_cutlass_moe_v2.pt + +import triton +import triton.language as tl + + +# === silu_mul_static_tensorwise_quant_triton_kernel_for_cutlass_moe === +@triton.jit +def silu_mul_static_tensorwise_quant_triton_kernel_for_cutlass_moe( + input_ptr, + output_ptr, + scale_ptr, + num_tokens_tensor_ptr, + intermediate_size, + BLOCK_SIZE: tl.constexpr, + NUM_STAGES: tl.constexpr, +): + OutDtype = output_ptr.dtype.element_ty + + num_tokens = tl.load(num_tokens_tensor_ptr) + numel = num_tokens * intermediate_size + gate_ptr = input_ptr + up_ptr = input_ptr + intermediate_size + scale = 1.0 / tl.load(scale_ptr) + + start_idx = tl.program_id(0) * BLOCK_SIZE + step = tl.num_programs(0) * BLOCK_SIZE + + for id in tl.range(start_idx, numel, step, num_stages=NUM_STAGES): + ids = id + tl.arange(0, BLOCK_SIZE) + token_ids = ids // intermediate_size + mask = ids < numel + + offs = ids + token_ids * intermediate_size + gate = tl.load(gate_ptr + offs, mask=mask, other=0.0).to(tl.float32) + up = tl.load(up_ptr + offs, mask=mask, other=0.0).to(tl.float32) + output = gate / (1 + tl.exp(-gate)) * up * scale + tl.store(output_ptr + ids, output.to(OutDtype), mask=mask) + diff --git a/benchmarks/TritonNPUBench/sglang/solve_tril_16x16_kernel/solve_tril_16x16_kernel.py b/benchmarks/TritonNPUBench/sglang/solve_tril_16x16_kernel/solve_tril_16x16_kernel.py new file mode 100644 index 00000000..5c9a408d --- /dev/null +++ b/benchmarks/TritonNPUBench/sglang/solve_tril_16x16_kernel/solve_tril_16x16_kernel.py @@ -0,0 +1,58 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/sglang_operator_cases/newtest_cases/test_solve_tril_16x16_kernel.py +# Main kernel: solve_tril_16x16_kernel +# PT file: test_solve_tril_16x16_kernel_v2.pt + +import triton +import triton.language as tl + + +# === solve_tril_16x16_kernel === +@triton.jit(do_not_specialize=["T"]) +def solve_tril_16x16_kernel( + A, + Ad, + cu_seqlens, + chunk_indices, + T, + H: tl.constexpr, + BT: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_t, i_bh = tl.program_id(0), tl.program_id(1) + i_b, i_h = i_bh // H, i_bh % H + if IS_VARLEN: + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load( + chunk_indices + i_t * 2 + 1 + ).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load( + cu_seqlens + i_n + 1 + ).to(tl.int32) + T = eos - bos + else: + bos, eos = i_b * T, i_b * T + T + + A = A + (bos * H + i_h) * BT + Ad = Ad + (bos * H + i_h) * 16 + + offset = (i_t * 16) % BT + p_A = tl.make_block_ptr( + A, (T, BT), (H * BT, 1), (i_t * 16, offset), (16, 16), (1, 0) + ) + p_Ai = tl.make_block_ptr(Ad, (T, 16), (H * 16, 1), (i_t * 16, 0), (16, 16), (1, 0)) + b_A = tl.load(p_A, boundary_check=(0, 1)).to(tl.float32) + b_A = -tl.where(tl.arange(0, 16)[:, None] > tl.arange(0, 16)[None, :], b_A, 0) + + o_i = tl.arange(0, 16) + for i in range(1, min(16, T - i_t * 16)): + b_a = -tl.load(A + (i_t * 16 + i) * H * BT + o_i + offset) + b_a = b_a + tl.sum(b_a[:, None] * b_A, 0) + mask = o_i == i + b_A = tl.where(mask[:, None], b_a, b_A) + b_A += o_i[:, None] == o_i[None, :] + tl.store( + p_Ai, + b_A.to(p_Ai.dtype.element_ty, fp_downcast_rounding="rtne"), + boundary_check=(0, 1), + ) + diff --git a/benchmarks/TritonNPUBench/sglang/tanh_kernel/tanh_kernel.py b/benchmarks/TritonNPUBench/sglang/tanh_kernel/tanh_kernel.py new file mode 100644 index 00000000..68f094d3 --- /dev/null +++ b/benchmarks/TritonNPUBench/sglang/tanh_kernel/tanh_kernel.py @@ -0,0 +1,25 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/sglang_operator_cases/newtest_cases/test_tanh.py +# Main kernel: tanh_kernel +# PT file: test_tanh_kernel_v2.pt + +import triton +import triton.language as tl + + +# === tanh_kernel === +@triton.jit +def tanh_kernel( + x_ptr, + y_ptr, + n_elements, + BLOCK_SIZE: tl.constexpr +): + pid = tl.program_id(axis=0) + block_start = pid * BLOCK_SIZE + offsets = block_start + tl.arange(0, BLOCK_SIZE) + mask = offsets < n_elements + x = tl.load(x_ptr + offsets, mask=mask, other=0.0) + y = 2 * tl.sigmoid(2 * x) - 1 + tl.store(y_ptr + offsets, y, mask=mask) + diff --git a/benchmarks/TritonNPUBench/sglang/track_mamba_state_if_needed_kernel/track_mamba_state_if_needed_kernel.py b/benchmarks/TritonNPUBench/sglang/track_mamba_state_if_needed_kernel/track_mamba_state_if_needed_kernel.py new file mode 100644 index 00000000..d4500ef3 --- /dev/null +++ b/benchmarks/TritonNPUBench/sglang/track_mamba_state_if_needed_kernel/track_mamba_state_if_needed_kernel.py @@ -0,0 +1,69 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/sglang_operator_cases/newtest_cases/test_track_mamba_state_if_needed_kernel.py +# Main kernel: track_mamba_state_if_needed_kernel +# PT file: test_track_mamba_state_if_needed_kernel_v2.pt + +import triton +import triton.language as tl + + +# === track_mamba_state_if_needed_kernel === +@triton.jit +def track_mamba_state_if_needed_kernel( + conv_states_ptr, + ssm_states_ptr, + cache_indices_ptr, + mamba_track_mask_ptr, + mamba_track_indices_ptr, + conv_state_stride_0, # stride for first dimension (batch/pool index) + ssm_state_stride_0, # stride for first dimension (batch/pool index) + conv_state_numel_per_row: tl.constexpr, # total elements per row + ssm_state_numel_per_row: tl.constexpr, # total elements per row + BLOCK_SIZE: tl.constexpr, +): + """ + Track conv_states and ssm_states rows based on track mask. + + This kernel replaces a Python loop that copies state tensors for mamba attention. + For each batch element, if the track mask is True, it copies the entire row from + the source index (cache_indices[i]) to the destination index (mamba_track_indices[i]). + + Grid: (batch_size,) + Each block handles one batch element, using multiple threads to copy data in parallel. + """ + batch_idx = tl.program_id(0) + + # Load the copy mask for this batch element + track_mask = tl.load(mamba_track_mask_ptr + batch_idx) + + # Early exit if we don't need to track + # if not track_mask: + # return + + # Load source and destination indices + src_idx = tl.load(cache_indices_ptr + batch_idx) + dst_idx = tl.load(mamba_track_indices_ptr + batch_idx) + + # Copy conv_states + # Each thread handles BLOCK_SIZE elements + for offset in range(0, conv_state_numel_per_row, BLOCK_SIZE): + element_indices = offset + tl.arange(0, BLOCK_SIZE) + mask = element_indices < conv_state_numel_per_row + + src_ptr = conv_states_ptr + src_idx * conv_state_stride_0 + element_indices + dst_ptr = conv_states_ptr + dst_idx * conv_state_stride_0 + element_indices + + data = tl.load(src_ptr, mask=mask, other=0.0) + tl.store(dst_ptr, data, mask=mask) + + # Copy ssm_states + for offset in range(0, ssm_state_numel_per_row, BLOCK_SIZE): + element_indices = offset + tl.arange(0, BLOCK_SIZE) + mask = element_indices < ssm_state_numel_per_row + + src_ptr = ssm_states_ptr + src_idx * ssm_state_stride_0 + element_indices + dst_ptr = ssm_states_ptr + dst_idx * ssm_state_stride_0 + element_indices + + data = tl.load(src_ptr, mask=mask, other=0.0) + tl.store(dst_ptr, data, mask=mask) + diff --git a/benchmarks/TritonNPUBench/sglang/transform_index_page_table_decode_kernel/transform_index_page_table_decode_kernel.py b/benchmarks/TritonNPUBench/sglang/transform_index_page_table_decode_kernel/transform_index_page_table_decode_kernel.py new file mode 100644 index 00000000..6ff6f8b0 --- /dev/null +++ b/benchmarks/TritonNPUBench/sglang/transform_index_page_table_decode_kernel/transform_index_page_table_decode_kernel.py @@ -0,0 +1,31 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/sglang_operator_cases/newtest_cases/test_transform_index_page_table_decode_kernel.py +# Main kernel: transform_index_page_table_decode_kernel +# PT file: test_transform_index_page_table_decode_kernel_v2.pt + +import triton +import triton.language as tl + + +# === transform_index_page_table_decode_kernel === +@triton.jit +def transform_index_page_table_decode_kernel( + page_table_ptr: torch.Tensor, + topk_indices_ptr: torch.Tensor, + result_ptr: torch.Tensor, + page_size: tl.constexpr, + max_seqlen_k: tl.constexpr, +): + TOPK: tl.constexpr = 2048 + req_id = tl.program_id(0) + page_table_ptr = page_table_ptr + req_id * max_seqlen_k + topk_indices_ptr = topk_indices_ptr + req_id * TOPK + result_ptr = result_ptr + req_id * TOPK + + offset = tl.arange(0, TOPK) # topk should be 2048 + loaded_topk_indices = tl.load(topk_indices_ptr + offset) + mask = loaded_topk_indices >= 0 + loaded_kv_indices = tl.load(page_table_ptr + loaded_topk_indices, mask=mask) + tl.store(result_ptr + offset, loaded_kv_indices, mask=mask) + tl.store(result_ptr + offset, -1, mask=~mask) + diff --git a/benchmarks/TritonNPUBench/sglang/unpad_draft_extend_output_kernel/unpad_draft_extend_output_kernel.py b/benchmarks/TritonNPUBench/sglang/unpad_draft_extend_output_kernel/unpad_draft_extend_output_kernel.py new file mode 100644 index 00000000..1f9b40d6 --- /dev/null +++ b/benchmarks/TritonNPUBench/sglang/unpad_draft_extend_output_kernel/unpad_draft_extend_output_kernel.py @@ -0,0 +1,80 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/sglang_operator_cases/newtest_cases/test_unpad_draft_extend_output_kernel.py +# Main kernel: unpad_draft_extend_output_kernel +# PT file: test_unpad_draft_extend_output_kernel_v2.pt + +import triton +import triton.language as tl + + +# === unpad_draft_extend_output_kernel === +@triton.jit +def unpad_draft_extend_output_kernel( + raw_out_ptr, # Input raw output tensor (batch_size, token_per_batch, tp_q_head_num, v_head_dim) + output_ptr, # Output tensor (-1, tp_q_head_num, v_head_dim) + accept_length_ptr, # Accept lengths for each sequence [batch_size] + cumsum_ptr, # Cumulative sum of accept lengths [batch_size + 1] + batch_size, + token_per_batch, + tp_q_head_num, + v_head_dim, + BLOCK_SIZE: tl.constexpr, +): + """Triton kernel for unpadding draft extended output tensor with parallelized head and dim processing.""" + batch_seq_pid = tl.program_id(0) + head_pid = tl.program_id(1) + dim_pid = tl.program_id(2) + + batch_id = batch_seq_pid // token_per_batch + seq_pos = batch_seq_pid % token_per_batch + + if batch_id >= batch_size: + return + + # Load accept length for this batch + accept_len = tl.load(accept_length_ptr + batch_id) + + if seq_pos >= accept_len: + return + + # Load cumulative sum to get start position in output tensor + output_start = tl.load(cumsum_ptr + batch_id) + output_pos = output_start + seq_pos + + # Calculate head and dim block ranges + head_start = head_pid * BLOCK_SIZE + head_end = tl.minimum(head_start + BLOCK_SIZE, tp_q_head_num) + head_mask = tl.arange(0, BLOCK_SIZE) < (head_end - head_start) + + dim_start = dim_pid * BLOCK_SIZE + dim_end = tl.minimum(dim_start + BLOCK_SIZE, v_head_dim) + dim_mask = tl.arange(0, BLOCK_SIZE) < (dim_end - dim_start) + + # Calculate input offset: (batch_id, seq_pos, head_id, dim_id) + input_offset = ( + batch_id * token_per_batch * tp_q_head_num * v_head_dim + + seq_pos * tp_q_head_num * v_head_dim + + (head_start + tl.arange(0, BLOCK_SIZE))[:, None] * v_head_dim + + (dim_start + tl.arange(0, BLOCK_SIZE))[None, :] + ) + + # Load data + data = tl.load( + raw_out_ptr + input_offset, + mask=head_mask[:, None] & dim_mask[None, :], + other=0.0, + ) + + output_offset = ( + output_pos * tp_q_head_num * v_head_dim + + (head_start + tl.arange(0, BLOCK_SIZE))[:, None] * v_head_dim + + (dim_start + tl.arange(0, BLOCK_SIZE))[None, :] + ) + + # Store data + tl.store( + output_ptr + output_offset, + data, + mask=head_mask[:, None] & dim_mask[None, :], + ) + diff --git a/benchmarks/TritonNPUBench/sglang/write_req_to_token_pool_triton/write_req_to_token_pool_triton.py b/benchmarks/TritonNPUBench/sglang/write_req_to_token_pool_triton/write_req_to_token_pool_triton.py new file mode 100644 index 00000000..37d714e4 --- /dev/null +++ b/benchmarks/TritonNPUBench/sglang/write_req_to_token_pool_triton/write_req_to_token_pool_triton.py @@ -0,0 +1,46 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/sglang_operator_cases/newtest_cases/test_write_req_to_token_pool_triton.py +# Main kernel: write_req_to_token_pool_triton +# PT file: test_write_req_to_token_pool_triton_unit_v2.pt + +import triton +import triton.language as tl + + +# === write_req_to_token_pool_triton === +@triton.jit +def write_req_to_token_pool_triton( + req_to_token_ptr, # [max_batch, max_context_len] + req_pool_indices, + pre_lens, + seq_lens, + extend_lens, + out_cache_loc, + req_to_token_ptr_stride: tl.constexpr, +): + BLOCK_SIZE: tl.constexpr = 512 + pid = tl.program_id(0) + + req_pool_index = tl.load(req_pool_indices + pid) + pre_len = tl.load(pre_lens + pid) + seq_len = tl.load(seq_lens + pid) + + # TODO: optimize this? + cumsum_start = 0 + for i in range(pid): + cumsum_start += tl.load(extend_lens + i) + + num_loop = tl.cdiv(seq_len - pre_len, BLOCK_SIZE) + for i in range(num_loop): + offset = tl.arange(0, BLOCK_SIZE) + i * BLOCK_SIZE + mask = offset < (seq_len - pre_len) + value = tl.load(out_cache_loc + cumsum_start + offset, mask=mask) + tl.store( + req_to_token_ptr + + req_pool_index * req_to_token_ptr_stride + + offset + + pre_len, + value, + mask=mask, + ) + diff --git a/benchmarks/TritonNPUBench/sglang/write_req_to_token_pool_triton_optimize/write_req_to_token_pool_triton_optimize.py b/benchmarks/TritonNPUBench/sglang/write_req_to_token_pool_triton_optimize/write_req_to_token_pool_triton_optimize.py new file mode 100644 index 00000000..f32b0b5a --- /dev/null +++ b/benchmarks/TritonNPUBench/sglang/write_req_to_token_pool_triton_optimize/write_req_to_token_pool_triton_optimize.py @@ -0,0 +1,52 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/sglang_operator_cases/newtest_cases/test_write_req_to_token_pool_triton_optimize.py +# Main kernel: write_req_to_token_pool_triton_optimize +# PT file: test_write_req_to_token_pool_triton_optimize_v2.pt + +import triton +import triton.language as tl + + +# === write_req_to_token_pool_triton_optimize === +@triton.jit +def write_req_to_token_pool_triton_optimize( + req_to_token_ptr, # [max_batch, max_context_len] + req_pool_indices, + pre_lens, + seq_lens, + extend_lens, + out_cache_loc, + req_to_token_ptr_stride: tl.constexpr, + BLOCK_SIZE: tl.constexpr, +): + pid_batch = tl.program_id(0) + pid_token = tl.program_id(1) + + req_pool_index = tl.load(req_pool_indices + pid_batch) + pre_len = tl.load(pre_lens + pid_batch) + seq_len = tl.load(seq_lens + pid_batch) + extend_len = seq_len - pre_len + + cumsum_start = 0 + for i in range(pid_batch): + cumsum_start += tl.load(extend_lens + i) + + token_start = pid_token * BLOCK_SIZE + + offset = tl.arange(0, BLOCK_SIZE) + actual_offset = token_start + offset + mask = actual_offset < extend_len + + src_ptr = out_cache_loc + cumsum_start + actual_offset + src_ptr = tl.max_contiguous(tl.multiple_of(src_ptr, BLOCK_SIZE), BLOCK_SIZE) + value = tl.load(src_ptr, mask=mask) + dst_ptr = ( + req_to_token_ptr + + req_pool_index * req_to_token_ptr_stride + + actual_offset + + pre_len + ) + dst_ptr = tl.max_contiguous(tl.multiple_of(dst_ptr, BLOCK_SIZE), BLOCK_SIZE) + + tl.store(dst_ptr, value, mask=mask) + From 5002512a90af85451f7cfc89cb603b3b6cf325df Mon Sep 17 00:00:00 2001 From: zhumingming Date: Fri, 17 Apr 2026 21:36:41 +0800 Subject: [PATCH 3/3] =?UTF-8?q?=E3=80=90benchmark=E3=80=91=E6=96=B0?= =?UTF-8?q?=E5=A2=9Efbgemm=E4=BB=9332=E4=B8=AA=E7=AE=97=E5=AD=90=E6=8F=8F?= =?UTF-8?q?=E8=BF=B0?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../_bwd_preprocess_do_o_dot.py | 49 ++++++ .../_fbgemm_gather_scale_dense_tokens.py | 67 ++++++++ ...er_scale_fp8_rowwise_quant_dense_tokens.py | 110 +++++++++++++ .../_fbgemm_scatter_add_dense_tokens.py | 49 ++++++ .../_fbgemm_scatter_add_padded_tokens.py | 74 +++++++++ .../_fbgemm_silu_mul/_fbgemm_silu_mul.py | 53 ++++++ .../_fbgemm_silu_mul_quant.py | 85 ++++++++++ .../fbgemm/_floor_log2/_floor_log2.py | 32 ++++ ...h_attention_bwd_preprocess_basic_kernel.py | 52 ++++++ .../_kernel_dequantize_fp8_block.py | 43 +++++ .../_kernel_dequantize_fp8_packed_row.py | 62 +++++++ .../_kernel_dequantize_fp8_row.py | 59 +++++++ .../_kernel_quantize_fp8_block.py | 96 +++++++++++ .../_kernel_quantize_fp8_group.py | 152 ++++++++++++++++++ .../_kernel_scale_fp8_row.py | 59 +++++++ ...d_flash_attention_bwd_preprocess_kernel.py | 65 ++++++++ .../fbgemm/cumsum_kernel/cumsum_kernel.py | 37 +++++ .../dense_jagged_cat_jagged_out_kernel.py | 38 +++++ .../fused_single_block_kernel.py | 107 ++++++++++++ .../jagged2_to_padded_dense_kernel.py | 50 ++++++ .../jagged_2_softmax_backward_kernel.py | 91 +++++++++++ .../jagged_2_softmax_kernel.py | 97 +++++++++++ ...dense_elementwise_mul_jagged_out_kernel.py | 66 ++++++++ ...agged_jagged_elementwise_arithmetic_ops.py | 60 +++++++ ...ged_self_substraction_jagged_out_kernel.py | 42 +++++ .../jagged_softmax_backward_kernel.py | 65 ++++++++ .../jagged_softmax_kernel.py | 57 +++++++ .../padded_dense_to_jagged2_kernel.py | 51 ++++++ ...iton_batched_dense_vec_jagged_2d_matmul.py | 74 +++++++++ .../triton_dense_to_jagged.py | 101 ++++++++++++ .../triton_jagged_to_dense.py | 113 +++++++++++++ .../triton_jagged_to_dense_optimization_2d.py | 100 ++++++++++++ 32 files changed, 2256 insertions(+) create mode 100644 benchmarks/TritonNPUBench/fbgemm/_bwd_preprocess_do_o_dot/_bwd_preprocess_do_o_dot.py create mode 100644 benchmarks/TritonNPUBench/fbgemm/_fbgemm_gather_scale_dense_tokens/_fbgemm_gather_scale_dense_tokens.py create mode 100644 benchmarks/TritonNPUBench/fbgemm/_fbgemm_gather_scale_fp8_rowwise_quant_dense_tokens/_fbgemm_gather_scale_fp8_rowwise_quant_dense_tokens.py create mode 100644 benchmarks/TritonNPUBench/fbgemm/_fbgemm_scatter_add_dense_tokens/_fbgemm_scatter_add_dense_tokens.py create mode 100644 benchmarks/TritonNPUBench/fbgemm/_fbgemm_scatter_add_padded_tokens/_fbgemm_scatter_add_padded_tokens.py create mode 100644 benchmarks/TritonNPUBench/fbgemm/_fbgemm_silu_mul/_fbgemm_silu_mul.py create mode 100644 benchmarks/TritonNPUBench/fbgemm/_fbgemm_silu_mul_quant/_fbgemm_silu_mul_quant.py create mode 100644 benchmarks/TritonNPUBench/fbgemm/_floor_log2/_floor_log2.py create mode 100644 benchmarks/TritonNPUBench/fbgemm/_jagged_flash_attention_bwd_preprocess_basic_kernel/_jagged_flash_attention_bwd_preprocess_basic_kernel.py create mode 100644 benchmarks/TritonNPUBench/fbgemm/_kernel_dequantize_fp8_block/_kernel_dequantize_fp8_block.py create mode 100644 benchmarks/TritonNPUBench/fbgemm/_kernel_dequantize_fp8_packed_row/_kernel_dequantize_fp8_packed_row.py create mode 100644 benchmarks/TritonNPUBench/fbgemm/_kernel_dequantize_fp8_row/_kernel_dequantize_fp8_row.py create mode 100644 benchmarks/TritonNPUBench/fbgemm/_kernel_quantize_fp8_block/_kernel_quantize_fp8_block.py create mode 100644 benchmarks/TritonNPUBench/fbgemm/_kernel_quantize_fp8_group/_kernel_quantize_fp8_group.py create mode 100644 benchmarks/TritonNPUBench/fbgemm/_kernel_scale_fp8_row/_kernel_scale_fp8_row.py create mode 100644 benchmarks/TritonNPUBench/fbgemm/_multi_head_jagged_flash_attention_bwd_preprocess_kernel/_multi_head_jagged_flash_attention_bwd_preprocess_kernel.py create mode 100644 benchmarks/TritonNPUBench/fbgemm/cumsum_kernel/cumsum_kernel.py create mode 100644 benchmarks/TritonNPUBench/fbgemm/dense_jagged_cat_jagged_out_kernel/dense_jagged_cat_jagged_out_kernel.py create mode 100644 benchmarks/TritonNPUBench/fbgemm/fused_single_block_kernel/fused_single_block_kernel.py create mode 100644 benchmarks/TritonNPUBench/fbgemm/jagged2_to_padded_dense_kernel/jagged2_to_padded_dense_kernel.py create mode 100644 benchmarks/TritonNPUBench/fbgemm/jagged_2_softmax_backward_kernel/jagged_2_softmax_backward_kernel.py create mode 100644 benchmarks/TritonNPUBench/fbgemm/jagged_2_softmax_kernel/jagged_2_softmax_kernel.py create mode 100644 benchmarks/TritonNPUBench/fbgemm/jagged_dense_elementwise_mul_jagged_out_kernel/jagged_dense_elementwise_mul_jagged_out_kernel.py create mode 100644 benchmarks/TritonNPUBench/fbgemm/jagged_jagged_elementwise_arithmetic_ops/jagged_jagged_elementwise_arithmetic_ops.py create mode 100644 benchmarks/TritonNPUBench/fbgemm/jagged_self_substraction_jagged_out_kernel/jagged_self_substraction_jagged_out_kernel.py create mode 100644 benchmarks/TritonNPUBench/fbgemm/jagged_softmax_backward_kernel/jagged_softmax_backward_kernel.py create mode 100644 benchmarks/TritonNPUBench/fbgemm/jagged_softmax_kernel/jagged_softmax_kernel.py create mode 100644 benchmarks/TritonNPUBench/fbgemm/padded_dense_to_jagged2_kernel/padded_dense_to_jagged2_kernel.py create mode 100644 benchmarks/TritonNPUBench/fbgemm/triton_batched_dense_vec_jagged_2d_matmul/triton_batched_dense_vec_jagged_2d_matmul.py create mode 100644 benchmarks/TritonNPUBench/fbgemm/triton_dense_to_jagged/triton_dense_to_jagged.py create mode 100644 benchmarks/TritonNPUBench/fbgemm/triton_jagged_to_dense/triton_jagged_to_dense.py create mode 100644 benchmarks/TritonNPUBench/fbgemm/triton_jagged_to_dense_optimization_2d/triton_jagged_to_dense_optimization_2d.py diff --git a/benchmarks/TritonNPUBench/fbgemm/_bwd_preprocess_do_o_dot/_bwd_preprocess_do_o_dot.py b/benchmarks/TritonNPUBench/fbgemm/_bwd_preprocess_do_o_dot/_bwd_preprocess_do_o_dot.py new file mode 100644 index 00000000..3e7311d6 --- /dev/null +++ b/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)) + diff --git a/benchmarks/TritonNPUBench/fbgemm/_fbgemm_gather_scale_dense_tokens/_fbgemm_gather_scale_dense_tokens.py b/benchmarks/TritonNPUBench/fbgemm/_fbgemm_gather_scale_dense_tokens/_fbgemm_gather_scale_dense_tokens.py new file mode 100644 index 00000000..5ec75c56 --- /dev/null +++ b/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 + diff --git a/benchmarks/TritonNPUBench/fbgemm/_fbgemm_gather_scale_fp8_rowwise_quant_dense_tokens/_fbgemm_gather_scale_fp8_rowwise_quant_dense_tokens.py b/benchmarks/TritonNPUBench/fbgemm/_fbgemm_gather_scale_fp8_rowwise_quant_dense_tokens/_fbgemm_gather_scale_fp8_rowwise_quant_dense_tokens.py new file mode 100644 index 00000000..8681630e --- /dev/null +++ b/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 + diff --git a/benchmarks/TritonNPUBench/fbgemm/_fbgemm_scatter_add_dense_tokens/_fbgemm_scatter_add_dense_tokens.py b/benchmarks/TritonNPUBench/fbgemm/_fbgemm_scatter_add_dense_tokens/_fbgemm_scatter_add_dense_tokens.py new file mode 100644 index 00000000..2bce192b --- /dev/null +++ b/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 + diff --git a/benchmarks/TritonNPUBench/fbgemm/_fbgemm_scatter_add_padded_tokens/_fbgemm_scatter_add_padded_tokens.py b/benchmarks/TritonNPUBench/fbgemm/_fbgemm_scatter_add_padded_tokens/_fbgemm_scatter_add_padded_tokens.py new file mode 100644 index 00000000..72f3a6b8 --- /dev/null +++ b/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 + diff --git a/benchmarks/TritonNPUBench/fbgemm/_fbgemm_silu_mul/_fbgemm_silu_mul.py b/benchmarks/TritonNPUBench/fbgemm/_fbgemm_silu_mul/_fbgemm_silu_mul.py new file mode 100644 index 00000000..86d12432 --- /dev/null +++ b/benchmarks/TritonNPUBench/fbgemm/_fbgemm_silu_mul/_fbgemm_silu_mul.py @@ -0,0 +1,53 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/fbgemm_operator_cases/newtest_cases/test__fbgemm_silu_mul_npu.py +# Main kernel: _fbgemm_silu_mul +# PT file: test__fbgemm_silu_mul_v2.pt + +import triton +import triton.language as tl + + +# === _fbgemm_silu_mul === +@triton.jit +def _fbgemm_silu_mul( + y_ptr, + x0_ptr, + x1_ptr, + stride_0, + stride_1, + valid_token_count, + D: tl.constexpr, + BLOCK_D_OUTER: tl.constexpr, + BLOCK_D_INNER: tl.constexpr, +) -> None: + token_index = tl.program_id(0) + 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 token_index >= valid_token_count: + return + + for _ in tl.range(0, BLOCK_D_OUTER // BLOCK_D_INNER, num_stages=3): + x0 = tl.load( + x0_ptr + token_index * stride_0 + feature_offset, + None, + eviction_policy="evict_first", + ).to(tl.float32) + x1 = tl.load( + x1_ptr + token_index * stride_1 + feature_offset, + None, + eviction_policy="evict_first", + ).to(tl.float32) + + y = x0 * tl.sigmoid(x0) * x1 + + tl.store( + y_ptr + token_index * D + feature_offset, + y, + None, + ) + feature_offset += BLOCK_D_INNER + diff --git a/benchmarks/TritonNPUBench/fbgemm/_fbgemm_silu_mul_quant/_fbgemm_silu_mul_quant.py b/benchmarks/TritonNPUBench/fbgemm/_fbgemm_silu_mul_quant/_fbgemm_silu_mul_quant.py new file mode 100644 index 00000000..a049f778 --- /dev/null +++ b/benchmarks/TritonNPUBench/fbgemm/_fbgemm_silu_mul_quant/_fbgemm_silu_mul_quant.py @@ -0,0 +1,85 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/fbgemm_operator_cases/newtest_cases/test__fbgemm_silu_mul_quant_npu.py +# Main kernel: _fbgemm_silu_mul_quant +# PT file: test__fbgemm_silu_mul_quant_v2.pt + +import triton +import triton.language as tl + + +# === _fbgemm_silu_mul_quant === +@triton.jit +def _fbgemm_silu_mul_quant( + y_ptr, + y_inv_scale_ptr, + x0_ptr, + x1_ptr, + scale_ub_ptr, + stride_0, + stride_1, + valid_token_count, + T, + D: tl.constexpr, + BLOCK_T: tl.constexpr, + TL_FP8_DTYPE: tl.constexpr, + MAX_FP8: tl.constexpr, + EPS: tl.constexpr, + CLAMP_MAX: tl.constexpr, +) -> None: + PADDED_D: tl.constexpr = triton.next_power_of_2(D) # pyre-ignore + + tidx = tl.program_id(0) + start_idx = tidx * BLOCK_T + end_idx = tl.minimum(start_idx + BLOCK_T, T) + + if valid_token_count is not None: + valid_token_count = tl.load( + valid_token_count, None, eviction_policy="evict_last" + ) + if start_idx >= valid_token_count: + return + + offsets = tl.arange(0, PADDED_D)[:] + mask = offsets < D + + if CLAMP_MAX: + ub = tl.load(scale_ub_ptr, eviction_policy="evict_last") + else: + ub = float("inf") + + for token_index in tl.range(start_idx, end_idx, 1, num_stages=2): + x0 = tl.load( + x0_ptr + token_index * stride_0 + offsets, + mask, + eviction_policy="evict_first", + ).to(tl.float32) + x1 = tl.load( + x1_ptr + token_index * stride_1 + offsets, + mask, + eviction_policy="evict_first", + ).to(tl.float32) + + y = x0 * tl.sigmoid(x0) * x1 + + # Masked values are set to 0.0. + row_max = tl.max(tl.where(mask, tl.abs(y), 0.0)) + if CLAMP_MAX: + row_max = tl.clamp(row_max, EPS, ub) + else: + row_max = tl.maximum(row_max, EPS) + + y_scale = MAX_FP8 / row_max + tl.store(y_inv_scale_ptr + token_index, 1.0 / y_scale) + + y = y * y_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. + y_fp8 = tl.clamp(y, -MAX_FP8, MAX_FP8).to(TL_FP8_DTYPE) + + tl.store( + y_ptr + token_index * D + offsets, + y_fp8, + mask, + ) + diff --git a/benchmarks/TritonNPUBench/fbgemm/_floor_log2/_floor_log2.py b/benchmarks/TritonNPUBench/fbgemm/_floor_log2/_floor_log2.py new file mode 100644 index 00000000..2b843adf --- /dev/null +++ b/benchmarks/TritonNPUBench/fbgemm/_floor_log2/_floor_log2.py @@ -0,0 +1,32 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/fbgemm_operator_cases/newtest_cases/test__floor_log2_npu.py +# Main kernel: _floor_log2 +# PT file: test__floor_log2_v2.pt + +import triton +import triton.language as tl + + +# === _floor_log2 === +@triton.jit +def _floor_log2(x): + """Helper function to efficiently compute floor(log2(x)) + + Args: + x (Tensor): FP32 Input tensor to operate on. + + Returns: + Tensor: Floor of log2(x). + """ + # Helpful bit constants. + FP32_EXP_MASK: tl.constexpr = 0x7F800000 # type: ignore[Incompatible variable type] + FP32_EXP_OFFSET: tl.constexpr = 23 # type: ignore[Incompatible variable type] + FP32_EXP_BIAS: tl.constexpr = 127 # type: ignore[Incompatible variable type] + + # View x as an integer and extract its exponent. + x = x.to(tl.int32, bitcast=True) & FP32_EXP_MASK + # Shift exponent down to bottom bits. + x = x >> FP32_EXP_OFFSET + # Remove FP32 exponent bias and return. + return (x - FP32_EXP_BIAS).to(tl.float32) + diff --git a/benchmarks/TritonNPUBench/fbgemm/_jagged_flash_attention_bwd_preprocess_basic_kernel/_jagged_flash_attention_bwd_preprocess_basic_kernel.py b/benchmarks/TritonNPUBench/fbgemm/_jagged_flash_attention_bwd_preprocess_basic_kernel/_jagged_flash_attention_bwd_preprocess_basic_kernel.py new file mode 100644 index 00000000..fe788c34 --- /dev/null +++ b/benchmarks/TritonNPUBench/fbgemm/_jagged_flash_attention_bwd_preprocess_basic_kernel/_jagged_flash_attention_bwd_preprocess_basic_kernel.py @@ -0,0 +1,52 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/fbgemm_operator_cases/newtest_cases/test__jagged_flash_attention_bwd_preprocess_basic_kernel_npu.py +# Main kernel: _jagged_flash_attention_bwd_preprocess_basic_kernel +# PT file: _jagged_flash_attention_bwd_preprocess_basic_kernel_v2.pt + +import triton +import triton.language as tl + + +# === _jagged_flash_attention_bwd_preprocess_basic_kernel === +@triton.jit +def _jagged_flash_attention_bwd_preprocess_basic_kernel( + o_ptr, + o_offset_ptr, + do_ptr, + delta_ptr, + stride_om, + stride_od, + max_seq_len, + D: tl.constexpr, + BLOCK_SIZE_M: tl.constexpr, + BLOCK_SIZE_D: tl.constexpr, +): + pid_m = tl.program_id(axis=0) + pid_batch = tl.program_id(axis=1) + + begin_o = tl.load(o_offset_ptr + pid_batch) + end_o = tl.load(o_offset_ptr + pid_batch + 1) + + M = end_o - begin_o + M = tl.minimum(M, max_seq_len) + + offs_om = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) + offs_od = tl.arange(0, BLOCK_SIZE_D) + + o_offsets = ( + offs_om[:, None] * stride_om + + offs_od[None, :] * stride_od + + begin_o * stride_om + ) + o_ptrs = o_ptr + o_offsets + do_ptrs = do_ptr + o_offsets + o_mask = (offs_om[:, None] < M) & (offs_od[None, :] < D) + + # Load O + o = tl.load(o_ptrs, mask=o_mask) + do = tl.load(do_ptrs, mask=o_mask) + + delta = tl.sum(o * do, axis=1) + + tl.store(delta_ptr + begin_o + offs_om, delta, mask=offs_om < M) + diff --git a/benchmarks/TritonNPUBench/fbgemm/_kernel_dequantize_fp8_block/_kernel_dequantize_fp8_block.py b/benchmarks/TritonNPUBench/fbgemm/_kernel_dequantize_fp8_block/_kernel_dequantize_fp8_block.py new file mode 100644 index 00000000..19d9348d --- /dev/null +++ b/benchmarks/TritonNPUBench/fbgemm/_kernel_dequantize_fp8_block/_kernel_dequantize_fp8_block.py @@ -0,0 +1,43 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/fbgemm_operator_cases/newtest_cases/test__kernel_dequantize_fp8_block_npu.py +# Main kernel: _kernel_dequantize_fp8_block +# PT file: test__kernel_dequantize_fp8_block_v2.pt + +import triton +import triton.language as tl + + +# === _kernel_dequantize_fp8_block === +@triton.jit +def _kernel_dequantize_fp8_block( + xq_ptr, + x_scale_ptr, + x_dequant_ptr, + M, + K, + BLOCK_M: tl.constexpr, + BLOCK_K: tl.constexpr, +): + """ + Kernel to dequantize FP8 tensor to BF16 tensor. + Args: + xq_ptr (tl.constexpr): Pointer to FP8 tensor. + x_scale_ptr (tl.constexpr): Pointer to FP8 scale tensor. + x_dequant_ptr (tl.constexpr): Pointer to BF16 tensor. + M (tl.constexpr): M dimension of input tensor. + K (tl.constexpr): K dimension of input tensor. + BLOCK_M (tl.constexpr): Block size for the M dimension. + BLOCK_K (tl.constexpr): Block size for the K dimension. + """ + pid_m = tl.program_id(axis=0) + pid_k = tl.program_id(axis=1) + k = tl.cdiv(K, BLOCK_K) + offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) + offs_k = pid_k * BLOCK_K + tl.arange(0, BLOCK_K) + offs = offs_m[:, None] * K + offs_k[None, :] + mask = (offs_m[:, None] < M) & (offs_k[None, :] < K) + xq = tl.load(xq_ptr + offs, mask=mask).to(tl.bfloat16) + x_scale = tl.load(x_scale_ptr + pid_m * k + pid_k) + x_dequant = xq * x_scale + tl.store(x_dequant_ptr + offs, x_dequant, mask=mask) + diff --git a/benchmarks/TritonNPUBench/fbgemm/_kernel_dequantize_fp8_packed_row/_kernel_dequantize_fp8_packed_row.py b/benchmarks/TritonNPUBench/fbgemm/_kernel_dequantize_fp8_packed_row/_kernel_dequantize_fp8_packed_row.py new file mode 100644 index 00000000..134148ff --- /dev/null +++ b/benchmarks/TritonNPUBench/fbgemm/_kernel_dequantize_fp8_packed_row/_kernel_dequantize_fp8_packed_row.py @@ -0,0 +1,62 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/fbgemm_operator_cases/newtest_cases/test__kernel_dequantize_fp8_packed_row_npu.py +# Main kernel: _kernel_dequantize_fp8_packed_row +# PT file: test__kernel_dequantize_fp8_packed_row_v2.pt + +import triton +import triton.language as tl + + +# === _kernel_dequantize_fp8_packed_row === +@triton.jit +def _kernel_dequantize_fp8_packed_row( + xq_ptr, + x_scale_ptr, + x_dequant_ptr, + M, + K, + stride_xm, + stride_xk, + stride_xdqm, + stride_xdqk, + BLOCK_M: tl.constexpr, + BLOCK_K: tl.constexpr, + NUM_STAGES: tl.constexpr, + USE_INT64: tl.constexpr, +): + """ + Kernel to dequantize FP8 tensor to BF16 tensor. + Args: + xq_ptr (tl.constexpr): Pointer to FP8 tensor. + x_scale_ptr (tl.constexpr): Pointer to FP8 scale tensor. + x_dequant_ptr (tl.constexpr): Pointer to BF16 tensor. + M (tl.constexpr): M dimension of input tensor. + K (tl.constexpr): K dimension of input tensor (along which scales are applied) + BLOCK_SIZE (tl.constexpr): Block size for the K dimension. + """ + pid = tl.program_id(axis=0) + if USE_INT64: + pid = pid.to(tl.int64) + offs_m = pid * BLOCK_M + tl.arange(0, BLOCK_M) + offs_k = tl.arange(0, BLOCK_K) + scales = tl.load(x_scale_ptr + offs_m) + + for _k in tl.range(0, tl.cdiv(K, BLOCK_K), num_stages=NUM_STAGES): + mask = (offs_m[:, None] < M) & (offs_k[None, :] < K) + + xq = tl.load( + xq_ptr + offs_m[:, None] * stride_xm + offs_k[None, :] * stride_xk, + mask=mask, + other=0.0, + ) + x_dq = xq * scales[:, None] + + tl.store( + x_dequant_ptr + + offs_m[:, None] * stride_xdqm + + offs_k[None, :] * stride_xdqk, + x_dq, + mask=mask, + ) + offs_k += BLOCK_K + diff --git a/benchmarks/TritonNPUBench/fbgemm/_kernel_dequantize_fp8_row/_kernel_dequantize_fp8_row.py b/benchmarks/TritonNPUBench/fbgemm/_kernel_dequantize_fp8_row/_kernel_dequantize_fp8_row.py new file mode 100644 index 00000000..949b96bf --- /dev/null +++ b/benchmarks/TritonNPUBench/fbgemm/_kernel_dequantize_fp8_row/_kernel_dequantize_fp8_row.py @@ -0,0 +1,59 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/fbgemm_operator_cases/newtest_cases/test__kernel_dequantize_fp8_row_npu.py +# Main kernel: _kernel_dequantize_fp8_row +# PT file: test__kernel_dequantize_fp8_row_v2.pt + +import triton +import triton.language as tl + + +# === _kernel_dequantize_fp8_row === +@triton.jit +def _kernel_dequantize_fp8_row( + xq_ptr, + x_scale_ptr, + x_dequant_ptr, + M, + K, + stride_xm, + stride_xk, + stride_xdqm, + stride_xdqk, + BLOCK_M: tl.constexpr, + BLOCK_K: tl.constexpr, + NUM_STAGES: tl.constexpr, + USE_INT64: tl.constexpr, +): + """ + Kernel to dequantize FP8 tensor to BF16 tensor. + Args: + xq_ptr (tl.constexpr): Pointer to FP8 tensor. + x_scale_ptr (tl.constexpr): Pointer to FP8 scale tensor. + x_dequant_ptr (tl.constexpr): Pointer to BF16 tensor. + M (tl.constexpr): M dimension of input tensor. + K (tl.constexpr): K dimension of input tensor (along which scales are applied) + BLOCK_SIZE (tl.constexpr): Block size for the K dimension. + """ + pid = tl.program_id(axis=0) + if USE_INT64: + pid = pid.to(tl.int64) + offs_m = pid * BLOCK_M + tl.arange(0, BLOCK_M) + offs_k = tl.arange(0, BLOCK_K) + scales = tl.load(x_scale_ptr + offs_m) + + for _k in tl.range(0, tl.cdiv(K, BLOCK_K), num_stages=NUM_STAGES): + mask = (offs_m[:, None] < M) & (offs_k[None, :] < K) + xq = tl.load( + xq_ptr + offs_m[:, None] * stride_xm + offs_k[None, :] * stride_xk, + mask=mask, + ) + x_dq = xq * scales[:, None] + tl.store( + x_dequant_ptr + + offs_m[:, None] * stride_xdqm + + offs_k[None, :] * stride_xdqk, + x_dq, + mask=mask, + ) + offs_k += BLOCK_K + diff --git a/benchmarks/TritonNPUBench/fbgemm/_kernel_quantize_fp8_block/_kernel_quantize_fp8_block.py b/benchmarks/TritonNPUBench/fbgemm/_kernel_quantize_fp8_block/_kernel_quantize_fp8_block.py new file mode 100644 index 00000000..7f4d6241 --- /dev/null +++ b/benchmarks/TritonNPUBench/fbgemm/_kernel_quantize_fp8_block/_kernel_quantize_fp8_block.py @@ -0,0 +1,96 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/fbgemm_operator_cases/newtest_cases/test__kernel_quantize_fp8_block_npu.py +# Main kernel: _kernel_quantize_fp8_block +# PT file: test__kernel_quantize_fp8_block_v2.pt + +import triton +import triton.language as tl + + +# === _kernel_quantize_fp8_block === +@triton.jit +def _kernel_quantize_fp8_block( + A, + A_scale, + A_fp8, + scale_ub, + M, + K, + stride_am, + stride_ak, + stride_om, + stride_ok, + stride_a_scale_m, + stride_a_scale_k, + TL_FP8_DTYPE: tl.constexpr, + MAX_FP8: tl.constexpr, + EPS: tl.constexpr, + CLAMP_MAX: tl.constexpr, + BLOCK_M: tl.constexpr, + BLOCK_K: tl.constexpr, + K_MAJOR: tl.constexpr, +) -> None: + """Quantize and scale each [BLOCK_M, BLOCK_K] block. + + Scale per block i, j is computed as 1 / (MAX_FP8 / max(abs(A[i:i+BLOCK_M, j:j+BLOCK_K]))) + + Kernel naively iterates through matrix with [BLOCK_M, BLOCK_K] tiles. + + Todo: + * Better tiling and ordering schemes. + + Args: + A (Tensor): [M, K] higher precision input tensor. + A_scale (Tensor): [cdiv(M, BLOCK_M), cdiv(K, BLOCK_K)] reciprocal scale tensor per block. + A_fp8 (Tensor): [M, K] fp8 scaled tensor. A_fp8 = A * a_scale + scale_ub (Tensor): [1] Maximum allowed value for scale. + M (int): Number of rows. + K (int): Number of columns. + stride_am (int): Stride of m dimension of A. + stride_ak (int): Stride of k dimension of A. + stride_om (int): Stride of m dimension of output. + stride_ok (int): Stride of k dimension of output. + stride_a_scale_m (int): Stride of m dimension of A_scale. + stride_a_scale_k (int): Stride of k dimension of A_scale. + TL_FP8_DTYPE (tl.dtype): Target fp8 datatype. + MAX_FP8 (float): Maxmimum expressible value for FP8. + EPS (float): Epsilon value for numerical stability. + CLAMP_MAX (bool): Whether to apply scale_ub. + BLOCK_M (int): Block size for M dimension of A_scale and kernel. + BLOCK_K (int): Block size for K dimension of A_scale and kernel. + K_MAJOR (bool): Whether output scales should be K major (True) or MN major (False). + """ + pid = tl.program_id(0) + grid_k = tl.cdiv(K, BLOCK_K) + block_m = pid // grid_k + block_k = pid % grid_k + rm = block_m * BLOCK_M + tl.arange(0, BLOCK_M) + rk = block_k * BLOCK_K + tl.arange(0, BLOCK_K) + a_offset = rm[:, None] * stride_am + rk[None, :] * stride_ak + out_offset = rm[:, None] * stride_om + rk[None, :] * stride_ok + a_mask = (rm < M)[:, None] & (rk < K)[None, :] + a_block = tl.load(A + a_offset, mask=a_mask, other=0.0) + + block_max = tl.max(tl.abs(a_block)) + # Apply appropriate clamping. + if CLAMP_MAX: + ub = tl.load(scale_ub) + block_max = tl.clamp(block_max, EPS, ub) + else: + block_max = tl.maximum(block_max, EPS) + scale = MAX_FP8 / block_max + + # Write in transposed order if specified. + if K_MAJOR: + scale_offset = block_m * stride_a_scale_m + block_k * stride_a_scale_k + else: + scale_offset = block_k * stride_a_scale_m + block_m * stride_a_scale_k + tl.store(A_scale + scale_offset, 1.0 / scale) + a_fp8 = a_block * 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. + a_fp8 = tl.clamp(a_fp8, -MAX_FP8, MAX_FP8) + a_fp8.to(TL_FP8_DTYPE) + tl.store(A_fp8 + out_offset, a_fp8, mask=a_mask) + diff --git a/benchmarks/TritonNPUBench/fbgemm/_kernel_quantize_fp8_group/_kernel_quantize_fp8_group.py b/benchmarks/TritonNPUBench/fbgemm/_kernel_quantize_fp8_group/_kernel_quantize_fp8_group.py new file mode 100644 index 00000000..473506d3 --- /dev/null +++ b/benchmarks/TritonNPUBench/fbgemm/_kernel_quantize_fp8_group/_kernel_quantize_fp8_group.py @@ -0,0 +1,152 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/fbgemm_operator_cases/newtest_cases/test__kernel_quantize_fp8_group_npu.py +# Main kernel: _kernel_quantize_fp8_group +# PT file: test__kernel_quantize_fp8_group_v2.pt + +import triton +import triton.language as tl + + +# === _kernel_quantize_fp8_group === +@triton.jit +def _kernel_quantize_fp8_group( + A, + A_scale, + A_fp8, + scale_ub, + m_sizes, + M, + K, + stride_am, + stride_ak, + stride_om, + stride_ok, + stride_a_scale_m, + stride_a_scale_k, + TL_FP8_DTYPE: tl.constexpr, + MAX_FP8: tl.constexpr, + EPS: tl.constexpr, + CLAMP_MAX: tl.constexpr, + USE_INT64: tl.constexpr, + GROUP_SIZE: tl.constexpr, + USE_M_MAJOR: tl.constexpr, + G: tl.constexpr, + GROUP_LOAD: tl.constexpr, +): + """Quantize and scale each GROUP_SIZE chunk of each row. + + Scale per group i is computed as 1 / (MAX_FP8 / max(abs(A[i:i+GROUP_SIZE]))) + + Each kernel thread is responsible for one row and loads and processes a tunable + number of groups at once. + + Args: + A (Tensor): [M, K] higher precision input tensor. + A_scale (Tensor): [M, cdiv(K, GROUP_SIZE)] reciprocal scale tensor per group. + A_fp8 (Tensor): [M, K] fp8 scaled tensor. A_fp8 = A * a + scale_ub (Tensor): [1] Maximum allowed value for scale. + m_sizes (Optional[Tensor]): [G] Number of rows in each group. + M (int): Number of rows. + K (int): Number of columns. + stride_am (int): Stride of m dimension of A. + stride_ak (int): Stride of k dimension of A. + stride_om (int): Stride of m dimension of output. + stride_ok (int): Stride of k dimension of output. + stride_a_scale_m (int): Stride of m dimension of A_scale. + stride_a_scale_k (int): Stride of k dimension of A_scale. + TL_FP8_DTYPE (tl.dtype): Target fp8 datatype. + MAX_FP8 (float): Maxmimum expressible value for FP8. + EPS (float): Epsilon value for numerical stability. + CLAMP_MAX (bool): Whether to apply scale_ub. + USE_INT64 (bool): Whether to index using int64, which may be needed for large tensors. + GROUP_SIZE (int): Group size for K dimension of A_scale and kernel. + USE_M_MAJOR (bool): Whether to use grouped M-major layout for A_scale. + G (int): Number of groups in A_scale, only relevant when m_sizes is provided. + GROUP_LOAD (int): Number of groups to load and process simultaneously. + """ + pid = tl.program_id(0) + if USE_INT64: + pid = pid.to(tl.int64) + # We load group_size * group_load chunks at a time. + row_offset = pid * stride_am + out_offset = pid * stride_om + scale_row_offset = pid * stride_a_scale_m + k_offset = tl.arange(0, GROUP_LOAD * GROUP_SIZE) + scale_k_offset = tl.arange(0, GROUP_LOAD) + NUM_GROUPS: tl.constexpr = K // GROUP_SIZE + + # When dealing with an M-major grouped gemm, we need to figure out + # which group this thread corresponds to and figure out the corresponding + # scale offset. + group_offset = 0 + group_cumsum = 0 + group_M = 0 + stop = False + if USE_M_MAJOR and G > 0: + # Iterate over groups to both compute the cumulative sum and find which group we are in. + for i in range(G): + if not stop: + group_M = tl.cast(tl.load(m_sizes + i), pid.dtype) + if (group_cumsum + group_M) <= pid: + group_cumsum += group_M + else: + # Indicate we are finished computing cumsum. + stop = True + + group_offset = group_cumsum * NUM_GROUPS + + for k in range(0, tl.cdiv(K, (GROUP_LOAD * GROUP_SIZE))): + # Load groups of the input. + chunk_offset = k_offset + k * GROUP_LOAD * GROUP_SIZE + a = tl.load( + A + row_offset + chunk_offset * stride_ak, mask=chunk_offset < K, other=0.0 + ) + # View loaded chunk as a set of groups. + a_grouped = tl.reshape(a, [GROUP_LOAD, GROUP_SIZE]) + # Reduce over groups. + group_max = tl.max(tl.abs(a_grouped), axis=1) + # Apply clamping if specified. + if CLAMP_MAX: + ub = tl.load(scale_ub) + group_max = tl.clamp(group_max, EPS, ub) + else: + group_max = tl.maximum(group_max, EPS) + # Scale and quantize. + a_scale = MAX_FP8 / group_max + scale_chunk_offset = scale_k_offset + k * GROUP_LOAD + + if USE_M_MAJOR and G > 0: + tl.store( + A_scale + + group_offset + + (pid - group_cumsum) * stride_a_scale_k + + (scale_chunk_offset * group_M), + 1.0 / a_scale, + mask=scale_chunk_offset < NUM_GROUPS, + ) + else: + if USE_M_MAJOR: + tl.store( + A_scale + + pid * stride_a_scale_k + + scale_chunk_offset * stride_a_scale_m, + 1.0 / a_scale, + mask=scale_chunk_offset < NUM_GROUPS, + ) + else: + tl.store( + A_scale + scale_row_offset + scale_chunk_offset * stride_a_scale_k, + 1.0 / a_scale, + mask=scale_chunk_offset < NUM_GROUPS, + ) + # Apply scale to input. + a_fp8 = a_grouped * a_scale[:, None] + # Clamp to FP8 range to avoid overflow + a_fp8 = tl.clamp(a_fp8, -MAX_FP8, MAX_FP8).to(TL_FP8_DTYPE) + # Write to output. + tl.store( + A_fp8 + out_offset + chunk_offset * stride_ok, + tl.ravel(a_fp8), + mask=chunk_offset < K, + ) + diff --git a/benchmarks/TritonNPUBench/fbgemm/_kernel_scale_fp8_row/_kernel_scale_fp8_row.py b/benchmarks/TritonNPUBench/fbgemm/_kernel_scale_fp8_row/_kernel_scale_fp8_row.py new file mode 100644 index 00000000..92af4e3f --- /dev/null +++ b/benchmarks/TritonNPUBench/fbgemm/_kernel_scale_fp8_row/_kernel_scale_fp8_row.py @@ -0,0 +1,59 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/fbgemm_operator_cases/newtest_cases/test__kernel_scale_fp8_row_npu.py +# Main kernel: _kernel_scale_fp8_row +# PT file: test__kernel_scale_fp8_row_v2.pt + +import triton +import triton.language as tl + + +# === _kernel_scale_fp8_row === +@triton.jit +def _kernel_scale_fp8_row( + A, + x_scale, + w_scale, + scaled_out, + M, + N, + stride_am, + stride_an, + stride_om, + stride_on, + BLOCK_SIZE: tl.constexpr, +) -> None: + """ + Scale each row of A by x_scale and each column of A by w_scale. + + Args: + A (Tensor): [m, n] Input tensor to scale. + x_scale (Tensor): [m] Row-wise scale tensor. + w_scale (Tensor): [n] Col-wise scale tensor. + scaled_out (Tensor): [m, n] Output tensor. + M (int): Number of rows. + N (int): Number of columns. + stride_am (int): Stride of m dimension of A. + stride_an (int): Stride of n dimension of A. + stride_om (int): Stride of m dimension of output. + stride_on (int): Stride of n dimension of output. + BLOCK_SIZE (int): Block size for data loads. + """ + pid = tl.program_id(0) + n_offset = tl.arange(0, BLOCK_SIZE) + # Load activation scale for this row. + row_scale = tl.load(x_scale + pid) + + # Iterate over chunks of the row and apply scales. + for _k in range(0, tl.cdiv(N, BLOCK_SIZE)): + a = tl.load( + A + pid * stride_am + n_offset * stride_an, mask=n_offset < N, other=0.0 + ) + col_scale = tl.load(w_scale + n_offset) + scaled_a = a * row_scale * col_scale + tl.store( + scaled_out + pid * stride_om + n_offset * stride_on, + scaled_a, + mask=n_offset < N, + ) + n_offset += BLOCK_SIZE + diff --git a/benchmarks/TritonNPUBench/fbgemm/_multi_head_jagged_flash_attention_bwd_preprocess_kernel/_multi_head_jagged_flash_attention_bwd_preprocess_kernel.py b/benchmarks/TritonNPUBench/fbgemm/_multi_head_jagged_flash_attention_bwd_preprocess_kernel/_multi_head_jagged_flash_attention_bwd_preprocess_kernel.py new file mode 100644 index 00000000..a833562c --- /dev/null +++ b/benchmarks/TritonNPUBench/fbgemm/_multi_head_jagged_flash_attention_bwd_preprocess_kernel/_multi_head_jagged_flash_attention_bwd_preprocess_kernel.py @@ -0,0 +1,65 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/fbgemm_operator_cases/newtest_cases/test__multi_head_jagged_flash_attention_bwd_preprocess_kernel_npu.py +# Main kernel: _multi_head_jagged_flash_attention_bwd_preprocess_kernel +# PT file: test__multi_head_jagged_flash_attention_bwd_preprocess_kernel_v2.pt + +import triton +import triton.language as tl + + +# === _multi_head_jagged_flash_attention_bwd_preprocess_kernel === +@triton.jit +def _multi_head_jagged_flash_attention_bwd_preprocess_kernel( + o_ptr, + o_offset_ptr, + do_ptr, + delta_ptr, + stride_oh, + stride_om, + stride_od, + stride_delta_h, + num_heads: tl.constexpr, + max_seq_len: tl.constexpr, + D: tl.constexpr, + BLOCK_M: tl.constexpr, + BLOCK_D: tl.constexpr, +): + pid_m = tl.program_id(axis=0) + pid_bh = tl.program_id(axis=1) + pid_batch = pid_bh // num_heads + pid_head = pid_bh % num_heads + + begin_o = tl.load(o_offset_ptr + pid_batch) + end_o = tl.load(o_offset_ptr + pid_batch + 1) + + M = end_o - begin_o + M = tl.minimum(M, max_seq_len) + + if M == 0: + return + + offs_om = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) + offs_od = tl.arange(0, BLOCK_D) + + o_offsets = ( + offs_om[:, None] * stride_om + + offs_od[None, :] * stride_od + + pid_head * stride_oh + + begin_o * stride_om + ) + o_ptrs = o_ptr + o_offsets + do_ptrs = do_ptr + o_offsets + o_mask = (offs_om[:, None] < M) & (offs_od[None, :] < D) + + # Load o and do + o = tl.load(o_ptrs, mask=o_mask) + do = tl.load(do_ptrs, mask=o_mask) + + delta = tl.sum(o * do, axis=1) + + tl.store( + delta_ptr + pid_head * stride_delta_h + begin_o + offs_om, + delta, + mask=offs_om < M, + ) + diff --git a/benchmarks/TritonNPUBench/fbgemm/cumsum_kernel/cumsum_kernel.py b/benchmarks/TritonNPUBench/fbgemm/cumsum_kernel/cumsum_kernel.py new file mode 100644 index 00000000..fd68b342 --- /dev/null +++ b/benchmarks/TritonNPUBench/fbgemm/cumsum_kernel/cumsum_kernel.py @@ -0,0 +1,37 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/fbgemm_operator_cases/newtest_cases/test_cumsum_kernel_gpu.py +# Main kernel: cumsum_kernel +# PT file: test_cumsum_kernel_v2.pt + +import triton +import triton.language as tl + + +# === cumsum_kernel === +@triton.jit +def cumsum_kernel( + m_sizes_ptr, + size_cumulative_ptr, + N: tl.constexpr, + BLOCK_SIZE: tl.constexpr, +): + offs = tl.arange(0, BLOCK_SIZE) + mask = offs < N + + # Load m_sizes + m_sizes = tl.load(m_sizes_ptr + offs, mask=mask, other=0) + + # Compute inclusive cumsum + cumsum = tl.cumsum(m_sizes, axis=0) + + # Store cumsum at indices 1 through N + tl.store(size_cumulative_ptr + offs + 1, cumsum, mask=mask) + + # Set first element to zero + first_elem_mask = offs == 0 + tl.store( + size_cumulative_ptr + offs, + tl.zeros([BLOCK_SIZE], dtype=cumsum.dtype), + mask=first_elem_mask, + ) + diff --git a/benchmarks/TritonNPUBench/fbgemm/dense_jagged_cat_jagged_out_kernel/dense_jagged_cat_jagged_out_kernel.py b/benchmarks/TritonNPUBench/fbgemm/dense_jagged_cat_jagged_out_kernel/dense_jagged_cat_jagged_out_kernel.py new file mode 100644 index 00000000..ca6d045a --- /dev/null +++ b/benchmarks/TritonNPUBench/fbgemm/dense_jagged_cat_jagged_out_kernel/dense_jagged_cat_jagged_out_kernel.py @@ -0,0 +1,38 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/fbgemm_operator_cases/newtest_cases/test_dense_jagged_cat_jagged_out_kernel_npu.py +# Main kernel: dense_jagged_cat_jagged_out_kernel +# PT file: dense_jagged_cat_jagged_out_kernel_v2.pt + +import triton +import triton.language as tl + + +# === dense_jagged_cat_jagged_out_kernel === +@triton.jit +def dense_jagged_cat_jagged_out_kernel( + a_ptr, # dense + b_ptr, # jagged + c_ptr, # jagged + b_offsets_ptr, + c_offsets_ptr, + max_seq_len, + BLOCK_SIZE: tl.constexpr, +): + pid_batch = tl.program_id(0) + b_start = tl.load(b_offsets_ptr + pid_batch) + b_end = tl.load(b_offsets_ptr + pid_batch + 1) + c_start = b_start + pid_batch + N = b_end - b_start + N = tl.minimum(N, max_seq_len) + + a = tl.load(a_ptr + pid_batch) + tl.store(c_ptr + c_start, a) + + offs_k = tl.arange(0, BLOCK_SIZE) + for k in range(0, N, BLOCK_SIZE): + b_offset = k + offs_k + b_ptrs = b_ptr + b_start + b_offset + b = tl.load(b_ptrs, mask=b_offset < N, other=0.0) + tl.store(c_ptr + c_start + 1 + b_offset, b, mask=b_offset < N) + tl.store(c_offsets_ptr + pid_batch, b_start + pid_batch) + diff --git a/benchmarks/TritonNPUBench/fbgemm/fused_single_block_kernel/fused_single_block_kernel.py b/benchmarks/TritonNPUBench/fbgemm/fused_single_block_kernel/fused_single_block_kernel.py new file mode 100644 index 00000000..f947c260 --- /dev/null +++ b/benchmarks/TritonNPUBench/fbgemm/fused_single_block_kernel/fused_single_block_kernel.py @@ -0,0 +1,107 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/fbgemm_operator_cases/newtest_cases/test_fused_single_block_kernel_npu.py +# Main kernel: fused_single_block_kernel +# PT file: test_fused_single_block_kernel_v2.pt + +import triton +import triton.language as tl + + +# === fused_single_block_kernel === +@triton.jit +def fused_single_block_kernel( + m_sizes_ptr, # [num_segments] input sizes + size_cumulative_ptr, # [num_segments + 1] cumulative size sum + starting_row_after_padding_ptr, # [num_segments + 1] output: padded cumsum + belong_indices_ptr, # [N] output: segment index + row_within_tensor_ptr, # [N] output: position within segment + num_segments: tl.constexpr, + N: tl.constexpr, + BLOCK_SIZE: tl.constexpr, + prefix_num: tl.constexpr, +): + pid = tl.program_id(0) + NUM_BLOCKS = tl.num_programs(0) + + offs = tl.arange(0, prefix_num) + mask = offs < num_segments + + # Load m_sizes + m_sizes = tl.load(m_sizes_ptr + offs, mask=mask, other=0) + + # Compute inclusive cumsum + cumsum = tl.cumsum(m_sizes, axis=0) + + # Store cumsum at indices 1 through N + tl.store( + size_cumulative_ptr + offs + 1 + (num_segments + 1) * pid, cumsum, mask=mask + ) + + # Set first element to zero + tl.store( + size_cumulative_ptr + offs + (num_segments + 1) * pid, + tl.zeros([1], dtype=cumsum.dtype), + mask=(offs == 0), + ) + + if pid == 0: + # Part 1: Compute padded cumsum (only first block does this) + offs = tl.arange(0, prefix_num) + mask = offs < num_segments + + # Load m_sizes + m_sizes = tl.load(m_sizes_ptr + offs, mask=mask, other=0) + + # Compute padded sizes + padded_sizes = ((m_sizes + 128 - 1) // 128) * 128 + + # Compute inclusive cumsum + cumsum = tl.cumsum(padded_sizes, axis=0) + + # Store at indices 1 through num_segments + tl.store(starting_row_after_padding_ptr + offs + 1, cumsum, mask=mask) + + # Set first element to zero + tl.store( + starting_row_after_padding_ptr + offs, + tl.zeros([1], dtype=cumsum.dtype), + mask=(offs == 0), + ) + tl.debug_barrier() + # Part 2: Segmented arange - process N elements in chunks + new_offs = tl.arange(0, BLOCK_SIZE) + BLOCK_SIZE * pid + for start in range(0, N, BLOCK_SIZE * NUM_BLOCKS): + row_idx = start + new_offs + row_mask = row_idx < N + + # Binary search using the cumsum_regular we computed + left = tl.zeros([BLOCK_SIZE], dtype=tl.int32) + right = tl.zeros([BLOCK_SIZE], dtype=tl.int32) + num_segments + + for _ in range(64): # log2(num_segments) iterations + mid = (left + right) // 2 + + # Get cumsum value at mid position + # Since we need cumsum[0] = 0, cumsum[1] = m_sizes[0], etc. + mid_val = tl.load( + size_cumulative_ptr + mid + (num_segments + 1) * pid, + mask=row_mask, + other=0, + ) + + cond = mid_val <= row_idx + left = tl.where(cond, mid + 1, left) + right = tl.where(~cond, mid, right) + + belong_idx = left - 1 + tl.store(belong_indices_ptr + row_idx, belong_idx, mask=row_mask) + + # Compute row_within_tensor + segment_start = tl.load( + size_cumulative_ptr + (num_segments + 1) * pid + belong_idx, + mask=row_mask, + other=0, + ) + row_within = row_idx - segment_start + tl.store(row_within_tensor_ptr + row_idx, row_within, mask=row_mask) + diff --git a/benchmarks/TritonNPUBench/fbgemm/jagged2_to_padded_dense_kernel/jagged2_to_padded_dense_kernel.py b/benchmarks/TritonNPUBench/fbgemm/jagged2_to_padded_dense_kernel/jagged2_to_padded_dense_kernel.py new file mode 100644 index 00000000..b0fd6c01 --- /dev/null +++ b/benchmarks/TritonNPUBench/fbgemm/jagged2_to_padded_dense_kernel/jagged2_to_padded_dense_kernel.py @@ -0,0 +1,50 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/fbgemm_operator_cases/newtest_cases/test_jagged2_to_padded_dense_kernel_npu.py +# Main kernel: jagged2_to_padded_dense_kernel +# PT file: jagged2_to_padded_dense_kernel_v2.pt + +import triton +import triton.language as tl + + +# === jagged2_to_padded_dense_kernel === +@triton.jit +def jagged2_to_padded_dense_kernel( + x_ptr, + lengths_ptr, + offsets_ptr, + output_dense_ptr, + stride_b, + stride_m, + stride_n, + max_length, + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, +): + pid_batch = tl.program_id(2) + pid_m = tl.program_id(0) + pid_n = tl.program_id(1) + + begin = tl.load(offsets_ptr + pid_batch) + seqlen = tl.load(lengths_ptr + pid_batch) + + seqlen = tl.minimum(seqlen, max_length) + if seqlen == 0: + return + + offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) + offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) + + x_ptrs = x_ptr + begin + offs_m[:, None] * seqlen + offs_n[None, :] + x = tl.load(x_ptrs, mask=((offs_m[:, None] < seqlen) & (offs_n[None, :] < seqlen))) + + out_ptrs = ( + output_dense_ptr + + pid_batch * stride_b + + offs_m[:, None] * stride_m + + offs_n[None, :] * stride_n + ) + tl.store( + out_ptrs, x, mask=((offs_m[:, None] < seqlen) & (offs_n[None, :] < seqlen)) + ) + diff --git a/benchmarks/TritonNPUBench/fbgemm/jagged_2_softmax_backward_kernel/jagged_2_softmax_backward_kernel.py b/benchmarks/TritonNPUBench/fbgemm/jagged_2_softmax_backward_kernel/jagged_2_softmax_backward_kernel.py new file mode 100644 index 00000000..e1cf7077 --- /dev/null +++ b/benchmarks/TritonNPUBench/fbgemm/jagged_2_softmax_backward_kernel/jagged_2_softmax_backward_kernel.py @@ -0,0 +1,91 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/fbgemm_operator_cases/newtest_cases/test_jagged_2_softmax_backward_kernel_npu.py +# Main kernel: jagged_2_softmax_backward_kernel +# PT file: test_jagged_2_softmax_backward_kernel_v2.pt + +import triton +import triton.language as tl + + +# === jagged_2_softmax_backward_kernel === +@triton.jit +def jagged_2_softmax_backward_kernel( + grad_output_ptr, # input + softmax_output_ptr, + grad_input_ptr, # return value + offsets_row_ptr, + offsets_col_ptr, + offsets_overall_ptr, + grad_output_stride, + softmax_output_stride, + grad_input_stride, + transpose, # transpose + max_seq_len_row: tl.constexpr, + max_seq_len_col: tl.constexpr, + BLOCK_SIZE: tl.constexpr, +): + pid_batch = tl.program_id(0) + pid_head = tl.program_id(1) + begin = tl.load(offsets_overall_ptr + pid_batch) + # end = tl.load(offsets_overall_ptr + pid_batch + 1) # noqa F841 + + # softmax on row + if transpose: + N = tl.load(offsets_row_ptr + pid_batch + 1) - tl.load( + offsets_row_ptr + pid_batch + ) + H = tl.load(offsets_col_ptr + pid_batch + 1) - tl.load( + offsets_col_ptr + pid_batch + ) + stride_n = H + stride_h = H // H # 1 + # sometimes H is larger than max_seq_len_col + H = tl.minimum(max_seq_len_col, H) + N = tl.minimum(max_seq_len_row, N) + # softmax on col + else: + N = tl.load(offsets_col_ptr + pid_batch + 1) - tl.load( + offsets_col_ptr + pid_batch + ) + H = tl.load(offsets_row_ptr + pid_batch + 1) - tl.load( + offsets_row_ptr + pid_batch + ) + stride_h = N + stride_n = N // N # 1 + H = tl.minimum(max_seq_len_row, H) + N = tl.minimum(max_seq_len_col, N) + + if pid_head >= H: + return + if H == 0 or N == 0: + pass + + start_ptr = grad_output_ptr + begin * grad_output_stride + offsets = tl.arange(0, BLOCK_SIZE) + + grad_output_ptrs = ( + start_ptr + + offsets * grad_output_stride * stride_n + + pid_head * grad_output_stride * stride_h + ) + softmax_output_ptrs = ( + softmax_output_ptr + + begin * softmax_output_stride + + offsets * softmax_output_stride * stride_n + + pid_head * softmax_output_stride * stride_h + ) + + grad_output_row = tl.load(grad_output_ptrs, mask=offsets < N, other=0.0) + softmax_output_row = tl.load(softmax_output_ptrs, mask=offsets < N, other=0.0) + + sum_value = tl.sum(grad_output_row * softmax_output_row, axis=0) + grad_input_row = (grad_output_row - sum_value) * softmax_output_row + + grad_input_row_start_ptr = grad_input_ptr + begin * grad_input_stride + grad_input_ptrs = ( + grad_input_row_start_ptr + + offsets * grad_input_stride * stride_n + + pid_head * grad_input_stride * stride_h + ) + tl.store(grad_input_ptrs, grad_input_row, mask=offsets < N) + diff --git a/benchmarks/TritonNPUBench/fbgemm/jagged_2_softmax_kernel/jagged_2_softmax_kernel.py b/benchmarks/TritonNPUBench/fbgemm/jagged_2_softmax_kernel/jagged_2_softmax_kernel.py new file mode 100644 index 00000000..ec6a8aee --- /dev/null +++ b/benchmarks/TritonNPUBench/fbgemm/jagged_2_softmax_kernel/jagged_2_softmax_kernel.py @@ -0,0 +1,97 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/fbgemm_operator_cases/newtest_cases/test_jagged_2_softmax_kernel_npu.py +# Main kernel: jagged_2_softmax_kernel +# PT file: test_jagged_2_softmax_kernel_v2.pt + +import triton +import triton.language as tl + + +# === jagged_2_softmax_kernel === +@triton.jit +def jagged_2_softmax_kernel( + input_ptr, + output_ptr, + offsets_row_ptr, # seq + offsets_col_ptr, # head + offsets_overall_ptr, # offsets for overall matrix = seq_length_i * head_i + input_stride, + output_stride, + transpose, # one if a is transpose, otherwise zero + max_seq_len_row, # max_seq_len for row (seq) + max_seq_len_col, # max_seq_len for col (head) + BLOCK_SIZE: tl.constexpr, # BLOCK_SIZE > seq_length +): + """ + input shape is [sum_B(Ni * Hi)] + output shape is [sum_B(Ni * Hi)] + Padded version = [B, N, H] + Calculate softmax alone N dim + Each kernel calulates softmax for 1 sample and 1 head + offsets_row.size == offsets_col.size == offsets_overall.size + """ + + pid_batch = tl.program_id(0) + pid_head = tl.program_id(1) + # start location of current example + begin = tl.load(offsets_overall_ptr + pid_batch) + # end = tl.load(offsets_overall_ptr + pid_batch + 1) # noqa F841 + # end - begin = M_i * N_i + + # softmax on row + if transpose: + N = tl.load(offsets_row_ptr + pid_batch + 1) - tl.load( + offsets_row_ptr + pid_batch + ) + H = tl.load(offsets_col_ptr + pid_batch + 1) - tl.load( + offsets_col_ptr + pid_batch + ) + stride_n = H + stride_h = H // H # 1 + # sometimes H is larger than max_seq_len_col + H = tl.minimum(max_seq_len_col, H) + N = tl.minimum(max_seq_len_row, N) + # softmax on col + else: + N = tl.load(offsets_col_ptr + pid_batch + 1) - tl.load( + offsets_col_ptr + pid_batch + ) + H = tl.load(offsets_row_ptr + pid_batch + 1) - tl.load( + offsets_row_ptr + pid_batch + ) + stride_h = N + stride_n = N // N # 1 + H = tl.minimum(max_seq_len_row, H) + N = tl.minimum(max_seq_len_col, N) + + if pid_head >= H: # TODO double check the equal here + return + if H == 0 or N == 0: + return + + # start of the current example + start_ptr = input_ptr + begin * input_stride + # offset for n + offsets = tl.arange(0, BLOCK_SIZE) + + # Load a softmax row + input_ptrs = ( + start_ptr + + offsets * input_stride * stride_n + + pid_head * input_stride * stride_h + ) # start + n offsets + head offset + row = tl.load(input_ptrs, mask=offsets < N, other=-float("inf")) + row_mins_max = row - tl.max(row, axis=0) + numerator = tl.exp(row_mins_max) + denominator = tl.sum(numerator, axis=0) + softmax_output = numerator / denominator + + # calculate output ptr, should be similar to input + output_start_ptr = output_ptr + begin * output_stride + output_ptrs = ( + output_start_ptr + + offsets * output_stride * stride_n + + pid_head * output_stride * stride_h + ) + tl.store(output_ptrs, softmax_output, mask=offsets < N) + diff --git a/benchmarks/TritonNPUBench/fbgemm/jagged_dense_elementwise_mul_jagged_out_kernel/jagged_dense_elementwise_mul_jagged_out_kernel.py b/benchmarks/TritonNPUBench/fbgemm/jagged_dense_elementwise_mul_jagged_out_kernel/jagged_dense_elementwise_mul_jagged_out_kernel.py new file mode 100644 index 00000000..2e30f30e --- /dev/null +++ b/benchmarks/TritonNPUBench/fbgemm/jagged_dense_elementwise_mul_jagged_out_kernel/jagged_dense_elementwise_mul_jagged_out_kernel.py @@ -0,0 +1,66 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/fbgemm_operator_cases/newtest_cases/test_jagged_dense_elementwise_mul_jagged_out_kernel_npu.py +# Main kernel: jagged_dense_elementwise_mul_jagged_out_kernel +# PT file: jagged_dense_elementwise_mul_jagged_out_kernel_v2.pt + +import triton +import triton.language as tl + + +# === jagged_dense_elementwise_mul_jagged_out_kernel === +@triton.jit +def jagged_dense_elementwise_mul_jagged_out_kernel( + a_ptr, # 1d jagged + b_ptr, # dense + c_ptr, # 1d jagged + a_seq_lengths_ptr, + a_offsets_ptr, + stride_a, + stride_bm, + stride_bn, + max_seq_len, + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, +): + pid_batch = tl.program_id(0) + pid_row_block = tl.program_id(1) + + batch_offset = tl.load(a_offsets_ptr + pid_batch) + batch_seq_len = tl.load(a_seq_lengths_ptr + pid_batch) + truncated_seq_len = tl.minimum(batch_seq_len, max_seq_len) + + offs_row = tl.arange(0, BLOCK_M) + offs_col = tl.arange(0, BLOCK_N) + + rows = pid_row_block * BLOCK_M + offs_row + + # a start + batch offset + row offsets + initial col offsets + a_ptrs = ( + a_ptr + + batch_offset * stride_a + + rows[:, None] * truncated_seq_len + + offs_col[None, :] + ) + + # b start + row offsets + initial col offsets + b_ptrs = b_ptr + rows[:, None] * stride_bm + offs_col[None, :] * stride_bn + + # c start + batch offset + row offsets + initial col offsets + c_ptrs = ( + c_ptr + batch_offset + rows[:, None] * truncated_seq_len + offs_col[None, :] + ) + + for block_start in range(0, truncated_seq_len, BLOCK_N): + cols = block_start + offs_col + # pyre-fixme[16]: `int` has no attribute `__getitem__`. + mask = (rows[:, None] < truncated_seq_len) & (cols[None, :] < truncated_seq_len) + a = tl.load(a_ptrs, mask=mask) + a_ptrs += BLOCK_N + + b = tl.load(b_ptrs, mask=mask) + b_ptrs += BLOCK_N + + c = a * b + tl.store(c_ptrs, c, mask=mask) + c_ptrs += BLOCK_N + diff --git a/benchmarks/TritonNPUBench/fbgemm/jagged_jagged_elementwise_arithmetic_ops/jagged_jagged_elementwise_arithmetic_ops.py b/benchmarks/TritonNPUBench/fbgemm/jagged_jagged_elementwise_arithmetic_ops/jagged_jagged_elementwise_arithmetic_ops.py new file mode 100644 index 00000000..29b719bd --- /dev/null +++ b/benchmarks/TritonNPUBench/fbgemm/jagged_jagged_elementwise_arithmetic_ops/jagged_jagged_elementwise_arithmetic_ops.py @@ -0,0 +1,60 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/fbgemm_operator_cases/newtest_cases/test_jagged_jagged_elementwise_arithmetic_ops_npu.py +# Main kernel: jagged_jagged_elementwise_arithmetic_ops +# PT file: test_jagged_jagged_elementwise_arithmetic_ops_v2.pt + +import triton +import triton.language as tl + +@triton.jit +def tensor_elementwise_add(x, y): + return x + y + +@triton.jit +def tensor_elementwise_mul(x, y): + return x * y + +# === jagged_jagged_elementwise_arithmetic_ops === +@triton.jit +def jagged_jagged_elementwise_arithmetic_ops( + # pyre-fixme[2]: Parameter must be annotated. + x_ptr, # x_ptr and y_ptr is pointer of jagged tensor value + # pyre-fixme[2]: Parameter must be annotated. + y_ptr, + M: tl.constexpr, # M and N would be size of the tensor with (M , N) + N: tl.constexpr, + stride_row: tl.constexpr, # shared row stride for tensor + stride_col: tl.constexpr, # shared colume stride for tensor + # pyre-fixme[2]: Parameter must be annotated. + output, + thread_block_row_size: tl.constexpr, # row and colume size of current thread block with size (thread_block_row_size * thread_block_col_size) + thread_block_col_size: tl.constexpr, + ops_func: tl.constexpr, # function use for calculation either add or multiplication +) -> None: + pid = tl.program_id(0) + # number of col group need for total N col + num_group_n = (N + thread_block_col_size - 1) // thread_block_col_size + # pid position in col perspective in range(0,num_group_n) + pid_n = pid % num_group_n + # pid position in row perspective since everytime row increase when we have num_group_n iteration + pid_m = pid // num_group_n + + offset_m = pid_m * thread_block_row_size + tl.arange(0, thread_block_row_size) + offset_n = pid_n * thread_block_col_size + tl.arange(0, thread_block_col_size) + mask = (offset_m[:, None] < M) & (offset_n[None, :] < N) + offset = offset_m[:, None] * stride_row + offset_n[None, :] * stride_col + + x_ptr += offset + y_ptr += offset + + x = tl.load(x_ptr, mask=mask) + y = tl.load(y_ptr, mask=mask) + + if ops_func == "add": + z = tensor_elementwise_add(x, y) + else: + z = tensor_elementwise_mul(x, y) + + output += offset + tl.store(output, z, mask=mask) + diff --git a/benchmarks/TritonNPUBench/fbgemm/jagged_self_substraction_jagged_out_kernel/jagged_self_substraction_jagged_out_kernel.py b/benchmarks/TritonNPUBench/fbgemm/jagged_self_substraction_jagged_out_kernel/jagged_self_substraction_jagged_out_kernel.py new file mode 100644 index 00000000..8b2e479f --- /dev/null +++ b/benchmarks/TritonNPUBench/fbgemm/jagged_self_substraction_jagged_out_kernel/jagged_self_substraction_jagged_out_kernel.py @@ -0,0 +1,42 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/fbgemm_operator_cases/newtest_cases/test_jagged_self_substraction_jagged_out_kernel_npu.py +# Main kernel: jagged_self_substraction_jagged_out_kernel +# PT file: jagged_self_substraction_jagged_out_kernel_v2.pt + +import triton +import triton.language as tl + + +# === jagged_self_substraction_jagged_out_kernel === +@triton.jit +def jagged_self_substraction_jagged_out_kernel( + a_ptr, # jagged + b_ptr, # jagged + a_offsets_ptr, + b_offsets_ptr, + max_seq_len, + BLOCK_SIZE: tl.constexpr, +): + pid_batch = tl.program_id(0) + pid_index = tl.program_id(1) + + a_offset = tl.load(a_offsets_ptr + pid_batch) + a_length = tl.load(a_offsets_ptr + pid_batch + 1) - a_offset + a_length = tl.minimum(a_length, max_seq_len + 1) + + if a_length <= 1: + return + + N = a_length - 1 + if pid_index >= N: + return + + a_cur = tl.load(a_ptr + a_offset + pid_index) + offs = tl.arange(0, BLOCK_SIZE) + mask = offs < N + a_row = tl.load(a_ptr + a_offset + offs + 1, mask=mask) + b = a_cur - a_row + + b_offset = tl.load(b_offsets_ptr + pid_batch) + tl.store(b_ptr + b_offset + pid_index * N + offs, b, mask=mask) + diff --git a/benchmarks/TritonNPUBench/fbgemm/jagged_softmax_backward_kernel/jagged_softmax_backward_kernel.py b/benchmarks/TritonNPUBench/fbgemm/jagged_softmax_backward_kernel/jagged_softmax_backward_kernel.py new file mode 100644 index 00000000..5e032e00 --- /dev/null +++ b/benchmarks/TritonNPUBench/fbgemm/jagged_softmax_backward_kernel/jagged_softmax_backward_kernel.py @@ -0,0 +1,65 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/fbgemm_operator_cases/newtest_cases/test_jagged_softmax_backward_kernel_npu.py +# Main kernel: jagged_softmax_backward_kernel +# PT file: jagged_softmax_backward_kernel_v2.pt + +import triton +import triton.language as tl + + +# === jagged_softmax_backward_kernel === +@triton.jit +def jagged_softmax_backward_kernel( + grad_output_ptr, + softmax_output_ptr, + grad_input_ptr, # return value + input_offsets_ptr, + grad_output_row_stride, + grad_output_head_stride, + softmax_output_row_stride, + softmax_output_head_stride, + grad_input_row_stride, + grad_input_head_stride, + max_seq_len: tl.constexpr, + BLOCK_SIZE: tl.constexpr, +): + """ + grad_output_ptr shpae is [SUM_B, H] + softmax_output shape is [SUM_B, H] + grad_input shape is [SUM_B, H] + """ + + pid_batch = tl.program_id(0) + pid_head = tl.program_id(1) + row_begin = tl.load(input_offsets_ptr + pid_batch) + row_end = tl.load(input_offsets_ptr + pid_batch + 1) + N = tl.minimum( + max_seq_len, row_end - row_begin + ) # number of rows to consider softmax + + col_offsets = tl.arange(0, BLOCK_SIZE) + grad_output_ptrs = ( + grad_output_ptr + + row_begin * grad_output_row_stride + + col_offsets * grad_output_row_stride + + pid_head * grad_output_head_stride + ) + softmax_output_ptrs = ( + softmax_output_ptr + + row_begin * softmax_output_row_stride + + col_offsets * softmax_output_row_stride + + pid_head * softmax_output_head_stride + ) + grad_output_row = tl.load(grad_output_ptrs, mask=col_offsets < N, other=0.0) + softmax_output_row = tl.load(softmax_output_ptrs, mask=col_offsets < N, other=0.0) + + sum_value = tl.sum(grad_output_row * softmax_output_row, axis=0) + grad_input_row = (grad_output_row - sum_value) * softmax_output_row + grad_input_ptrs = ( + grad_input_ptr + + row_begin * grad_input_row_stride + + col_offsets * grad_input_row_stride + + pid_head * grad_input_head_stride + ) + tl.store(grad_input_ptrs, grad_input_row, mask=col_offsets < N) + diff --git a/benchmarks/TritonNPUBench/fbgemm/jagged_softmax_kernel/jagged_softmax_kernel.py b/benchmarks/TritonNPUBench/fbgemm/jagged_softmax_kernel/jagged_softmax_kernel.py new file mode 100644 index 00000000..ef58df1c --- /dev/null +++ b/benchmarks/TritonNPUBench/fbgemm/jagged_softmax_kernel/jagged_softmax_kernel.py @@ -0,0 +1,57 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/fbgemm_operator_cases/newtest_cases/test_jagged_softmax_kernel_npu.py +# Main kernel: jagged_softmax_kernel +# PT file: jagged_softmax_kernel_v2.pt + +import triton +import triton.language as tl + + +# === jagged_softmax_kernel === +@triton.jit +def jagged_softmax_kernel( + input_ptr, + output_ptr, + input_offsets_ptr, + input_row_stride, + input_head_stride, + output_row_stride, + output_head_stride, + max_seq_len: tl.constexpr, + BLOCK_SIZE: tl.constexpr, # BLOCK_SIZE > N (seq len) +): + """ + input shpae is [SUM_B, H] + output shape is [SUM_B, H] + """ + + pid_batch = tl.program_id(0) + pid_head = tl.program_id(1) + row_begin = tl.load(input_offsets_ptr + pid_batch) + row_end = tl.load(input_offsets_ptr + pid_batch + 1) + N = tl.minimum( + max_seq_len, row_end - row_begin + ) # number of rows to consider softmax + if N == 0: + return + + row_start_ptr = input_ptr + row_begin * input_row_stride + col_offsets = tl.arange(0, BLOCK_SIZE) + input_ptrs = ( + row_start_ptr + col_offsets * input_row_stride + pid_head * input_head_stride + ) + row = tl.load(input_ptrs, mask=col_offsets < N, other=-float("inf")) + row_mins_max = row - tl.max(row, axis=0) + numerator = tl.exp(row_mins_max) + denominator = tl.sum(numerator, axis=0) + softmax_output = numerator / denominator + + output_row_start_ptr = output_ptr + row_begin * output_row_stride + output_ptrs = ( + output_row_start_ptr + + col_offsets * output_row_stride + + pid_head * output_head_stride + ) + + tl.store(output_ptrs, softmax_output, mask=col_offsets < N) + diff --git a/benchmarks/TritonNPUBench/fbgemm/padded_dense_to_jagged2_kernel/padded_dense_to_jagged2_kernel.py b/benchmarks/TritonNPUBench/fbgemm/padded_dense_to_jagged2_kernel/padded_dense_to_jagged2_kernel.py new file mode 100644 index 00000000..bd146c0b --- /dev/null +++ b/benchmarks/TritonNPUBench/fbgemm/padded_dense_to_jagged2_kernel/padded_dense_to_jagged2_kernel.py @@ -0,0 +1,51 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/fbgemm_operator_cases/newtest_cases/test_padded_dense_to_jagged2_kernel_npu.py +# Main kernel: padded_dense_to_jagged2_kernel +# PT file: padded_dense_to_jagged2_kernel_v2.pt + +import triton +import triton.language as tl + + +# === padded_dense_to_jagged2_kernel === +@triton.jit +def padded_dense_to_jagged2_kernel( + x_ptr, + lengths_ptr, + offsets_ptr, + output_jagged_ptr, + stride_b, + stride_m, + stride_n, + max_length, + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, +): + pid_batch = tl.program_id(2) + pid_m = tl.program_id(0) + pid_n = tl.program_id(1) + + begin = tl.load(offsets_ptr + pid_batch) + # end = tl.load(offsets_ptr + pid_batch + 1) + seqlen = tl.load(lengths_ptr + pid_batch) + + seqlen = tl.minimum(seqlen, max_length) + + if seqlen == 0: + return + + offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) + offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) + + x_ptrs = ( + x_ptr + + pid_batch * stride_b + + offs_m[:, None] * stride_m + + offs_n[None, :] * stride_n + ) + x = tl.load(x_ptrs, mask=((offs_m[:, None] < seqlen) & (offs_n[None, :] < seqlen))) + out_ptrs = output_jagged_ptr + begin + offs_m[:, None] * seqlen + offs_n[None, :] + tl.store( + out_ptrs, x, mask=((offs_m[:, None] < seqlen) & (offs_n[None, :] < seqlen)) + ) + diff --git a/benchmarks/TritonNPUBench/fbgemm/triton_batched_dense_vec_jagged_2d_matmul/triton_batched_dense_vec_jagged_2d_matmul.py b/benchmarks/TritonNPUBench/fbgemm/triton_batched_dense_vec_jagged_2d_matmul/triton_batched_dense_vec_jagged_2d_matmul.py new file mode 100644 index 00000000..52b07f92 --- /dev/null +++ b/benchmarks/TritonNPUBench/fbgemm/triton_batched_dense_vec_jagged_2d_matmul/triton_batched_dense_vec_jagged_2d_matmul.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_triton_batched_dense_vec_jagged_2d_matmul_npu.py +# Main kernel: triton_batched_dense_vec_jagged_2d_matmul +# PT file: test_triton_batched_dense_vec_jagged_2d_matmul_v2.pt + +import triton +import triton.language as tl + + +# === triton_batched_dense_vec_jagged_2d_matmul === +@triton.jit +def triton_batched_dense_vec_jagged_2d_matmul( + # pyre-fixme[2]: Parameter must be annotated. + jagged_tensor_ptr, + # pyre-fixme[2]: Parameter must be annotated. + dense_ptr, + # pyre-fixme[2]: Parameter must be annotated. + jagged_offset, + thread_block_col_size: tl.constexpr, + # pyre-fixme[2]: Parameter must be annotated. + dense_row_stride, + # pyre-fixme[2]: Parameter must be annotated. + jagged_value_row_stride, + # pyre-fixme[2]: Parameter must be annotated. + D, + H: tl.constexpr, + # pyre-fixme[2]: Parameter must be annotated. + output_ptr, +) -> None: + + pid = tl.program_id(0) + + # number of kernel need for with matrix (N,D) calculated by D // thread_block_col_size + GRID_DIM_COL = (D + thread_block_col_size - 1) // thread_block_col_size + + # current output row index + output_row_idx = pid // GRID_DIM_COL + + # current jagged tensor offset index + jagged_offset_id = output_row_idx // H + + # current index with D reference since the real shape of jagged values is [B , N , H * D] + D_refer_idx = output_row_idx % H + + # current part of [N * D] id + group_id = pid % GRID_DIM_COL + + # size of tile + offset = group_id * thread_block_col_size + tl.arange(0, thread_block_col_size) + + # begin index and end index of values + begin = tl.load(jagged_offset + jagged_offset_id) + end = tl.load(jagged_offset + (jagged_offset_id + 1)) + + # update each pointer to the correct address + dense_ptr += output_row_idx * dense_row_stride + jagged_tensor_ptr += begin * jagged_value_row_stride + D_refer_idx * D + output_ptr += D * output_row_idx + + # Number of row each kernel will go through + num_row = tl.minimum(end - begin, dense_row_stride) + + # accumulation variable use for matmul + acc = tl.zeros((thread_block_col_size,), dtype=tl.float32) + mask = offset < D + for i in range(num_row): + val1 = tl.load(dense_ptr + i) + val2 = tl.load(jagged_tensor_ptr + offset, mask=mask, other=0.0) + result = val1 * val2 + acc += result + jagged_tensor_ptr += jagged_value_row_stride + + tl.store(output_ptr + offset, acc, mask=mask) + diff --git a/benchmarks/TritonNPUBench/fbgemm/triton_dense_to_jagged/triton_dense_to_jagged.py b/benchmarks/TritonNPUBench/fbgemm/triton_dense_to_jagged/triton_dense_to_jagged.py new file mode 100644 index 00000000..0c9173cc --- /dev/null +++ b/benchmarks/TritonNPUBench/fbgemm/triton_dense_to_jagged/triton_dense_to_jagged.py @@ -0,0 +1,101 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/fbgemm_operator_cases/newtest_cases/test_triton_dense_to_jagged_npu.py +# Main kernel: triton_dense_to_jagged +# PT file: test_triton_dense_to_jagged_v2.pt + +import triton +import triton.language as tl + + +# === triton_dense_to_jagged === +@triton.jit +def triton_dense_to_jagged( + # pyre-fixme[2]: Parameter must be annotated. + jagged_value_ptr, + # pyre-fixme[2]: Parameter must be annotated. + jagged_offsets_ptr, + jagged_value_row_stride: int, + # pyre-fixme[2]: Parameter must be annotated. + output_dense_ptr, + # pyre-fixme[2]: Parameter must be annotated. + dense_indices_ptr, + # pyre-fixme[2]: Parameter must be annotated. + dense_col_stride, # stride of output dense with dimension (z,y,x) + dense_row_stride: int, + # pyre-fixme[2]: Parameter must be annotated. + dense_matrix_stride, + JAGGED_DIM: tl.constexpr, # number of dimension of jagged tensor + thread_block_row_size: tl.constexpr, + thread_block_col_size: tl.constexpr, + operation_function: tl.constexpr, # fusion arithmetic opeartion function and it's input dense + # pyre-fixme[2]: Parameter must be annotated. + operation_jagged_value_ptr, +) -> None: + pid = tl.program_id(0) + + begin = tl.load(jagged_offsets_ptr + pid) + end = tl.load(jagged_offsets_ptr + (pid + 1)) + + # size of the current value offset range (M , N) + N = jagged_value_row_stride + M = end - begin + + dense_boundary_col = dense_row_stride + # tl.minimum will change the return type cased compile issue + # in that case use if statement instead + if N < dense_row_stride: + dense_boundary_col = N + + dense_boundary_row = tl.minimum(dense_matrix_stride // dense_row_stride, M) + + jagged_value_ptr += begin * jagged_value_row_stride + if JAGGED_DIM > 2: + dense_indice = tl.load(dense_indices_ptr + pid) + # if dense output range we set dense_boundary to -1 + # that mean dense values will not be use with mask + # since we still need the calculation of fusion step + # therefore we do not do return here + if dense_indice == -1: + dense_boundary_col = -1 + else: + output_dense_ptr += dense_indice + else: + output_dense_ptr += pid * dense_matrix_stride + + if operation_function is not None: + operation_jagged_value_ptr += begin * jagged_value_row_stride + + offset_row = tl.arange(0, thread_block_row_size) + + for _i in range(begin, end, thread_block_row_size): + offset_col = tl.arange(0, thread_block_col_size) + block_offset = ( + offset_row[:, None] * dense_row_stride + + offset_col[None, :] * dense_col_stride + ) + + for _j in range(0, N, thread_block_col_size): + dense_mask = (offset_row[:, None] < dense_boundary_row) & ( + offset_col[None, :] < dense_boundary_col + ) + jagged_mask = (offset_row[:, None] < M) & (offset_col[None, :] < N) + dense_values = tl.load( + output_dense_ptr + block_offset, mask=dense_mask, other=0 + ) + if operation_function is not None: + operation_jagged_value = tl.load( + operation_jagged_value_ptr + block_offset, mask=jagged_mask, other=0 + ) + if operation_function == "add": + dense_values = tensor_elementwise_add( + dense_values, operation_jagged_value + ) + else: + dense_values = tensor_elementwise_mul( + dense_values, operation_jagged_value + ) + tl.store(jagged_value_ptr + block_offset, dense_values, mask=jagged_mask) + offset_col += thread_block_col_size + block_offset += thread_block_col_size + offset_row += thread_block_row_size + diff --git a/benchmarks/TritonNPUBench/fbgemm/triton_jagged_to_dense/triton_jagged_to_dense.py b/benchmarks/TritonNPUBench/fbgemm/triton_jagged_to_dense/triton_jagged_to_dense.py new file mode 100644 index 00000000..72463f5a --- /dev/null +++ b/benchmarks/TritonNPUBench/fbgemm/triton_jagged_to_dense/triton_jagged_to_dense.py @@ -0,0 +1,113 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/fbgemm_operator_cases/newtest_cases/test_triton_jagged_to_dense_npu.py +# Main kernel: triton_jagged_to_dense +# PT file: test_triton_jagged_to_dense_v2.pt + +import triton +import triton.language as tl + + +@triton.jit +def tensor_elementwise_add(x, y): + return x + y + +@triton.jit +def tensor_elementwise_mul(x, y): + return x * y + +# === triton_jagged_to_dense === +@triton.jit +def triton_jagged_to_dense( + # only constexpr annotations support in triton now + # pyre-fixme[2]: Parameter must be annotated. + jagged_value_ptr, + # pyre-fixme[2]: Parameter must be annotated. + jagged_offsets_ptr, + # pyre-fixme[2]: Parameter must be annotated. + jagged_value_row_stride, + # pyre-fixme[2]: Parameter must be annotated. + output_dense_ptr, + # pyre-fixme[2]: Parameter must be annotated. + dense_indices_ptr, + # pyre-fixme[2]: Parameter must be annotated. + dense_col_stride, # stride of output dense with dimension (z,y,x) + # pyre-fixme[2]: Parameter must be annotated. + dense_row_stride, + # pyre-fixme[2]: Parameter must be annotated. + dense_matrix_stride, + JAGGED_DIM: tl.constexpr, # number of dimension of jagged tensor + thread_block_row_size: tl.constexpr, + thread_block_col_size: tl.constexpr, + operation_function: tl.constexpr, # fusion arithmetic operation function and it's input dense + # pyre-fixme[2]: Parameter must be annotated. + operation_dense, +) -> None: + pid = tl.program_id(0) + + # begin index and end index of jagged tensor Values + begin = tl.load(jagged_offsets_ptr + pid) + end = tl.load(jagged_offsets_ptr + (pid + 1)) + + # adjust the address of the jagged tensor Values to the correct address + jagged_value_ptr += begin * jagged_value_row_stride + + # if it's 2D (or 1D) Jagged tensor we can direct use the offset in offsets ( since there is only one offset ) + # else we actually need to use the preprocess index to found the correct address of dense + if JAGGED_DIM > 2: + # read the index for current kernel + dense_indice = tl.load(dense_indices_ptr + pid) + + # if the dense_indice is -1 which mean it's a truncation case + # in that case we don't need to do anything since the dense + # initialize with padded value + if dense_indice == -1: + return + + # adjust the address of output dense ptr to the correct address + output_dense_ptr += dense_indice + + # also need to update the operation function if exist + # notice dense_indice of two is same because we assume + # the two dense + dense are same size + if operation_function is not None: + operation_dense += dense_indice + else: + output_dense_ptr += pid * dense_matrix_stride + + if operation_function is not None: + operation_dense += pid * dense_matrix_stride + + offset_row = tl.arange(0, thread_block_row_size) + + # boundary need for the mask since it could be dense's size smaller than jagged tensor or revert case + N = tl.minimum(dense_row_stride, jagged_value_row_stride) + M = tl.minimum(dense_matrix_stride // dense_row_stride, end - begin) + + for _i in range(begin, end, thread_block_row_size): + offset_col = tl.arange(0, thread_block_col_size) + block_offset = ( + offset_row[:, None] * dense_row_stride + + offset_col[None, :] * dense_col_stride + ) + for _j in range(0, N, thread_block_col_size): + mask = (offset_row[:, None] < M) & (offset_col[None, :] < N) + jagged_val = tl.load(jagged_value_ptr + block_offset, mask=mask, other=0) + + # if there is some arithmetic operation we do the fusion computation + if operation_function is not None: + val1 = jagged_val + val2 = tl.load(operation_dense + block_offset, mask=mask, other=0) + # do the arithmetic operation + if operation_function == "add": + jagged_val = tensor_elementwise_add(val1, val2) + else: + jagged_val = tensor_elementwise_mul(val1, val2) + + # store the result + tl.store(output_dense_ptr + block_offset, jagged_val, mask=mask) + + # update the block offset + offset_col += thread_block_col_size + block_offset += thread_block_col_size + offset_row += thread_block_row_size + diff --git a/benchmarks/TritonNPUBench/fbgemm/triton_jagged_to_dense_optimization_2d/triton_jagged_to_dense_optimization_2d.py b/benchmarks/TritonNPUBench/fbgemm/triton_jagged_to_dense_optimization_2d/triton_jagged_to_dense_optimization_2d.py new file mode 100644 index 00000000..27be2372 --- /dev/null +++ b/benchmarks/TritonNPUBench/fbgemm/triton_jagged_to_dense_optimization_2d/triton_jagged_to_dense_optimization_2d.py @@ -0,0 +1,100 @@ +# Extracted Triton Kernels +# Test file: /home/z00841464/benchmark/Triton_Automation-br_a3_operator/ascend/test/fbgemm_operator_cases/newtest_cases/test_triton_jagged_to_dense_optimization_2d_npu.py +# Main kernel: triton_jagged_to_dense_optimization_2d +# PT file: test_triton_jagged_to_dense_optimization_2d_v2.pt + +import triton +import triton.language as tl + + +# === triton_jagged_to_dense_optimization_2d === +@triton.jit +def triton_jagged_to_dense_optimization_2d( + # pyre-fixme[2]: Parameter must be annotated. + input_jagged_values_ptr, + # pyre-fixme[2]: Parameter must be annotated. + input_jagged_offset_ptr, + # pyre-fixme[2]: Parameter must be annotated. + input_jagged_row_stride, + # pyre-fixme[2]: Parameter must be annotated. + output_dense_ptr, + # pyre-fixme[2]: Parameter must be annotated. + output_dense_row_stride, + # pyre-fixme[2]: Parameter must be annotated. + output_dense_matrix_stride, + thread_block_row_size: tl.constexpr, + thread_block_col_size: tl.constexpr, + # pyre-fixme[2]: Parameter must be annotated. + padded_value, + operation_function: tl.constexpr, + # pyre-fixme[2]: Parameter must be annotated. + operation_dense, +) -> None: + pid = tl.program_id(0) + + # Current corresponding offset indice + offset_idx = pid + + # begin index and end index of jagged tensor Values + begin = tl.load(input_jagged_offset_ptr + offset_idx) + end = tl.load(input_jagged_offset_ptr + offset_idx + 1) + + # row size of current sub tensor + cur_jagged_tensor_row_size = end - begin + + # update dense and jagged tensor Values to corresponding address + output_dense_ptr += pid * output_dense_matrix_stride + input_jagged_values_ptr += begin * input_jagged_row_stride + + # also need to update the operation function if exist + # notice dense_indice of two is same because we assume + # the two dense + dense are same size + if operation_function is not None: + operation_dense += pid * output_dense_matrix_stride + + # jagged tensor row block + offset_row = tl.arange(0, thread_block_row_size) + + # dense row and col block + # notice jagged tensor and dense share same col block since embedding dimension is same + dense_col_size = output_dense_row_stride + dense_row_size = output_dense_matrix_stride // output_dense_row_stride + + for _i in range(0, dense_row_size, thread_block_row_size): + offset_col = tl.arange(0, thread_block_col_size) + block_offset = ( + offset_row[:, None] * output_dense_row_stride + offset_col[None, :] + ) + + for _j in range(0, dense_col_size, thread_block_col_size): + + # create mask for dense and jagged tensor for boundary check + dense_mask = (offset_row[:, None] < dense_row_size) & ( + offset_col[None, :] < dense_col_size + ) + jagged_mask = (offset_row[:, None] < cur_jagged_tensor_row_size) & ( + offset_col[None, :] < input_jagged_row_stride + ) + + # get value from jagged tesnor + jagged_val = tl.load( + input_jagged_values_ptr + block_offset, + mask=jagged_mask, + other=padded_value, + ) + + # do fusion operation if need + if operation_function is not None: + operation_dense_val = tl.load( + operation_dense + block_offset, mask=dense_mask, other=0.0 + ) + jagged_val = operation_function(operation_dense_val, jagged_val) + + # load value into empty dense + tl.store(output_dense_ptr + block_offset, jagged_val, mask=dense_mask) + + # update each block + offset_col += thread_block_col_size + block_offset += thread_block_col_size + offset_row += thread_block_row_size +