From 1e8f16ca05b9aae1034680d606ef0d5c4b93ef0e Mon Sep 17 00:00:00 2001 From: Donato Capitella Date: Tue, 28 Jul 2026 13:33:07 +0100 Subject: [PATCH 1/5] Restore ROCm DeepSeek prequant decode kernels --- ds4.c | 9 ++ rocm/ds4_rocm_attention_launch.cuh | 78 +++++++++++--- rocm/ds4_rocm_current_api_compat.cuh | 2 +- rocm/ds4_rocm_matmul.cuh | 148 +++++++++++++++++++++------ rocm/ds4_rocm_runtime.cuh | 24 +++++ 5 files changed, 216 insertions(+), 45 deletions(-) diff --git a/ds4.c b/ds4.c index b767a78692..a3e1fd4898 100644 --- a/ds4.c +++ b/ds4.c @@ -55955,6 +55955,15 @@ static int ds4_engine_open_internal(ds4_engine **out, } e->metal_ready = true; ds4_gpu_set_quality(e->quality); +#ifdef DS4_ROCM_BUILD + /* + * The ROCm Q8 decode selector must know that a multi-tier model + * is GLM before any layer dispatch. The single-tier path sets + * the same model-family state below. + */ + ds4_gpu_set_glm_model( + DS4_MODEL_FAMILY == DS4_MODEL_FAMILY_GLM_DSA); +#endif (void)ds4_gpu_set_model_fd(e->model.fd); if (engine_install_gpu_placement(e) != 0) { diff --git a/rocm/ds4_rocm_attention_launch.cuh b/rocm/ds4_rocm_attention_launch.cuh index 56563b4506..4b2534bfc0 100644 --- a/rocm/ds4_rocm_attention_launch.cuh +++ b/rocm/ds4_rocm_attention_launch.cuh @@ -1303,28 +1303,78 @@ extern "C" int ds4_gpu_attention_output_low_q8_tensor( (uint32_t)low_dim); return cuda_ok(cudaGetLastError(), "attention_output_low_q8 splitk sum launch"); } - if ((group_dim & 31u) == 0u && group_dim <= 4096u && (rank % 64u) == 0u) { - const unsigned rows_per_block = 64u; - grouped_q8_0_a_f32_sharedx_rows_w32_2row_kernel<<< - (unsigned)((low_dim + rows_per_block - 1u) / rows_per_block), - 1024u, - (size_t)group_dim * sizeof(float)>>>( + if (!cuda_q8_prequant_decode_enabled()) { + if ((group_dim & 31u) == 0u && group_dim <= 4096u && + (rank % 64u) == 0u) { + const unsigned rows_per_block = 64u; + grouped_q8_0_a_f32_sharedx_rows_w32_2row_kernel<<< + (unsigned)((low_dim + rows_per_block - 1u) / + rows_per_block), + 1024u, + (size_t)group_dim * sizeof(float)>>>( + (float *)low->ptr, + out_a, + (const float *)heads->ptr, + n_groups, + (uint32_t)blocks_a, + rank, + blocks_a * 34u); + return cuda_ok(cudaGetLastError(), + "attention_output_low_q8 f32 sharedx launch"); + } + grouped_q8_0_a_f32_warp8_kernel<<< + ((unsigned)low_dim + 7u) / 8u, 256>>>( (float *)low->ptr, out_a, (const float *)heads->ptr, - n_groups, - (uint32_t)blocks_a, + group_dim, rank, - blocks_a * 34u); - return cuda_ok(cudaGetLastError(), "attention_output_low_q8 f32 sharedx launch"); + n_groups, + blocks_a); + return cuda_ok(cudaGetLastError(), + "attention_output_low_q8 f32 launch"); } - grouped_q8_0_a_f32_warp8_kernel<<<((unsigned)low_dim + 7u) / 8u, 256>>>( + + const uint64_t x_rows = (uint64_t)n_groups; + const uint64_t xq_bytes = x_rows * blocks_a * 32u; + const uint64_t scale_offset = (xq_bytes + 15u) & ~15ull; + const uint64_t tmp_bytes = + scale_offset + x_rows * blocks_a * sizeof(float); + void *tmp = cuda_tmp_alloc(tmp_bytes, + "attention output low q8 prequant"); + if (!tmp) return 0; + int8_t *xq = (int8_t *)tmp; + float *xscale = (float *)((char *)tmp + scale_offset); + const ds4_rocm_runtime_config *cfg = cuda_runtime_config(); + const int use_dp4a = 1; + dim3 qgrid((unsigned)blocks_a, (unsigned)x_rows, 1); + quantize_q8_0_f32_kernel<<>>( + xq, + xscale, + (const float *)heads->ptr, + group_dim, + blocks_a); + if (!cuda_ok(cudaGetLastError(), + "attention_output_low_q8 prequant launch")) { + return 0; + } + const uint32_t rows_per_block = cfg->attn_out_low_decode_rpb; + dim3 grid_a( + ((unsigned)low_dim + rows_per_block - 1u) / rows_per_block, + 1, + 1); + grouped_q8_0_a_preq_warp8_kernel<<>>( (float *)low->ptr, out_a, - (const float *)heads->ptr, + xq, + xscale, group_dim, rank, n_groups, - blocks_a); - return cuda_ok(cudaGetLastError(), "attention_output_low_q8 f32 launch"); + 1, + blocks_a, + use_dp4a); + return cuda_ok(cudaGetLastError(), + "attention_output_low_q8 launch"); } diff --git a/rocm/ds4_rocm_current_api_compat.cuh b/rocm/ds4_rocm_current_api_compat.cuh index e75cc5ae4a..1c072c02c0 100644 --- a/rocm/ds4_rocm_current_api_compat.cuh +++ b/rocm/ds4_rocm_current_api_compat.cuh @@ -126,7 +126,7 @@ extern "C" void ds4_gpu_set_ssd_streaming(bool enabled) { } extern "C" void ds4_gpu_set_glm_model(bool enabled) { - (void)enabled; + g_glm_model = enabled ? 1 : 0; } extern "C" void ds4_gpu_set_glm_streaming_prefill_full_layer(bool enabled) { diff --git a/rocm/ds4_rocm_matmul.cuh b/rocm/ds4_rocm_matmul.cuh index 8c05433b99..fc20f8b95b 100644 --- a/rocm/ds4_rocm_matmul.cuh +++ b/rocm/ds4_rocm_matmul.cuh @@ -329,7 +329,7 @@ static int cuda_matmul_q8_0_tensor_labeled(ds4_gpu_tensor *out, const void *mode } const char *wptr = cuda_model_range_ptr(model_map, weight_offset, weight_bytes, "q8_0"); if (!wptr) return 0; - if (n_tok == 1) { + if (n_tok == 1 && !cuda_q8_prequant_decode_enabled()) { const bool extended_sharedx = in_dim > 8192u && in_dim <= 16384u && @@ -478,10 +478,27 @@ static int cuda_matmul_q8_0_tensor_labeled(ds4_gpu_tensor *out, const void *mode if (!tmp) return 0; int8_t *xq = (int8_t *)tmp; float *xscale = (float *)((char *)tmp + scale_offset); + const ds4_rocm_runtime_config *cfg = cuda_runtime_config(); const int use_dp4a = 1; dim3 qgrid((unsigned)blocks, (unsigned)n_tok, 1); quantize_q8_0_f32_kernel<<>>(xq, xscale, (const float *)x->ptr, in_dim, blocks); if (!cuda_ok(cudaGetLastError(), "matmul_q8_0 quantize launch")) return 0; + if (n_tok == 1) { + const uint32_t rows_per_block = cfg->q8_decode_rpb; + matmul_q8_0_preq_rows_w32_kernel<<< + ((unsigned)out_dim + rows_per_block - 1u) / rows_per_block, + rows_per_block * 32u>>>( + (float *)out->ptr, + reinterpret_cast(wptr), + xq, + xscale, + in_dim, + out_dim, + blocks, + rows_per_block, + use_dp4a); + return cuda_ok(cudaGetLastError(), "matmul_q8_0 rows launch"); + } if (blocks <= 32u) { dim3 bgrid(((unsigned)out_dim + 7u) / 8u, (unsigned)n_tok, 1); matmul_q8_0_preq_batch_warp8_kernel<<>>( @@ -606,36 +623,68 @@ extern "C" int ds4_gpu_matmul_q8_0_pair_tensor( const char *w0 = cuda_model_range_ptr(model_map, weight0_offset, weight0_bytes, "q8_0_pair0"); const char *w1 = cuda_model_range_ptr(model_map, weight1_offset, weight1_bytes, "q8_0_pair1"); if (!w0 || !w1) return 0; - const uint64_t max_out = out0_dim > out1_dim ? out0_dim : out1_dim; - if ((in_dim & 31u) == 0u && in_dim <= 8192u) { - const unsigned rows_per_block = 32u; - const unsigned threads = rows_per_block * 32u; - matmul_q8_0_pair_f32_sharedx_warp_rows_w32_kernel<<< - (unsigned)((max_out + rows_per_block - 1u) / rows_per_block), - threads, - (size_t)in_dim * sizeof(float)>>>( + if (!cuda_q8_prequant_decode_enabled()) { + const uint64_t max_out = out0_dim > out1_dim ? out0_dim : out1_dim; + if ((in_dim & 31u) == 0u && in_dim <= 8192u) { + const unsigned rows_per_block = 32u; + const unsigned threads = rows_per_block * 32u; + matmul_q8_0_pair_f32_sharedx_warp_rows_w32_kernel<<< + (unsigned)((max_out + rows_per_block - 1u) / rows_per_block), + threads, + (size_t)in_dim * sizeof(float)>>>( + (float *)out0->ptr, + (float *)out1->ptr, + reinterpret_cast(w0), + reinterpret_cast(w1), + (const float *)x->ptr, + (uint32_t)blocks, + out0_dim, + out1_dim, + blocks * 34u); + return cuda_ok(cudaGetLastError(), "matmul_q8_0 pair f32 sharedx launch"); + } + matmul_q8_0_pair_f32_warp8_kernel<<<((unsigned)max_out + 7u) / 8u, 256>>>( (float *)out0->ptr, (float *)out1->ptr, reinterpret_cast(w0), reinterpret_cast(w1), (const float *)x->ptr, - (uint32_t)blocks, + in_dim, out0_dim, out1_dim, - blocks * 34u); - return cuda_ok(cudaGetLastError(), "matmul_q8_0 pair f32 sharedx launch"); + blocks); + return cuda_ok(cudaGetLastError(), "matmul_q8_0 pair f32 warp launch"); + } + + const uint64_t xq_bytes = blocks * 32u; + const uint64_t scale_offset = (xq_bytes + 15u) & ~15ull; + const uint64_t tmp_bytes = scale_offset + blocks * sizeof(float); + void *tmp = cuda_tmp_alloc(tmp_bytes, "q8_0 pair prequant"); + if (!tmp) return 0; + int8_t *xq = (int8_t *)tmp; + float *xscale = (float *)((char *)tmp + scale_offset); + const int use_dp4a = 1; + dim3 qgrid((unsigned)blocks, 1, 1); + quantize_q8_0_f32_kernel<<>>( + xq, xscale, (const float *)x->ptr, in_dim, blocks); + if (!cuda_ok(cudaGetLastError(), "matmul_q8_0 pair quantize launch")) { + return 0; } - matmul_q8_0_pair_f32_warp8_kernel<<<((unsigned)max_out + 7u) / 8u, 256>>>( + const uint64_t max_out = out0_dim > out1_dim ? out0_dim : out1_dim; + matmul_q8_0_pair_preq_warp8_kernel<<< + ((unsigned)max_out + 7u) / 8u, 256>>>( (float *)out0->ptr, (float *)out1->ptr, reinterpret_cast(w0), reinterpret_cast(w1), - (const float *)x->ptr, + xq, + xscale, in_dim, out0_dim, out1_dim, - blocks); - return cuda_ok(cudaGetLastError(), "matmul_q8_0 pair f32 warp launch"); + blocks, + use_dp4a); + return cuda_ok(cudaGetLastError(), "matmul_q8_0 pair warp launch"); } static int cuda_matmul_q8_0_hc_expand_tensor_labeled( @@ -674,13 +723,31 @@ static int cuda_matmul_q8_0_hc_expand_tensor_labeled( } const char *wptr = cuda_model_range_ptr(model_map, weight_offset, weight_bytes, label ? label : "q8_0_hc_expand"); if (!wptr) return 0; - if ((in_dim & 31u) == 0u && in_dim <= 8192u) { - const unsigned rows_per_block = 32u; - const unsigned threads = rows_per_block * 32u; - matmul_q8_0_hc_expand_f32_sharedx_warp_rows_w32_kernel<<< - (unsigned)((out_dim + rows_per_block - 1u) / rows_per_block), - threads, - (size_t)in_dim * sizeof(float)>>>( + if (!cuda_q8_prequant_decode_enabled()) { + if ((in_dim & 31u) == 0u && in_dim <= 8192u) { + const unsigned rows_per_block = 32u; + const unsigned threads = rows_per_block * 32u; + matmul_q8_0_hc_expand_f32_sharedx_warp_rows_w32_kernel<<< + (unsigned)((out_dim + rows_per_block - 1u) / rows_per_block), + threads, + (size_t)in_dim * sizeof(float)>>>( + (float *)out_hc->ptr, + (float *)block_out->ptr, + block_add ? (const float *)block_add->ptr : (const float *)block_out->ptr, + (const float *)residual_hc->ptr, + (const float *)split->ptr, + reinterpret_cast(wptr), + (const float *)x->ptr, + (uint32_t)blocks, + out_dim, + blocks * 34u, + n_embd, + n_hc, + block_add ? 1 : 0); + return cuda_ok(cudaGetLastError(), "matmul_q8_0_hc_expand f32 sharedx launch"); + } + matmul_q8_0_hc_expand_f32_warp8_kernel<<< + ((unsigned)out_dim + 7u) / 8u, 256>>>( (float *)out_hc->ptr, (float *)block_out->ptr, block_add ? (const float *)block_add->ptr : (const float *)block_out->ptr, @@ -688,29 +755,50 @@ static int cuda_matmul_q8_0_hc_expand_tensor_labeled( (const float *)split->ptr, reinterpret_cast(wptr), (const float *)x->ptr, - (uint32_t)blocks, + in_dim, out_dim, - blocks * 34u, n_embd, n_hc, + blocks, block_add ? 1 : 0); - return cuda_ok(cudaGetLastError(), "matmul_q8_0_hc_expand f32 sharedx launch"); + return cuda_ok(cudaGetLastError(), "matmul_q8_0_hc_expand f32 launch"); + } + + const uint64_t xq_bytes = blocks * 32u; + const uint64_t scale_offset = (xq_bytes + 15u) & ~15ull; + const uint64_t tmp_bytes = scale_offset + blocks * sizeof(float); + void *tmp = cuda_tmp_alloc(tmp_bytes, "q8_0 hc expand prequant"); + if (!tmp) return 0; + int8_t *xq = (int8_t *)tmp; + float *xscale = (float *)((char *)tmp + scale_offset); + const ds4_rocm_runtime_config *cfg = cuda_runtime_config(); + const int use_dp4a = 1; + quantize_q8_0_f32_kernel<<<(unsigned)blocks, 32>>>( + xq, xscale, (const float *)x->ptr, in_dim, blocks); + if (!cuda_ok(cudaGetLastError(), "matmul_q8_0_hc_expand quantize launch")) { + return 0; } - matmul_q8_0_hc_expand_f32_warp8_kernel<<<((unsigned)out_dim + 7u) / 8u, 256>>>( + const uint32_t rows_per_block = cfg->q8_hc_decode_rpb; + matmul_q8_0_hc_expand_preq_rows_w32_kernel<<< + ((unsigned)out_dim + rows_per_block - 1u) / rows_per_block, + rows_per_block * 32u>>>( (float *)out_hc->ptr, (float *)block_out->ptr, block_add ? (const float *)block_add->ptr : (const float *)block_out->ptr, (const float *)residual_hc->ptr, (const float *)split->ptr, reinterpret_cast(wptr), - (const float *)x->ptr, + xq, + xscale, in_dim, out_dim, n_embd, n_hc, blocks, - block_add ? 1 : 0); - return cuda_ok(cudaGetLastError(), "matmul_q8_0_hc_expand f32 launch"); + rows_per_block, + block_add ? 1 : 0, + use_dp4a); + return cuda_ok(cudaGetLastError(), "matmul_q8_0_hc_expand rows launch"); } extern "C" int ds4_gpu_matmul_f16_tensor(ds4_gpu_tensor *out, const void *model_map, uint64_t model_size, uint64_t weight_offset, uint64_t in_dim, uint64_t out_dim, const ds4_gpu_tensor *x, uint64_t n_tok) { diff --git a/rocm/ds4_rocm_runtime.cuh b/rocm/ds4_rocm_runtime.cuh index c4b8f338d9..92e78a4445 100644 --- a/rocm/ds4_rocm_runtime.cuh +++ b/rocm/ds4_rocm_runtime.cuh @@ -21,6 +21,7 @@ static int g_cublas_ready; #include "ds4_rocm_hipblaslt.cuh" #endif static int g_quality_mode; +static int g_glm_model; enum { DS4_ROCM_N_EXPERT = 256u, @@ -4748,6 +4749,7 @@ static uint32_t cuda_rows_per_block_env_or_default(const char *name, uint32_t de struct ds4_rocm_runtime_config { int initialized; + int q8_prequant_decode; int disable_splitk_attn_out_low; int disable_shared_gate_up_fused_w32; int attention_output_cublas_all; @@ -4756,6 +4758,9 @@ struct ds4_rocm_runtime_config { int glm_grouped_qk_low; int q8_decode_sharedx_64k; int graph_dump; + uint32_t q8_decode_rpb; + uint32_t q8_hc_decode_rpb; + uint32_t attn_out_low_decode_rpb; uint32_t moe_decode_rpb; uint32_t moe_decode_gate_rpb; uint32_t moe_decode_down_rpb; @@ -4766,6 +4771,18 @@ static ds4_rocm_runtime_config g_rocm_cfg; static const ds4_rocm_runtime_config *cuda_runtime_config(void) { if (!g_rocm_cfg.initialized) { + const char *dsv4_prequant_env = + getenv("DS4_ROCM_DSV4_PREQUANT_DECODE"); + /* + * DeepSeek V4 used the prequantized Q8 decode kernels before ROCm + * GLM support landed. They are substantially faster on gfx1151. + * Keep GLM and --quality on the current full-F32 activation path. + * An explicit =0 remains available as a diagnostic rollback. + */ + g_rocm_cfg.q8_prequant_decode = + !g_quality_mode && + (dsv4_prequant_env == NULL || + cuda_env_present(dsv4_prequant_env)); g_rocm_cfg.disable_splitk_attn_out_low = !g_quality_mode; g_rocm_cfg.disable_shared_gate_up_fused_w32 = !g_quality_mode; g_rocm_cfg.attention_output_cublas_all = !g_quality_mode; @@ -4805,6 +4822,9 @@ static const ds4_rocm_runtime_config *cuda_runtime_config(void) { cuda_env_present(getenv("DS4_ROCM_GRAPH_DUMP_NONINVASIVE")); g_rocm_cfg.graph_dump = graph_dump_requested && !graph_dump_noninvasive; + g_rocm_cfg.q8_decode_rpb = 1u; + g_rocm_cfg.q8_hc_decode_rpb = 16u; + g_rocm_cfg.attn_out_low_decode_rpb = 32u; const char *moe_decode_rpb_env = getenv("DS4_ROCM_MOE_DECODE_RPB"); const int moe_decode_rpb_env_present = moe_decode_rpb_env != NULL && moe_decode_rpb_env[0] != '\0'; @@ -4830,6 +4850,10 @@ static const ds4_rocm_runtime_config *cuda_runtime_config(void) { return &g_rocm_cfg; } +static bool cuda_q8_prequant_decode_enabled(void) { + return !g_glm_model && cuda_runtime_config()->q8_prequant_decode; +} + static uint64_t cuda_q8_f16_cache_limit_bytes(void) { if (!g_ssd_streaming_mode) return UINT64_MAX; const char *env = getenv("DS4_ROCM_STREAM_Q8_F16_CACHE_GB"); From 9e1988bcd15a857390f0c045adb15f109a0ab101 Mon Sep 17 00:00:00 2001 From: Donato Capitella Date: Tue, 28 Jul 2026 19:22:22 +0100 Subject: [PATCH 2/5] Fix ROCm DeepSeek Q4 SSD expert staging --- ds4.c | 24 +++++++++++++++++++++--- 1 file changed, 21 insertions(+), 3 deletions(-) diff --git a/ds4.c b/ds4.c index 8de85ca318..8be06ec25c 100644 --- a/ds4.c +++ b/ds4.c @@ -20072,7 +20072,25 @@ static bool metal_graph_use_streaming_iq2_cpu_router(void) { getenv("DS4_METAL_DISABLE_STREAMING_IQ2_CPU_ROUTER") == NULL; } -static bool metal_graph_use_q4_selected_shared_overlap(void) { +static bool metal_graph_use_q4_selected_shared_overlap( + const ds4_gpu_graph *g) { +#ifdef DS4_ROCM_BUILD + /* + * ROCm SSD streaming maps only the selected routed experts, not the full + * Q4 expert table. Stage that compact selection while the shared expert + * runs so decode reaches the existing pointer-backed MoE path. + * + * GLM owns a separate routed-MoE graph and must retain its validated + * selected-expert policy. + */ + if (g && + g->ssd_streaming && + DS4_MODEL_FAMILY != DS4_MODEL_FAMILY_GLM_DSA) { + return true; + } +#else + (void)g; +#endif static int cache = -1; return metal_graph_env_flag("DS4_METAL_Q4_SELECTED_OVERLAP_SHARED", &cache); } @@ -23257,7 +23275,7 @@ static bool metal_graph_encode_decode_layer_phase( * needs the unfused gate/up/swiglu/down sequence. */ const bool tp_split_shared = g->tp_world == 2; const bool q4_selected_shared_overlap = - metal_graph_use_q4_selected_shared_overlap() && + metal_graph_use_q4_selected_shared_overlap(g) && metal_graph_decode_q4_selected_slots_expected(g, layer, layer->ffn_gate_exps->bytes, @@ -59986,7 +60004,7 @@ static bool metal_graph_native_session_batch_shared_supported( !items || count < 2 || !e || e->tp.active || e->support_kind != DS4_SUPPORT_NONE || metal_graph_use_reference_shared_down_hc() || - metal_graph_use_q4_selected_shared_overlap() || + metal_graph_use_q4_selected_shared_overlap(NULL) || metal_graph_use_pro_q4_cpu_router() || getenv("DS4_METAL_ENABLE_Q8_DECODE_EXACT_VIEWS") != NULL) { return false; From 6882a0f2f42f7e865a888f4fed426bdaed419776 Mon Sep 17 00:00:00 2001 From: Donato Capitella Date: Tue, 28 Jul 2026 21:44:11 +0100 Subject: [PATCH 3/5] Tighten ROCm DeepSeek large model arenas --- rocm/ds4_rocm_runtime.cuh | 13 ++++++++++++- 1 file changed, 12 insertions(+), 1 deletion(-) diff --git a/rocm/ds4_rocm_runtime.cuh b/rocm/ds4_rocm_runtime.cuh index 92e78a4445..8eeca5c2ec 100644 --- a/rocm/ds4_rocm_runtime.cuh +++ b/rocm/ds4_rocm_runtime.cuh @@ -5530,7 +5530,18 @@ static int cuda_stream_model_cache_prepare_memory( } static uint64_t cuda_model_arena_chunk_bytes(uint64_t need) { - uint64_t bytes = 1792ull * 1048576ull; + const uint64_t default_bytes = 1792ull * 1048576ull; + /* + * Two allocations larger than half the default arena can never share it. + * Allocate those spans tightly for DeepSeek instead of stranding the + * remainder of every arena. This matters for Q4 expert spans (1152 MiB), + * where the default policy otherwise wastes 640 MiB per tensor span. + * + * Keep GLM on its existing allocation policy. + */ + if (!g_glm_model && need > default_bytes / 2u) return need; + + uint64_t bytes = default_bytes; if (bytes < need) { const uint64_t align = 256ull * 1048576ull; bytes = (need + align - 1u) & ~(align - 1u); From 2f49b27edd797c46bbef965bf7d4e9b2870abcc5 Mon Sep 17 00:00:00 2001 From: Donato Capitella Date: Wed, 29 Jul 2026 09:27:35 +0100 Subject: [PATCH 4/5] Fix ROCm IQ2 SSD prefill deadlock --- ds4.c | 11 +++++++++++ ds4_gpu.h | 1 + rocm/ds4_rocm_current_api_compat.cuh | 4 ++++ 3 files changed, 16 insertions(+) diff --git a/ds4.c b/ds4.c index 8be06ec25c..f17cadad06 100644 --- a/ds4.c +++ b/ds4.c @@ -33593,6 +33593,17 @@ static bool metal_graph_prefill_layer_major( if (show_progress) fputc('\n', stderr); metal_graph_stream_prefill_selected_profile_summary(g); #ifdef DS4_ROCM_BUILD + /* + * The final layer can use its fully mapped prefill table after the + * asynchronous selected-expert loader has already queued a resident-cache + * batch. Drain that completed batch before hotlist seeding starts another + * read-pool job set, otherwise both callers wait forever for ownership of + * the single ROCm streaming read pool. + */ + if (g->ssd_streaming && + !ds4_gpu_stream_expert_cache_finish_pending_batch()) { + return false; + } (void)ds4_gpu_stream_expert_cache_release_layer_cache(); if (g->ssd_streaming) ds4_gpu_release_q8_f16_cache(); #endif diff --git a/ds4_gpu.h b/ds4_gpu.h index 2000bba8bd..c2464df490 100644 --- a/ds4_gpu.h +++ b/ds4_gpu.h @@ -200,6 +200,7 @@ int ds4_gpu_stream_expert_cache_seed_from_layer_selected( uint32_t n_tokens, uint32_t n_seed_tokens, uint32_t n_selected); +int ds4_gpu_stream_expert_cache_finish_pending_batch(void); int ds4_gpu_stream_expert_cache_release_layer_cache(void); #endif int ds4_gpu_stream_expert_cache_seed_experts( diff --git a/rocm/ds4_rocm_current_api_compat.cuh b/rocm/ds4_rocm_current_api_compat.cuh index 1c072c02c0..d24668ce7e 100644 --- a/rocm/ds4_rocm_current_api_compat.cuh +++ b/rocm/ds4_rocm_current_api_compat.cuh @@ -281,6 +281,10 @@ extern "C" int ds4_gpu_stream_expert_cache_seed_from_layer_selected( table->down_expert_bytes); } +extern "C" int ds4_gpu_stream_expert_cache_finish_pending_batch(void) { + return cuda_stream_batch_selected_finish_pending_missing(); +} + extern "C" int ds4_gpu_stream_expert_cache_release_layer_cache(void) { cuda_stream_layer_expert_cache_release(); return 1; From d250a7c07c6beb753e9b0a33951d8c00d6ef30ee Mon Sep 17 00:00:00 2001 From: Donato Capitella Date: Sat, 1 Aug 2026 08:29:22 +0100 Subject: [PATCH 5/5] Optimize ROCm two-token MTP verification --- ds4.c | 27 ++++++++++++++++++++------- 1 file changed, 20 insertions(+), 7 deletions(-) diff --git a/ds4.c b/ds4.c index f17cadad06..a2b58186d7 100644 --- a/ds4.c +++ b/ds4.c @@ -64435,15 +64435,28 @@ int ds4_session_eval_speculative_argmax(ds4_session *s, int first_token, } /* - * The useful N=2 verifier is the tiny batch path: it verifies two target - * positions in one layer-major pass and commits prefix-1 directly on a - * partial accept. Like the rest of the non-quality Metal path, it may pick - * a different greedy token when batched reductions perturb nearly-tied - * logits. --quality / DS4_MTP_STRICT selects the exact decode verifier, - * which preserves the one-token target stream but is not a speed win. + * Metal normally benefits from the tiny-batch verifier: it checks two + * target positions in one layer-major pass and commits prefix-1 directly + * on a partial accept. Like the rest of the non-quality Metal path, it may + * pick a different greedy token when batched reductions perturb nearly + * tied logits. + * + * ROCm is different. Its DeepSeek one-token graph uses the prequantized Q8 + * decode kernels, while the generic N=2 batch graph uses full-F32 + * activations and loses the paired and HC-expand decode fusions. Running + * the exact two-row verifier reuses those Q8 kernels and is faster. + * DS4_MTP_BATCH_VERIFY remains an explicit diagnostic rollback to the + * generic batch verifier. */ +#ifdef DS4_ROCM_BUILD + const bool prefer_decode2_exact = true; +#else + const bool prefer_decode2_exact = false; +#endif const bool use_decode2_exact = - draft_n == 2 && strict_mtp && getenv("DS4_MTP_BATCH_VERIFY") == NULL; + draft_n == 2 && + (strict_mtp || prefer_decode2_exact) && + getenv("DS4_MTP_BATCH_VERIFY") == NULL; if (use_decode2_exact) { ds4_spec_frontier frontier; memset(&frontier, 0, sizeof(frontier));