From aff897ca97e5790b8bf9322f2669fda1e21d3ee5 Mon Sep 17 00:00:00 2001 From: iamanishx Date: Thu, 22 Oct 2026 21:37:57 +0530 Subject: [PATCH 1/2] feat(gguf): read NVFP4 token_embd rows in the vendored get_rows Signed-off-by: iamanishx --- python/freetoken/kernel/csrc/ggml/VERSION | 1 + .../kernel/csrc/ggml/freetoken.patch | 31 +++++++++++ .../kernel/csrc/ggml/ggml-cuda/dequantize.cuh | 12 +++++ .../kernel/csrc/ggml/ggml-cuda/getrows.cu | 4 ++ .../layers/quantization/embedding/gguf.py | 6 +-- tests/kernels/test_ggml.py | 52 +++++++++++++++++++ 6 files changed, 103 insertions(+), 3 deletions(-) diff --git a/python/freetoken/kernel/csrc/ggml/VERSION b/python/freetoken/kernel/csrc/ggml/VERSION index f7b5c01705..c4a9671dd5 100644 --- a/python/freetoken/kernel/csrc/ggml/VERSION +++ b/python/freetoken/kernel/csrc/ggml/VERSION @@ -11,5 +11,6 @@ Local changes, kept in freetoken.patch (already applied): - 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. +- getrows.cu, dequantize.cuh: get_rows reads NVFP4 rows, e.g. an NVFP4 token_embd. 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 595498dd15..e968a07829 100644 --- a/python/freetoken/kernel/csrc/ggml/freetoken.patch +++ b/python/freetoken/kernel/csrc/ggml/freetoken.patch @@ -34,6 +34,37 @@ ggml_glu_op glu_op; float glu_limit = 0.0f; const void * shared_up = nullptr; +--- a/ggml/src/ggml-cuda/dequantize.cuh ++++ b/ggml/src/ggml-cuda/dequantize.cuh +@@ -450,3 +450,15 @@ + y[j+16] = ggml_cuda_cast(d * kvalues_mxfp4[q4[j] >> 4]*0.5f); + } + } ++ ++// FreeToken: for k_get_rows with qk = QK_NVFP4_SUB, so ib indexes a 16-element sub-block and the nibble pair lands at iqs and iqs + 8 ++static __device__ __forceinline__ void dequantize_nvfp4(const void * vx, const int64_t ib, const int iqs, float2 & v){ ++ const block_nvfp4 * x = (const block_nvfp4 *) vx + ib/(QK_NVFP4/QK_NVFP4_SUB); ++ const int sub = ib % (QK_NVFP4/QK_NVFP4_SUB); ++ ++ const float d = ggml_cuda_ue4m3_to_fp32(x->d[sub]); ++ const int vui = x->qs[sub*(QK_NVFP4_SUB/2) + iqs]; ++ ++ v.x = d * kvalues_mxfp4[vui & 0xF]; ++ v.y = d * kvalues_mxfp4[vui >> 4]; ++} +--- a/ggml/src/ggml-cuda/getrows.cu ++++ b/ggml/src/ggml-cuda/getrows.cu +@@ -404,6 +404,10 @@ + get_rows_cuda_kq<32, dst_t, dequantize_mxfp4>(src0_d, src1_d, dst_d, + ne00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb1, nb2, nb3, stream); + break; ++ case GGML_TYPE_NVFP4: ++ get_rows_cuda_q(src0_d, src1_d, dst_d, ++ ne00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb1, nb2, nb3, stream); ++ break; + default: + GGML_ABORT("%s: unsupported src0 type: %s\n", __func__, ggml_type_name(src0_type)); + } --- a/ggml/src/ggml-cuda/mmid.cu +++ b/ggml/src/ggml-cuda/mmid.cu @@ -177,3 +177,45 @@ diff --git a/python/freetoken/kernel/csrc/ggml/ggml-cuda/dequantize.cuh b/python/freetoken/kernel/csrc/ggml/ggml-cuda/dequantize.cuh index 651524be7b..de6dd857d1 100644 --- a/python/freetoken/kernel/csrc/ggml/ggml-cuda/dequantize.cuh +++ b/python/freetoken/kernel/csrc/ggml/ggml-cuda/dequantize.cuh @@ -450,3 +450,15 @@ static __device__ __forceinline__ void dequantize_mxfp4(const void * vx, const i y[j+16] = ggml_cuda_cast(d * kvalues_mxfp4[q4[j] >> 4]*0.5f); } } + +// FreeToken: for k_get_rows with qk = QK_NVFP4_SUB, so ib indexes a 16-element sub-block and the nibble pair lands at iqs and iqs + 8 +static __device__ __forceinline__ void dequantize_nvfp4(const void * vx, const int64_t ib, const int iqs, float2 & v){ + const block_nvfp4 * x = (const block_nvfp4 *) vx + ib/(QK_NVFP4/QK_NVFP4_SUB); + const int sub = ib % (QK_NVFP4/QK_NVFP4_SUB); + + const float d = ggml_cuda_ue4m3_to_fp32(x->d[sub]); + const int vui = x->qs[sub*(QK_NVFP4_SUB/2) + iqs]; + + v.x = d * kvalues_mxfp4[vui & 0xF]; + v.y = d * kvalues_mxfp4[vui >> 4]; +} diff --git a/python/freetoken/kernel/csrc/ggml/ggml-cuda/getrows.cu b/python/freetoken/kernel/csrc/ggml/ggml-cuda/getrows.cu index d97e8fdb0f..98a95c0c6b 100644 --- a/python/freetoken/kernel/csrc/ggml/ggml-cuda/getrows.cu +++ b/python/freetoken/kernel/csrc/ggml/ggml-cuda/getrows.cu @@ -404,6 +404,10 @@ static void ggml_cuda_get_rows_switch_src0_type( get_rows_cuda_kq<32, dst_t, dequantize_mxfp4>(src0_d, src1_d, dst_d, ne00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb1, nb2, nb3, stream); break; + case GGML_TYPE_NVFP4: + get_rows_cuda_q(src0_d, src1_d, dst_d, + ne00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb1, nb2, nb3, stream); + break; default: GGML_ABORT("%s: unsupported src0 type: %s\n", __func__, ggml_type_name(src0_type)); } diff --git a/python/freetoken/layers/quantization/embedding/gguf.py b/python/freetoken/layers/quantization/embedding/gguf.py index 98f12dd9cc..536c6f1643 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 NVFP4, Segment, pad_rows, row_bytes +from ..ggml_types import Segment, pad_rows, row_bytes from ..linear.gguf import GgmlLinearKernel from ..registry import LayerKind, register_method from ..scheme import QuantKind, gguf_types @@ -20,8 +20,8 @@ class GgmlEmbeddingKernel(EmbeddingKernel): def unusable_reason(self, cfg: EmbeddingConfig) -> str | 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 + # llama.cpp reads no token_embd scale and refuses a tied NVFP4 head that has one + return "GGUF tables with a weight scale are not supported" if cfg.scheme.weight.scale is not None else None def embed(self, layer: Any, ids: torch.Tensor) -> torch.Tensor: from freetoken.kernel import ggml diff --git a/tests/kernels/test_ggml.py b/tests/kernels/test_ggml.py index ea6a5d6830..c9c8e71089 100644 --- a/tests/kernels/test_ggml.py +++ b/tests/kernels/test_ggml.py @@ -1,11 +1,18 @@ """Packed GGUF ops against gguf-py's CPU dequantization and PyTorch math.""" +from types import SimpleNamespace + import numpy as np import pytest import torch import torch.nn.functional as F from freetoken.kernel import ggml +from freetoken.layers.quantization import KernelSelectionError, select_kernel +from freetoken.layers.quantization.embedding.base import EmbeddingConfig +from freetoken.layers.quantization.embedding.gguf import GGUFEmbeddingMethod, GgmlEmbeddingKernel +from freetoken.layers.quantization.ggml_types import pad_rows +from freetoken.layers.quantization.scheme import gguf_scheme from freetoken.third_party.gguf import GGMLQuantizationType as QType from freetoken.third_party.gguf import quants @@ -23,6 +30,20 @@ def _q4_weights(shape, seed=690): return storage[..., :-1, :], reference +def _nvfp4_weights(shape, max_scale_code=0x7E, seed=765): + """Random NVFP4 blocks over every fp4 code; UE4M3 scale codes up to ``max_scale_code`` (0x7F is NaN), padded for MMQ.""" + rng = np.random.default_rng(seed) + rows, k = shape + blocks = k // 64 + d = rng.integers(0, max_scale_code + 1, (rows, blocks, 4), dtype=np.uint8) + qs = rng.integers(0, 256, (rows, blocks, 32), dtype=np.uint8) + packed = np.concatenate([d, qs], axis=-1).reshape(rows, blocks * 36) + reference = torch.from_numpy(quants.dequantize(packed, QType.NVFP4)) + storage = torch.zeros(rows + pad_rows(int(QType.NVFP4), k), packed.shape[-1], dtype=torch.uint8, device="cuda") + storage[:rows].copy_(torch.from_numpy(packed)) + return storage, reference + + def _activation(shape, dtype, seed=19): rng = np.random.default_rng(seed) half = rng.integers(-127, 128, (*shape[:-1], shape[-1] // 32, 16)) @@ -82,3 +103,34 @@ def test_q4_nondefault_stream(): copied = actual.clone() stream.synchronize() torch.testing.assert_close(copied.cpu(), F.linear(x.cpu(), ref), atol=2e-4, rtol=1e-3) + + +@pytest.mark.parametrize("dtype", DTYPES) +def test_nvfp4_dequantize_and_rows(dtype): + # 320 values: five NVFP4 blocks, not a multiple of the 256-value super-block the k-quant get_rows path needs + w, ref = _nvfp4_weights((64, 320)) + actual = ggml.dequantize(w[:64], int(QType.NVFP4), 320, dtype) + torch.testing.assert_close(actual.cpu(), ref.to(dtype), atol=0, rtol=0) + ids = torch.tensor([63, 0, 17, 17], dtype=torch.int32, device="cuda") + rows = ggml.get_rows(w, int(QType.NVFP4), 320, ids, dtype) + torch.testing.assert_close(rows.cpu(), ref[ids.cpu().long()].to(dtype), atol=0, rtol=0) + + +def test_nvfp4_table_embed_and_tied_logits(): + w, ref = _nvfp4_weights((64, 320), max_scale_code=0x3F) + cfg = EmbeddingConfig(64, 320, 1, torch.float32, gguf_scheme([QType.NVFP4])) + layer = SimpleNamespace(quant_method=SimpleNamespace(cfg=cfg), weight=w) + kernel = GgmlEmbeddingKernel() + ids = torch.tensor([[63, 0], [17, 17]], dtype=torch.int32, device="cuda") + torch.testing.assert_close(kernel.embed(layer, ids).cpu(), ref[ids.cpu().long()], atol=0, rtol=0) + x = _activation((3, 320), torch.float32) + torch.testing.assert_close(kernel.logits(layer, x, None).cpu(), F.linear(x.cpu(), ref), atol=1e-3, rtol=1e-3) + + +def test_nvfp4_table_kernel_selection(): + cfg = EmbeddingConfig(64, 320, 1, torch.bfloat16, gguf_scheme([QType.NVFP4])) + assert select_kernel(GGUFEmbeddingMethod.candidates, "auto", cfg).name == "ggml" + # llama.cpp reads no token_embd scale and refuses a tied NVFP4 head that has one + scaled = EmbeddingConfig(64, 320, 1, torch.bfloat16, gguf_scheme([QType.NVFP4], scaled=True)) + with pytest.raises(KernelSelectionError): + select_kernel(GGUFEmbeddingMethod.candidates, "auto", scaled) From 2c9eb3f3b342d58ff1b94d7642fbba2524c0f1a6 Mon Sep 17 00:00:00 2001 From: iamanishx Date: Fri, 23 Oct 2026 07:25:06 +0530 Subject: [PATCH 2/2] fix(tests) Signed-off-by: iamanishx --- tests/kernels/test_ggml.py | 31 +++---------------------------- 1 file changed, 3 insertions(+), 28 deletions(-) diff --git a/tests/kernels/test_ggml.py b/tests/kernels/test_ggml.py index c9c8e71089..8616a17560 100644 --- a/tests/kernels/test_ggml.py +++ b/tests/kernels/test_ggml.py @@ -1,18 +1,12 @@ """Packed GGUF ops against gguf-py's CPU dequantization and PyTorch math.""" -from types import SimpleNamespace - import numpy as np import pytest import torch import torch.nn.functional as F from freetoken.kernel import ggml -from freetoken.layers.quantization import KernelSelectionError, select_kernel -from freetoken.layers.quantization.embedding.base import EmbeddingConfig -from freetoken.layers.quantization.embedding.gguf import GGUFEmbeddingMethod, GgmlEmbeddingKernel from freetoken.layers.quantization.ggml_types import pad_rows -from freetoken.layers.quantization.scheme import gguf_scheme from freetoken.third_party.gguf import GGMLQuantizationType as QType from freetoken.third_party.gguf import quants @@ -30,12 +24,12 @@ def _q4_weights(shape, seed=690): return storage[..., :-1, :], reference -def _nvfp4_weights(shape, max_scale_code=0x7E, seed=765): - """Random NVFP4 blocks over every fp4 code; UE4M3 scale codes up to ``max_scale_code`` (0x7F is NaN), padded for MMQ.""" +def _nvfp4_weights(shape, seed=765): + """Random NVFP4 blocks over every fp4 code and every UE4M3 scale code below NaN (0x7F), padded for MMQ.""" rng = np.random.default_rng(seed) rows, k = shape blocks = k // 64 - d = rng.integers(0, max_scale_code + 1, (rows, blocks, 4), dtype=np.uint8) + d = rng.integers(0, 0x7F, (rows, blocks, 4), dtype=np.uint8) qs = rng.integers(0, 256, (rows, blocks, 32), dtype=np.uint8) packed = np.concatenate([d, qs], axis=-1).reshape(rows, blocks * 36) reference = torch.from_numpy(quants.dequantize(packed, QType.NVFP4)) @@ -115,22 +109,3 @@ def test_nvfp4_dequantize_and_rows(dtype): rows = ggml.get_rows(w, int(QType.NVFP4), 320, ids, dtype) torch.testing.assert_close(rows.cpu(), ref[ids.cpu().long()].to(dtype), atol=0, rtol=0) - -def test_nvfp4_table_embed_and_tied_logits(): - w, ref = _nvfp4_weights((64, 320), max_scale_code=0x3F) - cfg = EmbeddingConfig(64, 320, 1, torch.float32, gguf_scheme([QType.NVFP4])) - layer = SimpleNamespace(quant_method=SimpleNamespace(cfg=cfg), weight=w) - kernel = GgmlEmbeddingKernel() - ids = torch.tensor([[63, 0], [17, 17]], dtype=torch.int32, device="cuda") - torch.testing.assert_close(kernel.embed(layer, ids).cpu(), ref[ids.cpu().long()], atol=0, rtol=0) - x = _activation((3, 320), torch.float32) - torch.testing.assert_close(kernel.logits(layer, x, None).cpu(), F.linear(x.cpu(), ref), atol=1e-3, rtol=1e-3) - - -def test_nvfp4_table_kernel_selection(): - cfg = EmbeddingConfig(64, 320, 1, torch.bfloat16, gguf_scheme([QType.NVFP4])) - assert select_kernel(GGUFEmbeddingMethod.candidates, "auto", cfg).name == "ggml" - # llama.cpp reads no token_embd scale and refuses a tied NVFP4 head that has one - scaled = EmbeddingConfig(64, 320, 1, torch.bfloat16, gguf_scheme([QType.NVFP4], scaled=True)) - with pytest.raises(KernelSelectionError): - select_kernel(GGUFEmbeddingMethod.candidates, "auto", scaled)