Skip to content

[KernelGen][Kunlunxin] Add digamma operator - #290

Merged
Dongxu-H merged 3 commits into
flagos-ai:infra-cifrom
NineAnnAnn:kunlunxin-digamma-opt
Aug 18, 2026
Merged

[KernelGen][Kunlunxin] Add digamma operator#290
Dongxu-H merged 3 commits into
flagos-ai:infra-cifrom
NineAnnAnn:kunlunxin-digamma-opt

Conversation

@NineAnnAnn

Copy link
Copy Markdown
Contributor

Summary

Adds a Kunlunxin (P800) backend specialization for digamma. Computes the logarithmic derivative of the gamma function element-wise on the input tensor. Adds a dedicated Triton kernel for this backend and registers it in the backend ops module.

Testing

  • Validated against reference on device (Kunlunxin (P800))
  • Accuracy: PASS (36 cases) over 3 dtypes x 6 shapes

Performance

Test command: pytest benchmark/test_digamma.py --level core (Kunlunxin (P800))

digamma

dtype Size Torch Latency (ms) Gems Latency (ms) Speedup TFLOPS
float16 [64, 64] 0.184460 0.008181 22.547 0.001
float16 [4096, 4096] 59.264362 1.440451 41.143 0.012
float16 [64, 512, 512] 58.909332 1.440531 40.894 0.012
float32 [64, 64] 0.156772 0.006792 23.082 0.001
float32 [4096, 4096] 80.619667 1.347136 59.845 0.012
float32 [64, 512, 512] 80.892136 1.347404 60.036 0.012
bfloat16 [64, 64] 0.433782 0.007068 61.373 0.001
bfloat16 [4096, 4096] 63.654781 1.565451 40.662 0.011
bfloat16 [64, 512, 512] 62.773350 1.564299 40.129 0.011
Operator Geometric Mean Speedup
digamma 40.825

Multi-backend Testing

Backend Accuracy Test Speedup (mean) Notes
Kunlunxin (P800) PASS (36 cases) 40.825 Primary

Files Changed

  • src/flag_gems/runtime/backend/_kunlunxin/ops/digamma.py: New Kunlunxin (P800) kernel implementation
  • src/flag_gems/runtime/backend/_kunlunxin/ops/__init__.py: Register import and __all__ entry

Add a kunlunxin backend specialization for digamma with a tuned Triton kernel and register it in the backend ops module. Accuracy validated on device; benchmark measured on Kunlunxin (P800).
@Dongxu-H

Copy link
Copy Markdown
Collaborator

/test | digamma:kunlunxin

@github-actions

Copy link
Copy Markdown

❌ On-demand test failed

Operator: digamma
Runner: kunlunxin
Backend: kunlunxin

The test failed to complete. Check the workflow run for details.

@Dongxu-H Dongxu-H left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

code-style CI failed
all 中的 digamma 放在 zeros_like 之后不符合字母顺序,应放在 digamma_ 之前
logger 应使用 Kunlunxin backend 约定写法而非主文件夹写法

@Dongxu-H
Dongxu-H merged commit d73b9c3 into flagos-ai:infra-ci Aug 18, 2026
9 checks passed
103yiran pushed a commit that referenced this pull request Aug 18, 2026
Co-authored-by: NineAnnAnn <NineAnnAnn@users.noreply.github.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants