Skip to content

Commit 8a15116

Browse files
SumanthRHclaudegemini-code-assist[bot]
authored
[chore] Upgrade flashinfer to 0.6.14 (#2135)
# What does this PR do? Bumps flashinfer packages (`flashinfer-python`, `flashinfer-cubin` and `flashinfer-jit-cache`) to 0.6.14. For `flashinfer-cubin` we need to use the wheel on flashinfer's index instead of the wheel on PyPI <!-- CURSOR_SUMMARY --> --- > [!NOTE] > **Medium Risk** > Touches the pinned GPU inference dependency graph (flashinfer + lockfile) used by vLLM on Linux; install/resolution regressions are the main risk, with no production training code changes beyond a small test env fix. > > **Overview** > Aligns the Linux x86_64 **flashinfer** stack (`flashinfer-python`, `flashinfer-jit-cache`, `flashinfer-cubin`) with **0.6.14**, matching what **vLLM 0.26.0** expects instead of the previous 0.6.13 pin. Override comments in `pyproject.toml` are updated to reflect that rationale (including the ≥0.6.13 MNNVL fusion requirement). > > Because **flashinfer-cubin** is not published on PyPI beyond 0.6.13, the PR adds an explicit **`flashinfer`** uv index (`https://flashinfer.ai/whl/`) and wires **`flashinfer-cubin`** through `[tool.uv.sources]` while **jit-cache** stays on the cu130 index. **`uv.lock`** is refreshed accordingly (including transitive bumps such as **`cuda-tile`** 1.5.0 tied to the new flashinfer release). > > GPU test Ray init only forwards **`LD_LIBRARY_PATH`** into the runtime env when it is already set, avoiding propagating a missing/`None` value. > > <sup>Reviewed by [Cursor Bugbot](https://cursor.com/bugbot) for commit 9b86de8. Bugbot is set up for automated code reviews on this repo. Configure [here](https://www.cursor.com/dashboard/bugbot).</sup> <!-- /CURSOR_SUMMARY --> --------- Signed-off-by: SumanthRH <sumanthrh@anyscale.com> Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com> Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
1 parent 25afc27 commit 8a15116

3 files changed

Lines changed: 101 additions & 103 deletions

File tree

‎pyproject.toml‎

Lines changed: 28 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -125,9 +125,9 @@ fsdp = [
125125
"causal-conv1d==1.6.2.post1+cu.13.0.torch.2.11; sys_platform == 'linux'",
126126
"flash-attn==2.8.3+cu.13.0.torch.2.11; sys_platform == 'linux'",
127127
"torch==2.11.0; sys_platform == 'linux'",
128-
"flashinfer-python==0.6.13; sys_platform == 'linux' and platform_machine == 'x86_64'",
129-
"flashinfer-jit-cache==0.6.13; sys_platform == 'linux' and platform_machine == 'x86_64'",
130-
"flashinfer-cubin==0.6.13; sys_platform == 'linux' and platform_machine == 'x86_64'",
128+
"flashinfer-python==0.6.14; sys_platform == 'linux' and platform_machine == 'x86_64'",
129+
"flashinfer-jit-cache==0.6.14; sys_platform == 'linux' and platform_machine == 'x86_64'",
130+
"flashinfer-cubin==0.6.14; sys_platform == 'linux' and platform_machine == 'x86_64'",
131131
"torchvision; sys_platform == 'linux'",
132132
]
133133

@@ -158,14 +158,14 @@ megatron = [
158158
# into the image for a module the shim would never load).
159159
"nixl; sys_platform == 'linux'",
160160
"torch==2.11.0; sys_platform == 'linux'",
161-
"flashinfer-python==0.6.13; sys_platform == 'linux' and platform_machine == 'x86_64'",
161+
"flashinfer-python==0.6.14; sys_platform == 'linux' and platform_machine == 'x86_64'",
162162
"torchvision; sys_platform == 'linux'",
163163
# megatron-bridge requires Python 3.12+; pin megatron-core to the same
164164
# constraint so both packages are consistently available (or absent).
165165
"megatron-bridge; sys_platform == 'linux' and python_version >= '3.12'",
166166
"megatron-core; sys_platform == 'linux' and python_version >= '3.12'",
167-
"flashinfer-jit-cache==0.6.13; sys_platform == 'linux' and platform_machine == 'x86_64'",
168-
"flashinfer-cubin==0.6.13; sys_platform == 'linux' and platform_machine == 'x86_64'",
167+
"flashinfer-jit-cache==0.6.14; sys_platform == 'linux' and platform_machine == 'x86_64'",
168+
"flashinfer-cubin==0.6.14; sys_platform == 'linux' and platform_machine == 'x86_64'",
169169
"nvidia-modelopt; sys_platform == 'linux'",
170170
"fast-hadamard-transform; sys_platform == 'linux'",
171171
]
@@ -250,15 +250,19 @@ override-dependencies = [
250250
# so it doesn't pull a second, unused CUDA runtime in alongside the cu13 one.
251251
"nixl-cu12; sys_platform == 'never'",
252252
# Megatron-Bridge pins flashinfer-python==0.6.8.post1, which conflicts with
253-
# our pin. vLLM 0.26.0's metadata asks for 0.6.14, but flashinfer-cubin was
254-
# never released at 0.6.14 and flashinfer hard-errors when cubin's version
255-
# differs from its own, so all three are held at 0.6.13 -- the version vLLM
256-
# 0.26.0 was developed against. It must be >=0.6.13: older flashinfer rejects
257-
# the `layout_code` vLLM's allreduce+RMS fusion pass hands the MNNVL backend
258-
# ("MNNVL AllReduce does not support quantization fusion").
259-
"flashinfer-python==0.6.13; sys_platform == 'linux' and platform_machine == 'x86_64'",
260-
"flashinfer-jit-cache==0.6.13; sys_platform == 'linux' and platform_machine == 'x86_64'",
261-
"flashinfer-cubin==0.6.13; sys_platform == 'linux' and platform_machine == 'x86_64'",
253+
# our pin, so all three flashinfer packages are held here at 0.6.14 -- the
254+
# version vLLM 0.26.0's own metadata asks for. flashinfer hard-errors when
255+
# cubin's version differs from its own, so the three must move together.
256+
# It must be >=0.6.13: older flashinfer rejects the `layout_code` vLLM's
257+
# allreduce+RMS fusion pass hands the MNNVL backend ("MNNVL AllReduce does
258+
# not support quantization fusion").
259+
#
260+
# flashinfer-cubin is not published to PyPI past 0.6.13, so it (like
261+
# flashinfer-jit-cache) comes from flashinfer's own index -- see the
262+
# `flashinfer` index and [tool.uv.sources] below.
263+
"flashinfer-python==0.6.14; sys_platform == 'linux' and platform_machine == 'x86_64'",
264+
"flashinfer-jit-cache==0.6.14; sys_platform == 'linux' and platform_machine == 'x86_64'",
265+
"flashinfer-cubin==0.6.14; sys_platform == 'linux' and platform_machine == 'x86_64'",
262266
# mamba-ssm 2.3.2.post1 pins tilelang==0.1.8 and apache-tvm-ffi<=0.1.9, while
263267
# vLLM 0.26.0 pins tilelang==0.1.9 and apache-tvm-ffi==0.1.10. Both overrides
264268
# hold these at the versions the lock already resolved to before mamba-ssm moved
@@ -305,6 +309,14 @@ name = "flashinfer-cu130"
305309
url = "https://flashinfer.ai/whl/cu130"
306310
explicit = true
307311

312+
# flashinfer's CUDA-agnostic index. flashinfer-cubin is a pure-python wheel
313+
# (py3-none-any) that stopped being published to PyPI after 0.6.13, so it is
314+
# sourced here instead; the cu130 index above only carries flashinfer-jit-cache.
315+
[[tool.uv.index]]
316+
name = "flashinfer"
317+
url = "https://flashinfer.ai/whl/"
318+
explicit = true
319+
308320
# Astral's prebuilt GPU wheels (https://wheels.astral.sh). cu130 to match torch's
309321
# CUDA variant. Wheels are published per (CUDA, torch) pair and carry that pair in
310322
# a local version segment (e.g. `+cu.13.0.torch.2.11`), so the exact build must be
@@ -318,6 +330,7 @@ explicit = true
318330
skyrl-gym = { path = "./skyrl-gym", editable = true }
319331
# Match torch's CUDA variant (cu130).
320332
flashinfer-jit-cache = { index = "flashinfer-cu130", marker = "sys_platform == 'linux'" }
333+
flashinfer-cubin = { index = "flashinfer", marker = "sys_platform == 'linux'" }
321334

322335
# torch-2.11 builds of these aren't published on PyPI, so they come from Astral's
323336
# prebuilt GPU index instead of wheels built by hand on personal forks. Astral

‎tests/backends/skyrl_train/gpu/utils.py‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -391,7 +391,8 @@ def ray_init_for_tests():
391391
env_vars["PYTHONPATH"] = os.environ.get("PYTHONPATH")
392392
env_vars["CUDA_DEVICE_MAX_CONNECTIONS"] = "1"
393393
env_vars["NVTE_FUSED_ATTN"] = "0"
394-
env_vars["LD_LIBRARY_PATH"] = os.environ.get("LD_LIBRARY_PATH")
394+
if "LD_LIBRARY_PATH" in os.environ:
395+
env_vars["LD_LIBRARY_PATH"] = os.environ["LD_LIBRARY_PATH"]
395396
ray.init(runtime_env={"env_vars": env_vars})
396397

397398

0 commit comments

Comments
 (0)