Skip to content

[WIP] Try new Triton package 20260608 #392

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
antiagainst wants to merge 1 commit into from
Closed

[WIP] Try new Triton package 20260608 #392

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
16 changes: 16 additions & 0 deletions test/ci_system/install_deps.sh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -30,6 +30,7 @@ export CPLUS_INCLUDE_PATH="/usr/local/cuda/include/cccl"
export C_INCLUDE_PATH="/usr/local/cuda/include/cccl"

WORKSPACE=${WORKSPACE:-$(pwd)}
TOKENSPEED_TESTPYPI_INDEX=${TOKENSPEED_TESTPYPI_INDEX:-https://test.pypi.org/simple}

# Wrap pip install in a retry loop. PyPI's CDN occasionally returns a
# bad Content-Type for /simple/<pkg>/ pages (most recently observed for
Expand All @@ -55,6 +56,19 @@ pip_install_with_retry() {
done
}

preinstall_tokenspeed_testpypi_packages() {
local requirements_file="${WORKSPACE}/tokenspeed-kernel/python/requirements/common.txt"
local vendor_requirements=()
mapfile -t vendor_requirements < <(grep -E '^tokenspeed-(triton|proton)==' "${requirements_file}")
if [ "${#vendor_requirements[@]}" -eq 0 ]; then
echo "No tokenspeed vendor requirements found in ${requirements_file}" >&2
return 1
fi
pip_install_with_retry pip3 install --no-deps \
--index-url "${TOKENSPEED_TESTPYPI_INDEX}" \
"${vendor_requirements[@]}"
}

echo "=========================================="
echo "SM=${SM}, CUDA_VERSION=${CUDA_VERSION}"
echo "WORKSPACE=${WORKSPACE}"
Expand Down Expand Up @@ -91,6 +105,8 @@ python3 -m pip install --upgrade pip setuptools wheel
# ============================================================
echo "=== Step 3: Install tokenspeed-kernel ==="
cd ${WORKSPACE}
echo "Preinstalling staged TokenSpeed vendor packages from ${TOKENSPEED_TESTPYPI_INDEX}"
preinstall_tokenspeed_testpypi_packages
export PIP_EXTRA_INDEX_URL="https://download.pytorch.org/whl/cu${CUINDEX}"
TOKENSPEED_KERNEL_BACKEND=cuda FLASHINFER_CUDA_ARCH_LIST="${FI_ARCH}" \
pip_install_with_retry pip3 install tokenspeed-kernel/python/ --no-build-isolation -v
Expand Down
15 changes: 15 additions & 0 deletions test/ci_system/install_deps_rocm.sh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,19 @@ ROCM_INDEX="https://download.pytorch.org/whl/rocm${ROCM_VERSION}"

export MAX_JOBS=${BUILD_AND_DOWNLOAD_PARALLEL}
WORKSPACE=${WORKSPACE:-$(pwd)}
TOKENSPEED_TESTPYPI_INDEX=${TOKENSPEED_TESTPYPI_INDEX:-https://test.pypi.org/simple}

preinstall_tokenspeed_testpypi_packages() {
local requirements_file="${WORKSPACE}/tokenspeed-kernel/python/requirements/common.txt"
local vendor_requirements=()
mapfile -t vendor_requirements < <(grep -E '^tokenspeed-(triton|proton)==' "${requirements_file}")
if [ "${#vendor_requirements[@]}" -eq 0 ]; then
echo "No tokenspeed vendor requirements found in ${requirements_file}" >&2
return 1
fi
pip3 install --no-deps --index-url "${TOKENSPEED_TESTPYPI_INDEX}" \
"${vendor_requirements[@]}"
}

echo "=========================================="
echo "GFX_ARCH=${GFX_ARCH}"
Expand All @@ -27,6 +40,8 @@ python3 -m pip install --upgrade pip setuptools wheel

echo "=== Step 3: Install tokenspeed-kernel ==="
cd "${WORKSPACE}"
echo "Preinstalling staged TokenSpeed vendor packages from ${TOKENSPEED_TESTPYPI_INDEX}"
preinstall_tokenspeed_testpypi_packages
export PIP_EXTRA_INDEX_URL="${ROCM_INDEX}"
TOKENSPEED_KERNEL_BACKEND=rocm pip3 install tokenspeed-kernel/python/ \
--no-build-isolation -v
Expand Down
4 changes: 2 additions & 2 deletions tokenspeed-kernel/python/requirements/common.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,7 @@ apache-tvm-ffi>=0.1.5
numpy
packaging
PyYAML>=6.0
tokenspeed-proton>=3.7.10.post20260531
tokenspeed-triton>=3.7.10.post20260531
tokenspeed-proton==3.7.10.post20260608
tokenspeed-triton==3.7.10.post20260608
Comment on lines +5 to +6

@chatgpt-codex-connector chatgpt-codex-connector Bot Jun 9, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P1 Badge Keep staged pins out of release requirements

When this exact post20260608 pin is still only available from TestPyPI, the tokenspeed-kernel release build has no way to resolve it: .github/workflows/release-tokenspeed-kernel.yml builds wheels with only the PyTorch extra index, and setup.py's build path installs requirements/cuda.txt, which includes this file. The new CI install script preinstalls from TestPyPI only for test/ci_system/install_deps.sh, so release wheels and any source install outside that script fail before the kernel build starts. Either make the release/source build use the same index explicitly or avoid committing TestPyPI-only exact pins in the shared requirements.

Useful? React with 👍 / 👎.

wheel
psutil
2 changes: 1 addition & 1 deletion tokenspeed-mla/pyproject.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,7 @@ license = { file = "LICENSE" }
dependencies = [
"apache-tvm-ffi>=0.1.5",
"nvidia-cutlass-dsl",
"tokenspeed-triton>=3.7.10.post20260531",
"tokenspeed-triton==3.7.10.post20260608",
"torch",
]

Expand Down
Loading