@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 0899f436aa
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 4, 13:45 UTC

[New Model] Laguna XS.2 implementation (#41129)

Signed-off-by: Joe Rowell <joerowell4@gmail.com> Signed-off-by: Robert Shaw <robertgshaw2@gmail.com> Co-authored-by: Robert Shaw <robertgshaw2@gmail.com>

Joe Rowell committed Apr 28, 2026 at 20:23 UTC 0899f436aab42f798fb8e728872334c83aaebb79
10 files changed +1701 -16
tests/models/registry.py
+1
@@ -372,6 +372,7 @@ _TEXT_GENERATION_EXAMPLE_MODELS = {
372 "KimiLinearForCausalLM": _HfExamplesInfo(
373 "moonshotai/Kimi-Linear-48B-A3B-Instruct", trust_remote_code=True
374 ),
375 + "LagunaForCausalLM": _HfExamplesInfo("poolside/Laguna-XS.2"),
376 "Lfm2ForCausalLM": _HfExamplesInfo("LiquidAI/LFM2-1.2B"),
377 "Lfm2MoeForCausalLM": _HfExamplesInfo(
378 "LiquidAI/LFM2-8B-A1B",
vllm/model_executor/models/laguna.py new
+886
@@ -0,0 +1,886 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Inference-only Laguna model compatible with HuggingFace weights."""
4 +
5 +import typing
6 +from collections.abc import Callable, Iterable
7 +from itertools import islice
8 +
9 +import torch
10 +import torch.nn.functional as F
11 +from torch import nn
12 +
13 +from vllm.compilation.decorators import support_torch_compile
14 +from vllm.config import CacheConfig, VllmConfig, get_current_vllm_config
15 +from vllm.distributed import (
16 + get_ep_group,
17 + get_pp_group,
18 + get_tensor_model_parallel_rank,
19 + get_tensor_model_parallel_world_size,
20 +)
21 +from vllm.logger import init_logger
22 +from vllm.model_executor.layers.attention import Attention
23 +from vllm.model_executor.layers.fused_moe import FusedMoE
24 +from vllm.model_executor.layers.layernorm import RMSNorm
25 +from vllm.model_executor.layers.linear import (
26 + ColumnParallelLinear,
27 + QKVParallelLinear,
28 + ReplicatedLinear,
29 + RowParallelLinear,
30 +)
31 +from vllm.model_executor.layers.logits_processor import LogitsProcessor
32 +from vllm.model_executor.layers.quantization import QuantizationConfig
33 +from vllm.model_executor.layers.rotary_embedding import get_rope
34 +from vllm.model_executor.layers.vocab_parallel_embedding import (
35 + ParallelLMHead,
36 + VocabParallelEmbedding,
37 +)
38 +from vllm.model_executor.model_loader.weight_utils import (
39 + default_weight_loader,
40 + maybe_remap_kv_scale_name,
41 +)
42 +from vllm.model_executor.models.interfaces import SupportsLoRA, SupportsPP
43 +from vllm.model_executor.models.utils import (
44 + AutoWeightsLoader,
45 + PPMissingLayer,
46 + extract_layer_index,
47 + is_pp_missing_parameter,
48 + make_empty_intermediate_tensors_factory,
49 + make_layers,
50 + maybe_prefix,
51 +)
52 +from vllm.sequence import IntermediateTensors
53 +
54 +logger = init_logger(__name__)
55 +
56 +
57 +class LagunaMLP(nn.Module):
58 + """Dense MLP for Laguna (used in mlp_only_layers)."""
59 +
60 + def __init__(
61 + self,
62 + hidden_size: int,
63 + intermediate_size: int,
64 + hidden_act: str,
65 + quant_config: QuantizationConfig | None = None,
66 + reduce_results: bool = True,
67 + prefix: str = "",
68 + ) -> None:
69 + super().__init__()
70 + # gate_proj and up_proj are kept as separate ColumnParallelLinear
71 + # rather than merged via MergedColumnParallelLinear. The merged form
72 + # requires per-partition NVFP4 global scales (weight_global_scale,
73 + # input_global_scale) to be packed into a length-2 PerTensorScaleParameter
74 + # and then collapsed via .max() in process_weights_after_loading; this
75 + # doesn't round-trip cleanly through Marlin's NVFP4 stacked-layer code
76 + # path. Splitting yields one global scale per Linear, exactly matching
77 + # the standard compressed-tensors per-Linear schema on disk.
78 + self.gate_proj = ColumnParallelLinear(
79 + hidden_size,
80 + intermediate_size,
81 + bias=False,
82 + quant_config=quant_config,
83 + prefix=f"{prefix}.gate_proj",
84 + )
85 + self.up_proj = ColumnParallelLinear(
86 + hidden_size,
87 + intermediate_size,
88 + bias=False,
89 + quant_config=quant_config,
90 + prefix=f"{prefix}.up_proj",
91 + )
92 + self.down_proj = RowParallelLinear(
93 + intermediate_size,
94 + hidden_size,
95 + bias=False,
96 + quant_config=quant_config,
97 + reduce_results=reduce_results,
98 + prefix=f"{prefix}.down_proj",
99 + )
100 + if hidden_act != "silu":
101 + raise ValueError(
102 + f"Unsupported activation: {hidden_act}. Only silu is supported."
103 + )
104 +
105 + def forward(self, x: torch.Tensor) -> torch.Tensor:
106 + gate, _ = self.gate_proj(x)
107 + up, _ = self.up_proj(x)
108 + x, _ = self.down_proj(F.silu(gate) * up)
109 + return x
110 +
111 +
112 +class LagunaMoE(nn.Module):
113 + """Sparse MoE block for Laguna with optional shared expert and sigmoid routing.
114 +
115 + Key differences from other MoE implementations:
116 + - Uses SIGMOID routing activation (not softmax)
117 + - Shared expert runs in parallel with routed experts (when enabled)
118 + - Matches HF reference: modular_laguna.py LagunaSparseMoeBlock
119 + """
120 +
121 + def __init__(
122 + self,
123 + config,
124 + quant_config: QuantizationConfig | None = None,
125 + prefix: str = "",
126 + enable_eplb: bool = False,
127 + ):
128 + super().__init__()
129 + self.config = config
130 + self.num_experts = config.num_experts
131 + self.top_k = config.num_experts_per_tok
132 +
133 + self.tp_size = get_tensor_model_parallel_world_size()
134 + self.ep_group = get_ep_group().device_group
135 + self.ep_rank = self.ep_group.rank()
136 + self.ep_size = self.ep_group.size()
137 +
138 + self.n_routed_experts = config.num_experts
139 + self.n_shared_experts = 1 if config.shared_expert_intermediate_size > 0 else 0
140 + self.routed_scaling_factor = float(
141 + getattr(config, "moe_routed_scaling_factor", 1.0)
142 + )
143 +
144 + if self.tp_size > config.num_experts:
145 + raise ValueError(
146 + f"Tensor parallel size {self.tp_size} is greater than "
147 + f"the number of experts {config.num_experts}."
148 + )
149 +
150 + # Load balancing settings.
151 + vllm_config = get_current_vllm_config()
152 + eplb_config = vllm_config.parallel_config.eplb_config
153 + self.enable_eplb = enable_eplb
154 + eplb_config.num_redundant_experts = (
155 + eplb_config.num_redundant_experts
156 + if eplb_config.num_redundant_experts is not None
157 + else 0
158 + )
159 + self.n_redundant_experts = eplb_config.num_redundant_experts
160 + self.n_logical_experts = self.n_routed_experts
161 + self.n_physical_experts = self.n_logical_experts + self.n_redundant_experts
162 + self.n_local_physical_experts = self.n_physical_experts // self.ep_size
163 + self.physical_expert_start = self.ep_rank * self.n_local_physical_experts
164 + self.physical_expert_end = (
165 + self.physical_expert_start + self.n_local_physical_experts
166 + )
167 +
168 + # Router gate
169 + self.gate = ReplicatedLinear(
170 + config.hidden_size,
171 + config.num_experts,
172 + bias=False,
173 + quant_config=None,
174 + prefix=f"{prefix}.gate",
175 + )
176 +
177 + # Shared expert (optional) - passed to FusedMoE for overlap optimization
178 + self.shared_expert: LagunaMLP | None
179 + if config.shared_expert_intermediate_size > 0:
180 + self.shared_expert = LagunaMLP(
181 + hidden_size=config.hidden_size,
182 + intermediate_size=config.shared_expert_intermediate_size,
183 + hidden_act=config.hidden_act,
184 + quant_config=quant_config,
185 + reduce_results=False, # Reduce after shared+routed combine
186 + prefix=f"{prefix}.shared_expert",
187 + )
188 + else:
189 + self.shared_expert = None
190 +
191 + # Auxiliary-loss-free load-balancing bias (arXiv:2408.15664). The
192 + # checkpoint stores one [num_experts] tensor per MoE layer at
193 + # `mlp.experts.e_score_correction_bias`; registering it as a Parameter
194 + # on the FusedMoE lets the weight loader pick it up and the router
195 + # add it during top-k selection. The fused top-k bias router requires
196 + # float32 regardless of model dtype.
197 + e_score_correction_bias = torch.nn.Parameter(
198 + torch.zeros(config.num_experts, dtype=torch.float32),
199 + requires_grad=False,
200 + )
201 +
202 + # FusedMoE with SIGMOID routing. Passing `shared_experts=` lets the
203 + # layer overlap the shared-expert compute with the all2all dispatch.
204 + # `apply_routed_scale_to_output=True` makes FusedMoE handle the
205 + # routed_scaling_factor, shared+routed combine, and TP all-reduce
206 + # internally, so forward() just returns the final hidden states.
207 + self.experts = FusedMoE(
208 + shared_experts=self.shared_expert,
209 + num_experts=config.num_experts,
210 + top_k=config.num_experts_per_tok,
211 + hidden_size=config.hidden_size,
212 + intermediate_size=config.moe_intermediate_size,
213 + renormalize=config.norm_topk_prob,
214 + quant_config=quant_config,
215 + prefix=f"{prefix}.experts",
216 + scoring_func="sigmoid",
217 + use_grouped_topk=False,
218 + apply_router_weight_on_input=bool(config.moe_apply_router_weight_on_input),
219 + e_score_correction_bias=e_score_correction_bias,
220 + enable_eplb=self.enable_eplb,
221 + num_redundant_experts=self.n_redundant_experts,
222 + routed_scaling_factor=self.routed_scaling_factor,
223 + apply_routed_scale_to_output=True,
224 + )
225 +
226 + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
227 + orig_shape = hidden_states.shape
228 + hidden_dim = hidden_states.shape[-1]
229 + hidden_states = hidden_states.view(-1, hidden_dim)
230 +
231 + router_logits, _ = self.gate(hidden_states)
232 + router_logits = router_logits.float()
233 + softcap = getattr(self.config, "moe_router_logit_softcapping", 0.0) or 0.0
234 + if softcap > 0.0:
235 + router_logits = torch.tanh(router_logits / softcap) * softcap
236 +
237 + final_hidden_states = self.experts(hidden_states, router_logits)
238 + return final_hidden_states.view(orig_shape)
239 +
240 +
241 +class LagunaAttention(nn.Module):
242 + """Laguna attention with optional softplus output gating.
243 +
244 + Supports per-layer sliding window attention when ``config.layer_types``
245 + is present. Layers whose type is ``"sliding_attention"`` use
246 + ``config.sliding_window``; all other layers (typically labelled
247 + ``"full_attention"``) use full attention. When ``layer_types`` is
248 + absent every layer defaults to full attention for backwards
249 + compatibility.
250 + """
251 +
252 + def __init__(
253 + self,
254 + config,
255 + hidden_size: int,
256 + num_heads: int,
257 + num_kv_heads: int,
258 + max_position_embeddings: int = 131072,
259 + head_dim: int | None = None,
260 + cache_config: CacheConfig | None = None,
261 + quant_config: QuantizationConfig | None = None,
262 + prefix: str = "",
263 + attention_sink: bool = False,
264 + ) -> None:
265 + super().__init__()
266 + self.hidden_size = hidden_size
267 + tp_size = get_tensor_model_parallel_world_size()
268 + self.total_num_heads = num_heads
269 + assert self.total_num_heads % tp_size == 0
270 + self.num_heads = self.total_num_heads // tp_size
271 + self.total_num_kv_heads = num_kv_heads
272 + if self.total_num_kv_heads >= tp_size:
273 + assert self.total_num_kv_heads % tp_size == 0
274 + else:
275 + assert tp_size % self.total_num_kv_heads == 0
276 + self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size)
277 + self.head_dim = head_dim or (hidden_size // self.total_num_heads)
278 + self.q_size = self.num_heads * self.head_dim
279 + self.kv_size = self.num_kv_heads * self.head_dim
280 + self.scaling = self.head_dim**-0.5
281 + self.max_position_embeddings = max_position_embeddings
282 +
283 + # Gating flag
284 + self.gating = config.gating
285 +
286 + # Per-layer sliding window (follows Gemma2/Cohere2 convention)
287 + layer_types = getattr(config, "layer_types", None)
288 + if layer_types is not None:
289 + layer_idx = extract_layer_index(prefix)
290 + is_sliding = layer_types[layer_idx] == "sliding_attention"
291 + self.sliding_window = config.sliding_window if is_sliding else None
292 + else:
293 + self.sliding_window = None
294 +
295 + # QKV projection (no bias for Laguna)
296 + self.qkv_proj = QKVParallelLinear(
297 + self.hidden_size,
298 + self.head_dim,
299 + self.total_num_heads,
300 + self.total_num_kv_heads,
301 + bias=config.qkv_bias,
302 + quant_config=quant_config,
303 + prefix=f"{prefix}.qkv_proj",
304 + )
305 +
306 + # Output projection
307 + self.o_proj = RowParallelLinear(
308 + self.total_num_heads * self.head_dim,
309 + self.hidden_size,
310 + bias=config.attention_bias,
311 + quant_config=quant_config,
312 + prefix=f"{prefix}.o_proj",
313 + )
314 +
315 + # Gating projection (Laguna-specific, optional)
316 + # config.gating may be:
317 + # - True / "per-element": one gate per (head, head_dim) channel
318 + # - "per-head": one gate per head, broadcast across head_dim
319 + if self.gating:
320 + # v5 LagunaConfig uses ``gating=True`` for per-head; older configs
321 + # used ``"per-head"``. Accept both. ``"per-element"`` (or legacy
322 + # ``True``) means per-element gating with output size num_heads ×
323 + # head_dim.
324 + gate_per_head = self.gating is True or self.gating == "per-head"
325 + g_out = (
326 + self.total_num_heads
327 + if gate_per_head
328 + else self.total_num_heads * self.head_dim
329 + )
330 + self.g_proj = ColumnParallelLinear(
331 + hidden_size,
332 + g_out,
333 + bias=False,
334 + quant_config=quant_config,
335 + prefix=f"{prefix}.g_proj",
336 + )
337 + self.gate_per_head = gate_per_head
338 + else:
339 + self.g_proj = None
340 + self.gate_per_head = False
341 +
342 + # Attention sinks (learnable per-head bias for SWA layers)
343 + sinks = None
344 + if attention_sink:
345 + self.sink = torch.nn.Parameter(
346 + torch.empty(self.total_num_heads // tp_size, requires_grad=False)
347 + )
348 + sinks = self.sink
349 +
350 + # Resolve rope params per-layer-type. ``config.rope_parameters`` is
351 + # either a flat dict (legacy) or a nested ``{layer_type: rope_dict}``
352 + # (v5 Laguna-XS schema). The v5 form is unhashable as-is and would
353 + # crash `get_rope`'s cache lookup, so always pull out the layer's
354 + # sub-dict before forwarding.
355 + layer_type = (
356 + layer_types[extract_layer_index(prefix)]
357 + if layer_types is not None
358 + else "full_attention"
359 + )
360 + is_sliding = layer_type == "sliding_attention"
361 +
362 + top_rope = getattr(config, "rope_parameters", None) or {}
363 + if any(isinstance(v, dict) for v in top_rope.values()):
364 + # Nested per-layer-type form.
365 + base_rope = top_rope.get(layer_type) or top_rope.get("full_attention") or {}
366 + else:
367 + base_rope = top_rope
368 +
369 + # Older flat-rope ckpts can carry a separate `swa_rope_parameters`
370 + # for SWA layers. Prefer it when present; otherwise the nested
371 + # rope dict above already supplies the correct sub-config.
372 + swa_rope = getattr(config, "swa_rope_parameters", None)
373 + if (
374 + is_sliding
375 + and swa_rope is None
376 + and not any(isinstance(v, dict) for v in top_rope.values())
377 + ):
378 + logger.warning_once(
379 + "Laguna config has sliding_attention layers but neither "
380 + "`swa_rope_parameters` nor a nested per-layer-type "
381 + "`rope_parameters` — SWA layers will reuse the global rope. "
382 + "If the checkpoint was trained with distinct SWA rope "
383 + "(theta / partial_rotary_factor), regenerate its HF config "
384 + "to include either form."
385 + )
386 + rope_params = swa_rope if (is_sliding and swa_rope is not None) else base_rope
387 + # `partial_rotary_factor` may live on the top-level config (main attention)
388 + # or on the per-layer rope dict itself (e.g. SWA can differ). Inject the
389 + # top-level value into `rope_params` if the dict doesn't already set it.
390 + top_partial = getattr(config, "partial_rotary_factor", None)
391 + if top_partial is not None and "partial_rotary_factor" not in rope_params:
392 + rope_params = {**rope_params, "partial_rotary_factor": top_partial}
393 +
394 + # Rotary embeddings (YaRN)
395 + self.rotary_emb = get_rope(
396 + head_size=self.head_dim,
397 + max_position=max_position_embeddings,
398 + is_neox_style=True,
399 + rope_parameters=rope_params,
400 + )
401 +
402 + self.attn = Attention(
403 + self.num_heads,
404 + self.head_dim,
405 + self.scaling,
406 + num_kv_heads=self.num_kv_heads,
407 + cache_config=cache_config,
408 + quant_config=quant_config,
409 + per_layer_sliding_window=self.sliding_window,
410 + prefix=f"{prefix}.attn",
411 + sinks=sinks,
412 + )
413 +
414 + # QK normalization (like Qwen3)
415 + self.q_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps)
416 + self.k_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps)
417 +
418 + def forward(
419 + self,
420 + positions: torch.Tensor,
421 + hidden_states: torch.Tensor,
422 + ) -> torch.Tensor:
423 + qkv, _ = self.qkv_proj(hidden_states)
424 + q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
425 +
426 + q_by_head = q.view(*q.shape[:-1], q.shape[-1] // self.head_dim, self.head_dim)
427 + q_by_head = self.q_norm(q_by_head)
428 + q = q_by_head.view(q.shape)
429 +
430 + k_by_head = k.view(*k.shape[:-1], k.shape[-1] // self.head_dim, self.head_dim)
431 + k_by_head = self.k_norm(k_by_head)
432 + k = k_by_head.view(k.shape)
433 +
434 + q, k = self.rotary_emb(positions, q, k)
435 + attn_output = self.attn(q, k, v)
436 +
437 + # Apply gating if enabled (compute softplus in float32 for precision)
438 + if self.gating and self.g_proj is not None:
439 + gate, _ = self.g_proj(hidden_states)
440 + gate = F.softplus(gate.float()).type_as(attn_output)
441 + if self.gate_per_head:
442 + # gate: [..., num_heads]; broadcast across head_dim
443 + attn_shape = attn_output.shape
444 + attn_output = (
445 + attn_output.view(*attn_shape[:-1], self.num_heads, self.head_dim)
446 + * gate.unsqueeze(-1)
447 + ).view(attn_shape)
448 + else:
449 + attn_output = attn_output * gate
450 +
451 + output, _ = self.o_proj(attn_output)
452 + return output
453 +
454 +
455 +class LagunaDecoderLayer(nn.Module):
456 + def __init__(
457 + self,
458 + config,
459 + cache_config: CacheConfig | None = None,
460 + quant_config: QuantizationConfig | None = None,
461 + prefix: str = "",
462 + enable_eplb: bool = False,
463 + ) -> None:
464 + super().__init__()
465 + self.hidden_size = config.hidden_size
466 + layer_idx = extract_layer_index(prefix)
467 +
468 + # Determine if this layer uses sliding window attention
469 + layer_types = getattr(config, "layer_types", None)
470 + is_sliding = (
471 + layer_types is not None and layer_types[layer_idx] == "sliding_attention"
472 + )
473 +
474 + # Enable attention sinks on SWA layers when configured
475 + attention_sink = is_sliding and getattr(
476 + config, "swa_attention_sink_enabled", False
477 + )
478 +
479 + # Optional per-layer override of head count (Laguna-XS).
480 + per_layer_heads = getattr(config, "num_attention_heads_per_layer", None)
481 + layer_num_heads = (
482 + per_layer_heads[layer_idx]
483 + if per_layer_heads is not None
484 + else config.num_attention_heads
485 + )
486 +
487 + self.self_attn = LagunaAttention(
488 + config=config,
489 + hidden_size=self.hidden_size,
490 + num_heads=layer_num_heads,
491 + num_kv_heads=config.num_key_value_heads,
492 + max_position_embeddings=config.max_position_embeddings,
493 + head_dim=getattr(config, "head_dim", None),
494 + cache_config=cache_config,
495 + quant_config=quant_config,
496 + prefix=f"{prefix}.self_attn",
497 + attention_sink=attention_sink,
498 + )
499 +
500 + # Check if this layer uses MoE or dense MLP (matches Qwen2/Qwen3 convention)
501 + mlp_only_layers = (
502 + [] if not hasattr(config, "mlp_only_layers") else config.mlp_only_layers
503 + )
504 + self.is_moe_layer = (
505 + (layer_idx not in mlp_only_layers)
506 + and (config.num_experts > 0)
507 + and ((layer_idx + 1) % config.decoder_sparse_step == 0)
508 + )
509 +
510 + if self.is_moe_layer:
511 + self.mlp = LagunaMoE(
512 + config=config,
513 + quant_config=quant_config,
514 + prefix=f"{prefix}.mlp",
515 + enable_eplb=enable_eplb,
516 + )
517 + else:
518 + self.mlp = LagunaMLP(
519 + hidden_size=config.hidden_size,
520 + intermediate_size=config.intermediate_size,
521 + hidden_act=config.hidden_act,
522 + quant_config=quant_config,
523 + prefix=f"{prefix}.mlp",
524 + )
525 +
526 + self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
527 + self.post_attention_layernorm = RMSNorm(
528 + config.hidden_size, eps=config.rms_norm_eps
529 + )
530 +
531 + def forward(
532 + self,
533 + positions: torch.Tensor,
534 + hidden_states: torch.Tensor,
535 + residual: torch.Tensor | None,
536 + ) -> tuple[torch.Tensor, torch.Tensor]:
537 + # Self Attention
538 + if residual is None:
539 + residual = hidden_states
540 + hidden_states = self.input_layernorm(hidden_states)
541 + else:
542 + hidden_states, residual = self.input_layernorm(hidden_states, residual)
543 +
544 + hidden_states = self.self_attn(
545 + positions=positions,
546 + hidden_states=hidden_states,
547 + )
548 +
549 + # Fully Connected
550 + hidden_states, residual = self.post_attention_layernorm(hidden_states, residual)
551 + hidden_states = self.mlp(hidden_states)
552 +
553 + return hidden_states, residual
554 +
555 +
556 +@support_torch_compile
557 +class LagunaModel(nn.Module):
558 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
559 + super().__init__()
560 +
561 + config = vllm_config.model_config.hf_config
562 + cache_config = vllm_config.cache_config
563 + quant_config = vllm_config.quant_config
564 + enable_eplb = vllm_config.parallel_config.enable_eplb
565 + eplb_config = vllm_config.parallel_config.eplb_config
566 + self.num_redundant_experts = eplb_config.num_redundant_experts
567 + self.config = config
568 + self.quant_config = quant_config
569 +
570 + # Disable the model-level sliding-window fallback in Attention.__init__.
571 + # Laguna drives SWA per-layer via `layer_types`, passing
572 + # `per_layer_sliding_window=self.sliding_window` (None for global
573 + # layers). Without this, global layers whose `per_layer_sliding_window`
574 + # is None would pick up `cache_config.sliding_window`
575 + # (populated from `config.sliding_window`) as a fallback, silently
576 + # applying a 512-token window to full-attention layers.
577 + if cache_config is not None:
578 + cache_config.sliding_window = None
579 +
580 + self.vocab_size = config.vocab_size
581 +
582 + if get_pp_group().is_first_rank or (
583 + config.tie_word_embeddings and get_pp_group().is_last_rank
584 + ):
585 + self.embed_tokens = VocabParallelEmbedding(
586 + config.vocab_size,
587 + config.hidden_size,
588 + quant_config=quant_config,
589 + prefix=f"{prefix}.embed_tokens",
590 + )
591 + else:
592 + self.embed_tokens = PPMissingLayer()
593 +
594 + self.start_layer, self.end_layer, self.layers = make_layers(
595 + config.num_hidden_layers,
596 + lambda prefix: LagunaDecoderLayer(
597 + config=config,
598 + cache_config=cache_config,
599 + quant_config=quant_config,
600 + prefix=prefix,
601 + enable_eplb=enable_eplb,
602 + ),
603 + prefix=f"{prefix}.layers",
604 + )
605 +
606 + if get_pp_group().is_last_rank:
607 + self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
608 + else:
609 + self.norm = PPMissingLayer()
610 +
611 + self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory(
612 + ["hidden_states", "residual"], config.hidden_size
613 + )
614 +
615 + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
616 + return self.embed_tokens(input_ids)
617 +
618 + def forward(
619 + self,
620 + input_ids: torch.Tensor,
621 + positions: torch.Tensor,
622 + intermediate_tensors: IntermediateTensors | None = None,
623 + inputs_embeds: torch.Tensor | None = None,
624 + ) -> torch.Tensor | IntermediateTensors:
625 + if get_pp_group().is_first_rank:
626 + if inputs_embeds is not None:
627 + hidden_states = inputs_embeds
628 + else:
629 + hidden_states = self.embed_tokens(input_ids)
630 + residual = None
631 + else:
632 + assert intermediate_tensors is not None
633 + hidden_states = intermediate_tensors["hidden_states"]
634 + residual = intermediate_tensors["residual"]
635 +
636 + for layer in islice(self.layers, self.start_layer, self.end_layer):
637 + hidden_states, residual = layer(positions, hidden_states, residual)
638 +
639 + if not get_pp_group().is_last_rank:
640 + return IntermediateTensors(
641 + {"hidden_states": hidden_states, "residual": residual}
642 + )
643 +
644 + hidden_states, _ = self.norm(hidden_states, residual)
645 + return hidden_states
646 +
647 + def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
648 + """Get expert parameter mapping for weight loading.
649 +
650 + Returns mapping tuples of (param_name, weight_name, expert_id, shard_id)
651 + that handle both weights and quantization scales.
652 + """
653 + return FusedMoE.make_expert_params_mapping(
654 + self,
655 + ckpt_gate_proj_name="gate_proj",
656 + ckpt_down_proj_name="down_proj",
657 + ckpt_up_proj_name="up_proj",
658 + num_experts=self.config.num_experts,
659 + num_redundant_experts=self.num_redundant_experts,
660 + )
661 +
662 + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
663 + stacked_params_mapping = [
664 + # (param_name, shard_name, shard_id)
665 + ("qkv_proj", "q_proj", "q"),
666 + ("qkv_proj", "k_proj", "k"),
667 + ("qkv_proj", "v_proj", "v"),
668 + # gate_proj and up_proj are loaded as separate Linears (see
669 + # LagunaMLP) so no merge entry is needed here.
670 + ]
671 +
672 + # Suffixes to skip for GPTQ/modelopt models if param doesn't exist
673 + ignore_suffixes = (
674 + ".bias",
675 + "_bias",
676 + ".k_scale",
677 + "_k_scale",
678 + ".v_scale",
679 + "_v_scale",
680 + ".weight_scale",
681 + "_weight_scale",
682 + ".input_scale",
683 + "_input_scale",
684 + )
685 +
686 + params_dict = dict(self.named_parameters())
687 + loaded_params: set[str] = set()
688 + expert_params_mapping = self.get_expert_mapping()
689 +
690 + tp_rank = get_tensor_model_parallel_rank()
691 +
692 + for name, loaded_weight in weights:
693 + # Handle attention sinks (distributed across ranks). Derive the
694 + # per-rank slice from the parameter's own shape so per-layer
695 + # variations in head count are handled correctly.
696 + if "sink" in name:
697 + param = params_dict.get(name)
698 + if param is not None:
699 + layer_heads_per_rank = param.shape[0]
700 + layer_head_start = tp_rank * layer_heads_per_rank
701 + narrow_weight = loaded_weight.narrow(
702 + 0, layer_head_start, layer_heads_per_rank
703 + )
704 + param.data.copy_(narrow_weight)
705 + loaded_params.add(name)
706 + continue
707 +
708 + # Handle KV cache quantization scales
709 + if self.quant_config is not None and (
710 + scale_name := self.quant_config.get_cache_scale(name)
711 + ):
712 + param = params_dict[scale_name]
713 + weight_loader = getattr(param, "weight_loader", default_weight_loader)
714 + assert loaded_weight.numel() == 1, (
715 + f"KV scale numel {loaded_weight.numel()} != 1"
716 + )
717 + loaded_weight = loaded_weight.squeeze()
718 + weight_loader(param, loaded_weight)
719 + loaded_params.add(scale_name)
720 + continue
721 +
722 + # Handle stacked params (QKV, gate_up for
723 + # non-expert layers and shared_expert)
724 + for param_name, weight_name, shard_id in stacked_params_mapping:
725 + if weight_name not in name:
726 + continue
727 + # Skip expert weights - handled below via expert_params_mapping
728 + if "mlp.experts" in name and "shared_expert" not in name:
729 + continue
730 + name = name.replace(weight_name, param_name)
731 +
732 + if name.endswith(ignore_suffixes) and name not in params_dict:
733 + continue
734 + if is_pp_missing_parameter(name, self):
735 + continue
736 + # Remap FP8 kv_scale names for backwards compatibility
737 + if name.endswith("scale"):
738 + name = maybe_remap_kv_scale_name(name, params_dict)
739 + if name is None:
740 + continue
741 + if name not in params_dict:
742 + continue
743 +
744 + param = params_dict[name]
745 + weight_loader = getattr(param, "weight_loader", default_weight_loader)
746 + if weight_loader == default_weight_loader:
747 + weight_loader(param, loaded_weight)
748 + else:
749 + weight_loader(param, loaded_weight, shard_id)
750 + loaded_params.add(name)
751 + break
752 + else:
753 + # Try expert params mapping (handles weights + quantization scales)
754 + is_expert_weight = False
755 + for mapping in expert_params_mapping:
756 + param_name, weight_name, expert_id, shard_id = mapping
757 + if weight_name not in name:
758 + continue
759 +
760 + # Mark as expert weight so we skip regular loading below
761 + is_expert_weight = True
762 +
763 + # Create mapped name without modifying original
764 + name_mapped = name.replace(weight_name, param_name)
765 +
766 + if is_pp_missing_parameter(name_mapped, self):
767 + continue
768 + if (
769 + name_mapped.endswith(ignore_suffixes)
770 + and name_mapped not in params_dict
771 + ):
772 + continue
773 + if name_mapped not in params_dict:
774 + continue
775 +
776 + param = params_dict[name_mapped]
777 + # Use return_success to handle expert parallelism correctly
778 + weight_loader = typing.cast(
779 + Callable[..., bool], param.weight_loader
780 + )
781 + success = weight_loader(
782 + param,
783 + loaded_weight,
784 + name_mapped,
785 + shard_id=shard_id,
786 + expert_id=expert_id,
787 + return_success=True,
788 + )
789 + if success:
790 + loaded_params.add(name_mapped)
791 + break
792 + else:
793 + # Expert weight not mapped to this rank - skip
794 + if is_expert_weight:
795 + continue
796 +
797 + # Remap kv_scale names before the ignore_suffixes filter:
798 + # the suffix list includes .k_scale/.v_scale, so filtering
799 + # first drops the checkpoint key before remap can rewrite
800 + # it to the .attn.* name that exists in params_dict.
801 + name = maybe_remap_kv_scale_name(name, params_dict)
802 + if name is None:
803 + continue
804 +
805 + if name.endswith(ignore_suffixes) and name not in params_dict:
806 + continue
807 +
808 + if is_pp_missing_parameter(name, self):
809 + continue
810 +
811 + if name not in params_dict:
812 + continue
813 +
814 + param = params_dict[name]
815 + weight_loader = getattr(
816 + param, "weight_loader", default_weight_loader
817 + )
818 + weight_loader(param, loaded_weight)
819 + loaded_params.add(name)
820 +
821 + return loaded_params
822 +
823 +
824 +class LagunaForCausalLM(nn.Module, SupportsPP, SupportsLoRA):
825 + fall_back_to_pt_during_load = False
826 +
827 + packed_modules_mapping = {
828 + "qkv_proj": ["q_proj", "k_proj", "v_proj"],
829 + }
830 +
831 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
832 + super().__init__()
833 + config = vllm_config.model_config.hf_config
834 + quant_config = vllm_config.quant_config
835 + self.config = config
836 + self.quant_config = quant_config
837 +
838 + self.model = LagunaModel(
839 + vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
840 + )
841 +
842 + if get_pp_group().is_last_rank:
843 + self.lm_head = ParallelLMHead(
844 + config.vocab_size,
845 + config.hidden_size,
846 + quant_config=quant_config,
847 + prefix=maybe_prefix(prefix, "lm_head"),
848 + )
849 + if self.config.tie_word_embeddings:
850 + self.lm_head = self.lm_head.tie_weights(self.model.embed_tokens)
851 + else:
852 + self.lm_head = PPMissingLayer()
853 +
854 + self.logits_processor = LogitsProcessor(config.vocab_size)
855 + self.make_empty_intermediate_tensors = (
856 + self.model.make_empty_intermediate_tensors
857 + )
858 +
859 + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
860 + return self.model.embed_input_ids(input_ids)
861 +
862 + def forward(
863 + self,
864 + input_ids: torch.Tensor,
865 + positions: torch.Tensor,
866 + intermediate_tensors: IntermediateTensors | None = None,
867 + inputs_embeds: torch.Tensor | None = None,
868 + ) -> torch.Tensor | IntermediateTensors:
869 + hidden_states = self.model(
870 + input_ids, positions, intermediate_tensors, inputs_embeds
871 + )
872 + return hidden_states
873 +
874 + def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor | None:
875 + logits = self.logits_processor(self.lm_head, hidden_states)
876 + return logits
877 +
878 + def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
879 + return self.model.get_expert_mapping()
880 +
881 + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
882 + loader = AutoWeightsLoader(
883 + self,
884 + skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None),
885 + )
886 + return loader.load_weights(weights)
vllm/model_executor/models/registry.py
+1
@@ -150,6 +150,7 @@ _TEXT_GENERATION_MODELS = {
150 "KimiLinearForCausalLM": ("kimi_linear", "KimiLinearForCausalLM"),
151 "Lfm2ForCausalLM": ("lfm2", "Lfm2ForCausalLM"),
152 "Lfm2MoeForCausalLM": ("lfm2_moe", "Lfm2MoeForCausalLM"),
153 + "LagunaForCausalLM": ("laguna", "LagunaForCausalLM"),
154 "LlamaForCausalLM": ("llama", "LlamaForCausalLM"),
155 "Llama4ForCausalLM": ("llama4", "Llama4ForCausalLM"),
156 # For decapoda-research/llama-*
vllm/reasoning/__init__.py
+4
@@ -32,6 +32,10 @@ _REASONING_PARSERS_TO_REGISTER = {
32 "deepseek_v3_reasoning_parser",
33 "DeepSeekV3ReasoningParser",
34 ),
35 + "poolside_v1": (
36 + "poolside_v1_reasoning_parser",
37 + "PoolsideV1ReasoningParser",
38 + ),
39 "ernie45": (
40 "ernie45_reasoning_parser",
41 "Ernie45ReasoningParser",
vllm/reasoning/poolside_v1_reasoning_parser.py new
+72
@@ -0,0 +1,72 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""
4 +Laguna reasoning parser.
5 +
6 +``DeepSeekV3ReasoningParser.is_reasoning_end`` walks the entire
7 +token sequence backwards and returns ``True`` on the first ``</think>`` it
8 +sees. When called on ``prompt_token_ids`` that mistakes any stray
9 +``</think>`` in conversation history, few-shot examples or tool descriptions
10 +for a template-injected "thinking already ended" marker. In the streaming
11 +path (see ``vllm/entrypoints/openai/chat_completion/serving.py``,
12 +``prompt_is_reasoning_end_arr``) that false positive short-circuits the
13 +reasoning parser for the whole response, so any ``<think>...</think>`` the
14 +model emits itself ends up in the content field instead of the reasoning
15 +field.
16 +
17 +As we have more flexible templates, we instead scope
18 +the backward search to the current assistant turn: the
19 +walk terminates as soon as we hit the ``<assistant>`` start-of-message
20 +token. A ``</think>`` in a prior user turn or few-shot example is no longer
21 +visible.
22 +"""
23 +
24 +from collections.abc import Sequence
25 +
26 +from transformers import PreTrainedTokenizerBase
27 +
28 +from vllm.reasoning.deepseek_r1_reasoning_parser import DeepSeekR1ReasoningParser
29 +from vllm.reasoning.deepseek_v3_reasoning_parser import DeepSeekV3ReasoningParser
30 +from vllm.reasoning.identity_reasoning_parser import IdentityReasoningParser
31 +
32 +
33 +class PoolsideV1ReasoningParser(DeepSeekV3ReasoningParser):
34 + """Drop-in replacement for ``deepseek_v3`` that tolerates ``</think>``
35 + tokens appearing anywhere in the prompt other than the generation prefix.
36 + """
37 +
38 + _start_of_assistant_message = "<assistant>"
39 +
40 + def __init__(self, tokenizer: PreTrainedTokenizerBase, *args, **kwargs):
41 + super().__init__(tokenizer, *args, **kwargs)
42 +
43 + if self._start_of_assistant_message not in self.vocab:
44 + raise ValueError(
45 + f"Tokenizer must contain {self._start_of_assistant_message!r} token"
46 + )
47 + self._start_of_assistant_message_token_id = self.vocab[
48 + self._start_of_assistant_message
49 + ]
50 +
51 + def is_reasoning_end(self, input_ids: Sequence[int]) -> bool:
52 + # IdentityReasoningParser always returns True: no reasoning to parse.
53 + if isinstance(self._parser, IdentityReasoningParser):
54 + return True
55 +
56 + assert isinstance(self._parser, DeepSeekR1ReasoningParser)
57 + for tok_id in reversed(input_ids):
58 + # <think>: reasoning is not yet ended.
59 + if tok_id == self._parser.start_token_id:
60 + return False
61 + # </think>: reasoning has ended.
62 + if tok_id == self._parser.end_token_id:
63 + return True
64 + # <assistant>: reached the start of the current assistant turn
65 + # without seeing either marker. Anything further back belongs to
66 + # the prior conversation and should be ignored.
67 + if tok_id == self._start_of_assistant_message_token_id:
68 + return False
69 + return False
70 +
71 +
72 +__all__ = ["PoolsideV1ReasoningParser"]
vllm/tool_parsers/__init__.py
+4
@@ -66,6 +66,10 @@ _TOOL_PARSERS_TO_REGISTER = {
66 "hermes_tool_parser",
67 "Hermes2ProToolParser",
68 ),
69 + "poolside_v1": (
70 + "poolside_v1_tool_parser",
71 + "PoolsideV1ToolParser",
72 + ),
73 "hunyuan_a13b": (
74 "hunyuan_a13b_tool_parser",
75 "HunyuanA13BToolParser",
vllm/tool_parsers/poolside_v1_tool_parser.py new
+583
@@ -0,0 +1,583 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""
4 +GLM-4 Tool Call Parser with incremental string streaming support.
5 +
6 +This parser fixes the streaming issue reported in Issue #32829 where long string
7 +parameters (e.g., file content with 4000+ characters of code) are buffered until
8 +complete, causing multi-second delays before the user sees any content.
9 +
10 +The fix streams string values incrementally as they arrive, providing a true
11 +streaming experience for long content.
12 +"""
13 +
14 +import ast
15 +import json
16 +from collections.abc import Sequence
17 +from typing import Any
18 +
19 +import partial_json_parser.core.complete
20 +import regex as re
21 +from partial_json_parser.core.options import Allow
22 +
23 +from vllm.entrypoints.chat_utils import make_tool_call_id
24 +from vllm.entrypoints.openai.chat_completion.protocol import (
25 + ChatCompletionRequest,
26 +)
27 +from vllm.entrypoints.openai.engine.protocol import (
28 + DeltaFunctionCall,
29 + DeltaMessage,
30 + DeltaToolCall,
31 + ExtractedToolCallInformation,
32 + FunctionCall,
33 + ToolCall,
34 +)
35 +from vllm.entrypoints.openai.responses.protocol import (
36 + ResponsesRequest,
37 +)
38 +from vllm.logger import init_logger
39 +from vllm.tokenizers import TokenizerLike
40 +from vllm.tool_parsers.abstract_tool_parser import (
41 + Tool,
42 + ToolParser,
43 +)
44 +
45 +logger = init_logger(__name__)
46 +
47 +
48 +class PoolsideV1ToolParser(ToolParser):
49 + """Tool parser for GLM-4 models with incremental string streaming.
50 +
51 + This parser emits tool-call deltas incrementally as arguments arrive.
52 + For string-type parameters, content is streamed character-by-character
53 + rather than waiting for the complete </arg_value> tag.
54 + """
55 +
56 + def __init__(self, tokenizer: TokenizerLike, tools: list[Tool] | None = None):
57 + super().__init__(tokenizer, tools)
58 + # Stateful streaming fields
59 + self.current_tool_name_sent: bool = False
60 + self.prev_tool_call_arr: list[dict[str, Any]] = []
61 + self.current_tool_id: int = -1
62 + self.streamed_args_for_tool: list[str] = []
63 +
64 + self.tool_call_start_token: str = "<tool_call>"
65 + self.tool_call_end_token: str = "</tool_call>"
66 + self.arg_key_start: str = "<arg_key>"
67 + self.arg_key_end: str = "</arg_key>"
68 + self.arg_val_start: str = "<arg_value>"
69 + self.arg_val_end: str = "</arg_value>"
70 +
71 + self.tool_calls_start_token = self.tool_call_start_token
72 +
73 + self.func_call_regex = re.compile(r"<tool_call>.*?</tool_call>", re.DOTALL)
74 + self.func_detail_regex = re.compile(
75 + r"<tool_call>([^\n]*)\n(.*)</tool_call>", re.DOTALL
76 + )
77 + self.func_arg_regex = re.compile(
78 + r"<arg_key>(.*?)</arg_key>\s*<arg_value>(.*?)</arg_value>", re.DOTALL
79 + )
80 +
81 + if not self.model_tokenizer:
82 + raise ValueError(
83 + "The model tokenizer must be passed to the ToolParser "
84 + "constructor during construction."
85 + )
86 +
87 + self.tool_call_start_token_id = self.vocab.get(self.tool_call_start_token)
88 + self.tool_call_end_token_id = self.vocab.get(self.tool_call_end_token)
89 + self._buffer: str = ""
90 +
91 + # Streaming state for incremental tool-call streaming
92 + self._in_tool_call: bool = False
93 + self._current_tool_name: str | None = None
94 + self._pending_key: str | None = None
95 + self._streaming_string_value: bool = False
96 + self._tool_call_ids: list[str] = []
97 + self._args_started: list[bool] = []
98 + self._args_closed: list[bool] = []
99 + self._seen_keys: list[set[str]] = []
100 +
101 + @staticmethod
102 + def _deserialize(value: str) -> Any:
103 + try:
104 + return json.loads(value)
105 + except json.JSONDecodeError:
106 + pass
107 +
108 + try:
109 + return ast.literal_eval(value)
110 + except (ValueError, SyntaxError):
111 + pass
112 +
113 + return value
114 +
115 + @staticmethod
116 + def _json_escape_string_content(s: str) -> str:
117 + """JSON-escape string content for incremental streaming.
118 +
119 + This escapes the content that goes INSIDE a JSON string (between quotes),
120 + not including the surrounding quotes themselves.
121 + """
122 + if not s:
123 + return ""
124 + return json.dumps(s, ensure_ascii=False)[1:-1]
125 +
126 + @staticmethod
127 + def _is_string_type(
128 + tool_name: str,
129 + arg_name: str,
130 + tools: list[Tool] | None,
131 + ) -> bool:
132 + if tools is None:
133 + return False
134 + for tool in tools:
135 + if tool.function.name != tool_name:
136 + continue
137 + if tool.function.parameters is None:
138 + return False
139 + arg_type = (
140 + tool.function.parameters.get("properties", {})
141 + .get(arg_name, {})
142 + .get("type", None)
143 + )
144 + return arg_type == "string"
145 + logger.debug("No tool named '%s'.", tool_name)
146 + return False
147 +
148 + @staticmethod
149 + def _tools_enabled(request: ChatCompletionRequest) -> bool:
150 + """Return whether tool parsing should be applied for this request."""
151 + try:
152 + tools = getattr(request, "tools", None)
153 + tool_choice = getattr(request, "tool_choice", None)
154 + return bool(tools) and tool_choice != "none"
155 + except Exception:
156 + logger.exception("Failed to determine if tools are enabled.")
157 + return False
158 +
159 + def adjust_request(
160 + self, request: ChatCompletionRequest | ResponsesRequest
161 + ) -> ChatCompletionRequest | ResponsesRequest:
162 + """Adjust request parameters for tool call token handling."""
163 + request = super().adjust_request(request)
164 + if request.tools and request.tool_choice != "none":
165 + # Ensure tool call tokens (<tool_call>, </tool_call>) are not skipped
166 + # during decoding. Even though they are not marked as special tokens,
167 + # setting skip_special_tokens=False ensures proper handling in
168 + # transformers 5.x where decoding behavior may have changed.
169 + request.skip_special_tokens = False
170 + return request
171 +
172 + def extract_tool_calls(
173 + self,
174 + model_output: str,
175 + request: ChatCompletionRequest,
176 + ) -> ExtractedToolCallInformation:
177 + matched_tool_calls = self.func_call_regex.findall(model_output)
178 + logger.debug("model_output: %s", model_output)
179 + try:
180 + tool_calls: list[ToolCall] = []
181 + for match in matched_tool_calls:
182 + tc_detail = self.func_detail_regex.search(match)
183 + if not tc_detail:
184 + logger.warning(
185 + "Failed to parse tool call details from: %s",
186 + match,
187 + )
188 + continue
189 + tc_name = tc_detail.group(1).strip()
190 + tc_args = tc_detail.group(2)
191 + pairs = self.func_arg_regex.findall(tc_args) if tc_args else []
192 + arg_dct: dict[str, Any] = {}
193 + for key, value in pairs:
194 + arg_key = key.strip()
195 + arg_val = value.strip()
196 + if not self._is_string_type(tc_name, arg_key, request.tools):
197 + arg_val = self._deserialize(arg_val)
198 + logger.debug("arg_key = %s, arg_val = %s", arg_key, arg_val)
199 + arg_dct[arg_key] = arg_val
200 + tool_calls.append(
201 + ToolCall(
202 + type="function",
203 + function=FunctionCall(
204 + name=tc_name,
205 + arguments=json.dumps(arg_dct, ensure_ascii=False),
206 + ),
207 + )
208 + )
209 + except Exception:
210 + logger.exception("Failed to extract tool call spec")
211 + return ExtractedToolCallInformation(
212 + tools_called=False, tool_calls=[], content=model_output
213 + )
214 + else:
215 + if len(tool_calls) > 0:
216 + content: str | None = model_output[
217 + : model_output.find(self.tool_calls_start_token)
218 + ]
219 + # Normalize empty/whitespace-only content to None
220 + if not content or not content.strip():
221 + content = None
222 + return ExtractedToolCallInformation(
223 + tools_called=True, tool_calls=tool_calls, content=content
224 + )
225 + return ExtractedToolCallInformation(
226 + tools_called=False, tool_calls=[], content=model_output
227 + )
228 +
229 + def extract_tool_calls_streaming(
230 + self,
231 + previous_text: str,
232 + current_text: str,
233 + delta_text: str,
234 + previous_token_ids: Sequence[int],
235 + current_token_ids: Sequence[int],
236 + delta_token_ids: Sequence[int],
237 + request: ChatCompletionRequest,
238 + ) -> DeltaMessage | None:
239 + if not self._tools_enabled(request):
240 + return DeltaMessage(content=delta_text) if delta_text else None
241 +
242 + self._buffer += delta_text
243 +
244 + pending_deltas: dict[int, DeltaToolCall] = {}
245 + content: str | None = None
246 +
247 + while True:
248 + if not self._in_tool_call:
249 + start_idx = self._buffer.find(self.tool_call_start_token)
250 + if start_idx == -1:
251 + # Check for partial start token at end of buffer
252 + for i in range(1, len(self.tool_call_start_token)):
253 + if self._buffer.endswith(self.tool_call_start_token[:i]):
254 + out = self._buffer[:-i]
255 + self._buffer = self._buffer[-i:]
256 + if out:
257 + content = (content or "") + out
258 + break
259 + else:
260 + out = self._buffer
261 + self._buffer = ""
262 + if out:
263 + content = (content or "") + out
264 + break
265 +
266 + if start_idx > 0:
267 + content = (content or "") + self._buffer[:start_idx]
268 + self._buffer = self._buffer[start_idx:]
269 +
270 + self._buffer = self._buffer[len(self.tool_call_start_token) :]
271 + self._begin_tool_call()
272 + continue
273 +
274 + # Parse tool name first
275 + if not self.current_tool_name_sent:
276 + nl = self._buffer.find("\n")
277 + ak = self._buffer.find(self.arg_key_start)
278 + end = self._buffer.find(self.tool_call_end_token)
279 + candidates = [i for i in [nl, ak, end] if i != -1]
280 + if not candidates:
281 + break
282 + cut = min(candidates)
283 + tool_name = self._buffer[:cut].strip()
284 + if tool_name == "" and cut == end:
285 + # Handle empty tool call like `<tool_call></tool_call>`.
286 + # Consume the tokens and reset state to avoid infinite loop.
287 + self._buffer = self._buffer[end + len(self.tool_call_end_token) :]
288 + self._finish_tool_call()
289 + self._revert_last_tool_call_state()
290 + continue
291 +
292 + if cut == nl:
293 + self._buffer = self._buffer[nl + 1 :]
294 + else:
295 + self._buffer = self._buffer[cut:]
296 +
297 + self._current_tool_name = tool_name
298 + self.current_tool_name_sent = True
299 + self._update_tool_name(pending_deltas, tool_name)
300 + continue
301 +
302 + assert self._current_tool_name is not None
303 +
304 + # Handle incremental string value streaming
305 + if self._streaming_string_value:
306 + val_end = self._buffer.find(self.arg_val_end)
307 + if val_end != -1:
308 + raw_content = self._buffer[:val_end]
309 + self._buffer = self._buffer[val_end + len(self.arg_val_end) :]
310 + self._streaming_string_value = False
311 + self._pending_key = None
312 +
313 + escaped = self._json_escape_string_content(raw_content)
314 + frag = escaped + '"'
315 + self.streamed_args_for_tool[self.current_tool_id] += frag
316 + self._update_tool_args(pending_deltas, frag)
317 + continue
318 +
319 + # Check for partial </arg_value> at end
320 + safe_len = len(self._buffer)
321 + for i in range(1, len(self.arg_val_end)):
322 + if self._buffer.endswith(self.arg_val_end[:i]):
323 + safe_len = len(self._buffer) - i
324 + break
325 +
326 + if safe_len > 0:
327 + to_emit = self._buffer[:safe_len]
328 + self._buffer = self._buffer[safe_len:]
329 + escaped = self._json_escape_string_content(to_emit)
330 + if escaped:
331 + self.streamed_args_for_tool[self.current_tool_id] += escaped
332 + self._update_tool_args(pending_deltas, escaped)
333 + break
334 +
335 + # If we have a pending key, parse its value
336 + if self._pending_key is not None:
337 + val_pos = self._buffer.find(self.arg_val_start)
338 + if val_pos == -1:
339 + break
340 + if val_pos > 0:
341 + self._buffer = self._buffer[val_pos:]
342 +
343 + key = (self._pending_key or "").strip()
344 +
345 + is_string = self._is_string_type(
346 + self._current_tool_name, key, request.tools
347 + )
348 +
349 + if is_string:
350 + # String type: stream incrementally
351 + self._buffer = self._buffer[len(self.arg_val_start) :]
352 +
353 + if key in self._seen_keys[self.current_tool_id]:
354 + self._pending_key = None
355 + continue
356 +
357 + self._seen_keys[self.current_tool_id].add(key)
358 + key_json = json.dumps(key, ensure_ascii=False)
359 +
360 + if not self._args_started[self.current_tool_id]:
361 + frag = "{" + key_json + ': "'
362 + self._args_started[self.current_tool_id] = True
363 + else:
364 + frag = ", " + key_json + ': "'
365 +
366 + self.streamed_args_for_tool[self.current_tool_id] += frag
367 + self._streaming_string_value = True
368 + self._update_tool_args(pending_deltas, frag)
369 + continue
370 +
371 + # Non-string type: wait for complete value
372 + val_end = self._buffer.find(self.arg_val_end)
373 + if val_end == -1:
374 + break
375 +
376 + raw_val = self._buffer[len(self.arg_val_start) : val_end].strip()
377 + self._buffer = self._buffer[val_end + len(self.arg_val_end) :]
378 + self._pending_key = None
379 +
380 + frag_or_none = self._append_arg_fragment(key=key, raw_val=raw_val)
381 + if frag_or_none:
382 + self._update_tool_args(pending_deltas, frag_or_none)
383 + continue
384 +
385 + # Parse next arg or close
386 + end_pos = self._buffer.find(self.tool_call_end_token)
387 + key_pos = self._buffer.find(self.arg_key_start)
388 + if end_pos != -1 and (key_pos == -1 or end_pos < key_pos):
389 + self._buffer = self._buffer[end_pos + len(self.tool_call_end_token) :]
390 + frag_or_none = self._close_args_if_needed()
391 + # Finalize prev_tool_call_arr with complete parsed arguments
392 + if self._current_tool_name:
393 + try:
394 + full_args_str = self.streamed_args_for_tool[
395 + self.current_tool_id
396 + ]
397 + args_dict = json.loads(full_args_str)
398 + self.prev_tool_call_arr[self.current_tool_id] = {
399 + "name": self._current_tool_name,
400 + "arguments": args_dict,
401 + }
402 + except (json.JSONDecodeError, IndexError) as e:
403 + logger.warning(
404 + "Failed to finalize tool call state for tool %d: %s",
405 + self.current_tool_id,
406 + e,
407 + )
408 + self._finish_tool_call()
409 + if frag_or_none:
410 + self._update_tool_args(pending_deltas, frag_or_none)
411 + continue
412 +
413 + if key_pos == -1:
414 + break
415 + if key_pos > 0:
416 + self._buffer = self._buffer[key_pos:]
417 + key_end = self._buffer.find(self.arg_key_end)
418 + if key_end == -1:
419 + break
420 + key = self._buffer[len(self.arg_key_start) : key_end]
421 + self._buffer = self._buffer[key_end + len(self.arg_key_end) :]
422 + self._pending_key = key
423 + continue
424 +
425 + tool_calls = list(pending_deltas.values())
426 + if content is None and len(tool_calls) == 0:
427 + if request.logprobs:
428 + return DeltaMessage(content="")
429 + return None
430 + return DeltaMessage(content=content, tool_calls=tool_calls)
431 +
432 + def _ensure_tool_state(self) -> None:
433 + while len(self._tool_call_ids) <= self.current_tool_id:
434 + self._tool_call_ids.append(
435 + make_tool_call_id(id_type="random", func_name=None, idx=None)
436 + )
437 + while len(self.streamed_args_for_tool) <= self.current_tool_id:
438 + self.streamed_args_for_tool.append("")
439 + while len(self.prev_tool_call_arr) <= self.current_tool_id:
440 + self.prev_tool_call_arr.append({})
441 + while len(self._args_started) <= self.current_tool_id:
442 + self._args_started.append(False)
443 + while len(self._args_closed) <= self.current_tool_id:
444 + self._args_closed.append(False)
445 + while len(self._seen_keys) <= self.current_tool_id:
446 + self._seen_keys.append(set())
447 +
448 + def _begin_tool_call(self) -> None:
449 + if self.current_tool_id == -1:
450 + self.current_tool_id = 0
451 + else:
452 + self.current_tool_id += 1
453 + self._ensure_tool_state()
454 + self.current_tool_name_sent = False
455 + self._current_tool_name = None
456 + self._pending_key = None
457 + self._streaming_string_value = False
458 + self._in_tool_call = True
459 +
460 + def _finish_tool_call(self) -> None:
461 + self._in_tool_call = False
462 + self._current_tool_name = None
463 + self._pending_key = None
464 + self._streaming_string_value = False
465 +
466 + def _revert_last_tool_call_state(self) -> None:
467 + """Revert the state allocation for the last tool call."""
468 + if self.current_tool_id < 0:
469 + return
470 + self._tool_call_ids.pop()
471 + self.streamed_args_for_tool.pop()
472 + self.prev_tool_call_arr.pop()
473 + self._args_started.pop()
474 + self._args_closed.pop()
475 + self._seen_keys.pop()
476 + self.current_tool_id -= 1
477 +
478 + def _get_or_create_delta(self, pending: dict[int, DeltaToolCall]) -> DeltaToolCall:
479 + idx = self.current_tool_id
480 + if idx not in pending:
481 + pending[idx] = DeltaToolCall(
482 + index=idx,
483 + function=DeltaFunctionCall(),
484 + )
485 + delta = pending[idx]
486 + assert delta.function is not None
487 + return delta
488 +
489 + def _update_tool_name(
490 + self, pending: dict[int, DeltaToolCall], tool_name: str
491 + ) -> None:
492 + self.prev_tool_call_arr[self.current_tool_id] = {
493 + "name": self._current_tool_name,
494 + "arguments": {},
495 + }
496 + delta = self._get_or_create_delta(pending)
497 + delta.id = self._tool_call_ids[self.current_tool_id]
498 + delta.type = "function"
499 + assert delta.function is not None
500 + delta.function.name = tool_name
501 + if delta.function.arguments is None:
502 + delta.function.arguments = ""
503 +
504 + @staticmethod
505 + def _complete_json_prefix(
506 + json_prefix: str,
507 + allowed_partial_types: Allow,
508 + ) -> dict | None:
509 + """Complete a partial JSON prefix into a valid JSON object.
510 +
511 + Returns (formatted_prefix, parsed_dict) or None on failure.
512 +
513 + Note: ``partial_json_parser`` strips trailing whitespace before
514 + parsing (``complete.py:20``), which means the returned slice is
515 + shorter than ``json_prefix`` when it has trailing whitespace.
516 + Since the parser controls the construction of the json_prefix value,
517 + this code relies on it being a valid prefix and we only use the fix for
518 + the completion of the JSON object.
519 + """
520 + try:
521 + _, partial_str_completion = partial_json_parser.core.complete.fix(
522 + json_prefix,
523 + allowed_partial_types,
524 + )
525 + return json.loads(json_prefix + partial_str_completion)
526 + except Exception:
527 + return None
528 +
529 + def _update_tool_args(
530 + self, pending: dict[int, DeltaToolCall], fragment: str
531 + ) -> None:
532 + result = self._complete_json_prefix(
533 + self.streamed_args_for_tool[self.current_tool_id],
534 + Allow.ALL,
535 + )
536 + if result is not None:
537 + self.prev_tool_call_arr[self.current_tool_id]["arguments"] = result
538 + delta = self._get_or_create_delta(pending)
539 + assert delta.function is not None
540 + if delta.function.arguments is None:
541 + delta.function.arguments = ""
542 + delta.function.arguments += fragment
543 +
544 + def _append_arg_fragment(
545 + self,
546 + *,
547 + key: str,
548 + raw_val: str,
549 + ) -> str | None:
550 + key = key.strip()
551 + if not key:
552 + return None
553 + if key in self._seen_keys[self.current_tool_id]:
554 + return None
555 +
556 + # This function is only called for non-string types (already checked
557 + # by _is_string_type in the caller), so we always deserialize.
558 + val_obj: Any = self._deserialize(raw_val)
559 +
560 + key_json = json.dumps(key, ensure_ascii=False)
561 + val_json = json.dumps(val_obj, ensure_ascii=False)
562 +
563 + if not self._args_started[self.current_tool_id]:
564 + fragment = "{" + key_json + ": " + val_json
565 + self._args_started[self.current_tool_id] = True
566 + else:
567 + fragment = ", " + key_json + ": " + val_json
568 +
569 + self._seen_keys[self.current_tool_id].add(key)
570 + self.streamed_args_for_tool[self.current_tool_id] += fragment
571 + return fragment
572 +
573 + def _close_args_if_needed(self) -> str | None:
574 + if self._args_closed[self.current_tool_id]:
575 + return None
576 + self._args_closed[self.current_tool_id] = True
577 + if not self._args_started[self.current_tool_id]:
578 + fragment = "{}"
579 + self.streamed_args_for_tool[self.current_tool_id] = fragment
580 + else:
581 + fragment = "}"
582 + self.streamed_args_for_tool[self.current_tool_id] += fragment
583 + return fragment
vllm/transformers_utils/config.py
+28 -16
@@ -127,6 +127,7 @@ _CONFIG_REGISTRY: dict[str, type[PretrainedConfig]] = LazyConfigDict(
127 qwen3_next="Qwen3NextConfig",
128 qwen3_5="Qwen3_5Config",
129 qwen3_5_moe="Qwen3_5MoeConfig",
130 + laguna="LagunaConfig",
131 lfm2_moe="Lfm2MoeConfig",
132 tarsier2="Tarsier2Config",
133 )
@@ -409,22 +410,33 @@ def patch_rope_parameters(config: PretrainedConfig) -> None:
410 ompe = getattr(config, "original_max_position_embeddings", None)
411
412 if Version(version("transformers")) < Version("5.0.0"):
412 - # Transformers v4 installed, legacy config fields may be present
413 - if (rope_scaling := getattr(config, "rope_scaling", None)) is not None:
414 - config.rope_parameters = rope_scaling
415 - if (
416 - rope_theta is not None
417 - or partial_rotary_factor is not None
418 - or ompe is not None
419 - ) and not getattr(config, "rope_parameters", None):
420 - config.rope_parameters = {"rope_type": "default"}
421 - # Patch legacy fields into rope_parameters
422 - if rope_theta is not None:
423 - config.rope_parameters["rope_theta"] = rope_theta
424 - if partial_rotary_factor is not None:
425 - config.rope_parameters["partial_rotary_factor"] = partial_rotary_factor
426 - if ompe is not None:
427 - config.rope_parameters["original_max_position_embeddings"] = ompe
413 + # Transformers v4 installed, legacy config fields may be present.
414 + existing_rp = getattr(config, "rope_parameters", None)
415 + if isinstance(existing_rp, dict) and is_rope_parameters_nested(existing_rp):
416 + # Interleaved-attention models (e.g. Laguna-XS.2) ship a nested
417 + # {layer_type: {...}} rope_parameters that the model code indexes
418 + # by layer_type. The per-layer-type sub-dicts already carry the
419 + # correct rope_theta / partial_rotary_factor / ompe (the converter
420 + # places top-level legacy fields inside full_attention), so don't
421 + # merge top-level fields here — that would shadow the per-type
422 + # values and break sliding-attention layers.
423 + pass
424 + else:
425 + if (rope_scaling := getattr(config, "rope_scaling", None)) is not None:
426 + config.rope_parameters = rope_scaling
427 + if (
428 + rope_theta is not None
429 + or partial_rotary_factor is not None
430 + or ompe is not None
431 + ) and not getattr(config, "rope_parameters", None):
432 + config.rope_parameters = {"rope_type": "default"}
433 + # Patch legacy fields into rope_parameters
434 + if rope_theta is not None:
435 + config.rope_parameters["rope_theta"] = rope_theta
436 + if partial_rotary_factor is not None:
437 + config.rope_parameters["partial_rotary_factor"] = partial_rotary_factor
438 + if ompe is not None:
439 + config.rope_parameters["original_max_position_embeddings"] = ompe
440 elif rope_theta is not None or getattr(config, "rope_parameters", None):
441 # Transformers v5 installed
442 # Patch these fields in case they used non-standard names
vllm/transformers_utils/configs/__init__.py
+2
@@ -45,6 +45,7 @@ _CLASS_TO_MODULE: dict[str, str] = {
45 # `FalconConfig` class from the official HuggingFace transformers library.
46 "RWConfig": "vllm.transformers_utils.configs.falcon",
47 "JAISConfig": "vllm.transformers_utils.configs.jais",
48 + "LagunaConfig": "vllm.transformers_utils.configs.laguna",
49 "Lfm2MoeConfig": "vllm.transformers_utils.configs.lfm2_moe",
50 "MedusaConfig": "vllm.transformers_utils.configs.medusa",
51 "MiDashengLMConfig": "vllm.transformers_utils.configs.midashenglm",
@@ -105,6 +106,7 @@ __all__ = [
106 "IsaacConfig",
107 "RWConfig",
108 "JAISConfig",
109 + "LagunaConfig",
110 "Lfm2MoeConfig",
111 "MedusaConfig",
112 "MiDashengLMConfig",
vllm/transformers_utils/configs/laguna.py new
+120
@@ -0,0 +1,120 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +from transformers.configuration_utils import PretrainedConfig
5 +
6 +
7 +class LagunaConfig(PretrainedConfig):
8 + model_type = "laguna"
9 + keys_to_ignore_at_inference = ["past_key_values"]
10 + base_model_tp_plan = {
11 + "layers.*.self_attn.q_proj": "colwise",
12 + "layers.*.self_attn.k_proj": "colwise",
13 + "layers.*.self_attn.v_proj": "colwise",
14 + "layers.*.self_attn.g_proj": "colwise",
15 + "layers.*.self_attn.o_proj": "rowwise",
16 + "layers.*.mlp.gate_proj": "colwise",
17 + "layers.*.mlp.up_proj": "colwise",
18 + "layers.*.mlp.down_proj": "rowwise",
19 + }
20 + base_model_pp_plan = {
21 + "embed_tokens": (["input_ids"], ["inputs_embeds"]),
22 + "layers": (["hidden_states", "attention_mask"], ["hidden_states"]),
23 + "norm": (["hidden_states"], ["hidden_states"]),
24 + }
25 +
26 + def __init__(
27 + self,
28 + vocab_size: int = 100352,
29 + hidden_size: int = 2048,
30 + intermediate_size: int = 8192,
31 + num_hidden_layers: int = 40,
32 + num_attention_heads: int = 48,
33 + num_key_value_heads: int = 8,
34 + head_dim: int = 128,
35 + qkv_bias: bool = False,
36 + attention_bias: bool = False,
37 + gating: bool | str = True,
38 + hidden_act: str = "silu",
39 + max_position_embeddings: int = 131072,
40 + initializer_range: float = 0.02,
41 + rms_norm_eps: float = 1e-6,
42 + use_cache: bool = True,
43 + tie_word_embeddings: bool = False,
44 + rope_theta: float = 500000.0,
45 + rope_scaling: dict | None = None,
46 + rope_parameters: dict | None = None,
47 + partial_rotary_factor: float = 1.0,
48 + attention_dropout: float = 0.0,
49 + sliding_window: int | None = None,
50 + layer_types: list[str] | None = None,
51 + swa_attention_sink_enabled: bool = False,
52 + swa_rope_parameters: dict | None = None,
53 + num_attention_heads_per_layer: list[int] | None = None,
54 + num_experts: int = 256,
55 + num_experts_per_tok: int = 8,
56 + moe_intermediate_size: int = 512,
57 + shared_expert_intermediate_size: int = 512,
58 + norm_topk_prob: bool = True,
59 + decoder_sparse_step: int = 1,
60 + mlp_only_layers: list[int] | None = None,
61 + router_aux_loss_coef: float = 0.001,
62 + output_router_logits: bool = False,
63 + moe_routed_scaling_factor: float = 1.0,
64 + moe_apply_router_weight_on_input: bool = False,
65 + **kwargs,
66 + ):
67 + if mlp_only_layers is None:
68 + mlp_only_layers = [0]
69 +
70 + # Accept either v4-style (rope_theta + rope_scaling) or v5-style
71 + # (rope_parameters). Translate v5 → v4 so downstream code has one path.
72 + if rope_parameters is not None:
73 + rp = dict(rope_parameters)
74 + rope_theta = float(rp.pop("rope_theta", rope_theta))
75 + rt = rp.pop("rope_type", None)
76 + if rt is not None and rt != "default":
77 + rope_scaling = {"rope_type": rt, **rp}
78 + elif rp and rope_scaling is None:
79 + rope_scaling = {"rope_type": "default", **rp}
80 +
81 + self.vocab_size = vocab_size
82 + self.hidden_size = hidden_size
83 + self.intermediate_size = intermediate_size
84 + self.num_hidden_layers = num_hidden_layers
85 + self.num_attention_heads = num_attention_heads
86 + self.num_key_value_heads = num_key_value_heads
87 + self.head_dim = head_dim
88 + self.qkv_bias = qkv_bias
89 + self.attention_bias = attention_bias
90 + self.gating = gating
91 + self.hidden_act = hidden_act
92 + self.max_position_embeddings = max_position_embeddings
93 + self.initializer_range = initializer_range
94 + self.rms_norm_eps = rms_norm_eps
95 + self.use_cache = use_cache
96 + self.rope_theta = rope_theta
97 + self.rope_scaling = rope_scaling
98 + self.partial_rotary_factor = partial_rotary_factor
99 + self.attention_dropout = attention_dropout
100 + self.sliding_window = sliding_window
101 + self.layer_types = layer_types
102 + self.swa_attention_sink_enabled = swa_attention_sink_enabled
103 + self.swa_rope_parameters = swa_rope_parameters
104 + self.num_attention_heads_per_layer = num_attention_heads_per_layer
105 + self.num_experts = num_experts
106 + self.num_experts_per_tok = num_experts_per_tok
107 + self.moe_intermediate_size = moe_intermediate_size
108 + self.shared_expert_intermediate_size = shared_expert_intermediate_size
109 + self.norm_topk_prob = norm_topk_prob
110 + self.decoder_sparse_step = decoder_sparse_step
111 + self.mlp_only_layers = mlp_only_layers
112 + self.router_aux_loss_coef = router_aux_loss_coef
113 + self.output_router_logits = output_router_logits
114 + self.moe_routed_scaling_factor = moe_routed_scaling_factor
115 + self.moe_apply_router_weight_on_input = moe_apply_router_weight_on_input
116 +
117 + super().__init__(tie_word_embeddings=tie_word_embeddings, **kwargs)
118 +
119 +
120 +__all__ = ["LagunaConfig"]