nemo_automodel.components.models.gemma4_moe.cp_attention

Inputs for Gemma4 manual ring CP attention (built by the run_cp_manual_attention seam).

Bases: Function

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

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
)

Bases: Function

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

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
)

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

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

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.

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

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.

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

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

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

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

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

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.

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.

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.

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

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.

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.

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

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

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.

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.

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

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.

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.

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.

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

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

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

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

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.

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