From e4ba03338aec9eac9fa33df7a2e3afc5dd02b82f Mon Sep 17 00:00:00 2001 From: i3wanna2 <15910307812@163.com> Date: Tue, 7 Jul 2026 09:03:46 +0000 Subject: [PATCH 01/18] Add TLE raw external source to Triton kernel cache key. Inject @dialect file contents and kwargs into fn.cache_key via __triton_tle_raw_source_cache_key__, so editing .cu files invalidates the disk cache without changing kernel Python source. Co-authored-by: Cursor --- .../test/tle/unit/test_tle_raw_cache_key.py | 199 ++++++++++++++++++ .../triton/experimental/tle/raw/cache_key.py | 109 ++++++++++ python/triton/experimental/tle/raw/runtime.py | 2 + python/triton/runtime/jit.py | 5 + 4 files changed, 315 insertions(+) create mode 100644 python/test/tle/unit/test_tle_raw_cache_key.py create mode 100644 python/triton/experimental/tle/raw/cache_key.py diff --git a/python/test/tle/unit/test_tle_raw_cache_key.py b/python/test/tle/unit/test_tle_raw_cache_key.py new file mode 100644 index 000000000..4ce75c7be --- /dev/null +++ b/python/test/tle/unit/test_tle_raw_cache_key.py @@ -0,0 +1,199 @@ +import ast +import textwrap +from pathlib import Path + +import pytest +import triton +import triton.language as tl +from triton.experimental.tle.raw import dialect +from triton.experimental.tle.raw.cache_key import ( + TLE_RAW_SOURCE_CACHE_KEY_ATTR, + compute_tle_raw_source_cache_key, +) +from triton.runtime.jit import DependenciesFinder + + +def test_compute_tle_raw_source_cache_key_reads_file_content(tmp_path): + cu_file = tmp_path / "kernel.cu" + cu_file.write_text("__device__ void foo() {}\n") + kwargs = {"name": "cuda", "file": cu_file} + key1 = compute_tle_raw_source_cache_key(kwargs) + + cu_file.write_text("__device__ void bar() {}\n") + key2 = compute_tle_raw_source_cache_key(kwargs) + + assert key1 != key2 + + +def test_compute_tle_raw_source_cache_key_includes_extern_file(tmp_path): + main = tmp_path / "main.cu" + extern = tmp_path / "extern.py" + main.write_text("__device__ void foo() {}\n") + extern.write_text("helper = 1\n") + + with_extern = compute_tle_raw_source_cache_key({ + "name": "cuda", + "file": main, + "extern_file": extern, + }) + without_extern = compute_tle_raw_source_cache_key({ + "name": "cuda", + "file": main, + }) + assert with_extern != without_extern + + +def test_compute_tle_raw_source_cache_key_includes_dialect_kwargs(tmp_path): + cu_file = tmp_path / "kernel.cu" + cu_file.write_text("__device__ void foo() {}\n") + + bc_key = compute_tle_raw_source_cache_key({ + "name": "cuda", + "compiler": "clang", + "target": "bc", + "file": cu_file, + "extern_func_name": "foo", + }) + llvm_key = compute_tle_raw_source_cache_key({ + "name": "cuda", + "compiler": "clang", + "target": "llvm", + "file": cu_file, + "extern_func_name": "foo", + }) + assert bc_key != llvm_key + + +def test_bind_tle_raw_source_cache_key_rereads_source_file(tmp_path): + cu_file = tmp_path / "kernel.cu" + cu_file.write_text("__device__ void foo() {}\n") + + @dialect(name="cuda", file=cu_file) + def edsl(*args, **kwargs): + ... + + key1 = getattr(edsl, TLE_RAW_SOURCE_CACHE_KEY_ATTR)() + + cu_file.write_text("__device__ void bar() {}\n") + key2 = getattr(edsl, TLE_RAW_SOURCE_CACHE_KEY_ATTR)() + + assert key1 != key2 + + +def test_jit_cache_key_includes_raw_source(tmp_path): + cu_file = tmp_path / "vector-add.cu" + cu_file.write_text(textwrap.dedent("""\ + __device__ void VectorAdd(float *C, const float *A, const float *B, const int N) { + const int idx = blockIdx.x * blockDim.x + threadIdx.x; + for (int i = idx; i < N; i += blockDim.x * gridDim.x) { + C[i] = A[i] + B[i]; + } + } + """)) + + @dialect(name="cuda", file=cu_file) + def edsl(*args, **kwargs): + ... + + kernel_src = textwrap.dedent("""\ + def add_kernel(x_ptr, y_ptr, output_ptr, n_elements, BLOCK_SIZE: tl.constexpr): + tle_raw.call(edsl, [output_ptr, x_ptr, y_ptr, n_elements]) + """) + + def make_cache_key() -> str: + finder = DependenciesFinder( + name="add_kernel", + globals={"edsl": edsl, "tl": tl, "tle_raw": __import__("triton.experimental.tle.language.raw", fromlist=["raw"])}, + nonlocals={}, + src=kernel_src, + ) + finder.visit(ast.parse(kernel_src)) + return finder.ret + + key1 = make_cache_key() + + cu_file.write_text(textwrap.dedent("""\ + __device__ void VectorAdd(float *C, const float *A, const float *B, const int N) { + const int idx = blockIdx.x * blockDim.x + threadIdx.x; + for (int i = idx; i < N; i += blockDim.x * gridDim.x) { + C[i] = A[i] + B[i] + 1.0f; + } + } + """)) + key2 = make_cache_key() + + assert key1 != key2 + + +def test_mlir_dialect_cache_key_changes_with_edsl_source(): + @dialect(name="mlir") + def edsl_v1(): + pass + + @dialect(name="mlir") + def edsl_v2(): + return 1 + + assert getattr(edsl_v1, TLE_RAW_SOURCE_CACHE_KEY_ATTR)() != getattr(edsl_v2, TLE_RAW_SOURCE_CACHE_KEY_ATTR)() + + +@pytest.mark.skipif(not triton.runtime.driver.active.get_current_target().backend == "cuda", reason="requires cuda") +def test_cuda_vector_add_cache_miss_after_cu_change(tmp_path): + import torch + import triton.experimental.tle.language.raw as tle_raw + + cu_file = tmp_path / "vector-add.cu" + cu_file.write_text(textwrap.dedent("""\ + __device__ void VectorAdd(__attribute__((address_space(1))) float *C, + __attribute__((address_space(1))) const float *A, + __attribute__((address_space(1))) const float *B, + const int N) { + const int idx = blockIdx.x * blockDim.x + threadIdx.x; + for (int i = idx; i < N; i += blockDim.x * gridDim.x) { + C[i] = A[i] + B[i]; + } + } + """)) + + def build_kernel(): + @dialect(name="cuda", file=cu_file) + def edsl(*args, **kwargs): + ... + + @triton.jit + def add_kernel(x_ptr, y_ptr, output_ptr, n_elements, BLOCK_SIZE: tl.constexpr): + tle_raw.call(edsl, [output_ptr, x_ptr, y_ptr, n_elements]) + + return add_kernel + + add_kernel_v1 = build_kernel() + + device = triton.runtime.driver.active.get_active_torch_device() + x = torch.randn(128, device=device) + y = torch.randn(128, device=device) + output = torch.empty_like(x) + n_elements = output.numel() + grid = lambda meta: (triton.cdiv(n_elements, meta["BLOCK_SIZE"]), ) + + cache_key_before = add_kernel_v1.cache_key + add_kernel_v1[grid](x, y, output, n_elements, BLOCK_SIZE=128) + assert torch.allclose(output, x + y) + + cu_file.write_text(textwrap.dedent("""\ + __device__ void VectorAdd(__attribute__((address_space(1))) float *C, + __attribute__((address_space(1))) const float *A, + __attribute__((address_space(1))) const float *B, + const int N) { + const int idx = blockIdx.x * blockDim.x + threadIdx.x; + for (int i = idx; i < N; i += blockDim.x * gridDim.x) { + C[i] = A[i] + B[i] + 1.0f; + } + } + """)) + + add_kernel_v2 = build_kernel() + assert add_kernel_v2.cache_key != cache_key_before + + output2 = torch.empty_like(x) + add_kernel_v2[grid](x, y, output2, n_elements, BLOCK_SIZE=128) + assert torch.allclose(output2, x + y + 1.0) diff --git a/python/triton/experimental/tle/raw/cache_key.py b/python/triton/experimental/tle/raw/cache_key.py new file mode 100644 index 000000000..7dfbc6140 --- /dev/null +++ b/python/triton/experimental/tle/raw/cache_key.py @@ -0,0 +1,109 @@ +from __future__ import annotations + +import hashlib +import inspect +import json +from pathlib import Path +from typing import Any, Optional, Union + +# Protocol attribute checked by triton.runtime.jit.DependenciesFinder. +# Value may be a str or a zero-arg callable returning the current key fragment. +TLE_RAW_SOURCE_CACHE_KEY_ATTR = "__triton_tle_raw_source_cache_key__" + +_DIALECT_KWARG_KEYS = ( + "name", + "compiler", + "target", + "extern_func_name", + "deferred", +) +_DIALECT_PATH_KEYS = ( + "file", + "extern_file", +) + + +def _read_source(path: Union[str, Path]) -> str: + return Path(path).read_text() + + +def _normalize_dialect_kwargs(dialect_kwargs: dict) -> dict: + normalized: dict[str, Any] = {} + for key in _DIALECT_KWARG_KEYS: + if key not in dialect_kwargs: + continue + value = dialect_kwargs[key] + if value is not None: + normalized[key] = value + for key in _DIALECT_PATH_KEYS: + if key not in dialect_kwargs: + continue + value = dialect_kwargs[key] + if value is not None: + normalized[key] = str(Path(value).resolve()) + if "extra_source_files" in dialect_kwargs: + extra = dialect_kwargs["extra_source_files"] + if extra: + normalized["extra_source_files"] = [str(Path(path).resolve()) for path in extra] + return normalized + + +def _serialize_dialect_kwargs(dialect_kwargs: dict) -> str: + return json.dumps(dialect_kwargs, sort_keys=True, separators=(",", ":")) + + +def _collect_source_file_paths(dialect_kwargs: dict) -> list[str]: + paths: list[str] = [] + for key in _DIALECT_PATH_KEYS: + path = dialect_kwargs.get(key) + if path is not None: + paths.append(str(path)) + for path in dialect_kwargs.get("extra_source_files", ()): + paths.append(str(path)) + return sorted(set(paths)) + + +def _read_inline_edsl_source(edsl: Any) -> Optional[str]: + fn = getattr(edsl, "fn", None) + if fn is None: + return None + return inspect.getsource(fn) + + +def compute_tle_raw_source_cache_key( + dialect_kwargs: dict, + *, + edsl: Any | None = None, +) -> str: + """Hash @dialect kwargs plus the contents of referenced source files.""" + normalized = _normalize_dialect_kwargs(dialect_kwargs) + hasher = hashlib.sha256() + hasher.update(_serialize_dialect_kwargs(normalized).encode()) + + source_paths = _collect_source_file_paths(normalized) + for path in source_paths: + hasher.update(path.encode()) + hasher.update(_read_source(path).encode()) + + if not source_paths and edsl is not None: + inline_source = _read_inline_edsl_source(edsl) + if inline_source is not None: + hasher.update(b"inline_edsl_source") + hasher.update(inline_source.encode()) + + return hasher.hexdigest() + + +def bind_tle_raw_source_cache_key(edsl: Any, **dialect_kwargs) -> None: + """Attach __triton_tle_raw_source_cache_key__ to a @dialect edsl object.""" + if getattr(edsl, TLE_RAW_SOURCE_CACHE_KEY_ATTR, None) is not None: + return + + normalized = _normalize_dialect_kwargs(dialect_kwargs) + if not normalized and getattr(edsl, "fn", None) is None: + return + + def tle_raw_source_cache_key() -> str: + return compute_tle_raw_source_cache_key(normalized, edsl=edsl) + + setattr(edsl, TLE_RAW_SOURCE_CACHE_KEY_ATTR, tle_raw_source_cache_key) diff --git a/python/triton/experimental/tle/raw/runtime.py b/python/triton/experimental/tle/raw/runtime.py index 6e9918b3c..62d8746ee 100644 --- a/python/triton/experimental/tle/raw/runtime.py +++ b/python/triton/experimental/tle/raw/runtime.py @@ -1,3 +1,4 @@ +from .cache_key import bind_tle_raw_source_cache_key from .cuda import CUDAJITFunction from .mlir import MLIRJITFunction @@ -19,6 +20,7 @@ def dialect( def decorator(fn): edsl = registry[name](fn, **kwargs) + bind_tle_raw_source_cache_key(edsl, name=name, **kwargs) return edsl return decorator diff --git a/python/triton/runtime/jit.py b/python/triton/runtime/jit.py index 33b98fd12..79331c207 100644 --- a/python/triton/runtime/jit.py +++ b/python/triton/runtime/jit.py @@ -122,6 +122,11 @@ def record_reference(self, val, var_dict=None, name=None): if val is None or type(val) is ModuleType: return + tle_raw_source_cache_key = getattr(val, "__triton_tle_raw_source_cache_key__", None) + if tle_raw_source_cache_key is not None: + part = (tle_raw_source_cache_key() if callable(tle_raw_source_cache_key) else tle_raw_source_cache_key) + self.hasher.update(str(part).encode("utf-8")) + if getattr(val, "__triton_aggregate__", False): for attr in val.hash_attrs: self.record_reference(attr) From ad08934114d76ceaaf92c77cc8fb9eef76600833 Mon Sep 17 00:00:00 2001 From: flagtree-bot Date: Tue, 7 Jul 2026 09:09:49 +0000 Subject: [PATCH 02/18] Apply code-format changes --- .../test/tle/unit/test_tle_raw_cache_key.py | 19 +++++++++++++------ 1 file changed, 13 insertions(+), 6 deletions(-) diff --git a/python/test/tle/unit/test_tle_raw_cache_key.py b/python/test/tle/unit/test_tle_raw_cache_key.py index 4ce75c7be..03861177a 100644 --- a/python/test/tle/unit/test_tle_raw_cache_key.py +++ b/python/test/tle/unit/test_tle_raw_cache_key.py @@ -1,6 +1,5 @@ import ast import textwrap -from pathlib import Path import pytest import triton @@ -82,7 +81,8 @@ def edsl(*args, **kwargs): def test_jit_cache_key_includes_raw_source(tmp_path): cu_file = tmp_path / "vector-add.cu" - cu_file.write_text(textwrap.dedent("""\ + cu_file.write_text( + textwrap.dedent("""\ __device__ void VectorAdd(float *C, const float *A, const float *B, const int N) { const int idx = blockIdx.x * blockDim.x + threadIdx.x; for (int i = idx; i < N; i += blockDim.x * gridDim.x) { @@ -103,7 +103,9 @@ def add_kernel(x_ptr, y_ptr, output_ptr, n_elements, BLOCK_SIZE: tl.constexpr): def make_cache_key() -> str: finder = DependenciesFinder( name="add_kernel", - globals={"edsl": edsl, "tl": tl, "tle_raw": __import__("triton.experimental.tle.language.raw", fromlist=["raw"])}, + globals={ + "edsl": edsl, "tl": tl, "tle_raw": __import__("triton.experimental.tle.language.raw", fromlist=["raw"]) + }, nonlocals={}, src=kernel_src, ) @@ -112,7 +114,8 @@ def make_cache_key() -> str: key1 = make_cache_key() - cu_file.write_text(textwrap.dedent("""\ + cu_file.write_text( + textwrap.dedent("""\ __device__ void VectorAdd(float *C, const float *A, const float *B, const int N) { const int idx = blockIdx.x * blockDim.x + threadIdx.x; for (int i = idx; i < N; i += blockDim.x * gridDim.x) { @@ -126,6 +129,7 @@ def make_cache_key() -> str: def test_mlir_dialect_cache_key_changes_with_edsl_source(): + @dialect(name="mlir") def edsl_v1(): pass @@ -143,7 +147,8 @@ def test_cuda_vector_add_cache_miss_after_cu_change(tmp_path): import triton.experimental.tle.language.raw as tle_raw cu_file = tmp_path / "vector-add.cu" - cu_file.write_text(textwrap.dedent("""\ + cu_file.write_text( + textwrap.dedent("""\ __device__ void VectorAdd(__attribute__((address_space(1))) float *C, __attribute__((address_space(1))) const float *A, __attribute__((address_space(1))) const float *B, @@ -156,6 +161,7 @@ def test_cuda_vector_add_cache_miss_after_cu_change(tmp_path): """)) def build_kernel(): + @dialect(name="cuda", file=cu_file) def edsl(*args, **kwargs): ... @@ -179,7 +185,8 @@ def add_kernel(x_ptr, y_ptr, output_ptr, n_elements, BLOCK_SIZE: tl.constexpr): add_kernel_v1[grid](x, y, output, n_elements, BLOCK_SIZE=128) assert torch.allclose(output, x + y) - cu_file.write_text(textwrap.dedent("""\ + cu_file.write_text( + textwrap.dedent("""\ __device__ void VectorAdd(__attribute__((address_space(1))) float *C, __attribute__((address_space(1))) const float *A, __attribute__((address_space(1))) const float *B, From e0ca43736b650bea8f8fbc6a4974e1bd63293e34 Mon Sep 17 00:00:00 2001 From: i3wanna2 <15910307812@163.com> Date: Wed, 8 Jul 2026 03:01:53 +0000 Subject: [PATCH 03/18] Add flagtree tle raw marker comment in jit cache key hook. Co-authored-by: Cursor --- python/triton/runtime/jit.py | 1 + 1 file changed, 1 insertion(+) diff --git a/python/triton/runtime/jit.py b/python/triton/runtime/jit.py index 79331c207..93ac1dd2e 100644 --- a/python/triton/runtime/jit.py +++ b/python/triton/runtime/jit.py @@ -122,6 +122,7 @@ def record_reference(self, val, var_dict=None, name=None): if val is None or type(val) is ModuleType: return + # flagtree tle raw tle_raw_source_cache_key = getattr(val, "__triton_tle_raw_source_cache_key__", None) if tle_raw_source_cache_key is not None: part = (tle_raw_source_cache_key() if callable(tle_raw_source_cache_key) else tle_raw_source_cache_key) From 5220cc66dd8886e56141ae032644381908454923 Mon Sep 17 00:00:00 2001 From: sunnycase Date: Wed, 8 Jul 2026 03:09:19 +0000 Subject: [PATCH 04/18] Debug PR760 cluster GEMM candidate --- .github/workflows/hopper-build-and-test.yml | 125 +++++++++++--------- 1 file changed, 66 insertions(+), 59 deletions(-) diff --git a/.github/workflows/hopper-build-and-test.yml b/.github/workflows/hopper-build-and-test.yml index d8b2eb243..27aba1306 100644 --- a/.github/workflows/hopper-build-and-test.yml +++ b/.github/workflows/hopper-build-and-test.yml @@ -156,62 +156,69 @@ jobs: run: | set -x source ~/env.sh - # python tutorials - python3 python/tutorials/01-vector-add.py --only_unit_test - python3 python/tutorials/02-fused-softmax.py --only_unit_test - python3 python/tutorials/03-matrix-multiplication.py --only_unit_test - python3 python/tutorials/04-low-memory-dropout.py --only_unit_test - python3 python/tutorials/05-layer-norm.py --only_unit_test - python3 python/tutorials/06-fused-attention.py --only_unit_test - python3 python/tutorials/07-extern-functions.py --only_unit_test - python3 python/tutorials/08-grouped-gemm.py --only_unit_test - python3 python/tutorials/09-persistent-matmul.py --only_unit_test - python3 python/tutorials/11-programmatic-dependent-launch.py --only_unit_test - # python unit test - python3 -m pytest -s python/test/unit/cuda - python3 -m pytest -s python/test/unit/instrumentation - python3 -m pytest -s python/test/unit/language \ - --ignore=python/test/unit/language/test_line_info.py - python3 -m pytest -s python/test/unit/runtime - if [ -d "python/test/operators" ]; then - python3 -m pytest -s python/test/operators - fi - # flagtree tle - # python tutorials - python3 python/tutorials/tle/01-fft.py - python3 python/tutorials/tle/02-moe_align_block_size.py - python3 python/tutorials/tle/03-topk.py - python3 python/tutorials/tle/04-cluster-gemm.py - python3 python/tutorials/tle/deepseek_v32/01-topk_selector.py - python3 python/tutorials/tle/deepseek_v32/02-sparse-mla.py - # python unit test - python3 -m pytest -s python/test/tle/integration - python3 -m pytest -s python/test/tle/unit \ - --ignore=python/test/tle/unit/test_tle_distributed_d2d.py \ - --ignore=python/test/tle/unit/test_tle_get_local_pe.py - CUDA_LAUNCH_BLOCKING=1 bash python/test/tle/unit/test_tle_distributed_d2d.sh - CUDA_LAUNCH_BLOCKING=1 bash python/test/tle/unit/test_tle_get_local_pe.sh - # flagtree hints - # python tutorials - python3 python/tutorials/hints/01/01-vector-add.py --only_unit_test - # python3 python/tutorials/hints/02/02-fused-softmax.py --only_unit_test - python3 python/tutorials/hints/03/03-matrix-multiplication.py --only_unit_test - python3 python/tutorials/hints/04/04-low-memory-dropout.py --only_unit_test - python3 python/tutorials/hints/05/05-layer-norm.py --only_unit_test - # python3 python/tutorials/hints/06/06-fused-attention.py --only_unit_test # Error on 3.6 - python3 python/tutorials/hints/07/07-extern-functions.py --only_unit_test - python3 python/tutorials/hints/08/08-grouped-gemm.py --only_unit_test - python3 python/tutorials/hints/11/11-programmatic-dependent-launch.py --only_unit_test - # flagtree tle raw - python3 python/tutorials/tle/raw/mlir/01-vector-add.py - python3 python/tutorials/tle/raw/mlir/02-fused-softmax.py - python3 python/tutorials/tle/raw/mlir/03-matrix-multiplication.py - python3 python/tutorials/tle/raw/mlir/04-hello-world.py - python3 python/tutorials/tle/raw/mlir/05-topk.py --kernel triton - python3 python/tutorials/tle/raw/mlir/05-topk.py --kernel tle - python3 python/tutorials/tle/raw/mlir/06-test-vassert.py - # flagtree tle cuda - # TODO: These tests are currently skipped because the CLANG environment variable cannot be set. - # python3 python/tutorials/tle/raw/cuda/01-vector-add.py - # python3 python/tutorials/tle/raw/cuda/02-fused-softmax.py - # python3 python/tutorials/tle/raw/cuda/03-matrix-multiplication.py + CUDA_LAUNCH_BLOCKING=1 python3 - <<'PY' + import importlib.util + import pathlib + import sys + + import torch + + path = pathlib.Path("python/tutorials/tle/04-cluster-gemm.py") + spec = importlib.util.spec_from_file_location("cluster_gemm_debug", path) + mod = importlib.util.module_from_spec(spec) + sys.modules[spec.name] = mod + spec.loader.exec_module(mod) + + skip_reason = mod._cluster_remote_support_skip_reason() + if skip_reason is not None: + print(f"SKIP: {skip_reason}", flush=True) + raise SystemExit(0) + + M = N = K = 4096 + torch.manual_seed(0) + a = torch.randn((M, K), device="cuda", dtype=torch.float16) + b = torch.randn((K, N), device="cuda", dtype=torch.float16) + c_remote = torch.empty((M, N), device="cuda", dtype=torch.float16) + print( + f"[debug] ptr mod256 a={a.data_ptr() % 256} " + f"b={b.data_ptr() % 256} c={c_remote.data_ptr() % 256}", + flush=True, + ) + + cases = [ + ("remote_candidate_0", mod.KernelConfig(32, 512, 256, 4, 2)), + ("remote_candidate_1_suspected_stage3", mod.KernelConfig(32, 512, 128, 4, 3)), + ] + progress = pathlib.Path("/tmp/cluster-gemm-candidate.txt") + + def run_case(name, cfg): + progress.write_text(f"{name}: {cfg}\n") + print(f"[debug] running {name}: {cfg}", flush=True) + mod._run_cluster_remote(a, b, c_remote, cfg.bm, cfg.bn, cfg.bk, cfg.num_warps, cfg.num_stages) + torch.cuda.synchronize() + ms = mod.triton.testing.do_bench( + lambda: mod._run_cluster_remote( + a, + b, + c_remote, + cfg.bm, + cfg.bn, + cfg.bk, + cfg.num_warps, + cfg.num_stages, + ), + warmup=1, + rep=1, + ) + torch.cuda.synchronize() + progress.write_text(f"{name}: {cfg} PASS ms={ms}\n") + print(f"[debug] {name}: PASS ms={ms}", flush=True) + + try: + for name, cfg in cases: + run_case(name, cfg) + except BaseException: + print("[debug] last cluster-gemm candidate:", flush=True) + print(progress.read_text(), flush=True) + raise + PY From 44dea6f88f435124a1f1e6c766e3c991ef409b69 Mon Sep 17 00:00:00 2001 From: sunnycase Date: Wed, 8 Jul 2026 03:29:43 +0000 Subject: [PATCH 05/18] Expand cluster GEMM debug candidates --- .github/workflows/hopper-build-and-test.yml | 40 ++++++++++++++++----- 1 file changed, 31 insertions(+), 9 deletions(-) diff --git a/.github/workflows/hopper-build-and-test.yml b/.github/workflows/hopper-build-and-test.yml index 27aba1306..4197db627 100644 --- a/.github/workflows/hopper-build-and-test.yml +++ b/.github/workflows/hopper-build-and-test.yml @@ -178,19 +178,39 @@ jobs: torch.manual_seed(0) a = torch.randn((M, K), device="cuda", dtype=torch.float16) b = torch.randn((K, N), device="cuda", dtype=torch.float16) - c_remote = torch.empty((M, N), device="cuda", dtype=torch.float16) + c_triton = torch.empty((M, N), device="cuda", dtype=torch.float16) + c_remote = torch.empty_like(c_triton) print( f"[debug] ptr mod256 a={a.data_ptr() % 256} " f"b={b.data_ptr() % 256} c={c_remote.data_ptr() % 256}", flush=True, ) - cases = [ - ("remote_candidate_0", mod.KernelConfig(32, 512, 256, 4, 2)), - ("remote_candidate_1_suspected_stage3", mod.KernelConfig(32, 512, 128, 4, 3)), - ] progress = pathlib.Path("/tmp/cluster-gemm-candidate.txt") + def run_triton_case(idx, cfg): + progress.write_text(f"triton {idx}: {cfg}\n") + print(f"[debug] running triton {idx}: {cfg}", flush=True) + mod._run_triton(a, b, c_triton, cfg.bm, cfg.bn, cfg.bk, cfg.num_warps, cfg.num_stages) + torch.cuda.synchronize() + ms = mod.triton.testing.do_bench( + lambda: mod._run_triton( + a, + b, + c_triton, + cfg.bm, + cfg.bn, + cfg.bk, + cfg.num_warps, + cfg.num_stages, + ), + warmup=10, + rep=30, + ) + torch.cuda.synchronize() + progress.write_text(f"triton {idx}: {cfg} PASS ms={ms}\n") + print(f"[debug] triton {idx}: PASS ms={ms}", flush=True) + def run_case(name, cfg): progress.write_text(f"{name}: {cfg}\n") print(f"[debug] running {name}: {cfg}", flush=True) @@ -207,16 +227,18 @@ jobs: cfg.num_warps, cfg.num_stages, ), - warmup=1, - rep=1, + warmup=10, + rep=30, ) torch.cuda.synchronize() progress.write_text(f"{name}: {cfg} PASS ms={ms}\n") print(f"[debug] {name}: PASS ms={ms}", flush=True) try: - for name, cfg in cases: - run_case(name, cfg) + for idx, cfg in enumerate(mod.TRITON_TUNE_CONFIGS): + run_triton_case(idx, cfg) + for idx, cfg in enumerate(mod.REMOTE_TUNE_CONFIGS): + run_case(f"remote_candidate_{idx}", cfg) except BaseException: print("[debug] last cluster-gemm candidate:", flush=True) print(progress.read_text(), flush=True) From 9e4209986fa90a411c344c43a0ff710a992e9297 Mon Sep 17 00:00:00 2001 From: sunnycase Date: Wed, 8 Jul 2026 04:22:20 +0000 Subject: [PATCH 06/18] Test original cluster GEMM failure merge --- .github/workflows/hopper-build-and-test.yml | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/.github/workflows/hopper-build-and-test.yml b/.github/workflows/hopper-build-and-test.yml index 1ea1b79fa..df08cf108 100644 --- a/.github/workflows/hopper-build-and-test.yml +++ b/.github/workflows/hopper-build-and-test.yml @@ -51,6 +51,12 @@ jobs: shell: bash run: | set -x + # Reproduce the original PR #760 attempt-3 merge commit that hit + # cluster-gemm misaligned address before the base branch moved. + git fetch origin dd79a2328a578057bd9dc8e0a1a2f6351bbf635a + git checkout --force FETCH_HEAD + git clean -ffdx + git log -1 --format='%H %s' pip uninstall -y triton source ~/env-3.6.sh env | grep -E '^(LLVM_SYSPATH)=' >> $GITHUB_ENV || true @@ -63,7 +69,8 @@ jobs: run: | set -x source ~/env.sh - CUDA_LAUNCH_BLOCKING=1 python3 - <<'PY' + git log -1 --format='[debug] workspace commit %H %s' + python3 - <<'PY' import importlib.util import pathlib import sys From 8e24f45187a49ab8b374421093503ef79aaf2d5d Mon Sep 17 00:00:00 2001 From: sunnycase Date: Wed, 8 Jul 2026 04:37:01 +0000 Subject: [PATCH 07/18] Clear Triton cache for cluster GEMM repro --- .github/workflows/hopper-build-and-test.yml | 2 ++ 1 file changed, 2 insertions(+) diff --git a/.github/workflows/hopper-build-and-test.yml b/.github/workflows/hopper-build-and-test.yml index df08cf108..e2f5ed1fc 100644 --- a/.github/workflows/hopper-build-and-test.yml +++ b/.github/workflows/hopper-build-and-test.yml @@ -58,6 +58,7 @@ jobs: git clean -ffdx git log -1 --format='%H %s' pip uninstall -y triton + rm -rf ~/.triton/cache source ~/env-3.6.sh env | grep -E '^(LLVM_SYSPATH)=' >> $GITHUB_ENV || true export USE_FLAGCX=ON @@ -70,6 +71,7 @@ jobs: set -x source ~/env.sh git log -1 --format='[debug] workspace commit %H %s' + rm -rf ~/.triton/cache python3 - <<'PY' import importlib.util import pathlib From 4c2da28376d45fae8b2a07a16ece848e4b7e5586 Mon Sep 17 00:00:00 2001 From: sunnycase Date: Wed, 8 Jul 2026 04:49:44 +0000 Subject: [PATCH 08/18] Stress old cluster GEMM failure case --- .github/workflows/hopper-build-and-test.yml | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/.github/workflows/hopper-build-and-test.yml b/.github/workflows/hopper-build-and-test.yml index e2f5ed1fc..1c72630a8 100644 --- a/.github/workflows/hopper-build-and-test.yml +++ b/.github/workflows/hopper-build-and-test.yml @@ -71,9 +71,12 @@ jobs: set -x source ~/env.sh git log -1 --format='[debug] workspace commit %H %s' - rm -rf ~/.triton/cache - python3 - <<'PY' + for trial in $(seq 1 20); do + echo "[debug] trial ${trial}/20" + rm -rf ~/.triton/cache + CLUSTER_GEMM_TRIAL="${trial}" python3 - <<'PY' import importlib.util + import os import pathlib import sys @@ -90,6 +93,7 @@ jobs: print(f"SKIP: {skip_reason}", flush=True) raise SystemExit(0) + trial = os.environ.get("CLUSTER_GEMM_TRIAL", "?") M = N = K = 4096 torch.manual_seed(0) a = torch.randn((M, K), device="cuda", dtype=torch.float16) @@ -97,7 +101,7 @@ jobs: c_triton = torch.empty((M, N), device="cuda", dtype=torch.float16) c_remote = torch.empty_like(c_triton) print( - f"[debug] ptr mod256 a={a.data_ptr() % 256} " + f"[debug] trial {trial} ptr mod256 a={a.data_ptr() % 256} " f"b={b.data_ptr() % 256} c={c_remote.data_ptr() % 256}", flush=True, ) @@ -160,3 +164,4 @@ jobs: print(progress.read_text(), flush=True) raise PY + done From 266a4a588b4cee1b380ac04534a16840e8295a0d Mon Sep 17 00:00:00 2001 From: sunnycase Date: Wed, 8 Jul 2026 05:23:46 +0000 Subject: [PATCH 09/18] Stress cluster GEMM without clearing cache --- .github/workflows/hopper-build-and-test.yml | 2 -- 1 file changed, 2 deletions(-) diff --git a/.github/workflows/hopper-build-and-test.yml b/.github/workflows/hopper-build-and-test.yml index 1c72630a8..08efa5c9c 100644 --- a/.github/workflows/hopper-build-and-test.yml +++ b/.github/workflows/hopper-build-and-test.yml @@ -58,7 +58,6 @@ jobs: git clean -ffdx git log -1 --format='%H %s' pip uninstall -y triton - rm -rf ~/.triton/cache source ~/env-3.6.sh env | grep -E '^(LLVM_SYSPATH)=' >> $GITHUB_ENV || true export USE_FLAGCX=ON @@ -73,7 +72,6 @@ jobs: git log -1 --format='[debug] workspace commit %H %s' for trial in $(seq 1 20); do echo "[debug] trial ${trial}/20" - rm -rf ~/.triton/cache CLUSTER_GEMM_TRIAL="${trial}" python3 - <<'PY' import importlib.util import os From 512d2c739ab40daf8b537996fb3d7708829b309e Mon Sep 17 00:00:00 2001 From: sunnycase Date: Wed, 8 Jul 2026 05:40:48 +0000 Subject: [PATCH 10/18] Run original CI prefix before cluster GEMM --- .github/workflows/hopper-build-and-test.yml | 131 +++++++++----------- 1 file changed, 58 insertions(+), 73 deletions(-) diff --git a/.github/workflows/hopper-build-and-test.yml b/.github/workflows/hopper-build-and-test.yml index 08efa5c9c..57d652144 100644 --- a/.github/workflows/hopper-build-and-test.yml +++ b/.github/workflows/hopper-build-and-test.yml @@ -60,8 +60,8 @@ jobs: pip uninstall -y triton source ~/env-3.6.sh env | grep -E '^(LLVM_SYSPATH)=' >> $GITHUB_ENV || true - export USE_FLAGCX=ON - MAX_JOBS=32 python3 -m pip install . --no-build-isolation + rm -rf ~/.triton/cache + USE_FLAGCX=ON MAX_JOBS=32 python3 -m pip install . --no-build-isolation - name: FlagTree Test on NVidia (triton_v3.6.x branch) if: ${{ steps.check_backend.outputs.should_skip != 'true' && env.TARGET_BRANCH == 'triton_v3.6.x' }} @@ -70,11 +70,36 @@ jobs: set -x source ~/env.sh git log -1 --format='[debug] workspace commit %H %s' - for trial in $(seq 1 20); do - echo "[debug] trial ${trial}/20" - CLUSTER_GEMM_TRIAL="${trial}" python3 - <<'PY' + echo "[debug] runner=${RUNNER_NAME:-unset} cuda_visible=${CUDA_VISIBLE_DEVICES:-unset}" + nvidia-smi --query-gpu=index,name,uuid,memory.used,memory.total --format=csv || true + # Re-run the same CI prefix that completed before the original + # PR #760 attempt-3 cluster-gemm failure. + python3 python/tutorials/01-vector-add.py --only_unit_test + python3 python/tutorials/02-fused-softmax.py --only_unit_test + python3 python/tutorials/03-matrix-multiplication.py --only_unit_test + python3 python/tutorials/04-low-memory-dropout.py --only_unit_test + python3 python/tutorials/05-layer-norm.py --only_unit_test + python3 python/tutorials/06-fused-attention.py --only_unit_test + python3 python/tutorials/07-extern-functions.py --only_unit_test + python3 python/tutorials/08-grouped-gemm.py --only_unit_test + python3 python/tutorials/09-persistent-matmul.py --only_unit_test + python3 python/tutorials/11-programmatic-dependent-launch.py --only_unit_test + python3 -m pytest -s python/test/unit/cuda + python3 -m pytest -s python/test/unit/instrumentation + python3 -m pytest -s python/test/unit/language \ + --ignore=python/test/unit/language/test_line_info.py + python3 -m pytest -s python/test/unit/runtime + if [ -d "python/test/operators" ]; then + python3 -m pytest -s python/test/operators + fi + python3 python/tutorials/tle/01-fft.py + python3 python/tutorials/tle/02-moe_align_block_size.py + python3 python/tutorials/tle/03-topk.py + echo "[debug] triton cache before cluster-gemm" + du -sh ~/.triton/cache || true + find ~/.triton/cache -type f 2>/dev/null | wc -l || true + python3 - <<'PY' import importlib.util - import os import pathlib import sys @@ -86,80 +111,40 @@ jobs: sys.modules[spec.name] = mod spec.loader.exec_module(mod) - skip_reason = mod._cluster_remote_support_skip_reason() - if skip_reason is not None: - print(f"SKIP: {skip_reason}", flush=True) - raise SystemExit(0) - - trial = os.environ.get("CLUSTER_GEMM_TRIAL", "?") - M = N = K = 4096 - torch.manual_seed(0) - a = torch.randn((M, K), device="cuda", dtype=torch.float16) - b = torch.randn((K, N), device="cuda", dtype=torch.float16) - c_triton = torch.empty((M, N), device="cuda", dtype=torch.float16) - c_remote = torch.empty_like(c_triton) - print( - f"[debug] trial {trial} ptr mod256 a={a.data_ptr() % 256} " - f"b={b.data_ptr() % 256} c={c_remote.data_ptr() % 256}", - flush=True, - ) - progress = pathlib.Path("/tmp/cluster-gemm-candidate.txt") - def run_triton_case(idx, cfg): - progress.write_text(f"triton {idx}: {cfg}\n") - print(f"[debug] running triton {idx}: {cfg}", flush=True) - mod._run_triton(a, b, c_triton, cfg.bm, cfg.bn, cfg.bk, cfg.num_warps, cfg.num_stages) - torch.cuda.synchronize() - ms = mod.triton.testing.do_bench( - lambda: mod._run_triton( - a, - b, - c_triton, - cfg.bm, - cfg.bn, - cfg.bk, - cfg.num_warps, - cfg.num_stages, - ), - warmup=10, - rep=30, + def logged_autotune(name, candidates, run_fn, M, N, K, warmup, rep): + best_cfg = candidates[0] + best_ms = float("inf") + for idx, cfg in enumerate(candidates): + progress.write_text(f"{name} candidate_{idx}: {cfg}\n") + print(f"[debug] running {name} candidate_{idx}: {cfg}", flush=True) + run_fn(cfg) + torch.cuda.synchronize() + ms = mod.triton.testing.do_bench( + lambda: run_fn(cfg), + warmup=warmup, + rep=rep, + ) + torch.cuda.synchronize() + progress.write_text(f"{name} candidate_{idx}: {cfg} PASS ms={ms}\n") + print(f"[debug] {name} candidate_{idx}: PASS ms={ms}", flush=True) + if ms < best_ms: + best_ms = ms + best_cfg = cfg + print( + f"[autotune] {name}: best cfg={best_cfg} " + f"ms={best_ms:.3f} tflops={mod._tflops(M, N, K, best_ms):.2f}" ) - torch.cuda.synchronize() - progress.write_text(f"triton {idx}: {cfg} PASS ms={ms}\n") - print(f"[debug] triton {idx}: PASS ms={ms}", flush=True) + return best_cfg - def run_case(name, cfg): - progress.write_text(f"{name}: {cfg}\n") - print(f"[debug] running {name}: {cfg}", flush=True) - mod._run_cluster_remote(a, b, c_remote, cfg.bm, cfg.bn, cfg.bk, cfg.num_warps, cfg.num_stages) - torch.cuda.synchronize() - ms = mod.triton.testing.do_bench( - lambda: mod._run_cluster_remote( - a, - b, - c_remote, - cfg.bm, - cfg.bn, - cfg.bk, - cfg.num_warps, - cfg.num_stages, - ), - warmup=10, - rep=30, - ) - torch.cuda.synchronize() - progress.write_text(f"{name}: {cfg} PASS ms={ms}\n") - print(f"[debug] {name}: PASS ms={ms}", flush=True) + mod._autotune_config = logged_autotune try: - for idx, cfg in enumerate(mod.TRITON_TUNE_CONFIGS): - run_triton_case(idx, cfg) - for idx, cfg in enumerate(mod.REMOTE_TUNE_CONFIGS): - run_case(f"remote_candidate_{idx}", cfg) + mod.main([]) except BaseException: print("[debug] last cluster-gemm candidate:", flush=True) - print(progress.read_text(), flush=True) + if progress.exists(): + print(progress.read_text(), flush=True) raise PY - done From 7133d06ad7ea074c8c076328779b749d3ee0ce6c Mon Sep 17 00:00:00 2001 From: sunnycase Date: Wed, 8 Jul 2026 06:21:37 +0000 Subject: [PATCH 11/18] Rerun original prefix cluster GEMM From a06c3bae2d7e79222c7c4158a4ad24cbbb85778b Mon Sep 17 00:00:00 2001 From: sunnycase Date: Wed, 8 Jul 2026 07:01:57 +0000 Subject: [PATCH 12/18] Stress cluster GEMM after original prefix --- .github/workflows/hopper-build-and-test.yml | 18 ++++++++++++------ 1 file changed, 12 insertions(+), 6 deletions(-) diff --git a/.github/workflows/hopper-build-and-test.yml b/.github/workflows/hopper-build-and-test.yml index 57d652144..c65ae611f 100644 --- a/.github/workflows/hopper-build-and-test.yml +++ b/.github/workflows/hopper-build-and-test.yml @@ -98,8 +98,11 @@ jobs: echo "[debug] triton cache before cluster-gemm" du -sh ~/.triton/cache || true find ~/.triton/cache -type f 2>/dev/null | wc -l || true - python3 - <<'PY' + for trial in $(seq 1 50); do + echo "[debug] cluster-gemm trial ${trial}/50" + CLUSTER_GEMM_TRIAL="${trial}" python3 - <<'PY' import importlib.util + import os import pathlib import sys @@ -111,14 +114,15 @@ jobs: sys.modules[spec.name] = mod spec.loader.exec_module(mod) - progress = pathlib.Path("/tmp/cluster-gemm-candidate.txt") + trial = os.environ.get("CLUSTER_GEMM_TRIAL", "?") + progress = pathlib.Path(f"/tmp/cluster-gemm-candidate-{trial}.txt") def logged_autotune(name, candidates, run_fn, M, N, K, warmup, rep): best_cfg = candidates[0] best_ms = float("inf") for idx, cfg in enumerate(candidates): - progress.write_text(f"{name} candidate_{idx}: {cfg}\n") - print(f"[debug] running {name} candidate_{idx}: {cfg}", flush=True) + progress.write_text(f"trial {trial} {name} candidate_{idx}: {cfg}\n") + print(f"[debug] trial {trial} running {name} candidate_{idx}: {cfg}", flush=True) run_fn(cfg) torch.cuda.synchronize() ms = mod.triton.testing.do_bench( @@ -127,8 +131,8 @@ jobs: rep=rep, ) torch.cuda.synchronize() - progress.write_text(f"{name} candidate_{idx}: {cfg} PASS ms={ms}\n") - print(f"[debug] {name} candidate_{idx}: PASS ms={ms}", flush=True) + progress.write_text(f"trial {trial} {name} candidate_{idx}: {cfg} PASS ms={ms}\n") + print(f"[debug] trial {trial} {name} candidate_{idx}: PASS ms={ms}", flush=True) if ms < best_ms: best_ms = ms best_cfg = cfg @@ -141,6 +145,7 @@ jobs: mod._autotune_config = logged_autotune try: + print(f"[debug] cluster-gemm trial {trial} start", flush=True) mod.main([]) except BaseException: print("[debug] last cluster-gemm candidate:", flush=True) @@ -148,3 +153,4 @@ jobs: print(progress.read_text(), flush=True) raise PY + done From 9ca5eb946aa51621f605cd75af8717543a4d697e Mon Sep 17 00:00:00 2001 From: sunnycase Date: Wed, 8 Jul 2026 07:56:06 +0000 Subject: [PATCH 13/18] Stress cluster GEMM with original autotune sync --- .github/workflows/hopper-build-and-test.yml | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/.github/workflows/hopper-build-and-test.yml b/.github/workflows/hopper-build-and-test.yml index c65ae611f..49c7d1e43 100644 --- a/.github/workflows/hopper-build-and-test.yml +++ b/.github/workflows/hopper-build-and-test.yml @@ -98,8 +98,8 @@ jobs: echo "[debug] triton cache before cluster-gemm" du -sh ~/.triton/cache || true find ~/.triton/cache -type f 2>/dev/null | wc -l || true - for trial in $(seq 1 50); do - echo "[debug] cluster-gemm trial ${trial}/50" + for trial in $(seq 1 200); do + echo "[debug] cluster-gemm trial ${trial}/200" CLUSTER_GEMM_TRIAL="${trial}" python3 - <<'PY' import importlib.util import os @@ -130,7 +130,6 @@ jobs: warmup=warmup, rep=rep, ) - torch.cuda.synchronize() progress.write_text(f"trial {trial} {name} candidate_{idx}: {cfg} PASS ms={ms}\n") print(f"[debug] trial {trial} {name} candidate_{idx}: PASS ms={ms}", flush=True) if ms < best_ms: From 4075b89604057a15f98736fa4cc51db9501f4305 Mon Sep 17 00:00:00 2001 From: sunnycase Date: Wed, 8 Jul 2026 09:35:56 +0000 Subject: [PATCH 14/18] Audit cluster GEMM remote shared loads --- .github/workflows/hopper-build-and-test.yml | 117 +++++++++++++++++++- 1 file changed, 113 insertions(+), 4 deletions(-) diff --git a/.github/workflows/hopper-build-and-test.yml b/.github/workflows/hopper-build-and-test.yml index 49c7d1e43..572404230 100644 --- a/.github/workflows/hopper-build-and-test.yml +++ b/.github/workflows/hopper-build-and-test.yml @@ -98,9 +98,119 @@ jobs: echo "[debug] triton cache before cluster-gemm" du -sh ~/.triton/cache || true find ~/.triton/cache -type f 2>/dev/null | wc -l || true - for trial in $(seq 1 200); do - echo "[debug] cluster-gemm trial ${trial}/200" - CLUSTER_GEMM_TRIAL="${trial}" python3 - <<'PY' + echo "[debug] audit remote cluster-gemm PTX" + python3 - <<'PY' + import collections + import importlib.util + import pathlib + import re + import sys + + import torch + + path = pathlib.Path("python/tutorials/tle/04-cluster-gemm.py") + spec = importlib.util.spec_from_file_location("cluster_gemm_audit", path) + mod = importlib.util.module_from_spec(spec) + sys.modules[spec.name] = mod + spec.loader.exec_module(mod) + + torch.manual_seed(0) + M = N = K = 4096 + dtype = torch.float16 + a = torch.randn((M, K), device="cuda", dtype=dtype) + b = torch.randn((K, N), device="cuda", dtype=dtype) + c = torch.empty((M, N), device="cuda", dtype=dtype) + + op_re = re.compile(r"\b(?:ld|st)\.shared::cluster(?:\.v\d+)?\.b\d+") + for idx, cfg in enumerate(mod.REMOTE_TUNE_CONFIGS): + dot_k = mod._select_remote_dot_k(cfg.bk) + use_mask = (M % cfg.bm != 0) or (N % cfg.bn != 0) or (K % cfg.bk != 0) + use_nv_mma_smem_layout = (cfg.bk == 32) or (cfg.bk == 64 and cfg.num_stages <= 2) + compiled = mod._cluster_remote_gemm_kernel.warmup( + a, + b, + c, + M, + N, + K, + a.stride(0), + a.stride(1), + b.stride(0), + b.stride(1), + c.stride(0), + c.stride(1), + mesh=mod.BLOCK_CLUSTER_MESH, + BM=cfg.bm, + BN=cfg.bn, + BK=cfg.bk, + DOT_K=dot_k, + CLUSTER_SIZE=2, + USE_MASK=use_mask, + A_SLOTS=2, + USE_NV_MMA_SMEM_LAYOUT=use_nv_mma_smem_layout, + grid=mod._grid_cluster_remote(M, N, cfg.bm, cfg.bn), + num_ctas=1, + num_warps=cfg.num_warps, + num_stages=cfg.num_stages, + ) + ptx = compiled.asm.get("ptx", "") + ops = collections.Counter(op_re.findall(ptx)) + print(f"[audit] remote candidate_{idx}: {cfg}", flush=True) + print( + f"[audit] cluster_dims={tuple(compiled.metadata.cluster_dims)} " + f"use_nv_mma_smem_layout={use_nv_mma_smem_layout}", + flush=True, + ) + print( + f"[audit] mapa.shared::cluster={ptx.count('mapa.shared::cluster')} " + f"barrier.cluster={ptx.count('barrier.cluster')}", + flush=True, + ) + for op, count in sorted(ops.items()): + print(f"[audit] {op}: {count}", flush=True) + PY + if command -v compute-sanitizer >/dev/null 2>&1; then + echo "[debug] compute-sanitizer memcheck for remote candidates" + compute-sanitizer --tool memcheck --error-exitcode=86 python3 - <<'PY' + import importlib.util + import pathlib + import sys + + import torch + + path = pathlib.Path("python/tutorials/tle/04-cluster-gemm.py") + spec = importlib.util.spec_from_file_location("cluster_gemm_sanitizer", path) + mod = importlib.util.module_from_spec(spec) + sys.modules[spec.name] = mod + spec.loader.exec_module(mod) + + torch.manual_seed(0) + M, N, K = 64, 1024, 256 + dtype = torch.float16 + a = torch.randn((M, K), device="cuda", dtype=dtype) + b = torch.randn((K, N), device="cuda", dtype=dtype) + c = torch.empty((M, N), device="cuda", dtype=dtype) + + for idx, cfg in enumerate(mod.REMOTE_TUNE_CONFIGS): + print(f"[sanitizer] running remote candidate_{idx}: {cfg}", flush=True) + mod._run_cluster_remote( + a, + b, + c, + cfg.bm, + cfg.bn, + cfg.bk, + cfg.num_warps, + cfg.num_stages, + ) + torch.cuda.synchronize() + print(f"[sanitizer] remote candidate_{idx}: PASS", flush=True) + PY + else + echo "[debug] compute-sanitizer not found; skipping memcheck" + fi + echo "[debug] cluster-gemm logged diagnostic trial" + CLUSTER_GEMM_TRIAL="diagnostic" python3 - <<'PY' import importlib.util import os import pathlib @@ -152,4 +262,3 @@ jobs: print(progress.read_text(), flush=True) raise PY - done From c63e34f8496273745a4d83ff616d5af25ce7db9e Mon Sep 17 00:00:00 2001 From: sunnycase Date: Wed, 8 Jul 2026 10:37:52 +0000 Subject: [PATCH 15/18] Match original MoE external download failure --- .github/workflows/hopper-build-and-test.yml | 25 ++++++++++++++++++++- 1 file changed, 24 insertions(+), 1 deletion(-) diff --git a/.github/workflows/hopper-build-and-test.yml b/.github/workflows/hopper-build-and-test.yml index 572404230..3072ea548 100644 --- a/.github/workflows/hopper-build-and-test.yml +++ b/.github/workflows/hopper-build-and-test.yml @@ -93,7 +93,30 @@ jobs: python3 -m pytest -s python/test/operators fi python3 python/tutorials/tle/01-fft.py - python3 python/tutorials/tle/02-moe_align_block_size.py + python3 - <<'PY' + import importlib.util + import pathlib + import sys + import urllib.error + + path = pathlib.Path("python/tutorials/tle/02-moe_align_block_size.py") + spec = importlib.util.spec_from_file_location("moe_align_force_429", path) + mod = importlib.util.module_from_spec(spec) + sys.modules[spec.name] = mod + spec.loader.exec_module(mod) + + def fail_urlopen(url, *args, **kwargs): + raise urllib.error.HTTPError( + url, + 429, + "Too Many Requests", + hdrs=None, + fp=None, + ) + + mod.urllib.request.urlopen = fail_urlopen + mod.main([]) + PY python3 python/tutorials/tle/03-topk.py echo "[debug] triton cache before cluster-gemm" du -sh ~/.triton/cache || true From 6274c07570da80e4dbcb2f586357edece4d7308f Mon Sep 17 00:00:00 2001 From: sunnycase Date: Wed, 8 Jul 2026 11:11:22 +0000 Subject: [PATCH 16/18] Run original cluster GEMM loop after forced MoE 429 --- .github/workflows/hopper-build-and-test.yml | 170 +------------------- 1 file changed, 5 insertions(+), 165 deletions(-) diff --git a/.github/workflows/hopper-build-and-test.yml b/.github/workflows/hopper-build-and-test.yml index 3072ea548..84a8550a5 100644 --- a/.github/workflows/hopper-build-and-test.yml +++ b/.github/workflows/hopper-build-and-test.yml @@ -118,170 +118,10 @@ jobs: mod.main([]) PY python3 python/tutorials/tle/03-topk.py - echo "[debug] triton cache before cluster-gemm" + echo "[debug] triton cache before cluster-gemm original loop" du -sh ~/.triton/cache || true find ~/.triton/cache -type f 2>/dev/null | wc -l || true - echo "[debug] audit remote cluster-gemm PTX" - python3 - <<'PY' - import collections - import importlib.util - import pathlib - import re - import sys - - import torch - - path = pathlib.Path("python/tutorials/tle/04-cluster-gemm.py") - spec = importlib.util.spec_from_file_location("cluster_gemm_audit", path) - mod = importlib.util.module_from_spec(spec) - sys.modules[spec.name] = mod - spec.loader.exec_module(mod) - - torch.manual_seed(0) - M = N = K = 4096 - dtype = torch.float16 - a = torch.randn((M, K), device="cuda", dtype=dtype) - b = torch.randn((K, N), device="cuda", dtype=dtype) - c = torch.empty((M, N), device="cuda", dtype=dtype) - - op_re = re.compile(r"\b(?:ld|st)\.shared::cluster(?:\.v\d+)?\.b\d+") - for idx, cfg in enumerate(mod.REMOTE_TUNE_CONFIGS): - dot_k = mod._select_remote_dot_k(cfg.bk) - use_mask = (M % cfg.bm != 0) or (N % cfg.bn != 0) or (K % cfg.bk != 0) - use_nv_mma_smem_layout = (cfg.bk == 32) or (cfg.bk == 64 and cfg.num_stages <= 2) - compiled = mod._cluster_remote_gemm_kernel.warmup( - a, - b, - c, - M, - N, - K, - a.stride(0), - a.stride(1), - b.stride(0), - b.stride(1), - c.stride(0), - c.stride(1), - mesh=mod.BLOCK_CLUSTER_MESH, - BM=cfg.bm, - BN=cfg.bn, - BK=cfg.bk, - DOT_K=dot_k, - CLUSTER_SIZE=2, - USE_MASK=use_mask, - A_SLOTS=2, - USE_NV_MMA_SMEM_LAYOUT=use_nv_mma_smem_layout, - grid=mod._grid_cluster_remote(M, N, cfg.bm, cfg.bn), - num_ctas=1, - num_warps=cfg.num_warps, - num_stages=cfg.num_stages, - ) - ptx = compiled.asm.get("ptx", "") - ops = collections.Counter(op_re.findall(ptx)) - print(f"[audit] remote candidate_{idx}: {cfg}", flush=True) - print( - f"[audit] cluster_dims={tuple(compiled.metadata.cluster_dims)} " - f"use_nv_mma_smem_layout={use_nv_mma_smem_layout}", - flush=True, - ) - print( - f"[audit] mapa.shared::cluster={ptx.count('mapa.shared::cluster')} " - f"barrier.cluster={ptx.count('barrier.cluster')}", - flush=True, - ) - for op, count in sorted(ops.items()): - print(f"[audit] {op}: {count}", flush=True) - PY - if command -v compute-sanitizer >/dev/null 2>&1; then - echo "[debug] compute-sanitizer memcheck for remote candidates" - compute-sanitizer --tool memcheck --error-exitcode=86 python3 - <<'PY' - import importlib.util - import pathlib - import sys - - import torch - - path = pathlib.Path("python/tutorials/tle/04-cluster-gemm.py") - spec = importlib.util.spec_from_file_location("cluster_gemm_sanitizer", path) - mod = importlib.util.module_from_spec(spec) - sys.modules[spec.name] = mod - spec.loader.exec_module(mod) - - torch.manual_seed(0) - M, N, K = 64, 1024, 256 - dtype = torch.float16 - a = torch.randn((M, K), device="cuda", dtype=dtype) - b = torch.randn((K, N), device="cuda", dtype=dtype) - c = torch.empty((M, N), device="cuda", dtype=dtype) - - for idx, cfg in enumerate(mod.REMOTE_TUNE_CONFIGS): - print(f"[sanitizer] running remote candidate_{idx}: {cfg}", flush=True) - mod._run_cluster_remote( - a, - b, - c, - cfg.bm, - cfg.bn, - cfg.bk, - cfg.num_warps, - cfg.num_stages, - ) - torch.cuda.synchronize() - print(f"[sanitizer] remote candidate_{idx}: PASS", flush=True) - PY - else - echo "[debug] compute-sanitizer not found; skipping memcheck" - fi - echo "[debug] cluster-gemm logged diagnostic trial" - CLUSTER_GEMM_TRIAL="diagnostic" python3 - <<'PY' - import importlib.util - import os - import pathlib - import sys - - import torch - - path = pathlib.Path("python/tutorials/tle/04-cluster-gemm.py") - spec = importlib.util.spec_from_file_location("cluster_gemm_debug", path) - mod = importlib.util.module_from_spec(spec) - sys.modules[spec.name] = mod - spec.loader.exec_module(mod) - - trial = os.environ.get("CLUSTER_GEMM_TRIAL", "?") - progress = pathlib.Path(f"/tmp/cluster-gemm-candidate-{trial}.txt") - - def logged_autotune(name, candidates, run_fn, M, N, K, warmup, rep): - best_cfg = candidates[0] - best_ms = float("inf") - for idx, cfg in enumerate(candidates): - progress.write_text(f"trial {trial} {name} candidate_{idx}: {cfg}\n") - print(f"[debug] trial {trial} running {name} candidate_{idx}: {cfg}", flush=True) - run_fn(cfg) - torch.cuda.synchronize() - ms = mod.triton.testing.do_bench( - lambda: run_fn(cfg), - warmup=warmup, - rep=rep, - ) - progress.write_text(f"trial {trial} {name} candidate_{idx}: {cfg} PASS ms={ms}\n") - print(f"[debug] trial {trial} {name} candidate_{idx}: PASS ms={ms}", flush=True) - if ms < best_ms: - best_ms = ms - best_cfg = cfg - print( - f"[autotune] {name}: best cfg={best_cfg} " - f"ms={best_ms:.3f} tflops={mod._tflops(M, N, K, best_ms):.2f}" - ) - return best_cfg - - mod._autotune_config = logged_autotune - - try: - print(f"[debug] cluster-gemm trial {trial} start", flush=True) - mod.main([]) - except BaseException: - print("[debug] last cluster-gemm candidate:", flush=True) - if progress.exists(): - print(progress.read_text(), flush=True) - raise - PY + for trial in $(seq 1 50); do + echo "[debug] original cluster-gemm trial ${trial}/50" + python3 python/tutorials/tle/04-cluster-gemm.py + done From e03469996c031220330b22f30ffee6b59d4d8d35 Mon Sep 17 00:00:00 2001 From: sunnycase Date: Thu, 9 Jul 2026 03:20:39 +0000 Subject: [PATCH 17/18] Stress original cluster GEMM loop after MoE 429 --- .github/workflows/hopper-build-and-test.yml | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/.github/workflows/hopper-build-and-test.yml b/.github/workflows/hopper-build-and-test.yml index 84a8550a5..7e87ab43c 100644 --- a/.github/workflows/hopper-build-and-test.yml +++ b/.github/workflows/hopper-build-and-test.yml @@ -121,7 +121,7 @@ jobs: echo "[debug] triton cache before cluster-gemm original loop" du -sh ~/.triton/cache || true find ~/.triton/cache -type f 2>/dev/null | wc -l || true - for trial in $(seq 1 50); do - echo "[debug] original cluster-gemm trial ${trial}/50" + for trial in $(seq 1 200); do + echo "[debug] original cluster-gemm trial ${trial}/200" python3 python/tutorials/tle/04-cluster-gemm.py done From adce48fd656756814b54460501fef60b980fb7aa Mon Sep 17 00:00:00 2001 From: sunnycase Date: Thu, 9 Jul 2026 05:36:51 +0000 Subject: [PATCH 18/18] Log cluster GEMM autotune candidates in CI repro --- .github/workflows/hopper-build-and-test.yml | 28 +++++++++++++++++++++ 1 file changed, 28 insertions(+) diff --git a/.github/workflows/hopper-build-and-test.yml b/.github/workflows/hopper-build-and-test.yml index 7e87ab43c..2316e6e71 100644 --- a/.github/workflows/hopper-build-and-test.yml +++ b/.github/workflows/hopper-build-and-test.yml @@ -118,6 +118,34 @@ jobs: mod.main([]) PY python3 python/tutorials/tle/03-topk.py + python3 - <<'PY' + from pathlib import Path + + path = Path("python/tutorials/tle/04-cluster-gemm.py") + text = path.read_text() + old = "\n".join([ + " for cfg in candidates:", + " run_fn(cfg)", + " torch.cuda.synchronize()", + " ms = triton.testing.do_bench(lambda: run_fn(cfg), warmup=warmup, rep=rep)", + " if ms < best_ms:", + "", + ]) + new = "\n".join([ + " for cfg in candidates:", + " print(f\"[autotune-candidate] {name}: start cfg={cfg}\", flush=True)", + " run_fn(cfg)", + " torch.cuda.synchronize()", + " print(f\"[autotune-candidate] {name}: initial_sync_ok cfg={cfg}\", flush=True)", + " ms = triton.testing.do_bench(lambda: run_fn(cfg), warmup=warmup, rep=rep)", + " print(f\"[autotune-candidate] {name}: bench_ok cfg={cfg} ms={ms:.3f}\", flush=True)", + " if ms < best_ms:", + "", + ]) + if old not in text: + raise SystemExit("cluster-gemm autotune patch target not found") + path.write_text(text.replace(old, new)) + PY echo "[debug] triton cache before cluster-gemm original loop" du -sh ~/.triton/cache || true find ~/.triton/cache -type f 2>/dev/null | wc -l || true