[PR #101] [CLOSED] [feat](kt): load STOCK ModelOpt MIXED_PRECISION NVFP4 checkpoints #102

New issue

Closed

opened 2026-10-29 21:07:14 +00:00 by bsakaguchi · 0 comments

bsakaguchi commented

2026-10-29 21:07:14 +00:00

Owner

📋 Pull Request Information

Original PR: https://github.com/kvcache-ai/sglang/pull/101
Author: @BlivionIaG
Created: 10/6/2026
Status: ❌ Closed

Base: main ← Head: feat/kt-native-mixed-precision-nvfp4

📝 Commits (5)

a700676 mixed precision config
5be27f8 per prefix dispatch
1fb14a7 nvfp4 lm head
e360687 sharded scalar scales
66642cf kt expert counts

📊 Changes

7 files changed (+313 additions, -14 deletions)

View changed files

📝 python/sglang/srt/layers/linear.py (+5 -0)
📝 python/sglang/srt/layers/logits_processor.py (+64 -0)
📝 python/sglang/srt/layers/moe/kt_ep_wrapper.py (+17 -3)
📝 python/sglang/srt/layers/parameter.py (+26 -0)
📝 python/sglang/srt/layers/quantization/modelopt_quant.py (+159 -4)
📝 python/sglang/srt/layers/vocab_parallel_embedding.py (+6 -1)
📝 python/sglang/srt/models/qwen3_5.py (+36 -6)

📄 Description

What does this PR do?

Teaches the ModelOpt config to load a MIXED_PRECISION checkpoint, and adds the
dispatch needed to run one: per-prefix FP8/NVFP4 routing, a quantised lm_head, and
scalar per-tensor scales on sharded parameters.

MIXED_PRECISION is what nvidia's Qwen3.6-35B-A3B-NVFP4 release declares. Its
hf_quant_config.json maps per layer — 130 prefixes FP8, 161 NVFP4, and lm_head
among the NVFP4 ones — so the top-level quant_algo alone cannot express which
method any given layer needs. Today such a checkpoint is rejected outright:

ValueError: ModelOpt currently only supports: FP8, NVFP4 quantizations in sglang.

None of this requires editing the model files; the checkpoint is loaded as shipped.

The five commits

Each is independently reviewable and revertable.

mixed precision config — accept MIXED_PRECISION; keep the checkpoint's
per-prefix scheme map instead of discarding it; derive group_size and the
exclusion list from the flat config.json shape, where group_size lives per
group inside config_groups and exclusions live under ignore.
per prefix dispatch — route each LinearBase to the method the checkpoint
declares for it. Two non-obvious cases: a fused module name (qkv_proj,
gate_up_proj) matches no key because the checkpoint declares the unfused
leaves, so it is resolved through them; and a prefix absent from the enumeration
means the producer left it at full precision, so it must stay unquantised
rather than take this config's NVFP4 default.
nvfp4 lm head — a ParallelLMHead is not a LinearBase, so it never reached
dispatch and was allocated at full width, which a nibble-packed weight cannot
fill. It now takes the NVFP4 method, and its projection runs on a dense CPU NVFP4
GEMM (companion ktransformers PR).
sharded scalar scales — a per-tensor scale is stored as a scalar while a
sharded parameter holds one entry per shard, so the strict shape assert rejects a
copy that is a valid broadcast. One helper replaces four duplicated call sites.
kt expert counts — num_local_experts must describe the GPU-resident subset
because ModelOpt's create_weights sizes its tensors from that attribute; and
the GPU MoE post-load is skipped when no expert is GPU-resident, since those
scales are allocated but only the CPU path ever fills them.

How was this tested?

On the stock checkpoint, unmodified (config.json / hf_quant_config.json keep
their original mtimes):

loads with zero not found in params_dict entries, where the same checkpoint
previously left lm_head.input_scale unreferenced;
serves coherent output — The capital of France is → "…situated on the River
Seine, in the region of Île de France", and 17 + 25 = → 42;
no load-time transformation of any tensor: the head's packed U8 weight and all
three of its scale tensors load as stored.

The change is +313/−14 across 7 files and applies cleanly to main.

Notes

_kt_-prefixed helpers follow the naming already used in these files (26 such
names exist upstream, including module-level functions like
_kt_ep_predicate), so the additions match their surroundings.
The MIXED_PRECISION path is the general case; a homogeneous FP8 or NVFP4
checkpoint is unaffected, since the scheme map is empty for those and dispatch
falls back to today's behaviour.

_{🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.}

## 📋 Pull Request Information **Original PR:** https://github.com/kvcache-ai/sglang/pull/101 **Author:** [@BlivionIaG](https://github.com/BlivionIaG) **Created:** 10/6/2026 **Status:** ❌ Closed **Base:** `main` ← **Head:** `feat/kt-native-mixed-precision-nvfp4` --- ### 📝 Commits (5) - [`a700676`](https://github.com/kvcache-ai/sglang/commit/a700676cab93da2569586572fe5ccaa1f29dc03a) mixed precision config - [`5be27f8`](https://github.com/kvcache-ai/sglang/commit/5be27f8bb81cc177a80fd60946461bfb104a313f) per prefix dispatch - [`1fb14a7`](https://github.com/kvcache-ai/sglang/commit/1fb14a770b84079fb4a2c94297ae15eb05f08c11) nvfp4 lm head - [`e360687`](https://github.com/kvcache-ai/sglang/commit/e360687a69bb86208b35fcbcf7680ecfe82fc99f) sharded scalar scales - [`66642cf`](https://github.com/kvcache-ai/sglang/commit/66642cfcfdc89d7f9598355c272fe517855a0722) kt expert counts ### 📊 Changes **7 files changed** (+313 additions, -14 deletions) <details> <summary>View changed files</summary> 📝 `python/sglang/srt/layers/linear.py` (+5 -0) 📝 `python/sglang/srt/layers/logits_processor.py` (+64 -0) 📝 `python/sglang/srt/layers/moe/kt_ep_wrapper.py` (+17 -3) 📝 `python/sglang/srt/layers/parameter.py` (+26 -0) 📝 `python/sglang/srt/layers/quantization/modelopt_quant.py` (+159 -4) 📝 `python/sglang/srt/layers/vocab_parallel_embedding.py` (+6 -1) 📝 `python/sglang/srt/models/qwen3_5.py` (+36 -6) </details> ### 📄 Description ## What does this PR do? Teaches the ModelOpt config to load a **`MIXED_PRECISION`** checkpoint, and adds the dispatch needed to run one: per-prefix FP8/NVFP4 routing, a quantised `lm_head`, and scalar per-tensor scales on sharded parameters. `MIXED_PRECISION` is what nvidia's `Qwen3.6-35B-A3B-NVFP4` release declares. Its `hf_quant_config.json` maps per layer — 130 prefixes FP8, 161 NVFP4, and `lm_head` among the NVFP4 ones — so the top-level `quant_algo` alone cannot express which method any given layer needs. Today such a checkpoint is rejected outright: ValueError: ModelOpt currently only supports: FP8, NVFP4 quantizations in sglang. None of this requires editing the model files; the checkpoint is loaded as shipped. ## The five commits Each is independently reviewable and revertable. 1. **mixed precision config** — accept `MIXED_PRECISION`; keep the checkpoint's per-prefix scheme map instead of discarding it; derive `group_size` and the exclusion list from the flat `config.json` shape, where `group_size` lives per group inside `config_groups` and exclusions live under `ignore`. 2. **per prefix dispatch** — route each `LinearBase` to the method the checkpoint declares for it. Two non-obvious cases: a fused module name (`qkv_proj`, `gate_up_proj`) matches no key because the checkpoint declares the *unfused* leaves, so it is resolved through them; and a prefix absent from the enumeration means the producer left it at full precision, so it must stay **unquantised** rather than take this config's NVFP4 default. 3. **nvfp4 lm head** — a `ParallelLMHead` is not a `LinearBase`, so it never reached dispatch and was allocated at full width, which a nibble-packed weight cannot fill. It now takes the NVFP4 method, and its projection runs on a dense CPU NVFP4 GEMM (companion ktransformers PR). 4. **sharded scalar scales** — a per-tensor scale is stored as a scalar while a sharded parameter holds one entry per shard, so the strict shape assert rejects a copy that is a valid broadcast. One helper replaces four duplicated call sites. 5. **kt expert counts** — `num_local_experts` must describe the GPU-resident subset because ModelOpt's `create_weights` sizes its tensors from that attribute; and the GPU MoE post-load is skipped when no expert is GPU-resident, since those scales are allocated but only the CPU path ever fills them. ## How was this tested? On the stock checkpoint, unmodified (`config.json` / `hf_quant_config.json` keep their original mtimes): - loads with **zero** `not found in params_dict` entries, where the same checkpoint previously left `lm_head.input_scale` unreferenced; - serves coherent output — `The capital of France is` → *"…situated on the River Seine, in the region of Île de France"*, and `17 + 25 =` → `42`; - no load-time transformation of any tensor: the head's packed U8 weight and all three of its scale tensors load as stored. The change is +313/−14 across 7 files and applies cleanly to `main`. ## Notes - `_kt_`-prefixed helpers follow the naming already used in these files (26 such names exist upstream, including module-level functions like `_kt_ep_predicate`), so the additions match their surroundings. - The `MIXED_PRECISION` path is the general case; a homogeneous FP8 or NVFP4 checkpoint is unaffected, since the scheme map is empty for those and dispatch falls back to today's behaviour. --- <sub>🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.</sub>

bsakaguchi added the

pull-request

label

2026-10-29 21:07:14 +00:00

bsakaguchi referenced this issue

2026-10-29 21:07:14 +00:00

[PR #102] [Fix] Fence all TP readers before FP8 staging-slot reuse #103

bsakaguchi closed this issue

2026-10-29 21:07:14 +00:00