diff --git a/Makefile b/Makefile index f3c7916f..556235dc 100644 --- a/Makefile +++ b/Makefile @@ -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 diff --git a/cuda/mmq/ds4_mmq.cu b/cuda/mmq/ds4_mmq.cu index e37ff2c9..4b41f509 100644 --- a/cuda/mmq/ds4_mmq.cu +++ b/cuda/mmq/ds4_mmq.cu @@ -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<<>>( W, (const block_q8_1 *)src1_q8_1_ptr, ids, out_f32, diff --git a/ds4.c b/ds4.c index d9717dba..66c8ab30 100644 --- a/ds4.c +++ b/ds4.c @@ -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; } diff --git a/ds4_rocm.cu b/ds4_rocm.cu index beddc4cb..8fab0953 100644 --- a/ds4_rocm.cu +++ b/ds4_rocm.cu @@ -1,6 +1,7 @@ #ifdef __HIP_PLATFORM_AMD__ #include "ds4_rocm.h" #include +#include #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 diff --git a/rocm/ds4_rocm_attention.cuh b/rocm/ds4_rocm_attention.cuh index 9ac028ae..939bb292 100644 --- a/rocm/ds4_rocm_attention.cuh +++ b/rocm/ds4_rocm_attention.cuh @@ -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, diff --git a/rocm/ds4_rocm_attention_launch.cuh b/rocm/ds4_rocm_attention_launch.cuh index fed7c694..825d0495 100644 --- a/rocm/ds4_rocm_attention_launch.cuh +++ b/rocm/ds4_rocm_attention_launch.cuh @@ -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) { diff --git a/rocm/ds4_rocm_matmul.cuh b/rocm/ds4_rocm_matmul.cuh index 353a557e..42ec6c03 100644 --- a/rocm/ds4_rocm_matmul.cuh +++ b/rocm/ds4_rocm_matmul.cuh @@ -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, diff --git a/rocm/ds4_rocm_moe_launch.cuh b/rocm/ds4_rocm_moe_launch.cuh index 7891671e..c1420409 100644 --- a/rocm/ds4_rocm_moe_launch.cuh +++ b/rocm/ds4_rocm_moe_launch.cuh @@ -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; diff --git a/rocm/ds4_rocm_runtime.cuh b/rocm/ds4_rocm_runtime.cuh index 8eeca5c2..a4cfa91d 100644 --- a/rocm/ds4_rocm_runtime.cuh +++ b/rocm/ds4_rocm_runtime.cuh @@ -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;