nemo_automodel.components.moe.optimized_ops

Memory-optimized MoE elementwise ops extracted from experts.py.

Chunked custom-autograd router-weight fp32 multiply: computes the identical fp32 math in row chunks and saves only low-precision inputs, removing the full-size fp32 intermediates that otherwise pin ~7 GiB blocks per MoE layer under activation checkpointing.

With BackendConfig.compile_router_weight the same Function runs each pass as one torch.compiled kernel over the whole tensor instead of the eager chunk loop (see _compile_router_weight_cores).

Module Contents

Classes

Name	Description
`_RouterWeightMulFunction`	Chunked fp32 router-weight multiply that saves only the raw inputs.

Functions

Name	Description
`_apply_router_weight_fp32`	Apply routing probabilities to expert outputs with fp32 arithmetic.
`_compile_router_weight_cores`	Wrap the router-weight multiply cores with `torch.compile`.
`_rw_bwd_probs_core`	Gradient w.r.t. `probs`: `(grad_out * x).sum(-1)` in fp32, cast to `probs_dtype`.
`_rw_bwd_x_core`	Gradient w.r.t. `x`: `(grad_out * probs)` in fp32, cast to `x_dtype`.
`_rw_fwd_core`	fp32 router-weight multiply: `(x * probs)` in fp32, cast to `out_dtype`.

Data

_RW_CHUNK_ROWS

_RW_CHUNK_THRESHOLD

_RW_CORES_COMPILED

_rw_bwd_probs_dispatch

_rw_bwd_x_dispatch

_rw_fwd_dispatch

API

class nemo_automodel.components.moe.optimized_ops._RouterWeightMulFunction()

Bases: Function

Chunked fp32 router-weight multiply that saves only the raw inputs.

The plain (x.float() * probs.float()).to(dtype) lets autograd keep full-size fp32 [tokens, hidden] intermediates alive for backward (the upcast input, the product, and a recompute copy under activation checkpointing). This Function computes the same fp32 multiply in row chunks and saves only the low-precision inputs. The forward is bitwise-identical; the backward matches autograd’s fp32 chain (grad_x = (g_f32 * probs_f32).to(x.dtype), grad_probs = (g_f32 * x_f32).sum(-1, keepdim=True) cast to probs.dtype).

nemo_automodel.components.moe.optimized_ops._RouterWeightMulFunction.backward(
    ctx: typing.Any,
    grad_out: torch.Tensor
) -> tuple[torch.Tensor | None, torch.Tensor | None, None, None]

staticmethod

Compute chunked fp32 gradients for the router-weight multiply.

Parameters:

ctx

Any

Autograd context holding the saved inputs.

grad_out

torch.Tensor

Upstream gradient of shape [tokens, hidden].

Returns: torch.Tensor | None

Tuple (grad_x, grad_probs, None, None) where grad_x has

nemo_automodel.components.moe.optimized_ops._RouterWeightMulFunction.forward(
    ctx: typing.Any,
    x: torch.Tensor,
    probs: torch.Tensor,
    out_dtype: torch.dtype,
    save_x: bool
) -> torch.Tensor

staticmethod

Multiply expert outputs by routing probabilities in fp32.

Parameters:

ctx

Any

Autograd context.

torch.Tensor

Expert outputs of shape [tokens, hidden].

probs

torch.Tensor

Routing probabilities of shape [tokens, 1].

out_dtype

torch.dtype

Output dtype (the dispatcher’s expected activation dtype, or fp32 for the scatter-add reduction path).

save_x

bool

Whether backward needs x. x is only consumed by the probs gradient; when probs carries no grad (e.g. FakeBalancedGate emits constant weights), saving it would pin a full-size [tokens, hidden] tensor per MoE layer across the activation-checkpointing backward window for nothing. Callers pass probs.requires_grad.

Returns: torch.Tensor

Tensor of shape [tokens, hidden] and dtype out_dtype.

nemo_automodel.components.moe.optimized_ops._apply_router_weight_fp32(
    output2: torch.Tensor,
    permuted_probs: torch.Tensor,
    compute_dtype: torch.dtype
) -> torch.Tensor

Apply routing probabilities to expert outputs with fp32 arithmetic.

Large 2-D row-aligned inputs go through the chunked custom Function so autograd does not retain full-size fp32 intermediates; every other shape keeps the plain eager multiply (bitwise-identical result).

Parameters:

output2

torch.Tensor

Expert down-projection outputs of shape [tokens, hidden].

permuted_probs

torch.Tensor

Routing probabilities broadcastable against output2, typically of shape [tokens, 1].

compute_dtype

torch.dtype

Output dtype.

Returns: torch.Tensor

(output2 * permuted_probs) computed in fp32 and cast to

nemo_automodel.components.moe.optimized_ops._compile_router_weight_cores() -> None

Wrap the router-weight multiply cores with torch.compile.

Runs once per process, same lazy pattern as Kimi K3’s compile_situ: the compiled functions replace the module-level dispatch targets, so every expert module shares one compiled kernel per pass and repeated model construction does not recompile. Fused kernels keep no fp32 intermediates, so the chunk loop (which only exists to bound eager fp32 transients) is skipped on this path. Numerics are allclose to eager (same fp32 math, one rounding to the output dtype), not guaranteed bitwise-identical.

nemo_automodel.components.moe.optimized_ops._rw_bwd_probs_core(
    grad_out: torch.Tensor,
    x: torch.Tensor,
    probs_dtype: torch.dtype
) -> torch.Tensor

Gradient w.r.t. probs: (grad_out * x).sum(-1) in fp32, cast to probs_dtype.

nemo_automodel.components.moe.optimized_ops._rw_bwd_x_core(
    grad_out: torch.Tensor,
    probs: torch.Tensor,
    x_dtype: torch.dtype
) -> torch.Tensor

Gradient w.r.t. x: (grad_out * probs) in fp32, cast to x_dtype.

nemo_automodel.components.moe.optimized_ops._rw_fwd_core(
    x: torch.Tensor,
    probs: torch.Tensor,
    out_dtype: torch.dtype
) -> torch.Tensor

fp32 router-weight multiply: (x * probs) in fp32, cast to out_dtype.

nemo_automodel.components.moe.optimized_ops._RW_CHUNK_ROWS = 8192

nemo_automodel.components.moe.optimized_ops._RW_CHUNK_THRESHOLD = 12288

nemo_automodel.components.moe.optimized_ops._RW_CORES_COMPILED = False

nemo_automodel.components.moe.optimized_ops._rw_bwd_probs_dispatch = _rw_bwd_probs_core

nemo_automodel.components.moe.optimized_ops._rw_bwd_x_dispatch = _rw_bwd_x_core

nemo_automodel.components.moe.optimized_ops._rw_fwd_dispatch = _rw_fwd_core