pytorch - 💡(How to fix) Fix [torch.compile] InductorError: both a fallback and a decomp for `aten.index_add.default` with bfloat16 [2 comments, 2 participants]

torch._inductor.exc.InductorError: AssertionError: both a fallback and a decomp for same op: aten.index_add.default

---

assert op not in check_decomps or override_decomp, (
    f"both a fallback and a decomp for same op: {op}"
)

---

python repro.py

---

import torch
import torch.nn as nn


class Block(nn.Module):
    def __init__(self, dim=256, drop=0.1):
        super().__init__()
        self.norm = nn.LayerNorm(dim)
        self.proj = nn.Linear(dim, dim)
        self.gamma = nn.Parameter(1e-5 * torch.ones(dim))
        self.drop = drop

    def forward(self, x):
        B = x.shape[0]
        k = max(int(B * (1 - self.drop)), 1)
        idx = torch.randperm(B, device=x.device)[:k]
        res = self.proj(self.norm(x[idx])) * self.gamma
        return torch.index_add(x, 0, idx, res)


device = "cuda"
model = nn.Sequential(*[Block() for _ in range(4)])
DTYPE = torch.bfloat16
model = model.to(device, dtype=DTYPE).train()

for m in model:
    m.compile()

x = torch.randn(16, 197, 256, device=device, dtype=DTYPE)
model(x)  # crashes here

print("Success")

---

Traceback (most recent call last):
  File "/storage/home/huyvvo/fairvit_balanced_patches/repro_index_add_inductor_bug.py", line 39, in <module>
    model(x)  # crashes here
    │     └ tensor([[[ 0.1855,  0.7305,  0.9844,  ...,  0.4453, -0.3379, -0.1377],
         [-0.3633,  0.2051, -0.6250,  ..., -2.2812,  0.50...
    └ Sequential(
  (0): Block(
    (norm): LayerNorm((256,), eps=1e-05, elementwise_affine=True, bias=True)
    (proj): Linear(in_fea...
  File "/storage/home/huyvvo/.local/share/mamba/envs/fairvit-py312-ptnightly-xformers-20260404/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1778, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           │                │       └ {}
           │                └ (tensor([[[ 0.1855,  0.7305,  0.9844,  ...,  0.4453, -0.3379, -0.1377],
         [-0.3633,  0.2051, -0.6250,  ..., -2.2812,  0.5...
           └ Sequential(
  (0): Block(
    (norm): LayerNorm((256,), eps=1e-05, elementwise_affine=True, bias=True)
    (proj): Linear(in_fea...
  File "/storage/home/huyvvo/.local/share/mamba/envs/fairvit-py312-ptnightly-xformers-20260404/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1789, in _call_impl
    return forward_call(*args, **kwargs)
           │             │       └ {}
           │             └ (tensor([[[ 0.1855,  0.7305,  0.9844,  ...,  0.4453, -0.3379, -0.1377],
         [-0.3633,  0.2051, -0.6250,  ..., -2.2812,  0.5...
           └ <bound method Sequential.forward of Sequential(
  (0): Block(
    (norm): LayerNorm((256,), eps=1e-05, elementwise_affine=True, ...
  File "/storage/home/huyvvo/.local/share/mamba/envs/fairvit-py312-ptnightly-xformers-20260404/lib/python3.12/site-packages/torch/nn/modules/container.py", line 253, in forward
    input = module(input)
    │       │      └ tensor([[[ 0.1855,  0.7305,  0.9844,  ...,  0.4453, -0.3379, -0.1377],
         [-0.3633,  0.2051, -0.6250,  ..., -2.2812,  0.50...
    │       └ Block(
  (norm): LayerNorm((256,), eps=1e-05, elementwise_affine=True, bias=True)
  (proj): Linear(in_features=256, out_features...
    └ tensor([[[ 0.1855,  0.7305,  0.9844,  ...,  0.4453, -0.3379, -0.1377],
         [-0.3633,  0.2051, -0.6250,  ..., -2.2812,  0.50...
  File "/storage/home/huyvvo/.local/share/mamba/envs/fairvit-py312-ptnightly-xformers-20260404/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1776, in _wrapped_call_impl
    return self._compiled_call_impl(*args, **kwargs)  # type: ignore[misc]
           │                         │       └ {}
           │                         └ (tensor([[[ 0.1855,  0.7305,  0.9844,  ...,  0.4453, -0.3379, -0.1377],
         [-0.3633,  0.2051, -0.6250,  ..., -2.2812,  0.5...
           └ Block(
  (norm): LayerNorm((256,), eps=1e-05, elementwise_affine=True, bias=True)
  (proj): Linear(in_features=256, out_features...
  File "/storage/home/huyvvo/.local/share/mamba/envs/fairvit-py312-ptnightly-xformers-20260404/lib/python3.12/site-packages/torch/_dynamo/eval_frame.py", line 1050, in compile_wrapper
    raise e.remove_dynamo_frames() from None  # see TORCHDYNAMO_VERBOSE=1
  File "/storage/home/huyvvo/.local/share/mamba/envs/fairvit-py312-ptnightly-xformers-20260404/lib/python3.12/site-packages/torch/_dynamo/eval_frame.py", line 1036, in compile_wrapper
    return fn(*args, **kwargs)
           │   │       └ {}
           │   └ (tensor([[[ 0.1855,  0.7305,  0.9844,  ...,  0.4453, -0.3379, -0.1377],
         [-0.3633,  0.2051, -0.6250,  ..., -2.2812,  0.5...
           └ <bound method Module._call_impl of Block(
  (norm): LayerNorm((256,), eps=1e-05, elementwise_affine=True, bias=True)
  (proj): L...
  File "/storage/home/huyvvo/.local/share/mamba/envs/fairvit-py312-ptnightly-xformers-20260404/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1789, in _call_impl
    return forward_call(*args, **kwargs)
           │             │       └ {}
           │             └ (tensor([[[ 0.1855,  0.7305,  0.9844,  ...,  0.4453, -0.3379, -0.1377],
         [-0.3633,  0.2051, -0.6250,  ..., -2.2812,  0.5...
           └ <bound method Block.forward of Block(
  (norm): LayerNorm((256,), eps=1e-05, elementwise_affine=True, bias=True)
  (proj): Linea...
  File "/storage/home/huyvvo/.local/share/mamba/envs/fairvit-py312-ptnightly-xformers-20260404/lib/python3.12/site-packages/torch/_dynamo/convert_frame.py", line 2432, in __call__
    result = self._torchdynamo_orig_backend(
  File "/storage/home/huyvvo/.local/share/mamba/envs/fairvit-py312-ptnightly-xformers-20260404/lib/python3.12/site-packages/torch/_dynamo/convert_frame.py", line 2150, in __call__
    result = self._inner_convert(
  File "/storage/home/huyvvo/.local/share/mamba/envs/fairvit-py312-ptnightly-xformers-20260404/lib/python3.12/site-packages/torch/_dynamo/convert_frame.py", line 735, in __call__
    result = _compile(
  File "/storage/home/huyvvo/.local/share/mamba/envs/fairvit-py312-ptnightly-xformers-20260404/lib/python3.12/site-packages/torch/_dynamo/convert_frame.py", line 1919, in _compile
    guarded_code, tracer_output = compile_inner(code, one_graph, hooks)
    │             │               │             │     │          └ Hooks(guard_export_fn=None, guard_fail_fn=None, guard_filter_fn=None)
    │             │               │             │     └ False
    │             │               │             └ <code object forward at 0x6227fd74f340, file "/storage/home/huyvvo/fairvit_balanced_patches/repro_index_add_inductor_bug.py", li...
    │             │               └ <function _compile.<locals>.compile_inner at 0x7c32b56fbec0>
    │             └ <torch._dynamo.output_graph.DynamoTracerOutput object at 0x7c32b5708500>
    └ None
  File "/storage/home/huyvvo/.local/share/mamba/envs/fairvit-py312-ptnightly-xformers-20260404/lib/python3.12/site-packages/torch/_utils_internal.py", line 96, in wrapper_function
    return function(*args, **kwargs)
           │         │       └ {}
           │         └ (<code object forward at 0x6227fd74f340, file "/storage/home/huyvvo/fairvit_balanced_patches/repro_index_add_inductor_bug.py", l...
           └ <function _compile.<locals>.compile_inner at 0x7c32b56fbe20>
  File "/storage/home/huyvvo/.local/share/mamba/envs/fairvit-py312-ptnightly-xformers-20260404/lib/python3.12/site-packages/torch/_dynamo/convert_frame.py", line 1535, in compile_inner
    result = _compile_inner(code, one_graph, hooks)
             │              │     │          └ Hooks(guard_export_fn=None, guard_fail_fn=None, guard_filter_fn=None)
             │              │     └ False
             │              └ <code object forward at 0x6227fd74f340, file "/storage/home/huyvvo/fairvit_balanced_patches/repro_index_add_inductor_bug.py", li...
             └ <function _compile.<locals>._compile_inner at 0x7c32b56fba60>
  File "/storage/home/huyvvo/.local/share/mamba/envs/fairvit-py312-ptnightly-xformers-20260404/lib/python3.12/site-packages/torch/_dynamo/convert_frame.py", line 1594, in _compile_inner
    dynamo_output = compile_frame(
  File "/storage/home/huyvvo/.local/share/mamba/envs/fairvit-py312-ptnightly-xformers-20260404/lib/python3.12/site-packages/torch/_dynamo/convert_frame.py", line 1442, in compile_frame
    bytecode, tracer_output = transform_code_object(code, transform)
                              │                     │     └ <function compile_frame.<locals>.transform at 0x7c32b570f560>
                              │                     └ <code object forward at 0x6227fd74f340, file "/storage/home/huyvvo/fairvit_balanced_patches/repro_index_add_inductor_bug.py", li...
                              └ <function transform_code_object at 0x7c32b86fd6c0>
  File "/storage/home/huyvvo/.local/share/mamba/envs/fairvit-py312-ptnightly-xformers-20260404/lib/python3.12/site-packages/torch/_dynamo/bytecode_transformation.py", line 1626, in transform_code_object
    tracer_output = transformations(instructions, code_options)
                    │               │             └ {'co_argcount': 2, 'co_posonlyargcount': 0, 'co_kwonlyargcount': 0, 'co_nlocals': 6, 'co_stacksize': 7, 'co_flags': 3, 'co_code'...
                    │               └ [Instruction(opcode=151, opname='RESUME', arg=0, argval=0, offset=0, starts_line=22, is_jump_target=False, positions=Positions(l...
                    └ <function compile_frame.<locals>.transform at 0x7c32b570f560>
  File "/storage/home/huyvvo/.local/share/mamba/envs/fairvit-py312-ptnightly-xformers-20260404/lib/python3.12/site-packages/torch/_dynamo/convert_frame.py", line 1414, in transform
    tracer_output = trace_frame(
  File "/storage/home/huyvvo/.local/share/mamba/envs/fairvit-py312-ptnightly-xformers-20260404/lib/python3.12/site-packages/torch/_dynamo/convert_frame.py", line 342, in _fn
    return fn(*args, **kwargs)
           │   │       └ {'export': False, 'export_constraints': None, 'frame_state': {'_id': 0}, 'distributed_state': None, 'package': None}
           │   └ (<code object forward at 0x6227fd74f340, file "/storage/home/huyvvo/fairvit_balanced_patches/repro_index_add_inductor_bug.py", l...
           └ <function trace_frame at 0x7c32b7808680>
  File "/storage/home/huyvvo/.local/share/mamba/envs/fairvit-py312-ptnightly-xformers-20260404/lib/python3.12/site-packages/torch/_dynamo/convert_frame.py", line 875, in trace_frame
    run_tracer()
    └ <function trace_frame.<locals>.run_tracer at 0x7c32b56fa480>
  File "/storage/home/huyvvo/.local/share/mamba/envs/fairvit-py312-ptnightly-xformers-20260404/lib/python3.12/site-packages/torch/_dynamo/convert_frame.py", line 856, in run_tracer
    tracer.run()
    └ <torch._dynamo.symbolic_convert.InstructionTranslator object at 0x7c32b570a2a0>
  File "/storage/home/huyvvo/.local/share/mamba/envs/fairvit-py312-ptnightly-xformers-20260404/lib/python3.12/site-packages/torch/_dynamo/symbolic_convert.py", line 1705, in run
    while self.step():
  File "/storage/home/huyvvo/.local/share/mamba/envs/fairvit-py312-ptnightly-xformers-20260404/lib/python3.12/site-packages/torch/_dynamo/symbolic_convert.py", line 1372, in step
    self.dispatch_table[inst.opcode](self, inst)
    │                   │            │     └ Instruction(opcode=83, opname='RETURN_VALUE', arg=None, argval=None, offset=324, starts_line=27, is_jump_target=False, positions...
    │                   │            └ <torch._dynamo.symbolic_convert.InstructionTranslator object at 0x7c32b570a2a0>
    │                   └ Instruction(opcode=83, opname='RETURN_VALUE', arg=None, argval=None, offset=324, starts_line=27, is_jump_target=False, positions...
    └ <torch._dynamo.symbolic_convert.InstructionTranslator object at 0x7c32b570a2a0>
  File "/storage/home/huyvvo/.local/share/mamba/envs/fairvit-py312-ptnightly-xformers-20260404/lib/python3.12/site-packages/torch/_dynamo/symbolic_convert.py", line 5035, in RETURN_VALUE
    self._return(inst)
    │            └ Instruction(opcode=83, opname='RETURN_VALUE', arg=None, argval=None, offset=324, starts_line=27, is_jump_target=False, positions...
    └ <torch._dynamo.symbolic_convert.InstructionTranslator object at 0x7c32b570a2a0>
  File "/storage/home/huyvvo/.local/share/mamba/envs/fairvit-py312-ptnightly-xformers-20260404/lib/python3.12/site-packages/torch/_dynamo/symbolic_convert.py", line 5017, in _return
    all_stack_locals_metadata = self.output.compile_subgraph(
  File "/storage/home/huyvvo/.local/share/mamba/envs/fairvit-py312-ptnightly-xformers-20260404/lib/python3.12/site-packages/torch/_dynamo/output_graph.py", line 2053, in compile_subgraph
    self.compile_and_call_fx_graph(tx, pass2.graph_output_vars(), root)
    │                              │   │                          └ FakeRootModule(...)
    │                              │   └ <torch._dynamo.codegen.PyCodegen object at 0x7c32b5510b00>
    │                              └ <torch._dynamo.symbolic_convert.InstructionTranslator object at 0x7c32b570a2a0>
    └ OutputGraph(local_scope={'self': Block(
  (norm): LayerNorm((256,), eps=1e-05, elementwise_affine=True, bias=True)
  (proj): Lin...
  File "/storage/home/huyvvo/.local/share/mamba/envs/fairvit-py312-ptnightly-xformers-20260404/lib/python3.12/site-packages/torch/_dynamo/output_graph.py", line 2700, in compile_and_call_fx_graph
    compiled_fn = self.call_user_compiler(gm, self.example_inputs())
                  │                       │   └ OutputGraph(local_scope={'self': Block(
  (norm): LayerNorm((256,), eps=1e-05, elementwise_affine=True, bias=True)
  (proj): Lin...
                  │                       └ GraphModule()
                  └ OutputGraph(local_scope={'self': Block(
  (norm): LayerNorm((256,), eps=1e-05, elementwise_affine=True, bias=True)
  (proj): Lin...
  File "/storage/home/huyvvo/.local/share/mamba/envs/fairvit-py312-ptnightly-xformers-20260404/lib/python3.12/site-packages/torch/_dynamo/output_graph.py", line 2867, in call_user_compiler
    return self._call_user_compiler(gm, example_inputs)
           │                        │   └ [tensor([[[ 0.1855,  0.7305,  0.9844,  ...,  0.4453, -0.3379, -0.1377],
         [-0.3633,  0.2051, -0.6250,  ..., -2.2812,  0.5...
           │                        └ GraphModule()
           └ OutputGraph(local_scope={'self': Block(
  (norm): LayerNorm((256,), eps=1e-05, elementwise_affine=True, bias=True)
  (proj): Lin...
  File "/storage/home/huyvvo/.local/share/mamba/envs/fairvit-py312-ptnightly-xformers-20260404/lib/python3.12/site-packages/torch/_dynamo/output_graph.py", line 2925, in _call_user_compiler
    compiled_fn = compiler_fn(gm, example_inputs)
                  │           │   └ [tensor([[[ 0.1855,  0.7305,  0.9844,  ...,  0.4453, -0.3379, -0.1377],
         [-0.3633,  0.2051, -0.6250,  ..., -2.2812,  0.5...
                  │           └ GraphModule()
                  └ <torch._dynamo.repro.after_dynamo.WrapBackendDebug object at 0x7c32b7ec4260>
  File "/storage/home/huyvvo/.local/share/mamba/envs/fairvit-py312-ptnightly-xformers-20260404/lib/python3.12/site-packages/torch/_dynamo/repro/after_dynamo.py", line 156, in __call__
    compiled_gm = compiler_fn(gm, example_inputs)
                  │           │   └ [tensor([[[ 0.1855,  0.7305,  0.9844,  ...,  0.4453, -0.3379, -0.1377],
         [-0.3633,  0.2051, -0.6250,  ..., -2.2812,  0.5...
                  │           └ GraphModule()
                  └ functools.partial(<torch._TorchCompileInductorWrapper object at 0x7c34f74431d0>)
  File "/storage/home/huyvvo/.local/share/mamba/envs/fairvit-py312-ptnightly-xformers-20260404/lib/python3.12/site-packages/torch/__init__.py", line 2477, in __call__
    return compile_fx(model_, inputs_, config_patches=all_patches)
           │          │       │                       └ {}
           │          │       └ [tensor([[[ 0.1855,  0.7305,  0.9844,  ...,  0.4453, -0.3379, -0.1377],
         [-0.3633,  0.2051, -0.6250,  ..., -2.2812,  0.5...
           │          └ GraphModule()
           └ <function compile_fx at 0x7c32b56f80e0>
  File "/storage/home/huyvvo/.local/share/mamba/envs/fairvit-py312-ptnightly-xformers-20260404/lib/python3.12/site-packages/torch/_inductor/compile_fx.py", line 2682, in compile_fx
    return _maybe_wrap_and_compile_fx_main(
  File "/storage/home/huyvvo/.local/share/mamba/envs/fairvit-py312-ptnightly-xformers-20260404/lib/python3.12/site-packages/torch/_inductor/compile_fx.py", line 2760, in _maybe_wrap_and_compile_fx_main
    return _compile_fx_main(
  File "/storage/home/huyvvo/.local/share/mamba/envs/fairvit-py312-ptnightly-xformers-20260404/lib/python3.12/site-packages/torch/_inductor/compile_fx.py", line 2972, in _compile_fx_main
    raise e.remove_dynamo_frames() from None  # see TORCHDYNAMO_VERBOSE=1
  File "/storage/home/huyvvo/.local/share/mamba/envs/fairvit-py312-ptnightly-xformers-20260404/lib/python3.12/site-packages/torch/_inductor/compile_fx.py", line 2957, in _compile_fx_main
    return aot_autograd(
  File "/storage/home/huyvvo/.local/share/mamba/envs/fairvit-py312-ptnightly-xformers-20260404/lib/python3.12/site-packages/torch/_dynamo/backends/common.py", line 123, in __call__
    cg = aot_module_simplified(gm, example_inputs, **self.kwargs)
         │                     │   │                 └ <torch._dynamo.backends.common.AotAutograd object at 0x7c32b546dbe0>
         │                     │   └ [tensor([[[ 0.1855,  0.7305,  0.9844,  ...,  0.4453, -0.3379, -0.1377],
         [-0.3633,  0.2051, -0.6250,  ..., -2.2812,  0.5...
         │                     └ GraphModule()
         └ <function aot_module_simplified at 0x7c32b67244a0>
  File "/storage/home/huyvvo/.local/share/mamba/envs/fairvit-py312-ptnightly-xformers-20260404/lib/python3.12/site-packages/torch/_functorch/aot_autograd.py", line 1147, in aot_module_simplified
    compiled_fn, _ = aot_stage2_compile(
  File "/storage/home/huyvvo/.local/share/mamba/envs/fairvit-py312-ptnightly-xformers-20260404/lib/python3.12/site-packages/torch/_functorch/_aot_autograd/graph_compile.py", line 385, in aot_stage2_compile
    return aot_stage2_autograd(aot_state, aot_graph_capture)
           │                   │          └ AOTGraphCapture(wrappers=[AOTDedupeWrapper(keep_arg_mask=[], add_dupe_map=[], old_input_metadata=[], needs_post_compile=False), ...
           │                   └ AOTState(needs_autograd=True, flat_args=[FakeTensor(..., device='cuda:0', size=(16, 197, 256), dtype=torch.bfloat16), Parameter(...
           └ <function aot_stage2_autograd at 0x7c32b7020540>
  File "/storage/home/huyvvo/.local/share/mamba/envs/fairvit-py312-ptnightly-xformers-20260404/lib/python3.12/site-packages/torch/_functorch/_aot_autograd/graph_compile.py", line 2219, in aot_stage2_autograd
    fwd_output_strides, compiled_fw_func = _aot_stage2b_fw_compile(
  File "/storage/home/huyvvo/.local/share/mamba/envs/fairvit-py312-ptnightly-xformers-20260404/lib/python3.12/site-packages/torch/_functorch/_aot_autograd/graph_compile.py", line 2008, in _aot_stage2b_fw_compile
    return _aot_stage2b_compile_forward_or_inference(
  File "/storage/home/huyvvo/.local/share/mamba/envs/fairvit-py312-ptnightly-xformers-20260404/lib/python3.12/site-packages/torch/_functorch/_aot_autograd/graph_compile.py", line 2536, in _aot_stage2b_compile_forward_or_inference
    compiled_fw_func = compiler(fw_module, adjusted_flat_args)
                       │        │          └ [FakeTensor(..., device='cuda:0', size=(16, 197, 256), dtype=torch.bfloat16), FakeTensor(..., device='cuda:0', size=(256,), dtyp...
                       │        └ GraphModule()
                       └ <torch._functorch._aot_autograd.schemas.SerializableAOTDispatchCompiler object at 0x7c32b546f9b0>
  File "/storage/home/huyvvo/.local/share/mamba/envs/fairvit-py312-ptnightly-xformers-20260404/lib/python3.12/site-packages/torch/_functorch/_aot_autograd/schemas.py", line 1416, in __call__
    output_code = self.compiler_fn(gm, example_inputs)
                  │                │   └ [FakeTensor(..., device='cuda:0', size=(16, 197, 256), dtype=torch.bfloat16), FakeTensor(..., device='cuda:0', size=(256,), dtyp...
                  │                └ GraphModule()
                  └ <torch._functorch._aot_autograd.schemas.SerializableAOTDispatchCompiler object at 0x7c32b546f9b0>
  File "/storage/home/huyvvo/.local/share/mamba/envs/fairvit-py312-ptnightly-xformers-20260404/lib/python3.12/site-packages/torch/_inductor/compile_fx.py", line 2820, in fw_compiler_base
    return compile_fx_forward(
  File "/storage/home/huyvvo/.local/share/mamba/envs/fairvit-py312-ptnightly-xformers-20260404/lib/python3.12/site-packages/torch/_inductor/compile_fx.py", line 2461, in compile_fx_forward
    result = inner_compile(
  File "/storage/home/huyvvo/.local/share/mamba/envs/fairvit-py312-ptnightly-xformers-20260404/lib/python3.12/site-packages/torch/_inductor/compile_fx.py", line 826, in compile_fx_inner
    return wrap_compiler_debug(_compile_fx_inner, compiler_name="inductor")(
  File "/storage/home/huyvvo/.local/share/mamba/envs/fairvit-py312-ptnightly-xformers-20260404/lib/python3.12/site-packages/torch/_dynamo/repro/after_aot.py", line 309, in debug_wrapper
    inner_compiled_fn = compiler_fn(gm, example_inputs)
                        │           │   └ [FakeTensor(..., device='cuda:0', size=(16, 197, 256), dtype=torch.bfloat16), FakeTensor(..., device='cuda:0', size=(256,), dtyp...
                        │           └ GraphModule()
                        └ functools.partial(<function _compile_fx_inner at 0x7c32b56f2de0>, get_decomp_fn=<function select_decomp_table at 0x7c32b5df3ce0>...
  File "/storage/home/huyvvo/.local/share/mamba/envs/fairvit-py312-ptnightly-xformers-20260404/lib/python3.12/site-packages/torch/_inductor/compile_fx.py", line 1053, in _compile_fx_inner
    raise InductorError(e, currentframe()).with_traceback(
  File "/storage/home/huyvvo/.local/share/mamba/envs/fairvit-py312-ptnightly-xformers-20260404/lib/python3.12/site-packages/torch/_inductor/compile_fx.py", line 1037, in _compile_fx_inner
    mb_compiled_graph = fx_codegen_and_compile(
  File "/storage/home/huyvvo/.local/share/mamba/envs/fairvit-py312-ptnightly-xformers-20260404/lib/python3.12/site-packages/torch/_inductor/compile_fx.py", line 1802, in fx_codegen_and_compile
    return scheme.codegen_and_compile(gm, example_inputs, inputs_to_check, graph_kwargs)
           │                          │   │               │                └ {'get_decomp_fn': <function select_decomp_table at 0x7c32b5df3ce0>, 'static_input_idxs': [1, 2, 3, 4, 5], 'cudagraphs': BoxedBoo...
           │                          │   │               └ [0]
           │                          │   └ [FakeTensor(..., device='cuda:0', size=(16, 197, 256), dtype=torch.bfloat16), FakeTensor(..., device='cuda:0', size=(256,), dtyp...
           │                          └ GraphModule()
           └ <torch._inductor.compile_fx._InProcessFxCompile object at 0x7c32aeaea720>
  File "/storage/home/huyvvo/.local/share/mamba/envs/fairvit-py312-ptnightly-xformers-20260404/lib/python3.12/site-packages/torch/_inductor/compile_fx.py", line 1489, in codegen_and_compile
    graph.run(*example_inputs)
    │          └ [FakeTensor(..., device='cuda:0', size=(16, 197, 256), dtype=torch.bfloat16), FakeTensor(..., device='cuda:0', size=(256,), dtyp...
    └ <torch._inductor.graph.GraphLowering object at 0x7c32aeaea150>
  File "/storage/home/huyvvo/.local/share/mamba/envs/fairvit-py312-ptnightly-xformers-20260404/lib/python3.12/site-packages/torch/_inductor/graph.py", line 1051, in run
    return super().run(*args)
                        └ (FakeTensor(..., device='cuda:0', size=(16, 197, 256), dtype=torch.bfloat16), FakeTensor(..., device='cuda:0', size=(256,), dtyp...
  File "/storage/home/huyvvo/.local/share/mamba/envs/fairvit-py312-ptnightly-xformers-20260404/lib/python3.12/site-packages/torch/fx/interpreter.py", line 200, in run
    self.env[node] = self.run_node(node)
    │        │       │             └ index_add
    │        │       └ <torch._inductor.graph.GraphLowering object at 0x7c32aeaea150>
    │        └ index_add
    └ <torch._inductor.graph.GraphLowering object at 0x7c32aeaea150>
  File "/storage/home/huyvvo/.local/share/mamba/envs/fairvit-py312-ptnightly-xformers-20260404/lib/python3.12/site-packages/torch/_inductor/graph.py", line 1896, in run_node
    result = super().run_node(n)
                              └ index_add
  File "/storage/home/huyvvo/.local/share/mamba/envs/fairvit-py312-ptnightly-xformers-20260404/lib/python3.12/site-packages/torch/fx/interpreter.py", line 297, in run_node
    return getattr(self, n.op)(n.target, args, kwargs)
                   │     │     │         │     └ {}
                   │     │     │         └ (TensorBox(StorageBox(
  InputBuffer(name='primals_1', layout=FixedLayout('cuda:0', torch.bfloat16, size=[16, 197, 256], stride=...
                   │     │     └ index_add
                   │     └ index_add
                   └ <torch._inductor.graph.GraphLowering object at 0x7c32aeaea150>
  File "/storage/home/huyvvo/.local/share/mamba/envs/fairvit-py312-ptnightly-xformers-20260404/lib/python3.12/site-packages/torch/_inductor/graph.py", line 1336, in call_function
    make_fallback(target, warn=False, get_decomp_fn=self.get_decomp_fn)
    │             │                                 └ <torch._inductor.graph.GraphLowering object at 0x7c32aeaea150>
    │             └ <OpOverload(op='aten.index_add', overload='default')>
    └ <function make_fallback at 0x7c32b5dfe200>
  File "/storage/home/huyvvo/.local/share/mamba/envs/fairvit-py312-ptnightly-xformers-20260404/lib/python3.12/site-packages/torch/_inductor/lowering.py", line 2511, in make_fallback
    assert op not in check_decomps or override_decomp, (
torch._inductor.exc.InductorError: AssertionError: both a fallback and a decomp for same op: aten.index_add.default