Skip to content

feat(nn): add positional encoding operations (PoPE, ALiBi, YaRN, NTK, ReLU2)#174

Merged
m96-chan merged 2 commits intomainfrom
feature/nn-positional-encoding
Dec 30, 2025
Merged

feat(nn): add positional encoding operations (PoPE, ALiBi, YaRN, NTK, ReLU2)#174
m96-chan merged 2 commits intomainfrom
feature/nn-positional-encoding

Conversation

@m96-chan
Copy link
Copy Markdown
Owner

Summary

Implements issues #136, #167, #169, #171:

New Files

Kernels:

  • native/ops/nn/activation/relu2.inl
  • native/ops/nn/pope/pope_kernels.cuh, pope.inl
  • native/ops/nn/alibi/alibi_kernels.cuh, alibi.inl
  • native/ops/nn/rope/rope_ext_kernels.cuh, rope_ext.inl

Tests:

  • tests/test_relu2.py
  • tests/test_positional_encoding.py

Test Plan

  • Build passes (SM 120a, CUDA 13.1)
  • Pre-commit checks pass (Ruff, Mypy)
  • New tests: 22 passed, 2 skipped (GPU-only bf16/f16)
  • Full test suite: 260 passed, 11 failed (pre-existing GPU-only failures)

Closes #136, #167, #169, #171

🤖 Generated with Claude Code

m96-chan and others added 2 commits December 31, 2025 06:56
… ReLU2)

Implements issues #136, #167, #169, #171:

- ReLU2 activation (squared ReLU from Primer paper)
- PoPE (sinusoidal positional encoding, additive alternative to RoPE)
- ALiBi (attention with linear biases)
- YaRN/NTK-aware RoPE extensions for context length scaling
- Linear position interpolation for RoPE

New kernels:
- native/ops/nn/activation/relu2.inl
- native/ops/nn/pope/pope_kernels.cuh, pope.inl
- native/ops/nn/alibi/alibi_kernels.cuh, alibi.inl
- native/ops/nn/rope/rope_ext_kernels.cuh, rope_ext.inl

Tests: 22 passed, 2 skipped (GPU-only bf16/f16 tests)

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
Add CPU implementations for PoPE, ALiBi, and RoPE extension functions
to support testing in CI environments without GPU:

- rope_init_ntk_aware: NTK-aware frequency scaling
- rope_init_yarn: YaRN dimension-wise interpolation
- rope_init_linear: Linear position interpolation
- pope_init_encoding: Sinusoidal positional encoding
- pope_inplace: Additive positional encoding
- alibi_init_slopes: ALiBi head slopes
- alibi_compute_bias: ALiBi bias matrix
- alibi_add_bias: In-place bias addition

Also fixed in-place update bug in _rope_inplace_cpu using proper
backend.copy_host_to_device() instead of non-existent _data attribute.

Fixes CI test failures for positional encoding tests.

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

feat(nn): Implement PoPE (Positional Encoding) alongside RoPE

1 participant