Spaces:
Runtime error
Runtime error
Remove per-call torch.compile JIT from gated_deltanet helpers (ZeroGPU warmup)
Browse files
kairos/third_party/fla/layers/gated_deltanet.py
CHANGED
|
@@ -21,12 +21,16 @@ if TYPE_CHECKING:
|
|
| 21 |
from fla.models.utils import Cache
|
| 22 |
|
| 23 |
|
| 24 |
-
@torch.compile
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 25 |
def elu_p1(x):
|
| 26 |
return (F.elu(x, 1., False) + 1.).to(x)
|
| 27 |
|
| 28 |
|
| 29 |
-
@torch.compile
|
| 30 |
def sum_norm(x):
|
| 31 |
return (x / x.sum(-1, keepdim=True)).to(x)
|
| 32 |
|
|
|
|
| 21 |
from fla.models.utils import Cache
|
| 22 |
|
| 23 |
|
| 24 |
+
# NOTE: These are intentionally left as eager functions (no @torch.compile).
|
| 25 |
+
# On ZeroGPU each @spaces.GPU call runs in a fresh forked worker where
|
| 26 |
+
# torch.compile (JIT) is not supported / would re-compile on every call,
|
| 27 |
+
# re-triggering the per-call warmup tax. They are tiny elementwise ops, so
|
| 28 |
+
# eager execution costs nothing at steady state while avoiding the first-use
|
| 29 |
+
# compilation on every fresh worker process.
|
| 30 |
def elu_p1(x):
|
| 31 |
return (F.elu(x, 1., False) + 1.).to(x)
|
| 32 |
|
| 33 |
|
|
|
|
| 34 |
def sum_norm(x):
|
| 35 |
return (x / x.sum(-1, keepdim=True)).to(x)
|
| 36 |
|