Repository navigation
[Refactor] Decouple torch/triton backend from base deps - #937
Conversation
setup.py picks the torch+triton flavor from FLA_TARGET_DEVICE or the installed torch, so pip install -e . on a rocm/xpu box no longer drags in CUDA wheels. extras [cuda]/[rocm]/[xpu]/[cpu] let you pick explicitly. also adds an amd-mi300 workflow (rocm/pytorch docker image, no fresh venv) and pins FLA_TARGET_DEVICE in setup-ci-env.
no fla-owned amd runner yet — cron would queue forever.
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces dynamic dependency resolution for different hardware backends (CUDA, ROCm, XPU, and CPU) during installation. It updates the CI configuration, documentation, pyproject.toml, and setup.py to support target device auto-detection and explicit backend selection via the FLA_TARGET_DEVICE environment variable or package extras. Feedback highlights a critical packaging limitation where dynamic dependencies do not behave as expected for pre-built wheels, as well as a robustness improvement to catch all exceptions when importing torch in setup.py to prevent installation crashes.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 82e9a725f2
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
kashif
changed the title
dynamic install_requires got baked into wheel METADATA at build time, so pip install fla[rocm] still pulled the cuda triton. drop torch/triton from base deps entirely — backend lives only in the extras now, which wheels handle correctly. also drops FLA_TARGET_DEVICE since it can't see the user's torch through pep517 build isolation anyway.
|
@codex review |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 9151b7b061
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
triton is imported at fla module load so [cpu] without it crashes. setuptools>=64 for PEP 621 [project] support.
|
@yzhangcs this is ready for a review if you can kindly have a look |
|
hey, thanks for tackling this — backend dependency coupling has been a pain for rocm/xpu users for a while. a couple of thoughts before merging: first, i'd lean against adding the also the installation section in README got pretty chunky. what do you think about moving the full backend table + per-platform commands into a separate on the correctness side — lastly, the behavior change is worth flagging: with torch/triton moved out of base deps, a plain happy to help iterate on any of the above if useful! |
|
thanks let me fix it up! |
keeps the install pattern documented; flip the guard when a runner is wired in.
setup.py no longer holds deps. backend extras ride with fla-core, the rest with flash-linear-attention.
|
hey @zhiyuan1i thanks for the careful review, all four points addressed:
pr description updated with the behavior change note. |
|
fixed @zhiyuan1i |
|
note as an aside running the amd CI manually has already surfaced issues in triton which has now been fixed in main: triton-lang/triton#10381 so I hope this will be useful moving forward |
|
@codex review |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 696b82e905
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
flash-linear-attention[cuda] now forwards to fla-core[cuda] so the released wheel actually pulls torch+triton. install docs split non-cuda into 'install backend torch from pytorch index, then install fla' so pip doesn't fall back to PyPI for torch.
|
@codex review |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: e639f65474
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
pytorch index uses different triton names on stable (pytorch-triton-rocm) vs nightly (triton-rocm), so pinning either brittles one path. drop the pin and follow the amd-recommended pattern of installing torch from --index-url, which transitively pulls the right triton.
|
@codex review |
|
Codex Review: Didn't find any major issues. 👍 ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
If Codex has suggestions, it will comment; otherwise it will react with 👍. Codex can also answer questions or update the PR. Try commenting "@codex address that feedback". |
Optional layers/models imports now skip missing extension or runtime deps without hiding fla internal import errors. Export the public API only when both extension modules are available, and extend split-package/smoke tests to cover fla-core-only imports from the repo root plus backend-extra metadata.
|
@codex review |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: cef22c0711
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
* decouple torch/triton backend from base deps setup.py picks the torch+triton flavor from FLA_TARGET_DEVICE or the installed torch, so pip install -e . on a rocm/xpu box no longer drags in CUDA wheels. extras [cuda]/[rocm]/[xpu]/[cpu] let you pick explicitly. also adds an amd-mi300 workflow (rocm/pytorch docker image, no fresh venv) and pins FLA_TARGET_DEVICE in setup-ci-env. * amd-mi300: drop cron, keep workflow_dispatch only no fla-owned amd runner yet — cron would queue forever. * address bot review: extras-only, no dynamic deps dynamic install_requires got baked into wheel METADATA at build time, so pip install fla[rocm] still pulled the cuda triton. drop torch/triton from base deps entirely — backend lives only in the extras now, which wheels handle correctly. also drops FLA_TARGET_DEVICE since it can't see the user's torch through pep517 build isolation anyway. * cpu extra needs triton, bump setuptools floor triton is imported at fla module load so [cpu] without it crashes. setuptools>=64 for PEP 621 [project] support. * no em-dashes * fix cpu triton flavor in table * drop amd ci until runner exists * amd ci: restore behind if:false guard keeps the install pattern documented; flip the guard when a runner is wired in. * add npu extra * build_packages: read pyproject.toml setup.py no longer holds deps. backend extras ride with fla-core, the rest with flash-linear-attention. * slim install section, move details to INSTALL.md * amd ci: drop if:always so a failed verify stops the run * npu: pin to torch 2.6.0 + triton-ascend 3.2.0 * keep backend extras on extension wheel, two-step non-cuda install flash-linear-attention[cuda] now forwards to fla-core[cuda] so the released wheel actually pulls torch+triton. install docs split non-cuda into 'install backend torch from pytorch index, then install fla' so pip doesn't fall back to PyPI for torch. * contributing: install with a backend extra * rocm/xpu: let torch pull its own triton flavor pytorch index uses different triton names on stable (pytorch-triton-rocm) vs nightly (triton-rocm), so pinning either brittles one path. drop the pin and follow the amd-recommended pattern of installing torch from --index-url, which transitively pulls the right triton. * install: match intro to extras shape * [Fix] Keep fla-core importable without extension deps Optional layers/models imports now skip missing extension or runtime deps without hiding fla internal import errors. Export the public API only when both extension modules are available, and extend split-package/smoke tests to cover fla-core-only imports from the repo root plus backend-extra metadata. --------- Co-authored-by: zhiyuan1i <uniartisan2017@gmail.com>
On rocm/xpu/npu boxes
pip install -e .happily pulls CUDA wheels becausetorch>=2.7.0+triton>=3.3get re-resolved against PyPI's default index. This makes the backend explicit.pyproject.tomlnow keepstorchand the matchingtritonflavor in backend extras ([cuda]/[rocm]/[xpu]/[npu]/[cpu]) instead of base deps. The base wheel METADATA carries onlytransformersandeinops, sopip install flash-linear-attention[rocm] --extra-index-url https://download.pytorch.org/whl/rocm7.2resolves to torch (rocm) + pytorch-triton-rocm with no CUDA carryover, and the same shape works on prebuilt wheels (which addresses the earlier bot reviews about wheel metadata being baked at build time).scripts/build_packages.pynow reads deps and extras frompyproject.tomlvia tomllib. The split-package layout keeps backend extras withfla-core(since the kernels need torch + triton) and ergonomics extras (tilelang/conv1d/benchmark/test) withflash-linear-attention.Install docs: short snippet in README, full backend table + NPU (Ascend) flow + source installs +
--no-depspath in a newINSTALL.md.CI: nvidia and intel paths unchanged in behavior. An
amd-mi300-ciworkflow is included as documentation behind a job-levelif: falseguard, ready to flip once fla has a runner.Behavior change (worth flagging): bare
pip install flash-linear-attentionno longer pullstorchortriton. Users must pick a backend extra (e.g.[cuda]). This is the trade-off for not silently overlaying CUDA wheels on ROCm / XPU / NPU machines. Existing CI install lines that go through--index-urlfor torch first and thenpip install .still work because the base no longer churns those packages.Tested on a Strix Halo box (ROCm 7.2 + torch nightly): auto-detect picks the right flavor, egg-info comes out clean, no
nvidia-*collection.