diff --git a/python/freetoken/engine/engine.py b/python/freetoken/engine/engine.py index 326c467d6d..ee6558014d 100644 --- a/python/freetoken/engine/engine.py +++ b/python/freetoken/engine/engine.py @@ -1770,6 +1770,9 @@ def override(attr: str, value: Any): # this is dangerous, use with caution "--dtype float16 with MXFP8 resident weights is unsupported (the " "W8A16 fold is only validated exact in bfloat16); use bfloat16." ) + if _dtype == torch.float16 and getattr(getattr(model_config, "quant", None), "weight_scales", False): + # GGUF NVFP4 results are stored in the activation dtype before weight_scale_2 (~1e-4) multiplies them, which overflows fp16. + raise ValueError("--dtype float16 with NVFP4 GGUF weight scales is unsupported (the unscaled results overflow float16); use bfloat16.") if config.attention_backend == "auto": override( "attention_backend", 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 c4887efbff..8162138ceb 100644 --- a/python/freetoken/kernel/csrc/cpu_moe/cpu_moe_ext.cpp +++ b/python/freetoken/kernel/csrc/cpu_moe/cpu_moe_ext.cpp @@ -48,6 +48,7 @@ std::unique_ptr make_mxfp4_format(const MoeParams& p, const FormatArg std::unique_ptr make_ds_fp4_format(const MoeParams& p, const FormatArgs& a); std::unique_ptr make_fp8_block_format(const MoeParams& p, const FormatArgs& a); std::unique_ptr make_gguf_format(const MoeParams& p, const FormatArgs& a); +std::unique_ptr make_gguf_scaled_format(const MoeParams& p, const FormatArgs& a); void set_ggml_cpu(uintptr_t addr); namespace { @@ -97,6 +98,7 @@ inline MoeTask from_word(uint64_t w) { const std::pair (*)(const MoeParams&, const FormatArgs&)> kFormats[] = { {"bf16", make_bf16_format}, {"nvfp4", make_nvfp4_format}, {"mxfp4_triton", make_mxfp4_format}, {"ds_fp4", make_ds_fp4_format}, {"fp8_block", make_fp8_block_format}, {"gguf", make_gguf_format}, + {"gguf_scaled", make_gguf_scaled_format}, }; std::unique_ptr make_format(const std::string& name, const MoeParams& p, const FormatArgs& a) { diff --git a/python/freetoken/kernel/csrc/cpu_moe/kernels/gguf.cpp b/python/freetoken/kernel/csrc/cpu_moe/kernels/gguf.cpp index c1641af51d..63c8107dd3 100644 --- a/python/freetoken/kernel/csrc/cpu_moe/kernels/gguf.cpp +++ b/python/freetoken/kernel/csrc/cpu_moe/kernels/gguf.cpp @@ -31,10 +31,11 @@ void set_ggml_cpu(uintptr_t addr) { namespace { // Packed ggml rows of the layers whose experts have this gate_up type and this down type, dotted by ggml-cpu against the activation converted to each type's vec_dot input (q8_0, q8_1 or q8_K, or the float type itself). +// With scales, each expert's (gate, up, down, 1) NVFP4 weight_scale_2 multiplies its gate and up rows before the glu and its down rows before the routing weight, as llama.cpp's build_moe_ffn does. class GgufFormat final : public MoeFormat { public: - GgufFormat(const MoeParams& p, const FormatArgs& a) - : MoeFormat(p), ggml_(ggml_api()), gate_up_(a.bank("gate_up")), down_(a.bank("down")) { + GgufFormat(const MoeParams& p, const FormatArgs& a, bool scaled) + : MoeFormat(p), ggml_(ggml_api()), gate_up_(a.bank("gate_up")), down_(a.bank("down")), scales_(scaled ? a.bank("global") : nullptr) { static std::once_flag tables; std::call_once(tables, ggml_.init); gu_ = kernel(static_cast(a.option("gate_up_type")), p.H, "gate_up"); @@ -63,10 +64,11 @@ class GgufFormat final : public MoeFormat { void gate_up(int layer, int e, int tok, int i0, int i1, const bf16_t* /*x*/, float w_in, bf16_t* g) const override { const uint8_t* w = static_cast(tbl_at(gate_up_, layer)) + static_cast(e) * 2 * p.I * gu_.row_bytes; const uint8_t* xq = xq_.data() + static_cast(tok) * gu_.act_row_bytes; + const float* sc = expert_scales(layer, e); gated_rows(p, i0, i1, w_in, g, [&](int row) { float s; gu_.vec_dot(p.H, &s, 0, w + static_cast(row) * gu_.row_bytes, 0, xq, 0, 1); - return s; + return sc == nullptr ? s : s * sc[row < p.I ? 0 : 1]; }); } @@ -76,9 +78,19 @@ class GgufFormat final : public MoeFormat { for (int j = 0; j < n; ++j) dn_.vec_dot(p.I, &out[j * kRowBlock + h - h0], 0, w + (static_cast(routes[j].e) * p.H + h) * dn_.row_bytes, 0, gq_.data() + static_cast(routes[j].r) * dn_.act_row_bytes, 0, 1); + if (scales_ == nullptr) return; + for (int j = 0; j < n; ++j) { + const float s = expert_scales(layer, routes[j].e)[2]; + for (int h = h0; h < h1; ++h) out[j * kRowBlock + h - h0] *= s; + } } private: + // the expert's (gate, up, down, 1) row of the layer's [E, 4] fp32 scale bank, or nullptr without scales + const float* expert_scales(int layer, int e) const { + return scales_ == nullptr ? nullptr : static_cast(tbl_at(scales_, layer)) + static_cast(e) * 4; + } + ft_ggml_row_kernel kernel(int type, int k, const char* role) const { ft_ggml_row_kernel out{}; const int rc = ggml_.row_kernel(type, k, &out); @@ -92,6 +104,7 @@ class GgufFormat final : public MoeFormat { const ft_ggml_cpu_api& ggml_; const uint64_t* gate_up_; const uint64_t* down_; + const uint64_t* scales_; ft_ggml_row_kernel gu_{}, dn_{}; std::vector xf_; // the activation in the vec_dot input type: x per token, g per route @@ -101,7 +114,11 @@ class GgufFormat final : public MoeFormat { } // namespace std::unique_ptr make_gguf_format(const MoeParams& p, const FormatArgs& a) { - return std::make_unique(p, a); + return std::make_unique(p, a, false); +} + +std::unique_ptr make_gguf_scaled_format(const MoeParams& p, const FormatArgs& a) { + return std::make_unique(p, a, true); } } // namespace freetoken::cpu_moe diff --git a/python/freetoken/kernel/csrc/ggml/VERSION b/python/freetoken/kernel/csrc/ggml/VERSION index 3c66f9b7d0..f7b5c01705 100644 --- a/python/freetoken/kernel/csrc/ggml/VERSION +++ b/python/freetoken/kernel/csrc/ggml/VERSION @@ -10,5 +10,6 @@ Local changes, kept in freetoken.patch (already applied): - mmvq.cu/.cuh: MMVQ takes bf16 or fp16 activations, can write bf16 or fp16 results, and can reuse activations already quantized to q8_1 (src1_q8_1_pre). - mmq.cu/.cuh, mmid.cu/.cuh: MoE MMQ runs a flat grid over only the (expert, column tile) pairs that hold tokens, listed by mm_ids_tile_map; ncols_max > 0 keeps the regular grid. - common.cuh: the native fp4 MMA (BLACKWELL_MMA_AVAILABLE, blackwell_mma_available) is off unless GGML_CUDA_FP4_MMA is defined; it needs an sm_120a/sm_120f build. +- mmvq.cu, common.cuh: the NVFP4 scale fusion (x_scale, gate_scale) reads strided scale vectors, e.g. columns of a per-expert scale bank. To update: copy the same files from a newer llama.cpp, apply freetoken.patch, and update this file. diff --git a/python/freetoken/kernel/csrc/ggml/freetoken.patch b/python/freetoken/kernel/csrc/ggml/freetoken.patch index 3dc715fe0f..595498dd15 100644 --- a/python/freetoken/kernel/csrc/ggml/freetoken.patch +++ b/python/freetoken/kernel/csrc/ggml/freetoken.patch @@ -24,6 +24,16 @@ } // Checks whether the tensor's base data pointer and higher-dimensional strides are byte-aligned to `alignment` bytes. +@@ -1582,6 +1588,9 @@ + const void * gate_bias = nullptr; + const void * x_scale = nullptr; + const void * gate_scale = nullptr; ++ // FreeToken: element strides of x_scale / gate_scale, so they can be columns of a per-expert scale bank ++ uint32_t x_scale_stride = 1; ++ uint32_t gate_scale_stride = 1; + ggml_glu_op glu_op; + float glu_limit = 0.0f; + const void * shared_up = nullptr; --- a/ggml/src/ggml-cuda/mmid.cu +++ b/ggml/src/ggml-cuda/mmid.cu @@ -177,3 +177,45 @@ @@ -572,6 +582,19 @@ stride_col_dst = fusion.shared_stride_col_dst; } +@@ -702,10 +706,10 @@ + } + if constexpr (type == GGML_TYPE_NVFP4) { + if (use_scale) { +- x_scales = x_scale[ids ? channel_x : 0]; ++ x_scales = x_scale[(ids ? channel_x : 0) * fusion.x_scale_stride]; + } + if (use_gate_scale) { +- gate_scales = gate_scale[ids ? channel_x : 0]; ++ gate_scales = gate_scale[(ids ? channel_x : 0) * fusion.gate_scale_stride]; + } + } + } @@ -837,7 +841,7 @@ } } @@ -615,6 +638,24 @@ stride_col_dst = fusion.shared_stride_col_dst; } +@@ -963,7 +969,7 @@ + + if constexpr (type == GGML_TYPE_NVFP4) { + if (x_scale) { +- result *= x_scale[channel_x]; ++ result *= x_scale[channel_x * fusion.x_scale_stride]; + } + } + if (x_bias) { +@@ -973,7 +979,7 @@ + float gate_value = tmp_gate[threadIdx.x]; + if constexpr (type == GGML_TYPE_NVFP4) { + if (gate_scale) { +- gate_value *= gate_scale[channel_x]; ++ gate_value *= gate_scale[channel_x * fusion.gate_scale_stride]; + } + } + if (gate_bias) { @@ -998,7 +1004,7 @@ } } @@ -718,7 +759,30 @@ GGML_ASSERT(!ids || ids->type == GGML_TYPE_I32); // Optional, used for batched GGML_MUL_MAT_ID. GGML_TENSOR_BINARY_OP_LOCALS; -@@ -1526,12 +1535,22 @@ +@@ -1498,17 +1507,20 @@ + GGML_ASSERT(!ids || fusion->gate_bias->ne[1] == src0->ne[2]); + fusion_local.gate_bias = fusion->gate_bias->data; + } ++ // FreeToken: a scale may be a strided vector, read with its element stride + if (fusion->x_scale) { + GGML_ASSERT(fusion->x_scale->type == GGML_TYPE_F32); +- GGML_ASSERT(ggml_is_contiguous(fusion->x_scale)); ++ GGML_ASSERT(ggml_nelements(fusion->x_scale) == fusion->x_scale->ne[0] && fusion->x_scale->nb[0] % sizeof(float) == 0); + GGML_ASSERT(ggml_nelements(fusion->x_scale) == (ids ? src0->ne[2] : 1)); + fusion_local.x_scale = fusion->x_scale->data; ++ fusion_local.x_scale_stride = fusion->x_scale->nb[0] / sizeof(float); + } + if (fusion->gate_scale) { + GGML_ASSERT(fusion->gate_scale->type == GGML_TYPE_F32); +- GGML_ASSERT(ggml_is_contiguous(fusion->gate_scale)); ++ GGML_ASSERT(ggml_nelements(fusion->gate_scale) == fusion->gate_scale->ne[0] && fusion->gate_scale->nb[0] % sizeof(float) == 0); + GGML_ASSERT(ggml_nelements(fusion->gate_scale) == (ids ? src0->ne[2] : 1)); + fusion_local.gate_scale = fusion->gate_scale->data; ++ fusion_local.gate_scale_stride = fusion->gate_scale->nb[0] / sizeof(float); + } + fusion_local.glu_op = fusion->glu_op; + fusion_local.glu_limit = fusion->glu_limit; +@@ -1526,12 +1538,22 @@ } const int64_t ne10_padded = GGML_PAD(ne10, MATRIX_ROW_PADDING); @@ -744,7 +808,7 @@ } const int64_t s01 = src0->nb[1] / ts_src0; -@@ -1556,11 +1575,23 @@ +@@ -1556,11 +1578,23 @@ const int64_t ids_stride = ids ? ids->nb[1] / ggml_type_size(ids->type) : 0; diff --git a/python/freetoken/kernel/csrc/ggml/freetoken/ggml_ops.cu b/python/freetoken/kernel/csrc/ggml/freetoken/ggml_ops.cu index 69d9322736..170cbcac71 100644 --- a/python/freetoken/kernel/csrc/ggml/freetoken/ggml_ops.cu +++ b/python/freetoken/kernel/csrc/ggml/freetoken/ggml_ops.cu @@ -115,6 +115,14 @@ int64_t check_weight(TensorView w, ggml_type type, int64_t k) { return row_bytes; } +// one fp32 scale per expert, e.g. a column of a [E, n] scale bank; the MMVQ epilogue multiplies NVFP4 results by it +ggml_tensor describe_scale(TensorView s, ggml_type type, int64_t e) { + RuntimeCheck(type == GGML_TYPE_NVFP4, "per-expert scales apply to NVFP4 weights, not ", ggml_type_name(type)); + RuntimeCheck(is_dtype(s.dtype(), kDLFloat, 32) && s.ndim() == 1 && s.size(0) == e && s.stride(0) > 0, "a per-expert scale must be an fp32 [E] view"); + const size_t nb0 = sizeof(float) * s.stride(0); + return describe(GGML_TYPE_F32, s.data_ptr(), {e, 1, 1, 1}, {nb0, nb0 * e, nb0 * e, nb0 * e}); +} + int64_t q8_1_row_bytes(int64_t k) { return GGML_PAD(k, MATRIX_ROW_PADDING) / QK8_1 * (int64_t) sizeof(block_q8_1); } // MMQ tiles read up to MATRIX_ROW_PADDING values past the last row; the caller keeps this many bytes after it @@ -234,7 +242,8 @@ void op_mul_mat(TensorView w, int64_t type_id, TensorView x, TensorView out, Opt } // out[T, k, N]: row (t, j) multiplies x[t, j or 0] by expert ids[t, j] of W[E, N, K]; the matrix kernel runs a flat grid over the tiles that hold routes -void op_mul_mat_id(TensorView w, int64_t type_id, TensorView x, TensorView ids, TensorView out) { +// scale[ids[t, j]] multiplies row (t, j); only the vector kernel takes it +void op_mul_mat_id(TensorView w, int64_t type_id, TensorView x, TensorView ids, TensorView out, Optional scale) { RuntimeCheck(w.ndim() == 3, "mul_mat_id takes w[E, N, row_bytes]"); RuntimeCheck(is_dtype(ids.dtype(), kDLInt, 32) && ids.ndim() == 2 && ids.IsContiguous(), "ids must be a contiguous int32 [T, k] tensor"); const DeviceGuard guard(x.device()); @@ -264,8 +273,15 @@ void op_mul_mat_id(TensorView w, int64_t type_id, TensorView x, TensorView ids, ggml_tensor src1 = describe(a.type, a.data, {k, ky, t, 1}, {a.esz, a.esz * k, a.esz * k * ky, a.esz * k * ky * t}); const size_t oe = ggml_type_size(otype); ggml_tensor dst = describe(otype, out.data_ptr(), {n, top_k, t, 1}, {oe, oe * n, oe * n * top_k, oe * n * top_k * t}); + RuntimeCheck(mmvq || !scale.has_value(), "mul_mat_id takes a per-expert scale only on the vector kernel, T=", t); if (mmvq) { - ggml_cuda_mul_mat_vec_q(ctx, &src0, &src1, &tids, &dst); + ggml_cuda_mm_fusion_args_host fusion{}; + ggml_tensor tscale; + if (scale.has_value()) { + tscale = describe_scale(scale.value(), type, e); + fusion.x_scale = &tscale; + } + ggml_cuda_mul_mat_vec_q(ctx, &src0, &src1, &tids, &dst, scale.has_value() ? &fusion : nullptr); } else { ggml_cuda_mul_mat_q(ctx, &src0, &src1, &tids, &dst, 0); } @@ -291,7 +307,9 @@ ggml_glu_op fused_glu_op(const String & name) { } // namespace // out[T, k, I] = glu(x @ gate^T, x @ up^T) for small T in one MMVQ pass; w[E, 2I, row_bytes] holds each expert's gate rows then its up rows; limit is swiglu_clamp's -void op_mul_mat_id_glu(TensorView w, int64_t type_id, TensorView x, TensorView ids, TensorView out, String glu, double limit) { +// gate_scale / up_scale multiply the expert's gate / up results before the glu +void op_mul_mat_id_glu(TensorView w, int64_t type_id, TensorView x, TensorView ids, TensorView out, String glu, double limit, + Optional gate_scale, Optional up_scale) { const ggml_glu_op glu_op = fused_glu_op(glu); RuntimeCheck(w.ndim() == 3 && w.size(1) % 2 == 0, "mul_mat_id_glu takes w[E, 2I, row_bytes]"); RuntimeCheck(is_dtype(ids.dtype(), kDLInt, 32) && ids.ndim() == 2 && ids.IsContiguous(), "ids must be a contiguous int32 [T, k] tensor"); @@ -325,6 +343,15 @@ void op_mul_mat_id_glu(TensorView w, int64_t type_id, TensorView x, TensorView i fusion.gate = &gate; fusion.glu_op = glu_op; fusion.glu_limit = (float) limit; + ggml_tensor tgate_scale, tup_scale; + if (gate_scale.has_value()) { + tgate_scale = describe_scale(gate_scale.value(), type, e); + fusion.gate_scale = &tgate_scale; + } + if (up_scale.has_value()) { + tup_scale = describe_scale(up_scale.value(), type, e); + fusion.x_scale = &tup_scale; + } ggml_cuda_mul_mat_vec_q(ctx, &up, &src1, &tids, &dst, &fusion); } diff --git a/python/freetoken/kernel/csrc/ggml/ggml-cuda/common.cuh b/python/freetoken/kernel/csrc/ggml/ggml-cuda/common.cuh index 36912ce4b6..a767c688f5 100644 --- a/python/freetoken/kernel/csrc/ggml/ggml-cuda/common.cuh +++ b/python/freetoken/kernel/csrc/ggml/ggml-cuda/common.cuh @@ -1588,6 +1588,9 @@ struct ggml_cuda_mm_fusion_args_device { const void * gate_bias = nullptr; const void * x_scale = nullptr; const void * gate_scale = nullptr; + // FreeToken: element strides of x_scale / gate_scale, so they can be columns of a per-expert scale bank + uint32_t x_scale_stride = 1; + uint32_t gate_scale_stride = 1; ggml_glu_op glu_op; float glu_limit = 0.0f; const void * shared_up = nullptr; diff --git a/python/freetoken/kernel/csrc/ggml/ggml-cuda/mmvq.cu b/python/freetoken/kernel/csrc/ggml/ggml-cuda/mmvq.cu index f3a3ccc683..8487a035c1 100644 --- a/python/freetoken/kernel/csrc/ggml/ggml-cuda/mmvq.cu +++ b/python/freetoken/kernel/csrc/ggml/ggml-cuda/mmvq.cu @@ -706,10 +706,10 @@ static __global__ void mul_mat_vec_q( } if constexpr (type == GGML_TYPE_NVFP4) { if (use_scale) { - x_scales = x_scale[ids ? channel_x : 0]; + x_scales = x_scale[(ids ? channel_x : 0) * fusion.x_scale_stride]; } if (use_gate_scale) { - gate_scales = gate_scale[ids ? channel_x : 0]; + gate_scales = gate_scale[(ids ? channel_x : 0) * fusion.gate_scale_stride]; } } } @@ -969,7 +969,7 @@ static __global__ void mul_mat_vec_q_moe( if constexpr (type == GGML_TYPE_NVFP4) { if (x_scale) { - result *= x_scale[channel_x]; + result *= x_scale[channel_x * fusion.x_scale_stride]; } } if (x_bias) { @@ -979,7 +979,7 @@ static __global__ void mul_mat_vec_q_moe( float gate_value = tmp_gate[threadIdx.x]; if constexpr (type == GGML_TYPE_NVFP4) { if (gate_scale) { - gate_value *= gate_scale[channel_x]; + gate_value *= gate_scale[channel_x * fusion.gate_scale_stride]; } } if (gate_bias) { @@ -1507,17 +1507,20 @@ void ggml_cuda_mul_mat_vec_q( GGML_ASSERT(!ids || fusion->gate_bias->ne[1] == src0->ne[2]); fusion_local.gate_bias = fusion->gate_bias->data; } + // FreeToken: a scale may be a strided vector, read with its element stride if (fusion->x_scale) { GGML_ASSERT(fusion->x_scale->type == GGML_TYPE_F32); - GGML_ASSERT(ggml_is_contiguous(fusion->x_scale)); + GGML_ASSERT(ggml_nelements(fusion->x_scale) == fusion->x_scale->ne[0] && fusion->x_scale->nb[0] % sizeof(float) == 0); GGML_ASSERT(ggml_nelements(fusion->x_scale) == (ids ? src0->ne[2] : 1)); fusion_local.x_scale = fusion->x_scale->data; + fusion_local.x_scale_stride = fusion->x_scale->nb[0] / sizeof(float); } if (fusion->gate_scale) { GGML_ASSERT(fusion->gate_scale->type == GGML_TYPE_F32); - GGML_ASSERT(ggml_is_contiguous(fusion->gate_scale)); + GGML_ASSERT(ggml_nelements(fusion->gate_scale) == fusion->gate_scale->ne[0] && fusion->gate_scale->nb[0] % sizeof(float) == 0); GGML_ASSERT(ggml_nelements(fusion->gate_scale) == (ids ? src0->ne[2] : 1)); fusion_local.gate_scale = fusion->gate_scale->data; + fusion_local.gate_scale_stride = fusion->gate_scale->nb[0] / sizeof(float); } fusion_local.glu_op = fusion->glu_op; fusion_local.glu_limit = fusion->glu_limit; diff --git a/python/freetoken/kernel/ggml.py b/python/freetoken/kernel/ggml.py index b672d58cf2..7a2765c5c0 100644 --- a/python/freetoken/kernel/ggml.py +++ b/python/freetoken/kernel/ggml.py @@ -95,24 +95,37 @@ def uses_mmvq(ggml_type: int, m: int) -> bool: return _module().uses_mmvq(ggml_type, m) -def mul_mat_id(w: torch.Tensor, ggml_type: int, x: torch.Tensor, ids: torch.Tensor) -> torch.Tensor: - """Per-route expert GEMM over ``w[E, N, row_bytes]``: returns ``[T, k, N]`` in ``x``'s dtype; ``x`` is ``[T, K]`` (shared by the k routes) or ``[T, k, K]``.""" +def mul_mat_id(w: torch.Tensor, ggml_type: int, x: torch.Tensor, ids: torch.Tensor, scale: torch.Tensor | None = None) -> torch.Tensor: + """Per-route expert GEMM over ``w[E, N, row_bytes]``: returns ``[T, k, N]`` in ``x``'s dtype; ``x`` is ``[T, K]`` (shared by the k routes) or ``[T, k, K]``. + + ``scale`` is an fp32 ``[E]`` view (strided is fine) that multiplies each NVFP4 expert's rows; it needs ``T <= mmvq_max_tokens``. + """ t, top_k = ids.shape out = torch.empty(t, top_k, w.shape[1], dtype=x.dtype, device=x.device) if t > mmvq_max_tokens(ggml_type): _check_mmq_tail(w, ggml_type, x.shape[-1]) - _module().mul_mat_id(w, ggml_type, (x.unsqueeze(1) if x.dim() == 2 else x).contiguous(), ids, out) + _module().mul_mat_id(w, ggml_type, (x.unsqueeze(1) if x.dim() == 2 else x).contiguous(), ids, out, scale) return out -def mul_mat_id_glu(w: torch.Tensor, ggml_type: int, x: torch.Tensor, ids: torch.Tensor, glu: str, limit: float = 0.0) -> torch.Tensor: +def mul_mat_id_glu( + w: torch.Tensor, + ggml_type: int, + x: torch.Tensor, + ids: torch.Tensor, + glu: str, + limit: float = 0.0, + gate_scale: torch.Tensor | None = None, + up_scale: torch.Tensor | None = None, +) -> torch.Tensor: """``glu(gate, up)`` per route in one pass for small ``T``: ``w[E, 2I, row_bytes]`` holds gate rows then up rows; returns ``[T, k, I]`` in ``x``'s dtype. ``glu`` is one of ``swiglu``, ``geglu`` (tanh GELU), ``swiglu_oai`` (alpha 1.702, limit 7) and ``swiglu_clamp`` (alpha 1, clamped at ``limit``). + ``gate_scale`` / ``up_scale`` are fp32 ``[E]`` views that multiply each NVFP4 expert's gate / up results before the glu. """ t, top_k = ids.shape out = torch.empty(t, top_k, w.shape[1] // 2, dtype=x.dtype, device=x.device) - _module().mul_mat_id_glu(w, ggml_type, x.contiguous(), ids, out, glu, limit) + _module().mul_mat_id_glu(w, ggml_type, x.contiguous(), ids, out, glu, limit, gate_scale, up_scale) return out diff --git a/python/freetoken/layers/quantization/configs/gguf.py b/python/freetoken/layers/quantization/configs/gguf.py index bbc24f137a..77859cd1f7 100644 --- a/python/freetoken/layers/quantization/configs/gguf.py +++ b/python/freetoken/layers/quantization/configs/gguf.py @@ -4,7 +4,7 @@ from typing import Any, Callable, ClassVar -from ..ggml_types import is_quantized, type_name +from ..ggml_types import has_expert_scales, has_weight_scale, is_quantized, type_name from ..names import NameMap from ..quant_backend import get_quant_backend from ..registry import LayerKind, method_class, register_dialect @@ -29,19 +29,24 @@ def __init__( super().__init__(name_map, unquantized) self.tensor_types = tensor_types self.to_gguf = to_gguf + self.expert_scales = has_expert_scales(tensor_types) + self.weight_scales = any(n.endswith(".scale") and has_weight_scale(tensor_types, n.removesuffix(".scale")) for n in tensor_types) @classmethod def claims(cls, q: dict[str, Any]) -> bool: return False - def _part_types(self, name: str) -> tuple[int, ...]: - types = [] + def _part_tensors(self, name: str) -> tuple[str, ...]: + keys = [] for tensor in self.to_gguf(name, self.tensor_types): key = f"{tensor}.weight" if f"{tensor}.weight" in self.tensor_types else tensor if key not in self.tensor_types: raise KeyError(f"GGUF has no tensor {key!r} for {name!r}") - types.append(self.tensor_types[key]) - return tuple(types) + keys.append(key) + return tuple(keys) + + def _part_types(self, name: str) -> tuple[int, ...]: + return tuple(self.tensor_types[key] for key in self._part_tensors(name)) def scheme_for_name(self, name: str) -> QuantScheme | None: types = self._part_types(name) @@ -55,7 +60,8 @@ def scheme_for(self, prefix: str) -> QuantScheme | None: if prefix in self._schemes: return self._schemes[prefix] names = self.name_map.to_checkpoint(prefix) - parts = [() if self.unquantized(n) else self._part_types(n) for n in names] + tensors = [() if self.unquantized(n) else self._part_tensors(n) for n in names] + parts = [tuple(self.tensor_types[key] for key in ts) for ts in tensors] quantized = [any(is_quantized(t) for t in p) for p in parts] if not any(quantized): scheme = None @@ -63,7 +69,8 @@ def scheme_for(self, prefix: str) -> QuantScheme | None: found = {n: [type_name(t) for t in p] for n, p in zip(names, parts)} raise NotImplementedError(f"fused module {prefix!r} needs one GGUF tensor per part, got {found}") else: - scheme = gguf_scheme(p[0] for p in parts) + scaled = any(has_weight_scale(self.tensor_types, ts[0]) for ts in tensors) + scheme = gguf_scheme((p[0] for p in parts), scaled=scaled) self._schemes[prefix] = scheme return scheme @@ -74,7 +81,8 @@ def get_quant_method(self, layer: Any, prefix: str): parts = [self._part_types(n) for n in self.name_map.to_checkpoint(prefix)] if any(len(p) != 1 for p in parts): raise NotImplementedError(f"experts {prefix!r} are stored as {parts} GGUF tensors per projection") - scheme = gguf_scheme(p[0] for p in parts) + # the scale bank is per file, so every MoE layer keeps one bank set + scheme = gguf_scheme((p[0] for p in parts), scaled=self.expert_scales) cfg = self.layer_config(layer, LayerKind.MOE, scheme) return method_class(QuantKind.GGUF, LayerKind.MOE)(cfg, get_quant_backend().select(LayerKind.MOE, QuantKind.GGUF)) diff --git a/python/freetoken/layers/quantization/embedding/gguf.py b/python/freetoken/layers/quantization/embedding/gguf.py index f41b7b6a4b..56c72c2d6a 100644 --- a/python/freetoken/layers/quantization/embedding/gguf.py +++ b/python/freetoken/layers/quantization/embedding/gguf.py @@ -7,7 +7,7 @@ import torch -from ..ggml_types import Segment, pad_rows, row_bytes +from ..ggml_types import NVFP4, Segment, pad_rows, row_bytes from ..linear.gguf import GgmlLinearKernel from ..registry import LayerKind, register_method from ..scheme import QuantKind, gguf_types @@ -18,7 +18,10 @@ class GgmlEmbeddingKernel(EmbeddingKernel): name = "ggml" def unusable_reason(self, cfg: EmbeddingConfig) -> str | None: - return "GGUF tables are not sharded over TP" if cfg.tp_size > 1 else None + if cfg.tp_size > 1: + return "GGUF tables are not sharded over TP" + # the vendored get_rows has no NVFP4 case, and the table's weight_scale_2 would need applying too + return "NVFP4 GGUF tables are not supported" if NVFP4 in gguf_types(cfg.scheme) else None def embed(self, layer: Any, ids: torch.Tensor) -> torch.Tensor: from freetoken.kernel import ggml diff --git a/python/freetoken/layers/quantization/ggml_types.py b/python/freetoken/layers/quantization/ggml_types.py index 8b257cd8c7..262c5668d3 100644 --- a/python/freetoken/layers/quantization/ggml_types.py +++ b/python/freetoken/layers/quantization/ggml_types.py @@ -2,6 +2,7 @@ from __future__ import annotations +import re from dataclasses import dataclass import torch @@ -85,4 +86,33 @@ def segment_param(index: int) -> str: return "weight" if index == 0 else f"weight_{index}" -__all__ = ["MATRIX_ROW_PADDING", "Segment", "float_dtype", "is_quantized", "mmq_tail_bytes", "pad_rows", "row_bytes", "segment_param", "segments", "type_name"] +NVFP4 = int(gguf.GGMLQuantizationType.NVFP4) + +# per-expert fp32 weight scales llama.cpp reads next to separate gate / up / down expert tensors (llama-model.cpp ffn_*_exps_s) +_EXPERT_SCALE_RE = re.compile(r"^blk\.\d+\.ffn_(gate|up|down)_exps\.scale$") + + +def scale_name(tensor: str) -> str: + """``.scale`` of the GGUF tensor ```` or ``.weight``.""" + return f"{tensor.removesuffix('.weight')}.scale" + + +def has_weight_scale(tensor_types: dict[str, int], tensor: str) -> bool: + """Whether llama.cpp's converter stored an NVFP4 tensor's weight_scale_2 as ``.scale``; other types keep no weight scale.""" + weight = tensor if tensor.endswith(".weight") else f"{tensor}.weight" + return tensor_types.get(weight, tensor_types.get(tensor)) == NVFP4 and scale_name(tensor) in tensor_types + + +def has_expert_scales(tensor_types: dict[str, int]) -> bool: + """Whether any NVFP4 expert tensor of the file carries per-expert scales, which then take a scale bank in every MoE layer.""" + merged = [n for n in tensor_types if n.endswith("ffn_gate_up_exps.scale") and tensor_types.get(n.replace(".scale", ".weight")) == NVFP4] + if merged: + # llama.cpp scales only separate gate / up expert tensors and would drop this one + raise NotImplementedError(f"per-expert scales of merged gate_up NVFP4 experts are not supported (e.g. {merged[0]})") + return any(_EXPERT_SCALE_RE.match(n) and has_weight_scale(tensor_types, n.removesuffix(".scale")) for n in tensor_types) + + +__all__ = [ + "MATRIX_ROW_PADDING", "NVFP4", "Segment", "float_dtype", "has_expert_scales", "has_weight_scale", "is_quantized", "mmq_tail_bytes", + "pad_rows", "row_bytes", "scale_name", "segment_param", "segments", "type_name", +] diff --git a/python/freetoken/layers/quantization/linear/gguf.py b/python/freetoken/layers/quantization/linear/gguf.py index ac08dbe287..41f61edf88 100644 --- a/python/freetoken/layers/quantization/linear/gguf.py +++ b/python/freetoken/layers/quantization/linear/gguf.py @@ -35,6 +35,9 @@ def apply(self, layer: Any, x: torch.Tensor) -> torch.Tensor: ggml.mul_mat(w, s.ggml_type, x2, out=out, xq=xq) else: torch.mm(x2, w.t(), out=out) + # llama.cpp multiplies an NVFP4 result by the tensor's weight_scale_2 (build_lora_mm) + if getattr(layer, "ggml_scale", None) is not None: + y.mul_(layer.ggml_scale) if layer.bias is not None: y = y + layer.bias return y.reshape(*x.shape[:-1], y.shape[-1]) @@ -54,3 +57,6 @@ def create_weights(self, layer: Any) -> None: setattr(layer, segment_param(i), torch.empty(s.rows + pad_rows(k), row_bytes(s.ggml_type, k), dtype=torch.uint8)) else: setattr(layer, segment_param(i), torch.empty(s.rows, k)) + if g.scheme.weight.scale is not None: + # one value per output row: each part's weight_scale_2, 1 for parts without one + layer.ggml_scale = torch.empty(sum(s.rows for s in segs), dtype=torch.float32) diff --git a/python/freetoken/layers/quantization/moe/gguf.py b/python/freetoken/layers/quantization/moe/gguf.py index 416c2f4b58..01e754d644 100644 --- a/python/freetoken/layers/quantization/moe/gguf.py +++ b/python/freetoken/layers/quantization/moe/gguf.py @@ -6,7 +6,7 @@ import torch -from ..ggml_types import float_dtype, is_quantized, row_bytes, type_name +from ..ggml_types import NVFP4, float_dtype, is_quantized, row_bytes, type_name from ..registry import LayerKind, register_method from ..scheme import QuantKind, gguf_types from .base import BankSpec, ExpertView, MoEConfig, MoEKernel, MoEMethod, gated_epilogue_reason, is_resident, limit_or_inf @@ -80,7 +80,11 @@ def unusable_reason(self, cfg: MoEConfig) -> str | None: def layout(self, cfg: MoEConfig) -> dict[str, BankSpec]: gate_up, down = expert_types(cfg) i, h = cfg.local_intermediate, cfg.hidden - return {"gate_up": _bank_spec(gate_up, 2 * i, h), "down": _bank_spec(down, h, i)} + banks = {"gate_up": _bank_spec(gate_up, 2 * i, h), "down": _bank_spec(down, h, i)} + if cfg.scheme.weight.scale is not None: + # each expert's NVFP4 weight_scale_2 as (gate, up, down, 1): a 16-byte row, which the slot cache's copies need + banks["global"] = BankSpec((4,), torch.float32) + return banks def pack(self, pieces, cfg: MoEConfig, out): i = cfg.local_intermediate @@ -113,17 +117,29 @@ def mul_mat_id(w: torch.Tensor, ggml_type: int, a: torch.Tensor) -> torch.Tensor return ggml.mul_mat_id(w, ggml_type, a, ids) return float_mul_mat_id(w, a, ids, topk_weights, config) + # llama.cpp's build_moe_ffn: each NVFP4 GEMM result times the expert's weight_scale_2, gate and up before the glu, down before the routing weight + scales = view.tensors.get("global") + gate_up_scales = scales if scales is not None and gate_up_type == NVFP4 else None + down_scales = scales if scales is not None and down_type == NVFP4 else None glu = fused_glu(layer) small = x.shape[0] <= min(ggml.mmvq_max_tokens(t) for t in (gate_up_type, down_type) if is_quantized(t)) if glu is not None and small and is_quantized(gate_up_type): - inter = ggml.mul_mat_id_glu(gate_up, gate_up_type, x, ids, *glu) + gu = (None, None) if gate_up_scales is None else (gate_up_scales[:, 0], gate_up_scales[:, 1]) + inter = ggml.mul_mat_id_glu(gate_up, gate_up_type, x, ids, *glu, *gu) else: h = mul_mat_id(gate_up, gate_up_type, x) t, k, two_i = h.shape + if gate_up_scales is not None: + h.view(t, k, 2, two_i // 2).mul_(gate_up_scales[ids, :2].unsqueeze(-1)) inter = torch.empty(t, k, two_i // 2, dtype=h.dtype, device=h.device) gated_act_and_mul(layer.activation, h.view(-1, two_i), inter.view(-1, two_i // 2), alpha=float(layer.alpha), limit=limit_or_inf(layer)) del h # free the gate_up output before the down GEMM allocates its own - y = mul_mat_id(view.tensors["down"], down_type, inter) + if down_scales is not None and x.shape[0] <= ggml.mmvq_max_tokens(down_type): + y = ggml.mul_mat_id(view.tensors["down"], down_type, inter, ids, scale=down_scales[:, 2]) + else: + y = mul_mat_id(view.tensors["down"], down_type, inter) + if down_scales is not None: + topk_weights = topk_weights * down_scales[ids, 2] return ggml.moe_weighted_sum(y, topk_weights) @staticmethod @@ -147,4 +163,6 @@ def cpu_spec(self) -> tuple[str | None, dict[str, int]]: # BF16 experts keep the executor's bf16 kernels, which use AVX512-BF16 where the CPU has it; ggml-cpu is built for AVX2 only if float_dtype(gate_up) is torch.bfloat16 and float_dtype(down) is torch.bfloat16: return "bf16", {} - return "gguf", {"gate_up_type": gate_up, "down_type": down} + # an NVFP4 layer of a file with weight scales reads the layout's global bank too + scaled = self.cfg.scheme.weight.scale is not None and NVFP4 in (gate_up, down) + return "gguf_scaled" if scaled else "gguf", {"gate_up_type": gate_up, "down_type": down} diff --git a/python/freetoken/layers/quantization/scheme.py b/python/freetoken/layers/quantization/scheme.py index ab43814e6a..a3aff0bc0f 100644 --- a/python/freetoken/layers/quantization/scheme.py +++ b/python/freetoken/layers/quantization/scheme.py @@ -108,9 +108,11 @@ def mxfp4_scheme() -> QuantScheme: return QuantScheme(QuantKind.MXFP4, WeightDesc("e2m1", (1, MX_GROUP), "e8m0"), {"weight", "weight_scale"}) -def gguf_scheme(types: Iterable[int]) -> QuantScheme: - """Packed ggml blocks, one ggml type per fused part (``types`` follows the parts' order).""" - return QuantScheme(QuantKind.GGUF, WeightDesc("ggml:" + ",".join(str(int(t)) for t in types), None, None), {"weight"}) +def gguf_scheme(types: Iterable[int], scaled: bool = False) -> QuantScheme: + """Packed ggml blocks, one ggml type per fused part (``types`` follows the parts' order). + + ``scaled``: the NVFP4 tensors carry the fp32 ``.scale`` (weight_scale_2) that llama.cpp's converter writes.""" + return QuantScheme(QuantKind.GGUF, WeightDesc("ggml:" + ",".join(str(int(t)) for t in types), None, "f32" if scaled else None), {"weight"}) def gguf_types(scheme: QuantScheme) -> tuple[int, ...]: diff --git a/python/freetoken/models/gemma4/gguf.py b/python/freetoken/models/gemma4/gguf.py index 2b86cca127..3d4de920b8 100644 --- a/python/freetoken/models/gemma4/gguf.py +++ b/python/freetoken/models/gemma4/gguf.py @@ -24,7 +24,7 @@ SWAAttentionGroupConfig, ) from freetoken.models.gguf.reader import gguf_tensor_index -from freetoken.models.gguf.tensors import float_tensor, fused, matrix +from freetoken.models.gguf.tensors import check_scales, float_tensor, fused, matrix, weight_scale if TYPE_CHECKING: from freetoken.models.gguf.config import GgufConfigShim @@ -48,6 +48,7 @@ def _full_rotary_dim(shim: "GgufConfigShim", full_head_dim: int) -> int: def parse_gguf_config(shim: "GgufConfigShim") -> ModelConfig: + check_scales(shim.model_path, known=("ffn_gate_inp.scale", "ffn_down_exps.scale")) m = shim.metadata def g(key: str): @@ -246,6 +247,9 @@ def t(name: str): blk, base = f"blk.{layer}", f"model.layers.{layer}" for suffix, rel in scalar_map.items(): tensor = float_tensor(index[f"{blk}.{suffix}"]) + if suffix == "ffn_down_exps.scale" and weight_scale(index[f"{blk}.ffn_down_exps.weight"]) is not None: + # the converter folds per_expert_scale into the NVFP4 down weight_scale_2, which the expert kernel applies + tensor = torch.ones_like(tensor) yield f"{base}.{rel}", tensor.reshape(1) if rel.endswith("layer_scalar") else tensor if f"{blk}.attn_qkv.weight" in index: # convert_hf_to_gguf.py --fuse-qkv stores q, k and v as one attn_qkv matrix diff --git a/python/freetoken/models/gguf/experts.py b/python/freetoken/models/gguf/experts.py index 0da56ceb4f..c21fff9fa2 100644 --- a/python/freetoken/models/gguf/experts.py +++ b/python/freetoken/models/gguf/experts.py @@ -1,7 +1,8 @@ """Routed-expert banks straight from a GGUF's ``blk.N.ffn_{gate,up,down}_exps`` tensors, in their packed ggml rows: pinned host banks for the offload cache, or device banks for resident experts. Each layer keeps its own ggml types, so each layer gets its own bank shapes: ``gate_up`` [E, 2I, row_bytes] (expert e's gate rows, then its up rows) and ``down`` [E, H, row_bytes]; float experts keep their dtype ([E, 2I, H] and [E, H, I]). -A GGUF that stores ``ffn_gate_up_exps`` (gemma4) already holds the gate_up rows in that order.""" +A GGUF that stores ``ffn_gate_up_exps`` (gemma4) already holds the gate_up rows in that order. +When NVFP4 experts carry per-expert ``.scale`` tensors, every layer also gets ``global`` [E, 4] fp32: the (gate, up, down, 1) weight_scale_2, 1 where a tensor has none.""" from __future__ import annotations @@ -11,10 +12,11 @@ import torch from freetoken.layers.quantization import QuantKind -from freetoken.layers.quantization.ggml_types import float_dtype, mmq_tail_bytes, pad_rows, type_name +from freetoken.layers.quantization.ggml_types import float_dtype, has_expert_scales, mmq_tail_bytes, pad_rows, type_name from freetoken.utils import init_logger from .reader import GgufTensorInfo, gguf_tensor_index +from .tensors import weight_scale logger = init_logger(__name__) @@ -46,14 +48,30 @@ def _device_bank(shape: tuple[int, ...], k: int, dtype: torch.dtype, device: tor return flat[: experts * rows].view(shape) +def _expert_scales(gate_up: list[GgufTensorInfo], down: GgufTensorInfo, experts: int) -> torch.Tensor: + """[E, 4] fp32 (gate, up, down, 1) weight_scale_2 of one layer's experts.""" + gate, up = gate_up if len(gate_up) == 2 else (gate_up[0], gate_up[0]) + out = torch.ones(experts, 4, dtype=torch.float32) + for col, info in enumerate((gate, up, down)): + s = weight_scale(info) + if s is None: + continue + if s.numel() != experts: + raise NotImplementedError(f"{info.name}: {s.numel()} scales for {experts} experts") + out[:, col] = s + return out + + def load_gguf_expert_banks(model_path: str, num_layers: int, *, first_layer: int = 0, dummy: bool = False, device: torch.device | None = None): """``ExpertBanks`` with one bank per (layer, role): pinned host banks for the offload cache, filled from the file and pinned while later layers load, or resident banks on ``device``.""" from freetoken.moe.expert_banks import ExpertBanks from freetoken.moe.host_banks import HostBank, PinPipeline index = gguf_tensor_index(model_path) - sources: dict[str, list[torch.Tensor]] = {"gate_up": [], "down": []} - tails: dict[str, list[int]] = {"gate_up": [], "down": []} + scaled = has_expert_scales({name: info.ggml_type for name, info in index.items()}) + roles = ("gate_up", "down", "global") if scaled else ("gate_up", "down") + sources: dict[str, list[torch.Tensor]] = {role: [] for role in roles} + tails: dict[str, list[int]] = {role: [] for role in roles} tags = [] with PinPipeline() if device is None else contextlib.nullcontext() as pins: for layer in range(num_layers): @@ -70,17 +88,26 @@ def load_gguf_expert_banks(model_path: str, num_layers: int, *, first_layer: int tags.append(("ggml", type_name(gate.ggml_type), type_name(down.ggml_type))) tails["gate_up"].append(mmq_tail_bytes(gate.ggml_type, gate.ne[0])) tails["down"].append(mmq_tail_bytes(down.ggml_type, down.ne[0])) + if scaled: + tails["global"].append(0) rows = [_expert_rows(info) for info in parts] d = _expert_rows(down) gate_up_shape = (rows[0].shape[0], sum(r.shape[1] for r in rows), rows[0].shape[2]) + experts = rows[0].shape[0] if device is None: host = {"gate_up": HostBank(gate_up_shape, rows[0].dtype), "down": HostBank(tuple(d.shape), d.dtype)} + if scaled: + host["global"] = HostBank((experts, 4), torch.float32) banks = {role: bank.tensor for role, bank in host.items()} else: banks = { "gate_up": _device_bank(gate_up_shape, gate.ne[0], rows[0].dtype, device), "down": _device_bank(tuple(d.shape), down.ne[0], d.dtype, device), } + if scaled: + banks["global"] = torch.empty(experts, 4, dtype=torch.float32, device=device) + if scaled: + banks["global"].copy_(_expert_scales(parts, down, experts)) if dummy: banks["gate_up"].zero_() banks["down"].zero_() @@ -98,7 +125,7 @@ def load_gguf_expert_banks(model_path: str, num_layers: int, *, first_layer: int sources[role].append(bank) total = sum(t.numel() * t.element_size() for ts in sources.values() for t in ts) logger.info_rank0(f"GGUF expert banks: {num_layers} layers, {total / 2**30:.1f} GiB {'on ' + str(device) if device is not None else 'in host RAM'}") - return ExpertBanks("gguf", sources, kind=QuantKind.GGUF, kernel="ggml", layout=None, layer_tags=tags, row_tail_bytes=tails) + return ExpertBanks("gguf_scaled" if scaled else "gguf", sources, kind=QuantKind.GGUF, kernel="ggml", layout=None, layer_tags=tags, row_tail_bytes=tails) __all__ = ["load_gguf_expert_banks"] diff --git a/python/freetoken/models/gguf/reader.py b/python/freetoken/models/gguf/reader.py index 2715f12909..ff7350941e 100644 --- a/python/freetoken/models/gguf/reader.py +++ b/python/freetoken/models/gguf/reader.py @@ -147,6 +147,8 @@ class GgufTensorInfo: nbytes: int ggml_type: int ne: tuple[int, ...] + # the model path the index was built from, which resolves tensors in any split + root: str = "" def data(self) -> np.ndarray: """Zero-copy uint8 view of the bytes (memory-mapped).""" @@ -161,7 +163,7 @@ def gguf_tensor_index(model_path: str) -> dict[str, GgufTensorInfo]: for path in gguf_split_paths(model_path): for t in _reader(path).tensors: index[t.name] = GgufTensorInfo( - t.name, path, int(t.data_offset), int(t.n_bytes), int(t.tensor_type), tuple(int(x) for x in t.shape) + t.name, path, int(t.data_offset), int(t.n_bytes), int(t.tensor_type), tuple(int(x) for x in t.shape), model_path ) return index diff --git a/python/freetoken/models/gguf/tensors.py b/python/freetoken/models/gguf/tensors.py index 4cba1d082e..83452d957a 100644 --- a/python/freetoken/models/gguf/tensors.py +++ b/python/freetoken/models/gguf/tensors.py @@ -7,9 +7,9 @@ import numpy as np import torch -from freetoken.layers.quantization.ggml_types import is_quantized, pad_rows, segment_param, segments +from freetoken.layers.quantization.ggml_types import NVFP4, is_quantized, pad_rows, scale_name, segment_param, segments -from .reader import GgufTensorInfo +from .reader import GgufTensorInfo, gguf_tensor_index _FLOAT_DTYPES = {0: torch.float32, 1: torch.float16, 30: torch.bfloat16} @@ -29,8 +29,49 @@ def matrix(info: GgufTensorInfo) -> torch.Tensor: return packed_rows(info) if is_quantized(info.ggml_type) else float_tensor(info) +def check_scales(model_path: str, *, nvfp4: bool = True, known: tuple[str, ...] = ()) -> None: + """Refuse ``.scale`` / ``.input_scale`` tensors that nothing reads. + + With ``nvfp4``, an NVFP4 tensor's ``.scale`` (its weight_scale_2) is read and its ``.input_scale`` goes unused, as in llama.cpp. + ``known`` lists tensor-name suffixes the family's loader reads itself.""" + types = {name: info.ggml_type for name, info in gguf_tensor_index(model_path).items()} + unread = [] + for name in types: + weight = name.removesuffix(".input_scale") if name.endswith(".input_scale") else name.removesuffix(".scale") + if weight == name or name.endswith(known): + continue + if not (nvfp4 and types.get(f"{weight}.weight", types.get(weight)) == NVFP4): + unread.append(name) + if unread: + what = "scale tensors outside NVFP4 weights" if nvfp4 else "tensors with per-tensor scales" + raise NotImplementedError(f"{model_path}: GGUF {what} are not supported (e.g. {sorted(unread)[0]})") + + +def weight_scale(info: GgufTensorInfo) -> torch.Tensor | None: + """The fp32 weight_scale_2 llama.cpp's converter writes next to an NVFP4 tensor as ``.scale`` (one value, or one per expert); None if there is none.""" + if info.ggml_type != NVFP4: + return None + scale = gguf_tensor_index(info.root or info.path).get(scale_name(info.name)) + return None if scale is None else float_tensor(scale).float().reshape(-1) + + +def _row_scales(prefix: str, infos: list[GgufTensorInfo], parts: list[torch.Tensor], pad: int) -> torch.Tensor | None: + scales = [weight_scale(i) for i in infos] + if all(s is None for s in scales): + return None + rows = [] + for info, part, s in zip(infos, parts, scales): + # llama.cpp reads a dense tensor's scale as one value + if s is not None and s.numel() != 1: + raise NotImplementedError(f"{prefix}: {info.name} has {s.numel()} scales, expected one") + rows.append(torch.ones(part.shape[0]) if s is None else s.expand(part.shape[0])) + return torch.cat([*rows, torch.ones(pad)]) + + def fused(prefix: str, infos: list[GgufTensorInfo], parts: list[torch.Tensor], k: int, pad_to: int = 0) -> Iterator[tuple[str, torch.Tensor]]: - """One fused module: float parts concatenated into ``weight``; with any quantized part, one matrix per ggml-type run (float runs as bf16 rows, quantized runs packed).""" + """One fused module: float parts concatenated into ``weight``; with any quantized part, one matrix per ggml-type run (float runs as bf16 rows, quantized runs packed). + + NVFP4 parts with a ``.scale`` add ``ggml_scale``: each part's scale on its output rows, fp32.""" types = [i.ggml_type for i in infos] pad = (-sum(p.shape[0] for p in parts)) % pad_to if pad_to else 0 if not any(is_quantized(t) for t in types): @@ -51,6 +92,9 @@ def fused(prefix: str, infos: list[GgufTensorInfo], parts: list[torch.Tensor], k dtype = torch.bfloat16 tail = [torch.zeros(zeros, rows[0].shape[1], dtype=dtype)] if zeros else [] yield f"{prefix}.{segment_param(i)}", torch.cat([*rows, *tail], dim=0) + scales = _row_scales(prefix, infos, parts, pad) + if scales is not None: + yield f"{prefix}.ggml_scale", scales -__all__ = ["float_tensor", "fused", "matrix", "packed_rows"] +__all__ = ["check_scales", "float_tensor", "fused", "matrix", "packed_rows", "weight_scale"] diff --git a/python/freetoken/models/llama/gguf.py b/python/freetoken/models/llama/gguf.py index 09d02627ae..833a5d66b5 100644 --- a/python/freetoken/models/llama/gguf.py +++ b/python/freetoken/models/llama/gguf.py @@ -13,7 +13,7 @@ import torch from freetoken.models.gguf.reader import GgufTensorInfo, gguf_tensor_index -from freetoken.models.gguf.tensors import float_tensor, fused, matrix +from freetoken.models.gguf.tensors import check_scales, float_tensor, fused, matrix from .config import parse_config @@ -53,10 +53,8 @@ def to_gguf(name: str, tensors: dict[str, int]) -> tuple[str, ...]: def _check_supported(shim: GgufConfigShim) -> None: if any(k.startswith("prism.hadamard.") for k in shim.metadata): raise NotImplementedError(f"{shim.model_path}: Hadamard-rotated (prism) GGUF weights are not supported") + check_scales(shim.model_path) names = sorted(gguf_tensor_index(shim.model_path)) - scaled = [n for n in names if n.endswith((".scale", ".input_scale"))] - if scaled: - raise NotImplementedError(f"{shim.model_path}: GGUF tensors with per-tensor scales are not supported (e.g. {scaled[0]})") # the llama arch also carries Mixtral experts and optional biases, which the llama model does not have extra = [n for n in names if n.endswith(".bias") or ".ffn_gate_inp." in n] if extra: diff --git a/python/freetoken/models/qwen3/gguf.py b/python/freetoken/models/qwen3/gguf.py index 06245ce4c2..8e0388fe3c 100644 --- a/python/freetoken/models/qwen3/gguf.py +++ b/python/freetoken/models/qwen3/gguf.py @@ -12,7 +12,7 @@ import torch from freetoken.models.gguf.reader import GgufTensorInfo, gguf_tensor_index -from freetoken.models.gguf.tensors import float_tensor, fused, matrix +from freetoken.models.gguf.tensors import check_scales, float_tensor, fused, matrix from .config import parse_config @@ -52,9 +52,7 @@ def to_gguf(name: str, tensors: dict[str, int]) -> tuple[str, ...]: def _check_supported(shim: GgufConfigShim) -> None: if any(k.startswith("prism.hadamard.") for k in shim.metadata): raise NotImplementedError(f"{shim.model_path}: Hadamard-rotated (prism) GGUF weights are not supported") - scaled = sorted(n for n in gguf_tensor_index(shim.model_path) if n.endswith((".scale", ".input_scale"))) - if scaled: - raise NotImplementedError(f"{shim.model_path}: GGUF tensors with per-tensor scales are not supported (e.g. {scaled[0]})") + check_scales(shim.model_path) def _hf_config(shim: GgufConfigShim) -> SimpleNamespace: diff --git a/python/freetoken/models/qwen3_5_moe/gguf.py b/python/freetoken/models/qwen3_5_moe/gguf.py index 85af8ab85f..7f72061e00 100644 --- a/python/freetoken/models/qwen3_5_moe/gguf.py +++ b/python/freetoken/models/qwen3_5_moe/gguf.py @@ -14,7 +14,7 @@ from freetoken.layers.quantization.ggml_types import is_quantized from freetoken.models.gguf.reader import GgufTensorInfo, gguf_tensor_index -from freetoken.models.gguf.tensors import float_tensor, fused, matrix +from freetoken.models.gguf.tensors import check_scales, float_tensor, fused, matrix from .config import parse_config @@ -67,9 +67,7 @@ def to_gguf(name: str, tensors: dict[str, int]) -> tuple[str, ...]: def _check_supported(shim: GgufConfigShim) -> None: if any(k.startswith("prism.hadamard.") for k in shim.metadata): raise NotImplementedError(f"{shim.model_path}: Hadamard-rotated (prism) GGUF weights are not supported") - scaled = sorted(n for n in gguf_tensor_index(shim.model_path) if n.endswith((".scale", ".input_scale"))) - if scaled: - raise NotImplementedError(f"{shim.model_path}: GGUF tensors with per-tensor scales are not supported (e.g. {scaled[0]})") + check_scales(shim.model_path) def parse_gguf_config(shim: GgufConfigShim) -> ModelConfig: diff --git a/python/freetoken/models/qwen3_moe/gguf.py b/python/freetoken/models/qwen3_moe/gguf.py index 7bf016292e..3a8ef98a7a 100644 --- a/python/freetoken/models/qwen3_moe/gguf.py +++ b/python/freetoken/models/qwen3_moe/gguf.py @@ -13,7 +13,7 @@ import torch from freetoken.models.gguf.reader import GgufTensorInfo, gguf_tensor_index -from freetoken.models.gguf.tensors import float_tensor, fused, matrix +from freetoken.models.gguf.tensors import check_scales, float_tensor, fused, matrix from .config import parse_config @@ -56,9 +56,7 @@ def to_gguf(name: str, tensors: dict[str, int]) -> tuple[str, ...]: def _check_supported(shim: GgufConfigShim) -> None: if any(k.startswith("prism.hadamard.") for k in shim.metadata): raise NotImplementedError(f"{shim.model_path}: Hadamard-rotated (prism) GGUF weights are not supported") - scaled = sorted(n for n in gguf_tensor_index(shim.model_path) if n.endswith((".scale", ".input_scale"))) - if scaled: - raise NotImplementedError(f"{shim.model_path}: GGUF tensors with per-tensor scales are not supported (e.g. {scaled[0]})") + check_scales(shim.model_path) def _hf_config(shim: GgufConfigShim) -> SimpleNamespace: diff --git a/python/freetoken/models/qwen4_exp/gguf.py b/python/freetoken/models/qwen4_exp/gguf.py index 7909324b2c..0c1bf98059 100644 --- a/python/freetoken/models/qwen4_exp/gguf.py +++ b/python/freetoken/models/qwen4_exp/gguf.py @@ -14,7 +14,7 @@ from freetoken.layers.quantization.ggml_types import is_quantized from freetoken.models.gguf.reader import GgufTensorInfo, gguf_tensor_index -from freetoken.models.gguf.tensors import float_tensor, fused, matrix +from freetoken.models.gguf.tensors import check_scales, float_tensor, fused, matrix from .config import parse_config @@ -83,6 +83,8 @@ def to_gguf(name: str, tensors: dict[str, int]) -> tuple[str, ...]: def parse_gguf_config(shim: GgufConfigShim) -> ModelConfig: + # NVFP4 weight scales are unvalidated on this arch, so refuse them instead of risking wrong output + check_scales(shim.model_path, nvfp4=False) md = shim.metadata def kv(key: str): diff --git a/python/freetoken/moe/cpu_executor.py b/python/freetoken/moe/cpu_executor.py index 4c0aee7ae5..bb65146b45 100644 --- a/python/freetoken/moe/cpu_executor.py +++ b/python/freetoken/moe/cpu_executor.py @@ -149,11 +149,11 @@ def resolve_threads_and_affinity(requested: int) -> tuple[int, list[int]]: def _format_label(name: str, opts: tuple) -> str: """One format instance as the startup log names it, e.g. ``gguf(gate_up=IQ2_XXS, down=IQ4_NL)``.""" - if name == "gguf": + if name in ("gguf", "gguf_scaled"): from freetoken.layers.quantization.ggml_types import type_name gate_up, down = dict(opts)["gate_up_type"], dict(opts)["down_type"] - return f"gguf(gate_up={type_name(gate_up)}, down={type_name(down)})" + return f"{name}(gate_up={type_name(gate_up)}, down={type_name(down)})" return name + (f"({', '.join(f'{k}={v}' for k, v in opts)})" if opts else "") @@ -192,7 +192,7 @@ def __init__( f"{sorted(_cpu_moe.weight_formats())} formats, but this checkpoint's experts are " f"{f!r}; use --moe-strategy offload (GPU-side dequant) instead." ) - if f == "gguf": + if f in ("gguf", "gguf_scaled"): from freetoken.kernel import ggml_cpu if (variant := ggml_cpu.select()) is None: @@ -423,8 +423,8 @@ def _resolve_banks(self, banks: dict, fmt: str, rep: int = 0, opts: dict | None (see ``_make_table``), not a single bank's -- the C++ side resolves ``tbl[layer_id]`` per task. """ - if fmt == "gguf": - return self._resolve_gguf_banks(banks, rep, int(opts["gate_up_type"]), int(opts["down_type"])) + if fmt in ("gguf", "gguf_scaled"): + return self._resolve_gguf_banks(banks, rep, int(opts["gate_up_type"]), int(opts["down_type"]), scaled=fmt == "gguf_scaled") if fmt == "bf16": gate_up = banks["gate_up"] down = banks["down"] @@ -465,8 +465,10 @@ def _resolve_banks(self, banks: dict, fmt: str, rep: int = 0, opts: dict | None banks = self._tables(gate_up=gup, down=dnp, gate_up_scale=gus, gate_up_global=gug, down_scale=dns, down_global=dng) return banks, {}, (H, I) - def _resolve_gguf_banks(self, banks: dict, rep: int, gate_up_type: int, down_type: int) -> tuple[dict, dict, tuple[int, int]]: - """GGUF rows read in place: gate_up [E, 2I, row_bytes(H)] (gate rows, then up rows) and down [E, H, row_bytes(I)], one ggml type each; float rows keep their dtype.""" + def _resolve_gguf_banks(self, banks: dict, rep: int, gate_up_type: int, down_type: int, scaled: bool = False) -> tuple[dict, dict, tuple[int, int]]: + """GGUF rows read in place: gate_up [E, 2I, row_bytes(H)] (gate rows, then up rows) and down [E, H, row_bytes(I)], one ggml type each; float rows keep their dtype. + + ``scaled`` adds global [E, 4] fp32, each expert's (gate, up, down, 1) NVFP4 weight_scale_2.""" from freetoken.layers.quantization.ggml_types import float_dtype, row_bytes, type_name gate_up, down = banks["gate_up"], banks["down"] @@ -475,8 +477,13 @@ def _resolve_gguf_banks(self, banks: dict, rep: int, gate_up_type: int, down_typ dtype = float_dtype(ggml_type) if float_dtype(ggml_type) is not None else torch.uint8 if t.dtype != dtype or t.shape[-1] * t.element_size() != row_bytes(ggml_type, k): raise ValueError(f"MoE layer {rep} {role} bank {tuple(t.shape)} {t.dtype} is not {type_name(ggml_type)} rows of {k} values") - banks = self._tables(gate_up=gate_up, down=down) - return banks, {"gate_up_type": gate_up_type, "down_type": down_type}, (H, I) + types = {"gate_up_type": gate_up_type, "down_type": down_type} + if not scaled: + return self._tables(gate_up=gate_up, down=down), types, (H, I) + scales = banks["global"] + if scales[rep].dtype != torch.float32 or tuple(scales[rep].shape) != (gate_up[rep].shape[0], 4) or not scales[rep].is_contiguous(): + raise ValueError(f"MoE layer {rep} global bank {tuple(scales[rep].shape)} {scales[rep].dtype} is not contiguous [experts, 4] fp32") + return self._tables(gate_up=gate_up, down=down, **{"global": scales}), types, (H, I) def _resolve_mxfp4_banks(self, banks: dict, rep: int = 0) -> tuple[dict, dict, tuple[int, int]]: """gpt-oss mxfp4 ``mxfp4_triton`` schema: transposed split-K blocks/scales diff --git a/python/freetoken/moe/offload_cache.py b/python/freetoken/moe/offload_cache.py index 23ba325106..5291e621f8 100644 --- a/python/freetoken/moe/offload_cache.py +++ b/python/freetoken/moe/offload_cache.py @@ -48,6 +48,8 @@ "fp8_block": ("gate_up", "gate_up_scale", "down", "down_scale"), # GGUF experts in their own ggml type per layer: gate_up [L*E, 2I, row_bytes] (gate rows, then up rows), down [L*E, H, row_bytes]; float experts keep their dtype "gguf": ("gate_up", "down"), + # GGUF experts whose NVFP4 tensors carry per-expert scales: global [L*E, 4] fp32 (gate, up, down, 1) weight_scale_2 + "gguf_scaled": ("gate_up", "down", "global"), # native ModelOpt rows for the Triton inline-dequant kernels: packed e2m1 codes + # fp8-e4m3 per-16 block scales + per-output-row fp16 globals (w1/w3 carry distinct # globals, and folding them into the e4m3 block scales would underflow)