> 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.ulysses_attention

Ulysses (all-to-all) context-parallel attention for the EAGLE-3 draft.

Context parallelism shards the sequence: each rank holds a contiguous `S/cp`
slice of the tokens. Ulysses turns that sequence shard into a *head* shard for
block-0 attention with an all-to-all, so a rank sees the **full** sequence for a
subset of the heads:

`[B, S_local, H, D]`  --all-to-all-->  `[B, S_full, H/cp, D]`

With the full sequence local, block-0 (`Q @ K_0^T`, causal over the whole
sequence) runs a single dense or packed `varlen` FlashAttention, so
`cu_seqlens` document boundaries carry through and CP composes with sequence
packing. The output is all-to-all'd back to the sequence shard, and the
per-position EAGLE-3 TTT diagonals (blocks `i >= 1`) merge into it via the
online-softmax identity, shard-local and comm-free.

This is a hand-written `autograd.Function` (like the ring): the backward runs
the FlashAttention backward against the **merged** joint-softmax out/lse (not the
block-0-only softmax), which is required for the block-0 `q`/`k` gradients to
be correct whenever any diagonal step is present. The all-to-all is done with the
raw collective in forward and backward (grads are threaded by hand).

## Module Contents

### Classes

| Name                                                                                                                | Description                                                                 |
| ------------------------------------------------------------------------------------------------------------------- | --------------------------------------------------------------------------- |
| [`_CachedUlyssesAttention`](#nemo_automodel-components-speculative-eagle-ulysses_attention-_CachedUlyssesAttention) | EAGLE-3 mixed causal-ring-free attention under Ulysses context parallelism. |

### Functions

| Name                                                                                                                    | Description                                                                               |
| ----------------------------------------------------------------------------------------------------------------------- | ----------------------------------------------------------------------------------------- |
| [`_all_to_all_single`](#nemo_automodel-components-speculative-eagle-ulysses_attention-_all_to_all_single)               | Blocking `all_to_all_single` splitting/joining along dim 0 (equal splits).                |
| [`_block0_backward`](#nemo_automodel-components-speculative-eagle-ulysses_attention-_block0_backward)                   | FlashAttention backward over the gathered sequence, using the MERGED out/lse.             |
| [`_block0_forward`](#nemo_automodel-components-speculative-eagle-ulysses_attention-_block0_forward)                     | Causal FlashAttention over the gathered sequence.                                         |
| [`_gather_lse`](#nemo_automodel-components-speculative-eagle-ulysses_attention-_gather_lse)                             | Sequence-sharded -> head-sharded log-sum-exp: `[B, H, S_local] -> [B, H/uly, S_full]`. |
| [`_gather_seq_scatter_heads`](#nemo_automodel-components-speculative-eagle-ulysses_attention-_gather_seq_scatter_heads) | Sequence-sharded -> head-sharded: `[B, S_local, H, D] -> [B, S_full, H/uly, D]`.       |
| [`_scatter_lse`](#nemo_automodel-components-speculative-eagle-ulysses_attention-_scatter_lse)                           | Head-sharded -> sequence-sharded log-sum-exp: `[B, H/uly, S_full] -> [B, H, S_local]`. |
| [`_scatter_seq_gather_heads`](#nemo_automodel-components-speculative-eagle-ulysses_attention-_scatter_seq_gather_heads) | Head-sharded -> sequence-sharded: `[B, S_full, H/uly, D] -> [B, S_local, H, D]`.       |
| [`cached_ulysses_attention`](#nemo_automodel-components-speculative-eagle-ulysses_attention-cached_ulysses_attention)   | EAGLE-3 mixed causal-block-0 + TTT-diagonal attention under Ulysses CP.                   |

### API

```python
class nemo_automodel.components.speculative.eagle.ulysses_attention._CachedUlyssesAttention()
```

**Bases:** `Function`

EAGLE-3 mixed causal-ring-free attention under Ulysses context parallelism.

Block 0 (`cache_k[:, 0]` / `cache_v[:, 0]`) is the causal sequence
attention run over the all-to-all-gathered full sequence; blocks `i >= 1` are
per-position TTT diagonals (same position, shard-local). Both are fused into one
softmax via the online-softmax merge. The backward re-runs the block-0
FlashAttention backward with the **merged** output/lse so it produces the
correct joint-softmax gradient (the same identity `_CachedRingAttention` uses);
the diagonal grads are added in closed form.

Layout: `q` is `[B, T_local, H, D]` (FlashAttention layout); `cache_k` /
`cache_v` carry a block axis `[B, num_blocks, T_local, H, D]`.

```python
nemo_automodel.components.speculative.eagle.ulysses_attention._CachedUlyssesAttention.backward(
    ctx,
    grad_out
)
```

staticmethod

```python
nemo_automodel.components.speculative.eagle.ulysses_attention._CachedUlyssesAttention.forward(
    ctx,
    q,
    cache_k,
    cache_v,
    group,
    scale,
    cu_seqlens,
    max_seqlen
)
```

staticmethod

```python
nemo_automodel.components.speculative.eagle.ulysses_attention._all_to_all_single(
    x: torch.Tensor,
    group: torch.distributed.ProcessGroup
) -> torch.Tensor
```

Blocking `all_to_all_single` splitting/joining along dim 0 (equal splits).

```python
nemo_automodel.components.speculative.eagle.ulysses_attention._block0_backward(
    dout_g: torch.Tensor,
    q_g: torch.Tensor,
    k_g: torch.Tensor,
    v_g: torch.Tensor,
    out_g: torch.Tensor,
    lse_g: torch.Tensor,
    scale: float,
    cu_seqlens: torch.Tensor | None,
    max_seqlen: int | None
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]
```

FlashAttention backward over the gathered sequence, using the MERGED out/lse.

**Parameters:**

**`dout_g`** `torch.Tensor`

Upstream grad, `[batch, seq_full, heads_local, head_dim]`.

---

**`q_g, k_g, v_g`**

Gathered q/k0/v0, same layout.

---

**`out_g`** `torch.Tensor`

The MERGED joint-softmax output gathered back, same layout / dtype as q.

---

**`lse_g`** `torch.Tensor`

The MERGED joint-softmax log-sum-exp, `[batch, heads_local, seq_full]`.

---

**`scale`** `float`

Softmax scale.

---

**`cu_seqlens, max_seqlen`**

Packed-document boundaries / longest doc, or `None`.

---

**Returns:** `tuple[torch.Tensor, torch.Tensor, torch.Tensor]`

`(dq, dk0, dv0)`, each `[batch, seq_full, heads_local, head_dim]`.

```python
nemo_automodel.components.speculative.eagle.ulysses_attention._block0_forward(
    q_g: torch.Tensor,
    k_g: torch.Tensor,
    v_g: torch.Tensor,
    scale: float,
    cu_seqlens: torch.Tensor | None,
    max_seqlen: int | None
) -> tuple[torch.Tensor, torch.Tensor]
```

Causal FlashAttention over the gathered sequence.

**Parameters:**

**`q_g, k_g, v_g`**

Gathered `[batch, seq_full, heads_local, head_dim]`.

---

**`scale`** `float`

Softmax scale.

---

**`cu_seqlens`** `torch.Tensor | None`

GLOBAL packed-document boundaries `[num_docs + 1]` (int32) for
the varlen path, or `None` for dense causal over the whole sequence.

---

**`max_seqlen`** `int | None`

Longest document length for the varlen path.

---

**Returns:** `torch.Tensor`

`(out, lse)` with `out` `[batch, seq_full, heads_local, head_dim]` and

```python
nemo_automodel.components.speculative.eagle.ulysses_attention._gather_lse(
    lse: torch.Tensor,
    group: torch.distributed.ProcessGroup
) -> torch.Tensor
```

Sequence-sharded -> head-sharded log-sum-exp: `[B, H, S_local] -> [B, H/uly, S_full]`.

The result is contiguous: the FlashAttention backward passes `softmax_lse`
straight to the kernel without a `maybe_contiguous` (unlike q/k/v/out/dout),
so a transposed view would be read with the wrong strides.

```python
nemo_automodel.components.speculative.eagle.ulysses_attention._gather_seq_scatter_heads(
    x: torch.Tensor,
    group: torch.distributed.ProcessGroup
) -> torch.Tensor
```

Sequence-sharded -> head-sharded: `[B, S_local, H, D] -> [B, S_full, H/uly, D]`.

**Parameters:**

**`x`** `torch.Tensor`

FlashAttention-layout tensor `[batch, seq_local, heads, head_dim]` --
this rank's contiguous sequence shard with all heads; `heads` must be
divisible by the Ulysses degree `uly` (the `group` world size).

---

**Returns:** `torch.Tensor`

Tensor `[batch, seq_full, heads // uly, head_dim]` with the full sequence

```python
nemo_automodel.components.speculative.eagle.ulysses_attention._scatter_lse(
    lse: torch.Tensor,
    group: torch.distributed.ProcessGroup
) -> torch.Tensor
```

Head-sharded -> sequence-sharded log-sum-exp: `[B, H/uly, S_full] -> [B, H, S_local]`.

```python
nemo_automodel.components.speculative.eagle.ulysses_attention._scatter_seq_gather_heads(
    x: torch.Tensor,
    group: torch.distributed.ProcessGroup
) -> torch.Tensor
```

Head-sharded -> sequence-sharded: `[B, S_full, H/uly, D] -> [B, S_local, H, D]`.

Inverse of :func:`_gather_seq_scatter_heads`.

**Parameters:**

**`x`** `torch.Tensor`

Tensor `[batch, seq_full, heads_local, head_dim]` -- the full sequence
with this rank's slice of the heads; `seq_full` must be divisible by
the Ulysses degree `uly`.

---

**Returns:** `torch.Tensor`

Tensor `[batch, seq_full // uly, heads_local * uly, head_dim]` -- this

```python
nemo_automodel.components.speculative.eagle.ulysses_attention.cached_ulysses_attention(
    q: torch.Tensor,
    cache_k: list[torch.Tensor],
    cache_v: list[torch.Tensor],
    group: torch.distributed.ProcessGroup,
    scale: float,
    cu_seqlens: torch.Tensor | None = None,
    max_seqlen: int | None = None
) -> torch.Tensor
```

EAGLE-3 mixed causal-block-0 + TTT-diagonal attention under Ulysses CP.

**Parameters:**

**`q`** `torch.Tensor`

This step's query, `[batch, seq_local, heads, head_dim]` (FlashAttention
layout; the sequence is CP-sharded).

---

**`cache_k`** `list[torch.Tensor]`

Per-TTT-step keys, each `[batch, seq_local, heads, head_dim]`;
index 0 is the step-0 sequence key, `i >= 1` the diagonal steps.

---

**`cache_v`** `list[torch.Tensor]`

Per-TTT-step values, same layout as `cache_k`.

---

**`group`** `dist.ProcessGroup`

The Ulysses (context-parallel) process group.

---

**`scale`** `float`

Softmax scale (`head_dim ** -0.5`).

---

**`cu_seqlens`** `torch.Tensor | None` — default: None

GLOBAL (un-sharded) packed-document boundaries `[num_docs + 1]`
(int32), or `None` for a single causal stream.

---

**`max_seqlen`** `int | None` — default: None

Longest document length for the packed path.

---

**Returns:** `torch.Tensor`

Attention output `[batch, seq_local, heads, head_dim]` in `q`'s dtype.