From 24f2c83d5007c5f866106f85d4b387c3c5344208 Mon Sep 17 00:00:00 2001 From: Donato Capitella Date: Wed, 22 Jul 2026 02:33:15 +0100 Subject: [PATCH 01/33] Add ROCm GLM routed IQ2 down support --- rocm/ds4_rocm_moe.cuh | 78 ++++++++++++++++++++++++++++++ rocm/ds4_rocm_moe_launch.cuh | 92 ++++++++++++++++++++++++++++-------- 2 files changed, 151 insertions(+), 19 deletions(-) diff --git a/rocm/ds4_rocm_moe.cuh b/rocm/ds4_rocm_moe.cuh index 74fbbbfed..df74b736d 100644 --- a/rocm/ds4_rocm_moe.cuh +++ b/rocm/ds4_rocm_moe.cuh @@ -2144,6 +2144,84 @@ __global__ static void moe_down_sum6_qwarp32_ptrs_batch_kernel( if (lane == 0) out[(uint64_t)tok * out_dim + row] = total; } +/* GLM's routed-IQ2 layout uses IQ2_XXS for the down projection as well as + * gate/up. One 8-lane subgroup owns an output row and accumulates all routed + * experts for a token. This mirrors the Q2_K direct-sum path above while + * using the existing IQ2_XXS x Q8_K dot primitive. */ +__global__ static void moe_down_iq2_sum_qwarp32_batch_kernel( + float *out, + const char *down_base, + const cuda_block_q8_K *midq, + const int32_t *selected, + uint64_t down_expert_bytes, + uint64_t down_row_bytes, + uint32_t midq_blocks, + uint32_t out_dim, + uint32_t n_expert, + uint32_t n_tokens) { + const uint32_t lane = threadIdx.x & 7u; + const uint32_t row = blockIdx.x * 32u + (threadIdx.x >> 3u); + const uint32_t tok = blockIdx.y; + if (row >= out_dim || tok >= n_tokens) return; + float total = 0.0f; +#pragma unroll + for (uint32_t slot = 0; slot < DS4_ROCM_N_EXPERT_USED; slot++) { + if (slot >= n_expert) continue; + int32_t expert_i = selected[(uint64_t)tok * n_expert + slot]; + if (expert_i < 0) expert_i = 0; + const cuda_block_iq2_xxs *wr = + (const cuda_block_iq2_xxs *)(down_base + + (uint64_t)(uint32_t)expert_i * down_expert_bytes + + (uint64_t)row * down_row_bytes); + const cuda_block_q8_K *xq = + midq + ((uint64_t)tok * n_expert + slot) * midq_blocks; + float acc = 0.0f; + for (uint32_t b = lane; b < midq_blocks; b += 8u) { + acc += dev_dot_iq2_xxs_q8_K_block(wr + b, xq + b); + } + acc = quarter_warp_sum_f32(acc, lane); + if (lane == 0u) total += acc; + } + if (lane == 0u) out[(uint64_t)tok * out_dim + row] = total; +} + +__global__ static void moe_down_iq2_sum_qwarp32_ptrs_batch_kernel( + float *out, + const char * const *down_slots, + const cuda_block_q8_K *midq, + const int32_t *selected, + uint64_t down_row_bytes, + uint32_t midq_blocks, + uint32_t out_dim, + uint32_t n_expert, + uint32_t n_tokens) { + const uint32_t lane = threadIdx.x & 7u; + const uint32_t row = blockIdx.x * 32u + (threadIdx.x >> 3u); + const uint32_t tok = blockIdx.y; + if (row >= out_dim || tok >= n_tokens) return; + float total = 0.0f; +#pragma unroll + for (uint32_t slot = 0; slot < DS4_ROCM_N_EXPERT_USED; slot++) { + if (slot >= n_expert) continue; + int32_t compact_i = selected[(uint64_t)tok * n_expert + slot]; + if (compact_i < 0) compact_i = 0; + const char *down = down_slots[(uint32_t)compact_i]; + if (!down) continue; + const cuda_block_iq2_xxs *wr = + (const cuda_block_iq2_xxs *)(down + + (uint64_t)row * down_row_bytes); + const cuda_block_q8_K *xq = + midq + ((uint64_t)tok * n_expert + slot) * midq_blocks; + float acc = 0.0f; + for (uint32_t b = lane; b < midq_blocks; b += 8u) { + acc += dev_dot_iq2_xxs_q8_K_block(wr + b, xq + b); + } + acc = quarter_warp_sum_f32(acc, lane); + if (lane == 0u) total += acc; + } + if (lane == 0u) out[(uint64_t)tok * out_dim + row] = total; +} + __global__ static void moe_down_q4K_sum6_qwarp32_kernel( float *out, const char *down_base, diff --git a/rocm/ds4_rocm_moe_launch.cuh b/rocm/ds4_rocm_moe_launch.cuh index c8ce8dfc3..3a3d417eb 100644 --- a/rocm/ds4_rocm_moe_launch.cuh +++ b/rocm/ds4_rocm_moe_launch.cuh @@ -441,6 +441,7 @@ static int routed_moe_q2_float_down_launch( typedef struct { int q4k_path; int iq2_path; + int iq2_iq2_path; int q2k_path; uint64_t gate_bytes; uint64_t down_bytes; @@ -490,8 +491,10 @@ static int routed_moe_build_plan( } plan->q4k_path = (gate_type == 12u && down_type == 12u); plan->iq2_path = (gate_type == 16u && down_type == 10u); + plan->iq2_iq2_path = (gate_type == 16u && down_type == 16u); plan->q2k_path = (gate_type == 10u && down_type == 10u); - if (!plan->q4k_path && !plan->iq2_path && !plan->q2k_path) return 0; + if (!plan->q4k_path && !plan->iq2_path && + !plan->iq2_iq2_path && !plan->q2k_path) return 0; if (!cuda_u64_mul_checked(n_total_expert, gate_expert_bytes, &plan->gate_bytes) || !cuda_u64_mul_checked(n_total_expert, down_expert_bytes, &plan->down_bytes) || !cuda_model_range_fits(model_size, gate_offset, plan->gate_bytes) || @@ -553,6 +556,8 @@ static int routed_moe_launch( } const int q4k_path = plan.q4k_path; const int iq2_path = plan.iq2_path; + const int iq2_iq2_path = plan.iq2_iq2_path; + const int iq2_gate_path = iq2_path || iq2_iq2_path; const int q2k_path = plan.q2k_path; const uint64_t gate_bytes = plan.gate_bytes; const uint64_t down_bytes = plan.down_bytes; @@ -604,7 +609,7 @@ static int routed_moe_launch( !stream_full_layer && !full_table_cached && n_tokens > 1u && - (iq2_path || q2k_path) && + (iq2_gate_path || q2k_path) && n_expert <= DS4_ROCM_N_EXPERT_USED && cuda_stream_batch_selected_apply_split(model_map, layer_index, @@ -631,7 +636,7 @@ static int routed_moe_launch( !full_table_cached && !batch_stream_split_selected && n_tokens > 1u && - (iq2_path || q2k_path) && + (iq2_gate_path || q2k_path) && n_expert <= DS4_ROCM_N_EXPERT_USED && cuda_stream_batch_selected_prepare(model_map, model_size, @@ -848,17 +853,33 @@ static int routed_moe_launch( } if (ok) { dim3 dgrid((out_dim + 31u) / 32u, n_tokens, 1); - moe_down_sum6_qwarp32_ptrs_batch_kernel<<>>( - (float *)out->ptr, - down_slot_ptrs, - midq, - (const int32_t *)selected_exec->ptr, - down_row_bytes, - midq_blocks, - out_dim, - n_expert, - n_tokens); - ok = cuda_ok(cudaGetLastError(), "routed_moe streaming batch down launch"); + if (iq2_iq2_path) { + moe_down_iq2_sum_qwarp32_ptrs_batch_kernel<<>>( + (float *)out->ptr, + down_slot_ptrs, + midq, + (const int32_t *)selected_exec->ptr, + down_row_bytes, + midq_blocks, + out_dim, + n_expert, + n_tokens); + ok = cuda_ok(cudaGetLastError(), + "routed_moe streaming batch iq2 down launch"); + } else { + moe_down_sum6_qwarp32_ptrs_batch_kernel<<>>( + (float *)out->ptr, + down_slot_ptrs, + midq, + (const int32_t *)selected_exec->ptr, + down_row_bytes, + midq_blocks, + out_dim, + n_expert, + n_tokens); + ok = cuda_ok(cudaGetLastError(), + "routed_moe streaming batch down launch"); + } } if (ok) ok = cuda_stream_batch_selected_mark_inflight(); return ok; @@ -958,7 +979,7 @@ static int routed_moe_launch( const uint32_t iq2_down_hot_threshold = 8u; uint32_t h_iq2_gate_hot[DS4_ROCM_MAX_N_EXPERT] = {0}; const uint32_t use_iq2_gate_wmma = - ok && iq2_path && n_tokens > 1u && n_expert == 6u && !write_gate_up && + ok && iq2_gate_path && n_tokens > 1u && n_expert == 6u && !write_gate_up && sorted_pairs && sorted_offsets && sorted_counts && tile_experts && iq2_gate_hot_dev && use_expert_tiles && (expert_in_dim % 16u) == 0u && (expert_mid_dim % 16u) == 0u && !g_quality_mode; @@ -999,7 +1020,7 @@ static int routed_moe_launch( int split_gateup_done = 0; if (ok && split_selected) { const int split_supported = - iq2_path && + iq2_gate_path && n_tokens == 1u && n_expert <= DS4_ROCM_N_EXPERT_USED && !q4k_path && @@ -1348,8 +1369,38 @@ static int routed_moe_launch( q8_K_quantize_kernel<<>>(midq, (const float *)mid->ptr, expert_mid_dim, pair_count); ok = cuda_ok(cudaGetLastError(), "routed_moe mid quantize launch"); } + int direct_iq2_down_done = 0; + if (ok && iq2_iq2_path) { + dim3 dgrid((out_dim + 31u) / 32u, n_tokens, 1); + if (split_gateup_done) { + moe_down_iq2_sum_qwarp32_ptrs_batch_kernel<<>>( + (float *)out->ptr, + down_slot_ptrs, + midq, + (const int32_t *)selected_exec->ptr, + down_row_bytes, + midq_blocks, + out_dim, + n_expert, + n_tokens); + } else { + moe_down_iq2_sum_qwarp32_batch_kernel<<>>( + (float *)out->ptr, + down_w, + midq, + (const int32_t *)selected_exec->ptr, + down_expert_bytes, + down_row_bytes, + midq_blocks, + out_dim, + n_expert, + n_tokens); + } + ok = cuda_ok(cudaGetLastError(), "routed_moe iq2 down launch"); + direct_iq2_down_done = ok; + } int split_ptr_down_done = 0; - if (ok && split_gateup_done) { + if (ok && !direct_iq2_down_done && split_gateup_done) { moe_down_sum6_qwarp32_ptrs_kernel<<<(out_dim + 31u) / 32u, 256>>>( (float *)out->ptr, down_slot_ptrs, @@ -1362,7 +1413,9 @@ static int routed_moe_launch( split_ptr_down_done = ok; } if (ok) { - if (split_ptr_down_done) { + if (direct_iq2_down_done) { + /* The IQ2 direct-sum kernel writes final token rows. */ + } 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( @@ -1544,7 +1597,8 @@ static int routed_moe_launch( ok = cuda_ok(cudaGetLastError(), "routed_moe down launch"); } } - if (ok && !use_atomic_down && !use_direct_down_sum6 && !use_iq2_q2_float_down) { + if (ok && !direct_iq2_down_done && !use_atomic_down && + !use_direct_down_sum6 && !use_iq2_q2_float_down) { uint64_t n = (uint64_t)n_tokens * out_dim; moe_sum_kernel<<<(n + 255) / 256, 256>>>((float *)out->ptr, (const float *)down->ptr, out_dim, n_expert, n_tokens); ok = cuda_ok(cudaGetLastError(), "routed_moe sum launch"); From d7a89093fa6f26c7cf7ec54266c1addb79b9962e Mon Sep 17 00:00:00 2001 From: Donato Capitella Date: Wed, 22 Jul 2026 03:40:37 +0100 Subject: [PATCH 02/33] Fix ROCm distributed GLM slice residency --- ds4.c | 423 ++++++++++++++++++++++++++++++-------- ds4_rocm.cu | 1 + rocm/ds4_rocm_runtime.cuh | 147 +++++++++---- 3 files changed, 452 insertions(+), 119 deletions(-) diff --git a/ds4.c b/ds4.c index 876cfe27f..e9f0bf82d 100644 --- a/ds4.c +++ b/ds4.c @@ -4472,40 +4472,66 @@ static bool ds4_streaming_routed_expert_bytes( if (per_expert_bytes_out) *per_expert_bytes_out = 0; if (!weights || !per_expert_bytes_out) return false; + /* Mixed-precision models can put an outlier quant at the first routed + * layer owned by a distributed slice. Choosing that first layer as the + * slab class makes every ordinary layer bypass the cache. Use the most + * common local size class instead (ties retain the earliest class). */ + uint64_t best_bytes = 0; + uint32_t best_count = 0; for (uint32_t il = 0; il < DS4_N_LAYER; il++) { - if (streaming_layer_routed_expert_bytes(&weights->layer[il], - per_expert_bytes_out)) { - return true; + uint64_t candidate = 0; + if (!streaming_layer_routed_expert_bytes(&weights->layer[il], + &candidate)) { + continue; + } + uint32_t count = 0; + for (uint32_t jl = 0; jl < DS4_N_LAYER; jl++) { + uint64_t bytes = 0; + if (streaming_layer_routed_expert_bytes(&weights->layer[jl], + &bytes) && + bytes == candidate) { + count++; + } + } + if (count > best_count) { + best_bytes = candidate; + best_count = count; } } - return false; + if (best_count == 0) return false; + *per_expert_bytes_out = best_bytes; + return true; } enum { DS4_STREAMING_PREFILL_HEADROOM_LAYERS = 2 }; -static bool ds4_streaming_max_routed_layer_bytes( +static bool ds4_streaming_cacheable_expert_count( const ds4_weights *weights, - uint64_t *layer_bytes_out) { - if (layer_bytes_out) *layer_bytes_out = 0; - if (!weights || !layer_bytes_out || DS4_N_EXPERT == 0) return false; + uint64_t *experts_out, + uint32_t *layers_out) { + if (experts_out) *experts_out = 0; + if (layers_out) *layers_out = 0; + if (!weights || !experts_out || DS4_N_EXPERT == 0) return false; - uint64_t max_bytes = 0; + uint64_t slab_bytes = 0; + if (!ds4_streaming_routed_expert_bytes(weights, &slab_bytes)) return false; + + uint32_t layers = 0; for (uint32_t il = 0; il < DS4_N_LAYER; il++) { uint64_t per_expert_bytes = 0; if (!streaming_layer_routed_expert_bytes(&weights->layer[il], &per_expert_bytes)) { continue; } - if (per_expert_bytes > UINT64_MAX / (uint64_t)DS4_N_EXPERT) { - return false; - } - const uint64_t layer_bytes = - per_expert_bytes * (uint64_t)DS4_N_EXPERT; - if (layer_bytes > max_bytes) max_bytes = layer_bytes; + if (per_expert_bytes == slab_bytes) layers++; } - if (max_bytes == 0) return false; - *layer_bytes_out = max_bytes; + if (layers == 0 || + (uint64_t)layers > UINT64_MAX / (uint64_t)DS4_N_EXPERT) { + return false; + } + *experts_out = (uint64_t)layers * (uint64_t)DS4_N_EXPERT; + if (layers_out) *layers_out = layers; return true; } @@ -4515,27 +4541,38 @@ static bool ds4_streaming_prefill_headroom_bytes( if (bytes_out) *bytes_out = 0; if (!weights || !bytes_out) return false; - uint64_t layer_bytes = 0; - if (!ds4_streaming_max_routed_layer_bytes(weights, &layer_bytes)) { + uint64_t per_expert_bytes = 0; + uint64_t cacheable_experts = 0; + uint32_t cacheable_layers = 0; + if (!ds4_streaming_routed_expert_bytes(weights, &per_expert_bytes) || + !ds4_streaming_cacheable_expert_count(weights, + &cacheable_experts, + &cacheable_layers)) { return false; } - if (layer_bytes > - UINT64_MAX / (uint64_t)DS4_STREAMING_PREFILL_HEADROOM_LAYERS) { + (void)cacheable_experts; + const uint32_t reserve_layers = + cacheable_layers < DS4_STREAMING_PREFILL_HEADROOM_LAYERS ? + cacheable_layers : DS4_STREAMING_PREFILL_HEADROOM_LAYERS; + if (per_expert_bytes > UINT64_MAX / (uint64_t)DS4_N_EXPERT) { return false; } + const uint64_t layer_bytes = + per_expert_bytes * (uint64_t)DS4_N_EXPERT; + if (reserve_layers != 0 && + layer_bytes > UINT64_MAX / (uint64_t)reserve_layers) return false; - *bytes_out = - layer_bytes * (uint64_t)DS4_STREAMING_PREFILL_HEADROOM_LAYERS; + *bytes_out = layer_bytes * (uint64_t)reserve_layers; return true; } /* * Mixed-precision ("boosted") GGUFs upcast a few layers' routed experts to a * bigger quant (e.g. Q4_K among IQ2 layers). The streaming expert cache is a - * single-size-class slab allocator sized from the FIRST routed layer, so those - * layers can never be served from it: they must read expert weights through the - * mapped-model views instead. A layer is "uniform" iff its per-expert bytes - * match the slab class. + * single-size-class slab allocator sized from the dominant local routed-layer + * size class, so other layers can never be served from it: they must read + * expert weights through the mapped-model views instead. A layer is "uniform" + * iff its per-expert bytes match the slab class. */ static DS4_MAYBE_UNUSED bool weights_streaming_layer_experts_uniform( const ds4_weights *w, @@ -5718,21 +5755,37 @@ static void config_validate_model(const ds4_model *m) { config_validate_deepseek4_model(m); } -static void weights_bind_output(ds4_weights *w, const ds4_model *m, bool required) { +static void weights_bind_output( + ds4_weights *w, + const ds4_model *m, + bool required, + bool optional) { if (DS4_MODEL_FAMILY == DS4_MODEL_FAMILY_GLM_DSA) { if (required) { w->output_norm = required_tensor(m, "output_norm.weight"); w->output = required_tensor(m, "output.weight"); + } else if (optional) { + w->output_norm = model_find_tensor(m, "output_norm.weight"); + w->output = model_find_tensor(m, "output.weight"); } - return; - } - - if (required) { + } else if (required) { w->output_hc_base = required_tensor(m, "output_hc_base.weight"); w->output_hc_fn = required_tensor(m, "output_hc_fn.weight"); w->output_hc_scale = required_tensor(m, "output_hc_scale.weight"); w->output_norm = required_tensor(m, "output_norm.weight"); w->output = required_tensor(m, "output.weight"); + } else if (optional) { + w->output_hc_base = model_find_tensor(m, "output_hc_base.weight"); + w->output_hc_fn = model_find_tensor(m, "output_hc_fn.weight"); + w->output_hc_scale = model_find_tensor(m, "output_hc_scale.weight"); + w->output_norm = model_find_tensor(m, "output_norm.weight"); + w->output = model_find_tensor(m, "output.weight"); + } + + if (optional && + weights_have_partial_output_head(w) && + !weights_have_output_head(w)) { + ds4_die("partial output head in GGUF"); } } @@ -5838,7 +5891,8 @@ static void weights_bind( bool load_slice, uint32_t load_layer_start, uint32_t load_layer_end, - bool require_output) { + bool require_output, + bool optional_output) { memset(w, 0, sizeof(*w)); uint32_t executable_layers = DS4_N_LAYER; @@ -5857,6 +5911,7 @@ static void weights_bind( require_token_embd = start == 0; } else { require_output = true; + optional_output = false; } if (require_token_embd) { @@ -5864,14 +5919,15 @@ static void weights_bind( } else { w->token_embd = model_find_tensor(m, "token_embd.weight"); } - weights_bind_output(w, m, require_output); + weights_bind_output(w, m, require_output, optional_output); for (uint32_t il = start; il <= end; il++) { weights_bind_layer(&w->layer[il], m, il); } /* GLM nextn/MTP block(s): excluded from the executable pass but bound * so the drafter can run them. Only when the full model is loaded. */ - if (DS4_MODEL_FAMILY == DS4_MODEL_FAMILY_GLM_DSA && + if (!load_slice && + DS4_MODEL_FAMILY == DS4_MODEL_FAMILY_GLM_DSA && start == 0 && end == executable_layers - 1u) { for (uint32_t il = executable_layers; il < DS4_N_LAYER; il++) { weights_bind_layer(&w->layer[il], m, il); @@ -6090,6 +6146,7 @@ static bool glm_stream_resident_decode_layer_supported( l->ffn_gate_exps->type == DS4_TENSOR_Q4_K); } +static uint32_t g_glm_streaming_full_resident_start; static uint32_t g_glm_streaming_full_resident_layers; static bool glm_stream_resident_decode_layer_enabled( @@ -6097,7 +6154,9 @@ static bool glm_stream_resident_decode_layer_enabled( uint32_t il) { if (!glm_stream_resident_decode_layer_supported(l, il)) return false; return g_glm_streaming_full_resident_layers != 0 && - il - DS4_N_LEADING_DENSE < g_glm_streaming_full_resident_layers; + il >= g_glm_streaming_full_resident_start && + il - g_glm_streaming_full_resident_start < + g_glm_streaming_full_resident_layers; } static bool glm_stream_expert_cache_addr_layout_supported( @@ -6396,7 +6455,12 @@ static DS4_MAYBE_UNUSED bool weights_streaming_non_routed_bytes( if (!w || !bytes_out) return false; ds4_model_map_span_vec spans; - if (!weights_model_map_decode_static_spans(w, true, true, &spans)) { + const bool include_token = + weights_layer_has_required(&w->layer[0], 0); + if (!weights_model_map_decode_static_spans(w, + include_token, + weights_have_output_head(w), + &spans)) { return false; } *bytes_out = model_map_span_vec_total_bytes(&spans); @@ -16795,6 +16859,7 @@ static bool metal_graph_alloc_raw_cap( g->prefill_cap = prefill_cap; uint32_t min_ratio = UINT32_MAX; for (uint32_t il = 0; il < DS4_N_LAYER; il++) { + if (!weights_layer_has_required(&weights->layer[il], il)) continue; const uint32_t ratio = ds4_layer_compress_ratio(il); if (ratio != 0 && ratio < min_ratio) min_ratio = ratio; } @@ -16806,6 +16871,10 @@ static bool metal_graph_alloc_raw_cap( if (g->attn_comp_stage_cap < 2u) g->attn_comp_stage_cap = 2u; } for (uint32_t il = 0; il < DS4_N_LAYER; il++) { + if (!weights_layer_has_required(&weights->layer[il], il)) { + g->layer_comp_cap[il] = 0; + continue; + } const uint32_t ratio = ds4_layer_compress_ratio(il); if (ratio == 0) { g->layer_comp_cap[il] = 0; @@ -16823,7 +16892,11 @@ static bool metal_graph_alloc_raw_cap( const uint64_t group_dim = (uint64_t)DS4_N_HEAD_DIM * (DS4_N_HEAD / DS4_N_OUT_GROUP); const uint64_t shared_dim = layer->ffn_gate_shexp->dim[1]; const uint64_t routed_mid_dim = layer->ffn_gate_exps->dim[1]; - const uint64_t vocab_dim = weights->output->dim[1]; + /* Distributed coordinators do not normally own the output head. The + * logits workspace still has a fixed model-vocabulary shape, while the + * actual head is encoded only on a node that bound its tensors. */ + const uint64_t vocab_dim = + weights->output ? weights->output->dim[1] : DS4_N_VOCAB; const uint64_t comp_width_max = 2ull * (DS4_N_HEAD_DIM > DS4_N_INDEXER_HEAD_DIM ? DS4_N_HEAD_DIM : DS4_N_INDEXER_HEAD_DIM); @@ -16909,6 +16982,10 @@ static bool metal_graph_alloc_raw_cap( } bool state_init_ok = true; for (uint32_t il = 0; il < DS4_N_LAYER; il++) { + /* A distributed process owns only its bound layer slice. Persistent + * KV state must follow that ownership just like the model tensors; + * allocating every model layer here defeats split-model residency. */ + if (!weights_layer_has_required(&weights->layer[il], il)) continue; /* per-layer Class L allocations land on the layer's * home tier. placement is NULL on single-tier / diagnostic paths * (all-tier-0); non-NULL on the engine path that opted into @@ -17204,6 +17281,7 @@ static bool metal_graph_alloc_raw_cap( bool layer_cache_ok = true; for (uint32_t il = 0; layer_cache_ok && il < DS4_N_LAYER; il++) { + if (!weights_layer_has_required(&weights->layer[il], il)) continue; layer_cache_ok = g->layer_raw_cache[il] != NULL; if (layer_cache_ok && g->cuda_tp_attn_cache_dup) { layer_cache_ok = g->layer_raw_cache_tp[il] != NULL; @@ -32579,6 +32657,7 @@ static bool metal_graph_reset_prefill_state(ds4_gpu_graph *g) { metal_graph_dspark_cache_reset(g); metal_graph_dspark_capture_invalidate(g); for (uint32_t il = 0; il < DS4_N_LAYER; il++) { + if (!g->layer_raw_cache[il]) continue; const uint32_t ratio = ds4_layer_compress_ratio(il); if (ratio == 0) continue; const uint32_t coff = ratio == 4 ? 2u : 1u; @@ -34422,12 +34501,6 @@ static uint32_t glm_graph_full_indexer_layer_count_range(uint32_t layer_start, return n; } -static uint32_t glm_graph_full_indexer_layer_count(uint32_t normal_layers) { - return normal_layers ? - glm_graph_full_indexer_layer_count_range(0, normal_layers - 1u) : - 0; -} - static uint64_t glm_graph_full_kv_cache_elem_bytes(void) { return sizeof(uint16_t); } @@ -34773,13 +34846,21 @@ static uint64_t glm_graph_workspace_bytes_for_cap( return bytes; } -static ds4_context_memory glm_graph_context_memory_estimate_for_compact_cap( +static ds4_context_memory glm_graph_context_memory_estimate_for_compact_cap_slice( uint32_t ctx, uint32_t work_ctx, uint32_t compact_cap, - bool ssd_streaming) { + bool ssd_streaming, + uint32_t layer_start, + uint32_t layer_end) { ds4_context_memory m = {0}; const uint32_t normal_layers = glm_graph_normal_layer_count(); + if (normal_layers == 0 || layer_start >= normal_layers || + layer_end < layer_start) { + return m; + } + if (layer_end >= normal_layers) layer_end = normal_layers - 1u; + const uint32_t layer_count = layer_end - layer_start + 1u; if (compact_cap > ctx) compact_cap = ctx; const bool expanded_kv = glm_graph_expanded_kv_cache_enabled(ssd_streaming); @@ -34793,7 +34874,7 @@ static ds4_context_memory glm_graph_context_memory_estimate_for_compact_cap( m.prefill_cap = batch_rows; m.raw_cap = expanded_kv ? work_ctx : 0; if (expanded_kv) { - m.raw_bytes = (uint64_t)normal_layers * + m.raw_bytes = (uint64_t)layer_count * work_ctx * ((uint64_t)DS4_N_HEAD * (DS4_N_KEY_MLA + DS4_N_VALUE_MLA)) * glm_graph_full_kv_cache_elem_bytes(); @@ -34806,14 +34887,34 @@ static ds4_context_memory glm_graph_context_memory_estimate_for_compact_cap( m.comp_cap = compact_cap; m.compressed_bytes = glm_graph_compact_cache_bytes_for_cap( - normal_layers, - glm_graph_full_indexer_layer_count(normal_layers), + layer_count, + glm_graph_full_indexer_layer_count_range(layer_start, + layer_end), compact_cap); } m.total_bytes = m.raw_bytes + m.compressed_bytes + m.scratch_bytes; return m; } +static ds4_context_memory glm_graph_context_memory_estimate_for_compact_cap( + uint32_t ctx, + uint32_t work_ctx, + uint32_t compact_cap, + bool ssd_streaming) { + const uint32_t normal_layers = glm_graph_normal_layer_count(); + if (normal_layers == 0) { + const ds4_context_memory empty = {0}; + return empty; + } + return glm_graph_context_memory_estimate_for_compact_cap_slice( + ctx, + work_ctx, + compact_cap, + ssd_streaming, + 0, + normal_layers - 1u); +} + ds4_context_memory ds4_context_memory_estimate_with_prefill_mode( ds4_backend backend, int ctx_size, @@ -35363,11 +35464,36 @@ static void ds4_engine_print_startup_memory( int ctx_size) { if (!e || ctx_size <= 0) return; - const ds4_context_memory mem = - ds4_context_memory_estimate_with_prefill_mode(e->backend, - ctx_size, - e->prefill_chunk, - e->ssd_streaming); + ds4_context_memory mem; +#ifndef DS4_NO_GPU + if (DS4_MODEL_FAMILY == DS4_MODEL_FAMILY_GLM_DSA && + e->distributed.role != DS4_DISTRIBUTED_NONE && + e->distributed.layers.set) { + const uint32_t normal_layers = glm_graph_normal_layer_count(); + const uint32_t layer_end = e->distributed.layers.has_output ? + (normal_layers ? normal_layers - 1u : 0u) : + e->distributed.layers.end; + const uint32_t ctx = (uint32_t)ctx_size; + const uint32_t work_ctx = + glm_graph_full_attention_cap(ctx, e->ssd_streaming); + const uint32_t compact_cap = + glm_graph_compact_cache_initial_cap(ctx, work_ctx); + mem = glm_graph_context_memory_estimate_for_compact_cap_slice( + ctx, + work_ctx, + compact_cap, + e->ssd_streaming, + e->distributed.layers.start, + layer_end); + } else { +#endif + mem = ds4_context_memory_estimate_with_prefill_mode(e->backend, + ctx_size, + e->prefill_chunk, + e->ssd_streaming); +#ifndef DS4_NO_GPU + } +#endif const uint64_t kv_bytes = ds4_add_sat_u64(mem.raw_bytes, mem.compressed_bytes); const uint64_t dynamic_expert_cache_bytes = @@ -37353,11 +37479,13 @@ static uint64_t glm_graph_streaming_active_model_bytes( uint64_t max_bytes = 0; ds4_model_map_span_vec spans; - if (weights_model_map_token_spans(weights, &spans)) { + if (weights_layer_has_required(&weights->layer[0], 0) && + weights_model_map_token_spans(weights, &spans)) { max_bytes = model_map_span_vec_total_bytes(&spans); free(spans.v); } - if (weights_model_map_output_spans(weights, &spans)) { + if (weights_have_output_head(weights) && + weights_model_map_output_spans(weights, &spans)) { const uint64_t bytes = model_map_span_vec_total_bytes(&spans); if (bytes > max_bytes) max_bytes = bytes; free(spans.v); @@ -37478,11 +37606,19 @@ static bool glm_graph_memory_guard_for_compact_cap( const uint32_t work_ctx = glm_graph_full_attention_cap(ctx_size, ssd_streaming); - const ds4_context_memory mem = - glm_graph_context_memory_estimate_for_compact_cap(ctx_size, - work_ctx, - compact_cap, - ssd_streaming); + const ds4_context_memory mem = load_slice ? + glm_graph_context_memory_estimate_for_compact_cap_slice( + ctx_size, + work_ctx, + compact_cap, + ssd_streaming, + layer_start, + layer_end) : + glm_graph_context_memory_estimate_for_compact_cap( + ctx_size, + work_ctx, + compact_cap, + ssd_streaming); const uint64_t graph_bytes = mem.total_bytes; const uint64_t model_bytes = glm_graph_model_bytes_for_guard(model, @@ -53272,12 +53408,31 @@ static bool ds4_engine_configure_streaming_auto_cache(ds4_engine *e) { uint64_t per_expert_bytes = 0; if (!ds4_streaming_routed_expert_bytes(&e->weights, &per_expert_bytes)) { + /* A valid distributed GLM slice can contain only the leading dense + * layers. It has no routed weights to stream or cache. */ + if (DS4_MODEL_FAMILY == DS4_MODEL_FAMILY_GLM_DSA && + e->distributed.role != DS4_DISTRIBUTED_NONE && + e->distributed.layers.set) { + fprintf(stderr, + "ds4: SSD streaming layer slice has no routed experts; " + "expert cache disabled\n"); + e->ssd_streaming_cache_experts = 0; + e->ssd_streaming_cache_bytes = 0; + return true; + } fprintf(stderr, "ds4: SSD streaming auto cache could not measure routed expert size\n"); return false; } - const uint64_t max_model_experts = (uint64_t)DS4_N_LAYER * (uint64_t)DS4_N_EXPERT; + uint64_t max_model_experts = 0; + if (!ds4_streaming_cacheable_expert_count(&e->weights, + &max_model_experts, + NULL)) { + fprintf(stderr, + "ds4: SSD streaming auto cache could not count local cacheable experts\n"); + return false; + } ds4_ssd_cache_plan plan; if (!ds4_ssd_auto_cache_plan(recommended, non_routed_bytes, @@ -53323,9 +53478,24 @@ static bool ds4_engine_configure_streaming_auto_cache(ds4_engine *e) { glm_graph_full_attention_cap(guard_ctx, true); const uint32_t compact_cap = glm_graph_compact_cache_initial_cap(guard_ctx, work_ctx); - const ds4_context_memory graph_mem = - glm_graph_context_memory_estimate_for_compact_cap( + ds4_context_memory graph_mem; + if (e->distributed.role != DS4_DISTRIBUTED_NONE && + e->distributed.layers.set) { + const uint32_t normal_layers = glm_graph_normal_layer_count(); + const uint32_t layer_end = e->distributed.layers.has_output ? + (normal_layers ? normal_layers - 1u : 0u) : + e->distributed.layers.end; + graph_mem = glm_graph_context_memory_estimate_for_compact_cap_slice( + guard_ctx, + work_ctx, + compact_cap, + true, + e->distributed.layers.start, + layer_end); + } else { + graph_mem = glm_graph_context_memory_estimate_for_compact_cap( guard_ctx, work_ctx, compact_cap, true); + } uint64_t active_model_bytes = glm_graph_streaming_active_model_bytes(&e->weights); if (non_routed_bytes > active_model_bytes) { @@ -53438,24 +53608,36 @@ static uint32_t ds4_glm_streaming_normal_layer_count(void) { } static uint32_t ds4_glm_streaming_supported_resident_prefix_layers( - const ds4_weights *weights) { + const ds4_weights *weights, + uint32_t *layer_start_out) { + if (layer_start_out) *layer_start_out = 0; if (!weights) return 0; const uint32_t normal_layers = ds4_glm_streaming_normal_layer_count(); if (normal_layers <= DS4_N_LEADING_DENSE) return 0; + uint32_t layer_start = DS4_N_LEADING_DENSE; + while (layer_start < normal_layers && + !glm_stream_resident_decode_layer_supported( + &weights->layer[layer_start], layer_start)) { + layer_start++; + } + if (layer_start == normal_layers) return 0; + uint32_t n = 0; - for (uint32_t il = DS4_N_LEADING_DENSE; il < normal_layers; il++) { + for (uint32_t il = layer_start; il < normal_layers; il++) { if (!glm_stream_resident_decode_layer_supported(&weights->layer[il], il)) { break; } n++; } + if (layer_start_out) *layer_start_out = layer_start; return n; } static bool ds4_glm_streaming_resident_prefix_bytes( const ds4_weights *weights, + uint32_t layer_start, uint32_t layers, uint64_t *bytes_out) { if (bytes_out) *bytes_out = 0; @@ -53463,7 +53645,7 @@ static bool ds4_glm_streaming_resident_prefix_bytes( uint64_t total = 0; for (uint32_t i = 0; i < layers; i++) { - const uint32_t il = DS4_N_LEADING_DENSE + i; + const uint32_t il = layer_start + i; if (!glm_stream_resident_decode_layer_supported(&weights->layer[il], il)) { return false; @@ -53486,6 +53668,7 @@ static bool ds4_glm_streaming_resident_prefix_bytes( static uint32_t ds4_glm_streaming_auto_full_layers( const ds4_weights *weights, + uint32_t layer_start, uint32_t supported_layers, uint64_t total_budget_bytes) { if (!weights || supported_layers == 0 || total_budget_bytes == 0) { @@ -53500,7 +53683,10 @@ static uint32_t ds4_glm_streaming_auto_full_layers( uint32_t best = 0; for (uint32_t n = 1; n <= supported_layers; n++) { uint64_t bytes = 0; - if (!ds4_glm_streaming_resident_prefix_bytes(weights, n, &bytes)) { + if (!ds4_glm_streaming_resident_prefix_bytes(weights, + layer_start, + n, + &bytes)) { break; } if (bytes > target_bytes) break; @@ -53510,6 +53696,7 @@ static uint32_t ds4_glm_streaming_auto_full_layers( } static bool ds4_engine_configure_streaming_cache_budget(ds4_engine *e) { + g_glm_streaming_full_resident_start = 0; g_glm_streaming_full_resident_layers = 0; #ifdef DS4_NO_GPU (void)e; @@ -53531,21 +53718,54 @@ static bool ds4_engine_configure_streaming_cache_budget(ds4_engine *e) { uint64_t per_expert_bytes = 0; const bool need_expert_bytes = + e->ssd_streaming_cache_experts != 0 || e->ssd_streaming_cache_bytes != 0 || (glm_full_layer_streaming && !e->ssd_streaming_full_layers_set) || e->ssd_streaming_full_layers != 0; if (need_expert_bytes && !ds4_streaming_routed_expert_bytes(&e->weights, &per_expert_bytes)) { + const bool no_routed_slice = + DS4_MODEL_FAMILY == DS4_MODEL_FAMILY_GLM_DSA && + e->distributed.role != DS4_DISTRIBUTED_NONE && + e->distributed.layers.set; + if (no_routed_slice && + e->ssd_streaming_cache_experts == 0 && + e->ssd_streaming_cache_bytes == 0 && + e->ssd_streaming_full_layers == 0) { + return true; + } fprintf(stderr, "ds4: SSD streaming could not measure routed expert size\n"); return false; } + uint64_t max_cache_experts_u64 = 0; + if (need_expert_bytes && + !ds4_streaming_cacheable_expert_count(&e->weights, + &max_cache_experts_u64, + NULL)) { + fprintf(stderr, + "ds4: SSD streaming could not count local cacheable experts\n"); + return false; + } + const uint32_t max_cache_experts = max_cache_experts_u64 > UINT32_MAX ? + UINT32_MAX : (uint32_t)max_cache_experts_u64; + if (e->ssd_streaming_cache_bytes == 0 && + max_cache_experts != 0 && + e->ssd_streaming_cache_experts > max_cache_experts) { + fprintf(stderr, + "ds4: SSD streaming expert cache capped from %u to %u " + "experts cacheable by this layer slice\n", + e->ssd_streaming_cache_experts, + max_cache_experts); + e->ssd_streaming_cache_experts = max_cache_experts; + } uint32_t full_layers = 0; uint64_t full_layers_bytes = 0; bool full_layers_auto = false; uint32_t supported = 0; + uint32_t supported_start = 0; uint64_t total_cache_bytes = e->ssd_streaming_cache_bytes; uint64_t prefill_headroom_bytes = 0; uint64_t budget_after_prefill_headroom = total_cache_bytes; @@ -53559,7 +53779,7 @@ static bool ds4_engine_configure_streaming_cache_budget(ds4_engine *e) { if (prefill_headroom_bytes >= total_cache_bytes) { fprintf(stderr, "ds4: --ssd-streaming-cache-experts byte budget %.2f GiB " - "is too small: two routed prefill layers need %.2f GiB\n", + "is too small: routed prefill headroom needs %.2f GiB\n", (double)total_cache_bytes / 1073741824.0, (double)prefill_headroom_bytes / 1073741824.0); return false; @@ -53570,7 +53790,8 @@ static bool ds4_engine_configure_streaming_cache_budget(ds4_engine *e) { e->ssd_streaming_prefill_headroom_bytes = prefill_headroom_bytes; if (glm_full_layer_streaming) { supported = - ds4_glm_streaming_supported_resident_prefix_layers(&e->weights); + ds4_glm_streaming_supported_resident_prefix_layers( + &e->weights, &supported_start); if (!e->ssd_streaming_full_layers_set && e->ssd_streaming_cache_bytes != 0) { /* @@ -53587,6 +53808,7 @@ static bool ds4_engine_configure_streaming_cache_budget(ds4_engine *e) { e->ssd_streaming_full_layers = ds4_glm_streaming_auto_full_layers( &e->weights, + supported_start, supported, budget_after_prefill_headroom); } @@ -53596,8 +53818,9 @@ static bool ds4_engine_configure_streaming_cache_budget(ds4_engine *e) { if (e->ssd_streaming_full_layers != 0) { const uint32_t requested = e->ssd_streaming_full_layers; - const uint32_t supported = - ds4_glm_streaming_supported_resident_prefix_layers(&e->weights); + supported = + ds4_glm_streaming_supported_resident_prefix_layers( + &e->weights, &supported_start); full_layers = requested < supported ? requested : supported; if (full_layers != requested) { fprintf(stderr, @@ -53626,6 +53849,7 @@ static bool ds4_engine_configure_streaming_cache_budget(ds4_engine *e) { while (full_layers != 0) { uint64_t bytes = 0; if (!ds4_glm_streaming_resident_prefix_bytes(&e->weights, + supported_start, full_layers, &bytes)) { fprintf(stderr, @@ -53650,6 +53874,7 @@ static bool ds4_engine_configure_streaming_cache_budget(ds4_engine *e) { } } else if (full_layers != 0 && !ds4_glm_streaming_resident_prefix_bytes(&e->weights, + supported_start, full_layers, &full_layers_bytes)) { fprintf(stderr, @@ -53670,11 +53895,14 @@ static bool ds4_engine_configure_streaming_cache_budget(ds4_engine *e) { } uint64_t budget_expert_bytes = 0; - const uint32_t budget = + uint32_t budget = ds4_streaming_cache_experts_for_byte_budget( &e->weights, dynamic_cache_bytes, &budget_expert_bytes); + if (max_cache_experts != 0 && budget > max_cache_experts) { + budget = max_cache_experts; + } if (budget == 0 || budget_expert_bytes == 0) { fprintf(stderr, "ds4: --ssd-streaming-cache-experts byte budget is too small or invalid for this model\n"); @@ -53743,6 +53971,7 @@ static bool ds4_engine_configure_streaming_cache_budget(ds4_engine *e) { e->ssd_streaming_full_layers = full_layers; e->ssd_streaming_full_layer_bytes = full_layers_bytes; + g_glm_streaming_full_resident_start = supported_start; g_glm_streaming_full_resident_layers = full_layers; return true; #endif @@ -55238,17 +55467,19 @@ static int ds4_engine_open_internal(ds4_engine **out, uint32_t load_layer_start = opt->load_layer_start; uint32_t load_layer_end = opt->load_layer_end; bool load_output = opt->load_output; + bool load_output_optional = false; if (opt->distributed.role != DS4_DISTRIBUTED_NONE && opt->distributed.layers.set) { load_slice = true; load_layer_start = opt->distributed.layers.start; - load_layer_end = - (DS4_MODEL_FAMILY != DS4_MODEL_FAMILY_GLM_DSA && - opt->distributed.layers.has_output) ? - UINT32_MAX : - opt->distributed.layers.end; + load_layer_end = opt->distributed.layers.end; load_output = opt->distributed.layers.has_output; + /* A coordinator may need to apply the output head locally when the + * last worker advertises N:M rather than N:output. Bind it when the + * local GGUF contains it, without requiring split GGUFs to do so. */ + load_output_optional = + opt->distributed.role == DS4_DISTRIBUTED_COORDINATOR; } const bool graph_backend = ds4_backend_uses_graph(opt->backend); @@ -55256,6 +55487,20 @@ static int ds4_engine_open_internal(ds4_engine **out, model_open(&e->model, opt->model_path, graph_backend, !opt->inspect_only); if (opt->warm_weights) model_warm_weights(&e->model); config_validate_model(&e->model); + if (load_slice && load_layer_end == UINT32_MAX) { + const uint32_t normal_layers = ds4_model_normal_layer_count(); + if (normal_layers == 0) { + fprintf(stderr, "ds4: model reports no executable transformer layers\n"); + ds4_engine_close(e); + *out = NULL; + return 1; + } + load_layer_end = normal_layers - 1u; + } + if (e->distributed.role != DS4_DISTRIBUTED_NONE && + e->distributed.layers.set) { + e->distributed.layers.end = load_layer_end; + } if (e->cuda_tensor_parallel && DS4_MODEL_FAMILY != DS4_MODEL_FAMILY_DEEPSEEK4) { fprintf(stderr, @@ -55290,7 +55535,8 @@ static int ds4_engine_open_internal(ds4_engine **out, load_slice, load_layer_start, load_layer_end, - load_output); + load_output, + load_output_optional); /* TP always maps one contiguous routed-expert half per rank. Decide * immediately after binding so memory guards account only the bytes this @@ -55674,7 +55920,9 @@ static int ds4_engine_open_internal(ds4_engine **out, load_slice, load_layer_start, load_layer_end, - load_output, + load_output || + (load_output_optional && + weights_have_output_head(&e->weights)), opt->context_size, "after GLM streaming cache budget")) { ds4_engine_close(e); @@ -55747,7 +55995,9 @@ static int ds4_engine_open_internal(ds4_engine **out, uint32_t load_span_count = 0; if (e->ssd_streaming) { const bool map_output = load_slice && - load_output; + (load_output || + (load_output_optional && + weights_have_output_head(&e->weights))); ds4_model_map_span_vec spans; bool spans_ok = false; if (load_slice) { @@ -55755,7 +56005,7 @@ static int ds4_engine_open_internal(ds4_engine **out, &e->weights, load_layer_start, load_layer_end, - true, + load_layer_start == 0, map_output, &spans); } else { @@ -55789,8 +56039,10 @@ static int ds4_engine_open_internal(ds4_engine **out, snprintf(load_end, sizeof(load_end), "%u", load_layer_end); } fprintf(stderr, - "ds4: SSD streaming initial %s model map restricted to token + non-routed layers %u:%s (%u spans, %.2f GiB tensor span)\n", + "ds4: SSD streaming initial %s model map restricted to " + "%snon-routed layers %u:%s (%u spans, %.2f GiB tensor span)\n", ds4_backend_name(e->backend), + load_layer_start == 0 ? "token + " : "", load_layer_start, load_end, spans.len, @@ -55810,7 +56062,10 @@ static int ds4_engine_open_internal(ds4_engine **out, spans.max_tensor_bytes); free(spans.v); } else if (load_slice) { - const bool map_output = load_output; + const bool map_output = + load_output || + (load_output_optional && + weights_have_output_head(&e->weights)); char load_end[32]; if (map_output && load_layer_end == UINT32_MAX) { snprintf(load_end, sizeof(load_end), "output"); diff --git a/ds4_rocm.cu b/ds4_rocm.cu index 640e00209..2230f6a58 100644 --- a/ds4_rocm.cu +++ b/ds4_rocm.cu @@ -34,6 +34,7 @@ #include #include #include +#include #include #include diff --git a/rocm/ds4_rocm_runtime.cuh b/rocm/ds4_rocm_runtime.cuh index 07152280c..730f91a3c 100644 --- a/rocm/ds4_rocm_runtime.cuh +++ b/rocm/ds4_rocm_runtime.cuh @@ -57,6 +57,7 @@ struct cuda_model_image { const void *host_base; uint64_t size; char *device_ptr; + uint64_t device_offset; }; struct cuda_q8_f16_range { @@ -597,11 +598,28 @@ static int cuda_model_image_find(const void *model_map) { } static const char *cuda_model_image_ptr(const void *model_map, uint64_t offset) { - const int idx = cuda_model_image_find(model_map); - if (idx < 0) return NULL; - const cuda_model_image &img = g_model_images[(size_t)idx]; - if (offset > img.size) return NULL; - return img.device_ptr + offset; + for (const cuda_model_image &img : g_model_images) { + if (img.host_base != model_map || offset < img.device_offset) continue; + const uint64_t rel = offset - img.device_offset; + if (rel >= img.size) continue; + return img.device_ptr + rel; + } + return NULL; +} + +static const char *cuda_model_image_range_ptr( + const void *model_map, + uint64_t offset, + uint64_t bytes) { + if (bytes == 0) return cuda_model_image_ptr(model_map, offset); + for (const cuda_model_image &img : g_model_images) { + if (img.host_base != model_map || offset < img.device_offset) continue; + const uint64_t rel = offset - img.device_offset; + if (rel <= img.size && bytes <= img.size - rel) { + return img.device_ptr + rel; + } + } + return NULL; } static int cuda_model_image_owned(const void *model_map) { @@ -4570,7 +4588,9 @@ static const char *cuda_model_range_copy_uncached( static const char *cuda_model_range_ptr(const void *model_map, uint64_t offset, uint64_t bytes, const char *what) { if (bytes == 0) return cuda_model_ptr(model_map, offset); - if (cuda_model_image_owned(model_map)) return cuda_model_ptr(model_map, offset); + const char *image_ptr = + cuda_model_image_range_ptr(model_map, offset, bytes); + if (image_ptr) return image_ptr; const uint64_t end = offset + bytes; auto exact = g_model_range_by_offset.find(offset); @@ -4669,7 +4689,7 @@ static const char *cuda_model_range_ptr(const void *model_map, uint64_t offset, static int cuda_model_range_is_cached(const void *model_map, uint64_t offset, uint64_t bytes) { if (bytes == 0) return 1; - if (cuda_model_image_owned(model_map)) return 1; + if (cuda_model_image_range_ptr(model_map, offset, bytes)) return 1; const uint64_t end = offset + bytes; if (end < offset) return 0; @@ -5481,13 +5501,31 @@ static char *cuda_model_arena_alloc(uint64_t bytes, const char *what) { void *dev = NULL; cudaError_t err = cudaMalloc(&dev, (size_t)chunk); if (err != cudaSuccess) { - fprintf(stderr, DS4_GPU_LOG_PREFIX "model arena alloc failed for %s (%.2f MiB chunk): %s\n", - what ? what : "weights", - (double)chunk / 1048576.0, - cudaGetErrorString(err)); (void)cudaGetLastError(); - g_model_cache_full = 1; - return NULL; + uint64_t fallback = chunk / 2u; + while (fallback >= aligned) { + err = cudaMalloc(&dev, (size_t)fallback); + if (err == cudaSuccess) break; + (void)cudaGetLastError(); + fallback /= 2u; + } + if (err != cudaSuccess) { + err = cudaMalloc(&dev, (size_t)aligned); + if (err != cudaSuccess) { + fprintf(stderr, + DS4_GPU_LOG_PREFIX "model arena alloc failed for %s " + "(%.2f MiB request): %s\n", + what ? what : "weights", + (double)aligned / 1048576.0, + cudaGetErrorString(err)); + (void)cudaGetLastError(); + g_model_cache_full = 1; + return NULL; + } + fallback = aligned; + } + g_model_arenas.push_back({(char *)dev, fallback, aligned}); + return (char *)dev; } g_model_arenas.push_back({(char *)dev, chunk, aligned}); return (char *)dev; @@ -5595,17 +5633,12 @@ static const char *cuda_model_range_ptr_from_fd( static int cuda_model_copy_chunked(const void *model_map, uint64_t model_size, uint64_t map_offset, uint64_t map_size) { if (!model_map || model_size == 0 || map_offset > model_size || map_size > model_size - map_offset) return 0; - if (cuda_model_image_owned(model_map)) { - g_model_host_base = model_map; - g_model_device_base = cuda_model_image_ptr(model_map, 0); - g_model_registered_size = model_size; - g_model_device_owned = 1; - return 1; - } + if (map_size == 0) return 0; + if (cuda_model_image_range_ptr(model_map, map_offset, map_size)) return 1; void *dev = NULL; const double t0 = cuda_wall_sec(); - cudaError_t err = cudaMalloc(&dev, (size_t)model_size); + cudaError_t err = cudaMalloc(&dev, (size_t)map_size); if (err != cudaSuccess) { fprintf(stderr, DS4_GPU_LOG_PREFIX "model allocation skipped: %s\n", cudaGetErrorString(err)); (void)cudaGetLastError(); @@ -5613,7 +5646,7 @@ static int cuda_model_copy_chunked(const void *model_map, uint64_t model_size, u } fprintf(stderr, DS4_GPU_LOG_PREFIX "chunk-copying %.2f GiB model image\n", - (double)model_size / 1073741824.0); + (double)map_size / 1073741824.0); const uint64_t chunk = cuda_model_copy_chunk_bytes(); const uint64_t stage_bytes = chunk + (g_model_direct_align > 1 ? g_model_direct_align : 1); @@ -5624,8 +5657,8 @@ static int cuda_model_copy_chunked(const void *model_map, uint64_t model_size, u uint64_t copied = 0; uint64_t chunk_idx = 0; - while (copied < model_size) { - const uint64_t n = (model_size - copied < chunk) ? (model_size - copied) : chunk; + while (copied < map_size) { + const uint64_t n = (map_size - copied < chunk) ? (map_size - copied) : chunk; const uint64_t bi = chunk_idx % 4u; if (chunk_idx >= 4u) { err = cudaEventSynchronize(g_model_stage_event[bi]); @@ -5638,7 +5671,7 @@ static int cuda_model_copy_chunked(const void *model_map, uint64_t model_size, u } const char *payload = NULL; if (!cuda_model_stage_read(g_model_stage[bi], g_model_stage_bytes, - copied, n, &payload)) { + map_offset + copied, n, &payload)) { fprintf(stderr, DS4_GPU_LOG_PREFIX "model staged read failed at %.2f GiB: %s\n", (double)copied / 1073741824.0, strerror(errno)); (void)cudaFree(dev); @@ -5660,11 +5693,12 @@ static int cuda_model_copy_chunked(const void *model_map, uint64_t model_size, u (void)cudaGetLastError(); return 0; } - cuda_model_drop_file_pages(copied, n); - cuda_model_discard_source_pages(model_map, model_size, copied, n); + cuda_model_drop_file_pages(map_offset + copied, n); + cuda_model_discard_source_pages(model_map, model_size, + map_offset + copied, n); copied += n; chunk_idx++; - cuda_model_load_progress_note(copied > map_offset ? copied - map_offset : 0); + cuda_model_load_progress_note(copied); } err = cudaStreamSynchronize(g_model_upload_stream); if (err != cudaSuccess) { @@ -5673,9 +5707,12 @@ static int cuda_model_copy_chunked(const void *model_map, uint64_t model_size, u (void)cudaGetLastError(); return 0; } - g_model_images.push_back({model_map, model_size, (char *)dev}); + g_model_images.push_back({model_map, map_size, (char *)dev, map_offset}); g_model_host_base = model_map; - g_model_device_base = (const char *)dev; + /* Sparse layer slices may create several disjoint device images, so no + * single base pointer can represent this model. Tensor lookups scan the + * image table and subtract each image's file offset. */ + g_model_device_base = NULL; g_model_registered_size = model_size; g_model_device_owned = 1; const double t1 = cuda_wall_sec(); @@ -6119,11 +6156,51 @@ extern "C" int ds4_gpu_set_model_map_spans( } return 1; } - /* - * The spans can be sparse distributed layer slices. Materializing their - * min..max envelope can be much larger than the actual selected tensors. - * Leave the precise per-tensor preparation to accelerator_cache_model_tensors(). - */ + + uint64_t span_bytes = 0; + uint64_t min_offset = offsets[0]; + uint64_t max_end = offsets[0] + sizes[0]; + for (uint32_t i = 0; i < count; i++) { + if (span_bytes > UINT64_MAX - sizes[i]) return 0; + span_bytes += sizes[i]; + if (offsets[i] < min_offset) min_offset = offsets[i]; + const uint64_t end = offsets[i] + sizes[i]; + if (end > max_end) max_end = end; + } + + /* Preserve the working rocm-multi-node residency policy: copy tight + * slices once, but split sparse coordinator layer+head selections into + * separate device images instead of allocating their huge file envelope. */ + const uint64_t bbox = max_end - min_offset; + if (bbox <= span_bytes + span_bytes / 10u) { + return cuda_model_copy_chunked(model_map, model_size, + min_offset, bbox); + } + + std::vector> sorted(count); + for (uint32_t i = 0; i < count; i++) { + sorted[i] = {offsets[i], offsets[i] + sizes[i]}; + } + std::sort(sorted.begin(), sorted.end()); + uint64_t group_offset = sorted[0].first; + uint64_t group_end = sorted[0].second; + const uint64_t merge_gap = 64ull * 1024ull; + for (uint32_t i = 1; i <= count; i++) { + if (i < count && + (sorted[i].first <= group_end || + sorted[i].first - group_end <= merge_gap)) { + if (sorted[i].second > group_end) group_end = sorted[i].second; + continue; + } + if (!cuda_model_copy_chunked(model_map, model_size, + group_offset, group_end - group_offset)) { + return 0; + } + if (i < count) { + group_offset = sorted[i].first; + group_end = sorted[i].second; + } + } return 1; } From 84d2de4b3fa1ced9f812cc30857a393c1737ba77 Mon Sep 17 00:00:00 2001 From: Donato Capitella Date: Wed, 22 Jul 2026 04:40:25 +0100 Subject: [PATCH 03/33] Allow resident ROCm GLM distributed slices --- ds4.c | 25 ++++++++++++++++++++----- 1 file changed, 20 insertions(+), 5 deletions(-) diff --git a/ds4.c b/ds4.c index e9f0bf82d..39e056894 100644 --- a/ds4.c +++ b/ds4.c @@ -55603,15 +55603,27 @@ static int ds4_engine_open_internal(ds4_engine **out, } bool glm_backend_supported = ds4_backend_uses_graph(e->backend); #ifdef DS4_ROCM_BUILD - if (e->backend == DS4_BACKEND_CUDA && !e->ssd_streaming) { + const bool rocm_full_model_requires_streaming = + e->backend == DS4_BACKEND_CUDA && + !e->ssd_streaming && + !load_slice; + if (rocm_full_model_requires_streaming) { glm_backend_supported = false; } #endif if (!glm_backend_supported) { #ifdef DS4_ROCM_BUILD - fprintf(stderr, - "ds4: GLM 5.2 ROCm inference requires --ssd-streaming; " - "use --inspect or --cpu --first-token-test for CPU diagnostics\n"); + if (rocm_full_model_requires_streaming) { + fprintf(stderr, + "ds4: full-model GLM 5.2 ROCm inference requires " + "--ssd-streaming; distributed layer slices can run " + "fully resident\n"); + } else { + fprintf(stderr, + "ds4: GLM 5.2 inference requires the ROCm graph " + "backend; use --inspect or --cpu --first-token-test " + "for CPU diagnostics\n"); + } #else fprintf(stderr, "ds4: GLM 5.2 inference requires the Metal or CUDA graph " @@ -55632,7 +55644,10 @@ static int ds4_engine_open_internal(ds4_engine **out, load_layer_start, load_layer_end, load_layer_start == 0, - load_output, + load_output || + (load_output_optional && + weights_have_output_head( + &e->weights)), guard_ctx) : glm_graph_memory_guard(&e->model, &e->weights, From fd6d056f5c6dbcca64aca8c6d8639bc605daff1b Mon Sep 17 00:00:00 2001 From: Donato Capitella Date: Wed, 22 Jul 2026 12:33:40 +0100 Subject: [PATCH 04/33] Add non-invasive ROCm graph dumps --- rocm/ds4_rocm_runtime.cuh | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/rocm/ds4_rocm_runtime.cuh b/rocm/ds4_rocm_runtime.cuh index 730f91a3c..91752ca49 100644 --- a/rocm/ds4_rocm_runtime.cuh +++ b/rocm/ds4_rocm_runtime.cuh @@ -4767,9 +4767,17 @@ static const ds4_rocm_runtime_config *cuda_runtime_config(void) { g_rocm_cfg.disable_shared_gate_up_fused_w32 = !g_quality_mode; g_rocm_cfg.attention_output_cublas_all = !g_quality_mode; g_rocm_cfg.shared_down_cublas = !g_quality_mode; - g_rocm_cfg.graph_dump = + const int graph_dump_requested = cuda_env_present(getenv("DS4_ROCM_GRAPH_DUMP_PREFIX")) || cuda_env_present(getenv("DS4_METAL_GRAPH_DUMP_PREFIX")); + /* Graph dumps traditionally select conservative kernels so their + * intermediate tensors are easier to inspect. Correctness bisects + * sometimes need the opposite: observe the exact production kernel + * path without the diagnostic changing it. */ + const int graph_dump_noninvasive = + cuda_env_present(getenv("DS4_ROCM_GRAPH_DUMP_NONINVASIVE")); + g_rocm_cfg.graph_dump = + graph_dump_requested && !graph_dump_noninvasive; 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'; From 34d8046bbf6791fb70a333cbe81cc884b418aa2b Mon Sep 17 00:00:00 2001 From: Donato Capitella Date: Wed, 22 Jul 2026 13:21:33 +0100 Subject: [PATCH 05/33] Add GLM batch FFN stage dumps --- ds4.c | 50 ++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 50 insertions(+) diff --git a/ds4.c b/ds4.c index 39e056894..9ab9cfcac 100644 --- a/ds4.c +++ b/ds4.c @@ -41699,6 +41699,13 @@ static bool glm_graph_encode_ffn_batch( pos0, n_tokens, stage_t0); + if (ok) { + metal_graph_debug_dump_tensor("glm_ffn_norm", + g->batch_ffn_norm, + (uint64_t)n_tokens * DS4_N_EMBD, + il, + pos0); + } if (!ok) return false; if (il < DS4_N_LEADING_DENSE) { @@ -41842,6 +41849,28 @@ static bool glm_graph_encode_ffn_batch( pos0, n_tokens, stage_t0); + if (ok) { + metal_graph_debug_dump_tensor("glm_ffn_router_logits", + g->batch_router_logits, + (uint64_t)n_tokens * DS4_N_EXPERT, + il, + pos0); + metal_graph_debug_dump_tensor("glm_ffn_router_probs", + g->batch_router_probs, + (uint64_t)n_tokens * DS4_N_EXPERT, + il, + pos0); + metal_graph_debug_dump_i32_tensor("glm_ffn_router_selected", + g->batch_router_selected, + (uint64_t)n_tokens * DS4_N_EXPERT_USED, + il, + pos0); + metal_graph_debug_dump_tensor("glm_ffn_router_weights", + g->batch_router_weights, + (uint64_t)n_tokens * DS4_N_EXPERT_USED, + il, + pos0); + } if (ok) ok = glm_graph_profile_router_selection_batch(g, l, il, @@ -42016,8 +42045,22 @@ static bool glm_graph_encode_ffn_batch( pos0, n_tokens, stage_t0); + if (ok) { + metal_graph_debug_dump_tensor("glm_ffn_routed_out", + g->batch_ffn_out, + (uint64_t)n_tokens * DS4_N_EMBD, + il, + pos0); + } if (ok && !shared_done) DS4_GLM_ENCODE_FFN_BATCH_SHARED(); #undef DS4_GLM_ENCODE_FFN_BATCH_SHARED + if (ok) { + metal_graph_debug_dump_tensor("glm_ffn_shared_out", + g->batch_attn_out, + (uint64_t)n_tokens * DS4_N_EMBD, + il, + pos0); + } if (ok && !glm_graph_disable_add3_residual()) { ok = ds4_gpu_add3_tensor(next, after_attn, @@ -42042,6 +42085,13 @@ static bool glm_graph_encode_ffn_batch( pos0, n_tokens, stage_t0); + if (ok) { + metal_graph_debug_dump_tensor("glm_ffn_next", + next, + (uint64_t)n_tokens * DS4_N_EMBD, + il, + pos0); + } return ok; } From 784bc77b21954e663d187794d00d44bc94e4649b Mon Sep 17 00:00:00 2001 From: Donato Capitella Date: Wed, 22 Jul 2026 13:51:29 +0100 Subject: [PATCH 06/33] Add resident IQ2 MoE sorted-path isolation --- rocm/ds4_rocm_moe_launch.cuh | 19 +++++++++++++++++-- 1 file changed, 17 insertions(+), 2 deletions(-) diff --git a/rocm/ds4_rocm_moe_launch.cuh b/rocm/ds4_rocm_moe_launch.cuh index 3a3d417eb..af4739a97 100644 --- a/rocm/ds4_rocm_moe_launch.cuh +++ b/rocm/ds4_rocm_moe_launch.cuh @@ -742,8 +742,23 @@ static int routed_moe_launch( if (!q2k_path && down->bytes >= xq_bytes && gate->bytes >= midq_bytes) { cuda_block_q8_K *xq = (cuda_block_q8_K *)down->ptr; cuda_block_q8_K *midq = (cuda_block_q8_K *)gate->ptr; - const uint32_t use_sorted_pairs = n_tokens > 1u && - (!q4k_path || n_tokens >= 32u); + const uint32_t disable_resident_iq2_sorted = + full_table_cached && + iq2_gate_path && + getenv("DS4_ROCM_DISABLE_RESIDENT_IQ2_SORTED") != NULL; + if (disable_resident_iq2_sorted) { + static int warned; + if (!warned) { + fprintf(stderr, + DS4_GPU_LOG_PREFIX "resident IQ2 sorted routed-MoE " + "disabled for correctness isolation\n"); + warned = 1; + } + } + const uint32_t use_sorted_pairs = + n_tokens > 1u && + (!q4k_path || n_tokens >= 32u) && + !disable_resident_iq2_sorted; const uint32_t use_expert_tiles = use_sorted_pairs; const uint32_t expert_tile_m = 8u; const uint32_t write_gate_up = 0u; From 507fd7c6a9373c3178db2bfc43e87b74865fd4be Mon Sep 17 00:00:00 2001 From: Donato Capitella Date: Wed, 22 Jul 2026 14:18:05 +0100 Subject: [PATCH 07/33] Fix resident ROCm IQ2 routed MoE correctness --- rocm/ds4_rocm_moe_launch.cuh | 17 +++++------------ 1 file changed, 5 insertions(+), 12 deletions(-) diff --git a/rocm/ds4_rocm_moe_launch.cuh b/rocm/ds4_rocm_moe_launch.cuh index af4739a97..488209d5a 100644 --- a/rocm/ds4_rocm_moe_launch.cuh +++ b/rocm/ds4_rocm_moe_launch.cuh @@ -742,19 +742,12 @@ static int routed_moe_launch( if (!q2k_path && down->bytes >= xq_bytes && gate->bytes >= midq_bytes) { cuda_block_q8_K *xq = (cuda_block_q8_K *)down->ptr; cuda_block_q8_K *midq = (cuda_block_q8_K *)gate->ptr; + /* The ROCm resident IQ2 sorted/expert-tile batch path corrupts the + * routed output. The ordinary selected-pair qwarp path below uses + * the same resident expert table and is correct; decode is unchanged + * because sorting only applies to multi-token batches. */ const uint32_t disable_resident_iq2_sorted = - full_table_cached && - iq2_gate_path && - getenv("DS4_ROCM_DISABLE_RESIDENT_IQ2_SORTED") != NULL; - if (disable_resident_iq2_sorted) { - static int warned; - if (!warned) { - fprintf(stderr, - DS4_GPU_LOG_PREFIX "resident IQ2 sorted routed-MoE " - "disabled for correctness isolation\n"); - warned = 1; - } - } + full_table_cached && iq2_gate_path; const uint32_t use_sorted_pairs = n_tokens > 1u && (!q4k_path || n_tokens >= 32u) && From 80afbe783687cf4ad35950ee92483c62f6f5d9b2 Mon Sep 17 00:00:00 2001 From: Donato Capitella Date: Wed, 22 Jul 2026 14:23:46 +0100 Subject: [PATCH 08/33] Fix ROCm IQ2 expert tile LUT initialization --- rocm/ds4_rocm_moe.cuh | 40 +++++++++++++++++++++--------------- rocm/ds4_rocm_moe_launch.cuh | 7 ++----- 2 files changed, 25 insertions(+), 22 deletions(-) diff --git a/rocm/ds4_rocm_moe.cuh b/rocm/ds4_rocm_moe.cuh index df74b736d..684f86729 100644 --- a/rocm/ds4_rocm_moe.cuh +++ b/rocm/ds4_rocm_moe.cuh @@ -883,11 +883,11 @@ __global__ static void moe_gate_up_mid_decode_lut_qwarp32_kernel( __shared__ uint8_t s_iq2_signs[128]; if (xq_blocks <= 16u) { for (uint32_t i = threadIdx.x; i < xq_blocks; i += blockDim.x) sxq[i] = xqb[i]; - for (uint32_t i = threadIdx.x; i < 256u; i += blockDim.x) s_iq2_grid[i] = cuda_iq2xxs_grid[i]; - for (uint32_t i = threadIdx.x; i < 128u; i += blockDim.x) s_iq2_signs[i] = cuda_ksigns_iq2xs[i]; - __syncthreads(); - xqb = sxq; } + for (uint32_t i = threadIdx.x; i < 256u; i += blockDim.x) s_iq2_grid[i] = cuda_iq2xxs_grid[i]; + for (uint32_t i = threadIdx.x; i < 128u; i += blockDim.x) s_iq2_signs[i] = cuda_ksigns_iq2xs[i]; + __syncthreads(); + if (xq_blocks <= 16u) xqb = sxq; for (uint32_t rr = 0; rr < 4u; rr++) { uint32_t row = blockIdx.x * 128u + row_lane + rr * 32u; if (row >= expert_mid_dim) continue; @@ -951,11 +951,11 @@ __global__ static void moe_gate_up_mid_decode_lut_qwarp32_ptrs_kernel( __shared__ uint8_t s_iq2_signs[128]; if (xq_blocks <= 16u) { for (uint32_t i = threadIdx.x; i < xq_blocks; i += blockDim.x) sxq[i] = xqb[i]; - for (uint32_t i = threadIdx.x; i < 256u; i += blockDim.x) s_iq2_grid[i] = cuda_iq2xxs_grid[i]; - for (uint32_t i = threadIdx.x; i < 128u; i += blockDim.x) s_iq2_signs[i] = cuda_ksigns_iq2xs[i]; - __syncthreads(); - xqb = sxq; } + for (uint32_t i = threadIdx.x; i < 256u; i += blockDim.x) s_iq2_grid[i] = cuda_iq2xxs_grid[i]; + for (uint32_t i = threadIdx.x; i < 128u; i += blockDim.x) s_iq2_signs[i] = cuda_ksigns_iq2xs[i]; + __syncthreads(); + if (xq_blocks <= 16u) xqb = sxq; for (uint32_t rr = 0; rr < 4u; rr++) { uint32_t row = blockIdx.x * 128u + row_lane + rr * 32u; if (row >= expert_mid_dim) continue; @@ -1327,9 +1327,11 @@ __global__ static void moe_gate_up_mid_expert_tile8_row32_kernel( uint32_t b = i - p * xq_blocks; sxq[p][b] = xqb[p][b]; } - for (uint32_t i = threadIdx.x; i < 256u; i += blockDim.x) s_iq2_grid[i] = cuda_iq2xxs_grid[i]; - for (uint32_t i = threadIdx.x; i < 128u; i += blockDim.x) s_iq2_signs[i] = cuda_ksigns_iq2xs[i]; - __syncthreads(); + } + for (uint32_t i = threadIdx.x; i < 256u; i += blockDim.x) s_iq2_grid[i] = cuda_iq2xxs_grid[i]; + for (uint32_t i = threadIdx.x; i < 128u; i += blockDim.x) s_iq2_signs[i] = cuda_ksigns_iq2xs[i]; + __syncthreads(); + if (xq_blocks <= 16u) { for (uint32_t p = 0; p < np; p++) xqb[p] = sxq[p]; } if (row >= expert_mid_dim) return; @@ -1420,9 +1422,11 @@ __global__ static void moe_gate_up_mid_expert_tile8_row2048_kernel( uint32_t b = i - p * xq_blocks; sxq[p][b] = xqb[p][b]; } - for (uint32_t i = threadIdx.x; i < 256u; i += blockDim.x) s_iq2_grid[i] = cuda_iq2xxs_grid[i]; - for (uint32_t i = threadIdx.x; i < 128u; i += blockDim.x) s_iq2_signs[i] = cuda_ksigns_iq2xs[i]; - __syncthreads(); + } + for (uint32_t i = threadIdx.x; i < 256u; i += blockDim.x) s_iq2_grid[i] = cuda_iq2xxs_grid[i]; + for (uint32_t i = threadIdx.x; i < 128u; i += blockDim.x) s_iq2_signs[i] = cuda_ksigns_iq2xs[i]; + __syncthreads(); + if (xq_blocks <= 16u) { for (uint32_t p = 0; p < np; p++) xqb[p] = sxq[p]; } for (uint32_t rr = 0; rr < 64u; rr++) { @@ -1517,9 +1521,11 @@ __global__ static void moe_gate_up_mid_expert_tile8_rowspan_kernel( uint32_t b = i - p * xq_blocks; sxq[p][b] = xqb[p][b]; } - for (uint32_t i = threadIdx.x; i < 256u; i += blockDim.x) s_iq2_grid[i] = cuda_iq2xxs_grid[i]; - for (uint32_t i = threadIdx.x; i < 128u; i += blockDim.x) s_iq2_signs[i] = cuda_ksigns_iq2xs[i]; - __syncthreads(); + } + for (uint32_t i = threadIdx.x; i < 256u; i += blockDim.x) s_iq2_grid[i] = cuda_iq2xxs_grid[i]; + for (uint32_t i = threadIdx.x; i < 128u; i += blockDim.x) s_iq2_signs[i] = cuda_ksigns_iq2xs[i]; + __syncthreads(); + if (xq_blocks <= 16u) { for (uint32_t p = 0; p < np; p++) xqb[p] = sxq[p]; } for (uint32_t rr = 0; rr < ROW_SPAN / 32u; rr++) { diff --git a/rocm/ds4_rocm_moe_launch.cuh b/rocm/ds4_rocm_moe_launch.cuh index 488209d5a..51fafe3d3 100644 --- a/rocm/ds4_rocm_moe_launch.cuh +++ b/rocm/ds4_rocm_moe_launch.cuh @@ -742,12 +742,9 @@ static int routed_moe_launch( if (!q2k_path && down->bytes >= xq_bytes && gate->bytes >= midq_bytes) { cuda_block_q8_K *xq = (cuda_block_q8_K *)down->ptr; cuda_block_q8_K *midq = (cuda_block_q8_K *)gate->ptr; - /* The ROCm resident IQ2 sorted/expert-tile batch path corrupts the - * routed output. The ordinary selected-pair qwarp path below uses - * the same resident expert table and is correct; decode is unchanged - * because sorting only applies to multi-token batches. */ + /* Correctness rollback for the optimized resident IQ2 prefill path. */ const uint32_t disable_resident_iq2_sorted = - full_table_cached && iq2_gate_path; + iq2_gate_path && getenv("DS4_ROCM_DISABLE_RESIDENT_IQ2_SORTED") != NULL; const uint32_t use_sorted_pairs = n_tokens > 1u && (!q4k_path || n_tokens >= 32u) && From 2dc022d90d4265e5fbcde3eebf20344998bac852 Mon Sep 17 00:00:00 2001 From: Donato Capitella Date: Wed, 22 Jul 2026 15:10:34 +0100 Subject: [PATCH 09/33] Tune GLM resident ROCm slice memory reserve --- ds4.c | 16 +++++++++++++++- 1 file changed, 15 insertions(+), 1 deletion(-) diff --git a/ds4.c b/ds4.c index 9ab9cfcac..14053f770 100644 --- a/ds4.c +++ b/ds4.c @@ -37634,8 +37634,22 @@ static bool glm_graph_memory_guard_for_compact_cap( const double fraction = glm_graph_env_double("DS4_GLM_MEMORY_GUARD_FRACTION", 0.99, 0.50, 1.00); - const double default_reserve_gib = + double default_reserve_gib = glm_graph_memory_guard_default_reserve_gib(budget_base, model_bytes); +#ifdef DS4_ROCM_BUILD + if (load_slice && !ssd_streaming) { + /* The original fixed reserve protects Metal's shared host/GPU heap. + * A resident ROCm layer slice already accounts its exact model spans + * and owned graph state above. Keep proportional backend headroom for + * driver and temporary allocations without rejecting viable UMA + * slices merely because the heap is smaller than a high-memory Mac. */ + double rocm_reserve_gib = glm_graph_bytes_to_gib(budget_base) / 16.0; + if (rocm_reserve_gib < 8.0) rocm_reserve_gib = 8.0; + if (rocm_reserve_gib < default_reserve_gib) { + default_reserve_gib = rocm_reserve_gib; + } + } +#endif const double reserve_gib = glm_graph_env_double("DS4_GLM_MEMORY_GUARD_RESERVE_GB", default_reserve_gib, From 7a424a456a31ecb857e3260cc2a6546592c550c5 Mon Sep 17 00:00:00 2001 From: Donato Capitella Date: Sat, 25 Jul 2026 13:44:57 +0100 Subject: [PATCH 10/33] Add grouped GLM value projection path --- rocm/ds4_rocm_glm.cuh | 21 +++++++++++++++++++++ rocm/ds4_rocm_runtime.cuh | 3 +++ 2 files changed, 24 insertions(+) diff --git a/rocm/ds4_rocm_glm.cuh b/rocm/ds4_rocm_glm.cuh index 2cd689d88..ad13a4836 100644 --- a/rocm/ds4_rocm_glm.cuh +++ b/rocm/ds4_rocm_glm.cuh @@ -1931,6 +1931,27 @@ extern "C" int ds4_gpu_glm_value_project_q8_0_batch_heads_tensor( "glm_value_project", &w, &row_bytes)) { return 0; } + if (n_tokens > 1u && + (kv_lora_dim & 31u) == 0u && + cuda_runtime_config()->glm_grouped_value_project) { + const uint32_t rows_per_block = 32u; + const uint32_t token_tile = 32u; + const uint32_t block_tile = 16u; + cuda_launch_grouped_q8_a_sharedx( + (float *)heads->ptr, + w, + (const float *)lora->ptr, + n_tokens, + n_head, + kv_lora_dim >> 5u, + value_dim, + row_bytes, + rows_per_block, + token_tile, + block_tile); + return cuda_ok(cudaGetLastError(), + "glm grouped value project batch heads launch"); + } dim3 grid(n_head, n_tokens, 1); glm_q8_project_head_kernel<<>>( (float *)heads->ptr, diff --git a/rocm/ds4_rocm_runtime.cuh b/rocm/ds4_rocm_runtime.cuh index 91752ca49..28ccf33a5 100644 --- a/rocm/ds4_rocm_runtime.cuh +++ b/rocm/ds4_rocm_runtime.cuh @@ -4752,6 +4752,7 @@ struct ds4_rocm_runtime_config { int disable_shared_gate_up_fused_w32; int attention_output_cublas_all; int shared_down_cublas; + int glm_grouped_value_project; int graph_dump; uint32_t moe_decode_rpb; uint32_t moe_decode_gate_rpb; @@ -4767,6 +4768,8 @@ static const ds4_rocm_runtime_config *cuda_runtime_config(void) { g_rocm_cfg.disable_shared_gate_up_fused_w32 = !g_quality_mode; g_rocm_cfg.attention_output_cublas_all = !g_quality_mode; g_rocm_cfg.shared_down_cublas = !g_quality_mode; + g_rocm_cfg.glm_grouped_value_project = + cuda_env_present(getenv("DS4_ROCM_GLM_GROUPED_VALUE_PROJECT")); const int graph_dump_requested = cuda_env_present(getenv("DS4_ROCM_GRAPH_DUMP_PREFIX")) || cuda_env_present(getenv("DS4_METAL_GRAPH_DUMP_PREFIX")); From fa8b0b9301e4d7aaefba0228049c7c5d1d6cbca0 Mon Sep 17 00:00:00 2001 From: Donato Capitella Date: Sat, 25 Jul 2026 15:26:19 +0100 Subject: [PATCH 11/33] Optimize ROCm indexer top-k 2048 --- ds4_rocm.cu | 3 +-- rocm/ds4_rocm_indexer.cuh | 52 +++++++++++++++++++++++++++++++++------ 2 files changed, 46 insertions(+), 9 deletions(-) diff --git a/ds4_rocm.cu b/ds4_rocm.cu index 2230f6a58..46ba223df 100644 --- a/ds4_rocm.cu +++ b/ds4_rocm.cu @@ -53,8 +53,7 @@ enum { * decode calls to the online attention kernel so this fixed buffer never * becomes an out-of-bounds write at long context. */ DS4_ROCM_ATTENTION_SCORE_CAP = 8192u, - DS4_ROCM_ATTENTION_RAW_SCORE_CAP = 256u, - DS4_ROCM_TOPK_MERGE_GROUP = 8u + DS4_ROCM_ATTENTION_RAW_SCORE_CAP = 256u }; struct ds4_gpu_tensor { diff --git a/rocm/ds4_rocm_indexer.cuh b/rocm/ds4_rocm_indexer.cuh index 26b483d1c..a5e4eb93d 100644 --- a/rocm/ds4_rocm_indexer.cuh +++ b/rocm/ds4_rocm_indexer.cuh @@ -982,14 +982,52 @@ extern "C" int ds4_gpu_indexer_topk_tensor( n_comp, n_tokens, top_k); return cuda_ok(cudaGetLastError(), "indexer topk 8192x1024 launch"); } - if (top_k == 512u) { + if (top_k == 2048u && n_comp <= 4096u) { + indexer_topk_pow2_kernel<4096><<>>((uint32_t *)selected->ptr, + (const float *)scores->ptr, + n_comp, n_tokens, top_k); + return cuda_ok(cudaGetLastError(), "indexer topk 4096x2048 launch"); + } + if (top_k == 2048u && n_comp <= 8192u) { + if (n_comp > 4096u) { + using TopkCubSort = cub::BlockRadixSort; + const int smem = (int)sizeof(typename TopkCubSort::TempStorage); + int dev = 0; + int max_optin_smem = 0; + cudaError_t attr_err = cudaGetDevice(&dev); + if (attr_err == cudaSuccess) { + attr_err = cudaDeviceGetAttribute(&max_optin_smem, + cudaDevAttrMaxSharedMemoryPerBlockOptin, + dev); + } + if (attr_err == cudaSuccess && max_optin_smem >= smem) { + attr_err = cudaFuncSetAttribute(indexer_topk_8192_cub_kernel, + cudaFuncAttributeMaxDynamicSharedMemorySize, + smem); + if (attr_err == cudaSuccess) { + indexer_topk_8192_cub_kernel<<>>((uint32_t *)selected->ptr, + (const float *)scores->ptr, + n_comp, n_tokens, top_k); + return cuda_ok(cudaGetLastError(), "indexer topk 8192x2048 cub launch"); + } + } + } + indexer_topk_pow2_u16_kernel<8192><<>>((uint32_t *)selected->ptr, + (const float *)scores->ptr, + n_comp, n_tokens, top_k); + return cuda_ok(cudaGetLastError(), "indexer topk 8192x2048 launch"); + } + if (top_k == 512u || top_k == 1024u || top_k == 2048u) { const uint32_t chunk_n = 4096u; const uint32_t n_chunks = (n_comp + chunk_n - 1u) / chunk_n; - const uint32_t candidate_stride = n_chunks * top_k; + const uint32_t merge_group = chunk_n / top_k; + const uint64_t candidate_stride64 = (uint64_t)n_chunks * top_k; + if (candidate_stride64 > UINT32_MAX) return 0; + const uint32_t candidate_stride = (uint32_t)candidate_stride64; uint32_t n_sets = n_chunks; uint64_t scratch_u32_per_token = candidate_stride; - while (n_sets > DS4_ROCM_TOPK_MERGE_GROUP) { - n_sets = (n_sets + DS4_ROCM_TOPK_MERGE_GROUP - 1u) / DS4_ROCM_TOPK_MERGE_GROUP; + while (n_sets > merge_group) { + n_sets = (n_sets + merge_group - 1u) / merge_group; scratch_u32_per_token += (uint64_t)n_sets * top_k; } if (scratch_u32_per_token > UINT64_MAX / n_tokens / sizeof(uint32_t)) return 0; @@ -1009,8 +1047,8 @@ extern "C" int ds4_gpu_indexer_topk_tensor( candidate_stride); if (!cuda_ok(cudaGetLastError(), "indexer topk chunk launch")) return 0; - while (n_sets > DS4_ROCM_TOPK_MERGE_GROUP) { - const uint32_t next_sets = (n_sets + DS4_ROCM_TOPK_MERGE_GROUP - 1u) / DS4_ROCM_TOPK_MERGE_GROUP; + while (n_sets > merge_group) { + const uint32_t next_sets = (n_sets + merge_group - 1u) / merge_group; const uint32_t next_stride = next_sets * top_k; uint32_t *next = cur + (uint64_t)n_tokens * cur_stride; dim3 grid_merge(n_tokens, next_sets, 1); @@ -1022,7 +1060,7 @@ extern "C" int ds4_gpu_indexer_topk_tensor( n_tokens, top_k, n_sets, - DS4_ROCM_TOPK_MERGE_GROUP, + merge_group, cur_stride, next_stride); if (!cuda_ok(cudaGetLastError(), "indexer topk tree merge launch")) return 0; From f6f5b926316d9de877ce3e8b9d802cf4cb1d4e6e Mon Sep 17 00:00:00 2001 From: Donato Capitella Date: Sat, 25 Jul 2026 17:28:46 +0100 Subject: [PATCH 12/33] Optimize GLM ROCm prefill projections --- rocm/ds4_rocm_glm.cuh | 23 +++++++++++ rocm/ds4_rocm_matmul.cuh | 81 +++++++++++++++++++++++++++++++++++++ rocm/ds4_rocm_q8.cuh | 85 +++++++++++++++++++++++++++++++++++++++ rocm/ds4_rocm_runtime.cuh | 12 +++++- 4 files changed, 200 insertions(+), 1 deletion(-) diff --git a/rocm/ds4_rocm_glm.cuh b/rocm/ds4_rocm_glm.cuh index ad13a4836..108783ead 100644 --- a/rocm/ds4_rocm_glm.cuh +++ b/rocm/ds4_rocm_glm.cuh @@ -1891,6 +1891,29 @@ extern "C" int ds4_gpu_glm_qk_lowrank_q8_0_batch_tensor( "glm_qk_lowrank_batch", &w, &row_bytes)) { return 0; } + if (n_tokens > 1u && + (qk_nope & 31u) == 0u && + cuda_runtime_config()->glm_grouped_qk_low) { + const uint32_t rows_per_block = 32u; + const uint32_t token_tile = 32u; + const uint32_t block_tile = 8u; + cuda_launch_grouped_q8_a_sharedx_strided( + (float *)qk_low->ptr, + w, + (const float *)q->ptr, + n_tokens, + n_head, + qk_nope >> 5u, + kv_lora_dim, + (uint32_t)x_stride64, + qk_dim, + row_bytes, + rows_per_block, + token_tile, + block_tile); + return cuda_ok(cudaGetLastError(), + "glm grouped qk lowrank batch launch"); + } dim3 grid(n_head, n_tokens, 1); glm_q8_project_head_kernel<<>>( (float *)qk_low->ptr, diff --git a/rocm/ds4_rocm_matmul.cuh b/rocm/ds4_rocm_matmul.cuh index 8a62dcd16..31292e7d0 100644 --- a/rocm/ds4_rocm_matmul.cuh +++ b/rocm/ds4_rocm_matmul.cuh @@ -99,6 +99,87 @@ static void cuda_launch_grouped_q8_a_sharedx( } } +template +static void cuda_launch_grouped_q8_a_sharedx_strided_bt( + float *low, + const unsigned char *w, + const float *heads, + uint32_t n_tokens, + uint32_t n_groups, + uint32_t n_blocks, + uint32_t rank, + uint32_t x_token_stride, + uint32_t x_group_stride, + uint64_t row_bytes, + dim3 grid, + uint32_t rows_per_block, + uint32_t tile) { + const size_t shmem = (size_t)tile * BT * 32u * sizeof(float); + if (tile == 2u) { + grouped_q8_0_a_f32_batch_sharedx_chunked_strided_w32_kernel<2u, BT> + <<>>( + low, w, heads, n_tokens, n_groups, n_blocks, rank, + x_token_stride, x_group_stride, row_bytes); + } else if (tile == 4u) { + grouped_q8_0_a_f32_batch_sharedx_chunked_strided_w32_kernel<4u, BT> + <<>>( + low, w, heads, n_tokens, n_groups, n_blocks, rank, + x_token_stride, x_group_stride, row_bytes); + } else if (tile == 8u) { + grouped_q8_0_a_f32_batch_sharedx_chunked_strided_w32_kernel<8u, BT> + <<>>( + low, w, heads, n_tokens, n_groups, n_blocks, rank, + x_token_stride, x_group_stride, row_bytes); + } else if (tile == 16u) { + grouped_q8_0_a_f32_batch_sharedx_chunked_strided_w32_kernel<16u, BT> + <<>>( + low, w, heads, n_tokens, n_groups, n_blocks, rank, + x_token_stride, x_group_stride, row_bytes); + } else { + grouped_q8_0_a_f32_batch_sharedx_chunked_strided_w32_kernel<32u, BT> + <<>>( + low, w, heads, n_tokens, n_groups, n_blocks, rank, + x_token_stride, x_group_stride, row_bytes); + } +} + +static void cuda_launch_grouped_q8_a_sharedx_strided( + float *low, + const unsigned char *w, + const float *heads, + uint32_t n_tokens, + uint32_t n_groups, + uint32_t n_blocks, + uint32_t rank, + uint32_t x_token_stride, + uint32_t x_group_stride, + uint64_t row_bytes, + uint32_t rows_per_block, + uint32_t tile, + uint32_t block_tile) { + const uint32_t row_blocks = + (rank + rows_per_block - 1u) / rows_per_block; + const dim3 grid(n_groups * row_blocks, + (n_tokens + tile - 1u) / tile, + 1u); + if (block_tile == 8u) { + cuda_launch_grouped_q8_a_sharedx_strided_bt<8u>( + low, w, heads, n_tokens, n_groups, n_blocks, rank, + x_token_stride, x_group_stride, row_bytes, grid, + rows_per_block, tile); + } else if (block_tile == 32u) { + cuda_launch_grouped_q8_a_sharedx_strided_bt<32u>( + low, w, heads, n_tokens, n_groups, n_blocks, rank, + x_token_stride, x_group_stride, row_bytes, grid, + rows_per_block, tile); + } else { + cuda_launch_grouped_q8_a_sharedx_strided_bt<16u>( + low, w, heads, n_tokens, n_groups, n_blocks, rank, + x_token_stride, x_group_stride, row_bytes, grid, + rows_per_block, tile); + } +} + static int cuda_matmul_q8_0_tensor_f16_gemm( ds4_gpu_tensor *out, const void *model_map, diff --git a/rocm/ds4_rocm_q8.cuh b/rocm/ds4_rocm_q8.cuh index 5b2423de3..365d239e4 100644 --- a/rocm/ds4_rocm_q8.cuh +++ b/rocm/ds4_rocm_q8.cuh @@ -1568,6 +1568,91 @@ __global__ static void grouped_q8_0_a_f32_batch_sharedx_chunked_w32_kernel( } } +/* + * Variant of the grouped shared-X kernel for inputs whose logical groups are + * slices of a wider physical row. GLM QK-low projects only qk_nope values + * from each q head, while consecutive heads remain qk_dim values apart. + */ +template +__global__ static void grouped_q8_0_a_f32_batch_sharedx_chunked_strided_w32_kernel( + float *low, + const unsigned char *w, + const float *heads, + uint32_t n_tokens, + uint32_t n_groups, + uint32_t n_blocks, + uint32_t rank, + uint32_t x_token_stride, + uint32_t x_group_stride, + uint64_t row_bytes) { + extern __shared__ float shx[]; + const uint32_t tid = threadIdx.x; + const uint32_t lane = tid & 31u; + const uint32_t wave = tid >> 5u; + const uint32_t rows_per_block = blockDim.x >> 5u; + const uint32_t row_blocks = (rank + rows_per_block - 1u) / rows_per_block; + const uint32_t g = blockIdx.x / row_blocks; + const uint32_t row0 = (blockIdx.x - g * row_blocks) * rows_per_block + wave; + const uint32_t t0 = blockIdx.y * TOK_TILE; + if (g >= n_groups || t0 >= n_tokens) return; + const bool row_valid = row0 < rank; + const unsigned char *wr = + w + ((uint64_t)g * rank + (row_valid ? row0 : 0u)) * row_bytes; + float acc[TOK_TILE]; +#pragma unroll + for (uint32_t u = 0; u < TOK_TILE; u++) acc[u] = 0.0f; + + for (uint32_t b0 = 0; b0 < n_blocks; b0 += BLOCKS_TILE) { + const uint32_t b_count = + ((b0 + BLOCKS_TILE) <= n_blocks) ? BLOCKS_TILE : (n_blocks - b0); + for (uint32_t j = tid; + j < TOK_TILE * BLOCKS_TILE * 32u; + j += blockDim.x) { + const uint32_t u = j / (BLOCKS_TILE * 32u); + const uint32_t r = j - u * (BLOCKS_TILE * 32u); + const uint32_t bb = r >> 5u; + const uint32_t k = r & 31u; + const uint32_t t = t0 + u; + const uint64_t xoff = + (uint64_t)t * x_token_stride + + (uint64_t)g * x_group_stride + + ((uint64_t)(b0 + bb) << 5u) + k; + shx[j] = + (t < n_tokens && bb < b_count) ? heads[xoff] : 0.0f; + } + __syncthreads(); + if (row_valid) { + for (uint32_t bb = 0; bb < b_count; bb++) { + const unsigned char *blk = + wr + (uint64_t)(b0 + bb) * 34u; + const float d = q8_0_scale_broadcast_w32(blk); + const int8_t q = ((const int8_t *)(blk + 2u))[lane]; + const float wv = d * (float)q; +#pragma unroll + for (uint32_t u = 0; u < TOK_TILE; u++) { + acc[u] += + wv * shx[(u * BLOCKS_TILE + bb) * 32u + lane]; + } + } + } + __syncthreads(); + } + +#pragma unroll + for (uint32_t u = 0; u < TOK_TILE; u++) { + acc[u] = warp_sum_f32(acc[u]); + } + if (lane == 0u && row_valid) { +#pragma unroll + for (uint32_t u = 0; u < TOK_TILE; u++) { + const uint32_t t = t0 + u; + if (t < n_tokens) { + low[((uint64_t)t * n_groups + g) * rank + row0] = acc[u]; + } + } + } +} + #if defined(__HIP_PLATFORM_AMD__) || defined(__HIPCC__) template __global__ static void grouped_q8_0_a_f32_batch_wmma_onthefly_kernel( diff --git a/rocm/ds4_rocm_runtime.cuh b/rocm/ds4_rocm_runtime.cuh index 28ccf33a5..6d726e56b 100644 --- a/rocm/ds4_rocm_runtime.cuh +++ b/rocm/ds4_rocm_runtime.cuh @@ -4753,6 +4753,7 @@ struct ds4_rocm_runtime_config { int attention_output_cublas_all; int shared_down_cublas; int glm_grouped_value_project; + int glm_grouped_qk_low; int graph_dump; uint32_t moe_decode_rpb; uint32_t moe_decode_gate_rpb; @@ -4768,8 +4769,17 @@ static const ds4_rocm_runtime_config *cuda_runtime_config(void) { g_rocm_cfg.disable_shared_gate_up_fused_w32 = !g_quality_mode; g_rocm_cfg.attention_output_cublas_all = !g_quality_mode; g_rocm_cfg.shared_down_cublas = !g_quality_mode; + const char *glm_grouped_value_project_env = + getenv("DS4_ROCM_GLM_GROUPED_VALUE_PROJECT"); + /* + * The grouped batch kernel is the validated production path. Keep + * the environment variable as an explicit =0 correctness fallback. + */ g_rocm_cfg.glm_grouped_value_project = - cuda_env_present(getenv("DS4_ROCM_GLM_GROUPED_VALUE_PROJECT")); + glm_grouped_value_project_env == NULL || + cuda_env_present(glm_grouped_value_project_env); + g_rocm_cfg.glm_grouped_qk_low = + cuda_env_present(getenv("DS4_ROCM_GLM_GROUPED_QK_LOW")); const int graph_dump_requested = cuda_env_present(getenv("DS4_ROCM_GRAPH_DUMP_PREFIX")) || cuda_env_present(getenv("DS4_METAL_GRAPH_DUMP_PREFIX")); From 9463fd73d913f34d88e4924908202329d5b40fb9 Mon Sep 17 00:00:00 2001 From: Donato Capitella Date: Sat, 25 Jul 2026 18:29:37 +0100 Subject: [PATCH 13/33] Rename DS4_ROCM_GLM_GROUPED_QK_LOW to DS4_ROCM_GLM_GROUPED_QK_LOW --- rocm/ds4_rocm_runtime.cuh | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/rocm/ds4_rocm_runtime.cuh b/rocm/ds4_rocm_runtime.cuh index 6d726e56b..26688e108 100644 --- a/rocm/ds4_rocm_runtime.cuh +++ b/rocm/ds4_rocm_runtime.cuh @@ -4778,8 +4778,11 @@ static const ds4_rocm_runtime_config *cuda_runtime_config(void) { g_rocm_cfg.glm_grouped_value_project = glm_grouped_value_project_env == NULL || cuda_env_present(glm_grouped_value_project_env); + const char *glm_grouped_qk_low_env = + getenv("DS4_ROCM_GLM_GROUPED_QK_LOW"); g_rocm_cfg.glm_grouped_qk_low = - cuda_env_present(getenv("DS4_ROCM_GLM_GROUPED_QK_LOW")); + glm_grouped_qk_low_env == NULL || + cuda_env_present(glm_grouped_qk_low_env); const int graph_dump_requested = cuda_env_present(getenv("DS4_ROCM_GRAPH_DUMP_PREFIX")) || cuda_env_present(getenv("DS4_METAL_GRAPH_DUMP_PREFIX")); From 26e8e88665a63c10fccbf2e0634047e3f3a0e67d Mon Sep 17 00:00:00 2001 From: Donato Capitella Date: Sat, 25 Jul 2026 18:45:22 +0100 Subject: [PATCH 14/33] Add ROCm GLM causal prefill GEMM path --- rocm/ds4_rocm_glm.cuh | 413 ++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 413 insertions(+) diff --git a/rocm/ds4_rocm_glm.cuh b/rocm/ds4_rocm_glm.cuh index 108783ead..881405fa5 100644 --- a/rocm/ds4_rocm_glm.cuh +++ b/rocm/ds4_rocm_glm.cuh @@ -1990,6 +1990,380 @@ extern "C" int ds4_gpu_glm_value_project_q8_0_batch_heads_tensor( return cuda_ok(cudaGetLastError(), "glm value project batch heads launch"); } +__global__ static void glm_causal_gemm_cache_to_f16_kernel( + __half *dst, + const char *src, + uint64_t n, + bool src_f16) { + const uint64_t i = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x; + if (i < n) dst[i] = __float2half(glm_rocm_cache_load(src, i, src_f16)); +} + +__global__ static void glm_causal_gemm_rope_to_f16_kernel( + __half *dst, + const char *src, + uint32_t n_rows, + uint32_t qk_rope, + bool src_f16, + uint32_t n_ctx_orig, + float freq_base, + float freq_scale, + float ext_factor, + float attn_factor, + float beta_fast, + float beta_slow) { + const uint64_t pair = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x; + const uint64_t pairs = (uint64_t)n_rows * (qk_rope >> 1u); + if (pair >= pairs) return; + const uint32_t row = (uint32_t)(pair / (qk_rope >> 1u)); + const uint32_t r = + (uint32_t)(pair - (uint64_t)row * (qk_rope >> 1u)) << 1u; + const float2 y = glm_rocm_rotated_cache_rope_pair(src, + (uint64_t)row * qk_rope, + r, + row, + qk_rope, + src_f16, + n_ctx_orig, + freq_base, + freq_scale, + ext_factor, + attn_factor, + beta_fast, + beta_slow); + dst[(uint64_t)row * qk_rope + r] = __float2half(y.x); + dst[(uint64_t)row * qk_rope + r + 1u] = __float2half(y.y); +} + +__global__ static void glm_causal_gemm_gather_head_f16_kernel( + __half *low_h, + __half *qrope_h, + const float *low, + const float *q, + uint32_t n_tokens, + uint32_t head, + uint32_t n_head, + uint32_t kv_lora_dim, + uint32_t qk_nope, + uint32_t qk_rope) { + const uint64_t total = (uint64_t)n_tokens * kv_lora_dim; + for (uint64_t i = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x; + i < total; + i += (uint64_t)gridDim.x * blockDim.x) { + const uint32_t token = (uint32_t)(i / kv_lora_dim); + const uint32_t j = + (uint32_t)(i - (uint64_t)token * kv_lora_dim); + low_h[i] = __float2half( + low[((uint64_t)token * n_head + head) * kv_lora_dim + j]); + } + const uint64_t rope_total = (uint64_t)n_tokens * qk_rope; + for (uint64_t i = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x; + i < rope_total; + i += (uint64_t)gridDim.x * blockDim.x) { + const uint32_t token = (uint32_t)(i / qk_rope); + const uint32_t r = (uint32_t)(i - (uint64_t)token * qk_rope); + const uint32_t qk_dim = qk_nope + qk_rope; + qrope_h[i] = __float2half( + q[((uint64_t)token * n_head + head) * qk_dim + qk_nope + r]); + } +} + +__global__ static void glm_causal_gemm_softmax_f16_kernel( + __half *probs, + float *scores, + uint32_t n_tokens, + uint32_t n_selected, + uint32_t pos0, + float scale) { + const uint32_t token = blockIdx.x; + if (token >= n_tokens) return; + float *row = scores + (uint64_t)token * n_selected; + __shared__ float reduce[256]; + const uint32_t visible = min(n_selected, pos0 + token + 1u); + float vmax = -INFINITY; + for (uint32_t s = threadIdx.x; s < n_selected; s += blockDim.x) { + const float v = s < visible ? row[s] * scale : -INFINITY; + row[s] = v; + vmax = fmaxf(vmax, v); + } + reduce[threadIdx.x] = vmax; + __syncthreads(); + for (uint32_t stride = blockDim.x >> 1u; stride != 0u; stride >>= 1u) { + if (threadIdx.x < stride) { + reduce[threadIdx.x] = + fmaxf(reduce[threadIdx.x], reduce[threadIdx.x + stride]); + } + __syncthreads(); + } + const float max_score = reduce[0]; + float sum = 0.0f; + for (uint32_t s = threadIdx.x; s < n_selected; s += blockDim.x) { + const float v = s < visible ? expf(row[s] - max_score) : 0.0f; + row[s] = v; + sum += v; + } + reduce[threadIdx.x] = sum; + __syncthreads(); + for (uint32_t stride = blockDim.x >> 1u; stride != 0u; stride >>= 1u) { + if (threadIdx.x < stride) { + reduce[threadIdx.x] += reduce[threadIdx.x + stride]; + } + __syncthreads(); + } + const float inv = 1.0f / fmaxf(reduce[0], 1.0e-20f); + for (uint32_t s = threadIdx.x; s < n_selected; s += blockDim.x) { + probs[(uint64_t)token * n_selected + s] = + __float2half(row[s] * inv); + } +} + +__global__ static void glm_causal_gemm_scatter_head_kernel( + float *out, + const float *head_out, + uint32_t n_tokens, + uint32_t head, + uint32_t n_head, + uint32_t kv_lora_dim) { + const uint64_t i = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x; + const uint64_t n = (uint64_t)n_tokens * kv_lora_dim; + if (i >= n) return; + const uint32_t token = (uint32_t)(i / kv_lora_dim); + const uint32_t j = (uint32_t)(i - (uint64_t)token * kv_lora_dim); + out[((uint64_t)token * n_head + head) * kv_lora_dim + j] = + head_out[i]; +} + +static int glm_causal_gemm_scratch_part( + uint64_t *offset, + uint64_t bytes, + uint64_t *part_offset) { + uint64_t aligned = 0; + if (!offset || !part_offset || + !cuda_u64_add_checked(*offset, 255u, &aligned)) { + return 0; + } + aligned &= ~255ull; + if (!cuda_u64_add_checked(aligned, bytes, offset)) return 0; + *part_offset = aligned; + return 1; +} + +static int glm_attention_indexed_lora_causal_gemm( + ds4_gpu_tensor *lora_out, + const ds4_gpu_tensor *q, + const ds4_gpu_tensor *qk_low, + const ds4_gpu_tensor *kv_lora_cache, + const ds4_gpu_tensor *k_rope_cache, + uint32_t n_tokens, + uint32_t pos0, + uint32_t n_selected, + bool cache_f16, + uint32_t n_head, + uint32_t kv_lora_dim, + uint32_t qk_nope, + uint32_t qk_rope, + uint32_t n_ctx_orig, + float freq_base, + float freq_scale, + float ext_factor, + float attn_factor, + float beta_fast, + float beta_slow) { + if (!g_cublas_ready || !lora_out || !q || !qk_low || + !kv_lora_cache || !k_rope_cache || + n_tokens < 256u || n_selected == 0u || + n_head == 0u || kv_lora_dim == 0u || + qk_rope == 0u || (qk_rope & 1u) != 0u || + n_tokens > INT_MAX || n_selected > INT_MAX || + kv_lora_dim > INT_MAX || qk_rope > INT_MAX) { + return 0; + } + + uint64_t kv_count = 0, rope_count = 0; + uint64_t low_count = 0, qrope_count = 0; + uint64_t score_count = 0, head_count = 0; + if (!cuda_u64_mul_checked(n_selected, kv_lora_dim, &kv_count) || + !cuda_u64_mul_checked(n_selected, qk_rope, &rope_count) || + !cuda_u64_mul_checked(n_tokens, kv_lora_dim, &low_count) || + !cuda_u64_mul_checked(n_tokens, qk_rope, &qrope_count) || + !cuda_u64_mul_checked(n_tokens, n_selected, &score_count) || + !cuda_u64_mul_checked(n_tokens, kv_lora_dim, &head_count)) { + return 0; + } + + uint64_t kv_bytes = 0, rope_bytes = 0; + uint64_t low_bytes = 0, qrope_bytes = 0; + uint64_t score_bytes = 0, prob_bytes = 0, head_bytes = 0; + if (!cuda_u64_mul_checked(kv_count, sizeof(__half), &kv_bytes) || + !cuda_u64_mul_checked(rope_count, sizeof(__half), &rope_bytes) || + !cuda_u64_mul_checked(low_count, sizeof(__half), &low_bytes) || + !cuda_u64_mul_checked(qrope_count, sizeof(__half), &qrope_bytes) || + !cuda_u64_mul_checked(score_count, sizeof(float), &score_bytes) || + !cuda_u64_mul_checked(score_count, sizeof(__half), &prob_bytes) || + !cuda_u64_mul_checked(head_count, sizeof(float), &head_bytes)) { + return 0; + } + + uint64_t scratch_bytes = 0; + uint64_t kv_offset = 0, rope_offset = 0; + uint64_t low_offset = 0, qrope_offset = 0; + uint64_t score_offset = 0, prob_offset = 0, head_offset = 0; + if (!glm_causal_gemm_scratch_part(&scratch_bytes, kv_bytes, &kv_offset) || + !glm_causal_gemm_scratch_part(&scratch_bytes, rope_bytes, &rope_offset) || + !glm_causal_gemm_scratch_part(&scratch_bytes, low_bytes, &low_offset) || + !glm_causal_gemm_scratch_part(&scratch_bytes, qrope_bytes, &qrope_offset) || + !glm_causal_gemm_scratch_part(&scratch_bytes, score_bytes, &score_offset) || + !glm_causal_gemm_scratch_part(&scratch_bytes, prob_bytes, &prob_offset) || + !glm_causal_gemm_scratch_part(&scratch_bytes, head_bytes, &head_offset)) { + return 0; + } + + char *scratch = + (char *)cuda_tmp_alloc(scratch_bytes, "glm causal attention gemm"); + if (!scratch) return 0; + __half *kv_h = (__half *)(scratch + kv_offset); + __half *rope_h = (__half *)(scratch + rope_offset); + __half *low_h = (__half *)(scratch + low_offset); + __half *qrope_h = (__half *)(scratch + qrope_offset); + float *scores = (float *)(scratch + score_offset); + __half *probs = (__half *)(scratch + prob_offset); + float *head_out = (float *)(scratch + head_offset); + + glm_causal_gemm_cache_to_f16_kernel<<< + (kv_count + 255u) / 256u, 256>>>( + kv_h, + (const char *)kv_lora_cache->ptr, + kv_count, + cache_f16); + if (!cuda_ok(cudaGetLastError(), "glm causal attention kv f16 launch")) { + return 0; + } + const uint64_t rope_pairs = (uint64_t)n_selected * (qk_rope >> 1u); + glm_causal_gemm_rope_to_f16_kernel<<< + (rope_pairs + 255u) / 256u, 256>>>( + rope_h, + (const char *)k_rope_cache->ptr, + n_selected, + qk_rope, + cache_f16, + n_ctx_orig, + freq_base, + freq_scale, + ext_factor, + attn_factor, + beta_fast, + beta_slow); + if (!cuda_ok(cudaGetLastError(), "glm causal attention rope f16 launch")) { + return 0; + } + + const float one = 1.0f; + const float zero = 0.0f; + const float scale = 1.0f / sqrtf((float)(qk_nope + qk_rope)); + const uint32_t gather_blocks = (uint32_t)((low_count + 255u) / 256u); + const uint32_t scatter_blocks = + (uint32_t)((head_count + 255u) / 256u); + for (uint32_t head = 0; head < n_head; head++) { + glm_causal_gemm_gather_head_f16_kernel<<>>( + low_h, + qrope_h, + (const float *)qk_low->ptr, + (const float *)q->ptr, + n_tokens, + head, + n_head, + kv_lora_dim, + qk_nope, + qk_rope); + if (!cuda_ok(cudaGetLastError(), + "glm causal attention query gather launch")) { + return 0; + } + cublasStatus_t st = cublasGemmEx(g_cublas, + CUBLAS_OP_T, + CUBLAS_OP_N, + (int)n_selected, + (int)n_tokens, + (int)kv_lora_dim, + &one, + kv_h, + CUDA_R_16F, + (int)kv_lora_dim, + low_h, + CUDA_R_16F, + (int)kv_lora_dim, + &zero, + scores, + CUDA_R_32F, + (int)n_selected, + CUBLAS_COMPUTE_32F, + CUBLAS_GEMM_DEFAULT); + if (!cublas_ok(st, "glm causal attention lora score gemm")) { + return 0; + } + st = cublasGemmEx(g_cublas, + CUBLAS_OP_T, + CUBLAS_OP_N, + (int)n_selected, + (int)n_tokens, + (int)qk_rope, + &one, + rope_h, + CUDA_R_16F, + (int)qk_rope, + qrope_h, + CUDA_R_16F, + (int)qk_rope, + &one, + scores, + CUDA_R_32F, + (int)n_selected, + CUBLAS_COMPUTE_32F, + CUBLAS_GEMM_DEFAULT); + if (!cublas_ok(st, "glm causal attention rope score gemm")) { + return 0; + } + glm_causal_gemm_softmax_f16_kernel<<>>( + probs, scores, n_tokens, n_selected, pos0, scale); + if (!cuda_ok(cudaGetLastError(), + "glm causal attention softmax launch")) { + return 0; + } + st = cublasGemmEx(g_cublas, + CUBLAS_OP_N, + CUBLAS_OP_N, + (int)kv_lora_dim, + (int)n_tokens, + (int)n_selected, + &one, + kv_h, + CUDA_R_16F, + (int)kv_lora_dim, + probs, + CUDA_R_16F, + (int)n_selected, + &zero, + head_out, + CUDA_R_32F, + (int)kv_lora_dim, + CUBLAS_COMPUTE_32F, + CUBLAS_GEMM_DEFAULT); + if (!cublas_ok(st, "glm causal attention value gemm")) return 0; + glm_causal_gemm_scatter_head_kernel<<>>( + (float *)lora_out->ptr, + head_out, + n_tokens, + head, + n_head, + kv_lora_dim); + if (!cuda_ok(cudaGetLastError(), + "glm causal attention head scatter launch")) { + return 0; + } + } + return 1; +} + static int glm_attention_indexed_lora_launch( ds4_gpu_tensor *lora_out, const ds4_gpu_tensor *q, @@ -2044,6 +2418,45 @@ static int glm_attention_indexed_lora_launch( !glm_rocm_tensor_has_cache2(k_rope_cache, cache_cap, qk_rope, elem)) { return 0; } + /* Diagnostic port of the heterogeneous branch's dense causal prefill + * path. Keep it opt-in until remote gfx1151 timing and logit comparisons + * establish that the FP16 GEMMs are a safe default. */ + if (causal_range && !has_selected && + cuda_env_present(getenv("DS4_ROCM_GLM_CAUSAL_ATTN_GEMM")) && + glm_attention_indexed_lora_causal_gemm(lora_out, + q, + qk_low, + kv_lora_cache, + k_rope_cache, + n_tokens, + pos0, + n_selected, + cache_f16, + n_head, + kv_lora_dim, + qk_nope, + qk_rope, + n_ctx_orig, + freq_base, + freq_scale, + ext_factor, + attn_factor, + beta_fast, + beta_slow)) { + static int notice_printed = 0; + if (!notice_printed) { + fprintf(stderr, + DS4_GPU_LOG_PREFIX + "GLM causal indexed prefill using fp16 " + DS4_GPU_BLAS_NAME + " attention GEMMs (tokens=%u rows=%u cache=%s)\n", + n_tokens, + n_selected, + cache_f16 ? "f16" : "f32"); + notice_printed = 1; + } + return 1; + } dim3 grid(n_head, n_tokens, 1); const size_t shmem = ((size_t)256u + n_selected) * sizeof(float); glm_attention_indexed_lora_kernel<<>>((float *)lora_out->ptr, From c9e1013e710a5184a0c0f48d31cc1197c4495c04 Mon Sep 17 00:00:00 2001 From: Donato Capitella Date: Sat, 25 Jul 2026 19:14:45 +0100 Subject: [PATCH 15/33] Add optimized ROCm GLM slice decode path --- ds4.c | 24 ++++++++++++++++++++---- 1 file changed, 20 insertions(+), 4 deletions(-) diff --git a/ds4.c b/ds4.c index 14053f770..b767a7869 100644 --- a/ds4.c +++ b/ds4.c @@ -39157,6 +39157,13 @@ static bool glm_graph_alloc_slice( fprintf(stderr, "ds4: GLM graph using compact DSA KV only; expanded full-attention KV cache is skipped\n"); } +#ifdef DS4_ROCM_BUILD + if (glm_graph_env_truthy( + getenv("DS4_ROCM_GLM_LAYER_SLICE_TOKEN_DECODE"))) { + fprintf(stderr, + "ds4: ROCm GLM one-token layer slices use the optimized token graph\n"); + } +#endif if (g->compact_cache_cap != 0) { const uint64_t compact_kv_total = (uint64_t)g->layer_count * (compact_kv_lora_bytes + compact_k_rope_bytes); @@ -57572,6 +57579,13 @@ int ds4_session_eval_layer_slice(ds4_session *s, } const uint64_t hidden_dim = DS4_N_EMBD; +#ifdef DS4_ROCM_BUILD + const bool rocm_layer_slice_token_decode = + glm_graph_env_truthy( + getenv("DS4_ROCM_GLM_LAYER_SLICE_TOKEN_DECODE")); +#else + const bool rocm_layer_slice_token_decode = false; +#endif uint32_t done = 0; while (done < n_tokens) { const uint32_t pos = pos0 + done; @@ -57582,11 +57596,13 @@ int ds4_session_eval_layer_slice(ds4_session *s, bool ok = false; /* - * Raw decode starts from a token embedding. Continuation slices - * receive a hidden vector from the previous node, so use the - * batch-equivalent path even for a single token. + * The token graph accepts both embeddings and inter-node hidden + * states. Keep the resident ROCm continuation path opt-in until + * remote output and timing tests validate it across GLM quants. */ - if (remaining == 1 && pos > 0 && !input_hc && !output_hc) { + if (remaining == 1 && pos > 0 && + ((!input_hc && !output_hc) || + rocm_layer_slice_token_decode)) { float *chunk_logits = output_logits ? logits : NULL; ok = glm_graph_forward_token(g, &e->model, From 0e609ab85b467f6be9bc06510ad71ae9a751d88f Mon Sep 17 00:00:00 2001 From: Donato Capitella Date: Sat, 25 Jul 2026 19:48:43 +0100 Subject: [PATCH 16/33] Add F32 ROCm GLM causal attention GEMMs --- rocm/ds4_rocm_glm.cuh | 356 +++++++++++++++++++++++++++++++++++++++++- 1 file changed, 354 insertions(+), 2 deletions(-) diff --git a/rocm/ds4_rocm_glm.cuh b/rocm/ds4_rocm_glm.cuh index 881405fa5..f332ae8fc 100644 --- a/rocm/ds4_rocm_glm.cuh +++ b/rocm/ds4_rocm_glm.cuh @@ -2133,6 +2133,122 @@ __global__ static void glm_causal_gemm_scatter_head_kernel( head_out[i]; } +__global__ static void glm_causal_gemm_rope_to_f32_kernel( + float *dst, + const char *src, + uint32_t n_rows, + uint32_t qk_rope, + uint32_t n_ctx_orig, + float freq_base, + float freq_scale, + float ext_factor, + float attn_factor, + float beta_fast, + float beta_slow) { + const uint64_t pair = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x; + const uint64_t pairs = (uint64_t)n_rows * (qk_rope >> 1u); + if (pair >= pairs) return; + const uint32_t row = (uint32_t)(pair / (qk_rope >> 1u)); + const uint32_t r = + (uint32_t)(pair - (uint64_t)row * (qk_rope >> 1u)) << 1u; + const float2 y = glm_rocm_rotated_cache_rope_pair(src, + (uint64_t)row * qk_rope, + r, + row, + qk_rope, + false, + n_ctx_orig, + freq_base, + freq_scale, + ext_factor, + attn_factor, + beta_fast, + beta_slow); + dst[(uint64_t)row * qk_rope + r] = y.x; + dst[(uint64_t)row * qk_rope + r + 1u] = y.y; +} + +__global__ static void glm_causal_gemm_gather_head_f32_kernel( + float *low_h, + float *qrope_h, + const float *low, + const float *q, + uint32_t n_tokens, + uint32_t head, + uint32_t n_head, + uint32_t kv_lora_dim, + uint32_t qk_nope, + uint32_t qk_rope) { + const uint64_t total = (uint64_t)n_tokens * kv_lora_dim; + for (uint64_t i = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x; + i < total; + i += (uint64_t)gridDim.x * blockDim.x) { + const uint32_t token = (uint32_t)(i / kv_lora_dim); + const uint32_t j = + (uint32_t)(i - (uint64_t)token * kv_lora_dim); + low_h[i] = + low[((uint64_t)token * n_head + head) * kv_lora_dim + j]; + } + const uint64_t rope_total = (uint64_t)n_tokens * qk_rope; + for (uint64_t i = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x; + i < rope_total; + i += (uint64_t)gridDim.x * blockDim.x) { + const uint32_t token = (uint32_t)(i / qk_rope); + const uint32_t r = (uint32_t)(i - (uint64_t)token * qk_rope); + const uint32_t qk_dim = qk_nope + qk_rope; + qrope_h[i] = + q[((uint64_t)token * n_head + head) * qk_dim + qk_nope + r]; + } +} + +__global__ static void glm_causal_gemm_softmax_f32_kernel( + float *probs, + float *scores, + uint32_t n_tokens, + uint32_t n_selected, + uint32_t pos0, + float scale) { + const uint32_t token = blockIdx.x; + if (token >= n_tokens) return; + float *row = scores + (uint64_t)token * n_selected; + __shared__ float reduce[256]; + const uint32_t visible = min(n_selected, pos0 + token + 1u); + float vmax = -INFINITY; + for (uint32_t s = threadIdx.x; s < n_selected; s += blockDim.x) { + const float v = s < visible ? row[s] * scale : -INFINITY; + row[s] = v; + vmax = fmaxf(vmax, v); + } + reduce[threadIdx.x] = vmax; + __syncthreads(); + for (uint32_t stride = blockDim.x >> 1u; stride != 0u; stride >>= 1u) { + if (threadIdx.x < stride) { + reduce[threadIdx.x] = + fmaxf(reduce[threadIdx.x], reduce[threadIdx.x + stride]); + } + __syncthreads(); + } + const float max_score = reduce[0]; + float sum = 0.0f; + for (uint32_t s = threadIdx.x; s < n_selected; s += blockDim.x) { + const float v = s < visible ? expf(row[s] - max_score) : 0.0f; + probs[(uint64_t)token * n_selected + s] = v; + sum += v; + } + reduce[threadIdx.x] = sum; + __syncthreads(); + for (uint32_t stride = blockDim.x >> 1u; stride != 0u; stride >>= 1u) { + if (threadIdx.x < stride) { + reduce[threadIdx.x] += reduce[threadIdx.x + stride]; + } + __syncthreads(); + } + const float inv = 1.0f / fmaxf(reduce[0], 1.0e-20f); + for (uint32_t s = threadIdx.x; s < n_selected; s += blockDim.x) { + probs[(uint64_t)token * n_selected + s] *= inv; + } +} + static int glm_causal_gemm_scratch_part( uint64_t *offset, uint64_t bytes, @@ -2148,6 +2264,201 @@ static int glm_causal_gemm_scratch_part( return 1; } +static int glm_attention_indexed_lora_causal_gemm_f32( + ds4_gpu_tensor *lora_out, + const ds4_gpu_tensor *q, + const ds4_gpu_tensor *qk_low, + const ds4_gpu_tensor *kv_lora_cache, + const ds4_gpu_tensor *k_rope_cache, + uint32_t n_tokens, + uint32_t pos0, + uint32_t n_selected, + bool cache_f16, + uint32_t n_head, + uint32_t kv_lora_dim, + uint32_t qk_nope, + uint32_t qk_rope, + uint32_t n_ctx_orig, + float freq_base, + float freq_scale, + float ext_factor, + float attn_factor, + float beta_fast, + float beta_slow) { + if (!g_cublas_ready || !lora_out || !q || !qk_low || + !kv_lora_cache || !k_rope_cache || cache_f16 || + n_tokens < 256u || n_selected == 0u || + n_head == 0u || kv_lora_dim == 0u || + qk_rope == 0u || (qk_rope & 1u) != 0u || + n_tokens > INT_MAX || n_selected > INT_MAX || + kv_lora_dim > INT_MAX || qk_rope > INT_MAX) { + return 0; + } + + uint64_t rope_count = 0, low_count = 0, qrope_count = 0; + uint64_t score_count = 0, head_count = 0; + if (!cuda_u64_mul_checked(n_selected, qk_rope, &rope_count) || + !cuda_u64_mul_checked(n_tokens, kv_lora_dim, &low_count) || + !cuda_u64_mul_checked(n_tokens, qk_rope, &qrope_count) || + !cuda_u64_mul_checked(n_tokens, n_selected, &score_count) || + !cuda_u64_mul_checked(n_tokens, kv_lora_dim, &head_count)) { + return 0; + } + + uint64_t rope_bytes = 0, low_bytes = 0, qrope_bytes = 0; + uint64_t score_bytes = 0, prob_bytes = 0, head_bytes = 0; + if (!cuda_u64_mul_checked(rope_count, sizeof(float), &rope_bytes) || + !cuda_u64_mul_checked(low_count, sizeof(float), &low_bytes) || + !cuda_u64_mul_checked(qrope_count, sizeof(float), &qrope_bytes) || + !cuda_u64_mul_checked(score_count, sizeof(float), &score_bytes) || + !cuda_u64_mul_checked(score_count, sizeof(float), &prob_bytes) || + !cuda_u64_mul_checked(head_count, sizeof(float), &head_bytes)) { + return 0; + } + + uint64_t scratch_bytes = 0; + uint64_t rope_offset = 0, low_offset = 0, qrope_offset = 0; + uint64_t score_offset = 0, prob_offset = 0, head_offset = 0; + if (!glm_causal_gemm_scratch_part( + &scratch_bytes, rope_bytes, &rope_offset) || + !glm_causal_gemm_scratch_part( + &scratch_bytes, low_bytes, &low_offset) || + !glm_causal_gemm_scratch_part( + &scratch_bytes, qrope_bytes, &qrope_offset) || + !glm_causal_gemm_scratch_part( + &scratch_bytes, score_bytes, &score_offset) || + !glm_causal_gemm_scratch_part( + &scratch_bytes, prob_bytes, &prob_offset) || + !glm_causal_gemm_scratch_part( + &scratch_bytes, head_bytes, &head_offset)) { + return 0; + } + + char *scratch = + (char *)cuda_tmp_alloc(scratch_bytes, "glm causal attention f32 gemm"); + if (!scratch) return 0; + float *rope_f = (float *)(scratch + rope_offset); + float *low_f = (float *)(scratch + low_offset); + float *qrope_f = (float *)(scratch + qrope_offset); + float *scores = (float *)(scratch + score_offset); + float *probs = (float *)(scratch + prob_offset); + float *head_out = (float *)(scratch + head_offset); + + const uint64_t rope_pairs = + (uint64_t)n_selected * (qk_rope >> 1u); + glm_causal_gemm_rope_to_f32_kernel<<< + (rope_pairs + 255u) / 256u, 256>>>( + rope_f, + (const char *)k_rope_cache->ptr, + n_selected, + qk_rope, + n_ctx_orig, + freq_base, + freq_scale, + ext_factor, + attn_factor, + beta_fast, + beta_slow); + if (!cuda_ok(cudaGetLastError(), "glm causal attention rope f32 launch")) { + return 0; + } + + const float one = 1.0f; + const float zero = 0.0f; + const float scale = 1.0f / sqrtf((float)(qk_nope + qk_rope)); + const uint32_t gather_blocks = + (uint32_t)((low_count + 255u) / 256u); + const uint32_t scatter_blocks = + (uint32_t)((head_count + 255u) / 256u); + const float *kv_f = (const float *)kv_lora_cache->ptr; + for (uint32_t head = 0; head < n_head; head++) { + glm_causal_gemm_gather_head_f32_kernel<<>>( + low_f, + qrope_f, + (const float *)qk_low->ptr, + (const float *)q->ptr, + n_tokens, + head, + n_head, + kv_lora_dim, + qk_nope, + qk_rope); + if (!cuda_ok(cudaGetLastError(), + "glm causal attention query f32 gather launch")) { + return 0; + } + cublasStatus_t st = cublasSgemm(g_cublas, + CUBLAS_OP_T, + CUBLAS_OP_N, + (int)n_selected, + (int)n_tokens, + (int)kv_lora_dim, + &one, + kv_f, + (int)kv_lora_dim, + low_f, + (int)kv_lora_dim, + &zero, + scores, + (int)n_selected); + if (!cublas_ok(st, "glm causal attention lora score f32 gemm")) { + return 0; + } + st = cublasSgemm(g_cublas, + CUBLAS_OP_T, + CUBLAS_OP_N, + (int)n_selected, + (int)n_tokens, + (int)qk_rope, + &one, + rope_f, + (int)qk_rope, + qrope_f, + (int)qk_rope, + &one, + scores, + (int)n_selected); + if (!cublas_ok(st, "glm causal attention rope score f32 gemm")) { + return 0; + } + glm_causal_gemm_softmax_f32_kernel<<>>( + probs, scores, n_tokens, n_selected, pos0, scale); + if (!cuda_ok(cudaGetLastError(), + "glm causal attention softmax f32 launch")) { + return 0; + } + st = cublasSgemm(g_cublas, + CUBLAS_OP_N, + CUBLAS_OP_N, + (int)kv_lora_dim, + (int)n_tokens, + (int)n_selected, + &one, + kv_f, + (int)kv_lora_dim, + probs, + (int)n_selected, + &zero, + head_out, + (int)kv_lora_dim); + if (!cublas_ok(st, "glm causal attention value f32 gemm")) { + return 0; + } + glm_causal_gemm_scatter_head_kernel<<>>( + (float *)lora_out->ptr, + head_out, + n_tokens, + head, + n_head, + kv_lora_dim); + if (!cuda_ok(cudaGetLastError(), + "glm causal attention f32 head scatter launch")) { + return 0; + } + } + return 1; +} + static int glm_attention_indexed_lora_causal_gemm( ds4_gpu_tensor *lora_out, const ds4_gpu_tensor *q, @@ -2418,9 +2729,50 @@ static int glm_attention_indexed_lora_launch( !glm_rocm_tensor_has_cache2(k_rope_cache, cache_cap, qk_rope, elem)) { return 0; } + /* + * The F32 path keeps the original cache, query, softmax, and value + * precision while replacing the scalar attention loops with BLAS GEMMs. + * Keep it separate from the reference branch's faster but lossy FP16 + * experiment so remote quality tests can isolate precision from layout. + */ + if (causal_range && !has_selected && + cuda_env_present(getenv("DS4_ROCM_GLM_CAUSAL_ATTN_GEMM_F32")) && + glm_attention_indexed_lora_causal_gemm_f32(lora_out, + q, + qk_low, + kv_lora_cache, + k_rope_cache, + n_tokens, + pos0, + n_selected, + cache_f16, + n_head, + kv_lora_dim, + qk_nope, + qk_rope, + n_ctx_orig, + freq_base, + freq_scale, + ext_factor, + attn_factor, + beta_fast, + beta_slow)) { + static int f32_notice_printed = 0; + if (!f32_notice_printed) { + fprintf(stderr, + DS4_GPU_LOG_PREFIX + "GLM causal indexed prefill using f32 " + DS4_GPU_BLAS_NAME + " attention GEMMs (tokens=%u rows=%u cache=f32)\n", + n_tokens, + n_selected); + f32_notice_printed = 1; + } + return 1; + } /* Diagnostic port of the heterogeneous branch's dense causal prefill - * path. Keep it opt-in until remote gfx1151 timing and logit comparisons - * establish that the FP16 GEMMs are a safe default. */ + * path. The FP16 output comparison failed on the resident IQ2 GLM model, + * so retain it only as an explicit approximation experiment. */ if (causal_range && !has_selected && cuda_env_present(getenv("DS4_ROCM_GLM_CAUSAL_ATTN_GEMM")) && glm_attention_indexed_lora_causal_gemm(lora_out, From b7f5cddc61015c43b6605afe2fcfbc37fdc864f7 Mon Sep 17 00:00:00 2001 From: Donato Capitella Date: Sat, 25 Jul 2026 21:05:31 +0100 Subject: [PATCH 17/33] remove FP32 hipblas path --- rocm/ds4_rocm_glm.cuh | 356 +----------------------------------------- 1 file changed, 2 insertions(+), 354 deletions(-) diff --git a/rocm/ds4_rocm_glm.cuh b/rocm/ds4_rocm_glm.cuh index f332ae8fc..881405fa5 100644 --- a/rocm/ds4_rocm_glm.cuh +++ b/rocm/ds4_rocm_glm.cuh @@ -2133,122 +2133,6 @@ __global__ static void glm_causal_gemm_scatter_head_kernel( head_out[i]; } -__global__ static void glm_causal_gemm_rope_to_f32_kernel( - float *dst, - const char *src, - uint32_t n_rows, - uint32_t qk_rope, - uint32_t n_ctx_orig, - float freq_base, - float freq_scale, - float ext_factor, - float attn_factor, - float beta_fast, - float beta_slow) { - const uint64_t pair = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x; - const uint64_t pairs = (uint64_t)n_rows * (qk_rope >> 1u); - if (pair >= pairs) return; - const uint32_t row = (uint32_t)(pair / (qk_rope >> 1u)); - const uint32_t r = - (uint32_t)(pair - (uint64_t)row * (qk_rope >> 1u)) << 1u; - const float2 y = glm_rocm_rotated_cache_rope_pair(src, - (uint64_t)row * qk_rope, - r, - row, - qk_rope, - false, - n_ctx_orig, - freq_base, - freq_scale, - ext_factor, - attn_factor, - beta_fast, - beta_slow); - dst[(uint64_t)row * qk_rope + r] = y.x; - dst[(uint64_t)row * qk_rope + r + 1u] = y.y; -} - -__global__ static void glm_causal_gemm_gather_head_f32_kernel( - float *low_h, - float *qrope_h, - const float *low, - const float *q, - uint32_t n_tokens, - uint32_t head, - uint32_t n_head, - uint32_t kv_lora_dim, - uint32_t qk_nope, - uint32_t qk_rope) { - const uint64_t total = (uint64_t)n_tokens * kv_lora_dim; - for (uint64_t i = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x; - i < total; - i += (uint64_t)gridDim.x * blockDim.x) { - const uint32_t token = (uint32_t)(i / kv_lora_dim); - const uint32_t j = - (uint32_t)(i - (uint64_t)token * kv_lora_dim); - low_h[i] = - low[((uint64_t)token * n_head + head) * kv_lora_dim + j]; - } - const uint64_t rope_total = (uint64_t)n_tokens * qk_rope; - for (uint64_t i = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x; - i < rope_total; - i += (uint64_t)gridDim.x * blockDim.x) { - const uint32_t token = (uint32_t)(i / qk_rope); - const uint32_t r = (uint32_t)(i - (uint64_t)token * qk_rope); - const uint32_t qk_dim = qk_nope + qk_rope; - qrope_h[i] = - q[((uint64_t)token * n_head + head) * qk_dim + qk_nope + r]; - } -} - -__global__ static void glm_causal_gemm_softmax_f32_kernel( - float *probs, - float *scores, - uint32_t n_tokens, - uint32_t n_selected, - uint32_t pos0, - float scale) { - const uint32_t token = blockIdx.x; - if (token >= n_tokens) return; - float *row = scores + (uint64_t)token * n_selected; - __shared__ float reduce[256]; - const uint32_t visible = min(n_selected, pos0 + token + 1u); - float vmax = -INFINITY; - for (uint32_t s = threadIdx.x; s < n_selected; s += blockDim.x) { - const float v = s < visible ? row[s] * scale : -INFINITY; - row[s] = v; - vmax = fmaxf(vmax, v); - } - reduce[threadIdx.x] = vmax; - __syncthreads(); - for (uint32_t stride = blockDim.x >> 1u; stride != 0u; stride >>= 1u) { - if (threadIdx.x < stride) { - reduce[threadIdx.x] = - fmaxf(reduce[threadIdx.x], reduce[threadIdx.x + stride]); - } - __syncthreads(); - } - const float max_score = reduce[0]; - float sum = 0.0f; - for (uint32_t s = threadIdx.x; s < n_selected; s += blockDim.x) { - const float v = s < visible ? expf(row[s] - max_score) : 0.0f; - probs[(uint64_t)token * n_selected + s] = v; - sum += v; - } - reduce[threadIdx.x] = sum; - __syncthreads(); - for (uint32_t stride = blockDim.x >> 1u; stride != 0u; stride >>= 1u) { - if (threadIdx.x < stride) { - reduce[threadIdx.x] += reduce[threadIdx.x + stride]; - } - __syncthreads(); - } - const float inv = 1.0f / fmaxf(reduce[0], 1.0e-20f); - for (uint32_t s = threadIdx.x; s < n_selected; s += blockDim.x) { - probs[(uint64_t)token * n_selected + s] *= inv; - } -} - static int glm_causal_gemm_scratch_part( uint64_t *offset, uint64_t bytes, @@ -2264,201 +2148,6 @@ static int glm_causal_gemm_scratch_part( return 1; } -static int glm_attention_indexed_lora_causal_gemm_f32( - ds4_gpu_tensor *lora_out, - const ds4_gpu_tensor *q, - const ds4_gpu_tensor *qk_low, - const ds4_gpu_tensor *kv_lora_cache, - const ds4_gpu_tensor *k_rope_cache, - uint32_t n_tokens, - uint32_t pos0, - uint32_t n_selected, - bool cache_f16, - uint32_t n_head, - uint32_t kv_lora_dim, - uint32_t qk_nope, - uint32_t qk_rope, - uint32_t n_ctx_orig, - float freq_base, - float freq_scale, - float ext_factor, - float attn_factor, - float beta_fast, - float beta_slow) { - if (!g_cublas_ready || !lora_out || !q || !qk_low || - !kv_lora_cache || !k_rope_cache || cache_f16 || - n_tokens < 256u || n_selected == 0u || - n_head == 0u || kv_lora_dim == 0u || - qk_rope == 0u || (qk_rope & 1u) != 0u || - n_tokens > INT_MAX || n_selected > INT_MAX || - kv_lora_dim > INT_MAX || qk_rope > INT_MAX) { - return 0; - } - - uint64_t rope_count = 0, low_count = 0, qrope_count = 0; - uint64_t score_count = 0, head_count = 0; - if (!cuda_u64_mul_checked(n_selected, qk_rope, &rope_count) || - !cuda_u64_mul_checked(n_tokens, kv_lora_dim, &low_count) || - !cuda_u64_mul_checked(n_tokens, qk_rope, &qrope_count) || - !cuda_u64_mul_checked(n_tokens, n_selected, &score_count) || - !cuda_u64_mul_checked(n_tokens, kv_lora_dim, &head_count)) { - return 0; - } - - uint64_t rope_bytes = 0, low_bytes = 0, qrope_bytes = 0; - uint64_t score_bytes = 0, prob_bytes = 0, head_bytes = 0; - if (!cuda_u64_mul_checked(rope_count, sizeof(float), &rope_bytes) || - !cuda_u64_mul_checked(low_count, sizeof(float), &low_bytes) || - !cuda_u64_mul_checked(qrope_count, sizeof(float), &qrope_bytes) || - !cuda_u64_mul_checked(score_count, sizeof(float), &score_bytes) || - !cuda_u64_mul_checked(score_count, sizeof(float), &prob_bytes) || - !cuda_u64_mul_checked(head_count, sizeof(float), &head_bytes)) { - return 0; - } - - uint64_t scratch_bytes = 0; - uint64_t rope_offset = 0, low_offset = 0, qrope_offset = 0; - uint64_t score_offset = 0, prob_offset = 0, head_offset = 0; - if (!glm_causal_gemm_scratch_part( - &scratch_bytes, rope_bytes, &rope_offset) || - !glm_causal_gemm_scratch_part( - &scratch_bytes, low_bytes, &low_offset) || - !glm_causal_gemm_scratch_part( - &scratch_bytes, qrope_bytes, &qrope_offset) || - !glm_causal_gemm_scratch_part( - &scratch_bytes, score_bytes, &score_offset) || - !glm_causal_gemm_scratch_part( - &scratch_bytes, prob_bytes, &prob_offset) || - !glm_causal_gemm_scratch_part( - &scratch_bytes, head_bytes, &head_offset)) { - return 0; - } - - char *scratch = - (char *)cuda_tmp_alloc(scratch_bytes, "glm causal attention f32 gemm"); - if (!scratch) return 0; - float *rope_f = (float *)(scratch + rope_offset); - float *low_f = (float *)(scratch + low_offset); - float *qrope_f = (float *)(scratch + qrope_offset); - float *scores = (float *)(scratch + score_offset); - float *probs = (float *)(scratch + prob_offset); - float *head_out = (float *)(scratch + head_offset); - - const uint64_t rope_pairs = - (uint64_t)n_selected * (qk_rope >> 1u); - glm_causal_gemm_rope_to_f32_kernel<<< - (rope_pairs + 255u) / 256u, 256>>>( - rope_f, - (const char *)k_rope_cache->ptr, - n_selected, - qk_rope, - n_ctx_orig, - freq_base, - freq_scale, - ext_factor, - attn_factor, - beta_fast, - beta_slow); - if (!cuda_ok(cudaGetLastError(), "glm causal attention rope f32 launch")) { - return 0; - } - - const float one = 1.0f; - const float zero = 0.0f; - const float scale = 1.0f / sqrtf((float)(qk_nope + qk_rope)); - const uint32_t gather_blocks = - (uint32_t)((low_count + 255u) / 256u); - const uint32_t scatter_blocks = - (uint32_t)((head_count + 255u) / 256u); - const float *kv_f = (const float *)kv_lora_cache->ptr; - for (uint32_t head = 0; head < n_head; head++) { - glm_causal_gemm_gather_head_f32_kernel<<>>( - low_f, - qrope_f, - (const float *)qk_low->ptr, - (const float *)q->ptr, - n_tokens, - head, - n_head, - kv_lora_dim, - qk_nope, - qk_rope); - if (!cuda_ok(cudaGetLastError(), - "glm causal attention query f32 gather launch")) { - return 0; - } - cublasStatus_t st = cublasSgemm(g_cublas, - CUBLAS_OP_T, - CUBLAS_OP_N, - (int)n_selected, - (int)n_tokens, - (int)kv_lora_dim, - &one, - kv_f, - (int)kv_lora_dim, - low_f, - (int)kv_lora_dim, - &zero, - scores, - (int)n_selected); - if (!cublas_ok(st, "glm causal attention lora score f32 gemm")) { - return 0; - } - st = cublasSgemm(g_cublas, - CUBLAS_OP_T, - CUBLAS_OP_N, - (int)n_selected, - (int)n_tokens, - (int)qk_rope, - &one, - rope_f, - (int)qk_rope, - qrope_f, - (int)qk_rope, - &one, - scores, - (int)n_selected); - if (!cublas_ok(st, "glm causal attention rope score f32 gemm")) { - return 0; - } - glm_causal_gemm_softmax_f32_kernel<<>>( - probs, scores, n_tokens, n_selected, pos0, scale); - if (!cuda_ok(cudaGetLastError(), - "glm causal attention softmax f32 launch")) { - return 0; - } - st = cublasSgemm(g_cublas, - CUBLAS_OP_N, - CUBLAS_OP_N, - (int)kv_lora_dim, - (int)n_tokens, - (int)n_selected, - &one, - kv_f, - (int)kv_lora_dim, - probs, - (int)n_selected, - &zero, - head_out, - (int)kv_lora_dim); - if (!cublas_ok(st, "glm causal attention value f32 gemm")) { - return 0; - } - glm_causal_gemm_scatter_head_kernel<<>>( - (float *)lora_out->ptr, - head_out, - n_tokens, - head, - n_head, - kv_lora_dim); - if (!cuda_ok(cudaGetLastError(), - "glm causal attention f32 head scatter launch")) { - return 0; - } - } - return 1; -} - static int glm_attention_indexed_lora_causal_gemm( ds4_gpu_tensor *lora_out, const ds4_gpu_tensor *q, @@ -2729,50 +2418,9 @@ static int glm_attention_indexed_lora_launch( !glm_rocm_tensor_has_cache2(k_rope_cache, cache_cap, qk_rope, elem)) { return 0; } - /* - * The F32 path keeps the original cache, query, softmax, and value - * precision while replacing the scalar attention loops with BLAS GEMMs. - * Keep it separate from the reference branch's faster but lossy FP16 - * experiment so remote quality tests can isolate precision from layout. - */ - if (causal_range && !has_selected && - cuda_env_present(getenv("DS4_ROCM_GLM_CAUSAL_ATTN_GEMM_F32")) && - glm_attention_indexed_lora_causal_gemm_f32(lora_out, - q, - qk_low, - kv_lora_cache, - k_rope_cache, - n_tokens, - pos0, - n_selected, - cache_f16, - n_head, - kv_lora_dim, - qk_nope, - qk_rope, - n_ctx_orig, - freq_base, - freq_scale, - ext_factor, - attn_factor, - beta_fast, - beta_slow)) { - static int f32_notice_printed = 0; - if (!f32_notice_printed) { - fprintf(stderr, - DS4_GPU_LOG_PREFIX - "GLM causal indexed prefill using f32 " - DS4_GPU_BLAS_NAME - " attention GEMMs (tokens=%u rows=%u cache=f32)\n", - n_tokens, - n_selected); - f32_notice_printed = 1; - } - return 1; - } /* Diagnostic port of the heterogeneous branch's dense causal prefill - * path. The FP16 output comparison failed on the resident IQ2 GLM model, - * so retain it only as an explicit approximation experiment. */ + * path. Keep it opt-in until remote gfx1151 timing and logit comparisons + * establish that the FP16 GEMMs are a safe default. */ if (causal_range && !has_selected && cuda_env_present(getenv("DS4_ROCM_GLM_CAUSAL_ATTN_GEMM")) && glm_attention_indexed_lora_causal_gemm(lora_out, From 312934c14a6f582f23d64f36b044dee890e2dcdc Mon Sep 17 00:00:00 2001 From: Donato Capitella Date: Sat, 25 Jul 2026 21:20:49 +0100 Subject: [PATCH 18/33] Add opt-in ROCm GLM F16 compact cache --- ds4.c | 18 +- ds4_rocm_compat.cu | 23 ++ speed-bench/README.md | 5 + .../ROCM_GLM_DISTRIBUTED_VALIDATION.md | 374 ++++++++++++++++++ speed-bench/compare_glm_validation.py | 172 ++++++++ 5 files changed, 587 insertions(+), 5 deletions(-) create mode 100644 speed-bench/ROCM_GLM_DISTRIBUTED_VALIDATION.md create mode 100644 speed-bench/compare_glm_validation.py diff --git a/ds4.c b/ds4.c index b767a7869..28ff51aa5 100644 --- a/ds4.c +++ b/ds4.c @@ -14888,7 +14888,8 @@ static void print_vec_stats(const char *name, const float *x, uint64_t n) { #define DS4_GPU_ATTN_COMP_CACHE_F16 0 #endif -#define DS4_GPU_GLM_COMPACT_CACHE_F16 DS4_GPU_ATTN_COMP_CACHE_F16 +/* GLM's DSA compact cache has a ROCm runtime opt-in below; it is independent + * from the general compressed-attention cache format. */ /* ========================================================================= * Metal Release Graph State. @@ -34466,12 +34467,19 @@ static uint32_t glm_graph_indexed_prefill_score_tokens( uint32_t indexed_prefill_cap, uint32_t compact_cap); -static uint64_t glm_graph_compact_cache_elem_bytes(void) { - return DS4_GPU_GLM_COMPACT_CACHE_F16 ? sizeof(uint16_t) : sizeof(float); +static uint32_t glm_graph_compact_cache_is_f16(void) { +#if defined(__APPLE__) + return 1u; +#elif defined(DS4_ROCM_BUILD) + return getenv("DS4_ROCM_GLM_COMPACT_CACHE_F16") != NULL ? 1u : 0u; +#else + return 0u; +#endif } -static uint32_t glm_graph_compact_cache_is_f16(void) { - return DS4_GPU_GLM_COMPACT_CACHE_F16 ? 1u : 0u; +static uint64_t glm_graph_compact_cache_elem_bytes(void) { + return glm_graph_compact_cache_is_f16() ? + sizeof(uint16_t) : sizeof(float); } static bool glm_graph_expanded_kv_cache_enabled(bool ssd_streaming) { diff --git a/ds4_rocm_compat.cu b/ds4_rocm_compat.cu index 79b587e3b..ac8d485f6 100644 --- a/ds4_rocm_compat.cu +++ b/ds4_rocm_compat.cu @@ -368,6 +368,29 @@ extern "C" int ds4_gpu_glm_attention_indexed_decode_typed_tensor( beta_fast, beta_slow); } +extern "C" int ds4_gpu_glm_attention_indexed_decode_split_group8_typed_tensor( + ds4_gpu_tensor *heads, ds4_gpu_tensor *partial_lora, + ds4_gpu_tensor *partial_ms, const ds4_gpu_tensor *q, + const ds4_gpu_tensor *qk_low, const ds4_gpu_tensor *kv_lora_cache, + const ds4_gpu_tensor *k_rope_cache, const void *model_map, + uint64_t model_size, uint64_t value_weight_offset, + uint32_t value_weight_type, const ds4_gpu_tensor *selected, + uint32_t n_selected, bool selected_rows_valid, uint32_t cache_cap, + bool cache_f16, uint32_t n_head, uint32_t kv_lora_dim, + uint32_t qk_nope, uint32_t qk_rope, uint32_t value_dim, + uint32_t n_ctx_orig, uint32_t block_rows, uint32_t n_blocks, + float freq_base, float freq_scale, float ext_factor, + float attn_factor, float beta_fast, float beta_slow) { + if (value_weight_type != 8u) return 0; + return ds4_gpu_glm_attention_indexed_decode_split_group8_tensor( + heads, partial_lora, partial_ms, q, qk_low, kv_lora_cache, + k_rope_cache, model_map, model_size, value_weight_offset, + selected, n_selected, selected_rows_valid, cache_cap, cache_f16, + n_head, kv_lora_dim, qk_nope, qk_rope, value_dim, n_ctx_orig, + block_rows, n_blocks, freq_base, freq_scale, ext_factor, + attn_factor, beta_fast, beta_slow); +} + extern "C" int ds4_gpu_glm_attention_indexed_batch_typed_tensor( ds4_gpu_tensor *heads, const ds4_gpu_tensor *q, const ds4_gpu_tensor *qk_low, const ds4_gpu_tensor *kv_lora_cache, diff --git a/speed-bench/README.md b/speed-bench/README.md index 32075fe18..b9f7aef09 100644 --- a/speed-bench/README.md +++ b/speed-bench/README.md @@ -26,3 +26,8 @@ python3 speed-bench/plot_speed.py speed-bench/m3_max.csv --title "M3 Max t/s" The script uses only the Python standard library. By default it writes a file next to the CSV using the `_ts.svg` suffix, such as `speed-bench/m3_max_ts.svg`. + +For ROCm GLM distributed optimization work, follow +[`ROCM_GLM_DISTRIBUTED_VALIDATION.md`](ROCM_GLM_DISTRIBUTED_VALIDATION.md). +It defines the fixed baseline/candidate protocol, two-node commands, graph +tensor and frontier-logit comparisons, API smoke test, and decision record. diff --git a/speed-bench/ROCM_GLM_DISTRIBUTED_VALIDATION.md b/speed-bench/ROCM_GLM_DISTRIBUTED_VALIDATION.md new file mode 100644 index 000000000..b3588043a --- /dev/null +++ b/speed-bench/ROCM_GLM_DISTRIBUTED_VALIDATION.md @@ -0,0 +1,374 @@ +# ROCm GLM Distributed Change Validation + +Use this playbook for every ROCm GLM distributed performance change. It keeps +speed measurements separate from numerical and model-quality checks, and it +changes one variable at a time. + +The commands below describe the current two-Strix-Halo test topology: + +- coordinator: `192.168.100.2`, layers `0:37` +- worker: `192.168.100.1`, layers `38:output` +- model: `GLM-5.2-UD-IQ2_XXS_RoutedIQ2XXS_blk78Q2K.gguf` +- image: `docker.io/kyuz0/strix-halo-ds4-toolbox:glm-rocm-7.2.4` +- distributed prefill: 256-token chunks, window 2 + +Change those constants deliberately if the hardware or model changes, and +record the new values with the results. + +Unless a command is explicitly labeled for the worker, run it on the +coordinator. The coordinator checkout used below is `~/ds4/ds4`. + +## Rules + +1. Build baseline and candidate from the same commit and image. +2. Keep the model file, SHA-256, layer split, prompt, context allocation, + chunk/window, generation length, power state, and background load fixed. +3. Hold previously accepted optimizations constant. Toggle only the change + under test. +4. Save worker and coordinator logs plus benchmark CSV files. +5. Test performance and numerical closeness separately. A coherent sampled + answer is not a numerical comparison, and a matching argmax is not enough. +6. Run the baseline and candidate at least once as a smoke test. Before + declaring a small improvement real or making a path default, repeat the + complete pair. Treat differences below roughly 2% as noise until repeated. +7. Do not compare `ds4-bench --show-output` text with chat answers. The former + is a raw continuation of the prompt file; use the server API for chat. + +Record the code revision and model hash before each series: + +```sh +git -C ~/ds4/ds4 rev-parse HEAD +sha256sum ~/ds4/GLM-5.2-UD-IQ2_XXS_RoutedIQ2XXS_blk78Q2K.gguf +``` + +## Build gate + +The ROCm/gfx1151 compile check does not require a GPU: + +```sh +podman run --rm --userns=keep-id \ + -v ~/ds4/ds4:/workspace:Z \ + -w /workspace \ + docker.io/rocm/dev-ubuntu-24.04:7.2.1-complete \ + make strix-halo ROCM_ARCH=gfx1151 +``` + +Require a successful warning-free build and run `git diff --check`. + +## Performance protocol + +The standard development point is a fresh 2048-token prefill followed by 128 +greedy generation tokens. Both nodes keep the accepted FP16 causal-attention +GEMM enabled. The candidate additionally enables the change under test. + +For the F16 compact-cache experiment, the isolated variable is: + +```text +baseline: DS4_ROCM_GLM_COMPACT_CACHE_F16 unset +candidate: DS4_ROCM_GLM_COMPACT_CACHE_F16=1 +``` + +### Baseline worker + +Run on `192.168.100.1`: + +```sh +podman run --rm -it \ + --name ds4-worker-cache-f32 \ + -e DS4_GLM_MEMORY_GUARD_RESERVE_GB=8 \ + -e DS4_ROCM_GLM_CAUSAL_ATTN_GEMM=1 \ + --device /dev/dri \ + --device /dev/kfd \ + --group-add keep-groups \ + --security-opt seccomp=unconfined \ + --ipc=host \ + --cap-add=SYS_PTRACE \ + --security-opt label=disable \ + --userns=keep-id \ + --network=host \ + -v /mnt/storage/ds4:/models:ro \ + docker.io/kyuz0/strix-halo-ds4-toolbox:glm-rocm-7.2.4 \ + ds4-server \ + -m /models/GLM-5.2-UD-IQ2_XXS_RoutedIQ2XXS_blk78Q2K.gguf \ + --ctx 2177 \ + --host 0.0.0.0 \ + --port 8000 \ + --role worker \ + --layers 38:output \ + --coordinator 192.168.100.2 8081 \ + 2>&1 | tee /tmp/glm-worker-cache-f32.log +``` + +### Baseline coordinator + +Run on `192.168.100.2`: + +```sh +env -u DS4_ROCM_GLM_COMPACT_CACHE_F16 \ +DS4_GLM_MEMORY_GUARD_RESERVE_GB=8 \ +DS4_ROCM_GLM_CAUSAL_ATTN_GEMM=1 \ +DS4_BENCH_DISABLE_SNAPSHOT=1 \ +ds4-bench \ + --prompt-file ~/ds4/ds4/speed-bench/promessi_sposi.txt \ + --ctx-start 2048 \ + --ctx-max 2048 \ + --ctx-alloc 2177 \ + --gen-tokens 128 \ + --show-output \ + --csv /tmp/glm-cache-f32.csv \ + -m ~/ds4/GLM-5.2-UD-IQ2_XXS_RoutedIQ2XXS_blk78Q2K.gguf \ + --dist-prefill-chunk 256 \ + --dist-prefill-window 2 \ + --role coordinator \ + --layers 0:37 \ + --listen 192.168.100.2 8081 \ + 2>&1 | tee /tmp/glm-coordinator-cache-f32.log +``` + +### Candidate worker + +Stop the baseline worker, then run on `192.168.100.1`: + +```sh +podman run --rm -it \ + --name ds4-worker-cache-f16 \ + -e DS4_GLM_MEMORY_GUARD_RESERVE_GB=8 \ + -e DS4_ROCM_GLM_CAUSAL_ATTN_GEMM=1 \ + -e DS4_ROCM_GLM_COMPACT_CACHE_F16=1 \ + --device /dev/dri \ + --device /dev/kfd \ + --group-add keep-groups \ + --security-opt seccomp=unconfined \ + --ipc=host \ + --cap-add=SYS_PTRACE \ + --security-opt label=disable \ + --userns=keep-id \ + --network=host \ + -v /mnt/storage/ds4:/models:ro \ + docker.io/kyuz0/strix-halo-ds4-toolbox:glm-rocm-7.2.4 \ + ds4-server \ + -m /models/GLM-5.2-UD-IQ2_XXS_RoutedIQ2XXS_blk78Q2K.gguf \ + --ctx 2177 \ + --host 0.0.0.0 \ + --port 8000 \ + --role worker \ + --layers 38:output \ + --coordinator 192.168.100.2 8081 \ + 2>&1 | tee /tmp/glm-worker-cache-f16.log +``` + +### Candidate coordinator + +Run on `192.168.100.2`: + +```sh +DS4_GLM_MEMORY_GUARD_RESERVE_GB=8 \ +DS4_ROCM_GLM_CAUSAL_ATTN_GEMM=1 \ +DS4_ROCM_GLM_COMPACT_CACHE_F16=1 \ +DS4_BENCH_DISABLE_SNAPSHOT=1 \ +ds4-bench \ + --prompt-file ~/ds4/ds4/speed-bench/promessi_sposi.txt \ + --ctx-start 2048 \ + --ctx-max 2048 \ + --ctx-alloc 2177 \ + --gen-tokens 128 \ + --show-output \ + --csv /tmp/glm-cache-f16.csv \ + -m ~/ds4/GLM-5.2-UD-IQ2_XXS_RoutedIQ2XXS_blk78Q2K.gguf \ + --dist-prefill-chunk 256 \ + --dist-prefill-window 2 \ + --role coordinator \ + --layers 0:37 \ + --listen 192.168.100.2 8081 \ + 2>&1 | tee /tmp/glm-coordinator-cache-f16.log +``` + +Confirm that the baseline logs report a compact `f32` cache and the candidate +logs report `f16`. The candidate compact KV allocation should be approximately +half the baseline size. Compare `prefill_tps`, `gen_tps`, and +`gen_first_ms`; do not infer a win from elapsed startup or model-copy time. + +## Layer-local numerical comparison + +This check catches wrong layouts, strides, masking, RoPE, scaling, and result +scattering near the operation that changed. It is stronger than comparing +generated prose. + +Use workers configured like the performance workers, but set `--ctx 641`. +Run the baseline first: + +```sh +mkdir -p /tmp/glm-cache-attn + +env -u DS4_ROCM_GLM_COMPACT_CACHE_F16 \ +DS4_GLM_MEMORY_GUARD_RESERVE_GB=8 \ +DS4_ROCM_GLM_CAUSAL_ATTN_GEMM=1 \ +DS4_ROCM_GRAPH_DUMP_PREFIX=/tmp/glm-cache-attn/baseline \ +DS4_ROCM_GRAPH_DUMP_NONINVASIVE=1 \ +DS4_ROCM_GRAPH_DUMP_LAYER=0 \ +DS4_ROCM_GRAPH_DUMP_POS=0 \ +DS4_ROCM_GRAPH_DUMP_NAME=glm_indexed_attn_out \ +DS4_BENCH_DISABLE_SNAPSHOT=1 \ +ds4-bench \ + --prompt-file ~/ds4/ds4/speed-bench/promessi_sposi.txt \ + --ctx-start 256 \ + --ctx-max 256 \ + --ctx-alloc 641 \ + --gen-tokens 0 \ + -m ~/ds4/GLM-5.2-UD-IQ2_XXS_RoutedIQ2XXS_blk78Q2K.gguf \ + --dist-prefill-chunk 256 \ + --dist-prefill-window 2 \ + --role coordinator \ + --layers 0:37 \ + --listen 192.168.100.2 8081 +``` + +Restart the worker with the candidate F16-cache environment and run: + +```sh +DS4_GLM_MEMORY_GUARD_RESERVE_GB=8 \ +DS4_ROCM_GLM_CAUSAL_ATTN_GEMM=1 \ +DS4_ROCM_GLM_COMPACT_CACHE_F16=1 \ +DS4_ROCM_GRAPH_DUMP_PREFIX=/tmp/glm-cache-attn/candidate \ +DS4_ROCM_GRAPH_DUMP_NONINVASIVE=1 \ +DS4_ROCM_GRAPH_DUMP_LAYER=0 \ +DS4_ROCM_GRAPH_DUMP_POS=0 \ +DS4_ROCM_GRAPH_DUMP_NAME=glm_indexed_attn_out \ +DS4_BENCH_DISABLE_SNAPSHOT=1 \ +ds4-bench \ + --prompt-file ~/ds4/ds4/speed-bench/promessi_sposi.txt \ + --ctx-start 256 \ + --ctx-max 256 \ + --ctx-alloc 641 \ + --gen-tokens 0 \ + -m ~/ds4/GLM-5.2-UD-IQ2_XXS_RoutedIQ2XXS_blk78Q2K.gguf \ + --dist-prefill-chunk 256 \ + --dist-prefill-window 2 \ + --role coordinator \ + --layers 0:37 \ + --listen 192.168.100.2 8081 +``` + +Compare the dumps: + +```sh +python3 ~/ds4/ds4/speed-bench/compare_glm_validation.py tensor \ + /tmp/glm-cache-attn/baseline_glm_indexed_attn_out-0_pos0.bin \ + /tmp/glm-cache-attn/candidate_glm_indexed_attn_out-0_pos0.bin \ + --hidden 6144 +``` + +Any non-finite value is a failure. For a single early layer, cosine should be +very close to 1. As a diagnostic heuristic, investigate cosine below `0.9999` +or relative RMSE above `0.01`; these are not universal model-quality limits. + +## Frontier-logit comparison + +Layer-local agreement can still accumulate into meaningful final-logit drift. +Collect complete vocabulary logits at a fixed 512-token frontier. Use workers +with `--ctx 641`. + +Baseline: + +```sh +mkdir -p /tmp/glm-cache-logits-f32 + +env -u DS4_ROCM_GLM_COMPACT_CACHE_F16 \ +DS4_GLM_MEMORY_GUARD_RESERVE_GB=8 \ +DS4_ROCM_GLM_CAUSAL_ATTN_GEMM=1 \ +DS4_BENCH_DISABLE_SNAPSHOT=1 \ +ds4-bench \ + --prompt-file ~/ds4/ds4/speed-bench/promessi_sposi.txt \ + --ctx-start 512 \ + --ctx-max 512 \ + --ctx-alloc 641 \ + --gen-tokens 0 \ + --dump-frontier-logits-dir /tmp/glm-cache-logits-f32 \ + -m ~/ds4/GLM-5.2-UD-IQ2_XXS_RoutedIQ2XXS_blk78Q2K.gguf \ + --dist-prefill-chunk 256 \ + --dist-prefill-window 2 \ + --role coordinator \ + --layers 0:37 \ + --listen 192.168.100.2 8081 +``` + +Candidate: + +```sh +mkdir -p /tmp/glm-cache-logits-f16 + +DS4_GLM_MEMORY_GUARD_RESERVE_GB=8 \ +DS4_ROCM_GLM_CAUSAL_ATTN_GEMM=1 \ +DS4_ROCM_GLM_COMPACT_CACHE_F16=1 \ +DS4_BENCH_DISABLE_SNAPSHOT=1 \ +ds4-bench \ + --prompt-file ~/ds4/ds4/speed-bench/promessi_sposi.txt \ + --ctx-start 512 \ + --ctx-max 512 \ + --ctx-alloc 641 \ + --gen-tokens 0 \ + --dump-frontier-logits-dir /tmp/glm-cache-logits-f16 \ + -m ~/ds4/GLM-5.2-UD-IQ2_XXS_RoutedIQ2XXS_blk78Q2K.gguf \ + --dist-prefill-chunk 256 \ + --dist-prefill-window 2 \ + --role coordinator \ + --layers 0:37 \ + --listen 192.168.100.2 8081 +``` + +Compare them: + +```sh +python3 ~/ds4/ds4/speed-bench/compare_glm_validation.py logits \ + /tmp/glm-cache-logits-f32/frontier_000512.logits.json \ + /tmp/glm-cache-logits-f16/frontier_000512.logits.json +``` + +Interpret the result as a group: + +- non-finite logits are always a failure; +- centered cosine and relative RMSE measure distribution shape without an + irrelevant constant logit shift; +- KL divergence measures probability-distribution movement; +- top-10/top-50 overlap and the reference top-token rank show whether ordering + changed near the useful head of the distribution; +- matching top-1 is encouraging, but does not by itself prove equivalence; +- a top-1 change is less concerning when the reference margin is tiny. + +There is no universal threshold that turns one frontier into a model-quality +proof. Investigate a material regression from the last accepted path and use +the official GLM continuation scorer before making an approximation the +unconditional release default. + +## API smoke test + +After tensor and logit checks pass, run a server and test the actual chat +template. This catches routing, KV replay, and API integration failures: + +```sh +wget --timeout=900 --tries=1 -qO- \ + --header='Content-Type: application/json' \ + --post-data='{"model":"glm-5.2","messages":[{"role":"user","content":"Reply with exactly: hello"}],"thinking":{"type":"disabled"},"temperature":0,"max_tokens":16}' \ + http://localhost:8000/v1/chat/completions | +jq -r '.choices[0].message.content // .' +``` + +Expect a coherent response, normally `hello`. This is a smoke test, not a +replacement for tensor/logit comparison or the 100-case GLM continuation gate +documented in `QA_BEFORE_RELEASES.md`. + +## Decision record + +For every candidate, retain: + +- commit and model SHA-256; +- exact baseline and candidate environment variables; +- coordinator and worker logs; +- baseline and candidate CSV rows; +- layer-local tensor comparison; +- frontier-logit comparison; +- API smoke response; +- decision: reject, keep diagnostic, retain opt-in, or make default. + +If a candidate is faster but its drift is unexplained, keep it diagnostic or +reject it. Explain and validate correctness before promoting performance code. diff --git a/speed-bench/compare_glm_validation.py b/speed-bench/compare_glm_validation.py new file mode 100644 index 000000000..48e1211be --- /dev/null +++ b/speed-bench/compare_glm_validation.py @@ -0,0 +1,172 @@ +#!/usr/bin/env python3 +"""Compare ROCm GLM validation dumps without third-party dependencies.""" + +import argparse +import array +import heapq +import json +import math +from pathlib import Path + + +def vector_metrics(reference, candidate): + if len(reference) != len(candidate): + raise SystemExit( + f"length mismatch: reference={len(reference)} " + f"candidate={len(candidate)}" + ) + if not reference: + raise SystemExit("input vectors are empty") + + nonfinite = sum( + not (math.isfinite(a) and math.isfinite(b)) + for a, b in zip(reference, candidate) + ) + diff2 = sum((a - b) ** 2 for a, b in zip(reference, candidate)) + ref2 = sum(a * a for a in reference) + cand2 = sum(b * b for b in candidate) + dot = sum(a * b for a, b in zip(reference, candidate)) + + return { + "values": len(reference), + "nonfinite_pairs": nonfinite, + "max_abs": max(abs(a - b) for a, b in zip(reference, candidate)), + "rmse": math.sqrt(diff2 / len(reference)), + "rel_rmse": math.sqrt(diff2 / ref2) if ref2 else math.inf, + "cosine": dot / math.sqrt(ref2 * cand2) + if ref2 and cand2 + else math.nan, + } + + +def print_metrics(metrics, indent=" "): + print(f"{indent}values: {metrics['values']}") + print(f"{indent}nonfinite_pairs: {metrics['nonfinite_pairs']}") + print(f"{indent}max_abs: {metrics['max_abs']:.12g}") + print(f"{indent}rmse: {metrics['rmse']:.12g}") + print(f"{indent}rel_rmse: {metrics['rel_rmse']:.12g}") + print(f"{indent}cosine: {metrics['cosine']:.12g}") + + +def load_f32(path): + values = array.array("f") + with path.open("rb") as fp: + values.frombytes(fp.read()) + return values + + +def compare_tensor(args): + reference = load_f32(args.reference) + candidate = load_f32(args.candidate) + if args.hidden <= 0: + raise SystemExit("--hidden must be positive") + if len(reference) % args.hidden: + raise SystemExit( + f"reference contains {len(reference)} values, which is not " + f"divisible by hidden size {args.hidden}" + ) + + print(f"tokens: {len(reference) // args.hidden}") + print("all tokens:") + print_metrics(vector_metrics(reference, candidate)) + print("final token:") + print_metrics( + vector_metrics(reference[-args.hidden :], candidate[-args.hidden :]) + ) + + +def load_logits(path): + with path.open(encoding="utf-8") as fp: + obj = json.load(fp) + logits = obj.get("logits") + if not isinstance(logits, list): + raise SystemExit(f"{path}: missing logits array") + return logits + + +def logsumexp(values): + maximum = max(values) + return maximum + math.log(sum(math.exp(value - maximum) for value in values)) + + +def compare_logits(args): + reference = load_logits(args.reference) + candidate = load_logits(args.candidate) + metrics = vector_metrics(reference, candidate) + count = len(reference) + + ref_mean = sum(reference) / count + cand_mean = sum(candidate) / count + ref_centered = [value - ref_mean for value in reference] + cand_centered = [value - cand_mean for value in candidate] + centered = vector_metrics(ref_centered, cand_centered) + + ref_top = heapq.nlargest(50, range(count), key=reference.__getitem__) + cand_top = heapq.nlargest(50, range(count), key=candidate.__getitem__) + ref_top1 = ref_top[0] + cand_top1 = cand_top[0] + ref_margin = reference[ref_top[0]] - reference[ref_top[1]] + cand_margin = candidate[cand_top[0]] - candidate[cand_top[1]] + ref_top1_rank = 1 + sum( + value > candidate[ref_top1] for value in candidate + ) + + ref_lse = logsumexp(reference) + cand_lse = logsumexp(candidate) + kl_divergence = sum( + math.exp(a - ref_lse) + * ((a - ref_lse) - (b - cand_lse)) + for a, b in zip(reference, candidate) + ) + + print("raw logits:") + print_metrics(metrics) + print(f" mean_shift: {cand_mean - ref_mean:.12g}") + print("centered logits:") + print_metrics(centered) + print(f"KL(reference || candidate): {kl_divergence:.12g}") + print(f"top1 IDs: {ref_top1} -> {cand_top1}") + print(f"reference top1 rank: {ref_top1_rank}") + print(f"top1 margins: {ref_margin:.12g} -> {cand_margin:.12g}") + print( + "top10 overlap: " + f"{len(set(ref_top[:10]) & set(cand_top[:10]))}/10" + ) + print( + "top50 overlap: " + f"{len(set(ref_top) & set(cand_top))}/50" + ) + + +def main(): + parser = argparse.ArgumentParser( + description="Compare ds4 ROCm GLM tensor or frontier-logit dumps." + ) + subparsers = parser.add_subparsers(dest="mode", required=True) + + tensor = subparsers.add_parser( + "tensor", help="compare two raw F32 graph tensor dumps" + ) + tensor.add_argument("reference", type=Path) + tensor.add_argument("candidate", type=Path) + tensor.add_argument( + "--hidden", + type=int, + default=6144, + help="hidden width used to report the final token (default: 6144)", + ) + tensor.set_defaults(func=compare_tensor) + + logits = subparsers.add_parser( + "logits", help="compare two ds4-bench frontier-logit JSON files" + ) + logits.add_argument("reference", type=Path) + logits.add_argument("candidate", type=Path) + logits.set_defaults(func=compare_logits) + + args = parser.parse_args() + args.func(args) + + +if __name__ == "__main__": + main() From 5e61ad5ca935112523133049df81304e9f3dbf82 Mon Sep 17 00:00:00 2001 From: Donato Capitella Date: Sun, 26 Jul 2026 10:10:48 +0100 Subject: [PATCH 19/33] Remove ineffective ROCm GLM F16 compact cache --- ds4.c | 18 +- ds4_rocm_compat.cu | 23 -- .../ROCM_GLM_DISTRIBUTED_VALIDATION.md | 214 +++++------------- 3 files changed, 59 insertions(+), 196 deletions(-) diff --git a/ds4.c b/ds4.c index 28ff51aa5..b767a7869 100644 --- a/ds4.c +++ b/ds4.c @@ -14888,8 +14888,7 @@ static void print_vec_stats(const char *name, const float *x, uint64_t n) { #define DS4_GPU_ATTN_COMP_CACHE_F16 0 #endif -/* GLM's DSA compact cache has a ROCm runtime opt-in below; it is independent - * from the general compressed-attention cache format. */ +#define DS4_GPU_GLM_COMPACT_CACHE_F16 DS4_GPU_ATTN_COMP_CACHE_F16 /* ========================================================================= * Metal Release Graph State. @@ -34467,19 +34466,12 @@ static uint32_t glm_graph_indexed_prefill_score_tokens( uint32_t indexed_prefill_cap, uint32_t compact_cap); -static uint32_t glm_graph_compact_cache_is_f16(void) { -#if defined(__APPLE__) - return 1u; -#elif defined(DS4_ROCM_BUILD) - return getenv("DS4_ROCM_GLM_COMPACT_CACHE_F16") != NULL ? 1u : 0u; -#else - return 0u; -#endif +static uint64_t glm_graph_compact_cache_elem_bytes(void) { + return DS4_GPU_GLM_COMPACT_CACHE_F16 ? sizeof(uint16_t) : sizeof(float); } -static uint64_t glm_graph_compact_cache_elem_bytes(void) { - return glm_graph_compact_cache_is_f16() ? - sizeof(uint16_t) : sizeof(float); +static uint32_t glm_graph_compact_cache_is_f16(void) { + return DS4_GPU_GLM_COMPACT_CACHE_F16 ? 1u : 0u; } static bool glm_graph_expanded_kv_cache_enabled(bool ssd_streaming) { diff --git a/ds4_rocm_compat.cu b/ds4_rocm_compat.cu index ac8d485f6..79b587e3b 100644 --- a/ds4_rocm_compat.cu +++ b/ds4_rocm_compat.cu @@ -368,29 +368,6 @@ extern "C" int ds4_gpu_glm_attention_indexed_decode_typed_tensor( beta_fast, beta_slow); } -extern "C" int ds4_gpu_glm_attention_indexed_decode_split_group8_typed_tensor( - ds4_gpu_tensor *heads, ds4_gpu_tensor *partial_lora, - ds4_gpu_tensor *partial_ms, const ds4_gpu_tensor *q, - const ds4_gpu_tensor *qk_low, const ds4_gpu_tensor *kv_lora_cache, - const ds4_gpu_tensor *k_rope_cache, const void *model_map, - uint64_t model_size, uint64_t value_weight_offset, - uint32_t value_weight_type, const ds4_gpu_tensor *selected, - uint32_t n_selected, bool selected_rows_valid, uint32_t cache_cap, - bool cache_f16, uint32_t n_head, uint32_t kv_lora_dim, - uint32_t qk_nope, uint32_t qk_rope, uint32_t value_dim, - uint32_t n_ctx_orig, uint32_t block_rows, uint32_t n_blocks, - float freq_base, float freq_scale, float ext_factor, - float attn_factor, float beta_fast, float beta_slow) { - if (value_weight_type != 8u) return 0; - return ds4_gpu_glm_attention_indexed_decode_split_group8_tensor( - heads, partial_lora, partial_ms, q, qk_low, kv_lora_cache, - k_rope_cache, model_map, model_size, value_weight_offset, - selected, n_selected, selected_rows_valid, cache_cap, cache_f16, - n_head, kv_lora_dim, qk_nope, qk_rope, value_dim, n_ctx_orig, - block_rows, n_blocks, freq_base, freq_scale, ext_factor, - attn_factor, beta_fast, beta_slow); -} - extern "C" int ds4_gpu_glm_attention_indexed_batch_typed_tensor( ds4_gpu_tensor *heads, const ds4_gpu_tensor *q, const ds4_gpu_tensor *qk_low, const ds4_gpu_tensor *kv_lora_cache, diff --git a/speed-bench/ROCM_GLM_DISTRIBUTED_VALIDATION.md b/speed-bench/ROCM_GLM_DISTRIBUTED_VALIDATION.md index b3588043a..fa45ef56b 100644 --- a/speed-bench/ROCM_GLM_DISTRIBUTED_VALIDATION.md +++ b/speed-bench/ROCM_GLM_DISTRIBUTED_VALIDATION.md @@ -1,6 +1,6 @@ # ROCm GLM Distributed Change Validation -Use this playbook for every ROCm GLM distributed performance change. It keeps +Use this playbook for every ROCm GLM distributed performance change. It keeps speed measurements separate from numerical and model-quality checks, and it changes one variable at a time. @@ -16,22 +16,22 @@ Change those constants deliberately if the hardware or model changes, and record the new values with the results. Unless a command is explicitly labeled for the worker, run it on the -coordinator. The coordinator checkout used below is `~/ds4/ds4`. +coordinator. The coordinator checkout used below is `~/ds4/ds4`. ## Rules 1. Build baseline and candidate from the same commit and image. 2. Keep the model file, SHA-256, layer split, prompt, context allocation, chunk/window, generation length, power state, and background load fixed. -3. Hold previously accepted optimizations constant. Toggle only the change +3. Hold previously accepted optimizations constant. Toggle only the change under test. 4. Save worker and coordinator logs plus benchmark CSV files. -5. Test performance and numerical closeness separately. A coherent sampled +5. Test performance and numerical closeness separately. A coherent sampled answer is not a numerical comparison, and a matching argmax is not enough. -6. Run the baseline and candidate at least once as a smoke test. Before +6. Run the baseline and candidate at least once as a smoke test. Before declaring a small improvement real or making a path default, repeat the - complete pair. Treat differences below roughly 2% as noise until repeated. -7. Do not compare `ds4-bench --show-output` text with chat answers. The former + complete pair. Treat differences below roughly 2% as noise until repeated. +7. Do not compare `ds4-bench --show-output` text with chat answers. The former is a raw continuation of the prompt file; use the server API for chat. Record the code revision and model hash before each series: @@ -55,86 +55,31 @@ podman run --rm --userns=keep-id \ Require a successful warning-free build and run `git diff --check`. -## Performance protocol +## Matched baseline and candidate -The standard development point is a fresh 2048-token prefill followed by 128 -greedy generation tokens. Both nodes keep the accepted FP16 causal-attention -GEMM enabled. The candidate additionally enables the change under test. +Write down the single environment setting, option, or code path being tested. +Run the commands below once without it for the baseline, then again with only +that setting changed for the candidate. If the setting affects both halves of +the model, apply it to both worker and coordinator. -For the F16 compact-cache experiment, the isolated variable is: +Keep the accepted FP16 causal-attention GEMM enabled in both runs: ```text -baseline: DS4_ROCM_GLM_COMPACT_CACHE_F16 unset -candidate: DS4_ROCM_GLM_COMPACT_CACHE_F16=1 +DS4_ROCM_GLM_CAUSAL_ATTN_GEMM=1 ``` -### Baseline worker +### Worker -Run on `192.168.100.1`: +Run on `192.168.100.1`. Change `RUN_NAME` to `baseline` or `candidate`, and add +the one candidate environment setting only for the candidate run. ```sh -podman run --rm -it \ - --name ds4-worker-cache-f32 \ - -e DS4_GLM_MEMORY_GUARD_RESERVE_GB=8 \ - -e DS4_ROCM_GLM_CAUSAL_ATTN_GEMM=1 \ - --device /dev/dri \ - --device /dev/kfd \ - --group-add keep-groups \ - --security-opt seccomp=unconfined \ - --ipc=host \ - --cap-add=SYS_PTRACE \ - --security-opt label=disable \ - --userns=keep-id \ - --network=host \ - -v /mnt/storage/ds4:/models:ro \ - docker.io/kyuz0/strix-halo-ds4-toolbox:glm-rocm-7.2.4 \ - ds4-server \ - -m /models/GLM-5.2-UD-IQ2_XXS_RoutedIQ2XXS_blk78Q2K.gguf \ - --ctx 2177 \ - --host 0.0.0.0 \ - --port 8000 \ - --role worker \ - --layers 38:output \ - --coordinator 192.168.100.2 8081 \ - 2>&1 | tee /tmp/glm-worker-cache-f32.log -``` - -### Baseline coordinator +RUN_NAME=baseline -Run on `192.168.100.2`: - -```sh -env -u DS4_ROCM_GLM_COMPACT_CACHE_F16 \ -DS4_GLM_MEMORY_GUARD_RESERVE_GB=8 \ -DS4_ROCM_GLM_CAUSAL_ATTN_GEMM=1 \ -DS4_BENCH_DISABLE_SNAPSHOT=1 \ -ds4-bench \ - --prompt-file ~/ds4/ds4/speed-bench/promessi_sposi.txt \ - --ctx-start 2048 \ - --ctx-max 2048 \ - --ctx-alloc 2177 \ - --gen-tokens 128 \ - --show-output \ - --csv /tmp/glm-cache-f32.csv \ - -m ~/ds4/GLM-5.2-UD-IQ2_XXS_RoutedIQ2XXS_blk78Q2K.gguf \ - --dist-prefill-chunk 256 \ - --dist-prefill-window 2 \ - --role coordinator \ - --layers 0:37 \ - --listen 192.168.100.2 8081 \ - 2>&1 | tee /tmp/glm-coordinator-cache-f32.log -``` - -### Candidate worker - -Stop the baseline worker, then run on `192.168.100.1`: - -```sh podman run --rm -it \ - --name ds4-worker-cache-f16 \ + --name "ds4-worker-${RUN_NAME}" \ -e DS4_GLM_MEMORY_GUARD_RESERVE_GB=8 \ -e DS4_ROCM_GLM_CAUSAL_ATTN_GEMM=1 \ - -e DS4_ROCM_GLM_COMPACT_CACHE_F16=1 \ --device /dev/dri \ --device /dev/kfd \ --group-add keep-groups \ @@ -154,17 +99,19 @@ podman run --rm -it \ --role worker \ --layers 38:output \ --coordinator 192.168.100.2 8081 \ - 2>&1 | tee /tmp/glm-worker-cache-f16.log + 2>&1 | tee "/tmp/glm-worker-${RUN_NAME}.log" ``` -### Candidate coordinator +### Coordinator performance run -Run on `192.168.100.2`: +Set `RUN_NAME` to match the worker. Add the one candidate environment setting +before `ds4-bench` only for the candidate run. ```sh +RUN_NAME=baseline + DS4_GLM_MEMORY_GUARD_RESERVE_GB=8 \ DS4_ROCM_GLM_CAUSAL_ATTN_GEMM=1 \ -DS4_ROCM_GLM_COMPACT_CACHE_F16=1 \ DS4_BENCH_DISABLE_SNAPSHOT=1 \ ds4-bench \ --prompt-file ~/ds4/ds4/speed-bench/promessi_sposi.txt \ @@ -173,63 +120,35 @@ ds4-bench \ --ctx-alloc 2177 \ --gen-tokens 128 \ --show-output \ - --csv /tmp/glm-cache-f16.csv \ + --csv "/tmp/glm-${RUN_NAME}.csv" \ -m ~/ds4/GLM-5.2-UD-IQ2_XXS_RoutedIQ2XXS_blk78Q2K.gguf \ --dist-prefill-chunk 256 \ --dist-prefill-window 2 \ --role coordinator \ --layers 0:37 \ --listen 192.168.100.2 8081 \ - 2>&1 | tee /tmp/glm-coordinator-cache-f16.log + 2>&1 | tee "/tmp/glm-coordinator-${RUN_NAME}.log" ``` -Confirm that the baseline logs report a compact `f32` cache and the candidate -logs report `f16`. The candidate compact KV allocation should be approximately -half the baseline size. Compare `prefill_tps`, `gen_tps`, and -`gen_first_ms`; do not infer a win from elapsed startup or model-copy time. +Compare `prefill_tps`, `gen_tps`, and `gen_first_ms`; do not infer a win from +elapsed startup or model-copy time. ## Layer-local numerical comparison -This check catches wrong layouts, strides, masking, RoPE, scaling, and result -scattering near the operation that changed. It is stronger than comparing -generated prose. +This catches wrong layouts, strides, masking, RoPE, scaling, and result +scattering near the operation that changed. Use workers configured like the +performance runs. -Use workers configured like the performance workers, but set `--ctx 641`. -Run the baseline first: +Run the command once as `baseline`, restart the worker with the candidate +configuration, then run it as `candidate`: ```sh -mkdir -p /tmp/glm-cache-attn - -env -u DS4_ROCM_GLM_COMPACT_CACHE_F16 \ -DS4_GLM_MEMORY_GUARD_RESERVE_GB=8 \ -DS4_ROCM_GLM_CAUSAL_ATTN_GEMM=1 \ -DS4_ROCM_GRAPH_DUMP_PREFIX=/tmp/glm-cache-attn/baseline \ -DS4_ROCM_GRAPH_DUMP_NONINVASIVE=1 \ -DS4_ROCM_GRAPH_DUMP_LAYER=0 \ -DS4_ROCM_GRAPH_DUMP_POS=0 \ -DS4_ROCM_GRAPH_DUMP_NAME=glm_indexed_attn_out \ -DS4_BENCH_DISABLE_SNAPSHOT=1 \ -ds4-bench \ - --prompt-file ~/ds4/ds4/speed-bench/promessi_sposi.txt \ - --ctx-start 256 \ - --ctx-max 256 \ - --ctx-alloc 641 \ - --gen-tokens 0 \ - -m ~/ds4/GLM-5.2-UD-IQ2_XXS_RoutedIQ2XXS_blk78Q2K.gguf \ - --dist-prefill-chunk 256 \ - --dist-prefill-window 2 \ - --role coordinator \ - --layers 0:37 \ - --listen 192.168.100.2 8081 -``` - -Restart the worker with the candidate F16-cache environment and run: +RUN_NAME=baseline +mkdir -p /tmp/glm-validation -```sh DS4_GLM_MEMORY_GUARD_RESERVE_GB=8 \ DS4_ROCM_GLM_CAUSAL_ATTN_GEMM=1 \ -DS4_ROCM_GLM_COMPACT_CACHE_F16=1 \ -DS4_ROCM_GRAPH_DUMP_PREFIX=/tmp/glm-cache-attn/candidate \ +DS4_ROCM_GRAPH_DUMP_PREFIX="/tmp/glm-validation/${RUN_NAME}" \ DS4_ROCM_GRAPH_DUMP_NONINVASIVE=1 \ DS4_ROCM_GRAPH_DUMP_LAYER=0 \ DS4_ROCM_GRAPH_DUMP_POS=0 \ @@ -253,27 +172,26 @@ Compare the dumps: ```sh python3 ~/ds4/ds4/speed-bench/compare_glm_validation.py tensor \ - /tmp/glm-cache-attn/baseline_glm_indexed_attn_out-0_pos0.bin \ - /tmp/glm-cache-attn/candidate_glm_indexed_attn_out-0_pos0.bin \ + /tmp/glm-validation/baseline_glm_indexed_attn_out-0_pos0.bin \ + /tmp/glm-validation/candidate_glm_indexed_attn_out-0_pos0.bin \ --hidden 6144 ``` -Any non-finite value is a failure. For a single early layer, cosine should be -very close to 1. As a diagnostic heuristic, investigate cosine below `0.9999` +Any non-finite value is a failure. For a single early layer, cosine should be +very close to 1. As a diagnostic heuristic, investigate cosine below `0.9999` or relative RMSE above `0.01`; these are not universal model-quality limits. ## Frontier-logit comparison Layer-local agreement can still accumulate into meaningful final-logit drift. -Collect complete vocabulary logits at a fixed 512-token frontier. Use workers -with `--ctx 641`. - -Baseline: +Collect complete vocabulary logits at a fixed 512-token frontier. Run once as +`baseline`, restart the worker with the candidate configuration, and run again +as `candidate`. ```sh -mkdir -p /tmp/glm-cache-logits-f32 +RUN_NAME=baseline +mkdir -p "/tmp/glm-logits-${RUN_NAME}" -env -u DS4_ROCM_GLM_COMPACT_CACHE_F16 \ DS4_GLM_MEMORY_GUARD_RESERVE_GB=8 \ DS4_ROCM_GLM_CAUSAL_ATTN_GEMM=1 \ DS4_BENCH_DISABLE_SNAPSHOT=1 \ @@ -283,7 +201,7 @@ ds4-bench \ --ctx-max 512 \ --ctx-alloc 641 \ --gen-tokens 0 \ - --dump-frontier-logits-dir /tmp/glm-cache-logits-f32 \ + --dump-frontier-logits-dir "/tmp/glm-logits-${RUN_NAME}" \ -m ~/ds4/GLM-5.2-UD-IQ2_XXS_RoutedIQ2XXS_blk78Q2K.gguf \ --dist-prefill-chunk 256 \ --dist-prefill-window 2 \ @@ -292,36 +210,12 @@ ds4-bench \ --listen 192.168.100.2 8081 ``` -Candidate: - -```sh -mkdir -p /tmp/glm-cache-logits-f16 - -DS4_GLM_MEMORY_GUARD_RESERVE_GB=8 \ -DS4_ROCM_GLM_CAUSAL_ATTN_GEMM=1 \ -DS4_ROCM_GLM_COMPACT_CACHE_F16=1 \ -DS4_BENCH_DISABLE_SNAPSHOT=1 \ -ds4-bench \ - --prompt-file ~/ds4/ds4/speed-bench/promessi_sposi.txt \ - --ctx-start 512 \ - --ctx-max 512 \ - --ctx-alloc 641 \ - --gen-tokens 0 \ - --dump-frontier-logits-dir /tmp/glm-cache-logits-f16 \ - -m ~/ds4/GLM-5.2-UD-IQ2_XXS_RoutedIQ2XXS_blk78Q2K.gguf \ - --dist-prefill-chunk 256 \ - --dist-prefill-window 2 \ - --role coordinator \ - --layers 0:37 \ - --listen 192.168.100.2 8081 -``` - -Compare them: +Compare the dumps: ```sh python3 ~/ds4/ds4/speed-bench/compare_glm_validation.py logits \ - /tmp/glm-cache-logits-f32/frontier_000512.logits.json \ - /tmp/glm-cache-logits-f16/frontier_000512.logits.json + /tmp/glm-logits-baseline/frontier_000512.logits.json \ + /tmp/glm-logits-candidate/frontier_000512.logits.json ``` Interpret the result as a group: @@ -336,14 +230,14 @@ Interpret the result as a group: - a top-1 change is less concerning when the reference margin is tiny. There is no universal threshold that turns one frontier into a model-quality -proof. Investigate a material regression from the last accepted path and use +proof. Investigate a material regression from the last accepted path and use the official GLM continuation scorer before making an approximation the unconditional release default. ## API smoke test After tensor and logit checks pass, run a server and test the actual chat -template. This catches routing, KV replay, and API integration failures: +template. This catches routing, KV replay, and API integration failures: ```sh wget --timeout=900 --tries=1 -qO- \ @@ -353,7 +247,7 @@ wget --timeout=900 --tries=1 -qO- \ jq -r '.choices[0].message.content // .' ``` -Expect a coherent response, normally `hello`. This is a smoke test, not a +Expect a coherent response, normally `hello`. This is a smoke test, not a replacement for tensor/logit comparison or the 100-case GLM continuation gate documented in `QA_BEFORE_RELEASES.md`. @@ -371,4 +265,4 @@ For every candidate, retain: - decision: reject, keep diagnostic, retain opt-in, or make default. If a candidate is faster but its drift is unexplained, keep it diagnostic or -reject it. Explain and validate correctness before promoting performance code. +reject it. Explain and validate correctness before promoting performance code. From 5920d6808b54c24e21b0308a1d7640aa9d407949 Mon Sep 17 00:00:00 2001 From: Donato Capitella Date: Sun, 26 Jul 2026 10:32:37 +0100 Subject: [PATCH 20/33] Add opt-in ROCm GLM performance experiments --- ds4.c | 118 +++-- rocm/ds4_rocm_glm.cuh | 423 ++++++++++++++++++ rocm/ds4_rocm_matmul.cuh | 36 +- rocm/ds4_rocm_runtime.cuh | 3 + .../ROCM_GLM_DISTRIBUTED_VALIDATION.md | 77 ++++ 5 files changed, 624 insertions(+), 33 deletions(-) diff --git a/ds4.c b/ds4.c index b767a7869..fb7553c48 100644 --- a/ds4.c +++ b/ds4.c @@ -45442,6 +45442,14 @@ static bool glm_graph_forward_token( g->ssd_streaming && !static_decode_map && glm_graph_streaming_decode_sync_each_layer(); +#ifdef DS4_ROCM_BUILD + const bool rocm_decode_qkv_pair = + !g->ssd_streaming && + glm_graph_env_truthy( + getenv("DS4_ROCM_GLM_DECODE_QKV_PAIR")); +#else + const bool rocm_decode_qkv_pair = false; +#endif bool ok = true; if (input_hc) { ok = ds4_gpu_tensor_write(g->cur, @@ -45521,7 +45529,46 @@ static bool glm_graph_forward_token( DS4_RMS_EPS) != 0; DS4_GLM_PROFILE_DECODE_STAGE("glm_decode_attn", "attn_norm"); const uint32_t decode_ablate = glm_decode_ablate_mask(); - if (ok && !(decode_ablate & DS4_GLM_ABLATE_QPATH)) { + bool qkv_pair_projected = false; + if (ok && + rocm_decode_qkv_pair && + !(decode_ablate & DS4_GLM_ABLATE_QPATH) && + l->attn_q_a->type == DS4_TENSOR_Q8_0 && + l->attn_kv_a_mqa->type == DS4_TENSOR_Q8_0) { + qkv_pair_projected = + ds4_gpu_matmul_q8_0_pair_tensor( + g->q_rank, + g->kv_raw, + model->map, + model->size, + l->attn_q_a->abs_offset, + l->attn_kv_a_mqa->abs_offset, + DS4_N_EMBD, + DS4_N_LORA_Q, + kv_raw_dim, + g->attn_norm, + 1) != 0; + if (qkv_pair_projected) { + static bool notice_printed = false; + if (!notice_printed) { + fprintf(stderr, + "ds4: ROCm GLM one-token Q/KV input " + "projections use the paired Q8 kernel\n"); + notice_printed = true; + } + } else { + static bool fallback_notice_printed = false; + if (!fallback_notice_printed) { + fprintf(stderr, + "ds4: ROCm GLM paired Q/KV projection " + "unavailable; using separate projections\n"); + fallback_notice_printed = true; + } + } + } + if (ok && + !qkv_pair_projected && + !(decode_ablate & DS4_GLM_ABLATE_QPATH)) { ok = glm_graph_matmul_q8_0_decode_profiled_tensor(g->q_rank, model, l->attn_q_a->abs_offset, @@ -45538,16 +45585,19 @@ static bool glm_graph_forward_token( g->compact_cache_cap != 0; const bool fuse_qkv_norm = !decode_stage_profile && !fuse_qkv_norm_store; if (ok && fuse_qkv_norm_store) { - ok = glm_graph_matmul_q8_0_decode_profiled_tensor(g->kv_raw, - model, - l->attn_kv_a_mqa->abs_offset, - DS4_N_EMBD, - kv_raw_dim, - g->attn_norm, - il, - pos, - "attn_kv_a_store", - g->ssd_streaming) != 0; + if (!qkv_pair_projected) { + ok = glm_graph_matmul_q8_0_decode_profiled_tensor( + g->kv_raw, + model, + l->attn_kv_a_mqa->abs_offset, + DS4_N_EMBD, + kv_raw_dim, + g->attn_norm, + il, + pos, + "attn_kv_a_store", + g->ssd_streaming) != 0; + } if (ok) { ok = ds4_gpu_glm_qkv_norm_store_compact_kv_tensor( g->q_rank_norm, @@ -45570,16 +45620,19 @@ static bool glm_graph_forward_token( DS4_RMS_EPS) != 0; } } else if (ok && fuse_qkv_norm) { - ok = glm_graph_matmul_q8_0_decode_profiled_tensor(g->kv_raw, - model, - l->attn_kv_a_mqa->abs_offset, - DS4_N_EMBD, - kv_raw_dim, - g->attn_norm, - il, - pos, - "attn_kv_a_norm", - g->ssd_streaming) != 0; + if (!qkv_pair_projected) { + ok = glm_graph_matmul_q8_0_decode_profiled_tensor( + g->kv_raw, + model, + l->attn_kv_a_mqa->abs_offset, + DS4_N_EMBD, + kv_raw_dim, + g->attn_norm, + il, + pos, + "attn_kv_a_norm", + g->ssd_streaming) != 0; + } if (ok) { ok = ds4_gpu_dsv4_qkv_rms_norm_rows_tensor(g->q_rank_norm, g->q_rank, @@ -45674,16 +45727,19 @@ static bool glm_graph_forward_token( } DS4_GLM_PROFILE_DECODE_STAGE("glm_decode_attn", "indexer_k"); if (ok && !fuse_qkv_norm && !fuse_qkv_norm_store) { - ok = glm_graph_matmul_q8_0_decode_profiled_tensor(g->kv_raw, - model, - l->attn_kv_a_mqa->abs_offset, - DS4_N_EMBD, - kv_raw_dim, - g->attn_norm, - il, - pos, - "attn_kv_a", - g->ssd_streaming) != 0; + if (!qkv_pair_projected) { + ok = glm_graph_matmul_q8_0_decode_profiled_tensor( + g->kv_raw, + model, + l->attn_kv_a_mqa->abs_offset, + DS4_N_EMBD, + kv_raw_dim, + g->attn_norm, + il, + pos, + "attn_kv_a", + g->ssd_streaming) != 0; + } if (ok) ok = ds4_gpu_glm_kv_lora_rms_norm_tensor(g->kv_norm, g->kv_raw, model->map, diff --git a/rocm/ds4_rocm_glm.cuh b/rocm/ds4_rocm_glm.cuh index 881405fa5..0172db49b 100644 --- a/rocm/ds4_rocm_glm.cuh +++ b/rocm/ds4_rocm_glm.cuh @@ -2133,6 +2133,130 @@ __global__ static void glm_causal_gemm_scatter_head_kernel( head_out[i]; } +__global__ static void glm_selected_gemm_kv_to_f16_kernel( + __half *dst, + const char *src, + const int32_t *selected, + uint32_t n_tokens, + uint32_t n_selected, + uint32_t kv_lora_dim, + uint32_t cache_cap, + bool src_f16) { + const uint64_t total = + (uint64_t)n_tokens * n_selected * kv_lora_dim; + for (uint64_t i = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x; + i < total; + i += (uint64_t)gridDim.x * blockDim.x) { + const uint64_t selected_row = i / kv_lora_dim; + const uint32_t j = + (uint32_t)(i - selected_row * kv_lora_dim); + const int32_t cache_row = selected[selected_row]; + dst[i] = + cache_row >= 0 && (uint32_t)cache_row < cache_cap ? + __float2half(glm_rocm_cache_load( + src, + (uint64_t)(uint32_t)cache_row * kv_lora_dim + j, + src_f16)) : + __float2half(0.0f); + } +} + +__global__ static void glm_selected_gemm_rope_to_f16_kernel( + __half *dst, + const char *src, + const int32_t *selected, + uint32_t n_tokens, + uint32_t n_selected, + uint32_t qk_rope, + uint32_t cache_cap, + bool src_f16, + uint32_t n_ctx_orig, + float freq_base, + float freq_scale, + float ext_factor, + float attn_factor, + float beta_fast, + float beta_slow) { + const uint32_t rope_pairs = qk_rope >> 1u; + const uint64_t total = + (uint64_t)n_tokens * n_selected * rope_pairs; + for (uint64_t i = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x; + i < total; + i += (uint64_t)gridDim.x * blockDim.x) { + const uint64_t selected_row = i / rope_pairs; + const uint32_t r = + (uint32_t)(i - selected_row * rope_pairs) << 1u; + const int32_t cache_row = selected[selected_row]; + float2 y = make_float2(0.0f, 0.0f); + if (cache_row >= 0 && (uint32_t)cache_row < cache_cap) { + y = glm_rocm_rotated_cache_rope_pair( + src, + (uint64_t)(uint32_t)cache_row * qk_rope, + r, + (uint32_t)cache_row, + qk_rope, + src_f16, + n_ctx_orig, + freq_base, + freq_scale, + ext_factor, + attn_factor, + beta_fast, + beta_slow); + } + const uint64_t out = selected_row * qk_rope + r; + dst[out] = __float2half(y.x); + dst[out + 1u] = __float2half(y.y); + } +} + +__global__ static void glm_selected_gemm_softmax_f16_kernel( + __half *probs, + float *scores, + uint32_t n_tokens, + uint32_t n_selected, + float scale) { + const uint32_t token = blockIdx.x; + if (token >= n_tokens) return; + float *row = scores + (uint64_t)token * n_selected; + __shared__ float reduce[256]; + float vmax = -INFINITY; + for (uint32_t s = threadIdx.x; s < n_selected; s += blockDim.x) { + const float v = row[s] * scale; + row[s] = v; + vmax = fmaxf(vmax, v); + } + reduce[threadIdx.x] = vmax; + __syncthreads(); + for (uint32_t stride = blockDim.x >> 1u; stride != 0u; stride >>= 1u) { + if (threadIdx.x < stride) { + reduce[threadIdx.x] = + fmaxf(reduce[threadIdx.x], reduce[threadIdx.x + stride]); + } + __syncthreads(); + } + const float max_score = reduce[0]; + float sum = 0.0f; + for (uint32_t s = threadIdx.x; s < n_selected; s += blockDim.x) { + const float v = expf(row[s] - max_score); + row[s] = v; + sum += v; + } + reduce[threadIdx.x] = sum; + __syncthreads(); + for (uint32_t stride = blockDim.x >> 1u; stride != 0u; stride >>= 1u) { + if (threadIdx.x < stride) { + reduce[threadIdx.x] += reduce[threadIdx.x + stride]; + } + __syncthreads(); + } + const float inv = 1.0f / fmaxf(reduce[0], 1.0e-20f); + for (uint32_t s = threadIdx.x; s < n_selected; s += blockDim.x) { + probs[(uint64_t)token * n_selected + s] = + __float2half(row[s] * inv); + } +} + static int glm_causal_gemm_scratch_part( uint64_t *offset, uint64_t bytes, @@ -2364,6 +2488,262 @@ static int glm_attention_indexed_lora_causal_gemm( return 1; } +static int glm_attention_indexed_lora_selected_gemm( + ds4_gpu_tensor *lora_out, + const ds4_gpu_tensor *q, + const ds4_gpu_tensor *qk_low, + const ds4_gpu_tensor *kv_lora_cache, + const ds4_gpu_tensor *k_rope_cache, + const ds4_gpu_tensor *selected, + uint32_t n_tokens, + uint32_t n_selected, + uint32_t cache_cap, + bool cache_f16, + uint32_t n_head, + uint32_t kv_lora_dim, + uint32_t qk_nope, + uint32_t qk_rope, + uint32_t n_ctx_orig, + float freq_base, + float freq_scale, + float ext_factor, + float attn_factor, + float beta_fast, + float beta_slow) { + if (!g_cublas_ready || !lora_out || !q || !qk_low || + !kv_lora_cache || !k_rope_cache || !selected || + n_tokens < 64u || n_tokens > 256u || + n_selected == 0u || n_selected > cache_cap || + n_head == 0u || kv_lora_dim == 0u || + qk_rope == 0u || (qk_rope & 1u) != 0u || + n_tokens > INT_MAX || n_selected > INT_MAX || + kv_lora_dim > INT_MAX || qk_rope > INT_MAX) { + return 0; + } + + uint64_t selected_rows = 0; + uint64_t kv_count = 0, rope_count = 0; + uint64_t low_count = 0, qrope_count = 0; + uint64_t score_count = 0, head_count = 0; + if (!cuda_u64_mul_checked(n_tokens, n_selected, &selected_rows) || + !cuda_u64_mul_checked(selected_rows, kv_lora_dim, &kv_count) || + !cuda_u64_mul_checked(selected_rows, qk_rope, &rope_count) || + !cuda_u64_mul_checked(n_tokens, kv_lora_dim, &low_count) || + !cuda_u64_mul_checked(n_tokens, qk_rope, &qrope_count) || + !cuda_u64_mul_checked(n_tokens, n_selected, &score_count) || + !cuda_u64_mul_checked(n_tokens, kv_lora_dim, &head_count)) { + return 0; + } + + uint64_t kv_bytes = 0, rope_bytes = 0; + uint64_t low_bytes = 0, qrope_bytes = 0; + uint64_t score_bytes = 0, prob_bytes = 0, head_bytes = 0; + if (!cuda_u64_mul_checked(kv_count, sizeof(__half), &kv_bytes) || + !cuda_u64_mul_checked(rope_count, sizeof(__half), &rope_bytes) || + !cuda_u64_mul_checked(low_count, sizeof(__half), &low_bytes) || + !cuda_u64_mul_checked(qrope_count, sizeof(__half), &qrope_bytes) || + !cuda_u64_mul_checked(score_count, sizeof(float), &score_bytes) || + !cuda_u64_mul_checked(score_count, sizeof(__half), &prob_bytes) || + !cuda_u64_mul_checked(head_count, sizeof(float), &head_bytes)) { + return 0; + } + + uint64_t scratch_bytes = 0; + uint64_t kv_offset = 0, rope_offset = 0; + uint64_t low_offset = 0, qrope_offset = 0; + uint64_t score_offset = 0, prob_offset = 0, head_offset = 0; + if (!glm_causal_gemm_scratch_part(&scratch_bytes, kv_bytes, &kv_offset) || + !glm_causal_gemm_scratch_part(&scratch_bytes, rope_bytes, &rope_offset) || + !glm_causal_gemm_scratch_part(&scratch_bytes, low_bytes, &low_offset) || + !glm_causal_gemm_scratch_part(&scratch_bytes, qrope_bytes, &qrope_offset) || + !glm_causal_gemm_scratch_part(&scratch_bytes, score_bytes, &score_offset) || + !glm_causal_gemm_scratch_part(&scratch_bytes, prob_bytes, &prob_offset) || + !glm_causal_gemm_scratch_part(&scratch_bytes, head_bytes, &head_offset)) { + return 0; + } + + char *scratch = + (char *)cuda_tmp_alloc(scratch_bytes, "glm selected attention gemm"); + if (!scratch) return 0; + __half *kv_h = (__half *)(scratch + kv_offset); + __half *rope_h = (__half *)(scratch + rope_offset); + __half *low_h = (__half *)(scratch + low_offset); + __half *qrope_h = (__half *)(scratch + qrope_offset); + float *scores = (float *)(scratch + score_offset); + __half *probs = (__half *)(scratch + prob_offset); + float *head_out = (float *)(scratch + head_offset); + + const uint32_t kv_blocks = + (uint32_t)min((kv_count + 255u) / 256u, 4096ull); + glm_selected_gemm_kv_to_f16_kernel<<>>( + kv_h, + (const char *)kv_lora_cache->ptr, + (const int32_t *)selected->ptr, + n_tokens, + n_selected, + kv_lora_dim, + cache_cap, + cache_f16); + if (!cuda_ok(cudaGetLastError(), + "glm selected attention kv gather launch")) { + return 0; + } + const uint64_t rope_pairs = selected_rows * (qk_rope >> 1u); + const uint32_t rope_blocks = + (uint32_t)min((rope_pairs + 255u) / 256u, 4096ull); + glm_selected_gemm_rope_to_f16_kernel<<>>( + rope_h, + (const char *)k_rope_cache->ptr, + (const int32_t *)selected->ptr, + n_tokens, + n_selected, + qk_rope, + cache_cap, + cache_f16, + n_ctx_orig, + freq_base, + freq_scale, + ext_factor, + attn_factor, + beta_fast, + beta_slow); + if (!cuda_ok(cudaGetLastError(), + "glm selected attention rope gather launch")) { + return 0; + } + + const float one = 1.0f; + const float zero = 0.0f; + const float scale = 1.0f / sqrtf((float)(qk_nope + qk_rope)); + const uint32_t gather_blocks = + (uint32_t)((low_count + 255u) / 256u); + const uint32_t scatter_blocks = + (uint32_t)((head_count + 255u) / 256u); + const long long kv_stride = + (long long)n_selected * (long long)kv_lora_dim; + const long long rope_stride = + (long long)n_selected * (long long)qk_rope; + const long long low_stride = (long long)kv_lora_dim; + const long long qrope_stride = (long long)qk_rope; + const long long score_stride = (long long)n_selected; + for (uint32_t head = 0; head < n_head; head++) { + glm_causal_gemm_gather_head_f16_kernel<<>>( + low_h, + qrope_h, + (const float *)qk_low->ptr, + (const float *)q->ptr, + n_tokens, + head, + n_head, + kv_lora_dim, + qk_nope, + qk_rope); + if (!cuda_ok(cudaGetLastError(), + "glm selected attention query gather launch")) { + return 0; + } + cublasStatus_t st = cublasGemmStridedBatchedEx( + g_cublas, + CUBLAS_OP_T, + CUBLAS_OP_N, + (int)n_selected, + 1, + (int)kv_lora_dim, + &one, + kv_h, + CUDA_R_16F, + (int)kv_lora_dim, + kv_stride, + low_h, + CUDA_R_16F, + (int)kv_lora_dim, + low_stride, + &zero, + scores, + CUDA_R_32F, + (int)n_selected, + score_stride, + (int)n_tokens, + CUBLAS_COMPUTE_32F, + CUBLAS_GEMM_DEFAULT); + if (!cublas_ok(st, "glm selected attention lora score gemm")) { + return 0; + } + st = cublasGemmStridedBatchedEx( + g_cublas, + CUBLAS_OP_T, + CUBLAS_OP_N, + (int)n_selected, + 1, + (int)qk_rope, + &one, + rope_h, + CUDA_R_16F, + (int)qk_rope, + rope_stride, + qrope_h, + CUDA_R_16F, + (int)qk_rope, + qrope_stride, + &one, + scores, + CUDA_R_32F, + (int)n_selected, + score_stride, + (int)n_tokens, + CUBLAS_COMPUTE_32F, + CUBLAS_GEMM_DEFAULT); + if (!cublas_ok(st, "glm selected attention rope score gemm")) { + return 0; + } + glm_selected_gemm_softmax_f16_kernel<<>>( + probs, scores, n_tokens, n_selected, scale); + if (!cuda_ok(cudaGetLastError(), + "glm selected attention softmax launch")) { + return 0; + } + st = cublasGemmStridedBatchedEx( + g_cublas, + CUBLAS_OP_N, + CUBLAS_OP_N, + (int)kv_lora_dim, + 1, + (int)n_selected, + &one, + kv_h, + CUDA_R_16F, + (int)kv_lora_dim, + kv_stride, + probs, + CUDA_R_16F, + (int)n_selected, + score_stride, + &zero, + head_out, + CUDA_R_32F, + (int)kv_lora_dim, + low_stride, + (int)n_tokens, + CUBLAS_COMPUTE_32F, + CUBLAS_GEMM_DEFAULT); + if (!cublas_ok(st, "glm selected attention value gemm")) { + return 0; + } + glm_causal_gemm_scatter_head_kernel<<>>( + (float *)lora_out->ptr, + head_out, + n_tokens, + head, + n_head, + kv_lora_dim); + if (!cuda_ok(cudaGetLastError(), + "glm selected attention head scatter launch")) { + return 0; + } + } + return 1; +} + static int glm_attention_indexed_lora_launch( ds4_gpu_tensor *lora_out, const ds4_gpu_tensor *q, @@ -2418,6 +2798,49 @@ static int glm_attention_indexed_lora_launch( !glm_rocm_tensor_has_cache2(k_rope_cache, cache_cap, qk_rope, elem)) { return 0; } + /* Once the visible context exceeds the indexer's top-k, each token carries + * its own selected row list. Gather those rows into per-token FP16 + * matrices, then batch the score and value products through BLAS. Keep + * the experiment opt-in and bounded: at 256 x 2048 selected rows the + * temporary workspace is already roughly 580 MiB. */ + if (!causal_range && has_selected && + cuda_env_present(getenv("DS4_ROCM_GLM_SELECTED_ATTN_GEMM")) && + glm_attention_indexed_lora_selected_gemm(lora_out, + q, + qk_low, + kv_lora_cache, + k_rope_cache, + selected, + n_tokens, + n_selected, + cache_cap, + cache_f16, + n_head, + kv_lora_dim, + qk_nope, + qk_rope, + n_ctx_orig, + freq_base, + freq_scale, + ext_factor, + attn_factor, + beta_fast, + beta_slow)) { + static int notice_printed = 0; + if (!notice_printed) { + fprintf(stderr, + DS4_GPU_LOG_PREFIX + "GLM selected indexed prefill using fp16 " + DS4_GPU_BLAS_NAME + " strided-batched attention GEMMs " + "(tokens=%u selected=%u cache=%s)\n", + n_tokens, + n_selected, + cache_f16 ? "f16" : "f32"); + notice_printed = 1; + } + return 1; + } /* Diagnostic port of the heterogeneous branch's dense causal prefill * path. Keep it opt-in until remote gfx1151 timing and logit comparisons * establish that the FP16 GEMMs are a safe default. */ diff --git a/rocm/ds4_rocm_matmul.cuh b/rocm/ds4_rocm_matmul.cuh index 31292e7d0..8c05433b9 100644 --- a/rocm/ds4_rocm_matmul.cuh +++ b/rocm/ds4_rocm_matmul.cuh @@ -330,7 +330,12 @@ 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 ((in_dim & 31u) == 0u && in_dim <= 8192u) { + const bool extended_sharedx = + in_dim > 8192u && + in_dim <= 16384u && + cuda_runtime_config()->q8_decode_sharedx_64k; + if ((in_dim & 31u) == 0u && + (in_dim <= 8192u || extended_sharedx)) { const unsigned rows_per_block = 32u; const unsigned threads = rows_per_block * 32u; matmul_q8_0_f32_sharedx_warp_rows_w32_kernel<<< @@ -343,7 +348,34 @@ static int cuda_matmul_q8_0_tensor_labeled(ds4_gpu_tensor *out, const void *mode (uint32_t)blocks, out_dim, blocks * 34u); - return cuda_ok(cudaGetLastError(), "matmul_q8_0 f32 sharedx launch"); + const cudaError_t launch_err = cudaGetLastError(); + if (launch_err == cudaSuccess) { + if (extended_sharedx) { + static int notice_printed = 0; + if (!notice_printed) { + fprintf(stderr, + DS4_GPU_LOG_PREFIX + "Q8 one-token shared-input kernel enabled " + "through 64 KiB LDS (in_dim=%llu)\n", + (unsigned long long)in_dim); + notice_printed = 1; + } + } + return 1; + } + if (!extended_sharedx) { + return cuda_ok(launch_err, + "matmul_q8_0 f32 sharedx launch"); + } + static int fallback_notice_printed = 0; + if (!fallback_notice_printed) { + fprintf(stderr, + DS4_GPU_LOG_PREFIX + "Q8 64 KiB shared-input launch unavailable " + "(%s); falling back to the warp-row kernel\n", + cudaGetErrorString(launch_err)); + fallback_notice_printed = 1; + } } matmul_q8_0_f32_warp8_kernel<<<((unsigned)out_dim + 7u) / 8u, 256>>>( (float *)out->ptr, diff --git a/rocm/ds4_rocm_runtime.cuh b/rocm/ds4_rocm_runtime.cuh index 26688e108..2443bc2b4 100644 --- a/rocm/ds4_rocm_runtime.cuh +++ b/rocm/ds4_rocm_runtime.cuh @@ -4754,6 +4754,7 @@ struct ds4_rocm_runtime_config { int shared_down_cublas; int glm_grouped_value_project; int glm_grouped_qk_low; + int q8_decode_sharedx_64k; int graph_dump; uint32_t moe_decode_rpb; uint32_t moe_decode_gate_rpb; @@ -4783,6 +4784,8 @@ static const ds4_rocm_runtime_config *cuda_runtime_config(void) { g_rocm_cfg.glm_grouped_qk_low = glm_grouped_qk_low_env == NULL || cuda_env_present(glm_grouped_qk_low_env); + g_rocm_cfg.q8_decode_sharedx_64k = + cuda_env_present(getenv("DS4_ROCM_Q8_DECODE_SHAREDX_64K")); const int graph_dump_requested = cuda_env_present(getenv("DS4_ROCM_GRAPH_DUMP_PREFIX")) || cuda_env_present(getenv("DS4_METAL_GRAPH_DUMP_PREFIX")); diff --git a/speed-bench/ROCM_GLM_DISTRIBUTED_VALIDATION.md b/speed-bench/ROCM_GLM_DISTRIBUTED_VALIDATION.md index fa45ef56b..8c0e90340 100644 --- a/speed-bench/ROCM_GLM_DISTRIBUTED_VALIDATION.md +++ b/speed-bench/ROCM_GLM_DISTRIBUTED_VALIDATION.md @@ -55,6 +55,83 @@ podman run --rm --userns=keep-id \ Require a successful warning-free build and run `git diff --check`. +## One-build experimental flag matrix + +The following ROCm experiments are independent and default to off. They can be +compiled into one image, enabled separately for attribution, and enabled +together for the final interaction check: + +| Experiment | Environment variable | Intended effect | +| --- | --- | --- | +| 64 KiB shared input | `DS4_ROCM_Q8_DECODE_SHAREDX_64K=1` | Let one-token Q8 projections with a 16,384-element input reuse the input through LDS. | +| Paired Q/KV projection | `DS4_ROCM_GLM_DECODE_QKV_PAIR=1` | Compute the one-token GLM Q and KV input projections together. | +| Selected attention GEMM | `DS4_ROCM_GLM_SELECTED_ATTN_GEMM=1` | Gather per-token selected KV rows and use FP16 strided-batched BLAS after the 2,048-row indexer boundary. | + +Apply an enabled flag to **both** worker and coordinator. For the Podman worker, +add it as `-e NAME=1`; for the coordinator, put `NAME=1 \` before `ds4-bench` +or `ds4-server`. + +Use this order so one deployment answers both attribution and interaction: + +1. no experimental flags; +2. `DS4_ROCM_Q8_DECODE_SHAREDX_64K=1`; +3. `DS4_ROCM_GLM_DECODE_QKV_PAIR=1`; +4. both decode flags; +5. `DS4_ROCM_GLM_SELECTED_ATTN_GEMM=1`; +6. all three flags. + +The expected activation messages are: + +```text +Q8 one-token shared-input kernel enabled through 64 KiB LDS +ROCm GLM one-token Q/KV input projections use the paired Q8 kernel +GLM selected indexed prefill using fp16 hipBLAS strided-batched attention GEMMs +``` + +Absence of a message means the tested workload did not reach that path. The +64 KiB path also reports and automatically falls back if the launch is not +supported by the device. + +The normal 2,048-token benchmark below exercises the two decode flags, but it +does **not** exercise selected attention: 2,048 visible rows still use the +dense causal path. Measure the prompt-processing cliff with a live 256-token +suffix curve: + +```sh +RUN_NAME=selected-baseline + +DS4_GLM_MEMORY_GUARD_RESERVE_GB=8 \ +DS4_ROCM_GLM_CAUSAL_ATTN_GEMM=1 \ +DS4_BENCH_DISABLE_SNAPSHOT=1 \ +ds4-bench \ + --prompt-file ~/ds4/ds4/speed-bench/promessi_sposi.txt \ + --ctx-start 2048 \ + --ctx-max 4096 \ + --step-incr 256 \ + --ctx-alloc 4097 \ + --gen-tokens 0 \ + --csv "/tmp/glm-${RUN_NAME}.csv" \ + -m ~/ds4/GLM-5.2-UD-IQ2_XXS_RoutedIQ2XXS_blk78Q2K.gguf \ + --dist-prefill-chunk 256 \ + --dist-prefill-window 2 \ + --role coordinator \ + --layers 0:37 \ + --listen 192.168.100.2 8081 \ + 2>&1 | tee "/tmp/glm-coordinator-${RUN_NAME}.log" +``` + +Run it once without selected attention and once with +`DS4_ROCM_GLM_SELECTED_ATTN_GEMM=1` on both nodes. The first 2,048-token row is +the causal control; the later 256-token suffix rows exercise selected +attention. The selected GEMM experiment accepts 64–256 tokens and uses about +580 MiB of temporary workspace at 256 tokens with 2,048 selected rows. + +For a layer-local selected-attention comparison, adapt the graph-dump command +below to `--ctx-start 2304 --ctx-max 2304 --ctx-alloc 2433` and set +`DS4_ROCM_GRAPH_DUMP_POS=2048`. For an end-to-end comparison, dump frontier +logits at 2,304 tokens and compare `frontier_002304.logits.json`. A 512-token +dump cannot validate this path. + ## Matched baseline and candidate Write down the single environment setting, option, or code path being tested. From 41a3c79d47d9bdea3146d85c70a290404e81ede6 Mon Sep 17 00:00:00 2001 From: Donato Capitella Date: Sun, 26 Jul 2026 11:29:24 +0100 Subject: [PATCH 21/33] Remove ineffective ROCm GLM QKV experiment --- ds4.c | 118 +++++------------- .../ROCM_GLM_DISTRIBUTED_VALIDATION.md | 16 +-- 2 files changed, 37 insertions(+), 97 deletions(-) diff --git a/ds4.c b/ds4.c index fb7553c48..b767a7869 100644 --- a/ds4.c +++ b/ds4.c @@ -45442,14 +45442,6 @@ static bool glm_graph_forward_token( g->ssd_streaming && !static_decode_map && glm_graph_streaming_decode_sync_each_layer(); -#ifdef DS4_ROCM_BUILD - const bool rocm_decode_qkv_pair = - !g->ssd_streaming && - glm_graph_env_truthy( - getenv("DS4_ROCM_GLM_DECODE_QKV_PAIR")); -#else - const bool rocm_decode_qkv_pair = false; -#endif bool ok = true; if (input_hc) { ok = ds4_gpu_tensor_write(g->cur, @@ -45529,46 +45521,7 @@ static bool glm_graph_forward_token( DS4_RMS_EPS) != 0; DS4_GLM_PROFILE_DECODE_STAGE("glm_decode_attn", "attn_norm"); const uint32_t decode_ablate = glm_decode_ablate_mask(); - bool qkv_pair_projected = false; - if (ok && - rocm_decode_qkv_pair && - !(decode_ablate & DS4_GLM_ABLATE_QPATH) && - l->attn_q_a->type == DS4_TENSOR_Q8_0 && - l->attn_kv_a_mqa->type == DS4_TENSOR_Q8_0) { - qkv_pair_projected = - ds4_gpu_matmul_q8_0_pair_tensor( - g->q_rank, - g->kv_raw, - model->map, - model->size, - l->attn_q_a->abs_offset, - l->attn_kv_a_mqa->abs_offset, - DS4_N_EMBD, - DS4_N_LORA_Q, - kv_raw_dim, - g->attn_norm, - 1) != 0; - if (qkv_pair_projected) { - static bool notice_printed = false; - if (!notice_printed) { - fprintf(stderr, - "ds4: ROCm GLM one-token Q/KV input " - "projections use the paired Q8 kernel\n"); - notice_printed = true; - } - } else { - static bool fallback_notice_printed = false; - if (!fallback_notice_printed) { - fprintf(stderr, - "ds4: ROCm GLM paired Q/KV projection " - "unavailable; using separate projections\n"); - fallback_notice_printed = true; - } - } - } - if (ok && - !qkv_pair_projected && - !(decode_ablate & DS4_GLM_ABLATE_QPATH)) { + if (ok && !(decode_ablate & DS4_GLM_ABLATE_QPATH)) { ok = glm_graph_matmul_q8_0_decode_profiled_tensor(g->q_rank, model, l->attn_q_a->abs_offset, @@ -45585,19 +45538,16 @@ static bool glm_graph_forward_token( g->compact_cache_cap != 0; const bool fuse_qkv_norm = !decode_stage_profile && !fuse_qkv_norm_store; if (ok && fuse_qkv_norm_store) { - if (!qkv_pair_projected) { - ok = glm_graph_matmul_q8_0_decode_profiled_tensor( - g->kv_raw, - model, - l->attn_kv_a_mqa->abs_offset, - DS4_N_EMBD, - kv_raw_dim, - g->attn_norm, - il, - pos, - "attn_kv_a_store", - g->ssd_streaming) != 0; - } + ok = glm_graph_matmul_q8_0_decode_profiled_tensor(g->kv_raw, + model, + l->attn_kv_a_mqa->abs_offset, + DS4_N_EMBD, + kv_raw_dim, + g->attn_norm, + il, + pos, + "attn_kv_a_store", + g->ssd_streaming) != 0; if (ok) { ok = ds4_gpu_glm_qkv_norm_store_compact_kv_tensor( g->q_rank_norm, @@ -45620,19 +45570,16 @@ static bool glm_graph_forward_token( DS4_RMS_EPS) != 0; } } else if (ok && fuse_qkv_norm) { - if (!qkv_pair_projected) { - ok = glm_graph_matmul_q8_0_decode_profiled_tensor( - g->kv_raw, - model, - l->attn_kv_a_mqa->abs_offset, - DS4_N_EMBD, - kv_raw_dim, - g->attn_norm, - il, - pos, - "attn_kv_a_norm", - g->ssd_streaming) != 0; - } + ok = glm_graph_matmul_q8_0_decode_profiled_tensor(g->kv_raw, + model, + l->attn_kv_a_mqa->abs_offset, + DS4_N_EMBD, + kv_raw_dim, + g->attn_norm, + il, + pos, + "attn_kv_a_norm", + g->ssd_streaming) != 0; if (ok) { ok = ds4_gpu_dsv4_qkv_rms_norm_rows_tensor(g->q_rank_norm, g->q_rank, @@ -45727,19 +45674,16 @@ static bool glm_graph_forward_token( } DS4_GLM_PROFILE_DECODE_STAGE("glm_decode_attn", "indexer_k"); if (ok && !fuse_qkv_norm && !fuse_qkv_norm_store) { - if (!qkv_pair_projected) { - ok = glm_graph_matmul_q8_0_decode_profiled_tensor( - g->kv_raw, - model, - l->attn_kv_a_mqa->abs_offset, - DS4_N_EMBD, - kv_raw_dim, - g->attn_norm, - il, - pos, - "attn_kv_a", - g->ssd_streaming) != 0; - } + ok = glm_graph_matmul_q8_0_decode_profiled_tensor(g->kv_raw, + model, + l->attn_kv_a_mqa->abs_offset, + DS4_N_EMBD, + kv_raw_dim, + g->attn_norm, + il, + pos, + "attn_kv_a", + g->ssd_streaming) != 0; if (ok) ok = ds4_gpu_glm_kv_lora_rms_norm_tensor(g->kv_norm, g->kv_raw, model->map, diff --git a/speed-bench/ROCM_GLM_DISTRIBUTED_VALIDATION.md b/speed-bench/ROCM_GLM_DISTRIBUTED_VALIDATION.md index 8c0e90340..66647f927 100644 --- a/speed-bench/ROCM_GLM_DISTRIBUTED_VALIDATION.md +++ b/speed-bench/ROCM_GLM_DISTRIBUTED_VALIDATION.md @@ -64,7 +64,6 @@ together for the final interaction check: | Experiment | Environment variable | Intended effect | | --- | --- | --- | | 64 KiB shared input | `DS4_ROCM_Q8_DECODE_SHAREDX_64K=1` | Let one-token Q8 projections with a 16,384-element input reuse the input through LDS. | -| Paired Q/KV projection | `DS4_ROCM_GLM_DECODE_QKV_PAIR=1` | Compute the one-token GLM Q and KV input projections together. | | Selected attention GEMM | `DS4_ROCM_GLM_SELECTED_ATTN_GEMM=1` | Gather per-token selected KV rows and use FP16 strided-batched BLAS after the 2,048-row indexer boundary. | Apply an enabled flag to **both** worker and coordinator. For the Podman worker, @@ -75,16 +74,13 @@ Use this order so one deployment answers both attribution and interaction: 1. no experimental flags; 2. `DS4_ROCM_Q8_DECODE_SHAREDX_64K=1`; -3. `DS4_ROCM_GLM_DECODE_QKV_PAIR=1`; -4. both decode flags; -5. `DS4_ROCM_GLM_SELECTED_ATTN_GEMM=1`; -6. all three flags. +3. `DS4_ROCM_GLM_SELECTED_ATTN_GEMM=1`; +4. both flags. The expected activation messages are: ```text Q8 one-token shared-input kernel enabled through 64 KiB LDS -ROCm GLM one-token Q/KV input projections use the paired Q8 kernel GLM selected indexed prefill using fp16 hipBLAS strided-batched attention GEMMs ``` @@ -92,10 +88,10 @@ Absence of a message means the tested workload did not reach that path. The 64 KiB path also reports and automatically falls back if the launch is not supported by the device. -The normal 2,048-token benchmark below exercises the two decode flags, but it -does **not** exercise selected attention: 2,048 visible rows still use the -dense causal path. Measure the prompt-processing cliff with a live 256-token -suffix curve: +The normal 2,048-token benchmark below exercises the decode flag, but it does +**not** exercise selected attention: 2,048 visible rows still use the dense +causal path. Measure the prompt-processing cliff with a live 256-token suffix +curve: ```sh RUN_NAME=selected-baseline From 7a84671c07bfb0636f12d9290028b023627cc2f3 Mon Sep 17 00:00:00 2001 From: Donato Capitella Date: Sun, 26 Jul 2026 14:56:43 +0100 Subject: [PATCH 22/33] Add opt-in GLM wave value projection --- rocm/ds4_rocm_glm.cuh | 73 +++++++++++++++++++ .../ROCM_GLM_DISTRIBUTED_VALIDATION.md | 12 ++- 2 files changed, 82 insertions(+), 3 deletions(-) diff --git a/rocm/ds4_rocm_glm.cuh b/rocm/ds4_rocm_glm.cuh index 0172db49b..ecec5bdf0 100644 --- a/rocm/ds4_rocm_glm.cuh +++ b/rocm/ds4_rocm_glm.cuh @@ -487,6 +487,46 @@ __global__ static void glm_q8_project_head_kernel( } } +/* Diagnostic port of the heterogeneous GLM branch's one-token Q8-head + * projection. One wave cooperates on each output row instead of assigning a + * complete serial dot product to one thread. The reduction changes FP32 + * summation order, so keep it opt-in until layer and frontier-logit comparisons + * establish the acceptable numerical envelope. */ +__global__ static void glm_q8_project_head_wave_kernel( + float *out, + const unsigned char *weight, + const float *x, + uint32_t n_head, + uint32_t in_dim, + uint32_t out_dim, + uint32_t x_stride, + uint32_t x_head_stride, + uint32_t row_bytes) { + const uint32_t lane = threadIdx.x & 31u; + const uint32_t wave = threadIdx.x >> 5u; + const uint32_t waves_per_block = blockDim.x >> 5u; + const uint32_t d = blockIdx.x * waves_per_block + wave; + const uint32_t head = blockIdx.y; + const uint32_t token = blockIdx.z; + if (d >= out_dim || head >= n_head) return; + + const float *xr = + x + (uint64_t)token * x_stride + (uint64_t)head * x_head_stride; + const unsigned char *row = + weight + ((uint64_t)head * out_dim + d) * row_bytes; + float acc = 0.0f; + for (uint32_t i = lane; i < in_dim; i += 32u) { + const unsigned char *blk = row + (uint64_t)(i >> 5u) * 34u; + const float scale = q8_0_scale_scalar(blk); + const int8_t q = ((const int8_t *)(blk + 2u))[i & 31u]; + acc += scale * (float)q * xr[i]; + } + acc = warp_sum_f32(acc); + if (lane == 0u) { + out[((uint64_t)token * n_head + head) * out_dim + d] = acc; + } +} + __global__ static void glm_build_kv_cache_kernel( char *key_cache, char *value_cache, @@ -1975,6 +2015,39 @@ extern "C" int ds4_gpu_glm_value_project_q8_0_batch_heads_tensor( return cuda_ok(cudaGetLastError(), "glm grouped value project batch heads launch"); } + if (n_tokens == 1u && + cuda_env_present( + getenv("DS4_ROCM_GLM_VALUE_PROJECT_WAVE_DECODE"))) { + const uint32_t threads = 256u; + const uint32_t waves_per_block = threads / 32u; + dim3 grid((value_dim + waves_per_block - 1u) / waves_per_block, + n_head, + n_tokens); + glm_q8_project_head_wave_kernel<<>>( + (float *)heads->ptr, + w, + (const float *)lora->ptr, + n_head, + kv_lora_dim, + value_dim, + (uint32_t)x_stride64, + kv_lora_dim, + row_bytes); + const cudaError_t launch_err = cudaGetLastError(); + if (launch_err == cudaSuccess) { + static int notice_printed = 0; + if (!notice_printed) { + fprintf(stderr, + DS4_GPU_LOG_PREFIX + "GLM one-token value projection using " + "wave-parallel Q8 rows\n"); + notice_printed = 1; + } + return 1; + } + return cuda_ok(launch_err, + "glm wave-parallel value project launch"); + } dim3 grid(n_head, n_tokens, 1); glm_q8_project_head_kernel<<>>( (float *)heads->ptr, diff --git a/speed-bench/ROCM_GLM_DISTRIBUTED_VALIDATION.md b/speed-bench/ROCM_GLM_DISTRIBUTED_VALIDATION.md index 66647f927..de75d9207 100644 --- a/speed-bench/ROCM_GLM_DISTRIBUTED_VALIDATION.md +++ b/speed-bench/ROCM_GLM_DISTRIBUTED_VALIDATION.md @@ -64,6 +64,7 @@ together for the final interaction check: | Experiment | Environment variable | Intended effect | | --- | --- | --- | | 64 KiB shared input | `DS4_ROCM_Q8_DECODE_SHAREDX_64K=1` | Let one-token Q8 projections with a 16,384-element input reuse the input through LDS. | +| Wave-parallel value projection | `DS4_ROCM_GLM_VALUE_PROJECT_WAVE_DECODE=1` | Give each one-token GLM value-projection output row a 32-lane wave instead of a serial thread. | | Selected attention GEMM | `DS4_ROCM_GLM_SELECTED_ATTN_GEMM=1` | Gather per-token selected KV rows and use FP16 strided-batched BLAS after the 2,048-row indexer boundary. | Apply an enabled flag to **both** worker and coordinator. For the Podman worker, @@ -74,19 +75,24 @@ Use this order so one deployment answers both attribution and interaction: 1. no experimental flags; 2. `DS4_ROCM_Q8_DECODE_SHAREDX_64K=1`; -3. `DS4_ROCM_GLM_SELECTED_ATTN_GEMM=1`; -4. both flags. +3. `DS4_ROCM_GLM_VALUE_PROJECT_WAVE_DECODE=1`; +4. both decode flags; +5. `DS4_ROCM_GLM_SELECTED_ATTN_GEMM=1`; +6. all flags. The expected activation messages are: ```text Q8 one-token shared-input kernel enabled through 64 KiB LDS +GLM one-token value projection using wave-parallel Q8 rows GLM selected indexed prefill using fp16 hipBLAS strided-batched attention GEMMs ``` Absence of a message means the tested workload did not reach that path. The 64 KiB path also reports and automatically falls back if the launch is not -supported by the device. +supported by the device. The wave-parallel value projection changes FP32 +summation order; validate its layer output and frontier logits even when its +generated text looks coherent. The normal 2,048-token benchmark below exercises the decode flag, but it does **not** exercise selected attention: 2,048 visible rows still use the dense From 34b3862d8038f2eb9aaa5ef5cbe9927346243603 Mon Sep 17 00:00:00 2001 From: Donato Capitella Date: Sun, 26 Jul 2026 18:33:29 +0100 Subject: [PATCH 23/33] Default GLM wave value projection --- rocm/ds4_rocm_glm.cuh | 9 +++-- .../ROCM_GLM_DISTRIBUTED_VALIDATION.md | 34 +++++++++---------- 2 files changed, 23 insertions(+), 20 deletions(-) diff --git a/rocm/ds4_rocm_glm.cuh b/rocm/ds4_rocm_glm.cuh index ecec5bdf0..3a1e6f6e1 100644 --- a/rocm/ds4_rocm_glm.cuh +++ b/rocm/ds4_rocm_glm.cuh @@ -2015,9 +2015,14 @@ extern "C" int ds4_gpu_glm_value_project_q8_0_batch_heads_tensor( return cuda_ok(cudaGetLastError(), "glm grouped value project batch heads launch"); } + const char *wave_decode_env = + getenv("DS4_ROCM_GLM_VALUE_PROJECT_WAVE_DECODE"); + /* + * The wave-per-output-row kernel is the validated production decode path. + * Keep an explicit =0 fallback for correctness diagnostics. + */ if (n_tokens == 1u && - cuda_env_present( - getenv("DS4_ROCM_GLM_VALUE_PROJECT_WAVE_DECODE"))) { + (wave_decode_env == NULL || cuda_env_present(wave_decode_env))) { const uint32_t threads = 256u; const uint32_t waves_per_block = threads / 32u; dim3 grid((value_dim + waves_per_block - 1u) / waves_per_block, diff --git a/speed-bench/ROCM_GLM_DISTRIBUTED_VALIDATION.md b/speed-bench/ROCM_GLM_DISTRIBUTED_VALIDATION.md index de75d9207..b20077d79 100644 --- a/speed-bench/ROCM_GLM_DISTRIBUTED_VALIDATION.md +++ b/speed-bench/ROCM_GLM_DISTRIBUTED_VALIDATION.md @@ -55,30 +55,28 @@ podman run --rm --userns=keep-id \ Require a successful warning-free build and run `git diff --check`. -## One-build experimental flag matrix +## One-build kernel flag matrix -The following ROCm experiments are independent and default to off. They can be -compiled into one image, enabled separately for attribution, and enabled -together for the final interaction check: +The following ROCm paths can be varied independently in one image for +attribution and interaction checks: -| Experiment | Environment variable | Intended effect | -| --- | --- | --- | -| 64 KiB shared input | `DS4_ROCM_Q8_DECODE_SHAREDX_64K=1` | Let one-token Q8 projections with a 16,384-element input reuse the input through LDS. | -| Wave-parallel value projection | `DS4_ROCM_GLM_VALUE_PROJECT_WAVE_DECODE=1` | Give each one-token GLM value-projection output row a 32-lane wave instead of a serial thread. | -| Selected attention GEMM | `DS4_ROCM_GLM_SELECTED_ATTN_GEMM=1` | Gather per-token selected KV rows and use FP16 strided-batched BLAS after the 2,048-row indexer boundary. | +| Path | Environment variable | Default | Intended effect | +| --- | --- | --- | --- | +| 64 KiB shared input | `DS4_ROCM_Q8_DECODE_SHAREDX_64K=1` | Off | Let one-token Q8 projections with a 16,384-element input reuse the input through LDS. | +| Wave-parallel value projection | `DS4_ROCM_GLM_VALUE_PROJECT_WAVE_DECODE=0` | On | Give each one-token GLM value-projection output row a 32-lane wave instead of a serial thread. Set `=0` only for a serial-row correctness baseline. | +| Selected attention GEMM | `DS4_ROCM_GLM_SELECTED_ATTN_GEMM=1` | Off | Gather per-token selected KV rows and use FP16 strided-batched BLAS after the 2,048-row indexer boundary. | -Apply an enabled flag to **both** worker and coordinator. For the Podman worker, -add it as `-e NAME=1`; for the coordinator, put `NAME=1 \` before `ds4-bench` -or `ds4-server`. +Apply any opt-in or fallback override to **both** worker and coordinator. For +the Podman worker, add it as `-e NAME=value`; for the coordinator, put +`NAME=value \` before `ds4-bench` or `ds4-server`. Use this order so one deployment answers both attribution and interaction: -1. no experimental flags; -2. `DS4_ROCM_Q8_DECODE_SHAREDX_64K=1`; -3. `DS4_ROCM_GLM_VALUE_PROJECT_WAVE_DECODE=1`; -4. both decode flags; -5. `DS4_ROCM_GLM_SELECTED_ATTN_GEMM=1`; -6. all flags. +1. `DS4_ROCM_GLM_VALUE_PROJECT_WAVE_DECODE=0` for the serial-row baseline; +2. no value-projection override for the production wave path; +3. production wave path plus `DS4_ROCM_Q8_DECODE_SHAREDX_64K=1`; +4. `DS4_ROCM_GLM_SELECTED_ATTN_GEMM=1`; +5. both opt-in flags with the production wave path. The expected activation messages are: From ea2e76681c531e3719b2f0154e31c6f1ef3898bf Mon Sep 17 00:00:00 2001 From: Donato Capitella Date: Sun, 26 Jul 2026 19:07:56 +0100 Subject: [PATCH 24/33] Default GLM selected attention GEMM --- rocm/ds4_rocm_glm.cuh | 10 +- .../ROCM_GLM_DISTRIBUTED_VALIDATION.md | 158 ++++++++++++++++-- 2 files changed, 148 insertions(+), 20 deletions(-) diff --git a/rocm/ds4_rocm_glm.cuh b/rocm/ds4_rocm_glm.cuh index 3a1e6f6e1..029461e5c 100644 --- a/rocm/ds4_rocm_glm.cuh +++ b/rocm/ds4_rocm_glm.cuh @@ -2879,10 +2879,14 @@ static int glm_attention_indexed_lora_launch( /* Once the visible context exceeds the indexer's top-k, each token carries * its own selected row list. Gather those rows into per-token FP16 * matrices, then batch the score and value products through BLAS. Keep - * the experiment opt-in and bounded: at 256 x 2048 selected rows the - * temporary workspace is already roughly 580 MiB. */ + * the path bounded: at 256 x 2048 selected rows the temporary workspace is + * already roughly 580 MiB. An explicit zero retains the scalar-kernel + * fallback for diagnosis and numerical comparisons. */ + const char *selected_attn_gemm_env = + getenv("DS4_ROCM_GLM_SELECTED_ATTN_GEMM"); if (!causal_range && has_selected && - cuda_env_present(getenv("DS4_ROCM_GLM_SELECTED_ATTN_GEMM")) && + (selected_attn_gemm_env == NULL || + cuda_env_present(selected_attn_gemm_env)) && glm_attention_indexed_lora_selected_gemm(lora_out, q, qk_low, diff --git a/speed-bench/ROCM_GLM_DISTRIBUTED_VALIDATION.md b/speed-bench/ROCM_GLM_DISTRIBUTED_VALIDATION.md index b20077d79..dbdfe0d7a 100644 --- a/speed-bench/ROCM_GLM_DISTRIBUTED_VALIDATION.md +++ b/speed-bench/ROCM_GLM_DISTRIBUTED_VALIDATION.md @@ -55,6 +55,75 @@ podman run --rm --userns=keep-id \ Require a successful warning-free build and run `git diff --check`. +## Publish and deploy the test image + +The image is built by +`kyuz0/strix-halo-ds4-toolbox/.github/workflows/build_and_publish.yml`. +Its `glm-rocm-7.2.4` Dockerfile clones +`https://github.com/kyuz0/ds4.git`, branch +`fix/rocm-distributed-glm`, so push the exact ds4 commit before dispatching the +workflow. The toolbox workflow itself runs from `main`. + +Dispatch only the GLM ROCm image and follow the run to completion: + +```sh +gh workflow run build_and_publish.yml \ + --repo kyuz0/strix-halo-ds4-toolbox \ + --ref main \ + -f backends=glm-rocm-7.2.4 + +gh run list \ + --repo kyuz0/strix-halo-ds4-toolbox \ + --workflow build_and_publish.yml \ + --limit 1 + +gh run watch RUN_ID \ + --repo kyuz0/strix-halo-ds4-toolbox \ + --exit-status +``` + +The workflow publishes both an immutable timestamped tag and the channel tag +used by this playbook: + +```text +docker.io/kyuz0/strix-halo-ds4-toolbox:glm-rocm-7.2.4 +``` + +Pull the channel tag on both hosts before testing. Do not assume a mutable tag +changed merely because `podman pull` succeeded; compare image IDs: + +```sh +ssh fw1 podman pull \ + docker.io/kyuz0/strix-halo-ds4-toolbox:glm-rocm-7.2.4 +ssh fw2 podman pull \ + docker.io/kyuz0/strix-halo-ds4-toolbox:glm-rocm-7.2.4 + +ssh fw1 podman image inspect \ + docker.io/kyuz0/strix-halo-ds4-toolbox:glm-rocm-7.2.4 \ + --format '{{.Id}} {{.Created}}' +ssh fw2 podman image inspect \ + docker.io/kyuz0/strix-halo-ds4-toolbox:glm-rocm-7.2.4 \ + --format '{{.Id}} {{.Created}}' +``` + +The IDs must match. The worker and coordinator examples below invoke the image +directly with Podman; entering a Toolbx shell is not required. Host `/tmp` +results must be mounted explicitly when the command writes inside the +container, for example: + +```sh +mkdir -p /tmp/glm-validation + +podman run --rm \ + --security-opt label=disable \ + -v /tmp/glm-validation:/results \ + IMAGE COMMAND --csv /results/run.csv +``` + +Keep logs outside the container with `2>&1 | tee /tmp/name.log`. Use distinct +container names and result directories for baseline and candidate, and stop +the worker after its coordinator run completes. + ## One-build kernel flag matrix The following ROCm paths can be varied independently in one image for @@ -64,7 +133,7 @@ attribution and interaction checks: | --- | --- | --- | --- | | 64 KiB shared input | `DS4_ROCM_Q8_DECODE_SHAREDX_64K=1` | Off | Let one-token Q8 projections with a 16,384-element input reuse the input through LDS. | | Wave-parallel value projection | `DS4_ROCM_GLM_VALUE_PROJECT_WAVE_DECODE=0` | On | Give each one-token GLM value-projection output row a 32-lane wave instead of a serial thread. Set `=0` only for a serial-row correctness baseline. | -| Selected attention GEMM | `DS4_ROCM_GLM_SELECTED_ATTN_GEMM=1` | Off | Gather per-token selected KV rows and use FP16 strided-batched BLAS after the 2,048-row indexer boundary. | +| Selected attention GEMM | `DS4_ROCM_GLM_SELECTED_ATTN_GEMM=0` | On | Gather per-token selected KV rows and use FP16 strided-batched BLAS after the 2,048-row indexer boundary. Set `=0` only for the scalar-kernel fallback. | Apply any opt-in or fallback override to **both** worker and coordinator. For the Podman worker, add it as `-e NAME=value`; for the coordinator, put @@ -75,8 +144,10 @@ Use this order so one deployment answers both attribution and interaction: 1. `DS4_ROCM_GLM_VALUE_PROJECT_WAVE_DECODE=0` for the serial-row baseline; 2. no value-projection override for the production wave path; 3. production wave path plus `DS4_ROCM_Q8_DECODE_SHAREDX_64K=1`; -4. `DS4_ROCM_GLM_SELECTED_ATTN_GEMM=1`; -5. both opt-in flags with the production wave path. +4. `DS4_ROCM_GLM_SELECTED_ATTN_GEMM=0` for the selected-attention scalar + fallback; +5. the production selected-attention GEMM path, with the shared-input + experiment either off or on as required. The expected activation messages are: @@ -120,11 +191,12 @@ ds4-bench \ 2>&1 | tee "/tmp/glm-coordinator-${RUN_NAME}.log" ``` -Run it once without selected attention and once with -`DS4_ROCM_GLM_SELECTED_ATTN_GEMM=1` on both nodes. The first 2,048-token row is -the causal control; the later 256-token suffix rows exercise selected -attention. The selected GEMM experiment accepts 64–256 tokens and uses about -580 MiB of temporary workspace at 256 tokens with 2,048 selected rows. +Run it once with `DS4_ROCM_GLM_SELECTED_ATTN_GEMM=0` on both nodes for the +scalar-kernel baseline and once without the override for the production GEMM +path. The first 2,048-token row is the causal control; the later 256-token +suffix rows exercise selected attention. The selected GEMM path accepts +64–256 tokens and uses about 580 MiB of temporary workspace at 256 tokens with +2,048 selected rows. For a layer-local selected-attention comparison, adapt the graph-dump command below to `--ctx-start 2304 --ctx-max 2304 --ctx-alloc 2433` and set @@ -153,7 +225,7 @@ the one candidate environment setting only for the candidate run. ```sh RUN_NAME=baseline -podman run --rm -it \ +podman run --rm \ --name "ds4-worker-${RUN_NAME}" \ -e DS4_GLM_MEMORY_GUARD_RESERVE_GB=8 \ -e DS4_ROCM_GLM_CAUSAL_ATTN_GEMM=1 \ @@ -182,23 +254,38 @@ podman run --rm -it \ ### Coordinator performance run Set `RUN_NAME` to match the worker. Add the one candidate environment setting -before `ds4-bench` only for the candidate run. +as another `-e NAME=value` only for the candidate run. ```sh RUN_NAME=baseline +mkdir -p "/tmp/glm-${RUN_NAME}" -DS4_GLM_MEMORY_GUARD_RESERVE_GB=8 \ -DS4_ROCM_GLM_CAUSAL_ATTN_GEMM=1 \ -DS4_BENCH_DISABLE_SNAPSHOT=1 \ -ds4-bench \ - --prompt-file ~/ds4/ds4/speed-bench/promessi_sposi.txt \ +podman run --rm \ + --name "ds4-coordinator-${RUN_NAME}" \ + -e DS4_GLM_MEMORY_GUARD_RESERVE_GB=8 \ + -e DS4_ROCM_GLM_CAUSAL_ATTN_GEMM=1 \ + -e DS4_BENCH_DISABLE_SNAPSHOT=1 \ + --device /dev/dri \ + --device /dev/kfd \ + --group-add keep-groups \ + --security-opt seccomp=unconfined \ + --ipc=host \ + --cap-add=SYS_PTRACE \ + --security-opt label=disable \ + --userns=keep-id \ + --network=host \ + -v /home/kyuz0/ds4:/models:ro \ + -v "/tmp/glm-${RUN_NAME}:/results" \ + docker.io/kyuz0/strix-halo-ds4-toolbox:glm-rocm-7.2.4 \ + ds4-bench \ + --prompt-file /models/ds4/speed-bench/promessi_sposi.txt \ --ctx-start 2048 \ --ctx-max 2048 \ --ctx-alloc 2177 \ --gen-tokens 128 \ --show-output \ - --csv "/tmp/glm-${RUN_NAME}.csv" \ - -m ~/ds4/GLM-5.2-UD-IQ2_XXS_RoutedIQ2XXS_blk78Q2K.gguf \ + --csv /results/run.csv \ + -m /models/GLM-5.2-UD-IQ2_XXS_RoutedIQ2XXS_blk78Q2K.gguf \ --dist-prefill-chunk 256 \ --dist-prefill-window 2 \ --role coordinator \ @@ -311,6 +398,43 @@ proof. Investigate a material regression from the last accepted path and use the official GLM continuation scorer before making an approximation the unconditional release default. +## Selected-attention promotion record + +On 2026-07-26, the selected-attention GEMM was compared with its scalar-kernel +fallback on the topology at the top of this document. Both runs used commit +`34b3862`, FP16 causal-attention GEMMs, a 2,304-token prompt, 256-token +distributed chunks, and a 2,433-token allocation. The only changed setting was +`DS4_ROCM_GLM_SELECTED_ATTN_GEMM=0` versus `=1`, on both nodes. + +The graph dump at layer 0 and position 2,048 produced: + +```text +all-token relative RMSE: 5.9441e-05 +all-token cosine: 0.999999998234 +final-token cosine: 0.999999997577 +non-finite pairs: 0 +``` + +The complete 154,880-value frontier-logit comparison produced: + +```text +KL(reference || candidate): 0.0664153 +top-1: 3956 -> 3956 +top-10 overlap: 10/10 +top-50 overlap: 46/50 +centered cosine: 0.9921320 +``` + +The dump-free performance pair improved from `19.11` to `27.48` prompt +tokens/s (`+43.8%`). The instrumented pair measured `18.85` versus `27.49` +tokens/s; do not compare either instrumented row with a normal benchmark +because serializing the full logit vector is included in the final chunk. + +Decision: make the selected-attention GEMM the ROCm GLM default and preserve +`DS4_ROCM_GLM_SELECTED_ATTN_GEMM=0` as the diagnostic rollback. The narrower +candidate top-1 margin (`1.2262` to `0.3320`) remains important context even +though the selected token and top-10 set matched. + ## API smoke test After tensor and logit checks pass, run a server and test the actual chat From 764987bc2cedaf1587ae9224c672a672b12ea499 Mon Sep 17 00:00:00 2001 From: Donato Capitella Date: Sun, 26 Jul 2026 19:27:15 +0100 Subject: [PATCH 25/33] Record GLM selected attention deployment --- .../ROCM_GLM_DISTRIBUTED_VALIDATION.md | 23 +++++++++++++++++++ 1 file changed, 23 insertions(+) diff --git a/speed-bench/ROCM_GLM_DISTRIBUTED_VALIDATION.md b/speed-bench/ROCM_GLM_DISTRIBUTED_VALIDATION.md index dbdfe0d7a..aaa54ac5c 100644 --- a/speed-bench/ROCM_GLM_DISTRIBUTED_VALIDATION.md +++ b/speed-bench/ROCM_GLM_DISTRIBUTED_VALIDATION.md @@ -435,6 +435,29 @@ Decision: make the selected-attention GEMM the ROCm GLM default and preserve candidate top-1 margin (`1.2262` to `0.3320`) remains important context even though the selected token and top-10 set matched. +The default-on change was published from ds4 commit `ea2e766` by toolbox +GitHub Actions run +[`30214076099`](https://github.com/kyuz0/strix-halo-ds4-toolbox/actions/runs/30214076099). +After pulling the channel tag, both hosts reported: + +```text +image ID: 379e27ad0439970bf72e04ac63fe2711e3471468f2171bad757ea03bca52d72c +digest: sha256:221bc3b031f5600094e40a1b0202fbd7f7057f31ab4854836bf97a5420dbc36f +created: 2026-07-26 18:13:38 UTC +``` + +A no-override post-deployment benchmark activated both the selected-attention +GEMM and wave-parallel one-token value projection and measured: + +```text +ctx=2304 prefill=27.45 t/s generation=2.30 t/s first-token=434.913 ms +``` + +The long API smoke used a 2,482-token chat prompt and 32 generated tokens. It +activated selected attention, completed prefill in `102.769 s`, decoded at +`2.29 t/s`, and returned a coherent description of the supplied passage as the +introduction to Manzoni's *I promessi sposi*. + ## API smoke test After tensor and logit checks pass, run a server and test the actual chat From dd837f033d8d99d3b1798d3536450027c412bfb2 Mon Sep 17 00:00:00 2001 From: Donato Capitella Date: Sun, 26 Jul 2026 19:35:26 +0100 Subject: [PATCH 26/33] Default GLM causal attention GEMM --- rocm/ds4_rocm_glm.cuh | 11 ++++--- .../ROCM_GLM_DISTRIBUTED_VALIDATION.md | 32 +++++++++---------- 2 files changed, 22 insertions(+), 21 deletions(-) diff --git a/rocm/ds4_rocm_glm.cuh b/rocm/ds4_rocm_glm.cuh index 029461e5c..801fee6d7 100644 --- a/rocm/ds4_rocm_glm.cuh +++ b/rocm/ds4_rocm_glm.cuh @@ -2923,11 +2923,14 @@ static int glm_attention_indexed_lora_launch( } return 1; } - /* Diagnostic port of the heterogeneous branch's dense causal prefill - * path. Keep it opt-in until remote gfx1151 timing and logit comparisons - * establish that the FP16 GEMMs are a safe default. */ + /* Dense causal prefill can use contiguous cache rows directly. Keep an + * explicit zero fallback for diagnosis and numerical comparisons with the + * scalar attention kernel. */ + const char *causal_attn_gemm_env = + getenv("DS4_ROCM_GLM_CAUSAL_ATTN_GEMM"); if (causal_range && !has_selected && - cuda_env_present(getenv("DS4_ROCM_GLM_CAUSAL_ATTN_GEMM")) && + (causal_attn_gemm_env == NULL || + cuda_env_present(causal_attn_gemm_env)) && glm_attention_indexed_lora_causal_gemm(lora_out, q, qk_low, diff --git a/speed-bench/ROCM_GLM_DISTRIBUTED_VALIDATION.md b/speed-bench/ROCM_GLM_DISTRIBUTED_VALIDATION.md index aaa54ac5c..e00df4974 100644 --- a/speed-bench/ROCM_GLM_DISTRIBUTED_VALIDATION.md +++ b/speed-bench/ROCM_GLM_DISTRIBUTED_VALIDATION.md @@ -133,20 +133,23 @@ attribution and interaction checks: | --- | --- | --- | --- | | 64 KiB shared input | `DS4_ROCM_Q8_DECODE_SHAREDX_64K=1` | Off | Let one-token Q8 projections with a 16,384-element input reuse the input through LDS. | | Wave-parallel value projection | `DS4_ROCM_GLM_VALUE_PROJECT_WAVE_DECODE=0` | On | Give each one-token GLM value-projection output row a 32-lane wave instead of a serial thread. Set `=0` only for a serial-row correctness baseline. | +| Causal attention GEMM | `DS4_ROCM_GLM_CAUSAL_ATTN_GEMM=0` | On | Use FP16 BLAS for contiguous causal attention during initial prefill. Set `=0` only for the scalar-kernel fallback. | | Selected attention GEMM | `DS4_ROCM_GLM_SELECTED_ATTN_GEMM=0` | On | Gather per-token selected KV rows and use FP16 strided-batched BLAS after the 2,048-row indexer boundary. Set `=0` only for the scalar-kernel fallback. | Apply any opt-in or fallback override to **both** worker and coordinator. For -the Podman worker, add it as `-e NAME=value`; for the coordinator, put -`NAME=value \` before `ds4-bench` or `ds4-server`. +Podman, add it as `-e NAME=value`; for a host binary, put `NAME=value \` +before `ds4-bench` or `ds4-server`. Use this order so one deployment answers both attribution and interaction: -1. `DS4_ROCM_GLM_VALUE_PROJECT_WAVE_DECODE=0` for the serial-row baseline; -2. no value-projection override for the production wave path; -3. production wave path plus `DS4_ROCM_Q8_DECODE_SHAREDX_64K=1`; -4. `DS4_ROCM_GLM_SELECTED_ATTN_GEMM=0` for the selected-attention scalar +1. `DS4_ROCM_GLM_CAUSAL_ATTN_GEMM=0` for the causal-attention scalar fallback; +2. no causal-attention override for the production causal GEMM path; +3. `DS4_ROCM_GLM_VALUE_PROJECT_WAVE_DECODE=0` for the serial-row baseline; +4. no value-projection override for the production wave path; +5. production wave path plus `DS4_ROCM_Q8_DECODE_SHAREDX_64K=1`; +6. `DS4_ROCM_GLM_SELECTED_ATTN_GEMM=0` for the selected-attention scalar fallback; -5. the production selected-attention GEMM path, with the shared-input +7. the production selected-attention GEMM path, with the shared-input experiment either off or on as required. The expected activation messages are: @@ -154,6 +157,7 @@ The expected activation messages are: ```text Q8 one-token shared-input kernel enabled through 64 KiB LDS GLM one-token value projection using wave-parallel Q8 rows +GLM causal indexed prefill using fp16 hipBLAS attention GEMMs GLM selected indexed prefill using fp16 hipBLAS strided-batched attention GEMMs ``` @@ -172,7 +176,6 @@ curve: RUN_NAME=selected-baseline DS4_GLM_MEMORY_GUARD_RESERVE_GB=8 \ -DS4_ROCM_GLM_CAUSAL_ATTN_GEMM=1 \ DS4_BENCH_DISABLE_SNAPSHOT=1 \ ds4-bench \ --prompt-file ~/ds4/ds4/speed-bench/promessi_sposi.txt \ @@ -211,11 +214,10 @@ Run the commands below once without it for the baseline, then again with only that setting changed for the candidate. If the setting affects both halves of the model, apply it to both worker and coordinator. -Keep the accepted FP16 causal-attention GEMM enabled in both runs: - -```text -DS4_ROCM_GLM_CAUSAL_ATTN_GEMM=1 -``` +The accepted FP16 causal-attention GEMM is enabled by default. Leave it without +an override in ordinary baseline/candidate comparisons. When validating the +causal GEMM itself, use `DS4_ROCM_GLM_CAUSAL_ATTN_GEMM=0` on both nodes for +the scalar-kernel reference and no override for the production candidate. ### Worker @@ -228,7 +230,6 @@ RUN_NAME=baseline podman run --rm \ --name "ds4-worker-${RUN_NAME}" \ -e DS4_GLM_MEMORY_GUARD_RESERVE_GB=8 \ - -e DS4_ROCM_GLM_CAUSAL_ATTN_GEMM=1 \ --device /dev/dri \ --device /dev/kfd \ --group-add keep-groups \ @@ -263,7 +264,6 @@ mkdir -p "/tmp/glm-${RUN_NAME}" podman run --rm \ --name "ds4-coordinator-${RUN_NAME}" \ -e DS4_GLM_MEMORY_GUARD_RESERVE_GB=8 \ - -e DS4_ROCM_GLM_CAUSAL_ATTN_GEMM=1 \ -e DS4_BENCH_DISABLE_SNAPSHOT=1 \ --device /dev/dri \ --device /dev/kfd \ @@ -311,7 +311,6 @@ RUN_NAME=baseline mkdir -p /tmp/glm-validation DS4_GLM_MEMORY_GUARD_RESERVE_GB=8 \ -DS4_ROCM_GLM_CAUSAL_ATTN_GEMM=1 \ DS4_ROCM_GRAPH_DUMP_PREFIX="/tmp/glm-validation/${RUN_NAME}" \ DS4_ROCM_GRAPH_DUMP_NONINVASIVE=1 \ DS4_ROCM_GRAPH_DUMP_LAYER=0 \ @@ -357,7 +356,6 @@ RUN_NAME=baseline mkdir -p "/tmp/glm-logits-${RUN_NAME}" DS4_GLM_MEMORY_GUARD_RESERVE_GB=8 \ -DS4_ROCM_GLM_CAUSAL_ATTN_GEMM=1 \ DS4_BENCH_DISABLE_SNAPSHOT=1 \ ds4-bench \ --prompt-file ~/ds4/ds4/speed-bench/promessi_sposi.txt \ From cfb1ff9a07870966a375711695f47d59c673eab9 Mon Sep 17 00:00:00 2001 From: Donato Capitella Date: Sun, 26 Jul 2026 19:57:09 +0100 Subject: [PATCH 27/33] Record GLM causal attention deployment --- .../ROCM_GLM_DISTRIBUTED_VALIDATION.md | 46 +++++++++++++++++++ 1 file changed, 46 insertions(+) diff --git a/speed-bench/ROCM_GLM_DISTRIBUTED_VALIDATION.md b/speed-bench/ROCM_GLM_DISTRIBUTED_VALIDATION.md index e00df4974..5b6741c33 100644 --- a/speed-bench/ROCM_GLM_DISTRIBUTED_VALIDATION.md +++ b/speed-bench/ROCM_GLM_DISTRIBUTED_VALIDATION.md @@ -456,6 +456,52 @@ activated selected attention, completed prefill in `102.769 s`, decoded at `2.29 t/s`, and returned a coherent description of the supplied passage as the introduction to Manzoni's *I promessi sposi*. +## Causal-attention default deployment record + +The causal-attention GEMM became default-on in ds4 commit `dd837f0`, published +by toolbox GitHub Actions run +[`30215053603`](https://github.com/kyuz0/strix-halo-ds4-toolbox/actions/runs/30215053603). +Both hosts pulled and verified: + +```text +image ID: 6f424ab78c427808d7d7cee3b4a0ea01b232223a0ca525ebfb033042dec4c4a0 +digest: sha256:9bda0b1d4c45c08bc06caffa90262a964d935534353595b199eab98ce9e6c5cb +created: 2026-07-26 18:41:10 UTC +``` + +The post-deployment benchmark used no kernel environment overrides, layers +`0:37` plus `38:output`, a 2,433-token allocation, 128 generated tokens, and +the production distributed prefill settings: + +```text +--dist-prefill-chunk 256 +--dist-prefill-window 2 +``` + +Chunk 256 is intentional. Earlier matched 2,048-token runs measured +`50.81 t/s` with chunk 256 versus `48.15 t/s` with chunk 512. + +The production curve measured: + +```text +ctx prefill tokens prefill t/s generation t/s first token +2048 2048 50.55 2.30 439.299 ms +2304 256 5.22 2.30 434.724 ms +``` + +The second row is the incremental selected-attention suffix at an existing +2,048-token frontier, not full-prompt throughput. It exposes the remaining +post-2,048 prefill cliff. A separate cold 2,304-token full-prompt run measured: + +```text +prefill=27.46 t/s generation=2.30 t/s first-token=439.220 ms +``` + +All three default activation messages were present: causal attention GEMM, +selected attention GEMM, and wave-parallel one-token value projection. Both +128-token raw continuations were coherent. Preserve +`DS4_ROCM_GLM_CAUSAL_ATTN_GEMM=0` as the independent scalar-kernel rollback. + ## API smoke test After tensor and logit checks pass, run a server and test the actual chat From 4d7fbbf5f3b4a32630879fbd2b2019864621b90d Mon Sep 17 00:00:00 2001 From: Donato Capitella Date: Sun, 26 Jul 2026 20:25:30 +0100 Subject: [PATCH 28/33] Add GLM selected attention head tiling --- rocm/ds4_rocm_glm.cuh | 411 ++++++++++++++++-- .../ROCM_GLM_DISTRIBUTED_VALIDATION.md | 13 +- 2 files changed, 393 insertions(+), 31 deletions(-) diff --git a/rocm/ds4_rocm_glm.cuh b/rocm/ds4_rocm_glm.cuh index 801fee6d7..e17bdef69 100644 --- a/rocm/ds4_rocm_glm.cuh +++ b/rocm/ds4_rocm_glm.cuh @@ -2335,6 +2335,135 @@ __global__ static void glm_selected_gemm_softmax_f16_kernel( } } +__global__ static void glm_selected_gemm_gather_heads_f16_kernel( + __half *low_h, + __half *qrope_h, + const float *low, + const float *q, + uint32_t n_tokens, + uint32_t head0, + uint32_t head_count, + uint32_t head_capacity, + uint32_t n_head, + uint32_t kv_lora_dim, + uint32_t qk_nope, + uint32_t qk_rope) { + const uint64_t low_per_token = + (uint64_t)head_capacity * kv_lora_dim; + const uint64_t low_total = (uint64_t)n_tokens * low_per_token; + for (uint64_t i = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x; + i < low_total; + i += (uint64_t)gridDim.x * blockDim.x) { + const uint32_t token = (uint32_t)(i / low_per_token); + const uint64_t token_i = i - (uint64_t)token * low_per_token; + const uint32_t head_local = (uint32_t)(token_i / kv_lora_dim); + const uint32_t j = + (uint32_t)(token_i - (uint64_t)head_local * kv_lora_dim); + if (head_local < head_count) { + low_h[i] = __float2half( + low[((uint64_t)token * n_head + head0 + head_local) * + kv_lora_dim + j]); + } + } + + const uint64_t rope_per_token = + (uint64_t)head_capacity * qk_rope; + const uint64_t rope_total = (uint64_t)n_tokens * rope_per_token; + const uint32_t qk_dim = qk_nope + qk_rope; + for (uint64_t i = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x; + i < rope_total; + i += (uint64_t)gridDim.x * blockDim.x) { + const uint32_t token = (uint32_t)(i / rope_per_token); + const uint64_t token_i = i - (uint64_t)token * rope_per_token; + const uint32_t head_local = (uint32_t)(token_i / qk_rope); + const uint32_t r = + (uint32_t)(token_i - (uint64_t)head_local * qk_rope); + if (head_local < head_count) { + qrope_h[i] = __float2half( + q[((uint64_t)token * n_head + head0 + head_local) * + qk_dim + qk_nope + r]); + } + } +} + +__global__ static void glm_selected_gemm_softmax_heads_f16_kernel( + __half *probs, + float *scores, + uint32_t n_tokens, + uint32_t n_selected, + uint32_t head_count, + uint32_t head_capacity, + float scale) { + const uint32_t token = blockIdx.x; + const uint32_t head_local = blockIdx.y; + if (token >= n_tokens || head_local >= head_count) return; + const uint64_t row_index = + (uint64_t)token * head_capacity + head_local; + float *row = scores + row_index * n_selected; + __shared__ float reduce[256]; + float vmax = -INFINITY; + for (uint32_t s = threadIdx.x; s < n_selected; s += blockDim.x) { + const float v = row[s] * scale; + row[s] = v; + vmax = fmaxf(vmax, v); + } + reduce[threadIdx.x] = vmax; + __syncthreads(); + for (uint32_t stride = blockDim.x >> 1u; stride != 0u; stride >>= 1u) { + if (threadIdx.x < stride) { + reduce[threadIdx.x] = + fmaxf(reduce[threadIdx.x], reduce[threadIdx.x + stride]); + } + __syncthreads(); + } + const float max_score = reduce[0]; + float sum = 0.0f; + for (uint32_t s = threadIdx.x; s < n_selected; s += blockDim.x) { + const float v = expf(row[s] - max_score); + row[s] = v; + sum += v; + } + reduce[threadIdx.x] = sum; + __syncthreads(); + for (uint32_t stride = blockDim.x >> 1u; stride != 0u; stride >>= 1u) { + if (threadIdx.x < stride) { + reduce[threadIdx.x] += reduce[threadIdx.x + stride]; + } + __syncthreads(); + } + const float inv = 1.0f / fmaxf(reduce[0], 1.0e-20f); + __half *prob_row = probs + row_index * n_selected; + for (uint32_t s = threadIdx.x; s < n_selected; s += blockDim.x) { + prob_row[s] = __float2half(row[s] * inv); + } +} + +__global__ static void glm_selected_gemm_scatter_heads_kernel( + float *out, + const float *head_out, + uint32_t n_tokens, + uint32_t head0, + uint32_t head_count, + uint32_t head_capacity, + uint32_t n_head, + uint32_t kv_lora_dim) { + const uint64_t n = + (uint64_t)n_tokens * head_count * kv_lora_dim; + const uint64_t i = + (uint64_t)blockIdx.x * blockDim.x + threadIdx.x; + if (i >= n) return; + const uint64_t per_token = (uint64_t)head_count * kv_lora_dim; + const uint32_t token = (uint32_t)(i / per_token); + const uint64_t token_i = i - (uint64_t)token * per_token; + const uint32_t head_local = (uint32_t)(token_i / kv_lora_dim); + const uint32_t j = + (uint32_t)(token_i - (uint64_t)head_local * kv_lora_dim); + out[((uint64_t)token * n_head + head0 + head_local) * + kv_lora_dim + j] = + head_out[((uint64_t)token * head_capacity + head_local) * + kv_lora_dim + j]; +} + static int glm_causal_gemm_scratch_part( uint64_t *offset, uint64_t bytes, @@ -2350,6 +2479,134 @@ static int glm_causal_gemm_scratch_part( return 1; } +enum glm_selected_gemm_profile_phase { + GLM_SELECTED_PROFILE_KV_GATHER = 0, + GLM_SELECTED_PROFILE_ROPE_GATHER, + GLM_SELECTED_PROFILE_QUERY_GATHER, + GLM_SELECTED_PROFILE_LORA_SCORE, + GLM_SELECTED_PROFILE_ROPE_SCORE, + GLM_SELECTED_PROFILE_SOFTMAX, + GLM_SELECTED_PROFILE_VALUE, + GLM_SELECTED_PROFILE_SCATTER, + GLM_SELECTED_PROFILE_PHASE_COUNT +}; + +struct glm_selected_gemm_profile { + int enabled; + float ms[GLM_SELECTED_PROFILE_PHASE_COUNT]; +}; + +static cudaEvent_t g_glm_selected_profile_begin; +static cudaEvent_t g_glm_selected_profile_end; +static int g_glm_selected_profile_events_ready = 0; +static int g_glm_selected_profile_claimed = 0; + +static uint32_t glm_selected_gemm_head_tile(void) { + static int initialized = 0; + static uint32_t tile = 1u; + if (initialized) return tile; + initialized = 1; + const char *env = getenv("DS4_ROCM_GLM_SELECTED_ATTN_HEAD_TILE"); + if (!env || !env[0]) return tile; + char *end = NULL; + errno = 0; + const unsigned long value = strtoul(env, &end, 10); + if (errno == 0 && end && *end == '\0' && + (value == 1ul || value == 2ul || + value == 4ul || value == 8ul)) { + tile = (uint32_t)value; + return tile; + } + fprintf(stderr, + "ds4: invalid DS4_ROCM_GLM_SELECTED_ATTN_HEAD_TILE='%s'; " + "using 1 (valid values: 1,2,4,8)\n", + env); + return tile; +} + +static void glm_selected_gemm_profile_init( + glm_selected_gemm_profile *profile) { + if (!profile) return; + memset(profile, 0, sizeof(*profile)); + if (!cuda_env_present( + getenv("DS4_ROCM_GLM_SELECTED_ATTN_PROFILE")) || + g_glm_selected_profile_claimed) { + return; + } + g_glm_selected_profile_claimed = 1; + if (!g_glm_selected_profile_events_ready) { + if (cudaEventCreate(&g_glm_selected_profile_begin) != cudaSuccess) { + return; + } + if (cudaEventCreate(&g_glm_selected_profile_end) != cudaSuccess) { + (void)cudaEventDestroy(g_glm_selected_profile_begin); + return; + } + g_glm_selected_profile_events_ready = 1; + } + profile->enabled = 1; +} + +static void glm_selected_gemm_profile_begin( + glm_selected_gemm_profile *profile) { + if (!profile || !profile->enabled) return; + if (cudaEventRecord(g_glm_selected_profile_begin, 0) != cudaSuccess) { + profile->enabled = 0; + } +} + +static void glm_selected_gemm_profile_end( + glm_selected_gemm_profile *profile, + glm_selected_gemm_profile_phase phase) { + if (!profile || !profile->enabled || + phase >= GLM_SELECTED_PROFILE_PHASE_COUNT) { + return; + } + float elapsed = 0.0f; + if (cudaEventRecord(g_glm_selected_profile_end, 0) != cudaSuccess || + cudaEventSynchronize(g_glm_selected_profile_end) != cudaSuccess || + cudaEventElapsedTime(&elapsed, + g_glm_selected_profile_begin, + g_glm_selected_profile_end) != cudaSuccess) { + profile->enabled = 0; + return; + } + profile->ms[phase] += elapsed; +} + +static void glm_selected_gemm_profile_report( + const glm_selected_gemm_profile *profile, + uint32_t n_tokens, + uint32_t n_selected, + uint32_t n_head, + uint32_t head_tile) { + if (!profile || !profile->enabled) return; + float total = 0.0f; + for (uint32_t i = 0; i < GLM_SELECTED_PROFILE_PHASE_COUNT; i++) { + total += profile->ms[i]; + } + fprintf(stderr, + "ds4: ROCm GLM selected attention profile " + "tokens=%u selected=%u heads=%u tile=%u " + "kv_gather=%.3fms rope_gather=%.3fms " + "query_gather=%.3fms lora_score=%.3fms " + "rope_score=%.3fms softmax=%.3fms " + "value=%.3fms scatter=%.3fms total=%.3fms\n", + n_tokens, + n_selected, + n_head, + head_tile, + profile->ms[GLM_SELECTED_PROFILE_KV_GATHER], + profile->ms[GLM_SELECTED_PROFILE_ROPE_GATHER], + profile->ms[GLM_SELECTED_PROFILE_QUERY_GATHER], + profile->ms[GLM_SELECTED_PROFILE_LORA_SCORE], + profile->ms[GLM_SELECTED_PROFILE_ROPE_SCORE], + profile->ms[GLM_SELECTED_PROFILE_SOFTMAX], + profile->ms[GLM_SELECTED_PROFILE_VALUE], + profile->ms[GLM_SELECTED_PROFILE_SCATTER], + total); +} + static int glm_attention_indexed_lora_causal_gemm( ds4_gpu_tensor *lora_out, const ds4_gpu_tensor *q, @@ -2599,6 +2856,18 @@ static int glm_attention_indexed_lora_selected_gemm( return 0; } + const uint32_t requested_head_tile = glm_selected_gemm_head_tile(); + const uint32_t head_tile = min(requested_head_tile, n_head); + if (head_tile > 1u) { + static int notice_printed = 0; + if (!notice_printed) { + fprintf(stderr, + "ds4: ROCm GLM selected attention using head tile %u\n", + head_tile); + notice_printed = 1; + } + } + uint64_t selected_rows = 0; uint64_t kv_count = 0, rope_count = 0; uint64_t low_count = 0, qrope_count = 0; @@ -2612,6 +2881,12 @@ static int glm_attention_indexed_lora_selected_gemm( !cuda_u64_mul_checked(n_tokens, kv_lora_dim, &head_count)) { return 0; } + if (!cuda_u64_mul_checked(low_count, head_tile, &low_count) || + !cuda_u64_mul_checked(qrope_count, head_tile, &qrope_count) || + !cuda_u64_mul_checked(score_count, head_tile, &score_count) || + !cuda_u64_mul_checked(head_count, head_tile, &head_count)) { + return 0; + } uint64_t kv_bytes = 0, rope_bytes = 0; uint64_t low_bytes = 0, qrope_bytes = 0; @@ -2651,8 +2926,12 @@ static int glm_attention_indexed_lora_selected_gemm( __half *probs = (__half *)(scratch + prob_offset); float *head_out = (float *)(scratch + head_offset); + glm_selected_gemm_profile profile; + glm_selected_gemm_profile_init(&profile); + const uint32_t kv_blocks = (uint32_t)min((kv_count + 255u) / 256u, 4096ull); + glm_selected_gemm_profile_begin(&profile); glm_selected_gemm_kv_to_f16_kernel<<>>( kv_h, (const char *)kv_lora_cache->ptr, @@ -2666,9 +2945,12 @@ static int glm_attention_indexed_lora_selected_gemm( "glm selected attention kv gather launch")) { return 0; } + glm_selected_gemm_profile_end( + &profile, GLM_SELECTED_PROFILE_KV_GATHER); const uint64_t rope_pairs = selected_rows * (qk_rope >> 1u); const uint32_t rope_blocks = (uint32_t)min((rope_pairs + 255u) / 256u, 4096ull); + glm_selected_gemm_profile_begin(&profile); glm_selected_gemm_rope_to_f16_kernel<<>>( rope_h, (const char *)k_rope_cache->ptr, @@ -2689,43 +2971,67 @@ static int glm_attention_indexed_lora_selected_gemm( "glm selected attention rope gather launch")) { return 0; } + glm_selected_gemm_profile_end( + &profile, GLM_SELECTED_PROFILE_ROPE_GATHER); const float one = 1.0f; const float zero = 0.0f; const float scale = 1.0f / sqrtf((float)(qk_nope + qk_rope)); const uint32_t gather_blocks = - (uint32_t)((low_count + 255u) / 256u); - const uint32_t scatter_blocks = - (uint32_t)((head_count + 255u) / 256u); + (uint32_t)min((low_count + 255u) / 256u, 4096ull); const long long kv_stride = (long long)n_selected * (long long)kv_lora_dim; const long long rope_stride = (long long)n_selected * (long long)qk_rope; - const long long low_stride = (long long)kv_lora_dim; - const long long qrope_stride = (long long)qk_rope; - const long long score_stride = (long long)n_selected; - for (uint32_t head = 0; head < n_head; head++) { - glm_causal_gemm_gather_head_f16_kernel<<>>( - low_h, - qrope_h, - (const float *)qk_low->ptr, - (const float *)q->ptr, - n_tokens, - head, - n_head, - kv_lora_dim, - qk_nope, - qk_rope); + const long long low_stride = + (long long)head_tile * (long long)kv_lora_dim; + const long long qrope_stride = + (long long)head_tile * (long long)qk_rope; + const long long score_stride = + (long long)head_tile * (long long)n_selected; + for (uint32_t head0 = 0; head0 < n_head; head0 += head_tile) { + const uint32_t tile_heads = min(head_tile, n_head - head0); + glm_selected_gemm_profile_begin(&profile); + if (head_tile == 1u) { + glm_causal_gemm_gather_head_f16_kernel<<>>( + low_h, + qrope_h, + (const float *)qk_low->ptr, + (const float *)q->ptr, + n_tokens, + head0, + n_head, + kv_lora_dim, + qk_nope, + qk_rope); + } else { + glm_selected_gemm_gather_heads_f16_kernel<<>>( + low_h, + qrope_h, + (const float *)qk_low->ptr, + (const float *)q->ptr, + n_tokens, + head0, + tile_heads, + head_tile, + n_head, + kv_lora_dim, + qk_nope, + qk_rope); + } if (!cuda_ok(cudaGetLastError(), "glm selected attention query gather launch")) { return 0; } + glm_selected_gemm_profile_end( + &profile, GLM_SELECTED_PROFILE_QUERY_GATHER); + glm_selected_gemm_profile_begin(&profile); cublasStatus_t st = cublasGemmStridedBatchedEx( g_cublas, CUBLAS_OP_T, CUBLAS_OP_N, (int)n_selected, - 1, + (int)tile_heads, (int)kv_lora_dim, &one, kv_h, @@ -2747,12 +3053,15 @@ static int glm_attention_indexed_lora_selected_gemm( if (!cublas_ok(st, "glm selected attention lora score gemm")) { return 0; } + glm_selected_gemm_profile_end( + &profile, GLM_SELECTED_PROFILE_LORA_SCORE); + glm_selected_gemm_profile_begin(&profile); st = cublasGemmStridedBatchedEx( g_cublas, CUBLAS_OP_T, CUBLAS_OP_N, (int)n_selected, - 1, + (int)tile_heads, (int)qk_rope, &one, rope_h, @@ -2774,18 +3083,37 @@ static int glm_attention_indexed_lora_selected_gemm( if (!cublas_ok(st, "glm selected attention rope score gemm")) { return 0; } - glm_selected_gemm_softmax_f16_kernel<<>>( - probs, scores, n_tokens, n_selected, scale); + glm_selected_gemm_profile_end( + &profile, GLM_SELECTED_PROFILE_ROPE_SCORE); + glm_selected_gemm_profile_begin(&profile); + if (head_tile == 1u) { + glm_selected_gemm_softmax_f16_kernel<<>>( + probs, scores, n_tokens, n_selected, scale); + } else { + const dim3 softmax_grid(n_tokens, tile_heads, 1u); + glm_selected_gemm_softmax_heads_f16_kernel<<< + softmax_grid, 256>>>( + probs, + scores, + n_tokens, + n_selected, + tile_heads, + head_tile, + scale); + } if (!cuda_ok(cudaGetLastError(), "glm selected attention softmax launch")) { return 0; } + glm_selected_gemm_profile_end( + &profile, GLM_SELECTED_PROFILE_SOFTMAX); + glm_selected_gemm_profile_begin(&profile); st = cublasGemmStridedBatchedEx( g_cublas, CUBLAS_OP_N, CUBLAS_OP_N, (int)kv_lora_dim, - 1, + (int)tile_heads, (int)n_selected, &one, kv_h, @@ -2807,18 +3135,41 @@ static int glm_attention_indexed_lora_selected_gemm( if (!cublas_ok(st, "glm selected attention value gemm")) { return 0; } - glm_causal_gemm_scatter_head_kernel<<>>( - (float *)lora_out->ptr, - head_out, - n_tokens, - head, - n_head, - kv_lora_dim); + glm_selected_gemm_profile_end( + &profile, GLM_SELECTED_PROFILE_VALUE); + const uint64_t scatter_count = + (uint64_t)n_tokens * tile_heads * kv_lora_dim; + const uint32_t scatter_blocks = + (uint32_t)((scatter_count + 255u) / 256u); + glm_selected_gemm_profile_begin(&profile); + if (head_tile == 1u) { + glm_causal_gemm_scatter_head_kernel<<>>( + (float *)lora_out->ptr, + head_out, + n_tokens, + head0, + n_head, + kv_lora_dim); + } else { + glm_selected_gemm_scatter_heads_kernel<<>>( + (float *)lora_out->ptr, + head_out, + n_tokens, + head0, + tile_heads, + head_tile, + n_head, + kv_lora_dim); + } if (!cuda_ok(cudaGetLastError(), "glm selected attention head scatter launch")) { return 0; } + glm_selected_gemm_profile_end( + &profile, GLM_SELECTED_PROFILE_SCATTER); } + glm_selected_gemm_profile_report( + &profile, n_tokens, n_selected, n_head, head_tile); return 1; } diff --git a/speed-bench/ROCM_GLM_DISTRIBUTED_VALIDATION.md b/speed-bench/ROCM_GLM_DISTRIBUTED_VALIDATION.md index 5b6741c33..fd6ea7760 100644 --- a/speed-bench/ROCM_GLM_DISTRIBUTED_VALIDATION.md +++ b/speed-bench/ROCM_GLM_DISTRIBUTED_VALIDATION.md @@ -135,6 +135,8 @@ attribution and interaction checks: | Wave-parallel value projection | `DS4_ROCM_GLM_VALUE_PROJECT_WAVE_DECODE=0` | On | Give each one-token GLM value-projection output row a 32-lane wave instead of a serial thread. Set `=0` only for a serial-row correctness baseline. | | Causal attention GEMM | `DS4_ROCM_GLM_CAUSAL_ATTN_GEMM=0` | On | Use FP16 BLAS for contiguous causal attention during initial prefill. Set `=0` only for the scalar-kernel fallback. | | Selected attention GEMM | `DS4_ROCM_GLM_SELECTED_ATTN_GEMM=0` | On | Gather per-token selected KV rows and use FP16 strided-batched BLAS after the 2,048-row indexer boundary. Set `=0` only for the scalar-kernel fallback. | +| Selected-attention head tile | `DS4_ROCM_GLM_SELECTED_ATTN_HEAD_TILE=1\|2\|4\|8` | `1` | Process several attention heads as BLAS matrix columns. Tile 1 preserves the production path. | +| Selected-attention phase profile | `DS4_ROCM_GLM_SELECTED_ATTN_PROFILE=1` | Off | Print HIP-event phase totals for the first selected-attention invocation in each process. Diagnostic only. | Apply any opt-in or fallback override to **both** worker and coordinator. For Podman, add it as `-e NAME=value`; for a host binary, put `NAME=value \` @@ -159,6 +161,8 @@ Q8 one-token shared-input kernel enabled through 64 KiB LDS GLM one-token value projection using wave-parallel Q8 rows GLM causal indexed prefill using fp16 hipBLAS attention GEMMs GLM selected indexed prefill using fp16 hipBLAS strided-batched attention GEMMs +ROCm GLM selected attention using head tile N +ROCm GLM selected attention profile tokens=... selected=... heads=... tile=... ``` Absence of a message means the tested workload did not reach that path. The @@ -199,7 +203,14 @@ scalar-kernel baseline and once without the override for the production GEMM path. The first 2,048-token row is the causal control; the later 256-token suffix rows exercise selected attention. The selected GEMM path accepts 64–256 tokens and uses about 580 MiB of temporary workspace at 256 tokens with -2,048 selected rows. +2,048 selected rows at head tile 1. Tiles 2, 4, and 8 use approximately 584, +591, and 606 MiB respectively. + +The phase profiler intentionally synchronizes at each measured boundary, so +never use its wall-clock benchmark row as a performance result. It claims only +the first selected-attention invocation per process. Run profiling separately +on worker and coordinator, then disable it for all matched performance and +correctness runs. For a layer-local selected-attention comparison, adapt the graph-dump command below to `--ctx-start 2304 --ctx-max 2304 --ctx-alloc 2433` and set From 367a64c1bc4bf50611ff3dff36f6c592f8928fe0 Mon Sep 17 00:00:00 2001 From: Donato Capitella Date: Sun, 26 Jul 2026 20:58:15 +0100 Subject: [PATCH 29/33] Extend GLM selected attention head tiling --- rocm/ds4_rocm_glm.cuh | 6 ++++-- speed-bench/ROCM_GLM_DISTRIBUTED_VALIDATION.md | 6 +++--- 2 files changed, 7 insertions(+), 5 deletions(-) diff --git a/rocm/ds4_rocm_glm.cuh b/rocm/ds4_rocm_glm.cuh index e17bdef69..256eab0f8 100644 --- a/rocm/ds4_rocm_glm.cuh +++ b/rocm/ds4_rocm_glm.cuh @@ -2513,13 +2513,15 @@ static uint32_t glm_selected_gemm_head_tile(void) { const unsigned long value = strtoul(env, &end, 10); if (errno == 0 && end && *end == '\0' && (value == 1ul || value == 2ul || - value == 4ul || value == 8ul)) { + value == 4ul || value == 8ul || + value == 16ul || value == 32ul || + value == 64ul)) { tile = (uint32_t)value; return tile; } fprintf(stderr, "ds4: invalid DS4_ROCM_GLM_SELECTED_ATTN_HEAD_TILE='%s'; " - "using 1 (valid values: 1,2,4,8)\n", + "using 1 (valid values: 1,2,4,8,16,32,64)\n", env); return tile; } diff --git a/speed-bench/ROCM_GLM_DISTRIBUTED_VALIDATION.md b/speed-bench/ROCM_GLM_DISTRIBUTED_VALIDATION.md index fd6ea7760..14e4507fd 100644 --- a/speed-bench/ROCM_GLM_DISTRIBUTED_VALIDATION.md +++ b/speed-bench/ROCM_GLM_DISTRIBUTED_VALIDATION.md @@ -135,7 +135,7 @@ attribution and interaction checks: | Wave-parallel value projection | `DS4_ROCM_GLM_VALUE_PROJECT_WAVE_DECODE=0` | On | Give each one-token GLM value-projection output row a 32-lane wave instead of a serial thread. Set `=0` only for a serial-row correctness baseline. | | Causal attention GEMM | `DS4_ROCM_GLM_CAUSAL_ATTN_GEMM=0` | On | Use FP16 BLAS for contiguous causal attention during initial prefill. Set `=0` only for the scalar-kernel fallback. | | Selected attention GEMM | `DS4_ROCM_GLM_SELECTED_ATTN_GEMM=0` | On | Gather per-token selected KV rows and use FP16 strided-batched BLAS after the 2,048-row indexer boundary. Set `=0` only for the scalar-kernel fallback. | -| Selected-attention head tile | `DS4_ROCM_GLM_SELECTED_ATTN_HEAD_TILE=1\|2\|4\|8` | `1` | Process several attention heads as BLAS matrix columns. Tile 1 preserves the production path. | +| Selected-attention head tile | `DS4_ROCM_GLM_SELECTED_ATTN_HEAD_TILE=1\|2\|4\|8\|16\|32\|64` | `1` | Process several attention heads as BLAS matrix columns. Tile 1 preserves the production path. | | Selected-attention phase profile | `DS4_ROCM_GLM_SELECTED_ATTN_PROFILE=1` | Off | Print HIP-event phase totals for the first selected-attention invocation in each process. Diagnostic only. | Apply any opt-in or fallback override to **both** worker and coordinator. For @@ -203,8 +203,8 @@ scalar-kernel baseline and once without the override for the production GEMM path. The first 2,048-token row is the causal control; the later 256-token suffix rows exercise selected attention. The selected GEMM path accepts 64–256 tokens and uses about 580 MiB of temporary workspace at 256 tokens with -2,048 selected rows at head tile 1. Tiles 2, 4, and 8 use approximately 584, -591, and 606 MiB respectively. +2,048 selected rows at head tile 1. Tiles 2, 4, 8, 16, 32, and 64 use +approximately 584, 591, 606, 637, 697, and 818 MiB respectively. The phase profiler intentionally synchronizes at each measured boundary, so never use its wall-clock benchmark row as a performance result. It claims only From 4e80dd7502acde7b99fe588e0d9c506e0a506f5e Mon Sep 17 00:00:00 2001 From: Donato Capitella Date: Sun, 26 Jul 2026 21:26:15 +0100 Subject: [PATCH 30/33] Default GLM selected attention to tile 16 --- rocm/ds4_rocm_glm.cuh | 4 ++-- speed-bench/ROCM_GLM_DISTRIBUTED_VALIDATION.md | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/rocm/ds4_rocm_glm.cuh b/rocm/ds4_rocm_glm.cuh index 256eab0f8..ea99dcc4d 100644 --- a/rocm/ds4_rocm_glm.cuh +++ b/rocm/ds4_rocm_glm.cuh @@ -2503,7 +2503,7 @@ static int g_glm_selected_profile_claimed = 0; static uint32_t glm_selected_gemm_head_tile(void) { static int initialized = 0; - static uint32_t tile = 1u; + static uint32_t tile = 16u; if (initialized) return tile; initialized = 1; const char *env = getenv("DS4_ROCM_GLM_SELECTED_ATTN_HEAD_TILE"); @@ -2521,7 +2521,7 @@ static uint32_t glm_selected_gemm_head_tile(void) { } fprintf(stderr, "ds4: invalid DS4_ROCM_GLM_SELECTED_ATTN_HEAD_TILE='%s'; " - "using 1 (valid values: 1,2,4,8,16,32,64)\n", + "using 16 (valid values: 1,2,4,8,16,32,64)\n", env); return tile; } diff --git a/speed-bench/ROCM_GLM_DISTRIBUTED_VALIDATION.md b/speed-bench/ROCM_GLM_DISTRIBUTED_VALIDATION.md index 14e4507fd..d700de5e2 100644 --- a/speed-bench/ROCM_GLM_DISTRIBUTED_VALIDATION.md +++ b/speed-bench/ROCM_GLM_DISTRIBUTED_VALIDATION.md @@ -135,7 +135,7 @@ attribution and interaction checks: | Wave-parallel value projection | `DS4_ROCM_GLM_VALUE_PROJECT_WAVE_DECODE=0` | On | Give each one-token GLM value-projection output row a 32-lane wave instead of a serial thread. Set `=0` only for a serial-row correctness baseline. | | Causal attention GEMM | `DS4_ROCM_GLM_CAUSAL_ATTN_GEMM=0` | On | Use FP16 BLAS for contiguous causal attention during initial prefill. Set `=0` only for the scalar-kernel fallback. | | Selected attention GEMM | `DS4_ROCM_GLM_SELECTED_ATTN_GEMM=0` | On | Gather per-token selected KV rows and use FP16 strided-batched BLAS after the 2,048-row indexer boundary. Set `=0` only for the scalar-kernel fallback. | -| Selected-attention head tile | `DS4_ROCM_GLM_SELECTED_ATTN_HEAD_TILE=1\|2\|4\|8\|16\|32\|64` | `1` | Process several attention heads as BLAS matrix columns. Tile 1 preserves the production path. | +| Selected-attention head tile | `DS4_ROCM_GLM_SELECTED_ATTN_HEAD_TILE=1\|2\|4\|8\|16\|32\|64` | `16` | Process several attention heads as BLAS matrix columns. Tile 1 is the bit-exact untiled baseline; tiles 32 and 64 remain experimental because they showed accumulated logit drift. | | Selected-attention phase profile | `DS4_ROCM_GLM_SELECTED_ATTN_PROFILE=1` | Off | Print HIP-event phase totals for the first selected-attention invocation in each process. Diagnostic only. | Apply any opt-in or fallback override to **both** worker and coordinator. For From 4ab75720ebd86d6393f8a24cf64c9f110853f408 Mon Sep 17 00:00:00 2001 From: Donato Capitella Date: Sun, 26 Jul 2026 22:25:05 +0100 Subject: [PATCH 31/33] Add GLM indexed decode tuning path --- ds4.c | 21 +++++++++++++++-- ds4_rocm_compat.cu | 23 +++++++++++++++++++ rocm/ds4_rocm_glm.cuh | 16 ++++++++----- rocm/ds4_rocm_runtime.cuh | 10 +++++++- .../ROCM_GLM_DISTRIBUTED_VALIDATION.md | 4 +++- 5 files changed, 64 insertions(+), 10 deletions(-) diff --git a/ds4.c b/ds4.c index b767a7869..52fe4789a 100644 --- a/ds4.c +++ b/ds4.c @@ -37943,13 +37943,30 @@ static uint32_t glm_graph_indexed_decode_split_blocks(void) { } static uint32_t glm_graph_indexed_decode_split_block_rows_for(uint32_t n_selected) { - return n_selected <= 1024u ? 32u : 128u; + const uint32_t default_rows = n_selected <= 1024u ? 32u : 128u; + const char *env = getenv("DS4_ROCM_GLM_INDEXED_DECODE_BLOCK_ROWS"); + if (!env || !env[0]) return default_rows; + char *end = NULL; + errno = 0; + const unsigned long rows = strtoul(env, &end, 10); + if (errno == 0 && end && *end == '\0' && + (rows == 32ul || rows == 64ul || + rows == 128ul || rows == 256ul)) { + return (uint32_t)rows; + } + return default_rows; } static bool glm_graph_indexed_decode_split_group8_available(uint32_t n_selected) { const uint32_t block_rows = glm_graph_indexed_decode_split_block_rows_for(n_selected); const uint32_t needed_blocks = block_rows != 0u ? (n_selected + block_rows - 1u) / block_rows : 0u; + const char *group8_f32_env = + getenv("DS4_ROCM_GLM_INDEXED_DECODE_GROUP8_F32"); + const bool compatible_cache = + glm_graph_compact_cache_is_f16() || + (group8_f32_env && group8_f32_env[0] && + strcmp(group8_f32_env, "0") != 0); return n_selected > 512u && block_rows > 0 && needed_blocks > 0 && @@ -37958,7 +37975,7 @@ static bool glm_graph_indexed_decode_split_group8_available(uint32_t n_selected) (DS4_N_HEAD % 8u) == 0 && DS4_N_KV_LORA == 512u && DS4_N_ROT == 64u && - glm_graph_compact_cache_is_f16(); + compatible_cache; } static bool glm_graph_prefill_stage_sync_boundary(void) { diff --git a/ds4_rocm_compat.cu b/ds4_rocm_compat.cu index 79b587e3b..ac8d485f6 100644 --- a/ds4_rocm_compat.cu +++ b/ds4_rocm_compat.cu @@ -368,6 +368,29 @@ extern "C" int ds4_gpu_glm_attention_indexed_decode_typed_tensor( beta_fast, beta_slow); } +extern "C" int ds4_gpu_glm_attention_indexed_decode_split_group8_typed_tensor( + ds4_gpu_tensor *heads, ds4_gpu_tensor *partial_lora, + ds4_gpu_tensor *partial_ms, const ds4_gpu_tensor *q, + const ds4_gpu_tensor *qk_low, const ds4_gpu_tensor *kv_lora_cache, + const ds4_gpu_tensor *k_rope_cache, const void *model_map, + uint64_t model_size, uint64_t value_weight_offset, + uint32_t value_weight_type, const ds4_gpu_tensor *selected, + uint32_t n_selected, bool selected_rows_valid, uint32_t cache_cap, + bool cache_f16, uint32_t n_head, uint32_t kv_lora_dim, + uint32_t qk_nope, uint32_t qk_rope, uint32_t value_dim, + uint32_t n_ctx_orig, uint32_t block_rows, uint32_t n_blocks, + float freq_base, float freq_scale, float ext_factor, + float attn_factor, float beta_fast, float beta_slow) { + if (value_weight_type != 8u) return 0; + return ds4_gpu_glm_attention_indexed_decode_split_group8_tensor( + heads, partial_lora, partial_ms, q, qk_low, kv_lora_cache, + k_rope_cache, model_map, model_size, value_weight_offset, + selected, n_selected, selected_rows_valid, cache_cap, cache_f16, + n_head, kv_lora_dim, qk_nope, qk_rope, value_dim, n_ctx_orig, + block_rows, n_blocks, freq_base, freq_scale, ext_factor, + attn_factor, beta_fast, beta_slow); +} + extern "C" int ds4_gpu_glm_attention_indexed_batch_typed_tensor( ds4_gpu_tensor *heads, const ds4_gpu_tensor *q, const ds4_gpu_tensor *qk_low, const ds4_gpu_tensor *kv_lora_cache, diff --git a/rocm/ds4_rocm_glm.cuh b/rocm/ds4_rocm_glm.cuh index ea99dcc4d..bfff45d5d 100644 --- a/rocm/ds4_rocm_glm.cuh +++ b/rocm/ds4_rocm_glm.cuh @@ -1029,6 +1029,7 @@ __global__ static void glm_attention_indexed_decode_split_group8_partial_valid_k const float *qk_low, const char *kv_lora_cache, const char *k_rope_cache, + bool cache_f16, const int32_t *selected, uint32_t n_selected, uint32_t n_head, @@ -1098,8 +1099,9 @@ __global__ static void glm_attention_indexed_decode_split_group8_partial_valid_k const uint32_t d = off - rr * kv_lora_dim; const uint32_t row = (uint32_t)selected[base + rr]; kv_shared[off] = - __half2float(((const __half *)kv_lora_cache) - [(uint64_t)row * kv_lora_dim + d]); + glm_rocm_cache_load(kv_lora_cache, + (uint64_t)row * kv_lora_dim + d, + cache_f16); } for (uint32_t off = tid; off < rows * rope_pairs; off += 256u) { const uint32_t rr = off / rope_pairs; @@ -1113,7 +1115,7 @@ __global__ static void glm_attention_indexed_decode_split_group8_partial_valid_k r, row, qk_rope, - true, + cache_f16, n_ctx_orig, freq_base, freq_scale, @@ -3651,6 +3653,8 @@ extern "C" int ds4_gpu_glm_attention_indexed_decode_split_group8_tensor( const unsigned char *value_weight = NULL; uint32_t row_bytes = 0; uint32_t qk_dim = 0; + const uint64_t cache_elem = + cache_f16 ? sizeof(__half) : sizeof(float); const uint32_t needed_blocks = block_rows != 0u ? (n_selected + block_rows - 1u) / block_rows : 0u; if (!glm_rocm_u32_add_checked(qk_nope, qk_rope, &qk_dim) || @@ -3663,7 +3667,6 @@ extern "C" int ds4_gpu_glm_attention_indexed_decode_split_group8_tensor( value_dim == 0u || qk_dim < qk_nope || block_rows == 0u || needed_blocks == 0u || n_blocks < needed_blocks || n_blocks > 64u || - !cache_f16 || !isfinite(freq_base) || freq_base <= 0.0f || !isfinite(freq_scale) || freq_scale <= 0.0f || !isfinite(ext_factor) || !isfinite(attn_factor) || @@ -3674,8 +3677,8 @@ extern "C" int ds4_gpu_glm_attention_indexed_decode_split_group8_tensor( !cuda_tensor_has_elems2(heads, n_head, value_dim, sizeof(float)) || !cuda_tensor_has_elems3(partial_lora, n_blocks, n_head, kv_lora_dim, sizeof(float)) || !cuda_tensor_has_elems3(partial_ms, n_blocks, n_head, 2u, sizeof(float)) || - !glm_rocm_tensor_has_cache2(kv_lora_cache, cache_cap, kv_lora_dim, sizeof(__half)) || - !glm_rocm_tensor_has_cache2(k_rope_cache, cache_cap, qk_rope, sizeof(__half)) || + !glm_rocm_tensor_has_cache2(kv_lora_cache, cache_cap, kv_lora_dim, cache_elem) || + !glm_rocm_tensor_has_cache2(k_rope_cache, cache_cap, qk_rope, cache_elem) || !glm_rocm_check_q8_rows(model_map, model_size, value_weight_offset, (uint64_t)n_head * value_dim, kv_lora_dim, "glm_value_project_split", &value_weight, &row_bytes)) { @@ -3697,6 +3700,7 @@ extern "C" int ds4_gpu_glm_attention_indexed_decode_split_group8_tensor( (const float *)qk_low->ptr, (const char *)kv_lora_cache->ptr, (const char *)k_rope_cache->ptr, + cache_f16, (const int32_t *)selected->ptr, n_selected, n_head, diff --git a/rocm/ds4_rocm_runtime.cuh b/rocm/ds4_rocm_runtime.cuh index 2443bc2b4..c4b8f338d 100644 --- a/rocm/ds4_rocm_runtime.cuh +++ b/rocm/ds4_rocm_runtime.cuh @@ -4784,8 +4784,16 @@ static const ds4_rocm_runtime_config *cuda_runtime_config(void) { g_rocm_cfg.glm_grouped_qk_low = glm_grouped_qk_low_env == NULL || cuda_env_present(glm_grouped_qk_low_env); + const char *q8_decode_sharedx_64k_env = + getenv("DS4_ROCM_Q8_DECODE_SHAREDX_64K"); + /* + * The 64 KiB shared-input path is bit-exact on gfx1151 and falls back + * automatically when a device cannot launch that much dynamic LDS. + * Keep an explicit =0 diagnostic rollback. + */ g_rocm_cfg.q8_decode_sharedx_64k = - cuda_env_present(getenv("DS4_ROCM_Q8_DECODE_SHAREDX_64K")); + q8_decode_sharedx_64k_env == NULL || + cuda_env_present(q8_decode_sharedx_64k_env); const int graph_dump_requested = cuda_env_present(getenv("DS4_ROCM_GRAPH_DUMP_PREFIX")) || cuda_env_present(getenv("DS4_METAL_GRAPH_DUMP_PREFIX")); diff --git a/speed-bench/ROCM_GLM_DISTRIBUTED_VALIDATION.md b/speed-bench/ROCM_GLM_DISTRIBUTED_VALIDATION.md index d700de5e2..6d6536e8c 100644 --- a/speed-bench/ROCM_GLM_DISTRIBUTED_VALIDATION.md +++ b/speed-bench/ROCM_GLM_DISTRIBUTED_VALIDATION.md @@ -131,8 +131,10 @@ attribution and interaction checks: | Path | Environment variable | Default | Intended effect | | --- | --- | --- | --- | -| 64 KiB shared input | `DS4_ROCM_Q8_DECODE_SHAREDX_64K=1` | Off | Let one-token Q8 projections with a 16,384-element input reuse the input through LDS. | +| 64 KiB shared input | `DS4_ROCM_Q8_DECODE_SHAREDX_64K=0` | On | Let one-token Q8 projections with a 16,384-element input reuse the input through LDS. Set `=0` for the warp-row fallback. | | Wave-parallel value projection | `DS4_ROCM_GLM_VALUE_PROJECT_WAVE_DECODE=0` | On | Give each one-token GLM value-projection output row a 32-lane wave instead of a serial thread. Set `=0` only for a serial-row correctness baseline. | +| F32-cache group-8 indexed decode | `DS4_ROCM_GLM_INDEXED_DECODE_GROUP8_F32=1` | Off | Test the split group-of-8 indexed-attention kernel without changing the validated F32 compact cache. | +| Indexed-decode block rows | `DS4_ROCM_GLM_INDEXED_DECODE_BLOCK_ROWS=32\|64\|128\|256` | Auto (`32` through 1,024 selected rows, otherwise `128`) | Tune split-attention rows per partial block. Relevant only when the group-8 path is active. | | Causal attention GEMM | `DS4_ROCM_GLM_CAUSAL_ATTN_GEMM=0` | On | Use FP16 BLAS for contiguous causal attention during initial prefill. Set `=0` only for the scalar-kernel fallback. | | Selected attention GEMM | `DS4_ROCM_GLM_SELECTED_ATTN_GEMM=0` | On | Gather per-token selected KV rows and use FP16 strided-batched BLAS after the 2,048-row indexer boundary. Set `=0` only for the scalar-kernel fallback. | | Selected-attention head tile | `DS4_ROCM_GLM_SELECTED_ATTN_HEAD_TILE=1\|2\|4\|8\|16\|32\|64` | `16` | Process several attention heads as BLAS matrix columns. Tile 1 is the bit-exact untiled baseline; tiles 32 and 64 remain experimental because they showed accumulated logit drift. | From 4f11251b5f700c99258867371ad57eea40820ab4 Mon Sep 17 00:00:00 2001 From: Donato Capitella Date: Sun, 26 Jul 2026 22:53:46 +0100 Subject: [PATCH 32/33] Remove ineffective GLM group8 decode experiment --- ds4.c | 21 ++--------------- ds4_rocm_compat.cu | 23 ------------------- rocm/ds4_rocm_glm.cuh | 16 +++++-------- .../ROCM_GLM_DISTRIBUTED_VALIDATION.md | 2 -- 4 files changed, 8 insertions(+), 54 deletions(-) diff --git a/ds4.c b/ds4.c index 52fe4789a..b767a7869 100644 --- a/ds4.c +++ b/ds4.c @@ -37943,30 +37943,13 @@ static uint32_t glm_graph_indexed_decode_split_blocks(void) { } static uint32_t glm_graph_indexed_decode_split_block_rows_for(uint32_t n_selected) { - const uint32_t default_rows = n_selected <= 1024u ? 32u : 128u; - const char *env = getenv("DS4_ROCM_GLM_INDEXED_DECODE_BLOCK_ROWS"); - if (!env || !env[0]) return default_rows; - char *end = NULL; - errno = 0; - const unsigned long rows = strtoul(env, &end, 10); - if (errno == 0 && end && *end == '\0' && - (rows == 32ul || rows == 64ul || - rows == 128ul || rows == 256ul)) { - return (uint32_t)rows; - } - return default_rows; + return n_selected <= 1024u ? 32u : 128u; } static bool glm_graph_indexed_decode_split_group8_available(uint32_t n_selected) { const uint32_t block_rows = glm_graph_indexed_decode_split_block_rows_for(n_selected); const uint32_t needed_blocks = block_rows != 0u ? (n_selected + block_rows - 1u) / block_rows : 0u; - const char *group8_f32_env = - getenv("DS4_ROCM_GLM_INDEXED_DECODE_GROUP8_F32"); - const bool compatible_cache = - glm_graph_compact_cache_is_f16() || - (group8_f32_env && group8_f32_env[0] && - strcmp(group8_f32_env, "0") != 0); return n_selected > 512u && block_rows > 0 && needed_blocks > 0 && @@ -37975,7 +37958,7 @@ static bool glm_graph_indexed_decode_split_group8_available(uint32_t n_selected) (DS4_N_HEAD % 8u) == 0 && DS4_N_KV_LORA == 512u && DS4_N_ROT == 64u && - compatible_cache; + glm_graph_compact_cache_is_f16(); } static bool glm_graph_prefill_stage_sync_boundary(void) { diff --git a/ds4_rocm_compat.cu b/ds4_rocm_compat.cu index ac8d485f6..79b587e3b 100644 --- a/ds4_rocm_compat.cu +++ b/ds4_rocm_compat.cu @@ -368,29 +368,6 @@ extern "C" int ds4_gpu_glm_attention_indexed_decode_typed_tensor( beta_fast, beta_slow); } -extern "C" int ds4_gpu_glm_attention_indexed_decode_split_group8_typed_tensor( - ds4_gpu_tensor *heads, ds4_gpu_tensor *partial_lora, - ds4_gpu_tensor *partial_ms, const ds4_gpu_tensor *q, - const ds4_gpu_tensor *qk_low, const ds4_gpu_tensor *kv_lora_cache, - const ds4_gpu_tensor *k_rope_cache, const void *model_map, - uint64_t model_size, uint64_t value_weight_offset, - uint32_t value_weight_type, const ds4_gpu_tensor *selected, - uint32_t n_selected, bool selected_rows_valid, uint32_t cache_cap, - bool cache_f16, uint32_t n_head, uint32_t kv_lora_dim, - uint32_t qk_nope, uint32_t qk_rope, uint32_t value_dim, - uint32_t n_ctx_orig, uint32_t block_rows, uint32_t n_blocks, - float freq_base, float freq_scale, float ext_factor, - float attn_factor, float beta_fast, float beta_slow) { - if (value_weight_type != 8u) return 0; - return ds4_gpu_glm_attention_indexed_decode_split_group8_tensor( - heads, partial_lora, partial_ms, q, qk_low, kv_lora_cache, - k_rope_cache, model_map, model_size, value_weight_offset, - selected, n_selected, selected_rows_valid, cache_cap, cache_f16, - n_head, kv_lora_dim, qk_nope, qk_rope, value_dim, n_ctx_orig, - block_rows, n_blocks, freq_base, freq_scale, ext_factor, - attn_factor, beta_fast, beta_slow); -} - extern "C" int ds4_gpu_glm_attention_indexed_batch_typed_tensor( ds4_gpu_tensor *heads, const ds4_gpu_tensor *q, const ds4_gpu_tensor *qk_low, const ds4_gpu_tensor *kv_lora_cache, diff --git a/rocm/ds4_rocm_glm.cuh b/rocm/ds4_rocm_glm.cuh index bfff45d5d..ea99dcc4d 100644 --- a/rocm/ds4_rocm_glm.cuh +++ b/rocm/ds4_rocm_glm.cuh @@ -1029,7 +1029,6 @@ __global__ static void glm_attention_indexed_decode_split_group8_partial_valid_k const float *qk_low, const char *kv_lora_cache, const char *k_rope_cache, - bool cache_f16, const int32_t *selected, uint32_t n_selected, uint32_t n_head, @@ -1099,9 +1098,8 @@ __global__ static void glm_attention_indexed_decode_split_group8_partial_valid_k const uint32_t d = off - rr * kv_lora_dim; const uint32_t row = (uint32_t)selected[base + rr]; kv_shared[off] = - glm_rocm_cache_load(kv_lora_cache, - (uint64_t)row * kv_lora_dim + d, - cache_f16); + __half2float(((const __half *)kv_lora_cache) + [(uint64_t)row * kv_lora_dim + d]); } for (uint32_t off = tid; off < rows * rope_pairs; off += 256u) { const uint32_t rr = off / rope_pairs; @@ -1115,7 +1113,7 @@ __global__ static void glm_attention_indexed_decode_split_group8_partial_valid_k r, row, qk_rope, - cache_f16, + true, n_ctx_orig, freq_base, freq_scale, @@ -3653,8 +3651,6 @@ extern "C" int ds4_gpu_glm_attention_indexed_decode_split_group8_tensor( const unsigned char *value_weight = NULL; uint32_t row_bytes = 0; uint32_t qk_dim = 0; - const uint64_t cache_elem = - cache_f16 ? sizeof(__half) : sizeof(float); const uint32_t needed_blocks = block_rows != 0u ? (n_selected + block_rows - 1u) / block_rows : 0u; if (!glm_rocm_u32_add_checked(qk_nope, qk_rope, &qk_dim) || @@ -3667,6 +3663,7 @@ extern "C" int ds4_gpu_glm_attention_indexed_decode_split_group8_tensor( value_dim == 0u || qk_dim < qk_nope || block_rows == 0u || needed_blocks == 0u || n_blocks < needed_blocks || n_blocks > 64u || + !cache_f16 || !isfinite(freq_base) || freq_base <= 0.0f || !isfinite(freq_scale) || freq_scale <= 0.0f || !isfinite(ext_factor) || !isfinite(attn_factor) || @@ -3677,8 +3674,8 @@ extern "C" int ds4_gpu_glm_attention_indexed_decode_split_group8_tensor( !cuda_tensor_has_elems2(heads, n_head, value_dim, sizeof(float)) || !cuda_tensor_has_elems3(partial_lora, n_blocks, n_head, kv_lora_dim, sizeof(float)) || !cuda_tensor_has_elems3(partial_ms, n_blocks, n_head, 2u, sizeof(float)) || - !glm_rocm_tensor_has_cache2(kv_lora_cache, cache_cap, kv_lora_dim, cache_elem) || - !glm_rocm_tensor_has_cache2(k_rope_cache, cache_cap, qk_rope, cache_elem) || + !glm_rocm_tensor_has_cache2(kv_lora_cache, cache_cap, kv_lora_dim, sizeof(__half)) || + !glm_rocm_tensor_has_cache2(k_rope_cache, cache_cap, qk_rope, sizeof(__half)) || !glm_rocm_check_q8_rows(model_map, model_size, value_weight_offset, (uint64_t)n_head * value_dim, kv_lora_dim, "glm_value_project_split", &value_weight, &row_bytes)) { @@ -3700,7 +3697,6 @@ extern "C" int ds4_gpu_glm_attention_indexed_decode_split_group8_tensor( (const float *)qk_low->ptr, (const char *)kv_lora_cache->ptr, (const char *)k_rope_cache->ptr, - cache_f16, (const int32_t *)selected->ptr, n_selected, n_head, diff --git a/speed-bench/ROCM_GLM_DISTRIBUTED_VALIDATION.md b/speed-bench/ROCM_GLM_DISTRIBUTED_VALIDATION.md index 6d6536e8c..78d1d2eed 100644 --- a/speed-bench/ROCM_GLM_DISTRIBUTED_VALIDATION.md +++ b/speed-bench/ROCM_GLM_DISTRIBUTED_VALIDATION.md @@ -133,8 +133,6 @@ attribution and interaction checks: | --- | --- | --- | --- | | 64 KiB shared input | `DS4_ROCM_Q8_DECODE_SHAREDX_64K=0` | On | Let one-token Q8 projections with a 16,384-element input reuse the input through LDS. Set `=0` for the warp-row fallback. | | Wave-parallel value projection | `DS4_ROCM_GLM_VALUE_PROJECT_WAVE_DECODE=0` | On | Give each one-token GLM value-projection output row a 32-lane wave instead of a serial thread. Set `=0` only for a serial-row correctness baseline. | -| F32-cache group-8 indexed decode | `DS4_ROCM_GLM_INDEXED_DECODE_GROUP8_F32=1` | Off | Test the split group-of-8 indexed-attention kernel without changing the validated F32 compact cache. | -| Indexed-decode block rows | `DS4_ROCM_GLM_INDEXED_DECODE_BLOCK_ROWS=32\|64\|128\|256` | Auto (`32` through 1,024 selected rows, otherwise `128`) | Tune split-attention rows per partial block. Relevant only when the group-8 path is active. | | Causal attention GEMM | `DS4_ROCM_GLM_CAUSAL_ATTN_GEMM=0` | On | Use FP16 BLAS for contiguous causal attention during initial prefill. Set `=0` only for the scalar-kernel fallback. | | Selected attention GEMM | `DS4_ROCM_GLM_SELECTED_ATTN_GEMM=0` | On | Gather per-token selected KV rows and use FP16 strided-batched BLAS after the 2,048-row indexer boundary. Set `=0` only for the scalar-kernel fallback. | | Selected-attention head tile | `DS4_ROCM_GLM_SELECTED_ATTN_HEAD_TILE=1\|2\|4\|8\|16\|32\|64` | `16` | Process several attention heads as BLAS matrix columns. Tile 1 is the bit-exact untiled baseline; tiles 32 and 64 remain experimental because they showed accumulated logit drift. | From 2510636cb0014076ed4519c095ca558757834f50 Mon Sep 17 00:00:00 2001 From: Donato Capitella Date: Mon, 27 Jul 2026 09:14:54 +0100 Subject: [PATCH 33/33] Remove internal GLM validation artifacts --- speed-bench/README.md | 5 - .../ROCM_GLM_DISTRIBUTED_VALIDATION.md | 547 ------------------ speed-bench/compare_glm_validation.py | 172 ------ 3 files changed, 724 deletions(-) delete mode 100644 speed-bench/ROCM_GLM_DISTRIBUTED_VALIDATION.md delete mode 100644 speed-bench/compare_glm_validation.py diff --git a/speed-bench/README.md b/speed-bench/README.md index b9f7aef09..32075fe18 100644 --- a/speed-bench/README.md +++ b/speed-bench/README.md @@ -26,8 +26,3 @@ python3 speed-bench/plot_speed.py speed-bench/m3_max.csv --title "M3 Max t/s" The script uses only the Python standard library. By default it writes a file next to the CSV using the `_ts.svg` suffix, such as `speed-bench/m3_max_ts.svg`. - -For ROCm GLM distributed optimization work, follow -[`ROCM_GLM_DISTRIBUTED_VALIDATION.md`](ROCM_GLM_DISTRIBUTED_VALIDATION.md). -It defines the fixed baseline/candidate protocol, two-node commands, graph -tensor and frontier-logit comparisons, API smoke test, and decision record. diff --git a/speed-bench/ROCM_GLM_DISTRIBUTED_VALIDATION.md b/speed-bench/ROCM_GLM_DISTRIBUTED_VALIDATION.md deleted file mode 100644 index 78d1d2eed..000000000 --- a/speed-bench/ROCM_GLM_DISTRIBUTED_VALIDATION.md +++ /dev/null @@ -1,547 +0,0 @@ -# ROCm GLM Distributed Change Validation - -Use this playbook for every ROCm GLM distributed performance change. It keeps -speed measurements separate from numerical and model-quality checks, and it -changes one variable at a time. - -The commands below describe the current two-Strix-Halo test topology: - -- coordinator: `192.168.100.2`, layers `0:37` -- worker: `192.168.100.1`, layers `38:output` -- model: `GLM-5.2-UD-IQ2_XXS_RoutedIQ2XXS_blk78Q2K.gguf` -- image: `docker.io/kyuz0/strix-halo-ds4-toolbox:glm-rocm-7.2.4` -- distributed prefill: 256-token chunks, window 2 - -Change those constants deliberately if the hardware or model changes, and -record the new values with the results. - -Unless a command is explicitly labeled for the worker, run it on the -coordinator. The coordinator checkout used below is `~/ds4/ds4`. - -## Rules - -1. Build baseline and candidate from the same commit and image. -2. Keep the model file, SHA-256, layer split, prompt, context allocation, - chunk/window, generation length, power state, and background load fixed. -3. Hold previously accepted optimizations constant. Toggle only the change - under test. -4. Save worker and coordinator logs plus benchmark CSV files. -5. Test performance and numerical closeness separately. A coherent sampled - answer is not a numerical comparison, and a matching argmax is not enough. -6. Run the baseline and candidate at least once as a smoke test. Before - declaring a small improvement real or making a path default, repeat the - complete pair. Treat differences below roughly 2% as noise until repeated. -7. Do not compare `ds4-bench --show-output` text with chat answers. The former - is a raw continuation of the prompt file; use the server API for chat. - -Record the code revision and model hash before each series: - -```sh -git -C ~/ds4/ds4 rev-parse HEAD -sha256sum ~/ds4/GLM-5.2-UD-IQ2_XXS_RoutedIQ2XXS_blk78Q2K.gguf -``` - -## Build gate - -The ROCm/gfx1151 compile check does not require a GPU: - -```sh -podman run --rm --userns=keep-id \ - -v ~/ds4/ds4:/workspace:Z \ - -w /workspace \ - docker.io/rocm/dev-ubuntu-24.04:7.2.1-complete \ - make strix-halo ROCM_ARCH=gfx1151 -``` - -Require a successful warning-free build and run `git diff --check`. - -## Publish and deploy the test image - -The image is built by -`kyuz0/strix-halo-ds4-toolbox/.github/workflows/build_and_publish.yml`. -Its `glm-rocm-7.2.4` Dockerfile clones -`https://github.com/kyuz0/ds4.git`, branch -`fix/rocm-distributed-glm`, so push the exact ds4 commit before dispatching the -workflow. The toolbox workflow itself runs from `main`. - -Dispatch only the GLM ROCm image and follow the run to completion: - -```sh -gh workflow run build_and_publish.yml \ - --repo kyuz0/strix-halo-ds4-toolbox \ - --ref main \ - -f backends=glm-rocm-7.2.4 - -gh run list \ - --repo kyuz0/strix-halo-ds4-toolbox \ - --workflow build_and_publish.yml \ - --limit 1 - -gh run watch RUN_ID \ - --repo kyuz0/strix-halo-ds4-toolbox \ - --exit-status -``` - -The workflow publishes both an immutable timestamped tag and the channel tag -used by this playbook: - -```text -docker.io/kyuz0/strix-halo-ds4-toolbox:glm-rocm-7.2.4 -``` - -Pull the channel tag on both hosts before testing. Do not assume a mutable tag -changed merely because `podman pull` succeeded; compare image IDs: - -```sh -ssh fw1 podman pull \ - docker.io/kyuz0/strix-halo-ds4-toolbox:glm-rocm-7.2.4 -ssh fw2 podman pull \ - docker.io/kyuz0/strix-halo-ds4-toolbox:glm-rocm-7.2.4 - -ssh fw1 podman image inspect \ - docker.io/kyuz0/strix-halo-ds4-toolbox:glm-rocm-7.2.4 \ - --format '{{.Id}} {{.Created}}' -ssh fw2 podman image inspect \ - docker.io/kyuz0/strix-halo-ds4-toolbox:glm-rocm-7.2.4 \ - --format '{{.Id}} {{.Created}}' -``` - -The IDs must match. The worker and coordinator examples below invoke the image -directly with Podman; entering a Toolbx shell is not required. Host `/tmp` -results must be mounted explicitly when the command writes inside the -container, for example: - -```sh -mkdir -p /tmp/glm-validation - -podman run --rm \ - --security-opt label=disable \ - -v /tmp/glm-validation:/results \ - IMAGE COMMAND --csv /results/run.csv -``` - -Keep logs outside the container with `2>&1 | tee /tmp/name.log`. Use distinct -container names and result directories for baseline and candidate, and stop -the worker after its coordinator run completes. - -## One-build kernel flag matrix - -The following ROCm paths can be varied independently in one image for -attribution and interaction checks: - -| Path | Environment variable | Default | Intended effect | -| --- | --- | --- | --- | -| 64 KiB shared input | `DS4_ROCM_Q8_DECODE_SHAREDX_64K=0` | On | Let one-token Q8 projections with a 16,384-element input reuse the input through LDS. Set `=0` for the warp-row fallback. | -| Wave-parallel value projection | `DS4_ROCM_GLM_VALUE_PROJECT_WAVE_DECODE=0` | On | Give each one-token GLM value-projection output row a 32-lane wave instead of a serial thread. Set `=0` only for a serial-row correctness baseline. | -| Causal attention GEMM | `DS4_ROCM_GLM_CAUSAL_ATTN_GEMM=0` | On | Use FP16 BLAS for contiguous causal attention during initial prefill. Set `=0` only for the scalar-kernel fallback. | -| Selected attention GEMM | `DS4_ROCM_GLM_SELECTED_ATTN_GEMM=0` | On | Gather per-token selected KV rows and use FP16 strided-batched BLAS after the 2,048-row indexer boundary. Set `=0` only for the scalar-kernel fallback. | -| Selected-attention head tile | `DS4_ROCM_GLM_SELECTED_ATTN_HEAD_TILE=1\|2\|4\|8\|16\|32\|64` | `16` | Process several attention heads as BLAS matrix columns. Tile 1 is the bit-exact untiled baseline; tiles 32 and 64 remain experimental because they showed accumulated logit drift. | -| Selected-attention phase profile | `DS4_ROCM_GLM_SELECTED_ATTN_PROFILE=1` | Off | Print HIP-event phase totals for the first selected-attention invocation in each process. Diagnostic only. | - -Apply any opt-in or fallback override to **both** worker and coordinator. For -Podman, add it as `-e NAME=value`; for a host binary, put `NAME=value \` -before `ds4-bench` or `ds4-server`. - -Use this order so one deployment answers both attribution and interaction: - -1. `DS4_ROCM_GLM_CAUSAL_ATTN_GEMM=0` for the causal-attention scalar fallback; -2. no causal-attention override for the production causal GEMM path; -3. `DS4_ROCM_GLM_VALUE_PROJECT_WAVE_DECODE=0` for the serial-row baseline; -4. no value-projection override for the production wave path; -5. production wave path plus `DS4_ROCM_Q8_DECODE_SHAREDX_64K=1`; -6. `DS4_ROCM_GLM_SELECTED_ATTN_GEMM=0` for the selected-attention scalar - fallback; -7. the production selected-attention GEMM path, with the shared-input - experiment either off or on as required. - -The expected activation messages are: - -```text -Q8 one-token shared-input kernel enabled through 64 KiB LDS -GLM one-token value projection using wave-parallel Q8 rows -GLM causal indexed prefill using fp16 hipBLAS attention GEMMs -GLM selected indexed prefill using fp16 hipBLAS strided-batched attention GEMMs -ROCm GLM selected attention using head tile N -ROCm GLM selected attention profile tokens=... selected=... heads=... tile=... -``` - -Absence of a message means the tested workload did not reach that path. The -64 KiB path also reports and automatically falls back if the launch is not -supported by the device. The wave-parallel value projection changes FP32 -summation order; validate its layer output and frontier logits even when its -generated text looks coherent. - -The normal 2,048-token benchmark below exercises the decode flag, but it does -**not** exercise selected attention: 2,048 visible rows still use the dense -causal path. Measure the prompt-processing cliff with a live 256-token suffix -curve: - -```sh -RUN_NAME=selected-baseline - -DS4_GLM_MEMORY_GUARD_RESERVE_GB=8 \ -DS4_BENCH_DISABLE_SNAPSHOT=1 \ -ds4-bench \ - --prompt-file ~/ds4/ds4/speed-bench/promessi_sposi.txt \ - --ctx-start 2048 \ - --ctx-max 4096 \ - --step-incr 256 \ - --ctx-alloc 4097 \ - --gen-tokens 0 \ - --csv "/tmp/glm-${RUN_NAME}.csv" \ - -m ~/ds4/GLM-5.2-UD-IQ2_XXS_RoutedIQ2XXS_blk78Q2K.gguf \ - --dist-prefill-chunk 256 \ - --dist-prefill-window 2 \ - --role coordinator \ - --layers 0:37 \ - --listen 192.168.100.2 8081 \ - 2>&1 | tee "/tmp/glm-coordinator-${RUN_NAME}.log" -``` - -Run it once with `DS4_ROCM_GLM_SELECTED_ATTN_GEMM=0` on both nodes for the -scalar-kernel baseline and once without the override for the production GEMM -path. The first 2,048-token row is the causal control; the later 256-token -suffix rows exercise selected attention. The selected GEMM path accepts -64–256 tokens and uses about 580 MiB of temporary workspace at 256 tokens with -2,048 selected rows at head tile 1. Tiles 2, 4, 8, 16, 32, and 64 use -approximately 584, 591, 606, 637, 697, and 818 MiB respectively. - -The phase profiler intentionally synchronizes at each measured boundary, so -never use its wall-clock benchmark row as a performance result. It claims only -the first selected-attention invocation per process. Run profiling separately -on worker and coordinator, then disable it for all matched performance and -correctness runs. - -For a layer-local selected-attention comparison, adapt the graph-dump command -below to `--ctx-start 2304 --ctx-max 2304 --ctx-alloc 2433` and set -`DS4_ROCM_GRAPH_DUMP_POS=2048`. For an end-to-end comparison, dump frontier -logits at 2,304 tokens and compare `frontier_002304.logits.json`. A 512-token -dump cannot validate this path. - -## Matched baseline and candidate - -Write down the single environment setting, option, or code path being tested. -Run the commands below once without it for the baseline, then again with only -that setting changed for the candidate. If the setting affects both halves of -the model, apply it to both worker and coordinator. - -The accepted FP16 causal-attention GEMM is enabled by default. Leave it without -an override in ordinary baseline/candidate comparisons. When validating the -causal GEMM itself, use `DS4_ROCM_GLM_CAUSAL_ATTN_GEMM=0` on both nodes for -the scalar-kernel reference and no override for the production candidate. - -### Worker - -Run on `192.168.100.1`. Change `RUN_NAME` to `baseline` or `candidate`, and add -the one candidate environment setting only for the candidate run. - -```sh -RUN_NAME=baseline - -podman run --rm \ - --name "ds4-worker-${RUN_NAME}" \ - -e DS4_GLM_MEMORY_GUARD_RESERVE_GB=8 \ - --device /dev/dri \ - --device /dev/kfd \ - --group-add keep-groups \ - --security-opt seccomp=unconfined \ - --ipc=host \ - --cap-add=SYS_PTRACE \ - --security-opt label=disable \ - --userns=keep-id \ - --network=host \ - -v /mnt/storage/ds4:/models:ro \ - docker.io/kyuz0/strix-halo-ds4-toolbox:glm-rocm-7.2.4 \ - ds4-server \ - -m /models/GLM-5.2-UD-IQ2_XXS_RoutedIQ2XXS_blk78Q2K.gguf \ - --ctx 2177 \ - --host 0.0.0.0 \ - --port 8000 \ - --role worker \ - --layers 38:output \ - --coordinator 192.168.100.2 8081 \ - 2>&1 | tee "/tmp/glm-worker-${RUN_NAME}.log" -``` - -### Coordinator performance run - -Set `RUN_NAME` to match the worker. Add the one candidate environment setting -as another `-e NAME=value` only for the candidate run. - -```sh -RUN_NAME=baseline -mkdir -p "/tmp/glm-${RUN_NAME}" - -podman run --rm \ - --name "ds4-coordinator-${RUN_NAME}" \ - -e DS4_GLM_MEMORY_GUARD_RESERVE_GB=8 \ - -e DS4_BENCH_DISABLE_SNAPSHOT=1 \ - --device /dev/dri \ - --device /dev/kfd \ - --group-add keep-groups \ - --security-opt seccomp=unconfined \ - --ipc=host \ - --cap-add=SYS_PTRACE \ - --security-opt label=disable \ - --userns=keep-id \ - --network=host \ - -v /home/kyuz0/ds4:/models:ro \ - -v "/tmp/glm-${RUN_NAME}:/results" \ - docker.io/kyuz0/strix-halo-ds4-toolbox:glm-rocm-7.2.4 \ - ds4-bench \ - --prompt-file /models/ds4/speed-bench/promessi_sposi.txt \ - --ctx-start 2048 \ - --ctx-max 2048 \ - --ctx-alloc 2177 \ - --gen-tokens 128 \ - --show-output \ - --csv /results/run.csv \ - -m /models/GLM-5.2-UD-IQ2_XXS_RoutedIQ2XXS_blk78Q2K.gguf \ - --dist-prefill-chunk 256 \ - --dist-prefill-window 2 \ - --role coordinator \ - --layers 0:37 \ - --listen 192.168.100.2 8081 \ - 2>&1 | tee "/tmp/glm-coordinator-${RUN_NAME}.log" -``` - -Compare `prefill_tps`, `gen_tps`, and `gen_first_ms`; do not infer a win from -elapsed startup or model-copy time. - -## Layer-local numerical comparison - -This catches wrong layouts, strides, masking, RoPE, scaling, and result -scattering near the operation that changed. Use workers configured like the -performance runs. - -Run the command once as `baseline`, restart the worker with the candidate -configuration, then run it as `candidate`: - -```sh -RUN_NAME=baseline -mkdir -p /tmp/glm-validation - -DS4_GLM_MEMORY_GUARD_RESERVE_GB=8 \ -DS4_ROCM_GRAPH_DUMP_PREFIX="/tmp/glm-validation/${RUN_NAME}" \ -DS4_ROCM_GRAPH_DUMP_NONINVASIVE=1 \ -DS4_ROCM_GRAPH_DUMP_LAYER=0 \ -DS4_ROCM_GRAPH_DUMP_POS=0 \ -DS4_ROCM_GRAPH_DUMP_NAME=glm_indexed_attn_out \ -DS4_BENCH_DISABLE_SNAPSHOT=1 \ -ds4-bench \ - --prompt-file ~/ds4/ds4/speed-bench/promessi_sposi.txt \ - --ctx-start 256 \ - --ctx-max 256 \ - --ctx-alloc 641 \ - --gen-tokens 0 \ - -m ~/ds4/GLM-5.2-UD-IQ2_XXS_RoutedIQ2XXS_blk78Q2K.gguf \ - --dist-prefill-chunk 256 \ - --dist-prefill-window 2 \ - --role coordinator \ - --layers 0:37 \ - --listen 192.168.100.2 8081 -``` - -Compare the dumps: - -```sh -python3 ~/ds4/ds4/speed-bench/compare_glm_validation.py tensor \ - /tmp/glm-validation/baseline_glm_indexed_attn_out-0_pos0.bin \ - /tmp/glm-validation/candidate_glm_indexed_attn_out-0_pos0.bin \ - --hidden 6144 -``` - -Any non-finite value is a failure. For a single early layer, cosine should be -very close to 1. As a diagnostic heuristic, investigate cosine below `0.9999` -or relative RMSE above `0.01`; these are not universal model-quality limits. - -## Frontier-logit comparison - -Layer-local agreement can still accumulate into meaningful final-logit drift. -Collect complete vocabulary logits at a fixed 512-token frontier. Run once as -`baseline`, restart the worker with the candidate configuration, and run again -as `candidate`. - -```sh -RUN_NAME=baseline -mkdir -p "/tmp/glm-logits-${RUN_NAME}" - -DS4_GLM_MEMORY_GUARD_RESERVE_GB=8 \ -DS4_BENCH_DISABLE_SNAPSHOT=1 \ -ds4-bench \ - --prompt-file ~/ds4/ds4/speed-bench/promessi_sposi.txt \ - --ctx-start 512 \ - --ctx-max 512 \ - --ctx-alloc 641 \ - --gen-tokens 0 \ - --dump-frontier-logits-dir "/tmp/glm-logits-${RUN_NAME}" \ - -m ~/ds4/GLM-5.2-UD-IQ2_XXS_RoutedIQ2XXS_blk78Q2K.gguf \ - --dist-prefill-chunk 256 \ - --dist-prefill-window 2 \ - --role coordinator \ - --layers 0:37 \ - --listen 192.168.100.2 8081 -``` - -Compare the dumps: - -```sh -python3 ~/ds4/ds4/speed-bench/compare_glm_validation.py logits \ - /tmp/glm-logits-baseline/frontier_000512.logits.json \ - /tmp/glm-logits-candidate/frontier_000512.logits.json -``` - -Interpret the result as a group: - -- non-finite logits are always a failure; -- centered cosine and relative RMSE measure distribution shape without an - irrelevant constant logit shift; -- KL divergence measures probability-distribution movement; -- top-10/top-50 overlap and the reference top-token rank show whether ordering - changed near the useful head of the distribution; -- matching top-1 is encouraging, but does not by itself prove equivalence; -- a top-1 change is less concerning when the reference margin is tiny. - -There is no universal threshold that turns one frontier into a model-quality -proof. Investigate a material regression from the last accepted path and use -the official GLM continuation scorer before making an approximation the -unconditional release default. - -## Selected-attention promotion record - -On 2026-07-26, the selected-attention GEMM was compared with its scalar-kernel -fallback on the topology at the top of this document. Both runs used commit -`34b3862`, FP16 causal-attention GEMMs, a 2,304-token prompt, 256-token -distributed chunks, and a 2,433-token allocation. The only changed setting was -`DS4_ROCM_GLM_SELECTED_ATTN_GEMM=0` versus `=1`, on both nodes. - -The graph dump at layer 0 and position 2,048 produced: - -```text -all-token relative RMSE: 5.9441e-05 -all-token cosine: 0.999999998234 -final-token cosine: 0.999999997577 -non-finite pairs: 0 -``` - -The complete 154,880-value frontier-logit comparison produced: - -```text -KL(reference || candidate): 0.0664153 -top-1: 3956 -> 3956 -top-10 overlap: 10/10 -top-50 overlap: 46/50 -centered cosine: 0.9921320 -``` - -The dump-free performance pair improved from `19.11` to `27.48` prompt -tokens/s (`+43.8%`). The instrumented pair measured `18.85` versus `27.49` -tokens/s; do not compare either instrumented row with a normal benchmark -because serializing the full logit vector is included in the final chunk. - -Decision: make the selected-attention GEMM the ROCm GLM default and preserve -`DS4_ROCM_GLM_SELECTED_ATTN_GEMM=0` as the diagnostic rollback. The narrower -candidate top-1 margin (`1.2262` to `0.3320`) remains important context even -though the selected token and top-10 set matched. - -The default-on change was published from ds4 commit `ea2e766` by toolbox -GitHub Actions run -[`30214076099`](https://github.com/kyuz0/strix-halo-ds4-toolbox/actions/runs/30214076099). -After pulling the channel tag, both hosts reported: - -```text -image ID: 379e27ad0439970bf72e04ac63fe2711e3471468f2171bad757ea03bca52d72c -digest: sha256:221bc3b031f5600094e40a1b0202fbd7f7057f31ab4854836bf97a5420dbc36f -created: 2026-07-26 18:13:38 UTC -``` - -A no-override post-deployment benchmark activated both the selected-attention -GEMM and wave-parallel one-token value projection and measured: - -```text -ctx=2304 prefill=27.45 t/s generation=2.30 t/s first-token=434.913 ms -``` - -The long API smoke used a 2,482-token chat prompt and 32 generated tokens. It -activated selected attention, completed prefill in `102.769 s`, decoded at -`2.29 t/s`, and returned a coherent description of the supplied passage as the -introduction to Manzoni's *I promessi sposi*. - -## Causal-attention default deployment record - -The causal-attention GEMM became default-on in ds4 commit `dd837f0`, published -by toolbox GitHub Actions run -[`30215053603`](https://github.com/kyuz0/strix-halo-ds4-toolbox/actions/runs/30215053603). -Both hosts pulled and verified: - -```text -image ID: 6f424ab78c427808d7d7cee3b4a0ea01b232223a0ca525ebfb033042dec4c4a0 -digest: sha256:9bda0b1d4c45c08bc06caffa90262a964d935534353595b199eab98ce9e6c5cb -created: 2026-07-26 18:41:10 UTC -``` - -The post-deployment benchmark used no kernel environment overrides, layers -`0:37` plus `38:output`, a 2,433-token allocation, 128 generated tokens, and -the production distributed prefill settings: - -```text ---dist-prefill-chunk 256 ---dist-prefill-window 2 -``` - -Chunk 256 is intentional. Earlier matched 2,048-token runs measured -`50.81 t/s` with chunk 256 versus `48.15 t/s` with chunk 512. - -The production curve measured: - -```text -ctx prefill tokens prefill t/s generation t/s first token -2048 2048 50.55 2.30 439.299 ms -2304 256 5.22 2.30 434.724 ms -``` - -The second row is the incremental selected-attention suffix at an existing -2,048-token frontier, not full-prompt throughput. It exposes the remaining -post-2,048 prefill cliff. A separate cold 2,304-token full-prompt run measured: - -```text -prefill=27.46 t/s generation=2.30 t/s first-token=439.220 ms -``` - -All three default activation messages were present: causal attention GEMM, -selected attention GEMM, and wave-parallel one-token value projection. Both -128-token raw continuations were coherent. Preserve -`DS4_ROCM_GLM_CAUSAL_ATTN_GEMM=0` as the independent scalar-kernel rollback. - -## API smoke test - -After tensor and logit checks pass, run a server and test the actual chat -template. This catches routing, KV replay, and API integration failures: - -```sh -wget --timeout=900 --tries=1 -qO- \ - --header='Content-Type: application/json' \ - --post-data='{"model":"glm-5.2","messages":[{"role":"user","content":"Reply with exactly: hello"}],"thinking":{"type":"disabled"},"temperature":0,"max_tokens":16}' \ - http://localhost:8000/v1/chat/completions | -jq -r '.choices[0].message.content // .' -``` - -Expect a coherent response, normally `hello`. This is a smoke test, not a -replacement for tensor/logit comparison or the 100-case GLM continuation gate -documented in `QA_BEFORE_RELEASES.md`. - -## Decision record - -For every candidate, retain: - -- commit and model SHA-256; -- exact baseline and candidate environment variables; -- coordinator and worker logs; -- baseline and candidate CSV rows; -- layer-local tensor comparison; -- frontier-logit comparison; -- API smoke response; -- decision: reject, keep diagnostic, retain opt-in, or make default. - -If a candidate is faster but its drift is unexplained, keep it diagnostic or -reject it. Explain and validate correctness before promoting performance code. diff --git a/speed-bench/compare_glm_validation.py b/speed-bench/compare_glm_validation.py deleted file mode 100644 index 48e1211be..000000000 --- a/speed-bench/compare_glm_validation.py +++ /dev/null @@ -1,172 +0,0 @@ -#!/usr/bin/env python3 -"""Compare ROCm GLM validation dumps without third-party dependencies.""" - -import argparse -import array -import heapq -import json -import math -from pathlib import Path - - -def vector_metrics(reference, candidate): - if len(reference) != len(candidate): - raise SystemExit( - f"length mismatch: reference={len(reference)} " - f"candidate={len(candidate)}" - ) - if not reference: - raise SystemExit("input vectors are empty") - - nonfinite = sum( - not (math.isfinite(a) and math.isfinite(b)) - for a, b in zip(reference, candidate) - ) - diff2 = sum((a - b) ** 2 for a, b in zip(reference, candidate)) - ref2 = sum(a * a for a in reference) - cand2 = sum(b * b for b in candidate) - dot = sum(a * b for a, b in zip(reference, candidate)) - - return { - "values": len(reference), - "nonfinite_pairs": nonfinite, - "max_abs": max(abs(a - b) for a, b in zip(reference, candidate)), - "rmse": math.sqrt(diff2 / len(reference)), - "rel_rmse": math.sqrt(diff2 / ref2) if ref2 else math.inf, - "cosine": dot / math.sqrt(ref2 * cand2) - if ref2 and cand2 - else math.nan, - } - - -def print_metrics(metrics, indent=" "): - print(f"{indent}values: {metrics['values']}") - print(f"{indent}nonfinite_pairs: {metrics['nonfinite_pairs']}") - print(f"{indent}max_abs: {metrics['max_abs']:.12g}") - print(f"{indent}rmse: {metrics['rmse']:.12g}") - print(f"{indent}rel_rmse: {metrics['rel_rmse']:.12g}") - print(f"{indent}cosine: {metrics['cosine']:.12g}") - - -def load_f32(path): - values = array.array("f") - with path.open("rb") as fp: - values.frombytes(fp.read()) - return values - - -def compare_tensor(args): - reference = load_f32(args.reference) - candidate = load_f32(args.candidate) - if args.hidden <= 0: - raise SystemExit("--hidden must be positive") - if len(reference) % args.hidden: - raise SystemExit( - f"reference contains {len(reference)} values, which is not " - f"divisible by hidden size {args.hidden}" - ) - - print(f"tokens: {len(reference) // args.hidden}") - print("all tokens:") - print_metrics(vector_metrics(reference, candidate)) - print("final token:") - print_metrics( - vector_metrics(reference[-args.hidden :], candidate[-args.hidden :]) - ) - - -def load_logits(path): - with path.open(encoding="utf-8") as fp: - obj = json.load(fp) - logits = obj.get("logits") - if not isinstance(logits, list): - raise SystemExit(f"{path}: missing logits array") - return logits - - -def logsumexp(values): - maximum = max(values) - return maximum + math.log(sum(math.exp(value - maximum) for value in values)) - - -def compare_logits(args): - reference = load_logits(args.reference) - candidate = load_logits(args.candidate) - metrics = vector_metrics(reference, candidate) - count = len(reference) - - ref_mean = sum(reference) / count - cand_mean = sum(candidate) / count - ref_centered = [value - ref_mean for value in reference] - cand_centered = [value - cand_mean for value in candidate] - centered = vector_metrics(ref_centered, cand_centered) - - ref_top = heapq.nlargest(50, range(count), key=reference.__getitem__) - cand_top = heapq.nlargest(50, range(count), key=candidate.__getitem__) - ref_top1 = ref_top[0] - cand_top1 = cand_top[0] - ref_margin = reference[ref_top[0]] - reference[ref_top[1]] - cand_margin = candidate[cand_top[0]] - candidate[cand_top[1]] - ref_top1_rank = 1 + sum( - value > candidate[ref_top1] for value in candidate - ) - - ref_lse = logsumexp(reference) - cand_lse = logsumexp(candidate) - kl_divergence = sum( - math.exp(a - ref_lse) - * ((a - ref_lse) - (b - cand_lse)) - for a, b in zip(reference, candidate) - ) - - print("raw logits:") - print_metrics(metrics) - print(f" mean_shift: {cand_mean - ref_mean:.12g}") - print("centered logits:") - print_metrics(centered) - print(f"KL(reference || candidate): {kl_divergence:.12g}") - print(f"top1 IDs: {ref_top1} -> {cand_top1}") - print(f"reference top1 rank: {ref_top1_rank}") - print(f"top1 margins: {ref_margin:.12g} -> {cand_margin:.12g}") - print( - "top10 overlap: " - f"{len(set(ref_top[:10]) & set(cand_top[:10]))}/10" - ) - print( - "top50 overlap: " - f"{len(set(ref_top) & set(cand_top))}/50" - ) - - -def main(): - parser = argparse.ArgumentParser( - description="Compare ds4 ROCm GLM tensor or frontier-logit dumps." - ) - subparsers = parser.add_subparsers(dest="mode", required=True) - - tensor = subparsers.add_parser( - "tensor", help="compare two raw F32 graph tensor dumps" - ) - tensor.add_argument("reference", type=Path) - tensor.add_argument("candidate", type=Path) - tensor.add_argument( - "--hidden", - type=int, - default=6144, - help="hidden width used to report the final token (default: 6144)", - ) - tensor.set_defaults(func=compare_tensor) - - logits = subparsers.add_parser( - "logits", help="compare two ds4-bench frontier-logit JSON files" - ) - logits.add_argument("reference", type=Path) - logits.add_argument("candidate", type=Path) - logits.set_defaults(func=compare_logits) - - args = parser.parse_args() - args.func(args) - - -if __name__ == "__main__": - main()