Add fast dispatch cache to Kernel.__call__ (#3009) #18634
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| name: Tests | |
| on: | |
| pull_request: | |
| push: | |
| branches: | |
| - main | |
| - release/* | |
| concurrency: | |
| group: ${{ github.workflow }}-${{ github.ref == 'refs/heads/main' && github.run_number || github.ref }} | |
| cancel-in-progress: true | |
| permissions: | |
| contents: read | |
| jobs: | |
| load-matrix: | |
| runs-on: ubuntu-latest | |
| outputs: | |
| matrix: ${{ steps.set-matrix.outputs.matrix }} | |
| steps: | |
| - name: Checkout repository | |
| uses: actions/checkout@v7 | |
| - name: Load matrix from file | |
| id: set-matrix | |
| run: | | |
| matrix=$(cat .github/matrix.json | jq -c .) | |
| echo "matrix=$matrix" >> $GITHUB_OUTPUT | |
| test: | |
| needs: load-matrix | |
| strategy: | |
| fail-fast: false | |
| matrix: ${{ fromJSON(needs.load-matrix.outputs.matrix) }} | |
| name: test-${{ matrix.runtime-version }}-py${{ matrix.python-version }}-${{ matrix.pytorch-version }}-${{ matrix.backend }}-${{ matrix.alias }} | |
| container: ${{ matrix.image != '' && fromJSON(format('{{"image":"{0}","options":"{1}"}}', matrix.image, matrix.container-options)) || '' }} | |
| runs-on: ${{ matrix.runner }} | |
| defaults: | |
| run: | |
| shell: bash -l {0} | |
| steps: | |
| - name: Run NVIDIA command | |
| if: startsWith(matrix.image, 'nvidia') | |
| run: | | |
| echo "Detected NVIDIA image" | |
| nvidia-smi || echo "nvidia-smi not found" | |
| - name: Run ROCm command | |
| if: startsWith(matrix.image, 'rocm') | |
| run: | | |
| echo "Detected ROCm image" | |
| rocminfo || echo "rocminfo not found" | |
| - name: Run XPU command | |
| if: startsWith(matrix.image, 'intel') | |
| run: | | |
| echo "Detected XPU image" | |
| xpu-smi discovery || echo "xpu-smi not found" | |
| - name: Check out code | |
| uses: actions/checkout@v7 | |
| - name: Install system dependencies | |
| if: runner.os == 'Linux' | |
| run: | | |
| set -eux | |
| SUDO=$(command -v sudo 2>/dev/null || true) | |
| retry () { | |
| "$@" || (sleep 10 && "$@") || (sleep 20 && "$@") || (sleep 40 && "$@") | |
| } | |
| retry $SUDO apt-get update | |
| retry $SUDO apt-get install -y --fix-missing libdw1 curl wget git pkg-config zlib1g-dev build-essential | |
| - name: Install NVSHMEM | |
| if: contains(matrix.alias, 'distributed') | |
| run: | | |
| set -euxo pipefail | |
| GPU_COUNT=$(nvidia-smi -L | wc -l) | |
| if [ "$GPU_COUNT" -ne 4 ]; then | |
| echo "Error: Expected 4 GPUs but found $GPU_COUNT" | |
| exit 1 | |
| fi | |
| curl -L https://raw.githubusercontent.com/pytorch/pytorch/main/.ci/docker/common/install_cuda.sh -o install_cuda.sh | |
| chmod +x install_cuda.sh | |
| source install_cuda.sh | |
| install_nvshmem 13 3.4.5 | |
| - name: Install uv | |
| uses: astral-sh/setup-uv@v7 | |
| with: | |
| python-version: ${{ matrix.python-version }} | |
| enable-cache: true | |
| - name: Create virtual environment | |
| run: | | |
| uv venv --python ${{ matrix.python-version }} | |
| - name: Get current month | |
| id: date | |
| run: echo "month=$(date +'%Y-%m')" >> $GITHUB_OUTPUT | |
| - name: Cache dependencies | |
| id: cache | |
| uses: actions/cache@v6 | |
| with: | |
| path: | | |
| ~/.cache/uv | |
| ~/.venv | |
| key: ${{ matrix.python-version }}-${{ matrix.runtime-version }}-${{ matrix.pytorch-version }}-${{ matrix.backend }}-${{ hashFiles('.github/workflows/test.yml') }}-${{ steps.date.outputs.month }} | |
| - name: Install PyTorch | |
| run: | | |
| set -eux | |
| source .venv/bin/activate | |
| if [[ "${{ matrix.pytorch-version }}" == pytorch-[0-9]* ]]; then | |
| VERSION="${{ matrix.pytorch-version }}" | |
| VERSION="${VERSION#pytorch-}" | |
| uv pip install -U "torch==${VERSION}.*" --index-url https://download.pytorch.org/whl/${{ matrix.runtime-version }} | |
| elif [[ "${{ matrix.runtime-version }}" == "tpu" ]]; then | |
| # Use the latest PyTorch wheel known to work with TorchTPU. | |
| uv pip install -U --pre "torch==2.14.0.dev20260707" --index-url https://download.pytorch.org/whl/nightly/cpu | |
| else | |
| # Default to nightly | |
| # On OSDC, route nightly through the pypi-cache /whl proxy (deps stay reachable behind the egress firewall) and clear the cache index vars so it wins over the cache default; off OSDC PYPI_CACHE_WHL_URL is unset so we hit download.pytorch.org. | |
| PYTORCH_INDEX="${PYPI_CACHE_WHL_URL:-https://download.pytorch.org/whl/cpu}" | |
| UV_DEFAULT_INDEX= UV_INDEX= uv pip install -U --pre torch --index-url "${PYTORCH_INDEX%/whl/*}/whl/nightly/${{ matrix.runtime-version }}" | |
| fi | |
| - name: Configure XPU runtime library path | |
| if: matrix.alias == 'xpu' | |
| run: | | |
| set -euo pipefail | |
| source .venv/bin/activate | |
| echo "LD_LIBRARY_PATH=$VIRTUAL_ENV/lib:${LD_LIBRARY_PATH:-}" >> "$GITHUB_ENV" | |
| - name: Install Triton from source | |
| if: matrix.triton-pin != '' | |
| run: | | |
| set -x | |
| source .venv/bin/activate | |
| apt-get update | |
| apt-get install -y git clang-20 clang++-20 zlib1g-dev | |
| export CC=clang-20 | |
| export CXX=clang++-20 | |
| mkdir -p /tmp/$USER | |
| cd /tmp/$USER | |
| uv pip uninstall triton pytorch-triton triton-xpu || true | |
| rm -rf triton/ || true | |
| REPO="${{ matrix.triton-repo || 'https://github.com/triton-lang/triton.git' }}" | |
| git clone --recursive "$REPO" triton | |
| git -C triton checkout ${{ matrix.triton-pin }} | |
| cd triton/ | |
| uv pip install -r python/requirements.txt | |
| # Disable proton on XPU: proton's NVIDIA sources (CuptiApi.cpp, NvtxApi.cpp) | |
| # require CUDA headers that are not present on Intel XPU runners. | |
| if [ "${{ matrix.alias }}" == "xpu" ]; then | |
| export TRITON_BUILD_PROTON=OFF | |
| fi | |
| MAX_JOBS=$(nproc) TRITON_PARALLEL_LINK_JOBS=2 uv pip install . | |
| cd /tmp/$USER | |
| rm -rf triton/ | |
| python -c "import triton; print(f'Triton version: {triton.__version__}')" | |
| - name: Pin networkx for Python 3.14 | |
| if: startsWith(matrix.python-version, '3.14') | |
| run: | | |
| source .venv/bin/activate | |
| uv pip install networkx==2.8.8 | |
| - name: Install Helion | |
| run: | | |
| source .venv/bin/activate | |
| uv pip install setuptools ninja | |
| SETUPTOOLS_SCM_PRETEND_VERSION="0.0.0" uv pip install -e .'[dev]' | |
| python -c "import helion; print(helion.__name__)" | |
| - name: Install TPU dependencies (Pallas) | |
| if: matrix.alias == 'tpu' | |
| run: | | |
| set -euxo pipefail | |
| source .venv/bin/activate | |
| uv pip install \ | |
| --extra-index-url https://us-python.pkg.dev/ml-oss-artifacts-published/jax/simple/ \ | |
| --find-links https://storage.googleapis.com/jax-releases/libtpu_releases.html \ | |
| --pre \ | |
| 'jax==0.10.1' 'jaxlib==0.10.1' 'libtpu==0.0.40' 'tpu-info==0.7.1' 'jaxtyping' 'frozendict' 'immutabledict' | |
| # Install Bazel | |
| if ! command -v bazel &> /dev/null; then | |
| sudo curl -L https://github.com/bazelbuild/bazelisk/releases/download/v1.27.0/bazelisk-linux-amd64 -o /usr/local/bin/bazel | |
| sudo chmod +x /usr/local/bin/bazel | |
| fi | |
| # Install gcloud CLI if not present (needed for Secret Manager) | |
| if ! command -v gcloud &> /dev/null; then | |
| sudo apt-get install -y apt-transport-https ca-certificates gnupg curl | |
| curl https://packages.cloud.google.com/apt/doc/apt-key.gpg | sudo gpg --dearmor -o /usr/share/keyrings/cloud.google.gpg | |
| echo "deb [signed-by=/usr/share/keyrings/cloud.google.gpg] https://packages.cloud.google.com/apt cloud-sdk main" | sudo tee /etc/apt/sources.list.d/google-cloud-sdk.list | |
| sudo apt-get update && sudo apt-get install -y google-cloud-cli | |
| fi | |
| # Clone torch_tpu via GCP Secret Manager SSH key (same as pytorch CI) | |
| TORCH_TPU_COMMIT=$(cat .github/ci_commit_pins/torch_tpu.txt) | |
| set +x | |
| gcloud secrets versions access latest \ | |
| --secret="torchtpu-read-key" \ | |
| --project="ml-velocity-actions-testing" > /tmp/torch_tpu_ssh_key | |
| set -x | |
| chmod 600 /tmp/torch_tpu_ssh_key | |
| GIT_SSH_COMMAND="ssh -i /tmp/torch_tpu_ssh_key -o IdentitiesOnly=yes -o StrictHostKeyChecking=no" \ | |
| git clone git@github.com:google-pytorch/torch_tpu.git /tmp/torch_tpu | |
| rm -f /tmp/torch_tpu_ssh_key | |
| cd /tmp/torch_tpu | |
| git checkout "${TORCH_TPU_COMMIT}" | |
| # Build torch_tpu wheel | |
| sudo apt-get install -y libpython3-dev | |
| export TORCH_SOURCE=$(python -c "import torch; import os; print(os.path.dirname(os.path.dirname(torch.__file__)))") | |
| bazel build -c opt //ci/wheel:torch_tpu_wheel \ | |
| --config=helion_public_caching_readwrite \ | |
| --define TORCH_SOURCE=local \ | |
| --repo_env=TORCH_SOURCE=$TORCH_SOURCE \ | |
| --action_env=JAX_PLATFORMS=cpu | |
| uv pip install bazel-bin/ci/wheel/*.whl | |
| # Ignore the torch_tpu wheel libtpu pin, and just use the most recent libtpu. | |
| uv pip install --upgrade libtpu | |
| cd - | |
| rm -rf /tmp/torch_tpu | |
| # Verify | |
| python -c "import torch, sys; print('Success') if torch.tpu.is_available() else (print('(Torch)TPU not available'), sys.exit(1))" | |
| - name: Install Pallas interpret dependencies | |
| if: matrix.alias == 'pallas-interpret' | |
| run: | | |
| source .venv/bin/activate | |
| uv pip install 'jax==0.10.1' 'jaxlib==0.10.1' 'absl-py' | |
| - name: Install CUTLASS CuTe DSL | |
| if: matrix.backend == 'cute' | |
| run: | | |
| source .venv/bin/activate | |
| ./scripts/install_cute.sh | |
| - name: CUDA Compute Check | |
| if: startsWith(matrix.image, 'nvidia') | |
| run: | | |
| source .venv/bin/activate | |
| python -c " | |
| import torch, sys | |
| assert torch.cuda.is_available(), 'FATAL: CUDA not available' | |
| n = torch.cuda.device_count() | |
| assert n > 0, 'FATAL: No CUDA devices found' | |
| print(f'CUDA devices: {n}') | |
| for i in range(n): | |
| dev = torch.device('cuda', i) | |
| a = torch.randn(256, 256, device=dev) | |
| b = (a @ a).sum().item() | |
| print(f' Device {i} ({torch.cuda.get_device_name(i)}): OK') | |
| print(f'All {n} devices healthy') | |
| " | |
| - name: Inductor Worker Check | |
| if: startsWith(matrix.image, 'nvidia') | |
| run: | | |
| source .venv/bin/activate | |
| python -c " | |
| import torch | |
| @torch.compile | |
| def f(x): | |
| return (x + 1).relu() | |
| x = torch.randn(8, device='cuda') | |
| f(x).sum().item() | |
| print('Inductor worker subprocess: OK') | |
| " | |
| - name: XPU Compute Check | |
| if: startsWith(matrix.image, 'intel') | |
| run: | | |
| source .venv/bin/activate | |
| python -c " | |
| import torch | |
| assert torch.xpu.is_available(), 'FATAL: XPU not available' | |
| n = torch.xpu.device_count() | |
| assert n > 0, 'FATAL: No XPU devices found' | |
| print(f'XPU devices: {n}') | |
| for i in range(n): | |
| dev = torch.device('xpu', i) | |
| a = torch.randn(256, 256, device=dev) | |
| b = (a @ a).sum().item() | |
| print(f' Device {i} ({torch.xpu.get_device_name(i)}): OK') | |
| print(f'All {n} devices healthy') | |
| " | |
| - name: MPS Availability Check | |
| if: matrix.backend == 'metal' | |
| run: | | |
| source .venv/bin/activate | |
| python -c "import torch; assert torch.backends.mps.is_available(), 'MPS not available'" | |
| - name: Run Tests | |
| run: | | |
| set -o pipefail | |
| source .venv/bin/activate | |
| uv pip install pytest-xdist | |
| uv pip install pytest-rerunfailures | |
| # Conditionally enable ref-eager and golden-accept/dtype-assert test modes | |
| if [[ "${{ matrix.dtype-asserts }}" == "true" ]]; then export HELION_DEBUG_DTYPE_ASSERTS=1; fi | |
| if [[ "${{ matrix.expecttest-accept }}" == "true" ]]; then export EXPECTTEST_ACCEPT=1; fi | |
| if [[ "${{ matrix.ref-eager }}" == "true" ]]; then export HELION_INTERPRET=1; fi | |
| if [[ "${{ matrix.backend }}" == "tileir" ]]; then export ENABLE_TILE=1; fi | |
| if [[ "${{ matrix.alias }}" == "pallas-interpret" ]]; then export HELION_PALLAS_INTERPRET=1; fi | |
| export HELION_BACKEND=${{ matrix.backend }} | |
| # Disable TorchTPU build cache to prevent running out of disk space. | |
| export TORCH_TPU_INTERNAL_TIER2_COMPILATION_CACHE=disabled | |
| # -rf: print failed tests | |
| # --timeout: max allowed time for each test | |
| PARALLEL="-n4" | |
| if [[ "${{ contains(matrix.alias, 'distributed') }}" == "true" ]]; then | |
| TEST_PATH="test/test_examples_dist.py" | |
| EXTRA_FLAGS="-rs" | |
| elif [[ "${{ matrix.alias }}" == "tpu" || "${{ matrix.alias }}" == "pallas-interpret" || "$HELION_BACKEND" == "metal" ]]; then | |
| TEST_PATH="." | |
| EXTRA_FLAGS="--ignore=test/test_examples_dist.py" | |
| PARALLEL="" | |
| else | |
| TEST_PATH="." | |
| EXTRA_FLAGS="--ignore=test/test_examples_dist.py" | |
| fi | |
| # A10G runners expose one 22 GiB GPU. Four xdist workers can collectively | |
| # exhaust VRAM with CUDA/Triton/Inductor state even when an individual test | |
| # only needs a small allocation. | |
| if [[ -n "$PARALLEL" && "${{ matrix.alias }}" == *a10g* ]]; then | |
| PARALLEL="-n2" | |
| fi | |
| # For distributed tests, fail if any test is skipped, failed, or has an error | |
| SKIP_CHECK=$([[ "${{ contains(matrix.alias, 'distributed') }}" == "true" ]] && echo "! grep -qE '(SKIPPED|FAILED|ERROR)'" || echo "cat > /dev/null") | |
| if [[ "${{ matrix.alias }}" == *xpu* ]]; then | |
| export TRITON_XPU_GEN_NATIVE_CODE=1 | |
| fi | |
| TIMEOUT=60 | |
| if [[ "${{ matrix.alias }}" == "tpu" ]]; then | |
| TIMEOUT=300 | |
| elif [[ "${{ matrix.alias }}" == "pallas-interpret" ]]; then | |
| TIMEOUT=300 | |
| elif [[ "${{ matrix.alias }}" == *a10g* ]]; then | |
| # A10G (sm86) ptxas cold-compiles some backward kernels slowly | |
| # (~75s for the softmax autodiff backward), so the default 60s | |
| # per-test timeout false-positives on a cold Triton cache and the | |
| # worker gets killed mid-compile. Give it headroom. | |
| TIMEOUT=180 | |
| fi | |
| pytest $PARALLEL -rf --timeout=$TIMEOUT --reruns=2 --timeout-method=thread $EXTRA_FLAGS $TEST_PATH | tee >(eval $SKIP_CHECK) | |
| test-notebooks: | |
| name: test-notebooks-cu130-py3.12-pytorch-2.9-a10g | |
| container: | |
| image: nvidia/cuda:13.1.0-devel-ubuntu24.04 | |
| options: --gpus all | |
| runs-on: linux.g5.4xlarge.nvidia.gpu | |
| defaults: | |
| run: | |
| shell: bash -l {0} | |
| steps: | |
| - name: Run NVIDIA command | |
| run: | | |
| echo "Detected NVIDIA image" | |
| nvidia-smi || echo "nvidia-smi not found" | |
| - name: Check out code | |
| uses: actions/checkout@v7 | |
| - name: Install uv | |
| uses: astral-sh/setup-uv@v7 | |
| with: | |
| python-version: "3.12" | |
| enable-cache: true | |
| - name: Create virtual environment | |
| run: | | |
| uv venv --python 3.12 | |
| - name: Install pip in venv | |
| run: | | |
| source .venv/bin/activate | |
| uv pip install pip | |
| - name: Get current month | |
| id: date | |
| run: echo "month=$(date +'%Y-%m')" >> $GITHUB_OUTPUT | |
| - name: Cache dependencies | |
| id: cache | |
| uses: actions/cache@v6 | |
| with: | |
| path: | | |
| ~/.cache/uv | |
| ~/.venv | |
| key: notebooks-3.12-cu130-${{ hashFiles('.github/workflows/test.yml') }}-${{ steps.date.outputs.month }} | |
| - name: Install notebook execution tools | |
| run: | | |
| source .venv/bin/activate | |
| # Install jupyter for executing notebooks | |
| uv pip install jupyter nbconvert pytest numpy "nbclient<0.10" | |
| - name: Run Notebook Tests | |
| run: | | |
| source .venv/bin/activate | |
| # Execute notebook using jupyter nbconvert | |
| # The notebook's subprocess pip install will install torch and helion | |
| jupyter nbconvert --to notebook --execute --inplace \ | |
| --ExecutePreprocessor.timeout=600 \ | |
| notebooks/softmax.ipynb |