From cffab3b8f0b2b0e537bd9d6247c639802dbc3cea Mon Sep 17 00:00:00 2001 From: Asher Feldman <59994+asher@users.noreply.github.com> Date: Tue, 11 Aug 2026 06:56:46 -0700 Subject: [PATCH] (metal): dsa_kv_qat f16_round flag for the compressor emit path --- CHANGELOG.md | 5 ++ bindings.cpp | 8 ++ docs/kernels.md | 4 +- metal/kq_dsa_qat.metal | 6 +- metal/mlx/backend/metal/kernels/kq_dsa_qat.h | 12 ++- src/kquant.h | 15 +++- src/kquant_dsa_qat.cpp | 11 ++- tests/test_dsa_kv_qat.py | 82 +++++++++++++++----- 8 files changed, 110 insertions(+), 33 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index b22906f..d476007 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/bindings.cpp b/bindings.cpp index 5acfadb..bf80d85 100644 --- a/bindings.cpp +++ b/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``. diff --git a/docs/kernels.md b/docs/kernels.md index 15900be..2192f11 100644 --- a/docs/kernels.md +++ b/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): diff --git a/metal/kq_dsa_qat.metal b/metal/kq_dsa_qat.metal index 2c3d10c..30946d5 100644 --- a/metal/kq_dsa_qat.metal +++ b/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); diff --git a/metal/mlx/backend/metal/kernels/kq_dsa_qat.h b/metal/mlx/backend/metal/kernels/kq_dsa_qat.h index 328de95..4ae2bdf 100644 --- a/metal/mlx/backend/metal/kernels/kq_dsa_qat.h +++ b/metal/mlx/backend/metal/kernels/kq_dsa_qat.h @@ -298,7 +298,12 @@ template // 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 +// +// 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 [[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 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]; } } diff --git a/src/kquant.h b/src/kquant.h index f330cfa..b0edd25 100644 --- a/src/kquant.h +++ b/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& 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. diff --git a/src/kquant_dsa_qat.cpp b/src/kquant_dsa_qat.cpp index 0f71700..8db39df 100644 --- a/src/kquant_dsa_qat.cpp +++ b/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 KQDsaKvQat::output_shapes( } bool KQDsaKvQat::is_equivalent(const mx::Primitive& other) const { - return n_rot_ == static_cast(other).n_rot_; + const auto& o = static_cast(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(s, n_rot), + std::make_shared(s, n_rot, f16_round), std::move(inputs)); } diff --git a/tests/test_dsa_kv_qat.py b/tests/test_dsa_kv_qat.py index b059dc5..dd886ac 100644 --- a/tests/test_dsa_kv_qat.py +++ b/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