Add two-section causal backward for shared K/V

[foraxe]

Summary

Add an experimental CuTe SM100/SM110 backward path for two causal query
sections that share one physical K/V owner.

_flash_attn_bwd_two_section_causal accepts two adjacent batch-one Q/output/
dO/LSE views, two causal K/V prefix lengths, and one physical K/V tensor. It
returns dQ for both logical query sections and compact dK/dV in the physical
K/V owner's layout. This avoids materializing two independent K/V-gradient
buffers before a context-parallel owner-return stage.

The existing forward and backward APIs are unchanged. Unsupported inputs fail
closed.

Supported contract

• SM100 or SM110
• causal attention
• exactly two equal-length Q sections
• one shared fixed-storage K/V tensor
• MHA with head dimension 128
• 128x128 tiles and the 2-CTA path
• aligned K/V prefixes with 0 < prefix1 <= prefix2

Changes

• add the two-section backward implementation in flash_bwd_sm100.py;
• expose a narrow internal interface for the operation;
• add a focused BF16 correctness test that skips on unsupported GPUs.

Validation

The focused test passed on one GB200 using:

Q sections: [1, 1024, 20, 128] x 2
K/V owner:  [1, 2048, 20, 128]
prefixes:   1024, 2048
reference:  two stock backward calls plus explicit dK/dV addition
tolerance:  atol=0.02, rtol=0.02

The reduced wheel was also validated in an opt-in Megatron CP=4 training stack
at sequence length 262144. After eight updates, all 73 checkpoint tensors
matched the TE p2p reference with worst max-absolute difference
1.9073486328125e-6 (atol=2e-6, rtol=0). Median iteration time over
updates 4-7 was 1707.45 ms for TE p2p and 1413.65 ms for the NVSHMEM candidate,
17.21% lower for that single clean comparison.

This PR does not claim a standalone kernel speedup: at the focused small shape,
one two-section launch did not beat two stock backward launches. Its intended
benefit is avoiding intermediate K/V-gradient materialization in the integrated
context-parallel path.

Stack relationship

This is a dependency of an experimental Megatron NVSHMEM context-parallel
backend:

• NVIDIA/TransformerEngine#3235
  provides stream-ordered CP owner-gradient return;
• NVIDIA/Megatron-LM#5969
  owns the opt-in backend, symmetric workspaces, and lifecycle.

Rebase status

The branch now merges cleanly with current main after a semantic rebase over
the SM100 block-sparse 2-CTA changes from #2661 and the varlen score-mod
backward changes from #2547.

CPU-only CuTe fake compilation passed for:

• the two-section SM100 backward target;
• stock dense SM100 backward at head dimensions 128 and 192;
• SM100 2-CTA block-sparse backward at head dimensions 128 and 192.

The PR remains draft until the focused correctness test and the integrated
Megatron stack are rerun on a real GB200.

[foraxe]

GB200 validation update for current head 5b9d334:

• Fixed a semantic-rebase regression in the SM100 block-sparse path by restoring the upstream None initialization contract for optional load_dOt, load_Qt, and load_Kt loaders.
• Focused two-section causal backward: 1 passed (BF16, two 1K query sections sharing 2K K/V, head dim 128).
• Block-sparse dispatch/normalization coverage: 7 passed.
• Stock SM100 dense causal backward on the pre-fix rebased tree: 2 passed (BF16 head dims 128 and 192); the three-line sentinel fix does not change those enabled-loader paths.
• Integrated four-GB200 CP=4 training gate at sequence length 262144: 73/73 checkpoint tensors matched, worst max-absolute difference 1.9073486328125e-06 at tolerance 2e-06.

The first upstream block-sparse actual-kernel case found the missing-loader bug immediately on the prior head. After the fix, both a 1CTA case and a 2CTA head-dim-192 case advanced past that error but did not complete CUTLASS/Flex compilation within 15-minute and 10-minute caps, respectively. I am keeping the PR in draft until an actual-kernel block-sparse case completes with a correctness result; the completed tests above should not be read as covering that unresolved gate.