Skip to content

[KernelGen][Nvidia] Add addr_ operator with Triton kernel - #5440

Open
LoserCheems wants to merge 1 commit into
flagos-ai:masterfrom
flash-algo:pr/nv-addr_
Open

[KernelGen][Nvidia] Add addr_ operator with Triton kernel#5440
LoserCheems wants to merge 1 commit into
flagos-ai:masterfrom
flash-algo:pr/nv-addr_

Conversation

@LoserCheems

Copy link
Copy Markdown
Contributor

Summary

Adds a Triton kernel for addr_ (in-place variant). Performs the outer-product of vectors vec1 and vec2, scales by alpha, and adds to beta * input in-place using a single fused kernel.

Testing

  • Parametrized tests over (M, N) shapes and float16/float32/bfloat16 dtypes
  • Validated against PyTorch reference on device side
  • Tested on: Nvidia, Tianshu, Muxi, Ascend, Hygon

Performance

Test command: pytest benchmark/test_addr_.py --level core --mode cudagraph (NVIDIA H20)

float16:

Shape Torch (ms) Gems (ms) Speedup
[2, 384] 0.0037 0.0010 3.71x
[2, 4096] 0.0024 0.0012 1.93x
[16, 1024] 0.0024 0.0011 2.23x
[16, 2048] 0.0047 0.0023 2.05x
[16, 4096] 0.0048 0.0011 4.29x
Arithmetic Mean 2.84x

float32:

Shape Torch (ms) Gems (ms) Speedup
[2, 384] 0.0015 0.0010 1.47x
[2, 4096] 0.0014 0.0012 1.13x
[16, 1024] 0.0017 0.0012 1.34x
[16, 2048] 0.0017 0.0013 1.34x
[16, 4096] 0.0017 0.0013 1.37x
Arithmetic Mean 1.33x

bfloat16:

Shape Torch (ms) Gems (ms) Speedup
[2, 384] 0.0020 0.0010 2.05x
[2, 4096] 0.0047 0.0022 2.11x
[16, 1024] 0.0048 0.0011 4.46x
[16, 2048] 0.0023 0.0012 1.92x
[16, 4096] 0.0047 0.0023 2.02x
Arithmetic Mean 2.51x

Multi-backend Testing

Backend Accuracy Test Benchmark Speedup (mean) Notes
Nvidia (H20) PASS PASS (15 cases, --level core) 2.23x Primary
Tianshu Not tested
Muxi Not tested
Ascend Not tested
Hygon Not tested

Files Changed

  • src/flag_gems/ops/addr_.py: Triton kernel implementation
  • tests/test_addr_.py: Accuracy test
  • benchmark/test_addr_.py: Performance benchmark
  • src/flag_gems/ops/__init__.py: Register import and __all__
  • src/flag_gems/__init__.py: Register to _FULL_CONFIG
  • conf/operators.yaml: Add operator entry (kind: BLAS, stage: stable 4.0)

Implements the in-place variant of addr using a single fused Triton
kernel. The kernel computes beta * input + alpha * outer(vec1, vec2)
and writes the result back to the input tensor directly.

Benchmark results on H20 (vs PyTorch native addr_):
- float16 average speedup: 1.89x
- float32 average speedup: 1.34x
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.

1 participant