feat: add fp8 training via cublasLt dispatch

- fp8_mm kernel (csrc): cublasLt fp8 e4m3 gemm, TN layout mapped zero-copy
- custom::fp8_mm custom op: meta/cuda/cpu kernels + scale-corrected bf16 autograd
- aten::linear and linear_backward dispatch on CUDA key, zero model changes
- per-tensor scale or raw cast; single-GPU smoke loss matches bf16

csrc/CMakeLists.txt

@@ -48,7 +48,7 @@ set(TORCH_LIBS
 
 set(CMAKE_CUDA_ARCHITECTURES "${ASTRAI_CUDA_ARCH}")
 
-set(KERNELS attn_decode attn_prefill attn_paged_decode attn_paged_prefill rotary_emb)
+set(KERNELS attn_decode attn_prefill attn_paged_decode attn_paged_prefill rotary_emb fp8_mm)
 
 foreach(name ${KERNELS})
     add_library(${name} MODULE "${CMAKE_CURRENT_SOURCE_DIR}/kernels/${name}.cu")
@@ -61,6 +61,9 @@ foreach(name ${KERNELS})
         "${PYTHON_INCLUDE_DIR}")
 
     target_link_libraries(${name} PRIVATE ${TORCH_LIBS})
+    if(${name} STREQUAL "fp8_mm")
+        target_link_libraries(${name} PRIVATE CUDA::cublasLt)
+    endif()
     target_link_options(${name} PRIVATE "-Wl,-rpath,${TORCH_LIB_DIR}")
 
     target_compile_options(${name} PRIVATE