lrnnx.ops.triton package

Triton kernel implementations.

Submodules

• lrnnx.ops.triton.layer_norm module
  • layer_norm_ref()
  • rms_norm_ref()
  • config_prune()
  • LayerNormFn
    • LayerNormFn.forward()
    • LayerNormFn.backward()
  • layer_norm_fn()
  • rms_norm_fn()
  • RMSNorm
    • RMSNorm.reset_parameters()
    • RMSNorm.forward()
  • LayerNormLinearFn
    • LayerNormLinearFn.forward()
    • LayerNormLinearFn.backward()
  • layer_norm_linear_fn()
• lrnnx.ops.triton.selective_state_update module
  • selective_state_update()
  • selective_state_update_ref()
• lrnnx.ops.triton.simplified_state_update module
  • simplified_state_update()
  • simplified_state_update_ref()
• lrnnx.ops.triton.softplus module
  • softplus()