kvcache-ai / kvcache-ai/ktransformers

[Bug] Kimi-K2.6 RAWINT4: corrupted output with `--tensor-parallel-size 4`, correct output with TP1 — same weights, same build (4× H200 NVL, AMD EPYC 9555)

Open
#2,076 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

### Reminder

- [x] I have read the above rules and searched the existing issues.

### System Info

- CPU: 2× AMD EPYC 9555 (Zen 5, 128 physical cores, AVX-512/VNNI/BF16, **no AMX**), 2 NUMA nodes
- RAM: 1.5 TB DDR5
- GPU: 4× NVIDIA H200 NVL, NVLink fully active (NV6 all-to-all per `nvidia-smi topo -m`, 18 links @ 26.56 GB/s each)
- OS: RHEL 9 (kernel 5.14.0-687.13.1.el9_8), GPFS filesystem
- Python 3.12.3, torch 2.9.1+cu128, nvidia-cudnn-cu12 9.16.0.29
- sglang-kt 0.6.3 (PyPI), transformers-kt 5.6.0.post1, sgl-kernel 0.3.21
- kt-kernel 0.6.3.post1, **built from source** at tag v0.6.3.post1 with gcc 12.2, binutils 2.44,
`CPUINFER_CPU_INSTRUCT=AVX512`, `CPUINFER_ENABLE_AMX=OFF` (EPYC has no AMX)
- Model: moonshotai/Kimi-K2.6, native INT4 (~595 GB), download verified complete (hf re-download; all 64 shards + all .py files present)

### Reproduction

## Command

```bash
python -m sglang.launch_server \
--host 0.0.0.0 --port 30000 \
--model $MODEL_PATH --kt-weight-path $MODEL_PATH \
--kt-cpuinfer 120 --kt-threadpool-count 2 \
--kt-num-gpu-experts 30 --kt-method RAWINT4 \
--kt-gpu-prefill-token-threshold 400 \
--trust-remote-code --tensor-parallel-size 4 \
--attention-backend flashinfer \
--mem-fraction-static 0.94 --enable-mixed-chunk \
--enable-p2p-check --disable-shared-experts-fusion \
--chunked-prefill-size 32768 --max-total-tokens 65536 \
--max-running-requests 32 \
--tool-call-parser kimi_k2 --reasoning-parser kimi_k2
```

## Symptom

Server starts cleanly, loads all weights, serves at ~321 tok/s prefill / ~33 tok/s decode (single request). Output is corrupted from the first token: random multilingual tokens (mixed Latin/CJK/Burmese/code fragments) degrading within a few hundred tokens into a stable repetition attractor (`" | | | | ..."` or `****...`). Example (`reasoning_content`, truncated):

```
Transient REQUIRE whales必需的includeing_execute不情愿-scrollbar in Fre蔡徐坤 multitudenob由美国nombreex | lightly Controlled:: Unifiedsymmogan | ... | | | | | | | | | | | | ...
```

Identical corruption signature across every configuration variation tested (below), which suggests deterministic corruption in a component common to all of them.

## Ruled out (all produce byte-similar garbage)

| Variation | Result |
|---|---|
| kt-kernel 0.6.1.post1 (source) + sglang-kt 0.6.1 | garbage (note: 0.6.1 reports "no x86 RAWINT4 backend available" in test suite, so this pairing is expected-invalid) |
| kt-kernel 0.6.3.post1 (source) + sglang-kt 0.6.3, transformers-kt 5.6.0.post1 | garbage |
| `--fp8-gemm-backend auto` (default) vs `triton` (per #1961) | garbage both |
| `--disable-custom-all-reduce` | garbage |
| cudnn 9.10 vs 9.16 | garbage both (9.16 required by startup check) |
| Weights re-verified via `hf download` resume | complete, no change |

## What passes

`kt-kernel/test/per_commit/test_moe_rawint4_accuracy.py` (main) against the same source-built kt-kernel 0.6.3.post1 binary on the same node: **3/3 PASSED**. Both `AVX2RawInt4_MOE` and `AVXVNNI256RawInt4_MOE` instantiate across 2 NUMA pools; diffs vs reference 0.003–0.016 at qlen 1 and 16. So the isolated CPU kernel arithmetic appears correct on this CPU.

## Decisive experiment: TP1 works, TP4 does not

Changing **only** `--tensor-parallel-size 4` → `--tensor-parallel-size 1` (with `--kt-num-gpu-experts` reduced to fit a single GPU), same weights, same kt-kernel binary, same sglang-kt, same node:

- **TP4:** corrupted output from token one (sample above), 100% reproducible
- **TP2:** corrupted output, same signature (confirmed 2026-07-07)
- **TP1:** fully correct, coherent output; clean `finish_reason: "stop"`; verified on domain-specific physics prompts

So the CPU RAWINT4 kernels, weights, and single-GPU path are all correct; the corruption is introduced by the TP>1 path — presumably the TP slicing/packing of RAWINT4 expert weights/scales or the GPU-expert path (`CompressedTensorsWNA16MarlinMoEMethod` per the logs) under TP>1. This is consistent with all working community reports we could find being TP1 (#1816, #1820), and the official benchmark being 8× L20 (we could not find a validated TP4/NVLink RAWINT4 report).

## Corroborating datapoint: DeepSeek-V4-Flash MXFP4 crashes at TP=2 with matching geometry

On a different machine (2× H100 NVL 94 GB), `--kt-method MXFP4` with `deepseek-ai/DeepSeek-V4-Flash` and `--tensor-parallel-size 2` fails during weight loading with

```
RuntimeError: The size of tensor a (2048) must match the size of tensor b (4096) at non-singleton dimension 1
```

— an exact factor-2 shape mismatch on a byte-packed 4-bit expert tensor at TP=2. This looks like the same root cause surfacing loudly (TP-shard size computed on packed bytes vs logical elements): where the packed-width arithmetic happens to produce *valid* shapes (Kimi RAWINT4), the weights load misaligned and the output is entropy; where it produces *invalid* shapes (V4 MXFP4), loading crashes. Both models work / proceed further at TP=1.

Question: is `--kt-method RAWINT4`/`MXFP4` with `--tensor-parallel-size > 1` a validated configuration? We are happy to run further diagnostics on this machine (dual-Turin + 4× H200 NVL, NV6 all-to-all) — it is exactly the AMD-AVX512 configuration your docs advertise, and we can reproduce both the failing and working case on demand.

### Others

_No response_

Contributor guide

Open the contributing guide

Research direction

Start with kt-kernel/test/per_commit/test_moe_rawint4_accuracy.py, then trace the TP>1 weight-loading and slicing path for RAWINT4 and MXFP4, including CompressedTensorsWNA16MarlinMoEMethod. Compare TP1, TP2, and TP4 behavior using the provided commands and logs. Done means TP>1 loads correctly and produces coherent output without the reported packed-tensor shape mismatch.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, pytorch
Domain
distributed-systems, machine-learning
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Quiet
Clarity
Mostly clear
Newbie friendliness
48/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.