nicolas
/
ik_llama_opt
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity
ik_llama_opt/src/llama-dflash-profile.h

341 lines
13 KiB
C++
Raw Blame History

#pragma once
#include <cstdint>
#include <cstring>
enum llama_dflash_kv_node_kind {
LLAMA_DFLASH_KV_NODE_NONE = 0,
LLAMA_DFLASH_KV_NODE_FUSED_TARGET,
LLAMA_DFLASH_KV_NODE_K_PROJ,
LLAMA_DFLASH_KV_NODE_K_NORM,
LLAMA_DFLASH_KV_NODE_K_ROPE,
LLAMA_DFLASH_KV_NODE_V_PROJ,
LLAMA_DFLASH_KV_NODE_K_STORE,
LLAMA_DFLASH_KV_NODE_V_STORE,
};
enum llama_dflash_main_node_kind {
LLAMA_DFLASH_MAIN_NODE_NONE = 0,
LLAMA_DFLASH_MAIN_NODE_QCUR,
LLAMA_DFLASH_MAIN_NODE_K_DRAFT,
LLAMA_DFLASH_MAIN_NODE_V_DRAFT,
LLAMA_DFLASH_MAIN_NODE_K_CTX_VIEW,
LLAMA_DFLASH_MAIN_NODE_V_CTX_VIEW,
LLAMA_DFLASH_MAIN_NODE_K_CONCAT,
LLAMA_DFLASH_MAIN_NODE_V_CONCAT,
LLAMA_DFLASH_MAIN_NODE_K_PAD,
LLAMA_DFLASH_MAIN_NODE_V_PAD,
LLAMA_DFLASH_MAIN_NODE_K_PERM_CONT,
LLAMA_DFLASH_MAIN_NODE_V_PERM_CONT,
LLAMA_DFLASH_MAIN_NODE_FLASH_ATTN,
LLAMA_DFLASH_MAIN_NODE_ATTN_OUT,
LLAMA_DFLASH_MAIN_NODE_FFN,
LLAMA_DFLASH_MAIN_NODE_RESULT_ROWS,
LLAMA_DFLASH_MAIN_NODE_RESULT_NORM,
LLAMA_DFLASH_MAIN_NODE_RESULT,
};
struct llama_dflash_kv_node_profiler {
llama_dflash_profile_stats * profile = nullptr;
int64_t t_start_us = 0;
llama_dflash_kv_node_kind active_kind = LLAMA_DFLASH_KV_NODE_NONE;
};
struct llama_dflash_main_node_profiler {
llama_dflash_profile_stats * profile = nullptr;
ggml_backend_sched_eval_callback prev_callback = nullptr;
void * prev_user_data = nullptr;
bool prev_active = false;
int64_t t_start_us = 0;
llama_dflash_main_node_kind active_kind = LLAMA_DFLASH_MAIN_NODE_NONE;
};
static inline bool llama_dflash_tensor_name_has_prefix(const struct ggml_tensor * tensor, const char * prefix) {
if (tensor == nullptr || prefix == nullptr || prefix[0] == '\0') {
return false;
}
return std::strncmp(tensor->name, prefix, std::strlen(prefix)) == 0;
}
static inline bool llama_dflash_tensor_name_matches_label(const struct ggml_tensor * tensor, const char * label) {
if (!llama_dflash_tensor_name_has_prefix(tensor, label)) {
return false;
}
const size_t label_len = std::strlen(label);
const char next = tensor->name[label_len];
return next == '\0' || next == '-';
}
static inline llama_dflash_kv_node_kind llama_dflash_kv_node_kind_from_tensor(const struct ggml_tensor * tensor) {
if (llama_dflash_tensor_name_has_prefix(tensor, "dflash_kv_fused_target")) {
return LLAMA_DFLASH_KV_NODE_FUSED_TARGET;
}
if (llama_dflash_tensor_name_has_prefix(tensor, "dflash_kv_k_proj")) {
return LLAMA_DFLASH_KV_NODE_K_PROJ;
}
if (llama_dflash_tensor_name_has_prefix(tensor, "dflash_kv_k_norm")) {
return LLAMA_DFLASH_KV_NODE_K_NORM;
}
if (llama_dflash_tensor_name_has_prefix(tensor, "dflash_kv_k_rope")) {
return LLAMA_DFLASH_KV_NODE_K_ROPE;
}
if (llama_dflash_tensor_name_has_prefix(tensor, "dflash_kv_v_proj")) {
return LLAMA_DFLASH_KV_NODE_V_PROJ;
}
if (llama_dflash_tensor_name_has_prefix(tensor, "dflash_kv_k_store")) {
return LLAMA_DFLASH_KV_NODE_K_STORE;
}
if (llama_dflash_tensor_name_has_prefix(tensor, "dflash_kv_v_store")) {
return LLAMA_DFLASH_KV_NODE_V_STORE;
}
return LLAMA_DFLASH_KV_NODE_NONE;
}
static inline void llama_dflash_kv_node_profile_add(
llama_dflash_profile_stats & profile,
llama_dflash_kv_node_kind kind,
uint64_t elapsed_us) {
switch (kind) {
case LLAMA_DFLASH_KV_NODE_FUSED_TARGET:
profile.graph_kv_node_fused_target_calls++;
profile.graph_kv_node_fused_target_us += elapsed_us;
break;
case LLAMA_DFLASH_KV_NODE_K_PROJ:
profile.graph_kv_node_k_proj_calls++;
profile.graph_kv_node_k_proj_us += elapsed_us;
break;
case LLAMA_DFLASH_KV_NODE_K_NORM:
profile.graph_kv_node_k_norm_calls++;
profile.graph_kv_node_k_norm_us += elapsed_us;
break;
case LLAMA_DFLASH_KV_NODE_K_ROPE:
profile.graph_kv_node_k_rope_calls++;
profile.graph_kv_node_k_rope_us += elapsed_us;
break;
case LLAMA_DFLASH_KV_NODE_V_PROJ:
profile.graph_kv_node_v_proj_calls++;
profile.graph_kv_node_v_proj_us += elapsed_us;
break;
case LLAMA_DFLASH_KV_NODE_K_STORE:
profile.graph_kv_node_k_store_calls++;
profile.graph_kv_node_k_store_us += elapsed_us;
break;
case LLAMA_DFLASH_KV_NODE_V_STORE:
profile.graph_kv_node_v_store_calls++;
profile.graph_kv_node_v_store_us += elapsed_us;
break;
case LLAMA_DFLASH_KV_NODE_NONE:
break;
}
}
static inline llama_dflash_main_node_kind llama_dflash_main_node_kind_from_tensor(const struct ggml_tensor * tensor) {
if (llama_dflash_tensor_name_has_prefix(tensor, "Qcur")) {
return LLAMA_DFLASH_MAIN_NODE_QCUR;
}
if (llama_dflash_tensor_name_has_prefix(tensor, "Kcur_noise")) {
return LLAMA_DFLASH_MAIN_NODE_K_DRAFT;
}
if (llama_dflash_tensor_name_has_prefix(tensor, "Vcur_noise")) {
return LLAMA_DFLASH_MAIN_NODE_V_DRAFT;
}
if (llama_dflash_tensor_name_has_prefix(tensor, "Kcur_ctx_cache")) {
return LLAMA_DFLASH_MAIN_NODE_K_CTX_VIEW;
}
if (llama_dflash_tensor_name_has_prefix(tensor, "Vcur_ctx_cache")) {
return LLAMA_DFLASH_MAIN_NODE_V_CTX_VIEW;
}
if (llama_dflash_tensor_name_has_prefix(tensor, "dflash_main_k_concat")) {
return LLAMA_DFLASH_MAIN_NODE_K_CONCAT;
}
if (llama_dflash_tensor_name_has_prefix(tensor, "dflash_main_v_concat")) {
return LLAMA_DFLASH_MAIN_NODE_V_CONCAT;
}
if (llama_dflash_tensor_name_has_prefix(tensor, "dflash_main_k_pad")) {
return LLAMA_DFLASH_MAIN_NODE_K_PAD;
}
if (llama_dflash_tensor_name_has_prefix(tensor, "dflash_main_v_pad")) {
return LLAMA_DFLASH_MAIN_NODE_V_PAD;
}
if (llama_dflash_tensor_name_has_prefix(tensor, "dflash_main_k_perm_cont")) {
return LLAMA_DFLASH_MAIN_NODE_K_PERM_CONT;
}
if (llama_dflash_tensor_name_has_prefix(tensor, "dflash_main_v_perm_cont")) {
return LLAMA_DFLASH_MAIN_NODE_V_PERM_CONT;
}
if (llama_dflash_tensor_name_has_prefix(tensor, "flash_attn_reshaped")) {
return LLAMA_DFLASH_MAIN_NODE_NONE;
}
if (llama_dflash_tensor_name_matches_label(tensor, "flash_attn")) {
return LLAMA_DFLASH_MAIN_NODE_FLASH_ATTN;
}
if (llama_dflash_tensor_name_has_prefix(tensor, "kqv_out")) {
return LLAMA_DFLASH_MAIN_NODE_ATTN_OUT;
}
if (llama_dflash_tensor_name_has_prefix(tensor, "ffn_out")) {
return LLAMA_DFLASH_MAIN_NODE_FFN;
}
if (llama_dflash_tensor_name_matches_label(tensor, "result_output_rows")) {
return LLAMA_DFLASH_MAIN_NODE_RESULT_ROWS;
}
if (llama_dflash_tensor_name_matches_label(tensor, "result_norm")) {
return LLAMA_DFLASH_MAIN_NODE_RESULT_NORM;
}
if (llama_dflash_tensor_name_matches_label(tensor, "output")) {
return LLAMA_DFLASH_MAIN_NODE_RESULT;
}
if (llama_dflash_tensor_name_matches_label(tensor, "result_output")) {
return LLAMA_DFLASH_MAIN_NODE_RESULT;
}
return LLAMA_DFLASH_MAIN_NODE_NONE;
}
static inline void llama_dflash_main_node_profile_add(
llama_dflash_profile_stats & profile,
llama_dflash_main_node_kind kind,
uint64_t elapsed_us) {
switch (kind) {
case LLAMA_DFLASH_MAIN_NODE_QCUR:
profile.graph_main_node_qcur_calls++;
profile.graph_main_node_qcur_us += elapsed_us;
break;
case LLAMA_DFLASH_MAIN_NODE_K_DRAFT:
profile.graph_main_node_k_draft_calls++;
profile.graph_main_node_k_draft_us += elapsed_us;
break;
case LLAMA_DFLASH_MAIN_NODE_V_DRAFT:
profile.graph_main_node_v_draft_calls++;
profile.graph_main_node_v_draft_us += elapsed_us;
break;
case LLAMA_DFLASH_MAIN_NODE_K_CTX_VIEW:
profile.graph_main_node_k_ctx_view_calls++;
profile.graph_main_node_k_ctx_view_us += elapsed_us;
break;
case LLAMA_DFLASH_MAIN_NODE_V_CTX_VIEW:
profile.graph_main_node_v_ctx_view_calls++;
profile.graph_main_node_v_ctx_view_us += elapsed_us;
break;
case LLAMA_DFLASH_MAIN_NODE_K_CONCAT:
profile.graph_main_node_k_concat_calls++;
profile.graph_main_node_k_concat_us += elapsed_us;
break;
case LLAMA_DFLASH_MAIN_NODE_V_CONCAT:
profile.graph_main_node_v_concat_calls++;
profile.graph_main_node_v_concat_us += elapsed_us;
break;
case LLAMA_DFLASH_MAIN_NODE_K_PAD:
profile.graph_main_node_k_pad_calls++;
profile.graph_main_node_k_pad_us += elapsed_us;
break;
case LLAMA_DFLASH_MAIN_NODE_V_PAD:
profile.graph_main_node_v_pad_calls++;
profile.graph_main_node_v_pad_us += elapsed_us;
break;
case LLAMA_DFLASH_MAIN_NODE_K_PERM_CONT:
profile.graph_main_node_k_perm_cont_calls++;
profile.graph_main_node_k_perm_cont_us += elapsed_us;
break;
case LLAMA_DFLASH_MAIN_NODE_V_PERM_CONT:
profile.graph_main_node_v_perm_cont_calls++;
profile.graph_main_node_v_perm_cont_us += elapsed_us;
break;
case LLAMA_DFLASH_MAIN_NODE_FLASH_ATTN:
profile.graph_main_node_flash_attn_calls++;
profile.graph_main_node_flash_attn_us += elapsed_us;
break;
case LLAMA_DFLASH_MAIN_NODE_ATTN_OUT:
profile.graph_main_node_attn_out_calls++;
profile.graph_main_node_attn_out_us += elapsed_us;
break;
case LLAMA_DFLASH_MAIN_NODE_FFN:
profile.graph_main_node_ffn_calls++;
profile.graph_main_node_ffn_us += elapsed_us;
break;
case LLAMA_DFLASH_MAIN_NODE_RESULT_ROWS:
profile.graph_main_node_result_rows_calls++;
profile.graph_main_node_result_rows_us += elapsed_us;
break;
case LLAMA_DFLASH_MAIN_NODE_RESULT_NORM:
profile.graph_main_node_result_norm_calls++;
profile.graph_main_node_result_norm_us += elapsed_us;
break;
case LLAMA_DFLASH_MAIN_NODE_RESULT:
profile.graph_main_node_result_calls++;
profile.graph_main_node_result_us += elapsed_us;
break;
case LLAMA_DFLASH_MAIN_NODE_NONE:
break;
}
}
static inline bool llama_dflash_kv_node_eval_callback(struct ggml_tensor * tensor, bool ask, void * user_data) {
auto * profiler = static_cast<llama_dflash_kv_node_profiler *>(user_data);
if (profiler == nullptr || profiler->profile == nullptr) {
return false;
}
const llama_dflash_kv_node_kind kind = llama_dflash_kv_node_kind_from_tensor(tensor);
if (ask) {
if (kind == LLAMA_DFLASH_KV_NODE_NONE) {
return false;
}
profiler->active_kind = kind;
profiler->t_start_us = ggml_time_us();
return true;
}
if (kind != LLAMA_DFLASH_KV_NODE_NONE && profiler->active_kind == kind && profiler->t_start_us > 0) {
llama_dflash_kv_node_profile_add(*profiler->profile, kind, (uint64_t) (ggml_time_us() - profiler->t_start_us));
}
profiler->active_kind = LLAMA_DFLASH_KV_NODE_NONE;
profiler->t_start_us = 0;
return true;
}
static inline bool llama_dflash_main_node_eval_callback(struct ggml_tensor * tensor, bool ask, void * user_data) {
auto * profiler = static_cast<llama_dflash_main_node_profiler *>(user_data);
if (profiler == nullptr || profiler->profile == nullptr) {
return false;
}
const llama_dflash_main_node_kind kind = llama_dflash_main_node_kind_from_tensor(tensor);
if (ask) {
profiler->prev_active = profiler->prev_callback != nullptr
? profiler->prev_callback(tensor, ask, profiler->prev_user_data)
: false;
if (kind == LLAMA_DFLASH_MAIN_NODE_NONE) {
profiler->active_kind = LLAMA_DFLASH_MAIN_NODE_NONE;
profiler->t_start_us = 0;
return profiler->prev_active;
}
profiler->active_kind = kind;
profiler->t_start_us = ggml_time_us();
return true;
}
bool prev_result = false;
if (profiler->prev_active && profiler->prev_callback != nullptr) {
prev_result = profiler->prev_callback(tensor, ask, profiler->prev_user_data);
}
const bool tracked = kind != LLAMA_DFLASH_MAIN_NODE_NONE &&
profiler->active_kind == kind &&
profiler->t_start_us > 0;
if (tracked) {
llama_dflash_main_node_profile_add(*profiler->profile, kind, (uint64_t) (ggml_time_us() - profiler->t_start_us));
}
profiler->prev_active = false;
profiler->active_kind = LLAMA_DFLASH_MAIN_NODE_NONE;
profiler->t_start_us = 0;
return prev_result || tracked;
}
Reference in New Issue View Git Blame Copy Permalink