test(attention): anchor the context-parallel suite to an independent reference

[nvegesna-netizen]

Problem

tests/pytorch/attention/run_attention_with_cp.py grades a context-parallel run against a non-CP
run of the same backend (tensors_no_cp vs tensors_cp). That validates the context-parallel
plumbing and nothing about the kernel beneath it: any error affecting both sides equally cancels,
and the comparison passes.

This is not hypothetical. Measured on B200 against the FA4 defect in #3528, where the backend
returns an all-zero output for causal attention:

cp_comm_type	before this change	after
all_gather	PASS	FAIL — non-CP output is all zeros
a2a	PASS	FAIL — non-CP output is all zeros
p2p	FAIL	FAIL

Two of three communication types reported PASS on a backend returning nothing at all. Zeros on one
side, zeros on the other, agreement. Only p2p failed, and only incidentally: its ring merge turned
the accompanying -inf LSE into a NaN that the zeros on the other side did not match.

The same structure also inverts the sign of a partial fix. An earlier attempt at #3528 corrected the
CP side only, which made the two sides disagree — so the fix was reported as a regression in
all_gather and a2a, while the fully broken state had been reported as a pass.

Change

Check one (batch, head) slice of the non-CP output against an independently computed float64
reference, before the CP/no-CP comparison, so a systematically wrong kernel cannot cancel itself out.

It is a sanity bound, not a precision test — the existing CP/no-CP comparison already covers
precision. It targets catastrophic wrongness: zeros, NaNs, a mask applied in the wrong place. An
all-zero output is asserted separately from the tolerance so it is named rather than surfacing as a
large numeric error.

float64 rather than float32: torch computes fp32 matmuls in TF32 on Ampere and newer, whose
significand is 11 bits — the same as fp16 — so an fp32 reference cannot judge a bf16 kernel.

Configurations the reference does not model are skipped rather than approximated: fp8, attention
bias, non-vanilla softmax, thd, padding masks, and softcap. Modelling softcap would mean
duplicating its tanh rescaling in the reference, which is how a reference starts reimplementing
the thing it is meant to check independently.

Verification (B200)

The negative control is the real defect rather than a synthetic one: this branch is cut from main,
which still carries the FA4 zero-output bug, so the check can be tested against a backend that is
genuinely returning nothing.

check	result
bug present on this tree	out.abs().max() = 0.0 confirmed
anchor + broken backend	all_gather, a2a, p2p all fail, naming the all-zero output
anchor + correct backend (FA4 off)	all_gather, a2a, p2p all pass
full FlashAttention CP suite, correct backend	18 passed, 0 failed

The third row matters as much as the second. A reference that rejects correct kernels would be worse
than the gap it closes, and it is the more likely way to get this wrong — the mask convention has to
match at bottom-right alignment, sliding windows and GQA. An earlier revision failed three softcap
cases for exactly that reason; those are now skipped.

Note

#3532 fixes the FA4 defect used as the negative control here. Once it lands, that particular
backend will no longer return zeros — but the gap this closes is structural and independent of it:
the suite cannot detect any error affecting CP and non-CP equally.

[nvegesna-netizen]

One note on merge order, since these two are related but independent.

Please land #3532 before this one. #3532 fixes the FA4 defect that this change uses as its
negative control. If this merges while that bug is still on main, any lane enabling FA4 on SM100
starts failing — correctly, since the backend really is returning an all-zero output, but it would
present as this change breaking CI rather than as the pre-existing bug it is surfacing. Today's
lanes export NVTE_FLASH_ATTN_V4=0 so it would not fire in practice, but that is a thin thing to
rely on.

Kept separate from #3532 deliberately rather than folded in:

• that one is a correctness fix to shipping code; this one modifies a harness every context-parallel
  test runs through, which is a different kind of review
• this adds a float64 computation per test plus skip conditions and a tolerance, all reasonably
  debatable — and a silent wrong-answer fix should not wait on that debate
• it is not FA4-specific: it protects every backend under context parallelism, and filing it under an
  FA4 fix would undersell it

No dependency between them otherwise — different files, cut from main, no conflict either way.

[greptile-apps]

Confidence Score: 5/5

The PR appears safe to merge, with no concrete correctness, security, or repository-rule failures identified.

Summary

This PR anchors context-parallel attention tests to an independently computed float64 reference so systematic backend failures cannot pass merely because CP and non-CP executions fail identically.

• Checks one batch/head slice to limit reference-computation memory.
• Models supported vanilla, bias-free bshd/sbhd attention, including causal and sliding-window masks.
• Skips configurations the independent reference does not model.
• Explicitly diagnoses non-finite and all-zero backend output before the existing CP/no-CP comparison.

Diagram

%%{init: {'theme': 'neutral'}}%%
flowchart TD
  A[Run non-CP attention backend] --> B{Configuration modeled by reference?}
  B -- No --> E[Continue to existing CP comparison]
  B -- Yes --> C[Compute one-head float64 reference]
  C --> D{Finite, nonzero, and within sanity bound?}
  D -- No --> F[Fail with targeted backend diagnostic]
  D -- Yes --> E
  E --> G[Run context-parallel attention]
  G --> H[Compare CP output with non-CP output]

Loading

_{Reviews (1) · Last reviewed commit: "test(attention): skip softcap configs in..."}