Commit History
-
Merge pull request #22 from MotifTechnologies/jangwoong/mla-rope-fa4-port
Jangwoong KimUnverified signatureoptimize: MLA RoPE + FA4 wrapper Triton kernels (port from llm-training)5adea7d
-
cleanup: drop k_pe RoPE custom kernel (caller uses PyTorch native)
Remove _rope_fwd_kernel, _rope_bwd_kernel, their custom_op wrappers, FusedRoPE autograd class, and the fused_apply_rope user-facing wrapper (plus the activation.__init__ export). Rationale: k_pe in MLA is head-shared (H=1), so the kernel handles a trivially small problem (B*S*1*rope_dim ≈ 8 MB at B=8, S=4096). At that size the kernel is launch-bound, peak HBM utilization sits around 5%, and a standalone Triton launch costs more than letting inductor inline the work into the surrounding graph. The caller now uses the same view_as_complex + reorder_headdim_elements_rope path as the upstream PyTorch implementation. Verified at B=8, S=4096, n_layers=4 on B200: step-10 loss matches (0.0001), aggregate RoPE-path savings vs vanilla unchanged at -68%. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>7e86d2e -
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
-
fix: unify all backward kernels to input-based math + fix test import
- Always save input in setup_context (no output-based path) - bwd_fused: input-based math (d_sum = sum(dy*x*w), grad = scale*dy*w - dxx*x) - bwd_large_input_grad: input-based math - bwd_large_weight_grad: input-based (weight_grad = dy*x*scale, needs inv_rms) - bwd_scalar: input-based (unchanged) - pytest.ini: add pythonpath to prevent activation/ CUDA dir namespace shadow Eliminates bf16 precision loss from output-based x recovery (y/(w*scale)). All 48 test_rms_norm tests pass. No performance regression. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>09ecd67 -
style: fix yapf/isort/clang-format for CI --all-files
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>9dcee96 -
feat: update RMSNorm Python interface for optimized kernels
- RMSNormFunction.forward returns (output, inv_rms) tuple - setup_context saves output (dim > 8192) or input (dim ≤ 8192) - backward passes saved tensor + inv_rms to C++ kernel - Updated fake tensor registration for tuple return - Updated DTensor sharding strategies for new backward signature Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>4bb42a5 -
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
-
feat: dedicated _kv_rope_bwd_kernel (register-sum + copy-fused)
FusedKVRoPE.backward used to chain three inductor-generated kernels: slice (grad_key nope) + cat (grad_k_nope, grad_value) + sum-over-H (grad_key rope section → grad_k_pe) measuring 452 µs (cat_slice_5) + 346 µs (slice_sum_4) = 798 µs for 5 backward calls at B=8, S=4096, n_layers=4 on B200. Replace with a single Triton kernel that, per token: - copies grad_key[:, :, :k_nope_dim] → grad_kv_latent nope section - copies grad_value → grad_kv_latent v section - sums grad_key[:, :, k_nope_dim:] across H (fp32 accum) → grad_k_pe Grid is one program per token; BLOCK_H=16 covers H_kv=16 in a single block so the head-reduction stays register-level (no atomics). The reduction runs in fp32 to avoid bf16 precision loss when summing 16 heads. Static config (BLOCK_H=16, warps=4, stages=2) picked for this shape; other shapes would need re-measurement, documented at top of file alongside the other static configs. Dropped `ctx.save_for_backward(k_pe)` from FusedKVRoPE.fo…35a25ee -
perf: remove autotune, hard-code per-kernel configs from live dump
Autotune isn't buying anything for these kernels: at runtime the shapes are fixed by the model config (H ∈ {1, 80}, rope_dim=64), and live-dump (TRITON_PRINT_AUTOTUNING=1) on motif3_seq B=8 S=4096 n_layers=4 showed every kernel's best config sitting in a tight cluster that doesn't change across reasonable workload sizes. Replace @triton.autotune + config helpers with a small dict of static (BLOCK_H, num_warps, num_stages) per kernel, passed at each kernel launch. Values are the plurality-pick from the live dump: _rope_fwd_kernel BLOCK_H=1, warps=1, stages=3 # k_pe, H=1 _rope_bwd_kernel BLOCK_H=2, warps=1, stages=1 # k_pe bwd _kv_rope_fwd_kernel BLOCK_H=32, warps=8, stages=2 # H=80 _q_rope_inplace_fwd_kernel BLOCK_H=8, warps=1, stages=2 # H=80 _q_rope_bwd_kernel BLOCK_H=4, warps=2, stages=2 # H=80 Validation (same shape, autotune-full-sweep vs no-autotune, both B=8 n_layers=4): step 10 loss matches (0.0001), per-kernel t…1e2bc2b -
cleanup: remove dead Phase 3 Q kernel + shrink autotune to hand-picked configs
Changes: - Delete _q_rope_fwd_kernel, _q_rope_fwd custom_op, FusedQRoPE autograd class, fused_q_rope wrapper. These were Phase 3 (out-of-place Q nope passthrough + RoPE) kernels, reverted in llm-training@4761dc9 for being slower than Phase 2, kept as "reference but unused" and then carried over in the activation port. Phase 4 in-place kernel fully subsumes them. Verified via TRITON_PRINT_AUTOTUNING=1: the kernel is never triggered during a full training step, and nothing outside fused_rope.py references it (activation/__init__.py exports only fused_q_rope_inplace). - _q_rope_bwd_kernel + _q_rope_bwd custom_op stay; they are shared with FusedQRoPEInplace.backward. - Shrink autotune search space to hand-picked configs. Dumping the live best-config per kernel showed tight clusters: fwd ∈ {BLOCK_H=2/4/8, warps=2/4, stages=2}, bwd ∈ {BLOCK_H=2/4, warps=2, stages=2}. Trials per shape 328 → 18 (~94% fewer); every observed best config is still in the list. _rope_bwd…4d94a7d -
review fixups: stride asserts, autotune split, intent comments
- assert x.stride(-1) == 1 on all 6 custom_ops (M1): enforce last-dim unit stride assumption used by all RoPE kernels. - Split _rope_autotune_configs() into fwd (18 cfgs) / bwd (128 cfgs) (M2): fwd sweep gain is ±noise, bwd gain is ~-28%. Total warm-up configs 6×128 → 4×18+2×128 (~57% fewer). - assert k_pe.shape[1] == 1 in _kv_rope_fwd (N3): enforce MLA head-shared convention; block GQA/MQA variant misuse. - assert q.is_contiguous() in FusedQRoPEInplace.forward (N4): guarantee reshape is a view so in-place mutation reaches the original q that mark_dirty targets. - Intent comments for autotune key design (N1), `-> None` in-place signature (N2a), and `return q` after mutation (N2b). Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>2712745
-
feat: MLA RoPE Triton kernels (port from llm-training)
Ports Motif3 GDLAttention RoPE optimization kernels from llm-training to the activation package. Python/Triton-only addition — no CUDA recompilation needed. ## Kernels added (fused_rope.py) ### _rope_fwd_kernel / _rope_bwd_kernel (Phase 1) Generic RoPE on [B*S, H, rope_dim] with sin/cos direct multiplication, interleaved → contiguous layout change inside kernel. Eliminates view_as_complex + complex_mul + reorder_headdim pattern. Used for k_pe (1 head, standalone) after Phase 4 landed (Q path uses the in-place variant below). ### _kv_rope_fwd_kernel (Phase 2) Fused KV split + k_pe RoPE broadcast + cat. Reads kv_latent (k_nope+v) and k_pe (1 head, already RoPE'd), outputs key=[k_nope||k_pe_broadcast] and value. k_pe broadcast across heads happens in register — no memory copy. ### _q_rope_inplace_fwd_kernel / _q_rope_bwd_kernel (Phase 4) In-place Q RoPE. Modifies q[..., nope_dim:] from interleaved to contiguous format without touching the nope section. Eliminates the subsequent cat(q_…f61868b -
style: fix yapf/isort formatting for CI --all-files check
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>3f2678c
-
grouped polynorm with padding aware (#19)
TaehyunKimUnverified signature972d63b
-
chore: remove pre-built binaries and add local build loader shim (#18)
Unverified signatureDrop the entire build/ tree of pre-built .abi3.so binaries and per-variant Python wrappers. They were committed via Git LFS for every PyTorch/CUDA/ROCm combo and bloated the repo without serving the local pip-install workflow. Add torch-ext/activation/_ops.py as a small loader shim that works with both kernel-builder builds (_activation_HASH module) and local setup.py builds (_activation), so `pip install -e .` users can import the package without a committed build/ tree. Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>1e08296
-
style: apply yapf, isort, and clang-format
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>6436ad6 -
style: fix clang-format on torch_binding.h
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>344ed39 -
fix: rename stale references and clean up Triton remnants
- Fix broken imports in benchmarks (grouped_fused_mul_poly_norm → fused_mul_grouped_poly_norm) - Rename GroupedTritonModule → GroupedCUDAModule in benchmarks - Rename _run_triton → _run_cuda in tests, update docstrings - Remove stale Triton comments in tests - Fix return type annotations in __init__.py (None → torch.Tensor) - Update README with new function name - Remove TRITON_PRINT_AUTOTUNING env var from profile_bwd.py Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>5a9d09d -
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: rename grouped_fused_mul_poly_norm → fused_mul_grouped_poly_norm
Align naming convention with existing ops: - fused_mul_poly_norm (poly_norm + mul fusion) - fused_mul_grouped_poly_norm (grouped_poly_norm + mul fusion) Also renames the ref impl: fused_mul_grouped_poly_norm_ref Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>60a628a
-
feat: add GroupedFusedMulPolyNorm Triton kernel for MoE models (#16)
TaehyunKim +2: Claude Opus 4.6, github-actions[bot]Unverified signature* feat: add GroupedFusedMulPolyNorm Triton kernel for MoE models Fuses the full PolyNorm computation into two Triton kernels (fwd + bwd) with per-expert weights/bias and in-kernel binary search for expert mapping. Includes benchmarks, tests, and README documentation with B200 results. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> * Add built binary [skip-build] --------- Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com> Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>e195bbb
-
fix: support PyTorch 2.10 register_op_strategy import path change
register_op_strategy moved from torch.distributed.tensor._ops.utils to torch.distributed.tensor._ops.registration in PyTorch 2.10. Use try/except fallback to support both 2.9.x and 2.10.x. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>ad23c2a
-
fix: fix fused add rms norm sharding strategy
a35a092
-
fix: fix rms norm sharding strategy
138159c
-
fix(rms_norm.py): add assertion for input gradients to handle unsupported cases in backward pass
f19f8f4 -
feat: support sequence parallel with fused_add_rms_norm
151bb5a -
refactor(activation): change fused_add_rms_norm and fused_add_rms_norm_backward to out-place operations
7e4334d -
refactor(rms_norm): move RMS normalization logic to a new module for better organization and maintainability
66b3c5e
-
feat: support sequence parallel with rms_norm
06d6367 -
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: support reset_parameters()
605f22e -
feat(rms-norm): Impl fused RMSNorm
f3b99fb
-
refactor(poly-norm): use const for immutable args
e85ecc9 -
chore: use latest build image and misc
f5a7d38
-
feat(poly-norm): add default value for eps argument
afd2a56 -
feat(poly-norm): Add PolyNorm
44e9845