From 7671c9ae2b804f12f102c3480a00039b2b2fb73d Mon Sep 17 00:00:00 2001 From: alpineq Date: Fri, 4 Sep 2026 18:39:41 +0300 Subject: [PATCH] FIX(qwen35): chunk dense-TP prefill at the rank's prefill batch scaled by the TP degree --- crates/hipfire-arch-qwen35/map.md | 10 ++++---- crates/hipfire-arch-qwen35/src/qwen35.rs | 3 ++- .../hipfire-arch-qwen35/src/qwen35/forward.rs | 5 +++- .../hipfire-arch-qwen35/src/qwen35/prefill.rs | 20 +++++++++++++++ crates/hipfire-generate/map.md | 4 +-- crates/hipfire-generate/src/qwen.rs | 25 +++++++++++++++++-- 6 files changed, 56 insertions(+), 11 deletions(-) diff --git a/crates/hipfire-arch-qwen35/map.md b/crates/hipfire-arch-qwen35/map.md index 94a8f3efd..ceff02d19 100644 --- a/crates/hipfire-arch-qwen35/map.md +++ b/crates/hipfire-arch-qwen35/map.md @@ -40,11 +40,11 @@ _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,258 | 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/prefill.rs`](src/qwen35/prefill.rs) | 9,332 | 12 | 48 | | [`src/qwen35/weights.rs`](src/qwen35/weights.rs) | 1,971 | 43 | 10 | -| [`src/qwen35.rs`](src/qwen35.rs) | 63 | 7 | 0 | +| [`src/qwen35.rs`](src/qwen35.rs) | 64 | 7 | 0 | | [`src/scheduler.rs`](src/scheduler.rs) | 142 | 3 | 4 | | [`src/serve_engine.rs`](src/serve_engine.rs) | 1,273 | 8 | 2 | | [`src/slot_batch.rs`](src/slot_batch.rs) | 123 | 4 | 6 | @@ -74,7 +74,7 @@ _Generated by `scripts/check-crate-maps.py` from the tree — do not edit inside - [`src/qwen35/ep_batch.rs`](src/qwen35/ep_batch.rs): `validate_ep_batch_compatibility`, `Qwen35DecodeBatchEpState`, `max_batch`, `lane_capacity`, `epoch`, `poison_mask`, `lane_state`, `new`, `reset_all`, `reset_lane`, `prefill_lane`, `forward_tick`, +8 more - [`src/qwen35/forward.rs`](src/qwen35/forward.rs): `dump_expert_stats`, `forward`, `Qwen35Scratch`, `new`, `new_with_kv_max`, `free_gpu`, `Qwen35ScratchSet`, `new_with_kv_max_multi`, `free_gpu_multi`, `forward_scratch`, `prepare_scratch_inputs`, `forward_scratch_with_hidden`, +19 more - [`src/qwen35/load.rs`](src/qwen35/load.rs): `hipfire_runtime`, `load_weights`, `HfqSource`, `new`, `ParoSource`, `preflight_weights_dense_tp`, `load_weights_dense_tp_rank`, `set_ep_expert_shard`, `EpShardGuard`, `load_weights_ep_rank` -- [`src/qwen35/prefill.rs`](src/qwen35/prefill.rs): `PREFILL_MAX_BATCH`, `prefill_max_batch`, `upload_prefill_batch_inputs`, `forward_prefill_batch_single_chunk_captured`, `forward_prefill_batch_single_chunk_captured_opts`, `forward_prefill_batch`, `forward_prefill_batch_capped`, `forward_prefill_batch_with_pbs`, `forward_prefill_batch_with_pbs_opts`, `qwen35_layer_batch_admissible`, `prefill_batch_pbs_eligible` +- [`src/qwen35/prefill.rs`](src/qwen35/prefill.rs): `PREFILL_MAX_BATCH`, `prefill_max_batch`, `prefill_max_batch_tp`, `upload_prefill_batch_inputs`, `forward_prefill_batch_single_chunk_captured`, `forward_prefill_batch_single_chunk_captured_opts`, `forward_prefill_batch`, `forward_prefill_batch_capped`, `forward_prefill_batch_with_pbs`, `forward_prefill_batch_with_pbs_opts`, `qwen35_layer_batch_admissible`, `prefill_batch_pbs_eligible` - [`src/qwen35/weights.rs`](src/qwen35/weights.rs): `DeltaNetLayerWeights`, `FullAttnLayerWeights`, `ExpertWeights`, `mixed_expert_tag`, `SharedExpertWeights`, `MoeFfnWeights`, `MoeParoSidecars`, `DeltaNetMoeLayerWeights`, `FullAttnMoeLayerWeights`, `LayerWeights`, `Qwen35HfqSourceIdentity`, `capture`, +31 more - [`src/qwen35.rs`](src/qwen35.rs): `batch`, `config`, `ep_batch`, `forward`, `load`, `prefill`, `weights` - [`src/scheduler.rs`](src/scheduler.rs): `Scheduler`, `PendingWork`, `next_batch` @@ -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,305 lines · 437 public items · 189 tests · 4 examples diff --git a/crates/hipfire-arch-qwen35/src/qwen35.rs b/crates/hipfire-arch-qwen35/src/qwen35.rs index 31fda38b1..ee4530a2a 100644 --- a/crates/hipfire-arch-qwen35/src/qwen35.rs +++ b/crates/hipfire-arch-qwen35/src/qwen35.rs @@ -46,7 +46,8 @@ pub use prefill::{ forward_prefill_batch, forward_prefill_batch_capped, forward_prefill_batch_single_chunk_captured, forward_prefill_batch_single_chunk_captured_opts, forward_prefill_batch_with_pbs, forward_prefill_batch_with_pbs_opts, - prefill_batch_pbs_eligible, prefill_max_batch, qwen35_layer_batch_admissible, + prefill_batch_pbs_eligible, prefill_max_batch, prefill_max_batch_tp, + qwen35_layer_batch_admissible, upload_prefill_batch_inputs, PREFILL_MAX_BATCH, }; pub(crate) use prefill::{ diff --git a/crates/hipfire-arch-qwen35/src/qwen35/forward.rs b/crates/hipfire-arch-qwen35/src/qwen35/forward.rs index 00c86bd15..016f9c588 100644 --- a/crates/hipfire-arch-qwen35/src/qwen35/forward.rs +++ b/crates/hipfire-arch-qwen35/src/qwen35/forward.rs @@ -4308,7 +4308,10 @@ pub fn forward_prefill_dense_tp( _ => return Err(HipError::new(0, "dense TP received a MoE/mismatched layer")), } } - let cap = crate::qwen35::prefill::prefill_max_batch(&gpus.devices[0]); + // Size the scratch to the call, not to the arch ceiling: an 8-token verify + // would otherwise allocate and free a 512-row PBS per rank on every cycle. + let cap = crate::qwen35::prefill::prefill_max_batch_tp(&gpus.devices[0], tp) + .min(tokens.len().max(1)); if cap == 0 { return Err(HipError::new(0, "prefill_max_batch is zero")); } diff --git a/crates/hipfire-arch-qwen35/src/qwen35/prefill.rs b/crates/hipfire-arch-qwen35/src/qwen35/prefill.rs index 418e23b44..f72154781 100644 --- a/crates/hipfire-arch-qwen35/src/qwen35/prefill.rs +++ b/crates/hipfire-arch-qwen35/src/qwen35/prefill.rs @@ -420,6 +420,26 @@ pub fn prefill_max_batch(gpu: &Gpu) -> usize { explicit_prefill_max_batch().unwrap_or_else(|| prefill_max_batch_for_arch(gpu.arch.as_str())) } +/// Ceiling on the TP compensation below; beyond it the prefill kernels stop +/// gaining and the per-chunk scratch keeps growing. +const PREFILL_TP_BATCH_CAP: usize = 2048; + +/// Prefill chunk for one rank of a `tp`-way split. +/// +/// The prefill attention grid is `local_heads x chunk / M_TILE`, and TP divides +/// the heads, so a rank running the arch default launches `tp` times fewer +/// workgroups than a single card and starves an already latency-bound kernel. +/// Scaling the chunk by `tp` restores the single-card workgroup count. +/// Measured on gfx1100, Qwen3.8-27B, 33k prompt, tp=2: 649 -> 773 tok/s with +/// byte-identical output. An explicit `HIPFIRE_PREFILL_MAX_BATCH` wins. +pub fn prefill_max_batch_tp(gpu: &Gpu, tp: usize) -> usize { + if let Some(explicit) = explicit_prefill_max_batch() { + return explicit; + } + let base = prefill_max_batch_for_arch(gpu.arch.as_str()); + base.saturating_mul(tp.max(1)).min(PREFILL_TP_BATCH_CAP) +} + /// Effective per-chunk capacity for one prefill call. /// /// Never form a chunk larger than the configured/capped max, the PBS diff --git a/crates/hipfire-generate/map.md b/crates/hipfire-generate/map.md index f59f344a3..8f6066335 100644 --- a/crates/hipfire-generate/map.md +++ b/crates/hipfire-generate/map.md @@ -27,7 +27,7 @@ _Generated by `scripts/check-crate-maps.py` from the tree — do not edit inside | [`src/common.rs`](src/common.rs) | 1,569 | 48 | 1 | | [`src/dense.rs`](src/dense.rs) | 8,455 | 91 | 4 | | [`src/lib.rs`](src/lib.rs) | 58 | 7 | 0 | -| [`src/qwen.rs`](src/qwen.rs) | 6,312 | 60 | 1 | +| [`src/qwen.rs`](src/qwen.rs) | 6,333 | 60 | 1 | | [`src/redline.rs`](src/redline.rs) | 4,361 | 49 | 1 | | [`src/vision.rs`](src/vision.rs) | 2,942 | 9 | 8 | @@ -55,6 +55,6 @@ _Generated by `scripts/check-crate-maps.py` from the tree — do not edit inside ### Totals -- 8 modules · 31,785 lines · 321 public items · 224 tests · 0 examples +- 8 modules · 31,806 lines · 321 public items · 224 tests · 0 examples diff --git a/crates/hipfire-generate/src/qwen.rs b/crates/hipfire-generate/src/qwen.rs index 5bea5b5ba..82e84701e 100644 --- a/crates/hipfire-generate/src/qwen.rs +++ b/crates/hipfire-generate/src/qwen.rs @@ -361,7 +361,28 @@ pub fn ep_serve_qwen35_dense_tp( ); let t_prefill = Instant::now(); - for (chunk_index, chunk) in prompt_ids.chunks(32).enumerate() { + // Chunk at the rank's prefill batch, not a fixed 32: M=32 misses the MMQ + // batch floor (128) and fires the per-layer collectives 16x more often + // than the chunk the TP prefill re-chunks to internally (8k prompt on + // 2x gfx1100: 72.1 -> 54.7 s wall). + let tp_prefill_chunk = match m.ep.as_ref() { + Some(EpState { gpus, .. }) => { + qwen35::prefill_max_batch_tp(&gpus.devices[0], gpus.devices.len()).max(1) + } + None => { + emit_active_attempt_error( + stdout, + Some(id), + "dense TP serve without EP state", + "validation", + false, + false, + ); + let _ = stdout.flush(); + return; + } + }; + for (chunk_index, chunk) in prompt_ids.chunks(tp_prefill_chunk).enumerate() { if check_abort(id) { ep_emit_abort(stdout, id, m, 0); return; @@ -395,7 +416,7 @@ pub fn ep_serve_qwen35_dense_tp( weights, configs, chunk, - chunk_index * 32, + chunk_index * tp_prefill_chunk, kv_caches, dn_states, scratches,