Multihead Attention does not work with jagged tensors due to __torch_function__

[SamuelGabriel]

🐛 Describe the bug

When I create a nested tensor in jagged format and try to use it with MHA, it throws

AssertionError: MultiheadAttention does not support NestedTensor outside of its fast path. The fast path was not hit because some Tensor argument has_torch_function

Reproduce with:

import torch

N=512
nheads=8

query = torch.nested.nested_tensor([
        torch.randn(100, N, device=device)
        for _ in range(2)
    ], layout=torch.jagged)

mha = torch.nn.MultiheadAttention(N,nheads,batch_first=True).eval()
mha(query,query,query)

It works with strided nested tensors, though, but the documentation says that jagged tensors should be supported and the underlying scaled_dot_product_attention does support it.

Versions

I am on python 3.12 and torch 2.8.0a0(+fb).

cc @cpuhrsch @jbschlosser @bhosmer @drisspg @soulitzer @davidberard98 @YuqingJ

[jbschlosser]

Note that the torch.nn maintainers are slowly pushing for usage of SDPA / FlexAttention rather than the torch.nn-provided Transformer modules (including nn.Transformer, nn.MHA, etc.) since the former are more flexible and easier to maintain.

That said, it wouldn't hurt to support MHA + NJT as MHA is likely not going away anytime soon, and it's arguably silly to require manual SDPA usage to use NJT. We'd accept a PR for this if you'd like to contribute.

As a workaround, note the minimal MHA implementation for NJT in the tutorial here.

[singhmeet11]

Hi @jbschlosser, this is my first time trying to contribute to PyTorch. Would it be okay if I took this?

[jbschlosser]

@singhmeet11 sure thing, go for it :)

[singhmeet11]

I have a few questions regarding supporting Nested Jagged Tensors (NJTs) in torch.nn.MultiheadAttention.

From what I understand, the core issue is that NJTs trigger __torch_function__, which disables the fast path in MHA — and currently, MHA requires the fast path to support NestedTensors at all.

Here are a few possible directions to address this:

1. Suppress has_torch_function for NJTs

One idea is to bypass the check that disables the fast path when __torch_function__ is present. However, this seems unsafe — the current implementation of torch._native_multi_head_attention likely isn’t designed to support tensors with custom overrides. Forcing it could lead to undefined behavior.

2. Modify the Backend (_native_multi_head_attention)

Another option is to extend the C++ backend to handle NJTs natively — supporting jagged shapes and attention computation without requiring flattening or padding. This would be the most robust solution long-term, but it’s also the most complex and would require careful autograd support.

3. Use a Custom Python MHA

As demonstrated in the Transformer Building Blocks tutorial, we can upgrade the implementation of MHA to operate on unbound sequences from the NJT and reassemble the output. This approach avoids the fast path entirely and works during training, though it comes with some performance tradeoffs.

I don't think I can do the second solution right now. However, 3rd seems possible. Do let me know if this is the correct direction. Thanks @jbschlosser

[jbschlosser]

@singhmeet11 I'll try to clear up some things, as the mess in MHA can certainly be confusing.

• The fast path uses an older version of nested tensors ("NST" AKA Nested Strided Tensor) for perf improvements. This is not supported for the newer NJT and likely never will be.
  • As you mentioned, torch._native_multi_head_attention is not designed to work with custom tensor types, and this includes NJT. Thus the fast path -should- be disabled when a non-standard tensor is encountered, as the __torch_function__ check does. We don't want to change this behavior or suppress __torch_function__() for NJT.
• I don't expect the fast path to be updated, and in fact, we'd prefer to move away from it towards more composable / maintainable methods for high performance. This includes usage of accelerated SDPA / Flex attention kernels. Your option (2) isn't something we want to pursue.
• The "non-fast-path" of MHA uses SDPA underneath, and SDPA supports NJT. So it shouldn't be too much lift to support NJT for MHa's "non-fast-path".
  • Note that "non-fast-path" is somewhat of a misnomer. Accelerated SDPA (e.g. backed by flash attention) was introduced after the "fast path" of torch._native_multihead_attention(), and in fact the former may be faster than the latter.

So I think my recommendation aligns most closely with your (3) above: allow NJT inputs to MHA for the "non-fast-path" case, dispatching to SDPA underneath. This should mostly involve removing unnecessary asserts, testing with NJT inputs, and ensuring general forward / backward correctness.

[rusdes]

@jbschlosser @singhmeet11 Hi! Since this issue has been inactive for a while, I’d be happy to take it on. Let me know!

[jbschlosser]

@rusdes sounds great :)

[rusdes]

@jbschlosser Hey, I spent some time digging into the code and wanted to share what I think is going on, along with a couple of ideas that might help. Would love to hear your thoughts.

The main issue seems to be a shape mismatch. NJTs use batch size as the first dimension, but the non-fast path in MultiheadAttention.forward calls F.multi_head_attention_forward, which expects inputs in (L, N, E) format so batch size is in the second dimension. That mismatch makes things break when using NJTs.

Also, F.multi_head_attention_forward is optimized for dense tensors, and the lower-level functions it uses (like _in_projection_packed) assume dense input too, so they’re not compatible with NJTs out of the box.

Ideas

1. Modify F.multi_head_attention_forward to handle NJTs
   This would mean updating a bunch of internals to support batch-first inputs and jagged shapes. Definitely a big change.

2. Add a separate NJT path in MultiheadAttention.forward
   We could check if the input is an NJT and, if so, go to the "non-fast-path", skip F.multi_head_attention_forward entirely and use SDPA directly instead. It’d duplicate a bit of logic, but would keep things clean and avoid touching the method optimized for dense tensors.

I think the second option makes most sense but wanted to hear what you think and make sure I'm on the right track.

Also, can you assign this issue to me? Thanks!

[rusdes]

Soft ping @jbschlosser

[jbschlosser]

The main issue seems to be a shape mismatch. NJTs use batch size as the first dimension, but the non-fast path in MultiheadAttention.forward calls F.multi_head_attention_forward, which expects inputs in (L, N, E) format so batch size is in the second dimension. That mismatch makes things break when using NJTs.

Good point and thanks for calling this out. NJT does not support transposing the batch dimension, so it can't match the expected interface for F.multi_head_attention(), which expects batch dim == 1.

Also, F.multi_head_attention_forward is optimized for dense tensors, and the lower-level functions it uses (like _in_projection_packed) assume dense input too, so they’re not compatible with NJTs out of the box.

can you provide more detail into which ops are failing with NJT? AFAICT _in_packed_projection() is composed of a bunch of standard ops that should be supported by NJT (e.g. linear(), unflatten(), transpose(), squeeze(), etc.). If there's a particular op gap, a better fix could be to address that.

I think the second option makes most sense but wanted to hear what you think and make sure I'm on the right track.

Agreed that option 2 sounds cleanest. This should be provided via a private API (named starting with an underscore). I'd still hope that we can reuse some helpers like _in_packed_projection() mentioned above. Thanks!

Also, can you assign this issue to me? Thanks!

@rusdes we don't generally do issue assignments; anyone is welcome to contribute at any time.

I'll also mention that this work is low priority now as both MHA and NJT are not being actively developed. I'll still make an effort to get a PR reviewed if you put one together.

[rusdes]

Hi @jbschlosser ,

I’ve just wrapped up the changes and will be opening a PR shortly. Here’s what I did:

• Context: multi_head_attention_forward uses __torch_function__ dispatch, but tested tensor doesn't implement it. This triggers DisableTorchFunctionSubclass() as a backup plan, which disables NJT logic and leads to SDPA failures.

• Solution: I added a separate NJT-specific multi_head_attention_forward implementation and registered it with the dispatcher for nested tensors. This keeps the original function untouched and avoids complicating the main path with NJT-specific logic.

A few notes about my changes so that it'll be easier for you to review:

• Supports separate query, key, and value as Nested Tensors (not just self-attention), as long as dimensions align.
• key_padding_mask is not applicable since NJTs don't have padding.
• attn_mask isn't supported, consistent with current NJT + SDPA limitations. (See here _validate_sdpa_input)
• need_weights = True isn't supported since it doesn't use SDPA to calculate.
• is_causal=True throws RuntimeError: No viable backend for scaled_dot_product_attention was found. I’d appreciate any thoughts on whether this should be handled or explicitly unsupported.
• I added unit tests to cover basic cases. Happy to expand if you have anything specific in mind.

Would love to hear your thoughts! Let me know if anything should be changed or improved.

P.S. Figuring out the dispatch path and how to cleanly extend it for NJTs was trickier than I expected haha. But it was definitely a fun exercise! :)