@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 905d76b51d
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 7, 14:00 UTC

[Model] Add huggingface skt/A.X-K1 model (#32407)

Signed-off-by: Sungwan(Alex) Kim <sw0726.kim@sktelecom.com> Signed-off-by: fort726 <38447663+fort726@users.noreply.github.com> Co-authored-by: Sungwan(Alex) Kim <sw0726.kim@sktelecom.com> Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> Co-authored-by: TJian <tunjian.tan@embeddedllm.com>

fort726 committed Feb 28, 2026 at 02:26 UTC 905d76b51dc3b98c4d0ee35317493f21f9e6b5d0
7 files changed +1396 -1
docs/models/supported_models.md
+1
@@ -369,6 +369,7 @@ th {
369 | `AquilaForCausalLM` | Aquila, Aquila2 | `BAAI/Aquila-7B`, `BAAI/AquilaChat-7B`, etc. | ✅︎ | ✅︎ |
370 | `ArceeForCausalLM` | Arcee (AFM) | `arcee-ai/AFM-4.5B-Base`, etc. | ✅︎ | ✅︎ |
371 | `ArcticForCausalLM` | Arctic | `Snowflake/snowflake-arctic-base`, `Snowflake/snowflake-arctic-instruct`, etc. | | ✅︎ |
372 +| `AXK1ForCausalLM` | A.X-K1 | `skt/A.X-K1`, etc. | | ✅︎ |
373 | `BaiChuanForCausalLM` | Baichuan2, Baichuan | `baichuan-inc/Baichuan2-13B-Chat`, `baichuan-inc/Baichuan-7B`, etc. | ✅︎ | ✅︎ |
374 | `BailingMoeForCausalLM` | Ling | `inclusionAI/Ling-lite-1.5`, `inclusionAI/Ling-plus`, etc. | ✅︎ | ✅︎ |
375 | `BailingMoeV2ForCausalLM` | Ling | `inclusionAI/Ling-mini-2.0`, etc. | ✅︎ | ✅︎ |
tests/models/registry.py
+1
@@ -194,6 +194,7 @@ _TEXT_GENERATION_EXAMPLE_MODELS = {
194 "ArcticForCausalLM": _HfExamplesInfo(
195 "Snowflake/snowflake-arctic-instruct", trust_remote_code=True
196 ),
197 + "AXK1ForCausalLM": _HfExamplesInfo("skt/A.X-K1", trust_remote_code=True),
198 "BaiChuanForCausalLM": _HfExamplesInfo(
199 "baichuan-inc/Baichuan-7B", trust_remote_code=True
200 ),
vllm/model_executor/models/AXK1.py new
+1168
@@ -0,0 +1,1168 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +# Adapted from
5 +# https://github.com/huggingface/transformers/blob/v4.28.0/src/transformers/models/llama/modeling_llama.py
6 +# Copyright 2023 The vLLM team.
7 +#
8 +# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX
9 +# and OPT implementations in this library. It has been modified from its
10 +# original forms to accommodate minor architectural differences compared
11 +# to GPT-NeoX and OPT used by the Meta AI team that trained the model.
12 +#
13 +# Licensed under the Apache License, Version 2.0 (the "License");
14 +# you may not use this file except in compliance with the License.
15 +# You may obtain a copy of the License at
16 +#
17 +# http://www.apache.org/licenses/LICENSE-2.0
18 +#
19 +# Unless required by applicable law or agreed to in writing, software
20 +# distributed under the License is distributed on an "AS IS" BASIS,
21 +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
22 +# See the License for the specific language governing permissions and
23 +# limitations under the License.
24 +"""Inference-only A.X K1 model."""
25 +
26 +import typing
27 +from collections.abc import Callable, Iterable
28 +from itertools import islice
29 +
30 +import torch
31 +from torch import nn
32 +
33 +from vllm._aiter_ops import rocm_aiter_ops
34 +from vllm.compilation.decorators import support_torch_compile
35 +from vllm.config import CacheConfig, ParallelConfig, VllmConfig
36 +from vllm.distributed import (
37 + get_ep_group,
38 + get_pp_group,
39 + get_tensor_model_parallel_rank,
40 + get_tensor_model_parallel_world_size,
41 + tensor_model_parallel_all_gather,
42 +)
43 +from vllm.logger import init_logger
44 +from vllm.model_executor.layers.attention import Attention
45 +from vllm.model_executor.layers.fused_moe import SharedFusedMoE
46 +from vllm.model_executor.layers.layernorm import RMSNorm
47 +from vllm.model_executor.layers.linear import (
48 + ColumnParallelLinear,
49 + MergedColumnParallelLinear,
50 + ReplicatedLinear,
51 + RowParallelLinear,
52 +)
53 +from vllm.model_executor.layers.logits_processor import LogitsProcessor
54 +from vllm.model_executor.layers.mla import MLAModules, MultiHeadLatentAttentionWrapper
55 +from vllm.model_executor.layers.quantization import QuantizationConfig
56 +from vllm.model_executor.layers.rotary_embedding import get_rope
57 +from vllm.model_executor.layers.vocab_parallel_embedding import (
58 + ParallelLMHead,
59 + VocabParallelEmbedding,
60 +)
61 +from vllm.model_executor.model_loader.weight_utils import (
62 + default_weight_loader,
63 + maybe_remap_kv_scale_name,
64 +)
65 +from vllm.model_executor.models.deepseek_v2 import (
66 + DeepseekAttention,
67 + DeepseekV2MLP,
68 + yarn_get_mscale,
69 +)
70 +from vllm.model_executor.models.utils import sequence_parallel_chunk
71 +from vllm.platforms import current_platform
72 +from vllm.sequence import IntermediateTensors
73 +from vllm.transformers_utils.configs.AXK1 import AXK1Config
74 +
75 +from .interfaces import MixtureOfExperts, SupportsEagle, SupportsLoRA, SupportsPP
76 +from .utils import (
77 + PPMissingLayer,
78 + is_pp_missing_parameter,
79 + make_empty_intermediate_tensors_factory,
80 + make_layers,
81 + maybe_prefix,
82 +)
83 +
84 +logger = init_logger(__name__)
85 +
86 +
87 +class AXK1MLP(DeepseekV2MLP):
88 + pass
89 +
90 +
91 +class AXK1MoE(nn.Module):
92 + def __init__(
93 + self,
94 + config: AXK1Config,
95 + parallel_config: ParallelConfig,
96 + quant_config: QuantizationConfig | None = None,
97 + prefix: str = "",
98 + ):
99 + super().__init__()
100 + self.tp_size = get_tensor_model_parallel_world_size()
101 + self.tp_rank = get_tensor_model_parallel_rank()
102 +
103 + self.routed_scaling_factor = config.routed_scaling_factor
104 +
105 + self.ep_group = get_ep_group().device_group
106 + self.ep_rank = get_ep_group().rank_in_group
107 + self.ep_size = self.ep_group.size()
108 + self.n_routed_experts: int = config.n_routed_experts
109 + self.n_shared_experts: int = config.n_shared_experts
110 +
111 + self.is_sequence_parallel = parallel_config.use_sequence_parallel_moe
112 +
113 + if config.hidden_act != "silu":
114 + raise ValueError(
115 + f"Unsupported activation: {config.hidden_act}. "
116 + "Only silu is supported for now."
117 + )
118 +
119 + self.gate = ReplicatedLinear(
120 + config.hidden_size,
121 + config.n_routed_experts,
122 + bias=False,
123 + quant_config=None,
124 + prefix=f"{prefix}.gate",
125 + )
126 + if config.topk_method == "noaux_tc":
127 + self.gate.e_score_correction_bias = nn.Parameter(
128 + torch.empty(config.n_routed_experts, dtype=torch.float32)
129 + )
130 + else:
131 + self.gate.e_score_correction_bias = None
132 +
133 + # Load balancing settings.
134 + eplb_config = parallel_config.eplb_config
135 + self.enable_eplb = parallel_config.enable_eplb
136 +
137 + self.n_redundant_experts = eplb_config.num_redundant_experts
138 + self.n_logical_experts = self.n_routed_experts
139 + self.n_physical_experts = self.n_logical_experts + self.n_redundant_experts
140 + self.n_local_physical_experts = self.n_physical_experts // self.ep_size
141 +
142 + self.physical_expert_start = self.ep_rank * self.n_local_physical_experts
143 + self.physical_expert_end = (
144 + self.physical_expert_start + self.n_local_physical_experts
145 + )
146 +
147 + self.is_rocm_aiter_moe_enabled = rocm_aiter_ops.is_fused_moe_enabled()
148 + self.is_fusion_moe_shared_experts_enabled = (
149 + rocm_aiter_ops.is_fusion_moe_shared_experts_enabled()
150 + )
151 + if config.n_shared_experts is None or self.is_fusion_moe_shared_experts_enabled:
152 + self.shared_experts = None
153 + else:
154 + intermediate_size = config.moe_intermediate_size * config.n_shared_experts
155 +
156 + self.shared_experts = AXK1MLP(
157 + hidden_size=config.hidden_size,
158 + intermediate_size=intermediate_size,
159 + hidden_act=config.hidden_act,
160 + quant_config=quant_config,
161 + is_sequence_parallel=self.is_sequence_parallel,
162 + reduce_results=False,
163 + prefix=f"{prefix}.shared_experts",
164 + )
165 +
166 + self.experts = SharedFusedMoE(
167 + shared_experts=self.shared_experts,
168 + gate=self.gate,
169 + num_experts=config.n_routed_experts,
170 + top_k=config.num_experts_per_tok,
171 + hidden_size=config.hidden_size,
172 + intermediate_size=config.moe_intermediate_size,
173 + reduce_results=False,
174 + renormalize=config.norm_topk_prob,
175 + quant_config=quant_config,
176 + use_grouped_topk=True,
177 + num_expert_group=config.n_group,
178 + topk_group=config.topk_group,
179 + prefix=f"{prefix}.experts",
180 + scoring_func=config.scoring_func,
181 + # we do scaling outside, set factor to 1.0 to avoid double mul
182 + # aiter applies routed_scaling_factor internally
183 + routed_scaling_factor=1.0
184 + if not self.is_rocm_aiter_moe_enabled
185 + else self.routed_scaling_factor,
186 + e_score_correction_bias=self.gate.e_score_correction_bias,
187 + enable_eplb=self.enable_eplb,
188 + num_redundant_experts=self.n_redundant_experts,
189 + is_sequence_parallel=self.is_sequence_parallel,
190 + n_shared_experts=config.n_shared_experts
191 + if self.is_fusion_moe_shared_experts_enabled
192 + else None,
193 + )
194 +
195 + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
196 + num_tokens, hidden_dim = hidden_states.shape
197 + hidden_states = hidden_states.view(-1, hidden_dim)
198 +
199 + # Chunk the hidden states so they aren't replicated across TP ranks.
200 + # This avoids duplicate computation in self.experts.
201 + # TODO: We can replace the all_reduce at the end of attn with a
202 + # reduce_scatter instead of chunking here.
203 + if self.is_sequence_parallel:
204 + hidden_states = sequence_parallel_chunk(hidden_states)
205 +
206 + if self.experts.is_internal_router:
207 + # In this case, the gate/router runs inside the FusedMoE class
208 + fused_moe_out = self.experts(
209 + hidden_states=hidden_states, router_logits=hidden_states
210 + )
211 + else:
212 + # router_logits: (num_tokens, n_experts)
213 + router_logits, _ = self.gate(hidden_states)
214 + fused_moe_out = self.experts(
215 + hidden_states=hidden_states, router_logits=router_logits
216 + )
217 +
218 + shared_output, final_hidden_states = fused_moe_out
219 + if self.shared_experts is None:
220 + assert shared_output is None
221 +
222 + # Fix FP16 overflow
223 + # See AXK1DecoderLayer for more details.
224 + if hidden_states.dtype != torch.float16:
225 + if not self.is_rocm_aiter_moe_enabled:
226 + final_hidden_states *= self.routed_scaling_factor
227 + elif self.shared_experts is not None:
228 + assert shared_output is not None
229 + shared_output *= 1.0 / self.routed_scaling_factor
230 +
231 + if self.shared_experts is not None:
232 + assert shared_output is not None
233 + final_hidden_states += shared_output
234 +
235 + if self.is_sequence_parallel:
236 + final_hidden_states = tensor_model_parallel_all_gather(
237 + final_hidden_states, 0
238 + )
239 + final_hidden_states = final_hidden_states[:num_tokens]
240 + elif self.tp_size > 1:
241 + final_hidden_states = self.experts.maybe_all_reduce_tensor_model_parallel(
242 + final_hidden_states
243 + )
244 +
245 + return final_hidden_states.view(num_tokens, hidden_dim)
246 +
247 +
248 +def _get_llama_4_scaling(
249 + original_max_position_embeddings: int, scaling_beta: float, positions: torch.Tensor
250 +) -> torch.Tensor:
251 + scaling = 1 + scaling_beta * torch.log(
252 + 1 + torch.floor(positions / original_max_position_embeddings)
253 + )
254 + # Broadcast over num_heads and head_dim
255 + return scaling[..., None, None]
256 +
257 +
258 +class AXK1Attention(nn.Module):
259 + def __init__(
260 + self,
261 + vllm_config: VllmConfig,
262 + config: AXK1Config,
263 + hidden_size: int,
264 + num_heads: int,
265 + qk_nope_head_dim: int,
266 + qk_rope_head_dim: int,
267 + v_head_dim: int,
268 + q_lora_rank: int,
269 + kv_lora_rank: int,
270 + max_position_embeddings: int = 8192,
271 + cache_config: CacheConfig | None = None,
272 + quant_config: QuantizationConfig | None = None,
273 + topk_indices_buffer: torch.Tensor | None = None,
274 + prefix: str = "",
275 + ) -> None:
276 + super().__init__()
277 + self.hidden_size = hidden_size
278 + self.qk_nope_head_dim = qk_nope_head_dim
279 + self.qk_rope_head_dim = qk_rope_head_dim
280 + self.qk_head_dim = qk_nope_head_dim + qk_rope_head_dim
281 + self.v_head_dim = v_head_dim
282 + self.q_lora_rank = q_lora_rank
283 + self.kv_lora_rank = kv_lora_rank
284 + self.num_heads = num_heads
285 + tp_size = get_tensor_model_parallel_world_size()
286 + assert num_heads % tp_size == 0
287 + self.num_local_heads = num_heads // tp_size
288 + self.scaling = self.qk_head_dim**-0.5
289 + self.max_position_embeddings = max_position_embeddings
290 + assert topk_indices_buffer is None, (
291 + "topk_indices_buffer is not \
292 + supported for AXK1Attention"
293 + )
294 +
295 + if self.q_lora_rank is not None:
296 + self.q_a_proj = ReplicatedLinear(
297 + self.hidden_size,
298 + self.q_lora_rank,
299 + bias=False,
300 + quant_config=quant_config,
301 + prefix=f"{prefix}.q_a_proj",
302 + )
303 + self.q_a_layernorm = RMSNorm(self.q_lora_rank, eps=config.rms_norm_eps)
304 + self.q_b_proj = ColumnParallelLinear(
305 + q_lora_rank,
306 + self.num_heads * self.qk_head_dim,
307 + bias=False,
308 + quant_config=quant_config,
309 + prefix=f"{prefix}.q_b_proj",
310 + )
311 + else:
312 + self.q_proj = ColumnParallelLinear(
313 + self.hidden_size,
314 + self.num_heads * self.qk_head_dim,
315 + bias=False,
316 + quant_config=quant_config,
317 + prefix=f"{prefix}.q_proj",
318 + )
319 +
320 + self.kv_a_proj_with_mqa = ReplicatedLinear(
321 + self.hidden_size,
322 + self.kv_lora_rank + self.qk_rope_head_dim,
323 + bias=False,
324 + quant_config=quant_config,
325 + prefix=f"{prefix}.kv_a_proj_with_mqa",
326 + )
327 + self.kv_a_layernorm = RMSNorm(self.kv_lora_rank, eps=config.rms_norm_eps)
328 + self.kv_b_proj = ColumnParallelLinear(
329 + self.kv_lora_rank,
330 + self.num_heads * (self.qk_nope_head_dim + self.v_head_dim),
331 + bias=False,
332 + quant_config=quant_config,
333 + prefix=f"{prefix}.kv_b_proj",
334 + )
335 + # O projection.
336 + self.o_proj = RowParallelLinear(
337 + self.num_heads * self.v_head_dim,
338 + self.hidden_size,
339 + bias=False,
340 + quant_config=quant_config,
341 + prefix=f"{prefix}.o_proj",
342 + )
343 + if config.rope_parameters["rope_type"] != "default":
344 + config.rope_parameters["rope_type"] = (
345 + "deepseek_yarn"
346 + if config.rope_parameters.get("apply_yarn_scaling", True)
347 + else "deepseek_llama_scaling"
348 + )
349 +
350 + self.rotary_emb = get_rope(
351 + qk_rope_head_dim,
352 + max_position=max_position_embeddings,
353 + rope_parameters=config.rope_parameters,
354 + is_neox_style=False,
355 + )
356 +
357 + if config.rope_parameters["rope_type"] == "deepseek_yarn":
358 + mscale_all_dim = config.rope_parameters.get("mscale_all_dim", False)
359 + scaling_factor = config.rope_parameters["factor"]
360 + mscale = yarn_get_mscale(scaling_factor, float(mscale_all_dim))
361 + self.scaling = self.scaling * mscale * mscale
362 +
363 + self.attn = Attention(
364 + self.num_local_heads,
365 + self.qk_head_dim,
366 + self.scaling,
367 + num_kv_heads=self.num_local_heads,
368 + cache_config=cache_config,
369 + quant_config=quant_config,
370 + prefix=f"{prefix}.attn",
371 + )
372 +
373 + def forward(
374 + self,
375 + positions: torch.Tensor,
376 + hidden_states: torch.Tensor,
377 + llama_4_scaling: torch.Tensor | None,
378 + ) -> torch.Tensor:
379 + if self.q_lora_rank is not None:
380 + q = self.q_a_proj(hidden_states)[0]
381 + q = self.q_a_layernorm(q)
382 + q = self.q_b_proj(q)[0].view(-1, self.num_local_heads, self.qk_head_dim)
383 + else:
384 + q = self.q_proj(hidden_states)[0].view(
385 + -1, self.num_local_heads, self.qk_head_dim
386 + )
387 + q_nope, q_pe = q.split([self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1)
388 + latent_cache = self.kv_a_proj_with_mqa(hidden_states)[0]
389 + kv_a, _ = latent_cache.split([self.kv_lora_rank, self.qk_rope_head_dim], dim=-1)
390 + latent_cache = latent_cache.unsqueeze(1)
391 + kv_a = self.kv_a_layernorm(kv_a)
392 + kv = self.kv_b_proj(kv_a)[0]
393 + kv = kv.view(-1, self.num_local_heads, self.qk_nope_head_dim + self.v_head_dim)
394 + k_nope, v = kv.split([self.qk_nope_head_dim, self.v_head_dim], dim=-1)
395 + k_pe = latent_cache[:, :, self.kv_lora_rank :]
396 +
397 + q_pe, k_pe = self.rotary_emb(positions, q_pe, k_pe)
398 +
399 + q[..., self.qk_nope_head_dim :] = q_pe
400 + k = torch.empty_like(q)
401 + k[..., : self.qk_nope_head_dim] = k_nope
402 + k[..., self.qk_nope_head_dim :] = k_pe
403 +
404 + # Apply llama 4 scaling if provided
405 + if llama_4_scaling is not None:
406 + q *= llama_4_scaling
407 +
408 + # padding value to qk_head_dim for alignment
409 + v = torch.nn.functional.pad(
410 + v, [0, self.qk_head_dim - self.v_head_dim], value=0
411 + ).view(-1, self.num_local_heads * self.qk_head_dim)
412 + attn_output = self.attn(q, k, v)
413 + attn_output = attn_output.view(-1, self.num_local_heads, self.qk_head_dim)[
414 + ..., : self.v_head_dim
415 + ].reshape(-1, self.num_local_heads * self.v_head_dim)
416 + output, _ = self.o_proj(attn_output)
417 + return output
418 +
419 +
420 +class AXK1MLAAttention(nn.Module):
421 + """
422 + Main reference: DeepseekV2 paper, and FlashInfer Implementation
423 + (https://arxiv.org/abs/2405.04434 and https://github.com/flashinfer-ai/flashinfer/pull/551).
424 +
425 + For more info see MLACommonImpl in:
426 + vllm/v1/attention/backends/mla/utils.py
427 + """
428 +
429 + def __init__(
430 + self,
431 + vllm_config: VllmConfig,
432 + config: AXK1Config,
433 + hidden_size: int,
434 + num_heads: int,
435 + qk_nope_head_dim: int,
436 + qk_rope_head_dim: int,
437 + v_head_dim: int,
438 + q_lora_rank: int | None,
439 + kv_lora_rank: int,
440 + max_position_embeddings: int = 8192,
441 + cache_config: CacheConfig | None = None,
442 + quant_config: QuantizationConfig | None = None,
443 + prefix: str = "",
444 + topk_indices_buffer: torch.Tensor | None = None,
445 + ) -> None:
446 + super().__init__()
447 + self.hidden_size = hidden_size
448 + self.qk_nope_head_dim = qk_nope_head_dim
449 + self.qk_rope_head_dim = qk_rope_head_dim
450 + self.qk_head_dim = qk_nope_head_dim + qk_rope_head_dim
451 + self.v_head_dim = v_head_dim
452 +
453 + self.q_lora_rank = q_lora_rank
454 + self.kv_lora_rank = kv_lora_rank
455 +
456 + self.num_heads = num_heads
457 + tp_size = get_tensor_model_parallel_world_size()
458 + assert num_heads % tp_size == 0
459 + self.num_local_heads = num_heads // tp_size
460 +
461 + self.scaling = self.qk_head_dim**-0.5
462 + self.max_position_embeddings = max_position_embeddings
463 +
464 + if self.q_lora_rank is not None:
465 + self.fused_qkv_a_proj = MergedColumnParallelLinear(
466 + self.hidden_size,
467 + [self.q_lora_rank, self.kv_lora_rank + self.qk_rope_head_dim],
468 + bias=False,
469 + quant_config=quant_config,
470 + prefix=f"{prefix}.fused_qkv_a_proj",
471 + disable_tp=True,
472 + )
473 + else:
474 + self.kv_a_proj_with_mqa = ReplicatedLinear(
475 + self.hidden_size,
476 + self.kv_lora_rank + self.qk_rope_head_dim,
477 + bias=False,
478 + quant_config=quant_config,
479 + prefix=f"{prefix}.kv_a_proj_with_mqa",
480 + )
481 +
482 + if self.q_lora_rank is not None:
483 + self.q_a_layernorm = RMSNorm(self.q_lora_rank, eps=config.rms_norm_eps)
484 + self.q_b_proj = ColumnParallelLinear(
485 + self.q_lora_rank,
486 + self.num_heads * self.qk_head_dim,
487 + bias=False,
488 + quant_config=quant_config,
489 + prefix=f"{prefix}.q_b_proj",
490 + )
491 + else:
492 + self.q_proj = ColumnParallelLinear(
493 + self.hidden_size,
494 + self.num_heads * self.qk_head_dim,
495 + bias=False,
496 + quant_config=quant_config,
497 + prefix=f"{prefix}.q_proj",
498 + )
499 + self.kv_a_layernorm = RMSNorm(self.kv_lora_rank, eps=config.rms_norm_eps)
500 + self.kv_b_proj = ColumnParallelLinear(
501 + self.kv_lora_rank,
502 + self.num_heads * (self.qk_nope_head_dim + self.v_head_dim),
503 + bias=False,
504 + quant_config=quant_config,
505 + prefix=f"{prefix}.kv_b_proj",
506 + )
507 + self.o_proj = RowParallelLinear(
508 + self.num_heads * self.v_head_dim,
509 + self.hidden_size,
510 + bias=False,
511 + quant_config=quant_config,
512 + prefix=f"{prefix}.o_proj",
513 + )
514 +
515 + if config.rope_parameters["rope_type"] != "default":
516 + config.rope_parameters["rope_type"] = (
517 + "deepseek_yarn"
518 + if config.rope_parameters.get("apply_yarn_scaling", True)
519 + else "deepseek_llama_scaling"
520 + )
521 +
522 + self.rotary_emb = get_rope(
523 + qk_rope_head_dim,
524 + max_position=max_position_embeddings,
525 + rope_parameters=config.rope_parameters,
526 + is_neox_style=False,
527 + )
528 +
529 + if config.rope_parameters["rope_type"] == "deepseek_yarn":
530 + mscale_all_dim = config.rope_parameters.get("mscale_all_dim", False)
531 + scaling_factor = config.rope_parameters["factor"]
532 + mscale = yarn_get_mscale(scaling_factor, float(mscale_all_dim))
533 + self.scaling = self.scaling * mscale * mscale
534 +
535 + mla_modules = MLAModules(
536 + kv_a_layernorm=self.kv_a_layernorm,
537 + kv_b_proj=self.kv_b_proj,
538 + rotary_emb=self.rotary_emb,
539 + o_proj=self.o_proj,
540 + fused_qkv_a_proj=self.fused_qkv_a_proj
541 + if self.q_lora_rank is not None
542 + else None,
543 + kv_a_proj_with_mqa=self.kv_a_proj_with_mqa
544 + if self.q_lora_rank is None
545 + else None,
546 + q_a_layernorm=self.q_a_layernorm if self.q_lora_rank is not None else None,
547 + q_b_proj=self.q_b_proj if self.q_lora_rank is not None else None,
548 + q_proj=self.q_proj if self.q_lora_rank is None else None,
549 + indexer=None,
550 + indexer_rotary_emb=None,
551 + is_sparse=False,
552 + topk_indices_buffer=topk_indices_buffer,
553 + )
554 +
555 + self.mla_attn = MultiHeadLatentAttentionWrapper(
556 + self.hidden_size,
557 + self.num_local_heads,
558 + self.scaling,
559 + self.qk_nope_head_dim,
560 + self.qk_rope_head_dim,
561 + self.v_head_dim,
562 + self.q_lora_rank,
563 + self.kv_lora_rank,
564 + mla_modules,
565 + cache_config,
566 + quant_config,
567 + prefix,
568 + )
569 +
570 + def forward(
571 + self,
572 + positions: torch.Tensor,
573 + hidden_states: torch.Tensor,
574 + llama_4_scaling: torch.Tensor | None,
575 + ) -> torch.Tensor:
576 + return self.mla_attn(positions, hidden_states, llama_4_scaling)
577 +
578 +
579 +class AXK1DecoderLayer(nn.Module):
580 + def __init__(
581 + self,
582 + vllm_config: VllmConfig,
583 + prefix: str,
584 + config: AXK1Config | None = None,
585 + ) -> None:
586 + super().__init__()
587 +
588 + if config is None:
589 + config = vllm_config.model_config.hf_config
590 + model_config = vllm_config.model_config
591 + cache_config = vllm_config.cache_config
592 + quant_config = vllm_config.quant_config
593 + parallel_config = vllm_config.parallel_config
594 + self.config = config
595 +
596 + self.hidden_size = config.hidden_size
597 + max_position_embeddings = config.max_position_embeddings
598 + # DecoderLayers are created with `make_layers` which passes the prefix
599 + # with the layer's index.
600 + layer_idx = int(prefix.split(sep=".")[-1])
601 + self.layer_idx = layer_idx
602 +
603 + # verify MLA attention specific fields
604 + qk_nope_head_dim = config.qk_nope_head_dim
605 + qk_rope_head_dim = config.qk_rope_head_dim
606 + v_head_dim = config.v_head_dim
607 + kv_lora_rank = config.kv_lora_rank
608 + use_mha = all(dim == 0 for dim in (qk_nope_head_dim, qk_rope_head_dim))
609 + self.use_mha = use_mha
610 +
611 + if use_mha:
612 + attn_cls = DeepseekAttention
613 + elif model_config.use_mla:
614 + attn_cls = AXK1MLAAttention
615 + else:
616 + attn_cls = AXK1Attention
617 + self.self_attn = attn_cls(
618 + vllm_config=vllm_config,
619 + config=config,
620 + hidden_size=self.hidden_size,
621 + num_heads=config.num_attention_heads,
622 + qk_nope_head_dim=qk_nope_head_dim,
623 + qk_rope_head_dim=qk_rope_head_dim,
624 + v_head_dim=v_head_dim,
625 + q_lora_rank=config.q_lora_rank,
626 + kv_lora_rank=kv_lora_rank,
627 + max_position_embeddings=max_position_embeddings,
628 + cache_config=cache_config,
629 + quant_config=quant_config,
630 + prefix=f"{prefix}.self_attn",
631 + topk_indices_buffer=None,
632 + )
633 +
634 + self.is_layer_sparse = self._is_layer_sparse()
635 + if self.is_layer_sparse:
636 + self.mlp = AXK1MoE(
637 + config=config,
638 + parallel_config=parallel_config,
639 + quant_config=quant_config,
640 + prefix=f"{prefix}.mlp",
641 + )
642 + else:
643 + self.mlp = AXK1MLP(
644 + hidden_size=config.hidden_size,
645 + intermediate_size=config.intermediate_size,
646 + hidden_act=config.hidden_act,
647 + quant_config=quant_config,
648 + prefix=f"{prefix}.mlp",
649 + )
650 + self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
651 + self.post_attention_layernorm = RMSNorm(
652 + config.hidden_size, eps=config.rms_norm_eps
653 + )
654 + self.post_mlp_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
655 + self.routed_scaling_factor = config.routed_scaling_factor
656 +
657 + def _is_layer_sparse(self) -> bool:
658 + return (
659 + self.config.n_routed_experts is not None
660 + and self.layer_idx >= self.config.first_k_dense_replace
661 + and self.layer_idx % self.config.moe_layer_freq == 0
662 + )
663 +
664 + def forward(
665 + self,
666 + positions: torch.Tensor,
667 + hidden_states: torch.Tensor,
668 + residual: torch.Tensor | None,
669 + llama_4_scaling: torch.Tensor | None = None,
670 + ) -> tuple[torch.Tensor, torch.Tensor]:
671 + # Self Attention
672 + if residual is None:
673 + residual = hidden_states.clone()
674 + hidden_states = self.input_layernorm(hidden_states)
675 + else:
676 + hidden_states, residual = self.input_layernorm(hidden_states, residual)
677 +
678 + attn_kwargs = {
679 + "positions": positions,
680 + "hidden_states": hidden_states,
681 + }
682 + if not self.use_mha:
683 + attn_kwargs["llama_4_scaling"] = llama_4_scaling
684 + hidden_states = self.self_attn(**attn_kwargs)
685 +
686 + if (
687 + not isinstance(self.self_attn, DeepseekAttention)
688 + and hidden_states.dtype == torch.float16
689 + ):
690 + # Fix FP16 overflow
691 + # We scale both hidden_states and residual before
692 + # rmsnorm, and rmsnorm result would not affect by scale.
693 + hidden_states *= 1.0 / self.routed_scaling_factor
694 + if self.layer_idx == 0:
695 + # The residual is shared by all layers, we only scale it on
696 + # first layer.
697 + residual *= 1.0 / self.routed_scaling_factor
698 +
699 + # Fully Connected
700 + hidden_states, residual = self.post_attention_layernorm(hidden_states, residual)
701 + hidden_states = self.mlp(hidden_states)
702 +
703 + if self.is_layer_sparse:
704 + hidden_states = self.post_mlp_layernorm(hidden_states)
705 +
706 + if isinstance(self.mlp, AXK1MLP) and hidden_states.dtype == torch.float16:
707 + # Fix FP16 overflow
708 + # Scaling the AXK1MLP output, it is the input of
709 + # input_layernorm of next decoder layer.
710 + # The scaling of AXK1MOE output would be done in the forward
711 + # of AXK1MOE
712 + hidden_states *= 1.0 / self.routed_scaling_factor
713 +
714 + return hidden_states, residual
715 +
716 +
717 +@support_torch_compile
718 +class AXK1Model(nn.Module):
719 + fall_back_to_pt_during_load = False
720 +
721 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
722 + super().__init__()
723 +
724 + config: AXK1Config = vllm_config.model_config.hf_config
725 + quant_config = vllm_config.quant_config
726 + self.config = config
727 + self.device = current_platform.device_type
728 + self.vocab_size = config.vocab_size
729 +
730 + if get_pp_group().is_first_rank:
731 + self.embed_tokens = VocabParallelEmbedding(
732 + config.vocab_size,
733 + config.hidden_size,
734 + quant_config=quant_config,
735 + prefix=f"{prefix}.embed_tokens",
736 + )
737 + else:
738 + self.embed_tokens = PPMissingLayer()
739 + self.start_layer, self.end_layer, self.layers = make_layers(
740 + config.num_hidden_layers,
741 + lambda prefix: AXK1DecoderLayer(vllm_config, prefix),
742 + prefix=f"{prefix}.layers",
743 + )
744 +
745 + if get_pp_group().is_last_rank:
746 + self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
747 + else:
748 + self.norm = PPMissingLayer()
749 + self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory(
750 + ["hidden_states", "residual"], config.hidden_size
751 + )
752 +
753 + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
754 + return self.embed_tokens(input_ids)
755 +
756 + def forward(
757 + self,
758 + input_ids: torch.Tensor | None,
759 + positions: torch.Tensor,
760 + intermediate_tensors: IntermediateTensors | None,
761 + inputs_embeds: torch.Tensor | None = None,
762 + ) -> torch.Tensor | IntermediateTensors:
763 + if get_pp_group().is_first_rank:
764 + if inputs_embeds is not None:
765 + hidden_states = inputs_embeds
766 + else:
767 + hidden_states = self.embed_input_ids(input_ids)
768 + residual = None
769 + else:
770 + assert intermediate_tensors is not None
771 + hidden_states = intermediate_tensors["hidden_states"]
772 + residual = intermediate_tensors["residual"]
773 +
774 + # Compute llama 4 scaling once per forward pass if enabled
775 + llama_4_scaling_config = getattr(self.config, "llama_4_scaling", None)
776 + llama_4_scaling: torch.Tensor | None
777 + if llama_4_scaling_config is not None:
778 + llama_4_scaling = _get_llama_4_scaling(
779 + original_max_position_embeddings=llama_4_scaling_config[
780 + "original_max_position_embeddings"
781 + ],
782 + scaling_beta=llama_4_scaling_config["beta"],
783 + positions=positions,
784 + )
785 + else:
786 + llama_4_scaling = None
787 +
788 + for layer in islice(self.layers, self.start_layer, self.end_layer):
789 + hidden_states, residual = layer(
790 + positions, hidden_states, residual, llama_4_scaling
791 + )
792 +
793 + if not get_pp_group().is_last_rank:
794 + return IntermediateTensors(
795 + {"hidden_states": hidden_states, "residual": residual}
796 + )
797 +
798 + hidden_states, _ = self.norm(hidden_states, residual)
799 + return hidden_states
800 +
801 +
802 +class AXK1MixtureOfExperts(MixtureOfExperts):
803 + moe_mlp_layers: list[AXK1MoE]
804 + """
805 + List of MoE MLP layers in the model.
806 + """
807 +
808 + def extract_moe_parameters(self, example_moe: AXK1MoE | None):
809 + if example_moe is None:
810 + self.num_moe_layers = 0
811 + self.num_expert_groups = 0
812 + self.num_logical_experts = 0
813 + self.num_physical_experts = 0
814 + self.num_local_physical_experts = 0
815 + self.num_routed_experts = 0
816 + self.num_shared_experts = 0
817 + self.num_redundant_experts = 0
818 + logger.warning("AXK1: No AXK1MoE layer found in model.layers.")
819 + else:
820 + self.num_logical_experts = example_moe.n_logical_experts
821 + self.num_physical_experts = example_moe.n_physical_experts
822 + self.num_local_physical_experts = example_moe.n_local_physical_experts
823 + self.num_routed_experts = example_moe.n_routed_experts
824 + self.num_shared_experts = example_moe.n_shared_experts
825 + self.num_redundant_experts = example_moe.n_redundant_experts
826 +
827 + def update_physical_experts_metadata(
828 + self,
829 + num_physical_experts: int,
830 + num_local_physical_experts: int,
831 + ) -> None:
832 + assert self.num_local_physical_experts == num_local_physical_experts
833 + self.num_physical_experts = num_physical_experts
834 + self.num_local_physical_experts = num_local_physical_experts
835 + self.num_redundant_experts = num_physical_experts - self.num_logical_experts
836 + for moe in self.moe_mlp_layers:
837 + moe.n_local_physical_experts = num_local_physical_experts
838 + moe.n_physical_experts = num_physical_experts
839 + moe.n_redundant_experts = self.num_redundant_experts
840 + moe.experts.update_expert_map()
841 +
842 +
843 +class AXK1ForCausalLM(
844 + nn.Module, SupportsPP, AXK1MixtureOfExperts, SupportsLoRA, SupportsEagle
845 +):
846 + packed_modules_mapping = {
847 + "gate_up_proj": ["gate_proj", "up_proj"],
848 + }
849 + model_cls = AXK1Model
850 +
851 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
852 + super().__init__()
853 + config: AXK1Config = vllm_config.model_config.hf_config
854 + quant_config = vllm_config.quant_config
855 + self.config = config
856 + self.quant_config = quant_config
857 +
858 + qk_nope_head_dim = config.qk_nope_head_dim
859 + qk_rope_head_dim = config.qk_rope_head_dim
860 + self.use_mha = all(dim == 0 for dim in (qk_nope_head_dim, qk_rope_head_dim))
861 +
862 + if self.use_mha:
863 + self.packed_modules_mapping["qkv_proj"] = ["q_proj", "k_proj", "v_proj"]
864 +
865 + # `packed_modules_mapping` needs to be modified before
866 + # initializing AXK1Model, as it is passed inplace to
867 + # quantization config init and may be used to select the
868 + # quant_method for relevant layers during initialization.
869 + self.fuse_qkv_a_proj = config.q_lora_rank is not None
870 + if self.fuse_qkv_a_proj:
871 + self.packed_modules_mapping["fused_qkv_a_proj"] = [
872 + "q_a_proj",
873 + "kv_a_proj_with_mqa",
874 + ]
875 +
876 + self.model = self.model_cls(
877 + vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
878 + )
879 + if get_pp_group().is_last_rank:
880 + self.lm_head = ParallelLMHead(
881 + config.vocab_size,
882 + config.hidden_size,
883 + quant_config=quant_config,
884 + prefix=maybe_prefix(prefix, "lm_head"),
885 + )
886 + else:
887 + self.lm_head = PPMissingLayer()
888 + self.logits_processor = LogitsProcessor(config.vocab_size)
889 + self.make_empty_intermediate_tensors = (
890 + self.model.make_empty_intermediate_tensors
891 + )
892 + # Set MoE hyperparameters
893 + self.num_moe_layers = (
894 + self.config.num_hidden_layers - self.config.first_k_dense_replace
895 + )
896 + self.set_moe_parameters()
897 +
898 + def set_moe_parameters(self):
899 + self.expert_weights = []
900 +
901 + self.num_expert_groups = getattr(self.config, "n_group", 1)
902 +
903 + self.moe_layers = []
904 + self.moe_mlp_layers = []
905 + example_moe = None
906 + for layer in self.model.layers:
907 + if isinstance(layer, PPMissingLayer):
908 + continue
909 +
910 + assert isinstance(layer, AXK1DecoderLayer)
911 + if isinstance(layer.mlp, AXK1MoE):
912 + # Pick last one layer since the first ones may be dense layers.
913 + example_moe = layer.mlp
914 + self.moe_mlp_layers.append(layer.mlp)
915 + self.moe_layers.append(layer.mlp.experts)
916 +
917 + self.extract_moe_parameters(example_moe)
918 +
919 + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
920 + return self.model.embed_input_ids(input_ids)
921 +
922 + def forward(
923 + self,
924 + input_ids: torch.Tensor | None,
925 + positions: torch.Tensor,
926 + intermediate_tensors: IntermediateTensors | None = None,
927 + inputs_embeds: torch.Tensor | None = None,
928 + ) -> torch.Tensor | IntermediateTensors:
929 + hidden_states = self.model(
930 + input_ids, positions, intermediate_tensors, inputs_embeds
931 + )
932 + return hidden_states
933 +
934 + def compute_logits(
935 + self,
936 + hidden_states: torch.Tensor,
937 + ) -> torch.Tensor | None:
938 + logits = self.logits_processor(self.lm_head, hidden_states)
939 + return logits
940 +
941 + def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
942 + # Params for weights, fp8 weight scales, fp8 activation scales
943 + # (param_name, weight_name, expert_id, shard_id)
944 + return SharedFusedMoE.make_expert_params_mapping(
945 + self,
946 + ckpt_gate_proj_name="gate_proj",
947 + ckpt_down_proj_name="down_proj",
948 + ckpt_up_proj_name="up_proj",
949 + num_experts=self.config.n_routed_experts,
950 + num_redundant_experts=0,
951 + )
952 +
953 + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
954 + rocm_aiter_moe_shared_expert_enabled = (
955 + rocm_aiter_ops.is_fusion_moe_shared_experts_enabled()
956 + )
957 + stacked_params_mapping = [
958 + # (param_name, shard_name, shard_id)
959 + ("gate_up_proj", "gate_proj", 0),
960 + ("gate_up_proj", "up_proj", 1),
961 + ]
962 + mla_params_mapping = [
963 + ("fused_qkv_a_proj", "q_a_proj", 0),
964 + ("fused_qkv_a_proj", "kv_a_proj_with_mqa", 1),
965 + ]
966 + mha_params_mapping = [
967 + ("qkv_proj", "q_proj", "q"),
968 + ("qkv_proj", "k_proj", "k"),
969 + ("qkv_proj", "v_proj", "v"),
970 + ]
971 + if self.use_mha:
972 + stacked_params_mapping.extend(mha_params_mapping)
973 + else:
974 + stacked_params_mapping.extend(mla_params_mapping)
975 +
976 + # Params for weights, fp8 weight scales, fp8 activation scales
977 + # (param_name, weight_name, expert_id, shard_id)
978 + expert_params_mapping = SharedFusedMoE.make_expert_params_mapping(
979 + self,
980 + ckpt_gate_proj_name="gate_proj",
981 + ckpt_down_proj_name="down_proj",
982 + ckpt_up_proj_name="up_proj",
983 + num_experts=self.config.n_routed_experts
984 + + (
985 + self.config.n_shared_experts
986 + if rocm_aiter_moe_shared_expert_enabled
987 + else 0
988 + ),
989 + num_redundant_experts=self.num_redundant_experts,
990 + )
991 +
992 + params_dict = dict(self.named_parameters())
993 + loaded_params: set[str] = set()
994 + for name, loaded_weight in weights:
995 + if "rotary_emb.inv_freq" in name:
996 + continue
997 +
998 + spec_layer = get_spec_layer_idx_from_weight_name(self.config, name)
999 + if spec_layer is not None:
1000 + continue # skip spec decode layers for main model
1001 +
1002 + is_fusion_moe_shared_experts_layer = (
1003 + rocm_aiter_moe_shared_expert_enabled and ("mlp.shared_experts" in name)
1004 + )
1005 +
1006 + for param_name, weight_name, shard_id in stacked_params_mapping:
1007 + # Skip non-stacked layers and experts (experts handled below).
1008 + if weight_name not in name:
1009 + continue
1010 + # We have mlp.experts[0].gate_proj in the checkpoint.
1011 + # Since we handle the experts below in expert_params_mapping,
1012 + # we need to skip here BEFORE we update the name, otherwise
1013 + # name will be updated to mlp.experts[0].gate_up_proj, which
1014 + # will then be updated below in expert_params_mapping
1015 + # for mlp.experts[0].gate_gate_up_proj, which breaks load.
1016 + if ("mlp.experts." in name) and name not in params_dict:
1017 + continue
1018 + if is_fusion_moe_shared_experts_layer:
1019 + continue
1020 + name_mapped = name.replace(weight_name, param_name)
1021 +
1022 + # QKV fusion is optional, fall back to normal
1023 + # weight loading if it's not enabled
1024 + # if go with fusion option, then update name
1025 + if (
1026 + param_name == "fused_qkv_a_proj"
1027 + ) and name_mapped not in params_dict:
1028 + continue
1029 + else:
1030 + name = name_mapped
1031 + # Skip loading extra bias for GPTQ models.
1032 + if name.endswith(".bias") and name not in params_dict:
1033 + continue
1034 +
1035 + if is_pp_missing_parameter(name, self):
1036 + continue
1037 +
1038 + param = params_dict[name]
1039 + weight_loader = param.weight_loader
1040 + weight_loader(param, loaded_weight, shard_id)
1041 + break
1042 + else:
1043 + is_expert_weight = False
1044 +
1045 + # Special handling: when AITER fusion_shared_experts is enabled,
1046 + # checkpoints may provide a single widened shared_experts tensor
1047 + # without explicit expert indices
1048 + # (e.g. ...mlp.shared_experts.gate_proj.weight).
1049 + # For models with multiple shared experts, split that tensor
1050 + # evenly into per-shared-expert slices and load them into
1051 + # appended expert slots mlp.experts.{n_routed_experts + j}.*
1052 + # accordingly.
1053 + num_chunks = 1
1054 + if is_fusion_moe_shared_experts_layer:
1055 + num_chunks = getattr(self.config, "n_shared_experts", 1) or 1
1056 + # Determine split axis based on op type
1057 + # gate/up: ColumnParallel → split along dim 0
1058 + # down: RowParallel → split along dim 1
1059 + split_dim = (
1060 + 1
1061 + if ("down_proj.weight" in name and loaded_weight.ndim > 1)
1062 + else 0
1063 + )
1064 + total = loaded_weight.shape[split_dim]
1065 + assert total % num_chunks == 0, (
1066 + f"Shared expert weight dim {total} "
1067 + f"not divisible by num_chunks {num_chunks}"
1068 + )
1069 + chunk_size = total // num_chunks
1070 +
1071 + for j in range(num_chunks):
1072 + chunk_name = name
1073 + weight_to_load = loaded_weight
1074 +
1075 + if is_fusion_moe_shared_experts_layer:
1076 + chunk_slice = slice(j * chunk_size, (j + 1) * chunk_size)
1077 + if loaded_weight.ndim == 1:
1078 + weight_to_load = loaded_weight[chunk_slice]
1079 + elif split_dim == 0:
1080 + weight_to_load = loaded_weight[chunk_slice, :]
1081 + else:
1082 + weight_to_load = loaded_weight[:, chunk_slice]
1083 + # Synthesize an expert-style name so expert mapping
1084 + # can route it
1085 + chunk_name = name.replace(
1086 + "mlp.shared_experts",
1087 + f"mlp.experts.{self.config.n_routed_experts + j}",
1088 + )
1089 +
1090 + # Use expert_params_mapping to locate the destination
1091 + # param and delegate to its expert-aware weight_loader
1092 + # with expert_id.
1093 + for mapping in expert_params_mapping:
1094 + param_name, weight_name, expert_id, shard_id = mapping
1095 + if weight_name not in chunk_name:
1096 + continue
1097 +
1098 + # Anyway, this is an expert weight and should not be
1099 + # attempted to load as other weights later
1100 + is_expert_weight = True
1101 +
1102 + # Do not modify `name` since the loop may continue here
1103 + # Instead, create a new variable
1104 + name_mapped = chunk_name.replace(weight_name, param_name)
1105 +
1106 + if is_pp_missing_parameter(name_mapped, self):
1107 + continue
1108 +
1109 + param = params_dict[name_mapped]
1110 + # We should ask the weight loader to return success or
1111 + # not here since otherwise we may skip experts with
1112 + # other available replicas.
1113 + weight_loader = typing.cast(
1114 + Callable[..., bool], param.weight_loader
1115 + )
1116 + success = weight_loader(
1117 + param,
1118 + weight_to_load,
1119 + name_mapped,
1120 + shard_id=shard_id,
1121 + expert_id=expert_id,
1122 + return_success=True,
1123 + )
1124 + if success:
1125 + if not is_fusion_moe_shared_experts_layer:
1126 + name = name_mapped
1127 + else:
1128 + loaded_params.add(name_mapped)
1129 + break
1130 + else:
1131 + if is_expert_weight:
1132 + # We've checked that this is an expert weight
1133 + # However it's not mapped locally to this rank
1134 + # So we simply skip it
1135 + continue
1136 +
1137 + # Skip loading extra bias for GPTQ models.
1138 + if name.endswith(".bias") and name not in params_dict:
1139 + continue
1140 +
1141 + # Remapping the name of FP8 kv-scale.
1142 + name = maybe_remap_kv_scale_name(name, params_dict)
1143 + if name is None:
1144 + continue
1145 +
1146 + if is_pp_missing_parameter(name, self):
1147 + continue
1148 +
1149 + param = params_dict[name]
1150 + weight_loader = getattr(
1151 + param, "weight_loader", default_weight_loader
1152 + )
1153 + weight_loader(param, loaded_weight)
1154 + if not is_fusion_moe_shared_experts_layer:
1155 + loaded_params.add(name)
1156 +
1157 + return loaded_params
1158 +
1159 +
1160 +def get_spec_layer_idx_from_weight_name(
1161 + config: AXK1Config, weight_name: str
1162 +) -> int | None:
1163 + if config.num_nextn_predict_layers and config.num_nextn_predict_layers > 0:
1164 + layer_idx = config.num_hidden_layers
1165 + for i in range(config.num_nextn_predict_layers):
1166 + if weight_name.startswith(f"model.layers.{layer_idx + i}."):
1167 + return layer_idx + i
1168 + return None
vllm/model_executor/models/registry.py
+1
@@ -75,6 +75,7 @@ _TEXT_GENERATION_MODELS = {
75 "AquilaForCausalLM": ("llama", "LlamaForCausalLM"), # AquilaChat2
76 "ArceeForCausalLM": ("arcee", "ArceeForCausalLM"),
77 "ArcticForCausalLM": ("arctic", "ArcticForCausalLM"),
78 + "AXK1ForCausalLM": ("AXK1", "AXK1ForCausalLM"),
79 # baichuan-7b, upper case 'C' in the class name
80 "BaiChuanForCausalLM": ("baichuan", "BaiChuanForCausalLM"),
81 # baichuan-13b, lower case 'c' in the class name
vllm/transformers_utils/configs/AXK1.py new
+215
@@ -0,0 +1,215 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +from typing import Any
4 +
5 +from transformers import PretrainedConfig
6 +
7 +
8 +class AXK1Config(PretrainedConfig):
9 + r"""
10 + This is the configuration class to store the configuration of a [`AXK1Model`].
11 + It is used to instantiate an A.X model according to the specified arguments,
12 + defining the model architecture. Instantiating a configuration with the defaults
13 + will yield a similar configuration to that of the A.X K1.
14 + Configuration objects inherit from [`PretrainedConfig`] and can be used to control
15 + the model outputs. Read the documentation from [`PretrainedConfig`] for more
16 + information.
17 + Args:
18 + vocab_size (`int`, *optional*, defaults to 163840):
19 + Vocabulary size of the A.X K1 model. Defines the number of different
20 + tokens that can be represented by the `inputs_ids` passed when calling
21 + [`AXK1Model`]
22 + hidden_size (`int`, *optional*, defaults to 7168):
23 + Dimension of the hidden representations.
24 + intermediate_size (`int`, *optional*, defaults to 18432):
25 + Dimension of the MLP representations.
26 + moe_intermediate_size (`int`, *optional*, defaults to 2048):
27 + Dimension of the MoE representations.
28 + num_hidden_layers (`int`, *optional*, defaults to 61):
29 + Number of hidden layers in the Transformer decoder.
30 + num_nextn_predict_layers (`int`, *optional*, defaults to 1):
31 + Number of nextn predict layers in the AXK1 Model.
32 + num_attention_heads (`int`, *optional*, defaults to 64):
33 + Number of attention heads for each attention layer in the Transformer
34 + decoder.
35 + n_shared_experts (`int`, *optional*, defaults to 1):
36 + Number of shared experts, None means dense model.
37 + n_routed_experts (`int`, *optional*, defaults to 192):
38 + Number of routed experts, None means dense model.
39 + routed_scaling_factor (`float`, *optional*, defaults to 2.5):
40 + Scaling factor or routed experts.
41 + topk_method (`str`, *optional*, defaults to `noaux_tc`):
42 + Topk method used in routed gate.
43 + n_group (`int`, *optional*, defaults to 8):
44 + Number of groups for routed experts.
45 + topk_group (`int`, *optional*, defaults to 4):
46 + Number of selected groups for each token(for each token, ensuring the
47 + selected experts is only within `topk_group` groups).
48 + num_experts_per_tok (`int`, *optional*, defaults to 8):
49 + Number of selected experts, None means dense model.
50 + moe_layer_freq (`int`, *optional*, defaults to 1):
51 + The frequency of the MoE layer: one expert layer for every
52 + `moe_layer_freq - 1` dense layers.
53 + first_k_dense_replace (`int`, *optional*, defaults to 1):
54 + Number of dense layers in shallow layers
55 + (embed->dense->dense->...->dense->moe->moe...->lm_head).
56 + \--k dense layers--/
57 + norm_topk_prob (`bool`, *optional*, defaults to True):
58 + Whether to normalize the weights of the routed experts.
59 + scoring_func (`str`, *optional*, defaults to 'sigmoid'):
60 + Method of computing expert weights.
61 + aux_loss_alpha (`float`, *optional*, defaults to 0.0001):
62 + Auxiliary loss weight coefficient.
63 + seq_aux = (`bool`, *optional*, defaults to True):
64 + Whether to compute the auxiliary loss for each individual sample.
65 + num_key_value_heads (`int`, *optional*):
66 + This is the number of key_value heads that should be used to implement
67 + Grouped Query Attention. If `num_key_value_heads=num_attention_heads`,
68 + the model will use Multi Head Attention (MHA), if `num_key_value_heads=1
69 + the model will use Multi Query Attention (MQA) otherwise GQA is used. When
70 + converting a multi-head checkpoint to a GQA checkpoint, each group key and
71 + value head should be constructed by meanpooling all the original heads
72 + within that group. For more details checkout
73 + [this paper](https://arxiv.org/pdf/2305.13245.pdf).
74 + If it is not specified, will default to `num_attention_heads`.
75 + hidden_act (`str` or `function`, *optional*, defaults to `"silu"`):
76 + The non-linear activation function (function or string) in the decoder.
77 + max_position_embeddings (`int`, *optional*, defaults to 131072):
78 + The maximum sequence length that this model might ever be used with.
79 + initializer_range (`float`, *optional*, defaults to 0.02):
80 + The standard deviation of the truncated_normal_initializer for
81 + initializing all weight matrices.
82 + rms_norm_eps (`float`, *optional*, defaults to 1e-06):
83 + The epsilon used by the rms normalization layers.
84 + use_cache (`bool`, *optional*, defaults to `True`):
85 + Whether or not the model should return the last key/values attentions
86 + (not used by all models). Only relevant if `config.is_decoder=True`.
87 + pad_token_id (`int`, *optional*):
88 + Padding token id.
89 + bos_token_id (`int`, *optional*, defaults to 163691):
90 + Beginning of stream token id.
91 + eos_token_id (`int`, *optional*, defaults to 163691):
92 + End of stream token id.
93 + pretraining_tp (`int`, *optional*, defaults to 1):
94 + Experimental feature. Tensor parallelism rank used during pretraining.
95 + Please refer to
96 + [this document](https://huggingface.co/docs/transformers/parallelism)
97 + to understand more about it. This value is necessary to ensure exact
98 + reproducibility of the pretraining results. Please refer to
99 + [this issue](https://github.com/pytorch/pytorch/issues/76232).
100 + tie_word_embeddings (`bool`, *optional*, defaults to `False`):
101 + Whether to tie weight embeddings
102 + rope_theta (`float`, *optional*, defaults to 10000.0):
103 + The base period of the RoPE embeddings.
104 + rope_scaling (`Dict`, *optional*):
105 + Dictionary containing the scaling configuration for the RoPE embeddings.
106 + Currently supports two scaling strategies: linear and dynamic.
107 + Their scaling factor must be a float greater than 1. The expected format
108 + is `{"type": strategy name, "factor": scaling factor}`. When using this
109 + flag, don't update `max_position_embeddings` to the expected new maximum.
110 + attention_bias (`bool`, defaults to `False`, *optional*, defaults to `False`):
111 + Whether to use a bias in the query, key, value and output projection
112 + layers during self-attention.
113 + attention_dropout (`float`, *optional*, defaults to 0.0):
114 + The dropout ratio for the attention probabilities.
115 + """
116 +
117 + model_type = "AXK1"
118 + keys_to_ignore_at_inference = ["past_key_values"]
119 +
120 + def __init__(
121 + self,
122 + vocab_size: int = 163840,
123 + hidden_size: int = 7168,
124 + intermediate_size: int = 18432,
125 + moe_intermediate_size: int = 2048,
126 + num_hidden_layers: int = 61,
127 + num_nextn_predict_layers: int | None = 1,
128 + num_attention_heads: int = 64,
129 + num_key_value_heads: int = 64,
130 + n_shared_experts: int | None = 1,
131 + n_routed_experts: int | None = 192,
132 + ep_size: int | None = 8, ## Ignored - Expert parallel size
133 + routed_scaling_factor: float | None = 2.5,
134 + kv_lora_rank: int | None = 512,
135 + q_lora_rank: int | None = 1536,
136 + qk_rope_head_dim: int | None = 64,
137 + v_head_dim: int | None = 128,
138 + qk_nope_head_dim: int | None = 128,
139 + topk_method: str | None = "noaux_tc",
140 + n_group: int | None = 8,
141 + topk_group: int | None = 4,
142 + num_experts_per_tok: int | None = 8,
143 + moe_layer_freq: int | None = 1,
144 + first_k_dense_replace: int = 1,
145 + norm_topk_prob: bool = True,
146 + scoring_func: str | None = "sigmoid",
147 + aux_loss_alpha: float | None = 0.0001,
148 + seq_aux: float | None = True,
149 + hidden_act: str | None = "silu",
150 + max_position_embeddings: int | None = 131072,
151 + initializer_range: float | None = 0.02,
152 + rms_norm_eps: float = 1e-6,
153 + use_cache: bool | None = True,
154 + pad_token_id: int | None = None,
155 + bos_token_id: int | None = 163691,
156 + eos_token_id: int | None = 163691,
157 + pretraining_tp: int | None = 1,
158 + tie_word_embeddings: bool | None = False,
159 + rope_theta: float | None = 10000.0,
160 + rope_scaling: dict[str, Any] | None = None,
161 + rope_parameters: dict[str, Any] | None = None,
162 + attention_bias: bool | None = False,
163 + attention_dropout: float | None = 0.0,
164 + **kwargs,
165 + ):
166 + self.vocab_size = vocab_size
167 + self.max_position_embeddings = max_position_embeddings
168 + self.hidden_size = hidden_size
169 + self.intermediate_size = intermediate_size
170 + self.moe_intermediate_size = moe_intermediate_size
171 + self.num_hidden_layers = num_hidden_layers
172 + self.num_nextn_predict_layers = num_nextn_predict_layers
173 + self.num_attention_heads = num_attention_heads
174 + self.n_shared_experts = n_shared_experts
175 + self.n_routed_experts = n_routed_experts
176 + self.ep_size = ep_size
177 + self.routed_scaling_factor = routed_scaling_factor
178 + self.kv_lora_rank = kv_lora_rank
179 + self.q_lora_rank = q_lora_rank
180 + self.qk_rope_head_dim = qk_rope_head_dim
181 + self.v_head_dim = v_head_dim
182 + self.qk_nope_head_dim = qk_nope_head_dim
183 + self.topk_method = topk_method
184 + self.n_group = n_group
185 + self.topk_group = topk_group
186 + self.num_experts_per_tok = num_experts_per_tok
187 + self.moe_layer_freq = moe_layer_freq
188 + self.first_k_dense_replace = first_k_dense_replace
189 + self.norm_topk_prob = norm_topk_prob
190 + self.scoring_func = scoring_func
191 + self.aux_loss_alpha = aux_loss_alpha
192 + self.seq_aux = seq_aux
193 + # for backward compatibility
194 + if num_key_value_heads is None:
195 + num_key_value_heads = num_attention_heads
196 +
197 + self.num_key_value_heads = num_key_value_heads
198 + self.hidden_act = hidden_act
199 + self.initializer_range = initializer_range
200 + self.rms_norm_eps = rms_norm_eps
201 + self.pretraining_tp = pretraining_tp
202 + self.use_cache = use_cache
203 + self.rope_theta = rope_theta
204 + self.rope_scaling = rope_scaling
205 + self.rope_parameters = rope_parameters
206 + self.attention_bias = attention_bias
207 + self.attention_dropout = attention_dropout
208 +
209 + super().__init__(
210 + pad_token_id=pad_token_id,
211 + bos_token_id=bos_token_id,
212 + eos_token_id=eos_token_id,
213 + tie_word_embeddings=tie_word_embeddings,
214 + **kwargs,
215 + )
vllm/transformers_utils/configs/__init__.py
+2
@@ -16,6 +16,7 @@ import importlib
16
17 _CLASS_TO_MODULE: dict[str, str] = {
18 "AfmoeConfig": "vllm.transformers_utils.configs.afmoe",
19 + "AXK1Config": "vllm.transformers_utils.configs.AXK1",
20 "BagelConfig": "vllm.transformers_utils.configs.bagel",
21 "ChatGLMConfig": "vllm.transformers_utils.configs.chatglm",
22 "ColModernVBertConfig": "vllm.transformers_utils.configs.colmodernvbert",
@@ -70,6 +71,7 @@ _CLASS_TO_MODULE: dict[str, str] = {
71
72 __all__ = [
73 "AfmoeConfig",
74 + "AXK1Config",
75 "BagelConfig",
76 "ChatGLMConfig",
77 "ColModernVBertConfig",
vllm/transformers_utils/model_arch_config_convertor.py
+8 -1
@@ -233,6 +233,7 @@ class ModelArchConfigConvertorBase:
233 if not hasattr(self.hf_text_config, "model_type"):
234 return False
235 elif self.hf_text_config.model_type in (
236 + "AXK1",
237 "deepseek_v2",
238 "deepseek_v3",
239 "deepseek_v32",
@@ -253,7 +254,13 @@ class ModelArchConfigConvertorBase:
254 # underlying architecture
255 return (
256 self.hf_text_config.model.model_type
256 - in ("deepseek_v2", "deepseek_v3", "deepseek_v32", "deepseek_mtp")
257 + in (
258 + "AXK1",
259 + "deepseek_v2",
260 + "deepseek_v3",
261 + "deepseek_v32",
262 + "deepseek_mtp",
263 + )
264 and self.hf_text_config.kv_lora_rank is not None
265 )
266 return False