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 10, 11:30 UTC
[Model] Support IQuestCoder model (#31575)
Signed-off-by: yxing <yxing@iquestlab.com>
yxing-bj committed
Jan 8, 2026 at 22:42 UTC
fe86be66c555af2110f65318d006dc94f2bba2d2
4 files changed
+605
docs/models/supported_models.md
+2
@@ -406,6 +406,8 @@ th {
406
| `InternLMForCausalLM` | InternLM | `internlm/internlm-7b`, `internlm/internlm-chat-7b`, etc. | ✅︎ | ✅︎ |
407
| `InternLM2ForCausalLM` | InternLM2 | `internlm/internlm2-7b`, `internlm/internlm2-chat-7b`, etc. | ✅︎ | ✅︎ |
408
| `InternLM3ForCausalLM` | InternLM3 | `internlm/internlm3-8b-instruct`, etc. | ✅︎ | ✅︎ |
409
+| `IQuestCoderForCausalLM` | IQuestCoderV1 | `IQuestLab/IQuest-Coder-V1-40B-Instruct`, etc. | | |
410
+| `IQuestLoopCoderForCausalLM` | IQuestLoopCoderV1 | `IQuestLab/IQuest-Coder-V1-40B-Loop-Instruct`, etc. | | |
411
| `JAISLMHeadModel` | Jais | `inceptionai/jais-13b`, `inceptionai/jais-13b-chat`, `inceptionai/jais-30b-v3`, `inceptionai/jais-30b-chat-v3`, etc. | | ✅︎ |
412
| `Jais2ForCausalLM` | Jais2 | `inceptionai/Jais-2-8B-Chat`, `inceptionai/Jais-2-70B-Chat`, etc. | | ✅︎ |
413
| `JambaForCausalLM` | Jamba | `ai21labs/AI21-Jamba-1.5-Large`, `ai21labs/AI21-Jamba-1.5-Mini`, `ai21labs/Jamba-v0.1`, etc. | ✅︎ | ✅︎ |
tests/models/registry.py
+6
@@ -306,6 +306,12 @@ _TEXT_GENERATION_EXAMPLE_MODELS = {
306
"InternLM3ForCausalLM": _HfExamplesInfo(
307
"internlm/internlm3-8b-instruct", trust_remote_code=True
308
),
309
+ "IQuestCoderForCausalLM": _HfExamplesInfo(
310
+ "IQuestLab/IQuest-Coder-V1-40B-Instruct", trust_remote_code=True
311
+ ),
312
+ "IQuestLoopCoderForCausalLM": _HfExamplesInfo(
313
+ "IQuestLab/IQuest-Coder-V1-40B-Loop-Instruct", trust_remote_code=True
314
+ ),
315
"JAISLMHeadModel": _HfExamplesInfo("inceptionai/jais-13b-chat"),
316
"Jais2ForCausalLM": _HfExamplesInfo(
317
"inceptionai/Jais-2-8B-Chat", min_transformers_version="4.58"
vllm/model_executor/models/iquest_loopcoder.py
new
+595
@@ -0,0 +1,595 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+#
4
+# Licensed under the Apache License, Version 2.0 (the "License");
5
+# you may not use this file except in compliance with the License.
6
+# You may obtain a copy of the License at
7
+#
8
+# http://www.apache.org/licenses/LICENSE-2.0
9
+#
10
+# Unless required by applicable law or agreed to in writing, software
11
+# distributed under the License is distributed on an "AS IS" BASIS,
12
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13
+# See the License for the specific language governing permissions and
14
+# limitations under the License.
15
+"""Inference-only LoopCoder model compatible with HuggingFace weights."""
16
+
17
+from __future__ import annotations
18
+
19
+from collections.abc import Iterable
20
+from dataclasses import replace
21
+from typing import Any
22
+
23
+import torch
24
+from torch import nn
25
+from transformers import PretrainedConfig
26
+
27
+from vllm.attention.backends.abstract import AttentionType
28
+from vllm.attention.layer import Attention
29
+from vllm.compilation.decorators import support_torch_compile
30
+from vllm.config import CacheConfig, VllmConfig
31
+from vllm.distributed import get_tensor_model_parallel_world_size
32
+from vllm.model_executor.layers.layernorm import RMSNorm
33
+from vllm.model_executor.layers.linear import (
34
+ ColumnParallelLinear,
35
+ QKVParallelLinear,
36
+ RowParallelLinear,
37
+)
38
+from vllm.model_executor.layers.logits_processor import LogitsProcessor
39
+from vllm.model_executor.layers.quantization import QuantizationConfig
40
+from vllm.model_executor.layers.rotary_embedding import get_rope
41
+from vllm.model_executor.layers.vocab_parallel_embedding import (
42
+ ParallelLMHead,
43
+ VocabParallelEmbedding,
44
+)
45
+from vllm.model_executor.model_loader.weight_utils import (
46
+ default_weight_loader,
47
+ maybe_remap_kv_scale_name,
48
+)
49
+from vllm.model_executor.models.llama import LlamaMLP
50
+from vllm.sequence import IntermediateTensors
51
+
52
+from .utils import (
53
+ AutoWeightsLoader,
54
+ extract_layer_index,
55
+ make_layers,
56
+ maybe_prefix,
57
+)
58
+
59
+
60
+class LoopCoderAttention(nn.Module):
61
+ def __init__(
62
+ self,
63
+ config: PretrainedConfig,
64
+ hidden_size: int,
65
+ num_heads: int,
66
+ num_kv_heads: int,
67
+ max_position: int = 4096 * 32,
68
+ cache_config: CacheConfig | None = None,
69
+ quant_config: QuantizationConfig | None = None,
70
+ prefix: str = "",
71
+ attn_type: str = AttentionType.DECODER,
72
+ dual_chunk_attention_config: dict[str, Any] | None = None,
73
+ layer_idx: int = 0,
74
+ ) -> None:
75
+ super().__init__()
76
+ self.layer_idx = layer_idx
77
+ self.hidden_size = hidden_size
78
+ tp_size = get_tensor_model_parallel_world_size()
79
+ self.total_num_heads = num_heads
80
+ assert self.total_num_heads % tp_size == 0
81
+ self.num_heads = self.total_num_heads // tp_size
82
+ self.total_num_kv_heads = num_kv_heads
83
+ if self.total_num_kv_heads >= tp_size:
84
+ # Number of KV heads is greater than TP size, so we partition
85
+ # the KV heads across multiple tensor parallel GPUs.
86
+ assert self.total_num_kv_heads % tp_size == 0
87
+ else:
88
+ # Number of KV heads is less than TP size, so we replicate
89
+ # the KV heads across multiple tensor parallel GPUs.
90
+ assert tp_size % self.total_num_kv_heads == 0
91
+ self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size)
92
+ self.head_dim = hidden_size // self.total_num_heads
93
+ self.q_size = self.num_heads * self.head_dim
94
+ self.kv_size = self.num_kv_heads * self.head_dim
95
+ self.scaling = self.head_dim**-0.5
96
+ self.dual_chunk_attention_config = dual_chunk_attention_config
97
+
98
+ # Get loop_num from config, default to 2 if not specified
99
+ self.loop_num = getattr(config, "loop_num", 2)
100
+
101
+ self.loop_window_size = getattr(config, "loop_window_size", 64)
102
+
103
+ # Use total number of hidden layers instead of hardcoded 24
104
+ total_layers = config.num_hidden_layers
105
+
106
+ self.qkv_proj = QKVParallelLinear(
107
+ hidden_size,
108
+ self.head_dim,
109
+ self.total_num_heads,
110
+ self.total_num_kv_heads,
111
+ bias=False,
112
+ quant_config=quant_config,
113
+ prefix=f"{prefix}.qkv_proj",
114
+ )
115
+ self.o_proj = RowParallelLinear(
116
+ self.total_num_heads * self.head_dim,
117
+ hidden_size,
118
+ bias=False,
119
+ quant_config=quant_config,
120
+ prefix=f"{prefix}.o_proj",
121
+ )
122
+
123
+ self.rotary_emb = get_rope(
124
+ self.head_dim,
125
+ max_position=max_position,
126
+ rope_parameters=config.rope_parameters,
127
+ dual_chunk_attention_config=dual_chunk_attention_config,
128
+ )
129
+ self.attn = nn.ModuleList()
130
+
131
+ base_cache_config = cache_config
132
+
133
+ for loop_idx in range(self.loop_num):
134
+ base_layer_idx = extract_layer_index(prefix)
135
+ unique_layer_idx = loop_idx * total_layers + base_layer_idx
136
+
137
+ unique_prefix = prefix.replace(
138
+ f"layers.{base_layer_idx}", f"layers.{unique_layer_idx}"
139
+ )
140
+
141
+ if loop_idx == 0:
142
+ loop_cache_config = cache_config
143
+ else:
144
+ if base_cache_config is not None:
145
+ loop_cache_config = replace(
146
+ base_cache_config,
147
+ sliding_window=self.loop_window_size,
148
+ )
149
+ else:
150
+ loop_cache_config = CacheConfig(
151
+ sliding_window=self.loop_window_size,
152
+ cache_dtype="auto",
153
+ )
154
+
155
+ self.attn.append(
156
+ Attention(
157
+ self.num_heads,
158
+ self.head_dim,
159
+ self.scaling,
160
+ num_kv_heads=self.num_kv_heads,
161
+ cache_config=loop_cache_config,
162
+ quant_config=quant_config,
163
+ attn_type=attn_type,
164
+ prefix=f"{unique_prefix}.attn",
165
+ **{
166
+ "layer_idx": unique_layer_idx,
167
+ "dual_chunk_attention_config": dual_chunk_attention_config,
168
+ }
169
+ if dual_chunk_attention_config and loop_idx == 0
170
+ else {},
171
+ )
172
+ )
173
+
174
+ def forward(
175
+ self,
176
+ positions: torch.Tensor,
177
+ hidden_states: torch.Tensor,
178
+ loop_idx: int,
179
+ gate_proj: LoopGateProjection | None = None,
180
+ ) -> torch.Tensor:
181
+ if loop_idx == 0:
182
+ attn = self.attn[0]
183
+ qkv, _ = self.qkv_proj(hidden_states)
184
+ q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
185
+ q, k = self.rotary_emb(positions, q, k)
186
+ attn_output = attn(q, k, v)
187
+ output, _ = self.o_proj(attn_output)
188
+ return output
189
+ else:
190
+ global_attn = self.attn[0]
191
+ local_attn = self.attn[loop_idx]
192
+ qkv, _ = self.qkv_proj(hidden_states)
193
+ q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
194
+ q, k = self.rotary_emb(positions, q, k)
195
+ num_tokens, _ = q.shape
196
+ num_heads = self.num_heads
197
+ head_dim = self.head_dim
198
+
199
+ q_reshaped = q.view(num_tokens, num_heads, head_dim).transpose(0, 1)
200
+
201
+ global_attn_output = global_attn(q, None, None)
202
+ local_attn_output = local_attn(q, k, v)
203
+ assert gate_proj is not None, "gate_proj must be provided for loop_idx > 0"
204
+ gate = gate_proj(q_reshaped)
205
+ output = global_attn_output * gate + local_attn_output * (1 - gate)
206
+ output, _ = self.o_proj(output)
207
+ return output
208
+
209
+
210
+class LoopCoderDecoderLayer(nn.Module):
211
+ def __init__(
212
+ self,
213
+ config: PretrainedConfig,
214
+ cache_config: CacheConfig | None = None,
215
+ quant_config: QuantizationConfig | None = None,
216
+ prefix: str = "",
217
+ layer_idx: int = 0,
218
+ ) -> None:
219
+ super().__init__()
220
+ self.hidden_size = config.hidden_size
221
+ dual_chunk_attention_config = getattr(
222
+ config, "dual_chunk_attention_config", None
223
+ )
224
+ self.layer_idx = layer_idx
225
+ if getattr(config, "is_causal", True):
226
+ attn_type = AttentionType.DECODER
227
+ else:
228
+ attn_type = AttentionType.ENCODER_ONLY
229
+
230
+ self.self_attn = LoopCoderAttention(
231
+ config=config,
232
+ hidden_size=self.hidden_size,
233
+ num_heads=config.num_attention_heads,
234
+ max_position=config.max_position_embeddings,
235
+ num_kv_heads=config.num_key_value_heads,
236
+ cache_config=cache_config,
237
+ quant_config=quant_config,
238
+ prefix=f"{prefix}.self_attn",
239
+ attn_type=attn_type,
240
+ dual_chunk_attention_config=dual_chunk_attention_config,
241
+ layer_idx=self.layer_idx,
242
+ )
243
+ self.mlp = LlamaMLP(
244
+ hidden_size=self.hidden_size,
245
+ intermediate_size=config.intermediate_size,
246
+ hidden_act=config.hidden_act,
247
+ quant_config=quant_config,
248
+ prefix=f"{prefix}.mlp",
249
+ )
250
+ self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
251
+ self.post_attention_layernorm = RMSNorm(
252
+ config.hidden_size, eps=config.rms_norm_eps
253
+ )
254
+
255
+ def forward(
256
+ self,
257
+ positions: torch.Tensor,
258
+ hidden_states: torch.Tensor,
259
+ loop_idx: int,
260
+ gate_proj: LoopGateProjection | None = None,
261
+ ) -> torch.Tensor:
262
+ residual = hidden_states
263
+ hidden_states = self.input_layernorm(hidden_states)
264
+ hidden_states = self.self_attn(
265
+ positions=positions,
266
+ hidden_states=hidden_states,
267
+ loop_idx=loop_idx,
268
+ gate_proj=gate_proj,
269
+ )
270
+ hidden_states = hidden_states + residual
271
+ residual = hidden_states
272
+ hidden_states = self.post_attention_layernorm(hidden_states)
273
+ hidden_states = self.mlp(hidden_states)
274
+ hidden_states = hidden_states + residual
275
+
276
+ return hidden_states
277
+
278
+
279
+class LoopGateProjection(nn.Module):
280
+ """Gate projection for mixed attention in Loop 2+.
281
+
282
+ Computes: g = sigmoid(linear(Q)) for each head independently.
283
+ This gate determines how much to use Loop1's KV (global) vs current
284
+ loop's KV (local).
285
+
286
+ Supports tensor parallelism: each GPU handles a subset of heads.
287
+ The weight matrix has shape [num_heads, head_dim] and is split along
288
+ the head dimension.
289
+ """
290
+
291
+ def __init__(
292
+ self,
293
+ total_num_heads: int,
294
+ head_dim: int,
295
+ quant_config: QuantizationConfig | None = None,
296
+ prefix: str = "",
297
+ ):
298
+ super().__init__()
299
+ self.total_num_heads = total_num_heads
300
+ self.head_dim = head_dim
301
+ tp_size = get_tensor_model_parallel_world_size()
302
+ assert self.total_num_heads % tp_size == 0
303
+ self.num_heads = self.total_num_heads // tp_size
304
+
305
+ self.gate_proj = ColumnParallelLinear(
306
+ head_dim,
307
+ self.total_num_heads,
308
+ bias=True,
309
+ gather_output=False,
310
+ quant_config=quant_config,
311
+ prefix=f"{prefix}.gate_proj",
312
+ )
313
+
314
+ def forward(self, query: torch.Tensor) -> torch.Tensor:
315
+ """Compute gate values from query tensor.
316
+
317
+ Args:
318
+ query: [num_heads, num_tokens, head_dim] (vLLM flattened format)
319
+ where num_heads is the number of heads on this TP rank
320
+ and num_tokens = batch * seq_len
321
+
322
+ Returns:
323
+ gate: [num_tokens, num_heads * head_dim] (flattened format matching q shape)
324
+ """
325
+ num_heads, num_tokens, head_dim = query.shape
326
+
327
+ assert num_heads == self.num_heads, (
328
+ f"Expected {self.num_heads} heads, got {num_heads}"
329
+ )
330
+
331
+ query_flat = query.reshape(-1, head_dim)
332
+
333
+ gate_logits_flat, _ = self.gate_proj(query_flat)
334
+
335
+ gate_logits = gate_logits_flat.reshape(
336
+ num_heads, num_tokens, self.num_heads
337
+ ) # [num_heads, num_tokens, num_heads]
338
+
339
+ # Extract diagonal: each head h's query should use output column h
340
+ # gate_logits[h, :, h] gives the output for head h at each token
341
+ gate_logits = torch.diagonal(
342
+ gate_logits, dim1=0, dim2=2
343
+ ) # [num_tokens, num_heads]
344
+ gate_logits = gate_logits.transpose(0, 1) # [num_heads, num_tokens]
345
+ gate_logits = gate_logits.unsqueeze(-1) # [num_heads, num_tokens, 1]
346
+
347
+ # Apply sigmoid
348
+ gate = torch.sigmoid(gate_logits) # [num_heads, num_tokens, 1]
349
+
350
+ # Expand and reshape to match q shape: [num_tokens, num_heads * head_dim]
351
+ gate = gate.transpose(0, 1) # [num_tokens, num_heads, 1]
352
+ gate = gate.expand(-1, -1, head_dim) # [num_tokens, num_heads, head_dim]
353
+ gate = gate.reshape(
354
+ num_tokens, num_heads * head_dim
355
+ ) # [num_tokens, num_heads * head_dim]
356
+
357
+ return gate
358
+
359
+
360
+@support_torch_compile(
361
+ dynamic_arg_dims={
362
+ "input_ids": 0,
363
+ "positions": -1,
364
+ "intermediate_tensors": 0,
365
+ "inputs_embeds": 0,
366
+ }
367
+)
368
+class IQuestLoopCoderModel(nn.Module):
369
+ def __init__(
370
+ self,
371
+ *,
372
+ vllm_config: VllmConfig,
373
+ prefix: str = "",
374
+ decoder_layer_type: type[nn.Module] = LoopCoderDecoderLayer,
375
+ ):
376
+ super().__init__()
377
+
378
+ config = vllm_config.model_config.hf_config
379
+ cache_config = vllm_config.cache_config
380
+ quant_config = vllm_config.quant_config
381
+
382
+ # TODO (@robertgshaw2): see if this can be moved out
383
+ if cache_config.sliding_window is not None and hasattr(
384
+ config, "max_window_layers"
385
+ ):
386
+ assert config.max_window_layers == config.num_hidden_layers, (
387
+ "Sliding window for some but all layers is not supported. "
388
+ "This model uses sliding window but `max_window_layers` = {} "
389
+ "is less than `num_hidden_layers` = {}. Please open an issue "
390
+ "to discuss this feature.".format(
391
+ config.max_window_layers,
392
+ config.num_hidden_layers,
393
+ )
394
+ )
395
+
396
+ self.config = config
397
+ self.quant_config = quant_config
398
+ self.vocab_size = config.vocab_size
399
+
400
+ self.embed_tokens = VocabParallelEmbedding(
401
+ config.vocab_size,
402
+ config.hidden_size,
403
+ quant_config=quant_config,
404
+ prefix=f"{prefix}.embed_tokens",
405
+ )
406
+
407
+ self.loop_num = getattr(self.config, "loop_num", 2)
408
+ self.window_size = getattr(self.config, "loop_window_size", 64)
409
+
410
+ # Gate projections for Loop 2+ (one per layer)
411
+ head_dim = config.hidden_size // config.num_attention_heads
412
+ _, _, self.gate_projections = make_layers(
413
+ config.num_hidden_layers,
414
+ lambda prefix: LoopGateProjection(
415
+ total_num_heads=config.num_attention_heads,
416
+ head_dim=head_dim,
417
+ quant_config=quant_config,
418
+ prefix=prefix,
419
+ ),
420
+ prefix=f"{prefix}.gate_projections",
421
+ )
422
+
423
+ self.start_layer, self.end_layer, self.layers = make_layers(
424
+ config.num_hidden_layers,
425
+ lambda prefix: LoopCoderDecoderLayer(
426
+ config=config,
427
+ cache_config=cache_config,
428
+ quant_config=quant_config,
429
+ prefix=prefix,
430
+ layer_idx=extract_layer_index(prefix),
431
+ ),
432
+ prefix=f"{prefix}.layers",
433
+ )
434
+ self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
435
+
436
+ def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
437
+ return self.embed_tokens(input_ids)
438
+
439
+ def forward(
440
+ self,
441
+ input_ids: torch.Tensor,
442
+ positions: torch.Tensor,
443
+ intermediate_tensors: IntermediateTensors | None = None,
444
+ inputs_embeds: torch.Tensor | None = None,
445
+ ) -> torch.Tensor | IntermediateTensors:
446
+ if inputs_embeds is not None:
447
+ hidden_states = inputs_embeds
448
+ else:
449
+ hidden_states = self.embed_input_ids(input_ids)
450
+
451
+ for loop_idx in range(self.loop_num):
452
+ for layer_idx, layer in enumerate(
453
+ self.layers[self.start_layer : self.end_layer]
454
+ ):
455
+ # Get the actual layer index (accounting for pipeline parallelism)
456
+ actual_layer_idx = self.start_layer + layer_idx
457
+ # Get gate_proj for this layer (only for loop_idx > 0)
458
+ gate_proj = (
459
+ self.gate_projections[actual_layer_idx] if loop_idx > 0 else None
460
+ )
461
+ hidden_states = layer(positions, hidden_states, loop_idx, gate_proj)
462
+ hidden_states = self.norm(hidden_states)
463
+ return hidden_states
464
+
465
+ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
466
+ stacked_params_mapping = [
467
+ # (param_name, shard_name, shard_id)
468
+ ("qkv_proj", "q_proj", "q"),
469
+ ("qkv_proj", "k_proj", "k"),
470
+ ("qkv_proj", "v_proj", "v"),
471
+ ("gate_up_proj", "gate_proj", 0),
472
+ ("gate_up_proj", "up_proj", 1),
473
+ ]
474
+ params_dict = dict(self.named_parameters(remove_duplicate=False))
475
+ loaded_params: set[str] = set()
476
+ for name, loaded_weight in weights:
477
+ if "rotary_emb.inv_freq" in name:
478
+ continue
479
+ if self.quant_config is not None and (
480
+ scale_name := self.quant_config.get_cache_scale(name)
481
+ ):
482
+ # Loading kv cache quantization scales
483
+ param = params_dict[scale_name]
484
+ weight_loader = getattr(param, "weight_loader", default_weight_loader)
485
+ loaded_weight = (
486
+ loaded_weight if loaded_weight.dim() == 0 else loaded_weight[0]
487
+ )
488
+ weight_loader(param, loaded_weight)
489
+ loaded_params.add(scale_name)
490
+ continue
491
+ for param_name, weight_name, shard_id in stacked_params_mapping:
492
+ if "gate_projections" in name:
493
+ continue
494
+ if weight_name not in name:
495
+ continue
496
+ name = name.replace(weight_name, param_name)
497
+ # Skip loading extra bias for GPTQ models.
498
+ if name.endswith(".bias") and name not in params_dict:
499
+ continue
500
+ if name.endswith("scale"):
501
+ # Remapping the name of FP8 kv-scale.
502
+ name = maybe_remap_kv_scale_name(name, params_dict)
503
+ if name is None:
504
+ continue
505
+ param = params_dict[name]
506
+ weight_loader = getattr(param, "weight_loader", default_weight_loader)
507
+ if weight_loader == default_weight_loader:
508
+ weight_loader(param, loaded_weight)
509
+ else:
510
+ weight_loader(param, loaded_weight, shard_id)
511
+ break
512
+ else:
513
+ if name.startswith("gate_projections."):
514
+ if name.endswith(".weight"):
515
+ vllm_name = name.replace(".weight", ".gate_proj.weight")
516
+ elif name.endswith(".bias"):
517
+ vllm_name = name.replace(".bias", ".gate_proj.bias")
518
+ else:
519
+ continue
520
+
521
+ if vllm_name in params_dict:
522
+ param = params_dict[vllm_name]
523
+ weight_loader = getattr(
524
+ param, "weight_loader", default_weight_loader
525
+ )
526
+ weight_loader(param, loaded_weight)
527
+ loaded_params.add(vllm_name)
528
+ continue
529
+ continue
530
+ if name.endswith(".bias") and name not in params_dict:
531
+ continue
532
+ # Remapping the name of FP8 kv-scale.
533
+ name = maybe_remap_kv_scale_name(name, params_dict)
534
+ if name is None:
535
+ continue
536
+ param = params_dict[name]
537
+ weight_loader = getattr(param, "weight_loader", default_weight_loader)
538
+ weight_loader(param, loaded_weight)
539
+ loaded_params.add(name)
540
+ return loaded_params
541
+
542
+
543
+class IQuestLoopCoderForCausalLM(nn.Module):
544
+ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
545
+ super().__init__()
546
+ config = vllm_config.model_config.hf_config
547
+ quant_config = vllm_config.quant_config
548
+
549
+ self.config = config
550
+
551
+ self.quant_config = quant_config
552
+ self.model = IQuestLoopCoderModel(
553
+ vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
554
+ )
555
+
556
+ if config.tie_word_embeddings:
557
+ self.lm_head = self.model.embed_tokens
558
+ else:
559
+ self.lm_head = ParallelLMHead(
560
+ config.vocab_size,
561
+ config.hidden_size,
562
+ quant_config=quant_config,
563
+ prefix=maybe_prefix(prefix, "lm_head"),
564
+ )
565
+
566
+ self.logits_processor = LogitsProcessor(config.vocab_size)
567
+
568
+ def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
569
+ return self.model.embed_input_ids(input_ids)
570
+
571
+ def forward(
572
+ self,
573
+ input_ids: torch.Tensor,
574
+ positions: torch.Tensor,
575
+ intermediate_tensors: IntermediateTensors | None = None,
576
+ inputs_embeds: torch.Tensor | None = None,
577
+ ) -> torch.Tensor | IntermediateTensors:
578
+ hidden_states = self.model(
579
+ input_ids, positions, intermediate_tensors, inputs_embeds
580
+ )
581
+ return hidden_states
582
+
583
+ def compute_logits(
584
+ self,
585
+ hidden_states: torch.Tensor,
586
+ ) -> torch.Tensor | None:
587
+ logits = self.logits_processor(self.lm_head, hidden_states)
588
+ return logits
589
+
590
+ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
591
+ loader = AutoWeightsLoader(
592
+ self,
593
+ skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None),
594
+ )
595
+ return loader.load_weights(weights)
vllm/model_executor/models/registry.py
+2
@@ -128,6 +128,8 @@ _TEXT_GENERATION_MODELS = {
128
"InternLM2ForCausalLM": ("internlm2", "InternLM2ForCausalLM"),
129
"InternLM2VEForCausalLM": ("internlm2_ve", "InternLM2VEForCausalLM"),
130
"InternLM3ForCausalLM": ("llama", "LlamaForCausalLM"),
131
+ "IQuestCoderForCausalLM": ("llama", "LlamaForCausalLM"),
132
+ "IQuestLoopCoderForCausalLM": ("iquest_loopcoder", "IQuestLoopCoderForCausalLM"),
133
"JAISLMHeadModel": ("jais", "JAISLMHeadModel"),
134
"Jais2ForCausalLM": ("jais2", "Jais2ForCausalLM"),
135
"JambaForCausalLM": ("jamba", "JambaForCausalLM"),