multimodalart HF Staff commited on
Commit
64cffdd
·
verified ·
1 Parent(s): 9d3c7ff

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