h34v3nzc0dex tested PR 5434 on Strix Halo (gfx1151, ROCm 7.13,
torch 2.11.0+rocm7.13.0) and hit a hard regression:
File ".../fla/ops/common/backends/tilelang/__init__.py", line 92,
in chunk_bwd_dqkwg
File ".../tilelang/jit/kernel.py", line 137, in __init__
File ".../tilelang/tileop/gemm/__init__.py", line 143,
in _select_gemm_instruction
tvm.error.InternalError: Check failed: (0) is false:
Unsupported target for gemm:
hip -keys=hip,gpu -mcpu=gfx1151 ...
`tilelang==0.1.8` ships no HIP GEMM instruction; `_select_gemm_instruction`
raises at lower-time, not import-time. So:
- pip install succeeds
- `import tilelang` succeeds
- `TileLangBackend.is_available()` returns True
- FLA's dispatcher picks TileLang for `chunk_bwd_dqkwg`
- training subprocess dies at first GDN backward, no graceful fallback
The PR's existing platform gate (`_tilelang_platform_supported`)
checked only `sys.platform == "linux"` and `platform.machine()`, both
of which look identical on a ROCm box.
Fix has two layers:
1. INSTALL GATE: new `_torch_has_hip()` helper checks
`torch.version.hip is not None`. `_tilelang_platform_supported`
now returns False on HIP torch, so the install never fires.
2. RUNTIME GATE: even with the install skipped, a user could have
tilelang already present (e.g. venv carried over from a CUDA box).
`_install_fast_path_hooks` now calls
`os.environ.setdefault("FLA_TILELANG", "0")` when HIP is detected,
which is the env-var FLA's `TileLangBackend` already honors. Users
who know they have a HIP-aware tilelang fork can override by
setting `FLA_TILELANG=1` explicitly.
This costs nothing on CUDA (the gate is a no-op when
`torch.version.hip is None`), and removes the crash for AMD users.
The benchmark numbers in the PR description (1.43x on B200 sm_100)
are not affected.
The other halves of the PR are confirmed working on gfx1151 by the
same report:
- `flash-linear-attention 0.5.0` runs at production scale
(B=1 T=8192 H=16 K=128 V=128 and others) with no patches.
- `causal-conv1d` runs at the shapes the fast-path gate cares
about. (A separate Ubuntu 24.04 `--gcc-install-dir` build
workaround is needed for the source-build path; that mirrors
bbf004c's llama.cpp fix and is out of scope here.)
Tests added:
- test_tilelang_platform_unsupported_on_hip_torch
- test_tilelang_install_skipped_on_hip_torch
- test_install_fast_path_hooks_sets_fla_tilelang_zero_on_hip
- test_install_fast_path_hooks_respects_user_fla_tilelang_override
- test_install_fast_path_hooks_does_not_set_fla_tilelang_on_cuda
Total 50 passing (was 45).
10b50c84df