[GH-ISSUE #2218] RAWINT4 direct-pointer MoE frees safetensors mmap owners and segfaults on first inference (AVX2) #1306

New issue

Open

opened 2026-10-29 14:51:04 +00:00 by bsakaguchi · 2 comments

bsakaguchi commented

2026-10-29 14:51:04 +00:00

Owner

Originally created by @f-wunderlich on GitHub (Sep 30, 2026).
Original GitHub issue: https://github.com/kvcache-ai/ktransformers/issues/2218

Environment

KTransformers/kt-kernel v0.7.0.post4, release commit
5632f7197e7a9898f0723755c2fc187a946a85a6
confirmed still present on upstream main at
c40722bf04c494f2492b7eb9e86ef01a4ede45b3
moonshotai/Kimi-K2.6, revision
7eb5002f6aadc958aed6a9177b7ed26bb94011bb
RAWINT4, TP1, one CPU pool, 24 CPU threads, zero GPU experts
dual AMD EPYC 7443 (Milan, AVX2), NVIDIA Blackwell GPU
CUDA graphs, server warmup and overlap scheduling disabled

Symptom

Model loading completes, the API starts, and the scheduler segfaults on its
first RAWINT4 GEMM. Disabling CUDA graphs, overlap scheduling and GPU experts
does not prevent it.

The 136 GB core is usable. Its crashing thread is a KTransformers task-queue
worker:

Program terminated with signal SIGSEGV
#0  0x00007f024e467703 in ??
    from kt_kernel/_kt_kernel_ext_avx2.cpython-311-x86_64-linux-gnu.so
#1  0x00007f024e5261a5 in the same module
#2  0x00007f024e5268bb in the same module
#3  libstdc++.so.6

The module mapping base is 0x7f024e3cb000, giving offsets 0x9c703,
0x15b1a5 and 0x15b8bb. The latter frames are in the internal
TaskQueue::worker region.

At the fault:

r13 = 0x7ef9f0c70bb8
rdx = 0

vpbroadcastd 0x0(%r13,%rdx,1), %ymm0

GDB reports the address around r13 as unmapped.

Root cause

operators/avx2/rawint4-moe.hpp intentionally assigns BufferB.b directly to
the per-expert safetensor tensor:

// Direct-pointer mode: BufferB.b is set to point into the mmap'd safetensor data
gate_bb_[expert_idx]->b =
    (uint8_t*)config_.gate_projs[0][logical_expert_id] + gate_tp_byte_offset;

After submitting and synchronizing that load task,
kt-kernel/python/utils/amx.py deletes the tensor lists and releases the shared
loader:

del self.gate_weights
del self.up_weights
del self.down_weights
NativeMoEWrapper._release_loader(layer_idx=self.layer_idx)

_release_loader() clears SafeTensorLoader.file_handle_map and calls
gc.collect(). Consequently no Python object owns the mmap-backed storage, but
the native BufferB.b pointers remain. The first GEMM dereferences one of those
dangling pointers.

Confirmed workaround

Retain self.gate_weights, self.up_weights and self.down_weights when
self.method == "RAWINT4". Scale tensors may still be released because the
native loader converts/copies them. Loader handles may also be cleared: the
retained tensor storage objects keep the mappings alive.

With that one change, the same image, checkpoint and launch arguments:

loaded all 60 MoE layers;
completed native generation and OpenAI chat requests;
returned the correct final answer for a deterministic prompt;
remained Ready with zero restarts.

No model-weight copy is required; retaining tensor owners preserves the mmap
design and avoids allocating a second copy of the checkpoint.

Suggested fix

Make mmap ownership explicit for every native direct-pointer backend. At
minimum, RAWINT4 wrappers must retain the source tensors for at least as long
as their native MoE object can access BufferB.b. A small regression test can
load a direct-pointer RAWINT4 layer, release the loader and run one GEMM after
forcing Python garbage collection.

Originally created by @f-wunderlich on GitHub (Sep 30, 2026). Original GitHub issue: https://github.com/kvcache-ai/ktransformers/issues/2218 ## Environment - KTransformers/kt-kernel `v0.7.0.post4`, release commit `5632f7197e7a9898f0723755c2fc187a946a85a6` - confirmed still present on upstream `main` at `c40722bf04c494f2492b7eb9e86ef01a4ede45b3` - `moonshotai/Kimi-K2.6`, revision `7eb5002f6aadc958aed6a9177b7ed26bb94011bb` - `RAWINT4`, TP1, one CPU pool, 24 CPU threads, zero GPU experts - dual AMD EPYC 7443 (Milan, AVX2), NVIDIA Blackwell GPU - CUDA graphs, server warmup and overlap scheduling disabled ## Symptom Model loading completes, the API starts, and the scheduler segfaults on its first RAWINT4 GEMM. Disabling CUDA graphs, overlap scheduling and GPU experts does not prevent it. The 136 GB core is usable. Its crashing thread is a KTransformers task-queue worker: ```text Program terminated with signal SIGSEGV #0 0x00007f024e467703 in ?? from kt_kernel/_kt_kernel_ext_avx2.cpython-311-x86_64-linux-gnu.so #1 0x00007f024e5261a5 in the same module #2 0x00007f024e5268bb in the same module #3 libstdc++.so.6 ``` The module mapping base is `0x7f024e3cb000`, giving offsets `0x9c703`, `0x15b1a5` and `0x15b8bb`. The latter frames are in the internal `TaskQueue::worker` region. At the fault: ```text r13 = 0x7ef9f0c70bb8 rdx = 0 vpbroadcastd 0x0(%r13,%rdx,1), %ymm0 ``` GDB reports the address around `r13` as unmapped. ## Root cause `operators/avx2/rawint4-moe.hpp` intentionally assigns `BufferB.b` directly to the per-expert safetensor tensor: ```cpp // Direct-pointer mode: BufferB.b is set to point into the mmap'd safetensor data gate_bb_[expert_idx]->b = (uint8_t*)config_.gate_projs[0][logical_expert_id] + gate_tp_byte_offset; ``` After submitting and synchronizing that load task, `kt-kernel/python/utils/amx.py` deletes the tensor lists and releases the shared loader: ```python del self.gate_weights del self.up_weights del self.down_weights NativeMoEWrapper._release_loader(layer_idx=self.layer_idx) ``` `_release_loader()` clears `SafeTensorLoader.file_handle_map` and calls `gc.collect()`. Consequently no Python object owns the mmap-backed storage, but the native `BufferB.b` pointers remain. The first GEMM dereferences one of those dangling pointers. ## Confirmed workaround Retain `self.gate_weights`, `self.up_weights` and `self.down_weights` when `self.method == "RAWINT4"`. Scale tensors may still be released because the native loader converts/copies them. Loader handles may also be cleared: the retained tensor storage objects keep the mappings alive. With that one change, the same image, checkpoint and launch arguments: - loaded all 60 MoE layers; - completed native generation and OpenAI chat requests; - returned the correct final answer for a deterministic prompt; - remained Ready with zero restarts. No model-weight copy is required; retaining tensor owners preserves the mmap design and avoids allocating a second copy of the checkpoint. ## Suggested fix Make mmap ownership explicit for every native direct-pointer backend. At minimum, RAWINT4 wrappers must retain the source tensors for at least as long as their native MoE object can access `BufferB.b`. A small regression test can load a direct-pointer RAWINT4 layer, release the loader and run one GEMM after forcing Python garbage collection.

bsakaguchi commented

2026-10-29 14:51:05 +00:00

Author

Owner

@yyj6666667 commented on GitHub (Oct 1, 2026):

Thanks for feedback. i will check this out

@yyj6666667 commented on GitHub (Oct 1, 2026): Thanks for feedback. i will check this out

bsakaguchi commented

2026-10-29 14:51:05 +00:00

Author

Owner

@yyj6666667 commented on GitHub (Oct 6, 2026):

Confirmed. I reproduced this issue on qj5090 and pushed a minimal fix to fix/rawint4-mmap-lifetime-2218.

The fix retains the source weight tensors specifically for the AVX2 RAWINT4 backend. Scale tensors and loader handles are still released, and copying backends keep their existing cleanup behavior. The branch changes only kt-kernel/python/utils/amx.py.

Verified on qj5090 with AVX2 explicitly selected:

Before the fix, both uint8 and int32 packing crashed on the first forward pass.
After the fix, both formats passed 21 consecutive forwards each for sequence lengths 1 and 16.
Lifecycle checks confirmed that weights remain alive during inference and are released after wrapper destruction.
The AVX512 copying-backend control and existing RAWINT4 loading tests also passed.

@yyj6666667 commented on GitHub (Oct 6, 2026): Confirmed. I reproduced this issue on qj5090 and pushed a minimal fix to [fix/rawint4-mmap-lifetime-2218](https://github.com/yyj6666667/ktransformers/tree/fix/rawint4-mmap-lifetime-2218). The fix retains the source weight tensors specifically for the AVX2 RAWINT4 backend. Scale tensors and loader handles are still released, and copying backends keep their existing cleanup behavior. The branch changes only `kt-kernel/python/utils/amx.py`. Verified on qj5090 with AVX2 explicitly selected: - Before the fix, both uint8 and int32 packing crashed on the first forward pass. - After the fix, both formats passed 21 consecutive forwards each for sequence lengths 1 and 16. - Lifecycle checks confirmed that weights remain alive during inference and are released after wrapper destruction. - The AVX512 copying-backend control and existing RAWINT4 loading tests also passed.