> 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_attention

Gemma4-specific context-parallel attention helpers.

## Module Contents

### Classes

| Name                                                                                                                         | Description                                                                                |
| ---------------------------------------------------------------------------------------------------------------------------- | ------------------------------------------------------------------------------------------ |
| [`CPRingAttentionContext`](#nemo_automodel-components-models-gemma4_moe-cp_attention-CPRingAttentionContext)                 | Inputs for Gemma4 manual ring CP attention (built by the run\_cp\_manual\_attention seam). |
| [`_Gemma4FFPAVarlenRingAttention`](#nemo_automodel-components-models-gemma4_moe-cp_attention-_Gemma4FFPAVarlenRingAttention) | -                                                                                          |
| [`_Gemma4FlexRingAttention`](#nemo_automodel-components-models-gemma4_moe-cp_attention-_Gemma4FlexRingAttention)             | -                                                                                          |

### Functions

| Name                                                                                                                                           | Description                                                                               |
| ---------------------------------------------------------------------------------------------------------------------------------------------- | ----------------------------------------------------------------------------------------- |
| [`_base_gemma4_cp_mask`](#nemo_automodel-components-models-gemma4_moe-cp_attention-_base_gemma4_cp_mask)                                       | -                                                                                         |
| [`_block_mask_set_generation`](#nemo_automodel-components-models-gemma4_moe-cp_attention-_block_mask_set_generation)                           | Reset the per-step block-mask cache when a new batch (new metadata) arrives.              |
| [`_build_packed_ring_segments`](#nemo_automodel-components-models-gemma4_moe-cp_attention-_build_packed_ring_segments)                         | Per-document varlen segments shared between a query shard and one KV chunk.               |
| [`_cached_block_mask`](#nemo_automodel-components-models-gemma4_moe-cp_attention-_cached_block_mask)                                           | -                                                                                         |
| [`_cached_ring_segments`](#nemo_automodel-components-models-gemma4_moe-cp_attention-_cached_ring_segments)                                     | Per-step cache around :func:`_build_packed_ring_segments`.                                |
| [`_chunk_dense_eligible`](#nemo_automodel-components-models-gemma4_moe-cp_attention-_chunk_dense_eligible)                                     | Select this chunk's dense-kernel eligibility for its role in the ring.                    |
| [`_collect_ring_kv_chunks`](#nemo_automodel-components-models-gemma4_moe-cp_attention-_collect_ring_kv_chunks)                                 | -                                                                                         |
| [`_compiled_flex_attention`](#nemo_automodel-components-models-gemma4_moe-cp_attention-_compiled_flex_attention)                               | -                                                                                         |
| [`_detach_metadata`](#nemo_automodel-components-models-gemma4_moe-cp_attention-_detach_metadata)                                               | -                                                                                         |
| [`_direct_exchange`](#nemo_automodel-components-models-gemma4_moe-cp_attention-_direct_exchange)                                               | -                                                                                         |
| [`_duck_shape_disabled`](#nemo_automodel-components-models-gemma4_moe-cp_attention-_duck_shape_disabled)                                       | Locally disable flex duck-shape specialization for the wrapped flex call.                 |
| [`_ffpa_varlen_backward_chunk`](#nemo_automodel-components-models-gemma4_moe-cp_attention-_ffpa_varlen_backward_chunk)                         | Per-chunk FFPA varlen backward using the *global* merged out/lse (packed).                |
| [`_ffpa_varlen_forward_chunk`](#nemo_automodel-components-models-gemma4_moe-cp_attention-_ffpa_varlen_forward_chunk)                           | One ring KV chunk via the FFPA varlen forward. Returns `(out[T,Hq,D], lse[Hq,T])`.        |
| [`_ffpa_varlen_ring_available`](#nemo_automodel-components-models-gemma4_moe-cp_attention-_ffpa_varlen_ring_available)                         | Whether the FFPA CuTeDSL *varlen* ops are ready (CPU-test monkeypatch seam).              |
| [`_gather_lse`](#nemo_automodel-components-models-gemma4_moe-cp_attention-_gather_lse)                                                         | Dense `[B, Hq, S]` -> packed `[Hq, T]` (inverse of :func:`_scatter_lse`).                 |
| [`_gather_thd`](#nemo_automodel-components-models-gemma4_moe-cp_attention-_gather_thd)                                                         | `[B, H, S, D]` -> packed `[T, H, D]` at flat `b*S + pos` indices.                         |
| [`_gemma4_cp_manual_attention`](#nemo_automodel-components-models-gemma4_moe-cp_attention-_gemma4_cp_manual_attention)                         | Gemma4-owned manual ring CP attention entry.                                              |
| [`_grad_add_live_`](#nemo_automodel-components-models-gemma4_moe-cp_attention-_grad_add_live_)                                                 | Add `grad_pack[T, H, D]` into `grad_bhsd[B, H, S, D]` at the `flat_index` rows only.      |
| [`_install_gemma4_cp_ring_sdpa`](#nemo_automodel-components-models-gemma4_moe-cp_attention-_install_gemma4_cp_ring_sdpa)                       | Swap `F.scaled_dot_product_attention` -> Gemma4 ring CP attention on this module.         |
| [`_merge_ffpa_packed_chunk`](#nemo_automodel-components-models-gemma4_moe-cp_attention-_merge_ffpa_packed_chunk)                               | `-inf`-safe online-softmax merge: like :func:`_merge_flex_chunk` but a row where both     |
| [`_merge_flex_chunk`](#nemo_automodel-components-models-gemma4_moe-cp_attention-_merge_flex_chunk)                                             | -                                                                                         |
| [`_merge_live_chunk_`](#nemo_automodel-components-models-gemma4_moe-cp_attention-_merge_live_chunk_)                                           | In-place online-softmax merge of one varlen chunk at its `flat_index` live rows only.     |
| [`_metadata_like`](#nemo_automodel-components-models-gemma4_moe-cp_attention-_metadata_like)                                                   | -                                                                                         |
| [`_patch_fsdp_accumulated_grad_guard`](#nemo_automodel-components-models-gemma4_moe-cp_attention-_patch_fsdp_accumulated_grad_guard)           | Guard `FSDPParam.to_accumulated_grad_if_needed` against uninitialized params.             |
| [`_ring_exchange`](#nemo_automodel-components-models-gemma4_moe-cp_attention-_ring_exchange)                                                   | One ring rotation step: send to `cp_rank+1`, receive from `cp_rank-1` (p2p).              |
| [`_ring_segment_set_generation`](#nemo_automodel-components-models-gemma4_moe-cp_attention-_ring_segment_set_generation)                       | Clear the ring-segment cache when a new batch (new q doc map) arrives.                    |
| [`_ring_use_ffpa_varlen`](#nemo_automodel-components-models-gemma4_moe-cp_attention-_ring_use_ffpa_varlen)                                     | Whether this ring attention call may use the FFPA *varlen* ring path.                     |
| [`_route_kv_grads_to_owners`](#nemo_automodel-components-models-gemma4_moe-cp_attention-_route_kv_grads_to_owners)                             | Sum each KV owner's dK/dV from every rank whose queries attended it.                      |
| [`_row_single_doc_full`](#nemo_automodel-components-models-gemma4_moe-cp_attention-_row_single_doc_full)                                       | Row is exactly one document with no padding (every position the same `>0` id).         |
| [`_row_single_doc_prefix`](#nemo_automodel-components-models-gemma4_moe-cp_attention-_row_single_doc_prefix)                                   | Row's real (`>0`) tokens are one document forming a contiguous prefix (pad, if any, is |
| [`_run_gemma4_cp_ffpa_varlen_ring_forward`](#nemo_automodel-components-models-gemma4_moe-cp_attention-_run_gemma4_cp_ffpa_varlen_ring_forward) | Forward of the FFPA ring: rotate K/V, per-chunk dense-or-varlen FFPA, online-merge.       |
| [`_run_gemma4_cp_ring_attention`](#nemo_automodel-components-models-gemma4_moe-cp_attention-_run_gemma4_cp_ring_attention)                     | Run Gemma4 local-query/ring-key CP attention.                                             |
| [`_run_gemma4_cp_ring_attention_forward`](#nemo_automodel-components-models-gemma4_moe-cp_attention-_run_gemma4_cp_ring_attention_forward)     | Run Gemma4 local-query/ring-key CP attention forward with FlexAttention.                  |
| [`_run_gemma4_flex_chunk`](#nemo_automodel-components-models-gemma4_moe-cp_attention-_run_gemma4_flex_chunk)                                   | -                                                                                         |
| [`_scatter_lse`](#nemo_automodel-components-models-gemma4_moe-cp_attention-_scatter_lse)                                                       | Varlen LSE `[Hq, T]` -> dense `[B, Hq, S]` (`-inf` for non-participating rows).           |
| [`_scatter_thd`](#nemo_automodel-components-models-gemma4_moe-cp_attention-_scatter_thd)                                                       | Inverse of :func:`_gather_thd`: `[T, H, D]` -> `[B, H, S, D]` (0 elsewhere).              |
| [`_use_live_row_merge`](#nemo_automodel-components-models-gemma4_moe-cp_attention-_use_live_row_merge)                                         | True when the live-token subset is small enough for in-place updates to beat scatter.     |
| [`_zero_if_none`](#nemo_automodel-components-models-gemma4_moe-cp_attention-_zero_if_none)                                                     | -                                                                                         |
| [`attach_gemma4_cp_ring_attention`](#nemo_automodel-components-models-gemma4_moe-cp_attention-attach_gemma4_cp_ring_attention)                 | Register Gemma4's model-owned p2p ring CP attention on a self-attention module.           |
| [`gemma4_vision_group_ids`](#nemo_automodel-components-models-gemma4_moe-cp_attention-gemma4_vision_group_ids)                                 | Return per-image-block ids for Gemma4 vision tokens, or -1 for text/padding.              |

### Data

[`_BLOCK_MASK_CACHE`](#nemo_automodel-components-models-gemma4_moe-cp_attention-_BLOCK_MASK_CACHE)

[`_BLOCK_MASK_GEN`](#nemo_automodel-components-models-gemma4_moe-cp_attention-_BLOCK_MASK_GEN)

[`_GEMMA4_CP_FLEX_RING_OK_LOGGED`](#nemo_automodel-components-models-gemma4_moe-cp_attention-_GEMMA4_CP_FLEX_RING_OK_LOGGED)

[`_LIVE_ROW_MERGE_MAX_FRACTION`](#nemo_automodel-components-models-gemma4_moe-cp_attention-_LIVE_ROW_MERGE_MAX_FRACTION)

[`_RING_SEGMENT_CACHE`](#nemo_automodel-components-models-gemma4_moe-cp_attention-_RING_SEGMENT_CACHE)

[`_RING_SEGMENT_GEN`](#nemo_automodel-components-models-gemma4_moe-cp_attention-_RING_SEGMENT_GEN)

[`logger`](#nemo_automodel-components-models-gemma4_moe-cp_attention-logger)

### API

```python
class nemo_automodel.components.models.gemma4_moe.cp_attention.CPRingAttentionContext(
    module: torch.nn.Module,
    query: torch.Tensor,
    key: torch.Tensor,
    value: torch.Tensor,
    cp_mesh: typing.Any,
    cp_group: typing.Any,
    cp_size: int,
    cp_rank: int,
    seq_local: int,
    seq_full: int,
    seq_global_start: int,
    attn_mask: typing.Any,
    dropout_p: float,
    is_causal: bool,
    scale: typing.Any,
    enable_gqa: bool,
    kwargs: dict[str, typing.Any],
    metadata: dict[str, torch.Tensor | None],
    metadata_seq_dims: dict[str, int]
)
```

Dataclass

Inputs for Gemma4 manual ring CP attention (built by the run\_cp\_manual\_attention seam).

```python
class nemo_automodel.components.models.gemma4_moe.cp_attention._Gemma4FFPAVarlenRingAttention()
```

**Bases:** `Function`

```python
nemo_automodel.components.models.gemma4_moe.cp_attention._Gemma4FFPAVarlenRingAttention.backward(
    autograd_ctx,
    grad_output: torch.Tensor
)
```

staticmethod

```python
nemo_automodel.components.models.gemma4_moe.cp_attention._Gemma4FFPAVarlenRingAttention.forward(
    autograd_ctx,
    query: torch.Tensor,
    key: torch.Tensor,
    value: torch.Tensor,
    ring_ctx: typing.Any
)
```

staticmethod

```python
class nemo_automodel.components.models.gemma4_moe.cp_attention._Gemma4FlexRingAttention()
```

**Bases:** `Function`

```python
nemo_automodel.components.models.gemma4_moe.cp_attention._Gemma4FlexRingAttention.backward(
    autograd_ctx,
    grad_output: torch.Tensor
)
```

staticmethod

```python
nemo_automodel.components.models.gemma4_moe.cp_attention._Gemma4FlexRingAttention.forward(
    autograd_ctx,
    query: torch.Tensor,
    key: torch.Tensor,
    value: torch.Tensor,
    ring_ctx: typing.Any
)
```

staticmethod

```python
nemo_automodel.components.models.gemma4_moe.cp_attention._base_gemma4_cp_mask(
    attention_module: torch.nn.Module,
    ctx: typing.Any,
    q_idx,
    kv_idx,
    kv_global_start: int = 0
)
```

```python
nemo_automodel.components.models.gemma4_moe.cp_attention._block_mask_set_generation(
    gen_tensor
) -> None
```

Reset the per-step block-mask cache when a new batch (new metadata) arrives.

```python
nemo_automodel.components.models.gemma4_moe.cp_attention._build_packed_ring_segments(
    q_ids: torch.Tensor,
    k_ids: torch.Tensor
) -> dict[str, typing.Any] | None
```

Per-document varlen segments shared between a query shard and one KV chunk.

`q_ids`/`k_ids` are the `[B, S]` `_packed_seq_ids` (`0` = pad, `>0` = doc id).
For each doc in *both* shards, pair its query/key tokens into one segment; returns flat
gather indices into `B*S` + int32 `cu_seqlens` (`cu_q[i]` pairs `cu_k[i]`), or
`None` when no doc is shared.

Dense-routing flags (consumed by :func:`_chunk_dense_eligible`): `dense_local` /
`dense_cross` say whether the chunk can skip the THD gather/scatter and feed the raw
`[B, H, S, D]` tensors to the dense FFPA kernel when used as the local (causal) or cross
(non-causal) chunk; `pad_rows` is the `[B, S]` pad-query mask to zero after a dense run
(`None` when unpadded).

```python
nemo_automodel.components.models.gemma4_moe.cp_attention._cached_block_mask(
    key,
    build
)
```

```python
nemo_automodel.components.models.gemma4_moe.cp_attention._cached_ring_segments(
    q_ids: torch.Tensor,
    k_ids: torch.Tensor,
    cp_rank: int,
    owner: int,
    cp_size: int
) -> dict[str, typing.Any] | None
```

Per-step cache around :func:`_build_packed_ring_segments`.

`owner` uniquely identifies the rotated k doc map within a step (each rank's
shard is visited exactly once), so `(cp_rank, owner, cp_size, B, Sq, Sk)`
keyed under the current q-doc-map generation uniquely identifies the segment.
`None` (no shared document) is cached too.

```python
nemo_automodel.components.models.gemma4_moe.cp_attention._chunk_dense_eligible(
    seg: dict[str, typing.Any],
    causal: bool
) -> bool
```

Select this chunk's dense-kernel eligibility for its role in the ring.

A local chunk runs causally (`causal=True`), a cross chunk non-causally; the two precomputed
flags from :func:`_build_packed_ring_segments` differ because only the cross chunk lacks a
causal mask to hide a second document or a pad tail in its KV. Dense-eligible chunks still
zero their pad query rows afterwards via `seg["pad_rows"]`.

```python
nemo_automodel.components.models.gemma4_moe.cp_attention._collect_ring_kv_chunks(
    ctx: typing.Any
) -> list[tuple[int, torch.Tensor, torch.Tensor, dict[str, torch.Tensor | None]]]
```

```python
nemo_automodel.components.models.gemma4_moe.cp_attention._compiled_flex_attention(
    attention_module: torch.nn.Module
)
```

```python
nemo_automodel.components.models.gemma4_moe.cp_attention._detach_metadata(
    metadata: dict[str, torch.Tensor | None]
) -> dict[str, torch.Tensor | None]
```

```python
nemo_automodel.components.models.gemma4_moe.cp_attention._direct_exchange(
    tensors: list[tuple[torch.Tensor, torch.Tensor]],
    cp_group: typing.Any,
    cp_rank: int,
    send_cp_rank: int,
    recv_cp_rank: int
) -> None
```

```python
nemo_automodel.components.models.gemma4_moe.cp_attention._duck_shape_disabled()
```

Locally disable flex duck-shape specialization for the wrapped flex call.

Otherwise the compiled flex kernel guards on incidental dim-equalities and recompiles on
every new sequence length. Dynamo reads `use_duck_shape` at (re)trace time inside the
flex call, so scoping it to the call window avoids mutating the process-global fx config.

```python
nemo_automodel.components.models.gemma4_moe.cp_attention._ffpa_varlen_backward_chunk(
    grad_out_pack: torch.Tensor,
    q_pack: torch.Tensor,
    k_pack: torch.Tensor,
    v_pack: torch.Tensor,
    out_pack: torch.Tensor,
    lse_pack: torch.Tensor,
    seg: dict[str, typing.Any],
    scale: float,
    causal: bool
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]
```

Per-chunk FFPA varlen backward using the *global* merged out/lse (packed).

```python
nemo_automodel.components.models.gemma4_moe.cp_attention._ffpa_varlen_forward_chunk(
    q_pack: torch.Tensor,
    k_pack: torch.Tensor,
    v_pack: torch.Tensor,
    seg: dict[str, typing.Any],
    scale: float,
    causal: bool
) -> tuple[torch.Tensor, torch.Tensor]
```

One ring KV chunk via the FFPA varlen forward. Returns `(out[T,Hq,D], lse[Hq,T])`.

```python
nemo_automodel.components.models.gemma4_moe.cp_attention._ffpa_varlen_ring_available() -> bool
```

Whether the FFPA CuTeDSL *varlen* ops are ready (CPU-test monkeypatch seam).

```python
nemo_automodel.components.models.gemma4_moe.cp_attention._gather_lse(
    lse_dense: torch.Tensor,
    flat_index: torch.Tensor
) -> torch.Tensor
```

Dense `[B, Hq, S]` -> packed `[Hq, T]` (inverse of :func:`_scatter_lse`).

```python
nemo_automodel.components.models.gemma4_moe.cp_attention._gather_thd(
    t_bhsd: torch.Tensor,
    flat_index: torch.Tensor
) -> torch.Tensor
```

`[B, H, S, D]` -> packed `[T, H, D]` at flat `b*S + pos` indices.

```python
nemo_automodel.components.models.gemma4_moe.cp_attention._gemma4_cp_manual_attention(
    attention_module: torch.nn.Module,
    query: torch.Tensor,
    key: torch.Tensor,
    value: torch.Tensor,
    cp_mesh,
    attn_mask,
    dropout_p,
    is_causal,
    scale,
    enable_gqa,
    kwargs
) -> torch.Tensor
```

Gemma4-owned manual ring CP attention entry.

Plugs into cp\_utils' generic `run_cp_manual_attention` seam: receives the
raw local (un-gathered) Q/K/V plus `cp_mesh`, builds the ring context, and
runs the p2p ring FlexAttention. K/V are rotated across CP ranks inside the
ring autograd function -- they are never all-gathered.

```python
nemo_automodel.components.models.gemma4_moe.cp_attention._grad_add_live_(
    grad_bhsd: torch.Tensor,
    grad_pack: torch.Tensor,
    flat_index: torch.Tensor
) -> None
```

Add `grad_pack[T, H, D]` into `grad_bhsd[B, H, S, D]` at the `flat_index` rows only.

The in-place advanced-index update on the permuted view keeps `grad_bhsd` BSHD-contiguous
(dense backward adds full tensors into it unchanged) and avoids the `torch.zeros` +
full-tensor add of `grad = grad + _scatter_thd(...)`. `flat_index` is duplicate-free, so
the assign is a true add.

```python
nemo_automodel.components.models.gemma4_moe.cp_attention._install_gemma4_cp_ring_sdpa(
    attention_module: torch.nn.Module,
    cp_mesh
) -> None
```

Swap `F.scaled_dot_product_attention` -> Gemma4 ring CP attention on this module.

Gemma4 owns its CP attention end-to-end (it does not use cp\_utils' generic CP
SDPA hooks). It installs its own `@torch._dynamo.disable` SDPA wrapper -- on
the inner attention module so it also fires during gradient-checkpointing
recompute -- that runs the p2p ring FlexAttention. The per-forward attention
kwargs the ring needs (mm\_token\_type\_ids, packed-seq ids, padding/vision masks)
are captured off the forward kwargs into `_cp_manual_metadata` here, since the
swapped SDPA only receives Q/K/V.

```python
nemo_automodel.components.models.gemma4_moe.cp_attention._merge_ffpa_packed_chunk(
    out_acc: torch.Tensor | None,
    lse_acc: torch.Tensor | None,
    out_step: torch.Tensor,
    lse_step: torch.Tensor
) -> tuple[torch.Tensor, torch.Tensor]
```

`-inf`-safe online-softmax merge: like :func:`_merge_flex_chunk` but a row where both
inputs are `-inf` (pad query no chunk covered) keeps 0 instead of `NaN`.

On finite rows `old_scale + new_scale == 1`, so the combine is a single `lerp`; the only
`NaN` is the both-`-inf` row, which `nan_to_num_` zeroes (leaving its already-0 output).
The fused `lerp` is \~2.5x cheaper than the explicit `isneginf`/`where` form.

```python
nemo_automodel.components.models.gemma4_moe.cp_attention._merge_flex_chunk(
    out_acc: torch.Tensor | None,
    lse_acc: torch.Tensor | None,
    out_step: torch.Tensor,
    lse_step: torch.Tensor
) -> tuple[torch.Tensor, torch.Tensor]
```

```python
nemo_automodel.components.models.gemma4_moe.cp_attention._merge_live_chunk_(
    out_acc: torch.Tensor,
    lse_acc: torch.Tensor,
    out_pack: torch.Tensor,
    lse_pack: torch.Tensor,
    flat_index: torch.Tensor
) -> None
```

In-place online-softmax merge of one varlen chunk at its `flat_index` live rows only.

Like :func:`_merge_ffpa_packed_chunk` but updates only the `T` live rows of the running
`[B, Hq, S, D]`/`[B, Hq, S]` accumulators. The cheap `lerp` (no `-inf` guard) is safe:
a cross chunk targets only real query tokens, all of which the local chunk already covered,
so `lse_acc` is finite at every live row.

```python
nemo_automodel.components.models.gemma4_moe.cp_attention._metadata_like(
    metadata: dict[str, torch.Tensor | None]
) -> dict[str, torch.Tensor | None]
```

```python
nemo_automodel.components.models.gemma4_moe.cp_attention._patch_fsdp_accumulated_grad_guard() -> None
```

Guard `FSDPParam.to_accumulated_grad_if_needed` against uninitialized params.

On some torch builds that method reads `self._unsharded_param` (the lazily
set unsharded tensor) without first checking it exists. In FSDP2 post-backward
under fp32 grad-reduce, frozen / never-unsharded params (e.g. the frozen Gemma4
vision tower and embeddings) have no `_unsharded_param` yet and it raises
`AttributeError`. Such params carry no grad to upcast anyway, so wrap the
method to skip them when uninitialized. No-op once applied / on fixed builds.

```python
nemo_automodel.components.models.gemma4_moe.cp_attention._ring_exchange(
    tensors: list[tuple[torch.Tensor, torch.Tensor]],
    cp_group: typing.Any,
    cp_rank: int,
    cp_size: int
) -> None
```

One ring rotation step: send to `cp_rank+1`, receive from `cp_rank-1` (p2p).

```python
nemo_automodel.components.models.gemma4_moe.cp_attention._ring_segment_set_generation(
    gen_tensor: torch.Tensor
) -> None
```

Clear the ring-segment cache when a new batch (new q doc map) arrives.

```python
nemo_automodel.components.models.gemma4_moe.cp_attention._ring_use_ffpa_varlen(
    attention_module: torch.nn.Module,
    ctx: typing.Any
) -> bool
```

Whether this ring attention call may use the FFPA *varlen* ring path.

The decision gates collective p2p exchanges so it must be rank-uniform: it
depends only on the per-layer config, the head\_dim / dtype / scale, and
*whether* `_packed_seq_ids` is present (a batch-level fact) -- never on
per-rank slice content. This path *requires* `_packed_seq_ids` (the document
map drives the varlen `cu_seqlens`); Gemma4's manual CP batch always attaches
one, so it is the sole FFPA path real CP training takes.

```python
nemo_automodel.components.models.gemma4_moe.cp_attention._route_kv_grads_to_owners(
    grad_key_by_owner: dict[int, torch.Tensor],
    grad_value_by_owner: dict[int, torch.Tensor],
    cp_group: typing.Any,
    cp_rank: int,
    cp_size: int
) -> tuple[torch.Tensor, torch.Tensor]
```

Sum each KV owner's dK/dV from every rank whose queries attended it.

In the ring forward, rank `r`'s queries attend the chunks owned by ranks
`r, r-1, ..., 0`, so backward produces a dK/dV contribution for each of those
owners on rank `r`. This sends each owner's contribution back to that owner
(p2p) and sums them, returning the local rank's accumulated `(grad_key,
grad_value)`. Shared by the Flex and FFPA-varlen ring backward passes.

```python
nemo_automodel.components.models.gemma4_moe.cp_attention._row_single_doc_full(
    row: list[int]
) -> bool
```

Row is exactly one document with no padding (every position the same `>0` id).

```python
nemo_automodel.components.models.gemma4_moe.cp_attention._row_single_doc_prefix(
    row: list[int]
) -> bool
```

Row's real (`>0`) tokens are one document forming a contiguous prefix (pad, if any, is
a trailing suffix) -- the shape a single-document shard takes after tail padding.

```python
nemo_automodel.components.models.gemma4_moe.cp_attention._run_gemma4_cp_ffpa_varlen_ring_forward(
    ctx: typing.Any,
    seg_sink: dict[int, typing.Any] | None = None
) -> tuple[torch.Tensor, torch.Tensor]
```

Forward of the FFPA ring: rotate K/V, per-chunk dense-or-varlen FFPA, online-merge.

Returns `(out_final[B,Hq,S,D] fp32, lse_final[B,Hq,S] fp32)`. `owner > cp_rank` chunks
are causally skipped; chunks sharing no document contribute nothing; pad query rows stay 0.
When `seg_sink` is given, each processed `owner`'s segment is recorded (`owner -> seg`)
so the backward reuses the same path choice and indices.

```python
nemo_automodel.components.models.gemma4_moe.cp_attention._run_gemma4_cp_ring_attention(
    attention_module: torch.nn.Module,
    ctx: typing.Any
) -> torch.Tensor
```

Run Gemma4 local-query/ring-key CP attention.

Full-attention (global) head\_dim=512 layers run their per-chunk attention
through the FFPA CuTeDSL kernel (`_ring_use_ffpa_varlen` gate), choosing per
chunk between the *dense* `_fwd_cute` path (single full document per row -- the
unpacked / synthesized-single-document case -- on the raw `[B, H, S, D]` tensors,
zero gather/scatter) and the *varlen* `_varlen_fwd_cute` path (genuinely packed /
straddling shards, which need cross-document masking via THD `cu_seqlens`). The
manual CP batch always attaches a `_packed_seq_ids` map
(`cp_batch._synthesize_single_document_seq_ids` injects a trivial single-document
one when the batch is not packed), so this is the path real CP training -- packed or
unpacked -- always takes. Every other layer / batch (sliding-window, no kernel,
wrong dtype/head\_dim) keeps using compiled FlexAttention.

```python
nemo_automodel.components.models.gemma4_moe.cp_attention._run_gemma4_cp_ring_attention_forward(
    attention_module: torch.nn.Module,
    ctx: typing.Any
) -> torch.Tensor
```

Run Gemma4 local-query/ring-key CP attention forward with FlexAttention.

```python
nemo_automodel.components.models.gemma4_moe.cp_attention._run_gemma4_flex_chunk(
    attention_module: torch.nn.Module,
    ctx: typing.Any,
    key_chunk: torch.Tensor,
    value_chunk: torch.Tensor,
    metadata_chunk: dict[str, torch.Tensor | None],
    kv_global_start: int
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None, int]
```

```python
nemo_automodel.components.models.gemma4_moe.cp_attention._scatter_lse(
    lse_pack: torch.Tensor,
    flat_index: torch.Tensor,
    B: int,
    Hq: int,
    S: int
) -> torch.Tensor
```

Varlen LSE `[Hq, T]` -> dense `[B, Hq, S]` (`-inf` for non-participating rows).

```python
nemo_automodel.components.models.gemma4_moe.cp_attention._scatter_thd(
    packed: torch.Tensor,
    flat_index: torch.Tensor,
    B: int,
    H: int,
    S: int
) -> torch.Tensor
```

Inverse of :func:`_gather_thd`: `[T, H, D]` -> `[B, H, S, D]` (0 elsewhere).

```python
nemo_automodel.components.models.gemma4_moe.cp_attention._use_live_row_merge(
    num_live: int,
    num_rows: int
) -> bool
```

True when the live-token subset is small enough for in-place updates to beat scatter.

```python
nemo_automodel.components.models.gemma4_moe.cp_attention._zero_if_none(
    grad: torch.Tensor | None,
    like: torch.Tensor
) -> torch.Tensor
```

```python
nemo_automodel.components.models.gemma4_moe.cp_attention.attach_gemma4_cp_ring_attention(
    attention_module: torch.nn.Module,
    use_ffpa: bool = False
) -> None
```

Register Gemma4's model-owned p2p ring CP attention on a self-attention module.

Declares the metadata keys the ring needs and exposes `setup_cp_attention(cp_mesh)`
\-- the model-owned CP-attention seam the parallelizer calls (with the CP mesh)
instead of cp\_utils' generic SDPA hooks. `run_cp_manual_attention` is also bound
as the ring entry point.

`use_ffpa` opts the (full-attention, head\_dim=512) ring chunks into the FFPA
CuTeDSL kernel; `_ring_use_ffpa_varlen` still verifies per-call eligibility, so
this is a no-op for sliding-window layers, non-512 head\_dim, wrong dtype, or when
the FFPA kernel is unavailable.

```python
nemo_automodel.components.models.gemma4_moe.cp_attention.gemma4_vision_group_ids(
    mm_token_type_ids: torch.Tensor
) -> torch.Tensor
```

Return per-image-block ids for Gemma4 vision tokens, or -1 for text/padding.

```python
nemo_automodel.components.models.gemma4_moe.cp_attention._BLOCK_MASK_CACHE: dict = {}
```

```python
nemo_automodel.components.models.gemma4_moe.cp_attention._BLOCK_MASK_GEN: list = [None, None]
```

```python
nemo_automodel.components.models.gemma4_moe.cp_attention._GEMMA4_CP_FLEX_RING_OK_LOGGED = False
```

```python
nemo_automodel.components.models.gemma4_moe.cp_attention._LIVE_ROW_MERGE_MAX_FRACTION = 0.25
```

```python
nemo_automodel.components.models.gemma4_moe.cp_attention._RING_SEGMENT_CACHE: dict = {}
```

```python
nemo_automodel.components.models.gemma4_moe.cp_attention._RING_SEGMENT_GEN: list = [None, None]
```

```python
nemo_automodel.components.models.gemma4_moe.cp_attention.logger = logging.getLogger(__name__)
```