[Triton Kernel] Add varlen segment mean triton kernel

[l1cacheDell]

PR types

New features

PR changes

Others

Description

Add segment mean triton kernel. Implemented for segment mean operation for varlen qkv.

For example, the k tensor is: [total_seqlen, num_head, head_dim], where total_seqlen = seqlen 1 + seqlen 2 + ... + seqlen n.

So the segment mean triton kernel will do mean operation along the seqlen dim.

It will finally generate a [bsz, num_head, head_dim] shape-like result, as the result of mean value of each seqlen segment.

[paddle-bot]

Thanks for your contribution!

[codecov]

Codecov Report

Attention: Patch coverage is 17.74194% with 51 lines in your changes missing coverage. Please review.

Project coverage is 48.97%. Comparing base (f53e113) to head (b630601).
Report is 171 commits behind head on develop.

Files with missing lines	Patch %	Lines
paddlenlp/ops/triton_ops/segment_mean.py	16.39%	51 Missing ⚠️

Additional details and impacted files

@@             Coverage Diff             @@
##           develop   #10369      +/-   ##
===========================================
- Coverage    48.99%   48.97%   -0.02%     
===========================================
  Files          765      766       +1     
  Lines       125974   126036      +62     
===========================================
+ Hits         61720    61731      +11     
- Misses       64254    64305      +51     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.