refactor: standardize packed 3d inference

- keep training attention on dense 4d tensors
- use packed 3d tensors with KV cache for inference
- extend CUDA rotary embedding to packed 3d inputs
- adapt torch, CUDA and FlashAttention backend dispatch

astrai/model/components/decoder_block.py

@@ -54,6 +54,7 @@ class DecoderBlock(nn.Module):
         attention_mask: Optional[Tensor] = None,
         kv_cache: Optional[KVCache] = None,
         is_causal: bool = False,
+        fwd: Optional[str] = None,
     ) -> DecoderOutput:
         attn_output = self.attention(
             self.input_norm(x),
@@ -61,6 +62,7 @@ class DecoderBlock(nn.Module):
             attention_mask,
             kv_cache,
             is_causal,
+            fwd,
         )
         x = attn_output + x
         normalized = self.post_attention_norm(x)