fla: fix relative import depths in the vendored source - #1129
fla: fix relative import depths in the vendored source#1129qgallouedec wants to merge 2 commits into
Conversation
Ten imports in `fla/torch-ext/fla` use a dot count that is one or two levels too shallow. In an `__init__.py` the current package is the directory itself, one level deeper than for a plain module, so `from ...ops.common...` inside `ops/common/backends/tilelang/__init__.py` resolves against `fla.ops.common` and the original path is appended to it, giving `fla.ops.common.ops.common...`. The imports sit inside functions, so nothing fails at import time. The gated_delta_rule backward hits three of them and raises ModuleNotFoundError, but only when tilelang is installed, otherwise it falls through to triton.
|
/kernel-bot build-and-stage fla |
|
Build request processed. Command: Dispatched (1): Hub uploads: |
vasqu
left a comment
vasqu
left a comment
There was a problem hiding this comment.
Thanks for the fix 🫡
Not fixed: ops/deltaformer/parallel.py:26, from ...utilslayers.utils import pad_input, unpad_input. No dot count works, layers.py here is a slim shim and those two functions aren't vendored. Either pull in layers/utils.py or drop deltaformer, your call.
Imo, we could just copy those functions into the file. I'd like to avoid dropping it for now
|
Oh, there was already an existing PR that just got merged: #1074 Let me try to resolve the conflicts, because I think this PR has more fixes. |
|
/kernel-bot build-and-stage fla |
|
Build request processed. Command: Dispatched (1): Hub uploads: |
4 participants
Fixes #1128.
Sets the correct dot count on the 10 imports that have one. I found them by resolving every relative import in
fla/torch-ext/flaagainst the file layout, and each of these has exactly one depth that resolves.Verified on the published build (
aa89c60) with the repro from the issue: thechunk_gated_delta_rulebackward runs, and the gradients match this repo's ownnaive_chunk_gated_delta_ruleto 1.9e-3 relative in fp32. That exercisesops/common/backends/tilelang/__init__.py:99only, the other nine I can't reach from here (NPU-only, multi-card CP,TYPE_CHECKING-only).chunk_kdafwd+bwd still passes.Not fixed:
ops/deltaformer/parallel.py:26,from ...utilslayers.utils import pad_input, unpad_input. No dot count works,layers.pyhere is a slim shim and those two functions aren't vendored. Either pull inlayers/utils.pyor drop deltaformer, your call.