mirror of
https://github.com/kvcache-ai/ktransformers.git
synced 2026-10-18 22:55:21 +08:00
yyj
aa2cd3072f
* [feat](sft): add K2 MoE SFT backend
* [feat](sft): enable packed KGroup backward for K2 SFT
* [feat](serving): add KT LoRA to SGLang adapter converter
* [test]: cover KT SFT and composite LoRA serving paths
* [perf](bench): add reusable KT E2E benchmark harness
* [test](sft): add K2 MoE AMX and real Kimi load smoke cases
* [perf](bench): add K2 SFT packed backward profile harness
* [test](sglang): cover Qwen3.5 LoRA shape filters
* [perf](sft): add K2 forward profiling logs
* [docs](sft): clarify K2 packed backward flow
* [docs](sft): number K2 backward comment steps
* chore: mirror sap4 loader baseline
* Add K2 SFT profiling ranges
* Add backward NVTX ranges for K2 SFT profiling
* perf(sft): add rank8 LoRA backward vector paths
* perf(sft): reuse down LoRA bprop cache
* [perf]: add K2 RAWINT4 prefill mat-mat dispatch
* [perf] Refactor K2 RAWINT4 prefill accumulation
* [perf]: add blocked RAWINT4 prefill backend
Cherry-picked from be40f8e with local amx.py conflict resolution.
* Optimize blocked RAWINT4 decode vecmul
* perf(sft): checkpoint qj5090 packed backward baseline
* perf(sft): batch expert LoRA backward as matmul
* perf(sft): align K2 base backward with BF16 matmul
* Optimize K2 SFT backward workspace path
* docs(sft): hand off RAWINT4 KGroup build and Kimi configs
* [feat](kt-kernel): add Kimi RAWINT4 LoRA SFT kernels
* [feat](kt-kernel): integrate Kimi RAWINT4 LoRA SFT
* [fix](sft): preserve input gradients through frozen routers
Build the routing graph when either activations or router parameters require gradients, while respecting the ambient grad mode. Keep frozen router parameters frozen and reject detached routing weights whenever gradients are needed.
Correct the existing frozen-router expectation and cover input-gradient equivalence for linear, DeepSeek/Kimi and GLM routing, checkpoint reuse, grad modes and detached-output failures.
Validated 62 routing/shared-expert/checkpoint tests on PyTorch 2.9.1 and a two-step Kimi K2.5 RAWINT4 four-GPU S512 LoRA smoke with per-rank routing-input gradient audits.
* [perf](kt-kernel): accelerate RAWINT4 prefill and backward
Reuse decoded weight columns across prefill rows and prepack BF16 backward coefficients into bounded per-worker panels. Use AMX BF16 tiles when available, retaining the AVX512-BF16 path and frozen base weights.
Balance expert work, parallelize cache copies and router reductions, and correct the signed INT4 -8 dot-product boundary. Add a CPU replay benchmark and forward/backward panel regressions.
Validated on sap4 with Kimi K2.5, TP1/TP2 replay, AMX and AVX-512 multi-GPU training, gradient windows, and saved LoRA artifacts. User-facing clean builds and AMD hardware remain separate validation work.
* [fix](sft): limit default OpenMP pools to the CPU owner
Keep the affinity-visible physical-core budget on rank 0 and default non-owner distributed ranks to one thread. Preserve explicit generic and KT-specific thread overrides, including intentional single-thread settings.
Cover single-process behavior, non-owner ranks, repeated configuration, and explicit overrides. The 109 configuration, activation-policy, and backend-dispatch tests pass.
---------
Co-authored-by: yyj <yyj@qj5090.local>
Co-authored-by: djh <jd4136@columbia.edu>
932 lines
40 KiB
Python
932 lines
40 KiB
Python
# AMX SFT MoE Wrapper implementation
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
"""
|
|
AMX-based SFT MoE Wrapper. Forward/backward buffer management is in base class;
|
|
this file handles weight loading, LoRA init, and C++ task construction.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import ctypes
|
|
import logging
|
|
import os
|
|
import glob as _glob
|
|
import torch
|
|
from typing import Optional, List
|
|
|
|
from kt_kernel_ext.moe import MOESFTConfig
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
from ..utils.loader import BF16SafeTensorLoader, SafeTensorLoader
|
|
|
|
try:
|
|
from kt_kernel_ext.moe import (
|
|
AMXBF16_SFT_MOE,
|
|
AMXInt8_SFT_MOE,
|
|
AMXInt4_SFT_MOE,
|
|
AMXBF16_SFT_MOE_SkipLoRA,
|
|
AMXInt8_SFT_MOE_SkipLoRA,
|
|
AMXInt4_SFT_MOE_SkipLoRA,
|
|
)
|
|
|
|
_HAS_AMX_SFT_SUPPORT = True
|
|
except (ImportError, AttributeError):
|
|
_HAS_AMX_SFT_SUPPORT = False
|
|
AMXBF16_SFT_MOE = None
|
|
AMXInt8_SFT_MOE = None
|
|
AMXInt4_SFT_MOE = None
|
|
AMXBF16_SFT_MOE_SkipLoRA = None
|
|
AMXInt8_SFT_MOE_SkipLoRA = None
|
|
AMXInt4_SFT_MOE_SkipLoRA = None
|
|
|
|
try:
|
|
from kt_kernel_ext.moe import AMXFP8_SFT_MOE
|
|
except (ImportError, AttributeError):
|
|
AMXFP8_SFT_MOE = None
|
|
|
|
try:
|
|
from kt_kernel_ext.moe import AMXInt4_KGroup_SFT_MOE
|
|
except (ImportError, AttributeError):
|
|
AMXInt4_KGroup_SFT_MOE = None
|
|
|
|
from .base import BaseSFTMoEWrapper, KExpertsSFTBuffer, _supports_authoritative_optimizer_grads
|
|
from .backend import is_fp8_sft_method, is_int8_sft_method, is_rawint4_sft_method
|
|
from .weights import BlockFP8ExpertWeights, RAWINT4ExpertWeights
|
|
|
|
_AMX_M_STEP = 32
|
|
|
|
|
|
# Mapping from method string to C++ SFT MOE class
|
|
_SFT_METHOD_TO_CLASS = {
|
|
"AMXBF16_SFT": AMXBF16_SFT_MOE,
|
|
"AMXFP8_SFT": AMXFP8_SFT_MOE,
|
|
"INT8_SFT": AMXInt8_SFT_MOE,
|
|
"AMXINT8_SFT": AMXInt8_SFT_MOE,
|
|
"AMXINT4_SFT": AMXInt4_SFT_MOE,
|
|
"RAWINT4_SFT": AMXInt4_KGroup_SFT_MOE,
|
|
"AMXINT4_KGroup_SFT": AMXInt4_KGroup_SFT_MOE,
|
|
"AMXBF16_SFT_SkipLoRA": AMXBF16_SFT_MOE_SkipLoRA,
|
|
"AMXINT8_SFT_SkipLoRA": AMXInt8_SFT_MOE_SkipLoRA,
|
|
"AMXINT4_SFT_SkipLoRA": AMXInt4_SFT_MOE_SkipLoRA,
|
|
}
|
|
|
|
|
|
class AMXSFTMoEWrapper(BaseSFTMoEWrapper):
|
|
"""
|
|
AMX-based SFT MoE wrapper.
|
|
|
|
Supports BF16, native block-FP8, INT8, INT4, and SkipLoRA variants.
|
|
Forward/backward buffer management is in BaseSFTMoEWrapper;
|
|
this class implements weight loading and C++ task construction.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
layer_idx: int,
|
|
num_experts: int,
|
|
num_experts_per_tok: int,
|
|
hidden_size: int,
|
|
moe_intermediate_size: int,
|
|
num_gpu_experts: int,
|
|
cpuinfer_threads: int,
|
|
threadpool_count: int,
|
|
weight_path: str,
|
|
chunked_prefill_size: int,
|
|
lora_rank: int = 16,
|
|
lora_alpha: float = 32.0,
|
|
lora_dropout: float = 0.0,
|
|
max_cache_depth: int = 1,
|
|
method: str = "AMXBF16_SFT",
|
|
group_size: int = 128,
|
|
zero_point: bool = True,
|
|
full_weight_grad: bool = False,
|
|
):
|
|
if not _HAS_AMX_SFT_SUPPORT:
|
|
raise RuntimeError(
|
|
"AMX SFT backend not available. kt_kernel_ext was not compiled with AMX SFT support.\n"
|
|
"Please recompile with AMX SFT enabled."
|
|
)
|
|
|
|
super().__init__(
|
|
layer_idx=layer_idx,
|
|
num_experts=num_experts,
|
|
num_experts_per_tok=num_experts_per_tok,
|
|
hidden_size=hidden_size,
|
|
moe_intermediate_size=moe_intermediate_size,
|
|
num_gpu_experts=num_gpu_experts,
|
|
cpuinfer_threads=cpuinfer_threads,
|
|
threadpool_count=threadpool_count,
|
|
weight_path=weight_path,
|
|
chunked_prefill_size=chunked_prefill_size,
|
|
lora_rank=lora_rank,
|
|
lora_alpha=lora_alpha,
|
|
lora_dropout=lora_dropout,
|
|
max_cache_depth=max_cache_depth,
|
|
full_weight_grad=full_weight_grad,
|
|
)
|
|
|
|
self.method = method
|
|
self._is_rawint4 = is_rawint4_sft_method(method)
|
|
self._is_skip_lora = "SkipLoRA" in method
|
|
self._uses_authoritative_optimizer_grads = _supports_authoritative_optimizer_grads(
|
|
method,
|
|
self.num_gpu_experts,
|
|
full_weight_grad=self._full_weight_grad,
|
|
lora_rank=self.lora_rank,
|
|
)
|
|
self.group_size = group_size
|
|
self.zero_point = zero_point
|
|
|
|
if is_fp8_sft_method(method):
|
|
if self._full_weight_grad or self.lora_rank <= 0:
|
|
raise ValueError("AMXFP8_SFT supports frozen-base LoRA only")
|
|
if self.num_gpu_experts != 0:
|
|
raise ValueError("AMXFP8_SFT requires all routed experts on CPU")
|
|
if self.hidden_size % 128 or self.moe_intermediate_size % 128:
|
|
raise ValueError(
|
|
"AMXFP8_SFT requires hidden_size and moe_intermediate_size "
|
|
"divisible by 128"
|
|
)
|
|
if self.threadpool_count < 1 or self.moe_intermediate_size % self.threadpool_count:
|
|
raise ValueError(
|
|
"AMXFP8_SFT requires moe_intermediate_size divisible by "
|
|
"threadpool_count"
|
|
)
|
|
if (self.moe_intermediate_size // self.threadpool_count) % 128:
|
|
raise ValueError(
|
|
"AMXFP8_SFT requires every TP intermediate slice divisible by 128"
|
|
)
|
|
self.group_size = 128
|
|
self.zero_point = False
|
|
|
|
if is_rawint4_sft_method(self.method):
|
|
if self._full_weight_grad:
|
|
raise ValueError("RAWINT4 SFT supports frozen-base LoRA only")
|
|
if self.lora_rank != 8:
|
|
raise ValueError("RAWINT4 SFT v1 requires LoRA rank 8")
|
|
if self.lora_dropout != 0.0:
|
|
raise ValueError("RAWINT4 SFT v1 requires LoRA dropout 0")
|
|
if self._is_skip_lora:
|
|
raise ValueError("RAWINT4 SFT v1 does not support SkipLoRA")
|
|
if self.num_gpu_experts != 0:
|
|
raise ValueError("RAWINT4 SFT requires all routed experts on CPU")
|
|
if self.threadpool_count not in {1, 2}:
|
|
raise ValueError("RAWINT4 SFT v1 supports TP1/TP2 only")
|
|
if self.hidden_size % 32 or self.moe_intermediate_size % 32:
|
|
raise ValueError("RAWINT4 SFT requires hidden_size and moe_intermediate_size " "divisible by 32")
|
|
if (
|
|
self.moe_intermediate_size % self.threadpool_count
|
|
or (self.moe_intermediate_size // self.threadpool_count) % 32
|
|
):
|
|
raise ValueError("RAWINT4 SFT requires every TP intermediate slice divisible by 32")
|
|
self.group_size = 32
|
|
self.zero_point = False
|
|
|
|
if method not in _SFT_METHOD_TO_CLASS:
|
|
raise ValueError(f"Unknown SFT method: {method}. Supported: {list(_SFT_METHOD_TO_CLASS.keys())}")
|
|
|
|
moe_class = _SFT_METHOD_TO_CLASS[method]
|
|
if moe_class is None:
|
|
raise RuntimeError(f"AMX SFT method '{method}' not available in current build.")
|
|
|
|
self.gate_proj: Optional[torch.Tensor] = None
|
|
self.up_proj: Optional[torch.Tensor] = None
|
|
self.down_proj: Optional[torch.Tensor] = None
|
|
self.gate_scale: Optional[torch.Tensor] = None
|
|
self.up_scale: Optional[torch.Tensor] = None
|
|
self.down_scale: Optional[torch.Tensor] = None
|
|
|
|
self._moe_class = moe_class
|
|
|
|
# ========== Template method: C++ task construction ==========
|
|
|
|
def _make_forward_task(self, buffer: KExpertsSFTBuffer, save_for_backward: bool):
|
|
return self.moe.forward_sft_task(
|
|
buffer.bsz_tensor.data_ptr(),
|
|
self.num_experts_per_tok,
|
|
buffer.expert_ids_cpu.data_ptr(),
|
|
buffer.weights_cpu.data_ptr(),
|
|
buffer.input_cpu.data_ptr(),
|
|
buffer.output_cpu.data_ptr(),
|
|
save_for_backward,
|
|
)
|
|
|
|
def _make_backward_task(
|
|
self,
|
|
buffer: KExpertsSFTBuffer,
|
|
accumulate_optimizer_grads: bool = False,
|
|
optimizer_grad_scale: float = 1.0,
|
|
):
|
|
if self._is_skip_lora:
|
|
return self.moe.backward_task(
|
|
buffer.grad_output_cpu.data_ptr(),
|
|
buffer.grad_input_cpu.data_ptr(),
|
|
0,
|
|
0,
|
|
0,
|
|
0,
|
|
0,
|
|
0,
|
|
buffer.grad_weights.data_ptr(),
|
|
0,
|
|
0,
|
|
0, # grad_gate_proj, grad_up_proj, grad_down_proj
|
|
)
|
|
|
|
# Base weight grad pointers (nullptr if not in full mode)
|
|
grad_gate_proj_ptr = (
|
|
self.grad_gate_proj_buf.data_ptr() if self._full_weight_grad and self.grad_gate_proj_buf is not None else 0
|
|
)
|
|
grad_up_proj_ptr = (
|
|
self.grad_up_proj_buf.data_ptr() if self._full_weight_grad and self.grad_up_proj_buf is not None else 0
|
|
)
|
|
grad_down_proj_ptr = (
|
|
self.grad_down_proj_buf.data_ptr() if self._full_weight_grad and self.grad_down_proj_buf is not None else 0
|
|
)
|
|
|
|
backward_args = (
|
|
buffer.grad_output_cpu.data_ptr(),
|
|
buffer.grad_input_cpu.data_ptr(),
|
|
self.grad_gate_lora_a.data_ptr() if self.lora_rank > 0 else 0,
|
|
self.grad_gate_lora_b.data_ptr() if self.lora_rank > 0 else 0,
|
|
self.grad_up_lora_a.data_ptr() if self.lora_rank > 0 else 0,
|
|
self.grad_up_lora_b.data_ptr() if self.lora_rank > 0 else 0,
|
|
self.grad_down_lora_a.data_ptr() if self.lora_rank > 0 else 0,
|
|
self.grad_down_lora_b.data_ptr() if self.lora_rank > 0 else 0,
|
|
buffer.grad_weights.data_ptr(),
|
|
grad_gate_proj_ptr,
|
|
grad_up_proj_ptr,
|
|
grad_down_proj_ptr,
|
|
)
|
|
if (
|
|
self._uses_authoritative_optimizer_grads
|
|
or bool(accumulate_optimizer_grads)
|
|
or float(optimizer_grad_scale) != 1.0
|
|
):
|
|
return self.moe.backward_task(
|
|
*backward_args,
|
|
bool(accumulate_optimizer_grads),
|
|
float(optimizer_grad_scale),
|
|
)
|
|
return self.moe.backward_task(*backward_args)
|
|
|
|
# ========== Weight loading ==========
|
|
|
|
def load_weights(self, physical_to_logical_map_cpu: torch.Tensor) -> None:
|
|
if self._weights_loaded:
|
|
return
|
|
|
|
if physical_to_logical_map_cpu is None:
|
|
physical_to_logical_map_cpu = torch.arange(self.num_experts, dtype=torch.int64)
|
|
strict_rawint4_map = bool(getattr(self, "_is_rawint4", False))
|
|
if strict_rawint4_map:
|
|
if not isinstance(physical_to_logical_map_cpu, torch.Tensor):
|
|
raise TypeError("physical_to_logical_map_cpu must be a torch.Tensor")
|
|
if physical_to_logical_map_cpu.dtype not in {torch.int32, torch.int64}:
|
|
raise ValueError(
|
|
"physical_to_logical_map_cpu must use an integer dtype, "
|
|
f"got {physical_to_logical_map_cpu.dtype}."
|
|
)
|
|
self._physical_to_logical_map_cpu = physical_to_logical_map_cpu.to(
|
|
dtype=torch.int64, device="cpu"
|
|
).reshape(-1).contiguous()
|
|
if strict_rawint4_map and self._physical_to_logical_map_cpu.numel() != self.num_experts:
|
|
raise ValueError(
|
|
"physical_to_logical_map_cpu must contain exactly "
|
|
f"{self.num_experts} entries, got {self._physical_to_logical_map_cpu.numel()}."
|
|
)
|
|
if not strict_rawint4_map and self._physical_to_logical_map_cpu.numel() < self.num_experts:
|
|
raise ValueError(
|
|
"physical_to_logical_map_cpu must contain at least "
|
|
f"{self.num_experts} entries, got {self._physical_to_logical_map_cpu.numel()}."
|
|
)
|
|
expected_experts = torch.arange(self.num_experts, dtype=torch.int64)
|
|
if strict_rawint4_map and not torch.equal(
|
|
self._physical_to_logical_map_cpu.sort().values, expected_experts
|
|
):
|
|
raise ValueError(
|
|
"physical_to_logical_map_cpu must be a permutation of "
|
|
f"[0, {self.num_experts}); got {self._physical_to_logical_map_cpu.tolist()}."
|
|
)
|
|
|
|
if self.gate_proj is None and not getattr(self, "_use_projs_path", False):
|
|
self._load_base_weights_from_file()
|
|
|
|
config = MOESFTConfig()
|
|
config.expert_num = self.num_experts
|
|
config.num_experts_per_tok = self.num_experts_per_tok
|
|
config.hidden_size = self.hidden_size
|
|
config.intermediate_size = self.moe_intermediate_size
|
|
config.lora_rank = self.lora_rank
|
|
config.lora_alpha = self.lora_alpha
|
|
config.lora_dropout = self.lora_dropout
|
|
config.max_cache_depth = self.max_cache_depth
|
|
config.max_len = self._aligned_max_len()
|
|
config.layer_idx = self.layer_idx
|
|
config.share_backward_bb = getattr(self, "share_backward_bb", False)
|
|
config.share_cache_pool = getattr(self, "share_cache_pool", False)
|
|
config.full_weight_grad = self._full_weight_grad
|
|
config.authoritative_optimizer_grads = self._uses_authoritative_optimizer_grads
|
|
config.physical_to_logical_map = self._physical_to_logical_map_cpu.data_ptr()
|
|
|
|
if getattr(self, "_use_kt_direct_load", False):
|
|
config.load = True
|
|
config.path = self.weight_path
|
|
elif getattr(self, "_use_projs_path", False):
|
|
config.gate_projs = self._gate_projs_ptrs
|
|
config.up_projs = self._up_projs_ptrs
|
|
config.down_projs = self._down_projs_ptrs
|
|
config.gate_scales = self._gate_scale_ptrs
|
|
config.up_scales = self._up_scale_ptrs
|
|
config.down_scales = self._down_scale_ptrs
|
|
if getattr(self, "_bf16_gate_proj", None) is not None:
|
|
config.gate_proj = self._bf16_gate_proj.data_ptr()
|
|
config.up_proj = self._bf16_up_proj.data_ptr()
|
|
config.down_proj = self._bf16_down_proj.data_ptr()
|
|
if getattr(self, "_has_bwd_projs", False):
|
|
config.gate_bwd_projs = self._gate_bwd_projs_ptrs
|
|
config.up_bwd_projs = self._up_bwd_projs_ptrs
|
|
config.down_bwd_projs = self._down_bwd_projs_ptrs
|
|
config.gate_bwd_scales = self._gate_bwd_scale_ptrs
|
|
config.up_bwd_scales = self._up_bwd_scale_ptrs
|
|
config.down_bwd_scales = self._down_bwd_scale_ptrs
|
|
else:
|
|
config.gate_proj = self.gate_proj.data_ptr()
|
|
config.up_proj = self.up_proj.data_ptr()
|
|
config.down_proj = self.down_proj.data_ptr()
|
|
if getattr(self, "_use_rawint4_tensor_path", False):
|
|
config.gate_scale = self.gate_scale.data_ptr()
|
|
config.up_scale = self.up_scale.data_ptr()
|
|
config.down_scale = self.down_scale.data_ptr()
|
|
|
|
if self._lora_initialized:
|
|
config.gate_lora_a = self.gate_lora_a.data_ptr()
|
|
config.gate_lora_b = self.gate_lora_b.data_ptr()
|
|
config.up_lora_a = self.up_lora_a.data_ptr()
|
|
config.up_lora_b = self.up_lora_b.data_ptr()
|
|
config.down_lora_a = self.down_lora_a.data_ptr()
|
|
config.down_lora_b = self.down_lora_b.data_ptr()
|
|
|
|
config.pool = self.cpu_infer.backend_
|
|
|
|
if is_rawint4_sft_method(self.method):
|
|
config.quant_config.bits = 4
|
|
config.quant_config.group_size = self.group_size
|
|
config.quant_config.zero_point = self.zero_point
|
|
elif is_fp8_sft_method(self.method):
|
|
config.quant_config.bits = 8
|
|
config.quant_config.group_size = 128
|
|
config.quant_config.zero_point = False
|
|
|
|
# Release old C++ MOE object before creating a new one to avoid memory leak
|
|
old_moe = getattr(self, "moe", None)
|
|
if old_moe is not None:
|
|
del old_moe
|
|
import gc
|
|
gc.collect()
|
|
|
|
self.moe = self._moe_class(config)
|
|
|
|
self.cpu_infer.submit(self.moe.load_weights_task())
|
|
self.cpu_infer.sync()
|
|
|
|
if os.environ.get("KT_SFT_ENABLE_WARMUP", "0") == "1":
|
|
self.cpu_infer.submit(self.moe.warm_up_task())
|
|
self.cpu_infer.sync()
|
|
|
|
# Release Python-side weight tensors (C++ copied them)
|
|
# In full_weight_grad mode, keep them for nn.Parameter initialization
|
|
if not self._full_weight_grad:
|
|
self.gate_proj = None
|
|
self.up_proj = None
|
|
self.down_proj = None
|
|
self.gate_scale = None
|
|
self.up_scale = None
|
|
self.down_scale = None
|
|
|
|
if getattr(self, "_bf16_gate_proj", None) is not None:
|
|
self._bf16_gate_proj = None
|
|
self._bf16_up_proj = None
|
|
self._bf16_down_proj = None
|
|
|
|
if getattr(self, "_use_projs_path", False):
|
|
for attr in [
|
|
"_gate_weights_per_numa",
|
|
"_up_weights_per_numa",
|
|
"_down_weights_per_numa",
|
|
"_gate_scales_per_numa",
|
|
"_up_scales_per_numa",
|
|
"_down_scales_per_numa",
|
|
"_gate_projs_ptrs",
|
|
"_up_projs_ptrs",
|
|
"_down_projs_ptrs",
|
|
"_gate_scale_ptrs",
|
|
"_up_scale_ptrs",
|
|
"_down_scale_ptrs",
|
|
]:
|
|
setattr(self, attr, None)
|
|
if getattr(self, "_has_bwd_projs", False):
|
|
for attr in [
|
|
"_gate_bwd_weights_per_numa",
|
|
"_up_bwd_weights_per_numa",
|
|
"_down_bwd_weights_per_numa",
|
|
"_gate_bwd_scales_per_numa",
|
|
"_up_bwd_scales_per_numa",
|
|
"_down_bwd_scales_per_numa",
|
|
"_gate_bwd_projs_ptrs",
|
|
"_up_bwd_projs_ptrs",
|
|
"_down_bwd_projs_ptrs",
|
|
"_gate_bwd_scale_ptrs",
|
|
"_up_bwd_scale_ptrs",
|
|
"_down_bwd_scale_ptrs",
|
|
]:
|
|
setattr(self, attr, None)
|
|
|
|
self._weights_loaded = True
|
|
|
|
def load_weights_from_tensors(
|
|
self,
|
|
gate_proj: torch.Tensor,
|
|
up_proj: torch.Tensor,
|
|
down_proj: torch.Tensor,
|
|
physical_to_logical_map_cpu: torch.Tensor,
|
|
) -> None:
|
|
if is_rawint4_sft_method(self.method):
|
|
raise ValueError(
|
|
"RAWINT4 SFT accepts packed weights plus BF16 group scales only; " "use load_rawint4_weights()"
|
|
)
|
|
if is_fp8_sft_method(self.method):
|
|
raise ValueError(
|
|
"AMXFP8_SFT accepts raw per-expert E4M3 weights only; "
|
|
"use load_block_fp8_weights()"
|
|
)
|
|
if is_int8_sft_method(self.method):
|
|
raise ValueError(
|
|
"INT8_SFT accepts pre-quantized .kt weights only; "
|
|
"online tensor conversion is not supported"
|
|
)
|
|
self.gate_proj = gate_proj.contiguous()
|
|
self.up_proj = up_proj.contiguous()
|
|
self.down_proj = down_proj.contiguous()
|
|
self.load_weights(physical_to_logical_map_cpu)
|
|
del gate_proj, up_proj, down_proj
|
|
|
|
def load_rawint4_weights(
|
|
self,
|
|
weights: RAWINT4ExpertWeights,
|
|
physical_to_logical_map_cpu: torch.Tensor,
|
|
) -> None:
|
|
"""Synchronously load strict signed group-32 packed expert weights."""
|
|
|
|
if not is_rawint4_sft_method(self.method):
|
|
raise ValueError(f"load_rawint4_weights() requires RAWINT4_SFT, got {self.method!r}")
|
|
if int(weights.group_size) != 32:
|
|
raise ValueError(f"RAWINT4 SFT requires group_size=32, got {weights.group_size}")
|
|
|
|
projections = {
|
|
"gate": (
|
|
weights.gate_proj,
|
|
weights.gate_scale,
|
|
(self.num_experts, self.moe_intermediate_size, self.hidden_size // 2),
|
|
(self.num_experts, self.moe_intermediate_size, self.hidden_size // 32),
|
|
),
|
|
"up": (
|
|
weights.up_proj,
|
|
weights.up_scale,
|
|
(self.num_experts, self.moe_intermediate_size, self.hidden_size // 2),
|
|
(self.num_experts, self.moe_intermediate_size, self.hidden_size // 32),
|
|
),
|
|
"down": (
|
|
weights.down_proj,
|
|
weights.down_scale,
|
|
(self.num_experts, self.hidden_size, self.moe_intermediate_size // 2),
|
|
(self.num_experts, self.hidden_size, self.moe_intermediate_size // 32),
|
|
),
|
|
}
|
|
for name, (packed, scale, packed_shape, scale_shape) in projections.items():
|
|
if (
|
|
packed.device.type != "cpu"
|
|
or packed.dtype != torch.uint8
|
|
or not packed.is_contiguous()
|
|
or tuple(packed.shape) != packed_shape
|
|
):
|
|
raise ValueError(f"{name} RAWINT4 weight must be contiguous CPU uint8 with " f"shape {packed_shape}")
|
|
if (
|
|
scale.device.type != "cpu"
|
|
or scale.dtype != torch.bfloat16
|
|
or not scale.is_contiguous()
|
|
or tuple(scale.shape) != scale_shape
|
|
):
|
|
raise ValueError(f"{name} RAWINT4 scale must be contiguous CPU BF16 with " f"shape {scale_shape}")
|
|
|
|
self.gate_proj = weights.gate_proj
|
|
self.gate_scale = weights.gate_scale
|
|
self.up_proj = weights.up_proj
|
|
self.up_scale = weights.up_scale
|
|
self.down_proj = weights.down_proj
|
|
self.down_scale = weights.down_scale
|
|
self._use_rawint4_tensor_path = True
|
|
self.load_weights(physical_to_logical_map_cpu)
|
|
|
|
def load_block_fp8_weights(
|
|
self,
|
|
weights: BlockFP8ExpertWeights,
|
|
physical_to_logical_map_cpu: torch.Tensor,
|
|
) -> None:
|
|
"""Synchronously pack raw per-expert FP8 checkpoint tensors in C++."""
|
|
|
|
if not is_fp8_sft_method(self.method):
|
|
raise ValueError(
|
|
f"load_block_fp8_weights() requires AMXFP8_SFT, got {self.method!r}"
|
|
)
|
|
if tuple(weights.block_size) != (128, 128):
|
|
raise ValueError(
|
|
"AMXFP8_SFT requires block_size=(128, 128), "
|
|
f"got {weights.block_size}"
|
|
)
|
|
|
|
projections = {
|
|
"gate": (weights.gate_proj, weights.gate_scale, (self.moe_intermediate_size, self.hidden_size)),
|
|
"up": (weights.up_proj, weights.up_scale, (self.moe_intermediate_size, self.hidden_size)),
|
|
"down": (weights.down_proj, weights.down_scale, (self.hidden_size, self.moe_intermediate_size)),
|
|
}
|
|
for name, (expert_weights, expert_scales, expected_shape) in projections.items():
|
|
if len(expert_weights) != self.num_experts or len(expert_scales) != self.num_experts:
|
|
raise ValueError(
|
|
f"{name} FP8 tensors must contain {self.num_experts} experts, got "
|
|
f"weights={len(expert_weights)}, scales={len(expert_scales)}"
|
|
)
|
|
expected_scale_shape = (expected_shape[0] // 128, expected_shape[1] // 128)
|
|
for expert_idx, (weight, scale) in enumerate(zip(expert_weights, expert_scales)):
|
|
if (
|
|
weight.device.type != "cpu"
|
|
or weight.dtype != torch.float8_e4m3fn
|
|
or not weight.is_contiguous()
|
|
or tuple(weight.shape) != expected_shape
|
|
):
|
|
raise ValueError(
|
|
f"{name} expert {expert_idx} must be contiguous CPU E4M3FN "
|
|
f"with shape {expected_shape}"
|
|
)
|
|
if (
|
|
scale.device.type != "cpu"
|
|
or scale.dtype != torch.float32
|
|
or not scale.is_contiguous()
|
|
or tuple(scale.shape) != expected_scale_shape
|
|
):
|
|
raise ValueError(
|
|
f"{name} scale {expert_idx} must be contiguous CPU FP32 "
|
|
f"with shape {expected_scale_shape}"
|
|
)
|
|
|
|
self._gate_weights_per_numa = [weights.gate_proj]
|
|
self._up_weights_per_numa = [weights.up_proj]
|
|
self._down_weights_per_numa = [weights.down_proj]
|
|
self._gate_scales_per_numa = [weights.gate_scale]
|
|
self._up_scales_per_numa = [weights.up_scale]
|
|
self._down_scales_per_numa = [weights.down_scale]
|
|
self._gate_projs_ptrs = [[tensor.data_ptr() for tensor in weights.gate_proj]]
|
|
self._up_projs_ptrs = [[tensor.data_ptr() for tensor in weights.up_proj]]
|
|
self._down_projs_ptrs = [[tensor.data_ptr() for tensor in weights.down_proj]]
|
|
self._gate_scale_ptrs = [[tensor.data_ptr() for tensor in weights.gate_scale]]
|
|
self._up_scale_ptrs = [[tensor.data_ptr() for tensor in weights.up_scale]]
|
|
self._down_scale_ptrs = [[tensor.data_ptr() for tensor in weights.down_scale]]
|
|
self._has_bwd_projs = False
|
|
self._use_projs_path = True
|
|
self.gate_proj = None
|
|
self.up_proj = None
|
|
self.down_proj = None
|
|
self.load_weights(physical_to_logical_map_cpu)
|
|
|
|
def _load_base_weights_from_file(self) -> None:
|
|
if not hasattr(self, "weight_path") or self.weight_path is None:
|
|
raise RuntimeError(
|
|
"weight_path not set. Cannot load weights from file. "
|
|
"Either set weight_path or call load_weights_from_tensors() instead."
|
|
)
|
|
|
|
kt_layer_dir = os.path.join(self.weight_path, f"_layer_{self.layer_idx}")
|
|
if os.path.isdir(kt_layer_dir):
|
|
kt_files = _glob.glob(os.path.join(kt_layer_dir, "_numa_0", "*.kt"))
|
|
if kt_files:
|
|
self._use_kt_direct_load = True
|
|
return
|
|
if is_fp8_sft_method(self.method):
|
|
raise RuntimeError(
|
|
"AMXFP8_SFT requires raw per-expert HF checkpoint tensors; "
|
|
"legacy merged weight files are not supported"
|
|
)
|
|
if is_rawint4_sft_method(self.method):
|
|
raise RuntimeError(
|
|
"RAWINT4 SFT requires packed compressed-tensors checkpoint tensors; "
|
|
"legacy merged weight files are not supported"
|
|
)
|
|
if is_int8_sft_method(self.method):
|
|
raise RuntimeError(
|
|
f"INT8_SFT requires pre-quantized .kt files under "
|
|
f"{kt_layer_dir}/_numa_*/; merged safetensors and online "
|
|
"conversion are not supported"
|
|
)
|
|
|
|
if "BF16" in self.method:
|
|
loader = BF16SafeTensorLoader(self.weight_path)
|
|
base_key = f"model.layers.{self.layer_idx}"
|
|
else:
|
|
loader = SafeTensorLoader(self.weight_path)
|
|
base_key = f"blk.{self.layer_idx}"
|
|
|
|
experts_data = loader.load_experts(base_key, device="cpu")
|
|
|
|
gate_weights: List[torch.Tensor] = experts_data["gate"]
|
|
up_weights: List[torch.Tensor] = experts_data["up"]
|
|
down_weights: List[torch.Tensor] = experts_data["down"]
|
|
|
|
if "BF16" in self.method:
|
|
self.gate_proj = torch.stack(gate_weights, dim=0).contiguous()
|
|
self.up_proj = torch.stack(up_weights, dim=0).contiguous()
|
|
self.down_proj = torch.stack(down_weights, dim=0).contiguous()
|
|
else:
|
|
|
|
def _make_ptrs(arrays_per_numa):
|
|
return [
|
|
[
|
|
ctypes.addressof(ctypes.cast(et.ctypes.data, ctypes.POINTER(ctypes.c_uint64)).contents)
|
|
for et in numa_array
|
|
]
|
|
for numa_array in arrays_per_numa
|
|
]
|
|
|
|
self._gate_weights_per_numa = gate_weights
|
|
self._up_weights_per_numa = up_weights
|
|
self._down_weights_per_numa = down_weights
|
|
self._gate_scales_per_numa = experts_data["gate_scale"]
|
|
self._up_scales_per_numa = experts_data["up_scale"]
|
|
self._down_scales_per_numa = experts_data["down_scale"]
|
|
self._validate_prepartitioned_weights()
|
|
|
|
self._gate_projs_ptrs = _make_ptrs(gate_weights)
|
|
self._up_projs_ptrs = _make_ptrs(up_weights)
|
|
self._down_projs_ptrs = _make_ptrs(down_weights)
|
|
self._gate_scale_ptrs = _make_ptrs(experts_data["gate_scale"])
|
|
self._up_scale_ptrs = _make_ptrs(experts_data["up_scale"])
|
|
self._down_scale_ptrs = _make_ptrs(experts_data["down_scale"])
|
|
|
|
if "gate_bwd" in experts_data:
|
|
self._gate_bwd_weights_per_numa = experts_data["gate_bwd"]
|
|
self._up_bwd_weights_per_numa = experts_data["up_bwd"]
|
|
self._down_bwd_weights_per_numa = experts_data["down_bwd"]
|
|
self._gate_bwd_scales_per_numa = experts_data["gate_bwd_scale"]
|
|
self._up_bwd_scales_per_numa = experts_data["up_bwd_scale"]
|
|
self._down_bwd_scales_per_numa = experts_data["down_bwd_scale"]
|
|
|
|
self._gate_bwd_projs_ptrs = _make_ptrs(experts_data["gate_bwd"])
|
|
self._up_bwd_projs_ptrs = _make_ptrs(experts_data["up_bwd"])
|
|
self._down_bwd_projs_ptrs = _make_ptrs(experts_data["down_bwd"])
|
|
self._gate_bwd_scale_ptrs = _make_ptrs(experts_data["gate_bwd_scale"])
|
|
self._up_bwd_scale_ptrs = _make_ptrs(experts_data["up_bwd_scale"])
|
|
self._down_bwd_scale_ptrs = _make_ptrs(experts_data["down_bwd_scale"])
|
|
self._has_bwd_projs = True
|
|
else:
|
|
self._has_bwd_projs = False
|
|
|
|
self.gate_proj = None
|
|
self.up_proj = None
|
|
self.down_proj = None
|
|
self._use_projs_path = True
|
|
|
|
loader.close_all_handles()
|
|
|
|
def _aligned_max_len(self) -> int:
|
|
return ((self.chunked_prefill_size + _AMX_M_STEP - 1) // _AMX_M_STEP) * _AMX_M_STEP
|
|
|
|
def _validate_prepartitioned_weights(self) -> None:
|
|
numa_count = len(self._gate_weights_per_numa)
|
|
if self.moe_intermediate_size % self.threadpool_count != 0:
|
|
raise ValueError(
|
|
f"moe_intermediate_size={self.moe_intermediate_size} must be divisible by "
|
|
f"threadpool_count={self.threadpool_count} for {self.method} SFT."
|
|
)
|
|
if numa_count != self.threadpool_count:
|
|
raise ValueError(
|
|
f"{self.method} SFT pre-partitioned expert weights have {numa_count} NUMA partitions, "
|
|
f"but CPUInfer was created with threadpool_count={self.threadpool_count}. "
|
|
f"Use --kt-threadpool-count {numa_count} for this weight directory, or convert weights "
|
|
"for the requested threadpool count."
|
|
)
|
|
|
|
collections = {
|
|
"gate": self._gate_weights_per_numa,
|
|
"up": self._up_weights_per_numa,
|
|
"down": self._down_weights_per_numa,
|
|
"gate_scale": self._gate_scales_per_numa,
|
|
"up_scale": self._up_scales_per_numa,
|
|
"down_scale": self._down_scales_per_numa,
|
|
}
|
|
for name, per_numa in collections.items():
|
|
if len(per_numa) != numa_count:
|
|
raise ValueError(f"{name} has {len(per_numa)} NUMA partitions, expected {numa_count}.")
|
|
for numa_id, entries in enumerate(per_numa):
|
|
if len(entries) != self.num_experts:
|
|
raise ValueError(
|
|
f"{name}[numa={numa_id}] has {len(entries)} experts, expected {self.num_experts}."
|
|
)
|
|
|
|
for numa_id in range(numa_count):
|
|
gate_scale_len = self._gate_scales_per_numa[numa_id][0].size
|
|
up_scale_len = self._up_scales_per_numa[numa_id][0].size
|
|
down_scale_len = self._down_scales_per_numa[numa_id][0].size
|
|
expected_intermediate = self.moe_intermediate_size // self.threadpool_count
|
|
if gate_scale_len != expected_intermediate or up_scale_len != expected_intermediate:
|
|
raise ValueError(
|
|
f"{self.method} gate/up scale length for NUMA {numa_id} is "
|
|
f"{gate_scale_len}/{up_scale_len}, expected {expected_intermediate}."
|
|
)
|
|
if down_scale_len != self.hidden_size:
|
|
raise ValueError(
|
|
f"{self.method} down scale length for NUMA {numa_id} is "
|
|
f"{down_scale_len}, expected {self.hidden_size}."
|
|
)
|
|
|
|
# ========== LoRA ==========
|
|
|
|
def init_lora_weights(
|
|
self,
|
|
gate_lora_a: torch.Tensor,
|
|
gate_lora_b: torch.Tensor,
|
|
up_lora_a: torch.Tensor,
|
|
up_lora_b: torch.Tensor,
|
|
down_lora_a: torch.Tensor,
|
|
down_lora_b: torch.Tensor,
|
|
grad_gate_lora_a: torch.Tensor,
|
|
grad_gate_lora_b: torch.Tensor,
|
|
grad_up_lora_a: torch.Tensor,
|
|
grad_up_lora_b: torch.Tensor,
|
|
grad_down_lora_a: torch.Tensor,
|
|
grad_down_lora_b: torch.Tensor,
|
|
) -> None:
|
|
expected_shapes = {
|
|
"gate_lora_a": (self.num_experts, self.lora_rank, self.hidden_size),
|
|
"gate_lora_b": (self.num_experts, self.moe_intermediate_size, self.lora_rank),
|
|
"up_lora_a": (self.num_experts, self.lora_rank, self.hidden_size),
|
|
"up_lora_b": (self.num_experts, self.moe_intermediate_size, self.lora_rank),
|
|
"down_lora_a": (self.num_experts, self.lora_rank, self.moe_intermediate_size),
|
|
"down_lora_b": (self.num_experts, self.hidden_size, self.lora_rank),
|
|
}
|
|
provided = {
|
|
"gate_lora_a": gate_lora_a,
|
|
"gate_lora_b": gate_lora_b,
|
|
"up_lora_a": up_lora_a,
|
|
"up_lora_b": up_lora_b,
|
|
"down_lora_a": down_lora_a,
|
|
"down_lora_b": down_lora_b,
|
|
}
|
|
for name, tensor in provided.items():
|
|
expected = expected_shapes[name]
|
|
if tensor.shape != expected:
|
|
raise ValueError(f"{name} shape mismatch: expected {expected}, got {tuple(tensor.shape)}")
|
|
if tensor.device.type != "cpu":
|
|
raise ValueError(
|
|
f"{name} must be a CPU tensor for {self.method} SFT, got {tensor.device}."
|
|
)
|
|
if tensor.dtype != torch.bfloat16:
|
|
raise ValueError(
|
|
f"{name} must use torch.bfloat16 for {self.method} SFT, got {tensor.dtype}."
|
|
)
|
|
if not tensor.is_contiguous():
|
|
raise ValueError(
|
|
f"{name} must be contiguous for stable {self.method} pointer registration."
|
|
)
|
|
|
|
grad_provided = {
|
|
"grad_gate_lora_a": grad_gate_lora_a,
|
|
"grad_gate_lora_b": grad_gate_lora_b,
|
|
"grad_up_lora_a": grad_up_lora_a,
|
|
"grad_up_lora_b": grad_up_lora_b,
|
|
"grad_down_lora_a": grad_down_lora_a,
|
|
"grad_down_lora_b": grad_down_lora_b,
|
|
}
|
|
for name, tensor in grad_provided.items():
|
|
expected = expected_shapes[name.removeprefix("grad_")]
|
|
if tensor.shape != expected:
|
|
raise ValueError(
|
|
f"{name} shape mismatch: expected {expected}, got {tuple(tensor.shape)}"
|
|
)
|
|
if tensor.device.type != "cpu" or tensor.dtype != torch.bfloat16:
|
|
raise ValueError(
|
|
f"{name} must be a CPU torch.bfloat16 tensor for {self.method} SFT, "
|
|
f"got {tensor.dtype} on {tensor.device}."
|
|
)
|
|
if not tensor.is_contiguous():
|
|
raise ValueError(
|
|
f"{name} must be contiguous for stable {self.method} pointer registration."
|
|
)
|
|
|
|
self.gate_lora_a = gate_lora_a
|
|
self.gate_lora_b = gate_lora_b
|
|
self.up_lora_a = up_lora_a
|
|
self.up_lora_b = up_lora_b
|
|
self.down_lora_a = down_lora_a
|
|
self.down_lora_b = down_lora_b
|
|
|
|
self.grad_gate_lora_a = grad_gate_lora_a
|
|
self.grad_gate_lora_b = grad_gate_lora_b
|
|
self.grad_up_lora_a = grad_up_lora_a
|
|
self.grad_up_lora_b = grad_up_lora_b
|
|
self.grad_down_lora_a = grad_down_lora_a
|
|
self.grad_down_lora_b = grad_down_lora_b
|
|
|
|
self._lora_initialized = True
|
|
|
|
if self._weights_loaded and self.moe is not None:
|
|
self.update_lora_weights()
|
|
|
|
def update_lora_weights(self) -> None:
|
|
if not self._weights_loaded:
|
|
raise RuntimeError("Weights not loaded. Call load_weights() first.")
|
|
if self._is_skip_lora:
|
|
return
|
|
if not self._lora_initialized:
|
|
if self.lora_rank <= 0:
|
|
return # Full mode without LoRA — no LoRA weights to update
|
|
raise RuntimeError("LoRA weights not initialized. Call init_lora_weights() first.")
|
|
|
|
# Weight pointer updates are load-time synchronous work. Calling the
|
|
# direct binding avoids nesting an update task inside CPUInfer's queue
|
|
# while SGLang is still in distributed model-loading barriers.
|
|
self.moe.update_lora_weights(
|
|
self.gate_lora_a.data_ptr(),
|
|
self.gate_lora_b.data_ptr(),
|
|
self.up_lora_a.data_ptr(),
|
|
self.up_lora_b.data_ptr(),
|
|
self.down_lora_a.data_ptr(),
|
|
self.down_lora_b.data_ptr(),
|
|
)
|
|
|
|
def update_base_weights(self) -> None:
|
|
"""Sync updated base weight parameters back to C++ kernel after optimizer step."""
|
|
if not self._weights_loaded:
|
|
raise RuntimeError("Weights not loaded. Call load_weights() first.")
|
|
if not self._full_weight_grad:
|
|
return # No base weights to update in LoRA mode
|
|
if self.gate_proj_buf is None:
|
|
raise RuntimeError("Base weight buffers not initialized. Call init_full_weight_grad_buffers() first.")
|
|
|
|
logger.info(f"Layer {self.layer_idx}: update_base_weights() - syncing updated weights to C++ kernel")
|
|
|
|
# Preferred path: update config pointers on existing C++ object and re-quantize.
|
|
# This avoids full C++ MOE object recreation (~0.6s/layer vs ~1.9s/layer).
|
|
if hasattr(self.moe, "set_base_weight_pointers"):
|
|
self.moe.set_base_weight_pointers(
|
|
self.gate_proj_buf.data.data_ptr(),
|
|
self.up_proj_buf.data.data_ptr(),
|
|
self.down_proj_buf.data.data_ptr(),
|
|
)
|
|
self.cpu_infer.submit(self.moe.load_weights_task())
|
|
self.cpu_infer.sync()
|
|
logger.info(f"Layer {self.layer_idx}: update_base_weights() - re-quantized existing kernel")
|
|
return
|
|
|
|
# Fallback: full reload path (creates new C++ MOE object)
|
|
# This is slower but works without C++ set_base_weight_pointers support.
|
|
logger.warning(
|
|
f"Layer {self.layer_idx}: set_base_weight_pointers not available, "
|
|
f"falling back to full C++ MOE object recreation"
|
|
)
|
|
old_moe = getattr(self, "moe", None)
|
|
if old_moe is not None:
|
|
del old_moe
|
|
|
|
self.gate_proj = self.gate_proj_buf.data
|
|
self.up_proj = self.up_proj_buf.data
|
|
self.down_proj = self.down_proj_buf.data
|
|
physical_to_logical_map = torch.arange(self.num_experts, dtype=torch.int64, device="cpu")
|
|
self._weights_loaded = False # Allow re-load
|
|
self.load_weights_from_tensors(
|
|
gate_proj=self.gate_proj,
|
|
up_proj=self.up_proj,
|
|
down_proj=self.down_proj,
|
|
physical_to_logical_map_cpu=physical_to_logical_map,
|
|
)
|
|
|
|
def save_backward_weights_from_tensors(
|
|
self,
|
|
gate_proj: torch.Tensor,
|
|
up_proj: torch.Tensor,
|
|
down_proj: torch.Tensor,
|
|
physical_to_logical_map: torch.Tensor,
|
|
output_path: str,
|
|
) -> None:
|
|
if not self._weights_loaded:
|
|
raise RuntimeError("Weights not loaded. Call load_weights() first.")
|
|
gate_proj = gate_proj.contiguous()
|
|
up_proj = up_proj.contiguous()
|
|
down_proj = down_proj.contiguous()
|
|
self.moe.prepare_and_save_bwd(
|
|
gate_proj.data_ptr(),
|
|
up_proj.data_ptr(),
|
|
down_proj.data_ptr(),
|
|
output_path,
|
|
)
|