> 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.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`](#nemo_automodel-components-speculative-eagle-zigzag_ring_attention-RingComm)                                           | P2P ring: send to `rank+1`, receive from `rank-1` (cp group order). |
| [`ZigZagRingFlashAttnVarlenFunc`](#nemo_automodel-components-speculative-eagle-zigzag_ring_attention-ZigZagRingFlashAttnVarlenFunc) | Autograd wrapper around the zig-zag ring varlen forward/backward.   |

### Functions

| Name                                                                                                                                                  | Description                                                                               |
| ----------------------------------------------------------------------------------------------------------------------------------------------------- | ----------------------------------------------------------------------------------------- |
| [`_block_backward`](#nemo_automodel-components-speculative-eagle-zigzag_ring_attention-_block_backward)                                               | -                                                                                         |
| [`_block_forward`](#nemo_automodel-components-speculative-eagle-zigzag_ring_attention-_block_forward)                                                 | One flash varlen forward block; front/back half uses the halved cu\_seqlens.              |
| [`get_half_index`](#nemo_automodel-components-speculative-eagle-zigzag_ring_attention-get_half_index)                                                 | Index selecting the front (early) or back (late) half of each varlen document.            |
| [`get_half_lse`](#nemo_automodel-components-speculative-eagle-zigzag_ring_attention-get_half_lse)                                                     | Front/back half of a `[num_heads, seqlen]` lse, per varlen document.                      |
| [`squeeze_batch`](#nemo_automodel-components-speculative-eagle-zigzag_ring_attention-squeeze_batch)                                                   | Drop a leading batch dim of size 1 from each tensor (varlen tensors are unbatched).       |
| [`update_out_and_lse`](#nemo_automodel-components-speculative-eagle-zigzag_ring_attention-update_out_and_lse)                                         | Online-softmax merge; also returns `sigmoid(block_lse - lse)` for the backward.           |
| [`zigzag_ring_flash_attn_varlen_backward`](#nemo_automodel-components-speculative-eagle-zigzag_ring_attention-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`](#nemo_automodel-components-speculative-eagle-zigzag_ring_attention-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`](#nemo_automodel-components-speculative-eagle-zigzag_ring_attention-zigzag_ring_flash_attn_varlen_func)         | Load-balanced causal ring FlashAttention. q/k/v are `[1, S_local, H, D]` (varlen packed). |

### API

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

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

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

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

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

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

**Bases:** `Function`

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

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

staticmethod

```python
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

```python
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
)
```

```python
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.

```python
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.

```python
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.

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

```python
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.

```python
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.

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

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