[Bug][TLE-Raw][nvshmem] NVSHMEM JIT post-compile hook incorrectly applies nvshmemx_cumodule_init to plain Triton kernels

[Fucov]

Description

When the NVSHMEM TLE Raw CUDA dialect is registered, a plain Triton kernel that does not contain any NVSHMEM calls may fail during its first JIT compilation with:

Cuda failure 'named symbol not found'
AssertionError: nvshmemx_cumodule_init failed: 7

The issue appears to be related to the NVSHMEM jit_post_compile_hook installed by TLE Raw.

After registering an NVSHMEM dialect, the hook seems to be applied globally to subsequently compiled Triton CUDA modules. As a result, nvshmemx_cumodule_init() is also invoked for ordinary Triton kernels whose modules do not contain the NVSHMEM device symbols expected by NVSHMEM.

A minimal reproducer is provided below.

---

Environment

GPU: NVIDIA H20, 2 GPUs used
FlagTree commit: <COMMIT_HASH>

Python: 3.12
CUDA Runtime: 13.0
CUDA Driver: 13.2

NVSHMEM: 3.4.5
NCCL: 2.28.9+cuda13.0

Relevant NVSHMEM initialization output:

NVSHMEM v3.4.5
NVSHMEM configuration:
  CUDA API                     13000
  CUDA Runtime                 13000
  CUDA Driver                  13020
  NVSHMEM_NVTX=ON
  NVSHMEM_USE_NCCL=ON

Both NVSHMEM ranks initialize successfully before the failure.

---

Minimal Reproducer

The following program does two things:

1. Registers an NVSHMEM TLE Raw dialect.
2. JIT-compiles a completely ordinary Triton add kernel.

The Triton kernel itself contains no tle_raw.call() and no NVSHMEM operation.

import os
from pathlib import Path

import torch
import torch.distributed as dist
import triton
import triton.language as tl

from triton.experimental.tle.raw import dialect
from triton.experimental.tle.raw.nvshmem.utils import (
    init_nvshmem_by_torch_pg,
    init_torch_distributed,
    load_common_host,
)


# Registering this dialect installs the NVSHMEM CUDA runtime hook.
# The function is intentionally never called by the Triton kernel below.
@dialect(
    name="cuda",
    library="nvshmem",
    compiler="clang",
    file=Path(__file__).parent / "ring-lab-device.cu",
    extern_func_name="ring_put_device",
)
def unused_nvshmem_device_function(*args, **kwargs):
    ...


# Completely plain Triton kernel:
# no TLE Raw call and no NVSHMEM operation.
@triton.jit
def plain_add_kernel(
    x_ptr,
    y_ptr,
    out_ptr,
    N: tl.constexpr,
):
    offsets = tl.arange(0, N)

    x = tl.load(x_ptr + offsets)
    y = tl.load(y_ptr + offsets)

    tl.store(out_ptr + offsets, x + y)


def main():
    rank = int(os.environ["RANK"])
    local_rank = int(os.environ["LOCAL_RANK"])

    torch.cuda.set_device(local_rank)

    group = init_torch_distributed()

    common = load_common_host()
    init_nvshmem_by_torch_pg(common, group)

    dist.barrier(group=group)

    N = 256

    x = torch.ones(
        N,
        device=f"cuda:{local_rank}",
        dtype=torch.float32,
    )
    y = torch.full(
        (N,),
        2.0,
        device=f"cuda:{local_rank}",
        dtype=torch.float32,
    )
    out = torch.empty_like(x)

    dist.barrier(group=group)

    print(
        f"[rank {rank}] launching plain_add_kernel...",
        flush=True,
    )

    # Failure happens during the first JIT compilation here.
    plain_add_kernel[(1,)](
        x,
        y,
        out,
        N=N,
        num_warps=1,
    )

    torch.cuda.synchronize()

    torch.testing.assert_close(
        out,
        torch.full_like(out, 3.0),
    )

    print(f"[rank {rank}] PASS", flush=True)


if __name__ == "__main__":
    main()

A minimal CUDA source can be used only to make the NVSHMEM dialect registration valid:

extern "C" __device__ void ring_put_device(
    float *destination,
    const float *source,
    int nelems) {
}

---

Reproduction Command

Force a fresh Triton cache so that plain_add_kernel is actually JIT-compiled:

export CUDA_VISIBLE_DEVICES=0,1
export GLOO_SOCKET_IFNAME=lo

export TRITON_CACHE_DIR="/tmp/triton-nvshmem-hook-repro"
rm -rf "${TRITON_CACHE_DIR}"
mkdir -p "${TRITON_CACHE_DIR}"

torchrun \
  --nnodes=1 \
  --node-rank=0 \
  --nproc-per-node=2 \
  --master-addr=127.0.0.1 \
  --master-port=29531 \
  repro_nvshmem_hook.py

---

Actual Behavior

NVSHMEM initialization succeeds for both ranks.

When plain_add_kernel is JIT-compiled, the following error is reported:

Cuda failure 'named symbol not found'

The traceback shows that the failure happens in the TLE Raw NVSHMEM JIT post-compile hook:

plain_add_kernel
  -> triton/runtime/jit.py
  -> jit_post_compile_hook
  -> triton/experimental/tle/raw/cuda/runtime.py
  -> nvshmemx_cumodule_init

Representative traceback:

File ".../triton/runtime/jit.py", line 871, in _do_compile
    self._call_hook(
        knobs.runtime.jit_post_compile_hook,
        ...
    )

File ".../triton/experimental/tle/raw/cuda/runtime.py", line 84, in hook
    assert result == 0, f"nvshmemx_cumodule_init failed: {result}"

AssertionError: nvshmemx_cumodule_init failed: 7

The CUDA runtime also prints:

Cuda failure 'named symbol not found'

---

Expected Behavior

A Triton kernel that does not use NVSHMEM should compile and run normally even if another TLE Raw NVSHMEM dialect has been registered in the same Python process.

For the reproducer above:

plain_add_kernel

should simply produce:

out = x + y = 3

on both ranks.

---

Suspected Cause

The NVSHMEM CUDA dialect appears to install a process-wide Triton post-compile hook.

The current behavior seems roughly equivalent to:

register NVSHMEM dialect
        |
        v
install global jit_post_compile_hook
        |
        v
compile any later Triton kernel
        |
        v
nvshmemx_cumodule_init(kernel.module)

This works for Triton modules that actually contain NVSHMEM device code.

However, a plain Triton module does not contain the NVSHMEM device symbols expected by nvshmemx_cumodule_init(), so module initialization fails with:

CUDA_ERROR_NOT_FOUND
named symbol not found

The hook may need to distinguish between:

Triton module using NVSHMEM
        -> nvshmemx_cumodule_init()

plain Triton module
        -> skip NVSHMEM module initialization

rather than applying NVSHMEM module initialization to every subsequently compiled Triton module.

---

Workaround

As a local workaround, plain Triton kernels can be precompiled while temporarily disabling the NVSHMEM post-compile hook, and the original hook can then be restored before compiling NVSHMEM kernels.

For example:

from triton import knobs

saved_hook = knobs.runtime.jit_post_compile_hook

try:
    knobs.runtime.jit_post_compile_hook = lambda *args, **kwargs: None

    plain_add_kernel[(1,)](
        x,
        y,
        out,
        N=256,
    )

    torch.cuda.synchronize()

finally:
    knobs.runtime.jit_post_compile_hook = saved_hook

After doing this, the process can successfully run both the plain Triton compute kernel and the NVSHMEM TLE Raw communication kernel.

This workaround also suggests that the failure is related to the scope of the NVSHMEM JIT post-compile hook rather than the NVSHMEM runtime initialization itself.