[BUG] TLE WS thread budget exceeds CUDA function thread limit with 3 warp-specialized roles

[Zhang-kg]

Describe the bug

Summary

This issue adds two self-contained repro cases for the TLE warp-specialization thread-budget issue, origin code link:

python/test/tle/integration/isolated_repros/oor384_receiver_w4
python/test/tle/integration/isolated_repros/oor384_receiver_w4_allwarps

Both cases reproduce:

triton.runtime.errors.OutOfResources:
out of resource: threads, Required: 384, Hardware limit: 256.

The repros are intentionally isolated. They do not depend on the previous external megakernel-moe experiment directory.

Environment

Validated locally on:

GPU: NVIDIA H100 80GB HBM3
CUDA: 12.8
Python: 3.10
NVSHMEM: 3.4.5 (nvidia-nvshmem-cu12==3.4.5)
MPI launcher: Open MPI 4.1.2
Triton: FlagTree PR682-based Triton with TLE raw NVSHMEM support

Repro Cases

1. oor384_receiver_w4

This case configures:

default / dispatch partition: 4 warps
receiver worker partition:   4 warps
compute worker partition:    4 warps
total:                       12 warps = 384 threads

The receiver partition is allocated 4 warps, but only warp_id == 0 executes the receiver logic.

Result:

Required: 384, Hardware limit: 256

2. oor384_receiver_w4_allwarps

This case removes the warp_id == 0 guard and lets the 4 receiver warps participate in receiver work distribution.

Result is still:

Required: 384, Hardware limit: 256

This shows that the OOR is independent of whether receiver internally uses only warp0 or all 4 receiver warps. The issue is the
total WS role thread budget:

4 + 4 + 4 warps = 12 warps = 384 threads

while the compiled function reports:

CU_FUNC_ATTRIBUTE_MAX_THREADS_PER_BLOCK = 256

How To Reproduce

Set local environment variables:

export CUDA_HOME=/usr/local/cuda-12.8
export NVSHMEM_HOME=/path/to/nvshmem
export LD_LIBRARY_PATH="$NVSHMEM_HOME/lib:${CUDA_HOME}/lib64:${LD_LIBRARY_PATH:-}"
export CPATH="${CUDA_HOME}/targets/x86_64-linux/include:$NVSHMEM_HOME/include:${CPATH:-}"
export PYTHON_BIN=/path/to/python

Run receiver-w4:

cd python/test/tle/integration/isolated_repros/oor384_receiver_w4

PYTHONNOUSERSITE=1 \
PYTHONDONTWRITEBYTECODE=1 \
PYTHONPATH=/path/to/FlagTree/python:$PWD \
"$PYTHON_BIN" repro_receiver_w4.py

Run receiver-w4-allwarps:

cd python/test/tle/integration/isolated_repros/oor384_receiver_w4_allwarps

PYTHONNOUSERSITE=1 \
PYTHONDONTWRITEBYTECODE=1 \
PYTHONPATH=/path/to/FlagTree/python:$PWD \
NVCC=$PWD/nvcc_flock_wrapper.sh \
TRITON_CACHE_DIR=/tmp/tle_ws_oor384_receiver_w4_allwarps_selfcontained \
"$PYTHON_BIN" repro_receiver_w4_allwarps.py

Expected result for both:

triton.runtime.errors.OutOfResources:
out of resource: threads, Required: 384, Hardware limit: 256.

Notes

This PR does not claim this is a compiler bug by itself. The repro documents a concrete TLE WS thread-budget boundary:

dispatch/default 4 warps + receiver 4 warps + compute 4 warps

currently results in a 384-thread kernel requirement, while the compiled function is limited to 256 threads per block.

Environment details

Environment

GPU: NVIDIA H100 80GB HBM3
CUDA: 12.8
Python: 3.10
NVSHMEM: 3.4.5 (nvidia-nvshmem-cu12==3.4.5)
MPI: Open MPI 4.1.2
Triton: FlagTree PR682-based Triton with TLE raw NVSHMEM support

[i3wanna2]

Could you provide a minimal reproduction that does not rely on distributed features and only uses the existing TLE primitives on the triton_v3.6.x branch? Alternatively, is this issue specific to the combination of TLE and TLE-Raw?

[Zhang-kg]

Could you provide a minimal reproduction that does not rely on distributed features and only uses the existing TLE primitives on the triton_v3.6.x branch? Alternatively, is this issue specific to the combination of TLE and TLE-Raw?

I have isolated a minimal standalone repro here: Zhang-kg@1b1ecc5

The repro does not depend on NVSHMEM, TLE-Raw, mpirun, distributed runtime, raw
extern calls, or any MegaMoE code. It only uses:

• single process / single GPU
• normal CUDA tensor
• tle.gpu.warp_specialize
• regular tl.load / tl.store

I also copied and ran the same repro script in a clean triton_v3.6.x environment. The same error can still be reproduced

[i3wanna2]

Thank you for the repro. This looks like a Triton issue rather than something caused by the TLE extension. In Triton, ttg.warp_specialize adds extra warps on top of the launch num_warps (treating them as a default warpgroup plus separate worker warpgroups), rather than splitting the default warpgroup into different roles. That can introduce additional register pressure compared to a hand-written CUDA kernel, so the same algorithm/config may not map one-to-one from CUDA to Triton. This case can be avoided by using smaller worker warp counts in tle.gpu.warp_specialize — for example [2, 2] in repro_ws_444_no_nvshmem.py. For finer-grained control, TLE Raw offers a lower-level interface when the high-level warp_specialize model is too restrictive.