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}

[cyanguwa]

Could we make this _pad_qkv_head_dim work for all three executions below: if use_flash/fused/unfused_attention:, so we only have one pair of _pad_qkv_head_dim/_trim_output for all three backends instead of duplicating the code multiple times.

[cyanguwa]

Can we make the comments throughout the PR a bit more concise, with possibly just one or two sentences? Thanks.

Also, please make it clear that we're extending the head dims to the max(head_dim_qk, head_dim_v) and not to the next supportable head_dim, because the max could still be unsupported by Flash/FusedAttention possibly.

[cyanguwa]

Apart from your use case, we're also trying to enable FA2 which does not support head_dim_qk != head_dim_v, so please include that logic when deciding if padding or not.

[cyanguwa]

I feel the test file is a little overly verbose - some of the functionalities are pretty obvious, for example, with the caching mechanism. Do you think we can fold/reduce the testing into one test that calls test_dot_product_attention, like with the other tests, and add it to test_attention.py? Essentially, what we want to see is that for a given ModelConfig such as head_dim_qk=96 and head_dim_v=128, it can cleverly pad and take advantage of one of the faster backends? Thanks!