Commit History
-
refactor: replace warp shuffle with CUB BlockReduce
- Replace hand-written warp_reduce_sum/block_reduce_sum with CUB BlockReduce - Add BLOCK_SIZE template param to fwd_small for compile-time CUB sizing - Dispatch fwd_small with matched BlockReduce<float, 32/64/128/256> - Simpler code, no performance regression (small dims slightly faster) Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>79a877a
-
perf: optimize RMSNorm CUDA kernels for all dims
Forward: - Single-pass register caching (dim ≤ 2048) - 2-pass with L1 re-read (dim > 2048) - Warp shuffle reduction replacing CUB (98% shared mem reduction) - Separate kernel functions for compiler-independent optimization Backward: - Save inv_rms from forward, eliminate sum_square recomputation - Multi-token block (32 tok/block) with shared mem weight grad (dim ≤ 8192) - 2-token batched float2 reduction (halves syncthreads) - Output-based backward (dim > 8192): save y instead of x, 33% traffic reduction - Column-parallel vec8 weight_grad kernel with token chunking (dim > 8192) - __launch_bounds__ tuning for register allocation Interface change: rms_norm returns (output, inv_rms) tuple. Backward takes output tensor + inv_rms instead of input. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>dc88599
-
style: apply yapf, isort, and clang-format
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>6436ad6 -
refactor: remove Triton kernels, add hidden_clamp to unscored ops
- Remove all Triton kernel code (fwd/bwd kernels, autotune configs, triton import) — replaced by CUDA kernels in grouped_poly_norm.cu - Add hidden_clamp parameter to unscored C++ ops (forward/backward) so both scored and unscored paths support clamping - Update register_fake, autograd Function, and dispatch for unscored ops - Replace HAS_TRITON with _has_cuda_ops in tests Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>906e125 -
feat: add grouped poly norm CUDA kernel with scores and hidden_clamp fusion
Hand-written CUDA kernel for GroupedFusedMulPolyNorm (MoE). Fuses polynomial normalization, mul, scores multiplication, and hidden_clamp into a single kernel launch per forward/backward. - 4 kernel variants: fwd/bwd × vectorized(width=8)/scalar - scores: nullable, always fp32, fused into fwd output and bwd gradients - hidden_clamp: < 0 disabled, >= 0 clamps input/mul/output with correct backward gradient masking (recompute output for mask, no extra memory) - ROCm compatible (64-bit warp sync mask) - C++ op registration for torch.compile (register_fake in Python layer) - ~3x faster than torch.compile'd PyTorch reference on B200 Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>0045757
-
refactor(activation): change fused_add_rms_norm and fused_add_rms_norm_backward to out-place operations
7e4334d
-
feat: make rms_norm as out-place
9d0a235
-
Fix fused add rms norm (#4)
TaehyunKim +1: TaehyunKimMotif Unverified signature* make fused add rms norm return 2 outputs * add builds * add plot --------- Co-authored-by: taehyun <taehyun.kim@motiftech.io>a1e5ca8
-
Add fusion (#3)
TaehyunKim +1: TaehyunKimMotif Unverified signature* vectorize and optimized block reduce * add benchmark test (w/o readme update) * implemented fused_mul_poly_norm Signed-off-by: taehyun <taehyun.kim@motiftech.io> * add_rms_norm added * deleted backward pass on fused add rms norm, split test and benchmarks Signed-off-by: taehyun <taehyun.kim@motiftech.io> * refactored benchmarks * add readme * fix readme * add build * fix readme * fix readme2 * add mi250 results * highlight used our kernel for baseline in fused performance * applied yapf --------- Signed-off-by: taehyun <taehyun.kim@motiftech.io> Co-authored-by: taehyun <taehyun.kim@motiftech.io>e5e2eeb
-
Optimize kernel (#2)
Verified signature- add poly_norm_naive.cui for temp test (18ec195c33e9fc3d259cbef6b6e6cdaf077c4921) - optimize poly norm kernel (f8a7e6f9a378e4860688af1c856fc285f7fdb5bc) - add pre-commit (a73a0c0cfd8960f3041692d9e2b70e93550277a3) - add readme with precommit hooks and applied pre commit to all files (f517c97d5ea3aa5479d47cfb5a11ca070972aad4) - add build (56776acf3c7345b3a5b71ea5aec28ee8ee59d0fa)97825b8
-
feat(rms-norm): Impl fused RMSNorm
f3b99fb
-
chore: use latest build image and misc
f5a7d38
-
feat(poly-norm): Add PolyNorm
44e9845