TRT/TRT-LLM support for NestedTensor / jagged or ragged tensors for doing away with padding

[vadimkantorov]

Does TRT / TRT-LLM support pytorch's concept of NestedTensor or is there some similar concept?

The idea is to not compute FeedForward's on padded-out positions and skip attention computation for them as well

I guess the question also becomes of whether underlying flash attention implementation supports such jagged compact formats.

In PyTorch it seems that they support SDPA with NestedTensor for speedups: pytorch/pytorch#105913 (comment) (can be useful for batched BERT inference)

Thanks!

[Advait251206]

Short answer: no — TensorRT and TensorRT-LLM do not currently support PyTorch NestedTensor or true ragged/jagged tensor execution in the way PyTorch SDPA does.

But there are a few important nuances and partial alternatives.

---

1. Does TensorRT / TensorRT-LLM support NestedTensor?

❌ Native support: NO

• TensorRT operates on dense, fixed-shape tensors
• TensorRT-LLM assumes:
  • padded sequences
  • static or dynamic shapes within bounds
  • explicit attention masks

There is:

• no NestedTensor abstraction
• no per-sample variable-length kernel execution
• no graph-level ragged control flow

---

2. Why not? (core reason)

TensorRT is built around:

kernel fusion + shape specialization

That requires:

• contiguous memory layouts
• predictable tensor shapes
• vectorized execution across batch

Ragged tensors break this assumption because:

• each sequence has different length
• compute becomes irregular
• GPU efficiency drops without padding/alignment strategies

---

3. What TensorRT-LLM uses instead

Instead of NestedTensor, TensorRT-LLM uses:

✔ 1. Packed sequences + KV cache indexing

Internally, it supports:

• sequence length arrays (seq_len)
• position IDs
• attention masks
• KV cache “slot mapping”

This allows logical raggedness, but physical storage is still dense.

---

✔ 2. In-flight batching (important)

TensorRT-LLM runtime (executor) supports:

• dynamic batching of different sequence lengths
• token-level scheduling

But still:

everything is padded inside kernels or handled via masking

---

✔ 3. FlashAttention-style masked kernels

TensorRT-LLM integrates optimized attention kernels (FlashAttention-like):

• compute is skipped using masks
• but tensor layout is still dense

So:

• padding is not eliminated
• it is masked away computationally

---

4. Does FlashAttention support NestedTensor?

Partially yes (PyTorch ecosystem)

PyTorch SDPA + NestedTensor:

• can avoid materializing padding
• can fuse attention over variable-length sequences

BUT:

❌ TensorRT-LLM does NOT use PyTorch SDPA

❌ and does NOT consume NestedTensor inputs

Instead it uses:

• CUDA kernels (FlashAttention / paged attention variants)
• custom KV cache layouts

---

5. What about “true ragged execution”?

What you're describing:

skip FFN + attention entirely for padded tokens

This would require:

• per-token execution graphs
• dynamic kernel launches per sequence
• irregular memory access patterns

👉 This is currently not supported in TensorRT

---

6. Closest existing alternatives in TensorRT-LLM

✔ Paged KV cache (important)

This is the closest concept to “ragged efficiency”:

• KV cache is stored in pages
• sequences are not required to be contiguous
• memory reuse is efficient

But:

• still operates on padded tokens inside kernels

---

✔ Token-level scheduling

Instead of sequence-level compute:

• execution happens per token step
• batching happens across requests

This improves efficiency but does not remove padding entirely.

---

7. Why fully ragged execution is hard on GPUs

Even if theoretically possible:

• warp divergence increases massively
• memory coalescing breaks
• kernel efficiency collapses
• FlashAttention-style tiling becomes impossible

So industry compromise is:

“mask padding, don’t eliminate it structurally”

---

8. Bottom line

Feature	TensorRT / TRT-LLM
PyTorch NestedTensor	❌ Not supported
Ragged / jagged execution	❌ Not supported
Padding-free compute	❌ Not supported
Masked attention	✔ Yes
Variable-length batching	✔ Yes (via executor)
Paged KV cache	✔ Yes

---

Final takeaway

TensorRT-LLM chooses:

dense computation + masking + scheduling

instead of:

true ragged tensor execution

So you get most of the performance benefits (via FlashAttention + KV paging), but not full NestedTensor-style compute elimination.

---

If this answer helped or pointed you in the right direction, I'd appreciate it if you could mark it as the accepted answer so it's easier for others with the same issue to find.

Also, if you found my contribution useful, I'd appreciate it if you could check out my GitHub profile, follow me, and star any repositories you find interesting.

GitHub: https://github.com/Advait251206