ROCm: restore DeepSeek V4 Flash decode performance and Q4 SSD streaming

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,
@@ -33575,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
@@ -55958,6 +55987,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) {
@@ -59977,7 +60015,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;
@@ -64397,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));

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(

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<<<qgrid, 32>>>(
+            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<<<grid_a,
+                                       rows_per_block * 32u>>>(
             (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");
 }

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) {
@@ -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;