From 45be9fa7634a17e0d564236cddd9aca7fa85bca2 Mon Sep 17 00:00:00 2001 From: alpineq Date: Fri, 4 Sep 2026 18:38:27 +0300 Subject: [PATCH] PERF(qwen35): run the MTP head on the trunk flash policy with partials sized for the launch tile --- crates/hipfire-arch-qwen35/map.md | 6 ++-- crates/hipfire-arch-qwen35/src/mtp_head.rs | 35 +++++++++++-------- .../hipfire-arch-qwen35/src/qwen35/forward.rs | 14 +------- crates/hipfire-runtime/map.md | 6 ++-- crates/hipfire-runtime/src/llama.rs | 25 +++++++------ 5 files changed, 42 insertions(+), 44 deletions(-) diff --git a/crates/hipfire-arch-qwen35/map.md b/crates/hipfire-arch-qwen35/map.md index 94a8f3efd..dd72af3bb 100644 --- a/crates/hipfire-arch-qwen35/map.md +++ b/crates/hipfire-arch-qwen35/map.md @@ -32,7 +32,7 @@ _Generated by `scripts/check-crate-maps.py` from the tree — do not edit inside | [`src/layer_driver.rs`](src/layer_driver.rs) | 112 | 0 | 0 | | [`src/lib.rs`](src/lib.rs) | 121 | 19 | 0 | | [`src/mtp_compose.rs`](src/mtp_compose.rs) | 1,374 | 8 | 0 | -| [`src/mtp_head.rs`](src/mtp_head.rs) | 2,611 | 32 | 2 | +| [`src/mtp_head.rs`](src/mtp_head.rs) | 2,616 | 32 | 2 | | [`src/mtp_probe.rs`](src/mtp_probe.rs) | 464 | 8 | 0 | | [`src/mtp_spec.rs`](src/mtp_spec.rs) | 3,883 | 33 | 12 | | [`src/mtp_speculator.rs`](src/mtp_speculator.rs) | 522 | 3 | 0 | @@ -40,7 +40,7 @@ _Generated by `scripts/check-crate-maps.py` from the tree — do not edit inside | [`src/qwen35/batch.rs`](src/qwen35/batch.rs) | 1,665 | 16 | 0 | | [`src/qwen35/config.rs`](src/qwen35/config.rs) | 1,630 | 40 | 21 | | [`src/qwen35/ep_batch.rs`](src/qwen35/ep_batch.rs) | 4,800 | 20 | 7 | -| [`src/qwen35/forward.rs`](src/qwen35/forward.rs) | 6,255 | 31 | 12 | +| [`src/qwen35/forward.rs`](src/qwen35/forward.rs) | 6,243 | 31 | 12 | | [`src/qwen35/load.rs`](src/qwen35/load.rs) | 4,906 | 10 | 0 | | [`src/qwen35/prefill.rs`](src/qwen35/prefill.rs) | 9,312 | 11 | 48 | | [`src/qwen35/weights.rs`](src/qwen35/weights.rs) | 1,971 | 43 | 10 | @@ -97,6 +97,6 @@ _Generated by `scripts/check-crate-maps.py` from the tree — do not edit inside ### Totals -- 29 modules · 57,281 lines · 436 public items · 189 tests · 4 examples +- 29 modules · 57,274 lines · 436 public items · 189 tests · 4 examples diff --git a/crates/hipfire-arch-qwen35/src/mtp_head.rs b/crates/hipfire-arch-qwen35/src/mtp_head.rs index 68e4eb798..5d3b873c2 100644 --- a/crates/hipfire-arch-qwen35/src/mtp_head.rs +++ b/crates/hipfire-arch-qwen35/src/mtp_head.rs @@ -487,11 +487,16 @@ impl Qwen35MtpHeadScratch { logits: gpu.alloc_tensor(&[config.vocab_size], DType::F32)?, logits_compressed: None, flash_partials: { - // Same sizing as trunk's prefill_partials at qwen35.rs:2822 - // (TILE_SIZE=128) but with batch_mult=1 since MTP forward - // is single-token. Allocated per scratch instance, lives - // for the lifetime of the slot. - let tile_size = 128usize; + // Sized with the tile the flash launch picks for this shape: a + // fixed 128 under-allocates 4x where the arch picks 32 (gfx1100 + // at max_seq <= 8192) and the tile kernel page-faults. + let tile_size = rdna_compute::attention::q8_flash_tile_size( + &gpu.arch, + config.n_head, + config.n_head_kv, + config.head_dim, + config.max_seq, + ); let max_tiles = (config.max_seq + tile_size - 1) / tile_size; gpu.alloc_tensor( &[config.n_head * max_tiles * (2 + config.head_dim)], @@ -1615,21 +1620,21 @@ pub fn mtp_head_forward_block_only_with_pos_buf( // per Phase 1 fwht4 commit `c64c0e3f`). // KV write + attention via the shared KV-usage abstraction. kv.inner is // built per kv_mode (new_gpu_q8/asym3/fwht4), so kv.inner.tier_inputs() - // produces exactly the tier kv.kv_mode used to dispatch: Q8→AttnQ8_0Kv - // (non-flash), Asym3→AttnFlashAsym3, Fwht4→AttnFlashAsym4Fwht — byte- - // identical kernels (incl. the Givens cos/sin + v_mode_bits sub-plan). The - // dispatch arm computes seq_len = pos+1, so pos = seq_len_hint-1 reproduces - // the hand seq_len_hint exactly (the write position flows via pos_buf). - // SPEC-DECODE: draft logits stay byte-identical → τ unchanged (validated by - // coherence-gate-dflash.sh + a τ A/B). flash_partials is always Some (the Q8 - // non-flash arm ignores it; asym3/fwht4 require it). Q8 non-flash is - // unconditional → derive returns AttnQ8_0Kv at seq_len_hint<=15000 (the - // documented >15k Q8-fidelity edge). + // produces the tier kv.kv_mode dispatches on — byte-identical kernels + // (incl. the Givens cos/sin + v_mode_bits sub-plan). The dispatch arm + // computes seq_len = pos+1, so pos = seq_len_hint-1 reproduces the hand + // seq_len_hint exactly (the write position flows via pos_buf). + // `tier_inputs()` reports flash_mode 0, which would pin the Q8 tier to the + // non-flash AttnQ8_0Kv at any context (10 ms per draft step at 33k against + // 0.3 ms on the flash tile), so the head takes the trunk's flash policy. + // flash_partials is always Some and sized with the same tile the launch + // picks (see `Qwen35MtpHeadScratch::new`). let dispatch_pos = seq_len_hint - 1; let ctx = hipfire_dispatch::context::DispatchCtx::new(gpu); let plan = hipfire_dispatch::families::kv_tier::KvTierPlan::derive( hipfire_dispatch::families::kv_tier::KvTierInputs { pos: dispatch_pos, + flash_mode: hipfire_runtime::llama::attention_flash_mode(&gpu.arch), ..kv.inner.tier_inputs() }, ) diff --git a/crates/hipfire-arch-qwen35/src/qwen35/forward.rs b/crates/hipfire-arch-qwen35/src/qwen35/forward.rs index 00c86bd15..b0542468b 100644 --- a/crates/hipfire-arch-qwen35/src/qwen35/forward.rs +++ b/crates/hipfire-arch-qwen35/src/qwen35/forward.rs @@ -1149,19 +1149,7 @@ impl Qwen35Scratch { // Honors HIPFIRE_ATTN_FLASH=never|0|off as an explicit override // for users who prefer the non-flash kernel and don't intend // to use graph capture. - flash_mode: match hipfire_runtime::config::get().attention_flash_mode.as_str() { - "never" | "0" | "off" => 0, - "always" | "2" | "force" => 2, - _ => { - let graph_capable_arch = - gpu.arch.starts_with("gfx12") || gpu.arch.starts_with("gfx11"); - if graph_capable_arch { - 2 - } else { - 1 - } - } - }, + flash_mode: hipfire_runtime::llama::attention_flash_mode(&gpu.arch) as u8, moe_router_logits: None, moe_scalar_buf: None, diff --git a/crates/hipfire-runtime/map.md b/crates/hipfire-runtime/map.md index 1995fd991..9426f9e96 100644 --- a/crates/hipfire-runtime/map.md +++ b/crates/hipfire-runtime/map.md @@ -53,7 +53,7 @@ _Generated by `scripts/check-crate-maps.py` from the tree — do not edit inside | [`src/kv_backend.rs`](src/kv_backend.rs) | 129 | 1 | 7 | | [`src/kv_mode.rs`](src/kv_mode.rs) | 298 | 10 | 7 | | [`src/lib.rs`](src/lib.rs) | 80 | 55 | 0 | -| [`src/llama.rs`](src/llama.rs) | 8,738 | 83 | 42 | +| [`src/llama.rs`](src/llama.rs) | 8,743 | 85 | 42 | | [`src/llama_spec.rs`](src/llama_spec.rs) | 617 | 6 | 1 | | [`src/loader_api.rs`](src/loader_api.rs) | 256 | 10 | 4 | | [`src/loop_guard.rs`](src/loop_guard.rs) | 194 | 8 | 4 | @@ -114,7 +114,7 @@ _Generated by `scripts/check-crate-maps.py` from the tree — do not edit inside - [`src/kv_backend.rs`](src/kv_backend.rs): `saddle_core` - [`src/kv_mode.rs`](src/kv_mode.rs): `saddle_core`, `KvModePolicy`, `ResolveResult`, `QWEN35_HFQ_POLICY`, `QWEN35_PARO_POLICY`, `DIR_SAFETENSORS_POLICY`, `LLAMA_HFQ_POLICY`, `HFQ_Q8_ONLY_POLICY`, `QWEN35_PP_POLICY`, `resolve` - [`src/lib.rs`](src/lib.rs): `admission`, `arch`, `arch_mapping`, `arch_model`, `arch_spec`, `augmentor`, `bf16_loader`, `cache_plan`, `cask`, `config`, `cpu_router`, `ddtree`, +43 more -- [`src/llama.rs`](src/llama.rs): `ModelArch`, `LlamaConfig`, `from_gguf`, `dequantize_q4_0`, `dequantize_q8_0`, `f16_to_f32`, `f32_to_f16`, `dequantize_q4_k`, `convert_q4k_to_q4f16_g64`, `convert_q4k_to_q4f16_g32`, `dequantize_q6_k`, `ParoRotation`, +71 more +- [`src/llama.rs`](src/llama.rs): `ModelArch`, `LlamaConfig`, `from_gguf`, `dequantize_q4_0`, `dequantize_q8_0`, `f16_to_f32`, `f32_to_f16`, `dequantize_q4_k`, `convert_q4k_to_q4f16_g64`, `convert_q4k_to_q4f16_g32`, `dequantize_q6_k`, `ParoRotation`, +73 more - [`src/llama_spec.rs`](src/llama_spec.rs): `verify_block_argmax`, `verify_block_logits`, `verify_block_argmax_capture_gpu`, `verify_block_sampled_capture_gpu`, `verify_tree_logits`, `lm_head_logits_n_rows` - [`src/loader_api.rs`](src/loader_api.rs): `ModelSource`, `from_path`, `arch_id`, `is_dir`, `describe`, `LoadCtx`, `SpecLoadCfg`, `CaskConfig`, `physical_cap`, `physical_cap_with_override` - [`src/loop_guard.rs`](src/loop_guard.rs): `StopReason`, `LoopGuard`, `from_config`, `new`, `off`, `enabled`, `check`, `window_len` @@ -156,6 +156,6 @@ _Generated by `scripts/check-crate-maps.py` from the tree — do not edit inside ### Totals -- 58 modules · 53,079 lines · 893 public items · 617 tests · 132 examples +- 58 modules · 53,084 lines · 895 public items · 617 tests · 132 examples diff --git a/crates/hipfire-runtime/src/llama.rs b/crates/hipfire-runtime/src/llama.rs index e9e909107..4bfeba41d 100644 --- a/crates/hipfire-runtime/src/llama.rs +++ b/crates/hipfire-runtime/src/llama.rs @@ -3758,8 +3758,12 @@ fn llama_forward_lowered_enabled() -> bool { }) } +/// Numeric attention flash policy for `HIPFIRE_ATTN_FLASH` (`config.attention_flash_mode`): +/// `0` never, `1` auto (flash at long context), `2` always. `auto` resolves to +/// `2` on graph-capable archs (gfx11/gfx12) so direct and captured forwards run +/// the same kernel, and to `1` elsewhere. #[inline] -fn llama_attention_flash_mode_for(mode: &str, gpu_arch: &str) -> usize { +pub fn attention_flash_mode_for(mode: &str, gpu_arch: &str) -> usize { match mode { "never" | "0" | "off" => 0, "always" | "2" | "force" => 2, @@ -3768,9 +3772,10 @@ fn llama_attention_flash_mode_for(mode: &str, gpu_arch: &str) -> usize { } } +/// [`attention_flash_mode_for`] against the process configuration. #[inline] -fn llama_attention_flash_mode(gpu_arch: &str) -> usize { - llama_attention_flash_mode_for(crate::config::get().attention_flash_mode.as_str(), gpu_arch) +pub fn attention_flash_mode(gpu_arch: &str) -> usize { + attention_flash_mode_for(crate::config::get().attention_flash_mode.as_str(), gpu_arch) } #[inline] @@ -3810,7 +3815,7 @@ fn llama_kv_write_attend( let ctx = DispatchCtx::new(gpu); let plan = KvTierPlan::derive(KvTierInputs { pos, - flash_mode: llama_attention_flash_mode(&gpu.arch), + flash_mode: attention_flash_mode(&gpu.arch), ..kv_cache.tier_inputs() }) .map_err(|e| hip_bridge::HipError::new(0, &e.to_string()))?; @@ -4064,7 +4069,7 @@ fn forward_scratch_layers_lowered( config, scratch, kv_cache: &*kv_cache, - flash_mode: llama_attention_flash_mode(&gpu.arch), + flash_mode: attention_flash_mode(&gpu.arch), knobs, pos, }; @@ -7901,11 +7906,11 @@ mod tests { #[test] fn qwen3_flash_mode_policy_matches_rdna_generation() { - assert_eq!(llama_attention_flash_mode_for("auto", "gfx1100"), 2); - assert_eq!(llama_attention_flash_mode_for("auto", "gfx1201"), 2); - assert_eq!(llama_attention_flash_mode_for("auto", "gfx1030"), 1); - assert_eq!(llama_attention_flash_mode_for("never", "gfx1100"), 0); - assert_eq!(llama_attention_flash_mode_for("always", "gfx1030"), 2); + assert_eq!(attention_flash_mode_for("auto", "gfx1100"), 2); + assert_eq!(attention_flash_mode_for("auto", "gfx1201"), 2); + assert_eq!(attention_flash_mode_for("auto", "gfx1030"), 1); + assert_eq!(attention_flash_mode_for("never", "gfx1100"), 0); + assert_eq!(attention_flash_mode_for("always", "gfx1030"), 2); } #[test]