Skip to content

[fix]: fix compile error in torch.compile(fn) when fn is a wrapper of a triton kernel#680

Open
small-cat wants to merge 1 commit into
flagos-ai:triton_v3.2.xfrom
small-cat:fix/compatible-torch-compile
Open

[fix]: fix compile error in torch.compile(fn) when fn is a wrapper of a triton kernel#680
small-cat wants to merge 1 commit into
flagos-ai:triton_v3.2.xfrom
small-cat:fix/compatible-torch-compile

Conversation

@small-cat

Copy link
Copy Markdown
Collaborator

Description
When using torch.compile to compile a graph in vllm-ascend, TorchDynamo throws an exception due to a FlagTree-related issue.

Reproduction
Run the following test case (python test.py):

import torch
import triton
import triton.language as tl
import torch_npu
import torch_npu._inductor

device="npu"

@triton.heuristics({"N_ELEMS": lambda args: args["n_elements"]}) # just a test to reproduce the problem
@triton.jit
def add_kernel(x_ptr,  # *Pointer* to first input vector.
               y_ptr,  # *Pointer* to second input vector.
               output_ptr,  # *Pointer* to output vector.
               n_elements,  # Size of the vector.
               N_ELEMS: tl.constexpr,  # Number of elements each program should process.
               BLOCK_SIZE: tl.constexpr,  # Number of elements each program should process.
               # NOTE: `constexpr` so it can be used as a shape value.
               ):
    # There are multiple 'programs' processing different data. We identify which program
    # we are here:
    pid = tl.program_id(axis=0)  # We use a 1D launch grid so axis is 0.
    # This program will process inputs that are offset from the initial data.
    # For instance, if you had a vector of length 256 and block_size of 64, the programs
    # would each access the elements [0:64, 64:128, 128:192, 192:256].
    # Note that offsets is a list of pointers:
    block_start = pid * BLOCK_SIZE
    offsets = block_start + tl.arange(0, BLOCK_SIZE)
    # Create a mask to guard memory operations against out-of-bounds accesses.
    mask = offsets < N_ELEMS
    # Load x and y from DRAM, masking out any extra elements in case the input is not a
    # multiple of the block size.
    x = tl.load(x_ptr + offsets, mask=mask)
    y = tl.load(y_ptr + offsets, mask=mask)
    output = x + y
    # Write x + y back to DRAM.
    tl.store(output_ptr + offsets, output, mask=mask)


def add(x: torch.Tensor, y: torch.Tensor):
    # We need to preallocate the output.
    output = torch.empty_like(x)
    n_elements = output.numel()
    grid = lambda meta: (triton.cdiv(n_elements, meta["BLOCK_SIZE"]), )
    add_kernel[grid](x, y, output, n_elements, BLOCK_SIZE=1024)
    return output.cpu()


def test_vector_add():
    print(add_kernel)
    print(type(add_kernel))
    print(type(add_kernel.fn))
    print(type(add_kernel.fn.fn))

    torch.manual_seed(0)
    size = 4432
    x = torch.rand(size, device=device)
    y = torch.rand(size, device=device)
    output_torch = x.cpu() + y.cpu()

    # output_triton = add(x, y)
    add_fn = torch.compile(add, fullgraph=True)
    output_triton = add_fn(x, y)

    print(f"The maximum difference between torch and triton is "
          f"{torch.max(torch.abs(output_torch - output_triton))}")
    assert torch.allclose(output_triton, output_torch), (output_triton, output_torch)


if __name__ == "__main__":
    test_vector_add()

error

  File "/usr/local/python3.11.14/lib/python3.11/site-packages/torch/_dynamo/eval_frame.py", line 841, in compile_wrapper
    raise e.with_traceback(None) from e.__cause__  # User compiler error
    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
torch._dynamo.exc.Unsupported: Attempted to inline function marked as skipped
  Explanation: Dynamo developers have intentionally marked that the function `KernelInterface.__getitem__` should not be traced.
  Hint: Avoid calling the function `KernelInterface.__getitem__`.
  Hint: Apply `@torch._dynamo.dont_skip_tracing` to the function `KernelInterface.__getitem__` to force tracing into the function. More graph breaks may occur as a result of attempting to trace into the function.
  Hint: Please file an issue to PyTorch.

  Developer debug context: qualname: KernelInterface.__getitem__, name: __getitem__, filename: `/docker_shared/worktrees/FlagTree-Hotfix/python/triton/backends/ascend/spec/triton/runtime/jit.py`, skip reason: skipped according trace_rules.lookup SKIP_DIRS

Root Cause
The Triton Ascend backend implementation resides in triton/backends/ascend/spec. TorchDynamo's SKIP_DIRS rule marks the entire triton/backends directory as skipped, causing all definitions under it (including the required Ascend backend code) to be ignored. This leads to the unsupported inline error.

Solution
To improve FlagTree's compatibility with torch.compile, modify the build process to install triton/backends/ascend/spec into triton/ascend instead. This avoids the problematic SKIP_DIRS rule while keeping the code accessible.

@small-cat

Copy link
Copy Markdown
Collaborator Author

error in CI, but it is strange, because I can't find any code about triton_enable_libdevice_simt from the error message in branch triton_v3.2.x

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant