Download src/attention.rs from Snapkitty/macrogrok-harness: direct link, hf CLI and curl.
- Browser
- Download file 912 Bytes
-
https://huggingface.co/Snapkitty/macrogrok-harness/resolve/main/src/attention.rs
- Command line
-
hf download hf://Snapkitty/macrogrok-harness/src/attention.rs
-
curl -L -o attention.rs https://huggingface.co/Snapkitty/macrogrok-harness/resolve/main/src/attention.rs
912 Bytes
| use crate::memory::AttentionMemory; | |
| use cudarc::driver::{CudaDevice, CudaSlice}; | |
| use half::f16; | |
| use std::sync::Arc; | |
| pub struct TensorCoreAttention { device: Arc<CudaDevice> } | |
| impl TensorCoreAttention { | |
| pub fn new(device: Arc<CudaDevice>) -> anyhow::Result<Self> { Ok(Self { device }) } | |
| /// X-algorithm attention: hybrid/sparse + Tensor Core dense tiles + explicit memory | |
| pub fn forward(&self, q: &CudaSlice<f16>, _mem: &AttentionMemory, _layer: usize, _causal: bool, _scale: f32) -> anyhow::Result<CudaSlice<f16>> { | |
| // Real impl: load K/V tiles from AttentionMemory, launch WMMA kernel (attention_tc.cu) | |
| // with head_dim tiling, online softmax (FlashAttention), and X-algorithm hook. | |
| let out_elems = 0; // compute from q dims | |
| let out = unsafe { self.device.alloc::<f16>(out_elems)? }; | |
| let _ = q; // suppress unused | |
| Ok(out) | |
| } | |
| } | |