> 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.shared.tp_linear

Low-level linear graph-shaping helpers shared by tensor-parallel components.

Inductor's async tensor-parallel pass (`torch._inductor.config._micro_pipeline_tp`)
fuses collectives with matmuls by pattern-matching the reshape-mm-reshape graph
that `F.linear` produces for 3-D input.  The compile-safe `torch.bmm` path
used by `TPLinear`/`LinearLoRA` never matches, so async-TP fusion silently
fails to fire.  These helpers detect async-TP tracing and emit the native
linear graph in that mode only.

## Module Contents

### Functions

| Name                                                                                          | Description                                                               |
| --------------------------------------------------------------------------------------------- | ------------------------------------------------------------------------- |
| [`_async_tp_linear`](#nemo_automodel-shared-tp_linear-_async_tp_linear)                       | Emit the native linear graph recognized by async-TP fusion.               |
| [`_is_async_tp_linear_enabled`](#nemo_automodel-shared-tp_linear-_is_async_tp_linear_enabled) | Return whether Dynamo is tracing with Inductor's async-TP pass enabled.   |
| [`tp_linear_forward`](#nemo_automodel-shared-tp_linear-tp_linear_forward)                     | Dispatch a TP-safe linear between F.linear, async-TP shaping, and bmm/mm. |

### API

```python
nemo_automodel.shared.tp_linear._async_tp_linear(
    x: torch.Tensor,
    weight: torch.Tensor,
    bias: torch.Tensor | None = None
) -> torch.Tensor
```

Emit the native linear graph recognized by async-TP fusion.

**Parameters:**

Input activations of shape `[..., in_features]`; any number of
leading dimensions is accepted (typically `[B, S, in_features]`
with `B` = batch, `S` = sequence).  May be a DTensor with
tensor-parallel placements.

Weight of shape `[out_features, in_features]`; may be a
DTensor sharded for colwise (`Shard(0)`) or rowwise
(`Shard(1)`) tensor parallelism.

Optional bias of shape `[out_features]`.

**Returns:** `torch.Tensor`

Output of shape `[..., out_features]` with the same leading

```python
nemo_automodel.shared.tp_linear._is_async_tp_linear_enabled() -> bool
```

Return whether Dynamo is tracing with Inductor's async-TP pass enabled.

True only when both conditions hold: the caller is being traced by
torch.compile and `torch._inductor.config._micro_pipeline_tp` is set
(see `enable_async_tensor_parallel` in parallelizer.py).  Always False
in eager mode.

```python
nemo_automodel.shared.tp_linear.tp_linear_forward(
    x: torch.Tensor,
    weight: torch.Tensor,
    bias: torch.Tensor | None,
    mm_for_2d_compile: bool
) -> torch.Tensor
```

Dispatch a TP-safe linear between F.linear, async-TP shaping, and bmm/mm.

Shared by `TPLinear.forward` and `LinearLoRA.forward`.  Eager
unsharded/last-dimension-sharded inputs use `F.linear`; async-TP tracing emits the
fusable native linear graph via `_async_tp_linear`; every remaining
compile path and dim-0/1-sharded DTensor input keeps the DTensor-safe
matmul fallback (`aten.view` cannot flatten a sharded dimension).

**Parameters:**

Input activations of shape `[B, S, in_features]` (`B` = batch,
`S` = sequence) or `[N, in_features]` (`N` = flattened
tokens).  May be a DTensor: a 3-D DTensor sharded on dim 0 or 1
(e.g. `Shard(1)` from sequence parallelism) always takes the
`torch.bmm` path — the async-TP fast path assumes the input is
replicated or sharded only on the feature dim, so fusion simply
does not fire for such layers.

Weight of shape `[out_features, in_features]`; may be a
DTensor sharded for colwise (`Shard(0)`) or rowwise (`Shard(1)`)
tensor parallelism.

Optional bias of shape `[out_features]`; replicated if DTensor.

Numerics for 2-D `x` traced under torch.compile
without async-TP: `torch.mm` plus explicit bias add when True
(`TPLinear`), `F.linear` when False (`LinearLoRA`).

**Returns:** `torch.Tensor`

Output of shape `[..., out_features]` with the same leading