Optionally pad QK/V for better kernel selection by janEbert · Pull Request #3339 · NVIDIA/TransformerEngine

janEbert · 2026-08-10T22:47:58Z

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 · 2026-08-10T22:52:09Z

Confidence Score: 5/5

[Medium impact] The changes since the previous review appear safe to merge; no new blocking issue was found.

Summary

The PR compares backend choices with native and padded head widths, then pads Q/K/V only when padding improves the choice. It trims attention output back to the original value width.

Attention pads unequal query/key and value widths only when it improves backend choice.

Diagram

%%{init: {'theme': 'neutral'}}%%
flowchart TD
    A[Native Q/K/V widths] --> B[Compare native and padded backend choices]
    B --> C{Does padding improve the choice?}
    C -->|Yes| D[Pad Q/K/V]
    C -->|No| E[Keep native widths]
    D --> F[Run selected attention backend]
    E --> F
    F --> G[Trim output if padded]

_{Reviews (12) · Last reviewed commit: "Improve cache clear usability" · Reviewed by Greptile}

janEbert · 2026-08-13T23:46:33Z

PTAL.

janEbert · 2026-08-27T08:19:43Z

Hey, friendly ping.

cyanguwa

Please address the remaining couple of comments, and I'll start the CI. Thanks!

cyanguwa · 2026-09-04T16:09:15Z

+                return None
+
+            if qkv_head_pad:
+                attn_out = _trim_output(attn_out, num_attention_heads, head_dim_qk, orig_head_dim_v)


Would _trim_output(attn_out, num_attention_heads, head_dim_v, orig_head_dim_v) read a bit better, even though head_dim_qk = head_dim_v at the point?

You might also need to be careful when attn_out is a Tuple for example with return_max_logit=True. There might be some logic needed to hitting that path. Thanks.

For the former, I actually deliberately used qk to highlight that fact, but I'm fine with changing it to head_dim_v.

For the latter; thanks a lot! I didn't catch that because I was only looking at the type hints, which appear to be outdated. I'll update it as well if you don't mind. :)

Would you like tests to be added for the return_max_logit=True case as well?

I think we can leave the type hints alone in backends.py. Currently TE doesn't have any mypy, pyright, or pytype checks anywhere - maybe something we should do in the future, and there's a lot more places to change when we do. :)

I think it's ok without the specific tests for return_max_logit=True. That's just "one of" the peripheral features that could or could not be affected when combined with the padding/backend reselection. If you like, you can add this flag to one of the mla tests you added (just by adding a kwarg to one of the configs in model_configs_mla).

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. Signed-off-by: janEbert <janpabloe@nvidia.com>

Signed-off-by: janEbert <janpabloe@nvidia.com>

It could be that all `native_use_*_attention` values are False, so this needs to be handled. Signed-off-by: janEbert <janpabloe@nvidia.com>

Signed-off-by: janEbert <janpabloe@nvidia.com>

Improve variable names along the way. Signed-off-by: janEbert <janpabloe@nvidia.com>

janEbert · 2026-10-12T23:28:04Z

I'm sorry for the late reply; doesn't make it easier to keep this in your head, for sure.

Signed-off-by: janEbert <janpabloe@nvidia.com>

cyanguwa · 2026-10-13T23:19:38Z

        quantizers=None,
        fp8_output: bool = False,
-    ) -> torch.Tensor:
+    ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:


How about this to prevent it from going stale - we may have more output tensors as more features are getting added.

Union[torch.Tensor, Tuple[torch.Tensor, ...]]

I would argue it's good if it can go stale – in theory, a type checker should complain for all locations where a new n-tuple isn't correctly handled.
WDYT? Still change?

See comment above. I think we can remove the changes you made, for now.

I reverted the varlength tuple change, if that is what you meant. If it's not, then I'd still argue to keep the type hint updates in some form (even if in a new PR). Even if it's not used by any static analysis tooling, it's extremely helpful just for human code reading.

Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

janEbert · 2026-10-14T16:51:53Z

I've also been fighting with the caching. It's extremely difficult due to the various interactions. Since you seemed to have started working on it, would you like to take it over?

The problem I've found is that, on the one hand, the _attention_backends variable is modified in forward after and possibly due to the results of should_pad_qkv_head_dim. However, we also want to force a cache refresh in case the backend data changes. Importantly, _attention_backends can be modified from the outside (e.g., other places in this repo also do that) and I think we also need to observe the different _THD_POLICY_BACKEND_ENV_VARS envvars as well in terms of cache invalidation.

Also add batch size config. Signed-off-by: janEbert <janpabloe@nvidia.com>

Signed-off-by: janEbert <janpabloe@nvidia.com>

greptile-apps · 2026-10-14T17:01:54Z

Want your agent to iterate on Greptile's feedback? Try greploops.

This reverts commit 4029149. Signed-off-by: janEbert <janpabloe@nvidia.com>

Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

cyanguwa · 2026-10-14T18:47:36Z

@janEbert can you take a look at my recent changes and see if they follow your original logic? Thanks.

cyanguwa · 2026-10-14T18:48:52Z

/te-ci pytorch L0 L1 L3

Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

janEbert · 2026-10-19T13:57:28Z

I believe there is one case in which padding is activated even though it shouldn't be, but it may never be possible in the current state:

With native FlashAttention available and selected, padded FlashAttention not selected, but padded FusedAttention selected. In that case, we do pad, even though native FlashAttention would already be the best choice.

janEbert · 2026-10-20T13:48:39Z

~~/te-ci pytorch L0 L1 L3~~

Too bad, doesn't work. Guess I'd have to be a member of a team for this repo. :)

janEbert · 2026-10-21T08:22:59Z

+            if _attention_backends["backend_selection_requires_update"]:
+                dpa_utils.clear_should_pad_qkv_head_dim_cache()


I added another commit to allow passing an invalidate_cache boolean as an optional argument to should_pad_qkv_head_dim (in this case invalidate_cache=_attention_backends["backend_selection_requires_update"]). Then we can call the cache clear inside the function, which may be better in terms of user-facing visibility and API. Adjusted the test accordingly.

If you don't like the change, feel free to revert it.

Signed-off-by: janEbert <janpabloe@nvidia.com>

janEbert requested a review from cyanguwa as a code owner August 10, 2026 22:47

github-actions Bot added the community-contribution PRs from external contributor outside the core maintainers, representing community-driven work. label Aug 10, 2026

greptile-apps Bot reviewed Aug 10, 2026

View reviewed changes

Comment thread transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py

Comment thread tests/pytorch/attention/test_dpa_mla_qkv_head_dim_pad.py Outdated

Comment thread tests/pytorch/attention/test_dpa_mla_qkv_head_dim_pad.py Outdated

janEbert force-pushed the optional-mla-pad branch from 8d821bf to bdc487e Compare August 10, 2026 23:30

greptile-apps Bot reviewed Aug 10, 2026

View reviewed changes

Comment thread transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py

cyanguwa reviewed Aug 11, 2026

View reviewed changes

Comment thread transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py

cyanguwa reviewed Aug 11, 2026

View reviewed changes

Comment thread transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py Outdated

Comment thread transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py Outdated

Comment thread tests/pytorch/attention/test_dpa_mla_qkv_head_dim_pad.py Outdated

ptrendx assigned cyanguwa Aug 11, 2026

janEbert mentioned this pull request Aug 14, 2026

Let MLA skip the thd V pad on the packed path NVIDIA/Megatron-LM#6241

Open

5 of 6 tasks

cyanguwa reviewed Sep 4, 2026

View reviewed changes

cyanguwa added the attention label Sep 21, 2026

janEbert mentioned this pull request Oct 1, 2026

[BUG] Packed THD MLA value padding may affect attention backend availability on Blackwell NVIDIA/Megatron-LM#7681

Open

janEbert added 15 commits October 13, 2026 00:49

Add padding-related tests

6e5bf36

Signed-off-by: janEbert <janpabloe@nvidia.com>

Fix tests

748957e

Signed-off-by: janEbert <janpabloe@nvidia.com>

Add memoization

64c347f

Signed-off-by: janEbert <janpabloe@nvidia.com>

Add test to CI

6bdac9b

Signed-off-by: janEbert <janpabloe@nvidia.com>

Skip padding when using fused+packed FP8 path

cf0e401

Signed-off-by: janEbert <janpabloe@nvidia.com>

Refactor optional pad if-guard

0546cdf

Signed-off-by: janEbert <janpabloe@nvidia.com>

Merge output trim code paths

1416805

Signed-off-by: janEbert <janpabloe@nvidia.com>

Reduce verbosity

78c1348

Signed-off-by: janEbert <janpabloe@nvidia.com>

Explicitly mention padding target and limits

69f7044

Signed-off-by: janEbert <janpabloe@nvidia.com>

Merge input pad code paths

60a69a2

Signed-off-by: janEbert <janpabloe@nvidia.com>

Remove standalone MLA pad tests

7559cad

Signed-off-by: janEbert <janpabloe@nvidia.com>

Assert MLA pad selects a faster backend

ca7fd18

Signed-off-by: janEbert <janpabloe@nvidia.com>

Remove removed test from test runner script

e576da3

Signed-off-by: janEbert <janpabloe@nvidia.com>

Reduce verbosity

3702220

Signed-off-by: janEbert <janpabloe@nvidia.com>

janEbert added 3 commits October 13, 2026 01:25

Fix padding selection

2da4177

It could be that all `native_use_*_attention` values are False, so this needs to be handled. Signed-off-by: janEbert <janpabloe@nvidia.com>

Handle tuple Attn backend outputs

def8ad1

Signed-off-by: janEbert <janpabloe@nvidia.com>

Fold FP8 DSA+packed check into main padding check

4dcde58

Improve variable names along the way. Signed-off-by: janEbert <janpabloe@nvidia.com>

janEbert force-pushed the optional-mla-pad branch from ebcc35d to 4dcde58 Compare October 12, 2026 23:27

greptile-apps Bot reviewed Oct 12, 2026

View reviewed changes

Comment thread transformer_engine/pytorch/attention/dot_product_attention/utils.py

Comment thread tests/pytorch/attention/test_attention.py Outdated

cyanguwa added the 2.21 label Oct 12, 2026

Fix backend selection cache invalidation

3f00799

Signed-off-by: janEbert <janpabloe@nvidia.com>

greptile-apps Bot reviewed Oct 13, 2026

View reviewed changes

Comment thread transformer_engine/pytorch/attention/dot_product_attention/utils.py Outdated

Fix tests

0dc265d

Signed-off-by: janEbert <janpabloe@nvidia.com>

cyanguwa reviewed Oct 13, 2026

View reviewed changes

To-be-reviewed: minor changes

220c370

Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

greptile-apps Bot reviewed Oct 14, 2026

View reviewed changes

Comment thread tests/pytorch/attention/test_attention.py Outdated

janEbert added 2 commits October 14, 2026 18:56

Fix QKV layout

9739216

Also add batch size config. Signed-off-by: janEbert <janpabloe@nvidia.com>

Use varlength tuples with type hints

4029149

Signed-off-by: janEbert <janpabloe@nvidia.com>

janEbert and others added 3 commits October 14, 2026 19:13

Revert "Use varlength tuples with type hints"

f9a2e77

This reverts commit 4029149. Signed-off-by: janEbert <janpabloe@nvidia.com>

fix FA2 pad logic and should_pad overall logic

0620e6d

Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

Merge branch 'main' into optional-mla-pad

cc780bd

greptile-apps Bot reviewed Oct 14, 2026

View reviewed changes

Comment thread tests/pytorch/attention/test_attention.py Outdated

Update tests/pytorch/attention/test_attention.py

d69ff67

Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

janEbert commented Oct 21, 2026

View reviewed changes

Improve cache clear usability

f8455f1

Signed-off-by: janEbert <janpabloe@nvidia.com>

cyanguwa added 2.22 and removed 2.21 labels Oct 23, 2026

		if _attention_backends["backend_selection_requires_update"]:
		dpa_utils.clear_should_pad_qkv_head_dim_cache()

Conversation

janEbert commented Aug 10, 2026

Description

Type of change

Changes

Checklist:

Uh oh!

greptile-apps Bot commented Aug 10, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Confidence Score: 5/5

Summary

Diagram

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

janEbert commented Aug 13, 2026

Uh oh!

janEbert commented Aug 27, 2026

Uh oh!

cyanguwa left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

janEbert commented Oct 12, 2026

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

janEbert commented Oct 14, 2026

Uh oh!

greptile-apps Bot commented Oct 14, 2026

Uh oh!

cyanguwa commented Oct 14, 2026

Uh oh!

cyanguwa commented Oct 14, 2026

Uh oh!

Uh oh!

janEbert commented Oct 19, 2026

Uh oh!

janEbert commented Oct 20, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

This branch has not been deployed

Reviewers

Uh oh!

Assignees

Labels

greptile-apps Bot commented Aug 10, 2026 •

edited

Loading

cyanguwa left a comment •

edited

Loading

janEbert commented Oct 20, 2026 •

edited

Loading