feat(sft) Deepseek v4 flash LoRA (#2216)

* feat(sft): add native MXFP4 routed-expert LoRA kernels

* test(sft): validate DeepSeek V4 MXFP4 routed experts

* feat(sft): AVX2 tier for MXFP4 routed-expert LoRA SFT

MXFP4_SFT_MOE on non-AVX512 builds, same Python contract as the AMX /
AVX512-BF16 kernel from #2185 (KTMoEWrapper(method="MXFP4_SFT")), for
AVX2+FMA+F16C hosts (Zen 2/3 EPYC, older Xeons).

* operators/avx2/sft_mxfp4-moe.hpp: AVX2_SFT_MXFP4_MOE_TP part on top of
  AVX2_MXFP4_MOE_TP (frozen packed E2M1/UE8M0 base, LoRA gate/up/down,
  V4 swiglu clamp, transpose-free packed-weight dX, authoritative LoRA grads).
* operators/moe-sft-tp.hpp: AVX2 result merge; AVX-512 merge loops and the
  block-FP8 staging guarded so the wrapper builds without AVX-512.
* ext_bindings.cpp: bind MXFP4_SFT_MOE -> TP_MOE_SFT<AVX2_SFT_MXFP4_MOE_TP>
  when USE_AMX_AVX_KERNEL is off; MOESFTBindings no longer AMX-only.
* python/sft/backend.py: runtime reports cpu_variant=avx2, kernel avx2-mxfp4-g32.
* test_mxfp4_sft_numerical.py: inference regression uses AVX2MXFP4_MOE when
  AMXFP4_KGroup_MOE is absent.

Validated on 2x EPYC 7452 (AVX2, no AVX-512): test_mxfp4_sft_numerical.py
TP1/TP2 all stages rel-L2 <= 0.4% (limits 3-5%); validate_dsv4_mxfp4_sft_layer.py
on DeepSeek-V4-Flash-0731 layer 3, TP1 and TP2: passed, forward/grad_input/
LoRA grads rel-L2 0.10-0.29%, base weights unchanged.

Note: the build uses -ffast-math, so exp() past the float range is undefined;
the scalar sigmoid tails clamp the argument like the vector path.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_011Mgsduzz2fx1gSvfHbR8RZ

* fix(sft): clamp MXFP4 scalar-tail sigmoid exponent

* feat(sft): integrate native V4 experts with LoRA training and export

Wrap only routed experts while preserving V4 routing and shared branches. Validate native checkpoint provenance and stream a BF16 non-expert cache without dequantizing routed experts.

Add audited static SGLang export and reproducible training, resume and generation acceptance helpers. CPU retain/GPU recompute GAS1/GAS4 smoke and 265 KT regression tests pass; long-run quality and resume acceptance remain in progress.

* fix(sft): retain async inference buffers until CUDA work completes

* test(sft): isolate V4 adapter effects with matched-kernel controls

Allow baseline and non-expert-only exports to retain a zero-B expert adapter so all ablations execute the same native forward kernel. Add paired held-out generation reporting with explicitly diagnostic lexical metrics.

* test(sft): compare PEFT targets semantically during V4 resume audits

* test(sft): cancel queued V4 inference checks after a failed request

* test(sft): preserve raw V4 responses before finite-output validation

* feat: prepare V4 cache atomically and package static adapter export

* build: preserve accepted CUDA carriers across Python ABIs

* test: retain FP8 loading contracts after V4 wrapper integration

* test: exclude historical concatenated-sample smoke from delivery

* test(sft): audit unpacked batches and single-GPU resume state

* docs(sft): align April V4 recipe with approved unpacked delivery

* test: keep frozen-parameter audit indices exact for large tensors

* test: record the loaded native backend in acceptance evidence

* test: timestamp training audit events for startup measurements

* docs: pin V4 serving dependencies and account for reserved token slots

* [refactor](sft): keep V4 follow-up focused on runtime integration

---------

Co-authored-by: 05yuki <ceu.azul.samu@gmail.com>
Co-authored-by: Claude Opus 5 <noreply@anthropic.com>

kt-kernel/python/experts.py

@@ -53,6 +53,7 @@ SFT_METHODS = frozenset(
     [
         "AMXBF16_SFT",  # AMX BF16 training
         "AMXFP8_SFT",  # Block-wise E4M3 routed-expert LoRA training
+        "MXFP4_SFT",  # Native E2M1/UE8M0 group-32 routed-expert LoRA training
         "INT8_SFT",  # Hardware-neutral INT8 training (AMX or AVX512-VNNI)
         "AMXINT8_SFT",  # AMX INT8 training
         "AMXINT4_SFT",  # AMX INT4 training
@@ -114,7 +115,7 @@ class KTMoEWrapper:
             threadpool_count=4,
             weight_path="/path/to/weights",
             chunked_prefill_size=25600,
-            method="AMXBF16_SFT",  # or "AMXINT8_SFT", "AMXINT4_SFT"
+            method="AMXBF16_SFT",  # or "MXFP4_SFT", "INT8_SFT"
             mode="sft",
             lora_rank=16,
             lora_alpha=32.0,
@@ -234,14 +235,13 @@ class KTMoEWrapper:
                 swiglu_alpha=swiglu_alpha,
             )
         else:  # mode == "sft"
-            # SFT factory does not plumb swiglu_limit; reject non-zero
-            # rather than dropping it on the floor. Origin: kt-sglang 耦合.
-            if swiglu_limit != 0.0:
+            if swiglu_limit != 0.0 and method != "MXFP4_SFT":
                 raise ValueError(
-                    f"swiglu_limit={swiglu_limit} is not supported in "
-                    f"mode='sft' (method={method!r}); SFT backends do not "
-                    f"implement the V4-2604B clamp."
+                    f"swiglu_limit={swiglu_limit} is only supported by "
+                    f"method='MXFP4_SFT' in SFT mode, got {method!r}"
                 )
+            if method == "MXFP4_SFT" and swiglu_alpha != 0.0:
+                raise ValueError("MXFP4_SFT requires swiglu_alpha=0")
             return _create_sft_wrapper(
                 layer_idx=layer_idx,
                 num_experts=num_experts,
@@ -261,6 +261,7 @@ class KTMoEWrapper:
                 group_size=group_size,
                 zero_point=zero_point,
                 full_weight_grad=full_weight_grad,
+                swiglu_limit=swiglu_limit,
             )
 
     # Forward static methods to the base class
@@ -427,6 +428,7 @@ def _create_sft_wrapper(
     group_size: int,
     zero_point: bool,
     full_weight_grad: bool = False,
+    swiglu_limit: float = 0.0,
 ):
     """
     Create an SFT wrapper based on the method.
@@ -459,4 +461,5 @@ def _create_sft_wrapper(
         group_size=group_size,
         zero_point=zero_point,
         full_weight_grad=full_weight_grad,
+        swiglu_limit=swiglu_limit,
     )

kt-kernel/python/sft/__init__.py

@@ -24,6 +24,8 @@ from .backend import (
     get_int8_runtime,
     get_rawint4_checkpoint_contract,
     get_rawint4_runtime,
+    MXFP4_SFT_METHOD,
+    get_mxfp4_runtime,
 )
 from .weight_manifest import validate_persistent_int8_weights
 from .artifacts import (
@@ -125,6 +127,8 @@ __all__ = [
     "RAWINT4_WEIGHT_LAYOUT",
     "get_rawint4_checkpoint_contract",
     "get_rawint4_runtime",
+    "MXFP4_SFT_METHOD",
+    "get_mxfp4_runtime",
     "validate_persistent_int8_weights",
     "FUSED_EXPERT_LORA_NAME",
     "KT_ADAPTER_MANIFEST_NAME",

kt-kernel/python/sft/amx.py

@@ -10,6 +10,7 @@ from __future__ import annotations
 
 import ctypes
 import logging
+import math
 import os
 import glob as _glob
 import torch
@@ -19,7 +20,7 @@ from kt_kernel_ext.moe import MOESFTConfig
 
 logger = logging.getLogger(__name__)
 
-from ..utils.loader import BF16SafeTensorLoader, SafeTensorLoader
+from ..utils.loader import BF16SafeTensorLoader, MXFP4SafeTensorLoader, SafeTensorLoader
 
 try:
     from kt_kernel_ext.moe import (
@@ -51,8 +52,13 @@ try:
 except (ImportError, AttributeError):
     AMXInt4_KGroup_SFT_MOE = None
 
+try:
+    from kt_kernel_ext.moe import MXFP4_SFT_MOE
+except (ImportError, AttributeError):
+    MXFP4_SFT_MOE = None
+
 from .base import BaseSFTMoEWrapper, KExpertsSFTBuffer, _supports_authoritative_optimizer_grads
-from .backend import is_fp8_sft_method, is_int8_sft_method, is_rawint4_sft_method
+from .backend import get_mxfp4_runtime, is_fp8_sft_method, is_int8_sft_method, is_mxfp4_sft_method, is_rawint4_sft_method
 from .weights import BlockFP8ExpertWeights, RAWINT4ExpertWeights
 
 _AMX_M_STEP = 32
@@ -62,6 +68,7 @@ _AMX_M_STEP = 32
 _SFT_METHOD_TO_CLASS = {
     "AMXBF16_SFT": AMXBF16_SFT_MOE,
     "AMXFP8_SFT": AMXFP8_SFT_MOE,
+    "MXFP4_SFT": MXFP4_SFT_MOE,
     "INT8_SFT": AMXInt8_SFT_MOE,
     "AMXINT8_SFT": AMXInt8_SFT_MOE,
     "AMXINT4_SFT": AMXInt4_SFT_MOE,
@@ -77,7 +84,7 @@ class AMXSFTMoEWrapper(BaseSFTMoEWrapper):
     """
     AMX-based SFT MoE wrapper.
 
-    Supports BF16, native block-FP8, INT8, INT4, and SkipLoRA variants.
+    Supports BF16, native block-FP8/MXFP4, INT8, INT4, and SkipLoRA variants.
     Forward/backward buffer management is in BaseSFTMoEWrapper;
     this class implements weight loading and C++ task construction.
     """
@@ -102,12 +109,15 @@ class AMXSFTMoEWrapper(BaseSFTMoEWrapper):
         group_size: int = 128,
         zero_point: bool = True,
         full_weight_grad: bool = False,
+        swiglu_limit: float = 0.0,
     ):
-        if not _HAS_AMX_SFT_SUPPORT:
+        if not _HAS_AMX_SFT_SUPPORT and not is_mxfp4_sft_method(method):
             raise RuntimeError(
                 "AMX SFT backend not available. kt_kernel_ext was not compiled with AMX SFT support.\n"
                 "Please recompile with AMX SFT enabled."
             )
+        if is_mxfp4_sft_method(method):
+            get_mxfp4_runtime()
 
         super().__init__(
             layer_idx=layer_idx,
@@ -138,6 +148,35 @@ class AMXSFTMoEWrapper(BaseSFTMoEWrapper):
         )
         self.group_size = group_size
         self.zero_point = zero_point
+        self.swiglu_limit = float(swiglu_limit)
+
+        if not is_mxfp4_sft_method(method) and self.swiglu_limit != 0.0:
+            raise ValueError(
+                f"swiglu_limit is only supported by MXFP4_SFT, got {method!r}"
+            )
+        if is_mxfp4_sft_method(method):
+            if self._full_weight_grad or self.lora_rank <= 0:
+                raise ValueError("MXFP4_SFT supports frozen-base LoRA only")
+            if self.num_gpu_experts != 0:
+                raise ValueError("MXFP4_SFT requires all routed experts on CPU")
+            if self.group_size != 32 or self.zero_point:
+                raise ValueError("MXFP4_SFT requires group_size=32 and zero_point=False")
+            if self.hidden_size % 32 or self.moe_intermediate_size % 32:
+                raise ValueError(
+                    "MXFP4_SFT requires hidden_size and moe_intermediate_size divisible by 32"
+                )
+            if self.threadpool_count not in (1, 2):
+                raise ValueError("MXFP4_SFT currently supports threadpool_count 1 or 2")
+            if self.moe_intermediate_size % self.threadpool_count:
+                raise ValueError(
+                    "MXFP4_SFT requires moe_intermediate_size divisible by threadpool_count"
+                )
+            if (self.moe_intermediate_size // self.threadpool_count) % 32:
+                raise ValueError(
+                    "MXFP4_SFT requires every TP intermediate slice divisible by 32"
+                )
+            if not math.isfinite(self.swiglu_limit) or self.swiglu_limit <= 0.0:
+                raise ValueError("MXFP4_SFT requires a finite positive swiglu_limit")
 
         if is_fp8_sft_method(method):
             if self._full_weight_grad or self.lora_rank <= 0:
@@ -276,6 +315,8 @@ class AMXSFTMoEWrapper(BaseSFTMoEWrapper):
 
     def load_weights(self, physical_to_logical_map_cpu: torch.Tensor) -> None:
         if self._weights_loaded:
+            if is_mxfp4_sft_method(self.method):
+                raise RuntimeError("MXFP4_SFT weights are immutable and already loaded")
             return
 
         if physical_to_logical_map_cpu is None:
@@ -310,6 +351,13 @@ class AMXSFTMoEWrapper(BaseSFTMoEWrapper):
                 "physical_to_logical_map_cpu must be a permutation of "
                 f"[0, {self.num_experts}); got {self._physical_to_logical_map_cpu.tolist()}."
             )
+        if is_mxfp4_sft_method(self.method):
+            logical_experts = self._physical_to_logical_map_cpu[: self.num_experts].tolist()
+            if sorted(logical_experts) != list(range(self.num_experts)):
+                raise ValueError(
+                    "MXFP4_SFT physical_to_logical_map_cpu must be a permutation of "
+                    f"[0, {self.num_experts})."
+                )
 
         if self.gate_proj is None and not getattr(self, "_use_projs_path", False):
             self._load_base_weights_from_file()
@@ -330,6 +378,7 @@ class AMXSFTMoEWrapper(BaseSFTMoEWrapper):
         config.full_weight_grad = self._full_weight_grad
         config.authoritative_optimizer_grads = self._uses_authoritative_optimizer_grads
         config.physical_to_logical_map = self._physical_to_logical_map_cpu.data_ptr()
+        config.swiglu_limit = self.swiglu_limit
 
         if getattr(self, "_use_kt_direct_load", False):
             config.load = True
@@ -379,6 +428,10 @@ class AMXSFTMoEWrapper(BaseSFTMoEWrapper):
             config.quant_config.bits = 8
             config.quant_config.group_size = 128
             config.quant_config.zero_point = False
+        elif is_mxfp4_sft_method(self.method):
+            config.quant_config.bits = 4
+            config.quant_config.group_size = 32
+            config.quant_config.zero_point = False
 
         # Release old C++ MOE object before creating a new one to avoid memory leak
         old_moe = getattr(self, "moe", None)
@@ -462,6 +515,11 @@ class AMXSFTMoEWrapper(BaseSFTMoEWrapper):
                 "AMXFP8_SFT accepts raw per-expert E4M3 weights only; "
                 "use load_block_fp8_weights()"
             )
+        if is_mxfp4_sft_method(self.method):
+            raise ValueError(
+                "MXFP4_SFT accepts native packed E2M1 weights and UE8M0 scales only; "
+                "use load_mxfp4_weights()"
+            )
         if is_int8_sft_method(self.method):
             raise ValueError(
                 "INT8_SFT accepts pre-quantized .kt weights only; "
@@ -600,6 +658,96 @@ class AMXSFTMoEWrapper(BaseSFTMoEWrapper):
         self.down_proj = None
         self.load_weights(physical_to_logical_map_cpu)
 
+    def _stage_mxfp4_weights(self, experts_data: dict[str, List[torch.Tensor]]) -> None:
+        """Validate native MXFP4 tensors and retain them until C++ load completes."""
+
+        if not is_mxfp4_sft_method(self.method):
+            raise ValueError(f"MXFP4 weights require MXFP4_SFT, got {self.method!r}")
+        if not isinstance(experts_data, dict):
+            raise TypeError(
+                "MXFP4 expert weights must be the dict returned by "
+                "MXFP4SafeTensorLoader.load_experts()"
+            )
+
+        expected = {
+            "gate": (torch.uint8, (self.moe_intermediate_size, self.hidden_size // 2)),
+            "up": (torch.uint8, (self.moe_intermediate_size, self.hidden_size // 2)),
+            "down": (torch.uint8, (self.hidden_size, self.moe_intermediate_size // 2)),
+            "gate_scale": (
+                torch.bfloat16,
+                (self.moe_intermediate_size, self.hidden_size // 32),
+            ),
+            "up_scale": (
+                torch.bfloat16,
+                (self.moe_intermediate_size, self.hidden_size // 32),
+            ),
+            "down_scale": (
+                torch.bfloat16,
+                (self.hidden_size, self.moe_intermediate_size // 32),
+            ),
+        }
+        validated: dict[str, List[torch.Tensor]] = {}
+        for name, (dtype, shape) in expected.items():
+            tensors = experts_data.get(name)
+            if not isinstance(tensors, (list, tuple)) or len(tensors) != self.num_experts:
+                actual = len(tensors) if isinstance(tensors, (list, tuple)) else type(tensors).__name__
+                raise ValueError(
+                    f"MXFP4 {name} must contain {self.num_experts} expert tensors, got {actual}"
+                )
+            checked = []
+            for expert_idx, tensor in enumerate(tensors):
+                if not isinstance(tensor, torch.Tensor):
+                    raise TypeError(
+                        f"MXFP4 {name}[{expert_idx}] must be a torch.Tensor, got {type(tensor)!r}"
+                    )
+                if (
+                    tensor.device.type != "cpu"
+                    or tensor.dtype != dtype
+                    or tuple(tensor.shape) != shape
+                    or not tensor.is_contiguous()
+                ):
+                    raise ValueError(
+                        f"MXFP4 {name}[{expert_idx}] must be contiguous CPU {dtype} "
+                        f"with shape {shape}, got {tensor.dtype} {tuple(tensor.shape)} "
+                        f"on {tensor.device}"
+                    )
+                if name.endswith("_scale") and not torch.isfinite(tensor).all().item():
+                    raise ValueError(
+                        f"MXFP4 {name}[{expert_idx}] contains a non-finite scale"
+                    )
+                checked.append(tensor)
+            validated[name] = checked
+
+        self._gate_weights_per_numa = [validated["gate"]]
+        self._up_weights_per_numa = [validated["up"]]
+        self._down_weights_per_numa = [validated["down"]]
+        self._gate_scales_per_numa = [validated["gate_scale"]]
+        self._up_scales_per_numa = [validated["up_scale"]]
+        self._down_scales_per_numa = [validated["down_scale"]]
+        self._gate_projs_ptrs = [[tensor.data_ptr() for tensor in validated["gate"]]]
+        self._up_projs_ptrs = [[tensor.data_ptr() for tensor in validated["up"]]]
+        self._down_projs_ptrs = [[tensor.data_ptr() for tensor in validated["down"]]]
+        self._gate_scale_ptrs = [[tensor.data_ptr() for tensor in validated["gate_scale"]]]
+        self._up_scale_ptrs = [[tensor.data_ptr() for tensor in validated["up_scale"]]]
+        self._down_scale_ptrs = [[tensor.data_ptr() for tensor in validated["down_scale"]]]
+        self._has_bwd_projs = False
+        self._use_projs_path = True
+        self.gate_proj = None
+        self.up_proj = None
+        self.down_proj = None
+
+    def load_mxfp4_weights(
+        self,
+        experts_data: dict[str, List[torch.Tensor]],
+        physical_to_logical_map_cpu: torch.Tensor,
+    ) -> None:
+        """Load native E2M1/UE8M0 expert tensors without dequantizing the base."""
+
+        if self._weights_loaded:
+            raise RuntimeError("MXFP4_SFT weights are immutable and already loaded")
+        self._stage_mxfp4_weights(experts_data)
+        self.load_weights(physical_to_logical_map_cpu)
+
     def _load_base_weights_from_file(self) -> None:
         if not hasattr(self, "weight_path") or self.weight_path is None:
             raise RuntimeError(
@@ -607,6 +755,19 @@ class AMXSFTMoEWrapper(BaseSFTMoEWrapper):
                 "Either set weight_path or call load_weights_from_tensors() instead."
             )
 
+        if is_mxfp4_sft_method(self.method):
+            loader = MXFP4SafeTensorLoader(self.weight_path)
+            try:
+                experts_data = loader.load_experts(
+                    f"model.layers.{self.layer_idx}",
+                    device="cpu",
+                    reject_non_finite_scales=True,
+                )
+                self._stage_mxfp4_weights(experts_data)
+            finally:
+                loader.close_all_handles()
+            return
+
         kt_layer_dir = os.path.join(self.weight_path, f"_layer_{self.layer_idx}")
         if os.path.isdir(kt_layer_dir):
             kt_files = _glob.glob(os.path.join(kt_layer_dir, "_numa_0", "*.kt"))
@@ -837,6 +998,18 @@ class AMXSFTMoEWrapper(BaseSFTMoEWrapper):
         self.grad_down_lora_a = grad_down_lora_a
         self.grad_down_lora_b = grad_down_lora_b
 
+        if getattr(self, "_uses_authoritative_optimizer_grads", False):
+            parameter_count = sum(isinstance(tensor, torch.nn.Parameter) for tensor in provided.values())
+            if parameter_count not in (0, len(provided)):
+                raise ValueError(
+                    "authoritative LoRA ownership requires all six weights to be nn.Parameter objects or none"
+                )
+            if parameter_count:
+                for name, parameter in provided.items():
+                    self.register_authoritative_optimizer_grad(
+                        f"lora.{name}", parameter, grad_provided[f"grad_{name}"]
+                    )
+
         self._lora_initialized = True
 
         if self._weights_loaded and self.moe is not None:

kt-kernel/python/sft/arch.py

@@ -77,6 +77,19 @@ def get_moe_arch_config(config) -> MOEArchConfig:
     model_type = str(getattr(config, "model_type", ""))
     text_model_type = str(getattr(cfg, "model_type", ""))
 
+    if "DeepseekV4" in arch or model_type == "deepseek_v4":
+        return MOEArchConfig(
+            moe_layer_attr="mlp",
+            router_attr="gate",
+            experts_attr="experts",
+            weight_names=("gate_proj", "up_proj", "down_proj"),
+            expert_num=cfg.n_routed_experts,
+            intermediate_size=cfg.moe_intermediate_size,
+            num_experts_per_tok=cfg.num_experts_per_tok,
+            has_shared_experts=True,
+            router_type="precomputed",
+        )
+
     if (
         "KimiK2" in arch
         or model_type in {"kimi_k2", "kimi_k25", "kimi_k2_5", "kimi_k26", "kimi_k2_6"}
@@ -274,6 +287,8 @@ def move_non_experts_to_gpu(
         container.embed_tokens.to(device)
     if hasattr(container, "norm"):
         container.norm.to(device)
+    if hasattr(container, "hc_head"):
+        container.hc_head.to(device)
     output_embeddings = getattr(model, "get_output_embeddings", lambda: None)()
     if output_embeddings is None:
         output_embeddings = getattr(model, "lm_head", None)
@@ -281,6 +296,9 @@ def move_non_experts_to_gpu(
         output_embeddings.to(device)
 
     for layer in layers:
+        for name in ("attn_hc", "ffn_hc"):
+            if hasattr(layer, name):
+                getattr(layer, name).to(device)
         if hasattr(layer, "self_attn"):
             layer.self_attn.to(device)
 

kt-kernel/python/sft/artifacts.py

@@ -64,6 +64,7 @@ _RUNTIME_TENSOR_CONTRACTS = "_kt_routed_expert_runtime_tensor_contracts"
 _SUPPORTED_MOE_ARCHITECTURES = (
     "DeepseekV2",
     "DeepseekV3",
+    "DeepseekV4",
     "KimiK2",
     "Qwen2Moe",
     "Qwen3Moe",
@@ -71,7 +72,7 @@ _SUPPORTED_MOE_ARCHITECTURES = (
     "Glm4Moe",
     "Mixtral",
 )
-_EXPERT_WEIGHT_FORMATS = frozenset({"bf16", "int8", "fp8", "rawint4"})
+_EXPERT_WEIGHT_FORMATS = frozenset({"bf16", "int8", "fp8", "rawint4", "mxfp4"})
 _RAWINT4_BASE_IDENTITY_SCHEMA = "safetensors-sharded-sampled-v1"
 _RAWINT4_SAMPLE_BYTES = 64 * 1024
 _RAWINT4_SAMPLE_COUNT = 8
@@ -791,6 +792,7 @@ def _config_expert_weight_format(config: Any) -> str:
         "amxfp8": "fp8",
         "rawint4": "rawint4",
         "amxint4_kgroup": "rawint4",
+        "mxfp4": "mxfp4",
     }.get(backend.strip().lower(), "")
 
 
@@ -816,7 +818,7 @@ def is_kt_supported_moe_model(model: Any) -> bool:
         str(getattr(config, "model_type", "")),
         str(getattr(text_config, "model_type", "")),
     }
-    return bool(model_types.intersection({"kimi_k2", "kimi_k25", "kimi_k2_5", "kimi_k26", "kimi_k2_6"}))
+    return bool(model_types.intersection({"kimi_k2", "kimi_k25", "kimi_k2_5", "kimi_k26", "kimi_k2_6", "deepseek_v4"}))
 
 
 def should_disable_kt_source_quantizer(
@@ -832,7 +834,19 @@ def should_disable_kt_source_quantizer(
     nested checkpoint metadata remains available for KT's strict validation.
     """
 
-    rawint4_requested = _config_expert_weight_format(kt_config) == "rawint4"
+    weight_format = _config_expert_weight_format(kt_config)
+    if weight_format == "mxfp4":
+        if explicit_quantization_config is not None:
+            raise KTArtifactError("MXFP4 SFT cannot use an explicit framework quantizer")
+        if getattr(model_config, "model_type", None) != "deepseek_v4":
+            raise KTArtifactError("native MXFP4 SFT currently requires DeepSeek V4")
+        cache_path = _config_value(kt_config, "kt_non_expert_weight_path")
+        if not cache_path:
+            raise KTArtifactError("Set kt_non_expert_weight_path for the automatic V4 BF16 non-expert cache")
+        # Config loading precedes the pretrained artifact resolver, which owns
+        # first-use preparation and full validation before any tensor is loaded.
+        return True
+    rawint4_requested = weight_format == "rawint4"
     skip_loading = bool(_config_value(kt_config, "kt_skip_expert_loading", True))
     if not rawint4_requested or not skip_loading:
         return False
@@ -898,7 +912,7 @@ def validate_kt_prequantized_loading_info(
 
     weight_format = _config_expert_weight_format(kt_config)
     skip_loading = _config_value(kt_config, "kt_skip_expert_loading", True)
-    if weight_format not in {"int8", "fp8", "rawint4"} or not bool(skip_loading):
+    if weight_format not in {"int8", "fp8", "rawint4", "mxfp4"} or not bool(skip_loading):
         return
     config = getattr(model, "config", None)
     text_config = getattr(config, "text_config", None)
@@ -961,6 +975,10 @@ def resolve_kt_pretrained_artifacts(
     ``disable_source_quantizer`` on the returned plan instead.
     """
 
+    if _config_expert_weight_format(kt_config) == "mxfp4":
+        from .deepseek_v4 import resolve_native_load_plan
+
+        return resolve_native_load_plan(kt_config, pretrained_model_name_or_path, quantization_config)
     cache_path = _config_value(kt_config, "kt_non_expert_weight_path")
     if not cache_path:
         return None
@@ -1246,6 +1264,8 @@ def _enumerate_runtime_routed_modules(model: Any) -> tuple[tuple[str, Any], ...]
             raise KTArtifactError(f"missing routed-expert subtree {path!r}") from exc
         if registered is not experts:
             raise KTArtifactError(f"routed-expert subtree {path!r} does not preserve module identity")
+        if getattr(experts, "_is_kt_routed_experts_wrapper", False):
+            experts = getattr(experts, experts._experts_attr)
         if id(experts) in identities:
             raise KTArtifactError(f"routed-expert subtree {path!r} shares a module with another layer")
         identities.add(id(experts))
@@ -1838,9 +1858,9 @@ def _adapter_provenance(model: Any, plan: KTPretrainedLoadPlan | None) -> dict[s
             f"fused wrapper LoRA alpha {float(alpha)} does not match the pretrained plan alpha {plan.lora_alpha}"
         )
     expert_weight_format = _runtime_expert_weight_format(model, plan)
-    if plan is not None and expert_weight_format != "int8":
+    if plan is not None and expert_weight_format not in {"int8", "mxfp4"}:
         raise KTArtifactError(
-            "a KT non-expert load plan requires INT8 routed expert provenance"
+            "a KT non-expert load plan requires INT8 or MXFP4 routed expert provenance"
         )
     base_model_name = _base_model_name(model, plan)
     payload: dict[str, Any] = {
@@ -1868,7 +1888,8 @@ def _adapter_provenance(model: Any, plan: KTPretrainedLoadPlan | None) -> dict[s
             "path": plan.weight_path,
             "fingerprint": cache_fingerprint,
         }
-        payload["int8_experts"] = {
+        routed_field = "native_experts" if expert_weight_format == "mxfp4" else "int8_experts"
+        payload[routed_field] = {
             "path": plan.routed_weight_path,
             "manifest": routed_path.name,
             "manifest_sha256": _sha256_file(routed_path),
@@ -1976,7 +1997,7 @@ def _validate_adapter_manifest(model: Any, adapter_path: Path) -> KTAdapterManif
         raise KTArtifactError(f"{manifest_path}: base does not match the runtime")
     if payload.get("lora") != expected_provenance["lora"]:
         raise KTArtifactError(f"{manifest_path}: lora does not match the runtime")
-    for field in ("non_expert_cache", "int8_experts", "expert_quantization"):
+    for field in ("non_expert_cache", "int8_experts", "native_experts", "expert_quantization"):
         if payload.get(field) != expected_provenance.get(field):
             raise KTArtifactError(f"{manifest_path}: {field} does not match the runtime")
     artifacts = payload.get("artifacts")

kt-kernel/python/sft/backend.py

@@ -24,6 +24,7 @@ RAWINT4_KERNEL = "amx-int4-kgroup-g32"
 RAWINT4_GROUP_SIZE = 32
 _RAWINT4_SFT_METHOD_ALIASES = frozenset({RAWINT4_SFT_METHOD, "AMXINT4_KGroup_SFT"})
 
+MXFP4_BACKEND = "MXFP4"
 
 def is_int8_sft_method(method: str) -> bool:
     return str(method) in _INT8_SFT_METHOD_ALIASES
@@ -58,6 +59,8 @@ def normalize_sft_backend(
             return FP8_BACKEND
         if expert_weight_format == "rawint4":
             return RAWINT4_BACKEND
+        if expert_weight_format == "mxfp4":
+            return MXFP4_BACKEND
         return "AMXBF16"
     if normalized == "int8":
         return INT8_BACKEND
@@ -80,6 +83,8 @@ def normalize_sft_backend(
         return FP8_BACKEND
     if normalized == "rawint4":
         return RAWINT4_BACKEND
+    if normalized == "mxfp4":
+        return MXFP4_BACKEND
     if normalized == "amxint4_kgroup":
         warnings.warn(
             "kt_backend='AMXINT4_KGroup' is deprecated; use " "kt_backend='auto' or kt_backend='RAWINT4'",
@@ -313,3 +318,60 @@ def get_rawint4_runtime() -> RAWINT4Runtime:
         kernel=kernel,
         weight_layout=layout,
     )
+
+
+MXFP4_SFT_METHOD = "MXFP4_SFT"
+
+
+MXFP4_WEIGHT_LAYOUT = "mxfp4-e2m1-ue8m0-g32"
+
+
+MXFP4_KERNEL = "avx512-bf16-mxfp4-g32"
+
+
+MXFP4_AVX2_KERNEL = "avx2-mxfp4-g32"
+
+
+_MXFP4_SFT_METHOD_ALIASES = frozenset({MXFP4_SFT_METHOD})
+
+
+def is_mxfp4_sft_method(method: str) -> bool:
+    return str(method) in _MXFP4_SFT_METHOD_ALIASES
+
+
+@dataclass(frozen=True)
+class MXFP4Runtime:
+    cpu_variant: str
+    kernel: str
+    weight_layout: str
+
+
+def get_mxfp4_runtime() -> MXFP4Runtime:
+    """Validate and describe native group-32 MXFP4 routed-expert SFT."""
+
+    import kt_kernel
+
+    extension = kt_kernel.kt_kernel_ext
+    cpu_variant = str(getattr(extension, "__cpu_variant__", "")).lower()
+    # Two tiers ship MXFP4_SFT_MOE with the same Python contract: the AMX /
+    # AVX512-BF16 kernel and the AVX2 (FMA + F16C) kernel for Zen 2/3-class hosts.
+    if cpu_variant in {"amx", "avx512_bf16"}:
+        kernel = MXFP4_KERNEL
+    elif cpu_variant == "avx2":
+        kernel = MXFP4_AVX2_KERNEL
+    else:
+        raise RuntimeError(
+            "MXFP4 SFT requires an AVX512-BF16 or AVX2 extension; "
+            f"loaded cpu_variant={cpu_variant or 'unknown'!r}"
+        )
+    moe_extension = getattr(extension, "moe", None)
+    if moe_extension is None or not hasattr(moe_extension, "MXFP4_SFT_MOE"):
+        raise RuntimeError(
+            "The loaded kt-kernel extension does not provide MXFP4 SFT. "
+            "Install a wheel containing MXFP4_SFT_MOE."
+        )
+    return MXFP4Runtime(
+        cpu_variant=cpu_variant,
+        kernel=kernel,
+        weight_layout=MXFP4_WEIGHT_LAYOUT,
+    )

kt-kernel/python/sft/base.py

@@ -19,7 +19,7 @@ from typing import Optional, Tuple
 from abc import ABC, abstractmethod
 
 from ..experts_base import KExpertsCPUBuffer, _MoEBase
-from .backend import is_fp8_sft_method, is_int8_sft_method, is_rawint4_sft_method
+from .backend import is_fp8_sft_method, is_int8_sft_method, is_rawint4_sft_method, is_mxfp4_sft_method
 
 
 def _supports_authoritative_optimizer_grads(
@@ -32,14 +32,14 @@ def _supports_authoritative_optimizer_grads(
     """Whether this SFT configuration can use C++-authoritative gradients.
 
     BF16 supports both base and LoRA authoritative gradients. Quantized base
-    weights are frozen, so INT8/FP8 support the same lifecycle only for pure LoRA.
+    weights are frozen, so INT8/FP8/MXFP4 support the lifecycle only for pure LoRA.
     """
     if int(num_gpu_experts) != 0:
         return False
     if method == "AMXBF16_SFT":
         return True
     return (
-        (is_int8_sft_method(method) or is_fp8_sft_method(method) or is_rawint4_sft_method(method))
+        (is_int8_sft_method(method) or is_fp8_sft_method(method) or is_rawint4_sft_method(method) or is_mxfp4_sft_method(method))
         and not bool(full_weight_grad)
         and int(lora_rank) > 0
     )
@@ -269,6 +269,7 @@ class BaseSFTMoEWrapper(_MoEBase, ABC):
         self._checkpoint_output_cpu: Optional[torch.Tensor] = None
         self._checkpoint_output_qlen: int = 0
         self._backward_repack_pending: bool = False
+        self._inference_inflight: list[tuple[torch.cuda.Event, _SFTForwardBufferView]] = []
 
         self.moe = None
 
@@ -829,6 +830,12 @@ class BaseSFTMoEWrapper(_MoEBase, ABC):
             return
 
         self._validate_forward_inputs(hidden_states, expert_ids, weights)
+        if hasattr(self, "_pending_inference_buffer"):
+            raise RuntimeError("An inference forward is already pending; call sync_forward_inference() first.")
+        if not torch.cuda.is_current_stream_capturing():
+            self._inference_inflight = [
+                (event, buffer) for event, buffer in self._inference_inflight if not event.query()
+            ]
         flat_hidden_states = hidden_states.view(-1, hidden_states.shape[-1])
 
         (
@@ -863,6 +870,9 @@ class BaseSFTMoEWrapper(_MoEBase, ABC):
         )
 
         self._pending_inference_fallback = False
+        # The native task owns raw pointers, including a CPU-only batch-size scalar.
+        # Keep their tensors alive when another batch replaces the shared buffer.
+        self._pending_inference_buffer = buffer_view
         self._pending_inference_output_cpu = output_cpu[current_slot]
         self._pending_inference_output_gpu = output_gpu[current_slot]
 
@@ -897,8 +907,13 @@ class BaseSFTMoEWrapper(_MoEBase, ABC):
         self.cpu_infer.sync_with_cuda_stream(cuda_stream)
         with torch.cuda.stream(torch_stream):
             self._pending_inference_output_gpu.copy_(self._pending_inference_output_cpu, non_blocking=True)
+            self._pending_inference_output_gpu.record_stream(torch_stream)
+            completed = torch.cuda.Event()
+            completed.record(torch_stream)
+        self._inference_inflight.append((completed, self._pending_inference_buffer))
         output = self._pending_inference_output_gpu
 
+        del self._pending_inference_buffer
         del self._pending_inference_output_cpu
         del self._pending_inference_output_gpu
         return output

kt-kernel/python/sft/config.py

@@ -26,6 +26,7 @@ from typing import Any, Callable, Literal
 from .backend import (
     FP8_BACKEND,
     INT8_BACKEND,
+    MXFP4_BACKEND,
     RAWINT4_BACKEND,
     normalize_sft_backend,
 )
@@ -45,6 +46,7 @@ _KNOWN_SFT_BACKENDS = {
     "amxfp8",
     "amxint4",
     "rawint4",
+    "mxfp4",
     "amxint4_kgroup",
     "amxbf16_skiplora",
     "amxint8_skiplora",
@@ -52,7 +54,7 @@ _KNOWN_SFT_BACKENDS = {
 }
 
 ActivationRetention = Literal["retain", "recompute"]
-ExpertWeightFormat = Literal["bf16", "int8", "fp8", "rawint4"]
+ExpertWeightFormat = Literal["bf16", "int8", "fp8", "rawint4", "mxfp4"]
 WeightLifecycle = Literal["persistent", "ephemeral"]
 
 
@@ -269,7 +271,7 @@ class KTConfig:
     )
     kt_num_gpu_experts: int | None = None
     kt_skip_expert_loading: bool | None = None
-    kt_share_backward_bb: bool | None = None  # defaults off for RAWINT4, on otherwise
+    kt_share_backward_bb: bool | None = None  # defaults off for RAWINT4/MXFP4
     kt_share_cache_pool: bool | None = (
         None  # auto-set by trainer_config_process, not user-facing
     )
@@ -411,10 +413,10 @@ class KTConfig:
             self.kt_expert_weight_format = (
                 str(self.kt_expert_weight_format).strip().lower()
             )
-            if self.kt_expert_weight_format not in {"bf16", "int8", "fp8", "rawint4"}:
+            if self.kt_expert_weight_format not in {"bf16", "int8", "fp8", "rawint4", "mxfp4"}:
                 raise ValueError(
                     "kt_expert_weight_format must be one of "
-                    "['bf16', 'fp8', 'int8', 'rawint4'], "
+                    "['bf16', 'fp8', 'int8', 'rawint4', 'mxfp4'], "
                     f"got {self.kt_expert_weight_format!r}"
                 )
         if self.kt_weight_lifecycle is None:
@@ -469,7 +471,7 @@ class KTConfig:
         if self.kt_backend is None:
             if env_backend:
                 self.kt_backend = env_backend
-            elif self.kt_expert_weight_format in {"int8", "fp8", "rawint4"}:
+            elif self.kt_expert_weight_format in {"int8", "fp8", "rawint4", "mxfp4"}:
                 self.kt_backend = "auto"
             else:
                 self.kt_backend = "AMXBF16"
@@ -492,6 +494,8 @@ class KTConfig:
                 self.kt_expert_weight_format = "fp8"
             elif backend_lower == RAWINT4_BACKEND.lower():
                 self.kt_expert_weight_format = "rawint4"
+            elif backend_lower == MXFP4_BACKEND.lower():
+                self.kt_expert_weight_format = "mxfp4"
             elif backend_lower == "amxbf16":
                 self.kt_expert_weight_format = "bf16"
         expected_backend = {
@@ -499,6 +503,7 @@ class KTConfig:
             "int8": INT8_BACKEND.lower(),
             "fp8": FP8_BACKEND.lower(),
             "rawint4": RAWINT4_BACKEND.lower(),
+            "mxfp4": MXFP4_BACKEND.lower(),
         }.get(self.kt_expert_weight_format)
         if expected_backend is not None and backend_lower != expected_backend:
             source = "kt_backend" if explicit_backend else "ACCELERATE_KT_BACKEND"
@@ -529,7 +534,7 @@ class KTConfig:
         if self.kt_share_backward_bb is None:
             self.kt_share_backward_bb = _env_bool(
                 "ACCELERATE_KT_SHARE_BACKWARD_BB",
-                self.kt_expert_weight_format != "rawint4",
+                self.kt_expert_weight_format not in {"rawint4", "mxfp4"},
             )
         if self.kt_share_cache_pool is None:
             self.kt_share_cache_pool = False
@@ -571,11 +576,30 @@ class KTConfig:
                     "ACCELERATE_KT_SKIP_EXPERT_LOADING", True
                 )
 
-        if self.kt_non_expert_weight_path and self.kt_expert_weight_format != "int8":
+        if self.kt_non_expert_weight_path and self.kt_expert_weight_format not in {"int8", "mxfp4"}:
             raise ValueError(
                 "kt_non_expert_weight_path is supported only with "
-                "kt_expert_weight_format='int8'"
+                "kt_expert_weight_format='int8' or 'mxfp4'"
             )
+        if self.kt_expert_weight_format == "mxfp4":
+            if self.kt_skip_expert_loading is None:
+                self.kt_skip_expert_loading = True
+            if self.kt_train_mode != "lora" or self.kt_full_weight_grad or int(self.kt_lora_rank or 0) <= 0:
+                raise ValueError("MXFP4 SFT supports frozen-base LoRA only with kt_lora_rank > 0")
+            if self.kt_num_gpu_experts or self.kt_use_lora_experts:
+                raise ValueError("MXFP4 SFT requires all routed experts and their LoRA on CPU")
+            if self.kt_share_backward_bb:
+                raise ValueError("MXFP4 SFT is transpose-free; kt_share_backward_bb must be false")
+            if self.kt_lora_dropout != 0.0:
+                raise ValueError("MXFP4 SFT requires kt_lora_dropout=0")
+            if self.kt_weight_lifecycle != "persistent" or self.kt_expert_checkpoint_path:
+                raise ValueError("MXFP4 SFT requires a persistent native checkpoint at kt_weight_path")
+            if not self.kt_weight_path:
+                raise ValueError("MXFP4 SFT requires kt_weight_path")
+            if not self.kt_skip_expert_loading:
+                raise ValueError("MXFP4 SFT requires kt_skip_expert_loading=true")
+            if (int(self.kt_threadpool_count or 1) if self.kt_tp_enabled else 1) not in (1, 2):
+                raise ValueError("MXFP4 SFT supports CPU TP1/TP2 only")
         if self.kt_expert_weight_format == "int8":
             if str(self.kt_backend).lower() != INT8_BACKEND.lower():
                 raise ValueError(

kt-kernel/python/sft/convert_kt_to_sglang_adapter.py

@@ -0,0 +1,477 @@
+#!/usr/bin/env python3
+"""Convert KT fused expert LoRA checkpoints into an SGLang adapter directory."""
+
+from __future__ import annotations
+
+import argparse
+import json
+import os
+import re
+import shutil
+from pathlib import Path
+from typing import Dict, Iterable, Mapping
+
+import torch
+from safetensors.torch import load_file, save_file
+
+
+FUSED_EXPERT_LORA_FILE = "fused_expert_lora.safetensors"
+ADAPTER_MODEL_FILE = "adapter_model.safetensors"
+ADAPTER_CONFIG_FILE = "adapter_config.json"
+
+KT_NAME_MAP = {
+    "gate_lora_a": ("gate_proj", "lora_A", 1),
+    "gate_lora_b": ("gate_proj", "lora_B", 2),
+    "up_lora_a": ("up_proj", "lora_A", 1),
+    "up_lora_b": ("up_proj", "lora_B", 2),
+    "down_lora_a": ("down_proj", "lora_A", 1),
+    "down_lora_b": ("down_proj", "lora_B", 2),
+}
+
+TARGET_MODULE_ORDER = [
+    "q_proj",
+    "k_proj",
+    "v_proj",
+    "o_proj",
+    "gate_proj",
+    "up_proj",
+    "down_proj",
+    "in_proj_qkv",
+    "in_proj_z",
+    "in_proj_b",
+    "in_proj_a",
+    "out_proj",
+    "embed_tokens",
+    "lm_head",
+]
+
+KT_FUSED_KEY_RE = re.compile(r"^layers\.(\d+)\.experts\.([^.]+)$")
+
+
+def _load_json(path: Path) -> dict:
+    with path.open("r", encoding="utf-8") as f:
+        return json.load(f)
+
+
+def _write_json(path: Path, data: Mapping) -> None:
+    with path.open("w", encoding="utf-8") as f:
+        json.dump(data, f, indent=2, sort_keys=True)
+        f.write("\n")
+
+
+def _clean_adapter_key(key: str) -> str:
+    """Match the existing SGLang converter's PEFT key cleanup."""
+    key = key.replace("base_model.model.", "")
+    key = key.replace(".orig_module", "")
+    return key
+
+
+def _ordered_target_modules(modules: Iterable[str]) -> list[str]:
+    seen = set(modules)
+    ordered = [name for name in TARGET_MODULE_ORDER if name in seen]
+    ordered.extend(sorted(seen.difference(ordered)))
+    return ordered
+
+
+def _infer_target_module_from_key(key: str) -> str | None:
+    if "lora_embedding_A" in key or "lora_embedding_B" in key:
+        if "embed_tokens" in key:
+            return "embed_tokens"
+        if "lm_head" in key or "unembed_tokens" in key:
+            return "lm_head"
+
+    marker = ".lora_"
+    if marker not in key:
+        return None
+    prefix = key.split(marker, 1)[0]
+    if "." not in prefix:
+        return prefix
+    return prefix.rsplit(".", 1)[-1]
+
+
+def _merge_tensor(tensors: Dict[str, torch.Tensor], key: str, value: torch.Tensor) -> None:
+    if key in tensors:
+        raise ValueError(f"Duplicate output tensor key: {key}")
+    tensors[key] = value.detach().cpu()
+
+
+def _load_existing_adapter(input_dir: Path) -> tuple[dict[str, torch.Tensor], set[str]]:
+    adapter_path = input_dir / ADAPTER_MODEL_FILE
+    if not adapter_path.exists():
+        return {}, set()
+
+    tensors: dict[str, torch.Tensor] = {}
+    target_modules: set[str] = set()
+    for key, value in load_file(str(adapter_path)).items():
+        cleaned_key = _clean_adapter_key(key)
+        _merge_tensor(tensors, cleaned_key, value)
+        target_module = _infer_target_module_from_key(cleaned_key)
+        if target_module is not None:
+            target_modules.add(target_module)
+    return tensors, target_modules
+
+
+def _convert_fused_expert_lora(
+    fused_path: Path,
+) -> tuple[dict[str, torch.Tensor], int, set[str]]:
+    if not fused_path.exists():
+        raise FileNotFoundError(f"Missing {FUSED_EXPERT_LORA_FILE}: {fused_path}")
+
+    output: dict[str, torch.Tensor] = {}
+    ranks: set[int] = set()
+    expert_counts: set[int] = set()
+    target_modules: set[str] = set()
+
+    for key, tensor in sorted(load_file(str(fused_path)).items()):
+        match = KT_FUSED_KEY_RE.match(key)
+        if match is None:
+            raise ValueError(f"Unexpected key in {FUSED_EXPERT_LORA_FILE}: {key}")
+
+        layer_idx, kt_name = match.groups()
+        if kt_name not in KT_NAME_MAP:
+            raise ValueError(f"Unsupported KT fused expert LoRA tensor: {key}")
+        if tensor.dim() != 3:
+            raise ValueError(f"{key} must be 3D [E, ...], got shape {tuple(tensor.shape)}")
+
+        proj_name, lora_name, rank_dim = KT_NAME_MAP[kt_name]
+        expert_count = int(tensor.shape[0])
+        rank = int(tensor.shape[rank_dim])
+        expert_counts.add(expert_count)
+        ranks.add(rank)
+        target_modules.add(proj_name)
+
+        for expert_idx in range(expert_count):
+            output_key = (
+                f"model.layers.{layer_idx}.mlp.experts.{expert_idx}."
+                f"{proj_name}.{lora_name}.weight"
+            )
+            _merge_tensor(output, output_key, tensor[expert_idx].contiguous())
+
+    if not output:
+        raise ValueError(f"No tensors found in {fused_path}")
+    if len(expert_counts) != 1:
+        raise ValueError(f"Inconsistent expert counts in {FUSED_EXPERT_LORA_FILE}: {sorted(expert_counts)}")
+    if len(ranks) != 1:
+        raise ValueError(f"Inconsistent LoRA ranks in {FUSED_EXPERT_LORA_FILE}: {sorted(ranks)}")
+
+    return output, next(iter(ranks)), target_modules
+
+
+def _build_adapter_config(
+    input_dir: Path,
+    rank: int,
+    target_modules: set[str],
+    base_model_name_or_path: str,
+    lora_alpha: float | None,
+    *,
+    include_input_target_modules: bool = True,
+) -> dict:
+    config_path = input_dir / ADAPTER_CONFIG_FILE
+    config = _load_json(config_path) if config_path.exists() else {}
+
+    if "lora_alpha" in config:
+        final_alpha = config["lora_alpha"]
+    elif lora_alpha is not None:
+        final_alpha = lora_alpha
+    else:
+        raise ValueError(
+            f"No {ADAPTER_CONFIG_FILE} with lora_alpha found in {input_dir}; "
+            "pass --lora-alpha to preserve runtime scaling."
+        )
+
+    existing_targets = config.get("target_modules", [])
+    if include_input_target_modules and isinstance(existing_targets, list):
+        target_modules.update(str(name).split(".")[-1] for name in existing_targets)
+
+    config["peft_type"] = config.get("peft_type", "LORA")
+    config["r"] = rank
+    config["lora_alpha"] = final_alpha
+    config["target_modules"] = _ordered_target_modules(target_modules)
+    config["bias"] = config.get("bias", "none")
+    config["task_type"] = config.get("task_type", "CAUSAL_LM")
+    config["base_model_name_or_path"] = base_model_name_or_path
+
+    return config
+
+
+def _paths_have_ancestor_relationship(left: Path, right: Path) -> bool:
+    if left == right:
+        return True
+    try:
+        left.relative_to(right)
+        return True
+    except ValueError:
+        pass
+    try:
+        right.relative_to(left)
+        return True
+    except ValueError:
+        return False
+
+
+def _validate_no_ancestor_paths(
+    paths: Iterable[Path],
+    *,
+    label: str,
+) -> None:
+    resolved = list(paths)
+    for i, left in enumerate(resolved):
+        for right in resolved[i + 1 :]:
+            if _paths_have_ancestor_relationship(left, right):
+                raise ValueError(
+                    f"{label} cannot have ancestor/descendant relationships: "
+                    f"{left} and {right}."
+                )
+
+
+def _prepare_output_dir(output_path: Path, input_path: Path, overwrite: bool) -> None:
+    _validate_output_dir(output_path, input_path, overwrite)
+    if output_path.exists() and any(output_path.iterdir()):
+        shutil.rmtree(output_path)
+    output_path.mkdir(parents=True, exist_ok=True)
+
+
+def _validate_output_dir(output_path: Path, input_path: Path, overwrite: bool) -> None:
+    if output_path == input_path:
+        raise ValueError("Output directory must be different from input directory.")
+    if _paths_have_ancestor_relationship(output_path, input_path):
+        raise ValueError(
+            "Output and input directories cannot be ancestor/descendant of each other: "
+            f"output={output_path}, input={input_path}."
+        )
+    if output_path.exists() and not output_path.is_dir():
+        raise FileExistsError(f"Output path exists and is not a directory: {output_path}")
+    if output_path.exists() and any(output_path.iterdir()):
+        if not overwrite:
+            raise FileExistsError(f"Output directory is not empty: {output_path}")
+
+
+def _infer_lora_rank_from_tensor(key: str, tensor: torch.Tensor) -> int | None:
+    if ".lora_A." in key:
+        return int(tensor.shape[0])
+    if ".lora_B." in key:
+        return int(tensor.shape[1])
+    return None
+
+
+def _validate_nonexpert_rank(
+    existing_tensors: Mapping[str, torch.Tensor],
+    expert_rank: int,
+    input_dir: Path,
+) -> None:
+    if not existing_tensors:
+        return
+
+    config_path = input_dir / ADAPTER_CONFIG_FILE
+    if config_path.exists():
+        config_rank = _load_json(config_path).get("r")
+        if config_rank is not None and int(config_rank) != expert_rank:
+            raise ValueError(
+                f"Non-expert adapter rank mismatch: adapter_config.json r={config_rank}, "
+                f"but fused expert LoRA rank={expert_rank}."
+            )
+
+    for key, tensor in existing_tensors.items():
+        tensor_rank = _infer_lora_rank_from_tensor(key, tensor)
+        if tensor_rank is None:
+            continue
+        if tensor_rank != expert_rank:
+            raise ValueError(
+                f"Non-expert adapter tensor rank mismatch for {key}: "
+                f"tensor rank={tensor_rank}, fused expert LoRA rank={expert_rank}."
+            )
+
+
+def _write_adapter(
+    output_path: Path,
+    input_path: Path,
+    tensors: dict[str, torch.Tensor],
+    config: Mapping,
+    *,
+    overwrite: bool,
+) -> None:
+    _prepare_output_dir(output_path, input_path, overwrite)
+    save_file(tensors, str(output_path / ADAPTER_MODEL_FILE), metadata={"format": "pt"})
+    _write_json(output_path / ADAPTER_CONFIG_FILE, config)
+
+
+def convert_kt_to_sglang_adapter(
+    input_dir: str | os.PathLike,
+    output_dir: str | os.PathLike,
+    *,
+    base_model_name_or_path: str,
+    lora_alpha: float | None = None,
+    overwrite: bool = False,
+    expert_output_dir: str | os.PathLike | None = None,
+    nonexpert_output_dir: str | os.PathLike | None = None,
+) -> dict:
+    input_path = Path(input_dir).expanduser().resolve()
+    output_path = Path(output_dir).expanduser().resolve()
+    expert_output_path = (
+        Path(expert_output_dir).expanduser().resolve()
+        if expert_output_dir is not None
+        else None
+    )
+    nonexpert_output_path = (
+        Path(nonexpert_output_dir).expanduser().resolve()
+        if nonexpert_output_dir is not None
+        else None
+    )
+
+    if not input_path.is_dir():
+        raise FileNotFoundError(f"Input directory not found: {input_path}")
+
+    output_paths = [output_path]
+    output_paths.extend(path for path in (expert_output_path, nonexpert_output_path) if path is not None)
+    if len(set(output_paths)) != len(output_paths):
+        raise ValueError("Merged, expert, and non-expert output directories must be distinct.")
+    _validate_no_ancestor_paths(
+        output_paths,
+        label="Merged/expert/non-expert output directories",
+    )
+    for path in output_paths:
+        _validate_output_dir(path, input_path, overwrite)
+
+    existing_tensors, existing_targets = _load_existing_adapter(input_path)
+    fused_tensors, rank, fused_targets = _convert_fused_expert_lora(input_path / FUSED_EXPERT_LORA_FILE)
+    _validate_nonexpert_rank(existing_tensors, rank, input_path)
+    if nonexpert_output_path is not None and not existing_tensors:
+        raise ValueError(
+            f"Cannot write non-expert adapter: no {ADAPTER_MODEL_FILE} found in {input_path}."
+        )
+
+    tensors: dict[str, torch.Tensor] = {}
+    for key, value in existing_tensors.items():
+        _merge_tensor(tensors, key, value)
+    for key, value in fused_tensors.items():
+        _merge_tensor(tensors, key, value)
+
+    target_modules = set(existing_targets)
+    target_modules.update(fused_targets)
+    config = _build_adapter_config(
+        input_path,
+        rank,
+        target_modules,
+        base_model_name_or_path,
+        lora_alpha,
+    )
+
+    _write_adapter(output_path, input_path, tensors, config, overwrite=overwrite)
+
+    split_outputs: dict[str, dict] = {}
+    if expert_output_path is not None:
+        expert_config = _build_adapter_config(
+            input_path,
+            rank,
+            set(fused_targets),
+            base_model_name_or_path,
+            lora_alpha,
+            include_input_target_modules=False,
+        )
+        _write_adapter(
+            expert_output_path,
+            input_path,
+            fused_tensors,
+            expert_config,
+            overwrite=overwrite,
+        )
+        split_outputs["expert"] = {
+            "output_dir": str(expert_output_path),
+            "tensor_count": len(fused_tensors),
+            "target_modules": expert_config["target_modules"],
+        }
+
+    if nonexpert_output_path is not None:
+        nonexpert_config = _build_adapter_config(
+            input_path,
+            rank,
+            set(existing_targets),
+            base_model_name_or_path,
+            lora_alpha,
+            include_input_target_modules=False,
+        )
+        _write_adapter(
+            nonexpert_output_path,
+            input_path,
+            existing_tensors,
+            nonexpert_config,
+            overwrite=overwrite,
+        )
+        split_outputs["nonexpert"] = {
+            "output_dir": str(nonexpert_output_path),
+            "tensor_count": len(existing_tensors),
+            "target_modules": nonexpert_config["target_modules"],
+        }
+
+    return {
+        "input_dir": str(input_path),
+        "output_dir": str(output_path),
+        "tensor_count": len(tensors),
+        "rank": rank,
+        "target_modules": config["target_modules"],
+        "lora_alpha": config["lora_alpha"],
+        "split_outputs": split_outputs,
+    }
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(
+        description="Convert KT fused expert LoRA weights to an SGLang adapter directory."
+    )
+    parser.add_argument("input_dir", help="Directory containing fused_expert_lora.safetensors.")
+    parser.add_argument("output_dir", help="Destination adapter directory.")
+    parser.add_argument(
+        "--base-model-name-or-path",
+        required=True,
+        help="Base model path/name to write into adapter_config.json.",
+    )
+    parser.add_argument(
+        "--lora-alpha",
+        type=float,
+        default=None,
+        help="LoRA alpha to use when input adapter_config.json is absent.",
+    )
+    parser.add_argument(
+        "--overwrite",
+        action="store_true",
+        help="Remove and recreate output_dir if it already contains files.",
+    )
+    parser.add_argument(
+        "--expert-output-dir",
+        default=None,
+        help="Optional destination for a split expert-only adapter directory.",
+    )
+    parser.add_argument(
+        "--nonexpert-output-dir",
+        default=None,
+        help="Optional destination for a split non-expert-only adapter directory.",
+    )
+    return parser.parse_args()
+
+
+def main() -> None:
+    args = parse_args()
+    summary = convert_kt_to_sglang_adapter(
+        args.input_dir,
+        args.output_dir,
+        base_model_name_or_path=args.base_model_name_or_path,
+        lora_alpha=args.lora_alpha,
+        overwrite=args.overwrite,
+        expert_output_dir=args.expert_output_dir,
+        nonexpert_output_dir=args.nonexpert_output_dir,
+    )
+    print(
+        "Converted KT fused expert LoRA adapter: "
+        f"{summary['tensor_count']} tensors, rank={summary['rank']}, "
+        f"target_modules={summary['target_modules']}"
+    )
+    for name, split_summary in summary["split_outputs"].items():
+        print(
+            f"Wrote {name} adapter: {split_summary['tensor_count']} tensors, "
+            f"target_modules={split_summary['target_modules']}, "
+            f"output_dir={split_summary['output_dir']}"
+        )
+
+
+if __name__ == "__main__":
+    main()

kt-kernel/python/sft/deepseek_v4.py

@@ -0,0 +1,473 @@
+"""Selective loading of native DeepSeek V4 MXFP4 checkpoints for LoRA SFT."""
+
+from __future__ import annotations
+
+import argparse
+import fcntl
+import logging
+import shutil
+import stat
+import math
+import os
+from pathlib import Path
+
+import torch
+from safetensors import safe_open
+from safetensors.torch import save_file
+
+from .artifacts import (
+    KTArtifactError,
+    KTPretrainedLoadPlan,
+    KT_NON_EXPERT_INDEX_NAME,
+    KT_NON_EXPERT_MANIFEST_NAME,
+    _canonical_json_sha256,
+    _config_value,
+    _distributed_validation_context,
+    _fsync_directory,
+    _rawint4_shard_record,
+    _read_json,
+    _safe_root,
+    _sha256_file,
+    _synchronize_artifact_validation,
+    _write_json_atomic,
+)
+
+logger = logging.getLogger(__name__)
+
+_KIND = "deepseek-v4-non-expert-bf16"
+_PREPARATION_NAME = ".kt-v4-cache-owner.json"
+_FP8_DTYPES = {"F8_E4M3", "F8_E4M3FN"}
+
+
+def canonical_non_expert_key(key: str) -> str:
+    """Map the native April checkpoint to Transformers' V4 module names."""
+    if key.startswith("mtp.") or ".experts." in key:
+        raise KTArtifactError(f"not a non-expert base-model tensor: {key}")
+    heads = {
+        "embed.weight": "model.embed_tokens.weight",
+        "head.weight": "lm_head.weight",
+        "norm.weight": "model.norm.weight",
+        "hc_head_fn": "model.hc_head.hc_fn",
+        "hc_head_base": "model.hc_head.hc_base",
+        "hc_head_scale": "model.hc_head.hc_scale",
+    }
+    if key in heads:
+        return heads[key]
+    if not key.startswith("layers."):
+        raise KTArtifactError(f"unsupported V4 checkpoint key: {key}")
+    key = "model." + key
+    replacements = (
+        (".attn.", ".self_attn."),
+        (".ffn.", ".mlp."),
+        (".indexer.compressor.", ".compressor.indexer."),
+        (".indexer.weights_proj.", ".compressor.indexer.scorer.weights_proj."),
+        (".indexer.wq_b.", ".compressor.indexer.q_b_proj."),
+        (".attn_norm.", ".input_layernorm."),
+        (".ffn_norm.", ".post_attention_layernorm."),
+        (".attn_sink", ".sinks"),
+        (".q_norm.", ".q_a_norm."),
+        (".norm.", ".kv_norm."),
+        (".wq_a.", ".q_a_proj."),
+        (".wq_b.", ".q_b_proj."),
+        (".wkv.", ".kv_proj."),
+        (".wgate.", ".gate_proj."),
+        (".wo_a.", ".o_a_proj."),
+        (".wo_b.", ".o_b_proj."),
+        (".shared_experts.w1.", ".shared_experts.gate_proj."),
+        (".shared_experts.w2.", ".shared_experts.down_proj."),
+        (".shared_experts.w3.", ".shared_experts.up_proj."),
+        (".gate.bias", ".gate.e_score_correction_bias"),
+        (".ape", ".position_bias"),
+    )
+    for old, new in replacements:
+        key = key.replace(old, new)
+    for native, target in (("hc_attn", "attn_hc"), ("hc_ffn", "ffn_hc")):
+        for suffix in ("fn", "base", "scale"):
+            key = key.replace(f".{native}_{suffix}", f".{target}.{suffix}")
+    return key
+
+
+def dequantize_fp8_blocks(weight: torch.Tensor, scale: torch.Tensor) -> torch.Tensor:
+    """Decode 128x128 E4M3 blocks; UE8M0 stores the multiplier, not its inverse."""
+    if weight.dtype != torch.float8_e4m3fn or scale.dtype != torch.float8_e8m0fnu:
+        raise KTArtifactError("expected E4M3 weights and UE8M0 scales")
+    if weight.ndim != 2 or tuple(scale.shape) != tuple(
+        math.ceil(d / 128) for d in weight.shape
+    ):
+        raise KTArtifactError("invalid V4 non-expert FP8 weight/scale shape")
+    scales = scale.float()
+    if not torch.isfinite(scales).all() or not (scales > 0).all():
+        raise KTArtifactError("non-finite or non-positive FP8 scales")
+    result = torch.empty(weight.shape, dtype=torch.bfloat16)
+    for row in range(0, weight.shape[0], 128):
+        block = weight[row : row + 128].float()
+        multipliers = scales[row // 128].repeat_interleave(128)[: weight.shape[1]]
+        block.mul_(multipliers)
+        if not torch.isfinite(block).all():
+            raise KTArtifactError("non-finite FP8 non-expert weight")
+        result[row : row + 128].copy_(block)
+        if not torch.isfinite(result[row : row + 128]).all():
+            raise KTArtifactError("FP8 non-expert weight overflows BF16")
+    return result
+
+
+def inspect_native_checkpoint(source: str | os.PathLike[str]) -> dict:
+    """Validate expert headers without materializing the 148 GiB expert payload."""
+    root = _safe_root(source, "DeepSeek V4 checkpoint")
+    config = _read_json(root / "config.json", "DeepSeek V4 config")
+    if config.get("model_type") != "deepseek_v4":
+        raise KTArtifactError("native MXFP4 SFT currently requires DeepSeek V4")
+    quant = config.get("quantization_config", {})
+    if (
+        quant.get("quant_method") != "fp8"
+        or quant.get("fmt") != "e4m3"
+        or quant.get("weight_block_size") != [128, 128]
+        or quant.get("scale_fmt") != "ue8m0"
+    ):
+        raise KTArtifactError("expected the native V4 FP8/UE8M0 non-expert checkpoint")
+    index = _read_json(root / KT_NON_EXPERT_INDEX_NAME, "native safetensors index")
+    weight_map = index.get("weight_map", {})
+    if not weight_map or any(
+        not isinstance(v, str) or Path(v).name != v for v in weight_map.values()
+    ):
+        raise KTArtifactError("invalid native checkpoint shard inventory")
+    dimensions = [
+        config.get(k)
+        for k in (
+            "num_hidden_layers",
+            "n_routed_experts",
+            "hidden_size",
+            "moe_intermediate_size",
+        )
+    ]
+    if any(type(value) is not int or value <= 0 for value in dimensions):
+        raise KTArtifactError("native V4 dimensions must be positive integers")
+    layers, experts, hidden, inter = dimensions
+    if min(layers, experts, hidden, inter) <= 0 or hidden % 32 or inter % 32:
+        raise KTArtifactError("invalid MXFP4 routed-expert dimensions")
+    expected = {}
+    for layer in range(layers):
+        for expert in range(experts):
+            for proj, (n, k) in {
+                "w1": (inter, hidden),
+                "w3": (inter, hidden),
+                "w2": (hidden, inter),
+            }.items():
+                prefix = f"layers.{layer}.ffn.experts.{expert}.{proj}"
+                expected[prefix + ".weight"] = ([n, k // 2], {"I8", "U8"})
+                expected[prefix + ".scale"] = ([n, k // 32], {"F8_E8M0"})
+    inventory, shard_records = {}, []
+    for name in sorted(set(weight_map.values())):
+        path = root / name
+        # Header + sampled payload identity is also used by the native INT4 contract.
+        shard_records.append(_rawint4_shard_record(path))
+        with safe_open(path, framework="pt", device="cpu") as handle:
+            for key in handle.keys():
+                if weight_map.get(key) != name or key in inventory:
+                    raise KTArtifactError(f"native shard/index mismatch: {key}")
+                tensor = handle.get_slice(key)
+                shape, dtype = tensor.get_shape(), tensor.get_dtype()
+                inventory[key] = {"shape": shape, "dtype": dtype}
+                if key in expected:
+                    expected_shape, expected_dtypes = expected[key]
+                    if shape != expected_shape or dtype not in expected_dtypes:
+                        raise KTArtifactError(
+                            f"invalid native MXFP4 tensor {key}: {shape}, {dtype}"
+                        )
+                elif ".experts." in key and not key.startswith("mtp."):
+                    raise KTArtifactError(f"unexpected routed-expert tensor: {key}")
+    if set(inventory) != set(weight_map) or not set(expected).issubset(inventory):
+        raise KTArtifactError("incomplete native MXFP4 checkpoint")
+    identity = {
+        "config_sha256": _sha256_file(root / "config.json"),
+        "index_sha256": _sha256_file(root / KT_NON_EXPERT_INDEX_NAME),
+        "shards": shard_records,
+    }
+    return {
+        "model_name_or_path": str(root),
+        "fingerprint": _canonical_json_sha256(identity),
+        "identity": identity,
+        "config": config,
+        "inventory": inventory,
+        "weight_map": weight_map,
+    }
+
+
+def _expected_cache_inventory(source: dict) -> dict:
+    result = {}
+    for key, spec in source["inventory"].items():
+        if key.startswith("mtp.") or ".experts." in key:
+            continue
+        if key.endswith(".scale"):
+            weight = key.removesuffix(".scale") + ".weight"
+            if source["inventory"].get(weight, {}).get("dtype") not in _FP8_DTYPES:
+                raise KTArtifactError(f"orphan non-expert FP8 scale: {key}")
+            continue
+        dtype = spec["dtype"]
+        if dtype in _FP8_DTYPES:
+            scale = key.removesuffix(".weight") + ".scale"
+            if source["inventory"].get(scale, {}).get("dtype") != "F8_E8M0":
+                raise KTArtifactError(f"missing native FP8 scale: {key}")
+            dtype = "BF16"
+        if dtype not in {"BF16", "F32", "I64"}:
+            raise KTArtifactError(f"unsupported non-expert dtype {dtype}: {key}")
+        target = canonical_non_expert_key(key)
+        if target in result:
+            raise KTArtifactError(f"duplicate converted key: {target}")
+        result[target] = {"shape": spec["shape"], "dtype": dtype}
+    return result
+
+
+def _convert_non_expert_cache(source: dict, output: Path) -> dict:
+    """Write only inside a private staging directory owned by the cache lock."""
+    source_path = source["model_name_or_path"]
+    expected = _expected_cache_inventory(source)
+    weight_map, records = {}, []
+    root = Path(source["model_name_or_path"])
+    for shard in sorted(set(source["weight_map"].values())):
+        tensors = {}
+        with safe_open(root / shard, framework="pt", device="cpu") as handle:
+            for key in handle.keys():
+                if (
+                    key.startswith("mtp.")
+                    or ".experts." in key
+                    or key.endswith(".scale")
+                ):
+                    continue
+                target = canonical_non_expert_key(key)
+                tensor = handle.get_tensor(key)
+                if source["inventory"][key]["dtype"] in _FP8_DTYPES:
+                    scale_key = key.removesuffix(".weight") + ".scale"
+                    scale_shard = source["weight_map"][scale_key]
+                    with safe_open(
+                        root / scale_shard, framework="pt", device="cpu"
+                    ) as scales:
+                        tensor = dequantize_fp8_blocks(
+                            tensor, scales.get_tensor(scale_key)
+                        )
+                elif tensor.is_floating_point() and not torch.isfinite(tensor).all():
+                    raise KTArtifactError(f"non-finite source tensor: {key}")
+                tensors[target] = tensor.contiguous()
+        if not tensors:
+            continue
+        name = f"non-expert-{len(records):03d}.safetensors"
+        save_file(tensors, output / name, metadata={"format": "pt"})
+        records.append(
+            {
+                "name": name,
+                "size": (output / name).stat().st_size,
+                "sha256": _sha256_file(output / name),
+            }
+        )
+        weight_map.update({key: name for key in tensors})
+        print(f"converted {shard}: {len(tensors)} non-expert tensors", flush=True)
+    if set(weight_map) != set(expected):
+        raise KTArtifactError("converted cache is missing non-expert tensors")
+    _write_json_atomic(
+        output / KT_NON_EXPERT_INDEX_NAME, {"metadata": {}, "weight_map": weight_map}
+    )
+    if source["identity"] != inspect_native_checkpoint(source_path)["identity"]:
+        raise KTArtifactError("source checkpoint changed during conversion")
+    payload = {
+        "kind": _KIND,
+        "version": 1,
+        "status": "ready",
+        "source": {
+            key: source[key]
+            for key in ("model_name_or_path", "fingerprint", "identity")
+        },
+        "expert_weight_format": "mxfp4",
+        "files": records,
+        "tensors": expected,
+        "index_sha256": _sha256_file(output / KT_NON_EXPERT_INDEX_NAME),
+    }
+    payload["fingerprint"] = _canonical_json_sha256(payload)
+    _write_json_atomic(output / KT_NON_EXPERT_MANIFEST_NAME, payload)
+    return payload
+
+
+def _validate_non_expert_cache(source: dict, root: Path, *, verify_payload: bool = True) -> dict:
+    manifest_path = root / KT_NON_EXPERT_MANIFEST_NAME
+    manifest = _read_json(manifest_path, "V4 non-expert cache manifest")
+    if (manifest.get("kind"), manifest.get("version"), manifest.get("status")) != (
+        _KIND,
+        1,
+        "ready",
+    ):
+        raise KTArtifactError("invalid V4 cache ready manifest")
+    body = {key: value for key, value in manifest.items() if key != "fingerprint"}
+    if manifest.get("fingerprint") != _canonical_json_sha256(body):
+        raise KTArtifactError("V4 cache manifest fingerprint mismatch")
+    if manifest.get("source") != {
+        key: source[key]
+        for key in ("model_name_or_path", "fingerprint", "identity")
+    }:
+        raise KTArtifactError("native V4 checkpoint changed after cache conversion")
+    expected = _expected_cache_inventory(source)
+    if manifest.get("tensors") != expected:
+        raise KTArtifactError("V4 cache tensor inventory differs from its source")
+    index_path = root / KT_NON_EXPERT_INDEX_NAME
+    if manifest.get("index_sha256") != _sha256_file(index_path):
+        raise KTArtifactError("V4 cache index hash mismatch")
+    index = _read_json(index_path, "V4 cache index")["weight_map"]
+    observed, files = {}, []
+    for record in manifest["files"]:
+        name = record["name"]
+        if Path(name).name != name or name in files:
+            raise KTArtifactError("invalid V4 cache shard name")
+        path = root / name
+        if (
+            path.is_symlink()
+            or not path.is_file()
+            or path.stat().st_size != record["size"]
+        ):
+            raise KTArtifactError(f"invalid V4 cache shard: {path}")
+        if verify_payload and _sha256_file(path) != record["sha256"]:
+            raise KTArtifactError(f"V4 cache payload hash mismatch: {path}")
+        with safe_open(path, framework="pt", device="cpu") as handle:
+            for key in handle.keys():
+                tensor = handle.get_slice(key)
+                if key in observed or index.get(key) != name:
+                    raise KTArtifactError(
+                        f"duplicate or unindexed cache tensor: {key}"
+                    )
+                observed[key] = {
+                    "shape": tensor.get_shape(),
+                    "dtype": tensor.get_dtype(),
+                }
+        files.append(name)
+    if (
+        observed != expected
+        or set(index) != set(expected)
+        or set(index.values()) != set(files)
+    ):
+        raise KTArtifactError("V4 cache payload inventory mismatch")
+    return manifest
+
+
+def _ensure_non_expert_cache(source: dict, output_path: str, *, verify_payload: bool = True) -> dict:
+    if not output_path:
+        raise KTArtifactError("Set kt_non_expert_weight_path to a writable non-expert cache directory")
+    output = Path(output_path).expanduser().absolute()
+    source_root = Path(source["model_name_or_path"])
+    resolved_output = output.resolve()
+    if resolved_output.is_relative_to(source_root) or source_root.is_relative_to(resolved_output):
+        raise KTArtifactError("V4 non-expert cache must not overlap the source checkpoint")
+    if output.is_symlink():
+        raise KTArtifactError(f"V4 cache must be a real directory: {output}")
+    output.parent.mkdir(parents=True, exist_ok=True)
+    # Never unlink the lock: all processes must keep locking the same inode.
+    lock_path = output.with_name(f".{output.name}.kt-v4-cache.lock")
+    flags = os.O_CREAT | os.O_RDWR | os.O_NOFOLLOW
+    descriptor = os.open(lock_path, flags, 0o600)
+    with os.fdopen(descriptor, "r+") as lock:
+        if not stat.S_ISREG(os.fstat(lock.fileno()).st_mode):
+            raise KTArtifactError(f"V4 cache lock must be a regular file: {lock_path}")
+        logger.info("Waiting for V4 non-expert cache lock: %s", lock_path)
+        fcntl.flock(lock, fcntl.LOCK_EX)
+        if output.is_symlink() or (output.exists() and not output.is_dir()):
+            raise KTArtifactError(f"V4 cache must be a real directory: {output}")
+        if output.exists() and any(output.iterdir()):
+            # Existing caches are immutable. Never silently replace damaged or unrelated data.
+            return _validate_non_expert_cache(source, output, verify_payload=verify_payload)
+
+        staging = output.with_name(f".{output.name}.kt-v4-cache.partial")
+        owner = {"kind": _KIND, "target": str(output), "source": source["fingerprint"]}
+        if staging.is_symlink():
+            raise KTArtifactError(f"V4 cache staging directory must not be a symlink: {staging}")
+        if staging.exists():
+            if not staging.is_dir() or _read_json(staging / _PREPARATION_NAME, "V4 cache owner") != owner:
+                raise KTArtifactError(f"Refusing to replace an unowned V4 staging directory: {staging}")
+            logger.info("Removing interrupted V4 cache preparation: %s", staging)
+            shutil.rmtree(staging)
+
+        sizes = {"BF16": 2, "F32": 4, "I64": 8}
+        payload_bytes = sum(math.prod(v["shape"]) * sizes[v["dtype"]] for v in _expected_cache_inventory(source).values())
+        required_bytes = payload_bytes + max(16 * 1024**2, payload_bytes // 100)
+        free_bytes = shutil.disk_usage(output.parent).free
+        if free_bytes < required_bytes:
+            raise KTArtifactError(
+                f"Insufficient space for V4 non-expert cache at {output}: "
+                f"need at least {required_bytes} bytes, available {free_bytes} bytes"
+            )
+        staging.mkdir(mode=0o700)
+        _write_json_atomic(staging / _PREPARATION_NAME, owner)
+        logger.info("Preparing V4 non-expert cache at %s (%d tensor bytes)", output, payload_bytes)
+        try:
+            _convert_non_expert_cache(source, staging)
+            manifest = _validate_non_expert_cache(source, staging, verify_payload=verify_payload)
+            # rename replaces a pre-existing empty directory, never a non-empty cache.
+            os.replace(staging, output)
+            _fsync_directory(output.parent)
+        except Exception:
+            # We created this directory while holding the lock; no unrelated paths are removed.
+            if staging.exists() and not staging.is_symlink():
+                shutil.rmtree(staging)
+            raise
+        logger.info("V4 non-expert cache is ready: %s", output)
+        return manifest
+
+
+def prepare_non_expert_cache(source_path: str, output_path: str) -> dict:
+    """Atomically prepare or fully validate/reuse a native V4 non-expert cache."""
+    return _ensure_non_expert_cache(inspect_native_checkpoint(source_path), output_path)
+
+
+def resolve_native_load_plan(kt_config, source_path, explicit_quantization_config=None):
+    if explicit_quantization_config is not None:
+        raise KTArtifactError("native MXFP4 cannot use an explicit framework quantizer")
+    dist, rank, world_size = _distributed_validation_context()
+    error, signature, plan = None, None, None
+    try:
+        source = inspect_native_checkpoint(source_path)
+        native_root = _safe_root(
+            _config_value(kt_config, "kt_weight_path"), "native MXFP4 experts"
+        )
+        if str(native_root) != source["model_name_or_path"]:
+            raise KTArtifactError(
+                "MXFP4 experts and non-expert cache must use the same native base"
+            )
+        manifest = _ensure_non_expert_cache(
+            source, _config_value(kt_config, "kt_non_expert_weight_path"), verify_payload=rank == 0
+        )
+        root = _safe_root(
+            Path(_config_value(kt_config, "kt_non_expert_weight_path")).expanduser().absolute(),
+            "V4 non-expert cache",
+        )
+        expected = _expected_cache_inventory(source)
+        files = [record["name"] for record in manifest["files"]]
+        manifest_path = root / KT_NON_EXPERT_MANIFEST_NAME
+        signature = (manifest["fingerprint"], str(root), str(native_root))
+        plan = KTPretrainedLoadPlan(
+            source_model_name_or_path=str(native_root),
+            weight_path=str(root),
+            checkpoint_files=tuple(str(root / name) for name in files),
+            weight_keys=frozenset(expected),
+            manifest=manifest,
+            manifest_path=str(manifest_path),
+            routed_weight_path=str(native_root),
+            routed_manifest=manifest,
+            routed_manifest_path=str(manifest_path),
+            lora_rank=_config_value(kt_config, "kt_lora_rank"),
+            lora_alpha=_config_value(kt_config, "kt_lora_alpha"),
+        )
+    except Exception as exc:
+        error = (
+            exc
+            if isinstance(exc, KTArtifactError)
+            else KTArtifactError(f"native V4 artifact validation failed: {exc}")
+        )
+    _synchronize_artifact_validation(dist, rank, world_size, error, signature)
+    return plan
+
+
+def main():
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--source", required=True)
+    parser.add_argument("--output", required=True)
+    args = parser.parse_args()
+    prepare_non_expert_cache(args.source, args.output)
+
+
+if __name__ == "__main__":
+    main()

kt-kernel/python/sft/export_dsv4_sglang_adapter.py

@@ -0,0 +1,308 @@
+#!/usr/bin/env python3
+"""Export a single V4 adapter: merged BF16 non-experts + native MXFP4 expert LoRA.
+
+This is static deployment, not adapter hot swapping. Original checkpoint files
+are never modified. The model copy contains no routed-expert base weights.
+"""
+
+import argparse
+import json
+import re
+import shutil
+from pathlib import Path
+from types import SimpleNamespace
+
+import torch
+from safetensors import safe_open
+from safetensors.torch import load_file, save_file
+
+from kt_kernel.sft.artifacts import (
+    FUSED_EXPERT_LORA_NAME,
+    KT_ADAPTER_MANIFEST_NAME,
+    KTArtifactError,
+    _read_json,
+    _sha256_file,
+    _write_json_atomic,
+)
+from kt_kernel.sft.deepseek_v4 import (
+    canonical_non_expert_key,
+    inspect_native_checkpoint,
+    resolve_native_load_plan,
+)
+
+from .convert_kt_to_sglang_adapter import (
+    _build_adapter_config,
+    _convert_fused_expert_lora,
+)
+
+
+_STANDARD = re.compile(
+    r"^(?:base_model\.model\.)?(model\.layers\.\d+\..+)\.lora_([AB])(?:\.default)?\.weight$"
+)
+_TARGET = re.compile(
+    r"^model\.layers\.\d+\.(?:self_attn\.(?:q_a_proj|q_b_proj|kv_proj|o_b_proj|compressor\.(?:kv_proj|gate_proj))"
+    r"|mlp\.shared_experts\.(?:gate_proj|up_proj|down_proj))$"
+)
+
+
+def _validate_fused_experts(path, config, rank):
+    expected = {}
+    experts = config["n_routed_experts"]
+    hidden, inter = config["hidden_size"], config["moe_intermediate_size"]
+    for layer in range(config["num_hidden_layers"]):
+        for projection, (n, k) in {
+            "gate": (inter, hidden),
+            "up": (inter, hidden),
+            "down": (hidden, inter),
+        }.items():
+            prefix = f"layers.{layer}.experts.{projection}_lora_"
+            expected[prefix + "a"] = [experts, rank, k]
+            expected[prefix + "b"] = [experts, n, rank]
+    with safe_open(path, framework="pt", device="cpu") as handle:
+        if set(handle.keys()) != set(expected):
+            raise KTArtifactError(
+                "incomplete or unexpected fused expert LoRA inventory"
+            )
+        for name, shape in expected.items():
+            tensor = handle.get_slice(name)
+            if tensor.get_shape() != shape or tensor.get_dtype() not in {
+                "BF16",
+                "F16",
+                "F32",
+            }:
+                raise KTArtifactError(f"invalid fused expert LoRA shape/dtype: {name}")
+            if not torch.isfinite(handle.get_tensor(name)).all():
+                raise KTArtifactError(f"non-finite fused expert LoRA: {name}")
+
+
+def export(
+    source_path,
+    cache_path,
+    adapter_path,
+    output_path,
+    component="all",
+    match_expert_kernel=False,
+):
+    output, adapter = Path(output_path).absolute(), Path(adapter_path).absolute()
+    if component not in {"all", "base", "experts", "nonexperts"}:
+        raise KTArtifactError(f"invalid export component: {component}")
+    if output.exists():
+        raise KTArtifactError(f"export destination already exists: {output}")
+    for path in (source_path, cache_path, adapter_path):
+        source_root = Path(path).resolve()
+        if output.resolve().is_relative_to(source_root) or source_root.is_relative_to(
+            output.resolve()
+        ):
+            raise KTArtifactError(
+                "export destination must not overlap an input directory"
+            )
+    plan = resolve_native_load_plan(
+        SimpleNamespace(
+            kt_weight_path=str(Path(source_path).resolve()),
+            kt_non_expert_weight_path=cache_path,
+        ),
+        source_path,
+    )
+    manifest = _read_json(adapter / KT_ADAPTER_MANIFEST_NAME, "KT adapter manifest")
+    if (
+        manifest.get("status") != "ready"
+        or manifest.get("expert_weight_format") != "mxfp4"
+    ):
+        raise KTArtifactError("expected a ready native-MXFP4 training adapter")
+    if (
+        manifest.get("base", {}).get("fingerprint")
+        != plan.manifest["source"]["fingerprint"]
+    ):
+        raise KTArtifactError("adapter base does not match the native checkpoint")
+    required = {
+        "adapter_config.json",
+        "adapter_model.safetensors",
+        FUSED_EXPERT_LORA_NAME,
+    }
+    artifacts = manifest.get("artifacts", {})
+    if not required.issubset(artifacts):
+        raise KTArtifactError("adapter manifest is missing a required artifact")
+    for name, record in artifacts.items():
+        path = adapter / name
+        if Path(name).name != name or path.is_symlink() or not path.is_file():
+            raise KTArtifactError(f"invalid adapter artifact: {name}")
+        if (
+            path.stat().st_size != record["size"]
+            or _sha256_file(path) != record["sha256"]
+        ):
+            raise KTArtifactError(f"adapter artifact failed integrity check: {name}")
+    config = _read_json(adapter / "adapter_config.json", "PEFT config")
+    rank, alpha = int(config["r"]), float(config["lora_alpha"])
+    if rank != 8 or alpha != 16 or config.get("lora_dropout", 0.0) != 0:
+        raise KTArtifactError("April V4 static export requires rank=8, alpha=16 and dropout=0")
+    if (
+        config.get("use_dora")
+        or config.get("use_rslora")
+        or config.get("rank_pattern")
+        or config.get("alpha_pattern")
+    ):
+        raise KTArtifactError(
+            "V4 static export currently requires uniform standard LoRA"
+        )
+    if manifest.get("lora") != {"rank": rank, "alpha": alpha}:
+        raise KTArtifactError("adapter rank/alpha provenance mismatch")
+    standard = load_file(str(adapter / "adapter_model.safetensors"))
+    pairs, consumed = {}, []
+    for name, tensor in standard.items():
+        match = _STANDARD.fullmatch(name)
+        if match is None or not _TARGET.fullmatch(match[1]):
+            raise KTArtifactError(f"unsupported non-expert adapter tensor: {name}")
+        pair = pairs.setdefault(match[1], {})
+        if match[2] in pair or tensor.ndim != 2 or not torch.isfinite(tensor).all():
+            raise KTArtifactError(f"invalid or duplicate adapter tensor: {name}")
+        pair[match[2]] = tensor
+    if not pairs or any(set(pair) != {"A", "B"} for pair in pairs.values()):
+        raise KTArtifactError("incomplete standard LoRA A/B pairs")
+    source = inspect_native_checkpoint(source_path)
+    _validate_fused_experts(adapter / FUSED_EXPERT_LORA_NAME, source["config"], rank)
+    reverse = {
+        canonical_non_expert_key(key): key
+        for key in source["inventory"]
+        if not key.startswith("mtp.")
+        and ".experts." not in key
+        and not key.endswith(".scale")
+    }
+    output.mkdir(parents=True)
+    model_dir = output / "model"
+    model_dir.mkdir()
+    weight_map, files, covered = {}, [], set()
+    for index, checkpoint in enumerate(plan.checkpoint_files):
+        tensors, exported = load_file(checkpoint), {}
+        for name, weight in tensors.items():
+            module = name.removesuffix(".weight")
+            if module in pairs:
+                covered.add(module)
+                a, b = pairs[module]["A"], pairs[module]["B"]
+                if tuple(a.shape) != (rank, weight.shape[1]) or tuple(b.shape) != (
+                    weight.shape[0],
+                    rank,
+                ):
+                    raise KTArtifactError(f"LoRA shape mismatch: {module}")
+                if component in {"all", "nonexperts"}:
+                    weight = (
+                        weight.float() + (alpha / rank) * (b.float() @ a.float())
+                    ).to(weight.dtype)
+                    if not torch.isfinite(weight).all():
+                        raise KTArtifactError(f"non-finite merged weight: {module}")
+                    consumed.append(module)
+            exported[reverse[name]] = weight.contiguous()
+        filename = f"non-experts-{index:03d}.safetensors"
+        save_file(exported, model_dir / filename, metadata={"format": "pt"})
+        weight_map.update({key: filename for key in exported})
+        files.append({"name": filename, "sha256": _sha256_file(model_dir / filename)})
+    if covered != set(pairs):
+        raise KTArtifactError("standard LoRA pair has no matching base weight")
+    model_config = dict(source["config"])
+    model_config.pop("quantization_config", None)
+    model_config.update(
+        torch_dtype="bfloat16", num_nextn_predict_layers=0, use_cache=True
+    )
+    _write_json_atomic(model_dir / "config.json", model_config)
+    _write_json_atomic(
+        model_dir / "model.safetensors.index.json",
+        {"metadata": {}, "weight_map": weight_map},
+    )
+    for name in (
+        "tokenizer.json",
+        "tokenizer_config.json",
+        "special_tokens_map.json",
+        "generation_config.json",
+    ):
+        path = Path(source_path) / name
+        if path.is_file():
+            shutil.copyfile(path, model_dir / name)
+    expert_count = 0
+    expert_enabled = component in {"all", "experts"}
+    if expert_enabled or match_expert_kernel:
+        expert_tensors, expert_rank, targets = _convert_fused_expert_lora(
+            adapter / FUSED_EXPERT_LORA_NAME
+        )
+        if (
+            expert_rank != rank
+            or len(expert_tensors)
+            != source["config"]["num_hidden_layers"]
+            * source["config"]["n_routed_experts"]
+            * 6
+        ):
+            raise KTArtifactError("expert adapter rank or tensor inventory mismatch")
+        if not expert_enabled:
+            expert_tensors = {
+                name: torch.zeros_like(tensor) if ".lora_B." in name else tensor
+                for name, tensor in expert_tensors.items()
+            }
+        expert_dir = output / "experts"
+        expert_dir.mkdir()
+        save_file(expert_tensors, expert_dir / "adapter_model.safetensors")
+        _write_json_atomic(
+            expert_dir / "adapter_config.json",
+            _build_adapter_config(
+                adapter,
+                rank,
+                targets,
+                source_path,
+                alpha,
+                include_input_target_modules=False,
+            ),
+        )
+        expert_count = len(expert_tensors)
+    report = {
+        "version": 1,
+        "status": "ready",
+        "component": component,
+        "deployment": "static-merged-bf16-non-experts/native-mxfp4-expert-lora",
+        "source": source_path,
+        "source_fingerprint": source["fingerprint"],
+        "adapter": str(adapter),
+        "adapter_manifest_sha256": _sha256_file(adapter / KT_ADAPTER_MANIFEST_NAME),
+        "standard_pairs_consumed": sorted(consumed),
+        "standard_source_tensor_count": len(standard),
+        "expert_exported_tensor_count": expert_count,
+        "expert_effect_enabled": expert_enabled,
+        "zero_expert_adapter_control": bool(match_expert_kernel and not expert_enabled),
+        "model_files": files,
+        "rounding": "standard LoRA merged in FP32 then rounded to BF16",
+    }
+    _write_json_atomic(output / "deployment_manifest.json", report)
+    print(
+        json.dumps(
+            {
+                key: value
+                for key, value in report.items()
+                if key not in {"standard_pairs_consumed", "model_files"}
+            },
+            indent=2,
+        )
+    )
+    return report
+
+
+def main():
+    parser = argparse.ArgumentParser(description=__doc__)
+    for argument in ("source", "cache", "adapter", "output"):
+        parser.add_argument("--" + argument, required=True)
+    parser.add_argument(
+        "--component", choices=("all", "base", "experts", "nonexperts"), default="all"
+    )
+    parser.add_argument(
+        "--match-expert-kernel",
+        action="store_true",
+        help="Use zero-B expert LoRA in baseline ablations so every component uses the same native SFT forward kernel.",
+    )
+    args = parser.parse_args()
+    export(
+        args.source,
+        args.cache,
+        args.adapter,
+        args.output,
+        args.component,
+        args.match_expert_kernel,
+    )
+
+
+if __name__ == "__main__":
+    main()

kt-kernel/python/sft/layer.py

@@ -15,7 +15,7 @@ import logging
 import math
 import os
 from contextlib import nullcontext
-from dataclasses import dataclass
+from dataclasses import dataclass, replace
 from typing import TYPE_CHECKING, Any
 
 import torch
@@ -423,12 +423,14 @@ class KTMoELayerWrapper(nn.Module):
         return result
 
     def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        with torch.profiler.record_function("kt.sft.routing"):
+            topk_ids, topk_weights = self._compute_routing(hidden_states)
+        return self._forward_routed(hidden_states, topk_ids, topk_weights)
 
+    def _forward_routed(self, hidden_states, topk_ids, topk_weights):
         import torch.distributed as dist
 
         dist_on = dist.is_initialized() and dist.get_world_size() > 1
-        with torch.profiler.record_function("kt.sft.routing"):
-            topk_ids, topk_weights = self._compute_routing(hidden_states)
 
         train_lora = bool(
             self._kt_managed_lora_enabled
@@ -989,3 +991,23 @@ class KTMoELayerWrapper(nn.Module):
         # PEFT weights are views into wrapper's contiguous buffers —
         # optimizer.step() already updated them in-place, just re-sync to C++.
         self.wrapper.update_lora_weights()
+
+
+class KTRoutedExpertsWrapper(KTMoELayerWrapper):
+    """Execute only experts; routing and shared experts stay in the model."""
+
+    def __init__(self, original_moe, **kwargs):
+        config = kwargs["moe_config"]
+        expert_container = nn.Module()
+        expert_container.add_module(config.experts_attr, getattr(original_moe, config.experts_attr))
+        kwargs["moe_config"] = replace(config, has_shared_experts=False)
+        super().__init__(original_moe=expert_container, **kwargs)
+        self._is_kt_routed_experts_wrapper = True
+
+    def forward(self, hidden_states, top_k_index, top_k_weights):
+        if hidden_states.ndim != 2 or hidden_states.shape[-1] != self.hidden_size:
+            raise ValueError("KT routed experts require [tokens, hidden_size] activations")
+        expected = (hidden_states.shape[0], self.moe_config.num_experts_per_tok)
+        if tuple(top_k_index.shape) != expected or tuple(top_k_weights.shape) != expected:
+            raise ValueError(f"KT routed experts require routing tensors shaped {expected}")
+        return self._forward_routed(hidden_states.unsqueeze(0), top_k_index, top_k_weights).squeeze(0)

kt-kernel/python/sft/wrapper.py

@@ -22,18 +22,21 @@ from .arch import (
     get_moe_arch_config,
     get_moe_module,
 )
-from .layer import KTMoELayerWrapper
+from .layer import KTMoELayerWrapper, KTRoutedExpertsWrapper
 from .lora import LoRAExperts
 from .base import _supports_authoritative_optimizer_grads
 from .backend import (
     FP8_BACKEND,
     INT8_BACKEND,
+    MXFP4_BACKEND,
+    MXFP4_SFT_METHOD,
     RAWINT4_BACKEND,
     RAWINT4_GROUP_SIZE,
     RAWINT4_SFT_METHOD,
     RAWINT4_WEIGHT_LAYOUT,
     get_fp8_runtime,
     get_int8_runtime,
+    get_mxfp4_runtime,
     get_rawint4_checkpoint_contract,
     get_rawint4_runtime,
 )
@@ -370,6 +373,7 @@ def wrap_moe_layers_with_kt_wrapper(model: nn.Module, kt_plugin: Any) -> list[KT
         INT8_BACKEND: "INT8_SFT",
         "AMXINT8": "INT8_SFT",
         RAWINT4_BACKEND: RAWINT4_SFT_METHOD,
+        MXFP4_BACKEND: MXFP4_SFT_METHOD,
         "AMXINT4_KGroup": RAWINT4_SFT_METHOD,
         "AMXINT4": "AMXINT4_SFT",
         "AMXBF16_SkipLoRA": "AMXBF16_SFT_SkipLoRA",
@@ -462,13 +466,14 @@ def wrap_moe_layers_with_kt_wrapper(model: nn.Module, kt_plugin: Any) -> list[KT
             "AMXFP8_SFT",
             "INT8_SFT",
             RAWINT4_SFT_METHOD,
+            MXFP4_SFT_METHOD,
         }
         or requested_num_gpu_experts != 0
         or use_lora_experts
     ):
         raise KTAMXConfigError(
             "activation_policy.cpu=retain requires CPU-only AMXBF16 Full/LoRA "
-            "or frozen-base INT8/FP8/RAWINT4 LoRA; Hybrid, GPU-expert, "
+            "or frozen-base INT8/FP8/RAWINT4/MXFP4 LoRA; Hybrid, GPU-expert, "
             "LoRA-expert, generic INT4, and SkipLoRA paths are not supported"
         )
     _validate_rawint4_activation_policy(expert_weight_format, activation_policy)
@@ -807,7 +812,7 @@ def wrap_moe_layers_with_kt_wrapper(model: nn.Module, kt_plugin: Any) -> list[KT
 
         _layer_experts = getattr(moe_module, moe_config.experts_attr, None)
         _layer_is_fused = _detect_fused(_layer_experts)
-        if expert_weight_format in {"int8", "fp8", "rawint4"} and not _layer_is_fused and not force_fused_expert_lora:
+        if expert_weight_format in {"int8", "fp8", "rawint4", "mxfp4"} and not _layer_is_fused and not force_fused_expert_lora:
             raise KTAMXConfigError(
                 f"{expert_weight_format.upper()} LoRA with non-fused runtime experts requires "
                 "kt_force_fused_expert_lora=true"
@@ -864,6 +869,10 @@ def wrap_moe_layers_with_kt_wrapper(model: nn.Module, kt_plugin: Any) -> list[KT
                         layer_idx=layer_idx,
                         expected_shapes=expected_shapes,
                     )
+                elif expert_weight_format == "mxfp4":
+                    get_mxfp4_runtime()
+                    if moe_config.router_type != "precomputed":
+                        raise KTAMXConfigError("MXFP4 SFT currently supports DeepSeek V4 routed experts only")
                 elif use_kt_weight_path:
                     logger.debug(f"Layer {layer_idx}: forward + backward from kt_weight_path (.kt files)")
                 elif expert_weight_format == "fp8":
@@ -943,8 +952,9 @@ def wrap_moe_layers_with_kt_wrapper(model: nn.Module, kt_plugin: Any) -> list[KT
                     lora_alpha=lora_alpha,
                     lora_dropout=lora_dropout,
                     max_cache_depth=getattr(cfg, "kt_max_cache_depth", 2),
-                    group_size=RAWINT4_GROUP_SIZE if expert_weight_format == "rawint4" else 128,
-                    zero_point=False if expert_weight_format == "rawint4" else True,
+                    group_size=32 if expert_weight_format in {"rawint4", "mxfp4"} else 128,
+                    zero_point=expert_weight_format not in {"rawint4", "mxfp4"},
+                    swiglu_limit=float(getattr(_text_cfg, "swiglu_limit", 0.0)) if expert_weight_format == "mxfp4" else 0.0,
                     full_weight_grad=full_weight_grad,
                 )
             except BaseException as exc:
@@ -1067,7 +1077,8 @@ def wrap_moe_layers_with_kt_wrapper(model: nn.Module, kt_plugin: Any) -> list[KT
                 dtype=torch.bfloat16,
             )
 
-        layer_wrapper = KTMoELayerWrapper(
+        layer_wrapper_cls = KTRoutedExpertsWrapper if moe_config.router_type == "precomputed" else KTMoELayerWrapper
+        layer_wrapper = layer_wrapper_cls(
             original_moe=moe_module,
             wrapper=wrapper,
             lora_params=None,
@@ -1098,7 +1109,8 @@ def wrap_moe_layers_with_kt_wrapper(model: nn.Module, kt_plugin: Any) -> list[KT
                 "weight_layout": RAWINT4_WEIGHT_LAYOUT,
             }
 
-        setattr(layer, moe_config.moe_layer_attr, layer_wrapper)
+        if moe_config.router_type != "precomputed":
+            setattr(layer, moe_config.moe_layer_attr, layer_wrapper)
         # Base weights have been copied into the C++ kernel's internal BufferB format.
         # In full_weight_grad mode, the authoritative copies are gate_proj_buf etc.
         # Always release local references to save ~1 GB/layer.
@@ -1115,8 +1127,10 @@ def wrap_moe_layers_with_kt_wrapper(model: nn.Module, kt_plugin: Any) -> list[KT
             moe_module,
             moe_config,
             full_weight_grad=full_weight_grad,
-            empty_placeholders=expert_weight_format in {"int8", "fp8", "rawint4"},
+            empty_placeholders=expert_weight_format in {"int8", "fp8", "rawint4", "mxfp4"},
         )
+        if moe_config.router_type == "precomputed":
+            setattr(moe_module, moe_config.experts_attr, layer_wrapper)
 
     ephemeral_finish_error = None
     if ephemeral_store is not None:
@@ -1365,7 +1379,7 @@ def load_kt_model(
     expert_weight_format = getattr(cfg, "kt_expert_weight_format", None)
     native_fp8_experts = expert_weight_format == "fp8"
     native_rawint4_experts = expert_weight_format == "rawint4"
-    native_checkpoint_experts = native_fp8_experts or native_rawint4_experts
+    native_checkpoint_experts = native_fp8_experts or native_rawint4_experts or expert_weight_format == "mxfp4"
     if native_checkpoint_experts:
         # Native quantized kt_weight_path is raw-checkpoint provenance; KT owns
         # and loads the routed experts without materializing them in Transformers.

kt-kernel/python/utils/loader.py

@@ -1220,17 +1220,41 @@ class MXFP4SafeTensorLoader(SafeTensorLoader):
         return list(dict.fromkeys(candidates))
 
     @staticmethod
-    def _ue8m0_to_bf16(scale_t: torch.Tensor) -> torch.Tensor:
+    def _ue8m0_to_bf16(
+        scale_t: torch.Tensor, *, reject_non_finite: bool = False
+    ) -> torch.Tensor:
+        if reject_non_finite:
+            valid_dtypes = {torch.uint8}
+            native_ue8m0 = getattr(torch, "float8_e8m0fnu", None)
+            if native_ue8m0 is not None:
+                valid_dtypes.add(native_ue8m0)
+            if scale_t.dtype not in valid_dtypes:
+                raise TypeError(
+                    "MXFP4 SFT requires raw uint8/float8_e8m0fnu scales, "
+                    f"got {scale_t.dtype}"
+                )
         if scale_t.dtype != torch.uint8:
             scale_t = scale_t.view(torch.uint8)
-        # bf16 = [sign(1) | exp(8) | mant(7)]; setting mant=0, exp=e gives 2^(e-127),
-        # which is exactly the value encoded by ue8m0 for e ∈ [1, 254]. e=0 → bf16 +0
-        # (acceptable: ue8m0=0 represents 2^-127, below bf16 normal range), e=255 → +inf.
-        # Compute in int32 then narrow to int16 (max value is 255<<7=32640, fits int16),
-        # because torch CPU has no lshift kernel for uint16.
-        return (scale_t.to(torch.int32) << 7).to(torch.int16).view(torch.bfloat16).contiguous()
+        if reject_non_finite and torch.any(scale_t == 0xFF).item():
+            raise ValueError(
+                "MXFP4 SFT rejects reserved UE8M0 scale 0xff because it "
+                "would decode to a non-finite BF16 value"
+            )
+        # BF16 normal exponent bits reproduce UE8M0 codes 1..254 exactly.
+        # UE8M0 code 0 is 2^-127, represented by BF16 subnormal bits 0x0040.
+        # Build through int32 because torch CPU has no lshift kernel for uint16.
+        scale_i32 = scale_t.to(torch.int32)
+        bf16_bits = scale_i32 << 7
+        bf16_bits = torch.where(scale_i32 == 0, 0x0040, bf16_bits)
+        return bf16_bits.to(torch.int16).view(torch.bfloat16).contiguous()
 
-    def load_experts(self, base_key: str, device: str = "cpu"):
+    def load_experts(
+        self,
+        base_key: str,
+        device: str = "cpu",
+        *,
+        reject_non_finite_scales: bool = False,
+    ):
         gate_name, up_name, down_name = self.PROJ_NAMES
         prefix = None
         expert_count = 0
@@ -1270,7 +1294,9 @@ class MXFP4SafeTensorLoader(SafeTensorLoader):
                 (down_name, down_scales),
             ):
                 s = self.load_tensor(f"{prefix}.{exp_id}.{proj}.scale", device)
-                dst[exp_id] = self._ue8m0_to_bf16(s)
+                dst[exp_id] = self._ue8m0_to_bf16(
+                    s, reject_non_finite=reject_non_finite_scales
+                )
 
         print(f"[MXFP4SafeTensorLoader] Loaded {expert_count} experts from {prefix}")
         return {