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`	EAGLE-3 mixed causal-ring-free attention under Ulysses context parallelism.

Functions

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

API

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

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

staticmethod

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

staticmethod

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

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

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

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.

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:

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

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

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:

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

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:

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 | NoneDefaults to None

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

max_seqlen

int | NoneDefaults to None

Longest document length for the packed path.

Returns: torch.Tensor

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