perf(rocm): accelerate gfx1151 DSpark verification

2026-08-29 21:51:45 +01:00
1 parent 307dc093d9
commit 227fa190aa
@@ -56,7 +56,7 @@ HIPCC ?= $(shell command -v hipcc 2>/dev/null || echo /opt/rocm/bin/hipcc)
 ROCM_ARCH ?= gfx1151
 ROCM_HOST_CFLAGS ?= -fPIC
 ROCM_CFLAGS ?= -O3 -ffast-math -g -fno-finite-math-only -pthread -D__HIP_PLATFORM_AMD__ -Wno-unused-command-line-argument --offload-arch=$(ROCM_ARCH)
-ROCM_LDLIBS ?= -lm -pthread -lhipblas -lhipblaslt
+ROCM_LDLIBS ?= -lm -pthread -lhipblas -lhipblaslt -lrocblas
 ROCM_MMQ_Y ?= 64
 ROCM_MMQ_FLAGS := $(ROCM_CFLAGS) -std=c++17 -DGGML_USE_HIP -DDS4_HIP_MMQ_Y=$(ROCM_MMQ_Y) $(MMQ_INCLUDES)
 ROCM_MMQ_OBJS := cuda/mmq/ds4_ggml_stubs.rocm.o cuda/mmq/ds4_mmq.rocm.o cuda/mmq/quantize.rocm.o cuda/mmq/mmid.rocm.o cuda/mmq/mmvq.rocm.o cuda/mmq/d2r_stubs.rocm.o
@@ -3271,8 +3271,8 @@ static __global__ void ds4_mmq_moe_down_sum6_q8_1_qwarp32_kernel(
    const int kqs = vdr * (lane % lanes_per_k);

 #pragma unroll
-    for (uint32_t rr = 0; rr < 4u; ++rr) {
-        const uint32_t row = blockIdx.x * 64u + row_lane + rr * 16u;
+    for (uint32_t rr = 0; rr < 8u; ++rr) {
+        const uint32_t row = blockIdx.x * 64u + row_lane + rr * 8u;
        if (row >= nrows_x) continue;
        float total = 0.0f;
 #pragma unroll
@@ -3516,7 +3516,7 @@ int ds4_mmq_moe_down_sum6_vec_impl(
    const uint32_t stride_channel_x = (uint32_t)((int64_t)M * stride_row_x);

    const dim3 block_nums((M + 63) / 64, n_tokens);
-    const dim3 block_dims(256);
+    const dim3 block_dims(128);

    ds4_mmq_moe_down_sum6_q8_1_qwarp32_kernel<type><<<block_nums, block_dims, 0, stream>>>(
        W, (const block_q8_1 *)src1_q8_1_ptr, ids, out_f32,
@@ -32879,7 +32879,8 @@ static bool metal_graph_eval_dspark_base_logits(
        const ds4_model          *base_model,
        const ds4_weights        *base_weights,
        const ds4_model          *dspark_model,
-        const ds4_dspark_weights *dw) {
+        const ds4_dspark_weights *dw,
+        uint32_t                  draft_cap) {
    if (!g || !base_model || !base_weights || !dspark_model || !dw ||
        !dspark_final_head_ready(g, base_weights, dw)) {
        return false;
@@ -32887,7 +32888,9 @@ static bool metal_graph_eval_dspark_base_logits(

    const ds4_dspark_stage_weights *final =
        &dw->stage[dw->n_stages - 1u];
-    const uint32_t draft = dw->block_size;
+    const uint32_t draft =
+        draft_cap != 0 && draft_cap < dw->block_size ?
+            draft_cap : dw->block_size;
    const uint64_t hc_dim = (uint64_t)DS4_N_HC * DS4_N_EMBD;
    const uint64_t vocab_dim = base_weights->output->dim[1];
    ds4_gpu_tensor *stage_output_hc = metal_graph_dspark_final_output_hc(g);
@@ -33360,11 +33363,77 @@ static bool dspark_disable_reuse_confidence0_markov(void) {
    return cache != 0;
 }

+#ifndef __APPLE__
+static bool dspark_apply_markov_greedy_gpu_runtime(
+        ds4_gpu_graph            *g,
+        const ds4_model          *dspark_model,
+        const ds4_dspark_weights *dw,
+        int                       first_prev_token,
+        uint32_t                  draft_cap,
+        int32_t                   proposal[DS4_DSPARK_MAX_BLOCK_SIZE],
+        uint32_t                 *proposal_len) {
+    if (proposal_len) *proposal_len = 0;
+    if (!g || !g->spec_logits || !g->dspark_draft_tokens ||
+        !dspark_model || !dw || !proposal ||
+        first_prev_token < 0 ||
+        (uint32_t)first_prev_token >= DS4_N_VOCAB ||
+        !dspark_markov_probe_ready(dw)) {
+        return false;
+    }
+    const ds4_dspark_stage_weights *final =
+        &dw->stage[dw->n_stages - 1u];
+    if (dw->markov_rank == 0 || (dw->markov_rank & 31u) != 0 ||
+        final->markov_w1->type != DS4_TENSOR_Q8_0 ||
+        final->markov_w2->type != DS4_TENSOR_Q8_0) {
+        return false;
+    }
+
+    const uint64_t logits_bytes =
+        (uint64_t)DS4_N_VOCAB * sizeof(float);
+    int32_t prev_token = first_prev_token;
+    const uint32_t draft_count =
+        draft_cap != 0 && draft_cap < dw->block_size ?
+            draft_cap : dw->block_size;
+    for (uint32_t draft = 0; draft < draft_count; draft++) {
+        ds4_gpu_tensor *row_view =
+            ds4_gpu_tensor_view(g->spec_logits,
+                                (uint64_t)draft * logits_bytes,
+                                logits_bytes);
+        uint64_t gpu_key = 0;
+        const bool ok = row_view &&
+            ds4_gpu_dspark_markov_argmax_tensor(
+                g->dspark_draft_tokens,
+                row_view,
+                dspark_model->map,
+                dspark_model->size,
+                final->markov_w1->abs_offset,
+                final->markov_w2->abs_offset,
+                (uint32_t)prev_token,
+                DS4_N_VOCAB,
+                dw->markov_rank) != 0 &&
+            ds4_gpu_tensor_read(g->dspark_draft_tokens,
+                                0,
+                                &gpu_key,
+                                sizeof(gpu_key)) != 0;
+        ds4_gpu_tensor_free(row_view);
+        const uint32_t gpu_token = ~(uint32_t)(gpu_key & 0xffffffffu);
+        if (!ok || gpu_key == 0 || gpu_token >= DS4_N_VOCAB) {
+            return false;
+        }
+        proposal[draft] = (int32_t)gpu_token;
+        prev_token = (int32_t)gpu_token;
+    }
+    if (proposal_len) *proposal_len = draft_count;
+    return true;
+}
+#endif
+
 static bool dspark_apply_markov_greedy_probe(
        float                  *logits,
        const ds4_model        *dspark_model,
        const ds4_dspark_weights *dw,
        int                     first_prev_token,
+        uint32_t                draft_cap,
        float                  *markov_state,
        float                  *markov_bias,
        int32_t                 proposal[DS4_DSPARK_MAX_BLOCK_SIZE],
@@ -33386,7 +33455,10 @@ static bool dspark_apply_markov_greedy_probe(
        &dw->stage[dw->n_stages - 1u];
    const bool no_bias = dspark_markov_bias_disabled();
    int32_t prev_token = first_prev_token;
-    for (uint32_t draft = 0; draft < dw->block_size; draft++) {
+    const uint32_t draft_count =
+        draft_cap != 0 && draft_cap < dw->block_size ?
+            draft_cap : dw->block_size;
+    for (uint32_t draft = 0; draft < draft_count; draft++) {
        float *row = logits + (uint64_t)draft * DS4_N_VOCAB;
        if (!no_bias) {
            if (!dspark_dense_row_to_f32(markov_state,
@@ -33405,7 +33477,7 @@ static bool dspark_apply_markov_greedy_probe(
        prev_token = (int32_t)token;
    }

-    if (proposal_len) *proposal_len = dw->block_size;
+    if (proposal_len) *proposal_len = draft_count;
    return true;
 }

@@ -61236,6 +61308,11 @@ static bool ds4_session_prepare_dspark_draft_impl(ds4_session *s,
        const bool batch_capture_ok =
            ds4_session_dspark_capture_batch_current(s);
        const ds4_dspark_weights *dw = &s->engine->dspark_weights;
+        const uint32_t verify_cap =
+            ds4_dspark_env_u32("DS4_DSPARK_VERIFY_CAP", 0);
+        const uint32_t proposal_cap =
+            verify_cap != 0 && verify_cap < dw->block_size ?
+                verify_cap : dw->block_size;
        const bool stage0_ready = dspark_stage0_weights_ready(&s->graph, dw);
        const bool runtime_fused_stage0_setup =
            enabled && !fake_argmax_enabled && !probe_log;
@@ -61427,7 +61504,8 @@ static bool ds4_session_prepare_dspark_draft_impl(ds4_session *s,
                                                    &s->engine->model,
                                                    &s->engine->weights,
                                                    &s->engine->mtp_model,
-                                                    dw);
+                                                    dw,
+                                                    proposal_cap);
            DS4_DSPARK_PROP_ADD(propose_logits_ms, logits_t0);
        }
        const bool markov_ready =
@@ -61456,9 +61534,27 @@ static bool ds4_session_prepare_dspark_draft_impl(ds4_session *s,
            DS4_DSPARK_PROP_ADD(propose_markov_ms, markov_t0);
            confidence_ok = markov_ok;
        } else if (markov_ready) {
+#ifndef __APPLE__
+            const char *gpu_markov =
+                getenv("DS4_ROCM_DSPARK_GREEDY_GPU_MARKOV");
+            if (gpu_markov && gpu_markov[0] != '0' &&
+                !dspark_markov_bias_disabled() &&
+                getenv("DS4_DSPARK_NO_GPU_MARKOV") == NULL) {
+                const double markov_t0 = DS4_DSPARK_PROP_T0();
+                markov_ok = dspark_apply_markov_greedy_gpu_runtime(
+                    &s->graph,
+                    &s->engine->mtp_model,
+                    dw,
+                    token,
+                    proposal_cap,
+                    markov_proposal,
+                    &markov_proposal_len);
+                DS4_DSPARK_PROP_ADD(propose_markov_ms, markov_t0);
+            }
+#endif
            const uint64_t logits_count =
-                (uint64_t)dw->block_size * (uint64_t)DS4_N_VOCAB;
-            if (logits_count != 0 &&
+                (uint64_t)proposal_cap * (uint64_t)DS4_N_VOCAB;
+            if (!markov_ok && logits_count != 0 &&
                logits_count <= (uint64_t)SIZE_MAX / sizeof(float)) {
                const double markov_t0 = DS4_DSPARK_PROP_T0();
                const uint64_t logits_bytes = logits_count * sizeof(float);
@@ -61477,6 +61573,7 @@ static bool ds4_session_prepare_dspark_draft_impl(ds4_session *s,
                        &s->engine->mtp_model,
                        dw,
                        token,
+                        proposal_cap,
                        markov_state,
                        markov_bias,
                        markov_proposal,
@@ -61552,6 +61649,9 @@ static bool ds4_session_prepare_dspark_draft_impl(ds4_session *s,
                proposal_len = confidence_ok ? confidence_prefix_len : 0;
            }
            s->dspark_draft_len = proposal_len;
+            if (verify_cap != 0 && s->dspark_draft_len > verify_cap) {
+                s->dspark_draft_len = verify_cap;
+            }
            if (s->dspark_draft_len > DS4_DSPARK_MAX_BLOCK_SIZE) {
                s->dspark_draft_len = DS4_DSPARK_MAX_BLOCK_SIZE;
            }
@@ -1,6 +1,7 @@
 #ifdef __HIP_PLATFORM_AMD__
 #include "ds4_rocm.h"
 #include <hipblaslt/hipblaslt.h>
+#include <rocblas/rocblas.h>

 #define FULL_WARP_MASK 0xFFFFFFFFFFFFFFFFULL
 #define MASK_T uint64_t
@@ -46,7 +47,13 @@ extern "C" int ds4_mmq_iq2_xxs_moe_pair(
    const int32_t *ids, float *out_a, float *out_b,
    int M, int K, int n_tokens, int n_experts, int n_expert_used,
    cudaStream_t stream);
-
+extern "C" int ds4_mmq_q8_0_dense_vec(
+    const void *W, const float *X_f32, float *out_f32,
+    int M, int N, int K, cudaStream_t stream);
+extern "C" int ds4_mmq_q2_K_moe_down_sum6_vec(
+    const void *W, const float *X, const int32_t *ids, float *out,
+    int M, int K, int n_tokens, int n_experts, int n_expert_used,
+    cudaStream_t stream);
 #ifndef M_PI
 #define M_PI 3.14159265358979323846
 #endif
@@ -496,6 +496,22 @@ __global__ static void attention_pack_group_heads_f16_kernel(
    dst[gid] = __float2half(heads[((uint64_t)t * n_groups + g) * group_dim + d]);
 }

+__global__ static void attention_pack_group_heads_f32_kernel(
+        float *dst,
+        const float *heads,
+        uint32_t n_tokens,
+        uint32_t n_groups,
+        uint32_t group_dim) {
+    uint64_t gid = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x;
+    uint64_t n = (uint64_t)n_groups * n_tokens * group_dim;
+    if (gid >= n) return;
+    uint32_t d = gid % group_dim;
+    uint64_t q = gid / group_dim;
+    uint32_t t = q % n_tokens;
+    uint32_t g = q / n_tokens;
+    dst[gid] = heads[((uint64_t)t * n_groups + g) * group_dim + d];
+}
+
 __global__ static void attention_unpack_group_low_kernel(
        float *low,
        const float *tmp,
@@ -1454,6 +1454,65 @@ extern "C" int ds4_gpu_attention_output_q8_batch_tensor(
                rank);
        if (!cuda_ok(cudaGetLastError(), "attention_output_q8_a unpack launch")) return 0;
    } else {
+        const char *grouped_mmvq =
+            getenv("DS4_ROCM_DSPARK_ATTN_A_MMVQ");
+        if (!g_quality_mode && n_tokens == 4u &&
+            group_dim == 4096u && rank == 1024u && n_groups == 8u &&
+            grouped_mmvq && grouped_mmvq[0] != '0' &&
+            ds4_mmq_init(0) == 0) {
+            const uint64_t heads_count =
+                (uint64_t)n_groups * n_tokens * group_dim;
+            const uint64_t low_count =
+                (uint64_t)n_groups * n_tokens * rank;
+            const uint64_t heads_bytes = heads_count * sizeof(float);
+            const uint64_t low_offset = (heads_bytes + 255u) & ~255ull;
+            void *tmp = cuda_tmp_alloc(
+                low_offset + low_count * sizeof(float),
+                "attention output a q4 grouped MMVQ");
+            if (!tmp) return 0;
+            float *heads_packed = (float *)tmp;
+            float *low_packed = (float *)((char *)tmp + low_offset);
+            attention_pack_group_heads_f32_kernel
+                <<<(heads_count + 255u) / 256u, 256u>>>(
+                    heads_packed,
+                    (const float *)heads->ptr,
+                    n_tokens,
+                    n_groups,
+                    (uint32_t)group_dim);
+            if (!cuda_ok(cudaGetLastError(),
+                         "attention_output_q8_a q4 MMVQ pack")) {
+                return 0;
+            }
+            for (uint32_t group = 0; group < n_groups; ++group) {
+                const unsigned char *group_w =
+                    out_a + (uint64_t)group * rank * blocks_a * 34u;
+                const float *group_x =
+                    heads_packed + (uint64_t)group * n_tokens * group_dim;
+                float *group_out =
+                    low_packed + (uint64_t)group * n_tokens * rank;
+                if (ds4_mmq_q8_0_dense_vec(
+                        group_w,
+                        group_x,
+                        group_out,
+                        (int)rank,
+                        (int)n_tokens,
+                        (int)group_dim,
+                        (cudaStream_t)0) != 0) {
+                    return 0;
+                }
+            }
+            attention_unpack_group_low_kernel
+                <<<(low_count + 255u) / 256u, 256u>>>(
+                    (float *)low->ptr,
+                    low_packed,
+                    n_tokens,
+                    n_groups,
+                    (uint32_t)rank);
+            if (!cuda_ok(cudaGetLastError(),
+                         "attention_output_q8_a q4 MMVQ unpack")) {
+                return 0;
+            }
+        } else {
        const uint64_t x_rows = (uint64_t)n_tokens * n_groups;
        const uint64_t xq_bytes = x_rows * blocks_a * 32u;
        const uint64_t scale_offset = (xq_bytes + 15u) & ~15ull;
@@ -1482,6 +1541,7 @@ extern "C" int ds4_gpu_attention_output_q8_batch_tensor(
                                                          blocks_a,
                                                          use_dp4a);
        if (!cuda_ok(cudaGetLastError(), "attention_output_q8_a preq launch")) return 0;
+        }
    }

    if (attn_output_cublas && !g_quality_mode) {
@@ -495,6 +495,17 @@ static int cuda_matmul_q8_0_tensor_labeled(ds4_gpu_tensor *out, const void *mode
    }
    if (n_tok > 1) {
 #if defined(__HIP_PLATFORM_AMD__) || defined(__HIPCC__)
+        const char *tiny_q8_mmvq = getenv("DS4_ROCM_DSPARK_Q8_MMVQ");
+        if (!g_quality_mode && n_tok <= 4u &&
+            in_dim <= INT_MAX && out_dim <= INT_MAX &&
+            tiny_q8_mmvq && tiny_q8_mmvq[0] != '0' &&
+            ds4_mmq_init(0) == 0 &&
+            ds4_mmq_q8_0_dense_vec(
+                wptr, (const float *)x->ptr, (float *)out->ptr,
+                (int)out_dim, (int)n_tok, (int)in_dim,
+                (cudaStream_t)0) == 0) {
+            return 1;
+        }
        /* The 2048-row projection hits a 72-VGPR occupancy cliff in exact8;
         * its retained 24-VGPR kernel is 2.4x faster on gfx1151. */
        if (!g_quality_mode && (in_dim % 256u) == 0u &&
@@ -1028,6 +1039,54 @@ extern "C" int ds4_gpu_matmul_f16_tensor(ds4_gpu_tensor *out, const void *model_
        }
        const float alpha = 1.0f;
        const float beta = 0.0f;
+#ifdef __HIP_PLATFORM_AMD__
+        if (n_tok == 4u && g_rocblas_ready &&
+            ds4_rocm_gfx1151_flag("DS4_ROCM_F16_Q4_SOLUTIONS")) {
+            int32_t solution = 0;
+            if (in_dim == 4096u &&
+                (out_dim == 64u || out_dim == 256u ||
+                 out_dim == 512u || out_dim == 1024u)) {
+                solution = -217;
+            } else if (in_dim == 1024u && out_dim == 8192u) {
+                solution = -216;
+            }
+            if (solution != 0) {
+                const rocblas_status rst = rocblas_gemm_ex(
+                        g_rocblas,
+                        rocblas_operation_transpose,
+                        rocblas_operation_none,
+                        (rocblas_int)out_dim,
+                        (rocblas_int)n_tok,
+                        (rocblas_int)in_dim,
+                        &alpha,
+                        w,
+                        rocblas_datatype_f16_r,
+                        (rocblas_int)in_dim,
+                        xh,
+                        rocblas_datatype_f16_r,
+                        (rocblas_int)in_dim,
+                        &beta,
+                        out->ptr,
+                        rocblas_datatype_f32_r,
+                        (rocblas_int)out_dim,
+                        out->ptr,
+                        rocblas_datatype_f32_r,
+                        (rocblas_int)out_dim,
+                        rocblas_datatype_f32_r,
+                        rocblas_gemm_algo_solution_index,
+                        solution,
+                        0u);
+                if (rst == rocblas_status_success) return 1;
+                fprintf(stderr,
+                        "ds4: rocBLAS q4 F16 solution %d failed for %llux%llux%llu: status %d\n",
+                        solution,
+                        (unsigned long long)out_dim,
+                        (unsigned long long)n_tok,
+                        (unsigned long long)in_dim,
+                        (int)rst);
+            }
+        }
+#endif
        cublasStatus_t st = cublasGemmEx(g_cublas,
                                         CUBLAS_OP_T,
                                         CUBLAS_OP_N,
@@ -1765,11 +1765,30 @@ static int routed_moe_launch(
            } else if (split_ptr_down_done) {
                /* The split pointer-table path writes the final token row. */
            } else if (use_iq2_q2_float_down) {
-                ok = routed_moe_q2_float_down_launch(
-                        out, down, mid, iq2_hot_mid_h, use_iq2_hot_f16_mid, down_w,
-                        sorted_counts, sorted_offsets, sorted_pairs, tile_experts,
-                        n_tokens, n_total_expert, n_expert, expert_mid_dim, out_dim,
-                        down_expert_bytes, down_row_bytes);
+                const char *mmvq_down =
+                    getenv("DS4_ROCM_DSPARK_Q2_DOWN_MMVQ");
+                if (!g_quality_mode && n_tokens <= 4u &&
+                    mmvq_down && mmvq_down[0] != '0' &&
+                    ds4_mmq_init(0) == 0 &&
+                    ds4_mmq_q2_K_moe_down_sum6_vec(
+                        down_w,
+                        (const float *)mid->ptr,
+                        (const int32_t *)selected_exec->ptr,
+                        (float *)out->ptr,
+                        (int)out_dim,
+                        (int)expert_mid_dim,
+                        (int)n_tokens,
+                        (int)n_total_expert,
+                        (int)n_expert,
+                        (cudaStream_t)0) == 0) {
+                    ok = 1;
+                } else {
+                    ok = routed_moe_q2_float_down_launch(
+                            out, down, mid, iq2_hot_mid_h, use_iq2_hot_f16_mid, down_w,
+                            sorted_counts, sorted_offsets, sorted_pairs, tile_experts,
+                            n_tokens, n_total_expert, n_expert, expert_mid_dim, out_dim,
+                            down_expert_bytes, down_row_bytes);
+                }
            } else {
            dim3 dgrid((out_dim + 31u) / 32u, pair_count, 1);
            uint32_t *down_tile_total = tile_total;
@@ -18,6 +18,8 @@ static uint64_t g_selected_readback_event_value;
 static cublasHandle_t g_cublas;
 static int g_cublas_ready;
 #ifdef __HIP_PLATFORM_AMD__
+static rocblas_handle g_rocblas;
+static int g_rocblas_ready;
 #include "ds4_rocm_hipblaslt.cuh"
 #endif
 static int g_quality_mode;
@@ -5840,6 +5842,14 @@ extern "C" int ds4_gpu_init(void) {
        g_cublas_ready = 1;
    }
 #ifdef __HIP_PLATFORM_AMD__
+    if (!g_rocblas_ready) {
+        const rocblas_status st = rocblas_create_handle(&g_rocblas);
+        if (st != rocblas_status_success) {
+            fprintf(stderr, "ds4: rocBLAS create handle failed: status %d\n", (int)st);
+            return 0;
+        }
+        g_rocblas_ready = 1;
+    }
    if (!g_hipblaslt_ready) {
        if (hipblaslt_ok(hipblasLtCreate(&g_hipblaslt), "create handle")) {
            g_hipblaslt_ready = 1;
@@ -5862,6 +5872,11 @@ extern "C" void ds4_gpu_cleanup(void) {
        g_cublas = NULL;
    }
 #ifdef __HIP_PLATFORM_AMD__
+    if (g_rocblas_ready) {
+        (void)rocblas_destroy_handle(g_rocblas);
+        g_rocblas_ready = 0;
+        g_rocblas = NULL;
+    }
    if (g_hipblaslt_ready) {
        (void)hipblasLtDestroy(g_hipblaslt);
        g_hipblaslt_ready = 0;