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}