refactor: separate extension ops and backends

astrai/model/components/__init__.py

@@ -1,4 +1,4 @@
-from astrai.extension.rotary_backend import apply_rotary_emb
+from astrai.extension.backend.rotary import apply_rotary_emb
 from astrai.model.components.attention import GQA, MLA
 from astrai.model.components.decoder_block import DecoderBlock
 from astrai.model.components.embedding import Embedding

astrai/model/components/attention.py

@@ -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