Skip to content

[feat](kt-kernel): opt-in skip of CPU copies for GPU-resident experts #2085

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
gvr13n wants to merge 3 commits into
base: main
Choose a base branch
from
Open

[feat](kt-kernel): opt-in skip of CPU copies for GPU-resident experts #2085

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
54e0943
[feat](kt-kernel): opt-in skip of CPU copies for GPU-resident experts
gvr13n Jun 24, 2026
4652d91
[fix](kt-kernel): qualify should_skip_expert guards and gate them on …
gvr13n Jul 13, 2026
7ff2203
[feat](kt-kernel): reject skip_gpu_expert_cpu_copy in SFT mode
gvr13n Jul 13, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions kt-kernel/ext_bindings.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -709,6 +709,7 @@ PYBIND11_MODULE(kt_kernel_ext, m) {
"gpu_experts_mask",
[](const GeneralMOEConfig& self) { return reinterpret_cast<uintptr_t>(self.gpu_experts_mask); },
[](GeneralMOEConfig& self, uintptr_t val) { self.gpu_experts_mask = reinterpret_cast<uint8_t*>(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)
Expand Down
11 changes: 11 additions & 0 deletions kt-kernel/operators/amx/awq-moe.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -514,6 +514,7 @@ class AMX_AWQ_MOE_TP : public AMX_MOE_BASE<T, AMX_AWQ_MOE_TP<T>> {
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;

@gemini-code-assist gemini-code-assist Bot Jul 12, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

The config_.save branches in awq-moe.hpp (around lines 583 and 631) are missing should_skip_expert guards.

When config_.save is enabled, the loops iterate over all experts from 0 to config_.expert_num and call write_weights with gate_bb_[expert_idx].get(). If skip_gpu_expert_cpu_copy is true, gate_bb_[expert_idx] is nullptr for skipped experts, causing a segfault when write_weights attempts to access buffer->b.

Please add a guard to skip saving weights for skipped experts:

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((int64_t)expert_idx)) continue;
    write_weights(prefix, "_gate_", gate_bb_[expert_idx].get(), expert_idx, "OFFLINE");
    ...

@gvr13n gvr13n Jul 12, 2026

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Fixed in 54e0943: both config_.save loops in awq-moe.hpp (OFFLINE ~:583 and ONLINE ~:631) now continue on should_skip_expert(expert_idx).

uint64_t logical_expert_id = expert_map(physical_to_logical_map, expert_idx);
int ith = task_id % nth;
// gate part
Expand All @@ -534,6 +535,7 @@ class AMX_AWQ_MOE_TP : public AMX_MOE_BASE<T, AMX_AWQ_MOE_TP<T>> {
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
Expand All @@ -548,6 +550,7 @@ class AMX_AWQ_MOE_TP : public AMX_MOE_BASE<T, AMX_AWQ_MOE_TP<T>> {
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;
Expand Down Expand Up @@ -579,6 +582,9 @@ class AMX_AWQ_MOE_TP : public AMX_MOE_BASE<T, AMX_AWQ_MOE_TP<T>> {
// 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");
Expand All @@ -594,6 +600,7 @@ class AMX_AWQ_MOE_TP : public AMX_MOE_BASE<T, AMX_AWQ_MOE_TP<T>> {
[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(
Expand All @@ -613,6 +620,7 @@ class AMX_AWQ_MOE_TP : public AMX_MOE_BASE<T, AMX_AWQ_MOE_TP<T>> {
[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(
Expand All @@ -625,6 +633,9 @@ class AMX_AWQ_MOE_TP : public AMX_MOE_BASE<T, AMX_AWQ_MOE_TP<T>> {
// 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");
Expand Down
3 changes: 3 additions & 0 deletions kt-kernel/operators/amx/fp4-moe.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -442,6 +442,7 @@ class AMX_FP4_MOE_TP : public AMX_MOE_BASE<T, AMX_FP4_MOE_TP<T>> {
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(
Expand All @@ -459,6 +460,7 @@ class AMX_FP4_MOE_TP : public AMX_MOE_BASE<T, AMX_FP4_MOE_TP<T>> {
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(
Expand All @@ -472,6 +474,7 @@ class AMX_FP4_MOE_TP : public AMX_MOE_BASE<T, AMX_FP4_MOE_TP<T>> {
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,
Expand Down
6 changes: 6 additions & 0 deletions kt-kernel/operators/amx/moe.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -211,6 +211,7 @@ class AMX_MOE_TP : public AMX_MOE_BASE<T, AMX_MOE_TP<T>> {
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;

@gemini-code-assist gemini-code-assist Bot Jul 12, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

There are two other loader paths in moe.hpp that are missing the should_skip_expert guards and will SEGFAULT if skip_gpu_expert_cpu_copy is enabled:

  1. The config_.load branch (around line 256):
    When loading pre-quantized weights from disk, the loop iterates over all experts without checking if they are skipped. Since up_bb_[expert_idx] is nullptr for skipped experts, calling read_weights on it will dereference nullptr and crash.

  2. The config_.save branch (around line 331):
    When saving weights, the loop iterates over all experts and calls write_weights with up_bb_[expert_idx]->b, which will also dereference nullptr and crash.

Please add the same should_skip_expert guards to both of these branches. For example, in the config_.save branch:

[this, physical_to_logical_map, prefix](int task_id) {
  int64_t expert_idx = task_id / mat_type_all;
  if (config_.skip_gpu_expert_cpu_copy && config_.should_skip_expert((int64_t)expert_idx)) return;
  expert_idx = expert_map(physical_to_logical_map, expert_idx);
  ...

@gvr13n gvr13n Jul 12, 2026

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Good catch — fixed in 54e0943: the config_.load work-stealing lambda now early-returns via should_skip_expert(expert_idx) (the same physical index used to access gate/up/down_bb_) before any buffer deref.

uint64_t logical_expert_id = expert_map(physical_to_logical_map, expert_id);
{
size_t scale_size = config_.intermediate_size * sizeof(float);
Expand Down Expand Up @@ -254,6 +255,9 @@ class AMX_MOE_TP : public AMX_MOE_BASE<T, AMX_MOE_TP<T>> {
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;
Expand Down Expand Up @@ -291,6 +295,7 @@ class AMX_MOE_TP : public AMX_MOE_BASE<T, AMX_MOE_TP<T>> {
[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(
Expand All @@ -309,6 +314,7 @@ class AMX_MOE_TP : public AMX_MOE_BASE<T, AMX_MOE_TP<T>> {
[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(
Expand Down
7 changes: 7 additions & 0 deletions kt-kernel/operators/amx/moe_base.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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));
Expand Down
3 changes: 3 additions & 0 deletions kt-kernel/operators/amx/mxfp8-moe.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -566,6 +566,7 @@ class AMX_MXFP8_MOE_TP : public AMX_MOE_BASE<T, AMX_MXFP8_MOE_TP<T>> {
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;

@gemini-code-assist gemini-code-assist Bot Jul 12, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

In mxfp8-moe.hpp (and similarly in fp4-moe.hpp), write_weights_to_buffer is used to copy CPU expert weights to the GPU pinned buffer.

If skip_gpu_expert_cpu_copy is enabled, the CPU buffers (gate_bb_, up_bb_, down_bb_) for GPU-resident experts are nullptr. If write_weights_to_buffer is called for a skipped expert, it will dereference nullptr (e.g., bb->b) and SEGFAULT.

To prevent silent crashes, please add a safety check at the beginning of write_weights_to_buffer in both mxfp8-moe.hpp and fp4-moe.hpp:

if (gate_bb_[expert_id] == nullptr) {
  throw std::runtime_error("Cannot write weights to buffer for skipped/GPU-resident expert");
}

@gvr13n gvr13n Jul 12, 2026

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This path is safe by construction: write_weights_to_buffer is only reached from the sglang full-GPU-prefill pipeline, whose mask-aware callers (_prepare_weight_fp8/_prepare_weight_mxfp4) partition experts and submit ONLY cpu_expert_ids to the wrapper — a skipped (GPU-resident) expert never reaches it. Adding a null-check inside would silently mask a caller bug rather than surface it, so I have left it unguarded deliberately. Happy to add an assert instead if maintainers prefer.

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)
Expand All @@ -580,6 +581,7 @@ class AMX_MXFP8_MOE_TP : public AMX_MOE_BASE<T, AMX_MXFP8_MOE_TP<T>> {
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;
Expand All @@ -593,6 +595,7 @@ class AMX_MXFP8_MOE_TP : public AMX_MOE_BASE<T, AMX_MXFP8_MOE_TP<T>> {
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<size_t>(config_.intermediate_size) * config_.hidden_size) / config_.quant_config.group_size;
Expand Down
1 change: 1 addition & 0 deletions kt-kernel/operators/common.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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() {
Expand Down
20 changes: 19 additions & 1 deletion kt-kernel/operators/moe-sft-tp.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -215,9 +215,27 @@ class TP_MOE_SFT : public TP_MOE<T> {
std::vector<ggml_bf16_t*> part_grad_input_;
std::vector<float*> 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<const GeneralMOEConfig&>(config)), sft_config(config) {
TP_MOE_SFT(const MOESFTConfig& config)
: Base(reject_skip_gpu_expert_cpu_copy(static_cast<const GeneralMOEConfig&>(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);
Expand Down
14 changes: 14 additions & 0 deletions kt-kernel/python/experts.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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,
Expand All @@ -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,
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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
Comment on lines +379 to +380

@gemini-code-assist gemini-code-assist Bot Jul 12, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

While skip_gpu_expert_cpu_copy is correctly passed to NativeMoEWrapper, it is completely ignored for AMXMoEWrapper (which handles AMXINT4 and AMXINT8 quantization methods).

In kt-kernel/python/utils/amx.py, AMXMoEWrapper.__init__ does not accept skip_gpu_expert_cpu_copy, and its load_weights / load_weights_from_tensors methods do not set moe_config.skip_gpu_expert_cpu_copy. This means the feature will have no effect when using the default/most common AMXINT4 or AMXINT8 backends, even though moe.hpp was updated to support it.

To fix this, please:

  1. Add skip_gpu_expert_cpu_copy: bool = False to AMXMoEWrapper.__init__ and store it as self.skip_gpu_expert_cpu_copy.
  2. Set moe_config.skip_gpu_expert_cpu_copy = self.skip_gpu_expert_cpu_copy in both load_weights and load_weights_from_tensors of AMXMoEWrapper.
  3. Update experts.py to pass skip_gpu_expert_cpu_copy to AMXMoEWrapper as well.

@gvr13n gvr13n Jul 12, 2026

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Intentional scoping, not an oversight: the user-facing flag is resolved (sglang side, kvcache-ai/sglang#64) into an effective flag that is only true for mask-static paths — MXFP4, or MXFP8 with dynamic expert update off — which are NativeMoEWrapper paths. AMXINT4/8 can be demoted at runtime by the dynamic-update mask rewrite, so an expert whose CPU copy was skipped could later be needed on CPU; wiring the flag through AMXMoEWrapper without that guarantee would trade a safe no-op for a correctness hazard. If a mask-static AMX mode lands later, plumbing it through is a small follow-up.

return backend_cls(
layer_idx=layer_idx,
num_experts=num_experts,
Expand Down
3 changes: 3 additions & 0 deletions kt-kernel/python/utils/amx.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 耦合.
Expand Down Expand Up @@ -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
Expand Down
Loading