Repository navigation
build: pin tilelang 0.1.11 for the automodel and sglang extras on aarch64 too - #4755
linmuchuiyang wants to merge 2 commits into
Conversation
…ch64 too The `tilelang==0.1.11` pins of the `automodel` and `sglang` extras carry a `platform_machine == 'x86_64'` marker, so on aarch64 (GB200) the Automodel venv resolves tilelang 0.1.12 from the base `tilelang>=0.1.9,<0.1.13` range. `tile-kernels==1.0.0` (pinned by nemo-automodel[moe]) cannot import on 0.1.12 (`tilelang.utils.target` moved; deepseek-ai/TileKernels#24), Automodel's `_HAS_TILE_KERNELS_SINKHORN` becomes False, and DeepSeek-V4.1 -- which forces the TileLang sinkhorn backend under `backend.attn: tilelang` -- fails at the first log-prob forward with "dsv4 sinkhorn TileLang backend was requested, but the optional kernel is unavailable". The pin's rationale (FLA kernels, and Automodel's apache-tvm-ffi<=0.1.11 cap against the 0.1.12 SIGABRT) is not architecture-specific. Extend both markers to aarch64, matching the base dependency and the vllm extra, and regenerate uv.lock (uv 0.11.28 on aarch64): only the tilelang entries for platform_machine != 'x86_64' and the markers that depended on the same expression change. Verified on 16 x 4 GB200: with tilelang 0.1.11 in the Automodel venv the V4.1 DAPO recipe trains (31 steps, validations at steps 10/20/30). Signed-off-by: linmuchuiyang <6169320+linmuchuiyang@users.noreply.github.com>
svcnvidia-nemo-ci
added
the
waiting-on-maintainers
|
/ok to test cbd35bf |
sharonyu-115
left a comment
sharonyu-115
left a comment
There was a problem hiding this comment.
Thanks for contributing the ARM64 extension to the Automodel and SGLang TileLang 0.1.11 pins, addressing the Automodel import incompatibility.
(C1) Commit the formatter output to resolve the failing lint check. This is a formatting failure introduced by the change. GPU tests were skipped.
Generated by Codex
| "flashinfer-jit-cache==0.6.14", | ||
| "transformers==5.12.1", | ||
| "tilelang==0.1.11; sys_platform == 'linux' and platform_machine == 'x86_64'", | ||
| "tilelang==0.1.11; sys_platform == 'linux' and (platform_machine == 'x86_64' or platform_machine == 'aarch64')", |
There was a problem hiding this comment.
C1 · 1 action item.
TL;DR — P3: commit the Taplo formatting changes to unblock CI.
The longer platform marker changes Taplo's alignment of the inline comments on lines 192–193. The lint job fails because the formatter modifies those lines.
AI-1
Action: Run uv run --group dev pre-commit run taplo-format --files pyproject.toml and commit the resulting changes. The formatter edits lines 192–193 outside this PR's diff, so an in-place suggestion cannot cover the fix.
svcnvidia-nemo-ci
removed
the
waiting-on-maintainers
What does this PR do?
Closes #4754.
Extends the
platform_machine == 'x86_64'marker on the twotilelang==0.1.11pins (automodelandsglangextras)to
aarch64, and regeneratesuv.lockaccordingly.Why
On aarch64 the Automodel worker venv currently resolves tilelang 0.1.12 (base range
>=0.1.9,<0.1.13; the pin onlyapplies on x86).
tile-kernels==1.0.0, pinned bynemo-automodel[moe], cannot be imported on 0.1.12(
tilelang.utils.target.determine_targetmoved totilelang.backend.target; deepseek-ai/TileKernels#24). Automodelswallows the import error,
_HAS_TILE_KERNELS_SINKHORNbecomesFalse, and DeepSeek-V4.1 — which forcessinkhorn_backend="tilelang"wheneverbackend.attn == "tilelang"— fails at the first log-prob forward:The pin's rationale is not x86-specific (Automodel's FLA kernels, and Automodel
main'sapache-tvm-ffi<=0.1.11capagainst the 0.1.12 SIGABRT). The
vllmextra already pins its tilelang on both architectures.Validation
deepseek-v4.1-supportrecipe(
dapo-deepseek-v4.1-12n4g-fsdp2tp1-ep48-colocated-stream-adam.yaml): with tilelang 0.1.11 in the Automodel venv therecipe trains 31 steps with validations at steps 10/20/30 (AIME-2024 accuracy 0.479 → 0.496 → 0.579); with 0.1.12 it
fails at step 1 as above.
uv run --extra automodel python -c "import tile_kernels"succeeds on aarch64 after the change.Lock
uv.lockregenerated withuv lock(uv 0.11.28, on aarch64, againstmain@540d2568, 568 packages resolved).Changes are confined to the consequences of the two markers: the
tilelangentries of thenemo-rldependencies(
platform_machine != 'x86_64'with theautomodel/sglangextras: 0.1.12 → 0.1.11), thetile-kernelsandnemo-automodel[moe]requirement markers, and the now-simplified markers of six packages that were conditioned on thesame expression (
apache-tvm-ffi,cloudpickle,ml-dtypes,numpy,psutil,setuptools, plus thetilelang 0.1.11/0.1.12resolution-marker blocks). No package versions other than tilelang change.