kvcache-ai/ktransformers
1
0
Fork 0
mirror of https://github.com/kvcache-ai/ktransformers.git synced 2026-07-16 08:43:45 +00:00
Code Issues Packages Projects Releases Wiki Activity

[feat](kt-lora): add end-to-end Qwen3.5 MoE KT LoRA serving workflow (#2031)

Browse Source 
* [feat](kt-lora): add KT expert LoRA adapter serving

* [feat]: pin Qwen3.5 non-expert LoRA support

* [feat](kt-lora): add merged SGLang adapter workflow

Document the KT SFT to SGLang serving loop and extend the converter with optional split outputs so users can serve one merged adapter while retaining debug-friendly expert/non-expert artifacts.

Co-authored-by: Cursor <cursoragent@cursor.com>

* [fix](kt-lora): validate adapter conversion

Co-authored-by: Cursor <cursoragent@cursor.com>

---------

Co-authored-by: Cursor <cursoragent@cursor.com>
This commit is contained in:
Jiaheng Dai
2026-06-05 16:57:14 +08:00
committed by GitHub
parent d41f569e84
commit c9a915e6ac
16 changed files with 2435 additions and 656 deletions
Download Patch File Download Diff File Expand all files Collapse all files

97
kt-kernel/python/sft/amx.py
View File

@@ -40,6 +40,8 @@ except (ImportError, AttributeError):
from .base import BaseSFTMoEWrapper, KExpertsSFTBuffer
_AMX_M_STEP = 32
# Mapping from method string to C++ SFT MOE class
_SFT_METHOD_TO_CLASS = {
@@ -159,6 +161,17 @@ class AMXSFTMoEWrapper(BaseSFTMoEWrapper):
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)
self._physical_to_logical_map_cpu = physical_to_logical_map_cpu.to(
dtype=torch.int64, device="cpu"
).contiguous()
if 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()}."
)
if self.gate_proj is None and not getattr(self, "_use_projs_path", False):
self._load_base_weights_from_file()
@@ -170,10 +183,11 @@ class AMXSFTMoEWrapper(BaseSFTMoEWrapper):
config.lora_rank = self.lora_rank
config.lora_alpha = self.lora_alpha
config.max_cache_depth = self.max_cache_depth
config.max_len = self.chunked_prefill_size
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.physical_to_logical_map = self._physical_to_logical_map_cpu.data_ptr()
if getattr(self, "_use_kt_direct_load", False):
config.load = True
@@ -220,8 +234,9 @@ class AMXSFTMoEWrapper(BaseSFTMoEWrapper):
self.cpu_infer.submit(self.moe.load_weights_task())
self.cpu_infer.sync()
self.cpu_infer.submit(self.moe.warm_up_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)
self.gate_proj = None
@@ -312,6 +327,7 @@ class AMXSFTMoEWrapper(BaseSFTMoEWrapper):
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)
@@ -345,6 +361,57 @@ class AMXSFTMoEWrapper(BaseSFTMoEWrapper):
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(
@@ -373,6 +440,10 @@ class AMXSFTMoEWrapper(BaseSFTMoEWrapper):
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}."
)
self.gate_lora_a = gate_lora_a.contiguous()
self.gate_lora_b = gate_lora_b.contiguous()
@@ -401,17 +472,17 @@ class AMXSFTMoEWrapper(BaseSFTMoEWrapper):
if not self._lora_initialized:
raise RuntimeError("LoRA weights not initialized. Call init_lora_weights() first.")
self.cpu_infer.submit(
self.moe.update_lora_weights_task(
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(),
)
# 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(),
)
self.cpu_infer.sync()
def save_backward_weights_from_tensors(
self,

120
kt-kernel/python/sft/base.py
View File

@@ -15,7 +15,7 @@ import torch
from typing import Optional, Tuple
from abc import ABC, abstractmethod
from ..experts_base import _MoEBase
from ..experts_base import KExpertsCPUBuffer, _MoEBase
class KExpertsSFTBuffer:
@@ -98,6 +98,26 @@ class KExpertsSFTBuffer:
cls._shared_buffer = None
class _SFTForwardBufferView:
"""Minimal buffer view consumed by AMXSFTMoEWrapper._make_forward_task."""
__slots__ = ("bsz_tensor", "expert_ids_cpu", "weights_cpu", "input_cpu", "output_cpu")
def __init__(
self,
bsz_tensor: torch.Tensor,
expert_ids_cpu: torch.Tensor,
weights_cpu: torch.Tensor,
input_cpu: torch.Tensor,
output_cpu: torch.Tensor,
):
self.bsz_tensor = bsz_tensor
self.expert_ids_cpu = expert_ids_cpu
self.weights_cpu = weights_cpu
self.input_cpu = input_cpu
self.output_cpu = output_cpu
class BaseSFTMoEWrapper(_MoEBase, ABC):
"""
Base class for SFT MoE CPU operations with concrete buffer management.
@@ -357,6 +377,104 @@ class BaseSFTMoEWrapper(_MoEBase, ABC):
return self._return_output(buffer, qlen, output_device)
# ========== Inference-only async forward ==========
def submit_forward_inference(
self,
hidden_states: torch.Tensor,
expert_ids: torch.Tensor,
weights: torch.Tensor,
cuda_stream,
) -> None:
"""
Submit an SFT MoE forward pass for serving.
This path mirrors the normal KT inference wrapper: inputs are copied to
pinned CPU staging buffers, the CPUInfer task is enqueued with the
caller CUDA stream, and sync_forward_inference() returns a persistent
GPU output buffer. It deliberately avoids the training-oriented
torch.cuda.synchronize() in _copy_inputs_to_buffer().
"""
if not hasattr(self.cpu_infer, "submit_with_cuda_stream"):
self.submit_forward(hidden_states, expert_ids, weights, save_for_backward=False)
self._pending_inference_fallback = True
self._pending_inference_fallback_device = hidden_states.device
return
self._validate_forward_inputs(hidden_states, expert_ids, weights)
flat_hidden_states = hidden_states.view(-1, hidden_states.shape[-1])
(
input_tensor_cpu,
expert_ids_cpu,
_deferred_expert_ids_cpu,
weights_cpu,
output_cpu,
bsz_tensor_cpu,
output_gpu,
) = KExpertsCPUBuffer.get_buffer(flat_hidden_states, self.num_experts_per_tok)
current_slot = self.layer_idx % KExpertsCPUBuffer.buffer_depth
bsz_slot_tensor = bsz_tensor_cpu[current_slot]
torch_stream = (
cuda_stream
if isinstance(cuda_stream, torch.cuda.Stream)
else torch.cuda.ExternalStream(cuda_stream, device=flat_hidden_states.device)
)
with torch.cuda.stream(torch_stream):
input_tensor_cpu[current_slot].copy_(flat_hidden_states.to(torch.bfloat16), non_blocking=True)
expert_ids_cpu[current_slot].copy_(expert_ids.to(torch.int64), non_blocking=True)
weights_cpu[current_slot].copy_(weights.to(torch.float32), non_blocking=True)
buffer_view = _SFTForwardBufferView(
bsz_tensor=bsz_slot_tensor,
expert_ids_cpu=expert_ids_cpu[current_slot],
weights_cpu=weights_cpu[current_slot],
input_cpu=input_tensor_cpu[current_slot],
output_cpu=output_cpu[current_slot],
)
self._pending_inference_fallback = False
self._pending_inference_output_cpu = output_cpu[current_slot]
self._pending_inference_output_gpu = output_gpu[current_slot]
self.cpu_infer.submit_with_cuda_stream(
cuda_stream,
self._make_forward_task(buffer_view, save_for_backward=False),
)
def sync_forward_inference(self, cuda_stream) -> torch.Tensor:
"""
Synchronize a serving forward submitted by submit_forward_inference().
Returns a persistent GPU buffer matching the input batch shape. Consumers
on the same CUDA stream will naturally wait for the non-blocking D2H/H2D
staging work ordered through CPUInfer's stream synchronization.
"""
if getattr(self, "_pending_inference_fallback", False):
self._pending_inference_fallback = False
output_device = getattr(self, "_pending_inference_fallback_device", None)
self._pending_inference_fallback_device = None
return self.sync_forward(output_device=output_device)
if not hasattr(self, "_pending_inference_output_cpu"):
raise RuntimeError("No pending inference forward. Call submit_forward_inference() first.")
torch_stream = (
cuda_stream
if isinstance(cuda_stream, torch.cuda.Stream)
else torch.cuda.ExternalStream(cuda_stream, device=self._pending_inference_output_gpu.device)
)
self.cpu_infer.sync_with_cuda_stream(cuda_stream)
with torch.cuda.stream(torch_stream):
self._pending_inference_output_gpu.copy_(self._pending_inference_output_cpu, non_blocking=True)
output = self._pending_inference_output_gpu
del self._pending_inference_output_cpu
del self._pending_inference_output_gpu
return output
# ========== Async backward ==========
def submit_backward_async(

1
kt-kernel/python/utils/loader.py
View File

@@ -726,6 +726,7 @@ class CompressedSafeTensorLoader(SafeTensorLoader):
"down_scale": down_scales,
}
class GGUFLoader:
"""
GGUF format loader using the official gguf library (gguf.gguf_reader.GGUFReader)