"""Pad DeepSeek-V4.1 so TP=3 / EP=3 can load on 3 Sparks. Checkpoint dims that do not divide 3: vocab 129280, heads 64, o_groups 8, dspark experts 128. Pads: vocab via embedding pad_to=lcm(64,tp); heads 64→96 and groups 8→12 so local heads=32 (FlashInfer SM120 sparse-MLA only instantiates 8/16/32/64/128); draft experts 128→129. Checkpoint tensors are smaller; loaders copy the real slice and zero-pad. """ from __future__ import annotations import logging import math import os import torch logger = logging.getLogger(__name__) def _tp() -> int: return int(os.environ.get('TP_SIZE', os.environ.get('TP', '3'))) def enabled() -> bool: if os.environ.get('DSV41_TP_PAD', '1') in ('0', 'false', 'off'): return False return _tp() > 1 def _ceil_to(n: int, d: int) -> int: return ((n + d - 1) // d) * d def _pad_fill(param) -> float: """Fill value for padded slices: 1.0 for fp32 block scales, 0 for weights.""" return 1.0 if param.dtype == torch.float32 else 0.0 def pad_dsv41_hf_config(hf) -> None: if hf is None: return text = getattr(hf, 'text_config', None) or hf tp = _tp() heads = int(getattr(text, 'num_attention_heads', 0) or 0) groups = int(getattr(text, 'o_groups', 0) or 0) if heads and groups: per = max(heads // groups, 1) new_heads, new_groups = heads, groups if new_heads % tp: new_groups = _ceil_to(groups, tp) new_heads = new_groups * per # SM120 sparse-MLA kernels only exist for local head counts in this set. legal_local = (8, 16, 32, 64, 128) local = new_heads // tp if local not in legal_local: target_local = next((h for h in legal_local if h >= local), legal_local[-1]) new_heads = target_local * tp new_groups = new_heads // per if new_groups * per != new_heads: new_groups = _ceil_to(new_groups, tp) new_heads = new_groups * per if new_heads != heads or new_groups != groups: text.o_groups = new_groups text.num_attention_heads = new_heads if hasattr(hf, 'num_attention_heads'): hf.num_attention_heads = new_heads logger.warning( 'DSV41 TP pad: num_attention_heads %s→%s, o_groups %s→%s ' '(tp=%s, local_heads=%s)', heads, new_heads, groups, new_groups, tp, new_heads // tp, ) draft = int(getattr(text, 'dspark_n_routed_experts', 0) or 0) if draft and draft % tp: new_draft = _ceil_to(draft, tp) text.dspark_n_routed_experts = new_draft logger.warning( 'DSV41 TP pad: dspark_n_routed_experts %s→%s (tp=%s)', draft, new_draft, tp, ) vis = getattr(hf, 'vision_config', None) if vis is not None: vh = int(getattr(vis, 'num_attention_heads', 0) or getattr(vis, 'num_heads', 0) or 0) if vh and vh % tp: nv = _ceil_to(vh, tp) if hasattr(vis, 'num_attention_heads'): vis.num_attention_heads = nv if hasattr(vis, 'num_heads'): vis.num_heads = nv logger.warning('DSV41 TP pad: vision heads %s→%s', vh, nv) def _install_vocab_pad() -> None: import sglang.srt.layers.vocab_parallel_embedding as vpe orig = vpe.pad_vocab_size def pad_vocab_size(vocab_size: int, pad_to: int = vpe.DEFAULT_VOCAB_PADDING_SIZE) -> int: tp = _tp() if tp > 1: pad_to = math.lcm(int(pad_to), tp) return orig(vocab_size, pad_to) vpe.pad_vocab_size = pad_vocab_size def _install_config_pad() -> None: from sglang.srt.configs.model_config import ModelConfig orig = ModelConfig.__init__ def wrapped(self, *args, **kwargs): orig(self, *args, **kwargs) try: pad_dsv41_hf_config(getattr(self, 'hf_config', None)) text = getattr(getattr(self, 'hf_config', None), 'text_config', None) if text is not None: pad_dsv41_hf_config(text) if hasattr(self, 'hf_text_config') and self.hf_text_config is not None: self.hf_text_config.num_attention_heads = int(text.num_attention_heads) if hasattr(self.hf_text_config, 'o_groups'): self.hf_text_config.o_groups = int(text.o_groups) if hasattr(self, 'num_attention_heads'): self.num_attention_heads = int(text.num_attention_heads) except Exception: logger.exception('DSV41 TP config pad failed') ModelConfig.__init__ = wrapped def _install_column_pad() -> None: from sglang.srt.layers.utils import pad_or_narrow_weight from sglang.srt.layers.parameter import _ColumnvLLMParameter from sglang.srt.layers.linear import ColumnParallelLinear orig_col = _ColumnvLLMParameter.load_column_parallel_weight def load_column_parallel_weight(self, loaded_weight, tp_rank, use_presharded_weights=False): if use_presharded_weights: return orig_col(self, loaded_weight, tp_rank, use_presharded_weights) shard_size = self.data.shape[self.output_dim] start = tp_rank * shard_size end = start + shard_size dim = self.output_dim if end > loaded_weight.shape[dim] or start >= loaded_weight.shape[dim]: real = max(0, loaded_weight.shape[dim] - start) loaded_weight = pad_or_narrow_weight(loaded_weight, dim, start, shard_size) if loaded_weight.dtype == torch.float32 and real < loaded_weight.shape[dim]: # block scales: the padded region must stay a positive power of two loaded_weight = loaded_weight.clone() loaded_weight.narrow(dim, real, loaded_weight.shape[dim] - real).fill_(1.0) else: loaded_weight = loaded_weight.narrow(dim, start, shard_size) if self.data.shape != loaded_weight.shape: # Padded rows of a weight are zero; padded rows of a block scale # (fp32) are 1.0 so they stay positive powers of two and the # MXFP8 re-encoding accepts the layer (0 * 1 is still 0). self.data.fill_(_pad_fill(self.data)) # overlap copy if pad_or_narrow returned a short tensor slices = tuple(slice(0, min(a, b)) for a, b in zip(self.data.shape, loaded_weight.shape)) self.data[slices].copy_(loaded_weight[slices]) return self.data.copy_(loaded_weight) _ColumnvLLMParameter.load_column_parallel_weight = load_column_parallel_weight orig_lin = ColumnParallelLinear.weight_loader def column_weight_loader(self, param, loaded_weight): output_dim = getattr(param, 'output_dim', None) if ( output_dim is not None and not getattr(param, 'use_bitsandbytes_4bit', False) and not getattr(self, 'use_presharded_weights', False) ): shard_size = param.data.shape[output_dim] start = self.tp_rank * shard_size end = start + shard_size if end > loaded_weight.shape[output_dim]: real = loaded_weight.shape[output_dim] - start loaded_weight = pad_or_narrow_weight( loaded_weight, output_dim, start, shard_size ) if loaded_weight.dtype == torch.float32 and 0 <= real < shard_size: # block scales: padded region must stay a power of two loaded_weight.narrow(output_dim, real, shard_size - real).fill_(1.0) if param.data.shape == loaded_weight.shape: param.data.copy_(loaded_weight) return return orig_lin(self, param, loaded_weight) ColumnParallelLinear.weight_loader = column_weight_loader def _install_default_loader_pad() -> None: import sglang.srt.model_loader.weight_utils as wu orig = wu.default_weight_loader def default_weight_loader(param, loaded_weight): try: return orig(param, loaded_weight) except AssertionError: pdata = param.data pdata.zero_() if loaded_weight.numel() == 0: return slices = tuple( slice(0, min(a, b)) for a, b in zip(pdata.shape, loaded_weight.shape) ) pdata[slices].copy_(loaded_weight[slices]) wu.default_weight_loader = default_weight_loader def _install_moe_padded_loading() -> None: try: from sglang.srt.layers.moe.fused_moe_triton.layer import FusedMoE except Exception: return orig = FusedMoE.use_padded_loading.fget if hasattr(FusedMoE.use_padded_loading, 'fget') else None def use_padded_loading(self) -> bool: if orig is not None and orig(self): return True return True try: FusedMoE.use_padded_loading = property(use_padded_loading) except Exception: logger.warning('could not force FusedMoE.use_padded_loading') def install() -> None: if not enabled(): return _install_vocab_pad() _install_config_pad() _install_column_pad() _install_default_loader_pad() _install_moe_padded_loading() logger.warning('DSV41 TP pad installed (tp=%s)', _tp())