Attribute and retry transient CP pool NaNs

[sudhakarsingh27]

Description

Long-lived context-parallel workers can occasionally retain state across heterogeneous FP8 THD cases. A fresh pool can clear these transient assertion failures, while a failure repeated in the fresh pool must remain visible.

This change also makes the retry lifecycle explicit so recovered attempts are distinguishable from terminal test failures.

Type of change

• Bug fix
• Documentation change
• New feature
• Breaking change
• Infra/Build change
• Code refactoring

Changes

• Keep documented retryable assertion signatures in one allowlist.
• Retry allowlisted FP8 THD NaN assertions once in a fresh pool.
• Include the active pytest node in every retry record.
• Emit explicit recovered and terminal-failure records.

Validation

• The test change completed two full default CP attention suite runs on B200: 320 passed, 12,032 skipped, 0 failed in each run.
• After rebasing onto current main, the retry control-flow harness covered transient recovery, persistent failure, the existing NaN signature, unrelated NaN rejection, and non-THD rejection.
• The rebased result passes Python compilation and the project-configured Black formatting check.

Checklist

• I have read and followed the contributing guidelines.
• The functionality is complete.
• I have commented the non-obvious retry policy and attribution choice.
• Documentation changes are not required for this test-only change.
• My changes generate no new warnings.
• New and existing test behavior has been validated.

[greptile-apps]

Greptile Summary

This PR extends the context-parallel attention test pool’s narrowly scoped retry policy to recognize two FP8 THD NaN signatures and makes retry outcomes attributable in CI logs.

• Centralizes retryable NaN assertion signatures in an allowlist.
• Retries matching FP8 THD failures once after recreating the worker pool.
• Adds pytest node attribution and explicit retrying, recovered, and terminal-failure records.

Confidence Score: 5/5

The PR appears safe to merge because matching failures are retried in a newly spawned pool and a repeated failure is still propagated.

The retry gate is limited to the intended FP8 THD cases and documented signatures, successful retries are explicitly recorded, and all terminal second-attempt errors remain test failures.

Important Files Changed

Filename	Overview
tests/pytorch/attention/test_attention_with_cp.py	The retry classification and diagnostics are consistently scoped to FP8 THD assertions, while persistent second-attempt failures remain visible.

_{Reviews (1): Last reviewed commit: "Attribute and retry transient CP pool Na..." | Re-trigger Greptile}

[sudhakarsingh27]

/te-ci pytorch L1

[fheinecke]

/te-ci pytorch L1

[sudhakarsingh27]

/te-ci pytorch L1

[fheinecke]

/te-ci pytorch L1

[fheinecke]

LGTM pending CI