[PyT] Add per-sequence attention policies to packed THD attention

[desh2608]

Description

Add per-sequence attention policies to packed THD attention.

A scalar attn_mask_type, window_size, and bottom_right_diagonal apply one policy to an entire packed invocation. Mixed offline/causal training instead needs individual packed sequences to use different policies while preserving a single caller-facing API.

The new thd_attention_policies API accepts an ordered list of complete policies:

policies = [
    {
        "sequence_ids": offline_sequence_ids,
        "mask_type": "padding",
        "window_size": (-1, -1),
        "bottom_right_diagonal": True,  # optional
    },
    {
        "sequence_ids": causal_sequence_ids,
        "mask_type": "padding_causal",
        "window_size": (128, 0),
    },
]

output = transformer_layer(
    hidden_states,
    cu_seqlens_q=cu_seqlens,
    cu_seqlens_kv=cu_seqlens,
    max_seqlen_q=max_seqlen,
    max_seqlen_kv=max_seqlen,
    thd_attention_policies=policies,
    thd_attention_policy_dispatch="auto",
)

Every policy must contain sequence_ids, mask_type, and window_size. bottom_right_diagonal is optional and falls back to the forward-level or module-level setting. Mask types with explicit causal alignment retain their existing semantics: padding_causal forces top-left alignment and padding_causal_bottom_right forces bottom-right alignment.

Sequence IDs are one-dimensional, strictly ascending CUDA int32 or int64 tensors. The policy tensors must be disjoint and together cover every packed sequence exactly once. Repeated mask types are allowed, so separate sequence subsets can use the same mask with different windows or diagonal alignment. Policy IDs are validated once before CUDA graph capture and then treated as immutable.

The mixed-policy argument is mutually exclusive with scalar attn_mask_type and window_size arguments.

API accessibility

Both mixed-policy options are exposed at every relevant PyTorch layer:

TransformerLayer.forward
  -> MultiheadAttention.forward
    -> DotProductAttention.forward

• thd_attention_policies
• thd_attention_policy_dispatch

Callers using TransformerLayer or MultiheadAttention therefore have the same automatic/grouped control as direct DPA callers.

Runtime dispatch

thd_attention_policy_dispatch="auto" is intentionally padded-first. For each nonempty policy, Transformer Engine queries its ordinary backend selector with the complete runtime AttentionParams and pad_between_seqs=True:

• If an eligible backend supports that representation, the policy retains the original packed Q/K/V storage and represents sequences owned by other policies as inter-sequence padding.
• Otherwise, the policy's logical Q/K/V tokens are compacted before an ordinary scalar-mask THD attention call and its output is restored to the original physical layout.
• Different policies in one invocation may select different representations; their disjoint full-layout outputs are combined.
• A uniform policy uses the existing scalar fast path in automatic mode.

thd_attention_policy_dispatch="grouped" explicitly compacts every policy and bypasses padded-backend probing. The override also applies to a uniform one-policy batch; the scalar fast path does not bypass the requested dispatch.

Inter-sequence-padding representation

The padding path:

• Keeps Q/K/V in their original packed layout.
• Derives logical and physical cumulative sequence lengths for each policy.
• Calls the existing scalar-mask DPA path with pad_between_seqs=True.
• Masks inactive forward lanes before combining policy outputs.
• Uses an identity-forward autograd guard to zero inactive Q/K/V gradient lanes when a backend leaves padding lanes undefined.
• Avoids policy-dependent host synchronization and supports CUDA graph capture.

Grouped-compaction representation

The grouped path:

• Compacts the logical Q/K/V tokens for each policy.
• Invokes the existing scalar THD attention path with pad_between_seqs=False.
• Restores each policy's output to the original physical THD layout.
• Supports unequal Q/KV lengths, pre-existing physical padding, and per-policy windows and diagonal alignment.

Backend-selection caching

Mixed policies alternate among a small set of AttentionParams. Their complete backend-selection results are stored in a bounded cache and reused by the recursive scalar attention calls. The cache accounts for backend-control environment state and is invalidated when ordinary DPA backend selection is marked for refresh.

In a controlled same-node H100 microbenchmark, the steady-state selector calls per mixed forward changed from 4 to 0 for automatic dispatch and from 2 to 0 for grouped dispatch. The directly measured selector CPU time removed was 242.0 us/forward and 169.6 us/forward, respectively. Total host submission time decreased by 170.6 us (8.1%) for automatic dispatch and 109.5 us (5.3%) for grouped dispatch. These are DPA microbenchmark results, not full-model throughput estimates.

Supported scope

The new API supports:

• Packed THD self-attention and cross-attention.
• Direct DPA, MultiheadAttention, and TransformerLayer callers.
• FP16 and BF16 forward and backward.
• Mixed and uniform policies, including repeated mask types.
• Padding, causal-padding, and bottom-right causal-padding masks.
• Per-policy sliding windows and diagonal alignment.
• Different Q and KV sequence lengths.
• Existing physical padding between sequences.
• Automatic padded-first routing and an explicit grouped override.
• Empty packed batches with connected zero gradients.
• CUDA graph capture on the supported inter-sequence-padding path.

It is currently restricted to padding mask types and does not support:

• Non-THD layouts.
• Explicit attention masks.
• Attention bias or ALiBi.
• KV caching.
• Score modification.
• FP8 attention or max-logit output.
• Context Parallelism.

Type of change

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

Validation

Focused tests cover:

• FP16/BF16 output and Q/K/V gradient parity against independent scalar-mask attention.
• Mixed, uniform, causal, full-attention, and bottom-right policies.
• Repeated mask types with distinct windows and policy ordering.
• Per-policy diagonal alignment with unequal Q/KV lengths.
• Existing inter-sequence padding.
• Cross-attention with unequal Q/KV lengths.
• Per-policy backend selection and combined padded/grouped outputs.
• Automatic and grouped dispatch parity.
• Grouped override accessibility through DPA, MHA, and TransformerLayer.
• Uniform-policy grouped semantics.
• Backend-selection cache reuse and backend-control invalidation.
• Invalid policy schema, sequence coverage, and scalar/mixed conflicts.
• Empty batches.
• CUDA graph capture and replay on Hopper/FA3.

Platform	Focused tests
H100 / SM90	34 passed
A100 / SM80	33 passed, 1 expected skip

The A100 skip is the Hopper/FA3-only CUDA graph case. Python compilation, the repository's Black configuration, and whitespace checks pass for all changed files.

End-to-end 4B ASR profile

A dedicated one-node/eight-GPU profile ran the 4B/A770M model on ASR-only data for 100 optimizer steps per mode. All-causal, all-full-attention, and 50%-causal modes ran sequentially in the same exclusive allocation on each architecture. The table reports the steady-state mean over iterations 7-100 (n=94).

Hardware	All full attention	All causal	50% causal	50% causal vs full
H100, 8 GPUs	1,585.954 ms	1,589.284 ms	1,772.807 ms	+11.782%
A100, 8 GPUs	1,860.790 ms	1,877.415 ms	2,110.247 ms	+13.406%

All-causal and all-full attention are within 1% end to end on both systems. The heterogeneous path is slower because policy grouping and multiple backend invocations are paid in each audio layer. Since A100 and H100 used different cluster-local containers, the cross-architecture ratio should not be interpreted as a pure silicon comparison.

Checklist

• I have read and followed the contributing guidelines.
• The functionality is complete.
• I have commented the hard-to-understand areas.
• I have updated the public API documentation.
• My changes generate no new warnings in the focused validation.
• I have added tests covering the new behavior.
• The full repository unit-test suite was run locally.

The focused GPU suites, backend-selection tests, end-to-end 4B profiles, Python compilation, formatting, and whitespace checks pass as described above.

[greptile-apps]

Mixed THD tests bypass CI

The PyTorch QA job explicitly enumerates attention test modules and does not include this new file, so its forward, gradient, validation, and empty-batch coverage never runs in CI and regressions in the new dispatch path go undetected.

Knowledge Base Used: Tests and QA

[greptile-apps]

Greptile Summary

The PR adds per-sequence attention policies for packed THD attention and exposes policy dispatch through DotProductAttention, MultiheadAttention, and TransformerLayer.

• Supports automatic padded-first routing and explicit grouped compaction.
• Adds policy validation, backend-selection caching, gradient masking, and empty-batch handling.
• Adds the mixed-THD attention test module to the PyTorch QA test runner, resolving the previous CI-coverage finding.

Confidence Score: 5/5

The PR appears safe to merge because no blocking failure remains.

No blocking failure remains.

Important Files Changed

Filename	Overview
transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py	Implements mixed-policy validation, backend routing, padded and grouped execution, caching, and output composition without an eligible blocking follow-up issue.
transformer_engine/pytorch/attention/dot_product_attention/utils.py	Extends attention parameter handling to support policy-specific backend selection.
transformer_engine/pytorch/attention/multi_head_attention.py	Exposes and forwards mixed THD policy options through MultiheadAttention.
transformer_engine/pytorch/transformer.py	Exposes and forwards mixed THD policy options through TransformerLayer.
tests/pytorch/attention/test_mixed_thd_attention.py	Adds focused forward, backward, validation, dispatch, padding, cross-attention, empty-batch, and CUDA-graph coverage.
qa/L0_pytorch_unittest/test.sh	Adds the mixed THD attention test module to the enumerated PyTorch QA suite, fixing the prior CI omission.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
  TL[TransformerLayer] --> MHA[MultiheadAttention]
  MHA --> DPA[DotProductAttention]
  DPA --> V[Validate per-sequence policies]
  V --> D{Dispatch mode}
  D -->|grouped| G[Compact policy tokens]
  D -->|auto| B{Backend supports padding representation?}
  B -->|yes| P[Preserve packed storage and mask inactive lanes]
  B -->|no| G
  P --> C[Combine policy outputs]
  G --> C

Loading

_{Reviews (18): Last reviewed commit: "Merge branch 'main' into desh/mixed-thd-..." | Re-trigger Greptile}

[desh2608]

@cyanguwa Thanks for your detailed suggestions! I followed the advice for the new implementation in a88c9b6. However, it seems that pad_between_seqs is not supported in FlashAttention 2?

[desh2608]

@cyanguwa @KshitijLakhani I made the changes following @cyanguwa's suggestions. The main issue is that the pad_between_seqs option is only supported for FA3, so I have kept the previous "grouping" implementation for FA2. This design follows some existing arch-based implementations in the repo. Here is a pictorial summary along with some benchmark numbers. Overall, the two implementations achieve loss parity and roughly similar iteration times when used in a real training setup.

[cyanguwa]

Thanks for the pictorial summary :) @KshitijLakhani will help guide you through the rest of the PR.

Regarding the benchmarks, should we not expect more savings from pad_between_seqs in C, given that FA3 is a bit faster than FA2 and then the pad_between_seqs approach should also provide some benefits over grouped? Also, you might want to benchmark the FusedAttention backend by the way. It uses cuDNN and supports THD on sm90+, and pad_between_seqs=True.

[desh2608]

Also, you might want to benchmark the FusedAttention backend by the way. It uses cuDNN and supports THD on sm90+, and pad_between_seqs=True.

I previously ran this comparison on H100: grouped averaged 1802.289 ms and pad-between averaged 1779.596 ms over ~100 train steps.

[desh2608]

Ran some benchmarking again with the latest changes:

[KshitijLakhani]

General comments:

1. The new argument/feature currently appears to be exposed directly through DotProductAttention; I do not see corresponding plumbing through MultiheadAttention or TransformerLayer. Would be good if you could add that too - I'm guessing we wouldn't necessarily want to restrict users of this feature to the DPA API only. Ideally, would be good if you could just plumb it in this PR, however, if you are in a rush, I am happy to have this addressed in a separate PR. Please add a TODO so that we do not forget, in case you plan to address it separately

2. I might have missed this but if this PR does not explicitly reject CP for this policy based approach, we should. Users should no try to use it with CP. A very small / lightweight CP rejection test would help add confidence to this rejection code.

[KshitijLakhani]

bottom_right_diagonal is stored internally per policy, but looks like it is derived from the global argument and mask type and ussers cannot specify it independently per policy. Is my understanding right ?
Now, this may be intentional, particularly because the causal mask names encode much of the alignment behavior. Still, since mask and sliding window are now policy-specific, could you please confirm whether bottom-right alignment can ever legitimately vary between policies. If so, it belongs in the policy object too.

[desh2608]

For our use cases right now, we don't need policy-specific variation in bottom_right_diagonal, but I can imagine in some cross-attention/local attention use cases it may be useful. I will include it as per-policy field in thd_attention_policies.

[desh2608]

Added in d46d821.

[KshitijLakhani]

Choose a more generic name like : thd_attention_policies or earlier suggested attn_policies_per_seq so that it is not perceived as a policy for swa and attn masks only

[desh2608]

Changed in 156069f

[KshitijLakhani]

Are we validating/verifying these mentioned conditions in the comments anywhere in code?:

• Strictly ascending IDs within a policy.
• Duplicates within one policy.

[KshitijLakhani]

Since policy metadata is expected to be static, could we validate the sequence IDs once before CUDA graph capture and then treat them as immutable?
I'm thinking that for each policy, we should check that IDs are strictly ascending, and the sorted concatenation of all policy IDs should equal arange(batch_size). Now the issues is that these checks will synchronize when their result is converted to a Python boolean IIUC, but that should be acceptable as a one-time preparation/warm-up cost. They should not execute inside the captured forward path. A small validation step (or cached validation before capture though need to think a bit through this one's detils) may be the cleanest approach.

[desh2608]

Added in 2dca008

[KshitijLakhani]

The routing currently treats padded-backend availability as a preference: if FA3 or Fused Attention can run the policy with pad_between_seqs=True, padded execution is always selected. This makes sense and I agree with it as an initial policy, but get_attention_backend() only determines backend eligibility for that representation; it does not compare its end-to-end performance against compacting and using FA2/FA4.
It is more of a what is supported call rather than which is more performant call.

So, I'm wondering, to give the users of this feature better guidance, could you please document that mixed THD dispatch is intentionally padded-first .
Users can however, indirectly force grouped execution by disabling FA3 and Fused Attention while leaving FA2/FA4 enabled, but those environment variables affect all attention calls. This is documentation/comment should give enough guidance to the user to maybe experiment the different routes padded and grouped, (assuming they have them available) and then if they find that our approach of choosing padded is less performant for their case they can force TE to use grouped via the env vars.

Note for the future: I have not done any benchamarking on this but @desh2608 if you do any benchmarking and we learn that grouped remains materially faster for important workloads, it may be worth providing a direct padded/grouped preference or override rather than relying solely on global env var backend controls.

[desh2608]

I had benchmarked both in my training setup, and the iteration times are as follows:

Method	Attn. backend	Iteration time (ms)
Grouped	FA2	1669.9
Grouped	FA3	1548.2
Padded	FA3	1635.7

Given that grouped + FA3 seems faster, I will add an optional override so users can select this implementation.

[desh2608]

Added in 9f01ef8 along with tests.

[KshitijLakhani]

So seems like for the attention params, we have an uncache probe here, which is unlike what we do otherwise in DPA (see lines 2535 in this file)
It looks something like this:

                if (
                    _attention_backends["attention_params"] is None
                    or attention_params != _attention_backends["attention_params"]
                ):

_partition_thd_mask_policies() calls get_attention_backend() directly for every policy at line 1359, bypassing the normal _attention_backends cache used by the usual forward path. The subsequent recursive self.forward() can then perform backend selection again when it executes that policy. Was this repeated selection intentional @desh2608 ?
I'm concerned this maybe fine initially, but could you measure the CPU overhead for the latest implementation ? If it is substantial/materical engouth , perhaps the policy probe results could be cached by AttentionParams or reused during execution ? Please evaluate this

[desh2608]

No this was an oversight. Fixed in 0ced9fc. The CPU overhead was ~200us per call.

[KshitijLakhani]

Could the new tests reuse more of the existing attention test infrastructur? For e.g. DPA/config construction, THD input generation, backend forcing and comparison helpers?

I am less conflicted about whether the tests remain in a dedicated file or we find an easy/clean way to integrate into the test_attention file instead, however, since the feature now has enough cases to justify a separate file, I'm happy to consider that possibility. My main maintenance concern is avoiding duplicated
setup/reference machinery. Trying to focus on mainainability.

While you perform this restructure, please weight the possibility of just having this as a separate function/class in the fused attention test file itself if it relatively clean to do so, else stick wih a separate file and only address the re-use comment i have bove

[desh2608]

Updated to reuse the existing infrastructure in 9f01ef8. I kept the new tests in a separate file since there are many cases.

[greptile-apps]

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

[desh2608]

2. I might have missed this but if this PR does not explicitly reject CP for this policy based approach, we should. Users should no try to use it with CP. A very small / lightweight CP rejection test would help add confidence to this rejection code.

Added in 2dca008

[desh2608]

1. The new argument/feature currently appears to be exposed directly through DotProductAttention; I do not see corresponding plumbing through MultiheadAttention or TransformerLayer. Would be good if you could add that too - I'm guessing we wouldn't necessarily want to restrict users of this feature to the DPA API only. Ideally, would be good if you could just plumb it in this PR, however, if you are in a rush, I am happy to have this addressed in a separate PR. Please add a TODO so that we do not forget, in case you plan to address it separately

Added plumbing in 156069f

[desh2608]

I extracted these into public methods so they can be re-used in test_mixed_thd_attention.py.

[desh2608]

We are now re-using the same setup as in test_attention.py. This file only contains cases specific to mixed attention.

[KshitijLakhani]

/te-ci L0 L1 L2

[desh2608]

@KshitijLakhani FYI one of the tests was failing on H100 due to a backend selection bug which is now fixed in 98f375f.

[KshitijLakhani]

This fix follows the existing convention of using max_seqlen_q != max_seqlen_kv for cross-attention-like cases. For variable-length THD, can individual Q/KV lengths differ while their global maxima are equal? @cyanguwa
If so, the selector cannot distinguish that case today. I believe that is a broader existing limitation rather than something this PR needs to solve, but it may be worth a follow-up test or issue.

Not blocking the PR on this for now

[KshitijLakhani]

/te-ci L0 L1 L2

[KshitijLakhani]

LGTM !
Thanks for implementing this feature, Desh and taking it to the finish line :)