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

API

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.

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.

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

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

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

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

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.

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.

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

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.

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.