diff --git a/kt-kernel/ext_bindings.cpp b/kt-kernel/ext_bindings.cpp index 87074e417..38763b15d 100644 --- a/kt-kernel/ext_bindings.cpp +++ b/kt-kernel/ext_bindings.cpp @@ -709,6 +709,7 @@ PYBIND11_MODULE(kt_kernel_ext, m) { "gpu_experts_mask", [](const GeneralMOEConfig& self) { return reinterpret_cast(self.gpu_experts_mask); }, [](GeneralMOEConfig& self, uintptr_t val) { self.gpu_experts_mask = reinterpret_cast(val); }) + .def_readwrite("skip_gpu_expert_cpu_copy", &GeneralMOEConfig::skip_gpu_expert_cpu_copy) .DEF_PTR_PROPERTY(GeneralMOEConfig, physical_to_logical_map) .DEF_PTR_PROPERTY(GeneralMOEConfig, gate_proj) diff --git a/kt-kernel/operators/amx/awq-moe.hpp b/kt-kernel/operators/amx/awq-moe.hpp index 52dfbf395..033a06e48 100644 --- a/kt-kernel/operators/amx/awq-moe.hpp +++ b/kt-kernel/operators/amx/awq-moe.hpp @@ -514,6 +514,7 @@ class AMX_AWQ_MOE_TP : public AMX_MOE_BASE> { nth * config_.expert_num, nullptr, [this, nth, physical_to_logical_map](int task_id) { uint64_t expert_idx = task_id / nth; + if (config_.skip_gpu_expert_cpu_copy && config_.should_skip_expert((int64_t)expert_idx)) return; uint64_t logical_expert_id = expert_map(physical_to_logical_map, expert_idx); int ith = task_id % nth; // gate part @@ -534,6 +535,7 @@ class AMX_AWQ_MOE_TP : public AMX_MOE_BASE> { nth * config_.expert_num, nullptr, [this, nth, physical_to_logical_map](int task_id) { uint64_t expert_idx = task_id / nth; + if (config_.skip_gpu_expert_cpu_copy && config_.should_skip_expert((int64_t)expert_idx)) return; uint64_t logical_expert_id = expert_map(physical_to_logical_map, expert_idx); int ith = task_id % nth; // down part @@ -548,6 +550,7 @@ class AMX_AWQ_MOE_TP : public AMX_MOE_BASE> { config_.expert_num, nullptr, [this, physical_to_logical_map](int task_id) { uint64_t expert_idx = task_id; + if (config_.skip_gpu_expert_cpu_copy && config_.should_skip_expert((int64_t)expert_idx)) return; uint64_t logical_expert_id = expert_map(physical_to_logical_map, expert_idx); size_t scale_elem_count = (config_.hidden_size * config_.intermediate_size) / config_.quant_config.group_size; @@ -579,6 +582,9 @@ class AMX_AWQ_MOE_TP : public AMX_MOE_BASE> { // Save offline quantization data if requested if (config_.save) { for (int expert_idx = 0; expert_idx < config_.expert_num; expert_idx++) { + if (config_.skip_gpu_expert_cpu_copy && config_.should_skip_expert(expert_idx)) { + continue; + } write_weights(prefix, "_gate_", gate_bb_[expert_idx].get(), expert_idx, "OFFLINE"); write_weights(prefix, "_up_", up_bb_[expert_idx].get(), expert_idx, "OFFLINE"); write_weights(prefix, "_down_", down_bb_[expert_idx].get(), expert_idx, "OFFLINE"); @@ -594,6 +600,7 @@ class AMX_AWQ_MOE_TP : public AMX_MOE_BASE> { [this, nth, physical_to_logical_map](int task_id) { int64_t expert_idx = task_id / nth; uint64_t logical_expert_id = expert_map(physical_to_logical_map, expert_idx); + if (config_.skip_gpu_expert_cpu_copy && config_.should_skip_expert((int64_t)logical_expert_id)) return; int ith = task_id % nth; // gate part gate_bb_[logical_expert_id]->from_mat( @@ -613,6 +620,7 @@ class AMX_AWQ_MOE_TP : public AMX_MOE_BASE> { [this, nth, physical_to_logical_map](int task_id) { int64_t expert_idx = task_id / nth; uint64_t logical_expert_id = expert_map(physical_to_logical_map, expert_idx); + if (config_.skip_gpu_expert_cpu_copy && config_.should_skip_expert((int64_t)logical_expert_id)) return; int ith = task_id % nth; // down part down_bb_[logical_expert_id]->from_mat( @@ -625,6 +633,9 @@ class AMX_AWQ_MOE_TP : public AMX_MOE_BASE> { // Save online quantization data if requested if (config_.save) { for (int expert_idx = 0; expert_idx < config_.expert_num; expert_idx++) { + if (config_.skip_gpu_expert_cpu_copy && config_.should_skip_expert(expert_idx)) { + continue; + } write_weights(prefix, "_gate_", gate_bb_[expert_idx].get(), expert_idx, "ONLINE"); write_weights(prefix, "_up_", up_bb_[expert_idx].get(), expert_idx, "ONLINE"); write_weights(prefix, "_down_", down_bb_[expert_idx].get(), expert_idx, "ONLINE"); diff --git a/kt-kernel/operators/amx/fp4-moe.hpp b/kt-kernel/operators/amx/fp4-moe.hpp index b35698847..594ff0d23 100644 --- a/kt-kernel/operators/amx/fp4-moe.hpp +++ b/kt-kernel/operators/amx/fp4-moe.hpp @@ -442,6 +442,7 @@ class AMX_FP4_MOE_TP : public AMX_MOE_BASE> { nth * config_.expert_num, nullptr, [this, nth, physical_to_logical_map](int task_id) { uint64_t expert_idx = task_id / nth; + if (config_.skip_gpu_expert_cpu_copy && config_.should_skip_expert((int64_t)expert_idx)) return; uint64_t logical_expert_id = expert_map(physical_to_logical_map, expert_idx); int ith = task_id % nth; gate_bb_[expert_idx]->from_raw_mat( @@ -459,6 +460,7 @@ class AMX_FP4_MOE_TP : public AMX_MOE_BASE> { nth * config_.expert_num, nullptr, [this, nth, physical_to_logical_map](int task_id) { uint64_t expert_idx = task_id / nth; + if (config_.skip_gpu_expert_cpu_copy && config_.should_skip_expert((int64_t)expert_idx)) return; uint64_t logical_expert_id = expert_map(physical_to_logical_map, expert_idx); int ith = task_id % nth; down_bb_[expert_idx]->from_raw_mat( @@ -472,6 +474,7 @@ class AMX_FP4_MOE_TP : public AMX_MOE_BASE> { config_.expert_num, nullptr, [this, physical_to_logical_map](int task_id) { uint64_t expert_idx = task_id; + if (config_.skip_gpu_expert_cpu_copy && config_.should_skip_expert((int64_t)expert_idx)) return; uint64_t logical_expert_id = expert_map(physical_to_logical_map, expert_idx); size_t scale_elem_count = (config_.hidden_size * config_.intermediate_size) / config_.quant_config.group_size; convert_or_copy(gate_bb_[expert_idx]->d, diff --git a/kt-kernel/operators/amx/moe.hpp b/kt-kernel/operators/amx/moe.hpp index 4a8450df0..d259554ad 100644 --- a/kt-kernel/operators/amx/moe.hpp +++ b/kt-kernel/operators/amx/moe.hpp @@ -211,6 +211,7 @@ class AMX_MOE_TP : public AMX_MOE_BASE> { config_.expert_num, nullptr, [this, physical_to_logical_map](int expert_id) { // printf("Load layer %d [%d/%d]\n", config_.layer_idx, expert_id, config_.expert_num); + if (config_.skip_gpu_expert_cpu_copy && config_.should_skip_expert((int64_t)expert_id)) return; uint64_t logical_expert_id = expert_map(physical_to_logical_map, expert_id); { size_t scale_size = config_.intermediate_size * sizeof(float); @@ -254,6 +255,9 @@ class AMX_MOE_TP : public AMX_MOE_BASE> { config_.expert_num * mat_type_all * mat_split, [this, physical_to_logical_map, prefix, mat_type_all, mat_split](int task_id) { int64_t expert_idx = task_id / (mat_type_all * mat_split); + if (config_.skip_gpu_expert_cpu_copy && config_.should_skip_expert(expert_idx)) { + return; + } uint64_t logical_expert_id = expert_map(physical_to_logical_map, expert_idx); uint8_t mat_class = (task_id % (mat_type_all * mat_split)) / mat_split; uint8_t mat_split_idex = task_id % mat_split; @@ -291,6 +295,7 @@ class AMX_MOE_TP : public AMX_MOE_BASE> { [this, nth, physical_to_logical_map](int task_id) { int64_t expert_idx = task_id / nth; uint64_t logical_expert_id = expert_map(physical_to_logical_map, expert_idx); + if (config_.skip_gpu_expert_cpu_copy && config_.should_skip_expert((int64_t)logical_expert_id)) return; int ith = task_id % nth; // gate part gate_bb_[logical_expert_id]->from_mat( @@ -309,6 +314,7 @@ class AMX_MOE_TP : public AMX_MOE_BASE> { [this, nth, physical_to_logical_map](int task_id) { int64_t expert_idx = task_id / nth; uint64_t logical_expert_id = expert_map(physical_to_logical_map, expert_idx); + if (config_.skip_gpu_expert_cpu_copy && config_.should_skip_expert((int64_t)logical_expert_id)) return; int ith = task_id % nth; // down part down_bb_[logical_expert_id]->from_mat( diff --git a/kt-kernel/operators/amx/moe_base.hpp b/kt-kernel/operators/amx/moe_base.hpp index 48c4da9ee..afd9e2c74 100644 --- a/kt-kernel/operators/amx/moe_base.hpp +++ b/kt-kernel/operators/amx/moe_base.hpp @@ -117,6 +117,13 @@ class AMX_MOE_BASE { down_ba_.push_back(make_buffer_a(config_.max_len, config_.intermediate_size, nullptr)); down_bc_.push_back(make_buffer_c(config_.max_len, config_.hidden_size, nullptr)); + if (config_.skip_gpu_expert_cpu_copy && config_.should_skip_expert((int64_t)i)) { + gate_bb_.push_back(nullptr); + up_bb_.push_back(nullptr); + down_bb_.push_back(nullptr); + continue; + } + void* gate_bb_ptr = std::aligned_alloc(64, buffer_b_required_size(config_.intermediate_size, config_.hidden_size)); gate_bb_.push_back(make_buffer_b(config_.intermediate_size, config_.hidden_size, gate_bb_ptr)); diff --git a/kt-kernel/operators/amx/mxfp8-moe.hpp b/kt-kernel/operators/amx/mxfp8-moe.hpp index 3d571c514..dedd3d448 100644 --- a/kt-kernel/operators/amx/mxfp8-moe.hpp +++ b/kt-kernel/operators/amx/mxfp8-moe.hpp @@ -566,6 +566,7 @@ class AMX_MXFP8_MOE_TP : public AMX_MOE_BASE> { nth * config_.expert_num, nullptr, [this, nth, physical_to_logical_map](int task_id) { uint64_t expert_idx = task_id / nth; + if (config_.skip_gpu_expert_cpu_copy && config_.should_skip_expert((int64_t)expert_idx)) return; uint64_t logical_expert_id = expert_map(physical_to_logical_map, expert_idx); int ith = task_id % nth; // FP8: no >> 1 (1 byte/element, not nibble-packed) @@ -580,6 +581,7 @@ class AMX_MXFP8_MOE_TP : public AMX_MOE_BASE> { nth * config_.expert_num, nullptr, [this, nth, physical_to_logical_map](int task_id) { uint64_t expert_idx = task_id / nth; + if (config_.skip_gpu_expert_cpu_copy && config_.should_skip_expert((int64_t)expert_idx)) return; uint64_t logical_expert_id = expert_map(physical_to_logical_map, expert_idx); int ith = task_id % nth; size_t weight_offset = logical_expert_id * config_.hidden_size * config_.intermediate_size; @@ -593,6 +595,7 @@ class AMX_MXFP8_MOE_TP : public AMX_MOE_BASE> { config_.expert_num, nullptr, [this, physical_to_logical_map](int task_id) { uint64_t expert_idx = task_id; + if (config_.skip_gpu_expert_cpu_copy && config_.should_skip_expert((int64_t)expert_idx)) return; uint64_t logical_expert_id = expert_map(physical_to_logical_map, expert_idx); size_t scale_count = (static_cast(config_.intermediate_size) * config_.hidden_size) / config_.quant_config.group_size; diff --git a/kt-kernel/operators/common.hpp b/kt-kernel/operators/common.hpp index 86f3464c6..c2ae6711d 100644 --- a/kt-kernel/operators/common.hpp +++ b/kt-kernel/operators/common.hpp @@ -241,6 +241,7 @@ struct GeneralMOEConfig { int num_gpu_experts = 0; // Computed from gpu_experts_mask uint8_t* gpu_experts_mask = nullptr; // Bool mask: true = expert on GPU void* physical_to_logical_map = nullptr; + bool skip_gpu_expert_cpu_copy = false; // Compute num_gpu_experts from gpu_experts_mask void compute_num_gpu_experts() { diff --git a/kt-kernel/operators/moe-sft-tp.hpp b/kt-kernel/operators/moe-sft-tp.hpp index 4e8da6ad0..14c0c2fa0 100644 --- a/kt-kernel/operators/moe-sft-tp.hpp +++ b/kt-kernel/operators/moe-sft-tp.hpp @@ -215,9 +215,27 @@ class TP_MOE_SFT : public TP_MOE { std::vector part_grad_input_; std::vector part_grad_weights_; + // SFT trains every expert on CPU, so full host copies of all expert weights + // are mandatory (GPU-resident experts included). MOESFTConfig inherits + // skip_gpu_expert_cpu_copy from GeneralMOEConfig, so a mis-set flag would + // otherwise ride the sliced config into the TP_MOE base; strip it here so no + // current or future base-class loader can act on it in a training run. + static GeneralMOEConfig reject_skip_gpu_expert_cpu_copy(GeneralMOEConfig cfg) { + if (cfg.skip_gpu_expert_cpu_copy) { + printf( + "[TP_MOE_SFT] skip_gpu_expert_cpu_copy is not supported in SFT mode; " + "ignoring it and keeping full CPU copies of all expert weights (layer %d)\n", + cfg.layer_idx); + cfg.skip_gpu_expert_cpu_copy = false; + } + return cfg; + } + public: - TP_MOE_SFT(const MOESFTConfig& config) : Base(static_cast(config)), sft_config(config) { + TP_MOE_SFT(const MOESFTConfig& config) + : Base(reject_skip_gpu_expert_cpu_copy(static_cast(config))), sft_config(config) { printf("Creating TP_MOE_SFT layer %d\n", config.layer_idx); + sft_config.skip_gpu_expert_cpu_copy = false; backward_temp_pools_.assign(tp_count, nullptr); backward_temp_pool_bytes_.assign(tp_count, 0); diff --git a/kt-kernel/python/experts.py b/kt-kernel/python/experts.py index 318fe3170..c6c35ccf4 100644 --- a/kt-kernel/python/experts.py +++ b/kt-kernel/python/experts.py @@ -132,6 +132,7 @@ def __new__( # Inference-specific parameters cpu_save: bool = False, max_deferred_experts_per_token: Optional[int] = None, + skip_gpu_expert_cpu_copy: bool = False, # Mode and method selection method: str = "AMXINT4", numa_nodes: Optional[List[int]] = None, @@ -221,6 +222,7 @@ def __new__( chunked_prefill_size=chunked_prefill_size, cpu_save=cpu_save, max_deferred_experts_per_token=max_deferred_experts_per_token, + skip_gpu_expert_cpu_copy=skip_gpu_expert_cpu_copy, method=method, numa_nodes=numa_nodes, swiglu_limit=swiglu_limit, @@ -235,6 +237,15 @@ def __new__( f"mode='sft' (method={method!r}); SFT backends do not " f"implement the V4-2604B clamp." ) + # SFT trains every expert on CPU and needs full host copies of all + # expert weights, including GPU-resident ones; the inference-only + # CPU-copy skip must never reach a training run. + if skip_gpu_expert_cpu_copy: + raise ValueError( + "skip_gpu_expert_cpu_copy=True is not supported in " + "mode='sft'; SFT requires full CPU copies of all expert " + "weights (GPU-resident experts included)." + ) return _create_sft_wrapper( layer_idx=layer_idx, num_experts=num_experts, @@ -318,6 +329,7 @@ def _create_inference_wrapper( chunked_prefill_size: int, cpu_save: bool, max_deferred_experts_per_token: Optional[int], + skip_gpu_expert_cpu_copy: bool, method: str, numa_nodes: Optional[List[int]] = None, swiglu_limit: float = 0.0, @@ -364,6 +376,8 @@ def _create_inference_wrapper( f"environment while the current launch does not actually use " f"MXFP4/MXFP8 weights — either unset the env or pass --kt-method MXFP4/MXFP8." ) + if backend_cls is NativeMoEWrapper: + extra_kwargs["skip_gpu_expert_cpu_copy"] = skip_gpu_expert_cpu_copy return backend_cls( layer_idx=layer_idx, num_experts=num_experts, diff --git a/kt-kernel/python/utils/amx.py b/kt-kernel/python/utils/amx.py index 1b04e4def..90f0f3bd5 100644 --- a/kt-kernel/python/utils/amx.py +++ b/kt-kernel/python/utils/amx.py @@ -543,8 +543,10 @@ def __init__( numa_nodes: Optional[List[int]] = None, swiglu_limit: float = 0.0, swiglu_alpha: float = 0.0, + skip_gpu_expert_cpu_copy: bool = False, ): self._swiglu_alpha = float(swiglu_alpha) + self.skip_gpu_expert_cpu_copy = skip_gpu_expert_cpu_copy # Defence in depth: reject swiglu_limit on non-MXFP4/MXFP8 methods even # if the experts.py guard is bypassed (e.g., by a future caller # that constructs NativeMoEWrapper directly). Origin: kt-sglang è€Šćˆ. @@ -788,6 +790,7 @@ def load_weights(self, physical_to_logical_map_cpu: torch.Tensor): moe_config.layer_idx = self.layer_idx moe_config.pool = self.cpu_infer.backend_ moe_config.max_len = self.chunked_prefill_size + moe_config.skip_gpu_expert_cpu_copy = self.skip_gpu_expert_cpu_copy # V4-Flash 2604B SwiGLU clamp; 0.0 = disabled (default for non-MXFP4 # paths). Read by `act_fn` in operators/amx/la/amx.hpp via # `apply_activation` in operators/amx/moe_base.hpp. Re-checked here diff --git a/kt-kernel/test/per_commit/test_moe_amx_skip_gpu_expert_cpu_copy.py b/kt-kernel/test/per_commit/test_moe_amx_skip_gpu_expert_cpu_copy.py new file mode 100644 index 000000000..2221ade54 --- /dev/null +++ b/kt-kernel/test/per_commit/test_moe_amx_skip_gpu_expert_cpu_copy.py @@ -0,0 +1,217 @@ +#!/usr/bin/env python +# coding=utf-8 +"""Parity test for MOEConfig.skip_gpu_expert_cpu_copy. + +With skip_gpu_expert_cpu_copy set, the CPU weight buffers of GPU-resident experts +(those in gpu_experts_mask) are never allocated or loaded. This test runs a hybrid +split with the flag on and off and asserts the forward output is bitwise identical, +plus matches a torch reference that drops the masked experts. +""" + +import os +import sys +import pytest + +sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..")) +from ci.ci_register import register_cpu_ci + +register_cpu_ci(est_time=90, suite="default") + +try: + import torch + import kt_kernel + + kt_kernel_ext = kt_kernel.kt_kernel_ext + HAS_DEPS = True +except ImportError as e: + HAS_DEPS = False + import_error = str(e) + +expert_num = 64 +hidden_size = 2048 +intermediate_size = 1024 +max_len = 512 +num_experts_per_tok = 8 + + +def act_fn(x): + return x / (1.0 + torch.exp(-x)) + + +def mlp_torch(input, gate_proj, up_proj, down_proj): + gate_buf = torch.mm(input, gate_proj.t()) + up_buf = torch.mm(input, up_proj.t()) + intermediate = act_fn(gate_buf) * up_buf + return torch.mm(intermediate, down_proj.t()) + + +def moe_torch(input, expert_ids, weights, gate_proj, up_proj, down_proj, skip_mask=None): + """Reference MoE; experts flagged in skip_mask contribute zero (GPU-resident).""" + cnts = expert_ids.new_zeros((expert_ids.shape[0], expert_num)) + cnts.scatter_(1, expert_ids, 1) + tokens_per_expert = cnts.sum(dim=0) + idxs = expert_ids.view(-1).argsort() + sorted_tokens = input[idxs // expert_ids.shape[1]] + + outputs = [] + start_idx = 0 + for i, num_tokens in enumerate(tokens_per_expert): + end_idx = start_idx + num_tokens + if num_tokens == 0: + continue + tokens_for_this_expert = sorted_tokens[start_idx:end_idx] + if skip_mask is not None and bool(skip_mask[i]): + expert_out = torch.zeros_like(tokens_for_this_expert) + else: + expert_out = mlp_torch(tokens_for_this_expert, gate_proj[i], up_proj[i], down_proj[i]) + outputs.append(expert_out) + start_idx = end_idx + + outs = torch.cat(outputs, dim=0) if len(outputs) else sorted_tokens.new_empty(0) + new_x = torch.empty_like(outs) + new_x[idxs] = outs + return ( + new_x.view(*expert_ids.shape, -1) + .type(weights.dtype) + .mul_(weights.unsqueeze(dim=-1)) + .sum(dim=1) + .type(new_x.dtype) + ) + + +def _build_moe(CPUInfer, gate_proj, up_proj, down_proj, gpu_experts_mask, physical_to_logical_map, skip): + config = kt_kernel_ext.moe.MOEConfig(expert_num, num_experts_per_tok, hidden_size, intermediate_size) + config.max_len = max_len + config.gate_proj = gate_proj.data_ptr() + config.up_proj = up_proj.data_ptr() + config.down_proj = down_proj.data_ptr() + config.gate_scale = 0 + config.pool = CPUInfer.backend_ + config.gpu_experts_mask = gpu_experts_mask.data_ptr() + config.skip_gpu_expert_cpu_copy = skip + + moe = kt_kernel_ext.moe.AMXInt4_MOE(config) + CPUInfer.submit(moe.load_weights_task(physical_to_logical_map.data_ptr())) + CPUInfer.sync() + CPUInfer.submit(moe.warm_up_task()) + CPUInfer.sync() + return moe + + +def _forward(CPUInfer, moe, expert_ids, weights, input_data): + qlen = input_data.shape[0] + bsz_tensor = torch.tensor([qlen], device="cpu") + output = torch.empty((qlen, hidden_size), dtype=torch.bfloat16).contiguous() + CPUInfer.submit( + moe.forward_task( + bsz_tensor.data_ptr(), + num_experts_per_tok, + expert_ids.data_ptr(), + weights.data_ptr(), + input_data.data_ptr(), + output.data_ptr(), + False, + ) + ) + CPUInfer.sync() + return output + + +@pytest.mark.cpu +def test_moe_amx_skip_gpu_expert_cpu_copy_parity(): + if not HAS_DEPS: + pytest.skip(f"Dependencies not available: {import_error}") + + torch.manual_seed(1234) + + physical_to_logical_map = torch.tensor(data=range(expert_num), device="cpu", dtype=torch.int64).contiguous() + + gpu_experts_mask = torch.zeros(expert_num, dtype=torch.uint8).contiguous() + gpu_experts_mask[::2] = 1 + skip_mask_bool = gpu_experts_mask.bool() + + CPUInfer = kt_kernel_ext.CPUInfer(60) + + with torch.inference_mode(mode=True): + gate_proj = torch.randn((expert_num, intermediate_size, hidden_size), dtype=torch.bfloat16).contiguous() + up_proj = torch.randn((expert_num, intermediate_size, hidden_size), dtype=torch.bfloat16).contiguous() + down_proj = torch.randn((expert_num, hidden_size, intermediate_size), dtype=torch.bfloat16).contiguous() + + moe_off = _build_moe( + CPUInfer, gate_proj, up_proj, down_proj, gpu_experts_mask, physical_to_logical_map, skip=False + ) + moe_on = _build_moe( + CPUInfer, gate_proj, up_proj, down_proj, gpu_experts_mask, physical_to_logical_map, skip=True + ) + + for qlen in (1, 7): + expert_ids = torch.stack( + [torch.randperm(expert_num)[:num_experts_per_tok] for _ in range(qlen)] + ).contiguous() + weights = torch.rand((qlen, num_experts_per_tok), dtype=torch.float32).contiguous() + input_data = (torch.randn((qlen, hidden_size), dtype=torch.bfloat16) / 100).contiguous() + + routed_masked = int(skip_mask_bool[expert_ids.view(-1)].sum()) + assert routed_masked > 0, "test setup routes no tokens to GPU-masked experts; mask never fires" + + out_off = _forward(CPUInfer, moe_off, expert_ids, weights, input_data) + out_on = _forward(CPUInfer, moe_on, expert_ids, weights, input_data) + + assert torch.equal(out_off, out_on), ( + f"skip_gpu_expert_cpu_copy parity failed at qlen={qlen}: " + f"max abs diff = {(out_off.float() - out_on.float()).abs().max():.6f}" + ) + + t_output = moe_torch( + input_data, expert_ids, weights, gate_proj, up_proj, down_proj, skip_mask=skip_mask_bool + ) + denom = torch.mean(torch.abs(t_output)) + diff = torch.mean(torch.abs(out_on.float() - t_output.float())) / denom if denom > 0 else torch.tensor(0.0) + print(f"qlen={qlen}, routed_masked={routed_masked}, parity=exact, diff_vs_ref={float(diff):.6f}") + assert diff < 0.35, f"skip path accuracy failed at qlen={qlen}: diff={float(diff):.6f} >= 0.35" + + +def test_skip_gpu_expert_cpu_copy_rejected_in_sft_mode(): + """SFT trains every expert on CPU and needs full host copies; the factory + must refuse the inference-only skip before any backend is constructed.""" + if not HAS_DEPS: + pytest.skip(f"Dependencies not available: {import_error}") + from kt_kernel.experts import KTMoEWrapper + + with pytest.raises(ValueError, match="skip_gpu_expert_cpu_copy"): + KTMoEWrapper( + layer_idx=0, + num_experts=expert_num, + num_experts_per_tok=num_experts_per_tok, + hidden_size=hidden_size, + moe_intermediate_size=intermediate_size, + gpu_experts_mask=None, + cpuinfer_threads=2, + threadpool_count=1, + weight_path="/nonexistent", + chunked_prefill_size=max_len, + skip_gpu_expert_cpu_copy=True, + method="AMXBF16_SFT", + mode="sft", + ) + + +def run_all_tests(): + if not HAS_DEPS: + print(f"Dependencies not available: {import_error}") + return + try: + test_moe_amx_skip_gpu_expert_cpu_copy_parity() + print("skip_gpu_expert_cpu_copy parity test passed") + test_skip_gpu_expert_cpu_copy_rejected_in_sft_mode() + print("skip_gpu_expert_cpu_copy sft-mode rejection test passed") + except Exception as e: + print(f"Test failed: {e}") + import traceback + + traceback.print_exc() + sys.exit(1) + + +if __name__ == "__main__": + run_all_tests()