> 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.models.gemma4_moe.cp_local_ring

Local FlashAttention-2 compute for Gemma4 SLIDING-window CP layers.

Compute only: it takes the KV chunks the ring already collected
(:func:`cp_attention._collect_ring_kv_chunks`) and runs one FlashAttention forward+backward
over the causal window neighborhood. Ring communication and gradient routing stay in
`cp_attention`'s :class:`_Gemma4LocalKernelRingAttention`, so this reuses the flex/FFPA
rings' exact comm -- no new collectives.

The causally-valid chunks (owners `<= cp_rank`) concat into one neighborhood; the local
query shard is its suffix, so a single `flash_attn` call with bottom-right causal + left
window places the shard correctly (no online-softmax merge). Packed multi-doc uses THD
`cu_seqlens` from :func:`_build_packed_ring_segments` (cached per step). Backward is
no-recompute (see `_fa_ll_*` below).

Sliding layers only: global (all-chunk) layers keep FFPA; vision-bidirectional masks keep flex.

## Module Contents

### Functions

| Name                                                                                                                    | Description                                                                               |
| ----------------------------------------------------------------------------------------------------------------------- | ----------------------------------------------------------------------------------------- |
| [`_cached_sliding_segments`](#nemo_automodel-components-models-gemma4_moe-cp_local_ring-_cached_sliding_segments)       | Per-step cache around the neighborhood segment build.                                     |
| [`_concat_neighborhood`](#nemo_automodel-components-models-gemma4_moe-cp_local_ring-_concat_neighborhood)               | Concatenate the causally-valid collected chunks (owners `<= cp_rank`) ascending.       |
| [`_fa_ll_dense_bwd`](#nemo_automodel-components-models-gemma4_moe-cp_local_ring-_fa_ll_dense_bwd)                       | -                                                                                         |
| [`_fa_ll_dense_fwd`](#nemo_automodel-components-models-gemma4_moe-cp_local_ring-_fa_ll_dense_fwd)                       | -                                                                                         |
| [`_fa_ll_varlen_bwd`](#nemo_automodel-components-models-gemma4_moe-cp_local_ring-_fa_ll_varlen_bwd)                     | -                                                                                         |
| [`_fa_ll_varlen_fwd`](#nemo_automodel-components-models-gemma4_moe-cp_local_ring-_fa_ll_varlen_fwd)                     | -                                                                                         |
| [`_fold_pad`](#nemo_automodel-components-models-gemma4_moe-cp_local_ring-_fold_pad)                                     | Fold padding into the doc map (pad -> id 0) so `_build_packed_ring_segments`              |
| [`_owner_grad_dict`](#nemo_automodel-components-models-gemma4_moe-cp_local_ring-_owner_grad_dict)                       | Full `owner -> dK/dV` map the ring routing expects: every owner `(cp_rank-d)%cp_size`  |
| [`_scale_of`](#nemo_automodel-components-models-gemma4_moe-cp_local_ring-_scale_of)                                     | -                                                                                         |
| [`sliding_ring_compute_fa_bwd`](#nemo_automodel-components-models-gemma4_moe-cp_local_ring-sliding_ring_compute_fa_bwd) | Backward of :func:`sliding_ring_compute_fa_fwd` using FlashAttention's own backward (no   |
| [`sliding_ring_compute_fa_fwd`](#nemo_automodel-components-models-gemma4_moe-cp_local_ring-sliding_ring_compute_fa_fwd) | No-recompute FlashAttention sliding-window forward over the collected ring chunks, saving |

### API

```python
nemo_automodel.components.models.gemma4_moe.cp_local_ring._cached_sliding_segments(
    ctx: typing.Any,
    folded_q_ids: torch.Tensor,
    folded_nb_ids: torch.Tensor,
    lo: int
)
```

Per-step cache around the neighborhood segment build.

The segment depends only on the doc maps + ring geometry, IDENTICAL across every sliding
layer in a step (same `_packed_seq_ids`, same `sliding_window` => same `lo`) and
forward/backward -- so the `.tolist()` D->H sync + Python pairing in
:func:`_build_packed_ring_segments` runs ONCE per step, not once per sliding layer per
pass (incl. activation-checkpoint recompute). Anchored on the *persistent*
`_packed_seq_ids` tensor, shared with the FFPA ring's cache. `None` is cached too.

```python
nemo_automodel.components.models.gemma4_moe.cp_local_ring._concat_neighborhood(
    ctx: typing.Any,
    chunks: list
)
```

Concatenate the causally-valid collected chunks (owners `<= cp_rank`) ascending.

`chunks` is :func:`cp_attention._collect_ring_kv_chunks` output
`(owner, key, value, metadata)`. Returns `(k_nb, v_nb, nb_ids, valid_owners)` -- the
neighborhood K/V (BHSD), the concatenated padding-folded `_packed_seq_ids` (or `None`
if unpacked), and the ascending list of kept owner ranks (matching the chunk order in
`k_nb`). Owners `> cp_rank` (future / wrapped) are causally dropped here; their grads
are zero-routed by the ring Function.

```python
nemo_automodel.components.models.gemma4_moe.cp_local_ring._fa_ll_dense_bwd(
    grad_out,
    st,
    scale,
    win
)
```

```python
nemo_automodel.components.models.gemma4_moe.cp_local_ring._fa_ll_dense_fwd(
    q,
    k,
    v,
    scale,
    win
)
```

```python
nemo_automodel.components.models.gemma4_moe.cp_local_ring._fa_ll_varlen_bwd(
    dout_p,
    st,
    seg,
    scale,
    win
)
```

```python
nemo_automodel.components.models.gemma4_moe.cp_local_ring._fa_ll_varlen_fwd(
    qp,
    kp,
    vp,
    seg,
    scale,
    win
)
```

```python
nemo_automodel.components.models.gemma4_moe.cp_local_ring._fold_pad(
    ids: torch.Tensor,
    pm: torch.Tensor | None
) -> torch.Tensor
```

Fold padding into the doc map (pad -> id 0) so `_build_packed_ring_segments`
excludes it; replaces the flex path's separate padding\_mask handling. Returns `ids`
unchanged (same object, cache-stable) when there is no padding\_mask.

```python
nemo_automodel.components.models.gemma4_moe.cp_local_ring._owner_grad_dict(
    per_owner: dict,
    cp_rank,
    cp_size,
    n_prior,
    zeros_like
)
```

Full `owner -> dK/dV` map the ring routing expects: every owner `(cp_rank-d)%cp_size`
for `d in 0..n_prior`, with zeros for owners this shard didn't causally attend.

```python
nemo_automodel.components.models.gemma4_moe.cp_local_ring._scale_of(
    ctx: typing.Any
) -> float
```

```python
nemo_automodel.components.models.gemma4_moe.cp_local_ring.sliding_ring_compute_fa_bwd(
    saved: dict,
    grad_out: torch.Tensor
)
```

Backward of :func:`sliding_ring_compute_fa_fwd` using FlashAttention's own backward (no
recompute). Returns `(grad_q[B,Hq,sl,D], grad_key_by_owner, grad_value_by_owner)` --
per-owner dK/dV maps ready for :func:`cp_attention._route_kv_grads_to_owners`.

```python
nemo_automodel.components.models.gemma4_moe.cp_local_ring.sliding_ring_compute_fa_fwd(
    ctx: typing.Any,
    chunks: list
)
```

No-recompute FlashAttention sliding-window forward over the collected ring chunks, saving
the kernel's backward context so the ring Function needs no forward recompute. Returns
`(out[B,Hq,sl,D], saved)`; pass `saved` to :func:`sliding_ring_compute_fa_bwd`.