Hugging Face's logo

respanai
/
span-02

Image-Text-to-Text
Safetensors
gemma4
decision-index
typed-outputs
nvfp4
fp8
quantized
Mixture of Experts
custom-engine
conversational
8-bit precision
Model card Files
xet
 Community
span-02 / engine /crates /kernels /src /vision.rs
lihe07's picture
lihe07
span-02: single-source release folder (engine final path, soft blend)
19886eb verified
Raw History Blame Contribute Delete
6.02 kB
//! Vision tower kernel launchers (cuda/vision_*.cu). All pointers are u64 device addresses; bf16 unless stated.
//! GEMM: C[M,N] = A[M,K] . B[N,K]^T (B = nn.Linear weight layout). Tile choice by problem size (see `vis_gemm`).
use crate::backend::Arg::{F, I, L, P};
use crate::ops::Gpu;
use anyhow::Result;
fn cdiv(a: usize, b: usize) -> u32 {
((a + b - 1) / b) as u32
}
/// GEMM tile configs (BM x BN, warps); `sym` = kernel symbol in cuda/vision_gemm.cu
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum VisTile {
T128x128,
T64x128,
T64x64,
}
impl VisTile {
fn dims(self) -> (usize, usize, u32) {
match self {
VisTile::T128x128 => (128, 128, 256),
VisTile::T64x128 => (64, 128, 256),
VisTile::T64x64 => (64, 64, 128),
}
}
fn sym(self, glu: bool) -> &'static str {
match (self, glu) {
(VisTile::T128x128, true) => "vis_gemm_128x128_glu",
(VisTile::T64x128, true) => "vis_gemm_64x128_glu",
(VisTile::T64x64, true) => "vis_gemm_64x64_glu",
(VisTile::T64x128, false) => "vis_gemm_64x128",
(VisTile::T64x64, false) => "vis_gemm_64x64",
(VisTile::T128x128, false) => unreachable!("128x128 is only used for the gated GEMM"),
}
}
/// measured on the RTX 5090: small tiles win for N=1152 and for short M (more CTAs per SM); 128x128 only for the big gated-GELU GEMM at full size
fn pick(m: usize, n: usize, glu: bool) -> VisTile {
if glu {
if m >= 2000 {
VisTile::T128x128
} else if m >= 600 {
VisTile::T64x128
} else {
VisTile::T64x64
}
} else if n >= 3000 && m >= 2000 {
VisTile::T64x128
} else {
VisTile::T64x64
}
}
}
impl Gpu {
/// c[m,n] = a[m,k] . b[n,k]^T (k % 8 == 0, n even)
pub fn vis_gemm(&self, c: u64, a: u64, b: u64, m: usize, n: usize, k: usize) -> Result<()> {
assert!(k % 8 == 0 && n % 2 == 0, "vis_gemm: k%8 / n%2");
let t = VisTile::pick(m, n, false);
let (bm, bn, thr) = t.dims();
self.be.launch("vision_gemm", t.sym(false), [cdiv(n, bn), cdiv(m, bm), 1], [thr, 1, 1], &[P(c), P(a), P(b), I(m as i32), I(n as i32), I(k as i32)])
}
/// gated GELU GEMM: `b` rows are interleaved in blocks of 8 (8 gate rows, 8 up rows, ...; `n2` = 2*I rows in total); c[m, n2/2] = bf16(gelu_tanh(bf16 g) * bf16 u)
pub fn vis_gemm_glu(&self, c: u64, a: u64, b: u64, m: usize, n2: usize, k: usize) -> Result<()> {
assert!(k % 8 == 0 && n2 % 16 == 0, "vis_gemm_glu: k%8 / n2%16");
let t = VisTile::pick(m, n2, true);
let (bm, bn, thr) = t.dims();
self.be.launch("vision_gemm", t.sym(true), [cdiv(n2, bn), cdiv(m, bm), 1], [thr, 1, 1], &[P(c), P(a), P(b), I(m as i32), I(n2 as i32), I(k as i32)])
}
/// non-causal flash attention, head_dim 72 (padded 80). q,k,v [h][m][80]; o [m][h*72]
pub fn vis_attn(&self, o: u64, q: u64, k: u64, v: u64, m: usize, h: usize) -> Result<()> {
self.be.launch("vision_attn", "vis_attn_fwd", [cdiv(m, 64), h as u32, 1], [128, 1, 1], &[P(o), P(q), P(k), P(v), I(m as i32), I(h as i32)])
}
/// resid_out = resid + rmsnorm(x; w_post); normed = rmsnorm(resid_out; w_pre) (w_pre = 0 skips it); hidden must be 1152; resid_out may alias resid
#[allow(clippy::too_many_arguments)]
pub fn vis_sandwich(&self, resid_out: u64, normed: u64, resid: u64, x: u64, w_post: u64, w_pre: u64, rows: usize, eps: f32) -> Result<()> {
self.be.launch("vision_misc", "vis_sandwich_1152", [cdiv(rows, 8), 1, 1], [256, 1, 1], &[P(resid_out), P(normed), P(resid), P(x), P(w_post), P(w_pre), I(rows as i32), F(eps)])
}
/// 128-query-row variant of `vis_attn` (8 warps)
pub fn vis_attn128(&self, o: u64, q: u64, k: u64, v: u64, m: usize, h: usize) -> Result<()> {
self.be.launch("vision_attn", "vis_attn_fwd128", [cdiv(m, 128), h as u32, 1], [256, 1, 1], &[P(o), P(q), P(k), P(v), I(m as i32), I(h as i32)])
}
pub fn vis_pixels(&self, out: u64, patches_u8: u64, n: usize) -> Result<()> {
self.be.launch("vision_misc", "vis_pixels_bf16", [cdiv(n, 256).min(4096), 1, 1], [256, 1, 1], &[P(out), P(patches_u8), L(n as i64)])
}
pub fn vis_pos_add(&self, h: u64, table: u64, rows: usize, gw: usize, pos_size: usize, hidden: usize) -> Result<()> {
self.be.launch("vision_misc", "vis_pos_add", [rows as u32, 1, 1], [256, 1, 1], &[P(h), P(table), I(gw as i32), I(pos_size as i32), I(hidden as i32)])
}
#[allow(clippy::too_many_arguments)]
pub fn vis_qkv_prep(&self, q: u64, k: u64, v: u64, qkv: u64, q_w: u64, k_w: u64, cos_t: u64, sin_t: u64, m: usize, heads: usize, gw: usize, eps: f32) -> Result<()> {
self.be.launch("vision_misc", "vis_qkv_prep", [m as u32, 1, 1], [256, 1, 1], &[P(q), P(k), P(v), P(qkv), P(q_w), P(k_w), P(cos_t), P(sin_t), I(m as i32), I(heads as i32), I(gw as i32), F(eps)])
}
/// 3x3 pool + sqrt(H) + standardize (-> `pooled`, may be 0) + embed_vision RMSNorm (-> `normed`); out tokens = gh*gw/9
#[allow(clippy::too_many_arguments)]
pub fn vis_pool_std_norm(&self, pooled: u64, normed: u64, h: u64, std_bias: u64, std_scale: u64, n_out: usize, gw: usize, hidden: usize, root: f32, eps: f32) -> Result<()> {
assert!(hidden <= 1280);
self.be.launch("vision_misc", "vis_pool_std_norm", [n_out as u32, 1, 1], [256, 1, 1], &[P(pooled), P(normed), P(h), P(std_bias), P(std_scale), I(gw as i32), I(hidden as i32), F(root), F(eps)])
}
/// x[row_idx[i]] = src[src_idx[i]] for i < n (soft-token injection, no scale)
pub fn inject_soft(&self, x: u64, src: u64, row_idx: u64, src_idx: u64, n: usize, h: usize) -> Result<()> {
if n == 0 {
return Ok(());
}
self.be.launch("vision_inject", "inject_soft_bf16", [n as u32, 1, 1], [256, 1, 1], &[P(x), P(src), P(row_idx), P(src_idx), I(h as i32)])
}
}