feat: import and extend ML kernels from inferi - #11

Merged

sebcrozet merged 13 commits into

mainfrom

Aug 27, 2026

Commits on Aug 26, 2026

feat: move inferi shaders into an optional ml module for vortx
sebcrozet
committed
6aecbf9

Commits on Aug 27, 2026

chore: cargo fmt
sebcrozet
committed
514eb35
feat(ml): add tanh activation + Adam optimizer GPU ops

sebcrozet
and
haixuanTao
committed
115504c
refactor(ml): move the tanh-backward + Adam kernels into vortx::ml and drop the duplicated tanh forward
sebcrozet
committed
fd94b78
feat(ml): add ELU activation GPU ops (fwd/backward/vec4)

sebcrozet
and
haixuanTao
committed
e9f89f8
refactor(ml): keep only the ELU backward pass, the forward is already UnaryOp::Elu
sebcrozet
committed
d074535
feat(ml): add PPO loss-gradient GPU kernels

sebcrozet
and
haixuanTao
committed
3b8b06b
refactor(ml): wire the PPO kernels to vortx::ml and use StepRng for uniform control flow
sebcrozet
committed
e3b30c1
perf(linalg): vec4 GEMM (compute-FMA inner loop + vec4 global-load variant)

sebcrozet
and
haixuanTao
committed
2a5ad2d
perf(linalg): keep the vec4 FMA inner loop, drop the unmeasured vec4 global-load GEMM variant
sebcrozet
committed
7f46b2d
feat(ml): gpu_ppo_stage_batch, building the PPO minibatch on device
sebcrozet
committed
71f7cb1
chore(ml): comment cleanups
sebcrozet
committed
eb1d7d0
chore: CI fixes
sebcrozet
committed
bb38a3d