perf(ggml/hip): optimize ROCMFPX fp3/fp2 MMVQ on RDNA (perm-LUT + VDR=4 + fp2)

.github/workflows/ci.yml

@@ -79,12 +79,12 @@ jobs:
             -DCMAKE_BUILD_TYPE=Release
           cmake --build build --target \
             test_dflash test_generate test_flash_attn_sparse test_server_unit \
-            test_deepseek4_unit -j$(nproc)
+            test_deepseek4_unit test_rocmfpx -j$(nproc)
 
       - name: Run C++ server unit tests
         run: |
           cd server/build
-          ctest --output-on-failure -R "server_unit|deepseek4_unit" --no-tests=error
+          ctest --output-on-failure -R "server_unit|deepseek4_unit|rocmfpx" --no-tests=error
 
       - name: Populate venv with cu128 torch + setuptools
         # First pass: install the workspace's default deps. dflash declares

server/CMakeLists.txt

@@ -718,6 +718,22 @@ if(DFLASH27B_TESTS)
         target_link_libraries(test_anchor_params PRIVATE dflash_common)
         add_test(NAME anchor_params COMMAND test_anchor_params)
     endif()
+    # ROCMFPX fp3/fp2 MMVQ decode-contract unit test (host-only, no GPU): pins the
+    # code->value tables, the packed-bit unpacking, and the half-block scale split
+    # that vec_dot_rocmfpx_fp3/fp2_q8_1 decode from, against the CPU dequantizer.
+    # Self-contained (compiles rocmfpx.c directly) so it runs on the hosted CI job.
+    if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/deps/llama.cpp/ggml/rocmfpx/test_rocmfpx.c")
+        add_executable(test_rocmfpx
+            deps/llama.cpp/ggml/rocmfpx/test_rocmfpx.c
+            deps/llama.cpp/ggml/rocmfpx/rocmfpx.c)
+        target_include_directories(test_rocmfpx PRIVATE
+            deps/llama.cpp/ggml/include
+            deps/llama.cpp/ggml/rocmfpx)
+        if(NOT WIN32)
+            target_link_libraries(test_rocmfpx PRIVATE m)
+        endif()
+        add_test(NAME rocmfpx_unit COMMAND test_rocmfpx)
+    endif()
     if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/test/test_drafter_early_exit_score_range.cpp")
         add_executable(test_drafter_early_exit_score_range
             test/test_drafter_early_exit_score_range.cpp)

server/deps/llama.cpp/ggml/rocmfpx/test_rocmfpx.c

@@ -40,6 +40,166 @@ static float weighted_mse(const float * a, const float * b, const float * w, int
     return sum_w > 0.0f ? err / sum_w : 0.0f;
 }
 
+// ── Kernel-decode correctness ────────────────────────────────────────────
+// The GPU MMVQ kernels (vec_dot_rocmfpx_fp3_q8_1 / _fp2_q8_1) don't dequantize
+// to float — they decode each packed code to an int8 kvalue, integer-dp4a it
+// against the q8 activations per half-block, then scale by the half-block's
+// ue4m3 e[]. These host tests pin the three things that decode contract depends
+// on — the code->value table, the packed-bit unpacking, and the half-block scale
+// assignment — against the independent library dequantizer. The device perm-LUT
+// decode is in turn tied to this same code->value mapping by a compile-time
+// static_assert in vecdotq.cuh, so a divergence fails the GPU build too.
+
+// fp3 code -> signed kvalue, mirroring rocmfpx_decode_fp3_code_vec_cuda:
+// mag = (code&3)==3 ? 4 : code&3;  value = (code&4) ? -mag : mag.  => {0,1,2,4,0,-1,-2,-4}
+static int fp3_code_value(unsigned code) {
+    const unsigned mag_code = code & 3u;
+    const int mag = (mag_code == 3u) ? 4 : (int) mag_code;
+    return (code & 4u) ? -mag : mag;
+}
+
+static int fp2_code_value(unsigned code) {
+    static const int kv[4] = {
+        ROCMFP2_KVALUE_0_I8, ROCMFP2_KVALUE_1_I8, ROCMFP2_KVALUE_2_I8, ROCMFP2_KVALUE_3_I8
+    };
+    return kv[code & 3u];
+}
+
+// Unpack fp3 code i from the little-endian 3-bit stream (matches the kernel's
+// dword-splice extraction: (stream >> 3i) & 7). 16-bit window handles a code
+// straddling a byte boundary; the guard avoids reading past the 12-byte qs.
+static unsigned fp3_get_code(const uint8_t * qs, int i) {
+    const int bit = 3 * i;
+    const int byte = bit >> 3;
+    unsigned window = qs[byte];
+    if (byte + 1 < QS_ROCMFP3) {
+        window |= (unsigned) qs[byte + 1] << 8;
+    }
+    return (window >> (bit & 7)) & 7u;
+}
+
+// fp2 code i: byte i/4, 2 bits at 2*(i%4) (matches rocmfpx_pack4_fp2_vec_cuda).
+static unsigned fp2_get_code(const uint8_t * qs, int i) {
+    return (qs[i >> 2] >> (2 * (i & 3))) & 3u;
+}
+
+// Table pinned exhaustively — this is the contract the device perm-LUT encodes.
+static void check_decode_tables(void) {
+    const int fp3_expect[8] = { 0, 1, 2, 4, 0, -1, -2, -4 };
+    for (unsigned c = 0; c < 8; ++c) {
+        assert(fp3_code_value(c) == fp3_expect[c]);
+    }
+    const int fp2_expect[4] = { -1, 0, 1, 2 };
+    for (unsigned c = 0; c < 4; ++c) {
+        assert(fp2_code_value(c) == fp2_expect[c]);
+    }
+    printf("decode tables: fp3 {0,1,2,4,0,-1,-2,-4} fp2 {-1,0,1,2} OK\n");
+}
+
+// Reconstruct each weight from code->value * half-block scale using the kernel's
+// unpack, and assert it equals the library dequantizer bit-for-bit. This pins
+// the bit order, the code table, AND the e[i>=QK/2] half-block scale split.
+static void check_decode_contract_fp3(void) {
+    enum { N = 2 * QK_ROCMFP3 };
+    float src[N], ref[N];
+    block_rocmfp3 q[N / QK_ROCMFP3];
+    fill_row(src, N);
+    rocmfpx_quantize_fp3(src, q, 1, N, NULL);
+    rocmfpx_dequantize_row_fp3(q, ref, N);
+    for (int b = 0; b < N / QK_ROCMFP3; ++b) {
+        const block_rocmfp3 * blk = &q[b];
+        const float e0 = rocmfpx_ue4m3_to_fp32(blk->e[0]);
+        const float e1 = rocmfpx_ue4m3_to_fp32(blk->e[1]);
+        for (int i = 0; i < QK_ROCMFP3; ++i) {
+            const float e = (i >= QK_ROCMFP3 / 2) ? e1 : e0;
+            const float recon = (float) fp3_code_value(fp3_get_code(blk->qs, i)) * e;
+            assert(recon == ref[b * QK_ROCMFP3 + i]);
+        }
+    }
+    printf("fp3 decode contract: unpack + table + half-block scale == dequantize_row (exact) OK\n");
+}
+
+static void check_decode_contract_fp2(void) {
+    enum { N = 2 * QK_ROCMFP2 };
+    float src[N], ref[N];
+    block_rocmfp2 q[N / QK_ROCMFP2];
+    fill_row(src, N);
+    rocmfpx_quantize_fp2(src, q, 1, N, NULL);
+    rocmfpx_dequantize_row_fp2(q, ref, N);
+    for (int b = 0; b < N / QK_ROCMFP2; ++b) {
+        const block_rocmfp2 * blk = &q[b];
+        const float e0 = rocmfpx_ue4m3_to_fp32(blk->e[0]);
+        const float e1 = rocmfpx_ue4m3_to_fp32(blk->e[1]);
+        for (int i = 0; i < QK_ROCMFP2; ++i) {
+            const float e = (i >= QK_ROCMFP2 / 2) ? e1 : e0;
+            const float recon = (float) fp2_code_value(fp2_get_code(blk->qs, i)) * e;
+            assert(recon == ref[b * QK_ROCMFP2 + i]);
+        }
+    }
+    printf("fp2 decode contract: unpack + table + half-block scale == dequantize_row (exact) OK\n");
+}
+
+// The kernel's dot is db * sum_half( e_half * sum_i code_value[i]*q8[i] ) with an
+// exact integer inner sum per half-block (the VDR=4 half-block structure for fp3
+// / the two-chain fp2 form). Assert it matches the float reference dot to tol.
+static void check_dot_structure_fp3(void) {
+    enum { N = QK_ROCMFP3 };
+    float src[N], w[N];
+    block_rocmfp3 q[1];
+    fill_row(src, N);
+    rocmfpx_quantize_fp3(src, q, 1, N, NULL);
+    rocmfpx_dequantize_row_fp3(q, w, N);
+
+    int8_t a8[N];
+    for (int i = 0; i < N; ++i) { a8[i] = (int8_t) ((i * 37) % 251 - 125); }
+    const float db = 0.031f;
+
+    double ref = 0.0;
+    for (int i = 0; i < N; ++i) { ref += (double) w[i] * (double) (db * (float) a8[i]); }
+
+    const block_rocmfp3 * blk = &q[0];
+    long s0 = 0, s1 = 0;
+    for (int i = 0; i < N; ++i) {
+        const int c = fp3_code_value(fp3_get_code(blk->qs, i));
+        if (i < N / 2) { s0 += (long) c * a8[i]; } else { s1 += (long) c * a8[i]; }
+    }
+    const double e0 = rocmfpx_ue4m3_to_fp32(blk->e[0]);
+    const double e1 = rocmfpx_ue4m3_to_fp32(blk->e[1]);
+    const double kern = (double) db * (e0 * (double) s0 + e1 * (double) s1);
+    const double rel = fabs(kern - ref) / (fabs(ref) + 1e-9);
+    printf("fp3 dot structure: ref=%g kernel=%g rel=%g\n", ref, kern, rel);
+    assert(rel < 1e-5);
+}
+
+static void check_dot_structure_fp2(void) {
+    enum { N = QK_ROCMFP2 };
+    float src[N], w[N];
+    block_rocmfp2 q[1];
+    fill_row(src, N);
+    rocmfpx_quantize_fp2(src, q, 1, N, NULL);
+    rocmfpx_dequantize_row_fp2(q, w, N);
+
+    int8_t a8[N];
+    for (int i = 0; i < N; ++i) { a8[i] = (int8_t) ((i * 37) % 251 - 125); }
+    const float db = 0.031f;
+
+    double ref = 0.0;
+    for (int i = 0; i < N; ++i) { ref += (double) w[i] * (double) (db * (float) a8[i]); }
+
+    const block_rocmfp2 * blk = &q[0];
+    long s0 = 0, s1 = 0;
+    for (int i = 0; i < N; ++i) {
+        const int c = fp2_code_value(fp2_get_code(blk->qs, i));
+        if (i < N / 2) { s0 += (long) c * a8[i]; } else { s1 += (long) c * a8[i]; }
+    }
+    const double e0 = rocmfpx_ue4m3_to_fp32(blk->e[0]);
+    const double e1 = rocmfpx_ue4m3_to_fp32(blk->e[1]);
+    const double kern = (double) db * (e0 * (double) s0 + e1 * (double) s1);
+    const double rel = fabs(kern - ref) / (fabs(ref) + 1e-9);
+    printf("fp2 dot structure: ref=%g kernel=%g rel=%g\n", ref, kern, rel);
+    assert(rel < 1e-5);
+}
+
 static void check_weighted_imatrix_fp3(void) {
     enum { N = QK_ROCMFP3 };
 
@@ -160,5 +320,12 @@ int main(void) {
     check_weighted_imatrix_fp2();
     check_weighted_imatrix_fp3();
 
+    // Kernel-decode correctness (MMVQ fp3/fp2 vec-dot contract).
+    check_decode_tables();
+    check_decode_contract_fp3();
+    check_decode_contract_fp2();
+    check_dot_structure_fp3();
+    check_dot_structure_fp2();
+
     return 0;
 }

server/deps/llama.cpp/ggml/src/ggml-cuda/mmvq.cu

@@ -415,6 +415,51 @@ static constexpr __host__ __device__ int calc_rows_per_block(int ncols_dst, int
     return 1;
 }
 
+#if defined(GGML_USE_HIP)
+// ROCMFPX (fp2/fp3) cross-lane float reduction. On RDNA, HIP's warp_reduce_sum(float)
+// lowers its __shfl_xor butterfly to ds_bpermute, which builds a per-lane
+// address VGPR (lane ^ offset) that sits on the dependent shuffle chain. For an
+// xor mask inside a 32-lane group the SAME permutation is one ds_swizzle_b32
+// bitmask-mode op with the xor encoded as an instruction immediate
+// ((offset<<10)|0x1f: xor=offset, or=0, and=0x1f) — no address VGPR. It is
+// bit-identical to warp_reduce_sum's ladder: identical {16,8,4,2,1} reduction
+// order, identical float adds (ds_swizzle by mask m gives src lane (lane&0x1f)^m,
+// exactly __shfl_xor's partner; verified 0 lane mismatches on gfx1201 for those
+// masks). warp==32 on RDNA so every offset is a valid intra-group swizzle.
+template <int width>
+static __device__ __forceinline__ float warp_reduce_sum_rocmfpx_dsswizzle(float x) {
+    static_assert(width == 32, "ROCMFPX ds_swizzle reduce assumes a 32-lane RDNA warp");
+    x += __int_as_float(__builtin_amdgcn_ds_swizzle(__float_as_int(x), (16 << 10) | 0x1f));
+    x += __int_as_float(__builtin_amdgcn_ds_swizzle(__float_as_int(x), ( 8 << 10) | 0x1f));
+    x += __int_as_float(__builtin_amdgcn_ds_swizzle(__float_as_int(x), ( 4 << 10) | 0x1f));
+    x += __int_as_float(__builtin_amdgcn_ds_swizzle(__float_as_int(x), ( 2 << 10) | 0x1f));
+    x += __int_as_float(__builtin_amdgcn_ds_swizzle(__float_as_int(x), ( 1 << 10) | 0x1f));
+    return x;
+}
+#endif
+
+// MMVQ warp reduction, specialised to the ds_swizzle ladder for the ROCMFPX
+// fp2 (GGML_TYPE_Q2_0_ROCMFP2) and fp3 (GGML_TYPE_Q3_0_ROCMFPX) instantiations
+// only — every other quant type, and all non-HIP builds, keep the shared
+// warp_reduce_sum unchanged. Gating on `type` keeps the change local to these
+// kernels' codegen (no backend-wide effect). The `width == 32` guard falls back
+// to warp_reduce_sum for any non-wave32 HIP target (e.g. wave64 CDNA):
+// ds_swizzle is a 32-lane intra-group op, so the fast path only applies on RDNA
+// wave32 and everything else stays generic. The ladder is bit-identical to
+// warp_reduce_sum for every type (same {16,8,4,2,1} order, same float adds).
+template <int width, ggml_type type>
+static __device__ __forceinline__ float warp_reduce_sum_mmvq(float x) {
+#if defined(GGML_USE_HIP)
+    if constexpr ((type == GGML_TYPE_Q3_0_ROCMFPX || type == GGML_TYPE_Q2_0_ROCMFP2) && width == 32) {
+        return warp_reduce_sum_rocmfpx_dsswizzle<width>(x);
+    } else {
+        return warp_reduce_sum<width>(x);
+    }
+#else
+    return warp_reduce_sum<width>(x);
+#endif
+}
+
 template <ggml_type type, int ncols_dst, bool has_fusion, bool small_k = false>
 __launch_bounds__(calc_nwarps(type, ncols_dst, get_device_table_id())*ggml_cuda_get_physical_warp_size(), 1)
 static __global__ void mul_mat_vec_q(
@@ -585,10 +630,10 @@ static __global__ void mul_mat_vec_q(
                     }
                 }
             }
-            tmp[j][i] = warp_reduce_sum<warp_size>(tmp[j][i]);
+            tmp[j][i] = warp_reduce_sum_mmvq<warp_size, type>(tmp[j][i]);
             if constexpr (has_fusion) {
                 if (use_gate) {
-                    tmp_gate[j][i] = warp_reduce_sum<warp_size>(tmp_gate[j][i]);
+                    tmp_gate[j][i] = warp_reduce_sum_mmvq<warp_size, type>(tmp_gate[j][i]);
                 }
             }
         }