diff --git a/.github/workflows/hopper-build-and-test.yml b/.github/workflows/hopper-build-and-test.yml index 9e425d8f0..2316e6e71 100644 --- a/.github/workflows/hopper-build-and-test.yml +++ b/.github/workflows/hopper-build-and-test.yml @@ -51,9 +51,17 @@ 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' }} @@ -61,7 +69,11 @@ jobs: 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 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 a7f2607a7..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,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)