Merge branch 'main' into tle_raw_cuda_ipc_allreduce

Author: i3wanna2

.github/workflows/sunrise3.6-build-and-test.yml

@@ -71,7 +71,7 @@ jobs:
           export TRITON_OFFLINE_BUILD=1
           export TRITON_BUILD_PROTON=OFF
           export FLAGTREE_BACKEND=sunrise
-          MAX_JOBS=32 python3 -m pip install . --no-build-isolation -v
+          MAX_JOBS=32 python3 -m pip install . -v
 
       - name: Clear cache if ClearCache label present
         if: steps.check_backend.outputs.should_skip != 'true'

python/setup_tools/utils/sunrise.py

@@ -33,15 +33,15 @@ def configure_llvm(path):
     cache.store(
         file="sunrise_llvm22_dev_release",
         condition=is_sunrise,
-        url="https://baai-cp-web.ks3-cn-beijing.ksyuncs.com/trans/llvm-1fdc1dfa-triton-v3.6.x.tar.gz",
+        url="https://baai-cp-web.ks3-cn-beijing.ksyuncs.com/trans/llvm-9027ac27-triton-v3.6.x.tar.gz",
         pre_hook=lambda: check_env("LLVM_SYSPATH"),
         post_hook=configure_llvm,
     )
     cache.store(
         file="sunriseTritonPlugin.so",
         condition=is_sunrise and not os.environ.get("FLAGTREE_PLUGIN"),
-        url="https://baai-cp-web.ks3-cn-beijing.ksyuncs.com/trans/sunriseTritonPlugin_v0.6.0.tar.gz",
-        md5_digest="f3c65d44",
+        url="https://baai-cp-web.ks3-cn-beijing.ksyuncs.com/trans/sunriseTritonPlugin_v0.6.0.4.tar.gz",
+        md5_digest="4c77b8c0",
     )
 
 

python/test/tle/unit/test_tle_cumsum.py

@@ -27,10 +27,6 @@ def _is_hcu_backend():
 threads_per_warp = get_current_target().warp_size if is_hip() else 32
 
 
-def _is_amd_hip_backend():
-    return is_hip() and not FLAGTREE_BACKEND
-
-
 def _require_cuda():
     try:
         if _is_enflame_backend():
@@ -280,7 +276,6 @@ def test_tle_cumsum_amdgcn_fastpath_regression_guard():
         "Detected predicated ds_write: possible regression to generic path"
 
 
-@pytest.mark.skipif(_is_amd_hip_backend(), reason="requires AMD local-pointer lowering")
 def test_tle_cumsum_helper_preserves_adjacent_sentinel():
     block = 512
     num_warps = block // threads_per_warp
@@ -302,7 +297,6 @@ def test_tle_cumsum_helper_preserves_adjacent_sentinel():
     torch.testing.assert_close(sentinel, expected_sentinel)
 
 
-@pytest.mark.skipif(_is_amd_hip_backend(), reason="requires AMD local-pointer lowering")
 def test_tle_cumsum_scalar_base_addptr_alias_regression():
     block = 512
     num_warps = block // threads_per_warp

python/triton/experimental/tle/raw/cuda/runtime.py

@@ -234,6 +234,7 @@ def __init__(self, fn: Any, file: Path, *args, **kwargs) -> None:
         self.lowered_region_dialect: Final[str] = "llvm"
         self.arg_dialect: Final[str] = "llvm"
         self.source_file: Final[str] = str(file)
+        self.opt_level: Final[str] = kwargs.get("opt_level", "-O3")
 
         if self.library == "nvshmem":
             from triton.experimental.tle.raw.nvshmem.utils import enable_nvshmem_device_bc
@@ -259,7 +260,7 @@ def register_pending_source(self, *, hint: str = "") -> str:
             extern_func_name=self.extern_func_name,
             source=self.code,
             hint=hint,
-            extra={"source_file": self.source_file},
+            extra={"source_file": self.source_file, "opt_level": self.opt_level},
         )
 
     def create_region_by_llvm(self, builder, llvm: str, handles, alias_indices, hint: str = "",
@@ -285,7 +286,7 @@ def make_llvm(self, mlir_context) -> str:
                 "--cuda-device-only",
                 _get_cuda_gpu_arch(),
                 "-emit-llvm",
-                "-O2",
+                self.opt_level,
                 "-S",
                 "-",
                 "-o",
@@ -314,5 +315,6 @@ def read_text(self):
         file=_CudaSourceFile(),
         extern_func_name=entry.get("extern_func_name"),
         deferred=True,
+        opt_level=entry.get("opt_level"),
     )
     return cuda_fn.make_llvm(context)

python/triton/experimental/tle/raw/nvshmem/utils.py

@@ -348,6 +348,18 @@ def set_signal_cuda_ptr(signal_ptr, signal, stream):
     CUDA_CHECK(err)
 
 
+def copy_on_stream(dst_ptr, src_ptr, nbytes, stream):
+    _, cudart = _get_cuda_modules()
+    (err, ) = cudart.cudaMemcpyAsync(
+        dst_ptr,
+        src_ptr,
+        nbytes,
+        cudart.cudaMemcpyKind.cudaMemcpyDefault,
+        stream.cuda_stream,
+    )
+    CUDA_CHECK(err)
+
+
 def print_perf(
     name: str,
     value: float,

python/tutorials/tle/raw/nvshmem/02-allgather-gemm/ag-gemm-device.cu

@@ -1,10 +1,13 @@
-#include <cuda_fp16.h>
 #include <stddef.h>
 #include <stdint.h>
 
 extern "C" __device__ uint64_t nvshmem_signal_wait_until(uint64_t *sig_addr,
                                                          int cmp,
                                                          uint64_t cmp_value);
+extern "C" __device__ void
+nvshmemx_putmem_signal_block(void *dest, const void *source, size_t bytes,
+                             uint64_t *sig_addr, uint64_t signal, int sig_op,
+                             int pe);
 
 enum {
   NVSHMEM_CMP_GE = 5,
@@ -16,7 +19,7 @@ ag_mark_local_ready(__attribute__((address_space(1))) uint64_t *ready,
                     int rank) {
   if (threadIdx.x == 0) {
     __threadfence_system();
-    ready[(size_t)rank] = 1;
+    ready[rank] = 1;
   }
   __syncthreads();
 }
@@ -29,3 +32,13 @@ ag_wait_ready(__attribute__((address_space(1))) uint64_t *ready,
   }
   __syncthreads();
 }
+
+extern "C" __device__ __attribute__((always_inline)) void
+ag_putmem_signal_block(__attribute__((address_space(1))) void *dest,
+                       __attribute__((address_space(1))) const void *source,
+                       size_t bytes,
+                       __attribute__((address_space(1))) uint64_t *sig_addr,
+                       uint64_t signal, int peer) {
+  nvshmemx_putmem_signal_block(dest, source, bytes, sig_addr, signal,
+                               NVSHMEM_SIGNAL_SET, peer);
+}

python/tutorials/tle/raw/nvshmem/02-allgather-gemm/ag-gemm-host.cu

@@ -1,10 +1,8 @@
-#include <cuda_fp16.h>
 #include <cuda_runtime.h>
 #include <nvshmem.h>
 #include <nvshmemx.h>
 #include <stdint.h>
 #include <stdio.h>
-#include <stdlib.h>
 
 #define CUDA_CHECK(stmt)                                                       \
   do {                                                                         \
@@ -16,10 +14,12 @@
     }                                                                          \
   } while (0)
 
-extern "C" int ag_gemm_workspace_create(int elements_per_rank, void **workspace,
-                                        uint64_t **ready, int *mype, int *npes,
-                                        int *mype_in_node, int *npes_in_node) {
-  if (elements_per_rank <= 0 || workspace == nullptr || ready == nullptr) {
+extern "C" int ag_gemm_workspace_create(int elements_per_rank, int element_size,
+                                        void **workspace, uint64_t **ready,
+                                        int *mype, int *npes, int *mype_in_node,
+                                        int *npes_in_node) {
+  if (elements_per_rank <= 0 || element_size <= 0 || workspace == nullptr ||
+      ready == nullptr) {
     return -1;
   }
 
@@ -29,7 +29,7 @@ extern "C" int ag_gemm_workspace_create(int elements_per_rank, void **workspace,
   *npes_in_node = nvshmem_team_n_pes(NVSHMEMX_TEAM_NODE);
   CUDA_CHECK(cudaSetDevice(*mype_in_node));
 
-  size_t workspace_bytes = (size_t)(*npes) * elements_per_rank * sizeof(__half);
+  size_t workspace_bytes = (size_t)(*npes) * elements_per_rank * element_size;
   *workspace = nvshmem_malloc(workspace_bytes);
   *ready = (uint64_t *)nvshmem_calloc((size_t)(*npes), sizeof(uint64_t));
   if (*workspace == nullptr || *ready == nullptr) {
@@ -61,3 +61,13 @@ extern "C" void *ag_gemm_peer_workspace_ptr(void *workspace, int peer) {
 extern "C" uint64_t *ag_gemm_peer_ready_ptr(uint64_t *ready, int peer) {
   return (uint64_t *)nvshmem_ptr(ready, peer);
 }
+
+extern "C" void ag_gemm_barrier_all_on_stream(cudaStream_t stream) {
+  nvshmemx_barrier_all_on_stream(stream);
+}
+
+extern "C" void ag_gemm_signal_wait_until_on_stream(uint64_t *signal,
+                                                    uint64_t value,
+                                                    cudaStream_t stream) {
+  nvshmemx_signal_wait_until_on_stream(signal, NVSHMEM_CMP_GE, value, stream);
+}