ferrum_models/models/qwen3_moe/api.rs
1use super::*;
2
3impl<B: MoeLlmBackend + BackendPagedKv, K: KvDtypeKind> DecoderOnlyLLM for Qwen3MoeModel<B, K> {
4 fn config(&self) -> &LlmRuntimeConfig {
5 &self.runtime_cfg
6 }
7
8 fn prepare(&mut self, cache_id: &str, max_tokens: usize) {
9 // Eager scratch + KV cache grow + a 1-token forward warmup so
10 // the first real prefill / decode doesn't pay the cold-start
11 // ~25-MTLBuffer scratch alloc + ~96-MTLBuffer KV alloc + Metal
12 // pipeline-state first-bind costs (~265 ms total on Qwen3-MoE
13 // 30B-A3B / M1 Max). Mirrors what llama-bench's --warmup does
14 // (which runs a same-shape forward before the timer).
15 self.ensure_scratch(max_tokens);
16 self.ensure_kv(cache_id);
17
18 // Warmup forward through all 48 layers under a scratch cache_id
19 // so the real `cache_id` starts at pos_offset=0. Token 0 is
20 // valid for any tokenizer (BOS or pad).
21 const WARMUP_CACHE: &str = "__ferrum_warmup__";
22 let _ = self.prefill_internal(WARMUP_CACHE, &[0u32]);
23 // Drop the warmup KV cache slot — real cache_id is unaffected.
24 if let Some(caches) = self.kv_caches.remove(WARMUP_CACHE) {
25 self.kv_free_pool.push(caches);
26 }
27 }
28
29 fn kv_capacity(&self) -> usize {
30 // Mirror the bound `ensure_kv` will use when allocating the cache.
31 let model_max = self.cfg.base.max_seq_len;
32 self.runtime_env.kv_capacity(model_max)
33 }
34
35 fn prefill(&mut self, cache_id: &str, tokens: &[u32]) -> Vec<f32> {
36 self.prefill_internal(cache_id, tokens)
37 }
38
39 fn decode(&mut self, cache_id: &str, token: u32, pos: u32) -> Vec<f32> {
40 self.decode_internal(cache_id, token, pos)
41 }
42
43 // decode_batch is gated to use the batched path only when it's a
44 // measurable win. The crossover depends on M:
45 //
46 // - At low M (≤ ~8) the per-item `decode_internal` loop wins
47 // because: (a) it stays at scratch offset 0 (no copy_slice
48 // overhead), (b) it preserves the cross-layer rms_norm fusion
49 // fast path (`weighted_sum_residual_norm_stacked`).
50 // - At high M (≥ ~12) the batched path wins because the dense
51 // GEMM batching (qkv_proj, o_proj, router, lm_head at m=M) and
52 // the prefill-batched MoE dispatch (one `gemm_quant_moe_id` for
53 // all tokens) amortise the ~48-dispatch lost-fusion penalty.
54 //
55 // Default ON in 0.7.2+. On CUDA with paged KV + vLLM MoE, the
56 // crossover is now M=4: 2026-05-28/29 Vast RTX 4090 random-256/128
57 // probes saw the old threshold=8 stay on sequential per-token decode
58 // (~89-122 tok/s), while threshold=4 measured 425.6 ± 36.6 tok/s.
59 // `FERRUM_MOE_BATCHED=0` forces the
60 // legacy loop; `FERRUM_MOE_BATCH_THRESHOLD` remains an escape hatch
61 // for future hardware/backends.
62 fn decode_batch(&mut self, batch: &[(String, u32, u32)]) -> Vec<Vec<f32>> {
63 let m = batch.len();
64 let opted_in = self.runtime_env.moe_batched_enabled;
65 let threshold = self.runtime_env.moe_batch_threshold;
66 if opted_in && m >= threshold {
67 self.decode_batch_internal(batch)
68 } else {
69 batch
70 .iter()
71 .map(|(cid, tok, p)| self.decode(cid, *tok, *p))
72 .collect()
73 }
74 }
75
76 fn unified_forward(
77 &mut self,
78 items: &[(String, Vec<u32>, usize, bool)],
79 ) -> std::result::Result<Vec<Option<Vec<f32>>>, FerrumError> {
80 if items.is_empty() {
81 return Ok(Vec::new());
82 }
83 if self.runtime_env.qwen_unified_trace {
84 let lens: Vec<usize> = items.iter().map(|it| it.1.len()).collect();
85 let positions: Vec<usize> = items.iter().map(|it| it.2).collect();
86 let finals: Vec<bool> = items.iter().map(|it| it.3).collect();
87 eprintln!(
88 "[qwen-unified] items={} lens={:?} positions={:?} finals={:?} use_vllm_paged_attn={}",
89 items.len(),
90 lens,
91 positions,
92 finals,
93 self.use_vllm_paged_attn
94 );
95 }
96 if !B::supports_varlen_qkv() {
97 return Err(FerrumError::unsupported(
98 "Qwen3MoeModel::unified_forward: backend lacks varlen QKV kernels. \
99 Engine will fall back to legacy paths.",
100 ));
101 }
102 // Pure-decode shortcut: every item is q_len=1 + is_final_chunk.
103 // For this shape, ferrum's legacy `forward_layer_batched_decode`
104 // path (with FERRUM_MOE_GRAPH=1 graph capture + decode-tuned
105 // moe_forward_stacked) is faster than our generic varlen +
106 // bucketed-MoE unified path. Returning Unsupported routes the
107 // engine to the legacy decode_batch path via LlmExecutor's
108 // fallback partition.
109 let all_decode = items.iter().all(|it| it.1.len() == 1 && it.3);
110 if all_decode {
111 return Err(FerrumError::unsupported(
112 "Qwen3MoeModel::unified_forward: pure-decode batch — \
113 routed to legacy decode_batch (faster for q_len=1)",
114 ));
115 }
116 if items.len() == 1 && items[0].1.len() > 1 {
117 return Err(FerrumError::unsupported(
118 "Qwen3MoeModel::unified_forward: single-seq prefill — \
119 routed to specialized prefill path",
120 ));
121 }
122 if !self.runtime_env.qwen_unified_prefill && items.iter().any(|it| it.1.len() > 1) {
123 return Err(FerrumError::unsupported(
124 "Qwen3MoeModel::unified_forward: prefill disabled by \
125 FERRUM_QWEN_UNIFIED_PREFILL=0",
126 ));
127 }
128 // Any prefill chunk (q_len > 1) OR non-final-chunk item:
129 // unified path wins by collapsing N serial prefills into one
130 // [M_total, hidden] forward.
131 if self.paged_pools.is_none() {
132 return Err(FerrumError::unsupported(
133 "Qwen3MoeModel::unified_forward: paged KV required \
134 (set FERRUM_METAL_PAGED_KV=1).",
135 ));
136 }
137 let m_total: usize = items.iter().map(|it| it.1.len()).sum();
138 if m_total > self.scratch.max_tokens {
139 return Err(FerrumError::unsupported(format!(
140 "Qwen3MoeModel::unified_forward: m_total={} > scratch.max_tokens={}",
141 m_total, self.scratch.max_tokens,
142 )));
143 }
144 Ok(self.unified_forward_internal(items))
145 }
146
147 fn release(&mut self, cache_id: &str) {
148 // Mirror LlamaFamilyModel::release — do NOT reset the captured
149 // graphs here. Graphs reference paged_pool addresses (model-
150 // level + stable) and paged_batch_* scratch addresses (also
151 // model-level + stable); the per-cache_id state (paged_block_
152 // indices) lives in `kv_caches` and never appears in graph
153 // node args. Wiping graphs on release would invalidate them
154 // mid-flight (a release between capture and the next replay
155 // → CUDA_ERROR_INVALID_VALUE on cuGraphLaunch).
156 let mut ctx = B::new_context();
157 B::sync(&mut ctx);
158 if let Some(mut caches) = self.kv_caches.remove(cache_id) {
159 // Paged mode: return the cache_id's blocks to the shared
160 // allocator so other sequences can reuse them. Without this,
161 // every request consumes max_blocks_per_seq blocks
162 // permanently — pool exhausts after FERRUM_PAGED_MAX_SEQS
163 // requests and subsequent ensure_kv panics with
164 // "scratch residual missing" (the cascade panic from a
165 // failed ensure_kv path leaving scratch poisoned).
166 if let Some(alloc_arc) = self.paged_block_alloc.as_ref() {
167 let mut alloc = alloc_arc.lock().unwrap_or_else(|p| p.into_inner());
168 if let Some(c0) = caches.first() {
169 if !c0.paged_block_indices.is_empty() {
170 alloc.free(&c0.paged_block_indices);
171 }
172 }
173 for c in caches.iter_mut() {
174 c.paged_block_indices.clear();
175 }
176 }
177 self.kv_free_pool.push(caches);
178 }
179 }
180
181 fn reset(&mut self) {
182 let mut ctx = B::new_context();
183 B::sync(&mut ctx);
184 B::reset_all_graphs(&mut ctx);
185 self.batched_graph_keys_seen.clear();
186 self.batched_graph_warmup = 0;
187 self.batched_graph_failed = false;
188 B::sync(&mut ctx);
189 self.kv_caches.clear();
190 self.kv_free_pool.clear();
191 }
192}