lightseekorg--tokenspeed/python/tokenspeed/runtime/layers/moe/weights/loaders.py

# Copyright (c) 2026 LightSeek Foundation
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

from __future__ import annotations

from collections.abc import Callable
from functools import partial

import torch

from tokenspeed.runtime.layers.moe.types import MoELayerSpec


def preserve_e8m0_bytes_for_uint8_param(
    dst: torch.Tensor,
    src: torch.Tensor,
) -> torch.Tensor:
    e8m0_dtype = getattr(torch, "float8_e8m0fnu", None)
    if e8m0_dtype is not None and dst.dtype == torch.uint8 and src.dtype == e8m0_dtype:
        return src.view(torch.uint8)
    return src


def load_w13(
    expert_data: torch.Tensor,
    loaded_weight: torch.Tensor,
    shard_id: str,
    shard_dim: int,
    tp_rank: int,
    is_bias: bool,
    use_presharded_weights: bool,
    do_transpose: bool,
    tp_size: int = 1,
    load_up_proj_weight_first: bool = False,
) -> None:
    if shard_id not in {"w1", "w3", "w13"}:
        raise ValueError(f"Unexpected w13 shard_id: {shard_id}")

    if is_bias:
        shard_dim = -1

    if shard_id in {"w1", "w3"}:
        shard_size = expert_data.shape[shard_dim] // 2
    else:
        shard_size = expert_data.shape[shard_dim]

    switch_w13 = load_up_proj_weight_first
    if (switch_w13 and shard_id == "w1") or (not switch_w13 and shard_id == "w3"):
        start = shard_size
    else:
        start = 0

    if not use_presharded_weights:
        if not is_bias and do_transpose:
            loaded_weight = loaded_weight.transpose(-2, -1)
        if tp_size > 1:
            # Derive the unpadded shard size from the checkpoint tensor.
            # expert_data may be Blackwell-padded; checkpoint is not.
            load_shard = loaded_weight.shape[shard_dim] // tp_size
            loaded_weight = loaded_weight.narrow(
                shard_dim, load_shard * tp_rank, load_shard
            )
        else:
            loaded_weight = loaded_weight.narrow(
                shard_dim, shard_size * tp_rank, shard_size
            )

    expert_data = expert_data.narrow(shard_dim, start, shard_size)
    dst = expert_data.narrow(shard_dim, 0, loaded_weight.shape[shard_dim])
    loaded_weight = preserve_e8m0_bytes_for_uint8_param(dst, loaded_weight)
    dst.copy_(loaded_weight)


def load_w2(
    expert_data: torch.Tensor,
    loaded_weight: torch.Tensor,
    shard_id: str,
    shard_dim: int,
    tp_rank: int,
    is_bias: bool,
    use_presharded_weights: bool,
    do_transpose: bool,
    tp_size: int = 1,
) -> None:
    if not isinstance(expert_data, torch.Tensor) or not isinstance(
        loaded_weight, torch.Tensor
    ):
        raise ValueError("expert_data and loaded_weight must be torch.Tensor")

    if shard_id != "w2":
        raise ValueError(f"shard_id must be 'w2', got {shard_id}")

    if is_bias:
        shard_dim = -1
        shard_size = expert_data.shape[-1]
    else:
        shard_size = expert_data.shape[shard_dim]

    if not use_presharded_weights:
        if not is_bias and do_transpose:
            loaded_weight = loaded_weight.transpose(-2, -1)
        if is_bias:
            load_shard = shard_size
        elif tp_size > 1:
            load_shard = loaded_weight.shape[shard_dim] // tp_size
        else:
            load_shard = shard_size
        start = 0 if is_bias else load_shard * tp_rank
        loaded_weight = loaded_weight.narrow(shard_dim, start, load_shard)

    dst = expert_data.narrow(shard_dim, 0, loaded_weight.shape[shard_dim])
    loaded_weight = preserve_e8m0_bytes_for_uint8_param(dst, loaded_weight)
    dst.copy_(loaded_weight)


def get_shard_dim(param: torch.Tensor, shard_id: str, do_transpose: bool) -> int:
    is_transposed = getattr(param, "is_transposed", False)
    if do_transpose:
        is_transposed = True

    shard_dim = {"w1": 0, "w2": 1, "w3": 0, "w13": 0}[shard_id]
    if is_transposed:
        shard_dim = int(not shard_dim)
    return shard_dim


def load_model_weight(
    param: torch.Tensor,
    loaded_weight: torch.Tensor,
    shard_id: str,
    local_expert_id: int,
    tp_rank: int,
    is_bias: bool,
    use_presharded_weights: bool,
    do_transpose: bool,
    tp_size: int = 1,
) -> None:
    expert_data = param.data[local_expert_id]
    shard_dim = get_shard_dim(param, shard_id, do_transpose)
    if shard_id == "w2":
        load_w2(
            expert_data,
            loaded_weight,
            shard_id,
            shard_dim,
            tp_rank,
            is_bias,
            use_presharded_weights,
            do_transpose,
            tp_size=tp_size,
        )
    elif shard_id in {"w1", "w3", "w13"}:
        load_w13(
            expert_data,
            loaded_weight,
            shard_id,
            shard_dim,
            tp_rank,
            is_bias,
            use_presharded_weights,
            do_transpose,
            tp_size=tp_size,
        )
    else:
        raise ValueError(f"Unknown shard_id: {shard_id}")


def load_group_weight_scale(
    param: torch.Tensor,
    loaded_weight: torch.Tensor,
    local_expert_id: int,
    shard_id: str,
    tp_rank: int,
    do_transpose: bool,
    tp_size: int = 1,
) -> None:
    load_model_weight(
        param,
        loaded_weight,
        shard_id,
        local_expert_id,
        tp_rank,
        False,
        False,
        do_transpose,
        tp_size=tp_size,
    )


def load_per_tensor_weight_scale(
    param: torch.nn.Parameter,
    loaded_weight: torch.Tensor,
    shard_id: str,
    local_expert_id: int,
) -> None:
    if shard_id in {"w1", "w3"}:
        idx = 0 if shard_id == "w1" else 1
        param.data[local_expert_id][idx] = loaded_weight
    elif shard_id == "w2":
        param.data[local_expert_id] = loaded_weight
    else:
        raise ValueError(f"Unknown shard_id: {shard_id}")


def load_per_tensor_input_scale(
    param: torch.nn.Parameter,
    loaded_weight: torch.Tensor,
    shard_id: str,
    local_expert_id: int,
) -> None:
    value = loaded_weight.detach().to(torch.float32).reshape(())
    if shard_id in {"w1", "w3"}:
        prev = param.data[local_expert_id]
        param.data[local_expert_id] = torch.maximum(prev, value)
    elif shard_id == "w2":
        param.data[local_expert_id] = value
    else:
        raise ValueError(f"Unknown shard_id for input scale: {shard_id}")


def make_weight_loader(
    spec: MoELayerSpec,
    *,
    is_bias: bool = False,
    do_transpose: bool = False,
    use_presharded_weights: bool = False,
) -> Callable:
    return partial(
        load_model_weight,
        tp_rank=spec.tp_rank,
        is_bias=is_bias,
        use_presharded_weights=use_presharded_weights,
        do_transpose=do_transpose,
        tp_size=spec.tp_size,
    )


def make_group_scale_loader(
    spec: MoELayerSpec,
    *,
    do_transpose: bool = False,
) -> Callable:
    return partial(
        load_group_weight_scale,
        tp_rank=spec.tp_rank,
        do_transpose=do_transpose,
        tp_size=spec.tp_size,
    )


def per_tensor_scale_loader() -> Callable:
    return load_per_tensor_weight_scale


def round_up(value: int, multiple: int) -> int:
    return (value + multiple - 1) // multiple * multiple


__all__ = [
    "load_per_tensor_input_scale",
    "make_group_scale_loader",
    "make_weight_loader",
    "per_tensor_scale_loader",
    "round_up",
]