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