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..8616a17560 100644 --- a/tests/kernels/test_ggml.py +++ b/tests/kernels/test_ggml.py @@ -6,6 +6,7 @@ import torch.nn.functional as F from freetoken.kernel import ggml +from freetoken.layers.quantization.ggml_types import pad_rows from freetoken.third_party.gguf import GGMLQuantizationType as QType from freetoken.third_party.gguf import quants @@ -23,6 +24,20 @@ def _q4_weights(shape, seed=690): return storage[..., :-1, :], reference +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, 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)) + 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 +97,15 @@ 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) +