# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from collections.abc import Iterable from typing import ClassVar, Literal import os import torch from torch import nn from vllm.config import ParallelConfig, VllmConfig from vllm.distributed import ( get_ep_group, get_pp_group, get_tensor_model_parallel_rank, get_tensor_model_parallel_world_size, tensor_model_parallel_all_gather, ) from vllm.logger import init_logger from vllm.model_executor.layers.activation import SiluAndMul, SiluAndMulWithClamp from vllm.model_executor.layers.fused_moe import ( FusedMoEFactory, GateLinear, fused_moe_make_expert_params_mapping, ) from vllm.model_executor.layers.layernorm import RMSNorm from vllm.model_executor.layers.linear import ( MergedColumnParallelLinear, RowParallelLinear, ) from vllm.model_executor.layers.logits_processor import LogitsProcessor from vllm.model_executor.layers.mamba.mamba_utils import ( MambaStateCopyFunc, MambaStateCopyFuncCalculator, MambaStateDtypeCalculator, MambaStateShapeCalculator, ) from vllm.model_executor.layers.mhc import ( MHCFusedPostPreOp, MHCPostOp, MHCPreOp, hc_contract, hc_expand, ) from vllm.model_executor.layers.quantization import QuantizationConfig from vllm.model_executor.layers.quantization.utils.quant_utils import ( GroupShape, scaled_dequantize, ) from vllm.model_executor.layers.vocab_parallel_embedding import ( ParallelLMHead, VocabParallelEmbedding, ) from vllm.model_executor.model_loader.weight_utils import ( default_weight_loader, maybe_remap_kv_scale_name, ) from vllm.model_executor.models.deepseek_v2 import _get_moe_router_dtype from vllm.model_executor.models.glm4_1v import ( Glm4vDummyInputsBuilder, Glm4vForConditionalGeneration, ) from vllm.model_executor.models.interfaces import ( EagleModelMixin, HasInnerState, IsHybrid, MixtureOfExperts, SupportsEagle3, SupportsPP, ) from vllm.model_executor.models.utils import ( AutoWeightsLoader, PPMissingLayer, init_vllm_registered_model, is_pp_missing_parameter, make_layers, maybe_prefix, sequence_parallel_chunk, ) from vllm.models.common.ops.sequence_parallel import ( sp_all_gather, sp_reduce_scatter, sp_shard, ) # GB10: sequence parallelism over plain TP for prefill-sized batches # (VLLM_GLM_SP_TP=1). mHC, norms and residuals run on each rank's quarter of # the tokens; attention, the MoE and the dense MLPs all-gather their input and # reduce-scatter their partial output, the bytes of the all-reduce they # replace. Decode-sized batches keep plain TP: four small collectives per # layer cost more than the sharded work saves there (-10% decode measured). # GLM runs eagerly around CUDA graphs, so the choice is made per forward from # the batch size; graphs are captured below VLLM_GLM_SP_MIN_TOKENS, plain TP. _SP_TP = os.environ.get("VLLM_GLM_SP_TP") == "1" _SP_MIN_TOKENS = int(os.environ.get("VLLM_GLM_SP_MIN_TOKENS", "1024")) # Gather KDA attention inputs as per-token FP8 (half the bytes). The only # consumer is in_proj's FP8 GEMM, which would quantize the same rows the same way. _SP_FP8_GATHER = os.environ.get("VLLM_GLM_SP_FP8_GATHER") == "1" # Decode: each arx all-reduce asks L2 for the weights the next kernels read # (the MoE's gate and shared expert after attention, the next KDA in_proj # after the MLP) while it waits for the peers. _ARX_PREFETCH = os.environ.get("VLLM_GLM_ARX_PREFETCH") == "1" # SP prefill MoE layers: router, shared expert and moe_prefill's fc1/fc2 run # here, and one arxbig kernel writes shared + scaled routed sum straight into # the RDMA reduce-scatter while the network already sends finished rows. The # runner's separate scale and add passes and the NCCL reduce-scatter go away. _SP_MOE_FUSED = os.environ.get("VLLM_GLM_SP_MOE_FUSED") == "1" # Route each rank's own rows, then gather the MoE input already quantized: # NVFP4 for the routed experts and per-token FP8 for the shared expert, the # same values each would compute from the gathered bf16 rows. _SP_MOE_QUANT_GATHER = os.environ.get("VLLM_GLM_SP_MOE_QUANT_GATHER") == "1" _sp_moe_check = int(os.environ.get("VLLM_GLM_SP_MOE_CHECK", "0")) _sp_active = False # KDA spec verify from one checkpoint per request (recoverssm.py). The model # supplies the runner state that commits the accepted tokens after sampling. _RECOVERSSM = os.environ.get("VLLM_GLM5NEXT_RECOVERSSM") == "1" def _arxbig_rs(): """arxbig's extension if it was set up with reduce-scatter buffers, else None.""" from vllm.distributed import get_tp_group big = getattr(get_tp_group().device_communicator, "arxbig", None) return big.ext if big is not None and not big.disabled and big.rs else None def _arx_prefetch(tensors) -> None: from vllm.distributed.device_communicators import arx arx.set_prefetch(tensors) from vllm.multimodal import MULTIMODAL_REGISTRY from vllm.platforms import current_platform from vllm.sequence import IntermediateTensors from vllm.transformers_utils.configs.glm5_next import Glm5NextConfig from .attention import Glm5NextMLAAttention from .kda import Glm5NextLinearAttention from .multimodal import ( Glm5NextMultiModalProcessor, Glm5NextProcessingInfo, Glm5NextVisionTransformer, ) logger = init_logger(__name__) class Glm5NextMLP(nn.Module): def __init__( self, hidden_size: int, intermediate_size: int, hidden_act: str, quant_config: QuantizationConfig | None = None, reduce_results: bool = True, is_sequence_parallel=False, prefix: str = "", swiglu_limit: float | None = None, ) -> None: super().__init__() # If is_sequence_parallel, the input and output tensors are sharded # across the ranks within the tp_group. In this case the weights are # replicated and no collective ops are needed. # Otherwise we use standard TP with an allreduce at the end. self.gate_up_proj = MergedColumnParallelLinear( hidden_size, [intermediate_size] * 2, bias=False, quant_config=quant_config, disable_tp=is_sequence_parallel, prefix=f"{prefix}.gate_up_proj", ) self.down_proj = RowParallelLinear( intermediate_size, hidden_size, bias=False, quant_config=quant_config, reduce_results=reduce_results, disable_tp=is_sequence_parallel, prefix=f"{prefix}.down_proj", ) if hidden_act != "silu": raise ValueError( f"Unsupported activation: {hidden_act}. Only silu is supported for now." ) self.swiglu_limit = swiglu_limit if self.swiglu_limit is not None: self.act_fn = SiluAndMulWithClamp(swiglu_limit=self.swiglu_limit) else: self.act_fn = SiluAndMul() def forward(self, x): gate_up, _ = self.gate_up_proj(x) x = self.act_fn(gate_up) x, _ = self.down_proj(x) return x class Glm5NextMoE(nn.Module): def __init__( self, config: Glm5NextConfig, parallel_config: ParallelConfig, quant_config: QuantizationConfig | None = None, prefix: str = "", apply_routed_scale_to_output: bool = False, ): super().__init__() self.tp_size = get_tensor_model_parallel_world_size() self.tp_rank = get_tensor_model_parallel_rank() self.routed_scaling_factor = config.routed_scaling_factor self.ep_group = get_ep_group().device_group self.ep_rank = get_ep_group().rank_in_group self.ep_size = self.ep_group.size() self.n_routed_experts: int = config.n_routed_experts self.n_shared_experts: int = config.n_shared_experts self.is_sequence_parallel = parallel_config.use_sequence_parallel_moe if config.hidden_act != "silu": raise ValueError( f"Unsupported activation: {config.hidden_act}. " "Only silu is supported for now." ) self.router_dtype = _get_moe_router_dtype(config) self.gate = GateLinear( config.hidden_size, config.n_routed_experts, out_dtype=self.router_dtype, prefix=f"{prefix}.gate", ) if config.topk_method == "noaux_tc": self.gate.e_score_correction_bias = nn.Parameter( torch.empty(config.n_routed_experts, dtype=torch.float32) ) else: self.gate.e_score_correction_bias = None # Load balancing settings. eplb_config = parallel_config.eplb_config self.enable_eplb = parallel_config.enable_eplb self.n_redundant_experts = eplb_config.num_redundant_experts self.n_logical_experts = self.n_routed_experts self.n_physical_experts = self.n_logical_experts + self.n_redundant_experts self.n_local_physical_experts = self.n_physical_experts // self.ep_size self.physical_expert_start = self.ep_rank * self.n_local_physical_experts self.physical_expert_end = ( self.physical_expert_start + self.n_local_physical_experts ) swiglu_limit = config.swiglu_limit if config.n_shared_experts is None: self.shared_experts = None else: intermediate_size = config.moe_intermediate_size * config.n_shared_experts self.shared_experts = Glm5NextMLP( hidden_size=config.hidden_size, intermediate_size=intermediate_size, hidden_act=config.hidden_act, quant_config=quant_config, is_sequence_parallel=self.is_sequence_parallel, reduce_results=False, prefix=f"{prefix}.shared_experts", swiglu_limit=swiglu_limit, ) self.experts = FusedMoEFactory( shared_experts=self.shared_experts, gate=self.gate, num_experts=config.n_routed_experts, top_k=config.num_experts_per_token, hidden_size=config.hidden_size, intermediate_size=config.moe_intermediate_size, renormalize=config.moe_renormalize, quant_config=quant_config, use_grouped_topk=True, num_expert_group=config.n_group, topk_group=config.topk_group, prefix=f"{prefix}.experts", scoring_func=config.scoring_func, routed_scaling_factor=self.routed_scaling_factor, apply_routed_scale_to_output=apply_routed_scale_to_output, e_score_correction_bias=self.gate.e_score_correction_bias, enable_eplb=self.enable_eplb, num_redundant_experts=self.n_redundant_experts, is_sequence_parallel=self.is_sequence_parallel, n_shared_experts=None, router_logits_dtype=self.gate.out_dtype, swiglu_limit=swiglu_limit, ) def forward( self, hidden_states: torch.Tensor, already_sequence_parallel: bool = False, ) -> torch.Tensor: num_tokens, hidden_dim = hidden_states.shape # Chunk the hidden states so they aren't replicated across TP ranks. # This avoids duplicate computation in self.experts. if self.is_sequence_parallel and not already_sequence_parallel: hidden_states = sequence_parallel_chunk(hidden_states) # MoERunner holds the gate (passed to FusedMoEFactory) and computes # the router logits itself, so nothing is precomputed here (matches # DeepseekV2MoE; `router_logits` is a placeholder). final_hidden_states = self.experts( hidden_states=hidden_states, router_logits=hidden_states ) if self.is_sequence_parallel and not already_sequence_parallel: final_hidden_states = tensor_model_parallel_all_gather( final_hidden_states, 0 ) final_hidden_states = final_hidden_states[:num_tokens] return final_hidden_states.view(num_tokens, hidden_dim) class Glm5NextDecoderLayer(nn.Module): def __init__( self, vllm_config: VllmConfig, config: Glm5NextConfig, layer_idx: int, prefix: str = "", topk_indices_buffer: torch.Tensor | None = None, is_mtp_layer: bool = False, **kwargs, ) -> None: super().__init__() cache_config = vllm_config.cache_config quant_config = vllm_config.quant_config parallel_config = vllm_config.parallel_config self.hidden_size = config.hidden_size self.layer_idx = layer_idx self.is_moe = config.is_moe self.num_hidden_layers = config.num_hidden_layers self.rms_norm_eps = config.rms_norm_eps self.num_experts = config.n_routed_experts self.is_mtp_layer = is_mtp_layer self.mhc = config.mhc is_kda_layer = not is_mtp_layer and config.is_kda_layer(layer_idx) self.layer_kind = "kda" if is_kda_layer else "mla" self.is_sequence_parallel = parallel_config.use_sequence_parallel_moe if is_kda_layer: self.self_attn = Glm5NextLinearAttention( config=config, vllm_config=vllm_config, prefix=f"{prefix}.self_attn", ) else: # MLA layers require the latent head dims, which are guaranteed set # on MLA configs; narrow away the `int | None`. assert config.v_head_dim is not None assert config.kv_lora_rank is not None self.self_attn = Glm5NextMLAAttention( vllm_config=vllm_config, config=config, hidden_size=self.hidden_size, num_heads=config.num_attention_heads, qk_nope_head_dim=config.qk_nope_head_dim, qk_rope_head_dim=config.qk_rope_head_dim, v_head_dim=config.v_head_dim, q_lora_rank=config.q_lora_rank, kv_lora_rank=config.kv_lora_rank, max_position_embeddings=config.max_position_embeddings, cache_config=cache_config, quant_config=None, # MLA projections are BF16 in checkpoint prefix=f"{prefix}.self_attn", topk_indices_buffer=topk_indices_buffer, skip_rope=config.mla_nope, ) # MTP layers sit past the base model's hidden layers (layer_idx >= # num_hidden_layers), so they're outside mlp_layer_types; default them # to the last base layer's MLP type (sparse/MoE for these checkpoints). mlp_layer_types = config.mlp_layer_types mlp_type = ( mlp_layer_types[layer_idx] if layer_idx < len(mlp_layer_types) else (mlp_layer_types[-1] if mlp_layer_types else "sparse") ) if self.is_moe and self.num_experts is not None and mlp_type == "sparse": self.mlp = Glm5NextMoE( config=config, parallel_config=parallel_config, quant_config=quant_config, prefix=f"{prefix}.mlp", ) else: self.mlp = Glm5NextMLP( hidden_size=self.hidden_size, intermediate_size=config.intermediate_size, hidden_act=config.hidden_act, quant_config=quant_config, is_sequence_parallel=self.is_sequence_parallel, prefix=f"{prefix}.mlp", swiglu_limit=config.swiglu_limit, ) self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) # Cached for the hot forward path (isinstance per layer per step). self._mlp_is_moe = isinstance(self.mlp, Glm5NextMoE) self._in_proj = getattr(self.self_attn, "in_proj_qkvbfg_a", None) # In SP, the attention output projection leaves a partial sum; the # decoder-layer reduce_scatter after attention completes it (DSv4 pattern). # MTP layers use the non-mHC path which has no sp_reduce_scatter, so # their o_proj must still reduce normally. if self.is_sequence_parallel and not is_mtp_layer: self.self_attn.o_proj.reduce_results = False self.post_attention_layernorm = RMSNorm( config.hidden_size, eps=config.rms_norm_eps ) if self.mhc and not is_mtp_layer: # mhc config self.mhc_num_residual_streams = config.mhc_num_residual_streams self.mhc_tau = config.mhc_tau self.hc_eps = config.hc_eps self.mhc_sinkhorn_iterations = config.mhc_sinkhorn_iterations self.mhc_post_mult_value = config.mhc_post_mult_value n = config.mhc_num_residual_streams d_model = n * self.hidden_size mix_hc = (2 + n) * n self.n = n # attn hc self.hc_attn_fn = nn.Parameter( torch.empty(mix_hc, d_model, dtype=torch.float32) ) self.hc_attn_base = nn.Parameter(torch.empty(mix_hc, dtype=torch.float32)) self.hc_attn_scale = nn.Parameter(torch.empty(3, dtype=torch.float32)) # ffn hc self.hc_ffn_fn = nn.Parameter( torch.empty(mix_hc, d_model, dtype=torch.float32) ) self.hc_ffn_base = nn.Parameter(torch.empty(mix_hc, dtype=torch.float32)) self.hc_ffn_scale = nn.Parameter(torch.empty(3, dtype=torch.float32)) self.mhc_pre_op = MHCPreOp() self.mhc_post_op = MHCPostOp() self.mhc_fused_post_pre_op = MHCFusedPostPreOp() if vllm_config.kernel_config.enable_jit_warmup: from vllm.model_executor.kernels.mhc.tilelang_kernels import ( _HC_PRENORM_GEMM_TILELANG_KERNEL, _MHC_FUSED_TILELANG_KERNEL, _MHC_POST_TILELANG_KERNEL, _MHC_PRE_BIG_FUSE_TILELANG_KERNEL, ) from vllm.utils.deep_gemm import is_deep_gemm_supported include_pre_gemm_splits = is_deep_gemm_supported() _MHC_PRE_BIG_FUSE_TILELANG_KERNEL.register_warmup( vllm_config, hidden_size=self.hidden_size, hc_mult=self.n, use_norm_weight=True, include_pre_gemm_splits=include_pre_gemm_splits, include_broadcast_splits=False, rms_eps=self.rms_norm_eps, hc_pre_eps=self.hc_eps, hc_sinkhorn_eps=self.hc_eps, hc_post_mult_value=self.mhc_post_mult_value, sinkhorn_repeat=self.mhc_sinkhorn_iterations, norm_eps=( self.input_layernorm.variance_epsilon, self.post_attention_layernorm.variance_epsilon, ), ) if not include_pre_gemm_splits: _HC_PRENORM_GEMM_TILELANG_KERNEL.register_warmup( vllm_config, hidden_size=self.hidden_size, hc_mult=self.n, n_out=self.n * (2 + self.n), ) _MHC_POST_TILELANG_KERNEL.register_warmup( hidden_size=self.hidden_size, hc_mult=self.n, ) _MHC_FUSED_TILELANG_KERNEL.register_warmup( vllm_config, hidden_size=self.hidden_size, hc_mult=self.n, ) def forward( self, positions: torch.Tensor, hidden_states: torch.Tensor, residual: torch.Tensor | None = None, post: torch.Tensor | None = None, comb: torch.Tensor | None = None, ) -> tuple[ torch.Tensor, torch.Tensor | None, torch.Tensor | None, torch.Tensor | None, ]: # 70B or MTP layers: KDA + MoE without HC. if not self.mhc or self.is_mtp_layer: residual = hidden_states hidden_states = self.input_layernorm(hidden_states) attn_output = self.self_attn( hidden_states=hidden_states, positions=positions, ) hidden_states, residual = self.post_attention_layernorm( attn_output, residual=residual ) hidden_states = self.mlp(hidden_states) if self.is_mtp_layer: # Return the unsummed pair: the MTP caller feeds it straight # into shared_head's fused_add_rms_norm (one kernel instead of # a separate residual-add + norm). The sum itself is unchanged # (fp32-accumulated inside the fused kernel). return hidden_states, residual, None, None hidden_states = residual + hidden_states return hidden_states, residual, None, None # mHC start. `post`/`comb` carry the previous layer's deferred # hc_post inputs (its ffn-pre outputs); when present, fuse that # hc_post with this layer's attn hc_pre into one kernel (inter-layer # fusion). Layer 0 has no incoming state -> standalone hc_pre. x = hidden_states if post is None: if self.layer_idx == 0: x = hc_expand(x, self.n) residual = x post, comb, x = self.hc_pre( x, self.hc_attn_fn, self.hc_attn_scale, self.hc_attn_base, norm_weight=self.input_layernorm.weight.data, norm_eps=self.input_layernorm.variance_epsilon, ) else: residual, post, comb, x = self.hc_fused_post_pre( x, residual, post, comb, self.hc_attn_fn, self.hc_attn_scale, self.hc_attn_base, norm_weight=self.input_layernorm.weight.data, norm_eps=self.input_layernorm.variance_epsilon, ) # Attention needs the full token sequence; mHC above ran on the SP # shard. Gather for attention, scatter back afterward (DSv4 pattern). sp = self.is_sequence_parallel or _sp_active if _ARX_PREFETCH and not sp: _arx_prefetch(self._after_attn_weights()) if (sp and _SP_FP8_GATHER and _sp_active and self._in_proj is not None and type(self._in_proj.quant_method).__name__ in ("Fp8DenseLinearMethod", "W4DenseLinearMethod")): x = self._fp8_gather(x, positions.shape[0]) elif sp: x = sp_all_gather(x)[: positions.shape[0]] x = self.self_attn( hidden_states=x, positions=positions, ) if sp: x = sp_reduce_scatter(x) # Fuse post-attn hc_post + pre-FFN hc_pre (+ RMSNorm) into one kernel. residual, post, comb, x = self.hc_fused_post_pre( x, residual, post, comb, self.hc_ffn_fn, self.hc_ffn_scale, self.hc_ffn_base, norm_weight=self.post_attention_layernorm.weight.data, norm_eps=self.post_attention_layernorm.variance_epsilon, ) # Fully Connected if _ARX_PREFETCH and not sp: nxt = self.__dict__.get("_next_layer") _arx_prefetch(nxt._before_attn_weights() if nxt is not None else ()) if _sp_active and not self.is_sequence_parallel: big = _arxbig_rs() if _SP_MOE_FUSED and self._mlp_is_moe else None if big is not None: x = self._sp_moe_fused(big, x, positions.shape[0]) else: x = sp_all_gather(x)[: positions.shape[0]] x = self.mlp(x) x = sp_reduce_scatter(x) elif self._mlp_is_moe: x = self.mlp(x, already_sequence_parallel=self.is_sequence_parallel) else: x = self.mlp(x) # mHC end. The last mHC layer materializes its final hc_post (nothing # to fuse with) then contracts; every other layer defers its hc_post to # the next layer's fused pre, returning the state. if self.layer_idx == self.num_hidden_layers - 1: x = self.hc_post(x, residual, post, comb) x = hc_contract(x, self.n) return x, None, None, None return x, residual, post, comb def _sp_moe_fused(self, big, shard: torch.Tensor, n: int, quantized: bool = _SP_MOE_QUANT_GATHER) -> torch.Tensor: """SP MoE layer: gather, route, experts, and a reduce-scatter fed row by row.""" global _sp_moe_check from vllm.model_executor.layers.fused_moe import megamoe_vllm moe = self.mlp runner = moe.experts routed = self.__dict__.get("_moe_prefill_layer") if routed is None: routed = next(m for m in moe.modules() if hasattr(m, "_moe_prefill_experts")) self.__dict__["_moe_prefill_layer"] = routed def route(x): logits = moe.gate(x) logits = logits[0] if isinstance(logits, tuple) else logits return runner.router.select_experts( hidden_states=x, router_logits=logits, topk_indices_dtype=runner._quant_method.topk_indices_dtype) H = shard.shape[1] # The quantized gather hands the shared experts an uninitialized # placeholder and their FP8 input through dense_fp8.prequant. A linear # left in BF16 would read the placeholder, so gather BF16 for those. shared_gate_up = getattr(moe.shared_experts, "gate_up_proj", None) if quantized and type(getattr(shared_gate_up, "quant_method", None)).__name__ not in ( "Fp8DenseLinearMethod", "W4DenseLinearMethod"): quantized = False if quantized: from vllm import _custom_ops as ops from vllm.model_executor.model_loader import dense_fp8 shard = shard.contiguous() rows = shard.shape[0] topk_weights, topk_ids = route(shard) topk_weights = sp_all_gather(topk_weights.float().contiguous())[:n] topk_ids = sp_all_gather(topk_ids.contiguous())[:n] x4 = shard if rows % 128 == 0 else torch.nn.functional.pad(shard, (0, 0, 0, -rows % 128)) q4, s4 = ops.scaled_fp4_quant(x4, routed._moe_prefill_a1) q4 = sp_all_gather(q4) s4 = sp_all_gather(s4.view(torch.uint8).flatten()) xq, xs = ops.scaled_fp8_quant(shard, use_per_token_if_dynamic=True) xq = sp_all_gather(xq.view(torch.uint8))[:n].view(torch.float8_e4m3fn) xs = sp_all_gather(xs)[:n] placeholder = shard.new_empty(n, H) dense_fp8.prequant(placeholder, xq, xs) shared = moe.shared_experts(placeholder) y, pos, y8s = megamoe_vllm.prefill_routed( routed, placeholder, topk_weights, topk_ids, megamoe_vllm._PREFILL_Y8, x4=(q4, s4, rows)) else: x = sp_all_gather(shard)[:n] topk_weights, topk_ids = route(x) shared = moe.shared_experts(x) y, pos, y8s = megamoe_vllm.prefill_routed(routed, x, topk_weights, topk_ids, megamoe_vllm._PREFILL_Y8) w = (topk_weights.float() * runner.routed_scaling_factor).contiguous() n_pad = shard.shape[0] * get_tensor_model_parallel_world_size() seq = big.moe_finalize_rs(y, pos, w, shared.contiguous(), n, n_pad, y8s) out = big.rs_finish(seq, shard.shape[0], H) if _sp_moe_check > 0: left, _sp_moe_check = _sp_moe_check - 1, 0 x = sp_all_gather(shard)[:n] ref = sp_reduce_scatter(self.mlp(x)) rel = ((out.float() - ref.float()).norm() / ref.float().norm()).item() logger.warning("SP MoE fused check: %d tokens, rel diff vs runner %.5f", n, rel) if quantized: alt = self._sp_moe_fused(big, shard, n, False) rel = ((out.float() - alt.float()).norm() / alt.float().norm()).item() logger.warning("SP MoE fused check: rel diff vs the bf16 gather %.6f", rel) _sp_moe_check = left return out def _after_attn_weights(self) -> list: """Weights the MoE reads first: the router gate and the shared expert.""" if not self._mlp_is_moe: return [] shared = self.mlp.shared_experts out = [self.mlp.gate.weight] if shared is not None: out += [shared.gate_up_proj.weight, shared.down_proj.weight] return out def _before_attn_weights(self) -> list: """This layer's KDA in_proj (packed weight and block scales).""" if self._in_proj is None: return [] return [self._in_proj.weight, getattr(self._in_proj, "weight_scale", None)] def _fp8_gather(self, x: torch.Tensor, n: int) -> torch.Tensor: """All-gather the SP shard as per-token FP8 for in_proj; returns its placeholder.""" from vllm import _custom_ops as ops from vllm.model_executor.model_loader import dense_fp8 xq, xs = ops.scaled_fp8_quant(x.contiguous(), use_per_token_if_dynamic=True) xq = sp_all_gather(xq.view(torch.uint8))[:n].view(torch.float8_e4m3fn) xs = sp_all_gather(xs)[:n] placeholder = x.new_empty(n, x.shape[1]) dense_fp8.prequant(placeholder, xq, xs) return placeholder def hc_pre( self, x: torch.Tensor, hc_fn: torch.Tensor, hc_scale: torch.Tensor, hc_base: torch.Tensor, norm_weight: torch.Tensor | None = None, norm_eps: float = 0.0, ): post_mix, res_mix, layer_input = self.mhc_pre_op( residual=x, fn=hc_fn, hc_scale=hc_scale, hc_base=hc_base, rms_eps=self.rms_norm_eps, hc_pre_eps=self.hc_eps, hc_sinkhorn_eps=self.hc_eps, hc_post_mult_value=self.mhc_post_mult_value, sinkhorn_repeat=self.mhc_sinkhorn_iterations, norm_weight=norm_weight, norm_eps=norm_eps, ) return post_mix, res_mix, layer_input def hc_post( self, x: torch.Tensor, residual: torch.Tensor, post: torch.Tensor, comb: torch.Tensor, ): return self.mhc_post_op(x, residual, post, comb) def hc_fused_post_pre( self, x: torch.Tensor, residual: torch.Tensor, post: torch.Tensor, comb: torch.Tensor, hc_fn: torch.Tensor, hc_scale: torch.Tensor, hc_base: torch.Tensor, norm_weight: torch.Tensor | None = None, norm_eps: float = 0.0, ): return self.mhc_fused_post_pre_op( x=x, residual=residual, post_layer_mix=post, comb_res_mix=comb, fn=hc_fn, hc_scale=hc_scale, hc_base=hc_base, rms_eps=self.rms_norm_eps, hc_pre_eps=self.hc_eps, hc_sinkhorn_eps=self.hc_eps, hc_post_mult_value=self.mhc_post_mult_value, sinkhorn_repeat=self.mhc_sinkhorn_iterations, n_splits=1, tile_n=1, norm_weight=norm_weight, norm_eps=norm_eps, ) class Glm5NextModel(nn.Module, EagleModelMixin): def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): super().__init__() config = vllm_config.model_config.hf_config self.config = config self.vocab_size = config.vocab_size self.device = current_platform.device_type self.is_v32 = config.index_topk is not None if self.is_v32: topk_tokens = config.index_topk assert topk_tokens is not None # Reserve room for the incomplete pool tail. kpool = config.index_kpool assert kpool is not None buffer_width = topk_tokens + (kpool - 1 if kpool > 1 else 0) # Sparse MLA tiles top-k in 128 columns; padded slots remain masked. sparse_topk_block_n = 128 buffer_width = ( (buffer_width + sparse_topk_block_n - 1) // sparse_topk_block_n ) * sparse_topk_block_n topk_indices_buffer = torch.empty( vllm_config.scheduler_config.max_num_batched_tokens, buffer_width, dtype=torch.int32, device=self.device, ) else: # Full-MLA config (no kpool sparse indexer): no topk buffer. topk_indices_buffer = None if get_pp_group().is_first_rank: self.embed_tokens = VocabParallelEmbedding( config.vocab_size, config.hidden_size, prefix=f"{prefix}.embed_tokens", ) else: self.embed_tokens = PPMissingLayer() def get_layer(prefix: str): layer_idx = int(prefix.rsplit(".", 1)[1]) return Glm5NextDecoderLayer( vllm_config=vllm_config, config=config, layer_idx=layer_idx, prefix=prefix, topk_indices_buffer=topk_indices_buffer, ) self.start_layer, self.end_layer, self.layers = make_layers( config.num_hidden_layers, get_layer, prefix=f"{prefix}.layers", ) # The active slice is fixed after construction; cache it so forward # doesn't rebuild the slice (a fresh list) every step. self._active_layers = self.layers[self.start_layer : self.end_layer] for a, b in zip(self._active_layers, self._active_layers[1:]): a.__dict__["_next_layer"] = b # plain attribute: not a submodule if get_pp_group().is_last_rank: self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) else: self.norm = PPMissingLayer() self.is_sequence_parallel = ( vllm_config.parallel_config.use_sequence_parallel_moe ) world_size = get_tensor_model_parallel_world_size() assert config.num_attention_heads % world_size == 0, ( "num_attention_heads must be divisible by world_size" ) def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: return self.embed_tokens(input_ids) def _completed_layer_output( self, layer: Glm5NextDecoderLayer, hidden_states: torch.Tensor, residual: torch.Tensor | None, post: torch.Tensor | None, comb: torch.Tensor | None, ) -> torch.Tensor: """The tensor an aux consumer expects: this layer's completed output. `post is None` means there is nothing deferred -- a non-mHC layer, whose hidden_states already carries residual + mlp, or the final mHC layer, already hc_post'ed and contracted. """ if post is None: return hidden_states recon = layer.hc_post(hidden_states, residual, post, comb) return hc_contract(recon, layer.n) def forward( self, input_ids: torch.Tensor | None, positions: torch.Tensor, intermediate_tensors: IntermediateTensors | None, inputs_embeds: torch.Tensor | None = None, **kwargs, ) -> torch.Tensor: if get_pp_group().is_first_rank: if inputs_embeds is not None: hidden_states = inputs_embeds else: hidden_states = self.embed_input_ids(input_ids) residual = None post = None comb = None else: assert intermediate_tensors is not None hidden_states = intermediate_tensors["hidden_states"] residual = intermediate_tensors["residual"] # post/comb (deferred mHC hc_post state) are not propagated across # PP ranks; the receiving rank's first mHC layer uses standalone pre. post = None comb = None full_num_tokens = positions.shape[0] global _sp_active sp_now = _SP_TP and not self.is_sequence_parallel and full_num_tokens >= _SP_MIN_TOKENS if sp_now != _sp_active: self._set_sp_reductions(sp_now) _sp_active = sp_now sp = self.is_sequence_parallel or sp_now if sp: hidden_states = sp_shard(hidden_states) aux_hidden_states: list[torch.Tensor] = [] for idx, layer in enumerate(self._active_layers, start=self.start_layer): hidden_states, residual, post, comb = layer( positions, hidden_states, residual, post, comb ) if idx + 1 in self.aux_hidden_state_layers: aux = self._completed_layer_output( layer, hidden_states, residual, post, comb ) if sp: # clone: the gather may return a buffer later gathers reuse aux = sp_all_gather(aux)[:full_num_tokens].clone() aux_hidden_states.append(aux) if not get_pp_group().is_last_rank: # PP is gated off for GLM-5.3-Flash (no make_empty_intermediate_tensors), # so this branch is not exercised. post/comb are the deferred # hc_post state of this rank's last mHC layer; a future PP path # would need to propagate them, but for now they are dropped (the # receiving rank's first layer would fall back to standalone pre). return IntermediateTensors( {"hidden_states": hidden_states, "residual": residual} ) if sp: hidden_states = sp_all_gather(hidden_states)[:full_num_tokens] hidden_states = self.norm(hidden_states) if self.aux_hidden_state_layers: # The runner unpacks a tuple whenever use_aux_hidden_state_outputs # is set, which is exactly when the layers were configured. return hidden_states, aux_hidden_states return hidden_states def _set_sp_reductions(self, sp: bool) -> None: """Under SP the decoder layer reduce-scatters attention and MLP outputs itself, so their own final all-reduces must be off.""" for layer in self._active_layers: if getattr(layer, "is_mtp_layer", False) or layer.is_sequence_parallel: continue layer.self_attn.o_proj.reduce_results = not sp if layer._mlp_is_moe: for mod in layer.mlp.modules(): cfg = getattr(mod, "moe_config", None) if cfg is not None and hasattr(cfg, "skip_final_all_reduce"): cfg.skip_final_all_reduce = sp else: layer.mlp.down_proj.reduce_results = not sp def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: # TP=3: zero-pad heads / MoE intermediate to the served config's sizes. # Only tp3.yaml mounts tp3pad.py, so import it only for a padded config. if getattr(self.config, "tp_pad_orig", None): from .tp3pad import pad_weights weights = pad_weights(weights, self.config) stacked_params_mapping = [ # (param_name, shard_name, shard_id) (".gate_up_proj", ".gate_proj", 0), (".gate_up_proj", ".up_proj", 1), # MLA: fuse q_a_proj and kv_a_proj_with_mqa (".fused_qkv_a_proj", ".q_a_proj", 0), (".fused_qkv_a_proj", ".kv_a_proj_with_mqa", 1), # Indexer: fuse wk and weights_proj (".wk_weights_proj", ".wk", 0), (".wk_weights_proj", ".weights_proj", 1), # KDA: merge q, k, v, b, f_a, g_a projections into one GEMM (".in_proj_qkvbfg_a", ".q_proj", 0), (".in_proj_qkvbfg_a", ".k_proj", 1), (".in_proj_qkvbfg_a", ".v_proj", 2), (".in_proj_qkvbfg_a", ".b_proj", 3), (".in_proj_qkvbfg_a", ".f_a_proj", 4), (".in_proj_qkvbfg_a", ".g_a_proj", 5), ] if self.config.is_moe: # Params for weights, fp8 weight scales, fp8 activation scales # (param_name, weight_name, expert_id, shard_id) # EPLB: the mapping enumerates physical experts, so it must cover # the redundant replicas or their slots are never loaded. num_redundant_experts = next( ( layer.mlp.n_redundant_experts for layer in self.layers if isinstance(layer, Glm5NextDecoderLayer) and isinstance(layer.mlp, Glm5NextMoE) ), 0, ) expert_params_mapping = fused_moe_make_expert_params_mapping( self, ckpt_gate_proj_name="gate_proj", ckpt_down_proj_name="down_proj", ckpt_up_proj_name="up_proj", num_experts=self.config.n_routed_experts, num_redundant_experts=num_redundant_experts, ) else: expert_params_mapping = [] params_dict = dict(self.named_parameters()) loaded_params: set[str] = set() # GLM-5.3-Flash NoPE checkpoints omit the RoPE rows from # ``kv_a_proj_with_mqa``; pad them with zeros for the model shape. kv_a_pad_size = 0 if self.config.mla_nope and self.config.qk_rope_head_dim > 0: kv_a_pad_size = self.config.qk_rope_head_dim _pending_wk_fp8: dict = {} for args in weights: name, loaded_weight = args[:2] kwargs: dict = args[2] if len(args) > 2 else {} if "rotary_emb.inv_freq" in name: continue spec_layer = get_spec_layer_idx_from_weight_name(self.config, name) if spec_layer is not None: continue # skip spec decode layers for main model if "rotary_emb.cos_cached" in name or "rotary_emb.sin_cached" in name: # Models trained using ColossalAI may include these tensors in # the checkpoint. Skip them. continue # Handle FP8 indexer WK: dequantize to BF16 for fusion with # weights_proj into wk_weights_proj. if _try_load_fp8_indexer_wk( name, loaded_weight, _pending_wk_fp8, params_dict, loaded_params, ): continue # FP8 checkpoint: dequantize BF16-kept MLA projections # (q_a_proj / kv_a_proj_with_mqa / o_proj) to BF16. if _try_load_fp8_attn_proj( name, loaded_weight, _pending_wk_fp8, params_dict, loaded_params, kv_a_pad_size, ): continue # Pad kv_a_proj_with_mqa for NoPE models if kv_a_pad_size > 0 and ".kv_a_proj_with_mqa." in name: pad = torch.zeros( kv_a_pad_size, *loaded_weight.shape[1:], dtype=loaded_weight.dtype, device=loaded_weight.device, ) loaded_weight = torch.cat([loaded_weight, pad], dim=0) for param_name, weight_name, shard_id in stacked_params_mapping: if weight_name not in name: continue # We have mlp.experts[0].gate_proj in the checkpoint. # Since we handle the experts below in expert_params_mapping, # we need to skip here BEFORE we update the name, otherwise # name will be updated to mlp.experts[0].gate_up_proj, which # will then be updated below in expert_params_mapping # for mlp.experts[0].gate_gate_up_proj, which breaks load. if ("mlp.experts." in name) and name not in params_dict: continue name_mapped = name.replace(weight_name, param_name) # QKV fusion: skip if fused module doesn't exist in model if param_name == ".fused_qkv_a_proj" and name_mapped not in params_dict: continue name = name_mapped # Skip loading extra bias for GPTQ models. if name.endswith(".bias") and name not in params_dict: continue if is_pp_missing_parameter(name, self): continue param = params_dict[name] weight_loader = param.weight_loader weight_loader(param, loaded_weight, shard_id) break else: is_expert_weight = False for ( param_name, weight_name, expert_id, expert_shard_id, ) in expert_params_mapping: if weight_name not in name: continue # A checkpoint expert may map to several physical replicas # under EPLB; keep `name` intact and try the next entry # when this physical expert is not local to the rank. is_expert_weight = True name_mapped = name.replace(weight_name, param_name) if is_pp_missing_parameter(name_mapped, self): continue param = params_dict[name_mapped] weight_loader = param.weight_loader success = weight_loader( param, loaded_weight, name_mapped, expert_id=expert_id, shard_id=expert_shard_id, return_success=True, ) if success: name = name_mapped break else: if is_expert_weight: continue # Skip loading extra bias for GPTQ models. if ( name.endswith(".bias") and name not in params_dict and not self.config.is_linear_attn ): # noqa: E501 continue # Remapping the name of FP8 kv-scale. remapped_name = maybe_remap_kv_scale_name(name, params_dict) if remapped_name is None: continue name = remapped_name if is_pp_missing_parameter(name, self): continue param = params_dict[name] weight_loader = getattr( param, "weight_loader", default_weight_loader ) weight_loader(param, loaded_weight, **kwargs) loaded_params.add(name) return loaded_params class Glm5NextForCausalLM( nn.Module, HasInnerState, SupportsPP, SupportsEagle3, MixtureOfExperts, IsHybrid ): def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): super().__init__() self.model_config = vllm_config.model_config self.vllm_config = vllm_config self.config = self.model_config.hf_config quant_config = vllm_config.quant_config self.quant_config = quant_config self.model = Glm5NextModel( vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model") ) if get_pp_group().is_last_rank: self.lm_head = ParallelLMHead( self.config.vocab_size, self.config.hidden_size, quant_config=quant_config, prefix=maybe_prefix(prefix, "lm_head"), ) else: self.lm_head = PPMissingLayer() self.logits_processor = LogitsProcessor( self.config.vocab_size, scale=self.config.logit_scale ) def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: return self.model.embed_input_ids(input_ids) if _RECOVERSSM: @staticmethod def get_model_state_cls(): from .recoverssm import model_state_cls return model_state_cls() def forward( self, input_ids: torch.Tensor | None, positions: torch.Tensor, intermediate_tensors: IntermediateTensors | None = None, inputs_embeds: torch.Tensor | None = None, **kwargs, ) -> torch.Tensor | IntermediateTensors: hidden_states = self.model( input_ids, positions, intermediate_tensors, inputs_embeds, **kwargs ) return hidden_states @classmethod def get_mamba_state_dtype_from_config( cls, vllm_config: "VllmConfig", ) -> tuple[torch.dtype, torch.dtype]: return MambaStateDtypeCalculator.kda_state_dtype( vllm_config.model_config.dtype, vllm_config.cache_config.mamba_cache_dtype ) @classmethod def get_mamba_state_shape_from_config( cls, vllm_config: "VllmConfig" ) -> tuple[tuple[int, int], tuple[int, int, int]]: parallel_config = vllm_config.parallel_config hf_config = vllm_config.model_config.hf_config tp_size = parallel_config.tensor_parallel_size num_spec = ( vllm_config.speculative_config.num_speculative_tokens if vllm_config.speculative_config else 0 ) return MambaStateShapeCalculator.kda_state_shape( tp_size, hf_config.linear_num_heads, hf_config.linear_head_dim, conv_kernel_size=hf_config.linear_conv_kernel_dim, num_spec=num_spec, ) @classmethod def get_mamba_state_copy_func( cls, ) -> tuple[ MambaStateCopyFunc, MambaStateCopyFunc, MambaStateCopyFunc, MambaStateCopyFunc ]: return MambaStateCopyFuncCalculator.kda_state_copy_func() def compute_logits( self, hidden_states: torch.Tensor, ) -> torch.Tensor | None: logits = self.logits_processor(self.lm_head, hidden_states) return logits def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: loader = AutoWeightsLoader(self) return loader.load_weights(weights) @MULTIMODAL_REGISTRY.register_processor( Glm5NextMultiModalProcessor, info=Glm5NextProcessingInfo, dummy_inputs=Glm4vDummyInputsBuilder, ) class Glm5NextForConditionalGeneration( Glm4vForConditionalGeneration, HasInnerState, IsHybrid, MixtureOfExperts, SupportsEagle3 ): # The text model (KDA + dense-MLA + MoE) is a hybrid mamba model. The # multimodal wrapper must declare the same interfaces so vLLM treats it as # hybrid (auto-aligns mamba/attention block sizes, sizes the mamba state # cache); the mamba-state classmethods delegate to the text model. has_inner_state: ClassVar[Literal[True]] = True is_hybrid: ClassVar[Literal[True]] = True # GLM-5.3-Flash stores the dense-MLP gate/up as separate tensors (like # ``Glm4vMoeForConditionalGeneration``, ``glm4_moe`` and ``deepseek_v2``), # so the fused ``gate_up_proj`` must expand to its real shard names for # per-layer quant-scheme resolution. The identity ``gate_up_proj`` entry # inherited from ``Glm4vForConditionalGeneration`` (pre-fused gate_up_proj) # would otherwise route the module to ``global_quant_config`` and mismatch # at load for mixed-precision Quark checkpoints. packed_modules_mapping = { "gate_up_proj": ["gate_proj", "up_proj"], } # NOTE: weight-prefix mapping is inherited from Glm4vForConditionalGeneration # (``model.visual.`` -> ``visual.``, ``model.language_model.`` -> # ``language_model.model.``, ``lm_head.`` -> ``language_model.lm_head.``), # matching the GLM-OCR / GLM-4V serialization convention. If the real # checkpoint's safetensors keys differ (e.g. ``language_model.model.`` with # no outer ``model.``), override ``hf_to_vllm_mapper`` accordingly. @classmethod def get_mamba_state_dtype_from_config(cls, vllm_config: VllmConfig): from .model import Glm5NextForCausalLM return Glm5NextForCausalLM.get_mamba_state_dtype_from_config(vllm_config) @classmethod def get_mamba_state_shape_from_config(cls, vllm_config: VllmConfig): from .model import Glm5NextForCausalLM return Glm5NextForCausalLM.get_mamba_state_shape_from_config(vllm_config) @classmethod def get_mamba_state_copy_func(cls): from .model import Glm5NextForCausalLM return Glm5NextForCausalLM.get_mamba_state_copy_func() if _RECOVERSSM: @staticmethod def get_model_state_cls(): from .recoverssm import model_state_cls return model_state_cls() def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): super(Glm4vForConditionalGeneration, self).__init__() config = vllm_config.model_config.hf_config multimodal_config = vllm_config.model_config.multimodal_config assert multimodal_config is not None self.config = config self.model_config = vllm_config.model_config self.multimodal_config = multimodal_config self.use_data_parallel = multimodal_config.mm_encoder_tp_mode == "data" self.is_multimodal_pruning_enabled = ( multimodal_config.is_multimodal_pruning_enabled() ) with self._mark_tower_model(vllm_config, {"image", "video"}): self.visual = Glm5NextVisionTransformer( config.text_config, config.vision_config, # Read eps from the VISION sub-config, not the top-level # `config.rms_norm_eps`: Glm5NextConfig.__getattribute__ mirrors # the latter onto text_config (1e-5), silently ignoring the # vision tower's own (1e-6) rms_norm_eps. norm_eps=config.vision_config.rms_norm_eps, # Vision tower ships BF16 weights in this fp8 checkpoint (no # weight_scale_inv for visual.*), so it must NOT inherit the # global fp8 quant_config -- doing so incorrectly quantizes # the tower # and yields NaN image features. Mirrors the MLA/KDA proj # pattern (quant_config=None for BF16 submodules). quant_config=None, prefix=maybe_prefix(prefix, "visual"), ) with self._mark_language_model(vllm_config): self.language_model = init_vllm_registered_model( vllm_config=vllm_config, hf_config=config.text_config, prefix=maybe_prefix(prefix, "language_model"), architectures=["Glm5NextForCausalLM"], ) self.set_moe_parameters() # Glm5NextForCausalLM does not implement make_empty_intermediate_tensors, # so pipeline parallelism is gated off (consistent with the text-only # model) and we intentionally do not alias it here. def set_moe_parameters(self) -> None: self.moe_mlp_layers = [ layer.mlp for layer in self.language_model.model.layers if isinstance(layer, Glm5NextDecoderLayer) and isinstance(layer.mlp, Glm5NextMoE) ] self.moe_layers = [moe.experts for moe in self.moe_mlp_layers] self.num_moe_layers = len(self.moe_layers) if not self.num_moe_layers: return example_moe = self.moe_mlp_layers[0] self.num_expert_groups = self.config.text_config.n_group self.num_logical_experts = example_moe.n_logical_experts self.num_physical_experts = example_moe.n_physical_experts self.num_local_physical_experts = example_moe.n_local_physical_experts self.num_routed_experts = example_moe.n_routed_experts self.num_shared_experts = example_moe.n_shared_experts self.num_redundant_experts = example_moe.n_redundant_experts def update_physical_experts_metadata( self, num_physical_experts: int, num_local_physical_experts: int, ) -> None: if not self.num_moe_layers: return assert self.num_local_physical_experts == num_local_physical_experts self.num_physical_experts = num_physical_experts self.num_redundant_experts = num_physical_experts - self.num_logical_experts for moe in self.moe_mlp_layers: moe.n_physical_experts = num_physical_experts moe.n_redundant_experts = self.num_redundant_experts moe.experts.update_expert_map() def get_encoder_cudagraph_config(self): # This vision tower does not produce the absolute position embedding # buffer used by GLM4V. config = super().get_encoder_cudagraph_config() config.buffer_keys = [k for k in config.buffer_keys if k != "pos_embeds"] return config def get_spec_layer_idx_from_weight_name( config: Glm5NextConfig, weight_name: str ) -> int | None: if hasattr(config, "num_nextn_predict_layers") and ( config.num_nextn_predict_layers > 0 ): layer_idx = config.num_hidden_layers for i in range(config.num_nextn_predict_layers): if weight_name.startswith( f"model.layers.{layer_idx + i}." ) or weight_name.startswith(f"layers.{layer_idx + i}."): return layer_idx + i return None def _try_load_fp8_indexer_wk(name, tensor, buf, params_dict, loaded_params): if "indexer.wk." not in name or "wk_weights" in name: return False is_weight = name.endswith(".weight") and tensor.dtype == torch.float8_e4m3fn is_scale = "weight_scale_inv" in name if not is_weight and not is_scale: return False layer_prefix = name.rsplit(".wk.", 1)[0] entry = buf.setdefault(layer_prefix, {}) entry["weight" if is_weight else "scale"] = tensor if "weight" not in entry or "scale" not in entry: return True weight_fp8, scale_inv = entry["weight"], entry["scale"] del buf[layer_prefix] block_size = weight_fp8.shape[1] // scale_inv.shape[1] weight_bf16 = scaled_dequantize( weight_fp8, scale_inv, group_shape=GroupShape(block_size, block_size), out_dtype=torch.bfloat16, ) fused_name = f"{layer_prefix}.wk_weights_proj.weight" param = params_dict[fused_name] param.weight_loader(param, weight_bf16, 0) loaded_params.add(fused_name) return True def _dequant_fp8_block( weight_fp8: torch.Tensor, scale_inv: torch.Tensor, block_size: int = 128, ) -> torch.Tensor: """Dequantize a block-FP8 (e4m3) weight with per-block scale to BF16. Unlike ``scaled_dequantize`` this tolerates a non-divisible (partial last block) shape by zero-padding to a multiple of ``block_size`` before the scale broadcast and trimming back afterwards (e.g. kv_a_proj_with_mqa is 576 rows = 4*128 + 64). """ out_dim, in_dim = weight_fp8.shape pad_out = (-out_dim) % block_size pad_in = (-in_dim) % block_size w = weight_fp8 if pad_out or pad_in: w = torch.nn.functional.pad(w, (0, pad_in, 0, pad_out)) # scale_inv is (ceil(out/block), ceil(in/block)); broadcast to (out, in). s = scale_inv.to(torch.float32) s_full = s.repeat_interleave(block_size, dim=0).repeat_interleave(block_size, dim=1) out = (w.to(torch.float32) * s_full).to(torch.bfloat16) return out[:out_dim, :in_dim].contiguous() # FP8 checkpoint projections that the MODEL keeps in BF16, so the block-FP8 # (weight + weight_scale_inv) must be dequantized to BF16 on load. # Maps checkpoint proj-suffix -> (buffer key, model target base, fused shard id # or None for a direct projection, whether NoPE rope-padding applies). _FP8_ATTN_PROJS = { ".q_a_proj.": ("q_a", "fused_qkv_a_proj", 0, False), ".kv_a_proj_with_mqa.": ("kv_a", "fused_qkv_a_proj", 1, True), ".q_b_proj.": ("q_b", "q_b_proj", None, False), ".o_proj.": ("o_proj", "o_proj", None, False), } def _try_load_fp8_attn_proj( name, tensor, buf, params_dict, loaded_params, kv_a_pad_size: int, ) -> bool: """Dequantize FP8 q_a_proj / kv_a_proj_with_mqa / o_proj to BF16 on load. The FP8 checkpoint stores these as block-FP8 (weight + weight_scale_inv), but the model holds them in BF16 (``fused_qkv_a_proj`` is always BF16 via DeepSeekV2FusedQkvAProjLinear; ``o_proj`` is excluded by modules_to_not_convert). When the model target is BF16 (no ``weight_scale_inv`` param) we dequantize; otherwise we return False so the normal stacked/direct path loads the FP8 tensor as-is. """ matched = None for suffix, info in _FP8_ATTN_PROJS.items(): if suffix in name: matched = (suffix, info) break if matched is None: return False suffix, (key, target_base, shard_id, is_kva) = matched is_weight = name.endswith(".weight") and tensor.dtype == torch.float8_e4m3fn # Need to accept both the DeepSeek-native ``weight_scale_inv`` and the Quark # ``weight_scale`` names before feeding the shared block dequant below. is_scale = "weight_scale_inv" in name or name.endswith(".weight_scale") if not is_weight and not is_scale: return False layer_prefix = name.rsplit(suffix, 1)[0] target_w = f"{layer_prefix}.{target_base}.weight" target_s = f"{layer_prefix}.{target_base}.weight_scale_inv" # If the model actually kept this projection in FP8, let the normal path # handle it (it has a weight_scale_inv param). if target_s in params_dict: return False entry = buf.setdefault(layer_prefix, {}).setdefault(key, {}) entry["weight" if is_weight else "scale"] = tensor if "weight" not in entry or "scale" not in entry: return True weight_fp8, scale_inv = entry["weight"], entry["scale"] buf[layer_prefix].pop(key, None) block_size = weight_fp8.shape[1] // scale_inv.shape[1] weight_bf16 = _dequant_fp8_block(weight_fp8, scale_inv, block_size) # NoPE: pad kv_a rope portion (kv_lora_rank -> kv_lora_rank + qk_rope_head_dim). if is_kva and kv_a_pad_size > 0: pad = torch.zeros( kv_a_pad_size, weight_bf16.shape[1], dtype=weight_bf16.dtype, device=weight_bf16.device, ) weight_bf16 = torch.cat([weight_bf16, pad], dim=0) param = params_dict[target_w] if shard_id is None: param.weight_loader(param, weight_bf16) else: param.weight_loader(param, weight_bf16, shard_id) loaded_params.add(target_w) return True