perf(qwen35): run the MTP head on the trunk flash policy with partials sized for the launch tile

crates/hipfire-arch-qwen35/map.md

@@ -32,15 +32,15 @@ _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 |
 | [`src/paro_moe.rs`](src/paro_moe.rs) | 222 | 0 | 0 |
 | [`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
 
 <!-- crate-map:generated:end -->

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()
         },
     )

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,

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
 
 <!-- crate-map:generated:end -->

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]