Optionally pad QK/V for better kernel selection

[janEbert]

Description

While support for native QKV shapes has been added, some MLA kernels
still benefit from having QKV padded to the same head dimension. Since
it's not easy to automate this decision, we probe
get_attention_backend twice with and without padding to figure out
which version selects the better kernel.

If padding results in a better kernel than not padding, we pad;
otherwise, Q, K, and V stay native.

Ref NVIDIA/Megatron-LM#6240, ref NVIDIA/Megatron-LM#6241.

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

Please list the changes introduced in this PR:

• Automate MLA padding-related kernel selection
• Add corresponding tests

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

[greptile-apps]

Greptile Summary

The PR adds backend probing that optionally pads unequal MLA QK/V head dimensions, trims backend outputs to the original V width, and adds L0 tests for selection and numerical behavior.

• Adds cached native-versus-padded backend selection.
• Applies optional padding across attention backend dispatch paths.
• Adds MLA padding tests to the PyTorch L0 CI suite.

Confidence Score: 4/5

The PR is not yet safe to merge because packed QKV/KV inputs can still dispatch padded fused attention using unchanged unpadded packed storage.

The attempted packed-input fix only excludes Float8TensorStorage values; ordinary packed tensors still produce padded views while the original packed buffers are consumed by fused-attention quantization.

Files Needing Attention: transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py

Important Files Changed

Filename	Overview
transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py	Adds optional padding and output trimming across backend paths, but packed fused-attention storage can remain inconsistent with the padded views.
transformer_engine/pytorch/attention/dot_product_attention/utils.py	Adds cached native-versus-padded backend probing using copied AttentionParams.
tests/pytorch/attention/test_dpa_mla_qkv_head_dim_pad.py	Adds backend-selection, shape, backward, and padding-identity coverage.
qa/L0_pytorch_unittest/test.sh	Adds the new MLA padding test module to the explicit L0 PyTorch CI suite.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
  A[Build native AttentionParams] --> B[Probe native backend]
  B --> C{Native uses unfused attention?}
  C -- No --> D[Keep native QK/V dimensions]
  C -- Yes --> E[Probe equal padded dimensions]
  E --> F{Fused or Flash available?}
  F -- No --> D
  F -- Yes --> G[Pad Q, K, and V views]
  G --> H[Dispatch selected backend]
  D --> H
  H --> I{Padding applied?}
  I -- Yes --> J[Trim output to original V width]
  I -- No --> K[Return native-width output]
  J --> K

Loading

_{Reviews (2): Last reviewed commit: "[pre-commit.ci] auto fixes from pre-comm..." | Re-trigger Greptile}