kvcache-ai / kvcache-ai/ktransformers

kt-kernel loads every expert into CPU RAM even when they're GPU-resident (follow-up to #2023)

Open
#2,084 3 comments 0 reactions 0 assignees View on GitHub
Dominant language
Python
Stars
19.5k
Forks
1.6k
Avg merge
19h 32m
Merged PRs (30d)
27

Description

This is the same thing #2023 raised: with `--kt-num-gpu-experts N`, kt-kernel still loads all
256 experts per layer into CPU RAM, including the N that get uploaded to and run on the GPU.
That issue was closed as "feasible but", and I think the two concerns raised there are
answerable, so I wanted to revisit it with the actual code paths and some numbers.

For anyone landing here fresh: I'm running DeepSeek-V4-Flash (MXFP4) on a single RTX PRO 6000
(96GB) + 123GB RAM through the sglang fork. With `--kt-num-gpu-experts 110` I expected host RAM
to come down as experts moved onto the GPU, but it doesn't move at all. The scheduler's
anonymous memory sits around 170GB (RAM+swap) no matter what N I pass, and decode thrashes swap
the whole time. willbed hit the same wall in #2023 and was reaching for SSD swap to cope.

The reason is that the GPU experts' CPU copies are dead weight. `forward_*` skips them through
`should_skip_expert()`, so they get allocated and filled at load and then never read again:

- `AMX_MOE_BASE::init()` (operators/amx/moe_base.hpp) allocates gate_bb_/up_bb_/down_bb_ for the
full expert_num, with no mask check
- the per-backend `load_weights()` (operators/amx/fp4-moe.hpp on my path) copies and converts
all of them in
- gpu_experts_mask / should_skip_expert (operators/common.hpp) are only consulted at compute
time, never at alloc or load

So `--kt-num-gpu-experts` decides GPU placement but has no effect on host memory.

At N=110 that's about 63GB of redundant CPU copies. I have a patch that skips the allocation
and the load for the masked experts. On my box it brings the scheduler's resident set from
~172GB down to ~99GB, which finally fits, and decode goes from swap-thrashing to zero swap at
the same tok/s with identical output. The state it produces (a null buffer for a GPU expert) is
already what `forward_*` expects, since those experts end up with m_local_num_==0 and are never
dereferenced.

## the two concerns from #2023

When #2023 was closed, @yyj6666667 raised two fair points. I think both are handled, at least on
the mask-aware paths:

1. During GPU prefill you have to keep a layer's experts fully loaded. That's true for the
full-GPU prefill fallback, but for MXFP4 that path is already mask-aware:
`_prepare_weight_mxfp4` -> `_prepare_weight_fp8` (kt_ep_wrapper.py around :1255/:1054) copies
the GPU experts GPU->GPU from original_layer and only submits the CPU experts to
write_weight_scale_to_buffer. It never reads the CPU copies I'd be skipping. (Paths whose
fallback isn't mask-aware would, which is why this is gated; see below.)

2. Loading experts over SSD would slow the pipeline. That was a response to willbed's SSD-swap
workaround, and it's a real concern for SSD tiering, but it doesn't apply here. There's no
SSD and no tiering. The experts that actually run on the CPU stay resident exactly as they do
now. The only thing that goes away is the duplicate in-RAM copy of experts that already live
on the GPU. The load pipeline is untouched.

## where it isn't safe, and how I'd gate it

It's not unconditionally safe, so it should be opt-in and off by default. Two paths would break
a naive version:

- with kt_enable_dynamic_expert_update on a non-MXFP4 path, update_kt_wrapper_masks
(kt_ep_wrapper.py:2395) rewrites the CPU wrapper's mask in place, so an expert that started on
the GPU can get demoted to CPU compute and would then need the weights I skipped
- the INT4-Marlin fallback `_prepare_weight_int4` (def :866) submits all experts
unconditionally

So the plan would be a user-facing opt-in (off by default), folded internally into an effective
flag that's only true when dynamic update is off and the path's full-GPU fallback is mask-aware
(MXFP4/MXFP8, or gpu_prefill_token_threshold 0). If someone enables it alongside dynamic update,
dynamic update wins: the skip quietly disables, you get stock behavior, and it logs why. That
keeps it safe by construction rather than by user discipline.

## how this differs from the other memory work

There's clearly demand for this. #1421 asks for SSD offload of inactive experts, and PR #2003
(MESH) is building a full NVMe/io_uring tiered residency system. This is deliberately none of
that. It's a static, zero-overhead skip for the common hybrid case where the CPU working set
fits fine once you stop duplicating the GPU experts. It's complementary to MESH, since it
shrinks the set MESH would otherwise have to manage, and it's a much smaller change.

## if there's interest

Happy to send a PR: the opt-in resolved in Python and passed to kt-kernel as a config bool (no
new env vars), limited to the mask-aware paths, with a per_commit test that runs a hybrid split
with and without the skip and checks the outputs match. There's no test covering that case
today as far as I can see.

Env: ktransformers main (v0.6.3.post1) + sglang fork main, RTX PRO 6000 (96GB) + Ryzen 9 9950X
(AVX-512, no AMX), 123GB DDR5, CUDA sm_120.

Contributor guide

Open the contributing guide

Research direction

Start with AMX_MOE_BASE::init() in operators/amx/moe_base.hpp and the per-backend load_weights() in operators/amx/fp4-moe.hpp, then trace gpu_experts_mask and should_skip_expert in operators/common.hpp. Review the mask-aware paths and gating points in kt_ep_wrapper.py, especially _prepare_weight_mxfp4, _prepare_weight_fp8, _prepare_weight_int4, and update_kt_wrapper_masks. Done means an opt-in hybrid split avoids redundant CPU copies safely and a per_commit test confirms matching outputs.

Written by the indexing model from the issue text.

Assessment

Tech stack
cpp, python
Domain
backend, machine-learning, performance
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Quiet
Clarity
Clearly specified
Newbie friendliness
55/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.