[PyTorch] Torch.compile support

Make TE layers compatible with torch.compile.

@sbhavani

A few months ago, I also had trouble compiling TE fused RoPE in TorchTitan nightly:

        # Apply RoPE.
        # xq, xk = apply_rotary_emb(xq, xk, freqs_cis=freqs_cis)
        xq = apply_rotary_pos_emb_te(
            xq,
            # Needs to be real-valued for TE fused RoPE.
            freqs_cis.unsqueeze(1).unsqueeze(1),
            tensor_format="bshd",
            interleaved=False,
            fused=True,
        )
        xk = apply_rotary_pos_emb_te(
            xk,
            freqs_cis.unsqueeze(1).unsqueeze(1),
            tensor_format="bshd",
            interleaved=False,
            fused=True,
        )

@sbhavani I think I may or may not have talked to you about this in our many conversations around that time, but this blocked me from getting Llama 405B metrics with FSDP2 + Torch Compile with this fused RoPE (for parity with Megatron-LM's fused RoPE) so we had to math out the best-case improvement instead.

I tried this adding it as a custom op, but didn't have time to dig further.

Here's a simple repro

import torch
import transformer_engine.pytorch
from transformer_engine.pytorch.attention.rope import RotaryPositionEmbedding

layer = transformer_engine.pytorch.TransformerLayer(128, 128, 8)
rope_layer = RotaryPositionEmbedding(128 // 8)

inputs = torch.randn(8, 128).to("cuda")
freqs = rope_layer.forward(128)

inputs.to("cuda")
freqs.to("cuda")
layer.to("cuda")

layer = torch.compile(layer, fullgraph=True)
output = layer(inputs, rotary_pos_emb=freqs)

related issue: #1241

There is draft PR to add torch.compile support for te.Sequential #2608 - te.sequential will be the first TE layers will torch.compile support.

Is this not relatively easy to add, and I think would be a big win. All that is needed is to register some ops and remove the explicit no dynamo tags. For me I see no improvements using TE as we compile our whole graph, with all the graph breaks its not faster or more memory efficient.

@ben-da6

@ben-da6 we would welcome a PR! It'd be great to see even in a small scale test where torch.compile could speed up

We are actively working on that, this is much harder that it seems - there are multiple issues reported by us to PyTorch which will need to be closed before full torch.compile support will be merged. They are mostly related to new PyTorch OpaqueObject API, which allows to passing arbitrary objects as an inputs to torch custom operators. For example one merged PR here.

There are some parts of the TE which can be converted simply - like MoE related permutation ops which we will merge soon. Then we work on te.Linear and te.Sequential - which are more complex, but we already have some prototypes and we work towards exposing them as non-draft PRs. The attention is the most complex part, we have some issues there - for example, but not limited to, choice of the attention backend based on tensor dptr() parameters. It will be merged probably as the last part.

Any updates on this?

We have draft PR #3053 which is actively developed and the reviews will start soon.

I think the hf kernels framework is compatible with torch.compile
I think the best options would be to just expose the custom kernels in this repo there instead

Navigation Menu

Repository navigation

[PyTorch] Torch.compile support #2590

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

[PyTorch] Torch.compile support #2590

Description

Activity

sbhavani commented on Jan 13, 2026

cspades commented on Jan 13, 2026

pstjohn commented on Jan 13, 2026

ntadimeti commented on Jan 14, 2026

pggPL commented on Jan 21, 2026

ben-da6 commented on Feb 23, 2026

sbhavani commented on Feb 23, 2026

pggPL commented on Feb 23, 2026

omaralvarez commented on May 20, 2026

Abhishekghosh1998 commented on Jun 2, 2026

pggPL commented on Jun 3, 2026

honeyspoon commented on Jun 3, 2026

Metadata

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Issue actions