| 25238 | model, |
| 25239 | weights->output, |
| 25240 | DS4_N_EMBD, |
| 25241 | vocab_dim, |
| 25242 | metal_graph_output_norm(g), |
| 25243 | 1); |
| 25244 | } |
| 25245 | if (ok) { |
| 25246 | metal_graph_debug_dump_tensor("result_output", metal_graph_logits(g), vocab_dim, DS4_N_LAYER, 0); |
| 25247 | } |
| 25248 | #undef DS4_METAL_PROFILE_OUTPUT_STAGE |
| 25249 | return ok; |
| 25250 | } |
| 25251 | |
| 25252 | /* Greedy-only output head: compute one local top-1 candidate per output TP |
| 25253 | * split and leave the full split logits on their owning tiers. This avoids |
| 25254 | * gathering the whole vocabulary row back to the head tier when the caller only |
| 25255 | * needs the next argmax token. */ |
| 25256 | static bool metal_graph_encode_output_head_split_top1( |
| 25257 | ds4_gpu_graph *g, |
| 25258 | const ds4_model *model, |
| 25259 | const ds4_weights *weights, |
| 25260 | uint64_t vocab_dim, |
| 25261 | int cuda_tp_output_tiers[DS4_MAX_GPUS], |
| 25262 | uint32_t *cuda_tp_output_ways_out) { |
| 25263 | if (!g || !model || !weights || |
| 25264 | !cuda_tp_output_tiers || !cuda_tp_output_ways_out || |
| 25265 | vocab_dim > UINT32_MAX) { |
| 25266 | return false; |
| 25267 | } |
| 25268 | *cuda_tp_output_ways_out = 0; |
| 25269 | |
| 25270 | if (g->placement) { |
| 25271 | if (!metal_graph_set_active_tier_decode(g, g->head_tier)) return false; |
| 25272 | } |
| 25273 | const uint64_t hc_dim = (uint64_t)DS4_N_HC * DS4_N_EMBD; |
| 25274 | bool ok = ds4_gpu_rms_norm_plain_tensor(metal_graph_flat_hc(g), |
| 25275 | metal_graph_cur_hc(g), |
| 25276 | (uint32_t)hc_dim, |
| 25277 | DS4_RMS_EPS) != 0; |
| 25278 | if (ok) ok = ds4_gpu_matmul_f16_tensor(metal_graph_output_pre(g), |
| 25279 | model->map, |
| 25280 | model->size, |
| 25281 | weights->output_hc_fn->abs_offset, |
| 25282 | hc_dim, |
| 25283 | DS4_N_HC, |
| 25284 | metal_graph_flat_hc(g), |
| 25285 | 1) != 0; |
| 25286 | if (ok) ok = ds4_gpu_output_hc_weights_tensor(metal_graph_output_weights(g), |
| 25287 | metal_graph_output_pre(g), |
| 25288 | model->map, |
| 25289 | model->size, |
| 25290 | weights->output_hc_scale->abs_offset, |
| 25291 | weights->output_hc_base->abs_offset, |
| 25292 | DS4_N_HC, |
| 25293 | DS4_HC_EPS) != 0; |
| 25294 | if (ok) ok = ds4_gpu_hc_weighted_sum_tensor(metal_graph_output_embd(g), |
| 25295 | metal_graph_cur_hc(g), |
| 25296 | metal_graph_output_weights(g), |
| 25297 | DS4_N_EMBD, |
no test coverage detected