[0z5a]

Incremental review: Files changed against main.

Addresses #2541.

What was broken

softmax_type="learnable" (the learnable attention sink) with cp_comm_type="p2p" was rejected by an assertion. Behind that assertion, the p2p fused-attention helpers never passed softmax_type to fused_attn_fwd / fused_attn_bwd, so the kernels fell back to "vanilla" and ignored softmax_offset.

What changed

• cp_p2p_fwd_fused_attn / cp_p2p_bwd_fused_attn take softmax_type and softmax_offset. The offset joins the auxiliary tensor pack in the order the forward returns it ([stats, rng, (bias), (offset)]). The previously unconditional attn_bias = rest[0] is now guarded by whether a bias was passed, because with a sink the pack would otherwise return the offset as the bias.
• The sink enters each query row's softmax exactly once. cuDNN counts the sink as one more softmax column, so the softmax stats of a tile include e^offset. The p2p forward merges the tiles' stats, so passing the offset to every tile would add the sink once per tile that a row visits. Only step 0 gets the sink: it is the local KV tile, and it covers every local query row (the diagonal tile under a causal mask, the full tile without one). The other tiles run a plain softmax. In the backward, only the mirror tile passes the sink and returns its gradient. The other tiles' dq/dk/dv use the merged stats, which already include the sink.
• Sink gradient reduction. Each rank's sink gradient covers only its own query rows. It is all-reduced over the CP group, so every rank returns the full-sequence gradient. This matches the a2a path, which all-gathers the per-head gradients, and it is what run_attention_with_cp.py compares against.
• utils.py and attn_forward_func_with_cp: fused attention is no longer disabled for non-vanilla softmax with p2p, except with FP8 attention. That case is unverified, and the FP8 branch of the per-tile forward unpacks a two-entry aux pack, so it stays off FusedAttention.
• Tests:
  • test_attention_with_cp.py now runs the non-vanilla softmax cases (cp_4_1 off-by-one, cp_4_2 learnable) for p2p as well as a2a.
  • run_attention_with_cp.py kept the reference sink gradient as an alias of softmax_offset.grad, then zeroed it before the CP run accumulated into the same tensor. The sink-gradient check therefore compared the CP result with itself, for a2a as well. It now clones the reference.
  • The learnable offsets are drawn as randn * 5 from a fixed CPU generator. With the default zero offset, the sink is almost invisible next to thousands of keys.

An earlier version of this PR passed the sink to every tile and did not reduce the gradient. Both errors cancel without a mask at cp=2: every row was counted twice, and each rank held half the rows. Under a causal mask they don't cancel: the sum of the per-rank sink gradients was off by up to 0.84 (reference max 6.36). The CP test skipped non-vanilla softmax for p2p, and its sink-gradient check compared against an alias, so neither caught it.

Verification

Two GH200 nodes (sm_90), one GPU each, cross-node over InfiniBand, cp=2, BF16, FusedAttention. The run used TE main fac8d4a9 with this PR's Python files on top, cuDNN 9.24.0, and torch 2.14.1+cu130. run_attention_with_cp.py was run with the changes above:

softmax	comm	layout	TE Python	result
learnable (cp_4_2)	p2p	bshd	this PR	pass
learnable (cp_4_2)	p2p	sbhd	this PR	pass
off-by-one (cp_4_1)	p2p	bshd	this PR	pass
vanilla (cp_4_0)	p2p	bshd	this PR	pass
learnable (cp_4_2)	a2a	bshd	main	pass
off-by-one (cp_4_1)	a2a	bshd	main	pass
learnable (cp_4_2)	p2p	bshd	previous PR head 649969bb	fail: sink gradient, 12/64 heads, max abs diff 1.55

"pass" means out, dq, dk, dv and the sink gradient are all within the runner's tolerances. All configs are causal GQA with s=4096. The two a2a rows show that the runner changes keep the existing a2a cases passing.

Sink gradient, per rank, max abs diff against the non-CP reference (cp_4_2, offsets randn * 5):

mask	reference max	previous PR head	this PR
causal	6.36	0.84 (sum over ranks)	0.0032
no_mask	2.55	0.51 (sum over ranks)	0.0030

An FP64 emulation of the p2p tile schedule shows the same pattern. With the sink on every tile, the ratio of the CP sum to the reference ranges from 0.5× to 2× per head at cp=2, and from 0.25× to 4× at cp=4. With the sink on step 0 only, it matches to 1e-15 at both sizes.

Not covered

The test changes enable more CI cases than I ran. I ran BF16 bshd/sbhd for p2p and BF16 bshd for a2a, on sm_90 with cuDNN 9.24. These newly enabled or changed cases were not run here:

• p2p with qkv_format="thd": main requires cuDNN ≥ 9.26 for a ragged layout with a sink in the backward.
• p2p and a2a with FP16.
• a2a sbhd/thd and cp_4_3 (sliding window, learnable) with the new non-zero offsets.
• FP8 a2a with a sink, which needs sm100 and cuDNN ≥ 9.21.
• cp_size > 2 on GPUs (only the emulation covers cp=4), and the flash-attention backend.

cp_comm_type="a2a+p2p" is still rejected for non-vanilla softmax.

Question for reviewers

Both a2a and p2p now return the full-sequence sink gradient on every CP rank. Other parameters get only their own tokens' gradient on each CP rank. Megatron-LM scales the loss by the CP size and scales dense gradients by 1 / dp_cp_group.size() before the DP×CP reduction. As far as I can tell, it treats softmax_offset like any other parameter. In that setup, the sink gradient would come out cp_size times larger than the weights' gradients. Is the full gradient on every rank the intended contract, or should both paths return the per-rank partial sum? I kept the existing a2a contract here so that the two paths agree.

[greptile-apps]

Confidence Score: 4/5

[Medium impact] The change reaches P2P context-parallel attention users who choose non-vanilla softmax.

The PR still needs the existing bias-gradient shape bug fixed before merging.

Findings

1.  Offset Corrupts Bias Gradient Shape ▶

Summary

Adds learnable-softmax support to fused P2P attention.

• P2P context-parallel attention now applies the learnable sink.

Diagram

%%{init: {'theme': 'neutral'}}%%
flowchart LR
  A[Local tile with sink] --> C[Merge tile outputs and statistics]
  B[Other tiles without sink] --> C
  C --> D[Backward with merged output and statistics]
  D --> E[Local tile returns sink gradient]
  E --> F[Sum sink gradients across CP ranks]

Loading

_{Reviews (4) · Last reviewed commit: "Test learnable and off-by-one softmax wi..." · Reviewed by Greptile}

[greptile-apps]

Offset Corrupts Bias Gradient Shape

With post_scale_bias and learnable softmax, P2P reconstructs the backward auxiliary pack as [stats, rng, bias, offset]. The fused backward binding allocates dBias using the shape of the last auxiliary tensor, which is now the [1, h, 1, 1] offset rather than the tile bias. This configuration cannot return a correctly shaped bias gradient and may fail backward shape validation.

[0z5a]

I pushed a correction to this PR: 462355e ("Count the attention sink once per query row in CP p2p") and aa5a177 (tests).

The earlier version was wrong. cuDNN counts the sink as one more softmax column, so it is included in every p2p tile's softmax stats. Passing the offset to every tile made the LSE merge count the sink once per tile that a query row visits. I also claimed that every rank already holds the full sink gradient. That is not true for p2p, where each rank computes only its own query rows. Without a mask at cp=2, the two errors cancel, which is what my original check measured. Under a causal mask they don't.

The sink now goes only to the local KV tile (step 0), and the sink gradient is all-reduced over the CP group to match a2a. The CP test now covers p2p for learnable and off-by-one softmax. It also fixes an aliased reference that made the sink-gradient check compare the CP result with itself. I updated the description with the cross-node GH200 results, and the updated test fails on the previous head.

There is one open question at the end of the description, about whether every CP rank should return the full sink gradient (the current a2a contract) or only its partial sum.