[KMCompiler][TLERaw] Support nvshmem by lizhangyu258 · Pull Request #787 · flagos-ai/FlagTree

lizhangyu258 · 2026-07-14T00:00:50Z

Reuse the dsl_region implementation of tle_raw to implement the nvshmem extension, and use anonymous compilation to defer the linking of CUDA device-side nvshmem calls until after dsl_region_inline.

The user needs to complete three parts when writing an nvshmem program:

*.py
*-device.cu
*-host.cu

run

export NVSHMEM_HOME=/data/zyuli/miniconda3/envs/flagtree_triton_v3.6.x/lib/python3.12/site-packages/nvidia/nvshmem

cd ${NVSHMEM_HOME}/lib
ln -s libnvshmem_host.so.3 libnvshmem_host.so

# 01-simple-shift
torchrun --nproc_per_node=2 python/tutorials/tle/raw/nvshmem/01-simple-shift/simple-shift.py
# 02-ag-gemm
torchrun --nproc_per_node=2 python/tutorials/tle/raw/nvshmem/02-allgather-gemm/ag-gemm.py

performance

config

1 node
8 GPUs, NVIDIA H800
dtype: torch.float16
allgather: copy engine
gemm kernel
- triton-distributed: autotune
- flagtree: use triton-distributed best config
The latency reported here is the average across all ranks.

FlagTree

(m_per_rank, n_per_rank, k)	GPUs	flagtree(ms)	torch(ms)	speedup
(1024, 1024, 1024)	8	0.418	0.257	0.615 x
(4096, 1024, 8192)	8	3.041	3.803	1.251 x
(4096, 8192, 8192)	8	7.043	9.270	1.316 x
(4096, 28672, 8192)	8	23.440	25.165	1.074 x
(4096, 8192, 28672)	8	23.450	32.963	1.406 x
(8192, 1024, 8192)	8	5.868	7.493	1.277 x
(8192, 8192, 8192)	8	13.783	18.052	1.310 x
(8192, 28672, 8192)	8	45.405	55.212	1.216 x
(8192, 8192, 28672)	8	46.205	66.344	1.436 x

Triton-distributed

(m_per_rank, n_per_rank, k)	GPUs	td autotune(ms)	torch(ms)	speedup
(1024, 1024, 1024)	8	0.557	0.326	0.585 x
(4096, 1024, 8192)	8	3.015	3.699	1.227 x
(4096, 8192, 8192)	8	6.368	9.048	1.421 x
(4096, 28672, 8192)	8	22.771	25.241	1.108 x
(4096, 8192, 28672)	8	23.467	32.639	1.391 x
(8192, 1024, 8192)	8	5.949	7.282	1.224 x
(8192, 8192, 8192)	8	13.003	17.979	1.383 x
(8192, 28672, 8192)	8	46.779	52.209	1.116 x
(8192, 8192, 28672)	8	47.092	66.463	1.411 x

i3wanna2 · 2026-07-14T04:41:14Z

An alternative approach is in PR #682: compile TLE-raw NVSHMEM device code to LLVM bitcode (.bc) and invoke it via extern_call. That path is frozen and kept for reference only; official NVSHMEM support will use this PR.

…ctions

lizhangyu258 · 2026-07-16T09:18:59Z

All cases uniformly use torchrun for launching. If you want to use nvshmrun to launch, you need to install it by referring to https://developer.nvidia.com/nvshmem-downloads and then adjust the corresponding case accordingly.

…t registry Avoid pulling cuda-python/nvshmem on import dialect or ordinary CUDA compile paths (fixes nvidia unit/cuda and enflame collect failures). Register cuda backend with try/except like other dialects, and align tops_mlir with extern_func_name API.

Align tops dialect with the multi-function raw call API so tle_raw.call does not AttributeError on Enflame tops tutorials.

Match deferred handling so backends without the attribute still work. Drop redundant extern_func_name fields from tops/tops_mlir; keep the create_region_by_llvm parameter for API compatibility.

Share extern_func_name/deferred and default create_region_by_llvm so cuda/mlir/tops stay aligned without duplicating dialect kwargs.

i3wanna2

LGTM

Discover clang>=20 and pip NVSHMEM when unset, wire knobs for optional overrides, and extend nv3.6 workflow with cuda/nvshmem tutorials including smem.

…ments nvcc rejects versioned .so as a positional input; use -L/-l: and rpath. Add optional nvidia-nvshmem-cu12/cu13 requirement files for tle raw.

NVSHMEM 3.7+ wheels drop unified libnvshmem_device.bc; fall back to libnvshmem_device_sm_XX.bc using the compile arch.

Only link libnvshmem_device when library=nvshmem enables a utils flag, so ordinary CUDA compiles do not inject nvshmem into extern_libs.

@dialect

* [KMCompiler][TLERaw] Fix cuda device-side not supporting multiple functions * [KMCompiler][TLERaw] Support nvshmem * [KMCompiler][TLERaw] Unify launcher to torchrun * [KMCompiler][TLERaw] Lazy-load nvshmem/cuda imports and harden dialect registry Avoid pulling cuda-python/nvshmem on import dialect or ordinary CUDA compile paths (fixes nvidia unit/cuda and enflame collect failures). Register cuda backend with try/except like other dialects, and align tops_mlir with extern_func_name API. * [KMCompiler][TLERaw] Fix extra blank line in nvshmem utils * [KMCompiler][TLERaw] Add extern_func_name to TOPSJITFunction Align tops dialect with the multi-function raw call API so tle_raw.call does not AttributeError on Enflame tops tutorials. * [KMCompiler][TLERaw] Use getattr for extern_func_name in raw call Match deferred handling so backends without the attribute still work. Drop redundant extern_func_name fields from tops/tops_mlir; keep the create_region_by_llvm parameter for API compatibility. * [KMCompiler][TLERaw] Extract RawJITFunction base for dialect backends Share extern_func_name/deferred and default create_region_by_llvm so cuda/mlir/tops stay aligned without duplicating dialect kwargs. * [KMCompiler][TLERaw] Add tle cuda and nvshmem test * [KMCompiler][TLERaw] Replace os.getenv with knobs * [KMCompiler][TLERaw] Auto-resolve clang and NVSHMEM for raw CUDA CI Discover clang>=20 and pip NVSHMEM when unset, wire knobs for optional overrides, and extend nv3.6 workflow with cuda/nvshmem tutorials including smem. * [KMCompiler][TLERaw] Address review: knobs comment and CI clang-22 path * [KMCompiler][TLERaw] Shorten flagtree knobs comment * [KMCompiler][TLERaw] Link nvshmem host via -L/-l and add CUDA requirements nvcc rejects versioned .so as a positional input; use -L/-l: and rpath. Add optional nvidia-nvshmem-cu12/cu13 requirement files for tle raw. * [KMCompiler][TLERaw] Resolve per-SM NVSHMEM device bitcode NVSHMEM 3.7+ wheels drop unified libnvshmem_device.bc; fall back to libnvshmem_device_sm_XX.bc using the compile arch. * [KMCompiler][TLERaw] Gate nvshmem device bc on @dialect library Only link libnvshmem_device when library=nvshmem enables a utils flag, so ordinary CUDA compiles do not inject nvshmem into extern_libs. --------- Co-authored-by: i3wanna2 <15910307812@163.com> Co-authored-by: iwanna-lxy <37344393+i3wanna2@users.noreply.github.com>

@dialect

* [KMCompiler][TLERaw] Fix cuda device-side not supporting multiple functions * [KMCompiler][TLERaw] Support nvshmem * [KMCompiler][TLERaw] Unify launcher to torchrun * [KMCompiler][TLERaw] Lazy-load nvshmem/cuda imports and harden dialect registry Avoid pulling cuda-python/nvshmem on import dialect or ordinary CUDA compile paths (fixes nvidia unit/cuda and enflame collect failures). Register cuda backend with try/except like other dialects, and align tops_mlir with extern_func_name API. * [KMCompiler][TLERaw] Fix extra blank line in nvshmem utils * [KMCompiler][TLERaw] Add extern_func_name to TOPSJITFunction Align tops dialect with the multi-function raw call API so tle_raw.call does not AttributeError on Enflame tops tutorials. * [KMCompiler][TLERaw] Use getattr for extern_func_name in raw call Match deferred handling so backends without the attribute still work. Drop redundant extern_func_name fields from tops/tops_mlir; keep the create_region_by_llvm parameter for API compatibility. * [KMCompiler][TLERaw] Extract RawJITFunction base for dialect backends Share extern_func_name/deferred and default create_region_by_llvm so cuda/mlir/tops stay aligned without duplicating dialect kwargs. * [KMCompiler][TLERaw] Add tle cuda and nvshmem test * [KMCompiler][TLERaw] Replace os.getenv with knobs * [KMCompiler][TLERaw] Auto-resolve clang and NVSHMEM for raw CUDA CI Discover clang>=20 and pip NVSHMEM when unset, wire knobs for optional overrides, and extend nv3.6 workflow with cuda/nvshmem tutorials including smem. * [KMCompiler][TLERaw] Address review: knobs comment and CI clang-22 path * [KMCompiler][TLERaw] Shorten flagtree knobs comment * [KMCompiler][TLERaw] Link nvshmem host via -L/-l and add CUDA requirements nvcc rejects versioned .so as a positional input; use -L/-l: and rpath. Add optional nvidia-nvshmem-cu12/cu13 requirement files for tle raw. * [KMCompiler][TLERaw] Resolve per-SM NVSHMEM device bitcode NVSHMEM 3.7+ wheels drop unified libnvshmem_device.bc; fall back to libnvshmem_device_sm_XX.bc using the compile arch. * [KMCompiler][TLERaw] Gate nvshmem device bc on @dialect library Only link libnvshmem_device when library=nvshmem enables a utils flag, so ordinary CUDA compiles do not inject nvshmem into extern_libs. --------- Co-authored-by: i3wanna2 <15910307812@163.com> Co-authored-by: iwanna-lxy <37344393+i3wanna2@users.noreply.github.com>

lizhangyu258 requested review from Galaxy1458, i3wanna2, menchunlei, sunnycase and zhzhcookie as code owners July 14, 2026 00:00

github-actions Bot added nvidia tle labels Jul 14, 2026

lizhangyu258 marked this pull request as draft July 14, 2026 00:01

github-actions Bot added the main label Jul 14, 2026

lizhangyu258 force-pushed the triton_v3.6.x_support_nvshmem_by_region branch from 48c25a8 to 52da041 Compare July 14, 2026 00:16

[KMCompiler][TLERaw] Fix cuda device-side not supporting multiple fun…

33d73bf

…ctions

lizhangyu258 force-pushed the triton_v3.6.x_support_nvshmem_by_region branch from 52da041 to a3ace68 Compare July 15, 2026 07:30

lizhangyu258 changed the title ~~Triton v3.6.x support nvshmem by region~~ [KMCompiler][TLERaw] Support nvshmem Jul 15, 2026

[KMCompiler][TLERaw] Support nvshmem

23bc351

lizhangyu258 force-pushed the triton_v3.6.x_support_nvshmem_by_region branch from a3ace68 to 23bc351 Compare July 16, 2026 00:28

[KMCompiler][TLERaw] Unify launcher to torchrun

ac3948a

lizhangyu258 force-pushed the triton_v3.6.x_support_nvshmem_by_region branch from 5cb3b98 to ac3948a Compare July 16, 2026 09:26

i3wanna2 marked this pull request as ready for review July 16, 2026 15:09

i3wanna2 added 5 commits July 16, 2026 16:57

[KMCompiler][TLERaw] Fix extra blank line in nvshmem utils

459f691

[KMCompiler][TLERaw] Add extern_func_name to TOPSJITFunction

63af249

Align tops dialect with the multi-function raw call API so tle_raw.call does not AttributeError on Enflame tops tutorials.

[KMCompiler][TLERaw] Use getattr for extern_func_name in raw call

e915385

Match deferred handling so backends without the attribute still work. Drop redundant extern_func_name fields from tops/tops_mlir; keep the create_region_by_llvm parameter for API compatibility.

[KMCompiler][TLERaw] Extract RawJITFunction base for dialect backends

5af081a

Share extern_func_name/deferred and default create_region_by_llvm so cuda/mlir/tops stay aligned without duplicating dialect kwargs.

i3wanna2 previously approved these changes Jul 17, 2026

View reviewed changes

zhzhcookie reviewed Jul 17, 2026

View reviewed changes

Comment thread third_party/nvidia/backend/compiler.py Outdated

lizhangyu258 added 2 commits July 17, 2026 02:46

[KMCompiler][TLERaw] Add tle cuda and nvshmem test

c5f67bd

[KMCompiler][TLERaw] Replace os.getenv with knobs

289ef3d

lizhangyu258 dismissed i3wanna2’s stale review via 289ef3d July 17, 2026 02:58

[KMCompiler][TLERaw] Auto-resolve clang and NVSHMEM for raw CUDA CI

51fa18c

Discover clang>=20 and pip NVSHMEM when unset, wire knobs for optional overrides, and extend nv3.6 workflow with cuda/nvshmem tutorials including smem.

zhzhcookie reviewed Jul 17, 2026

View reviewed changes

Comment thread python/triton/knobs.py

zhzhcookie reviewed Jul 17, 2026

View reviewed changes

Comment thread .github/workflows/nv3.6-build-and-test.yml

i3wanna2 added 2 commits July 17, 2026 06:00

[KMCompiler][TLERaw] Address review: knobs comment and CI clang-22 path

5afb5e2

[KMCompiler][TLERaw] Shorten flagtree knobs comment

f4a6389

i3wanna2 self-requested a review July 17, 2026 06:10

[KMCompiler][TLERaw] Link nvshmem host via -L/-l and add CUDA require…

81e5678

…ments nvcc rejects versioned .so as a positional input; use -L/-l: and rpath. Add optional nvidia-nvshmem-cu12/cu13 requirement files for tle raw.

zhzhcookie reviewed Jul 17, 2026

View reviewed changes

Comment thread third_party/nvidia/backend/compiler.py

i3wanna2 and others added 3 commits July 17, 2026 11:38

[KMCompiler][TLERaw] Resolve per-SM NVSHMEM device bitcode

5e7c499

NVSHMEM 3.7+ wheels drop unified libnvshmem_device.bc; fall back to libnvshmem_device_sm_XX.bc using the compile arch.

[KMCompiler][TLERaw] Gate nvshmem device bc on @dialect library

acfda82

Only link libnvshmem_device when library=nvshmem enables a utils flag, so ordinary CUDA compiles do not inject nvshmem into extern_libs.

Merge branch 'main' into triton_v3.6.x_support_nvshmem_by_region

4e3a407

i3wanna2 merged commit 98ad17c into flagos-ai:main Jul 18, 2026
15 of 22 checks passed

Zhang-kg mentioned this pull request Jul 21, 2026

[BUG] SMEM atomic (atom.shared) deadlocks inside a warp_specialize worker partition #828

Closed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

[KMCompiler][TLERaw] Support nvshmem#787

[KMCompiler][TLERaw] Support nvshmem#787
i3wanna2 merged 17 commits into
flagos-ai:mainfrom
lizhangyu258:triton_v3.6.x_support_nvshmem_by_region

lizhangyu258 commented Jul 14, 2026 •

edited

Loading

Uh oh!

i3wanna2 commented Jul 14, 2026

Uh oh!

lizhangyu258 commented Jul 16, 2026

Uh oh!

i3wanna2 left a comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Uh oh!

Conversation

lizhangyu258 commented Jul 14, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

run

performance

config

FlagTree

Triton-distributed

Uh oh!

i3wanna2 commented Jul 14, 2026

Uh oh!

lizhangyu258 commented Jul 16, 2026

Uh oh!

i3wanna2 left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

lizhangyu258 commented Jul 14, 2026 •

edited

Loading