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 dispatch for packed THD self-attention.

• Validates the restricted FP16/BF16 feature configuration.
• Groups packed sequences by causal policy, invokes existing scalar attention paths, and restores token order.
• Adds numerical, gradient, validation, and empty-batch tests, though the new test module is not wired into QA.

Confidence Score: 4/5

The implementation appears safe to merge, with the non-blocking exception that its new regression tests are not currently executed by repository CI.

The mixed-policy dispatcher validates unsupported configurations and preserves the existing scalar attention path, while the only accepted concern is that the newly added test module is absent from the explicitly enumerated PyTorch QA suite.

Files Needing Attention: tests/pytorch/attention/test_mixed_thd_attention.py and qa/L0_pytorch_unittest/test.sh

Important Files Changed

Filename	Overview
transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py	Adds constrained mixed-policy THD validation, grouping, scalar dispatch, and output reordering without an accepted functional defect.
tests/pytorch/attention/test_mixed_thd_attention.py	Adds focused numerical and validation coverage, but the standalone test module is absent from the explicit QA invocation list.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
    A[Packed THD Q/K/V and per-sequence policy] --> B{All policies equal?}
    B -->|Yes| C[Invoke scalar-policy attention once]
    B -->|No| D[Group token indices by sequence 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 attention output]
    G --> H

Loading

_{Reviews (1): Last reviewed commit: "Add per-sequence causal policy to THD at..." | Re-trigger Greptile}