refactor: replace magic layout ints with TensorLayout enum

- Add TensorLayout enum (C++ + Python) to replace magic layout ints
- Add C10_CUDA_CHECK post-launch error checking to all kernel entries
- Add CUDAGuard + freqs_cis shape validation to rotary_emb.cu
- Cache SM count to eliminate per-call cudaDeviceGetAttribute
- Add DISPATCH_CAUSAL_MASK macro to deduplicate dispatcher if/else
- Convert mask type hints from X|None to Optional[X]

astrai/extension/__init__.py

@@ -25,6 +25,7 @@ from astrai.extension.attention_backend import (
     get_backend,
 )
 from astrai.extension.attention_ops import (
+    TensorLayout,
     attn_decode,
     attn_paged_decode,
     attn_prefill,
@@ -37,6 +38,7 @@ __all__ = [
     "AttentionBackend",
     "CudaBackend",
     "TorchNativeBackend",
+    "TensorLayout",
     "attention",
     "attn_backend",
     "get_backend",