nemo_automodel.components.speculative.eagle.zigzag_ring_attention

Load-balanced (zig-zag) ring FlashAttention over a context-parallel group.

The sequence is chunked into 2 * cp pieces; rank r owns chunk r (early) and chunk 2*cp-1-r (late). Pairing an early + late chunk balances the causal triangle so every rank does equal work — unlike a contiguous shard where the last rank does ~2x. Varlen (cu_seqlens) throughout, matching the FlashAttention varlen kernels; causal=True only (zig-zag is meaningless otherwise).

Module Contents

Classes

Name	Description
`RingComm`	P2P ring: send to `rank+1`, receive from `rank-1` (cp group order).
`ZigZagRingFlashAttnVarlenFunc`	Autograd wrapper around the zig-zag ring varlen forward/backward.

Functions

Name	Description
`_block_backward`	-
`_block_forward`	One flash varlen forward block; front/back half uses the halved cu_seqlens.
`get_half_index`	Index selecting the front (early) or back (late) half of each varlen document.
`get_half_lse`	Front/back half of a `[num_heads, seqlen]` lse, per varlen document.
`squeeze_batch`	Drop a leading batch dim of size 1 from each tensor (varlen tensors are unbatched).
`update_out_and_lse`	Online-softmax merge; also returns `sigmoid(block_lse - lse)` for the backward.
`zigzag_ring_flash_attn_varlen_backward`	Backward via the merged-lse trick (zhuzilin): each block’s flash backward is fed
`zigzag_ring_flash_attn_varlen_forward`	Zig-zag ring forward: rotate K/V, per-block flash, online-softmax merge -> (out, lse).
`zigzag_ring_flash_attn_varlen_func`	Load-balanced causal ring FlashAttention. q/k/v are `[1, S_local, H, D]` (varlen packed).

API

class nemo_automodel.components.speculative.eagle.zigzag_ring_attention.RingComm(
    process_group
)

P2P ring: send to rank+1, receive from rank-1 (cp group order).

_ops

list = []

rank

= dist.get_rank(process_group)

world_size

= dist.get_world_size(process_group)

nemo_automodel.components.speculative.eagle.zigzag_ring_attention.RingComm.commit()

nemo_automodel.components.speculative.eagle.zigzag_ring_attention.RingComm.send_recv(
    to_send: torch.Tensor,
    recv_tensor: typing.Optional[torch.Tensor] = None
) -> torch.Tensor

nemo_automodel.components.speculative.eagle.zigzag_ring_attention.RingComm.send_recv_kv(
    k,
    v,
    k_buffer = None,
    v_buffer = None
)

nemo_automodel.components.speculative.eagle.zigzag_ring_attention.RingComm.wait()

class nemo_automodel.components.speculative.eagle.zigzag_ring_attention.ZigZagRingFlashAttnVarlenFunc()

Bases: Function

Autograd wrapper around the zig-zag ring varlen forward/backward.

nemo_automodel.components.speculative.eagle.zigzag_ring_attention.ZigZagRingFlashAttnVarlenFunc.backward(
    ctx,
    dout
)

staticmethod

nemo_automodel.components.speculative.eagle.zigzag_ring_attention.ZigZagRingFlashAttnVarlenFunc.forward(
    ctx,
    q,
    k,
    v,
    cu_seqlens,
    max_seqlen,
    softmax_scale,
    causal,
    window_size,
    group
)

staticmethod

nemo_automodel.components.speculative.eagle.zigzag_ring_attention._block_backward(
    dout,
    q,
    k,
    v,
    out,
    lse,
    causal,
    cu_seqlens,
    max_seqlen,
    block_seq_len,
    dq_buf,
    dk_buf,
    dv_buf,
    scale,
    det,
    window
)

nemo_automodel.components.speculative.eagle.zigzag_ring_attention._block_forward(
    q,
    k,
    v,
    causal,
    cu_seqlens,
    max_seqlen,
    block_seq_len,
    softmax_scale,
    window_size
)

One flash varlen forward block; front/back half uses the halved cu_seqlens.

nemo_automodel.components.speculative.eagle.zigzag_ring_attention.get_half_index(
    cu_seqlens,
    front: bool
)

Index selecting the front (early) or back (late) half of each varlen document.

nemo_automodel.components.speculative.eagle.zigzag_ring_attention.get_half_lse(
    lse,
    cu_seqlens,
    front: bool
)

Front/back half of a [num_heads, seqlen] lse, per varlen document.

nemo_automodel.components.speculative.eagle.zigzag_ring_attention.squeeze_batch(
    t = ()
)

Drop a leading batch dim of size 1 from each tensor (varlen tensors are unbatched).

nemo_automodel.components.speculative.eagle.zigzag_ring_attention.update_out_and_lse(
    out,
    lse,
    block_out,
    block_lse
)

Online-softmax merge; also returns sigmoid(block_lse - lse) for the backward.

nemo_automodel.components.speculative.eagle.zigzag_ring_attention.zigzag_ring_flash_attn_varlen_backward(
    process_group,
    dout,
    q,
    k,
    v,
    out,
    softmax_lse,
    cu_seqlens,
    max_seqlen,
    half_index0,
    half_index1,
    softmax_scale,
    causal = True,
    window_size = (-1, -1),
    deterministic = False
)

Backward via the merged-lse trick (zhuzilin): each block’s flash backward is fed the FINAL merged out/lse, so the joint-softmax gradient (both the output and the normalization/lse paths) is captured without a separate merge backward. ms-swift’s lse_grad variant drops the lse path, which skews dq/dk.

nemo_automodel.components.speculative.eagle.zigzag_ring_attention.zigzag_ring_flash_attn_varlen_forward(
    process_group,
    q,
    k,
    v,
    cu_seqlens,
    max_seqlen,
    half_index0,
    half_index1,
    softmax_scale,
    causal = True,
    window_size = (-1, -1)
)

Zig-zag ring forward: rotate K/V, per-block flash, online-softmax merge -> (out, lse).

nemo_automodel.components.speculative.eagle.zigzag_ring_attention.zigzag_ring_flash_attn_varlen_func(
    q,
    k,
    v,
    cu_seqlens,
    max_seqlen,
    softmax_scale = None,
    causal = True,
    window_size = (-1, -1),
    group = None
)

Load-balanced causal ring FlashAttention. q/k/v are [1, S_local, H, D] (varlen packed).