refactor: separate extension ops and backends

2026-08-16 21:15:52 +08:00
parent 3406157431
commit 6ac3b51496
18 changed files with 82 additions and 54 deletions
@@ -5,8 +5,7 @@ import torch.nn as nn
 import torch.nn.functional as F
 from torch import Tensor

-from astrai.extension import attention
-from astrai.extension.rotary_backend import apply_rotary_emb
+from astrai.extension.backend import apply_rotary_emb, attention
 from astrai.factory import BaseFactory
 from astrai.inference.cache import KVCache
 from astrai.model.components.linear import Linear