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`	EAGLE-3 mixed attention under context parallelism.
`_CachedZigZagRingAttention`	Load-balanced variant of :class:`_CachedRingAttention`.

Functions

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

API

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

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

staticmethod

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

staticmethod

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.

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

staticmethod

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

staticmethod

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.

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

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

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

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

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

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

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.

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.

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