Agent/pr760 cluster gemm case

.github/workflows/hopper-build-and-test.yml

@@ -51,17 +51,29 @@ 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
-          export USE_FLAGCX=ON
-          MAX_JOBS=32 python3 -m pip install . --no-build-isolation
+          env | grep -E '^(LLVM_SYSPATH)=' >> $GITHUB_ENV || true
+          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' }}
         shell: bash
         run: |
           set -x
           source ~/env.sh
-          ## python tutorials
+          git log -1 --format='[debug] workspace commit %H %s'
+          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
@@ -72,7 +84,6 @@ jobs:
           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 \
@@ -81,40 +92,64 @@ jobs:
           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 - <<'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
-          # python3 python/tutorials/tle/04-cluster-gemm.py  # TODO: sometimes failed
-          python3 python/tutorials/tle/deepseek_v32/01-topk_selector.py
-          python3 python/tutorials/tle/deepseek_v32/02-sparse-mla.py
-          ## flagtree tle 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
+          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
+          for trial in $(seq 1 200); do
+            echo "[debug] original cluster-gemm trial ${trial}/200"
+            python3 python/tutorials/tle/04-cluster-gemm.py
+          done

python/test/tle/unit/test_tle_raw_cache_key.py

@@ -1,6 +1,8 @@
 import ast
 import textwrap
 
+import pytest
+import triton
 import triton.language as tl
 from triton.experimental.tle.raw import dialect
 from triton.experimental.tle.raw.cache_key import (
@@ -137,3 +139,68 @@ 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)