diff --git a/ds4.c b/ds4.c index 876cfe27f..b767a7869 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, @@ -37498,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, @@ -39007,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); @@ -41563,6 +41720,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) { @@ -41706,6 +41870,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, @@ -41880,8 +42066,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, @@ -41906,6 +42106,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; } @@ -53272,12 +53479,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 +53549,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 +53679,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 +53716,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 +53739,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 +53754,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 +53767,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 +53789,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 +53850,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 +53861,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 +53879,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 +53889,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 +53920,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 +53945,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 +53966,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 +54042,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 +55538,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 +55558,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 +55606,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 @@ -55357,15 +55674,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 " @@ -55386,7 +55715,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, @@ -55674,7 +56006,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 +56081,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 +56091,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 +56125,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 +56148,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"); @@ -57238,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; @@ -57248,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, diff --git a/ds4_rocm.cu b/ds4_rocm.cu index 640e00209..46ba223df 100644 --- a/ds4_rocm.cu +++ b/ds4_rocm.cu @@ -34,6 +34,7 @@ #include #include #include +#include #include #include @@ -52,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_glm.cuh b/rocm/ds4_rocm_glm.cuh index 2cd689d88..ea99dcc4d 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, @@ -1891,6 +1931,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, @@ -1931,6 +1994,65 @@ 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"); + } + 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 && + (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, + 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, @@ -1946,6 +2068,1113 @@ 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]; +} + +__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); + } +} + +__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, + 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; +} + +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 = 16u; + 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 || + 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 16 (valid values: 1,2,4,8,16,32,64)\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, + 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_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; + } + + 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; + 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; + } + 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; + 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); + + 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, + (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; + } + 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, + (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; + } + 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)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)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, + (int)tile_heads, + (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; + } + 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, + (int)tile_heads, + (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_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, + (int)tile_heads, + (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_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; +} + static int glm_attention_indexed_lora_launch( ds4_gpu_tensor *lora_out, const ds4_gpu_tensor *q, @@ -2000,6 +3229,95 @@ 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 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 && + (selected_attn_gemm_env == NULL || + cuda_env_present(selected_attn_gemm_env)) && + 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; + } + /* 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 && + (causal_attn_gemm_env == NULL || + cuda_env_present(causal_attn_gemm_env)) && + 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, 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; diff --git a/rocm/ds4_rocm_matmul.cuh b/rocm/ds4_rocm_matmul.cuh index 8a62dcd16..8c05433b9 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, @@ -249,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<<< @@ -262,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_moe.cuh b/rocm/ds4_rocm_moe.cuh index 74fbbbfed..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++) { @@ -2144,6 +2150,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..51fafe3d3 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, @@ -737,8 +742,13 @@ 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); + /* Correctness rollback for the optimized resident IQ2 prefill path. */ + const uint32_t disable_resident_iq2_sorted = + 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) && + !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; @@ -848,17 +858,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 +984,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 +1025,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 +1374,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 +1418,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 +1602,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"); 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 07152280c..c4b8f338d 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; @@ -4732,6 +4752,9 @@ 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 glm_grouped_qk_low; + int q8_decode_sharedx_64k; int graph_dump; uint32_t moe_decode_rpb; uint32_t moe_decode_gate_rpb; @@ -4747,9 +4770,41 @@ 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 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 = + 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 = + 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 = + 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")); + /* 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'; @@ -5481,13 +5536,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 +5668,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 +5681,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 +5692,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 +5706,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 +5728,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 +5742,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 +6191,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; }