diff --git a/crates/hipfire-arch-gemma4/map.md b/crates/hipfire-arch-gemma4/map.md index b96c398e9..f7ae9e3dc 100644 --- a/crates/hipfire-arch-gemma4/map.md +++ b/crates/hipfire-arch-gemma4/map.md @@ -31,7 +31,7 @@ _Generated by `scripts/check-crate-maps.py` from the tree — do not edit inside | [`src/gemma4.rs`](src/gemma4.rs) | 1,088 | 13 | 0 | | [`src/gemma4_vision.rs`](src/gemma4_vision.rs) | 16 | 3 | 0 | | [`src/lib.rs`](src/lib.rs) | 51 | 8 | 0 | -| [`src/lowered.rs`](src/lowered.rs) | 7,427 | 41 | 8 | +| [`src/lowered.rs`](src/lowered.rs) | 7,457 | 41 | 9 | | [`src/speculative.rs`](src/speculative.rs) | 252 | 6 | 0 | ### Public API surface @@ -60,6 +60,6 @@ _Generated by `scripts/check-crate-maps.py` from the tree — do not edit inside ### Totals -- 10 modules · 13,923 lines · 120 public items · 31 tests · 8 examples +- 10 modules · 13,953 lines · 120 public items · 32 tests · 8 examples diff --git a/crates/hipfire-arch-gemma4/src/lowered.rs b/crates/hipfire-arch-gemma4/src/lowered.rs index ca9c1ae80..e046a45d3 100644 --- a/crates/hipfire-arch-gemma4/src/lowered.rs +++ b/crates/hipfire-arch-gemma4/src/lowered.rs @@ -2179,12 +2179,21 @@ pub const GEMMA4_FLASH_TILE: usize = 128; /// scale linearly with this. pub const GEMMA4_MAX_PREFILL_BATCH: usize = 128; +#[inline] +fn flash_partials_len_for_tile( + max_seq: usize, + n_heads: usize, + head_dim: usize, + tile_size: usize, +) -> usize { + n_heads * max_seq.div_ceil(tile_size) * (2 + head_dim) +} + /// Pure geometry: single-query flash partial length for `max_seq`. /// `n_heads * ceil(max_seq / TILE) * (2 + head_dim)` floats. #[inline] pub fn gemma4_flash_partials_len(max_seq: usize, n_heads: usize, full_head_dim: usize) -> usize { - let tiles = max_seq.div_ceil(GEMMA4_FLASH_TILE); - n_heads * tiles * (2 + full_head_dim) + flash_partials_len_for_tile(max_seq, n_heads, full_head_dim, GEMMA4_FLASH_TILE) } /// Pure geometry: batched flash partial length for `max_seq`. @@ -2443,13 +2452,24 @@ impl Gemma4Scratch { let i_sample_buf = alloc!(&[2], DType::F32); let i_repeat_buf = alloc!(&[1024], DType::F32); - // Flash partials sizing. Per-head × max_tiles × (2 + head_dim) floats. - // Sized for FULL attn (head_dim=512 stride 514, vs sliding 256 stride 258); - // sliding-layer dispatches use part of the buffer, full-layer dispatches - // use all of it. `max_seq` is the single authority shared with both KV - // caches — no independent `HIPFIRE_KV_SEQ` env var. - let flash_partials_sz = + // gfx1100 sliding Q8 uses a smaller tile than full Asym3 attention, so + // its larger tile count can outweigh the smaller per-tile stride. + let sliding_tile = rdna_compute::attention::q8_flash_tile_size( + &gpu.arch, + config.n_heads, + config.sliding_n_kv_heads, + config.sliding_head_dim, + max_seq, + ); + let sliding_flash_partials_sz = flash_partials_len_for_tile( + max_seq, + config.n_heads, + config.sliding_head_dim, + sliding_tile, + ); + let full_flash_partials_sz = gemma4_flash_partials_len(max_seq, config.n_heads, config.full_head_dim); + let flash_partials_sz = sliding_flash_partials_sz.max(full_flash_partials_sz); let i_flash_partials = alloc!(&[flash_partials_sz], DType::F32); // (Note 2026-05-19): removed the precomputed sliding/full cos+sin @@ -2671,6 +2691,16 @@ impl Gemma4Scratch { mod scratch_geometry_tests { use super::*; + #[test] + fn sliding_tile32_partials_can_exceed_full_tile128_partials() { + let max_seq = 1025; + let sliding = flash_partials_len_for_tile(max_seq, 16, 256, 32); + let full = flash_partials_len_for_tile(max_seq, 16, 512, 128); + assert!(sliding > full); + assert_eq!(sliding, 16 * 33 * 258); + assert_eq!(full, 16 * 9 * 514); + } + fn dummy_cfg_31b() -> Gemma4Config { // Minimal config mirroring 31B/26B shapes: n_heads=32, full_head_dim=512 Gemma4Config { diff --git a/crates/hipfire-dispatch/map.md b/crates/hipfire-dispatch/map.md index 9d480b1df..02e548c71 100644 --- a/crates/hipfire-dispatch/map.md +++ b/crates/hipfire-dispatch/map.md @@ -24,7 +24,7 @@ _Generated by `scripts/check-crate-maps.py` from the tree — do not edit inside |---|---:|---:|---:| | [`src/context.rs`](src/context.rs) | 67 | 5 | 0 | | [`src/coverage_tests.rs`](src/coverage_tests.rs) | 1,841 | 0 | 21 | -| [`src/families/attention.rs`](src/families/attention.rs) | 2,463 | 9 | 10 | +| [`src/families/attention.rs`](src/families/attention.rs) | 2,471 | 9 | 10 | | [`src/families/fused_qkv.rs`](src/families/fused_qkv.rs) | 1,579 | 8 | 1 | | [`src/families/gemm.rs`](src/families/gemm.rs) | 647 | 7 | 2 | | [`src/families/gemv.rs`](src/families/gemv.rs) | 608 | 20 | 0 | @@ -105,6 +105,6 @@ _Generated by `scripts/check-crate-maps.py` from the tree — do not edit inside ### Totals -- 33 modules · 22,681 lines · 223 public items · 242 tests · 0 examples +- 33 modules · 22,689 lines · 223 public items · 242 tests · 0 examples diff --git a/crates/hipfire-dispatch/src/families/attention.rs b/crates/hipfire-dispatch/src/families/attention.rs index f7fcad484..57a9ecb26 100644 --- a/crates/hipfire-dispatch/src/families/attention.rs +++ b/crates/hipfire-dispatch/src/families/attention.rs @@ -1187,6 +1187,14 @@ fn dispatch_attend( let st = io.givens_sin.unwrap(); let fp = io.flash_partials.unwrap(); if io.head_dim == 512 { + if io.output_gate.is_some() { + return Err(DispatchError::UnsupportedVariant { + family: "attention/attend", + variant: "D512 Asym3 output gate is unsupported", + arch: "", + quant: "asym3", + }); + } return hip!(gpu.attention_flash_asym3_hd512( io.q, io.k_cache,