diff --git a/ggml/CMakeLists.txt b/ggml/CMakeLists.txt index 5381c213..b5d9a3f7 100644 --- a/ggml/CMakeLists.txt +++ b/ggml/CMakeLists.txt @@ -208,6 +208,9 @@ option(GGML_CUDA_FA "ggml: compile ggml FlashAttention C option(GGML_CUDA_FA_ALL_QUANTS "ggml: compile all quants for FlashAttention" OFF) option(GGML_CUDA_GRAPHS "ggml: use CUDA graphs (llama.cpp only)" ${GGML_CUDA_GRAPHS_DEFAULT}) option(GGML_CUDA_NCCL "ggml: use NVIDIA Collective Comm. Library" ON) +option(GGML_CUDA_PALW_TRACE "ggml: enable experimental request-local PALW MMVQ tracing" OFF) +set (PALW_CUDA_RUNTIME_DIR "" CACHE PATH + "ggml: runtime-palw/cuda source directory for PALW tracing") set (GGML_CUDA_COMPRESSION_MODE "size" CACHE STRING "ggml: cuda link binary compression mode; requires cuda 12.8+") set_property(CACHE GGML_CUDA_COMPRESSION_MODE PROPERTY STRINGS "none;speed;balance;size") diff --git a/ggml/src/ggml-cuda/CMakeLists.txt b/ggml/src/ggml-cuda/CMakeLists.txt index d3953eee..3a991741 100644 --- a/ggml/src/ggml-cuda/CMakeLists.txt +++ b/ggml/src/ggml-cuda/CMakeLists.txt @@ -129,6 +129,21 @@ if (CUDAToolkit_FOUND) ${GGML_SOURCES_CUDA} ) + if (GGML_CUDA_PALW_TRACE) + if (NOT IS_DIRECTORY "${PALW_CUDA_RUNTIME_DIR}" OR + NOT EXISTS "${PALW_CUDA_RUNTIME_DIR}/palw_cuda_llama_mmvq.h" OR + NOT EXISTS "${PALW_CUDA_RUNTIME_DIR}/producer_accumulator.cu") + message(FATAL_ERROR + "GGML_CUDA_PALW_TRACE requires PALW_CUDA_RUNTIME_DIR to point to runtime-palw/cuda") + endif() + target_sources(ggml-cuda PRIVATE + "${PALW_CUDA_RUNTIME_DIR}/producer_accumulator.cu" + "${PALW_CUDA_RUNTIME_DIR}/producer_contract.cpp" + "${PALW_CUDA_RUNTIME_DIR}/record_codec.cpp") + target_include_directories(ggml-cuda PUBLIC "${PALW_CUDA_RUNTIME_DIR}") + target_compile_definitions(ggml-cuda PUBLIC GGML_CUDA_PALW_TRACE=1) + endif() + add_compile_definitions(GGML_CUDA_PEER_MAX_BATCH_SIZE=${GGML_CUDA_PEER_MAX_BATCH_SIZE}) if (GGML_CUDA_GRAPHS) diff --git a/ggml/src/ggml-cuda/common.cuh b/ggml/src/ggml-cuda/common.cuh index 290dc4af..e2733885 100644 --- a/ggml/src/ggml-cuda/common.cuh +++ b/ggml/src/ggml-cuda/common.cuh @@ -40,6 +40,10 @@ #include "vendors/cuda.h" #endif // defined(GGML_USE_HIP) +#if defined(GGML_CUDA_PALW_TRACE) +#include "palw_cuda_llama_mmvq.h" +#endif + #define STRINGIZE_IMPL(...) #__VA_ARGS__ #define STRINGIZE(...) STRINGIZE_IMPL(__VA_ARGS__) @@ -1388,6 +1392,48 @@ struct ggml_cuda_stream_context { } }; +#if defined(GGML_CUDA_PALW_TRACE) +struct ggml_cuda_palw_trace_state { + const palw_cuda_llama_mmvq_request_v1 * request = nullptr; + const palw_cuda_llama_trace_request_v2 * request_v2 = nullptr; + palw_cuda_llama_mmvq_association_v1 association = {}; + palw_cuda_llama_trace_association_v2 association_v2 = {}; + palw_cuda_producer_launch_v3 expected_launch = {}; + palw_cuda_trace_status status = PALW_CUDA_TRACE_OK; + palw_cuda_llama_mmvq_fault_v1 fault = PALW_CUDA_LLAMA_MMVQ_FAULT_NONE; + uint64_t accepted_launch_count = 0; + bool association_active = false; + bool association_launched = false; + bool failed = false; + + bool active() const { + return (request != nullptr) != (request_v2 != nullptr); + } + + bool mixed_stream() const { + return request_v2 != nullptr; + } + + uint64_t expected_launch_count() const { + return request_v2 != nullptr + ? request_v2->expected_launch_count + : (request != nullptr ? request->expected_launch_count : 0); + } + + palw_cuda_producer_trace_context_v3 * producer_context() const { + return request_v2 != nullptr + ? request_v2->producer_context + : (request != nullptr ? request->producer_context : nullptr); + } + + const void * operation_token() const { + return request_v2 != nullptr + ? association_v2.operation_token + : association.operation_token; + } +}; +#endif + struct ggml_backend_cuda_context { int device; std::string name; @@ -1398,6 +1444,10 @@ struct ggml_backend_cuda_context { int curr_stream_no = 0; +#if defined(GGML_CUDA_PALW_TRACE) + ggml_cuda_palw_trace_state palw_trace; +#endif + #ifdef USE_CUDA_GRAPH // Map from first_node_ptr to cuda_graph - allows multiple graphs per context // when the computation is split across CPU/GPU (e.g., with --n-cpu-moe) @@ -1501,6 +1551,264 @@ struct ggml_backend_cuda_context { } }; +#if defined(GGML_CUDA_PALW_TRACE) +static inline bool ggml_cuda_palw_trace_active(const ggml_backend_cuda_context * ctx) { + return ctx != nullptr && ctx->palw_trace.active(); +} + +static inline bool ggml_cuda_palw_trace_failed(const ggml_backend_cuda_context * ctx) { + return ggml_cuda_palw_trace_active(ctx) && ctx->palw_trace.failed; +} + +static inline void ggml_cuda_palw_trace_fail( + ggml_backend_cuda_context * ctx, + palw_cuda_trace_status status, + palw_cuda_llama_mmvq_fault_v1 fault) { + if (!ggml_cuda_palw_trace_active(ctx) || ctx->palw_trace.failed) { + return; + } + ctx->palw_trace.failed = true; + ctx->palw_trace.status = status; + ctx->palw_trace.fault = fault; + if (ctx->palw_trace.request_v2 != nullptr) { + const auto * request = ctx->palw_trace.request_v2; + if (request->fault != nullptr) { + request->fault( + request->user_data, + status, + fault, + ctx->palw_trace.association_active ? &ctx->palw_trace.association_v2 : nullptr); + } + } else { + const auto * request = ctx->palw_trace.request; + if (request != nullptr && request->fault != nullptr) { + request->fault( + request->user_data, + status, + fault, + ctx->palw_trace.association_active ? &ctx->palw_trace.association : nullptr); + } + } +} + +static inline void ggml_cuda_palw_trace_launch_accepted( + ggml_backend_cuda_context * ctx) { + if (!ggml_cuda_palw_trace_active(ctx)) { + return; + } + auto & trace = ctx->palw_trace; + trace.association_launched = true; + ++trace.accepted_launch_count; + if (trace.request_v2 != nullptr) { + if (trace.request_v2->launch_accepted != nullptr) { + trace.request_v2->launch_accepted( + trace.request_v2->user_data, + &trace.association_v2, + &trace.expected_launch); + } + } else if (trace.request != nullptr && trace.request->launch_accepted != nullptr) { + trace.request->launch_accepted( + trace.request->user_data, + &trace.association, + &trace.expected_launch); + } +} + +struct ggml_cuda_palw_attention_work_launch_v1 { + const void * entry_point = nullptr; + palw_cuda_producer_launch_dimensions_v1 dimensions = {}; + uint8_t work_tensor_type = PALW_CUDA_LLAMA_TENSOR_TYPE_NONE; + uint8_t mask_tensor_type = PALW_CUDA_LLAMA_TENSOR_TYPE_NONE; + uint64_t output_row_stride = 0; +}; + +static inline uint8_t ggml_cuda_palw_tensor_type(const ggml_tensor * tensor) { + if (tensor == nullptr) { + return PALW_CUDA_LLAMA_TENSOR_TYPE_NONE; + } + switch (tensor->type) { + case GGML_TYPE_F32: return PALW_CUDA_LLAMA_TENSOR_TYPE_F32; + case GGML_TYPE_F16: return PALW_CUDA_LLAMA_TENSOR_TYPE_F16; + case GGML_TYPE_Q4_K: return PALW_CUDA_LLAMA_TENSOR_TYPE_Q4_K; + case GGML_TYPE_Q6_K: return PALW_CUDA_LLAMA_TENSOR_TYPE_Q6_K; + default: return UINT8_MAX; + } +} + +static inline bool ggml_cuda_palw_tensor_view_matches( + const palw_cuda_llama_tensor_view_v1 & expected, + const ggml_tensor * actual) { + if (actual == nullptr) { + return expected.type == PALW_CUDA_LLAMA_TENSOR_TYPE_NONE && + expected.data == nullptr; + } + if (expected.data != actual->data || + expected.type != ggml_cuda_palw_tensor_type(actual)) { + return false; + } + for (uint32_t i = 0; i < 4; ++i) { + if (actual->ne[i] <= 0 || expected.ne[i] != (uint64_t) actual->ne[i] || + expected.nb[i] != (uint64_t) actual->nb[i]) { + return false; + } + } + return true; +} + +static inline bool ggml_cuda_palw_trace_v2_attention_active( + const ggml_backend_cuda_context * ctx) { + return ggml_cuda_palw_trace_active(ctx) && + ctx->palw_trace.request_v2 != nullptr && + ctx->palw_trace.association_active && + ctx->palw_trace.association_v2.attention_stage != + PALW_CUDA_TRACE_ATTENTION_STAGE_NONE; +} + +static inline bool ggml_cuda_palw_trace_attention_preflight( + ggml_backend_cuda_context * ctx, + const ggml_tensor * src0, + const ggml_tensor * src1, + const ggml_tensor * src2, + const ggml_tensor * dst, + uint8_t attention_stage, + uint8_t precision, + uint8_t attention_mask) { + if (!ggml_cuda_palw_trace_v2_attention_active(ctx)) { + return false; + } + auto & trace = ctx->palw_trace; + const auto & association = trace.association_v2; + if (trace.failed || trace.association_launched || + association.operation_token != dst || + association.attention_stage != attention_stage || + association.precision != precision || + association.attention_mask != attention_mask || + !ggml_cuda_palw_tensor_view_matches(association.src0, src0) || + !ggml_cuda_palw_tensor_view_matches(association.src1, src1) || + !ggml_cuda_palw_tensor_view_matches(association.src2, src2) || + !ggml_cuda_palw_tensor_view_matches(association.dst, dst) || + ctx->curr_stream_no != 0 || + !ctx->stream_context().concurrent_events.empty() || + ggml_cuda_info().devices[ctx->device].cc != GGML_CUDA_CC_ADA_LOVELACE) { + ggml_cuda_palw_trace_fail( + ctx, + PALW_CUDA_TRACE_IDENTITY_MISMATCH, + PALW_CUDA_LLAMA_MMVQ_FAULT_ASSOCIATION_MISMATCH); + return false; + } + return true; +} + +static inline palw_cuda_producer_actual_identity_v3 +ggml_cuda_palw_attention_actual_identity_v3(uint8_t attention_stage) { + palw_cuda_producer_actual_identity_v3 identity = {}; + static constexpr uint8_t qk_variant[PALW_CUDA_PRODUCER_ID_SIZE] = { + 0x0a, 0x29, 0x91, 0x34, 0x08, 0x5c, 0x7b, 0x67, + 0xbf, 0xee, 0x93, 0x51, 0x42, 0x7f, 0xa4, 0x5a, + 0x40, 0x9c, 0x68, 0x4b, 0x65, 0x33, 0xa7, 0x0d, + 0x9f, 0x41, 0xea, 0xcd, 0xcf, 0x80, 0x5b, 0xb4, + }; + static constexpr uint8_t qk_work[PALW_CUDA_PRODUCER_ID_SIZE] = { + 0xf1, 0x11, 0x28, 0x82, 0xf6, 0xd4, 0x5f, 0xe2, + 0x97, 0x5c, 0x00, 0x12, 0x26, 0x92, 0x29, 0xbb, + 0xb3, 0xf5, 0xe2, 0x1c, 0xde, 0xdf, 0xd0, 0xc0, + 0x85, 0xcc, 0xf3, 0xf5, 0x36, 0x94, 0xb1, 0x83, + }; + static constexpr uint8_t softmax_variant[PALW_CUDA_PRODUCER_ID_SIZE] = { + 0xe7, 0x4e, 0x46, 0x18, 0x5a, 0x11, 0x64, 0x79, + 0x24, 0x46, 0x40, 0x43, 0x86, 0xc2, 0xcf, 0x2e, + 0xd3, 0xae, 0x7f, 0x6c, 0x67, 0x26, 0xd7, 0x2e, + 0x26, 0xb1, 0x00, 0x9d, 0x7e, 0x89, 0xc0, 0xcc, + }; + static constexpr uint8_t softmax_work[PALW_CUDA_PRODUCER_ID_SIZE] = { + 0xeb, 0x58, 0x35, 0x7f, 0x1b, 0x03, 0x97, 0x04, + 0xe6, 0x8c, 0x9c, 0x71, 0xd9, 0xfd, 0xa4, 0x9f, + 0x35, 0x6b, 0x69, 0xbc, 0x23, 0x4d, 0x18, 0xa6, + 0x49, 0x14, 0xf5, 0x97, 0xca, 0xf4, 0xcf, 0xfb, + }; + static constexpr uint8_t pv_variant[PALW_CUDA_PRODUCER_ID_SIZE] = { + 0xe0, 0x28, 0xf4, 0x56, 0x07, 0x14, 0xce, 0x35, + 0xde, 0x86, 0xdc, 0x93, 0x46, 0xf7, 0xa2, 0x33, + 0x22, 0x95, 0x1b, 0xff, 0xaa, 0x6d, 0x96, 0xd0, + 0xf3, 0x8c, 0xcb, 0xe6, 0x76, 0x1f, 0x6a, 0x61, + }; + static constexpr uint8_t pv_work[PALW_CUDA_PRODUCER_ID_SIZE] = { + 0xe9, 0xd7, 0xc2, 0x4f, 0xaf, 0x71, 0x67, 0xa4, + 0xac, 0xe7, 0xec, 0x61, 0xbf, 0x72, 0x1d, 0x8e, + 0xb0, 0x2e, 0x1f, 0x36, 0x8a, 0x02, 0x93, 0xc7, + 0x98, 0xfc, 0xf1, 0x25, 0x96, 0xc1, 0x2d, 0x72, + }; + const uint8_t * variant = nullptr; + const uint8_t * work = nullptr; + if (attention_stage == PALW_CUDA_TRACE_ATTENTION_STAGE_EAGER_QK_SCORE_MMVF) { + variant = qk_variant; + work = qk_work; + } else if (attention_stage == + PALW_CUDA_TRACE_ATTENTION_STAGE_EAGER_MASKED_SCALED_SOFTMAX) { + variant = softmax_variant; + work = softmax_work; + } else if (attention_stage == + PALW_CUDA_TRACE_ATTENTION_STAGE_EAGER_VALUE_AGGREGATION_MMVF) { + variant = pv_variant; + work = pv_work; + } + if (variant != nullptr) { + memcpy(identity.producer_variant_id, variant, PALW_CUDA_PRODUCER_ID_SIZE); + memcpy(identity.work_entry_point_id, work, PALW_CUDA_PRODUCER_ID_SIZE); + const uint8_t capture[PALW_CUDA_PRODUCER_ID_SIZE] = { + PALW_CUDA_PRODUCER_GROUPED_CAPTURE_ID_BYTES_V1, + }; + memcpy(identity.capture_implementation_id, capture, PALW_CUDA_PRODUCER_ID_SIZE); + } + return identity; +} + +static inline void ggml_cuda_palw_trace_enqueue_attention( + ggml_backend_cuda_context * ctx, + const ggml_tensor * dst, + const ggml_cuda_palw_attention_work_launch_v1 & work) { + if (!ggml_cuda_palw_trace_v2_attention_active(ctx) || + ctx->palw_trace.failed) { + return; + } + auto & trace = ctx->palw_trace; + const auto & association = trace.association_v2; + const auto identity = ggml_cuda_palw_attention_actual_identity_v3( + association.attention_stage); + if (work.entry_point == nullptr || dst == nullptr || + work.output_row_stride == 0 || + work.output_row_stride != association.dst.nb[1] / sizeof(float) || + work.work_tensor_type != association.src0.type || + (association.attention_stage == + PALW_CUDA_TRACE_ATTENTION_STAGE_EAGER_MASKED_SCALED_SOFTMAX && + work.mask_tensor_type != association.src1.type)) { + ggml_cuda_palw_trace_fail( + ctx, + PALW_CUDA_TRACE_IDENTITY_MISMATCH, + PALW_CUDA_LLAMA_MMVQ_FAULT_ASSOCIATION_MISMATCH); + return; + } + const palw_cuda_trace_status status = + palw_cuda_producer_trace_enqueue_grouped_final_output_f32_v3( + trace.producer_context(), + ctx->stream(), + &trace.expected_launch, + work.entry_point, + &work.dimensions, + static_cast(dst->data), + work.output_row_stride, + &identity); + if (status != PALW_CUDA_TRACE_OK) { + ggml_cuda_palw_trace_fail( + ctx, + status, + PALW_CUDA_LLAMA_MMVQ_FAULT_LAUNCH); + return; + } + ggml_cuda_palw_trace_launch_accepted(ctx); +} +#endif + struct ggml_cuda_mm_fusion_args_host { const ggml_tensor * x_bias = nullptr; const ggml_tensor * gate = nullptr; @@ -1642,4 +1950,3 @@ static __inline__ void ggml_cuda_kernel_launch(Kernel kernel, const ggml_cuda_ke kernel<<>>(std::forward(args)... ); CUDA_CHECK(cudaGetLastError()); } - diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu index 0878ab9c..c7bf90d9 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu @@ -3038,6 +3038,12 @@ static bool ggml_cuda_can_fuse(const struct ggml_cgraph * cgraph, // try and fuse nodes and return the number of nodes to skip static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph * cgraph, int i) { +#if defined(GGML_CUDA_PALW_TRACE) + if (ggml_cuda_palw_trace_active(cuda_ctx)) { + return 0; + } +#endif + static bool disable_fusion = getenv("GGML_CUDA_DISABLE_FUSION") != nullptr && std::atoi(getenv("GGML_CUDA_DISABLE_FUSION")); if (disable_fusion) { return 0; @@ -3926,6 +3932,12 @@ static void ggml_cuda_graph_evaluate_and_capture(ggml_backend_cuda_context * cud } GGML_ASSERT(ok); +#if defined(GGML_CUDA_PALW_TRACE) + if (ggml_cuda_palw_trace_failed(cuda_ctx)) { + return; + } +#endif + if (!is_concurrent_event_active) { try_launch_concurrent_event(node); } @@ -3933,7 +3945,10 @@ static void ggml_cuda_graph_evaluate_and_capture(ggml_backend_cuda_context * cud } #ifdef USE_CUDA_GRAPH - ggml_cuda_graph * graph = cuda_ctx->cuda_graph(graph_key); + // A request-local PALW trace always executes directly. In that mode + // use_cuda_graph is false and even looking up the null graph key would + // mutate the backend's graph map, violating the no-graph bind contract. + ggml_cuda_graph * graph = use_cuda_graph ? cuda_ctx->cuda_graph(graph_key) : nullptr; if (use_cuda_graph && cuda_graph_update_required) { // End CUDA graph capture if (graph->graph != nullptr) { CUDA_CHECK(cudaGraphDestroy(graph->graph)); @@ -3991,39 +4006,56 @@ static enum ggml_status ggml_backend_cuda_graph_compute(ggml_backend_t backend, ggml_cuda_set_device(cuda_ctx->device); +#if defined(GGML_CUDA_PALW_TRACE) + const bool palw_trace_active = ggml_cuda_palw_trace_active(cuda_ctx); + if (palw_trace_active && + (cuda_ctx->curr_stream_no != 0 || + !cuda_ctx->stream_context().concurrent_events.empty())) { + ggml_cuda_palw_trace_fail( + cuda_ctx, + PALW_CUDA_TRACE_INVALID_ARGUMENT, + PALW_CUDA_LLAMA_MMVQ_FAULT_CONCURRENCY); + return GGML_STATUS_FAILED; + } +#else + constexpr bool palw_trace_active = false; +#endif + bool use_cuda_graph = false; bool cuda_graph_update_required = false; const void * graph_key = nullptr; #ifdef USE_CUDA_GRAPH - graph_key = ggml_cuda_graph_get_key(cgraph); + if (!palw_trace_active) { + graph_key = ggml_cuda_graph_get_key(cgraph); - ggml_cuda_graph_set_enabled(cuda_ctx, graph_key); + ggml_cuda_graph_set_enabled(cuda_ctx, graph_key); - ggml_cuda_graph * graph = cuda_ctx->cuda_graph(graph_key); - if (graph->is_enabled()) { - const bool graph_compatible = ggml_cuda_graph_check_compability(cgraph); - if (graph_compatible) { - const bool properties_changed = ggml_cuda_graph_update_required(cuda_ctx, cgraph); - - if (!graph->warmup_complete) { - // Warmup: need at least 2 calls with no property change on the 2nd call - if (!properties_changed) { - graph->warmup_complete = true; - GGML_LOG_DEBUG("%s: CUDA graph warmup complete\n", __func__); - use_cuda_graph = true; - cuda_graph_update_required = true; - } - // else: properties changed or first call - execute directly (use_cuda_graph stays false) - } else { - // Post-warmup: normal CUDA graph operation - if (properties_changed) { - // Properties changed - reset warmup, execute directly until stable again - graph->warmup_complete = false; - GGML_LOG_DEBUG("%s: CUDA graph warmup reset\n", __func__); + ggml_cuda_graph * graph = cuda_ctx->cuda_graph(graph_key); + if (graph->is_enabled()) { + const bool graph_compatible = ggml_cuda_graph_check_compability(cgraph); + if (graph_compatible) { + const bool properties_changed = ggml_cuda_graph_update_required(cuda_ctx, cgraph); + + if (!graph->warmup_complete) { + // Warmup: need at least 2 calls with no property change on the 2nd call + if (!properties_changed) { + graph->warmup_complete = true; + GGML_LOG_DEBUG("%s: CUDA graph warmup complete\n", __func__); + use_cuda_graph = true; + cuda_graph_update_required = true; + } + // else: properties changed or first call - execute directly (use_cuda_graph stays false) } else { - use_cuda_graph = true; - cuda_graph_update_required = graph->instance == nullptr; + // Post-warmup: normal CUDA graph operation + if (properties_changed) { + // Properties changed - reset warmup, execute directly until stable again + graph->warmup_complete = false; + GGML_LOG_DEBUG("%s: CUDA graph warmup reset\n", __func__); + } else { + use_cuda_graph = true; + cuda_graph_update_required = graph->instance == nullptr; + } } } } @@ -4042,6 +4074,12 @@ static enum ggml_status ggml_backend_cuda_graph_compute(ggml_backend_t backend, ggml_cuda_graph_evaluate_and_capture(cuda_ctx, cgraph, use_cuda_graph, cuda_graph_update_required, graph_key); +#if defined(GGML_CUDA_PALW_TRACE) + if (ggml_cuda_palw_trace_failed(cuda_ctx)) { + return GGML_STATUS_FAILED; + } +#endif + return GGML_STATUS_SUCCESS; } @@ -4073,6 +4111,19 @@ static void ggml_backend_cuda_event_wait(ggml_backend_t backend, ggml_backend_ev static void ggml_backend_cuda_graph_optimize(ggml_backend_t backend, ggml_cgraph * cgraph) { ggml_backend_cuda_context * cuda_ctx = (ggml_backend_cuda_context *) backend->context; +#if defined(GGML_CUDA_PALW_TRACE) + if (ggml_cuda_palw_trace_active(cuda_ctx)) { + if (cuda_ctx->curr_stream_no != 0 || + !cuda_ctx->stream_context().concurrent_events.empty()) { + ggml_cuda_palw_trace_fail( + cuda_ctx, + PALW_CUDA_TRACE_INVALID_ARGUMENT, + PALW_CUDA_LLAMA_MMVQ_FAULT_CONCURRENCY); + } + return; + } +#endif + #ifdef USE_CUDA_GRAPH const void * graph_key = ggml_cuda_graph_get_key(cgraph); const bool use_cuda_graph = ggml_cuda_graph_set_enabled(cuda_ctx, graph_key); @@ -5137,6 +5188,498 @@ static ggml_backend_feature * ggml_backend_cuda_get_features(ggml_backend_reg_t GGML_UNUSED(reg); } +#if defined(GGML_CUDA_PALW_TRACE) +static void ggml_backend_cuda_palw_report_unbound_fault( + const palw_cuda_llama_mmvq_request_v1 * request, + palw_cuda_trace_status status, + palw_cuda_llama_mmvq_fault_v1 fault) { + if (request != nullptr && request->fault != nullptr) { + request->fault(request->user_data, status, fault, nullptr); + } +} + +static void ggml_backend_cuda_palw_report_unbound_fault_v2( + const palw_cuda_llama_trace_request_v2 * request, + palw_cuda_trace_status status, + palw_cuda_llama_mmvq_fault_v1 fault) { + if (request != nullptr && request->fault != nullptr) { + request->fault(request->user_data, status, fault, nullptr); + } +} + +static ggml_backend_cuda_context * ggml_backend_cuda_palw_context(void * backend_ptr) { + auto backend = static_cast(backend_ptr); + if (!ggml_backend_is_cuda(backend)) { + return nullptr; + } + return static_cast(backend->context); +} + +static palw_cuda_trace_status ggml_backend_cuda_palw_mmvq_kernel_info_v1( + void * backend_ptr, + uint8_t quantization, + palw_cuda_llama_mmvq_kernel_info_v1 * out_info) { + ggml_backend_cuda_context * ctx = ggml_backend_cuda_palw_context(backend_ptr); + if (ctx == nullptr || out_info == nullptr) { + return PALW_CUDA_TRACE_INVALID_ARGUMENT; + } + return ggml_cuda_palw_mmvq_kernel_info_v1( + ctx->device, quantization, out_info); +} + +static palw_cuda_trace_status ggml_backend_cuda_palw_trace_bind_v1( + void * backend_ptr, + const palw_cuda_llama_mmvq_request_v1 * request) { + ggml_backend_cuda_context * ctx = ggml_backend_cuda_palw_context(backend_ptr); + if (ctx == nullptr || + !palw_cuda_llama_mmvq_request_is_compatible_v1(request)) { + ggml_backend_cuda_palw_report_unbound_fault( + request, + PALW_CUDA_TRACE_INVALID_ARGUMENT, + PALW_CUDA_LLAMA_MMVQ_FAULT_INVALID_REQUEST); + return PALW_CUDA_TRACE_INVALID_ARGUMENT; + } + if (ctx->palw_trace.active()) { + ggml_backend_cuda_palw_report_unbound_fault( + request, + PALW_CUDA_TRACE_PENDING, + PALW_CUDA_LLAMA_MMVQ_FAULT_INVALID_REQUEST); + return PALW_CUDA_TRACE_PENDING; + } + if (ctx->curr_stream_no != 0 || !ctx->stream_context().concurrent_events.empty()) { + ggml_backend_cuda_palw_report_unbound_fault( + request, + PALW_CUDA_TRACE_INVALID_ARGUMENT, + PALW_CUDA_LLAMA_MMVQ_FAULT_CONCURRENCY); + return PALW_CUDA_TRACE_INVALID_ARGUMENT; + } +#ifdef USE_CUDA_GRAPH + if (ctx->any_cuda_graph_has_instance()) { + ggml_backend_cuda_palw_report_unbound_fault( + request, + PALW_CUDA_TRACE_INVALID_ARGUMENT, + PALW_CUDA_LLAMA_MMVQ_FAULT_GRAPH); + return PALW_CUDA_TRACE_INVALID_ARGUMENT; + } +#endif + ctx->stream_context().reset(); + ctx->palw_trace = {}; + ctx->palw_trace.request = request; + return PALW_CUDA_TRACE_OK; +} + +static palw_cuda_trace_status ggml_backend_cuda_palw_trace_associate_v1( + void * backend_ptr, + const palw_cuda_llama_mmvq_association_v1 * actual) { + ggml_backend_cuda_context * ctx = ggml_backend_cuda_palw_context(backend_ptr); + if (ctx == nullptr || !ggml_cuda_palw_trace_active(ctx) || actual == nullptr) { + return PALW_CUDA_TRACE_INVALID_ARGUMENT; + } + if (ggml_cuda_palw_trace_failed(ctx)) { + return ctx->palw_trace.status; + } + if (ctx->curr_stream_no != 0 || !ctx->stream_context().concurrent_events.empty()) { + ggml_cuda_palw_trace_fail( + ctx, + PALW_CUDA_TRACE_INVALID_ARGUMENT, + PALW_CUDA_LLAMA_MMVQ_FAULT_CONCURRENCY); + return ctx->palw_trace.status; + } + if (ctx->palw_trace.association_active || + ctx->palw_trace.accepted_launch_count >= + ctx->palw_trace.request->expected_launch_count) { + ggml_cuda_palw_trace_fail( + ctx, + PALW_CUDA_TRACE_SEQUENCE_MISMATCH, + PALW_CUDA_LLAMA_MMVQ_FAULT_ASSOCIATION_ORDER); + return ctx->palw_trace.status; + } + if (!palw_cuda_llama_mmvq_association_is_supported_v1(actual)) { + ctx->palw_trace.association = *actual; + ctx->palw_trace.association_active = true; + ggml_cuda_palw_trace_fail( + ctx, + PALW_CUDA_TRACE_INVALID_ARGUMENT, + PALW_CUDA_LLAMA_MMVQ_FAULT_UNSUPPORTED_DISPATCH); + return ctx->palw_trace.status; + } + + ctx->palw_trace.association = *actual; + ctx->palw_trace.association_active = true; + palw_cuda_producer_launch_v3 expected = {}; + const palw_cuda_trace_status approval = ctx->palw_trace.request->approve( + ctx->palw_trace.request->user_data, + &ctx->palw_trace.association, + &expected); + if (approval != PALW_CUDA_TRACE_OK) { + ggml_cuda_palw_trace_fail( + ctx, + approval, + PALW_CUDA_LLAMA_MMVQ_FAULT_ASSOCIATION_MISMATCH); + return ctx->palw_trace.status; + } + + uint64_t required_records = 0; + const palw_cuda_trace_record & base = expected.semantic_template.base.base; + if (palw_cuda_producer_validate_launch_v3(&expected, &required_records) != + PALW_CUDA_TRACE_OK || + required_records == 0 || base.m != actual->m || base.n != actual->n || + base.k != actual->k || base.batch != actual->batch || + base.output_rows != actual->output_rows || + base.output_columns != actual->output_columns || + base.quantization != actual->quantization) { + ggml_cuda_palw_trace_fail( + ctx, + PALW_CUDA_TRACE_IDENTITY_MISMATCH, + PALW_CUDA_LLAMA_MMVQ_FAULT_ASSOCIATION_MISMATCH); + return ctx->palw_trace.status; + } + ctx->palw_trace.expected_launch = expected; + return PALW_CUDA_TRACE_OK; +} + +static palw_cuda_trace_status ggml_backend_cuda_palw_trace_complete_v1( + void * backend_ptr, + const void * operation_token) { + ggml_backend_cuda_context * ctx = ggml_backend_cuda_palw_context(backend_ptr); + if (ctx == nullptr || !ggml_cuda_palw_trace_active(ctx)) { + return PALW_CUDA_TRACE_INVALID_ARGUMENT; + } + if (!ctx->palw_trace.association_active || + ctx->palw_trace.association.operation_token != operation_token || + !ctx->palw_trace.association_launched) { + ggml_cuda_palw_trace_fail( + ctx, + PALW_CUDA_TRACE_INCOMPLETE, + PALW_CUDA_LLAMA_MMVQ_FAULT_INCOMPLETE); + } + const palw_cuda_trace_status result = ctx->palw_trace.failed + ? ctx->palw_trace.status + : PALW_CUDA_TRACE_OK; + ctx->palw_trace.association = {}; + ctx->palw_trace.expected_launch = {}; + ctx->palw_trace.association_active = false; + ctx->palw_trace.association_launched = false; + return result; +} + +static palw_cuda_trace_status ggml_backend_cuda_palw_trace_unbind_v1( + void * backend_ptr, + const palw_cuda_llama_mmvq_request_v1 * request, + uint64_t * accepted_launch_count) { + ggml_backend_cuda_context * ctx = ggml_backend_cuda_palw_context(backend_ptr); + if (ctx == nullptr || !ggml_cuda_palw_trace_active(ctx) || request == nullptr || + ctx->palw_trace.request != request || accepted_launch_count == nullptr) { + return PALW_CUDA_TRACE_INVALID_ARGUMENT; + } + if (!ctx->palw_trace.failed && + (ctx->palw_trace.association_active || + ctx->palw_trace.accepted_launch_count != request->expected_launch_count || + palw_cuda_producer_trace_poisoned_v3(request->producer_context) != 0)) { + ggml_cuda_palw_trace_fail( + ctx, + PALW_CUDA_TRACE_INCOMPLETE, + PALW_CUDA_LLAMA_MMVQ_FAULT_INCOMPLETE); + } + *accepted_launch_count = ctx->palw_trace.accepted_launch_count; + const palw_cuda_trace_status result = ctx->palw_trace.failed + ? ctx->palw_trace.status + : PALW_CUDA_TRACE_OK; + ctx->palw_trace = {}; + return result; +} + +static palw_cuda_trace_status ggml_backend_cuda_palw_trace_bind_v2( + void * backend_ptr, + const palw_cuda_llama_trace_request_v2 * request) { + ggml_backend_cuda_context * ctx = ggml_backend_cuda_palw_context(backend_ptr); + if (ctx == nullptr || + !palw_cuda_llama_trace_request_is_compatible_v2(request)) { + ggml_backend_cuda_palw_report_unbound_fault_v2( + request, + PALW_CUDA_TRACE_INVALID_ARGUMENT, + PALW_CUDA_LLAMA_MMVQ_FAULT_INVALID_REQUEST); + return PALW_CUDA_TRACE_INVALID_ARGUMENT; + } + if (ctx->palw_trace.active()) { + ggml_backend_cuda_palw_report_unbound_fault_v2( + request, + PALW_CUDA_TRACE_PENDING, + PALW_CUDA_LLAMA_MMVQ_FAULT_INVALID_REQUEST); + return PALW_CUDA_TRACE_PENDING; + } + if (ctx->curr_stream_no != 0 || !ctx->stream_context().concurrent_events.empty()) { + ggml_backend_cuda_palw_report_unbound_fault_v2( + request, + PALW_CUDA_TRACE_INVALID_ARGUMENT, + PALW_CUDA_LLAMA_MMVQ_FAULT_CONCURRENCY); + return PALW_CUDA_TRACE_INVALID_ARGUMENT; + } +#ifdef USE_CUDA_GRAPH + if (ctx->any_cuda_graph_has_instance()) { + ggml_backend_cuda_palw_report_unbound_fault_v2( + request, + PALW_CUDA_TRACE_INVALID_ARGUMENT, + PALW_CUDA_LLAMA_MMVQ_FAULT_GRAPH); + return PALW_CUDA_TRACE_INVALID_ARGUMENT; + } +#endif + ctx->stream_context().reset(); + ctx->palw_trace = {}; + ctx->palw_trace.request_v2 = request; + return PALW_CUDA_TRACE_OK; +} + +static bool ggml_backend_cuda_palw_expected_matches_v2( + const palw_cuda_llama_trace_association_v2 & actual, + const palw_cuda_producer_launch_v3 & expected) { + const palw_cuda_trace_record_v3 & semantic = expected.semantic_template; + const palw_cuda_trace_record & base = semantic.base.base; + if (base.kind != actual.kind || base.decode_step != actual.decode_step || + base.m != actual.m || base.n != actual.n || base.k != actual.k || + base.batch != actual.batch || base.output_rows != actual.output_rows || + base.output_columns != actual.output_columns || base.phase != actual.phase || + base.quantization != actual.quantization || base.causal != actual.causal || + base.query_tokens != actual.query_tokens || + base.key_value_tokens != actual.key_value_tokens || + semantic.layer_present != actual.layer_present || + (actual.layer_present != 0 && base.layer_id != actual.layer_id)) { + return false; + } + if (actual.attention_stage == PALW_CUDA_TRACE_ATTENTION_STAGE_NONE) { + return semantic.attention_group_present == 0 && + semantic.attention_stage == PALW_CUDA_TRACE_ATTENTION_STAGE_NONE; + } + return semantic.attention_group_present == 1 && + semantic.attention_stage == actual.attention_stage && + semantic.query_heads == actual.query_heads && + semantic.key_value_heads == actual.key_value_heads && + semantic.head_dim == actual.head_dim && + semantic.logical_batch == actual.logical_batch && + semantic.physical_key_value_tokens == actual.physical_key_value_tokens; +} + +static palw_cuda_trace_status ggml_backend_cuda_palw_trace_associate_v2( + void * backend_ptr, + const palw_cuda_llama_trace_association_v2 * actual) { + ggml_backend_cuda_context * ctx = ggml_backend_cuda_palw_context(backend_ptr); + if (ctx == nullptr || !ggml_cuda_palw_trace_active(ctx) || + ctx->palw_trace.request_v2 == nullptr || actual == nullptr) { + return PALW_CUDA_TRACE_INVALID_ARGUMENT; + } + if (ggml_cuda_palw_trace_failed(ctx)) { + return ctx->palw_trace.status; + } + if (ctx->curr_stream_no != 0 || !ctx->stream_context().concurrent_events.empty()) { + ggml_cuda_palw_trace_fail( + ctx, + PALW_CUDA_TRACE_INVALID_ARGUMENT, + PALW_CUDA_LLAMA_MMVQ_FAULT_CONCURRENCY); + return ctx->palw_trace.status; + } + auto & trace = ctx->palw_trace; + if (trace.association_active || + trace.accepted_launch_count >= trace.expected_launch_count()) { + ggml_cuda_palw_trace_fail( + ctx, + PALW_CUDA_TRACE_SEQUENCE_MISMATCH, + PALW_CUDA_LLAMA_MMVQ_FAULT_ASSOCIATION_ORDER); + return trace.status; + } + trace.association_v2 = *actual; + trace.association_active = true; + if (!palw_cuda_llama_trace_association_is_supported_v2(actual)) { + ggml_cuda_palw_trace_fail( + ctx, + PALW_CUDA_TRACE_INVALID_ARGUMENT, + PALW_CUDA_LLAMA_MMVQ_FAULT_UNSUPPORTED_DISPATCH); + return trace.status; + } + if (actual->attention_stage == PALW_CUDA_TRACE_ATTENTION_STAGE_NONE) { + trace.association.abi_version = + PALW_CUDA_LLAMA_MMVQ_ASSOCIATION_ABI_VERSION_V1; + trace.association.struct_size = sizeof(trace.association); + trace.association.operation_token = actual->operation_token; + trace.association.weight_data = actual->src0.data; + trace.association.m = actual->m; + trace.association.n = actual->n; + trace.association.k = actual->k; + trace.association.batch = actual->batch; + trace.association.output_rows = actual->output_rows; + trace.association.output_columns = actual->output_columns; + trace.association.quantization = actual->quantization; + } + + palw_cuda_producer_launch_v3 expected = {}; + const palw_cuda_trace_status approval = trace.request_v2->approve( + trace.request_v2->user_data, + &trace.association_v2, + &expected); + if (approval != PALW_CUDA_TRACE_OK) { + ggml_cuda_palw_trace_fail( + ctx, + approval, + PALW_CUDA_LLAMA_MMVQ_FAULT_ASSOCIATION_MISMATCH); + return trace.status; + } + + uint64_t required_records = 0; + const palw_cuda_trace_status validation = + actual->attention_stage != PALW_CUDA_TRACE_ATTENTION_STAGE_NONE + ? palw_cuda_producer_validate_grouped_final_output_launch_v3( + &expected, &required_records) + : palw_cuda_producer_validate_launch_v3(&expected, &required_records); + if (validation != PALW_CUDA_TRACE_OK || required_records == 0 || + !ggml_backend_cuda_palw_expected_matches_v2(*actual, expected)) { + ggml_cuda_palw_trace_fail( + ctx, + PALW_CUDA_TRACE_IDENTITY_MISMATCH, + PALW_CUDA_LLAMA_MMVQ_FAULT_ASSOCIATION_MISMATCH); + return trace.status; + } + trace.expected_launch = expected; + return PALW_CUDA_TRACE_OK; +} + +static palw_cuda_trace_status ggml_backend_cuda_palw_trace_complete_v2( + void * backend_ptr, + const void * operation_token) { + ggml_backend_cuda_context * ctx = ggml_backend_cuda_palw_context(backend_ptr); + if (ctx == nullptr || !ggml_cuda_palw_trace_active(ctx) || + ctx->palw_trace.request_v2 == nullptr) { + return PALW_CUDA_TRACE_INVALID_ARGUMENT; + } + auto & trace = ctx->palw_trace; + if (!trace.association_active || + trace.association_v2.operation_token != operation_token || + !trace.association_launched) { + ggml_cuda_palw_trace_fail( + ctx, + PALW_CUDA_TRACE_INCOMPLETE, + PALW_CUDA_LLAMA_MMVQ_FAULT_INCOMPLETE); + } + const palw_cuda_trace_status result = trace.failed + ? trace.status + : PALW_CUDA_TRACE_OK; + trace.association = {}; + trace.association_v2 = {}; + trace.expected_launch = {}; + trace.association_active = false; + trace.association_launched = false; + return result; +} + +static palw_cuda_trace_status ggml_backend_cuda_palw_trace_unbind_v2( + void * backend_ptr, + const palw_cuda_llama_trace_request_v2 * request, + uint64_t * accepted_launch_count) { + ggml_backend_cuda_context * ctx = ggml_backend_cuda_palw_context(backend_ptr); + if (ctx == nullptr || !ggml_cuda_palw_trace_active(ctx) || request == nullptr || + ctx->palw_trace.request_v2 != request || accepted_launch_count == nullptr) { + return PALW_CUDA_TRACE_INVALID_ARGUMENT; + } + auto & trace = ctx->palw_trace; + if (!trace.failed && + (trace.association_active || + trace.accepted_launch_count != request->expected_launch_count || + palw_cuda_producer_trace_poisoned_v3(request->producer_context) != 0)) { + ggml_cuda_palw_trace_fail( + ctx, + PALW_CUDA_TRACE_INCOMPLETE, + PALW_CUDA_LLAMA_MMVQ_FAULT_INCOMPLETE); + } + *accepted_launch_count = trace.accepted_launch_count; + const palw_cuda_trace_status result = trace.failed + ? trace.status + : PALW_CUDA_TRACE_OK; + trace = {}; + return result; +} + +static palw_cuda_trace_status ggml_backend_cuda_palw_attention_kernel_info_v1( + void * backend_ptr, + const palw_cuda_llama_trace_association_v2 * actual, + palw_cuda_llama_attention_kernel_info_v1 * out_info) { + ggml_backend_cuda_context * ctx = ggml_backend_cuda_palw_context(backend_ptr); + if (ctx == nullptr || actual == nullptr || out_info == nullptr || + !palw_cuda_llama_trace_association_is_supported_v2(actual) || + actual->attention_stage == PALW_CUDA_TRACE_ATTENTION_STAGE_NONE) { + return PALW_CUDA_TRACE_INVALID_ARGUMENT; + } + if (actual->attention_stage == + PALW_CUDA_TRACE_ATTENTION_STAGE_EAGER_MASKED_SCALED_SOFTMAX) { + return ggml_cuda_palw_softmax_attention_kernel_info_v1( + ctx->device, actual, out_info); + } + return ggml_cuda_palw_mmvf_attention_kernel_info_v1( + ctx->device, actual, out_info); +} + +static palw_cuda_trace_status ggml_backend_cuda_palw_grouped_capture_info_v1( + void * backend_ptr, + palw_cuda_producer_grouped_capture_info_v3 * out_info) { + ggml_backend_cuda_context * ctx = ggml_backend_cuda_palw_context(backend_ptr); + if (ctx == nullptr || out_info == nullptr) { + return PALW_CUDA_TRACE_INVALID_ARGUMENT; + } + ggml_cuda_set_device(ctx->device); + return palw_cuda_producer_query_grouped_capture_info_v3(out_info); +} + +static palw_cuda_trace_status ggml_backend_cuda_palw_producer_create_v1( + void * backend_ptr, + uint64_t capacity, + palw_cuda_producer_trace_context_v3 ** out_context) { + ggml_backend_cuda_context * ctx = ggml_backend_cuda_palw_context(backend_ptr); + if (ctx == nullptr) { + return PALW_CUDA_TRACE_INVALID_ARGUMENT; + } + ggml_cuda_set_device(ctx->device); + return palw_cuda_producer_trace_create_v3(capacity, out_context); +} + +static void ggml_backend_cuda_palw_producer_destroy_v1( + void * backend_ptr, + palw_cuda_producer_trace_context_v3 * context) { + ggml_backend_cuda_context * ctx = ggml_backend_cuda_palw_context(backend_ptr); + if (ctx != nullptr) { + ggml_cuda_set_device(ctx->device); + } + palw_cuda_producer_trace_destroy_v3(context); +} + +static palw_cuda_trace_status ggml_backend_cuda_palw_producer_finalize_v1( + void * backend_ptr, + palw_cuda_producer_trace_context_v3 * context, + palw_cuda_trace_record_v3 * host_records, + uint64_t host_capacity) { + ggml_backend_cuda_context * ctx = ggml_backend_cuda_palw_context(backend_ptr); + if (ctx == nullptr || ctx->curr_stream_no != 0 || ggml_cuda_palw_trace_active(ctx)) { + return PALW_CUDA_TRACE_INVALID_ARGUMENT; + } + ggml_cuda_set_device(ctx->device); + return palw_cuda_producer_trace_finalize_v3( + context, ctx->stream(ctx->device, 0), host_records, host_capacity); +} + +static const palw_cuda_llama_mmvq_backend_api_v1 * +ggml_backend_cuda_palw_trace_api_v1() { + static const palw_cuda_llama_mmvq_backend_api_v1 api = { + PALW_CUDA_LLAMA_MMVQ_BACKEND_API_VERSION_V1, + sizeof(palw_cuda_llama_mmvq_backend_api_v1), + ggml_backend_cuda_palw_producer_create_v1, + ggml_backend_cuda_palw_producer_destroy_v1, + ggml_backend_cuda_palw_producer_finalize_v1, + palw_cuda_producer_trace_count_v3, + palw_cuda_producer_trace_committed_count_v3, + palw_cuda_producer_trace_faults_v3, + palw_cuda_producer_trace_poisoned_v3, + palw_cuda_trace_encode_record_v3, + }; + return &api; +} +#endif + static void * ggml_backend_cuda_reg_get_proc_address(ggml_backend_reg_t reg, const char * name) { GGML_UNUSED(reg); if (strcmp(name, "ggml_backend_comm_init") == 0) { @@ -5157,6 +5700,44 @@ static void * ggml_backend_cuda_reg_get_proc_address(ggml_backend_reg_t reg, con if (strcmp(name, "ggml_backend_get_features") == 0) { return (void *)ggml_backend_cuda_get_features; } +#if defined(GGML_CUDA_PALW_TRACE) + if (strcmp(name, PALW_CUDA_LLAMA_MMVQ_GET_API_PROC_V1) == 0) { + return (void *) ggml_backend_cuda_palw_trace_api_v1; + } + if (strcmp(name, PALW_CUDA_LLAMA_MMVQ_KERNEL_INFO_PROC_V1) == 0) { + return (void *) ggml_backend_cuda_palw_mmvq_kernel_info_v1; + } + if (strcmp(name, PALW_CUDA_LLAMA_MMVQ_BIND_PROC_V1) == 0) { + return (void *) ggml_backend_cuda_palw_trace_bind_v1; + } + if (strcmp(name, PALW_CUDA_LLAMA_MMVQ_ASSOCIATE_PROC_V1) == 0) { + return (void *) ggml_backend_cuda_palw_trace_associate_v1; + } + if (strcmp(name, PALW_CUDA_LLAMA_MMVQ_COMPLETE_PROC_V1) == 0) { + return (void *) ggml_backend_cuda_palw_trace_complete_v1; + } + if (strcmp(name, PALW_CUDA_LLAMA_MMVQ_UNBIND_PROC_V1) == 0) { + return (void *) ggml_backend_cuda_palw_trace_unbind_v1; + } + if (strcmp(name, PALW_CUDA_LLAMA_TRACE_BIND_PROC_V2) == 0) { + return (void *) ggml_backend_cuda_palw_trace_bind_v2; + } + if (strcmp(name, PALW_CUDA_LLAMA_TRACE_ASSOCIATE_PROC_V2) == 0) { + return (void *) ggml_backend_cuda_palw_trace_associate_v2; + } + if (strcmp(name, PALW_CUDA_LLAMA_TRACE_COMPLETE_PROC_V2) == 0) { + return (void *) ggml_backend_cuda_palw_trace_complete_v2; + } + if (strcmp(name, PALW_CUDA_LLAMA_TRACE_UNBIND_PROC_V2) == 0) { + return (void *) ggml_backend_cuda_palw_trace_unbind_v2; + } + if (strcmp(name, PALW_CUDA_LLAMA_ATTENTION_KERNEL_INFO_PROC_V1) == 0) { + return (void *) ggml_backend_cuda_palw_attention_kernel_info_v1; + } + if (strcmp(name, PALW_CUDA_LLAMA_GROUPED_CAPTURE_INFO_PROC_V1) == 0) { + return (void *) ggml_backend_cuda_palw_grouped_capture_info_v1; + } +#endif return nullptr; } diff --git a/ggml/src/ggml-cuda/mmvf.cu b/ggml/src/ggml-cuda/mmvf.cu index d7dbc8b9..13e3323a 100644 --- a/ggml/src/ggml-cuda/mmvf.cu +++ b/ggml/src/ggml-cuda/mmvf.cu @@ -409,6 +409,64 @@ static void mul_mat_vec_f_switch_fusion( } +#if defined(GGML_CUDA_PALW_TRACE) +template +static void ggml_cuda_palw_launch_attention_mmvf( + ggml_backend_cuda_context & ctx, + const half * x, + const float * y, + float * dst, + const ggml_tensor * dst_tensor, + const int64_t ncols, + const int64_t nrows, + const int64_t stride_row, + const int64_t stride_col_y, + const int64_t stride_col_dst, + const int64_t nchannels_x, + const int64_t nchannels_y, + const int64_t nchannels_dst, + const int64_t stride_channel_x, + const int64_t stride_channel_y, + const int64_t stride_channel_dst, + const int64_t nsamples_x, + const int64_t nsamples_dst, + const int64_t stride_sample_x, + const int64_t stride_sample_y, + const int64_t stride_sample_dst) { + GGML_UNUSED(nchannels_y); + const uint3 nchannels_y_fd = make_uint3(0, 0, 0); + const uint3 channel_ratio_fd = init_fastdiv_values(nchannels_dst / nchannels_x); + const uint3 sample_ratio_fd = init_fastdiv_values(nsamples_dst / nsamples_x); + const ggml_cuda_mm_fusion_args_device fusion = {}; + const dim3 block_nums(nrows, nchannels_dst, nsamples_dst); + const dim3 block_dims(block_size, 1, 1); + constexpr int nbytes_shared = 32 * sizeof(float); + cudaStream_t stream = ctx.stream(); + + mul_mat_vec_f_switch_fusion( + x, y, nullptr, fusion, dst, ncols / 2, nchannels_y_fd, + stride_row, stride_col_y / 2, stride_col_dst, + channel_ratio_fd, stride_channel_x, stride_channel_y, stride_channel_dst, + sample_ratio_fd, stride_sample_x, stride_sample_y, stride_sample_dst, + block_dims, block_nums, nbytes_shared, 0, stream); + + ggml_cuda_palw_attention_work_launch_v1 work = {}; + work.entry_point = reinterpret_cast( + mul_mat_vec_f); + work.dimensions.grid_x = block_nums.x; + work.dimensions.grid_y = block_nums.y; + work.dimensions.grid_z = block_nums.z; + work.dimensions.block_x = block_dims.x; + work.dimensions.block_y = block_dims.y; + work.dimensions.block_z = block_dims.z; + work.dimensions.dynamic_shared_memory_bytes = nbytes_shared; + work.work_tensor_type = PALW_CUDA_LLAMA_TENSOR_TYPE_F16; + work.mask_tensor_type = PALW_CUDA_LLAMA_TENSOR_TYPE_NONE; + work.output_row_stride = dst_tensor->nb[1] / sizeof(float); + ggml_cuda_palw_trace_enqueue_attention(&ctx, dst_tensor, work); +} +#endif + template void launch_mul_mat_vec_f_cuda( const T * x, const float * y, const int32_t * ids, const ggml_cuda_mm_fusion_args_device fusion, float * dst, @@ -698,6 +756,55 @@ void ggml_cuda_mul_mat_vec_f(ggml_backend_cuda_context & ctx, const ggml_tensor const int64_t ids_stride = ids ? ids->nb[1] / ggml_type_size(ids->type) : 0; +#if defined(GGML_CUDA_PALW_TRACE) + if (ggml_cuda_palw_trace_v2_attention_active(&ctx)) { + const auto & association = ctx.palw_trace.association_v2; + const uint8_t stage = association.attention_stage; + const bool qk = stage == + PALW_CUDA_TRACE_ATTENTION_STAGE_EAGER_QK_SCORE_MMVF; + const bool pv = stage == + PALW_CUDA_TRACE_ATTENTION_STAGE_EAGER_VALUE_AGGREGATION_MMVF; + const uint8_t expected_precision = qk + ? PALW_CUDA_LLAMA_WORK_PRECISION_F32 + : PALW_CUDA_LLAMA_WORK_PRECISION_DEFAULT; + const bool preflight = (qk || pv) && + ggml_cuda_palw_trace_attention_preflight( + &ctx, src0, src1, nullptr, dst, stage, + expected_precision, PALW_CUDA_LLAMA_ATTENTION_MASK_NONE); + const bool exact = preflight && ids == nullptr && fusion == nullptr && + src0->type == GGML_TYPE_F16 && ncols_dst == 1 && + ne02 == 8 && ne03 == 1 && nchannels_y == 32 && + nchannels_dst == 32 && ne3 == 1 && dst->nb[1] == dst->nb[2] && + association.physical_key_value_tokens == 256 && + ((qk && prec == GGML_PREC_F32 && ne00 == 128 && ne01 == 256) || + (pv && prec == GGML_PREC_DEFAULT && ne00 == 256 && ne01 == 128)); + if (!exact) { + if (!ctx.palw_trace.failed) { + ggml_cuda_palw_trace_fail( + &ctx, + PALW_CUDA_TRACE_IDENTITY_MISMATCH, + PALW_CUDA_LLAMA_MMVQ_FAULT_UNSUPPORTED_DISPATCH); + } + return; + } + const half * src0_d = static_cast(src0->data); + if (qk) { + ggml_cuda_palw_launch_attention_mmvf( + ctx, src0_d, src1_d, dst_d, dst, ne00, ne01, s01, + stride_col_y, stride_col_dst, ne02, nchannels_y, nchannels_dst, + s02, stride_channel_y, stride_channel_dst, ne03, ne3, + s03, s13, s3); + } else { + ggml_cuda_palw_launch_attention_mmvf( + ctx, src0_d, src1_d, dst_d, dst, ne00, ne01, s01, + stride_col_y, stride_col_dst, ne02, nchannels_y, nchannels_dst, + s02, stride_channel_y, stride_channel_dst, ne03, ne3, + s03, s13, s3); + } + return; + } +#endif + switch (src0->type) { case GGML_TYPE_F32: { const float * src0_d = (const float *) src0->data; @@ -722,6 +829,70 @@ void ggml_cuda_mul_mat_vec_f(ggml_backend_cuda_context & ctx, const ggml_tensor } } +#if defined(GGML_CUDA_PALW_TRACE) +palw_cuda_trace_status ggml_cuda_palw_mmvf_attention_kernel_info_v1( + int device, + const palw_cuda_llama_trace_association_v2 * actual, + palw_cuda_llama_attention_kernel_info_v1 * out_info) { + if (device < 0 || actual == nullptr || out_info == nullptr || + actual->kind != PALW_CUDA_TRACE_GEMM || + actual->physical_key_value_tokens != 256 || + (actual->attention_stage != + PALW_CUDA_TRACE_ATTENTION_STAGE_EAGER_QK_SCORE_MMVF && + actual->attention_stage != + PALW_CUDA_TRACE_ATTENTION_STAGE_EAGER_VALUE_AGGREGATION_MMVF)) { + return PALW_CUDA_TRACE_INVALID_ARGUMENT; + } + *out_info = {}; + const bool qk = actual->attention_stage == + PALW_CUDA_TRACE_ATTENTION_STAGE_EAGER_QK_SCORE_MMVF; + const void * entry_point = qk + ? reinterpret_cast( + mul_mat_vec_f) + : reinterpret_cast( + mul_mat_vec_f); + cudaDeviceProp properties = {}; + cudaFuncAttributes attributes = {}; + ggml_cuda_set_device(device); + cudaError_t error = cudaGetDeviceProperties(&properties, device); + if (error == cudaSuccess) { + error = cudaFuncGetAttributes(&attributes, entry_point); + } + if (error != cudaSuccess) { + (void) cudaGetLastError(); + return PALW_CUDA_TRACE_CUDA_ERROR; + } + palw_cuda_llama_attention_kernel_info_v1 result = {}; + result.abi_version = PALW_CUDA_LLAMA_ATTENTION_KERNEL_INFO_ABI_VERSION_V1; + result.struct_size = sizeof(result); + result.attention_stage = actual->attention_stage; + result.work_tensor_type = PALW_CUDA_LLAMA_TENSOR_TYPE_F16; + result.mask_tensor_type = PALW_CUDA_LLAMA_TENSOR_TYPE_NONE; + result.actual_identity = ggml_cuda_palw_attention_actual_identity_v3( + actual->attention_stage); + result.dimensions.grid_x = qk ? 256 : 128; + result.dimensions.grid_y = 32; + result.dimensions.grid_z = 1; + result.dimensions.block_x = qk ? 64 : 128; + result.dimensions.block_y = 1; + result.dimensions.block_z = 1; + result.dimensions.dynamic_shared_memory_bytes = 128; + result.sm_arch = static_cast(properties.major * 10 + properties.minor); + result.binary_version = attributes.binaryVersion; + result.ptx_version = attributes.ptxVersion; + result.num_regs = attributes.numRegs; + result.max_threads_per_block = attributes.maxThreadsPerBlock; + result.static_shared_memory_bytes = static_cast(attributes.sharedSizeBytes); + result.local_memory_bytes = static_cast(attributes.localSizeBytes); + result.output_row_stride = actual->dst.nb[1] / sizeof(float); + if (!palw_cuda_llama_attention_kernel_info_is_compatible_v1(&result)) { + return PALW_CUDA_TRACE_IDENTITY_MISMATCH; + } + *out_info = result; + return PALW_CUDA_TRACE_OK; +} +#endif + void ggml_cuda_op_mul_mat_vec_f( ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst, const char * src0_dd_i, const float * src1_ddf_i, diff --git a/ggml/src/ggml-cuda/mmvf.cuh b/ggml/src/ggml-cuda/mmvf.cuh index a50f7c02..3baabfdf 100644 --- a/ggml/src/ggml-cuda/mmvf.cuh +++ b/ggml/src/ggml-cuda/mmvf.cuh @@ -12,3 +12,10 @@ void ggml_cuda_op_mul_mat_vec_f( const int64_t src1_padded_row_size, cudaStream_t stream); bool ggml_cuda_should_use_mmvf(enum ggml_type type, int cc, const int64_t * src0_ne, const size_t * src0_nb, int64_t ne11); + +#if defined(GGML_CUDA_PALW_TRACE) +palw_cuda_trace_status ggml_cuda_palw_mmvf_attention_kernel_info_v1( + int device, + const palw_cuda_llama_trace_association_v2 * actual, + palw_cuda_llama_attention_kernel_info_v1 * out_info); +#endif diff --git a/ggml/src/ggml-cuda/mmvq.cu b/ggml/src/ggml-cuda/mmvq.cu index e18ada53..c84b5420 100644 --- a/ggml/src/ggml-cuda/mmvq.cu +++ b/ggml/src/ggml-cuda/mmvq.cu @@ -3,6 +3,10 @@ #include "unary.cuh" #include "vecdotq.cuh" +#if defined(GGML_CUDA_PALW_TRACE) +#include "palw_cuda_producer_trace.cuh" +#endif + #include typedef float (*vec_dot_q_cuda_t)(const void * __restrict__ vbq, const block_q8_1 * __restrict__ bq8_1, const int & kbx, const int & iqs); @@ -475,6 +479,75 @@ static constexpr __host__ __device__ int calc_rows_per_block(int ncols_dst, int return 1; } +#if defined(GGML_CUDA_PALW_TRACE) +static_assert(PALW_CUDA_PRODUCER_MMVQ_KERNEL_ARGUMENT_COUNT_V3 == 20U); +static_assert(PALW_CUDA_PRODUCER_MMVQ_VIEW_ARGUMENT_INDEX_V3 == 19U); +static_assert( + PALW_CUDA_PRODUCER_MMVQ_VIEW_ARGUMENT_INDEX_V3 + 1U == + PALW_CUDA_PRODUCER_MMVQ_KERNEL_ARGUMENT_COUNT_V3); + +template +static __host__ __device__ __forceinline__ palw_cuda_producer_actual_identity_v3 +palw_mmvq_actual_identity_v3() { + palw_cuda_producer_actual_identity_v3 identity = {}; + if constexpr (type == GGML_TYPE_Q4_K && ncols_dst == 1 && !has_fusion && !small_k) { + const uint8_t variant[PALW_CUDA_PRODUCER_ID_SIZE] = { + 0xbb, 0x83, 0x34, 0xa8, 0xf1, 0x31, 0xc5, 0xe9, + 0xba, 0xd2, 0x90, 0x2a, 0xd4, 0xd2, 0xe4, 0x8b, + 0xbc, 0xd2, 0xee, 0x6a, 0x62, 0x2f, 0xee, 0xf4, + 0xb5, 0xbb, 0x1b, 0xb2, 0x47, 0x84, 0xe8, 0xcb, + }; + const uint8_t work[PALW_CUDA_PRODUCER_ID_SIZE] = { + 0xe6, 0x75, 0x72, 0x58, 0x7d, 0x81, 0xeb, 0xab, + 0x61, 0x08, 0x81, 0xfd, 0x51, 0x05, 0x99, 0x98, + 0x70, 0x2f, 0xab, 0x2f, 0xbd, 0x1f, 0x51, 0x02, + 0x81, 0x2b, 0xe4, 0x61, 0x81, 0xfb, 0x79, 0xb5, + }; +#if defined(__CUDA_ARCH__) +#pragma unroll +#endif + for (uint32_t i = 0; i < PALW_CUDA_PRODUCER_ID_SIZE; ++i) { + identity.producer_variant_id[i] = variant[i]; + identity.work_entry_point_id[i] = work[i]; + } + } else if constexpr (type == GGML_TYPE_Q6_K && ncols_dst == 1 && !has_fusion && !small_k) { + const uint8_t variant[PALW_CUDA_PRODUCER_ID_SIZE] = { + 0x85, 0x5c, 0x0d, 0x09, 0xac, 0xbf, 0x49, 0x4b, + 0x38, 0x8b, 0x17, 0x2d, 0x67, 0x02, 0x6e, 0xa3, + 0x57, 0x60, 0xd3, 0x7b, 0xfb, 0xce, 0x7e, 0x60, + 0x72, 0x59, 0x43, 0x85, 0x0f, 0x84, 0x70, 0x40, + }; + const uint8_t work[PALW_CUDA_PRODUCER_ID_SIZE] = { + 0xf7, 0xc9, 0x5c, 0xeb, 0x94, 0x5a, 0xf4, 0xfc, + 0x2c, 0x51, 0x0c, 0xe2, 0xb9, 0xe8, 0x79, 0x36, + 0x2f, 0xa4, 0xfc, 0x2e, 0xff, 0xae, 0x66, 0xb1, + 0xb7, 0xd0, 0xf3, 0xb3, 0x27, 0x2c, 0x32, 0x50, + }; +#if defined(__CUDA_ARCH__) +#pragma unroll +#endif + for (uint32_t i = 0; i < PALW_CUDA_PRODUCER_ID_SIZE; ++i) { + identity.producer_variant_id[i] = variant[i]; + identity.work_entry_point_id[i] = work[i]; + } + } + const uint8_t capture[PALW_CUDA_PRODUCER_ID_SIZE] = { + 0xb9, 0xd5, 0xfd, 0xda, 0x45, 0xdb, 0x7c, 0xec, + 0x8c, 0x86, 0x52, 0xfc, 0x60, 0xee, 0x97, 0xa0, + 0x2b, 0xb0, 0x9e, 0x83, 0xdf, 0xee, 0x57, 0xf2, + 0xfc, 0x2e, 0xc5, 0x2c, 0x51, 0xab, 0xde, 0xb7, + }; +#if defined(__CUDA_ARCH__) +#pragma unroll +#endif + for (uint32_t i = 0; i < PALW_CUDA_PRODUCER_ID_SIZE; ++i) { + identity.capture_implementation_id[i] = capture[i]; + } + return identity; +} + +#endif + template __launch_bounds__(calc_nwarps(type, ncols_dst, get_device_table_id())*ggml_cuda_get_physical_warp_size(), 1) static __global__ void mul_mat_vec_q( @@ -483,7 +556,11 @@ static __global__ void mul_mat_vec_q( const uint32_t stride_col_dst, const uint3 channel_ratio, const uint32_t stride_channel_x, const uint32_t stride_channel_y, const uint32_t stride_channel_dst, const uint3 sample_ratio, const uint32_t stride_sample_x, const uint32_t stride_sample_y, const uint32_t stride_sample_dst, - const uint32_t ids_stride) { + const uint32_t ids_stride +#if defined(GGML_CUDA_PALW_TRACE) + , const palw_cuda_producer_device_view_v3 palw_trace_view +#endif + ) { const void * GGML_CUDA_RESTRICT vx = vx_ptr; const void * GGML_CUDA_RESTRICT vy = vy_ptr; const int32_t * GGML_CUDA_RESTRICT ids = ids_ptr; @@ -657,6 +734,18 @@ static __global__ void mul_mat_vec_q( } if (threadIdx.x == i && (rows_per_cuda_block == 1 || uint32_t(row0 + i) < stride_col_dst)) { +#if defined(GGML_CUDA_PALW_TRACE) + if (palw_trace_view.abi_version != 0U) { + const auto actual_identity = + palw_mmvq_actual_identity_v3(); + palw_cuda_producer_capture_accumulator_f32_v3( + palw_trace_view, + static_cast(j), + static_cast(row0 + i), + tmp[j][i], + actual_identity); + } +#endif float result = tmp[j][i]; if constexpr (has_fusion) { if constexpr (type == GGML_TYPE_NVFP4) { @@ -698,6 +787,66 @@ static __global__ void mul_mat_vec_q( } } +#if defined(GGML_CUDA_PALW_TRACE) +template +static palw_cuda_trace_status palw_mmvq_kernel_info_for_type_v1( + int device, + uint8_t quantization, + palw_cuda_llama_mmvq_kernel_info_v1 * out_info) { + cudaDeviceProp properties = {}; + cudaFuncAttributes attributes = {}; + ggml_cuda_set_device(device); + cudaError_t error = cudaGetDeviceProperties(&properties, device); + if (error == cudaSuccess) { + const void * entry_point = reinterpret_cast( + mul_mat_vec_q); + error = cudaFuncGetAttributes(&attributes, entry_point); + } + if (error != cudaSuccess) { + (void) cudaGetLastError(); + return PALW_CUDA_TRACE_CUDA_ERROR; + } + + palw_cuda_llama_mmvq_kernel_info_v1 result = {}; + result.abi_version = PALW_CUDA_LLAMA_MMVQ_KERNEL_INFO_ABI_VERSION_V1; + result.struct_size = sizeof(result); + result.quantization = quantization; + result.actual_identity = palw_mmvq_actual_identity_v3(); + result.sm_arch = static_cast(properties.major * 10 + properties.minor); + result.binary_version = attributes.binaryVersion; + result.ptx_version = attributes.ptxVersion; + result.num_regs = attributes.numRegs; + result.max_threads_per_block = attributes.maxThreadsPerBlock; + result.static_shared_memory_bytes = static_cast(attributes.sharedSizeBytes); + result.local_memory_bytes = static_cast(attributes.localSizeBytes); + if (!palw_cuda_llama_mmvq_kernel_info_is_compatible_v1(&result)) { + return PALW_CUDA_TRACE_IDENTITY_MISMATCH; + } + *out_info = result; + return PALW_CUDA_TRACE_OK; +} + +palw_cuda_trace_status ggml_cuda_palw_mmvq_kernel_info_v1( + int device, + uint8_t quantization, + palw_cuda_llama_mmvq_kernel_info_v1 * out_info) { + if (device < 0 || out_info == nullptr) { + return PALW_CUDA_TRACE_INVALID_ARGUMENT; + } + *out_info = {}; + switch (quantization) { + case PALW_CUDA_TRACE_QUANTIZATION_Q4_K: + return palw_mmvq_kernel_info_for_type_v1( + device, quantization, out_info); + case PALW_CUDA_TRACE_QUANTIZATION_Q6_K: + return palw_mmvq_kernel_info_for_type_v1( + device, quantization, out_info); + default: + return PALW_CUDA_TRACE_INVALID_ARGUMENT; + } +} +#endif + // Dedicated MoE multi-token kernel. // Grid: (ceil(nrows_x / c_rows_per_block), nchannels_dst) // Block: (warp_size, ncols_dst) - each warp handles one token independently. @@ -788,17 +937,140 @@ static void mul_mat_vec_q_switch_fusion( const uint32_t stride_channel_y, const uint32_t stride_channel_dst, const uint3 sample_ratio, const uint32_t stride_sample_x, const uint32_t stride_sample_y, const uint32_t stride_sample_dst, const dim3 & block_nums, const dim3 & block_dims, const int nbytes_shared, - const uint32_t ids_stride, cudaStream_t stream) { + const uint32_t ids_stride, cudaStream_t stream +#if defined(GGML_CUDA_PALW_TRACE) + , ggml_backend_cuda_context * palw_ctx, const void * palw_operation_token +#endif + ) { const bool has_fusion = fusion.gate != nullptr || fusion.x_bias != nullptr || fusion.gate_bias != nullptr || fusion.x_scale != nullptr || fusion.gate_scale != nullptr; + +#if defined(GGML_CUDA_PALW_TRACE) + if (ggml_cuda_palw_trace_active(palw_ctx)) { + if (ggml_cuda_palw_trace_failed(palw_ctx)) { + return; + } + if constexpr ((type != GGML_TYPE_Q4_K && type != GGML_TYPE_Q6_K) || + c_ncols_dst != 1 || small_k) { + ggml_cuda_palw_trace_fail( + palw_ctx, + PALW_CUDA_TRACE_INVALID_ARGUMENT, + PALW_CUDA_LLAMA_MMVQ_FAULT_UNSUPPORTED_DISPATCH); + return; + } else { + auto & trace = palw_ctx->palw_trace; + const auto actual_identity = + palw_mmvq_actual_identity_v3(); + bool identity_matches = true; + for (uint32_t i = 0; i < PALW_CUDA_PRODUCER_ID_SIZE; ++i) { + identity_matches = identity_matches && + trace.expected_launch.semantic_template.base.producer_variant_id[i] == + actual_identity.producer_variant_id[i] && + trace.expected_launch.semantic_template.work_entry_point_id[i] == + actual_identity.work_entry_point_id[i] && + trace.expected_launch.semantic_template.capture_implementation_id[i] == + actual_identity.capture_implementation_id[i]; + } + if (!trace.association_active || trace.association_launched || + trace.operation_token() != palw_operation_token || + trace.association.weight_data != vx || has_fusion || ids != nullptr || + ggml_cuda_info().devices[palw_ctx->device].cc != GGML_CUDA_CC_ADA_LOVELACE || + block_nums.x != trace.association.output_columns || block_nums.y != 1U || + block_nums.z != 1U || block_dims.x != 32U || block_dims.y != 4U || + block_dims.z != 1U || nbytes_shared != 0 || + ncols_x != trace.association.k || !identity_matches) { + ggml_cuda_palw_trace_fail( + palw_ctx, + PALW_CUDA_TRACE_IDENTITY_MISMATCH, + PALW_CUDA_LLAMA_MMVQ_FAULT_ASSOCIATION_MISMATCH); + return; + } + + palw_cuda_producer_launch_dimensions_v1 dimensions = {}; + dimensions.grid_x = block_nums.x; + dimensions.grid_y = block_nums.y; + dimensions.grid_z = block_nums.z; + dimensions.block_x = block_dims.x; + dimensions.block_y = block_dims.y; + dimensions.block_z = block_dims.z; + dimensions.dynamic_shared_memory_bytes = static_cast(nbytes_shared); + palw_cuda_producer_device_view_v3 prepared_view = {}; + const void * entry_point = reinterpret_cast( + mul_mat_vec_q); + palw_cuda_trace_status status = palw_cuda_producer_trace_prepare_v3( + trace.producer_context(), + stream, + &trace.expected_launch, + entry_point, + &dimensions, + PALW_CUDA_PRODUCER_MMVQ_KERNEL_ARGUMENT_COUNT_V3, + PALW_CUDA_PRODUCER_MMVQ_VIEW_ARGUMENT_INDEX_V3, + &prepared_view); + if (status != PALW_CUDA_TRACE_OK) { + ggml_cuda_palw_trace_fail( + palw_ctx, + status, + PALW_CUDA_LLAMA_MMVQ_FAULT_PREPARE); + return; + } + + auto argument_pointer = [](const auto * value) -> void * { + return const_cast(static_cast(value)); + }; + palw_cuda_producer_device_view_v3 placeholder_view = {}; + void * kernel_arguments[PALW_CUDA_PRODUCER_MMVQ_KERNEL_ARGUMENT_COUNT_V3] = { + argument_pointer(&vx), + argument_pointer(&vy), + argument_pointer(&ids), + argument_pointer(&fusion), + argument_pointer(&dst), + argument_pointer(&ncols_x), + argument_pointer(&nchannels_y), + argument_pointer(&stride_row_x), + argument_pointer(&stride_col_y), + argument_pointer(&stride_col_dst), + argument_pointer(&channel_ratio), + argument_pointer(&stride_channel_x), + argument_pointer(&stride_channel_y), + argument_pointer(&stride_channel_dst), + argument_pointer(&sample_ratio), + argument_pointer(&stride_sample_x), + argument_pointer(&stride_sample_y), + argument_pointer(&stride_sample_dst), + argument_pointer(&ids_stride), + argument_pointer(&placeholder_view), + }; + status = palw_cuda_producer_trace_launch_v3( + trace.producer_context(), + stream, + entry_point, + kernel_arguments); + if (status != PALW_CUDA_TRACE_OK) { + ggml_cuda_palw_trace_fail( + palw_ctx, + status, + PALW_CUDA_LLAMA_MMVQ_FAULT_LAUNCH); + return; + } + ggml_cuda_palw_trace_launch_accepted(palw_ctx); + return; + } + } + const palw_cuda_producer_device_view_v3 palw_trace_view = {}; +#endif + if constexpr (c_ncols_dst == 1) { if (has_fusion) { const ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params(block_nums, block_dims, nbytes_shared, stream); ggml_cuda_kernel_launch(mul_mat_vec_q, launch_params, vx, vy, ids, fusion, dst, ncols_x, nchannels_y, stride_row_x, stride_col_y, stride_col_dst, channel_ratio, stride_channel_x, stride_channel_y, stride_channel_dst, - sample_ratio, stride_sample_x, stride_sample_y, stride_sample_dst, ids_stride); + sample_ratio, stride_sample_x, stride_sample_y, stride_sample_dst, ids_stride +#if defined(GGML_CUDA_PALW_TRACE) + , palw_trace_view +#endif + ); return; } } @@ -809,7 +1081,11 @@ static void mul_mat_vec_q_switch_fusion( ggml_cuda_kernel_launch(mul_mat_vec_q, launch_params, vx, vy, ids, fusion, dst, ncols_x, nchannels_y, stride_row_x, stride_col_y, stride_col_dst, channel_ratio, stride_channel_x, stride_channel_y, stride_channel_dst, - sample_ratio, stride_sample_x, stride_sample_y, stride_sample_dst, ids_stride); + sample_ratio, stride_sample_x, stride_sample_y, stride_sample_dst, ids_stride +#if defined(GGML_CUDA_PALW_TRACE) + , palw_trace_view +#endif + ); } template @@ -842,7 +1118,12 @@ static void mul_mat_vec_q_switch_ncols_dst( const int nchannels_x, const int nchannels_y, const int nchannels_dst, const int stride_channel_x, const int stride_channel_y, const int stride_channel_dst, const int nsamples_x, const int nsamples_dst, const int stride_sample_x, const int stride_sample_y, const int stride_sample_dst, - const int ids_stride, cudaStream_t stream) { + const int ids_stride, cudaStream_t stream +#if defined(GGML_CUDA_PALW_TRACE) + , ggml_backend_cuda_context * palw_ctx = nullptr, + const void * palw_operation_token = nullptr +#endif + ) { GGML_ASSERT(ncols_x % ggml_blck_size(type) == 0); GGML_ASSERT(ncols_dst <= MMVQ_MAX_BATCH_SIZE); @@ -858,6 +1139,26 @@ static void mul_mat_vec_q_switch_ncols_dst( const bool has_ids = ids != nullptr; +#if defined(GGML_CUDA_PALW_TRACE) + if (ggml_cuda_palw_trace_active(palw_ctx)) { + if (ggml_cuda_palw_trace_failed(palw_ctx)) { + return; + } + if (has_ids || ncols_dst != 1 || + nchannels_x != 1 || nchannels_y != 1 || nchannels_dst != 1 || + nsamples_x != 1 || nsamples_dst != 1 || + cc != GGML_CUDA_CC_ADA_LOVELACE || warp_size != 32 || + table_id != MMVQ_PARAMETERS_GENERIC || + (type != GGML_TYPE_Q4_K && type != GGML_TYPE_Q6_K)) { + ggml_cuda_palw_trace_fail( + palw_ctx, + PALW_CUDA_TRACE_INVALID_ARGUMENT, + PALW_CUDA_LLAMA_MMVQ_FAULT_UNSUPPORTED_DISPATCH); + return; + } + } +#endif + const auto should_use_small_k = [&](int c_ncols_dst) { // When K is small, increase rows_per_block to match nwarps so each warp has more work to do // Trigger when the full thread block covers all K blocks in a single loop iteration and few threads remain idle. @@ -923,7 +1224,11 @@ static void mul_mat_vec_q_switch_ncols_dst( vx, vy, ids, fusion, dst, ncols_x, nchannels_y_fd, stride_row_x, stride_col_y, stride_col_dst, channel_ratio_fd, stride_channel_x, stride_channel_y, stride_channel_dst, sample_ratio_fd, stride_sample_x, stride_sample_y, stride_sample_dst, dims.first, dims.second, 0, ids_stride, - stream); + stream +#if defined(GGML_CUDA_PALW_TRACE) + , palw_ctx, palw_operation_token +#endif + ); } else { std::pair dims = calc_launch_params(c_ncols_dst, nrows_x, nchannels_dst, nsamples_dst, warp_size, table_id); @@ -931,7 +1236,11 @@ static void mul_mat_vec_q_switch_ncols_dst( vx, vy, ids, fusion, dst, ncols_x, nchannels_y_fd, stride_row_x, stride_col_y, stride_col_dst, channel_ratio_fd, stride_channel_x, stride_channel_y, stride_channel_dst, sample_ratio_fd, stride_sample_x, stride_sample_y, stride_sample_dst, dims.first, dims.second, 0, ids_stride, - stream); + stream +#if defined(GGML_CUDA_PALW_TRACE) + , palw_ctx, palw_operation_token +#endif + ); } } break; case 2: { @@ -940,7 +1249,11 @@ static void mul_mat_vec_q_switch_ncols_dst( mul_mat_vec_q_switch_fusion(vx, vy, ids, fusion, dst, ncols_x, nchannels_y_fd, stride_row_x, stride_col_y, stride_col_dst, channel_ratio_fd, stride_channel_x, stride_channel_y, stride_channel_dst, sample_ratio_fd, stride_sample_x, stride_sample_y, stride_sample_dst, - dims.first, dims.second, 0, ids_stride, stream); + dims.first, dims.second, 0, ids_stride, stream +#if defined(GGML_CUDA_PALW_TRACE) + , palw_ctx, palw_operation_token +#endif + ); } break; case 3: { constexpr int c_ncols_dst = 3; @@ -948,7 +1261,11 @@ static void mul_mat_vec_q_switch_ncols_dst( mul_mat_vec_q_switch_fusion(vx, vy, ids, fusion, dst, ncols_x, nchannels_y_fd, stride_row_x, stride_col_y, stride_col_dst, channel_ratio_fd, stride_channel_x, stride_channel_y, stride_channel_dst, sample_ratio_fd, stride_sample_x, stride_sample_y, stride_sample_dst, - dims.first, dims.second, 0, ids_stride, stream); + dims.first, dims.second, 0, ids_stride, stream +#if defined(GGML_CUDA_PALW_TRACE) + , palw_ctx, palw_operation_token +#endif + ); } break; case 4: { constexpr int c_ncols_dst = 4; @@ -956,7 +1273,11 @@ static void mul_mat_vec_q_switch_ncols_dst( mul_mat_vec_q_switch_fusion(vx, vy, ids, fusion, dst, ncols_x, nchannels_y_fd, stride_row_x, stride_col_y, stride_col_dst, channel_ratio_fd, stride_channel_x, stride_channel_y, stride_channel_dst, sample_ratio_fd, stride_sample_x, stride_sample_y, stride_sample_dst, - dims.first, dims.second, 0, ids_stride, stream); + dims.first, dims.second, 0, ids_stride, stream +#if defined(GGML_CUDA_PALW_TRACE) + , palw_ctx, palw_operation_token +#endif + ); } break; case 5: { constexpr int c_ncols_dst = 5; @@ -964,7 +1285,11 @@ static void mul_mat_vec_q_switch_ncols_dst( mul_mat_vec_q_switch_fusion(vx, vy, ids, fusion, dst, ncols_x, nchannels_y_fd, stride_row_x, stride_col_y, stride_col_dst, channel_ratio_fd, stride_channel_x, stride_channel_y, stride_channel_dst, sample_ratio_fd, stride_sample_x, stride_sample_y, stride_sample_dst, - dims.first, dims.second, 0, ids_stride, stream); + dims.first, dims.second, 0, ids_stride, stream +#if defined(GGML_CUDA_PALW_TRACE) + , palw_ctx, palw_operation_token +#endif + ); } break; case 6: { constexpr int c_ncols_dst = 6; @@ -972,7 +1297,11 @@ static void mul_mat_vec_q_switch_ncols_dst( mul_mat_vec_q_switch_fusion(vx, vy, ids, fusion, dst, ncols_x, nchannels_y_fd, stride_row_x, stride_col_y, stride_col_dst, channel_ratio_fd, stride_channel_x, stride_channel_y, stride_channel_dst, sample_ratio_fd, stride_sample_x, stride_sample_y, stride_sample_dst, - dims.first, dims.second, 0, ids_stride, stream); + dims.first, dims.second, 0, ids_stride, stream +#if defined(GGML_CUDA_PALW_TRACE) + , palw_ctx, palw_operation_token +#endif + ); } break; case 7: { constexpr int c_ncols_dst = 7; @@ -980,7 +1309,11 @@ static void mul_mat_vec_q_switch_ncols_dst( mul_mat_vec_q_switch_fusion(vx, vy, ids, fusion, dst, ncols_x, nchannels_y_fd, stride_row_x, stride_col_y, stride_col_dst, channel_ratio_fd, stride_channel_x, stride_channel_y, stride_channel_dst, sample_ratio_fd, stride_sample_x, stride_sample_y, stride_sample_dst, - dims.first, dims.second, 0, ids_stride, stream); + dims.first, dims.second, 0, ids_stride, stream +#if defined(GGML_CUDA_PALW_TRACE) + , palw_ctx, palw_operation_token +#endif + ); } break; case 8: { constexpr int c_ncols_dst = 8; @@ -988,7 +1321,11 @@ static void mul_mat_vec_q_switch_ncols_dst( mul_mat_vec_q_switch_fusion(vx, vy, ids, fusion, dst, ncols_x, nchannels_y_fd, stride_row_x, stride_col_y, stride_col_dst, channel_ratio_fd, stride_channel_x, stride_channel_y, stride_channel_dst, sample_ratio_fd, stride_sample_x, stride_sample_y, stride_sample_dst, - dims.first, dims.second, 0, ids_stride, stream); + dims.first, dims.second, 0, ids_stride, stream +#if defined(GGML_CUDA_PALW_TRACE) + , palw_ctx, palw_operation_token +#endif + ); } break; default: GGML_ABORT("fatal error"); @@ -1002,7 +1339,26 @@ static void mul_mat_vec_q_switch_type( const int nchannels_x, const int nchannels_y, const int nchannels_dst, const int stride_channel_x, const int stride_channel_y, const int stride_channel_dst, const int nsamples_x, const int nsamples_dst, const int stride_sample_x, const int stride_sample_y, const int stride_sample_dst, - const int ids_stride, cudaStream_t stream) { + const int ids_stride, cudaStream_t stream +#if defined(GGML_CUDA_PALW_TRACE) + , ggml_backend_cuda_context * palw_ctx = nullptr, + const void * palw_operation_token = nullptr +#endif + ) { +#if defined(GGML_CUDA_PALW_TRACE) + if (ggml_cuda_palw_trace_active(palw_ctx)) { + if (ggml_cuda_palw_trace_failed(palw_ctx)) { + return; + } + if (type_x != GGML_TYPE_Q4_K && type_x != GGML_TYPE_Q6_K) { + ggml_cuda_palw_trace_fail( + palw_ctx, + PALW_CUDA_TRACE_INVALID_ARGUMENT, + PALW_CUDA_LLAMA_MMVQ_FAULT_UNSUPPORTED_DISPATCH); + return; + } + } +#endif switch (type_x) { case GGML_TYPE_Q1_0: mul_mat_vec_q_switch_ncols_dst @@ -1068,7 +1424,11 @@ static void mul_mat_vec_q_switch_type( mul_mat_vec_q_switch_ncols_dst (vx, vy, ids, fusion, dst, ncols_x, nrows_x, ncols_dst, stride_row_x, stride_col_y, stride_col_dst, nchannels_x, nchannels_y, nchannels_dst, stride_channel_x, stride_channel_y, stride_channel_dst, - nsamples_x, nsamples_dst, stride_sample_x, stride_sample_y, stride_sample_dst, ids_stride, stream); + nsamples_x, nsamples_dst, stride_sample_x, stride_sample_y, stride_sample_dst, ids_stride, stream +#if defined(GGML_CUDA_PALW_TRACE) + , palw_ctx, palw_operation_token +#endif + ); break; case GGML_TYPE_Q5_K: mul_mat_vec_q_switch_ncols_dst @@ -1080,7 +1440,11 @@ static void mul_mat_vec_q_switch_type( mul_mat_vec_q_switch_ncols_dst (vx, vy, ids, fusion, dst, ncols_x, nrows_x, ncols_dst, stride_row_x, stride_col_y, stride_col_dst, nchannels_x, nchannels_y, nchannels_dst, stride_channel_x, stride_channel_y, stride_channel_dst, - nsamples_x, nsamples_dst, stride_sample_x, stride_sample_y, stride_sample_dst, ids_stride, stream); + nsamples_x, nsamples_dst, stride_sample_x, stride_sample_y, stride_sample_dst, ids_stride, stream +#if defined(GGML_CUDA_PALW_TRACE) + , palw_ctx, palw_operation_token +#endif + ); break; case GGML_TYPE_IQ2_XXS: mul_mat_vec_q_switch_ncols_dst @@ -1149,6 +1513,28 @@ void ggml_cuda_mul_mat_vec_q( GGML_ASSERT( dst->type == GGML_TYPE_F32); GGML_ASSERT(!ids || ids->type == GGML_TYPE_I32); // Optional, used for batched GGML_MUL_MAT_ID. +#if defined(GGML_CUDA_PALW_TRACE) + if (ggml_cuda_palw_trace_active(&ctx) && ctx.palw_trace.request_v2 != nullptr) { + if (ggml_cuda_palw_trace_failed(&ctx)) { + return; + } + const auto & trace = ctx.palw_trace; + const auto & association = trace.association_v2; + if (!trace.association_active || + association.attention_stage != PALW_CUDA_TRACE_ATTENTION_STAGE_NONE || + !ggml_cuda_palw_tensor_view_matches(association.src0, src0) || + !ggml_cuda_palw_tensor_view_matches(association.src1, src1) || + !ggml_cuda_palw_tensor_view_matches(association.src2, ids) || + !ggml_cuda_palw_tensor_view_matches(association.dst, dst)) { + ggml_cuda_palw_trace_fail( + &ctx, + PALW_CUDA_TRACE_IDENTITY_MISMATCH, + PALW_CUDA_LLAMA_MMVQ_FAULT_ASSOCIATION_MISMATCH); + return; + } + } +#endif + GGML_TENSOR_BINARY_OP_LOCALS; cudaStream_t stream = ctx.stream(); @@ -1254,7 +1640,11 @@ void ggml_cuda_mul_mat_vec_q( src0->data, src0->type, src1_q8_1.get(), ids_d, fusion_local, dst_d, ne00, ne01, ncols_dst, s01, stride_col_y, stride_col_dst, ne02, nchannels_y, nchannels_dst, s02, stride_channel_y, stride_channel_dst, - ne03, ne3, s03, s13, s3, ids_stride, stream); + ne03, ne3, s03, s13, s3, ids_stride, stream +#if defined(GGML_CUDA_PALW_TRACE) + , &ctx, dst +#endif + ); } void ggml_cuda_op_mul_mat_vec_q( @@ -1283,7 +1673,11 @@ void ggml_cuda_op_mul_mat_vec_q( ggml_cuda_mm_fusion_args_device fusion_local{}; mul_mat_vec_q_switch_type( src0_dd_i, src0->type, src1_ddq_i, nullptr, fusion_local, dst_dd_i, ne00, row_diff, src1_ncols, stride_row_x, stride_col_y, nrows_dst, - 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, stream); + 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, stream +#if defined(GGML_CUDA_PALW_TRACE) + , &ctx, dst +#endif + ); GGML_UNUSED_VARS(src1, dst, src1_ddf_i, src1_ncols, src1_padded_row_size); } diff --git a/ggml/src/ggml-cuda/mmvq.cuh b/ggml/src/ggml-cuda/mmvq.cuh index 5605bf7a..80fc7a91 100644 --- a/ggml/src/ggml-cuda/mmvq.cuh +++ b/ggml/src/ggml-cuda/mmvq.cuh @@ -16,3 +16,12 @@ void ggml_cuda_op_mul_mat_vec_q( const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst, const char * src0_dd_i, const float * src1_ddf_i, const char * src1_ddq_i, float * dst_dd_i, const int64_t row_low, const int64_t row_high, const int64_t src1_ncols, const int64_t src1_padded_row_size, cudaStream_t stream); + +#if defined(GGML_CUDA_PALW_TRACE) +// Diagnostic-only query for the exact specialization later passed to the V3 +// launch wrapper. The release binary digest remains a caller/manifest input. +palw_cuda_trace_status ggml_cuda_palw_mmvq_kernel_info_v1( + int device, + uint8_t quantization, + palw_cuda_llama_mmvq_kernel_info_v1 * out_info); +#endif diff --git a/ggml/src/ggml-cuda/softmax.cu b/ggml/src/ggml-cuda/softmax.cu index 285c0e95..1b9ad656 100644 --- a/ggml/src/ggml-cuda/softmax.cu +++ b/ggml/src/ggml-cuda/softmax.cu @@ -436,6 +436,62 @@ void ggml_cuda_op_soft_max(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { params.m0 = m0; params.m1 = m1; +#if defined(GGML_CUDA_PALW_TRACE) + if (ggml_cuda_palw_trace_v2_attention_active(&ctx)) { + const bool preflight = ggml_cuda_palw_trace_attention_preflight( + &ctx, + src0, + src1, + src2, + dst, + PALW_CUDA_TRACE_ATTENTION_STAGE_EAGER_MASKED_SCALED_SOFTMAX, + PALW_CUDA_LLAMA_WORK_PRECISION_F32, + PALW_CUDA_LLAMA_ATTENTION_MASK_F32_CAUSAL); + uint32_t scale_bits = 0; + uint32_t max_bias_bits = 0; + memcpy(&scale_bits, &scale, sizeof(scale_bits)); + memcpy(&max_bias_bits, &max_bias, sizeof(max_bias_bits)); + const bool exact = preflight && !use_f16 && src1 != nullptr && src2 == nullptr && + params.ncols == 256 && params.ne01 == 1 && params.ne02 == 32 && + params.ne03 == 1 && scale_bits == 0x3db504f3U && + max_bias_bits == 0U && dst->nb[0] == sizeof(float) && + dst->nb[1] == 256 * sizeof(float) && dst->nb[1] == dst->nb[2]; + if (!exact) { + if (!ctx.palw_trace.failed) { + ggml_cuda_palw_trace_fail( + &ctx, + PALW_CUDA_TRACE_IDENTITY_MISMATCH, + PALW_CUDA_LLAMA_MMVQ_FAULT_UNSUPPORTED_DISPATCH); + } + return; + } + const dim3 block_dims(256, 1, 1); + const dim3 block_nums(1, 32, 1); + constexpr size_t nbytes_shared = (256 + 32) * sizeof(float); + const size_t smpbo = ggml_cuda_info().devices[ggml_cuda_get_device()].smpbo; + CUDA_SET_SHARED_MEMORY_LIMIT((soft_max_f32), smpbo); + soft_max_f32 + <<>>( + src0_d, static_cast(src1_d), nullptr, dst_d, params); + + ggml_cuda_palw_attention_work_launch_v1 work = {}; + work.entry_point = reinterpret_cast( + soft_max_f32); + work.dimensions.grid_x = block_nums.x; + work.dimensions.grid_y = block_nums.y; + work.dimensions.grid_z = block_nums.z; + work.dimensions.block_x = block_dims.x; + work.dimensions.block_y = block_dims.y; + work.dimensions.block_z = block_dims.z; + work.dimensions.dynamic_shared_memory_bytes = nbytes_shared; + work.work_tensor_type = PALW_CUDA_LLAMA_TENSOR_TYPE_F32; + work.mask_tensor_type = PALW_CUDA_LLAMA_TENSOR_TYPE_F32; + work.output_row_stride = dst->nb[1] / sizeof(float); + ggml_cuda_palw_trace_enqueue_attention(&ctx, dst, work); + return; + } +#endif + if (use_f16) { soft_max_f32_cuda(src0_d, (const half *) src1_d, (const float *) src2_d, dst_d, params, stream, ctx); } else { @@ -443,6 +499,65 @@ void ggml_cuda_op_soft_max(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { } } +#if defined(GGML_CUDA_PALW_TRACE) +palw_cuda_trace_status ggml_cuda_palw_softmax_attention_kernel_info_v1( + int device, + const palw_cuda_llama_trace_association_v2 * actual, + palw_cuda_llama_attention_kernel_info_v1 * out_info) { + if (device < 0 || actual == nullptr || out_info == nullptr || + actual->kind != PALW_CUDA_TRACE_ATTENTION || + actual->attention_stage != + PALW_CUDA_TRACE_ATTENTION_STAGE_EAGER_MASKED_SCALED_SOFTMAX || + actual->physical_key_value_tokens != 256 || + actual->src0.type != PALW_CUDA_LLAMA_TENSOR_TYPE_F32 || + actual->src1.type != PALW_CUDA_LLAMA_TENSOR_TYPE_F32) { + return PALW_CUDA_TRACE_INVALID_ARGUMENT; + } + *out_info = {}; + const void * entry_point = reinterpret_cast( + soft_max_f32); + cudaDeviceProp properties = {}; + cudaFuncAttributes attributes = {}; + ggml_cuda_set_device(device); + cudaError_t error = cudaGetDeviceProperties(&properties, device); + if (error == cudaSuccess) { + error = cudaFuncGetAttributes(&attributes, entry_point); + } + if (error != cudaSuccess) { + (void) cudaGetLastError(); + return PALW_CUDA_TRACE_CUDA_ERROR; + } + palw_cuda_llama_attention_kernel_info_v1 result = {}; + result.abi_version = PALW_CUDA_LLAMA_ATTENTION_KERNEL_INFO_ABI_VERSION_V1; + result.struct_size = sizeof(result); + result.attention_stage = actual->attention_stage; + result.work_tensor_type = PALW_CUDA_LLAMA_TENSOR_TYPE_F32; + result.mask_tensor_type = PALW_CUDA_LLAMA_TENSOR_TYPE_F32; + result.actual_identity = ggml_cuda_palw_attention_actual_identity_v3( + actual->attention_stage); + result.dimensions.grid_x = 1; + result.dimensions.grid_y = 32; + result.dimensions.grid_z = 1; + result.dimensions.block_x = 256; + result.dimensions.block_y = 1; + result.dimensions.block_z = 1; + result.dimensions.dynamic_shared_memory_bytes = (256 + 32) * sizeof(float); + result.sm_arch = static_cast(properties.major * 10 + properties.minor); + result.binary_version = attributes.binaryVersion; + result.ptx_version = attributes.ptxVersion; + result.num_regs = attributes.numRegs; + result.max_threads_per_block = attributes.maxThreadsPerBlock; + result.static_shared_memory_bytes = static_cast(attributes.sharedSizeBytes); + result.local_memory_bytes = static_cast(attributes.localSizeBytes); + result.output_row_stride = actual->dst.nb[1] / sizeof(float); + if (!palw_cuda_llama_attention_kernel_info_is_compatible_v1(&result)) { + return PALW_CUDA_TRACE_IDENTITY_MISMATCH; + } + *out_info = result; + return PALW_CUDA_TRACE_OK; +} +#endif + void ggml_cuda_op_soft_max_back(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { const ggml_tensor * src0 = dst->src[0]; // grad const ggml_tensor * src1 = dst->src[1]; // forward pass output diff --git a/ggml/src/ggml-cuda/softmax.cuh b/ggml/src/ggml-cuda/softmax.cuh index 93dfee83..997a0ffb 100644 --- a/ggml/src/ggml-cuda/softmax.cuh +++ b/ggml/src/ggml-cuda/softmax.cuh @@ -5,3 +5,10 @@ void ggml_cuda_op_soft_max(ggml_backend_cuda_context & ctx, ggml_tensor * dst); void ggml_cuda_op_soft_max_back(ggml_backend_cuda_context & ctx, ggml_tensor * dst); + +#if defined(GGML_CUDA_PALW_TRACE) +palw_cuda_trace_status ggml_cuda_palw_softmax_attention_kernel_info_v1( + int device, + const palw_cuda_llama_trace_association_v2 * actual, + palw_cuda_llama_attention_kernel_info_v1 * out_info); +#endif diff --git a/include/llama.h b/include/llama.h index a311ac20..a7a36ec2 100644 --- a/include/llama.h +++ b/include/llama.h @@ -7,6 +7,10 @@ #include "ggml-opt.h" #include "gguf.h" +#if defined(GGML_CUDA_PALW_TRACE) +#include "palw_cuda_llama_mmvq.h" +#endif + #include #include #include @@ -396,6 +400,7 @@ extern "C" { // a source/target/parent context // can be utilized in various ways, for example by sharing results or llama_memory between 2 contexts struct llama_context * ctx_other; + }; struct llama_model_tensor_override { @@ -527,6 +532,70 @@ extern "C" { // Frees all allocated memory LLAMA_API void llama_free(struct llama_context * ctx); +#if defined(GGML_CUDA_PALW_TRACE) + // Creates the V3 producer on this context's exact CUDA backend and binds a + // request before its first decode. request->producer_context must be NULL. + // The request value is copied; callback state must live through unbind. + // The attached producer is context-owned until destroy or context release. + // The context must have been created without a cb_eval callback. + LLAMA_API palw_cuda_trace_status llama_palw_cuda_trace_attach( + struct llama_context * ctx, + uint64_t record_capacity, + const palw_cuda_llama_mmvq_request_v1 * request); + + // Additive mixed-stream binding for the 253 MMVQ launches and the 36 + // FA-off QK/softmax/PV triples. V1 remains unchanged and independently + // selectable through llama_palw_cuda_trace_attach. + LLAMA_API palw_cuda_trace_status llama_palw_cuda_trace_attach_v2( + struct llama_context * ctx, + uint64_t record_capacity, + const palw_cuda_llama_trace_request_v2 * request); + + // Queries cudaFuncGetAttributes and compile-time producer IDs for the exact + // Q4_K/Q6_K specialization in the bound backend. This is diagnostic + // self-description, not release-manifest or authority evidence. + LLAMA_API palw_cuda_trace_status llama_palw_cuda_trace_kernel_info( + struct llama_context * ctx, + uint8_t quantization, + palw_cuda_llama_mmvq_kernel_info_v1 * out_info); + + LLAMA_API palw_cuda_trace_status llama_palw_cuda_trace_attention_kernel_info( + struct llama_context * ctx, + const palw_cuda_llama_trace_association_v2 * actual, + palw_cuda_llama_attention_kernel_info_v1 * out_info); + + LLAMA_API palw_cuda_trace_status llama_palw_cuda_trace_grouped_capture_info( + struct llama_context * ctx, + palw_cuda_producer_grouped_capture_info_v3 * out_info); + + // Ends the request-local CUDA trace binding. This must be called before + // finalizing the producer trace; accepted_launch_count may not be NULL. + LLAMA_API palw_cuda_trace_status llama_palw_cuda_trace_unbind( + struct llama_context * ctx, + uint64_t * accepted_launch_count); + + // Finalizes the attached producer context on the exact + // CUDA backend stream used by this llama context. Call only after unbind. + LLAMA_API palw_cuda_trace_status llama_palw_cuda_trace_finalize( + struct llama_context * ctx, + palw_cuda_trace_record_v3 * host_records, + uint64_t host_capacity); + + // Returns producer counters without synchronizing or changing ownership. + LLAMA_API palw_cuda_trace_status llama_palw_cuda_trace_diagnostics( + struct llama_context * ctx, + uint64_t * count, + uint64_t * committed_count, + uint32_t * faults, + int * poisoned); + + // Releases a producer allocated by llama_palw_cuda_trace_attach. It is + // valid after unbind on both success and failure; finalized data remains in + // caller-owned host memory. + LLAMA_API palw_cuda_trace_status llama_palw_cuda_trace_destroy( + struct llama_context * ctx); +#endif + LLAMA_API int64_t llama_time_us(void); LLAMA_API size_t llama_max_devices(void); diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index 320784c3..42a11f9c 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -55,6 +55,13 @@ target_compile_features (llama PRIVATE cxx_std_17) # don't bump target_link_libraries(llama PUBLIC ggml) +if (GGML_CUDA_PALW_TRACE) + target_compile_definitions(llama PUBLIC + GGML_CUDA_PALW_TRACE=1 + PALW_CUDA_TRACE_NO_CUDA_HEADERS=1) + target_include_directories(llama PUBLIC "${PALW_CUDA_RUNTIME_DIR}") +endif() + if (BUILD_SHARED_LIBS) set_target_properties(llama PROPERTIES POSITION_INDEPENDENT_CODE ON) target_compile_definitions(llama PRIVATE LLAMA_BUILD) diff --git a/src/llama-context.cpp b/src/llama-context.cpp index 3a469bc9..4e2df4b8 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -19,6 +19,14 @@ #include #include +#if defined(GGML_CUDA_PALW_TRACE) +static_assert(PALW_CUDA_TRACE_PRODUCTION_CAPABLE == 0U, "PALW CUDA trace must remain non-production"); +static_assert(PALW_CUDA_PRODUCER_PRODUCTION_CAPABLE == 0U, "PALW CUDA producer must remain non-production"); +static_assert(PALW_CUDA_PRODUCER_RECEIPT_MAPPING_AVAILABLE == 0U, "PALW receipt mapping must remain disabled"); +static_assert(PALW_CUDA_LLAMA_MMVQ_PRODUCTION_CAPABLE == 0U, "PALW llama bridge must remain non-production"); +static_assert(PALW_CUDA_LLAMA_MMVQ_RECEIPT_MAPPING_AVAILABLE == 0U, "PALW llama receipt mapping must remain disabled"); +#endif + // // llama_context // @@ -119,6 +127,10 @@ llama_context::llama_context( cparams.cb_eval = params.cb_eval; cparams.cb_eval_user_data = params.cb_eval_user_data; +#if defined(GGML_CUDA_PALW_TRACE) + cparams.palw_cuda_trace_request = nullptr; +#endif + cparams.ctx_other = nullptr; // TODO: more generic @@ -448,9 +460,34 @@ llama_context::llama_context( sampling.token_ids_full_vocab[i] = i; } } + } llama_context::~llama_context() { +#if defined(GGML_CUDA_PALW_TRACE) + if (palw_cuda_trace_bound) { + uint64_t accepted_launch_count = 0; + const palw_cuda_trace_status status = palw_cuda_trace_unbind(&accepted_launch_count); + if (status != PALW_CUDA_TRACE_OK) { + LLAMA_LOG_ERROR( + "%s: PALW CUDA trace unbind failed with status %d after %" PRIu64 " accepted launches\n", + __func__, static_cast(status), accepted_launch_count); + } + } + if (palw_cuda_trace_producer_owned && palw_cuda_trace_backend_api != nullptr && + palw_cuda_trace_backend != nullptr && palw_cuda_trace_producer_context != nullptr) { + palw_cuda_trace_backend_api->producer_destroy( + palw_cuda_trace_backend, palw_cuda_trace_producer_context); + palw_cuda_trace_producer_context = nullptr; + palw_cuda_trace_owned_request = {}; + palw_cuda_trace_owned_request_v2 = {}; + palw_cuda_trace_attention_metadata.clear(); + palw_cuda_trace_producer_owned = false; + palw_cuda_trace_mixed_stream = false; + cparams.palw_cuda_trace_request = nullptr; + } +#endif + if (!model.hparams.no_alloc) { for (size_t i = 0; i < backend_ptrs.size(); ++i) { ggml_backend_t backend = backend_ptrs[i]; @@ -470,6 +507,786 @@ llama_context::~llama_context() { ggml_opt_free(opt_ctx); } +#if defined(GGML_CUDA_PALW_TRACE) +void llama_context::palw_cuda_trace_record_status(palw_cuda_trace_status status) { + if (palw_cuda_trace_status_ == PALW_CUDA_TRACE_OK && status != PALW_CUDA_TRACE_OK) { + palw_cuda_trace_status_ = status; + } +} + +palw_cuda_trace_status llama_context::palw_cuda_trace_attach( + uint64_t record_capacity, + const palw_cuda_llama_mmvq_request_v1 * request) { + const auto fail_unbound = [request](palw_cuda_trace_status status) { + if (request != nullptr && request->fault != nullptr) { + request->fault( + request->user_data, + status, + PALW_CUDA_LLAMA_MMVQ_FAULT_INVALID_REQUEST, + nullptr); + } + return status; + }; + + if (request == nullptr || request->abi_version != PALW_CUDA_LLAMA_MMVQ_REQUEST_ABI_VERSION_V1 || + request->struct_size != sizeof(*request) || request->expected_launch_count == 0 || + record_capacity < request->expected_launch_count || request->producer_context != nullptr || + request->approve == nullptr || request->fault == nullptr || + cparams.palw_cuda_trace_request != nullptr || palw_cuda_trace_bound || + palw_cuda_trace_producer_owned || has_evaluated_once || n_queued_tokens != 0 || + cparams.cb_eval != nullptr || cparams.pipeline_parallel || + cparams.flash_attn || cparams.n_batch != 1 || cparams.n_ubatch != 1 || + cparams.n_seq_max != 1) { + return fail_unbound(PALW_CUDA_TRACE_INVALID_ARGUMENT); + } + + size_t candidate_count = 0; + ggml_backend_t candidate_backend = nullptr; + palw_cuda_llama_mmvq_get_backend_api_v1 get_api_fn = nullptr; + for (const auto & backend : backends) { + ggml_backend_dev_t device = ggml_backend_get_device(backend.get()); + ggml_backend_reg_t registry = device != nullptr ? ggml_backend_dev_backend_reg(device) : nullptr; + if (registry == nullptr) { + continue; + } + auto bind_fn = reinterpret_cast( + ggml_backend_reg_get_proc_address(registry, PALW_CUDA_LLAMA_MMVQ_BIND_PROC_V1)); + if (bind_fn == nullptr) { + continue; + } + ++candidate_count; + if (candidate_count == 1) { + candidate_backend = backend.get(); + get_api_fn = reinterpret_cast( + ggml_backend_reg_get_proc_address(registry, PALW_CUDA_LLAMA_MMVQ_GET_API_PROC_V1)); + } + } + if (candidate_count != 1 || candidate_backend == nullptr || get_api_fn == nullptr) { + return fail_unbound(PALW_CUDA_TRACE_INVALID_ARGUMENT); + } + + const palw_cuda_llama_mmvq_backend_api_v1 * backend_api = get_api_fn(); + if (!palw_cuda_llama_mmvq_backend_api_is_compatible_v1(backend_api)) { + return fail_unbound(PALW_CUDA_TRACE_INVALID_ARGUMENT); + } + + palw_cuda_producer_trace_context_v3 * producer_context = nullptr; + palw_cuda_trace_status status = backend_api->producer_create( + candidate_backend, record_capacity, &producer_context); + if (status != PALW_CUDA_TRACE_OK || producer_context == nullptr) { + return fail_unbound(status != PALW_CUDA_TRACE_OK + ? status + : PALW_CUDA_TRACE_INVALID_ARGUMENT); + } + + palw_cuda_trace_owned_request = *request; + palw_cuda_trace_owned_request.producer_context = producer_context; + cparams.palw_cuda_trace_request = &palw_cuda_trace_owned_request; + palw_cuda_trace_status_ = PALW_CUDA_TRACE_OK; + palw_cuda_trace_finalized = false; + palw_cuda_trace_producer_owned = true; + palw_cuda_trace_mixed_stream = false; + try { + palw_cuda_trace_bind(); + } catch (const std::exception &) { + const bool report_fault = palw_cuda_trace_status_ == PALW_CUDA_TRACE_OK; + status = palw_cuda_trace_status_ != PALW_CUDA_TRACE_OK + ? palw_cuda_trace_status_ + : PALW_CUDA_TRACE_INVALID_ARGUMENT; + backend_api->producer_destroy(candidate_backend, producer_context); + cparams.palw_cuda_trace_request = nullptr; + palw_cuda_trace_backend = nullptr; + palw_cuda_trace_backend_api = nullptr; + palw_cuda_trace_associate_fn = nullptr; + palw_cuda_trace_complete_fn = nullptr; + palw_cuda_trace_unbind_fn = nullptr; + palw_cuda_trace_kernel_info_fn = nullptr; + palw_cuda_trace_producer_context = nullptr; + palw_cuda_trace_owned_request = {}; + palw_cuda_trace_pending_operation = nullptr; + palw_cuda_trace_status_ = PALW_CUDA_TRACE_OK; + palw_cuda_trace_bound = false; + palw_cuda_trace_finalized = false; + palw_cuda_trace_producer_owned = false; + if (report_fault) { + (void) fail_unbound(status); + } + return status; + } + return PALW_CUDA_TRACE_OK; +} + +palw_cuda_trace_status llama_context::palw_cuda_trace_attach_v2( + uint64_t record_capacity, + const palw_cuda_llama_trace_request_v2 * request) { + const auto fail_unbound = [request](palw_cuda_trace_status status) { + if (request != nullptr && request->fault != nullptr) { + request->fault( + request->user_data, + status, + PALW_CUDA_LLAMA_MMVQ_FAULT_INVALID_REQUEST, + nullptr); + } + return status; + }; + // This public boundary accepts the unbound request. The stricter + // `is_compatible_v2` helper is reserved for the owned copy after its + // producer context has been created below. + if (request == nullptr || + request->abi_version != PALW_CUDA_LLAMA_TRACE_REQUEST_ABI_VERSION_V2 || + request->struct_size != sizeof(*request) || + request->expected_launch_count != 361 || + record_capacity < request->expected_launch_count || + request->producer_context != nullptr || request->approve == nullptr || + request->fault == nullptr || palw_cuda_trace_bound || + palw_cuda_trace_producer_owned || has_evaluated_once || n_queued_tokens != 0 || + palw_cuda_state_imported || + cparams.palw_cuda_trace_request != nullptr || cparams.cb_eval != nullptr || + cparams.pipeline_parallel || cparams.flash_attn || cparams.n_batch != 1 || + cparams.n_ubatch != 1 || cparams.n_seq_max != 1 || + model.arch != LLM_ARCH_QWEN3 || model.hparams.n_layer() != 36 || + model.hparams.n_head() != 32 || model.hparams.n_head_kv() != 8 || + model.hparams.n_embd_head_k() != 128 || + model.hparams.n_embd_head_v() != 128) { + return fail_unbound(PALW_CUDA_TRACE_INVALID_ARGUMENT); + } + + size_t candidate_count = 0; + ggml_backend_t candidate_backend = nullptr; + palw_cuda_llama_mmvq_get_backend_api_v1 get_api_fn = nullptr; + for (const auto & backend : backends) { + ggml_backend_dev_t device = ggml_backend_get_device(backend.get()); + ggml_backend_reg_t registry = device != nullptr + ? ggml_backend_dev_backend_reg(device) + : nullptr; + if (registry == nullptr || ggml_backend_reg_get_proc_address( + registry, PALW_CUDA_LLAMA_TRACE_BIND_PROC_V2) == nullptr) { + continue; + } + ++candidate_count; + if (candidate_count == 1) { + candidate_backend = backend.get(); + get_api_fn = reinterpret_cast( + ggml_backend_reg_get_proc_address( + registry, PALW_CUDA_LLAMA_MMVQ_GET_API_PROC_V1)); + } + } + if (candidate_count != 1 || candidate_backend == nullptr || get_api_fn == nullptr) { + return fail_unbound(PALW_CUDA_TRACE_INVALID_ARGUMENT); + } + const palw_cuda_llama_mmvq_backend_api_v1 * backend_api = get_api_fn(); + if (!palw_cuda_llama_mmvq_backend_api_is_compatible_v1(backend_api)) { + return fail_unbound(PALW_CUDA_TRACE_INVALID_ARGUMENT); + } + palw_cuda_producer_trace_context_v3 * producer_context = nullptr; + palw_cuda_trace_status status = backend_api->producer_create( + candidate_backend, record_capacity, &producer_context); + if (status != PALW_CUDA_TRACE_OK || producer_context == nullptr) { + return fail_unbound(status != PALW_CUDA_TRACE_OK + ? status + : PALW_CUDA_TRACE_INVALID_ARGUMENT); + } + + palw_cuda_trace_owned_request_v2 = *request; + palw_cuda_trace_owned_request_v2.producer_context = producer_context; + palw_cuda_trace_status_ = PALW_CUDA_TRACE_OK; + palw_cuda_trace_finalized = false; + palw_cuda_trace_producer_owned = true; + palw_cuda_trace_mixed_stream = true; + try { + palw_cuda_trace_bind(); + } catch (const std::exception &) { + const bool report_fault = palw_cuda_trace_status_ == PALW_CUDA_TRACE_OK; + status = palw_cuda_trace_status_ != PALW_CUDA_TRACE_OK + ? palw_cuda_trace_status_ + : PALW_CUDA_TRACE_INVALID_ARGUMENT; + backend_api->producer_destroy(candidate_backend, producer_context); + palw_cuda_trace_backend = nullptr; + palw_cuda_trace_backend_api = nullptr; + palw_cuda_trace_associate_fn_v2 = nullptr; + palw_cuda_trace_complete_fn_v2 = nullptr; + palw_cuda_trace_unbind_fn_v2 = nullptr; + palw_cuda_trace_attention_kernel_info_fn = nullptr; + palw_cuda_trace_grouped_capture_info_fn = nullptr; + palw_cuda_trace_producer_context = nullptr; + palw_cuda_trace_owned_request_v2 = {}; + palw_cuda_trace_pending_operation = nullptr; + palw_cuda_trace_status_ = PALW_CUDA_TRACE_OK; + palw_cuda_trace_bound = false; + palw_cuda_trace_finalized = false; + palw_cuda_trace_producer_owned = false; + palw_cuda_trace_mixed_stream = false; + if (report_fault) { + (void) fail_unbound(status); + } + return status; + } + return PALW_CUDA_TRACE_OK; +} + +void llama_context::palw_cuda_trace_bind() { + if (palw_cuda_trace_mixed_stream) { + const palw_cuda_llama_trace_request_v2 * request = + &palw_cuda_trace_owned_request_v2; + if (!palw_cuda_llama_trace_request_is_compatible_v2(request)) { + throw std::runtime_error("invalid PALW CUDA mixed trace request"); + } + size_t candidate_count = 0; + ggml_backend_t candidate_backend = nullptr; + palw_cuda_llama_mmvq_get_backend_api_v1 get_api_fn = nullptr; + palw_cuda_llama_trace_bind_backend_v2 bind_fn = nullptr; + palw_cuda_llama_trace_associate_backend_v2 associate_fn = nullptr; + palw_cuda_llama_trace_complete_backend_v2 complete_fn = nullptr; + palw_cuda_llama_trace_unbind_backend_v2 unbind_fn = nullptr; + palw_cuda_llama_attention_kernel_info_backend_v1 kernel_info_fn = nullptr; + palw_cuda_llama_mmvq_kernel_info_backend_v1 mmvq_kernel_info_fn = nullptr; + palw_cuda_llama_grouped_capture_info_backend_v1 capture_info_fn = nullptr; + for (const auto & backend : backends) { + ggml_backend_dev_t device = ggml_backend_get_device(backend.get()); + ggml_backend_reg_t registry = device != nullptr + ? ggml_backend_dev_backend_reg(device) + : nullptr; + if (registry == nullptr) { + continue; + } + auto candidate_bind_fn = + reinterpret_cast( + ggml_backend_reg_get_proc_address( + registry, PALW_CUDA_LLAMA_TRACE_BIND_PROC_V2)); + if (candidate_bind_fn == nullptr) { + continue; + } + ++candidate_count; + if (candidate_count == 1) { + candidate_backend = backend.get(); + bind_fn = candidate_bind_fn; + get_api_fn = reinterpret_cast( + ggml_backend_reg_get_proc_address( + registry, PALW_CUDA_LLAMA_MMVQ_GET_API_PROC_V1)); + associate_fn = reinterpret_cast( + ggml_backend_reg_get_proc_address( + registry, PALW_CUDA_LLAMA_TRACE_ASSOCIATE_PROC_V2)); + complete_fn = reinterpret_cast( + ggml_backend_reg_get_proc_address( + registry, PALW_CUDA_LLAMA_TRACE_COMPLETE_PROC_V2)); + unbind_fn = reinterpret_cast( + ggml_backend_reg_get_proc_address( + registry, PALW_CUDA_LLAMA_TRACE_UNBIND_PROC_V2)); + kernel_info_fn = + reinterpret_cast( + ggml_backend_reg_get_proc_address( + registry, PALW_CUDA_LLAMA_ATTENTION_KERNEL_INFO_PROC_V1)); + mmvq_kernel_info_fn = + reinterpret_cast( + ggml_backend_reg_get_proc_address( + registry, PALW_CUDA_LLAMA_MMVQ_KERNEL_INFO_PROC_V1)); + capture_info_fn = + reinterpret_cast( + ggml_backend_reg_get_proc_address( + registry, PALW_CUDA_LLAMA_GROUPED_CAPTURE_INFO_PROC_V1)); + } + } + if (candidate_count != 1 || get_api_fn == nullptr || bind_fn == nullptr || + associate_fn == nullptr || complete_fn == nullptr || + unbind_fn == nullptr || kernel_info_fn == nullptr || + mmvq_kernel_info_fn == nullptr || capture_info_fn == nullptr) { + throw std::runtime_error("PALW CUDA mixed backend bridge is incomplete"); + } + const auto * backend_api = get_api_fn(); + if (!palw_cuda_llama_mmvq_backend_api_is_compatible_v1(backend_api)) { + throw std::runtime_error("incompatible PALW CUDA mixed backend API"); + } + const palw_cuda_trace_status status = bind_fn(candidate_backend, request); + palw_cuda_trace_record_status(status); + if (status != PALW_CUDA_TRACE_OK) { + throw std::runtime_error("failed to bind PALW CUDA mixed trace request"); + } + palw_cuda_trace_backend = candidate_backend; + palw_cuda_trace_backend_api = backend_api; + palw_cuda_trace_associate_fn_v2 = associate_fn; + palw_cuda_trace_complete_fn_v2 = complete_fn; + palw_cuda_trace_unbind_fn_v2 = unbind_fn; + palw_cuda_trace_attention_kernel_info_fn = kernel_info_fn; + palw_cuda_trace_kernel_info_fn = mmvq_kernel_info_fn; + palw_cuda_trace_grouped_capture_info_fn = capture_info_fn; + palw_cuda_trace_producer_context = request->producer_context; + palw_cuda_trace_pending_operation = nullptr; + palw_cuda_trace_bound = true; + palw_cuda_trace_finalized = false; + return; + } + const palw_cuda_llama_mmvq_request_v1 * request = cparams.palw_cuda_trace_request; + if (!palw_cuda_llama_mmvq_request_is_compatible_v1(request)) { + throw std::runtime_error("invalid PALW CUDA MMVQ trace request"); + } + if (cparams.pipeline_parallel) { + throw std::runtime_error("PALW CUDA MMVQ tracing does not support pipeline parallelism"); + } + + size_t candidate_count = 0; + ggml_backend_t candidate_backend = nullptr; + palw_cuda_llama_mmvq_get_backend_api_v1 get_api_fn = nullptr; + palw_cuda_llama_mmvq_bind_backend_v1 bind_fn = nullptr; + palw_cuda_llama_mmvq_associate_backend_v1 associate_fn = nullptr; + palw_cuda_llama_mmvq_complete_backend_v1 complete_fn = nullptr; + palw_cuda_llama_mmvq_unbind_backend_v1 unbind_fn = nullptr; + palw_cuda_llama_mmvq_kernel_info_backend_v1 kernel_info_fn = nullptr; + + for (const auto & backend : backends) { + ggml_backend_dev_t device = ggml_backend_get_device(backend.get()); + ggml_backend_reg_t registry = device != nullptr ? ggml_backend_dev_backend_reg(device) : nullptr; + if (registry == nullptr) { + continue; + } + + auto candidate_bind_fn = reinterpret_cast( + ggml_backend_reg_get_proc_address(registry, PALW_CUDA_LLAMA_MMVQ_BIND_PROC_V1)); + if (candidate_bind_fn == nullptr) { + continue; + } + + ++candidate_count; + if (candidate_count == 1) { + candidate_backend = backend.get(); + bind_fn = candidate_bind_fn; + get_api_fn = reinterpret_cast( + ggml_backend_reg_get_proc_address(registry, PALW_CUDA_LLAMA_MMVQ_GET_API_PROC_V1)); + associate_fn = reinterpret_cast( + ggml_backend_reg_get_proc_address(registry, PALW_CUDA_LLAMA_MMVQ_ASSOCIATE_PROC_V1)); + complete_fn = reinterpret_cast( + ggml_backend_reg_get_proc_address(registry, PALW_CUDA_LLAMA_MMVQ_COMPLETE_PROC_V1)); + unbind_fn = reinterpret_cast( + ggml_backend_reg_get_proc_address(registry, PALW_CUDA_LLAMA_MMVQ_UNBIND_PROC_V1)); + kernel_info_fn = reinterpret_cast( + ggml_backend_reg_get_proc_address(registry, PALW_CUDA_LLAMA_MMVQ_KERNEL_INFO_PROC_V1)); + } + } + + if (candidate_count != 1) { + throw std::runtime_error( + "PALW CUDA MMVQ tracing requires exactly one CUDA backend, found " + + std::to_string(candidate_count)); + } + if (get_api_fn == nullptr || bind_fn == nullptr || associate_fn == nullptr || + complete_fn == nullptr || unbind_fn == nullptr || kernel_info_fn == nullptr) { + throw std::runtime_error("PALW CUDA MMVQ backend bridge is incomplete"); + } + + const palw_cuda_llama_mmvq_backend_api_v1 * backend_api = get_api_fn(); + if (!palw_cuda_llama_mmvq_backend_api_is_compatible_v1(backend_api)) { + throw std::runtime_error("incompatible PALW CUDA MMVQ backend API"); + } + + const palw_cuda_trace_status status = bind_fn(candidate_backend, request); + palw_cuda_trace_record_status(status); + if (status != PALW_CUDA_TRACE_OK) { + throw std::runtime_error( + "failed to bind PALW CUDA MMVQ trace request, status " + + std::to_string(static_cast(status))); + } + + palw_cuda_trace_backend = candidate_backend; + palw_cuda_trace_backend_api = backend_api; + palw_cuda_trace_associate_fn = associate_fn; + palw_cuda_trace_complete_fn = complete_fn; + palw_cuda_trace_unbind_fn = unbind_fn; + palw_cuda_trace_kernel_info_fn = kernel_info_fn; + palw_cuda_trace_producer_context = request->producer_context; + palw_cuda_trace_pending_operation = nullptr; + palw_cuda_trace_bound = true; + palw_cuda_trace_finalized = false; +} + +palw_cuda_trace_status llama_context::palw_cuda_trace_kernel_info( + uint8_t quantization, + palw_cuda_llama_mmvq_kernel_info_v1 * out_info) { + if (!palw_cuda_trace_bound || palw_cuda_trace_backend == nullptr || + palw_cuda_trace_kernel_info_fn == nullptr || out_info == nullptr) { + return PALW_CUDA_TRACE_INVALID_ARGUMENT; + } + const palw_cuda_trace_status status = palw_cuda_trace_kernel_info_fn( + palw_cuda_trace_backend, quantization, out_info); + if (status != PALW_CUDA_TRACE_OK || + !palw_cuda_llama_mmvq_kernel_info_is_compatible_v1(out_info)) { + return status != PALW_CUDA_TRACE_OK + ? status + : PALW_CUDA_TRACE_IDENTITY_MISMATCH; + } + return PALW_CUDA_TRACE_OK; +} + +palw_cuda_trace_status llama_context::palw_cuda_trace_attention_kernel_info( + const palw_cuda_llama_trace_association_v2 * actual, + palw_cuda_llama_attention_kernel_info_v1 * out_info) { + if (!palw_cuda_trace_bound || !palw_cuda_trace_mixed_stream || + palw_cuda_trace_backend == nullptr || + palw_cuda_trace_attention_kernel_info_fn == nullptr || + actual == nullptr || out_info == nullptr) { + return PALW_CUDA_TRACE_INVALID_ARGUMENT; + } + const palw_cuda_trace_status status = + palw_cuda_trace_attention_kernel_info_fn( + palw_cuda_trace_backend, actual, out_info); + if (status != PALW_CUDA_TRACE_OK || + !palw_cuda_llama_attention_kernel_info_is_compatible_v1(out_info)) { + return status != PALW_CUDA_TRACE_OK + ? status + : PALW_CUDA_TRACE_IDENTITY_MISMATCH; + } + return PALW_CUDA_TRACE_OK; +} + +palw_cuda_trace_status llama_context::palw_cuda_trace_grouped_capture_info( + palw_cuda_producer_grouped_capture_info_v3 * out_info) { + if (!palw_cuda_trace_bound || !palw_cuda_trace_mixed_stream || + palw_cuda_trace_backend == nullptr || + palw_cuda_trace_grouped_capture_info_fn == nullptr || out_info == nullptr) { + return PALW_CUDA_TRACE_INVALID_ARGUMENT; + } + return palw_cuda_trace_grouped_capture_info_fn( + palw_cuda_trace_backend, out_info); +} + +bool llama_context::palw_cuda_trace_candidate(const ggml_tensor * tensor) const { + if (tensor == nullptr) { + return false; + } + + if (palw_cuda_trace_mixed_stream && + palw_cuda_trace_attention_metadata.find(tensor) != + palw_cuda_trace_attention_metadata.end()) { + return true; + } + if ((tensor->op != GGML_OP_MUL_MAT && tensor->op != GGML_OP_MUL_MAT_ID) || + tensor->src[0] == nullptr) { + return false; + } + + // Select by operation and quantized source type only. Shape, dtype, device, + // dispatch, and expected-slot mismatches are intentionally passed to the + // bound backend so they become sticky failures instead of silent bypasses. + return ggml_is_quantized(tensor->src[0]->type); +} + +static palw_cuda_llama_tensor_view_v1 palw_cuda_trace_tensor_view( + const ggml_tensor * tensor) { + palw_cuda_llama_tensor_view_v1 view = {}; + if (tensor == nullptr) { + return view; + } + view.data = tensor->data; + switch (tensor->type) { + case GGML_TYPE_F32: + view.type = PALW_CUDA_LLAMA_TENSOR_TYPE_F32; + break; + case GGML_TYPE_F16: + view.type = PALW_CUDA_LLAMA_TENSOR_TYPE_F16; + break; + case GGML_TYPE_Q4_K: + view.type = PALW_CUDA_LLAMA_TENSOR_TYPE_Q4_K; + break; + case GGML_TYPE_Q6_K: + view.type = PALW_CUDA_LLAMA_TENSOR_TYPE_Q6_K; + break; + default: + view.type = UINT8_MAX; + break; + } + for (uint32_t i = 0; i < 4; ++i) { + view.ne[i] = tensor->ne[i] > 0 ? (uint64_t) tensor->ne[i] : 0; + view.nb[i] = (uint64_t) tensor->nb[i]; + } + return view; +} + +bool llama_context::palw_cuda_trace_eval(ggml_tensor * tensor, bool ask) { + const bool trace_need = palw_cuda_trace_bound && palw_cuda_trace_candidate(tensor); + + if (ask) { + if (trace_need) { + const auto positive_dimension = [](int64_t value) -> uint64_t { + return value > 0 ? static_cast(value) : 0U; + }; + if (palw_cuda_trace_pending_operation == nullptr) { + palw_cuda_trace_pending_operation = tensor; + } + palw_cuda_trace_status status = PALW_CUDA_TRACE_INVALID_ARGUMENT; + if (palw_cuda_trace_mixed_stream) { + palw_cuda_llama_trace_association_v2 actual = {}; + actual.abi_version = PALW_CUDA_LLAMA_TRACE_ASSOCIATION_ABI_VERSION_V2; + actual.struct_size = sizeof(actual); + actual.operation_token = tensor; + actual.src0 = palw_cuda_trace_tensor_view(tensor->src[0]); + actual.src1 = palw_cuda_trace_tensor_view(tensor->src[1]); + actual.src2 = palw_cuda_trace_tensor_view(tensor->src[2]); + actual.dst = palw_cuda_trace_tensor_view(tensor); + actual.phase = PALW_CUDA_TRACE_PREFILL; + actual.decode_step = 0; + const auto metadata_it = palw_cuda_trace_attention_metadata.find(tensor); + if (metadata_it == palw_cuda_trace_attention_metadata.end()) { + const ggml_tensor * weight = tensor->src[0]; + actual.kind = PALW_CUDA_TRACE_GEMM; + actual.m = positive_dimension(tensor->ne[0]); + actual.n = positive_dimension(tensor->ne[1]); + actual.k = positive_dimension(weight->ne[0]); + const uint64_t batch_2 = positive_dimension(tensor->ne[2]); + const uint64_t batch_3 = positive_dimension(tensor->ne[3]); + actual.batch = batch_2 != 0 && + batch_3 <= std::numeric_limits::max() / batch_2 + ? batch_2 * batch_3 + : 0U; + actual.output_rows = actual.n <= std::numeric_limits::max() + ? static_cast(actual.n) + : 0U; + actual.output_columns = actual.m <= std::numeric_limits::max() + ? static_cast(actual.m) + : 0U; + actual.quantization = weight->type == GGML_TYPE_Q4_K + ? PALW_CUDA_TRACE_QUANTIZATION_Q4_K + : weight->type == GGML_TYPE_Q6_K + ? PALW_CUDA_TRACE_QUANTIZATION_Q6_K + : PALW_CUDA_TRACE_QUANTIZATION_NONE; + // MUL_MAT_ID remains an explicit unsupported-dispatch + // failure in the backend kernel hook; IDs are not part of + // the pinned batch-one association ABI. + actual.src2 = {}; + } else { + const auto & metadata = metadata_it->second; + actual.attention_stage = metadata.attention_stage; + actual.kind = metadata.attention_stage == + PALW_CUDA_TRACE_ATTENTION_STAGE_EAGER_MASKED_SCALED_SOFTMAX + ? PALW_CUDA_TRACE_ATTENTION + : PALW_CUDA_TRACE_GEMM; + actual.layer_id = metadata.layer_id; + actual.layer_present = metadata.layer_present; + actual.query_tokens = metadata.query_tokens; + actual.key_value_tokens = metadata.key_value_tokens; + actual.physical_key_value_tokens = + metadata.physical_key_value_tokens; + actual.query_heads = metadata.query_heads; + actual.key_value_heads = metadata.key_value_heads; + actual.head_dim = metadata.head_dim; + actual.logical_batch = metadata.logical_batch; + actual.decode_step = metadata.decode_step; + actual.phase = metadata.phase; + actual.causal = metadata.causal; + actual.quantization = PALW_CUDA_TRACE_QUANTIZATION_NONE; + actual.n = metadata.query_tokens; + actual.batch = (uint64_t) metadata.query_heads * + metadata.logical_batch; + actual.output_rows = metadata.query_heads * + metadata.logical_batch; + if (metadata.attention_stage == + PALW_CUDA_TRACE_ATTENTION_STAGE_EAGER_QK_SCORE_MMVF) { + actual.m = positive_dimension(tensor->ne[0]); + actual.k = positive_dimension(tensor->src[0]->ne[0]); + actual.precision = PALW_CUDA_LLAMA_WORK_PRECISION_F32; + } else if (metadata.attention_stage == + PALW_CUDA_TRACE_ATTENTION_STAGE_EAGER_MASKED_SCALED_SOFTMAX) { + actual.m = positive_dimension(tensor->ne[0]); + actual.k = 1; + actual.precision = PALW_CUDA_LLAMA_WORK_PRECISION_F32; + actual.attention_mask = + tensor->src[1] != nullptr && + tensor->src[1]->type == GGML_TYPE_F32 + ? PALW_CUDA_LLAMA_ATTENTION_MASK_F32_CAUSAL + : PALW_CUDA_LLAMA_ATTENTION_MASK_F16_CAUSAL; + memcpy(&actual.scale_bits, tensor->op_params, sizeof(uint32_t)); + memcpy(&actual.max_bias_bits, + (const float *) tensor->op_params + 1, + sizeof(uint32_t)); + } else { + actual.m = positive_dimension(tensor->ne[0]); + actual.k = positive_dimension(tensor->src[0]->ne[0]); + actual.precision = PALW_CUDA_LLAMA_WORK_PRECISION_DEFAULT; + } + actual.output_columns = actual.m <= + std::numeric_limits::max() + ? static_cast(actual.m) + : 0U; + } + status = palw_cuda_trace_associate_fn_v2( + palw_cuda_trace_backend, &actual); + } else { + const ggml_tensor * weight = tensor->src[0]; + palw_cuda_llama_mmvq_association_v1 actual = {}; + actual.abi_version = PALW_CUDA_LLAMA_MMVQ_ASSOCIATION_ABI_VERSION_V1; + actual.struct_size = sizeof(actual); + actual.operation_token = tensor; + actual.weight_data = weight->data; + actual.m = positive_dimension(tensor->ne[0]); + actual.n = positive_dimension(tensor->ne[1]); + actual.k = positive_dimension(weight->ne[0]); + const uint64_t batch_2 = positive_dimension(tensor->ne[2]); + const uint64_t batch_3 = positive_dimension(tensor->ne[3]); + actual.batch = batch_2 != 0 && + batch_3 <= std::numeric_limits::max() / batch_2 + ? batch_2 * batch_3 + : 0U; + actual.output_rows = actual.n <= std::numeric_limits::max() + ? static_cast(actual.n) + : 0U; + actual.output_columns = actual.m <= std::numeric_limits::max() + ? static_cast(actual.m) + : 0U; + actual.quantization = weight->type == GGML_TYPE_Q4_K + ? PALW_CUDA_TRACE_QUANTIZATION_Q4_K + : weight->type == GGML_TYPE_Q6_K + ? PALW_CUDA_TRACE_QUANTIZATION_Q6_K + : PALW_CUDA_TRACE_QUANTIZATION_NONE; + status = palw_cuda_trace_associate_fn( + palw_cuda_trace_backend, &actual); + } + palw_cuda_trace_record_status(status); + } + + const bool user_need = cparams.cb_eval != nullptr + ? cparams.cb_eval(tensor, true, cparams.cb_eval_user_data) + : false; + return trace_need || user_need; + } + + bool trace_ok = true; + if (trace_need) { + // tensor is the exact graph operation token; the CUDA MMVQ host path + // passes the same destination tensor pointer to the backend hook. + const palw_cuda_trace_status status = + palw_cuda_trace_mixed_stream + ? palw_cuda_trace_complete_fn_v2(palw_cuda_trace_backend, tensor) + : palw_cuda_trace_complete_fn(palw_cuda_trace_backend, tensor); + palw_cuda_trace_record_status(status); + trace_ok = status == PALW_CUDA_TRACE_OK && palw_cuda_trace_pending_operation == tensor; + if (!trace_ok && status == PALW_CUDA_TRACE_OK) { + palw_cuda_trace_record_status(PALW_CUDA_TRACE_SEQUENCE_MISMATCH); + } + palw_cuda_trace_pending_operation = nullptr; + } + + const bool user_ok = cparams.cb_eval != nullptr + ? cparams.cb_eval(tensor, false, cparams.cb_eval_user_data) + : true; + return trace_ok && user_ok; +} + +bool llama_context::palw_cuda_trace_eval_callback( + ggml_tensor * tensor, + bool ask, + void * user_data) { + if (user_data == nullptr) { + return false; + } + return static_cast(user_data)->palw_cuda_trace_eval(tensor, ask); +} + +palw_cuda_trace_status llama_context::palw_cuda_trace_unbind(uint64_t * accepted_launch_count) { + if (accepted_launch_count == nullptr) { + return PALW_CUDA_TRACE_INVALID_ARGUMENT; + } + *accepted_launch_count = 0; + if (!palw_cuda_trace_bound || palw_cuda_trace_backend == nullptr || + (palw_cuda_trace_mixed_stream + ? palw_cuda_trace_unbind_fn_v2 == nullptr + : palw_cuda_trace_unbind_fn == nullptr || + cparams.palw_cuda_trace_request == nullptr)) { + return PALW_CUDA_TRACE_INVALID_ARGUMENT; + } + + synchronize(); + const palw_cuda_trace_status status = palw_cuda_trace_mixed_stream + ? palw_cuda_trace_unbind_fn_v2( + palw_cuda_trace_backend, + &palw_cuda_trace_owned_request_v2, + accepted_launch_count) + : palw_cuda_trace_unbind_fn( + palw_cuda_trace_backend, + cparams.palw_cuda_trace_request, + accepted_launch_count); + palw_cuda_trace_record_status(status); + palw_cuda_trace_bound = false; + palw_cuda_trace_pending_operation = nullptr; + return status == PALW_CUDA_TRACE_OK ? palw_cuda_trace_status_ : status; +} + +palw_cuda_trace_status llama_context::palw_cuda_trace_finalize( + palw_cuda_trace_record_v3 * host_records, + uint64_t host_capacity) { + if (palw_cuda_trace_bound || palw_cuda_trace_finalized || + palw_cuda_trace_status_ != PALW_CUDA_TRACE_OK || + palw_cuda_trace_backend == nullptr || palw_cuda_trace_backend_api == nullptr || + palw_cuda_trace_producer_context == nullptr || host_records == nullptr || host_capacity == 0) { + return palw_cuda_trace_status_ != PALW_CUDA_TRACE_OK + ? palw_cuda_trace_status_ + : PALW_CUDA_TRACE_INVALID_ARGUMENT; + } + + const palw_cuda_trace_status status = palw_cuda_trace_backend_api->producer_finalize( + palw_cuda_trace_backend, + palw_cuda_trace_producer_context, + host_records, + host_capacity); + palw_cuda_trace_record_status(status); + if (status == PALW_CUDA_TRACE_OK) { + palw_cuda_trace_finalized = true; + } + return status; +} + +palw_cuda_trace_status llama_context::palw_cuda_trace_diagnostics( + uint64_t * count, + uint64_t * committed_count, + uint32_t * faults, + int * poisoned) const { + if (count == nullptr || committed_count == nullptr || faults == nullptr || + poisoned == nullptr || palw_cuda_trace_backend_api == nullptr || + palw_cuda_trace_producer_context == nullptr) { + return PALW_CUDA_TRACE_INVALID_ARGUMENT; + } + *count = palw_cuda_trace_backend_api->producer_count(palw_cuda_trace_producer_context); + *committed_count = palw_cuda_trace_backend_api->producer_committed_count( + palw_cuda_trace_producer_context); + *faults = palw_cuda_trace_backend_api->producer_faults(palw_cuda_trace_producer_context); + *poisoned = palw_cuda_trace_backend_api->producer_poisoned( + palw_cuda_trace_producer_context); + return PALW_CUDA_TRACE_OK; +} + +palw_cuda_trace_status llama_context::palw_cuda_trace_destroy() { + if (palw_cuda_trace_bound) { + return PALW_CUDA_TRACE_PENDING; + } + if (!palw_cuda_trace_producer_owned || palw_cuda_trace_backend == nullptr || + palw_cuda_trace_backend_api == nullptr || + palw_cuda_trace_producer_context == nullptr) { + return PALW_CUDA_TRACE_INVALID_ARGUMENT; + } + + palw_cuda_trace_backend_api->producer_destroy( + palw_cuda_trace_backend, palw_cuda_trace_producer_context); + cparams.palw_cuda_trace_request = nullptr; + palw_cuda_trace_backend = nullptr; + palw_cuda_trace_backend_api = nullptr; + palw_cuda_trace_associate_fn = nullptr; + palw_cuda_trace_complete_fn = nullptr; + palw_cuda_trace_unbind_fn = nullptr; + palw_cuda_trace_kernel_info_fn = nullptr; + palw_cuda_trace_associate_fn_v2 = nullptr; + palw_cuda_trace_complete_fn_v2 = nullptr; + palw_cuda_trace_unbind_fn_v2 = nullptr; + palw_cuda_trace_attention_kernel_info_fn = nullptr; + palw_cuda_trace_grouped_capture_info_fn = nullptr; + palw_cuda_trace_producer_context = nullptr; + palw_cuda_trace_owned_request = {}; + palw_cuda_trace_owned_request_v2 = {}; + palw_cuda_trace_attention_metadata.clear(); + palw_cuda_trace_pending_operation = nullptr; + palw_cuda_trace_status_ = PALW_CUDA_TRACE_OK; + palw_cuda_trace_finalized = false; + palw_cuda_trace_producer_owned = false; + palw_cuda_trace_mixed_stream = false; + return PALW_CUDA_TRACE_OK; +} +#endif + void llama_context::resolve_fused_ops(const llama_memory_context_i * mctx, uint32_t n_seqs) { const char * func = __func__; auto resolve = [&](const llm_fused_op_probe & probe, bool & enabled) { @@ -1312,7 +2129,19 @@ llm_graph_result * llama_context::process_ubatch(const llama_ubatch & ubatch, ll res->reset(); ggml_backend_sched_reset(sched.get()); +#if defined(GGML_CUDA_PALW_TRACE) + if (palw_cuda_trace_mixed_stream) { + palw_cuda_trace_attention_metadata.clear(); + } + if (palw_cuda_trace_bound) { + ggml_backend_sched_set_eval_callback( + sched.get(), palw_cuda_trace_eval_callback, this); + } else { + ggml_backend_sched_set_eval_callback(sched.get(), cparams.cb_eval, cparams.cb_eval_user_data); + } +#else ggml_backend_sched_set_eval_callback(sched.get(), cparams.cb_eval, cparams.cb_eval_user_data); +#endif //const auto t_start_us = ggml_time_us(); @@ -2424,6 +3253,12 @@ ggml_status llama_context::graph_compute( } auto status = ggml_backend_sched_graph_compute_async(sched.get(), gf); +#if defined(GGML_CUDA_PALW_TRACE) + if (status == GGML_STATUS_SUCCESS && palw_cuda_trace_bound && + palw_cuda_trace_status_ != PALW_CUDA_TRACE_OK) { + status = GGML_STATUS_FAILED; + } +#endif if (status != GGML_STATUS_SUCCESS) { LLAMA_LOG_ERROR("%s: ggml_backend_sched_graph_compute_async failed with error %d\n", __func__, status); } @@ -2435,6 +3270,49 @@ ggml_status llama_context::graph_compute( llm_graph_cb llama_context::graph_get_cb() const { return [&](const llama_ubatch & ubatch, ggml_tensor * cur, const char * name, int il) { +#if defined(GGML_CUDA_PALW_TRACE) + if (palw_cuda_trace_mixed_stream && cur != nullptr && name != nullptr && il >= 0) { + uint8_t stage = PALW_CUDA_TRACE_ATTENTION_STAGE_NONE; + if (strcmp(name, "kq") == 0) { + stage = PALW_CUDA_TRACE_ATTENTION_STAGE_EAGER_QK_SCORE_MMVF; + } else if (strcmp(name, "kq_soft_max") == 0) { + stage = PALW_CUDA_TRACE_ATTENTION_STAGE_EAGER_MASKED_SCALED_SOFTMAX; + } else if (strcmp(name, "kqv") == 0) { + stage = PALW_CUDA_TRACE_ATTENTION_STAGE_EAGER_VALUE_AGGREGATION_MMVF; + } + if (stage != PALW_CUDA_TRACE_ATTENTION_STAGE_NONE) { + const bool pinned_fresh_ubatch = + ubatch.n_tokens == 1 && ubatch.n_seq_tokens == 1 && + ubatch.n_seqs == 1 && ubatch.n_seqs_unq == 1 && + ubatch.n_pos == 1 && ubatch.pos != nullptr && + ubatch.pos[0] == 0 && ubatch.n_seq_id != nullptr && + ubatch.n_seq_id[0] == 1 && ubatch.seq_id != nullptr && + ubatch.seq_id[0] != nullptr && ubatch.seq_id[0][0] == 0; + palw_cuda_attention_graph_metadata metadata = {}; + metadata.layer_id = static_cast(il); + metadata.query_tokens = ubatch.n_tokens; + metadata.key_value_tokens = pinned_fresh_ubatch ? 1U : 0U; + // The pinned one-token Qwen3 path uses QWEN_KV_PADDING=256. + // Runtime tensor shapes and strides are independently copied + // into association_v2 and rechecked by the work hook. An + // imported state or nonzero/ambiguous input position is not + // silently re-described as this fresh-KV profile: the zero + // values make association validation fail closed. + metadata.physical_key_value_tokens = + pinned_fresh_ubatch ? 256U : 0U; + metadata.query_heads = 32; + metadata.key_value_heads = 8; + metadata.head_dim = 128; + metadata.logical_batch = ubatch.n_seqs; + metadata.decode_step = 0; + metadata.attention_stage = stage; + metadata.phase = PALW_CUDA_TRACE_PREFILL; + metadata.causal = 1; + metadata.layer_present = 1; + palw_cuda_trace_attention_metadata.insert_or_assign(cur, metadata); + } + } +#endif if (il >= 0) { ggml_format_name(cur, "%s-%d", name, il); } else { @@ -3114,6 +3992,15 @@ size_t llama_context::state_write_data(llama_io_write_i & io) { size_t llama_context::state_read_data(llama_io_read_i & io) { LLAMA_LOG_DEBUG("%s: reading state\n", __func__); +#if defined(GGML_CUDA_PALW_TRACE) + if (palw_cuda_trace_bound || palw_cuda_trace_producer_owned) { + throw std::runtime_error("cannot import context state while a PALW CUDA trace exists"); + } + // Set this before parsing: a failed or partial import must also keep the + // pinned empty-KV attention profile unavailable. + palw_cuda_state_imported = true; +#endif + // read model info { LLAMA_LOG_DEBUG("%s: - reading model info\n", __func__); @@ -3150,6 +4037,13 @@ size_t llama_context::state_seq_write_data(llama_io_write_i & io, llama_seq_id s size_t llama_context::state_seq_read_data(llama_io_read_i & io, llama_seq_id seq_id, llama_state_seq_flags flags) { GGML_UNUSED(seq_id); +#if defined(GGML_CUDA_PALW_TRACE) + if (palw_cuda_trace_bound || palw_cuda_trace_producer_owned) { + throw std::runtime_error("cannot import sequence state while a PALW CUDA trace exists"); + } + palw_cuda_state_imported = true; +#endif + if (memory) { memory->state_read(io, seq_id, flags); } @@ -3568,6 +4462,96 @@ void llama_free(llama_context * ctx) { delete ctx; } +#if defined(GGML_CUDA_PALW_TRACE) +palw_cuda_trace_status llama_palw_cuda_trace_attach( + llama_context * ctx, + uint64_t record_capacity, + const palw_cuda_llama_mmvq_request_v1 * request) { + if (ctx == nullptr) { + return PALW_CUDA_TRACE_INVALID_ARGUMENT; + } + return ctx->palw_cuda_trace_attach(record_capacity, request); +} + +palw_cuda_trace_status llama_palw_cuda_trace_attach_v2( + llama_context * ctx, + uint64_t record_capacity, + const palw_cuda_llama_trace_request_v2 * request) { + if (ctx == nullptr) { + return PALW_CUDA_TRACE_INVALID_ARGUMENT; + } + return ctx->palw_cuda_trace_attach_v2(record_capacity, request); +} + +palw_cuda_trace_status llama_palw_cuda_trace_kernel_info( + llama_context * ctx, + uint8_t quantization, + palw_cuda_llama_mmvq_kernel_info_v1 * out_info) { + if (ctx == nullptr) { + return PALW_CUDA_TRACE_INVALID_ARGUMENT; + } + return ctx->palw_cuda_trace_kernel_info(quantization, out_info); +} + +palw_cuda_trace_status llama_palw_cuda_trace_attention_kernel_info( + llama_context * ctx, + const palw_cuda_llama_trace_association_v2 * actual, + palw_cuda_llama_attention_kernel_info_v1 * out_info) { + if (ctx == nullptr) { + return PALW_CUDA_TRACE_INVALID_ARGUMENT; + } + return ctx->palw_cuda_trace_attention_kernel_info(actual, out_info); +} + +palw_cuda_trace_status llama_palw_cuda_trace_grouped_capture_info( + llama_context * ctx, + palw_cuda_producer_grouped_capture_info_v3 * out_info) { + if (ctx == nullptr) { + return PALW_CUDA_TRACE_INVALID_ARGUMENT; + } + return ctx->palw_cuda_trace_grouped_capture_info(out_info); +} + +palw_cuda_trace_status llama_palw_cuda_trace_unbind( + llama_context * ctx, + uint64_t * accepted_launch_count) { + if (ctx == nullptr) { + return PALW_CUDA_TRACE_INVALID_ARGUMENT; + } + return ctx->palw_cuda_trace_unbind(accepted_launch_count); +} + +palw_cuda_trace_status llama_palw_cuda_trace_finalize( + llama_context * ctx, + palw_cuda_trace_record_v3 * host_records, + uint64_t host_capacity) { + if (ctx == nullptr) { + return PALW_CUDA_TRACE_INVALID_ARGUMENT; + } + return ctx->palw_cuda_trace_finalize(host_records, host_capacity); +} + +palw_cuda_trace_status llama_palw_cuda_trace_diagnostics( + llama_context * ctx, + uint64_t * count, + uint64_t * committed_count, + uint32_t * faults, + int * poisoned) { + if (ctx == nullptr) { + return PALW_CUDA_TRACE_INVALID_ARGUMENT; + } + return ctx->palw_cuda_trace_diagnostics( + count, committed_count, faults, poisoned); +} + +palw_cuda_trace_status llama_palw_cuda_trace_destroy(llama_context * ctx) { + if (ctx == nullptr) { + return PALW_CUDA_TRACE_INVALID_ARGUMENT; + } + return ctx->palw_cuda_trace_destroy(); +} +#endif + uint32_t llama_n_ctx(const llama_context * ctx) { return ctx->n_ctx(); } diff --git a/src/llama-context.h b/src/llama-context.h index bf91daa8..a9753b00 100644 --- a/src/llama-context.h +++ b/src/llama-context.h @@ -12,6 +12,7 @@ #include "ggml-opt.h" #include +#include #include struct llama_model; @@ -247,6 +248,33 @@ public: // returns the result of ggml_backend_sched_graph_compute_async execution ggml_status graph_compute(ggml_cgraph * gf, bool batched); +#if defined(GGML_CUDA_PALW_TRACE) + palw_cuda_trace_status palw_cuda_trace_attach( + uint64_t record_capacity, + const palw_cuda_llama_mmvq_request_v1 * request); + palw_cuda_trace_status palw_cuda_trace_attach_v2( + uint64_t record_capacity, + const palw_cuda_llama_trace_request_v2 * request); + palw_cuda_trace_status palw_cuda_trace_kernel_info( + uint8_t quantization, + palw_cuda_llama_mmvq_kernel_info_v1 * out_info); + palw_cuda_trace_status palw_cuda_trace_attention_kernel_info( + const palw_cuda_llama_trace_association_v2 * actual, + palw_cuda_llama_attention_kernel_info_v1 * out_info); + palw_cuda_trace_status palw_cuda_trace_grouped_capture_info( + palw_cuda_producer_grouped_capture_info_v3 * out_info); + palw_cuda_trace_status palw_cuda_trace_unbind(uint64_t * accepted_launch_count); + palw_cuda_trace_status palw_cuda_trace_finalize( + palw_cuda_trace_record_v3 * host_records, + uint64_t host_capacity); + palw_cuda_trace_status palw_cuda_trace_diagnostics( + uint64_t * count, + uint64_t * committed_count, + uint32_t * faults, + int * poisoned) const; + palw_cuda_trace_status palw_cuda_trace_destroy(); +#endif + // reserve a graph with a dummy ubatch of the specified size ggml_cgraph * graph_reserve( uint32_t n_tokens, uint32_t n_seqs, uint32_t n_outputs, const llama_memory_context_i * mctx, bool split_only = false, size_t * sizes = nullptr); @@ -262,6 +290,18 @@ private: llm_graph_cb graph_get_cb() const; +#if defined(GGML_CUDA_PALW_TRACE) + void palw_cuda_trace_bind(); + bool palw_cuda_trace_eval(ggml_tensor * tensor, bool ask); + bool palw_cuda_trace_candidate(const ggml_tensor * tensor) const; + void palw_cuda_trace_record_status(palw_cuda_trace_status status); + + static bool palw_cuda_trace_eval_callback( + ggml_tensor * tensor, + bool ask, + void * user_data); +#endif + // disable auto fused ops (Flash Attention, Gated Delta Net) whose op lands on a device // that differs from the layer it belongs to (usually due to missing backend support) void resolve_fused_ops(const llama_memory_context_i * mctx, uint32_t n_seqs); @@ -348,6 +388,49 @@ private: ggml_backend_t backend_cpu = nullptr; std::vector backends; +#if defined(GGML_CUDA_PALW_TRACE) + ggml_backend_t palw_cuda_trace_backend = nullptr; + const palw_cuda_llama_mmvq_backend_api_v1 * palw_cuda_trace_backend_api = nullptr; + palw_cuda_llama_mmvq_associate_backend_v1 palw_cuda_trace_associate_fn = nullptr; + palw_cuda_llama_mmvq_complete_backend_v1 palw_cuda_trace_complete_fn = nullptr; + palw_cuda_llama_mmvq_unbind_backend_v1 palw_cuda_trace_unbind_fn = nullptr; + palw_cuda_llama_mmvq_kernel_info_backend_v1 palw_cuda_trace_kernel_info_fn = nullptr; + palw_cuda_producer_trace_context_v3 * palw_cuda_trace_producer_context = nullptr; + palw_cuda_llama_mmvq_request_v1 palw_cuda_trace_owned_request = {}; + palw_cuda_llama_trace_associate_backend_v2 palw_cuda_trace_associate_fn_v2 = nullptr; + palw_cuda_llama_trace_complete_backend_v2 palw_cuda_trace_complete_fn_v2 = nullptr; + palw_cuda_llama_trace_unbind_backend_v2 palw_cuda_trace_unbind_fn_v2 = nullptr; + palw_cuda_llama_attention_kernel_info_backend_v1 + palw_cuda_trace_attention_kernel_info_fn = nullptr; + palw_cuda_llama_grouped_capture_info_backend_v1 + palw_cuda_trace_grouped_capture_info_fn = nullptr; + palw_cuda_llama_trace_request_v2 palw_cuda_trace_owned_request_v2 = {}; + struct palw_cuda_attention_graph_metadata { + uint32_t layer_id = 0; + uint32_t query_tokens = 0; + uint32_t key_value_tokens = 0; + uint32_t physical_key_value_tokens = 0; + uint32_t query_heads = 0; + uint32_t key_value_heads = 0; + uint32_t head_dim = 0; + uint32_t logical_batch = 0; + uint64_t decode_step = 0; + uint8_t attention_stage = PALW_CUDA_TRACE_ATTENTION_STAGE_NONE; + uint8_t phase = PALW_CUDA_TRACE_PREFILL; + uint8_t causal = 0; + uint8_t layer_present = 0; + }; + mutable std::unordered_map + palw_cuda_trace_attention_metadata; + const void * palw_cuda_trace_pending_operation = nullptr; + palw_cuda_trace_status palw_cuda_trace_status_ = PALW_CUDA_TRACE_OK; + bool palw_cuda_trace_bound = false; + bool palw_cuda_trace_finalized = false; + bool palw_cuda_trace_producer_owned = false; + bool palw_cuda_trace_mixed_stream = false; + bool palw_cuda_state_imported = false; +#endif + // training ggml_opt_context_t opt_ctx = nullptr; diff --git a/src/llama-cparams.h b/src/llama-cparams.h index 58520caa..7140b114 100644 --- a/src/llama-cparams.h +++ b/src/llama-cparams.h @@ -58,4 +58,8 @@ struct llama_cparams { void * cb_eval_user_data; llama_context * ctx_other; + +#if defined(GGML_CUDA_PALW_TRACE) + const palw_cuda_llama_mmvq_request_v1 * palw_cuda_trace_request; +#endif }; diff --git a/tools/CMakeLists.txt b/tools/CMakeLists.txt index 780df326..d822ba52 100644 --- a/tools/CMakeLists.txt +++ b/tools/CMakeLists.txt @@ -28,6 +28,10 @@ else() endif() add_subdirectory(tokenize) add_subdirectory(parser) + add_subdirectory(palw-observer) + if (GGML_CUDA AND GGML_CUDA_PALW_TRACE) + add_subdirectory(palw-mmvq-smoke) + endif() add_subdirectory(tts) add_subdirectory(mtmd) if (GGML_RPC) diff --git a/tools/palw-mmvq-smoke/CMakeLists.txt b/tools/palw-mmvq-smoke/CMakeLists.txt new file mode 100644 index 00000000..7213ef33 --- /dev/null +++ b/tools/palw-mmvq-smoke/CMakeLists.txt @@ -0,0 +1,6 @@ +set(TARGET llama-palw-mmvq-smoke) + +add_executable(${TARGET} palw-mmvq-smoke.cpp) +target_link_libraries(${TARGET} PRIVATE llama ${CMAKE_THREAD_LIBS_INIT}) +target_compile_features(${TARGET} PRIVATE cxx_std_17) + diff --git a/tools/palw-mmvq-smoke/palw-mmvq-smoke.cpp b/tools/palw-mmvq-smoke/palw-mmvq-smoke.cpp new file mode 100644 index 00000000..e601e32e --- /dev/null +++ b/tools/palw-mmvq-smoke/palw-mmvq-smoke.cpp @@ -0,0 +1,1983 @@ +#include "ggml-backend.h" +#include "ggml.h" +#include "llama.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#if !defined(GGML_CUDA_PALW_TRACE) +#error "llama-palw-mmvq-smoke requires GGML_CUDA_PALW_TRACE" +#endif + +namespace { + +constexpr uint64_t EXPECTED_LAUNCH_COUNT = 253; +constexpr uint64_t EXPECTED_MIXED_LAUNCH_COUNT = 361; +constexpr uint64_t EXPECTED_ATTENTION_COUNT = 108; +constexpr uint64_t EXPECTED_Q4_K_COUNT = 216; +constexpr uint64_t EXPECTED_Q6_K_COUNT = 37; +constexpr uint32_t CONTEXT_TOKENS = 512; +constexpr uint64_t FNV64_OFFSET = UINT64_C(14695981039346656037); +constexpr uint64_t FNV64_PRIME = UINT64_C(1099511628211); +constexpr const char * Q4_EXACT_SYMBOL = + "_Z13mul_mat_vec_qIL9ggml_type12ELi1ELb0ELb0EEvPKvS2_PKi31ggml_cuda_mm_fusion_args_devicePfj5uint3jjjS7_jjjS7_jjjj33palw_cuda_producer_device_view_v3"; +constexpr const char * Q6_EXACT_SYMBOL = + "_Z13mul_mat_vec_qIL9ggml_type14ELi1ELb0ELb0EEvPKvS2_PKi31ggml_cuda_mm_fusion_args_devicePfj5uint3jjjS7_jjjS7_jjjj33palw_cuda_producer_device_view_v3"; +constexpr const char * QK_EXACT_SYMBOL = + "_Z13mul_mat_vec_fI6__halffLi1ELi64ELb0ELb0EEvPKT_PKfPKi31ggml_cuda_mm_fusion_args_devicePfi5uint3iiiSA_iiiSA_iiii"; +constexpr const char * PV_EXACT_SYMBOL = + "_Z13mul_mat_vec_fI6__halfS0_Li1ELi128ELb0ELb0EEvPKT_PKfPKi31ggml_cuda_mm_fusion_args_devicePfi5uint3iiiSA_iiiSA_iiii"; +constexpr const char * SOFTMAX_EXACT_SYMBOL = + "_Z12soft_max_f32ILb1ELi256ELi256EfEvPKfPKT2_S1_Pf15soft_max_params"; + +struct options { + std::string model_path; + std::string binary_digest_hex; + std::array binary_digest{}; + std::array q4_cubin_digest{}; + std::array q6_cubin_digest{}; + std::array qk_cubin_digest{}; + std::array softmax_cubin_digest{}; + std::array pv_cubin_digest{}; + std::string q4_cubin_digest_hex; + std::string q6_cubin_digest_hex; + std::string qk_cubin_digest_hex; + std::string softmax_cubin_digest_hex; + std::string pv_cubin_digest_hex; + // Optional path to write the raw canonical V3 record stream for the Rust + // authority verifier. Diagnostic-only: does not change any computed record + // or the canonical fingerprint. + std::string emit_records_path; + // Optional path to a 361*128-byte authority-derived canonical identity table + // that the live callback consumes in place of self-derived diagnostic IDs + // (R13). Order: per launch, operation/plan/physical-layout/attention-group. + std::string authority_ids_path; + uint64_t reject_index = UINT64_MAX; + bool mixed = false; +}; + +struct candidate { + const void * weight_data = nullptr; + uint64_t m = 0; + uint64_t n = 0; + uint64_t k = 0; + uint64_t batch = 0; + uint8_t quantization = PALW_CUDA_TRACE_QUANTIZATION_NONE; +}; + +struct discovery_state { + std::vector candidates; + bool failed = false; +}; + +struct mixed_candidate { + const void * weight_data = nullptr; + uint64_t m = 0; + uint64_t n = 0; + uint64_t k = 0; + uint64_t batch = 0; + uint32_t output_rows = 0; + uint32_t output_columns = 0; + uint32_t layer_id = 0; + uint8_t kind = PALW_CUDA_TRACE_GEMM; + uint8_t quantization = PALW_CUDA_TRACE_QUANTIZATION_NONE; + uint8_t attention_stage = PALW_CUDA_TRACE_ATTENTION_STAGE_NONE; +}; + +struct mixed_discovery_state { + std::vector candidates; + uint32_t attention_seen = 0; + bool failed = false; +}; + +enum class callback_error : uint32_t { + none = 0, + invalid_argument = 1, + unexpected_count = 2, + association_mismatch = 3, + duplicate_operation_token = 4, + launch_notification_mismatch = 5, +}; + +struct trace_state { + const std::vector * candidates = nullptr; + std::array operation_tokens{}; + std::array binary_digest{}; + palw_cuda_llama_mmvq_kernel_info_v1 q4_info{}; + palw_cuda_llama_mmvq_kernel_info_v1 q6_info{}; + uint64_t reject_index = UINT64_MAX; + uint64_t approved = 0; + uint64_t launch_notifications = 0; + uint64_t fault_callbacks = 0; + palw_cuda_trace_status last_fault_status = PALW_CUDA_TRACE_OK; + palw_cuda_llama_mmvq_fault_v1 last_fault = PALW_CUDA_LLAMA_MMVQ_FAULT_NONE; + callback_error error = callback_error::none; + bool rejection_triggered = false; +}; + +struct mixed_trace_state { + const std::vector * candidates = nullptr; + std::array operation_tokens{}; + std::array q4_cubin_digest{}; + std::array q6_cubin_digest{}; + std::array qk_cubin_digest{}; + std::array softmax_cubin_digest{}; + std::array pv_cubin_digest{}; + palw_cuda_llama_mmvq_kernel_info_v1 q4_info{}; + palw_cuda_llama_mmvq_kernel_info_v1 q6_info{}; + palw_cuda_llama_attention_kernel_info_v1 qk_info{}; + palw_cuda_llama_attention_kernel_info_v1 softmax_info{}; + palw_cuda_llama_attention_kernel_info_v1 pv_info{}; + palw_cuda_producer_grouped_capture_info_v3 collector_info{}; + uint64_t reject_index = UINT64_MAX; + uint64_t approved = 0; + uint64_t launch_notifications = 0; + uint64_t fault_callbacks = 0; + palw_cuda_trace_status last_fault_status = PALW_CUDA_TRACE_OK; + palw_cuda_llama_mmvq_fault_v1 last_fault = PALW_CUDA_LLAMA_MMVQ_FAULT_NONE; + callback_error error = callback_error::none; + bool rejection_triggered = false; + // Optional authority-derived canonical identity table (R13): 361 launches x + // (operation, plan, physical-layout, attention-group) x 32 bytes, in launch + // order. When non-null the callback stamps these authority identities in + // place of the self-derived `misaka.palw.diagnostic.*` ones. + const uint8_t * authority_ids = nullptr; +}; + +void print_usage(const char * argv0) { + std::fprintf(stderr, + "usage: %s --model MODEL [--kernel-binary-sha256 64_HEX | --mixed " + "--q4-cubin-sha256 HEX --q6-cubin-sha256 HEX --qk-cubin-sha256 HEX " + "--softmax-cubin-sha256 HEX --pv-cubin-sha256 HEX] [--reject-index INDEX] " + "[--emit-records PATH] [--authority-ids PATH]\n", + argv0); +} + +bool take_value(int argc, char ** argv, int & index, const char * option, const char *& value) { + if (index + 1 >= argc) { + std::fprintf(stderr, "error: %s requires a value\n", option); + return false; + } + value = argv[++index]; + return true; +} + +uint8_t hex_nibble(char value) { + if (value >= '0' && value <= '9') { + return static_cast(value - '0'); + } + if (value >= 'a' && value <= 'f') { + return static_cast(value - 'a' + 10); + } + if (value >= 'A' && value <= 'F') { + return static_cast(value - 'A' + 10); + } + return 0xff; +} + +bool parse_digest_value( + std::string_view value, + std::array & digest, + std::string & digest_hex) { + if (value.size() != PALW_CUDA_PRODUCER_ID_SIZE * 2) { + return false; + } + static constexpr char HEX[] = "0123456789abcdef"; + bool nonzero = false; + digest_hex.resize(value.size()); + for (size_t index = 0; index < digest.size(); ++index) { + const uint8_t high = hex_nibble(value[index * 2]); + const uint8_t low = hex_nibble(value[index * 2 + 1]); + if (high == 0xff || low == 0xff) { + return false; + } + const uint8_t byte = static_cast((high << 4) | low); + digest[index] = byte; + digest_hex[index * 2] = HEX[high]; + digest_hex[index * 2 + 1] = HEX[low]; + nonzero = nonzero || byte != 0; + } + return nonzero; +} + +bool parse_options(int argc, char ** argv, options & opts, bool & help) { + help = false; + bool digest_seen = false; + for (int index = 1; index < argc; ++index) { + const std::string_view argument(argv[index]); + const char * value = nullptr; + if (argument == "--help" || argument == "-h") { + help = true; + return true; + } + if (argument == "--model" || argument == "-m") { + if (!take_value(argc, argv, index, argv[index], value)) { + return false; + } + opts.model_path = value; + continue; + } + if (argument == "--kernel-binary-sha256") { + if (digest_seen || !take_value(argc, argv, index, argv[index], value) || + !parse_digest_value( + value, opts.binary_digest, opts.binary_digest_hex)) { + std::fprintf(stderr, + "error: --kernel-binary-sha256 must be one nonzero 64-hex digest\n"); + return false; + } + digest_seen = true; + continue; + } + const auto take_entry_digest = [&](auto & digest, std::string & hex) { + return take_value(argc, argv, index, argv[index], value) && + parse_digest_value(value, digest, hex); + }; + if (argument == "--q4-cubin-sha256") { + if (!opts.q4_cubin_digest_hex.empty() || + !take_entry_digest(opts.q4_cubin_digest, opts.q4_cubin_digest_hex)) { + return false; + } + continue; + } + if (argument == "--q6-cubin-sha256") { + if (!opts.q6_cubin_digest_hex.empty() || + !take_entry_digest(opts.q6_cubin_digest, opts.q6_cubin_digest_hex)) { + return false; + } + continue; + } + if (argument == "--qk-cubin-sha256") { + if (!opts.qk_cubin_digest_hex.empty() || + !take_entry_digest(opts.qk_cubin_digest, opts.qk_cubin_digest_hex)) { + return false; + } + continue; + } + if (argument == "--softmax-cubin-sha256") { + if (!opts.softmax_cubin_digest_hex.empty() || + !take_entry_digest( + opts.softmax_cubin_digest, opts.softmax_cubin_digest_hex)) { + return false; + } + continue; + } + if (argument == "--pv-cubin-sha256") { + if (!opts.pv_cubin_digest_hex.empty() || + !take_entry_digest(opts.pv_cubin_digest, opts.pv_cubin_digest_hex)) { + return false; + } + continue; + } + if (argument == "--mixed") { + if (opts.mixed) { + return false; + } + opts.mixed = true; + continue; + } + if (argument == "--emit-records") { + if (!opts.emit_records_path.empty() || + !take_value(argc, argv, index, argv[index], value)) { + return false; + } + opts.emit_records_path = value; + continue; + } + if (argument == "--authority-ids") { + if (!opts.authority_ids_path.empty() || + !take_value(argc, argv, index, argv[index], value)) { + return false; + } + opts.authority_ids_path = value; + continue; + } + if (argument == "--reject-index") { + if (opts.reject_index != UINT64_MAX || + !take_value(argc, argv, index, argv[index], value)) { + return false; + } + const char * end = value + std::strlen(value); + const auto converted = std::from_chars(value, end, opts.reject_index); + if (converted.ec != std::errc() || converted.ptr != end || + opts.reject_index >= EXPECTED_MIXED_LAUNCH_COUNT) { + std::fprintf(stderr, + "error: --reject-index exceeds the selected stream length\n"); + return false; + } + continue; + } + std::fprintf(stderr, "error: unknown argument: %s\n", argv[index]); + return false; + } + const bool mixed_digests = !opts.q4_cubin_digest_hex.empty() && + !opts.q6_cubin_digest_hex.empty() && !opts.qk_cubin_digest_hex.empty() && + !opts.softmax_cubin_digest_hex.empty() && !opts.pv_cubin_digest_hex.empty(); + if (opts.model_path.empty() || (opts.mixed ? !mixed_digests : !digest_seen)) { + std::fprintf(stderr, + "error: --model plus the selected legacy or five-entry digest set are required\n"); + return false; + } + if (!opts.mixed && opts.reject_index >= EXPECTED_LAUNCH_COUNT && + opts.reject_index != UINT64_MAX) { + std::fprintf(stderr, "error: legacy --reject-index must be in [0, 252]\n"); + return false; + } + return true; +} + +void log_callback(ggml_log_level level, const char * text, void *) { + if (level == GGML_LOG_LEVEL_WARN || level == GGML_LOG_LEVEL_ERROR) { + std::fprintf(stderr, "%s", text != nullptr ? text : ""); + } +} + +uint8_t quantization_for_type(ggml_type type) { + if (type == GGML_TYPE_Q4_K) { + return PALW_CUDA_TRACE_QUANTIZATION_Q4_K; + } + if (type == GGML_TYPE_Q6_K) { + return PALW_CUDA_TRACE_QUANTIZATION_Q6_K; + } + return PALW_CUDA_TRACE_QUANTIZATION_NONE; +} + +uint64_t positive_dimension(int64_t value) { + return value > 0 ? static_cast(value) : 0; +} + +bool discover_callback(ggml_tensor * tensor, bool ask, void * user_data) noexcept { + auto * state = static_cast(user_data); + if (state == nullptr || state->failed || !ask || tensor == nullptr || + (tensor->op != GGML_OP_MUL_MAT && tensor->op != GGML_OP_MUL_MAT_ID) || + tensor->src[0] == nullptr || !ggml_is_quantized(tensor->src[0]->type)) { + return false; + } + + try { + const ggml_tensor * weight = tensor->src[0]; + candidate value{}; + value.weight_data = weight->data; + value.m = positive_dimension(tensor->ne[0]); + value.n = positive_dimension(tensor->ne[1]); + value.k = positive_dimension(weight->ne[0]); + const uint64_t batch_2 = positive_dimension(tensor->ne[2]); + const uint64_t batch_3 = positive_dimension(tensor->ne[3]); + value.batch = batch_2 != 0 && batch_3 <= UINT64_MAX / batch_2 + ? batch_2 * batch_3 + : 0; + value.quantization = quantization_for_type(weight->type); + state->candidates.push_back(value); + } catch (...) { + state->failed = true; + } + return false; +} + +uint8_t discover_attention_stage(const ggml_tensor * tensor) { + if (tensor == nullptr || tensor->type != GGML_TYPE_F32 || + tensor->src[0] == nullptr) { + return PALW_CUDA_TRACE_ATTENTION_STAGE_NONE; + } + if (tensor->op == GGML_OP_SOFT_MAX && tensor->src[0]->type == GGML_TYPE_F32 && + tensor->src[1] != nullptr && tensor->src[1]->type == GGML_TYPE_F32 && + tensor->src[2] == nullptr && tensor->ne[0] == 256 && + tensor->ne[1] == 1 && tensor->ne[2] == 32 && tensor->ne[3] == 1) { + return PALW_CUDA_TRACE_ATTENTION_STAGE_EAGER_MASKED_SCALED_SOFTMAX; + } + if (tensor->op != GGML_OP_MUL_MAT || tensor->src[0]->type != GGML_TYPE_F16 || + tensor->src[1] == nullptr || tensor->src[1]->type != GGML_TYPE_F32 || + tensor->ne[1] != 1 || tensor->ne[2] != 32 || tensor->ne[3] != 1 || + tensor->src[0]->ne[2] != 8 || tensor->src[0]->ne[3] != 1) { + return PALW_CUDA_TRACE_ATTENTION_STAGE_NONE; + } + if (tensor->src[0]->ne[0] == 128 && tensor->src[0]->ne[1] == 256 && + tensor->ne[0] == 256) { + return PALW_CUDA_TRACE_ATTENTION_STAGE_EAGER_QK_SCORE_MMVF; + } + if (tensor->src[0]->ne[0] == 256 && tensor->src[0]->ne[1] == 128 && + tensor->ne[0] == 128) { + return PALW_CUDA_TRACE_ATTENTION_STAGE_EAGER_VALUE_AGGREGATION_MMVF; + } + return PALW_CUDA_TRACE_ATTENTION_STAGE_NONE; +} + +bool discover_mixed_callback(ggml_tensor * tensor, bool ask, void * user_data) noexcept { + auto * state = static_cast(user_data); + if (state == nullptr || state->failed || !ask || tensor == nullptr) { + return false; + } + try { + mixed_candidate value{}; + value.attention_stage = discover_attention_stage(tensor); + if (value.attention_stage != PALW_CUDA_TRACE_ATTENTION_STAGE_NONE) { + const uint8_t expected_stage = static_cast( + state->attention_seen % 3U + 1U); + if (value.attention_stage != expected_stage) { + state->failed = true; + return false; + } + value.layer_id = state->attention_seen / 3U; + value.kind = value.attention_stage == + PALW_CUDA_TRACE_ATTENTION_STAGE_EAGER_MASKED_SCALED_SOFTMAX + ? PALW_CUDA_TRACE_ATTENTION + : PALW_CUDA_TRACE_GEMM; + value.m = positive_dimension(tensor->ne[0]); + value.n = 1; + value.k = value.attention_stage == + PALW_CUDA_TRACE_ATTENTION_STAGE_EAGER_MASKED_SCALED_SOFTMAX + ? 1 + : positive_dimension(tensor->src[0]->ne[0]); + value.batch = 32; + value.output_rows = 32; + value.output_columns = static_cast(value.m); + ++state->attention_seen; + state->candidates.push_back(value); + return false; + } + if ((tensor->op == GGML_OP_MUL_MAT || tensor->op == GGML_OP_MUL_MAT_ID) && + tensor->src[0] != nullptr && ggml_is_quantized(tensor->src[0]->type)) { + const ggml_tensor * weight = tensor->src[0]; + value.weight_data = weight->data; + value.m = positive_dimension(tensor->ne[0]); + value.n = positive_dimension(tensor->ne[1]); + value.k = positive_dimension(weight->ne[0]); + value.batch = positive_dimension(tensor->ne[2]) * + positive_dimension(tensor->ne[3]); + value.output_rows = static_cast(value.n); + value.output_columns = static_cast(value.m); + value.quantization = quantization_for_type(weight->type); + state->candidates.push_back(value); + } + } catch (...) { + state->failed = true; + } + return false; +} + +bool same_candidate( + const candidate & expected, + const palw_cuda_llama_mmvq_association_v1 & actual) { + return expected.weight_data == actual.weight_data && expected.m == actual.m && + expected.n == actual.n && expected.k == actual.k && expected.batch == actual.batch && + expected.quantization == actual.quantization; +} + +uint64_t fnv_byte(uint64_t hash, uint8_t byte) { + return (hash ^ byte) * FNV64_PRIME; +} + +uint64_t fnv_bytes(uint64_t hash, const void * data, size_t size) { + const auto * bytes = static_cast(data); + for (size_t index = 0; index < size; ++index) { + hash = fnv_byte(hash, bytes[index]); + } + return hash; +} + +uint64_t fnv_u64_be(uint64_t hash, uint64_t value) { + for (int shift = 56; shift >= 0; shift -= 8) { + hash = fnv_byte(hash, static_cast(value >> shift)); + } + return hash; +} + +void store_u64_be(uint8_t * output, uint64_t value) { + for (int shift = 56; shift >= 0; shift -= 8) { + *output++ = static_cast(value >> shift); + } +} + +void derive_diagnostic_id( + const char * domain, + uint64_t sequence, + const palw_cuda_llama_mmvq_association_v1 * actual, + uint8_t output[PALW_CUDA_PRODUCER_ID_SIZE]) { + for (uint64_t lane = 0; lane < 4; ++lane) { + uint64_t hash = fnv_bytes(FNV64_OFFSET, domain, std::strlen(domain)); + hash = fnv_u64_be(hash, lane); + hash = fnv_u64_be(hash, sequence); + if (actual != nullptr) { + hash = fnv_u64_be(hash, actual->m); + hash = fnv_u64_be(hash, actual->n); + hash = fnv_u64_be(hash, actual->k); + hash = fnv_u64_be(hash, actual->batch); + hash = fnv_byte(hash, actual->quantization); + } + store_u64_be(output + lane * sizeof(uint64_t), hash); + } +} + +const palw_cuda_llama_mmvq_kernel_info_v1 * kernel_info_for( + const trace_state & state, + uint8_t quantization) { + if (quantization == PALW_CUDA_TRACE_QUANTIZATION_Q4_K) { + return &state.q4_info; + } + if (quantization == PALW_CUDA_TRACE_QUANTIZATION_Q6_K) { + return &state.q6_info; + } + return nullptr; +} + +void derive_diagnostic_id_v2( + const char * domain, + uint64_t sequence, + const palw_cuda_llama_trace_association_v2 * actual, + uint8_t output[PALW_CUDA_PRODUCER_ID_SIZE]) { + for (uint64_t lane = 0; lane < 4; ++lane) { + uint64_t hash = fnv_bytes(FNV64_OFFSET, domain, std::strlen(domain)); + hash = fnv_u64_be(hash, lane); + hash = fnv_u64_be(hash, sequence); + if (actual != nullptr) { + hash = fnv_u64_be(hash, actual->m); + hash = fnv_u64_be(hash, actual->n); + hash = fnv_u64_be(hash, actual->k); + hash = fnv_u64_be(hash, actual->batch); + hash = fnv_u64_be(hash, actual->layer_id); + hash = fnv_byte(hash, actual->kind); + hash = fnv_byte(hash, actual->quantization); + hash = fnv_byte(hash, actual->attention_stage); + } + store_u64_be(output + lane * sizeof(uint64_t), hash); + } +} + +bool same_mixed_candidate( + const mixed_candidate & expected, + const palw_cuda_llama_trace_association_v2 & actual) { + if (expected.m != actual.m || expected.n != actual.n || expected.k != actual.k || + expected.batch != actual.batch || expected.output_rows != actual.output_rows || + expected.output_columns != actual.output_columns || + expected.kind != actual.kind || expected.quantization != actual.quantization || + expected.attention_stage != actual.attention_stage) { + return false; + } + if (expected.attention_stage == PALW_CUDA_TRACE_ATTENTION_STAGE_NONE) { + return expected.weight_data == actual.src0.data; + } + return expected.layer_id == actual.layer_id && actual.layer_present == 1 && + actual.query_tokens == 1 && actual.key_value_tokens == 1 && + actual.physical_key_value_tokens == 256 && actual.query_heads == 32 && + actual.key_value_heads == 8 && actual.head_dim == 128 && + actual.logical_batch == 1 && actual.phase == PALW_CUDA_TRACE_PREFILL && + actual.decode_step == 0 && actual.causal == 1; +} + +const palw_cuda_llama_mmvq_kernel_info_v1 * mixed_mmvq_info_for( + const mixed_trace_state & state, + uint8_t quantization) { + if (quantization == PALW_CUDA_TRACE_QUANTIZATION_Q4_K) { + return &state.q4_info; + } + if (quantization == PALW_CUDA_TRACE_QUANTIZATION_Q6_K) { + return &state.q6_info; + } + return nullptr; +} + +const palw_cuda_llama_attention_kernel_info_v1 * mixed_attention_info_for( + const mixed_trace_state & state, + uint8_t stage) { + if (stage == PALW_CUDA_TRACE_ATTENTION_STAGE_EAGER_QK_SCORE_MMVF) { + return &state.qk_info; + } + if (stage == PALW_CUDA_TRACE_ATTENTION_STAGE_EAGER_MASKED_SCALED_SOFTMAX) { + return &state.softmax_info; + } + if (stage == PALW_CUDA_TRACE_ATTENTION_STAGE_EAGER_VALUE_AGGREGATION_MMVF) { + return &state.pv_info; + } + return nullptr; +} + +palw_cuda_trace_status approve_callback_v2( + void * user_data, + const palw_cuda_llama_trace_association_v2 * actual, + palw_cuda_producer_launch_v3 * expected_launch) noexcept { + auto * state = static_cast(user_data); + if (state == nullptr || state->candidates == nullptr || actual == nullptr || + expected_launch == nullptr || state->error != callback_error::none || + !palw_cuda_llama_trace_association_is_supported_v2(actual)) { + if (state != nullptr) { + state->error = callback_error::invalid_argument; + } + return PALW_CUDA_TRACE_INVALID_ARGUMENT; + } + const uint64_t index = state->approved; + if (index >= state->candidates->size() || index >= EXPECTED_MIXED_LAUNCH_COUNT) { + state->error = callback_error::unexpected_count; + return PALW_CUDA_TRACE_SEQUENCE_MISMATCH; + } + if (index == state->reject_index) { + state->rejection_triggered = true; + return PALW_CUDA_TRACE_IDENTITY_MISMATCH; + } + if (!same_mixed_candidate((*state->candidates)[static_cast(index)], *actual)) { + state->error = callback_error::association_mismatch; + return PALW_CUDA_TRACE_IDENTITY_MISMATCH; + } + for (uint64_t prior = 0; prior < index; ++prior) { + if (state->operation_tokens[static_cast(prior)] == actual->operation_token) { + state->error = callback_error::duplicate_operation_token; + return PALW_CUDA_TRACE_SEQUENCE_MISMATCH; + } + } + + *expected_launch = {}; + expected_launch->launch_nonce = index + 1; + palw_cuda_trace_record_v3 & semantic = expected_launch->semantic_template; + palw_cuda_trace_record & base = semantic.base.base; + base.kernel_sequence = index; + base.decode_step = actual->decode_step; + base.m = actual->m; + base.n = actual->n; + base.k = actual->k; + base.batch = actual->batch; + base.query_tokens = actual->query_tokens; + base.key_value_tokens = actual->key_value_tokens; + base.layer_id = actual->layer_id; + base.tile_rows = actual->attention_stage == PALW_CUDA_TRACE_ATTENTION_STAGE_NONE + ? 1 + : 32; + base.tile_columns = actual->output_columns; + base.output_rows = actual->output_rows; + base.output_columns = actual->output_columns; + base.schema_version = PALW_CUDA_TRACE_SCHEMA_VERSION_V3; + base.kind = actual->kind; + base.phase = actual->phase; + base.dtype = PALW_CUDA_TRACE_DTYPE_FP32; + base.quantization = actual->quantization; + base.causal = actual->causal; + semantic.covered_schedule_index = index; + + const palw_cuda_producer_actual_identity_v3 * identity = nullptr; + palw_cuda_producer_actual_identity_v3 selected_identity{}; + const std::array * cubin_digest = nullptr; + if (actual->attention_stage == PALW_CUDA_TRACE_ATTENTION_STAGE_NONE) { + const auto * info = mixed_mmvq_info_for(*state, actual->quantization); + if (!palw_cuda_llama_mmvq_kernel_info_is_compatible_v1(info)) { + state->error = callback_error::invalid_argument; + return PALW_CUDA_TRACE_IDENTITY_MISMATCH; + } + selected_identity = info->actual_identity; + identity = &selected_identity; + cubin_digest = actual->quantization == PALW_CUDA_TRACE_QUANTIZATION_Q4_K + ? &state->q4_cubin_digest + : &state->q6_cubin_digest; + semantic.base.declared_origin = PALW_CUDA_TRACE_ORIGIN_PRODUCER_ACCUMULATOR; + semantic.base.accumulator_stage = + PALW_CUDA_TRACE_ACCUMULATOR_STAGE_FULL_K_PRE_EPILOGUE; + semantic.base.accumulator_dtype = PALW_CUDA_TRACE_ACCUMULATOR_DTYPE_FP32; + semantic.base.sketch_scheme = PALW_CUDA_TRACE_SKETCH_SCALAR_ACCUMULATOR_F32_V1; + semantic.base.reduction_segment_count = 1; + semantic.sublaunch_count = 1; + semantic.grid_x = actual->output_columns; + semantic.grid_y = 1; + semantic.grid_z = 1; + semantic.block_x = 32; + semantic.block_y = 4; + semantic.block_z = 1; + auto & binding = expected_launch->kernel_binding; + binding.sm_arch = info->sm_arch; + binding.binary_version = info->binary_version; + binding.ptx_version = info->ptx_version; + binding.num_regs = info->num_regs; + binding.max_threads_per_block = info->max_threads_per_block; + binding.static_shared_memory_bytes = info->static_shared_memory_bytes; + binding.local_memory_bytes = info->local_memory_bytes; + } else { + const auto * info = mixed_attention_info_for(*state, actual->attention_stage); + if (!palw_cuda_llama_attention_kernel_info_is_compatible_v1(info) || + info->attention_stage != actual->attention_stage) { + state->error = callback_error::invalid_argument; + return PALW_CUDA_TRACE_IDENTITY_MISMATCH; + } + selected_identity = info->actual_identity; + identity = &selected_identity; + cubin_digest = actual->attention_stage == + PALW_CUDA_TRACE_ATTENTION_STAGE_EAGER_QK_SCORE_MMVF + ? &state->qk_cubin_digest + : actual->attention_stage == + PALW_CUDA_TRACE_ATTENTION_STAGE_EAGER_MASKED_SCALED_SOFTMAX + ? &state->softmax_cubin_digest + : &state->pv_cubin_digest; + semantic.base.declared_origin = PALW_CUDA_TRACE_ORIGIN_FINAL_OUTPUT; + semantic.base.accumulator_stage = PALW_CUDA_TRACE_ACCUMULATOR_STAGE_NONE; + semantic.base.accumulator_dtype = PALW_CUDA_TRACE_ACCUMULATOR_DTYPE_NONE; + semantic.base.sketch_scheme = PALW_CUDA_TRACE_SKETCH_FINAL_OUTPUT_TILE_F32_V1; + semantic.base.reduction_segment_count = 0; + semantic.layer_present = 1; + semantic.attention_group_present = 1; + semantic.attention_stage = actual->attention_stage; + semantic.grouping_version = PALW_CUDA_TRACE_GROUPING_VERSION_V1; + semantic.sublaunch_index = actual->attention_stage - 1U; + semantic.sublaunch_count = 3; + semantic.attention_owner_schedule_index = + actual->attention_stage == + PALW_CUDA_TRACE_ATTENTION_STAGE_EAGER_QK_SCORE_MMVF + ? index + 1 + : actual->attention_stage == + PALW_CUDA_TRACE_ATTENTION_STAGE_EAGER_MASKED_SCALED_SOFTMAX + ? index + : index - 1; + semantic.query_heads = actual->query_heads; + semantic.key_value_heads = actual->key_value_heads; + semantic.head_dim = actual->head_dim; + semantic.logical_batch = actual->logical_batch; + semantic.physical_key_value_tokens = actual->physical_key_value_tokens; + semantic.grid_x = info->dimensions.grid_x; + semantic.grid_y = info->dimensions.grid_y; + semantic.grid_z = info->dimensions.grid_z; + semantic.block_x = info->dimensions.block_x; + semantic.block_y = info->dimensions.block_y; + semantic.block_z = info->dimensions.block_z; + semantic.dynamic_shared_memory_bytes = static_cast( + info->dimensions.dynamic_shared_memory_bytes); + derive_diagnostic_id_v2( + "misaka.palw.diagnostic.attention_group.v1", + semantic.attention_owner_schedule_index, + nullptr, + semantic.attention_group_instance_id); + auto & binding = expected_launch->kernel_binding; + binding.sm_arch = info->sm_arch; + binding.binary_version = info->binary_version; + binding.ptx_version = info->ptx_version; + binding.num_regs = info->num_regs; + binding.max_threads_per_block = info->max_threads_per_block; + binding.static_shared_memory_bytes = info->static_shared_memory_bytes; + binding.local_memory_bytes = info->local_memory_bytes; + } + std::memcpy( + semantic.base.producer_variant_id, + identity->producer_variant_id, + PALW_CUDA_PRODUCER_ID_SIZE); + std::memcpy( + semantic.work_entry_point_id, + identity->work_entry_point_id, + PALW_CUDA_PRODUCER_ID_SIZE); + std::memcpy( + semantic.capture_implementation_id, + identity->capture_implementation_id, + PALW_CUDA_PRODUCER_ID_SIZE); + derive_diagnostic_id_v2( + "misaka.palw.diagnostic.operation.v2", + index, + actual, + semantic.operation_instance_id); + if (actual->attention_stage == PALW_CUDA_TRACE_ATTENTION_STAGE_NONE) { + derive_diagnostic_id_v2( + "misaka.palw.diagnostic.plan.v2", + index, + actual, + semantic.decomposition_plan_id); + } else { + // All three stages are one grouped decomposition. The producer + // requires a group-wide plan ID while operation/layout/work IDs stay + // stage-distinct. + derive_diagnostic_id_v2( + "misaka.palw.diagnostic.attention_plan.v2", + semantic.attention_owner_schedule_index, + nullptr, + semantic.decomposition_plan_id); + } + derive_diagnostic_id_v2( + "misaka.palw.diagnostic.physical_layout.v2", + index, + actual, + semantic.physical_layout_id); + if (state->authority_ids != nullptr) { + // R13: the live callback consumes the authority-derived canonical + // operation/plan/physical-layout/attention-group identities instead of + // the self-derived diagnostic ones. Bytes are stamped verbatim; no + // computed launch field (shape, sketch, kernel attributes) is affected. + const uint8_t * row = + state->authority_ids + static_cast(index) * 128U; + std::memcpy(semantic.operation_instance_id, row, PALW_CUDA_PRODUCER_ID_SIZE); + std::memcpy( + semantic.decomposition_plan_id, row + 32, PALW_CUDA_PRODUCER_ID_SIZE); + std::memcpy( + semantic.physical_layout_id, row + 64, PALW_CUDA_PRODUCER_ID_SIZE); + std::memcpy( + semantic.attention_group_instance_id, row + 96, PALW_CUDA_PRODUCER_ID_SIZE); + } + std::memcpy( + expected_launch->kernel_binding.kernel_binary_digest, + cubin_digest->data(), + cubin_digest->size()); + + state->operation_tokens[static_cast(index)] = actual->operation_token; + ++state->approved; + return PALW_CUDA_TRACE_OK; +} + +void launch_accepted_callback_v2( + void * user_data, + const palw_cuda_llama_trace_association_v2 * actual, + const palw_cuda_producer_launch_v3 * expected_launch) noexcept { + auto * state = static_cast(user_data); + if (state == nullptr || state->candidates == nullptr || actual == nullptr || + expected_launch == nullptr || state->launch_notifications >= state->approved || + state->launch_notifications >= state->candidates->size() || + state->operation_tokens[static_cast(state->launch_notifications)] != + actual->operation_token || + !same_mixed_candidate( + (*state->candidates)[static_cast(state->launch_notifications)], + *actual) || + expected_launch->launch_nonce != state->launch_notifications + 1 || + expected_launch->semantic_template.base.base.kernel_sequence != + state->launch_notifications) { + if (state != nullptr && state->error == callback_error::none) { + state->error = callback_error::launch_notification_mismatch; + } + return; + } + ++state->launch_notifications; +} + +void fault_callback_v2( + void * user_data, + palw_cuda_trace_status status, + palw_cuda_llama_mmvq_fault_v1 fault, + const palw_cuda_llama_trace_association_v2 *) noexcept { + auto * state = static_cast(user_data); + if (state != nullptr) { + ++state->fault_callbacks; + state->last_fault_status = status; + state->last_fault = fault; + } +} + +palw_cuda_trace_status approve_callback( + void * user_data, + const palw_cuda_llama_mmvq_association_v1 * actual, + palw_cuda_producer_launch_v3 * expected_launch) noexcept { + auto * state = static_cast(user_data); + if (state == nullptr || state->candidates == nullptr || actual == nullptr || + expected_launch == nullptr || state->error != callback_error::none || + !palw_cuda_llama_mmvq_association_is_supported_v1(actual)) { + if (state != nullptr) { + state->error = callback_error::invalid_argument; + } + return PALW_CUDA_TRACE_INVALID_ARGUMENT; + } + + const uint64_t index = state->approved; + if (index >= state->candidates->size() || index >= EXPECTED_LAUNCH_COUNT) { + state->error = callback_error::unexpected_count; + return PALW_CUDA_TRACE_SEQUENCE_MISMATCH; + } + if (index == state->reject_index) { + state->rejection_triggered = true; + return PALW_CUDA_TRACE_IDENTITY_MISMATCH; + } + if (!same_candidate((*state->candidates)[static_cast(index)], *actual)) { + state->error = callback_error::association_mismatch; + return PALW_CUDA_TRACE_IDENTITY_MISMATCH; + } + for (uint64_t prior = 0; prior < index; ++prior) { + if (state->operation_tokens[static_cast(prior)] == actual->operation_token) { + state->error = callback_error::duplicate_operation_token; + return PALW_CUDA_TRACE_SEQUENCE_MISMATCH; + } + } + + const palw_cuda_llama_mmvq_kernel_info_v1 * info = + kernel_info_for(*state, actual->quantization); + if (!palw_cuda_llama_mmvq_kernel_info_is_compatible_v1(info) || + info->quantization != actual->quantization) { + state->error = callback_error::invalid_argument; + return PALW_CUDA_TRACE_IDENTITY_MISMATCH; + } + + *expected_launch = {}; + expected_launch->launch_nonce = index + 1; + palw_cuda_trace_record_v3 & semantic = expected_launch->semantic_template; + palw_cuda_trace_record & base = semantic.base.base; + base.kernel_sequence = index; + base.decode_step = 0; + base.m = actual->m; + base.n = actual->n; + base.k = actual->k; + base.batch = actual->batch; + base.layer_id = 0; + base.tile_rows = 1; + base.tile_columns = actual->output_columns; + base.output_rows = actual->output_rows; + base.output_columns = actual->output_columns; + base.schema_version = PALW_CUDA_TRACE_SCHEMA_VERSION_V3; + base.kind = PALW_CUDA_TRACE_GEMM; + base.phase = PALW_CUDA_TRACE_PREFILL; + base.dtype = PALW_CUDA_TRACE_DTYPE_FP32; + base.quantization = actual->quantization; + semantic.base.declared_origin = PALW_CUDA_TRACE_ORIGIN_PRODUCER_ACCUMULATOR; + semantic.base.accumulator_stage = + PALW_CUDA_TRACE_ACCUMULATOR_STAGE_FULL_K_PRE_EPILOGUE; + semantic.base.accumulator_dtype = PALW_CUDA_TRACE_ACCUMULATOR_DTYPE_FP32; + semantic.base.sketch_scheme = PALW_CUDA_TRACE_SKETCH_SCALAR_ACCUMULATOR_F32_V1; + semantic.base.reduction_segment_count = 1; + semantic.covered_schedule_index = index; + semantic.sublaunch_count = 1; + semantic.grid_x = actual->output_columns; + semantic.grid_y = 1; + semantic.grid_z = 1; + semantic.block_x = 32; + semantic.block_y = 4; + semantic.block_z = 1; + + std::memcpy( + semantic.base.producer_variant_id, + info->actual_identity.producer_variant_id, + PALW_CUDA_PRODUCER_ID_SIZE); + std::memcpy( + semantic.work_entry_point_id, + info->actual_identity.work_entry_point_id, + PALW_CUDA_PRODUCER_ID_SIZE); + std::memcpy( + semantic.capture_implementation_id, + info->actual_identity.capture_implementation_id, + PALW_CUDA_PRODUCER_ID_SIZE); + derive_diagnostic_id( + "misaka.palw.diagnostic.operation.v1", index, actual, semantic.operation_instance_id); + derive_diagnostic_id( + "misaka.palw.diagnostic.direct_plan.v1", 0, nullptr, semantic.decomposition_plan_id); + derive_diagnostic_id( + "misaka.palw.diagnostic.physical_layout.v1", index, actual, semantic.physical_layout_id); + + auto & binding = expected_launch->kernel_binding; + std::memcpy( + binding.kernel_binary_digest, + state->binary_digest.data(), + state->binary_digest.size()); + binding.sm_arch = info->sm_arch; + binding.binary_version = info->binary_version; + binding.ptx_version = info->ptx_version; + binding.num_regs = info->num_regs; + binding.max_threads_per_block = info->max_threads_per_block; + binding.static_shared_memory_bytes = info->static_shared_memory_bytes; + binding.local_memory_bytes = info->local_memory_bytes; + + state->operation_tokens[static_cast(index)] = actual->operation_token; + ++state->approved; + return PALW_CUDA_TRACE_OK; +} + +void launch_accepted_callback( + void * user_data, + const palw_cuda_llama_mmvq_association_v1 * actual, + const palw_cuda_producer_launch_v3 * expected_launch) noexcept { + auto * state = static_cast(user_data); + if (state == nullptr || state->candidates == nullptr || actual == nullptr || + expected_launch == nullptr || state->launch_notifications >= state->approved || + state->launch_notifications >= state->candidates->size() || + state->operation_tokens[static_cast(state->launch_notifications)] != + actual->operation_token || + !same_candidate( + (*state->candidates)[static_cast(state->launch_notifications)], *actual) || + expected_launch->launch_nonce != state->launch_notifications + 1 || + expected_launch->semantic_template.base.base.kernel_sequence != + state->launch_notifications) { + if (state != nullptr && state->error == callback_error::none) { + state->error = callback_error::launch_notification_mismatch; + } + return; + } + ++state->launch_notifications; +} + +void fault_callback( + void * user_data, + palw_cuda_trace_status status, + palw_cuda_llama_mmvq_fault_v1 fault, + const palw_cuda_llama_mmvq_association_v1 *) noexcept { + auto * state = static_cast(user_data); + if (state != nullptr) { + ++state->fault_callbacks; + state->last_fault_status = status; + state->last_fault = fault; + } +} + +llama_context_params make_context_params( + ggml_backend_sched_eval_callback callback, + void * user_data) { + llama_context_params params = llama_context_default_params(); + params.n_ctx = CONTEXT_TOKENS; + params.n_batch = 1; + params.n_ubatch = 1; + params.n_seq_max = 1; + params.n_outputs_max = 1; + params.n_threads = 1; + params.n_threads_batch = 1; + params.flash_attn_type = LLAMA_FLASH_ATTN_TYPE_DISABLED; + params.type_k = GGML_TYPE_F16; + params.type_v = GGML_TYPE_F16; + params.cb_eval = callback; + params.cb_eval_user_data = user_data; + params.embeddings = false; + params.offload_kqv = true; + params.no_perf = true; + params.op_offload = true; + params.kv_unified = false; + return params; +} + +bool valid_context_policy(llama_context * context) { + return context != nullptr && llama_n_batch(context) == 1 && + llama_n_ubatch(context) == 1 && llama_n_seq_max(context) == 1 && + llama_n_threads(context) == 1 && llama_n_threads_batch(context) == 1; +} + +bool get_metadata_value( + const llama_model * model, + const char * key, + std::string & output) { + std::array buffer{}; + const int32_t length = llama_model_meta_val_str(model, key, buffer.data(), buffer.size()); + if (length < 0 || static_cast(length) >= buffer.size()) { + return false; + } + output.assign(buffer.data(), static_cast(length)); + return true; +} + +bool valid_qwen3_8b_q4_k_m(const llama_model * model) { + std::string architecture; + return model != nullptr && get_metadata_value(model, "general.architecture", architecture) && + architecture == "qwen3" && llama_model_ftype(model) == LLAMA_FTYPE_MOSTLY_Q4_K_M && + llama_model_n_layer(model) == 36 && llama_model_n_embd(model) == 4096 && + llama_model_n_head(model) == 32 && llama_model_n_head_kv(model) == 8 && + llama_vocab_n_tokens(llama_model_get_vocab(model)) == 151936 && + !llama_model_has_encoder(model) && llama_model_has_decoder(model); +} + +llama_token select_token(const llama_model * model) { + llama_token token = llama_model_decoder_start_token(model); + const int32_t vocab_size = llama_vocab_n_tokens(llama_model_get_vocab(model)); + if (token < 0 || token >= vocab_size) { + token = llama_vocab_bos(llama_model_get_vocab(model)); + } + return token >= 0 && token < vocab_size ? token : 0; +} + +uint8_t * write_u16_be(uint8_t * output, uint16_t value) { + *output++ = static_cast(value >> 8); + *output++ = static_cast(value); + return output; +} + +uint8_t * write_u32_be(uint8_t * output, uint32_t value) { + for (int shift = 24; shift >= 0; shift -= 8) { + *output++ = static_cast(value >> shift); + } + return output; +} + +uint8_t * write_u64_be(uint8_t * output, uint64_t value) { + for (int shift = 56; shift >= 0; shift -= 8) { + *output++ = static_cast(value >> shift); + } + return output; +} + +uint8_t * write_id(uint8_t * output, const uint8_t * identity) { + std::memcpy(output, identity, PALW_CUDA_PRODUCER_ID_SIZE); + return output + PALW_CUDA_PRODUCER_ID_SIZE; +} + +bool encode_record_v3( + const palw_cuda_trace_record_v3 & record, + std::array & encoded) { + uint8_t * output = encoded.data(); + const palw_cuda_trace_record & base = record.base.base; + output = write_u64_be(output, base.kernel_sequence); + output = write_u64_be(output, base.tile_index); + output = write_u64_be(output, base.decode_step); + output = write_u64_be(output, base.m); + output = write_u64_be(output, base.n); + output = write_u64_be(output, base.k); + output = write_u64_be(output, base.batch); + output = write_u64_be(output, base.query_tokens); + output = write_u64_be(output, base.key_value_tokens); + for (uint64_t lane : base.sketch) { + output = write_u64_be(output, lane); + } + output = write_u32_be(output, base.layer_id); + output = write_u32_be(output, base.tile_row); + output = write_u32_be(output, base.tile_column); + output = write_u32_be(output, base.tile_rows); + output = write_u32_be(output, base.tile_columns); + output = write_u32_be(output, base.output_rows); + output = write_u32_be(output, base.output_columns); + output = write_u16_be(output, base.schema_version); + *output++ = base.kind; + *output++ = base.phase; + *output++ = base.dtype; + *output++ = base.quantization; + *output++ = base.causal; + *output++ = record.base.declared_origin; + *output++ = record.base.accumulator_stage; + *output++ = record.base.accumulator_dtype; + output = write_u16_be(output, record.base.sketch_scheme); + output = write_u32_be(output, record.base.reduction_segment_index); + output = write_u32_be(output, record.base.reduction_segment_count); + output = write_id(output, record.base.producer_variant_id); + output = write_u64_be(output, record.covered_schedule_index); + output = write_u64_be(output, record.attention_owner_schedule_index); + output = write_u32_be(output, record.sublaunch_index); + output = write_u32_be(output, record.sublaunch_count); + *output++ = record.layer_present; + *output++ = record.attention_group_present; + *output++ = record.attention_stage; + *output++ = record.grouping_version; + output = write_u32_be(output, record.query_heads); + output = write_u32_be(output, record.key_value_heads); + output = write_u32_be(output, record.head_dim); + output = write_u32_be(output, record.logical_batch); + output = write_u32_be(output, record.physical_key_value_tokens); + output = write_u32_be(output, record.grid_x); + output = write_u32_be(output, record.grid_y); + output = write_u32_be(output, record.grid_z); + output = write_u32_be(output, record.block_x); + output = write_u32_be(output, record.block_y); + output = write_u32_be(output, record.block_z); + output = write_u32_be(output, record.dynamic_shared_memory_bytes); + output = write_id(output, record.operation_instance_id); + output = write_id(output, record.attention_group_instance_id); + output = write_id(output, record.decomposition_plan_id); + output = write_id(output, record.physical_layout_id); + output = write_id(output, record.work_entry_point_id); + output = write_id(output, record.capture_implementation_id); + return static_cast(output - encoded.data()) == encoded.size(); +} + +bool verify_record( + const palw_cuda_trace_record_v3 & record, + uint64_t index, + const candidate & expected, + const palw_cuda_llama_mmvq_kernel_info_v1 & info) { + const palw_cuda_trace_record & base = record.base.base; + return base.kernel_sequence == index && record.covered_schedule_index == index && + base.tile_index == 0 && base.tile_row == 0 && base.tile_column == 0 && + base.m == expected.m && base.n == expected.n && base.k == expected.k && + base.batch == expected.batch && base.quantization == expected.quantization && + base.output_rows == 1 && base.output_columns == expected.m && + std::memcmp( + record.base.producer_variant_id, + info.actual_identity.producer_variant_id, + PALW_CUDA_PRODUCER_ID_SIZE) == 0 && + std::memcmp( + record.work_entry_point_id, + info.actual_identity.work_entry_point_id, + PALW_CUDA_PRODUCER_ID_SIZE) == 0 && + std::memcmp( + record.capture_implementation_id, + info.actual_identity.capture_implementation_id, + PALW_CUDA_PRODUCER_ID_SIZE) == 0; +} + +void print_kernel_info(const palw_cuda_llama_mmvq_kernel_info_v1 & info) { + std::printf( + "{\"sm\":%u,\"binary\":%d,\"ptx\":%d,\"regs\":%d," + "\"max_threads\":%d,\"static_shared\":%llu,\"local\":%llu}", + info.sm_arch, + info.binary_version, + info.ptx_version, + info.num_regs, + info.max_threads_per_block, + static_cast(info.static_shared_memory_bytes), + static_cast(info.local_memory_bytes)); +} + +palw_cuda_llama_tensor_view_v1 make_synthetic_view( + uint8_t type, + const void * data, + uint64_t ne0, + uint64_t ne1, + uint64_t ne2, + uint64_t ne3) { + palw_cuda_llama_tensor_view_v1 view{}; + view.data = data; + view.type = type; + view.ne[0] = ne0; + view.ne[1] = ne1; + view.ne[2] = ne2; + view.ne[3] = ne3; + const uint64_t size = type == PALW_CUDA_LLAMA_TENSOR_TYPE_F16 ? 2 : 4; + view.nb[0] = size; + view.nb[1] = size * ne0; + view.nb[2] = view.nb[1] * ne1; + view.nb[3] = view.nb[2] * ne2; + return view; +} + +palw_cuda_llama_trace_association_v2 make_synthetic_attention(uint8_t stage) { + static int tokens[3]{}; + static int allocations[9]{}; + const size_t index = static_cast(stage - 1U); + palw_cuda_llama_trace_association_v2 actual{}; + actual.abi_version = PALW_CUDA_LLAMA_TRACE_ASSOCIATION_ABI_VERSION_V2; + actual.struct_size = sizeof(actual); + actual.operation_token = &tokens[index]; + actual.decode_step = 0; + actual.n = 1; + actual.batch = 32; + actual.output_rows = 32; + actual.layer_id = 0; + actual.query_tokens = 1; + actual.key_value_tokens = 1; + actual.physical_key_value_tokens = 256; + actual.query_heads = 32; + actual.key_value_heads = 8; + actual.head_dim = 128; + actual.logical_batch = 1; + actual.attention_stage = stage; + actual.kind = stage == + PALW_CUDA_TRACE_ATTENTION_STAGE_EAGER_MASKED_SCALED_SOFTMAX + ? PALW_CUDA_TRACE_ATTENTION + : PALW_CUDA_TRACE_GEMM; + actual.layer_present = 1; + actual.phase = PALW_CUDA_TRACE_PREFILL; + actual.causal = 1; + if (stage == PALW_CUDA_TRACE_ATTENTION_STAGE_EAGER_QK_SCORE_MMVF) { + actual.m = 256; + actual.k = 128; + actual.output_columns = 256; + actual.src0 = make_synthetic_view( + PALW_CUDA_LLAMA_TENSOR_TYPE_F16, &allocations[0], 128, 256, 8, 1); + actual.src1 = make_synthetic_view( + PALW_CUDA_LLAMA_TENSOR_TYPE_F32, &allocations[1], 128, 1, 32, 1); + actual.dst = make_synthetic_view( + PALW_CUDA_LLAMA_TENSOR_TYPE_F32, &allocations[2], 256, 1, 32, 1); + actual.precision = PALW_CUDA_LLAMA_WORK_PRECISION_F32; + } else if (stage == + PALW_CUDA_TRACE_ATTENTION_STAGE_EAGER_MASKED_SCALED_SOFTMAX) { + actual.m = 256; + actual.k = 1; + actual.output_columns = 256; + actual.src0 = make_synthetic_view( + PALW_CUDA_LLAMA_TENSOR_TYPE_F32, &allocations[3], 256, 1, 32, 1); + actual.src1 = make_synthetic_view( + PALW_CUDA_LLAMA_TENSOR_TYPE_F32, &allocations[4], 256, 1, 1, 1); + actual.dst = make_synthetic_view( + PALW_CUDA_LLAMA_TENSOR_TYPE_F32, &allocations[5], 256, 1, 32, 1); + actual.scale_bits = 0x3db504f3U; + actual.attention_mask = PALW_CUDA_LLAMA_ATTENTION_MASK_F32_CAUSAL; + actual.precision = PALW_CUDA_LLAMA_WORK_PRECISION_F32; + } else { + actual.m = 128; + actual.k = 256; + actual.output_columns = 128; + actual.src0 = make_synthetic_view( + PALW_CUDA_LLAMA_TENSOR_TYPE_F16, &allocations[6], 256, 128, 8, 1); + actual.src1 = make_synthetic_view( + PALW_CUDA_LLAMA_TENSOR_TYPE_F32, &allocations[7], 256, 1, 32, 1); + actual.dst = make_synthetic_view( + PALW_CUDA_LLAMA_TENSOR_TYPE_F32, &allocations[8], 128, 1, 32, 1); + actual.precision = PALW_CUDA_LLAMA_WORK_PRECISION_DEFAULT; + } + return actual; +} + +bool verify_mixed_record( + const palw_cuda_trace_record_v3 & record, + uint64_t index, + const mixed_candidate & expected, + const mixed_trace_state & state) { + const palw_cuda_trace_record & base = record.base.base; + if (base.kernel_sequence != index || record.covered_schedule_index != index || + base.tile_index != 0 || base.tile_row != 0 || base.tile_column != 0 || + base.m != expected.m || base.n != expected.n || base.k != expected.k || + base.batch != expected.batch || base.kind != expected.kind || + base.quantization != expected.quantization || + base.output_rows != expected.output_rows || + base.output_columns != expected.output_columns) { + return false; + } + const palw_cuda_producer_actual_identity_v3 * identity = nullptr; + palw_cuda_producer_actual_identity_v3 selected{}; + if (expected.attention_stage == PALW_CUDA_TRACE_ATTENTION_STAGE_NONE) { + const auto * info = mixed_mmvq_info_for(state, expected.quantization); + if (info == nullptr) { + return false; + } + selected = info->actual_identity; + identity = &selected; + if (record.attention_group_present != 0 || record.sublaunch_count != 1) { + return false; + } + } else { + const auto * info = mixed_attention_info_for(state, expected.attention_stage); + if (info == nullptr) { + return false; + } + selected = info->actual_identity; + identity = &selected; + if (record.layer_present != 1 || base.layer_id != expected.layer_id || + record.attention_group_present != 1 || + record.attention_stage != expected.attention_stage || + record.sublaunch_index + 1U != expected.attention_stage || + record.sublaunch_count != 3 || record.query_heads != 32 || + record.key_value_heads != 8 || record.head_dim != 128 || + record.logical_batch != 1 || record.physical_key_value_tokens != 256 || + record.grid_x != info->dimensions.grid_x || + record.grid_y != info->dimensions.grid_y || + record.block_x != info->dimensions.block_x) { + return false; + } + } + return std::memcmp( + record.base.producer_variant_id, + identity->producer_variant_id, + PALW_CUDA_PRODUCER_ID_SIZE) == 0 && + std::memcmp( + record.work_entry_point_id, + identity->work_entry_point_id, + PALW_CUDA_PRODUCER_ID_SIZE) == 0 && + std::memcmp( + record.capture_implementation_id, + identity->capture_implementation_id, + PALW_CUDA_PRODUCER_ID_SIZE) == 0; +} + +void print_id(const uint8_t id[PALW_CUDA_PRODUCER_ID_SIZE]) { + for (size_t i = 0; i < PALW_CUDA_PRODUCER_ID_SIZE; ++i) { + std::printf("%02x", static_cast(id[i])); + } +} + +void print_actual_identity(const palw_cuda_producer_actual_identity_v3 & identity) { + std::printf("{\"producer_variant_id\":\""); + print_id(identity.producer_variant_id); + std::printf("\",\"work_entry_point_id\":\""); + print_id(identity.work_entry_point_id); + std::printf("\",\"capture_implementation_id\":\""); + print_id(identity.capture_implementation_id); + std::printf("\"}"); +} + +void print_attention_kernel_info(const palw_cuda_llama_attention_kernel_info_v1 & info) { + std::printf( + "{\"sm_arch\":%u,\"binary_version\":%d,\"ptx_version\":%d," + "\"num_regs\":%d,\"max_threads_per_block\":%d," + "\"static_shared_memory_bytes\":%llu,\"local_memory_bytes\":%llu," + "\"actual_identity\":", + info.sm_arch, + info.binary_version, + info.ptx_version, + info.num_regs, + info.max_threads_per_block, + static_cast(info.static_shared_memory_bytes), + static_cast(info.local_memory_bytes)); + print_actual_identity(info.actual_identity); + std::printf("}"); +} + +void print_collector_info(const palw_cuda_producer_grouped_capture_info_v3 & info) { + std::printf( + "{\"sm_arch\":%u,\"binary_version\":%d,\"ptx_version\":%d," + "\"num_regs\":%d,\"max_threads_per_block\":%d," + "\"static_shared_memory_bytes\":%llu,\"local_memory_bytes\":%llu," + "\"capture_implementation_id\":\"", + info.sm_arch, + info.binary_version, + info.ptx_version, + info.num_regs, + info.max_threads_per_block, + static_cast(info.static_shared_memory_bytes), + static_cast(info.local_memory_bytes)); + print_id(info.capture_implementation_id); + std::printf("\"}"); +} + +void print_mmvq_kernel_info_exact( + const palw_cuda_llama_mmvq_kernel_info_v1 & info, + const char * symbol) { + std::printf( + "{\"exact_mangled_symbol\":\"%s\",\"sm_arch\":%u," + "\"binary_version\":%d,\"ptx_version\":%d,\"num_regs\":%d," + "\"max_threads_per_block\":%d,\"static_shared_memory_bytes\":%llu," + "\"local_memory_bytes\":%llu,\"actual_identity\":", + symbol, + info.sm_arch, + info.binary_version, + info.ptx_version, + info.num_regs, + info.max_threads_per_block, + static_cast(info.static_shared_memory_bytes), + static_cast(info.local_memory_bytes)); + print_actual_identity(info.actual_identity); + std::printf("}"); +} + +void print_attention_kernel_info_exact( + const palw_cuda_llama_attention_kernel_info_v1 & info, + const char * symbol) { + std::printf("{\"exact_mangled_symbol\":\"%s\",\"kernel_info\":", symbol); + print_attention_kernel_info(info); + std::printf("}"); +} + +struct backend_guard { + ~backend_guard() { + llama_backend_free(); + } +}; + +using model_ptr = std::unique_ptr; +using context_ptr = std::unique_ptr; + +int fail(const char * message, int code = 3) { + std::fprintf(stderr, "error: %s\n", message); + return code; +} + +int run_mixed( + const options & opts, + llama_model * model, + llama_token token) { + mixed_discovery_state discovery; + discovery.candidates.reserve(EXPECTED_MIXED_LAUNCH_COUNT); + { + context_ptr context( + llama_init_from_model( + model, make_context_params(discover_mixed_callback, &discovery)), + llama_free); + if (!valid_context_policy(context.get())) { + return fail("failed to create the mixed discovery context"); + } + llama_token mutable_token = token; + if (llama_decode(context.get(), llama_batch_get_one(&mutable_token, 1)) != 0 || + discovery.failed) { + return fail("mixed one-token discovery decode failed"); + } + } + + uint64_t q4_count = 0; + uint64_t q6_count = 0; + std::array stage_counts{}; + std::array first_index = { + UINT64_MAX, UINT64_MAX, UINT64_MAX, UINT64_MAX, UINT64_MAX, + }; + for (size_t candidate_index = 0; + candidate_index < discovery.candidates.size(); ++candidate_index) { + const mixed_candidate & value = discovery.candidates[candidate_index]; + const uint64_t index = static_cast(candidate_index); + if (value.attention_stage == PALW_CUDA_TRACE_ATTENTION_STAGE_NONE) { + if (value.m == 0 || value.n != 1 || value.k == 0 || value.batch != 1 || + value.output_rows != 1 || value.output_columns != value.m) { + return fail("mixed discovery produced unsupported MMVQ geometry"); + } + if (value.quantization == PALW_CUDA_TRACE_QUANTIZATION_Q4_K) { + ++q4_count; + first_index[0] = std::min(first_index[0], index); + } else if (value.quantization == PALW_CUDA_TRACE_QUANTIZATION_Q6_K) { + ++q6_count; + first_index[1] = std::min(first_index[1], index); + } else { + return fail("mixed discovery produced unsupported quantization"); + } + } else if (value.attention_stage < stage_counts.size()) { + ++stage_counts[value.attention_stage]; + first_index[static_cast(value.attention_stage) + 1U] = std::min( + first_index[static_cast(value.attention_stage) + 1U], index); + } else { + return fail("mixed discovery produced unsupported attention stage"); + } + } + if (discovery.candidates.size() != EXPECTED_MIXED_LAUNCH_COUNT || + discovery.attention_seen != EXPECTED_ATTENTION_COUNT || + q4_count != EXPECTED_Q4_K_COUNT || q6_count != EXPECTED_Q6_K_COUNT || + stage_counts[1] != 36 || stage_counts[2] != 36 || stage_counts[3] != 36 || + std::find(first_index.begin(), first_index.end(), UINT64_MAX) != + first_index.end()) { + std::fprintf(stderr, + "error: mixed discovery expected 361 = 253 + 108, got %zu = %llu + %llu + %u\n", + discovery.candidates.size(), + static_cast(q4_count), + static_cast(q6_count), + discovery.attention_seen); + return 3; + } + + mixed_trace_state state; + state.candidates = &discovery.candidates; + state.q4_cubin_digest = opts.q4_cubin_digest; + state.q6_cubin_digest = opts.q6_cubin_digest; + state.qk_cubin_digest = opts.qk_cubin_digest; + state.softmax_cubin_digest = opts.softmax_cubin_digest; + state.pv_cubin_digest = opts.pv_cubin_digest; + state.reject_index = opts.reject_index; + std::vector authority_ids_buffer; + if (!opts.authority_ids_path.empty()) { + std::FILE * ids_in = std::fopen(opts.authority_ids_path.c_str(), "rb"); + if (ids_in == nullptr) { + return fail("failed to open authority identity table", 5); + } + authority_ids_buffer.resize( + static_cast(EXPECTED_MIXED_LAUNCH_COUNT) * 128U); + const size_t got = std::fread( + authority_ids_buffer.data(), 1, authority_ids_buffer.size(), ids_in); + const bool at_eof = std::fgetc(ids_in) == EOF; + (void) std::fclose(ids_in); + if (got != authority_ids_buffer.size() || !at_eof) { + return fail( + "authority identity table must be exactly 361*128 bytes", 5); + } + state.authority_ids = authority_ids_buffer.data(); + } + context_ptr traced( + llama_init_from_model(model, make_context_params(nullptr, nullptr)), + llama_free); + if (!valid_context_policy(traced.get())) { + return fail("failed to create the mixed traced context"); + } + + palw_cuda_llama_trace_request_v2 request{}; + request.abi_version = PALW_CUDA_LLAMA_TRACE_REQUEST_ABI_VERSION_V2; + request.struct_size = sizeof(request); + request.expected_launch_count = EXPECTED_MIXED_LAUNCH_COUNT; + request.user_data = &state; + request.approve = approve_callback_v2; + request.launch_accepted = launch_accepted_callback_v2; + request.fault = fault_callback_v2; + if (llama_palw_cuda_trace_attach_v2( + traced.get(), EXPECTED_MIXED_LAUNCH_COUNT, &request) != PALW_CUDA_TRACE_OK) { + return fail("failed to attach the mixed PALW producer"); + } + const auto destroy_trace = [&]() { + return llama_palw_cuda_trace_destroy(traced.get()); + }; + + const auto qk = make_synthetic_attention( + PALW_CUDA_TRACE_ATTENTION_STAGE_EAGER_QK_SCORE_MMVF); + const auto softmax = make_synthetic_attention( + PALW_CUDA_TRACE_ATTENTION_STAGE_EAGER_MASKED_SCALED_SOFTMAX); + const auto pv = make_synthetic_attention( + PALW_CUDA_TRACE_ATTENTION_STAGE_EAGER_VALUE_AGGREGATION_MMVF); + if (llama_palw_cuda_trace_kernel_info( + traced.get(), PALW_CUDA_TRACE_QUANTIZATION_Q4_K, &state.q4_info) != + PALW_CUDA_TRACE_OK || + llama_palw_cuda_trace_kernel_info( + traced.get(), PALW_CUDA_TRACE_QUANTIZATION_Q6_K, &state.q6_info) != + PALW_CUDA_TRACE_OK || + llama_palw_cuda_trace_attention_kernel_info( + traced.get(), &qk, &state.qk_info) != PALW_CUDA_TRACE_OK || + llama_palw_cuda_trace_attention_kernel_info( + traced.get(), &softmax, &state.softmax_info) != PALW_CUDA_TRACE_OK || + llama_palw_cuda_trace_attention_kernel_info( + traced.get(), &pv, &state.pv_info) != PALW_CUDA_TRACE_OK || + llama_palw_cuda_trace_grouped_capture_info( + traced.get(), &state.collector_info) != PALW_CUDA_TRACE_OK || + state.collector_info.abi_version != + PALW_CUDA_PRODUCER_GROUPED_CAPTURE_INFO_ABI_VERSION_V3 || + state.collector_info.struct_size != sizeof(state.collector_info) || + std::memcmp( + state.collector_info.capture_implementation_id, + state.qk_info.actual_identity.capture_implementation_id, + PALW_CUDA_PRODUCER_ID_SIZE) != 0) { + uint64_t ignored = 0; + (void) llama_palw_cuda_trace_unbind(traced.get(), &ignored); + (void) destroy_trace(); + return fail("failed to query mixed work/collector kernel diagnostics"); + } + + llama_token mutable_token = token; + const int32_t decode_status = + llama_decode(traced.get(), llama_batch_get_one(&mutable_token, 1)); + uint64_t accepted = 0; + const palw_cuda_trace_status unbind_status = + llama_palw_cuda_trace_unbind(traced.get(), &accepted); + uint64_t count = 0; + uint64_t committed = 0; + uint32_t faults = 0; + int poisoned = 0; + const palw_cuda_trace_status diagnostics_status = + llama_palw_cuda_trace_diagnostics( + traced.get(), &count, &committed, &faults, &poisoned); + + if (opts.reject_index != UINT64_MAX) { + const bool expected_failure = decode_status != 0 && + unbind_status != PALW_CUDA_TRACE_OK && + diagnostics_status == PALW_CUDA_TRACE_OK && state.rejection_triggered && + state.approved == opts.reject_index && accepted == opts.reject_index && + count == opts.reject_index && + state.launch_notifications == opts.reject_index && committed == 0 && + state.fault_callbacks == 1 && + state.last_fault_status == PALW_CUDA_TRACE_IDENTITY_MISMATCH && + state.last_fault == PALW_CUDA_LLAMA_MMVQ_FAULT_ASSOCIATION_MISMATCH && + faults == 0 && poisoned == 0; + const palw_cuda_trace_status destroy_status = destroy_trace(); + if (!expected_failure || destroy_status != PALW_CUDA_TRACE_OK) { + return fail("mixed negative smoke did not stop at requested index", 5); + } + std::printf( + "{\"schema\":\"misaka.palw.mixed_hook_smoke\",\"version\":2," + "\"status\":\"expected_failure\",\"reject_index\":%llu," + "\"rejected_class\":\"%s\",\"accepted\":%llu," + "\"records\":%llu,\"poisoned\":false}\n", + static_cast(opts.reject_index), + discovery.candidates[static_cast(opts.reject_index)].attention_stage == + PALW_CUDA_TRACE_ATTENTION_STAGE_EAGER_QK_SCORE_MMVF + ? "attention_qk" + : discovery.candidates[static_cast(opts.reject_index)].attention_stage == + PALW_CUDA_TRACE_ATTENTION_STAGE_EAGER_MASKED_SCALED_SOFTMAX + ? "attention_softmax" + : discovery.candidates[static_cast(opts.reject_index)].attention_stage == + PALW_CUDA_TRACE_ATTENTION_STAGE_EAGER_VALUE_AGGREGATION_MMVF + ? "attention_pv" + : discovery.candidates[static_cast(opts.reject_index)].quantization == + PALW_CUDA_TRACE_QUANTIZATION_Q4_K + ? "q4_k" + : "q6_k", + static_cast(accepted), + static_cast(count)); + return 0; + } + + if (decode_status != 0 || unbind_status != PALW_CUDA_TRACE_OK || + diagnostics_status != PALW_CUDA_TRACE_OK || + state.error != callback_error::none || + state.approved != EXPECTED_MIXED_LAUNCH_COUNT || + state.launch_notifications != EXPECTED_MIXED_LAUNCH_COUNT || + state.fault_callbacks != 0 || accepted != EXPECTED_MIXED_LAUNCH_COUNT || + count != EXPECTED_MIXED_LAUNCH_COUNT || committed != 0 || + faults != 0 || poisoned != 0) { + (void) destroy_trace(); + return fail("positive mixed decode or pre-finalize diagnostics failed", 5); + } + + std::vector records(EXPECTED_MIXED_LAUNCH_COUNT); + if (llama_palw_cuda_trace_finalize( + traced.get(), records.data(), records.size()) != PALW_CUDA_TRACE_OK || + llama_palw_cuda_trace_diagnostics( + traced.get(), &count, &committed, &faults, &poisoned) != + PALW_CUDA_TRACE_OK || + count != EXPECTED_MIXED_LAUNCH_COUNT || + committed != EXPECTED_MIXED_LAUNCH_COUNT || faults != 0 || poisoned != 0) { + (void) destroy_trace(); + return fail("mixed producer finalize failed", 5); + } + uint64_t record_fingerprint = FNV64_OFFSET; + std::array encoded{}; + std::vector emit_buffer; + if (!opts.emit_records_path.empty()) { + emit_buffer.reserve(EXPECTED_MIXED_LAUNCH_COUNT * encoded.size()); + } + for (uint64_t index = 0; index < EXPECTED_MIXED_LAUNCH_COUNT; ++index) { + if (!verify_mixed_record( + records[static_cast(index)], + index, + discovery.candidates[static_cast(index)], + state) || + !encode_record_v3(records[static_cast(index)], encoded)) { + (void) destroy_trace(); + return fail("mixed record order, identity, or encoding mismatch", 5); + } + record_fingerprint = fnv_bytes( + record_fingerprint, encoded.data(), encoded.size()); + if (!opts.emit_records_path.empty()) { + emit_buffer.insert(emit_buffer.end(), encoded.begin(), encoded.end()); + } + } + if (destroy_trace() != PALW_CUDA_TRACE_OK || request.producer_context != nullptr) { + return fail("failed to destroy mixed producer", 5); + } + if (!opts.emit_records_path.empty()) { + std::FILE * emit_out = std::fopen(opts.emit_records_path.c_str(), "wb"); + if (emit_out == nullptr || + std::fwrite(emit_buffer.data(), 1, emit_buffer.size(), emit_out) != + emit_buffer.size() || + std::fclose(emit_out) != 0) { + return fail("failed to write emitted record stream", 5); + } + } + + std::printf( + "{\"schema\":\"misaka.palw.mixed_hook_smoke\",\"version\":2," + "\"status\":\"ok\",\"diagnostic_only\":true," + "\"receipt_authority\":false,\"model\":\"qwen3-8b-q4_k_m\"," + "\"token\":%d,\"fa\":\"off\",\"launches\":361," + "\"mmvq\":253,\"attention\":108,\"q4_k\":216,\"q6_k\":37," + "\"attention_qk\":36,\"attention_softmax\":36,\"attention_pv\":36," + "\"first_index\":{\"q4_k\":%llu,\"q6_k\":%llu," + "\"attention_qk\":%llu,\"attention_softmax\":%llu," + "\"attention_pv\":%llu}," + "\"accepted\":%llu,\"records\":%llu,\"committed\":%llu," + "\"canonical_record_fnv1a64\":\"%016llx\",\"entry_cubin_sha256\":{" + "\"q4_k\":\"%s\",\"q6_k\":\"%s\",\"attention_qk\":\"%s\"," + "\"attention_softmax\":\"%s\",\"attention_pv\":\"%s\"}," + "\"kernel_info\":{\"q4_k\":", + token, + static_cast(first_index[0]), + static_cast(first_index[1]), + static_cast(first_index[2]), + static_cast(first_index[3]), + static_cast(first_index[4]), + static_cast(accepted), + static_cast(count), + static_cast(committed), + static_cast(record_fingerprint), + opts.q4_cubin_digest_hex.c_str(), + opts.q6_cubin_digest_hex.c_str(), + opts.qk_cubin_digest_hex.c_str(), + opts.softmax_cubin_digest_hex.c_str(), + opts.pv_cubin_digest_hex.c_str()); + print_mmvq_kernel_info_exact(state.q4_info, Q4_EXACT_SYMBOL); + std::printf(",\"q6_k\":"); + print_mmvq_kernel_info_exact(state.q6_info, Q6_EXACT_SYMBOL); + std::printf(",\"attention_qk\":"); + print_attention_kernel_info_exact(state.qk_info, QK_EXACT_SYMBOL); + std::printf(",\"attention_softmax\":"); + print_attention_kernel_info_exact(state.softmax_info, SOFTMAX_EXACT_SYMBOL); + std::printf(",\"attention_pv\":"); + print_attention_kernel_info_exact(state.pv_info, PV_EXACT_SYMBOL); + std::printf(",\"grouped_collector\":"); + print_collector_info(state.collector_info); + std::printf("}}\n"); + return 0; +} + +} // namespace + +int main(int argc, char ** argv) { + options opts; + bool help = false; + if (!parse_options(argc, argv, opts, help)) { + print_usage(argv[0]); + return 2; + } + if (help) { + print_usage(argv[0]); + return 0; + } + + llama_log_set(log_callback, nullptr); + llama_backend_init(); + backend_guard backend_cleanup; + ggml_backend_load_all(); + + llama_model_params model_params = llama_model_default_params(); + model_params.n_gpu_layers = std::numeric_limits::max(); + model_params.split_mode = LLAMA_SPLIT_MODE_NONE; + model_params.main_gpu = 0; + model_params.use_mmap = true; + model_params.use_mlock = false; + model_params.check_tensors = true; + model_params.use_extra_bufts = false; + + model_ptr model(llama_model_load_from_file(opts.model_path.c_str(), model_params), llama_model_free); + if (!model) { + return fail("failed to load model"); + } + if (!valid_qwen3_8b_q4_k_m(model.get())) { + return fail("model is not the dense Qwen3-8B Q4_K_M profile"); + } + const llama_token token = select_token(model.get()); + + if (opts.mixed) { + return run_mixed(opts, model.get(), token); + } + + discovery_state discovery; + discovery.candidates.reserve(EXPECTED_LAUNCH_COUNT); + { + context_ptr context( + llama_init_from_model(model.get(), make_context_params(discover_callback, &discovery)), + llama_free); + if (!valid_context_policy(context.get())) { + return fail("failed to create the one-token discovery context"); + } + llama_token mutable_token = token; + if (llama_decode(context.get(), llama_batch_get_one(&mutable_token, 1)) != 0 || + discovery.failed) { + return fail("one-token discovery decode failed"); + } + } + + uint64_t q4_count = 0; + uint64_t q6_count = 0; + for (const candidate & value : discovery.candidates) { + if (value.weight_data == nullptr || value.m == 0 || value.n != 1 || value.k == 0 || + value.batch != 1 || value.m > UINT32_MAX) { + return fail("discovery produced an unsupported quantized GEMM shape"); + } + if (value.quantization == PALW_CUDA_TRACE_QUANTIZATION_Q4_K) { + ++q4_count; + } else if (value.quantization == PALW_CUDA_TRACE_QUANTIZATION_Q6_K) { + ++q6_count; + } else { + return fail("discovery produced a quantization outside Q4_K/Q6_K"); + } + } + if (discovery.candidates.size() != EXPECTED_LAUNCH_COUNT || + q4_count != EXPECTED_Q4_K_COUNT || q6_count != EXPECTED_Q6_K_COUNT) { + std::fprintf(stderr, + "error: expected 253 candidates (216 Q4_K, 37 Q6_K), got %zu (%llu, %llu)\n", + discovery.candidates.size(), + static_cast(q4_count), + static_cast(q6_count)); + return 3; + } + + trace_state state; + state.candidates = &discovery.candidates; + state.binary_digest = opts.binary_digest; + state.reject_index = opts.reject_index; + context_ptr traced( + llama_init_from_model(model.get(), make_context_params(nullptr, nullptr)), + llama_free); + if (!valid_context_policy(traced.get())) { + return fail("failed to create the one-token traced context"); + } + + palw_cuda_llama_mmvq_request_v1 request{}; + request.abi_version = PALW_CUDA_LLAMA_MMVQ_REQUEST_ABI_VERSION_V1; + request.struct_size = sizeof(request); + request.expected_launch_count = EXPECTED_LAUNCH_COUNT; + request.user_data = &state; + request.approve = approve_callback; + request.launch_accepted = launch_accepted_callback; + request.fault = fault_callback; + + if (llama_palw_cuda_trace_attach( + traced.get(), EXPECTED_LAUNCH_COUNT, &request) != PALW_CUDA_TRACE_OK) { + return fail("failed to attach the request-local PALW producer"); + } + const auto destroy_trace = [&]() { + return llama_palw_cuda_trace_destroy(traced.get()); + }; + + if (llama_palw_cuda_trace_kernel_info( + traced.get(), PALW_CUDA_TRACE_QUANTIZATION_Q4_K, &state.q4_info) != + PALW_CUDA_TRACE_OK || + llama_palw_cuda_trace_kernel_info( + traced.get(), PALW_CUDA_TRACE_QUANTIZATION_Q6_K, &state.q6_info) != + PALW_CUDA_TRACE_OK || + !palw_cuda_llama_mmvq_kernel_info_is_compatible_v1(&state.q4_info) || + !palw_cuda_llama_mmvq_kernel_info_is_compatible_v1(&state.q6_info)) { + uint64_t ignored = 0; + (void) llama_palw_cuda_trace_unbind(traced.get(), &ignored); + (void) destroy_trace(); + return fail("failed to query exact Q4_K/Q6_K kernel diagnostics"); + } + + llama_token mutable_token = token; + const int32_t decode_status = + llama_decode(traced.get(), llama_batch_get_one(&mutable_token, 1)); + uint64_t accepted = 0; + const palw_cuda_trace_status unbind_status = + llama_palw_cuda_trace_unbind(traced.get(), &accepted); + + uint64_t count = 0; + uint64_t committed = 0; + uint32_t faults = 0; + int poisoned = 0; + const palw_cuda_trace_status diagnostics_status = llama_palw_cuda_trace_diagnostics( + traced.get(), &count, &committed, &faults, &poisoned); + + if (opts.reject_index != UINT64_MAX) { + const bool expected_failure = decode_status != 0 && + unbind_status != PALW_CUDA_TRACE_OK && + diagnostics_status == PALW_CUDA_TRACE_OK && state.rejection_triggered && + state.approved == opts.reject_index && accepted == opts.reject_index && + count == opts.reject_index && state.launch_notifications == opts.reject_index && + committed == 0 && state.fault_callbacks == 1 && + state.last_fault_status == PALW_CUDA_TRACE_IDENTITY_MISMATCH && + state.last_fault == PALW_CUDA_LLAMA_MMVQ_FAULT_ASSOCIATION_MISMATCH && + faults == 0 && poisoned == 0; + const palw_cuda_trace_status destroy_status = destroy_trace(); + if (!expected_failure || destroy_status != PALW_CUDA_TRACE_OK) { + return fail("negative fail-closed smoke did not stop at the requested index", 5); + } + std::printf( + "{\"schema\":\"misaka.palw.mmvq_hook_smoke\",\"version\":1," + "\"status\":\"expected_failure\",\"diagnostic_only\":true," + "\"receipt_authority\":false,\"reject_index\":%llu," + "\"decode_status\":%d,\"unbind_status\":%d,\"accepted\":%llu," + "\"records\":%llu,\"fault_callbacks\":%llu,\"producer_faults\":%u," + "\"poisoned\":%s}\n", + static_cast(opts.reject_index), + decode_status, + static_cast(unbind_status), + static_cast(accepted), + static_cast(count), + static_cast(state.fault_callbacks), + faults, + poisoned != 0 ? "true" : "false"); + return 0; + } + + if (decode_status != 0 || unbind_status != PALW_CUDA_TRACE_OK || + diagnostics_status != PALW_CUDA_TRACE_OK || + state.error != callback_error::none || state.approved != EXPECTED_LAUNCH_COUNT || + state.launch_notifications != EXPECTED_LAUNCH_COUNT || state.fault_callbacks != 0 || + accepted != EXPECTED_LAUNCH_COUNT || count != EXPECTED_LAUNCH_COUNT || + committed != 0 || faults != 0 || poisoned != 0) { + (void) destroy_trace(); + return fail("positive traced decode or pre-finalize diagnostics failed", 5); + } + + std::vector records(EXPECTED_LAUNCH_COUNT); + if (llama_palw_cuda_trace_finalize( + traced.get(), records.data(), records.size()) != PALW_CUDA_TRACE_OK || + llama_palw_cuda_trace_diagnostics( + traced.get(), &count, &committed, &faults, &poisoned) != PALW_CUDA_TRACE_OK || + count != EXPECTED_LAUNCH_COUNT || committed != EXPECTED_LAUNCH_COUNT || + faults != 0 || poisoned != 0) { + (void) destroy_trace(); + return fail("producer finalize or post-finalize diagnostics failed", 5); + } + + uint64_t record_fingerprint = FNV64_OFFSET; + std::array encoded{}; + for (uint64_t index = 0; index < EXPECTED_LAUNCH_COUNT; ++index) { + const candidate & expected = discovery.candidates[static_cast(index)]; + const auto * info = kernel_info_for(state, expected.quantization); + if (info == nullptr || !verify_record( + records[static_cast(index)], index, expected, *info) || + !encode_record_v3(records[static_cast(index)], encoded)) { + (void) destroy_trace(); + return fail("finalized record order, identity, or canonical encoding mismatch", 5); + } + record_fingerprint = fnv_bytes(record_fingerprint, encoded.data(), encoded.size()); + } + + if (destroy_trace() != PALW_CUDA_TRACE_OK || request.producer_context != nullptr) { + return fail("failed to destroy the request-local producer", 5); + } + + std::printf( + "{\"schema\":\"misaka.palw.mmvq_hook_smoke\",\"version\":1," + "\"status\":\"ok\",\"diagnostic_only\":true,\"receipt_authority\":false," + "\"model\":\"qwen3-8b-q4_k_m\",\"token\":%d,\"fa\":\"off\"," + "\"discovered\":%llu,\"q4_k\":%llu,\"q6_k\":%llu," + "\"approved\":%llu,\"accepted\":%llu,\"records\":%llu," + "\"committed\":%llu,\"producer_faults\":%u,\"poisoned\":false," + "\"canonical_record_fnv1a64\":\"%016llx\"," + "\"kernel_binary_sha256\":\"%s\",\"q4_kernel\":", + token, + static_cast(discovery.candidates.size()), + static_cast(q4_count), + static_cast(q6_count), + static_cast(state.approved), + static_cast(accepted), + static_cast(count), + static_cast(committed), + faults, + static_cast(record_fingerprint), + opts.binary_digest_hex.c_str()); + print_kernel_info(state.q4_info); + std::printf(",\"q6_kernel\":"); + print_kernel_info(state.q6_info); + std::printf("}\n"); + return 0; +} diff --git a/tools/palw-observer/CMakeLists.txt b/tools/palw-observer/CMakeLists.txt new file mode 100644 index 00000000..40e3ffd8 --- /dev/null +++ b/tools/palw-observer/CMakeLists.txt @@ -0,0 +1,9 @@ +set(TARGET llama-palw-observer) + +add_executable(${TARGET} palw-observer.cpp) +target_link_libraries(${TARGET} PRIVATE llama llama-common-base ${CMAKE_THREAD_LIBS_INIT}) +target_compile_features(${TARGET} PRIVATE cxx_std_17) + +if(LLAMA_TOOLS_INSTALL) + install(TARGETS ${TARGET} RUNTIME) +endif() diff --git a/tools/palw-observer/README.md b/tools/palw-observer/README.md new file mode 100644 index 00000000..2c7731b1 --- /dev/null +++ b/tools/palw-observer/README.md @@ -0,0 +1,61 @@ +# PALW native graph observer + +`llama-palw-observer` is a non-interactive, single-request Qwen3-8B runner. It +uses only the public llama and ggml APIs and writes versioned JSONL records to +stdout. Runtime and model logs are written to stderr. + +## Build and run + +```sh +cmake --build build-palw --target llama-palw-observer -j +build-palw/bin/llama-palw-observer \ + --model /path/to/Qwen3-8B-Q4_K_M.gguf \ + --prompt "Hello" \ + --n-predict 8 \ + --observer graph +``` + +The runner fixes greedy sampling, request batch, logical/physical token batch, +parallel sequences, tensor parallelism, and CPU thread counts to one. It also +disables context shifting, speculation, Flash Attention, and split-model tensor +parallelism. The prompt is therefore evaluated one token at a time. The context +is sized before execution and the run fails instead of shifting when the prompt +and prediction bound exceed the model training context. + +The accepted model profile is dense Qwen3-8B: architecture `qwen3`, 36 layers, +4096 hidden elements, 32 query heads, 8 KV heads, and a 151936-token vocabulary. +Other profiles fail before a header or inference result is emitted. + +## Observer modes + +- `off` installs no scheduler callback. +- `graph` emits metadata for every graph node from the callback's `ask` stage. + It never requests tensor data. +- `sketch` emits the same metadata and requests post-compute data only for + `MUL_MAT` and `MUL_MAT_ID` nodes. It copies at most the first 64 contiguous + output elements through `ggml_backend_tensor_get`. Each element becomes one + hexadecimal nibble: one sign bit and a fixed three-bit magnitude bucket. The + resulting 64 hex digits are a 256-bit sketch. Raw activation values are never + serialized. + +Nibble bit 3 is the sign bit. Bits 0-2 use these absolute-value buckets: zero, +`(0, 2^-8)`, `[2^-8, 2^-4)`, `[2^-4, 2^-2)`, `[2^-2, 1)`, `[1, 4)`, +`[4, 16)`, and `[16, infinity)`. Fewer than 64 available elements are padded +with zero nibbles. + +Every stdout line has `schema`, `schema_version`, and `record`. A successful run +contains one `header`, zero or more `event` records, and one `result`. The result +contains prompt token IDs, sampled token IDs (including an EOG token when one is +sampled), the non-special generated output bytes, and the stop reason. + +## Limits + +This target is a runtime-observation prototype, not a PALW receipt generator. +It does not calculate canonical compute units, sign receipts, form commitments, +or trace CUDA kernels. The sketch tile is explicitly identified as +`graph_fallback_logical_prefix_v1`, and `kernel_trace.claim` is +`not_a_cuda_kernel_trace`. Scheduler callbacks add synchronization and can +change timing, but they must not change token IDs or output bytes. The sketch is +lossy observation data, not a cryptographic proof. Graph topology and floating +point results can vary across backends or hardware, so byte identity is only +claimed for repeated runs of the same pinned model, runtime, and backend class. diff --git a/tools/palw-observer/palw-observer.cpp b/tools/palw-observer/palw-observer.cpp new file mode 100644 index 00000000..5e2646b1 --- /dev/null +++ b/tools/palw-observer/palw-observer.cpp @@ -0,0 +1,1243 @@ +#include "ggml-backend.h" +#include "ggml.h" +#include "llama.h" +#include "build-info.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace { + +constexpr const char * SCHEMA_NAME = "misaka.palw.runtime_observer"; +constexpr int SCHEMA_VERSION = 1; +constexpr size_t SKETCH_SAMPLES = 64; +constexpr uint32_t PALW_CONTEXT_TOKENS = 4096; + +enum class observer_mode { + off, + graph, + sketch, +}; + +struct options { + std::string model_path; + std::string prompt; + int32_t n_predict = -1; + int32_t n_gpu_layers = 999; + observer_mode mode = observer_mode::off; + bool prompt_set = false; + bool prompt_stdin = false; + bool emit_output_bytes = false; +}; + +struct model_metadata { + std::vector> entries; +}; + +static const char * mode_name(observer_mode mode) { + switch (mode) { + case observer_mode::off: return "off"; + case observer_mode::graph: return "graph"; + case observer_mode::sketch: return "sketch"; + } + return "invalid"; +} + +static void log_callback(ggml_log_level level, const char * text, void *) { + const char * label = "unknown"; + switch (level) { + case GGML_LOG_LEVEL_NONE: label = "none"; break; + case GGML_LOG_LEVEL_DEBUG: label = "debug"; break; + case GGML_LOG_LEVEL_INFO: label = "info"; break; + case GGML_LOG_LEVEL_WARN: label = "warn"; break; + case GGML_LOG_LEVEL_ERROR: label = "error"; break; + case GGML_LOG_LEVEL_CONT: label = "cont"; break; + } + std::fprintf(stderr, "[llama:%s] %s", label, text ? text : ""); +} + +static void print_usage(const char * argv0) { + std::fprintf(stderr, + "usage: %s --model MODEL (--prompt TEXT|--prompt-stdin) --n-predict N " + "[--observer off|graph|sketch] [--n-gpu-layers N] [--emit-output-bytes]\n", + argv0); +} + +static bool parse_i32(const char * value, int32_t min_value, int32_t max_value, int32_t & result) { + if (!value || value[0] == '\0') { + return false; + } + int32_t parsed = 0; + const char * end = value + std::strlen(value); + const auto converted = std::from_chars(value, end, parsed); + if (converted.ec != std::errc() || converted.ptr != end || parsed < min_value || parsed > max_value) { + return false; + } + result = parsed; + return true; +} + +static bool take_value(int argc, char ** argv, int & i, const char * option, const char *& value) { + if (i + 1 >= argc) { + std::fprintf(stderr, "error: %s requires a value\n", option); + return false; + } + value = argv[++i]; + return true; +} + +static bool parse_options(int argc, char ** argv, options & opts, bool & help, bool & version) { + help = false; + version = false; + for (int i = 1; i < argc; ++i) { + const std::string_view arg(argv[i]); + const char * value = nullptr; + if (arg == "--help" || arg == "-h") { + help = true; + return true; + } else if (arg == "--version") { + version = true; + return true; + } else if (arg == "--model" || arg == "-m") { + if (!take_value(argc, argv, i, argv[i], value)) { + return false; + } + opts.model_path = value; + } else if (arg == "--prompt" || arg == "-p") { + if (opts.prompt_stdin) { + std::fprintf(stderr, "error: --prompt and --prompt-stdin are mutually exclusive\n"); + return false; + } + if (!take_value(argc, argv, i, argv[i], value)) { + return false; + } + opts.prompt = value; + opts.prompt_set = true; + } else if (arg == "--prompt-stdin") { + if (opts.prompt_set || opts.prompt_stdin) { + std::fprintf(stderr, "error: prompt input may be selected only once\n"); + return false; + } + opts.prompt_stdin = true; + } else if (arg == "--emit-output-bytes") { + opts.emit_output_bytes = true; + } else if (arg == "--n-predict" || arg == "-n") { + if (!take_value(argc, argv, i, argv[i], value) || + !parse_i32(value, 0, 65536, opts.n_predict)) { + std::fprintf(stderr, "error: --n-predict must be an integer in [0, 65536]\n"); + return false; + } + } else if (arg == "--n-gpu-layers" || arg == "-ngl") { + if (!take_value(argc, argv, i, argv[i], value) || + !parse_i32(value, -1, 100000, opts.n_gpu_layers)) { + std::fprintf(stderr, "error: --n-gpu-layers must be an integer in [-1, 100000]\n"); + return false; + } + } else if (arg == "--observer") { + if (!take_value(argc, argv, i, argv[i], value)) { + return false; + } + const std::string_view mode(value); + if (mode == "off") { + opts.mode = observer_mode::off; + } else if (mode == "graph") { + opts.mode = observer_mode::graph; + } else if (mode == "sketch") { + opts.mode = observer_mode::sketch; + } else { + std::fprintf(stderr, "error: --observer must be off, graph, or sketch\n"); + return false; + } + } else { + std::fprintf(stderr, "error: unknown argument: %s\n", argv[i]); + return false; + } + } + + if (opts.model_path.empty()) { + std::fprintf(stderr, "error: --model is required\n"); + return false; + } + if (opts.prompt_stdin) { + std::array buffer = {}; + while (true) { + const size_t count = std::fread(buffer.data(), 1, buffer.size(), stdin); + if (count != 0) { + if (opts.prompt.size() > static_cast(std::numeric_limits::max()) - count) { + std::fprintf(stderr, "error: stdin prompt exceeds the tokenizer API bound\n"); + return false; + } + opts.prompt.append(buffer.data(), count); + } + if (count != buffer.size()) { + if (std::ferror(stdin)) { + std::fprintf(stderr, "error: failed to read prompt from stdin\n"); + return false; + } + break; + } + } + opts.prompt_set = true; + } + if (!opts.prompt_set) { + std::fprintf(stderr, "error: --prompt is required (an explicitly empty prompt is allowed)\n"); + return false; + } + if (opts.n_predict < 0) { + std::fprintf(stderr, "error: --n-predict is required\n"); + return false; + } + if (opts.prompt.size() > static_cast(std::numeric_limits::max())) { + std::fprintf(stderr, "error: prompt exceeds the tokenizer API bound\n"); + return false; + } + return true; +} + +static void append_json_string(std::string & out, std::string_view value) { + static constexpr char hex[] = "0123456789abcdef"; + out.push_back('"'); + for (const unsigned char c : value) { + switch (c) { + case '"': out += "\\\""; break; + case '\\': out += "\\\\"; break; + case '\b': out += "\\b"; break; + case '\f': out += "\\f"; break; + case '\n': out += "\\n"; break; + case '\r': out += "\\r"; break; + case '\t': out += "\\t"; break; + default: + if (c < 0x20) { + out += "\\u00"; + out.push_back(hex[c >> 4]); + out.push_back(hex[c & 0x0f]); + } else { + out.push_back(static_cast(c)); + } + break; + } + } + out.push_back('"'); +} + +static void append_bool(std::string & out, bool value) { + out += value ? "true" : "false"; +} + +template +static void append_integer(std::string & out, T value) { + out += std::to_string(value); +} + +static void append_float(std::string & out, double value) { + if (!std::isfinite(value)) { + out += "null"; + return; + } + char buffer[64]; + const int n = std::snprintf(buffer, sizeof(buffer), "%.9g", value); + if (n <= 0 || static_cast(n) >= sizeof(buffer)) { + out += "null"; + return; + } + out.append(buffer, static_cast(n)); +} + +static bool write_json_line(const std::string & line) { + if (std::fwrite(line.data(), 1, line.size(), stdout) != line.size() || + std::fputc('\n', stdout) == EOF || std::fflush(stdout) != 0) { + std::fprintf(stderr, "error: failed to write JSONL to stdout\n"); + return false; + } + return true; +} + +static void append_record_prefix(std::string & out, const char * record) { + out += "{\"schema\":\""; + out += SCHEMA_NAME; + out += "\",\"schema_version\":"; + append_integer(out, SCHEMA_VERSION); + out += ",\"record\":"; + append_json_string(out, record); +} + +static size_t bounded_name_length(const char * name) { + size_t length = 0; + while (length < GGML_MAX_NAME && name[length] != '\0') { + ++length; + } + return length; +} + +static std::string_view tensor_name(const ggml_tensor * tensor) { + if (!tensor) { + return {}; + } + return std::string_view(tensor->name, bounded_name_length(tensor->name)); +} + +static std::string lowercase(std::string_view value) { + std::string result; + result.reserve(value.size()); + for (const unsigned char c : value) { + if (c >= 'A' && c <= 'Z') { + result.push_back(static_cast(c - 'A' + 'a')); + } else { + result.push_back(static_cast(c)); + } + } + return result; +} + +static bool has_text(std::string_view value, std::string_view needle) { + return value.find(needle) != std::string_view::npos; +} + +static void collect_related_names(const ggml_tensor * tensor, std::vector & names) { + if (!tensor) { + return; + } + if (!tensor_name(tensor).empty()) { + names.push_back(tensor_name(tensor)); + } + if (tensor->view_src && !tensor_name(tensor->view_src).empty()) { + names.push_back(tensor_name(tensor->view_src)); + } + for (int i = 0; i < GGML_MAX_SRC; ++i) { + const ggml_tensor * src = tensor->src[i]; + if (!src) { + continue; + } + if (!tensor_name(src).empty()) { + names.push_back(tensor_name(src)); + } + if (src->view_src && !tensor_name(src->view_src).empty()) { + names.push_back(tensor_name(src->view_src)); + } + } +} + +static bool related_names_contain(const ggml_tensor * tensor, std::string_view needle) { + std::vector names; + collect_related_names(tensor, names); + for (const auto name : names) { + if (has_text(lowercase(name), needle)) { + return true; + } + } + return false; +} + +static bool parse_decimal_at(std::string_view name, size_t offset, int & value) { + if (offset >= name.size() || name[offset] < '0' || name[offset] > '9') { + return false; + } + uint64_t parsed = 0; + size_t i = offset; + while (i < name.size() && name[i] >= '0' && name[i] <= '9') { + parsed = parsed * 10 + static_cast(name[i] - '0'); + if (parsed > static_cast(std::numeric_limits::max())) { + return false; + } + ++i; + } + value = static_cast(parsed); + return true; +} + +static bool parse_layer_from_name(std::string_view name, int & layer) { + const size_t dash = name.rfind('-'); + if (dash != std::string_view::npos && dash + 1 < name.size()) { + int parsed = -1; + if (parse_decimal_at(name, dash + 1, parsed)) { + size_t end = dash + 1; + while (end < name.size() && name[end] >= '0' && name[end] <= '9') { + ++end; + } + if (end == name.size()) { + layer = parsed; + return true; + } + } + } + + const size_t block = name.find("blk."); + if (block != std::string_view::npos && parse_decimal_at(name, block + 4, layer)) { + return true; + } + + if (has_text(lowercase(name), "cache_")) { + const size_t marker = name.rfind("_l"); + if (marker != std::string_view::npos && parse_decimal_at(name, marker + 2, layer)) { + return true; + } + } + return false; +} + +static int tensor_layer(const ggml_tensor * tensor) { + int layer = -1; + if (parse_layer_from_name(tensor_name(tensor), layer)) { + return layer; + } + std::vector names; + collect_related_names(tensor, names); + for (const auto name : names) { + if (parse_layer_from_name(name, layer)) { + return layer; + } + } + return -1; +} + +static void append_shape(std::string & out, const ggml_tensor * tensor) { + out.push_back('['); + for (int i = 0; i < GGML_MAX_DIMS; ++i) { + if (i != 0) { + out.push_back(','); + } + append_integer(out, tensor->ne[i]); + } + out.push_back(']'); +} + +static void append_token_ids(std::string & out, const std::vector & tokens) { + out.push_back('['); + for (size_t i = 0; i < tokens.size(); ++i) { + if (i != 0) { + out.push_back(','); + } + append_integer(out, tokens[i]); + } + out.push_back(']'); +} + +class graph_observer { +public: + graph_observer(observer_mode mode, int32_t n_layers) : mode_(mode), n_layers_(n_layers) {} + + void set_phase(const char * phase, int64_t step) { + phase_ = phase; + phase_step_ = step; + } + + bool failed() const { + return failed_; + } + + const std::string & error() const { + return error_; + } + + uint64_t event_count() const { + return event_sequence_; + } + + static bool callback(ggml_tensor * tensor, bool ask, void * user_data) noexcept { + auto * observer = static_cast(user_data); + try { + return observer->on_tensor(tensor, ask); + } catch (const std::exception & exception) { + observer->fail(std::string("observer exception: ") + exception.what()); + } catch (...) { + observer->fail("observer exception: unknown"); + } + return false; + } + +private: + observer_mode mode_; + int32_t n_layers_; + std::string phase_ = "uninitialized"; + int64_t phase_step_ = -1; + uint64_t event_sequence_ = 0; + bool failed_ = false; + std::string error_; + + void fail(std::string message) { + if (!failed_) { + failed_ = true; + error_ = std::move(message); + std::fprintf(stderr, "observer error: %s\n", error_.c_str()); + } + } + + bool validate_tensor(const ggml_tensor * tensor, bool validate_op) { + if (!tensor) { + fail("null graph tensor"); + return false; + } + if (bounded_name_length(tensor->name) == GGML_MAX_NAME) { + fail("graph tensor name is not terminated"); + return false; + } + const int type = static_cast(tensor->type); + if (type < 0 || type >= static_cast(GGML_TYPE_COUNT)) { + fail("graph tensor type is outside public ggml bounds"); + return false; + } + if (validate_op) { + const int op = static_cast(tensor->op); + if (op < 0 || op >= static_cast(GGML_OP_COUNT)) { + fail("graph op is outside public ggml bounds"); + return false; + } + } + + uint64_t elements = 1; + for (int i = 0; i < GGML_MAX_DIMS; ++i) { + if (tensor->ne[i] <= 0) { + fail("graph tensor has a non-positive dimension"); + return false; + } + const uint64_t dimension = static_cast(tensor->ne[i]); + if (elements > static_cast(std::numeric_limits::max()) / dimension) { + fail("graph tensor element count overflows int64"); + return false; + } + elements *= dimension; + } + if (elements != static_cast(ggml_nelements(tensor)) || ggml_nbytes(tensor) == 0) { + fail("graph tensor storage bounds are inconsistent"); + return false; + } + return true; + } + + bool validate_node(const ggml_tensor * tensor) { + if (!validate_tensor(tensor, true)) { + return false; + } + for (int i = 0; i < GGML_MAX_SRC; ++i) { + if (tensor->src[i] && !validate_tensor(tensor->src[i], false)) { + return false; + } + } + const int layer = tensor_layer(tensor); + if (layer >= n_layers_) { + fail("graph layer id exceeds model layer count"); + return false; + } + if (is_gemm(tensor)) { + if (!tensor->src[0] || !tensor->src[1]) { + fail("GEMM node is missing an input"); + return false; + } + if (tensor->src[0]->ne[0] != tensor->src[1]->ne[0] || + tensor->ne[0] != tensor->src[0]->ne[1] || + tensor->ne[1] != tensor->src[1]->ne[1]) { + fail("GEMM dimensions are inconsistent"); + return false; + } + } + return true; + } + + static bool is_gemm(const ggml_tensor * tensor) { + return tensor->op == GGML_OP_MUL_MAT || tensor->op == GGML_OP_MUL_MAT_ID; + } + + std::vector categories(const ggml_tensor * tensor) const { + std::vector result; + if (is_gemm(tensor)) { + result.emplace_back("gemm"); + } + if (tensor->op == GGML_OP_NORM || tensor->op == GGML_OP_RMS_NORM || + tensor->op == GGML_OP_GROUP_NORM || tensor->op == GGML_OP_L2_NORM) { + result.emplace_back("norm"); + } + if (tensor->op == GGML_OP_ROPE || tensor->op == GGML_OP_ROPE_BACK) { + result.emplace_back("rope"); + } + if (tensor->op == GGML_OP_FLASH_ATTN_EXT || tensor->op == GGML_OP_FLASH_ATTN_BACK || + related_names_contain(tensor, "attn") || related_names_contain(tensor, "qcur") || + related_names_contain(tensor, "kcur") || related_names_contain(tensor, "vcur") || + related_names_contain(tensor, "kq")) { + result.emplace_back("attention"); + } + if (related_names_contain(tensor, "cache_") || related_names_contain(tensor, "_cache")) { + result.emplace_back("kv_cache"); + } + if (result.empty()) { + result.emplace_back("other"); + } + return result; + } + + static int magnitude_bucket(float value) { + const float magnitude = std::fabs(value); + if (magnitude == 0.0f) return 0; + if (magnitude < 0.00390625f) return 1; + if (magnitude < 0.0625f) return 2; + if (magnitude < 0.25f) return 3; + if (magnitude < 1.0f) return 4; + if (magnitude < 4.0f) return 5; + if (magnitude < 16.0f) return 6; + return 7; + } + + bool make_sketch(const ggml_tensor * tensor, std::string & sketch, size_t & copied_bytes, size_t & sampled) { + if (!tensor->buffer || !ggml_is_contiguous(tensor)) { + fail("GEMM output is not backed by a contiguous public backend tensor"); + return false; + } + + size_t element_size = 0; + switch (tensor->type) { + case GGML_TYPE_F32: element_size = sizeof(float); break; + case GGML_TYPE_F16: element_size = sizeof(ggml_fp16_t); break; + case GGML_TYPE_BF16: element_size = sizeof(ggml_bf16_t); break; + default: + fail("GEMM output type is unsupported by sign/bucket sketch v1"); + return false; + } + + const uint64_t n_elements = static_cast(ggml_nelements(tensor)); + sampled = static_cast(std::min(SKETCH_SAMPLES, n_elements)); + if (sampled > std::numeric_limits::max() / element_size) { + fail("GEMM sketch copy bound overflows size_t"); + return false; + } + copied_bytes = sampled * element_size; + if (copied_bytes > ggml_nbytes(tensor)) { + fail("GEMM sketch copy exceeds tensor storage"); + return false; + } + + std::vector host(copied_bytes); + ggml_backend_tensor_get(tensor, host.data(), 0, copied_bytes); + + static constexpr char hex[] = "0123456789abcdef"; + sketch.assign(SKETCH_SAMPLES, '0'); + for (size_t i = 0; i < sampled; ++i) { + float value = 0.0f; + if (tensor->type == GGML_TYPE_F32) { + std::memcpy(&value, host.data() + i * element_size, sizeof(value)); + } else if (tensor->type == GGML_TYPE_F16) { + ggml_fp16_t packed; + std::memcpy(&packed, host.data() + i * element_size, sizeof(packed)); + value = ggml_fp16_to_fp32(packed); + } else { + ggml_bf16_t packed; + std::memcpy(&packed, host.data() + i * element_size, sizeof(packed)); + value = ggml_bf16_to_fp32(packed); + } + if (!std::isfinite(value)) { + fail("GEMM sketch encountered a non-finite output"); + return false; + } + const unsigned nibble = (std::signbit(value) ? 8u : 0u) | + static_cast(magnitude_bucket(value)); + sketch[i] = hex[nibble]; + } + return true; + } + + void append_tensor_metadata(std::string & out, const ggml_tensor * tensor) const { + out += "\"name\":"; + append_json_string(out, tensor_name(tensor)); + out += ",\"op\":"; + append_json_string(out, ggml_op_name(tensor->op)); + out += ",\"type\":"; + append_json_string(out, ggml_type_name(tensor->type)); + out += ",\"shape\":"; + append_shape(out, tensor); + out += ",\"n_bytes\":"; + append_integer(out, ggml_nbytes(tensor)); + } + + bool emit_event(const ggml_tensor * tensor, const char * stage, + const std::string * sketch, size_t copied_bytes, size_t sampled) { + std::string line; + line.reserve(2048); + append_record_prefix(line, "event"); + line += ",\"event_seq\":"; + append_integer(line, event_sequence_++); + line += ",\"observer\":"; + append_json_string(line, mode_name(mode_)); + line += ",\"phase\":"; + append_json_string(line, phase_); + line += ",\"phase_step\":"; + append_integer(line, phase_step_); + line += ",\"stage\":"; + append_json_string(line, stage); + + const int layer = tensor_layer(tensor); + line += ",\"layer\":"; + if (layer < 0) { + line += "null"; + } else { + append_integer(line, layer); + } + + line += ",\"categories\":["; + const auto node_categories = categories(tensor); + for (size_t i = 0; i < node_categories.size(); ++i) { + if (i != 0) { + line.push_back(','); + } + append_json_string(line, node_categories[i]); + } + line += "],\"tensor\":{"; + append_tensor_metadata(line, tensor); + line += "},\"sources\":["; + bool first_source = true; + for (int i = 0; i < GGML_MAX_SRC; ++i) { + if (!tensor->src[i]) { + continue; + } + if (!first_source) { + line.push_back(','); + } + first_source = false; + line.push_back('{'); + append_tensor_metadata(line, tensor->src[i]); + line.push_back('}'); + } + line.push_back(']'); + + if (is_gemm(tensor)) { + line += ",\"gemm\":{\"variant\":\"ggml_graph_op_v1\",\"m\":"; + append_integer(line, tensor->src[0]->ne[1]); + line += ",\"n\":"; + append_integer(line, tensor->src[1]->ne[1]); + line += ",\"k\":"; + append_integer(line, tensor->src[0]->ne[0]); + line += ",\"batch_shape\":["; + append_integer(line, tensor->ne[2]); + line.push_back(','); + append_integer(line, tensor->ne[3]); + line += "]}"; + } + + if (sketch) { + line += ",\"sketch\":{\"version\":\"sign_bucket_256_v1\",\"bits\":256,"; + line += "\"encoding\":\"hex\",\"probe\":\"contiguous_prefix_64_v1\","; + line += "\"sample_count\":"; + append_integer(line, sampled); + line += ",\"copied_bytes\":"; + append_integer(line, copied_bytes); + line += ",\"value\":"; + append_json_string(line, *sketch); + line += "},\"tile\":{\"variant\":\"graph_fallback_logical_prefix_v1\","; + line += "\"linear_offset\":0,\"linear_elements\":"; + append_integer(line, sampled); + line += "},\"kernel_trace\":{\"available\":false,\"backend\":\"none\","; + line += "\"claim\":\"not_a_cuda_kernel_trace\"}"; + } + + line.push_back('}'); + if (!write_json_line(line)) { + fail("stdout JSONL write failed"); + return false; + } + return true; + } + + bool on_tensor(ggml_tensor * tensor, bool ask) { + if (failed_ || mode_ == observer_mode::off) { + return false; + } + if (!validate_node(tensor)) { + return false; + } + + if (ask) { + if (mode_ == observer_mode::graph) { + emit_event(tensor, "ask_metadata", nullptr, 0, 0); + return false; + } + if (is_gemm(tensor)) { + return true; + } + emit_event(tensor, "ask_metadata", nullptr, 0, 0); + return false; + } + + if (mode_ != observer_mode::sketch || !is_gemm(tensor)) { + fail("unexpected post-compute callback"); + return false; + } + + std::string sketch; + size_t copied_bytes = 0; + size_t sampled = 0; + if (!make_sketch(tensor, sketch, copied_bytes, sampled)) { + return false; + } + return emit_event(tensor, "post_compute_sketch", &sketch, copied_bytes, sampled); + } +}; + +static bool get_model_string(const llama_model * model, int32_t index, bool key, std::string & value) { + std::vector buffer(256); + for (int attempt = 0; attempt < 3; ++attempt) { + const int32_t length = key + ? llama_model_meta_key_by_index(model, index, buffer.data(), buffer.size()) + : llama_model_meta_val_str_by_index(model, index, buffer.data(), buffer.size()); + if (length < 0) { + return false; + } + if (static_cast(length) < buffer.size()) { + value.assign(buffer.data(), static_cast(length)); + return true; + } + buffer.resize(static_cast(length) + 1); + } + return false; +} + +static bool read_model_metadata(const llama_model * model, model_metadata & metadata) { + const int32_t count = llama_model_meta_count(model); + if (count < 0 || count > 100000) { + std::fprintf(stderr, "error: model metadata count is outside bounds\n"); + return false; + } + metadata.entries.reserve(static_cast(count)); + for (int32_t i = 0; i < count; ++i) { + std::string key; + std::string value; + if (!get_model_string(model, i, true, key) || !get_model_string(model, i, false, value)) { + std::fprintf(stderr, "error: failed to read model metadata index %d\n", i); + return false; + } + metadata.entries.emplace_back(std::move(key), std::move(value)); + } + return true; +} + +static const std::string * find_metadata(const model_metadata & metadata, const char * key) { + for (const auto & entry : metadata.entries) { + if (entry.first == key) { + return &entry.second; + } + } + return nullptr; +} + +static bool validate_qwen3_8b_profile(const llama_model * model, const model_metadata & metadata) { + const std::string * architecture = find_metadata(metadata, "general.architecture"); + const llama_vocab * vocab = llama_model_get_vocab(model); + if (!architecture || *architecture != "qwen3" || + llama_model_n_layer(model) != 36 || + llama_model_n_embd(model) != 4096 || + llama_model_n_head(model) != 32 || + llama_model_n_head_kv(model) != 8 || + llama_vocab_n_tokens(vocab) != 151936 || + llama_model_has_encoder(model) || !llama_model_has_decoder(model)) { + std::fprintf(stderr, + "error: model is not the supported dense Qwen3-8B profile " + "(qwen3, 36 layers, 4096 hidden, 32/8 heads, 151936 vocab)\n"); + return false; + } + return true; +} + +static const char * device_type_name(enum ggml_backend_dev_type type) { + switch (type) { + case GGML_BACKEND_DEVICE_TYPE_CPU: return "cpu"; + case GGML_BACKEND_DEVICE_TYPE_GPU: return "gpu"; + case GGML_BACKEND_DEVICE_TYPE_IGPU: return "igpu"; + case GGML_BACKEND_DEVICE_TYPE_ACCEL: return "accelerator"; + case GGML_BACKEND_DEVICE_TYPE_META: return "meta"; + } + return "unknown"; +} + +static bool emit_header(const options & opts, const llama_model * model, const llama_context * context, + const model_metadata & metadata, size_t prompt_tokens) { + char description[1024] = {}; + const int32_t description_length = llama_model_desc(model, description, sizeof(description)); + if (description_length < 0 || static_cast(description_length) >= sizeof(description)) { + std::fprintf(stderr, "error: model description exceeds header bound\n"); + return false; + } + + std::string line; + line.reserve(32768); + append_record_prefix(line, "header"); + line += ",\"observer\":"; + append_json_string(line, mode_name(opts.mode)); + line += ",\"trace_variant\":"; + if (opts.mode == observer_mode::off) { + append_json_string(line, "none"); + } else if (opts.mode == observer_mode::graph) { + append_json_string(line, "ggml_sched_ask_metadata_v1"); + } else { + append_json_string(line, "ggml_sched_fixed_prefix_sketch_v1"); + } + line += ",\"cuda_kernel_trace\":false"; + + line += ",\"model\":{\"path\":"; + append_json_string(line, opts.model_path); + line += ",\"description\":"; + append_json_string(line, description); + line += ",\"tensor_size_bytes\":"; + append_integer(line, llama_model_size(model)); + line += ",\"parameter_count\":"; + append_integer(line, llama_model_n_params(model)); + line += ",\"file_type\":"; + append_integer(line, static_cast(llama_model_ftype(model))); + line += ",\"n_ctx_train\":"; + append_integer(line, llama_model_n_ctx_train(model)); + line += ",\"n_embd\":"; + append_integer(line, llama_model_n_embd(model)); + line += ",\"n_layer\":"; + append_integer(line, llama_model_n_layer(model)); + line += ",\"n_head\":"; + append_integer(line, llama_model_n_head(model)); + line += ",\"n_head_kv\":"; + append_integer(line, llama_model_n_head_kv(model)); + line += ",\"n_vocab\":"; + append_integer(line, llama_vocab_n_tokens(llama_model_get_vocab(model))); + line += ",\"rope_type\":"; + append_integer(line, static_cast(llama_model_rope_type(model))); + line += ",\"rope_freq_scale_train\":"; + append_float(line, llama_model_rope_freq_scale_train(model)); + line += ",\"metadata\":["; + for (size_t i = 0; i < metadata.entries.size(); ++i) { + if (i != 0) { + line.push_back(','); + } + line += "{\"key\":"; + append_json_string(line, metadata.entries[i].first); + line += ",\"value\":"; + append_json_string(line, metadata.entries[i].second); + line.push_back('}'); + } + line += "]}"; + + line += ",\"runtime\":{\"ggml_version\":"; + append_json_string(line, ggml_version()); + line += ",\"ggml_commit\":"; + append_json_string(line, ggml_commit()); + line += ",\"system_info\":"; + append_json_string(line, llama_print_system_info()); + line += ",\"supports_gpu_offload\":"; + append_bool(line, llama_supports_gpu_offload()); + line += ",\"requested_gpu_layers\":"; + append_integer(line, opts.n_gpu_layers); + line += ",\"devices\":["; + const size_t device_count = ggml_backend_dev_count(); + for (size_t i = 0; i < device_count; ++i) { + if (i != 0) { + line.push_back(','); + } + const ggml_backend_dev_t device = ggml_backend_dev_get(i); + ggml_backend_dev_props properties = {}; + ggml_backend_dev_get_props(device, &properties); + line += "{\"name\":"; + append_json_string(line, properties.name ? properties.name : ""); + line += ",\"description\":"; + append_json_string(line, properties.description ? properties.description : ""); + line += ",\"type\":"; + append_json_string(line, device_type_name(properties.type)); + line += ",\"memory_free_observed\":"; + append_integer(line, properties.memory_free); + line += ",\"memory_total\":"; + append_integer(line, properties.memory_total); + line.push_back('}'); + } + line += "]}"; + + line += ",\"execution_policy\":{\"sampling\":\"greedy\",\"temperature\":0,"; + line += "\"top_p\":1,\"top_k\":0,\"batch\":1,\"request_batch\":1,\"n_batch\":"; + append_integer(line, llama_n_batch(context)); + line += ",\"n_ubatch\":"; + append_integer(line, llama_n_ubatch(context)); + line += ",\"parallel\":1,\"parallel_sequences\":"; + append_integer(line, llama_n_seq_max(context)); + line += ",\"tensor_parallel\":1,\"split_mode\":\"none\",\"scheduler_parallel\":false,"; + line += "\"tensor_repack\":false,"; + line += "\"context_shift\":false,\"speculation\":false,\"flash_attention\":false,"; + line += "\"threads\":1,\"threads_batch\":1,\"n_predict\":"; + append_integer(line, opts.n_predict); + line += ",\"n_ctx\":"; + append_integer(line, llama_n_ctx(context)); + line += ",\"prompt_tokens\":"; + append_integer(line, prompt_tokens); + line += ",\"tokenization\":{\"add_special\":true,\"parse_special\":true}}"; + + line += ",\"observation_policy\":{\"read_only\":true,"; + line += "\"graph_metadata_stage\":\"ask\",\"sketch_bits\":256,"; + line += "\"sketch_probe\":\"gemm_output_contiguous_prefix_64\","; + line += "\"raw_activation_values_published\":false,"; + line += "\"tile_variant\":\"graph_fallback_logical_prefix_v1\","; + line += "\"kernel_trace_claim\":\"none\"}"; + line.push_back('}'); + return write_json_line(line); +} + +static bool tokenize_prompt(const llama_vocab * vocab, const std::string & prompt, + std::vector & tokens) { + const int32_t length = static_cast(prompt.size()); + const int32_t required = llama_tokenize(vocab, prompt.data(), length, nullptr, 0, true, true); + if (required == std::numeric_limits::min()) { + std::fprintf(stderr, "error: tokenizer result overflow\n"); + return false; + } + const int32_t count = required < 0 ? -required : required; + if (count <= 0) { + std::fprintf(stderr, "error: prompt tokenization produced no tokens\n"); + return false; + } + tokens.resize(static_cast(count)); + const int32_t actual = llama_tokenize(vocab, prompt.data(), length, + tokens.data(), count, true, true); + if (actual != count) { + std::fprintf(stderr, "error: prompt tokenization was not stable across sizing calls\n"); + return false; + } + return true; +} + +static bool append_token_piece(const llama_vocab * vocab, llama_token token, + std::vector & output) { + std::array small = {}; + int32_t length = llama_token_to_piece(vocab, token, small.data(), small.size(), 0, false); + if (length >= 0) { + output.insert(output.end(), small.begin(), small.begin() + length); + return true; + } + if (length == std::numeric_limits::min()) { + return false; + } + const int32_t required = -length; + std::vector buffer(static_cast(required)); + length = llama_token_to_piece(vocab, token, buffer.data(), required, 0, false); + if (length != required) { + return false; + } + output.insert(output.end(), buffer.begin(), buffer.end()); + return true; +} + +static bool emit_result(const options & opts, const char * status, const char * stop_reason, + const std::vector & prompt_tokens, + const std::vector & generated_tokens, + const std::vector & output_bytes, + const graph_observer & observer, const std::string & error) { + std::string line; + line.reserve(1024 + prompt_tokens.size() * 12 + generated_tokens.size() * 12 + output_bytes.size() * 4); + append_record_prefix(line, "result"); + line += ",\"status\":"; + append_json_string(line, status); + line += ",\"observer\":"; + append_json_string(line, mode_name(opts.mode)); + line += ",\"stop_reason\":"; + append_json_string(line, stop_reason); + line += ",\"prompt_token_ids\":"; + append_token_ids(line, prompt_tokens); + line += ",\"generated_token_ids\":"; + append_token_ids(line, generated_tokens); + line += ",\"output_bytes\":["; + const size_t published_output_bytes = opts.emit_output_bytes ? output_bytes.size() : 0; + for (size_t i = 0; i < published_output_bytes; ++i) { + if (i != 0) { + line.push_back(','); + } + append_integer(line, static_cast(output_bytes[i])); + } + line += "],\"output_n_bytes\":"; + append_integer(line, published_output_bytes); + line += ",\"event_count\":"; + append_integer(line, observer.event_count()); + if (!error.empty()) { + line += ",\"error\":"; + append_json_string(line, error); + } + line.push_back('}'); + return write_json_line(line); +} + +struct backend_guard { + ~backend_guard() { + llama_backend_free(); + } +}; + +} // namespace + +int main(int argc, char ** argv) { + std::setlocale(LC_NUMERIC, "C"); + llama_log_set(log_callback, nullptr); + + options opts; + bool help = false; + bool version = false; + if (!parse_options(argc, argv, opts, help, version)) { + print_usage(argv[0]); + return 2; + } + if (version) { + std::fprintf(stderr, "version: %d (%s)\n", llama_build_number(), llama_commit()); + return 0; + } + if (help) { + print_usage(argv[0]); + return 0; + } + + llama_backend_init(); + backend_guard backend_cleanup; + ggml_backend_load_all(); + + llama_model_params model_params = llama_model_default_params(); + model_params.n_gpu_layers = opts.n_gpu_layers; + model_params.split_mode = LLAMA_SPLIT_MODE_NONE; + model_params.main_gpu = 0; + model_params.use_mmap = true; + model_params.use_mlock = false; + model_params.check_tensors = true; + model_params.use_extra_bufts = false; + + using model_ptr = std::unique_ptr; + model_ptr model(llama_model_load_from_file(opts.model_path.c_str(), model_params), llama_model_free); + if (!model) { + std::fprintf(stderr, "error: failed to load model\n"); + return 3; + } + + model_metadata metadata; + if (!read_model_metadata(model.get(), metadata) || !validate_qwen3_8b_profile(model.get(), metadata)) { + return 3; + } + + const llama_vocab * vocab = llama_model_get_vocab(model.get()); + std::vector prompt_tokens; + if (!tokenize_prompt(vocab, opts.prompt, prompt_tokens)) { + return 3; + } + + const uint64_t required_context = prompt_tokens.size() + static_cast(opts.n_predict); + if (required_context > PALW_CONTEXT_TOKENS || + PALW_CONTEXT_TOKENS > static_cast(llama_model_n_ctx_train(model.get()))) { + std::fprintf(stderr, "error: prompt plus n-predict exceeds the fixed PALW context bound\n"); + return 3; + } + + graph_observer observer(opts.mode, llama_model_n_layer(model.get())); + llama_context_params context_params = llama_context_default_params(); + context_params.n_ctx = PALW_CONTEXT_TOKENS; + context_params.n_batch = 1; + context_params.n_ubatch = 1; + context_params.n_seq_max = 1; + context_params.n_outputs_max = 1; + context_params.n_threads = 1; + context_params.n_threads_batch = 1; + context_params.flash_attn_type = LLAMA_FLASH_ATTN_TYPE_DISABLED; + context_params.type_k = GGML_TYPE_F16; + context_params.type_v = GGML_TYPE_F16; + context_params.embeddings = false; + context_params.offload_kqv = true; + context_params.no_perf = true; + context_params.op_offload = true; + context_params.kv_unified = false; + if (opts.mode != observer_mode::off) { + context_params.cb_eval = graph_observer::callback; + context_params.cb_eval_user_data = &observer; + } + + using context_ptr = std::unique_ptr; + context_ptr context(llama_init_from_model(model.get(), context_params), llama_free); + if (!context) { + std::fprintf(stderr, "error: failed to initialize context\n"); + return 3; + } + if (llama_n_batch(context.get()) != 1 || llama_n_ubatch(context.get()) != 1 || + llama_n_seq_max(context.get()) != 1 || llama_n_ctx(context.get()) != PALW_CONTEXT_TOKENS || + llama_n_threads(context.get()) != 1 || llama_n_threads_batch(context.get()) != 1) { + std::fprintf(stderr, "error: runtime did not honor the deterministic execution policy\n"); + return 3; + } + + if (!emit_header(opts, model.get(), context.get(), metadata, prompt_tokens.size())) { + return 4; + } + + std::vector generated_tokens; + std::vector output_bytes; + generated_tokens.reserve(static_cast(opts.n_predict)); + + for (size_t i = 0; i < prompt_tokens.size(); ++i) { + observer.set_phase("prefill", static_cast(i)); + llama_token token = prompt_tokens[i]; + const int32_t decode_status = llama_decode(context.get(), llama_batch_get_one(&token, 1)); + if (decode_status != 0) { + const std::string error = "llama_decode prefill status " + std::to_string(decode_status); + emit_result(opts, "error", "decode_error", prompt_tokens, generated_tokens, + output_bytes, observer, error); + return 5; + } + if (observer.failed()) { + emit_result(opts, "error", "observer_error", prompt_tokens, generated_tokens, + output_bytes, observer, observer.error()); + return 5; + } + } + + using sampler_ptr = std::unique_ptr; + sampler_ptr sampler(llama_sampler_init_greedy(), llama_sampler_free); + if (!sampler) { + emit_result(opts, "error", "sampler_error", prompt_tokens, generated_tokens, + output_bytes, observer, "failed to initialize greedy sampler"); + return 5; + } + + const char * stop_reason = "n_predict"; + for (int32_t i = 0; i < opts.n_predict; ++i) { + const llama_token token = llama_sampler_sample(sampler.get(), context.get(), -1); + if (token < 0 || token >= llama_vocab_n_tokens(vocab)) { + emit_result(opts, "error", "sampler_error", prompt_tokens, generated_tokens, + output_bytes, observer, "greedy sampler returned an out-of-range token"); + return 5; + } + generated_tokens.push_back(token); + if (llama_vocab_is_eog(vocab, token)) { + stop_reason = "eog"; + break; + } + if (!append_token_piece(vocab, token, output_bytes)) { + emit_result(opts, "error", "detokenize_error", prompt_tokens, generated_tokens, + output_bytes, observer, "failed to convert generated token to bytes"); + return 5; + } + if (i + 1 == opts.n_predict) { + break; + } + + observer.set_phase("decode", i); + llama_token mutable_token = token; + const int32_t decode_status = llama_decode(context.get(), llama_batch_get_one(&mutable_token, 1)); + if (decode_status != 0) { + const std::string error = "llama_decode generation status " + std::to_string(decode_status); + emit_result(opts, "error", "decode_error", prompt_tokens, generated_tokens, + output_bytes, observer, error); + return 5; + } + if (observer.failed()) { + emit_result(opts, "error", "observer_error", prompt_tokens, generated_tokens, + output_bytes, observer, observer.error()); + return 5; + } + } + + if (!emit_result(opts, "ok", stop_reason, prompt_tokens, generated_tokens, + output_bytes, observer, {})) { + return 4; + } + return 0; +} diff --git a/src/models/qwen35moe.cpp b/src/models/qwen35moe.cpp index 7b0876cb..4d955261 100644 --- a/src/models/qwen35moe.cpp +++ b/src/models/qwen35moe.cpp @@ -6,7 +6,18 @@ void llama_model_qwen35moe::load_arch_hparams(llama_model_loader & ml) { ml.get_key(LLM_KV_EXPERT_SHARED_FEED_FORWARD_LENGTH, hparams.n_ff_shexp, false); ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps); - ml.get_key_or_arr(LLM_KV_ROPE_DIMENSION_SECTIONS, hparams.rope_sections, 4, true); + // PALW: newer HF->GGUF conversions write 3 mrope sections ([t, h, w]) and + // omit the trailing zero; older conversions write 4. Accept both forms and + // zero-pad so the pinned runtime loads current upstream GGUF artifacts. + { + std::vector sections; + ml.get_arr(LLM_KV_ROPE_DIMENSION_SECTIONS, sections, true); + if (sections.size() != 3 && sections.size() != 4) { + throw std::runtime_error("rope.dimension_sections must have 3 or 4 entries"); + } + std::fill(hparams.rope_sections.begin(), hparams.rope_sections.end(), 0); + std::copy(sections.begin(), sections.end(), hparams.rope_sections.begin()); + } // Load linear attention (gated delta net) parameters ml.get_key(LLM_KV_SSM_CONV_KERNEL, hparams.ssm_d_conv); @@ -73,8 +84,13 @@ void llama_model_qwen35moe::load_arch_tensors(llama_model_loader & ml) { layer.attn_post_norm = create_tensor(tn(LLM_TENSOR_ATTN_POST_NORM, "weight", il), { n_embd }, flags); if (!hparams.is_recr(il)) { - // Attention layers - create_tensor_qkv(layer, il, n_embd, n_embd_head_k * n_head * 2, n_embd_k_gqa, n_embd_v_gqa, flags); + // Attention layers. The LLAMA_LOAD_LOCALS macro derives the global + // n_embd_k_gqa/n_embd_v_gqa from layer 0, which in this MoE is a + // linear-attention (recurrent) layer with n_head_kv == 0. Use the + // uniform full-attention KV projection width instead. + const int64_t n_embd_k_gqa_attn = hparams.n_embd_k_gqa_max(); + const int64_t n_embd_v_gqa_attn = hparams.n_embd_v_gqa_max(); + create_tensor_qkv(layer, il, n_embd, n_embd_head_k * n_head * 2, n_embd_k_gqa_attn, n_embd_v_gqa_attn, flags); layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", il), { n_embd_head_k * n_head, n_embd }, flags); // Q/K normalization for attention layers @@ -86,7 +102,12 @@ void llama_model_qwen35moe::load_arch_tensors(llama_model_loader & ml) { layer.wqkv = create_tensor(tn(LLM_TENSOR_ATTN_QKV, "weight", il), { n_embd, key_dim * 2 + value_dim }, TENSOR_NOT_REQUIRED); layer.wqkv_gate = create_tensor(tn(LLM_TENSOR_ATTN_GATE, "weight", il), { n_embd, value_dim }, TENSOR_NOT_REQUIRED); layer.ssm_conv1d = create_tensor(tn(LLM_TENSOR_SSM_CONV1D, "weight", il), { hparams.ssm_d_conv, conv_dim }, flags); - layer.ssm_dt = create_tensor(tn(LLM_TENSOR_SSM_DT, "bias", il), { hparams.ssm_dt_rank }, flags); + // PALW: some conversions store the delta-time bias without the + // ".bias" suffix ("blk.N.ssm_dt"). Accept both spellings. + layer.ssm_dt = create_tensor(tn(LLM_TENSOR_SSM_DT, "bias", il), { hparams.ssm_dt_rank }, flags | TENSOR_NOT_REQUIRED); + if (!layer.ssm_dt) { + layer.ssm_dt = create_tensor(tn(LLM_TENSOR_SSM_DT, il), { hparams.ssm_dt_rank }, flags); + } layer.ssm_a = create_tensor(tn(LLM_TENSOR_SSM_A_NOSCAN, il), { hparams.ssm_dt_rank }, flags); layer.ssm_beta = create_tensor(tn(LLM_TENSOR_SSM_BETA, "weight", il), { n_embd, n_v_heads }, flags); layer.ssm_alpha = create_tensor(tn(LLM_TENSOR_SSM_ALPHA, "weight", il), { n_embd, n_v_heads }, flags); @@ -116,7 +137,11 @@ void llama_model_qwen35moe::load_arch_tensors(llama_model_loader & ml) { layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", il), { n_embd }, 0); layer.attn_post_norm = create_tensor(tn(LLM_TENSOR_ATTN_POST_NORM, "weight", il), { n_embd }, 0); - create_tensor_qkv(layer, il, n_embd, n_embd_head_k * n_head * 2, n_embd_k_gqa, n_embd_v_gqa, 0); + // The MTP layer index has no entry in the per-layer n_head_kv array, so + // derive the KV projection width from the full-attention trunk layers. + const int64_t n_embd_k_gqa_mtp = hparams.n_embd_k_gqa_max(); + const int64_t n_embd_v_gqa_mtp = hparams.n_embd_v_gqa_max(); + create_tensor_qkv(layer, il, n_embd, n_embd_head_k * n_head * 2, n_embd_k_gqa_mtp, n_embd_v_gqa_mtp, 0); layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", il), { n_embd_head_k * n_head, n_embd }, 0); layer.attn_q_norm = create_tensor(tn(LLM_TENSOR_ATTN_Q_NORM, "weight", il), { n_embd_head_k }, 0); layer.attn_k_norm = create_tensor(tn(LLM_TENSOR_ATTN_K_NORM, "weight", il), { n_embd_head_k }, 0); @@ -147,6 +172,24 @@ void llama_model_qwen35moe::load_arch_tensors(llama_model_loader & ml) { for (int i = n_layer; i < n_layer_all; ++i) { load_block_mtp(i); } + + // PALW: the Ollama-packaged Qwen3.6-35B-A3B GGUF bundles the multimodal + // vision tower ("v.*") and, when the next-token/MTP head is not enabled by + // hyperparameters, the MTP sub-model ("mtp.*") in the same file. The pinned + // text runtime never builds those sibling sub-models, so account for their + // tensors here; otherwise done_getting_tensors() rejects the load for + // having created fewer tensors than the file contains. This does not weaken + // the check for a genuinely missing text tensor: only tensors that exist in + // the file under these sibling prefixes are counted, exactly once each. + for (const auto & entry : ml.weights_map) { + const std::string & name = entry.first; + const bool is_vision = name.rfind("v.", 0) == 0; + const bool is_mtp = name.rfind("mtp.", 0) == 0; + if (is_vision || is_mtp) { + ml.size_data -= ggml_nbytes(entry.second.tensor); + ml.n_created++; + } + } } std::unique_ptr llama_model_qwen35moe::build_arch_graph(const llm_graph_params & params) const { @@ -287,6 +330,12 @@ ggml_tensor * llama_model_qwen35moe::graph::build_layer_attn( const int64_t n_embd_head = hparams.n_embd_head_v(); GGML_ASSERT(n_embd_head == hparams.n_embd_head_k()); + // PALW: the llm_graph_context base initializes n_head/n_head_kv from layer + // 0, which in this hybrid MoE is a linear-attention layer with n_head_kv 0. + // Use the per-layer counts for the full-attention reshapes below. + const int64_t n_head = hparams.n_head(il); + const int64_t n_head_kv = hparams.n_head_kv(il); + // Order: joint QG projection, QG split, Q norm, KV projection, K norm, RoPE, attention // Qwen3Next uses a single Q projection that outputs query + gate diff --git a/ggml/include/ggml-metal.h b/ggml/include/ggml-metal.h index 433838f0..213e148d 100644 --- a/ggml/include/ggml-metal.h +++ b/ggml/include/ggml-metal.h @@ -56,6 +56,21 @@ GGML_BACKEND_API void ggml_backend_metal_capture_next_compute(ggml_backend_t bac GGML_BACKEND_API ggml_backend_reg_t ggml_backend_metal_reg(void); +// PALW kernel-level dispatch hook. +// +// When set (process-global), the callback is invoked for every Metal compute +// dispatch with the bound pipeline (kernel) name and the launch geometry +// (threadgroup grid tg0..2 and threads-per-threadgroup tptg0..2). This lets the +// PALW observer produce a kernel-level execution trace bound to the actual GPU +// kernel dispatches rather than to graph-node outputs. Pass NULL to disable. +typedef void (*ggml_metal_palw_dispatch_cb)( + void * user_data, + const char * pipeline, + int tg0, int tg1, int tg2, + int tptg0, int tptg1, int tptg2); + +GGML_BACKEND_API void ggml_metal_palw_set_dispatch_hook(ggml_metal_palw_dispatch_cb cb, void * user_data); + #ifdef __cplusplus } #endif diff --git a/ggml/src/ggml-metal/ggml-metal-device.m b/ggml/src/ggml-metal/ggml-metal-device.m index 80e47f2c..3e36c08d 100644 --- a/ggml/src/ggml-metal/ggml-metal-device.m +++ b/ggml/src/ggml-metal/ggml-metal-device.m @@ -1,5 +1,6 @@ #import "ggml-metal-device.h" +#import "ggml-metal.h" #import "ggml-impl.h" #import "ggml-backend-impl.h" @@ -72,6 +73,9 @@ void ggml_metal_cv_set_bool(ggml_metal_cv_t cv, bool value, int32_t idx) { struct ggml_metal_pipeline { id obj; + + // PALW: stable kernel (pipeline) name, captured for kernel-level tracing. + char name[128]; }; ggml_metal_pipeline_t ggml_metal_pipeline_init(void) { @@ -79,6 +83,7 @@ ggml_metal_pipeline_t ggml_metal_pipeline_init(void) { *res = (struct ggml_metal_pipeline) { /*.obj =*/ nil, + /*.name =*/ {0}, }; return res; @@ -443,6 +448,8 @@ struct ggml_metal_pipeline_with_params ggml_metal_library_compile_pipeline(ggml_ res.pipeline = ggml_metal_pipeline_init(); res.pipeline->obj = obj; + // PALW: record the stable kernel name for kernel-level dispatch tracing. + snprintf(res.pipeline->name, sizeof(res.pipeline->name), "%s", name); ggml_metal_pipelines_add(lib->pipelines, name, res.pipeline); } @@ -458,8 +465,23 @@ struct ggml_metal_pipeline_with_params ggml_metal_library_compile_pipeline(ggml_ struct ggml_metal_encoder { id obj; + + // PALW: name of the pipeline currently bound, for kernel-level dispatch tracing. + const char * cur_pipeline; }; +// PALW: process-global kernel-dispatch hook. When set, it is invoked for every +// Metal compute dispatch with the bound pipeline (kernel) name and the launch +// geometry (threadgroup grid + threads-per-threadgroup). Used by the PALW +// observer to produce a kernel-level (not graph-fallback) execution trace. +static ggml_metal_palw_dispatch_cb g_palw_dispatch_cb = NULL; +static void * g_palw_dispatch_ud = NULL; + +void ggml_metal_palw_set_dispatch_hook(ggml_metal_palw_dispatch_cb cb, void * user_data) { + g_palw_dispatch_cb = cb; + g_palw_dispatch_ud = user_data; +} + ggml_metal_encoder_t ggml_metal_encoder_init(ggml_metal_cmd_buf_t cmd_buf_raw, bool concurrent) { ggml_metal_encoder_t res = calloc(1, sizeof(struct ggml_metal_encoder)); @@ -491,6 +513,8 @@ void ggml_metal_encoder_debug_group_pop (ggml_metal_encoder_t encoder) { void ggml_metal_encoder_set_pipeline(ggml_metal_encoder_t encoder, struct ggml_metal_pipeline_with_params pipeline) { [encoder->obj setComputePipelineState:pipeline.pipeline->obj]; + // PALW: remember the bound kernel name so dispatches can be attributed to it. + encoder->cur_pipeline = pipeline.pipeline->name; } void ggml_metal_encoder_set_bytes(ggml_metal_encoder_t encoder, void * data, size_t size, int idx) { @@ -506,6 +530,11 @@ void ggml_metal_encoder_set_threadgroup_memory_size(ggml_metal_encoder_t encoder } void ggml_metal_encoder_dispatch_threadgroups(ggml_metal_encoder_t encoder, int tg0, int tg1, int tg2, int tptg0, int tptg1, int tptg2) { + // PALW: report the kernel-level dispatch (bound pipeline + launch geometry). + if (g_palw_dispatch_cb) { + g_palw_dispatch_cb(g_palw_dispatch_ud, encoder->cur_pipeline ? encoder->cur_pipeline : "", + tg0, tg1, tg2, tptg0, tptg1, tptg2); + } [encoder->obj dispatchThreadgroups:MTLSizeMake(tg0, tg1, tg2) threadsPerThreadgroup:MTLSizeMake(tptg0, tptg1, tptg2)]; } diff --git a/tools/palw-observer/PALW_PATCHSET b/tools/palw-observer/PALW_PATCHSET new file mode 100644 index 00000000..887d5502 --- /dev/null +++ b/tools/palw-observer/PALW_PATCHSET @@ -0,0 +1,4 @@ +palw-full-patch-v1 +base=12127defda4f41b7679cb2477a4b0d65ee6a0c8f +scope=cuda-trace,metal-dispatch,qwen35-gguf-compat,observer +receipt-authority=false