> This page is for version Nightly (default).
> For other versions, use one of these documentation indexes:
> - Nightly (default): https://docs.nvidia.com/nemo/automodel/nightly/llms.txt
> - Latest: https://docs.nvidia.com/nemo/automodel/latest/llms.txt
> - 0.7.0 · 26.10: https://docs.nvidia.com/nemo/automodel/v0.7/llms.txt
> - 0.5.0 · 26.06: https://docs.nvidia.com/nemo/automodel/v0.5/llms.txt
> - 0.4.0 · 26.04: https://docs.nvidia.com/nemo/automodel/v0.4/llms.txt

> For clean Markdown of any page, append .md to the page URL.
> For a complete documentation index, see https://docs.nvidia.com/nemo/automodel/llms.txt.
> For AI client integration (Claude Code, Cursor, etc.), connect to the MCP server at https://docs.nvidia.com/nemo/automodel/_mcp/server.

# nemo_automodel.components.optim.muown_triton

Fused local Muown row operations; Newton-Schulz and collectives stay in Muown.

Only contiguous FP32 local tensors with complete input-feature axes use these
kernels. Matrix batches and output-neuron shards are independent. Public DTensor
ownership and autograd version counters remain the optimizer's responsibility.

## Module Contents

### Functions

| Name                                                                                             | Description                                                           |
| ------------------------------------------------------------------------------------------------ | --------------------------------------------------------------------- |
| [`_direction_kernel`](#nemo_automodel-components-optim-muown_triton-_direction_kernel)           | -                                                                     |
| [`_finish`](#nemo_automodel-components-optim-muown_triton-_finish)                               | Apply direction/magnitude updates, normalize and decay local weights. |
| [`_finish_kernel`](#nemo_automodel-components-optim-muown_triton-_finish_kernel)                 | -                                                                     |
| [`_magnitude_grad_kernel`](#nemo_automodel-components-optim-muown_triton-_magnitude_grad_kernel) | -                                                                     |
| [`_prepare`](#nemo_automodel-components-optim-muown_triton-_prepare)                             | Prepare complete local rows and mutate momentum.                      |
| [`_prepare_kernel`](#nemo_automodel-components-optim-muown_triton-_prepare_kernel)               | -                                                                     |
| [`_supported`](#nemo_automodel-components-optim-muown_triton-_supported)                         | Check local dtype/strides and the bounded row-reduction footprint.    |

### Data

[`HAVE_TRITON`](#nemo_automodel-components-optim-muown_triton-HAVE_TRITON)

### API

```python
nemo_automodel.components.optim.muown_triton._direction_kernel(
    W,
    G,
    N,
    Grad,
    Mom,
    D,
    U,
    GG,
    SIZE: tl.constexpr,
    R: tl.constexpr,
    INPUTS: tl.constexpr,
    MU: tl.constexpr,
    EPS: tl.constexpr,
    NESTEROV: tl.constexpr,
    BLOCK: tl.constexpr
)
```

```python
nemo_automodel.components.optim.muown_triton._finish(
    weight: torch.Tensor,
    magnitude: torch.Tensor,
    norm: torch.Tensor,
    first_moment: torch.Tensor,
    second_moment: torch.Tensor,
    direction: torch.Tensor,
    update: torch.Tensor,
    magnitude_grad: torch.Tensor,
    transposed: bool,
    lr: float,
    scale: float,
    beta1: float,
    beta2: float,
    step: int,
    epsilon: float,
    weight_decay: float
) -> None
```

Apply direction/magnitude updates, normalize and decay local weights.

Tensor inputs follow \_prepare's contiguous local-row contract. Update is
BF16 with weight's shape; direction is FP32 with that shape. All other
tensors are FP32 with magnitude's row shape. Weight, magnitude, norm and
both moments are mutated in place; direction/update/gradients are preserved.
Scale must be computed from the GLOBAL matrix shape, including on FSDP
output shards. LR and bias corrections are runtime scalars.

```python
nemo_automodel.components.optim.muown_triton._finish_kernel(
    W,
    G,
    N,
    M,
    V,
    D,
    U,
    GG,
    R: tl.constexpr,
    INPUTS: tl.constexpr,
    TRANSPOSED: tl.constexpr,
    LR,
    DIRECTION_LR,
    DECAY,
    B1: tl.constexpr,
    B2: tl.constexpr,
    BC1,
    BC2,
    EPS: tl.constexpr,
    WD: tl.constexpr,
    BR: tl.constexpr,
    BI: tl.constexpr
)
```

```python
nemo_automodel.components.optim.muown_triton._magnitude_grad_kernel(
    W,
    G,
    Grad,
    GG,
    R: tl.constexpr,
    INPUTS: tl.constexpr,
    EPS: tl.constexpr,
    BR: tl.constexpr,
    BI: tl.constexpr
)
```

```python
nemo_automodel.components.optim.muown_triton._prepare(
    weight: torch.Tensor,
    gradient: torch.Tensor,
    momentum: torch.Tensor,
    magnitude: torch.Tensor,
    norm: torch.Tensor,
    transposed: bool,
    mu: float,
    epsilon: float,
    nesterov: bool
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]
```

Prepare complete local rows and mutate momentum.

**Parameters:**

**`weight`** `Tensor`

Contiguous FP32 CUDA \[..., output, input], or \[..., input,
output] when transposed. Leading dimensions batch independent
matrices; the input axis must be complete on this rank.

---

**`gradient`** `Tensor`

Same shape/dtype/device/layout as weight, preserved.

---

**`momentum`** `Tensor`

Same shape/dtype/device/layout as weight, updated in place.

---

**`magnitude`** `Tensor`

Contiguous FP32 tensor with the input axis reduced to 1.

---

**`norm`** `Tensor`

Same shape/layout as magnitude; direction norms.

---

**`transposed`** `bool`

Whether output neurons occupy the final physical axis.

---

**`mu`** `float`

Momentum coefficient.

---

**`epsilon`** `float`

Signed magnitude division guard.

---

**`nesterov`** `bool`

Whether to use Nesterov momentum.

---

**Returns:** `Tensor`

Independent direction (FP32), update (BF16), and magnitude-gradient

```python
nemo_automodel.components.optim.muown_triton._prepare_kernel(
    W,
    G,
    N,
    Grad,
    Mom,
    D,
    U,
    GG,
    R: tl.constexpr,
    INPUTS: tl.constexpr,
    TRANSPOSED: tl.constexpr,
    MU: tl.constexpr,
    EPS: tl.constexpr,
    NESTEROV: tl.constexpr,
    BR: tl.constexpr,
    BI: tl.constexpr
)
```

```python
nemo_automodel.components.optim.muown_triton._supported(
    tensors: tuple[torch.Tensor, ...],
    transposed: bool
) -> bool
```

Check local dtype/strides and the bounded row-reduction footprint.

```python
nemo_automodel.components.optim.muown_triton.HAVE_TRITON = _HAS_TRITON and _HAS_LANGUAGE
```