Drop the rotated-out internal ROCm apt repo before apt-get update

.github/workflows/iris-performance-regression-test.yml

@@ -27,7 +27,9 @@ jobs:
           - example_name: "GEMM All-Scatter WG Specialization"
             example_path: "10_gemm_all_scatter_wg_specialization"
             tflops_threshold: 1440  # Actual: ~1802 TFLOPs (80% regression threshold)
-            benchmark_args: "-m 16384 -n 16384 -k 16384 --BLK_M 128 --BLK_N 128 --BLK_K 64 --gsize_m 6 --gemm_sms 256"
+            # gemm_sms must stay below the CU count -- the remainder is what runs the
+            # communication path. 256 left zero of it on a 256-CU part.
+            benchmark_args: "-m 16384 -n 16384 -k 16384 --BLK_M 128 --BLK_N 128 --BLK_K 64 --gsize_m 6 --gemm_sms 192"
 
           - example_name: "GEMM All-Scatter"
             example_path: "07_gemm_all_scatter"

apptainer/iris.def

@@ -13,6 +13,8 @@ From: rocm/pytorch:rocm7.2.1_ubuntu24.04_py3.14_pytorch_2.10.0
     export PATH=\"\$ROCM_PATH/bin:\$PATH\"
 
     # Install system packages
+    # see docker/Dockerfile -- base image pins a rotated-out internal ROCm apt repo
+    rm -f /etc/apt/sources.list.d/rocm.list && \
     apt-get update && \
     DEBIAN_FRONTEND=noninteractive apt-get install -y \
         git wget ninja-build cmake python3-pip python3-dev build-essential jq libdwarf-dev && \

docker/Dockerfile

@@ -20,7 +20,13 @@ ENV OMPI_ALLOW_RUN_AS_ROOT_CONFIRM=1 \
     OMPI_ALLOW_RUN_AS_ROOT=1
 
 # Install system packages
-RUN apt-get update && \
+# The rocm/pytorch base images ship /etc/apt/sources.list.d/rocm.list pinned to a
+# specific internal build (e.g. compute-rocm-rel-7.1 20) on compute-artifactory.amd.com.
+# Once that build is rotated out upstream the index 404s, apt-get update returns 100,
+# and the whole layer fails -- taking every downstream CI job with it. ROCm and torch
+# are already installed in the image, so nothing here needs that repo.
+RUN rm -f /etc/apt/sources.list.d/rocm.list && \
+    apt-get update && \
     DEBIAN_FRONTEND=noninteractive apt-get install -y \
     git wget ninja-build cmake python3-pip python3-dev build-essential libdwarf-dev && \
     rm -rf /var/lib/apt/lists/*

docker/Dockerfile.ccl

@@ -19,7 +19,13 @@ ENV OMPI_ALLOW_RUN_AS_ROOT_CONFIRM=1 \
     OMPI_ALLOW_RUN_AS_ROOT=1
 
 # Install system packages
-RUN apt-get update && \
+# The rocm/pytorch base images ship /etc/apt/sources.list.d/rocm.list pinned to a
+# specific internal build (e.g. compute-rocm-rel-7.1 20) on compute-artifactory.amd.com.
+# Once that build is rotated out upstream the index 404s, apt-get update returns 100,
+# and the whole layer fails -- taking every downstream CI job with it. ROCm and torch
+# are already installed in the image, so nothing here needs that repo.
+RUN rm -f /etc/apt/sources.list.d/rocm.list && \
+    apt-get update && \
     DEBIAN_FRONTEND=noninteractive apt-get install -y \
     git wget ninja-build cmake python3-pip python3-dev build-essential libdwarf-dev && \
     rm -rf /var/lib/apt/lists/*

examples/10_gemm_all_scatter_wg_specialization/benchmark.py

@@ -93,6 +93,18 @@ def _worker(local_rank: int, world_size: int, init_url: str, args: dict):
     if args["gemm_sms"] is None:
         # For wg_specialized: use next smaller power of 2
         args["gemm_sms"] = 2 ** int(math.log2(cu_count)) if cu_count > 0 else 1
+        # The kernel launches num_sms workgroups and gives whatever is left after
+        # gemm_sms to the communication path. When the CU count is itself a power
+        # of two (256 on MI350X) that leaves nothing, so step down once.
+        if args["gemm_sms"] >= args["num_sms"]:
+            args["gemm_sms"] //= 2
+
+    if args["gemm_sms"] >= args["num_sms"]:
+        raise ValueError(
+            f"gemm_sms ({args['gemm_sms']}) must be less than num_sms ({args['num_sms']}): "
+            "workgroup specialization needs the remainder for the communication path, "
+            "and with none of it the all-scatter never runs."
+        )
 
     # GEMM
     datatype = torch.float32

examples/22_gemm_one_shot_reduce_scatter_wg_specialization/benchmark.py

@@ -89,6 +89,18 @@ def _worker(local_rank: int, world_size: int, init_url: str, args: dict):
     if args["gemm_sms"] is None:
         # Use next smaller power of 2 for GEMM SMs
         args["gemm_sms"] = 2 ** int(math.log2(cu_count)) if cu_count > 0 else 1
+        # The kernel launches num_sms workgroups and gives whatever is left after
+        # gemm_sms to the communication path. When the CU count is itself a power
+        # of two (256 on MI350X) that leaves nothing, so step down once.
+        if args["gemm_sms"] >= args["num_sms"]:
+            args["gemm_sms"] //= 2
+
+    if args["gemm_sms"] >= args["num_sms"]:
+        raise ValueError(
+            f"gemm_sms ({args['gemm_sms']}) must be less than num_sms ({args['num_sms']}): "
+            "workgroup specialization needs the remainder for the communication path, "
+            "and with none of it the reduce-scatter never runs."
+        )
 
     datatype = torch.float16
     if args["datatype"] == "fp16":

iris/mem/utils.py

@@ -13,16 +13,29 @@
 import triton.language as tl
 from triton.language.target_info import is_hip_cdna3, is_hip_cdna4
 
+# Probe each intrinsic separately. They were added to
+# ``triton.language.extra.hip`` at different times, so a build can export one
+# and not the other -- upstream Triton at the pinned commit ships
+# ``memrealtime`` but not ``smid``. Guarding both behind a single try/except
+# meant a missing ``smid`` also disabled ``read_realtime``, and the fallback
+# was ``static_assert(False)``, which fails compilation of every kernel that
+# records a trace event rather than degrading the trace.
 try:
     from triton.language.extra.hip import memrealtime as _memrealtime
+
+    _HAS_MEMREALTIME = True
+except ImportError:
+    _HAS_MEMREALTIME = False
+
+try:
     from triton.language.extra.hip import smid as _smid
 
-    _HAS_HIP_INTRINSICS = True
+    _HAS_SMID = True
 except ImportError:
-    _HAS_HIP_INTRINSICS = False
+    _HAS_SMID = False
 
 
-if _HAS_HIP_INTRINSICS:
+if _HAS_MEMREALTIME:
 
     @triton.jit
     def read_realtime():
@@ -39,6 +52,34 @@ def read_realtime():
             int64: Current timestamp in cycles (100 MHz constant clock)
         """
         return _memrealtime()
+else:
+
+    @triton.jit
+    def read_realtime():
+        """
+        Read GPU wall clock timestamp on builds without ``tl.extra.hip.memrealtime``.
+
+        Emits the same instruction the intrinsic would. gfx11/gfx12 use a
+        message rather than ``s_memrealtime``; there is no portable fallback for
+        those here, so they report 0 and timestamps are simply unavailable.
+
+        Returns:
+            int64: Timestamp in cycles, or 0 where unsupported
+        """
+        if is_hip_cdna3() or is_hip_cdna4():
+            return tl.inline_asm_elementwise(
+                asm="s_memrealtime $0\n\ts_waitcnt vmcnt(0)",
+                constraints=("=s"),
+                args=[],
+                dtype=tl.int64,
+                is_pure=False,
+                pack=1,
+            )
+        else:
+            return tl.cast(0, tl.int64)
+
+
+if _HAS_SMID:
 
     @triton.jit
     def get_cu_id():
@@ -55,17 +96,35 @@ def get_cu_id():
         return _smid()
 else:
 
-    @triton.jit
-    def read_realtime():
-        """Fallback stub when HIP intrinsics are missing."""
-        tl.static_assert(False, "memrealtime is unavailable in this Triton build")
-        return tl.cast(0, tl.int64)
-
     @triton.jit
     def get_cu_id():
-        """Fallback stub when HIP intrinsics are missing."""
-        tl.static_assert(False, "smid is unavailable in this Triton build")
-        return tl.cast(0, tl.int32)
+        """
+        Get compute-unit ID on builds without ``tl.extra.hip.smid``.
+
+        Reads CU_ID out of ``HW_REG_HW_ID`` directly, the same mechanism
+        ``get_xcc_id`` below uses. The field is 4 bits, so this identifies the
+        CU within its shader engine rather than globally; pair it with
+        ``get_xcc_id`` for a fuller picture. Verified on gfx950, where a
+        256-workgroup launch reports CU_ID 0-8 alongside XCC_ID 0-7.
+
+        Other architectures report 0 rather than failing to compile: tracing is
+        diagnostic, and losing CU attribution is preferable to breaking every
+        traced kernel.
+
+        Returns:
+            int32: CU ID within the shader engine, or 0 where unsupported
+        """
+        if is_hip_cdna3() or is_hip_cdna4():
+            return tl.inline_asm_elementwise(
+                asm="s_getreg_b32 $0, hwreg(HW_REG_HW_ID, 8, 4)",
+                constraints=("=s"),
+                args=[],
+                dtype=tl.int32,
+                is_pure=False,
+                pack=1,
+            )
+        else:
+            return tl.cast(0, tl.int32)
 
 
 @triton.jit