| 25544 | weights->output->type == DS4_TENSOR_Q8_0 && |
| 25545 | weights->output->ndim == 2 && |
| 25546 | weights->output->dim[0] == DS4_N_EMBD && |
| 25547 | weights->output->dim[1] == vocab_dim && |
| 25548 | head_rows <= DS4_DSPARK_MAX_BLOCK_SIZE && |
| 25549 | getenv("DS4_DSPARK_VERIFY_HEAD_NO_TP") == NULL; |
| 25550 | for (uint32_t i = 0; tp_ok && i < tp_ways; i++) { |
| 25551 | const int t = tp_tiers[i]; |
| 25552 | tp_ok = t >= 0 && t < DS4_MAX_GPUS && |
| 25553 | g->logits_by_tier[t] && |
| 25554 | ds4_gpu_tensor_bytes(g->logits_by_tier[t]) >= |
| 25555 | (uint64_t)head_rows * |
| 25556 | ((vocab_dim + tp_ways - 1u) / tp_ways) * |
| 25557 | sizeof(float) && |
| 25558 | (t == g->active_tier || |
| 25559 | (g->batch_ffn_norm_by_tier[t] && |
| 25560 | ds4_gpu_tensor_bytes(g->batch_ffn_norm_by_tier[t]) >= |
| 25561 | (uint64_t)head_rows * DS4_N_EMBD * sizeof(float))); |
| 25562 | } |
| 25563 | if (tp_ok) { |
| 25564 | const uint64_t row_bytes = metal_graph_q8_0_row_bytes(DS4_N_EMBD); |
| 25565 | const int home_tier = g->active_tier; |
| 25566 | uint64_t split_start[DS4_MAX_GPUS] = {0}; |
| 25567 | uint64_t split_count[DS4_MAX_GPUS] = {0}; |
| 25568 | for (uint32_t i = 0; ok && i < tp_ways; i++) { |
| 25569 | const int t = tp_tiers[i]; |
| 25570 | split_start[i] = (vocab_dim * (uint64_t)i) / tp_ways; |
| 25571 | const uint64_t split_end = |
| 25572 | (vocab_dim * (uint64_t)(i + 1u)) / tp_ways; |
| 25573 | split_count[i] = split_end - split_start[i]; |
| 25574 | if (split_count[i] == 0) { ok = false; break; } |
| 25575 | if (t != home_tier) { |
| 25576 | ok = ds4_gpu_tensor_copy_xdev( |
| 25577 | g->batch_ffn_norm_by_tier[t], |
| 25578 | output_norm, |
| 25579 | (uint64_t)head_rows * DS4_N_EMBD * |
| 25580 | sizeof(float)) != 0; |
| 25581 | } |
| 25582 | } |
| 25583 | for (uint32_t i = 0; ok && i < tp_ways; i++) { |
| 25584 | const int t = tp_tiers[i]; |
| 25585 | ok = ds4_gpu_set_current_device(t) == 0; |
| 25586 | if (!ok) break; |
| 25587 | ds4_gpu_tensor *shard_out = |
| 25588 | ds4_gpu_tensor_view(g->logits_by_tier[t], |
| 25589 | 0, |
| 25590 | (uint64_t)head_rows * split_count[i] * |
| 25591 | sizeof(float)); |
| 25592 | ds4_gpu_tensor *shard_in = t == home_tier ? |
| 25593 | NULL : |
| 25594 | ds4_gpu_tensor_view(g->batch_ffn_norm_by_tier[t], |
| 25595 | 0, |
| 25596 | (uint64_t)head_rows * DS4_N_EMBD * |
| 25597 | sizeof(float)); |
| 25598 | ok = shard_out && (t == home_tier || shard_in) && |
| 25599 | ds4_gpu_matmul_q8_0_tensor(shard_out, |
| 25600 | model->map, |
| 25601 | model->size, |
| 25602 | weights->output->abs_offset + |
| 25603 | split_start[i] * row_bytes, |