| use crate::attention::TensorCoreAttention;
|
| use crate::memory::AttentionMemory;
|
| use cudarc::driver::{CudaDevice, CudaSlice};
|
| use half::f16;
|
| use std::sync::Arc;
|
|
|
| pub struct TransformerConfig { pub hidden_size: usize, pub num_layers: usize, pub num_heads: usize, pub head_dim: usize, pub max_seq_len: usize, pub vocab_size: usize }
|
| pub struct TransformerBlock { attn: TensorCoreAttention }
|
| impl TransformerBlock {
|
| pub fn forward(&self, x: &CudaSlice<f16>, mem: &mut AttentionMemory, layer_idx: usize) -> anyhow::Result<CudaSlice<f16>> {
|
|
|
| let _ = (x, mem, layer_idx, &self.attn);
|
| Err(anyhow::anyhow!("wire QKV + MLP Tensor Core GEMMs via CUTLASS or hand-written mma.sync"))
|
| }
|
| }
|
| pub struct Transformer { pub config: TransformerConfig, pub blocks: Vec<TransformerBlock>, pub embed: CudaSlice<f16>, pub lm_head: CudaSlice<f16>, pub device: Arc<CudaDevice> }
|
| impl Transformer {
|
| pub fn forward(&self, _input_ids: &[u32], _mem: &mut AttentionMemory) -> anyhow::Result<CudaSlice<f16>> {
|
| Err(anyhow::anyhow!("embed -> blocks -> norm -> lm_head"))
|
| }
|
| }
|
|
|
