perf(cuda): add warp-cooperative TurboQuant FWHT

server/CMakeLists.txt

@@ -977,6 +977,17 @@ if(DFLASH27B_TESTS)
             ggml-base)
         list(APPEND _raw_unit_test_targets test_mmq_streamk_iq4_xs)
     endif()
+    if(DFLASH27B_GPU_BACKEND STREQUAL "cuda" AND EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/test/test_turbo_wht_warp.cu")
+        find_package(CUDAToolkit REQUIRED)
+        add_executable(test_turbo_wht_warp test/test_turbo_wht_warp.cu)
+        set_target_properties(test_turbo_wht_warp PROPERTIES CUDA_ARCHITECTURES "${_dflash_archs}")
+        target_include_directories(test_turbo_wht_warp PRIVATE
+            ${CMAKE_CURRENT_SOURCE_DIR}/deps/llama.cpp/ggml/include
+            ${CMAKE_CURRENT_SOURCE_DIR}/deps/llama.cpp/ggml/src
+            ${CMAKE_CURRENT_SOURCE_DIR}/deps/llama.cpp/ggml/src/ggml-cuda)
+        target_link_libraries(test_turbo_wht_warp PRIVATE CUDA::cudart)
+        list(APPEND _raw_unit_test_targets test_turbo_wht_warp)
+    endif()
     if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/test/test_drafter_tail_capture_guard.cpp")
         # RED phase binary: same source WITHOUT the fix flag — documents the bug.
         add_executable(test_drafter_tail_capture_guard_red

server/deps/llama.cpp/ggml/src/ggml-cuda/turbo-wht.cu

@@ -2,11 +2,13 @@
 #include "tq3-quant.cuh"
 #include "cpy-utils.cuh"
 
-// Each thread independently transforms one 128-element group.
 // Supports non-contiguous src via separate src/dst strides (dim0 must be
 // contiguous in both). This lets us skip ggml_cont before turbo_wht when
 // the input comes from ggml_permute with dim0 unchanged.
-static __global__ void k_turbo_wht(
+#if defined(GGML_USE_HIP)
+// Keep the established scalar implementation on HIP. The cooperative helper
+// below is a 32-lane CUDA-warp primitive and has not been qualified on wave64.
+static __global__ void k_turbo_wht_scalar(
         const char * __restrict__ src_base,
         char       * __restrict__ dst_base,
         const int64_t ne00,
@@ -41,6 +43,55 @@ static __global__ void k_turbo_wht(
 
     for (int i = 0; i < 128; i++) out_row[i] = x[i];
 }
+#else
+// One CUDA warp transforms one 128-element group (four values per lane).
+// The previous one-thread implementation kept a 128-float local array per
+// thread, causing register spills and 23-42 us launches for only 1-3 blocks on
+// sm_86. This uses the same warp-cooperative primitive already exercised by
+// the chunked-attention path.
+static __global__ void k_turbo_wht_warp(
+        const char * __restrict__ src_base,
+        char       * __restrict__ dst_base,
+        const int64_t ne00,
+        const int64_t ne01,
+        const int64_t ne02,
+        const int64_t src_nb1,
+        const int64_t src_nb2,
+        const int64_t dst_nb1,
+        const int64_t dst_nb2,
+        const int64_t total_groups,
+        const int64_t groups_per_row,
+        int           direction) {
+    constexpr int warp_size_local = 32;
+    const int warp = threadIdx.x / warp_size_local;
+    const int lane = threadIdx.x & (warp_size_local - 1);
+    const int64_t gid = (int64_t)blockIdx.x * (blockDim.x / warp_size_local) + warp;
+    if (gid >= total_groups) return;
+
+    const int64_t g   = gid % groups_per_row;
+    const int64_t rem = gid / groups_per_row;
+    const int64_t i01 = rem % ne01;
+    const int64_t i02 = rem / ne01;
+
+    const float * row = (const float *)(src_base + i01 * src_nb1 + i02 * src_nb2) + g * QK_TQ3_0_GROUP;
+    float * out_row   = (float *)(dst_base + i01 * dst_nb1 + i02 * dst_nb2) + g * QK_TQ3_0_GROUP;
+    const int base = lane * 4;
+
+    float v0 = row[base + 0];
+    float v1 = row[base + 1];
+    float v2 = row[base + 2];
+    float v3 = row[base + 3];
+    if (direction == 0) {
+        warp_tq3_rotate_forward(v0, v1, v2, v3);
+    } else {
+        warp_tq3_rotate_inverse(v0, v1, v2, v3);
+    }
+    out_row[base + 0] = v0;
+    out_row[base + 1] = v1;
+    out_row[base + 2] = v2;
+    out_row[base + 3] = v3;
+}
+#endif
 
 // Fused kernel: FWHT-rotate a non-contiguous F32 source and quantize directly
 // to Q4_0 (or Q8_0). Eliminates the intermediate F32 buffer and two kernel
@@ -101,13 +152,22 @@ void ggml_cuda_op_turbo_wht(ggml_backend_cuda_context & ctx, ggml_tensor * dst)
     const int64_t total_groups   = groups_per_row * ne01 * ne02;
 
     constexpr int THREADS_PER_BLOCK = 128;
+#if defined(GGML_USE_HIP)
     const int n_blocks = (int)((total_groups + THREADS_PER_BLOCK - 1) / THREADS_PER_BLOCK);
+#else
+    constexpr int GROUPS_PER_BLOCK = THREADS_PER_BLOCK / 32;
+    const int n_blocks = (int)((total_groups + GROUPS_PER_BLOCK - 1) / GROUPS_PER_BLOCK);
+#endif
 
     // Destination strides are always contiguous
     const int64_t dst_nb1 = ne00 * sizeof(float);
     const int64_t dst_nb2 = ne00 * ne01 * sizeof(float);
 
-    k_turbo_wht<<<n_blocks, THREADS_PER_BLOCK, 0, ctx.stream()>>>(
+#if defined(GGML_USE_HIP)
+    k_turbo_wht_scalar<<<n_blocks, THREADS_PER_BLOCK, 0, ctx.stream()>>>(
+#else
+    k_turbo_wht_warp<<<n_blocks, THREADS_PER_BLOCK, 0, ctx.stream()>>>(
+#endif
         (const char *)src0->data, (char *)dst->data,
         ne00, ne01, ne02,
         src0->nb[1], src0->nb[2],

server/test/test_turbo_wht_warp.cu

@@ -0,0 +1,187 @@
+#include "common.cuh"
+#include "tq3-quant.cuh"
+
+#include <cuda_runtime.h>
+
+#include <cstdint>
+#include <cstdio>
+#include <cstdlib>
+#include <vector>
+
+#define TEST_CUDA_CHECK(expr) do { \
+    const cudaError_t err = (expr); \
+    if (err != cudaSuccess) { \
+        std::fprintf(stderr, "%s:%d: %s\n", __FILE__, __LINE__, cudaGetErrorString(err)); \
+        std::exit(1); \
+    } \
+} while (0)
+
+static __global__ void scalar_reference(
+        const float * input, float * output, int64_t groups, int direction) {
+    const int64_t group = (int64_t) blockIdx.x * blockDim.x + threadIdx.x;
+    if (group >= groups) {
+        return;
+    }
+    float values[128];
+    for (int i = 0; i < 128; ++i) {
+        values[i] = input[group * 128 + i];
+    }
+    if (direction == 0) {
+        tq3_rotate_forward(values);
+    } else {
+        tq3_rotate_inverse(values);
+    }
+    for (int i = 0; i < 128; ++i) {
+        output[group * 128 + i] = values[i];
+    }
+}
+
+static __global__ void warp_candidate(
+        const float * input, float * output, int64_t groups, int direction) {
+    constexpr int warp_size = 32;
+    const int warp = threadIdx.x / warp_size;
+    const int lane = threadIdx.x & (warp_size - 1);
+    const int64_t group = (int64_t) blockIdx.x * (blockDim.x / warp_size) + warp;
+    if (group >= groups) {
+        return;
+    }
+    const int64_t base = group * 128 + lane * 4;
+    float v0 = input[base + 0];
+    float v1 = input[base + 1];
+    float v2 = input[base + 2];
+    float v3 = input[base + 3];
+    if (direction == 0) {
+        warp_tq3_rotate_forward(v0, v1, v2, v3);
+    } else {
+        warp_tq3_rotate_inverse(v0, v1, v2, v3);
+    }
+    output[base + 0] = v0;
+    output[base + 1] = v1;
+    output[base + 2] = v2;
+    output[base + 3] = v3;
+}
+
+static bool run_case(int64_t groups, int direction) {
+    const int64_t count = groups * 128;
+    std::vector<float> input((size_t) count);
+    for (int64_t i = 0; i < count; ++i) {
+        input[(size_t) i] = (float) ((i * 37 + 11) % 257 - 128) / 64.0f;
+    }
+
+    float * d_input = nullptr;
+    float * d_reference = nullptr;
+    float * d_candidate = nullptr;
+    TEST_CUDA_CHECK(cudaMalloc(&d_input, count * sizeof(float)));
+    TEST_CUDA_CHECK(cudaMalloc(&d_reference, count * sizeof(float)));
+    TEST_CUDA_CHECK(cudaMalloc(&d_candidate, count * sizeof(float)));
+    TEST_CUDA_CHECK(cudaMemcpy(d_input, input.data(), count * sizeof(float), cudaMemcpyHostToDevice));
+
+    scalar_reference<<<(groups + 127) / 128, 128>>>(d_input, d_reference, groups, direction);
+    warp_candidate<<<(groups + 3) / 4, 128>>>(d_input, d_candidate, groups, direction);
+    TEST_CUDA_CHECK(cudaGetLastError());
+    TEST_CUDA_CHECK(cudaDeviceSynchronize());
+
+    std::vector<float> reference((size_t) count);
+    std::vector<float> candidate((size_t) count);
+    TEST_CUDA_CHECK(cudaMemcpy(reference.data(), d_reference, count * sizeof(float), cudaMemcpyDeviceToHost));
+    TEST_CUDA_CHECK(cudaMemcpy(candidate.data(), d_candidate, count * sizeof(float), cudaMemcpyDeviceToHost));
+
+    int64_t mismatches = 0;
+    for (int64_t i = 0; i < count; ++i) {
+        if (reference[(size_t) i] != candidate[(size_t) i]) {
+            ++mismatches;
+            if (mismatches <= 4) {
+                std::fprintf(stderr, "groups=%lld direction=%d i=%lld ref=%a got=%a\n",
+                    (long long) groups, direction, (long long) i,
+                    reference[(size_t) i], candidate[(size_t) i]);
+            }
+        }
+    }
+
+    cudaFree(d_candidate);
+    cudaFree(d_reference);
+    cudaFree(d_input);
+    std::printf("[%s] groups=%lld direction=%d mismatches=%lld\n",
+        mismatches == 0 ? "PASS" : "FAIL", (long long) groups,
+        direction, (long long) mismatches);
+    return mismatches == 0;
+}
+
+static double benchmark_case(int64_t groups, bool warp_kernel, int iterations) {
+    const int64_t count = groups * 128;
+    std::vector<float> input((size_t) count, 0.125f);
+    float * d_input = nullptr;
+    float * d_output = nullptr;
+    cudaEvent_t start = nullptr;
+    cudaEvent_t stop = nullptr;
+    TEST_CUDA_CHECK(cudaMalloc(&d_input, count * sizeof(float)));
+    TEST_CUDA_CHECK(cudaMalloc(&d_output, count * sizeof(float)));
+    TEST_CUDA_CHECK(cudaMemcpy(d_input, input.data(), count * sizeof(float), cudaMemcpyHostToDevice));
+    TEST_CUDA_CHECK(cudaEventCreate(&start));
+    TEST_CUDA_CHECK(cudaEventCreate(&stop));
+
+    for (int i = 0; i < 100; ++i) {
+        if (warp_kernel) {
+            warp_candidate<<<(groups + 3) / 4, 128>>>(d_input, d_output, groups, 0);
+        } else {
+            scalar_reference<<<(groups + 127) / 128, 128>>>(d_input, d_output, groups, 0);
+        }
+    }
+    TEST_CUDA_CHECK(cudaDeviceSynchronize());
+    TEST_CUDA_CHECK(cudaEventRecord(start));
+    for (int i = 0; i < iterations; ++i) {
+        if (warp_kernel) {
+            warp_candidate<<<(groups + 3) / 4, 128>>>(d_input, d_output, groups, 0);
+        } else {
+            scalar_reference<<<(groups + 127) / 128, 128>>>(d_input, d_output, groups, 0);
+        }
+    }
+    TEST_CUDA_CHECK(cudaEventRecord(stop));
+    TEST_CUDA_CHECK(cudaEventSynchronize(stop));
+    float elapsed_ms = 0.0f;
+    TEST_CUDA_CHECK(cudaEventElapsedTime(&elapsed_ms, start, stop));
+
+    cudaEventDestroy(stop);
+    cudaEventDestroy(start);
+    cudaFree(d_output);
+    cudaFree(d_input);
+    return (double) elapsed_ms * 1000.0 / iterations;
+}
+
+int main() {
+    int device_count = 0;
+    const cudaError_t device_status = cudaGetDeviceCount(&device_count);
+    if (device_status == cudaErrorNoDevice) {
+        std::puts("SKIP: no CUDA device");
+        return 0;
+    }
+    TEST_CUDA_CHECK(device_status);
+    if (device_count == 0) {
+        std::puts("SKIP: no CUDA device");
+        return 0;
+    }
+    TEST_CUDA_CHECK(cudaSetDevice(0));
+
+    const int64_t group_counts[] = {1, 4, 32, 128, 384};
+    int failures = 0;
+    for (const int direction : {0, 1}) {
+        for (const int64_t groups : group_counts) {
+            if (!run_case(groups, direction)) {
+                ++failures;
+            }
+        }
+    }
+    if (failures != 0) {
+        std::fprintf(stderr, "FAILED: %d cases\n", failures);
+        return 1;
+    }
+    std::puts("ALL PASS: scalar and warp FWHT outputs are bit-identical");
+    for (const int64_t groups : {128LL, 384LL}) {
+        const double scalar_us = benchmark_case(groups, false, 10000);
+        const double warp_us = benchmark_case(groups, true, 10000);
+        std::printf("BENCH groups=%lld scalar=%.3f us warp=%.3f us speedup=%.2fx reduction=%.1f%%\n",
+            (long long) groups, scalar_us, warp_us, scalar_us / warp_us,
+            100.0 * (1.0 - warp_us / scalar_us));
+    }
+    return 0;
+}