Support learnable softmax (attention sink) with cp_comm_type="p2p"

The p2p fused-attention helpers never passed softmax_type to fused_attn_fwd or
fused_attn_bwd, so the kernels fell back to their "vanilla" default and discarded
softmax_offset on the floor. Nothing failed loudly: the forward returned the
vanilla answer and the sink gradient came back as an exact zero. The a2a path has
always passed it; p2p never did.

Thread softmax_type through both helpers, add the offset to the auxiliary tensor
pack in the order the forward returns it, and accumulate the per-tile sink gradient
into the p2p backward's return tuple. The offset is replicated across CP ranks and
every rank computes all query rows, so each already holds the full gradient and the
outer framework's DDP average reproduces it; no extra reduction is needed here.

The aux pack is [stats, rng, (bias), (offset)], so the previously unconditional
attn_bias = rest[0] is now guarded -- otherwise it would hand the offset back as the
attention bias once a sink is active.

Also bind d_sink_ before the section dispatch: it is assigned only in the fused
backward branch, so a vanilla run on the flash backend would have raised
UnboundLocalError at the accumulation.

Verified against an FP64 oracle at cp=2, BF16, 1024 tokens: the p2p forward matches
to 9.5e-3 relative error and dSink to 1.7e-3, on par with the already-supported a2a
path. The oracle's own autograd and finite-difference answers agree to 2.3e-08.
Not covered: qkv_format="thd", cp_size > 2, and the flash backend.

Signed-off-by: 0z5a <dezhen.lu@student.uni-tuebingen.de>

Author: 0z5a

transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py

@@ -994,6 +994,8 @@ def cp_p2p_fwd_fused_attn(
     step,
     cp_size,
     return_max_logit,
+    softmax_type,
+    softmax_offset,
     q_part,
     k_part,
     v_part,
@@ -1110,6 +1112,8 @@ def cp_p2p_fwd_fused_attn(
         cu_seqlens_kv_padded=cu_seqlens_kv_padded_,
         **fp8_meta_kwargs,
         return_max_logit=return_max_logit,
+        softmax_type=softmax_type,
+        softmax_offset=softmax_offset,
         cuda_graph=is_graph_capturing(),
         qkv_scale_inv_format=qkv_scale_inv_format,
     )
@@ -1118,7 +1122,10 @@ def cp_p2p_fwd_fused_attn(
         softmax_lse_per_step, rng_states = aux_ctx_tensors
     else:
         softmax_lse_per_step, rng_states, *rest = aux_ctx_tensors
-        attn_bias = rest[0] if len(rest) > 0 else None
+        # The aux pack is [stats, rng, (bias), (offset)], so an unconditional
+        # rest[0] would return the offset as the bias once a sink is active.
+        has_attn_bias = attn_bias_type not in ["no_bias", "alibi"] and attn_bias_inputs is not None
+        attn_bias = rest[0] if has_attn_bias else None
 
     if return_max_logit:
         return out_per_step, softmax_lse_per_step, rng_states, attn_bias, *max_logit
@@ -1314,6 +1321,8 @@ def cp_p2p_bwd_fused_attn(
     softmax_lse,
     softmax_lse_,
     rng_states,
+    softmax_type,
+    softmax_offset,
     attn_dbias,
     attn_biases,
     max_seqlen_q,
@@ -1374,6 +1383,11 @@ def cp_p2p_bwd_fused_attn(
     if attn_dbias is not None:
         aux_tensors += [attn_biases[cp_size - step - 1]]
 
+    # Aux pack order is [softmax_stats, rng_state, (bias), (softmax_offset)],
+    # matching what the forward returned.
+    if softmax_offset is not None:
+        aux_tensors += [softmax_offset]
+
     fp8_meta_kwargs = {}
     qkv_scale_inv_format = None
     do_scale_inv_format = None
@@ -1410,7 +1424,7 @@ def cp_p2p_bwd_fused_attn(
         fp8_meta_kwargs["dp_quantizer"] = dP_quantizer_per_step
         fp8_meta_kwargs["dqkv_quantizer"] = dQKV_quantizer_per_step
 
-    dq, dk, dv, dbias, *_ = fused_attn_bwd(
+    fused_bwd_out = fused_attn_bwd(
         max_seqlen_q_,
         max_seqlen_kv_,
         cu_seqlens_q_per_step[cp_size - step - 1],
@@ -1433,14 +1447,16 @@ def cp_p2p_bwd_fused_attn(
         dqkv_layout=dqkv_layout,
         attn_mask_type=attn_mask_type_,
         attn_bias_type=attn_bias_type,
+        softmax_type=softmax_type,
         deterministic=deterministic,
         cuda_graph=is_graph_capturing(),
         qkv_scale_inv_format=qkv_scale_inv_format,
         do_scale_inv_format=do_scale_inv_format,
         **fp8_meta_kwargs,
     )
 
-    return dq, dk, dv, dbias
+    dq, dk, dv, dbias, d_softmax_offset = fused_bwd_out
+    return dq, dk, dv, dbias, d_softmax_offset
 
 
 def cp_p2p_bwd_flash_attn(
@@ -1618,6 +1634,7 @@ def forward(
         use_fused_attention,
         return_max_logit,
         softcap,
+        softmax_offset,
         fp8,
         fp8_meta,
         cp_group,
@@ -1629,6 +1646,7 @@ def forward(
         use_flash_attn_4,
         fp8_output,
         layer_number,
+        softmax_type,
     ):
         # pylint: disable=missing-function-docstring
 
@@ -2007,6 +2025,8 @@ def forward(
                             i,
                             cp_size,
                             return_max_logit,
+                            softmax_type,
+                            softmax_offset,
                         ]
                     else:
                         flash_attn_inputs = [
@@ -2396,6 +2416,8 @@ def forward(
         ctx.max_seqlen_q = max_seqlen_q
         ctx.max_seqlen_kv = max_seqlen_kv
         ctx.softmax_scale = softmax_scale
+        ctx.softmax_offset = softmax_offset
+        ctx.softmax_type = softmax_type
         ctx.attn_mask_type = attn_mask_type
         ctx.attn_bias_type = attn_bias_type
         ctx.attn_bias_shape = None if attn_bias is None else attn_bias.shape
@@ -2520,6 +2542,12 @@ def backward(ctx, dout, *_args):
             attn_dbias = None
             attn_dbias_ = None
 
+        # Sink gradient accumulator. The offset is replicated on every rank and every
+        # per-tile fused call sees the same tensor, so tile contributions sum.
+        d_softmax_offset = (
+            torch.zeros_like(ctx.softmax_offset) if ctx.softmax_offset is not None else None
+        )
+
         # set up softmax_lse
         softmax_lse_ = None
         if causal and ctx.second_half_lse_seqlen is not None:
@@ -2779,6 +2807,8 @@ def backward(ctx, dout, *_args):
                     softmax_lse,
                     softmax_lse_,
                     rng_states,
+                    ctx.softmax_type,
+                    ctx.softmax_offset,
                     attn_dbias,
                     attn_biases,
                     ctx.max_seqlen_q,
@@ -2831,12 +2861,13 @@ def backward(ctx, dout, *_args):
             # Reverse the steps in forward. In the cp_size x cp_size (i.e. GPU x step) matrix,
             # there are still three sections in these tiles based on their attention pattern
             # for attn_mask_type = causal, and one for attn_mask_type != causal.
+            d_sink_ = None
             if causal:
                 if i == (cp_size - 1):
                     section = "diagonal"
                     prepare_outputs = cp_p2p_bwd_prepare_qkv(*prepare_inputs, section)
                     if ctx.use_fused_attention:
-                        dq_, dk_, dv_, dbias_ = cp_p2p_bwd_fused_attn(
+                        dq_, dk_, dv_, dbias_, d_sink_ = cp_p2p_bwd_fused_attn(
                             *fused_attn_inputs, *prepare_outputs, section
                         )
                     else:
@@ -2849,7 +2880,7 @@ def backward(ctx, dout, *_args):
                     section = "lower-triangle"
                     prepare_outputs = cp_p2p_bwd_prepare_qkv(*prepare_inputs, section)
                     if ctx.use_fused_attention:
-                        dq_, dk_, dv_, dbias_ = cp_p2p_bwd_fused_attn(
+                        dq_, dk_, dv_, dbias_, d_sink_ = cp_p2p_bwd_fused_attn(
                             *fused_attn_inputs, *prepare_outputs, section
                         )
                     else:
@@ -2862,7 +2893,7 @@ def backward(ctx, dout, *_args):
                     section = "upper-triangle"
                     prepare_outputs = cp_p2p_bwd_prepare_qkv(*prepare_inputs, section)
                     if ctx.use_fused_attention:
-                        dq_, dk_, dv_, dbias_ = cp_p2p_bwd_fused_attn(
+                        dq_, dk_, dv_, dbias_, d_sink_ = cp_p2p_bwd_fused_attn(
                             *fused_attn_inputs, *prepare_outputs, section
                         )
                     else:
@@ -2875,7 +2906,7 @@ def backward(ctx, dout, *_args):
                 section = "all"
                 prepare_outputs = cp_p2p_bwd_prepare_qkv(*prepare_inputs, section)
                 if ctx.use_fused_attention:
-                    dq_, dk_, dv_, dbias_ = cp_p2p_bwd_fused_attn(
+                    dq_, dk_, dv_, dbias_, d_sink_ = cp_p2p_bwd_fused_attn(
                         *fused_attn_inputs, *prepare_outputs, section
                     )
                 else:
@@ -2885,6 +2916,13 @@ def backward(ctx, dout, *_args):
                         section,
                     )
 
+            # Accumulate this tile's sink gradient. Every rank computes all query rows,
+            # so each already holds the full gradient for this replicated parameter and
+            # the outer framework's DDP average reproduces it; reducing here as well
+            # would double-count.
+            if d_sink_ is not None:
+                d_softmax_offset += d_sink_
+
             # dq, dk, dv are reduced across steps in higher precision
             # DelayedScaling: collect all results in uint8 to one tensor, dequantize to float32, then reduce
             # CurrentScaling: dequantize partial results from each step to float32, then reduce
@@ -3204,6 +3242,8 @@ def backward(ctx, dout, *_args):
             None,
             None,
             None,
+            d_softmax_offset,
+            None,
             None,
             None,
             None,
@@ -5715,9 +5755,10 @@ def attn_forward_func_with_cp(
     assert (
         softmax_type == "vanilla" or use_fused_attention
     ), f"Context parallelism only supports {softmax_type=} with FusedAttention backend!"
-    assert (
-        softmax_type == "vanilla" or cp_comm_type == "a2a"
-    ), f"Context parallelism only supports {softmax_type=} with cp_comm_type = 'a2a'!"
+    assert softmax_type == "vanilla" or cp_comm_type in (
+        "a2a",
+        "p2p",
+    ), f"Context parallelism supports {softmax_type=} with cp_comm_type in ('a2a', 'p2p')!"
     if get_cudnn_version() < (9, 18, 0):
         assert softmax_type == "vanilla" or qkv_format != "thd", (
             f"Before cuDNN 9.18.0, context parallelism does not support {softmax_type=} with"
@@ -5749,6 +5790,7 @@ def attn_forward_func_with_cp(
 
     if cp_comm_type in ["p2p", "a2a+p2p"]:
         args += [
+            softmax_offset,
             fp8,
             fp8_meta,
             cp_group,
@@ -5760,6 +5802,7 @@ def attn_forward_func_with_cp(
             use_flash_attn_4,
             fp8_output,
             layer_number,
+            softmax_type,
         ]
         out = AttnFuncWithCPAndKVP2P.apply(*args)
     elif cp_comm_type == "all_gather":

transformer_engine/pytorch/attention/dot_product_attention/utils.py

@@ -1245,7 +1245,7 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt
                 softmax_type,
             )
             use_fused_attention = False
-        if context_parallel and cp_comm_type != "a2a":
+        if context_parallel and cp_comm_type not in ("a2a", "p2p"):
             logger.debug(
                 "Disabling FusedAttention for context parallelism with softmax_type = %s and"
                 " cp_comm_type = %s",