> 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.speculative.eagle.ring_attention

Differentiable ring FlashAttention over a context-parallel process group.

Each rank holds a contiguous shard `S/cp` of the sequence. K/V are rotated
around the cp ranks p2p (`RingComm`); each incoming K/V block is attended
against the local Q with FlashAttention and merged into the running output via
the online-softmax log-sum-exp identity (`_update_out_and_lse`). The whole
thing is a plain autograd chain (forward rotates K/V, backward rotates the K/V
grads back to their owners), so it composes with FSDP2 like any other module.

The FlashAttention kernels come from the optional `flash_attn` package; import
is guarded so this module never breaks import of the package when it is absent.

## Module Contents

### Classes

| Name                                                                                                                   | Description                                             |
| ---------------------------------------------------------------------------------------------------------------------- | ------------------------------------------------------- |
| [`_CachedRingAttention`](#nemo_automodel-components-speculative-eagle-ring_attention-_CachedRingAttention)             | EAGLE-3 mixed attention under context parallelism.      |
| [`_CachedZigZagRingAttention`](#nemo_automodel-components-speculative-eagle-ring_attention-_CachedZigZagRingAttention) | Load-balanced variant of :class:`_CachedRingAttention`. |

### Functions

| Name                                                                                                                       | Description                                                                             |
| -------------------------------------------------------------------------------------------------------------------------- | --------------------------------------------------------------------------------------- |
| [`_fa_backward`](#nemo_automodel-components-speculative-eagle-ring_attention-_fa_backward)                                 | flash\_attn 2.8.3 bwd: writes `dq/dk/dv` in place.                                      |
| [`_fa_forward`](#nemo_automodel-components-speculative-eagle-ring_attention-_fa_forward)                                   | flash\_attn 2.8.3 fwd: returns `(out[B,S,H,D], lse[B,H,S])`.                            |
| [`_init_out_and_lse`](#nemo_automodel-components-speculative-eagle-ring_attention-_init_out_and_lse)                       | -                                                                                       |
| [`_merge_diag`](#nemo_automodel-components-speculative-eagle-ring_attention-_merge_diag)                                   | Merge a per-position diagonal block `(block_out, block_lse)` into `(out, lse)`.         |
| [`_update_out_and_lse`](#nemo_automodel-components-speculative-eagle-ring_attention-_update_out_and_lse)                   | Online-softmax merge of a new attention block into the running `(out, lse)`.            |
| [`cached_ring_attention`](#nemo_automodel-components-speculative-eagle-ring_attention-cached_ring_attention)               | EAGLE-3 mixed causal-ring + TTT-diagonal attention (see :class:`_CachedRingAttention`). |
| [`cached_zigzag_ring_attention`](#nemo_automodel-components-speculative-eagle-ring_attention-cached_zigzag_ring_attention) | Load-balanced :func:`cached_ring_attention` (zig-zag block-0 ring).                     |
| [`require_flash_attn_version`](#nemo_automodel-components-speculative-eagle-ring_attention-require_flash_attn_version)     | Refuse flash-attn releases the ring was not written against.                            |
| [`ring_flash_attn_backward`](#nemo_automodel-components-speculative-eagle-ring_attention-ring_flash_attn_backward)         | Ring FlashAttention backward. Rotates K/V forward and the K/V grads back to owners.     |
| [`ring_flash_attn_forward`](#nemo_automodel-components-speculative-eagle-ring_attention-ring_flash_attn_forward)           | Ring FlashAttention forward. Q/K/V are `[B, S_local, H, D]`.                            |

### API

```python
class nemo_automodel.components.speculative.eagle.ring_attention._CachedRingAttention()
```

**Bases:** `Function`

EAGLE-3 mixed attention under context parallelism.

`cache_k[:, 0]` / `cache_v[:, 0]` are the step-0 sequence K/V: Q attends
to them **causally over the full (cp-sharded) sequence** via the ring. Every
later block `i >= 1` is a TTT cache step contributing a **per-position
diagonal** `lse_i[t] = (Q_t . K_i_t) * scale` with value `V_i_t` (same
position, no cross-rank comms). Both are fused in one softmax via the
online-softmax log-sum-exp merge. The block-0 backward reuses the *merged*
output/lse so it produces the correct joint-softmax gradient (the standard
FlashAttention backward identity); the diagonal grads are added in closed form.

Layout: `q`/`cache_k`/`cache_v` are FlashAttention-style `[B, T, H, D]`
(`cache_*` carry a block axis: `[B, num_blocks, T, H, D]`).

```python
nemo_automodel.components.speculative.eagle.ring_attention._CachedRingAttention.backward(
    ctx,
    grad_out
)
```

staticmethod

```python
nemo_automodel.components.speculative.eagle.ring_attention._CachedRingAttention.forward(
    ctx,
    q,
    cache_k,
    cache_v,
    process_group,
    scale
)
```

staticmethod

```python
class nemo_automodel.components.speculative.eagle.ring_attention._CachedZigZagRingAttention()
```

**Bases:** `Function`

Load-balanced variant of :class:`_CachedRingAttention`.

Identical joint softmax -- block 0 is the causal sequence attention, blocks
`i >= 1` are per-position TTT diagonals -- but block 0 runs the zig-zag ring
so every cp rank does equal causal work (a contiguous shard leaves the last
rank doing \~2x). This requires the sequence to be sharded in zig-zag order
(rank `r` owns chunks `r` and `2*cp-1-r`); the diagonals are layout
agnostic (`q` and `cache_k[i]` share the shard) so they merge unchanged.

Layout matches :class:`_CachedRingAttention` (`q` is `[1, T, H, D]`,
`cache_*` are `[1, num_blocks, T, H, D]`); the zig-zag ring is varlen and
unbatched, so `B == 1`.

```python
nemo_automodel.components.speculative.eagle.ring_attention._CachedZigZagRingAttention.backward(
    ctx,
    grad_out
)
```

staticmethod

```python
nemo_automodel.components.speculative.eagle.ring_attention._CachedZigZagRingAttention.forward(
    ctx,
    q,
    cache_k,
    cache_v,
    process_group,
    scale
)
```

staticmethod

```python
nemo_automodel.components.speculative.eagle.ring_attention._fa_backward(
    dout,
    q,
    k,
    v,
    out,
    lse,
    dq,
    dk,
    dv,
    softmax_scale,
    causal,
    deterministic = False
)
```

flash\_attn 2.8.3 bwd: writes `dq/dk/dv` in place.

```python
nemo_automodel.components.speculative.eagle.ring_attention._fa_forward(
    q,
    k,
    v,
    softmax_scale,
    causal,
    dropout_p = 0.0
)
```

flash\_attn 2.8.3 fwd: returns `(out[B,S,H,D], lse[B,H,S])`.

```python
nemo_automodel.components.speculative.eagle.ring_attention._init_out_and_lse(
    block_out: torch.Tensor,
    block_lse: torch.Tensor
) -> tuple[torch.Tensor, torch.Tensor]
```

```python
nemo_automodel.components.speculative.eagle.ring_attention._merge_diag(
    out,
    lse,
    block_out,
    block_lse
)
```

Merge a per-position diagonal block `(block_out, block_lse)` into `(out, lse)`.

`out`/`block_out` are `[B, T, H, D]` fp32; `lse`/`block_lse` are `[B, T, H]`.

```python
nemo_automodel.components.speculative.eagle.ring_attention._update_out_and_lse(
    out: torch.Tensor,
    lse: torch.Tensor,
    block_out: torch.Tensor,
    block_lse: torch.Tensor
) -> tuple[torch.Tensor, torch.Tensor]
```

Online-softmax merge of a new attention block into the running `(out, lse)`.

`out` is `[B, S, H, D]` fp32, `lse` is `[B, S, H, 1]` fp32; `block_lse`
arrives as the kernel's `[B, H, S]` and is transposed in. Numerically-stable
sigmoid form (see zhuzilin/ring-flash-attention#34).

```python
nemo_automodel.components.speculative.eagle.ring_attention.cached_ring_attention(
    q: torch.Tensor,
    cache_k: list[torch.Tensor],
    cache_v: list[torch.Tensor],
    process_group,
    scale: float
) -> torch.Tensor
```

EAGLE-3 mixed causal-ring + TTT-diagonal attention (see :class:`_CachedRingAttention`).

`q` is `[B, T, H, D]`; `cache_k`/`cache_v` are lists of `[B, T, H, D]`
(index 0 = sequence, 1.. = TTT cache steps). Returns `[B, T, H, D]`.

```python
nemo_automodel.components.speculative.eagle.ring_attention.cached_zigzag_ring_attention(
    q: torch.Tensor,
    cache_k: list[torch.Tensor],
    cache_v: list[torch.Tensor],
    process_group,
    scale: float
) -> torch.Tensor
```

Load-balanced :func:`cached_ring_attention` (zig-zag block-0 ring).

Inputs must be in zig-zag-sharded layout (see :class:`_CachedZigZagRingAttention`).
`q` is `[1, T, H, D]`; `cache_k`/`cache_v` are lists of `[1, T, H, D]`.

```python
nemo_automodel.components.speculative.eagle.ring_attention.require_flash_attn_version() -> None
```

Refuse flash-attn releases the ring was not written against.

`_fa_forward` / `_fa_backward` call the private `_flash_attn_forward` /
`_flash_attn_backward` by POSITIONAL args pinned to the 2.8.x layout
(`q, k, v, dropout, scale, causal, win_left, win_right, softcap, alibi,
return_softmax`). Other 2.x releases reorder or insert params, which would
silently bind `causal` / `softmax_scale` to the wrong slots and train on
garbage, so gate the version rather than the mere presence of the package.

```python
nemo_automodel.components.speculative.eagle.ring_attention.ring_flash_attn_backward(
    process_group,
    dout: torch.Tensor,
    q: torch.Tensor,
    k: torch.Tensor,
    v: torch.Tensor,
    out: torch.Tensor,
    softmax_lse: torch.Tensor,
    softmax_scale: float,
    causal: bool = True
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]
```

Ring FlashAttention backward. Rotates K/V forward and the K/V grads back to owners.

```python
nemo_automodel.components.speculative.eagle.ring_attention.ring_flash_attn_forward(
    process_group,
    q: torch.Tensor,
    k: torch.Tensor,
    v: torch.Tensor,
    softmax_scale: float,
    causal: bool = True
) -> tuple[torch.Tensor, torch.Tensor]
```

Ring FlashAttention forward. Q/K/V are `[B, S_local, H, D]`.

Returns `(out[B, S_local, H, D], lse[B, H, S_local])`. For `causal=True` a
rank only attends to K/V blocks from itself and earlier ranks (`step <= rank`),
and applies the causal mask only to its own block (`step == 0`).