@tduncan / Nvidia-Optimization-Stack-VLLM / commits / ba1cdcfcf0

`[Model][Quantization] Add Ling-3.0-flash-fp8 support` (#51265)

Signed-off-by: Isotr0py <Isotr0py@outlook.com> Co-authored-by: Isotr0py <Isotr0py@outlook.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>

zexplorerhj committed Aug 10, 2026 at 14:09 UTC ba1cdcfcf05f581d1f58cc911644fed62317c8c6

6 files changed +298 -36

tests/quantization/test_quark.py

+23 -2

@@ -701,11 +701,32 @@ def test_non_fused_layer_unaffected():

701

702

703 def test_substr_match_on_fused_name():

704 - # skip_with_substr=True path: fused-name substring match should also

704 + # Substring matching: a fused-name match should also

705 # short-circuit before shard expansion.

706 assert is_layer_skipped(

707 prefix="model.layers.0.self_attn.qkv_proj",

708 ignored_layers=["self_attn.qkv_proj"],

709 fused_mapping=FUSED_MAPPING,

710 - skip_with_substr=True,

710 + match_mode="substring",

711 + )

712 +

713 +

714 +@pytest.mark.parametrize(

715 + ("prefix", "ignored_layer", "expected"),

716 + [

717 + ("model.layers.0.self_attn.b_proj", "b_proj", True),

718 + ("model.layers.0.self_attn.q_b_proj", "b_proj", False),

719 + ("model.layers.0.self_attn.kv_b_proj", "b_proj", False),

720 + ("model.layers.5.self_attn.g_proj", "5.self_attn.g_proj", True),

721 + ("model.layers.6.self_attn.g_proj", "5.self_attn.g_proj", False),

722 + ],

723 +)

724 +def test_suffix_match_at_module_boundary(prefix, ignored_layer, expected):

725 + assert (

726 + is_layer_skipped(

727 + prefix=prefix,

728 + ignored_layers=[ignored_layer],

729 + match_mode="suffix",

730 + )

731 + is expected

732 )

vllm/model_executor/layers/quantization/auto_awq.py

+2 -2

@@ -292,7 +292,7 @@ class AutoAWQConfig(QuantizationConfig):

292 prefix,

293 self.modules_to_not_convert,

294 self.packed_modules_mapping,

295 - skip_with_substr=True,

295 + match_mode="substring",

296 ):

297 return UnquantizedLinearMethod()

298

@@ -335,7 +335,7 @@ class AutoAWQConfig(QuantizationConfig):

335 if is_layer_skipped(

336 prefix,

337 getattr(self, "modules_to_not_convert", []),

338 - skip_with_substr=True,

338 + match_mode="substring",

339 ):

340 return UnquantizedFusedMoEMethod(layer.moe_config)

341

vllm/model_executor/layers/quantization/fp8.py

+6 -1

@@ -1,7 +1,7 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

4 -from typing import TYPE_CHECKING, Any

4 +from typing import TYPE_CHECKING, Any, Literal

6 import torch

@@ -108,6 +108,9 @@ class Fp8Config(QuantizationConfig):

108 raise ValueError(f"Unsupported activation scheme {activation_scheme}")

109 self.activation_scheme = activation_scheme

110 self.ignored_layers = ignored_layers or []

111 + self.ignored_layers_match_mode: Literal["exact", "substring", "suffix"] = (

112 + "exact"

113 + )

114 self.store_dtype = store_dtype

115 if weight_block_size is not None:

116 if not is_checkpoint_fp8_serialized:

@@ -177,6 +180,7 @@ class Fp8Config(QuantizationConfig):

180 prefix=prefix,

181 ignored_layers=self.ignored_layers,

182 fused_mapping=self.packed_modules_mapping,

183 + match_mode=self.ignored_layers_match_mode,

184 ):

185 return UnquantizedLinearMethod()

186 if not self.is_checkpoint_fp8_serialized:

@@ -196,6 +200,7 @@ class Fp8Config(QuantizationConfig):

200 prefix=prefix,

201 ignored_layers=self.ignored_layers,

202 fused_mapping=self.packed_modules_mapping,

203 + match_mode=self.ignored_layers_match_mode,

204 ):

205 return UnquantizedFusedMoEMethod(layer.moe_config)

206 if self.store_dtype == "mxfp4":

vllm/model_executor/layers/quantization/utils/quant_utils.py

+20 -8

@@ -5,7 +5,7 @@

5 from collections.abc import Callable, Mapping

6 from dataclasses import dataclass

7 from types import MappingProxyType

8 -from typing import TYPE_CHECKING, ClassVar, NamedTuple

8 +from typing import TYPE_CHECKING, ClassVar, Literal, NamedTuple

10 import numpy

11 import torch

@@ -604,7 +604,7 @@ def is_layer_skipped(

604 ignored_layers: list[str],

605 fused_mapping: Mapping[str, list[str]] = MappingProxyType({}),

606 *,

607 - skip_with_substr: bool = False,

607 + match_mode: Literal["exact", "substring", "suffix"] = "exact",

608 ) -> bool:

609 def prefix_full_match(prefix: str, ignored_layers: list[str]) -> bool:

610 return prefix in ignored_layers

@@ -613,7 +613,22 @@ def is_layer_skipped(

613 def substr_match(prefix: str, ignored_layers: list[str]) -> bool:

614 return any(layer in prefix for layer in ignored_layers)

615

616 - match_func = substr_match if skip_with_substr else prefix_full_match

616 + # Match abbreviated module paths at a component boundary. For example,

617 + # ``b_proj`` matches ``model.layers.0.self_attn.b_proj`` but not

618 + # ``model.layers.0.self_attn.q_b_proj``.

619 + def suffix_match(prefix: str, ignored_layers: list[str]) -> bool:

620 + return any(

621 + prefix == layer or prefix.endswith(f".{layer}") for layer in ignored_layers

622 + )

623 +

624 + if match_mode == "exact":

625 + match_func = prefix_full_match

626 + elif match_mode == "substring":

627 + match_func = substr_match

628 + elif match_mode == "suffix":

629 + match_func = suffix_match

630 + else:

631 + raise ValueError(f"Unsupported layer skip match mode: {match_mode}")

632

633 # prefix: model.layers.0.self_attn.q_proj

634 # proj_name: q_proj

@@ -649,14 +664,11 @@ def is_layer_skipped(

664 "are quantized. All shards of fused layers "

665 "to have the same precision."

666 )

652 - elif "experts" in prefix and not skip_with_substr:

667 + elif "experts" in prefix and match_mode == "exact":

668 expert_ignore_layers = filter(

669 lambda layer_name: "experts" in layer_name, ignored_layers

670 )

656 - return any(

657 - prefix in layer_name if not skip_with_substr else layer_name in prefix

658 - for layer_name in expert_ignore_layers

659 - )

671 + return any(prefix in layer_name for layer_name in expert_ignore_layers)

672 else:

673 is_skipped = match_func(prefix, ignored_layers)

674

vllm/model_executor/models/bailing_moe_v3.py

+224 -21

@@ -10,6 +10,8 @@ vLLM's parallel linear layers, MLA kernel, KDA kernel and fused MoE loader.

11 import copy

12 from collections.abc import Callable, Iterable

13 +from math import lcm

14 +from typing import TypeGuard

16 import torch

17 import torch.nn as nn

@@ -55,6 +57,8 @@ from vllm.model_executor.layers.mla import (

57 MultiHeadLatentAttentionWrapper,

58 )

59 from vllm.model_executor.layers.quantization.base_config import QuantizationConfig

60 +from vllm.model_executor.layers.quantization.fp8 import Fp8Config

61 +from vllm.model_executor.layers.quantization.utils.quant_utils import is_layer_skipped

62 from vllm.model_executor.layers.rotary_embedding import get_rope

63 from vllm.model_executor.layers.vocab_parallel_embedding import (

64 ParallelLMHead,

@@ -79,7 +83,13 @@ from vllm.v1.attention.backends.gdn_attn import GDNAttentionMetadata

83 from vllm.v1.attention.backends.registry import MambaAttentionBackendEnum

85 from .interfaces import HasInnerState, IsHybrid, SupportsPP

82 -from .utils import PPMissingLayer, is_pp_missing_parameter, make_layers, maybe_prefix

86 +from .utils import (

87 + PPMissingLayer,

88 + WeightsMapper,

89 + is_pp_missing_parameter,

90 + make_layers,

91 + maybe_prefix,

92 +)

95 def bailing_v3_kda_attention(

@@ -193,6 +203,155 @@ def _load_a_log(param: torch.nn.Parameter, loaded_weight: torch.Tensor) -> None:

203 sharded_weight_loader(2)(param, loaded_weight)

204

205

206 +def _is_block_fp8_config(

207 + quant_config: QuantizationConfig | None,

208 +) -> TypeGuard[Fp8Config]:

209 + return (

210 + isinstance(quant_config, Fp8Config)

211 + and quant_config.weight_block_size is not None

212 + )

213 +

214 +

215 +def _configure_ling_fp8_quant_config(

216 + quant_config: QuantizationConfig | None,

217 +) -> None:

218 + if _is_block_fp8_config(quant_config):

219 + quant_config.ignored_layers_match_mode = "suffix"

220 +

221 +

222 +def _is_fp8_module_excluded(

223 + quant_config: QuantizationConfig | None,

224 + prefix: str,

225 +) -> bool:

226 + """Match Ling's abbreviated FP8 exclusions against mapped vLLM prefixes."""

227 + if not _is_block_fp8_config(quant_config):

228 + return False

229 +

230 + return is_layer_skipped(

231 + prefix=prefix,

232 + ignored_layers=quant_config.ignored_layers,

233 + fused_mapping=quant_config.packed_modules_mapping,

234 + match_mode=quant_config.ignored_layers_match_mode,

235 + )

236 +

237 +

238 +def _get_block_fp8_mlp_padded_intermediate_size(

239 + quant_config: QuantizationConfig | None,

240 + intermediate_size: int,

241 + prefix: str,

242 +) -> int:

243 + """Pad a block-FP8 MLP so each TP shard contains whole quant blocks."""

244 + if not _is_block_fp8_config(quant_config):

245 + return intermediate_size

246 + block_size = quant_config.weight_block_size

247 + assert block_size is not None

248 +

249 + alignments: list[int] = []

250 + if not _is_fp8_module_excluded(quant_config, f"{prefix}.gate_up_proj"):

251 + alignments.append(int(block_size[0]))

252 + if not _is_fp8_module_excluded(quant_config, f"{prefix}.down_proj"):

253 + alignments.append(int(block_size[1]))

254 + if not alignments:

255 + return intermediate_size

256 +

257 + tp_size = get_tensor_model_parallel_world_size()

258 + tp_alignment = tp_size * lcm(*alignments)

259 + return (intermediate_size + tp_alignment - 1) // tp_alignment * tp_alignment

260 +

261 +

262 +def _pad_block_fp8_mlp_checkpoint_tensor(

263 + quant_config: QuantizationConfig,

264 + name: str,

265 + loaded_weight: torch.Tensor,

266 + intermediate_size: int,

267 + padded_intermediate_size: int,

268 +) -> torch.Tensor:

269 + """Zero-pad an MLP checkpoint tensor on its intermediate dimension."""

270 + if padded_intermediate_size == intermediate_size:

271 + return loaded_weight

272 +

273 + block_size = getattr(quant_config, "weight_block_size", None)

274 + assert block_size is not None

275 +

276 + if ".down_proj." in name:

277 + if name.endswith(".bias"):

278 + return loaded_weight

279 + dim = 1

280 + block = int(block_size[1])

281 + logical_shards = 1

282 + elif any(

283 + projection in name

284 + for projection in (".gate_proj.", ".up_proj.", ".gate_up_proj.")

285 + ):

286 + dim = 0

287 + block = int(block_size[0])

288 + logical_shards = 2 if ".gate_up_proj." in name else 1

289 + else:

290 + return loaded_weight

291 +

292 + if name.endswith(".weight_scale_inv"):

293 + expected_shard_size = (intermediate_size + block - 1) // block

294 + target_shard_size = (padded_intermediate_size + block - 1) // block

295 + elif name.endswith((".weight", ".bias")):

296 + expected_shard_size = intermediate_size

297 + target_shard_size = padded_intermediate_size

298 + else:

299 + return loaded_weight

300 +

301 + current_size = loaded_weight.shape[dim]

302 + if current_size % logical_shards != 0:

303 + raise ValueError(

304 + f"Cannot split {name} dimension {current_size} into "

305 + f"{logical_shards} logical shards."

306 + )

307 + current_shard_size = current_size // logical_shards

308 + if current_shard_size == target_shard_size:

309 + return loaded_weight

310 + if current_shard_size != expected_shard_size:

311 + raise ValueError(

312 + f"Cannot pad {name}: expected each logical intermediate dimension "

313 + f"to be {expected_shard_size}, but got {current_shard_size}."

314 + )

315 +

316 + padded_shards: list[torch.Tensor] = []

317 + for shard in loaded_weight.split(current_shard_size, dim=dim):

318 + pad_shape = list(shard.shape)

319 + pad_shape[dim] = target_shard_size - current_shard_size

320 + padded_shards.extend([shard, shard.new_zeros(pad_shape)])

321 + return torch.cat(padded_shards, dim=dim)

322 +

323 +

324 +def _maybe_pad_block_fp8_shared_expert_checkpoint_tensor(

325 + quant_config: QuantizationConfig | None,

326 + config: PretrainedConfig,

327 + name: str,

328 + loaded_weight: torch.Tensor,

329 +) -> torch.Tensor:

330 + if ".mlp.shared_experts." not in name:

331 + return loaded_weight

332 +

333 + shared_prefix = name.split(".shared_experts.", 1)[0] + ".shared_experts"

334 + shared_intermediate = (

335 + config.moe_shared_expert_intermediate_size * config.num_shared_experts

336 + )

337 + padded_shared_intermediate = _get_block_fp8_mlp_padded_intermediate_size(

338 + quant_config,

339 + shared_intermediate,

340 + shared_prefix,

341 + )

342 + if padded_shared_intermediate == shared_intermediate:

343 + return loaded_weight

344 +

345 + assert quant_config is not None

346 + return _pad_block_fp8_mlp_checkpoint_tensor(

347 + quant_config,

348 + name,

349 + loaded_weight,

350 + shared_intermediate,

351 + padded_shared_intermediate,

352 + )

353 +

354 +

355 class _IdentityProjection(nn.Module):

356 def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, None]:

357 return x, None

@@ -446,13 +605,36 @@ class BailingMoeV3KimiDeltaAttention(PluggableLayer, MambaBase):

605

606 projection_size = self.head_dim * self.num_heads

607 self.projection_size_per_partition = projection_size // self.tp_size

449 - self.qkvb_proj = MergedColumnParallelLinear(

450 - self.hidden_size,

451 - [projection_size, projection_size, projection_size, self.num_heads],

452 - bias=False,

453 - quant_config=quant_config,

454 - prefix=f"{prefix}.qkvb_proj",

455 - )

608 + self.separate_b_proj = _is_fp8_module_excluded(quant_config, f"{prefix}.b_proj")

609 + self.qkv_proj: MergedColumnParallelLinear | None

610 + self.b_proj: ColumnParallelLinear | None

611 + self.qkvb_proj: MergedColumnParallelLinear | None

612 + if self.separate_b_proj:

613 + self.qkv_proj = MergedColumnParallelLinear(

614 + self.hidden_size,

615 + [projection_size, projection_size, projection_size],

616 + bias=False,

617 + quant_config=quant_config,

618 + prefix=f"{prefix}.qkv_proj",

619 + )

620 + self.b_proj = ColumnParallelLinear(

621 + self.hidden_size,

622 + self.num_heads,

623 + bias=False,

624 + quant_config=quant_config,

625 + prefix=f"{prefix}.b_proj",

626 + )

627 + self.qkvb_proj = None

628 + else:

629 + self.qkvb_proj = MergedColumnParallelLinear(

630 + self.hidden_size,

631 + [projection_size, projection_size, projection_size, self.num_heads],

632 + bias=False,

633 + quant_config=quant_config,

634 + prefix=f"{prefix}.qkvb_proj",

635 + )

636 + self.qkv_proj = None

637 + self.b_proj = None

638 self.f_proj = ColumnParallelLinear(

639 self.hidden_size,

640 projection_size,

@@ -532,16 +714,23 @@ class BailingMoeV3KimiDeltaAttention(PluggableLayer, MambaBase):

714 ) -> None:

715 del positions

716 num_tokens = hidden_states.size(0)

535 - qkvb = self.qkvb_proj(hidden_states)[0]

536 - q, k, v, beta_logits = qkvb.split(

537 - [

538 - self.projection_size_per_partition,

539 - self.projection_size_per_partition,

540 - self.projection_size_per_partition,

541 - self.local_num_heads,

542 - ],

543 - dim=-1,

544 - )

717 + if self.separate_b_proj:

718 + assert self.qkv_proj is not None and self.b_proj is not None

719 + qkv = self.qkv_proj(hidden_states)[0]

720 + q, k, v = qkv.split(self.projection_size_per_partition, dim=-1)

721 + beta_logits = self.b_proj(hidden_states)[0]

722 + else:

723 + assert self.qkvb_proj is not None

724 + qkvb = self.qkvb_proj(hidden_states)[0]

725 + q, k, v, beta_logits = qkvb.split(

726 + [

727 + self.projection_size_per_partition,

728 + self.projection_size_per_partition,

729 + self.projection_size_per_partition,

730 + self.local_num_heads,

731 + ],

732 + dim=-1,

733 + )

734 beta = beta_logits.float().sigmoid().unsqueeze(0)

735 g1 = self.f_proj(hidden_states)[0]

736 g1 = fused_kda_gate(

@@ -910,7 +1099,6 @@ class BailingMoeV3MoE(nn.Module):

1099 prefix: str = "",

1100 ) -> None:

1101 super().__init__()

913 - self.tp_size = get_tensor_model_parallel_world_size()

1102 self.num_experts = config.num_experts

1103 self.top_k = config.num_experts_per_tok

1104 self.hidden_size = config.hidden_size

@@ -926,7 +1114,9 @@ class BailingMoeV3MoE(nn.Module):

1114 config.moe_shared_expert_intermediate_size * config.num_shared_experts

1115 )

1116 self.shared_experts = BailingMoeV3MLP(

929 - intermediate_size=shared_intermediate,

1117 + intermediate_size=_get_block_fp8_mlp_padded_intermediate_size(

1118 + quant_config, shared_intermediate, f"{prefix}.shared_experts"

1119 + ),

1120 config=config,

1121 quant_config=quant_config,

1122 reduce_results=False,

@@ -1149,8 +1339,11 @@ class BailingMoeV3Model(nn.Module):

1339

1340

1341 class BailingMoeV3ForCausalLM(nn.Module, HasInnerState, IsHybrid, SupportsPP):

1342 + hf_to_vllm_mapper = WeightsMapper(orig_to_new_substr={".attention.": ".self_attn."})

1343 +

1344 packed_modules_mapping = {

1345 "gate_up_proj": ["gate_proj", "up_proj"],

1346 + "qkv_proj": ["q_proj", "k_proj", "v_proj"],

1347 "qkvb_proj": ["q_proj", "k_proj", "v_proj", "b_proj"],

1348 }

1349

@@ -1158,6 +1351,7 @@ class BailingMoeV3ForCausalLM(nn.Module, HasInnerState, IsHybrid, SupportsPP):

1351 super().__init__()

1352 config = vllm_config.model_config.hf_config

1353 quant_config = vllm_config.quant_config

1354 + _configure_ling_fp8_quant_config(quant_config)

1355 self.config = config

1356 self.quant_config = quant_config

1357 self.model = BailingMoeV3Model(

@@ -1226,11 +1420,15 @@ class BailingMoeV3ForCausalLM(nn.Module, HasInnerState, IsHybrid, SupportsPP):

1420 return self.model.get_expert_mapping()

1421

1422 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:

1423 + weights = self.hf_to_vllm_mapper.apply(weights)

1424 params_dict = dict(self.named_parameters(remove_duplicate=False))

1425 loaded_params: set[str] = set()

1426 stacked_mappings = [

1427 (".gate_up_proj", ".gate_proj", 0),

1428 (".gate_up_proj", ".up_proj", 1),

1429 + (".qkv_proj", ".q_proj", 0),

1430 + (".qkv_proj", ".k_proj", 1),

1431 + (".qkv_proj", ".v_proj", 2),

1432 (".qkvb_proj", ".q_proj", 0),

1433 (".qkvb_proj", ".k_proj", 1),

1434 (".qkvb_proj", ".v_proj", 2),

@@ -1263,13 +1461,18 @@ class BailingMoeV3ForCausalLM(nn.Module, HasInnerState, IsHybrid, SupportsPP):

1461 layer_idx = int(name.split("model.layers.")[1].split(".")[0])

1462 if layer_idx >= self.config.num_hidden_layers:

1463 return None

1266 - name = name.replace("attention.", "self_attn.")

1464 return maybe_remap_kv_scale_name(name, params_dict)

1465

1466 for orig_name, weight in weights:

1467 name = normalize_name(orig_name)

1468 if name is None:

1469 continue

1470 + weight = _maybe_pad_block_fp8_shared_expert_checkpoint_tensor(

1471 + self.quant_config,

1472 + self.config,

1473 + name,

1474 + weight,

1475 + )

1476 loaded = False

1477 for param_suf, weight_suf, stacked_shard_id in stacked_mappings:

1478 if weight_suf in name:

vllm/model_executor/models/bailing_moe_v3_mtp.py

+23 -2

@@ -15,6 +15,7 @@ from vllm.model_executor.layers.fused_moe import (

15 fused_moe_make_expert_params_mapping,

16 )

17 from vllm.model_executor.layers.layernorm import RMSNorm

18 +from vllm.model_executor.layers.linear import ReplicatedLinear

19 from vllm.model_executor.layers.logits_processor import LogitsProcessor

20 from vllm.model_executor.layers.vocab_parallel_embedding import (

21 ParallelLMHead,

@@ -25,8 +26,11 @@ from vllm.model_executor.model_loader.weight_utils import (

26 maybe_remap_kv_scale_name,

27 )

28 from vllm.model_executor.models.bailing_moe_v3 import (

29 + BailingMoeV3ForCausalLM,

30 BailingMoeV3MLAAttention,

31 BailingMoeV3MoE,

32 + _configure_ling_fp8_quant_config,

33 + _maybe_pad_block_fp8_shared_expert_checkpoint_tensor,

34 )

35 from vllm.model_executor.models.interfaces import SupportsPP

36 from vllm.model_executor.models.utils import (

@@ -79,7 +83,14 @@ class BailingMoeV3MultiTokenPredictorLayer(nn.Module):

83 self.layer_id = layer_id

84 self.enorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)

85 self.hnorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)

82 - self.eh_proj = nn.Linear(config.hidden_size * 2, config.hidden_size, bias=False)

86 + self.eh_proj = ReplicatedLinear(

87 + config.hidden_size * 2,

88 + config.hidden_size,

89 + bias=False,

90 + quant_config=vllm_config.quant_config,

91 + prefix=maybe_prefix(prefix, "eh_proj"),

92 + return_bias=False,

93 + )

94 self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)

95 self.self_attn = BailingMoeV3MLAAttention(

96 config,

@@ -204,6 +215,8 @@ class BailingMoeV3MultiTokenPredictor(nn.Module):

215

216 @support_torch_compile

217 class BailingMoeV3MTPModel(nn.Module, SupportsPP):

218 + hf_to_vllm_mapper = BailingMoeV3ForCausalLM.hf_to_vllm_mapper

219 +

220 packed_modules_mapping = {

221 "gate_up_proj": ["gate_proj", "up_proj"],

222 "fused_qkv_a_proj": ["q_a_proj", "kv_a_proj_with_mqa"],

@@ -217,6 +230,8 @@ class BailingMoeV3MTPModel(nn.Module, SupportsPP):

230 ) -> None:

231 super().__init__()

232 self.config = _get_draft_hf_config(vllm_config)

233 + self.quant_config = vllm_config.quant_config

234 + _configure_ling_fp8_quant_config(self.quant_config)

235 self.model = BailingMoeV3MultiTokenPredictor(

236 vllm_config=vllm_config,

237 prefix=maybe_prefix(prefix, "model"),

@@ -265,6 +280,7 @@ class BailingMoeV3MTPModel(nn.Module, SupportsPP):

280 )

281

282 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:

283 + weights = self.hf_to_vllm_mapper.apply(weights)

284 stacked_params_mapping = [

285 (".fused_qkv_a_proj", ".q_a_proj", 0),

286 (".fused_qkv_a_proj", ".kv_a_proj_with_mqa", 1),

@@ -320,7 +336,6 @@ class BailingMoeV3MTPModel(nn.Module, SupportsPP):

336 return None

337

338 def normalize_name(name: str) -> str:

323 - name = name.replace(".attention.", ".self_attn.")

339 return name.replace(

340 "mlp.gate.e_score_correction_bias",

341 "mlp.gate.expert_bias",

@@ -349,6 +364,12 @@ class BailingMoeV3MTPModel(nn.Module, SupportsPP):

364 if spec_layer is None:

365 continue

366 name = normalize_name(name)

367 + loaded_weight = _maybe_pad_block_fp8_shared_expert_checkpoint_tensor(

368 + self.quant_config,

369 + self.config,

370 + name,

371 + loaded_weight,

372 + )

373

374 loaded = False

375 for param_name, weight_name, stacked_shard_id in stacked_params_mapping: