Repository: ggml-org/llama.cpp
Pull request: gguf-py : add Numpy MXFP4 de/quantization support

MXFP4 support was added in #15091, but Python-based (de)quantization is missing, which this should fix.

Other change: I've removed the FMA workaround for `Q4_0` and `Q5_0` quantization functions, since the C versions don't use FMA since a while ago (not sure exactly when). 

All correctness tests pass.

```console
$ python3 gguf-py/tests/test_quants.py --type MXFP4
INFO:test-quants:Testing MXFP4
DEBUG:test-quants:Quantizing to MXFP4 with Python
DEBUG:test-quants:Quantizing to MXFP4 with C
INFO:test-quants:Quantization to MXFP4 matches exactly ✅
DEBUG:test-quants:Dequantizing from MXFP4 with Python
DEBUG:test-quants:Dequantizing from MXFP4 with C
INFO:test-quants:Dequantization from MXFP4 matches exactly ✅
DEBUG:test-quants:Dequantizing random f16 data as MXFP4 with Python
gguf-py/gguf/quants.py:704: RuntimeWarning: overflow encountered in multiply
  return (d * qs.astype(np.float32))
DEBUG:test-quants:Dequantizing random f16 data as MXFP4 with C
INFO:test-quants:Dequantization from random f16 data as MXFP4 matches exactly ✅
```

Note that there's a RuntimeWarning from Numpy above, but the equality check still passes. The warning is likely caused by invalid E8M0 values from the random data which is dequantized in that test.

<details><summary>Full test results</summary>

```console
$ python3 gguf-py/tests/test_quants.py
INFO:test-quants:Testing F16
DEBUG:test-quants:Quantizing to F16 with Python
DEBUG:test-quants:Quantizing to F16 with C
INFO:test-quants:Quantization to F16 matches exactly ✅
DEBUG:test-quants:Dequantizing from F16 with Python
DEBUG:test-quants:Dequantizing from F16 with C
INFO:test-quants:Dequantization from F16 matches exactly ✅
DEBUG:test-quants:Dequantizing random f16 data as F16 with Python
DEBUG:test-quants:Dequantizing random f16 data as F16 with C
INFO:test-quants:Dequantization from random f16 data as F16 matches exactly ✅
INFO:test-quants:Testing BF16
DEBUG:test-quants:Quantizing to BF16 with Python
DEBUG:test-quants:Quantizing to BF16 with C
INFO:test-quants:Quantization to BF16 matches exactly ✅
DEBUG:test-quants:Dequantizing from BF16 with Python
DEBUG:test-quants:Dequantizing from BF16 with C
INFO:test-quants:Dequantization from BF16 matches exactly ✅
DEBUG:test-quants:Dequantizing random f16 data as BF16 with Python
DEBUG:test-quants:Dequantizing random f16 data as BF16 with C
INFO:test-quants:Dequantization from random f16 data as BF16 matches exactly ✅
INFO:test-quants:Testing Q4_0
DEBUG:test-quants:Quantizing to Q4_0 with Python
DEBUG:test-quants:Quantizing to Q4_0 with C
INFO:test-quants:Quantization to Q4_0 matches exactly ✅
DEBUG:test-quants:Dequantizing from Q4_0 with Python
DEBUG:test-quants:Dequantizing from Q4_0 with C
INFO:test-quants:Dequantization from Q4_0 matches exactly ✅
DEBUG:test-quants:Dequantizing random f16 data as Q4_0 with Python
DEBUG:test-quants:Dequantizing random f16 data as Q4_0 with C
INFO:test-quants:Dequantization from random f16 data as Q4_0 matches exactly ✅
INFO:test-quants:Testing Q4_1
DEBUG:test-quants:Quantizing to Q4_1 with Python
DEBUG:test-quants:Quantizing to Q4_1 with C
INFO:test-quants:Quantization to Q4_1 matches exactly ✅
DEBUG:test-quants:Dequantizing from Q4_1 with Python
DEBUG:test-quants:Dequantizing from Q4_1 with C
INFO:test-quants:Dequantization from Q4_1 matches exactly ✅
DEBUG:test-quants:Dequantizing random f16 data as Q4_1 with Python
DEBUG:test-quants:Dequantizing random f16 data as Q4_1 with C
INFO:test-quants:Dequantization from random f16 data as Q4_1 matches exactly ✅
INFO:test-quants:Testing Q5_0
DEBUG:test-quants:Quantizing to Q5_0 with Python
DEBUG:test-quants:Quantizing to Q5_0 with C
INFO:test-quants:Quantization to Q5_0 matches exactly ✅
DEBUG:test-quants:Dequantizing from Q5_0 with Python
DEBUG:test-quants:Dequantizing from Q5_0 with C
INFO:test-quants:Dequantization from Q5_0 matches exactly ✅
DEBUG:test-quants:Dequantizing random f16 data as Q5_0 with Python
DEBUG:test-quants:Dequantizing random f16 data as Q5_0 with C
INFO:test-quants:Dequantization from random f16 data as Q5_0 matches exactly ✅
INFO:test-quants:Testing Q5_1
DEBUG:test-quants:Quantizing to Q5_1 with Python
DEBUG:test-quants:Quantizing to Q5_1 with C
INFO:test-quants:Quantization to Q5_1 matches exactly ✅
DEBUG:test-quants:Dequantizing from Q5_1 with Python
DEBUG:test-quants:Dequantizing from Q5_1 with C
INFO:test-quants:Dequantization from Q5_1 matches exactly ✅
DEBUG:test-quants:Dequantizing random f16 data as Q5_1 with Python
DEBUG:test-quants:Dequantizing random f16 data as Q5_1 with C
INFO:test-quants:Dequantization from random f16 data as Q5_1 matches exactly ✅
INFO:test-quants:Testing Q8_0
DEBUG:test-quants:Quantizing to Q8_0 with Python
DEBUG:test-quants:Quantizing to Q8_0 with C
INFO:test-quants:Quantization to Q8_0 matches exactly ✅
DEBUG:test-quants:Dequantizing from Q8_0 with Python
DEBUG:test-quants:Dequantizing from Q8_0 with C
INFO:test-quants:Dequantization from Q8_0 matches exactly ✅
DEBUG:test-quants:Dequantizing random f16 data as Q8_0 with Python
DEBUG:test-quants:Dequantizing random f16 data as Q8_0 with C
INFO:test-quants:Dequantization from random f16 data as Q8_0 matches exactly ✅
INFO:test-quants:Testing Q2_K
DEBUG:test-quants:Quantizing to Q2_K with C
DEBUG:test-quants:Dequantizing from Q2_K with Python
DEBUG:test-quants:Dequantizing from Q2_K with C
INFO:test-quants:Dequantization from Q2_K matches exactly ✅
DEBUG:test-quants:Dequantizing random f16 data as Q2_K with Python
DEBUG:test-quants:Dequantizing random f16 data as Q2_K with C
INFO:test-quants:Dequantization from random f16 data as Q2_K matches exactly ✅
INFO:test-quants:Testing Q3_K
DEBUG:test-quants:Quantizing to Q3_K with C
DEBUG:test-quants:Dequantizing from Q3_K with Python
DEBUG:test-quants:Dequantizing from Q3_K with C
INFO:test-quants:Dequantization from Q3_K matches exactly ✅
DEBUG:test-quants:Dequantizing random f16 data as Q3_K with Python
DEBUG:test-quants:Dequantizing random f16 data as Q3_K with C
INFO:test-quants:Dequantization from random f16 data as Q3_K matches exactly ✅
INFO:test-quants:Testing Q4_K
DEBUG:test-quants:Quantizing to Q4_K with C
DEBUG:test-quants:Dequantizing from Q4_K with Python
DEBUG:test-quants:Dequantizing from Q4_K with C
INFO:test-quants:Dequantization from Q4_K matches exactly ✅
DEBUG:test-quants:Dequantizing random f16 data as Q4_K with Python
DEBUG:test-quants:Dequantizing random f16 data as Q4_K with C
INFO:test-quants:Dequantization from random f16 data as Q4_K matches exactly ✅
INFO:test-quants:Testing Q5_K
DEBUG:test-quants:Quantizing to Q5_K with C
DEBUG:test-quants:Dequantizing from Q5_K with Python
DEBUG:test-quants:Dequantizing from Q5_K with C
INFO:test-quants:Dequantization from Q5_K matches exactly ✅
DEBUG:test-quants:Dequantizing random f16 data as Q5_K with Python
DEBUG:test-quants:Dequantizing random f16 data as Q5_K with C
INFO:test-quants:Dequantization from random f16 data as Q5_K matches exactly ✅
INFO:test-quants:Testing Q6_K
DEBUG:test-quants:Quantizing to Q6_K with C
DEBUG:test-quants:Dequantizing from Q6_K with Python
DEBUG:test-quants:Dequantizing from Q6_K with C
INFO:test-quants:Dequantization from Q6_K matches exactly ✅
DEBUG:test-quants:Dequantizing random f16 data as Q6_K with Python
DEBUG:test-quants:Dequantizing random f16 data as Q6_K with C
INFO:test-quants:Dequantization from random f16 data as Q6_K matches exactly ✅
INFO:test-quants:Testing TQ1_0
DEBUG:test-quants:Quantizing to TQ1_0 with Python
DEBUG:test-quants:Quantizing to TQ1_0 with C
INFO:test-quants:Quantization to TQ1_0 matches exactly ✅
DEBUG:test-quants:Dequantizing from TQ1_0 with Python
DEBUG:test-quants:Dequantizing from TQ1_0 with C
INFO:test-quants:Dequantization from TQ1_0 matches exactly ✅
DEBUG:test-quants:Dequantizing random f16 data as TQ1_0 with Python
DEBUG:test-quants:Dequantizing random f16 data as TQ1_0 with C
INFO:test-quants:Dequantization from random f16 data as TQ1_0 matches exactly ✅
INFO:test-quants:Testing TQ2_0
DEBUG:test-quants:Quantizing to TQ2_0 with Python
DEBUG:test-quants:Quantizing to TQ2_0 with C
INFO:test-quants:Quantization to TQ2_0 matches exactly ✅
DEBUG:test-quants:Dequantizing from TQ2_0 with Python
DEBUG:test-quants:Dequantizing from TQ2_0 with C
INFO:test-quants:Dequantization from TQ2_0 matches exactly ✅
DEBUG:test-quants:Dequantizing random f16 data as TQ2_0 with Python
DEBUG:test-quants:Dequantizing random f16 data as TQ2_0 with C
INFO:test-quants:Dequantization from random f16 data as TQ2_0 matches exactly ✅
INFO:test-quants:Testing MXFP4
DEBUG:test-quants:Quantizing to MXFP4 with Python
DEBUG:test-quants:Quantizing to MXFP4 with C
INFO:test-quants:Quantization to MXFP4 matches exactly ✅
DEBUG:test-quants:Dequantizing from MXFP4 with Python
DEBUG:test-quants:Dequantizing from MXFP4 with C
INFO:test-quants:Dequantization from MXFP4 matches exactly ✅
DEBUG:test-quants:Dequantizing random f16 data as MXFP4 with Python
gguf-py/gguf/quants.py:704: RuntimeWarning: overflow encountered in multiply
  return (d * qs.astype(np.float32))
DEBUG:test-quants:Dequantizing random f16 data as MXFP4 with C
INFO:test-quants:Dequantization from random f16 data as MXFP4 matches exactly ✅
INFO:test-quants:Testing IQ2_XXS
DEBUG:test-quants:Quantizing to IQ2_XXS with C
DEBUG:test-quants:Dequantizing from IQ2_XXS with Python
DEBUG:test-quants:Dequantizing from IQ2_XXS with C
INFO:test-quants:Dequantization from IQ2_XXS matches exactly ✅
DEBUG:test-quants:Dequantizing random f16 data as IQ2_XXS with Python
DEBUG:test-quants:Dequantizing random f16 data as IQ2_XXS with C
INFO:test-quants:Dequantization from random f16 data as IQ2_XXS matches exactly ✅
INFO:test-quants:Testing IQ2_XS
DEBUG:test-quants:Quantizing to IQ2_XS with C
DEBUG:test-quants:Dequantizing from IQ2_XS with Python
DEBUG:test-quants:Dequantizing from IQ2_XS with C
INFO:test-quants:Dequantization from IQ2_XS matches exactly ✅
DEBUG:test-quants:Dequantizing random f16 data as IQ2_XS with Python
DEBUG:test-quants:Dequantizing random f16 data as IQ2_XS with C
INFO:test-quants:Dequantization from random f16 data as IQ2_XS matches exactly ✅
INFO:test-quants:Testing IQ2_S
DEBUG:test-quants:Quantizing to IQ2_S with C
DEBUG:test-quants:Dequantizing from IQ2_S with Python
DEBUG:test-quants:Dequantizing from IQ2_S with C
INFO:test-quants:Dequantization from IQ2_S matches exactly ✅
DEBUG:test-quants:Dequantizing random f16 data as IQ2_S with Python
DEBUG:test-quants:Dequantizing random f16 data as IQ2_S with C
INFO:test-quants:Dequantization from random f16 data as IQ2_S matches exactly ✅
INFO:test-quants:Testing IQ3_XXS
DEBUG:test-quants:Quantizing to IQ3_XXS with C
DEBUG:test-quants:Dequantizing from IQ3_XXS with Python
DEBUG:test-quants:Dequantizing from IQ3_XXS with C
INFO:test-quants:Dequantization from IQ3_XXS matches exactly ✅
DEBUG:test-quants:Dequantizing random f16 data as IQ3_XXS with Python
DEBUG:test-quants:Dequantizing random f16 data as IQ3_XXS with C
INFO:test-quants:Dequantization from random f16 data as IQ3_XXS matches exactly ✅
INFO:test-quants:Testing IQ3_S
DEBUG:test-quants:Quantizing to IQ3_S with C
DEBUG:test-quants:Dequantizing from IQ3_S with Python
DEBUG:test-quants:Dequantizing from IQ3_S with C
INFO:test-quants:Dequantization from IQ3_S matches exactly ✅
DEBUG:test-quants:Dequantizing random f16 data as IQ3_S with Python
DEBUG:test-quants:Dequantizing random f16 data as IQ3_S with C
INFO:test-quants:Dequantization from random f16 data as IQ3_S matches exactly ✅
INFO:test-quants:Testing IQ1_S
DEBUG:test-quants:Quantizing to IQ1_S with C
DEBUG:test-quants:Dequantizing from IQ1_S with Python
DEBUG:test-quants:Dequantizing from IQ1_S with C
INFO:test-quants:Dequantization from IQ1_S matches exactly ✅
DEBUG:test-quants:Dequantizing random f16 data as IQ1_S with Python
DEBUG:test-quants:Dequantizing random f16 data as IQ1_S with C
INFO:test-quants:Dequantization from random f16 data as IQ1_S matches exactly ✅
INFO:test-quants:Testing IQ1_M
DEBUG:test-quants:Quantizing to IQ1_M with C
DEBUG:test-quants:Dequantizing from IQ1_M with Python
DEBUG:test-quants:Dequantizing from IQ1_M with C
INFO:test-quants:Dequantization from IQ1_M matches exactly ✅
DEBUG:test-quants:Dequantizing random f16 data as IQ1_M with Python
DEBUG:test-quants:Dequantizing random f16 data as IQ1_M with C
INFO:test-quants:Dequantization from random f16 data as IQ1_M matches exactly ✅
INFO:test-quants:Testing IQ4_NL
DEBUG:test-quants:Quantizing to IQ4_NL with C
DEBUG:test-quants:Dequantizing from IQ4_NL with Python
DEBUG:test-quants:Dequantizing from IQ4_NL with C
INFO:test-quants:Dequantization from IQ4_NL matches exactly ✅
DEBUG:test-quants:Dequantizing random f16 data as IQ4_NL with Python
DEBUG:test-quants:Dequantizing random f16 data as IQ4_NL with C
INFO:test-quants:Dequantization from random f16 data as IQ4_NL matches exactly ✅
INFO:test-quants:Testing IQ4_XS
DEBUG:test-quants:Quantizing to IQ4_XS with C
DEBUG:test-quants:Dequantizing from IQ4_XS with Python
DEBUG:test-quants:Dequantizing from IQ4_XS with C
INFO:test-quants:Dequantization from IQ4_XS matches exactly ✅
DEBUG:test-quants:Dequantizing random f16 data as IQ4_XS with Python
DEBUG:test-quants:Dequantizing random f16 data as IQ4_XS with C
INFO:test-quants:Dequantization from random f16 data as IQ4_XS matches exactly ✅
```
</details>

---

*Make sure to read the [contributing guidelines](https://github.com/ggml-org/llama.cpp/blob/master/CONTRIBUTING.md) before submitting a PR*


Changed files:
filename: ggml/src/ggml-quants.c
patch: @@ -288,7 +288,7 @@ void quantize_row_mxfp4_ref(const float * GGML_RESTRICT x, block_mxfp4 * GGML_RE
             }
         }
 
-        const uint8_t e = (uint8_t) (floorf(log2f(amax)) - 2 + 127);
+        const uint8_t e = amax > 0.0f ? (uint8_t) (floorf(log2f(amax)) - 2 + 127) : 0;
 
         const float d = GGML_E8M0_TO_FP32_HALF(e);
 

filename: gguf-py/gguf/quants.py
patch: @@ -228,8 +228,7 @@ def quantize_blocks(cls, blocks: np.ndarray) -> np.ndarray:
         d = max / -8
         with np.errstate(divide="ignore"):
             id = np.where(d == 0, 0, 1 / d)
-        # FIXME: Q4_0's reference rounding is cursed and depends on FMA
-        qs = np.trunc((np.float64(blocks) * np.float64(id)) + np.float64(8.5), dtype=np.float32).astype(np.uint8).clip(0, 15)
+        qs = np.trunc((blocks * id) + np.float32(8.5), dtype=np.float32).astype(np.uint8).clip(0, 15)
 
         qs = qs.reshape((n_blocks, 2, cls.block_size // 2))
         qs = qs[..., 0, :] | (qs[..., 1, :] << np.uint8(4))
@@ -300,8 +299,7 @@ def quantize_blocks(cls, blocks: np.ndarray) -> np.ndarray:
         d = max / -16
         with np.errstate(divide="ignore"):
             id = np.where(d == 0, 0, 1 / d)
-        # FIXME: Q5_0's reference rounding is cursed and depends on FMA
-        q = np.trunc((np.float64(blocks) * np.float64(id)) + np.float64(16.5), dtype=np.float32).astype(np.uint8).clip(0, 31)
+        q = np.trunc((blocks * id) + np.float32(16.5), dtype=np.float32).astype(np.uint8).clip(0, 31)
 
         qs = q.reshape((n_blocks, 2, cls.block_size // 2))
         qs = (qs[..., 0, :] & np.uint8(0x0F)) | (qs[..., 1, :] << np.uint8(4))
@@ -655,6 +653,57 @@ def dequantize_blocks(cls, blocks: np.ndarray) -> np.ndarray:
         return (d * qs.astype(np.float32))
 
 
+class MXFP4(__Quant, qtype=GGMLQuantizationType.MXFP4):
+    # e2m1 values (doubled)
+    # ref: https://www.opencompute.org/documents/ocp-microscaling-formats-mx-v1-0-spec-final-pdf
+    kvalues = (0, 1, 2, 3, 4, 6, 8, 12, 0, -1, -2, -3, -4, -6, -8, -12)
+
+    @staticmethod
+    # see ggml_e8m0_to_fp32_half in ggml-impl.h
+    def e8m0_to_fp32_half(x: np.ndarray) -> np.ndarray:
+        bits = np.where(x < 2, np.uint32(0x00200000) << np.uint32(x), np.uint32(x - 1) << np.uint32(23))
+        return bits.view(np.float32)
+
+    @classmethod
+    def quantize_blocks(cls, blocks: np.ndarray) -> np.ndarray:
+        n_blocks = blocks.shape[0]
+
+        d = abs(blocks).max(axis=-1, keepdims=True)
+
+        with np.errstate(divide="ignore"):
+            e = np.where(d > 0, np.floor(np.log2(d)) - 2 + 127, 0).astype(np.uint8)
+
+        d = cls.e8m0_to_fp32_half(e)
+
+        kvalues = np.array(cls.kvalues, dtype=np.int8).reshape((1, 1, 16))
+
+        errs = np.abs(d.reshape((n_blocks, 1, 1)) * kvalues.astype(np.float32) - blocks.reshape((n_blocks, cls.block_size, 1)))
+        best = np.argmin(errs, axis=-1, keepdims=True)
+
+        qs = best.reshape(n_blocks, 2, cls.block_size // 2).astype(np.uint8)
+        qs = qs[:, 0] | (qs[:, 1] << np.uint8(4))
+
+        qs = qs.reshape((n_blocks, cls.block_size // 2))
+
+        return np.concatenate([e, qs], axis=-1)
+
+    @classmethod
+    def dequantize_blocks(cls, blocks: np.ndarray) -> np.ndarray:
+        n_blocks = blocks.shape[0]
+
+        e, qs = np.hsplit(blocks, [1])
+
+        d = cls.e8m0_to_fp32_half(e)
+
+        qs = qs.reshape((n_blocks, 1, cls.block_size // 2)) >> np.array([0, 4], dtype=np.uint8).reshape((1, 2, 1))
+        qs = (qs & np.uint8(0x0F)).view(np.int8)
+
+        kvalues = np.array(cls.kvalues, dtype=np.int8).reshape(1, 1, 16)
+        qs = np.take_along_axis(kvalues, qs, axis=-1).reshape((n_blocks, cls.block_size))
+
+        return (d * qs.astype(np.float32))
+
+
 class IQ2_XXS(__Quant, qtype=GGMLQuantizationType.IQ2_XXS):
     ksigns: bytes = (
         b"\x00\x81\x82\x03\x84\x05\x06\x87\x88\x09\x0a\x8b\x0c\x8d\x8e\x0f"

filename: gguf-py/tests/test_quants.py
patch: @@ -67,6 +67,7 @@ def __init__(self, libggml: Path):
             "q4_0", "q4_1", "q5_0", "q5_1", "q8_0",
             "q2_K", "q3_K", "q4_K", "q5_K", "q6_K",
             "tq1_0", "tq2_0",
+            "mxfp4",
             "iq2_xxs", "iq2_xs", "iq2_s", "iq3_xxs", "iq3_s", "iq1_s", "iq1_m",
             "iq4_nl", "iq4_xs",
         ):
@@ -140,14 +141,21 @@ def compare_tensors(t1: np.ndarray, t2: np.ndarray, qtype: GGMLQuantizationType)
         return False
 
 
-def do_test(libggml_path: Path, quick: bool = False):
+def do_test(libggml_path: Path, quick: bool = False, user_type: GGMLQuantizationType | None = None):
     ggml_quants = GGMLQuants(libggml_path)
 
     np.set_printoptions(precision=None, threshold=(4 * 256) + 1, formatter={"int": lambda n: "0x%02X" % n})
 
     r = np.random.randn(8, 1024, 1024).astype(np.float32, copy=False)
-
-    for qtype in (GGMLQuantizationType.F16, *gguf.quants._type_traits.keys()):
+    # test zero blocks
+    r[0, 0, :] = 0
+    ## Maybe test infinities? (can make NANs, not really useful in practice)
+    # r[0, 1, 0] = np.inf
+    # r[0, 2, 0] = -np.inf
+    # r[0, 3, 0] = np.inf
+    # r[0, 3, 1] = -np.inf
+
+    for qtype in ((GGMLQuantizationType.F16, *gguf.quants._type_traits.keys()) if user_type is None else (user_type,)):
         has_dequantize = False
         has_quantize = False
 
@@ -228,11 +236,12 @@ def do_test(libggml_path: Path, quick: bool = False):
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser(description="Test Python (de)quantization against the reference C implementation")
-    parser.add_argument("--libggml", type=Path, default=Path(__file__).parent.parent.parent / "build" / "ggml" / "src" / "libggml.so", help="The path to libggml.so")
+    parser.add_argument("--libggml", type=Path, default=Path(__file__).parent.parent.parent / "build" / "bin" / "libggml.so", help="The path to libggml.so")
     parser.add_argument("--quick", action="store_true", help="Don't quantize with C when it's not strictly necessary")
+    parser.add_argument("--type", type=str, help="The quant type to test (all by default)")
 
     args = parser.parse_args()
 
     logging.basicConfig(level=logging.DEBUG)
 
-    do_test(args.libggml, args.quick)
+    do_test(args.libggml, args.quick, GGMLQuantizationType[args.type.upper()] if args.type is not None else None)