diff --git a/.devin/config.local.json b/.devin/config.local.json new file mode 100644 index 000000000..3054c4bf5 --- /dev/null +++ b/.devin/config.local.json @@ -0,0 +1,12 @@ +{ + "permissions": { + "allow": [ + "Exec(git status)", + "Exec(gh pr)", + "Exec(gh run)", + "Exec(cargo fmt)", + "Exec(git diff)", + "Exec(cargo clippy)" + ] + } +} \ No newline at end of file diff --git a/fusor-ml/core/src/quantized/matmul/mod.rs b/fusor-ml/core/src/quantized/matmul/mod.rs index 96fd9bb95..44cd6340b 100644 --- a/fusor-ml/core/src/quantized/matmul/mod.rs +++ b/fusor-ml/core/src/quantized/matmul/mod.rs @@ -911,8 +911,9 @@ fn qgemv_cols_per_workgroup_for_direct(format: tile_ir::GgmlQuantFormat, k: u32, if format.is_q6k_family() && n <= 4096 && k > 4096 { return 4; // was Q6KLargeNarrow4 } - // Q8_0 wide. - if format.is_q8_0_family() && k <= 1024 && n >= 8192 { + // Q8_0 wide. Keep this aligned with qgemv_tile_with_epilogue, which + // switches Q8_0/Q8_0Native to 4x8 output columns for wide outputs. + if format.is_q8_0_family() && n >= 8192 { return 32; // was Q8WideAccelerated32 } // FormatAccelerated: Q5_0 mid (K,N both 2048..=4096), Q4K/Q6K general, @@ -929,7 +930,10 @@ fn qgemv_cols_per_workgroup_for_direct(format: tile_ir::GgmlQuantFormat, k: u32, if format.is_q5_0_family() && k <= 1024 && n <= 4096 { return 8; // was Q5Small8 } - 4 // was Default4 + // The shader-side qgemv shape table is the source of truth for generic + // formats. Returning a smaller value here over-dispatches masked columns, + // and every extra workgroup still pays the full K-loop cost. + tile_ir_kernels::qgemv_cols_per_workgroup_for_shape(format, k, n) } fn qmatmul_m_pad_target_for_caps(m: usize, n: usize, caps: KernelDeviceCaps) -> Option { diff --git a/fusor-ml/core/src/quantized/mod.rs b/fusor-ml/core/src/quantized/mod.rs index a63d3f20e..050005840 100644 --- a/fusor-ml/core/src/quantized/mod.rs +++ b/fusor-ml/core/src/quantized/mod.rs @@ -255,21 +255,26 @@ fn qmatrix_storage_layout_for_parts_with_env( // the raw-word ggml qgemv dot (`qgemv_q6k_ggml`) cannot address. The // f32-scale layout is 212 bytes (word-aligned) and feeds that amortized // decode; the +2 bytes/block (~0.9%) is paid back many times over by the - // faster kernel. Q4K/Q5K keep their native f16-scale layout (their native - // blocks are word-aligned, and native f16 scales read less memory). + // faster kernel. if ty == GgmlType::Q6K { return QMatrixStorageLayout::GpuF32Scales; } - if matches!( - ty, - GgmlType::Q4_0 | GgmlType::Q5_0 | GgmlType::Q8_0 | GgmlType::Q4K | GgmlType::Q5K - ) && native_half_scale_storage_enabled(shader_f16_supported, env_override) + // The small 32-element block formats save only two bytes/block with native + // f16 scales, but native Q4_0/Q5_0/Q8_0 blocks are byte-addressed. Decode is + // much faster with word-aligned f32-scale storage on Apple GPUs. + if matches!(ty, GgmlType::Q4_0 | GgmlType::Q5_0 | GgmlType::Q8_0) { + if env_override.is_some() + && native_half_scale_storage_enabled(shader_f16_supported, env_override) + { + QMatrixStorageLayout::Native + } else { + QMatrixStorageLayout::GpuF32Scales + } + } else if matches!(ty, GgmlType::Q4K | GgmlType::Q5K) + && native_half_scale_storage_enabled(shader_f16_supported, env_override) { QMatrixStorageLayout::Native - } else if matches!( - ty, - GgmlType::Q4_0 | GgmlType::Q5_0 | GgmlType::Q8_0 | GgmlType::Q4K | GgmlType::Q5K - ) { + } else if matches!(ty, GgmlType::Q4K | GgmlType::Q5K) { QMatrixStorageLayout::GpuF32Scales } else { QMatrixStorageLayout::Native @@ -536,14 +541,22 @@ mod tests { use super::*; #[test] - fn native_capable_quant_storage_defaults_to_native_when_shader_f16_is_supported() { - for ty in [ - GgmlType::Q4_0, - GgmlType::Q5_0, - GgmlType::Q8_0, - GgmlType::Q4K, - GgmlType::Q5K, - ] { + fn small_block_quant_storage_defaults_to_f32_scales() { + for ty in [GgmlType::Q4_0, GgmlType::Q5_0, GgmlType::Q8_0] { + assert_eq!( + qmatrix_storage_layout_for_parts_with_env(ty, true, None), + QMatrixStorageLayout::GpuF32Scales + ); + assert_eq!( + qmatrix_storage_layout_for_parts_with_env(ty, false, None), + QMatrixStorageLayout::GpuF32Scales + ); + } + } + + #[test] + fn k_quant_storage_defaults_to_native_when_shader_f16_is_supported() { + for ty in [GgmlType::Q4K, GgmlType::Q5K] { assert_eq!( qmatrix_storage_layout_for_parts_with_env(ty, true, None), QMatrixStorageLayout::Native @@ -564,7 +577,15 @@ mod tests { } #[test] - fn q4k_storage_env_can_force_native_or_f32_expanded() { + fn quant_storage_env_can_force_native_or_f32_expanded() { + assert_eq!( + qmatrix_storage_layout_for_parts_with_env(GgmlType::Q4_0, true, Some("1")), + QMatrixStorageLayout::Native + ); + assert_eq!( + qmatrix_storage_layout_for_parts_with_env(GgmlType::Q4_0, true, Some("0")), + QMatrixStorageLayout::GpuF32Scales + ); assert_eq!( qmatrix_storage_layout_for_parts_with_env(GgmlType::Q4K, false, Some("1")), QMatrixStorageLayout::Native diff --git a/fusor-ml/core/src/row_program.rs b/fusor-ml/core/src/row_program.rs index 03fc02fa8..6d8ee5e09 100644 --- a/fusor-ml/core/src/row_program.rs +++ b/fusor-ml/core/src/row_program.rs @@ -46,11 +46,6 @@ use crate::{ const BLOCK: u32 = 256; -/// Below this many rows a long axis is split across workgroups (one tile -/// each) with a combine kernel folding the spans — decode has too few rows -/// to fill the device with one workgroup per row. -const SPLIT_ROWS_TARGET: u32 = 256; - /// Workgroup buckets for dynamic-axis row programs, smallest first. The /// kernel monomorphizes per bucket; the active axis length rides in the /// params input. @@ -606,10 +601,12 @@ fn build_row_program_kernel( let block = workgroup_shape.x(); let lanes_own_axis = operation.dynamic_axis.is_some(); - // Long axes with few rows fan out across workgroups: each split runs - // the online body over one tile, writing its unnormalized accumulator - // and softmax statistics to scratch; a combine kernel folds the spans - // with the online monoid. + // Long reducing outputs fan out across workgroups: each split owns one + // axis tile, writes its unnormalized accumulator and softmax statistics + // to scratch, then a combine kernel folds the spans with the online + // monoid. This avoids the single-workgroup multi-tile loop for attention + // prefill, which is both slower to fill the GPU and less robust at Gemma's + // unscaled vision-attention logits. let output_kind = operation.output_step().clone(); let phase_steps = operation.phase_steps().to_vec(); let free_dim_out = match &output_kind { @@ -619,11 +616,7 @@ fn build_row_program_kernel( let tiles = k.div_ceil(block); let splits: u32 = match free_dim_out { Some(free) - if lanes_own_axis - && rows < SPLIT_ROWS_TARGET - && tiles > 1 - && tiles <= block - && (free as u32 + 2) <= block => + if lanes_own_axis && tiles > 1 && tiles <= block && (free as u32 + 2) <= block => { tiles } diff --git a/fusor-ml/core/tests/attention.rs b/fusor-ml/core/tests/attention.rs index 2aee33924..fd114d2a3 100644 --- a/fusor-ml/core/tests/attention.rs +++ b/fusor-ml/core/tests/attention.rs @@ -9,6 +9,15 @@ fn values(len: usize, scale: f32) -> Vec { (0..len).map(|i| ((i as f32) * scale).sin()).collect() } +fn high_variance_values(len: usize, scale: f32) -> Vec { + (0..len) + .map(|i| { + let x = i as f32; + ((x * 0.173).sin() * 1.7 + (x * 0.071).cos() * 0.9) * scale + }) + .collect() +} + struct AttentionCase { batch: usize, heads: usize, @@ -81,6 +90,24 @@ fn cpu_attention( } fn check_attention(case: AttentionCase, tolerance: f32) { + check_attention_with_scale(case, tolerance, None); +} + +fn check_attention_with_scale(case: AttentionCase, tolerance: f32, scale_override: Option) { + check_attention_impl(case, tolerance, scale_override, false, false); +} + +fn check_attention_full_high_variance(case: AttentionCase, tolerance: f32, scale_override: f32) { + check_attention_impl(case, tolerance, Some(scale_override), true, true); +} + +fn check_attention_impl( + case: AttentionCase, + tolerance: f32, + scale_override: Option, + full_compare: bool, + high_variance: bool, +) { pollster::block_on(async { let Ok(device) = Device::new().await else { return; @@ -95,15 +122,27 @@ fn check_attention(case: AttentionCase, tolerance: f32) { causal, masked, } = case; - let q_data = values(batch * heads * q_len * head_dim, 0.13); - let k_data = values(batch * kv_heads * kv_len * head_dim, 0.07); - let v_data = values(batch * kv_heads * kv_len * head_dim, 0.11); + let q_data = if high_variance { + high_variance_values(batch * heads * q_len * head_dim, 1.0) + } else { + values(batch * heads * q_len * head_dim, 0.13) + }; + let k_data = if high_variance { + high_variance_values(batch * kv_heads * kv_len * head_dim, 0.8) + } else { + values(batch * kv_heads * kv_len * head_dim, 0.07) + }; + let v_data = if high_variance { + high_variance_values(batch * kv_heads * kv_len * head_dim, 1.1) + } else { + values(batch * kv_heads * kv_len * head_dim, 0.11) + }; let mask_data = masked.then(|| { (0..q_len * kv_len) .map(|i| if i % 7 == 0 { -1.5 } else { 0.25 }) .collect::>() }); - let scale = 1.0 / (head_dim as f32).sqrt(); + let scale = scale_override.unwrap_or_else(|| 1.0 / (head_dim as f32).sqrt()); let q = Tensor::from_slice(&device, [batch, heads, q_len, head_dim], &q_data); let k = Tensor::from_slice(&device, [batch, kv_heads, kv_len, head_dim], &k_data); @@ -132,16 +171,44 @@ fn check_attention(case: AttentionCase, tolerance: f32) { mask_data.as_deref(), scale, ); - for b in 0..batch { - for h in [0, heads - 1] { - for qi in [0, q_len / 2, q_len - 1] { - for d in [0, head_dim / 2, head_dim - 1] { - let want = expected[((b * heads + h) * q_len + qi) * head_dim + d]; - let got = actual[[b, h, qi, d]]; - assert!( - (got - want).abs() < tolerance, - "b={b} h={h} q={qi} d={d}: got {got}, expected {want}" - ); + if full_compare { + let mut worst = 0.0f32; + let mut worst_index = (0, 0, 0, 0); + for b in 0..batch { + for h in 0..heads { + for qi in 0..q_len { + for d in 0..head_dim { + let want = expected[((b * heads + h) * q_len + qi) * head_dim + d]; + let got = actual[[b, h, qi, d]]; + let err = (got - want).abs(); + if err > worst { + worst = err; + worst_index = (b, h, qi, d); + } + } + } + } + } + assert!( + worst < tolerance, + "worst attention diff {worst} at {worst_index:?}: got {}, expected {}, tolerance {tolerance}", + actual[[worst_index.0, worst_index.1, worst_index.2, worst_index.3]], + expected[((worst_index.0 * heads + worst_index.1) * q_len + worst_index.2) + * head_dim + + worst_index.3] + ); + } else { + for b in 0..batch { + for h in [0, heads - 1] { + for qi in [0, q_len / 2, q_len - 1] { + for d in [0, head_dim / 2, head_dim - 1] { + let want = expected[((b * heads + h) * q_len + qi) * head_dim + d]; + let got = actual[[b, h, qi, d]]; + assert!( + (got - want).abs() < tolerance, + "b={b} h={h} q={qi} d={d}: got {got}, expected {want}" + ); + } } } } @@ -221,6 +288,80 @@ fn attention_causal_prefill_streams_long_kv() { ); } +#[test] +fn attention_unmasked_prefill_streams_long_kv() { + // Gemma 4 vision uses full, non-causal image self-attention over thousands + // of patch tokens. This exercises the streaming row-program path without + // the causal axis bound. + check_attention( + AttentionCase { + batch: 1, + heads: 4, + kv_heads: 4, + q_len: 1024, + kv_len: 1024, + head_dim: 64, + causal: false, + masked: false, + }, + 1e-4, + ); +} + +#[test] +fn attention_unmasked_prefill_streams_long_kv_scale_one() { + check_attention_full_high_variance( + AttentionCase { + batch: 1, + heads: 4, + kv_heads: 4, + q_len: 1024, + kv_len: 1024, + head_dim: 64, + causal: false, + masked: false, + }, + 3e-4, + 1.0, + ); +} + +#[test] +fn attention_unmasked_prefill_many_tiles_scale_one() { + check_attention_full_high_variance( + AttentionCase { + batch: 1, + heads: 2, + kv_heads: 2, + q_len: 384, + kv_len: 2304, + head_dim: 64, + causal: false, + masked: false, + }, + 3e-4, + 1.0, + ); +} + +#[test] +fn attention_masked_prefill_streams_long_kv_scale_one() { + check_attention_full_high_variance( + AttentionCase { + batch: 1, + heads: 4, + kv_heads: 4, + q_len: 512, + kv_len: 768, + head_dim: 64, + causal: false, + masked: true, + }, + 3e-4, + 1.0, + ); +} + #[test] fn attention_odd_sized_causal_prefill() { // Odd extents that don't align with any tile or bucket boundary. @@ -256,6 +397,113 @@ fn attention_masked_prefill() { ); } +#[test] +fn attention_offset_causal_mask_prefill() { + check_offset_causal_mask_prefill(20, 280, 4, 4); +} + +#[test] +fn attention_offset_causal_mask_prefill_single_tile() { + check_offset_causal_mask_prefill(20, 256, 4, 4); +} + +#[test] +fn attention_offset_causal_mask_prefill_gqa() { + check_offset_causal_mask_prefill(20, 256, 8, 4); +} + +fn check_offset_causal_mask_prefill(q_len: usize, kv_len: usize, heads: usize, kv_heads: usize) { + pollster::block_on(async { + let Ok(device) = Device::new().await else { + return; + }; + let case = AttentionCase { + batch: 1, + heads, + kv_heads, + q_len, + kv_len, + head_dim: 64, + causal: false, + masked: true, + }; + let q_data = + high_variance_values(case.batch * case.heads * case.q_len * case.head_dim, 1.0); + let k_data = high_variance_values( + case.batch * case.kv_heads * case.kv_len * case.head_dim, + 0.8, + ); + let v_data = high_variance_values( + case.batch * case.kv_heads * case.kv_len * case.head_dim, + 1.1, + ); + let prefix = case.kv_len - case.q_len; + let mask_data = (0..case.q_len * case.kv_len) + .map(|i| { + let q = i / case.kv_len; + let kv = i % case.kv_len; + if kv <= prefix + q { + 0.0 + } else { + f32::NEG_INFINITY + } + }) + .collect::>(); + let q = Tensor::from_slice( + &device, + [case.batch, case.heads, case.q_len, case.head_dim], + &q_data, + ); + let k = Tensor::from_slice( + &device, + [case.batch, case.kv_heads, case.kv_len, case.head_dim], + &k_data, + ); + let v = Tensor::from_slice( + &device, + [case.batch, case.kv_heads, case.kv_len, case.head_dim], + &v_data, + ); + let mask = Tensor::from_slice(&device, [case.q_len, case.kv_len], &mask_data); + let scale = 1.0; + + let out = q.flash_attention(&k, &v, scale, Some(&mask)); + assert_eq!( + out.count_kernels_to_resolve(), + 1, + "offset-causal mask attention should lower as one row-program kernel" + ); + let actual = out.as_slice::<4, f32>().await.unwrap(); + let expected = cpu_attention(&case, &q_data, &k_data, &v_data, Some(&mask_data), scale); + let mut worst = 0.0f32; + let mut worst_index = (0, 0, 0, 0); + for b in 0..case.batch { + for h in 0..case.heads { + for qi in 0..case.q_len { + for d in 0..case.head_dim { + let want = + expected[((b * case.heads + h) * case.q_len + qi) * case.head_dim + d]; + let got = actual[[b, h, qi, d]]; + let err = (got - want).abs(); + if err > worst { + worst = err; + worst_index = (b, h, qi, d); + } + } + } + } + } + assert!( + worst < 3e-4, + "worst offset-causal attention diff {worst} at {worst_index:?}: got {}, expected {}", + actual[[worst_index.0, worst_index.1, worst_index.2, worst_index.3]], + expected[((worst_index.0 * case.heads + worst_index.1) * case.q_len + worst_index.2) + * case.head_dim + + worst_index.3] + ); + }); +} + #[test] fn attention_f16_io() { pollster::block_on(async { diff --git a/fusor-ml/fusor/src/cache/mask_cache.rs b/fusor-ml/fusor/src/cache/mask_cache.rs index f05db019f..4b1707061 100644 --- a/fusor-ml/fusor/src/cache/mask_cache.rs +++ b/fusor-ml/fusor/src/cache/mask_cache.rs @@ -60,7 +60,12 @@ where } else { // Create the mask based on whether we have a sliding window let mask = if let Some(sliding_window_size) = sliding_window_size { - Self::create_sliding_window_mask(device, seq_len, sliding_window_size) + let mask = Self::create_sliding_window_mask(device, seq_len, sliding_window_size); + if seq_len <= sliding_window_size { + mask.mark_strict_causal() + } else { + mask + } } else { AttentionMask::causal(device, seq_len) }; diff --git a/fusor-ml/fusor/src/composite/flash_attention.rs b/fusor-ml/fusor/src/composite/flash_attention.rs index 32568638c..00aaba64f 100644 --- a/fusor-ml/fusor/src/composite/flash_attention.rs +++ b/fusor-ml/fusor/src/composite/flash_attention.rs @@ -68,12 +68,20 @@ where (Tensor::Gpu(q), Tensor::Gpu(k_gpu), Tensor::Gpu(v_gpu)) if !matches!(mask, Some((_, MaskKind::BatchKeyMask))) => { - // Decode (q_seq_len == 1) runs the DecodeSmall attention kernel, - // which uses workgroup reductions and needs no subgroups, so it - // works on browser adapters that report no subgroup support. Only - // the prefill/streaming kernels require subgroups — keep the - // fallback for those (q_seq_len > 1). - if !q.device().subgroups_supported() && self.shape()[2] != 1 { + let q_seq_len = self.shape()[2]; + let use_cpu_qk_mask_prefill = q_seq_len != 1 + && matches!(mask, Some((_, MaskKind::QKMask))) + && std::env::var_os("FUSOR_ALLOW_GPU_QK_MASK_PREFILL").is_none(); + let force_cpu_prefill = q_seq_len != 1 + && (std::env::var_os("FUSOR_FORCE_CPU_PREFILL_ATTENTION").is_some() + || (mask.is_none() + && std::env::var_os("FUSOR_FORCE_CPU_UNMASKED_PREFILL_ATTENTION") + .is_some()) + || (matches!(mask, Some((_, MaskKind::QKMask))) + && std::env::var_os("FUSOR_FORCE_CPU_QK_MASK_PREFILL_ATTENTION") + .is_some()) + || use_cpu_qk_mask_prefill); + if force_cpu_prefill { #[cfg(target_arch = "wasm32")] { return self.flash_attention_composite_impl(k, v, scale, mask); diff --git a/fusor-ml/fusor/src/layers/rms_norm.rs b/fusor-ml/fusor/src/layers/rms_norm.rs index e902f77ed..2f5bc29ea 100644 --- a/fusor-ml/fusor/src/layers/rms_norm.rs +++ b/fusor-ml/fusor/src/layers/rms_norm.rs @@ -98,6 +98,21 @@ where T: CastTo + CastTensor, f32: CastTo + CastTensor, { + /// Forward pass for 2D input with generic type. + /// Converts input to f32 for computation, then converts back. + pub fn forward_generic_2d(&self, input: &Tensor<2, T, B>) -> Tensor<2, T> + where + B: Fusion<2, T>, + { + let input_f32 = input.cast::(); + let weight_f32: Tensor<1, f32> = self.weight.cast(); + let bias_f32: Option> = self.bias.as_ref().map(|b| b.cast()); + + let result_f32 = input_f32.rms_norm_fused::<1, 1>(&weight_f32, bias_f32.as_ref(), self.eps); + + result_f32.cast() + } + /// Forward pass for 3D input with generic type. /// Converts input to f32 for computation, then converts back. pub fn forward_generic(&self, input: &Tensor<3, T, B>) -> Tensor<3, T> diff --git a/fusor-ml/tile-ir-kernels/src/dispatch.rs b/fusor-ml/tile-ir-kernels/src/dispatch.rs index 03e4d6a5e..63d167152 100644 --- a/fusor-ml/tile-ir-kernels/src/dispatch.rs +++ b/fusor-ml/tile-ir-kernels/src/dispatch.rs @@ -42,6 +42,10 @@ pub const fn qgemv_cols_per_workgroup(format: GgmlQuantFormat) -> u32 { /// This includes the Q4K/Q6K GGML specializations whose column grouping /// depends on both K (`rows`) and N (`cols`). pub fn qgemv_cols_per_workgroup_for_shape(format: GgmlQuantFormat, rows: u32, cols: u32) -> u32 { + if matches!(format, GgmlQuantFormat::Q4_0 | GgmlQuantFormat::Q4_0Native) { + return q4_0_override(q4_0_default(rows, cols)).cols_per_workgroup(); + } + if is_q4k_family(format) && rows <= 4096 && (4096..8192).contains(&cols) { return q4k_mid_override(q4k_default_mid(rows, cols)).cols_per_workgroup(); } @@ -95,12 +99,15 @@ pub(crate) const fn qgemv_subgroups_per_workgroup(format: GgmlQuantFormat) -> u3 } /// Shape-aware subgroup count used by the qgemv dispatch policy. -pub const fn qgemv_subgroups_per_workgroup_for_shape( +pub fn qgemv_subgroups_per_workgroup_for_shape( format: GgmlQuantFormat, rows: u32, - _cols: u32, + cols: u32, ) -> u32 { match format { + GgmlQuantFormat::Q4_0 | GgmlQuantFormat::Q4_0Native => { + q4_0_override(q4_0_default(rows, cols)).subgroups + } format if format.is_q6k_family() && rows > 4096 => 8, _ => qgemv_subgroups_per_workgroup(format), } @@ -235,6 +242,20 @@ const STANDARD_8_TILES: &[(&str, QgemvShape)] = &[ ("ggml_8x4", QgemvShape::new(8, 4)), ]; +const Q4_0_TILES: &[(&str, QgemvShape)] = &[ + ("ggml_1x4", QgemvShape::new(1, 4)), + ("ggml_1x8", QgemvShape::new(1, 8)), + ("ggml_2x2", QgemvShape::new(2, 2)), + ("ggml_2x4", QgemvShape::new(2, 4)), + ("ggml_2x8", QgemvShape::new(2, 8)), + ("ggml_4x2", QgemvShape::new(4, 2)), + ("ggml_4x4", QgemvShape::new(4, 4)), + ("ggml_4x8", QgemvShape::new(4, 8)), + ("ggml_8x1", QgemvShape::new(8, 1)), + ("ggml_8x2", QgemvShape::new(8, 2)), + ("ggml_8x4", QgemvShape::new(8, 4)), +]; + fn env_tile_override(var: &str, table: &[(&str, QgemvShape)], default: QgemvShape) -> QgemvShape { let Ok(value) = std::env::var(var) else { return default; @@ -250,6 +271,23 @@ pub(crate) fn q4k_mid_override(default: QgemvShape) -> QgemvShape { env_tile_override("FUSOR_Q4K_MID_TILE", Q4K_MID_TILES, default) } +pub(crate) const fn q4_0_default(_rows: u32, _cols: u32) -> QgemvShape { + QgemvShape::new(4, 4) +} + +pub(crate) fn q4_0_override(default: QgemvShape) -> QgemvShape { + env_tile_override("FUSOR_Q4_0_TILE", Q4_0_TILES, default) +} + +pub(crate) fn q4_0_values_per_lane(default: u32) -> u32 { + match std::env::var("FUSOR_Q4_0_VALUES_PER_LANE").ok().as_deref() { + Some("8") => 8, + Some("16") => 16, + Some("32") => 32, + _ => default, + } +} + // ----- Q4K large (rows<=4096, cols>=8192) ----- /// Default Q4K large-shape: cols<=16_384 → 8x4, else 2x4. diff --git a/fusor-ml/tile-ir-kernels/src/kernels/qgemv.rs b/fusor-ml/tile-ir-kernels/src/kernels/qgemv.rs index b2133e3cb..6d33a236b 100644 --- a/fusor-ml/tile-ir-kernels/src/kernels/qgemv.rs +++ b/fusor-ml/tile-ir-kernels/src/kernels/qgemv.rs @@ -19,8 +19,9 @@ use fusor_tile_ir::tile::{range, Mask, Program, Storage, Tile, TileBlock}; use fusor_tile_ir::{GgmlQuantFormat, QuantizedMatrix, TileLiteral}; use crate::dispatch::{ - q4k_default_large, q4k_default_mid, q4k_default_tall, q4k_large_override, q4k_mid_override, - q4k_tall_override, q6k_default_large, q6k_default_tall, q6k_large_override, q6k_tall_override, + q4_0_default, q4_0_override, q4_0_values_per_lane, q4k_default_large, q4k_default_mid, + q4k_default_tall, q4k_large_override, q4k_mid_override, q4k_tall_override, q6k_default_large, + q6k_default_tall, q6k_large_override, q6k_tall_override, qgemv_subgroups_per_workgroup_for_shape, QgemvShape, SubgroupConfig, }; use crate::grid::{ @@ -216,19 +217,30 @@ pub(crate) fn qgemv_tile_with_epilogue( qgemv_shape(2, 4), 16, ), - GgmlQuantFormat::Q4_0 - | GgmlQuantFormat::Q4_0Native - | GgmlQuantFormat::Q4_1 - | GgmlQuantFormat::Q5_1 - | GgmlQuantFormat::Q2K => qgemv_perf_with_epilogue( - program, - tensors, - workgroups_x, - subgroups, - ep, - qgemv_shape(2, 4), - 8, - ), + GgmlQuantFormat::Q4_0 | GgmlQuantFormat::Q4_0Native => { + let shape = q4_0_override(q4_0_default(b.rows, output_cols)); + let values_per_lane = q4_0_values_per_lane(8); + qgemv_perf_with_epilogue( + program, + tensors, + workgroups_x, + subgroups, + ep, + shape, + values_per_lane, + ) + } + GgmlQuantFormat::Q4_1 | GgmlQuantFormat::Q5_1 | GgmlQuantFormat::Q2K => { + qgemv_perf_with_epilogue( + program, + tensors, + workgroups_x, + subgroups, + ep, + qgemv_shape(2, 4), + 8, + ) + } GgmlQuantFormat::Q3K | GgmlQuantFormat::Q8K => qgemv_perf_with_epilogue( program, tensors, @@ -336,6 +348,11 @@ fn select_qgemv_dot( if format.is_q8_0_family() && values_per_lane == 8 { return QgemvDot::F32Vec; } + if format.is_q4_0_family() + && (values_per_lane == 8 || values_per_lane == 16 || values_per_lane == 32) + { + return QgemvDot::F32Vec; + } if format.is_q4k_family() && (values_per_lane == 8 || values_per_lane == 16 || values_per_lane == 32) { diff --git a/fusor-ml/tile-ir/src/lower/quantized/helpers.rs b/fusor-ml/tile-ir/src/lower/quantized/helpers.rs index a5a708be4..98154c2aa 100644 --- a/fusor-ml/tile-ir/src/lower/quantized/helpers.rs +++ b/fusor-ml/tile-ir/src/lower/quantized/helpers.rs @@ -139,6 +139,19 @@ impl<'a> Lowerer<'a> { } } + pub(in crate::lower) fn q4_0_data_byte_offset( + &self, + format: GgmlQuantFormat, + ) -> Result { + match format { + GgmlQuantFormat::Q4_0 => Ok(4), + GgmlQuantFormat::Q4_0Native => Ok(2), + _ => Err(LowerError::UnsupportedOperation( + "q4_0 data offset requires a Q4_0 format", + )), + } + } + pub(in crate::lower) fn q4k_data_word_offset( &self, format: GgmlQuantFormat, diff --git a/fusor-ml/tile-ir/src/lower/quantized/values.rs b/fusor-ml/tile-ir/src/lower/quantized/values.rs index 8df0c445f..1cdb6f259 100644 --- a/fusor-ml/tile-ir/src/lower/quantized/values.rs +++ b/fusor-ml/tile-ir/src/lower/quantized/values.rs @@ -103,6 +103,82 @@ impl<'a> Lowerer<'a> { Ok(self.mul(expressions, body, sum, parts.scale)) } + pub(in crate::lower) fn q4_0_f32_dot( + &self, + expressions: &mut Arena, + matrix: &QuantizedMatrix, + k_base: Handle, + col: Handle, + a: &[Handle], + body: &mut Block, + ) -> Result, LowerError> { + if !matrix.format.is_q4_0_family() || !matches!(a.len(), 8 | 16 | 32) { + return Err(LowerError::UnsupportedOperation( + "q4_0 f32 dot requires a Q4_0 format and 8, 16, or 32 activation values", + )); + } + + let (base, q_base) = self.quantized_flat_block_base_and_q( + expressions, + matrix, + k_base, + col, + matrix.format.block_words(), + body, + ); + let scale = self.load_affine_scale_f32(expressions, matrix, base, 0, body)?; + let data_offset = self.q4_0_data_byte_offset(matrix.format)?; + let q_local = self.and_lit(expressions, body, q_base, 15); + let high = self.cmp_lit(expressions, body, BinaryOperator::GreaterEqual, q_base, 16); + let word_count = if a.len() == 32 { 4 } else { a.len() / 4 }; + let mut words = Vec::with_capacity(word_count); + for word_index in 0..word_count { + let byte_offset = self.add_lit( + expressions, + body, + q_local, + data_offset + (word_index * 4) as u32, + ); + words.push(self.load_word_at_block_dynamic_byte_offset( + expressions, + matrix, + base, + byte_offset, + body, + )?); + } + + let mut quants = Vec::with_capacity(a.len()); + for lane in 0..a.len() { + let byte_lane = self.u32(expressions, (lane % 4) as u32); + let word = if a.len() == 32 { + words[(lane % 16) / 4] + } else { + words[lane / 4] + }; + let byte = self.byte_at(expressions, body, word, byte_lane); + let low = self.and_lit(expressions, body, byte, 0x0f); + let high4 = self.shr_lit(expressions, body, byte, 4); + let quant = if a.len() == 32 { + if lane >= 16 { + high4 + } else { + low + } + } else { + self.select(expressions, body, high, high4, low) + }; + quants.push(quant); + } + + let weighted_sum = self.dot_quant_vec4_chunks(expressions, body, a, &quants); + let activation_sum = self.sum_values(expressions, body, a); + let center = self.f32(expressions, 8.0); + let center_term = self.mul(expressions, body, activation_sum, center); + let centered = self.sub(expressions, body, weighted_sum, center_term); + Ok(self.mul(expressions, body, centered, scale)) + } + pub(in crate::lower) fn q8_0_block_parts8( &self, expressions: &mut Arena, diff --git a/fusor-ml/tile-ir/src/lower/tile_program/quantized.rs b/fusor-ml/tile-ir/src/lower/tile_program/quantized.rs index ef7b9bf66..8c432f427 100644 --- a/fusor-ml/tile-ir/src/lower/tile_program/quantized.rs +++ b/fusor-ml/tile-ir/src/lower/tile_program/quantized.rs @@ -135,6 +135,9 @@ impl<'a> Lowerer<'a> { let c = self.lower_expr_lane(expressions, block, col, spill_depth)?; match packing { QuantActivation::F32 => match (src.format, block_n) { + (GgmlQuantFormat::Q4_0 | GgmlQuantFormat::Q4_0Native, 8 | 16 | 32) => { + self.q4_0_f32_dot(expressions, src, k, c, &a_handles, block) + } (GgmlQuantFormat::Q8_0 | GgmlQuantFormat::Q8_0Native, 8) => { let a8 = Self::expect_dot8(&a_handles)?; self.dequantize_q8_0_dot8(expressions, src, k, c, &a8, block) @@ -147,7 +150,7 @@ impl<'a> Lowerer<'a> { self.q4k_f32_dot(expressions, src, k, c, &a_handles, block) } _ => Err(LowerError::UnsupportedOperation( - "f32 activation dot only supports Q8_0/Q6K dot8 or Q4K dot8/16/32", + "f32 activation dot only supports Q4_0 dot8/16, Q8_0/Q6K dot8, or Q4K dot8/16/32", )), }, QuantActivation::Q8 => { diff --git a/fusor-ml/tile-ir/src/quantized.rs b/fusor-ml/tile-ir/src/quantized.rs index b9fdebe24..eebf4ed44 100644 --- a/fusor-ml/tile-ir/src/quantized.rs +++ b/fusor-ml/tile-ir/src/quantized.rs @@ -111,6 +111,10 @@ impl GgmlQuantFormat { matches!(self, Self::Q4K | Self::Q4KNative) } + pub const fn is_q4_0_family(self) -> bool { + matches!(self, Self::Q4_0 | Self::Q4_0Native) + } + pub const fn is_q8_0_family(self) -> bool { matches!(self, Self::Q8_0 | Self::Q8_0Native) } diff --git a/interfaces/kalosm/examples/vision.rs b/interfaces/kalosm/examples/vision.rs index 3ccdc50a0..0e5ac5cff 100644 --- a/interfaces/kalosm/examples/vision.rs +++ b/interfaces/kalosm/examples/vision.rs @@ -5,36 +5,44 @@ use std::time::Instant; async fn main() { tracing_subscriber::fmt::init(); let t_load_start = Instant::now(); - let model = Llama::builder() - .with_source(LlamaSource::qwen_2_5_3b_vl_chat_q4()) - .build() - .await - .unwrap(); + let mut builder = Llama::builder().with_source(LlamaSource::gemma_4_e2b_it_qat_chat()); + builder = if std::env::var_os("KALOSM_VISION_CPU").is_some() { + builder.with_device(Device::Cpu) + } else { + builder.with_device(Device::gpu().await.expect( + "The vision example requires a GPU by default; set KALOSM_VISION_CPU=1 to run the slow CPU path.", + )) + }; + let model = builder.build().await.unwrap(); tracing::info!("[timing] model load: {:.2?}", t_load_start.elapsed()); let mut chat = model.chat(); let max_tokens = std::env::var("KALOSM_VISION_MAX_TOKENS") .ok() .and_then(|value| value.parse::().ok()) - .unwrap_or(64); + .unwrap_or(1024); let t_total = Instant::now(); let image_source = if let Ok(url) = std::env::var("KALOSM_VISION_URL") { MediaSource::url(url) } else if let Ok(path) = std::env::var("KALOSM_VISION_IMAGE") { MediaSource::file(path).unwrap() } else { - MediaSource::url("https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen-VL/assets/demo.jpeg") + MediaSource::bytes(include_bytes!("landscape.jpg").as_slice()) }; - let mut response = chat(&( - MediaChunk::new(image_source, MediaType::Image), - "Describe this image.", - )); + let mut sampler = GenerationParameters::new() + .with_standard_sampler() + .with_temperature(0.0); if let Some(seed) = std::env::var("KALOSM_VISION_SEED") .ok() .and_then(|seed| seed.parse::().ok()) { - response = response.with_sampler(GenerationParameters::new().with_seed(seed)); + sampler = sampler.with_seed(seed); } + let mut response = chat(&( + MediaChunk::new(image_source, MediaType::Image), + "Describe this image.", + )) + .with_sampler(sampler); let mut first_token_at: Option = None; let mut token_count = 0u64; let t_prefill = Instant::now(); @@ -48,6 +56,7 @@ async fn main() { } token_count += 1; print!("{}", token); + std::io::Write::flush(&mut std::io::stdout()).unwrap(); if token_count >= max_tokens { break; } diff --git a/models/kalosm-llama/examples/vision.rs b/models/kalosm-llama/examples/vision.rs index 421634e4b..425ccb572 100644 --- a/models/kalosm-llama/examples/vision.rs +++ b/models/kalosm-llama/examples/vision.rs @@ -1,27 +1,141 @@ +#![recursion_limit = "256"] + use kalosm_llama::prelude::*; use kalosm_streams::text_stream::TextStream; +use std::time::Instant; + +async fn collect_or_bench(mut response: impl TextStream + Unpin, max_tokens: u32) { + let bench = std::env::var_os("KALOSM_VISION_BENCH").is_some(); + let total_start = Instant::now(); + let mut first = None; + let mut last = None; + let mut tokens = 0usize; + let mut output = String::new(); + while let Some(token) = response.next().await { + let now = Instant::now(); + first.get_or_insert(now); + last = Some(now); + tokens += 1; + output.push_str(&token); + } + + if !bench { + print!("{output}"); + println!("\n"); + return; + } -// The demo image is fetched over HTTP, which needs a tokio reactor. + let total_elapsed = total_start.elapsed().as_secs_f64(); + let steady_elapsed = first + .zip(last) + .map(|(first, last)| (last - first).as_secs_f64()) + .unwrap_or_default(); + let steady_tokens = tokens.saturating_sub(1); + let total_tps = if total_elapsed > 0.0 { + tokens as f64 / total_elapsed + } else { + 0.0 + }; + let steady_tps = if steady_elapsed > 0.0 { + steady_tokens as f64 / steady_elapsed + } else { + 0.0 + }; + + print!("{output}"); + println!("\n"); + eprintln!( + "vision_bench tokens={tokens} max_tokens={max_tokens} total_s={total_elapsed:.3} total_tok_s={total_tps:.2} steady_tok_s={steady_tps:.2}" + ); +} + +// The demo image may be fetched over HTTP, which needs a tokio reactor. #[tokio::main] async fn main() { tracing_subscriber::fmt::init(); - let model = Llama::builder() - .with_source(LlamaSource::qwen_2_5_3b_vl_chat_q4()) - .build() - .await - .unwrap(); + let source = match std::env::var("KALOSM_VISION_SOURCE").as_deref() { + Ok("gemma4") => { + let mut source = LlamaSource::gemma_4_e2b_it_qat_chat().with_vision_model( + kalosm_model_types::FileSource::HuggingFace { + model_id: "unsloth/gemma-4-E2B-it-qat-GGUF".into(), + revision: "main".into(), + file: "mmproj-F16.gguf".into(), + }, + ); + if std::env::var_os("KALOSM_LLAMA_MTP_DRAFT").is_some() { + let file = std::env::var("KALOSM_LLAMA_MTP_FILE") + .unwrap_or_else(|_| "MTP/gemma-4-E2B-it-Q4_0-MTP.gguf".into()); + source = source.with_mtp_model(kalosm_model_types::FileSource::HuggingFace { + model_id: "unsloth/gemma-4-E2B-it-qat-GGUF".into(), + revision: "main".into(), + file, + }); + } + source + } + _ => LlamaSource::qwen_2_5_3b_vl_chat_q4(), + }; + let mut builder = Llama::builder().with_source(source); + if std::env::var_os("KALOSM_VISION_CPU").is_some() { + builder = builder.with_device(kalosm_llama::Device::Cpu); + } + let model = builder.build().await.unwrap(); let mut chat = model.chat(); + let max_tokens = std::env::var("KALOSM_VISION_MAX_TOKENS") + .ok() + .and_then(|value| value.parse::().ok()) + .unwrap_or(64); + let mut sampler = GenerationParameters::new().with_max_length(max_tokens); + if let Ok(seed) = std::env::var("KALOSM_VISION_SEED") { + if let Ok(seed) = seed.parse::() { + sampler = sampler.with_seed(seed); + } + } + if std::env::var_os("KALOSM_VISION_STANDARD").is_some() { + sampler = sampler.with_standard_sampler(); + } + if let Ok(temperature) = std::env::var("KALOSM_VISION_TEMPERATURE") { + if let Ok(temperature) = temperature.parse::() { + sampler = sampler.with_temperature(temperature); + if temperature <= 0.0 { + sampler = sampler.with_standard_sampler(); + } + } + } + if std::env::var_os("KALOSM_VISION_TEXT_ONLY").is_some() { + let prompt = std::env::var("KALOSM_VISION_TEXT_PROMPT").unwrap_or_else(|_| { + "Answer in one short sentence: what color is a ripe banana?".into() + }); + let mut response = chat(&prompt).with_sampler(sampler); + if std::env::var_os("KALOSM_VISION_COLLECT").is_some() { + collect_or_bench(response, max_tokens).await; + return; + } + response.to_std_out().await.unwrap(); + println!("\n"); + return; + } + + let image_source = if let Ok(url) = std::env::var("KALOSM_VISION_URL") { + MediaSource::url(url) + } else if let Ok(path) = std::env::var("KALOSM_VISION_IMAGE") { + MediaSource::file(path).unwrap() + } else { + MediaSource::url("https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen-VL/assets/demo.jpeg") + }; + let prompt = + std::env::var("KALOSM_VISION_PROMPT").unwrap_or_else(|_| "Describe this image.".into()); + let mut response = chat(&( - MediaChunk::new( - MediaSource::url( - "https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen-VL/assets/demo.jpeg", - ), - MediaType::Image, - ), - "Describe this image.", - )); + MediaChunk::new(image_source, MediaType::Image), + prompt.as_str(), + )) + .with_sampler(sampler); + if std::env::var_os("KALOSM_VISION_COLLECT").is_some() { + collect_or_bench(response, max_tokens).await; + return; + } response.to_std_out().await.unwrap(); - response.await.unwrap(); println!("\n"); } diff --git a/models/kalosm-llama/src/chat_template.rs b/models/kalosm-llama/src/chat_template.rs index bfb11b54a..59f1648ae 100644 --- a/models/kalosm-llama/src/chat_template.rs +++ b/models/kalosm-llama/src/chat_template.rs @@ -1,6 +1,6 @@ -use std::fmt::Display; +use std::{collections::BTreeMap, fmt::Display}; -use kalosm_language_model::{ChatMessage, ContentChunk}; +use kalosm_language_model::{ChatMessage, ContentChunk, MediaType}; use minijinja::{context, Environment, ErrorKind, Value}; use minijinja_contrib::pycompat; @@ -11,11 +11,14 @@ use pretty_assertions::assert_eq; pub(crate) struct HuggingFaceChatTemplate { environment: Environment<'static>, + enable_thinking_by_default: bool, } impl HuggingFaceChatTemplate { pub(crate) fn create(chat_template: impl Display) -> Result { let chat_template = chat_template.to_string(); + let enable_thinking_by_default = + chat_template.contains("enable_thinking") && chat_template.contains("<|think|>"); let mut environment = Environment::new(); // enable python compatibility methods because most models are tested with python @@ -40,7 +43,10 @@ impl HuggingFaceChatTemplate { // compile the template expression in the environment environment.add_template_owned("main", chat_template)?; - Ok(Self { environment }) + Ok(Self { + environment, + enable_thinking_by_default, + }) } pub(crate) fn format( @@ -64,11 +70,17 @@ impl HuggingFaceChatTemplate { .iter() .map(|chunk| match chunk { ContentChunk::Text(text) => { - context! { text } - } - ContentChunk::Media(_) => { - context! { image => "" } + chunk_context([("type", "text"), ("text", text)]) } + ContentChunk::Media(media) => match media.media_type() { + MediaType::Image => { + chunk_context([("type", "image"), ("image", "")]) + } + MediaType::Video => { + chunk_context([("type", "video"), ("video", "")]) + } + _ => chunk_context([("type", "media")]), + }, }) .collect::>(); chunks.into() @@ -76,13 +88,18 @@ impl HuggingFaceChatTemplate { context! { role, content } }) .collect::>(); - let ctx = context! { bos_token, eos_token, messages, add_generation_prompt, tools }; + let enable_thinking = self.enable_thinking_by_default; + let ctx = context! { bos_token, eos_token, messages, add_generation_prompt, tools, enable_thinking }; let template = self.environment.get_template("main")?; let result = template.render(&ctx)?; Ok(result) } } +fn chunk_context(entries: [(&str, &str); N]) -> Value { + Value::from_serialize(BTreeMap::from(entries)) +} + #[test] fn test_qwen_chat_template() { let template = r#"{%- if tools %} @@ -238,6 +255,49 @@ I'd like to show off how chat templating works!<|vision_start|><|image_pad|><|vi ); } +#[test] +fn test_gemma_vl_chat_template_media_type() { + use kalosm_language_model::{MediaChunk, MediaSource}; + + let template = "{% for message in messages %}{% for item in message['content'] %}{% if item['type'] == 'image' %}<|image|>{% elif item['type'] == 'text' %}{{ item['text'] }}{% endif %}{% endfor %}{% endfor %}"; + let template = HuggingFaceChatTemplate::create(template).unwrap(); + let inputs = [ChatMessage::new( + MessageType::UserMessage, + ( + MediaChunk::new( + MediaSource::url("https://example.com/image.png"), + kalosm_language_model::MediaType::Image, + ), + "Describe this image.", + ), + )]; + let result = template.format("", "", &inputs, false).unwrap(); + assert_eq!(result, "<|image|>Describe this image."); +} + +#[test] +fn test_gemma_4_chat_template_enables_thinking() { + use kalosm_language_model::{MediaChunk, MediaSource}; + + let template = "{{- bos_token -}}{%- if enable_thinking is defined and enable_thinking -%}{{- '<|turn>system\n<|think|>\n\n' -}}{%- endif -%}{%- for message in messages -%}{{- '<|turn>' + message['role'] + '\n' -}}{%- for item in message['content'] -%}{%- if item['type'] == 'image' -%}{{- '<|image|>' -}}{%- elif item['type'] == 'text' -%}{{- item['text'] | trim -}}{%- endif -%}{%- endfor -%}{{- '\n' -}}{%- endfor -%}{%- if add_generation_prompt -%}{{- '<|turn>model\n' -}}{%- endif -%}"; + let template = HuggingFaceChatTemplate::create(template).unwrap(); + let inputs = [ChatMessage::new( + MessageType::UserMessage, + ( + MediaChunk::new( + MediaSource::url("https://example.com/image.png"), + kalosm_language_model::MediaType::Image, + ), + "Describe this image.", + ), + )]; + let result = template.format("", "", &inputs, true).unwrap(); + assert_eq!( + result, + "<|turn>system\n<|think|>\n\n<|turn>user\n<|image|>Describe this image.\n<|turn>model\n" + ); +} + #[test] fn test_llama_chat_template() { let template = "{% set loop_messages = messages %}{% for message in loop_messages %}{% set content = '<|start_header_id|>' + message['role'] + '<|end_header_id|>\n\n'+ message['content'] | trim + '<|eot_id|>' %}{% if loop.index0 == 0 %}{% set content = bos_token + content %}{% endif %}{{ content }}{% endfor %}{% if add_generation_prompt %}{{ '<|start_header_id|>assistant<|end_header_id|>\n\n' }}{% endif %}"; diff --git a/models/kalosm-llama/src/gguf_tokenizer.rs b/models/kalosm-llama/src/gguf_tokenizer.rs index 95d2c037a..a962782ae 100644 --- a/models/kalosm-llama/src/gguf_tokenizer.rs +++ b/models/kalosm-llama/src/gguf_tokenizer.rs @@ -21,6 +21,7 @@ enum PreTokenizerType { Default, Exaone, Falcon, + Gemma4, Gpt2, Gpt3Finnish, Jais, @@ -69,6 +70,7 @@ impl PreTokenizerType { "'s|'t|'re|'ve|'m|'ll|'d| ?\\p{L}+| ?\\p{N}+| ?[^\\s\\p{L}\\p{N}]+|\\s+(?!\\S)", "[0-9][0-9][0-9]", ], + Self::Gemma4 => &["[^\\n]+|[\\n]+"], Self::Starcoder | Self::Refact | Self::CommandR @@ -261,6 +263,8 @@ fn sanitize_regex(regex: &str) -> String { #[derive(Clone, Copy)] pub(crate) struct GGUFPreTokenizerConfig { add_bos: bool, + byte_level: bool, + escape_whitespaces: bool, ignore_merges: bool, ty: PreTokenizerType, } @@ -298,7 +302,11 @@ impl GGUFPreTokenizerConfig { let merges = merges .into_iter() .map(|(left, right)| format!("{left} {right}")); - let bpe = FastBpe::from_vocab_and_merges(vocab, merges, self.ignore_merges)?; + let bpe = if self.byte_level { + FastBpe::from_vocab_and_merges(vocab, merges, self.ignore_merges)? + } else { + FastBpe::from_raw_utf8_vocab_and_merges(vocab, merges, self.ignore_merges)? + }; let pre_tokenizer = PreTokenizer::new(self.ty)?; Ok(GgufTokenizer { @@ -308,6 +316,7 @@ impl GGUFPreTokenizerConfig { special_token_matches, bos_token, add_bos: self.add_bos, + escape_whitespaces: self.escape_whitespaces, }) } } @@ -316,6 +325,8 @@ impl Default for GGUFPreTokenizerConfig { fn default() -> Self { Self { add_bos: true, + byte_level: true, + escape_whitespaces: false, ignore_merges: false, ty: PreTokenizerType::Default, } @@ -330,6 +341,7 @@ pub(crate) struct GgufTokenizer { special_token_matches: Vec, u32)>>, bos_token: Option, add_bos: bool, + escape_whitespaces: bool, } impl GgufTokenizer { @@ -394,7 +406,12 @@ impl GgufTokenizer { bytes.extend_from_slice(token_bytes); } } - String::from_utf8_lossy(&bytes).into_owned() + let text = String::from_utf8_lossy(&bytes).into_owned(); + if self.escape_whitespaces { + text.replace('▁', " ") + } else { + text + } } pub(crate) fn is_special_token(&self, token: u32) -> bool { @@ -419,6 +436,14 @@ impl GgufTokenizer { tokenization, } = buffers; + let normalized; + let text = if self.escape_whitespaces { + normalized = text.replace(' ', "▁"); + normalized.as_str() + } else { + text + }; + self.pre_tokenizer.split_into_ranges(text, pre_tokenizer); for piece in pre_tokenizer.pieces.iter().copied() { if piece.is_empty() { @@ -460,6 +485,14 @@ pub(crate) fn get_pre_tokenizer( ignore_merges: true, add_bos: true, ty: PreTokenizerType::Llama3, + ..Default::default() + }, + "gemma4" => GGUFPreTokenizerConfig { + add_bos: true, + byte_level: false, + escape_whitespaces: true, + ty: PreTokenizerType::Gemma4, + ..Default::default() }, "deepseek-llm" => GGUFPreTokenizerConfig { ty: PreTokenizerType::DeepseekLlm, @@ -534,6 +567,7 @@ pub(crate) fn get_pre_tokenizer( ignore_merges: true, add_bos: true, ty: PreTokenizerType::Tekken, + ..Default::default() }, "smollm" => GGUFPreTokenizerConfig { ty: PreTokenizerType::Smollm, @@ -569,7 +603,7 @@ pub(crate) fn get_pre_tokenizer( }; if let Some(add_bos) = add_bos { - tokenizer.add_bos = add_bos; + tokenizer.add_bos = add_bos || pre_tokenizer_type == "gemma4"; } tokenizer diff --git a/models/kalosm-llama/src/language_model.rs b/models/kalosm-llama/src/language_model.rs index ac8d71a8c..c659530b4 100644 --- a/models/kalosm-llama/src/language_model.rs +++ b/models/kalosm-llama/src/language_model.rs @@ -54,7 +54,12 @@ where .tokenizer .as_ref() .is_some_and(|tokenizer| !cache.exists(tokenizer)); - model_missing || tokenizer_missing + let mtp_missing = self + .source + .mtp_model + .as_ref() + .is_some_and(|mtp| !cache.exists(mtp)); + model_missing || tokenizer_missing || mtp_missing } } @@ -118,6 +123,13 @@ where } Vec::new() }; + if std::env::var_os("KALOSM_TRACE_PROMPT").is_some() { + eprintln!( + "[stream_text] prompt_len={} image_count={}", + text.len(), + images.len() + ); + } self.inner .sender .unbounded_send(Task::UnstructuredGeneration(UnstructuredGenerationTask::< diff --git a/models/kalosm-llama/src/model/forward.rs b/models/kalosm-llama/src/model/forward.rs index fc93bf559..fcbb6e7b1 100644 --- a/models/kalosm-llama/src/model/forward.rs +++ b/models/kalosm-llama/src/model/forward.rs @@ -46,7 +46,11 @@ where }; let token_start = trace_enabled.then(Instant::now); let build_start = trace_enabled.then(Instant::now); - let logits = model.forward(tokens, images, device, cache); + let logits = if !images.is_empty() && model.should_chunk_multimodal_prompt() { + model.forward_chunked_multimodal(tokens, images, device, cache) + } else { + model.forward(tokens, images, device, cache) + }; if let Some(start) = build_start { tracing::info!( "forward_graph_build path={path} decode_eligible={decode_eligible} elapsed={:?}", @@ -229,7 +233,9 @@ where ); } - if gpu_fused_logits_sampling_enabled() { + if gpu_fused_logits_sampling_enabled() + && (images.is_empty() || !model.should_chunk_multimodal_prompt()) + { return Self::forward_sample_token_fused_logits( ForwardInputs { model, @@ -327,6 +333,18 @@ where return Ok(None); } + if !images.is_empty() && model.should_chunk_multimodal_prompt() { + let logits = model.forward_chunked_multimodal(tokens, images, device, cache)?; + let logits: fusor::Tensor<1, f32> = logits.squeeze(0).cast(); + return Self::forward_sample_token_pending_from_logits( + logits, + sampler, + previous_tokens, + None, + top_k, + ); + } + Self::forward_sample_token_pending_from_hidden( model.forward_last_hidden_f32(tokens, images, device, cache)?, model, @@ -337,6 +355,37 @@ where ) } + fn forward_sample_token_pending_from_logits( + logits: fusor::Tensor<1, f32>, + sampler: &mut LlamaGpuSamplerState, + previous_tokens: Vec, + previous_gpu_token: Option<&fusor::GpuSampledToken>, + top_k: usize, + ) -> Result, LlamaModelError> { + match sampler.config.sampling_strategy { + kalosm_language_model::SamplingStrategy::Mirostat2 => { + let params = sampler.mirostat_params(top_k); + let mirostat = sampler.mirostat.as_mut().ok_or_else(|| { + LlamaModelError::SamplerError("missing Mirostat GPU sampler".into()) + })?; + logits + .sample_mirostat2_token_pending( + mirostat, + &previous_tokens, + previous_gpu_token, + params, + ) + .map_err(LlamaModelError::from) + } + kalosm_language_model::SamplingStrategy::Standard => { + let params = sampler.standard_params(top_k); + logits + .sample_standard_token_pending(&previous_tokens, previous_gpu_token, params) + .map_err(LlamaModelError::from) + } + } + } + pub(crate) fn forward_sample_token_from_gpu_token_pending( model: &Model, device: &Device, @@ -377,32 +426,14 @@ where previous_gpu_token: Option<&fusor::GpuSampledToken>, top_k: usize, ) -> Result, LlamaModelError> { - let logits = hidden - .squeeze(0) - .to_concrete() - .q_mat_mul(model.output_matrix()); - match sampler.config.sampling_strategy { - kalosm_language_model::SamplingStrategy::Mirostat2 => { - let params = sampler.mirostat_params(top_k); - let mirostat = sampler.mirostat.as_mut().ok_or_else(|| { - LlamaModelError::SamplerError("missing Mirostat GPU sampler".into()) - })?; - logits - .sample_mirostat2_token_pending( - mirostat, - &previous_tokens, - previous_gpu_token, - params, - ) - .map_err(LlamaModelError::from) - } - kalosm_language_model::SamplingStrategy::Standard => { - let params = sampler.standard_params(top_k); - logits - .sample_standard_token_pending(&previous_tokens, previous_gpu_token, params) - .map_err(LlamaModelError::from) - } - } + let logits = model.logits_from_hidden_f32(hidden.squeeze(0).to_concrete()); + Self::forward_sample_token_pending_from_logits( + logits, + sampler, + previous_tokens, + previous_gpu_token, + top_k, + ) } fn forward_sample_token_fused_logits<'a>( @@ -443,10 +474,7 @@ where Ok(hidden) => hidden, Err(err) => return Box::pin(async move { Err(err.into()) }), }; - let logits = hidden - .squeeze(0) - .to_concrete() - .q_mat_mul(model.output_matrix()); + let logits = model.logits_from_hidden_f32(hidden.squeeze(0).to_concrete()); let mut kernels = 0; if trace { if let Some(gpu_logits) = logits.as_gpu() { diff --git a/models/kalosm-llama/src/model/inference.rs b/models/kalosm-llama/src/model/inference.rs index ae36773a1..8713c0ea1 100644 --- a/models/kalosm-llama/src/model/inference.rs +++ b/models/kalosm-llama/src/model/inference.rs @@ -1,5 +1,22 @@ use super::*; +/// Yield once to the async runtime so long generation loops stay cooperative +/// between GPU dispatches without spinning the executor. +async fn yield_once() { + use std::sync::atomic::{AtomicBool, Ordering}; + let yielded = AtomicBool::new(false); + std::future::poll_fn(|cx| { + if yielded.load(Ordering::Relaxed) { + std::task::Poll::Ready(()) + } else { + yielded.store(true, Ordering::Relaxed); + cx.waker().wake_by_ref(); + std::task::Poll::Pending + } + }) + .await; +} + impl LlamaModel where F: CastTo + CastTensor + WasmNotSend + WasmNotSync + 'static, @@ -8,6 +25,21 @@ where AddOp: SimdBinaryOp, SumOp: SimdReduceOp, { + /// Emit a `[sampled_token]` trace line when `KALOSM_TRACE_SAMPLED_TOKEN` is + /// set. No-op (beyond an env lookup) otherwise. + fn trace_sampled_token(&self, index: impl std::fmt::Display, token: u32, stop_tokens: &[u32]) { + if std::env::var_os("KALOSM_TRACE_SAMPLED_TOKEN").is_none() { + return; + } + let decoded = self + .tokenizer + .decode(&[token], false) + .unwrap_or_else(|err| format!("")); + eprintln!( + "[sampled_token] index={index} id={token} text={decoded:?} stop_tokens={stop_tokens:?}" + ); + } + pub(crate) async fn _infer( &mut self, settings: InferenceSettings, @@ -28,6 +60,16 @@ where .tokenizer .encode(&prompt, false) .map_err(LlamaModelError::Tokenizer)?; + if std::env::var_os("KALOSM_TRACE_PROMPT").is_some() { + let decoded = self + .tokenizer + .decode(&tokens, false) + .unwrap_or_else(|err| format!("")); + eprintln!( + "[prompt] len={} text={prompt:?} tokens={tokens:?} decoded={decoded:?}", + tokens.len() + ); + } let mut text_stream = TokenOutputStream::new(self.tokenizer.clone()); for &token in &tokens { text_stream @@ -35,13 +77,34 @@ where .map_err(LlamaModelError::TokenOutputStreamError)?; } + if mtp_speculative_enabled() + && images.is_empty() + && stop_on.is_none() + && sampler.sampling_strategy == kalosm_language_model::SamplingStrategy::Standard + && sampler.temperature <= 0.0 + && gpu_token_sampling_enabled() + && self.mtp.is_some() + { + if let Some(gpu_sampler) = LlamaGpuSamplerState::new(&self.device, sampler, seed) { + return self + .infer_mtp_speculative( + &tokens, + &images, + text_stream, + session, + max_tokens, + gpu_sampler, + on_token, + finished, + ) + .await; + } + } + if gpu_token_sampling_enabled() && stop_on.is_none() { if let Some(mut gpu_sampler) = LlamaGpuSamplerState::new(&self.device, sampler, seed) { let top_k = gpu_sample_top_k(&gpu_sampler.config); - if gpu_run_ahead_enabled() - && images.is_empty() - && self.model.supports_gpu_token_run_ahead() - { + if gpu_run_ahead_enabled() && self.model.supports_gpu_token_run_ahead() { let next_token = { let previous_tokens = gpu_sampler.previous_tokens(&text_stream); let mut session_lock = session @@ -64,7 +127,7 @@ where }; if let Some(mut next_token) = next_token { - let stop_token = self.model.config.stop_token; + let stop_tokens = &self.model.config.stop_tokens; let mut tokens_generated = 0; while !finished.is_canceled() && tokens_generated < max_tokens { let scheduled_next = if tokens_generated + 1 < max_tokens { @@ -102,7 +165,8 @@ where "pending GPU sampler refused slow fallback".into(), ) })?; - if new_token == stop_token { + self.trace_sampled_token(tokens_generated, new_token, stop_tokens); + if stop_tokens.contains(&new_token) { tracing::trace!("Stopping on stop token"); break; } @@ -153,20 +217,7 @@ where break; } - { - use std::sync::atomic::{AtomicBool, Ordering}; - let yielded = AtomicBool::new(false); - std::future::poll_fn(|cx| { - if yielded.load(Ordering::Relaxed) { - std::task::Poll::Ready(()) - } else { - yielded.store(true, Ordering::Relaxed); - cx.waker().wake_by_ref(); - std::task::Poll::Pending - } - }) - .await; - } + yield_once().await; } return Ok(()); @@ -196,11 +247,12 @@ where } .await?; - let stop_token = self.model.config.stop_token; + let stop_tokens = &self.model.config.stop_tokens; let mut tokens_generated = 0; while !finished.is_canceled() && tokens_generated < max_tokens { let new_token = next_token; - if new_token == stop_token { + self.trace_sampled_token(tokens_generated, new_token, stop_tokens); + if stop_tokens.contains(&new_token) { tracing::trace!("Stopping on stop token"); break; } @@ -239,20 +291,7 @@ where } .await?; - { - use std::sync::atomic::{AtomicBool, Ordering}; - let yielded = AtomicBool::new(false); - std::future::poll_fn(|cx| { - if yielded.load(Ordering::Relaxed) { - std::task::Poll::Ready(()) - } else { - yielded.store(true, Ordering::Relaxed); - cx.waker().wake_by_ref(); - std::task::Poll::Pending - } - }) - .await; - } + yield_once().await; } return Ok(()); @@ -282,14 +321,15 @@ where let mut queued_text_matching_stop_on = String::new(); let stop_on_lowercase = stop_on.as_ref().map(|s| s.to_lowercase()); let stop_on_lowercase = stop_on_lowercase.as_deref(); - let stop_token = self.model.config.stop_token; + let stop_tokens = &self.model.config.stop_tokens; let mut tokens_generated = 0; 'generate: while !finished.is_canceled() && tokens_generated < max_tokens { let new_token = text_stream .sample_token(&mut cpu_sampler, logits, stop_on.as_deref(), sample_top_k) .map_err(LlamaModelError::TokenOutputStreamError)?; - if new_token == stop_token { + self.trace_sampled_token(tokens_generated, new_token, stop_tokens); + if stop_tokens.contains(&new_token) { tracing::trace!("Stopping on stop token"); break; } @@ -369,20 +409,7 @@ where .await?; logits = logits_from_sorted_top_k(logit_probs); // Yield control to allow the stream to deliver tokens - { - use std::sync::atomic::{AtomicBool, Ordering}; - let yielded = AtomicBool::new(false); - std::future::poll_fn(|cx| { - if yielded.load(Ordering::Relaxed) { - std::task::Poll::Ready(()) - } else { - yielded.store(true, Ordering::Relaxed); - cx.waker().wake_by_ref(); - std::task::Poll::Pending - } - }) - .await; - } + yield_once().await; } // Flush the queued text @@ -394,4 +421,504 @@ where Ok(()) } + + #[allow(clippy::too_many_arguments)] + async fn infer_mtp_speculative( + &self, + prompt_tokens: &[u32], + images: &[LlamaImage], + mut text_stream: TokenOutputStream, + session: crate::LlamaSession, + max_tokens: u32, + mut gpu_sampler: LlamaGpuSamplerState, + mut on_token: crate::BoxedTokenCallback, + finished: &futures_channel::oneshot::Sender>, + ) -> Result<(), LlamaModelError> { + let mtp = self.mtp.as_ref().ok_or_else(|| { + LlamaModelError::SamplerError("Gemma4 MTP assistant is not loaded".into()) + })?; + let top_k = gpu_sample_top_k(&gpu_sampler.config); + let target = { + let mut session_lock = session + .cache + .write() + .map_err(|err| LlamaModelError::Session(err.to_string()))?; + self.model + .forward_logits_and_nextn_f32( + prompt_tokens, + images, + &self.device, + Some(&mut session_lock), + ) + .map_err(LlamaModelError::from)? + }; + let last_row = target.logits.shape()[0].saturating_sub(1); + let mut pending_h = Self::h_nextn_row(&target.h_nextn, last_row); + let history = Self::mtp_previous_tokens(&gpu_sampler.config, text_stream.tokens(), &[]); + let mut next_token = Self::sample_standard_logits_row( + target.logits, + last_row, + &mut gpu_sampler, + history, + top_k, + ) + .await?; + + let stop_tokens = &self.model.config.stop_tokens; + let mut tokens_generated = 0usize; + let mut drafted_total = 0usize; + let mut accepted_total = 0usize; + while !finished.is_canceled() && tokens_generated < max_tokens as usize { + let new_token = next_token; + self.trace_sampled_token(tokens_generated, new_token, stop_tokens); + if stop_tokens.contains(&new_token) { + tracing::trace!("Stopping on stop token"); + break; + } + tokens_generated += 1; + if let Some(new_text) = text_stream + .next_token(new_token) + .map_err(LlamaModelError::TokenOutputStreamError)? + { + on_token(new_text)?; + } + if finished.is_canceled() || tokens_generated >= max_tokens as usize { + break; + } + + let remaining = max_tokens as usize - tokens_generated; + let mut draft_limit = mtp_draft_limit(mtp.draft_n()).min(remaining); + if mtp_auto_fallback_enabled() && drafted_total < mtp_fallback_probe_drafts() { + draft_limit = + draft_limit.min(mtp_fallback_probe_drafts().saturating_sub(drafted_total)); + } + let mut draft_tokens = Vec::with_capacity(draft_limit); + let mut assistant_h = pending_h.clone(); + let mut assistant_token = new_token; + let draft_position = { + let session_lock = session + .cache + .read() + .map_err(|err| LlamaModelError::Session(err.to_string()))?; + session_lock.tokens.len() + }; + for _ in 0..draft_limit { + let draft_position = mtp_draft_position(draft_position, draft_tokens.len()); + let step = { + let session_lock = session + .cache + .read() + .map_err(|err| LlamaModelError::Session(err.to_string()))?; + mtp.draft_step( + &self.model, + assistant_token, + &assistant_h, + &session_lock, + &self.device, + draft_position, + ) + .map_err(LlamaModelError::from)? + }; + let history = Self::mtp_previous_tokens( + &gpu_sampler.config, + text_stream.tokens(), + &draft_tokens, + ); + assistant_token = + Self::sample_standard_logits(step.logits, &mut gpu_sampler, history, top_k) + .await?; + assistant_h = step.h_nextn; + draft_tokens.push(assistant_token); + } + drafted_total += draft_tokens.len(); + + let mut accepted = 0usize; + let mut verify_input = new_token; + let mut mismatched = false; + for draft in draft_tokens.iter().copied() { + let history = + Self::mtp_previous_tokens(&gpu_sampler.config, text_stream.tokens(), &[]); + let (verified, h_nextn) = self + .mtp_target_step(verify_input, &session, &mut gpu_sampler, history, top_k) + .await?; + pending_h = h_nextn; + if verified != draft { + next_token = verified; + mismatched = true; + break; + } + + accepted += 1; + if stop_tokens.contains(&draft) { + tracing::trace!("Stopping on accepted MTP stop token"); + if std::env::var_os("KALOSM_TRACE_MTP").is_some() { + tracing::info!( + "mtp_summary drafted={drafted_total} accepted={accepted_total}" + ); + } + return Ok(()); + } + + tokens_generated += 1; + if let Some(new_text) = text_stream + .next_token(draft) + .map_err(LlamaModelError::TokenOutputStreamError)? + { + on_token(new_text)?; + } + + verify_input = draft; + if tokens_generated >= max_tokens as usize || finished.is_canceled() { + break; + } + } + accepted_total += accepted; + + if !mismatched + && accepted == draft_tokens.len() + && tokens_generated < max_tokens as usize + && !finished.is_canceled() + { + let history = + Self::mtp_previous_tokens(&gpu_sampler.config, text_stream.tokens(), &[]); + let (bonus, h_nextn) = self + .mtp_target_step(verify_input, &session, &mut gpu_sampler, history, top_k) + .await?; + pending_h = h_nextn; + next_token = bonus; + } + + if Self::should_mtp_fallback(drafted_total, accepted_total) { + if std::env::var_os("KALOSM_TRACE_MTP").is_some() { + tracing::info!( + "mtp_auto_fallback drafted={drafted_total} accepted={accepted_total}" + ); + } + return self + .infer_target_only_from_pending( + next_token, + text_stream, + session, + max_tokens, + tokens_generated, + gpu_sampler, + on_token, + finished, + ) + .await; + } + + yield_once().await; + } + if std::env::var_os("KALOSM_TRACE_MTP").is_some() { + tracing::info!("mtp_summary drafted={drafted_total} accepted={accepted_total}"); + } + Ok(()) + } + + #[allow(clippy::too_many_arguments)] + async fn infer_target_only_from_pending( + &self, + mut next_token: u32, + mut text_stream: TokenOutputStream, + session: crate::LlamaSession, + max_tokens: u32, + mut tokens_generated: usize, + mut gpu_sampler: LlamaGpuSamplerState, + mut on_token: crate::BoxedTokenCallback, + finished: &futures_channel::oneshot::Sender>, + ) -> Result<(), LlamaModelError> { + let stop_tokens = &self.model.config.stop_tokens; + let top_k = gpu_sample_top_k(&gpu_sampler.config); + while !finished.is_canceled() && tokens_generated < max_tokens as usize { + let new_token = next_token; + self.trace_sampled_token(tokens_generated, new_token, stop_tokens); + if stop_tokens.contains(&new_token) { + tracing::trace!("Stopping on stop token"); + break; + } + tokens_generated += 1; + if let Some(new_text) = text_stream + .next_token(new_token) + .map_err(LlamaModelError::TokenOutputStreamError)? + { + on_token(new_text)?; + } + if finished.is_canceled() || tokens_generated >= max_tokens as usize { + break; + } + + let previous_tokens = gpu_sampler.previous_tokens(&text_stream); + let pending = { + let mut session_lock = session + .cache + .write() + .map_err(|err| LlamaModelError::Session(err.to_string()))?; + Self::forward_sample_token_pending( + ForwardInputs { + model: &self.model, + device: &self.device, + tokens: &[new_token], + images: &[], + cache: Some(&mut session_lock), + tokenizer: &self.tokenizer, + }, + &mut gpu_sampler, + previous_tokens.clone(), + top_k, + )? + }; + if let Some(next_pending) = pending { + return self + .infer_target_only_from_gpu_pending( + next_pending, + text_stream, + session, + max_tokens, + tokens_generated, + gpu_sampler, + on_token, + finished, + ) + .await; + } + + next_token = { + let mut session_lock = session + .cache + .write() + .map_err(|err| LlamaModelError::Session(err.to_string()))?; + Self::forward_sample_token( + ForwardInputs { + model: &self.model, + device: &self.device, + tokens: &[new_token], + images: &[], + cache: Some(&mut session_lock), + tokenizer: &self.tokenizer, + }, + &mut gpu_sampler, + previous_tokens, + top_k, + ) + } + .await?; + + yield_once().await; + } + Ok(()) + } + + #[allow(clippy::too_many_arguments)] + async fn infer_target_only_from_gpu_pending( + &self, + mut next_token: fusor::GpuSampledToken, + mut text_stream: TokenOutputStream, + session: crate::LlamaSession, + max_tokens: u32, + mut tokens_generated: usize, + mut gpu_sampler: LlamaGpuSamplerState, + mut on_token: crate::BoxedTokenCallback, + finished: &futures_channel::oneshot::Sender>, + ) -> Result<(), LlamaModelError> { + let stop_tokens = &self.model.config.stop_tokens; + let top_k = gpu_sample_top_k(&gpu_sampler.config); + while !finished.is_canceled() && tokens_generated < max_tokens as usize { + let scheduled_next = if tokens_generated + 1 < max_tokens as usize { + let previous_tokens = gpu_sampler.previous_tokens(&text_stream); + let mut speculative_cache = session + .cache + .read() + .map_err(|err| LlamaModelError::Session(err.to_string()))? + .clone(); + if speculative_cache.tokens.len() < self.model.config.context_length { + Self::forward_sample_token_from_gpu_token_pending( + &self.model, + &self.device, + &next_token, + &mut speculative_cache, + &mut gpu_sampler, + previous_tokens, + top_k, + )? + .map(|(next, cache_slot)| (next, speculative_cache, cache_slot)) + } else { + None + } + } else { + None + }; + + let new_token = next_token + .read_token() + .await + .map_err(|err| LlamaModelError::Fusor(fusor::Error::Gpu(err)))? + .ok_or_else(|| { + LlamaModelError::SamplerError( + "pending GPU sampler refused slow fallback".into(), + ) + })?; + self.trace_sampled_token(tokens_generated, new_token, stop_tokens); + if stop_tokens.contains(&new_token) { + tracing::trace!("Stopping on stop token"); + break; + } + + tokens_generated += 1; + if let Some(new_text) = text_stream + .next_token(new_token) + .map_err(LlamaModelError::TokenOutputStreamError)? + { + on_token(new_text)?; + } + + if let Some((scheduled_token, mut speculative_cache, cache_slot)) = scheduled_next { + if let Some(slot) = speculative_cache.tokens.get_mut(cache_slot) { + *slot = new_token; + } + *session + .cache + .write() + .map_err(|err| LlamaModelError::Session(err.to_string()))? = speculative_cache; + next_token = scheduled_token; + } else if !finished.is_canceled() && tokens_generated < max_tokens as usize { + let previous_tokens = gpu_sampler.previous_tokens(&text_stream); + let pending = { + let mut session_lock = session + .cache + .write() + .map_err(|err| LlamaModelError::Session(err.to_string()))?; + Self::forward_sample_token_pending( + ForwardInputs { + model: &self.model, + device: &self.device, + tokens: &[new_token], + images: &[], + cache: Some(&mut session_lock), + tokenizer: &self.tokenizer, + }, + &mut gpu_sampler, + previous_tokens, + top_k, + )? + }; + match pending { + Some(sampled) => next_token = sampled, + None => break, + } + } else { + break; + } + + yield_once().await; + } + Ok(()) + } + + async fn mtp_target_step( + &self, + token: u32, + session: &crate::LlamaSession, + sampler: &mut LlamaGpuSamplerState, + previous_tokens: Vec, + top_k: usize, + ) -> Result<(u32, fusor::Tensor<2, f32>), LlamaModelError> { + let target = { + let mut session_lock = session + .cache + .write() + .map_err(|err| LlamaModelError::Session(err.to_string()))?; + self.model + .forward_logits_and_nextn_f32(&[token], &[], &self.device, Some(&mut session_lock)) + .map_err(LlamaModelError::from)? + }; + let h_nextn = Self::h_nextn_row(&target.h_nextn, 0); + let token = + Self::sample_standard_logits_row(target.logits, 0, sampler, previous_tokens, top_k) + .await?; + Ok((token, h_nextn)) + } + + fn should_mtp_fallback(drafted_total: usize, accepted_total: usize) -> bool { + if !mtp_auto_fallback_enabled() { + return false; + } + if drafted_total < mtp_fallback_probe_drafts() { + return false; + } + accepted_total * 100 < drafted_total * mtp_fallback_min_accept_percent() + } + + fn h_nextn_row(h_nextn: &fusor::Tensor<2, f32>, row: usize) -> fusor::Tensor<2, f32> { + let row: fusor::Tensor<1, f32> = h_nextn.i((row, ..)).to_concrete(); + row.unsqueeze(0).to_concrete() + } + + fn mtp_previous_tokens( + config: &GpuSamplerConfig, + base_tokens: &[u32], + extra_tokens: &[u32], + ) -> Vec { + let range = config.repetition_penalty_range; + if range == 0 { + return Vec::new(); + } + let total_len = base_tokens.len() + extra_tokens.len(); + let keep_from = total_len.saturating_sub(range); + let mut result = Vec::with_capacity(range.min(total_len)); + if keep_from < base_tokens.len() { + result.extend_from_slice(&base_tokens[keep_from..]); + result.extend_from_slice(extra_tokens); + } else { + result.extend_from_slice(&extra_tokens[keep_from - base_tokens.len()..]); + } + result + } + + async fn sample_standard_logits( + logits: fusor::Tensor<1, f32>, + sampler: &mut LlamaGpuSamplerState, + previous_tokens: Vec, + top_k: usize, + ) -> Result { + let params = sampler.standard_params(top_k); + logits + .sample_standard_token(&previous_tokens, params) + .await + .map_err(LlamaModelError::from) + } + + async fn sample_standard_logits_row( + logits: fusor::Tensor<2, f32>, + row: usize, + sampler: &mut LlamaGpuSamplerState, + previous_tokens: Vec, + top_k: usize, + ) -> Result { + let row_logits: fusor::Tensor<1, f32> = logits.i((row, ..)).to_concrete(); + Self::sample_standard_logits(row_logits, sampler, previous_tokens, top_k).await + } +} + +fn mtp_draft_position(base_position: usize, draft_offset: usize) -> usize { + base_position + draft_offset +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn multi_token_mtp_draft_positions_advance_by_draft_offset() { + let base_position = 10; + let positions: Vec<_> = (0..3) + .map(|draft_offset| mtp_draft_position(base_position, draft_offset)) + .collect(); + + assert_eq!( + positions, + vec![10, 11, 12], + "each MTP draft token must be evaluated at its future position" + ); + } } diff --git a/models/kalosm-llama/src/model/mod.rs b/models/kalosm-llama/src/model/mod.rs index def415db7..ace2570b8 100644 --- a/models/kalosm-llama/src/model/mod.rs +++ b/models/kalosm-llama/src/model/mod.rs @@ -1,6 +1,6 @@ use crate::gguf_tokenizer::get_pre_tokenizer; use crate::raw::cache::LlamaCache; -use crate::raw::{LlamaVarSource, Model, RopeScalingConfig}; +use crate::raw::{Gemma4MtpAssistant, LlamaVarSource, Model, RopeScalingConfig}; use crate::token_stream::TokenOutputStream; use crate::token_stream::TokenOutputStreamError; use crate::tokenizer::{LlamaTokenizer, LlamaTokenizerError}; @@ -109,6 +109,42 @@ fn gpu_run_ahead_enabled() -> bool { .unwrap_or(true) } +fn mtp_speculative_enabled() -> bool { + std::env::var_os("KALOSM_LLAMA_MTP") + .map(|value| value != "0") + .unwrap_or(true) +} + +fn mtp_draft_limit(default: usize) -> usize { + std::env::var("KALOSM_LLAMA_MTP_DRAFT_N") + .ok() + .and_then(|value| value.parse::().ok()) + .unwrap_or(default) + .max(1) +} + +fn mtp_auto_fallback_enabled() -> bool { + std::env::var_os("KALOSM_LLAMA_MTP_AUTO_FALLBACK") + .map(|value| value != "0") + .unwrap_or(true) +} + +fn mtp_fallback_probe_drafts() -> usize { + std::env::var("KALOSM_LLAMA_MTP_FALLBACK_PROBE_DRAFTS") + .ok() + .and_then(|value| value.parse::().ok()) + .unwrap_or(1) + .max(1) +} + +fn mtp_fallback_min_accept_percent() -> usize { + std::env::var("KALOSM_LLAMA_MTP_FALLBACK_MIN_ACCEPT_PERCENT") + .ok() + .and_then(|value| value.parse::().ok()) + .unwrap_or(60) + .min(100) +} + fn decode_trace_enabled() -> bool { std::env::var_os("KALOSM_TRACE_DECODE_TIMING").is_some() || std::env::var_os("FUSOR_TRACE_DECODE").is_some() @@ -116,16 +152,26 @@ fn decode_trace_enabled() -> bool { } fn gpu_sample_top_k(config: &GpuSamplerConfig) -> usize { + if let Some(top_k) = std::env::var("KALOSM_LLAMA_GPU_SAMPLE_TOP_K") + .ok() + .and_then(|value| value.parse::().ok()) + { + return top_k.max(1); + } + // Greedy standard sampling (temperature <= 0) is argmax, so a top-1 cut is + // both exact and cheaper. This intentionally ignores any configured `top_k`, + // which has no effect once the distribution collapses to its max; the + // `KALOSM_LLAMA_GPU_SAMPLE_TOP_K` override above still wins if set. + if config.sampling_strategy == kalosm_language_model::SamplingStrategy::Standard + && config.temperature <= 0.0 + { + return 1; + } let default_top_k = match config.sampling_strategy { kalosm_language_model::SamplingStrategy::Mirostat2 => 16, kalosm_language_model::SamplingStrategy::Standard => 64, }; - std::env::var("KALOSM_LLAMA_GPU_SAMPLE_TOP_K") - .ok() - .and_then(|value| value.parse().ok()) - .or(config.top_k) - .unwrap_or(default_top_k) - .max(1) + config.top_k.unwrap_or(default_top_k).max(1) } fn parse_external_tokenizer( @@ -179,15 +225,17 @@ fn tokenizer_from_gguf_source( .clone() .try_into() .map_err(|_| LlamaSourceError::NoTokenizer)?; - if &*tokenizer_model != "gpt2" { + if !matches!(&*tokenizer_model, "gpt2" | "gemma4") { return Err(LlamaSourceError::NoTokenizer); } let pre: Box = source .get("tokenizer.ggml.pre") + .ok() + .cloned() + .map(|v| v.try_into()) + .transpose() .map_err(|_| LlamaSourceError::NoTokenizer)? - .clone() - .try_into() - .map_err(|_| LlamaSourceError::NoTokenizer)?; + .unwrap_or_else(|| tokenizer_model.clone()); let add_bos_token = source .get("tokenizer.ggml.add_bos_token") .ok() @@ -425,7 +473,35 @@ fn trim_previous_tokens_for_gpu_tail( #[cfg(test)] mod tests { - use super::trim_previous_tokens_for_gpu_tail; + use super::{gpu_sample_top_k, trim_previous_tokens_for_gpu_tail}; + use crate::GpuSamplerConfig; + use kalosm_language_model::SamplingStrategy; + use std::sync::{Mutex, OnceLock}; + + fn env_lock() -> &'static Mutex<()> { + static LOCK: OnceLock> = OnceLock::new(); + LOCK.get_or_init(|| Mutex::new(())) + } + + fn with_gpu_sample_top_k_env(value: Option<&str>, f: impl FnOnce() -> R) -> R { + let _guard = env_lock().lock().unwrap_or_else(|err| err.into_inner()); + let prior = std::env::var("KALOSM_LLAMA_GPU_SAMPLE_TOP_K").ok(); + // SAFETY: these tests serialize access to this process-wide env var. + unsafe { + match value { + Some(value) => std::env::set_var("KALOSM_LLAMA_GPU_SAMPLE_TOP_K", value), + None => std::env::remove_var("KALOSM_LLAMA_GPU_SAMPLE_TOP_K"), + } + } + let result = f(); + unsafe { + match prior { + Some(value) => std::env::set_var("KALOSM_LLAMA_GPU_SAMPLE_TOP_K", value), + None => std::env::remove_var("KALOSM_LLAMA_GPU_SAMPLE_TOP_K"), + } + } + result + } #[test] fn gpu_tail_history_preserves_repetition_window() { @@ -446,6 +522,31 @@ mod tests { (vec![1, 2], true) ); } + + #[test] + fn greedy_standard_gpu_sampling_uses_top_one() { + let mut config = GpuSamplerConfig::new(0.0, 5.0, 0.1, 10.0, 1.0, 64, Some(64)); + config.sampling_strategy = SamplingStrategy::Standard; + with_gpu_sample_top_k_env(None, || { + assert_eq!(gpu_sample_top_k(&config), 1); + }); + + config.temperature = 0.8; + with_gpu_sample_top_k_env(None, || { + assert_eq!(gpu_sample_top_k(&config), 64); + }); + with_gpu_sample_top_k_env(Some("8"), || { + assert_eq!(gpu_sample_top_k(&config), 8); + }); + } + + #[test] + fn greedy_top_k_keeps_mirostat_default() { + let config = GpuSamplerConfig::new(0.0, 5.0, 0.1, 10.0, 1.0, 64, None); + with_gpu_sample_top_k_env(None, || { + assert_eq!(gpu_sample_top_k(&config), 16); + }); + } } struct ForwardTrace { @@ -610,6 +711,7 @@ impl From for LlamaModelError { /// The inner, synchronous Llama model. pub(crate) struct LlamaModel { pub(crate) model: Model, + pub(crate) mtp: Option>, pub(crate) device: Device, pub(crate) tokenizer: Arc, } @@ -658,8 +760,8 @@ where random: 0.5, }; - let _ = hidden - .q_mat_mul(model.output_matrix()) + let _ = model + .logits_from_hidden_f32(hidden) .sample_mirostat2_token(&mut sampler, &[], params) .await; if let Some(start) = start { @@ -755,6 +857,29 @@ where None }; + #[cfg(not(target_arch = "wasm32"))] + let mtp_model_bytes = match &builder.source.mtp_model { + Some(mtp_model) => { + let mtp_model_source = format!("MTP Model ({mtp_model})"); + let mut create_progress = + ModelLoadingProgress::downloading_progress(mtp_model_source); + let mtp_model_bytes = builder + .source + .cache + .get_bytes(mtp_model, |progress| handler(create_progress(progress))) + .await?; + Some(mtp_model_bytes) + } + None => None, + }; + #[cfg(target_arch = "wasm32")] + let mtp_model_bytes: Option> = { + if builder.source.mtp_model.is_some() { + tracing::warn!("Gemma4 MTP speculative decoding is not loaded on wasm targets"); + } + None + }; + let source = format!("Model ({})", builder.source.model[0]); let mut create_progress = ModelLoadingProgress::downloading_progress(source); #[cfg(target_arch = "wasm32")] @@ -784,7 +909,7 @@ where let override_chat_template = builder.source.override_chat_template.clone(); #[cfg(target_arch = "wasm32")] - let (model, tokenizer) = { + let (model, tokenizer, mtp) = { let files_with_metadata = match model_source { WasmModelSource::Opfs(model_files) => { if model_files.is_empty() { @@ -840,13 +965,15 @@ where ) .await?; - (model, tokenizer) + (model, tokenizer, None) }; #[cfg(not(target_arch = "wasm32"))] - let (model, tokenizer) = { + let (model, tokenizer, mtp) = { let device = device.clone(); - let load_model = move || -> Result<(Model, LlamaTokenizer), LlamaSourceError> { + #[allow(clippy::type_complexity)] + let load_model = + move || -> Result<(Model, LlamaTokenizer, Option>), LlamaSourceError> { let tokenizer = parse_external_tokenizer(tokenizer_source)?; let config = parse_external_config(config_bytes)?; if model_bytes.is_empty() { @@ -885,7 +1012,16 @@ where override_chat_template, config, )?; - Ok((model, tokenizer)) + let mtp = match mtp_model_bytes { + Some(bytes) => { + let mut cursor = std::io::Cursor::new(bytes); + let metadata = GgufMetadata::read(&mut cursor)?; + let mut mtp_source = ShardedVarBuilder::new(vec![(metadata, cursor)]); + Some(Gemma4MtpAssistant::from_gguf(&mut mtp_source, &device)?) + } + None => None, + }; + Ok((model, tokenizer, mtp)) }; load_model()? @@ -894,6 +1030,7 @@ where Ok(Self { model, + mtp, tokenizer: Arc::new(tokenizer), device, }) diff --git a/models/kalosm-llama/src/raw/attention_layer.rs b/models/kalosm-llama/src/raw/attention_layer.rs index 104a3721b..0a1ad1ade 100644 --- a/models/kalosm-llama/src/raw/attention_layer.rs +++ b/models/kalosm-llama/src/raw/attention_layer.rs @@ -16,6 +16,12 @@ pub enum FeedForwardVariant { Phi(PhiFeedForward), } +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub enum FeedForwardActivation { + Silu, + Gelu, +} + impl FeedForwardVariant where F: CastTo + CastTensor, @@ -83,6 +89,7 @@ impl PhiFeedForward { } pub struct LlamaFeedForward { + activation: FeedForwardActivation, gate: QMatrix, gate_up: Option, gate_bias: Option>, @@ -93,9 +100,15 @@ pub struct LlamaFeedForward { } impl LlamaFeedForward { - pub(crate) fn new(gate: QMatrix, down: QMatrix, up: QMatrix) -> Self { + pub(crate) fn new_with_activation( + gate: QMatrix, + down: QMatrix, + up: QMatrix, + activation: FeedForwardActivation, + ) -> Self { let gate_up = QMatrix::concat_rows(&[&gate, &up]); Self { + activation, gate, gate_up, down, @@ -117,6 +130,7 @@ impl LlamaFeedForward { ) -> Self { let gate_up = QMatrix::concat_rows(&[&gate, &up]); Self { + activation: FeedForwardActivation::Silu, gate, gate_up, gate_bias, @@ -192,7 +206,7 @@ impl LlamaFeedForward { let up = projected .narrow(D::Minus1, pair_len, pair_len) .to_concrete(); - (gate.silu() * up).to_concrete() + (self.activate(gate) * up).to_concrete() } Some(gate_up) => { let gate_width = self.gate.shape()[0]; @@ -215,7 +229,7 @@ impl LlamaFeedForward { up_states = up_states.add_(&bias_f32); } - (gate_states.silu() * up_states).to_concrete() + (self.activate(gate_states) * up_states).to_concrete() } None => { let mut w1 = x_f32.q_mat_mul(&self.gate); @@ -223,7 +237,7 @@ impl LlamaFeedForward { let bias_f32: Tensor<1, f32> = bias.cast(); w1 = w1.add_(&bias_f32); } - let w1 = w1.silu(); + let w1 = self.activate(w1); let mut w3 = x_f32.q_mat_mul(&self.up); if let Some(ref bias) = self.up_bias { @@ -235,6 +249,13 @@ impl LlamaFeedForward { } } } + + fn activate(&self, x: Tensor<3, f32>) -> Tensor<3, f32> { + match self.activation { + FeedForwardActivation::Silu => x.silu(), + FeedForwardActivation::Gelu => x.gelu(), + } + } } pub enum AttentionVariant { @@ -265,9 +286,10 @@ pub struct SeparateAttention { pub attention_wq: QMatrix, pub attention_qkv: Option, pub attention_q_norm: Option>, - pub attention_wk: QMatrix, + pub attention_wk: Option, pub attention_k_norm: Option>, - pub attention_wv: QMatrix, + pub attention_wv: Option, + pub attention_v_norm: Option>, pub bias: Option>, pub interleaved_rope: bool, } @@ -343,11 +365,16 @@ where .narrow(D::Minus1, query_width + key_width, value_width) .to_concrete(); - value_states + let value_states: Tensor<4, F> = value_states .reshape([b_sz, seq_len, num_key_value_heads, head_dim]) .transpose(1, 2) .to_concrete() - .cast() + .cast(); + if let Some(norm) = &self.attention_v_norm { + norm.forward_generic_4d(&value_states) + } else { + value_states + } }; let (query_states, key_states) = rope_cache.forward( @@ -381,7 +408,11 @@ where } }; let key_states: Tensor<4, F> = { - let mut key_states = hidden_f32.q_mat_mul(&self.attention_wk); + let attention_wk = self + .attention_wk + .as_ref() + .expect("separate attention without K weights must use a shared KV cache"); + let mut key_states = hidden_f32.q_mat_mul(attention_wk); if let Some(bias) = &self.bias { let bias_f32: Tensor<1, f32> = bias.bias_k.cast(); @@ -401,18 +432,27 @@ where } }; let value_states: Tensor<4, F> = { - let mut value_states = hidden_f32.q_mat_mul(&self.attention_wv); + let attention_wv = self + .attention_wv + .as_ref() + .expect("separate attention without V weights must use a shared KV cache"); + let mut value_states = hidden_f32.q_mat_mul(attention_wv); if let Some(bias) = &self.bias { let bias_f32: Tensor<1, f32> = bias.bias_v.cast(); value_states = value_states.add_(&bias_f32); } - value_states + let value_states: Tensor<4, F> = value_states .reshape([b_sz, seq_len, num_key_value_heads, head_dim]) .transpose(1, 2) .to_concrete() - .cast() + .cast(); + if let Some(norm) = &self.attention_v_norm { + norm.forward_generic_4d(&value_states) + } else { + value_states + } }; let (query_states, key_states) = rope_cache.forward( @@ -424,6 +464,73 @@ where ); (query_states, key_states, value_states) } + + #[allow(clippy::too_many_arguments)] + fn forward_query( + &self, + num_heads: usize, + head_dim: usize, + hidden_states: &Tensor<3, F, B>, + rope_cache: &RopeImplementation, + start_pos: usize, + pos_ids: Option<&Tensor<2, F>>, + ) -> Tensor<4, F> + where + B: Fusion<3, F>, + { + let [b_sz, seq_len, _] = hidden_states.shape(); + let hidden_f32 = hidden_states.cast::(); + + let query_states: Tensor<4, F> = if let Some(attention_qkv) = &self.attention_qkv { + // Shared-KV callers only need Q, but a fused QKV weight forces us to + // project K/V as well and discard them. No current shared-KV model + // (Gemma 4 or the MTP assistant) uses a fused QKV weight, so this + // branch is effectively unreachable today; if a future one does, it + // pays a ~3x projection here and should grow a Q-only weight slice. + let query_width = num_heads * head_dim; + let mut qkv = hidden_f32.q_mat_mul(attention_qkv); + if let Some(bias) = &self.bias { + let bias_f32: Tensor<1, f32> = bias.bias_qkv.cast(); + qkv = qkv.add_(&bias_f32); + } + let query_states = qkv.narrow(D::Minus1, 0, query_width).to_concrete(); + let query = query_states + .reshape([b_sz, seq_len, num_heads, head_dim]) + .transpose(1, 2) + .to_concrete(); + let query: Tensor<4, F> = query.cast(); + if let Some(norm) = &self.attention_q_norm { + norm.forward_generic_4d(&query) + } else { + query + } + } else { + let mut query_states = hidden_f32.q_mat_mul(&self.attention_wq); + if let Some(bias) = &self.bias { + let bias_f32: Tensor<1, f32> = bias.bias_q.cast(); + query_states = query_states.add_(&bias_f32); + } + let query = query_states + .reshape([b_sz, seq_len, num_heads, head_dim]) + .transpose(1, 2) + .to_concrete(); + let query: Tensor<4, F> = query.cast(); + if let Some(norm) = &self.attention_q_norm { + norm.forward_generic_4d(&query) + } else { + query + } + }; + + let (query_states, _) = rope_cache.forward( + &query_states, + &query_states, + start_pos, + pos_ids, + self.interleaved_rope, + ); + query_states + } } pub struct GroupedAttention { @@ -504,6 +611,12 @@ pub struct LlamaAttention { pub hidden_size: usize, pub rope_cache: RopeImplementation, pub(crate) sliding_window_size: Option, + pub(crate) attention_scale: f32, + pub(crate) shared_kv_layer: Option, + pub(crate) per_layer_inp_gate: Option, + pub(crate) per_layer_proj: Option, + pub(crate) per_layer_post_norm: Option>, + pub(crate) layer_output_scale: Option>, } impl LlamaAttention @@ -511,6 +624,41 @@ where F: CastTo + CastTensor, f32: CastTo + CastTensor, { + fn logical_kv_len(&self, start_pos: usize, q_len: usize) -> usize { + let len = start_pos + q_len; + self.sliding_window_size + .map(|window| len.min(window)) + .unwrap_or(len) + } + + fn narrow_kv_to_logical_len( + &self, + key_states: Tensor<4, f32>, + value_states: Tensor<4, f32>, + logical_kv_len: usize, + ) -> (Tensor<4, f32>, Tensor<4, f32>) { + if key_states.shape()[2] <= logical_kv_len { + return (key_states, value_states); + } + + ( + key_states + .narrow( + 2, + kv_narrow_start(key_states.shape()[2], logical_kv_len), + logical_kv_len, + ) + .to_concrete(), + value_states + .narrow( + 2, + kv_narrow_start(value_states.shape()[2], logical_kv_len), + logical_kv_len, + ) + .to_concrete(), + ) + } + pub(crate) fn forward( &self, hidden_states: &Tensor<3, F, B>, @@ -558,6 +706,11 @@ where None => (key_f32, value_f32), Some(cache) => cache.append(&query_f32.device(), &key_f32, &value_f32), }; + let (key_f32, value_f32) = self.narrow_kv_to_logical_len( + key_f32, + value_f32, + self.logical_kv_len(start_pos, q_len), + ); forward_attention_qkv_f32( &query_f32, @@ -565,10 +718,56 @@ where &value_f32, &self.attention_wo, attention_mask, - head_dim, b_sz, q_len, hidden_size, + self.attention_scale, + ) + } + + pub(crate) fn forward_with_shared_kv( + &self, + hidden_states: &Tensor<3, F, B>, + attention_mask: Option<&AttentionMask>, + start_pos: usize, + pos_ids: Option<&Tensor<2, F>>, + key_states: &Tensor<4, f32>, + value_states: &Tensor<4, f32>, + ) -> Tensor<3, F> + where + B: Fusion<3, F>, + { + let [b_sz, q_len, _] = hidden_states.shape(); + let query_states = match self.attention_variant { + AttentionVariant::Separate(ref attention) => attention.forward_query( + self.n_head, + self.head_dim, + hidden_states, + &self.rope_cache, + start_pos, + pos_ids, + ), + AttentionVariant::Grouped(_) => { + panic!("grouped attention cannot reuse a shared KV cache") + } + }; + let query_f32: Tensor<4, f32> = query_states.cast(); + let (key_states, value_states) = self.narrow_kv_to_logical_len( + key_states.clone(), + value_states.clone(), + self.logical_kv_len(start_pos, q_len), + ); + + forward_attention_qkv_f32( + &query_f32, + &key_states, + &value_states, + &self.attention_wo, + attention_mask, + b_sz, + q_len, + self.hidden_size, + self.attention_scale, ) } @@ -624,15 +823,23 @@ where None => (key_f32, value_f32), Some(cache) => cache.append(&query_f32.device(), &key_f32, &value_f32), }; + let (key_f32, value_f32) = self.narrow_kv_to_logical_len( + key_f32, + value_f32, + self.logical_kv_len(start_pos, q_len), + ); crate::raw::debug_check_nan_f32(&key_f32, layer_idx, "K_cache_view", start_pos); crate::raw::debug_check_nan_f32(&value_f32, layer_idx, "V_cache_view", start_pos); - let scale = 1. / (head_dim as f64).sqrt(); + let scale = self.attention_scale; + let padded_attention_mask = + attention_mask.and_then(|m| pad_attention_mask_to_kv_len(m, key_f32.shape()[2])); + let attention_mask = padded_attention_mask.as_ref().or(attention_mask); let attn_raw = query_f32.flash_attention( &key_f32, &value_f32, - scale as f32, + scale, attention_mask.map(|m| { let kind = if m.is_strict_causal() { fusor::MaskKind::Causal @@ -653,6 +860,33 @@ where } } +fn kv_narrow_start(kv_seq_len: usize, logical_kv_len: usize) -> usize { + kv_seq_len.saturating_sub(logical_kv_len) +} + +fn pad_attention_mask_to_kv_len( + attention_mask: &AttentionMask, + kv_seq_len: usize, +) -> Option> { + let mask = attention_mask.mask(); + let [rows, cols] = mask.shape(); + if cols == kv_seq_len { + return None; + } + + if cols > kv_seq_len { + let start_col = cols - kv_seq_len; + return Some(AttentionMask::new( + mask.narrow(1, start_col, kv_seq_len).to_concrete(), + )); + } + + let padded = Tensor::full(&mask.device(), [rows, kv_seq_len], f32::NEG_INFINITY); + Some(AttentionMask::new( + padded.slice_assign([0..rows, 0..cols], mask), + )) +} + /// Forward attention QKV computation in f32 for SIMD compatibility. /// All intermediate computation happens in f32, with the final result cast back to F. #[allow(clippy::too_many_arguments)] @@ -662,20 +896,22 @@ pub(crate) fn forward_attention_qkv_f32( value_states: &Tensor<4, f32>, attention_wo: &Linear, attention_mask: Option<&AttentionMask>, - head_dim: usize, b_sz: usize, q_len: usize, hidden_size: usize, + attention_scale: f32, ) -> Tensor<3, F> where F: FloatDataType + SimdElement + Default + CastTo + CastTensor, f32: CastTo + CastTensor, { - let scale = 1. / (head_dim as f64).sqrt(); + let padded_attention_mask = + attention_mask.and_then(|m| pad_attention_mask_to_kv_len(m, key_states.shape()[2])); + let attention_mask = padded_attention_mask.as_ref().or(attention_mask); let attn_output = query_states.flash_attention( key_states, value_states, - scale as f32, + attention_scale, attention_mask.map(|m| { let kind = if m.is_strict_causal() { fusor::MaskKind::Causal @@ -692,3 +928,42 @@ where attention_wo.forward_generic(&attn_output.cast()) } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn sliding_window_prefill_kv_slice_keeps_newest_entries() { + let kv_seq_len = 8; + let logical_kv_len = 4; + + assert_eq!( + kv_narrow_start(kv_seq_len, logical_kv_len), + kv_seq_len - logical_kv_len, + "sliding-window prefill must keep the newest KV entries" + ); + } + + #[test] + fn sliding_window_prefill_mask_crop_targets_newest_columns() { + let device = fusor::Device::Cpu; + let mask = AttentionMask::new( + Tensor::from_slice( + &device, + [1, 8], + &[-7.0f32, -6.0, -5.0, -4.0, -3.0, -2.0, -1.0, 0.0], + ) + .to_concrete(), + ); + + let cropped = pad_attention_mask_to_kv_len(&mask, 4).expect("mask should be cropped"); + let values = pollster::block_on(cropped.mask().as_slice()).unwrap(); + + assert_eq!( + values.as_slice(), + &[-3.0, -2.0, -1.0, 0.0], + "attention mask cropping keeps the newest columns, so KV slicing must match it" + ); + } +} diff --git a/models/kalosm-llama/src/raw/cache.rs b/models/kalosm-llama/src/raw/cache.rs index dd26510d6..94627496f 100644 --- a/models/kalosm-llama/src/raw/cache.rs +++ b/models/kalosm-llama/src/raw/cache.rs @@ -24,7 +24,13 @@ impl LlamaCache { let max_seq_len = config.context_length; let mut blocks = Vec::with_capacity(config.n_layer); for layer_idx in 0..config.n_layer { - let max_seq_len = if let (Some(sliding_window_type), Some(sliding_window_size)) = + let max_seq_len = if let Some(layer_window_size) = config + .layer_sliding_window_sizes + .as_ref() + .and_then(|sizes| sizes.get(layer_idx).copied().flatten()) + { + layer_window_size + } else if let (Some(sliding_window_type), Some(sliding_window_size)) = (config.sliding_window_type, config.sliding_window_size) { let is_sliding = (layer_idx + 1) % sliding_window_type != 0; diff --git a/models/kalosm-llama/src/raw/mod.rs b/models/kalosm-llama/src/raw/mod.rs index c2d8bbfe3..db0f88c02 100644 --- a/models/kalosm-llama/src/raw/mod.rs +++ b/models/kalosm-llama/src/raw/mod.rs @@ -1,4 +1,6 @@ +use std::collections::HashMap; use std::future::Future; +use std::ops::Range; use std::pin::Pin; use std::sync::Arc; #[cfg(not(target_arch = "wasm32"))] @@ -57,6 +59,68 @@ pub(crate) fn debug_check_nan_f32( } } +#[cfg(feature = "vision")] +pub(crate) fn debug_tensor_stats_f32(t: &fusor::Tensor, label: &str) { + #[cfg(target_arch = "wasm32")] + { + let _ = (t, label); + return; + } + #[cfg(not(target_arch = "wasm32"))] + { + if std::env::var_os("KALOSM_TRACE_VISION_STATS").is_none() { + return; + } + let Ok(slice) = pollster::block_on(t.as_slice()) else { + tracing::warn!( + "[vision_stats] {label} shape={:?} readback failed", + t.shape() + ); + return; + }; + let values = slice.as_slice(); + let mut nan = 0usize; + let mut pos_inf = 0usize; + let mut neg_inf = 0usize; + let mut min = f32::INFINITY; + let mut max = f32::NEG_INFINITY; + let mut sum = 0.0f64; + for value in values.iter().copied() { + if value.is_nan() { + nan += 1; + } else if value == f32::INFINITY { + pos_inf += 1; + } else if value == f32::NEG_INFINITY { + neg_inf += 1; + } else { + min = min.min(value); + max = max.max(value); + sum += value as f64; + } + } + let finite = values.len().saturating_sub(nan + pos_inf + neg_inf); + let mean = if finite == 0 { + 0.0 + } else { + (sum / finite as f64) as f32 + }; + let sample = values.iter().copied().take(8).collect::>(); + tracing::info!( + "[vision_stats] {label} shape={:?} len={} finite={} nan={} +inf={} -inf={} min={} max={} mean={} sample={:?}", + t.shape(), + values.len(), + finite, + nan, + pos_inf, + neg_inf, + min, + max, + mean, + sample + ); + } +} + #[cfg(not(feature = "vision"))] pub(crate) fn debug_check_nan_f32( _: &fusor::Tensor, @@ -66,25 +130,65 @@ pub(crate) fn debug_check_nan_f32( ) { } +#[cfg(not(target_arch = "wasm32"))] +fn resolve_intermediate_hidden_f32(tensor: &fusor::Tensor<2, f32>) { + let marker = tensor.clone().mul_scalar(1.0).to_concrete(); + if let Some(gpu_marker) = marker.as_gpu() { + gpu_marker.materialize_sync(); + } else { + std::mem::drop(marker.to_concrete()); + } +} + +#[cfg(target_arch = "wasm32")] +fn resolve_intermediate_hidden_f32(_: &fusor::Tensor<2, f32>) {} + +#[cfg(all(feature = "vision", not(target_arch = "wasm32")))] +fn copy_image_embeddings_to_device( + tensor: Tensor<2, f32>, + device: &Device, +) -> Result> { + let shape = tensor.shape(); + let expected_len = shape.iter().product::(); + let values = pollster::block_on(tensor.as_slice())?; + let values = values.as_slice(); + if values.len() != expected_len { + return Err(fusor::Error::msg(format!( + "Image embedding transfer expected {expected_len} values for shape {shape:?}, got {}", + values.len() + ))); + } + Ok(Tensor::from_slice(device, shape, values)) +} + +#[cfg(all(feature = "vision", target_arch = "wasm32"))] +fn copy_image_embeddings_to_device( + tensor: Tensor<2, f32>, + _device: &Device, +) -> Result> { + Ok(tensor) +} + use crate::chat_template::HuggingFaceChatTemplate; use crate::raw::attention_layer::LlamaAttention; use crate::raw::rope::RopeImplementation; use crate::LlamaSourceError; use attention_layer::AttentionBias; use attention_layer::AttentionVariant; +use attention_layer::FeedForwardActivation; use attention_layer::FeedForwardVariant; use attention_layer::GroupedAttention; use attention_layer::LlamaFeedForward; use attention_layer::PhiFeedForward; use attention_layer::SeparateAttention; -use fusor::cache::MaskCache; +use fusor::cache::{AttentionMask, MaskCache}; use fusor::layers::Embedding; use fusor::layers::Linear; use fusor::layers::RmsNorm; use fusor::QMatrix; use fusor::ShardedVarBuilder; use fusor::{ - AddOp, CastTensor, CastTo, FloatDataType, FloatOps, MatmulImpl, MulOp, SimdBinaryOp, + AddOp, CastTensor, CastTo, FloatDataType, FloatOps, Fusion, MatmulImpl, MulOp, SimdBinaryOp, SimdElement, SimdReduceOp, SumOp, }; use fusor::{AsyncReadRange, AsyncShardedVarBuilder}; @@ -94,47 +198,85 @@ use fusor_gguf::GgufValue; mod attention_layer; pub mod cache; +mod mtp; mod rope; #[cfg(feature = "vision")] mod vision; use crate::LlamaImage; use cache::LlamaCache; +pub(crate) use mtp::Gemma4MtpAssistant; pub const DEFAULT_ROPE_FREQUENCY: f32 = 1_000_000.; pub const GEMMA_DEFAULT_SLIDING_WINDOW_TYPE: usize = 6; pub const GEMMA_DEFAULT_ROPE_FREQUENCY_SLIDING: f32 = 10_000.; +/// Build the additive attention-mask values (`0.0` allowed, `-inf` blocked) +/// for a `[seq_len, index_pos + seq_len]` score matrix. +/// +/// Tokens are causal by default. Any query/key position that falls inside the +/// same entry of `non_causal_token_ranges` may attend to its peers regardless +/// of ordering (this is how image-token blocks attend bidirectionally), and an +/// optional sliding `window` blocks keys older than `window` positions. +fn non_causal_mask_data( + seq_len: usize, + index_pos: usize, + sliding_window_size: Option, + non_causal_token_ranges: &[Range], +) -> Vec { + let cols = index_pos + seq_len; + let mut mask_data = vec![0.0_f32; seq_len * cols]; + for row in 0..seq_len { + let global_row = index_pos + row; + for col in 0..cols { + let same_non_causal_range = col >= index_pos + && non_causal_token_ranges + .iter() + .any(|range| range.contains(&row) && range.contains(&(col - index_pos))); + let future = col > global_row && !same_non_causal_range; + let outside_window = sliding_window_size + .map(|window| col + window <= global_row) + .unwrap_or(false); + if future || outside_window { + mask_data[row * cols + col] = f32::NEG_INFINITY; + } + } + } + mask_data +} + /// The configuration of a Llama model. pub struct LlamaConfig { pub(crate) rope_freq_weight: Option>, pub(crate) rope_theta: f32, pub(crate) context_length: usize, pub(crate) head_dimension: usize, - n_head: usize, pub(crate) n_layer: usize, pub(crate) start_token_string: String, - pub(crate) stop_token: u32, + pub(crate) stop_tokens: Vec, pub(crate) stop_token_string: String, pub(crate) chat_template: Option, pub(crate) rope_scaling: Option, pub(crate) sliding_window_type: Option, pub(crate) sliding_window_size: Option, + pub(crate) layer_sliding_window_sizes: Option>>, + pub(crate) final_logit_softcapping: Option, + pub(crate) per_layer_embedding_length: Option, #[cfg_attr(not(feature = "vision"), allow(dead_code))] pub(crate) vision_start_token: Option, pub(crate) _vision_end_token: Option, #[cfg_attr(not(feature = "vision"), allow(dead_code))] pub(crate) image_pad_token: Option, #[cfg_attr(not(feature = "vision"), allow(dead_code))] + pub(crate) image_start_token: Option, + #[cfg_attr(not(feature = "vision"), allow(dead_code))] + pub(crate) image_end_token: Option, + #[cfg_attr(not(feature = "vision"), allow(dead_code))] pub(crate) video_pad_token: Option, pub(crate) mrope_sections: Option>, } impl LlamaConfig { - fn hidden_size(&self) -> usize { - self.head_dimension * self.n_head - } - #[cfg(test)] pub(crate) fn mock_test() -> Self { Self { @@ -142,18 +284,22 @@ impl LlamaConfig { rope_theta: 5000., context_length: 6, head_dimension: 2, - n_head: 0, n_layer: 0, start_token_string: "<|startoftext|>".to_string(), - stop_token: 0, + stop_tokens: vec![0], stop_token_string: "<|endoftext|>".to_string(), sliding_window_type: None, sliding_window_size: None, + layer_sliding_window_sizes: None, + final_logit_softcapping: None, + per_layer_embedding_length: None, chat_template: None, rope_scaling: None, vision_start_token: None, _vision_end_token: None, image_pad_token: None, + image_start_token: None, + image_end_token: None, video_pad_token: None, mrope_sections: None, } @@ -171,9 +317,12 @@ pub struct RopeScalingConfig { pub struct Model { pub(crate) config: Arc>, #[cfg(feature = "vision")] - vision_encoder: Option>, + vision_encoder: Option>, tok_embeddings: Embedding, tok_embedding_scale: Option, + per_layer_tok_embeddings: Option>, + per_layer_model_proj: Option, + per_layer_proj_norm: Option>, layers: Vec>, norm: RmsNorm<1, F>, output: QMatrix, @@ -181,13 +330,25 @@ pub struct Model { masks: MaskCache, } +pub(crate) struct TargetBatchOutput { + pub(crate) logits: Tensor<2, f32>, + pub(crate) h_nextn: Tensor<2, f32>, +} + +struct PreNormForwardOutput { + hidden: Tensor<3, F>, + seq_len: usize, +} + /// The embedded token inputs produced by [`Model::encode_tokens`], ready to be /// run through the transformer layers. pub(crate) struct EncodedTokens { embeddings: Tensor<3, F>, + per_layer_inputs: Option>, seq_len: usize, index_pos: usize, pos_ids: Option>, + non_causal_token_ranges: Vec>, } pub(crate) trait LlamaVarSource { @@ -342,17 +503,22 @@ where .get("tokenizer.ggml.bos_token_id") .ok() .and_then(|v| v.try_into().ok()); + let eos_token: u32 = source + .get("tokenizer.ggml.eos_token_id")? + .clone() + .try_into()?; let stop_token = if let Some(override_stop_token_string) = override_stop_token_string { tokens .iter() .position(|v| **v == override_stop_token_string) .unwrap_or(0) as u32 } else { - source - .get("tokenizer.ggml.eos_token_id")? - .clone() - .try_into()? + eos_token }; + let mut stop_tokens = vec![stop_token]; + if eos_token != stop_token { + stop_tokens.push(eos_token); + } let start_token_string = start_token .map(|v| tokens[v as usize].to_string()) .unwrap_or_default(); @@ -375,12 +541,22 @@ where // Parameter extraction from metadata. let architecture = source.get("general.architecture")?.to_string()?.clone(); + let is_gemma4 = architecture.as_ref() == "gemma4"; let head_count = source.get(".attention.head_count")?.to_u32()? as usize; let head_count_kv = source.get(".attention.head_count_kv")?.to_u32()? as usize; let block_count = source.get(".block_count")?.to_u32()? as usize; let embedding_length = source.get(".embedding_length")?.to_u32()? as usize; + let per_layer_embedding_length = source + .get(".embedding_length_per_layer_input") + .and_then(|m| Ok(m.to_u32()?)) + .ok() + .map(|x| x as usize); // Strangely this value is generally 1e-6 in GGUF file but used to be 1e-5 by default. let rms_norm_eps = source.get(".attention.layer_norm_rms_epsilon")?.to_f32()? as f64; + let final_logit_softcapping = source + .get(".final_logit_softcapping") + .and_then(|m| Ok(m.to_f32()?)) + .ok(); let rope_freq_base = source .get(".rope.freq_base") @@ -399,20 +575,68 @@ where .or_else(|| (&*architecture == "gemma3").then_some(GEMMA_DEFAULT_SLIDING_WINDOW_TYPE)); let rope_freq_base_sliding = source - .get(".rope.local_freq_base") + .get(".rope.freq_base_swa") .and_then(|m| Ok(m.to_f32()?)) .ok() .or_else(|| { - (&*architecture == "gemma3").then_some(GEMMA_DEFAULT_ROPE_FREQUENCY_SLIDING) + source + .get(".rope.local_freq_base") + .and_then(|m| Ok(m.to_f32()?)) + .ok() + }) + .or_else(|| { + (&*architecture == "gemma3" || is_gemma4) + .then_some(GEMMA_DEFAULT_ROPE_FREQUENCY_SLIDING) }); - let context_length = source.get(".context_length")?.to_u32()? as usize; + let sliding_window_pattern = source + .get(".attention.sliding_window_pattern") + .ok() + .and_then(|m| { + let values = m.to_array().ok()?; + values + .iter() + .map(|value| value.to_bool().ok()) + .collect::>>() + }); + let layer_is_sliding: Vec = if let Some(pattern) = sliding_window_pattern { + pattern + } else if let Some(sliding_window_type) = sliding_window_type { + (0..block_count) + .map(|layer_idx| (layer_idx + 1) % sliding_window_type != 0) + .collect() + } else { + vec![false; block_count] + }; + let layer_sliding_window_sizes = sliding_window_size.map(|sliding_window_size| { + layer_is_sliding + .iter() + .map(|is_sliding| is_sliding.then_some(sliding_window_size)) + .collect::>() + }); + + let shared_kv_layers = source + .get(".attention.shared_kv_layers") + .and_then(|m| Ok(m.to_u32()?)) + .ok() + .map(|x| x as usize) + .unwrap_or_default(); + let n_layer_kv_from_start = block_count.saturating_sub(shared_kv_layers); + let head_dim = source .get(".attention.key_length") .and_then(|v| Ok(v.to_u32()?)) .ok() .map(|x| x as usize) .unwrap_or_else(|| embedding_length / head_count); + let head_dim_swa = source + .get(".attention.key_length_swa") + .and_then(|v| Ok(v.to_u32()?)) + .ok() + .map(|x| x as usize) + .unwrap_or(head_dim); + + let context_length = source.get(".context_length")?.to_u32()? as usize; let rope_freq_weight: Option> = source .tensor("rope_freqs.weight", device) @@ -425,15 +649,17 @@ where rope_theta: rope_freq_base, context_length, head_dimension: head_dim, - n_head: head_count, n_layer: block_count, start_token_string, - stop_token, + stop_tokens, stop_token_string, chat_template, rope_scaling, sliding_window_type, sliding_window_size, + layer_sliding_window_sizes, + final_logit_softcapping, + per_layer_embedding_length, vision_start_token: tokens .iter() .position(|v| &**v == "<|vision_start|>") @@ -445,6 +671,15 @@ where image_pad_token: tokens .iter() .position(|v| &**v == "<|image_pad|>") + .or_else(|| tokens.iter().position(|v| &**v == "<|image|>")) + .map(|v| v as u32), + image_start_token: tokens + .iter() + .position(|v| &**v == "<|image>") + .map(|v| v as u32), + image_end_token: tokens + .iter() + .position(|v| &**v == "") .map(|v| v as u32), video_pad_token: tokens .iter() @@ -462,19 +697,48 @@ where }; let config = Arc::new(config); - let rope: RopeImplementation = - rope::RopeImplementation::new(&config, config.rope_theta, device)?; + let rope: RopeImplementation = rope::RopeImplementation::new_with_head_dimension( + &config, + head_dim, + config.rope_freq_weight.as_ref(), + config.rope_theta, + device, + )?; let sliding_rope: Option> = rope_freq_base_sliding + .filter(|_| layer_is_sliding.iter().any(|is_sliding| *is_sliding)) .map(|rope_freq_base_sliding| { - RopeImplementation::new(&config, rope_freq_base_sliding, device) + let rope_freq_weight = (!is_gemma4) + .then_some(config.rope_freq_weight.as_ref()) + .flatten(); + RopeImplementation::new_with_head_dimension( + &config, + head_dim_swa, + rope_freq_weight, + rope_freq_base_sliding, + device, + ) }) .transpose()?; let tok_embeddings_q = source.tensor("token_embd.weight", device).await?; let tok_embedding_scale = - (&*architecture == "gemma3").then(|| (embedding_length as f32).sqrt()); + (&*architecture == "gemma3" || is_gemma4).then(|| (embedding_length as f32).sqrt()); let tok_embeddings = Embedding::new(tok_embeddings_q.clone()); + let (per_layer_tok_embeddings, per_layer_model_proj, per_layer_proj_norm) = + if per_layer_embedding_length.is_some() { + let embeddings = source.tensor("per_layer_token_embd.weight", device).await?; + let model_proj = source.tensor("per_layer_model_proj.weight", device).await?; + let proj_norm = source.tensor("per_layer_proj_norm.weight", device).await?; + ( + Some(Embedding::new(embeddings)), + Some(model_proj), + Some(decode_norm(proj_norm, rms_norm_eps)?), + ) + } else { + (None, None, None) + }; + let norm = source.tensor("output_norm.weight", device).await?; let norm = decode_norm(norm, rms_norm_eps)?; let output = match source.tensor("output.weight", device).await { @@ -484,11 +748,43 @@ where tok_embeddings_q.clone() } }; + let mut layers = Vec::with_capacity(block_count); let interleaved_rope = architecture.as_ref() != "qwen2" && architecture.as_ref() != "qwen3" - && architecture.as_ref() != "gemma3"; + && architecture.as_ref() != "gemma3" + && !is_gemma4; + for layer_idx in 0..block_count { + let layer_is_sliding = layer_is_sliding.get(layer_idx).copied().unwrap_or(false); + let layer_head_dim = if is_gemma4 && layer_is_sliding { + head_dim_swa + } else { + head_dim + }; + let layer_attention_width = head_count * layer_head_dim; + let layer_sliding_window_size = config + .layer_sliding_window_sizes + .as_ref() + .and_then(|sizes| sizes.get(layer_idx).copied().flatten()); + + let has_kv = !is_gemma4 || layer_idx < n_layer_kv_from_start; + let shared_kv_layer = if is_gemma4 && !has_kv { + let offset = if layer_is_sliding { 2 } else { 1 }; + Some(n_layer_kv_from_start.saturating_sub(offset)) + } else { + None + }; + + let rope_cache = if layer_is_sliding { + sliding_rope + .as_ref() + .cloned() + .unwrap_or_else(|| rope.clone()) + } else { + rope.clone() + }; + let prefix = format!("blk.{layer_idx}"); let attention_variant = if let Ok(attention_qkv) = source .tensor(&format!("{prefix}.attn_qkv.weight"), device) @@ -502,13 +798,29 @@ where let q = source .tensor(&format!("{prefix}.attn_q.weight"), device) .await?; - let k = source - .tensor(&format!("{prefix}.attn_k.weight"), device) - .await?; - let v = source - .tensor(&format!("{prefix}.attn_v.weight"), device) - .await?; - let qkv = QMatrix::concat_rows(&[&q, &k, &v]); + let k = if has_kv { + Some( + source + .tensor(&format!("{prefix}.attn_k.weight"), device) + .await?, + ) + } else { + None + }; + let v = if has_kv { + Some( + source + .tensor(&format!("{prefix}.attn_v.weight"), device) + .await?, + ) + } else { + None + }; + let qkv = if let (Some(k), Some(v)) = (&k, &v) { + QMatrix::concat_rows(&[&q, k, v]) + } else { + None + }; let bias_q = source .tensor(&format!("{prefix}.attn_q.bias"), device) .await; @@ -535,6 +847,15 @@ where .tensor(&format!("{prefix}.attn_k_norm.weight"), device) .await .ok(); + let attention_v_norm = if is_gemma4 && has_kv { + Some(RmsNorm::new( + Tensor::ones(device, [layer_head_dim]), + None, + rms_norm_eps as f32, + )) + } else { + None + }; let separate = SeparateAttention { attention_wq: q, attention_qkv: qkv, @@ -546,6 +867,7 @@ where .map(|norm| decode_norm(norm, rms_norm_eps)) .transpose()?, attention_wv: v, + attention_v_norm, interleaved_rope, bias, }; @@ -566,10 +888,16 @@ where let feed_forward_w3 = source .tensor(&format!("{prefix}.ffn_up.weight"), device) .await?; - FeedForwardVariant::Llama(Box::new(LlamaFeedForward::new( + let activation = if is_gemma4 { + FeedForwardActivation::Gelu + } else { + FeedForwardActivation::Silu + }; + FeedForwardVariant::Llama(Box::new(LlamaFeedForward::new_with_activation( feed_forward_w1, feed_forward_w2, feed_forward_w3, + activation, ))) } else { // Otherwise, try to read from the up, and down weights @@ -603,26 +931,45 @@ where .await .ok(); - let mut layer_sliding_window_size = None; - - let rope_cache = if let ( - Some(rope_sliding), - Some(sliding_window_type), - Some(sliding_window_size), - ) = ( - sliding_rope.as_ref(), - sliding_window_type, - sliding_window_size, - ) { - let is_sliding = (layer_idx + 1) % sliding_window_type != 0; - if is_sliding { - layer_sliding_window_size = Some(sliding_window_size); - rope_sliding.clone() - } else { - rope.clone() - } + let per_layer_inp_gate = if per_layer_embedding_length.is_some() { + Some( + source + .tensor(&format!("{prefix}.inp_gate.weight"), device) + .await?, + ) } else { - rope.clone() + None + }; + let per_layer_proj = if per_layer_embedding_length.is_some() { + Some( + source + .tensor(&format!("{prefix}.proj.weight"), device) + .await?, + ) + } else { + None + }; + let per_layer_post_norm = if per_layer_embedding_length.is_some() { + let norm = source + .tensor(&format!("{prefix}.post_norm.weight"), device) + .await?; + Some(decode_norm(norm, rms_norm_eps)?) + } else { + None + }; + let layer_output_scale = source + .tensor(&format!("{prefix}.layer_output_scale.weight"), device) + .await + .ok() + .map(&dequantize_1d); + // Gemma 4 folds the query pre-attention scaling into the exported + // `attn_q_norm` weights, so the softmax logits are already scaled and + // flash-attention must run with a unit scale. Every other supported + // architecture applies the usual 1/sqrt(head_dim) here. + let attention_scale = if is_gemma4 { + 1.0 + } else { + 1.0 / (layer_head_dim as f32).sqrt() }; layers.push(LlamaAttention { @@ -639,10 +986,16 @@ where .transpose()?, n_head: head_count, n_kv_head: head_count_kv, - head_dim, - hidden_size: config.hidden_size(), + head_dim: layer_head_dim, + hidden_size: layer_attention_width, rope_cache, sliding_window_size: layer_sliding_window_size, + attention_scale, + shared_kv_layer, + per_layer_inp_gate, + per_layer_proj, + per_layer_post_norm, + layer_output_scale, }) } @@ -650,11 +1003,11 @@ where #[cfg(feature = "vision")] let vision_encoder = if let (Some(vision_ct), Some(vision_bytes)) = (vision_ct, vision_bytes) { - Some(vision::QwenVisionTransformer::from_gguf( + Some(vision::VisionTransformer::from_gguf( vision_ct, &vision_bytes, device, - )) + )?) } else { None }; @@ -662,12 +1015,15 @@ where config, tok_embeddings, tok_embedding_scale, + per_layer_tok_embeddings, + per_layer_model_proj, + per_layer_proj_norm, layers, norm, output, masks: Default::default(), #[cfg(feature = "vision")] - vision_encoder: vision_encoder.transpose()?, + vision_encoder, }) } } @@ -684,6 +1040,10 @@ where #[cfg(feature = "vision")] { self.vision_encoder.is_none() + || matches!( + self.vision_encoder, + Some(vision::VisionTransformer::Gemma(_)) + ) } #[cfg(not(feature = "vision"))] @@ -692,6 +1052,73 @@ where } } + /// Compute the Gemma "per-layer input" embeddings that are blended into each + /// decoder layer (the `inp_gate`/`proj`/`post_norm` path). Returns `None` + /// for models without per-layer embeddings. + /// + /// `per_layer_token_ids` is invoked lazily, so models without per-layer + /// embeddings never pay for building the token id tensor. It yields the + /// `[batch, positions]` ids used for the per-layer token lookup, with + /// image/control tokens already zeroed by the caller. A single position is + /// broadcast across the whole sequence (image chunks share one zeroed + /// per-layer token). + fn compute_per_layer_inputs( + &self, + embeddings_f32: &Tensor<3, f32>, + per_layer_token_ids: impl FnOnce() -> Tensor<2, u32, B>, + ) -> Option> + where + B: Fusion<2, u32>, + { + let ( + per_layer_tok_embeddings, + per_layer_model_proj, + per_layer_proj_norm, + per_layer_embedding_length, + ) = match ( + &self.per_layer_tok_embeddings, + &self.per_layer_model_proj, + &self.per_layer_proj_norm, + self.config.per_layer_embedding_length, + ) { + (Some(embeddings), Some(model_proj), Some(proj_norm), Some(length)) => { + (embeddings, model_proj, proj_norm, length) + } + _ => return None, + }; + + let [batch, seq, embedding_dim] = embeddings_f32.shape(); + let n_layer = self.config.n_layer; + let per_layer_token_ids = per_layer_token_ids(); + let positions = per_layer_token_ids.shape()[1]; + + let token_inputs = per_layer_tok_embeddings.forward::<2, 3, _>(&per_layer_token_ids) + * (per_layer_embedding_length as f32).sqrt(); + let token_inputs = + token_inputs.reshape([batch, positions, n_layer, per_layer_embedding_length]); + let token_inputs: Tensor<4, f32> = if positions == seq { + token_inputs.to_concrete() + } else { + token_inputs + .broadcast_as([batch, seq, n_layer, per_layer_embedding_length]) + .to_concrete() + }; + + let projected_inputs = + embeddings_f32.q_mat_mul(per_layer_model_proj) * (1.0 / (embedding_dim as f32).sqrt()); + let projected_inputs = projected_inputs + .reshape([batch, seq, n_layer, per_layer_embedding_length]) + .to_concrete(); + let projected_inputs: Tensor<4, F> = + per_layer_proj_norm.forward_generic_4d(&projected_inputs.cast()); + + Some( + ((projected_inputs.cast::() + token_inputs) * (1.0 / 2.0_f32.sqrt())) + .to_concrete() + .cast(), + ) + } + pub fn encode_tokens( &self, raw_tokens: &[u32], @@ -714,43 +1141,25 @@ where images.push(image); grid_thw.push(thw) } + } else if !raw_images.is_empty() { + return Err(fusor::Error::msg( + "Media inputs require a loaded vision encoder.", + )); } - // Add any image padding tokens to the tokens if needed - let tokens = if let (Some(image_pad_token), Some(vision_start_token), Some(vision)) = ( - self.config.image_pad_token, - self.config.vision_start_token, - &self.vision_encoder, - ) { - let mut tokens = Vec::new(); - let mut token_iter = raw_tokens.iter().copied(); - let mut image_iter = grid_thw.iter(); - while let Some(token) = token_iter.next() { - tokens.push(token); - let start_index = tokens.len(); - if token == vision_start_token { - match token_iter.next() { - Some(next) if next == image_pad_token => { - // Push a pad token for every image token - let grid = image_iter.next().ok_or_else(|| { - fusor::Error::msg( - "Image pad token found without matching image.", - ) - })?; - for _ in 0..grid.iter().product::() - / (vision.spacial_merge_size as u32).pow(2) - { - tokens.push(image_pad_token); - } - image_token_ranges.push(start_index..tokens.len()); - } - Some(next) => { - tokens.push(next); - } - None => break, - } - } - } + // Add image padding tokens for any placeholders in the prompt. + let tokens = if let (Some(image_pad_token), Some(vision)) = + (self.config.image_pad_token, &self.vision_encoder) + { + let (tokens, ranges) = vision.expand_image_tokens( + raw_tokens, + image_pad_token, + self.config.vision_start_token, + self.config.image_start_token, + self.config.image_end_token, + &grid_thw, + )?; + image_token_ranges = ranges; tokens } else { raw_tokens.to_vec() @@ -809,42 +1218,86 @@ where embeddings_f32 = (embeddings_f32 * scale).to_concrete(); } #[cfg(feature = "vision")] - let mut embeddings: Tensor<3, F> = embeddings_f32.cast(); - #[cfg(not(feature = "vision"))] - let embeddings: Tensor<3, F> = embeddings_f32.cast(); - #[cfg(feature = "vision")] let mut pos_ids = None; #[cfg(not(feature = "vision"))] let pos_ids = None; - #[cfg(feature = "vision")] - let batch_size = embeddings.shape()[0]; - #[cfg(feature = "vision")] - let embed_dim = embeddings.shape()[2]; #[cfg(feature = "vision")] if let Some(vision_encoder) = &self.vision_encoder { - for ((pixels, grid), range) in images.iter().zip(&grid_thw).zip(image_token_ranges) { + let batch_size = embeddings_f32.shape()[0]; + let embed_dim = embeddings_f32.shape()[2]; + for ((pixels, grid), range) in + images.iter().zip(&grid_thw).zip(image_token_ranges.iter()) + { let pixels_f: Tensor<2, F> = pixels.cast(); let image_embeds = vision_encoder.forward_image(&pixels_f, *grid)?; - let image_embeds_3d = image_embeds.unsqueeze(0); - embeddings = - embeddings.slice_assign([0..batch_size, range, 0..embed_dim], &image_embeds_3d); + let mut image_embeds_f32: Tensor<2, f32> = image_embeds.cast(); + if vision_encoder.outputs_on_isolated_device() { + image_embeds_f32 = copy_image_embeddings_to_device(image_embeds_f32, device)?; + } + debug_tensor_stats_f32(&image_embeds_f32, "image_embeds_projected"); + let image_embeds_3d: Tensor<3, f32> = image_embeds_f32.unsqueeze(0).to_concrete(); + embeddings_f32 = embeddings_f32.slice_assign( + [0..batch_size, range.clone(), 0..embed_dim], + &image_embeds_3d, + ); } - let (new_pos_ids, new_start_time) = - vision_encoder.get_rope_index(&tokens, &grid_thw, &self.config, start_time)?; - if let Some(cache) = cache.as_mut() { - cache.start_time = new_start_time; + if let Some((new_pos_ids, new_start_time)) = + vision_encoder.get_rope_index(&tokens, &grid_thw, &self.config, start_time)? + { + if let Some(cache) = cache.as_mut() { + cache.start_time = new_start_time; + } + let pos_f32: Tensor<2, f32> = new_pos_ids.cast(); + let pos_f: Tensor<2, F> = pos_f32.cast(); + pos_ids = Some(pos_f); } - let pos_f32: Tensor<2, f32> = new_pos_ids.cast(); - let pos_f: Tensor<2, F> = pos_f32.cast(); - pos_ids = Some(pos_f); } + let per_layer_inputs = self.compute_per_layer_inputs(&embeddings_f32, || { + #[cfg(feature = "vision")] + { + let mut per_layer_tokens = tokens.clone(); + for range in &image_token_ranges { + per_layer_tokens[range.clone()].fill(0); + } + if let Some(image_start_token) = self.config.image_start_token { + for token in &mut per_layer_tokens { + if *token == image_start_token { + *token = 0; + } + } + } + if let Some(image_end_token) = self.config.image_end_token { + for token in &mut per_layer_tokens { + if *token == image_end_token { + *token = 0; + } + } + } + Tensor::from_slice( + device, + [1, per_layer_tokens.len()], + per_layer_tokens.as_slice(), + ) + } + #[cfg(not(feature = "vision"))] + { + x.clone() + } + }); + let embeddings: Tensor<3, F> = embeddings_f32.cast(); + Ok(EncodedTokens { embeddings, + per_layer_inputs, seq_len, index_pos, pos_ids, + #[cfg(feature = "vision")] + non_causal_token_ranges: image_token_ranges, + #[cfg(not(feature = "vision"))] + non_causal_token_ranges: Vec::new(), }) } @@ -860,8 +1313,246 @@ where f32: CastTo + CastTensor, { let x_f32 = self.forward_last_hidden_f32(tokens, images, device, cache)?; - let result_f32 = x_f32.q_mat_mul(&self.output); - Ok(result_f32.cast()) + self.forward_logits_from_hidden_f32(x_f32) + } + + fn forward_logits_from_hidden_f32(&self, x_f32: Tensor<2, f32>) -> Result> + where + f32: CastTo + CastTensor, + { + Ok(self.logits_from_hidden_f32(x_f32).cast()) + } + + pub(crate) fn logits_from_hidden_f32( + &self, + x_f32: Tensor, + ) -> Tensor { + self.apply_final_logit_softcap(x_f32.q_mat_mul(&self.output)) + } + + pub(crate) fn apply_final_logit_softcap( + &self, + logits: Tensor, + ) -> Tensor { + if let Some(softcap) = self.config.final_logit_softcapping { + logits.mul_scalar(1.0 / softcap).tanh().mul_scalar(softcap) + } else { + logits + } + } + + pub(crate) fn should_chunk_multimodal_prompt(&self) -> bool { + if std::env::var_os("KALOSM_LLAMA_DISABLE_MULTIMODAL_CHUNK").is_some() { + return false; + } + if std::env::var_os("KALOSM_LLAMA_FORCE_MULTIMODAL_CHUNK").is_some() { + return self.config.image_pad_token.is_some(); + } + #[cfg(feature = "vision")] + { + matches!( + self.vision_encoder, + Some(vision::VisionTransformer::Gemma(_)) + ) && self.config.image_pad_token.is_some() + } + #[cfg(not(feature = "vision"))] + { + false + } + } + + pub(crate) fn forward_chunked_multimodal( + &self, + tokens: &[u32], + images: &[LlamaImage], + device: &Device, + mut cache: Option<&mut LlamaCache>, + ) -> Result> + where + F: CastTo + CastTensor + Default, + f32: CastTo + CastTensor, + { + let Some(image_pad_token) = self.config.image_pad_token else { + return self.forward(tokens, images, device, cache); + }; + if images.is_empty() { + return self.forward(tokens, images, device, cache); + } + + let mut image_index = 0; + let mut segment_start = 0; + let mut text_prefix = Vec::new(); + let mut last_logits = None; + for (index, token) in tokens.iter().copied().enumerate() { + if token != image_pad_token || image_index >= images.len() { + continue; + } + let mut text_tokens = std::mem::take(&mut text_prefix); + text_tokens.extend_from_slice(&tokens[segment_start..index]); + if let Some(image_start_token) = self.config.image_start_token { + text_tokens.push(image_start_token); + } + if !text_tokens.is_empty() { + self.forward_text_chunk_for_multimodal( + &text_tokens, + device, + cache.as_deref_mut(), + false, + )?; + } + let image_hidden = self.forward_image_embeddings_only_hidden_f32( + &images[image_index..image_index + 1], + device, + cache.as_deref_mut(), + )?; + image_index += 1; + segment_start = index + 1; + if let Some(image_end_token) = self.config.image_end_token { + text_prefix.push(image_end_token); + } + if segment_start < tokens.len() || !text_prefix.is_empty() { + resolve_intermediate_hidden_f32(&image_hidden); + } else { + last_logits = Some(self.forward_logits_from_hidden_f32(image_hidden)?); + } + } + + if segment_start < tokens.len() || !text_prefix.is_empty() { + let mut text_tokens = text_prefix; + text_tokens.extend_from_slice(&tokens[segment_start..]); + last_logits = + self.forward_text_chunk_for_multimodal(&text_tokens, device, cache, true)?; + } + + last_logits.ok_or_else(|| fusor::Error::msg("No tokens to forward")) + } + + fn forward_text_chunk_for_multimodal( + &self, + text_tokens: &[u32], + device: &Device, + mut cache: Option<&mut LlamaCache>, + return_logits: bool, + ) -> Result>> + where + F: CastTo + CastTensor + Default, + f32: CastTo + CastTensor, + { + if text_tokens.is_empty() { + return Ok(None); + } + + let has_cache_prefix = cache + .as_ref() + .map(|cache| !cache.tokens.is_empty()) + .unwrap_or(false); + let incremental = device.is_gpu() + && has_cache_prefix + && text_tokens.len() > 1 + && std::env::var_os("KALOSM_LLAMA_ENABLE_MULTIMODAL_TEXT_INCREMENTAL").is_some(); + + if incremental { + let last = text_tokens.len() - 1; + for (index, token) in text_tokens.iter().copied().enumerate() { + let one = [token]; + if return_logits && index == last { + return Ok(Some(self.forward( + &one, + &[], + device, + cache.as_deref_mut(), + )?)); + } + let hidden = + self.forward_last_hidden_f32(&one, &[], device, cache.as_deref_mut())?; + resolve_intermediate_hidden_f32(&hidden); + } + return Ok(None); + } + + if return_logits { + Ok(Some(self.forward(text_tokens, &[], device, cache)?)) + } else { + let hidden = self.forward_last_hidden_f32(text_tokens, &[], device, cache)?; + resolve_intermediate_hidden_f32(&hidden); + Ok(None) + } + } + + #[cfg(feature = "vision")] + fn forward_image_embeddings_only_hidden_f32( + &self, + raw_images: &[LlamaImage], + device: &Device, + mut cache: Option<&mut LlamaCache>, + ) -> Result> + where + F: CastTo + CastTensor + Default, + f32: CastTo + CastTensor, + { + let vision_encoder = self + .vision_encoder + .as_ref() + .ok_or_else(|| fusor::Error::msg("Image chunk requires a vision encoder"))?; + let image_pad_token = self + .config + .image_pad_token + .ok_or_else(|| fusor::Error::msg("Image chunk requires an image token"))?; + let (image, hints) = raw_images + .first() + .ok_or_else(|| fusor::Error::msg("Image chunk requires an image"))?; + let t_encode = Instant::now(); + let (pixels, grid) = + vision_encoder.preprocess_image(image, hints.min_tokens(), hints.max_tokens())?; + let pixels_f: Tensor<2, F> = pixels.cast(); + let image_embeds = vision_encoder.forward_image(&pixels_f, grid)?; + let mut embeddings_f32: Tensor<2, f32> = image_embeds.cast(); + if vision_encoder.outputs_on_isolated_device() { + embeddings_f32 = copy_image_embeddings_to_device(embeddings_f32, device)?; + } + debug_tensor_stats_f32(&embeddings_f32, "image_embeds_projected"); + let seq_len = embeddings_f32.shape()[0]; + let embeddings_f32 = embeddings_f32.unsqueeze(0).to_concrete(); + + let index_pos = cache.as_ref().map(|c| c.tokens.len()).unwrap_or_default(); + if let Some(cache) = cache.as_mut() { + cache + .tokens + .extend(std::iter::repeat_n(image_pad_token, seq_len)); + } + + // Image chunks carry no text tokens, so every position shares a single + // zeroed per-layer token that the helper broadcasts across the chunk. + let per_layer_inputs = self.compute_per_layer_inputs(&embeddings_f32, || { + Tensor::from_slice(device, [1, 1], &[0u32]) + }); + + let encoded = EncodedTokens { + embeddings: embeddings_f32.cast(), + per_layer_inputs, + seq_len, + index_pos, + pos_ids: None, + // The whole image chunk attends bidirectionally. + non_causal_token_ranges: std::iter::once(0..seq_len).collect(), + }; + self.forward_last_hidden_from_embeddings(encoded, device, cache, Some(t_encode.elapsed())) + } + + #[cfg(not(feature = "vision"))] + fn forward_image_embeddings_only_hidden_f32( + &self, + _raw_images: &[LlamaImage], + _device: &Device, + _cache: Option<&mut LlamaCache>, + ) -> Result> + where + F: CastTo + CastTensor + Default, + f32: CastTo + CastTensor, + { + Err(fusor::Error::msg( + "Image chunks require the `vision` feature", + )) } pub(crate) fn forward_last_hidden_f32( @@ -880,6 +1571,32 @@ where self.forward_last_hidden_from_embeddings(encoded, device, cache, Some(t_encode.elapsed())) } + pub(crate) fn forward_logits_and_nextn_f32( + &self, + tokens: &[u32], + images: &[LlamaImage], + device: &Device, + mut cache: Option<&mut LlamaCache>, + ) -> Result + where + F: CastTo + CastTensor + Default, + f32: CastTo + CastTensor, + { + let t_encode = Instant::now(); + let encoded = self.encode_tokens(tokens, images, device, cache.as_deref_mut())?; + let pre_norm = self.forward_pre_norm_hidden_from_embeddings( + encoded, + device, + cache, + Some(t_encode.elapsed()), + )?; + let normed = self.norm.forward_generic(&pre_norm.hidden); + let h_nextn: Tensor<2, f32> = normed.clone().cast::().squeeze(0).to_concrete(); + let logits = self.logits_from_hidden_f32(normed.cast::()); + let logits: Tensor<2, f32> = logits.squeeze(0).to_concrete(); + Ok(TargetBatchOutput { logits, h_nextn }) + } + pub(crate) fn forward_last_hidden_f32_gpu_token( &self, token: &Tensor<1, u32>, @@ -891,9 +1608,9 @@ where f32: CastTo + CastTensor, { #[cfg(feature = "vision")] - if self.vision_encoder.is_some() { + if !self.supports_gpu_token_run_ahead() { return Err(fusor::Error::msg( - "GPU token run-ahead is only available for text-only models", + "GPU token run-ahead is not available for this vision model", )); } @@ -910,34 +1627,106 @@ where if let Some(scale) = self.tok_embedding_scale { embeddings_f32 = (embeddings_f32 * scale).to_concrete(); } + let per_layer_inputs = self.compute_per_layer_inputs(&embeddings_f32, || x.clone()); let embeddings: Tensor<3, F> = embeddings_f32.cast(); let encoded = EncodedTokens { embeddings, + per_layer_inputs, seq_len: 1, index_pos: cache_slot, pos_ids: None, + non_causal_token_ranges: Vec::new(), }; let hidden = self.forward_last_hidden_from_embeddings(encoded, device, Some(cache), None)?; Ok((hidden, cache_slot)) } + fn get_attention_mask( + &self, + seq_len: usize, + index_pos: usize, + sliding_window_size: Option, + non_causal_token_ranges: &[Range], + device: &Device, + ) -> AttentionMask { + if non_causal_token_ranges.is_empty() { + return self + .masks + .get_mask(seq_len, index_pos, sliding_window_size, device); + } + + let cols = index_pos + seq_len; + let mask_data = non_causal_mask_data( + seq_len, + index_pos, + sliding_window_size, + non_causal_token_ranges, + ); + let mask: Tensor<2, f32> = Tensor::new(device, mask_data.as_slice()) + .reshape([seq_len, cols]) + .to_concrete(); + AttentionMask::::new(mask) + } + + fn can_skip_attention_mask( + seq_len: usize, + index_pos: usize, + sliding_window_size: Option, + non_causal_token_ranges: &[Range], + ) -> bool { + if std::env::var_os("KALOSM_LLAMA_DISABLE_NON_CAUSAL_MASK_SKIP").is_some() { + return false; + } + let all_query_tokens_are_non_causal = non_causal_token_ranges + .iter() + .any(|range| range.start == 0 && range.end >= seq_len); + if !all_query_tokens_are_non_causal { + return false; + } + + match sliding_window_size { + Some(window) => window >= index_pos + seq_len, + None => true, + } + } + fn forward_last_hidden_from_embeddings( &self, encoded: EncodedTokens, device: &Device, - mut cache: Option<&mut LlamaCache>, + cache: Option<&mut LlamaCache>, encode_elapsed: Option, ) -> Result> + where + F: CastTo + CastTensor + Default, + f32: CastTo + CastTensor, + { + let output = + self.forward_pre_norm_hidden_from_embeddings(encoded, device, cache, encode_elapsed)?; + let x = output.hidden.i((.., output.seq_len - 1, ..)); + let x = self.norm.forward_generic_2d(&x); + Ok(x.cast::()) + } + + fn forward_pre_norm_hidden_from_embeddings( + &self, + encoded: EncodedTokens, + device: &Device, + mut cache: Option<&mut LlamaCache>, + encode_elapsed: Option, + ) -> Result> where F: CastTo + CastTensor + Default, f32: CastTo + CastTensor, { let EncodedTokens { embeddings: mut layer_in, + per_layer_inputs, seq_len, index_pos, pos_ids, + non_causal_token_ranges, } = encoded; let _trace_text_prefill = seq_len > 1 && std::env::var_os("KALOSM_TRACE_TEXT").is_some(); let trace_forward_timing = @@ -958,6 +1747,7 @@ where debug_check_nan_f32(&probe, usize::MAX, "embed", index_pos); } + let mut non_causal_masks: HashMap, AttentionMask> = HashMap::new(); for (i, layer) in self.layers.iter().enumerate() { let x = layer_in; let residual: Tensor<3, f32> = x.cast(); @@ -966,23 +1756,95 @@ where let probe: fusor::Tensor<3, f32> = x.clone().cast(); debug_check_nan_f32(&probe, i, "post_attn_norm", index_pos); } - let mask = (seq_len > 1).then(|| { - self.masks - .get_mask(seq_len, index_pos, layer.sliding_window_size, device) - }); + let mask = if seq_len > 1 { + if Self::can_skip_attention_mask( + seq_len, + index_pos, + layer.sliding_window_size, + &non_causal_token_ranges, + ) { + None + } else if non_causal_token_ranges.is_empty() { + Some(self.get_attention_mask( + seq_len, + index_pos, + layer.sliding_window_size, + &non_causal_token_ranges, + device, + )) + } else { + Some( + non_causal_masks + .entry(layer.sliding_window_size) + .or_insert_with(|| { + self.get_attention_mask( + seq_len, + index_pos, + layer.sliding_window_size, + &non_causal_token_ranges, + device, + ) + }) + .clone(), + ) + } + } else { + None + }; + let shared_kv = if let Some(shared_kv_layer) = layer.shared_kv_layer { + let cache_ref = cache.as_deref().ok_or_else(|| { + fusor::Error::msg("Gemma 4 shared KV attention requires a populated cache") + })?; + let key = cache_ref.blocks[shared_kv_layer] + .k() + .cloned() + .ok_or_else(|| { + fusor::Error::msg("Gemma 4 shared KV source key cache is empty") + })?; + let value = cache_ref.blocks[shared_kv_layer] + .v() + .cloned() + .ok_or_else(|| { + fusor::Error::msg("Gemma 4 shared KV source value cache is empty") + })?; + Some((key, value)) + } else { + None + }; let mut attn = { - #[cfg(feature = "vision")] - { - if trace_layer_nan { - layer.forward_with_trace( - &x, - mask.as_ref(), - index_pos, - pos_ids.as_ref(), - cache.as_mut().map(|c| &mut c.blocks[i]), - i, - ) - } else { + if let Some((shared_key, shared_value)) = shared_kv.as_ref() { + layer.forward_with_shared_kv( + &x, + mask.as_ref(), + index_pos, + pos_ids.as_ref(), + shared_key, + shared_value, + ) + } else { + #[cfg(feature = "vision")] + { + if trace_layer_nan { + layer.forward_with_trace( + &x, + mask.as_ref(), + index_pos, + pos_ids.as_ref(), + cache.as_mut().map(|c| &mut c.blocks[i]), + i, + ) + } else { + layer.forward( + &x, + mask.as_ref(), + index_pos, + pos_ids.as_ref(), + cache.as_mut().map(|c| &mut c.blocks[i]), + ) + } + } + #[cfg(not(feature = "vision"))] + { layer.forward( &x, mask.as_ref(), @@ -992,16 +1854,6 @@ where ) } } - #[cfg(not(feature = "vision"))] - { - layer.forward( - &x, - mask.as_ref(), - index_pos, - pos_ids.as_ref(), - cache.as_mut().map(|c| &mut c.blocks[i]), - ) - } }; if trace_layer_nan { let probe: fusor::Tensor<3, f32> = attn.clone().cast(); @@ -1021,19 +1873,52 @@ where .forward_add_residuals(&x, &attn_f32, &residual) { layer_in = layer_out; - if trace_layer_nan { - let probe: fusor::Tensor<3, f32> = layer_in.cast(); - debug_check_nan_f32(&probe, i, "ffn_fused", index_pos); - } - continue; + } else { + let x = layer.feed_forward_variant.forward(&x); + let x_f32: Tensor<3, f32> = x.cast(); + layer_in = (x_f32 + attn_f32 + residual).cast(); } + } else { + let mut x = layer.feed_forward_variant.forward(&x); + if let Some(post_ffn_norm) = &layer.post_ffn_norm { + x = post_ffn_norm.forward_generic(&x); + } + let x_f32: Tensor<3, f32> = x.cast(); + layer_in = (x_f32 + attn_f32 + residual).cast(); } - let mut x = layer.feed_forward_variant.forward(&x); - if let Some(post_ffn_norm) = &layer.post_ffn_norm { - x = post_ffn_norm.forward_generic(&x); + if let ( + Some(per_layer_inputs), + Some(per_layer_inp_gate), + Some(per_layer_proj), + Some(per_layer_post_norm), + ) = ( + per_layer_inputs.as_ref(), + layer.per_layer_inp_gate.as_ref(), + layer.per_layer_proj.as_ref(), + layer.per_layer_post_norm.as_ref(), + ) { + let pe_in: Tensor<3, f32> = layer_in.cast(); + let gate = pe_in.q_mat_mul(per_layer_inp_gate).gelu(); + let layer_input: Tensor<3, F> = per_layer_inputs + .narrow(2, i, 1) + .squeeze::<3>(2) + .to_concrete(); + let projected = (gate.cast::() * layer_input) + .to_concrete() + .cast::() + .q_mat_mul(per_layer_proj) + .cast(); + let projected = per_layer_post_norm.forward_generic(&projected); + let projected_f32: Tensor<3, f32> = projected.cast(); + layer_in = (pe_in + projected_f32).cast(); + } + if let Some(layer_output_scale) = &layer.layer_output_scale { + let scale = layer_output_scale + .reshape([1, 1, 1]) + .broadcast_as(layer_in.shape()) + .to_concrete(); + layer_in = (layer_in * scale).to_concrete(); } - let x_f32: Tensor<3, f32> = x.cast(); - layer_in = (x_f32 + attn_f32 + residual).cast(); if trace_layer_nan { let probe: fusor::Tensor<3, f32> = layer_in.cast(); debug_check_nan_f32(&probe, i, "ffn_unfused", index_pos); @@ -1042,10 +1927,10 @@ where if trace_forward_timing { tracing::info!("[timing] text layer loop: {:.2?}", t_text_layers.elapsed()); } - let x = self.norm.forward_generic(&layer_in); - let x = x.i((.., seq_len - 1, ..)); - let out = x.cast::(); - Ok(out) + Ok(PreNormForwardOutput { + hidden: layer_in, + seq_len, + }) } pub(crate) fn output_matrix(&self) -> &QMatrix { diff --git a/models/kalosm-llama/src/raw/mtp.rs b/models/kalosm-llama/src/raw/mtp.rs new file mode 100644 index 000000000..aa5a9eaff --- /dev/null +++ b/models/kalosm-llama/src/raw/mtp.rs @@ -0,0 +1,429 @@ +use super::*; + +pub(crate) struct Gemma4MtpAssistant { + config: Arc>, + pre_projection: QMatrix, + post_projection: QMatrix, + layers: Vec>, + norm: RmsNorm<1, F>, + output: QMatrix, + layer_is_sliding: Vec, +} + +pub(crate) struct Gemma4MtpStep { + pub(crate) logits: Tensor<1, f32>, + pub(crate) h_nextn: Tensor<2, f32>, +} + +impl Gemma4MtpAssistant +where + MulOp: SimdBinaryOp, + AddOp: SimdBinaryOp, + SumOp: SimdReduceOp, +{ + #[cfg(not(target_arch = "wasm32"))] + pub(crate) fn from_gguf( + source: &mut ShardedVarBuilder, + device: &Device, + ) -> std::result::Result + where + f32: CastTensor + CastTo, + F: CastTensor + CastTo, + { + super::block_on_ready(Self::from_var_source(source, device)) + } + + pub(crate) async fn from_var_source( + source: &mut S, + device: &Device, + ) -> std::result::Result + where + f32: CastTensor + CastTo, + F: CastTensor + CastTo, + { + let dequantize_1d = |qmatrix: QMatrix| -> Tensor<1, F> { + let shape = qmatrix.shape(); + if shape.len() == 1 { + let w1d: Tensor<1, f32> = qmatrix.dequantize(); + w1d.cast() + } else if shape.len() == 2 { + let w2d: Tensor<2, f32> = qmatrix.dequantize(); + w2d.reshape([w2d.shape()[0] * w2d.shape()[1]]) + .to_concrete() + .cast() + } else { + panic!( + "Expected 1D or 2D tensor for dequantize_1d, got {}D", + shape.len() + ) + } + }; + let decode_norm = |qmatrix: QMatrix, eps: f64| -> Result> { + let weight = dequantize_1d(qmatrix); + Ok(RmsNorm::new(weight, None, eps as f32)) + }; + + let architecture = source.get("general.architecture")?.to_string()?.clone(); + if architecture.as_ref() != "gemma4-assistant" { + return Err(fusor::Error::msg(format!( + "MTP assistant architecture must be gemma4-assistant, got {architecture}" + )) + .into()); + } + + let block_count = source.get(".block_count")?.to_u32()? as usize; + let context_length = source.get(".context_length")?.to_u32()? as usize; + let embedding_length = source.get(".embedding_length")?.to_u32()? as usize; + let target_embedding_length = source.get(".embedding_length_out")?.to_u32()? as usize; + let head_count = source.get(".attention.head_count")?.to_u32()? as usize; + let head_count_kv = source.get(".attention.head_count_kv")?.to_u32()? as usize; + let rms_norm_eps = source.get(".attention.layer_norm_rms_epsilon")?.to_f32()? as f64; + let rope_freq_base = source + .get(".rope.freq_base") + .and_then(|m| Ok(m.to_f32()?)) + .unwrap_or(DEFAULT_ROPE_FREQUENCY); + let rope_freq_base_sliding = source + .get(".rope.freq_base_swa") + .and_then(|m| Ok(m.to_f32()?)) + .ok() + .or_else(|| { + source + .get(".rope.local_freq_base") + .and_then(|m| Ok(m.to_f32()?)) + .ok() + }) + .unwrap_or(GEMMA_DEFAULT_ROPE_FREQUENCY_SLIDING); + let sliding_window_size = source + .get(".attention.sliding_window") + .and_then(|m| Ok(m.to_u32()?)) + .ok() + .map(|x| x as usize); + let head_dim = source + .get(".attention.key_length") + .and_then(|m| Ok(m.to_u32()?)) + .ok() + .map(|x| x as usize) + .unwrap_or_else(|| target_embedding_length / head_count); + let head_dim_swa = source + .get(".attention.key_length_swa") + .and_then(|m| Ok(m.to_u32()?)) + .ok() + .map(|x| x as usize) + .unwrap_or(head_dim); + let sliding_window_pattern = source + .get(".attention.sliding_window_pattern") + .ok() + .and_then(|m| { + let values = m.to_array().ok()?; + values + .iter() + .map(|value| value.to_bool().ok()) + .collect::>>() + }); + let layer_is_sliding = sliding_window_pattern.unwrap_or_else(|| { + (0..block_count) + .map(|idx| sliding_window_size.is_some() && idx + 1 < block_count) + .collect() + }); + let layer_sliding_window_sizes = sliding_window_size.map(|window| { + layer_is_sliding + .iter() + .map(|is_sliding| is_sliding.then_some(window)) + .collect::>() + }); + let rope_freq_weight: Option> = source + .tensor("rope_freqs.weight", device) + .await + .ok() + .map(&dequantize_1d); + + let config = Arc::new(LlamaConfig { + rope_freq_weight, + rope_theta: rope_freq_base, + context_length, + head_dimension: head_dim, + n_layer: block_count, + start_token_string: String::new(), + stop_tokens: Vec::new(), + stop_token_string: String::new(), + chat_template: None, + rope_scaling: None, + sliding_window_type: None, + sliding_window_size, + layer_sliding_window_sizes, + final_logit_softcapping: None, + per_layer_embedding_length: None, + vision_start_token: None, + _vision_end_token: None, + image_pad_token: None, + image_start_token: None, + image_end_token: None, + video_pad_token: None, + mrope_sections: None, + }); + + let rope = RopeImplementation::new_with_head_dimension( + &config, + head_dim, + config.rope_freq_weight.as_ref(), + rope_freq_base, + device, + )?; + let sliding_rope = RopeImplementation::new_with_head_dimension( + &config, + head_dim_swa, + None, + rope_freq_base_sliding, + device, + )?; + + let pre_projection = source.tensor("nextn.pre_projection.weight", device).await?; + let post_projection = source + .tensor("nextn.post_projection.weight", device) + .await?; + let output_norm = source.tensor("output_norm.weight", device).await?; + let norm = decode_norm(output_norm, rms_norm_eps)?; + let token_embd = source.tensor("token_embd.weight", device).await?; + let output = source + .tensor("output.weight", device) + .await + .unwrap_or_else(|_| token_embd.clone()); + + let mut layers = Vec::with_capacity(block_count); + for layer_idx in 0..block_count { + let layer_is_sliding = layer_is_sliding.get(layer_idx).copied().unwrap_or(false); + let layer_head_dim = if layer_is_sliding { + head_dim_swa + } else { + head_dim + }; + let layer_attention_width = head_count * layer_head_dim; + let layer_sliding_window_size = config + .layer_sliding_window_sizes + .as_ref() + .and_then(|sizes| sizes.get(layer_idx).copied().flatten()); + let rope_cache = if layer_is_sliding { + sliding_rope.clone() + } else { + rope.clone() + }; + let prefix = format!("blk.{layer_idx}"); + let q = source + .tensor(&format!("{prefix}.attn_q.weight"), device) + .await?; + let q_norm = source + .tensor(&format!("{prefix}.attn_q_norm.weight"), device) + .await + .ok(); + let attention_variant = AttentionVariant::Separate(Box::new(SeparateAttention { + attention_wq: q, + attention_qkv: None, + attention_q_norm: q_norm + .map(|norm| decode_norm(norm, rms_norm_eps)) + .transpose()?, + attention_wk: None, + attention_k_norm: None, + attention_wv: None, + attention_v_norm: None, + interleaved_rope: false, + bias: None, + })); + + let attention_wo = source + .tensor(&format!("{prefix}.attn_output.weight"), device) + .await?; + let feed_forward_w1 = source + .tensor(&format!("{prefix}.ffn_gate.weight"), device) + .await?; + let feed_forward_w2 = source + .tensor(&format!("{prefix}.ffn_down.weight"), device) + .await?; + let feed_forward_w3 = source + .tensor(&format!("{prefix}.ffn_up.weight"), device) + .await?; + let attention_norm = source + .tensor(&format!("{prefix}.attn_norm.weight"), device) + .await?; + let post_attention_norm = source + .tensor(&format!("{prefix}.post_attention_norm.weight"), device) + .await + .ok(); + let ffn_norm = source + .tensor(&format!("{prefix}.ffn_norm.weight"), device) + .await?; + let ffn_post_norm = source + .tensor(&format!("{prefix}.post_ffw_norm.weight"), device) + .await + .ok(); + let layer_output_scale = source + .tensor(&format!("{prefix}.layer_output_scale.weight"), device) + .await + .ok() + .map(&dequantize_1d); + + layers.push(LlamaAttention { + attention_variant, + attention_wo: Linear::new(attention_wo, None), + attention_norm: decode_norm(attention_norm, rms_norm_eps)?, + post_attention_norm: post_attention_norm + .map(|norm| decode_norm(norm, rms_norm_eps)) + .transpose()?, + feed_forward_variant: FeedForwardVariant::Llama(Box::new( + LlamaFeedForward::new_with_activation( + feed_forward_w1, + feed_forward_w2, + feed_forward_w3, + FeedForwardActivation::Gelu, + ), + )), + ffn_norm: decode_norm(ffn_norm, rms_norm_eps)?, + post_ffn_norm: ffn_post_norm + .map(|norm| decode_norm(norm, rms_norm_eps)) + .transpose()?, + n_head: head_count, + n_kv_head: head_count_kv, + head_dim: layer_head_dim, + hidden_size: layer_attention_width, + rope_cache, + sliding_window_size: layer_sliding_window_size, + // Unit scale: like the Gemma 4 target model, the assistant's + // query pre-attention scaling is baked into its `attn_q_norm` + // weights (see `Model::from_gguf`). + attention_scale: 1.0, + shared_kv_layer: None, + per_layer_inp_gate: None, + per_layer_proj: None, + per_layer_post_norm: None, + layer_output_scale, + }); + } + + if pre_projection.shape().get(1).copied() != Some(target_embedding_length * 2) { + return Err(fusor::Error::msg(format!( + "unexpected Gemma4 MTP pre_projection input width {:?}, target hidden {target_embedding_length}", + pre_projection.shape() + )) + .into()); + } + if post_projection.shape().first().copied() != Some(target_embedding_length) { + return Err(fusor::Error::msg(format!( + "unexpected Gemma4 MTP post_projection output width {:?}, target hidden {target_embedding_length}", + post_projection.shape() + )) + .into()); + } + if embedding_length == 0 { + return Err(fusor::Error::msg("Gemma4 MTP assistant hidden size is zero").into()); + } + + Ok(Self { + config, + pre_projection, + post_projection, + layers, + norm, + output, + layer_is_sliding, + }) + } +} + +impl Gemma4MtpAssistant +where + F: CastTo + CastTensor, + f32: CastTo + CastTensor, + MulOp: SimdBinaryOp, + AddOp: SimdBinaryOp, + SumOp: SimdReduceOp, +{ + pub(crate) fn draft_step( + &self, + target: &Model, + token: u32, + h_nextn: &Tensor<2, f32>, + target_cache: &LlamaCache, + device: &Device, + position: usize, + ) -> Result { + let token_tensor: Tensor<2, u32> = + Tensor::from_slice(device, [1, 1], &[token]).to_concrete(); + let mut token_embedding = target.tok_embeddings.forward::<2, 3, _>(&token_tensor); + if let Some(scale) = target.tok_embedding_scale { + token_embedding = (token_embedding * scale).to_concrete(); + } + let h_nextn = h_nextn.unsqueeze(0).to_concrete(); + let projected = fusor::cat([token_embedding, h_nextn], 2) + .to_concrete() + .q_mat_mul(&self.pre_projection); + let mut layer_in: Tensor<3, F> = projected.cast(); + + for (layer_idx, layer) in self.layers.iter().enumerate() { + let target_kv_layer = self.target_kv_layer(target, layer_idx).ok_or_else(|| { + fusor::Error::msg("Gemma4 MTP could not find a matching target KV cache layer") + })?; + let key = target_cache.blocks[target_kv_layer] + .k() + .cloned() + .ok_or_else(|| fusor::Error::msg("Gemma4 MTP source key cache is empty"))?; + let value = target_cache.blocks[target_kv_layer] + .v() + .cloned() + .ok_or_else(|| fusor::Error::msg("Gemma4 MTP source value cache is empty"))?; + + let x = layer_in; + let residual: Tensor<3, f32> = x.cast(); + let x = layer.attention_norm.forward_generic(&x); + let mut attn = layer.forward_with_shared_kv(&x, None, position, None, &key, &value); + if let Some(post_attention_norm) = &layer.post_attention_norm { + attn = post_attention_norm.forward_generic(&attn); + } + let attn_f32: Tensor<3, f32> = attn.cast(); + let x = layer.ffn_norm.forward_residual_f32(&attn_f32, &residual); + let mut x = layer.feed_forward_variant.forward(&x); + if let Some(post_ffn_norm) = &layer.post_ffn_norm { + x = post_ffn_norm.forward_generic(&x); + } + let x_f32: Tensor<3, f32> = x.cast(); + layer_in = (x_f32 + attn_f32 + residual).cast(); + if let Some(layer_output_scale) = &layer.layer_output_scale { + let scale = layer_output_scale + .reshape([1, 1, 1]) + .broadcast_as(layer_in.shape()) + .to_concrete(); + layer_in = (layer_in * scale).to_concrete(); + } + } + + let normed = self.norm.forward_generic(&layer_in); + let logits: Tensor<1, f32> = target + .apply_final_logit_softcap(normed.cast::().q_mat_mul(&self.output)) + .squeeze(0) + .squeeze(0) + .to_concrete(); + let h_nextn: Tensor<2, f32> = normed + .cast::() + .q_mat_mul(&self.post_projection) + .squeeze(0) + .to_concrete(); + Ok(Gemma4MtpStep { logits, h_nextn }) + } + + fn target_kv_layer(&self, target: &Model, assistant_layer_idx: usize) -> Option { + let wants_sliding = self + .layer_is_sliding + .get(assistant_layer_idx) + .copied() + .unwrap_or(false); + target + .layers + .iter() + .enumerate() + .rev() + .find(|(_, layer)| layer.sliding_window_size.is_some() == wants_sliding) + .map(|(idx, layer)| layer.shared_kv_layer.unwrap_or(idx)) + } + + pub(crate) fn draft_n(&self) -> usize { + self.config.n_layer.max(1) + } +} diff --git a/models/kalosm-llama/src/raw/rope.rs b/models/kalosm-llama/src/raw/rope.rs index 6bbfe717b..0128bb0cf 100644 --- a/models/kalosm-llama/src/raw/rope.rs +++ b/models/kalosm-llama/src/raw/rope.rs @@ -60,15 +60,21 @@ where F: CastTo + CastTensor, f32: CastTo + CastTensor, { - pub fn new(config: &LlamaConfig, rope_theta: f32, device: &Device) -> fusor::Result { + pub fn new_with_head_dimension( + config: &LlamaConfig, + head_dimension: usize, + rope_freq_weight: Option<&Tensor<1, F>>, + rope_theta: f32, + device: &Device, + ) -> fusor::Result { if let Some(mrope_sections) = &config.mrope_sections { let cache = QwenVLRopeCache::new(config, rope_theta, mrope_sections, device)?; Ok(Self::QwenVL(cache)) } else { let inverse_frequency: Tensor<2, F> = create_inverse_frequency( config.rope_scaling.as_ref(), - config.rope_freq_weight.as_ref(), - config.head_dimension, + rope_freq_weight, + head_dimension, rope_theta, device, ); diff --git a/models/kalosm-llama/src/raw/vision/gemma.rs b/models/kalosm-llama/src/raw/vision/gemma.rs new file mode 100644 index 000000000..05b7ed03b --- /dev/null +++ b/models/kalosm-llama/src/raw/vision/gemma.rs @@ -0,0 +1,785 @@ +use fusor::{ + layers::RmsNorm, AddOp, CastTensor, CastTo, Device, DivOp, ExpOp, FloatDataType, FloatOps, + Fusion, MatmulImpl, MulOp, NegOp, QMatrix, Result, SimdBinaryOp, SimdElement, SimdUnaryOp, + Tensor, VarBuilder, +}; +use fusor_gguf::GgufMetadata; + +use crate::raw::rope::create_inverse_frequency; + +pub(crate) struct GemmaVisionTransformer { + patch_size: usize, + merge_size: usize, + patch_embed: GemmaVisionPatchEmbed, + position_embeddings: Tensor<3, F>, + blocks: Vec>, + projector_norm: RmsNorm<1, F>, + projector: GemmaClippedLinear, + std_bias: Option>, + std_scale: Option>, + image_mean: Vec, + image_std: Vec, + rope_theta: f32, + device: Device, + uses_isolated_device: bool, +} + +impl GemmaVisionTransformer +where + F: FloatDataType + + SimdElement + + FloatOps + + MatmulImpl + + Default + + CastTo + + CastTensor, + f32: CastTo + CastTensor, + fusor::MulOp: fusor::SimdBinaryOp, + fusor::AddOp: fusor::SimdBinaryOp, + fusor::SumOp: fusor::SimdReduceOp, +{ + pub(crate) fn from_gguf( + vision_ct: GgufMetadata, + vision_bytes: &[u8], + device: &Device, + ) -> Result { + let uses_isolated_device = cfg!(not(target_arch = "wasm32")) + && device.is_gpu() + && std::env::var_os("KALOSM_GEMMA4_VISION_DISABLE_SUBGROUPS").is_some(); + let vision_device = if uses_isolated_device { + // Debug fallback: use the no-subgroup sibling while sharing the + // same GPU adapter. The normal path is faster and covered by the + // split/combine long-attention tests. + device.without_subgroups() + } else { + device.clone() + }; + let device = &vision_device; + let block_count = metadata_usize(&vision_ct, "clip.vision.block_count", 16); + let head_count = metadata_usize(&vision_ct, "clip.vision.attention.head_count", 12); + let patch_size = metadata_usize(&vision_ct, "clip.vision.patch_size", 16); + let hidden_size = metadata_usize(&vision_ct, "clip.vision.embedding_length", 768); + let merge_size = metadata_usize(&vision_ct, "clip.vision.projector.scale_factor", 3); + let rope_theta = vision_ct + .metadata + .get("clip.vision.rope_theta") + .and_then(|x| x.to_f64().ok()) + .unwrap_or(100.0) as f32; + let layer_norm_eps = vision_ct + .metadata + .get("clip.vision.attention.layer_norm_epsilon") + .and_then(|x| x.to_f64().ok()) + .unwrap_or(1e-6) as f32; + let image_mean = metadata_f32_array(&vision_ct, "clip.vision.image_mean") + .unwrap_or_else(|| vec![0.0, 0.0, 0.0]); + let image_std = metadata_f32_array(&vision_ct, "clip.vision.image_std") + .unwrap_or_else(|| vec![1.0, 1.0, 1.0]); + + let mut cursor = std::io::Cursor::new(vision_bytes); + let mut vb = VarBuilder::from_gguf(&mut cursor)?; + let patch_embed = GemmaVisionPatchEmbed::new( + patch_size, + hidden_size, + &mut vb.pp("v.patch_embd"), + device, + )?; + let position_embeddings: Tensor<3, F> = vb + .get("v.position_embd.weight", device)? + .dequantize() + .cast(); + let head_dim = hidden_size / head_count; + let blocks = (0..block_count) + .map(|i| { + GemmaVisionBlock::new( + &mut vb.pp(format!("v.blk.{i}")), + device, + head_count, + head_dim, + hidden_size, + layer_norm_eps, + ) + }) + .collect::>>()?; + let projector_norm = + RmsNorm::new(Tensor::ones(device, [hidden_size]), None, layer_norm_eps); + let projector = clipped_linear(&mut vb, device, "mm.input_projection")?; + let std_bias: Option> = + vb.get("v.std_bias", device).ok().map(|x| x.dequantize()); + let std_scale: Option> = + vb.get("v.std_scale", device).ok().map(|x| x.dequantize()); + + Ok(Self { + patch_size, + merge_size, + patch_embed, + position_embeddings, + blocks, + projector_norm, + projector, + std_bias, + std_scale, + image_mean, + image_std, + rope_theta, + device: device.clone(), + uses_isolated_device, + }) + } + + pub(crate) fn uses_isolated_device(&self) -> bool { + self.uses_isolated_device + } + + pub(crate) fn preprocess_image( + &self, + image: &image::DynamicImage, + min_pixels: Option, + max_pixels: Option, + ) -> Result<(Tensor<2, f32>, [u32; 3])> { + let (target_width, target_height) = self.target_image_size(image, min_pixels, max_pixels); + let resized = image.resize_exact( + target_width as u32, + target_height as u32, + image::imageops::FilterType::Triangle, + ); + let rgb = image_to_rgb(&resized, &self.device)?; + let mean_tensor = Tensor::new(&self.device, &self.image_mean); + let mean = mean_tensor.reshape([1, 3, 1, 1]); + let std_tensor = Tensor::new(&self.device, &self.image_std); + let std = std_tensor.reshape([1, 3, 1, 1]); + let rgb = rgb + .sub_(&mean) + .div_(&std) + .mul_scalar(2.0) + .add_scalar(-1.0) + .to_concrete(); + let grid_h = target_height / self.patch_size; + let grid_w = target_width / self.patch_size; + if std::env::var_os("KALOSM_TRACE_VISION_STATS").is_some() { + tracing::info!( + "[vision_stats] preprocess original={}x{} target={}x{} grid=[1,{grid_h},{grid_w}] pooled_tokens={}", + image.width(), + image.height(), + target_width, + target_height, + (grid_h / self.merge_size) * (grid_w / self.merge_size) + ); + } + let patches = rgb + .reshape([1, 3, grid_h, self.patch_size, grid_w, self.patch_size]) + .permute([0, 2, 4, 1, 3, 5]) + .reshape([grid_h * grid_w, 3 * self.patch_size * self.patch_size]) + .to_concrete(); + crate::raw::debug_tensor_stats_f32(&patches, "image_patches"); + + Ok((patches, [1, grid_h as u32, grid_w as u32])) + } + + pub(crate) fn image_token_count(&self, grid: [u32; 3]) -> u32 { + grid[0] * (grid[1] / self.merge_size as u32) * (grid[2] / self.merge_size as u32) + } + + pub(crate) fn forward_image( + &self, + pixels: &Tensor<2, F>, + grid: [u32; 3], + ) -> Result> { + let [pos_x, pos_y] = self.patch_positions(grid)?; + let (cos_x, sin_x) = self.rope_sin_cos(&pos_x)?; + let (cos_y, sin_y) = self.rope_sin_cos(&pos_y)?; + let mut hidden_states = self.patch_embed.forward(pixels); + let hidden_f32: Tensor<2, f32> = hidden_states.cast(); + crate::raw::debug_tensor_stats_f32(&hidden_f32, "patch_embeds"); + hidden_states = self.add_position_embeddings(&hidden_states, &pos_x, &pos_y, grid)?; + let hidden_f32: Tensor<2, f32> = hidden_states.cast(); + crate::raw::debug_tensor_stats_f32(&hidden_f32, "patch_plus_position"); + let mut hidden_states = hidden_states.unsqueeze(0).to_concrete(); + for block in &self.blocks { + hidden_states = block.forward(&hidden_states, &cos_x, &sin_x, &cos_y, &sin_y); + } + let hidden_f32: Tensor<3, f32> = hidden_states.cast(); + crate::raw::debug_tensor_stats_f32(&hidden_f32, "vision_blocks_out"); + let mut hidden_states = self.pool(hidden_states, grid)?; + let hidden_f32: Tensor<3, f32> = hidden_states.cast(); + crate::raw::debug_tensor_stats_f32(&hidden_f32, "vision_pooled"); + if let (Some(std_bias), Some(std_scale)) = (&self.std_bias, &self.std_scale) { + let hidden_f32: Tensor<3, f32> = hidden_states.cast(); + hidden_states = hidden_f32.sub_(std_bias).mul_(std_scale).cast(); + } + let hidden_states = self.projector_norm.forward_generic(&hidden_states); + let hidden_f32: Tensor<3, f32> = hidden_states.cast(); + crate::raw::debug_tensor_stats_f32(&hidden_f32, "vision_projector_norm"); + Ok(self + .projector + .forward(&hidden_states) + .squeeze::<2>(0) + .to_concrete()) + } + + fn target_image_size( + &self, + image: &image::DynamicImage, + min_pixels: Option, + max_pixels: Option, + ) -> (usize, usize) { + let align = self.patch_size * self.merge_size; + let (min_pixels, max_pixels) = gemma_image_pixel_bounds(align, min_pixels, max_pixels); + smart_resize( + image.width() as usize, + image.height() as usize, + align, + min_pixels, + max_pixels, + ) + } + + fn add_position_embeddings( + &self, + hidden_states: &Tensor<2, F, B>, + x_ids: &Tensor<1, u32>, + y_ids: &Tensor<1, u32>, + grid: [u32; 3], + ) -> Result> + where + B: Fusion<2, F>, + { + let [grid_t, _grid_h, _grid_w] = grid; + if grid_t != 1 { + return Err(fusor::Error::msg( + "Gemma 4 vision currently supports image inputs, not video frames.", + )); + } + + let pos_x_table = self.position_embeddings.i((0, .., ..)).to_concrete(); + let pos_y_table = self.position_embeddings.i((1, .., ..)).to_concrete(); + let pos_x = pos_x_table.index_select(0, x_ids); + let pos_y = pos_y_table.index_select(0, y_ids); + let position_embeddings = (pos_x + pos_y).to_concrete(); + Ok((hidden_states.to_concrete() + position_embeddings).to_concrete()) + } + + fn patch_positions(&self, grid: [u32; 3]) -> Result<[Tensor<1, u32>; 2]> { + let [grid_t, grid_h, grid_w] = grid; + if grid_t != 1 { + return Err(fusor::Error::msg( + "Gemma 4 vision currently supports image inputs, not video frames.", + )); + } + let mut y_ids = Vec::with_capacity((grid_h * grid_w) as usize); + let mut x_ids = Vec::with_capacity((grid_h * grid_w) as usize); + for y in 0..grid_h { + for x in 0..grid_w { + y_ids.push(y); + x_ids.push(x); + } + } + Ok([ + Tensor::new(&self.device, &x_ids), + Tensor::new(&self.device, &y_ids), + ]) + } + + fn rope_sin_cos(&self, positions: &Tensor<1, u32>) -> Result<(Tensor<2, f32>, Tensor<2, f32>)> { + let half_head_dim = self + .blocks + .first() + .map(|block| block.head_dim() / 2) + .ok_or_else(|| { + fusor::Error::msg("Gemma 4 vision transformer must have at least one block") + })?; + let positions: Tensor<2, f32> = positions + .cast::() + .reshape([positions.shape()[0], 1]) + .to_concrete(); + let inv_freq: Tensor<2, f32> = create_inverse_frequency::( + None, + None, + half_head_dim, + self.rope_theta, + &self.device, + ); + let freqs = positions.matmul(&inv_freq); + Ok((freqs.cos().to_concrete(), freqs.sin().to_concrete())) + } + + fn pool(&self, hidden_states: Tensor<3, F>, grid: [u32; 3]) -> Result> { + let [batch, seq, hidden] = hidden_states.shape(); + let grid_h = grid[1] as usize; + let grid_w = grid[2] as usize; + if batch != 1 || seq != grid_h * grid_w { + return Err(fusor::Error::msg( + "Gemma 4 vision grid does not match hidden states", + )); + } + if grid_h % self.merge_size != 0 || grid_w % self.merge_size != 0 { + return Err(fusor::Error::msg( + "Gemma 4 vision grid must be divisible by the merge size", + )); + } + + let out_h = grid_h / self.merge_size; + let out_w = grid_w / self.merge_size; + let pooled = hidden_states + .reshape([ + batch, + out_h, + self.merge_size, + out_w, + self.merge_size, + hidden, + ]) + .sum::<5>(4) + .mul_scalar(F::from_f32(1.0 / self.merge_size as f32)) + .sum::<4>(2) + .mul_scalar(F::from_f32(1.0 / self.merge_size as f32)) + .reshape([batch, out_h * out_w, hidden]) + .mul_scalar(F::from_f32((hidden as f32).sqrt())) + .to_concrete(); + Ok(pooled) + } +} + +fn gemma_image_pixel_bounds( + align: usize, + min_pixels: Option, + max_pixels: Option, +) -> (usize, usize) { + ( + min_pixels + .map(|pixels| pixels as usize) + .unwrap_or(40 * align * align), + max_pixels + .map(|pixels| pixels as usize) + .unwrap_or(280 * align * align), + ) +} + +struct GemmaVisionPatchEmbed { + weight: Tensor<2, F>, +} + +impl GemmaVisionPatchEmbed +where + F: FloatDataType + SimdElement + FloatOps + MatmulImpl + CastTo + CastTensor, + f32: CastTo + CastTensor, +{ + fn new( + patch_size: usize, + hidden_size: usize, + vb: &mut VarBuilder, + device: &Device, + ) -> Result { + let weight: Tensor<4, F> = vb.get("weight", device)?.dequantize().cast(); + let weight = weight + .permute([1, 2, 3, 0]) + .reshape([3 * patch_size * patch_size, hidden_size]) + .to_concrete(); + Ok(Self { weight }) + } + + fn forward(&self, pixels: &Tensor<2, F, B>) -> Tensor<2, F> + where + B: Fusion<2, F>, + { + pixels.matmul(&self.weight) + } +} + +struct GemmaVisionBlock { + norm1: RmsNorm<1, F>, + norm2: RmsNorm<1, F>, + attn: GemmaVisionAttention, + attn_post_norm: RmsNorm<1, F>, + mlp: GemmaVisionFeedForward, + ffn_post_norm: RmsNorm<1, F>, +} + +impl GemmaVisionBlock +where + F: FloatDataType + SimdElement + Default + CastTo + CastTensor, + f32: CastTo + CastTensor, +{ + fn new( + vb: &mut VarBuilder, + device: &Device, + head_count: usize, + head_dim: usize, + hidden_size: usize, + layer_norm_eps: f32, + ) -> Result { + let norm1 = rms_norm(vb.get("ln1.weight", device)?, layer_norm_eps); + let norm2 = rms_norm(vb.get("ln2.weight", device)?, layer_norm_eps); + let attn_post_norm = rms_norm(vb.get("attn_post_norm.weight", device)?, layer_norm_eps); + let ffn_post_norm = rms_norm(vb.get("ffn_post_norm.weight", device)?, layer_norm_eps); + let attn = GemmaVisionAttention::new(vb, device, head_count, head_dim, hidden_size)?; + let mlp = GemmaVisionFeedForward::new(vb, device)?; + + Ok(Self { + norm1, + norm2, + attn, + attn_post_norm, + mlp, + ffn_post_norm, + }) + } + + fn head_dim(&self) -> usize { + self.attn.head_dim + } + + fn forward( + &self, + xs: &Tensor<3, F, B>, + cos_x: &Tensor<2, f32>, + sin_x: &Tensor<2, f32>, + cos_y: &Tensor<2, f32>, + sin_y: &Tensor<2, f32>, + ) -> Tensor<3, F> + where + B: Fusion<3, F>, + { + let residual: Tensor<3, f32> = xs.cast(); + let x = self.norm1.forward_generic(xs); + let attn = self.attn.forward(&x, cos_x, sin_x, cos_y, sin_y); + let attn = self.attn_post_norm.forward_generic(&attn); + let attn_f32: Tensor<3, f32> = attn.cast(); + let x = self.norm2.forward_residual_f32(&attn_f32, &residual); + let ffn = self.mlp.forward(&x); + let ffn = self.ffn_post_norm.forward_generic(&ffn); + let ffn_f32: Tensor<3, f32> = ffn.cast(); + (ffn_f32 + attn_f32 + residual).cast() + } +} + +struct GemmaVisionAttention { + q: GemmaClippedLinear, + k: GemmaClippedLinear, + v: GemmaClippedLinear, + out: GemmaClippedLinear, + q_norm: RmsNorm<1, F>, + k_norm: RmsNorm<1, F>, + v_norm: RmsNorm<1, F>, + head_count: usize, + head_dim: usize, + hidden_size: usize, +} + +impl GemmaVisionAttention +where + F: FloatDataType + SimdElement + Default + CastTo + CastTensor, + f32: CastTo + CastTensor, +{ + fn new( + vb: &mut VarBuilder, + device: &Device, + head_count: usize, + head_dim: usize, + hidden_size: usize, + ) -> Result { + Ok(Self { + q: clipped_linear(vb, device, "attn_q")?, + k: clipped_linear(vb, device, "attn_k")?, + v: clipped_linear(vb, device, "attn_v")?, + out: clipped_linear(vb, device, "attn_out")?, + q_norm: rms_norm(vb.get("attn_q_norm.weight", device)?, 1e-6), + k_norm: rms_norm(vb.get("attn_k_norm.weight", device)?, 1e-6), + v_norm: RmsNorm::new(Tensor::ones(device, [head_dim]), None, 1e-6), + head_count, + head_dim, + hidden_size, + }) + } + + fn forward( + &self, + xs: &Tensor<3, F, B>, + cos_x: &Tensor<2, f32>, + sin_x: &Tensor<2, f32>, + cos_y: &Tensor<2, f32>, + sin_y: &Tensor<2, f32>, + ) -> Tensor<3, F> + where + B: Fusion<3, F>, + { + let [batch, seq_len, _] = xs.shape(); + let q = self + .q + .forward(xs) + .reshape([batch, seq_len, self.head_count, self.head_dim]) + .transpose(1, 2) + .to_concrete(); + let k = self + .k + .forward(xs) + .reshape([batch, seq_len, self.head_count, self.head_dim]) + .transpose(1, 2) + .to_concrete(); + let v = self + .v + .forward(xs) + .reshape([batch, seq_len, self.head_count, self.head_dim]) + .transpose(1, 2) + .to_concrete(); + let q: Tensor<4, f32> = self.q_norm.forward_generic_4d(&q).cast(); + let k: Tensor<4, f32> = self.k_norm.forward_generic_4d(&k).cast(); + let v: Tensor<4, f32> = self.v_norm.forward_generic_4d(&v).cast(); + let half = self.head_dim / 2; + let q_x = q.narrow(3, 0, half).to_concrete(); + let k_x = k.narrow(3, 0, half).to_concrete(); + let q_y = q.narrow(3, half, half).to_concrete(); + let k_y = k.narrow(3, half, half).to_concrete(); + let (q_x, k_x) = q_x.rope_normal_pair_fused(&k_x, cos_x, sin_x); + let (q_y, k_y) = q_y.rope_normal_pair_fused(&k_y, cos_y, sin_y); + let q = Tensor::cat([q_x, q_y], 3).to_concrete(); + let k = Tensor::cat([k_x, k_y], 3).to_concrete(); + let attn = q.flash_attention(&k, &v, 1.0, None); + let attn = attn + .transpose(1, 2) + .reshape([batch, seq_len, self.hidden_size]) + .cast(); + self.out.forward(&attn) + } +} + +struct GemmaVisionFeedForward { + gate: GemmaClippedLinear, + down: GemmaClippedLinear, + up: GemmaClippedLinear, +} + +impl GemmaVisionFeedForward { + fn new(vb: &mut VarBuilder, device: &Device) -> Result { + Ok(Self { + gate: clipped_linear(vb, device, "ffn_gate")?, + down: clipped_linear(vb, device, "ffn_down")?, + up: clipped_linear(vb, device, "ffn_up")?, + }) + } + + fn forward(&self, x: &Tensor<3, F, B>) -> Tensor<3, F> + where + F: FloatDataType + SimdElement + Default + CastTo + CastTensor, + f32: CastTo + CastTensor, + B: Fusion<3, F>, + { + let gate = quick_gelu(&self.gate.forward_f32(x)); + let up = self.up.forward_f32(x); + let hidden = (gate * up).to_concrete(); + self.down.forward_from_f32(&hidden).cast() + } +} + +fn quick_gelu(x: &Tensor) -> Tensor +where + B: Fusion, + AddOp: SimdBinaryOp, + DivOp: SimdBinaryOp, + ExpOp: SimdUnaryOp, + MulOp: SimdBinaryOp, + NegOp: SimdUnaryOp, +{ + let x = x.to_concrete(); + let scaled = (x.clone() * 1.702).to_concrete(); + let sigmoid = scaled.sigmoid(); + (x * sigmoid).to_concrete() +} + +#[derive(Clone, Copy)] +struct ClampInfo { + input_min: f32, + input_max: f32, + output_min: f32, + output_max: f32, +} + +impl ClampInfo { + fn from_tensors(vb: &mut VarBuilder, device: &Device, prefix: &str) -> Self { + Self { + input_min: scalar_tensor(vb, device, &format!("{prefix}.input_min"), -f32::MAX), + input_max: scalar_tensor(vb, device, &format!("{prefix}.input_max"), f32::MAX), + output_min: scalar_tensor(vb, device, &format!("{prefix}.output_min"), -f32::MAX), + output_max: scalar_tensor(vb, device, &format!("{prefix}.output_max"), f32::MAX), + } + } + + fn has_input_clamp(self) -> bool { + self.input_min > -f32::MAX || self.input_max < f32::MAX + } + + fn has_output_clamp(self) -> bool { + self.output_min > -f32::MAX || self.output_max < f32::MAX + } +} + +struct GemmaClippedLinear { + weight: QMatrix, + clamp: ClampInfo, +} + +impl GemmaClippedLinear { + fn new(weight: QMatrix, clamp: ClampInfo) -> Self { + Self { weight, clamp } + } + + fn forward(&self, input: &Tensor<3, F, B>) -> Tensor<3, F> + where + F: FloatDataType + SimdElement + Default + CastTo + CastTensor, + f32: CastTo + CastTensor, + B: Fusion<3, F>, + { + self.forward_f32(input).cast() + } + + fn forward_f32(&self, input: &Tensor<3, F, B>) -> Tensor<3, f32> + where + F: FloatDataType + SimdElement + Default + CastTo + CastTensor, + B: Fusion<3, F>, + { + self.forward_from_f32(&input.cast::()) + } + + fn forward_from_f32(&self, input: &Tensor<3, f32, B>) -> Tensor<3, f32> + where + B: Fusion<3, f32>, + { + let input = if self.clamp.has_input_clamp() { + input + .clamp(self.clamp.input_min, self.clamp.input_max) + .to_concrete() + } else { + input.to_concrete() + }; + let output = input.q_mat_mul(&self.weight); + if self.clamp.has_output_clamp() { + output + .clamp(self.clamp.output_min, self.clamp.output_max) + .to_concrete() + } else { + output + } + } +} + +fn clipped_linear( + vb: &mut VarBuilder, + device: &Device, + prefix: &str, +) -> Result { + let weight = vb.get(&format!("{prefix}.weight"), device)?; + let clamp = ClampInfo::from_tensors(vb, device, prefix); + Ok(GemmaClippedLinear::new(weight, clamp)) +} + +fn scalar_tensor(vb: &mut VarBuilder, device: &Device, name: &str, default: f32) -> f32 { + let Ok(tensor) = vb.get(name, device) else { + return default; + }; + let tensor: Tensor<1, f32> = tensor.dequantize(); + first_f32(&tensor).unwrap_or(default) +} + +#[cfg(not(target_arch = "wasm32"))] +fn first_f32(tensor: &Tensor<1, f32>) -> Option { + let slice = pollster::block_on(tensor.as_slice()).ok()?; + slice.as_slice().first().copied() +} + +#[cfg(target_arch = "wasm32")] +fn first_f32(_tensor: &Tensor<1, f32>) -> Option { + None +} + +fn rms_norm(weight: QMatrix, eps: f32) -> RmsNorm<1, F> +where + F: FloatDataType + SimdElement + Default + CastTo + CastTensor, + f32: CastTo + CastTensor, +{ + let weight: Tensor<1, F> = weight.dequantize().cast(); + RmsNorm::new(weight, None, eps) +} + +fn metadata_usize(metadata: &GgufMetadata, key: &str, default: usize) -> usize { + metadata + .metadata + .get(key) + .and_then(|x| x.to_u64().ok()) + .map(|x| x as usize) + .unwrap_or(default) +} + +fn smart_resize( + width: usize, + height: usize, + align: usize, + min_pixels: usize, + max_pixels: usize, +) -> (usize, usize) { + let round_by_factor = |x: f64| ((x / align as f64).round() as usize).max(1) * align; + let ceil_by_factor = |x: f64| ((x / align as f64).ceil() as usize).max(1) * align; + let floor_by_factor = |x: f64| ((x / align as f64).floor() as usize).max(1) * align; + + let mut target_h = round_by_factor(height as f64); + let mut target_w = round_by_factor(width as f64); + let pixels = (width * height) as f64; + if target_h * target_w > max_pixels { + let beta = (pixels / max_pixels as f64).sqrt(); + target_h = floor_by_factor(height as f64 / beta); + target_w = floor_by_factor(width as f64 / beta); + } else if target_h * target_w < min_pixels { + let beta = (min_pixels as f64 / pixels).sqrt(); + target_h = ceil_by_factor(height as f64 * beta); + target_w = ceil_by_factor(width as f64 * beta); + } + + (target_w, target_h) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn gemma_image_pixel_bounds_honor_media_hints() { + let align = 8; + let requested_min = 10 * align * align; + let requested_max = 20 * align * align; + + assert_eq!( + gemma_image_pixel_bounds( + align, + Some(requested_min as u32), + Some(requested_max as u32) + ), + (requested_min, requested_max), + "Gemma preprocessing must honor caller-provided MediaHints image budgets" + ); + } +} + +fn metadata_f32_array(metadata: &GgufMetadata, key: &str) -> Option> { + metadata.metadata.get(key).and_then(|value| { + value.to_array().ok().map(|values| { + values + .iter() + .filter_map(|value| value.to_f32().ok()) + .collect() + }) + }) +} + +fn image_to_rgb(image: &image::DynamicImage, device: &Device) -> Result> { + let height = image.height() as usize; + let width = image.width() as usize; + let rgb = image.to_rgb8(); + let as_u32 = rgb + .into_raw() + .into_iter() + .map(|x| x as u32) + .collect::>(); + let data_tensor = Tensor::new(device, &as_u32); + let data = data_tensor.reshape([height, width, 3]); + let img = data.permute([2, 0, 1]).cast::() * (1.0 / 255.0); + + Ok(img.unsqueeze(0).to_concrete()) +} diff --git a/models/kalosm-llama/src/raw/vision/mod.rs b/models/kalosm-llama/src/raw/vision/mod.rs index f3671fb95..f0ef4c77c 100644 --- a/models/kalosm-llama/src/raw/vision/mod.rs +++ b/models/kalosm-llama/src/raw/vision/mod.rs @@ -1,3 +1,4 @@ +mod gemma; mod qwen; mod qwen_image_processing; mod qwen_patch_merger; @@ -6,6 +7,177 @@ mod qwen_vision; mod qwen_vision_block; mod qwen_vision_embed; +use std::ops::Range; + +use fusor::{CastTensor, CastTo, Device, FloatDataType, Result, SimdElement, Tensor}; +use fusor_gguf::GgufMetadata; + +pub(crate) use gemma::GemmaVisionTransformer; pub(crate) use qwen::QwenVisionTransformer; pub const QWEN_EPS: f64 = 1e-6; + +pub(crate) enum VisionTransformer { + Qwen(QwenVisionTransformer), + Gemma(GemmaVisionTransformer), +} + +impl VisionTransformer +where + F: CastTo + CastTensor + fusor::FloatOps + fusor::MatmulImpl + Default, + f32: CastTo + CastTensor, + fusor::MulOp: fusor::SimdBinaryOp, + fusor::AddOp: fusor::SimdBinaryOp, + fusor::SumOp: fusor::SimdReduceOp, +{ + pub(crate) fn from_gguf( + vision_ct: GgufMetadata, + vision_bytes: &[u8], + device: &Device, + ) -> Result { + let projector_type = vision_ct + .metadata + .get("clip.vision.projector_type") + .and_then(|x| x.to_string().ok()) + .map(|x| x.to_string()); + + match projector_type.as_deref() { + Some("gemma4v") => Ok(Self::Gemma(GemmaVisionTransformer::from_gguf( + vision_ct, + vision_bytes, + device, + )?)), + _ => Ok(Self::Qwen(QwenVisionTransformer::from_gguf( + vision_ct, + vision_bytes, + device, + )?)), + } + } + + pub(crate) fn preprocess_image( + &self, + image: &image::DynamicImage, + min_pixels: Option, + max_pixels: Option, + ) -> Result<(Tensor<2, f32>, [u32; 3])> { + match self { + Self::Qwen(vision) => vision.preprocess_image(image, min_pixels, max_pixels), + Self::Gemma(vision) => vision.preprocess_image(image, min_pixels, max_pixels), + } + } + + pub(crate) fn image_token_count(&self, grid: [u32; 3]) -> u32 { + match self { + Self::Qwen(vision) => { + grid.iter().product::() / (vision.spacial_merge_size as u32).pow(2) + } + Self::Gemma(vision) => vision.image_token_count(grid), + } + } + + pub(crate) fn expand_image_tokens( + &self, + raw_tokens: &[u32], + image_pad_token: u32, + vision_start_token: Option, + image_start_token: Option, + image_end_token: Option, + grid_thw: &[[u32; 3]], + ) -> Result<(Vec, Vec>)> { + match self { + Self::Qwen(_) => { + let Some(vision_start_token) = vision_start_token else { + return Ok((raw_tokens.to_vec(), Vec::new())); + }; + let mut tokens = Vec::new(); + let mut token_iter = raw_tokens.iter().copied(); + let mut image_iter = grid_thw.iter(); + let mut image_token_ranges = Vec::new(); + while let Some(token) = token_iter.next() { + tokens.push(token); + let start_index = tokens.len(); + if token == vision_start_token { + match token_iter.next() { + Some(next) if next == image_pad_token => { + let grid = *image_iter.next().ok_or_else(|| { + fusor::Error::msg( + "Image pad token found without matching image.", + ) + })?; + for _ in 0..self.image_token_count(grid) { + tokens.push(image_pad_token); + } + image_token_ranges.push(start_index..tokens.len()); + } + Some(next) => { + tokens.push(next); + } + None => break, + } + } + } + Ok((tokens, image_token_ranges)) + } + Self::Gemma(_) => { + let mut tokens = Vec::new(); + let mut image_iter = grid_thw.iter(); + let mut image_token_ranges = Vec::new(); + for token in raw_tokens.iter().copied() { + if token == image_pad_token { + let grid = *image_iter.next().ok_or_else(|| { + fusor::Error::msg("Image token found without matching image.") + })?; + if let Some(image_start_token) = image_start_token { + tokens.push(image_start_token); + } + let start_index = tokens.len(); + for _ in 0..self.image_token_count(grid) { + tokens.push(image_pad_token); + } + image_token_ranges.push(start_index..tokens.len()); + if let Some(image_end_token) = image_end_token { + tokens.push(image_end_token); + } + } else { + tokens.push(token); + } + } + Ok((tokens, image_token_ranges)) + } + } + } + + pub(crate) fn get_rope_index( + &self, + input_ids: &[u32], + grid_thw: &[[u32; 3]], + config: &crate::raw::LlamaConfig, + start_time: u32, + ) -> Result, u32)>> { + match self { + Self::Qwen(vision) => vision + .get_rope_index(input_ids, grid_thw, config, start_time) + .map(Some), + Self::Gemma(_) => Ok(None), + } + } + + pub(crate) fn forward_image( + &self, + pixels: &Tensor<2, F>, + grid: [u32; 3], + ) -> Result> { + match self { + Self::Qwen(vision) => vision.forward_image(pixels, grid), + Self::Gemma(vision) => vision.forward_image(pixels, grid), + } + } + + pub(crate) fn outputs_on_isolated_device(&self) -> bool { + match self { + Self::Qwen(_) => false, + Self::Gemma(vision) => vision.uses_isolated_device(), + } + } +} diff --git a/models/kalosm-llama/src/source.rs b/models/kalosm-llama/src/source.rs index 610990eee..8c23fb820 100644 --- a/models/kalosm-llama/src/source.rs +++ b/models/kalosm-llama/src/source.rs @@ -56,6 +56,7 @@ pub(crate) struct LlamaConfigJson { pub struct LlamaSource { pub(crate) model: Vec, pub(crate) vision_model: Option, + pub(crate) mtp_model: Option, pub(crate) tokenizer: Option, pub(crate) config: Option, pub(crate) group_query_attention: u8, @@ -125,6 +126,7 @@ impl LlamaSource { override_stop_token_string: None, override_chat_template: None, vision_model: None, + mtp_model: None, } } @@ -139,6 +141,7 @@ impl LlamaSource { override_stop_token_string: None, override_chat_template: None, vision_model: None, + mtp_model: None, } } @@ -223,6 +226,13 @@ impl LlamaSource { self } + /// Set the Gemma4 MTP assistant model to use for opt-in speculative decoding. + pub fn with_mtp_model(mut self, model: FileSource) -> Self { + self.mtp_model = Some(model); + + self + } + #[cfg(not(target_arch = "wasm32"))] pub(crate) async fn model( &self, @@ -976,6 +986,28 @@ impl LlamaSource { .with_override_stop_token_string("") } + /// A preset for Unsloth's Gemma 4 E2B instruction QAT GGUF. + /// + /// Note: The gemma model series does not support system prompts. + pub fn gemma_4_e2b_it_qat_chat() -> Self { + Self::new(FileSource::huggingface( + "unsloth/gemma-4-E2B-it-qat-GGUF".to_string(), + "main".to_string(), + "gemma-4-E2B-it-qat-UD-Q4_K_XL.gguf".to_string(), + )) + .with_mtp_model(FileSource::huggingface( + "unsloth/gemma-4-E2B-it-qat-GGUF".to_string(), + "main".to_string(), + "MTP/gemma-4-E2B-it-Q4_0-MTP.gguf".to_string(), + )) + .with_vision_model(FileSource::huggingface( + "unsloth/gemma-4-E2B-it-qat-GGUF".to_string(), + "main".to_string(), + "mmproj-F16.gguf".to_string(), + )) + .with_override_stop_token_string("") + } + /// A preset for qwen 2.5 3b VL chat in f16 precision pub fn qwen_2_5_3b_vl_chat_f16() -> Self { Self::new(kalosm_model_types::FileSource::HuggingFace { diff --git a/models/kalosm-tokenizer/src/lib.rs b/models/kalosm-tokenizer/src/lib.rs index a42dd0671..ee72fbfed 100644 --- a/models/kalosm-tokenizer/src/lib.rs +++ b/models/kalosm-tokenizer/src/lib.rs @@ -57,6 +57,7 @@ pub struct FastBpe { byte_to_token: [u32; 256], all_bytes_present: bool, byte_token_mapping: ByteTokenMapping, + raw_utf8_initial_tokens: bool, token_bytes: FxHashMap, u32>, id_to_bytes: Vec>>, merges: FxHashMap, @@ -71,25 +72,76 @@ impl FastBpe { vocab: impl IntoIterator, merges: impl IntoIterator, ignore_merges: bool, + ) -> Result { + Self::from_vocab_and_merges_with_decoder( + vocab, + merges, + ignore_merges, + TokenByteDecoder::ByteLevel, + ) + } + + /// Build a BPE tokenizer whose vocabulary tokens are raw UTF-8 strings. + /// + /// Some GGUF tokenizers, including Gemma 4, are BPE tokenizers but do not + /// use the GPT-2 byte-level character mapping. Their byte fallback tokens + /// are encoded as `<0xNN>` strings in the vocabulary. + pub fn from_raw_utf8_vocab_and_merges( + vocab: impl IntoIterator, + merges: impl IntoIterator, + ignore_merges: bool, + ) -> Result { + Self::from_vocab_and_merges_with_decoder( + vocab, + merges, + ignore_merges, + TokenByteDecoder::RawUtf8WithHexBytes, + ) + } + + fn from_vocab_and_merges_with_decoder( + vocab: impl IntoIterator, + merges: impl IntoIterator, + ignore_merges: bool, + decoder: TokenByteDecoder, ) -> Result { let mut byte_to_token = [MISSING_BYTE_TOKEN; 256]; let mut token_bytes = FxHashMap::default(); + let mut token_byte_preference = FxHashMap::default(); + let mut exact_token_ids = FxHashMap::default(); let mut id_to_bytes = Vec::new(); + let mut vocab = vocab.into_iter().collect::>(); + vocab.sort_unstable_by_key(|(_, id)| *id); for (token, id) in vocab { - let bytes = decode_token_bytes(&token); + let bytes = decoder.decode_token_bytes(&token); + let is_byte_fallback = decoder.is_byte_fallback_token(&token); if bytes.len() == 1 { if id == MISSING_BYTE_TOKEN { return Err(TokenizerError::ReservedByteTokenId(id)); } - byte_to_token[bytes[0] as usize] = id; + if decoder == TokenByteDecoder::RawUtf8WithHexBytes { + if is_byte_fallback { + byte_to_token[bytes[0] as usize] = id; + } + } else { + byte_to_token[bytes[0] as usize] = id; + } } let id = id as usize; if id >= id_to_bytes.len() { id_to_bytes.resize(id + 1, None); } id_to_bytes[id] = Some(bytes.clone()); - token_bytes.insert(bytes, id as u32); + exact_token_ids.insert(token, id as u32); + let preference = TokenBytePreference::new(is_byte_fallback, id as u32); + if token_byte_preference + .get(&bytes) + .is_none_or(|existing: &TokenBytePreference| preference < *existing) + { + token_byte_preference.insert(bytes.clone(), preference); + token_bytes.insert(bytes, id as u32); + } } let all_bytes_present = byte_to_token @@ -100,7 +152,9 @@ impl FastBpe { let raw_merges = merges .into_iter() .enumerate() - .map(|(rank, merge)| parse_merge(rank as u32, &merge, &token_bytes)) + .map(|(rank, merge)| { + parse_merge(rank as u32, &merge, &token_bytes, &exact_token_ids, decoder) + }) .collect::, _>>()?; let merge_levels = assign_merge_levels(&raw_merges); let level_count = merge_levels @@ -154,6 +208,7 @@ impl FastBpe { byte_to_token, all_bytes_present, byte_token_mapping, + raw_utf8_initial_tokens: decoder == TokenByteDecoder::RawUtf8WithHexBytes, token_bytes, id_to_bytes, merges: merges_by_pair, @@ -271,6 +326,10 @@ impl FastBpe { } fn encode_bytes_into(&self, input: &[u8], out: &mut Vec) -> Result<(), TokenizerError> { + if self.raw_utf8_initial_tokens { + return self.encode_raw_utf8_chars_into(input, out); + } + out.clear(); out.resize(input.len(), 0); @@ -283,6 +342,28 @@ impl FastBpe { Ok(()) } + fn encode_raw_utf8_chars_into( + &self, + input: &[u8], + out: &mut Vec, + ) -> Result<(), TokenizerError> { + out.clear(); + let text = std::str::from_utf8(input).map_err(|_| TokenizerError::InvalidUtf8)?; + out.reserve(text.len()); + for ch in text.chars() { + let mut buf = [0; 4]; + let bytes = ch.encode_utf8(&mut buf).as_bytes(); + if let Some(token) = self.token_bytes.get(bytes) { + out.push(*token); + } else { + for byte in bytes { + out.push(lookup_byte_token(&self.byte_to_token, *byte)?); + } + } + } + Ok(()) + } + fn encode_bytes_and_apply_single_merge( &self, input: &[u8], @@ -291,6 +372,12 @@ impl FastBpe { right: u32, new_token: u32, ) -> Result<(), TokenizerError> { + if self.raw_utf8_initial_tokens { + self.encode_raw_utf8_chars_into(input, out)?; + apply_single_merge(out, left, right, new_token); + return Ok(()); + } + if self.all_bytes_present { self.encode_bytes_and_apply_single_merge_unchecked(input, out, left, right, new_token); Ok(()) @@ -457,6 +544,9 @@ pub enum TokenizerError { /// A byte was missing from the vocabulary. #[error("byte 0x{0:02x} is missing from the vocabulary")] MissingByteToken(u8), + /// Raw UTF-8 tokenization received invalid UTF-8. + #[error("raw UTF-8 input is invalid")] + InvalidUtf8, } #[inline(always)] @@ -625,27 +715,52 @@ struct RawMerge { new_token: u32, } +#[derive(Clone, Copy, Debug, Eq, PartialEq, Ord, PartialOrd)] +struct TokenBytePreference { + byte_fallback_rank: u8, + id: u32, +} + +impl TokenBytePreference { + fn new(is_byte_fallback: bool, id: u32) -> Self { + Self { + byte_fallback_rank: u8::from(is_byte_fallback), + id, + } + } +} + fn parse_merge( rank: u32, merge: &str, token_bytes: &FxHashMap, u32>, + exact_token_ids: &FxHashMap, + decoder: TokenByteDecoder, ) -> Result { - let (left, right) = merge - .split_once(' ') + let split_at = merge + .char_indices() + .skip(1) + .find_map(|(index, ch)| (ch == ' ').then_some(index)) .ok_or_else(|| TokenizerError::InvalidMerge(merge.to_owned()))?; - let left_bytes = decode_token_bytes(left); - let right_bytes = decode_token_bytes(right); + let left = &merge[..split_at]; + let right = &merge[split_at + 1..]; + let left_bytes = decoder.decode_token_bytes(left); + let right_bytes = decoder.decode_token_bytes(right); let new_bytes = left_bytes .iter() .chain(right_bytes.iter()) .copied() .collect::>(); - let left = *token_bytes - .get(&left_bytes) + let left = exact_token_ids + .get(left) + .or_else(|| token_bytes.get(&left_bytes)) + .copied() .ok_or_else(|| TokenizerError::MissingToken(left.to_owned()))?; - let right = *token_bytes - .get(&right_bytes) + let right = exact_token_ids + .get(right) + .or_else(|| token_bytes.get(&right_bytes)) + .copied() .ok_or_else(|| TokenizerError::MissingToken(right.to_owned()))?; let new_token = *token_bytes .get(&new_bytes) @@ -665,6 +780,25 @@ fn apply_greedy_merges(merges: &FxHashMap, tokens: &mut Vec< } } +fn apply_single_merge(tokens: &mut Vec, left: u32, right: u32, new_token: u32) { + let len = tokens.len(); + let mut read = 0; + let mut write = 0; + + while read < len { + if read + 1 < len && tokens[read] == left && tokens[read + 1] == right { + tokens[write] = new_token; + read += 2; + } else { + tokens[write] = tokens[read]; + read += 1; + } + write += 1; + } + + tokens.truncate(write); +} + fn rebuild_merge_candidates( merges: &MergeLookup, tokens: &[u32], @@ -1381,6 +1515,59 @@ fn decode_token_bytes(token: &str) -> Vec { bytes } +#[derive(Clone, Copy, Eq, PartialEq)] +enum TokenByteDecoder { + ByteLevel, + RawUtf8WithHexBytes, +} + +impl TokenByteDecoder { + fn decode_token_bytes(self, token: &str) -> Vec { + match self { + Self::ByteLevel => decode_token_bytes(token), + Self::RawUtf8WithHexBytes => decode_raw_utf8_token_bytes(token), + } + } + + fn is_byte_fallback_token(self, token: &str) -> bool { + self == Self::RawUtf8WithHexBytes && decode_raw_utf8_byte_fallback(token).is_some() + } +} + +fn decode_raw_utf8_token_bytes(token: &str) -> Vec { + if let Some(byte) = decode_raw_utf8_byte_fallback(token) { + return vec![byte]; + } + + token.as_bytes().to_vec() +} + +fn decode_raw_utf8_byte_fallback(token: &str) -> Option { + if token.len() == 6 + && token.as_bytes()[0] == b'<' + && token.as_bytes()[1] == b'0' + && matches!(token.as_bytes()[2], b'x' | b'X') + && token.as_bytes()[5] == b'>' + { + if let (Some(high), Some(low)) = ( + hex_value(token.as_bytes()[3]), + hex_value(token.as_bytes()[4]), + ) { + return Some((high << 4) | low); + } + } + None +} + +fn hex_value(byte: u8) -> Option { + match byte { + b'0'..=b'9' => Some(byte - b'0'), + b'a'..=b'f' => Some(byte - b'a' + 10), + b'A'..=b'F' => Some(byte - b'A' + 10), + _ => None, + } +} + fn byte_level_char_to_byte(ch: char) -> Option { let codepoint = ch as u32; if (33..=126).contains(&codepoint) @@ -1474,6 +1661,46 @@ mod tests { assert_eq!(decode_token_bytes("Ā"), &[0]); } + #[test] + fn raw_utf8_hex_byte_tokens_decode_to_original_bytes() { + let vocab = [("<0x20>", 0), ("<0xC3>", 1), ("<0xA9>", 2), ("é", 3)] + .into_iter() + .map(|(token, id)| (token.to_owned(), id)); + let tokenizer = + FastBpe::from_raw_utf8_vocab_and_merges(vocab, ["<0xC3> <0xA9>".to_owned()], false) + .unwrap(); + + assert_eq!(tokenizer.tokenize(" é".as_bytes()).unwrap(), vec![0, 3]); + assert_eq!(tokenizer.token_bytes(0), Some(b" ".as_slice())); + assert_eq!(tokenizer.token_bytes(3), Some("é".as_bytes())); + } + + #[test] + fn raw_utf8_single_merge_fast_path_keeps_direct_char_tokens() { + let vocab = [ + ("<0x20>", 0), + ("<0xC3>", 1), + ("<0xA9>", 2), + ("é", 3), + ("a", 4), + ("b", 5), + ("ab", 6), + ] + .into_iter() + .map(|(token, id)| (token.to_owned(), id)); + let tokenizer = + FastBpe::from_raw_utf8_vocab_and_merges(vocab, ["a b".to_owned()], false).unwrap(); + + assert_eq!( + tokenizer.tokenize(" éab".as_bytes()).unwrap(), + vec![0, 3, 6] + ); + assert_eq!( + tokenizer.tokenize(" éab".as_bytes()).unwrap(), + tokenizer.tokenize_reference(" éab".as_bytes()).unwrap() + ); + } + #[test] fn levelized_tokenization_matches_reference_for_long_input() { let tokenizer = small_tokenizer(false);