Activation

• Currently implemented
  • PolyNorm
  • RMSNorm
  • FusedAddRMSNorm
    A fused operator that combines residual addition (x + residual) with RMSNorm in a single kernel.
    • Instead of:
      python

      1y = x + residual 2hidden_state = rms_norm(y, weight, eps) 3out = y + some_op(hidden_state)
    • Fused as:
      python

      1hidden_state, y = fused_add_rms_norm(x, residual, weight, eps) 2out = y + some_op(hidden_state)
  • FusedMulPolyNorm
    A fused operator that combines PolyNorm with an element-wise multiplication by a Tensor.
    • Instead of:
      python

      1y = poly_norm(x, weight, bias, eps) 2out = y * a
    • Fused as:
      out = fused_mul_poly_norm(x, a, weight, bias, eps)
  • FusedMulGroupedPolyNorm (CUDA)
    A CUDA-accelerated grouped variant of FusedMulPolyNorm for MoE (Mixture of Experts) models. Fuses the entire PolyNorm computation into CUDA kernels (fwd + bwd), with per-expert weights/bias, in-kernel binary search for expert mapping, optional routing scores multiplication, and hidden_clamp fusion.
    • Instead of:
      python

      1for i, expert in enumerate(experts): 2 out[start:end] = fused_mul_poly_norm(x[start:end], mul[start:end], weight[i], bias[i], eps)
    • Fused as:
      python

      1out = fused_mul_grouped_poly_norm(x, mul, weight, bias, offsets, eps, 2 scores=scores, hidden_clamp=10.0)

Installation

Usage

Performance

• Test cases are from the Motif LLM
• The results can be reproduced using the provided benchmarking tools.
• For details on how to use the benchmarking tools, please refer to the benchmarks README.
• The benchmark results may show fluctuations, especially in the backward pass and when the dimension size is small.

RMSNorm

H100 Results

Forward Performance

Backward Performance

MI250 Results

Forward Performance

Backward Performance

FusedAddRMSNorm

[!NOTE]
For fusion case performance, the non-fused baseline was implemented with our custom kernels.

H100 Results

Forward Performance

Backward Performance

MI250 Results

Forward Performance

Backward Performance

PolyNorm

H100 Results

Forward Performance

Backward Performance

MI250 Results

Forward Performance

Backward Performance

FusedMulPolyNorm

[!NOTE]
For fusion case performance, the non-fused baseline was implemented with our custom kernels.

H100 Results

Forward Performance

Backward Performance

MI250 Results

Forward Performance

Backward Performance

FusedMulGroupedPolyNorm (CUDA)

[!NOTE] This kernel is implemented in CUDA C++ (compiled via setup.py). Benchmarks compare three variants: Naive (raw PyTorch reference), Compiled (torch.compile'd reference), and CUDA (fused CUDA kernel). Benchmark dimension: 1280, 384 experts.

Training profile (B200, motif3_seq, lbs=8, seqlen=4K):

	CUDA kernel	torch.compile	Speedup
Forward	0.7 ms	2.1 ms	3.0x
Backward	1.4 ms	3.7 ms	2.6x

Pre-commit Hooks

Setup

1. Install pre-commit:
   pip install pre-commit
2. Install the git hooks:

Included Hooks

• yapf – Python code formatter
• typos – Spell checker for common typos
• isort – Organizes and sorts Python imports
• clang-format – Formats C++/CUDA code (--style=file)
• pymarkdown – Lints and auto-fixes Markdown files
• actionlint – Validates GitHub Actions workflows

Usage

• Run all checks on the entire codebase:
  pre-commit run --all-files
• Run a specific hook (example: isort):