(metal): dsa_kv_qat f16_round flag for the compressor emit path

CHANGELOG.md

@@ -6,6 +6,11 @@ adhere to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
 ## [Unreleased]
 
+### Added
+- `dsa_kv_qat` takes `f16_round=False`, which stops at the fp8 result and
+  copies the RoPE tail through unchanged. Fuses the DeepSeek-V4 compressor
+  emit-path quantization, which has no f16 cache step, into one dispatch.
+
 ## [0.3.11]
 
 ### Changed

bindings.cpp

@@ -855,6 +855,7 @@ NB_MODULE(_ext, m) {
       "x"_a,
       "n_rot"_a,
       nb::kw_only(),
+      "f16_round"_a = true,
       "stream"_a = nb::none(),
       R"(
         DeepSeek-V4-Flash main-attention KV QAT round-trip, fused: the
@@ -865,10 +866,17 @@ NB_MODULE(_ext, m) {
         One kernel in place of the split + fp8-core + concat + astype
         chain, bit-identically.
 
+        With ``f16_round`` false the fp16 step is dropped: the fp8 result
+        stays in the storage dtype and the RoPE tail is copied through
+        unchanged. That is the compressor emit-path form, where the pooled
+        row is quantized but never passes through the f16 KV cache; it
+        replaces the split + fp8-core + concat chain on its own.
+
         Args:
             x (array): any shape with trailing dim D,
                 (D - n_rot) % 64 == 0; float16/bfloat16/float32.
             n_rot (int): trailing RoPE dims excluded from the fp8 step.
+            f16_round (bool): apply the trailing fp16 round. Default True.
 
         Returns:
             array: same shape and dtype as ``x``.

docs/kernels.md

@@ -136,7 +136,9 @@ modifications, from omlx's `glm_moe_dsa` custom kernels (see the
   per-head attention sinks, in one flash-softmax dispatch (f32 accumulation).
 - **`dsa_kv_qat`** / **`dsa_indexer_qat`** - the fused quantization-aware round-trips DeepSeek-V4 does
   on its main-attention KV (per-64-block FP8-E4M3FN) and indexer activations (128-wide Hadamard then
-  per-32-block FP4-E2M1), each bit-identical to the equivalent MLX graph.
+  per-32-block FP4-E2M1), each bit-identical to the equivalent MLX graph. `dsa_kv_qat(...,
+  f16_round=False)` drops the trailing fp16 round for the compressor emit path, whose pooled rows are
+  quantized but never stored in the f16 KV cache.
 
 Tuning levers (defaults are right for normal use):
 

metal/kq_dsa_qat.metal

@@ -28,8 +28,10 @@ instantiate_kq_dsa_indexer_qat_pack(float, float);
 // clang-format on
 
 // clang-format off
-#define instantiate_kq_dsa_kv_qat(tname, dtype)                         \
-  instantiate_kernel("kq_dsa_kv_qat_" #tname, kq_dsa_kv_qat, dtype)
+#define instantiate_kq_dsa_kv_qat(tname, dtype)                            \
+  instantiate_kernel("kq_dsa_kv_qat_" #tname, kq_dsa_kv_qat, dtype, true)  \
+  instantiate_kernel(                                                      \
+      "kq_dsa_kv_qat_nof16_" #tname, kq_dsa_kv_qat, dtype, false)
 
 instantiate_kq_dsa_kv_qat(float16_t, half);
 instantiate_kq_dsa_kv_qat(bfloat16_t, bfloat16_t);

metal/mlx/backend/metal/kernels/kq_dsa_qat.h

@@ -298,7 +298,12 @@ template <typename T>
 // threads per row: one simdgroup per 64-block (lane owns elements
 // 64b + lane and 64b + 32 + lane), lanes re-derive the block scale from
 // the simd_max broadcast.
-template <typename T>
+//
+// F16R selects the trailing fp16 round. With F16R off the kernel stops at
+// the storage-dtype fp8 result and copies the RoPE tail through unchanged,
+// which is the compressor emit-path form (ds4.c compressor site: the
+// pooled row is quantized but never passes through the f16 KV cache).
+template <typename T, bool F16R>
 [[kernel, max_total_threads_per_threadgroup(256)]] void kq_dsa_kv_qat(
     const device T* X [[buffer(0)]],
     device T* O [[buffer(1)]],
@@ -335,11 +340,12 @@ template <typename T>
       e = metal::clamp(e, -6.0f, 8.0f);
       const float q = metal::ldexp(1.0f, int(e) - 3);
       const float r = sgn * metal::rint(a / q) * q * scale;
-      o_row[64 * b + 32 * j + int(simd_lid)] = T(half(float(T(r))));
+      o_row[64 * b + 32 * j + int(simd_lid)] =
+          F16R ? T(half(float(T(r)))) : T(r);
     }
   }
   for (int i = int(lid); i < NROT; i += 256) {
     const int e = (D - NROT) + i;
-    o_row[e] = T(half(float(x_row[e])));
+    o_row[e] = F16R ? T(half(float(x_row[e]))) : x_row[e];
   }
 }

src/kquant.h

@@ -493,8 +493,14 @@ mx::array dsa_indexer_scores_q(
 // n_rot RoPE dims fp8-exempt, then the whole row rounded through fp16 (the
 // f16 KV-cache step). x is any shape with trailing dim D where
 // (D - n_rot) % 64 == 0; returns the same shape and dtype. Bit-compatible
-// with the split + fp8-core + concat + astype chain. Metal-only.
-mx::array dsa_kv_qat(mx::array x, int n_rot, mx::StreamOrDevice s = {});
+// with the split + fp8-core + concat + astype chain. Set f16_round false
+// for the compressor emit-path form, which stops at the fp8 result and
+// passes the RoPE tail through unchanged. Metal-only.
+mx::array dsa_kv_qat(
+    mx::array x,
+    int n_rot,
+    bool f16_round = true,
+    mx::StreamOrDevice s = {});
 
 // K-quant gathered matvec (down projection), same wire layout. x [T, R, K]
 // (one row per expert slot), indices [T, R]. Optional per-(expert, out_dim)
@@ -1293,8 +1299,8 @@ class KQDsaIndexerScoresQ : public mx::Primitive {
 // dsa_kv_qat). Inference-only, Metal-only.
 class KQDsaKvQat : public mx::Primitive {
  public:
-  explicit KQDsaKvQat(mx::Stream stream, int n_rot)
-      : mx::Primitive(stream), n_rot_(n_rot) {}
+  explicit KQDsaKvQat(mx::Stream stream, int n_rot, bool f16_round)
+      : mx::Primitive(stream), n_rot_(n_rot), f16_round_(f16_round) {}
 
   void eval_cpu(
       const std::vector<mx::array>& inputs,
@@ -1313,6 +1319,7 @@ class KQDsaKvQat : public mx::Primitive {
 
  private:
   int n_rot_;
+  bool f16_round_;
 };
 
 // K-quant fused MoE GLU gather (see moe_glu_gather_kq). Inference-only.

src/kquant_dsa_qat.cpp

@@ -135,7 +135,8 @@ void KQDsaKvQat::eval_gpu(
   const int rows = int(x.size() / D);
   const int n_rot = n_rot_;
 
-  const std::string kname = "kq_dsa_kv_qat_" + kq_type_string(x.dtype());
+  const std::string kname = "kq_dsa_kv_qat_" +
+      std::string(f16_round_ ? "" : "nof16_") + kq_type_string(x.dtype());
   auto kernel = kq_get_kernel(d, kname, kname, {});
   auto& ce = mx::metal::get_command_encoder(s);
   ce.set_compute_pipeline_state(kernel);
@@ -246,10 +247,12 @@ std::vector<mx::Shape> KQDsaKvQat::output_shapes(
 }
 
 bool KQDsaKvQat::is_equivalent(const mx::Primitive& other) const {
-  return n_rot_ == static_cast<const KQDsaKvQat&>(other).n_rot_;
+  const auto& o = static_cast<const KQDsaKvQat&>(other);
+  return n_rot_ == o.n_rot_ && f16_round_ == o.f16_round_;
 }
 
-mx::array dsa_kv_qat(mx::array x, int n_rot, mx::StreamOrDevice s_) {
+mx::array
+dsa_kv_qat(mx::array x, int n_rot, bool f16_round, mx::StreamOrDevice s_) {
   auto s = mx::to_stream(s_);
 
   if (x.ndim() < 1) {
@@ -278,7 +281,7 @@ mx::array dsa_kv_qat(mx::array x, int n_rot, mx::StreamOrDevice s_) {
   return mx::array(
       std::move(out_shape),
       x.dtype(),
-      std::make_shared<KQDsaKvQat>(s, n_rot),
+      std::make_shared<KQDsaKvQat>(s, n_rot, f16_round),
       std::move(inputs));
 }
 

tests/test_dsa_kv_qat.py

@@ -11,6 +11,13 @@
 outputs are asserted BIT-identical, not just close: any drift would move
 KV-cache contents.
 
+`f16_round=False` is the compressor emit-path form: same fp8 arithmetic,
+no trailing fp16 round, RoPE tail copied through. Its reference is the
+gguf-mlx compressor site (`Compressor.__call__`, `_qat == "fp8"`), which
+is the same chain minus the astype pair, and it is asserted bit-identical
+against that: pooled rows land in the indexer pool and in attention keys,
+so drift there moves the top-k selection.
+
 Metal-only kernel (eval_cpu throws): skipped under KQUANT_FORCE_CPU.
 
 Usage: test_dsa_kv_qat.py
@@ -64,6 +71,12 @@ def _ref(kv, n_rot):
     return kv.astype(mx.float16).astype(orig)
 
 
+def _ref_nof16(kv, n_rot):
+    """gguf-mlx Compressor.__call__ emit path (_qat == "fp8")."""
+    nope, rot = kv[..., : kv.shape[-1] - n_rot], kv[..., kv.shape[-1] - n_rot :]
+    return mx.concatenate([_fp8_roundtrip(nope), rot], axis=-1)
+
+
 def _bits(a):
     bits_dtype = {2: mx.uint16, 4: mx.uint32}[a.itemsize]
     return np.array(a.view(bits_dtype))
@@ -95,24 +108,53 @@ def _bits(a):
 ]
 
 
+@pytest.mark.parametrize("f16_round", [True, False], ids=["f16", "nof16"])
 @pytest.mark.parametrize("dtype", [mx.float16, mx.bfloat16, mx.float32])
 @pytest.mark.parametrize("case", CASES, ids=[c[0] for c in CASES])
-def test_dsa_kv_qat_bit_identity(case, dtype):
+def test_dsa_kv_qat_bit_identity(case, dtype, f16_round):
     name, gen = case
     rng = np.random.default_rng(7)
     x = mx.array(gen(rng, (2048, 576), dtype)).astype(dtype)
     mx.eval(x)
-    got = kq.dsa_kv_qat(x, 64)
-    ref = _ref(x, 64)
+    got = kq.dsa_kv_qat(x, 64, f16_round=f16_round)
+    ref = (_ref if f16_round else _ref_nof16)(x, 64)
     mx.eval(got, ref)
     gb, rb = _bits(got), _bits(ref)
     mismatch = int((gb != rb).sum())
     assert mismatch == 0, (
-        f"{name} {dtype}: {mismatch}/{gb.size} words differ "
-        f"(first at {np.argwhere(gb != rb)[:4].tolist()})"
+        f"{name} {dtype} f16_round={f16_round}: {mismatch}/{gb.size} words "
+        f"differ (first at {np.argwhere(gb != rb)[:4].tolist()})"
     )
 
 
+@pytest.mark.parametrize("dtype", [mx.float16, mx.bfloat16, mx.float32])
+def test_dsa_kv_qat_nof16_emit_geometry(dtype):
+    """The V4-Flash compressor emit shape: head_dim 512, rope tail 64, one
+    or two pooled rows per call (ratio-4 overlap emits two, then trims)."""
+    rng = np.random.default_rng(11)
+    for rows in (1, 2):
+        x = mx.array(rng.standard_normal((1, rows, 512))).astype(dtype)
+        mx.eval(x)
+        got = kq.dsa_kv_qat(x, 64, f16_round=False)
+        ref = _ref_nof16(x, 64)
+        mx.eval(got, ref)
+        assert got.shape == x.shape
+        assert not int((_bits(got) != _bits(ref)).sum()), f"{dtype} rows={rows}"
+
+
+def test_dsa_kv_qat_nof16_differs_from_f16():
+    """Guard against the flag being ignored: bf16 values that survive the
+    fp8 step unchanged still move when they are re-rounded through fp16."""
+    xt = mx.full((4, 576), 1e20, dtype=mx.bfloat16)
+    mx.eval(xt)
+    a = kq.dsa_kv_qat(xt, 64, f16_round=True)
+    b = kq.dsa_kv_qat(xt, 64, f16_round=False)
+    mx.eval(a, b)
+    assert int((_bits(a) != _bits(b)).sum()) > 0
+    # rope tail: f16 saturates 1e20 to inf, the emit form leaves it alone
+    assert float(b[0, -1]) == float(xt[0, -1])
+
+
 def test_dsa_kv_qat_shapes_and_rejects():
     # 4-D decode shape [B, 1, L, D] and a non-576 geometry (128 + 64)
     x4 = mx.random.normal((2, 1, 3, 576)).astype(mx.bfloat16)
@@ -148,20 +190,22 @@ def test_dsa_kv_qat_shapes_and_rejects():
 
 def main() -> int:
     fails = 0
-    for name, gen in CASES:
-        for dtype in (mx.float16, mx.bfloat16, mx.float32):
-            rng = np.random.default_rng(7)
-            x = mx.array(gen(rng, (2048, 576), dtype)).astype(dtype)
-            mx.eval(x)
-            got = kq.dsa_kv_qat(x, 64)
-            ref = _ref(x, 64)
-            mx.eval(got, ref)
-            n = int((_bits(got) != _bits(ref)).sum())
-            fails += n > 0
-            print(
-                f"  {name:<16} {str(dtype):<18} "
-                f"{'bit-identical' if n == 0 else f'{n} words differ'}"
-            )
+    for f16_round in (True, False):
+        print(f"f16_round={f16_round}")
+        for name, gen in CASES:
+            for dtype in (mx.float16, mx.bfloat16, mx.float32):
+                rng = np.random.default_rng(7)
+                x = mx.array(gen(rng, (2048, 576), dtype)).astype(dtype)
+                mx.eval(x)
+                got = kq.dsa_kv_qat(x, 64, f16_round=f16_round)
+                ref = (_ref if f16_round else _ref_nof16)(x, 64)
+                mx.eval(got, ref)
+                n = int((_bits(got) != _bits(ref)).sum())
+                fails += n > 0
+                print(
+                    f"  {name:<16} {str(dtype):<18} "
+                    f"{'bit-identical' if n == 0 else f'{n} words differ'}"
+                )
     print("ALL OK" if not fails else f"FAILURES: {fails}")
     return 1 if fails else 0