diff --git a/python/freetoken/engine/engine.py b/python/freetoken/engine/engine.py index 018e9622..c386e28b 100644 --- a/python/freetoken/engine/engine.py +++ b/python/freetoken/engine/engine.py @@ -1144,6 +1144,21 @@ def _cpu_moe_executor_viable(model_config) -> bool: return False expert_quant = getattr(model_config, "expert_quant", "none") fmt = expert_quant if expert_quant != "none" else (moe_wfmt or "bf16") + if fmt == "gguf": + # "gguf" is a container tag, not a layout: the checkpoint picks a ggml type per + # tensor and the concrete CPU format has to be recovered from the bank types. + # Testing the tag against _WFMT_IDS answers False for EVERY GGUF checkpoint, which + # silently disables the automatic residency split on hosts where CUDA pinning is + # quota-capped (WSL caps it near half of RAM). The symptom is not a clear refusal + # but cudaHostRegister failing partway through the banks. + from freetoken.moe.cpu_executor import _GGML_TO_CPU_FMT + + types = getattr(model_config, "gguf_expert_types", None) + if not types: + return False + gate_up, down = int(types[0]), int(types[1]) + # one weight_format serves both banks, so mixed types cannot run on the CPU path + return gate_up == down and gate_up in _GGML_TO_CPU_FMT return fmt == "mxfp4" or fmt in _WFMT_IDS diff --git a/python/freetoken/kernel/csrc/cpu_moe/cpu_moe_ext.cpp b/python/freetoken/kernel/csrc/cpu_moe/cpu_moe_ext.cpp index 68e6324c..48210b9d 100644 --- a/python/freetoken/kernel/csrc/cpu_moe/cpu_moe_ext.cpp +++ b/python/freetoken/kernel/csrc/cpu_moe/cpu_moe_ext.cpp @@ -1650,6 +1650,13 @@ struct CpuMoeExecutor { const uint8_t* w = gu_packed_l + ((size_t)e * (2 * I) + row) * (size_t)q4_gu_row_bytes; return q6_k_dot_f32_scalar(w, x, H); // W4A16: bf16 activations, K-quant dequant } + // Anything that reaches here is assumed NVFP4 and dereferences the scale/global + // pointers, which are null for formats that do not have them (the GGUF banks pass 0). + // Falling through with an unhandled format therefore segfaults inside the worker + // thread rather than reporting anything useful, so reject it here instead. + TORCH_CHECK(fmt == WF_NVFP4 || fmt == WF_DSFP4, + "cpu_moe gemm1_dot: unhandled weight_format ", fmt, + " (handled: bf16=0, nvfp4=1, mxfp4=2, dsfp4=3, q4_0=4, q4_k=5, q6_k=6)"); const size_t r = (size_t)e * (2 * I) + row; if (use_vnni) return nvi8dot(gu_packed_l + r * (size_t)(H / 2), gu_scale_l + r * (size_t)(H / 16), diff --git a/python/freetoken/layers/gguf.py b/python/freetoken/layers/gguf.py index bc4754e8..5936d34d 100644 --- a/python/freetoken/layers/gguf.py +++ b/python/freetoken/layers/gguf.py @@ -280,7 +280,12 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: flat = x.flatten() rows = self.qweight.index_select(0, flat) # [n, row_bytes] packed - y = ggml_dequantize(rows, self._quant_type, flat.shape[0], self.embedding_dim, torch.bfloat16) + if self._quant_type in GGML_UNQUANTIZED: + # Raw value bytes, not blocks: there is no dequant kernel for the unquantized + # types (ggml_dequantize rejects type 1), so reinterpret the gathered rows. + y = rows.view(_UNQUANTIZED_DTYPE[self._quant_type]).to(torch.bfloat16) + else: + y = ggml_dequantize(rows, self._quant_type, flat.shape[0], self.embedding_dim, torch.bfloat16) y = y.view(*x.shape, self.embedding_dim) if self._embed_scale is not None: if self._embed_scale_t is None: