[RPalmr]

Fixes #3487.

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

[greptile-apps]

Confidence Score: 4/5

The attention fixes appear sound, but the PR is not safe to merge until the release wheel path mismatch caused by the version spelling is corrected.

Summary

This PR corrects THD inter-sequence padding detection and limits padding cleanup to cases that actually contain inter-sequence padding.

• Caches eager padding decisions for CUDA graph capture using tensor identity and mutation versions.
• Uses capture-safe masks while retaining efficient eager tail-zeroing.
• Gates context-parallel internal-padding cleanup on pad_between_seqs.
• Also changes the development version spelling, which breaks assumptions in the release wheel repackaging script.

Diagram

%%{init: {'theme': 'neutral'}}%%
flowchart TD
    A[Read raw VERSION.txt<br/>2.21.0-dev0] --> B[setuptools builds wheel]
    B --> C[Packaging normalizes version<br/>2.21.0.dev0]
    C --> D[wheel unpack creates normalized directory]
    A --> E[build_wheels.sh constructs raw hyphenated path]
    D --> F{Paths match?}
    E --> F
    F -->|No| G[METADATA edit or directory move fails]

Loading

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

[timmoon10]

I don't understand this code very well, but at a glance it looks reasonable aside from some nits.

[timmoon10]

functools.lru_cache would be simpler and less error-prone than implementing the cache from scratch. That said, we already follow that anti-pattern with _thd_policy_validation_cache, so there's something to be said about maintaining consistency.

[timmoon10]

Is it safe to assume that two tensors with the same buffer will have the same padding requirements? If cu_seqlens_* are not persistent, the same buffers may be reused in different steps or different layers. Perhaps it would be safer to do something like:

return (id(t), version)

[timmoon10]

Noting that the cache is not populated during graph capture. In order to take advantage of the cache, we assume that the cache is populated in the warmup steps.

[RPalmr]

Thanks for looking. Context: #3487 regressed THD P2P when padded cu_seqlens are value equal but separate objects. Identity only detection forced the padded path. This restores value comparison in eager runs. It reuses that answer during capture with a small id plus version cache. Warmup hits reuse it. Misses use the safe fallback. It also skips tail and padding cleanup when there is nothing to clear. Your nits are addressed in the staged update.

[greptile-apps]

Want your agent to iterate on Greptile's feedback? Try greploops.