compile error with custom softmax triton kernel

[mayank31398]

🐛 Describe the bug

The following code works fine if we commend out torch.compile

# **************************************************
# Copyright (c) 2025, Mayank Mishra
# **************************************************

import torch
import triton
import triton.language as tl


@triton.jit
def _load_x(x_ptr, h, H, BLOCK_SIZE_H, indices_b, mask_b, other=None):
    indices_h = h * BLOCK_SIZE_H + tl.arange(0, BLOCK_SIZE_H)
    mask_h = indices_h < H

    indices = indices_b[:, None] * H + indices_h[None, :]
    mask_bh = mask_b[:, None] & mask_h[None, :]

    x = tl.load(x_ptr + indices, mask=mask_bh, other=other)

    return x, indices, mask_bh


@triton.jit
def softmax_forward_triton_kernel(
    x_ptr,
    output_ptr,
    logits_multiplier,
    B,
    H,
    BLOCK_SIZE_B: tl.constexpr,
    BLOCK_SIZE_H: tl.constexpr,
):
    pid = tl.program_id(axis=0)

    indices_b = pid * BLOCK_SIZE_B + tl.arange(0, BLOCK_SIZE_B)
    mask_b = indices_b < B

    Z = tl.zeros((BLOCK_SIZE_B, 1), dtype=tl.float32)
    M = tl.full((BLOCK_SIZE_B, 1), -float("inf"), dtype=tl.float32)

    num_blocks_h = tl.cdiv(H, BLOCK_SIZE_H)

    for h in range(num_blocks_h):
        x, indices, mask_bh = _load_x(
            x_ptr=x_ptr, h=h, H=H, BLOCK_SIZE_H=BLOCK_SIZE_H, indices_b=indices_b, mask_b=mask_b, other=-float("inf")
        )

        x = x.to(tl.float32)
        if logits_multiplier is not None:
            x *= logits_multiplier

        prev_m = M
        m = tl.max(x, axis=1, keep_dims=True)
        M = max(M, m)

        x -= M
        x = tl.exp(x)
        Z = Z * tl.exp(prev_m - M) + tl.sum(x, axis=1, keep_dims=True)

    for h in range(num_blocks_h):
        x, indices, mask_bh = _load_x(
            x_ptr=x_ptr, h=h, H=H, BLOCK_SIZE_H=BLOCK_SIZE_H, indices_b=indices_b, mask_b=mask_b
        )

        x = x.to(tl.float32)
        if logits_multiplier is not None:
            x *= logits_multiplier

        x -= M
        x = tl.exp(x)
        x /= Z

        tl.store(output_ptr + indices, x, mask=mask_bh)


@triton.jit
def _load_output_output_grad(output_ptr, output_grad_ptr, h, H, BLOCK_SIZE_H, indices_b, mask_b):
    indices_h = h * BLOCK_SIZE_H + tl.arange(0, BLOCK_SIZE_H)
    mask_h = indices_h < H

    indices = indices_b[:, None] * H + indices_h[None, :]
    mask_bh = mask_b[:, None] & mask_h[None, :]

    output = tl.load(output_ptr + indices, mask=mask_bh)
    output_grad = tl.load(output_grad_ptr + indices, mask=mask_bh)

    return output, output_grad, indices, mask_bh


@triton.jit
def softmax_backward_triton_kernel(
    output_ptr,
    output_grad_ptr,
    x_grad_ptr,
    logits_multiplier,
    B,
    H,
    BLOCK_SIZE_B: tl.constexpr,
    BLOCK_SIZE_H: tl.constexpr,
):
    pid = tl.program_id(axis=0)

    indices_b = pid * BLOCK_SIZE_B + tl.arange(0, BLOCK_SIZE_B)
    mask_b = indices_b < B

    accumulator = tl.zeros((BLOCK_SIZE_B, 1), dtype=tl.float32)
    num_blocks_h = tl.cdiv(H, BLOCK_SIZE_H)

    for h in range(num_blocks_h):
        output, output_grad, indices, mask_bh = _load_output_output_grad(
            output_ptr=output_ptr,
            output_grad_ptr=output_grad_ptr,
            h=h,
            H=H,
            BLOCK_SIZE_H=BLOCK_SIZE_H,
            indices_b=indices_b,
            mask_b=mask_b,
        )

        acc = output_grad * output
        acc = acc.to(tl.float32)
        accumulator += tl.sum(acc, axis=1, keep_dims=True)

    for h in range(num_blocks_h):
        output, output_grad, indices, mask_bh = _load_output_output_grad(
            output_ptr=output_ptr,
            output_grad_ptr=output_grad_ptr,
            h=h,
            H=H,
            BLOCK_SIZE_H=BLOCK_SIZE_H,
            indices_b=indices_b,
            mask_b=mask_b,
        )

        output_grad -= accumulator
        output *= output_grad
        if logits_multiplier is not None:
            output *= logits_multiplier

        tl.store(x_grad_ptr + indices, output, mask=mask_bh)


class _Softmax(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x: torch.Tensor, logits_multiplier: float | None) -> torch.Tensor:
        output = torch.empty_like(x)

        B, H = x.size()

        BLOCK_SIZE_B = 1
        BLOCK_SIZE_H = min(triton.next_power_of_2(H), 4096 if x.dtype == torch.float32 else 8192)

        softmax_forward_triton_kernel[(B + BLOCK_SIZE_B - 1) // BLOCK_SIZE_B,](
            x_ptr=x,
            output_ptr=output,
            logits_multiplier=logits_multiplier,
            B=B,
            H=H,
            BLOCK_SIZE_B=BLOCK_SIZE_B,
            BLOCK_SIZE_H=BLOCK_SIZE_H,
        )

        ctx.save_for_backward(output)
        ctx.logits_multiplier = logits_multiplier

        return output

    @staticmethod
    def backward(ctx, output_grad: torch.Tensor) -> tuple[torch.Tensor | None]:
        output = ctx.saved_tensors[0]
        x_grad = torch.empty_like(output)

        B, H = x_grad.size()

        BLOCK_SIZE_B = 1
        BLOCK_SIZE_H = min(triton.next_power_of_2(H), 4096 if output.dtype == torch.float32 else 8192)


        softmax_backward_triton_kernel[(B + BLOCK_SIZE_B - 1) // BLOCK_SIZE_B,](
            output_ptr=output,
            output_grad_ptr=output_grad,
            x_grad_ptr=x_grad,
            logits_multiplier=ctx.logits_multiplier,
            B=B,
            H=H,
            BLOCK_SIZE_B=BLOCK_SIZE_B,
            BLOCK_SIZE_H=BLOCK_SIZE_H,
        )

        return x_grad, None

@torch.compile(fullgraph=True)
def softmax(
    x: torch.Tensor,
    logits_multiplier: float | None = None
) -> torch.Tensor:
    """computes softmax activation

    Args:
        x (torch.Tensor): input activation tensor
        logits_multiplier (float, optional): pre-multiplies `x` with `logits_multiplier` before computing softmax.
            Defaults to None.
        kernel_backend (KernelBackend | CutoTuneParameter, optional): kernel backend to prioritize.
            Defaults to KernelBackend.triton.

    Returns:
        torch.Tensor: output tensor
    """

    return _Softmax.apply(x, logits_multiplier)


x = torch.randn(5, 4, device=torch.cuda.current_device(), requires_grad=True)
print(softmax(x))

error:

Traceback (most recent call last):
  File "/proj/checkpoints/shawntan/mayank/flash-model-architectures/a.py", line 214, in <module>
    print(softmax(x))
          ^^^^^^^^^^
  File "/u/shawntan/.conda/envs/ai-mayank/lib/python3.12/site-packages/torch/_dynamo/eval_frame.py", line 900, in compile_wrapper
    raise e.remove_dynamo_frames() from None  # see TORCHDYNAMO_VERBOSE=1
    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/u/shawntan/.conda/envs/ai-mayank/lib/python3.12/site-packages/torch/_dynamo/output_graph.py", line 2213, in _call_user_compiler
    raise BackendCompilerFailed(
  File "/u/shawntan/.conda/envs/ai-mayank/lib/python3.12/site-packages/torch/_dynamo/output_graph.py", line 2188, in _call_user_compiler
    compiled_fn = compiler_fn(gm, example_inputs)
                  ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/u/shawntan/.conda/envs/ai-mayank/lib/python3.12/site-packages/torch/_dynamo/repro/after_dynamo.py", line 156, in __call__
    compiled_gm = compiler_fn(gm, example_inputs)
                  ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/u/shawntan/.conda/envs/ai-mayank/lib/python3.12/site-packages/torch/__init__.py", line 2388, in __call__
    return compile_fx(model_, inputs_, config_patches=self.config)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/u/shawntan/.conda/envs/ai-mayank/lib/python3.12/site-packages/torch/_inductor/compile_fx.py", line 2681, in compile_fx
    return aot_autograd(
           ^^^^^^^^^^^^^
  File "/u/shawntan/.conda/envs/ai-mayank/lib/python3.12/site-packages/torch/_dynamo/backends/common.py", line 117, in __call__
    cg = aot_module_simplified(gm, example_inputs, **self.kwargs)
         ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/u/shawntan/.conda/envs/ai-mayank/lib/python3.12/site-packages/torch/_functorch/aot_autograd.py", line 1107, in aot_module_simplified
    compiled_fn, _ = aot_stage2_compile(aot_state, aot_graph_capture)
                     ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/u/shawntan/.conda/envs/ai-mayank/lib/python3.12/site-packages/torch/_functorch/_aot_autograd/graph_compile.py", line 240, in aot_stage2_compile
    return aot_stage2_autograd(aot_state, aot_graph_capture)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/u/shawntan/.conda/envs/ai-mayank/lib/python3.12/site-packages/torch/_functorch/_aot_autograd/graph_compile.py", line 1393, in aot_stage2_autograd
    fw_module, bw_module = aot_config.partition_fn(
                           ^^^^^^^^^^^^^^^^^^^^^^^^
  File "/u/shawntan/.conda/envs/ai-mayank/lib/python3.12/site-packages/torch/_inductor/compile_fx.py", line 2118, in partition_fn
    return min_cut_rematerialization_partition(
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/u/shawntan/.conda/envs/ai-mayank/lib/python3.12/site-packages/torch/_functorch/partitioners.py", line 2792, in min_cut_rematerialization_partition
    node_info = classify_nodes(joint_module, static_lifetime_input_indices)
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/u/shawntan/.conda/envs/ai-mayank/lib/python3.12/site-packages/torch/_functorch/partitioners.py", line 2759, in classify_nodes
    forward_only_graph = _extract_graph_with_inputs_outputs(
                         ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/u/shawntan/.conda/envs/ai-mayank/lib/python3.12/site-packages/torch/_functorch/partitioners.py", line 238, in _extract_graph_with_inputs_outputs
    assert not isinstance(env[x], InvalidNodeBase), (
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
torch._dynamo.exc.BackendCompilerFailed: backend='inductor' raised:
AssertionError: Node getitem_2 was invalid, but is output

Set TORCHDYNAMO_VERBOSE=1 for the internal stack trace (please do this especially if you're reporting a bug to PyTorch). For even more developer context, set TORCH_LOGS="+dynamo"

Versions

Collecting environment information...
PyTorch version: 2.10.0.dev20250916+cu128
Is debug build: False
CUDA used to build PyTorch: 12.8
ROCM used to build PyTorch: N/A

OS: Red Hat Enterprise Linux 9.4 (Plow) (x86_64)
GCC version: (GCC) 11.4.1 20231218 (Red Hat 11.4.1-3)
Clang version: Could not collect
CMake version: version 3.31.2
Libc version: glibc-2.34

Python version: 3.12.11 | packaged by Anaconda, Inc. | (main, Jun  5 2025, 13:09:17) [GCC 11.2.0] (64-bit runtime)
Python platform: Linux-5.14.0-427.42.1.el9_4.x86_64-x86_64-with-glibc2.34
Is CUDA available: True
CUDA runtime version: 12.1.105
CUDA_MODULE_LOADING set to: 
GPU models and configuration: GPU 0: NVIDIA H100 80GB HBM3
Nvidia driver version: 560.35.03
cuDNN version: Could not collect
Is XPU available: False
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

CPU:
Architecture:                         x86_64
CPU op-mode(s):                       32-bit, 64-bit
Address sizes:                        46 bits physical, 57 bits virtual
Byte Order:                           Little Endian
CPU(s):                               96
On-line CPU(s) list:                  0-95
Vendor ID:                            GenuineIntel
Model name:                           Intel(R) Xeon(R) Platinum 8468
CPU family:                           6
Model:                                143
Thread(s) per core:                   1
Core(s) per socket:                   48
Socket(s):                            2
Stepping:                             8
CPU(s) scaling MHz:                   100%
CPU max MHz:                          3800.0000
CPU min MHz:                          800.0000
BogoMIPS:                             4200.00
Flags:                                fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush dts acpi mmx fxsr sse sse2 ss ht tm pbe syscall nx pdpe1gb rdtscp lm constant_tsc art arch_perfmon pebs bts rep_good nopl xtopology nonstop_tsc cpuid aperfmperf tsc_known_freq pni pclmulqdq dtes64 monitor ds_cpl smx est tm2 ssse3 sdbg fma cx16 xtpr pdcm pcid dca sse4_1 sse4_2 x2apic movbe popcnt tsc_deadline_timer aes xsave avx f16c rdrand lahf_lm abm 3dnowprefetch cpuid_fault epb cat_l3 cat_l2 cdp_l3 cdp_l2 ssbd mba ibrs ibpb stibp ibrs_enhanced fsgsbase tsc_adjust bmi1 avx2 smep bmi2 erms invpcid cqm rdt_a avx512f avx512dq rdseed adx smap avx512ifma clflushopt clwb intel_pt avx512cd sha_ni avx512bw avx512vl xsaveopt xsavec xgetbv1 xsaves cqm_llc cqm_occup_llc cqm_mbm_total cqm_mbm_local split_lock_detect avx_vnni avx512_bf16 wbnoinvd dtherm ida arat pln pts avx512vbmi umip pku ospke waitpkg avx512_vbmi2 gfni vaes vpclmulqdq avx512_vnni avx512_bitalg tme avx512_vpopcntdq la57 rdpid bus_lock_detect cldemote movdiri movdir64b enqcmd fsrm md_clear serialize tsxldtrk pconfig arch_lbr ibt amx_bf16 avx512_fp16 amx_tile amx_int8 flush_l1d arch_capabilities
L1d cache:                            4.5 MiB (96 instances)
L1i cache:                            3 MiB (96 instances)
L2 cache:                             192 MiB (96 instances)
L3 cache:                             210 MiB (2 instances)
NUMA node(s):                         2
NUMA node0 CPU(s):                    0-47
NUMA node1 CPU(s):                    48-95
Vulnerability Gather data sampling:   Not affected
Vulnerability Itlb multihit:          Not affected
Vulnerability L1tf:                   Not affected
Vulnerability Mds:                    Not affected
Vulnerability Meltdown:               Not affected
Vulnerability Mmio stale data:        Not affected
Vulnerability Reg file data sampling: Not affected
Vulnerability Retbleed:               Not affected
Vulnerability Spec rstack overflow:   Not affected
Vulnerability Spec store bypass:      Mitigation; Speculative Store Bypass disabled via prctl
Vulnerability Spectre v1:             Mitigation; usercopy/swapgs barriers and __user pointer sanitization
Vulnerability Spectre v2:             Mitigation; Enhanced / Automatic IBRS; IBPB conditional; RSB filling; PBRSB-eIBRS SW sequence; BHI BHI_DIS_S
Vulnerability Srbds:                  Not affected
Vulnerability Tsx async abort:        Not affected

Versions of relevant libraries:
[pip3] numpy==2.3.3
[pip3] nvidia-cublas-cu12==12.8.4.1
[pip3] nvidia-cuda-cupti-cu12==12.8.90
[pip3] nvidia-cuda-nvrtc-cu12==12.8.93
[pip3] nvidia-cuda-runtime-cu12==12.8.90
[pip3] nvidia-cudnn-cu12==9.10.2.21
[pip3] nvidia-cufft-cu12==11.3.3.83
[pip3] nvidia-curand-cu12==10.3.9.90
[pip3] nvidia-cusolver-cu12==11.7.3.90
[pip3] nvidia-cusparse-cu12==12.5.8.93
[pip3] nvidia-cusparselt-cu12==0.7.1
[pip3] nvidia-nccl-cu12==2.27.5
[pip3] nvidia-nvjitlink-cu12==12.8.93
[pip3] nvidia-nvtx-cu12==12.8.90
[pip3] pytorch-triton==3.5.0+git5ae38bdb
[pip3] torch==2.10.0.dev20250916+cu128
[pip3] triton==3.4.0
[conda] numpy                     2.3.3                    pypi_0    pypi
[conda] nvidia-cublas-cu12        12.8.4.1                 pypi_0    pypi
[conda] nvidia-cuda-cupti-cu12    12.8.90                  pypi_0    pypi
[conda] nvidia-cuda-nvrtc-cu12    12.8.93                  pypi_0    pypi
[conda] nvidia-cuda-runtime-cu12  12.8.90                  pypi_0    pypi
[conda] nvidia-cudnn-cu12         9.10.2.21                pypi_0    pypi
[conda] nvidia-cufft-cu12         11.3.3.83                pypi_0    pypi
[conda] nvidia-curand-cu12        10.3.9.90                pypi_0    pypi
[conda] nvidia-cusolver-cu12      11.7.3.90                pypi_0    pypi
[conda] nvidia-cusparse-cu12      12.5.8.93                pypi_0    pypi
[conda] nvidia-cusparselt-cu12    0.7.1                    pypi_0    pypi
[conda] nvidia-nccl-cu12          2.27.5                   pypi_0    pypi
[conda] nvidia-nvjitlink-cu12     12.8.93                  pypi_0    pypi
[conda] nvidia-nvtx-cu12          12.8.90                  pypi_0    pypi
[conda] pytorch-triton            3.5.0+git5ae38bdb          pypi_0    pypi
[conda] torch                     2.10.0.dev20250916+cu128          pypi_0    pypi
[conda] triton                    3.4.0                    pypi_0    pypi

cc @chauhang @penguinwu @zou3519 @bdhirsh

[pythongiant]

What if you detach (perferably clone) the tensor before saving?

also, you set @torch.compile(fullgraph=True), Try removing fullgraph=True or set it to False. The fullgraph path forces AOT-autograd over the whole forward/backward, with fullgraph off you are less likely to hit the AOT partitioner assertion.

essentially replace:
ctx.save_for_backward(output)
with
ctx.save_for_backward(output.detach().clone())

[mayank31398]

What if you detach (perferably clone) the tensor before saving?

also, you set @torch.compile(fullgraph=True), Try removing fullgraph=True or set it to False. The fullgraph path forces AOT-autograd over the whole forward/backward, with fullgraph off you are less likely to hit the AOT partitioner assertion.

essentially replace: ctx.save_for_backward(output) with ctx.save_for_backward(output.detach().clone())

the clone with detach works.
But I would prefer not cloning, its just wasting memory.
It seems like a bug to me.
Also what is the assertion?
Why is this happening?

I would also prefer to have fullgraph

[pythongiant]

You don’t actually need the clone. The core problem is that Dynamo/AOT sees a graph node being saved so the fix is to save something that isn’t a traced node. detach() by itself is usually enough, if the tensor is materialized outside the FX graph.

[mayank31398]

the detach() by itself doesnt work for me though.
It requires a clone() :/

[mayank31398]

Also, https://github.com/open-lm-engine/flash-model-architectures/blob/7dadcd0704498169d6f585bc2d80bede7384f002/fma/ops/swiglu/__init__.py#L15-L59
works with fullgraph. The error is only happening for this softmax kernel.
Also if I wrap the softmax_forward and backward kernels inside a torch custom op, then everything works fine.

[pythongiant]

the detach() by itself doesnt work for me though. It requires a clone() :/

try it with a detach().contiguous() ?

Also, https://github.com/open-lm-engine/flash-model-architectures/blob/7dadcd0704498169d6f585bc2d80bede7384f002/fma/ops/swiglu/__init__.py#L15-L59
works with fullgraph. The error is only happening for this softmax kernel.
Also if I wrap the softmax_forward and backward kernels inside a torch custom op, then everything works fine.

So it's not Triton per se; it's how the forward output is represented to the tracer and whether the tracer ends up needing to treat some internal node as an output/saved tensor.

If you want to probe what Dynamo sees, add small prints (run without compile or with TORCHDYNAMO_VERBOSE=1) to check these for output:

print("grad_fn:", getattr(output, "grad_fn", None))
print("is_leaf:", output.is_leaf)
print("is_view:", getattr(output, "_is_view", lambda: False)())

Compare the values between:

i. normal (no torch.compile)
ii. torch.compile(fullgraph=True)
iii. wrapping with torch._dynamo.disable() around ctx.save_for_backward

That will show whether Dynamo is seeing a node-like saved tensor.

TLDR;

Try ctx.save_for_backward(output.detach.contiguous())
If not, use the torch._dynamo.disable() wrapper around the ctx.save_for_backward(output) call.

import torch._dynamo

# ... inside forward
with torch._dynamo.disable():
    ctx.save_for_backward(output)

[pythongiant]

• Your output is produced by writing into a pre-allocated tensor (output = torch.empty_like(x) + tl.store(...)) That means the runtime identity/aliasing and how the tracer represents output can be different from typical aten-created tensors (which go through regular ATen ops and produce clean FX nodes).

• The tracer/AOT-autograd tries to capture the forward graph nodes and also the saved tensors used in backward. If a saved tensor is (or is represented by) a graph node / view / getitem-like node rather than a plain runtime value, the partitioner can end up treating that saved tensor as a graph output, which it then finds invalid during partitioning and raises InvalidNodeBase / Node ... was invalid, but is output.

• Your swiglu/other implementation likely creates output via aten ops (or otherwise in a tracer-friendly way) so the traced graph and saved tensors are in a form Inductor expects.

• Your output is produced by writing into a pre-allocated tensor (output = torch.empty_like(x) + tl.store(...)). That means the runtime identity/aliasing and how the tracer represents output can be different from typical aten-created tensors (which go through regular ATen ops and produce clean FX nodes).

• The tracer/AOT-autograd tries to capture the forward graph nodes and also the saved tensors used in backward. If a saved tensor is (or is represented by) a graph node / view / getitem-like node rather than a plain runtime value, the partitioner can end up treating that saved tensor as a graph outpus, which it then finds invalid during partitioning and raises InvalidNodeBase / Node ... was invalid, but is output.

• Your swiglu/other implementation likely creates output via aten ops (or otherwise in a tracer-friendly way) so the traced graph and saved tensors are in a form Inductor expects.

• Wrapping kernels in a custom op hides the internals from Dynamo/FX (it becomes an opaque ATen operator). That prevents the tracer from trying to treat internal intermediate nodes as graph outputs, which is why the problem goes away when you put the kernels behind a custom op.

[mayank31398]

wait @pythongiant are you telling me dynamo is tracing into the triton kernel?
and failing to understand what how i am handling the output tensor?
and custom_op is a workaround (prevents dynamo from looking into the triton kernel) thats why it works?

[pythongiant]

wait @pythongiant are you telling me dynamo is tracing into the triton kernel? and failing to understand what how i am handling the output tensor? and custom_op is a workaround (prevents dynamo from looking into the triton kernel) thats why it works?

Exactly.
Dynamo does trace the Python wrapper logic around it (the torch.autograd.Function.forward where you allocate output = torch.empty_like(x) and then call your Triton kernel to fill it). To Dynamo, the tl.store inside the Triton kernel is invisible. All it sees is: you allocated an empty tensor, you handed it into some opaque call.

and then you return that same tensor.

AOTAutograd wants saved tensors to be “valid graph values” (things it can replay or rematerialize).

But here output is from a pre-allocated buffer, mutated in-place by a kernel Dynamo doesn’t know about.

So the FX graph treats output as an “invalid node” (basically “I don’t know how this tensor was produced”). When that “invalid node” gets marked as both “output of the forward graph” and “something backward needs,” the partitioner asserts → your error.

When you wrap the Triton kernel in a torch.library.custom_op:

• Dynamo sees a single opaque “op node” with inputs and outputs.

• It never has to reason about how the outputs were produced, they’re just results of a legal op.

• The saved tensor is now considered a valid graph output → partitioner is happy. This is why the swiglu code you linked works fine: its forward path ultimately returns tensors from known aten ops (or is wrapped as a proper op), so Dynamo never sees “invalid” outputs.

Wrapping in a custom_op (or hiding the save) is a workaround that prevents Dynamo from needing to reason about your kernel’s internals.

[mayank31398]

right, makes sense now!
But I dont understand why it works for the RMSNorm kernel.
I do it the same way there.
If I remove custom_op from rmsnorm, it works fine there.

[StrongerXi]

cc @zou3519