Add per-sequence causal policy to packed THD attention

[desh2608]

Description

Add an optional per-sequence causal policy to DotProductAttention for packed THD self-attention. A scalar attn_mask_type currently applies one policy to the entire packed invocation, while mixed offline/causal training needs each packed sequence to retain its own policy.

When thd_sequence_is_causal is supplied, Transformer Engine groups whole sequences by policy, invokes the existing scalar-policy attention path once per non-empty group, and restores the original packed token order. Uniform all-causal and all-offline batches retain the existing one-call path.

This is a logical dispatcher over existing attention calls, not a new heterogeneous-mask fused kernel. Initial support is intentionally limited to plain FP16/BF16 THD self-attention; options whose tensors or semantics cannot be safely regrouped are rejected explicitly.

Type of change

• Documentation change (change only to the documentation, either a fix or a new content)
• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• Infra/Build change
• Code refactoring

Changes

• Add thd_sequence_is_causal, a boolean tensor with one value per packed sequence, to DotProductAttention.forward.
• Group causal and full-context sequences without changing the surrounding packed token order or autograd connectivity.
• Preserve the single scalar-policy call for uniform batches.
• Reject unsupported packed metadata, dropout, FP8, context parallelism, sliding windows, bottom-right causal alignment, cache inputs, bias/ALiBi, score modifications, padding gaps, and attention checkpointing.
• Add focused FP16/BF16 forward and input-gradient parity tests for mixed, all-offline, and all-causal policies, plus validation and empty-batch tests.

Validation

• python -m pytest --import-mode=importlib tests/pytorch/attention/test_mixed_thd_attention.py -q -vv
  • 9 passed on an H100 with no skips
  • covers FP16/BF16 forward and Q/K/V gradient parity against scalar-policy
    attention
• Full Python Pylint scope: 10.00/10
• Black 24.4.2, python -m py_compile, and git diff --check

Checklist:

• I have read and followed the contributing guidelines
• The functionality is complete
• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• New and existing unit tests pass locally with my changes

The focused GPU suite and repository Python lint pass as described above; the full repository unit-test suite was not run locally.

[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

Adds per-sequence causal-policy dispatch for packed THD self-attention.

• Groups sequences by causal policy, runs the existing scalar-policy attention path, and restores packed token order.
• Validates and rejects unsupported tensor layouts, attention options, and execution modes.
• Adds forward, gradient, validation, and empty-batch tests to the PyTorch L0 CI suite.

Confidence Score: 5/5

The PR appears safe to merge.

No blocking failure remains.

Important Files Changed

Filename	Overview
transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py	Adds guarded mixed-policy THD dispatch by regrouping complete sequences around the existing scalar attention implementation.
tests/pytorch/attention/test_mixed_thd_attention.py	Adds FP16/BF16 forward and gradient parity, unsupported-input validation, and empty-batch coverage.
qa/L0_pytorch_unittest/test.sh	Registers the new mixed-THD attention test in the existing single-GPU PyTorch CI job, resolving the previous coverage omission.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
  A["Packed THD Q/K/V + per-sequence policy"] --> B{"Policy uniform?"}
  B -->|Yes| C["Run scalar attention once"]
  B -->|No| D["Group token indices by policy"]
  D --> E["Run full-context scalar attention"]
  D --> F["Run causal scalar attention"]
  E --> G["Restore original packed token order"]
  F --> G
  C --> H["Return output"]
  G --> H

Loading

_{Reviews (2): Last reviewed commit: "Fix mixed THD test CI coverage" | Re-trigger Greptile}

[cyanguwa]

The miscellaneous kernels generated in this approach might degrade the overall performance quite a bit. TE supports an option called pad_between_seqs=True, which allows attention backends to skip through the padded tokens and only process the "intended" tokens/sequences. In your case, I think we can perform two attention calls, one viewing the causal sequences as the "intended" sequences (and the rest "pad tokens"), and the other one viewing the non-causal sequences as the "intended" sequences. Could you please take a look at this note and this test see if it makes sense for your use case?

For example, we have a batch of 7 sequences, [111aaa222bb3344cccc], where we would like to perform "padding_causal" on the numbered sequences, {111, 222, 33, 44}, and "padding" on the lettered ones, {aaa, bb, cccc}, we can do something like this:

mask_type='padding_causal'
cu_seqlens=[0,3,6,8,10]
cu_seqlens_padded=[0,6,11,13,19]
out_causal = self.forward(
    q, k, v,  # full batch, not split
    qkv_format='thd',
    attn_mask_type=mask_type,
    cu_seqlens_q=cu_seqlens,
    cu_seqlens_kv=cu_seqlens,
    cu_seqlens_q_padded=cu_seqlens_padded,
    cu_seqlens_kv_padded=cu_seqlens_padded,
) # shape [19, h, d] with 0s on pad positions

mask_type='padding'
cu_seqlens=[0,3,5,9]
cu_seqlens_padded=[3,9,15,19]
out_non_causal = self.forward(
    q, k, v,  # full batch, not split
    qkv_format='thd',
    attn_mask_type=mask_type,
    cu_seqlens_q=cu_seqlens,
    cu_seqlens_kv=cu_seqlens,
    cu_seqlens_q_padded=cu_seqlens_padded,
    cu_seqlens_kv_padded=cu_seqlens_padded,
) # shape [19, h, d] with 0s on pad positions

out = out_causal + out_non_causal

[cyanguwa]

Nit: cacheing -> caching

But more importantly, can we rename the argument to something more generic and extensible, for example, attn_mask_type_per_seq? We can have the type as a dict, for example, users can pass in {"padding": torch.Tensor, "padding_causal": xxx, "padding_causal_bottom_right": xxx, xxx} to specify which SeqIDs need to run with "padding", and which with "padding_causal", etc. The tensor needs to be on-device.

[cyanguwa]

Please try to avoid GPU-CPU syncs (e.g. item()). I think if we pursue the approach suggested above, we can achieve CUDA graph and torch.compile compatibility.

[cyanguwa]

Some of the runtime checks can stay here, but others, regarding what's supported and what's not, can probably go to get_attention_backend. For example, the feature only supports THD, non-CP, etc. If you pursue the out=out_causal+out_non_causal approach above, the support matrix will expand naturally I think, for example, with self/cross-attention, bottom_right_diagonal=T/F. Also, pad_between_seqs=T already has some constraints in get_attention_backend which can be reused for this feature.

[cyanguwa]

The entire if thd_sequence_is_causal is not None code block can sit here, to take advantage of the existing mask/window_size checks, as well as be in the self.prepare_forward_ctx context for FP8 (if FP8 is ever available for this feature).