Fix THD P2P pad detection and tail-zero gating

[RPalmr]

Fixes #3487.

Signed-off-by: Reece Palmer rxp528@student.bham.ac.uk

[greptile-apps]

Confidence Score: 5/5

The PR appears safe to merge; no actionable new correctness, security, or repository-rule issue remains.

Summary

• Detects inter-sequence padding by comparing cumulative-length values during eager execution and caches the result for graph capture.
• Includes tensor mutation versions in padding-cache keys.
• Restricts context-parallel inter-sequence gradient cleanup to inputs that actually contain such padding.
• Uses direct eager slices while retaining capture-safe masks for THD gradient tail zeroing.

Diagram

%%{init: {'theme': 'neutral'}}%%
flowchart TD
    A[THD attention input] --> B{pad_between_seqs explicitly set?}
    B -- Yes --> F[Use explicit value]
    B -- No --> C{CUDA graph capture?}
    C -- No --> D[Compare non-tail cumulative lengths]
    D --> E[Cache value-sensitive decision]
    C -- Yes --> G{Matching eager cache entry?}
    G -- Yes --> H[Reuse eager decision]
    G -- No --> I[Use capture-safe conservative inference]
    E --> J[Backend selection]
    H --> J
    I --> J
    F --> J
    J --> K[Backward gradients]
    K --> L[Zero THD tail padding]
    K --> M{Inter-sequence padding?}
    M -- Yes --> N[Zero inter-sequence padding]
    M -- No --> O[Skip inter-sequence cleanup]

Loading

_{Reviews (4) · Last reviewed commit: "Merge branch 'main' into fix/thd-p2p-pad..."}