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

[gvr13n]

Implements the proposal in #2084 (follow-up to #2023).

With --kt-num-gpu-experts N, kt-kernel still allocates and loads all experts per layer into host RAM, including the N that run on the GPU. Those CPU copies are dead weight: forward_* skips them via should_skip_expert(), so they are filled at load and never read. At N=110 on DeepSeek-V4-Flash that is ~63GB of redundant host RAM; on a 123GB box it is the difference between fitting and swap-thrashing (scheduler RSS ~172GB → ~99GB, decode swap io → 0, identical output).

Design (never-allocate, gated):

• AMX_MOE_BASE::init() pushes null gate_bb_/up_bb_/down_bb_ placeholders for masked experts instead of aligned_alloc — the same state forward_* already tolerates (m_local_num_==0).
• Every AMX backend loader that inherits init() guards with should_skip_expert on the same index it uses to access the buffers: fp4-moe.hpp, mxfp8-moe.hpp, moe.hpp (int4/int8/bf16), awq-moe.hpp. (mxfp8 previously had no guard and would have segfaulted on the null buffers — fixed here.)
• Exposed as MOEConfig.skip_gpu_expert_cpu_copy (C++ bool, default false, pybind def_readwrite). The sglang side (companion PR: feat(kt): add --kt-skip-gpu-expert-cpu-copy to reclaim host RAM for GPU experts sglang#64) resolves the user-facing flag into an effective flag that is only true when the path's full-GPU prefill fallback is mask-aware (MXFP4, or MXFP8 with dynamic expert update off); otherwise it warns once and falls back to stock behavior.

Test: test/per_commit/test_moe_amx_skip_gpu_expert_cpu_copy.py — deterministic (seed 1234), CPU-only, runs a hybrid split with a gpu_experts_mask and asserts exact bitwise parity skip-on vs skip-off on both decode (qlen=1) and prefill (qlen=7), plus correctness vs a torch reference. Full per_commit suite passes (the pre-existing unseeded flake in test_moe_amx_accuracy_int4_1.py is unrelated).

Diff is pure additions (+221, 10 files), rebased on current main.

[gemini-code-assist]

Code Review

This pull request introduces the skip_gpu_expert_cpu_copy feature to optimize memory usage by skipping the allocation and copying of CPU weight buffers for GPU-resident experts. However, the code review highlights several critical issues: the feature is currently ignored in AMXMoEWrapper, making it ineffective for the default AMX backends, and multiple weight loading, saving, and writing paths in moe.hpp, awq-moe.hpp, mxfp8-moe.hpp, and fp4-moe.hpp lack proper guards or safety checks, which will lead to null pointer dereferences and segfaults when the feature is enabled.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[gemini-code-assist]

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]

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.

[gemini-code-assist]

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]

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.

[gemini-code-assist]

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]

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

[gemini-code-assist]

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]

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.

[yyj6666667]

I can see your point : for those experts already on gpu , the cpu side copies are unnecessary.

I have 2 question:

1. How do you handle the situation when launch option --kt-enable-dynamic-expert-update is turned on
2. Does the layerwise-prefill works well right now ? since it may load unexisted experts from cpu ram.

[gvr13n]

Both cases are handled in the sglang integration side: kvcache-ai/sglang PR #64 — the kt-kernel skip_gpu_expert_cpu_copy parameter is only ever set after a policy gate there. Details:

1. --kt-enable-dynamic-expert-update

The skip is only activated when the GPU-expert set is provably static (kt_ep_wrapper.py#L2662-L2678):

_mask_static = _method == "MXFP4" or (_method == "MXFP8" and not _dyn)
_kt_skip_cpu_copy = self.kt_config.kt_skip_gpu_expert_cpu_copy and _mask_static

• MXFP8 + dynamic update ON → the flag is ignored (with a one-time layer-0 warning), so full CPU copies are retained and _update_gpu_experts_from_batch has host weights available for promote/demote.
• MXFP4 → the dynamic-update path is already skipped for MXFP4 independently of this PR (kt_ep_wrapper.py#L2999-L3007): copy_experts_weights_int4 hardcodes int4 attribute names (w13_weight_packed etc.) and the full-GPU prefill fallback reloads all experts on every fire anyway, making dynamic promotion a no-op there. Static set → skipping the CPU copies is safe.

2. Layerwise / full-GPU prefill

It never reads GPU-resident experts from CPU RAM. SharedFullContext._prepare_weight_fp8 (reused by the MXFP4 path) partitions experts by the same gpu_experts_mask that decided the skip (kt_ep_wrapper.py#L1091-L1115):

• Phase 1: GPU-resident experts are copied GPU→GPU from the persistent layer tensors (original_layer). Precondition: the model loader keeps the flat tensors alive when kt_gpu_prefill_token_threshold > 0 (the post-swizzle deletes in mxfp4_deepseek.py are gated on that).
• Phase 2: only CPU-resident experts (mask == False) stream through the KT double-buffer pipeline — those are exactly the experts that still exist in host RAM.

This combination (MXFP4, skip ON, --kt-gpu-prefill-token-threshold 4096, 100 GPU experts / 156 CPU experts) is my daily production config on an RTX PRO 6000 Blackwell + 128 GB host: the full-GPU fallback fires on every ≥4096-token chunk and has run hundreds of prepare-weight cycles across benchmarks with coherent output (~350 tok/s prefill at 40k ctx, re-verified today after rebasing onto current main).

[yyj6666667]

Both cases are handled in the sglang integration side: kvcache-ai/sglang PR #64 — the kt-kernel skip_gpu_expert_cpu_copy parameter is only ever set after a policy gate there. Details:

1. --kt-enable-dynamic-expert-update

The skip is only activated when the GPU-expert set is provably static (kt_ep_wrapper.py#L2662-L2678):

_mask_static = _method == "MXFP4" or (_method == "MXFP8" and not _dyn)
_kt_skip_cpu_copy = self.kt_config.kt_skip_gpu_expert_cpu_copy and _mask_static

• MXFP8 + dynamic update ON → the flag is ignored (with a one-time layer-0 warning), so full CPU copies are retained and _update_gpu_experts_from_batch has host weights available for promote/demote.
• MXFP4 → the dynamic-update path is already skipped for MXFP4 independently of this PR (kt_ep_wrapper.py#L2999-L3007): copy_experts_weights_int4 hardcodes int4 attribute names (w13_weight_packed etc.) and the full-GPU prefill fallback reloads all experts on every fire anyway, making dynamic promotion a no-op there. Static set → skipping the CPU copies is safe.

2. Layerwise / full-GPU prefill

It never reads GPU-resident experts from CPU RAM. SharedFullContext._prepare_weight_fp8 (reused by the MXFP4 path) partitions experts by the same gpu_experts_mask that decided the skip (kt_ep_wrapper.py#L1091-L1115):

• Phase 1: GPU-resident experts are copied GPU→GPU from the persistent layer tensors (original_layer). Precondition: the model loader keeps the flat tensors alive when kt_gpu_prefill_token_threshold > 0 (the post-swizzle deletes in mxfp4_deepseek.py are gated on that).
• Phase 2: only CPU-resident experts (mask == False) stream through the KT double-buffer pipeline — those are exactly the experts that still exist in host RAM.

This combination (MXFP4, skip ON, --kt-gpu-prefill-token-threshold 4096, 100 GPU experts / 156 CPU experts) is my daily production config on an RTX PRO 6000 Blackwell + 128 GB host: the full-GPU fallback fires on every ≥4096-token chunk and has run hundreds of prepare-weight cycles across benchmarks with coherent output (~350 tok/s prefill at 40k ctx, re-verified today after rebasing onto current main).

Here is one more thing that I confirmed this afternoon: ktransformers's SFT part recently has shared same pipeline with inference path, for which a full weight on cpu ram is neccessary. Nice work as it is, it's a little pity that this PR will not be merged right away.

（By the way, 350 tok/s is usually not the speed when you turned on layerwise prefill, especially for powerful machine like PRO 6000 with your description. But I believe your PR has no problem with that !

[gvr13n]

@yyj6666667 - First off, thank you for taking the time to review this and for confirming the SFT detail. I get the maintenance cost: this is an opt-in path nobody on the team asked for, and every future loading-pipeline change would have to be checked against it. With SFT now sharing the pipeline with inference, and needing full weights in host RAM, deferring makes sense. No hard feelings.

To make the branch safe to revisit later, I pushed a guard: the KTMoEWrapper factory now raises if mode="sft" is combined with the flag, and TP_MOE_SFT strips it (with a warning) before the sliced config reaches the shared TP_MOE base, since MOESFTConfig inherits the field. Test included. So a mis-set flag can't corrupt a training run, now or after future refactors.

On the 350 tok/s: that is with layerwise prefill on (threshold 4096, chunk 8192). The bottleneck is the platform, not the GPU. It's a consumer desktop, Ryzen 9950X with dual-channel DDR5-5200 (~54 GB/s measured), streaming 156 of 256 experts per layer from host RAM each chunk, measured cold. On a real multi-channel server it should be several times faster.

[gvr13n]

@yyj6666667 - You were right about the 350 tok/s, and thanks for the nudge. I dug in and it was not the layer-wise reload at all: on sm_120 the triton_kernels opt-flags heuristic picks a tile too big for consumer Blackwell's shared memory, the stage count collapses to 1, and the simulated-MXFP GEMM runs un-pipelined at about 7 TFLOPS.

Constraining block_m / block_k / num_stages for large-M calls takes it from 342ms to 13ms per layer, and end-to-end cold prefill went from ~350 tok/s to ~910 tok/s on the same box. The weight reload is the bottleneck only now that the GEMM is fixed. Happy to send that as a separate PR if useful, since it should help any GPU outside the SM_100 whitelist.