Merge branch 'main' into feature/tle_fix_extract_tile

Author: lzllx123

.github/new-prs-labeler.yml

@@ -30,7 +30,6 @@ CORE:
           - 'cmake/**'
           - 'include/**'
           - 'lib/**'
-          - 'python/**'
           - 'scripts/**'
           - 'test/**'
           - 'unittest/**'
@@ -45,6 +44,19 @@ CORE:
           - 'setup.py'
           - '.github/actions/**'
           - '.github/workflows/code-format-check.yml'
+      - all-globs-to-any-file:
+          - 'python/**'
+          - '!python/setup_tools/utils/cambricon.py'
+          - '!python/setup_tools/utils/enflame.py'
+          - '!python/setup_tools/utils/hcu.py'
+          - '!python/setup_tools/utils/iluvatar.py'
+          - '!python/setup_tools/utils/metax.py'
+          - '!python/setup_tools/utils/mthreads.py'
+          - '!python/setup_tools/utils/ppu.py'
+          - '!python/setup_tools/utils/rpu.py'
+          - '!python/setup_tools/utils/sunrise.py'
+          - '!python/setup_tools/utils/thrive.py'
+          - '!python/setup_tools/utils/xpu.py'
 
 # ==========================================
 # 2. CI/CD Labels
@@ -128,6 +140,7 @@ cambricon:
           - 'third_party/cambricon/**'
           - '.github/workflows/cambricon*'
           - '.github/workflows/benchmark/cambricon/**'
+          - 'python/setup_tools/utils/cambricon.py'
 
 cpu:
   - changed-files:
@@ -142,40 +155,45 @@ enflame:
           - 'third_party/enflame/**'
           - '.github/workflows/enflame*'
           - '.github/workflows/benchmark/enflame/**'
+          - 'python/setup_tools/utils/enflame.py'
 
 hcu:
   - changed-files:
       - any-glob-to-any-file:
           - 'third_party/hcu/**'
           - '.github/workflows/hcu*'
           - '.github/workflows/benchmark/hcu/**'
+          - 'python/setup_tools/utils/hcu.py'
 
 iluvatar:
   - changed-files:
       - any-glob-to-any-file:
           - 'third_party/iluvatar/**'
           - '.github/workflows/iluvatar*'
           - '.github/workflows/benchmark/iluvatar/**'
+          - 'python/setup_tools/utils/iluvatar.py'
 
 metax:
   - changed-files:
       - any-glob-to-any-file:
           - 'third_party/metax/**'
           - '.github/workflows/metax*'
           - '.github/workflows/benchmark/metax/**'
+          - 'python/setup_tools/utils/metax.py'
 
 mthreads:
   - changed-files:
       - any-glob-to-any-file:
           - 'third_party/mthreads/**'
           - '.github/workflows/mthreads*'
           - '.github/workflows/benchmark/mthreads/**'
+          - 'python/setup_tools/utils/mthreads.py'
 
 nvidia:
   - changed-files:
       - any-glob-to-any-file:
           - 'third_party/nvidia/**'
-          - '.github/workflows/nv*'
+          - '.github/workflows/nvidia*'
           - '.github/workflows/benchmark/nvidia/**'
 
 ppu:
@@ -184,27 +202,31 @@ ppu:
           - 'third_party/ppu/**'
           - '.github/workflows/ppu*'
           - '.github/workflows/benchmark/ppu/**'
+          - 'python/setup_tools/utils/ppu.py'
 
 rpu:
   - changed-files:
       - any-glob-to-any-file:
           - 'third_party/rpu/**'
           - '.github/workflows/rpu*'
           - '.github/workflows/benchmark/rpu/**'
+          - 'python/setup_tools/utils/rpu.py'
 
 sunrise:
   - changed-files:
       - any-glob-to-any-file:
           - 'third_party/sunrise/**'
           - '.github/workflows/sunrise*'
           - '.github/workflows/benchmark/sunrise/**'
+          - 'python/setup_tools/utils/sunrise.py'
 
 thrive:
   - changed-files:
       - any-glob-to-any-file:
           - 'third_party/thrive/**'
           - '.github/workflows/thrive*'
           - '.github/workflows/benchmark/thrive/**'
+          - 'python/setup_tools/utils/thrive.py'
 
 tileir:
   - changed-files:
@@ -233,6 +255,7 @@ xpu:
           - 'third_party/xpu/**'
           - '.github/workflows/xpu*'
           - '.github/workflows/benchmark/xpu/**'
+          - 'python/setup_tools/utils/xpu.py'
 
 # ====================================================================================
 # 6. TLE Labels

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

@@ -39,7 +39,7 @@ concurrency:
   # (cancel-in-progress: false) so an in-flight board test is never interrupted
   # mid-run (a cancelled smoke can orphan a launch_kernel_runner that wedges
   # /dev/rpu for the next job).
-  group: rpu3.6-board
+  group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }}
   cancel-in-progress: false
 
 # The RPU CI runner is a small x86 host that orchestrates the build/test on a

.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

@@ -222,6 +222,7 @@ def _compile_cuda_host_to_cache(
         f"-L{lib_dir}",
         f"-l:{host_lib.name}",
         "-lnvshmem_device",
+        "-lcuda",
         "-Xlinker",
         "-rpath",
         "-Xlinker",
@@ -313,12 +314,23 @@ def init_nvshmem_by_torch_pg(common, group):
     torch.distributed.barrier(group=group)
 
 
-def tensor_from_pointer(pointer, shape, dtype, device):
+def tensor_from_pointer(
+    pointer: int | ctypes.c_void_p,
+    shape: tuple[int, ...],
+    dtype: torch.dtype,
+    device: torch.device,
+) -> torch.Tensor:
+    """Create a non-owning Torch tensor view over a CUDA allocation."""
+    address = pointer.value if isinstance(pointer, ctypes.c_void_p) else pointer
+    if address is not None and not isinstance(address, int):
+        raise TypeError(f"pointer must be int or ctypes.c_void_p, got {type(pointer).__name__}")
+    if not address:
+        raise ValueError("pointer cannot be null; CUDA memory must be allocated")
     elements = 1
     for extent in shape:
         elements *= extent
     storage = torch._C._construct_storage_from_data_pointer(
-        pointer.value,
+        address,
         device,
         elements * dtype.itemsize,
     )
@@ -336,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/triton/runtime/autotuner.py

@@ -310,7 +310,7 @@ def prune_configs(self, kwargs: Dict) -> List[Config]:
         pruned_configs = copy.deepcopy(self.configs)
         # pruned_configs = self.configs
         if self.early_config_prune:
-            pruned_configs = self.early_config_prune(self.configs, self.nargs, **kwargs)
+            pruned_configs = self.early_config_prune(pruned_configs, self.nargs, **kwargs)
             if not pruned_configs:
                 raise AutotunerError(
                     "No valid autotuner configs after pruning. `early_config_prune` should return at least one config.")

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);
+}