feat: Add template definitions for JD.com's core retail operations.-catlass-AtomGit

已合并

feat: Add template definitions for JD.com's core retail operations. #1161

ware2009创建于 3 天前从ware2009/catlass： master 合入到 cann/catlass： master

feat: Add template definitions for JD.com's core retail operations. #1161

已合并

ware2009创建于 3 天前

从ware2009/catlass master合入到cann/catlass master

ware2009

3 天前

描述

增加京东使用的一些自定义模板类型

关联的Issue

原因

1.主要增加x_flash_attention_infer模板和x_attention模板，支撑推荐推理和传统推荐。

测试

使用实例：
https://github.com/xLLM-AI/xllm-ops/tree/main/xllm_ops/x_attention
https://github.com/xLLM-AI/xllm-ops/tree/main/xllm_ops/x_flash_attention_infer

测试用例运行情况：[root@A03-R40-I191-99-4100006 python_test]# pytest test_x_attention_with_pa.py
============================================================================================================================================================== test session starts ==============================================================================================================================================================
platform linux -- Python 3.11.15, pytest-8.3.2, pluggy-1.6.0
rootdir: /export/home/ext.wangguangcai1/code/github/xllm-ops/test/python_test
plugins: forked-1.6.0, xdist-3.6.1, anyio-4.14.2
collected 1920 items

test_x_attention_with_pa.py ............................................................................................................................................................................................................................................................................................................. [ 15%]
......................................................................................................................................................................................................................................................................................................................................... [ 32%]
......................................................................................................................................................................................................................................................................................................................................... [ 49%]
......................................................................................................................................................................................................................................................................................................................................... [ 67%]
......................................................................................................................................................................................................................................................................................................................................... [ 84%]
............................................................................................................................................................................................................................................................................................................... [100%]

======================================================================================================================================================= 1920 passed in 1516.38s (0:25:16) =======================================================================================================================================================
[root@A03-R40-I191-99-4100006 python_test]# ls
RegisterOps.cpp pytorch_npu_helper.hpp test_causal_conv1d_qkv.py test_gamma_add_rms_norm.py test_layer_norm_fwd.py test_moe_gating_top_k_hash.py test_multi_latent_attention.py test_replace_token.py test_x_attention.py
pycache setup.py test_causal_conv1d_update.py test_hc_post.py test_lightning_indexer_quant.py test_moe_grouped_matmul.py test_onerec_final_beam_select.py test_reshape_and_cache_a5.py test_x_attention_with_pa.py
build test_add_rms_norm_bias.py test_compressor.py test_hc_pre_inv_rms.py test_lightning_indexer_quant_metadata.py test_moe_grouped_matmul_swiglu_quant.py test_pp_matmul_opt.py test_rms_norm_dynamic_quant.py test_x_flash_attention_infer.py
build_and_run.sh test_beam_search.py test_convert_kv_cache_format.py test_hc_pre_sinkhorn.py test_mega_chunk_gdn.py test_moe_grouped_matmul_swiglu_quant_v2.py test_quant_lightning_indexer.py test_scatter_nd_update_v2.py
custom_ops.egg-info test_beam_search_group.py test_dequant_swiglu_quant.py test_index_group_matmul.py test_moe_fused_add_topk.py test_moe_init_routing_custom.py test_quant_lightning_indexer_metadata.py test_select_unshared_kv.py
custom_ops.py test_cache_unshared_kv.py test_dispatch_ffn_combine.py test_inplace_partial_rotary_mul.py test_moe_fused_reducesum_div.py test_moe_init_routing_v3.py test_rec_constrained_topk.py test_sparse_attn_sharedkv.py
dist test_causal_conv1d.py test_dispatch_gmm_combine_decode.py test_laser_attention.py test_moe_gating_top_k.py test_mtp_prepare_next_draft.py test_recurrent_gated_delta_rule.py test_sparse_attn_sharedkv_metadata.py
[root@A03-R40-I191-99-4100006 python_test]#
[root@A03-R40-I191-99-4100006 python_test]#
[root@A03-R40-I191-99-4100006 python_test]# ls
RegisterOps.cpp pytorch_npu_helper.hpp test_causal_conv1d_qkv.py test_gamma_add_rms_norm.py test_layer_norm_fwd.py test_moe_gating_top_k_hash.py test_multi_latent_attention.py test_replace_token.py test_x_attention.py
pycache setup.py test_causal_conv1d_update.py test_hc_post.py test_lightning_indexer_quant.py test_moe_grouped_matmul.py test_onerec_final_beam_select.py test_reshape_and_cache_a5.py test_x_attention_with_pa.py
build test_add_rms_norm_bias.py test_compressor.py test_hc_pre_inv_rms.py test_lightning_indexer_quant_metadata.py test_moe_grouped_matmul_swiglu_quant.py test_pp_matmul_opt.py test_rms_norm_dynamic_quant.py test_x_flash_attention_infer.py
build_and_run.sh test_beam_search.py test_convert_kv_cache_format.py test_hc_pre_sinkhorn.py test_mega_chunk_gdn.py test_moe_grouped_matmul_swiglu_quant_v2.py test_quant_lightning_indexer.py test_scatter_nd_update_v2.py
custom_ops.egg-info test_beam_search_group.py test_dequant_swiglu_quant.py test_index_group_matmul.py test_moe_fused_add_topk.py test_moe_init_routing_custom.py test_quant_lightning_indexer_metadata.py test_select_unshared_kv.py
custom_ops.py test_cache_unshared_kv.py test_dispatch_ffn_combine.py test_inplace_partial_rotary_mul.py test_moe_fused_reducesum_div.py test_moe_init_routing_v3.py test_rec_constrained_topk.py test_sparse_attn_sharedkv.py
dist test_causal_conv1d.py test_dispatch_gmm_combine_decode.py test_laser_attention.py test_moe_gating_top_k.py test_mtp_prepare_next_draft.py test_recurrent_gated_delta_rule.py test_sparse_attn_sharedkv_metadata.py
[root@A03-R40-I191-99-4100006 python_test]# pytest test_x_flash_attention_infer.py
============================================================================================================================================================== test session starts ==============================================================================================================================================================
platform linux -- Python 3.11.15, pytest-8.3.2, pluggy-1.6.0
rootdir: /export/home/ext.wangguangcai1/code/github/xllm-ops/test/python_test
plugins: forked-1.6.0, xdist-3.6.1, anyio-4.14.2
collected 3 items

test_x_flash_attention_infer.py ... [100%]

=============================================================================================================================================================== 3 passed in 6.89s ===============================================================================================================================================================
[root@A03-R40-I191-99-4100006 python_test]#
[root@A03-R40-I191-99-4100006 python_test]#
[root@A03-R40-I191-99-4100006 python_test]# pytest test_x_attention.py
============================================================================================================================================================== test session starts ==============================================================================================================================================================
platform linux -- Python 3.11.15, pytest-8.3.2, pluggy-1.6.0
rootdir: /export/home/ext.wangguangcai1/code/github/xllm-ops/test/python_test
plugins: forked-1.6.0, xdist-3.6.1, anyio-4.14.2
collected 18 items

test_x_attention.py .................. [100%]

========================================================================================================================================================= 18 passed in 69.31s (0:01:09) =========================================================================================================================================================
[root@A03-R40-I191-99-4100006 python_test]# ls

文档更新

类型标签

Bug修复
新特性
性能优化
文档更新
其他，请描述：

Pull Request已成功合入, 合并人@CANN-robot

（感谢 ware2009 的贡献）

ware2009

3 天前创建了 pull request，commit 2802bd2e

atomgit-bot

3 天前评论：

🤖 正在生成合并请求摘要，请稍候…

atomgit-bot

3 天前评论：

🤖 AI 代码检视正在进行中，请稍候…

CANN-robot成员

3 天前添加了label：cann-cla/yes

CANN-robot成员

3 天前评论：

CLA Signature Pass

ware2009, thanks for your pull request. All authors of the commits have signed the CLA. 👍

此处折叠了110条消息查看更多

sunhao_hw成员

1 天前评论：

/lgtm
/approve

CANN-robot成员

1 天前添加了label：approved

yuantao成员

1 天前评论：

/lgtm

CANN-robot成员

1 天前添加了label：lgtm

CANN-robot成员

1 天前合入了pull request