`core.ops.attention.csa.output_rope`#

CSA output rotation and alias-safe backward reconstruction.

The attention Function owns and rotates its output in place. Backward reconstructs both the raw output and its gradient into private storage. Reconstruction trades one additional low-precision rotation for not retaining a second output activation.

Module Contents#

Classes#

OutputRopeParams

Inverse-RoPE applied to the attention output inside the fused Functions.

Functions#

`_mla_rope_reconstruct_output_kernel`	Undo the output rotation and copy its non-rotary prefix into OUT.
`_raw_output_rope`
`mla_rope_apply_raw_`	Rotate an attention-owned output without adding an autograd node.
`mla_rope_unapply_raw`	Undo output RoPE, fusing non-rotary copy when rebuilding saved output.
`undo_output_rope_for_backward`	Reconstruct cuDNN inputs without mutating saved output or upstream gradient.

Data#

_autotuned_mla_rope_reconstruct_output_kernel

API#

core.ops.attention.csa.output_rope._mla_rope_reconstruct_output_kernel( DO, OUT, COS, SIN, nope_dim, emb_dim: megatron.core.fusions.fused_mla_yarn_rope_apply.tl.constexpr, head_num: megatron.core.fusions.fused_mla_yarn_rope_apply.tl.constexpr, batch_size, seq_num, cu_seqlens_q, position_ids, stride_x_seq, stride_x_nheads, stride_cos_seq, stride_sin_seq, cp_rank, cp_size, INVERSE: megatron.core.fusions.fused_mla_yarn_rope_apply.tl.constexpr, REMOVE_INTERLEAVING: megatron.core.fusions.fused_mla_yarn_rope_apply.tl.constexpr, BLOCK_H: megatron.core.fusions.fused_mla_yarn_rope_apply.tl.constexpr, BLOCK_NOPE: megatron.core.fusions.fused_mla_yarn_rope_apply.tl.constexpr, )#: Undo the output rotation and copy its non-rotary prefix into OUT.

core.ops.attention.csa.output_rope._autotuned_mla_rope_reconstruct_output_kernel#: ‘(…)’

core.ops.attention.csa.output_rope._raw_output_rope( t, cos, sin, nope_dim, emb_dim, cu_seqlens_q=None, cp_rank=0, cp_size=1, inverse=False, remove_interleaving=False, position_ids=None, undo=False, )#

core.ops.attention.csa.output_rope.mla_rope_apply_raw_(t, cos, sin, nope_dim, emb_dim, **kwargs)#: Rotate an attention-owned output without adding an autograd node.

core.ops.attention.csa.output_rope.mla_rope_unapply_raw( t, cos, sin, nope_dim, emb_dim, out=None, cu_seqlens_q=None, cp_rank=0, cp_size=1, inverse=False, remove_interleaving=False, position_ids=None, )#: Undo output RoPE, fusing non-rotary copy when rebuilding saved output.

class core.ops.attention.csa.output_rope.OutputRopeParams#

Inverse-RoPE applied to the attention output inside the fused Functions.

DSv4 undoes the query-side RoPE on the attention output. Folding that into the sparse-attention Function lets the rotation happen in place on the tensor the Function already saves for backward, instead of on a private clone, which removes one full (rows, heads, v_head_dim) buffer per layer from the forward-to-backward live set.

The rotation is orthogonal only when the cos/sin were built with mscale=1.0 (DSv4’s “pure rotation” contract). That is what makes

Meth:: undo an inverse up to floating-point rounding, so backward can recover the pre-RoPE output and gradient.

sbhd_batch_size is required for the sbhd layout, where the flat (sq * b, ...) rows are sq-major and the kernel derives the token index as row // batch_size. For thd it must be None and cu_seqlens_q supplies the segment boundaries instead.

cos: torch.Tensor#: None

sin: torch.Tensor#: None

nope_dim: int#: None

pos_dim: int#: None

cu_seqlens_q: Optional[torch.Tensor]#: None

cp_rank: int#: 0

cp_size: int#: 1

position_ids: Optional[torch.Tensor]#: None

sbhd_batch_size: Optional[int]#: None

_kernel_view(out_flat: torch.Tensor) → torch.Tensor#: View a flat (rows, heads, head_dim) output the way the kernel wants it.

_kernel_kwargs() → dict#

apply_(out_flat: torch.Tensor) → torch.Tensor#: Inverse-RoPE a (rows, heads, head_dim) attention output in place.

undo( t: torch.Tensor, out: Optional[torch.Tensor] = None, ) → torch.Tensor#

Recover the pre-RoPE values of t, in place or into out.

Pass out whenever t is aliased by a tensor that another autograd node still expects to hold the rotated values, which is the case for the saved attention output.

core.ops.attention.csa.output_rope.undo_output_rope_for_backward( out_rope: Optional[core.ops.attention.csa.output_rope.OutputRopeParams], out_flat: torch.Tensor, dO_flat: torch.Tensor, ) → Tuple[torch.Tensor, torch.Tensor]#

Reconstruct cuDNN inputs without mutating saved output or upstream gradient.

Either input may be shared with another autograd consumer. contiguous() alone does not establish exclusive ownership, so use separate destinations.

core.ops.attention.csa.output_rope#