diff --git a/ggml/CMakeLists.txt b/ggml/CMakeLists.txt index 159da3afa0b0..6f2230a562a7 100644 --- a/ggml/CMakeLists.txt +++ b/ggml/CMakeLists.txt @@ -213,6 +213,7 @@ set (GGML_CUDA_COMPRESSION_MODE "size" CACHE STRING set_property(CACHE GGML_CUDA_COMPRESSION_MODE PROPERTY STRINGS "none;speed;balance;size") option(GGML_HIP "ggml: use HIP" OFF) +option(GGML_HRX "ggml: use HRX" OFF) option(GGML_HIP_GRAPHS "ggml: use HIP graph" ON) option(GGML_HIP_RCCL "ggml: use ROCm Collective Comm. Library" OFF) option(GGML_HIP_NO_VMM "ggml: do not try to use HIP VMM" ON) diff --git a/ggml/include/ggml-hrx.h b/ggml/include/ggml-hrx.h new file mode 100644 index 000000000000..d40c095b75bc --- /dev/null +++ b/ggml/include/ggml-hrx.h @@ -0,0 +1,26 @@ +#pragma once + +#include "ggml-backend.h" + +#ifdef __cplusplus +extern "C" { +#endif + +struct ggml_backend_hrx_cache_stats { + uint64_t graph_program_builds; + uint64_t graph_program_hits; + uint64_t prepared_program_builds; + uint64_t prepared_program_hits; +}; + +GGML_BACKEND_API ggml_backend_t ggml_backend_hrx_init(size_t device); +GGML_BACKEND_API bool ggml_backend_is_hrx(ggml_backend_t backend); +GGML_BACKEND_API int ggml_backend_hrx_get_device_count(void); +GGML_BACKEND_API ggml_backend_buffer_type_t ggml_backend_hrx_buffer_type(size_t device); +GGML_BACKEND_API bool ggml_backend_hrx_get_cache_stats(ggml_backend_t backend, + struct ggml_backend_hrx_cache_stats * stats); +GGML_BACKEND_API ggml_backend_reg_t ggml_backend_hrx_reg(void); + +#ifdef __cplusplus +} +#endif diff --git a/ggml/src/CMakeLists.txt b/ggml/src/CMakeLists.txt index 82e9480c2f24..785a9ed2790a 100644 --- a/ggml/src/CMakeLists.txt +++ b/ggml/src/CMakeLists.txt @@ -475,6 +475,7 @@ ggml_add_backend(CANN) ggml_add_backend(CUDA) ggml_add_backend(ET) ggml_add_backend(HIP) +ggml_add_backend(HRX) ggml_add_backend(METAL) ggml_add_backend(MUSA) ggml_add_backend(RPC) diff --git a/ggml/src/ggml-backend-reg.cpp b/ggml/src/ggml-backend-reg.cpp index e5959467071d..d33e2dd4d99d 100644 --- a/ggml/src/ggml-backend-reg.cpp +++ b/ggml/src/ggml-backend-reg.cpp @@ -34,6 +34,10 @@ #include "ggml-cuda.h" #endif +#ifdef GGML_USE_HRX +#include "ggml-hrx.h" +#endif + #ifdef GGML_USE_METAL #include "ggml-metal.h" #endif @@ -120,6 +124,9 @@ struct ggml_backend_registry { #ifdef GGML_USE_CUDA register_backend(ggml_backend_cuda_reg()); #endif +#ifdef GGML_USE_HRX + register_backend(ggml_backend_hrx_reg()); +#endif #ifdef GGML_USE_METAL register_backend(ggml_backend_metal_reg()); #endif diff --git a/ggml/src/ggml-hrx/CMakeLists.txt b/ggml/src/ggml-hrx/CMakeLists.txt new file mode 100644 index 000000000000..b20301189ce0 --- /dev/null +++ b/ggml/src/ggml-hrx/CMakeLists.txt @@ -0,0 +1,245 @@ +set(HRX_SOURCE_DIR "" CACHE PATH "Optional HRX source tree to build instead of using installed hrx and loomc packages") + +if(HRX_SOURCE_DIR) + include(ExternalProject) + include(GNUInstallDirs) + + get_filename_component(HRX_SOURCE_DIR "${HRX_SOURCE_DIR}" ABSOLUTE) + if(NOT EXISTS "${HRX_SOURCE_DIR}/CMakeLists.txt") + message(FATAL_ERROR "HRX_SOURCE_DIR does not contain a CMakeLists.txt: ${HRX_SOURCE_DIR}") + endif() + + set(GGML_HRX_PREFIX "${CMAKE_CURRENT_BINARY_DIR}/hrx") + set(GGML_HRX_BUILD_DIR "${GGML_HRX_PREFIX}/src/ggml-hrx-deps-build") + set(GGML_HRX_LIB "${GGML_HRX_BUILD_DIR}/libhrx/src/libhrx/${CMAKE_SHARED_LIBRARY_PREFIX}hrx${CMAKE_SHARED_LIBRARY_SUFFIX}") + set(GGML_LOOMC_LIB "${GGML_HRX_BUILD_DIR}/loom/binding/c/${CMAKE_SHARED_LIBRARY_PREFIX}loomc${CMAKE_SHARED_LIBRARY_SUFFIX}") + set(GGML_HRX_LOOM_LINK "${GGML_HRX_BUILD_DIR}/loom/src/loom/tools/loom-link/loom-link${CMAKE_EXECUTABLE_SUFFIX}") + set(GGML_HRX_LOOM_FORMAT "${GGML_HRX_BUILD_DIR}/loom/src/loom/tools/loom-format/loom-format${CMAKE_EXECUTABLE_SUFFIX}") + set(GGML_HRX_DEPS_TARGET ggml-hrx-deps) + + set(GGML_HRX_CMAKE_ARGS + -DCMAKE_BUILD_TYPE=${CMAKE_BUILD_TYPE} + -DCMAKE_C_COMPILER=${CMAKE_C_COMPILER} + -DCMAKE_CXX_COMPILER=${CMAKE_CXX_COMPILER} + -DIREE_BUILD_TESTS=OFF + -DIREE_BUILD_BENCHMARKS=OFF + -DIREE_HAL_DRIVER_DEFAULTS=OFF + -DIREE_HAL_DRIVER_AMDGPU=ON + -DIREE_HAL_DRIVER_TASK=ON + + -DIREE_HAL_DRIVER_NULL=ON + -DLIBHRX_BUILD_CTS=OFF + ) + if(IREE_ROCM_PATH) + list(APPEND GGML_HRX_CMAKE_ARGS -DIREE_ROCM_PATH=${IREE_ROCM_PATH}) + endif() + if(FETCHCONTENT_BASE_DIR) + list(APPEND GGML_HRX_CMAKE_ARGS -DFETCHCONTENT_BASE_DIR=${FETCHCONTENT_BASE_DIR}) + endif() + + ExternalProject_Add(ggml-hrx-deps + SOURCE_DIR "${HRX_SOURCE_DIR}" + PREFIX "${GGML_HRX_PREFIX}" + CMAKE_ARGS ${GGML_HRX_CMAKE_ARGS} + BUILD_COMMAND ${CMAKE_COMMAND} --build . --target hrx loomc_shared loom_tools_loom-link_loom-link loom_tools_loom-format_loom-format --config ${CMAKE_BUILD_TYPE} + INSTALL_COMMAND "" + BUILD_BYPRODUCTS "${GGML_HRX_LIB}" "${GGML_LOOMC_LIB}" "${GGML_HRX_LOOM_LINK}" "${GGML_HRX_LOOM_FORMAT}" + UPDATE_COMMAND "" + ) + + add_library(hrx::hrx SHARED IMPORTED GLOBAL) + set_target_properties(hrx::hrx PROPERTIES + IMPORTED_LOCATION "${GGML_HRX_LIB}" + INTERFACE_INCLUDE_DIRECTORIES "${HRX_SOURCE_DIR}/libhrx/include") + add_dependencies(hrx::hrx ggml-hrx-deps) + + add_library(loomc::loomc SHARED IMPORTED GLOBAL) + set_target_properties(loomc::loomc PROPERTIES + IMPORTED_LOCATION "${GGML_LOOMC_LIB}" + INTERFACE_INCLUDE_DIRECTORIES "${HRX_SOURCE_DIR}/loom/binding/c/include" + INTERFACE_COMPILE_DEFINITIONS LOOMC_USING_SHARED_LIBRARY) + add_dependencies(loomc::loomc ggml-hrx-deps) +else() + find_package(hrx CONFIG REQUIRED) + find_package(loomc CONFIG REQUIRED) + find_program(GGML_HRX_LOOM_LINK NAMES loom-link) + find_program(GGML_HRX_LOOM_FORMAT NAMES loom-format) +endif() + +find_package(Python3 REQUIRED COMPONENTS Interpreter) + +set(GGML_HRX_KERNEL_CORPUS_SOURCE_FORMAT "binary" CACHE STRING "Embedded Loom corpus source format: text or binary") +set_property(CACHE GGML_HRX_KERNEL_CORPUS_SOURCE_FORMAT PROPERTY STRINGS text binary) +if(GGML_HRX_KERNEL_CORPUS_SOURCE_FORMAT STREQUAL "binary") + if(NOT GGML_HRX_LOOM_LINK OR NOT GGML_HRX_LOOM_FORMAT) + message(FATAL_ERROR "GGML_HRX_KERNEL_CORPUS_SOURCE_FORMAT=binary requires loom-link and loom-format") + endif() + set(GGML_HRX_KERNEL_CORPUS_TOOL_ARGS + --loom-link "${GGML_HRX_LOOM_LINK}" + --loom-format "${GGML_HRX_LOOM_FORMAT}" + ) + set(GGML_HRX_KERNEL_CORPUS_TOOL_DEPENDS + "${GGML_HRX_LOOM_LINK}" + "${GGML_HRX_LOOM_FORMAT}" + ) +elseif(NOT GGML_HRX_KERNEL_CORPUS_SOURCE_FORMAT STREQUAL "text") + message(FATAL_ERROR "Unsupported GGML_HRX_KERNEL_CORPUS_SOURCE_FORMAT: ${GGML_HRX_KERNEL_CORPUS_SOURCE_FORMAT}") +endif() + +set(GGML_HRX_QWEN_KERNEL_CORPUS_DIR "${CMAKE_CURRENT_SOURCE_DIR}/kernel-corpus/kernels/qwen_moe") +set(GGML_HRX_QWEN_KERNEL_CORPUS_MANIFEST "${GGML_HRX_QWEN_KERNEL_CORPUS_DIR}/manifest.json") +set(GGML_HRX_KERNEL_CORPUS_SOURCES_INC "${CMAKE_CURRENT_BINARY_DIR}/kernel-corpus-sources.inc") +set(GGML_HRX_KERNEL_CORPUS_QWEN_INC "${CMAKE_CURRENT_BINARY_DIR}/kernel-corpus-qwen.inc") +set(GGML_HRX_KERNEL_CORPUS_CATALOG_INC "${CMAKE_CURRENT_BINARY_DIR}/kernel-corpus-catalog.inc") +set(GGML_HRX_KERNEL_CORPUS_DEPFILE "${CMAKE_CURRENT_BINARY_DIR}/kernel-corpus.d") + +add_custom_command( + OUTPUT + "${GGML_HRX_KERNEL_CORPUS_SOURCES_INC}" + "${GGML_HRX_KERNEL_CORPUS_QWEN_INC}" + "${GGML_HRX_KERNEL_CORPUS_CATALOG_INC}" + COMMAND ${Python3_EXECUTABLE} + "${CMAKE_CURRENT_SOURCE_DIR}/tools/generate_kernel_corpus.py" + --source-output "${GGML_HRX_KERNEL_CORPUS_SOURCES_INC}" + --corpus-output "${GGML_HRX_KERNEL_CORPUS_QWEN_INC}" + --catalog-output "${GGML_HRX_KERNEL_CORPUS_CATALOG_INC}" + --manifest "${GGML_HRX_QWEN_KERNEL_CORPUS_MANIFEST}" + --corpus-dir "${GGML_HRX_QWEN_KERNEL_CORPUS_DIR}" + --source-format "${GGML_HRX_KERNEL_CORPUS_SOURCE_FORMAT}" + ${GGML_HRX_KERNEL_CORPUS_TOOL_ARGS} + --depfile "${GGML_HRX_KERNEL_CORPUS_DEPFILE}" + DEPENDS + "${CMAKE_CURRENT_SOURCE_DIR}/tools/generate_kernel_corpus.py" + "${GGML_HRX_QWEN_KERNEL_CORPUS_MANIFEST}" + ${GGML_HRX_KERNEL_CORPUS_TOOL_DEPENDS} + DEPFILE "${GGML_HRX_KERNEL_CORPUS_DEPFILE}" + VERBATIM +) + +option(GGML_HRX_BUNDLE_RUNTIME_LIBS "Bundle HRX/ROCm runtime libraries next to the HRX backend" OFF) +set(GGML_HRX_BUNDLE_LIBRARY_DIRS "" CACHE STRING "Library directories to scan when GGML_HRX_BUNDLE_RUNTIME_LIBS=ON") + +add_library(ggml-hrx-kernel-corpus STATIC + status.h + kernel-corpus/kernel-corpus-json.cpp + kernel-corpus/kernel-corpus-json.h + kernel-corpus/kernel-corpus-catalog-verify.h + kernel-corpus/kernel-corpus-catalog.h + kernel-corpus/kernel-corpus.cpp + kernel-corpus/kernel-corpus.h + kernel-corpus/kernel-types.h + "${GGML_HRX_KERNEL_CORPUS_SOURCES_INC}" + "${GGML_HRX_KERNEL_CORPUS_QWEN_INC}" + "${GGML_HRX_KERNEL_CORPUS_CATALOG_INC}" +) +if(GGML_HRX_DEPS_TARGET) + add_dependencies(ggml-hrx-kernel-corpus ${GGML_HRX_DEPS_TARGET}) +endif() +target_include_directories(ggml-hrx-kernel-corpus PUBLIC . PRIVATE "${CMAKE_CURRENT_BINARY_DIR}" ../../../vendor) +target_compile_features(ggml-hrx-kernel-corpus PRIVATE cxx_std_17) +set_target_properties(ggml-hrx-kernel-corpus PROPERTIES POSITION_INDEPENDENT_CODE ON) + +ggml_add_backend_library(ggml-hrx + backend-buffer-binding.cpp + backend-buffer-binding.h + backend-context.h + dispatch/command-plan-metadata.cpp + dispatch/command-plan-metadata.h + dispatch/command-plan.h + dispatch/command-program-bindings.cpp + dispatch/command-program-bindings.h + dispatch/command-program-diagnostics.cpp + dispatch/command-program-diagnostics.h + dispatch/command-program.cpp + dispatch/command-program.h + dispatch/command-program-resolver.cpp + dispatch/command-program-resolver.h + dispatch/dispatch-scheduler.cpp + dispatch/dispatch-scheduler.h + dispatch/dispatch.h + dispatch/transient-allocator.cpp + dispatch/transient-allocator.h + dispatch_registration/dispatch-add.cpp + dispatch_registration/dispatch-add.h + dispatch_registration/dispatch-gather-add.cpp + dispatch_registration/dispatch-get-rows.cpp + dispatch_registration/dispatch-get-rows.h + dispatch_registration/dispatch-gather-add.h + dispatch_registration/dispatch-llm-matmul.cpp + dispatch_registration/dispatch-llm-matmul.h + dispatch_registration/dispatch-llm-profiles.h + dispatch_registration/dispatch-llm-shapes.h + dispatch_registration/dispatch-qwen-attention-postprocess.cpp + dispatch_registration/dispatch-qwen-attention-postprocess.h + dispatch_registration/dispatch-qwen-flash-attention.cpp + dispatch_registration/dispatch-qwen-flash-attention.h + dispatch_registration/dispatch-qwen-matmul.cpp + dispatch_registration/dispatch-qwen-matmul.h + dispatch_registration/dispatch-moe-router.cpp + dispatch_registration/dispatch-moe-router.h + dispatch_registration/dispatch-qwen-preamble.cpp + dispatch_registration/dispatch-qwen-preamble.h + dispatch_registration/dispatch-registry.cpp + dispatch_registration/dispatch-registry.h + dispatch_registration/dispatch-rmsnorm.cpp + dispatch_registration/dispatch-rmsnorm.h + dispatch_registration/dispatch-routed-ffn.cpp + dispatch_registration/dispatch-routed-ffn.h + status.h + graph/graph.cpp + graph/graph-diagnostics.cpp + graph/graph-diagnostics.h + graph/graph.h + graph/graph-matcher.cpp + graph/graph-matcher.h + graph/graph-traversal.cpp + graph/graph-traversal.h + graph/op-params.cpp + graph/op-params.h + ggml-hrx.cpp + loom-jit.cpp + graph/value-map.cpp + graph/value-map.h + runtime/command-program-executor.cpp + runtime/command-program-executor.h + runtime/graph-executor.cpp + runtime/graph-executor.h + runtime/graph-program-cache.cpp + runtime/graph-program-cache.h + runtime/host-memory.cpp + runtime/host-memory.h + runtime/kernel-executable-cache.cpp + runtime/kernel-executable-cache.h + runtime/loom-kernel-jit.cpp + runtime/loom-kernel-jit.h + runtime/prepared-command-program-cache.cpp + runtime/prepared-command-program-cache.h + runtime/transient-arena.cpp + runtime/transient-arena.h + runtime/host-buffer-registry.cpp + runtime/host-buffer-registry.h +) +target_link_libraries(ggml-hrx PRIVATE ggml-hrx-kernel-corpus hrx::hrx loomc::loomc) +target_include_directories(ggml-hrx PRIVATE . "${CMAKE_CURRENT_BINARY_DIR}" ../../../vendor) +target_compile_definitions(ggml-hrx PRIVATE GGML_USE_HRX) + +if (GGML_HRX_BUNDLE_RUNTIME_LIBS) + include("${CMAKE_CURRENT_SOURCE_DIR}/cmake/BundleRuntime.cmake") + ggml_hrx_bundle_runtime(ggml-hrx) +endif() + +add_executable(ggml-hrx-compile-kernel + tools/compile-kernel.cpp + tools/tool-utils.h + loom-jit.cpp +) +target_link_libraries(ggml-hrx-compile-kernel PRIVATE hrx::hrx loomc::loomc) +target_include_directories(ggml-hrx-compile-kernel PRIVATE .) +target_compile_features(ggml-hrx-compile-kernel PRIVATE cxx_std_17) + +add_executable(ggml-hrx-analyze-graph + tools/analyze-graph.cpp +) +target_link_libraries(ggml-hrx-analyze-graph PRIVATE ggml-hrx ggml ggml-hrx-kernel-corpus) +target_include_directories(ggml-hrx-analyze-graph PRIVATE . ../../../vendor) +target_compile_features(ggml-hrx-analyze-graph PRIVATE cxx_std_17) diff --git a/ggml/src/ggml-hrx/backend-buffer-binding.cpp b/ggml/src/ggml-hrx/backend-buffer-binding.cpp new file mode 100644 index 000000000000..bab5f12b1af7 --- /dev/null +++ b/ggml/src/ggml-hrx/backend-buffer-binding.cpp @@ -0,0 +1,86 @@ +#include "backend-buffer-binding.h" + +#include "ggml-backend-impl.h" +#include "ggml.h" + +ggml_backend_hrx_buffer_context * ggml_backend_hrx_buffer_context_from_buffer(ggml_backend_buffer_t buffer) { + return static_cast(buffer->context); +} + +size_t ggml_backend_hrx_tensor_offset(const ggml_backend_hrx_buffer_context * context, const ggml_tensor * tensor) { + return static_cast(static_cast(tensor->data) - context->base); +} + +void * ggml_backend_hrx_buffer_base(ggml_backend_buffer_t buffer) { + return ggml_backend_hrx_buffer_context_from_buffer(buffer)->base; +} + +bool ggml_backend_hrx_tensor_binding(const ggml_tensor * tensor, + ggml_backend_hrx_buffer_context ** out_context, + size_t * out_offset) { + if (tensor == nullptr) { + return false; + } + ggml_backend_buffer_t buffer = tensor->view_src != nullptr ? tensor->view_src->buffer : tensor->buffer; + if (buffer == nullptr || buffer->iface.get_base != ggml_backend_hrx_buffer_base) { + return false; + } + auto * context = ggml_backend_hrx_buffer_context_from_buffer(buffer); + const size_t offset = ggml_backend_hrx_tensor_offset(context, tensor); + if (context->buffer == nullptr || offset > buffer->size || ggml_nbytes(tensor) > buffer->size - offset) { + return false; + } + *out_context = context; + *out_offset = offset; + return true; +} + +bool ggml_backend_hrx_resolve_value_buffer(const ggml_tensor * tensor, ggml::hrx::ValueBufferBinding & binding) { + ggml_backend_hrx_buffer_context * context = nullptr; + size_t offset = 0; + if (!ggml_backend_hrx_tensor_binding(tensor, &context, &offset)) { + if (tensor == nullptr) { + return false; + } + const ggml_tensor * root = tensor->view_src != nullptr ? tensor->view_src : tensor; + ggml_backend_buffer_t buffer = root->buffer; + if (buffer == nullptr || !ggml_backend_buffer_is_host(buffer)) { + return false; + } + void * base = ggml_backend_buffer_get_base(buffer); + const size_t capacity = ggml_backend_buffer_get_size(buffer); + if (base == nullptr || tensor->data == nullptr || + static_cast(tensor->data) < static_cast(base)) { + return false; + } + const size_t host_offset = + static_cast(static_cast(tensor->data) - static_cast(base)); + if (host_offset > capacity || ggml_nbytes(tensor) > capacity - host_offset) { + return false; + } + const uint64_t buffer_address = static_cast(reinterpret_cast(buffer)); + const uint64_t base_address = static_cast(reinterpret_cast(base)); + binding.host_data = base; + binding.offset = host_offset; + binding.length = ggml_nbytes(tensor); + binding.identity = + buffer_address ^ (base_address + 0x9e3779b97f4a7c15ull + (buffer_address << 6) + (buffer_address >> 2)); + binding.generation = 1; + binding.capacity = capacity; + binding.weight = ggml_backend_buffer_get_usage(buffer) == GGML_BACKEND_BUFFER_USAGE_WEIGHTS; + return true; + } + ggml_backend_buffer_t buffer = tensor->view_src != nullptr ? tensor->view_src->buffer : tensor->buffer; + const bool directly_bindable = !ggml_backend_buffer_is_host(buffer) || context->direct_host_binding; + // Coherent HRX host allocations are directly device-addressable. Represent them with an HRX buffer handle so + // command-program preparation bypasses host materialization. Noncoherent host allocations remain host data. + binding.buffer = directly_bindable ? context->buffer : nullptr; + binding.host_data = directly_bindable ? nullptr : context->base; + binding.offset = offset; + binding.length = ggml_nbytes(tensor); + binding.identity = context->identity; + binding.generation = context->generation; + binding.capacity = buffer != nullptr ? buffer->size : 0; + binding.weight = buffer != nullptr && ggml_backend_buffer_get_usage(buffer) == GGML_BACKEND_BUFFER_USAGE_WEIGHTS; + return true; +} diff --git a/ggml/src/ggml-hrx/backend-buffer-binding.h b/ggml/src/ggml-hrx/backend-buffer-binding.h new file mode 100644 index 000000000000..59b925873f63 --- /dev/null +++ b/ggml/src/ggml-hrx/backend-buffer-binding.h @@ -0,0 +1,18 @@ +#pragma once + +#include "backend-context.h" +#include "ggml-backend.h" +#include "graph/value-map.h" + +#include + +struct ggml_tensor; + +ggml_backend_hrx_buffer_context * ggml_backend_hrx_buffer_context_from_buffer(ggml_backend_buffer_t buffer); +size_t ggml_backend_hrx_tensor_offset(const ggml_backend_hrx_buffer_context * context, const ggml_tensor * tensor); +void * ggml_backend_hrx_buffer_base(ggml_backend_buffer_t buffer); + +bool ggml_backend_hrx_tensor_binding(const ggml_tensor * tensor, + ggml_backend_hrx_buffer_context ** out_context, + size_t * out_offset); +bool ggml_backend_hrx_resolve_value_buffer(const ggml_tensor * tensor, ggml::hrx::ValueBufferBinding & binding); diff --git a/ggml/src/ggml-hrx/backend-context.h b/ggml/src/ggml-hrx/backend-context.h new file mode 100644 index 000000000000..3632037a7f1e --- /dev/null +++ b/ggml/src/ggml-hrx/backend-context.h @@ -0,0 +1,78 @@ +// Copyright 2026 The HRX Authors +// SPDX-License-Identifier: Apache-2.0 + +#pragma once + +#include "ggml-backend-impl.h" +#include "graph/value-map.h" +#include "runtime/graph-program-cache.h" +#include "runtime/host-buffer-registry.h" +#include "runtime/host-memory.h" +#include "runtime/kernel-executable-cache.h" +#include "runtime/prepared-command-program-cache.h" +#include "runtime/transient-arena.h" + +#include +#include +#include +#include +#include +#include +#include + +struct ggml_tensor; + +struct ggml_backend_hrx_device_context; + +struct ggml_backend_hrx_buffer_type_context { + ggml_backend_hrx_device_context * device; + std::string name; + bool host_visible = false; +}; + +struct ggml_backend_hrx_buffer_context { + ggml_backend_hrx_device_context * device; + hrx_buffer_t buffer; + uint8_t * base; + uint64_t identity; + uint64_t generation; + bool direct_host_binding; +}; + +struct ggml_backend_hrx_device_context { + hrx_device_t device = nullptr; + std::string name; + std::string description; + std::string architecture; + size_t memory_total = 0; + bool use_direct_host_bindings = false; + ggml_backend_buffer_type buft = {}; + ggml_backend_hrx_buffer_type_context buft_context = {}; + ggml_backend_buffer_type host_buft = {}; + ggml_backend_hrx_buffer_type_context host_buft_context = {}; + ggml::hrx::HostBufferRegistry host_buffers; + std::atomic synchronous_upload_fallbacks{ 0 }; + std::atomic synchronous_download_fallbacks{ 0 }; + std::mutex buffer_stream_mutex; + hrx_stream_t buffer_stream = nullptr; +}; + +struct ggml_backend_hrx_context { + ggml_backend_hrx_device_context * device; + hrx_stream_t stream; + ggml::hrx::KernelExecutableCache kernel_executables; + ggml::hrx::GraphProgramCache graph_programs; + ggml::hrx::PreparedCommandProgramCache prepared_programs; + ggml::hrx::TransientArena transient_arena; + ggml::hrx::HostTransferManager host_transfers; + ggml::hrx::HostWeightCache host_weights; + std::string name; +}; + +struct ggml_backend_hrx_reg_context { + bool initialized = false; + std::vector> device_contexts; + std::vector devices; + + ~ggml_backend_hrx_reg_context(); +}; diff --git a/ggml/src/ggml-hrx/cmake/BundleRuntime.cmake b/ggml/src/ggml-hrx/cmake/BundleRuntime.cmake new file mode 100644 index 000000000000..91853b38ad18 --- /dev/null +++ b/ggml/src/ggml-hrx/cmake/BundleRuntime.cmake @@ -0,0 +1,204 @@ +# Capture the module directory while it is the active list file. On CMake 3.14, +# CMAKE_CURRENT_LIST_DIR inside a function refers to the function's call site. +set(_GGML_HRX_BUNDLE_RUNTIME_DIR "${CMAKE_CURRENT_LIST_DIR}") + +# Find every entry matching PATTERN in exactly one of SEARCH_DIRS. +# +# OUT_PATHS receives the sorted matching paths. Missing families are fatal. +function(_ggml_hrx_find_bundle_entries OUT_PATHS LABEL PATTERN SEARCH_DIRS) + set(GGML_HRX_SELECTED_PATHS) + set(GGML_HRX_SELECTED_DIR "") + foreach(GGML_HRX_SEARCH_DIR IN LISTS SEARCH_DIRS) + file(GLOB GGML_HRX_DIR_MATCHES + CONFIGURE_DEPENDS + LIST_DIRECTORIES FALSE + "${GGML_HRX_SEARCH_DIR}/${PATTERN}") + if (GGML_HRX_DIR_MATCHES) + if (NOT GGML_HRX_SELECTED_DIR STREQUAL "") + message(FATAL_ERROR "GGML_HRX_BUNDLE_RUNTIME_LIBS found ${LABEL} in multiple source directories: ${GGML_HRX_SELECTED_DIR};${GGML_HRX_SEARCH_DIR}") + endif() + set(GGML_HRX_SELECTED_DIR "${GGML_HRX_SEARCH_DIR}") + set(GGML_HRX_SELECTED_PATHS ${GGML_HRX_DIR_MATCHES}) + endif() + endforeach() + + if (NOT GGML_HRX_SELECTED_PATHS) + message(FATAL_ERROR "GGML_HRX_BUNDLE_RUNTIME_LIBS could not find required ${LABEL} matching ${PATTERN}. Searched: ${SEARCH_DIRS}") + endif() + + foreach(GGML_HRX_SELECTED_PATH IN LISTS GGML_HRX_SELECTED_PATHS) + if (IS_SYMLINK "${GGML_HRX_SELECTED_PATH}") + if (NOT EXISTS "${GGML_HRX_SELECTED_PATH}") + message(FATAL_ERROR "GGML_HRX_BUNDLE_RUNTIME_LIBS matched a broken symlink: ${GGML_HRX_SELECTED_PATH}") + endif() + elseif(NOT EXISTS "${GGML_HRX_SELECTED_PATH}") + message(FATAL_ERROR "GGML_HRX_BUNDLE_RUNTIME_LIBS matched a nonexistent file: ${GGML_HRX_SELECTED_PATH}") + endif() + endforeach() + list(LENGTH GGML_HRX_SELECTED_PATHS GGML_HRX_SELECTED_COUNT) + message(STATUS " ${LABEL}: ${GGML_HRX_SELECTED_DIR} (${GGML_HRX_SELECTED_COUNT} entries)") + + set(${OUT_PATHS} "${GGML_HRX_SELECTED_PATHS}" PARENT_SCOPE) +endfunction() + +# Keep the target's existing relative build and install RPATH entries and drop +# absolute entries. Append the paths needed by the adjacent runtime bundle. The +# resulting list is returned through OUT_VAR. +function(_ggml_hrx_collect_portable_rpath TARGET_NAME OUT_VAR) + set(GGML_HRX_PORTABLE_RPATH) + foreach(GGML_HRX_RPATH_PROPERTY BUILD_RPATH INSTALL_RPATH) + get_target_property(GGML_HRX_RPATH_ENTRIES "${TARGET_NAME}" "${GGML_HRX_RPATH_PROPERTY}") + if (NOT GGML_HRX_RPATH_ENTRIES) + continue() + endif() + foreach(GGML_HRX_RPATH_ENTRY IN LISTS GGML_HRX_RPATH_ENTRIES) + if (GGML_HRX_RPATH_ENTRY STREQUAL "") + continue() + endif() + if (IS_ABSOLUTE "${GGML_HRX_RPATH_ENTRY}") + continue() + endif() + list(APPEND GGML_HRX_PORTABLE_RPATH "${GGML_HRX_RPATH_ENTRY}") + endforeach() + endforeach() + list(APPEND GGML_HRX_PORTABLE_RPATH + "$ORIGIN" + "$ORIGIN/rocm_sysdeps/lib") + list(REMOVE_DUPLICATES GGML_HRX_PORTABLE_RPATH) + set(${OUT_VAR} "${GGML_HRX_PORTABLE_RPATH}" PARENT_SCOPE) +endfunction() + +# Add matching build-tree copy and install rules for SOURCES. +# RELATIVE_DESTINATION is appended below both backend destinations. +function(_ggml_hrx_add_bundle_rules TARGET_NAME SOURCES RELATIVE_DESTINATION INSTALL_BASE COPY_SCRIPT) + set(GGML_HRX_BUILD_DESTINATION "$") + set(GGML_HRX_INSTALL_DESTINATION "${INSTALL_BASE}") + if (NOT RELATIVE_DESTINATION STREQUAL "") + string(APPEND GGML_HRX_BUILD_DESTINATION "/${RELATIVE_DESTINATION}") + string(APPEND GGML_HRX_INSTALL_DESTINATION "/${RELATIVE_DESTINATION}") + endif() + + if (NOT RELATIVE_DESTINATION STREQUAL "") + add_custom_command(TARGET "${TARGET_NAME}" POST_BUILD + COMMAND "${CMAKE_COMMAND}" -E make_directory + "${GGML_HRX_BUILD_DESTINATION}" + VERBATIM) + endif() + add_custom_command(TARGET "${TARGET_NAME}" POST_BUILD + COMMAND "${CMAKE_COMMAND}" + "-DGGML_HRX_BUNDLE_SOURCES=${SOURCES}" + "-DGGML_HRX_BUNDLE_DESTINATION=${GGML_HRX_BUILD_DESTINATION}" + -P "${COPY_SCRIPT}" + VERBATIM) + install(FILES ${SOURCES} + DESTINATION "${GGML_HRX_INSTALL_DESTINATION}") +endfunction() + +# Discover, copy, and install the HRX runtime dependencies for TARGET_NAME. +function(ggml_hrx_bundle_runtime TARGET_NAME) + set(GGML_HRX_COPY_SCRIPT "${_GGML_HRX_BUNDLE_RUNTIME_DIR}/copy_bundle_entry.cmake") + if (NOT TARGET "${TARGET_NAME}") + message(FATAL_ERROR "ggml_hrx_bundle_runtime target does not exist: ${TARGET_NAME}") + endif() + if (NOT CMAKE_SYSTEM_NAME STREQUAL "Linux") + message(FATAL_ERROR "GGML_HRX_BUNDLE_RUNTIME_LIBS is currently implemented for Linux only") + endif() + if (NOT BUILD_SHARED_LIBS) + message(FATAL_ERROR "GGML_HRX_BUNDLE_RUNTIME_LIBS requires BUILD_SHARED_LIBS=ON") + endif() + + set(GGML_HRX_BUNDLE_SEARCH_DIRS) + foreach(GGML_HRX_BUNDLE_SEARCH_DIR IN LISTS GGML_HRX_BUNDLE_LIBRARY_DIRS) + if (GGML_HRX_BUNDLE_SEARCH_DIR STREQUAL "") + message(FATAL_ERROR "GGML_HRX_BUNDLE_LIBRARY_DIRS contains an empty directory entry") + endif() + get_filename_component(GGML_HRX_BUNDLE_SEARCH_DIR_ABSOLUTE "${GGML_HRX_BUNDLE_SEARCH_DIR}" ABSOLUTE BASE_DIR "${CMAKE_CURRENT_SOURCE_DIR}") + if (NOT IS_DIRECTORY "${GGML_HRX_BUNDLE_SEARCH_DIR_ABSOLUTE}") + message(FATAL_ERROR "GGML_HRX_BUNDLE_LIBRARY_DIRS contains a nonexistent directory: ${GGML_HRX_BUNDLE_SEARCH_DIR}") + endif() + get_filename_component(GGML_HRX_BUNDLE_SEARCH_DIR_CANONICAL "${GGML_HRX_BUNDLE_SEARCH_DIR_ABSOLUTE}" REALPATH) + list(APPEND GGML_HRX_BUNDLE_SEARCH_DIRS "${GGML_HRX_BUNDLE_SEARCH_DIR_CANONICAL}") + endforeach() + list(REMOVE_DUPLICATES GGML_HRX_BUNDLE_SEARCH_DIRS) + if (NOT GGML_HRX_BUNDLE_SEARCH_DIRS) + message(FATAL_ERROR "GGML_HRX_BUNDLE_LIBRARY_DIRS must list at least one directory when GGML_HRX_BUNDLE_RUNTIME_LIBS=ON") + endif() + + message(STATUS "HRX runtime bundling search directories: ${GGML_HRX_BUNDLE_SEARCH_DIRS}") + message(STATUS "HRX runtime bundle selection:") + _ggml_hrx_find_bundle_entries( + GGML_HRX_BUNDLE_HRX_LIBS + "libhrx" "libhrx.so*" "${GGML_HRX_BUNDLE_SEARCH_DIRS}") + _ggml_hrx_find_bundle_entries( + GGML_HRX_BUNDLE_LOOMC_LIBS + "libloomc" "libloomc.so*" "${GGML_HRX_BUNDLE_SEARCH_DIRS}") + _ggml_hrx_find_bundle_entries( + GGML_HRX_BUNDLE_HSA_RUNTIME_LIBS + "libhsa-runtime64" "libhsa-runtime64.so*" "${GGML_HRX_BUNDLE_SEARCH_DIRS}") + _ggml_hrx_find_bundle_entries( + GGML_HRX_BUNDLE_HSA_AQLPROFILE_LIBS + "libhsa-amd-aqlprofile64" "libhsa-amd-aqlprofile64.so*" "${GGML_HRX_BUNDLE_SEARCH_DIRS}") + _ggml_hrx_find_bundle_entries( + GGML_HRX_BUNDLE_ROCPROFILER_REGISTER_LIBS + "librocprofiler-register" "librocprofiler-register.so*" "${GGML_HRX_BUNDLE_SEARCH_DIRS}") + _ggml_hrx_find_bundle_entries( + GGML_HRX_BUNDLE_OMP_LIBS + "libomp" "libomp.so*" "${GGML_HRX_BUNDLE_SEARCH_DIRS}") + # HRX runtime bundles require the ROCm sysdeps overlay. + _ggml_hrx_find_bundle_entries( + GGML_HRX_BUNDLE_SYSDEP_LIBS + "rocm_sysdeps/lib overlay" "rocm_sysdeps/lib/*.so*" "${GGML_HRX_BUNDLE_SEARCH_DIRS}") + + set(GGML_HRX_BUNDLE_MAIN_LIBS + ${GGML_HRX_BUNDLE_HRX_LIBS} + ${GGML_HRX_BUNDLE_LOOMC_LIBS} + ${GGML_HRX_BUNDLE_HSA_RUNTIME_LIBS} + ${GGML_HRX_BUNDLE_HSA_AQLPROFILE_LIBS} + ${GGML_HRX_BUNDLE_ROCPROFILER_REGISTER_LIBS} + ${GGML_HRX_BUNDLE_OMP_LIBS} + ) + set_property(TARGET "${TARGET_NAME}" APPEND PROPERTY LINK_DEPENDS + ${GGML_HRX_BUNDLE_MAIN_LIBS} + ${GGML_HRX_BUNDLE_SYSDEP_LIBS} + "${GGML_HRX_COPY_SCRIPT}") + + # GGML_BACKEND_DL builds backends as runtime-loaded modules instead of + # normally linked libraries. Install dependencies next to the backend: + # modules use GGML_BACKEND_DIR or bin; linked backends use the standard + # library directory. + if (GGML_BACKEND_DL) + if (GGML_BACKEND_DIR) + set(GGML_HRX_BUNDLE_INSTALL_DIR "${GGML_BACKEND_DIR}") + else() + set(GGML_HRX_BUNDLE_INSTALL_DIR "${CMAKE_INSTALL_BINDIR}") + endif() + else() + set(GGML_HRX_BUNDLE_INSTALL_DIR "${CMAKE_INSTALL_LIBDIR}") + endif() + message(STATUS "HRX runtime bundle install directory: ${GGML_HRX_BUNDLE_INSTALL_DIR}") + + # Give build and install artifacts the same portable RUNPATH. Keep relative + # entries, add adjacent bundle directories, and prevent absolute HRX/ROCm + # link directories. + _ggml_hrx_collect_portable_rpath("${TARGET_NAME}" GGML_HRX_PORTABLE_RPATH) + set_target_properties("${TARGET_NAME}" PROPERTIES + BUILD_RPATH "${GGML_HRX_PORTABLE_RPATH}" + INSTALL_RPATH "${GGML_HRX_PORTABLE_RPATH}" + BUILD_WITH_INSTALL_RPATH TRUE + INSTALL_RPATH_USE_LINK_PATH FALSE + ) + message(STATUS "HRX backend RUNPATH: ${GGML_HRX_PORTABLE_RPATH}") + + _ggml_hrx_add_bundle_rules( + "${TARGET_NAME}" + "${GGML_HRX_BUNDLE_MAIN_LIBS}" + "" + "${GGML_HRX_BUNDLE_INSTALL_DIR}" + "${GGML_HRX_COPY_SCRIPT}") + _ggml_hrx_add_bundle_rules( + "${TARGET_NAME}" + "${GGML_HRX_BUNDLE_SYSDEP_LIBS}" + "rocm_sysdeps/lib" + "${GGML_HRX_BUNDLE_INSTALL_DIR}" + "${GGML_HRX_COPY_SCRIPT}") +endfunction() diff --git a/ggml/src/ggml-hrx/cmake/copy_bundle_entry.cmake b/ggml/src/ggml-hrx/cmake/copy_bundle_entry.cmake new file mode 100644 index 000000000000..70bb5bbe52ad --- /dev/null +++ b/ggml/src/ggml-hrx/cmake/copy_bundle_entry.cmake @@ -0,0 +1,33 @@ +if (NOT DEFINED GGML_HRX_BUNDLE_SOURCES OR NOT DEFINED GGML_HRX_BUNDLE_DESTINATION) + message(FATAL_ERROR "HRX bundle copy requires sources and destination") +endif() + +foreach(GGML_HRX_BUNDLE_SOURCE IN LISTS GGML_HRX_BUNDLE_SOURCES) + get_filename_component(GGML_HRX_BUNDLE_NAME "${GGML_HRX_BUNDLE_SOURCE}" NAME) + set(GGML_HRX_BUNDLE_DEST "${GGML_HRX_BUNDLE_DESTINATION}/${GGML_HRX_BUNDLE_NAME}") + if (IS_SYMLINK "${GGML_HRX_BUNDLE_SOURCE}") + if (NOT EXISTS "${GGML_HRX_BUNDLE_SOURCE}") + message(FATAL_ERROR "HRX bundle source is a broken symlink: ${GGML_HRX_BUNDLE_SOURCE}") + endif() + file(READ_SYMLINK "${GGML_HRX_BUNDLE_SOURCE}" GGML_HRX_BUNDLE_LINK_TARGET) + elseif (EXISTS "${GGML_HRX_BUNDLE_SOURCE}") + if (IS_DIRECTORY "${GGML_HRX_BUNDLE_SOURCE}") + message(FATAL_ERROR "HRX bundle source is not a regular file: ${GGML_HRX_BUNDLE_SOURCE}") + endif() + else() + message(FATAL_ERROR "HRX bundle source no longer exists: ${GGML_HRX_BUNDLE_SOURCE}") + endif() + + get_filename_component(GGML_HRX_BUNDLE_SOURCE_ABSOLUTE "${GGML_HRX_BUNDLE_SOURCE}" ABSOLUTE) + get_filename_component(GGML_HRX_BUNDLE_DEST_ABSOLUTE "${GGML_HRX_BUNDLE_DEST}" ABSOLUTE) + if ("${GGML_HRX_BUNDLE_SOURCE_ABSOLUTE}" STREQUAL "${GGML_HRX_BUNDLE_DEST_ABSOLUTE}") + continue() + endif() + + file(REMOVE "${GGML_HRX_BUNDLE_DEST}") + if (IS_SYMLINK "${GGML_HRX_BUNDLE_SOURCE}") + file(CREATE_LINK "${GGML_HRX_BUNDLE_LINK_TARGET}" "${GGML_HRX_BUNDLE_DEST}" SYMBOLIC) + else() + file(COPY "${GGML_HRX_BUNDLE_SOURCE}" DESTINATION "${GGML_HRX_BUNDLE_DESTINATION}") + endif() +endforeach() diff --git a/ggml/src/ggml-hrx/dispatch/command-plan-metadata.cpp b/ggml/src/ggml-hrx/dispatch/command-plan-metadata.cpp new file mode 100644 index 000000000000..eac956c644ac --- /dev/null +++ b/ggml/src/ggml-hrx/dispatch/command-plan-metadata.cpp @@ -0,0 +1,143 @@ +#include "command-plan-metadata.h" + +#include +#include + +namespace ggml::hrx { +namespace { + +static bool metadata_matches(const CommandPlanResourceMetadata & lhs, const CommandPlanResourceMetadata & rhs) { + return lhs.kind == rhs.kind && lhs.size == rhs.size && + std::memcmp(lhs.bytes.data(), rhs.bytes.data(), lhs.size) == 0; +} + +static bool generated_resource_matches(const CommandPlanGeneratedResource & lhs, + const CommandPlanGeneratedResource & rhs) { + return lhs.source_value == rhs.source_value && lhs.role == rhs.role && lhs.generated_value == rhs.generated_value && + lhs.byte_count == rhs.byte_count && metadata_matches(lhs.metadata, rhs.metadata); +} + +static bool alternate_value_matches(const CommandPlanAlternateValue & lhs, const CommandPlanAlternateValue & rhs) { + return lhs.graph_value == rhs.graph_value && lhs.alternate_value == rhs.alternate_value && lhs.type == rhs.type && + lhs.byte_count == rhs.byte_count && lhs.name == rhs.name; +} + +static bool moe_routing_bundle_matches(const CommandPlanMoeRoutingBundle & lhs, + const CommandPlanMoeRoutingBundle & rhs) { + return lhs.route_ids == rhs.route_ids && lhs.route_weights == rhs.route_weights && + lhs.expert_table == rhs.expert_table && lhs.partition_table == rhs.partition_table && + lhs.expert_table_byte_count == rhs.expert_table_byte_count && + lhs.partition_table_byte_count == rhs.partition_table_byte_count && lhs.token_count == rhs.token_count && + lhs.route_count == rhs.route_count && lhs.route_stride == rhs.route_stride && + lhs.expert_count == rhs.expert_count; +} + +} // namespace + +void CommandPlanMetadata::clear() { + generated_resources_.clear(); + alternate_values_.clear(); + moe_routing_bundles_.clear(); +} + +bool CommandPlanMetadata::append(CommandPlanMetadata && other, Status & status) { + for (CommandPlanGeneratedResource & resource : other.generated_resources_) { + if (!append_generated_resource(std::move(resource), status)) { + return false; + } + } + for (CommandPlanAlternateValue & alternate : other.alternate_values_) { + if (!append_alternate_value(std::move(alternate), status)) { + return false; + } + } + for (CommandPlanMoeRoutingBundle & bundle : other.moe_routing_bundles_) { + if (!append_moe_routing_bundle(std::move(bundle), status)) { + return false; + } + } + return true; +} + +bool CommandPlanMetadata::append_generated_resource(CommandPlanGeneratedResource resource, Status & status) { + for (const CommandPlanGeneratedResource & existing : generated_resources_) { + if (existing.source_value == resource.source_value && existing.role == resource.role) { + if (generated_resource_matches(existing, resource)) { + return true; + } + status.log("conflicting generated resource for source value %d role %d", resource.source_value.value, + static_cast(resource.role)); + return false; + } + } + generated_resources_.push_back(std::move(resource)); + return true; +} + +bool CommandPlanMetadata::append_alternate_value(CommandPlanAlternateValue alternate, Status & status) { + for (const CommandPlanAlternateValue & existing : alternate_values_) { + if (existing.graph_value == alternate.graph_value) { + if (alternate_value_matches(existing, alternate)) { + return true; + } + status.log("conflicting alternate value for graph value %d", alternate.graph_value.value); + return false; + } + } + alternate_values_.push_back(std::move(alternate)); + return true; +} + +bool CommandPlanMetadata::append_moe_routing_bundle(CommandPlanMoeRoutingBundle bundle, Status & status) { + for (const CommandPlanMoeRoutingBundle & existing : moe_routing_bundles_) { + if (existing.route_ids == bundle.route_ids) { + if (moe_routing_bundle_matches(existing, bundle)) { + return true; + } + status.log("conflicting MoE routing bundle for route ids value %d", bundle.route_ids.value); + return false; + } + } + moe_routing_bundles_.push_back(std::move(bundle)); + return true; +} + +const CommandPlanGeneratedResource * CommandPlanMetadata::find_generated_resource(ValueId source_value, + GeneratedResourceRole role) const { + for (const CommandPlanGeneratedResource & resource : generated_resources_) { + if (resource.source_value == source_value && resource.role == role) { + return &resource; + } + } + return nullptr; +} + +const CommandPlanAlternateValue * CommandPlanMetadata::find_alternate_value(ValueId graph_value) const { + for (const CommandPlanAlternateValue & alternate : alternate_values_) { + if (alternate.graph_value == graph_value) { + return &alternate; + } + } + return nullptr; +} + +const CommandPlanAlternateValue * CommandPlanMetadata::find_alternate_value(ValueId graph_value, + ggml_type type, + size_t byte_count) const { + const CommandPlanAlternateValue * alternate = find_alternate_value(graph_value); + if (alternate == nullptr || alternate->type != type || alternate->byte_count != byte_count) { + return nullptr; + } + return alternate; +} + +const CommandPlanMoeRoutingBundle * CommandPlanMetadata::find_moe_routing_bundle(ValueId route_ids) const { + for (const CommandPlanMoeRoutingBundle & bundle : moe_routing_bundles_) { + if (bundle.route_ids == route_ids) { + return &bundle; + } + } + return nullptr; +} + +} // namespace ggml::hrx diff --git a/ggml/src/ggml-hrx/dispatch/command-plan-metadata.h b/ggml/src/ggml-hrx/dispatch/command-plan-metadata.h new file mode 100644 index 000000000000..f6d469c2cdaa --- /dev/null +++ b/ggml/src/ggml-hrx/dispatch/command-plan-metadata.h @@ -0,0 +1,135 @@ +#pragma once + +#include "dispatch.h" +#include "status.h" + +#include +#include +#include +#include +#include +#include +#include + +namespace ggml::hrx { + +enum class GeneratedResourceRole { + MoeExpertTable, + MoePartitionTable, +}; + +enum class CommandPlanResourceMetadataKind { + None, + MoeRoutingResource, +}; + +struct MoeRoutingResourceMetadata { + int64_t token_count = 0; + int64_t route_count = 0; + int64_t route_stride = 0; + int64_t expert_count = 0; +}; + +template constexpr CommandPlanResourceMetadataKind command_plan_resource_metadata_kind() { + static_assert(sizeof(T) == 0, "unsupported command plan resource metadata type"); + return CommandPlanResourceMetadataKind::None; +} + +template <> +constexpr CommandPlanResourceMetadataKind command_plan_resource_metadata_kind() { + return CommandPlanResourceMetadataKind::MoeRoutingResource; +} + +struct CommandPlanResourceMetadata { + static constexpr size_t kMaxBytes = 64; + + CommandPlanResourceMetadataKind kind = CommandPlanResourceMetadataKind::None; + size_t size = 0; + alignas(std::max_align_t) std::array bytes = {}; + + template bool read(T & value) const { + static_assert(std::is_trivially_copyable::value, "metadata payload must be trivially copyable"); + if (kind != command_plan_resource_metadata_kind() || size != sizeof(T)) { + return false; + } + std::memcpy(&value, bytes.data(), sizeof(T)); + return true; + } +}; + +template CommandPlanResourceMetadata make_command_plan_resource_metadata(const T & value) { + static_assert(std::is_trivially_copyable::value, "metadata payload must be trivially copyable"); + static_assert(sizeof(T) <= CommandPlanResourceMetadata::kMaxBytes, "metadata payload is too large"); + + CommandPlanResourceMetadata metadata; + metadata.kind = command_plan_resource_metadata_kind(); + metadata.size = sizeof(T); + std::memcpy(metadata.bytes.data(), &value, sizeof(T)); + return metadata; +} + +struct CommandPlanGeneratedResource { + ValueId source_value; + GeneratedResourceRole role = GeneratedResourceRole::MoeExpertTable; + ValueId generated_value; + size_t byte_count = 0; + CommandPlanResourceMetadata metadata; +}; + +struct CommandPlanAlternateValue { + ValueId graph_value; + ValueId alternate_value; + ggml_type type = GGML_TYPE_COUNT; + size_t byte_count = 0; + std::string name; +}; + +struct CommandPlanMoeRoutingBundle { + ValueId route_ids; + ValueId route_weights; + ValueId expert_table; + ValueId partition_table; + size_t expert_table_byte_count = 0; + size_t partition_table_byte_count = 0; + int64_t token_count = 0; + int64_t route_count = 0; + int64_t route_stride = 0; + int64_t expert_count = 0; +}; + +class CommandPlanMetadata { + public: + void clear(); + + bool append(CommandPlanMetadata && other, Status & status); + + bool append_generated_resource(CommandPlanGeneratedResource resource, Status & status); + + bool append_alternate_value(CommandPlanAlternateValue alternate, Status & status); + + bool append_moe_routing_bundle(CommandPlanMoeRoutingBundle bundle, Status & status); + + const CommandPlanGeneratedResource * find_generated_resource(ValueId source_value, + GeneratedResourceRole role) const; + + const CommandPlanAlternateValue * find_alternate_value(ValueId graph_value) const; + + const CommandPlanAlternateValue * find_alternate_value(ValueId graph_value, + ggml_type type, + size_t byte_count) const; + + const CommandPlanMoeRoutingBundle * find_moe_routing_bundle(ValueId route_ids) const; + + const std::vector & generated_resources() const { return generated_resources_; } + + const std::vector & alternate_values() const { return alternate_values_; } + + const std::vector & moe_routing_bundles() const { return moe_routing_bundles_; } + + private: + std::vector generated_resources_; + std::vector alternate_values_; + std::vector moe_routing_bundles_; +}; + +} // namespace ggml::hrx diff --git a/ggml/src/ggml-hrx/dispatch/command-plan.h b/ggml/src/ggml-hrx/dispatch/command-plan.h new file mode 100644 index 000000000000..a1f1abf62168 --- /dev/null +++ b/ggml/src/ggml-hrx/dispatch/command-plan.h @@ -0,0 +1,115 @@ +#pragma once + +#include "command-plan-metadata.h" +#include "dispatch.h" +#include "graph/graph.h" +#include "status.h" + +#include +#include +#include +#include + +namespace ggml::hrx { + +struct CommandPlanTransient { + ValueId value; + std::string name; + size_t size = 0; + size_t alignment = 256; +}; + +struct CommandPlanConstantInitialization { + ValueId value; + std::string name; + size_t offset = 0; + std::vector data; +}; + +struct CommandPlanCompletionCounterRequest { + ValueId value; + std::string name; + uint32_t count = 0; +}; + +struct CommandPlan { + std::vector initialization_dispatches; + std::vector dispatches; + std::vector transients; + std::vector constant_initializations; + std::vector completion_counter_requests; + CommandPlanMetadata metadata; + Status status; + + bool valid() const { return status.success(); } +}; + +inline const CommandPlanAlternateValue * find_alternate_value(const CommandPlan & plan, ValueId graph_value) { + return plan.metadata.find_alternate_value(graph_value); +} + +inline const CommandPlanAlternateValue * find_alternate_value(const CommandPlan & plan, + ValueId graph_value, + ggml_type type, + size_t byte_count) { + return plan.metadata.find_alternate_value(graph_value, type, byte_count); +} + +inline bool same_full_value_range(const Value & lhs, const Value & rhs) { + return lhs.storage == rhs.storage && lhs.storage_offset == rhs.storage_offset && lhs.byte_count == rhs.byte_count; +} + +inline const CommandPlanAlternateValue * find_alternate_value(const Graph & graph, + const CommandPlan & plan, + ValueId graph_value, + ggml_type type, + size_t byte_count) { + const CommandPlanAlternateValue * exact = find_alternate_value(plan, graph_value, type, byte_count); + if (exact != nullptr) { + return exact; + } + + const Value * value = graph.values().find(graph_value); + if (value == nullptr) { + return nullptr; + } + + auto find_if_same_range = [&](ValueId candidate_id) -> const CommandPlanAlternateValue * { + const Value * candidate = graph.values().find(candidate_id); + if (candidate == nullptr || !same_full_value_range(*value, *candidate)) { + return nullptr; + } + return find_alternate_value(plan, candidate_id, type, byte_count); + }; + + ValueId alias = value->alias_source; + for (size_t i = 0; alias.value >= 0 && i < graph.values().size(); ++i) { + const CommandPlanAlternateValue * alternate = find_if_same_range(alias); + if (alternate != nullptr) { + return alternate; + } + const Value * alias_value = graph.values().find(alias); + if (alias_value == nullptr) { + break; + } + alias = alias_value->alias_source; + } + + const CommandPlanAlternateValue * root_alternate = find_if_same_range(value->storage_root); + if (root_alternate != nullptr) { + return root_alternate; + } + + for (const CommandPlanAlternateValue & alternate : plan.metadata.alternate_values()) { + if (alternate.type != type || alternate.byte_count != byte_count) { + continue; + } + const Value * alternate_value = graph.values().find(alternate.graph_value); + if (alternate_value != nullptr && same_full_value_range(*value, *alternate_value)) { + return &alternate; + } + } + return nullptr; +} + +} // namespace ggml::hrx diff --git a/ggml/src/ggml-hrx/dispatch/command-program-bindings.cpp b/ggml/src/ggml-hrx/dispatch/command-program-bindings.cpp new file mode 100644 index 000000000000..ddaf460a293a --- /dev/null +++ b/ggml/src/ggml-hrx/dispatch/command-program-bindings.cpp @@ -0,0 +1,90 @@ +#include "command-program-bindings.h" + +#include +#include + +namespace ggml::hrx { +namespace { + +static void mix_hash(uint64_t & hash, uint64_t value) { + hash ^= value; + hash *= UINT64_C(1099511628211); +} + +} // namespace + +CommandProgramBindings CommandProgramBindings::from_value_map(const ValueMap & values) { + std::vector bindings; + CommandProgramBindings result; + for (const ValueId id : values.external_value_ids()) { + const Value * value = values.find(id); + if (value == nullptr) { + result.status.log("external value %d does not exist", id.value); + continue; + } + const std::optional buffer = values.resolve_buffer_binding(id); + if (!buffer.has_value()) { + result.status.log("external value %d is not bound", id.value); + continue; + } + bindings.push_back({ value->id, buffer->buffer, buffer->offset, buffer->length, buffer->identity, + buffer->generation, buffer->capacity, buffer->host_data, buffer->weight }); + } + return from_bindings(std::move(bindings), result.status); +} + +CommandProgramBindings CommandProgramBindings::from_bindings(std::vector bindings, + const Status & errors) { + CommandProgramBindings result; + result.status.append(errors); + result.bindings_ = std::move(bindings); + for (const CommandProgramBinding & binding : result.bindings_) { + if (binding.buffer == nullptr && binding.host_data == nullptr) { + result.status.log("external value %d has a null binding", binding.value.value); + } + // A zero-byte binding is not an error: ggml legitimately produces empty tensors + // (e.g. a GET_ROWS over an empty id list during decode), and a dispatch that only + // touches them is a no-op. + } + return result; +} + +const CommandProgramBinding * CommandProgramBindings::find(ValueId value) const { + for (const CommandProgramBinding & binding : bindings_) { + if (binding.value == value) { + return &binding; + } + } + return nullptr; +} + +CommandProgramBindingsHash command_program_bindings_hash(const CommandProgramBindings & bindings) { + uint64_t hash = UINT64_C(1469598103934665603); + mix_hash(hash, UINT64_C(0x6872782d62696e64)); + for (const CommandProgramBinding & binding : bindings.bindings()) { + mix_hash(hash, static_cast(static_cast(binding.value.value))); + mix_hash(hash, binding.host_data != nullptr ? 1 : 0); + mix_hash(hash, binding.identity); + mix_hash(hash, binding.generation); + mix_hash(hash, static_cast(binding.capacity)); + mix_hash(hash, static_cast(binding.offset)); + mix_hash(hash, static_cast(binding.length)); + mix_hash(hash, binding.weight ? 1 : 0); + } + mix_hash(hash, static_cast(bindings.bindings().size())); + return { hash }; +} + +CommandProgramBindingsFingerprint command_program_bindings_fingerprint(const CommandProgramBindings & bindings) { + std::ostringstream out; + out << "hrx-bindings-v1"; + for (const CommandProgramBinding & binding : bindings.bindings()) { + out << "|value=" << binding.value.value << "|kind=" << (binding.host_data != nullptr ? "host" : "device") + << "|identity=" << binding.identity << "|generation=" << binding.generation + << "|capacity=" << binding.capacity << "|offset=" << binding.offset << "|length=" << binding.length + << "|weight=" << (binding.weight ? 1 : 0); + } + return { out.str() }; +} + +} // namespace ggml::hrx diff --git a/ggml/src/ggml-hrx/dispatch/command-program-bindings.h b/ggml/src/ggml-hrx/dispatch/command-program-bindings.h new file mode 100644 index 000000000000..5f71ab57f985 --- /dev/null +++ b/ggml/src/ggml-hrx/dispatch/command-program-bindings.h @@ -0,0 +1,58 @@ +#pragma once + +#include "graph/value-map.h" +#include "status.h" + +#include +#include +#include +#include + +namespace ggml::hrx { + +struct CommandProgramBinding { + // A buffer is directly bindable by an HRX command program. Host data requires residency or staging before + // execution. These are alternate storage forms and should not both be populated. + ValueId value; + hrx_buffer_t buffer = nullptr; + size_t offset = 0; + size_t length = 0; + uint64_t identity = 0; + uint64_t generation = 0; + size_t capacity = 0; + void * host_data = nullptr; + bool weight = false; + + bool requires_materialization() const { return host_data != nullptr; } +}; + +struct CommandProgramBindingsFingerprint { + std::string value; +}; + +struct CommandProgramBindingsHash { + uint64_t value = 0; +}; + +class CommandProgramBindings { + public: + static CommandProgramBindings from_value_map(const ValueMap & values); + static CommandProgramBindings from_bindings(std::vector bindings, + const Status & errors = {}); + + const CommandProgramBinding * find(ValueId value) const; + + const std::vector & bindings() const { return bindings_; } + + bool valid() const { return status.success(); } + + Status status; + + private: + std::vector bindings_; +}; + +CommandProgramBindingsHash command_program_bindings_hash(const CommandProgramBindings & bindings); +CommandProgramBindingsFingerprint command_program_bindings_fingerprint(const CommandProgramBindings & bindings); + +} // namespace ggml::hrx diff --git a/ggml/src/ggml-hrx/dispatch/command-program-diagnostics.cpp b/ggml/src/ggml-hrx/dispatch/command-program-diagnostics.cpp new file mode 100644 index 000000000000..433204ca7910 --- /dev/null +++ b/ggml/src/ggml-hrx/dispatch/command-program-diagnostics.cpp @@ -0,0 +1,92 @@ +#include "command-program-diagnostics.h" + +#include + +namespace ggml::hrx { +namespace { + +template static std::string unknown_enum_name(Enum value) { + std::ostringstream out; + out << "Unknown(" << static_cast(value) << ")"; + return out.str(); +} + +static const char * binding_name(const CommandBinding & binding) { + return binding.name.empty() ? "" : binding.name.c_str(); +} + +} // namespace + +std::string command_kind_name(CommandKind kind) { + switch (kind) { + case CommandKind::Invalid: + return "Invalid"; + case CommandKind::Kernel: + return "Kernel"; + } + return unknown_enum_name(kind); +} + +std::string command_binding_origin_name(CommandBindingOrigin origin) { + switch (origin) { + case CommandBindingOrigin::GraphValue: + return "GraphValue"; + case CommandBindingOrigin::Transient: + return "Transient"; + case CommandBindingOrigin::ProgramConstant: + return "ProgramConstant"; + } + return unknown_enum_name(origin); +} + +std::string resource_access_name(ResourceAccess access) { + switch (access) { + case ResourceAccess::Read: + return "Read"; + case ResourceAccess::Write: + return "Write"; + case ResourceAccess::ReadWrite: + return "ReadWrite"; + } + return unknown_enum_name(access); +} + +std::string format_command_binding(const CommandBinding & binding) { + std::ostringstream out; + out << "binding " << binding_name(binding) << " value=" << binding.value.value + << " origin=" << command_binding_origin_name(binding.origin) + << " access=" << resource_access_name(binding.access) << " range=[" << binding.offset << ", " + << binding.offset + binding.length << ")"; + return out.str(); +} + +std::string format_command(const Command & command) { + std::ostringstream out; + out << "command " << command.ordinal << " kind=" << command_kind_name(command.kind) + << " kernel_id=" << command.kernel.kernel_id << " bindings=" << command.bindings.size() + << " deps=" << command.dependencies.size(); + return out.str(); +} + +std::string format_command_program(const CommandProgram & program) { + std::ostringstream out; + out << "command_program commands=" << program.commands.size() + << " transient_arena=" << program.transients.arena_size + << " transient_allocations=" << program.transients.allocations.size() + << " completion_counters=" << program.completion_counters.count << " completion_counter_range=[" + << program.completion_counters.arena_offset << ", " + << program.completion_counters.arena_offset + program.completion_counters.byte_count << ")"; + for (const Command & command : program.commands) { + out << '\n' << format_command(command); + for (const CommandBinding & binding : command.bindings) { + out << "\n " << format_command_binding(binding); + } + } + for (const TransientAllocation & allocation : program.transients.allocations) { + out << "\ntransient value=" << allocation.value.value << " range=[" << allocation.arena_offset << ", " + << allocation.arena_offset + allocation.size << ") alignment=" << allocation.alignment; + } + return out.str(); +} + +} // namespace ggml::hrx diff --git a/ggml/src/ggml-hrx/dispatch/command-program-diagnostics.h b/ggml/src/ggml-hrx/dispatch/command-program-diagnostics.h new file mode 100644 index 000000000000..c961db905115 --- /dev/null +++ b/ggml/src/ggml-hrx/dispatch/command-program-diagnostics.h @@ -0,0 +1,17 @@ +#pragma once + +#include "command-program.h" + +#include + +namespace ggml::hrx { + +std::string command_kind_name(CommandKind kind); +std::string command_binding_origin_name(CommandBindingOrigin origin); +std::string resource_access_name(ResourceAccess access); + +std::string format_command_binding(const CommandBinding & binding); +std::string format_command(const Command & command); +std::string format_command_program(const CommandProgram & program); + +} // namespace ggml::hrx diff --git a/ggml/src/ggml-hrx/dispatch/command-program-resolver.cpp b/ggml/src/ggml-hrx/dispatch/command-program-resolver.cpp new file mode 100644 index 000000000000..18805390c903 --- /dev/null +++ b/ggml/src/ggml-hrx/dispatch/command-program-resolver.cpp @@ -0,0 +1,128 @@ +#include "command-program-resolver.h" + +#include "command-program-diagnostics.h" + +#include + +namespace ggml::hrx { +namespace { + +static Status resolve_command_binding(const Command & command, + const CommandProgram & program, + const CommandBinding & binding, + const CommandProgramBindings & bindings, + const TransientArenaAllocationRef * transient_arena, + ResolvedBufferRef & ref) { + Status status; + const std::string command_context = format_command(command); + const std::string binding_context = format_command_binding(binding); + if (binding.length == 0) { + status.log("%s %s has an empty range", command_context.c_str(), binding_context.c_str()); + return status; + } + switch (binding.origin) { + case CommandBindingOrigin::GraphValue: + { + const CommandProgramBinding * concrete = bindings.find(binding.value); + if (concrete == nullptr) { + status.log("%s %s is not bound", command_context.c_str(), binding_context.c_str()); + return status; + } + if (concrete->buffer == nullptr) { + status.log("%s %s has a null buffer", command_context.c_str(), binding_context.c_str()); + return status; + } + if (binding.offset > concrete->length || binding.length > concrete->length - binding.offset) { + status.log("%s %s is outside runtime binding length %zu", command_context.c_str(), + binding_context.c_str(), concrete->length); + return status; + } + ref = { concrete->buffer, concrete->offset + binding.offset, binding.length }; + return status; + } + case CommandBindingOrigin::Transient: + { + const TransientAllocation * allocation = find_transient_allocation(program.transients, binding.value); + if (allocation == nullptr) { + status.log("%s %s has no transient allocation", command_context.c_str(), binding_context.c_str()); + return status; + } + if (transient_arena == nullptr || transient_arena->buffer == nullptr) { + status.log("%s %s has no transient arena", command_context.c_str(), binding_context.c_str()); + return status; + } + if (transient_arena->allocation_id == kInvalidTransientArenaAllocationId) { + status.log("%s %s has no transient arena allocation id", command_context.c_str(), + binding_context.c_str()); + return status; + } + if (program.transients.arena_size > transient_arena->capacity) { + status.log("%s %s requires transient arena size %zu but only %zu bytes are available", + command_context.c_str(), binding_context.c_str(), program.transients.arena_size, + transient_arena->capacity); + return status; + } + if (binding.offset > allocation->size || binding.length > allocation->size - binding.offset) { + status.log("%s %s is outside transient allocation length %zu", command_context.c_str(), + binding_context.c_str(), allocation->size); + return status; + } + ref = { transient_arena->buffer, allocation->arena_offset + binding.offset, binding.length }; + return status; + } + case CommandBindingOrigin::ProgramConstant: + break; + } + status.log("%s %s has an unsupported binding origin", command_context.c_str(), binding_context.c_str()); + return status; +} + +} // namespace + +static void resolve_command_list(const CommandProgram & program, + const std::vector & commands, + const CommandProgramBindings & bindings, + const TransientArenaAllocationRef * transient_arena, + std::vector & resolved_commands, + Status & status) { + resolved_commands.reserve(commands.size()); + for (const Command & command : commands) { + ResolvedCommand resolved_command; + resolved_command.ordinal = command.ordinal; + resolved_command.kind = command.kind; + resolved_command.kernel = command.kernel; + resolved_command.bindings.reserve(command.bindings.size()); + + for (const CommandBinding & binding : command.bindings) { + ResolvedCommandBinding resolved_binding; + resolved_binding.binding = binding; + Status binding_status = + resolve_command_binding(command, program, binding, bindings, transient_arena, resolved_binding.ref); + if (binding_status.success()) { + resolved_command.bindings.push_back(resolved_binding); + } else { + status.append(binding_status); + } + } + resolved_commands.push_back(resolved_command); + } +} + +ResolvedCommandProgram resolve_command_program_bindings(const CommandProgram & program, + const CommandProgramBindings & bindings, + const TransientArenaAllocationRef * transient_arena) { + ResolvedCommandProgram result; + if (!program.valid()) { + result.status.append(program.status); + } + if (!bindings.valid()) { + result.status.append(bindings.status); + } + + resolve_command_list(program, program.initialization_commands, bindings, transient_arena, + result.initialization_commands, result.status); + resolve_command_list(program, program.commands, bindings, transient_arena, result.commands, result.status); + return result; +} + +} // namespace ggml::hrx diff --git a/ggml/src/ggml-hrx/dispatch/command-program-resolver.h b/ggml/src/ggml-hrx/dispatch/command-program-resolver.h new file mode 100644 index 000000000000..b4a517fbb0eb --- /dev/null +++ b/ggml/src/ggml-hrx/dispatch/command-program-resolver.h @@ -0,0 +1,51 @@ +#pragma once + +#include "command-program-bindings.h" +#include "command-program.h" +#include "status.h" + +#include +#include +#include + +namespace ggml::hrx { + +struct ResolvedBufferRef { + hrx_buffer_t buffer = nullptr; + size_t offset = 0; + size_t length = 0; +}; + +static constexpr uint64_t kInvalidTransientArenaAllocationId = 0; + +struct TransientArenaAllocationRef { + hrx_buffer_t buffer = nullptr; + size_t capacity = 0; + uint64_t allocation_id = kInvalidTransientArenaAllocationId; +}; + +struct ResolvedCommandBinding { + CommandBinding binding; + ResolvedBufferRef ref; +}; + +struct ResolvedCommand { + uint32_t ordinal = 0; + CommandKind kind = CommandKind::Kernel; + KernelSpecialization kernel; + std::vector bindings; +}; + +struct ResolvedCommandProgram { + std::vector initialization_commands; + std::vector commands; + Status status; + + bool valid() const { return status.success(); } +}; + +ResolvedCommandProgram resolve_command_program_bindings(const CommandProgram & program, + const CommandProgramBindings & bindings, + const TransientArenaAllocationRef * transient_arena = nullptr); + +} // namespace ggml::hrx diff --git a/ggml/src/ggml-hrx/dispatch/command-program.cpp b/ggml/src/ggml-hrx/dispatch/command-program.cpp new file mode 100644 index 000000000000..642d8589f7f9 --- /dev/null +++ b/ggml/src/ggml-hrx/dispatch/command-program.cpp @@ -0,0 +1,409 @@ +#include "command-program.h" + +#include "command-program-diagnostics.h" +#include "transient-allocator.h" + +#include +#include +#include +#include +#include +#include + +namespace ggml::hrx { +namespace { + +static bool string_equal(const char * lhs, const char * rhs) { + return std::strcmp(lhs != nullptr ? lhs : "", rhs != nullptr ? rhs : "") == 0; +} + +static std::string string_value(const char * value) { + return value != nullptr ? value : ""; +} + +static CommandBindingOrigin command_binding_origin(const Graph & graph, const Value & value) { + if (value.kind == ValueKind::External) { + return CommandBindingOrigin::GraphValue; + } + const Value * root = graph.values().find(value.storage_root); + switch (root != nullptr ? root->kind : value.kind) { + case ValueKind::External: + return CommandBindingOrigin::GraphValue; + case ValueKind::Transient: + return CommandBindingOrigin::Transient; + } + return CommandBindingOrigin::GraphValue; +} + +struct StorageBindingTarget { + ValueId value; + size_t offset = 0; +}; + +static StorageBindingTarget storage_binding_target(const Graph & graph, ValueId value) { + StorageBindingTarget target; + target.value = value; + const Value * graph_value = graph.values().find(value); + if (graph_value == nullptr || graph_value->kind != ValueKind::Transient) { + return target; + } + const Value * root = graph.values().find(graph_value->storage_root); + if (root == nullptr) { + return target; + } + target.value = root->id; + target.offset = graph_value->storage_offset; + return target; +} + +static const CommandPlanTransient * find_plan_transient(const CommandPlan & plan, ValueId value) { + const auto found = std::find_if(plan.transients.begin(), plan.transients.end(), + [&](const CommandPlanTransient & transient) { return transient.value == value; }); + return found == plan.transients.end() ? nullptr : &*found; +} + +static const CommandPlanCompletionCounterRequest * find_plan_completion_counter_request(const CommandPlan & plan, + ValueId value) { + const auto found = + std::find_if(plan.completion_counter_requests.begin(), plan.completion_counter_requests.end(), + [&](const CommandPlanCompletionCounterRequest & request) { return request.value == value; }); + return found == plan.completion_counter_requests.end() ? nullptr : &*found; +} + +static void append_command(const Graph & graph, + const CommandPlan & plan, + const KernelCorpus & corpus, + const std::string & target, + const Dispatch & dispatch, + bool linear_dependency, + std::vector & commands, + Status & status) { + Command command; + command.ordinal = static_cast(commands.size()); + command.kind = CommandKind::Kernel; + command.kernel = dispatch.kernel; + // TODO: replace this linear ordinal dependency with real graph/resource dependency analysis. + if (linear_dependency && command.ordinal > 0) { + command.dependencies.push_back(command.ordinal - 1); + } + const KernelResolveResult resolved = resolve_kernel_definition(corpus, target, command.kernel.kernel_id); + const KernelDefinition * definition = resolved.definition; + if (!resolved.found()) { + status.log("%s", format_kernel_resolve_error(resolved, command.kernel.kernel_id).c_str()); + definition = nullptr; + } else if (dispatch.bindings.size() != definition->bindings.size()) { + status.log("command %u kernel %s has %zu bindings but its ABI requires %zu", command.ordinal, + kernel_definition_name(*definition).c_str(), dispatch.bindings.size(), definition->bindings.size()); + } + command.bindings.reserve(dispatch.bindings.size()); + for (size_t binding_index = 0; binding_index < dispatch.bindings.size(); ++binding_index) { + const DispatchBinding & binding = dispatch.bindings[binding_index]; + CommandBinding command_binding; + command_binding.value = binding.value; + command_binding.offset = binding.offset; + command_binding.length = binding.length; + const Value * value = graph.values().find(command_binding.value); + const CommandPlanTransient * plan_transient = find_plan_transient(plan, command_binding.value); + const CommandPlanCompletionCounterRequest * completion_counter = + find_plan_completion_counter_request(plan, command_binding.value); + if (value == nullptr && plan_transient == nullptr && completion_counter == nullptr) { + status.log("command %u binding %zu references missing value %d", command.ordinal, binding_index, + command_binding.value.value); + } else if (value != nullptr) { + command_binding.origin = command_binding_origin(graph, *value); + } else { + command_binding.origin = CommandBindingOrigin::Transient; + } + const StorageBindingTarget binding_target = storage_binding_target(graph, command_binding.value); + command_binding.value = binding_target.value; + if (binding_target.offset > 0) { + if (binding_target.offset > std::numeric_limits::max() - command_binding.offset) { + status.log("command %u binding %zu storage alias offset overflows", command.ordinal, binding_index); + } else { + command_binding.offset += binding_target.offset; + } + } + if (definition != nullptr && binding_index < definition->bindings.size()) { + command_binding.name = string_value(definition->bindings[binding_index].name); + command_binding.access = definition->bindings[binding_index].access; + } + command.bindings.push_back(std::move(command_binding)); + } + commands.push_back(std::move(command)); +} + +static void verify_command_list(const std::vector & commands, + const TransientPlan & transients, + const KernelCorpus & corpus, + const std::string & target, + Status & status) { + for (size_t i = 0; i < commands.size(); ++i) { + const Command & command = commands[i]; + const std::string command_context = format_command(command); + if (command.ordinal != i) { + status.log("%s has non-contiguous ordinal at index %zu", command_context.c_str(), i); + } + if (command.kind != CommandKind::Kernel) { + status.log("%s is not a kernel command", command_context.c_str()); + } + KernelResolveResult resolved; + const KernelDefinition * definition = nullptr; + if (command.kind == CommandKind::Kernel) { + resolved = resolve_kernel_definition(corpus, target, command.kernel.kernel_id); + definition = resolved.definition; + } + if (command.kind == CommandKind::Kernel && !resolved.found()) { + status.log("%s: %s", command_context.c_str(), + format_kernel_resolve_error(resolved, command.kernel.kernel_id).c_str()); + } else if (definition != nullptr) { + if (command.bindings.size() != definition->bindings.size()) { + status.log("%s kernel %s has %zu bindings but its ABI requires %zu", command_context.c_str(), + kernel_definition_name(*definition).c_str(), command.bindings.size(), + definition->bindings.size()); + } + const size_t shared_count = std::min(command.bindings.size(), definition->bindings.size()); + for (size_t binding_index = 0; binding_index < shared_count; ++binding_index) { + const CommandBinding & binding = command.bindings[binding_index]; + const KernelBindingDefinition & abi = definition->bindings[binding_index]; + if (!string_equal(binding.name.c_str(), abi.name) || binding.access != abi.access) { + status.log("%s %s does not match ABI binding %zu", command_context.c_str(), + format_command_binding(binding).c_str(), binding_index); + } + } + } + if (command.bindings.empty()) { + status.log("%s has no bindings", command_context.c_str()); + } + for (uint32_t dependency : command.dependencies) { + if (dependency >= command.ordinal) { + status.log("%s has forward dependency %u", command_context.c_str(), dependency); + } + } + for (const CommandBinding & binding : command.bindings) { + const std::string binding_context = format_command_binding(binding); + if (binding.origin != CommandBindingOrigin::GraphValue && + binding.origin != CommandBindingOrigin::Transient) { + status.log("%s %s has an unsupported binding origin", command_context.c_str(), binding_context.c_str()); + } + if (binding.origin == CommandBindingOrigin::Transient) { + const TransientAllocation * allocation = find_transient_allocation(transients, binding.value); + if (allocation == nullptr) { + status.log("%s %s has no transient allocation", command_context.c_str(), binding_context.c_str()); + } else if (binding.offset > allocation->size || binding.length > allocation->size - binding.offset) { + status.log("%s %s is outside transient allocation length %zu", command_context.c_str(), + binding_context.c_str(), allocation->size); + } + } + if (binding.value.value < 0) { + status.log("%s %s has an invalid value id", command_context.c_str(), binding_context.c_str()); + } + if (binding.length == 0) { + status.log("%s %s has an empty binding", command_context.c_str(), binding_context.c_str()); + } + } + } +} + +struct VerifyTransientLifetime { + bool reserved = false; + bool has_lifetime = false; + uint32_t first_command = 0; + uint32_t last_command = 0; +}; + +static size_t saturated_range_end(size_t offset, size_t size) { + if (offset > std::numeric_limits::max() - size) { + return std::numeric_limits::max(); + } + return offset + size; +} + +static bool allocation_overlaps_region(const TransientAllocation & allocation, size_t offset, size_t size) { + if (size == 0) { + return false; + } + return allocation.arena_offset < saturated_range_end(offset, size) && + offset < saturated_range_end(allocation.arena_offset, allocation.size); +} + +static std::unordered_map collect_verify_transient_lifetimes( + const CommandProgram & program) { + std::unordered_map lifetimes; + lifetimes.reserve(program.transients.allocations.size()); + for (const TransientAllocation & allocation : program.transients.allocations) { + VerifyTransientLifetime & lifetime = lifetimes[allocation.value.value]; + if (allocation_overlaps_region(allocation, program.completion_counters.arena_offset, + program.completion_counters.byte_count)) { + lifetime.reserved = true; + } + } + for (const Command & command : program.initialization_commands) { + for (const CommandBinding & binding : command.bindings) { + if (binding.origin == CommandBindingOrigin::Transient) { + lifetimes[binding.value.value].reserved = true; + } + } + } + for (const ConstantInitialization & initialization : program.constant_initializations) { + lifetimes[initialization.value.value].reserved = true; + } + for (const Command & command : program.commands) { + for (const CommandBinding & binding : command.bindings) { + if (binding.origin != CommandBindingOrigin::Transient) { + continue; + } + VerifyTransientLifetime & lifetime = lifetimes[binding.value.value]; + if (lifetime.has_lifetime) { + lifetime.first_command = std::min(lifetime.first_command, command.ordinal); + lifetime.last_command = std::max(lifetime.last_command, command.ordinal); + } else { + lifetime.has_lifetime = true; + lifetime.first_command = command.ordinal; + lifetime.last_command = command.ordinal; + } + } + } + return lifetimes; +} + +static bool verify_transient_allocations_can_overlap( + const std::unordered_map & lifetimes, + const TransientAllocation & lhs, + const TransientAllocation & rhs) { + const auto lhs_lifetime = lifetimes.find(lhs.value.value); + const auto rhs_lifetime = lifetimes.find(rhs.value.value); + if (lhs_lifetime == lifetimes.end() || rhs_lifetime == lifetimes.end() || lhs_lifetime->second.reserved || + rhs_lifetime->second.reserved || !lhs_lifetime->second.has_lifetime || !rhs_lifetime->second.has_lifetime) { + return false; + } + return lhs_lifetime->second.last_command < rhs_lifetime->second.first_command || + rhs_lifetime->second.last_command < lhs_lifetime->second.first_command; +} + +static void verify_transient_allocations(const CommandProgram & program, Status & status) { + const std::unordered_map lifetimes = collect_verify_transient_lifetimes(program); + std::vector allocations_by_offset; + allocations_by_offset.reserve(program.transients.allocations.size()); + for (const TransientAllocation & allocation : program.transients.allocations) { + if (allocation.value.value < 0 || allocation.size == 0 || allocation.alignment == 0 || + allocation.arena_offset % allocation.alignment != 0 || + allocation.arena_offset > std::numeric_limits::max() - allocation.size || + allocation.arena_offset + allocation.size > program.transients.arena_size) { + status.log("invalid transient allocation for value %d", allocation.value.value); + } + allocations_by_offset.push_back(&allocation); + } + std::sort(allocations_by_offset.begin(), allocations_by_offset.end(), + [](const TransientAllocation * lhs, const TransientAllocation * rhs) { + if (lhs->arena_offset != rhs->arena_offset) { + return lhs->arena_offset < rhs->arena_offset; + } + return lhs->value.value < rhs->value.value; + }); + for (size_t i = 0; i < allocations_by_offset.size(); ++i) { + const TransientAllocation & allocation = *allocations_by_offset[i]; + const size_t end = saturated_range_end(allocation.arena_offset, allocation.size); + for (size_t j = i + 1; j < allocations_by_offset.size(); ++j) { + const TransientAllocation & other = *allocations_by_offset[j]; + if (other.arena_offset >= end) { + break; + } + if (!verify_transient_allocations_can_overlap(lifetimes, allocation, other)) { + status.log("transient allocations overlap"); + } + } + } +} + +} // namespace + +const TransientAllocation * find_transient_allocation(const TransientPlan & plan, ValueId value) { + const auto found = std::find_if(plan.allocations.begin(), plan.allocations.end(), + [&](const TransientAllocation & allocation) { return allocation.value == value; }); + return found == plan.allocations.end() ? nullptr : &*found; +} + +CommandProgram build_command_program(const Graph & graph, + const CommandPlan & plan, + const KernelCorpus & corpus, + const std::string & target) { + CommandProgram result; + if (!plan.valid()) { + result.status.append(plan.status); + return result; + } + + result.initialization_commands.reserve(plan.initialization_dispatches.size()); + for (const Dispatch & dispatch : plan.initialization_dispatches) { + append_command(graph, plan, corpus, target, dispatch, false, result.initialization_commands, result.status); + } + for (const Dispatch & dispatch : plan.dispatches) { + append_command(graph, plan, corpus, target, dispatch, true, result.commands, result.status); + } + result.transients = TransientAllocator::allocate(graph, plan, result.initialization_commands, result.commands, + result.completion_counters, result.status); + result.constant_initializations.reserve(plan.constant_initializations.size()); + for (const CommandPlanConstantInitialization & initialization : plan.constant_initializations) { + result.constant_initializations.push_back({ + initialization.value, + initialization.name, + initialization.offset, + initialization.data, + }); + } + return result; +} + +VerificationResult verify_command_program(const CommandProgram & program, + const KernelCorpus & corpus, + const std::string & target) { + VerificationResult result; + if (!program.valid()) { + result.status.append(program.status); + } + verify_command_list(program.initialization_commands, program.transients, corpus, target, result.status); + verify_command_list(program.commands, program.transients, corpus, target, result.status); + if (program.transients.arena_alignment == 0) { + result.status.log("transient arena has zero alignment"); + } + if (program.completion_counters.count == 0) { + if (program.completion_counters.byte_count != 0) { + result.status.log("completion counter region has bytes but no counters"); + } + } else { + if (program.completion_counters.byte_count == 0) { + result.status.log("completion counter region has counters but no bytes"); + } + if (program.completion_counters.byte_count % sizeof(int32_t) != 0) { + result.status.log("completion counter region byte count is not i32 aligned"); + } + if (program.completion_counters.arena_offset % 16 != 0) { + result.status.log("completion counter region is not 16-byte aligned"); + } + if (program.completion_counters.arena_offset > program.transients.arena_size || + program.completion_counters.byte_count > + program.transients.arena_size - program.completion_counters.arena_offset) { + result.status.log("completion counter region is outside transient arena"); + } + } + verify_transient_allocations(program, result.status); + for (const ConstantInitialization & initialization : program.constant_initializations) { + const TransientAllocation * allocation = find_transient_allocation(program.transients, initialization.value); + if (allocation == nullptr) { + result.status.log("constant initialization %s references missing transient value %d", + initialization.name.c_str(), initialization.value.value); + continue; + } + if (initialization.data.empty()) { + result.status.log("constant initialization %s has no data", initialization.name.c_str()); + } + if (initialization.offset > allocation->size || + initialization.data.size() > allocation->size - initialization.offset) { + result.status.log("constant initialization %s is outside transient allocation length %zu", + initialization.name.c_str(), allocation->size); + } + } + return result; +} + +} // namespace ggml::hrx diff --git a/ggml/src/ggml-hrx/dispatch/command-program.h b/ggml/src/ggml-hrx/dispatch/command-program.h new file mode 100644 index 000000000000..980fe8c3e534 --- /dev/null +++ b/ggml/src/ggml-hrx/dispatch/command-program.h @@ -0,0 +1,91 @@ +#pragma once + +#include "command-plan.h" +#include "graph/graph.h" +#include "kernel-corpus/kernel-corpus.h" +#include "kernel-corpus/kernel-types.h" +#include "status.h" + +#include +#include +#include +#include + +namespace ggml::hrx { + +enum class CommandKind : uint8_t { + Invalid, + Kernel, +}; + +enum class CommandBindingOrigin : uint8_t { + GraphValue, + Transient, + ProgramConstant, +}; + +struct CommandBinding { + std::string name; + ValueId value; + CommandBindingOrigin origin = CommandBindingOrigin::GraphValue; + size_t offset = 0; + size_t length = 0; + ResourceAccess access = ResourceAccess::Read; +}; + +struct Command { + uint32_t ordinal = 0; + CommandKind kind = CommandKind::Kernel; + KernelSpecialization kernel; + std::vector bindings; + std::vector dependencies; +}; + +struct TransientAllocation { + ValueId value; + size_t size = 0; + size_t alignment = 1; + size_t arena_offset = 0; +}; + +struct TransientPlan { + size_t arena_size = 0; + size_t arena_alignment = 1; + std::vector allocations; +}; + +struct ConstantInitialization { + ValueId value; + std::string name; + size_t offset = 0; + std::vector data; +}; + +struct CompletionCounterPlan { + size_t arena_offset = 0; + size_t byte_count = 0; + uint32_t count = 0; +}; + +struct CommandProgram { + std::vector initialization_commands; + std::vector commands; + TransientPlan transients; + CompletionCounterPlan completion_counters; + std::vector constant_initializations; + Status status; + + bool valid() const { return status.success(); } +}; + +const TransientAllocation * find_transient_allocation(const TransientPlan & plan, ValueId value); + +CommandProgram build_command_program(const Graph & graph, + const CommandPlan & plan, + const KernelCorpus & corpus, + const std::string & target); +VerificationResult verify_command_program(const CommandProgram & program, + const KernelCorpus & corpus, + const std::string & target); + +} // namespace ggml::hrx diff --git a/ggml/src/ggml-hrx/dispatch/dispatch-scheduler.cpp b/ggml/src/ggml-hrx/dispatch/dispatch-scheduler.cpp new file mode 100644 index 000000000000..e1e512a0b3bd --- /dev/null +++ b/ggml/src/ggml-hrx/dispatch/dispatch-scheduler.cpp @@ -0,0 +1,299 @@ +// Copyright 2026 The HRX Authors +// SPDX-License-Identifier: Apache-2.0 + +#include "dispatch-scheduler.h" + +#include "ggml.h" +#include "graph/graph-traversal.h" + +#include +#include +#include +#include +#include + +namespace ggml::hrx { +namespace { + +static bool match_covers_root(const DispatchMatch & match, size_t root_index) { + return std::find(match.covered_nodes.begin(), match.covered_nodes.end(), root_index) != match.covered_nodes.end(); +} + +static bool match_overlaps_covered_nodes(const DispatchMatch & match, const std::vector & covered_nodes) { + for (const size_t node_index : match.covered_nodes) { + if (node_index >= covered_nodes.size() || covered_nodes[node_index]) { + return true; + } + } + return false; +} + +static bool try_match_registration(const Graph & graph, + const GraphNode * node, + size_t node_index, + const std::vector & covered_nodes, + const CommandPlan & plan, + const DispatchRegistry & registry, + ValueId next_plan_value, + DispatchMatch & match, + DispatchMatchDiagnostics * diagnostics) { + const DispatchMatchContext context = { + graph, node, node_index, covered_nodes, plan, next_plan_value, + }; + return registry.match(context, match, diagnostics); +} + +static void clear_plan_results(CommandPlan & plan) { + plan.initialization_dispatches.clear(); + plan.dispatches.clear(); + plan.transients.clear(); + plan.constant_initializations.clear(); + plan.completion_counter_requests.clear(); + plan.metadata.clear(); +} + +static void append_value_summary(std::ostringstream & stream, const Graph & graph, ValueId value_id) { + const Value * value = graph.values().find(value_id); + if (value == nullptr) { + stream << value_id.value << ":missing"; + return; + } + stream << value_id.value << ":" << ggml_type_name(value->type) << "[" << value->ne[0] << "," << value->ne[1] << "," + << value->ne[2] << "," << value->ne[3] << "]"; + const GraphNode * producer = graph.index().producer(value_id); + if (producer != nullptr) { + stream << "<-" << ggml_op_name(producer->op); + } +} + +static void append_node_summary(std::ostringstream & stream, const Graph & graph, const GraphNode * node) { + if (node == nullptr) { + stream << "null"; + return; + } + size_t node_index = 0; + if (graph.index().node_index(node, node_index)) { + stream << node_index << ":"; + } + stream << ggml_op_name(node->op); +} + +static std::string unsupported_node_message(const Graph & graph, size_t index, const GraphNode & node) { + std::ostringstream stream; + stream << "unsupported HRX node " << index << ": " << ggml_op_name(node.op) << " output="; + append_value_summary(stream, graph, node.output); + stream << " inputs=["; + for (size_t i = 0; i < node.inputs.size(); ++i) { + if (i > 0) { + stream << ", "; + } + append_value_summary(stream, graph, node.inputs[i]); + } + stream << "]"; + stream << " consumers=["; + const std::vector & consumers = graph.index().consumers(node.output); + for (size_t i = 0; i < consumers.size(); ++i) { + if (i > 0) { + stream << ", "; + } + append_node_summary(stream, graph, consumers[i]); + } + stream << "]"; + return stream.str(); +} + +static bool value_is_available(const Graph & graph, ValueId value, const std::vector & covered_nodes) { + const GraphNode * producer = graph.index().producer(value); + if (producer == nullptr) { + return true; + } + size_t producer_index = 0; + return graph.index().node_index(producer, producer_index) && producer_index < covered_nodes.size() && + covered_nodes[producer_index]; +} + +static bool can_elide_layout_alias_node(const Graph & graph, + const GraphNode & node, + const std::vector & covered_nodes) { + return is_layout_alias_node(graph, node) && value_is_available(graph, node.inputs[0], covered_nodes); +} + +static bool apply_value_aliases(Graph & graph, const DispatchMatch & match, Status & status) { + for (const DispatchValueAliasRequest & alias : match.value_aliases) { + Status alias_status = graph.values().alias_storage(alias.target_value, alias.source_value); + if (!alias_status.success()) { + status.append(alias_status); + return false; + } + } + return true; +} + +} // namespace + +bool DispatchScheduler::schedule_graph(Graph & graph, const DispatchTarget & target) { + return this->schedule_graph(graph, target, nullptr); +} + +bool DispatchScheduler::schedule_graph(Graph & graph, + const DispatchTarget & target, + DispatchScheduleDiagnostics * diagnostics) { + plan_ = {}; + if (diagnostics != nullptr) { + *diagnostics = {}; + } + const DispatchRegistry * registry = find_dispatch_registry(target); + if (registry == nullptr) { + plan_.status.log("no HRX dispatch registry for target %s", target.architecture.c_str()); + return false; + } + const std::vector & nodes = graph.nodes(); + if (!graph.has_index()) { + plan_.status.log("HRX graph is missing graph index"); + return false; + } + std::vector covered_nodes(nodes.size(), false); + Status pending_diagnostics; + const GraphTraversalOrder traversal = GraphTraversalOrder::build(graph); + for (const GraphNode * node : traversal.nodes()) { + size_t i = 0; + if (node == nullptr || !graph.index().node_index(node, i)) { + plan_.status.log("HRX traversal references a node outside the graph"); + clear_plan_results(plan_); + return false; + } + if (covered_nodes[i]) { + continue; + } + if (node->op == GGML_OP_NONE) { + // Leaf/parameter nodes have no producer command: their storage is bound + // from outside the graph and is already available to consumers. Treat them + // as covered instead of failing the whole graph with an "unsupported HRX + // node" diagnostic (ggml_build_forward_expand emits leaves as nodes). + covered_nodes[i] = true; + continue; + } + DispatchMatch match; + const ValueId next_plan_value(static_cast(graph.values().size() + plan_.transients.size() + + plan_.completion_counter_requests.size())); + DispatchMatchDiagnostics match_diagnostics; + if (!try_match_registration(graph, node, i, covered_nodes, plan_, *registry, next_plan_value, match, + &match_diagnostics)) { + if (can_elide_layout_alias_node(graph, *node, covered_nodes)) { + pending_diagnostics.append(match.status); + covered_nodes[i] = true; + continue; + } + plan_.status.append(pending_diagnostics); + plan_.status.append(match.status); + const std::string message = unsupported_node_message(graph, i, *node); + plan_.status.log("%s", message.c_str()); + if (diagnostics != nullptr) { + diagnostics->unsupported_node_index = i; + diagnostics->unsupported_node = node; + diagnostics->unsupported_message = message; + diagnostics->match = std::move(match_diagnostics); + } + clear_plan_results(plan_); + return false; + } + if (match.covered_nodes.empty() || match.dispatches.empty() || !match_covers_root(match, i) || + match_overlaps_covered_nodes(match, covered_nodes)) { + plan_.status.log("invalid HRX dispatch match for node %zu: %s", i, ggml_op_name(node->op)); + if (diagnostics != nullptr) { + diagnostics->unsupported_node_index = i; + diagnostics->unsupported_node = node; + diagnostics->unsupported_message = "invalid HRX dispatch match"; + diagnostics->match = std::move(match_diagnostics); + } + clear_plan_results(plan_); + return false; + } + if (!apply_value_aliases(graph, match, plan_.status)) { + if (diagnostics != nullptr) { + diagnostics->unsupported_node_index = i; + diagnostics->unsupported_node = node; + diagnostics->unsupported_message = "invalid HRX value alias"; + diagnostics->match = std::move(match_diagnostics); + } + clear_plan_results(plan_); + return false; + } + for (Dispatch & dispatch : match.initialization_dispatches) { + plan_.initialization_dispatches.push_back(std::move(dispatch)); + } + for (Dispatch & dispatch : match.dispatches) { + plan_.dispatches.push_back(std::move(dispatch)); + } + for (CommandPlanTransient & transient : match.transients) { + plan_.transients.push_back(std::move(transient)); + } + for (CommandPlanConstantInitialization & initialization : match.constant_initializations) { + plan_.constant_initializations.push_back(std::move(initialization)); + } + for (CommandPlanCompletionCounterRequest & request : match.completion_counter_requests) { + plan_.completion_counter_requests.push_back(std::move(request)); + } + if (!plan_.metadata.append(std::move(match.metadata), plan_.status)) { + clear_plan_results(plan_); + return false; + } + for (const size_t covered_node : match.covered_nodes) { + covered_nodes[covered_node] = true; + } + } + for (size_t i = 0; i < nodes.size(); ++i) { + if (!covered_nodes[i]) { + if (can_elide_layout_alias_node(graph, nodes[i], covered_nodes)) { + covered_nodes[i] = true; + continue; + } + plan_.status.append(pending_diagnostics); + const std::string message = unsupported_node_message(graph, i, nodes[i]); + plan_.status.log("%s", message.c_str()); + if (diagnostics != nullptr) { + DispatchMatch match; + const ValueId next_plan_value(static_cast(graph.values().size() + plan_.transients.size() + + plan_.completion_counter_requests.size())); + DispatchMatchDiagnostics match_diagnostics; + try_match_registration(graph, &nodes[i], i, covered_nodes, plan_, *registry, next_plan_value, match, + &match_diagnostics); + diagnostics->unsupported_node_index = i; + diagnostics->unsupported_node = &nodes[i]; + diagnostics->unsupported_message = message; + diagnostics->match = std::move(match_diagnostics); + } + clear_plan_results(plan_); + return false; + } + } + return true; +} + +bool DispatchScheduler::supports_node(const Graph & graph, const GraphNode * node, const DispatchTarget & target) { + const DispatchRegistry * registry = find_dispatch_registry(target); + if (registry == nullptr) { + return false; + } + if (node == nullptr || !graph.has_index()) { + return false; + } + size_t node_index = 0; + if (!graph.index().node_index(node, node_index)) { + return false; + } + const std::vector covered_nodes(graph.nodes().size(), false); + DispatchMatch match; + const ValueId next_plan_value(static_cast(graph.values().size())); + const CommandPlan plan; + return try_match_registration(graph, node, node_index, covered_nodes, plan, *registry, next_plan_value, match, + nullptr); +} + +bool DispatchScheduler::can_schedule_graph(const Graph & graph, const DispatchTarget & target) { + Graph graph_copy = graph; + DispatchScheduler scheduler; + return scheduler.schedule_graph(graph_copy, target); +} + +} // namespace ggml::hrx diff --git a/ggml/src/ggml-hrx/dispatch/dispatch-scheduler.h b/ggml/src/ggml-hrx/dispatch/dispatch-scheduler.h new file mode 100644 index 000000000000..b533c0ed78d2 --- /dev/null +++ b/ggml/src/ggml-hrx/dispatch/dispatch-scheduler.h @@ -0,0 +1,40 @@ +#pragma once + +#include "command-plan.h" +#include "dispatch_registration/dispatch-registry.h" +#include "graph/graph.h" + +#include +#include + +namespace ggml::hrx { + +struct DispatchScheduleDiagnostics { + size_t unsupported_node_index = 0; + const GraphNode * unsupported_node = nullptr; + std::string unsupported_message; + DispatchMatchDiagnostics match; +}; + +class DispatchScheduler { + public: + bool schedule_graph(Graph & graph, const DispatchTarget & target); + bool schedule_graph(Graph & graph, const DispatchTarget & target, DispatchScheduleDiagnostics * diagnostics); + + const CommandPlan & plan() const { return plan_; } + + const std::vector & dispatches() const { return plan_.dispatches; } + + const std::string & error() const { + static const std::string empty; + return plan_.status.errors().empty() ? empty : plan_.status.errors().front(); + } + + static bool supports_node(const Graph & graph, const GraphNode * node, const DispatchTarget & target); + static bool can_schedule_graph(const Graph & graph, const DispatchTarget & target); + + private: + CommandPlan plan_; +}; + +} // namespace ggml::hrx diff --git a/ggml/src/ggml-hrx/dispatch/dispatch.h b/ggml/src/ggml-hrx/dispatch/dispatch.h new file mode 100644 index 000000000000..a5e388628800 --- /dev/null +++ b/ggml/src/ggml-hrx/dispatch/dispatch.h @@ -0,0 +1,38 @@ +#pragma once + +#include "graph/value-map.h" +#include "kernel-corpus/kernel-corpus-catalog.h" +#include "kernel-corpus/kernel-types.h" + +#include +#include +#include +#include +#include + +namespace ggml::hrx { + +struct KernelSpecialization { + uint64_t kernel_id = kUncatalogedKernelId; + std::map integer_parameters; + std::map compile_parameters; +}; + +inline KernelSpecialization make_kernel_specialization(KernelCatalogRef ref) { + KernelSpecialization kernel; + kernel.kernel_id = ref.id; + return kernel; +} + +struct DispatchBinding { + ValueId value; + size_t offset = 0; + size_t length = 0; +}; + +struct Dispatch { + KernelSpecialization kernel; + std::vector bindings; +}; + +} // namespace ggml::hrx diff --git a/ggml/src/ggml-hrx/dispatch/transient-allocator.cpp b/ggml/src/ggml-hrx/dispatch/transient-allocator.cpp new file mode 100644 index 000000000000..421fd2c9bf60 --- /dev/null +++ b/ggml/src/ggml-hrx/dispatch/transient-allocator.cpp @@ -0,0 +1,461 @@ +#include "transient-allocator.h" + +#include +#include +#include +#include +#include + +namespace ggml::hrx { +namespace { + +static size_t align_up(size_t value, size_t alignment) { + return alignment == 0 ? value : (value + alignment - 1) / alignment * alignment; +} + +struct TransientAllocationRequest { + ValueId value; + size_t required_size = 0; +}; + +struct TransientAllocationInterval { + TransientAllocation allocation; + uint32_t first_use = 0; + uint32_t last_use = 0; +}; + +struct FreeTransientBlock { + size_t offset = 0; + size_t size = 0; +}; + +struct ActiveTransientAllocation { + size_t offset = 0; + size_t size = 0; + uint32_t last_use = 0; +}; + +static const CommandPlanTransient * find_plan_transient(const CommandPlan & plan, ValueId value) { + const auto found = std::find_if(plan.transients.begin(), plan.transients.end(), + [&](const CommandPlanTransient & transient) { return transient.value == value; }); + return found == plan.transients.end() ? nullptr : &*found; +} + +static const CommandPlanCompletionCounterRequest * find_plan_completion_counter_request(const CommandPlan & plan, + ValueId value) { + const auto found = + std::find_if(plan.completion_counter_requests.begin(), plan.completion_counter_requests.end(), + [&](const CommandPlanCompletionCounterRequest & request) { return request.value == value; }); + return found == plan.completion_counter_requests.end() ? nullptr : &*found; +} + +static void add_transient_allocation_request(std::vector & requests, + ValueId value, + size_t required_size) { + for (TransientAllocationRequest & request : requests) { + if (request.value == value) { + request.required_size = std::max(request.required_size, required_size); + return; + } + } + requests.push_back({ value, required_size }); +} + +static const TransientAllocationRequest * find_transient_allocation_request( + const std::vector & requests, + ValueId value) { + const auto found = std::find_if(requests.begin(), requests.end(), + [&](const TransientAllocationRequest & request) { return request.value == value; }); + return found == requests.end() ? nullptr : &*found; +} + +static bool contains_value(const std::vector & values, ValueId value) { + return std::find(values.begin(), values.end(), value) != values.end(); +} + +static bool make_transient_allocation(const Graph & graph, + const CommandPlan & command_plan, + const TransientAllocationRequest & request, + TransientAllocation & allocation, + Status & errors) { + const Value * graph_value = graph.values().find(request.value); + const CommandPlanTransient * plan_transient = find_plan_transient(command_plan, request.value); + if (graph_value == nullptr && plan_transient == nullptr) { + errors.log("transient value %d is missing from graph values and command plan transients", request.value.value); + return false; + } + if (graph_value != nullptr && graph_value->kind != ValueKind::Transient) { + errors.log("transient value %d aliases a non-transient graph value", request.value.value); + return false; + } + + allocation.value = request.value; + allocation.size = + std::max(graph_value != nullptr ? graph_value->byte_count : plan_transient->size, request.required_size); + allocation.alignment = plan_transient != nullptr ? plan_transient->alignment : 256; + allocation.arena_offset = 0; + return true; +} + +static void add_transient_allocation(const Graph & graph, + const CommandPlan & command_plan, + const TransientAllocationRequest & request, + TransientPlan & plan, + Status & errors) { + TransientAllocation allocation; + if (!make_transient_allocation(graph, command_plan, request, allocation, errors)) { + return; + } + allocation.arena_offset = align_up(plan.arena_size, allocation.alignment); + plan.arena_size = allocation.arena_offset + allocation.size; + plan.allocations.push_back(allocation); +} + +static void add_transient_interval(std::vector & intervals, + TransientAllocation allocation, + uint32_t command_ordinal) { + for (TransientAllocationInterval & interval : intervals) { + if (interval.allocation.value == allocation.value) { + interval.allocation.size = std::max(interval.allocation.size, allocation.size); + interval.allocation.alignment = std::max(interval.allocation.alignment, allocation.alignment); + interval.first_use = std::min(interval.first_use, command_ordinal); + interval.last_use = std::max(interval.last_use, command_ordinal); + return; + } + } + intervals.push_back({ allocation, command_ordinal, command_ordinal }); +} + +static void add_free_transient_block(std::vector & free_blocks, size_t offset, size_t size) { + if (size > 0) { + free_blocks.push_back({ offset, size }); + } +} + +static void coalesce_free_transient_blocks(std::vector & free_blocks) { + std::sort(free_blocks.begin(), free_blocks.end(), + [](const FreeTransientBlock & lhs, const FreeTransientBlock & rhs) { return lhs.offset < rhs.offset; }); + size_t write_index = 0; + for (const FreeTransientBlock & block : free_blocks) { + if (block.size == 0) { + continue; + } + if (write_index > 0) { + FreeTransientBlock & previous = free_blocks[write_index - 1]; + const size_t previous_end = previous.offset + previous.size; + if (previous_end == block.offset) { + previous.size += block.size; + continue; + } + } + free_blocks[write_index++] = block; + } + free_blocks.resize(write_index); +} + +static void release_completed_transient_allocations(std::vector & active, + std::vector & free_blocks, + uint32_t first_use) { + size_t write_index = 0; + for (size_t read_index = 0; read_index < active.size(); ++read_index) { + if (active[read_index].last_use < first_use) { + add_free_transient_block(free_blocks, active[read_index].offset, active[read_index].size); + continue; + } + if (write_index != read_index) { + active[write_index] = active[read_index]; + } + ++write_index; + } + active.resize(write_index); +} + +static bool try_allocate_from_free_blocks(std::vector & free_blocks, + size_t size, + size_t alignment, + size_t & offset) { + coalesce_free_transient_blocks(free_blocks); + for (size_t i = 0; i < free_blocks.size(); ++i) { + const size_t block_begin = free_blocks[i].offset; + const size_t block_end = free_blocks[i].offset + free_blocks[i].size; + const size_t aligned_offset = align_up(block_begin, alignment); + const bool aligned_in_block = aligned_offset >= block_begin && aligned_offset <= block_end; + if (!aligned_in_block || size > block_end - aligned_offset) { + continue; + } + + offset = aligned_offset; + const FreeTransientBlock block = free_blocks[i]; + free_blocks.erase(free_blocks.begin() + static_cast(i)); + add_free_transient_block(free_blocks, block.offset, aligned_offset - block.offset); + add_free_transient_block(free_blocks, aligned_offset + size, block_end - (aligned_offset + size)); + return true; + } + return false; +} + +static void pack_transient_intervals(std::vector & intervals, TransientPlan & plan) { + std::sort(intervals.begin(), intervals.end(), + [](const TransientAllocationInterval & lhs, const TransientAllocationInterval & rhs) { + if (lhs.first_use != rhs.first_use) { + return lhs.first_use < rhs.first_use; + } + if (lhs.last_use != rhs.last_use) { + return lhs.last_use < rhs.last_use; + } + return lhs.allocation.value.value < rhs.allocation.value.value; + }); + + std::vector active; + std::vector free_blocks; + for (TransientAllocationInterval & interval : intervals) { + release_completed_transient_allocations(active, free_blocks, interval.first_use); + + size_t offset = 0; + if (!try_allocate_from_free_blocks(free_blocks, interval.allocation.size, interval.allocation.alignment, + offset)) { + offset = align_up(plan.arena_size, interval.allocation.alignment); + plan.arena_size = offset + interval.allocation.size; + } + + interval.allocation.arena_offset = offset; + active.push_back({ offset, interval.allocation.size, interval.last_use }); + plan.allocations.push_back(interval.allocation); + } +} + +static void add_completion_counter_allocations(const CommandPlan & command_plan, + const std::vector & binding_requests, + TransientPlan & plan, + CompletionCounterPlan & completion_counters, + Status & errors) { + for (size_t i = 0; i < command_plan.completion_counter_requests.size(); ++i) { + const CommandPlanCompletionCounterRequest & request = command_plan.completion_counter_requests[i]; + if (request.value.value < 0) { + errors.log("completion counter request %s has invalid value %d", request.name.c_str(), request.value.value); + continue; + } + if (request.count == 0) { + errors.log("completion counter request %s has zero counters", request.name.c_str()); + continue; + } + for (size_t j = i + 1; j < command_plan.completion_counter_requests.size(); ++j) { + if (request.value == command_plan.completion_counter_requests[j].value) { + errors.log("duplicate completion counter request value %d", request.value.value); + } + } + const size_t byte_count = static_cast(request.count) * sizeof(int32_t); + const TransientAllocationRequest * binding_request = + find_transient_allocation_request(binding_requests, request.value); + if (binding_request != nullptr && binding_request->required_size > byte_count) { + errors.log("completion counter request %s requires %zu bytes but binding uses %zu bytes", + request.name.c_str(), byte_count, binding_request->required_size); + continue; + } + if (completion_counters.count > std::numeric_limits::max() - request.count) { + errors.log("completion counter count overflows"); + continue; + } + + TransientAllocation allocation; + allocation.value = request.value; + allocation.size = byte_count; + allocation.alignment = 16; + allocation.arena_offset = align_up(plan.arena_size, allocation.alignment); + if (completion_counters.byte_count == 0) { + completion_counters.arena_offset = allocation.arena_offset; + } + plan.arena_size = allocation.arena_offset + allocation.size; + completion_counters.byte_count = + plan.arena_size > completion_counters.arena_offset ? plan.arena_size - completion_counters.arena_offset : 0; + completion_counters.count += request.count; + plan.allocations.push_back(allocation); + } +} + +static void record_transient_lifetime(std::vector & lifetimes, + const CommandBinding & binding, + uint32_t command_ordinal) { + for (TransientAllocationInterval & lifetime : lifetimes) { + if (lifetime.allocation.value == binding.value) { + lifetime.first_use = std::min(lifetime.first_use, command_ordinal); + lifetime.last_use = std::max(lifetime.last_use, command_ordinal); + return; + } + } + TransientAllocation allocation; + allocation.value = binding.value; + lifetimes.push_back({ allocation, command_ordinal, command_ordinal }); +} + +static const TransientAllocationInterval * find_transient_lifetime( + const std::vector & lifetimes, + ValueId value) { + const auto found = + std::find_if(lifetimes.begin(), lifetimes.end(), + [&](const TransientAllocationInterval & lifetime) { return lifetime.allocation.value == value; }); + return found == lifetimes.end() ? nullptr : &*found; +} + +static bool transient_allocation_overlaps_region(const TransientAllocation & allocation, + size_t region_offset, + size_t region_size) { + if (region_size == 0) { + return false; + } + return allocation.arena_offset < region_offset + region_size && + region_offset < allocation.arena_offset + allocation.size; +} + +static bool transient_allocation_has_reserved_lifetime(const CommandProgram & program, + const TransientAllocation & allocation) { + if (transient_allocation_overlaps_region(allocation, program.completion_counters.arena_offset, + program.completion_counters.byte_count)) { + return true; + } + for (const Command & command : program.initialization_commands) { + for (const CommandBinding & binding : command.bindings) { + if (binding.origin == CommandBindingOrigin::Transient && binding.value == allocation.value) { + return true; + } + } + } + for (const ConstantInitialization & initialization : program.constant_initializations) { + if (initialization.value == allocation.value) { + return true; + } + } + return false; +} + +static std::vector collect_main_transient_lifetimes(const CommandProgram & program) { + std::vector lifetimes; + for (const Command & command : program.commands) { + for (const CommandBinding & binding : command.bindings) { + if (binding.origin == CommandBindingOrigin::Transient) { + record_transient_lifetime(lifetimes, binding, command.ordinal); + } + } + } + return lifetimes; +} + +static bool transient_lifetimes_disjoint(const std::vector & lifetimes, + const TransientAllocation & lhs, + const TransientAllocation & rhs) { + const TransientAllocationInterval * lhs_lifetime = find_transient_lifetime(lifetimes, lhs.value); + const TransientAllocationInterval * rhs_lifetime = find_transient_lifetime(lifetimes, rhs.value); + if (lhs_lifetime == nullptr || rhs_lifetime == nullptr) { + return false; + } + return lhs_lifetime->last_use < rhs_lifetime->first_use || rhs_lifetime->last_use < lhs_lifetime->first_use; +} + +} // namespace + +TransientPlan TransientAllocator::allocate(const Graph & graph, + const CommandPlan & command_plan, + const std::vector & initialization_commands, + const std::vector & commands, + CompletionCounterPlan & completion_counters, + Status & errors) { + TransientPlan plan; + plan.arena_alignment = 256; + std::vector reserved_values; + std::vector reserved_requests; + std::vector packable_requests; + std::vector completion_counter_binding_requests; + + for (const Command & command : initialization_commands) { + for (const CommandBinding & binding : command.bindings) { + if (binding.origin == CommandBindingOrigin::Transient && + find_plan_completion_counter_request(command_plan, binding.value) == nullptr && + !contains_value(reserved_values, binding.value)) { + reserved_values.push_back(binding.value); + } + } + } + for (const CommandPlanConstantInitialization & initialization : command_plan.constant_initializations) { + const bool completion_counter = + find_plan_completion_counter_request(command_plan, initialization.value) != nullptr; + if (!completion_counter && !contains_value(reserved_values, initialization.value)) { + reserved_values.push_back(initialization.value); + } + if (initialization.offset > std::numeric_limits::max() - initialization.data.size()) { + errors.log("constant initialization %s range overflows", initialization.name.c_str()); + continue; + } + if (completion_counter) { + add_transient_allocation_request(completion_counter_binding_requests, initialization.value, + initialization.offset + initialization.data.size()); + } else { + add_transient_allocation_request(reserved_requests, initialization.value, + initialization.offset + initialization.data.size()); + } + } + + auto append_command_bindings = [&](const std::vector & command_list, bool packable) { + for (const Command & command : command_list) { + for (const CommandBinding & binding : command.bindings) { + if (binding.origin != CommandBindingOrigin::Transient) { + continue; + } + if (binding.offset > std::numeric_limits::max() - binding.length) { + errors.log("transient value %d binding range overflows", binding.value.value); + continue; + } + if (find_plan_completion_counter_request(command_plan, binding.value) != nullptr) { + add_transient_allocation_request(completion_counter_binding_requests, binding.value, + binding.offset + binding.length); + } else if (!packable || contains_value(reserved_values, binding.value)) { + add_transient_allocation_request(reserved_requests, binding.value, binding.offset + binding.length); + } else { + add_transient_allocation_request(packable_requests, binding.value, binding.offset + binding.length); + } + } + } + }; + append_command_bindings(initialization_commands, false); + append_command_bindings(commands, true); + add_completion_counter_allocations(command_plan, completion_counter_binding_requests, plan, completion_counters, + errors); + for (const TransientAllocationRequest & request : reserved_requests) { + add_transient_allocation(graph, command_plan, request, plan, errors); + } + + std::vector intervals; + for (const Command & command : commands) { + for (const CommandBinding & binding : command.bindings) { + if (binding.origin != CommandBindingOrigin::Transient || + find_plan_completion_counter_request(command_plan, binding.value) != nullptr || + contains_value(reserved_values, binding.value)) { + continue; + } + const TransientAllocationRequest * request = + find_transient_allocation_request(packable_requests, binding.value); + if (request == nullptr) { + continue; + } + TransientAllocation allocation; + if (make_transient_allocation(graph, command_plan, *request, allocation, errors)) { + add_transient_interval(intervals, allocation, command.ordinal); + } + } + } + pack_transient_intervals(intervals, plan); + plan.arena_size = align_up(plan.arena_size, plan.arena_alignment); + return plan; +} + +bool TransientAllocator::allocations_can_overlap(const CommandProgram & program, + const TransientAllocation & lhs, + const TransientAllocation & rhs) { + if (transient_allocation_has_reserved_lifetime(program, lhs) || + transient_allocation_has_reserved_lifetime(program, rhs)) { + return false; + } + const std::vector lifetimes = collect_main_transient_lifetimes(program); + return transient_lifetimes_disjoint(lifetimes, lhs, rhs); +} + +} // namespace ggml::hrx diff --git a/ggml/src/ggml-hrx/dispatch/transient-allocator.h b/ggml/src/ggml-hrx/dispatch/transient-allocator.h new file mode 100644 index 000000000000..ad059b7de60f --- /dev/null +++ b/ggml/src/ggml-hrx/dispatch/transient-allocator.h @@ -0,0 +1,23 @@ +#pragma once + +#include "command-program.h" + +#include + +namespace ggml::hrx { + +class TransientAllocator { + public: + static TransientPlan allocate(const Graph & graph, + const CommandPlan & command_plan, + const std::vector & initialization_commands, + const std::vector & commands, + CompletionCounterPlan & completion_counters, + Status & errors); + + static bool allocations_can_overlap(const CommandProgram & program, + const TransientAllocation & lhs, + const TransientAllocation & rhs); +}; + +} // namespace ggml::hrx diff --git a/ggml/src/ggml-hrx/dispatch_registration/dispatch-add.cpp b/ggml/src/ggml-hrx/dispatch_registration/dispatch-add.cpp new file mode 100644 index 000000000000..5d6713eaad49 --- /dev/null +++ b/ggml/src/ggml-hrx/dispatch_registration/dispatch-add.cpp @@ -0,0 +1,91 @@ +// Copyright 2026 The HRX Authors +// SPDX-License-Identifier: Apache-2.0 + +#include "dispatch-add.h" + +#include "ggml.h" +#include "kernel-corpus/kernel-corpus-catalog-verify.h" + +#include +#include +#include + +namespace ggml::hrx { + +static constexpr KernelCatalogRef kAddF32Kernel = GGML_HRX_KERNEL_REF("qwen3_moe", "ggml_add_f32"); + +// The HRX/loom runtime submits a dispatch through a fixed-size kernarg ring and +// fails with OUT_OF_RANGE when the grid does not fit (see +// amdgpu/host_queue_submission.c); that failed submission also leaves the device +// stream unusable for every later submission. Only claim shapes whose grid stays +// inside the ring. The add kernel launches one 256-wide workgroup per 256 +// elements, so 4096 workgroups is ~1M elements - far above any activation this +// route is used for, and well below the ~1.9M-element point where the measured +// submission starts failing. +static constexpr uint64_t kAddMaxWorkgroups = 4096; + +static bool same_shape(const Value & lhs, const Value & rhs) { + for (int i = 0; i < GGML_MAX_DIMS; ++i) { + if (lhs.ne[i] != rhs.ne[i]) { + return false; + } + } + return true; +} + +static const Value * graph_value(const Graph & graph, ValueId id) { + return graph.values().find(id); +} + +static bool supports_add_f32_dispatch(const Graph & graph, const GraphNode * node) { + if (node == nullptr || node->op != GGML_OP_ADD || node->inputs.size() != 2) { + return false; + } + const Value * output = graph_value(graph, node->output); + const Value * a = graph_value(graph, node->inputs[0]); + const Value * b = graph_value(graph, node->inputs[1]); + if (output == nullptr || a == nullptr || b == nullptr) { + return false; + } + return output->type == GGML_TYPE_F32 && a->type == GGML_TYPE_F32 && b->type == GGML_TYPE_F32 && + same_shape(*output, *a) && same_shape(*output, *b) && output->contiguous && a->contiguous && b->contiguous && + output->element_count > 0 && + static_cast(output->element_count) <= std::numeric_limits::max() && + (static_cast(output->element_count) + 255) / 256 <= kAddMaxWorkgroups; +} + +static bool match_add_f32_dispatch(const DispatchMatchContext & context, DispatchMatch & match) { + if (!supports_add_f32_dispatch(context.graph, context.root_node)) { + return false; + } + const Value * output = graph_value(context.graph, context.root_node->output); + const Value * a = graph_value(context.graph, context.root_node->inputs[0]); + const Value * b = graph_value(context.graph, context.root_node->inputs[1]); + if (output == nullptr || a == nullptr || b == nullptr) { + return false; + } + + Dispatch dispatch; + dispatch.kernel = make_kernel_specialization(kAddF32Kernel); + dispatch.kernel.integer_parameters.emplace("element_count", output->element_count); + dispatch.bindings.push_back({ a->id, 0, a->byte_count }); + dispatch.bindings.push_back({ b->id, 0, b->byte_count }); + dispatch.bindings.push_back({ output->id, 0, output->byte_count }); + + match.covered_nodes.push_back(context.root_index); + match.dispatches.push_back(std::move(dispatch)); + return true; +} + +void register_add_dispatch(DispatchRegistryBuilder & registry) { + registry.add({ + "common.add_f32", + GGML_OP_ADD, + DispatchMatchKind::SingleOp, + 0, + DispatchSource::Common, + match_add_f32_dispatch, + }); +} + +} // namespace ggml::hrx diff --git a/ggml/src/ggml-hrx/dispatch_registration/dispatch-add.h b/ggml/src/ggml-hrx/dispatch_registration/dispatch-add.h new file mode 100644 index 000000000000..3054b942958c --- /dev/null +++ b/ggml/src/ggml-hrx/dispatch_registration/dispatch-add.h @@ -0,0 +1,9 @@ +#pragma once + +#include "dispatch-registry.h" + +namespace ggml::hrx { + +void register_add_dispatch(DispatchRegistryBuilder & registry); + +} // namespace ggml::hrx diff --git a/ggml/src/ggml-hrx/dispatch_registration/dispatch-gather-add.cpp b/ggml/src/ggml-hrx/dispatch_registration/dispatch-gather-add.cpp new file mode 100644 index 000000000000..2bcd8b4d4616 --- /dev/null +++ b/ggml/src/ggml-hrx/dispatch_registration/dispatch-gather-add.cpp @@ -0,0 +1,200 @@ +#include "dispatch-gather-add.h" + +#include "ggml.h" +#include "kernel-corpus/kernel-corpus-catalog-verify.h" + +#include +#include + +namespace ggml::hrx { +namespace { + +static constexpr KernelCatalogRef kGatherAddF32Kernel = GGML_HRX_KERNEL_REF("qwen3_moe", "ggml_gather_add_f32"); + +static const Value * graph_value(const Graph & graph, ValueId id) { + return graph.values().find(id); +} + +static bool same_shape(const Value & lhs, const Value & rhs) { + for (int i = 0; i < GGML_MAX_DIMS; ++i) { + if (lhs.ne[i] != rhs.ne[i]) { + return false; + } + } + return true; +} + +static bool is_2d_f32(const Value & value) { + return value.type == GGML_TYPE_F32 && value.contiguous && value.ne[0] > 0 && value.ne[1] > 0 && value.ne[2] == 1 && + value.ne[3] == 1; +} + +static bool is_row_id_tensor(const Value & value, int64_t output_token_count) { + return value.type == GGML_TYPE_I32 && value.contiguous && value.element_count == output_token_count && + output_token_count > 0; +} + +static bool is_supported_hidden_size(int64_t hidden_size) { + return hidden_size >= 128 && hidden_size <= 32768 && hidden_size % 128 == 0; +} + +static bool is_supported_token_count(int64_t token_count) { + return token_count >= 1 && token_count <= 2048; +} + +struct GatherAddMatch { + const GraphNode * first_get_rows = nullptr; + const GraphNode * second_get_rows = nullptr; + const GraphNode * add_node = nullptr; + const Value * first_source = nullptr; + const Value * second_source = nullptr; + const Value * row_ids = nullptr; + const Value * output = nullptr; + size_t first_get_rows_index = 0; + size_t second_get_rows_index = 0; + size_t add_node_index = 0; + int64_t source_token_count = 0; + int64_t output_token_count = 0; + int64_t hidden_size = 0; + + bool matched() const { + return first_get_rows != nullptr && second_get_rows != nullptr && add_node != nullptr && + first_source != nullptr && second_source != nullptr && row_ids != nullptr && output != nullptr; + } +}; + +static const GraphNode * find_single_add_consumer(const Graph & graph, const GraphNode & get_rows) { + const std::vector & consumers = graph.index().consumers(get_rows.output); + if (consumers.size() != 1) { + return nullptr; + } + const GraphNode * add = consumers.front(); + return add != nullptr && add->op == GGML_OP_ADD && add->inputs.size() == 2 ? add : nullptr; +} + +static const GraphNode * peer_get_rows_input(const Graph & graph, + const GraphNode & add, + const GraphNode & root_get_rows) { + ValueId peer_output; + if (add.inputs[0] == root_get_rows.output) { + peer_output = add.inputs[1]; + } else if (add.inputs[1] == root_get_rows.output) { + peer_output = add.inputs[0]; + } else { + return nullptr; + } + + const GraphNode * peer = graph.index().producer(peer_output); + return peer != nullptr && peer->op == GGML_OP_GET_ROWS && peer->inputs.size() == 2 ? peer : nullptr; +} + +static GatherAddMatch match_gather_add_f32(const Graph & graph, const GraphNode * node, size_t node_index) { + GatherAddMatch match; + if (node == nullptr || node->op != GGML_OP_GET_ROWS || node->inputs.size() != 2 || !graph.has_index()) { + return match; + } + + const GraphNode * add_node = find_single_add_consumer(graph, *node); + const GraphNode * peer = add_node != nullptr ? peer_get_rows_input(graph, *add_node, *node) : nullptr; + if (add_node == nullptr || peer == nullptr) { + return {}; + } + + size_t peer_index = 0; + size_t add_index = 0; + if (!graph.index().node_index(peer, peer_index) || !graph.index().node_index(add_node, add_index)) { + return {}; + } + + const Value * first_source = graph_value(graph, node->inputs[0]); + const Value * first_ids = graph_value(graph, node->inputs[1]); + const Value * first_output = graph_value(graph, node->output); + const Value * second_source = graph_value(graph, peer->inputs[0]); + const Value * second_ids = graph_value(graph, peer->inputs[1]); + const Value * second_output = graph_value(graph, peer->output); + const Value * output = graph_value(graph, add_node->output); + if (first_source == nullptr || first_ids == nullptr || first_output == nullptr || second_source == nullptr || + second_ids == nullptr || second_output == nullptr || output == nullptr) { + return {}; + } + if (node->inputs[1] != peer->inputs[1]) { + return {}; + } + if (!is_2d_f32(*first_source) || !is_2d_f32(*second_source) || !is_2d_f32(*first_output) || + !is_2d_f32(*second_output) || !is_2d_f32(*output)) { + return {}; + } + if (!same_shape(*first_source, *second_source) || !same_shape(*first_output, *second_output) || + !same_shape(*first_output, *output)) { + return {}; + } + + const int64_t hidden_size = first_source->ne[0]; + const int64_t source_token_count = first_source->ne[1]; + const int64_t output_token_count = first_output->ne[1]; + if (first_output->ne[0] != hidden_size || !is_row_id_tensor(*first_ids, output_token_count) || + !is_row_id_tensor(*second_ids, output_token_count)) { + return {}; + } + if (!is_supported_hidden_size(hidden_size) || !is_supported_token_count(source_token_count) || + !is_supported_token_count(output_token_count)) { + return {}; + } + + match.first_get_rows = node; + match.second_get_rows = peer; + match.add_node = add_node; + match.first_source = first_source; + match.second_source = second_source; + match.row_ids = first_ids; + match.output = output; + match.first_get_rows_index = node_index; + match.second_get_rows_index = peer_index; + match.add_node_index = add_index; + match.source_token_count = source_token_count; + match.output_token_count = output_token_count; + match.hidden_size = hidden_size; + return match; +} + +} // namespace + +static bool match_gather_add_f32_dispatch(const DispatchMatchContext & context, DispatchMatch & match) { + const GatherAddMatch gather_add = match_gather_add_f32(context.graph, context.root_node, context.root_index); + if (!gather_add.matched() || gather_add.first_get_rows_index >= context.covered_nodes.size() || + gather_add.second_get_rows_index >= context.covered_nodes.size() || + gather_add.add_node_index >= context.covered_nodes.size() || + context.covered_nodes[gather_add.first_get_rows_index] || + context.covered_nodes[gather_add.second_get_rows_index] || context.covered_nodes[gather_add.add_node_index]) { + return false; + } + + Dispatch dispatch; + dispatch.kernel = make_kernel_specialization(kGatherAddF32Kernel); + dispatch.kernel.integer_parameters.emplace("source_token_count", gather_add.source_token_count); + dispatch.kernel.integer_parameters.emplace("output_token_count", gather_add.output_token_count); + dispatch.kernel.integer_parameters.emplace("hidden_size", gather_add.hidden_size); + dispatch.bindings.push_back({ gather_add.first_source->id, 0, gather_add.first_source->byte_count }); + dispatch.bindings.push_back({ gather_add.second_source->id, 0, gather_add.second_source->byte_count }); + dispatch.bindings.push_back({ gather_add.row_ids->id, 0, gather_add.row_ids->byte_count }); + dispatch.bindings.push_back({ gather_add.output->id, 0, gather_add.output->byte_count }); + + match.covered_nodes.push_back(gather_add.first_get_rows_index); + match.covered_nodes.push_back(gather_add.second_get_rows_index); + match.covered_nodes.push_back(gather_add.add_node_index); + match.dispatches.push_back(std::move(dispatch)); + return true; +} + +void register_gather_add_dispatch(DispatchRegistryBuilder & registry) { + registry.add({ + "common.gather_add_f32", + GGML_OP_GET_ROWS, + DispatchMatchKind::Fused, + 1000, + DispatchSource::Common, + match_gather_add_f32_dispatch, + }); +} + +} // namespace ggml::hrx diff --git a/ggml/src/ggml-hrx/dispatch_registration/dispatch-gather-add.h b/ggml/src/ggml-hrx/dispatch_registration/dispatch-gather-add.h new file mode 100644 index 000000000000..b8d43035350f --- /dev/null +++ b/ggml/src/ggml-hrx/dispatch_registration/dispatch-gather-add.h @@ -0,0 +1,9 @@ +#pragma once + +#include "dispatch-registry.h" + +namespace ggml::hrx { + +void register_gather_add_dispatch(DispatchRegistryBuilder & registry); + +} // namespace ggml::hrx diff --git a/ggml/src/ggml-hrx/dispatch_registration/dispatch-get-rows.cpp b/ggml/src/ggml-hrx/dispatch_registration/dispatch-get-rows.cpp new file mode 100644 index 000000000000..68541b61c5af --- /dev/null +++ b/ggml/src/ggml-hrx/dispatch_registration/dispatch-get-rows.cpp @@ -0,0 +1,153 @@ +// Copyright 2026 The HRX Authors +// SPDX-License-Identifier: Apache-2.0 + +#include "dispatch-get-rows.h" + +#include "ggml.h" +#include "kernel-corpus/kernel-corpus-catalog-verify.h" + +#include +#include +#include + +namespace ggml::hrx { + +static constexpr KernelCatalogRef kGetRowsF32Kernel = GGML_HRX_KERNEL_REF("qwen3_moe", "ggml_get_rows_f32"); + +static const Value * graph_value(const Graph & graph, ValueId id) { + return graph.values().find(id); +} + +// GET_ROWS: output[nrows, width] = source[ids[nrows], width], ids is i32. +// +// The kernel's domain is fully pinned by the Loom source and by the ABI, and the +// claim has to match it exactly - a claimed shape that the dispatcher cannot run +// either fails to compile or (worse) fails at submission time and leaves the device +// stream unusable: +// +// - `index.assume` in get_rows_f32.loom caps width to [32, 16384], the row counts to +// [1, 131072] / [1, 2048]; specialising a constant outside those ranges is a +// compile error, not a fallback. +// - The kernel indexes a flat 2D view with (output_row, column) and derives the +// launch grid as (ceil(width/256), output_row_count, 1), so batched (ne[2]/ne[3] > +// 1) or higher-rank get_rows is not covered. +// - The HRX runtime charges the dispatch's kernarg footprint against a fixed +// 262144-entry ring; measured on gfx1151 the footprint is about +// 16 * ceil(width/256) * (source_row_count + output_row_count) blocks. Crossing +// the ring aborts the submission and corrupts the stream, so claim only shapes +// that stay comfortably inside it. +static constexpr uint64_t kGetRowsMinWidth = 32; +static constexpr uint64_t kGetRowsMaxWidth = 16384; +static constexpr uint64_t kGetRowsMaxOutputRows = 2048; +static constexpr uint64_t kGetRowsMaxSourceRows = 131072; +// 0.75 x the measured 262144-entry limit, leaving headroom for the per-dispatch +// constant terms the estimate above does not model. +static constexpr uint64_t kGetRowsMaxKernargBlocks = 196608; + +static bool supports_get_rows_f32_dispatch(const Graph & graph, const GraphNode * node) { + if (node == nullptr || node->op != GGML_OP_GET_ROWS || node->inputs.size() != 2) { + return false; + } + const Value * output = graph_value(graph, node->output); + const Value * source = graph_value(graph, node->inputs[0]); + const Value * ids = graph_value(graph, node->inputs[1]); + if (output == nullptr || source == nullptr || ids == nullptr) { + return false; + } + // Source may be f32 or q8_0 (the kernel dequantizes q8_0 inline); any + // other quantized embed type (q4k etc.) is covered by the fused dispatch. + if (output->type != GGML_TYPE_F32) { + return false; + } + const bool source_f32 = source->type == GGML_TYPE_F32; + const bool source_q8 = source->type == GGML_TYPE_Q8_0; + if (!source_f32 && !source_q8) { + return false; + } + if (ids->type != GGML_TYPE_I32) { + return false; + } + if (!output->contiguous || !source->contiguous || !ids->contiguous) { + return false; + } + // Flat 2D gather only: the kernel has no batch dimensions. + if (source->ne[2] != 1 || source->ne[3] != 1 || output->ne[2] != 1 || output->ne[3] != 1 || ids->ne[1] != 1 || + ids->ne[2] != 1) { + return false; + } + // source is [width, source_rows]; output is [width, output_rows] == [ids->ne[0], width]. + if (source->ne[0] != output->ne[0] || output->ne[1] != ids->ne[0]) { + return false; + } + const uint64_t width = static_cast(output->ne[0]); + const uint64_t output_rows = static_cast(output->ne[1]); + const uint64_t source_rows = static_cast(source->ne[1]); + if (output_rows == 0 || source_rows == 0) { + // Degenerate gather (e.g. an empty id list during decode): nothing to dispatch. + return false; + } + if (width < kGetRowsMinWidth || width > kGetRowsMaxWidth) { + return false; + } + if (output_rows > kGetRowsMaxOutputRows || source_rows > kGetRowsMaxSourceRows) { + return false; + } + // q8_0 rows are whole 32-value blocks. + if (source_q8 && width % 32 != 0) { + return false; + } + const uint64_t column_workgroups = (width + 255) / 256; + const uint64_t kernarg_blocks = 16 * column_workgroups * (source_rows + output_rows) + 256; + if (kernarg_blocks > kGetRowsMaxKernargBlocks) { + return false; + } + return true; +} + +static bool match_get_rows_f32_dispatch(const DispatchMatchContext & context, DispatchMatch & match) { + if (!supports_get_rows_f32_dispatch(context.graph, context.root_node)) { + return false; + } + const Value * output = graph_value(context.graph, context.root_node->output); + const Value * source = graph_value(context.graph, context.root_node->inputs[0]); + const Value * ids = graph_value(context.graph, context.root_node->inputs[1]); + if (output == nullptr || source == nullptr || ids == nullptr) { + return false; + } + + Dispatch dispatch; + dispatch.kernel = make_kernel_specialization(kGetRowsF32Kernel); + // source is [width, row_count]; the row count is ne[1] (ne[0] is the width). + dispatch.kernel.integer_parameters.emplace("source_row_count", static_cast(source->ne[1])); + dispatch.kernel.integer_parameters.emplace("output_row_count", static_cast(output->ne[1])); + dispatch.kernel.integer_parameters.emplace("width", static_cast(output->ne[0])); + // source_format: 0 = f32, 1 = q8_0 (mirrors the kernel's format branch). + const int64_t source_format = source->type == GGML_TYPE_Q8_0 ? 1 : 0; + dispatch.kernel.integer_parameters.emplace("source_format", source_format); + dispatch.bindings.push_back({ source->id, 0, source->byte_count }); + dispatch.bindings.push_back({ ids->id, 0, ids->byte_count }); + dispatch.bindings.push_back({ output->id, 0, output->byte_count }); + + match.covered_nodes.push_back(context.root_index); + match.dispatches.push_back(std::move(dispatch)); + return true; +} + +void register_get_rows_dispatch(DispatchRegistryBuilder & registry) { + // Registered for the shapes the standalone kernel can actually run. Anything + // outside supports_get_rows_f32_dispatch() - batched/higher-rank gathers, widths + // outside [32, 16384], vocabularies or row counts that would overflow the + // runtime's kernarg ring - is left to the CPU backend, which needs no copies + // because the HRX buffer type is host-visible. + registry.add({ + "common.get_rows_f32", + GGML_OP_GET_ROWS, + DispatchMatchKind::SingleOp, + 0, + DispatchSource::Common, + match_get_rows_f32_dispatch, + }); + GGML_UNUSED(registry); +} + +} // namespace ggml::hrx diff --git a/ggml/src/ggml-hrx/dispatch_registration/dispatch-get-rows.h b/ggml/src/ggml-hrx/dispatch_registration/dispatch-get-rows.h new file mode 100644 index 000000000000..cb623a68d710 --- /dev/null +++ b/ggml/src/ggml-hrx/dispatch_registration/dispatch-get-rows.h @@ -0,0 +1,9 @@ +#pragma once + +#include "dispatch-registry.h" + +namespace ggml::hrx { + +void register_get_rows_dispatch(DispatchRegistryBuilder & registry); + +} // namespace ggml::hrx diff --git a/ggml/src/ggml-hrx/dispatch_registration/dispatch-llm-matmul.cpp b/ggml/src/ggml-hrx/dispatch_registration/dispatch-llm-matmul.cpp new file mode 100644 index 000000000000..66706333e202 --- /dev/null +++ b/ggml/src/ggml-hrx/dispatch_registration/dispatch-llm-matmul.cpp @@ -0,0 +1,339 @@ +// Copyright 2026 The HRX Authors +// SPDX-License-Identifier: Apache-2.0 + +#include "dispatch-llm-matmul.h" + +#include "dispatch-llm-shapes.h" +#include "ggml.h" +#include "graph/graph-matcher.h" +#include "kernel-corpus/kernel-corpus-catalog-verify.h" + +#include +#include +#include +#include + +namespace ggml::hrx { +namespace { + +static constexpr KernelCatalogRef kLlmDenseLinearQ4KF16WmmaKernel = + GGML_HRX_KERNEL_REF("qwen3_moe", "qwen3_moe_dense_linear_q4k_f16_wmma"); +static constexpr KernelCatalogRef kLlmDenseLinearQ6KF16WmmaKernel = + GGML_HRX_KERNEL_REF("qwen3_moe", "qwen3_moe_dense_linear_q6k_f16_wmma"); +static constexpr KernelCatalogRef kLlmIq3XxsMatVecKernel = + GGML_HRX_KERNEL_REF("qwen3_moe", "ggml_mul_mat_vec_iq3xxs_f32"); + +// IQ3_XXS codebook: 256x u32 grid (LE) | 128 sign bytes | 8 sign-mask bytes. +static const uint8_t kIq3XxsTables[1168] = { + 0x04, 0x04, 0x04, 0x04, 0x14, 0x04, 0x04, 0x04, 0x24, 0x04, 0x04, 0x04, + 0x0c, 0x0c, 0x04, 0x04, 0x1c, 0x0c, 0x04, 0x04, 0x3e, 0x0c, 0x04, 0x04, + 0x04, 0x14, 0x04, 0x04, 0x14, 0x14, 0x04, 0x04, 0x0c, 0x1c, 0x04, 0x04, + 0x14, 0x24, 0x04, 0x04, 0x1c, 0x3e, 0x04, 0x04, 0x2c, 0x3e, 0x04, 0x04, + 0x0c, 0x04, 0x0c, 0x04, 0x1c, 0x04, 0x0c, 0x04, 0x04, 0x0c, 0x0c, 0x04, + 0x14, 0x0c, 0x0c, 0x04, 0x0c, 0x14, 0x0c, 0x04, 0x2c, 0x14, 0x0c, 0x04, + 0x04, 0x1c, 0x0c, 0x04, 0x14, 0x1c, 0x0c, 0x04, 0x0c, 0x24, 0x0c, 0x04, + 0x24, 0x2c, 0x0c, 0x04, 0x04, 0x3e, 0x0c, 0x04, 0x04, 0x04, 0x14, 0x04, + 0x14, 0x04, 0x14, 0x04, 0x24, 0x04, 0x14, 0x04, 0x0c, 0x0c, 0x14, 0x04, + 0x04, 0x14, 0x14, 0x04, 0x14, 0x14, 0x14, 0x04, 0x0c, 0x1c, 0x14, 0x04, + 0x1c, 0x1c, 0x14, 0x04, 0x3e, 0x1c, 0x14, 0x04, 0x0c, 0x2c, 0x14, 0x04, + 0x3e, 0x2c, 0x14, 0x04, 0x2c, 0x3e, 0x14, 0x04, 0x0c, 0x04, 0x1c, 0x04, + 0x3e, 0x04, 0x1c, 0x04, 0x04, 0x0c, 0x1c, 0x04, 0x14, 0x0c, 0x1c, 0x04, + 0x2c, 0x14, 0x1c, 0x04, 0x04, 0x3e, 0x1c, 0x04, 0x1c, 0x0c, 0x24, 0x04, + 0x3e, 0x1c, 0x24, 0x04, 0x24, 0x24, 0x24, 0x04, 0x3e, 0x2c, 0x24, 0x04, + 0x1c, 0x3e, 0x24, 0x04, 0x2c, 0x3e, 0x24, 0x04, 0x0c, 0x04, 0x2c, 0x04, + 0x3e, 0x04, 0x2c, 0x04, 0x14, 0x1c, 0x2c, 0x04, 0x14, 0x2c, 0x2c, 0x04, + 0x2c, 0x1c, 0x34, 0x04, 0x24, 0x34, 0x34, 0x04, 0x04, 0x0c, 0x3e, 0x04, + 0x24, 0x0c, 0x3e, 0x04, 0x34, 0x0c, 0x3e, 0x04, 0x1c, 0x24, 0x3e, 0x04, + 0x0c, 0x34, 0x3e, 0x04, 0x0c, 0x04, 0x04, 0x0c, 0x1c, 0x04, 0x04, 0x0c, + 0x04, 0x0c, 0x04, 0x0c, 0x14, 0x0c, 0x04, 0x0c, 0x0c, 0x14, 0x04, 0x0c, + 0x1c, 0x14, 0x04, 0x0c, 0x04, 0x1c, 0x04, 0x0c, 0x14, 0x1c, 0x04, 0x0c, + 0x24, 0x1c, 0x04, 0x0c, 0x3e, 0x24, 0x04, 0x0c, 0x04, 0x2c, 0x04, 0x0c, + 0x04, 0x04, 0x0c, 0x0c, 0x14, 0x04, 0x0c, 0x0c, 0x0c, 0x0c, 0x0c, 0x0c, + 0x04, 0x14, 0x0c, 0x0c, 0x14, 0x14, 0x0c, 0x0c, 0x0c, 0x04, 0x14, 0x0c, + 0x1c, 0x04, 0x14, 0x0c, 0x04, 0x0c, 0x14, 0x0c, 0x14, 0x0c, 0x14, 0x0c, + 0x0c, 0x14, 0x14, 0x0c, 0x04, 0x1c, 0x14, 0x0c, 0x14, 0x3e, 0x14, 0x0c, + 0x04, 0x04, 0x1c, 0x0c, 0x14, 0x04, 0x1c, 0x0c, 0x04, 0x14, 0x1c, 0x0c, + 0x0c, 0x1c, 0x1c, 0x0c, 0x34, 0x24, 0x1c, 0x0c, 0x34, 0x34, 0x1c, 0x0c, + 0x0c, 0x04, 0x24, 0x0c, 0x2c, 0x04, 0x24, 0x0c, 0x04, 0x2c, 0x24, 0x0c, + 0x04, 0x14, 0x2c, 0x0c, 0x24, 0x14, 0x2c, 0x0c, 0x34, 0x24, 0x2c, 0x0c, + 0x0c, 0x3e, 0x2c, 0x0c, 0x2c, 0x04, 0x34, 0x0c, 0x14, 0x14, 0x3e, 0x0c, + 0x04, 0x24, 0x3e, 0x0c, 0x04, 0x04, 0x04, 0x14, 0x14, 0x04, 0x04, 0x14, + 0x0c, 0x0c, 0x04, 0x14, 0x1c, 0x0c, 0x04, 0x14, 0x04, 0x14, 0x04, 0x14, + 0x14, 0x14, 0x04, 0x14, 0x34, 0x14, 0x04, 0x14, 0x0c, 0x1c, 0x04, 0x14, + 0x14, 0x24, 0x04, 0x14, 0x0c, 0x04, 0x0c, 0x14, 0x1c, 0x04, 0x0c, 0x14, + 0x2c, 0x04, 0x0c, 0x14, 0x04, 0x0c, 0x0c, 0x14, 0x14, 0x0c, 0x0c, 0x14, + 0x0c, 0x14, 0x0c, 0x14, 0x04, 0x1c, 0x0c, 0x14, 0x1c, 0x34, 0x0c, 0x14, + 0x3e, 0x34, 0x0c, 0x14, 0x04, 0x3e, 0x0c, 0x14, 0x04, 0x04, 0x14, 0x14, + 0x14, 0x04, 0x14, 0x14, 0x0c, 0x0c, 0x14, 0x14, 0x3e, 0x0c, 0x14, 0x14, + 0x04, 0x14, 0x14, 0x14, 0x14, 0x14, 0x14, 0x14, 0x3e, 0x1c, 0x14, 0x14, + 0x04, 0x24, 0x14, 0x14, 0x2c, 0x2c, 0x14, 0x14, 0x0c, 0x04, 0x1c, 0x14, + 0x04, 0x0c, 0x1c, 0x14, 0x24, 0x0c, 0x1c, 0x14, 0x04, 0x3e, 0x1c, 0x14, + 0x24, 0x3e, 0x1c, 0x14, 0x2c, 0x1c, 0x24, 0x14, 0x1c, 0x2c, 0x24, 0x14, + 0x1c, 0x04, 0x2c, 0x14, 0x3e, 0x14, 0x2c, 0x14, 0x0c, 0x24, 0x2c, 0x14, + 0x24, 0x3e, 0x2c, 0x14, 0x0c, 0x04, 0x3e, 0x14, 0x1c, 0x04, 0x3e, 0x14, + 0x34, 0x0c, 0x3e, 0x14, 0x2c, 0x24, 0x3e, 0x14, 0x0c, 0x04, 0x04, 0x1c, + 0x04, 0x0c, 0x04, 0x1c, 0x14, 0x0c, 0x04, 0x1c, 0x0c, 0x14, 0x04, 0x1c, + 0x1c, 0x14, 0x04, 0x1c, 0x04, 0x2c, 0x04, 0x1c, 0x2c, 0x34, 0x04, 0x1c, + 0x14, 0x3e, 0x04, 0x1c, 0x04, 0x04, 0x0c, 0x1c, 0x14, 0x04, 0x0c, 0x1c, + 0x04, 0x14, 0x0c, 0x1c, 0x0c, 0x1c, 0x0c, 0x1c, 0x24, 0x24, 0x0c, 0x1c, + 0x34, 0x24, 0x0c, 0x1c, 0x0c, 0x04, 0x14, 0x1c, 0x1c, 0x04, 0x14, 0x1c, + 0x04, 0x0c, 0x14, 0x1c, 0x2c, 0x14, 0x14, 0x1c, 0x14, 0x2c, 0x14, 0x1c, + 0x14, 0x3e, 0x14, 0x1c, 0x0c, 0x0c, 0x1c, 0x1c, 0x1c, 0x1c, 0x1c, 0x1c, + 0x04, 0x1c, 0x24, 0x1c, 0x3e, 0x24, 0x24, 0x1c, 0x14, 0x3e, 0x24, 0x1c, + 0x04, 0x04, 0x2c, 0x1c, 0x34, 0x04, 0x2c, 0x1c, 0x14, 0x14, 0x2c, 0x1c, + 0x2c, 0x2c, 0x2c, 0x1c, 0x24, 0x0c, 0x34, 0x1c, 0x34, 0x1c, 0x34, 0x1c, + 0x1c, 0x34, 0x34, 0x1c, 0x1c, 0x1c, 0x3e, 0x1c, 0x04, 0x34, 0x3e, 0x1c, + 0x24, 0x04, 0x04, 0x24, 0x3e, 0x0c, 0x04, 0x24, 0x2c, 0x1c, 0x04, 0x24, + 0x3e, 0x1c, 0x04, 0x24, 0x1c, 0x2c, 0x04, 0x24, 0x3e, 0x2c, 0x04, 0x24, + 0x24, 0x3e, 0x0c, 0x24, 0x04, 0x14, 0x14, 0x24, 0x3e, 0x1c, 0x14, 0x24, + 0x04, 0x24, 0x14, 0x24, 0x04, 0x34, 0x14, 0x24, 0x34, 0x34, 0x14, 0x24, + 0x3e, 0x04, 0x1c, 0x24, 0x2c, 0x24, 0x1c, 0x24, 0x24, 0x04, 0x24, 0x24, + 0x0c, 0x2c, 0x24, 0x24, 0x24, 0x34, 0x24, 0x24, 0x2c, 0x14, 0x2c, 0x24, + 0x1c, 0x24, 0x2c, 0x24, 0x04, 0x3e, 0x2c, 0x24, 0x2c, 0x04, 0x3e, 0x24, + 0x04, 0x0c, 0x3e, 0x24, 0x14, 0x0c, 0x3e, 0x24, 0x04, 0x1c, 0x3e, 0x24, + 0x14, 0x0c, 0x04, 0x2c, 0x0c, 0x24, 0x04, 0x2c, 0x04, 0x3e, 0x04, 0x2c, + 0x04, 0x04, 0x0c, 0x2c, 0x34, 0x04, 0x0c, 0x2c, 0x34, 0x14, 0x0c, 0x2c, + 0x2c, 0x2c, 0x0c, 0x2c, 0x24, 0x0c, 0x14, 0x2c, 0x14, 0x1c, 0x14, 0x2c, + 0x14, 0x3e, 0x14, 0x2c, 0x14, 0x04, 0x1c, 0x2c, 0x1c, 0x2c, 0x1c, 0x2c, + 0x04, 0x0c, 0x24, 0x2c, 0x1c, 0x14, 0x24, 0x2c, 0x3e, 0x14, 0x24, 0x2c, + 0x14, 0x3e, 0x24, 0x2c, 0x14, 0x04, 0x2c, 0x2c, 0x0c, 0x1c, 0x2c, 0x2c, + 0x04, 0x2c, 0x34, 0x2c, 0x24, 0x14, 0x3e, 0x2c, 0x14, 0x24, 0x3e, 0x2c, + 0x24, 0x14, 0x04, 0x34, 0x24, 0x24, 0x04, 0x34, 0x34, 0x24, 0x04, 0x34, + 0x24, 0x34, 0x04, 0x34, 0x0c, 0x14, 0x0c, 0x34, 0x0c, 0x34, 0x0c, 0x34, + 0x3e, 0x0c, 0x14, 0x34, 0x24, 0x34, 0x14, 0x34, 0x04, 0x1c, 0x1c, 0x34, + 0x34, 0x1c, 0x1c, 0x34, 0x24, 0x24, 0x24, 0x34, 0x2c, 0x04, 0x2c, 0x34, + 0x14, 0x2c, 0x2c, 0x34, 0x1c, 0x1c, 0x34, 0x34, 0x1c, 0x04, 0x3e, 0x34, + 0x0c, 0x14, 0x3e, 0x34, 0x1c, 0x04, 0x04, 0x3e, 0x2c, 0x04, 0x04, 0x3e, + 0x3e, 0x04, 0x04, 0x3e, 0x04, 0x0c, 0x04, 0x3e, 0x14, 0x1c, 0x04, 0x3e, + 0x14, 0x2c, 0x04, 0x3e, 0x34, 0x14, 0x0c, 0x3e, 0x04, 0x24, 0x0c, 0x3e, + 0x14, 0x0c, 0x14, 0x3e, 0x2c, 0x24, 0x14, 0x3e, 0x14, 0x2c, 0x14, 0x3e, + 0x04, 0x04, 0x1c, 0x3e, 0x2c, 0x0c, 0x1c, 0x3e, 0x1c, 0x1c, 0x1c, 0x3e, + 0x04, 0x34, 0x1c, 0x3e, 0x0c, 0x14, 0x24, 0x3e, 0x0c, 0x24, 0x24, 0x3e, + 0x04, 0x04, 0x2c, 0x3e, 0x14, 0x04, 0x2c, 0x3e, 0x24, 0x14, 0x2c, 0x3e, + 0x04, 0x1c, 0x34, 0x3e, 0x00, 0x81, 0x82, 0x03, 0x84, 0x05, 0x06, 0x87, + 0x88, 0x09, 0x0a, 0x8b, 0x0c, 0x8d, 0x8e, 0x0f, 0x90, 0x11, 0x12, 0x93, + 0x14, 0x95, 0x96, 0x17, 0x18, 0x99, 0x9a, 0x1b, 0x9c, 0x1d, 0x1e, 0x9f, + 0xa0, 0x21, 0x22, 0xa3, 0x24, 0xa5, 0xa6, 0x27, 0x28, 0xa9, 0xaa, 0x2b, + 0xac, 0x2d, 0x2e, 0xaf, 0x30, 0xb1, 0xb2, 0x33, 0xb4, 0x35, 0x36, 0xb7, + 0xb8, 0x39, 0x3a, 0xbb, 0x3c, 0xbd, 0xbe, 0x3f, 0xc0, 0x41, 0x42, 0xc3, + 0x44, 0xc5, 0xc6, 0x47, 0x48, 0xc9, 0xca, 0x4b, 0xcc, 0x4d, 0x4e, 0xcf, + 0x50, 0xd1, 0xd2, 0x53, 0xd4, 0x55, 0x56, 0xd7, 0xd8, 0x59, 0x5a, 0xdb, + 0x5c, 0xdd, 0xde, 0x5f, 0x60, 0xe1, 0xe2, 0x63, 0xe4, 0x65, 0x66, 0xe7, + 0xe8, 0x69, 0x6a, 0xeb, 0x6c, 0xed, 0xee, 0x6f, 0xf0, 0x71, 0x72, 0xf3, + 0x74, 0xf5, 0xf6, 0x77, 0x78, 0xf9, 0xfa, 0x7b, 0xfc, 0x7d, 0x7e, 0xff, + 0x01, 0x02, 0x04, 0x08, 0x10, 0x20, 0x40, 0x80, 0x00, 0x00, 0x00, 0x00, + 0x00, 0x00, 0x00, 0x00, +}; + +static const Value * graph_value(const Graph & graph, ValueId id) { + return graph.values().find(id); +} + +static bool is_2d(const Value & value) { + return value.ne[0] > 0 && value.ne[1] > 0 && value.ne[2] == 1 && value.ne[3] == 1; +} + +static bool is_supported_dense_input_size(int64_t input_size) { + return input_size >= 256 && input_size <= 32768 && input_size % 256 == 0; +} + +static bool is_supported_dense_output_size(int64_t output_size) { + return output_size >= 1 && output_size <= 262144; +} + +enum class LlmDenseMatmulRoute { + Q4K, + Q6K, +}; + +static std::string to_config_value(int64_t value) { + return std::to_string(value); +} + +struct LlmDenseMatmulMatch { + const Value * input = nullptr; + const Value * weight = nullptr; + const Value * output = nullptr; + KernelCatalogRef kernel = {}; + int64_t input_size = 0; + int64_t output_size = 0; + int64_t token_count = 0; + + bool matched() const { + return input != nullptr && weight != nullptr && output != nullptr && kernel.id != kUncatalogedKernelId; + } +}; + +// The HRX runtime charges every dispatch against a fixed 262144-entry kernarg ring and +// aborts the submission (leaving the device stream unusable) when it does not fit. +// Measured on gfx1151 the graph-replay footprint of the dense wmma route grows with +// output_size x token_count: with input_size 1024 a 16384x64 tile already needs 284957 +// entries, 8192x512 needs 402833, while everything up to 4096x512 stays inside the ring. +// The route exists for LLM projections (hidden/ffn sizes), not for the vocab-sized LM +// head, so cap the claim inside the measured-safe region. +static constexpr int64_t kLlmDenseMaxOutputSize = 4096; +static constexpr int64_t kLlmDenseMaxTokenCount = 256; + +static LlmDenseMatmulMatch match_llm_dense_matmul(const Graph & graph, + const GraphNode * node, + LlmDenseMatmulRoute route) { + LlmDenseMatmulMatch match; + if (node == nullptr || node->op != GGML_OP_MUL_MAT || node->inputs.size() != 2) { + return match; + } + + const Value * weight = graph_value(graph, node->inputs[0]); + const Value * input = graph_value(graph, node->inputs[1]); + const Value * output = graph_value(graph, node->output); + if (weight == nullptr || input == nullptr || output == nullptr || !is_2d(*weight) || !is_2d(*input) || + !is_2d(*output) || !weight->contiguous || !input->contiguous || !output->contiguous || + input->type != GGML_TYPE_F32 || output->type != GGML_TYPE_F32) { + return {}; + } + + const int64_t input_size = weight->ne[0]; + const int64_t output_size = weight->ne[1]; + const int64_t token_count = input->ne[1]; + if (input->ne[0] != input_size || output->ne[0] != output_size || output->ne[1] != token_count || + !is_llm_prefill_query_length(kActiveLlmMoeDispatchProfile, token_count) || + !is_supported_dense_input_size(input_size) || !is_supported_dense_output_size(output_size) || + output_size > kLlmDenseMaxOutputSize || token_count > kLlmDenseMaxTokenCount) { + return {}; + } + + if (route == LlmDenseMatmulRoute::Q4K && weight->type == GGML_TYPE_Q4_K) { + match.kernel = kLlmDenseLinearQ4KF16WmmaKernel; + } else if (route == LlmDenseMatmulRoute::Q6K && weight->type == GGML_TYPE_Q6_K) { + match.kernel = kLlmDenseLinearQ6KF16WmmaKernel; + } else { + return {}; + } + + match.input = input; + match.weight = weight; + match.output = output; + match.input_size = input_size; + match.output_size = output_size; + match.token_count = token_count; + return match; +} + +} // namespace + +static void build_llm_dense_matmul_dispatch(const LlmDenseMatmulMatch & match, + DispatchMatch & dispatch_match, + size_t root_index) { + Dispatch dispatch; + dispatch.kernel = make_kernel_specialization(match.kernel); + dispatch.kernel.integer_parameters.emplace("token_count", match.token_count); + dispatch.kernel.compile_parameters.emplace("qwen3_moe.workload.token_capacity", to_config_value(match.token_count)); + dispatch.kernel.compile_parameters.emplace("qwen3_moe.dense_quantized.input_size", + to_config_value(match.input_size)); + dispatch.kernel.compile_parameters.emplace("qwen3_moe.dense_quantized.output_size", + to_config_value(match.output_size)); + dispatch.kernel.compile_parameters.emplace("qwen3_moe.dense_quantized.output_accumulation", "0"); + dispatch.bindings.push_back({ match.input->id, 0, match.input->byte_count }); + dispatch.bindings.push_back({ match.weight->id, 0, match.weight->byte_count }); + dispatch.bindings.push_back({ match.output->id, 0, match.output->byte_count }); + + dispatch_match.covered_nodes.push_back(root_index); + dispatch_match.dispatches.push_back(std::move(dispatch)); +} + +static bool match_llm_dense_q4k_dispatch(const DispatchMatchContext & context, DispatchMatch & dispatch_match) { + const LlmDenseMatmulMatch match = + match_llm_dense_matmul(context.graph, context.root_node, LlmDenseMatmulRoute::Q4K); + if (!match.matched()) { + return false; + } + build_llm_dense_matmul_dispatch(match, dispatch_match, context.root_index); + return true; +} + +static bool match_llm_dense_q6k_dispatch(const DispatchMatchContext & context, DispatchMatch & dispatch_match) { + const LlmDenseMatmulMatch match = + match_llm_dense_matmul(context.graph, context.root_node, LlmDenseMatmulRoute::Q6K); + if (!match.matched()) { + return false; + } + build_llm_dense_matmul_dispatch(match, dispatch_match, context.root_index); + return true; +} + +// Sub-4-bit IQ3_XXS weight matvec. This is the kernel that unblocks the +// unsloth "UD" GGUF files: the Loom corpus has a codebook decoder for the IQ3_XXS +// grid, so those MUL_MAT nodes can run on the NPU instead of falling back to CPU. +// It is a plain matrix-vector product (one workitem per output row), so it matches +// any 2D f32 activation / f32 output MUL_MAT whose weight is IQ3_XXS. +static bool match_llm_iq3xxs_matvec_dispatch(const DispatchMatchContext & context, DispatchMatch & dispatch_match) { + const GraphNode * node = context.root_node; + if (node == nullptr || node->op != GGML_OP_MUL_MAT || node->inputs.size() != 2) { + return false; + } + const Value * weight = graph_value(context.graph, node->inputs[0]); + const Value * input = graph_value(context.graph, node->inputs[1]); + const Value * output = graph_value(context.graph, node->output); + if (weight == nullptr || input == nullptr || output == nullptr) { + return false; + } + if (weight->type != GGML_TYPE_IQ3_XXS || input->type != GGML_TYPE_F32 || output->type != GGML_TYPE_F32) { + return false; + } + if (!is_2d(*weight) || !is_2d(*input) || !is_2d(*output) || !weight->contiguous || !input->contiguous || + !output->contiguous) { + return false; + } + const int64_t input_size = weight->ne[0]; + const int64_t output_size = weight->ne[1]; + if (input_size <= 0 || input_size % 256 != 0 || input->ne[0] != input_size || output->ne[0] != output_size) { + return false; + } + Dispatch dispatch; + dispatch.kernel = make_kernel_specialization(kLlmIq3XxsMatVecKernel); + dispatch.kernel.integer_parameters.emplace("input_size", input_size); + dispatch.kernel.integer_parameters.emplace("output_size", output_size); + dispatch.kernel.integer_parameters.emplace("token_count", input->ne[1]); + const ValueId tables = ValueId(context.next_plan_value.value + + static_cast(dispatch_match.transients.size()) + + static_cast(dispatch_match.completion_counter_requests.size())); + dispatch_match.transients.push_back({ tables, "llm.matmul.iq3xxs.tables", sizeof(kIq3XxsTables), 16 }); + dispatch_match.constant_initializations.push_back({ tables, "llm.matmul.iq3xxs.tables", 0, + std::vector(kIq3XxsTables, + kIq3XxsTables + sizeof(kIq3XxsTables)) }); + dispatch.bindings.push_back({ input->id, 0, input->byte_count }); + dispatch.bindings.push_back({ weight->id, 0, weight->byte_count }); + dispatch.bindings.push_back({ tables, 0, sizeof(kIq3XxsTables) }); + dispatch.bindings.push_back({ output->id, 0, output->byte_count }); + dispatch_match.covered_nodes.push_back(context.root_index); + dispatch_match.dispatches.push_back(std::move(dispatch)); + return true; +} + +void register_llm_matmul_dispatches(DispatchRegistryBuilder & registry) { + registry.add({ + "llm.matmul.iq3xxs_matvec_f32", + GGML_OP_MUL_MAT, + DispatchMatchKind::SingleOp, + 200, + DispatchSource::Llm, + match_llm_iq3xxs_matvec_dispatch, + }); + registry.add({ + "llm.matmul.dense_q4k_f16_wmma", + GGML_OP_MUL_MAT, + DispatchMatchKind::SingleOp, + 100, + DispatchSource::Llm, + match_llm_dense_q4k_dispatch, + }); + registry.add({ + "llm.matmul.dense_q6k_f16_wmma", + GGML_OP_MUL_MAT, + DispatchMatchKind::SingleOp, + 100, + DispatchSource::Llm, + match_llm_dense_q6k_dispatch, + }); +} + +} // namespace ggml::hrx diff --git a/ggml/src/ggml-hrx/dispatch_registration/dispatch-llm-matmul.h b/ggml/src/ggml-hrx/dispatch_registration/dispatch-llm-matmul.h new file mode 100644 index 000000000000..65136aaff293 --- /dev/null +++ b/ggml/src/ggml-hrx/dispatch_registration/dispatch-llm-matmul.h @@ -0,0 +1,9 @@ +#pragma once + +#include "dispatch-registry.h" + +namespace ggml::hrx { + +void register_llm_matmul_dispatches(DispatchRegistryBuilder & registry); + +} // namespace ggml::hrx diff --git a/ggml/src/ggml-hrx/dispatch_registration/dispatch-llm-profiles.h b/ggml/src/ggml-hrx/dispatch_registration/dispatch-llm-profiles.h new file mode 100644 index 000000000000..ef9a06f9f8b0 --- /dev/null +++ b/ggml/src/ggml-hrx/dispatch_registration/dispatch-llm-profiles.h @@ -0,0 +1,36 @@ +#pragma once + +#include + +namespace ggml::hrx { + +struct LlmMoeDispatchProfile { + const char * name = ""; + int64_t hidden_size = 0; + int64_t expert_hidden_size = 0; + int64_t expert_count = 0; + int64_t route_count = 0; + int64_t max_token_count = 0; + float rms_norm_epsilon = 0.0f; +}; + +constexpr LlmMoeDispatchProfile kLlmMoeQwen30BDispatchProfile = { + "qwen30b", 2048, 768, 128, 8, 2048, 0.000001f, +}; + +static constexpr const LlmMoeDispatchProfile & kActiveLlmMoeDispatchProfile = kLlmMoeQwen30BDispatchProfile; +static constexpr const LlmMoeDispatchProfile & kQwen30BMoeDispatchProfile = kLlmMoeQwen30BDispatchProfile; + +constexpr bool is_llm_supported_query_length(const LlmMoeDispatchProfile & profile, int64_t query_length) { + return query_length >= 1 && query_length <= profile.max_token_count; +} + +constexpr bool is_llm_decode_query_length(int64_t query_length) { + return query_length == 1; +} + +constexpr bool is_llm_prefill_query_length(const LlmMoeDispatchProfile & profile, int64_t query_length) { + return query_length > 1 && query_length <= profile.max_token_count; +} + +} // namespace ggml::hrx diff --git a/ggml/src/ggml-hrx/dispatch_registration/dispatch-llm-shapes.h b/ggml/src/ggml-hrx/dispatch_registration/dispatch-llm-shapes.h new file mode 100644 index 000000000000..d24637a2c37b --- /dev/null +++ b/ggml/src/ggml-hrx/dispatch_registration/dispatch-llm-shapes.h @@ -0,0 +1,29 @@ +#pragma once + +#include "dispatch-llm-profiles.h" + +#include + +namespace ggml::hrx { + +constexpr bool is_llm_prefill_512_query_length(const LlmMoeDispatchProfile & profile, int64_t query_length) { + return is_llm_prefill_query_length(profile, query_length) && query_length == 512; +} + +constexpr bool is_qwen_supported_query_length(int64_t query_length) { + return is_llm_supported_query_length(kQwen30BMoeDispatchProfile, query_length); +} + +constexpr bool is_qwen_decode_query_length(int64_t query_length) { + return is_llm_decode_query_length(query_length); +} + +constexpr bool is_qwen_prefill_query_length(int64_t query_length) { + return is_llm_prefill_query_length(kQwen30BMoeDispatchProfile, query_length); +} + +constexpr bool is_qwen_prefill_512_query_length(int64_t query_length) { + return is_llm_prefill_512_query_length(kQwen30BMoeDispatchProfile, query_length); +} + +} // namespace ggml::hrx diff --git a/ggml/src/ggml-hrx/dispatch_registration/dispatch-moe-router.cpp b/ggml/src/ggml-hrx/dispatch_registration/dispatch-moe-router.cpp new file mode 100644 index 000000000000..3b88b7d26ca6 --- /dev/null +++ b/ggml/src/ggml-hrx/dispatch_registration/dispatch-moe-router.cpp @@ -0,0 +1,680 @@ +#include "dispatch-moe-router.h" + +#include "dispatch-llm-shapes.h" +#include "ggml.h" +#include "graph/graph-matcher.h" +#include "kernel-corpus/kernel-corpus-catalog-verify.h" + +#include +#include +#include +#include +#include +#include + +namespace ggml::hrx { +namespace { + +static constexpr KernelCatalogRef kQwenRouterTop8F32Kernel = + GGML_HRX_KERNEL_REF("qwen3_moe", "qwen3_moe_router_top8_f32"); +static constexpr KernelCatalogRef kQwenRouterProjectionTop8FusedDecodeF32Kernel = + GGML_HRX_KERNEL_REF("qwen3_moe", "qwen3_moe_router_projection_top8_fused_decode_f32"); +static constexpr KernelCatalogRef kQwenRouterProjectionF32FourRowWave32Kernel = + GGML_HRX_KERNEL_REF("qwen3_moe", "qwen3_moe_router_projection_f32_four_row_wave32"); +static constexpr KernelCatalogRef kQwenBuildExpertTableKernel = + GGML_HRX_KERNEL_REF("qwen3_moe", "qwen3_moe_build_expert_table"); +static constexpr KernelCatalogRef kQwenBuildExpertPartitionTableKernel = + GGML_HRX_KERNEL_REF("qwen3_moe", "qwen3_moe_build_expert_partition_table"); +static constexpr KernelCatalogRef kQwenBuildExpertTablePartitionPrefill512Kernel = + GGML_HRX_KERNEL_REF("qwen3_moe", "qwen3_moe_build_expert_table_partition_prefill_512"); + +static constexpr const LlmMoeDispatchProfile & kMoeRouterProfile = kActiveLlmMoeDispatchProfile; +static constexpr size_t kMoeRouterPlanTransientAlignment = 256; + +static const Value * graph_value(const Graph & graph, ValueId id) { + return graph.values().find(id); +} + +static bool nearly_equal(float lhs, float rhs) { + return std::fabs(lhs - rhs) <= 1.0e-12f; +} + +static const GraphNode * find_consumer_with_op(const Graph & graph, ValueId value, ggml_op op) { + const std::vector consumers = consumers_with_op_through_layout_aliases(graph, value, op); + return consumers.empty() ? nullptr : consumers.front(); +} + +static const GraphNode * find_single_consumer_with_op(const Graph & graph, ValueId value, ggml_op op) { + const std::vector consumers = consumers_with_op_through_layout_aliases(graph, value, op); + return consumers.size() == 1 ? consumers.front() : nullptr; +} + +static const GraphNode * find_consumer_with_op_and_input(const Graph & graph, + ValueId value, + ggml_op op, + ValueId input) { + for (const GraphNode * consumer : consumers_with_op_through_layout_aliases(graph, value, op)) { + if (consumer != nullptr && node_has_input_or_alias(graph, *consumer, input)) { + return consumer; + } + } + return nullptr; +} + +static bool is_shape(const Value & value, int64_t ne0, int64_t ne1, int64_t ne2, int64_t ne3) { + return value.ne[0] == ne0 && value.ne[1] == ne1 && value.ne[2] == ne2 && value.ne[3] == ne3; +} + +static bool is_2d(const Value & value) { + return value.ne[0] > 0 && value.ne[1] > 0 && value.ne[2] == 1 && value.ne[3] == 1; +} + +static bool same_shape(const Value & lhs, const Value & rhs) { + for (int i = 0; i < GGML_MAX_DIMS; ++i) { + if (lhs.ne[i] != rhs.ne[i]) { + return false; + } + } + return true; +} + +static bool is_supported_expert_count(int64_t expert_count) { + return expert_count >= 32 && expert_count <= 512 && expert_count % 32 == 0; +} + +static bool is_supported_route_count(int64_t route_count, int64_t expert_count) { + return route_count >= 1 && route_count <= 32 && route_count <= expert_count; +} + +static bool is_supported_route_stride(int64_t route_stride, int64_t route_count, int64_t expert_count) { + return route_stride >= route_count && route_stride <= expert_count; +} + +static bool is_default_scale_softmax(const GraphNode & node) { + const SoftMaxParams * params = op_params_as(node.params); + return params != nullptr && nearly_equal(params->scale, 1.0f) && nearly_equal(params->max_bias, 0.0f); +} + +static bool is_descending_argsort(const GraphNode & node) { + const ArgsortParams * params = op_params_as(node.params); + return params != nullptr && params->order == GGML_SORT_ORDER_DESC; +} + +static bool is_topk_normalization_clamp(const GraphNode & node) { + const ClampParams * params = op_params_as(node.params); + return params != nullptr && params->min >= 0.0f && params->min <= 1.0e-4f && std::isinf(params->max) && + params->max > 0.0f; +} + +static std::string to_config_value(int64_t value) { + return std::to_string(value); +} + +static size_t expert_table_size(int64_t token_count, int64_t expert_count) { + return static_cast(expert_count + expert_count * token_count) * sizeof(int32_t); +} + +static size_t partition_table_size(int64_t token_count, int64_t route_count, int64_t expert_count) { + const int64_t assignment_count = token_count * route_count; + const int64_t assignment_partition_count = (assignment_count + 31) / 32; + return static_cast(1 + assignment_partition_count + expert_count) * sizeof(int32_t); +} + +static std::string value_summary(const Graph & graph, const Value * value) { + if (value == nullptr) { + return "missing"; + } + std::ostringstream stream; + stream << value->id.value << ":" << ggml_type_name(value->type) << "[" << value->ne[0] << "," << value->ne[1] << "," + << value->ne[2] << "," << value->ne[3] << "] nb=[" << value->nb[0] << "," << value->nb[1] << "," + << value->nb[2] << "," << value->nb[3] << "]"; + const GraphNode * producer = graph.index().producer(value->id); + if (producer != nullptr) { + stream << "<-" << ggml_op_name(producer->op); + } + return stream.str(); +} + +static bool is_moe_router_candidate_root(const Graph & graph, const GraphNode * softmax_node) { + if (softmax_node == nullptr || softmax_node->op != GGML_OP_SOFT_MAX || softmax_node->inputs.size() != 1 || + !graph.has_index() || !is_default_scale_softmax(*softmax_node)) { + return false; + } + const Value * logits = graph_value(graph, softmax_node->inputs[0]); + const Value * probs = graph_value(graph, softmax_node->output); + if (logits == nullptr || probs == nullptr || logits->type != GGML_TYPE_F32 || probs->type != GGML_TYPE_F32) { + return false; + } + return find_consumer_with_op(graph, softmax_node->output, GGML_OP_RESHAPE) != nullptr && + find_consumer_with_op(graph, softmax_node->output, GGML_OP_ARGSORT) != nullptr; +} + +static void log_router_reject(Status * status, + const Graph & graph, + const GraphNode * node, + const std::string & reason) { + if (status == nullptr || !is_moe_router_candidate_root(graph, node)) { + return; + } + const Value * logits = node == nullptr || node->inputs.empty() ? nullptr : graph_value(graph, node->inputs[0]); + const Value * probs = node == nullptr ? nullptr : graph_value(graph, node->output); + status->log("MoE router top-k matcher rejected node: %s logits=%s probs=%s", reason.c_str(), + value_summary(graph, logits).c_str(), value_summary(graph, probs).c_str()); +} + +static bool append_covered_node(const DispatchMatchContext & context, const GraphNode * node, DispatchMatch & match) { + return append_covered_node_index_once(context.graph, context.covered_nodes, node, match.covered_nodes); +} + +struct RouterTop8Match { + const Value * logits = nullptr; + const Value * route_ids = nullptr; + const Value * route_weights = nullptr; + int64_t token_count = 0; + int64_t expert_count = 0; + int64_t route_count = 0; + int64_t route_stride = 0; + + bool matched() const { + return logits != nullptr && route_ids != nullptr && route_weights != nullptr && token_count > 0 && + is_supported_expert_count(expert_count) && is_supported_route_count(route_count, expert_count) && + is_supported_route_stride(route_stride, route_count, expert_count); + } +}; + +static bool supports_fused_prefill_expert_table_partition(const RouterTop8Match & router_match) { + // Matches the reference prefill recipe gate; q=1 uses decode routing paths. + return is_llm_prefill_512_query_length(kMoeRouterProfile, router_match.token_count) && + router_match.route_count == kMoeRouterProfile.route_count && + router_match.route_stride == router_match.route_count && + router_match.expert_count == kMoeRouterProfile.expert_count; +} + +static RouterTop8Match match_moe_router_top8(const Graph & graph, const GraphNode * softmax_node, Status * status) { + RouterTop8Match match; + if (softmax_node == nullptr || softmax_node->op != GGML_OP_SOFT_MAX || softmax_node->inputs.size() != 1 || + !graph.has_index() || !is_default_scale_softmax(*softmax_node)) { + return match; + } + + const Value * logits = graph_value(graph, softmax_node->inputs[0]); + const Value * probs = graph_value(graph, softmax_node->output); + if (logits == nullptr || probs == nullptr || logits->type != GGML_TYPE_F32 || probs->type != GGML_TYPE_F32 || + !same_shape(*logits, *probs) || !logits->contiguous || !probs->contiguous) { + log_router_reject(status, graph, softmax_node, "logits/probs must be same contiguous F32 shape"); + return {}; + } + const int64_t expert_count = logits->ne[0]; + const int64_t token_count = logits->ne[1]; + if (!is_shape(*logits, expert_count, token_count, 1, 1) || !is_supported_expert_count(expert_count) || + !is_llm_supported_query_length(kMoeRouterProfile, token_count)) { + log_router_reject(status, graph, softmax_node, "unsupported logits expert/token shape"); + return {}; + } + + const GraphNode * probs_reshape = find_consumer_with_op(graph, softmax_node->output, GGML_OP_RESHAPE); + const GraphNode * argsort = find_consumer_with_op(graph, softmax_node->output, GGML_OP_ARGSORT); + if (probs_reshape == nullptr || argsort == nullptr || argsort->inputs.size() != 1 || + !is_descending_argsort(*argsort)) { + log_router_reject(status, graph, softmax_node, "missing probability reshape or descending argsort"); + return {}; + } + + const Value * probs_reshaped = graph_value(graph, probs_reshape->output); + const Value * argsort_output = graph_value(graph, argsort->output); + if (probs_reshaped == nullptr || argsort_output == nullptr || probs_reshaped->type != GGML_TYPE_F32 || + argsort_output->type != GGML_TYPE_I32 || !is_shape(*probs_reshaped, 1, expert_count, token_count, 1) || + !is_shape(*argsort_output, expert_count, token_count, 1, 1)) { + log_router_reject(status, graph, softmax_node, "probability reshape or argsort output shape is incompatible"); + return {}; + } + + const GraphNode * topk_view = find_consumer_with_op(graph, argsort->output, GGML_OP_VIEW); + if (topk_view == nullptr || topk_view->inputs.size() != 1) { + return {}; + } + const Value * route_ids = graph_value(graph, topk_view->output); + const int64_t route_count = route_ids == nullptr ? 0 : route_ids->ne[0]; + if (route_ids == nullptr || route_ids->type != GGML_TYPE_I32 || + !is_shape(*route_ids, route_count, token_count, 1, 1) || !is_supported_route_count(route_count, expert_count) || + route_ids->nb[0] != sizeof(int32_t) || route_ids->nb[1] % sizeof(int32_t) != 0) { + log_router_reject(status, graph, softmax_node, "top-k route id view shape or stride is incompatible"); + return {}; + } + const int64_t route_stride = static_cast(route_ids->nb[1] / sizeof(int32_t)); + if (!is_supported_route_stride(route_stride, route_count, expert_count)) { + log_router_reject(status, graph, softmax_node, "top-k route id stride is outside supported bounds"); + return {}; + } + + const GraphNode * get_rows = + find_consumer_with_op_and_input(graph, probs_reshape->output, GGML_OP_GET_ROWS, topk_view->output); + if (get_rows == nullptr || get_rows->inputs.size() != 2) { + log_router_reject(status, graph, softmax_node, "missing GET_ROWS from reshaped probabilities and top-k ids"); + return {}; + } + const Value * selected_weights = graph_value(graph, get_rows->output); + if (selected_weights == nullptr || selected_weights->type != GGML_TYPE_F32 || + !is_shape(*selected_weights, 1, route_count, token_count, 1)) { + log_router_reject(status, graph, softmax_node, "selected route weight shape is incompatible"); + return {}; + } + + const GraphNode * weights_reshape = find_consumer_with_op(graph, get_rows->output, GGML_OP_RESHAPE); + const Value * weights_flat = weights_reshape == nullptr ? nullptr : graph_value(graph, weights_reshape->output); + if (weights_flat == nullptr || weights_flat->type != GGML_TYPE_F32 || + !is_shape(*weights_flat, route_count, token_count, 1, 1)) { + log_router_reject(status, graph, softmax_node, "flattened route weight shape is incompatible"); + return {}; + } + + const GraphNode * sum_rows = find_consumer_with_op(graph, weights_reshape->output, GGML_OP_SUM_ROWS); + const Value * sum = sum_rows == nullptr ? nullptr : graph_value(graph, sum_rows->output); + if (sum == nullptr || sum->type != GGML_TYPE_F32 || !is_shape(*sum, 1, token_count, 1, 1)) { + log_router_reject(status, graph, softmax_node, "missing SUM_ROWS over selected route weights"); + return {}; + } + + const GraphNode * clamp = find_consumer_with_op(graph, sum_rows->output, GGML_OP_CLAMP); + const Value * clamped_sum = clamp == nullptr ? nullptr : graph_value(graph, clamp->output); + if (clamped_sum == nullptr || clamped_sum->type != GGML_TYPE_F32 || !is_shape(*clamped_sum, 1, token_count, 1, 1) || + !is_topk_normalization_clamp(*clamp)) { + log_router_reject(status, graph, softmax_node, "missing supported CLAMP on selected route weight sum"); + return {}; + } + + const GraphNode * div = find_consumer_with_op_and_input(graph, weights_reshape->output, GGML_OP_DIV, clamp->output); + const Value * normalized = div == nullptr ? nullptr : graph_value(graph, div->output); + if (normalized == nullptr || normalized->type != GGML_TYPE_F32 || + !is_shape(*normalized, route_count, token_count, 1, 1)) { + log_router_reject(status, graph, softmax_node, "missing DIV normalization for selected route weights"); + return {}; + } + + const GraphNode * output_reshape = find_consumer_with_op(graph, div->output, GGML_OP_RESHAPE); + const Value * route_weights = output_reshape == nullptr ? nullptr : graph_value(graph, output_reshape->output); + if (route_weights == nullptr || route_weights->type != GGML_TYPE_F32 || + !is_shape(*route_weights, 1, route_count, token_count, 1) || !route_weights->contiguous) { + log_router_reject(status, graph, softmax_node, "route weight output shape is incompatible"); + return {}; + } + + match.logits = logits; + match.route_ids = route_ids; + match.route_weights = route_weights; + match.token_count = token_count; + match.expert_count = expert_count; + match.route_count = route_count; + match.route_stride = route_stride; + return match; +} + +static bool append_moe_router_top8_coverage(const DispatchMatchContext & context, DispatchMatch & match) { + const GraphNode * softmax = context.root_node; + const GraphNode * probs_reshape = find_consumer_with_op(context.graph, softmax->output, GGML_OP_RESHAPE); + const GraphNode * argsort = find_consumer_with_op(context.graph, softmax->output, GGML_OP_ARGSORT); + const GraphNode * topk_view = + argsort == nullptr ? nullptr : find_consumer_with_op(context.graph, argsort->output, GGML_OP_VIEW); + const GraphNode * get_rows = + probs_reshape == nullptr || topk_view == nullptr ? + nullptr : + find_consumer_with_op_and_input(context.graph, probs_reshape->output, GGML_OP_GET_ROWS, topk_view->output); + const GraphNode * weights_reshape = + get_rows == nullptr ? nullptr : find_consumer_with_op(context.graph, get_rows->output, GGML_OP_RESHAPE); + const GraphNode * sum_rows = weights_reshape == nullptr ? + nullptr : + find_consumer_with_op(context.graph, weights_reshape->output, GGML_OP_SUM_ROWS); + const GraphNode * clamp = + sum_rows == nullptr ? nullptr : find_consumer_with_op(context.graph, sum_rows->output, GGML_OP_CLAMP); + const GraphNode * div = + weights_reshape == nullptr || clamp == nullptr ? + nullptr : + find_consumer_with_op_and_input(context.graph, weights_reshape->output, GGML_OP_DIV, clamp->output); + const GraphNode * output_reshape = + div == nullptr ? nullptr : find_consumer_with_op(context.graph, div->output, GGML_OP_RESHAPE); + + return append_covered_node(context, softmax, match) && append_covered_node(context, probs_reshape, match) && + append_covered_node(context, argsort, match) && append_covered_node(context, topk_view, match) && + append_covered_node(context, get_rows, match) && append_covered_node(context, weights_reshape, match) && + append_covered_node(context, sum_rows, match) && append_covered_node(context, clamp, match) && + append_covered_node(context, div, match) && append_covered_node(context, output_reshape, match); +} + +static bool append_moe_router_top8_coverage_from_softmax(const DispatchMatchContext & context, + const GraphNode * softmax, + DispatchMatch & match) { + if (softmax == nullptr) { + return false; + } + DispatchMatchContext softmax_context = context; + softmax_context.root_node = softmax; + if (!context.graph.index().node_index(softmax, softmax_context.root_index)) { + return false; + } + return append_moe_router_top8_coverage(softmax_context, match); +} + +static void add_routed_gate_up_compile_parameters(Dispatch & dispatch, const RouterTop8Match & router_match) { + dispatch.kernel.compile_parameters.emplace("qwen3_moe.routed_gate_up.expert_count", + to_config_value(router_match.expert_count)); + dispatch.kernel.compile_parameters.emplace("qwen3_moe.routed_gate_up.route_count", + to_config_value(router_match.route_count)); +} + +struct RouterProjectionTop8Match { + const GraphNode * projection = nullptr; + const GraphNode * softmax = nullptr; + const Value * input = nullptr; + const Value * weight = nullptr; + const Value * logits = nullptr; + RouterTop8Match top8; + + bool matched() const { + return projection != nullptr && softmax != nullptr && input != nullptr && weight != nullptr && + logits != nullptr && top8.matched(); + } +}; + +struct RouterProjectionMatch { + const Value * input = nullptr; + const Value * weight = nullptr; + const Value * output = nullptr; + int64_t token_count = 0; + + bool matched() const { return input != nullptr && weight != nullptr && output != nullptr && token_count > 0; } +}; + +static RouterProjectionTop8Match match_moe_router_projection_top8_decode(const DispatchMatchContext & context, + Status * status) { + RouterProjectionTop8Match match; + const GraphNode * projection = context.root_node; + if (projection == nullptr || projection->op != GGML_OP_MUL_MAT || projection->inputs.size() != 2 || + !context.graph.has_index()) { + return match; + } + + const Value * weight = graph_value(context.graph, projection->inputs[0]); + const Value * input = graph_value(context.graph, projection->inputs[1]); + const Value * logits = graph_value(context.graph, projection->output); + if (weight == nullptr || input == nullptr || logits == nullptr || weight->type != GGML_TYPE_F32 || + input->type != GGML_TYPE_F32 || logits->type != GGML_TYPE_F32 || !weight->contiguous || !input->contiguous || + !logits->contiguous || !is_shape(*input, kMoeRouterProfile.hidden_size, 1, 1, 1) || + !is_shape(*weight, kMoeRouterProfile.hidden_size, kMoeRouterProfile.expert_count, 1, 1) || + !is_shape(*logits, kMoeRouterProfile.expert_count, 1, 1, 1)) { + return {}; + } + + const GraphNode * softmax = find_single_consumer_with_op(context.graph, projection->output, GGML_OP_SOFT_MAX); + RouterTop8Match top8 = match_moe_router_top8(context.graph, softmax, status); + if (!top8.matched() || top8.token_count != 1) { + return {}; + } + + match.projection = projection; + match.softmax = softmax; + match.input = input; + match.weight = weight; + match.logits = logits; + match.top8 = top8; + return match; +} + +static RouterProjectionMatch match_moe_router_projection_f32(const DispatchMatchContext & context) { + RouterProjectionMatch match; + const GraphNode * projection = context.root_node; + if (projection == nullptr || projection->op != GGML_OP_MUL_MAT || projection->inputs.size() != 2) { + return match; + } + + const Value * weight = graph_value(context.graph, projection->inputs[0]); + const Value * input = graph_value(context.graph, projection->inputs[1]); + const Value * output = graph_value(context.graph, projection->output); + if (weight == nullptr || input == nullptr || output == nullptr || !is_2d(*weight) || !is_2d(*input) || + !is_2d(*output) || !weight->contiguous || !input->contiguous || !output->contiguous || + weight->type != GGML_TYPE_F32 || input->type != GGML_TYPE_F32 || output->type != GGML_TYPE_F32) { + return {}; + } + + const int64_t input_size = weight->ne[0]; + const int64_t output_size = weight->ne[1]; + const int64_t token_count = input->ne[1]; + if (input_size != kMoeRouterProfile.hidden_size || output_size != kMoeRouterProfile.expert_count || + input->ne[0] != input_size || output->ne[0] != output_size || output->ne[1] != token_count || + !is_llm_supported_query_length(kMoeRouterProfile, token_count)) { + return {}; + } + + match.input = input; + match.weight = weight; + match.output = output; + match.token_count = token_count; + return match; +} + +} // namespace + +static bool match_moe_router_projection_f32_dispatch(const DispatchMatchContext & context, + DispatchMatch & dispatch_match) { + const RouterProjectionMatch match = match_moe_router_projection_f32(context); + if (!match.matched()) { + return false; + } + + Dispatch dispatch; + dispatch.kernel = make_kernel_specialization(kQwenRouterProjectionF32FourRowWave32Kernel); + dispatch.kernel.integer_parameters.emplace("token_count", match.token_count); + dispatch.kernel.compile_parameters.emplace("qwen3_moe.workload.token_capacity", to_config_value(match.token_count)); + dispatch.kernel.compile_parameters.emplace("qwen3_moe.model.hidden_size", + to_config_value(kMoeRouterProfile.hidden_size)); + dispatch.kernel.compile_parameters.emplace("qwen3_moe.router.expert_count", + to_config_value(kMoeRouterProfile.expert_count)); + dispatch.bindings.push_back({ match.input->id, 0, match.input->byte_count }); + dispatch.bindings.push_back({ match.weight->id, 0, match.weight->byte_count }); + dispatch.bindings.push_back({ match.output->id, 0, match.output->byte_count }); + + dispatch_match.covered_nodes.push_back(context.root_index); + dispatch_match.dispatches.push_back(std::move(dispatch)); + return true; +} + +static bool match_moe_router_projection_top8_fused_decode_dispatch(const DispatchMatchContext & context, + DispatchMatch & dispatch_match) { + const RouterProjectionTop8Match match = match_moe_router_projection_top8_decode(context, &dispatch_match.status); + if (!match.matched()) { + return false; + } + + const ValueId completion_counter_value = context.next_plan_value; + Dispatch dispatch; + dispatch.kernel = make_kernel_specialization(kQwenRouterProjectionTop8FusedDecodeF32Kernel); + dispatch.kernel.integer_parameters.emplace("token_count", match.top8.token_count); + dispatch.kernel.integer_parameters.emplace("route_id_stride", match.top8.route_stride); + dispatch.kernel.compile_parameters.emplace("qwen3_moe.model.hidden_size", + to_config_value(kMoeRouterProfile.hidden_size)); + dispatch.kernel.compile_parameters.emplace("qwen3_moe.router.expert_count", + to_config_value(match.top8.expert_count)); + dispatch.kernel.compile_parameters.emplace("qwen3_moe.router.route_count", to_config_value(match.top8.route_count)); + dispatch.kernel.compile_parameters.emplace("qwen3_moe.workload.token_capacity", + to_config_value(match.top8.token_count)); + + const size_t route_id_length = + static_cast(match.top8.token_count * match.top8.route_stride) * sizeof(int32_t); + dispatch.bindings.push_back({ match.input->id, 0, match.input->byte_count }); + dispatch.bindings.push_back({ match.weight->id, 0, match.weight->byte_count }); + dispatch.bindings.push_back({ match.logits->id, 0, match.logits->byte_count }); + dispatch.bindings.push_back({ completion_counter_value, 0, sizeof(int32_t) }); + dispatch.bindings.push_back({ match.top8.route_ids->id, 0, route_id_length }); + dispatch.bindings.push_back({ match.top8.route_weights->id, 0, match.top8.route_weights->byte_count }); + + dispatch_match.completion_counter_requests.push_back({ + completion_counter_value, + "qwen.router.decode_projection_top8_completion_counter", + 1, + }); + if (!append_covered_node(context, match.projection, dispatch_match) || + !append_moe_router_top8_coverage_from_softmax(context, match.softmax, dispatch_match)) { + return false; + } + dispatch_match.dispatches.push_back(std::move(dispatch)); + return true; +} + +static bool match_moe_router_top8_dispatch(const DispatchMatchContext & context, DispatchMatch & dispatch_match) { + const RouterTop8Match router_match = + match_moe_router_top8(context.graph, context.root_node, &dispatch_match.status); + if (!router_match.matched()) { + return false; + } + + Dispatch dispatch; + dispatch.kernel = make_kernel_specialization(kQwenRouterTop8F32Kernel); + dispatch.kernel.integer_parameters.emplace("token_count", router_match.token_count); + dispatch.kernel.integer_parameters.emplace("route_id_stride", router_match.route_stride); + dispatch.kernel.compile_parameters.emplace("qwen3_moe.router.expert_count", + to_config_value(router_match.expert_count)); + dispatch.kernel.compile_parameters.emplace("qwen3_moe.router.route_count", + to_config_value(router_match.route_count)); + dispatch.kernel.compile_parameters.emplace("qwen3_moe.workload.token_capacity", + to_config_value(router_match.token_count)); + + const size_t route_id_length = + static_cast(router_match.token_count * router_match.route_stride) * sizeof(int32_t); + dispatch.bindings.push_back({ router_match.logits->id, 0, router_match.logits->byte_count }); + dispatch.bindings.push_back({ router_match.route_ids->id, 0, route_id_length }); + dispatch.bindings.push_back({ router_match.route_weights->id, 0, router_match.route_weights->byte_count }); + + if (!append_moe_router_top8_coverage(context, dispatch_match)) { + return false; + } + dispatch_match.dispatches.push_back(std::move(dispatch)); + + const ValueId expert_table_value(context.next_plan_value.value); + const ValueId partition_table_value(context.next_plan_value.value + 1); + const ValueId completion_counter_value(context.next_plan_value.value + 2); + const size_t expert_table_bytes = expert_table_size(router_match.token_count, router_match.expert_count); + const size_t partition_table_bytes = + partition_table_size(router_match.token_count, router_match.route_count, router_match.expert_count); + const bool use_fused_prefill_expert_table_partition = supports_fused_prefill_expert_table_partition(router_match); + dispatch_match.transients.push_back( + { expert_table_value, "qwen.router.expert_table", expert_table_bytes, kMoeRouterPlanTransientAlignment }); + dispatch_match.transients.push_back({ partition_table_value, "qwen.router.partition_table", partition_table_bytes, + kMoeRouterPlanTransientAlignment }); + if (use_fused_prefill_expert_table_partition) { + dispatch_match.completion_counter_requests.push_back({ + completion_counter_value, + "qwen.router.prefill_expert_table_partition_completion_counter", + 1, + }); + } + const CommandPlanResourceMetadata routing_metadata = make_command_plan_resource_metadata(MoeRoutingResourceMetadata{ + router_match.token_count, + router_match.route_count, + router_match.route_stride, + router_match.expert_count, + }); + Status metadata_status; + if (!dispatch_match.metadata.append_generated_resource( + { + router_match.route_ids->id, + GeneratedResourceRole::MoeExpertTable, + expert_table_value, + expert_table_bytes, + routing_metadata, + }, + metadata_status) || + !dispatch_match.metadata.append_generated_resource( + { + router_match.route_ids->id, + GeneratedResourceRole::MoePartitionTable, + partition_table_value, + partition_table_bytes, + routing_metadata, + }, + metadata_status) || + !dispatch_match.metadata.append_moe_routing_bundle( + { + router_match.route_ids->id, + router_match.route_weights->id, + expert_table_value, + partition_table_value, + expert_table_bytes, + partition_table_bytes, + router_match.token_count, + router_match.route_count, + router_match.route_stride, + router_match.expert_count, + }, + metadata_status)) { + return false; + } + + if (use_fused_prefill_expert_table_partition) { + Dispatch expert_table_partition_dispatch; + expert_table_partition_dispatch.kernel = + make_kernel_specialization(kQwenBuildExpertTablePartitionPrefill512Kernel); + expert_table_partition_dispatch.kernel.integer_parameters.emplace("token_count", router_match.token_count); + expert_table_partition_dispatch.kernel.integer_parameters.emplace("route_count", router_match.route_count); + expert_table_partition_dispatch.kernel.integer_parameters.emplace("route_stride", router_match.route_stride); + expert_table_partition_dispatch.kernel.integer_parameters.emplace("expert_count", router_match.expert_count); + expert_table_partition_dispatch.bindings.push_back({ router_match.route_ids->id, 0, route_id_length }); + expert_table_partition_dispatch.bindings.push_back({ expert_table_value, 0, expert_table_bytes }); + expert_table_partition_dispatch.bindings.push_back({ partition_table_value, 0, partition_table_bytes }); + expert_table_partition_dispatch.bindings.push_back({ completion_counter_value, 0, sizeof(int32_t) }); + dispatch_match.dispatches.push_back(std::move(expert_table_partition_dispatch)); + } else { + Dispatch expert_table_dispatch; + expert_table_dispatch.kernel = make_kernel_specialization(kQwenBuildExpertTableKernel); + expert_table_dispatch.kernel.integer_parameters.emplace("token_count", router_match.token_count); + expert_table_dispatch.kernel.integer_parameters.emplace("route_count", router_match.route_count); + expert_table_dispatch.kernel.integer_parameters.emplace("route_stride", router_match.route_stride); + expert_table_dispatch.kernel.integer_parameters.emplace("expert_count", router_match.expert_count); + expert_table_dispatch.kernel.compile_parameters.emplace("qwen3_moe.workload.token_capacity", + to_config_value(router_match.token_count)); + add_routed_gate_up_compile_parameters(expert_table_dispatch, router_match); + expert_table_dispatch.bindings.push_back({ router_match.route_ids->id, 0, route_id_length }); + expert_table_dispatch.bindings.push_back({ expert_table_value, 0, expert_table_bytes }); + dispatch_match.dispatches.push_back(std::move(expert_table_dispatch)); + + Dispatch partition_table_dispatch; + partition_table_dispatch.kernel = make_kernel_specialization(kQwenBuildExpertPartitionTableKernel); + partition_table_dispatch.kernel.integer_parameters.emplace("token_count", router_match.token_count); + partition_table_dispatch.kernel.integer_parameters.emplace("route_count", router_match.route_count); + partition_table_dispatch.kernel.integer_parameters.emplace("expert_count", router_match.expert_count); + partition_table_dispatch.kernel.compile_parameters.emplace("qwen3_moe.workload.token_capacity", + to_config_value(router_match.token_count)); + add_routed_gate_up_compile_parameters(partition_table_dispatch, router_match); + partition_table_dispatch.bindings.push_back({ expert_table_value, 0, expert_table_bytes }); + partition_table_dispatch.bindings.push_back({ partition_table_value, 0, partition_table_bytes }); + dispatch_match.dispatches.push_back(std::move(partition_table_dispatch)); + } + return true; +} + +void register_moe_router_dispatches(DispatchRegistryBuilder & registry) { + registry.add({ + "llm.moe_router.projection_top8_fused_decode", + GGML_OP_MUL_MAT, + DispatchMatchKind::Fused, + 1200, + DispatchSource::Llm, + match_moe_router_projection_top8_fused_decode_dispatch, + }); + registry.add({ + "llm.moe_router.top8_f32", + GGML_OP_SOFT_MAX, + DispatchMatchKind::Fused, + 1000, + DispatchSource::Llm, + match_moe_router_top8_dispatch, + }); + registry.add({ + "llm.moe_router.projection_f32_four_row_wave32", + GGML_OP_MUL_MAT, + DispatchMatchKind::SingleOp, + 90, + DispatchSource::Llm, + match_moe_router_projection_f32_dispatch, + }); +} + +} // namespace ggml::hrx diff --git a/ggml/src/ggml-hrx/dispatch_registration/dispatch-moe-router.h b/ggml/src/ggml-hrx/dispatch_registration/dispatch-moe-router.h new file mode 100644 index 000000000000..296282072a7f --- /dev/null +++ b/ggml/src/ggml-hrx/dispatch_registration/dispatch-moe-router.h @@ -0,0 +1,9 @@ +#pragma once + +#include "dispatch-registry.h" + +namespace ggml::hrx { + +void register_moe_router_dispatches(DispatchRegistryBuilder & registry); + +} // namespace ggml::hrx diff --git a/ggml/src/ggml-hrx/dispatch_registration/dispatch-qwen-attention-postprocess.cpp b/ggml/src/ggml-hrx/dispatch_registration/dispatch-qwen-attention-postprocess.cpp new file mode 100644 index 000000000000..28a34bc01836 --- /dev/null +++ b/ggml/src/ggml-hrx/dispatch_registration/dispatch-qwen-attention-postprocess.cpp @@ -0,0 +1,992 @@ +#include "dispatch-qwen-attention-postprocess.h" + +#include "dispatch-llm-shapes.h" +#include "ggml.h" +#include "graph/graph-matcher.h" +#include "kernel-corpus/kernel-corpus-catalog-verify.h" + +#include +#include +#include +#include +#include +#include +#include + +namespace ggml::hrx { +namespace { + +static constexpr KernelCatalogRef kQwenAttentionPostprocessF32F16Kernel = + GGML_HRX_KERNEL_REF("qwen3_moe", "qwen3_moe_attention_postprocess_f32_f16"); +static constexpr KernelCatalogRef kQwenAttentionQkvPostprocessFusedDecodeKernel = + GGML_HRX_KERNEL_REF("qwen3_moe", "qwen3_moe_attention_qkv_postprocess_fused_decode"); +static constexpr KernelCatalogRef kQwenAttentionContextBaseCaptureKernel = + GGML_HRX_KERNEL_REF("qwen3_moe", "qwen_attention_context_base_capture"); +static constexpr KernelCatalogRef kQwenAttentionMetadataKernel = + GGML_HRX_KERNEL_REF("qwen3_moe", "qwen_attention_metadata"); +static constexpr int64_t kQwenAttentionHeadSize = 128; + +static const Value * graph_value(const Graph & graph, ValueId id) { + return graph.values().find(id); +} + +static bool is_shape(const Value & value, int64_t ne0, int64_t ne1, int64_t ne2, int64_t ne3) { + return value.ne[0] == ne0 && value.ne[1] == ne1 && value.ne[2] == ne2 && value.ne[3] == ne3; +} + +static bool is_2d(const Value & value) { + return value.ne[0] > 0 && value.ne[1] > 0 && value.ne[2] == 1 && value.ne[3] == 1; +} + +static size_t q8_1_x4_byte_count(int64_t token_count, int64_t hidden_size) { + if (token_count <= 0 || hidden_size <= 0) { + return 0; + } + return static_cast(token_count) * ggml_row_size(GGML_TYPE_Q8_1, hidden_size); +} + +static bool is_supported_token_count(int64_t token_count) { + return token_count >= 1 && token_count <= 2048; +} + +static bool is_supported_head_count(int64_t head_count) { + return head_count >= 1 && head_count <= 64; +} + +static bool is_qwen_rms_norm_epsilon(float eps) { + return eps >= 0.0000009f && eps <= 0.0000011f; +} + +static bool is_qwen_implicit_rope_contract(const RopeParams & params) { + return params.n_dims == kQwenAttentionHeadSize && params.mode == GGML_ROPE_TYPE_NEOX && + std::isfinite(params.freq_base) && params.freq_base > 0.0f && std::isfinite(params.freq_scale) && + params.freq_scale > 0.0f && params.ext_factor == 0.0f && params.attn_factor == 1.0f; +} + +static bool build_inverse_frequency_table(const GraphNode & rope, std::vector & data) { + const RopeParams * params = op_params_as(rope.params); + if (params == nullptr || !is_qwen_implicit_rope_contract(*params)) { + return false; + } + + data.resize(static_cast(params->n_dims / 2) * sizeof(float)); + const float theta_scale = std::pow(params->freq_base, -2.0f / static_cast(params->n_dims)); + float theta = params->freq_scale; + for (int i = 0; i < params->n_dims / 2; ++i) { + std::memcpy(data.data() + static_cast(i) * sizeof(float), &theta, sizeof(theta)); + theta *= theta_scale; + } + return true; +} + +static bool is_supported_cache_index_type(ggml_type type) { + return type == GGML_TYPE_I64; +} + +static const GraphNode * find_single_consumer_with_op(const Graph & graph, ValueId value, ggml_op op) { + const GraphNode * match = nullptr; + for (const GraphNode * consumer : graph.index().consumers(value)) { + if (consumer == nullptr || consumer->op != op) { + continue; + } + if (match != nullptr) { + return nullptr; + } + match = consumer; + } + return match; +} + +static const GraphNode * producer_with_op(const Graph & graph, ValueId value, ggml_op op) { + const GraphNode * producer = graph.index().producer(value); + return producer != nullptr && producer->op == op ? producer : nullptr; +} + +static bool append_covered_node(const DispatchMatchContext & context, const GraphNode * node, DispatchMatch & match) { + return append_covered_node_index_once(context.graph, context.covered_nodes, node, match.covered_nodes); +} + +static std::string to_config_value(int64_t value) { + return std::to_string(value); +} + +static std::string value_summary(const Graph & graph, const Value * value) { + if (value == nullptr) { + return "missing"; + } + std::ostringstream stream; + stream << value->id.value << ":" << ggml_type_name(value->type) << "[" << value->ne[0] << "," << value->ne[1] << "," + << value->ne[2] << "," << value->ne[3] << "] nb=[" << value->nb[0] << "," << value->nb[1] << "," + << value->nb[2] << "," << value->nb[3] << "]"; + const GraphNode * producer = graph.index().producer(value->id); + if (producer != nullptr) { + stream << "<-" << ggml_op_name(producer->op); + } + if (value->alias_source.value >= 0) { + stream << " alias=" << value->alias_source.value << " storage_root=" << value->storage_root.value + << " storage_offset=" << value->storage_offset; + } + return stream.str(); +} + +static std::string node_summary(const Graph & graph, const GraphNode * node) { + if (node == nullptr) { + return "missing"; + } + std::ostringstream stream; + size_t index = 0; + if (graph.index().node_index(node, index)) { + stream << index << ":"; + } + stream << ggml_op_name(node->op) << " output=" << value_summary(graph, graph_value(graph, node->output)); + return stream.str(); +} + +static std::string rope_params_summary(const GraphNode & node) { + const RopeParams * params = op_params_as(node.params); + if (params == nullptr) { + return "missing"; + } + std::ostringstream stream; + stream << "n_dims=" << params->n_dims << " mode=" << params->mode << " n_ctx_orig=" << params->n_ctx_orig + << " freq_base=" << params->freq_base << " freq_scale=" << params->freq_scale + << " ext_factor=" << params->ext_factor << " attn_factor=" << params->attn_factor + << " beta_fast=" << params->beta_fast << " beta_slow=" << params->beta_slow; + return stream.str(); +} + +static bool is_attention_postprocess_candidate_root(const Graph & graph, const GraphNode * root) { + if (root == nullptr || root->op != GGML_OP_RESHAPE || root->inputs.size() != 1 || !graph.has_index()) { + return false; + } + const GraphNode * projection = producer_with_op(graph, root->inputs[0], GGML_OP_MUL_MAT); + const Value * raw_input = graph_value(graph, root->inputs[0]); + const Value * reshaped = graph_value(graph, root->output); + return projection != nullptr && raw_input != nullptr && reshaped != nullptr && raw_input->type == GGML_TYPE_F32 && + reshaped->type == GGML_TYPE_F32 && raw_input->ne[1] > 0 && reshaped->ne[0] == kQwenAttentionHeadSize && + reshaped->ne[2] == raw_input->ne[1]; +} + +static void log_attention_reject(Status * status, + const Graph & graph, + const GraphNode * root, + const std::string & reason) { + if (status == nullptr || !is_attention_postprocess_candidate_root(graph, root)) { + return; + } + const Value * input = root->inputs.empty() ? nullptr : graph_value(graph, root->inputs[0]); + const Value * output = graph_value(graph, root->output); + status->log("qwen attention postprocess matcher rejected node: %s root=%s input=%s output=%s", reason.c_str(), + node_summary(graph, root).c_str(), value_summary(graph, input).c_str(), + value_summary(graph, output).c_str()); +} + +static bool append_postprocess_node(const DispatchMatchContext & context, + const GraphNode * root, + const char * role, + const GraphNode * node, + DispatchMatch & dispatch_match, + Status * status) { + if (append_covered_node(context, node, dispatch_match)) { + return true; + } + std::string reason = std::string("cannot cover ") + role + " node " + node_summary(context.graph, node); + log_attention_reject(status, context.graph, root, reason); + return false; +} + +struct NormRopeChain { + const GraphNode * projection_node = nullptr; + const GraphNode * reshape_node = nullptr; + const GraphNode * rms_node = nullptr; + const GraphNode * mul_node = nullptr; + const GraphNode * rope_node = nullptr; + const Value * projection_input = nullptr; + const Value * raw_input = nullptr; + const Value * reshaped = nullptr; + const Value * norm_weight = nullptr; + const Value * positions = nullptr; + const Value * inverse_freqs = nullptr; + size_t inverse_freqs_byte_count = 0; + std::vector inverse_freqs_data; + const Value * output = nullptr; + int64_t token_count = 0; + int64_t head_count = 0; + + bool matched() const { + return projection_node != nullptr && reshape_node != nullptr && rms_node != nullptr && mul_node != nullptr && + rope_node != nullptr && projection_input != nullptr && raw_input != nullptr && reshaped != nullptr && + norm_weight != nullptr && positions != nullptr && + (inverse_freqs != nullptr || !inverse_freqs_data.empty()) && inverse_freqs_byte_count > 0 && + output != nullptr && token_count > 0 && head_count > 0; + } +}; + +struct CachePublishChain { + NormRopeChain key; + const GraphNode * layout_node = nullptr; + const GraphNode * set_rows_node = nullptr; + const Value * cache_indices = nullptr; + const Value * cache = nullptr; + int64_t cache_row_count = 0; + bool key_publish_path = false; + + bool matched_key() const { + return key_publish_path && key.matched() && layout_node != nullptr && set_rows_node != nullptr && + cache_indices != nullptr && cache != nullptr && cache_row_count > 0; + } +}; + +struct ValuePublishChain { + const GraphNode * projection_node = nullptr; + const GraphNode * reshape_node = nullptr; + const GraphNode * layout_node = nullptr; + const GraphNode * set_rows_node = nullptr; + const Value * projection_input = nullptr; + const Value * raw_input = nullptr; + const Value * cache_indices = nullptr; + const Value * cache = nullptr; + int64_t token_count = 0; + int64_t head_count = 0; + int64_t cache_row_count = 0; + + bool matched() const { + return projection_node != nullptr && reshape_node != nullptr && layout_node != nullptr && + set_rows_node != nullptr && projection_input != nullptr && raw_input != nullptr && + cache_indices != nullptr && cache != nullptr && token_count > 0 && head_count > 0 && cache_row_count > 0; + } +}; + +struct FlashInputLayoutChain { + const GraphNode * query_layout = nullptr; + const GraphNode * query_permute = nullptr; + const GraphNode * key_layout = nullptr; + const GraphNode * key_permute = nullptr; + const GraphNode * value_layout = nullptr; + const GraphNode * value_permute = nullptr; + const GraphNode * flash = nullptr; + + bool matched() const { + return query_layout != nullptr && query_permute != nullptr && key_layout != nullptr && key_permute != nullptr && + value_layout != nullptr && value_permute != nullptr && flash != nullptr; + } +}; + +struct AttentionPostprocessMatch { + NormRopeChain query; + CachePublishChain key; + ValuePublishChain value; + FlashInputLayoutChain flash_layouts; + + bool matched() const { return query.matched() && key.matched_key() && value.matched(); } +}; + +static bool matching_inverse_frequencies(const NormRopeChain & lhs, const NormRopeChain & rhs) { + if (lhs.inverse_freqs != nullptr || rhs.inverse_freqs != nullptr) { + return lhs.inverse_freqs != nullptr && rhs.inverse_freqs != nullptr && + lhs.inverse_freqs->id == rhs.inverse_freqs->id; + } + return lhs.inverse_freqs_data == rhs.inverse_freqs_data; +} + +static bool has_qwen_rope_params(const GraphNode & node) { + const RopeParams * params = op_params_as(node.params); + return params != nullptr && params->n_dims == kQwenAttentionHeadSize && params->mode == GGML_ROPE_TYPE_NEOX; +} + +static bool has_qwen_rms_params(const GraphNode & node) { + const RmsNormParams * params = op_params_as(node.params); + return params != nullptr && is_qwen_rms_norm_epsilon(params->eps); +} + +static bool is_norm_weight(const Value & value) { + return value.type == GGML_TYPE_F32 && is_shape(value, kQwenAttentionHeadSize, 1, 1, 1); +} + +static bool is_inverse_frequency_table(const Value & value) { + return value.type == GGML_TYPE_F32 && is_shape(value, kQwenAttentionHeadSize / 2, 1, 1, 1); +} + +static bool match_projection_reshape(const Graph & graph, + const GraphNode * reshape, + NormRopeChain & chain, + Status * status, + const std::string & label) { + if (reshape == nullptr || reshape->op != GGML_OP_RESHAPE || reshape->inputs.size() != 1) { + log_attention_reject(status, graph, reshape, label + " projection reshape is not a single-input RESHAPE"); + return false; + } + + const GraphNode * projection = producer_with_op(graph, reshape->inputs[0], GGML_OP_MUL_MAT); + const Value * raw_input = graph_value(graph, reshape->inputs[0]); + const Value * reshaped = graph_value(graph, reshape->output); + const Value * projection_input = + projection == nullptr || projection->inputs.size() != 2 ? nullptr : graph_value(graph, projection->inputs[1]); + if (projection == nullptr || projection_input == nullptr || raw_input == nullptr || reshaped == nullptr) { + log_attention_reject(status, graph, reshape, label + " projection producer or values are missing"); + return false; + } + if (raw_input->type != GGML_TYPE_F32 || reshaped->type != GGML_TYPE_F32 || !is_2d(*raw_input) || + reshaped->ne[0] != kQwenAttentionHeadSize || reshaped->ne[3] != 1) { + log_attention_reject(status, graph, reshape, + label + " projection reshape has incompatible type, rank, or head size"); + return false; + } + + const int64_t head_count = reshaped->ne[1]; + const int64_t token_count = reshaped->ne[2]; + if (!is_supported_head_count(head_count) || !is_supported_token_count(token_count) || + raw_input->ne[0] != head_count * kQwenAttentionHeadSize || raw_input->ne[1] != token_count) { + log_attention_reject(status, graph, reshape, label + " projection reshape has unsupported head/token shape"); + return false; + } + + chain.projection_node = projection; + chain.reshape_node = reshape; + chain.projection_input = projection_input; + chain.raw_input = raw_input; + chain.reshaped = reshaped; + chain.token_count = token_count; + chain.head_count = head_count; + return true; +} + +static bool match_norm_rope_chain_from_reshape(const Graph & graph, + const GraphNode * reshape, + NormRopeChain & chain, + Status * status = nullptr, + const std::string & label = "attention") { + if (!match_projection_reshape(graph, reshape, chain, status, label)) { + return false; + } + + const GraphNode * rms = find_single_consumer_with_op(graph, chain.reshaped->id, GGML_OP_RMS_NORM); + if (rms == nullptr || rms->inputs.size() != 1 || !has_qwen_rms_params(*rms)) { + log_attention_reject(status, graph, reshape, label + " chain is missing supported RMS_NORM"); + return false; + } + + const GraphNode * mul = find_single_consumer_with_op(graph, rms->output, GGML_OP_MUL); + if (mul == nullptr || mul->inputs.size() != 2) { + log_attention_reject(status, graph, reshape, label + " chain is missing norm-weight MUL"); + return false; + } + ValueId weight_id; + if (mul->inputs[0] == rms->output) { + weight_id = mul->inputs[1]; + } else if (mul->inputs[1] == rms->output) { + weight_id = mul->inputs[0]; + } else { + log_attention_reject(status, graph, reshape, label + " norm-weight MUL does not consume RMS output"); + return false; + } + const Value * norm_weight = graph_value(graph, weight_id); + if (norm_weight == nullptr || !is_norm_weight(*norm_weight)) { + log_attention_reject(status, graph, reshape, label + " norm weight shape is incompatible"); + return false; + } + + const GraphNode * rope = find_single_consumer_with_op(graph, mul->output, GGML_OP_ROPE); + if (rope == nullptr || rope->inputs.size() < 2 || rope->inputs.size() > 3 || rope->inputs[0] != mul->output || + !has_qwen_rope_params(*rope)) { + std::string reason = label + " chain is missing supported ROPE"; + if (rope != nullptr) { + reason += " params=" + rope_params_summary(*rope); + } + log_attention_reject(status, graph, reshape, reason); + return false; + } + const Value * positions = graph_value(graph, rope->inputs[1]); + const Value * output = graph_value(graph, rope->output); + size_t inverse_freqs_byte_count = 0; + const Value * inverse_freqs = nullptr; + std::vector inverse_freqs_data; + if (rope->inputs.size() == 3) { + inverse_freqs = graph_value(graph, rope->inputs[2]); + if (inverse_freqs == nullptr || !is_inverse_frequency_table(*inverse_freqs)) { + log_attention_reject(status, graph, reshape, label + " explicit inverse-frequency table is incompatible"); + return false; + } + inverse_freqs_byte_count = inverse_freqs->byte_count; + } else if (build_inverse_frequency_table(*rope, inverse_freqs_data)) { + inverse_freqs_byte_count = inverse_freqs_data.size(); + } else { + log_attention_reject(status, graph, reshape, + label + " implicit inverse-frequency table cannot be derived from ROPE params=" + + rope_params_summary(*rope)); + return false; + } + if (positions == nullptr || output == nullptr || positions->type != GGML_TYPE_I32 || + !is_shape(*positions, chain.token_count, 1, 1, 1) || output->type != GGML_TYPE_F32 || + !is_shape(*output, kQwenAttentionHeadSize, chain.head_count, chain.token_count, 1)) { + log_attention_reject(status, graph, reshape, label + " positions or ROPE output shape is incompatible"); + return false; + } + + chain.rms_node = rms; + chain.mul_node = mul; + chain.rope_node = rope; + chain.norm_weight = norm_weight; + chain.positions = positions; + chain.inverse_freqs = inverse_freqs; + chain.inverse_freqs_byte_count = inverse_freqs_byte_count; + chain.inverse_freqs_data = std::move(inverse_freqs_data); + chain.output = output; + return true; +} + +static const GraphNode * find_cache_read_layout(const Graph & graph, const Value & cache, int64_t head_count) { + const GraphNode * match = nullptr; + for (const GraphNode * consumer : layout_alias_consumers(graph, cache.id)) { + const Value * output = graph_value(graph, consumer->output); + if (output == nullptr || output->type != GGML_TYPE_F16 || output->ne[0] != kQwenAttentionHeadSize || + output->ne[1] != head_count || output->ne[3] != 1) { + continue; + } + if (match != nullptr) { + return nullptr; + } + match = consumer; + } + return match; +} + +static int64_t cache_row_count_for_value(const Value & cache, int64_t head_count) { + if (cache.type != GGML_TYPE_F16 || head_count <= 0) { + return 0; + } + if (cache.ne[0] == kQwenAttentionHeadSize * head_count && cache.ne[2] == 1 && cache.ne[3] == 1) { + return cache.ne[1]; + } + if (cache.ne[0] == kQwenAttentionHeadSize && cache.ne[2] == head_count && cache.ne[3] == 1) { + return cache.ne[1]; + } + return 0; +} + +static bool match_key_publish_chain(const Graph & graph, const GraphNode * set_rows, CachePublishChain & chain) { + if (set_rows == nullptr || set_rows->op != GGML_OP_SET_ROWS || set_rows->inputs.size() != 3) { + return false; + } + const GraphNode * layout = graph.index().producer(set_rows->inputs[0]); + if (layout == nullptr || !is_layout_alias_node(graph, *layout) || layout->inputs.size() != 1) { + return false; + } + + const GraphNode * rope = producer_with_op(graph, layout->inputs[0], GGML_OP_ROPE); + if (rope == nullptr) { + return false; + } + const GraphNode * mul = producer_with_op(graph, rope->inputs.empty() ? ValueId() : rope->inputs[0], GGML_OP_MUL); + if (mul == nullptr || mul->inputs.size() != 2) { + return false; + } + const GraphNode * rms = nullptr; + if (mul->inputs[0] != rope->inputs[0]) { + rms = producer_with_op(graph, mul->inputs[0], GGML_OP_RMS_NORM); + } + if (rms == nullptr && mul->inputs[1] != rope->inputs[0]) { + rms = producer_with_op(graph, mul->inputs[1], GGML_OP_RMS_NORM); + } + if (rms == nullptr || rms->inputs.size() != 1) { + return false; + } + const GraphNode * reshape = producer_with_op(graph, rms->inputs[0], GGML_OP_RESHAPE); + NormRopeChain key_chain; + if (!match_norm_rope_chain_from_reshape(graph, reshape, key_chain) || key_chain.rope_node != rope) { + return false; + } + + const Value * cache_indices = graph_value(graph, set_rows->inputs[1]); + const Value * cache = graph_value(graph, set_rows->inputs[2]); + if (cache_indices == nullptr || cache == nullptr || !is_supported_cache_index_type(cache_indices->type) || + !is_shape(*cache_indices, key_chain.token_count, 1, 1, 1)) { + return false; + } + const int64_t cache_row_count = cache_row_count_for_value(*cache, key_chain.head_count); + if (cache_row_count <= 0) { + return false; + } + + chain.key = key_chain; + chain.layout_node = layout; + chain.set_rows_node = set_rows; + chain.cache_indices = cache_indices; + chain.cache = cache; + chain.cache_row_count = cache_row_count; + chain.key_publish_path = true; + return true; +} + +static bool match_value_publish_chain(const Graph & graph, const GraphNode * set_rows, ValuePublishChain & chain) { + if (set_rows == nullptr || set_rows->op != GGML_OP_SET_ROWS || set_rows->inputs.size() != 3) { + return false; + } + const GraphNode * layout = graph.index().producer(set_rows->inputs[0]); + if (layout == nullptr || !is_layout_alias_node(graph, *layout) || layout->inputs.size() != 1) { + return false; + } + + const GraphNode * reshape = producer_with_op(graph, layout->inputs[0], GGML_OP_RESHAPE); + if (reshape == nullptr) { + reshape = layout; + } + if (reshape == nullptr || reshape->op != GGML_OP_RESHAPE || reshape->inputs.size() != 1) { + return false; + } + + NormRopeChain projection_shape; + if (!match_projection_reshape(graph, reshape, projection_shape, nullptr, "value")) { + return false; + } + + const Value * cache_indices = graph_value(graph, set_rows->inputs[1]); + const Value * cache = graph_value(graph, set_rows->inputs[2]); + if (cache_indices == nullptr || cache == nullptr || !is_supported_cache_index_type(cache_indices->type) || + !is_shape(*cache_indices, projection_shape.token_count, 1, 1, 1)) { + return false; + } + const int64_t cache_row_count = cache_row_count_for_value(*cache, projection_shape.head_count); + if (cache_row_count <= 0) { + return false; + } + + chain.projection_node = projection_shape.projection_node; + chain.reshape_node = projection_shape.reshape_node; + chain.layout_node = layout; + chain.set_rows_node = set_rows; + chain.projection_input = projection_shape.projection_input; + chain.raw_input = projection_shape.raw_input; + chain.cache_indices = cache_indices; + chain.cache = cache; + chain.token_count = projection_shape.token_count; + chain.head_count = projection_shape.head_count; + chain.cache_row_count = cache_row_count; + return true; +} + +static bool same_projection_input(const NormRopeChain & lhs, const NormRopeChain & rhs) { + return lhs.projection_input != nullptr && rhs.projection_input != nullptr && + lhs.projection_input->id == rhs.projection_input->id; +} + +static bool same_projection_input(const NormRopeChain & lhs, const ValuePublishChain & rhs) { + return lhs.projection_input != nullptr && rhs.projection_input != nullptr && + lhs.projection_input->id == rhs.projection_input->id; +} + +static FlashInputLayoutChain match_flash_input_layouts(const Graph & graph, const AttentionPostprocessMatch & match) { + FlashInputLayoutChain layouts; + const GraphNode * query_layout = find_single_layout_alias_consumer(graph, match.query.output->id); + if (query_layout == nullptr) { + return layouts; + } + const GraphNode * query_permute = find_single_consumer_with_op(graph, query_layout->output, GGML_OP_PERMUTE); + if (query_permute == nullptr) { + return {}; + } + + const GraphNode * key_layout = find_cache_read_layout(graph, *match.key.cache, match.key.key.head_count); + if (key_layout == nullptr) { + return {}; + } + const GraphNode * key_permute = find_single_consumer_with_op(graph, key_layout->output, GGML_OP_PERMUTE); + if (key_permute == nullptr) { + return {}; + } + + const GraphNode * value_layout = find_cache_read_layout(graph, *match.value.cache, match.value.head_count); + if (value_layout == nullptr) { + return {}; + } + const GraphNode * value_permute = find_single_consumer_with_op(graph, value_layout->output, GGML_OP_PERMUTE); + if (value_permute == nullptr) { + return {}; + } + + const GraphNode * flash = find_single_consumer_with_op(graph, query_permute->output, GGML_OP_FLASH_ATTN_EXT); + if (flash == nullptr || flash->inputs.size() != 4 || flash->inputs[0] != query_permute->output || + flash->inputs[1] != key_permute->output || flash->inputs[2] != value_permute->output) { + return {}; + } + + layouts.query_layout = query_layout; + layouts.query_permute = query_permute; + layouts.key_layout = key_layout; + layouts.key_permute = key_permute; + layouts.value_layout = value_layout; + layouts.value_permute = value_permute; + layouts.flash = flash; + return layouts; +} + +static AttentionPostprocessMatch match_qwen_attention_postprocess(const Graph & graph, + const GraphNode * root, + Status * status) { + AttentionPostprocessMatch match; + if (root == nullptr || root->op != GGML_OP_RESHAPE || !graph.has_index()) { + return match; + } + if (!match_norm_rope_chain_from_reshape(graph, root, match.query, status, "query")) { + return {}; + } + + for (const GraphNode & node : graph.nodes()) { + if (node.op != GGML_OP_SET_ROWS) { + continue; + } + CachePublishChain key; + if (!match.key.matched_key() && match_key_publish_chain(graph, &node, key) && + same_projection_input(match.query, key.key)) { + match.key = key; + continue; + } + ValuePublishChain value; + if (!match.value.matched() && match_value_publish_chain(graph, &node, value) && + same_projection_input(match.query, value)) { + match.value = value; + } + } + + if (!match.matched()) { + if (!match.key.matched_key()) { + log_attention_reject(status, graph, root, "no matching key ROPE cache publish chain found"); + } + if (!match.value.matched()) { + log_attention_reject(status, graph, root, "no matching value cache publish chain found"); + } + return {}; + } + if (match.query.token_count != match.key.key.token_count || match.query.token_count != match.value.token_count || + match.key.key.head_count != match.value.head_count || + match.query.positions->id != match.key.key.positions->id || + !matching_inverse_frequencies(match.query, match.key.key) || + match.key.cache_row_count != match.value.cache_row_count) { + log_attention_reject(status, graph, root, "query/key/value postprocess invariants are incompatible"); + return {}; + } + match.flash_layouts = match_flash_input_layouts(graph, match); + return match; +} + +static bool append_postprocess_covered_nodes(const DispatchMatchContext & context, + const AttentionPostprocessMatch & postprocess, + DispatchMatch & dispatch_match, + Status * status) { + // TODO: move fused matcher coverage into a shared builder that records GraphNode pointers during matching and + // materializes scheduler indices once. This is constant-size today, but the explicit list will not scale well as + // Qwen fused patterns grow. + const GraphNode * root = postprocess.query.reshape_node; + if (!append_postprocess_node(context, root, "query reshape", postprocess.query.reshape_node, dispatch_match, + status) || + !append_postprocess_node(context, root, "query rms", postprocess.query.rms_node, dispatch_match, status) || + !append_postprocess_node(context, root, "query mul", postprocess.query.mul_node, dispatch_match, status) || + !append_postprocess_node(context, root, "query rope", postprocess.query.rope_node, dispatch_match, status) || + !append_postprocess_node(context, root, "key reshape", postprocess.key.key.reshape_node, dispatch_match, + status) || + !append_postprocess_node(context, root, "key rms", postprocess.key.key.rms_node, dispatch_match, status) || + !append_postprocess_node(context, root, "key mul", postprocess.key.key.mul_node, dispatch_match, status) || + !append_postprocess_node(context, root, "key rope", postprocess.key.key.rope_node, dispatch_match, status) || + !append_postprocess_node(context, root, "key layout", postprocess.key.layout_node, dispatch_match, status) || + !append_postprocess_node(context, root, "key set rows", postprocess.key.set_rows_node, dispatch_match, + status) || + !append_postprocess_node(context, root, "value reshape", postprocess.value.reshape_node, dispatch_match, + status) || + !append_postprocess_node(context, root, "value layout", postprocess.value.layout_node, dispatch_match, + status) || + !append_postprocess_node(context, root, "value set rows", postprocess.value.set_rows_node, dispatch_match, + status)) { + return false; + } + if (postprocess.flash_layouts.matched() && + (!append_postprocess_node(context, root, "flash query layout", postprocess.flash_layouts.query_layout, + dispatch_match, status) || + !append_postprocess_node(context, root, "flash query permute", postprocess.flash_layouts.query_permute, + dispatch_match, status) || + !append_postprocess_node(context, root, "flash key layout", postprocess.flash_layouts.key_layout, + dispatch_match, status) || + !append_postprocess_node(context, root, "flash key permute", postprocess.flash_layouts.key_permute, + dispatch_match, status) || + !append_postprocess_node(context, root, "flash value layout", postprocess.flash_layouts.value_layout, + dispatch_match, status) || + !append_postprocess_node(context, root, "flash value permute", postprocess.flash_layouts.value_permute, + dispatch_match, status))) { + return false; + } + return true; +} + +static bool has_attention_metadata_initialization(const CommandPlan & plan) { + for (const Dispatch & dispatch : plan.initialization_dispatches) { + if (dispatch.kernel.kernel_id == kQwenAttentionMetadataKernel.id) { + return true; + } + } + return false; +} + +static ValueId next_match_transient_value(const DispatchMatchContext & context, const DispatchMatch & dispatch_match) { + return ValueId(context.next_plan_value.value + static_cast(dispatch_match.transients.size()) + + static_cast(dispatch_match.completion_counter_requests.size())); +} + +static bool append_attention_metadata_initialization(const DispatchMatchContext & context, + const AttentionPostprocessMatch & match, + DispatchMatch & dispatch_match) { + if (has_attention_metadata_initialization(context.plan)) { + return true; + } + if (!match.flash_layouts.matched() || match.flash_layouts.flash->inputs.size() != 4) { + return true; + } + + const Value * mask = graph_value(context.graph, match.flash_layouts.flash->inputs[3]); + if (mask == nullptr || mask->type != GGML_TYPE_F16 || mask->ne[0] <= 0 || mask->ne[1] != match.query.token_count) { + return false; + } + + const ValueId control = next_match_transient_value(context, dispatch_match); + dispatch_match.transients.push_back({ control, "qwen.attention.control", sizeof(int32_t), 16 }); + + Dispatch context_capture; + context_capture.kernel = make_kernel_specialization(kQwenAttentionContextBaseCaptureKernel); + context_capture.bindings.push_back({ match.query.positions->id, 0, match.query.positions->byte_count }); + context_capture.bindings.push_back({ control, 0, sizeof(int32_t) }); + dispatch_match.initialization_dispatches.push_back(std::move(context_capture)); + + Dispatch metadata; + metadata.kernel = make_kernel_specialization(kQwenAttentionMetadataKernel); + metadata.kernel.integer_parameters.emplace("token_count", match.query.token_count); + metadata.kernel.integer_parameters.emplace("context_capacity", mask->ne[0]); + metadata.bindings.push_back({ control, 0, sizeof(int32_t) }); + metadata.bindings.push_back({ match.query.positions->id, 0, match.query.positions->byte_count }); + metadata.bindings.push_back({ match.key.cache_indices->id, 0, match.key.cache_indices->byte_count }); + metadata.bindings.push_back({ match.value.cache_indices->id, 0, match.value.cache_indices->byte_count }); + metadata.bindings.push_back({ mask->id, 0, mask->byte_count }); + dispatch_match.initialization_dispatches.push_back(std::move(metadata)); + return true; +} + +static const Value * projection_weight(const Graph & graph, const GraphNode * projection) { + return projection == nullptr || projection->inputs.size() != 2 ? nullptr : + graph_value(graph, projection->inputs[0]); +} + +static bool is_qwen_attention_projection_weight(const Value & weight, int64_t input_size, int64_t output_size) { + return (weight.type == GGML_TYPE_Q4_K || weight.type == GGML_TYPE_Q6_K) && weight.contiguous && + is_shape(weight, input_size, output_size, 1, 1); +} + +static uint32_t attention_qkv_completion_counter_count(const AttentionPostprocessMatch & match) { + return static_cast(match.query.head_count + 2 * match.key.key.head_count); +} + +} // namespace + +static bool match_qwen_attention_qkv_postprocess_fused_decode_dispatch(const DispatchMatchContext & context, + DispatchMatch & dispatch_match) { + const GraphNode * root = context.root_node; + if (root != nullptr && root->op == GGML_OP_MUL_MAT) { + root = find_single_consumer_with_op(context.graph, root->output, GGML_OP_RESHAPE); + } + const AttentionPostprocessMatch match = + match_qwen_attention_postprocess(context.graph, root, &dispatch_match.status); + if (!match.matched() || !is_qwen_decode_query_length(match.query.token_count)) { + return false; + } + if (match.query.projection_input->id != match.key.key.projection_input->id || + match.query.projection_input->id != match.value.projection_input->id) { + return false; + } + + const int64_t hidden_size = match.query.projection_input->ne[0]; + const int64_t query_size = match.query.head_count * kQwenAttentionHeadSize; + const int64_t key_value_size = match.key.key.head_count * kQwenAttentionHeadSize; + if (hidden_size != 2048 || match.query.head_count != 32 || match.key.key.head_count != 4 || + match.value.head_count != match.key.key.head_count) { + return false; + } + + const Value * query_weight = projection_weight(context.graph, match.query.projection_node); + const Value * key_weight = projection_weight(context.graph, match.key.key.projection_node); + const Value * value_weight = projection_weight(context.graph, match.value.projection_node); + if (query_weight == nullptr || key_weight == nullptr || value_weight == nullptr || + query_weight->type != GGML_TYPE_Q4_K || key_weight->type != GGML_TYPE_Q4_K || + (value_weight->type != GGML_TYPE_Q4_K && value_weight->type != GGML_TYPE_Q6_K) || + !is_qwen_attention_projection_weight(*query_weight, hidden_size, query_size) || + !is_qwen_attention_projection_weight(*key_weight, hidden_size, key_value_size) || + !is_qwen_attention_projection_weight(*value_weight, hidden_size, key_value_size)) { + return false; + } + + const size_t q8_input_bytes = q8_1_x4_byte_count(match.query.token_count, hidden_size); + const CommandPlanAlternateValue * q8_input = + find_alternate_value(context.plan, match.query.projection_input->id, GGML_TYPE_Q8_1, q8_input_bytes); + if (q8_input == nullptr) { + return false; + } + if (!append_postprocess_covered_nodes(context, match, dispatch_match, &dispatch_match.status)) { + return false; + } + if (!append_postprocess_node(context, root, "query projection", match.query.projection_node, dispatch_match, + &dispatch_match.status) || + !append_postprocess_node(context, root, "key projection", match.key.key.projection_node, dispatch_match, + &dispatch_match.status) || + !append_postprocess_node(context, root, "value projection", match.value.projection_node, dispatch_match, + &dispatch_match.status)) { + return false; + } + if (!append_attention_metadata_initialization(context, match, dispatch_match)) { + return false; + } + + const bool synthetic_inverse_frequencies = match.query.inverse_freqs == nullptr; + const ValueId inverse_frequencies_value = synthetic_inverse_frequencies ? + next_match_transient_value(context, dispatch_match) : + match.query.inverse_freqs->id; + const size_t inverse_frequencies_size = match.query.inverse_freqs_byte_count; + if (synthetic_inverse_frequencies) { + dispatch_match.transients.push_back({ inverse_frequencies_value, + "qwen.attention_qkv_decode.inverse_frequencies", inverse_frequencies_size, + 256 }); + dispatch_match.constant_initializations.push_back({ + inverse_frequencies_value, + "qwen.attention_qkv_decode.inverse_frequencies", + 0, + match.query.inverse_freqs_data, + }); + } + + const ValueId completion_counters = next_match_transient_value(context, dispatch_match); + const uint32_t completion_counter_count = attention_qkv_completion_counter_count(match); + + Dispatch dispatch; + dispatch.kernel = make_kernel_specialization(kQwenAttentionQkvPostprocessFusedDecodeKernel); + dispatch.kernel.integer_parameters.emplace("token_count", match.query.token_count); + dispatch.kernel.integer_parameters.emplace("cache_row_count", match.key.cache_row_count); + dispatch.kernel.compile_parameters.emplace("qwen3_moe.model.hidden_size", to_config_value(hidden_size)); + dispatch.kernel.compile_parameters.emplace("qwen3_moe.attention.query_size", to_config_value(query_size)); + dispatch.kernel.compile_parameters.emplace("qwen3_moe.attention.key_value_size", to_config_value(key_value_size)); + dispatch.kernel.compile_parameters.emplace("qwen3_moe.attention.value_uses_q6", + value_weight->type == GGML_TYPE_Q6_K ? "1" : "0"); + dispatch.kernel.compile_parameters.emplace("qwen3_moe.attention.head_size", + to_config_value(kQwenAttentionHeadSize)); + dispatch.kernel.compile_parameters.emplace("qwen3_moe.model.rms_epsilon", "0.000001"); + dispatch.kernel.compile_parameters.emplace("qwen3_moe.workload.token_capacity", + to_config_value(match.query.token_count)); + + dispatch.bindings.push_back({ q8_input->alternate_value, 0, q8_input->byte_count }); + dispatch.bindings.push_back({ query_weight->id, 0, query_weight->byte_count }); + dispatch.bindings.push_back({ key_weight->id, 0, key_weight->byte_count }); + dispatch.bindings.push_back({ value_weight->id, 0, value_weight->byte_count }); + dispatch.bindings.push_back({ match.query.positions->id, 0, match.query.positions->byte_count }); + dispatch.bindings.push_back({ match.key.cache_indices->id, 0, match.key.cache_indices->byte_count }); + dispatch.bindings.push_back({ match.value.cache_indices->id, 0, match.value.cache_indices->byte_count }); + dispatch.bindings.push_back({ match.query.raw_input->id, 0, match.query.raw_input->byte_count }); + dispatch.bindings.push_back({ match.key.key.raw_input->id, 0, match.key.key.raw_input->byte_count }); + dispatch.bindings.push_back({ match.value.raw_input->id, 0, match.value.raw_input->byte_count }); + dispatch.bindings.push_back({ match.query.norm_weight->id, 0, match.query.norm_weight->byte_count }); + dispatch.bindings.push_back({ match.key.key.norm_weight->id, 0, match.key.key.norm_weight->byte_count }); + dispatch.bindings.push_back({ inverse_frequencies_value, 0, inverse_frequencies_size }); + dispatch.bindings.push_back({ match.query.output->id, 0, match.query.output->byte_count }); + dispatch.bindings.push_back({ match.key.cache->id, 0, match.key.cache->byte_count }); + dispatch.bindings.push_back({ match.value.cache->id, 0, match.value.cache->byte_count }); + dispatch.bindings.push_back({ completion_counters, 0, completion_counter_count * sizeof(int32_t) }); + + dispatch_match.completion_counter_requests.push_back({ + completion_counters, + "qwen.attention_qkv_decode.completion_counters", + completion_counter_count, + }); + dispatch_match.dispatches.push_back(std::move(dispatch)); + return true; +} + +static bool match_qwen_attention_postprocess_dispatch(const DispatchMatchContext & context, + DispatchMatch & dispatch_match) { + const AttentionPostprocessMatch match = + match_qwen_attention_postprocess(context.graph, context.root_node, &dispatch_match.status); + if (!match.matched()) { + return false; + } + if (!append_postprocess_covered_nodes(context, match, dispatch_match, &dispatch_match.status)) { + return false; + } + if (!append_attention_metadata_initialization(context, match, dispatch_match)) { + return false; + } + + Dispatch dispatch; + const bool synthetic_inverse_frequencies = match.query.inverse_freqs == nullptr; + const ValueId inverse_frequencies_value = synthetic_inverse_frequencies ? + next_match_transient_value(context, dispatch_match) : + match.query.inverse_freqs->id; + const size_t inverse_frequencies_size = match.query.inverse_freqs_byte_count; + dispatch.kernel = make_kernel_specialization(kQwenAttentionPostprocessF32F16Kernel); + dispatch.kernel.integer_parameters.emplace("token_count", match.query.token_count); + dispatch.kernel.integer_parameters.emplace("cache_row_count", match.key.cache_row_count); + dispatch.kernel.compile_parameters.emplace("qwen3_moe.model.rms_epsilon", "0.000001"); + dispatch.kernel.compile_parameters.emplace("qwen3_moe.attention.head_size", + to_config_value(kQwenAttentionHeadSize)); + dispatch.kernel.compile_parameters.emplace("qwen3_moe.attention.query_size", + to_config_value(match.query.head_count * kQwenAttentionHeadSize)); + dispatch.kernel.compile_parameters.emplace("qwen3_moe.attention.key_value_size", + to_config_value(match.key.key.head_count * kQwenAttentionHeadSize)); + dispatch.kernel.compile_parameters.emplace("qwen3_moe.workload.token_capacity", + to_config_value(match.query.token_count)); + dispatch.bindings.push_back({ match.query.positions->id, 0, match.query.positions->byte_count }); + dispatch.bindings.push_back({ match.key.cache_indices->id, 0, match.key.cache_indices->byte_count }); + dispatch.bindings.push_back({ match.value.cache_indices->id, 0, match.value.cache_indices->byte_count }); + dispatch.bindings.push_back({ match.query.raw_input->id, 0, match.query.raw_input->byte_count }); + dispatch.bindings.push_back({ match.key.key.raw_input->id, 0, match.key.key.raw_input->byte_count }); + dispatch.bindings.push_back({ match.value.raw_input->id, 0, match.value.raw_input->byte_count }); + dispatch.bindings.push_back({ match.query.norm_weight->id, 0, match.query.norm_weight->byte_count }); + dispatch.bindings.push_back({ match.key.key.norm_weight->id, 0, match.key.key.norm_weight->byte_count }); + dispatch.bindings.push_back({ inverse_frequencies_value, 0, inverse_frequencies_size }); + dispatch.bindings.push_back({ match.query.output->id, 0, match.query.output->byte_count }); + dispatch.bindings.push_back({ match.key.cache->id, 0, match.key.cache->byte_count }); + dispatch.bindings.push_back({ match.value.cache->id, 0, match.value.cache->byte_count }); + + dispatch_match.dispatches.push_back(std::move(dispatch)); + if (synthetic_inverse_frequencies) { + dispatch_match.transients.push_back({ inverse_frequencies_value, + "qwen.attention_postprocess.inverse_frequencies", + inverse_frequencies_size, 256 }); + dispatch_match.constant_initializations.push_back({ + inverse_frequencies_value, + "qwen.attention_postprocess.inverse_frequencies", + 0, + match.query.inverse_freqs_data, + }); + } + return true; +} + +void register_qwen_attention_postprocess_dispatches(DispatchRegistryBuilder & registry) { + registry.add({ + "qwen.attention_qkv_postprocess_fused_decode", + GGML_OP_MUL_MAT, + DispatchMatchKind::Fused, + 300, + DispatchSource::Qwen, + match_qwen_attention_qkv_postprocess_fused_decode_dispatch, + }); + registry.add({ + "qwen.attention_qkv_postprocess_fused_decode", + GGML_OP_RESHAPE, + DispatchMatchKind::Fused, + 200, + DispatchSource::Qwen, + match_qwen_attention_qkv_postprocess_fused_decode_dispatch, + }); + registry.add({ + "qwen.attention_postprocess_f32_f16", + GGML_OP_RESHAPE, + DispatchMatchKind::Fused, + 100, + DispatchSource::Qwen, + match_qwen_attention_postprocess_dispatch, + }); +} + +} // namespace ggml::hrx diff --git a/ggml/src/ggml-hrx/dispatch_registration/dispatch-qwen-attention-postprocess.h b/ggml/src/ggml-hrx/dispatch_registration/dispatch-qwen-attention-postprocess.h new file mode 100644 index 000000000000..0c3a8b6be354 --- /dev/null +++ b/ggml/src/ggml-hrx/dispatch_registration/dispatch-qwen-attention-postprocess.h @@ -0,0 +1,9 @@ +#pragma once + +#include "dispatch-registry.h" + +namespace ggml::hrx { + +void register_qwen_attention_postprocess_dispatches(DispatchRegistryBuilder & registry); + +} // namespace ggml::hrx diff --git a/ggml/src/ggml-hrx/dispatch_registration/dispatch-qwen-flash-attention.cpp b/ggml/src/ggml-hrx/dispatch_registration/dispatch-qwen-flash-attention.cpp new file mode 100644 index 000000000000..a75509a7e88d --- /dev/null +++ b/ggml/src/ggml-hrx/dispatch_registration/dispatch-qwen-flash-attention.cpp @@ -0,0 +1,441 @@ +#include "dispatch-qwen-flash-attention.h" + +#include "dispatch-llm-shapes.h" +#include "ggml.h" +#include "graph/graph-matcher.h" +#include "kernel-corpus/kernel-corpus-catalog-verify.h" + +#include +#include +#include +#include + +namespace ggml::hrx { +namespace { + +static constexpr KernelCatalogRef kQwenFlashAttentionF32F16WmmaKernel = + GGML_HRX_KERNEL_REF("qwen3_moe", "qwen3_moe_flash_attention_f32_f16_wmma"); +static constexpr KernelCatalogRef kQwenFlashAttentionDecodeSplitNextQ8Kernel = + GGML_HRX_KERNEL_REF("qwen3_moe", "qwen3_moe_flash_attention_decode_split_f32_f16_wmma_next_q8"); +static constexpr int64_t kQwenAttentionHeadSize = 128; +static constexpr int64_t kQwenQueryHeadCount = 32; +static constexpr int64_t kQwenKeyValueHeadCount = 4; +static constexpr int64_t kQwenDecodeRowCapacity = 16; +static constexpr int64_t kQwenDecodeKvTileSize = 64; + +static const Value * graph_value(const Graph & graph, ValueId id) { + return graph.values().find(id); +} + +static bool nearly_equal(float lhs, float rhs) { + return std::fabs(lhs - rhs) <= 1.0e-6f; +} + +static bool is_supported_key_value_token_count(int64_t token_count) { + return token_count >= 1 && token_count <= 32768; +} + +static bool is_supported_decode_key_value_token_count(int64_t token_count) { + return token_count >= 1 && token_count <= 2048; +} + +static bool is_qwen_flash_decode_query_length(int64_t query_length) { + return query_length >= 1 && query_length < kQwenDecodeRowCapacity; +} + +static bool is_supported_head_count(int64_t head_count) { + return head_count >= 1 && head_count <= 64; +} + +static bool has_query_layout(const Value & value, int64_t query_head_count) { + const size_t element_size = sizeof(float); + return value.nb[0] == element_size && + value.nb[1] == static_cast(query_head_count * kQwenAttentionHeadSize) * element_size && + (value.ne[2] == 1 || value.nb[2] == static_cast(kQwenAttentionHeadSize) * element_size); +} + +static bool has_key_value_layout(const Value & value, int64_t key_value_head_count) { + const size_t element_size = sizeof(ggml_fp16_t); + return value.nb[0] == element_size && + value.nb[1] == static_cast(key_value_head_count * kQwenAttentionHeadSize) * element_size && + (value.ne[2] == 1 || value.nb[2] == static_cast(kQwenAttentionHeadSize) * element_size); +} + +static bool has_mask_layout(const Value & value, int64_t key_value_token_count) { + const size_t element_size = sizeof(ggml_fp16_t); + return value.nb[0] == element_size && value.nb[1] == static_cast(key_value_token_count) * element_size; +} + +static bool has_output_layout(const Value & value, int64_t query_head_count) { + const size_t element_size = sizeof(float); + return value.nb[0] == element_size && value.nb[1] == static_cast(kQwenAttentionHeadSize) * element_size && + value.nb[2] == static_cast(query_head_count * kQwenAttentionHeadSize) * element_size; +} + +static std::string to_config_value(int64_t value) { + return std::to_string(value); +} + +static size_t attention_mask_byte_count(int64_t query_token_count, int64_t key_value_token_count) { + if (query_token_count <= 0 || key_value_token_count <= 0) { + return 0; + } + return static_cast(query_token_count) * static_cast(key_value_token_count) * sizeof(ggml_fp16_t); +} + +static size_t q8_1_x4_byte_count(int64_t row_count, int64_t hidden_size) { + if (row_count <= 0 || hidden_size <= 0) { + return 0; + } + return static_cast(row_count) * ggml_row_size(GGML_TYPE_Q8_1, hidden_size); +} + +static int64_t ceil_div(int64_t value, int64_t divisor) { + return (value + divisor - 1) / divisor; +} + +static ValueId match_value(const DispatchMatchContext & context, const DispatchMatch & dispatch_match, int32_t offset) { + return ValueId(context.next_plan_value.value + static_cast(dispatch_match.transients.size()) + + static_cast(dispatch_match.completion_counter_requests.size()) + offset); +} + +struct QwenFlashAttentionMatch { + const Value * query = nullptr; + const Value * key = nullptr; + const Value * value = nullptr; + const Value * mask = nullptr; + const Value * output = nullptr; + const GraphNode * output_layout = nullptr; + ValueId mask_binding_value; + size_t mask_binding_bytes = 0; + int64_t query_token_count = 0; + int64_t key_value_token_count = 0; + int64_t query_head_count = 0; + int64_t key_value_head_count = 0; + + bool matched() const { + return query != nullptr && key != nullptr && value != nullptr && mask != nullptr && output != nullptr; + } +}; + +struct QwenDecodeSplitFlashAttentionMatch { + const Value * query = nullptr; + const Value * key = nullptr; + const Value * value = nullptr; + const Value * mask = nullptr; + const Value * output = nullptr; + const GraphNode * output_layout = nullptr; + int64_t query_token_count = 0; + int64_t key_value_token_count = 0; + int64_t key_value_capacity = 0; + int64_t query_head_count = 0; + int64_t key_value_head_count = 0; + + bool matched() const { + return query != nullptr && key != nullptr && value != nullptr && mask != nullptr && output != nullptr; + } +}; + +static bool has_qwen_flash_attention_params(const GraphNode & node) { + const FlashAttnExtParams * params = op_params_as(node.params); + if (params == nullptr) { + return false; + } + const float expected_scale = 1.0f / std::sqrt(static_cast(kQwenAttentionHeadSize)); + return nearly_equal(params->scale, expected_scale) && nearly_equal(params->max_bias, 0.0f) && + nearly_equal(params->logit_softcap, 0.0f) && + (params->prec == GGML_PREC_DEFAULT || params->prec == GGML_PREC_F32); +} + +static QwenFlashAttentionMatch match_qwen_flash_attention(const Graph & graph, + const CommandPlan & plan, + const GraphNode * node) { + QwenFlashAttentionMatch match; + if (node == nullptr || node->op != GGML_OP_FLASH_ATTN_EXT || node->inputs.size() != 4 || + !has_qwen_flash_attention_params(*node)) { + return match; + } + + const Value * query = graph_value(graph, node->inputs[0]); + const Value * key = graph_value(graph, node->inputs[1]); + const Value * value = graph_value(graph, node->inputs[2]); + const Value * mask = graph_value(graph, node->inputs[3]); + const Value * output = graph_value(graph, node->output); + if (query == nullptr || key == nullptr || value == nullptr || mask == nullptr || output == nullptr) { + return {}; + } + if (query->type != GGML_TYPE_F32 || key->type != GGML_TYPE_F16 || value->type != GGML_TYPE_F16 || + mask->type != GGML_TYPE_F16 || output->type != GGML_TYPE_F32) { + return {}; + } + if (query->ne[0] != kQwenAttentionHeadSize || key->ne[0] != kQwenAttentionHeadSize || + value->ne[0] != kQwenAttentionHeadSize || output->ne[0] != kQwenAttentionHeadSize) { + return {}; + } + if (query->ne[3] != 1 || key->ne[3] != 1 || value->ne[3] != 1 || output->ne[3] != 1 || mask->ne[2] != 1 || + mask->ne[3] != 1) { + return {}; + } + + const int64_t query_token_count = query->ne[1]; + const int64_t query_head_count = query->ne[2]; + const int64_t key_value_capacity = key->ne[1]; + const int64_t key_value_head_count = key->ne[2]; + if (!is_qwen_prefill_query_length(query_token_count) || key_value_capacity < query_token_count || + !is_supported_head_count(query_head_count) || !is_supported_head_count(key_value_head_count) || + query_head_count % key_value_head_count != 0) { + return {}; + } + if (value->ne[1] != key_value_capacity || value->ne[2] != key_value_head_count || + mask->ne[0] > key_value_capacity || mask->ne[1] != query_token_count || output->ne[1] != query_head_count || + output->ne[2] != query_token_count) { + return {}; + } + + int64_t key_value_token_count = key_value_capacity; + ValueId mask_binding_value = mask->id; + size_t mask_binding_bytes = mask->byte_count; + const size_t compact_mask_bytes = attention_mask_byte_count(query_token_count, query_token_count); + const auto * compact_mask = find_alternate_value(plan, mask->id, GGML_TYPE_F16, compact_mask_bytes); + const bool mask_is_compact = mask->ne[0] == query_token_count; + const bool mask_is_capacity = mask->ne[0] == key_value_capacity; + if (compact_mask != nullptr && mask_is_capacity && mask->ne[0] > query_token_count) { + key_value_token_count = query_token_count; + mask_binding_value = compact_mask->alternate_value; + mask_binding_bytes = compact_mask->byte_count; + } else if (!mask_is_compact && !mask_is_capacity) { + return {}; + } + if (!is_supported_key_value_token_count(key_value_token_count)) { + return {}; + } + + if (!has_query_layout(*query, query_head_count) || !has_key_value_layout(*key, key_value_head_count) || + !has_key_value_layout(*value, key_value_head_count) || !has_output_layout(*output, query_head_count)) { + return {}; + } + if (mask_binding_value == mask->id && !has_mask_layout(*mask, key_value_token_count)) { + return {}; + } + + match.query = query; + match.key = key; + match.value = value; + match.mask = mask; + match.output = output; + match.output_layout = find_single_layout_alias_consumer(graph, output->id); + match.mask_binding_value = mask_binding_value; + match.mask_binding_bytes = mask_binding_bytes; + match.query_token_count = query_token_count; + match.key_value_token_count = key_value_token_count; + match.query_head_count = query_head_count; + match.key_value_head_count = key_value_head_count; + return match; +} + +static QwenDecodeSplitFlashAttentionMatch match_qwen_decode_split_flash_attention(const Graph & graph, + const GraphNode * node) { + QwenDecodeSplitFlashAttentionMatch match; + if (node == nullptr || node->op != GGML_OP_FLASH_ATTN_EXT || node->inputs.size() != 4 || + !has_qwen_flash_attention_params(*node)) { + return match; + } + + const Value * query = graph_value(graph, node->inputs[0]); + const Value * key = graph_value(graph, node->inputs[1]); + const Value * value = graph_value(graph, node->inputs[2]); + const Value * mask = graph_value(graph, node->inputs[3]); + const Value * output = graph_value(graph, node->output); + if (query == nullptr || key == nullptr || value == nullptr || mask == nullptr || output == nullptr) { + return {}; + } + if (query->type != GGML_TYPE_F32 || key->type != GGML_TYPE_F16 || value->type != GGML_TYPE_F16 || + mask->type != GGML_TYPE_F16 || output->type != GGML_TYPE_F32) { + return {}; + } + if (query->ne[0] != kQwenAttentionHeadSize || key->ne[0] != kQwenAttentionHeadSize || + value->ne[0] != kQwenAttentionHeadSize || output->ne[0] != kQwenAttentionHeadSize) { + return {}; + } + if (query->ne[3] != 1 || key->ne[3] != 1 || value->ne[3] != 1 || output->ne[3] != 1 || mask->ne[2] != 1 || + mask->ne[3] != 1) { + return {}; + } + + const int64_t query_token_count = query->ne[1]; + const int64_t query_head_count = query->ne[2]; + const int64_t key_value_capacity = key->ne[1]; + const int64_t key_value_head_count = key->ne[2]; + const int64_t key_value_token_count = mask->ne[0]; + if (!is_qwen_flash_decode_query_length(query_token_count) || + !is_supported_decode_key_value_token_count(key_value_token_count) || query_head_count != kQwenQueryHeadCount || + key_value_head_count != kQwenKeyValueHeadCount || key_value_capacity < key_value_token_count || + value->ne[1] != key_value_capacity || value->ne[2] != key_value_head_count || + mask->ne[1] != query_token_count || output->ne[1] != query_head_count || output->ne[2] != query_token_count) { + return {}; + } + if (!has_query_layout(*query, query_head_count) || !has_key_value_layout(*key, key_value_head_count) || + !has_key_value_layout(*value, key_value_head_count) || !has_mask_layout(*mask, key_value_token_count) || + !has_output_layout(*output, query_head_count)) { + return {}; + } + + match.query = query; + match.key = key; + match.value = value; + match.mask = mask; + match.output = output; + match.output_layout = find_single_layout_alias_consumer(graph, output->id); + match.query_token_count = query_token_count; + match.key_value_token_count = key_value_token_count; + match.key_value_capacity = ceil_div(key_value_token_count, kQwenDecodeKvTileSize) * kQwenDecodeKvTileSize; + match.query_head_count = query_head_count; + match.key_value_head_count = key_value_head_count; + return match; +} + +} // namespace + +static bool match_qwen_decode_split_flash_attention_next_q8_dispatch(const DispatchMatchContext & context, + DispatchMatch & dispatch_match) { + const QwenDecodeSplitFlashAttentionMatch match = + match_qwen_decode_split_flash_attention(context.graph, context.root_node); + if (!match.matched()) { + return false; + } + + const int64_t key_value_block_count = ceil_div(match.key_value_capacity, kQwenDecodeKvTileSize); + const size_t partial_scalar_count = static_cast(match.key_value_head_count) * + static_cast(key_value_block_count) * + static_cast(kQwenDecodeRowCapacity); + const size_t partial_value_count = partial_scalar_count * static_cast(kQwenAttentionHeadSize); + const size_t partial_scalar_bytes = partial_scalar_count * sizeof(float); + const size_t partial_output_bytes = partial_value_count * sizeof(ggml_fp16_t); + const int64_t hidden_size = match.query_head_count * kQwenAttentionHeadSize; + const size_t q8_row_bytes = q8_1_x4_byte_count(1, hidden_size); + const size_t q8_output_bytes = q8_1_x4_byte_count(match.query_token_count, hidden_size); + if (partial_scalar_bytes == 0 || partial_output_bytes == 0 || q8_row_bytes == 0 || q8_output_bytes == 0) { + return false; + } + + const ValueId partial_max = match_value(context, dispatch_match, 0); + const ValueId partial_sum = match_value(context, dispatch_match, 1); + const ValueId partial_output = match_value(context, dispatch_match, 2); + const ValueId completion_counter = match_value(context, dispatch_match, 3); + const ValueId q8_output = match_value(context, dispatch_match, 4); + + dispatch_match.transients.push_back( + { partial_max, "qwen.decode.flash_attention.partial_max", partial_scalar_bytes, 256 }); + dispatch_match.transients.push_back( + { partial_sum, "qwen.decode.flash_attention.partial_sum", partial_scalar_bytes, 256 }); + dispatch_match.transients.push_back( + { partial_output, "qwen.decode.flash_attention.partial_output", partial_output_bytes, 256 }); + dispatch_match.transients.push_back( + { q8_output, "qwen.decode.flash_attention.next_q8_output", q8_output_bytes, 256 }); + dispatch_match.completion_counter_requests.push_back({ + completion_counter, + "qwen.decode.flash_attention.completion_counter", + static_cast(match.key_value_head_count), + }); + + Status metadata_status; + if (!dispatch_match.metadata.append_alternate_value({ match.output->id, q8_output, GGML_TYPE_Q8_1, q8_output_bytes, + "qwen.decode.flash_attention.next_q8_output" }, + metadata_status)) { + dispatch_match.status.append(metadata_status); + return false; + } + + const size_t query_row_bytes = static_cast(hidden_size) * sizeof(float); + const size_t mask_row_bytes = static_cast(match.key_value_token_count) * sizeof(ggml_fp16_t); + const size_t output_row_bytes = query_row_bytes; + for (int64_t row = 0; row < match.query_token_count; ++row) { + Dispatch dispatch; + dispatch.kernel = make_kernel_specialization(kQwenFlashAttentionDecodeSplitNextQ8Kernel); + dispatch.kernel.integer_parameters.emplace("key_value_token_count", match.key_value_token_count); + dispatch.kernel.compile_parameters.emplace("qwen3_moe.attention.query_head_count", + to_config_value(match.query_head_count)); + dispatch.kernel.compile_parameters.emplace("qwen3_moe.attention.key_value_head_count", + to_config_value(match.key_value_head_count)); + dispatch.kernel.compile_parameters.emplace("qwen3_moe.attention.key_value_token_capacity", + to_config_value(match.key_value_capacity)); + dispatch.kernel.compile_parameters.emplace("qwen3_moe.workload.token_capacity", "1"); + dispatch.bindings.push_back( + { match.query->id, static_cast(row) * match.query->nb[1], query_row_bytes }); + dispatch.bindings.push_back({ match.key->id, 0, match.key->byte_count }); + dispatch.bindings.push_back({ match.value->id, 0, match.value->byte_count }); + dispatch.bindings.push_back({ match.mask->id, static_cast(row) * match.mask->nb[1], mask_row_bytes }); + dispatch.bindings.push_back({ partial_max, 0, partial_scalar_bytes }); + dispatch.bindings.push_back({ partial_sum, 0, partial_scalar_bytes }); + dispatch.bindings.push_back({ partial_output, 0, partial_output_bytes }); + dispatch.bindings.push_back( + { completion_counter, 0, static_cast(match.key_value_head_count) * sizeof(int32_t) }); + dispatch.bindings.push_back( + { match.output->id, static_cast(row) * match.output->nb[2], output_row_bytes }); + dispatch.bindings.push_back({ q8_output, static_cast(row) * q8_row_bytes, q8_row_bytes }); + dispatch_match.dispatches.push_back(std::move(dispatch)); + } + + dispatch_match.covered_nodes.push_back(context.root_index); + if (match.output_layout != nullptr) { + if (!append_covered_node_index_once(context.graph, context.covered_nodes, match.output_layout, + dispatch_match.covered_nodes)) { + return false; + } + } + return true; +} + +static bool match_qwen_flash_attention_dispatch(const DispatchMatchContext & context, DispatchMatch & dispatch_match) { + const QwenFlashAttentionMatch match = match_qwen_flash_attention(context.graph, context.plan, context.root_node); + if (!match.matched()) { + return false; + } + + Dispatch dispatch; + dispatch.kernel = make_kernel_specialization(kQwenFlashAttentionF32F16WmmaKernel); + dispatch.kernel.integer_parameters.emplace("query_token_count", match.query_token_count); + dispatch.kernel.integer_parameters.emplace("key_value_token_count", match.key_value_token_count); + dispatch.kernel.compile_parameters.emplace("qwen3_moe.attention.query_head_count", + to_config_value(match.query_head_count)); + dispatch.kernel.compile_parameters.emplace("qwen3_moe.attention.key_value_head_count", + to_config_value(match.key_value_head_count)); + dispatch.kernel.compile_parameters.emplace("qwen3_moe.workload.token_capacity", + to_config_value(match.query_token_count)); + dispatch.bindings.push_back({ match.query->id, 0, match.query->byte_count }); + dispatch.bindings.push_back({ match.key->id, 0, match.key->byte_count }); + dispatch.bindings.push_back({ match.value->id, 0, match.value->byte_count }); + dispatch.bindings.push_back({ match.mask_binding_value, 0, match.mask_binding_bytes }); + dispatch.bindings.push_back({ match.output->id, 0, match.output->byte_count }); + + dispatch_match.covered_nodes.push_back(context.root_index); + if (match.output_layout != nullptr) { + if (!append_covered_node_index_once(context.graph, context.covered_nodes, match.output_layout, + dispatch_match.covered_nodes)) { + return false; + } + } + dispatch_match.dispatches.push_back(std::move(dispatch)); + return true; +} + +void register_qwen_flash_attention_dispatches(DispatchRegistryBuilder & registry) { + registry.add({ + "qwen.flash_attention_decode_split_next_q8", + GGML_OP_FLASH_ATTN_EXT, + DispatchMatchKind::SingleOp, + 200, + DispatchSource::Qwen, + match_qwen_decode_split_flash_attention_next_q8_dispatch, + }); + registry.add({ + "qwen.flash_attention_f32_f16_wmma", + GGML_OP_FLASH_ATTN_EXT, + DispatchMatchKind::SingleOp, + 100, + DispatchSource::Qwen, + match_qwen_flash_attention_dispatch, + }); +} + +} // namespace ggml::hrx diff --git a/ggml/src/ggml-hrx/dispatch_registration/dispatch-qwen-flash-attention.h b/ggml/src/ggml-hrx/dispatch_registration/dispatch-qwen-flash-attention.h new file mode 100644 index 000000000000..5288f3b656d2 --- /dev/null +++ b/ggml/src/ggml-hrx/dispatch_registration/dispatch-qwen-flash-attention.h @@ -0,0 +1,9 @@ +#pragma once + +#include "dispatch-registry.h" + +namespace ggml::hrx { + +void register_qwen_flash_attention_dispatches(DispatchRegistryBuilder & registry); + +} // namespace ggml::hrx diff --git a/ggml/src/ggml-hrx/dispatch_registration/dispatch-qwen-matmul.cpp b/ggml/src/ggml-hrx/dispatch_registration/dispatch-qwen-matmul.cpp new file mode 100644 index 000000000000..715751f5521a --- /dev/null +++ b/ggml/src/ggml-hrx/dispatch_registration/dispatch-qwen-matmul.cpp @@ -0,0 +1,627 @@ +#include "dispatch-qwen-matmul.h" + +#include "dispatch-llm-shapes.h" +#include "ggml.h" +#include "graph/graph-matcher.h" +#include "kernel-corpus/kernel-corpus-catalog-verify.h" + +#include +#include +#include + +namespace ggml::hrx { +namespace { + +static constexpr KernelCatalogRef kQwenDenseLinearQ4KF16WmmaKernel = + GGML_HRX_KERNEL_REF("qwen3_moe", "qwen3_moe_dense_linear_q4k_f16_wmma"); +static constexpr KernelCatalogRef kQwenDenseLinearQ6KF16WmmaKernel = + GGML_HRX_KERNEL_REF("qwen3_moe", "qwen3_moe_dense_linear_q6k_f16_wmma"); +static constexpr KernelCatalogRef kQwenDenseLinearQ4KQ8NextQ8Kernel = + GGML_HRX_KERNEL_REF("qwen3_moe", "qwen3_moe_dense_linear_q4k_q8_1_x4_next_q8"); +static constexpr KernelCatalogRef kGgmlLinearQ6KQ8_1X4Kernel = + GGML_HRX_KERNEL_REF("qwen3_moe", "ggml_linear_q6k_q8_1_x4"); + +static constexpr int64_t kQwenHiddenSize = kQwen30BMoeDispatchProfile.hidden_size; +static constexpr int64_t kQwenVocabularyCount = 151936; + +static const Value * graph_value(const Graph & graph, ValueId id) { + return graph.values().find(id); +} + +static bool is_2d(const Value & value) { + return value.ne[0] > 0 && value.ne[1] > 0 && value.ne[2] == 1 && value.ne[3] == 1; +} + +static bool is_supported_dense_input_size(int64_t input_size) { + return input_size >= 256 && input_size <= 32768 && input_size % 256 == 0; +} + +static bool is_supported_dense_output_size(int64_t output_size) { + return output_size >= 1 && output_size <= 262144; +} + +static bool is_qwen_endpoint_projection(int64_t input_size, int64_t output_size) { + return input_size == kQwenHiddenSize && output_size == kQwenVocabularyCount; +} + +static std::string to_config_value(int64_t value) { + return std::to_string(value); +} + +struct QwenMatmulMatch { + const Value * input = nullptr; + const Value * weight = nullptr; + const Value * output = nullptr; + KernelCatalogRef kernel = {}; + ValueId input_value = {}; + size_t input_bytes = 0; + int64_t input_size = 0; + int64_t output_size = 0; + int64_t token_count = 0; + bool dense = false; + + bool matched() const { + return input != nullptr && weight != nullptr && output != nullptr && kernel.id != kUncatalogedKernelId; + } +}; + +struct QwenAttentionOutputNextQ8Match { + const Value * input = nullptr; + const CommandPlanAlternateValue * input_alternate = nullptr; + const Value * weight = nullptr; + const Value * projection_output = nullptr; + const Value * residual_input = nullptr; + const Value * residual_output = nullptr; + const Value * norm_weight = nullptr; + const Value * normalized_output = nullptr; + const GraphNode * projection_get_rows = nullptr; + const GraphNode * residual_get_rows = nullptr; + const GraphNode * add_node = nullptr; + const GraphNode * rms_node = nullptr; + const GraphNode * mul_node = nullptr; + int64_t input_size = 0; + int64_t output_size = 0; + int64_t token_count = 0; + + bool matched() const { + return input != nullptr && input_alternate != nullptr && weight != nullptr && projection_output != nullptr && + residual_input != nullptr && residual_output != nullptr && norm_weight != nullptr && + normalized_output != nullptr && add_node != nullptr && rms_node != nullptr && mul_node != nullptr; + } +}; + +struct QwenAttentionOutputAccumulateMatch { + const Value * input = nullptr; + const Value * weight = nullptr; + const Value * projection_output = nullptr; + const Value * residual_input = nullptr; + const Value * residual_output = nullptr; + const GraphNode * add_node = nullptr; + int64_t input_size = 0; + int64_t output_size = 0; + int64_t token_count = 0; + + bool matched() const { + return input != nullptr && weight != nullptr && projection_output != nullptr && residual_input != nullptr && + residual_output != nullptr && add_node != nullptr; + } +}; + +static size_t q8_1_x4_byte_count(int64_t token_count, int64_t input_size) { + if (token_count <= 0 || input_size <= 0) { + return 0; + } + return static_cast(token_count) * ggml_row_size(GGML_TYPE_Q8_1, input_size); +} + +static const GraphNode * find_single_consumer_with_op(const Graph & graph, ValueId value, ggml_op op) { + const GraphNode * match = nullptr; + for (const GraphNode * consumer : graph.index().consumers(value)) { + if (consumer == nullptr || consumer->op != op) { + continue; + } + if (match != nullptr) { + return nullptr; + } + match = consumer; + } + return match; +} + +static bool value_has_no_uncovered_consumers_except(const DispatchMatchContext & context, + ValueId value, + const GraphNode * expected_consumer) { + for (const GraphNode * consumer : context.graph.index().consumers(value)) { + if (consumer == expected_consumer) { + continue; + } + size_t consumer_index = 0; + if (!context.graph.index().node_index(consumer, consumer_index) || + consumer_index >= context.covered_nodes.size() || !context.covered_nodes[consumer_index]) { + return false; + } + } + return true; +} + +static bool same_value_layout(const Value & lhs, const Value & rhs) { + if (lhs.type != rhs.type || lhs.byte_count != rhs.byte_count || lhs.element_count != rhs.element_count || + lhs.contiguous != rhs.contiguous) { + return false; + } + for (int i = 0; i < GGML_MAX_DIMS; ++i) { + if (lhs.ne[i] != rhs.ne[i] || lhs.nb[i] != rhs.nb[i]) { + return false; + } + } + return true; +} + +static const Value * get_rows_source_with_same_layout(const Graph & graph, const GraphNode * node) { + if (node == nullptr || node->op != GGML_OP_GET_ROWS || node->inputs.size() != 2) { + return nullptr; + } + const Value * source = graph_value(graph, node->inputs[0]); + const Value * output = graph_value(graph, node->output); + if (source == nullptr || output == nullptr || !same_value_layout(*source, *output)) { + return nullptr; + } + return source; +} + +static QwenMatmulMatch match_qwen_q6k_q8_matmul(const Graph & graph, const GraphNode * node, const CommandPlan & plan) { + QwenMatmulMatch match; + if (node == nullptr || node->op != GGML_OP_MUL_MAT || node->inputs.size() != 2) { + return match; + } + + const Value * weight = graph_value(graph, node->inputs[0]); + const Value * input = graph_value(graph, node->inputs[1]); + const Value * output = graph_value(graph, node->output); + if (weight == nullptr || input == nullptr || output == nullptr) { + return {}; + } + if (!is_2d(*weight) || !is_2d(*input) || !is_2d(*output)) { + return {}; + } + if (!weight->contiguous || !input->contiguous || !output->contiguous) { + return {}; + } + if (weight->type != GGML_TYPE_Q6_K || input->type != GGML_TYPE_F32 || output->type != GGML_TYPE_F32) { + return {}; + } + + const int64_t input_size = weight->ne[0]; + const int64_t output_size = weight->ne[1]; + const int64_t token_count = input->ne[1]; + if (input->ne[0] != input_size || output->ne[0] != output_size || output->ne[1] != token_count) { + return {}; + } + if (!is_qwen_decode_query_length(token_count) || input_size != kQwenHiddenSize || + output_size != kQwenVocabularyCount || !is_supported_dense_input_size(input_size) || + !is_supported_dense_output_size(output_size)) { + return {}; + } + + const size_t q8_byte_count = q8_1_x4_byte_count(token_count, input_size); + const CommandPlanAlternateValue * alternate = + find_alternate_value(graph, plan, input->id, GGML_TYPE_Q8_1, q8_byte_count); + if (alternate == nullptr) { + return {}; + } + + match.input = input; + match.weight = weight; + match.output = output; + match.input_value = alternate->alternate_value; + match.input_bytes = alternate->byte_count; + match.kernel = kGgmlLinearQ6KQ8_1X4Kernel; + match.input_size = input_size; + match.output_size = output_size; + match.token_count = token_count; + return match; +} + +static QwenMatmulMatch match_qwen_decode_endpoint_q6k_matmul(const Graph & graph, const GraphNode * node) { + QwenMatmulMatch match; + if (node == nullptr || node->op != GGML_OP_MUL_MAT || node->inputs.size() != 2) { + return match; + } + + const Value * weight = graph_value(graph, node->inputs[0]); + const Value * input = graph_value(graph, node->inputs[1]); + const Value * output = graph_value(graph, node->output); + if (weight == nullptr || input == nullptr || output == nullptr || !is_2d(*weight) || !is_2d(*input) || + !is_2d(*output) || !weight->contiguous || !input->contiguous || !output->contiguous || + weight->type != GGML_TYPE_Q6_K || input->type != GGML_TYPE_F32 || output->type != GGML_TYPE_F32) { + return {}; + } + + const int64_t input_size = weight->ne[0]; + const int64_t output_size = weight->ne[1]; + const int64_t token_count = input->ne[1]; + if (input->ne[0] != input_size || output->ne[0] != output_size || output->ne[1] != token_count || + !is_qwen_decode_query_length(token_count) || !is_qwen_endpoint_projection(input_size, output_size) || + !is_supported_dense_input_size(input_size) || !is_supported_dense_output_size(output_size)) { + return {}; + } + + match.input = input; + match.weight = weight; + match.output = output; + match.input_value = input->id; + match.input_bytes = input->byte_count; + match.kernel = kQwenDenseLinearQ6KF16WmmaKernel; + match.input_size = input_size; + match.output_size = output_size; + match.token_count = token_count; + match.dense = true; + return match; +} + +static QwenAttentionOutputNextQ8Match match_qwen_attention_output_next_q8(const DispatchMatchContext & context) { + QwenAttentionOutputNextQ8Match match; + const Graph & graph = context.graph; + const GraphNode * node = context.root_node; + if (node == nullptr || node->op != GGML_OP_MUL_MAT || node->inputs.size() != 2 || !graph.has_index()) { + return match; + } + + const Value * weight = graph_value(graph, node->inputs[0]); + const Value * input = graph_value(graph, node->inputs[1]); + const Value * projection_output = graph_value(graph, node->output); + if (weight == nullptr || input == nullptr || projection_output == nullptr || !is_2d(*weight) || !is_2d(*input) || + !is_2d(*projection_output)) { + return {}; + } + if (weight->type != GGML_TYPE_Q4_K || input->type != GGML_TYPE_F32 || projection_output->type != GGML_TYPE_F32 || + !weight->contiguous || !input->contiguous || !projection_output->contiguous) { + return {}; + } + + const int64_t input_size = weight->ne[0]; + const int64_t output_size = weight->ne[1]; + const int64_t token_count = input->ne[1]; + if (token_count != 1 || input->ne[0] != input_size || projection_output->ne[0] != output_size || + projection_output->ne[1] != token_count || input_size != 4096 || output_size != kQwenHiddenSize) { + return {}; + } + + const CommandPlanAlternateValue * input_alternate = find_alternate_value( + graph, context.plan, input->id, GGML_TYPE_Q8_1, q8_1_x4_byte_count(token_count, input_size)); + if (input_alternate == nullptr) { + return {}; + } + + const Value * selected_projection = projection_output; + const GraphNode * projection_get_rows = nullptr; + const GraphNode * add_node = find_single_consumer_with_op(graph, projection_output->id, GGML_OP_ADD); + if (add_node == nullptr) { + projection_get_rows = find_single_consumer_with_op(graph, projection_output->id, GGML_OP_GET_ROWS); + if (get_rows_source_with_same_layout(graph, projection_get_rows) != projection_output) { + return {}; + } + selected_projection = graph_value(graph, projection_get_rows->output); + add_node = find_single_consumer_with_op(graph, projection_get_rows->output, GGML_OP_ADD); + } + if (add_node == nullptr || add_node->inputs.size() != 2) { + return {}; + } + const Value * residual_input = nullptr; + const Value * selected_residual = nullptr; + const GraphNode * residual_get_rows = nullptr; + const GraphNode * residual_consumer = add_node; + if (add_node->inputs[0] == selected_projection->id) { + selected_residual = graph_value(graph, add_node->inputs[1]); + } else if (add_node->inputs[1] == selected_projection->id) { + selected_residual = graph_value(graph, add_node->inputs[0]); + } + if (selected_residual == nullptr) { + return {}; + } + const GraphNode * selected_residual_producer = graph.index().producer(selected_residual->id); + residual_input = get_rows_source_with_same_layout(graph, selected_residual_producer); + if (residual_input != nullptr) { + residual_get_rows = selected_residual_producer; + residual_consumer = residual_get_rows; + } else { + residual_input = selected_residual; + } + const Value * residual_output = graph_value(graph, add_node->output); + if (residual_input == nullptr || residual_output == nullptr || residual_input->type != GGML_TYPE_F32 || + residual_output->type != GGML_TYPE_F32 || !same_value_layout(*selected_projection, *selected_residual) || + !same_value_layout(*selected_projection, *residual_input) || + !same_value_layout(*selected_projection, *residual_output) || + !value_has_no_uncovered_consumers_except(context, residual_input->id, residual_consumer)) { + return {}; + } + + const GraphNode * rms_node = find_single_consumer_with_op(graph, residual_output->id, GGML_OP_RMS_NORM); + if (rms_node == nullptr || rms_node->inputs.size() != 1) { + return {}; + } + const Value * rms_output = graph_value(graph, rms_node->output); + const GraphNode * mul_node = find_single_consumer_with_op(graph, rms_node->output, GGML_OP_MUL); + if (rms_output == nullptr || mul_node == nullptr || mul_node->inputs.size() != 2) { + return {}; + } + + const Value * norm_weight = nullptr; + if (mul_node->inputs[0] == rms_node->output) { + norm_weight = graph_value(graph, mul_node->inputs[1]); + } else if (mul_node->inputs[1] == rms_node->output) { + norm_weight = graph_value(graph, mul_node->inputs[0]); + } + const Value * normalized_output = graph_value(graph, mul_node->output); + if (norm_weight == nullptr || normalized_output == nullptr || norm_weight->type != GGML_TYPE_F32 || + normalized_output->type != GGML_TYPE_F32 || !norm_weight->contiguous || !normalized_output->contiguous || + norm_weight->ne[0] != output_size || normalized_output->ne[0] != output_size || + normalized_output->ne[1] != token_count) { + return {}; + } + + match.input = input; + match.input_alternate = input_alternate; + match.weight = weight; + match.projection_output = projection_output; + match.residual_input = residual_input; + match.residual_output = residual_output; + match.norm_weight = norm_weight; + match.normalized_output = normalized_output; + match.projection_get_rows = projection_get_rows; + match.residual_get_rows = residual_get_rows; + match.add_node = add_node; + match.rms_node = rms_node; + match.mul_node = mul_node; + match.input_size = input_size; + match.output_size = output_size; + match.token_count = token_count; + return match; +} + +static QwenAttentionOutputAccumulateMatch match_qwen_attention_output_accumulate(const DispatchMatchContext & context) { + QwenAttentionOutputAccumulateMatch match; + const Graph & graph = context.graph; + const GraphNode * node = context.root_node; + if (node == nullptr || node->op != GGML_OP_MUL_MAT || node->inputs.size() != 2 || !graph.has_index()) { + return match; + } + + const Value * weight = graph_value(graph, node->inputs[0]); + const Value * input = graph_value(graph, node->inputs[1]); + const Value * projection_output = graph_value(graph, node->output); + if (weight == nullptr || input == nullptr || projection_output == nullptr || !is_2d(*weight) || !is_2d(*input) || + !is_2d(*projection_output)) { + return {}; + } + if (weight->type != GGML_TYPE_Q4_K || input->type != GGML_TYPE_F32 || projection_output->type != GGML_TYPE_F32 || + !weight->contiguous || !input->contiguous || !projection_output->contiguous) { + return {}; + } + + const int64_t input_size = weight->ne[0]; + const int64_t output_size = weight->ne[1]; + const int64_t token_count = input->ne[1]; + if (!is_qwen_prefill_query_length(token_count) || input->ne[0] != input_size || + projection_output->ne[0] != output_size || projection_output->ne[1] != token_count || input_size != 4096 || + output_size != kQwenHiddenSize) { + return {}; + } + + const GraphNode * add_node = find_single_consumer_with_op(graph, projection_output->id, GGML_OP_ADD); + if (add_node == nullptr || add_node->inputs.size() != 2) { + return {}; + } + const Value * residual_input = nullptr; + if (add_node->inputs[0] == projection_output->id) { + residual_input = graph_value(graph, add_node->inputs[1]); + } else if (add_node->inputs[1] == projection_output->id) { + residual_input = graph_value(graph, add_node->inputs[0]); + } + const Value * residual_output = graph_value(graph, add_node->output); + if (residual_input == nullptr || residual_output == nullptr || residual_input->type != GGML_TYPE_F32 || + residual_output->type != GGML_TYPE_F32 || !same_value_layout(*projection_output, *residual_input) || + !same_value_layout(*projection_output, *residual_output) || + !value_has_no_uncovered_consumers_except(context, residual_input->id, add_node)) { + return {}; + } + + match.input = input; + match.weight = weight; + match.projection_output = projection_output; + match.residual_input = residual_input; + match.residual_output = residual_output; + match.add_node = add_node; + match.input_size = input_size; + match.output_size = output_size; + match.token_count = token_count; + return match; +} + +} // namespace + +static void build_qwen_matmul_dispatch(const QwenMatmulMatch & match, + DispatchMatch & dispatch_match, + size_t root_index) { + Dispatch dispatch; + dispatch.kernel = make_kernel_specialization(match.kernel); + dispatch.kernel.integer_parameters.emplace("token_count", match.token_count); + if (match.kernel.id == kGgmlLinearQ6KQ8_1X4Kernel.id) { + dispatch.kernel.integer_parameters.emplace("input_size", match.input_size); + dispatch.kernel.integer_parameters.emplace("output_size", match.output_size); + dispatch.kernel.compile_parameters.emplace("ggml.linear_q6k_q8_1_x4.token_capacity", + to_config_value(match.token_count)); + dispatch.kernel.compile_parameters.emplace("ggml.linear_q6k_q8_1_x4.output_capacity", + to_config_value(match.output_size)); + } else { + dispatch.kernel.compile_parameters.emplace("qwen3_moe.workload.token_capacity", + to_config_value(match.token_count)); + } + if (match.dense) { + dispatch.kernel.compile_parameters.emplace("qwen3_moe.dense_quantized.input_size", + to_config_value(match.input_size)); + dispatch.kernel.compile_parameters.emplace("qwen3_moe.dense_quantized.output_size", + to_config_value(match.output_size)); + dispatch.kernel.compile_parameters.emplace("qwen3_moe.dense_quantized.output_accumulation", "0"); + } + dispatch.bindings.push_back({ match.input_value, 0, match.input_bytes }); + dispatch.bindings.push_back({ match.weight->id, 0, match.weight->byte_count }); + dispatch.bindings.push_back({ match.output->id, 0, match.output->byte_count }); + + dispatch_match.covered_nodes.push_back(root_index); + dispatch_match.dispatches.push_back(std::move(dispatch)); +} + +static bool match_qwen_q6k_q8_dispatch(const DispatchMatchContext & context, DispatchMatch & dispatch_match) { + const QwenMatmulMatch match = match_qwen_q6k_q8_matmul(context.graph, context.root_node, context.plan); + if (!match.matched()) { + return false; + } + build_qwen_matmul_dispatch(match, dispatch_match, context.root_index); + return true; +} + +static bool match_qwen_decode_endpoint_q6k_dispatch(const DispatchMatchContext & context, + DispatchMatch & dispatch_match) { + const QwenMatmulMatch match = match_qwen_decode_endpoint_q6k_matmul(context.graph, context.root_node); + if (!match.matched()) { + return false; + } + build_qwen_matmul_dispatch(match, dispatch_match, context.root_index); + return true; +} + +static bool match_qwen_attention_output_next_q8_dispatch(const DispatchMatchContext & context, + DispatchMatch & dispatch_match) { + const QwenAttentionOutputNextQ8Match match = match_qwen_attention_output_next_q8(context); + if (!match.matched()) { + return false; + } + + const ValueId completion_counter(context.next_plan_value.value); + const ValueId q8_output(context.next_plan_value.value + 1); + const size_t q8_output_bytes = q8_1_x4_byte_count(match.token_count, match.output_size); + + Dispatch dispatch; + dispatch.kernel = make_kernel_specialization(kQwenDenseLinearQ4KQ8NextQ8Kernel); + dispatch.kernel.integer_parameters.emplace("token_count", match.token_count); + dispatch.kernel.compile_parameters.emplace("qwen3_moe.dense_quantized.input_size", + to_config_value(match.input_size)); + dispatch.kernel.compile_parameters.emplace("qwen3_moe.dense_quantized.output_size", + to_config_value(match.output_size)); + dispatch.kernel.compile_parameters.emplace("qwen3_moe.dense_quantized.output_accumulation", "1"); + dispatch.kernel.compile_parameters.emplace("qwen3_moe.model.hidden_size", to_config_value(match.output_size)); + dispatch.kernel.compile_parameters.emplace("qwen3_moe.model.rms_epsilon", "0.000001"); + dispatch.kernel.compile_parameters.emplace("qwen3_moe.workload.token_capacity", to_config_value(match.token_count)); + dispatch.bindings.push_back({ match.input_alternate->alternate_value, 0, match.input_alternate->byte_count }); + dispatch.bindings.push_back({ match.weight->id, 0, match.weight->byte_count }); + dispatch.bindings.push_back({ match.residual_output->id, 0, match.residual_output->byte_count }); + dispatch.bindings.push_back({ match.norm_weight->id, 0, match.norm_weight->byte_count }); + dispatch.bindings.push_back({ match.normalized_output->id, 0, match.normalized_output->byte_count }); + dispatch.bindings.push_back({ completion_counter, 0, sizeof(int32_t) }); + dispatch.bindings.push_back({ q8_output, 0, q8_output_bytes }); + + dispatch_match.value_aliases.push_back({ match.residual_input->id, match.residual_output->id }); + dispatch_match.completion_counter_requests.push_back({ + completion_counter, + "qwen.decode.attention_output.completion_counter", + 1, + }); + dispatch_match.transients.push_back( + { q8_output, "qwen.decode.attention_output.next_q8_output", q8_output_bytes, 256 }); + Status metadata_status; + if (!dispatch_match.metadata.append_alternate_value( + { match.normalized_output->id, q8_output, GGML_TYPE_Q8_1, q8_output_bytes, + "qwen.decode.attention_output.next_q8_output" }, + metadata_status)) { + dispatch_match.status.append(metadata_status); + return false; + } + + if (!append_covered_node_index_once(context.graph, context.covered_nodes, context.root_node, + dispatch_match.covered_nodes) || + (match.projection_get_rows != nullptr && + !append_covered_node_index_once(context.graph, context.covered_nodes, match.projection_get_rows, + dispatch_match.covered_nodes)) || + (match.residual_get_rows != nullptr && + !append_covered_node_index_once(context.graph, context.covered_nodes, match.residual_get_rows, + dispatch_match.covered_nodes)) || + !append_covered_node_index_once(context.graph, context.covered_nodes, match.add_node, + dispatch_match.covered_nodes) || + !append_covered_node_index_once(context.graph, context.covered_nodes, match.rms_node, + dispatch_match.covered_nodes) || + !append_covered_node_index_once(context.graph, context.covered_nodes, match.mul_node, + dispatch_match.covered_nodes)) { + return false; + } + dispatch_match.dispatches.push_back(std::move(dispatch)); + return true; +} + +static bool match_qwen_attention_output_accumulate_dispatch(const DispatchMatchContext & context, + DispatchMatch & dispatch_match) { + const QwenAttentionOutputAccumulateMatch match = match_qwen_attention_output_accumulate(context); + if (!match.matched()) { + return false; + } + + Dispatch dispatch; + dispatch.kernel = make_kernel_specialization(kQwenDenseLinearQ4KF16WmmaKernel); + dispatch.kernel.integer_parameters.emplace("token_count", match.token_count); + dispatch.kernel.compile_parameters.emplace("qwen3_moe.dense_quantized.input_size", + to_config_value(match.input_size)); + dispatch.kernel.compile_parameters.emplace("qwen3_moe.dense_quantized.output_size", + to_config_value(match.output_size)); + dispatch.kernel.compile_parameters.emplace("qwen3_moe.dense_quantized.output_accumulation", "1"); + dispatch.kernel.compile_parameters.emplace("qwen3_moe.workload.token_capacity", to_config_value(match.token_count)); + dispatch.bindings.push_back({ match.input->id, 0, match.input->byte_count }); + dispatch.bindings.push_back({ match.weight->id, 0, match.weight->byte_count }); + dispatch.bindings.push_back({ match.residual_output->id, 0, match.residual_output->byte_count }); + + dispatch_match.value_aliases.push_back({ match.residual_input->id, match.residual_output->id }); + if (!append_covered_node_index_once(context.graph, context.covered_nodes, context.root_node, + dispatch_match.covered_nodes) || + !append_covered_node_index_once(context.graph, context.covered_nodes, match.add_node, + dispatch_match.covered_nodes)) { + return false; + } + dispatch_match.dispatches.push_back(std::move(dispatch)); + return true; +} + +void register_qwen_matmul_dispatches(DispatchRegistryBuilder & registry) { + registry.add({ + "qwen.matmul.attention_output_q4k_q8_1_x4_next_q8", + GGML_OP_MUL_MAT, + DispatchMatchKind::Fused, + 300, + DispatchSource::Qwen, + match_qwen_attention_output_next_q8_dispatch, + }); + registry.add({ + "qwen.matmul.attention_output_q4k_f16_accumulate", + GGML_OP_MUL_MAT, + DispatchMatchKind::Fused, + 250, + DispatchSource::Qwen, + match_qwen_attention_output_accumulate_dispatch, + }); + registry.add({ + "qwen.matmul.q6k_q8_1_x4", + GGML_OP_MUL_MAT, + DispatchMatchKind::SingleOp, + 200, + DispatchSource::Qwen, + match_qwen_q6k_q8_dispatch, + }); + registry.add({ + "qwen.matmul.decode_endpoint_q6k_f16_wmma", + GGML_OP_MUL_MAT, + DispatchMatchKind::SingleOp, + 100, + DispatchSource::Qwen, + match_qwen_decode_endpoint_q6k_dispatch, + }); +} + +} // namespace ggml::hrx diff --git a/ggml/src/ggml-hrx/dispatch_registration/dispatch-qwen-matmul.h b/ggml/src/ggml-hrx/dispatch_registration/dispatch-qwen-matmul.h new file mode 100644 index 000000000000..c2f9a6bb9783 --- /dev/null +++ b/ggml/src/ggml-hrx/dispatch_registration/dispatch-qwen-matmul.h @@ -0,0 +1,9 @@ +#pragma once + +#include "dispatch-registry.h" + +namespace ggml::hrx { + +void register_qwen_matmul_dispatches(DispatchRegistryBuilder & registry); + +} // namespace ggml::hrx diff --git a/ggml/src/ggml-hrx/dispatch_registration/dispatch-qwen-preamble.cpp b/ggml/src/ggml-hrx/dispatch_registration/dispatch-qwen-preamble.cpp new file mode 100644 index 000000000000..5f5be8d6d0c5 --- /dev/null +++ b/ggml/src/ggml-hrx/dispatch_registration/dispatch-qwen-preamble.cpp @@ -0,0 +1,124 @@ +#include "dispatch-qwen-preamble.h" + +#include "ggml.h" +#include "kernel-corpus/kernel-corpus-catalog-verify.h" + +#include +#include + +namespace ggml::hrx { +namespace { + +static constexpr KernelCatalogRef kQwenTokenEmbeddingQ4KKernel = + GGML_HRX_KERNEL_REF("qwen3_moe", "qwen_token_embedding_q4k"); + +static const Value * graph_value(const Graph & graph, ValueId id) { + return graph.values().find(id); +} + +static bool is_1d_or_2d_column(const Value & value) { + return value.ne[0] > 0 && value.ne[1] == 1 && value.ne[2] == 1 && value.ne[3] == 1; +} + +static bool is_2d(const Value & value) { + return value.ne[0] > 0 && value.ne[1] > 0 && value.ne[2] == 1 && value.ne[3] == 1; +} + +static bool is_supported_hidden_size(int64_t hidden_size) { + return hidden_size == 2048; +} + +static bool is_supported_token_count(int64_t token_count) { + return token_count >= 1 && token_count <= 2048; +} + +static bool is_supported_vocabulary_count(int64_t vocabulary_count) { + return vocabulary_count >= 1 && vocabulary_count <= 262144; +} + +struct QwenTokenEmbeddingMatch { + const Value * token_ids = nullptr; + const Value * weight = nullptr; + const Value * output = nullptr; + int64_t token_count = 0; + int64_t vocabulary_count = 0; + int64_t hidden_size = 0; + + bool matched() const { return token_ids != nullptr && weight != nullptr && output != nullptr; } +}; + +static QwenTokenEmbeddingMatch match_qwen_token_embedding(const Graph & graph, const GraphNode * node) { + QwenTokenEmbeddingMatch match; + if (node == nullptr || node->op != GGML_OP_GET_ROWS || node->inputs.size() != 2) { + return match; + } + + const Value * weight = graph_value(graph, node->inputs[0]); + const Value * token_ids = graph_value(graph, node->inputs[1]); + const Value * output = graph_value(graph, node->output); + if (weight == nullptr || token_ids == nullptr || output == nullptr) { + return {}; + } + if (weight->type != GGML_TYPE_Q4_K || token_ids->type != GGML_TYPE_I32 || output->type != GGML_TYPE_F32) { + return {}; + } + if (!weight->contiguous || !token_ids->contiguous || !output->contiguous) { + return {}; + } + if (!is_2d(*weight) || !is_1d_or_2d_column(*token_ids) || !is_2d(*output)) { + return {}; + } + + const int64_t hidden_size = weight->ne[0]; + const int64_t vocabulary_count = weight->ne[1]; + const int64_t token_count = token_ids->ne[0]; + if (output->ne[0] != hidden_size || output->ne[1] != token_count) { + return {}; + } + if (!is_supported_hidden_size(hidden_size) || !is_supported_vocabulary_count(vocabulary_count) || + !is_supported_token_count(token_count)) { + return {}; + } + + match.token_ids = token_ids; + match.weight = weight; + match.output = output; + match.token_count = token_count; + match.vocabulary_count = vocabulary_count; + match.hidden_size = hidden_size; + return match; +} + +} // namespace + +static bool match_qwen_token_embedding_dispatch(const DispatchMatchContext & context, DispatchMatch & dispatch_match) { + const QwenTokenEmbeddingMatch match = match_qwen_token_embedding(context.graph, context.root_node); + if (!match.matched()) { + return false; + } + + Dispatch dispatch; + dispatch.kernel = make_kernel_specialization(kQwenTokenEmbeddingQ4KKernel); + dispatch.kernel.integer_parameters.emplace("token_count", match.token_count); + dispatch.kernel.integer_parameters.emplace("vocabulary_count", match.vocabulary_count); + dispatch.bindings.push_back({ match.token_ids->id, 0, match.token_ids->byte_count }); + dispatch.bindings.push_back({ match.weight->id, 0, match.weight->byte_count }); + dispatch.bindings.push_back({ match.output->id, 0, match.output->byte_count }); + + dispatch_match.covered_nodes.push_back(context.root_index); + dispatch_match.dispatches.push_back(std::move(dispatch)); + return true; +} + +void register_qwen_preamble_dispatches(DispatchRegistryBuilder & registry) { + registry.add({ + "qwen.preamble.token_embedding_q4k", + GGML_OP_GET_ROWS, + DispatchMatchKind::SingleOp, + 100, + DispatchSource::Qwen, + match_qwen_token_embedding_dispatch, + }); +} + +} // namespace ggml::hrx diff --git a/ggml/src/ggml-hrx/dispatch_registration/dispatch-qwen-preamble.h b/ggml/src/ggml-hrx/dispatch_registration/dispatch-qwen-preamble.h new file mode 100644 index 000000000000..1aee68b7c3d7 --- /dev/null +++ b/ggml/src/ggml-hrx/dispatch_registration/dispatch-qwen-preamble.h @@ -0,0 +1,9 @@ +#pragma once + +#include "dispatch-registry.h" + +namespace ggml::hrx { + +void register_qwen_preamble_dispatches(DispatchRegistryBuilder & registry); + +} // namespace ggml::hrx diff --git a/ggml/src/ggml-hrx/dispatch_registration/dispatch-registry.cpp b/ggml/src/ggml-hrx/dispatch_registration/dispatch-registry.cpp new file mode 100644 index 000000000000..5df4c680dd08 --- /dev/null +++ b/ggml/src/ggml-hrx/dispatch_registration/dispatch-registry.cpp @@ -0,0 +1,159 @@ +#include "dispatch-registry.h" + +#include "dispatch-add.h" +#include "dispatch-get-rows.h" +#include "dispatch-gather-add.h" +#include "dispatch-llm-matmul.h" +#include "dispatch-moe-router.h" +#include "dispatch-qwen-attention-postprocess.h" +#include "dispatch-qwen-flash-attention.h" +#include "dispatch-qwen-matmul.h" +#include "dispatch-qwen-preamble.h" +#include "dispatch-rmsnorm.h" +#include "dispatch-routed-ffn.h" + +#include +#include + +namespace ggml::hrx { +namespace { + +static constexpr size_t kOpCount = static_cast(GGML_OP_COUNT); + +static const std::vector kEmptyRegistrations; + +static bool valid_root_op(ggml_op op) { + return op >= 0 && static_cast(op) < kOpCount; +} + +static void sort_registrations(std::vector & registrations) { + std::stable_sort( + registrations.begin(), registrations.end(), + [](const DispatchRegistration & lhs, const DispatchRegistration & rhs) { return lhs.priority > rhs.priority; }); +} + +static void register_llm_dispatches(DispatchRegistryBuilder & builder) { + register_qwen_attention_postprocess_dispatches(builder); + register_qwen_flash_attention_dispatches(builder); + register_qwen_matmul_dispatches(builder); + register_llm_matmul_dispatches(builder); + register_routed_ffn_dispatches(builder); + register_qwen_preamble_dispatches(builder); + register_qwen_rmsnorm_dispatches(builder); + register_moe_router_dispatches(builder); +} + +static DispatchRegistry build_llm_registry() { + DispatchRegistryBuilder builder; + register_add_dispatch(builder); + register_get_rows_dispatch(builder); + register_gather_add_dispatch(builder); + register_llm_dispatches(builder); + return builder.build(); +} + +} // namespace + +bool DispatchRegistry::match(const DispatchMatchContext & context, DispatchMatch & match) const { + return this->match(context, match, nullptr); +} + +bool DispatchRegistry::match(const DispatchMatchContext & context, + DispatchMatch & match, + DispatchMatchDiagnostics * diagnostics) const { + if (context.root_node == nullptr || !valid_root_op(context.root_node->op) || registrations_by_root_.empty()) { + return false; + } + if (diagnostics != nullptr) { + diagnostics->root_op = context.root_node->op; + diagnostics->attempts.clear(); + } + const std::vector & registrations = + registrations_by_root_[static_cast(context.root_node->op)].ordered; + for (const DispatchRegistration & registration : registrations) { + DispatchMatch candidate; + if (registration.matcher != nullptr && registration.matcher(context, candidate)) { + if (diagnostics != nullptr) { + diagnostics->attempts.push_back({ + registration.name != nullptr ? registration.name : "", + registration.root_op, + registration.kind, + registration.priority, + registration.source, + true, + candidate.covered_nodes, + candidate.status.errors(), + }); + } + match = std::move(candidate); + return true; + } + if (diagnostics != nullptr) { + diagnostics->attempts.push_back({ + registration.name != nullptr ? registration.name : "", + registration.root_op, + registration.kind, + registration.priority, + registration.source, + false, + candidate.covered_nodes, + candidate.status.errors(), + }); + } + match.status.append(candidate.status); + } + return false; +} + +const std::vector & DispatchRegistry::registrations_for_root(ggml_op root_op) const { + if (!valid_root_op(root_op) || registrations_by_root_.empty()) { + return kEmptyRegistrations; + } + return registrations_by_root_[static_cast(root_op)].ordered; +} + +void DispatchRegistryBuilder::add(DispatchRegistration registration) { + if (!valid_root_op(registration.root_op) || registration.matcher == nullptr) { + return; + } + if (registry_.registrations_by_root_.empty()) { + registry_.registrations_by_root_.resize(kOpCount); + } + DispatchRegistry::RegistrationGroup & group = + registry_.registrations_by_root_[static_cast(registration.root_op)]; + if (registration.kind == DispatchMatchKind::Fused) { + group.fused.push_back(std::move(registration)); + } else { + group.single_op.push_back(registration); + registry_.single_op_registrations_.push_back(std::move(registration)); + } +} + +DispatchRegistry DispatchRegistryBuilder::build() { + if (registry_.registrations_by_root_.empty()) { + registry_.registrations_by_root_.resize(kOpCount); + } + for (DispatchRegistry::RegistrationGroup & group : registry_.registrations_by_root_) { + sort_registrations(group.fused); + sort_registrations(group.single_op); + group.ordered = group.fused; + group.ordered.insert(group.ordered.end(), group.single_op.begin(), group.single_op.end()); + } + sort_registrations(registry_.single_op_registrations_); + return std::move(registry_); +} + +const DispatchRegistry * find_dispatch_registry(const DispatchTarget & target) { + static const DispatchRegistry gfx1100_registry = build_llm_registry(); + static const DispatchRegistry gfx1151_registry = build_llm_registry(); + + if (target.architecture == "gfx1100") { + return &gfx1100_registry; + } + if (target.architecture == "gfx1151") { + return &gfx1151_registry; + } + return nullptr; +} + +} // namespace ggml::hrx diff --git a/ggml/src/ggml-hrx/dispatch_registration/dispatch-registry.h b/ggml/src/ggml-hrx/dispatch_registration/dispatch-registry.h new file mode 100644 index 000000000000..bf92f9f82e1e --- /dev/null +++ b/ggml/src/ggml-hrx/dispatch_registration/dispatch-registry.h @@ -0,0 +1,117 @@ +#pragma once + +#include "dispatch/command-plan.h" +#include "ggml.h" +#include "graph/graph.h" +#include "status.h" + +#include +#include +#include + +namespace ggml::hrx { + +struct DispatchTarget { + std::string architecture; +}; + +enum class DispatchMatchKind { + Fused, + SingleOp, +}; + +enum class DispatchSource { + Common, + Llm, + Qwen, +}; + +struct DispatchMatchContext { + const Graph & graph; + const GraphNode * root_node = nullptr; + size_t root_index = 0; + const std::vector & covered_nodes; + const CommandPlan & plan; + ValueId next_plan_value; +}; + +struct DispatchValueAliasRequest { + ValueId source_value; + ValueId target_value; +}; + +struct DispatchMatch { + std::vector initialization_dispatches; + std::vector covered_nodes; + std::vector dispatches; + std::vector transients; + std::vector constant_initializations; + std::vector completion_counter_requests; + std::vector value_aliases; + CommandPlanMetadata metadata; + Status status; +}; + +using DispatchMatcher = bool (*)(const DispatchMatchContext & context, DispatchMatch & match); + +struct DispatchRegistration { + const char * name = ""; + ggml_op root_op = GGML_OP_NONE; + DispatchMatchKind kind = DispatchMatchKind::SingleOp; + int priority = 0; + DispatchSource source = DispatchSource::Common; + DispatchMatcher matcher = nullptr; +}; + +struct DispatchRegistrationAttempt { + std::string name; + ggml_op root_op = GGML_OP_NONE; + DispatchMatchKind kind = DispatchMatchKind::SingleOp; + int priority = 0; + DispatchSource source = DispatchSource::Common; + bool matched = false; + std::vector covered_nodes; + std::vector errors; +}; + +struct DispatchMatchDiagnostics { + ggml_op root_op = GGML_OP_NONE; + std::vector attempts; +}; + +class DispatchRegistry { + public: + bool match(const DispatchMatchContext & context, DispatchMatch & match) const; + bool match(const DispatchMatchContext & context, + DispatchMatch & match, + DispatchMatchDiagnostics * diagnostics) const; + + const std::vector & registrations_for_root(ggml_op root_op) const; + + const std::vector & single_op_registrations() const { return single_op_registrations_; } + + private: + friend class DispatchRegistryBuilder; + + struct RegistrationGroup { + std::vector fused; + std::vector single_op; + std::vector ordered; + }; + + std::vector registrations_by_root_; + std::vector single_op_registrations_; +}; + +class DispatchRegistryBuilder { + public: + void add(DispatchRegistration registration); + DispatchRegistry build(); + + private: + DispatchRegistry registry_; +}; + +const DispatchRegistry * find_dispatch_registry(const DispatchTarget & target); + +} // namespace ggml::hrx diff --git a/ggml/src/ggml-hrx/dispatch_registration/dispatch-rmsnorm.cpp b/ggml/src/ggml-hrx/dispatch_registration/dispatch-rmsnorm.cpp new file mode 100644 index 000000000000..a28032969d1f --- /dev/null +++ b/ggml/src/ggml-hrx/dispatch_registration/dispatch-rmsnorm.cpp @@ -0,0 +1,433 @@ +#include "dispatch-rmsnorm.h" + +#include "dispatch-llm-profiles.h" +#include "ggml.h" +#include "kernel-corpus/kernel-corpus-catalog-verify.h" + +#include +#include +#include +#include + +namespace ggml::hrx { +namespace { + +static constexpr KernelCatalogRef kQwenRmsNormF32Kernel = GGML_HRX_KERNEL_REF("qwen3_moe", "qwen3_moe_rmsnorm_f32"); +static constexpr KernelCatalogRef kQwenRmsNormF32QuantizeQ8_1X4Kernel = + GGML_HRX_KERNEL_REF("qwen3_moe", "qwen3_moe_rmsnorm_f32_quantize_q8_1_x4"); +static constexpr KernelCatalogRef kGgmlLinearQ6KQ8_1X4Kernel = + GGML_HRX_KERNEL_REF("qwen3_moe", "ggml_linear_q6k_q8_1_x4"); +static constexpr float kQwenRmsNormEpsilon = kQwen30BMoeDispatchProfile.rms_norm_epsilon; +static constexpr int64_t kQwenHiddenSize = kQwen30BMoeDispatchProfile.hidden_size; +static constexpr int64_t kQwenVocabularyCount = 151936; + +static const Value * graph_value(const Graph & graph, ValueId id) { + return graph.values().find(id); +} + +static bool same_shape(const Value & lhs, const Value & rhs) { + for (int i = 0; i < GGML_MAX_DIMS; ++i) { + if (lhs.ne[i] != rhs.ne[i]) { + return false; + } + } + return true; +} + +static bool is_qwen_rms_norm_epsilon(float eps) { + return std::fabs(eps - kQwenRmsNormEpsilon) <= 1.0e-12f; +} + +static bool is_supported_hidden_size(int64_t hidden_size) { + return hidden_size >= 128 && hidden_size <= 32768 && hidden_size % 128 == 0; +} + +static bool is_supported_token_count(int64_t token_count) { + return is_llm_supported_query_length(kQwen30BMoeDispatchProfile, token_count); +} + +static bool has_decode_q8_consumer(const Graph & graph, ValueId value) { + if (!graph.has_index()) { + return false; + } + for (const GraphNode * consumer : graph.index().consumers(value)) { + if (consumer != nullptr && (consumer->op == GGML_OP_MUL_MAT || consumer->op == GGML_OP_MUL_MAT_ID)) { + return true; + } + } + return false; +} + +static size_t q8_1_x4_byte_count(int64_t token_count, int64_t hidden_size) { + if (token_count <= 0 || hidden_size <= 0) { + return 0; + } + return static_cast(token_count) * ggml_row_size(GGML_TYPE_Q8_1, hidden_size); +} + +static bool is_weight_shape(const Value & weight, int64_t hidden_size) { + if (weight.ne[0] != hidden_size) { + return false; + } + for (int i = 1; i < GGML_MAX_DIMS; ++i) { + if (weight.ne[i] != 1) { + return false; + } + } + return true; +} + +struct RmsNormMatch { + const GraphNode * rms_node = nullptr; + const GraphNode * mul_node = nullptr; + const Value * input = nullptr; + const Value * weight = nullptr; + const Value * output = nullptr; + size_t rms_node_index = 0; + size_t mul_node_index = 0; + int64_t hidden_size = 0; + int64_t token_count = 0; + int64_t q8_group_count = 0; + + bool matched() const { + return rms_node != nullptr && mul_node != nullptr && input != nullptr && weight != nullptr && output != nullptr; + } +}; + +static RmsNormMatch match_qwen_rmsnorm_f32(const Graph & graph, const GraphNode * node, size_t node_index) { + RmsNormMatch match; + if (node == nullptr || node->op != GGML_OP_RMS_NORM || node->inputs.size() != 1 || !graph.has_index()) { + return match; + } + + const RmsNormParams * rms_params = op_params_as(node->params); + if (rms_params == nullptr || !is_qwen_rms_norm_epsilon(rms_params->eps)) { + return {}; + } + const std::vector & consumers = graph.index().consumers(node->output); + if (consumers.size() != 1) { + return {}; + } + const GraphNode * mul_node = consumers.front(); + size_t mul_node_index; + if (mul_node == nullptr || mul_node->op != GGML_OP_MUL || mul_node->inputs.size() != 2 || + !graph.index().node_index(mul_node, mul_node_index)) { + return {}; + } + + const Value * weight = nullptr; + for (ValueId input : mul_node->inputs) { + if (input != node->output) { + weight = graph_value(graph, input); + } + } + const Value * input = graph_value(graph, node->inputs[0]); + const Value * rms = graph_value(graph, node->output); + const Value * output = graph_value(graph, mul_node->output); + if (input == nullptr || rms == nullptr || weight == nullptr || output == nullptr) { + return {}; + } + if (input->type != GGML_TYPE_F32 || rms->type != GGML_TYPE_F32 || weight->type != GGML_TYPE_F32 || + output->type != GGML_TYPE_F32) { + return {}; + } + if (!input->contiguous || !rms->contiguous || !weight->contiguous || !output->contiguous) { + return {}; + } + if (!same_shape(*input, *rms) || !same_shape(*input, *output)) { + return {}; + } + + const int64_t hidden_size = output->ne[0]; + if (!is_supported_hidden_size(hidden_size) || !is_weight_shape(*weight, hidden_size)) { + return {}; + } + if (hidden_size == 0 || output->element_count <= 0 || output->element_count % hidden_size != 0) { + return {}; + } + const int64_t token_count = output->element_count / hidden_size; + if (!is_supported_token_count(token_count)) { + return {}; + } + + match.rms_node = node; + match.mul_node = mul_node; + match.input = input; + match.weight = weight; + match.output = output; + match.rms_node_index = node_index; + match.mul_node_index = mul_node_index; + match.hidden_size = hidden_size; + match.token_count = token_count; + match.q8_group_count = token_count * ((hidden_size + 127) / 128); + return match; +} + +struct QwenEndpointProjectionMatch { + const GraphNode * projection_node = nullptr; + const Value * weight = nullptr; + const Value * output = nullptr; + size_t projection_node_index = 0; + + bool matched() const { return projection_node != nullptr && weight != nullptr && output != nullptr; } +}; + +static QwenEndpointProjectionMatch match_qwen_endpoint_projection(const Graph & graph, const RmsNormMatch & match) { + QwenEndpointProjectionMatch projection_match; + if (!graph.has_index()) { + return projection_match; + } + const std::vector & consumers = graph.index().consumers(match.output->id); + if (consumers.size() != 1) { + return {}; + } + const GraphNode * consumer = consumers.front(); + if (consumer == nullptr || consumer->op != GGML_OP_MUL_MAT || consumer->inputs.size() != 2) { + return {}; + } + size_t consumer_index = 0; + if (!graph.index().node_index(consumer, consumer_index)) { + return {}; + } + const Value * weight = graph_value(graph, consumer->inputs[0]); + const Value * input = graph_value(graph, consumer->inputs[1]); + const Value * output = graph_value(graph, consumer->output); + if (weight == nullptr || input == nullptr || output == nullptr) { + return {}; + } + if (input->id != match.output->id || weight->type != GGML_TYPE_Q6_K || output->type != GGML_TYPE_F32 || + !weight->contiguous || !input->contiguous || !output->contiguous || match.hidden_size != kQwenHiddenSize || + !is_supported_token_count(match.token_count) || weight->ne[0] != match.hidden_size || + weight->ne[1] != kQwenVocabularyCount || output->ne[0] != kQwenVocabularyCount || + output->ne[1] != match.token_count || output->ne[2] != 1 || output->ne[3] != 1) { + return {}; + } + + projection_match.projection_node = consumer; + projection_match.projection_node_index = consumer_index; + projection_match.weight = weight; + projection_match.output = output; + return projection_match; +} + +static RmsNormMatch match_qwen_endpoint_rmsnorm_from_projection(const Graph & graph, const GraphNode * projection) { + if (projection == nullptr || projection->op != GGML_OP_MUL_MAT || projection->inputs.size() != 2 || + !graph.has_index()) { + return {}; + } + + const Value * input = graph_value(graph, projection->inputs[1]); + if (input == nullptr) { + return {}; + } + const GraphNode * mul_node = graph.index().producer(input->id); + size_t mul_node_index; + if (mul_node == nullptr || mul_node->op != GGML_OP_MUL || mul_node->inputs.size() != 2 || + !graph.index().node_index(mul_node, mul_node_index)) { + return {}; + } + + for (const ValueId mul_input : mul_node->inputs) { + const GraphNode * rms_node = graph.index().producer(mul_input); + size_t rms_node_index; + if (rms_node != nullptr && rms_node->op == GGML_OP_RMS_NORM && + graph.index().node_index(rms_node, rms_node_index)) { + return match_qwen_rmsnorm_f32(graph, rms_node, rms_node_index); + } + } + return {}; +} + +static std::string to_config_value(int64_t value) { + return std::to_string(value); +} + +} // namespace + +static bool match_qwen_rmsnorm_f32_dispatch(const DispatchMatchContext & context, DispatchMatch & match) { + const std::vector & nodes = context.graph.nodes(); + if (context.root_index >= nodes.size()) { + return false; + } + const RmsNormMatch rms_match = + match_qwen_rmsnorm_f32(context.graph, &nodes[context.root_index], context.root_index); + if (!rms_match.matched() || rms_match.rms_node_index >= context.covered_nodes.size() || + rms_match.mul_node_index >= context.covered_nodes.size() || context.covered_nodes[rms_match.rms_node_index] || + context.covered_nodes[rms_match.mul_node_index]) { + return false; + } + + Dispatch dispatch; + dispatch.kernel = make_kernel_specialization(kQwenRmsNormF32Kernel); + dispatch.kernel.integer_parameters.emplace("token_count", rms_match.token_count); + dispatch.kernel.compile_parameters.emplace("qwen3_moe.model.hidden_size", to_config_value(rms_match.hidden_size)); + dispatch.kernel.compile_parameters.emplace("qwen3_moe.model.rms_epsilon", "0.000001"); + dispatch.kernel.compile_parameters.emplace("qwen3_moe.workload.token_capacity", + to_config_value(rms_match.token_count)); + dispatch.kernel.compile_parameters.emplace("ggml.quantize_q8_1_x4.group_capacity", + to_config_value(rms_match.q8_group_count)); + dispatch.bindings.push_back({ rms_match.input->id, 0, rms_match.input->byte_count }); + dispatch.bindings.push_back({ rms_match.weight->id, 0, rms_match.weight->byte_count }); + dispatch.bindings.push_back({ rms_match.output->id, 0, rms_match.output->byte_count }); + + match.covered_nodes.push_back(rms_match.rms_node_index); + match.covered_nodes.push_back(rms_match.mul_node_index); + match.dispatches.push_back(std::move(dispatch)); + return true; +} + +static bool match_qwen_rmsnorm_f32_quantize_q8_1_x4_dispatch(const DispatchMatchContext & context, + DispatchMatch & match) { + const std::vector & nodes = context.graph.nodes(); + if (context.root_index >= nodes.size()) { + return false; + } + const RmsNormMatch rms_match = + context.root_node->op == GGML_OP_MUL_MAT ? + match_qwen_endpoint_rmsnorm_from_projection(context.graph, context.root_node) : + match_qwen_rmsnorm_f32(context.graph, &nodes[context.root_index], context.root_index); + const QwenEndpointProjectionMatch projection_match = + rms_match.matched() ? match_qwen_endpoint_projection(context.graph, rms_match) : QwenEndpointProjectionMatch{}; + if (!rms_match.matched() || rms_match.rms_node_index >= context.covered_nodes.size() || + rms_match.mul_node_index >= context.covered_nodes.size() || context.covered_nodes[rms_match.rms_node_index] || + context.covered_nodes[rms_match.mul_node_index] || !projection_match.matched() || + projection_match.projection_node_index >= context.covered_nodes.size() || + context.covered_nodes[projection_match.projection_node_index] || + (context.root_node->op == GGML_OP_MUL_MAT && projection_match.projection_node != context.root_node)) { + return false; + } + + const size_t q8_byte_count = q8_1_x4_byte_count(rms_match.token_count, rms_match.hidden_size); + if (q8_byte_count == 0) { + return false; + } + + const ValueId q8_value = context.next_plan_value; + + Dispatch rms_dispatch; + rms_dispatch.kernel = make_kernel_specialization(kQwenRmsNormF32QuantizeQ8_1X4Kernel); + rms_dispatch.kernel.integer_parameters.emplace("token_count", rms_match.token_count); + rms_dispatch.kernel.compile_parameters.emplace("qwen3_moe.model.hidden_size", + to_config_value(rms_match.hidden_size)); + rms_dispatch.kernel.compile_parameters.emplace("qwen3_moe.model.rms_epsilon", "0.000001"); + rms_dispatch.kernel.compile_parameters.emplace("qwen3_moe.workload.token_capacity", + to_config_value(rms_match.token_count)); + rms_dispatch.kernel.compile_parameters.emplace("ggml.quantize_q8_1_x4.group_capacity", + to_config_value(rms_match.q8_group_count)); + rms_dispatch.bindings.push_back({ rms_match.input->id, 0, rms_match.input->byte_count }); + rms_dispatch.bindings.push_back({ rms_match.weight->id, 0, rms_match.weight->byte_count }); + rms_dispatch.bindings.push_back({ rms_match.output->id, 0, rms_match.output->byte_count }); + rms_dispatch.bindings.push_back({ q8_value, 0, q8_byte_count }); + + Dispatch projection_dispatch; + projection_dispatch.kernel = make_kernel_specialization(kGgmlLinearQ6KQ8_1X4Kernel); + projection_dispatch.kernel.integer_parameters.emplace("token_count", rms_match.token_count); + projection_dispatch.kernel.integer_parameters.emplace("input_size", rms_match.hidden_size); + projection_dispatch.kernel.integer_parameters.emplace("output_size", kQwenVocabularyCount); + projection_dispatch.kernel.compile_parameters.emplace("ggml.linear_q6k_q8_1_x4.token_capacity", + to_config_value(rms_match.token_count)); + projection_dispatch.kernel.compile_parameters.emplace("ggml.linear_q6k_q8_1_x4.output_capacity", + to_config_value(kQwenVocabularyCount)); + projection_dispatch.bindings.push_back({ q8_value, 0, q8_byte_count }); + projection_dispatch.bindings.push_back({ projection_match.weight->id, 0, projection_match.weight->byte_count }); + projection_dispatch.bindings.push_back({ projection_match.output->id, 0, projection_match.output->byte_count }); + + match.covered_nodes.push_back(rms_match.rms_node_index); + match.covered_nodes.push_back(rms_match.mul_node_index); + match.covered_nodes.push_back(projection_match.projection_node_index); + match.dispatches.push_back(std::move(rms_dispatch)); + match.dispatches.push_back(std::move(projection_dispatch)); + match.transients.push_back({ q8_value, "qwen.rmsnorm.q8_1_x4", q8_byte_count, 256 }); + return match.status.success(); +} + +static bool match_qwen_decode_rmsnorm_f32_quantize_q8_1_x4_dispatch(const DispatchMatchContext & context, + DispatchMatch & match) { + const std::vector & nodes = context.graph.nodes(); + if (context.root_index >= nodes.size()) { + return false; + } + const RmsNormMatch rms_match = + match_qwen_rmsnorm_f32(context.graph, &nodes[context.root_index], context.root_index); + if (!rms_match.matched() || rms_match.token_count != 1 || rms_match.hidden_size != kQwenHiddenSize || + rms_match.rms_node_index >= context.covered_nodes.size() || + rms_match.mul_node_index >= context.covered_nodes.size() || context.covered_nodes[rms_match.rms_node_index] || + context.covered_nodes[rms_match.mul_node_index]) { + return false; + } + if (!has_decode_q8_consumer(context.graph, rms_match.output->id)) { + return false; + } + + const size_t q8_byte_count = q8_1_x4_byte_count(rms_match.token_count, rms_match.hidden_size); + if (q8_byte_count == 0) { + return false; + } + + const ValueId q8_value = context.next_plan_value; + + Dispatch dispatch; + dispatch.kernel = make_kernel_specialization(kQwenRmsNormF32QuantizeQ8_1X4Kernel); + dispatch.kernel.integer_parameters.emplace("token_count", rms_match.token_count); + dispatch.kernel.compile_parameters.emplace("qwen3_moe.model.hidden_size", to_config_value(rms_match.hidden_size)); + dispatch.kernel.compile_parameters.emplace("qwen3_moe.model.rms_epsilon", "0.000001"); + dispatch.kernel.compile_parameters.emplace("qwen3_moe.workload.token_capacity", + to_config_value(rms_match.token_count)); + dispatch.kernel.compile_parameters.emplace("ggml.quantize_q8_1_x4.group_capacity", + to_config_value(rms_match.q8_group_count)); + dispatch.bindings.push_back({ rms_match.input->id, 0, rms_match.input->byte_count }); + dispatch.bindings.push_back({ rms_match.weight->id, 0, rms_match.weight->byte_count }); + dispatch.bindings.push_back({ rms_match.output->id, 0, rms_match.output->byte_count }); + dispatch.bindings.push_back({ q8_value, 0, q8_byte_count }); + + Status metadata_status; + if (!match.metadata.append_alternate_value( + { rms_match.output->id, q8_value, GGML_TYPE_Q8_1, q8_byte_count, "qwen.decode.q8_hidden" }, + metadata_status)) { + match.status.append(metadata_status); + return false; + } + + match.covered_nodes.push_back(rms_match.rms_node_index); + match.covered_nodes.push_back(rms_match.mul_node_index); + match.dispatches.push_back(std::move(dispatch)); + match.transients.push_back({ q8_value, "qwen.decode.q8_hidden", q8_byte_count, 256 }); + return match.status.success(); +} + +void register_qwen_rmsnorm_dispatches(DispatchRegistryBuilder & registry) { + registry.add({ + "qwen.endpoint_rmsnorm_q6k_q8_1_x4", + GGML_OP_MUL_MAT, + DispatchMatchKind::Fused, + 1200, + DispatchSource::Qwen, + match_qwen_rmsnorm_f32_quantize_q8_1_x4_dispatch, + }); + registry.add({ + "qwen.decode_rmsnorm_f32_quantize_q8_1_x4", + GGML_OP_RMS_NORM, + DispatchMatchKind::Fused, + 1150, + DispatchSource::Qwen, + match_qwen_decode_rmsnorm_f32_quantize_q8_1_x4_dispatch, + }); + registry.add({ + "qwen.rmsnorm_f32_quantize_q8_1_x4", + GGML_OP_RMS_NORM, + DispatchMatchKind::Fused, + 1100, + DispatchSource::Qwen, + match_qwen_rmsnorm_f32_quantize_q8_1_x4_dispatch, + }); + registry.add({ + "qwen.rmsnorm_f32.mul_weight", + GGML_OP_RMS_NORM, + DispatchMatchKind::Fused, + 1000, + DispatchSource::Qwen, + match_qwen_rmsnorm_f32_dispatch, + }); +} + +} // namespace ggml::hrx diff --git a/ggml/src/ggml-hrx/dispatch_registration/dispatch-rmsnorm.h b/ggml/src/ggml-hrx/dispatch_registration/dispatch-rmsnorm.h new file mode 100644 index 000000000000..4a238a7e22b8 --- /dev/null +++ b/ggml/src/ggml-hrx/dispatch_registration/dispatch-rmsnorm.h @@ -0,0 +1,9 @@ +#pragma once + +#include "dispatch-registry.h" + +namespace ggml::hrx { + +void register_qwen_rmsnorm_dispatches(DispatchRegistryBuilder & registry); + +} // namespace ggml::hrx diff --git a/ggml/src/ggml-hrx/dispatch_registration/dispatch-routed-ffn.cpp b/ggml/src/ggml-hrx/dispatch_registration/dispatch-routed-ffn.cpp new file mode 100644 index 000000000000..46bf7025649c --- /dev/null +++ b/ggml/src/ggml-hrx/dispatch_registration/dispatch-routed-ffn.cpp @@ -0,0 +1,1228 @@ +#include "dispatch-routed-ffn.h" + +#include "dispatch-llm-shapes.h" +#include "ggml.h" +#include "graph/graph-matcher.h" +#include "kernel-corpus/kernel-corpus-catalog-verify.h" + +#include +#include +#include +#include +#include + +namespace ggml::hrx { +namespace { + +static constexpr KernelCatalogRef kQwenRoutedGateUpSwiGLUQ4KF16WmmaKernel = + GGML_HRX_KERNEL_REF("qwen3_moe", "qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma"); +static constexpr KernelCatalogRef kQwenRoutedGateUpSwiGLUQ4KQ8Kernel = + GGML_HRX_KERNEL_REF("qwen3_moe", "qwen3_moe_routed_gate_up_swiglu_q4k_q8"); +static constexpr KernelCatalogRef kQwenRoutedGateUpSwiGLUQ4KQ8NextQ8Kernel = + GGML_HRX_KERNEL_REF("qwen3_moe", "qwen3_moe_routed_gate_up_swiglu_q4k_q8_1_x4_next_q8"); +static constexpr KernelCatalogRef kQwenRoutedDownQ4KF16WmmaGroupedKernel = + GGML_HRX_KERNEL_REF("qwen3_moe", "qwen3_moe_routed_down_q4k_f16_wmma_grouped"); +static constexpr KernelCatalogRef kQwenRoutedDownQ6KF16WmmaGroupedKernel = + GGML_HRX_KERNEL_REF("qwen3_moe", "qwen3_moe_routed_down_q6k_f16_wmma_grouped"); +static constexpr KernelCatalogRef kQwenRoutedDownQ4KQ8NextQ8Kernel = + GGML_HRX_KERNEL_REF("qwen3_moe", "qwen3_moe_routed_down_q4k_q8_1_x4_next_q8"); +static constexpr KernelCatalogRef kQwenRoutedDownQ6KF32Wave64NextQ8Kernel = + GGML_HRX_KERNEL_REF("qwen3_moe", "qwen3_moe_routed_down_q6k_f32_wave64_next_q8"); +static constexpr KernelCatalogRef kQwenRoutedDownWeightedReduceF16F32Kernel = + GGML_HRX_KERNEL_REF("qwen3_moe", "qwen3_moe_routed_down_weighted_reduce_f16_f32"); +static constexpr KernelCatalogRef kQwenRoutedDownWeightedReduceNextRmsNormF32Kernel = + GGML_HRX_KERNEL_REF("qwen3_moe", "qwen3_moe_routed_down_weighted_reduce_next_rmsnorm_f32"); + +static constexpr const LlmMoeDispatchProfile & kRoutedFfnProfile = kActiveLlmMoeDispatchProfile; +static constexpr int64_t kRoutedFfnInputSize = kRoutedFfnProfile.hidden_size; +static constexpr int64_t kRoutedFfnExpertHiddenSize = kRoutedFfnProfile.expert_hidden_size; +static constexpr int64_t kRoutedFfnExpertCount = kRoutedFfnProfile.expert_count; +static constexpr int64_t kRoutedFfnRouteCount = kRoutedFfnProfile.route_count; +static constexpr size_t kRoutedFfnPlanTransientAlignment = 256; +static constexpr const char * kRoutedFfnF16GateUpOutputName = "qwen.moe.gate_up_swiglu_f16"; +static constexpr const char * kRoutedFfnF16RoutedDownOutputName = "qwen.moe.routed_down_f16"; +static constexpr const char * kRoutedFfnQ8GateUpOutputName = "qwen.decode.moe.gate_up_swiglu_q8"; +static constexpr const char * kRoutedFfnQ8HiddenOutputName = "qwen.decode.moe.hidden_q8"; + +static const Value * graph_value(const Graph & graph, ValueId id) { + return graph.values().find(id); +} + +static bool is_shape(const Value & value, int64_t ne0, int64_t ne1, int64_t ne2, int64_t ne3) { + return value.ne[0] == ne0 && value.ne[1] == ne1 && value.ne[2] == ne2 && value.ne[3] == ne3; +} + +static bool same_shape(const Value & lhs, const Value & rhs) { + for (int i = 0; i < GGML_MAX_DIMS; ++i) { + if (lhs.ne[i] != rhs.ne[i]) { + return false; + } + } + return true; +} + +static bool is_profile_rms_norm_epsilon(float eps) { + const float expected = kRoutedFfnProfile.rms_norm_epsilon; + return eps >= expected * 0.9f && eps <= expected * 1.1f; +} + +static bool is_routed_ffn_gate_up_weight(const Value & value) { + return value.type == GGML_TYPE_Q4_K && value.contiguous && + is_shape(value, kRoutedFfnInputSize, kRoutedFfnExpertHiddenSize, kRoutedFfnExpertCount, 1); +} + +static bool is_routed_ffn_down_weight(const Value & value) { + return (value.type == GGML_TYPE_Q4_K || value.type == GGML_TYPE_Q6_K) && value.contiguous && + is_shape(value, kRoutedFfnExpertHiddenSize, kRoutedFfnInputSize, kRoutedFfnExpertCount, 1); +} + +static bool is_routed_ffn_projection_output(const Value & value, int64_t token_count) { + return value.type == GGML_TYPE_F32 && value.contiguous && + is_shape(value, kRoutedFfnExpertHiddenSize, kRoutedFfnRouteCount, token_count, 1); +} + +static bool is_routed_ffn_down_output(const Value & value, int64_t token_count) { + return value.type == GGML_TYPE_F32 && value.contiguous && + is_shape(value, kRoutedFfnInputSize, kRoutedFfnRouteCount, token_count, 1); +} + +static const GraphNode * find_consumer_with_op(const Graph & graph, ValueId value, ggml_op op) { + for (const GraphNode * consumer : graph.index().consumers(value)) { + if (consumer != nullptr && consumer->op == op) { + return consumer; + } + } + return nullptr; +} + +static std::vector find_consumers_with_op(const Graph & graph, ValueId value, ggml_op op) { + std::vector matches; + for (const GraphNode * consumer : graph.index().consumers(value)) { + if (consumer != nullptr && consumer->op == op) { + matches.push_back(consumer); + } + } + return matches; +} + +static const GraphNode * find_single_consumer_with_op(const Graph & graph, ValueId value, ggml_op op) { + const std::vector consumers = find_consumers_with_op(graph, value, op); + return consumers.size() == 1 ? consumers.front() : nullptr; +} + +static const GraphNode * producer_with_op(const Graph & graph, ValueId value, ggml_op op) { + const GraphNode * producer = graph.index().producer(value); + return producer != nullptr && producer->op == op ? producer : nullptr; +} + +static bool append_covered_node(const DispatchMatchContext & context, const GraphNode * node, DispatchMatch & match) { + return append_covered_node_index_once(context.graph, context.covered_nodes, node, match.covered_nodes); +} + +static std::string to_config_value(int64_t value) { + return std::to_string(value); +} + +static size_t expert_table_size(int64_t token_count) { + return static_cast(kRoutedFfnExpertCount + kRoutedFfnExpertCount * token_count) * sizeof(int32_t); +} + +static size_t partition_table_size(int64_t token_count) { + const int64_t assignment_count = token_count * kRoutedFfnRouteCount; + const int64_t assignment_partition_count = (assignment_count + 31) / 32; + return static_cast(1 + assignment_partition_count + kRoutedFfnExpertCount) * sizeof(int32_t); +} + +static size_t f16_gate_up_output_size(int64_t token_count) { + return static_cast(token_count * kRoutedFfnRouteCount * kRoutedFfnExpertHiddenSize) * sizeof(ggml_fp16_t); +} + +static size_t f16_routed_down_output_size(int64_t token_count) { + return static_cast(token_count * kRoutedFfnRouteCount * kRoutedFfnInputSize) * sizeof(ggml_fp16_t); +} + +static size_t q8_1_x4_byte_count(int64_t row_count, int64_t input_size) { + if (row_count <= 0 || input_size <= 0) { + return 0; + } + return static_cast(row_count) * ggml_row_size(GGML_TYPE_Q8_1, input_size); +} + +static uint32_t gate_up_completion_counter_count(int64_t token_count) { + const int64_t physical_group_count = (kRoutedFfnExpertHiddenSize + 127) / 128; + return static_cast(token_count * kRoutedFfnRouteCount * physical_group_count); +} + +static void add_routed_down_compile_parameters(Dispatch & dispatch, int64_t token_count) { + dispatch.kernel.compile_parameters.emplace("qwen3_moe.routed_down.input_size", + to_config_value(kRoutedFfnExpertHiddenSize)); + dispatch.kernel.compile_parameters.emplace("qwen3_moe.routed_down.route_count", + to_config_value(kRoutedFfnRouteCount)); + dispatch.kernel.compile_parameters.emplace("qwen3_moe.routed_down.expert_count", + to_config_value(kRoutedFfnExpertCount)); + dispatch.kernel.compile_parameters.emplace("qwen3_moe.routed_down.output_size", + to_config_value(kRoutedFfnInputSize)); + dispatch.kernel.compile_parameters.emplace("qwen3_moe.workload.token_capacity", to_config_value(token_count)); +} + +struct RoutedGateUpMatch { + const Value * gate_weight = nullptr; + const Value * up_weight = nullptr; + const Value * input = nullptr; + const Value * route_ids = nullptr; + const Value * gate_output = nullptr; + const Value * up_output = nullptr; + const Value * glu_output = nullptr; + const CommandPlanMoeRoutingBundle * routing_bundle = nullptr; + const GraphNode * gate_node = nullptr; + const GraphNode * up_node = nullptr; + const GraphNode * glu_node = nullptr; + int64_t token_count = 0; + + bool matched() const { + return gate_weight != nullptr && up_weight != nullptr && input != nullptr && route_ids != nullptr && + gate_output != nullptr && up_output != nullptr && glu_output != nullptr && routing_bundle != nullptr && + gate_node != nullptr && up_node != nullptr && glu_node != nullptr && token_count > 0; + } +}; + +struct DecodeRoutedGateUpMatch { + const Value * gate_weight = nullptr; + const Value * up_weight = nullptr; + const Value * input = nullptr; + const Value * route_ids = nullptr; + const Value * glu_output = nullptr; + const CommandPlanAlternateValue * input_alternate = nullptr; + const GraphNode * down_node = nullptr; + const Value * down_weight = nullptr; + const GraphNode * gate_node = nullptr; + const GraphNode * up_node = nullptr; + const GraphNode * glu_node = nullptr; + int64_t token_count = 0; + int64_t route_stride = 0; + bool publish_q8 = false; + + bool matched() const { + return gate_weight != nullptr && up_weight != nullptr && input != nullptr && route_ids != nullptr && + glu_output != nullptr && input_alternate != nullptr && down_node != nullptr && down_weight != nullptr && + gate_node != nullptr && up_node != nullptr && glu_node != nullptr && token_count == 1 && + route_stride >= kRoutedFfnRouteCount; + } +}; + +struct RoutedDownMatch { + const Value * input_graph_value = nullptr; + const CommandPlanAlternateValue * input_alternate = nullptr; + const Value * weight = nullptr; + const Value * output = nullptr; + const Value * route_ids = nullptr; + const CommandPlanMoeRoutingBundle * routing_bundle = nullptr; + KernelCatalogRef kernel = {}; + int64_t token_count = 0; + + bool matched() const { + return input_graph_value != nullptr && input_alternate != nullptr && weight != nullptr && output != nullptr && + route_ids != nullptr && routing_bundle != nullptr && kernel.id != kUncatalogedKernelId && + token_count > 0; + } +}; + +struct WeightedReduceNextRmsNormMatch { + const GraphNode * rms_node = nullptr; + const GraphNode * mul_node = nullptr; + const Value * norm_weight = nullptr; + const Value * output = nullptr; + + bool matched() const { + return rms_node != nullptr && mul_node != nullptr && norm_weight != nullptr && output != nullptr; + } +}; + +struct WeightedReduceMatch { + const Value * route_weights = nullptr; + const Value * routed_output = nullptr; + const CommandPlanAlternateValue * routed_alternate = nullptr; + const Value * residual_input = nullptr; + const Value * output = nullptr; + const GraphNode * weighted_node = nullptr; + std::vector views; + std::vector reductions; + const GraphNode * residual = nullptr; + WeightedReduceNextRmsNormMatch next_rmsnorm; + int64_t token_count = 0; + + bool topology_matched() const { + return route_weights != nullptr && routed_output != nullptr && residual_input != nullptr && output != nullptr && + weighted_node != nullptr && !views.empty() && residual != nullptr && token_count > 0; + } + + bool matched() const { return topology_matched() && routed_alternate != nullptr; } +}; + +struct DecodeRoutedDownMatch { + const Value * input_graph_value = nullptr; + const CommandPlanAlternateValue * input_alternate = nullptr; + const Value * weight = nullptr; + const Value * output = nullptr; + const Value * route_ids = nullptr; + WeightedReduceMatch reduce; + KernelCatalogRef kernel = {}; + int64_t token_count = 0; + int64_t route_stride = 0; + bool input_is_q8 = false; + + bool matched() const { + return input_graph_value != nullptr && weight != nullptr && output != nullptr && route_ids != nullptr && + reduce.topology_matched() && reduce.next_rmsnorm.matched() && kernel.id != kUncatalogedKernelId && + token_count == 1 && route_stride >= kRoutedFfnRouteCount && (!input_is_q8 || input_alternate != nullptr); + } +}; + +static bool bundle_matches_moe_routing(const CommandPlanMoeRoutingBundle & bundle, + ValueId route_ids, + int64_t token_count) { + return bundle.route_ids == route_ids && bundle.route_weights.value >= 0 && bundle.expert_table.value >= 0 && + bundle.partition_table.value >= 0 && bundle.expert_table_byte_count == expert_table_size(token_count) && + bundle.partition_table_byte_count == partition_table_size(token_count) && + bundle.token_count == token_count && bundle.route_count == kRoutedFfnRouteCount && + bundle.expert_count == kRoutedFfnExpertCount && bundle.route_stride >= kRoutedFfnRouteCount; +} + +static bool match_same_route_projection(const Graph & graph, + const GraphNode & node, + ValueId expected_input, + ValueId expected_route_ids, + int64_t token_count, + const Value *& weight, + const Value *& output) { + if (node.op != GGML_OP_MUL_MAT_ID || node.inputs.size() != 3 || node.inputs[1] != expected_input || + node.inputs[2] != expected_route_ids) { + return false; + } + weight = graph_value(graph, node.inputs[0]); + output = graph_value(graph, node.output); + return weight != nullptr && output != nullptr && is_routed_ffn_gate_up_weight(*weight) && + is_routed_ffn_projection_output(*output, token_count); +} + +static RoutedDownMatch match_routed_ffn_down_grouped(const DispatchMatchContext & context) { + RoutedDownMatch match; + const GraphNode * root = context.root_node; + if (root == nullptr || root->op != GGML_OP_MUL_MAT_ID || root->inputs.size() != 3 || !context.graph.has_index()) { + return match; + } + + const Value * weight = graph_value(context.graph, root->inputs[0]); + const Value * input = graph_value(context.graph, root->inputs[1]); + const Value * route_ids = graph_value(context.graph, root->inputs[2]); + const Value * root_output = graph_value(context.graph, root->output); + if (weight == nullptr || input == nullptr || route_ids == nullptr || root_output == nullptr || + !is_routed_ffn_down_weight(*weight) || !is_routed_ffn_projection_output(*input, input->ne[2]) || + route_ids->type != GGML_TYPE_I32 || !is_shape(*route_ids, kRoutedFfnRouteCount, input->ne[2], 1, 1)) { + return {}; + } + + const int64_t token_count = input->ne[2]; + if (!is_llm_supported_query_length(kRoutedFfnProfile, token_count) || + !is_routed_ffn_down_output(*root_output, token_count)) { + return {}; + } + + const CommandPlanMoeRoutingBundle * routing_bundle = context.plan.metadata.find_moe_routing_bundle(route_ids->id); + if (routing_bundle == nullptr || !bundle_matches_moe_routing(*routing_bundle, route_ids->id, token_count)) { + return {}; + } + + const CommandPlanAlternateValue * input_alternate = + find_alternate_value(context.plan, input->id, GGML_TYPE_F16, f16_gate_up_output_size(token_count)); + if (input_alternate == nullptr) { + return {}; + } + + match.input_graph_value = input; + match.input_alternate = input_alternate; + match.weight = weight; + match.output = root_output; + match.route_ids = route_ids; + match.routing_bundle = routing_bundle; + match.kernel = weight->type == GGML_TYPE_Q4_K ? kQwenRoutedDownQ4KF16WmmaGroupedKernel : + kQwenRoutedDownQ6KF16WmmaGroupedKernel; + match.token_count = token_count; + return match; +} + +static RoutedGateUpMatch match_routed_ffn_gate_up_swiglu(const DispatchMatchContext & context) { + RoutedGateUpMatch match; + const GraphNode * root = context.root_node; + if (root == nullptr || root->op != GGML_OP_MUL_MAT_ID || root->inputs.size() != 3 || !context.graph.has_index()) { + return match; + } + + const Value * root_weight = graph_value(context.graph, root->inputs[0]); + const Value * input = graph_value(context.graph, root->inputs[1]); + const Value * route_ids = graph_value(context.graph, root->inputs[2]); + const Value * root_output = graph_value(context.graph, root->output); + if (root_weight == nullptr || input == nullptr || route_ids == nullptr || root_output == nullptr || + !is_routed_ffn_gate_up_weight(*root_weight) || input->type != GGML_TYPE_F32 || !input->contiguous || + !is_shape(*input, kRoutedFfnInputSize, 1, input->ne[2], 1) || route_ids->type != GGML_TYPE_I32 || + !is_shape(*route_ids, kRoutedFfnRouteCount, input->ne[2], 1, 1)) { + return {}; + } + + const int64_t token_count = input->ne[2]; + if (!is_llm_supported_query_length(kRoutedFfnProfile, token_count) || + !is_routed_ffn_projection_output(*root_output, token_count)) { + return {}; + } + + const CommandPlanMoeRoutingBundle * routing_bundle = context.plan.metadata.find_moe_routing_bundle(route_ids->id); + if (routing_bundle == nullptr || !bundle_matches_moe_routing(*routing_bundle, route_ids->id, token_count)) { + return {}; + } + + const GraphNode * glu_node = find_consumer_with_op(context.graph, root->output, GGML_OP_GLU); + if (glu_node == nullptr || glu_node->inputs.size() != 2) { + return {}; + } + const GluParams * glu_params = op_params_as(glu_node->params); + if (glu_params == nullptr || glu_params->op != GGML_GLU_OP_SWIGLU) { + return {}; + } + + const GraphNode * gate_node = producer_with_op(context.graph, glu_node->inputs[0], GGML_OP_MUL_MAT_ID); + const GraphNode * up_node = producer_with_op(context.graph, glu_node->inputs[1], GGML_OP_MUL_MAT_ID); + if (gate_node == nullptr || up_node == nullptr || gate_node == up_node || (gate_node != root && up_node != root)) { + return {}; + } + + const Value * gate_weight = nullptr; + const Value * gate_output = nullptr; + const Value * up_weight = nullptr; + const Value * up_output = nullptr; + if (!match_same_route_projection(context.graph, *gate_node, input->id, route_ids->id, token_count, gate_weight, + gate_output) || + !match_same_route_projection(context.graph, *up_node, input->id, route_ids->id, token_count, up_weight, + up_output)) { + return {}; + } + if (!same_shape(*gate_output, *up_output)) { + return {}; + } + + const Value * glu_output = graph_value(context.graph, glu_node->output); + if (glu_output == nullptr || glu_output->kind != ValueKind::Transient || + !is_routed_ffn_projection_output(*glu_output, token_count)) { + return {}; + } + + match.gate_weight = gate_weight; + match.up_weight = up_weight; + match.input = input; + match.route_ids = route_ids; + match.gate_output = gate_output; + match.up_output = up_output; + match.glu_output = glu_output; + match.routing_bundle = routing_bundle; + match.gate_node = gate_node; + match.up_node = up_node; + match.glu_node = glu_node; + match.token_count = token_count; + return match; +} + +static DecodeRoutedGateUpMatch match_decode_routed_ffn_gate_up_swiglu(const DispatchMatchContext & context) { + DecodeRoutedGateUpMatch match; + const GraphNode * root = context.root_node; + if (root == nullptr || root->op != GGML_OP_MUL_MAT_ID || root->inputs.size() != 3 || !context.graph.has_index()) { + return match; + } + + const Value * root_weight = graph_value(context.graph, root->inputs[0]); + const Value * input = graph_value(context.graph, root->inputs[1]); + const Value * route_ids = graph_value(context.graph, root->inputs[2]); + const Value * root_output = graph_value(context.graph, root->output); + if (root_weight == nullptr || input == nullptr || route_ids == nullptr || root_output == nullptr || + !is_routed_ffn_gate_up_weight(*root_weight) || input->type != GGML_TYPE_F32 || !input->contiguous || + !is_shape(*input, kRoutedFfnInputSize, 1, 1, 1) || route_ids->type != GGML_TYPE_I32 || + !is_shape(*route_ids, kRoutedFfnRouteCount, 1, 1, 1) || !is_routed_ffn_projection_output(*root_output, 1)) { + return {}; + } + + const int64_t route_stride = static_cast(route_ids->nb[1] / sizeof(int32_t)); + const CommandPlanAlternateValue * input_alternate = find_alternate_value( + context.graph, context.plan, input->id, GGML_TYPE_Q8_1, q8_1_x4_byte_count(1, kRoutedFfnInputSize)); + if (input_alternate == nullptr) { + return {}; + } + + const GraphNode * glu_node = find_consumer_with_op(context.graph, root->output, GGML_OP_GLU); + if (glu_node == nullptr || glu_node->inputs.size() != 2) { + return {}; + } + const GluParams * glu_params = op_params_as(glu_node->params); + if (glu_params == nullptr || glu_params->op != GGML_GLU_OP_SWIGLU) { + return {}; + } + + const GraphNode * gate_node = producer_with_op(context.graph, glu_node->inputs[0], GGML_OP_MUL_MAT_ID); + const GraphNode * up_node = producer_with_op(context.graph, glu_node->inputs[1], GGML_OP_MUL_MAT_ID); + if (gate_node == nullptr || up_node == nullptr || gate_node == up_node || (gate_node != root && up_node != root)) { + return {}; + } + + const Value * gate_weight = nullptr; + const Value * gate_output = nullptr; + const Value * up_weight = nullptr; + const Value * up_output = nullptr; + if (!match_same_route_projection(context.graph, *gate_node, input->id, route_ids->id, 1, gate_weight, + gate_output) || + !match_same_route_projection(context.graph, *up_node, input->id, route_ids->id, 1, up_weight, up_output) || + !same_shape(*gate_output, *up_output)) { + return {}; + } + + const Value * glu_output = graph_value(context.graph, glu_node->output); + if (glu_output == nullptr || glu_output->kind != ValueKind::Transient || + !is_routed_ffn_projection_output(*glu_output, 1)) { + return {}; + } + + const GraphNode * down_node = find_single_consumer_with_op(context.graph, glu_output->id, GGML_OP_MUL_MAT_ID); + const Value * down_weight = down_node == nullptr || down_node->inputs.size() != 3 ? + nullptr : + graph_value(context.graph, down_node->inputs[0]); + if (down_node == nullptr || down_weight == nullptr || !is_routed_ffn_down_weight(*down_weight)) { + return {}; + } + + match.gate_weight = gate_weight; + match.up_weight = up_weight; + match.input = input; + match.route_ids = route_ids; + match.glu_output = glu_output; + match.input_alternate = input_alternate; + match.down_node = down_node; + match.down_weight = down_weight; + match.gate_node = gate_node; + match.up_node = up_node; + match.glu_node = glu_node; + match.token_count = 1; + match.route_stride = route_stride; + match.publish_q8 = down_weight->type == GGML_TYPE_Q4_K; + return match; +} + +static WeightedReduceNextRmsNormMatch match_qwen_weighted_reduce_next_rmsnorm(const DispatchMatchContext & context, + const Value & residual) { + WeightedReduceNextRmsNormMatch match; + const GraphNode * rms_node = find_single_consumer_with_op(context.graph, residual.id, GGML_OP_RMS_NORM); + if (rms_node == nullptr || rms_node->inputs.size() != 1) { + return match; + } + const RmsNormParams * rms_params = op_params_as(rms_node->params); + if (rms_params == nullptr || !is_profile_rms_norm_epsilon(rms_params->eps)) { + return {}; + } + + const Value * rms = graph_value(context.graph, rms_node->output); + if (rms == nullptr || rms->type != GGML_TYPE_F32 || !same_shape(*rms, residual)) { + return {}; + } + + const GraphNode * mul_node = find_single_consumer_with_op(context.graph, rms_node->output, GGML_OP_MUL); + if (mul_node == nullptr || mul_node->inputs.size() != 2) { + return {}; + } + + const Value * norm_weight = nullptr; + for (ValueId input : mul_node->inputs) { + if (input != rms_node->output) { + norm_weight = graph_value(context.graph, input); + } + } + const Value * output = graph_value(context.graph, mul_node->output); + if (norm_weight == nullptr || output == nullptr || norm_weight->type != GGML_TYPE_F32 || + output->type != GGML_TYPE_F32 || !norm_weight->contiguous || !output->contiguous || + !is_shape(*norm_weight, kRoutedFfnInputSize, 1, 1, 1) || !same_shape(*output, residual)) { + return {}; + } + + match.rms_node = rms_node; + match.mul_node = mul_node; + match.norm_weight = norm_weight; + match.output = output; + return match; +} + +static bool append_node_if_uncovered(const DispatchMatchContext & context, + const GraphNode * node, + std::vector & nodes) { + size_t index = 0; + if (node == nullptr || !context.graph.index().node_index(node, index) || index >= context.covered_nodes.size() || + context.covered_nodes[index]) { + return false; + } + for (const GraphNode * existing : nodes) { + if (existing == node) { + return true; + } + } + nodes.push_back(node); + return true; +} + +static bool node_is_covered(const DispatchMatchContext & context, const GraphNode * node) { + size_t index = 0; + return node != nullptr && context.graph.index().node_index(node, index) && index < context.covered_nodes.size() && + context.covered_nodes[index]; +} + +static bool residual_input_is_safe_for_in_place(const DispatchMatchContext & context, + const WeightedReduceMatch & match) { + if (match.residual_input == nullptr || match.residual == nullptr) { + return false; + } + for (const GraphNode * consumer : context.graph.index().consumers(match.residual_input->id)) { + if (consumer == match.residual || node_is_covered(context, consumer)) { + continue; + } + return false; + } + return true; +} + +static const Value * find_qwen_route_weights_for_route_ids(const Graph & graph, + ValueId route_ids, + int64_t token_count) { + const GraphNode * get_rows = find_single_consumer_with_op(graph, route_ids, GGML_OP_GET_ROWS); + if (get_rows == nullptr || get_rows->inputs.size() != 2) { + return nullptr; + } + const Value * selected = graph_value(graph, get_rows->output); + if (selected == nullptr || selected->type != GGML_TYPE_F32 || + !is_shape(*selected, 1, kRoutedFfnRouteCount, token_count, 1)) { + return nullptr; + } + const GraphNode * reshape = find_single_consumer_with_op(graph, get_rows->output, GGML_OP_RESHAPE); + const Value * flat_weights = reshape == nullptr ? nullptr : graph_value(graph, reshape->output); + if (flat_weights == nullptr || flat_weights->type != GGML_TYPE_F32 || + !is_shape(*flat_weights, kRoutedFfnRouteCount, token_count, 1, 1)) { + return nullptr; + } + const GraphNode * sum_rows = find_single_consumer_with_op(graph, reshape->output, GGML_OP_SUM_ROWS); + const GraphNode * clamp = + sum_rows == nullptr ? nullptr : find_single_consumer_with_op(graph, sum_rows->output, GGML_OP_CLAMP); + const GraphNode * div = + clamp == nullptr ? nullptr : find_single_consumer_with_op(graph, clamp->output, GGML_OP_DIV); + const GraphNode * output_reshape = + div == nullptr ? nullptr : find_single_consumer_with_op(graph, div->output, GGML_OP_RESHAPE); + const Value * weights = output_reshape == nullptr ? nullptr : graph_value(graph, output_reshape->output); + if (weights == nullptr || weights->type != GGML_TYPE_F32 || + !is_shape(*weights, 1, kRoutedFfnRouteCount, token_count, 1)) { + return nullptr; + } + return weights; +} + +static WeightedReduceMatch match_routed_ffn_down_weighted_reduce_topology(const DispatchMatchContext & context, + const GraphNode * weighted, + const Value * routed_output, + const Value * route_weights) { + WeightedReduceMatch match; + if (weighted == nullptr || weighted->op != GGML_OP_MUL || weighted->inputs.size() != 2 || + routed_output == nullptr || route_weights == nullptr || !context.graph.has_index()) { + return match; + } + if (!node_has_input_or_alias(context.graph, *weighted, routed_output->id) || + !node_has_input_or_alias(context.graph, *weighted, route_weights->id)) { + return {}; + } + const Value * weighted_output = graph_value(context.graph, weighted->output); + if (!is_routed_ffn_down_output(*routed_output, routed_output->ne[2]) || route_weights->type != GGML_TYPE_F32 || + !route_weights->contiguous || !is_shape(*route_weights, 1, kRoutedFfnRouteCount, routed_output->ne[2], 1) || + weighted_output == nullptr || !same_shape(*weighted_output, *routed_output)) { + return {}; + } + + const int64_t token_count = routed_output->ne[2]; + if (!is_llm_supported_query_length(kRoutedFfnProfile, token_count)) { + return {}; + } + + std::vector views = + layout_alias_consumers_with_op(context.graph, weighted->output, GGML_OP_VIEW); + if (views.size() != kRoutedFfnRouteCount) { + return {}; + } + + std::set routed_values; + std::vector owned_views; + for (const GraphNode * view : views) { + const Value * value = view == nullptr ? nullptr : graph_value(context.graph, view->output); + if (value == nullptr || value->type != GGML_TYPE_F32 || + !is_shape(*value, kRoutedFfnInputSize, token_count, 1, 1) || + !append_node_if_uncovered(context, view, owned_views)) { + return {}; + } + routed_values.insert(view->output.value); + } + + std::vector reductions; + bool changed = true; + while (changed) { + changed = false; + const std::set values = routed_values; + for (int32_t value : values) { + for (const GraphNode * add : find_consumers_with_op(context.graph, ValueId(value), GGML_OP_ADD)) { + if (add == nullptr || add->inputs.size() != 2) { + continue; + } + bool already_owned = false; + for (const GraphNode * reduction : reductions) { + if (reduction == add) { + already_owned = true; + break; + } + } + if (already_owned) { + continue; + } + bool all_routed = true; + for (ValueId input : add->inputs) { + all_routed = all_routed && routed_values.count(input.value) != 0; + } + if (!all_routed) { + continue; + } + const Value * output = graph_value(context.graph, add->output); + if (output == nullptr || output->type != GGML_TYPE_F32 || + !is_shape(*output, kRoutedFfnInputSize, token_count, 1, 1) || + !append_node_if_uncovered(context, add, reductions)) { + return {}; + } + routed_values.insert(add->output.value); + changed = true; + } + } + } + + const GraphNode * residual = nullptr; + const Value * residual_input = nullptr; + for (int32_t value : routed_values) { + for (const GraphNode * add : find_consumers_with_op(context.graph, ValueId(value), GGML_OP_ADD)) { + if (add == nullptr || add->inputs.size() != 2) { + continue; + } + bool is_reduction = false; + for (const GraphNode * reduction : reductions) { + if (reduction == add) { + is_reduction = true; + break; + } + } + if (is_reduction) { + continue; + } + int routed_input_count = 0; + const Value * non_routed_input = nullptr; + for (ValueId input : add->inputs) { + if (routed_values.count(input.value) != 0) { + ++routed_input_count; + } else { + non_routed_input = graph_value(context.graph, input); + } + } + if (routed_input_count != 1 || non_routed_input == nullptr || residual != nullptr) { + return {}; + } + residual = add; + residual_input = non_routed_input; + } + } + if (residual == nullptr || reductions.size() + 1 != views.size()) { + return {}; + } + const Value * output = graph_value(context.graph, residual->output); + if (output == nullptr || residual_input == nullptr || output->type != GGML_TYPE_F32 || + residual_input->type != GGML_TYPE_F32 || !is_shape(*output, kRoutedFfnInputSize, token_count, 1, 1) || + !same_shape(*output, *residual_input) || output->byte_count != residual_input->byte_count || + !output->contiguous || !residual_input->contiguous) { + return {}; + } + + match.route_weights = route_weights; + match.routed_output = routed_output; + match.residual_input = residual_input; + match.output = output; + match.weighted_node = weighted; + match.views = std::move(owned_views); + match.reductions = std::move(reductions); + match.residual = residual; + match.next_rmsnorm = match_qwen_weighted_reduce_next_rmsnorm(context, *output); + match.token_count = token_count; + return match; +} + +static WeightedReduceMatch match_routed_ffn_down_weighted_reduce(const DispatchMatchContext & context) { + WeightedReduceMatch match; + const GraphNode * weighted = context.root_node; + if (weighted == nullptr || weighted->op != GGML_OP_MUL || weighted->inputs.size() != 2 || + !context.graph.has_index()) { + return match; + } + + const Value * routed_output = nullptr; + const Value * route_weights = nullptr; + for (ValueId input : weighted->inputs) { + const Value * value = graph_value(context.graph, input); + if (value == nullptr) { + return {}; + } + if (is_routed_ffn_down_output(*value, value->ne[2])) { + routed_output = value; + } else if (value->type == GGML_TYPE_F32 && value->contiguous && + is_shape(*value, 1, kRoutedFfnRouteCount, value->ne[2], 1)) { + route_weights = value; + } + } + match = match_routed_ffn_down_weighted_reduce_topology(context, weighted, routed_output, route_weights); + if (!match.topology_matched()) { + return {}; + } + const int64_t token_count = match.token_count; + const CommandPlanAlternateValue * routed_alternate = + find_alternate_value(context.plan, routed_output->id, GGML_TYPE_F16, f16_routed_down_output_size(token_count)); + if (routed_alternate == nullptr) { + return {}; + } + + bool known_route_weights = false; + for (const CommandPlanMoeRoutingBundle & bundle : context.plan.metadata.moe_routing_bundles()) { + if (bundle.route_weights == route_weights->id && + bundle_matches_moe_routing(bundle, bundle.route_ids, token_count)) { + known_route_weights = true; + break; + } + } + if (!known_route_weights) { + return {}; + } + + match.routed_alternate = routed_alternate; + return match; +} + +static bool match_routed_ffn_gate_up_swiglu_q4k_f16_wmma_dispatch(const DispatchMatchContext & context, + DispatchMatch & dispatch_match) { + const RoutedGateUpMatch match = match_routed_ffn_gate_up_swiglu(context); + if (!match.matched()) { + return false; + } + + const ValueId f16_output(context.next_plan_value.value); + const size_t f16_output_bytes = f16_gate_up_output_size(match.token_count); + dispatch_match.transients.push_back( + { f16_output, kRoutedFfnF16GateUpOutputName, f16_output_bytes, kRoutedFfnPlanTransientAlignment }); + Status metadata_status; + if (!dispatch_match.metadata.append_alternate_value( + { match.glu_output->id, f16_output, GGML_TYPE_F16, f16_output_bytes, kRoutedFfnF16GateUpOutputName }, + metadata_status)) { + return false; + } + + Dispatch dispatch; + dispatch.kernel = make_kernel_specialization(kQwenRoutedGateUpSwiGLUQ4KF16WmmaKernel); + dispatch.kernel.integer_parameters.emplace("token_count", match.token_count); + dispatch.kernel.compile_parameters.emplace("qwen3_moe.routed_gate_up.input_size", + to_config_value(kRoutedFfnInputSize)); + dispatch.kernel.compile_parameters.emplace("qwen3_moe.routed_gate_up.expert_count", + to_config_value(kRoutedFfnExpertCount)); + dispatch.kernel.compile_parameters.emplace("qwen3_moe.routed_gate_up.route_count", + to_config_value(kRoutedFfnRouteCount)); + dispatch.kernel.compile_parameters.emplace("qwen3_moe.routed_gate_up.output_size", + to_config_value(kRoutedFfnExpertHiddenSize)); + dispatch.kernel.compile_parameters.emplace("qwen3_moe.workload.token_capacity", to_config_value(match.token_count)); + + dispatch.bindings.push_back({ match.input->id, 0, match.input->byte_count }); + dispatch.bindings.push_back( + { match.routing_bundle->expert_table, 0, match.routing_bundle->expert_table_byte_count }); + dispatch.bindings.push_back( + { match.routing_bundle->partition_table, 0, match.routing_bundle->partition_table_byte_count }); + dispatch.bindings.push_back({ match.gate_weight->id, 0, match.gate_weight->byte_count }); + dispatch.bindings.push_back({ match.up_weight->id, 0, match.up_weight->byte_count }); + dispatch.bindings.push_back({ f16_output, 0, f16_output_bytes }); + + if (!append_covered_node(context, match.gate_node, dispatch_match) || + !append_covered_node(context, match.up_node, dispatch_match) || + !append_covered_node(context, match.glu_node, dispatch_match)) { + return false; + } + dispatch_match.dispatches.push_back(std::move(dispatch)); + return true; +} + +static bool match_decode_routed_ffn_gate_up_swiglu_q4k_q8_dispatch(const DispatchMatchContext & context, + DispatchMatch & dispatch_match) { + const DecodeRoutedGateUpMatch match = match_decode_routed_ffn_gate_up_swiglu(context); + if (!match.matched()) { + return false; + } + + const size_t q8_output_bytes = + q8_1_x4_byte_count(match.token_count * kRoutedFfnRouteCount, kRoutedFfnExpertHiddenSize); + const ValueId q8_output = match.publish_q8 ? ValueId(context.next_plan_value.value) : ValueId(); + const ValueId completion_counters = match.publish_q8 ? ValueId(context.next_plan_value.value + 1) : ValueId(); + + Dispatch dispatch; + dispatch.kernel = make_kernel_specialization(match.publish_q8 ? kQwenRoutedGateUpSwiGLUQ4KQ8NextQ8Kernel : + kQwenRoutedGateUpSwiGLUQ4KQ8Kernel); + dispatch.kernel.integer_parameters.emplace("token_count", match.token_count); + dispatch.kernel.integer_parameters.emplace("route_count", kRoutedFfnRouteCount); + dispatch.kernel.integer_parameters.emplace("route_stride", match.route_stride); + dispatch.kernel.integer_parameters.emplace("expert_count", kRoutedFfnExpertCount); + dispatch.kernel.integer_parameters.emplace("output_size", kRoutedFfnExpertHiddenSize); + dispatch.kernel.compile_parameters.emplace("qwen3_moe.routed_gate_up.input_size", + to_config_value(kRoutedFfnInputSize)); + dispatch.kernel.compile_parameters.emplace("qwen3_moe.routed_gate_up.expert_count", + to_config_value(kRoutedFfnExpertCount)); + dispatch.kernel.compile_parameters.emplace("qwen3_moe.routed_gate_up.route_count", + to_config_value(kRoutedFfnRouteCount)); + dispatch.kernel.compile_parameters.emplace("qwen3_moe.routed_gate_up.output_size", + to_config_value(kRoutedFfnExpertHiddenSize)); + dispatch.kernel.compile_parameters.emplace("qwen3_moe.workload.token_capacity", to_config_value(match.token_count)); + + const size_t route_id_length = static_cast(match.token_count * match.route_stride) * sizeof(int32_t); + dispatch.bindings.push_back({ match.input_alternate->alternate_value, 0, match.input_alternate->byte_count }); + dispatch.bindings.push_back({ match.route_ids->id, 0, route_id_length }); + dispatch.bindings.push_back({ match.gate_weight->id, 0, match.gate_weight->byte_count }); + dispatch.bindings.push_back({ match.up_weight->id, 0, match.up_weight->byte_count }); + dispatch.bindings.push_back({ match.glu_output->id, 0, match.glu_output->byte_count }); + if (match.publish_q8) { + dispatch.bindings.push_back( + { completion_counters, 0, gate_up_completion_counter_count(match.token_count) * sizeof(int32_t) }); + dispatch.bindings.push_back({ q8_output, 0, q8_output_bytes }); + dispatch_match.completion_counter_requests.push_back({ + completion_counters, + "qwen.decode.moe.gate_up_completion_counters", + gate_up_completion_counter_count(match.token_count), + }); + dispatch_match.transients.push_back( + { q8_output, kRoutedFfnQ8GateUpOutputName, q8_output_bytes, kRoutedFfnPlanTransientAlignment }); + Status metadata_status; + if (!dispatch_match.metadata.append_alternate_value( + { match.glu_output->id, q8_output, GGML_TYPE_Q8_1, q8_output_bytes, kRoutedFfnQ8GateUpOutputName }, + metadata_status)) { + dispatch_match.status.append(metadata_status); + return false; + } + } + + if (!append_covered_node(context, match.gate_node, dispatch_match) || + !append_covered_node(context, match.up_node, dispatch_match) || + !append_covered_node(context, match.glu_node, dispatch_match)) { + return false; + } + dispatch_match.dispatches.push_back(std::move(dispatch)); + return true; +} + +static DecodeRoutedDownMatch match_decode_routed_ffn_down_next_q8(const DispatchMatchContext & context) { + DecodeRoutedDownMatch match; + const GraphNode * root = context.root_node; + if (root == nullptr || root->op != GGML_OP_MUL_MAT_ID || root->inputs.size() != 3 || !context.graph.has_index()) { + return match; + } + + const Value * weight = graph_value(context.graph, root->inputs[0]); + const Value * input = graph_value(context.graph, root->inputs[1]); + const Value * route_ids = graph_value(context.graph, root->inputs[2]); + const Value * root_output = graph_value(context.graph, root->output); + if (weight == nullptr || input == nullptr || route_ids == nullptr || root_output == nullptr || + !is_routed_ffn_down_weight(*weight) || !is_routed_ffn_projection_output(*input, 1) || + !is_routed_ffn_down_output(*root_output, 1) || route_ids->type != GGML_TYPE_I32 || + route_ids->nb[0] != sizeof(int32_t) || route_ids->nb[1] % sizeof(int32_t) != 0 || + !is_shape(*route_ids, kRoutedFfnRouteCount, 1, 1, 1)) { + return {}; + } + + const int64_t route_stride = static_cast(route_ids->nb[1] / sizeof(int32_t)); + const GraphNode * glu_node = producer_with_op(context.graph, input->id, GGML_OP_GLU); + const GraphNode * gate_node = glu_node == nullptr || glu_node->inputs.size() != 2 ? + nullptr : + producer_with_op(context.graph, glu_node->inputs[0], GGML_OP_MUL_MAT_ID); + const Value * gate_input = gate_node == nullptr || gate_node->inputs.size() != 3 ? + nullptr : + graph_value(context.graph, gate_node->inputs[1]); + const CommandPlanAlternateValue * gate_input_q8 = + gate_input == nullptr ? nullptr : + find_alternate_value(context.graph, context.plan, gate_input->id, GGML_TYPE_Q8_1, + q8_1_x4_byte_count(1, kRoutedFfnInputSize)); + if (gate_input_q8 == nullptr) { + return {}; + } + + const Value * route_weights = find_qwen_route_weights_for_route_ids(context.graph, route_ids->id, 1); + const GraphNode * weighted = + find_single_consumer_with_op_through_layout_aliases(context.graph, root_output->id, GGML_OP_MUL); + WeightedReduceMatch reduce = + match_routed_ffn_down_weighted_reduce_topology(context, weighted, root_output, route_weights); + if (!reduce.topology_matched() || !reduce.next_rmsnorm.matched() || + !residual_input_is_safe_for_in_place(context, reduce)) { + return {}; + } + + const bool input_is_q8 = weight->type == GGML_TYPE_Q4_K; + const CommandPlanAlternateValue * input_alternate = + input_is_q8 ? find_alternate_value(context.graph, context.plan, input->id, GGML_TYPE_Q8_1, + q8_1_x4_byte_count(kRoutedFfnRouteCount, kRoutedFfnExpertHiddenSize)) : + nullptr; + if (input_is_q8 && input_alternate == nullptr) { + return {}; + } + + match.input_graph_value = input; + match.input_alternate = input_alternate; + match.weight = weight; + match.output = reduce.output; + match.route_ids = route_ids; + match.reduce = std::move(reduce); + match.kernel = + weight->type == GGML_TYPE_Q4_K ? kQwenRoutedDownQ4KQ8NextQ8Kernel : kQwenRoutedDownQ6KF32Wave64NextQ8Kernel; + match.token_count = 1; + match.route_stride = route_stride; + match.input_is_q8 = input_is_q8; + return match; +} + +static bool match_decode_routed_ffn_down_next_q8_dispatch(const DispatchMatchContext & context, + DispatchMatch & dispatch_match) { + const DecodeRoutedDownMatch match = match_decode_routed_ffn_down_next_q8(context); + if (!match.matched()) { + return false; + } + + const ValueId completion_counter_value(context.next_plan_value.value); + const ValueId q8_output(context.next_plan_value.value + 1); + const size_t q8_output_bytes = q8_1_x4_byte_count(match.token_count, kRoutedFfnInputSize); + + Dispatch dispatch; + dispatch.kernel = make_kernel_specialization(match.kernel); + dispatch.kernel.integer_parameters.emplace("token_count", match.token_count); + dispatch.kernel.integer_parameters.emplace("input_size", kRoutedFfnExpertHiddenSize); + dispatch.kernel.integer_parameters.emplace("route_count", kRoutedFfnRouteCount); + dispatch.kernel.integer_parameters.emplace("route_id_stride", match.route_stride); + dispatch.kernel.integer_parameters.emplace("expert_count", kRoutedFfnExpertCount); + dispatch.kernel.integer_parameters.emplace("output_size", kRoutedFfnInputSize); + add_routed_down_compile_parameters(dispatch, match.token_count); + dispatch.kernel.compile_parameters.emplace("qwen3_moe.model.hidden_size", to_config_value(kRoutedFfnInputSize)); + dispatch.kernel.compile_parameters.emplace("qwen3_moe.model.rms_epsilon", "0.000001"); + + const size_t route_id_length = static_cast(match.token_count * match.route_stride) * sizeof(int32_t); + if (match.input_is_q8) { + dispatch.bindings.push_back({ match.input_alternate->alternate_value, 0, match.input_alternate->byte_count }); + } else { + dispatch.bindings.push_back({ match.input_graph_value->id, 0, match.input_graph_value->byte_count }); + } + dispatch.bindings.push_back({ match.route_ids->id, 0, route_id_length }); + dispatch.bindings.push_back({ match.reduce.route_weights->id, 0, match.reduce.route_weights->byte_count }); + dispatch.bindings.push_back({ match.weight->id, 0, match.weight->byte_count }); + dispatch.bindings.push_back({ match.output->id, 0, match.output->byte_count }); + dispatch.bindings.push_back( + { match.reduce.next_rmsnorm.norm_weight->id, 0, match.reduce.next_rmsnorm.norm_weight->byte_count }); + dispatch.bindings.push_back({ completion_counter_value, 0, sizeof(int32_t) }); + dispatch.bindings.push_back({ q8_output, 0, q8_output_bytes }); + + dispatch_match.value_aliases.push_back({ match.reduce.residual_input->id, match.output->id }); + dispatch_match.completion_counter_requests.push_back({ + completion_counter_value, + "qwen.decode.moe.routed_down_completion_counter", + 1, + }); + dispatch_match.transients.push_back( + { q8_output, kRoutedFfnQ8HiddenOutputName, q8_output_bytes, kRoutedFfnPlanTransientAlignment }); + Status metadata_status; + if (!dispatch_match.metadata.append_alternate_value( + { match.reduce.next_rmsnorm.output->id, q8_output, GGML_TYPE_Q8_1, q8_output_bytes, + kRoutedFfnQ8HiddenOutputName }, + metadata_status)) { + dispatch_match.status.append(metadata_status); + return false; + } + + if (!append_covered_node(context, context.root_node, dispatch_match) || + !append_covered_node(context, match.reduce.weighted_node, dispatch_match)) { + return false; + } + for (const GraphNode * view : match.reduce.views) { + if (!append_covered_node(context, view, dispatch_match)) { + return false; + } + } + for (const GraphNode * reduction : match.reduce.reductions) { + if (!append_covered_node(context, reduction, dispatch_match)) { + return false; + } + } + if (!append_covered_node(context, match.reduce.residual, dispatch_match) || + !append_covered_node(context, match.reduce.next_rmsnorm.rms_node, dispatch_match) || + !append_covered_node(context, match.reduce.next_rmsnorm.mul_node, dispatch_match)) { + return false; + } + + dispatch_match.dispatches.push_back(std::move(dispatch)); + return true; +} + +static bool build_routed_ffn_down_grouped_dispatch(const DispatchMatchContext & context, + DispatchMatch & dispatch_match, + KernelCatalogRef expected_kernel) { + const RoutedDownMatch match = match_routed_ffn_down_grouped(context); + if (!match.matched() || match.kernel.id != expected_kernel.id) { + return false; + } + + const ValueId f16_output(context.next_plan_value.value); + const size_t f16_output_bytes = f16_routed_down_output_size(match.token_count); + dispatch_match.transients.push_back( + { f16_output, kRoutedFfnF16RoutedDownOutputName, f16_output_bytes, kRoutedFfnPlanTransientAlignment }); + Status metadata_status; + if (!dispatch_match.metadata.append_alternate_value( + { match.output->id, f16_output, GGML_TYPE_F16, f16_output_bytes, kRoutedFfnF16RoutedDownOutputName }, + metadata_status)) { + return false; + } + + Dispatch dispatch; + dispatch.kernel = make_kernel_specialization(match.kernel); + dispatch.kernel.integer_parameters.emplace("token_count", match.token_count); + add_routed_down_compile_parameters(dispatch, match.token_count); + dispatch.bindings.push_back({ match.input_alternate->alternate_value, 0, match.input_alternate->byte_count }); + dispatch.bindings.push_back( + { match.routing_bundle->expert_table, 0, match.routing_bundle->expert_table_byte_count }); + dispatch.bindings.push_back({ match.weight->id, 0, match.weight->byte_count }); + dispatch.bindings.push_back({ f16_output, 0, f16_output_bytes }); + + dispatch_match.covered_nodes.push_back(context.root_index); + dispatch_match.dispatches.push_back(std::move(dispatch)); + return true; +} + +static bool match_routed_ffn_down_weighted_reduce_dispatch(const DispatchMatchContext & context, + DispatchMatch & dispatch_match) { + const WeightedReduceMatch match = match_routed_ffn_down_weighted_reduce(context); + if (!match.matched()) { + return false; + } + const bool use_next_rmsnorm = match.next_rmsnorm.matched() && match.output->kind == ValueKind::Transient && + residual_input_is_safe_for_in_place(context, match); + + Dispatch dispatch; + dispatch.kernel = make_kernel_specialization(use_next_rmsnorm ? kQwenRoutedDownWeightedReduceNextRmsNormF32Kernel : + kQwenRoutedDownWeightedReduceF16F32Kernel); + dispatch.kernel.integer_parameters.emplace("token_count", match.token_count); + add_routed_down_compile_parameters(dispatch, match.token_count); + if (use_next_rmsnorm) { + dispatch_match.value_aliases.push_back({ match.residual_input->id, match.output->id }); + dispatch.kernel.compile_parameters.emplace("qwen3_moe.model.hidden_size", to_config_value(kRoutedFfnInputSize)); + dispatch.kernel.compile_parameters.emplace("qwen3_moe.model.rms_epsilon", "0.000001"); + dispatch.bindings.push_back({ match.route_weights->id, 0, match.route_weights->byte_count }); + dispatch.bindings.push_back({ match.routed_alternate->alternate_value, 0, match.routed_alternate->byte_count }); + dispatch.bindings.push_back({ match.output->id, 0, match.output->byte_count }); + dispatch.bindings.push_back( + { match.next_rmsnorm.norm_weight->id, 0, match.next_rmsnorm.norm_weight->byte_count }); + dispatch.bindings.push_back({ match.next_rmsnorm.output->id, 0, match.next_rmsnorm.output->byte_count }); + } else { + dispatch.bindings.push_back({ match.route_weights->id, 0, match.route_weights->byte_count }); + dispatch.bindings.push_back({ match.routed_alternate->alternate_value, 0, match.routed_alternate->byte_count }); + dispatch.bindings.push_back({ match.output->id, 0, match.output->byte_count }); + } + + if (!append_covered_node(context, match.weighted_node, dispatch_match)) { + return false; + } + for (const GraphNode * view : match.views) { + if (!append_covered_node(context, view, dispatch_match)) { + return false; + } + } + for (const GraphNode * reduction : match.reductions) { + if (!append_covered_node(context, reduction, dispatch_match)) { + return false; + } + } + if (!append_covered_node(context, match.residual, dispatch_match)) { + return false; + } + if (use_next_rmsnorm && (!append_covered_node(context, match.next_rmsnorm.rms_node, dispatch_match) || + !append_covered_node(context, match.next_rmsnorm.mul_node, dispatch_match))) { + return false; + } + + dispatch_match.dispatches.push_back(std::move(dispatch)); + return true; +} + +static bool match_routed_ffn_down_q4k_f16_wmma_grouped_dispatch(const DispatchMatchContext & context, + DispatchMatch & dispatch_match) { + return build_routed_ffn_down_grouped_dispatch(context, dispatch_match, kQwenRoutedDownQ4KF16WmmaGroupedKernel); +} + +static bool match_routed_ffn_down_q6k_f16_wmma_grouped_dispatch(const DispatchMatchContext & context, + DispatchMatch & dispatch_match) { + return build_routed_ffn_down_grouped_dispatch(context, dispatch_match, kQwenRoutedDownQ6KF16WmmaGroupedKernel); +} + +} // namespace + +void register_routed_ffn_dispatches(DispatchRegistryBuilder & registry) { + registry.add({ + "llm.routed_ffn.decode_gate_up_swiglu_q4k_q8", + GGML_OP_MUL_MAT_ID, + DispatchMatchKind::Fused, + 1200, + DispatchSource::Llm, + match_decode_routed_ffn_gate_up_swiglu_q4k_q8_dispatch, + }); + registry.add({ + "llm.routed_ffn.decode_down_next_q8", + GGML_OP_MUL_MAT_ID, + DispatchMatchKind::Fused, + 1150, + DispatchSource::Llm, + match_decode_routed_ffn_down_next_q8_dispatch, + }); + registry.add({ + "llm.routed_ffn.gate_up_swiglu_q4k_f16_wmma", + GGML_OP_MUL_MAT_ID, + DispatchMatchKind::Fused, + 1000, + DispatchSource::Llm, + match_routed_ffn_gate_up_swiglu_q4k_f16_wmma_dispatch, + }); + registry.add({ + "llm.routed_ffn.down_q4k_f16_wmma_grouped", + GGML_OP_MUL_MAT_ID, + DispatchMatchKind::Fused, + 900, + DispatchSource::Llm, + match_routed_ffn_down_q4k_f16_wmma_grouped_dispatch, + }); + registry.add({ + "llm.routed_ffn.down_q6k_f16_wmma_grouped", + GGML_OP_MUL_MAT_ID, + DispatchMatchKind::Fused, + 900, + DispatchSource::Llm, + match_routed_ffn_down_q6k_f16_wmma_grouped_dispatch, + }); + registry.add({ + "llm.routed_ffn.down_weighted_reduce", + GGML_OP_MUL, + DispatchMatchKind::Fused, + 800, + DispatchSource::Llm, + match_routed_ffn_down_weighted_reduce_dispatch, + }); +} + +} // namespace ggml::hrx diff --git a/ggml/src/ggml-hrx/dispatch_registration/dispatch-routed-ffn.h b/ggml/src/ggml-hrx/dispatch_registration/dispatch-routed-ffn.h new file mode 100644 index 000000000000..18cdabc5b51f --- /dev/null +++ b/ggml/src/ggml-hrx/dispatch_registration/dispatch-routed-ffn.h @@ -0,0 +1,9 @@ +#pragma once + +#include "dispatch-registry.h" + +namespace ggml::hrx { + +void register_routed_ffn_dispatches(DispatchRegistryBuilder & registry); + +} // namespace ggml::hrx diff --git a/ggml/src/ggml-hrx/ggml-hrx.cpp b/ggml/src/ggml-hrx/ggml-hrx.cpp new file mode 100644 index 000000000000..6263e48e3dc6 --- /dev/null +++ b/ggml/src/ggml-hrx/ggml-hrx.cpp @@ -0,0 +1,758 @@ +// Copyright 2026 The HRX Authors +// SPDX-License-Identifier: Apache-2.0 + +#include "ggml-hrx.h" + +#include "backend-buffer-binding.h" +#include "backend-context.h" +#include "ggml-backend-impl.h" +#include "ggml-impl.h" +#include "hrx_runtime.h" +#include "kernel-corpus/kernel-corpus.h" +#include "loom-jit.h" +#include "runtime/graph-executor.h" +#include "runtime/graph-program-cache.h" +#include "runtime/kernel-executable-cache.h" +#include "runtime/prepared-command-program-cache.h" +#include "runtime/transient-arena.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace { + +static constexpr size_t GGML_HRX_ALIGNMENT = 256; +// GGML represents tensor locations as host pointers and derives view/arena offsets with ordinary pointer arithmetic. +// Device-local HRX buffers have no host address to return, so expose a non-null sentinel base as an offset coordinate. +static constexpr uintptr_t GGML_HRX_FAKE_PTR_BASE = 0x1000; +static std::atomic g_allocation_generation{ 1 }; + +static bool environment_flag_enabled(const char * name) { + const char * value = std::getenv(name); + return value != nullptr && value[0] != '\0' && value[0] != '0'; +} + +static bool hrx_check(hrx_status_t status, const char * expression, const char * file, int line) { + if (hrx_status_is_ok(status)) { + return true; + } + char * message = nullptr; + size_t length = 0; + hrx_status_to_string(status, &message, &length); + GGML_LOG_ERROR("%s:%d: %s failed: %s\n", file, line, expression, + message != nullptr ? message : "unknown HRX error"); + hrx_status_free_message(message); + hrx_status_ignore(status); + return false; +} + +#define HRX_CHECK(expression) hrx_check((expression), #expression, __FILE__, __LINE__) + +} // namespace + +ggml_backend_hrx_reg_context::~ggml_backend_hrx_reg_context() { + for (auto & context : device_contexts) { + if (context->buffer_stream != nullptr) { + hrx_stream_release(context->buffer_stream); + } + if (context->device != nullptr) { + hrx_device_release(context->device); + } + } + if (initialized) { + hrx_status_t status = hrx_gpu_shutdown(); + if (!hrx_status_is_ok(status)) { + hrx_status_ignore(status); + } + } +} + +namespace { + +static std::optional device_string_property(hrx_device_t device, + hrx_device_property_t property, + const char * property_name) { + std::vector buffer(64); + while (buffer.size() <= 4096) { + hrx_status_t status = hrx_device_get_property(device, property, buffer.data(), buffer.size()); + if (hrx_status_is_ok(status)) { + return std::string(buffer.data()); + } + if (hrx_status_code(status) != HRX_STATUS_OUT_OF_RANGE) { + hrx_check(status, property_name, __FILE__, __LINE__); + return std::nullopt; + } + hrx_status_ignore(status); + buffer.resize(buffer.size() * 2); + } + GGML_LOG_ERROR("%s exceeds the maximum supported property string length\n", property_name); + return std::nullopt; +} + +static ggml_guid_t ggml_backend_hrx_guid() { + static ggml_guid guid = { + 0xd2, 0x3d, 0x72, 0x83, 0xb2, 0x82, 0x4d, 0xe0, 0x8a, 0x3e, 0x21, 0x1d, 0x68, 0x87, 0x2f, 0x4b, + }; + return &guid; +} + +static ggml_backend_hrx_device_context * device_context(ggml_backend_dev_t device) { + return static_cast(device->context); +} + +static ggml_backend_hrx_buffer_context * buffer_context(ggml_backend_buffer_t buffer) { + return ggml_backend_hrx_buffer_context_from_buffer(buffer); +} + +static size_t tensor_offset(const ggml_backend_hrx_buffer_context * context, const ggml_tensor * tensor) { + return ggml_backend_hrx_tensor_offset(context, tensor); +} + +static const char * buffer_type_name(ggml_backend_buffer_type_t buft) { + return static_cast(buft->context)->name.c_str(); +} + +static bool buffer_type_is_host(ggml_backend_buffer_type_t buft) { + return static_cast(buft->context)->host_visible; +} + +static bool buffer_submit_and_wait(ggml_backend_hrx_device_context * device, + hrx_status_t (*submit)(hrx_stream_t, void *), + void * user_data) { + std::lock_guard lock(device->buffer_stream_mutex); + if (!HRX_CHECK(submit(device->buffer_stream, user_data))) { + return false; + } + return HRX_CHECK(hrx_stream_synchronize(device->buffer_stream)); +} + +struct FillBufferArgs { + hrx_buffer_t buffer; + size_t offset; + size_t size; + uint8_t value; +}; + +static hrx_status_t submit_fill_buffer(hrx_stream_t stream, void * user_data) { + auto * args = static_cast(user_data); + return hrx_stream_fill_buffer(stream, args->buffer, args->offset, args->size, &args->value, sizeof(args->value)); +} + +struct CopyBufferArgs { + hrx_buffer_t source; + size_t source_offset; + hrx_buffer_t destination; + size_t destination_offset; + size_t size; +}; + +static hrx_status_t submit_copy_buffer(hrx_stream_t stream, void * user_data) { + auto * args = static_cast(user_data); + return hrx_stream_copy_buffer(stream, args->source, args->source_offset, args->destination, + args->destination_offset, args->size); +} + +static void buffer_free(ggml_backend_buffer_t buffer) { + auto * context = buffer_context(buffer); + if (context->base != reinterpret_cast(GGML_HRX_FAKE_PTR_BASE)) { + context->device->host_buffers.remove(context->buffer); + } + if (context->buffer != nullptr) { + hrx_buffer_release(context->buffer); + } + delete context; +} + +static void buffer_memset(ggml_backend_buffer_t buffer, + ggml_tensor * tensor, + uint8_t value, + size_t offset, + size_t size) { + if (size == 0) { + return; + } + auto * context = buffer_context(buffer); + const size_t destination_offset = tensor_offset(context, tensor) + offset; + GGML_ASSERT(destination_offset <= buffer->size && size <= buffer->size - destination_offset); + if (context->base != reinterpret_cast(GGML_HRX_FAKE_PTR_BASE)) { + std::memset(context->base + destination_offset, value, size); + return; + } + FillBufferArgs args{ context->buffer, destination_offset, size, value }; + if (!buffer_submit_and_wait(context->device, submit_fill_buffer, &args)) { + GGML_LOG_ERROR("%s: HRX buffer fill failed\n", __func__); + } +} + +static void buffer_set(ggml_backend_buffer_t buffer, + ggml_tensor * tensor, + const void * data, + size_t offset, + size_t size) { + if (size == 0) { + return; + } + auto * context = buffer_context(buffer); + const size_t destination_offset = tensor_offset(context, tensor) + offset; + GGML_ASSERT(destination_offset <= buffer->size && size <= buffer->size - destination_offset); + if (context->base != reinterpret_cast(GGML_HRX_FAKE_PTR_BASE)) { + std::memcpy(context->base + destination_offset, data, size); + return; + } + if (!HRX_CHECK(hrx_synchronous_h2d(context->device->device, data, context->buffer, destination_offset, size))) { + GGML_LOG_ERROR("%s: HRX buffer upload failed\n", __func__); + } +} + +static void buffer_get(ggml_backend_buffer_t buffer, + const ggml_tensor * tensor, + void * data, + size_t offset, + size_t size) { + if (size == 0) { + return; + } + auto * context = buffer_context(buffer); + const size_t source_offset = tensor_offset(context, tensor) + offset; + GGML_ASSERT(source_offset <= buffer->size && size <= buffer->size - source_offset); + if (context->base != reinterpret_cast(GGML_HRX_FAKE_PTR_BASE)) { + std::memcpy(data, context->base + source_offset, size); + return; + } + if (!HRX_CHECK(hrx_synchronous_d2h(context->device->device, context->buffer, source_offset, data, size))) { + GGML_LOG_ERROR("%s: HRX buffer download failed\n", __func__); + } +} + +static bool buffer_copy(ggml_backend_buffer_t buffer, const ggml_tensor * source, ggml_tensor * destination) { + ggml_backend_buffer_t source_buffer = source->view_src != nullptr ? source->view_src->buffer : source->buffer; + if (source_buffer == nullptr || source_buffer->iface.get_base != ggml_backend_hrx_buffer_base) { + return false; + } + auto * source_context = buffer_context(source_buffer); + auto * destination_context = buffer_context(buffer); + if (source_context->device != destination_context->device) { + return false; + } + const size_t source_offset = tensor_offset(source_context, source); + const size_t destination_offset = tensor_offset(destination_context, destination); + const size_t size = ggml_nbytes(source); + if (source_offset > source_buffer->size || size > source_buffer->size - source_offset || + destination_offset > buffer->size || size > buffer->size - destination_offset) { + return false; + } + CopyBufferArgs args{ source_context->buffer, source_offset, destination_context->buffer, destination_offset, size }; + return buffer_submit_and_wait(destination_context->device, submit_copy_buffer, &args); +} + +static void buffer_clear(ggml_backend_buffer_t buffer, uint8_t value) { + if (buffer->size == 0) { + return; + } + auto * context = buffer_context(buffer); + if (context->base != reinterpret_cast(GGML_HRX_FAKE_PTR_BASE)) { + std::memset(context->base, value, buffer->size); + return; + } + FillBufferArgs args{ context->buffer, 0, buffer->size, value }; + if (!buffer_submit_and_wait(context->device, submit_fill_buffer, &args)) { + GGML_LOG_ERROR("%s: HRX buffer clear failed\n", __func__); + } +} + +static const ggml_backend_buffer_i buffer_i = { + buffer_free, ggml_backend_hrx_buffer_base, + nullptr, buffer_memset, + buffer_set, buffer_get, + nullptr, nullptr, + buffer_copy, buffer_clear, + nullptr, +}; + +static ggml_backend_buffer_t buffer_alloc(ggml_backend_buffer_type_t buft, size_t size) { + auto * type_context = static_cast(buft->context); + const bool host_visible = type_context->host_visible; + const bool direct_host_binding = host_visible && type_context->device->use_direct_host_bindings; + hrx_memory_type_t memory_type = HRX_MEMORY_TYPE_DEVICE_LOCAL; + // Direct command-program bindings require coherent CPU/GPU visibility. Otherwise HRX host buffers are pinned + // transfer memory: DEVICE_VISIBLE permits handle-based stream copies without implying direct device access. + if (host_visible) { + memory_type = HRX_MEMORY_TYPE_HOST_LOCAL | HRX_MEMORY_TYPE_DEVICE_VISIBLE; + if (direct_host_binding) { + memory_type |= HRX_MEMORY_TYPE_HOST_COHERENT; + } + } + hrx_buffer_params_t params = { + memory_type, + HRX_MEMORY_ACCESS_ALL, + host_visible ? + HRX_BUFFER_USAGE_DEFAULT | HRX_BUFFER_USAGE_MAPPING_SCOPED | HRX_BUFFER_USAGE_MAPPING_PERSISTENT : + HRX_BUFFER_USAGE_DEFAULT, + 0, + }; + hrx_buffer_t allocation = nullptr; + if (size > 0 && !HRX_CHECK(hrx_allocator_allocate_buffer(hrx_device_allocator(type_context->device->device), params, + size, &allocation))) { + return nullptr; + } + uint8_t * base = reinterpret_cast(GGML_HRX_FAKE_PTR_BASE); + if (host_visible && size > 0) { + void * mapped = nullptr; + if (!HRX_CHECK(hrx_buffer_map(allocation, HRX_MAP_READ | HRX_MAP_WRITE, 0, size, &mapped))) { + hrx_buffer_release(allocation); + return nullptr; + } + base = static_cast(mapped); + } + const uint64_t generation = g_allocation_generation.fetch_add(1); + auto * context = new (std::nothrow) ggml_backend_hrx_buffer_context{ + type_context->device, allocation, base, generation, generation, direct_host_binding, + }; + if (context == nullptr) { + if (allocation != nullptr) { + hrx_buffer_release(allocation); + } + return nullptr; + } + if (host_visible && allocation != nullptr) { + type_context->device->host_buffers.add(allocation, base, size); + } + return ggml_backend_buffer_init(buft, buffer_i, context, size); +} + +static size_t buffer_alignment(ggml_backend_buffer_type_t buft) { + GGML_UNUSED(buft); + return GGML_HRX_ALIGNMENT; +} + +static size_t buffer_max_size(ggml_backend_buffer_type_t buft) { + return static_cast(buft->context)->device->memory_total; +} + +static const ggml_backend_buffer_type_i buffer_type_i = { + buffer_type_name, buffer_alloc, buffer_alignment, buffer_max_size, nullptr, buffer_type_is_host, +}; + +static const char * backend_name(ggml_backend_t backend) { + return static_cast(backend->context)->name.c_str(); +} + +static void backend_free(ggml_backend_t backend) { + auto * context = static_cast(backend->context); + HRX_CHECK(hrx_stream_synchronize(context->stream)); + context->prepared_programs.clear(); + context->graph_programs.clear(); + context->kernel_executables.clear(); + context->transient_arena.clear(); + context->host_weights.clear(); + context->host_transfers.clear(); + hrx_stream_release(context->stream); + delete context; + delete backend; +} + +static bool synchronous_upload_fallback(ggml_backend_hrx_context * backend, + const void * source, + hrx_buffer_t destination, + size_t destination_offset, + size_t size) { + const uint64_t fallback = + backend->device->synchronous_upload_fallbacks.fetch_add(1, std::memory_order_relaxed); + if (fallback == 0) { + GGML_LOG_WARN("ggml_hrx: synchronous upload fallback for an unregistered host pointer; use the HRX host " + "buffer type for asynchronous transfers\n"); + } + // Compatibility path for arbitrary GGML pointers. Keep the synchronization explicit until a bounded staging ring + // with transfer retirement is available. + const ggml::hrx::Status status = + backend->host_transfers.upload_synchronous(backend->stream, source, destination, destination_offset, size); + if (!status.success()) { + GGML_LOG_ERROR("%s: %s\n", __func__, status.errors().front().c_str()); + return false; + } + return true; +} + +static bool synchronous_download_fallback(ggml_backend_hrx_context * backend, + hrx_buffer_t source, + size_t source_offset, + void * destination, + size_t size) { + const uint64_t fallback = + backend->device->synchronous_download_fallbacks.fetch_add(1, std::memory_order_relaxed); + if (fallback == 0) { + GGML_LOG_WARN("ggml_hrx: synchronous download fallback for an unregistered host pointer; use the HRX host " + "buffer type for asynchronous transfers\n"); + } + // Compatibility path for arbitrary GGML pointers. Keep the synchronization explicit until a bounded staging ring + // with transfer retirement is available. + const ggml::hrx::Status status = + backend->host_transfers.download_synchronous(backend->stream, source, source_offset, destination, size); + if (!status.success()) { + GGML_LOG_ERROR("%s: %s\n", __func__, status.errors().front().c_str()); + return false; + } + return true; +} + +static void backend_set_tensor_async(ggml_backend_t backend, + ggml_tensor * tensor, + const void * data, + size_t offset, + size_t size) { + auto * backend_context = static_cast(backend->context); + ggml_backend_hrx_buffer_context * context = nullptr; + size_t tensor_base = 0; + if (!ggml_backend_hrx_tensor_binding(tensor, &context, &tensor_base) || offset > ggml_nbytes(tensor) || + size > ggml_nbytes(tensor) - offset) { + GGML_LOG_ERROR("%s: invalid HRX tensor upload\n", __func__); + return; + } + ggml::hrx::HostBufferRef source = backend_context->device->host_buffers.find(data, size); + if (source.valid()) { + // Registered host buffers can participate directly in the stream command buffer. + HRX_CHECK(hrx_stream_copy_buffer(backend_context->stream, source.buffer(), source.offset(), context->buffer, + tensor_base + offset, size)); + } else { + synchronous_upload_fallback(backend_context, data, context->buffer, tensor_base + offset, size); + } +} + +static void backend_get_tensor_async(ggml_backend_t backend, + const ggml_tensor * tensor, + void * data, + size_t offset, + size_t size) { + auto * backend_context = static_cast(backend->context); + ggml_backend_hrx_buffer_context * context = nullptr; + size_t tensor_base = 0; + if (!ggml_backend_hrx_tensor_binding(tensor, &context, &tensor_base) || offset > ggml_nbytes(tensor) || + size > ggml_nbytes(tensor) - offset) { + GGML_LOG_ERROR("%s: invalid HRX tensor download\n", __func__); + return; + } + ggml::hrx::HostBufferRef destination = backend_context->device->host_buffers.find(data, size); + if (destination.valid()) { + // Registered host buffers can participate directly in the stream command buffer. + HRX_CHECK(hrx_stream_copy_buffer(backend_context->stream, context->buffer, tensor_base + offset, + destination.buffer(), destination.offset(), size)); + } else { + synchronous_download_fallback(backend_context, context->buffer, tensor_base + offset, data, size); + } +} + +static bool backend_copy_tensor_async(ggml_backend_t backend_src, + ggml_backend_t backend_dst, + const ggml_tensor * source, + ggml_tensor * destination) { + GGML_UNUSED(backend_src); + auto * destination_backend = static_cast(backend_dst->context); + ggml_backend_hrx_buffer_context * destination_context = nullptr; + size_t destination_offset = 0; + if (!ggml_backend_hrx_tensor_binding(destination, &destination_context, &destination_offset)) { + return false; + } + ggml_backend_hrx_buffer_context * source_context = nullptr; + size_t source_offset = 0; + const size_t size = ggml_nbytes(source); + if (ggml_backend_hrx_tensor_binding(source, &source_context, &source_offset)) { + if (source_context->device != destination_context->device) { + return false; + } + return HRX_CHECK(hrx_stream_copy_buffer(destination_backend->stream, source_context->buffer, source_offset, + destination_context->buffer, destination_offset, size)); + } + ggml_backend_buffer_t source_buffer = source->view_src != nullptr ? source->view_src->buffer : source->buffer; + if (source_buffer != nullptr && ggml_backend_buffer_is_host(source_buffer)) { + return synchronous_upload_fallback( + destination_backend, source->data, destination_context->buffer, destination_offset, size); + } + return false; +} + +static void backend_synchronize(ggml_backend_t backend) { + auto * context = static_cast(backend->context); + HRX_CHECK(hrx_stream_synchronize(context->stream)); +} + +static const char * status_first_error(const ggml::hrx::Status & status) { + return status.errors().empty() ? "" : status.errors().front().c_str(); +} + +static enum ggml_status graph_compute(ggml_backend_t backend, ggml_cgraph * graph) { + auto * context = static_cast(backend->context); + const ggml::hrx::GraphExecutor executor = ggml::hrx::GraphExecutor(*context); + const ggml::hrx::GraphExecutionResult result = executor.execute(*graph); + if (!result.success()) { + GGML_LOG_ERROR("%s: %s\n", __func__, status_first_error(result.status)); + } + return result.code; +} + +static const ggml_backend_i backend_i = { + backend_name, + backend_free, + backend_set_tensor_async, + backend_get_tensor_async, + nullptr, + nullptr, + backend_copy_tensor_async, + backend_synchronize, + nullptr, + nullptr, + nullptr, + nullptr, + graph_compute, + nullptr, + nullptr, + nullptr, +}; + +static const char * device_name(ggml_backend_dev_t device) { + return device_context(device)->name.c_str(); +} + +static const char * device_description(ggml_backend_dev_t device) { + return device_context(device)->description.c_str(); +} + +static void device_memory(ggml_backend_dev_t device, size_t * free, size_t * total) { + *free = device_context(device)->memory_total; + *total = device_context(device)->memory_total; +} + +static enum ggml_backend_dev_type device_type(ggml_backend_dev_t device) { + GGML_UNUSED(device); + return GGML_BACKEND_DEVICE_TYPE_GPU; +} + +static void device_props(ggml_backend_dev_t device, ggml_backend_dev_props * props) { + props->name = device_name(device); + props->description = device_description(device); + device_memory(device, &props->memory_free, &props->memory_total); + props->type = GGML_BACKEND_DEVICE_TYPE_GPU; + props->device_id = nullptr; + props->caps = { true, true, false, false }; +} + +static ggml_backend_t device_init(ggml_backend_dev_t device, const char * parameters) { + GGML_UNUSED(parameters); + auto * device_ctx = device_context(device); + hrx_stream_t stream = nullptr; + if (!HRX_CHECK(hrx_stream_create(device_ctx->device, 0, &stream))) { + return nullptr; + } + auto * context = new (std::nothrow) ggml_backend_hrx_context; + if (context != nullptr) { + context->device = device_ctx; + context->stream = stream; + context->name = device_ctx->name; + } + auto * backend = context != nullptr ? new (std::nothrow) + ggml_backend{ ggml_backend_hrx_guid(), backend_i, device, context } : + nullptr; + if (backend == nullptr) { + delete context; + hrx_stream_release(stream); + } + return backend; +} + +static ggml_backend_buffer_type_t device_buffer_type(ggml_backend_dev_t device) { + // HRX weights/KV are placed in a host-visible buffer so that quant types and + // graph patterns the Loom corpus cannot dispatch can be executed by the CPU + // backend instead of fail-closing inside the HRX graph scheduler. + return &device_context(device)->host_buft; +} + +static ggml_backend_buffer_type_t device_host_buffer_type(ggml_backend_dev_t device) { + return &device_context(device)->host_buft; +} + +// The Loom-JIT backend is a fused-pattern dispatcher, and ggml_backend_sched +// assigns nodes per-op. Claim a node only when the dispatcher can actually execute +// it: either as a standalone dispatch or as the root of a fused pattern rooted at +// that node. Anything else is left to the CPU backend; the HRX buffer type is +// host-visible, so the CPU can read HRX tensors without a copy. +static bool device_supports_op(ggml_backend_dev_t device, const ggml_tensor * op) { + if (op == nullptr) { + return false; + } + auto * context = device_context(device); + if (op->op == GGML_OP_NONE) { + // Buffer-placement probe for a model weight or graph input. Any quant type + // may live in the host-visible HRX buffers. + return true; + } + return ggml::hrx::can_execute_standalone_op_as_graph(op, context->architecture); +} + +static bool device_supports_buffer_type(ggml_backend_dev_t device, ggml_backend_buffer_type_t buft) { + auto * context = device_context(device); + return buft == &context->buft || buft == &context->host_buft || ggml_backend_buft_is_host(buft); +} + +static const ggml_backend_device_i device_i = { + device_name, + device_description, + device_memory, + device_type, + device_props, + device_init, + device_buffer_type, + device_host_buffer_type, + nullptr, + device_supports_op, + device_supports_buffer_type, + nullptr, + nullptr, + nullptr, + nullptr, +}; + +static const char * registry_name(ggml_backend_reg_t registry) { + GGML_UNUSED(registry); + return "HRX"; +} + +static size_t registry_device_count(ggml_backend_reg_t registry) { + return static_cast(registry->context)->devices.size(); +} + +static ggml_backend_dev_t registry_device(ggml_backend_reg_t registry, size_t index) { + auto * context = static_cast(registry->context); + GGML_ASSERT(index < context->devices.size()); + return &context->devices[index]; +} + +static void * registry_proc(ggml_backend_reg_t registry, const char * name) { + GGML_UNUSED(registry); + GGML_UNUSED(name); + return nullptr; +} + +static const ggml_backend_reg_i registry_i = { registry_name, registry_device_count, registry_device, registry_proc }; + +static std::unique_ptr create_registry_context() { + auto context = std::make_unique(); + hrx_status_t status = hrx_gpu_initialize(0); + if (hrx_status_is_ok(status)) { + context->initialized = true; + } else if (hrx_status_code(status) == HRX_STATUS_ALREADY_EXISTS) { + hrx_status_ignore(status); + } else { + hrx_status_ignore(status); + return context; + } + int count = 0; + if (!HRX_CHECK(hrx_gpu_device_count(&count))) { + return context; + } + context->device_contexts.reserve(count); + context->devices.reserve(count); + for (int i = 0; i < count; ++i) { + hrx_device_t hrx_device = nullptr; + if (!HRX_CHECK(hrx_gpu_device_get(i, &hrx_device)) || hrx_device == nullptr) { + continue; + } + hrx_device_retain(hrx_device); + auto device_ctx = std::make_unique(); + device_ctx->device = hrx_device; + device_ctx->name = "HRX" + std::to_string(i); + device_ctx->use_direct_host_bindings = environment_flag_enabled("GGML_HRX_USE_UNIFIED_MEMORY"); + if (device_ctx->use_direct_host_bindings) { + GGML_LOG_INFO("ggml_hrx: direct coherent host bindings enabled by GGML_HRX_USE_UNIFIED_MEMORY\n"); + } + const std::optional name = + device_string_property(hrx_device, HRX_DEVICE_PROPERTY_NAME, "query HRX device name"); + const std::optional architecture = + device_string_property(hrx_device, HRX_DEVICE_PROPERTY_ARCHITECTURE, "query HRX device architecture"); + if (!name || !architecture) { + hrx_device_release(hrx_device); + continue; + } + uint64_t memory = 0; + if (!HRX_CHECK( + hrx_device_get_property(hrx_device, HRX_DEVICE_PROPERTY_TOTAL_MEMORY, &memory, sizeof(memory)))) { + hrx_device_release(hrx_device); + continue; + } + if (!HRX_CHECK(hrx_stream_create(hrx_device, 0, &device_ctx->buffer_stream))) { + hrx_device_release(hrx_device); + continue; + } + device_ctx->memory_total = static_cast(memory); + device_ctx->description = *name + " (" + *architecture + ")"; + device_ctx->architecture = *architecture; + device_ctx->buft_context = { device_ctx.get(), device_ctx->name, false }; + device_ctx->buft = { buffer_type_i, nullptr, &device_ctx->buft_context }; + device_ctx->host_buft_context = { device_ctx.get(), device_ctx->name + "_HOST", true }; + device_ctx->host_buft = { buffer_type_i, nullptr, &device_ctx->host_buft_context }; + context->device_contexts.emplace_back(std::move(device_ctx)); + context->devices.push_back({ device_i, nullptr, context->device_contexts.back().get() }); + context->device_contexts.back()->buft.device = &context->devices.back(); + context->device_contexts.back()->host_buft.device = &context->devices.back(); + } + return context; +} + +} // namespace + +ggml_backend_reg_t ggml_backend_hrx_reg() { + static std::unique_ptr context = create_registry_context(); + static ggml_backend_reg registry = { GGML_BACKEND_API_VERSION, registry_i, context.get() }; + for (auto & device : context->devices) { + device.reg = ®istry; + } + return ®istry; +} + +ggml_backend_t ggml_backend_hrx_init(size_t device) { + ggml_backend_reg_t registry = ggml_backend_hrx_reg(); + if (device >= ggml_backend_reg_dev_count(registry)) { + return nullptr; + } + return ggml_backend_dev_init(ggml_backend_reg_dev_get(registry, device), nullptr); +} + +bool ggml_backend_is_hrx(ggml_backend_t backend) { + return backend != nullptr && ggml_guid_matches(backend->guid, ggml_backend_hrx_guid()); +} + +bool ggml_backend_hrx_get_cache_stats(ggml_backend_t backend, ggml_backend_hrx_cache_stats * stats) { + if (!ggml_backend_is_hrx(backend) || stats == nullptr) { + return false; + } + auto * context = static_cast(backend->context); + const ggml::hrx::GraphProgramCacheStats graph_stats = context->graph_programs.stats(); + const ggml::hrx::PreparedCommandProgramCacheStats prepared_stats = context->prepared_programs.stats(); + stats->graph_program_builds = graph_stats.builds; + stats->graph_program_hits = graph_stats.hits; + stats->prepared_program_builds = graph_stats.prepared_program_builds + prepared_stats.builds; + stats->prepared_program_hits = graph_stats.prepared_program_hits + prepared_stats.hits; + return true; +} + +int ggml_backend_hrx_get_device_count() { + return static_cast(ggml_backend_reg_dev_count(ggml_backend_hrx_reg())); +} + +ggml_backend_buffer_type_t ggml_backend_hrx_buffer_type(size_t device) { + ggml_backend_reg_t registry = ggml_backend_hrx_reg(); + return device < ggml_backend_reg_dev_count(registry) ? + ggml_backend_dev_buffer_type(ggml_backend_reg_dev_get(registry, device)) : + nullptr; +} + +GGML_BACKEND_DL_IMPL(ggml_backend_hrx_reg) diff --git a/ggml/src/ggml-hrx/graph/graph-diagnostics.cpp b/ggml/src/ggml-hrx/graph/graph-diagnostics.cpp new file mode 100644 index 000000000000..23292feef5e0 --- /dev/null +++ b/ggml/src/ggml-hrx/graph/graph-diagnostics.cpp @@ -0,0 +1,512 @@ +#include "graph-diagnostics.h" + +#include "ggml-impl.h" +#include "ggml.h" + +#include +#include +#include +#include +#include +#include + +namespace ggml::hrx { +namespace { + +using json = nlohmann::ordered_json; + +const char * value_kind_name(ValueKind kind) { + switch (kind) { + case ValueKind::External: + return "external"; + case ValueKind::Transient: + return "transient"; + } + return "unknown"; +} + +bool parse_value_kind(const std::string & name, ValueKind & kind) { + if (name == "external") { + kind = ValueKind::External; + return true; + } + if (name == "transient") { + kind = ValueKind::Transient; + return true; + } + return false; +} + +const char * match_kind_name(DispatchMatchKind kind) { + switch (kind) { + case DispatchMatchKind::Fused: + return "fused"; + case DispatchMatchKind::SingleOp: + return "single_op"; + } + return "unknown"; +} + +const char * dispatch_source_name(DispatchSource source) { + switch (source) { + case DispatchSource::Common: + return "common"; + case DispatchSource::Llm: + return "llm"; + case DispatchSource::Qwen: + return "qwen"; + } + return "unknown"; +} + +json dims_json(const std::array & values) { + json result = json::array(); + for (int i = 0; i < GGML_MAX_DIMS; ++i) { + result.push_back(values[i]); + } + return result; +} + +json strides_json(const std::array & values) { + json result = json::array(); + for (int i = 0; i < GGML_MAX_DIMS; ++i) { + result.push_back(values[i]); + } + return result; +} + +Status read_dims(const json & item, const char * name, std::array & values) { + Status status; + if (!item.contains(name) || !item[name].is_array() || item[name].size() != GGML_MAX_DIMS) { + status.log("snapshot array %s must contain %d values", name, GGML_MAX_DIMS); + return status; + } + for (int i = 0; i < GGML_MAX_DIMS; ++i) { + values[i] = item[name][i].get(); + } + return status; +} + +Status read_strides(const json & item, const char * name, std::array & values) { + Status status; + if (!item.contains(name) || !item[name].is_array() || item[name].size() != GGML_MAX_DIMS) { + status.log("snapshot array %s must contain %d values", name, GGML_MAX_DIMS); + return status; + } + for (int i = 0; i < GGML_MAX_DIMS; ++i) { + values[i] = item[name][i].get(); + } + return status; +} + +json op_params_json(const OpParams & params) { + return std::visit( + [](const auto & value) -> json { + using T = std::decay_t; + if constexpr (std::is_same_v) { + return { + { "kind", "none" } + }; + } else if constexpr (std::is_same_v) { + return { + { "kind", "rms_norm" }, + { "eps", value.eps } + }; + } else if constexpr (std::is_same_v) { + return { + { "kind", "flash_attn_ext" }, + { "scale", value.scale }, + { "max_bias", value.max_bias }, + { "logit_softcap", value.logit_softcap }, + { "prec", static_cast(value.prec) }, + }; + } else if constexpr (std::is_same_v) { + return { + { "kind", "soft_max" }, + { "scale", value.scale }, + { "max_bias", value.max_bias } + }; + } else if constexpr (std::is_same_v) { + return { + { "kind", "argsort" }, + { "order", static_cast(value.order) } + }; + } else if constexpr (std::is_same_v) { + return { + { "kind", "clamp" }, + { "min", value.min }, + { "max", value.max } + }; + } else if constexpr (std::is_same_v) { + return { + { "kind", "glu" }, + { "op", static_cast(value.op) } + }; + } else if constexpr (std::is_same_v) { + return { + { "kind", "rope" }, + { "n_dims", value.n_dims }, + { "mode", value.mode }, + { "n_ctx_orig", value.n_ctx_orig }, + { "freq_base", value.freq_base }, + { "freq_scale", value.freq_scale }, + { "ext_factor", value.ext_factor }, + { "attn_factor", value.attn_factor }, + { "beta_fast", value.beta_fast }, + { "beta_slow", value.beta_slow }, + }; + } + }, + params); +} + +OpParams parse_op_params(const json & item) { + const std::string kind = item.value("kind", "none"); + if (kind == "rms_norm") { + return RmsNormParams{ item.value("eps", 0.0f) }; + } + if (kind == "flash_attn_ext") { + return FlashAttnExtParams{ + item.value("scale", 0.0f), + item.value("max_bias", 0.0f), + item.value("logit_softcap", 0.0f), + static_cast(item.value("prec", static_cast(GGML_PREC_DEFAULT))), + }; + } + if (kind == "soft_max") { + return SoftMaxParams{ item.value("scale", 0.0f), item.value("max_bias", 0.0f) }; + } + if (kind == "argsort") { + return ArgsortParams{ static_cast( + item.value("order", static_cast(GGML_SORT_ORDER_ASC))) }; + } + if (kind == "clamp") { + const float min = item.contains("min") && !item["min"].is_null() ? item["min"].get() : 0.0f; + const float max = item.contains("max") && !item["max"].is_null() ? item["max"].get() : + std::numeric_limits::infinity(); + return ClampParams{ min, max }; + } + if (kind == "glu") { + return GluParams{ static_cast(item.value("op", static_cast(GGML_GLU_OP_REGLU))) }; + } + if (kind == "rope") { + return RopeParams{ + item.value("n_dims", 0), item.value("mode", 0), item.value("n_ctx_orig", 0), + item.value("freq_base", 0.0f), item.value("freq_scale", 0.0f), item.value("ext_factor", 0.0f), + item.value("attn_factor", 0.0f), item.value("beta_fast", 0.0f), item.value("beta_slow", 0.0f), + }; + } + return std::monostate{}; +} + +json value_json(const Value & value) { + return { + { "id", value.id.value }, + { "kind", value_kind_name(value.kind) }, + { "storage", value.storage.value }, + { "storage_root", value.storage_root.value }, + { "alias_source", value.alias_source.value }, + { "storage_offset", value.storage_offset }, + { "storage_byte_count", value.storage_byte_count }, + { "type", static_cast(value.type) }, + { "type_name", ggml_type_name(value.type) }, + { "ne", dims_json(value.ne) }, + { "nb", strides_json(value.nb) }, + { "element_count", value.element_count }, + { "byte_count", value.byte_count }, + { "contiguous", value.contiguous }, + }; +} + +json storage_json(const ValueStorage & storage) { + return { + { "id", storage.id.value }, + { "root", storage.root.value }, + { "byte_count", storage.byte_count }, + }; +} + +json node_json(const GraphNode & node) { + json inputs = json::array(); + for (ValueId input : node.inputs) { + inputs.push_back(input.value); + } + return { + { "op", static_cast(node.op) }, + { "op_name", ggml_op_name(node.op) }, + { "output", node.output.value }, + { "inputs", std::move(inputs) }, + { "params", op_params_json(node.params) }, + }; +} + +std::string value_summary(const Graph & graph, ValueId id) { + std::ostringstream out; + const Value * value = graph.values().find(id); + if (value == nullptr) { + out << id.value << ":missing"; + return out.str(); + } + out << id.value << ":" << ggml_type_name(value->type) << "[" << value->ne[0] << "," << value->ne[1] << "," + << value->ne[2] << "," << value->ne[3] << "] " << value_kind_name(value->kind); + if (value->alias_source.value >= 0) { + out << " alias=" << value->alias_source.value << " storage_offset=" << value->storage_offset; + } + return out.str(); +} + +json attempts_json(const DispatchMatchDiagnostics & diagnostics) { + json attempts = json::array(); + for (const DispatchRegistrationAttempt & attempt : diagnostics.attempts) { + json covered = json::array(); + for (size_t node : attempt.covered_nodes) { + covered.push_back(node); + } + attempts.push_back({ + { "name", attempt.name }, + { "root_op", static_cast(attempt.root_op) }, + { "root_op_name", ggml_op_name(attempt.root_op) }, + { "kind", match_kind_name(attempt.kind) }, + { "priority", attempt.priority }, + { "source", dispatch_source_name(attempt.source) }, + { "matched", attempt.matched }, + { "covered_nodes", std::move(covered) }, + { "errors", attempt.errors }, + }); + } + return attempts; +} + +void write_file(const std::filesystem::path & path, const std::string & contents) { + std::filesystem::create_directories(path.parent_path()); + std::ofstream output(path, std::ios::binary | std::ios::trunc); + if (!output) { + throw std::runtime_error("cannot create " + path.string()); + } + output << contents; + if (contents.empty() || contents.back() != '\n') { + output << '\n'; + } +} + +} // namespace + +std::string serialize_graph_snapshot_json(const Graph & graph, const std::string & target, uint64_t uid) { + json values = json::array(); + for (const Value & value : graph.values().values()) { + values.push_back(value_json(value)); + } + json storages = json::array(); + for (const ValueStorage & storage : graph.values().storages()) { + storages.push_back(storage_json(storage)); + } + json nodes = json::array(); + for (const GraphNode & node : graph.nodes()) { + nodes.push_back(node_json(node)); + } + json root = { + { "schema", "ggml-hrx-graph-snapshot-v1" }, + { "uid", uid }, + { "target", target }, + { "storages", std::move(storages) }, + { "values", std::move(values) }, + { "nodes", std::move(nodes) }, + }; + return root.dump(2); +} + +std::string format_graph_snapshot_text(const Graph & graph, const std::string & target, uint64_t uid) { + std::ostringstream out; + out << "schema=ggml-hrx-graph-snapshot-v1\n"; + out << "uid=" << uid << "\ntarget=" << target << "\nvalues=" << graph.values().size() + << "\nnodes=" << graph.nodes().size() << '\n'; + for (size_t i = 0; i < graph.nodes().size(); ++i) { + const GraphNode & node = graph.nodes()[i]; + out << "node " << i << " " << ggml_op_name(node.op) << " output=" << value_summary(graph, node.output) + << " inputs=["; + for (size_t j = 0; j < node.inputs.size(); ++j) { + if (j > 0) { + out << ", "; + } + out << value_summary(graph, node.inputs[j]); + } + out << "] consumers=["; + const std::vector & consumers = graph.index().consumers(node.output); + for (size_t j = 0; j < consumers.size(); ++j) { + size_t consumer_index = 0; + if (j > 0) { + out << ", "; + } + if (graph.index().node_index(consumers[j], consumer_index)) { + out << consumer_index << ":" << ggml_op_name(consumers[j]->op); + } + } + out << "]\n"; + } + return out.str(); +} + +GraphSnapshotLoadResult load_graph_snapshot_json(const std::string & contents) { + GraphSnapshotLoadResult result; + try { + const json root = json::parse(contents); + if (root.value("schema", "") != "ggml-hrx-graph-snapshot-v1") { + result.status.log("unsupported graph snapshot schema"); + return result; + } + result.uid = root.value("uid", 0ULL); + result.target = root.value("target", ""); + + ValueMap & values = result.graph.values(); + for (const json & item : root.at("storages")) { + Status status = values.add_snapshot_storage({ + ValueStorageId(item.at("id").get()), + ValueId(item.at("root").get()), + item.at("byte_count").get(), + }); + if (!status.success()) { + result.status.append(status); + return result; + } + } + for (const json & item : root.at("values")) { + Value value = { + ValueId(item.at("id").get()), + ValueKind::External, + ValueStorageId(item.at("storage").get()), + ValueId(item.at("storage_root").get()), + ValueId(item.at("alias_source").get()), + item.at("storage_offset").get(), + item.at("storage_byte_count").get(), + static_cast(item.at("type").get()), + {}, + {}, + item.at("element_count").get(), + item.at("byte_count").get(), + item.at("contiguous").get(), + nullptr, + std::nullopt, + }; + if (!parse_value_kind(item.at("kind").get(), value.kind)) { + result.status.log("snapshot value %d has unknown kind", value.id.value); + return result; + } + Status dims_status = read_dims(item, "ne", value.ne); + if (!dims_status.success()) { + result.status.append(dims_status); + return result; + } + Status strides_status = read_strides(item, "nb", value.nb); + if (!strides_status.success()) { + result.status.append(strides_status); + return result; + } + Status status = values.add_snapshot_value(std::move(value)); + if (!status.success()) { + result.status.append(status); + return result; + } + } + for (const json & item : root.at("nodes")) { + std::vector inputs; + for (const json & input : item.at("inputs")) { + inputs.push_back(ValueId(input.get())); + } + GraphNode & node = result.graph.add_node(static_cast(item.at("op").get()), + ValueId(item.at("output").get()), std::move(inputs)); + node.params = parse_op_params(item.at("params")); + } + result.status.append(result.graph.build_index()); + } catch (const std::exception & error) { + result.status.log("failed to load graph snapshot: %s", error.what()); + } + return result; +} + +Status write_graph_snapshot(const std::filesystem::path & directory, + const Graph & graph, + const std::string & target, + uint64_t uid) { + Status status; + try { + static std::atomic sequence{ 0 }; + const uint64_t id = sequence.fetch_add(1); + std::ostringstream name; + name << "graph-" << id << "-uid-" << uid << "-" << target << "-" << graph.nodes().size() << "-nodes"; + const std::filesystem::path base = directory / name.str(); + write_file(base.string() + ".json", serialize_graph_snapshot_json(graph, target, uid)); + write_file(base.string() + ".txt", format_graph_snapshot_text(graph, target, uid)); + } catch (const std::exception & error) { + status.log("failed to write HRX graph snapshot: %s", error.what()); + } + return status; +} + +std::string format_schedule_diagnostics_text(const Graph & graph, + const CommandPlan & plan, + const DispatchScheduleDiagnostics & diagnostics) { + std::ostringstream out; + out << "valid=" << (plan.valid() ? "true" : "false") << '\n'; + for (const std::string & error : plan.status.errors()) { + out << "error=" << error << '\n'; + } + if (diagnostics.unsupported_node != nullptr) { + out << "unsupported_node=" << diagnostics.unsupported_node_index << ":" + << ggml_op_name(diagnostics.unsupported_node->op) << '\n'; + out << "unsupported_message=" << diagnostics.unsupported_message << '\n'; + out << "output=" << value_summary(graph, diagnostics.unsupported_node->output) << '\n'; + for (size_t i = 0; i < diagnostics.unsupported_node->inputs.size(); ++i) { + out << "input" << i << "=" << value_summary(graph, diagnostics.unsupported_node->inputs[i]) << '\n'; + } + } + out << "matcher_attempts=" << diagnostics.match.attempts.size() << '\n'; + for (const DispatchRegistrationAttempt & attempt : diagnostics.match.attempts) { + out << "attempt name=" << attempt.name << " kind=" << match_kind_name(attempt.kind) + << " source=" << dispatch_source_name(attempt.source) << " priority=" << attempt.priority + << " matched=" << (attempt.matched ? "true" : "false") << " covered=["; + for (size_t i = 0; i < attempt.covered_nodes.size(); ++i) { + if (i > 0) { + out << ","; + } + out << attempt.covered_nodes[i]; + } + out << "]\n"; + for (const std::string & error : attempt.errors) { + out << " error=" << error << '\n'; + } + } + return out.str(); +} + +std::string serialize_schedule_diagnostics_json(const Graph & graph, + const CommandPlan & plan, + const DispatchScheduleDiagnostics & diagnostics) { + json errors = json::array(); + for (const std::string & error : plan.status.errors()) { + errors.push_back(error); + } + json root = { + { "schema", "ggml-hrx-schedule-diagnostics-v1" }, + { "valid", plan.valid() }, + { "errors", std::move(errors) }, + { "matcher_attempts", attempts_json(diagnostics.match) }, + }; + if (diagnostics.unsupported_node != nullptr) { + json inputs = json::array(); + for (ValueId input : diagnostics.unsupported_node->inputs) { + inputs.push_back(input.value); + } + root["unsupported_node"] = { + { "index", diagnostics.unsupported_node_index }, + { "op", static_cast(diagnostics.unsupported_node->op) }, + { "op_name", ggml_op_name(diagnostics.unsupported_node->op) }, + { "output", diagnostics.unsupported_node->output.value }, + { "inputs", std::move(inputs) }, + { "message", diagnostics.unsupported_message }, + }; + } + GGML_UNUSED(graph); + return root.dump(2); +} + +} // namespace ggml::hrx diff --git a/ggml/src/ggml-hrx/graph/graph-diagnostics.h b/ggml/src/ggml-hrx/graph/graph-diagnostics.h new file mode 100644 index 000000000000..9956301ca9a8 --- /dev/null +++ b/ggml/src/ggml-hrx/graph/graph-diagnostics.h @@ -0,0 +1,39 @@ +#pragma once + +#include "dispatch/dispatch-scheduler.h" +#include "graph.h" +#include "status.h" + +#include +#include +#include + +namespace ggml::hrx { + +struct GraphSnapshotLoadResult { + uint64_t uid = 0; + std::string target; + Graph graph; + Status status; + + bool valid() const { return status.success(); } +}; + +std::string serialize_graph_snapshot_json(const Graph & graph, const std::string & target, uint64_t uid); +std::string format_graph_snapshot_text(const Graph & graph, const std::string & target, uint64_t uid); + +GraphSnapshotLoadResult load_graph_snapshot_json(const std::string & contents); + +Status write_graph_snapshot(const std::filesystem::path & directory, + const Graph & graph, + const std::string & target, + uint64_t uid); + +std::string format_schedule_diagnostics_text(const Graph & graph, + const CommandPlan & plan, + const DispatchScheduleDiagnostics & diagnostics); +std::string serialize_schedule_diagnostics_json(const Graph & graph, + const CommandPlan & plan, + const DispatchScheduleDiagnostics & diagnostics); + +} // namespace ggml::hrx diff --git a/ggml/src/ggml-hrx/graph/graph-matcher.cpp b/ggml/src/ggml-hrx/graph/graph-matcher.cpp new file mode 100644 index 000000000000..0d466f8d5e0f --- /dev/null +++ b/ggml/src/ggml-hrx/graph/graph-matcher.cpp @@ -0,0 +1,116 @@ +#include "graph-matcher.h" + +namespace ggml::hrx { +namespace { + +static bool node_in_list(const GraphNode * node, const std::vector & nodes) { + for (const GraphNode * candidate : nodes) { + if (candidate == node) { + return true; + } + } + return false; +} + +static void append_unique_node(std::vector & nodes, const GraphNode * node) { + if (node != nullptr && !node_in_list(node, nodes)) { + nodes.push_back(node); + } +} + +} // namespace + +std::vector layout_alias_consumers(const Graph & graph, ValueId value) { + std::vector matches; + for (const GraphNode * consumer : graph.index().consumers(value)) { + if (consumer != nullptr && is_layout_alias_node(graph, *consumer)) { + matches.push_back(consumer); + } + } + return matches; +} + +std::vector layout_alias_consumers_with_op(const Graph & graph, ValueId value, ggml_op op) { + std::vector matches; + for (const GraphNode * consumer : layout_alias_consumers(graph, value)) { + if (consumer->op == op) { + matches.push_back(consumer); + } + } + return matches; +} + +const GraphNode * find_single_layout_alias_consumer(const Graph & graph, ValueId value) { + const std::vector matches = layout_alias_consumers(graph, value); + return matches.size() == 1 ? matches.front() : nullptr; +} + +const GraphNode * find_single_layout_alias_consumer_with_op(const Graph & graph, ValueId value, ggml_op op) { + const std::vector matches = layout_alias_consumers_with_op(graph, value, op); + return matches.size() == 1 ? matches.front() : nullptr; +} + +std::vector consumers_with_op_through_layout_aliases(const Graph & graph, + ValueId value, + ggml_op op) { + std::vector matches; + for (const GraphNode * consumer : graph.index().consumers(value)) { + if (consumer == nullptr) { + continue; + } + if (consumer->op == op) { + append_unique_node(matches, consumer); + } + if (!is_layout_alias_node(graph, *consumer)) { + continue; + } + for (const GraphNode * alias_consumer : graph.index().consumers(consumer->output)) { + if (alias_consumer != nullptr && alias_consumer->op == op) { + append_unique_node(matches, alias_consumer); + } + } + } + return matches; +} + +const GraphNode * find_single_consumer_with_op_through_layout_aliases(const Graph & graph, ValueId value, ggml_op op) { + const std::vector matches = consumers_with_op_through_layout_aliases(graph, value, op); + return matches.size() == 1 ? matches.front() : nullptr; +} + +bool node_has_input_or_alias(const Graph & graph, const GraphNode & node, ValueId input) { + const Value * input_value = graph.values().find(input); + for (ValueId candidate : node.inputs) { + if (candidate == input) { + return true; + } + const Value * candidate_value = graph.values().find(candidate); + if (candidate_value != nullptr && candidate_value->alias_source == input) { + return true; + } + if (input_value != nullptr && input_value->alias_source == candidate) { + return true; + } + } + return false; +} + +bool append_covered_node_index_once(const Graph & graph, + const std::vector & covered_nodes, + const GraphNode * node, + std::vector & covered_indices) { + size_t index = 0; + if (node == nullptr || !graph.index().node_index(node, index) || index >= covered_nodes.size() || + covered_nodes[index]) { + return false; + } + for (const size_t covered : covered_indices) { + if (covered == index) { + return true; + } + } + covered_indices.push_back(index); + return true; +} + +} // namespace ggml::hrx diff --git a/ggml/src/ggml-hrx/graph/graph-matcher.h b/ggml/src/ggml-hrx/graph/graph-matcher.h new file mode 100644 index 000000000000..507165d45eff --- /dev/null +++ b/ggml/src/ggml-hrx/graph/graph-matcher.h @@ -0,0 +1,25 @@ +#pragma once + +#include "graph.h" + +#include +#include + +namespace ggml::hrx { + +std::vector layout_alias_consumers(const Graph & graph, ValueId value); +std::vector layout_alias_consumers_with_op(const Graph & graph, ValueId value, ggml_op op); + +const GraphNode * find_single_layout_alias_consumer(const Graph & graph, ValueId value); +const GraphNode * find_single_layout_alias_consumer_with_op(const Graph & graph, ValueId value, ggml_op op); + +std::vector consumers_with_op_through_layout_aliases(const Graph & graph, ValueId value, ggml_op op); +const GraphNode * find_single_consumer_with_op_through_layout_aliases(const Graph & graph, ValueId value, ggml_op op); + +bool node_has_input_or_alias(const Graph & graph, const GraphNode & node, ValueId input); +bool append_covered_node_index_once(const Graph & graph, + const std::vector & covered_nodes, + const GraphNode * node, + std::vector & covered_indices); + +} // namespace ggml::hrx diff --git a/ggml/src/ggml-hrx/graph/graph-traversal.cpp b/ggml/src/ggml-hrx/graph/graph-traversal.cpp new file mode 100644 index 000000000000..22edc6bba011 --- /dev/null +++ b/ggml/src/ggml-hrx/graph/graph-traversal.cpp @@ -0,0 +1,163 @@ +#include "graph-traversal.h" + +#include +#include + +namespace ggml::hrx { +namespace { + +static bool is_root_op(ggml_op op) { + switch (op) { + case GGML_OP_MUL_MAT: + case GGML_OP_MUL_MAT_ID: + case GGML_OP_FLASH_ATTN_EXT: + case GGML_OP_CONV_TRANSPOSE_1D: + case GGML_OP_CONV_2D: + case GGML_OP_CONV_3D: + case GGML_OP_CONV_2D_DW: + case GGML_OP_CONV_TRANSPOSE_2D: + case GGML_OP_SSM_CONV: + return true; + default: + return false; + } +} + +static bool is_fusable_followup_op(ggml_op op) { + switch (op) { + case GGML_OP_ADD: + case GGML_OP_MUL: + return true; + default: + return false; + } +} + +static void erase_ready(size_t node_index, std::set & root_queue, std::set & regular_queue) { + root_queue.erase(node_index); + regular_queue.erase(node_index); +} + +static void add_ready_node(const GraphNode & node, + size_t node_index, + const std::vector & selected, + std::set & root_queue, + std::set & regular_queue) { + if (node_index >= selected.size() || selected[node_index]) { + return; + } + if (is_root_op(node.op)) { + root_queue.insert(node_index); + } else { + regular_queue.insert(node_index); + } +} + +static bool select_merge_candidate(const Graph & graph, + size_t selected_node, + const std::vector & pending_inputs, + const std::vector & selected, + size_t & next_node) { + const std::vector & nodes = graph.nodes(); + if (selected_node >= nodes.size()) { + return false; + } + + bool found = false; + size_t best_index = 0; + for (const GraphNode * consumer : graph.index().consumers(nodes[selected_node].output)) { + size_t consumer_index = 0; + if (consumer == nullptr || !graph.index().node_index(consumer, consumer_index) || + consumer_index >= selected.size() || selected[consumer_index] || pending_inputs[consumer_index] != 0 || + !is_fusable_followup_op(consumer->op)) { + continue; + } + if (!found || consumer_index < best_index) { + found = true; + best_index = consumer_index; + } + } + if (!found) { + return false; + } + next_node = best_index; + return true; +} + +} // namespace + +GraphTraversalOrder GraphTraversalOrder::build(const Graph & graph) { + GraphTraversalOrder result; + const std::vector & nodes = graph.nodes(); + result.nodes_.reserve(nodes.size()); + if (!graph.has_index()) { + for (const GraphNode & node : nodes) { + result.nodes_.push_back(&node); + } + return result; + } + + std::vector pending_inputs(nodes.size(), 0); + for (size_t i = 0; i < nodes.size(); ++i) { + for (ValueId input : nodes[i].inputs) { + if (graph.index().producer(input) != nullptr) { + ++pending_inputs[i]; + } + } + } + + std::set root_queue; + std::set regular_queue; + std::vector selected(nodes.size(), false); + for (size_t i = 0; i < nodes.size(); ++i) { + if (pending_inputs[i] == 0) { + add_ready_node(nodes[i], i, selected, root_queue, regular_queue); + } + } + + bool has_previous = false; + size_t previous_node = 0; + while (result.nodes_.size() < nodes.size()) { + size_t selected_index = 0; + if (has_previous && select_merge_candidate(graph, previous_node, pending_inputs, selected, selected_index)) { + erase_ready(selected_index, root_queue, regular_queue); + } else if (!root_queue.empty()) { + selected_index = *root_queue.begin(); + root_queue.erase(root_queue.begin()); + } else if (!regular_queue.empty()) { + selected_index = *regular_queue.begin(); + regular_queue.erase(regular_queue.begin()); + } else { + break; + } + + if (selected[selected_index]) { + continue; + } + selected[selected_index] = true; + result.nodes_.push_back(&nodes[selected_index]); + has_previous = true; + previous_node = selected_index; + + for (const GraphNode * consumer : graph.index().consumers(nodes[selected_index].output)) { + size_t consumer_index = 0; + if (consumer == nullptr || !graph.index().node_index(consumer, consumer_index) || + consumer_index >= pending_inputs.size() || selected[consumer_index]) { + continue; + } + --pending_inputs[consumer_index]; + if (pending_inputs[consumer_index] == 0) { + add_ready_node(*consumer, consumer_index, selected, root_queue, regular_queue); + } + } + } + + for (size_t i = 0; i < nodes.size(); ++i) { + if (!selected[i]) { + result.nodes_.push_back(&nodes[i]); + } + } + return result; +} + +} // namespace ggml::hrx diff --git a/ggml/src/ggml-hrx/graph/graph-traversal.h b/ggml/src/ggml-hrx/graph/graph-traversal.h new file mode 100644 index 000000000000..c5cb1f27a751 --- /dev/null +++ b/ggml/src/ggml-hrx/graph/graph-traversal.h @@ -0,0 +1,21 @@ +#pragma once + +#include "graph.h" + +#include + +namespace ggml::hrx { + +class GraphTraversalOrder { + public: + GraphTraversalOrder() = default; + + static GraphTraversalOrder build(const Graph & graph); + + const std::vector & nodes() const { return nodes_; } + + private: + std::vector nodes_; +}; + +} // namespace ggml::hrx diff --git a/ggml/src/ggml-hrx/graph/graph.cpp b/ggml/src/ggml-hrx/graph/graph.cpp new file mode 100644 index 000000000000..002291b0f19b --- /dev/null +++ b/ggml/src/ggml-hrx/graph/graph.cpp @@ -0,0 +1,217 @@ +// Copyright 2026 The HRX Authors +// SPDX-License-Identifier: Apache-2.0 + +#include "graph.h" + +#include "ggml-impl.h" + +#include +#include +#include +#include + +namespace ggml::hrx { +namespace { + +// A value crosses the imported graph boundary (and therefore needs an external +// binding) in two cases: it is not produced by any node in the graph (a leaf such +// as a weight/token id, or an activation produced by another backend), or it is +// produced here but never consumed here (a graph output that must be downloaded). +// Only values produced and consumed inside the graph are true transients. +static bool tensor_is_external(const ggml_tensor * tensor, + const std::unordered_set & node_outputs, + const std::unordered_map & use_counts) { + if (tensor->op == GGML_OP_NONE) { + return true; + } + if (node_outputs.find(tensor) == node_outputs.end()) { + return true; + } + const auto found = use_counts.find(tensor); + return found == use_counts.end() || found->second == 0; +} + +} // namespace + +GraphIndex GraphIndex::build(const Graph & graph) { + GraphIndex index; + const std::vector & nodes = graph.nodes(); + for (size_t i = 0; i < nodes.size(); ++i) { + const GraphNode & node = nodes[i]; + index.node_indices_.emplace(&node, i); + index.producers_.emplace(node.output.value, &node); + for (ValueId input : node.inputs) { + index.consumers_[input.value].push_back(&node); + } + } + return index; +} + +const GraphNode * GraphIndex::producer(ValueId value) const { + const auto found = producers_.find(value.value); + return found == producers_.end() ? nullptr : found->second; +} + +const std::vector & GraphIndex::consumers(ValueId value) const { + static const std::vector empty; + const auto found = consumers_.find(value.value); + return found == consumers_.end() ? empty : found->second; +} + +bool GraphIndex::has_single_consumer(ValueId value) const { + return consumers(value).size() == 1; +} + +bool GraphIndex::node_index(const GraphNode * node, size_t & index) const { + const auto found = node_indices_.find(node); + if (found == node_indices_.end()) { + return false; + } + index = found->second; + return true; +} + +Graph::Graph(const Graph & other) : values_(other.values_), nodes_(other.nodes_) { + if (other.has_index()) { + index_ = GraphIndex::build(*this); + } +} + +Graph & Graph::operator=(const Graph & other) { + if (this == &other) { + return *this; + } + values_ = other.values_; + nodes_ = other.nodes_; + index_.reset(); + if (other.has_index()) { + index_ = GraphIndex::build(*this); + } + return *this; +} + +Graph::Graph(Graph && other) : values_(std::move(other.values_)), nodes_(std::move(other.nodes_)) { + if (other.has_index()) { + index_ = GraphIndex::build(*this); + } +} + +Graph & Graph::operator=(Graph && other) { + if (this == &other) { + return *this; + } + values_ = std::move(other.values_); + nodes_ = std::move(other.nodes_); + index_.reset(); + if (other.has_index()) { + index_ = GraphIndex::build(*this); + } + return *this; +} + +GraphNode & Graph::add_node(ggml_op op, ValueId output, std::vector inputs) { + index_.reset(); + GraphNode node; + node.op = op; + node.output = output; + node.inputs = std::move(inputs); + nodes_.push_back(std::move(node)); + return nodes_.back(); +} + +Status Graph::build_index() { + index_ = GraphIndex::build(*this); + return {}; +} + +const GraphIndex & Graph::index() const { + assert(index_.has_value()); + return *index_; +} + +GraphImportResult import_ggml_graph(const ggml_cgraph & graph) { + GraphImportResult result; + std::unordered_set node_outputs; + std::unordered_map use_counts; + for (int i = 0; i < graph.n_nodes; ++i) { + const ggml_tensor * node = graph.nodes[i]; + if (node == nullptr) { + result.status.log("ggml graph contains a null node"); + return result; + } + node_outputs.insert(node); + for (const ggml_tensor * source : node->src) { + if (source != nullptr) { + ++use_counts[source]; + } + } + } + + ValueMap & values = result.graph.values(); + + // Views share storage with the tensor they were created from, and the scheduler may + // hand the HRX dispatcher a graph whose terminal value is such a view (e.g. a + // standalone MUL_MAT followed by a RESHAPE that ggml_backend_sched kept on the same + // backend). The classification must be consistent for the whole storage: if any + // value in a storage is external (bound across the graph boundary), every value in + // that storage is external too. Otherwise the kernel would write its result into a + // transient arena slot while the consumer reads the host tensor the view aliases. + auto storage_root_of = [](const ggml_tensor * tensor) { + const ggml_tensor * root = tensor; + while (root->view_src != nullptr) { + root = root->view_src; + } + return root; + }; + std::unordered_map storage_is_external; + auto note_storage = [&](const ggml_tensor * tensor) { + const bool external = tensor_is_external(tensor, node_outputs, use_counts); + const ggml_tensor * root = storage_root_of(tensor); + bool & flag = storage_is_external[root]; + flag = flag || external; + }; + for (int i = 0; i < graph.n_nodes; ++i) { + const ggml_tensor * node = graph.nodes[i]; + note_storage(node); + for (const ggml_tensor * source : node->src) { + if (source != nullptr) { + note_storage(source); + } + } + } + auto value_kind = [&](const ggml_tensor * tensor) { + const bool external = + tensor_is_external(tensor, node_outputs, use_counts) || storage_is_external[storage_root_of(tensor)]; + return external ? ValueKind::External : ValueKind::Transient; + }; + + for (int i = 0; i < graph.n_nodes; ++i) { + const ggml_tensor * node = graph.nodes[i]; + std::vector inputs; + for (const ggml_tensor * source : node->src) { + if (source == nullptr) { + continue; + } + inputs.push_back(values.get_or_add_tensor_value(source, value_kind(source))); + } + + const ValueKind output_kind = value_kind(node); + const ValueId output = values.get_or_add_tensor_value(node, output_kind); + GraphNode & graph_node = result.graph.add_node(node->op, output, std::move(inputs)); + graph_node.params = import_op_params(*node); + } + + result.status.append(result.graph.build_index()); + return result; +} + +bool is_layout_alias_op(ggml_op op) { + return op == GGML_OP_VIEW || op == GGML_OP_RESHAPE || op == GGML_OP_PERMUTE || op == GGML_OP_TRANSPOSE; +} + +bool is_layout_alias_node(const Graph & graph, const GraphNode & node) { + return is_layout_alias_op(node.op) && node.inputs.size() == 1 && + graph.values().same_storage(node.output, node.inputs[0]); +} + +} // namespace ggml::hrx diff --git a/ggml/src/ggml-hrx/graph/graph.h b/ggml/src/ggml-hrx/graph/graph.h new file mode 100644 index 000000000000..60b26a59daa7 --- /dev/null +++ b/ggml/src/ggml-hrx/graph/graph.h @@ -0,0 +1,85 @@ +#pragma once + +#include "ggml.h" +#include "op-params.h" +#include "status.h" +#include "value-map.h" + +#include +#include +#include +#include +#include + +struct ggml_cgraph; +struct ggml_tensor; + +namespace ggml::hrx { + +struct GraphNode { + ggml_op op; + ValueId output; + std::vector inputs; + OpParams params; +}; + +class Graph; + +class GraphIndex { + public: + GraphIndex() = default; + + static GraphIndex build(const Graph & graph); + + const GraphNode * producer(ValueId value) const; + const std::vector & consumers(ValueId value) const; + bool has_single_consumer(ValueId value) const; + bool node_index(const GraphNode * node, size_t & index) const; + + private: + std::unordered_map producers_; + std::unordered_map> consumers_; + std::unordered_map node_indices_; +}; + +class Graph { + public: + Graph() = default; + Graph(const Graph & other); + Graph & operator=(const Graph & other); + Graph(Graph && other); + Graph & operator=(Graph && other); + + GraphNode & add_node(ggml_op op, ValueId output, std::vector inputs); + + Status build_index(); + + bool has_index() const { return index_.has_value(); } + + const GraphIndex & index() const; + + const std::vector & nodes() const { return nodes_; } + + const ValueMap & values() const { return values_; } + + ValueMap & values() { return values_; } + + private: + ValueMap values_; + std::vector nodes_; + std::optional index_; +}; + +struct GraphImportResult { + Graph graph; + Status status; + + bool valid() const { return status.success(); } +}; + +GraphImportResult import_ggml_graph(const ggml_cgraph & graph); + +bool is_layout_alias_op(ggml_op op); +bool is_layout_alias_node(const Graph & graph, const GraphNode & node); + +} // namespace ggml::hrx diff --git a/ggml/src/ggml-hrx/graph/op-params.cpp b/ggml/src/ggml-hrx/graph/op-params.cpp new file mode 100644 index 000000000000..f76097164e14 --- /dev/null +++ b/ggml/src/ggml-hrx/graph/op-params.cpp @@ -0,0 +1,135 @@ +#include "op-params.h" + +#include "ggml-impl.h" + +#include + +namespace ggml::hrx { +namespace { + +static bool nearly_equal(float lhs, float rhs) { + if (lhs == rhs) { + return true; + } + return std::fabs(lhs - rhs) <= 1.0e-12f; +} + +static bool rms_norm_params_equivalent(const OpParams & lhs, const OpParams & rhs) { + const RmsNormParams * lhs_params = op_params_as(lhs); + const RmsNormParams * rhs_params = op_params_as(rhs); + return lhs_params != nullptr && rhs_params != nullptr && nearly_equal(lhs_params->eps, rhs_params->eps); +} + +static bool flash_attn_ext_params_equivalent(const OpParams & lhs, const OpParams & rhs) { + const FlashAttnExtParams * lhs_params = op_params_as(lhs); + const FlashAttnExtParams * rhs_params = op_params_as(rhs); + return lhs_params != nullptr && rhs_params != nullptr && nearly_equal(lhs_params->scale, rhs_params->scale) && + nearly_equal(lhs_params->max_bias, rhs_params->max_bias) && + nearly_equal(lhs_params->logit_softcap, rhs_params->logit_softcap) && lhs_params->prec == rhs_params->prec; +} + +static bool soft_max_params_equivalent(const OpParams & lhs, const OpParams & rhs) { + const SoftMaxParams * lhs_params = op_params_as(lhs); + const SoftMaxParams * rhs_params = op_params_as(rhs); + return lhs_params != nullptr && rhs_params != nullptr && nearly_equal(lhs_params->scale, rhs_params->scale) && + nearly_equal(lhs_params->max_bias, rhs_params->max_bias); +} + +static bool argsort_params_equivalent(const OpParams & lhs, const OpParams & rhs) { + const ArgsortParams * lhs_params = op_params_as(lhs); + const ArgsortParams * rhs_params = op_params_as(rhs); + return lhs_params != nullptr && rhs_params != nullptr && lhs_params->order == rhs_params->order; +} + +static bool clamp_params_equivalent(const OpParams & lhs, const OpParams & rhs) { + const ClampParams * lhs_params = op_params_as(lhs); + const ClampParams * rhs_params = op_params_as(rhs); + return lhs_params != nullptr && rhs_params != nullptr && nearly_equal(lhs_params->min, rhs_params->min) && + nearly_equal(lhs_params->max, rhs_params->max); +} + +static bool glu_params_equivalent(const OpParams & lhs, const OpParams & rhs) { + const GluParams * lhs_params = op_params_as(lhs); + const GluParams * rhs_params = op_params_as(rhs); + return lhs_params != nullptr && rhs_params != nullptr && lhs_params->op == rhs_params->op; +} + +static bool rope_params_equivalent(const OpParams & lhs, const OpParams & rhs) { + const RopeParams * lhs_params = op_params_as(lhs); + const RopeParams * rhs_params = op_params_as(rhs); + return lhs_params != nullptr && rhs_params != nullptr && lhs_params->n_dims == rhs_params->n_dims && + lhs_params->mode == rhs_params->mode && lhs_params->n_ctx_orig == rhs_params->n_ctx_orig && + nearly_equal(lhs_params->freq_base, rhs_params->freq_base) && + nearly_equal(lhs_params->freq_scale, rhs_params->freq_scale) && + nearly_equal(lhs_params->ext_factor, rhs_params->ext_factor) && + nearly_equal(lhs_params->attn_factor, rhs_params->attn_factor) && + nearly_equal(lhs_params->beta_fast, rhs_params->beta_fast) && + nearly_equal(lhs_params->beta_slow, rhs_params->beta_slow); +} + +} // namespace + +OpParams import_op_params(const ggml_tensor & tensor) { + switch (tensor.op) { + case GGML_OP_RMS_NORM: + return RmsNormParams{ ggml_get_op_params_f32(&tensor, 0) }; + case GGML_OP_SOFT_MAX: + return SoftMaxParams{ + ggml_get_op_params_f32(&tensor, 0), + ggml_get_op_params_f32(&tensor, 1), + }; + case GGML_OP_FLASH_ATTN_EXT: + return FlashAttnExtParams{ + ggml_get_op_params_f32(&tensor, 0), + ggml_get_op_params_f32(&tensor, 1), + ggml_get_op_params_f32(&tensor, 2), + ggml_flash_attn_ext_get_prec(&tensor), + }; + case GGML_OP_ARGSORT: + return ArgsortParams{ static_cast(ggml_get_op_params_i32(&tensor, 0)) }; + case GGML_OP_CLAMP: + return ClampParams{ + ggml_get_op_params_f32(&tensor, 0), + ggml_get_op_params_f32(&tensor, 1), + }; + case GGML_OP_GLU: + return GluParams{ ggml_get_glu_op(&tensor) }; + case GGML_OP_ROPE: + return RopeParams{ + ggml_get_op_params_i32(&tensor, 1), ggml_get_op_params_i32(&tensor, 2), + ggml_get_op_params_i32(&tensor, 4), ggml_get_op_params_f32(&tensor, 5), + ggml_get_op_params_f32(&tensor, 6), ggml_get_op_params_f32(&tensor, 7), + ggml_get_op_params_f32(&tensor, 8), ggml_get_op_params_f32(&tensor, 9), + ggml_get_op_params_f32(&tensor, 10), + }; + default: + return std::monostate{}; + } +} + +bool op_params_equivalent(ggml_op op, const OpParams & lhs, const OpParams & rhs) { + switch (op) { + case GGML_OP_RMS_NORM: + return rms_norm_params_equivalent(lhs, rhs); + case GGML_OP_SOFT_MAX: + return soft_max_params_equivalent(lhs, rhs); + case GGML_OP_FLASH_ATTN_EXT: + return flash_attn_ext_params_equivalent(lhs, rhs); + case GGML_OP_ARGSORT: + return argsort_params_equivalent(lhs, rhs); + case GGML_OP_CLAMP: + return clamp_params_equivalent(lhs, rhs); + case GGML_OP_GLU: + return glu_params_equivalent(lhs, rhs); + case GGML_OP_ROPE: + return rope_params_equivalent(lhs, rhs); + default: + return lhs.index() == rhs.index(); + } +} + +bool op_params_equivalent(ggml_op op, const OpParams & lhs, const ggml_tensor & rhs) { + return op_params_equivalent(op, lhs, import_op_params(rhs)); +} + +} // namespace ggml::hrx diff --git a/ggml/src/ggml-hrx/graph/op-params.h b/ggml/src/ggml-hrx/graph/op-params.h new file mode 100644 index 000000000000..dd1a97ebbdf8 --- /dev/null +++ b/ggml/src/ggml-hrx/graph/op-params.h @@ -0,0 +1,72 @@ +#pragma once + +#include "ggml.h" + +#include + +struct ggml_tensor; + +namespace ggml::hrx { + +struct RmsNormParams { + float eps = 0.0f; +}; + +struct FlashAttnExtParams { + float scale = 0.0f; + float max_bias = 0.0f; + float logit_softcap = 0.0f; + ggml_prec prec = GGML_PREC_DEFAULT; +}; + +struct SoftMaxParams { + float scale = 0.0f; + float max_bias = 0.0f; +}; + +struct ArgsortParams { + ggml_sort_order order = GGML_SORT_ORDER_ASC; +}; + +struct ClampParams { + float min = 0.0f; + float max = 0.0f; +}; + +struct GluParams { + ggml_glu_op op = GGML_GLU_OP_REGLU; +}; + +struct RopeParams { + int n_dims = 0; + int mode = 0; + int n_ctx_orig = 0; + float freq_base = 0.0f; + float freq_scale = 0.0f; + float ext_factor = 0.0f; + float attn_factor = 0.0f; + float beta_fast = 0.0f; + float beta_slow = 0.0f; +}; + +// clang-format off +using OpParams = std::variant< + std::monostate, + RmsNormParams, + FlashAttnExtParams, + SoftMaxParams, + ArgsortParams, + ClampParams, + GluParams, + RopeParams>; +// clang-format on + +template const T * op_params_as(const OpParams & params) { + return std::get_if(¶ms); +} + +OpParams import_op_params(const ggml_tensor & tensor); +bool op_params_equivalent(ggml_op op, const OpParams & lhs, const OpParams & rhs); +bool op_params_equivalent(ggml_op op, const OpParams & lhs, const ggml_tensor & rhs); + +} // namespace ggml::hrx diff --git a/ggml/src/ggml-hrx/graph/value-map.cpp b/ggml/src/ggml-hrx/graph/value-map.cpp new file mode 100644 index 000000000000..8bc09d186877 --- /dev/null +++ b/ggml/src/ggml-hrx/graph/value-map.cpp @@ -0,0 +1,245 @@ +#include "value-map.h" + +#include "ggml-impl.h" + +#include + +namespace ggml::hrx { + +const Value * ValueMap::find_alias_source(const ggml_tensor * tensor) const { + if (tensor == nullptr || tensor->view_src == nullptr) { + return nullptr; + } + return find_tensor(tensor->view_src); +} + +ValueId ValueMap::get_or_add_tensor_value(const ggml_tensor * tensor, ValueKind kind) { + const auto found = tensor_values_.find(tensor); + if (found != tensor_values_.end()) { + Value & value = values_[found->second]; + if (kind == ValueKind::External) { + value.kind = ValueKind::External; + } + return value.id; + } + + const ValueId id(static_cast(values_.size())); + const Value * alias_source = find_alias_source(tensor); + ValueStorageId storage; + ValueId storage_root; + ValueId alias_source_id; + size_t storage_offset = 0; + size_t storage_byte_count = ggml_nbytes(tensor); + if (alias_source != nullptr) { + storage = alias_source->storage; + storage_root = alias_source->storage_root; + alias_source_id = alias_source->id; + storage_offset = tensor->view_offs; + storage_byte_count = alias_source->storage_byte_count; + const Value * root = find(storage_root); + if (root != nullptr && root->kind == ValueKind::External) { + kind = ValueKind::External; + } + } else { + storage = ValueStorageId(static_cast(storages_.size())); + storage_root = id; + storages_.push_back({ storage, storage_root, storage_byte_count }); + } + + Value value = { + id, + kind, + storage, + storage_root, + alias_source_id, + storage_offset, + storage_byte_count, + tensor->type, + {}, + {}, + ggml_nelements(tensor), + ggml_nbytes(tensor), + ggml_is_contiguous(tensor), + tensor, + std::nullopt, + }; + for (int i = 0; i < GGML_MAX_DIMS; ++i) { + value.ne[i] = tensor->ne[i]; + value.nb[i] = tensor->nb[i]; + } + + values_.push_back(std::move(value)); + tensor_values_.emplace(tensor, values_.size() - 1); + return values_.back().id; +} + +const Value * ValueMap::find(ValueId id) const { + if (id.value < 0 || static_cast(id.value) >= values_.size()) { + return nullptr; + } + return &values_[static_cast(id.value)]; +} + +const ValueStorage * ValueMap::find_storage(ValueStorageId id) const { + if (id.value < 0 || static_cast(id.value) >= storages_.size()) { + return nullptr; + } + return &storages_[static_cast(id.value)]; +} + +bool ValueMap::bind_buffer(ValueId id, ValueBufferBinding binding) { + if (id.value < 0 || static_cast(id.value) >= values_.size()) { + return false; + } + Value & value = values_[static_cast(id.value)]; + if (value.kind != ValueKind::External) { + return false; + } + value.buffer = std::move(binding); + return true; +} + +std::optional ValueMap::resolve_buffer_binding(ValueId id) const { + const Value * value = find(id); + if (value == nullptr) { + return std::nullopt; + } + if (value->buffer.has_value()) { + return value->buffer; + } + if (value->storage_root == value->id) { + return std::nullopt; + } + const Value * root = find(value->storage_root); + if (root == nullptr || !root->buffer.has_value()) { + return std::nullopt; + } + ValueBufferBinding binding = *root->buffer; + if (value->storage_offset > binding.length) { + return std::nullopt; + } + if (value->byte_count > binding.length - value->storage_offset) { + return std::nullopt; + } + binding.offset += value->storage_offset; + binding.length = value->byte_count; + return binding; +} + +std::vector ValueMap::external_value_ids() const { + std::vector ids; + for (const Value & value : values_) { + if (value.kind == ValueKind::External) { + ids.push_back(value.id); + } + } + return ids; +} + +Status ValueMap::alias_storage(ValueId target, ValueId source) { + Status status; + if (target.value < 0 || static_cast(target.value) >= values_.size()) { + status.log("value alias target %d does not exist", target.value); + return status; + } + if (source.value < 0 || static_cast(source.value) >= values_.size()) { + status.log("value alias source %d does not exist", source.value); + return status; + } + if (target == source) { + status.log("value alias target %d aliases itself", target.value); + return status; + } + + Value & target_value = values_[static_cast(target.value)]; + const Value & source_value = values_[static_cast(source.value)]; + if (target_value.storage == source_value.storage) { + return status; + } + if (target_value.kind != ValueKind::Transient) { + status.log("value alias target %d is not transient", target.value); + return status; + } + if (target_value.type != source_value.type || target_value.byte_count != source_value.byte_count || + target_value.element_count != source_value.element_count || + target_value.contiguous != source_value.contiguous) { + status.log("value alias target %d is incompatible with source %d", target.value, source.value); + return status; + } + for (int i = 0; i < GGML_MAX_DIMS; ++i) { + if (target_value.ne[i] != source_value.ne[i] || target_value.nb[i] != source_value.nb[i]) { + status.log("value alias target %d has a different layout than source %d", target.value, source.value); + return status; + } + } + if (target_value.alias_source.value >= 0 && target_value.alias_source != source) { + status.log("value alias target %d already aliases source %d", target.value, target_value.alias_source.value); + return status; + } + + target_value.storage = source_value.storage; + target_value.storage_root = source_value.storage_root; + target_value.alias_source = source; + target_value.storage_offset = source_value.storage_offset; + target_value.storage_byte_count = source_value.storage_byte_count; + return status; +} + +ValueId ValueMap::storage_root(ValueId id) const { + const Value * value = find(id); + return value == nullptr ? ValueId() : value->storage_root; +} + +bool ValueMap::same_storage(ValueId lhs, ValueId rhs) const { + const Value * lhs_value = find(lhs); + const Value * rhs_value = find(rhs); + return lhs_value != nullptr && rhs_value != nullptr && lhs_value->storage == rhs_value->storage; +} + +Status ValueMap::add_snapshot_storage(ValueStorage storage) { + Status status; + if (storage.id.value < 0 || static_cast(storage.id.value) != storages_.size()) { + status.log("snapshot storage id %d is not the next storage id %zu", storage.id.value, storages_.size()); + return status; + } + if (storage.root.value < 0) { + status.log("snapshot storage %d has invalid root value %d", storage.id.value, storage.root.value); + return status; + } + storages_.push_back(storage); + return status; +} + +Status ValueMap::add_snapshot_value(Value value) { + Status status; + if (value.id.value < 0 || static_cast(value.id.value) != values_.size()) { + status.log("snapshot value id %d is not the next value id %zu", value.id.value, values_.size()); + return status; + } + if (value.storage.value < 0 || static_cast(value.storage.value) >= storages_.size()) { + status.log("snapshot value %d references missing storage %d", value.id.value, value.storage.value); + return status; + } + if (value.storage_root.value < 0 || static_cast(value.storage_root.value) > values_.size()) { + status.log("snapshot value %d references invalid storage root %d", value.id.value, value.storage_root.value); + return status; + } + if (value.alias_source.value >= 0 && static_cast(value.alias_source.value) >= values_.size()) { + status.log("snapshot value %d references missing alias source %d", value.id.value, value.alias_source.value); + return status; + } + value.tensor = nullptr; + value.buffer.reset(); + values_.push_back(std::move(value)); + return status; +} + +const Value * ValueMap::find_tensor(const ggml_tensor * tensor) const { + const auto found = tensor_values_.find(tensor); + if (found == tensor_values_.end()) { + return nullptr; + } + return &values_[found->second]; +} + +} // namespace ggml::hrx diff --git a/ggml/src/ggml-hrx/graph/value-map.h b/ggml/src/ggml-hrx/graph/value-map.h new file mode 100644 index 000000000000..783ffd0519ea --- /dev/null +++ b/ggml/src/ggml-hrx/graph/value-map.h @@ -0,0 +1,127 @@ +#pragma once + +#include "ggml.h" +#include "status.h" + +#include +#include +#include +#include +#include +#include + +struct ggml_tensor; +typedef struct hrx_buffer_s * hrx_buffer_t; + +namespace ggml::hrx { + +struct ValueId { + ValueId() : value(-1) {} + + explicit ValueId(int32_t value) : value(value) {} + + int32_t value; +}; + +struct ValueStorageId { + ValueStorageId() : value(-1) {} + + explicit ValueStorageId(int32_t value) : value(value) {} + + int32_t value; +}; + +inline bool operator==(ValueId lhs, ValueId rhs) { + return lhs.value == rhs.value; +} + +inline bool operator!=(ValueId lhs, ValueId rhs) { + return !(lhs == rhs); +} + +inline bool operator==(ValueStorageId lhs, ValueStorageId rhs) { + return lhs.value == rhs.value; +} + +inline bool operator!=(ValueStorageId lhs, ValueStorageId rhs) { + return !(lhs == rhs); +} + +enum class ValueKind : uint8_t { + External, + Transient, +}; + +struct ValueBufferBinding { + // A buffer is directly bindable by an HRX command program. Host data requires residency or staging before + // execution. These are alternate storage forms and should not both be populated. + hrx_buffer_t buffer = nullptr; + size_t offset = 0; + size_t length = 0; + uint64_t identity = 0; + uint64_t generation = 0; + size_t capacity = 0; + void * host_data = nullptr; + bool weight = false; + + bool requires_materialization() const { return host_data != nullptr; } +}; + +struct Value { + ValueId id; + ValueKind kind; + ValueStorageId storage; + ValueId storage_root; + ValueId alias_source; + size_t storage_offset = 0; + size_t storage_byte_count = 0; + ggml_type type; + std::array ne; + std::array nb; + int64_t element_count = 0; + size_t byte_count = 0; + bool contiguous = false; + const ggml_tensor * tensor = nullptr; + std::optional buffer; +}; + +struct ValueStorage { + ValueStorageId id; + ValueId root; + size_t byte_count = 0; +}; + +class ValueMap { + public: + ValueMap() = default; + + ValueId get_or_add_tensor_value(const ggml_tensor * tensor, ValueKind kind); + + const Value * find(ValueId id) const; + const Value * find_tensor(const ggml_tensor * tensor) const; + const ValueStorage * find_storage(ValueStorageId id) const; + bool bind_buffer(ValueId id, ValueBufferBinding binding); + std::optional resolve_buffer_binding(ValueId id) const; + std::vector external_value_ids() const; + Status alias_storage(ValueId target, ValueId source); + ValueId storage_root(ValueId id) const; + bool same_storage(ValueId lhs, ValueId rhs) const; + + const std::vector & values() const { return values_; } + + const std::vector & storages() const { return storages_; } + + size_t size() const { return values_.size(); } + + Status add_snapshot_storage(ValueStorage storage); + Status add_snapshot_value(Value value); + + private: + const Value * find_alias_source(const ggml_tensor * tensor) const; + + std::vector values_; + std::vector storages_; + std::unordered_map tensor_values_; +}; + +} // namespace ggml::hrx diff --git a/ggml/src/ggml-hrx/hrx-interop-utils.h b/ggml/src/ggml-hrx/hrx-interop-utils.h new file mode 100644 index 000000000000..08d3c2190639 --- /dev/null +++ b/ggml/src/ggml-hrx/hrx-interop-utils.h @@ -0,0 +1,31 @@ +#pragma once + +#include "hrx_runtime.h" + +#include +#include + +namespace ggml::hrx { + +// Success has no payload; failure carries the diagnostic produced by HRX or +// by the caller. Keeping this distinct from an empty string makes status tests +// explicit at API boundaries. +using ErrorResult = std::optional; + +inline ErrorResult take_status(hrx_status_t status) { + if (hrx_status_is_ok(status)) { + return std::nullopt; + } + char * message = nullptr; + size_t length = 0; + hrx_status_t format_status = hrx_status_to_string(status, &message, &length); + if (!hrx_status_is_ok(format_status)) { + hrx_status_ignore(format_status); + } + std::string result = message != nullptr ? std::string(message, length) : std::string("unknown HRX error"); + hrx_status_free_message(message); + hrx_status_ignore(status); + return result; +} + +} // namespace ggml::hrx diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernel-corpus-catalog-verify.h b/ggml/src/ggml-hrx/kernel-corpus/kernel-corpus-catalog-verify.h new file mode 100644 index 000000000000..15af8e0bf4f8 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernel-corpus-catalog-verify.h @@ -0,0 +1,16 @@ +#pragma once + +#include "kernel-corpus-catalog.h" + +namespace ggml::hrx { + +#include "kernel-corpus-catalog.inc" + +} // namespace ggml::hrx + +#define GGML_HRX_KERNEL_REF(family_literal, name_literal) \ + ([] { \ + static_assert(::ggml::hrx::kernel_catalog_entry_exists(family_literal, name_literal), \ + "unknown HRX kernel catalog entry"); \ + return ::ggml::hrx::kernel_catalog_ref(family_literal, name_literal); \ + }()) diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernel-corpus-catalog.h b/ggml/src/ggml-hrx/kernel-corpus/kernel-corpus-catalog.h new file mode 100644 index 000000000000..11e480071261 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernel-corpus-catalog.h @@ -0,0 +1,56 @@ +#pragma once + +#include +#include + +namespace ggml::hrx { + +static constexpr uint64_t kUncatalogedKernelId = 0; + +constexpr bool kernel_catalog_name_equal(const char * lhs, const char * rhs) { + while (*lhs != 0 && *rhs != 0) { + if (*lhs != *rhs) { + return false; + } + ++lhs; + ++rhs; + } + return *lhs == *rhs; +} + +constexpr uint64_t kernel_catalog_id(const char * family, const char * name) { + uint64_t hash = UINT64_C(1469598103934665603); + while (*family != 0) { + hash ^= static_cast(*family); + hash *= UINT64_C(1099511628211); + ++family; + } + hash ^= 0; + hash *= UINT64_C(1099511628211); + while (*name != 0) { + hash ^= static_cast(*name); + hash *= UINT64_C(1099511628211); + ++name; + } + return hash; +} + +struct KernelCatalogRef { + const char * family = ""; + const char * name = ""; + uint64_t id = kUncatalogedKernelId; + + constexpr bool valid() const { + return id != kUncatalogedKernelId && family != nullptr && family[0] != 0 && name != nullptr && name[0] != 0; + } +}; + +constexpr KernelCatalogRef kernel_catalog_ref(const char * family, const char * name) { + return { + family, + name, + family != nullptr && name != nullptr ? kernel_catalog_id(family, name) : kUncatalogedKernelId, + }; +} + +} // namespace ggml::hrx diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernel-corpus-json.cpp b/ggml/src/ggml-hrx/kernel-corpus/kernel-corpus-json.cpp new file mode 100644 index 000000000000..222ee9e259d4 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernel-corpus-json.cpp @@ -0,0 +1,102 @@ +#include "kernel-corpus-json.h" + +#include + +namespace ggml::hrx { +namespace { + +const char * kernel_resource_access_name(ResourceAccess access) { + switch (access) { + case ResourceAccess::Read: + return "read"; + case ResourceAccess::Write: + return "write"; + case ResourceAccess::ReadWrite: + return "read_write"; + } + return "unknown"; +} + +nlohmann::ordered_json string_span_json(KernelSpan values) { + nlohmann::ordered_json result = nlohmann::ordered_json::array(); + for (const char * value : values) { + result.push_back(value != nullptr ? value : ""); + } + return result; +} + +nlohmann::ordered_json source_ref_span_json(KernelSpan values) { + nlohmann::ordered_json result = nlohmann::ordered_json::array(); + for (const KernelSourceRef & value : values) { + result.push_back(value.path != nullptr ? value.path : ""); + } + return result; +} + +nlohmann::ordered_json compile_config_json(KernelSpan values) { + nlohmann::ordered_json result = nlohmann::ordered_json::object(); + for (const KernelCompileConfig & value : values) { + result[value.key != nullptr ? value.key : ""] = value.value != nullptr ? value.value : ""; + } + return result; +} + +} // namespace + +std::string serialize_kernel_corpus_json(const KernelCorpus & corpus) { + nlohmann::ordered_json root = { + { "schema", corpus.schema }, + { "upstream_revision", corpus.upstream_revision }, + { "corpus_digest", corpus.corpus_digest }, + { "recipe_digest", corpus.recipe_digest }, + { "plan_case_count", corpus.plan_case_count }, + { "kernels", nlohmann::ordered_json::array() }, + }; + for (const KernelDefinition & kernel : corpus.kernels) { + nlohmann::ordered_json item = { + { "family", kernel.family }, + { "name", kernel.name }, + { "id", kernel.id }, + { "source", kernel.source }, + { "dependencies", string_span_json(kernel.dependencies) }, + { "symbol", kernel.symbol }, + { "backend", kernel.backend }, + { "target_selector", kernel.target_selector }, + { "compile_config", compile_config_json(kernel.compile_config) }, + { "scalar_parameters", string_span_json(kernel.scalar_parameters) }, + { "source_digest", kernel.source_digest }, + { "compile_recipe", + { + { "mode", kernel.compile_recipe.mode }, + { "link_module", kernel.compile_recipe.link_module }, + { "primary_sources", source_ref_span_json(kernel.compile_recipe.primary_sources) }, + { "library_sources", source_ref_span_json(kernel.compile_recipe.library_sources) }, + } }, + { "workload_parameters", nlohmann::ordered_json::array() }, + { "launch_parameters", nlohmann::ordered_json::array() }, + { "bindings", nlohmann::ordered_json::array() }, + }; + for (const KernelScalarDefinition & parameter : kernel.workload_parameters) { + item["workload_parameters"].push_back({ + { "name", parameter.name }, + { "type", parameter.type } + }); + } + for (const KernelScalarDefinition & parameter : kernel.launch_parameters) { + item["launch_parameters"].push_back({ + { "name", parameter.name }, + { "type", parameter.type } + }); + } + for (const KernelBindingDefinition & binding : kernel.bindings) { + item["bindings"].push_back({ + { "name", binding.name }, + { "access", kernel_resource_access_name(binding.access) } + }); + } + root["kernels"].push_back(std::move(item)); + } + return root.dump(); +} + +} // namespace ggml::hrx diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernel-corpus-json.h b/ggml/src/ggml-hrx/kernel-corpus/kernel-corpus-json.h new file mode 100644 index 000000000000..2edf760b5976 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernel-corpus-json.h @@ -0,0 +1,11 @@ +#pragma once + +#include "kernel-corpus.h" + +#include + +namespace ggml::hrx { + +std::string serialize_kernel_corpus_json(const KernelCorpus & corpus); + +} // namespace ggml::hrx diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernel-corpus.cpp b/ggml/src/ggml-hrx/kernel-corpus/kernel-corpus.cpp new file mode 100644 index 000000000000..5ae64c068dc0 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernel-corpus.cpp @@ -0,0 +1,287 @@ +#include "kernel-corpus.h" + +#include +#include +#include +#include +#include +#include +#include + +namespace ggml::hrx { +namespace { + +const char * kernel_resource_access_name(ResourceAccess access) { + switch (access) { + case ResourceAccess::Read: + return "read"; + case ResourceAccess::Write: + return "write"; + case ResourceAccess::ReadWrite: + return "read_write"; + } + return "unknown"; +} + +struct KernelSourceRecordEntry { + const char * source_path; + const KernelSource * source; +}; + +static bool string_equal(const char * lhs, const char * rhs) { + return std::strcmp(lhs != nullptr ? lhs : "", rhs != nullptr ? rhs : "") == 0; +} + +static bool string_empty(const char * value) { + return value == nullptr || value[0] == 0; +} + +static bool contains_source_ref(KernelSpan values, const char * path) { + return std::find_if(values.begin(), values.end(), + [&](const KernelSourceRef & item) { return string_equal(item.path, path); }) != values.end(); +} + +static bool string_span_equal(KernelSpan lhs, KernelSpan rhs) { + return lhs.size() == rhs.size() && std::equal(lhs.begin(), lhs.end(), rhs.begin(), + [](const char * a, const char * b) { return string_equal(a, b); }); +} + +static bool scalar_span_equal(KernelSpan lhs, KernelSpan rhs) { + return lhs.size() == rhs.size() && + std::equal(lhs.begin(), lhs.end(), rhs.begin(), + [](const KernelScalarDefinition & a, const KernelScalarDefinition & b) { + return string_equal(a.name, b.name) && string_equal(a.type, b.type); + }); +} + +static bool binding_span_equal(KernelSpan lhs, KernelSpan rhs) { + return lhs.size() == rhs.size() && + std::equal(lhs.begin(), lhs.end(), rhs.begin(), + [](const KernelBindingDefinition & a, const KernelBindingDefinition & b) { + return string_equal(a.name, b.name) && a.access == b.access; + }); +} + +static bool kernel_variant_contract_equal(const KernelDefinition & lhs, const KernelDefinition & rhs) { + return string_equal(lhs.backend, rhs.backend) && string_span_equal(lhs.scalar_parameters, rhs.scalar_parameters) && + scalar_span_equal(lhs.workload_parameters, rhs.workload_parameters) && + scalar_span_equal(lhs.launch_parameters, rhs.launch_parameters) && + binding_span_equal(lhs.bindings, rhs.bindings); +} + +// clang-format off +#include "kernel-corpus-sources.inc" +#include "kernel-corpus-qwen.inc" +// clang-format on + +} // namespace + +const KernelSource * get_kernel_source(const char * source_path) { + if (source_path == nullptr) { + return nullptr; + } + for (const KernelSourceRecordEntry & entry : kKernelSourceRecords) { + if (std::strcmp(source_path, entry.source_path) == 0) { + return entry.source; + } + } + return nullptr; +} + +const KernelCorpus & get_qwen_kernel_corpus() { + return kQwenKernelCorpus; +} + +KernelResolveResult resolve_kernel_definition(const KernelCorpus & corpus, + const std::string & target, + uint64_t kernel_id) { + if (kernel_id == kUncatalogedKernelId) { + return { KernelResolveStatus::UncatalogedKernel, nullptr }; + } + const KernelDefinition * first_match = nullptr; + const KernelDefinition * default_variant = nullptr; + bool target_mismatch = false; + for (const KernelDefinition & kernel : corpus.kernels) { + if (kernel.id != kernel_id) { + continue; + } + if (first_match == nullptr) { + first_match = &kernel; + } else if (!string_equal(first_match->family, kernel.family) || !string_equal(first_match->name, kernel.name)) { + return { KernelResolveStatus::HashCollision, nullptr }; + } + if (string_equal(kernel.target_selector, target.c_str())) { + return { KernelResolveStatus::Found, &kernel }; + } + if (string_empty(kernel.target_selector)) { + default_variant = &kernel; + } else { + target_mismatch = true; + } + } + if (default_variant != nullptr) { + return { KernelResolveStatus::Found, default_variant }; + } + if (target_mismatch) { + return { KernelResolveStatus::UnsupportedTarget, first_match }; + } + return { KernelResolveStatus::MissingActiveCorpusEntry, nullptr }; +} + +const char * kernel_resolve_status_name(KernelResolveStatus status) { + switch (status) { + case KernelResolveStatus::Found: + return "found"; + case KernelResolveStatus::UncatalogedKernel: + return "uncataloged_kernel"; + case KernelResolveStatus::MissingActiveCorpusEntry: + return "missing_active_corpus_entry"; + case KernelResolveStatus::HashCollision: + return "hash_collision"; + case KernelResolveStatus::UnsupportedTarget: + return "unsupported_target"; + } + return "unknown"; +} + +std::string kernel_definition_name(const KernelDefinition & definition) { + return std::string(definition.family != nullptr ? definition.family : "") + ":" + + (definition.name != nullptr ? definition.name : ""); +} + +std::string kernel_definition_name_or_id(const KernelDefinition * definition, uint64_t kernel_id) { + if (definition != nullptr) { + return kernel_definition_name(*definition); + } + return "kernel_id=" + std::to_string(kernel_id); +} + +std::string format_kernel_resolve_error(const KernelResolveResult & result, uint64_t kernel_id) { + const std::string label = kernel_definition_name_or_id(result.definition, kernel_id); + switch (result.status) { + case KernelResolveStatus::Found: + return ""; + case KernelResolveStatus::UncatalogedKernel: + return "uncataloged kernel " + label; + case KernelResolveStatus::MissingActiveCorpusEntry: + return "cataloged kernel " + label + " is not available in the active corpus"; + case KernelResolveStatus::HashCollision: + return "kernel catalog id collision while resolving " + label; + case KernelResolveStatus::UnsupportedTarget: + return "cataloged kernel " + label + " has no implementation for the requested target"; + } + return "unknown kernel resolution failure for " + label; +} + +VerificationResult verify_kernel_corpus(const KernelCorpus & corpus) { + VerificationResult result; + if (!string_equal(corpus.schema, "ggml-hrx-kernel-corpus-v2")) { + result.status.log("unsupported kernel corpus schema"); + } + if (string_empty(corpus.upstream_revision)) { + result.status.log("kernel corpus has no upstream revision"); + } + if (string_empty(corpus.corpus_digest)) { + result.status.log("kernel corpus has no digest"); + } + if (string_empty(corpus.recipe_digest)) { + result.status.log("kernel corpus has no BUILD.bazel recipe digest"); + } + if (corpus.plan_case_count == 0) { + result.status.log("kernel corpus has no compile plan cases"); + } + std::set variants; + std::map contracts; + for (const KernelDefinition & kernel : corpus.kernels) { + if (string_empty(kernel.family) || string_empty(kernel.name) || string_empty(kernel.source) || + string_empty(kernel.symbol) || string_empty(kernel.backend) || string_empty(kernel.source_digest)) { + result.status.log("kernel definition is incomplete"); + } + if (kernel.id != kernel_catalog_id(kernel.family != nullptr ? kernel.family : "", + kernel.name != nullptr ? kernel.name : "")) { + result.status.log("kernel %s has an invalid catalog id", kernel.name != nullptr ? kernel.name : ""); + } + const bool source_is_primary = contains_source_ref(kernel.compile_recipe.primary_sources, kernel.source); + const bool source_is_library = contains_source_ref(kernel.compile_recipe.library_sources, kernel.source); + if ((!string_equal(kernel.compile_recipe.mode, "direct") && + !string_equal(kernel.compile_recipe.mode, "archive")) || + kernel.compile_recipe.primary_sources.empty() || (!source_is_primary && !source_is_library) || + (string_equal(kernel.compile_recipe.mode, "archive") && string_empty(kernel.compile_recipe.link_module))) { + result.status.log("kernel %s has an invalid BUILD compile recipe", + kernel.name != nullptr ? kernel.name : ""); + } + for (const KernelSourceRef & source : kernel.compile_recipe.primary_sources) { + if (string_empty(source.path) || source.contents == nullptr) { + result.status.log("kernel %s has an invalid embedded primary source reference", + kernel.name != nullptr ? kernel.name : ""); + } + } + for (const KernelSourceRef & source : kernel.compile_recipe.library_sources) { + if (string_empty(source.path) || source.contents == nullptr) { + result.status.log("kernel %s has an invalid embedded library source reference", + kernel.name != nullptr ? kernel.name : ""); + } + } + const std::string full_name = std::string(kernel.family != nullptr ? kernel.family : "") + ":" + + std::string(kernel.name != nullptr ? kernel.name : ""); + const std::string target_selector = kernel.target_selector != nullptr ? kernel.target_selector : ""; + if (!variants.insert(full_name + "@" + target_selector).second) { + result.status.log("kernel corpus repeats target variant %s@%s", full_name.c_str(), + target_selector.empty() ? "default" : target_selector.c_str()); + } + const auto contract = contracts.emplace(full_name, &kernel); + if (!contract.second && !kernel_variant_contract_equal(*contract.first->second, kernel)) { + result.status.log("kernel target variants disagree on ABI for %s", full_name.c_str()); + } + std::set binding_names; + for (const KernelBindingDefinition & binding : kernel.bindings) { + if (string_empty(binding.name) || + !binding_names.insert(binding.name != nullptr ? binding.name : "").second) { + result.status.log("kernel %s has invalid binding names", kernel.name != nullptr ? kernel.name : ""); + } + } + if (kernel.bindings.size() == 0) { + result.status.log("kernel %s has no binding ABI", kernel.name != nullptr ? kernel.name : ""); + } + } + return result; +} + +std::string format_kernel_corpus(const KernelCorpus & corpus) { + std::ostringstream out; + out << "kernel-corpus " << corpus.schema << " revision=" << corpus.upstream_revision + << " digest=" << corpus.corpus_digest << " recipe=" << corpus.recipe_digest + << " kernels=" << corpus.kernels.size() << " plan_cases=" << corpus.plan_case_count << '\n'; + for (const KernelDefinition & kernel : corpus.kernels) { + out << " kernel " << kernel.family << ':' << kernel.name << " id=0x" << std::hex << kernel.id << std::dec + << " backend=" << kernel.backend + << " target=" << (string_empty(kernel.target_selector) ? "default" : kernel.target_selector) << " symbol=@" + << kernel.symbol << " source=" << kernel.source << " sha256=" << kernel.source_digest << '\n'; + out << " recipe " << kernel.compile_recipe.mode; + if (!string_empty(kernel.compile_recipe.link_module)) { + out << " module=" << kernel.compile_recipe.link_module; + } + out << " primary="; + for (size_t i = 0; i < kernel.compile_recipe.primary_sources.size(); ++i) { + out << (i ? "," : "") << kernel.compile_recipe.primary_sources[i].path; + } + out << " libraries="; + for (size_t i = 0; i < kernel.compile_recipe.library_sources.size(); ++i) { + out << (i ? "," : "") << kernel.compile_recipe.library_sources[i].path; + } + out << '\n'; + for (size_t i = 0; i < kernel.bindings.size(); ++i) { + out << " binding[" << i << "] " << kernel.bindings[i].name << ' ' + << kernel_resource_access_name(kernel.bindings[i].access) << '\n'; + } + for (const KernelScalarDefinition & parameter : kernel.workload_parameters) { + out << " workload " << parameter.name << ' ' << parameter.type << '\n'; + } + for (const KernelScalarDefinition & parameter : kernel.launch_parameters) { + out << " launch " << parameter.name << ' ' << parameter.type << '\n'; + } + } + return out.str(); +} + +} // namespace ggml::hrx diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernel-corpus.h b/ggml/src/ggml-hrx/kernel-corpus/kernel-corpus.h new file mode 100644 index 000000000000..501eae8c19fa --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernel-corpus.h @@ -0,0 +1,129 @@ +#pragma once + +#include "kernel-corpus-catalog.h" +#include "kernel-types.h" + +#include +#include +#include + +namespace ggml::hrx { + +template struct KernelSpan { + const T * items = nullptr; + size_t count = 0; + + const T * begin() const { return items; } + + const T * end() const { return items == nullptr ? nullptr : items + count; } + + const T * data() const { return items; } + + size_t size() const { return count; } + + bool empty() const { return count == 0; } + + const T & operator[](size_t index) const { return items[index]; } + + const T & front() const { return items[0]; } +}; + +struct KernelCompileConfig { + const char * key = ""; + const char * value = ""; +}; + +struct KernelBindingDefinition { + const char * name = ""; + ResourceAccess access = ResourceAccess::Read; +}; + +struct KernelScalarDefinition { + const char * name = ""; + const char * type = ""; +}; + +enum KernelSourceFormat { + KERNEL_SOURCE_FORMAT_TEXT, + KERNEL_SOURCE_FORMAT_BINARY, +}; + +struct KernelSourceSpan { + const char * data; + size_t length; + KernelSourceFormat format; +}; + +struct KernelSource { + KernelSourceSpan source; + const KernelSourceSpan * dependencies; + size_t dependency_count; +}; + +struct KernelSourceRef { + const char * path = ""; + const KernelSource * contents = nullptr; +}; + +struct KernelCompileRecipe { + const char * mode = ""; + const char * link_module = ""; + KernelSpan primary_sources; + KernelSpan library_sources; +}; + +struct KernelDefinition { + const char * family = ""; + const char * name = ""; + uint64_t id = kUncatalogedKernelId; + const char * source = ""; + KernelSpan dependencies; + const char * symbol = ""; + const char * backend = ""; + const char * target_selector = ""; + KernelSpan compile_config; + KernelSpan scalar_parameters; + KernelSpan bindings; + const char * source_digest = ""; + KernelSpan workload_parameters; + KernelSpan launch_parameters; + KernelCompileRecipe compile_recipe; +}; + +struct KernelCorpus { + const char * schema = "ggml-hrx-kernel-corpus-v2"; + const char * upstream_revision = ""; + const char * corpus_digest = ""; + const char * recipe_digest = ""; + size_t plan_case_count = 0; + KernelSpan kernels; +}; + +enum class KernelResolveStatus : uint8_t { + Found, + UncatalogedKernel, + MissingActiveCorpusEntry, + HashCollision, + UnsupportedTarget, +}; + +struct KernelResolveResult { + KernelResolveStatus status = KernelResolveStatus::MissingActiveCorpusEntry; + const KernelDefinition * definition = nullptr; + + bool found() const { return status == KernelResolveStatus::Found && definition != nullptr; } +}; + +const KernelSource * get_kernel_source(const char * source_path); +const KernelCorpus & get_qwen_kernel_corpus(); +KernelResolveResult resolve_kernel_definition(const KernelCorpus & corpus, + const std::string & target, + uint64_t kernel_id); +const char * kernel_resolve_status_name(KernelResolveStatus status); +std::string kernel_definition_name(const KernelDefinition & definition); +std::string kernel_definition_name_or_id(const KernelDefinition * definition, uint64_t kernel_id); +std::string format_kernel_resolve_error(const KernelResolveResult & result, uint64_t kernel_id); +VerificationResult verify_kernel_corpus(const KernelCorpus & corpus); +std::string format_kernel_corpus(const KernelCorpus & corpus); + +} // namespace ggml::hrx diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernel-types.h b/ggml/src/ggml-hrx/kernel-corpus/kernel-types.h new file mode 100644 index 000000000000..eb63504cb07f --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernel-types.h @@ -0,0 +1,22 @@ +#pragma once + +#include "status.h" + +#include +#include + +namespace ggml::hrx { + +enum class ResourceAccess : uint8_t { + Read, + Write, + ReadWrite, +}; + +struct VerificationResult { + Status status; + + bool valid() const { return status.success(); } +}; + +} // namespace ggml::hrx diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/hrx_owned/add_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/hrx_owned/add_f32.loom new file mode 100644 index 000000000000..1797b01935c7 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/hrx_owned/add_f32.loom @@ -0,0 +1,47 @@ +// Copyright 2026 The HRX Authors +// SPDX-License-Identifier: Apache-2.0 + +amdgpu.target @ggml_add_gfx11_wave64 {subgroup_size = 64} + +kernel.def target(@ggml_add_gfx11_wave64) export("ggml_add_f32") @ggml_add_f32(%element_count: index) { + %one = index.constant 1 : index + %twofiftysix = index.constant 256 : index + %rounding = index.constant 255 : index + %rounded = index.add %element_count, %rounding : index + %workgroup_count = index.div %rounded, %twofiftysix : index + kernel.launch.config workgroups(%workgroup_count, %one, %one) workgroup_size(%twofiftysix, %one, %one) : index +} launch(%element_count: index, %a: buffer, %b: buffer, %output: buffer) { + %count = index.assume %element_count [range(%element_count, 1, 134217728)] : index + %workgroup = kernel.workgroup.id : index + %workitem = kernel.workitem.id : index + %twofiftysix = index.constant 256 : index + %base0 = index.mul %workgroup, %twofiftysix : index + %linear0 = index.add %base0, %workitem : index + %linear = index.assume %linear0 [range(%linear0, 0, 134217983)] : index + %in_bounds = index.cmp ult, %linear, %count : index + %zero_offset = index.constant 0 : offset + %a_noalias, %b_noalias, %output_noalias = buffer.assume.noalias %a, %b, %output : buffer, buffer, buffer + %a_view = buffer.view %a_noalias[%zero_offset] : buffer -> view<[%count]xf32> + %b_view = buffer.view %b_noalias[%zero_offset] : buffer -> view<[%count]xf32> + %output_view = buffer.view %output_noalias[%zero_offset] : buffer -> view<[%count]xf32> + scf.if %in_bounds { + %a_value = view.load %a_view[%linear] : view<[%count]xf32> -> f32 + %b_value = view.load %b_view[%linear] : view<[%count]xf32> -> f32 + %sum = scalar.addf %a_value, %b_value : f32 + view.store %sum, %output_view[%linear] : f32, view<[%count]xf32> + } + kernel.return +} + +check.case public @ggml_add_f32_small_case { + %four = check.literal value(4) : index + %a = check.generate.iota offset(1) step(1) : tensor<4xf32> + %b = check.generate.fill value(10) : tensor<4xf32> + %output = check.generate.fill value(0) : tensor<4xf32> + %expected = check.generate.iota offset(11) step(1) : tensor<4xf32> + func.call @ggml_add_f32(%four, %a, %b, %output) : (index, tensor<4xf32>, tensor<4xf32>, tensor<4xf32>) + check.expect.close actual(%output) expected(%expected) atol(0.0) rtol(0.0) nan(same) : tensor<4xf32> + check.return +} + +check.benchmark<@ggml_add_f32_small_case> @ggml_add_f32_small diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/hrx_owned/dequant_iq3xxs_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/hrx_owned/dequant_iq3xxs_f32.loom new file mode 100644 index 000000000000..241f919bccf0 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/hrx_owned/dequant_iq3xxs_f32.loom @@ -0,0 +1,122 @@ +// Copyright 2026 The HRX Authors +// SPDX-License-Identifier: Apache-2.0 + +// Dequantize GGML IQ3_XXS weight blocks to F32 (see ggml dequantize_row_iq3_xxs). +// +// block_iq3_xxs = 98 bytes / 256 values: f16 d, qs[64] grid indices, +// scales_and_signs[32] (4 bytes per 32-value ib32: top nibble scale, 4x7-bit signs). +// One workgroup per block, one workitem per decoded value. +amdgpu.target @ggml_iq3xxs_dequant_gfx11_wave64 {subgroup_size = 64} + +global.rodata.def @iq3xxs_grid = align(16) bytes("0404040414040404240404040c0c04041c0c04043e0c040404140404141404040c1c0404142404041c3e04042c3e04040c040c041c040c04040c0c04140c0c040c140c042c140c04041c0c04141c0c040c240c04242c0c04043e0c040404140414041404240414040c0c140404141404141414040c1c14041c1c14043e1c14040c2c14043e2c14042c3e14040c041c043e041c04040c1c04140c1c042c141c04043e1c041c0c24043e1c2404242424043e2c24041c3e24042c3e24040c042c043e042c04141c2c04142c2c042c1c340424343404040c3e04240c3e04340c3e041c243e040c343e040c04040c1c04040c040c040c140c040c0c14040c1c14040c041c040c141c040c241c040c3e24040c042c040c04040c0c14040c0c0c0c0c0c04140c0c14140c0c0c04140c1c04140c040c140c140c140c0c14140c041c140c143e140c04041c0c14041c0c04141c0c0c1c1c0c34241c0c34341c0c0c04240c2c04240c042c240c04142c0c24142c0c34242c0c0c3e2c0c2c04340c14143e0c04243e0c04040414140404140c0c04141c0c04140414041414140414341404140c1c0414142404140c040c141c040c142c040c14040c0c14140c0c140c140c14041c0c141c340c143e340c14043e0c1404041414140414140c0c14143e0c141404141414141414143e1c1414042414142c2c14140c041c14040c1c14240c1c14043e1c14243e1c142c1c24141c2c24141c042c143e142c140c242c14243e2c140c043e141c043e14340c3e142c243e140c04041c040c041c140c041c0c14041c1c14041c042c041c2c34041c143e041c04040c1c14040c1c04140c1c0c1c0c1c24240c1c34240c1c0c04141c1c04141c040c141c2c14141c142c141c143e141c0c0c1c1c1c1c1c1c041c241c3e24241c143e241c04042c1c34042c1c14142c1c2c2c2c1c240c341c341c341c1c34341c1c1c3e1c04343e1c240404243e0c04242c1c04243e1c04241c2c04243e2c0424243e0c24041414243e1c14240424142404341424343414243e041c242c241c24240424240c2c2424243424242c142c241c242c24043e2c242c043e24040c3e24140c3e24041c3e24140c042c0c24042c043e042c04040c2c34040c2c34140c2c2c2c0c2c240c142c141c142c143e142c14041c2c1c2c1c2c040c242c1c14242c3e14242c143e242c14042c2c0c1c2c2c042c342c24143e2c14243e2c241404342424043434240434243404340c140c340c340c343e0c143424341434041c1c34341c1c34242424342c042c34142c2c341c1c34341c043e340c143e341c04043e2c04043e3e04043e040c043e141c043e142c043e34140c3e04240c3e140c143e2c24143e142c143e04041c3e2c0c1c3e1c1c1c3e04341c3e0c14243e0c24243e04042c3e14042c3e24142c3e041c343e") + +global.rodata.def @ksigns_iq2xs = align(16) bytes("008182038405068788090a8b0c8d8e0f901112931495961718999a1b9c1d1e9fa02122a324a5a62728a9aa2bac2d2eaf30b1b233b43536b7b8393abb3cbdbe3fc04142c344c5c64748c9ca4bcc4d4ecf50d1d253d45556d7d8595adb5cddde5f60e1e263e46566e7e8696aeb6cedee6ff07172f374f5f67778f9fa7bfc7d7eff") + +global.rodata.def @kmask_iq2xs = align(8) bytes("0102040810204080") + +kernel.def target(@ggml_iq3xxs_dequant_gfx11_wave64) export("ggml_dequant_iq3xxs_f32") @ggml_dequant_iq3xxs_f32(%block_count: index) { + %c1 = index.constant 1 : index + %c256 = index.constant 256 : index + kernel.launch.config workgroups(%block_count, %c1, %c1) workgroup_size(%c256, %c1, %c1) : index +} launch(%block_count: index, %weight: buffer, %output: buffer) { + %c0 = index.constant 0 : index + %c2 = index.constant 2 : index + %c4 = index.constant 4 : index + %c8 = index.constant 8 : index + %c7 = index.constant 7 : index + %c32 = index.constant 32 : index + %c256 = index.constant 256 : index + %o0 = index.constant 0 : offset + %o1 = index.constant 1 : offset + %o2 = index.constant 2 : offset + %o4 = index.constant 4 : offset + %o66 = index.constant 66 : offset + %o98 = index.constant 98 : offset + %c255 = index.constant 255 : index + %block = kernel.workgroup.id : index + %value = kernel.workitem.id : index + %weight_na, %output_na = buffer.assume.noalias %weight, %output : buffer, buffer + %block_off = index.scale %block, %o98 : index, offset -> offset + %qs_base = index.add %block_off, %o2 : offset + %ss_base = index.add %block_off, %o66 : offset + %d_view = buffer.view %weight_na[%block_off] : buffer -> view<1xf16> + %d_f16 = view.load %d_view[%c0] : view<1xf16> -> f16 + %d = scalar.extf %d_f16 : f16 to f32 + %ib32 = index.div %value, %c32 : index + %rem = index.rem %value, %c32 : index + %l = index.div %rem, %c8 : index + %p = index.rem %rem, %c8 : index + %half = index.div %p, %c4 : index + %jj = index.rem %p, %c4 : index + %qs_a = index.mul %ib32, %c8 : index + %qs_b = index.mul %l, %c2 : index + %qs_c = index.add %qs_a, %qs_b : index + %qs_index = index.add %qs_c, %half : index + %qs_off = index.scale %qs_index, %o1 : index, offset -> offset + %qs_ptr = index.add %qs_base, %qs_off : offset + %qs_view = buffer.view %weight_na[%qs_ptr] : buffer -> view<1xi8> + %grid_idx_i8 = view.load %qs_view[%c0] : view<1xi8> -> i8 + %grid_idx = index.cast %grid_idx_i8 : i8 to index + %g_a = index.mul %grid_idx, %c4 : index + %g_off_i = index.add %g_a, %jj : index + %g_off = index.scale %g_off_i, %o1 : index, offset -> offset + %grid_base = global.load @iq3xxs_grid : buffer + %grid_view = buffer.view %grid_base[%g_off] : buffer -> view<1xi8> + %grid_byte_i8 = view.load %grid_view[%c0] : view<1xi8> -> i8 + %grid_byte = scalar.uitofp %grid_byte_i8 : i8 to f32 + %ss_a = index.mul %ib32, %c4 : index + %ss_off = index.scale %ss_a, %o1 : index, offset -> offset + %ss_ptr = index.add %ss_base, %ss_off : offset + %aux_view = buffer.view %weight_na[%ss_ptr] : buffer -> view<1xi32> + %aux32 = vector.load %aux_view[%c0] : view<1xi32> -> vector<1xi32> + %c28 = vector.constant 28 : vector<1xi32> + %scale_nib = vector.shrui %aux32, %c28 : vector<1xi32> + %scale_i32 = vector.extract %scale_nib[0] : vector<1xi32> -> i32 + %scale_f = scalar.uitofp %scale_i32 : i32 to f32 + %half_c = scalar.constant 0.5 : f32 + %scale_plus = scalar.addf %half_c, %scale_f : f32 + %d_half = scalar.mulf %d, %half_c : f32 + %db = scalar.mulf %d_half, %scale_plus : f32 + %sh7 = index.mul %l, %c7 : index + %sh7_i32 = index.cast %sh7 : index to i32 + %sh_v = vector.splat %sh7_i32 : vector<1xi32> + %signs_shifted = vector.shrui %aux32, %sh_v : vector<1xi32> + %c127 = vector.constant 127 : vector<1xi32> + %signs_masked = vector.andi %signs_shifted, %c127 : vector<1xi32> + %signs_idx_i32 = vector.extract %signs_masked[0] : vector<1xi32> -> i32 + %signs_idx = index.cast %signs_idx_i32 : i32 to index + %signs_off = index.scale %signs_idx, %o1 : index, offset -> offset + %ks_base = global.load @ksigns_iq2xs : buffer + %ks_view = buffer.view %ks_base[%signs_off] : buffer -> view<1xi8> + %signs_i8 = view.load %ks_view[%c0] : view<1xi8> -> i8 + %signs_i32 = index.cast %signs_i8 : i8 to index + %km_off = index.scale %p, %o1 : index, offset -> offset + %km_base = global.load @kmask_iq2xs : buffer + %km_view = buffer.view %km_base[%km_off] : buffer -> view<1xi8> + %kmask_p_i8 = view.load %km_view[%c0] : view<1xi8> -> i8 + %kmask_p = index.cast %kmask_p_i8 : i8 to index + %and_bits = index.andi %signs_i32, %kmask_p : index + %is_neg = index.cmp ne, %and_bits, %c0 : index + %neg_one = scalar.constant -1.0 : f32 + %pos_one = scalar.constant 1.0 : f32 + %sign_f = scf.select %is_neg, %neg_one, %pos_one : f32 + %db_g = scalar.mulf %db, %grid_byte : f32 + %result = scalar.mulf %db_g, %sign_f : f32 + %out_index = index.mul %block, %c256 : index + %out_index2 = index.add %out_index, %value : index + %out_off = index.scale %out_index2, %o4 : index, offset -> offset + %out_f32 = buffer.view %output_na[%out_off] : buffer -> view<1xf32> + view.store %result, %out_f32[%c0] : f32, view<1xf32> + kernel.return +} + +check.case public @ggml_dequant_iq3xxs_f32_coverage { + %one = check.literal value(1) : index + %weight = check.generate.fill value(0) : tensor<1x98xi8> + %output = check.generate.fill value(0.0) : tensor<1x256xf32> + kernel.launch @ggml_dequant_iq3xxs_f32[%one](%one, %weight, %output) : [index](index, tensor<1x98xi8>, tensor<1x256xf32>) + check.expect.close actual(%output) expected(%output) atol(0.0) rtol(0.0) nan(same) : tensor<1x256xf32> + check.return +} + +check.benchmark<@ggml_dequant_iq3xxs_f32_coverage> @ggml_dequant_iq3xxs_f32_benchmark diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/hrx_owned/gather_add_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/hrx_owned/gather_add_f32.loom new file mode 100644 index 000000000000..6bef54908bc7 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/hrx_owned/gather_add_f32.loom @@ -0,0 +1,68 @@ +// Copyright 2026 The IREE Authors +// +// Licensed under the Apache License v2.0 with LLVM Exceptions. +// See https://llvm.org/LICENSE.txt for license information. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +// Publishes an indexed residual without assuming that the requested rows are +// contiguous or at the end of the source. This is the physical form of two +// GGML GET_ROWS operations followed by ADD. +amdgpu.target @ggml_gather_add_gfx11_wave64 {subgroup_size = 64} + +kernel.def target(@ggml_gather_add_gfx11_wave64) export("ggml_gather_add_f32") @ggml_gather_add_f32( + %source_token_count: index, %output_token_count: index, %hidden_size: index) { + %one = index.constant 1 : index + %twofiftysix = index.constant 256 : index + %rounding = index.constant 255 : index + %rounded_width = index.add %hidden_size, %rounding : index + %column_workgroup_count = index.div %rounded_width, %twofiftysix : index + kernel.launch.config workgroups(%column_workgroup_count, %output_token_count, %one) workgroup_size(%twofiftysix, %one, %one) : index +} launch(%source_token_count: index, %output_token_count: index, %hidden_size: index, + %attention: buffer, %residual: buffer, %output_ids: buffer, %output: buffer) { + %source_count = index.assume %source_token_count [range(%source_token_count, 1, 2048)] : index + %output_count = index.assume %output_token_count [range(%output_token_count, 1, 2048)] : index + %width = index.assume %hidden_size [range(%hidden_size, 1, 65536)] : index + %column_workgroup = kernel.workgroup.id : index + %output_row0 = kernel.workgroup.id : index + %workitem = kernel.workitem.id : index + %twofiftysix = index.constant 256 : index + %output_row, %launch_output_count = index.assume %output_row0, %output_count [lt(%output_row0, %output_count)] : index, index + %base0 = index.mul %column_workgroup, %twofiftysix : index + %column0 = index.add %base0, %workitem : index + %column = index.assume %column0 [range(%column0, 0, 65535)] : index + %in_bounds = index.cmp ult, %column, %width : index + %zero_offset = index.constant 0 : offset + %attention_noalias, %residual_noalias, %ids_noalias, %output_noalias = buffer.assume.noalias %attention, %residual, %output_ids, %output : buffer, buffer, buffer, buffer + %attention_view = buffer.view %attention_noalias[%zero_offset] : buffer -> view<[%source_count]x[%width]xf32> + %residual_view = buffer.view %residual_noalias[%zero_offset] : buffer -> view<[%source_count]x[%width]xf32> + %ids_view = buffer.view %ids_noalias[%zero_offset] : buffer -> view<[%launch_output_count]xi32> + %output_view = buffer.view %output_noalias[%zero_offset] : buffer -> view<[%launch_output_count]x[%width]xf32> + scf.if %in_bounds { + %safe_column = index.assume %column [lt(%column, %width)] : index + %source_row_i32 = view.load %ids_view[%output_row] : view<[%launch_output_count]xi32> -> i32 + %source_row0 = index.cast %source_row_i32 : i32 to index + %source_row = index.assume %source_row0 [range(%source_row0, 0, 2047), lt(%source_row0, %source_count)] : index + %attention_value = view.load %attention_view[%source_row, %safe_column] : view<[%source_count]x[%width]xf32> -> f32 + %residual_value = view.load %residual_view[%source_row, %safe_column] : view<[%source_count]x[%width]xf32> -> f32 + %sum = scalar.addf %attention_value, %residual_value : f32 + view.store %sum, %output_view[%output_row, %safe_column] : f32, view<[%launch_output_count]x[%width]xf32> + } + kernel.return +} + +// Select source row 2 rather than a positional tail inferred by the host. +check.case public @ggml_gather_add_noncontiguous_case { + %one = check.literal value(1) : index + %three = check.literal value(3) : index + %four = check.literal value(4) : index + %attention = check.generate.iota offset(0) step(1) : tensor<3x4xf32> + %residual = check.generate.fill value(100) : tensor<3x4xf32> + %output_ids = check.generate.fill value(2) : tensor<1xi32> + %output = check.generate.fill value(0) : tensor<1x4xf32> + %expected = check.generate.iota offset(108) step(1) : tensor<1x4xf32> + func.call @ggml_gather_add_f32(%three, %one, %four, %attention, %residual, %output_ids, %output) : (index, index, index, tensor<3x4xf32>, tensor<3x4xf32>, tensor<1xi32>, tensor<1x4xf32>) + check.expect.close actual(%output) expected(%expected) atol(0.0) rtol(0.0) nan(same) : tensor<1x4xf32> + check.return +} + +check.benchmark<@ggml_gather_add_noncontiguous_case> @ggml_gather_add_noncontiguous diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/hrx_owned/get_rows_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/hrx_owned/get_rows_f32.loom new file mode 100644 index 000000000000..69a5f8961d53 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/hrx_owned/get_rows_f32.loom @@ -0,0 +1,100 @@ +// Copyright 2026 The IREE Authors +// +// Licensed under the Apache License v2.0 with LLVM Exceptions. +// See https://llvm.org/LICENSE.txt for license information. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +// Generic GGML GET_ROWS: output[row, col] = source[ids[row], col] for an +// arbitrary (non-contiguous) row index list. This is the embedding lookup +// every model needs at the start of the graph; the fused Qwen3-MoE dispatch +// handles it inside fused kernels, and this standalone kernel covers the +// unfused (dense) graphs that otherwise fail-closed. +// +// source_format: 0 = f32 rows ([rows]x[width] f32), 1 = q8_0 rows (each row +// is width/32 block_q8_0: f16 scale + 32 int8 quants; width % 32 == 0). +amdgpu.target @ggml_get_rows_gfx11_wave64 {subgroup_size = 64} + +kernel.def target(@ggml_get_rows_gfx11_wave64) export("ggml_get_rows_f32") @ggml_get_rows_f32( + %source_row_count: index, %output_row_count: index, %width: index) { + %one = index.constant 1 : index + %twofiftysix = index.constant 256 : index + %rounding = index.constant 255 : index + %rounded_width = index.add %width, %rounding : index + %column_workgroup_count = index.div %rounded_width, %twofiftysix : index + kernel.launch.config workgroups(%column_workgroup_count, %output_row_count, %one) workgroup_size(%twofiftysix, %one, %one) : index +} launch(%source_row_count: index, %output_row_count: index, %width: index, %source_format: i32, + %source: buffer, %output_ids: buffer, %output: buffer) { + %source_count = index.assume %source_row_count [range(%source_row_count, 1, 131072)] : index + %output_count = index.assume %output_row_count [range(%output_row_count, 1, 2048)] : index + %cols = index.assume %width [range(%width, 32, 16384)] : index + %column_workgroup = kernel.workgroup.id : index + %output_row0 = kernel.workgroup.id : index + %workitem = kernel.workitem.id : index + %twofiftysix = index.constant 256 : index + %output_row, %launch_output_count = index.assume %output_row0, %output_count [lt(%output_row0, %output_count)] : index, index + %base0 = index.mul %column_workgroup, %twofiftysix : index + %column0 = index.add %base0, %workitem : index + %column = index.assume %column0 [range(%column0, 0, 65535)] : index + %in_bounds = index.cmp ult, %column, %cols : index + %zero_offset = index.constant 0 : offset + %source_noalias, %ids_noalias, %output_noalias = buffer.assume.noalias %source, %output_ids, %output : buffer, buffer, buffer + %ids_view = buffer.view %ids_noalias[%zero_offset] : buffer -> view<[%launch_output_count]xi32> + %output_view = buffer.view %output_noalias[%zero_offset] : buffer -> view<[%launch_output_count]x[%cols]xf32> + scf.if %in_bounds { + %safe_column = index.assume %column [lt(%column, %cols)] : index + %source_row_i32 = view.load %ids_view[%output_row] : view<[%launch_output_count]xi32> -> i32 + %source_row0 = index.cast %source_row_i32 : i32 to index + %source_row = index.assume %source_row0 [range(%source_row0, 0, 131071), lt(%source_row0, %source_count)] : index + %source_format_idx = index.cast %source_format : i32 to index + %c0_idx = index.constant 0 : index + %is_f32 = index.cmp eq, %source_format_idx, %c0_idx : index + scf.if %is_f32 { + %source_view = buffer.view %source_noalias[%zero_offset] : buffer -> view<[%source_count]x[%cols]xf32> + %value = view.load %source_view[%source_row, %safe_column] : view<[%source_count]x[%cols]xf32> -> f32 + view.store %value, %output_view[%output_row, %safe_column] : f32, view<[%launch_output_count]x[%cols]xf32> + } else { + // q8_0: row stride = width/32 blocks x (2 bytes f16 scale + 32 int8). + %c32 = index.constant 32 : index + %c2 = index.constant 2 : offset + %block_count = index.div %cols, %c32 : index + %c34_off = index.constant 34 : offset + %row_bytes_off = index.scale %block_count, %c34_off : index, offset -> offset + %row_base = index.scale %source_row, %row_bytes_off : index, offset -> offset + %block0 = index.div %safe_column, %c32 : index + %block = index.assume %block0 [lt(%block0, %block_count)] : index + %block_off = index.scale %block, %c34_off : index, offset -> offset + %block_base = index.add %row_base, %block_off : offset + %c0_idx = index.constant 0 : index + %dm_view = buffer.view %source_noalias[%block_base] : buffer -> view<2xf16> + %dm = vector.load %dm_view[%c0_idx] : view<2xf16> -> vector<2xf16> + %d_f16 = vector.extract %dm[0] : vector<2xf16> -> f16 + %d = scalar.extf %d_f16 : f16 to f32 + %col_in_block0 = index.rem %safe_column, %c32 : index + %col_in_block = index.assume %col_in_block0 [lt(%col_in_block0, %c32)] : index + %qs_offset = index.add %block_base, %c2 : offset + %qs_view = buffer.view %source_noalias[%qs_offset] : buffer -> view<32xi8> + %q_i8 = view.load %qs_view[%col_in_block] : view<32xi8> -> i8 + %q = scalar.sitofp %q_i8 : i8 to f32 + %value = scalar.mulf %d, %q : f32 + view.store %value, %output_view[%output_row, %safe_column] : f32, view<[%launch_output_count]x[%cols]xf32> + } + } + kernel.return +} + +// Check: f32 source, identity ids (output rows 0..1 = source rows 0..1). +check.case public @ggml_get_rows_f32_identity_case { + %two = check.literal value(2) : index + %three = check.literal value(3) : index + %four = check.literal value(4) : index + %fmt = check.literal value(0) : i32 + %source = check.generate.iota offset(0) step(1) : tensor<3x4xf32> + %output_ids = check.generate.iota offset(0) step(1) : tensor<2xi32> + %output = check.generate.fill value(0) : tensor<2x4xf32> + %expected = check.generate.iota offset(0) step(1) : tensor<2x4xf32> + func.call @ggml_get_rows_f32(%three, %two, %four, %fmt, %source, %output_ids, %output) : (index, index, index, i32, tensor<3x4xf32>, tensor<2xi32>, tensor<2x4xf32>) + check.expect.close actual(%output) expected(%expected) atol(0.0) rtol(0.0) nan(same) : tensor<2x4xf32> + check.return +} + +check.benchmark<@ggml_get_rows_f32_identity_case> @ggml_get_rows_f32_identity diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/hrx_owned/mul_mat_vec_iq3xxs_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/hrx_owned/mul_mat_vec_iq3xxs_f32.loom new file mode 100644 index 000000000000..ea93fc4c8246 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/hrx_owned/mul_mat_vec_iq3xxs_f32.loom @@ -0,0 +1,157 @@ +// Copyright 2026 The HRX Authors +// SPDX-License-Identifier: Apache-2.0 + +// IQ3_XXS matrix-vector product: out[row] = dot(weight[row,:], activation[:]). +// +// One workitem per output row; each row loops over its 256-value IQ3_XXS blocks +// and decodes them with the GGML grid/sign codebook (see dequantize_row_iq3_xxs). +amdgpu.target @ggml_mmviq3_gfx11_wave64 {subgroup_size = 64} + +kernel.def target(@ggml_mmviq3_gfx11_wave64) export("ggml_mul_mat_vec_iq3xxs_f32") @ggml_mul_mat_vec_iq3xxs_f32(%input_size: index, %output_size: index, %token_count: index) { + %c1 = index.constant 1 : index + %c63 = index.constant 63 : index + %c64 = index.constant 64 : index + %rows_plus = index.add %output_size, %c63 : index + %wgx = index.div %rows_plus, %c64 : index + kernel.launch.config workgroups(%wgx, %token_count, %c1) workgroup_size(%c64, %c1, %c1) : index +} launch(%input_size: index, %output_size: index, %token_count: index, %activation: buffer, %weight: buffer, %tables: buffer, %output: buffer) { + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c4 = index.constant 4 : index + %c8 = index.constant 8 : index + %c7 = index.constant 7 : index + %c32 = index.constant 32 : index + %c64 = index.constant 64 : index + %c98 = index.constant 98 : index + %c256 = index.constant 256 : index + %o1 = index.constant 1 : offset + %o2 = index.constant 2 : offset + %o4 = index.constant 4 : offset + %o66 = index.constant 66 : offset + %o98 = index.constant 98 : offset + %o1024 = index.constant 1024 : offset + %o1152 = index.constant 1152 : offset + %c255 = index.constant 255 : index + %zero = scalar.constant 0.0 : f32 + %halfc = scalar.constant 0.5 : f32 + %negone = scalar.constant -1.0 : f32 + %posone = scalar.constant 1.0 : f32 + %wg = kernel.workgroup.id : index + %tok = kernel.workgroup.id : index + %lane = kernel.workitem.id : index + %rowi = index.mul %wg, %c64 : index + %row = index.add %rowi, %lane : index + %row_ok = index.cmp ult, %row, %output_size : index + %activation_na, %weight_na, %output_na, %tables_na = buffer.assume.noalias %activation, %weight, %output, %tables : buffer, buffer, buffer, buffer + %blocks = index.div %input_size, %c256 : index + %row_bytes_i = index.mul %blocks, %c98 : index + %row_bytes = index.scale %row_bytes_i, %o1 : index, offset -> offset + %row_off = index.scale %row, %row_bytes : index, offset -> offset + scf.if %row_ok { + %acc = scf.for %kb = [%c0 to %blocks step %c1](%acc0 = %zero : f32) -> (f32) { + %kbo = index.scale %kb, %o98 : index, offset -> offset + %blk = index.add %row_off, %kbo : offset + %dview = buffer.view %weight_na[%blk] : buffer -> view<1xf16> + %dv = view.load %dview[%c0] : view<1xf16> -> f16 + %d = scalar.extf %dv : f16 to f32 + %acc1 = scf.for %v = [%c0 to %c256 step %c1](%accv = %acc0 : f32) -> (f32) { + // --- decode weight value %v of block %kb --- + %ib32 = index.div %v, %c32 : index + %remv = index.rem %v, %c32 : index + %lv = index.div %remv, %c8 : index + %pv = index.rem %remv, %c8 : index + %halfv = index.div %pv, %c4 : index + %jjv = index.rem %pv, %c4 : index + %qsa = index.mul %ib32, %c8 : index + %qsb = index.mul %lv, %c2 : index + %qsc = index.add %qsa, %qsb : index + %qsiv = index.add %qsc, %halfv : index + %qsov = index.scale %qsiv, %o1 : index, offset -> offset + %qsbase = index.add %blk, %o2 : offset + %qspv = index.add %qsbase, %qsov : offset + %qsview = buffer.view %weight_na[%qspv] : buffer -> view<1xi8> + %gidx_i8 = view.load %qsview[%c0] : view<1xi8> -> i8 + %gidx_s = index.cast %gidx_i8 : i8 to index + %gidx0 = index.andi %gidx_s, %c255 : index + %gidx = index.assume %gidx0 [range(%gidx0, 0, 255)] : index + %ga = index.mul %gidx, %c4 : index + %goi = index.add %ga, %jjv : index + %go = index.scale %goi, %o1 : index, offset -> offset + %gridview = buffer.view %tables_na[%go] : buffer -> view<1xi8> + %gbyte_i8 = view.load %gridview[%c0] : view<1xi8> -> i8 + %gbyte = scalar.uitofp %gbyte_i8 : i8 to f32 + %ssa = index.mul %ib32, %c4 : index + %sso = index.scale %ssa, %o1 : index, offset -> offset + %ssb = index.add %blk, %o66 : offset + %ssp = index.add %ssb, %sso : offset + %auxview = buffer.view %weight_na[%ssp] : buffer -> view<1xi32> + %aux32 = vector.load %auxview[%c0] : view<1xi32> -> vector<1xi32> + %c28v = vector.constant 28 : vector<1xi32> + %snib = vector.shrui %aux32, %c28v : vector<1xi32> + %sf_i32 = vector.extract %snib[0] : vector<1xi32> -> i32 + %sf = scalar.uitofp %sf_i32 : i32 to f32 + %scp = scalar.addf %halfc, %sf : f32 + %dhalf = scalar.mulf %d, %halfc : f32 + %dbv = scalar.mulf %dhalf, %scp : f32 + %sh7 = index.mul %lv, %c7 : index + %sh7_i32 = index.cast %sh7 : index to i32 + %shv = vector.splat %sh7_i32 : vector<1xi32> + %sgnsh = vector.shrui %aux32, %shv : vector<1xi32> + %c127v = vector.constant 127 : vector<1xi32> + %sgnmk = vector.andi %sgnsh, %c127v : vector<1xi32> + %sgnidx_i32 = vector.extract %sgnmk[0] : vector<1xi32> -> i32 + %sgnidx_s = index.cast %sgnidx_i32 : i32 to index + %sgnidx = index.assume %sgnidx_s [range(%sgnidx_s, 0, 127)] : index + %sgno = index.scale %sgnidx, %o1 : index, offset -> offset + %ksabs = index.add %o1024, %sgno : offset + %ksview = buffer.view %tables_na[%ksabs] : buffer -> view<1xi8> + %sgni8 = view.load %ksview[%c0] : view<1xi8> -> i8 + %sgni32 = index.cast %sgni8 : i8 to index + %kmo = index.scale %pv, %o1 : index, offset -> offset + %kmabs = index.add %o1152, %kmo : offset + %kmview = buffer.view %tables_na[%kmabs] : buffer -> view<1xi8> + %kmi8 = view.load %kmview[%c0] : view<1xi8> -> i8 + %kmi32 = index.cast %kmi8 : i8 to index + %andb = index.andi %sgni32, %kmi32 : index + %isneg = index.cmp ne, %andb, %c0 : index + %signf = scf.select %isneg, %negone, %posone : f32 + %dbg = scalar.mulf %dbv, %gbyte : f32 + %wval = scalar.mulf %dbg, %signf : f32 + // --- activation value x[kb*256 + v] --- + %xtok = index.mul %tok, %input_size : index + %xai = index.mul %kb, %c256 : index + %xai1 = index.add %xtok, %xai : index + %xai2 = index.add %xai1, %v : index + %xao = index.scale %xai2, %o4 : index, offset -> offset + %xview = buffer.view %activation_na[%xao] : buffer -> view<1xf32> + %xval = view.load %xview[%c0] : view<1xf32> -> f32 + %prod = scalar.mulf %wval, %xval : f32 + %acc2 = scalar.addf %accv, %prod : f32 + scf.yield %acc2 : f32 + } + scf.yield %acc1 : f32 + } + %otok = index.mul %tok, %output_size : index + %orow = index.add %otok, %row : index + %outo = index.scale %orow, %o4 : index, offset -> offset + %outv = buffer.view %output_na[%outo] : buffer -> view<1xf32> + view.store %acc, %outv[%c0] : f32, view<1xf32> + } + kernel.return +} + +check.case public @ggml_mul_mat_vec_iq3xxs_f32_coverage { + %in = check.literal value(256) : index + %out = check.literal value(64) : index + %tok = check.literal value(1) : index + %activation = check.generate.fill value(0.0) : tensor<256xf32> + %weight = check.generate.fill value(0) : tensor<64x98xi8> + %tables = check.generate.fill value(0) : tensor<1168xi8> + %output = check.generate.fill value(0.0) : tensor<64xf32> + kernel.launch @ggml_mul_mat_vec_iq3xxs_f32[%in, %out, %tok](%in, %out, %tok, %activation, %weight, %tables, %output) : [index, index, index](index, index, index, tensor<256xf32>, tensor<64x98xi8>, tensor<1168xi8>, tensor<64xf32>) + check.expect.close actual(%output) expected(%output) atol(0.0) rtol(0.0) nan(same) : tensor<64xf32> + check.return +} + +check.benchmark<@ggml_mul_mat_vec_iq3xxs_f32_coverage> @ggml_mul_mat_vec_iq3xxs_f32_benchmark diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/ggml/linear_q6k_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/ggml/linear_q6k_f32.loom new file mode 100644 index 000000000000..192b0443ae96 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/ggml/linear_q6k_f32.loom @@ -0,0 +1,423 @@ +// Contracts GGML Q6_K weight rows directly with F32 activation rows using +// the decode schedule selected by llama.cpp's Vulkan backend. +// +// One 64-workitem workgroup computes two adjacent output rows. Four cohorts +// of 16 lanes each process four Q6_K blocks in parallel, while each lane +// decodes four values from the 0, 32, 64, and 96 element quarters of its +// block. Per-group Q6_K scales are exchanged through two 256-byte LDS frames, +// one for each output row. The frames preserve the producer/consumer shape of +// the Vulkan oracle without repacking the persistent GGUF weights. +// +// This provider is intentionally independent of the Q8_1 integer-dot path. +// On gfx11, direct F32 loads and software Q6_K decode can beat activation +// quantization for decode-sized batches. Keeping both algorithms in the same +// corpus lets JIT selection depend on the specialized shape and target. +amdgpu.target @ggml_q6k_gfx11_wave64 {subgroup_size = 64} + +amdgpu.target @ggml_q6k_gfx11_wave32 {subgroup_size = 32} + +config.decl @ggml.linear_q6k_f32.token_capacity : %value: index where [range(%value, 1, 2048)] + +config.decl @ggml.linear_q6k_f32.output_capacity : %value: index where [range(%value, 1, 262144)] + +// Q8_1 reference declarations used only by the differential case. +kernel.decl @ggml_quantize_q8_1_x4_f32(%token_count: index, %input_size: index) launch(%token_count: index, %input_size: index, %input: buffer, %output: buffer) + +kernel.decl @ggml_linear_q6k_q8_1_x4(%token_count: index, %input_size: index, %output_size: index) launch(%token_count: index, %input_size: index, %output_size: index, %q8_input: buffer, %weight: buffer, %output: buffer) + +// Computes an ordered four-element F32 dot product. The scalar recurrence +// mirrors the GLSL oracle and gives target scheduling four independent dot +// chains per decoded Q6_K block. +func.def inline @ggml_q6k_dot4_f32(%lhs: vector<4xf32>, %rhs: vector<4xf32>) -> (f32) { + %c0 = scalar.constant 0.0 : f32 + %lhs0 = vector.extract %lhs[0] : vector<4xf32> -> f32 + %lhs1 = vector.extract %lhs[1] : vector<4xf32> -> f32 + %lhs2 = vector.extract %lhs[2] : vector<4xf32> -> f32 + %lhs3 = vector.extract %lhs[3] : vector<4xf32> -> f32 + %rhs0 = vector.extract %rhs[0] : vector<4xf32> -> f32 + %rhs1 = vector.extract %rhs[1] : vector<4xf32> -> f32 + %rhs2 = vector.extract %rhs[2] : vector<4xf32> -> f32 + %rhs3 = vector.extract %rhs[3] : vector<4xf32> -> f32 + %sum0 = scalar.fmaf %lhs0, %rhs0, %c0 : f32 + %sum1 = scalar.fmaf %lhs1, %rhs1, %sum0 : f32 + %sum2 = scalar.fmaf %lhs2, %rhs2, %sum1 : f32 + %sum3 = scalar.fmaf %lhs3, %rhs3, %sum2 : f32 + func.return %sum3 : f32 +} + +// Loads one row contraction's activation packet. The two output rows read the +// same packet, but each load remains adjacent to its consumer. Hoisting the +// packet across the intervening LDS barrier halves issued activation traffic +// while extending sixteen F32 live values across synchronization; on gfx1151 +// that longer live range is slower than the repeated cache-hot read. +func.def inline @ggml_q6k_load_f32_block(%token_count: index, %input_size: index, %token: index, %block: index, %lane: index, %input: buffer) -> (vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>) { + %c0 = index.constant 0 : index + %c4 = index.constant 4 : index + %c8 = index.constant 8 : index + %c16 = index.constant 16 : index + %c32 = index.constant 32 : index + %c64 = index.constant 64 : index + %c96 = index.constant 96 : index + %c128 = index.constant 128 : index + %c256 = index.constant 256 : index + %c0_offset = index.constant 0 : offset + %bounded_lane = index.assume %lane [range(%lane, 0, 63)] : index + %block_count = index.div %input_size, %c256 : index + %valid_block = index.cmp ult, %block, %block_count : index + %safe_block0 = scf.select %valid_block, %block, %c0 : index + %safe_block, %launch_block_count = index.assume %safe_block0, %block_count [lt(%safe_block0, %block_count)] : index, index + %itid = index.rem %bounded_lane, %c16 : index + %vector_half = index.div %itid, %c8 : index + %vector_index = index.rem %itid, %c8 : index + %input_view = buffer.view %input[%c0_offset] : buffer -> view<[%token_count]x[%launch_block_count]x256xf32> + %vector_half_base = index.mul %vector_half, %c128 : index + %vector_offset = index.mul %vector_index, %c4 : index + %input_index0 = index.add %vector_half_base, %vector_offset : index + %input_index1 = index.add %input_index0, %c32 : index + %input_index2 = index.add %input_index0, %c64 : index + %input_index3 = index.add %input_index0, %c96 : index + %input0 = vector.load %input_view[%token, %safe_block, %input_index0] : view<[%token_count]x[%launch_block_count]x256xf32> -> vector<4xf32> + %input1 = vector.load %input_view[%token, %safe_block, %input_index1] : view<[%token_count]x[%launch_block_count]x256xf32> -> vector<4xf32> + %input2 = vector.load %input_view[%token, %safe_block, %input_index2] : view<[%token_count]x[%launch_block_count]x256xf32> -> vector<4xf32> + %input3 = vector.load %input_view[%token, %safe_block, %input_index3] : view<[%token_count]x[%launch_block_count]x256xf32> -> vector<4xf32> + func.return %input0, %input1, %input2, %input3 : vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32> +} + +// Loads the signed group scale owned by one lane of a Q6_K block. +func.def inline @ggml_q6k_load_f32_scale(%input_size: index, %row: index, %block: index, %lane: index, %weight: buffer) -> (f32) { + %c0 = index.constant 0 : index + %c16 = index.constant 16 : index + %c192 = index.constant 192 : offset + %c210_bytes = index.constant 210 : offset + %c256 = index.constant 256 : index + %bounded_lane = index.assume %lane [range(%lane, 0, 63)] : index + %block_count = index.div %input_size, %c256 : index + %valid_block = index.cmp ult, %block, %block_count : index + %safe_block = scf.if %valid_block -> (index) { + scf.yield %block : index + } else { + scf.yield %c0 : index + } + %itid = index.rem %bounded_lane, %c16 : index + %weight_row_bytes = index.scale %block_count, %c210_bytes : index, offset -> offset + %weight_row_byte_base = index.scale %row, %weight_row_bytes : index, offset -> offset + %block_byte_add = index.scale %safe_block, %c210_bytes : index, offset -> offset + %weight_block_byte_base = index.add %weight_row_byte_base, %block_byte_add : offset + %scale_byte_base = index.add %weight_block_byte_base, %c192 : offset + %scale_view = buffer.view %weight[%scale_byte_base] : buffer -> view<16xi8> + %scale_i8 = view.load %scale_view[%itid] : view<16xi8> -> i8 + %scale = scalar.sitofp %scale_i8 : i8 to f32 + func.return %scale : f32 +} + +// Stages one row's group scales and synchronizes before contraction. +func.def inline @ggml_q6k_stage_f32_scales(%input_size: index, %row: index, %block: index, %frame_count: index, %frame: index, %lane: index, %weight: buffer, %scale_stage: buffer) { + %c0_offset = index.constant 0 : offset + %bounded_lane = index.assume %lane [range(%lane, 0, 63)] : index + %scale_stage_view = buffer.view %scale_stage[%c0_offset] : buffer -> view<[%frame_count]x64xf32> + %scale = func.call @ggml_q6k_load_f32_scale(%input_size, %row, %block, %bounded_lane, %weight) : (index, index, index, index, buffer) -> (f32) + view.store %scale, %scale_stage_view[%frame, %bounded_lane] : f32, view<[%frame_count]x64xf32> + kernel.barrier scope(workgroup) ordering(acq_rel) + func.return +} + +// Computes one lane contribution for one output row and one four-block +// cohort after its group scales have been staged. +func.def inline @ggml_q6k_f32_block_row(%input_size: index, %row: index, %block: index, %frame_count: index, %frame: index, %lane: index, %weight: buffer, %scale_stage: buffer, %input0: vector<4xf32>, %input1: vector<4xf32>, %input2: vector<4xf32>, %input3: vector<4xf32>) -> (f32) { + %c0 = index.constant 0 : index + %c2 = index.constant 2 : index + %c4 = index.constant 4 : index + %c6 = index.constant 6 : index + %c8 = index.constant 8 : index + %c16 = index.constant 16 : index + %c128 = index.constant 128 : offset + %c208 = index.constant 208 : offset + %c210_bytes = index.constant 210 : offset + %c256 = index.constant 256 : index + %c0_offset = index.constant 0 : offset + %c4_i32v = vector.constant 4 : vector<1xi32> + %c2_i32v = vector.constant 2 : vector<1xi32> + %nibble_mask = vector.constant 252645135 : vector<1xi32> + %high0_mask = vector.constant 50529027 : vector<1xi32> + %high2_mask = vector.constant 202116108 : vector<1xi32> + %high4_mask = vector.constant 808464432 : vector<1xi32> + %high6_mask = vector.constant -1061109568 : vector<1xi32> + %c32_f32v = vector.constant 32.0 : vector<4xf32> + %c0_f32 = scalar.constant 0.0 : f32 + %bounded_lane = index.assume %lane [range(%lane, 0, 63)] : index + %block_count = index.div %input_size, %c256 : index + %valid_block = index.cmp ult, %block, %block_count : index + %safe_block = scf.if %valid_block -> (index) { + scf.yield %block : index + } else { + scf.yield %c0 : index + } + %itid = index.rem %bounded_lane, %c16 : index + %cohort = index.div %bounded_lane, %c16 : index + %vector_half = index.div %itid, %c8 : index + %vector_index = index.rem %itid, %c8 : index + %vector_quarter = index.div %vector_index, %c4 : index + %weight_row_bytes = index.scale %block_count, %c210_bytes : index, offset -> offset + %weight_row_byte_base = index.scale %row, %weight_row_bytes : index, offset -> offset + %block_byte_add = index.scale %safe_block, %c210_bytes : index, offset -> offset + %weight_block_byte_base = index.add %weight_row_byte_base, %block_byte_add : offset + %qh_byte_base = index.add %weight_block_byte_base, %c128 : offset + %d_byte_base = index.add %weight_block_byte_base, %c208 : offset + %ql_view = buffer.view %weight[%weight_block_byte_base] : buffer -> view<32xi32> + %qh_view = buffer.view %weight[%qh_byte_base] : buffer -> view<16xi32> + %d_view = buffer.view %weight[%d_byte_base] : buffer -> view<1xf16> + %scale_stage_view = buffer.view %scale_stage[%c0_offset] : buffer -> view<[%frame_count]x64xf32> + %scale_cohort_base = index.mul %cohort, %c16 : index + %scale_half_base = index.mul %vector_half, %c8 : index + %scale_index0 = index.add %scale_half_base, %vector_quarter : index + %scale_index1 = index.add %scale_index0, %c2 : index + %scale_index2 = index.add %scale_index0, %c4 : index + %scale_index3 = index.add %scale_index0, %c6 : index + %stage_scale_index0 = index.add %scale_cohort_base, %scale_index0 : index + %stage_scale_index1 = index.add %scale_cohort_base, %scale_index1 : index + %stage_scale_index2 = index.add %scale_cohort_base, %scale_index2 : index + %stage_scale_index3 = index.add %scale_cohort_base, %scale_index3 : index + %scale0 = view.load %scale_stage_view[%frame, %stage_scale_index0] : view<[%frame_count]x64xf32> -> f32 + %scale1 = view.load %scale_stage_view[%frame, %stage_scale_index1] : view<[%frame_count]x64xf32> -> f32 + %scale2 = view.load %scale_stage_view[%frame, %stage_scale_index2] : view<[%frame_count]x64xf32> -> f32 + %scale3 = view.load %scale_stage_view[%frame, %stage_scale_index3] : view<[%frame_count]x64xf32> -> f32 + %ql_half_word_base = index.mul %vector_half, %c16 : index + %ql_word_index00 = index.add %ql_half_word_base, %vector_index : index + %ql_word_index10 = index.add %ql_word_index00, %c8 : index + %qh_half_word_base = index.mul %vector_half, %c8 : index + %qh_word_index0 = index.add %qh_half_word_base, %vector_index : index + %ql_word_index0, %ql_word_index1, %qh_word_index = index.assume %ql_word_index00, %ql_word_index10, %qh_word_index0 [range(%ql_word_index00, 0, 23), range(%ql_word_index10, 8, 31), range(%qh_word_index0, 0, 15)] : index, index, index + %ql_word0 = vector.load %ql_view[%ql_word_index0] : view<32xi32> -> vector<1xi32> + %ql_word1 = vector.load %ql_view[%ql_word_index1] : view<32xi32> -> vector<1xi32> + %qh_word = vector.load %qh_view[%qh_word_index] : view<16xi32> -> vector<1xi32> + %ql0 = vector.andi %ql_word0, %nibble_mask : vector<1xi32> + %ql1 = vector.andi %ql_word1, %nibble_mask : vector<1xi32> + %ql_word0_high = vector.shrui %ql_word0, %c4_i32v : vector<1xi32> + %ql_word1_high = vector.shrui %ql_word1, %c4_i32v : vector<1xi32> + %ql2 = vector.andi %ql_word0_high, %nibble_mask : vector<1xi32> + %ql3 = vector.andi %ql_word1_high, %nibble_mask : vector<1xi32> + %qh0_low = vector.andi %qh_word, %high0_mask : vector<1xi32> + %qh1_low = vector.andi %qh_word, %high2_mask : vector<1xi32> + %qh2 = vector.andi %qh_word, %high4_mask : vector<1xi32> + %qh3_high = vector.andi %qh_word, %high6_mask : vector<1xi32> + %qh0 = vector.shli %qh0_low, %c4_i32v : vector<1xi32> + %qh1 = vector.shli %qh1_low, %c2_i32v : vector<1xi32> + %qh3 = vector.shrui %qh3_high, %c2_i32v : vector<1xi32> + %code0 = vector.ori %ql0, %qh0 : vector<1xi32> + %code1 = vector.ori %ql1, %qh1 : vector<1xi32> + %code2 = vector.ori %ql2, %qh2 : vector<1xi32> + %code3 = vector.ori %ql3, %qh3 : vector<1xi32> + %code0_i8 = vector.bitcast %code0 : vector<1xi32> to vector<4xi8> + %code1_i8 = vector.bitcast %code1 : vector<1xi32> to vector<4xi8> + %code2_i8 = vector.bitcast %code2 : vector<1xi32> to vector<4xi8> + %code3_i8 = vector.bitcast %code3 : vector<1xi32> to vector<4xi8> + %code0_f32 = vector.uitofp %code0_i8 : vector<4xi8> to vector<4xf32> + %code1_f32 = vector.uitofp %code1_i8 : vector<4xi8> to vector<4xf32> + %code2_f32 = vector.uitofp %code2_i8 : vector<4xi8> to vector<4xf32> + %code3_f32 = vector.uitofp %code3_i8 : vector<4xi8> to vector<4xf32> + %q0 = vector.subf %code0_f32, %c32_f32v : vector<4xf32> + %q1 = vector.subf %code1_f32, %c32_f32v : vector<4xf32> + %q2 = vector.subf %code2_f32, %c32_f32v : vector<4xf32> + %q3 = vector.subf %code3_f32, %c32_f32v : vector<4xf32> + %dot0 = func.call @ggml_q6k_dot4_f32(%input0, %q0) : (vector<4xf32>, vector<4xf32>) -> (f32) + %dot1 = func.call @ggml_q6k_dot4_f32(%input1, %q1) : (vector<4xf32>, vector<4xf32>) -> (f32) + %dot2 = func.call @ggml_q6k_dot4_f32(%input2, %q2) : (vector<4xf32>, vector<4xf32>) -> (f32) + %dot3 = func.call @ggml_q6k_dot4_f32(%input3, %q3) : (vector<4xf32>, vector<4xf32>) -> (f32) + %scaled3 = scalar.mulf %dot3, %scale3 : f32 + %scaled2 = scalar.fmaf %dot2, %scale2, %scaled3 : f32 + %scaled1 = scalar.fmaf %dot1, %scale1, %scaled2 : f32 + %scaled0 = scalar.fmaf %dot0, %scale0, %scaled1 : f32 + %d_f16 = view.load %d_view[0] : view<1xf16> -> f16 + %d = scalar.extf %d_f16 : f16 to f32 + %contribution0 = scalar.mulf %scaled0, %d : f32 + %contribution = scf.if %valid_block -> (f32) { + scf.yield %contribution0 : f32 + } else { + scf.yield %c0_f32 : f32 + } + func.return %contribution : f32 +} + +// Shared direct-F32 matrix-vector schedule. Target entry points below vary +// subgroup width without duplicating the Q6_K decoding or memory schedule. +func.def inline @ggml_linear_q6k_f32_body(%publish_output: i1, %token_count: index, %token0: index, %pair: index, %input_size: index, %output_size: index, %input: buffer, %weight: buffer, %output: buffer) { + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 2048)] : index + %bounded_input_size = index.assume %input_size [range(%input_size, 256, 32768), mul(%input_size, 256)] : index + %bounded_output_size = index.assume %output_size [range(%output_size, 1, 262144)] : index + %lane0 = kernel.workitem.id : index + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c4 = index.constant 4 : index + %c16 = index.constant 16 : index + %c256 = index.constant 256 : index + %c0_f32 = scalar.constant 0.0 : f32 + %c0_offset = index.constant 0 : offset + %scale_stage_bytes = index.constant 512 : offset + %token, %launch_token_count = index.assume %token0, %bounded_token_count [lt(%token0, %bounded_token_count)] : index, index + %lane = index.assume %lane0 [range(%lane0, 0, 63)] : index + %row00 = index.mul %pair, %c2 : index + %row0, %launch_output_size = index.assume %row00, %bounded_output_size [lt(%row00, %bounded_output_size)] : index, index + %row1 = index.add %row0, %c1 : index + %row1_valid = index.cmp ult, %row1, %launch_output_size : index + %block_count = index.div %bounded_input_size, %c256 : index + %cohort = index.div %lane, %c16 : index + scf.if %publish_output { + %input_noalias, %weight_noalias, %output_noalias = buffer.assume.noalias %input, %weight, %output : buffer, buffer, buffer + %scale_stage = buffer.alloca align(16) %scale_stage_bytes : buffer + %acc0, %acc1 = scf.for %block_base = [%c0 to %block_count step %c4](%row_acc0 = %c0_f32 : f32, %row_acc1 = %c0_f32 : f32) -> (f32, f32) { + %block = index.add %block_base, %cohort : index + func.call @ggml_q6k_stage_f32_scales(%bounded_input_size, %row0, %block, %c2, %c0, %lane, %weight_noalias, %scale_stage) : (index, index, index, index, index, index, buffer, buffer) + %row0_input0, %row0_input1, %row0_input2, %row0_input3 = func.call @ggml_q6k_load_f32_block(%launch_token_count, %bounded_input_size, %token, %block, %lane, %input_noalias) : (index, index, index, index, index, buffer) -> (vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>) + %contribution0 = func.call @ggml_q6k_f32_block_row(%bounded_input_size, %row0, %block, %c2, %c0, %lane, %weight_noalias, %scale_stage, %row0_input0, %row0_input1, %row0_input2, %row0_input3) : (index, index, index, index, index, index, buffer, buffer, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>) -> (f32) + %contribution1 = scf.if %row1_valid -> (f32) { + func.call @ggml_q6k_stage_f32_scales(%bounded_input_size, %row1, %block, %c2, %c1, %lane, %weight_noalias, %scale_stage) : (index, index, index, index, index, index, buffer, buffer) + %row1_input0, %row1_input1, %row1_input2, %row1_input3 = func.call @ggml_q6k_load_f32_block(%launch_token_count, %bounded_input_size, %token, %block, %lane, %input_noalias) : (index, index, index, index, index, buffer) -> (vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>) + %row1_contribution = func.call @ggml_q6k_f32_block_row(%bounded_input_size, %row1, %block, %c2, %c1, %lane, %weight_noalias, %scale_stage, %row1_input0, %row1_input1, %row1_input2, %row1_input3) : (index, index, index, index, index, index, buffer, buffer, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>) -> (f32) + scf.yield %row1_contribution : f32 + } else { + scf.yield %c0_f32 : f32 + } + %next0 = scalar.addf %row_acc0, %contribution0 : f32 + %next1 = scalar.addf %row_acc1, %contribution1 : f32 + scf.yield %next0, %next1 : f32, f32 + } + %sum0 = kernel.workgroup.reduce %acc0 : f32 + %sum1 = kernel.workgroup.reduce %acc1 : f32 + %is_lane_zero = index.cmp eq, %lane, %c0 : index + %output_view = buffer.view %output_noalias[%c0_offset] : buffer -> view<[%launch_token_count]x[%launch_output_size]xf32> + scf.if %is_lane_zero { + view.store %sum0, %output_view[%token, %row0] : f32, view<[%launch_token_count]x[%launch_output_size]xf32> + } + scf.if %row1_valid { + scf.if %is_lane_zero { + view.store %sum1, %output_view[%token, %row1] : f32, view<[%launch_token_count]x[%launch_output_size]xf32> + } + } + } + func.return +} + +// Wave64 provider matching the Vulkan oracle's subgroup schedule. +kernel.def target(@ggml_q6k_gfx11_wave64) @ggml_linear_q6k_f32_wave64(%token_count: index, %input_size: index, %output_size: index) { + %token_capacity = config.get @ggml.linear_q6k_f32.token_capacity : index + %output_capacity = config.get @ggml.linear_q6k_f32.output_capacity : index + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c64 = index.constant 64 : index + %padded_output_size = index.add %output_capacity, %c1 : index + %output_pairs = index.div %padded_output_size, %c2 : index + kernel.launch.config workgroups(%output_pairs, %token_capacity, %c1) workgroup_size(%c64, %c1, %c1) : index +} launch(%token_count: index, %input_size: index, %output_size: index, %input: buffer, %weight: buffer, %output: buffer) { + %token_capacity = config.get @ggml.linear_q6k_f32.token_capacity : index + %output_capacity = config.get @ggml.linear_q6k_f32.output_capacity : index + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 2048), le(%token_count, %token_capacity)] : index + %bounded_output_size = index.assume %output_size [range(%output_size, 1, 262144), le(%output_size, %output_capacity)] : index + %c2 = index.constant 2 : index + %pair = kernel.workgroup.id : index + %token = kernel.workgroup.id : index + %row0 = index.mul %pair, %c2 : index + %valid_token = index.cmp ult, %token, %bounded_token_count : index + %valid_row = index.cmp ult, %row0, %bounded_output_size : index + %publish_output = scalar.andi %valid_token, %valid_row : i1 + %c0 = index.constant 0 : index + %safe_token = scf.select %valid_token, %token, %c0 : index + %safe_pair = scf.select %valid_row, %pair, %c0 : index + func.call @ggml_linear_q6k_f32_body(%publish_output, %bounded_token_count, %safe_token, %safe_pair, %input_size, %bounded_output_size, %input, %weight, %output) : (i1, index, index, index, index, index, buffer, buffer, buffer) + kernel.return +} + +// Wave32 provider testing whether two subgroups hide direct-F32 decode latency +// better on targets where wave64 is not the default execution mode. +kernel.def target(@ggml_q6k_gfx11_wave32) @ggml_linear_q6k_f32_wave32(%token_count: index, %input_size: index, %output_size: index) { + %token_capacity = config.get @ggml.linear_q6k_f32.token_capacity : index + %output_capacity = config.get @ggml.linear_q6k_f32.output_capacity : index + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c64 = index.constant 64 : index + %padded_output_size = index.add %output_capacity, %c1 : index + %output_pairs = index.div %padded_output_size, %c2 : index + kernel.launch.config workgroups(%output_pairs, %token_capacity, %c1) workgroup_size(%c64, %c1, %c1) : index +} launch(%token_count: index, %input_size: index, %output_size: index, %input: buffer, %weight: buffer, %output: buffer) { + %token_capacity = config.get @ggml.linear_q6k_f32.token_capacity : index + %output_capacity = config.get @ggml.linear_q6k_f32.output_capacity : index + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 2048), le(%token_count, %token_capacity)] : index + %bounded_output_size = index.assume %output_size [range(%output_size, 1, 262144), le(%output_size, %output_capacity)] : index + %c2 = index.constant 2 : index + %pair = kernel.workgroup.id : index + %token = kernel.workgroup.id : index + %row0 = index.mul %pair, %c2 : index + %valid_token = index.cmp ult, %token, %bounded_token_count : index + %valid_row = index.cmp ult, %row0, %bounded_output_size : index + %publish_output = scalar.andi %valid_token, %valid_row : i1 + %c0 = index.constant 0 : index + %safe_token = scf.select %valid_token, %token, %c0 : index + %safe_pair = scf.select %valid_row, %pair, %c0 : index + func.call @ggml_linear_q6k_f32_body(%publish_output, %bounded_token_count, %safe_token, %safe_pair, %input_size, %bounded_output_size, %input, %weight, %output) : (i1, index, index, index, index, index, buffer, buffer, buffer) + kernel.return +} + +// Exact-representable activations make the direct-F32 and Q8_1 contraction +// paths comparable while nonzero packed bytes exercise every Q6_K field. +check.case public @ggml_linear_q6k_f32_differential_case { + %token_count = check.literal value(2) : index + %input_size = check.literal value(2048) : index + %output_size = check.literal value(9) : index + %input = check.generate.fill value(0.00390625) : tensor<2x2048xf32> + %q8_input = check.generate.fill value(0) : tensor<2x2304xi8> + %weight = check.generate.fill value(-86) : tensor<9x8x210xi8> + %expected = check.generate.fill value(0.0) : tensor<2x9xf32> + %actual_wave64 = check.generate.fill value(1.0) : tensor<2x9xf32> + %actual_wave32 = check.generate.fill value(1.0) : tensor<2x9xf32> + kernel.launch @ggml_quantize_q8_1_x4_f32[%token_count, %input_size](%token_count, %input_size, %input, %q8_input) : [index, index](index, index, tensor<2x2048xf32>, tensor<2x2304xi8>) + kernel.launch @ggml_linear_q6k_q8_1_x4[%token_count, %input_size, %output_size](%token_count, %input_size, %output_size, %q8_input, %weight, %expected) : [index, index, index](index, index, index, tensor<2x2304xi8>, tensor<9x8x210xi8>, tensor<2x9xf32>) + kernel.launch @ggml_linear_q6k_f32_wave64[%token_count, %input_size, %output_size](%token_count, %input_size, %output_size, %input, %weight, %actual_wave64) : [index, index, index](index, index, index, tensor<2x2048xf32>, tensor<9x8x210xi8>, tensor<2x9xf32>) + kernel.launch @ggml_linear_q6k_f32_wave32[%token_count, %input_size, %output_size](%token_count, %input_size, %output_size, %input, %weight, %actual_wave32) : [index, index, index](index, index, index, tensor<2x2048xf32>, tensor<9x8x210xi8>, tensor<2x9xf32>) + check.expect.close actual(%actual_wave64) expected(%expected) atol(0.25) rtol(0.01) nan(same) : tensor<2x9xf32> + check.expect.close actual(%actual_wave32) expected(%expected) atol(0.25) rtol(0.01) nan(same) : tensor<2x9xf32> + check.return +} + +check.case public @ggml_linear_q6k_f32_wave64_dense_v_benchmark_case { + %token_count = check.param.choice values([1, 8, 32, 128, 512]) name("token_count") : index + %input_size = check.literal value(2048) : index + %output_size = check.literal value(512) : index + %input = check.generate.fill value(0.0) : tensor<[%token_count]x2048xf32> + %weight = check.generate.fill value(0) : tensor<512x8x210xi8> + %output = check.generate.fill value(1.0) : tensor<[%token_count]x512xf32> + %expected = check.generate.fill value(0.0) : tensor<[%token_count]x512xf32> + kernel.launch @ggml_linear_q6k_f32_wave64[%token_count, %input_size, %output_size](%token_count, %input_size, %output_size, %input, %weight, %output) : [index, index, index](index, index, index, tensor<[%token_count]x2048xf32>, tensor<512x8x210xi8>, tensor<[%token_count]x512xf32>) + check.expect.close actual(%output) expected(%expected) atol(0.0) rtol(0.0) nan(same) : tensor<[%token_count]x512xf32> + check.return +} + +check.case public @ggml_linear_q6k_f32_wave32_dense_v_benchmark_case { + %token_count = check.param.choice values([1, 8, 32, 128, 512]) name("token_count") : index + %input_size = check.literal value(2048) : index + %output_size = check.literal value(512) : index + %input = check.generate.fill value(0.0) : tensor<[%token_count]x2048xf32> + %weight = check.generate.fill value(0) : tensor<512x8x210xi8> + %output = check.generate.fill value(1.0) : tensor<[%token_count]x512xf32> + %expected = check.generate.fill value(0.0) : tensor<[%token_count]x512xf32> + kernel.launch @ggml_linear_q6k_f32_wave32[%token_count, %input_size, %output_size](%token_count, %input_size, %output_size, %input, %weight, %output) : [index, index, index](index, index, index, tensor<[%token_count]x2048xf32>, tensor<512x8x210xi8>, tensor<[%token_count]x512xf32>) + check.expect.close actual(%output) expected(%expected) atol(0.0) rtol(0.0) nan(same) : tensor<[%token_count]x512xf32> + check.return +} + +check.benchmark<@ggml_linear_q6k_f32_differential_case> @ggml_linear_q6k_f32_differential + +check.benchmark<@ggml_linear_q6k_f32_wave64_dense_v_benchmark_case> @ggml_linear_q6k_f32_wave64_dense_v_decode {token_count = 1} + +check.benchmark<@ggml_linear_q6k_f32_wave64_dense_v_benchmark_case> @ggml_linear_q6k_f32_wave64_dense_v_prefill_32 {token_count = 32} + +check.benchmark<@ggml_linear_q6k_f32_wave64_dense_v_benchmark_case> @ggml_linear_q6k_f32_wave64_dense_v_prefill_128 {token_count = 128} + +check.benchmark<@ggml_linear_q6k_f32_wave64_dense_v_benchmark_case> @ggml_linear_q6k_f32_wave64_dense_v_prefill_512 {token_count = 512} + +check.benchmark<@ggml_linear_q6k_f32_wave32_dense_v_benchmark_case> @ggml_linear_q6k_f32_wave32_dense_v_decode {token_count = 1} + +check.benchmark<@ggml_linear_q6k_f32_wave32_dense_v_benchmark_case> @ggml_linear_q6k_f32_wave32_dense_v_prefill_32 {token_count = 32} + +check.benchmark<@ggml_linear_q6k_f32_wave32_dense_v_benchmark_case> @ggml_linear_q6k_f32_wave32_dense_v_prefill_128 {token_count = 128} + +check.benchmark<@ggml_linear_q6k_f32_wave32_dense_v_benchmark_case> @ggml_linear_q6k_f32_wave32_dense_v_prefill_512 {token_count = 512} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/ggml/linear_q6k_q8_1_x4.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/ggml/linear_q6k_q8_1_x4.loom new file mode 100644 index 000000000000..508150a1b308 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/ggml/linear_q6k_q8_1_x4.loom @@ -0,0 +1,533 @@ +// Contracts GGML Q6_K weight rows directly with the Q8_1 x4 activation +// layout used by llama.cpp's Vulkan matrix path. One 210-byte Q6_K block +// represents 256 signed six-bit values: +// +// i8 ql[128]; // low four bits +// i8 qh[64]; // high two bits +// i8 scales[16]; // signed scale for each 16-value group +// f16 d; // block-wide scale +// +// A wave owns one output value. Each lane contracts eight values from every +// Q6_K block as two native signed dot4 packets, and the wave reduces those +// partials. The kernel is the raw-layout correctness and decode baseline shared +// by dense and routed projections; prefill schedules can reuse the inline +// packed-row primitive while staging weights across multiple activation rows. +// The shared packer is linked beside this module. +func.def inline @ggml_q8_1_x4_word(%q8_input: buffer, %row_byte_base: offset, %q8_block: index, %word_in_block: index) -> (vector<4xi8>, f32) { + %c2 = index.constant 2 : index + %c4 = index.constant 4 : index + %c8 = index.constant 8 : index + %group_bytes = index.constant 144 : offset + %payload_byte_add = index.constant 16 : offset + %group = index.div %q8_block, %c4 : index + %inner0 = index.rem %q8_block, %c4 : index + %inner = index.assume %inner0 [range(%inner0, 0, 3)] : index + %word0 = index.assume %word_in_block [range(%word_in_block, 0, 7)] : index + %group_byte_add = index.scale %group, %group_bytes : index, offset -> offset + %group_byte_base = index.add %row_byte_base, %group_byte_add : offset + %payload_byte_base = index.add %group_byte_base, %payload_byte_add : offset + %ds_view = buffer.view %q8_input[%group_byte_base] : buffer -> view<8xf16> + %payload_view = buffer.view %q8_input[%payload_byte_base] : buffer -> view<32xi32> + %d_index = index.mul %inner, %c2 : index + %inner_word_base = index.mul %inner, %c8 : index + %word_index = index.add %inner_word_base, %word0 : index + %d_f16 = view.load %ds_view[%d_index] : view<8xf16> -> f16 + %packed = vector.load %payload_view[%word_index] : view<32xi32> -> vector<1xi32> + %values = vector.bitcast %packed : vector<1xi32> to vector<4xi8> + %d = scalar.extf %d_f16 : f16 to f32 + func.return %values, %d : vector<4xi8>, f32 +} + +amdgpu.target @ggml_q6k_q8_gfx1151_wave64 {subgroup_size = 64} + +config.decl @ggml.linear_q6k_q8_1_x4.token_capacity : %value: index where [range(%value, 1, 2048)] + +config.decl @ggml.linear_q6k_q8_1_x4.output_capacity : %value: index where [range(%value, 1, 262144)] + +kernel.decl @ggml_quantize_q8_1_x4_f32(%token_count: index, %input_size: index) launch(%token_count: index, %input_size: index, %input: buffer, %output: buffer) + +func.decl @ggml_q8_1_x4_word(%q8_input: buffer, %row_byte_base: offset, %q8_block: index, %word_in_block: index) -> (vector<4xi8>, f32) + +// Sign-extends four packed six-bit values without scalar lane extraction. +// Each byte enters with bits [5:0] populated and leaves as signed i8. +func.def inline @ggml_q6k_sign_extend_dot4(%code: vector<1xi32>) -> (vector<4xi8>) { + %c1_i32v = vector.constant 1 : vector<1xi32> + %c2_i32v = vector.constant 2 : vector<1xi32> + %low5_mask = vector.constant 522133279 : vector<1xi32> + %bit5_mask = vector.constant 538976288 : vector<1xi32> + %sign_mask = vector.constant -522133280 : vector<1xi32> + %low5 = vector.andi %code, %low5_mask : vector<1xi32> + %bit5 = vector.andi %code, %bit5_mask : vector<1xi32> + %bit6 = vector.shli %bit5, %c1_i32v : vector<1xi32> + %bit7 = vector.shli %bit5, %c2_i32v : vector<1xi32> + %high01 = vector.ori %bit5, %bit6 : vector<1xi32> + %high = vector.ori %high01, %bit7 : vector<1xi32> + %sign = vector.xori %high, %sign_mask : vector<1xi32> + %signed_i32 = vector.ori %low5, %sign : vector<1xi32> + %signed = vector.bitcast %signed_i32 : vector<1xi32> to vector<4xi8> + func.return %signed : vector<4xi8> +} + +// Decodes four adjacent values from one 32-value Q6_K group for FP16 matrix +// staging. The group and packet coordinates match the contiguous K dimension +// consumed by WMMA tiles, unlike the lane/part mapping used by the Q8_1 dot +// contraction below. +func.def inline @ggml_q6k_f16_vector4(%weight: buffer, %weight_row_byte_base: offset, %q6_block: index, %q6_group: index, %packet: index) -> (vector<4xf16>) { + %c2 = index.constant 2 : index + %c4 = index.constant 4 : index + %c8 = index.constant 8 : index + %c16 = index.constant 16 : index + %block_bytes = index.constant 210 : offset + %qh_byte_add = index.constant 128 : offset + %scale_byte_add = index.constant 192 : offset + %d_byte_add = index.constant 208 : offset + %c4_i32v = vector.constant 4 : vector<1xi32> + %nibble_mask = vector.constant 252645135 : vector<1xi32> + %high_mask = vector.constant 50529027 : vector<1xi32> + %c32_f32v = vector.constant 32.0 : vector<4xf32> + %bounded_group = index.assume %q6_group [range(%q6_group, 0, 7)] : index + %bounded_packet = index.assume %packet [range(%packet, 0, 7)] : index + %block_byte_add = index.scale %q6_block, %block_bytes : index, offset -> offset + %block_byte_base = index.add %weight_row_byte_base, %block_byte_add : offset + %qh_byte_base = index.add %block_byte_base, %qh_byte_add : offset + %scale_byte_base = index.add %block_byte_base, %scale_byte_add : offset + %d_byte_base = index.add %block_byte_base, %d_byte_add : offset + %ql_view = buffer.view %weight[%block_byte_base] : buffer -> view<32xi32> + %qh_view = buffer.view %weight[%qh_byte_base] : buffer -> view<16xi32> + %scale_view = buffer.view %weight[%scale_byte_base] : buffer -> view<16xi8> + %d_view = buffer.view %weight[%d_byte_base] : buffer -> view<1xf16> + %group_in_half = index.rem %bounded_group, %c4 : index + %half = index.div %bounded_group, %c4 : index + %ql_side = index.rem %group_in_half, %c2 : index + %ql_half_word_base = index.mul %half, %c16 : index + %ql_side_word_add = index.mul %ql_side, %c8 : index + %ql_word_base = index.add %ql_half_word_base, %ql_side_word_add : index + %ql_word_index = index.add %ql_word_base, %bounded_packet : index + %qh_half_word_base = index.mul %half, %c8 : index + %qh_word_index = index.add %qh_half_word_base, %bounded_packet : index + %nibble = index.div %group_in_half, %c2 : index + %nibble_shift_index = index.mul %nibble, %c4 : index + %nibble_shift_i32 = index.cast %nibble_shift_index : index to i32 + %nibble_shift = vector.splat %nibble_shift_i32 : vector<1xi32> + %qh_shift_index = index.mul %group_in_half, %c2 : index + %qh_shift_i32 = index.cast %qh_shift_index : index to i32 + %qh_shift = vector.splat %qh_shift_i32 : vector<1xi32> + %scale_packet_half = index.div %bounded_packet, %c4 : index + %scale_group_base = index.mul %bounded_group, %c2 : index + %scale_index = index.add %scale_group_base, %scale_packet_half : index + %ql_word = vector.load %ql_view[%ql_word_index] : view<32xi32> -> vector<1xi32> + %qh_word = vector.load %qh_view[%qh_word_index] : view<16xi32> -> vector<1xi32> + %ql_shifted = vector.shrui %ql_word, %nibble_shift : vector<1xi32> + %ql = vector.andi %ql_shifted, %nibble_mask : vector<1xi32> + %qh_shifted = vector.shrui %qh_word, %qh_shift : vector<1xi32> + %qh_low = vector.andi %qh_shifted, %high_mask : vector<1xi32> + %qh = vector.shli %qh_low, %c4_i32v : vector<1xi32> + %code = vector.ori %ql, %qh : vector<1xi32> + %code_i8 = vector.bitcast %code : vector<1xi32> to vector<4xi8> + %code_f32 = vector.uitofp %code_i8 : vector<4xi8> to vector<4xf32> + %centered = vector.subf %code_f32, %c32_f32v : vector<4xf32> + %scale_i8 = view.load %scale_view[%scale_index] : view<16xi8> -> i8 + %d_f16 = view.load %d_view[0] : view<1xf16> -> f16 + %scale = scalar.sitofp %scale_i8 : i8 to f32 + %d = scalar.extf %d_f16 : f16 to f32 + %combined_scale = scalar.mulf %scale, %d : f32 + %combined_scale_vector = vector.splat %combined_scale : vector<4xf32> + %values_f32 = vector.mulf %centered, %combined_scale_vector : vector<4xf32> + %values = vector.fptrunc %values_f32 : vector<4xf32> to vector<4xf16> + func.return %values : vector<4xf16> +} + +// Contracts both four-value packets assigned to one lane in a Q6_K block. +// QL, QH, and the block scale are shared by the two packed parts, and their +// contributions are summed before entering the row-wide recurrence. +func.def inline @ggml_q6k_q8_1_x4_block_lane(%weight: buffer, %weight_row_byte_base: offset, %q8_input: buffer, %q8_row_byte_base: offset, %q6_block: index, %lane: index) -> (f32) { + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c4 = index.constant 4 : index + %c8 = index.constant 8 : index + %c16 = index.constant 16 : index + %block_bytes = index.constant 210 : offset + %qh_byte_add = index.constant 128 : offset + %scale_byte_add = index.constant 192 : offset + %d_byte_add = index.constant 208 : offset + %c4_i32v = vector.constant 4 : vector<1xi32> + %c0_i32v = vector.constant 0 : vector<1xi32> + %nibble_mask = vector.constant 252645135 : vector<1xi32> + %high_mask = vector.constant 808464432 : vector<1xi32> + %c0_f32 = scalar.constant 0.0 : f32 + %bounded_lane = index.assume %lane [range(%lane, 0, 31)] : index + %block_byte_add = index.scale %q6_block, %block_bytes : index, offset -> offset + %block_byte_base = index.add %weight_row_byte_base, %block_byte_add : offset + %qh_byte_base = index.add %block_byte_base, %qh_byte_add : offset + %scale_byte_base = index.add %block_byte_base, %scale_byte_add : offset + %d_byte_base = index.add %block_byte_base, %d_byte_add : offset + %ql_view = buffer.view %weight[%block_byte_base] : buffer -> view<32xi32> + %qh_view = buffer.view %weight[%qh_byte_base] : buffer -> view<16xi32> + %scale_view = buffer.view %weight[%scale_byte_base] : buffer -> view<16xi8> + %d_view = buffer.view %weight[%d_byte_base] : buffer -> view<1xf16> + %lane_mod8 = index.rem %bounded_lane, %c8 : index + %lane_mod16 = index.rem %bounded_lane, %c16 : index + %lane_div16 = index.div %bounded_lane, %c16 : index + %lane_div8_in_16 = index.div %lane_mod16, %c8 : index + %lane_div4_in_16 = index.div %lane_mod16, %c4 : index + %qh_high_base = index.mul %lane_div16, %c8 : index + %qh_index0 = index.add %qh_high_base, %lane_mod8 : index + %qh_index = index.assume %qh_index0 [range(%qh_index0, 0, 15)] : index + %ql_word = vector.load %ql_view[%bounded_lane] : view<32xi32> -> vector<1xi32> + %qh_word = vector.load %qh_view[%qh_index] : view<16xi32> -> vector<1xi32> + %qh_base_shift_index = index.mul %lane_div8_in_16, %c2 : index + %qh_base_shift_i32 = index.cast %qh_base_shift_index : index to i32 + %q8_block_base = index.mul %q6_block, %c8 : index + %q8_high_add = index.mul %lane_div16, %c4 : index + %q8_quadrant = index.add %q8_high_add, %lane_div8_in_16 : index + %scale_high_base = index.mul %lane_div16, %c8 : index + %scale_lane0 = index.add %scale_high_base, %lane_div4_in_16 : index + %d_f16 = view.load %d_view[0] : view<1xf16> -> f16 + %d = scalar.extf %d_f16 : f16 to f32 + %sum = scf.for %part = [%c0 to %c2 step %c1](%accumulator = %c0_f32 : f32) -> (f32) unroll { + %bounded_part = index.assume %part [range(%part, 0, 1)] : index + %part_shift_index = index.mul %bounded_part, %c4 : index + %part_shift_i32 = index.cast %part_shift_index : index to i32 + %part_shift = vector.splat %part_shift_i32 : vector<1xi32> + %ql_shifted = vector.shrui %ql_word, %part_shift : vector<1xi32> + %ql = vector.andi %ql_shifted, %nibble_mask : vector<1xi32> + %qh_shift_i32 = scalar.addi %qh_base_shift_i32, %part_shift_i32 : i32 + %qh_shift = vector.splat %qh_shift_i32 : vector<1xi32> + %qh_shifted = vector.shrui %qh_word, %qh_shift : vector<1xi32> + %qh_positioned = vector.shli %qh_shifted, %c4_i32v : vector<1xi32> + %qh = vector.andi %qh_positioned, %high_mask : vector<1xi32> + %code = vector.ori %ql, %qh : vector<1xi32> + %signed_weight = func.call @ggml_q6k_sign_extend_dot4(%code) : (vector<1xi32>) -> (vector<4xi8>) + %q8_part_add = index.mul %bounded_part, %c2 : index + %q8_block_part = index.add %q8_block_base, %q8_part_add : index + %q8_block = index.add %q8_block_part, %q8_quadrant : index + %q8_values, %q8_d = func.call @ggml_q8_1_x4_word(%q8_input, %q8_row_byte_base, %q8_block, %lane_mod8) : (buffer, offset, index, index) -> (vector<4xi8>, f32) + %scale_lane = index.add %scale_lane0, %part_shift_index : index + %scale_i8 = view.load %scale_view[%scale_lane] : view<16xi8> -> i8 + %scale = scalar.sitofp %scale_i8 : i8 to f32 + %dot = vector.dot4i %signed_weight, %q8_values, %c0_i32v : vector<4xi8>, vector<4xi8>, vector<1xi32> + %dot_i32 = vector.extract %dot[0] : vector<1xi32> -> i32 + %dot_f32 = scalar.sitofp %dot_i32 : i32 to f32 + %scaled0 = scalar.mulf %dot_f32, %scale : f32 + %scaled1 = scalar.mulf %scaled0, %d : f32 + %contribution = scalar.mulf %scaled1, %q8_d : f32 + %next = scalar.addf %accumulator, %contribution : f32 + scf.yield %next : f32 + } + func.return %sum : f32 +} + +// Computes one lane's partial for a complete Q6_K row. Callers choose how +// rows and activations are tiled, then reduce the returned value by subgroup. +func.def inline @ggml_q6k_q8_1_x4_row_lane(%input_size: index, %weight: buffer, %weight_row_byte_base: offset, %q8_input: buffer, %q8_row_byte_base: offset, %lane: index) -> (f32) { + %bounded_input_size = index.assume %input_size [range(%input_size, 256, 32768), mul(%input_size, 256)] : index + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c256 = index.constant 256 : index + %c0_f32 = scalar.constant 0.0 : f32 + %block_count = index.div %bounded_input_size, %c256 : index + %result = scf.for %block = [%c0 to %block_count step %c1](%block_acc = %c0_f32 : f32) -> (f32) { + %contribution = func.call @ggml_q6k_q8_1_x4_block_lane(%weight, %weight_row_byte_base, %q8_input, %q8_row_byte_base, %block, %lane) : (buffer, offset, buffer, offset, index, index) -> (f32) + %next = scalar.addf %block_acc, %contribution : f32 + scf.yield %next : f32 + } + func.return %result : f32 +} + +// Maps four Q6_K blocks across the four 16-lane partitions of a wave64. Each +// physical lane contracts the pair of virtual wave32 lanes that own the same +// packed positions in the low and high halves of one block. The subgroup +// reduction therefore combines four complete blocks per loop iteration while +// preserving the canonical wave32 unpacking and dot-product primitive. +func.def inline @ggml_q6k_q8_1_x4_row_lane_wave64_block4(%input_size: index, %weight: buffer, %weight_row_byte_base: offset, %q8_input: buffer, %q8_row_byte_base: offset, %lane: index) -> (f32) { + %bounded_input_size = index.assume %input_size [range(%input_size, 256, 32768), mul(%input_size, 256)] : index + %bounded_lane = index.assume %lane [range(%lane, 0, 63)] : index + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c3 = index.constant 3 : index + %c4 = index.constant 4 : index + %c16 = index.constant 16 : index + %c256 = index.constant 256 : index + %c0_f32 = scalar.constant 0.0 : f32 + %block_count = index.div %bounded_input_size, %c256 : index + %padded_block_count = index.add %block_count, %c3 : index + %iteration_count = index.div %padded_block_count, %c4 : index + %block_in_iteration = index.div %bounded_lane, %c16 : index + %virtual_lane0 = index.rem %bounded_lane, %c16 : index + %virtual_lane1 = index.add %virtual_lane0, %c16 : index + %result = scf.for %iteration = [%c0 to %iteration_count step %c1](%row_acc = %c0_f32 : f32) -> (f32) { + %block_base = index.mul %iteration, %c4 : index + %block0 = index.add %block_base, %block_in_iteration : index + %valid_block = index.cmp ult, %block0, %block_count : index + %block_acc = scf.if %valid_block -> (f32) { + %low = func.call @ggml_q6k_q8_1_x4_block_lane(%weight, %weight_row_byte_base, %q8_input, %q8_row_byte_base, %block0, %virtual_lane0) : (buffer, offset, buffer, offset, index, index) -> (f32) + %high = func.call @ggml_q6k_q8_1_x4_block_lane(%weight, %weight_row_byte_base, %q8_input, %q8_row_byte_base, %block0, %virtual_lane1) : (buffer, offset, buffer, offset, index, index) -> (f32) + %sum = scalar.addf %low, %high : f32 + scf.yield %sum : f32 + } else { + scf.yield %c0_f32 : f32 + } + %next = scalar.addf %row_acc, %block_acc : f32 + scf.yield %next : f32 + } + func.return %result : f32 +} + +// Contracts one output channel for the calling Q6_K projection kernel. The +// caller owns publication so the same canonical contraction can feed either a +// dense logits tensor or an endpoint reduction. +func.def inline @ggml_linear_q6k_q8_1_x4_body(%publish_output: i1, %token_count: index, %token0: index, %input_size: index, %output_size: index, %q8_input: buffer, %weight: buffer) -> (f32, index, index, i1, index, index) { + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 2048)] : index + %bounded_input_size = index.assume %input_size [range(%input_size, 256, 32768), mul(%input_size, 256)] : index + %bounded_output_size = index.assume %output_size [range(%output_size, 1, 262144)] : index + %channel_tile = kernel.workgroup.id : index + %subgroup0 = kernel.subgroup.id : index + %lane = kernel.subgroup.lane.id : index + %c0 = index.constant 0 : index + %c8 = index.constant 8 : index + %c128 = index.constant 128 : index + %c210_bytes = index.constant 210 : offset + %c256 = index.constant 256 : index + %c144_bytes = index.constant 144 : offset + %token, %launch_token_count = index.assume %token0, %bounded_token_count [lt(%token0, %bounded_token_count)] : index, index + %subgroup = index.assume %subgroup0 [range(%subgroup0, 0, 7)] : index + %channel_base = index.mul %channel_tile, %c8 : index + %channel = index.add %channel_base, %subgroup : index + %valid_channel = index.cmp ult, %channel, %bounded_output_size : index + %q6_block_count = index.div %bounded_input_size, %c256 : index + %weight_row_bytes = index.scale %q6_block_count, %c210_bytes : index, offset -> offset + %weight_row_byte_base = index.scale %channel, %weight_row_bytes : index, offset -> offset + %q8_group_count = index.div %bounded_input_size, %c128 : index + %q8_row_bytes = index.scale %q8_group_count, %c144_bytes : index, offset -> offset + %q8_row_byte_base = index.scale %token, %q8_row_bytes : index, offset -> offset + %lane_acc = scf.if %publish_output -> (f32) { + %channel_acc = scf.if %valid_channel -> (f32) { + %value = func.call @ggml_q6k_q8_1_x4_row_lane(%bounded_input_size, %weight, %weight_row_byte_base, %q8_input, %q8_row_byte_base, %lane) : (index, buffer, offset, buffer, offset, index) -> (f32) + scf.yield %value : f32 + } else { + %c0_f32 = scalar.constant 0.0 : f32 + scf.yield %c0_f32 : f32 + } + scf.yield %channel_acc : f32 + } else { + %c0_f32 = scalar.constant 0.0 : f32 + scf.yield %c0_f32 : f32 + } + %dot = kernel.subgroup.reduce %lane_acc : f32 + func.return %dot, %token, %channel, %valid_channel, %launch_token_count, %bounded_output_size : f32, index, index, i1, index, index +} + +// Dense raw-layout baseline. A 256-thread workgroup carries eight independent +// wave32 output rows so decode launches enough waves without duplicating Q6_K +// decoding within a row. +kernel.def export("ggml_linear_q6k_q8_1_x4") @ggml_linear_q6k_q8_1_x4(%token_count: index, %input_size: index, %output_size: index) { + %token_capacity = config.get @ggml.linear_q6k_q8_1_x4.token_capacity : index + %output_capacity = config.get @ggml.linear_q6k_q8_1_x4.output_capacity : index + %c8 = index.constant 8 : index + %c7 = index.constant 7 : index + %c1 = index.constant 1 : index + %workgroup_size = index.constant 256 : index + %padded_output_size = index.add %output_capacity, %c7 : index + %output_tiles = index.div %padded_output_size, %c8 : index + kernel.launch.config workgroups(%output_tiles, %token_capacity, %c1) workgroup_size(%workgroup_size, %c1, %c1) : index +} launch(%token_count: index, %input_size: index, %output_size: index, %q8_input: buffer, %weight: buffer, %output: buffer) { + %q8_noalias, %weight_noalias, %output_noalias = buffer.assume.noalias %q8_input, %weight, %output : buffer, buffer, buffer + %token_capacity = config.get @ggml.linear_q6k_q8_1_x4.token_capacity : index + %output_capacity = config.get @ggml.linear_q6k_q8_1_x4.output_capacity : index + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 2048), le(%token_count, %token_capacity)] : index + %bounded_output_size = index.assume %output_size [range(%output_size, 1, 262144), le(%output_size, %output_capacity)] : index + %c0 = index.constant 0 : index + %token0 = kernel.workgroup.id : index + %valid_token = index.cmp ult, %token0, %bounded_token_count : index + %safe_token = scf.select %valid_token, %token0, %c0 : index + %dot, %token, %channel, %valid_channel, %launch_token_count, %output_bound = func.call @ggml_linear_q6k_q8_1_x4_body(%valid_token, %bounded_token_count, %safe_token, %input_size, %bounded_output_size, %q8_noalias, %weight_noalias) : (i1, index, index, index, index, buffer, buffer) -> (f32, index, index, i1, index, index) + %c0_i32 = scalar.constant 0 : i32 + %c0_offset = index.constant 0 : offset + %lane = kernel.subgroup.lane.id : index + %lane_i32 = index.cast %lane : index to i32 + %is_lane_zero = scalar.cmpi eq, %lane_i32, %c0_i32 : i32 + scf.if %valid_token { + scf.if %valid_channel { + scf.if %is_lane_zero { + %output_view = buffer.view %output_noalias[%c0_offset] : buffer -> view<[%launch_token_count]x[%output_bound]xf32> + view.store %dot, %output_view[%token, %channel] : f32, view<[%launch_token_count]x[%output_bound]xf32> + } + } + } + kernel.return +} + +// Wave64 vocabulary schedule matching llama.cpp Vulkan's one-row workgroup and +// four-block concurrency on AMD targets with native wave64 execution. +kernel.def target(@ggml_q6k_q8_gfx1151_wave64) export("ggml_linear_q6k_q8_1_x4") @ggml_linear_q6k_q8_1_x4_wave64_block4(%token_count: index, %input_size: index, %output_size: index) { + %token_capacity = config.get @ggml.linear_q6k_q8_1_x4.token_capacity : index + %output_capacity = config.get @ggml.linear_q6k_q8_1_x4.output_capacity : index + %c1 = index.constant 1 : index + %subgroup_size = target.subgroup.size : index + kernel.launch.config workgroups(%output_capacity, %token_capacity, %c1) workgroup_size(%subgroup_size, %c1, %c1) : index +} launch(%token_count: index, %input_size: index, %output_size: index, %q8_input: buffer, %weight: buffer, %output: buffer) { + %q8_noalias, %weight_noalias, %output_noalias = buffer.assume.noalias %q8_input, %weight, %output : buffer, buffer, buffer + %token_capacity = config.get @ggml.linear_q6k_q8_1_x4.token_capacity : index + %output_capacity = config.get @ggml.linear_q6k_q8_1_x4.output_capacity : index + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 2048), le(%token_count, %token_capacity)] : index + %bounded_input_size = index.assume %input_size [range(%input_size, 256, 32768), mul(%input_size, 256)] : index + %bounded_output_size = index.assume %output_size [range(%output_size, 1, 262144), le(%output_size, %output_capacity)] : index + %token = kernel.workgroup.id : index + %channel = kernel.workgroup.id : index + %lane = kernel.subgroup.lane.id : index + %valid_token = index.cmp ult, %token, %bounded_token_count : index + %valid_channel = index.cmp ult, %channel, %bounded_output_size : index + %publish_output = scalar.andi %valid_token, %valid_channel : i1 + %c0 = index.constant 0 : index + %c128 = index.constant 128 : index + %c144_bytes = index.constant 144 : offset + %c210_bytes = index.constant 210 : offset + %c256 = index.constant 256 : index + %safe_token = scf.select %valid_token, %token, %c0 : index + %safe_channel = scf.select %valid_channel, %channel, %c0 : index + %q6_block_count = index.div %bounded_input_size, %c256 : index + %weight_row_bytes = index.scale %q6_block_count, %c210_bytes : index, offset -> offset + %weight_row_byte_base = index.scale %safe_channel, %weight_row_bytes : index, offset -> offset + %q8_group_count = index.div %bounded_input_size, %c128 : index + %q8_row_bytes = index.scale %q8_group_count, %c144_bytes : index, offset -> offset + %q8_row_byte_base = index.scale %safe_token, %q8_row_bytes : index, offset -> offset + %lane_acc = scf.if %publish_output -> (f32) { + %value = func.call @ggml_q6k_q8_1_x4_row_lane_wave64_block4(%bounded_input_size, %weight_noalias, %weight_row_byte_base, %q8_noalias, %q8_row_byte_base, %lane) : (index, buffer, offset, buffer, offset, index) -> (f32) + scf.yield %value : f32 + } else { + %c0_f32 = scalar.constant 0.0 : f32 + scf.yield %c0_f32 : f32 + } + %dot = kernel.subgroup.reduce %lane_acc : f32 + %is_lane_zero = index.cmp eq, %lane, %c0 : index + scf.if %publish_output { + scf.if %is_lane_zero { + %c0_offset = index.constant 0 : offset + %output_view = buffer.view %output_noalias[%c0_offset] : buffer -> view<[%bounded_token_count]x[%bounded_output_size]xf32> + view.store %dot, %output_view[%token, %channel] : f32, view<[%bounded_token_count]x[%bounded_output_size]xf32> + } + } + kernel.return +} + +// Uniform packed bytes exercise both low/high Q6 fields and signed per-group +// scales. The output width crosses the eight-wave tile boundary, while K=2048 +// covers the exact mixed-format dense V contraction depth. +check.case public @ggml_linear_q6k_q8_1_x4_nonzero_tail_case { + %token_count = check.literal value(2) : index + %input_size = check.literal value(2048) : index + %output_size = check.literal value(9) : index + %input = check.generate.fill value(0.00390625) : tensor<2x2048xf32> + %q8_input = check.generate.fill value(0) : tensor<2x2304xi8> + %weight = check.generate.fill value(-86) : tensor<9x8x210xi8> + %output = check.generate.fill value(0.0) : tensor<2x9xf32> + %expected = check.generate.fill value(358.1715087890625) : tensor<2x9xf32> + kernel.launch @ggml_quantize_q8_1_x4_f32[%token_count, %input_size](%token_count, %input_size, %input, %q8_input) : [index, index](index, index, tensor<2x2048xf32>, tensor<2x2304xi8>) + kernel.launch @ggml_linear_q6k_q8_1_x4[%token_count, %input_size, %output_size](%token_count, %input_size, %output_size, %q8_input, %weight, %output) : [index, index, index](index, index, index, tensor<2x2304xi8>, tensor<9x8x210xi8>, tensor<2x9xf32>) + check.expect.close actual(%output) expected(%expected) atol(0.25) rtol(9.9999999999999995e-07) nan(same) : tensor<2x9xf32> + check.return +} + +// The gfx1151 schedule has an exact target requirement, so its nonzero packed +// data case remains independently selectable from the generic gfx11 case. +check.case public @ggml_linear_q6k_q8_1_x4_wave64_block4_nonzero_tail_case { + %token_count = check.literal value(2) : index + %input_size = check.literal value(2048) : index + %output_size = check.literal value(9) : index + %input = check.generate.fill value(0.00390625) : tensor<2x2048xf32> + %q8_input = check.generate.fill value(0) : tensor<2x2304xi8> + %weight = check.generate.fill value(-86) : tensor<9x8x210xi8> + %output = check.generate.fill value(0.0) : tensor<2x9xf32> + %expected = check.generate.fill value(358.1715087890625) : tensor<2x9xf32> + kernel.launch @ggml_quantize_q8_1_x4_f32[%token_count, %input_size](%token_count, %input_size, %input, %q8_input) : [index, index](index, index, tensor<2x2048xf32>, tensor<2x2304xi8>) + kernel.launch @ggml_linear_q6k_q8_1_x4_wave64_block4[%token_count, %input_size, %output_size](%token_count, %input_size, %output_size, %q8_input, %weight, %output) : [index, index, index](index, index, index, tensor<2x2304xi8>, tensor<9x8x210xi8>, tensor<2x9xf32>) + check.expect.close actual(%output) expected(%expected) atol(0.25) rtol(9.9999999999999995e-07) nan(same) : tensor<2x9xf32> + check.return +} + +check.case public @ggml_linear_q6k_q8_1_x4_benchmark_case { + %token_count = check.param.choice values([1, 17, 32, 63, 128, 129, 512]) name("token_count") : index + %input_size = check.literal value(768) : index + %output_size = check.literal value(2048) : index + %q8_input = check.generate.fill value(0) : tensor<[%token_count]x864xi8> + %weight = check.generate.fill value(0) : tensor<2048x3x210xi8> + %output = check.generate.fill value(1.0) : tensor<[%token_count]x2048xf32> + %expected = check.generate.fill value(0.0) : tensor<[%token_count]x2048xf32> + kernel.launch @ggml_linear_q6k_q8_1_x4[%token_count, %input_size, %output_size](%token_count, %input_size, %output_size, %q8_input, %weight, %output) : [index, index, index](index, index, index, tensor<[%token_count]x864xi8>, tensor<2048x3x210xi8>, tensor<[%token_count]x2048xf32>) + check.expect.close actual(%output) expected(%expected) atol(0.0) rtol(0.0) nan(same) : tensor<[%token_count]x2048xf32> + check.return +} + +// Production vocabulary shape used to compare the eight-row wave32 baseline +// with the one-row wave64 block schedule without model-runtime overhead. +check.case public @ggml_linear_q6k_q8_1_x4_vocabulary_wave32_benchmark_case { + %token_count = check.literal value(1) : index + %input_size = check.literal value(2048) : index + %output_size = check.literal value(151936) : index + %q8_input = check.generate.fill value(1) : tensor<1x2304xi8> + %weight = check.generate.fill value(0) : tensor<151936x8x210xi8> + %output = check.generate.fill value(1.0) : tensor<1x151936xf32> + %expected = check.generate.fill value(0.0) : tensor<1x151936xf32> + kernel.launch @ggml_linear_q6k_q8_1_x4[%token_count, %input_size, %output_size](%token_count, %input_size, %output_size, %q8_input, %weight, %output) : [index, index, index](index, index, index, tensor<1x2304xi8>, tensor<151936x8x210xi8>, tensor<1x151936xf32>) + check.expect.close actual(%output) expected(%expected) atol(0.0) rtol(0.0) nan(same) : tensor<1x151936xf32> + check.return +} + +check.case public @ggml_linear_q6k_q8_1_x4_vocabulary_wave64_block4_benchmark_case { + %token_count = check.literal value(1) : index + %input_size = check.literal value(2048) : index + %output_size = check.literal value(151936) : index + %q8_input = check.generate.fill value(1) : tensor<1x2304xi8> + %weight = check.generate.fill value(0) : tensor<151936x8x210xi8> + %output = check.generate.fill value(1.0) : tensor<1x151936xf32> + %expected = check.generate.fill value(0.0) : tensor<1x151936xf32> + kernel.launch @ggml_linear_q6k_q8_1_x4_wave64_block4[%token_count, %input_size, %output_size](%token_count, %input_size, %output_size, %q8_input, %weight, %output) : [index, index, index](index, index, index, tensor<1x2304xi8>, tensor<151936x8x210xi8>, tensor<1x151936xf32>) + check.expect.close actual(%output) expected(%expected) atol(0.0) rtol(0.0) nan(same) : tensor<1x151936xf32> + check.return +} + +// Exercises the complete activation-pack and dense Q6_K projection boundary at +// the K=2048, M=512 shape used by mixed-format attention V weights. The token +// buckets cover decode, awkward prefill tails, llama.cpp's 512-token +// microbatch, and larger schedules available to callers without host routing. +check.case public @ggml_linear_q6k_q8_1_x4_dense_v_benchmark_case { + %token_count = check.param.choice values([1, 17, 32, 63, 128, 129, 512, 1024, 2048]) name("token_count") : index + %input_size = check.literal value(2048) : index + %output_size = check.literal value(512) : index + %input = check.generate.fill value(0.0) : tensor<[%token_count]x2048xf32> + %q8_input = check.generate.fill value(1) : tensor<[%token_count]x2304xi8> + %weight = check.generate.fill value(0) : tensor<512x8x210xi8> + %output = check.generate.fill value(1.0) : tensor<[%token_count]x512xf32> + %expected = check.generate.fill value(0.0) : tensor<[%token_count]x512xf32> + kernel.launch @ggml_quantize_q8_1_x4_f32[%token_count, %input_size](%token_count, %input_size, %input, %q8_input) : [index, index](index, index, tensor<[%token_count]x2048xf32>, tensor<[%token_count]x2304xi8>) + kernel.launch @ggml_linear_q6k_q8_1_x4[%token_count, %input_size, %output_size](%token_count, %input_size, %output_size, %q8_input, %weight, %output) : [index, index, index](index, index, index, tensor<[%token_count]x2304xi8>, tensor<512x8x210xi8>, tensor<[%token_count]x512xf32>) + check.expect.close actual(%output) expected(%expected) atol(0.0) rtol(0.0) nan(same) : tensor<[%token_count]x512xf32> + check.return +} + +check.benchmark<@ggml_linear_q6k_q8_1_x4_nonzero_tail_case> @ggml_linear_q6k_q8_1_x4_small + +check.benchmark<@ggml_linear_q6k_q8_1_x4_benchmark_case> @ggml_linear_q6k_q8_1_x4_decode {token_count = 1} + +check.benchmark<@ggml_linear_q6k_q8_1_x4_vocabulary_wave32_benchmark_case> @ggml_linear_q6k_q8_1_x4_vocabulary_wave32 + +check.benchmark<@ggml_linear_q6k_q8_1_x4_vocabulary_wave64_block4_benchmark_case> @ggml_linear_q6k_q8_1_x4_vocabulary_wave64_block4 + +check.benchmark<@ggml_linear_q6k_q8_1_x4_benchmark_case> @ggml_linear_q6k_q8_1_x4_prefill_32 {token_count = 32} + +check.benchmark<@ggml_linear_q6k_q8_1_x4_benchmark_case> @ggml_linear_q6k_q8_1_x4_prefill_128 {token_count = 128} + +check.benchmark<@ggml_linear_q6k_q8_1_x4_benchmark_case> @ggml_linear_q6k_q8_1_x4_prefill_512 {token_count = 512} + +check.benchmark<@ggml_linear_q6k_q8_1_x4_dense_v_benchmark_case> @ggml_linear_q6k_q8_1_x4_dense_v_decode {token_count = 1} + +check.benchmark<@ggml_linear_q6k_q8_1_x4_dense_v_benchmark_case> @ggml_linear_q6k_q8_1_x4_dense_v_prefill_32 {token_count = 32} + +check.benchmark<@ggml_linear_q6k_q8_1_x4_dense_v_benchmark_case> @ggml_linear_q6k_q8_1_x4_dense_v_prefill_128 {token_count = 128} + +check.benchmark<@ggml_linear_q6k_q8_1_x4_dense_v_benchmark_case> @ggml_linear_q6k_q8_1_x4_dense_v_prefill_512 {token_count = 512} + +check.benchmark<@ggml_linear_q6k_q8_1_x4_dense_v_benchmark_case> @ggml_linear_q6k_q8_1_x4_dense_v_prefill_1024 {token_count = 1024} + +check.benchmark<@ggml_linear_q6k_q8_1_x4_dense_v_benchmark_case> @ggml_linear_q6k_q8_1_x4_dense_v_prefill_2048 {token_count = 2048} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/ggml/quantize_q8_1_x4.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/ggml/quantize_q8_1_x4.loom new file mode 100644 index 000000000000..8ffd25089dfe --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/ggml/quantize_q8_1_x4.loom @@ -0,0 +1,274 @@ +// Defines GGML's block_q8_1_x4 physical-layout accessors and packs contiguous +// F32 activations into that layout. Four logical 32-element Q8_1 blocks share +// one 144-byte physical group: +// +// struct block_q8_1_x4 { +// f16 ds[4][2]; // per-block (scale, quantized_sum * scale) +// i32 qs[4][8]; // four signed i8 values per packed word +// }; +config.decl @ggml.quantize_q8_1_x4.group_capacity : %value: index where [range(%value, 1, 524288)] + +// Loads one logical 32-element Q8_1 block from the four-way physical packing. +// `row_byte_base` addresses the first x4 group for one activation row. +func.def inline @ggml_q8_1_x4_block(%q8_input: buffer, %row_byte_base: offset, %q8_block: index) -> (vector<32xi8>, f32, f32) { + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c4 = index.constant 4 : index + %c8 = index.constant 8 : index + %group_bytes = index.constant 144 : offset + %payload_byte_add = index.constant 16 : offset + %group = index.div %q8_block, %c4 : index + %inner0 = index.rem %q8_block, %c4 : index + %inner = index.assume %inner0 [range(%inner0, 0, 3)] : index + %group_byte_add = index.scale %group, %group_bytes : index, offset -> offset + %group_byte_base = index.add %row_byte_base, %group_byte_add : offset + %payload_byte_base = index.add %group_byte_base, %payload_byte_add : offset + %ds_view = buffer.view %q8_input[%group_byte_base] : buffer -> view<8xf16> + %payload_view = buffer.view %q8_input[%payload_byte_base] : buffer -> view<32xi32> + %d_index = index.mul %inner, %c2 : index + %s_index = index.add %d_index, %c1 : index + %word_index = index.mul %inner, %c8 : index + %d_f16 = view.load %ds_view[%d_index] : view<8xf16> -> f16 + %s_f16 = view.load %ds_view[%s_index] : view<8xf16> -> f16 + %words = vector.load %payload_view[%word_index] : view<32xi32> -> vector<8xi32> + %values = vector.bitcast %words : vector<8xi32> to vector<32xi8> + %d = scalar.extf %d_f16 : f16 to f32 + %s = scalar.extf %s_f16 : f16 to f32 + func.return %values, %d, %s : vector<32xi8>, f32, f32 +} + +// Loads one packed four-value word and its block scale. Dot-product kernels +// use this narrower form so a lane does not load the other seven words owned +// by its subgroup peers. +func.def inline @ggml_q8_1_x4_word(%q8_input: buffer, %row_byte_base: offset, %q8_block: index, %word_in_block: index) -> (vector<4xi8>, f32) { + %c2 = index.constant 2 : index + %c4 = index.constant 4 : index + %c8 = index.constant 8 : index + %group_bytes = index.constant 144 : offset + %payload_byte_add = index.constant 16 : offset + %group = index.div %q8_block, %c4 : index + %inner0 = index.rem %q8_block, %c4 : index + %inner = index.assume %inner0 [range(%inner0, 0, 3)] : index + %word0 = index.assume %word_in_block [range(%word_in_block, 0, 7)] : index + %group_byte_add = index.scale %group, %group_bytes : index, offset -> offset + %group_byte_base = index.add %row_byte_base, %group_byte_add : offset + %payload_byte_base = index.add %group_byte_base, %payload_byte_add : offset + %ds_view = buffer.view %q8_input[%group_byte_base] : buffer -> view<8xf16> + %payload_view = buffer.view %q8_input[%payload_byte_base] : buffer -> view<32xi32> + %d_index = index.mul %inner, %c2 : index + %inner_word_base = index.mul %inner, %c8 : index + %word_index = index.add %inner_word_base, %word0 : index + %d_f16 = view.load %ds_view[%d_index] : view<8xf16> -> f16 + %packed = vector.load %payload_view[%word_index] : view<32xi32> -> vector<1xi32> + %values = vector.bitcast %packed : vector<1xi32> to vector<4xi8> + %d = scalar.extf %d_f16 : f16 to f32 + func.return %values, %d : vector<4xi8>, f32 +} + +// Packs one explicit physical group. Callers provide the complete physical +// group domain and an in-range ordinal, and guarantee that one complete +// 32-lane wave executes after all 128 source values are visible. The uniform +// publication predicate suppresses only destination writes so every lane still +// reaches the workgroup barriers. +func.def inline @ggml_quantize_q8_1_x4_group_body(%publish_output: i1, %group_count0: index, %group0: index, %input: buffer, %output: buffer) { + %group_count, %group = index.assume %group_count0, %group0 [range(%group_count0, 1, 524288), lt(%group0, %group_count0)] : index, index + %lane = kernel.workitem.id : index + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c4 = index.constant 4 : index + %c8 = index.constant 8 : index + %c32 = index.constant 32 : index + %c128 = index.constant 128 : index + %group_bytes = index.constant 144 : offset + %payload_byte_add = index.constant 16 : offset + %scratch_d_byte_add = index.constant 128 : offset + %scratch_bytes = index.constant 144 : offset + %c0_f32 = scalar.constant 0.0 : f32 + %c1_f32 = scalar.constant 1.0 : f32 + %c127 = scalar.constant 127.0 : f32 + %c0_offset = index.constant 0 : offset + %launched_element_count = index.mul %group_count, %c128 : index + %block_in_group0 = index.div %lane, %c8 : index + %block_in_group = index.assume %block_in_group0 [range(%block_in_group0, 0, 3)] : index + %word_in_block0 = index.rem %lane, %c8 : index + %word_in_block = index.assume %word_in_block0 [range(%word_in_block0, 0, 7)] : index + %group_element_base = index.mul %group, %c128 : index + %block_element_add = index.mul %block_in_group, %c32 : index + %block_word_add = index.mul %block_in_group, %c8 : index + %word_element_add = index.mul %word_in_block, %c4 : index + %input_block_base = index.add %group_element_base, %block_element_add : index + %input_index = index.add %input_block_base, %word_element_add : index + %input_noalias, %output_noalias = buffer.assume.noalias %input, %output : buffer, buffer + %input_view = buffer.view %input_noalias[%c0_offset] : buffer -> view<[%launched_element_count]xf32> + %input_values = vector.load %input_view[%input_index] : view<[%launched_element_count]xf32> -> vector<4xf32> + %absolute_values = vector.absf %input_values : vector<4xf32> + %thread_max = vector.reduce %absolute_values, %c0_f32 : vector<4xf32>, f32 + %scratch = buffer.alloca align(16) %scratch_bytes : buffer + %scratch_values = buffer.view %scratch[%c0_offset] : buffer -> view<32xf32> + %scratch_d = buffer.view %scratch[%scratch_d_byte_add] : buffer -> view<4xf32> + view.store %thread_max, %scratch_values[%lane] : f32, view<32xf32> + kernel.barrier scope(workgroup) ordering(acq_rel) + %is_cohort_leader = index.cmp eq, %word_in_block, %c0 : index + scf.if %is_cohort_leader { + %cohort_base = index.mul %block_in_group, %c8 : index + %cohort_maxima = vector.load %scratch_values[%cohort_base] : view<32xf32> -> vector<8xf32> + %amax = vector.reduce %cohort_maxima, %c0_f32 : vector<8xf32>, f32 + %d = scalar.divf %amax, %c127 : f32 + view.store %d, %scratch_d[%block_in_group] : f32, view<4xf32> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + %d = view.load %scratch_d[%block_in_group] : view<4xf32> -> f32 + %d_nonzero = scalar.cmpf one, %d, %c0_f32 : f32 + %d_inverse = scf.if %d_nonzero -> (f32) { + %inverse = scalar.divf %c1_f32, %d : f32 + scf.yield %inverse : f32 + } else { + scf.yield %c0_f32 : f32 + } + %d_inverse_vector = vector.splat %d_inverse : vector<4xf32> + %scaled_values = vector.mulf %input_values, %d_inverse_vector : vector<4xf32> + %rounded_values = vector.roundf %scaled_values : vector<4xf32> + %quantized_values = vector.fptosi %rounded_values : vector<4xf32> to vector<4xi8> + %packed_word = vector.bitcast %quantized_values : vector<4xi8> to vector<1xi32> + %group_byte_offset = index.scale %group, %group_bytes : index, offset -> offset + %payload_byte_offset = index.add %group_byte_offset, %payload_byte_add : offset + %group_ds = buffer.view %output_noalias[%group_byte_offset] : buffer -> view<8xf16> + %group_qs = buffer.view %output_noalias[%payload_byte_offset] : buffer -> view<32xi32> + %packed_word_index0 = index.add %block_word_add, %word_in_block : index + %packed_word_index = index.assume %packed_word_index0 [range(%packed_word_index0, 0, 31)] : index + scf.if %publish_output { + vector.store %packed_word, %group_qs[%packed_word_index] : vector<1xi32>, view<32xi32> + } + %thread_sum = vector.reduce %rounded_values, %c0_f32 : vector<4xf32>, f32 + view.store %thread_sum, %scratch_values[%lane] : f32, view<32xf32> + kernel.barrier scope(workgroup) ordering(acq_rel) + %publishes_metadata = scalar.andi %is_cohort_leader, %publish_output : i1 + scf.if %publishes_metadata { + %cohort_base = index.mul %block_in_group, %c8 : index + %cohort_sums = vector.load %scratch_values[%cohort_base] : view<32xf32> -> vector<8xf32> + %quantized_sum = vector.reduce %cohort_sums, %c0_f32 : vector<8xf32>, f32 + %s = scalar.mulf %quantized_sum, %d : f32 + %d_f16 = scalar.fptrunc %d : f32 to f16 + %s_f16 = scalar.fptrunc %s : f32 to f16 + %ds_index = index.mul %block_in_group, %c2 : index + view.store %d_f16, %group_ds[%ds_index] : f16, view<8xf16> + %s_index = index.add %ds_index, %c1 : index + view.store %s_f16, %group_ds[%s_index] : f16, view<8xf16> + } + func.return +} + +// A 32-lane workgroup owns one physical group. Each eight-lane cohort loads +// four values, computes one Q8_1 block, and writes disjoint metadata and packed +// words. The LDS reductions make the independent eight-lane cohorts explicit. +kernel.def @ggml_quantize_q8_1_x4_f32(%token_count: index, %input_size: index) { + %group_capacity = config.get @ggml.quantize_q8_1_x4.group_capacity : index + %unit = index.constant 1 : index + %workgroup_size = index.constant 32 : index + kernel.launch.config workgroups(%group_capacity, %unit, %unit) workgroup_size(%workgroup_size, %unit, %unit) : index +} launch(%token_count: index, %input_size: index, %input: buffer, %output: buffer) { + %group_capacity = config.get @ggml.quantize_q8_1_x4.group_capacity : index + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 2048)] : index + %bounded_input_size = index.assume %input_size [range(%input_size, 128, 32768), mul(%input_size, 128)] : index + %elements_per_group = index.constant 128 : index + %c0 = index.constant 0 : index + %element_count = index.mul %bounded_token_count, %bounded_input_size : index + %group_count0 = index.div %element_count, %elements_per_group : index + %group_count = index.assume %group_count0 [range(%group_count0, 1, 524288), le(%group_count0, %group_capacity)] : index + %group0 = kernel.workgroup.id : index + %valid_group = index.cmp ult, %group0, %group_count : index + %safe_group = scf.select %valid_group, %group0, %c0 : index + %group = index.assume %safe_group [lt(%safe_group, %group_count)] : index + func.call @ggml_quantize_q8_1_x4_group_body(%valid_group, %group_count, %group, %input, %output) : (i1, index, index, buffer, buffer) + kernel.return +} + +// This check-only inspector keeps the production packer ABI exact while +// exposing packed words and heterogeneous metadata as comparable tensors. +kernel.def @ggml_q8_1_x4_inspect_one_group() { + %unit = index.constant 1 : index + %workgroup_size = index.constant 32 : index + kernel.launch.config workgroups(%unit, %unit, %unit) workgroup_size(%workgroup_size, %unit, %unit) : index +} launch(%packed: buffer, %words: buffer, %d_values: buffer, %s_values: buffer) { + %lane = kernel.workitem.id : index + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c4 = index.constant 4 : index + %payload_offset = index.constant 16 : offset + %c0_offset = index.constant 0 : offset + %packed_ds = buffer.view %packed[%c0_offset] : buffer -> view<8xf16> + %packed_words = buffer.view %packed[%payload_offset] : buffer -> view<32xi32> + %word_output = buffer.view %words[%c0_offset] : buffer -> view<32xi32> + %d_output = buffer.view %d_values[%c0_offset] : buffer -> view<4xf32> + %s_output = buffer.view %s_values[%c0_offset] : buffer -> view<4xf32> + %word = view.load %packed_words[%lane] : view<32xi32> -> i32 + view.store %word, %word_output[%lane] : i32, view<32xi32> + %is_metadata_lane = index.cmp ult, %lane, %c4 : index + scf.if %is_metadata_lane { + %ds_index = index.mul %lane, %c2 : index + %s_index = index.add %ds_index, %c1 : index + %d_f16 = view.load %packed_ds[%ds_index] : view<8xf16> -> f16 + %s_f16 = view.load %packed_ds[%s_index] : view<8xf16> -> f16 + %d = scalar.extf %d_f16 : f16 to f32 + %s = scalar.extf %s_f16 : f16 to f32 + view.store %d, %d_output[%lane] : f32, view<4xf32> + view.store %s, %s_output[%lane] : f32, view<4xf32> + } + kernel.return +} + +check.case public @ggml_quantize_q8_1_x4_f32_nonzero_case { + %token_count = check.literal value(1) : index + %input_size = check.literal value(128) : index + %positive_input = check.generate.fill value(1.0) : tensor<128xf32> + %negative_input = check.generate.fill value(-1.0) : tensor<128xf32> + %positive_packed = check.generate.fill value(0) : tensor<144xi8> + %negative_packed = check.generate.fill value(0) : tensor<144xi8> + %positive_words = check.generate.fill value(0) : tensor<32xi32> + %negative_words = check.generate.fill value(0) : tensor<32xi32> + %positive_d = check.generate.fill value(0.0) : tensor<4xf32> + %negative_d = check.generate.fill value(0.0) : tensor<4xf32> + %positive_s = check.generate.fill value(0.0) : tensor<4xf32> + %negative_s = check.generate.fill value(0.0) : tensor<4xf32> + %expected_positive_words = check.generate.fill value(2139062143) : tensor<32xi32> + %expected_negative_words = check.generate.fill value(-2122219135) : tensor<32xi32> + %expected_d = check.generate.fill value(0.00787353515625) : tensor<4xf32> + %expected_positive_s = check.generate.fill value(32.0) : tensor<4xf32> + %expected_negative_s = check.generate.fill value(-32.0) : tensor<4xf32> + kernel.launch @ggml_quantize_q8_1_x4_f32[%token_count, %input_size](%token_count, %input_size, %positive_input, %positive_packed) : [index, index](index, index, tensor<128xf32>, tensor<144xi8>) + kernel.launch @ggml_q8_1_x4_inspect_one_group(%positive_packed, %positive_words, %positive_d, %positive_s) : (tensor<144xi8>, tensor<32xi32>, tensor<4xf32>, tensor<4xf32>) + kernel.launch @ggml_quantize_q8_1_x4_f32[%token_count, %input_size](%token_count, %input_size, %negative_input, %negative_packed) : [index, index](index, index, tensor<128xf32>, tensor<144xi8>) + kernel.launch @ggml_q8_1_x4_inspect_one_group(%negative_packed, %negative_words, %negative_d, %negative_s) : (tensor<144xi8>, tensor<32xi32>, tensor<4xf32>, tensor<4xf32>) + check.expect.equal actual(%positive_words) expected(%expected_positive_words) : tensor<32xi32> + check.expect.equal actual(%negative_words) expected(%expected_negative_words) : tensor<32xi32> + check.expect.close actual(%positive_d) expected(%expected_d) atol(0.0) rtol(0.0) nan(same) : tensor<4xf32> + check.expect.close actual(%negative_d) expected(%expected_d) atol(0.0) rtol(0.0) nan(same) : tensor<4xf32> + check.expect.close actual(%positive_s) expected(%expected_positive_s) atol(0.0) rtol(0.0) nan(same) : tensor<4xf32> + check.expect.close actual(%negative_s) expected(%expected_negative_s) atol(0.0) rtol(0.0) nan(same) : tensor<4xf32> + check.return +} + +check.case public @ggml_quantize_q8_1_x4_f32_benchmark_case { + %token_count = check.param.choice values([1, 4, 8, 17, 32, 63, 128, 129, 512]) name("token_count") : index + %input_size = check.literal value(2048) : index + %input = check.generate.fill value(0.0) : tensor<[%token_count]x2048xf32> + %packed = check.generate.fill value(1) : tensor<[%token_count]x2304xi8> + %expected = check.generate.fill value(0) : tensor<[%token_count]x2304xi8> + kernel.launch @ggml_quantize_q8_1_x4_f32[%token_count, %input_size](%token_count, %input_size, %input, %packed) : [index, index](index, index, tensor<[%token_count]x2048xf32>, tensor<[%token_count]x2304xi8>) + check.expect.equal actual(%packed) expected(%expected) : tensor<[%token_count]x2304xi8> + check.return +} + +check.benchmark<@ggml_quantize_q8_1_x4_f32_nonzero_case> @ggml_quantize_q8_1_x4_f32_small + +check.benchmark<@ggml_quantize_q8_1_x4_f32_benchmark_case> @ggml_quantize_q8_1_x4_f32_decode {token_count = 1} + +check.benchmark<@ggml_quantize_q8_1_x4_f32_benchmark_case> @ggml_quantize_q8_1_x4_f32_small_batch_8 {token_count = 8} + +check.benchmark<@ggml_quantize_q8_1_x4_f32_benchmark_case> @ggml_quantize_q8_1_x4_f32_prefill_32 {token_count = 32} + +check.benchmark<@ggml_quantize_q8_1_x4_f32_benchmark_case> @ggml_quantize_q8_1_x4_f32_prefill_128 {token_count = 128} + +check.benchmark<@ggml_quantize_q8_1_x4_f32_benchmark_case> @ggml_quantize_q8_1_x4_f32_prefill_512 {token_count = 512} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/manifest.json b/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/manifest.json new file mode 100644 index 000000000000..fcbc5c57f4fe --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/manifest.json @@ -0,0 +1,5386 @@ +{ + "schema": "ggml-hrx-qwen-kernel-corpus-v2", + "upstream_repository": "https://github.com/rocm/hrx", + "upstream_revision": "d343e4ad5167dca6758cb0d2db08785e43e99070", + "source_subdirectory": "experimental/qwen_moe/kernels", + "qwen_endpoint_source_subdirectory": "experimental/qwen/kernels", + "corpus_sha256": "c173d047ce932c887a7540457eb7d967f0170161c0137db9234a096a5ae7d263", + "upstream_corpus_sha256": "76f504d2c19e85a1b51679bb2edc5e9a60b865f545ecaabacedd14ba7db15842", + "owned_corpus_sha256": "c173d047ce932c887a7540457eb7d967f0170161c0137db9234a096a5ae7d263", + "build_bazel_sha256": "598eeb34f8606182a732299c989387dd450c0d8f33e4275635b4a679c0f986d0", + "files": [ + { + "path": "ggml/linear_q6k_f32.loom", + "sha256": "e7d125de01614edee0ba41c6d89dbbc4fe32b2c5d12f8a5ce5db8b0f8b36236e", + "size": 28480 + }, + { + "path": "ggml/linear_q6k_q8_1_x4.loom", + "sha256": "57eda5e5b68fcb2ef6ab98ad41473351f4993a6e014c9a88d1fad9d9f8957961", + "size": 33965 + }, + { + "path": "ggml/quantize_q8_1_x4.loom", + "sha256": "7ec5f7a917d6053e3ee360cc6df07873c103a7622a87d2aa51bbd60bf9349c6d", + "size": 16671 + }, + { + "path": "qwen3_moe/attention_postprocess_f32_f16.loom", + "sha256": "15afee416d46accb1641743e9a573705c29ba2c789caecd30053bb47ef013b13", + "size": 25906 + }, + { + "path": "qwen3_moe/attention_prepare_quantized.loom", + "sha256": "c0f11099bddaa839c9ea02de66a0311dfb1da00122f1229961ba1c916f3ccd8e", + "size": 28982 + }, + { + "path": "qwen3_moe/attention_qkv_postprocess_fused.loom", + "sha256": "1994380981471568b29484be80ccb8504485df9782594efca73f067fce550763", + "size": 66529 + }, + { + "path": "qwen3_moe/attention_qkv_quantized.loom", + "sha256": "3b731f5deac992795a919ee7c34cbeb33f592747ff9bb4c9846ccc1fde08e283", + "size": 47466 + }, + { + "path": "qwen3_moe/attention_qkv_same_format_prefill.loom", + "sha256": "988571bf53b4b98290894209590ee6da5314cf51772a136389f3299c3e7e7953", + "size": 13057 + }, + { + "path": "qwen3_moe/batched_decode_expert_dispatch.loom", + "sha256": "d97e2590dff046b79248b4a4b64828a5d2a328488ef42d4db7f489f0b2fdc8e6", + "size": 43960 + }, + { + "path": "qwen3_moe/batched_decode_gate_up_q4k.loom", + "sha256": "84d79b15e6f60da6a1f12688a7125a4ef23bc3839ff97cb4a4300edaacf97cf0", + "size": 34641 + }, + { + "path": "qwen3_moe/dense_linear_quantized_f16_wmma.loom", + "sha256": "6f98adb100407e65e5eeeee847de96db2a4e55aea80432352368d058cd4a2352", + "size": 96510 + }, + { + "path": "qwen3_moe/expert_table_partition_fused.loom", + "sha256": "829eb55d5aff318c6872f0270fbad767748965a9943a229f6b17657b2a9b972d", + "size": 17122 + }, + { + "path": "qwen3_moe/flash_attention_decode_f32_f16_wmma.loom", + "sha256": "e02cb287df33b22d103117ce099310d9650061215b7df86f10b3ec3df5bf5e42", + "size": 55705 + }, + { + "path": "qwen3_moe/flash_attention_decode_q128_f32_f16_wmma.loom", + "sha256": "4a44c7d01dd37802656e250cae410593cc4cfb87e2d05ff6152bb10847e29903", + "size": 19360 + }, + { + "path": "qwen3_moe/flash_attention_decode_split_f32_f16_wmma.loom", + "sha256": "87f35d34836f360135d44ddaccfbdfca96ab812a19c25feeb5a08ae320259121", + "size": 94570 + }, + { + "path": "qwen3_moe/flash_attention_decode_split_next_q8_test.loom", + "sha256": "1a7a99fd45f351cae79e6574849558173c68532f3ca1766d2656515446095005", + "size": 14503 + }, + { + "path": "qwen3_moe/flash_attention_f32_f16_wmma.loom", + "sha256": "f5a745cee4cdc8d8287ce42398050a501f9f8b0c984654c4e53a8e89eba8c4c8", + "size": 71377 + }, + { + "path": "qwen3_moe/model_config.loom", + "sha256": "802d51e2d035461beb185cff93be4d602f98b408e385e0284b1096fe259a6095", + "size": 960 + }, + { + "path": "qwen3_moe/routed_down_q4k.loom", + "sha256": "84f47df526ebee5d095cf91084c634488475fe48428e92cd9536b413caa05903", + "size": 60068 + }, + { + "path": "qwen3_moe/routed_down_q6k.loom", + "sha256": "db5275587341a0d9da827d50642b00204d9c867c110d75c1ab35cc2ca3048211", + "size": 98823 + }, + { + "path": "qwen3_moe/routed_down_next_q8.loom", + "sha256": "c1b6b653d62f9d455562eaa87c375d14a5479e6976252ecd3b2bdade3674d1e6", + "size": 14864 + }, + { + "path": "qwen3_moe/routed_down_quantized_f16_wmma.loom", + "sha256": "cc8288c066bcd36c57fd1fc689c9a03e923597f23c02eb0cdf7030359279c12d", + "size": 81206 + }, + { + "path": "qwen3_moe/routed_down_weighted_reduce_next_rmsnorm_f32.loom", + "sha256": "351f7bcf6230058df9b78f59abb7011cf5690eedf5862ec0b75b537bf27e0e52", + "size": 11933 + }, + { + "path": "qwen3_moe/routed_down_weighted_reduce_next_rmsnorm_q8_1_x4.loom", + "sha256": "1ce06fd438cdebd63805b06066445cee717800ad1228ccfdc8294e8faccd6d95", + "size": 19387 + }, + { + "path": "qwen3_moe/routed_gate_up_swiglu_q4k.loom", + "sha256": "c8a1abd339be25eec7b5529853bf499475ebb488cb6641e81d8aef69580e15ea", + "size": 127267 + }, + { + "path": "qwen3_moe/routed_linear_q4k_f16_wmma.loom", + "sha256": "ec36cb939573e0969a4a3177ff1c3e0129d991503a3e810f9722488ce717acc0", + "size": 76615 + }, + { + "path": "qwen3_moe/router_projection_f32.loom", + "sha256": "0f4e45d2d8e1d72afdf7fadf84c889e377dc850b33d95eb82915033f7444f8bd", + "size": 18752 + }, + { + "path": "qwen3_moe/router_projection_top8_fused_f32.loom", + "sha256": "d887d3779998e84b33749d7b5da002dbb9edc48f576ed7c4dd808301fc94430f", + "size": 18152 + }, + { + "path": "qwen3_moe/router_top8_f32.loom", + "sha256": "eeeb7cdea09595829f8c6a575ad20ce08ea3c03054bd478ba4997d772d215238", + "size": 14620 + }, + { + "path": "../qwen_owned/token_embedding_q4k.loom", + "sha256": "5cb5eed0b07991e1af246965bfeff0034b066d1645d5d38cbcdd63757dee5c9a", + "size": 12428, + "upstream_path": "experimental/qwen/kernels/token_embedding_q4k.loom" + }, + { + "path": "../qwen_owned/attention_metadata.loom", + "sha256": "9614108d1680e6d0d86187c75ce0505e44fb1f53f3117f38baf56ab686b0632c", + "size": 13620, + "upstream_path": "experimental/qwen/kernels/attention_metadata.loom" + }, + { + "path": "../qwen_owned/token_embedding_bringup_workaround.loom", + "sha256": "13666225ffea241039106b33b6ed26000f9cc271b50249f5bdce61a44243e417", + "size": 15073, + "owner": "ggml-hrx" + }, + { + "path": "../qwen_owned/attention_state_initialize.loom", + "sha256": "e0492f884ea2bc0f0818e77d43b6639a81fdf994a1940416abaeac86e4165f55", + "size": 7222, + "owner": "ggml-hrx" + }, + { + "path": "../qwen_owned/attention_metadata_bringup_workaround.loom", + "sha256": "44a462abbda7e68dc84becc45501885d834b20d8445be1eb0bc7690d829b0935", + "size": 12346, + "owner": "ggml-hrx" + }, + { + "path": "../hrx_owned/gather_add_f32.loom", + "sha256": "c0f28e67b0f26d49664c1a307afb66367f3f05663e939953fba36885b7a2f6ac", + "size": 4497, + "owner": "ggml-hrx" + }, + { + "path": "../hrx_owned/add_f32.loom", + "sha256": "249b1d4c7167d0533277481cbb0e930207f14f9e2ba75fc1b8754eb0717ed094", + "size": 2444, + "owner": "ggml-hrx" + }, + { + "path": "../hrx_owned/get_rows_f32.loom", + "sha256": "06838b0226016790cdf1391a97ff321b43beb21b7743b32abe6cddce171d6300", + "size": 6080, + "owner": "ggml-hrx" + }, + { + "path": "../hrx_owned/dequant_iq3xxs_f32.loom", + "sha256": "5a66d5b63c221ffeb1471ed5c7da4d420c4657d377720675a548db709aabfdbb", + "size": 8227, + "owner": "ggml-hrx" + }, + { + "path": "../hrx_owned/mul_mat_vec_iq3xxs_f32.loom", + "sha256": "87c8320c4ae65655b68b9b66c72c5dab58ee1ba9dd40aa8b94e733b16cfe0d6c", + "size": 7790, + "owner": "ggml-hrx" + } + ], + "exports": [ + { + "name": "ggml_add_f32", + "symbol": "ggml_add_f32", + "source": "../hrx_owned/add_f32.loom", + "workload_parameters": [ + { + "name": "element_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "element_count", + "type": "index" + } + ], + "bindings": [ + "a", + "b", + "output" + ], + "binding_access": [ + "read", + "read", + "read_write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "../hrx_owned/add_f32.loom" + ], + "library_sources": [] + }, + "compile_dependencies": [] + }, + { + "name": "ggml_gather_add_f32", + "symbol": "ggml_gather_add_f32", + "source": "../hrx_owned/gather_add_f32.loom", + "workload_parameters": [ + { + "name": "source_token_count", + "type": "index" + }, + { + "name": "output_token_count", + "type": "index" + }, + { + "name": "hidden_size", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "source_token_count", + "type": "index" + }, + { + "name": "output_token_count", + "type": "index" + }, + { + "name": "hidden_size", + "type": "index" + } + ], + "bindings": [ + "attention", + "residual", + "output_ids", + "output" + ], + "binding_access": [ + "read", + "read", + "read", + "read_write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "../hrx_owned/gather_add_f32.loom" + ], + "library_sources": [] + }, + "compile_dependencies": [] + }, + { + "name": "ggml_linear_q6k_f32_wave32", + "symbol": "ggml_linear_q6k_f32_wave32", + "source": "ggml/linear_q6k_f32.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + }, + { + "name": "input_size", + "type": "index" + }, + { + "name": "output_size", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + }, + { + "name": "input_size", + "type": "index" + }, + { + "name": "output_size", + "type": "index" + } + ], + "bindings": [ + "input", + "weight", + "output" + ], + "binding_access": [ + "read", + "read", + "read_write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "archive", + "link_module": "linear_q6k_f32_linked", + "primary_sources": [ + "ggml/linear_q6k_f32.loom" + ], + "library_sources": [ + "ggml/linear_q6k_q8_1_x4.loom", + "ggml/quantize_q8_1_x4.loom" + ] + }, + "compile_dependencies": [ + "ggml/linear_q6k_q8_1_x4.loom", + "ggml/quantize_q8_1_x4.loom" + ] + }, + { + "name": "ggml_linear_q6k_f32_wave64", + "symbol": "ggml_linear_q6k_f32_wave64", + "source": "ggml/linear_q6k_f32.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + }, + { + "name": "input_size", + "type": "index" + }, + { + "name": "output_size", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + }, + { + "name": "input_size", + "type": "index" + }, + { + "name": "output_size", + "type": "index" + } + ], + "bindings": [ + "input", + "weight", + "output" + ], + "binding_access": [ + "read", + "read", + "read_write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "archive", + "link_module": "linear_q6k_f32_linked", + "primary_sources": [ + "ggml/linear_q6k_f32.loom" + ], + "library_sources": [ + "ggml/linear_q6k_q8_1_x4.loom", + "ggml/quantize_q8_1_x4.loom" + ] + }, + "compile_dependencies": [ + "ggml/linear_q6k_q8_1_x4.loom", + "ggml/quantize_q8_1_x4.loom" + ] + }, + { + "name": "ggml_linear_q6k_q8_1_x4", + "symbol": "ggml_linear_q6k_q8_1_x4", + "source": "ggml/linear_q6k_q8_1_x4.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + }, + { + "name": "input_size", + "type": "index" + }, + { + "name": "output_size", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + }, + { + "name": "input_size", + "type": "index" + }, + { + "name": "output_size", + "type": "index" + } + ], + "bindings": [ + "q8_input", + "weight", + "output" + ], + "binding_access": [ + "read", + "read", + "read_write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "archive", + "link_module": "linear_q6k_q8_1_x4_linked", + "primary_sources": [ + "ggml/linear_q6k_q8_1_x4.loom" + ], + "library_sources": [ + "ggml/quantize_q8_1_x4.loom" + ] + }, + "compile_dependencies": [ + "ggml/quantize_q8_1_x4.loom" + ] + }, + { + "name": "ggml_linear_q6k_q8_1_x4", + "symbol": "ggml_linear_q6k_q8_1_x4_wave64_block4", + "source": "ggml/linear_q6k_q8_1_x4.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + }, + { + "name": "input_size", + "type": "index" + }, + { + "name": "output_size", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + }, + { + "name": "input_size", + "type": "index" + }, + { + "name": "output_size", + "type": "index" + } + ], + "bindings": [ + "q8_input", + "weight", + "output" + ], + "binding_access": [ + "read", + "read", + "read_write" + ], + "target_selector": "gfx1151", + "compile_recipe": { + "mode": "archive", + "link_module": "linear_q6k_q8_1_x4_linked", + "primary_sources": [ + "ggml/linear_q6k_q8_1_x4.loom" + ], + "library_sources": [ + "ggml/quantize_q8_1_x4.loom" + ] + }, + "compile_dependencies": [ + "ggml/quantize_q8_1_x4.loom" + ] + }, + { + "name": "ggml_q8_1_x4_inspect_one_group", + "symbol": "ggml_q8_1_x4_inspect_one_group", + "source": "ggml/quantize_q8_1_x4.loom", + "workload_parameters": [], + "launch_parameters": [], + "bindings": [ + "packed", + "words", + "d_values", + "s_values" + ], + "binding_access": [ + "read", + "write", + "write", + "write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ggml/quantize_q8_1_x4.loom" + ], + "library_sources": [] + }, + "compile_dependencies": [] + }, + { + "name": "ggml_quantize_q8_1_x4_f32", + "symbol": "ggml_quantize_q8_1_x4_f32", + "source": "ggml/quantize_q8_1_x4.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + }, + { + "name": "input_size", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + }, + { + "name": "input_size", + "type": "index" + } + ], + "bindings": [ + "input", + "output" + ], + "binding_access": [ + "read", + "read_write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ggml/quantize_q8_1_x4.loom" + ], + "library_sources": [] + }, + "compile_dependencies": [] + }, + { + "name": "qwen3_moe_attention_key_q4_aggregate_prefill_512", + "symbol": "qwen3_moe_attention_key_q4_aggregate_prefill_512", + "source": "qwen3_moe/attention_qkv_same_format_prefill.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "bindings": [ + "input", + "combined_weight", + "combined_output" + ], + "binding_access": [ + "read", + "read", + "read" + ], + "target_selector": "", + "compile_recipe": { + "mode": "archive", + "link_module": "attention_qkv_same_format_prefill_linked", + "primary_sources": [ + "qwen3_moe/attention_qkv_same_format_prefill.loom" + ], + "library_sources": [ + "qwen3_moe/dense_linear_quantized_f16_wmma.loom", + "ggml/linear_q6k_q8_1_x4.loom", + "qwen3_moe/attention_prepare_quantized.loom", + "qwen3_moe/model_config.loom", + "qwen3_moe/routed_linear_q4k_f16_wmma.loom", + "qwen3_moe/routed_gate_up_swiglu_q4k.loom", + "ggml/quantize_q8_1_x4.loom" + ] + }, + "compile_dependencies": [ + "qwen3_moe/dense_linear_quantized_f16_wmma.loom", + "ggml/linear_q6k_q8_1_x4.loom", + "qwen3_moe/attention_prepare_quantized.loom", + "qwen3_moe/model_config.loom", + "qwen3_moe/routed_linear_q4k_f16_wmma.loom", + "qwen3_moe/routed_gate_up_swiglu_q4k.loom", + "ggml/quantize_q8_1_x4.loom" + ] + }, + { + "name": "qwen3_moe_attention_postprocess_f32_f16", + "symbol": "qwen3_moe_attention_postprocess_f32_f16", + "source": "qwen3_moe/attention_postprocess_f32_f16.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + }, + { + "name": "cache_row_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + }, + { + "name": "cache_row_count", + "type": "index" + } + ], + "bindings": [ + "positions", + "key_cache_indices", + "value_cache_indices", + "query_input", + "key_input", + "value_input", + "query_norm_weight", + "key_norm_weight", + "inverse_frequencies", + "query_output", + "key_cache", + "value_cache" + ], + "binding_access": [ + "read", + "read", + "read", + "read", + "read", + "read", + "read", + "read", + "read", + "read_write", + "read_write", + "read_write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "archive", + "link_module": "attention_postprocess_f32_f16_linked", + "primary_sources": [ + "qwen3_moe/attention_postprocess_f32_f16.loom" + ], + "library_sources": [ + "qwen3_moe/model_config.loom" + ] + }, + "compile_dependencies": [ + "qwen3_moe/model_config.loom" + ] + }, + { + "name": "qwen3_moe_attention_qkv_postprocess_fused_decode", + "symbol": "qwen3_moe_attention_qkv_postprocess_fused_decode", + "source": "qwen3_moe/attention_qkv_postprocess_fused.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + }, + { + "name": "cache_row_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + }, + { + "name": "cache_row_count", + "type": "index" + } + ], + "bindings": [ + "q8_input", + "query_weight", + "key_weight", + "value_weight", + "positions", + "key_cache_indices", + "value_cache_indices", + "query_output_raw", + "key_output_raw", + "value_output_raw", + "query_norm_weight", + "key_norm_weight", + "inverse_frequencies", + "query_output", + "key_cache", + "value_cache", + "completion_counters" + ], + "binding_access": [ + "read", + "read", + "read", + "read", + "read", + "read", + "read", + "read_write", + "read_write", + "read_write", + "read", + "read", + "read", + "read_write", + "read_write", + "read_write", + "read_write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "archive", + "link_module": "attention_qkv_postprocess_fused_linked", + "primary_sources": [ + "qwen3_moe/attention_qkv_postprocess_fused.loom" + ], + "library_sources": [ + "qwen3_moe/attention_postprocess_f32_f16.loom", + "qwen3_moe/model_config.loom", + "qwen3_moe/attention_qkv_quantized.loom", + "qwen3_moe/dense_linear_quantized_f16_wmma.loom", + "ggml/linear_q6k_q8_1_x4.loom", + "qwen3_moe/attention_prepare_quantized.loom", + "qwen3_moe/routed_linear_q4k_f16_wmma.loom", + "qwen3_moe/routed_gate_up_swiglu_q4k.loom", + "ggml/quantize_q8_1_x4.loom" + ] + }, + "compile_dependencies": [ + "qwen3_moe/attention_postprocess_f32_f16.loom", + "qwen3_moe/model_config.loom", + "qwen3_moe/attention_qkv_quantized.loom", + "qwen3_moe/dense_linear_quantized_f16_wmma.loom", + "ggml/linear_q6k_q8_1_x4.loom", + "qwen3_moe/attention_prepare_quantized.loom", + "qwen3_moe/routed_linear_q4k_f16_wmma.loom", + "qwen3_moe/routed_gate_up_swiglu_q4k.loom", + "ggml/quantize_q8_1_x4.loom" + ] + }, + { + "name": "qwen3_moe_attention_qkv_postprocess_fused_decode_q4", + "symbol": "qwen3_moe_attention_qkv_postprocess_fused_decode_q4", + "source": "qwen3_moe/attention_qkv_postprocess_fused.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + }, + { + "name": "cache_row_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + }, + { + "name": "cache_row_count", + "type": "index" + } + ], + "bindings": [ + "q8_input", + "query_weight", + "key_weight", + "value_weight", + "positions", + "key_cache_indices", + "value_cache_indices", + "query_output_raw", + "key_output_raw", + "value_output_raw", + "query_norm_weight", + "key_norm_weight", + "inverse_frequencies", + "query_output", + "key_cache", + "value_cache", + "completion_counters" + ], + "binding_access": [ + "read", + "read", + "read", + "read", + "read", + "read", + "read", + "read", + "read", + "read", + "read", + "read", + "read", + "read_write", + "read_write", + "read_write", + "read_write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "archive", + "link_module": "attention_qkv_postprocess_fused_linked", + "primary_sources": [ + "qwen3_moe/attention_qkv_postprocess_fused.loom" + ], + "library_sources": [ + "qwen3_moe/attention_postprocess_f32_f16.loom", + "qwen3_moe/model_config.loom", + "qwen3_moe/attention_qkv_quantized.loom", + "qwen3_moe/dense_linear_quantized_f16_wmma.loom", + "ggml/linear_q6k_q8_1_x4.loom", + "qwen3_moe/attention_prepare_quantized.loom", + "qwen3_moe/routed_linear_q4k_f16_wmma.loom", + "qwen3_moe/routed_gate_up_swiglu_q4k.loom", + "ggml/quantize_q8_1_x4.loom" + ] + }, + "compile_dependencies": [ + "qwen3_moe/attention_postprocess_f32_f16.loom", + "qwen3_moe/model_config.loom", + "qwen3_moe/attention_qkv_quantized.loom", + "qwen3_moe/dense_linear_quantized_f16_wmma.loom", + "ggml/linear_q6k_q8_1_x4.loom", + "qwen3_moe/attention_prepare_quantized.loom", + "qwen3_moe/routed_linear_q4k_f16_wmma.loom", + "qwen3_moe/routed_gate_up_swiglu_q4k.loom", + "ggml/quantize_q8_1_x4.loom" + ] + }, + { + "name": "qwen3_moe_attention_qkv_postprocess_fused_decode_q6", + "symbol": "qwen3_moe_attention_qkv_postprocess_fused_decode_q6", + "source": "qwen3_moe/attention_qkv_postprocess_fused.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + }, + { + "name": "cache_row_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + }, + { + "name": "cache_row_count", + "type": "index" + } + ], + "bindings": [ + "q8_input", + "query_weight", + "key_weight", + "value_weight", + "positions", + "key_cache_indices", + "value_cache_indices", + "query_output_raw", + "key_output_raw", + "value_output_raw", + "query_norm_weight", + "key_norm_weight", + "inverse_frequencies", + "query_output", + "key_cache", + "value_cache", + "completion_counters" + ], + "binding_access": [ + "read", + "read", + "read", + "read", + "read", + "read", + "read", + "read", + "read", + "read", + "read", + "read", + "read", + "read_write", + "read_write", + "read_write", + "read_write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "archive", + "link_module": "attention_qkv_postprocess_fused_linked", + "primary_sources": [ + "qwen3_moe/attention_qkv_postprocess_fused.loom" + ], + "library_sources": [ + "qwen3_moe/attention_postprocess_f32_f16.loom", + "qwen3_moe/model_config.loom", + "qwen3_moe/attention_qkv_quantized.loom", + "qwen3_moe/dense_linear_quantized_f16_wmma.loom", + "ggml/linear_q6k_q8_1_x4.loom", + "qwen3_moe/attention_prepare_quantized.loom", + "qwen3_moe/routed_linear_q4k_f16_wmma.loom", + "qwen3_moe/routed_gate_up_swiglu_q4k.loom", + "ggml/quantize_q8_1_x4.loom" + ] + }, + "compile_dependencies": [ + "qwen3_moe/attention_postprocess_f32_f16.loom", + "qwen3_moe/model_config.loom", + "qwen3_moe/attention_qkv_quantized.loom", + "qwen3_moe/dense_linear_quantized_f16_wmma.loom", + "ggml/linear_q6k_q8_1_x4.loom", + "qwen3_moe/attention_prepare_quantized.loom", + "qwen3_moe/routed_linear_q4k_f16_wmma.loom", + "qwen3_moe/routed_gate_up_swiglu_q4k.loom", + "ggml/quantize_q8_1_x4.loom" + ] + }, + { + "name": "qwen3_moe_attention_qkv_q4_prefill_512", + "symbol": "qwen3_moe_attention_qkv_q4_prefill_512", + "source": "qwen3_moe/attention_qkv_same_format_prefill.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "bindings": [ + "input", + "combined_weight", + "combined_output" + ], + "binding_access": [ + "read", + "read", + "read" + ], + "target_selector": "", + "compile_recipe": { + "mode": "archive", + "link_module": "attention_qkv_same_format_prefill_linked", + "primary_sources": [ + "qwen3_moe/attention_qkv_same_format_prefill.loom" + ], + "library_sources": [ + "qwen3_moe/dense_linear_quantized_f16_wmma.loom", + "ggml/linear_q6k_q8_1_x4.loom", + "qwen3_moe/attention_prepare_quantized.loom", + "qwen3_moe/model_config.loom", + "qwen3_moe/routed_linear_q4k_f16_wmma.loom", + "qwen3_moe/routed_gate_up_swiglu_q4k.loom", + "ggml/quantize_q8_1_x4.loom" + ] + }, + "compile_dependencies": [ + "qwen3_moe/dense_linear_quantized_f16_wmma.loom", + "ggml/linear_q6k_q8_1_x4.loom", + "qwen3_moe/attention_prepare_quantized.loom", + "qwen3_moe/model_config.loom", + "qwen3_moe/routed_linear_q4k_f16_wmma.loom", + "qwen3_moe/routed_gate_up_swiglu_q4k.loom", + "ggml/quantize_q8_1_x4.loom" + ] + }, + { + "name": "qwen3_moe_attention_qkv_quantized", + "symbol": "qwen3_moe_attention_qkv_quantized", + "source": "qwen3_moe/attention_qkv_quantized.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "bindings": [ + "q8_input", + "query_weight", + "key_weight", + "value_weight", + "query_output", + "key_output", + "value_output" + ], + "binding_access": [ + "read", + "read", + "read", + "read", + "read_write", + "read_write", + "read_write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "archive", + "link_module": "attention_qkv_quantized_linked", + "primary_sources": [ + "qwen3_moe/attention_qkv_quantized.loom" + ], + "library_sources": [ + "qwen3_moe/dense_linear_quantized_f16_wmma.loom", + "ggml/linear_q6k_q8_1_x4.loom", + "qwen3_moe/attention_prepare_quantized.loom", + "qwen3_moe/model_config.loom", + "qwen3_moe/routed_linear_q4k_f16_wmma.loom", + "qwen3_moe/routed_gate_up_swiglu_q4k.loom", + "ggml/quantize_q8_1_x4.loom" + ] + }, + "compile_dependencies": [ + "qwen3_moe/dense_linear_quantized_f16_wmma.loom", + "ggml/linear_q6k_q8_1_x4.loom", + "qwen3_moe/attention_prepare_quantized.loom", + "qwen3_moe/model_config.loom", + "qwen3_moe/routed_linear_q4k_f16_wmma.loom", + "qwen3_moe/routed_gate_up_swiglu_q4k.loom", + "ggml/quantize_q8_1_x4.loom" + ] + }, + { + "name": "qwen3_moe_attention_query_q4_aggregate_prefill_512", + "symbol": "qwen3_moe_attention_query_q4_aggregate_prefill_512", + "source": "qwen3_moe/attention_qkv_same_format_prefill.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "bindings": [ + "input", + "combined_weight", + "combined_output" + ], + "binding_access": [ + "read", + "read", + "read" + ], + "target_selector": "", + "compile_recipe": { + "mode": "archive", + "link_module": "attention_qkv_same_format_prefill_linked", + "primary_sources": [ + "qwen3_moe/attention_qkv_same_format_prefill.loom" + ], + "library_sources": [ + "qwen3_moe/dense_linear_quantized_f16_wmma.loom", + "ggml/linear_q6k_q8_1_x4.loom", + "qwen3_moe/attention_prepare_quantized.loom", + "qwen3_moe/model_config.loom", + "qwen3_moe/routed_linear_q4k_f16_wmma.loom", + "qwen3_moe/routed_gate_up_swiglu_q4k.loom", + "ggml/quantize_q8_1_x4.loom" + ] + }, + "compile_dependencies": [ + "qwen3_moe/dense_linear_quantized_f16_wmma.loom", + "ggml/linear_q6k_q8_1_x4.loom", + "qwen3_moe/attention_prepare_quantized.loom", + "qwen3_moe/model_config.loom", + "qwen3_moe/routed_linear_q4k_f16_wmma.loom", + "qwen3_moe/routed_gate_up_swiglu_q4k.loom", + "ggml/quantize_q8_1_x4.loom" + ] + }, + { + "name": "qwen3_moe_attention_rmsnorm_quantize_q8_1_x4", + "symbol": "qwen3_moe_attention_rmsnorm_quantize_q8_1_x4", + "source": "qwen3_moe/attention_prepare_quantized.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "bindings": [ + "input", + "weight", + "output" + ], + "binding_access": [ + "read", + "read", + "read_write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "archive", + "link_module": "attention_prepare_quantized_linked", + "primary_sources": [ + "qwen3_moe/attention_prepare_quantized.loom" + ], + "library_sources": [ + "qwen3_moe/model_config.loom", + "ggml/quantize_q8_1_x4.loom" + ] + }, + "compile_dependencies": [ + "qwen3_moe/model_config.loom", + "ggml/quantize_q8_1_x4.loom" + ] + }, + { + "name": "qwen3_moe_attention_value_q4_aggregate_prefill_512", + "symbol": "qwen3_moe_attention_value_q4_aggregate_prefill_512", + "source": "qwen3_moe/attention_qkv_same_format_prefill.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "bindings": [ + "input", + "combined_weight", + "combined_output" + ], + "binding_access": [ + "read", + "read", + "read" + ], + "target_selector": "", + "compile_recipe": { + "mode": "archive", + "link_module": "attention_qkv_same_format_prefill_linked", + "primary_sources": [ + "qwen3_moe/attention_qkv_same_format_prefill.loom" + ], + "library_sources": [ + "qwen3_moe/dense_linear_quantized_f16_wmma.loom", + "ggml/linear_q6k_q8_1_x4.loom", + "qwen3_moe/attention_prepare_quantized.loom", + "qwen3_moe/model_config.loom", + "qwen3_moe/routed_linear_q4k_f16_wmma.loom", + "qwen3_moe/routed_gate_up_swiglu_q4k.loom", + "ggml/quantize_q8_1_x4.loom" + ] + }, + "compile_dependencies": [ + "qwen3_moe/dense_linear_quantized_f16_wmma.loom", + "ggml/linear_q6k_q8_1_x4.loom", + "qwen3_moe/attention_prepare_quantized.loom", + "qwen3_moe/model_config.loom", + "qwen3_moe/routed_linear_q4k_f16_wmma.loom", + "qwen3_moe/routed_gate_up_swiglu_q4k.loom", + "ggml/quantize_q8_1_x4.loom" + ] + }, + { + "name": "qwen3_moe_batched_decode_gate_up_q4k_rows2", + "symbol": "qwen3_moe_batched_decode_gate_up_q4k_rows2", + "source": "qwen3_moe/batched_decode_gate_up_q4k.loom", + "workload_parameters": [ + { + "name": "descriptor_count", + "type": "index" + }, + { + "name": "queue_ordinal", + "type": "index" + }, + { + "name": "token_count", + "type": "index" + }, + { + "name": "route_count", + "type": "index" + }, + { + "name": "expert_count", + "type": "index" + }, + { + "name": "output_size", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "descriptor_count", + "type": "index" + }, + { + "name": "queue_ordinal", + "type": "index" + }, + { + "name": "token_count", + "type": "index" + }, + { + "name": "route_count", + "type": "index" + }, + { + "name": "expert_count", + "type": "index" + }, + { + "name": "output_size", + "type": "index" + } + ], + "bindings": [ + "queue_descriptors", + "assignment_ordinals", + "q8_input", + "gate_weight", + "up_weight", + "output" + ], + "binding_access": [ + "read", + "read", + "read", + "read", + "read", + "read_write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "archive", + "link_module": "batched_decode_gate_up_q4k_linked", + "primary_sources": [ + "qwen3_moe/batched_decode_gate_up_q4k.loom" + ], + "library_sources": [ + "ggml/quantize_q8_1_x4.loom", + "qwen3_moe/batched_decode_expert_dispatch.loom", + "qwen3_moe/routed_gate_up_swiglu_q4k.loom" + ] + }, + "compile_dependencies": [ + "ggml/quantize_q8_1_x4.loom", + "qwen3_moe/batched_decode_expert_dispatch.loom", + "qwen3_moe/routed_gate_up_swiglu_q4k.loom" + ] + }, + { + "name": "qwen3_moe_build_batched_decode_expert_dispatch", + "symbol": "qwen3_moe_build_batched_decode_expert_dispatch", + "source": "qwen3_moe/batched_decode_expert_dispatch.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + }, + { + "name": "route_count", + "type": "index" + }, + { + "name": "route_stride", + "type": "index" + }, + { + "name": "expert_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + }, + { + "name": "route_count", + "type": "index" + }, + { + "name": "route_stride", + "type": "index" + }, + { + "name": "expert_count", + "type": "index" + } + ], + "bindings": [ + "route_ids", + "assignment_ordinals", + "queue_counts", + "queue_descriptors" + ], + "binding_access": [ + "read", + "read", + "read", + "read" + ], + "target_selector": "", + "compile_recipe": { + "mode": "archive", + "link_module": "batched_decode_expert_dispatch_linked", + "primary_sources": [ + "qwen3_moe/batched_decode_expert_dispatch.loom" + ], + "library_sources": [ + "qwen3_moe/model_config.loom" + ] + }, + "compile_dependencies": [ + "qwen3_moe/model_config.loom" + ] + }, + { + "name": "qwen3_moe_build_batched_decode_expert_dispatch_reference", + "symbol": "qwen3_moe_build_batched_decode_expert_dispatch_reference", + "source": "qwen3_moe/batched_decode_expert_dispatch.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + }, + { + "name": "route_count", + "type": "index" + }, + { + "name": "route_stride", + "type": "index" + }, + { + "name": "expert_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + }, + { + "name": "route_count", + "type": "index" + }, + { + "name": "route_stride", + "type": "index" + }, + { + "name": "expert_count", + "type": "index" + } + ], + "bindings": [ + "route_ids", + "assignment_ordinals", + "queue_counts", + "queue_descriptors" + ], + "binding_access": [ + "read", + "read", + "read", + "read" + ], + "target_selector": "", + "compile_recipe": { + "mode": "archive", + "link_module": "batched_decode_expert_dispatch_linked", + "primary_sources": [ + "qwen3_moe/batched_decode_expert_dispatch.loom" + ], + "library_sources": [ + "qwen3_moe/model_config.loom" + ] + }, + "compile_dependencies": [ + "qwen3_moe/model_config.loom" + ] + }, + { + "name": "qwen3_moe_build_expert_partition_table", + "symbol": "qwen3_moe_build_expert_partition_table", + "source": "qwen3_moe/routed_gate_up_swiglu_q4k.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + }, + { + "name": "route_count", + "type": "index" + }, + { + "name": "expert_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + }, + { + "name": "route_count", + "type": "index" + }, + { + "name": "expert_count", + "type": "index" + } + ], + "bindings": [ + "expert_table", + "partition_table" + ], + "binding_access": [ + "read", + "write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "archive", + "link_module": "routed_gate_up_swiglu_q4k_linked", + "primary_sources": [ + "qwen3_moe/routed_gate_up_swiglu_q4k.loom" + ], + "library_sources": [ + "ggml/quantize_q8_1_x4.loom" + ] + }, + "compile_dependencies": [ + "ggml/quantize_q8_1_x4.loom" + ] + }, + { + "name": "qwen3_moe_build_expert_table", + "symbol": "qwen3_moe_build_expert_table", + "source": "qwen3_moe/routed_gate_up_swiglu_q4k.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + }, + { + "name": "route_count", + "type": "index" + }, + { + "name": "route_stride", + "type": "index" + }, + { + "name": "expert_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + }, + { + "name": "route_count", + "type": "index" + }, + { + "name": "route_stride", + "type": "index" + }, + { + "name": "expert_count", + "type": "index" + } + ], + "bindings": [ + "route_ids", + "expert_table" + ], + "binding_access": [ + "read", + "write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "archive", + "link_module": "routed_gate_up_swiglu_q4k_linked", + "primary_sources": [ + "qwen3_moe/routed_gate_up_swiglu_q4k.loom" + ], + "library_sources": [ + "ggml/quantize_q8_1_x4.loom" + ] + }, + "compile_dependencies": [ + "ggml/quantize_q8_1_x4.loom" + ] + }, + { + "name": "qwen3_moe_build_expert_table_partition_prefill_512", + "symbol": "qwen3_moe_build_expert_table_partition_prefill_512", + "source": "qwen3_moe/expert_table_partition_fused.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + }, + { + "name": "route_count", + "type": "index" + }, + { + "name": "route_stride", + "type": "index" + }, + { + "name": "expert_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + }, + { + "name": "route_count", + "type": "index" + }, + { + "name": "route_stride", + "type": "index" + }, + { + "name": "expert_count", + "type": "index" + } + ], + "bindings": [ + "route_ids", + "expert_table", + "partition_table", + "completion_counter" + ], + "binding_access": [ + "read", + "write", + "write", + "read_write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "archive", + "link_module": "expert_table_partition_fused_linked", + "primary_sources": [ + "qwen3_moe/expert_table_partition_fused.loom" + ], + "library_sources": [ + "qwen3_moe/routed_gate_up_swiglu_q4k.loom", + "ggml/quantize_q8_1_x4.loom" + ] + }, + "compile_dependencies": [ + "qwen3_moe/routed_gate_up_swiglu_q4k.loom", + "ggml/quantize_q8_1_x4.loom" + ] + }, + { + "name": "qwen3_moe_dense_linear_q4k_f16_wmma", + "symbol": "qwen3_moe_dense_linear_q4k_f16_wmma", + "source": "qwen3_moe/dense_linear_quantized_f16_wmma.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "bindings": [ + "input", + "weight", + "output" + ], + "binding_access": [ + "read", + "read", + "read_write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "archive", + "link_module": "dense_linear_quantized_f16_wmma_linked", + "primary_sources": [ + "qwen3_moe/dense_linear_quantized_f16_wmma.loom" + ], + "library_sources": [ + "ggml/linear_q6k_q8_1_x4.loom", + "qwen3_moe/attention_prepare_quantized.loom", + "qwen3_moe/model_config.loom", + "qwen3_moe/routed_linear_q4k_f16_wmma.loom", + "qwen3_moe/routed_gate_up_swiglu_q4k.loom", + "ggml/quantize_q8_1_x4.loom" + ] + }, + "compile_dependencies": [ + "ggml/linear_q6k_q8_1_x4.loom", + "qwen3_moe/attention_prepare_quantized.loom", + "qwen3_moe/model_config.loom", + "qwen3_moe/routed_linear_q4k_f16_wmma.loom", + "qwen3_moe/routed_gate_up_swiglu_q4k.loom", + "ggml/quantize_q8_1_x4.loom" + ] + }, + { + "name": "qwen3_moe_dense_linear_q4k_f16_wmma_parameterized", + "symbol": "qwen3_moe_dense_linear_q4k_f16_wmma_parameterized", + "source": "qwen3_moe/dense_linear_quantized_f16_wmma.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + }, + { + "name": "input_size", + "type": "index" + }, + { + "name": "output_size", + "type": "index" + }, + { + "name": "output_accumulation", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + }, + { + "name": "input_size", + "type": "index" + }, + { + "name": "output_size", + "type": "index" + }, + { + "name": "output_accumulation", + "type": "index" + } + ], + "bindings": [ + "input", + "weight", + "output" + ], + "binding_access": [ + "read", + "read", + "read_write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "archive", + "link_module": "dense_linear_quantized_f16_wmma_linked", + "primary_sources": [ + "qwen3_moe/dense_linear_quantized_f16_wmma.loom" + ], + "library_sources": [ + "ggml/linear_q6k_q8_1_x4.loom", + "qwen3_moe/attention_prepare_quantized.loom", + "qwen3_moe/model_config.loom", + "qwen3_moe/routed_linear_q4k_f16_wmma.loom", + "qwen3_moe/routed_gate_up_swiglu_q4k.loom", + "ggml/quantize_q8_1_x4.loom" + ] + }, + "compile_dependencies": [ + "ggml/linear_q6k_q8_1_x4.loom", + "qwen3_moe/attention_prepare_quantized.loom", + "qwen3_moe/model_config.loom", + "qwen3_moe/routed_linear_q4k_f16_wmma.loom", + "qwen3_moe/routed_gate_up_swiglu_q4k.loom", + "ggml/quantize_q8_1_x4.loom" + ] + }, + { + "name": "qwen3_moe_dense_linear_q4k_q8_1_x4", + "symbol": "qwen3_moe_dense_linear_q4k_q8_1_x4", + "source": "qwen3_moe/dense_linear_quantized_f16_wmma.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "bindings": [ + "q8_input", + "weight", + "output" + ], + "binding_access": [ + "read", + "read", + "read_write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "archive", + "link_module": "dense_linear_quantized_f16_wmma_linked", + "primary_sources": [ + "qwen3_moe/dense_linear_quantized_f16_wmma.loom" + ], + "library_sources": [ + "ggml/linear_q6k_q8_1_x4.loom", + "qwen3_moe/attention_prepare_quantized.loom", + "qwen3_moe/model_config.loom", + "qwen3_moe/routed_linear_q4k_f16_wmma.loom", + "qwen3_moe/routed_gate_up_swiglu_q4k.loom", + "ggml/quantize_q8_1_x4.loom" + ] + }, + "compile_dependencies": [ + "ggml/linear_q6k_q8_1_x4.loom", + "qwen3_moe/attention_prepare_quantized.loom", + "qwen3_moe/model_config.loom", + "qwen3_moe/routed_linear_q4k_f16_wmma.loom", + "qwen3_moe/routed_gate_up_swiglu_q4k.loom", + "ggml/quantize_q8_1_x4.loom" + ] + }, + { + "name": "qwen3_moe_dense_linear_q4k_q8_1_x4_next_q8", + "symbol": "qwen3_moe_dense_linear_q4k_q8_1_x4_next_q8", + "source": "qwen3_moe/dense_linear_quantized_f16_wmma.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "bindings": [ + "q8_input", + "weight", + "output", + "norm_weight", + "normalized_output", + "completion_counter", + "next_q8_output" + ], + "binding_access": [ + "read", + "read", + "read_write", + "read", + "read_write", + "read_write", + "read_write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "archive", + "link_module": "dense_linear_quantized_f16_wmma_linked", + "primary_sources": [ + "qwen3_moe/dense_linear_quantized_f16_wmma.loom" + ], + "library_sources": [ + "ggml/linear_q6k_q8_1_x4.loom", + "qwen3_moe/attention_prepare_quantized.loom", + "qwen3_moe/model_config.loom", + "qwen3_moe/routed_linear_q4k_f16_wmma.loom", + "qwen3_moe/routed_gate_up_swiglu_q4k.loom", + "ggml/quantize_q8_1_x4.loom" + ] + }, + "compile_dependencies": [ + "ggml/linear_q6k_q8_1_x4.loom", + "qwen3_moe/attention_prepare_quantized.loom", + "qwen3_moe/model_config.loom", + "qwen3_moe/routed_linear_q4k_f16_wmma.loom", + "qwen3_moe/routed_gate_up_swiglu_q4k.loom", + "ggml/quantize_q8_1_x4.loom" + ] + }, + { + "name": "qwen3_moe_dense_linear_q6k_f16_wmma", + "symbol": "qwen3_moe_dense_linear_q6k_f16_wmma", + "source": "qwen3_moe/dense_linear_quantized_f16_wmma.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "bindings": [ + "input", + "weight", + "output" + ], + "binding_access": [ + "read", + "read", + "read_write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "archive", + "link_module": "dense_linear_quantized_f16_wmma_linked", + "primary_sources": [ + "qwen3_moe/dense_linear_quantized_f16_wmma.loom" + ], + "library_sources": [ + "ggml/linear_q6k_q8_1_x4.loom", + "qwen3_moe/attention_prepare_quantized.loom", + "qwen3_moe/model_config.loom", + "qwen3_moe/routed_linear_q4k_f16_wmma.loom", + "qwen3_moe/routed_gate_up_swiglu_q4k.loom", + "ggml/quantize_q8_1_x4.loom" + ] + }, + "compile_dependencies": [ + "ggml/linear_q6k_q8_1_x4.loom", + "qwen3_moe/attention_prepare_quantized.loom", + "qwen3_moe/model_config.loom", + "qwen3_moe/routed_linear_q4k_f16_wmma.loom", + "qwen3_moe/routed_gate_up_swiglu_q4k.loom", + "ggml/quantize_q8_1_x4.loom" + ] + }, + { + "name": "qwen3_moe_dense_linear_q6k_f16_wmma_parameterized", + "symbol": "qwen3_moe_dense_linear_q6k_f16_wmma_parameterized", + "source": "qwen3_moe/dense_linear_quantized_f16_wmma.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + }, + { + "name": "input_size", + "type": "index" + }, + { + "name": "output_size", + "type": "index" + }, + { + "name": "output_accumulation", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + }, + { + "name": "input_size", + "type": "index" + }, + { + "name": "output_size", + "type": "index" + }, + { + "name": "output_accumulation", + "type": "index" + } + ], + "bindings": [ + "input", + "weight", + "output" + ], + "binding_access": [ + "read", + "read", + "read_write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "archive", + "link_module": "dense_linear_quantized_f16_wmma_linked", + "primary_sources": [ + "qwen3_moe/dense_linear_quantized_f16_wmma.loom" + ], + "library_sources": [ + "ggml/linear_q6k_q8_1_x4.loom", + "qwen3_moe/attention_prepare_quantized.loom", + "qwen3_moe/model_config.loom", + "qwen3_moe/routed_linear_q4k_f16_wmma.loom", + "qwen3_moe/routed_gate_up_swiglu_q4k.loom", + "ggml/quantize_q8_1_x4.loom" + ] + }, + "compile_dependencies": [ + "ggml/linear_q6k_q8_1_x4.loom", + "qwen3_moe/attention_prepare_quantized.loom", + "qwen3_moe/model_config.loom", + "qwen3_moe/routed_linear_q4k_f16_wmma.loom", + "qwen3_moe/routed_gate_up_swiglu_q4k.loom", + "ggml/quantize_q8_1_x4.loom" + ] + }, + { + "name": "qwen3_moe_flash_attention_decode_f32_f16_wmma", + "symbol": "qwen3_moe_flash_attention_decode_f32_f16_wmma", + "source": "qwen3_moe/flash_attention_decode_f32_f16_wmma.loom", + "workload_parameters": [ + { + "name": "key_value_token_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "key_value_token_count", + "type": "index" + } + ], + "bindings": [ + "query", + "key", + "value", + "mask", + "output" + ], + "binding_access": [ + "read", + "read", + "read", + "read", + "read_write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "qwen3_moe/flash_attention_decode_f32_f16_wmma.loom" + ], + "library_sources": [] + }, + "compile_dependencies": [] + }, + { + "name": "qwen3_moe_flash_attention_decode_q128_fused_f32_f16_wmma", + "symbol": "qwen3_moe_flash_attention_decode_q128_fused_f32_f16_wmma", + "source": "qwen3_moe/flash_attention_decode_q128_f32_f16_wmma.loom", + "workload_parameters": [ + { + "name": "key_value_token_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "key_value_token_count", + "type": "index" + } + ], + "bindings": [ + "query", + "key", + "value", + "mask", + "output" + ], + "binding_access": [ + "read", + "read", + "read", + "read", + "read_write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "qwen3_moe/flash_attention_decode_q128_f32_f16_wmma.loom" + ], + "library_sources": [] + }, + "compile_dependencies": [] + }, + { + "name": "qwen3_moe_flash_attention_decode_split_f32_f16_wmma", + "symbol": "qwen3_moe_flash_attention_decode_split_f32_f16_wmma", + "source": "qwen3_moe/flash_attention_decode_split_f32_f16_wmma.loom", + "workload_parameters": [ + { + "name": "key_value_token_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "key_value_token_count", + "type": "index" + } + ], + "bindings": [ + "query", + "key", + "value", + "mask", + "partial_max", + "partial_sum", + "partial_output", + "completion_counter", + "output" + ], + "binding_access": [ + "read", + "read", + "read", + "read", + "read_write", + "read_write", + "read_write", + "read_write", + "read_write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "qwen3_moe/flash_attention_decode_split_f32_f16_wmma.loom" + ], + "library_sources": [] + }, + "compile_dependencies": [] + }, + { + "name": "qwen3_moe_flash_attention_decode_split_f32_f16_wmma_next_q8", + "symbol": "qwen3_moe_flash_attention_decode_split_f32_f16_wmma_next_q8", + "source": "qwen3_moe/flash_attention_decode_split_f32_f16_wmma.loom", + "workload_parameters": [ + { + "name": "key_value_token_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "key_value_token_count", + "type": "index" + } + ], + "bindings": [ + "query", + "key", + "value", + "mask", + "partial_max", + "partial_sum", + "partial_output", + "completion_counter", + "output", + "next_q8_output" + ], + "binding_access": [ + "read", + "read", + "read", + "read", + "read_write", + "read_write", + "read_write", + "read_write", + "read_write", + "read_write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "qwen3_moe/flash_attention_decode_split_f32_f16_wmma.loom" + ], + "library_sources": [] + }, + "compile_dependencies": [] + }, + { + "name": "qwen3_moe_flash_attention_decode_split_mask_513_of_576", + "symbol": "qwen3_moe_flash_attention_decode_split_mask_513_of_576", + "source": "qwen3_moe/flash_attention_decode_split_next_q8_test.loom", + "workload_parameters": [], + "launch_parameters": [], + "bindings": [ + "mask" + ], + "binding_access": [ + "read" + ], + "target_selector": "", + "compile_recipe": { + "mode": "archive", + "link_module": "flash_attention_decode_split_f32_f16_wmma_next_q8_linked", + "primary_sources": [ + "qwen3_moe/flash_attention_decode_split_next_q8_test.loom" + ], + "library_sources": [ + "qwen3_moe/flash_attention_decode_split_f32_f16_wmma.loom", + "ggml/quantize_q8_1_x4.loom" + ] + }, + "compile_dependencies": [ + "qwen3_moe/flash_attention_decode_split_f32_f16_wmma.loom", + "ggml/quantize_q8_1_x4.loom" + ] + }, + { + "name": "qwen3_moe_flash_attention_decode_split_pack_completed_q8_test", + "symbol": "qwen3_moe_flash_attention_decode_split_pack_completed_q8_test", + "source": "qwen3_moe/flash_attention_decode_split_next_q8_test.loom", + "workload_parameters": [], + "launch_parameters": [], + "bindings": [ + "input", + "output" + ], + "binding_access": [ + "read", + "read_write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "archive", + "link_module": "flash_attention_decode_split_f32_f16_wmma_next_q8_linked", + "primary_sources": [ + "qwen3_moe/flash_attention_decode_split_next_q8_test.loom" + ], + "library_sources": [ + "qwen3_moe/flash_attention_decode_split_f32_f16_wmma.loom", + "ggml/quantize_q8_1_x4.loom" + ] + }, + "compile_dependencies": [ + "qwen3_moe/flash_attention_decode_split_f32_f16_wmma.loom", + "ggml/quantize_q8_1_x4.loom" + ] + }, + { + "name": "qwen3_moe_flash_attention_decode_split_produce_partials_f32_f16_wmma", + "symbol": "qwen3_moe_flash_attention_decode_split_produce_partials_f32_f16_wmma", + "source": "qwen3_moe/flash_attention_decode_split_f32_f16_wmma.loom", + "workload_parameters": [ + { + "name": "key_value_token_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "key_value_token_count", + "type": "index" + } + ], + "bindings": [ + "query", + "key", + "value", + "mask", + "partial_max", + "partial_sum", + "partial_output" + ], + "binding_access": [ + "read", + "read", + "read", + "read", + "read_write", + "read_write", + "read_write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "qwen3_moe/flash_attention_decode_split_f32_f16_wmma.loom" + ], + "library_sources": [] + }, + "compile_dependencies": [] + }, + { + "name": "qwen3_moe_flash_attention_decode_split_quantize_reference_4096", + "symbol": "qwen3_moe_flash_attention_decode_split_quantize_reference_4096", + "source": "qwen3_moe/flash_attention_decode_split_next_q8_test.loom", + "workload_parameters": [], + "launch_parameters": [], + "bindings": [ + "input", + "output" + ], + "binding_access": [ + "read", + "read_write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "archive", + "link_module": "flash_attention_decode_split_f32_f16_wmma_next_q8_linked", + "primary_sources": [ + "qwen3_moe/flash_attention_decode_split_next_q8_test.loom" + ], + "library_sources": [ + "qwen3_moe/flash_attention_decode_split_f32_f16_wmma.loom", + "ggml/quantize_q8_1_x4.loom" + ] + }, + "compile_dependencies": [ + "qwen3_moe/flash_attention_decode_split_f32_f16_wmma.loom", + "ggml/quantize_q8_1_x4.loom" + ] + }, + { + "name": "qwen3_moe_flash_attention_decode_split_reduce_f32", + "symbol": "qwen3_moe_flash_attention_decode_split_reduce_f32", + "source": "qwen3_moe/flash_attention_decode_split_f32_f16_wmma.loom", + "workload_parameters": [ + { + "name": "key_value_token_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "key_value_token_count", + "type": "index" + } + ], + "bindings": [ + "partial_max", + "partial_sum", + "partial_output", + "output" + ], + "binding_access": [ + "read_write", + "read_write", + "read_write", + "read_write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "qwen3_moe/flash_attention_decode_split_f32_f16_wmma.loom" + ], + "library_sources": [] + }, + "compile_dependencies": [] + }, + { + "name": "qwen3_moe_flash_attention_f32_f16_wmma", + "symbol": "qwen3_moe_flash_attention_f32_f16_wmma", + "source": "qwen3_moe/flash_attention_f32_f16_wmma.loom", + "workload_parameters": [ + { + "name": "query_token_count", + "type": "index" + }, + { + "name": "key_value_token_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "query_token_count", + "type": "index" + }, + { + "name": "key_value_token_count", + "type": "index" + } + ], + "bindings": [ + "query", + "key", + "value", + "mask", + "output" + ], + "binding_access": [ + "read", + "read", + "read", + "read", + "read_write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "qwen3_moe/flash_attention_f32_f16_wmma.loom" + ], + "library_sources": [] + }, + "compile_dependencies": [] + }, + { + "name": "qwen3_moe_flash_attention_test_extract_row", + "symbol": "qwen3_moe_flash_attention_test_extract_row", + "source": "qwen3_moe/flash_attention_f32_f16_wmma.loom", + "workload_parameters": [ + { + "name": "query_token_count", + "type": "index" + }, + { + "name": "context_count", + "type": "index" + }, + { + "name": "source_row", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "query_token_count", + "type": "index" + }, + { + "name": "context_count", + "type": "index" + }, + { + "name": "source_row", + "type": "index" + } + ], + "bindings": [ + "source_query", + "source_mask", + "target_query", + "target_mask" + ], + "binding_access": [ + "read", + "read", + "read", + "read" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "qwen3_moe/flash_attention_f32_f16_wmma.loom" + ], + "library_sources": [] + }, + "compile_dependencies": [] + }, + { + "name": "qwen3_moe_flash_attention_test_make_causal_mask", + "symbol": "qwen3_moe_flash_attention_test_make_causal_mask", + "source": "qwen3_moe/flash_attention_f32_f16_wmma.loom", + "workload_parameters": [], + "launch_parameters": [], + "bindings": [ + "mask" + ], + "binding_access": [ + "read" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "qwen3_moe/flash_attention_f32_f16_wmma.loom" + ], + "library_sources": [] + }, + "compile_dependencies": [] + }, + { + "name": "qwen3_moe_rmsnorm_f32", + "symbol": "qwen3_moe_rmsnorm_f32", + "source": "qwen3_moe/attention_prepare_quantized.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "bindings": [ + "input", + "weight", + "output" + ], + "binding_access": [ + "read", + "read", + "read_write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "archive", + "link_module": "attention_prepare_quantized_linked", + "primary_sources": [ + "qwen3_moe/attention_prepare_quantized.loom" + ], + "library_sources": [ + "qwen3_moe/model_config.loom", + "ggml/quantize_q8_1_x4.loom" + ] + }, + "compile_dependencies": [ + "qwen3_moe/model_config.loom", + "ggml/quantize_q8_1_x4.loom" + ] + }, + { + "name": "qwen3_moe_rmsnorm_f32_quantize_q8_1_x4", + "symbol": "qwen3_moe_rmsnorm_f32_quantize_q8_1_x4", + "source": "qwen3_moe/attention_prepare_quantized.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "bindings": [ + "input", + "weight", + "normalized_output", + "q8_output" + ], + "binding_access": [ + "read", + "read", + "read_write", + "read_write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "archive", + "link_module": "attention_prepare_quantized_linked", + "primary_sources": [ + "qwen3_moe/attention_prepare_quantized.loom" + ], + "library_sources": [ + "qwen3_moe/model_config.loom", + "ggml/quantize_q8_1_x4.loom" + ] + }, + "compile_dependencies": [ + "qwen3_moe/model_config.loom", + "ggml/quantize_q8_1_x4.loom" + ] + }, + { + "name": "qwen3_moe_rmsnorm_quantize_q8_1_x4_wave64_production_check", + "symbol": "qwen3_moe_rmsnorm_quantize_q8_1_x4_wave64_production_check", + "source": "qwen3_moe/routed_down_q6k.loom", + "workload_parameters": [], + "launch_parameters": [], + "bindings": [ + "input", + "norm_weight", + "q8_output" + ], + "binding_access": [ + "read", + "read", + "read_write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "archive", + "link_module": "routed_down_q6k_linked", + "primary_sources": [ + "qwen3_moe/routed_down_q6k.loom" + ], + "library_sources": [ + "ggml/linear_q6k_f32.loom", + "ggml/linear_q6k_q8_1_x4.loom", + "ggml/quantize_q8_1_x4.loom", + "qwen3_moe/attention_prepare_quantized.loom", + "qwen3_moe/model_config.loom", + "qwen3_moe/routed_down_next_q8.loom" + ] + }, + "compile_dependencies": [ + "ggml/linear_q6k_f32.loom", + "ggml/linear_q6k_q8_1_x4.loom", + "ggml/quantize_q8_1_x4.loom", + "qwen3_moe/attention_prepare_quantized.loom", + "qwen3_moe/model_config.loom", + "qwen3_moe/routed_down_next_q8.loom" + ] + }, + { + "name": "qwen3_moe_routed_down_q4k_f16_wmma_grouped", + "symbol": "qwen3_moe_routed_down_q4k_f16_wmma_grouped", + "source": "qwen3_moe/routed_down_quantized_f16_wmma.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "bindings": [ + "input", + "expert_table", + "weight", + "output" + ], + "binding_access": [ + "read", + "read", + "read", + "read_write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "archive", + "link_module": "routed_down_quantized_f16_wmma_linked", + "primary_sources": [ + "qwen3_moe/routed_down_quantized_f16_wmma.loom" + ], + "library_sources": [ + "qwen3_moe/routed_linear_q4k_f16_wmma.loom", + "qwen3_moe/routed_down_q4k.loom", + "qwen3_moe/routed_down_q6k.loom", + "qwen3_moe/routed_down_next_q8.loom", + "qwen3_moe/attention_prepare_quantized.loom", + "qwen3_moe/model_config.loom", + "ggml/linear_q6k_q8_1_x4.loom", + "qwen3_moe/routed_gate_up_swiglu_q4k.loom", + "ggml/quantize_q8_1_x4.loom" + ] + }, + "compile_dependencies": [ + "qwen3_moe/routed_linear_q4k_f16_wmma.loom", + "qwen3_moe/routed_down_q4k.loom", + "qwen3_moe/routed_down_q6k.loom", + "qwen3_moe/routed_down_next_q8.loom", + "qwen3_moe/attention_prepare_quantized.loom", + "qwen3_moe/model_config.loom", + "ggml/linear_q6k_q8_1_x4.loom", + "qwen3_moe/routed_gate_up_swiglu_q4k.loom", + "ggml/quantize_q8_1_x4.loom" + ] + }, + { + "name": "qwen3_moe_routed_down_q4k_q8_1_x4", + "symbol": "qwen3_moe_routed_down_q4k_q8_1_x4", + "source": "qwen3_moe/routed_down_q4k.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + }, + { + "name": "input_size", + "type": "index" + }, + { + "name": "route_count", + "type": "index" + }, + { + "name": "route_id_stride", + "type": "index" + }, + { + "name": "expert_count", + "type": "index" + }, + { + "name": "output_size", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + }, + { + "name": "input_size", + "type": "index" + }, + { + "name": "route_count", + "type": "index" + }, + { + "name": "route_id_stride", + "type": "index" + }, + { + "name": "expert_count", + "type": "index" + }, + { + "name": "output_size", + "type": "index" + } + ], + "bindings": [ + "q8_input", + "route_ids", + "route_weights", + "weight", + "output" + ], + "binding_access": [ + "read", + "read", + "read", + "read", + "read_write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "archive", + "link_module": "routed_down_q4k_linked", + "primary_sources": [ + "qwen3_moe/routed_down_q4k.loom" + ], + "library_sources": [ + "qwen3_moe/attention_prepare_quantized.loom", + "qwen3_moe/model_config.loom", + "qwen3_moe/routed_down_next_q8.loom", + "qwen3_moe/routed_gate_up_swiglu_q4k.loom", + "ggml/quantize_q8_1_x4.loom" + ] + }, + "compile_dependencies": [ + "qwen3_moe/attention_prepare_quantized.loom", + "qwen3_moe/model_config.loom", + "qwen3_moe/routed_down_next_q8.loom", + "qwen3_moe/routed_gate_up_swiglu_q4k.loom", + "ggml/quantize_q8_1_x4.loom" + ] + }, + { + "name": "qwen3_moe_routed_down_q4k_q8_1_x4_next_q8", + "symbol": "qwen3_moe_routed_down_q4k_q8_1_x4_next_q8", + "source": "qwen3_moe/routed_down_q4k.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + }, + { + "name": "input_size", + "type": "index" + }, + { + "name": "route_count", + "type": "index" + }, + { + "name": "route_id_stride", + "type": "index" + }, + { + "name": "expert_count", + "type": "index" + }, + { + "name": "output_size", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + }, + { + "name": "input_size", + "type": "index" + }, + { + "name": "route_count", + "type": "index" + }, + { + "name": "route_id_stride", + "type": "index" + }, + { + "name": "expert_count", + "type": "index" + }, + { + "name": "output_size", + "type": "index" + } + ], + "bindings": [ + "q8_input", + "route_ids", + "route_weights", + "weight", + "output", + "norm_weight", + "completion_counter", + "next_q8_output" + ], + "binding_access": [ + "read", + "read", + "read", + "read", + "read_write", + "read", + "read_write", + "read_write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "archive", + "link_module": "routed_down_q4k_linked", + "primary_sources": [ + "qwen3_moe/routed_down_q4k.loom" + ], + "library_sources": [ + "qwen3_moe/attention_prepare_quantized.loom", + "qwen3_moe/model_config.loom", + "qwen3_moe/routed_down_next_q8.loom", + "qwen3_moe/routed_gate_up_swiglu_q4k.loom", + "ggml/quantize_q8_1_x4.loom" + ] + }, + "compile_dependencies": [ + "qwen3_moe/attention_prepare_quantized.loom", + "qwen3_moe/model_config.loom", + "qwen3_moe/routed_down_next_q8.loom", + "qwen3_moe/routed_gate_up_swiglu_q4k.loom", + "ggml/quantize_q8_1_x4.loom" + ] + }, + { + "name": "qwen3_moe_routed_down_q6k_f16_wmma_grouped", + "symbol": "qwen3_moe_routed_down_q6k_f16_wmma_grouped", + "source": "qwen3_moe/routed_down_quantized_f16_wmma.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "bindings": [ + "input", + "expert_table", + "weight", + "output" + ], + "binding_access": [ + "read", + "read", + "read", + "read_write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "archive", + "link_module": "routed_down_quantized_f16_wmma_linked", + "primary_sources": [ + "qwen3_moe/routed_down_quantized_f16_wmma.loom" + ], + "library_sources": [ + "qwen3_moe/routed_linear_q4k_f16_wmma.loom", + "qwen3_moe/routed_down_q4k.loom", + "qwen3_moe/routed_down_q6k.loom", + "qwen3_moe/routed_down_next_q8.loom", + "qwen3_moe/attention_prepare_quantized.loom", + "qwen3_moe/model_config.loom", + "ggml/linear_q6k_q8_1_x4.loom", + "qwen3_moe/routed_gate_up_swiglu_q4k.loom", + "ggml/quantize_q8_1_x4.loom" + ] + }, + "compile_dependencies": [ + "qwen3_moe/routed_linear_q4k_f16_wmma.loom", + "qwen3_moe/routed_down_q4k.loom", + "qwen3_moe/routed_down_q6k.loom", + "qwen3_moe/routed_down_next_q8.loom", + "qwen3_moe/attention_prepare_quantized.loom", + "qwen3_moe/model_config.loom", + "ggml/linear_q6k_q8_1_x4.loom", + "qwen3_moe/routed_gate_up_swiglu_q4k.loom", + "ggml/quantize_q8_1_x4.loom" + ] + }, + { + "name": "qwen3_moe_routed_down_q6k_f32_wave64", + "symbol": "qwen3_moe_routed_down_q6k_f32_wave64", + "source": "qwen3_moe/routed_down_q6k.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + }, + { + "name": "input_size", + "type": "index" + }, + { + "name": "route_count", + "type": "index" + }, + { + "name": "route_id_stride", + "type": "index" + }, + { + "name": "expert_count", + "type": "index" + }, + { + "name": "output_size", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + }, + { + "name": "input_size", + "type": "index" + }, + { + "name": "route_count", + "type": "index" + }, + { + "name": "route_id_stride", + "type": "index" + }, + { + "name": "expert_count", + "type": "index" + }, + { + "name": "output_size", + "type": "index" + } + ], + "bindings": [ + "input", + "route_ids", + "route_weights", + "weight", + "output" + ], + "binding_access": [ + "read", + "read", + "read", + "read", + "read_write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "archive", + "link_module": "routed_down_q6k_linked", + "primary_sources": [ + "qwen3_moe/routed_down_q6k.loom" + ], + "library_sources": [ + "ggml/linear_q6k_f32.loom", + "ggml/linear_q6k_q8_1_x4.loom", + "ggml/quantize_q8_1_x4.loom", + "qwen3_moe/attention_prepare_quantized.loom", + "qwen3_moe/model_config.loom", + "qwen3_moe/routed_down_next_q8.loom" + ] + }, + "compile_dependencies": [ + "ggml/linear_q6k_f32.loom", + "ggml/linear_q6k_q8_1_x4.loom", + "ggml/quantize_q8_1_x4.loom", + "qwen3_moe/attention_prepare_quantized.loom", + "qwen3_moe/model_config.loom", + "qwen3_moe/routed_down_next_q8.loom" + ] + }, + { + "name": "qwen3_moe_routed_down_q6k_f32_wave64_next_q8", + "symbol": "qwen3_moe_routed_down_q6k_f32_wave64_next_q8", + "source": "qwen3_moe/routed_down_q6k.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + }, + { + "name": "input_size", + "type": "index" + }, + { + "name": "route_count", + "type": "index" + }, + { + "name": "route_id_stride", + "type": "index" + }, + { + "name": "expert_count", + "type": "index" + }, + { + "name": "output_size", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + }, + { + "name": "input_size", + "type": "index" + }, + { + "name": "route_count", + "type": "index" + }, + { + "name": "route_id_stride", + "type": "index" + }, + { + "name": "expert_count", + "type": "index" + }, + { + "name": "output_size", + "type": "index" + } + ], + "bindings": [ + "input", + "route_ids", + "route_weights", + "weight", + "output", + "norm_weight", + "completion_counter", + "next_q8_output" + ], + "binding_access": [ + "read", + "read", + "read", + "read", + "read_write", + "read", + "read_write", + "read_write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "archive", + "link_module": "routed_down_q6k_linked", + "primary_sources": [ + "qwen3_moe/routed_down_q6k.loom" + ], + "library_sources": [ + "ggml/linear_q6k_f32.loom", + "ggml/linear_q6k_q8_1_x4.loom", + "ggml/quantize_q8_1_x4.loom", + "qwen3_moe/attention_prepare_quantized.loom", + "qwen3_moe/model_config.loom", + "qwen3_moe/routed_down_next_q8.loom" + ] + }, + "compile_dependencies": [ + "ggml/linear_q6k_f32.loom", + "ggml/linear_q6k_q8_1_x4.loom", + "ggml/quantize_q8_1_x4.loom", + "qwen3_moe/attention_prepare_quantized.loom", + "qwen3_moe/model_config.loom", + "qwen3_moe/routed_down_next_q8.loom" + ] + }, + { + "name": "qwen3_moe_routed_down_q6k_q8_1_x4", + "symbol": "qwen3_moe_routed_down_q6k_q8_1_x4", + "source": "qwen3_moe/routed_down_q6k.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + }, + { + "name": "input_size", + "type": "index" + }, + { + "name": "route_count", + "type": "index" + }, + { + "name": "route_id_stride", + "type": "index" + }, + { + "name": "expert_count", + "type": "index" + }, + { + "name": "output_size", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + }, + { + "name": "input_size", + "type": "index" + }, + { + "name": "route_count", + "type": "index" + }, + { + "name": "route_id_stride", + "type": "index" + }, + { + "name": "expert_count", + "type": "index" + }, + { + "name": "output_size", + "type": "index" + } + ], + "bindings": [ + "q8_input", + "route_ids", + "route_weights", + "weight", + "output" + ], + "binding_access": [ + "read", + "read", + "read", + "read", + "read_write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "archive", + "link_module": "routed_down_q6k_linked", + "primary_sources": [ + "qwen3_moe/routed_down_q6k.loom" + ], + "library_sources": [ + "ggml/linear_q6k_f32.loom", + "ggml/linear_q6k_q8_1_x4.loom", + "ggml/quantize_q8_1_x4.loom", + "qwen3_moe/attention_prepare_quantized.loom", + "qwen3_moe/model_config.loom", + "qwen3_moe/routed_down_next_q8.loom" + ] + }, + "compile_dependencies": [ + "ggml/linear_q6k_f32.loom", + "ggml/linear_q6k_q8_1_x4.loom", + "ggml/quantize_q8_1_x4.loom", + "qwen3_moe/attention_prepare_quantized.loom", + "qwen3_moe/model_config.loom", + "qwen3_moe/routed_down_next_q8.loom" + ] + }, + { + "name": "qwen3_moe_routed_down_q6k_q8_1_x4_next_q8", + "symbol": "qwen3_moe_routed_down_q6k_q8_1_x4_next_q8", + "source": "qwen3_moe/routed_down_q6k.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + }, + { + "name": "input_size", + "type": "index" + }, + { + "name": "route_count", + "type": "index" + }, + { + "name": "route_id_stride", + "type": "index" + }, + { + "name": "expert_count", + "type": "index" + }, + { + "name": "output_size", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + }, + { + "name": "input_size", + "type": "index" + }, + { + "name": "route_count", + "type": "index" + }, + { + "name": "route_id_stride", + "type": "index" + }, + { + "name": "expert_count", + "type": "index" + }, + { + "name": "output_size", + "type": "index" + } + ], + "bindings": [ + "q8_input", + "route_ids", + "route_weights", + "weight", + "output", + "norm_weight", + "completion_counter", + "next_q8_output" + ], + "binding_access": [ + "read", + "read", + "read", + "read", + "read_write", + "read", + "read_write", + "read_write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "archive", + "link_module": "routed_down_q6k_linked", + "primary_sources": [ + "qwen3_moe/routed_down_q6k.loom" + ], + "library_sources": [ + "ggml/linear_q6k_f32.loom", + "ggml/linear_q6k_q8_1_x4.loom", + "ggml/quantize_q8_1_x4.loom", + "qwen3_moe/attention_prepare_quantized.loom", + "qwen3_moe/model_config.loom", + "qwen3_moe/routed_down_next_q8.loom" + ] + }, + "compile_dependencies": [ + "ggml/linear_q6k_f32.loom", + "ggml/linear_q6k_q8_1_x4.loom", + "ggml/quantize_q8_1_x4.loom", + "qwen3_moe/attention_prepare_quantized.loom", + "qwen3_moe/model_config.loom", + "qwen3_moe/routed_down_next_q8.loom" + ] + }, + { + "name": "qwen3_moe_routed_down_weighted_reduce_f16_f32", + "symbol": "qwen3_moe_routed_down_weighted_reduce_f16_f32", + "source": "qwen3_moe/routed_down_quantized_f16_wmma.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "bindings": [ + "route_weights", + "routed_output", + "output" + ], + "binding_access": [ + "read", + "read", + "read_write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "archive", + "link_module": "routed_down_quantized_f16_wmma_linked", + "primary_sources": [ + "qwen3_moe/routed_down_quantized_f16_wmma.loom" + ], + "library_sources": [ + "qwen3_moe/routed_linear_q4k_f16_wmma.loom", + "qwen3_moe/routed_down_q4k.loom", + "qwen3_moe/routed_down_q6k.loom", + "qwen3_moe/routed_down_next_q8.loom", + "qwen3_moe/attention_prepare_quantized.loom", + "qwen3_moe/model_config.loom", + "ggml/linear_q6k_q8_1_x4.loom", + "qwen3_moe/routed_gate_up_swiglu_q4k.loom", + "ggml/quantize_q8_1_x4.loom" + ] + }, + "compile_dependencies": [ + "qwen3_moe/routed_linear_q4k_f16_wmma.loom", + "qwen3_moe/routed_down_q4k.loom", + "qwen3_moe/routed_down_q6k.loom", + "qwen3_moe/routed_down_next_q8.loom", + "qwen3_moe/attention_prepare_quantized.loom", + "qwen3_moe/model_config.loom", + "ggml/linear_q6k_q8_1_x4.loom", + "qwen3_moe/routed_gate_up_swiglu_q4k.loom", + "ggml/quantize_q8_1_x4.loom" + ] + }, + { + "name": "qwen3_moe_routed_down_weighted_reduce_next_rmsnorm_f32", + "symbol": "qwen3_moe_routed_down_weighted_reduce_next_rmsnorm_f32", + "source": "qwen3_moe/routed_down_weighted_reduce_next_rmsnorm_f32.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "bindings": [ + "route_weights", + "routed_output", + "hidden_state", + "next_norm_weight", + "next_projection_input" + ], + "binding_access": [ + "read", + "read", + "read_write", + "read", + "write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "archive", + "link_module": "routed_down_weighted_reduce_next_rmsnorm_f32_linked", + "primary_sources": [ + "qwen3_moe/routed_down_weighted_reduce_next_rmsnorm_f32.loom" + ], + "library_sources": [ + "qwen3_moe/routed_down_quantized_f16_wmma.loom", + "qwen3_moe/routed_linear_q4k_f16_wmma.loom", + "qwen3_moe/routed_down_q4k.loom", + "qwen3_moe/routed_down_q6k.loom", + "qwen3_moe/routed_down_next_q8.loom", + "qwen3_moe/attention_prepare_quantized.loom", + "qwen3_moe/model_config.loom", + "ggml/linear_q6k_q8_1_x4.loom", + "qwen3_moe/routed_gate_up_swiglu_q4k.loom", + "ggml/quantize_q8_1_x4.loom" + ] + }, + "compile_dependencies": [ + "qwen3_moe/routed_down_quantized_f16_wmma.loom", + "qwen3_moe/routed_linear_q4k_f16_wmma.loom", + "qwen3_moe/routed_down_q4k.loom", + "qwen3_moe/routed_down_q6k.loom", + "qwen3_moe/routed_down_next_q8.loom", + "qwen3_moe/attention_prepare_quantized.loom", + "qwen3_moe/model_config.loom", + "ggml/linear_q6k_q8_1_x4.loom", + "qwen3_moe/routed_gate_up_swiglu_q4k.loom", + "ggml/quantize_q8_1_x4.loom" + ] + }, + { + "name": "qwen3_moe_routed_down_weighted_reduce_next_rmsnorm_q8_1_x4", + "symbol": "qwen3_moe_routed_down_weighted_reduce_next_rmsnorm_q8_1_x4", + "source": "qwen3_moe/routed_down_weighted_reduce_next_rmsnorm_q8_1_x4.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "bindings": [ + "route_weights", + "routed_output", + "hidden_state", + "next_norm_weight", + "next_projection_input" + ], + "binding_access": [ + "read", + "read", + "read", + "read", + "read" + ], + "target_selector": "", + "compile_recipe": { + "mode": "archive", + "link_module": "routed_down_weighted_reduce_next_rmsnorm_q8_1_x4_linked", + "primary_sources": [ + "qwen3_moe/routed_down_weighted_reduce_next_rmsnorm_q8_1_x4.loom" + ], + "library_sources": [ + "qwen3_moe/routed_down_quantized_f16_wmma.loom", + "qwen3_moe/routed_linear_q4k_f16_wmma.loom", + "qwen3_moe/routed_down_q4k.loom", + "qwen3_moe/routed_down_q6k.loom", + "qwen3_moe/routed_down_next_q8.loom", + "qwen3_moe/attention_prepare_quantized.loom", + "qwen3_moe/model_config.loom", + "ggml/linear_q6k_q8_1_x4.loom", + "qwen3_moe/routed_gate_up_swiglu_q4k.loom", + "ggml/quantize_q8_1_x4.loom" + ] + }, + "compile_dependencies": [ + "qwen3_moe/routed_down_quantized_f16_wmma.loom", + "qwen3_moe/routed_linear_q4k_f16_wmma.loom", + "qwen3_moe/routed_down_q4k.loom", + "qwen3_moe/routed_down_q6k.loom", + "qwen3_moe/routed_down_next_q8.loom", + "qwen3_moe/attention_prepare_quantized.loom", + "qwen3_moe/model_config.loom", + "ggml/linear_q6k_q8_1_x4.loom", + "qwen3_moe/routed_gate_up_swiglu_q4k.loom", + "ggml/quantize_q8_1_x4.loom" + ] + }, + { + "name": "qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma", + "symbol": "qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma", + "source": "qwen3_moe/routed_linear_q4k_f16_wmma.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "bindings": [ + "input", + "expert_table", + "partition_table", + "gate_weight", + "up_weight", + "output" + ], + "binding_access": [ + "read", + "read", + "read", + "read", + "read", + "read_write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "archive", + "link_module": "routed_gate_up_swiglu_q4k_f16_wmma_linked", + "primary_sources": [ + "qwen3_moe/routed_linear_q4k_f16_wmma.loom" + ], + "library_sources": [ + "qwen3_moe/routed_gate_up_swiglu_q4k.loom", + "ggml/quantize_q8_1_x4.loom" + ] + }, + "compile_dependencies": [ + "qwen3_moe/routed_gate_up_swiglu_q4k.loom", + "ggml/quantize_q8_1_x4.loom" + ] + }, + { + "name": "qwen3_moe_routed_gate_up_swiglu_q4k_q8", + "symbol": "qwen3_moe_routed_gate_up_swiglu_q4k_q8", + "source": "qwen3_moe/routed_gate_up_swiglu_q4k.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + }, + { + "name": "route_count", + "type": "index" + }, + { + "name": "route_stride", + "type": "index" + }, + { + "name": "expert_count", + "type": "index" + }, + { + "name": "output_size", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + }, + { + "name": "route_count", + "type": "index" + }, + { + "name": "route_stride", + "type": "index" + }, + { + "name": "expert_count", + "type": "index" + }, + { + "name": "output_size", + "type": "index" + } + ], + "bindings": [ + "q8_input", + "route_ids", + "gate_weight", + "up_weight", + "output" + ], + "binding_access": [ + "read", + "read", + "read", + "read", + "read_write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "archive", + "link_module": "routed_gate_up_swiglu_q4k_linked", + "primary_sources": [ + "qwen3_moe/routed_gate_up_swiglu_q4k.loom" + ], + "library_sources": [ + "ggml/quantize_q8_1_x4.loom" + ] + }, + "compile_dependencies": [ + "ggml/quantize_q8_1_x4.loom" + ] + }, + { + "name": "qwen3_moe_routed_gate_up_swiglu_q4k_q8_1_x4_next_q8", + "symbol": "qwen3_moe_routed_gate_up_swiglu_q4k_q8_1_x4_next_q8", + "source": "qwen3_moe/routed_gate_up_swiglu_q4k.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + }, + { + "name": "route_count", + "type": "index" + }, + { + "name": "route_stride", + "type": "index" + }, + { + "name": "expert_count", + "type": "index" + }, + { + "name": "output_size", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + }, + { + "name": "route_count", + "type": "index" + }, + { + "name": "route_stride", + "type": "index" + }, + { + "name": "expert_count", + "type": "index" + }, + { + "name": "output_size", + "type": "index" + } + ], + "bindings": [ + "q8_input", + "route_ids", + "gate_weight", + "up_weight", + "output", + "completion_counters", + "next_q8_output" + ], + "binding_access": [ + "read", + "read", + "read", + "read", + "read_write", + "read_write", + "read_write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "archive", + "link_module": "routed_gate_up_swiglu_q4k_linked", + "primary_sources": [ + "qwen3_moe/routed_gate_up_swiglu_q4k.loom" + ], + "library_sources": [ + "ggml/quantize_q8_1_x4.loom" + ] + }, + "compile_dependencies": [ + "ggml/quantize_q8_1_x4.loom" + ] + }, + { + "name": "qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped", + "symbol": "qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped", + "source": "qwen3_moe/routed_gate_up_swiglu_q4k.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + }, + { + "name": "route_count", + "type": "index" + }, + { + "name": "expert_count", + "type": "index" + }, + { + "name": "output_size", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + }, + { + "name": "route_count", + "type": "index" + }, + { + "name": "expert_count", + "type": "index" + }, + { + "name": "output_size", + "type": "index" + } + ], + "bindings": [ + "q8_input", + "expert_table", + "gate_weight", + "up_weight", + "output" + ], + "binding_access": [ + "read", + "read", + "read", + "read", + "read_write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "archive", + "link_module": "routed_gate_up_swiglu_q4k_linked", + "primary_sources": [ + "qwen3_moe/routed_gate_up_swiglu_q4k.loom" + ], + "library_sources": [ + "ggml/quantize_q8_1_x4.loom" + ] + }, + "compile_dependencies": [ + "ggml/quantize_q8_1_x4.loom" + ] + }, + { + "name": "qwen3_moe_routed_linear_q4k_f16_wmma", + "symbol": "qwen3_moe_routed_linear_q4k_f16_wmma", + "source": "qwen3_moe/routed_linear_q4k_f16_wmma.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "bindings": [ + "input", + "expert_table", + "weight", + "output" + ], + "binding_access": [ + "read", + "read", + "read", + "read_write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "archive", + "link_module": "routed_gate_up_swiglu_q4k_f16_wmma_linked", + "primary_sources": [ + "qwen3_moe/routed_linear_q4k_f16_wmma.loom" + ], + "library_sources": [ + "qwen3_moe/routed_gate_up_swiglu_q4k.loom", + "ggml/quantize_q8_1_x4.loom" + ] + }, + "compile_dependencies": [ + "qwen3_moe/routed_gate_up_swiglu_q4k.loom", + "ggml/quantize_q8_1_x4.loom" + ] + }, + { + "name": "qwen3_moe_routed_swiglu_f16", + "symbol": "qwen3_moe_routed_swiglu_f16", + "source": "qwen3_moe/routed_linear_q4k_f16_wmma.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "bindings": [ + "gate", + "up", + "output" + ], + "binding_access": [ + "read", + "read", + "read_write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "archive", + "link_module": "routed_gate_up_swiglu_q4k_f16_wmma_linked", + "primary_sources": [ + "qwen3_moe/routed_linear_q4k_f16_wmma.loom" + ], + "library_sources": [ + "qwen3_moe/routed_gate_up_swiglu_q4k.loom", + "ggml/quantize_q8_1_x4.loom" + ] + }, + "compile_dependencies": [ + "qwen3_moe/routed_gate_up_swiglu_q4k.loom", + "ggml/quantize_q8_1_x4.loom" + ] + }, + { + "name": "qwen3_moe_routed_swiglu_f32", + "symbol": "qwen3_moe_routed_swiglu_f32", + "source": "qwen3_moe/routed_linear_q4k_f16_wmma.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "bindings": [ + "gate", + "up", + "output" + ], + "binding_access": [ + "read", + "read", + "read_write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "archive", + "link_module": "routed_gate_up_swiglu_q4k_f16_wmma_linked", + "primary_sources": [ + "qwen3_moe/routed_linear_q4k_f16_wmma.loom" + ], + "library_sources": [ + "qwen3_moe/routed_gate_up_swiglu_q4k.loom", + "ggml/quantize_q8_1_x4.loom" + ] + }, + "compile_dependencies": [ + "qwen3_moe/routed_gate_up_swiglu_q4k.loom", + "ggml/quantize_q8_1_x4.loom" + ] + }, + { + "name": "qwen3_moe_router_projection_f32_four_row_wave32", + "symbol": "qwen3_moe_router_projection_f32_four_row_wave32", + "source": "qwen3_moe/router_projection_f32.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "bindings": [ + "input", + "weight", + "output" + ], + "binding_access": [ + "read", + "read", + "read_write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "archive", + "link_module": "router_projection_f32_linked", + "primary_sources": [ + "qwen3_moe/router_projection_f32.loom" + ], + "library_sources": [ + "qwen3_moe/model_config.loom" + ] + }, + "compile_dependencies": [ + "qwen3_moe/model_config.loom" + ] + }, + { + "name": "qwen3_moe_router_projection_f32_one_row_wave64", + "symbol": "qwen3_moe_router_projection_f32_one_row_wave64", + "source": "qwen3_moe/router_projection_f32.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "bindings": [ + "input", + "weight", + "output" + ], + "binding_access": [ + "read", + "read", + "read_write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "archive", + "link_module": "router_projection_f32_linked", + "primary_sources": [ + "qwen3_moe/router_projection_f32.loom" + ], + "library_sources": [ + "qwen3_moe/model_config.loom" + ] + }, + "compile_dependencies": [ + "qwen3_moe/model_config.loom" + ] + }, + { + "name": "qwen3_moe_router_projection_f32_reference", + "symbol": "qwen3_moe_router_projection_f32_reference", + "source": "qwen3_moe/router_projection_f32.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "bindings": [ + "input", + "weight", + "output" + ], + "binding_access": [ + "read", + "read", + "read_write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "archive", + "link_module": "router_projection_f32_linked", + "primary_sources": [ + "qwen3_moe/router_projection_f32.loom" + ], + "library_sources": [ + "qwen3_moe/model_config.loom" + ] + }, + "compile_dependencies": [ + "qwen3_moe/model_config.loom" + ] + }, + { + "name": "qwen3_moe_router_projection_top8_fused_decode_f32", + "symbol": "qwen3_moe_router_projection_top8_fused_decode_f32", + "source": "qwen3_moe/router_projection_top8_fused_f32.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + }, + { + "name": "route_id_stride", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + }, + { + "name": "route_id_stride", + "type": "index" + } + ], + "bindings": [ + "input", + "weight", + "logits", + "completion_counter", + "route_ids", + "route_weights" + ], + "binding_access": [ + "read", + "read", + "read_write", + "read_write", + "read_write", + "read_write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "archive", + "link_module": "router_projection_top8_fused_f32_linked", + "primary_sources": [ + "qwen3_moe/router_projection_top8_fused_f32.loom" + ], + "library_sources": [ + "qwen3_moe/model_config.loom", + "qwen3_moe/router_projection_f32.loom", + "qwen3_moe/router_top8_f32.loom" + ] + }, + "compile_dependencies": [ + "qwen3_moe/model_config.loom", + "qwen3_moe/router_projection_f32.loom", + "qwen3_moe/router_top8_f32.loom" + ] + }, + { + "name": "qwen3_moe_router_top8_f32", + "symbol": "qwen3_moe_router_top8_f32", + "source": "qwen3_moe/router_top8_f32.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + }, + { + "name": "route_id_stride", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + }, + { + "name": "route_id_stride", + "type": "index" + } + ], + "bindings": [ + "logits", + "route_ids", + "route_weights" + ], + "binding_access": [ + "read", + "write", + "write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "archive", + "link_module": "router_top8_f32_linked", + "primary_sources": [ + "qwen3_moe/router_top8_f32.loom" + ], + "library_sources": [ + "qwen3_moe/model_config.loom" + ] + }, + "compile_dependencies": [ + "qwen3_moe/model_config.loom" + ] + }, + { + "name": "qwen3_moe_router_top8_wide_stride_reference", + "symbol": "qwen3_moe_router_top8_wide_stride_reference", + "source": "qwen3_moe/router_top8_f32.loom", + "workload_parameters": [], + "launch_parameters": [], + "bindings": [ + "route_ids" + ], + "binding_access": [ + "read" + ], + "target_selector": "", + "compile_recipe": { + "mode": "archive", + "link_module": "router_top8_f32_linked", + "primary_sources": [ + "qwen3_moe/router_top8_f32.loom" + ], + "library_sources": [ + "qwen3_moe/model_config.loom" + ] + }, + "compile_dependencies": [ + "qwen3_moe/model_config.loom" + ] + }, + { + "name": "qwen_attention_context_base_capture", + "symbol": "qwen_attention_context_base_capture", + "source": "../qwen_owned/attention_state_initialize.loom", + "workload_parameters": [], + "launch_parameters": [], + "bindings": [ + "positions", + "control" + ], + "binding_access": [ + "read", + "write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "../qwen_owned/attention_state_initialize.loom" + ], + "library_sources": [] + }, + "compile_dependencies": [] + }, + { + "name": "qwen_attention_decode_state_initialize", + "symbol": "qwen_attention_decode_state_initialize", + "source": "../qwen_owned/attention_state_initialize.loom", + "workload_parameters": [ + { + "name": "context_capacity", + "type": "index" + }, + { + "name": "completion_counter_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "context_capacity", + "type": "index" + }, + { + "name": "completion_counter_count", + "type": "index" + } + ], + "bindings": [ + "positions", + "key_cache_indices", + "value_cache_indices", + "attention_mask", + "completion_counters" + ], + "binding_access": [ + "read", + "write", + "write", + "write", + "write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "../qwen_owned/attention_state_initialize.loom" + ], + "library_sources": [] + }, + "compile_dependencies": [] + }, + { + "name": "qwen_attention_metadata", + "symbol": "qwen_attention_metadata", + "source": "../qwen_owned/attention_metadata.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + }, + { + "name": "context_capacity", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + }, + { + "name": "context_capacity", + "type": "index" + } + ], + "bindings": [ + "control", + "positions", + "key_cache_indices", + "value_cache_indices", + "attention_mask" + ], + "binding_access": [ + "read", + "read_write", + "read_write", + "read_write", + "read_write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "../qwen_owned/attention_metadata.loom" + ], + "library_sources": [] + }, + "compile_dependencies": [] + }, + { + "name": "qwen_attention_metadata_bringup_workaround", + "symbol": "qwen_attention_metadata_bringup_workaround", + "source": "../qwen_owned/attention_metadata_bringup_workaround.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + }, + { + "name": "context_capacity", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + }, + { + "name": "context_capacity", + "type": "index" + } + ], + "bindings": [ + "control", + "positions", + "key_cache_indices", + "value_cache_indices", + "attention_mask" + ], + "binding_access": [ + "read", + "read_write", + "read_write", + "read_write", + "read_write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "../qwen_owned/attention_metadata_bringup_workaround.loom" + ], + "library_sources": [] + }, + "compile_dependencies": [] + }, + { + "name": "qwen_decode_attention_metadata", + "symbol": "qwen_decode_attention_metadata", + "source": "../qwen_owned/attention_metadata.loom", + "workload_parameters": [], + "launch_parameters": [], + "bindings": [ + "control", + "positions", + "key_cache_indices", + "value_cache_indices" + ], + "binding_access": [ + "read", + "read_write", + "read_write", + "read_write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "../qwen_owned/attention_metadata.loom" + ], + "library_sources": [] + }, + "compile_dependencies": [] + }, + { + "name": "qwen_token_embedding_q4k", + "symbol": "qwen_token_embedding_q4k", + "source": "../qwen_owned/token_embedding_q4k.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + }, + { + "name": "vocabulary_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + }, + { + "name": "vocabulary_count", + "type": "index" + } + ], + "bindings": [ + "token_ids", + "weight", + "output" + ], + "binding_access": [ + "read", + "read", + "read_write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "../qwen_owned/token_embedding_q4k.loom" + ], + "library_sources": [] + }, + "compile_dependencies": [] + }, + { + "name": "qwen_token_embedding_q4k_bringup_workaround", + "symbol": "qwen_token_embedding_q4k_bringup_workaround", + "source": "../qwen_owned/token_embedding_bringup_workaround.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + }, + { + "name": "vocabulary_count", + "type": "index" + }, + { + "name": "hidden_size", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + }, + { + "name": "vocabulary_count", + "type": "index" + }, + { + "name": "hidden_size", + "type": "index" + } + ], + "bindings": [ + "token_ids", + "weight", + "output" + ], + "binding_access": [ + "read", + "read", + "read_write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "../qwen_owned/token_embedding_bringup_workaround.loom" + ], + "library_sources": [] + }, + "compile_dependencies": [] + }, + { + "name": "ggml_get_rows_f32", + "symbol": "ggml_get_rows_f32", + "source": "../hrx_owned/get_rows_f32.loom", + "workload_parameters": [ + { + "name": "source_row_count", + "type": "index" + }, + { + "name": "output_row_count", + "type": "index" + }, + { + "name": "width", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "source_row_count", + "type": "index" + }, + { + "name": "output_row_count", + "type": "index" + }, + { + "name": "width", + "type": "index" + }, + { + "name": "source_format", + "type": "i32" + } + ], + "bindings": [ + "source", + "output_ids", + "output" + ], + "binding_access": [ + "read", + "read", + "read_write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "../hrx_owned/get_rows_f32.loom" + ], + "library_sources": [] + }, + "compile_dependencies": [] + }, + { + "name": "ggml_dequant_iq3xxs_f32", + "symbol": "ggml_dequant_iq3xxs_f32", + "source": "../hrx_owned/dequant_iq3xxs_f32.loom", + "workload_parameters": [ + { + "name": "block_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "block_count", + "type": "index" + } + ], + "bindings": [ + "weight", + "output" + ], + "binding_access": [ + "read", + "write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "../hrx_owned/dequant_iq3xxs_f32.loom" + ], + "library_sources": [] + }, + "compile_dependencies": [] + }, + { + "name": "ggml_mul_mat_vec_iq3xxs_f32", + "symbol": "ggml_mul_mat_vec_iq3xxs_f32", + "source": "../hrx_owned/mul_mat_vec_iq3xxs_f32.loom", + "workload_parameters": [ + { + "name": "input_size", + "type": "index" + }, + { + "name": "output_size", + "type": "index" + }, + { + "name": "token_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "input_size", + "type": "index" + }, + { + "name": "output_size", + "type": "index" + }, + { + "name": "token_count", + "type": "index" + } + ], + "bindings": [ + "activation", + "weight", + "tables", + "output" + ], + "binding_access": [ + "read", + "read", + "read", + "write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "../hrx_owned/mul_mat_vec_iq3xxs_f32.loom" + ], + "library_sources": [] + }, + "compile_dependencies": [] + } + ], + "link_modules": [ + { + "name": "router_projection_f32_linked", + "srcs": [ + "qwen3_moe/router_projection_f32.loom" + ], + "libraries": [ + "qwen3_moe/model_config.loom" + ] + }, + { + "name": "router_top8_f32_linked", + "srcs": [ + "qwen3_moe/router_top8_f32.loom" + ], + "libraries": [ + "qwen3_moe/model_config.loom" + ] + }, + { + "name": "router_projection_top8_fused_f32_linked", + "srcs": [ + "qwen3_moe/router_projection_top8_fused_f32.loom" + ], + "libraries": [ + "qwen3_moe/model_config.loom", + "qwen3_moe/router_projection_f32.loom", + "qwen3_moe/router_top8_f32.loom" + ] + }, + { + "name": "attention_postprocess_f32_f16_linked", + "srcs": [ + "qwen3_moe/attention_postprocess_f32_f16.loom" + ], + "libraries": [ + "qwen3_moe/model_config.loom" + ] + }, + { + "name": "attention_prepare_quantized_linked", + "srcs": [ + "qwen3_moe/attention_prepare_quantized.loom" + ], + "libraries": [ + "qwen3_moe/model_config.loom", + "ggml/quantize_q8_1_x4.loom" + ] + }, + { + "name": "attention_qkv_quantized_linked", + "srcs": [ + "qwen3_moe/attention_qkv_quantized.loom" + ], + "libraries": [ + ":dense_linear_quantized_f16_wmma_linked" + ] + }, + { + "name": "attention_qkv_postprocess_fused_linked", + "srcs": [ + "qwen3_moe/attention_qkv_postprocess_fused.loom" + ], + "libraries": [ + ":attention_postprocess_f32_f16_linked", + ":attention_qkv_quantized_linked" + ] + }, + { + "name": "attention_qkv_same_format_prefill_linked", + "srcs": [ + "qwen3_moe/attention_qkv_same_format_prefill.loom" + ], + "libraries": [ + ":dense_linear_quantized_f16_wmma_linked" + ] + }, + { + "name": "linear_q6k_q8_1_x4_linked", + "srcs": [ + "ggml/linear_q6k_q8_1_x4.loom" + ], + "libraries": [ + "ggml/quantize_q8_1_x4.loom" + ] + }, + { + "name": "linear_q6k_f32_linked", + "srcs": [ + "ggml/linear_q6k_f32.loom" + ], + "libraries": [ + ":linear_q6k_q8_1_x4_linked" + ] + }, + { + "name": "routed_down_quantized_f16_wmma_linked", + "srcs": [ + "qwen3_moe/routed_down_quantized_f16_wmma.loom" + ], + "libraries": [ + "qwen3_moe/routed_linear_q4k_f16_wmma.loom", + "qwen3_moe/routed_down_q4k.loom", + "qwen3_moe/routed_down_q6k.loom", + "qwen3_moe/routed_down_next_q8.loom", + "qwen3_moe/attention_prepare_quantized.loom", + "qwen3_moe/model_config.loom", + "ggml/linear_q6k_q8_1_x4.loom", + "qwen3_moe/routed_gate_up_swiglu_q4k.loom", + "ggml/quantize_q8_1_x4.loom" + ] + }, + { + "name": "routed_down_weighted_reduce_next_rmsnorm_f32_linked", + "srcs": [ + "qwen3_moe/routed_down_weighted_reduce_next_rmsnorm_f32.loom" + ], + "libraries": [ + "qwen3_moe/routed_down_quantized_f16_wmma.loom", + "qwen3_moe/routed_linear_q4k_f16_wmma.loom", + "qwen3_moe/routed_down_q4k.loom", + "qwen3_moe/routed_down_q6k.loom", + "qwen3_moe/routed_down_next_q8.loom", + "qwen3_moe/attention_prepare_quantized.loom", + "qwen3_moe/model_config.loom", + "ggml/linear_q6k_q8_1_x4.loom", + "qwen3_moe/routed_gate_up_swiglu_q4k.loom", + "ggml/quantize_q8_1_x4.loom" + ] + }, + { + "name": "routed_down_weighted_reduce_next_rmsnorm_q8_1_x4_linked", + "srcs": [ + "qwen3_moe/routed_down_weighted_reduce_next_rmsnorm_q8_1_x4.loom" + ], + "libraries": [ + "qwen3_moe/routed_down_quantized_f16_wmma.loom", + "qwen3_moe/routed_linear_q4k_f16_wmma.loom", + "qwen3_moe/routed_down_q4k.loom", + "qwen3_moe/routed_down_q6k.loom", + "qwen3_moe/routed_down_next_q8.loom", + "qwen3_moe/attention_prepare_quantized.loom", + "qwen3_moe/model_config.loom", + "ggml/linear_q6k_q8_1_x4.loom", + "qwen3_moe/routed_gate_up_swiglu_q4k.loom", + "ggml/quantize_q8_1_x4.loom" + ] + }, + { + "name": "routed_down_q4k_linked", + "srcs": [ + "qwen3_moe/routed_down_q4k.loom" + ], + "libraries": [ + "qwen3_moe/attention_prepare_quantized.loom", + "qwen3_moe/model_config.loom", + "qwen3_moe/routed_down_next_q8.loom", + "qwen3_moe/routed_gate_up_swiglu_q4k.loom", + "ggml/quantize_q8_1_x4.loom" + ] + }, + { + "name": "routed_down_q6k_linked", + "srcs": [ + "qwen3_moe/routed_down_q6k.loom" + ], + "libraries": [ + "ggml/linear_q6k_f32.loom", + "ggml/linear_q6k_q8_1_x4.loom", + "ggml/quantize_q8_1_x4.loom", + "qwen3_moe/attention_prepare_quantized.loom", + "qwen3_moe/model_config.loom", + "qwen3_moe/routed_down_next_q8.loom" + ] + }, + { + "name": "routed_gate_up_swiglu_q4k_linked", + "srcs": [ + "qwen3_moe/routed_gate_up_swiglu_q4k.loom" + ], + "libraries": [ + "ggml/quantize_q8_1_x4.loom" + ] + }, + { + "name": "batched_decode_expert_dispatch_linked", + "srcs": [ + "qwen3_moe/batched_decode_expert_dispatch.loom" + ], + "libraries": [ + "qwen3_moe/model_config.loom" + ] + }, + { + "name": "batched_decode_gate_up_q4k_linked", + "srcs": [ + "qwen3_moe/batched_decode_gate_up_q4k.loom" + ], + "libraries": [ + "ggml/quantize_q8_1_x4.loom", + "qwen3_moe/batched_decode_expert_dispatch.loom", + "qwen3_moe/routed_gate_up_swiglu_q4k.loom" + ] + }, + { + "name": "expert_table_partition_fused_linked", + "srcs": [ + "qwen3_moe/expert_table_partition_fused.loom" + ], + "libraries": [ + "qwen3_moe/routed_gate_up_swiglu_q4k.loom", + "ggml/quantize_q8_1_x4.loom" + ] + }, + { + "name": "routed_gate_up_swiglu_q4k_f16_wmma_linked", + "srcs": [ + "qwen3_moe/routed_linear_q4k_f16_wmma.loom" + ], + "libraries": [ + "qwen3_moe/routed_gate_up_swiglu_q4k.loom", + "ggml/quantize_q8_1_x4.loom" + ] + }, + { + "name": "dense_linear_quantized_f16_wmma_linked", + "srcs": [ + "qwen3_moe/dense_linear_quantized_f16_wmma.loom" + ], + "libraries": [ + "ggml/linear_q6k_q8_1_x4.loom", + "qwen3_moe/attention_prepare_quantized.loom", + "qwen3_moe/model_config.loom", + "qwen3_moe/routed_linear_q4k_f16_wmma.loom", + "qwen3_moe/routed_gate_up_swiglu_q4k.loom", + "ggml/quantize_q8_1_x4.loom" + ] + }, + { + "name": "flash_attention_decode_split_f32_f16_wmma_next_q8_linked", + "srcs": [ + "qwen3_moe/flash_attention_decode_split_next_q8_test.loom" + ], + "libraries": [ + "qwen3_moe/flash_attention_decode_split_f32_f16_wmma.loom", + "ggml/quantize_q8_1_x4.loom" + ] + } + ], + "plan_cases": [ + { + "name": "router_projection_f32_plan_test", + "args": [ + "$(location :router_projection_f32_linked)", + "--benchmark=@qwen3_moe_router_projection_f32_decode", + "--config=qwen3_moe.model.hidden_size=2048", + "--config=qwen3_moe.router.expert_count=128", + "--config=qwen3_moe.workload.token_capacity=1", + "--dry-run", + "--output-format=jsonl" + ], + "link_module": "router_projection_f32_linked" + }, + { + "name": "router_top8_f32_plan_test", + "args": [ + "$(location :router_top8_f32_linked)", + "--benchmark=@qwen3_moe_router_top8_f32_decode", + "--config=qwen3_moe.router.expert_count=128", + "--config=qwen3_moe.router.route_count=8", + "--config=qwen3_moe.workload.token_capacity=1", + "--dry-run", + "--output-format=jsonl" + ], + "link_module": "router_top8_f32_linked" + }, + { + "name": "router_projection_top8_fused_f32_plan_test", + "args": [ + "$(location :router_projection_top8_fused_f32_linked)", + "--benchmark=@qwen3_moe_router_projection_top8_fused_decode", + "--config=qwen3_moe.model.hidden_size=2048", + "--config=qwen3_moe.router.expert_count=128", + "--config=qwen3_moe.router.route_count=8", + "--config=qwen3_moe.workload.token_capacity=1", + "--dry-run", + "--output-format=jsonl" + ], + "link_module": "router_projection_top8_fused_f32_linked" + }, + { + "name": "attention_postprocess_f32_f16_plan_test", + "args": [ + "$(location :attention_postprocess_f32_f16_linked)", + "--benchmark=@qwen3_moe_attention_postprocess_decode", + "--config=qwen3_moe.attention.head_size=128", + "--config=qwen3_moe.attention.key_value_size=512", + "--config=qwen3_moe.attention.query_size=4096", + "--config=qwen3_moe.model.rms_epsilon=0.000001", + "--config=qwen3_moe.workload.token_capacity=1", + "--dry-run", + "--output-format=jsonl" + ], + "link_module": "attention_postprocess_f32_f16_linked" + }, + { + "name": "attention_prepare_quantized_plan_test", + "args": [ + "$(location :attention_prepare_quantized_linked)", + "--config=qwen3_moe.model.hidden_size=2048", + "--config=qwen3_moe.model.rms_epsilon=0.000001", + "--config=ggml.quantize_q8_1_x4.group_capacity=8192", + "--config=qwen3_moe.workload.token_capacity=512", + "--dry-run", + "--output-format=jsonl" + ], + "link_module": "attention_prepare_quantized_linked" + }, + { + "name": "attention_qkv_quantized_plan_test", + "args": [ + "$(location :attention_qkv_quantized_linked)", + "--benchmark=@qwen3_moe_attention_qkv_full_q6_decode", + "--config=qwen3_moe.attention.key_value_size=512", + "--config=qwen3_moe.attention.query_size=4096", + "--config=qwen3_moe.attention.value_uses_q6=1", + "--config=qwen3_moe.dense_quantized.input_size=2048", + "--config=qwen3_moe.dense_quantized.output_accumulation=0", + "--config=qwen3_moe.dense_quantized.output_size=512", + "--config=qwen3_moe.model.hidden_size=2048", + "--config=qwen3_moe.model.rms_epsilon=0.000001", + "--config=qwen3_moe.workload.token_capacity=32", + "--dry-run", + "--output-format=jsonl" + ], + "link_module": "attention_qkv_quantized_linked" + }, + { + "name": "attention_qkv_postprocess_fused_plan_test", + "args": [ + "$(location :attention_qkv_postprocess_fused_linked)", + "--benchmark=@qwen3_moe_attention_qkv_postprocess_fused_boundary_decode", + "--config=qwen3_moe.attention.head_size=128", + "--config=qwen3_moe.attention.key_value_size=512", + "--config=qwen3_moe.attention.query_size=4096", + "--config=qwen3_moe.attention.value_uses_q6=1", + "--config=qwen3_moe.dense_quantized.input_size=2048", + "--config=qwen3_moe.dense_quantized.output_accumulation=0", + "--config=qwen3_moe.dense_quantized.output_size=512", + "--config=qwen3_moe.model.hidden_size=2048", + "--config=qwen3_moe.model.rms_epsilon=0.000001", + "--config=ggml.quantize_q8_1_x4.group_capacity=16", + "--config=qwen3_moe.workload.token_capacity=1", + "--dry-run", + "--output-format=jsonl" + ], + "link_module": "attention_qkv_postprocess_fused_linked" + }, + { + "name": "attention_qkv_same_format_prefill_plan_test", + "args": [ + "$(location :attention_qkv_same_format_prefill_linked)", + "--benchmark=@qwen3_moe_attention_qkv_q4_prefill_512_fused", + "--config=qwen3_moe.attention.key_value_size=512", + "--config=qwen3_moe.attention.query_size=4096", + "--config=qwen3_moe.dense_quantized.output_accumulation=0", + "--config=qwen3_moe.model.hidden_size=2048", + "--config=qwen3_moe.workload.token_capacity=512", + "--dry-run", + "--output-format=jsonl" + ], + "link_module": "attention_qkv_same_format_prefill_linked" + }, + { + "name": "quantize_q8_1_x4_plan_test", + "args": [ + "$(location ggml/quantize_q8_1_x4.loom)", + "--config=ggml.quantize_q8_1_x4.group_capacity=8192", + "--dry-run", + "--output-format=jsonl" + ], + "source": "ggml/quantize_q8_1_x4.loom" + }, + { + "name": "routed_gate_up_swiglu_q4k_plan_test", + "args": [ + "$(location :routed_gate_up_swiglu_q4k_linked)", + "--benchmark=@qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped_prefill_512", + "--config=qwen3_moe.routed_gate_up.expert_count=128", + "--config=qwen3_moe.routed_gate_up.input_size=2048", + "--config=qwen3_moe.routed_gate_up.output_size=768", + "--config=qwen3_moe.routed_gate_up.route_count=8", + "--config=qwen3_moe.workload.token_capacity=512", + "--dry-run", + "--output-format=jsonl" + ], + "link_module": "routed_gate_up_swiglu_q4k_linked" + }, + { + "name": "linear_q6k_f32_plan_test", + "args": [ + "$(location :linear_q6k_f32_linked)", + "--benchmark=@ggml_linear_q6k_f32_wave64_dense_v_decode", + "--config=ggml.linear_q6k_f32.output_capacity=512", + "--config=ggml.linear_q6k_f32.token_capacity=2048", + "--dry-run", + "--output-format=jsonl" + ], + "link_module": "linear_q6k_f32_linked" + }, + { + "name": "linear_q6k_q8_1_x4_plan_test", + "args": [ + "$(location :linear_q6k_q8_1_x4_linked)", + "--config=ggml.linear_q6k_q8_1_x4.output_capacity=151936", + "--config=ggml.linear_q6k_q8_1_x4.token_capacity=2048", + "--dry-run", + "--output-format=jsonl" + ], + "link_module": "linear_q6k_q8_1_x4_linked" + }, + { + "name": "routed_down_weighted_reduce_next_rmsnorm_f32_plan_test", + "args": [ + "$(location :routed_down_weighted_reduce_next_rmsnorm_f32_linked)", + "--benchmark=@qwen3_moe_routed_down_weighted_reduce_next_rmsnorm_f32_fused_prefill_512", + "--config=qwen3_moe.model.hidden_size=2048", + "--config=qwen3_moe.model.rms_epsilon=0.000001", + "--config=qwen3_moe.routed_down.output_size=2048", + "--config=qwen3_moe.routed_down.route_count=8", + "--config=qwen3_moe.workload.token_capacity=512", + "--dry-run", + "--output-format=jsonl" + ], + "link_module": "routed_down_weighted_reduce_next_rmsnorm_f32_linked" + }, + { + "name": "routed_down_weighted_reduce_next_rmsnorm_q8_1_x4_plan_test", + "args": [ + "$(location :routed_down_weighted_reduce_next_rmsnorm_q8_1_x4_linked)", + "--benchmark=@qwen3_moe_routed_down_weighted_reduce_next_rmsnorm_q8_1_x4_fused_prefill_14", + "--config=qwen3_moe.model.hidden_size=2048", + "--config=qwen3_moe.model.rms_epsilon=0.000001", + "--config=qwen3_moe.routed_down.output_size=2048", + "--config=qwen3_moe.routed_down.route_count=8", + "--config=qwen3_moe.workload.token_capacity=512", + "--dry-run", + "--output-format=jsonl" + ], + "link_module": "routed_down_weighted_reduce_next_rmsnorm_q8_1_x4_linked" + }, + { + "name": "routed_down_quantized_f16_wmma_plan_test", + "args": [ + "$(location :routed_down_quantized_f16_wmma_linked)", + "--config=qwen3_moe.routed_down.expert_count=128", + "--config=qwen3_moe.routed_down.input_size=768", + "--config=qwen3_moe.routed_down.output_size=2048", + "--config=qwen3_moe.routed_down.route_count=8", + "--config=qwen3_moe.workload.token_capacity=512", + "--dry-run", + "--output-format=jsonl" + ], + "link_module": "routed_down_quantized_f16_wmma_linked" + }, + { + "name": "routed_down_q4k_plan_test", + "args": [ + "$(location :routed_down_q4k_linked)", + "--benchmark=@qwen3_moe_routed_down_q4k_q8_1_x4_prefill_512", + "--config=qwen3_moe.routed_down.expert_count=128", + "--config=qwen3_moe.routed_down.input_size=768", + "--config=qwen3_moe.routed_down.output_size=2048", + "--config=qwen3_moe.routed_down.route_count=8", + "--config=qwen3_moe.workload.token_capacity=512", + "--dry-run", + "--output-format=jsonl" + ], + "link_module": "routed_down_q4k_linked" + }, + { + "name": "routed_down_q6k_plan_test", + "args": [ + "$(location :routed_down_q6k_linked)", + "--benchmark=@qwen3_moe_routed_down_q6k_f32_wave64_decode", + "--config=qwen3_moe.routed_down.expert_count=128", + "--config=qwen3_moe.routed_down.input_size=768", + "--config=qwen3_moe.routed_down.output_size=2048", + "--config=qwen3_moe.routed_down.route_count=8", + "--config=qwen3_moe.workload.token_capacity=1", + "--dry-run", + "--output-format=jsonl" + ], + "link_module": "routed_down_q6k_linked" + }, + { + "name": "batched_decode_gate_up_q4k_plan_test", + "args": [ + "$(location :batched_decode_gate_up_q4k_linked)", + "--benchmark=@qwen3_moe_batched_decode_gate_up_q4k_rows2_benchmark", + "--config=qwen3_moe.batched_decode.rows2_descriptor_capacity=64", + "--config=qwen3_moe.batched_decode.schedule0_row_limit=1", + "--config=qwen3_moe.batched_decode.schedule1_row_limit=2", + "--config=qwen3_moe.batched_decode.schedule2_row_limit=4", + "--config=qwen3_moe.router.expert_count=128", + "--config=qwen3_moe.router.route_count=8", + "--config=qwen3_moe.routed_gate_up.expert_count=128", + "--config=qwen3_moe.routed_gate_up.input_size=2048", + "--config=qwen3_moe.routed_gate_up.output_size=768", + "--config=qwen3_moe.routed_gate_up.route_count=8", + "--config=qwen3_moe.workload.token_capacity=16", + "--dry-run", + "--output-format=jsonl" + ], + "link_module": "batched_decode_gate_up_q4k_linked" + }, + { + "name": "batched_decode_expert_dispatch_plan_test", + "args": [ + "$(location :batched_decode_expert_dispatch_linked)", + "--benchmark=@qwen3_moe_batched_decode_expert_dispatch_diverse", + "--config=qwen3_moe.batched_decode.schedule0_row_limit=1", + "--config=qwen3_moe.batched_decode.schedule1_row_limit=2", + "--config=qwen3_moe.batched_decode.schedule2_row_limit=4", + "--config=qwen3_moe.router.expert_count=128", + "--config=qwen3_moe.router.route_count=8", + "--config=qwen3_moe.workload.token_capacity=16", + "--dry-run", + "--output-format=jsonl" + ], + "link_module": "batched_decode_expert_dispatch_linked" + }, + { + "name": "batched_decode_expert_dispatch_configurable_plan_test", + "args": [ + "$(location :batched_decode_expert_dispatch_linked)", + "--benchmark=@qwen3_moe_batched_decode_expert_dispatch_configurable", + "--config=qwen3_moe.batched_decode.schedule0_row_limit=1", + "--config=qwen3_moe.batched_decode.schedule1_row_limit=3", + "--config=qwen3_moe.batched_decode.schedule2_row_limit=6", + "--config=qwen3_moe.router.expert_count=32", + "--config=qwen3_moe.router.route_count=4", + "--config=qwen3_moe.workload.token_capacity=8", + "--dry-run", + "--output-format=jsonl" + ], + "link_module": "batched_decode_expert_dispatch_linked" + }, + { + "name": "expert_table_partition_fused_plan_test", + "args": [ + "$(location :expert_table_partition_fused_linked)", + "--benchmark=@qwen3_moe_expert_table_partition_fused_prefill_512", + "--dry-run", + "--output-format=jsonl" + ], + "link_module": "expert_table_partition_fused_linked" + }, + { + "name": "dense_linear_q6k_f16_wmma_plan_test", + "args": [ + "$(location :dense_linear_quantized_f16_wmma_linked)", + "--benchmark=@qwen3_moe_dense_linear_q6k_f16_wmma_v_prefill_128", + "--config=qwen3_moe.dense_quantized.input_size=2048", + "--config=qwen3_moe.dense_quantized.output_accumulation=0", + "--config=qwen3_moe.dense_quantized.output_size=512", + "--config=qwen3_moe.workload.token_capacity=128", + "--dry-run", + "--output-format=jsonl" + ], + "link_module": "dense_linear_quantized_f16_wmma_linked" + }, + { + "name": "dense_linear_q4k_f16_wmma_o_plan_test", + "args": [ + "$(location :dense_linear_quantized_f16_wmma_linked)", + "--benchmark=@qwen3_moe_dense_linear_q4k_f16_wmma_o_prefill_512", + "--config=qwen3_moe.dense_quantized.input_size=4096", + "--config=qwen3_moe.dense_quantized.output_accumulation=1", + "--config=qwen3_moe.dense_quantized.output_size=2048", + "--config=qwen3_moe.workload.token_capacity=512", + "--dry-run", + "--output-format=jsonl" + ], + "link_module": "dense_linear_quantized_f16_wmma_linked" + }, + { + "name": "dense_linear_q4k_q8_1_x4_plan_test", + "args": [ + "$(location :dense_linear_quantized_f16_wmma_linked)", + "--benchmark=@qwen3_moe_dense_linear_q4k_q8_1_x4_k_decode", + "--config=qwen3_moe.dense_quantized.input_size=2048", + "--config=qwen3_moe.dense_quantized.output_accumulation=0", + "--config=qwen3_moe.dense_quantized.output_size=512", + "--config=qwen3_moe.workload.token_capacity=1", + "--dry-run", + "--output-format=jsonl" + ], + "link_module": "dense_linear_quantized_f16_wmma_linked" + }, + { + "name": "dense_linear_q4k_q8_1_x4_next_q8_plan_test", + "args": [ + "$(location :dense_linear_quantized_f16_wmma_linked)", + "--benchmark=@qwen3_moe_dense_linear_q4k_q8_1_x4_next_q8_decode", + "--config=qwen3_moe.dense_quantized.input_size=4096", + "--config=qwen3_moe.dense_quantized.output_accumulation=1", + "--config=qwen3_moe.dense_quantized.output_size=2048", + "--config=qwen3_moe.model.hidden_size=2048", + "--config=qwen3_moe.model.rms_epsilon=0.000001", + "--config=qwen3_moe.workload.token_capacity=1", + "--dry-run", + "--output-format=jsonl" + ], + "link_module": "dense_linear_quantized_f16_wmma_linked" + }, + { + "name": "routed_gate_up_swiglu_q4k_f16_wmma_plan_test", + "args": [ + "$(location :routed_gate_up_swiglu_q4k_f16_wmma_linked)", + "--config=qwen3_moe.routed_gate_up.expert_count=128", + "--config=qwen3_moe.routed_gate_up.input_size=2048", + "--config=qwen3_moe.routed_gate_up.output_size=768", + "--config=qwen3_moe.routed_gate_up.route_count=8", + "--config=qwen3_moe.workload.token_capacity=2048", + "--dry-run", + "--output-format=jsonl" + ], + "link_module": "routed_gate_up_swiglu_q4k_f16_wmma_linked" + }, + { + "name": "flash_attention_decode_f32_f16_wmma_plan_test", + "args": [ + "$(location qwen3_moe/flash_attention_decode_f32_f16_wmma.loom)", + "--config=qwen3_moe.attention.decode.output_partition_count=1", + "--config=qwen3_moe.attention.key_value_head_count=4", + "--config=qwen3_moe.attention.query_head_count=32", + "--dry-run", + "--output-format=jsonl" + ], + "source": "qwen3_moe/flash_attention_decode_f32_f16_wmma.loom" + }, + { + "name": "flash_attention_decode_q128_f32_f16_wmma_plan_test", + "args": [ + "$(location qwen3_moe/flash_attention_decode_q128_f32_f16_wmma.loom)", + "--config=qwen3_moe.attention.key_value_head_count=4", + "--config=qwen3_moe.attention.query_head_count=32", + "--dry-run", + "--output-format=jsonl" + ], + "source": "qwen3_moe/flash_attention_decode_q128_f32_f16_wmma.loom" + }, + { + "name": "flash_attention_decode_split_f32_f16_wmma_plan_test", + "args": [ + "$(location qwen3_moe/flash_attention_decode_split_f32_f16_wmma.loom)", + "--config=qwen3_moe.attention.key_value_head_count=4", + "--config=qwen3_moe.attention.key_value_token_capacity=2048", + "--config=qwen3_moe.attention.query_head_count=32", + "--dry-run", + "--output-format=jsonl" + ], + "source": "qwen3_moe/flash_attention_decode_split_f32_f16_wmma.loom" + }, + { + "name": "flash_attention_f32_f16_wmma_plan_test", + "args": [ + "$(location qwen3_moe/flash_attention_f32_f16_wmma.loom)", + "--config=qwen3_moe.attention.key_value_head_count=4", + "--config=qwen3_moe.attention.query_head_count=32", + "--config=qwen3_moe.workload.token_capacity=2048", + "--dry-run", + "--output-format=jsonl" + ], + "source": "qwen3_moe/flash_attention_f32_f16_wmma.loom" + }, + { + "name": "owned_token_embedding_decode_plan_test", + "args": [ + "$(location ../qwen_owned/token_embedding_q4k.loom)", + "--benchmark=@qwen_token_embedding_q4k_decode", + "--dry-run", + "--output-format=jsonl", + "--sample-compilation=per_sample" + ], + "source": "../qwen_owned/token_embedding_q4k.loom" + }, + { + "name": "owned_token_embedding_prefill_plan_test", + "args": [ + "$(location ../qwen_owned/token_embedding_q4k.loom)", + "--benchmark=@qwen_token_embedding_q4k_prefill_512", + "--dry-run", + "--output-format=jsonl", + "--sample-compilation=per_sample" + ], + "source": "../qwen_owned/token_embedding_q4k.loom" + }, + { + "name": "owned_attention_context_base_capture_plan_test", + "args": [ + "$(location ../qwen_owned/attention_state_initialize.loom)", + "--benchmark=@qwen_attention_context_base_capture_benchmark", + "--dry-run", + "--output-format=jsonl", + "--sample-compilation=per_sample" + ], + "source": "../qwen_owned/attention_state_initialize.loom", + "owner": "ggml-hrx" + }, + { + "name": "owned_attention_decode_state_initialize_plan_test", + "args": [ + "$(location ../qwen_owned/attention_state_initialize.loom)", + "--benchmark=@qwen_attention_decode_state_initialize_benchmark", + "--dry-run", + "--output-format=jsonl", + "--sample-compilation=per_sample" + ], + "source": "../qwen_owned/attention_state_initialize.loom", + "owner": "ggml-hrx" + }, + { + "name": "owned_attention_metadata_prefill_plan_test", + "args": [ + "$(location ../qwen_owned/attention_metadata.loom)", + "--benchmark=@qwen_attention_metadata_prefill_512", + "--dry-run", + "--output-format=jsonl", + "--sample-compilation=per_sample" + ], + "source": "../qwen_owned/attention_metadata.loom" + }, + { + "name": "owned_gather_add_plan_test", + "args": [ + "$(location ../hrx_owned/gather_add_f32.loom)", + "--benchmark=@ggml_gather_add_noncontiguous", + "--dry-run", + "--output-format=jsonl", + "--sample-compilation=per_sample" + ], + "source": "../hrx_owned/gather_add_f32.loom", + "owner": "ggml-hrx" + } + ], + "license": "Apache-2.0" +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/attention_postprocess_f32_f16.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/attention_postprocess_f32_f16.loom new file mode 100644 index 000000000000..63f81815127f --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/attention_postprocess_f32_f16.loom @@ -0,0 +1,296 @@ +// Qwen grouped-query postprocessing with direct F16 cache publication. +// +// Raw projection rows are physically `[token][head][channel]`: the reshape +// operations in the reference graph only expose that existing layout. Query +// and key rows receive their independent per-head RMSNorm and NEOX rotary +// transform. Queries remain F32 for attention, while keys and values are +// converted directly into their indexed F16 cache rows. No reshape, +// transpose, normalized-row, or rotated-key allocation is materialized. +// +// Logical positions and cache rows are separate runtime inputs. This preserves +// continuous batching, where a token's rotary position need not equal the +// physical cache row selected by the allocator. K and V indices also remain +// distinct so the kernel does not impose an ordering contract on the cache. +// The stage scheduler owns those indices and guarantees they select valid cache +// rows; violating that trusted contract is an error rather than a skipped write. +amdgpu.target @qwen3_moe_attention_postprocess_gfx11_wave32 {subgroup_size = 32} + +config.decl @qwen3_moe.model.rms_epsilon : f32 + +config.decl @qwen3_moe.attention.head_size : %value: index where [range(%value, 4, 1024), mul(%value, 4)] + +config.decl @qwen3_moe.attention.query_size : %value: index where [range(%value, 1, 262144)] + +config.decl @qwen3_moe.attention.key_value_size : %value: index where [range(%value, 1, 262144)] + +// Applies one RMSNorm and two adjacent NEOX pairs. NEOX pairs corresponding +// channels from the low and high halves, while keeping each half contiguous. +// A two-pair packet therefore gives both halves naturally coalesced loads and +// stores without changing the model's pairing rule. +func.def inline @qwen3_moe_rmsnorm_neox_packet(%row_sum: f32, %head_size: f32, %epsilon: f32, %position: f32, %inverse_frequencies: vector<2xf32>, %low_values: vector<2xf32>, %high_values: vector<2xf32>, %low_weights: vector<2xf32>, %high_weights: vector<2xf32>) -> (vector<2xf32>, vector<2xf32>) { + %mean = scalar.divf %row_sum, %head_size : f32 + %biased_mean = scalar.addf %mean, %epsilon : f32 + %scale = scalar.rsqrtf %biased_mean : f32 + %scale_vector = vector.splat %scale : vector<2xf32> + %position_vector = vector.splat %position : vector<2xf32> + %inverse_two_pi = scalar.constant 0.15915494309189535 : f32 + %inverse_two_pi_vector = vector.splat %inverse_two_pi : vector<2xf32> + %normalized_low = vector.mulf %low_values, %scale_vector : vector<2xf32> + %normalized_high = vector.mulf %high_values, %scale_vector : vector<2xf32> + %scaled_low = vector.mulf %normalized_low, %low_weights : vector<2xf32> + %scaled_high = vector.mulf %normalized_high, %high_weights : vector<2xf32> + %angles = vector.mulf %position_vector, %inverse_frequencies : vector<2xf32> + %turns = vector.mulf %angles, %inverse_two_pi_vector : vector<2xf32> + %cosines = vector.costurnsf %turns : vector<2xf32> + %sines = vector.sinturnsf %turns : vector<2xf32> + %low_cosines = vector.mulf %scaled_low, %cosines : vector<2xf32> + %high_sines = vector.mulf %scaled_high, %sines : vector<2xf32> + %low_sines = vector.mulf %scaled_low, %sines : vector<2xf32> + %high_cosines = vector.mulf %scaled_high, %cosines : vector<2xf32> + %rotated_low = vector.subf %low_cosines, %high_sines : vector<2xf32> + %rotated_high = vector.addf %low_sines, %high_cosines : vector<2xf32> + func.return %rotated_low, %rotated_high : vector<2xf32>, vector<2xf32> +} + +// Processes one head domain after its raw projection row is visible. Callers +// may use a larger workgroup than the packet count; inactive workitems +// contribute zero to normalization and perform no loads or stores. +func.def inline @qwen3_moe_attention_postprocess_head_body(%publish_output: i1, %token_count: index, %cache_row_count: index, %head_domain0: index, %token0: index, %positions: buffer, %key_cache_indices: buffer, %value_cache_indices: buffer, %query_input: buffer, %key_input: buffer, %value_input: buffer, %query_norm_weight: buffer, %key_norm_weight: buffer, %inverse_frequencies: buffer, %query_output: buffer, %key_cache: buffer, %value_cache: buffer) { + %query_size0 = config.get @qwen3_moe.attention.query_size : index + %key_value_size0 = config.get @qwen3_moe.attention.key_value_size : index + %head_size0 = config.get @qwen3_moe.attention.head_size : index + %epsilon = config.get @qwen3_moe.model.rms_epsilon : f32 + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 2048)] : index + %bounded_cache_row_count = index.assume %cache_row_count [range(%cache_row_count, 1, 1048576)] : index + %query_size, %key_value_size, %head_size = index.assume %query_size0, %key_value_size0, %head_size0 [range(%query_size0, 1, 262144), range(%key_value_size0, 1, 262144), range(%head_size0, 4, 1024), mul(%head_size0, 4), mul(%query_size0, %head_size0), mul(%key_value_size0, %head_size0)] : index, index, index + %channel0 = kernel.workitem.id : index + %c0 = index.constant 0 : index + %c2 = index.constant 2 : index + %c4 = index.constant 4 : index + %c0_f32 = scalar.constant 0.0 : f32 + %c0_offset = index.constant 0 : offset + %token, %launch_token_count = index.assume %token0, %bounded_token_count [lt(%token0, %bounded_token_count)] : index, index + %half_head_size = index.div %head_size, %c2 : index + %pair_packet_count = index.div %head_size, %c4 : index + %query_head_count = index.div %query_size, %head_size : index + %key_value_head_count = index.div %key_value_size, %head_size : index + %key_value_domain_count = index.mul %key_value_head_count, %c2 : index + %head_domain_count = index.add %query_head_count, %key_value_domain_count : index + %head_domain = index.assume %head_domain0 [lt(%head_domain0, %head_domain_count)] : index + %key_domain_end = index.add %query_head_count, %key_value_head_count : index + %is_query = index.cmp ult, %head_domain, %query_head_count : index + %is_query_or_key = index.cmp ult, %head_domain, %key_domain_end : index + %key_value_head = index.rem %head_domain, %key_value_head_count : index + %active_channel = index.cmp ult, %channel0, %pair_packet_count : index + %head_size_i32 = index.cast %head_size : index to i32 + %head_size_f32 = scalar.sitofp %head_size_i32 : i32 to f32 + %positions_noalias, %key_cache_indices_noalias, %value_cache_indices_noalias, %query_input_noalias, %key_input_noalias, %value_input_noalias, %query_norm_weight_noalias, %key_norm_weight_noalias, %inverse_frequencies_noalias, %query_output_noalias, %key_cache_noalias, %value_cache_noalias = buffer.assume.noalias %positions, %key_cache_indices, %value_cache_indices, %query_input, %key_input, %value_input, %query_norm_weight, %key_norm_weight, %inverse_frequencies, %query_output, %key_cache, %value_cache : buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer + %positions_view = buffer.view %positions_noalias[%c0_offset] : buffer -> view<[%launch_token_count]xi32> + %key_cache_indices_view = buffer.view %key_cache_indices_noalias[%c0_offset] : buffer -> view<[%launch_token_count]xi64> + %value_cache_indices_view = buffer.view %value_cache_indices_noalias[%c0_offset] : buffer -> view<[%launch_token_count]xi64> + %query_input_view = buffer.view %query_input_noalias[%c0_offset] : buffer -> view<[%launch_token_count]x[%query_head_count]x[%head_size]xf32> + %key_input_view = buffer.view %key_input_noalias[%c0_offset] : buffer -> view<[%launch_token_count]x[%key_value_head_count]x[%head_size]xf32> + %value_input_view = buffer.view %value_input_noalias[%c0_offset] : buffer -> view<[%launch_token_count]x[%key_value_head_count]x[%head_size]xf32> + %query_norm_weight_view = buffer.view %query_norm_weight_noalias[%c0_offset] : buffer -> view<[%head_size]xf32> + %key_norm_weight_view = buffer.view %key_norm_weight_noalias[%c0_offset] : buffer -> view<[%head_size]xf32> + %inverse_frequencies_view = buffer.view %inverse_frequencies_noalias[%c0_offset] : buffer -> view<[%half_head_size]xf32> + %query_output_view = buffer.view %query_output_noalias[%c0_offset] : buffer -> view<[%launch_token_count]x[%query_head_count]x[%head_size]xf32> + %key_cache_words_view = buffer.view %key_cache_noalias[%c0_offset] : buffer -> view<[%bounded_cache_row_count]x[%key_value_head_count]x[%half_head_size]xi32> + %value_cache_words_view = buffer.view %value_cache_noalias[%c0_offset] : buffer -> view<[%bounded_cache_row_count]x[%key_value_head_count]x[%half_head_size]xi32> + %row_sum = scf.if %is_query_or_key -> (f32) { + %partial_sum = scf.if %publish_output -> (f32) { + %channel_sum = scf.if %active_channel -> (f32) { + %channel = index.assume %channel0 [lt(%channel0, %pair_packet_count)] : index + %reduction_channel = index.mul %channel, %c4 : index + %reduction_values = scf.if %is_query -> (vector<4xf32>) { + %query_values = vector.load %query_input_view[%token, %head_domain, %reduction_channel] : view<[%launch_token_count]x[%query_head_count]x[%head_size]xf32> -> vector<4xf32> + scf.yield %query_values : vector<4xf32> + } else { + %key_values = vector.load %key_input_view[%token, %key_value_head, %reduction_channel] : view<[%launch_token_count]x[%key_value_head_count]x[%head_size]xf32> -> vector<4xf32> + scf.yield %key_values : vector<4xf32> + } + %squares = vector.mulf %reduction_values, %reduction_values : vector<4xf32> + %sum = vector.reduce %squares, %c0_f32 : vector<4xf32>, f32 + scf.yield %sum : f32 + } else { + scf.yield %c0_f32 : f32 + } + scf.yield %channel_sum : f32 + } else { + scf.yield %c0_f32 : f32 + } + %sum = kernel.workgroup.reduce %partial_sum : f32 + scf.yield %sum : f32 + } else { + scf.yield %c0_f32 : f32 + } + scf.if %publish_output { + scf.if %active_channel { + %channel = index.assume %channel0 [lt(%channel0, %pair_packet_count)] : index + %pair_channel = index.mul %channel, %c2 : index + %paired_channel = index.add %pair_channel, %half_head_size : index + %paired_word = index.add %channel, %pair_packet_count : index + %reduction_channel = index.mul %channel, %c4 : index + scf.if %is_query { + %low_values = vector.load %query_input_view[%token, %head_domain, %pair_channel] : view<[%launch_token_count]x[%query_head_count]x[%head_size]xf32> -> vector<2xf32> + %high_values = vector.load %query_input_view[%token, %head_domain, %paired_channel] : view<[%launch_token_count]x[%query_head_count]x[%head_size]xf32> -> vector<2xf32> + %low_weights = vector.load %query_norm_weight_view[%pair_channel] : view<[%head_size]xf32> -> vector<2xf32> + %high_weights = vector.load %query_norm_weight_view[%paired_channel] : view<[%head_size]xf32> -> vector<2xf32> + %inverse_frequencies_packet = vector.load %inverse_frequencies_view[%pair_channel] : view<[%half_head_size]xf32> -> vector<2xf32> + %position_i32 = view.load %positions_view[%token] : view<[%launch_token_count]xi32> -> i32 + %position = scalar.sitofp %position_i32 : i32 to f32 + %rotated_low, %rotated_high = func.call @qwen3_moe_rmsnorm_neox_packet(%row_sum, %head_size_f32, %epsilon, %position, %inverse_frequencies_packet, %low_values, %high_values, %low_weights, %high_weights) : (f32, f32, f32, f32, vector<2xf32>, vector<2xf32>, vector<2xf32>, vector<2xf32>, vector<2xf32>) -> (vector<2xf32>, vector<2xf32>) + vector.store %rotated_low, %query_output_view[%token, %head_domain, %pair_channel] : vector<2xf32>, view<[%launch_token_count]x[%query_head_count]x[%head_size]xf32> + vector.store %rotated_high, %query_output_view[%token, %head_domain, %paired_channel] : vector<2xf32>, view<[%launch_token_count]x[%query_head_count]x[%head_size]xf32> + } else { + scf.if %is_query_or_key { + %low_values = vector.load %key_input_view[%token, %key_value_head, %pair_channel] : view<[%launch_token_count]x[%key_value_head_count]x[%head_size]xf32> -> vector<2xf32> + %high_values = vector.load %key_input_view[%token, %key_value_head, %paired_channel] : view<[%launch_token_count]x[%key_value_head_count]x[%head_size]xf32> -> vector<2xf32> + %low_weights = vector.load %key_norm_weight_view[%pair_channel] : view<[%head_size]xf32> -> vector<2xf32> + %high_weights = vector.load %key_norm_weight_view[%paired_channel] : view<[%head_size]xf32> -> vector<2xf32> + %inverse_frequencies_packet = vector.load %inverse_frequencies_view[%pair_channel] : view<[%half_head_size]xf32> -> vector<2xf32> + %position_i32 = view.load %positions_view[%token] : view<[%launch_token_count]xi32> -> i32 + %position = scalar.sitofp %position_i32 : i32 to f32 + %cache_index_raw = view.load %key_cache_indices_view[%token] : view<[%launch_token_count]xi64> -> i64 + %cache_index_i64 = scalar.assume %cache_index_raw [range(%cache_index_raw, 0, 1048575)] : i64 + %cache_index0 = index.cast %cache_index_i64 : i64 to index + %cache_index = index.assume %cache_index0 [lt(%cache_index0, %bounded_cache_row_count)] : index + %rotated_low, %rotated_high = func.call @qwen3_moe_rmsnorm_neox_packet(%row_sum, %head_size_f32, %epsilon, %position, %inverse_frequencies_packet, %low_values, %high_values, %low_weights, %high_weights) : (f32, f32, f32, f32, vector<2xf32>, vector<2xf32>, vector<2xf32>, vector<2xf32>, vector<2xf32>) -> (vector<2xf32>, vector<2xf32>) + %half_low = vector.fptrunc %rotated_low : vector<2xf32> to vector<2xf16> + %half_high = vector.fptrunc %rotated_high : vector<2xf32> to vector<2xf16> + %packed_low = vector.bitcast %half_low : vector<2xf16> to vector<1xi32> + %packed_high = vector.bitcast %half_high : vector<2xf16> to vector<1xi32> + vector.store %packed_low, %key_cache_words_view[%cache_index, %key_value_head, %channel] : vector<1xi32>, view<[%bounded_cache_row_count]x[%key_value_head_count]x[%half_head_size]xi32> + vector.store %packed_high, %key_cache_words_view[%cache_index, %key_value_head, %paired_word] : vector<1xi32>, view<[%bounded_cache_row_count]x[%key_value_head_count]x[%half_head_size]xi32> + } else { + %cache_index_raw = view.load %value_cache_indices_view[%token] : view<[%launch_token_count]xi64> -> i64 + %cache_index_i64 = scalar.assume %cache_index_raw [range(%cache_index_raw, 0, 1048575)] : i64 + %cache_index0 = index.cast %cache_index_i64 : i64 to index + %cache_index = index.assume %cache_index0 [lt(%cache_index0, %bounded_cache_row_count)] : index + %values = vector.load %value_input_view[%token, %key_value_head, %reduction_channel] : view<[%launch_token_count]x[%key_value_head_count]x[%head_size]xf32> -> vector<4xf32> + %half_values = vector.fptrunc %values : vector<4xf32> to vector<4xf16> + %packed_values = vector.bitcast %half_values : vector<4xf16> to vector<2xi32> + vector.store %packed_values, %value_cache_words_view[%cache_index, %key_value_head, %pair_channel] : vector<2xi32>, view<[%bounded_cache_row_count]x[%key_value_head_count]x[%half_head_size]xi32> + } + } + } + } + func.return +} + +kernel.def target(@qwen3_moe_attention_postprocess_gfx11_wave32) @qwen3_moe_attention_postprocess_f32_f16(%token_count: index, %cache_row_count: index) { + %query_size = config.get @qwen3_moe.attention.query_size : index + %key_value_size = config.get @qwen3_moe.attention.key_value_size : index + %head_size = config.get @qwen3_moe.attention.head_size : index + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c4 = index.constant 4 : index + %pair_packet_count = index.div %head_size, %c4 : index + %query_head_count = index.div %query_size, %head_size : index + %key_value_head_count = index.div %key_value_size, %head_size : index + %key_value_domain_count = index.mul %key_value_head_count, %c2 : index + %head_domain_count = index.add %query_head_count, %key_value_domain_count : index + kernel.launch.config workgroups(%head_domain_count, %token_count, %c1) workgroup_size(%pair_packet_count, %c1, %c1) : index +} launch(%token_count: index, %cache_row_count: index, %positions: buffer, %key_cache_indices: buffer, %value_cache_indices: buffer, %query_input: buffer, %key_input: buffer, %value_input: buffer, %query_norm_weight: buffer, %key_norm_weight: buffer, %inverse_frequencies: buffer, %query_output: buffer, %key_cache: buffer, %value_cache: buffer) where [range(%token_count, 1, 2048)] { + %head_domain = kernel.workgroup.id : index + %token0 = kernel.workgroup.id : index + %c0 = index.constant 0 : index + %valid_token = index.cmp ult, %token0, %token_count : index + %safe_token = scf.select %valid_token, %token0, %c0 : index + func.call @qwen3_moe_attention_postprocess_head_body(%valid_token, %token_count, %cache_row_count, %head_domain, %safe_token, %positions, %key_cache_indices, %value_cache_indices, %query_input, %key_input, %value_input, %query_norm_weight, %key_norm_weight, %inverse_frequencies, %query_output, %key_cache, %value_cache) : (i1, index, index, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + kernel.return +} + +// At position zero NEOX is the identity. Uniform rows make the RMSNorm result +// analytically one to within epsilon, while the distinct V value proves that +// all three output domains reach their intended bindings. +check.case public @qwen3_moe_attention_postprocess_identity_case { + %token_count = check.literal value(2) : index + %cache_row_count = check.literal value(2) : index + %positions = check.generate.fill value(0) : tensor<2xi32> + %key_cache_indices = check.generate.iota offset(0) step(1) : tensor<2xi64> + %value_cache_indices = check.generate.iota offset(0) step(1) : tensor<2xi64> + %query_input = check.generate.fill value(1.0) : tensor<2x1x4xf32> + %key_input = check.generate.fill value(1.0) : tensor<2x1x4xf32> + %value_input = check.generate.fill value(2.0) : tensor<2x1x4xf32> + %query_norm_weight = check.generate.fill value(1.0) : tensor<4xf32> + %key_norm_weight = check.generate.fill value(1.0) : tensor<4xf32> + %inverse_frequencies = check.generate.fill value(1.0) : tensor<2xf32> + %query_output = check.generate.fill value(0.0) : tensor<2x1x4xf32> + %key_cache = check.generate.fill value(0.0) : tensor<2x1x4xf16> + %value_cache = check.generate.fill value(0.0) : tensor<2x1x4xf16> + %expected_query = check.generate.fill value(1.0) : tensor<2x1x4xf32> + %expected_key = check.generate.fill value(1.0) : tensor<2x1x4xf16> + %expected_value = check.generate.fill value(2.0) : tensor<2x1x4xf16> + kernel.launch @qwen3_moe_attention_postprocess_f32_f16[%token_count, %cache_row_count](%token_count, %cache_row_count, %positions, %key_cache_indices, %value_cache_indices, %query_input, %key_input, %value_input, %query_norm_weight, %key_norm_weight, %inverse_frequencies, %query_output, %key_cache, %value_cache) : [index, index](index, index, tensor<2xi32>, tensor<2xi64>, tensor<2xi64>, tensor<2x1x4xf32>, tensor<2x1x4xf32>, tensor<2x1x4xf32>, tensor<4xf32>, tensor<4xf32>, tensor<2xf32>, tensor<2x1x4xf32>, tensor<2x1x4xf16>, tensor<2x1x4xf16>) + check.expect.close actual(%query_output) expected(%expected_query) atol(1.0000000000000001e-05) rtol(1.0000000000000001e-05) nan(same) : tensor<2x1x4xf32> + check.expect.close actual(%key_cache) expected(%expected_key) atol(0.001) rtol(0.001) nan(same) : tensor<2x1x4xf16> + check.expect.close actual(%value_cache) expected(%expected_value) atol(0.0) rtol(0.0) nan(same) : tensor<2x1x4xf16> + check.return +} + +// Distinct angles make the two NEOX pairs rotate `[1, 3]` and `[2, 4]` +// into `[sqrt(5), sqrt(5)]` and `[sqrt(10), sqrt(10)]`. Their flattened +// `[A, B, A, B]` result catches adjacent-channel pairing and packet-order +// mistakes without embedding a second implementation in the test. +check.case public @qwen3_moe_attention_postprocess_neox_case { + %token_count = check.literal value(2) : index + %cache_row_count = check.literal value(2) : index + %positions = check.generate.fill value(1) : tensor<2xi32> + %key_cache_indices = check.generate.iota offset(0) step(1) : tensor<2xi64> + %value_cache_indices = check.generate.iota offset(0) step(1) : tensor<2xi64> + %query_input = check.generate.fill value(1.0) : tensor<2x1x4xf32> + %key_input = check.generate.fill value(1.0) : tensor<2x1x4xf32> + %value_input = check.generate.fill value(0.0) : tensor<2x1x4xf32> + %query_norm_weight = check.generate.iota offset(1.0) step(1.0) : tensor<4xf32> + %key_norm_weight = check.generate.iota offset(1.0) step(1.0) : tensor<4xf32> + %inverse_frequencies = check.generate.iota offset(-0.46364760900080609) step(0.14189705460416391) : tensor<2xf32> + %query_output = check.generate.fill value(0.0) : tensor<2x1x4xf32> + %key_cache = check.generate.fill value(0.0) : tensor<2x1x4xf16> + %value_cache = check.generate.fill value(0.0) : tensor<2x1x4xf16> + %expected_query = check.generate.iota offset(2.2360668182373047) step(0.92620921134948736) period(2) : tensor<2x1x4xf32> + %expected_key = check.generate.iota offset(2.236328125) step(0.92578125) period(2) : tensor<2x1x4xf16> + %expected_value = check.generate.fill value(0.0) : tensor<2x1x4xf16> + kernel.launch @qwen3_moe_attention_postprocess_f32_f16[%token_count, %cache_row_count](%token_count, %cache_row_count, %positions, %key_cache_indices, %value_cache_indices, %query_input, %key_input, %value_input, %query_norm_weight, %key_norm_weight, %inverse_frequencies, %query_output, %key_cache, %value_cache) : [index, index](index, index, tensor<2xi32>, tensor<2xi64>, tensor<2xi64>, tensor<2x1x4xf32>, tensor<2x1x4xf32>, tensor<2x1x4xf32>, tensor<4xf32>, tensor<4xf32>, tensor<2xf32>, tensor<2x1x4xf32>, tensor<2x1x4xf16>, tensor<2x1x4xf16>) + check.expect.close actual(%query_output) expected(%expected_query) atol(0.0001) rtol(0.0001) nan(same) : tensor<2x1x4xf32> + check.expect.close actual(%key_cache) expected(%expected_key) atol(0.002) rtol(0.002) nan(same) : tensor<2x1x4xf16> + check.expect.close actual(%value_cache) expected(%expected_value) atol(0.0) rtol(0.0) nan(same) : tensor<2x1x4xf16> + check.return +} + +check.case public @qwen3_moe_attention_postprocess_benchmark_case { + %token_count = check.param.choice values([1, 32, 128, 512]) name("token_count") : index + %cache_row_count = check.literal value(4096) : index + %positions = check.generate.iota offset(0) step(1) : tensor<[%token_count]xi32> + %key_cache_indices = check.generate.iota offset(0) step(1) : tensor<[%token_count]xi64> + %value_cache_indices = check.generate.iota offset(0) step(1) : tensor<[%token_count]xi64> + %query_input = check.generate.fill value(0.0) : tensor<[%token_count]x32x128xf32> + %key_input = check.generate.fill value(0.0) : tensor<[%token_count]x4x128xf32> + %value_input = check.generate.fill value(0.0) : tensor<[%token_count]x4x128xf32> + %query_norm_weight = check.generate.fill value(1.0) : tensor<128xf32> + %key_norm_weight = check.generate.fill value(1.0) : tensor<128xf32> + %inverse_frequencies = check.generate.fill value(1.0) : tensor<64xf32> + %query_output = check.generate.fill value(1.0) : tensor<[%token_count]x32x128xf32> + %key_cache = check.generate.fill value(0.0) : tensor<4096x4x128xf16> + %value_cache = check.generate.fill value(0.0) : tensor<4096x4x128xf16> + %expected_query = check.generate.fill value(0.0) : tensor<[%token_count]x32x128xf32> + %expected_key_cache = check.generate.fill value(0.0) : tensor<4096x4x128xf16> + %expected_value_cache = check.generate.fill value(0.0) : tensor<4096x4x128xf16> + kernel.launch @qwen3_moe_attention_postprocess_f32_f16[%token_count, %cache_row_count](%token_count, %cache_row_count, %positions, %key_cache_indices, %value_cache_indices, %query_input, %key_input, %value_input, %query_norm_weight, %key_norm_weight, %inverse_frequencies, %query_output, %key_cache, %value_cache) : [index, index](index, index, tensor<[%token_count]xi32>, tensor<[%token_count]xi64>, tensor<[%token_count]xi64>, tensor<[%token_count]x32x128xf32>, tensor<[%token_count]x4x128xf32>, tensor<[%token_count]x4x128xf32>, tensor<128xf32>, tensor<128xf32>, tensor<64xf32>, tensor<[%token_count]x32x128xf32>, tensor<4096x4x128xf16>, tensor<4096x4x128xf16>) + check.expect.close actual(%query_output) expected(%expected_query) atol(0.0) rtol(0.0) nan(same) : tensor<[%token_count]x32x128xf32> + check.expect.close actual(%key_cache) expected(%expected_key_cache) atol(0.0) rtol(0.0) nan(same) : tensor<4096x4x128xf16> + check.expect.close actual(%value_cache) expected(%expected_value_cache) atol(0.0) rtol(0.0) nan(same) : tensor<4096x4x128xf16> + check.return +} + +check.benchmark<@qwen3_moe_attention_postprocess_identity_case> @qwen3_moe_attention_postprocess_identity + +check.benchmark<@qwen3_moe_attention_postprocess_neox_case> @qwen3_moe_attention_postprocess_neox + +check.benchmark<@qwen3_moe_attention_postprocess_benchmark_case> @qwen3_moe_attention_postprocess_decode {token_count = 1} + +check.benchmark<@qwen3_moe_attention_postprocess_benchmark_case> @qwen3_moe_attention_postprocess_prefill_32 {token_count = 32} + +check.benchmark<@qwen3_moe_attention_postprocess_benchmark_case> @qwen3_moe_attention_postprocess_prefill_128 {token_count = 128} + +check.benchmark<@qwen3_moe_attention_postprocess_benchmark_case> @qwen3_moe_attention_postprocess_prefill_512 {token_count = 512} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/attention_prepare_quantized.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/attention_prepare_quantized.loom new file mode 100644 index 000000000000..497b8e49eaee --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/attention_prepare_quantized.loom @@ -0,0 +1,432 @@ +// Qwen attention preparation for raw GGUF quantized projections. +// +// Decode contractions consume a shared Q8_1 x4 activation row. Producing that +// row directly from the residual stream removes the materialized F32 attention +// RMSNorm tensor and avoids rereading it solely for quantization. The fused +// producer preserves GGML's physical Q8_1 x4 contract: +// +// struct block_q8_1_x4 { +// f16 ds[4][2]; +// i32 qs[4][8]; +// }; +// +// One 256-workitem workgroup owns one token. It first reduces the complete +// hidden row to one reciprocal RMS scale, then visits 1024-element stripes. +// Every workitem packs one four-value word per stripe; adjacent eight-lane +// cohorts reduce the maximum and quantized sum for one logical Q8_1 block. +// Scratch is fixed at 1152 bytes regardless of hidden size. +amdgpu.target @qwen3_moe_attention_prepare_gfx11_wave32 {subgroup_size = 32} + +config.decl @qwen3_moe.model.hidden_size : %value: index where [range(%value, 128, 32768), mul(%value, 128)] + +config.decl @qwen3_moe.model.rms_epsilon : f32 + +// Shared Q8_1 x4 packer used by the standalone differential path. +kernel.decl @ggml_quantize_q8_1_x4_f32(%token_count: index, %input_size: index) launch(%token_count: index, %input_size: index, %input: buffer, %output: buffer) + +// Materializes the ordinary Qwen RMSNorm boundary. This remains useful for +// prefill schedules that reuse one normalized row across many output tiles; +// decode uses the fused Q8_1 producer below. +kernel.def target(@qwen3_moe_attention_prepare_gfx11_wave32) @qwen3_moe_rmsnorm_f32(%token_count: index) { + %c1 = index.constant 1 : index + %c256 = index.constant 256 : index + kernel.launch.config workgroups(%token_count, %c1, %c1) workgroup_size(%c256, %c1, %c1) : index +} launch(%token_count: index, %input: buffer, %weight: buffer, %output: buffer) where [range(%token_count, 1, 2048)] { + %hidden_size0 = config.get @qwen3_moe.model.hidden_size : index + %epsilon = config.get @qwen3_moe.model.rms_epsilon : f32 + %hidden_size = index.assume %hidden_size0 [range(%hidden_size0, 128, 32768), mul(%hidden_size0, 128)] : index + %token0 = kernel.workgroup.id : index + %workitem = kernel.workitem.id : index + %subgroup0 = kernel.subgroup.id : index + %subgroup = index.assume %subgroup0 [range(%subgroup0, 0, 7)] : index + %lane = kernel.subgroup.lane.id : index + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c8 = index.constant 8 : index + %c256 = index.constant 256 : index + %scratch_bytes = index.constant 1024 : offset + %c0_offset = index.constant 0 : offset + %c0_f32 = scalar.constant 0.0 : f32 + %valid_token = index.cmp ult, %token0, %token_count : index + %safe_token0 = scf.select %valid_token, %token0, %c0 : index + %token, %launch_token_count = index.assume %safe_token0, %token_count [lt(%safe_token0, %token_count)] : index, index + %hidden_size_i32 = index.cast %hidden_size : index to i32 + %hidden_size_f32 = scalar.sitofp %hidden_size_i32 : i32 to f32 + %input_noalias, %weight_noalias, %output_noalias = buffer.assume.noalias %input, %weight, %output : buffer, buffer, buffer + %input_view = buffer.view %input_noalias[%c0_offset] : buffer -> view<[%launch_token_count]x[%hidden_size]xf32> + %weight_view = buffer.view %weight_noalias[%c0_offset] : buffer -> view<[%hidden_size]xf32> + %output_view = buffer.view %output_noalias[%c0_offset] : buffer -> view<[%launch_token_count]x[%hidden_size]xf32> + %thread_sum = scf.for %channel = [%workitem to %hidden_size step %c256](%running_sum = %c0_f32 : f32) -> (f32) { + %value = view.load %input_view[%token, %channel] : view<[%launch_token_count]x[%hidden_size]xf32> -> f32 + %square = scalar.mulf %value, %value : f32 + %next_sum = scalar.addf %running_sum, %square : f32 + scf.yield %next_sum : f32 + } + %subgroup_sum = kernel.subgroup.reduce %thread_sum : f32 + %scratch = buffer.alloca align(16) %scratch_bytes : buffer + %scratch_view = buffer.view %scratch[%c0_offset] : buffer -> view<256xf32> + %is_subgroup_leader = index.cmp eq, %lane, %c0 : index + scf.if %is_subgroup_leader { + view.store %subgroup_sum, %scratch_view[%subgroup] : f32, view<256xf32> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + %is_reduction_subgroup = index.cmp eq, %subgroup, %c0 : index + %is_reduction_lane = index.cmp ult, %lane, %c8 : index + %loads_subgroup_sum = scalar.andi %is_reduction_subgroup, %is_reduction_lane : i1 + %subgroup_partial = scf.if %loads_subgroup_sum -> (f32) { + %value = view.load %scratch_view[%lane] : view<256xf32> -> f32 + scf.yield %value : f32 + } else { + scf.yield %c0_f32 : f32 + } + %row_sum = kernel.subgroup.reduce %subgroup_partial : f32 + %writes_scale = scalar.andi %is_reduction_subgroup, %is_subgroup_leader : i1 + scf.if %writes_scale { + %mean = scalar.divf %row_sum, %hidden_size_f32 : f32 + %biased_mean = scalar.addf %mean, %epsilon : f32 + %scale = scalar.rsqrtf %biased_mean : f32 + view.store %scale, %scratch_view[%c0] : f32, view<256xf32> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + %scale = view.load %scratch_view[%c0] : view<256xf32> -> f32 + scf.if %valid_token { + scf.for %channel = [%workitem to %hidden_size step %c256] { + %value = view.load %input_view[%token, %channel] : view<[%launch_token_count]x[%hidden_size]xf32> -> f32 + %learned_weight = view.load %weight_view[%channel] : view<[%hidden_size]xf32> -> f32 + %normalized = scalar.mulf %value, %scale : f32 + %result = scalar.mulf %normalized, %learned_weight : f32 + view.store %result, %output_view[%token, %channel] : f32, view<[%launch_token_count]x[%hidden_size]xf32> + } + } + kernel.return +} + +// Shared RMSNorm and GGML Q8_1 x4 row producer. The caller owns launch geometry +// and supplies the logical token whose complete hidden row this workgroup owns. +// The attention export discards the ordinary F32 row, while feed-forward +// publishes it for the router without rereading and requantizing the normalized +// values. +func.def inline @qwen3_moe_rmsnorm_quantize_q8_1_x4_body(%publish_normalized: i1, %reduction_subgroup_count0: index, %token_count: index, %token0: index, %input: buffer, %weight: buffer, %normalized_output: buffer, %q8_output: buffer) { + %hidden_size0 = config.get @qwen3_moe.model.hidden_size : index + %epsilon = config.get @qwen3_moe.model.rms_epsilon : f32 + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 2048)] : index + %reduction_subgroup_count = index.assume %reduction_subgroup_count0 [range(%reduction_subgroup_count0, 1, 8)] : index + %hidden_size = index.assume %hidden_size0 [range(%hidden_size0, 128, 32768), mul(%hidden_size0, 128)] : index + %workitem = kernel.workitem.id : index + %subgroup0 = kernel.subgroup.id : index + %subgroup = index.assume %subgroup0 [range(%subgroup0, 0, 7)] : index + %lane = kernel.subgroup.lane.id : index + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c4 = index.constant 4 : index + %c8 = index.constant 8 : index + %c16 = index.constant 16 : index + %c32 = index.constant 32 : index + %c128 = index.constant 128 : index + %c256 = index.constant 256 : index + %c1024 = index.constant 1024 : index + %group_bytes = index.constant 144 : offset + %payload_byte_add = index.constant 16 : offset + %scratch_d_byte_add = index.constant 1024 : offset + %scratch_bytes = index.constant 1152 : offset + %c0_offset = index.constant 0 : offset + %c0_f32 = scalar.constant 0.0 : f32 + %c1_f32 = scalar.constant 1.0 : f32 + %c127 = scalar.constant 127.0 : f32 + %c0_f32x4 = vector.constant 0.0 : vector<4xf32> + %valid_token = index.cmp ult, %token0, %bounded_token_count : index + %safe_token0 = scf.select %valid_token, %token0, %c0 : index + %token, %launch_token_count = index.assume %safe_token0, %bounded_token_count [lt(%safe_token0, %bounded_token_count)] : index, index + %hidden_size_i32 = index.cast %hidden_size : index to i32 + %hidden_size_f32 = scalar.sitofp %hidden_size_i32 : i32 to f32 + %physical_group_count = index.div %hidden_size, %c128 : index + %row_bytes = index.scale %physical_group_count, %group_bytes : index, offset -> offset + %token_output_byte_base = index.scale %token, %row_bytes : index, offset -> offset + %input_noalias, %weight_noalias, %q8_output_noalias = buffer.assume.noalias %input, %weight, %q8_output : buffer, buffer, buffer + %input_view = buffer.view %input_noalias[%c0_offset] : buffer -> view<[%launch_token_count]x[%hidden_size]xf32> + %weight_view = buffer.view %weight_noalias[%c0_offset] : buffer -> view<[%hidden_size]xf32> + %normalized_output_view = buffer.view %normalized_output[%c0_offset] : buffer -> view<[%launch_token_count]x[%hidden_size]xf32> + %scratch = buffer.alloca align(16) %scratch_bytes : buffer + %scratch_values = buffer.view %scratch[%c0_offset] : buffer -> view<256xf32> + %scratch_d = buffer.view %scratch[%scratch_d_byte_add] : buffer -> view<32xf32> + // Reduce the complete row before any block-local quantization. + %thread_sum = scf.for %channel = [%workitem to %hidden_size step %c256](%running_sum = %c0_f32 : f32) -> (f32) { + %value = view.load %input_view[%token, %channel] : view<[%launch_token_count]x[%hidden_size]xf32> -> f32 + %square = scalar.mulf %value, %value : f32 + %next_sum = scalar.addf %running_sum, %square : f32 + scf.yield %next_sum : f32 + } + %subgroup_sum = kernel.subgroup.reduce %thread_sum : f32 + %is_subgroup_leader = index.cmp eq, %lane, %c0 : index + scf.if %is_subgroup_leader { + view.store %subgroup_sum, %scratch_values[%subgroup] : f32, view<256xf32> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + %is_reduction_subgroup = index.cmp eq, %subgroup, %c0 : index + %is_reduction_lane = index.cmp ult, %lane, %reduction_subgroup_count : index + %loads_subgroup_sum = scalar.andi %is_reduction_subgroup, %is_reduction_lane : i1 + %subgroup_partial = scf.if %loads_subgroup_sum -> (f32) { + %value = view.load %scratch_values[%lane] : view<256xf32> -> f32 + scf.yield %value : f32 + } else { + scf.yield %c0_f32 : f32 + } + %row_sum = kernel.subgroup.reduce %subgroup_partial : f32 + %writes_scale = scalar.andi %is_reduction_subgroup, %is_subgroup_leader : i1 + scf.if %writes_scale { + %mean = scalar.divf %row_sum, %hidden_size_f32 : f32 + %biased_mean = scalar.addf %mean, %epsilon : f32 + %scale = scalar.rsqrtf %biased_mean : f32 + view.store %scale, %scratch_values[%c0] : f32, view<256xf32> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + %row_scale = view.load %scratch_values[%c0] : view<256xf32> -> f32 + %row_scale_vector = vector.splat %row_scale : vector<4xf32> + %publishes_normalized = scalar.andi %publish_normalized, %valid_token : i1 + // Reuse one scratch frame for each 1024-element stripe. Hidden sizes need + // only be divisible by 128; inactive workitems in the final stripe carry + // zeros and never publish. + scf.for %stripe_base = [%c0 to %hidden_size step %c1024] { + %word_element_add = index.mul %workitem, %c4 : index + %channel = index.add %stripe_base, %word_element_add : index + %valid_word = index.cmp ult, %channel, %hidden_size : index + %mask = vector.mask.range [%channel to %hidden_size step %c1] : index -> vector<4xi1> + %input_values = vector.load.mask %input_view[%token, %channel], %mask, %c0_f32x4 : view<[%launch_token_count]x[%hidden_size]xf32>, vector<4xi1>, vector<4xf32> + %learned_weights = vector.load.mask %weight_view[%channel], %mask, %c0_f32x4 : view<[%hidden_size]xf32>, vector<4xi1>, vector<4xf32> + %normalized0 = vector.mulf %input_values, %row_scale_vector : vector<4xf32> + %normalized = vector.mulf %normalized0, %learned_weights : vector<4xf32> + scf.if %publishes_normalized { + vector.store.mask %normalized, %normalized_output_view[%token, %channel], %mask : vector<4xf32>, view<[%launch_token_count]x[%hidden_size]xf32>, vector<4xi1> + } + %absolute_values = vector.absf %normalized : vector<4xf32> + %thread_max = vector.reduce %absolute_values, %c0_f32 : vector<4xf32>, f32 + view.store %thread_max, %scratch_values[%workitem] : f32, view<256xf32> + kernel.barrier scope(workgroup) ordering(acq_rel) + %word_in_block = index.rem %workitem, %c8 : index + %block_in_stripe = index.div %workitem, %c8 : index + %is_block_leader = index.cmp eq, %word_in_block, %c0 : index + %writes_block_d = scalar.andi %valid_word, %is_block_leader : i1 + scf.if %writes_block_d { + %cohort_base = index.mul %block_in_stripe, %c8 : index + %cohort_maxima = vector.load %scratch_values[%cohort_base] : view<256xf32> -> vector<8xf32> + %amax = vector.reduce %cohort_maxima, %c0_f32 : vector<8xf32>, f32 + %d = scalar.divf %amax, %c127 : f32 + view.store %d, %scratch_d[%block_in_stripe] : f32, view<32xf32> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + %d = scf.if %valid_word -> (f32) { + %block_d = view.load %scratch_d[%block_in_stripe] : view<32xf32> -> f32 + scf.yield %block_d : f32 + } else { + scf.yield %c0_f32 : f32 + } + %d_nonzero = scalar.cmpf one, %d, %c0_f32 : f32 + %d_inverse = scf.if %d_nonzero -> (f32) { + %inverse = scalar.divf %c1_f32, %d : f32 + scf.yield %inverse : f32 + } else { + scf.yield %c0_f32 : f32 + } + %d_inverse_vector = vector.splat %d_inverse : vector<4xf32> + %scaled_values = vector.mulf %normalized, %d_inverse_vector : vector<4xf32> + %rounded_values = vector.roundf %scaled_values : vector<4xf32> + %quantized_values = vector.fptosi %rounded_values : vector<4xf32> to vector<4xi8> + %packed_word = vector.bitcast %quantized_values : vector<4xi8> to vector<1xi32> + %publishes_q8_word = scalar.andi %valid_word, %valid_token : i1 + scf.if %publishes_q8_word { + %q8_block = index.div %channel, %c32 : index + %physical_group = index.div %q8_block, %c4 : index + %block_in_group = index.rem %q8_block, %c4 : index + %group_byte_add = index.scale %physical_group, %group_bytes : index, offset -> offset + %group_byte_offset = index.add %token_output_byte_base, %group_byte_add : offset + %payload_byte_offset = index.add %group_byte_offset, %payload_byte_add : offset + %group_ds = buffer.view %q8_output_noalias[%group_byte_offset] : buffer -> view<8xf16> + %group_qs = buffer.view %q8_output_noalias[%payload_byte_offset] : buffer -> view<32xi32> + %block_word_base = index.mul %block_in_group, %c8 : index + %packed_word_index0 = index.add %block_word_base, %word_in_block : index + %packed_word_index = index.assume %packed_word_index0 [range(%packed_word_index0, 0, 31)] : index + vector.store %packed_word, %group_qs[%packed_word_index] : vector<1xi32>, view<32xi32> + } + %thread_quantized_sum = vector.reduce %rounded_values, %c0_f32 : vector<4xf32>, f32 + view.store %thread_quantized_sum, %scratch_values[%workitem] : f32, view<256xf32> + kernel.barrier scope(workgroup) ordering(acq_rel) + %publishes_block_ds = scalar.andi %writes_block_d, %valid_token : i1 + scf.if %publishes_block_ds { + %cohort_base = index.mul %block_in_stripe, %c8 : index + %cohort_sums = vector.load %scratch_values[%cohort_base] : view<256xf32> -> vector<8xf32> + %quantized_sum = vector.reduce %cohort_sums, %c0_f32 : vector<8xf32>, f32 + %s = scalar.mulf %quantized_sum, %d : f32 + %q8_block = index.div %channel, %c32 : index + %physical_group = index.div %q8_block, %c4 : index + %block_in_group = index.rem %q8_block, %c4 : index + %group_byte_add = index.scale %physical_group, %group_bytes : index, offset -> offset + %group_byte_offset = index.add %token_output_byte_base, %group_byte_add : offset + %group_ds = buffer.view %q8_output_noalias[%group_byte_offset] : buffer -> view<8xf16> + %d_f16 = scalar.fptrunc %d : f32 to f16 + %s_f16 = scalar.fptrunc %s : f32 to f16 + %ds_index = index.mul %block_in_group, %c2 : index + %s_index = index.add %ds_index, %c1 : index + view.store %d_f16, %group_ds[%ds_index] : f16, view<8xf16> + view.store %s_f16, %group_ds[%s_index] : f16, view<8xf16> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + } + func.return +} + +// Fuses attention RMSNorm directly into the Q8_1 x4 row consumed by decode +// projections. The output has hidden_size / 128 physical groups of 144 bytes. +kernel.def target(@qwen3_moe_attention_prepare_gfx11_wave32) @qwen3_moe_attention_rmsnorm_quantize_q8_1_x4(%token_count: index) { + %c1 = index.constant 1 : index + %c256 = index.constant 256 : index + kernel.launch.config workgroups(%token_count, %c1, %c1) workgroup_size(%c256, %c1, %c1) : index +} launch(%token_count: index, %input: buffer, %weight: buffer, %output: buffer) where [range(%token_count, 1, 2048)] { + %publish_normalized = scalar.constant false : i1 + %reduction_subgroup_count = index.constant 8 : index + %token = kernel.workgroup.id : index + %input_noalias, %weight_noalias, %output_noalias = buffer.assume.noalias %input, %weight, %output : buffer, buffer, buffer + func.call @qwen3_moe_rmsnorm_quantize_q8_1_x4_body(%publish_normalized, %reduction_subgroup_count, %token_count, %token, %input_noalias, %weight_noalias, %output_noalias, %output_noalias) : (i1, index, index, index, buffer, buffer, buffer, buffer) + kernel.return +} + +// Publishes both the ordinary F32 RMSNorm row required by routing and the Q8_1 +// x4 row consumed by direct decode contractions. +kernel.def target(@qwen3_moe_attention_prepare_gfx11_wave32) @qwen3_moe_rmsnorm_f32_quantize_q8_1_x4(%token_count: index) { + %c1 = index.constant 1 : index + %c256 = index.constant 256 : index + kernel.launch.config workgroups(%token_count, %c1, %c1) workgroup_size(%c256, %c1, %c1) : index +} launch(%token_count: index, %input: buffer, %weight: buffer, %normalized_output: buffer, %q8_output: buffer) where [range(%token_count, 1, 2048)] { + %publish_normalized = scalar.constant true : i1 + %reduction_subgroup_count = index.constant 8 : index + %token = kernel.workgroup.id : index + %input_noalias, %weight_noalias, %normalized_output_noalias, %q8_output_noalias = buffer.assume.noalias %input, %weight, %normalized_output, %q8_output : buffer, buffer, buffer, buffer + func.call @qwen3_moe_rmsnorm_quantize_q8_1_x4_body(%publish_normalized, %reduction_subgroup_count, %token_count, %token, %input_noalias, %weight_noalias, %normalized_output_noalias, %q8_output_noalias) : (i1, index, index, index, buffer, buffer, buffer, buffer) + kernel.return +} + +// A uniform residual makes RMSNorm equal to the learned weight up to the +// configured epsilon. The nonuniform weights cross two physical Q8_1 groups, +// exercise signed values, and compare every packed metadata and payload byte +// against the ordinary RMSNorm-plus-packer path. +check.case public @qwen3_moe_attention_rmsnorm_quantize_q8_1_x4_differential_case { + %token_count = check.literal value(1) : index + %hidden_size = check.literal value(256) : index + %input = check.generate.fill value(2.0) : tensor<256xf32> + %weight = check.generate.iota offset(-1.0) step(0.0078125) : tensor<256xf32> + %normalized = check.generate.fill value(0.0) : tensor<256xf32> + %expected = check.generate.fill value(0) : tensor<288xi8> + %actual = check.generate.fill value(1) : tensor<288xi8> + %dual_normalized = check.generate.fill value(1.0) : tensor<256xf32> + %dual_q8 = check.generate.fill value(1) : tensor<288xi8> + kernel.launch @qwen3_moe_rmsnorm_f32[%token_count](%token_count, %input, %weight, %normalized) : [index](index, tensor<256xf32>, tensor<256xf32>, tensor<256xf32>) + kernel.launch @ggml_quantize_q8_1_x4_f32[%token_count, %hidden_size](%token_count, %hidden_size, %normalized, %expected) : [index, index](index, index, tensor<256xf32>, tensor<288xi8>) + kernel.launch @qwen3_moe_attention_rmsnorm_quantize_q8_1_x4[%token_count](%token_count, %input, %weight, %actual) : [index](index, tensor<256xf32>, tensor<256xf32>, tensor<288xi8>) + kernel.launch @qwen3_moe_rmsnorm_f32_quantize_q8_1_x4[%token_count](%token_count, %input, %weight, %dual_normalized, %dual_q8) : [index](index, tensor<256xf32>, tensor<256xf32>, tensor<256xf32>, tensor<288xi8>) + check.expect.equal actual(%actual) expected(%expected) : tensor<288xi8> + check.expect.close actual(%dual_normalized) expected(%normalized) atol(0.0) rtol(0.0) nan(same) : tensor<256xf32> + check.expect.equal actual(%dual_q8) expected(%expected) : tensor<288xi8> + check.return +} + +// Fourteen distinct production-width rows cross two complete 1024-element +// scratch stripes. Comparing every packed byte with the unfused path locks +// both stripe reuse and row-byte addressing across independent workgroups. +check.case public @qwen3_moe_attention_rmsnorm_quantize_q8_1_x4_multistripe_case { + %token_count = check.literal value(14) : index + %hidden_size = check.literal value(2048) : index + %input_seed = check.param.seed base(0x514d4f4550524550) count(1) : i64 + %input = check.generate.random.uniform seed(%input_seed) range(-1.0 to 1.0) : tensor<14x2048xf32> + %weight = check.generate.iota offset(-1.0) step(0.0009765625) : tensor<2048xf32> + %normalized = check.generate.fill value(0.0) : tensor<14x2048xf32> + %expected = check.generate.fill value(0) : tensor<14x2304xi8> + %actual = check.generate.fill value(1) : tensor<14x2304xi8> + %dual_normalized = check.generate.fill value(1.0) : tensor<14x2048xf32> + %dual_q8 = check.generate.fill value(1) : tensor<14x2304xi8> + kernel.launch @qwen3_moe_rmsnorm_f32[%token_count](%token_count, %input, %weight, %normalized) : [index](index, tensor<14x2048xf32>, tensor<2048xf32>, tensor<14x2048xf32>) + kernel.launch @ggml_quantize_q8_1_x4_f32[%token_count, %hidden_size](%token_count, %hidden_size, %normalized, %expected) : [index, index](index, index, tensor<14x2048xf32>, tensor<14x2304xi8>) + kernel.launch @qwen3_moe_attention_rmsnorm_quantize_q8_1_x4[%token_count](%token_count, %input, %weight, %actual) : [index](index, tensor<14x2048xf32>, tensor<2048xf32>, tensor<14x2304xi8>) + kernel.launch @qwen3_moe_rmsnorm_f32_quantize_q8_1_x4[%token_count](%token_count, %input, %weight, %dual_normalized, %dual_q8) : [index](index, tensor<14x2048xf32>, tensor<2048xf32>, tensor<14x2048xf32>, tensor<14x2304xi8>) + check.expect.equal actual(%actual) expected(%expected) : tensor<14x2304xi8> + check.expect.close actual(%dual_normalized) expected(%normalized) atol(0.0) rtol(0.0) nan(same) : tensor<14x2048xf32> + check.expect.equal actual(%dual_q8) expected(%expected) : tensor<14x2304xi8> + check.return +} + +check.case public @qwen3_moe_attention_rmsnorm_f32_benchmark_case { + %token_count = check.param.choice values([1, 8, 32, 128, 512]) name("token_count") : index + %input = check.generate.fill value(0.0) : tensor<[%token_count]x2048xf32> + %weight = check.generate.fill value(1.0) : tensor<2048xf32> + %output = check.generate.fill value(1.0) : tensor<[%token_count]x2048xf32> + %expected = check.generate.fill value(0.0) : tensor<[%token_count]x2048xf32> + kernel.launch @qwen3_moe_rmsnorm_f32[%token_count](%token_count, %input, %weight, %output) : [index](index, tensor<[%token_count]x2048xf32>, tensor<2048xf32>, tensor<[%token_count]x2048xf32>) + check.expect.close actual(%output) expected(%expected) atol(0.0) rtol(0.0) nan(same) : tensor<[%token_count]x2048xf32> + check.return +} + +check.case public @qwen3_moe_attention_rmsnorm_quantize_q8_1_x4_benchmark_case { + %token_count = check.param.choice values([1, 8, 32, 128, 512]) name("token_count") : index + %input = check.generate.fill value(0.0) : tensor<[%token_count]x2048xf32> + %weight = check.generate.fill value(1.0) : tensor<2048xf32> + %output = check.generate.fill value(1) : tensor<[%token_count]x2304xi8> + %expected = check.generate.fill value(0) : tensor<[%token_count]x2304xi8> + kernel.launch @qwen3_moe_attention_rmsnorm_quantize_q8_1_x4[%token_count](%token_count, %input, %weight, %output) : [index](index, tensor<[%token_count]x2048xf32>, tensor<2048xf32>, tensor<[%token_count]x2304xi8>) + check.expect.equal actual(%output) expected(%expected) : tensor<[%token_count]x2304xi8> + check.return +} + +check.case public @qwen3_moe_rmsnorm_f32_quantize_q8_1_x4_benchmark_case { + %token_count = check.param.choice values([1, 8, 32, 128, 512]) name("token_count") : index + %input = check.generate.fill value(0.0) : tensor<[%token_count]x2048xf32> + %weight = check.generate.fill value(1.0) : tensor<2048xf32> + %normalized_output = check.generate.fill value(1.0) : tensor<[%token_count]x2048xf32> + %q8_output = check.generate.fill value(1) : tensor<[%token_count]x2304xi8> + %expected_normalized = check.generate.fill value(0.0) : tensor<[%token_count]x2048xf32> + %expected_q8 = check.generate.fill value(0) : tensor<[%token_count]x2304xi8> + kernel.launch @qwen3_moe_rmsnorm_f32_quantize_q8_1_x4[%token_count](%token_count, %input, %weight, %normalized_output, %q8_output) : [index](index, tensor<[%token_count]x2048xf32>, tensor<2048xf32>, tensor<[%token_count]x2048xf32>, tensor<[%token_count]x2304xi8>) + check.expect.close actual(%normalized_output) expected(%expected_normalized) atol(0.0) rtol(0.0) nan(same) : tensor<[%token_count]x2048xf32> + check.expect.equal actual(%q8_output) expected(%expected_q8) : tensor<[%token_count]x2304xi8> + check.return +} + +check.case public @qwen3_moe_attention_rmsnorm_then_quantize_q8_1_x4_benchmark_case { + %token_count = check.param.choice values([1, 8, 32, 128, 512]) name("token_count") : index + %hidden_size = check.literal value(2048) : index + %input = check.generate.fill value(0.0) : tensor<[%token_count]x2048xf32> + %weight = check.generate.fill value(1.0) : tensor<2048xf32> + %normalized = check.generate.fill value(1.0) : tensor<[%token_count]x2048xf32> + %output = check.generate.fill value(1) : tensor<[%token_count]x2304xi8> + %expected = check.generate.fill value(0) : tensor<[%token_count]x2304xi8> + kernel.launch @qwen3_moe_rmsnorm_f32[%token_count](%token_count, %input, %weight, %normalized) : [index](index, tensor<[%token_count]x2048xf32>, tensor<2048xf32>, tensor<[%token_count]x2048xf32>) + kernel.launch @ggml_quantize_q8_1_x4_f32[%token_count, %hidden_size](%token_count, %hidden_size, %normalized, %output) : [index, index](index, index, tensor<[%token_count]x2048xf32>, tensor<[%token_count]x2304xi8>) + check.expect.equal actual(%output) expected(%expected) : tensor<[%token_count]x2304xi8> + check.return +} + +check.benchmark<@qwen3_moe_attention_rmsnorm_quantize_q8_1_x4_differential_case> @qwen3_moe_attention_rmsnorm_quantize_q8_1_x4_differential + +check.benchmark<@qwen3_moe_attention_rmsnorm_f32_benchmark_case> @qwen3_moe_attention_rmsnorm_f32_decode {token_count = 1} + +check.benchmark<@qwen3_moe_attention_rmsnorm_f32_benchmark_case> @qwen3_moe_attention_rmsnorm_f32_prefill_32 {token_count = 32} + +check.benchmark<@qwen3_moe_attention_rmsnorm_f32_benchmark_case> @qwen3_moe_attention_rmsnorm_f32_prefill_128 {token_count = 128} + +check.benchmark<@qwen3_moe_attention_rmsnorm_f32_benchmark_case> @qwen3_moe_attention_rmsnorm_f32_prefill_512 {token_count = 512} + +check.benchmark<@qwen3_moe_attention_rmsnorm_quantize_q8_1_x4_benchmark_case> @qwen3_moe_attention_rmsnorm_quantize_q8_1_x4_decode {token_count = 1} + +check.benchmark<@qwen3_moe_rmsnorm_f32_quantize_q8_1_x4_benchmark_case> @qwen3_moe_rmsnorm_f32_quantize_q8_1_x4_decode {token_count = 1} + +check.benchmark<@qwen3_moe_attention_rmsnorm_quantize_q8_1_x4_benchmark_case> @qwen3_moe_attention_rmsnorm_quantize_q8_1_x4_prefill_32 {token_count = 32} + +check.benchmark<@qwen3_moe_attention_rmsnorm_quantize_q8_1_x4_benchmark_case> @qwen3_moe_attention_rmsnorm_quantize_q8_1_x4_prefill_128 {token_count = 128} + +check.benchmark<@qwen3_moe_attention_rmsnorm_quantize_q8_1_x4_benchmark_case> @qwen3_moe_attention_rmsnorm_quantize_q8_1_x4_prefill_512 {token_count = 512} + +check.benchmark<@qwen3_moe_attention_rmsnorm_then_quantize_q8_1_x4_benchmark_case> @qwen3_moe_attention_rmsnorm_then_quantize_q8_1_x4_decode {token_count = 1} + +check.benchmark<@qwen3_moe_attention_rmsnorm_then_quantize_q8_1_x4_benchmark_case> @qwen3_moe_attention_rmsnorm_then_quantize_q8_1_x4_prefill_32 {token_count = 32} + +check.benchmark<@qwen3_moe_attention_rmsnorm_then_quantize_q8_1_x4_benchmark_case> @qwen3_moe_attention_rmsnorm_then_quantize_q8_1_x4_prefill_128 {token_count = 128} + +check.benchmark<@qwen3_moe_attention_rmsnorm_then_quantize_q8_1_x4_benchmark_case> @qwen3_moe_attention_rmsnorm_then_quantize_q8_1_x4_prefill_512 {token_count = 512} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/attention_qkv_postprocess_fused.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/attention_qkv_postprocess_fused.loom new file mode 100644 index 000000000000..405e8dd69272 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/attention_qkv_postprocess_fused.loom @@ -0,0 +1,827 @@ +// Decode-only Qwen Q/K/V projection with last-arrival head postprocessing. +// +// The accepted quantized projection publishes eight adjacent raw F32 rows per +// 256-workitem workgroup. Head width is fixed at 128 for this model, so sixteen +// workgroups publish each Q, K, or V head. One device-scope completion counter +// per head forms a release sequence over those stores. The last arrival +// acquires the complete raw row and invokes the canonical per-head +// normalization, NEOX rotary, and cache-publication body. +// +// Raw Q/K/V rows remain explicit because they are the inter-workgroup +// communication surface. The fusion removes the command-buffer boundary and +// lets the postprocess consume freshly published rows; it does not impose a +// hidden cache-index or position relationship. Counters return to zero only +// after semantic outputs are visible, preserving reusable command buffers. +config.decl @qwen3_moe.model.hidden_size : %value: index where [range(%value, 128, 32768), mul(%value, 128)] +config.decl @qwen3_moe.model.rms_epsilon : f32 + +func.def inline @ggml_q6k_q8_1_x4_block_lane(%weight: buffer, %weight_row_byte_base: offset, %q8_input: buffer, %q8_row_byte_base: offset, %q6_block: index, %lane: index) -> (f32) { + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c4 = index.constant 4 : index + %c8 = index.constant 8 : index + %c16 = index.constant 16 : index + %block_bytes = index.constant 210 : offset + %qh_byte_add = index.constant 128 : offset + %scale_byte_add = index.constant 192 : offset + %d_byte_add = index.constant 208 : offset + %c4_i32v = vector.constant 4 : vector<1xi32> + %c0_i32v = vector.constant 0 : vector<1xi32> + %nibble_mask = vector.constant 252645135 : vector<1xi32> + %high_mask = vector.constant 808464432 : vector<1xi32> + %c0_f32 = scalar.constant 0.0 : f32 + %bounded_lane = index.assume %lane [range(%lane, 0, 31)] : index + %block_byte_add = index.scale %q6_block, %block_bytes : index, offset -> offset + %block_byte_base = index.add %weight_row_byte_base, %block_byte_add : offset + %qh_byte_base = index.add %block_byte_base, %qh_byte_add : offset + %scale_byte_base = index.add %block_byte_base, %scale_byte_add : offset + %d_byte_base = index.add %block_byte_base, %d_byte_add : offset + %ql_view = buffer.view %weight[%block_byte_base] : buffer -> view<32xi32> + %qh_view = buffer.view %weight[%qh_byte_base] : buffer -> view<16xi32> + %scale_view = buffer.view %weight[%scale_byte_base] : buffer -> view<16xi8> + %d_view = buffer.view %weight[%d_byte_base] : buffer -> view<1xf16> + %lane_mod8 = index.rem %bounded_lane, %c8 : index + %lane_mod16 = index.rem %bounded_lane, %c16 : index + %lane_div16 = index.div %bounded_lane, %c16 : index + %lane_div8_in_16 = index.div %lane_mod16, %c8 : index + %lane_div4_in_16 = index.div %lane_mod16, %c4 : index + %qh_high_base = index.mul %lane_div16, %c8 : index + %qh_index0 = index.add %qh_high_base, %lane_mod8 : index + %qh_index = index.assume %qh_index0 [range(%qh_index0, 0, 15)] : index + %ql_word = vector.load %ql_view[%bounded_lane] : view<32xi32> -> vector<1xi32> + %qh_word = vector.load %qh_view[%qh_index] : view<16xi32> -> vector<1xi32> + %qh_base_shift_index = index.mul %lane_div8_in_16, %c2 : index + %qh_base_shift_i32 = index.cast %qh_base_shift_index : index to i32 + %q8_block_base = index.mul %q6_block, %c8 : index + %q8_high_add = index.mul %lane_div16, %c4 : index + %q8_quadrant = index.add %q8_high_add, %lane_div8_in_16 : index + %scale_high_base = index.mul %lane_div16, %c8 : index + %scale_lane0 = index.add %scale_high_base, %lane_div4_in_16 : index + %d_f16 = view.load %d_view[0] : view<1xf16> -> f16 + %d = scalar.extf %d_f16 : f16 to f32 + %sum = scf.for %part = [%c0 to %c2 step %c1](%accumulator = %c0_f32 : f32) -> (f32) unroll { + %bounded_part = index.assume %part [range(%part, 0, 1)] : index + %part_shift_index = index.mul %bounded_part, %c4 : index + %part_shift_i32 = index.cast %part_shift_index : index to i32 + %part_shift = vector.splat %part_shift_i32 : vector<1xi32> + %ql_shifted = vector.shrui %ql_word, %part_shift : vector<1xi32> + %ql = vector.andi %ql_shifted, %nibble_mask : vector<1xi32> + %qh_shift_i32 = scalar.addi %qh_base_shift_i32, %part_shift_i32 : i32 + %qh_shift = vector.splat %qh_shift_i32 : vector<1xi32> + %qh_shifted = vector.shrui %qh_word, %qh_shift : vector<1xi32> + %qh_positioned = vector.shli %qh_shifted, %c4_i32v : vector<1xi32> + %qh = vector.andi %qh_positioned, %high_mask : vector<1xi32> + %code = vector.ori %ql, %qh : vector<1xi32> + %signed_weight = func.call @ggml_q6k_sign_extend_dot4(%code) : (vector<1xi32>) -> (vector<4xi8>) + %q8_part_add = index.mul %bounded_part, %c2 : index + %q8_block_part = index.add %q8_block_base, %q8_part_add : index + %q8_block = index.add %q8_block_part, %q8_quadrant : index + %q8_values, %q8_d = func.call @ggml_q8_1_x4_word(%q8_input, %q8_row_byte_base, %q8_block, %lane_mod8) : (buffer, offset, index, index) -> (vector<4xi8>, f32) + %scale_lane = index.add %scale_lane0, %part_shift_index : index + %scale_i8 = view.load %scale_view[%scale_lane] : view<16xi8> -> i8 + %scale = scalar.sitofp %scale_i8 : i8 to f32 + %dot = vector.dot4i %signed_weight, %q8_values, %c0_i32v : vector<4xi8>, vector<4xi8>, vector<1xi32> + %dot_i32 = vector.extract %dot[0] : vector<1xi32> -> i32 + %dot_f32 = scalar.sitofp %dot_i32 : i32 to f32 + %scaled0 = scalar.mulf %dot_f32, %scale : f32 + %scaled1 = scalar.mulf %scaled0, %d : f32 + %contribution = scalar.mulf %scaled1, %q8_d : f32 + %next = scalar.addf %accumulator, %contribution : f32 + scf.yield %next : f32 + } + func.return %sum : f32 +} + +func.def inline @ggml_q6k_q8_1_x4_row_lane(%input_size: index, %weight: buffer, %weight_row_byte_base: offset, %q8_input: buffer, %q8_row_byte_base: offset, %lane: index) -> (f32) { + %bounded_input_size = index.assume %input_size [range(%input_size, 256, 32768), mul(%input_size, 256)] : index + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c256 = index.constant 256 : index + %c0_f32 = scalar.constant 0.0 : f32 + %block_count = index.div %bounded_input_size, %c256 : index + %result = scf.for %block = [%c0 to %block_count step %c1](%block_acc = %c0_f32 : f32) -> (f32) { + %contribution = func.call @ggml_q6k_q8_1_x4_block_lane(%weight, %weight_row_byte_base, %q8_input, %q8_row_byte_base, %block, %lane) : (buffer, offset, buffer, offset, index, index) -> (f32) + %next = scalar.addf %block_acc, %contribution : f32 + scf.yield %next : f32 + } + func.return %result : f32 +} + +func.def inline @ggml_q6k_sign_extend_dot4(%code: vector<1xi32>) -> (vector<4xi8>) { + %c1_i32v = vector.constant 1 : vector<1xi32> + %c2_i32v = vector.constant 2 : vector<1xi32> + %low5_mask = vector.constant 522133279 : vector<1xi32> + %bit5_mask = vector.constant 538976288 : vector<1xi32> + %sign_mask = vector.constant -522133280 : vector<1xi32> + %low5 = vector.andi %code, %low5_mask : vector<1xi32> + %bit5 = vector.andi %code, %bit5_mask : vector<1xi32> + %bit6 = vector.shli %bit5, %c1_i32v : vector<1xi32> + %bit7 = vector.shli %bit5, %c2_i32v : vector<1xi32> + %high01 = vector.ori %bit5, %bit6 : vector<1xi32> + %high = vector.ori %high01, %bit7 : vector<1xi32> + %sign = vector.xori %high, %sign_mask : vector<1xi32> + %signed_i32 = vector.ori %low5, %sign : vector<1xi32> + %signed = vector.bitcast %signed_i32 : vector<1xi32> to vector<4xi8> + func.return %signed : vector<4xi8> +} + +func.def inline @ggml_q8_1_x4_word(%q8_input: buffer, %row_byte_base: offset, %q8_block: index, %word_in_block: index) -> (vector<4xi8>, f32) { + %c2 = index.constant 2 : index + %c4 = index.constant 4 : index + %c8 = index.constant 8 : index + %group_bytes = index.constant 144 : offset + %payload_byte_add = index.constant 16 : offset + %group = index.div %q8_block, %c4 : index + %inner0 = index.rem %q8_block, %c4 : index + %inner = index.assume %inner0 [range(%inner0, 0, 3)] : index + %word0 = index.assume %word_in_block [range(%word_in_block, 0, 7)] : index + %group_byte_add = index.scale %group, %group_bytes : index, offset -> offset + %group_byte_base = index.add %row_byte_base, %group_byte_add : offset + %payload_byte_base = index.add %group_byte_base, %payload_byte_add : offset + %ds_view = buffer.view %q8_input[%group_byte_base] : buffer -> view<8xf16> + %payload_view = buffer.view %q8_input[%payload_byte_base] : buffer -> view<32xi32> + %d_index = index.mul %inner, %c2 : index + %inner_word_base = index.mul %inner, %c8 : index + %word_index = index.add %inner_word_base, %word0 : index + %d_f16 = view.load %ds_view[%d_index] : view<8xf16> -> f16 + %packed = vector.load %payload_view[%word_index] : view<32xi32> -> vector<1xi32> + %values = vector.bitcast %packed : vector<1xi32> to vector<4xi8> + %d = scalar.extf %d_f16 : f16 to f32 + func.return %values, %d : vector<4xi8>, f32 +} + +func.def inline @qwen3_moe_attention_postprocess_head_body(%publish_output: i1, %token_count: index, %cache_row_count: index, %head_domain0: index, %token0: index, %positions: buffer, %key_cache_indices: buffer, %value_cache_indices: buffer, %query_input: buffer, %key_input: buffer, %value_input: buffer, %query_norm_weight: buffer, %key_norm_weight: buffer, %inverse_frequencies: buffer, %query_output: buffer, %key_cache: buffer, %value_cache: buffer) { + %query_size0 = config.get @qwen3_moe.attention.query_size : index + %key_value_size0 = config.get @qwen3_moe.attention.key_value_size : index + %head_size0 = config.get @qwen3_moe.attention.head_size : index + %epsilon = config.get @qwen3_moe.model.rms_epsilon : f32 + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 2048)] : index + %bounded_cache_row_count = index.assume %cache_row_count [range(%cache_row_count, 1, 1048576)] : index + %query_size, %key_value_size, %head_size = index.assume %query_size0, %key_value_size0, %head_size0 [range(%query_size0, 1, 262144), range(%key_value_size0, 1, 262144), range(%head_size0, 4, 1024), mul(%head_size0, 4), mul(%query_size0, %head_size0), mul(%key_value_size0, %head_size0)] : index, index, index + %channel0 = kernel.workitem.id : index + %c0 = index.constant 0 : index + %c2 = index.constant 2 : index + %c4 = index.constant 4 : index + %c0_f32 = scalar.constant 0.0 : f32 + %c0_offset = index.constant 0 : offset + %token, %launch_token_count = index.assume %token0, %bounded_token_count [lt(%token0, %bounded_token_count)] : index, index + %half_head_size = index.div %head_size, %c2 : index + %pair_packet_count = index.div %head_size, %c4 : index + %query_head_count = index.div %query_size, %head_size : index + %key_value_head_count = index.div %key_value_size, %head_size : index + %key_value_domain_count = index.mul %key_value_head_count, %c2 : index + %head_domain_count = index.add %query_head_count, %key_value_domain_count : index + %head_domain = index.assume %head_domain0 [lt(%head_domain0, %head_domain_count)] : index + %key_domain_end = index.add %query_head_count, %key_value_head_count : index + %is_query = index.cmp ult, %head_domain, %query_head_count : index + %is_query_or_key = index.cmp ult, %head_domain, %key_domain_end : index + %key_value_head = index.rem %head_domain, %key_value_head_count : index + %active_channel = index.cmp ult, %channel0, %pair_packet_count : index + %head_size_i32 = index.cast %head_size : index to i32 + %head_size_f32 = scalar.sitofp %head_size_i32 : i32 to f32 + %positions_noalias, %key_cache_indices_noalias, %value_cache_indices_noalias, %query_input_noalias, %key_input_noalias, %value_input_noalias, %query_norm_weight_noalias, %key_norm_weight_noalias, %inverse_frequencies_noalias, %query_output_noalias, %key_cache_noalias, %value_cache_noalias = buffer.assume.noalias %positions, %key_cache_indices, %value_cache_indices, %query_input, %key_input, %value_input, %query_norm_weight, %key_norm_weight, %inverse_frequencies, %query_output, %key_cache, %value_cache : buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer + %positions_view = buffer.view %positions_noalias[%c0_offset] : buffer -> view<[%launch_token_count]xi32> + %key_cache_indices_view = buffer.view %key_cache_indices_noalias[%c0_offset] : buffer -> view<[%launch_token_count]xi64> + %value_cache_indices_view = buffer.view %value_cache_indices_noalias[%c0_offset] : buffer -> view<[%launch_token_count]xi64> + %query_input_view = buffer.view %query_input_noalias[%c0_offset] : buffer -> view<[%launch_token_count]x[%query_head_count]x[%head_size]xf32> + %key_input_view = buffer.view %key_input_noalias[%c0_offset] : buffer -> view<[%launch_token_count]x[%key_value_head_count]x[%head_size]xf32> + %value_input_view = buffer.view %value_input_noalias[%c0_offset] : buffer -> view<[%launch_token_count]x[%key_value_head_count]x[%head_size]xf32> + %query_norm_weight_view = buffer.view %query_norm_weight_noalias[%c0_offset] : buffer -> view<[%head_size]xf32> + %key_norm_weight_view = buffer.view %key_norm_weight_noalias[%c0_offset] : buffer -> view<[%head_size]xf32> + %inverse_frequencies_view = buffer.view %inverse_frequencies_noalias[%c0_offset] : buffer -> view<[%half_head_size]xf32> + %query_output_view = buffer.view %query_output_noalias[%c0_offset] : buffer -> view<[%launch_token_count]x[%query_head_count]x[%head_size]xf32> + %key_cache_words_view = buffer.view %key_cache_noalias[%c0_offset] : buffer -> view<[%bounded_cache_row_count]x[%key_value_head_count]x[%half_head_size]xi32> + %value_cache_words_view = buffer.view %value_cache_noalias[%c0_offset] : buffer -> view<[%bounded_cache_row_count]x[%key_value_head_count]x[%half_head_size]xi32> + %row_sum = scf.if %is_query_or_key -> (f32) { + %partial_sum = scf.if %publish_output -> (f32) { + %channel_sum = scf.if %active_channel -> (f32) { + %channel = index.assume %channel0 [lt(%channel0, %pair_packet_count)] : index + %reduction_channel = index.mul %channel, %c4 : index + %reduction_values = scf.if %is_query -> (vector<4xf32>) { + %query_values = vector.load %query_input_view[%token, %head_domain, %reduction_channel] : view<[%launch_token_count]x[%query_head_count]x[%head_size]xf32> -> vector<4xf32> + scf.yield %query_values : vector<4xf32> + } else { + %key_values = vector.load %key_input_view[%token, %key_value_head, %reduction_channel] : view<[%launch_token_count]x[%key_value_head_count]x[%head_size]xf32> -> vector<4xf32> + scf.yield %key_values : vector<4xf32> + } + %squares = vector.mulf %reduction_values, %reduction_values : vector<4xf32> + %sum = vector.reduce %squares, %c0_f32 : vector<4xf32>, f32 + scf.yield %sum : f32 + } else { + scf.yield %c0_f32 : f32 + } + scf.yield %channel_sum : f32 + } else { + scf.yield %c0_f32 : f32 + } + %sum = kernel.workgroup.reduce %partial_sum : f32 + scf.yield %sum : f32 + } else { + scf.yield %c0_f32 : f32 + } + scf.if %publish_output { + scf.if %active_channel { + %channel = index.assume %channel0 [lt(%channel0, %pair_packet_count)] : index + %pair_channel = index.mul %channel, %c2 : index + %paired_channel = index.add %pair_channel, %half_head_size : index + %paired_word = index.add %channel, %pair_packet_count : index + %reduction_channel = index.mul %channel, %c4 : index + scf.if %is_query { + %low_values = vector.load %query_input_view[%token, %head_domain, %pair_channel] : view<[%launch_token_count]x[%query_head_count]x[%head_size]xf32> -> vector<2xf32> + %high_values = vector.load %query_input_view[%token, %head_domain, %paired_channel] : view<[%launch_token_count]x[%query_head_count]x[%head_size]xf32> -> vector<2xf32> + %low_weights = vector.load %query_norm_weight_view[%pair_channel] : view<[%head_size]xf32> -> vector<2xf32> + %high_weights = vector.load %query_norm_weight_view[%paired_channel] : view<[%head_size]xf32> -> vector<2xf32> + %inverse_frequencies_packet = vector.load %inverse_frequencies_view[%pair_channel] : view<[%half_head_size]xf32> -> vector<2xf32> + %position_i32 = view.load %positions_view[%token] : view<[%launch_token_count]xi32> -> i32 + %position = scalar.sitofp %position_i32 : i32 to f32 + %rotated_low, %rotated_high = func.call @qwen3_moe_rmsnorm_neox_packet(%row_sum, %head_size_f32, %epsilon, %position, %inverse_frequencies_packet, %low_values, %high_values, %low_weights, %high_weights) : (f32, f32, f32, f32, vector<2xf32>, vector<2xf32>, vector<2xf32>, vector<2xf32>, vector<2xf32>) -> (vector<2xf32>, vector<2xf32>) + vector.store %rotated_low, %query_output_view[%token, %head_domain, %pair_channel] : vector<2xf32>, view<[%launch_token_count]x[%query_head_count]x[%head_size]xf32> + vector.store %rotated_high, %query_output_view[%token, %head_domain, %paired_channel] : vector<2xf32>, view<[%launch_token_count]x[%query_head_count]x[%head_size]xf32> + } else { + scf.if %is_query_or_key { + %low_values = vector.load %key_input_view[%token, %key_value_head, %pair_channel] : view<[%launch_token_count]x[%key_value_head_count]x[%head_size]xf32> -> vector<2xf32> + %high_values = vector.load %key_input_view[%token, %key_value_head, %paired_channel] : view<[%launch_token_count]x[%key_value_head_count]x[%head_size]xf32> -> vector<2xf32> + %low_weights = vector.load %key_norm_weight_view[%pair_channel] : view<[%head_size]xf32> -> vector<2xf32> + %high_weights = vector.load %key_norm_weight_view[%paired_channel] : view<[%head_size]xf32> -> vector<2xf32> + %inverse_frequencies_packet = vector.load %inverse_frequencies_view[%pair_channel] : view<[%half_head_size]xf32> -> vector<2xf32> + %position_i32 = view.load %positions_view[%token] : view<[%launch_token_count]xi32> -> i32 + %position = scalar.sitofp %position_i32 : i32 to f32 + %cache_index_raw = view.load %key_cache_indices_view[%token] : view<[%launch_token_count]xi64> -> i64 + %cache_index_i64 = scalar.assume %cache_index_raw [range(%cache_index_raw, 0, 1048575)] : i64 + %cache_index0 = index.cast %cache_index_i64 : i64 to index + %cache_index = index.assume %cache_index0 [lt(%cache_index0, %bounded_cache_row_count)] : index + %rotated_low, %rotated_high = func.call @qwen3_moe_rmsnorm_neox_packet(%row_sum, %head_size_f32, %epsilon, %position, %inverse_frequencies_packet, %low_values, %high_values, %low_weights, %high_weights) : (f32, f32, f32, f32, vector<2xf32>, vector<2xf32>, vector<2xf32>, vector<2xf32>, vector<2xf32>) -> (vector<2xf32>, vector<2xf32>) + %half_low = vector.fptrunc %rotated_low : vector<2xf32> to vector<2xf16> + %half_high = vector.fptrunc %rotated_high : vector<2xf32> to vector<2xf16> + %packed_low = vector.bitcast %half_low : vector<2xf16> to vector<1xi32> + %packed_high = vector.bitcast %half_high : vector<2xf16> to vector<1xi32> + vector.store %packed_low, %key_cache_words_view[%cache_index, %key_value_head, %channel] : vector<1xi32>, view<[%bounded_cache_row_count]x[%key_value_head_count]x[%half_head_size]xi32> + vector.store %packed_high, %key_cache_words_view[%cache_index, %key_value_head, %paired_word] : vector<1xi32>, view<[%bounded_cache_row_count]x[%key_value_head_count]x[%half_head_size]xi32> + } else { + %cache_index_raw = view.load %value_cache_indices_view[%token] : view<[%launch_token_count]xi64> -> i64 + %cache_index_i64 = scalar.assume %cache_index_raw [range(%cache_index_raw, 0, 1048575)] : i64 + %cache_index0 = index.cast %cache_index_i64 : i64 to index + %cache_index = index.assume %cache_index0 [lt(%cache_index0, %bounded_cache_row_count)] : index + %values = vector.load %value_input_view[%token, %key_value_head, %reduction_channel] : view<[%launch_token_count]x[%key_value_head_count]x[%head_size]xf32> -> vector<4xf32> + %half_values = vector.fptrunc %values : vector<4xf32> to vector<4xf16> + %packed_values = vector.bitcast %half_values : vector<4xf16> to vector<2xi32> + vector.store %packed_values, %value_cache_words_view[%cache_index, %key_value_head, %pair_channel] : vector<2xi32>, view<[%bounded_cache_row_count]x[%key_value_head_count]x[%half_head_size]xi32> + } + } + } + } + func.return +} + +func.def inline @qwen3_moe_attention_qkv_quantized_body(%value_uses_q6_index: index, %publish_output: i1, %token_count: index, %token0: index, %q8_input: buffer, %query_weight: buffer, %key_weight: buffer, %value_weight: buffer, %query_output: buffer, %key_output: buffer, %value_output: buffer) { + %hidden_size0 = config.get @qwen3_moe.model.hidden_size : index + %query_size0 = config.get @qwen3_moe.attention.query_size : index + %key_value_size0 = config.get @qwen3_moe.attention.key_value_size : index + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 2048)] : index + %hidden_size = index.assume %hidden_size0 [range(%hidden_size0, 128, 32768), mul(%hidden_size0, 128)] : index + %query_size, %key_value_size = index.assume %query_size0, %key_value_size0 [range(%query_size0, 1, 262144), range(%key_value_size0, 1, 262144), mul(%query_size0, %key_value_size0)] : index, index + %channel_tile = kernel.workgroup.id : index + %subgroup0 = kernel.subgroup.id : index + %lane = kernel.subgroup.lane.id : index + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c8 = index.constant 8 : index + %c128 = index.constant 128 : index + %c144_bytes = index.constant 144 : offset + %c210_bytes = index.constant 210 : offset + %c256 = index.constant 256 : index + %c0_i32 = scalar.constant 0 : i32 + %c0_f32 = scalar.constant 0.0 : f32 + %c0_offset = index.constant 0 : offset + %token, %launch_token_count = index.assume %token0, %bounded_token_count [lt(%token0, %bounded_token_count)] : index, index + %subgroup = index.assume %subgroup0 [range(%subgroup0, 0, 7)] : index + %channel_base = index.mul %channel_tile, %c8 : index + %global_channel = index.add %channel_base, %subgroup : index + %key_value_end = index.add %query_size, %key_value_size : index + %key_value_output_size = index.mul %key_value_size, %c2 : index + %total_output_size = index.add %query_size, %key_value_output_size : index + %valid_channel = index.cmp ult, %global_channel, %total_output_size : index + %is_query = index.cmp ult, %global_channel, %query_size : index + %is_key = index.cmp ult, %global_channel, %key_value_end : index + %key_value_channel = index.rem %global_channel, %key_value_size : index + %value_uses_q6 = index.cmp eq, %value_uses_q6_index, %c1 : index + %lane_i32 = index.cast %lane : index to i32 + %is_lane_zero = scalar.cmpi eq, %lane_i32, %c0_i32 : i32 + %quant_block_count = index.div %hidden_size, %c256 : index + %q4_row_bytes = index.scale %quant_block_count, %c144_bytes : index, offset -> offset + %q6_row_bytes = index.scale %quant_block_count, %c210_bytes : index, offset -> offset + %q8_group_count = index.div %hidden_size, %c128 : index + %q8_row_bytes = index.scale %q8_group_count, %c144_bytes : index, offset -> offset + %q8_row_byte_base = index.scale %token, %q8_row_bytes : index, offset -> offset + %q8_noalias, %query_weight_noalias, %key_weight_noalias, %value_weight_noalias, %query_output_noalias, %key_output_noalias, %value_output_noalias = buffer.assume.noalias %q8_input, %query_weight, %key_weight, %value_weight, %query_output, %key_output, %value_output : buffer, buffer, buffer, buffer, buffer, buffer, buffer + %lane_sum = scf.if %publish_output -> (f32) { + %channel_sum = scf.if %valid_channel -> (f32) { + %projection_sum = scf.if %is_query -> (f32) { + %row_byte_base = index.scale %global_channel, %q4_row_bytes : index, offset -> offset + %sum = func.call @qwen3_moe_q4k_q8_1_x4_paired_row_lane(%hidden_size, %query_weight_noalias, %row_byte_base, %q8_noalias, %q8_row_byte_base, %lane) : (index, buffer, offset, buffer, offset, index) -> (f32) + scf.yield %sum : f32 + } else { + %key_or_value_sum = scf.if %is_key -> (f32) { + %row_byte_base = index.scale %key_value_channel, %q4_row_bytes : index, offset -> offset + %sum = func.call @qwen3_moe_q4k_q8_1_x4_paired_row_lane(%hidden_size, %key_weight_noalias, %row_byte_base, %q8_noalias, %q8_row_byte_base, %lane) : (index, buffer, offset, buffer, offset, index) -> (f32) + scf.yield %sum : f32 + } else { + %value_sum = scf.if %value_uses_q6 -> (f32) { + %row_byte_base = index.scale %key_value_channel, %q6_row_bytes : index, offset -> offset + %sum = func.call @ggml_q6k_q8_1_x4_row_lane(%hidden_size, %value_weight_noalias, %row_byte_base, %q8_noalias, %q8_row_byte_base, %lane) : (index, buffer, offset, buffer, offset, index) -> (f32) + scf.yield %sum : f32 + } else { + %row_byte_base = index.scale %key_value_channel, %q4_row_bytes : index, offset -> offset + %sum = func.call @qwen3_moe_q4k_q8_1_x4_paired_row_lane(%hidden_size, %value_weight_noalias, %row_byte_base, %q8_noalias, %q8_row_byte_base, %lane) : (index, buffer, offset, buffer, offset, index) -> (f32) + scf.yield %sum : f32 + } + scf.yield %value_sum : f32 + } + scf.yield %key_or_value_sum : f32 + } + scf.yield %projection_sum : f32 + } else { + scf.yield %c0_f32 : f32 + } + scf.yield %channel_sum : f32 + } else { + scf.yield %c0_f32 : f32 + } + %dot = kernel.subgroup.reduce %lane_sum : f32 + %query_output_view = buffer.view %query_output_noalias[%c0_offset] : buffer -> view<[%launch_token_count]x[%query_size]xf32> + %key_output_view = buffer.view %key_output_noalias[%c0_offset] : buffer -> view<[%launch_token_count]x[%key_value_size]xf32> + %value_output_view = buffer.view %value_output_noalias[%c0_offset] : buffer -> view<[%launch_token_count]x[%key_value_size]xf32> + scf.if %publish_output { + scf.if %valid_channel { + scf.if %is_lane_zero { + scf.if %is_query { + view.store %dot, %query_output_view[%token, %global_channel] : f32, view<[%launch_token_count]x[%query_size]xf32> + } else { + scf.if %is_key { + view.store %dot, %key_output_view[%token, %key_value_channel] : f32, view<[%launch_token_count]x[%key_value_size]xf32> + } else { + view.store %dot, %value_output_view[%token, %key_value_channel] : f32, view<[%launch_token_count]x[%key_value_size]xf32> + } + } + } + } + } + func.return +} + +func.def inline @qwen3_moe_q4k_chunk_pair_global(%weight: buffer, %row_byte_base: offset, %q4_block: index, %q4_group_pair: index, %q4_half: index, %header_words: vector<4xi32>) -> (vector<16xi8>, f32, f32, vector<16xi8>, f32, f32) { + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c4 = index.constant 4 : index + %c8 = index.constant 8 : index + %block_bytes = index.constant 144 : offset + %code_byte_add = index.constant 16 : offset + %c4_i32 = scalar.constant 4 : i32 + %nibble_mask = vector.constant 252645135 : vector<4xi32> + %bounded_pair = index.assume %q4_group_pair [range(%q4_group_pair, 0, 3)] : index + %bounded_half = index.assume %q4_half [range(%q4_half, 0, 1)] : index + %block_byte_add = index.scale %q4_block, %block_bytes : index, offset -> offset + %block_byte_base = index.add %row_byte_base, %block_byte_add : offset + %code_byte_base = index.add %block_byte_base, %code_byte_add : offset + %code_view = buffer.view %weight[%code_byte_base] : buffer -> view<32xi32> + %header_halves = vector.bitcast %header_words : vector<4xi32> to vector<8xf16> + %d_f16 = vector.extract %header_halves[0] : vector<8xf16> -> f16 + %dmin_f16 = vector.extract %header_halves[1] : vector<8xf16> -> f16 + %scale0 = vector.extract %header_words[1] : vector<4xi32> -> i32 + %scale1 = vector.extract %header_words[2] : vector<4xi32> -> i32 + %scale2 = vector.extract %header_words[3] : vector<4xi32> -> i32 + %d = scalar.extf %d_f16 : f16 to f32 + %dmin = scalar.extf %dmin_f16 : f16 to f32 + %pair_code_base = index.mul %bounded_pair, %c8 : index + %half_code_add = index.mul %bounded_half, %c4 : index + %code_index0 = index.add %pair_code_base, %half_code_add : index + %code_index = index.assume %code_index0 [range(%code_index0, 0, 28)] : index + %packed_codes = vector.load %code_view[%code_index] : view<32xi32> -> vector<4xi32> + %low_codes = vector.andi %packed_codes, %nibble_mask : vector<4xi32> + %c4_i32v = vector.splat %c4_i32 : vector<4xi32> + %high_shifted = vector.shrui %packed_codes, %c4_i32v : vector<4xi32> + %high_codes = vector.andi %high_shifted, %nibble_mask : vector<4xi32> + %q4_low = vector.bitcast %low_codes : vector<4xi32> to vector<16xi8> + %q4_high = vector.bitcast %high_codes : vector<4xi32> to vector<16xi8> + %low_group = index.mul %bounded_pair, %c2 : index + %high_group = index.add %low_group, %c1 : index + %low_scale, %low_minimum = func.call @qwen3_moe_q4k_scale_from_header(%scale0, %scale1, %scale2, %low_group) : (i32, i32, i32, index) -> (i32, i32) + %high_scale, %high_minimum = func.call @qwen3_moe_q4k_scale_from_header(%scale0, %scale1, %scale2, %high_group) : (i32, i32, i32, index) -> (i32, i32) + %low_scale_f32 = scalar.uitofp %low_scale : i32 to f32 + %low_minimum_f32 = scalar.uitofp %low_minimum : i32 to f32 + %high_scale_f32 = scalar.uitofp %high_scale : i32 to f32 + %high_minimum_f32 = scalar.uitofp %high_minimum : i32 to f32 + %low_d_scale = scalar.mulf %d, %low_scale_f32 : f32 + %low_dmin_scale = scalar.mulf %dmin, %low_minimum_f32 : f32 + %high_d_scale = scalar.mulf %d, %high_scale_f32 : f32 + %high_dmin_scale = scalar.mulf %dmin, %high_minimum_f32 : f32 + func.return %q4_low, %low_d_scale, %low_dmin_scale, %q4_high, %high_d_scale, %high_dmin_scale : vector<16xi8>, f32, f32, vector<16xi8>, f32, f32 +} + +func.def inline @qwen3_moe_q4k_q8_1_dot(%q4_values: vector<16xi8>, %d_scale: f32, %dmin_scale: f32, %q8_values: vector<16xi8>, %q8_d: f32, %q8_s: f32) -> (f32) { + %c0_i32 = scalar.constant 0 : i32 + %c0_i32v = vector.constant 0 : vector<4xi32> + %half_f32 = scalar.constant 0.5 : f32 + %partial_dots = vector.dot4i %q4_values, %q8_values, %c0_i32v : vector<16xi8>, vector<16xi8>, vector<4xi32> + %q_sum = vector.reduce %partial_dots, %c0_i32 : vector<4xi32>, i32 + %q_sum_f32 = scalar.sitofp %q_sum : i32 to f32 + %scaled_dot0 = scalar.mulf %q8_d, %d_scale : f32 + %scaled_dot = scalar.mulf %scaled_dot0, %q_sum_f32 : f32 + %q8_half_sum = scalar.mulf %q8_s, %half_f32 : f32 + %minimum_correction = scalar.mulf %dmin_scale, %q8_half_sum : f32 + %contribution = scalar.subf %scaled_dot, %minimum_correction : f32 + func.return %contribution : f32 +} + +func.def inline @qwen3_moe_q4k_q8_1_x4_paired_block_lane(%input_size: index, %weight: buffer, %weight_row_byte_base: offset, %q8_input: buffer, %q8_row_byte_base: offset, %q4_block: index, %block_lane: index) -> (f32) { + %bounded_input_size = index.assume %input_size [range(%input_size, 256, 32768), mul(%input_size, 256)] : index + %bounded_q4_block0 = index.assume %q4_block [range(%q4_block, 0, 127)] : index + %bounded_block_lane = index.assume %block_lane [range(%block_lane, 0, 7)] : index + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c4 = index.constant 4 : index + %c8 = index.constant 8 : index + %c128 = index.constant 128 : index + %c256 = index.constant 256 : index + %q4_block_bytes = index.constant 144 : offset + %q8_group_bytes = index.constant 144 : offset + %q8_payload_byte_add = index.constant 16 : offset + %q4_block_count = index.div %bounded_input_size, %c256 : index + %q8_group_count = index.div %bounded_input_size, %c128 : index + %bounded_q4_block, %bounded_q4_block_count = index.assume %bounded_q4_block0, %q4_block_count [lt(%bounded_q4_block0, %q4_block_count)] : index, index + %q4_group_pair0 = index.div %bounded_block_lane, %c2 : index + %q4_group_pair = index.assume %q4_group_pair0 [range(%q4_group_pair0, 0, 3)] : index + %q4_half0 = index.rem %bounded_block_lane, %c2 : index + %q4_half = index.assume %q4_half0 [range(%q4_half0, 0, 1)] : index + %q8_group_in_block0 = index.div %q4_group_pair, %c2 : index + %q8_group_in_block = index.assume %q8_group_in_block0 [range(%q8_group_in_block0, 0, 1)] : index + %pair_in_q8_group0 = index.rem %q4_group_pair, %c2 : index + %pair_in_q8_group = index.assume %pair_in_q8_group0 [range(%pair_in_q8_group0, 0, 1)] : index + %q8_low_inner_block0 = index.mul %pair_in_q8_group, %c2 : index + %q8_low_inner_block = index.assume %q8_low_inner_block0 [range(%q8_low_inner_block0, 0, 2)] : index + %q8_high_inner_block0 = index.add %q8_low_inner_block, %c1 : index + %q8_high_inner_block = index.assume %q8_high_inner_block0 [range(%q8_high_inner_block0, 1, 3)] : index + %q8_half_word_add = index.mul %q4_half, %c4 : index + %q8_low_inner_word_base = index.mul %q8_low_inner_block, %c8 : index + %q8_low_word_index0 = index.add %q8_low_inner_word_base, %q8_half_word_add : index + %q8_low_word_index = index.assume %q8_low_word_index0 [range(%q8_low_word_index0, 0, 20)] : index + %q8_high_inner_word_base = index.mul %q8_high_inner_block, %c8 : index + %q8_high_word_index0 = index.add %q8_high_inner_word_base, %q8_half_word_add : index + %q8_high_word_index = index.assume %q8_high_word_index0 [range(%q8_high_word_index0, 8, 28)] : index + %q8_low_ds_index0 = index.mul %q8_low_inner_block, %c2 : index + %q8_low_ds_index = index.assume %q8_low_ds_index0 [range(%q8_low_ds_index0, 0, 4)] : index + %q8_block_group_base = index.mul %bounded_q4_block, %c2 : index + %q8_group0 = index.add %q8_block_group_base, %q8_group_in_block : index + %q8_group, %bounded_q8_group_count = index.assume %q8_group0, %q8_group_count [lt(%q8_group0, %q8_group_count)] : index, index + %q8_group_byte_add = index.scale %q8_group, %q8_group_bytes : index, offset -> offset + %q8_group_byte_base = index.add %q8_row_byte_base, %q8_group_byte_add : offset + %q8_payload_byte_base = index.add %q8_group_byte_base, %q8_payload_byte_add : offset + %q8_ds_view = buffer.view %q8_input[%q8_group_byte_base] : buffer -> view<8xf16> + %q8_words_view = buffer.view %q8_input[%q8_payload_byte_base] : buffer -> view<32xi32> + %q8_ds = vector.load %q8_ds_view[%q8_low_ds_index] : view<8xf16> -> vector<4xf16> + %q8_low_d_f16 = vector.extract %q8_ds[0] : vector<4xf16> -> f16 + %q8_low_s_f16 = vector.extract %q8_ds[1] : vector<4xf16> -> f16 + %q8_high_d_f16 = vector.extract %q8_ds[2] : vector<4xf16> -> f16 + %q8_high_s_f16 = vector.extract %q8_ds[3] : vector<4xf16> -> f16 + %q8_low_d = scalar.extf %q8_low_d_f16 : f16 to f32 + %q8_low_s = scalar.extf %q8_low_s_f16 : f16 to f32 + %q8_high_d = scalar.extf %q8_high_d_f16 : f16 to f32 + %q8_high_s = scalar.extf %q8_high_s_f16 : f16 to f32 + %q8_low_words = vector.load %q8_words_view[%q8_low_word_index] : view<32xi32> -> vector<4xi32> + %q8_high_words = vector.load %q8_words_view[%q8_high_word_index] : view<32xi32> -> vector<4xi32> + %q8_low_values = vector.bitcast %q8_low_words : vector<4xi32> to vector<16xi8> + %q8_high_values = vector.bitcast %q8_high_words : vector<4xi32> to vector<16xi8> + %q4_block_byte_add = index.scale %bounded_q4_block, %q4_block_bytes : index, offset -> offset + %q4_block_byte_base = index.add %weight_row_byte_base, %q4_block_byte_add : offset + %q4_header_view = buffer.view %weight[%q4_block_byte_base] : buffer -> view<4xi32> + %q4_header_words = vector.load %q4_header_view[0] : view<4xi32> -> vector<4xi32> + %q4_low, %low_d_scale, %low_dmin_scale, %q4_high, %high_d_scale, %high_dmin_scale = func.call @qwen3_moe_q4k_chunk_pair_global(%weight, %weight_row_byte_base, %bounded_q4_block, %q4_group_pair, %q4_half, %q4_header_words) : (buffer, offset, index, index, index, vector<4xi32>) -> (vector<16xi8>, f32, f32, vector<16xi8>, f32, f32) + %low = func.call @qwen3_moe_q4k_q8_1_dot(%q4_low, %low_d_scale, %low_dmin_scale, %q8_low_values, %q8_low_d, %q8_low_s) : (vector<16xi8>, f32, f32, vector<16xi8>, f32, f32) -> (f32) + %high = func.call @qwen3_moe_q4k_q8_1_dot(%q4_high, %high_d_scale, %high_dmin_scale, %q8_high_values, %q8_high_d, %q8_high_s) : (vector<16xi8>, f32, f32, vector<16xi8>, f32, f32) -> (f32) + %pair = scalar.addf %low, %high : f32 + func.return %pair : f32 +} + +func.def inline @qwen3_moe_q4k_q8_1_x4_paired_row_lane(%input_size: index, %weight: buffer, %weight_row_byte_base: offset, %q8_input: buffer, %q8_row_byte_base: offset, %lane: index) -> (f32) { + %bounded_input_size = index.assume %input_size [range(%input_size, 256, 32768), mul(%input_size, 256)] : index + %bounded_lane = index.assume %lane [range(%lane, 0, 31)] : index + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c4 = index.constant 4 : index + %c8 = index.constant 8 : index + %c256 = index.constant 256 : index + %c1023 = index.constant 1023 : index + %c1024 = index.constant 1024 : index + %c0_f32 = scalar.constant 0.0 : f32 + %q4_block_count = index.div %bounded_input_size, %c256 : index + %padded_input_size = index.add %bounded_input_size, %c1023 : index + %iteration_count = index.div %padded_input_size, %c1024 : index + %lane_q4_block = index.div %bounded_lane, %c8 : index + %block_lane0 = index.rem %bounded_lane, %c8 : index + %block_lane = index.assume %block_lane0 [range(%block_lane0, 0, 7)] : index + %sum = scf.for %iteration = [%c0 to %iteration_count step %c1](%iteration_acc = %c0_f32 : f32) -> (f32) unroll { + %iteration_q4_block = index.mul %iteration, %c4 : index + %q4_block0 = index.add %iteration_q4_block, %lane_q4_block : index + %valid_q4_block = index.cmp ult, %q4_block0, %q4_block_count : index + %contribution = scf.if %valid_q4_block -> (f32) { + %q4_block, %bounded_q4_block_count = index.assume %q4_block0, %q4_block_count [lt(%q4_block0, %q4_block_count)] : index, index + %pair = func.call @qwen3_moe_q4k_q8_1_x4_paired_block_lane(%bounded_input_size, %weight, %weight_row_byte_base, %q8_input, %q8_row_byte_base, %q4_block, %block_lane) : (index, buffer, offset, buffer, offset, index, index) -> (f32) + scf.yield %pair : f32 + } else { + scf.yield %c0_f32 : f32 + } + %next = scalar.addf %iteration_acc, %contribution : f32 + scf.yield %next : f32 + } + func.return %sum : f32 +} + +func.def inline @qwen3_moe_q4k_scale_from_header(%scale0: i32, %scale1: i32, %scale2: i32, %q4_group: index) -> (i32, i32) { + %c4 = index.constant 4 : index + %c8 = index.constant 8 : index + %c2_i32 = scalar.constant 2 : i32 + %c4_i32 = scalar.constant 4 : i32 + %c15_i32 = scalar.constant 15 : i32 + %c48_i32 = scalar.constant 48 : i32 + %bounded_group = index.assume %q4_group [range(%q4_group, 0, 7)] : index + %is_low_group = index.cmp ult, %bounded_group, %c4 : index + %scale_lane = index.rem %bounded_group, %c4 : index + %scale_shift_index = index.mul %scale_lane, %c8 : index + %scale_shift = index.cast %scale_shift_index : index to i32 + %high_shift = scalar.addi %scale_shift, %c2_i32 : i32 + %minimum_shift = scalar.addi %scale_shift, %c4_i32 : i32 + %selected_scale_source = scf.select %is_low_group, %scale0, %scale2 : i32 + %selected_minimum_source = scf.select %is_low_group, %scale1, %scale2 : i32 + %selected_scale_high_shift = scf.select %is_low_group, %scale_shift, %high_shift : i32 + %selected_minimum_low_shift = scf.select %is_low_group, %scale_shift, %minimum_shift : i32 + %scale_low0 = scalar.shrui %selected_scale_source, %scale_shift : i32 + %scale_low = scalar.andi %scale_low0, %c15_i32 : i32 + %scale_high0 = scalar.shrui %scale0, %selected_scale_high_shift : i32 + %scale_high = scalar.andi %scale_high0, %c48_i32 : i32 + %scale = scalar.ori %scale_low, %scale_high : i32 + %minimum_low0 = scalar.shrui %selected_minimum_source, %selected_minimum_low_shift : i32 + %minimum_low = scalar.andi %minimum_low0, %c15_i32 : i32 + %minimum_high0 = scalar.shrui %scale1, %selected_scale_high_shift : i32 + %minimum_high = scalar.andi %minimum_high0, %c48_i32 : i32 + %minimum = scalar.ori %minimum_low, %minimum_high : i32 + func.return %scale, %minimum : i32, i32 +} + +func.def inline @qwen3_moe_rmsnorm_neox_packet(%row_sum: f32, %head_size: f32, %epsilon: f32, %position: f32, %inverse_frequencies: vector<2xf32>, %low_values: vector<2xf32>, %high_values: vector<2xf32>, %low_weights: vector<2xf32>, %high_weights: vector<2xf32>) -> (vector<2xf32>, vector<2xf32>) { + %mean = scalar.divf %row_sum, %head_size : f32 + %biased_mean = scalar.addf %mean, %epsilon : f32 + %scale = scalar.rsqrtf %biased_mean : f32 + %scale_vector = vector.splat %scale : vector<2xf32> + %position_vector = vector.splat %position : vector<2xf32> + %inverse_two_pi = scalar.constant 0.15915494309189535 : f32 + %inverse_two_pi_vector = vector.splat %inverse_two_pi : vector<2xf32> + %normalized_low = vector.mulf %low_values, %scale_vector : vector<2xf32> + %normalized_high = vector.mulf %high_values, %scale_vector : vector<2xf32> + %scaled_low = vector.mulf %normalized_low, %low_weights : vector<2xf32> + %scaled_high = vector.mulf %normalized_high, %high_weights : vector<2xf32> + %angles = vector.mulf %position_vector, %inverse_frequencies : vector<2xf32> + %turns = vector.mulf %angles, %inverse_two_pi_vector : vector<2xf32> + %cosines = vector.costurnsf %turns : vector<2xf32> + %sines = vector.sinturnsf %turns : vector<2xf32> + %low_cosines = vector.mulf %scaled_low, %cosines : vector<2xf32> + %high_sines = vector.mulf %scaled_high, %sines : vector<2xf32> + %low_sines = vector.mulf %scaled_low, %sines : vector<2xf32> + %high_cosines = vector.mulf %scaled_high, %cosines : vector<2xf32> + %rotated_low = vector.subf %low_cosines, %high_sines : vector<2xf32> + %rotated_high = vector.addf %low_sines, %high_cosines : vector<2xf32> + func.return %rotated_low, %rotated_high : vector<2xf32>, vector<2xf32> +} + +amdgpu.target @qwen3_moe_attention_qkv_postprocess_gfx11_wave32 {subgroup_size = 32} + +config.decl @qwen3_moe.attention.head_size : %value: index where [range(%value, 4, 1024), mul(%value, 4)] + +config.decl @qwen3_moe.attention.query_size : %value: index where [range(%value, 1, 262144)] + +config.decl @qwen3_moe.attention.key_value_size : %value: index where [range(%value, 1, 262144)] + +config.decl @qwen3_moe.attention.value_uses_q6 : %value: index where [range(%value, 0, 1)] + +// Reference entry points used by differential and benchmark cases. +kernel.decl @ggml_quantize_q8_1_x4_f32(%token_count: index, %input_size: index) launch(%token_count: index, %input_size: index, %input: buffer, %output: buffer) + +kernel.decl @qwen3_moe_attention_qkv_quantized(%token_count: index) launch(%token_count: index, %q8_input: buffer, %query_weight: buffer, %key_weight: buffer, %value_weight: buffer, %query_output: buffer, %key_output: buffer, %value_output: buffer) + +kernel.decl @qwen3_moe_attention_postprocess_f32_f16(%token_count: index, %cache_row_count: index) launch(%token_count: index, %cache_row_count: index, %positions: buffer, %key_cache_indices: buffer, %value_cache_indices: buffer, %query_input: buffer, %key_input: buffer, %value_input: buffer, %query_norm_weight: buffer, %key_norm_weight: buffer, %inverse_frequencies: buffer, %query_output: buffer, %key_cache: buffer, %value_cache: buffer) + +// Storage-specific exports pass a constant value format into this shared body, +// while the compatibility export continues to source it from model config. +func.def inline @qwen3_moe_attention_qkv_postprocess_fused_decode_body(%value_uses_q6_index: index, %token_count: index, %cache_row_count: index, %q8_input: buffer, %query_weight: buffer, %key_weight: buffer, %value_weight: buffer, %positions: buffer, %key_cache_indices: buffer, %value_cache_indices: buffer, %query_output_raw: buffer, %key_output_raw: buffer, %value_output_raw: buffer, %query_norm_weight: buffer, %key_norm_weight: buffer, %inverse_frequencies: buffer, %query_output: buffer, %key_cache: buffer, %value_cache: buffer, %completion_counters: buffer) { + %publish_projection_output = scalar.constant true : i1 + %projection_token = kernel.workgroup.id : index + func.call @qwen3_moe_attention_qkv_quantized_body(%value_uses_q6_index, %publish_projection_output, %token_count, %projection_token, %q8_input, %query_weight, %key_weight, %value_weight, %query_output_raw, %key_output_raw, %value_output_raw) : (index, i1, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + %query_size0 = config.get @qwen3_moe.attention.query_size : index + %key_value_size0 = config.get @qwen3_moe.attention.key_value_size : index + %head_size0 = config.get @qwen3_moe.attention.head_size : index + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 1)] : index + %query_size, %key_value_size, %head_size = index.assume %query_size0, %key_value_size0, %head_size0 [range(%query_size0, 128, 262144), range(%key_value_size0, 128, 262144), range(%head_size0, 128, 128), mul(%query_size0, %head_size0), mul(%key_value_size0, %head_size0)] : index, index, index + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c8 = index.constant 8 : index + %c0_i32 = scalar.constant 0 : i32 + %c1_i32 = scalar.constant 1 : i32 + %c0_offset = index.constant 0 : offset + %counter_scratch_bytes = index.constant 4 : offset + %channel_tile = kernel.workgroup.id : index + %workitem = kernel.workitem.id : index + %token = index.assume %c0 [lt(%c0, %bounded_token_count)] : index + %head_tile_count = index.div %head_size, %c8 : index + %head_domain0 = index.div %channel_tile, %head_tile_count : index + %query_head_count = index.div %query_size, %head_size : index + %key_value_head_count = index.div %key_value_size, %head_size : index + %key_value_domain_count = index.mul %key_value_head_count, %c2 : index + %head_domain_count = index.add %query_head_count, %key_value_domain_count : index + %head_domain, %completion_counter_count = index.assume %head_domain0, %head_domain_count [lt(%head_domain0, %head_domain_count)] : index, index + %is_arrival_workitem = index.cmp eq, %workitem, %c0 : index + %completion_counters_aligned = buffer.assume.alignment %completion_counters {minimum_alignment = 16} : buffer + %completion_counters_view = buffer.view %completion_counters_aligned[%c0_offset] : buffer -> view<[%completion_counter_count]xi32> + %counter_scratch = buffer.alloca align(4) %counter_scratch_bytes : buffer + %counter_scratch_view = buffer.view %counter_scratch[%c0_offset] : buffer -> view<1xi32> + // Publish every producer's projection stores before the leader advances one + // workgroup arrival. The last arrival then acquires the complete head. + kernel.barrier scope(workgroup) ordering(release) + scf.if %is_arrival_workitem { + %old_counter = view.atomic.rmw %c1_i32, %completion_counters_view[%head_domain] {ordering = acq_rel, scope = device} : i32, view<[%completion_counter_count]xi32> -> i32 + view.store %old_counter, %counter_scratch_view[%c0] : i32, view<1xi32> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + %old_counter = view.load %counter_scratch_view[%c0] : view<1xi32> -> i32 + %head_tile_count_i32 = index.cast %head_tile_count : index to i32 + %last_head_tile_i32 = scalar.subi %head_tile_count_i32, %c1_i32 : i32 + %negative_head_tile_count_i32 = scalar.subi %c0_i32, %head_tile_count_i32 : i32 + %is_last_head_tile = scalar.cmpi eq, %old_counter, %last_head_tile_i32 : i32 + scf.if %is_last_head_tile { + kernel.barrier scope(workgroup) ordering(acquire) + %publish_postprocess_output = scalar.constant true : i1 + func.call @qwen3_moe_attention_postprocess_head_body(%publish_postprocess_output, %bounded_token_count, %cache_row_count, %head_domain, %token, %positions, %key_cache_indices, %value_cache_indices, %query_output_raw, %key_output_raw, %value_output_raw, %query_norm_weight, %key_norm_weight, %inverse_frequencies, %query_output, %key_cache, %value_cache) : (i1, index, index, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + // The counter becomes reusable only after all semantic output stores. + kernel.barrier scope(workgroup) ordering(release) + scf.if %is_arrival_workitem { + view.atomic.reduce %negative_head_tile_count_i32, %completion_counters_view[%head_domain] {ordering = release, scope = device} : i32, view<[%completion_counter_count]xi32> + } + } + func.return +} + +template.decl @qwen3_moe_attention_qkv_postprocess_fused_decode_launch() -> (index, index, index) +template.def<@qwen3_moe_attention_qkv_postprocess_fused_decode_launch> @qwen3_moe_attention_qkv_postprocess_fused_decode_launch_impl() -> (index, index, index) { + %query_size = config.get @qwen3_moe.attention.query_size : index + %key_value_size = config.get @qwen3_moe.attention.key_value_size : index + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c7 = index.constant 7 : index + %c8 = index.constant 8 : index + %c256 = index.constant 256 : index + %key_value_output_size = index.mul %key_value_size, %c2 : index + %output_size = index.add %query_size, %key_value_output_size : index + %padded_output_size = index.add %output_size, %c7 : index + %output_tiles = index.div %padded_output_size, %c8 : index + template.return %output_tiles, %c1, %c256 : index, index, index +} + +kernel.def target(@qwen3_moe_attention_qkv_postprocess_gfx11_wave32) @qwen3_moe_attention_qkv_postprocess_fused_decode(%token_count: index, %cache_row_count: index) { + %output_tiles, %c1, %c256 = template.apply<@qwen3_moe_attention_qkv_postprocess_fused_decode_launch>() pure : () -> (index, index, index) + kernel.launch.config workgroups(%output_tiles, %c1, %c1) workgroup_size(%c256, %c1, %c1) : index +} launch(%token_count: index, %cache_row_count: index, %q8_input: buffer, %query_weight: buffer, %key_weight: buffer, %value_weight: buffer, %positions: buffer, %key_cache_indices: buffer, %value_cache_indices: buffer, %query_output_raw: buffer, %key_output_raw: buffer, %value_output_raw: buffer, %query_norm_weight: buffer, %key_norm_weight: buffer, %inverse_frequencies: buffer, %query_output: buffer, %key_cache: buffer, %value_cache: buffer, %completion_counters: buffer) { + %value_uses_q6_index = config.get @qwen3_moe.attention.value_uses_q6 : index + func.call @qwen3_moe_attention_qkv_postprocess_fused_decode_body(%value_uses_q6_index, %token_count, %cache_row_count, %q8_input, %query_weight, %key_weight, %value_weight, %positions, %key_cache_indices, %value_cache_indices, %query_output_raw, %key_output_raw, %value_output_raw, %query_norm_weight, %key_norm_weight, %inverse_frequencies, %query_output, %key_cache, %value_cache, %completion_counters) : (index, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + kernel.return +} + +kernel.def target(@qwen3_moe_attention_qkv_postprocess_gfx11_wave32) @qwen3_moe_attention_qkv_postprocess_fused_decode_q4(%token_count: index, %cache_row_count: index) { + %output_tiles, %c1, %c256 = template.apply<@qwen3_moe_attention_qkv_postprocess_fused_decode_launch>() pure : () -> (index, index, index) + kernel.launch.config workgroups(%output_tiles, %c1, %c1) workgroup_size(%c256, %c1, %c1) : index +} launch(%token_count: index, %cache_row_count: index, %q8_input: buffer, %query_weight: buffer, %key_weight: buffer, %value_weight: buffer, %positions: buffer, %key_cache_indices: buffer, %value_cache_indices: buffer, %query_output_raw: buffer, %key_output_raw: buffer, %value_output_raw: buffer, %query_norm_weight: buffer, %key_norm_weight: buffer, %inverse_frequencies: buffer, %query_output: buffer, %key_cache: buffer, %value_cache: buffer, %completion_counters: buffer) { + %q4 = index.constant 0 : index + func.call @qwen3_moe_attention_qkv_postprocess_fused_decode_body(%q4, %token_count, %cache_row_count, %q8_input, %query_weight, %key_weight, %value_weight, %positions, %key_cache_indices, %value_cache_indices, %query_output_raw, %key_output_raw, %value_output_raw, %query_norm_weight, %key_norm_weight, %inverse_frequencies, %query_output, %key_cache, %value_cache, %completion_counters) : (index, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + kernel.return +} + +kernel.def target(@qwen3_moe_attention_qkv_postprocess_gfx11_wave32) @qwen3_moe_attention_qkv_postprocess_fused_decode_q6(%token_count: index, %cache_row_count: index) { + %output_tiles, %c1, %c256 = template.apply<@qwen3_moe_attention_qkv_postprocess_fused_decode_launch>() pure : () -> (index, index, index) + kernel.launch.config workgroups(%output_tiles, %c1, %c1) workgroup_size(%c256, %c1, %c1) : index +} launch(%token_count: index, %cache_row_count: index, %q8_input: buffer, %query_weight: buffer, %key_weight: buffer, %value_weight: buffer, %positions: buffer, %key_cache_indices: buffer, %value_cache_indices: buffer, %query_output_raw: buffer, %key_output_raw: buffer, %value_output_raw: buffer, %query_norm_weight: buffer, %key_norm_weight: buffer, %inverse_frequencies: buffer, %query_output: buffer, %key_cache: buffer, %value_cache: buffer, %completion_counters: buffer) { + %q6 = index.constant 1 : index + func.call @qwen3_moe_attention_qkv_postprocess_fused_decode_body(%q6, %token_count, %cache_row_count, %q8_input, %query_weight, %key_weight, %value_weight, %positions, %key_cache_indices, %value_cache_indices, %query_output_raw, %key_output_raw, %value_output_raw, %query_norm_weight, %key_norm_weight, %inverse_frequencies, %query_output, %key_cache, %value_cache, %completion_counters) : (index, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + kernel.return +} + +// Uses the production head geometry, nonzero rotary position, and distinct K/V +// cache rows. Two fused calls reuse the same counters after comparison with +// the ordinary storage-selected composition. The value buffer has production +// Q6_K capacity and is intentionally oversized when the Q4_K config is tested. +check.case public @qwen3_moe_attention_qkv_postprocess_fused_differential_case { + %token_count = check.literal value(1) : index + %cache_row_count = check.literal value(4) : index + %hidden_size = check.literal value(2048) : index + %input = check.generate.fill value(0.00390625) : tensor<1x2048xf32> + %q8_input = check.generate.fill value(0) : tensor<2304xi8> + %query_weight = check.generate.fill value(34) : tensor<4096x8x144xi8> + %key_weight = check.generate.fill value(35) : tensor<512x8x144xi8> + %value_weight = check.generate.fill value(-86) : tensor<860160xi8> + %positions = check.generate.fill value(7) : tensor<1xi32> + %key_cache_indices = check.generate.fill value(1) : tensor<1xi64> + %value_cache_indices = check.generate.fill value(2) : tensor<1xi64> + %query_norm_weight = check.generate.iota offset(0.5) step(0.00390625) : tensor<128xf32> + %key_norm_weight = check.generate.iota offset(0.75) step(0.001953125) : tensor<128xf32> + %inverse_frequencies = check.generate.fill value(0.03125) : tensor<64xf32> + %expected_query_raw = check.generate.fill value(1.0) : tensor<1x4096xf32> + %expected_key_raw = check.generate.fill value(1.0) : tensor<1x512xf32> + %expected_value_raw = check.generate.fill value(1.0) : tensor<1x512xf32> + %expected_query = check.generate.fill value(0.0) : tensor<1x32x128xf32> + %expected_key_cache = check.generate.fill value(-1.0) : tensor<4x4x128xf16> + %expected_value_cache = check.generate.fill value(-1.0) : tensor<4x4x128xf16> + %actual_query_raw0 = check.generate.fill value(1.0) : tensor<1x4096xf32> + %actual_key_raw0 = check.generate.fill value(1.0) : tensor<1x512xf32> + %actual_value_raw0 = check.generate.fill value(1.0) : tensor<1x512xf32> + %actual_query0 = check.generate.fill value(0.0) : tensor<1x32x128xf32> + %actual_key_cache0 = check.generate.fill value(-1.0) : tensor<4x4x128xf16> + %actual_value_cache0 = check.generate.fill value(-1.0) : tensor<4x4x128xf16> + %actual_query_raw1 = check.generate.fill value(1.0) : tensor<1x4096xf32> + %actual_key_raw1 = check.generate.fill value(1.0) : tensor<1x512xf32> + %actual_value_raw1 = check.generate.fill value(1.0) : tensor<1x512xf32> + %actual_query1 = check.generate.fill value(0.0) : tensor<1x32x128xf32> + %actual_key_cache1 = check.generate.fill value(-1.0) : tensor<4x4x128xf16> + %actual_value_cache1 = check.generate.fill value(-1.0) : tensor<4x4x128xf16> + %completion_counters = check.generate.fill value(0) : tensor<40xi32> + %expected_counters = check.generate.fill value(0) : tensor<40xi32> + kernel.launch @ggml_quantize_q8_1_x4_f32[%token_count, %hidden_size](%token_count, %hidden_size, %input, %q8_input) : [index, index](index, index, tensor<1x2048xf32>, tensor<2304xi8>) + kernel.launch @qwen3_moe_attention_qkv_quantized[%token_count](%token_count, %q8_input, %query_weight, %key_weight, %value_weight, %expected_query_raw, %expected_key_raw, %expected_value_raw) : [index](index, tensor<2304xi8>, tensor<4096x8x144xi8>, tensor<512x8x144xi8>, tensor<860160xi8>, tensor<1x4096xf32>, tensor<1x512xf32>, tensor<1x512xf32>) + kernel.launch @qwen3_moe_attention_postprocess_f32_f16[%token_count, %cache_row_count](%token_count, %cache_row_count, %positions, %key_cache_indices, %value_cache_indices, %expected_query_raw, %expected_key_raw, %expected_value_raw, %query_norm_weight, %key_norm_weight, %inverse_frequencies, %expected_query, %expected_key_cache, %expected_value_cache) : [index, index](index, index, tensor<1xi32>, tensor<1xi64>, tensor<1xi64>, tensor<1x4096xf32>, tensor<1x512xf32>, tensor<1x512xf32>, tensor<128xf32>, tensor<128xf32>, tensor<64xf32>, tensor<1x32x128xf32>, tensor<4x4x128xf16>, tensor<4x4x128xf16>) + kernel.launch @qwen3_moe_attention_qkv_postprocess_fused_decode[%token_count, %cache_row_count](%token_count, %cache_row_count, %q8_input, %query_weight, %key_weight, %value_weight, %positions, %key_cache_indices, %value_cache_indices, %actual_query_raw0, %actual_key_raw0, %actual_value_raw0, %query_norm_weight, %key_norm_weight, %inverse_frequencies, %actual_query0, %actual_key_cache0, %actual_value_cache0, %completion_counters) : [index, index](index, index, tensor<2304xi8>, tensor<4096x8x144xi8>, tensor<512x8x144xi8>, tensor<860160xi8>, tensor<1xi32>, tensor<1xi64>, tensor<1xi64>, tensor<1x4096xf32>, tensor<1x512xf32>, tensor<1x512xf32>, tensor<128xf32>, tensor<128xf32>, tensor<64xf32>, tensor<1x32x128xf32>, tensor<4x4x128xf16>, tensor<4x4x128xf16>, tensor<40xi32>) + kernel.launch @qwen3_moe_attention_qkv_postprocess_fused_decode[%token_count, %cache_row_count](%token_count, %cache_row_count, %q8_input, %query_weight, %key_weight, %value_weight, %positions, %key_cache_indices, %value_cache_indices, %actual_query_raw1, %actual_key_raw1, %actual_value_raw1, %query_norm_weight, %key_norm_weight, %inverse_frequencies, %actual_query1, %actual_key_cache1, %actual_value_cache1, %completion_counters) : [index, index](index, index, tensor<2304xi8>, tensor<4096x8x144xi8>, tensor<512x8x144xi8>, tensor<860160xi8>, tensor<1xi32>, tensor<1xi64>, tensor<1xi64>, tensor<1x4096xf32>, tensor<1x512xf32>, tensor<1x512xf32>, tensor<128xf32>, tensor<128xf32>, tensor<64xf32>, tensor<1x32x128xf32>, tensor<4x4x128xf16>, tensor<4x4x128xf16>, tensor<40xi32>) + check.expect.close actual(%actual_query_raw0) expected(%expected_query_raw) atol(0.0) rtol(0.0) nan(same) : tensor<1x4096xf32> + check.expect.close actual(%actual_key_raw0) expected(%expected_key_raw) atol(0.0) rtol(0.0) nan(same) : tensor<1x512xf32> + check.expect.close actual(%actual_value_raw0) expected(%expected_value_raw) atol(0.0) rtol(0.0) nan(same) : tensor<1x512xf32> + check.expect.close actual(%actual_query0) expected(%expected_query) atol(0.0001) rtol(0.0001) nan(same) : tensor<1x32x128xf32> + check.expect.close actual(%actual_key_cache0) expected(%expected_key_cache) atol(0.002) rtol(0.002) nan(same) : tensor<4x4x128xf16> + check.expect.close actual(%actual_value_cache0) expected(%expected_value_cache) atol(0.0) rtol(0.0) nan(same) : tensor<4x4x128xf16> + check.expect.close actual(%actual_query_raw1) expected(%expected_query_raw) atol(0.0) rtol(0.0) nan(same) : tensor<1x4096xf32> + check.expect.close actual(%actual_key_raw1) expected(%expected_key_raw) atol(0.0) rtol(0.0) nan(same) : tensor<1x512xf32> + check.expect.close actual(%actual_value_raw1) expected(%expected_value_raw) atol(0.0) rtol(0.0) nan(same) : tensor<1x512xf32> + check.expect.close actual(%actual_query1) expected(%expected_query) atol(0.0001) rtol(0.0001) nan(same) : tensor<1x32x128xf32> + check.expect.close actual(%actual_key_cache1) expected(%expected_key_cache) atol(0.002) rtol(0.002) nan(same) : tensor<4x4x128xf16> + check.expect.close actual(%actual_value_cache1) expected(%expected_value_cache) atol(0.0) rtol(0.0) nan(same) : tensor<4x4x128xf16> + check.expect.equal actual(%completion_counters) expected(%expected_counters) : tensor<40xi32> + check.return +} + +check.case public @qwen3_moe_attention_qkv_postprocess_composed_benchmark_case { + %token_count = check.literal value(1) : index + %cache_row_count = check.literal value(1024) : index + %q8_input = check.generate.fill value(0) : tensor<2304xi8> + %query_weight = check.generate.fill value(0) : tensor<4096x8x144xi8> + %key_weight = check.generate.fill value(0) : tensor<512x8x144xi8> + %value_weight = check.generate.fill value(0) : tensor<512x8x210xi8> + %positions = check.generate.fill value(513) : tensor<1xi32> + %key_cache_indices = check.generate.fill value(513) : tensor<1xi64> + %value_cache_indices = check.generate.fill value(513) : tensor<1xi64> + %query_output_raw = check.generate.fill value(1.0) : tensor<1x4096xf32> + %key_output_raw = check.generate.fill value(1.0) : tensor<1x512xf32> + %value_output_raw = check.generate.fill value(1.0) : tensor<1x512xf32> + %query_norm_weight = check.generate.fill value(1.0) : tensor<128xf32> + %key_norm_weight = check.generate.fill value(1.0) : tensor<128xf32> + %inverse_frequencies = check.generate.fill value(1.0) : tensor<64xf32> + %query_output = check.generate.fill value(0.0) : tensor<1x32x128xf32> + %key_cache = check.generate.fill value(0.0) : tensor<1024x4x128xf16> + %value_cache = check.generate.fill value(0.0) : tensor<1024x4x128xf16> + kernel.launch @qwen3_moe_attention_qkv_quantized[%token_count](%token_count, %q8_input, %query_weight, %key_weight, %value_weight, %query_output_raw, %key_output_raw, %value_output_raw) : [index](index, tensor<2304xi8>, tensor<4096x8x144xi8>, tensor<512x8x144xi8>, tensor<512x8x210xi8>, tensor<1x4096xf32>, tensor<1x512xf32>, tensor<1x512xf32>) + kernel.launch @qwen3_moe_attention_postprocess_f32_f16[%token_count, %cache_row_count](%token_count, %cache_row_count, %positions, %key_cache_indices, %value_cache_indices, %query_output_raw, %key_output_raw, %value_output_raw, %query_norm_weight, %key_norm_weight, %inverse_frequencies, %query_output, %key_cache, %value_cache) : [index, index](index, index, tensor<1xi32>, tensor<1xi64>, tensor<1xi64>, tensor<1x4096xf32>, tensor<1x512xf32>, tensor<1x512xf32>, tensor<128xf32>, tensor<128xf32>, tensor<64xf32>, tensor<1x32x128xf32>, tensor<1024x4x128xf16>, tensor<1024x4x128xf16>) + check.return +} + +check.case public @qwen3_moe_attention_qkv_postprocess_fused_benchmark_case { + %token_count = check.literal value(1) : index + %cache_row_count = check.literal value(1024) : index + %q8_input = check.generate.fill value(0) : tensor<2304xi8> + %query_weight = check.generate.fill value(0) : tensor<4096x8x144xi8> + %key_weight = check.generate.fill value(0) : tensor<512x8x144xi8> + %value_weight = check.generate.fill value(0) : tensor<512x8x210xi8> + %positions = check.generate.fill value(513) : tensor<1xi32> + %key_cache_indices = check.generate.fill value(513) : tensor<1xi64> + %value_cache_indices = check.generate.fill value(513) : tensor<1xi64> + %query_output_raw = check.generate.fill value(1.0) : tensor<1x4096xf32> + %key_output_raw = check.generate.fill value(1.0) : tensor<1x512xf32> + %value_output_raw = check.generate.fill value(1.0) : tensor<1x512xf32> + %query_norm_weight = check.generate.fill value(1.0) : tensor<128xf32> + %key_norm_weight = check.generate.fill value(1.0) : tensor<128xf32> + %inverse_frequencies = check.generate.fill value(1.0) : tensor<64xf32> + %query_output = check.generate.fill value(0.0) : tensor<1x32x128xf32> + %key_cache = check.generate.fill value(0.0) : tensor<1024x4x128xf16> + %value_cache = check.generate.fill value(0.0) : tensor<1024x4x128xf16> + %completion_counters = check.generate.fill value(0) : tensor<40xi32> + kernel.launch @qwen3_moe_attention_qkv_postprocess_fused_decode[%token_count, %cache_row_count](%token_count, %cache_row_count, %q8_input, %query_weight, %key_weight, %value_weight, %positions, %key_cache_indices, %value_cache_indices, %query_output_raw, %key_output_raw, %value_output_raw, %query_norm_weight, %key_norm_weight, %inverse_frequencies, %query_output, %key_cache, %value_cache, %completion_counters) : [index, index](index, index, tensor<2304xi8>, tensor<4096x8x144xi8>, tensor<512x8x144xi8>, tensor<512x8x210xi8>, tensor<1xi32>, tensor<1xi64>, tensor<1xi64>, tensor<1x4096xf32>, tensor<1x512xf32>, tensor<1x512xf32>, tensor<128xf32>, tensor<128xf32>, tensor<64xf32>, tensor<1x32x128xf32>, tensor<1024x4x128xf16>, tensor<1024x4x128xf16>, tensor<40xi32>) + check.return +} + +check.benchmark<@qwen3_moe_attention_qkv_postprocess_composed_benchmark_case> @qwen3_moe_attention_qkv_postprocess_composed_decode + +check.benchmark<@qwen3_moe_attention_qkv_postprocess_fused_benchmark_case> @qwen3_moe_attention_qkv_postprocess_fused_boundary_decode diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/attention_qkv_quantized.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/attention_qkv_quantized.loom new file mode 100644 index 000000000000..00a7ab8ac84e --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/attention_qkv_quantized.loom @@ -0,0 +1,670 @@ +// Co-scheduled Qwen Q/K/V projections from one GGML Q8_1 x4 activation row. +// +// Query and key weights use their native GGUF Q4_K rows. Value weights are +// selected at JIT time between Q4_K and Q6_K because the checkpoint uses both +// layer contracts. One wave owns one output row, while a 256-workitem +// workgroup carries eight rows drawn from the concatenated Q/K/V output +// domain. The three logical outputs remain distinct bindings and layouts. +// +// This kernel intentionally starts after activation packing and ends at raw +// F32 projections. The adjacent preparation producer owns RMSNorm and Q8_1 +// packing; the following attention postprocess owns per-head normalization, +// RoPE, and K/V cache publication. +func.def inline @ggml_q6k_q8_1_x4_block_lane(%weight: buffer, %weight_row_byte_base: offset, %q8_input: buffer, %q8_row_byte_base: offset, %q6_block: index, %lane: index) -> (f32) { + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c4 = index.constant 4 : index + %c8 = index.constant 8 : index + %c16 = index.constant 16 : index + %block_bytes = index.constant 210 : offset + %qh_byte_add = index.constant 128 : offset + %scale_byte_add = index.constant 192 : offset + %d_byte_add = index.constant 208 : offset + %c4_i32v = vector.constant 4 : vector<1xi32> + %c0_i32v = vector.constant 0 : vector<1xi32> + %nibble_mask = vector.constant 252645135 : vector<1xi32> + %high_mask = vector.constant 808464432 : vector<1xi32> + %c0_f32 = scalar.constant 0.0 : f32 + %bounded_lane = index.assume %lane [range(%lane, 0, 31)] : index + %block_byte_add = index.scale %q6_block, %block_bytes : index, offset -> offset + %block_byte_base = index.add %weight_row_byte_base, %block_byte_add : offset + %qh_byte_base = index.add %block_byte_base, %qh_byte_add : offset + %scale_byte_base = index.add %block_byte_base, %scale_byte_add : offset + %d_byte_base = index.add %block_byte_base, %d_byte_add : offset + %ql_view = buffer.view %weight[%block_byte_base] : buffer -> view<32xi32> + %qh_view = buffer.view %weight[%qh_byte_base] : buffer -> view<16xi32> + %scale_view = buffer.view %weight[%scale_byte_base] : buffer -> view<16xi8> + %d_view = buffer.view %weight[%d_byte_base] : buffer -> view<1xf16> + %lane_mod8 = index.rem %bounded_lane, %c8 : index + %lane_mod16 = index.rem %bounded_lane, %c16 : index + %lane_div16 = index.div %bounded_lane, %c16 : index + %lane_div8_in_16 = index.div %lane_mod16, %c8 : index + %lane_div4_in_16 = index.div %lane_mod16, %c4 : index + %qh_high_base = index.mul %lane_div16, %c8 : index + %qh_index0 = index.add %qh_high_base, %lane_mod8 : index + %qh_index = index.assume %qh_index0 [range(%qh_index0, 0, 15)] : index + %ql_word = vector.load %ql_view[%bounded_lane] : view<32xi32> -> vector<1xi32> + %qh_word = vector.load %qh_view[%qh_index] : view<16xi32> -> vector<1xi32> + %qh_base_shift_index = index.mul %lane_div8_in_16, %c2 : index + %qh_base_shift_i32 = index.cast %qh_base_shift_index : index to i32 + %q8_block_base = index.mul %q6_block, %c8 : index + %q8_high_add = index.mul %lane_div16, %c4 : index + %q8_quadrant = index.add %q8_high_add, %lane_div8_in_16 : index + %scale_high_base = index.mul %lane_div16, %c8 : index + %scale_lane0 = index.add %scale_high_base, %lane_div4_in_16 : index + %d_f16 = view.load %d_view[0] : view<1xf16> -> f16 + %d = scalar.extf %d_f16 : f16 to f32 + %sum = scf.for %part = [%c0 to %c2 step %c1](%accumulator = %c0_f32 : f32) -> (f32) unroll { + %bounded_part = index.assume %part [range(%part, 0, 1)] : index + %part_shift_index = index.mul %bounded_part, %c4 : index + %part_shift_i32 = index.cast %part_shift_index : index to i32 + %part_shift = vector.splat %part_shift_i32 : vector<1xi32> + %ql_shifted = vector.shrui %ql_word, %part_shift : vector<1xi32> + %ql = vector.andi %ql_shifted, %nibble_mask : vector<1xi32> + %qh_shift_i32 = scalar.addi %qh_base_shift_i32, %part_shift_i32 : i32 + %qh_shift = vector.splat %qh_shift_i32 : vector<1xi32> + %qh_shifted = vector.shrui %qh_word, %qh_shift : vector<1xi32> + %qh_positioned = vector.shli %qh_shifted, %c4_i32v : vector<1xi32> + %qh = vector.andi %qh_positioned, %high_mask : vector<1xi32> + %code = vector.ori %ql, %qh : vector<1xi32> + %signed_weight = func.call @ggml_q6k_sign_extend_dot4(%code) : (vector<1xi32>) -> (vector<4xi8>) + %q8_part_add = index.mul %bounded_part, %c2 : index + %q8_block_part = index.add %q8_block_base, %q8_part_add : index + %q8_block = index.add %q8_block_part, %q8_quadrant : index + %q8_values, %q8_d = func.call @ggml_q8_1_x4_word(%q8_input, %q8_row_byte_base, %q8_block, %lane_mod8) : (buffer, offset, index, index) -> (vector<4xi8>, f32) + %scale_lane = index.add %scale_lane0, %part_shift_index : index + %scale_i8 = view.load %scale_view[%scale_lane] : view<16xi8> -> i8 + %scale = scalar.sitofp %scale_i8 : i8 to f32 + %dot = vector.dot4i %signed_weight, %q8_values, %c0_i32v : vector<4xi8>, vector<4xi8>, vector<1xi32> + %dot_i32 = vector.extract %dot[0] : vector<1xi32> -> i32 + %dot_f32 = scalar.sitofp %dot_i32 : i32 to f32 + %scaled0 = scalar.mulf %dot_f32, %scale : f32 + %scaled1 = scalar.mulf %scaled0, %d : f32 + %contribution = scalar.mulf %scaled1, %q8_d : f32 + %next = scalar.addf %accumulator, %contribution : f32 + scf.yield %next : f32 + } + func.return %sum : f32 +} + +func.def inline @ggml_q6k_q8_1_x4_row_lane(%input_size: index, %weight: buffer, %weight_row_byte_base: offset, %q8_input: buffer, %q8_row_byte_base: offset, %lane: index) -> (f32) { + %bounded_input_size = index.assume %input_size [range(%input_size, 256, 32768), mul(%input_size, 256)] : index + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c256 = index.constant 256 : index + %c0_f32 = scalar.constant 0.0 : f32 + %block_count = index.div %bounded_input_size, %c256 : index + %result = scf.for %block = [%c0 to %block_count step %c1](%block_acc = %c0_f32 : f32) -> (f32) { + %contribution = func.call @ggml_q6k_q8_1_x4_block_lane(%weight, %weight_row_byte_base, %q8_input, %q8_row_byte_base, %block, %lane) : (buffer, offset, buffer, offset, index, index) -> (f32) + %next = scalar.addf %block_acc, %contribution : f32 + scf.yield %next : f32 + } + func.return %result : f32 +} + +func.def inline @ggml_q6k_sign_extend_dot4(%code: vector<1xi32>) -> (vector<4xi8>) { + %c1_i32v = vector.constant 1 : vector<1xi32> + %c2_i32v = vector.constant 2 : vector<1xi32> + %low5_mask = vector.constant 522133279 : vector<1xi32> + %bit5_mask = vector.constant 538976288 : vector<1xi32> + %sign_mask = vector.constant -522133280 : vector<1xi32> + %low5 = vector.andi %code, %low5_mask : vector<1xi32> + %bit5 = vector.andi %code, %bit5_mask : vector<1xi32> + %bit6 = vector.shli %bit5, %c1_i32v : vector<1xi32> + %bit7 = vector.shli %bit5, %c2_i32v : vector<1xi32> + %high01 = vector.ori %bit5, %bit6 : vector<1xi32> + %high = vector.ori %high01, %bit7 : vector<1xi32> + %sign = vector.xori %high, %sign_mask : vector<1xi32> + %signed_i32 = vector.ori %low5, %sign : vector<1xi32> + %signed = vector.bitcast %signed_i32 : vector<1xi32> to vector<4xi8> + func.return %signed : vector<4xi8> +} + +func.def inline @ggml_q8_1_x4_word(%q8_input: buffer, %row_byte_base: offset, %q8_block: index, %word_in_block: index) -> (vector<4xi8>, f32) { + %c2 = index.constant 2 : index + %c4 = index.constant 4 : index + %c8 = index.constant 8 : index + %group_bytes = index.constant 144 : offset + %payload_byte_add = index.constant 16 : offset + %group = index.div %q8_block, %c4 : index + %inner0 = index.rem %q8_block, %c4 : index + %inner = index.assume %inner0 [range(%inner0, 0, 3)] : index + %word0 = index.assume %word_in_block [range(%word_in_block, 0, 7)] : index + %group_byte_add = index.scale %group, %group_bytes : index, offset -> offset + %group_byte_base = index.add %row_byte_base, %group_byte_add : offset + %payload_byte_base = index.add %group_byte_base, %payload_byte_add : offset + %ds_view = buffer.view %q8_input[%group_byte_base] : buffer -> view<8xf16> + %payload_view = buffer.view %q8_input[%payload_byte_base] : buffer -> view<32xi32> + %d_index = index.mul %inner, %c2 : index + %inner_word_base = index.mul %inner, %c8 : index + %word_index = index.add %inner_word_base, %word0 : index + %d_f16 = view.load %ds_view[%d_index] : view<8xf16> -> f16 + %packed = vector.load %payload_view[%word_index] : view<32xi32> -> vector<1xi32> + %values = vector.bitcast %packed : vector<1xi32> to vector<4xi8> + %d = scalar.extf %d_f16 : f16 to f32 + func.return %values, %d : vector<4xi8>, f32 +} + +func.def inline @qwen3_moe_q4k_chunk_pair_global(%weight: buffer, %row_byte_base: offset, %q4_block: index, %q4_group_pair: index, %q4_half: index, %header_words: vector<4xi32>) -> (vector<16xi8>, f32, f32, vector<16xi8>, f32, f32) { + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c4 = index.constant 4 : index + %c8 = index.constant 8 : index + %block_bytes = index.constant 144 : offset + %code_byte_add = index.constant 16 : offset + %c4_i32 = scalar.constant 4 : i32 + %nibble_mask = vector.constant 252645135 : vector<4xi32> + %bounded_pair = index.assume %q4_group_pair [range(%q4_group_pair, 0, 3)] : index + %bounded_half = index.assume %q4_half [range(%q4_half, 0, 1)] : index + %block_byte_add = index.scale %q4_block, %block_bytes : index, offset -> offset + %block_byte_base = index.add %row_byte_base, %block_byte_add : offset + %code_byte_base = index.add %block_byte_base, %code_byte_add : offset + %code_view = buffer.view %weight[%code_byte_base] : buffer -> view<32xi32> + %header_halves = vector.bitcast %header_words : vector<4xi32> to vector<8xf16> + %d_f16 = vector.extract %header_halves[0] : vector<8xf16> -> f16 + %dmin_f16 = vector.extract %header_halves[1] : vector<8xf16> -> f16 + %scale0 = vector.extract %header_words[1] : vector<4xi32> -> i32 + %scale1 = vector.extract %header_words[2] : vector<4xi32> -> i32 + %scale2 = vector.extract %header_words[3] : vector<4xi32> -> i32 + %d = scalar.extf %d_f16 : f16 to f32 + %dmin = scalar.extf %dmin_f16 : f16 to f32 + %pair_code_base = index.mul %bounded_pair, %c8 : index + %half_code_add = index.mul %bounded_half, %c4 : index + %code_index0 = index.add %pair_code_base, %half_code_add : index + %code_index = index.assume %code_index0 [range(%code_index0, 0, 28)] : index + %packed_codes = vector.load %code_view[%code_index] : view<32xi32> -> vector<4xi32> + %low_codes = vector.andi %packed_codes, %nibble_mask : vector<4xi32> + %c4_i32v = vector.splat %c4_i32 : vector<4xi32> + %high_shifted = vector.shrui %packed_codes, %c4_i32v : vector<4xi32> + %high_codes = vector.andi %high_shifted, %nibble_mask : vector<4xi32> + %q4_low = vector.bitcast %low_codes : vector<4xi32> to vector<16xi8> + %q4_high = vector.bitcast %high_codes : vector<4xi32> to vector<16xi8> + %low_group = index.mul %bounded_pair, %c2 : index + %high_group = index.add %low_group, %c1 : index + %low_scale, %low_minimum = func.call @qwen3_moe_q4k_scale_from_header(%scale0, %scale1, %scale2, %low_group) : (i32, i32, i32, index) -> (i32, i32) + %high_scale, %high_minimum = func.call @qwen3_moe_q4k_scale_from_header(%scale0, %scale1, %scale2, %high_group) : (i32, i32, i32, index) -> (i32, i32) + %low_scale_f32 = scalar.uitofp %low_scale : i32 to f32 + %low_minimum_f32 = scalar.uitofp %low_minimum : i32 to f32 + %high_scale_f32 = scalar.uitofp %high_scale : i32 to f32 + %high_minimum_f32 = scalar.uitofp %high_minimum : i32 to f32 + %low_d_scale = scalar.mulf %d, %low_scale_f32 : f32 + %low_dmin_scale = scalar.mulf %dmin, %low_minimum_f32 : f32 + %high_d_scale = scalar.mulf %d, %high_scale_f32 : f32 + %high_dmin_scale = scalar.mulf %dmin, %high_minimum_f32 : f32 + func.return %q4_low, %low_d_scale, %low_dmin_scale, %q4_high, %high_d_scale, %high_dmin_scale : vector<16xi8>, f32, f32, vector<16xi8>, f32, f32 +} + +func.def inline @qwen3_moe_q4k_q8_1_dot(%q4_values: vector<16xi8>, %d_scale: f32, %dmin_scale: f32, %q8_values: vector<16xi8>, %q8_d: f32, %q8_s: f32) -> (f32) { + %c0_i32 = scalar.constant 0 : i32 + %c0_i32v = vector.constant 0 : vector<4xi32> + %half_f32 = scalar.constant 0.5 : f32 + %partial_dots = vector.dot4i %q4_values, %q8_values, %c0_i32v : vector<16xi8>, vector<16xi8>, vector<4xi32> + %q_sum = vector.reduce %partial_dots, %c0_i32 : vector<4xi32>, i32 + %q_sum_f32 = scalar.sitofp %q_sum : i32 to f32 + %scaled_dot0 = scalar.mulf %q8_d, %d_scale : f32 + %scaled_dot = scalar.mulf %scaled_dot0, %q_sum_f32 : f32 + %q8_half_sum = scalar.mulf %q8_s, %half_f32 : f32 + %minimum_correction = scalar.mulf %dmin_scale, %q8_half_sum : f32 + %contribution = scalar.subf %scaled_dot, %minimum_correction : f32 + func.return %contribution : f32 +} + +func.def inline @qwen3_moe_q4k_q8_1_x4_paired_block_lane(%input_size: index, %weight: buffer, %weight_row_byte_base: offset, %q8_input: buffer, %q8_row_byte_base: offset, %q4_block: index, %block_lane: index) -> (f32) { + %bounded_input_size = index.assume %input_size [range(%input_size, 256, 32768), mul(%input_size, 256)] : index + %bounded_q4_block0 = index.assume %q4_block [range(%q4_block, 0, 127)] : index + %bounded_block_lane = index.assume %block_lane [range(%block_lane, 0, 7)] : index + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c4 = index.constant 4 : index + %c8 = index.constant 8 : index + %c128 = index.constant 128 : index + %c256 = index.constant 256 : index + %q4_block_bytes = index.constant 144 : offset + %q8_group_bytes = index.constant 144 : offset + %q8_payload_byte_add = index.constant 16 : offset + %q4_block_count = index.div %bounded_input_size, %c256 : index + %q8_group_count = index.div %bounded_input_size, %c128 : index + %bounded_q4_block, %bounded_q4_block_count = index.assume %bounded_q4_block0, %q4_block_count [lt(%bounded_q4_block0, %q4_block_count)] : index, index + %q4_group_pair0 = index.div %bounded_block_lane, %c2 : index + %q4_group_pair = index.assume %q4_group_pair0 [range(%q4_group_pair0, 0, 3)] : index + %q4_half0 = index.rem %bounded_block_lane, %c2 : index + %q4_half = index.assume %q4_half0 [range(%q4_half0, 0, 1)] : index + %q8_group_in_block0 = index.div %q4_group_pair, %c2 : index + %q8_group_in_block = index.assume %q8_group_in_block0 [range(%q8_group_in_block0, 0, 1)] : index + %pair_in_q8_group0 = index.rem %q4_group_pair, %c2 : index + %pair_in_q8_group = index.assume %pair_in_q8_group0 [range(%pair_in_q8_group0, 0, 1)] : index + %q8_low_inner_block0 = index.mul %pair_in_q8_group, %c2 : index + %q8_low_inner_block = index.assume %q8_low_inner_block0 [range(%q8_low_inner_block0, 0, 2)] : index + %q8_high_inner_block0 = index.add %q8_low_inner_block, %c1 : index + %q8_high_inner_block = index.assume %q8_high_inner_block0 [range(%q8_high_inner_block0, 1, 3)] : index + %q8_half_word_add = index.mul %q4_half, %c4 : index + %q8_low_inner_word_base = index.mul %q8_low_inner_block, %c8 : index + %q8_low_word_index0 = index.add %q8_low_inner_word_base, %q8_half_word_add : index + %q8_low_word_index = index.assume %q8_low_word_index0 [range(%q8_low_word_index0, 0, 20)] : index + %q8_high_inner_word_base = index.mul %q8_high_inner_block, %c8 : index + %q8_high_word_index0 = index.add %q8_high_inner_word_base, %q8_half_word_add : index + %q8_high_word_index = index.assume %q8_high_word_index0 [range(%q8_high_word_index0, 8, 28)] : index + %q8_low_ds_index0 = index.mul %q8_low_inner_block, %c2 : index + %q8_low_ds_index = index.assume %q8_low_ds_index0 [range(%q8_low_ds_index0, 0, 4)] : index + %q8_block_group_base = index.mul %bounded_q4_block, %c2 : index + %q8_group0 = index.add %q8_block_group_base, %q8_group_in_block : index + %q8_group, %bounded_q8_group_count = index.assume %q8_group0, %q8_group_count [lt(%q8_group0, %q8_group_count)] : index, index + %q8_group_byte_add = index.scale %q8_group, %q8_group_bytes : index, offset -> offset + %q8_group_byte_base = index.add %q8_row_byte_base, %q8_group_byte_add : offset + %q8_payload_byte_base = index.add %q8_group_byte_base, %q8_payload_byte_add : offset + %q8_ds_view = buffer.view %q8_input[%q8_group_byte_base] : buffer -> view<8xf16> + %q8_words_view = buffer.view %q8_input[%q8_payload_byte_base] : buffer -> view<32xi32> + %q8_ds = vector.load %q8_ds_view[%q8_low_ds_index] : view<8xf16> -> vector<4xf16> + %q8_low_d_f16 = vector.extract %q8_ds[0] : vector<4xf16> -> f16 + %q8_low_s_f16 = vector.extract %q8_ds[1] : vector<4xf16> -> f16 + %q8_high_d_f16 = vector.extract %q8_ds[2] : vector<4xf16> -> f16 + %q8_high_s_f16 = vector.extract %q8_ds[3] : vector<4xf16> -> f16 + %q8_low_d = scalar.extf %q8_low_d_f16 : f16 to f32 + %q8_low_s = scalar.extf %q8_low_s_f16 : f16 to f32 + %q8_high_d = scalar.extf %q8_high_d_f16 : f16 to f32 + %q8_high_s = scalar.extf %q8_high_s_f16 : f16 to f32 + %q8_low_words = vector.load %q8_words_view[%q8_low_word_index] : view<32xi32> -> vector<4xi32> + %q8_high_words = vector.load %q8_words_view[%q8_high_word_index] : view<32xi32> -> vector<4xi32> + %q8_low_values = vector.bitcast %q8_low_words : vector<4xi32> to vector<16xi8> + %q8_high_values = vector.bitcast %q8_high_words : vector<4xi32> to vector<16xi8> + %q4_block_byte_add = index.scale %bounded_q4_block, %q4_block_bytes : index, offset -> offset + %q4_block_byte_base = index.add %weight_row_byte_base, %q4_block_byte_add : offset + %q4_header_view = buffer.view %weight[%q4_block_byte_base] : buffer -> view<4xi32> + %q4_header_words = vector.load %q4_header_view[0] : view<4xi32> -> vector<4xi32> + %q4_low, %low_d_scale, %low_dmin_scale, %q4_high, %high_d_scale, %high_dmin_scale = func.call @qwen3_moe_q4k_chunk_pair_global(%weight, %weight_row_byte_base, %bounded_q4_block, %q4_group_pair, %q4_half, %q4_header_words) : (buffer, offset, index, index, index, vector<4xi32>) -> (vector<16xi8>, f32, f32, vector<16xi8>, f32, f32) + %low = func.call @qwen3_moe_q4k_q8_1_dot(%q4_low, %low_d_scale, %low_dmin_scale, %q8_low_values, %q8_low_d, %q8_low_s) : (vector<16xi8>, f32, f32, vector<16xi8>, f32, f32) -> (f32) + %high = func.call @qwen3_moe_q4k_q8_1_dot(%q4_high, %high_d_scale, %high_dmin_scale, %q8_high_values, %q8_high_d, %q8_high_s) : (vector<16xi8>, f32, f32, vector<16xi8>, f32, f32) -> (f32) + %pair = scalar.addf %low, %high : f32 + func.return %pair : f32 +} + +func.def inline @qwen3_moe_q4k_q8_1_x4_paired_row_lane(%input_size: index, %weight: buffer, %weight_row_byte_base: offset, %q8_input: buffer, %q8_row_byte_base: offset, %lane: index) -> (f32) { + %bounded_input_size = index.assume %input_size [range(%input_size, 256, 32768), mul(%input_size, 256)] : index + %bounded_lane = index.assume %lane [range(%lane, 0, 31)] : index + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c4 = index.constant 4 : index + %c8 = index.constant 8 : index + %c256 = index.constant 256 : index + %c1023 = index.constant 1023 : index + %c1024 = index.constant 1024 : index + %c0_f32 = scalar.constant 0.0 : f32 + %q4_block_count = index.div %bounded_input_size, %c256 : index + %padded_input_size = index.add %bounded_input_size, %c1023 : index + %iteration_count = index.div %padded_input_size, %c1024 : index + %lane_q4_block = index.div %bounded_lane, %c8 : index + %block_lane0 = index.rem %bounded_lane, %c8 : index + %block_lane = index.assume %block_lane0 [range(%block_lane0, 0, 7)] : index + %sum = scf.for %iteration = [%c0 to %iteration_count step %c1](%iteration_acc = %c0_f32 : f32) -> (f32) unroll { + %iteration_q4_block = index.mul %iteration, %c4 : index + %q4_block0 = index.add %iteration_q4_block, %lane_q4_block : index + %valid_q4_block = index.cmp ult, %q4_block0, %q4_block_count : index + %contribution = scf.if %valid_q4_block -> (f32) { + %q4_block, %bounded_q4_block_count = index.assume %q4_block0, %q4_block_count [lt(%q4_block0, %q4_block_count)] : index, index + %pair = func.call @qwen3_moe_q4k_q8_1_x4_paired_block_lane(%bounded_input_size, %weight, %weight_row_byte_base, %q8_input, %q8_row_byte_base, %q4_block, %block_lane) : (index, buffer, offset, buffer, offset, index, index) -> (f32) + scf.yield %pair : f32 + } else { + scf.yield %c0_f32 : f32 + } + %next = scalar.addf %iteration_acc, %contribution : f32 + scf.yield %next : f32 + } + func.return %sum : f32 +} + +func.def inline @qwen3_moe_q4k_scale_from_header(%scale0: i32, %scale1: i32, %scale2: i32, %q4_group: index) -> (i32, i32) { + %c4 = index.constant 4 : index + %c8 = index.constant 8 : index + %c2_i32 = scalar.constant 2 : i32 + %c4_i32 = scalar.constant 4 : i32 + %c15_i32 = scalar.constant 15 : i32 + %c48_i32 = scalar.constant 48 : i32 + %bounded_group = index.assume %q4_group [range(%q4_group, 0, 7)] : index + %is_low_group = index.cmp ult, %bounded_group, %c4 : index + %scale_lane = index.rem %bounded_group, %c4 : index + %scale_shift_index = index.mul %scale_lane, %c8 : index + %scale_shift = index.cast %scale_shift_index : index to i32 + %high_shift = scalar.addi %scale_shift, %c2_i32 : i32 + %minimum_shift = scalar.addi %scale_shift, %c4_i32 : i32 + %selected_scale_source = scf.select %is_low_group, %scale0, %scale2 : i32 + %selected_minimum_source = scf.select %is_low_group, %scale1, %scale2 : i32 + %selected_scale_high_shift = scf.select %is_low_group, %scale_shift, %high_shift : i32 + %selected_minimum_low_shift = scf.select %is_low_group, %scale_shift, %minimum_shift : i32 + %scale_low0 = scalar.shrui %selected_scale_source, %scale_shift : i32 + %scale_low = scalar.andi %scale_low0, %c15_i32 : i32 + %scale_high0 = scalar.shrui %scale0, %selected_scale_high_shift : i32 + %scale_high = scalar.andi %scale_high0, %c48_i32 : i32 + %scale = scalar.ori %scale_low, %scale_high : i32 + %minimum_low0 = scalar.shrui %selected_minimum_source, %selected_minimum_low_shift : i32 + %minimum_low = scalar.andi %minimum_low0, %c15_i32 : i32 + %minimum_high0 = scalar.shrui %scale1, %selected_scale_high_shift : i32 + %minimum_high = scalar.andi %minimum_high0, %c48_i32 : i32 + %minimum = scalar.ori %minimum_low, %minimum_high : i32 + func.return %scale, %minimum : i32, i32 +} + +amdgpu.target @qwen3_moe_attention_qkv_gfx11_wave32 {subgroup_size = 32} + +config.decl @qwen3_moe.model.hidden_size : %value: index where [range(%value, 128, 32768), mul(%value, 128)] + +config.decl @qwen3_moe.attention.query_size : %value: index where [range(%value, 1, 262144)] + +config.decl @qwen3_moe.attention.key_value_size : %value: index where [range(%value, 1, 262144)] + +config.decl @qwen3_moe.workload.token_capacity : %value: index where [range(%value, 1, 2048)] + +// Selects Q4_K (0) or Q6_K (1) storage for the value projection. +config.decl @qwen3_moe.attention.value_uses_q6 : %value: index where [range(%value, 0, 1)] + +// Paired-nibble Q4_K row contraction provider linked from the dense quantized +// library. +func.decl @qwen3_moe_q4k_q8_1_x4_paired_row_lane(%input_size: index, %weight: buffer, %weight_row_byte_base: offset, %q8_input: buffer, %q8_row_byte_base: offset, %lane: index) -> (f32) + +// Q6_K row contraction provider linked from the GGML quantized library. +func.decl @ggml_q6k_q8_1_x4_row_lane(%input_size: index, %weight: buffer, %weight_row_byte_base: offset, %q8_input: buffer, %q8_row_byte_base: offset, %lane: index) -> (f32) + +// Fused RMSNorm and Q8_1 x4 producer linked from attention preparation. +kernel.decl @qwen3_moe_attention_rmsnorm_quantize_q8_1_x4(%token_count: index) launch(%token_count: index, %input: buffer, %weight: buffer, %q8_output: buffer) + +// Reference entry points used only by differential cases. +kernel.decl @ggml_quantize_q8_1_x4_f32(%token_count: index, %input_size: index) launch(%token_count: index, %input_size: index, %input: buffer, %output: buffer) + +kernel.decl @qwen3_moe_dense_linear_q4k_q8_1_x4(%token_count: index) launch(%token_count: index, %q8_input: buffer, %weight: buffer, %output: buffer) + +kernel.decl @ggml_linear_q6k_q8_1_x4(%token_count: index, %input_size: index, %output_size: index) launch(%token_count: index, %input_size: index, %output_size: index, %q8_input: buffer, %weight: buffer, %output: buffer) + +// Device body shared by the ordinary projection and completion-fused +// postprocess exports. Each caller owns its boundary after the raw row stores. +func.def inline @qwen3_moe_attention_qkv_quantized_body(%value_uses_q6_index: index, %publish_output: i1, %token_count: index, %token0: index, %q8_input: buffer, %query_weight: buffer, %key_weight: buffer, %value_weight: buffer, %query_output: buffer, %key_output: buffer, %value_output: buffer) { + %hidden_size0 = config.get @qwen3_moe.model.hidden_size : index + %query_size0 = config.get @qwen3_moe.attention.query_size : index + %key_value_size0 = config.get @qwen3_moe.attention.key_value_size : index + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 2048)] : index + %hidden_size = index.assume %hidden_size0 [range(%hidden_size0, 128, 32768), mul(%hidden_size0, 128)] : index + %query_size, %key_value_size = index.assume %query_size0, %key_value_size0 [range(%query_size0, 1, 262144), range(%key_value_size0, 1, 262144), mul(%query_size0, %key_value_size0)] : index, index + %channel_tile = kernel.workgroup.id : index + %subgroup0 = kernel.subgroup.id : index + %lane = kernel.subgroup.lane.id : index + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c8 = index.constant 8 : index + %c128 = index.constant 128 : index + %c144_bytes = index.constant 144 : offset + %c210_bytes = index.constant 210 : offset + %c256 = index.constant 256 : index + %c0_i32 = scalar.constant 0 : i32 + %c0_f32 = scalar.constant 0.0 : f32 + %c0_offset = index.constant 0 : offset + %token, %launch_token_count = index.assume %token0, %bounded_token_count [lt(%token0, %bounded_token_count)] : index, index + %subgroup = index.assume %subgroup0 [range(%subgroup0, 0, 7)] : index + %channel_base = index.mul %channel_tile, %c8 : index + %global_channel = index.add %channel_base, %subgroup : index + %key_value_end = index.add %query_size, %key_value_size : index + %key_value_output_size = index.mul %key_value_size, %c2 : index + %total_output_size = index.add %query_size, %key_value_output_size : index + %valid_channel = index.cmp ult, %global_channel, %total_output_size : index + %is_query = index.cmp ult, %global_channel, %query_size : index + %is_key = index.cmp ult, %global_channel, %key_value_end : index + %key_value_channel = index.rem %global_channel, %key_value_size : index + %value_uses_q6 = index.cmp eq, %value_uses_q6_index, %c1 : index + %lane_i32 = index.cast %lane : index to i32 + %is_lane_zero = scalar.cmpi eq, %lane_i32, %c0_i32 : i32 + %quant_block_count = index.div %hidden_size, %c256 : index + %q4_row_bytes = index.scale %quant_block_count, %c144_bytes : index, offset -> offset + %q6_row_bytes = index.scale %quant_block_count, %c210_bytes : index, offset -> offset + %q8_group_count = index.div %hidden_size, %c128 : index + %q8_row_bytes = index.scale %q8_group_count, %c144_bytes : index, offset -> offset + %q8_row_byte_base = index.scale %token, %q8_row_bytes : index, offset -> offset + %q8_noalias, %query_weight_noalias, %key_weight_noalias, %value_weight_noalias, %query_output_noalias, %key_output_noalias, %value_output_noalias = buffer.assume.noalias %q8_input, %query_weight, %key_weight, %value_weight, %query_output, %key_output, %value_output : buffer, buffer, buffer, buffer, buffer, buffer, buffer + %lane_sum = scf.if %publish_output -> (f32) { + %channel_sum = scf.if %valid_channel -> (f32) { + %projection_sum = scf.if %is_query -> (f32) { + %row_byte_base = index.scale %global_channel, %q4_row_bytes : index, offset -> offset + %sum = func.call @qwen3_moe_q4k_q8_1_x4_paired_row_lane(%hidden_size, %query_weight_noalias, %row_byte_base, %q8_noalias, %q8_row_byte_base, %lane) : (index, buffer, offset, buffer, offset, index) -> (f32) + scf.yield %sum : f32 + } else { + %key_or_value_sum = scf.if %is_key -> (f32) { + %row_byte_base = index.scale %key_value_channel, %q4_row_bytes : index, offset -> offset + %sum = func.call @qwen3_moe_q4k_q8_1_x4_paired_row_lane(%hidden_size, %key_weight_noalias, %row_byte_base, %q8_noalias, %q8_row_byte_base, %lane) : (index, buffer, offset, buffer, offset, index) -> (f32) + scf.yield %sum : f32 + } else { + %value_sum = scf.if %value_uses_q6 -> (f32) { + %row_byte_base = index.scale %key_value_channel, %q6_row_bytes : index, offset -> offset + %sum = func.call @ggml_q6k_q8_1_x4_row_lane(%hidden_size, %value_weight_noalias, %row_byte_base, %q8_noalias, %q8_row_byte_base, %lane) : (index, buffer, offset, buffer, offset, index) -> (f32) + scf.yield %sum : f32 + } else { + %row_byte_base = index.scale %key_value_channel, %q4_row_bytes : index, offset -> offset + %sum = func.call @qwen3_moe_q4k_q8_1_x4_paired_row_lane(%hidden_size, %value_weight_noalias, %row_byte_base, %q8_noalias, %q8_row_byte_base, %lane) : (index, buffer, offset, buffer, offset, index) -> (f32) + scf.yield %sum : f32 + } + scf.yield %value_sum : f32 + } + scf.yield %key_or_value_sum : f32 + } + scf.yield %projection_sum : f32 + } else { + scf.yield %c0_f32 : f32 + } + scf.yield %channel_sum : f32 + } else { + scf.yield %c0_f32 : f32 + } + %dot = kernel.subgroup.reduce %lane_sum : f32 + %query_output_view = buffer.view %query_output_noalias[%c0_offset] : buffer -> view<[%launch_token_count]x[%query_size]xf32> + %key_output_view = buffer.view %key_output_noalias[%c0_offset] : buffer -> view<[%launch_token_count]x[%key_value_size]xf32> + %value_output_view = buffer.view %value_output_noalias[%c0_offset] : buffer -> view<[%launch_token_count]x[%key_value_size]xf32> + scf.if %publish_output { + scf.if %valid_channel { + scf.if %is_lane_zero { + scf.if %is_query { + view.store %dot, %query_output_view[%token, %global_channel] : f32, view<[%launch_token_count]x[%query_size]xf32> + } else { + scf.if %is_key { + view.store %dot, %key_output_view[%token, %key_value_channel] : f32, view<[%launch_token_count]x[%key_value_size]xf32> + } else { + view.store %dot, %value_output_view[%token, %key_value_channel] : f32, view<[%launch_token_count]x[%key_value_size]xf32> + } + } + } + } + } + func.return +} + +kernel.def target(@qwen3_moe_attention_qkv_gfx11_wave32) @qwen3_moe_attention_qkv_quantized(%token_count: index) { + %query_size = config.get @qwen3_moe.attention.query_size : index + %key_value_size = config.get @qwen3_moe.attention.key_value_size : index + %token_capacity = config.get @qwen3_moe.workload.token_capacity : index + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c7 = index.constant 7 : index + %c8 = index.constant 8 : index + %c256 = index.constant 256 : index + %key_value_output_size = index.mul %key_value_size, %c2 : index + %output_size = index.add %query_size, %key_value_output_size : index + %padded_output_size = index.add %output_size, %c7 : index + %output_tiles = index.div %padded_output_size, %c8 : index + kernel.launch.config workgroups(%output_tiles, %token_capacity, %c1) workgroup_size(%c256, %c1, %c1) : index +} launch(%token_count: index, %q8_input: buffer, %query_weight: buffer, %key_weight: buffer, %value_weight: buffer, %query_output: buffer, %key_output: buffer, %value_output: buffer) { + %value_uses_q6_index = config.get @qwen3_moe.attention.value_uses_q6 : index + %token_capacity = config.get @qwen3_moe.workload.token_capacity : index + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 2048), le(%token_count, %token_capacity)] : index + %token0 = kernel.workgroup.id : index + %c0 = index.constant 0 : index + %valid_token = index.cmp ult, %token0, %bounded_token_count : index + %safe_token = scf.select %valid_token, %token0, %c0 : index + func.call @qwen3_moe_attention_qkv_quantized_body(%value_uses_q6_index, %valid_token, %bounded_token_count, %safe_token, %q8_input, %query_weight, %key_weight, %value_weight, %query_output, %key_output, %value_output) : (index, i1, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + kernel.return +} + +// Distinct token rows and Q/K/V byte patterns make row ownership and +// binding-domain mistakes observable. Equal synthetic output widths let the +// two Q4_K references share the config-specialized direct provider while still +// crossing every domain edge. +check.case public @qwen3_moe_attention_qkv_q6_differential_case { + %token_count = check.literal value(14) : index + %hidden_size = check.literal value(512) : index + %output_size = check.literal value(65) : index + %input_seed = check.param.seed base(0x514d4f45514b5636) count(1) : i64 + %input = check.generate.random.uniform seed(%input_seed) range(-1.0 to 1.0) : tensor<14x512xf32> + %q8_input = check.generate.fill value(0) : tensor<14x576xi8> + %query_weight = check.generate.fill value(34) : tensor<65x2x144xi8> + %key_weight = check.generate.fill value(35) : tensor<65x2x144xi8> + %value_weight = check.generate.fill value(-86) : tensor<65x2x210xi8> + %expected_query = check.generate.fill value(0.0) : tensor<14x65xf32> + %expected_key = check.generate.fill value(0.0) : tensor<14x65xf32> + %expected_value = check.generate.fill value(0.0) : tensor<14x65xf32> + %actual_query = check.generate.fill value(1.0) : tensor<14x65xf32> + %actual_key = check.generate.fill value(1.0) : tensor<14x65xf32> + %actual_value = check.generate.fill value(1.0) : tensor<14x65xf32> + kernel.launch @ggml_quantize_q8_1_x4_f32[%token_count, %hidden_size](%token_count, %hidden_size, %input, %q8_input) : [index, index](index, index, tensor<14x512xf32>, tensor<14x576xi8>) + kernel.launch @qwen3_moe_dense_linear_q4k_q8_1_x4[%token_count](%token_count, %q8_input, %query_weight, %expected_query) : [index](index, tensor<14x576xi8>, tensor<65x2x144xi8>, tensor<14x65xf32>) + kernel.launch @qwen3_moe_dense_linear_q4k_q8_1_x4[%token_count](%token_count, %q8_input, %key_weight, %expected_key) : [index](index, tensor<14x576xi8>, tensor<65x2x144xi8>, tensor<14x65xf32>) + kernel.launch @ggml_linear_q6k_q8_1_x4[%token_count, %hidden_size, %output_size](%token_count, %hidden_size, %output_size, %q8_input, %value_weight, %expected_value) : [index, index, index](index, index, index, tensor<14x576xi8>, tensor<65x2x210xi8>, tensor<14x65xf32>) + kernel.launch @qwen3_moe_attention_qkv_quantized[%token_count](%token_count, %q8_input, %query_weight, %key_weight, %value_weight, %actual_query, %actual_key, %actual_value) : [index](index, tensor<14x576xi8>, tensor<65x2x144xi8>, tensor<65x2x144xi8>, tensor<65x2x210xi8>, tensor<14x65xf32>, tensor<14x65xf32>, tensor<14x65xf32>) + check.expect.close actual(%actual_query) expected(%expected_query) atol(0.25) rtol(0.01) nan(same) : tensor<14x65xf32> + check.expect.close actual(%actual_key) expected(%expected_key) atol(0.25) rtol(0.01) nan(same) : tensor<14x65xf32> + check.expect.close actual(%actual_value) expected(%expected_value) atol(0.25) rtol(0.01) nan(same) : tensor<14x65xf32> + check.return +} + +check.case public @qwen3_moe_attention_qkv_q4_differential_case { + %token_count = check.literal value(14) : index + %hidden_size = check.literal value(512) : index + %input_seed = check.param.seed base(0x514d4f45514b5634) count(1) : i64 + %input = check.generate.random.uniform seed(%input_seed) range(-1.0 to 1.0) : tensor<14x512xf32> + %q8_input = check.generate.fill value(0) : tensor<14x576xi8> + %query_weight = check.generate.fill value(34) : tensor<65x2x144xi8> + %key_weight = check.generate.fill value(35) : tensor<65x2x144xi8> + %value_weight = check.generate.fill value(36) : tensor<65x2x144xi8> + %expected_query = check.generate.fill value(0.0) : tensor<14x65xf32> + %expected_key = check.generate.fill value(0.0) : tensor<14x65xf32> + %expected_value = check.generate.fill value(0.0) : tensor<14x65xf32> + %actual_query = check.generate.fill value(1.0) : tensor<14x65xf32> + %actual_key = check.generate.fill value(1.0) : tensor<14x65xf32> + %actual_value = check.generate.fill value(1.0) : tensor<14x65xf32> + kernel.launch @ggml_quantize_q8_1_x4_f32[%token_count, %hidden_size](%token_count, %hidden_size, %input, %q8_input) : [index, index](index, index, tensor<14x512xf32>, tensor<14x576xi8>) + kernel.launch @qwen3_moe_dense_linear_q4k_q8_1_x4[%token_count](%token_count, %q8_input, %query_weight, %expected_query) : [index](index, tensor<14x576xi8>, tensor<65x2x144xi8>, tensor<14x65xf32>) + kernel.launch @qwen3_moe_dense_linear_q4k_q8_1_x4[%token_count](%token_count, %q8_input, %key_weight, %expected_key) : [index](index, tensor<14x576xi8>, tensor<65x2x144xi8>, tensor<14x65xf32>) + kernel.launch @qwen3_moe_dense_linear_q4k_q8_1_x4[%token_count](%token_count, %q8_input, %value_weight, %expected_value) : [index](index, tensor<14x576xi8>, tensor<65x2x144xi8>, tensor<14x65xf32>) + kernel.launch @qwen3_moe_attention_qkv_quantized[%token_count](%token_count, %q8_input, %query_weight, %key_weight, %value_weight, %actual_query, %actual_key, %actual_value) : [index](index, tensor<14x576xi8>, tensor<65x2x144xi8>, tensor<65x2x144xi8>, tensor<65x2x144xi8>, tensor<14x65xf32>, tensor<14x65xf32>, tensor<14x65xf32>) + check.expect.close actual(%actual_query) expected(%expected_query) atol(0.25) rtol(0.01) nan(same) : tensor<14x65xf32> + check.expect.close actual(%actual_key) expected(%expected_key) atol(0.25) rtol(0.01) nan(same) : tensor<14x65xf32> + check.expect.close actual(%actual_value) expected(%expected_value) atol(0.25) rtol(0.01) nan(same) : tensor<14x65xf32> + check.return +} + +check.case public @qwen3_moe_attention_qkv_q6_benchmark_case { + %token_count = check.param.choice values([1, 8, 32, 128, 512]) name("token_count") : index + %q8_input = check.generate.fill value(0) : tensor<[%token_count]x2304xi8> + %query_weight = check.generate.fill value(0) : tensor<4096x8x144xi8> + %key_weight = check.generate.fill value(0) : tensor<512x8x144xi8> + %value_weight = check.generate.fill value(0) : tensor<512x8x210xi8> + %query_output = check.generate.fill value(1.0) : tensor<[%token_count]x4096xf32> + %key_output = check.generate.fill value(1.0) : tensor<[%token_count]x512xf32> + %value_output = check.generate.fill value(1.0) : tensor<[%token_count]x512xf32> + %expected_query = check.generate.fill value(0.0) : tensor<[%token_count]x4096xf32> + %expected_key = check.generate.fill value(0.0) : tensor<[%token_count]x512xf32> + %expected_value = check.generate.fill value(0.0) : tensor<[%token_count]x512xf32> + kernel.launch @qwen3_moe_attention_qkv_quantized[%token_count](%token_count, %q8_input, %query_weight, %key_weight, %value_weight, %query_output, %key_output, %value_output) : [index](index, tensor<[%token_count]x2304xi8>, tensor<4096x8x144xi8>, tensor<512x8x144xi8>, tensor<512x8x210xi8>, tensor<[%token_count]x4096xf32>, tensor<[%token_count]x512xf32>, tensor<[%token_count]x512xf32>) + check.expect.close actual(%query_output) expected(%expected_query) atol(0.0) rtol(0.0) nan(same) : tensor<[%token_count]x4096xf32> + check.expect.close actual(%key_output) expected(%expected_key) atol(0.0) rtol(0.0) nan(same) : tensor<[%token_count]x512xf32> + check.expect.close actual(%value_output) expected(%expected_value) atol(0.0) rtol(0.0) nan(same) : tensor<[%token_count]x512xf32> + check.return +} + +check.case public @qwen3_moe_attention_qkv_q4_benchmark_case { + %token_count = check.param.choice values([1, 8, 32, 128, 512]) name("token_count") : index + %q8_input = check.generate.fill value(0) : tensor<[%token_count]x2304xi8> + %query_weight = check.generate.fill value(0) : tensor<4096x8x144xi8> + %key_weight = check.generate.fill value(0) : tensor<512x8x144xi8> + %value_weight = check.generate.fill value(0) : tensor<512x8x144xi8> + %query_output = check.generate.fill value(1.0) : tensor<[%token_count]x4096xf32> + %key_output = check.generate.fill value(1.0) : tensor<[%token_count]x512xf32> + %value_output = check.generate.fill value(1.0) : tensor<[%token_count]x512xf32> + %expected_query = check.generate.fill value(0.0) : tensor<[%token_count]x4096xf32> + %expected_key = check.generate.fill value(0.0) : tensor<[%token_count]x512xf32> + %expected_value = check.generate.fill value(0.0) : tensor<[%token_count]x512xf32> + kernel.launch @qwen3_moe_attention_qkv_quantized[%token_count](%token_count, %q8_input, %query_weight, %key_weight, %value_weight, %query_output, %key_output, %value_output) : [index](index, tensor<[%token_count]x2304xi8>, tensor<4096x8x144xi8>, tensor<512x8x144xi8>, tensor<512x8x144xi8>, tensor<[%token_count]x4096xf32>, tensor<[%token_count]x512xf32>, tensor<[%token_count]x512xf32>) + check.expect.close actual(%query_output) expected(%expected_query) atol(0.0) rtol(0.0) nan(same) : tensor<[%token_count]x4096xf32> + check.expect.close actual(%key_output) expected(%expected_key) atol(0.0) rtol(0.0) nan(same) : tensor<[%token_count]x512xf32> + check.expect.close actual(%value_output) expected(%expected_value) atol(0.0) rtol(0.0) nan(same) : tensor<[%token_count]x512xf32> + check.return +} + +// Production decode and small-prefill boundary. The two calls become two +// serialized dispatches in one reusable command buffer: the first publishes +// the only Q8_1 activation row and the second consumes it for all projections. +check.case public @qwen3_moe_attention_qkv_full_q6_benchmark_case { + %token_count = check.param.choice values([1, 8, 32]) name("token_count") : index + %input = check.generate.fill value(0.0) : tensor<[%token_count]x2048xf32> + %norm_weight = check.generate.fill value(1.0) : tensor<2048xf32> + %q8_input = check.generate.fill value(1) : tensor<[%token_count]x2304xi8> + %query_weight = check.generate.fill value(0) : tensor<4096x8x144xi8> + %key_weight = check.generate.fill value(0) : tensor<512x8x144xi8> + %value_weight = check.generate.fill value(0) : tensor<512x8x210xi8> + %query_output = check.generate.fill value(1.0) : tensor<[%token_count]x4096xf32> + %key_output = check.generate.fill value(1.0) : tensor<[%token_count]x512xf32> + %value_output = check.generate.fill value(1.0) : tensor<[%token_count]x512xf32> + %expected_query = check.generate.fill value(0.0) : tensor<[%token_count]x4096xf32> + %expected_key = check.generate.fill value(0.0) : tensor<[%token_count]x512xf32> + %expected_value = check.generate.fill value(0.0) : tensor<[%token_count]x512xf32> + kernel.launch @qwen3_moe_attention_rmsnorm_quantize_q8_1_x4[%token_count](%token_count, %input, %norm_weight, %q8_input) : [index](index, tensor<[%token_count]x2048xf32>, tensor<2048xf32>, tensor<[%token_count]x2304xi8>) + kernel.launch @qwen3_moe_attention_qkv_quantized[%token_count](%token_count, %q8_input, %query_weight, %key_weight, %value_weight, %query_output, %key_output, %value_output) : [index](index, tensor<[%token_count]x2304xi8>, tensor<4096x8x144xi8>, tensor<512x8x144xi8>, tensor<512x8x210xi8>, tensor<[%token_count]x4096xf32>, tensor<[%token_count]x512xf32>, tensor<[%token_count]x512xf32>) + check.expect.close actual(%query_output) expected(%expected_query) atol(0.0) rtol(0.0) nan(same) : tensor<[%token_count]x4096xf32> + check.expect.close actual(%key_output) expected(%expected_key) atol(0.0) rtol(0.0) nan(same) : tensor<[%token_count]x512xf32> + check.expect.close actual(%value_output) expected(%expected_value) atol(0.0) rtol(0.0) nan(same) : tensor<[%token_count]x512xf32> + check.return +} + +check.case public @qwen3_moe_attention_qkv_full_q4_benchmark_case { + %token_count = check.param.choice values([1, 8, 32]) name("token_count") : index + %input = check.generate.fill value(0.0) : tensor<[%token_count]x2048xf32> + %norm_weight = check.generate.fill value(1.0) : tensor<2048xf32> + %q8_input = check.generate.fill value(1) : tensor<[%token_count]x2304xi8> + %query_weight = check.generate.fill value(0) : tensor<4096x8x144xi8> + %key_weight = check.generate.fill value(0) : tensor<512x8x144xi8> + %value_weight = check.generate.fill value(0) : tensor<512x8x144xi8> + %query_output = check.generate.fill value(1.0) : tensor<[%token_count]x4096xf32> + %key_output = check.generate.fill value(1.0) : tensor<[%token_count]x512xf32> + %value_output = check.generate.fill value(1.0) : tensor<[%token_count]x512xf32> + %expected_query = check.generate.fill value(0.0) : tensor<[%token_count]x4096xf32> + %expected_key = check.generate.fill value(0.0) : tensor<[%token_count]x512xf32> + %expected_value = check.generate.fill value(0.0) : tensor<[%token_count]x512xf32> + kernel.launch @qwen3_moe_attention_rmsnorm_quantize_q8_1_x4[%token_count](%token_count, %input, %norm_weight, %q8_input) : [index](index, tensor<[%token_count]x2048xf32>, tensor<2048xf32>, tensor<[%token_count]x2304xi8>) + kernel.launch @qwen3_moe_attention_qkv_quantized[%token_count](%token_count, %q8_input, %query_weight, %key_weight, %value_weight, %query_output, %key_output, %value_output) : [index](index, tensor<[%token_count]x2304xi8>, tensor<4096x8x144xi8>, tensor<512x8x144xi8>, tensor<512x8x144xi8>, tensor<[%token_count]x4096xf32>, tensor<[%token_count]x512xf32>, tensor<[%token_count]x512xf32>) + check.expect.close actual(%query_output) expected(%expected_query) atol(0.0) rtol(0.0) nan(same) : tensor<[%token_count]x4096xf32> + check.expect.close actual(%key_output) expected(%expected_key) atol(0.0) rtol(0.0) nan(same) : tensor<[%token_count]x512xf32> + check.expect.close actual(%value_output) expected(%expected_value) atol(0.0) rtol(0.0) nan(same) : tensor<[%token_count]x512xf32> + check.return +} + +check.benchmark<@qwen3_moe_attention_qkv_q6_differential_case> @qwen3_moe_attention_qkv_q6_differential + +check.benchmark<@qwen3_moe_attention_qkv_q4_differential_case> @qwen3_moe_attention_qkv_q4_differential + +check.benchmark<@qwen3_moe_attention_qkv_q6_benchmark_case> @qwen3_moe_attention_qkv_q6_decode {token_count = 1} + +check.benchmark<@qwen3_moe_attention_qkv_q6_benchmark_case> @qwen3_moe_attention_qkv_q6_prefill_32 {token_count = 32} + +check.benchmark<@qwen3_moe_attention_qkv_q6_benchmark_case> @qwen3_moe_attention_qkv_q6_prefill_128 {token_count = 128} + +check.benchmark<@qwen3_moe_attention_qkv_q6_benchmark_case> @qwen3_moe_attention_qkv_q6_prefill_512 {token_count = 512} + +check.benchmark<@qwen3_moe_attention_qkv_q4_benchmark_case> @qwen3_moe_attention_qkv_q4_decode {token_count = 1} + +check.benchmark<@qwen3_moe_attention_qkv_q4_benchmark_case> @qwen3_moe_attention_qkv_q4_prefill_32 {token_count = 32} + +check.benchmark<@qwen3_moe_attention_qkv_q4_benchmark_case> @qwen3_moe_attention_qkv_q4_prefill_128 {token_count = 128} + +check.benchmark<@qwen3_moe_attention_qkv_q4_benchmark_case> @qwen3_moe_attention_qkv_q4_prefill_512 {token_count = 512} + +check.benchmark<@qwen3_moe_attention_qkv_full_q6_benchmark_case> @qwen3_moe_attention_qkv_full_q6_decode {token_count = 1} + +check.benchmark<@qwen3_moe_attention_qkv_full_q6_benchmark_case> @qwen3_moe_attention_qkv_full_q6_prefill_32 {token_count = 32} + +check.benchmark<@qwen3_moe_attention_qkv_full_q4_benchmark_case> @qwen3_moe_attention_qkv_full_q4_decode {token_count = 1} + +check.benchmark<@qwen3_moe_attention_qkv_full_q4_benchmark_case> @qwen3_moe_attention_qkv_full_q4_prefill_32 {token_count = 32} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/attention_qkv_same_format_prefill.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/attention_qkv_same_format_prefill.loom new file mode 100644 index 000000000000..cf3fdd148c78 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/attention_qkv_same_format_prefill.loom @@ -0,0 +1,193 @@ +// Same-format Q4_K attention projections over one contiguous Q/K/V tensor. +// +// The fixed parameter layout places Q, K, and Q4_K V rows consecutively. This +// bounded candidate pairs those [5120][2048] weights with a row-interleaved +// [token][5120] raw output so one uniform 80xN grid can reuse the canonical +// dense WMMA body without buffer-valued control flow, completion counters, or +// additional workgroup storage. +amdgpu.target @qwen3_moe_attention_qkv_same_format_prefill_gfx11_wave64 {subgroup_size = 64} + +config.decl @qwen3_moe.model.hidden_size : %value: index where [range(%value, 128, 32768), mul(%value, 128)] + +config.decl @qwen3_moe.attention.query_size : %value: index where [range(%value, 1, 262144)] + +config.decl @qwen3_moe.attention.key_value_size : %value: index where [range(%value, 1, 262144)] + +config.decl @qwen3_moe.workload.token_capacity : %value: index where [range(%value, 1, 2048)] + +// Runs every Q4_K Q/K/V tile in one launch over the aggregate physical rows. +func.decl @qwen3_moe_dense_linear_quantized_f16_wmma_body(%weight_format: index, %token_count: index, %input_size0: index, %output_size0: index, %output_accumulation: index, %channel_tile: index, %token_tile: index, %input: buffer, %weight: buffer, %output: buffer) + +kernel.def target(@qwen3_moe_attention_qkv_same_format_prefill_gfx11_wave64) @qwen3_moe_attention_qkv_q4_prefill_512(%token_count: index) { + %query_size = config.get @qwen3_moe.attention.query_size : index + %key_value_size = config.get @qwen3_moe.attention.key_value_size : index + %token_capacity = config.get @qwen3_moe.workload.token_capacity : index + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c31 = index.constant 31 : index + %c32 = index.constant 32 : index + %c63 = index.constant 63 : index + %c64 = index.constant 64 : index + %c128 = index.constant 128 : index + %key_value_output_size = index.mul %key_value_size, %c2 : index + %combined_output_size = index.add %query_size, %key_value_output_size : index + %padded_output_size = index.add %combined_output_size, %c63 : index + %output_tiles = index.div %padded_output_size, %c64 : index + %padded_token_count = index.add %token_capacity, %c31 : index + %token_tiles = index.div %padded_token_count, %c32 : index + kernel.launch.config workgroups(%output_tiles, %token_tiles, %c1) workgroup_size(%c128, %c1, %c1) : index +} launch(%token_count: index, %input: buffer, %combined_weight: buffer, %combined_output: buffer) { + %hidden_size = config.get @qwen3_moe.model.hidden_size : index + %query_size = config.get @qwen3_moe.attention.query_size : index + %key_value_size = config.get @qwen3_moe.attention.key_value_size : index + %channel_tile = kernel.workgroup.id : index + %token_tile = kernel.workgroup.id : index + %c2 = index.constant 2 : index + %q4 = index.constant 4 : index + %overwrite = index.constant 0 : index + %key_value_output_size = index.mul %key_value_size, %c2 : index + %combined_output_size = index.add %query_size, %key_value_output_size : index + func.call @qwen3_moe_dense_linear_quantized_f16_wmma_body(%q4, %token_count, %hidden_size, %combined_output_size, %overwrite, %channel_tile, %token_tile, %input, %combined_weight, %combined_output) : (index, index, index, index, index, index, index, buffer, buffer, buffer) + kernel.return +} + +// The three segment exports form the exact aggregate-layout baseline. Their +// workgroups and memory accesses match the combined launch; only the launch +// records differ. +kernel.def target(@qwen3_moe_attention_qkv_same_format_prefill_gfx11_wave64) @qwen3_moe_attention_query_q4_aggregate_prefill_512(%token_count: index) { + %query_size = config.get @qwen3_moe.attention.query_size : index + %token_capacity = config.get @qwen3_moe.workload.token_capacity : index + %c1 = index.constant 1 : index + %c31 = index.constant 31 : index + %c32 = index.constant 32 : index + %c63 = index.constant 63 : index + %c64 = index.constant 64 : index + %c128 = index.constant 128 : index + %padded_output_size = index.add %query_size, %c63 : index + %output_tiles = index.div %padded_output_size, %c64 : index + %padded_token_count = index.add %token_capacity, %c31 : index + %token_tiles = index.div %padded_token_count, %c32 : index + kernel.launch.config workgroups(%output_tiles, %token_tiles, %c1) workgroup_size(%c128, %c1, %c1) : index +} launch(%token_count: index, %input: buffer, %combined_weight: buffer, %combined_output: buffer) { + %hidden_size = config.get @qwen3_moe.model.hidden_size : index + %query_size = config.get @qwen3_moe.attention.query_size : index + %key_value_size = config.get @qwen3_moe.attention.key_value_size : index + %channel_tile = kernel.workgroup.id : index + %token_tile = kernel.workgroup.id : index + %c2 = index.constant 2 : index + %q4 = index.constant 4 : index + %overwrite = index.constant 0 : index + %key_value_output_size = index.mul %key_value_size, %c2 : index + %combined_output_size = index.add %query_size, %key_value_output_size : index + func.call @qwen3_moe_dense_linear_quantized_f16_wmma_body(%q4, %token_count, %hidden_size, %combined_output_size, %overwrite, %channel_tile, %token_tile, %input, %combined_weight, %combined_output) : (index, index, index, index, index, index, index, buffer, buffer, buffer) + kernel.return +} + +kernel.def target(@qwen3_moe_attention_qkv_same_format_prefill_gfx11_wave64) @qwen3_moe_attention_key_q4_aggregate_prefill_512(%token_count: index) { + %key_value_size = config.get @qwen3_moe.attention.key_value_size : index + %token_capacity = config.get @qwen3_moe.workload.token_capacity : index + %c1 = index.constant 1 : index + %c31 = index.constant 31 : index + %c32 = index.constant 32 : index + %c63 = index.constant 63 : index + %c64 = index.constant 64 : index + %c128 = index.constant 128 : index + %padded_output_size = index.add %key_value_size, %c63 : index + %output_tiles = index.div %padded_output_size, %c64 : index + %padded_token_count = index.add %token_capacity, %c31 : index + %token_tiles = index.div %padded_token_count, %c32 : index + kernel.launch.config workgroups(%output_tiles, %token_tiles, %c1) workgroup_size(%c128, %c1, %c1) : index +} launch(%token_count: index, %input: buffer, %combined_weight: buffer, %combined_output: buffer) { + %hidden_size = config.get @qwen3_moe.model.hidden_size : index + %query_size = config.get @qwen3_moe.attention.query_size : index + %key_value_size = config.get @qwen3_moe.attention.key_value_size : index + %local_channel_tile = kernel.workgroup.id : index + %token_tile = kernel.workgroup.id : index + %c2 = index.constant 2 : index + %q4 = index.constant 4 : index + %c64 = index.constant 64 : index + %overwrite = index.constant 0 : index + %query_channel_tile_count = index.div %query_size, %c64 : index + %channel_tile = index.add %query_channel_tile_count, %local_channel_tile : index + %key_value_output_size = index.mul %key_value_size, %c2 : index + %combined_output_size = index.add %query_size, %key_value_output_size : index + func.call @qwen3_moe_dense_linear_quantized_f16_wmma_body(%q4, %token_count, %hidden_size, %combined_output_size, %overwrite, %channel_tile, %token_tile, %input, %combined_weight, %combined_output) : (index, index, index, index, index, index, index, buffer, buffer, buffer) + kernel.return +} + +kernel.def target(@qwen3_moe_attention_qkv_same_format_prefill_gfx11_wave64) @qwen3_moe_attention_value_q4_aggregate_prefill_512(%token_count: index) { + %key_value_size = config.get @qwen3_moe.attention.key_value_size : index + %token_capacity = config.get @qwen3_moe.workload.token_capacity : index + %c1 = index.constant 1 : index + %c31 = index.constant 31 : index + %c32 = index.constant 32 : index + %c63 = index.constant 63 : index + %c64 = index.constant 64 : index + %c128 = index.constant 128 : index + %padded_output_size = index.add %key_value_size, %c63 : index + %output_tiles = index.div %padded_output_size, %c64 : index + %padded_token_count = index.add %token_capacity, %c31 : index + %token_tiles = index.div %padded_token_count, %c32 : index + kernel.launch.config workgroups(%output_tiles, %token_tiles, %c1) workgroup_size(%c128, %c1, %c1) : index +} launch(%token_count: index, %input: buffer, %combined_weight: buffer, %combined_output: buffer) { + %hidden_size = config.get @qwen3_moe.model.hidden_size : index + %query_size = config.get @qwen3_moe.attention.query_size : index + %key_value_size = config.get @qwen3_moe.attention.key_value_size : index + %local_channel_tile = kernel.workgroup.id : index + %token_tile = kernel.workgroup.id : index + %c2 = index.constant 2 : index + %q4 = index.constant 4 : index + %c64 = index.constant 64 : index + %overwrite = index.constant 0 : index + %query_channel_tile_count = index.div %query_size, %c64 : index + %key_value_channel_tile_count = index.div %key_value_size, %c64 : index + %key_channel_tile_end = index.add %query_channel_tile_count, %key_value_channel_tile_count : index + %channel_tile = index.add %key_channel_tile_end, %local_channel_tile : index + %key_value_output_size = index.mul %key_value_size, %c2 : index + %combined_output_size = index.add %query_size, %key_value_output_size : index + func.call @qwen3_moe_dense_linear_quantized_f16_wmma_body(%q4, %token_count, %hidden_size, %combined_output_size, %overwrite, %channel_tile, %token_tile, %input, %combined_weight, %combined_output) : (index, index, index, index, index, index, index, buffer, buffer, buffer) + kernel.return +} + +// The production shape crosses every token and channel tile. The aggregate +// output starts with different sentinels so missing or overlapping segment +// ownership is observable even though the packed bytes are uniform. +check.case public @qwen3_moe_attention_qkv_q4_prefill_512_differential_case { + %token_count = check.literal value(512) : index + %input = check.generate.fill value(0.00390625) : tensor<512x2048xf32> + %combined_weight = check.generate.fill value(34) : tensor<5120x8x144xi8> + %expected = check.generate.fill value(-1.0) : tensor<512x5120xf32> + %actual = check.generate.fill value(1.0) : tensor<512x5120xf32> + kernel.launch @qwen3_moe_attention_query_q4_aggregate_prefill_512[%token_count](%token_count, %input, %combined_weight, %expected) : [index](index, tensor<512x2048xf32>, tensor<5120x8x144xi8>, tensor<512x5120xf32>) + kernel.launch @qwen3_moe_attention_key_q4_aggregate_prefill_512[%token_count](%token_count, %input, %combined_weight, %expected) : [index](index, tensor<512x2048xf32>, tensor<5120x8x144xi8>, tensor<512x5120xf32>) + kernel.launch @qwen3_moe_attention_value_q4_aggregate_prefill_512[%token_count](%token_count, %input, %combined_weight, %expected) : [index](index, tensor<512x2048xf32>, tensor<5120x8x144xi8>, tensor<512x5120xf32>) + kernel.launch @qwen3_moe_attention_qkv_q4_prefill_512[%token_count](%token_count, %input, %combined_weight, %actual) : [index](index, tensor<512x2048xf32>, tensor<5120x8x144xi8>, tensor<512x5120xf32>) + check.expect.close actual(%actual) expected(%expected) atol(0.0) rtol(0.0) nan(same) : tensor<512x5120xf32> + check.return +} + +check.case public @qwen3_moe_attention_qkv_q4_prefill_512_composed_benchmark_case { + %token_count = check.literal value(512) : index + %input = check.generate.fill value(0.0) : tensor<512x2048xf32> + %combined_weight = check.generate.fill value(0) : tensor<5120x8x144xi8> + %output = check.generate.fill value(1.0) : tensor<512x5120xf32> + kernel.launch @qwen3_moe_attention_query_q4_aggregate_prefill_512[%token_count](%token_count, %input, %combined_weight, %output) : [index](index, tensor<512x2048xf32>, tensor<5120x8x144xi8>, tensor<512x5120xf32>) + kernel.launch @qwen3_moe_attention_key_q4_aggregate_prefill_512[%token_count](%token_count, %input, %combined_weight, %output) : [index](index, tensor<512x2048xf32>, tensor<5120x8x144xi8>, tensor<512x5120xf32>) + kernel.launch @qwen3_moe_attention_value_q4_aggregate_prefill_512[%token_count](%token_count, %input, %combined_weight, %output) : [index](index, tensor<512x2048xf32>, tensor<5120x8x144xi8>, tensor<512x5120xf32>) + check.return +} + +check.case public @qwen3_moe_attention_qkv_q4_prefill_512_fused_benchmark_case { + %token_count = check.literal value(512) : index + %input = check.generate.fill value(0.0) : tensor<512x2048xf32> + %combined_weight = check.generate.fill value(0) : tensor<5120x8x144xi8> + %output = check.generate.fill value(1.0) : tensor<512x5120xf32> + kernel.launch @qwen3_moe_attention_qkv_q4_prefill_512[%token_count](%token_count, %input, %combined_weight, %output) : [index](index, tensor<512x2048xf32>, tensor<5120x8x144xi8>, tensor<512x5120xf32>) + check.return +} + +check.benchmark<@qwen3_moe_attention_qkv_q4_prefill_512_differential_case> @qwen3_moe_attention_qkv_q4_prefill_512_differential + +check.benchmark<@qwen3_moe_attention_qkv_q4_prefill_512_composed_benchmark_case> @qwen3_moe_attention_qkv_q4_prefill_512_composed + +check.benchmark<@qwen3_moe_attention_qkv_q4_prefill_512_fused_benchmark_case> @qwen3_moe_attention_qkv_q4_prefill_512_fused diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/batched_decode_expert_dispatch.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/batched_decode_expert_dispatch.loom new file mode 100644 index 000000000000..b45ac4ea517d --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/batched_decode_expert_dispatch.loom @@ -0,0 +1,498 @@ +// Builds the GPU-resident expert work queues used by batched decode. +// +// One lane owns each configured expert, counts its assignments, and uses +// workgroup scans to assign deterministic compact offsets. Three configured +// row-count limits classify active experts into four target-selected schedule +// queues. The resulting counts are suitable for device-side indirect launch; +// route identities never cross the host boundary. +// +// assignment_ordinals is expert-contiguous. Each ordinal is the original +// flattened [token][route] position, so consumers can recover token and route +// with division and remainder by the configured route count and scatter +// results directly into the established compact routed-output layout. +// +// queue_descriptors is physically [4][expert_count][4xi32]. Each naturally +// aligned descriptor contains the expert ordinal, assignment base, row count, +// and a reserved zero field. The b128 representation keeps the ABI independent +// of model geometry while giving consumers one native-width descriptor load. +config.decl @qwen3_moe.router.expert_count : %value: index where [range(%value, 32, 512), mul(%value, 32)] + +config.decl @qwen3_moe.router.route_count : %value: index where [range(%value, 1, 32)] + +config.decl @qwen3_moe.workload.token_capacity : %value: index where [range(%value, 1, 2048)] + +config.decl @qwen3_moe.batched_decode.schedule0_row_limit : %value: index where [range(%value, 1, 2048)] + +config.decl @qwen3_moe.batched_decode.schedule1_row_limit : %value: index where [range(%value, 1, 2048)] + +config.decl @qwen3_moe.batched_decode.schedule2_row_limit : %value: index where [range(%value, 1, 2048)] + +amdgpu.target @qwen3_moe_batched_decode_dispatch_gfx11_wave32 {subgroup_size = 32} + +// Decodes the stable batched-decode expert descriptor representation. +func.def inline @qwen3_moe_unpack_batched_decode_expert_descriptor(%descriptor: vector<4xi32>) -> (index, index, index) { + %configured_expert_count = config.get @qwen3_moe.router.expert_count : index + %configured_route_count = config.get @qwen3_moe.router.route_count : index + %configured_token_capacity = config.get @qwen3_moe.workload.token_capacity : index + %assignment_capacity = index.mul %configured_token_capacity, %configured_route_count : index + %expert_i32 = vector.extract %descriptor[0] : vector<4xi32> -> i32 + %assignment_base_i32 = vector.extract %descriptor[1] : vector<4xi32> -> i32 + %row_count_i32 = vector.extract %descriptor[2] : vector<4xi32> -> i32 + %expert0 = index.cast %expert_i32 : i32 to index + %expert, %descriptor_expert_count = index.assume %expert0, %configured_expert_count [range(%expert0, 0, 511), lt(%expert0, %configured_expert_count)] : index, index + %assignment_base0 = index.cast %assignment_base_i32 : i32 to index + %assignment_base, %descriptor_assignment_capacity = index.assume %assignment_base0, %assignment_capacity [range(%assignment_base0, 0, 65535), lt(%assignment_base0, %assignment_capacity)] : index, index + %row_count0 = index.cast %row_count_i32 : i32 to index + %row_count, %descriptor_token_capacity = index.assume %row_count0, %configured_token_capacity [range(%row_count0, 1, 2048), le(%row_count0, %configured_token_capacity)] : index, index + func.return %expert, %assignment_base, %row_count : index, index, index +} + +kernel.def target(@qwen3_moe_batched_decode_dispatch_gfx11_wave32) @qwen3_moe_build_batched_decode_expert_dispatch(%token_count: index, %route_count: index, %route_stride: index, %expert_count: index) { + %configured_expert_count = config.get @qwen3_moe.router.expert_count : index + %c1 = index.constant 1 : index + kernel.launch.config workgroups(%c1, %c1, %c1) workgroup_size(%configured_expert_count, %c1, %c1) : index +} launch(%token_count: index, %route_count: index, %route_stride: index, %expert_count: index, %route_ids: buffer, %assignment_ordinals: buffer, %queue_counts: buffer, %queue_descriptors: buffer) { + %configured_token_capacity0 = config.get @qwen3_moe.workload.token_capacity : index + %bounded_token_count, %configured_token_capacity = index.assume %token_count, %configured_token_capacity0 [range(%token_count, 1, 2048), le(%token_count, %configured_token_capacity0)] : index, index + %configured_route_count0 = config.get @qwen3_moe.router.route_count : index + %bounded_route_count, %configured_route_count = index.assume %route_count, %configured_route_count0 [range(%route_count, 1, 32), eq(%route_count, %configured_route_count0)] : index, index + %bounded_route_stride, %route_row_width = index.assume %route_stride, %bounded_route_count [range(%route_stride, 1, 512), le(%bounded_route_count, %route_stride)] : index, index + %configured_expert_count0 = config.get @qwen3_moe.router.expert_count : index + %bounded_expert_count, %configured_expert_count = index.assume %expert_count, %configured_expert_count0 [range(%expert_count, 32, 512), eq(%expert_count, %configured_expert_count0)] : index, index + %schedule0_row_limit0 = config.get @qwen3_moe.batched_decode.schedule0_row_limit : index + %schedule1_row_limit0 = config.get @qwen3_moe.batched_decode.schedule1_row_limit : index + %schedule2_row_limit0 = config.get @qwen3_moe.batched_decode.schedule2_row_limit : index + %schedule0_row_limit, %schedule1_row_limit, %schedule2_row_limit = index.assume %schedule0_row_limit0, %schedule1_row_limit0, %schedule2_row_limit0 [lt(%schedule0_row_limit0, %schedule1_row_limit0), lt(%schedule1_row_limit0, %schedule2_row_limit0)] : index, index, index + %lane0 = kernel.workitem.id : index + %expert, %launch_expert_count = index.assume %lane0, %bounded_expert_count [lt(%lane0, %bounded_expert_count)] : index, index + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c0_i32 = scalar.constant 0 : i32 + %c1_i32 = scalar.constant 1 : i32 + %c0_offset = index.constant 0 : offset + %assignment_count = index.mul %bounded_token_count, %configured_route_count : index + %route_ids_noalias, %assignment_ordinals_noalias, %queue_counts_noalias, %queue_descriptors_noalias = buffer.assume.noalias %route_ids, %assignment_ordinals, %queue_counts, %queue_descriptors : buffer, buffer, buffer, buffer + %route_view = buffer.view %route_ids_noalias[%c0_offset] : buffer -> view<[%bounded_token_count]x[%bounded_route_stride]xi32> + %assignment_view = buffer.view %assignment_ordinals_noalias[%c0_offset] : buffer -> view<[%assignment_count]xi32> + %queue_count_view = buffer.view %queue_counts_noalias[%c0_offset] : buffer -> view<4xi32> + %queue_descriptor_view = buffer.view %queue_descriptors_noalias[%c0_offset] : buffer -> view<4x[%configured_expert_count]x4xi32> + %expert_i32 = index.cast %expert : index to i32 + + // Every lane walks the same tiny, cache-resident route plane. The serial + // comparisons are substantially cheaper than a host-visible sort or atomics + // and make one lane the sole owner of each expert's count. + %expert_assignment_count = scf.for %assignment = [%c0 to %assignment_count step %c1](%matched_count = %c0_i32 : i32) -> (i32) { + %token0 = index.div %assignment, %configured_route_count : index + %route0 = index.rem %assignment, %configured_route_count : index + %token, %route_token_count = index.assume %token0, %bounded_token_count [lt(%token0, %bounded_token_count)] : index, index + %route, %route_row_stride = index.assume %route0, %bounded_route_stride [lt(%route0, %bounded_route_stride)] : index, index + %route_expert_i32 = view.load %route_view[%token, %route] : view<[%bounded_token_count]x[%bounded_route_stride]xi32> -> i32 + %matches = scalar.cmpi eq, %route_expert_i32, %expert_i32 : i32 + %match_increment = scf.select %matches, %c1_i32, %c0_i32 : i32 + %next_matched_count = scalar.addi %matched_count, %match_increment : i32 + scf.yield %next_matched_count : i32 + } + + %expert_assignment_base_i32 = kernel.workgroup.scan %expert_assignment_count {direction = forward, mode = exclusive} : i32 + + // A second pass publishes stable original assignment ordinals into the + // expert-contiguous permutation. This metadata remains cache-resident for + // practical decode batches. + %published_count = scf.for %assignment = [%c0 to %assignment_count step %c1](%matched_count = %c0_i32 : i32) -> (i32) { + %token0 = index.div %assignment, %configured_route_count : index + %route0 = index.rem %assignment, %configured_route_count : index + %token, %route_token_count = index.assume %token0, %bounded_token_count [lt(%token0, %bounded_token_count)] : index, index + %route, %route_row_stride = index.assume %route0, %bounded_route_stride [lt(%route0, %bounded_route_stride)] : index, index + %route_expert_i32 = view.load %route_view[%token, %route] : view<[%bounded_token_count]x[%bounded_route_stride]xi32> -> i32 + %matches = scalar.cmpi eq, %route_expert_i32, %expert_i32 : i32 + scf.if %matches { + %compact_ordinal_i32 = scalar.addi %expert_assignment_base_i32, %matched_count : i32 + %compact_ordinal0 = index.cast %compact_ordinal_i32 : i32 to index + %bounded_compact_ordinal, %bounded_assignment_count = index.assume %compact_ordinal0, %assignment_count [range(%compact_ordinal0, 0, 65535), lt(%compact_ordinal0, %assignment_count)] : index, index + %assignment_i32 = index.cast %assignment : index to i32 + view.store %assignment_i32, %assignment_view[%bounded_compact_ordinal] : i32, view<[%assignment_count]xi32> + } + %match_increment = scf.select %matches, %c1_i32, %c0_i32 : i32 + %next_matched_count = scalar.addi %matched_count, %match_increment : i32 + scf.yield %next_matched_count : i32 + } + + // Top-k route IDs are unique within a token, so one expert owns at most the + // configured token capacity. Target-selected row limits partition that + // range without changing this producer or its ABI. + %published_count0 = index.cast %published_count : i32 to index + %published_row_count, %published_token_capacity = index.assume %published_count0, %configured_token_capacity [range(%published_count0, 0, 2048), le(%published_count0, %configured_token_capacity)] : index, index + %has_assignments = index.cmp ugt, %published_row_count, %c0 : index + %at_most_schedule0 = index.cmp ule, %published_row_count, %schedule0_row_limit : index + %above_schedule0 = index.cmp ugt, %published_row_count, %schedule0_row_limit : index + %at_most_schedule1 = index.cmp ule, %published_row_count, %schedule1_row_limit : index + %above_schedule1 = index.cmp ugt, %published_row_count, %schedule1_row_limit : index + %at_most_schedule2 = index.cmp ule, %published_row_count, %schedule2_row_limit : index + %is_schedule0 = scalar.andi %has_assignments, %at_most_schedule0 : i1 + %is_schedule1 = scalar.andi %above_schedule0, %at_most_schedule1 : i1 + %is_schedule2 = scalar.andi %above_schedule1, %at_most_schedule2 : i1 + %is_schedule3 = index.cmp ugt, %published_row_count, %schedule2_row_limit : index + %schedule0_i32 = scf.select %is_schedule0, %c1_i32, %c0_i32 : i32 + %schedule1_i32 = scf.select %is_schedule1, %c1_i32, %c0_i32 : i32 + %schedule2_i32 = scf.select %is_schedule2, %c1_i32, %c0_i32 : i32 + %schedule3_i32 = scf.select %is_schedule3, %c1_i32, %c0_i32 : i32 + + %schedule0_ordinal_i32 = kernel.workgroup.scan %schedule0_i32 {direction = forward, mode = exclusive} : i32 + %schedule1_ordinal_i32 = kernel.workgroup.scan %schedule1_i32 {direction = forward, mode = exclusive} : i32 + %schedule2_ordinal_i32 = kernel.workgroup.scan %schedule2_i32 {direction = forward, mode = exclusive} : i32 + %schedule3_ordinal_i32 = kernel.workgroup.scan %schedule3_i32 {direction = forward, mode = exclusive} : i32 + %schedule0_count = kernel.workgroup.reduce %schedule0_i32 : i32 + %schedule1_count = kernel.workgroup.reduce %schedule1_i32 : i32 + %schedule2_count = kernel.workgroup.reduce %schedule2_i32 : i32 + %schedule3_count = kernel.workgroup.reduce %schedule3_i32 : i32 + %descriptor = vector.from_elements %expert_i32, %expert_assignment_base_i32, %published_count, %c0_i32 : vector<4xi32> + + // Inactive lanes use ordinal zero so every physical origin is in bounds; + // the vector masks suppress their stores. This keeps the descriptor write a + // native b128 operation without placing it inside divergent control flow. + %safe_schedule0_ordinal_i32 = scf.select %is_schedule0, %schedule0_ordinal_i32, %c0_i32 : i32 + %schedule0_ordinal0 = index.cast %safe_schedule0_ordinal_i32 : i32 to index + %schedule0_ordinal, %schedule0_queue_capacity = index.assume %schedule0_ordinal0, %configured_expert_count [range(%schedule0_ordinal0, 0, 511), lt(%schedule0_ordinal0, %configured_expert_count)] : index, index + %schedule0_mask = vector.splat %is_schedule0 : vector<4xi1> + vector.store.mask %descriptor, %queue_descriptor_view[0, %schedule0_ordinal, 0], %schedule0_mask : vector<4xi32>, view<4x[%configured_expert_count]x4xi32>, vector<4xi1> + %safe_schedule1_ordinal_i32 = scf.select %is_schedule1, %schedule1_ordinal_i32, %c0_i32 : i32 + %schedule1_ordinal0 = index.cast %safe_schedule1_ordinal_i32 : i32 to index + %schedule1_ordinal, %schedule1_queue_capacity = index.assume %schedule1_ordinal0, %configured_expert_count [range(%schedule1_ordinal0, 0, 511), lt(%schedule1_ordinal0, %configured_expert_count)] : index, index + %schedule1_mask = vector.splat %is_schedule1 : vector<4xi1> + vector.store.mask %descriptor, %queue_descriptor_view[1, %schedule1_ordinal, 0], %schedule1_mask : vector<4xi32>, view<4x[%configured_expert_count]x4xi32>, vector<4xi1> + %safe_schedule2_ordinal_i32 = scf.select %is_schedule2, %schedule2_ordinal_i32, %c0_i32 : i32 + %schedule2_ordinal0 = index.cast %safe_schedule2_ordinal_i32 : i32 to index + %schedule2_ordinal, %schedule2_queue_capacity = index.assume %schedule2_ordinal0, %configured_expert_count [range(%schedule2_ordinal0, 0, 511), lt(%schedule2_ordinal0, %configured_expert_count)] : index, index + %schedule2_mask = vector.splat %is_schedule2 : vector<4xi1> + vector.store.mask %descriptor, %queue_descriptor_view[2, %schedule2_ordinal, 0], %schedule2_mask : vector<4xi32>, view<4x[%configured_expert_count]x4xi32>, vector<4xi1> + %safe_schedule3_ordinal_i32 = scf.select %is_schedule3, %schedule3_ordinal_i32, %c0_i32 : i32 + %schedule3_ordinal0 = index.cast %safe_schedule3_ordinal_i32 : i32 to index + %schedule3_ordinal, %schedule3_queue_capacity = index.assume %schedule3_ordinal0, %configured_expert_count [range(%schedule3_ordinal0, 0, 511), lt(%schedule3_ordinal0, %configured_expert_count)] : index, index + %schedule3_mask = vector.splat %is_schedule3 : vector<4xi1> + vector.store.mask %descriptor, %queue_descriptor_view[3, %schedule3_ordinal, 0], %schedule3_mask : vector<4xi32>, view<4x[%configured_expert_count]x4xi32>, vector<4xi1> + + %is_lane_zero = index.cmp eq, %expert, %c0 : index + scf.if %is_lane_zero { + view.store %schedule0_count, %queue_count_view[0] : i32, view<4xi32> + view.store %schedule1_count, %queue_count_view[1] : i32, view<4xi32> + view.store %schedule2_count, %queue_count_view[2] : i32, view<4xi32> + view.store %schedule3_count, %queue_count_view[3] : i32, view<4xi32> + } + kernel.return +} + +// Serial specification oracle for exact differential tests. Its ownership and +// control flow deliberately differ from the parallel scan implementation: +// one workitem visits experts in order and carries all compact offsets and +// queue tails explicitly. +kernel.def target(@qwen3_moe_batched_decode_dispatch_gfx11_wave32) @qwen3_moe_build_batched_decode_expert_dispatch_reference(%token_count: index, %route_count: index, %route_stride: index, %expert_count: index) { + %c1 = index.constant 1 : index + kernel.launch.config workgroups(%c1, %c1, %c1) workgroup_size(%c1, %c1, %c1) : index +} launch(%token_count: index, %route_count: index, %route_stride: index, %expert_count: index, %route_ids: buffer, %assignment_ordinals: buffer, %queue_counts: buffer, %queue_descriptors: buffer) { + %configured_token_capacity0 = config.get @qwen3_moe.workload.token_capacity : index + %bounded_token_count, %configured_token_capacity = index.assume %token_count, %configured_token_capacity0 [range(%token_count, 1, 2048), le(%token_count, %configured_token_capacity0)] : index, index + %configured_route_count0 = config.get @qwen3_moe.router.route_count : index + %bounded_route_count, %configured_route_count = index.assume %route_count, %configured_route_count0 [range(%route_count, 1, 32), eq(%route_count, %configured_route_count0)] : index, index + %bounded_route_stride, %route_row_width = index.assume %route_stride, %bounded_route_count [range(%route_stride, 1, 512), le(%bounded_route_count, %route_stride)] : index, index + %configured_expert_count0 = config.get @qwen3_moe.router.expert_count : index + %bounded_expert_count, %configured_expert_count = index.assume %expert_count, %configured_expert_count0 [range(%expert_count, 32, 512), eq(%expert_count, %configured_expert_count0)] : index, index + %schedule0_row_limit0 = config.get @qwen3_moe.batched_decode.schedule0_row_limit : index + %schedule1_row_limit0 = config.get @qwen3_moe.batched_decode.schedule1_row_limit : index + %schedule2_row_limit0 = config.get @qwen3_moe.batched_decode.schedule2_row_limit : index + %schedule0_row_limit, %schedule1_row_limit, %schedule2_row_limit = index.assume %schedule0_row_limit0, %schedule1_row_limit0, %schedule2_row_limit0 [lt(%schedule0_row_limit0, %schedule1_row_limit0), lt(%schedule1_row_limit0, %schedule2_row_limit0)] : index, index, index + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c0_i32 = scalar.constant 0 : i32 + %c1_i32 = scalar.constant 1 : i32 + %c0_offset = index.constant 0 : offset + %assignment_count = index.mul %bounded_token_count, %configured_route_count : index + %route_ids_noalias, %assignment_ordinals_noalias, %queue_counts_noalias, %queue_descriptors_noalias = buffer.assume.noalias %route_ids, %assignment_ordinals, %queue_counts, %queue_descriptors : buffer, buffer, buffer, buffer + %route_view = buffer.view %route_ids_noalias[%c0_offset] : buffer -> view<[%bounded_token_count]x[%bounded_route_stride]xi32> + %assignment_view = buffer.view %assignment_ordinals_noalias[%c0_offset] : buffer -> view<[%assignment_count]xi32> + %queue_count_view = buffer.view %queue_counts_noalias[%c0_offset] : buffer -> view<4xi32> + %queue_descriptor_view = buffer.view %queue_descriptors_noalias[%c0_offset] : buffer -> view<4x[%configured_expert_count]x4xi32> + + %final_assignment_base, %final_schedule0_count, %final_schedule1_count, %final_schedule2_count, %final_schedule3_count = scf.for %expert = [%c0 to %bounded_expert_count step %c1](%assignment_base = %c0_i32 : i32, %schedule0_count = %c0_i32 : i32, %schedule1_count = %c0_i32 : i32, %schedule2_count = %c0_i32 : i32, %schedule3_count = %c0_i32 : i32) -> (i32, i32, i32, i32, i32) { + %expert_i32 = index.cast %expert : index to i32 + %expert_assignment_count = scf.for %assignment = [%c0 to %assignment_count step %c1](%matched_count = %c0_i32 : i32) -> (i32) { + %token0 = index.div %assignment, %configured_route_count : index + %route0 = index.rem %assignment, %configured_route_count : index + %token, %route_token_count = index.assume %token0, %bounded_token_count [lt(%token0, %bounded_token_count)] : index, index + %route, %route_row_stride = index.assume %route0, %bounded_route_stride [lt(%route0, %bounded_route_stride)] : index, index + %route_expert_i32 = view.load %route_view[%token, %route] : view<[%bounded_token_count]x[%bounded_route_stride]xi32> -> i32 + %matches = scalar.cmpi eq, %route_expert_i32, %expert_i32 : i32 + scf.if %matches { + %compact_ordinal_i32 = scalar.addi %assignment_base, %matched_count : i32 + %compact_ordinal0 = index.cast %compact_ordinal_i32 : i32 to index + %bounded_compact_ordinal, %bounded_assignment_count = index.assume %compact_ordinal0, %assignment_count [range(%compact_ordinal0, 0, 65535), lt(%compact_ordinal0, %assignment_count)] : index, index + %assignment_i32 = index.cast %assignment : index to i32 + view.store %assignment_i32, %assignment_view[%bounded_compact_ordinal] : i32, view<[%assignment_count]xi32> + } + %match_increment = scf.select %matches, %c1_i32, %c0_i32 : i32 + %next_matched_count = scalar.addi %matched_count, %match_increment : i32 + scf.yield %next_matched_count : i32 + } + + %expert_assignment_count0 = index.cast %expert_assignment_count : i32 to index + %expert_row_count, %expert_token_capacity = index.assume %expert_assignment_count0, %configured_token_capacity [range(%expert_assignment_count0, 0, 2048), le(%expert_assignment_count0, %configured_token_capacity)] : index, index + %has_assignments = index.cmp ugt, %expert_row_count, %c0 : index + %at_most_schedule0 = index.cmp ule, %expert_row_count, %schedule0_row_limit : index + %above_schedule0 = index.cmp ugt, %expert_row_count, %schedule0_row_limit : index + %at_most_schedule1 = index.cmp ule, %expert_row_count, %schedule1_row_limit : index + %above_schedule1 = index.cmp ugt, %expert_row_count, %schedule1_row_limit : index + %at_most_schedule2 = index.cmp ule, %expert_row_count, %schedule2_row_limit : index + %is_schedule0 = scalar.andi %has_assignments, %at_most_schedule0 : i1 + %is_schedule1 = scalar.andi %above_schedule0, %at_most_schedule1 : i1 + %is_schedule2 = scalar.andi %above_schedule1, %at_most_schedule2 : i1 + %is_schedule3 = index.cmp ugt, %expert_row_count, %schedule2_row_limit : index + %descriptor = vector.from_elements %expert_i32, %assignment_base, %expert_assignment_count, %c0_i32 : vector<4xi32> + + // The scalar oracle retains its loop-carried queue tails. Masked b128 + // stores express conditional publication without divergent branch regions, + // while safe ordinal zero keeps inactive physical origins in bounds. + %safe_schedule0_count = scf.select %is_schedule0, %schedule0_count, %c0_i32 : i32 + %schedule0_ordinal0 = index.cast %safe_schedule0_count : i32 to index + %schedule0_ordinal, %schedule0_capacity = index.assume %schedule0_ordinal0, %configured_expert_count [range(%schedule0_ordinal0, 0, 511), lt(%schedule0_ordinal0, %configured_expert_count)] : index, index + %schedule0_mask = vector.splat %is_schedule0 : vector<4xi1> + vector.store.mask %descriptor, %queue_descriptor_view[0, %schedule0_ordinal, 0], %schedule0_mask : vector<4xi32>, view<4x[%configured_expert_count]x4xi32>, vector<4xi1> + %safe_schedule1_count = scf.select %is_schedule1, %schedule1_count, %c0_i32 : i32 + %schedule1_ordinal0 = index.cast %safe_schedule1_count : i32 to index + %schedule1_ordinal, %schedule1_capacity = index.assume %schedule1_ordinal0, %configured_expert_count [range(%schedule1_ordinal0, 0, 511), lt(%schedule1_ordinal0, %configured_expert_count)] : index, index + %schedule1_mask = vector.splat %is_schedule1 : vector<4xi1> + vector.store.mask %descriptor, %queue_descriptor_view[1, %schedule1_ordinal, 0], %schedule1_mask : vector<4xi32>, view<4x[%configured_expert_count]x4xi32>, vector<4xi1> + %safe_schedule2_count = scf.select %is_schedule2, %schedule2_count, %c0_i32 : i32 + %schedule2_ordinal0 = index.cast %safe_schedule2_count : i32 to index + %schedule2_ordinal, %schedule2_capacity = index.assume %schedule2_ordinal0, %configured_expert_count [range(%schedule2_ordinal0, 0, 511), lt(%schedule2_ordinal0, %configured_expert_count)] : index, index + %schedule2_mask = vector.splat %is_schedule2 : vector<4xi1> + vector.store.mask %descriptor, %queue_descriptor_view[2, %schedule2_ordinal, 0], %schedule2_mask : vector<4xi32>, view<4x[%configured_expert_count]x4xi32>, vector<4xi1> + %safe_schedule3_count = scf.select %is_schedule3, %schedule3_count, %c0_i32 : i32 + %schedule3_ordinal0 = index.cast %safe_schedule3_count : i32 to index + %schedule3_ordinal, %schedule3_capacity = index.assume %schedule3_ordinal0, %configured_expert_count [range(%schedule3_ordinal0, 0, 511), lt(%schedule3_ordinal0, %configured_expert_count)] : index, index + %schedule3_mask = vector.splat %is_schedule3 : vector<4xi1> + vector.store.mask %descriptor, %queue_descriptor_view[3, %schedule3_ordinal, 0], %schedule3_mask : vector<4xi32>, view<4x[%configured_expert_count]x4xi32>, vector<4xi1> + + %schedule0_increment = scf.select %is_schedule0, %c1_i32, %c0_i32 : i32 + %schedule1_increment = scf.select %is_schedule1, %c1_i32, %c0_i32 : i32 + %schedule2_increment = scf.select %is_schedule2, %c1_i32, %c0_i32 : i32 + %schedule3_increment = scf.select %is_schedule3, %c1_i32, %c0_i32 : i32 + %next_assignment_base = scalar.addi %assignment_base, %expert_assignment_count : i32 + %next_schedule0_count = scalar.addi %schedule0_count, %schedule0_increment : i32 + %next_schedule1_count = scalar.addi %schedule1_count, %schedule1_increment : i32 + %next_schedule2_count = scalar.addi %schedule2_count, %schedule2_increment : i32 + %next_schedule3_count = scalar.addi %schedule3_count, %schedule3_increment : i32 + scf.yield %next_assignment_base, %next_schedule0_count, %next_schedule1_count, %next_schedule2_count, %next_schedule3_count : i32, i32, i32, i32, i32 + } + + view.store %final_schedule0_count, %queue_count_view[0] : i32, view<4xi32> + view.store %final_schedule1_count, %queue_count_view[1] : i32, view<4xi32> + view.store %final_schedule2_count, %queue_count_view[2] : i32, view<4xi32> + view.store %final_schedule3_count, %queue_count_view[3] : i32, view<4xi32> + kernel.return +} + +check.case public @qwen3_moe_batched_decode_expert_dispatch_singleton_case { + %token_count = check.literal value(16) : index + %route_count = check.literal value(8) : index + %route_stride = check.literal value(8) : index + %expert_count = check.literal value(128) : index + %route_ids = check.generate.iota offset(0) step(1) period(128) : tensor<16x8xi32> + %actual_assignments = check.generate.fill value(-1) : tensor<128xi32> + %actual_counts = check.generate.fill value(-1) : tensor<4xi32> + %actual_descriptors = check.generate.fill value(-1) : tensor<4x128x4xi32> + %expected_assignments = check.generate.fill value(-1) : tensor<128xi32> + %expected_counts = check.generate.fill value(-1) : tensor<4xi32> + %expected_descriptors = check.generate.fill value(-1) : tensor<4x128x4xi32> + kernel.launch @qwen3_moe_build_batched_decode_expert_dispatch[%token_count, %route_count, %route_stride, %expert_count](%token_count, %route_count, %route_stride, %expert_count, %route_ids, %actual_assignments, %actual_counts, %actual_descriptors) : [index, index, index, index](index, index, index, index, tensor<16x8xi32>, tensor<128xi32>, tensor<4xi32>, tensor<4x128x4xi32>) + kernel.launch @qwen3_moe_build_batched_decode_expert_dispatch_reference[%token_count, %route_count, %route_stride, %expert_count](%token_count, %route_count, %route_stride, %expert_count, %route_ids, %expected_assignments, %expected_counts, %expected_descriptors) : [index, index, index, index](index, index, index, index, tensor<16x8xi32>, tensor<128xi32>, tensor<4xi32>, tensor<4x128x4xi32>) + check.expect.equal actual(%actual_assignments) expected(%expected_assignments) : tensor<128xi32> + check.expect.equal actual(%actual_counts) expected(%expected_counts) : tensor<4xi32> + check.expect.equal actual(%actual_descriptors) expected(%expected_descriptors) : tensor<4x128x4xi32> + check.return +} + +check.case public @qwen3_moe_batched_decode_expert_dispatch_pair_case { + %token_count = check.literal value(16) : index + %route_count = check.literal value(8) : index + %route_stride = check.literal value(8) : index + %expert_count = check.literal value(128) : index + %route_ids = check.generate.iota offset(0) step(1) period(64) : tensor<16x8xi32> + %actual_assignments = check.generate.fill value(-1) : tensor<128xi32> + %actual_counts = check.generate.fill value(-1) : tensor<4xi32> + %actual_descriptors = check.generate.fill value(-1) : tensor<4x128x4xi32> + %expected_assignments = check.generate.fill value(-1) : tensor<128xi32> + %expected_counts = check.generate.fill value(-1) : tensor<4xi32> + %expected_descriptors = check.generate.fill value(-1) : tensor<4x128x4xi32> + kernel.launch @qwen3_moe_build_batched_decode_expert_dispatch[%token_count, %route_count, %route_stride, %expert_count](%token_count, %route_count, %route_stride, %expert_count, %route_ids, %actual_assignments, %actual_counts, %actual_descriptors) : [index, index, index, index](index, index, index, index, tensor<16x8xi32>, tensor<128xi32>, tensor<4xi32>, tensor<4x128x4xi32>) + kernel.launch @qwen3_moe_build_batched_decode_expert_dispatch_reference[%token_count, %route_count, %route_stride, %expert_count](%token_count, %route_count, %route_stride, %expert_count, %route_ids, %expected_assignments, %expected_counts, %expected_descriptors) : [index, index, index, index](index, index, index, index, tensor<16x8xi32>, tensor<128xi32>, tensor<4xi32>, tensor<4x128x4xi32>) + check.expect.equal actual(%actual_assignments) expected(%expected_assignments) : tensor<128xi32> + check.expect.equal actual(%actual_counts) expected(%expected_counts) : tensor<4xi32> + check.expect.equal actual(%actual_descriptors) expected(%expected_descriptors) : tensor<4x128x4xi32> + check.return +} + +check.case public @qwen3_moe_batched_decode_expert_dispatch_three_row_case { + %token_count = check.literal value(16) : index + %route_count = check.literal value(8) : index + %route_stride = check.literal value(8) : index + %expert_count = check.literal value(128) : index + %route_ids = check.generate.iota offset(0) step(1) period(43) : tensor<16x8xi32> + %actual_assignments = check.generate.fill value(-1) : tensor<128xi32> + %actual_counts = check.generate.fill value(-1) : tensor<4xi32> + %actual_descriptors = check.generate.fill value(-1) : tensor<4x128x4xi32> + %expected_assignments = check.generate.fill value(-1) : tensor<128xi32> + %expected_counts = check.generate.fill value(-1) : tensor<4xi32> + %expected_descriptors = check.generate.fill value(-1) : tensor<4x128x4xi32> + kernel.launch @qwen3_moe_build_batched_decode_expert_dispatch[%token_count, %route_count, %route_stride, %expert_count](%token_count, %route_count, %route_stride, %expert_count, %route_ids, %actual_assignments, %actual_counts, %actual_descriptors) : [index, index, index, index](index, index, index, index, tensor<16x8xi32>, tensor<128xi32>, tensor<4xi32>, tensor<4x128x4xi32>) + kernel.launch @qwen3_moe_build_batched_decode_expert_dispatch_reference[%token_count, %route_count, %route_stride, %expert_count](%token_count, %route_count, %route_stride, %expert_count, %route_ids, %expected_assignments, %expected_counts, %expected_descriptors) : [index, index, index, index](index, index, index, index, tensor<16x8xi32>, tensor<128xi32>, tensor<4xi32>, tensor<4x128x4xi32>) + check.expect.equal actual(%actual_assignments) expected(%expected_assignments) : tensor<128xi32> + check.expect.equal actual(%actual_counts) expected(%expected_counts) : tensor<4xi32> + check.expect.equal actual(%actual_descriptors) expected(%expected_descriptors) : tensor<4x128x4xi32> + check.return +} + +check.case public @qwen3_moe_batched_decode_expert_dispatch_four_row_case { + %token_count = check.literal value(16) : index + %route_count = check.literal value(8) : index + %route_stride = check.literal value(8) : index + %expert_count = check.literal value(128) : index + %route_ids = check.generate.iota offset(0) step(1) period(32) : tensor<16x8xi32> + %actual_assignments = check.generate.fill value(-1) : tensor<128xi32> + %actual_counts = check.generate.fill value(-1) : tensor<4xi32> + %actual_descriptors = check.generate.fill value(-1) : tensor<4x128x4xi32> + %expected_assignments = check.generate.fill value(-1) : tensor<128xi32> + %expected_counts = check.generate.fill value(-1) : tensor<4xi32> + %expected_descriptors = check.generate.fill value(-1) : tensor<4x128x4xi32> + kernel.launch @qwen3_moe_build_batched_decode_expert_dispatch[%token_count, %route_count, %route_stride, %expert_count](%token_count, %route_count, %route_stride, %expert_count, %route_ids, %actual_assignments, %actual_counts, %actual_descriptors) : [index, index, index, index](index, index, index, index, tensor<16x8xi32>, tensor<128xi32>, tensor<4xi32>, tensor<4x128x4xi32>) + kernel.launch @qwen3_moe_build_batched_decode_expert_dispatch_reference[%token_count, %route_count, %route_stride, %expert_count](%token_count, %route_count, %route_stride, %expert_count, %route_ids, %expected_assignments, %expected_counts, %expected_descriptors) : [index, index, index, index](index, index, index, index, tensor<16x8xi32>, tensor<128xi32>, tensor<4xi32>, tensor<4x128x4xi32>) + check.expect.equal actual(%actual_assignments) expected(%expected_assignments) : tensor<128xi32> + check.expect.equal actual(%actual_counts) expected(%expected_counts) : tensor<4xi32> + check.expect.equal actual(%actual_descriptors) expected(%expected_descriptors) : tensor<4x128x4xi32> + check.return +} + +check.case public @qwen3_moe_batched_decode_expert_dispatch_five_to_six_row_case { + %token_count = check.literal value(16) : index + %route_count = check.literal value(8) : index + %route_stride = check.literal value(8) : index + %expert_count = check.literal value(128) : index + %route_ids = check.generate.iota offset(0) step(1) period(25) : tensor<16x8xi32> + %actual_assignments = check.generate.fill value(-1) : tensor<128xi32> + %actual_counts = check.generate.fill value(-1) : tensor<4xi32> + %actual_descriptors = check.generate.fill value(-1) : tensor<4x128x4xi32> + %expected_assignments = check.generate.fill value(-1) : tensor<128xi32> + %expected_counts = check.generate.fill value(-1) : tensor<4xi32> + %expected_descriptors = check.generate.fill value(-1) : tensor<4x128x4xi32> + kernel.launch @qwen3_moe_build_batched_decode_expert_dispatch[%token_count, %route_count, %route_stride, %expert_count](%token_count, %route_count, %route_stride, %expert_count, %route_ids, %actual_assignments, %actual_counts, %actual_descriptors) : [index, index, index, index](index, index, index, index, tensor<16x8xi32>, tensor<128xi32>, tensor<4xi32>, tensor<4x128x4xi32>) + kernel.launch @qwen3_moe_build_batched_decode_expert_dispatch_reference[%token_count, %route_count, %route_stride, %expert_count](%token_count, %route_count, %route_stride, %expert_count, %route_ids, %expected_assignments, %expected_counts, %expected_descriptors) : [index, index, index, index](index, index, index, index, tensor<16x8xi32>, tensor<128xi32>, tensor<4xi32>, tensor<4x128x4xi32>) + check.expect.equal actual(%actual_assignments) expected(%expected_assignments) : tensor<128xi32> + check.expect.equal actual(%actual_counts) expected(%expected_counts) : tensor<4xi32> + check.expect.equal actual(%actual_descriptors) expected(%expected_descriptors) : tensor<4x128x4xi32> + check.return +} + +check.case public @qwen3_moe_batched_decode_expert_dispatch_fourteen_to_fifteen_row_case { + %token_count = check.literal value(16) : index + %route_count = check.literal value(8) : index + %route_stride = check.literal value(8) : index + %expert_count = check.literal value(128) : index + %route_ids = check.generate.iota offset(0) step(1) period(9) : tensor<16x8xi32> + %actual_assignments = check.generate.fill value(-1) : tensor<128xi32> + %actual_counts = check.generate.fill value(-1) : tensor<4xi32> + %actual_descriptors = check.generate.fill value(-1) : tensor<4x128x4xi32> + %expected_assignments = check.generate.fill value(-1) : tensor<128xi32> + %expected_counts = check.generate.fill value(-1) : tensor<4xi32> + %expected_descriptors = check.generate.fill value(-1) : tensor<4x128x4xi32> + kernel.launch @qwen3_moe_build_batched_decode_expert_dispatch[%token_count, %route_count, %route_stride, %expert_count](%token_count, %route_count, %route_stride, %expert_count, %route_ids, %actual_assignments, %actual_counts, %actual_descriptors) : [index, index, index, index](index, index, index, index, tensor<16x8xi32>, tensor<128xi32>, tensor<4xi32>, tensor<4x128x4xi32>) + kernel.launch @qwen3_moe_build_batched_decode_expert_dispatch_reference[%token_count, %route_count, %route_stride, %expert_count](%token_count, %route_count, %route_stride, %expert_count, %route_ids, %expected_assignments, %expected_counts, %expected_descriptors) : [index, index, index, index](index, index, index, index, tensor<16x8xi32>, tensor<128xi32>, tensor<4xi32>, tensor<4x128x4xi32>) + check.expect.equal actual(%actual_assignments) expected(%expected_assignments) : tensor<128xi32> + check.expect.equal actual(%actual_counts) expected(%expected_counts) : tensor<4xi32> + check.expect.equal actual(%actual_descriptors) expected(%expected_descriptors) : tensor<4x128x4xi32> + check.return +} + +check.case public @qwen3_moe_batched_decode_expert_dispatch_coherent_case { + %token_count = check.literal value(16) : index + %route_count = check.literal value(8) : index + %route_stride = check.literal value(8) : index + %expert_count = check.literal value(128) : index + %route_ids = check.generate.iota offset(0) step(1) period(8) : tensor<16x8xi32> + %actual_assignments = check.generate.fill value(-1) : tensor<128xi32> + %actual_counts = check.generate.fill value(-1) : tensor<4xi32> + %actual_descriptors = check.generate.fill value(-1) : tensor<4x128x4xi32> + %expected_assignments = check.generate.fill value(-1) : tensor<128xi32> + %expected_counts = check.generate.fill value(-1) : tensor<4xi32> + %expected_descriptors = check.generate.fill value(-1) : tensor<4x128x4xi32> + kernel.launch @qwen3_moe_build_batched_decode_expert_dispatch[%token_count, %route_count, %route_stride, %expert_count](%token_count, %route_count, %route_stride, %expert_count, %route_ids, %actual_assignments, %actual_counts, %actual_descriptors) : [index, index, index, index](index, index, index, index, tensor<16x8xi32>, tensor<128xi32>, tensor<4xi32>, tensor<4x128x4xi32>) + kernel.launch @qwen3_moe_build_batched_decode_expert_dispatch_reference[%token_count, %route_count, %route_stride, %expert_count](%token_count, %route_count, %route_stride, %expert_count, %route_ids, %expected_assignments, %expected_counts, %expected_descriptors) : [index, index, index, index](index, index, index, index, tensor<16x8xi32>, tensor<128xi32>, tensor<4xi32>, tensor<4x128x4xi32>) + check.expect.equal actual(%actual_assignments) expected(%expected_assignments) : tensor<128xi32> + check.expect.equal actual(%actual_counts) expected(%expected_counts) : tensor<4xi32> + check.expect.equal actual(%actual_descriptors) expected(%expected_descriptors) : tensor<4x128x4xi32> + check.return +} + +// A second model geometry proves that expert, route, batch, descriptor, and +// schedule boundaries are configuration rather than Qwen3-30B constants. +check.case public @qwen3_moe_batched_decode_expert_dispatch_configurable_geometry_case { + %token_count = check.literal value(8) : index + %route_count = check.literal value(4) : index + %route_stride = check.literal value(4) : index + %expert_count = check.literal value(32) : index + %route_ids = check.generate.iota offset(0) step(1) period(13) : tensor<8x4xi32> + %actual_assignments = check.generate.fill value(-1) : tensor<32xi32> + %actual_counts = check.generate.fill value(-1) : tensor<4xi32> + %actual_descriptors = check.generate.fill value(-1) : tensor<4x32x4xi32> + %expected_assignments = check.generate.fill value(-1) : tensor<32xi32> + %expected_counts = check.generate.fill value(-1) : tensor<4xi32> + %expected_descriptors = check.generate.fill value(-1) : tensor<4x32x4xi32> + kernel.launch @qwen3_moe_build_batched_decode_expert_dispatch[%token_count, %route_count, %route_stride, %expert_count](%token_count, %route_count, %route_stride, %expert_count, %route_ids, %actual_assignments, %actual_counts, %actual_descriptors) : [index, index, index, index](index, index, index, index, tensor<8x4xi32>, tensor<32xi32>, tensor<4xi32>, tensor<4x32x4xi32>) + kernel.launch @qwen3_moe_build_batched_decode_expert_dispatch_reference[%token_count, %route_count, %route_stride, %expert_count](%token_count, %route_count, %route_stride, %expert_count, %route_ids, %expected_assignments, %expected_counts, %expected_descriptors) : [index, index, index, index](index, index, index, index, tensor<8x4xi32>, tensor<32xi32>, tensor<4xi32>, tensor<4x32x4xi32>) + check.expect.equal actual(%actual_assignments) expected(%expected_assignments) : tensor<32xi32> + check.expect.equal actual(%actual_counts) expected(%expected_counts) : tensor<4xi32> + check.expect.equal actual(%actual_descriptors) expected(%expected_descriptors) : tensor<4x32x4xi32> + check.return +} + +check.case public @qwen3_moe_batched_decode_expert_dispatch_diverse_benchmark_case { + %token_count = check.literal value(16) : index + %route_count = check.literal value(8) : index + %route_stride = check.literal value(8) : index + %expert_count = check.literal value(128) : index + %route_ids = check.generate.iota offset(0) step(1) period(128) : tensor<16x8xi32> + %assignment_ordinals = check.generate.fill value(-1) : tensor<128xi32> + %queue_counts = check.generate.fill value(-1) : tensor<4xi32> + %queue_descriptors = check.generate.fill value(-1) : tensor<4x128x4xi32> + kernel.launch @qwen3_moe_build_batched_decode_expert_dispatch[%token_count, %route_count, %route_stride, %expert_count](%token_count, %route_count, %route_stride, %expert_count, %route_ids, %assignment_ordinals, %queue_counts, %queue_descriptors) : [index, index, index, index](index, index, index, index, tensor<16x8xi32>, tensor<128xi32>, tensor<4xi32>, tensor<4x128x4xi32>) + check.return +} + +check.case public @qwen3_moe_batched_decode_expert_dispatch_coherent_benchmark_case { + %token_count = check.literal value(16) : index + %route_count = check.literal value(8) : index + %route_stride = check.literal value(8) : index + %expert_count = check.literal value(128) : index + %route_ids = check.generate.iota offset(0) step(1) period(8) : tensor<16x8xi32> + %assignment_ordinals = check.generate.fill value(-1) : tensor<128xi32> + %queue_counts = check.generate.fill value(-1) : tensor<4xi32> + %queue_descriptors = check.generate.fill value(-1) : tensor<4x128x4xi32> + kernel.launch @qwen3_moe_build_batched_decode_expert_dispatch[%token_count, %route_count, %route_stride, %expert_count](%token_count, %route_count, %route_stride, %expert_count, %route_ids, %assignment_ordinals, %queue_counts, %queue_descriptors) : [index, index, index, index](index, index, index, index, tensor<16x8xi32>, tensor<128xi32>, tensor<4xi32>, tensor<4x128x4xi32>) + check.return +} + +check.case public @qwen3_moe_batched_decode_expert_dispatch_configurable_benchmark_case { + %token_count = check.literal value(8) : index + %route_count = check.literal value(4) : index + %route_stride = check.literal value(4) : index + %expert_count = check.literal value(32) : index + %route_ids = check.generate.iota offset(0) step(1) period(13) : tensor<8x4xi32> + %assignment_ordinals = check.generate.fill value(-1) : tensor<32xi32> + %queue_counts = check.generate.fill value(-1) : tensor<4xi32> + %queue_descriptors = check.generate.fill value(-1) : tensor<4x32x4xi32> + kernel.launch @qwen3_moe_build_batched_decode_expert_dispatch[%token_count, %route_count, %route_stride, %expert_count](%token_count, %route_count, %route_stride, %expert_count, %route_ids, %assignment_ordinals, %queue_counts, %queue_descriptors) : [index, index, index, index](index, index, index, index, tensor<8x4xi32>, tensor<32xi32>, tensor<4xi32>, tensor<4x32x4xi32>) + check.return +} + +check.benchmark<@qwen3_moe_batched_decode_expert_dispatch_diverse_benchmark_case> @qwen3_moe_batched_decode_expert_dispatch_diverse + +check.benchmark<@qwen3_moe_batched_decode_expert_dispatch_coherent_benchmark_case> @qwen3_moe_batched_decode_expert_dispatch_coherent + +check.benchmark<@qwen3_moe_batched_decode_expert_dispatch_configurable_benchmark_case> @qwen3_moe_batched_decode_expert_dispatch_configurable diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/batched_decode_gate_up_q4k.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/batched_decode_gate_up_q4k.loom new file mode 100644 index 000000000000..97390fb6aaf8 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/batched_decode_gate_up_q4k.loom @@ -0,0 +1,427 @@ +// Expert-grouped batched-decode gate/up consumers for raw GGUF Q4_K weights. +// +// The assignment producer owns routing and publishes expert-contiguous +// assignment ordinals plus aligned descriptors. This file owns row-count +// schedules that consume those descriptors without inspecting route IDs. The +// first provider is exact M_e=2: each wave reuses one expert/channel weight +// row across two independently quantized activation rows and scatters the +// fused SwiGLU results back to their original [token][route] assignments. +func.def inline @qwen3_moe_q4k_chunk_pair_global(%weight: buffer, %row_byte_base: offset, %q4_block: index, %q4_group_pair: index, %q4_half: index, %header_words: vector<4xi32>) -> (vector<16xi8>, f32, f32, vector<16xi8>, f32, f32) { + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c4 = index.constant 4 : index + %c8 = index.constant 8 : index + %block_bytes = index.constant 144 : offset + %code_byte_add = index.constant 16 : offset + %c4_i32 = scalar.constant 4 : i32 + %nibble_mask = vector.constant 252645135 : vector<4xi32> + %bounded_pair = index.assume %q4_group_pair [range(%q4_group_pair, 0, 3)] : index + %bounded_half = index.assume %q4_half [range(%q4_half, 0, 1)] : index + %block_byte_add = index.scale %q4_block, %block_bytes : index, offset -> offset + %block_byte_base = index.add %row_byte_base, %block_byte_add : offset + %code_byte_base = index.add %block_byte_base, %code_byte_add : offset + %code_view = buffer.view %weight[%code_byte_base] : buffer -> view<32xi32> + %header_halves = vector.bitcast %header_words : vector<4xi32> to vector<8xf16> + %d_f16 = vector.extract %header_halves[0] : vector<8xf16> -> f16 + %dmin_f16 = vector.extract %header_halves[1] : vector<8xf16> -> f16 + %scale0 = vector.extract %header_words[1] : vector<4xi32> -> i32 + %scale1 = vector.extract %header_words[2] : vector<4xi32> -> i32 + %scale2 = vector.extract %header_words[3] : vector<4xi32> -> i32 + %d = scalar.extf %d_f16 : f16 to f32 + %dmin = scalar.extf %dmin_f16 : f16 to f32 + %pair_code_base = index.mul %bounded_pair, %c8 : index + %half_code_add = index.mul %bounded_half, %c4 : index + %code_index0 = index.add %pair_code_base, %half_code_add : index + %code_index = index.assume %code_index0 [range(%code_index0, 0, 28)] : index + %packed_codes = vector.load %code_view[%code_index] : view<32xi32> -> vector<4xi32> + %low_codes = vector.andi %packed_codes, %nibble_mask : vector<4xi32> + %c4_i32v = vector.splat %c4_i32 : vector<4xi32> + %high_shifted = vector.shrui %packed_codes, %c4_i32v : vector<4xi32> + %high_codes = vector.andi %high_shifted, %nibble_mask : vector<4xi32> + %q4_low = vector.bitcast %low_codes : vector<4xi32> to vector<16xi8> + %q4_high = vector.bitcast %high_codes : vector<4xi32> to vector<16xi8> + %low_group = index.mul %bounded_pair, %c2 : index + %high_group = index.add %low_group, %c1 : index + %low_scale, %low_minimum = func.call @qwen3_moe_q4k_scale_from_header(%scale0, %scale1, %scale2, %low_group) : (i32, i32, i32, index) -> (i32, i32) + %high_scale, %high_minimum = func.call @qwen3_moe_q4k_scale_from_header(%scale0, %scale1, %scale2, %high_group) : (i32, i32, i32, index) -> (i32, i32) + %low_scale_f32 = scalar.uitofp %low_scale : i32 to f32 + %low_minimum_f32 = scalar.uitofp %low_minimum : i32 to f32 + %high_scale_f32 = scalar.uitofp %high_scale : i32 to f32 + %high_minimum_f32 = scalar.uitofp %high_minimum : i32 to f32 + %low_d_scale = scalar.mulf %d, %low_scale_f32 : f32 + %low_dmin_scale = scalar.mulf %dmin, %low_minimum_f32 : f32 + %high_d_scale = scalar.mulf %d, %high_scale_f32 : f32 + %high_dmin_scale = scalar.mulf %dmin, %high_minimum_f32 : f32 + func.return %q4_low, %low_d_scale, %low_dmin_scale, %q4_high, %high_d_scale, %high_dmin_scale : vector<16xi8>, f32, f32, vector<16xi8>, f32, f32 +} + +func.def inline @qwen3_moe_q4k_q8_1_dot(%q4_values: vector<16xi8>, %d_scale: f32, %dmin_scale: f32, %q8_values: vector<16xi8>, %q8_d: f32, %q8_s: f32) -> (f32) { + %c0_i32 = scalar.constant 0 : i32 + %c0_i32v = vector.constant 0 : vector<4xi32> + %half_f32 = scalar.constant 0.5 : f32 + %partial_dots = vector.dot4i %q4_values, %q8_values, %c0_i32v : vector<16xi8>, vector<16xi8>, vector<4xi32> + %q_sum = vector.reduce %partial_dots, %c0_i32 : vector<4xi32>, i32 + %q_sum_f32 = scalar.sitofp %q_sum : i32 to f32 + %scaled_dot0 = scalar.mulf %q8_d, %d_scale : f32 + %scaled_dot = scalar.mulf %scaled_dot0, %q_sum_f32 : f32 + %q8_half_sum = scalar.mulf %q8_s, %half_f32 : f32 + %minimum_correction = scalar.mulf %dmin_scale, %q8_half_sum : f32 + %contribution = scalar.subf %scaled_dot, %minimum_correction : f32 + func.return %contribution : f32 +} + +func.def inline @qwen3_moe_q4k_scale_from_header(%scale0: i32, %scale1: i32, %scale2: i32, %q4_group: index) -> (i32, i32) { + %c4 = index.constant 4 : index + %c8 = index.constant 8 : index + %c2_i32 = scalar.constant 2 : i32 + %c4_i32 = scalar.constant 4 : i32 + %c15_i32 = scalar.constant 15 : i32 + %c48_i32 = scalar.constant 48 : i32 + %bounded_group = index.assume %q4_group [range(%q4_group, 0, 7)] : index + %is_low_group = index.cmp ult, %bounded_group, %c4 : index + %scale_lane = index.rem %bounded_group, %c4 : index + %scale_shift_index = index.mul %scale_lane, %c8 : index + %scale_shift = index.cast %scale_shift_index : index to i32 + %high_shift = scalar.addi %scale_shift, %c2_i32 : i32 + %minimum_shift = scalar.addi %scale_shift, %c4_i32 : i32 + %selected_scale_source = scf.select %is_low_group, %scale0, %scale2 : i32 + %selected_minimum_source = scf.select %is_low_group, %scale1, %scale2 : i32 + %selected_scale_high_shift = scf.select %is_low_group, %scale_shift, %high_shift : i32 + %selected_minimum_low_shift = scf.select %is_low_group, %scale_shift, %minimum_shift : i32 + %scale_low0 = scalar.shrui %selected_scale_source, %scale_shift : i32 + %scale_low = scalar.andi %scale_low0, %c15_i32 : i32 + %scale_high0 = scalar.shrui %scale0, %selected_scale_high_shift : i32 + %scale_high = scalar.andi %scale_high0, %c48_i32 : i32 + %scale = scalar.ori %scale_low, %scale_high : i32 + %minimum_low0 = scalar.shrui %selected_minimum_source, %selected_minimum_low_shift : i32 + %minimum_low = scalar.andi %minimum_low0, %c15_i32 : i32 + %minimum_high0 = scalar.shrui %scale1, %selected_scale_high_shift : i32 + %minimum_high = scalar.andi %minimum_high0, %c48_i32 : i32 + %minimum = scalar.ori %minimum_low, %minimum_high : i32 + func.return %scale, %minimum : i32, i32 +} + +func.def inline @qwen3_moe_unpack_batched_decode_expert_descriptor(%descriptor: vector<4xi32>) -> (index, index, index) { + %configured_expert_count = config.get @qwen3_moe.router.expert_count : index + %configured_route_count = config.get @qwen3_moe.router.route_count : index + %configured_token_capacity = config.get @qwen3_moe.workload.token_capacity : index + %assignment_capacity = index.mul %configured_token_capacity, %configured_route_count : index + %expert_i32 = vector.extract %descriptor[0] : vector<4xi32> -> i32 + %assignment_base_i32 = vector.extract %descriptor[1] : vector<4xi32> -> i32 + %row_count_i32 = vector.extract %descriptor[2] : vector<4xi32> -> i32 + %expert0 = index.cast %expert_i32 : i32 to index + %expert, %descriptor_expert_count = index.assume %expert0, %configured_expert_count [range(%expert0, 0, 511), lt(%expert0, %configured_expert_count)] : index, index + %assignment_base0 = index.cast %assignment_base_i32 : i32 to index + %assignment_base, %descriptor_assignment_capacity = index.assume %assignment_base0, %assignment_capacity [range(%assignment_base0, 0, 65535), lt(%assignment_base0, %assignment_capacity)] : index, index + %row_count0 = index.cast %row_count_i32 : i32 to index + %row_count, %descriptor_token_capacity = index.assume %row_count0, %configured_token_capacity [range(%row_count0, 1, 2048), le(%row_count0, %configured_token_capacity)] : index, index + func.return %expert, %assignment_base, %row_count : index, index, index +} + +config.decl @qwen3_moe.router.expert_count : %value: index where [range(%value, 32, 512), mul(%value, 32)] + +config.decl @qwen3_moe.router.route_count : %value: index where [range(%value, 1, 32)] + +config.decl @qwen3_moe.routed_gate_up.input_size : %value: index where [range(%value, 512, 32768), mul(%value, 512)] + +config.decl @qwen3_moe.routed_gate_up.expert_count : %value: index where [range(%value, 1, 512)] + +config.decl @qwen3_moe.routed_gate_up.route_count : %value: index where [range(%value, 1, 8)] + +config.decl @qwen3_moe.routed_gate_up.output_size : %value: index where [range(%value, 1, 4096)] + +config.decl @qwen3_moe.workload.token_capacity : %value: index where [range(%value, 1, 2048)] + +// Static fallback capacity for the exact-two-row queue. Device-side indirect +// launch replaces this upper bound with the producer's exact queue count. +config.decl @qwen3_moe.batched_decode.rows2_descriptor_capacity : %value: index where [range(%value, 1, 512)] + +kernel.decl @ggml_quantize_q8_1_x4_f32(%token_count: index, %input_size: index) launch(%token_count: index, %input_size: index, %input: buffer, %output: buffer) + +func.decl @qwen3_moe_q4k_chunk_pair_global(%weight: buffer, %row_byte_base: offset, %q4_block: index, %q4_group_pair: index, %q4_half: index, %header_words: vector<4xi32>) -> (vector<16xi8>, f32, f32, vector<16xi8>, f32, f32) + +func.decl @qwen3_moe_q4k_q8_1_dot(%q4_values: vector<16xi8>, %d_scale: f32, %dmin_scale: f32, %q8_values: vector<16xi8>, %q8_d: f32, %q8_s: f32) -> (f32) + +func.decl @qwen3_moe_unpack_batched_decode_expert_descriptor(%descriptor: vector<4xi32>) -> (index, index, index) + +kernel.decl @qwen3_moe_build_batched_decode_expert_dispatch(%token_count: index, %route_count: index, %route_stride: index, %expert_count: index) launch(%token_count: index, %route_count: index, %route_stride: index, %expert_count: index, %route_ids: buffer, %assignment_ordinals: buffer, %queue_counts: buffer, %queue_descriptors: buffer) + +kernel.decl @qwen3_moe_routed_gate_up_swiglu_q4k_q8(%token_count: index, %route_count: index, %route_stride: index, %expert_count: index, %output_size: index) launch(%token_count: index, %route_count: index, %route_stride: index, %expert_count: index, %output_size: index, %q8_input: buffer, %route_ids: buffer, %gate_weight: buffer, %up_weight: buffer, %output: buffer) + +// Contracts one expert/channel Q4_K row against two independently quantized +// activation rows. Splitting gate and up into separate contractions keeps each +// recurrence small while preserving the important invariant: every weight +// byte loaded by this helper contributes to both rows. +func.def inline @qwen3_moe_batched_decode_q4k_rows2_lane(%input_size: index, %weight: buffer, %q8_input: buffer, %weight_row_byte_base: offset, %q8_row0_byte_base: offset, %q8_row1_byte_base: offset, %lane: index) -> (f32, f32) { + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c4 = index.constant 4 : index + %c8 = index.constant 8 : index + %c128 = index.constant 128 : index + %c256 = index.constant 256 : index + %c1023 = index.constant 1023 : index + %c1024 = index.constant 1024 : index + %q4_block_bytes = index.constant 144 : offset + %q8_group_bytes = index.constant 144 : offset + %q8_payload_byte_add = index.constant 16 : offset + %c0_f32 = scalar.constant 0.0 : f32 + %q4_block_count = index.div %input_size, %c256 : index + %q8_group_count = index.div %input_size, %c128 : index + %padded_input_size = index.add %input_size, %c1023 : index + %iteration_count = index.div %padded_input_size, %c1024 : index + %q4_group_pair0 = index.div %lane, %c2 : index + %q4_group_pair1 = index.rem %q4_group_pair0, %c4 : index + %q4_group_pair = index.assume %q4_group_pair1 [range(%q4_group_pair1, 0, 3)] : index + %q4_half0 = index.rem %lane, %c2 : index + %q4_half = index.assume %q4_half0 [range(%q4_half0, 0, 1)] : index + %lane_q4_block = index.div %lane, %c8 : index + %q8_group_in_block0 = index.div %q4_group_pair, %c2 : index + %q8_group_in_block = index.assume %q8_group_in_block0 [range(%q8_group_in_block0, 0, 1)] : index + %pair_in_q8_group0 = index.rem %q4_group_pair, %c2 : index + %pair_in_q8_group = index.assume %pair_in_q8_group0 [range(%pair_in_q8_group0, 0, 1)] : index + %q8_low_inner_block0 = index.mul %pair_in_q8_group, %c2 : index + %q8_low_inner_block = index.assume %q8_low_inner_block0 [range(%q8_low_inner_block0, 0, 2)] : index + %q8_high_inner_block0 = index.add %q8_low_inner_block, %c1 : index + %q8_high_inner_block = index.assume %q8_high_inner_block0 [range(%q8_high_inner_block0, 1, 3)] : index + %q8_half_word_add = index.mul %q4_half, %c4 : index + %q8_low_inner_word_base = index.mul %q8_low_inner_block, %c8 : index + %q8_low_word_index0 = index.add %q8_low_inner_word_base, %q8_half_word_add : index + %q8_low_word_index = index.assume %q8_low_word_index0 [range(%q8_low_word_index0, 0, 20)] : index + %q8_high_inner_word_base = index.mul %q8_high_inner_block, %c8 : index + %q8_high_word_index0 = index.add %q8_high_inner_word_base, %q8_half_word_add : index + %q8_high_word_index = index.assume %q8_high_word_index0 [range(%q8_high_word_index0, 8, 28)] : index + %q8_low_ds_index0 = index.mul %q8_low_inner_block, %c2 : index + %q8_low_ds_index = index.assume %q8_low_ds_index0 [range(%q8_low_ds_index0, 0, 4)] : index + + %row0_acc, %row1_acc = scf.for %iteration = [%c0 to %iteration_count step %c1](%row0_iter = %c0_f32 : f32, %row1_iter = %c0_f32 : f32) -> (f32, f32) unroll { + %iteration_q4_block = index.mul %iteration, %c4 : index + %q4_block0 = index.add %iteration_q4_block, %lane_q4_block : index + %valid_q4_block = index.cmp ult, %q4_block0, %q4_block_count : index + %row0_contribution, %row1_contribution = scf.if %valid_q4_block -> (f32, f32) { + %q4_block, %bounded_q4_block_count = index.assume %q4_block0, %q4_block_count [lt(%q4_block0, %q4_block_count)] : index, index + %q8_block_group_base = index.mul %q4_block, %c2 : index + %q8_group0 = index.add %q8_block_group_base, %q8_group_in_block : index + %q8_group, %bounded_q8_group_count = index.assume %q8_group0, %q8_group_count [lt(%q8_group0, %q8_group_count)] : index, index + %q8_group_byte_add = index.scale %q8_group, %q8_group_bytes : index, offset -> offset + %q8_group0_byte_base = index.add %q8_row0_byte_base, %q8_group_byte_add : offset + %q8_group1_byte_base = index.add %q8_row1_byte_base, %q8_group_byte_add : offset + %q8_payload0_byte_base = index.add %q8_group0_byte_base, %q8_payload_byte_add : offset + %q8_payload1_byte_base = index.add %q8_group1_byte_base, %q8_payload_byte_add : offset + %q8_ds0_view = buffer.view %q8_input[%q8_group0_byte_base] : buffer -> view<8xf16> + %q8_ds1_view = buffer.view %q8_input[%q8_group1_byte_base] : buffer -> view<8xf16> + %q8_words0_view = buffer.view %q8_input[%q8_payload0_byte_base] : buffer -> view<32xi32> + %q8_words1_view = buffer.view %q8_input[%q8_payload1_byte_base] : buffer -> view<32xi32> + %q8_ds0 = vector.load %q8_ds0_view[%q8_low_ds_index] : view<8xf16> -> vector<4xf16> + %q8_ds1 = vector.load %q8_ds1_view[%q8_low_ds_index] : view<8xf16> -> vector<4xf16> + %q8_low_d0_f16 = vector.extract %q8_ds0[0] : vector<4xf16> -> f16 + %q8_low_s0_f16 = vector.extract %q8_ds0[1] : vector<4xf16> -> f16 + %q8_high_d0_f16 = vector.extract %q8_ds0[2] : vector<4xf16> -> f16 + %q8_high_s0_f16 = vector.extract %q8_ds0[3] : vector<4xf16> -> f16 + %q8_low_d1_f16 = vector.extract %q8_ds1[0] : vector<4xf16> -> f16 + %q8_low_s1_f16 = vector.extract %q8_ds1[1] : vector<4xf16> -> f16 + %q8_high_d1_f16 = vector.extract %q8_ds1[2] : vector<4xf16> -> f16 + %q8_high_s1_f16 = vector.extract %q8_ds1[3] : vector<4xf16> -> f16 + %q8_low_d0 = scalar.extf %q8_low_d0_f16 : f16 to f32 + %q8_low_s0 = scalar.extf %q8_low_s0_f16 : f16 to f32 + %q8_high_d0 = scalar.extf %q8_high_d0_f16 : f16 to f32 + %q8_high_s0 = scalar.extf %q8_high_s0_f16 : f16 to f32 + %q8_low_d1 = scalar.extf %q8_low_d1_f16 : f16 to f32 + %q8_low_s1 = scalar.extf %q8_low_s1_f16 : f16 to f32 + %q8_high_d1 = scalar.extf %q8_high_d1_f16 : f16 to f32 + %q8_high_s1 = scalar.extf %q8_high_s1_f16 : f16 to f32 + %q8_low_words0 = vector.load %q8_words0_view[%q8_low_word_index] : view<32xi32> -> vector<4xi32> + %q8_high_words0 = vector.load %q8_words0_view[%q8_high_word_index] : view<32xi32> -> vector<4xi32> + %q8_low_words1 = vector.load %q8_words1_view[%q8_low_word_index] : view<32xi32> -> vector<4xi32> + %q8_high_words1 = vector.load %q8_words1_view[%q8_high_word_index] : view<32xi32> -> vector<4xi32> + %q8_low_values0 = vector.bitcast %q8_low_words0 : vector<4xi32> to vector<16xi8> + %q8_high_values0 = vector.bitcast %q8_high_words0 : vector<4xi32> to vector<16xi8> + %q8_low_values1 = vector.bitcast %q8_low_words1 : vector<4xi32> to vector<16xi8> + %q8_high_values1 = vector.bitcast %q8_high_words1 : vector<4xi32> to vector<16xi8> + %q4_block_byte_add = index.scale %q4_block, %q4_block_bytes : index, offset -> offset + %q4_block_byte_base = index.add %weight_row_byte_base, %q4_block_byte_add : offset + %header_view = buffer.view %weight[%q4_block_byte_base] : buffer -> view<4xi32> + %header_words = vector.load %header_view[0] : view<4xi32> -> vector<4xi32> + %q4_low, %low_d_scale, %low_dmin_scale, %q4_high, %high_d_scale, %high_dmin_scale = func.call @qwen3_moe_q4k_chunk_pair_global(%weight, %weight_row_byte_base, %q4_block, %q4_group_pair, %q4_half, %header_words) : (buffer, offset, index, index, index, vector<4xi32>) -> (vector<16xi8>, f32, f32, vector<16xi8>, f32, f32) + %row0_low = func.call @qwen3_moe_q4k_q8_1_dot(%q4_low, %low_d_scale, %low_dmin_scale, %q8_low_values0, %q8_low_d0, %q8_low_s0) : (vector<16xi8>, f32, f32, vector<16xi8>, f32, f32) -> (f32) + %row0_high = func.call @qwen3_moe_q4k_q8_1_dot(%q4_high, %high_d_scale, %high_dmin_scale, %q8_high_values0, %q8_high_d0, %q8_high_s0) : (vector<16xi8>, f32, f32, vector<16xi8>, f32, f32) -> (f32) + %row1_low = func.call @qwen3_moe_q4k_q8_1_dot(%q4_low, %low_d_scale, %low_dmin_scale, %q8_low_values1, %q8_low_d1, %q8_low_s1) : (vector<16xi8>, f32, f32, vector<16xi8>, f32, f32) -> (f32) + %row1_high = func.call @qwen3_moe_q4k_q8_1_dot(%q4_high, %high_d_scale, %high_dmin_scale, %q8_high_values1, %q8_high_d1, %q8_high_s1) : (vector<16xi8>, f32, f32, vector<16xi8>, f32, f32) -> (f32) + %row0_pair = scalar.addf %row0_low, %row0_high : f32 + %row1_pair = scalar.addf %row1_low, %row1_high : f32 + scf.yield %row0_pair, %row1_pair : f32, f32 + } else { + scf.yield %c0_f32, %c0_f32 : f32, f32 + } + %row0_next = scalar.addf %row0_iter, %row0_contribution : f32 + %row1_next = scalar.addf %row1_iter, %row1_contribution : f32 + scf.yield %row0_next, %row1_next : f32, f32 + } + func.return %row0_acc, %row1_acc : f32, f32 +} + +amdgpu.target @qwen3_moe_batched_decode_gate_up_gfx11_wave32 {subgroup_size = 32} + +// Processes the exact-two-row descriptor queue. The static launch covers a +// config-derived safe descriptor capacity and guards inactive workgroups. A +// composed device launcher instead supplies the exact descriptor count as +// both launch geometry and an ABI fact. +kernel.def target(@qwen3_moe_batched_decode_gate_up_gfx11_wave32) @qwen3_moe_batched_decode_gate_up_q4k_rows2(%descriptor_count: index, %queue_ordinal: index, %token_count: index, %route_count: index, %expert_count: index, %output_size: index) { + %configured_output_size = config.get @qwen3_moe.routed_gate_up.output_size : index + %descriptor_capacity = config.get @qwen3_moe.batched_decode.rows2_descriptor_capacity : index + %c1 = index.constant 1 : index + %c3 = index.constant 3 : index + %c4 = index.constant 4 : index + %workgroup_size = index.constant 128 : index + %padded_output_size = index.add %configured_output_size, %c3 : index + %channel_workgroup_count = index.div %padded_output_size, %c4 : index + kernel.launch.config workgroups(%channel_workgroup_count, %descriptor_capacity, %c1) workgroup_size(%workgroup_size, %c1, %c1) : index +} launch(%descriptor_count: index, %queue_ordinal: index, %token_count: index, %route_count: index, %expert_count: index, %output_size: index, %queue_descriptors: buffer, %assignment_ordinals: buffer, %q8_input: buffer, %gate_weight: buffer, %up_weight: buffer, %output: buffer) { + %configured_token_capacity0 = config.get @qwen3_moe.workload.token_capacity : index + %bounded_token_count, %configured_token_capacity = index.assume %token_count, %configured_token_capacity0 [range(%token_count, 1, 2048), eq(%token_count, %configured_token_capacity0)] : index, index + %configured_route_count0 = config.get @qwen3_moe.router.route_count : index + %bounded_route_count, %configured_route_count = index.assume %route_count, %configured_route_count0 [range(%route_count, 1, 32), eq(%route_count, %configured_route_count0)] : index, index + %configured_expert_count0 = config.get @qwen3_moe.router.expert_count : index + %bounded_expert_count, %configured_expert_count = index.assume %expert_count, %configured_expert_count0 [range(%expert_count, 32, 512), eq(%expert_count, %configured_expert_count0)] : index, index + %configured_weight_expert_count0 = config.get @qwen3_moe.routed_gate_up.expert_count : index + %configured_weight_expert_count, %model_expert_count = index.assume %configured_weight_expert_count0, %configured_expert_count [eq(%configured_weight_expert_count0, %configured_expert_count)] : index, index + %configured_output_size0 = config.get @qwen3_moe.routed_gate_up.output_size : index + %bounded_output_size, %configured_output_size = index.assume %output_size, %configured_output_size0 [range(%output_size, 1, 4096), eq(%output_size, %configured_output_size0)] : index, index + %bounded_queue_ordinal = index.assume %queue_ordinal [range(%queue_ordinal, 0, 3)] : index + %input_size = config.get @qwen3_moe.routed_gate_up.input_size : index + %descriptor_ordinal0 = kernel.workgroup.id : index + %channel_workgroup = kernel.workgroup.id : index + %subgroup = kernel.subgroup.id : index + %lane0 = kernel.subgroup.lane.id : index + %lane = index.assume %lane0 [range(%lane0, 0, 31)] : index + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c4 = index.constant 4 : index + %c128 = index.constant 128 : index + %c256 = index.constant 256 : index + %c1_byte = index.constant 1 : offset + %q4_block_bytes = index.constant 144 : index + %c0_i32 = scalar.constant 0 : i32 + %c0_offset = index.constant 0 : offset + %assignment_count = index.mul %bounded_token_count, %configured_route_count : index + %assignment_capacity = index.mul %configured_token_capacity, %configured_route_count : index + %assignment_limited_descriptor_capacity0 = index.div %assignment_capacity, %c2 : index + %descriptor_capacity0 = index.min %configured_expert_count, %assignment_limited_descriptor_capacity0 : index + %configured_descriptor_capacity0 = config.get @qwen3_moe.batched_decode.rows2_descriptor_capacity : index + %descriptor_capacity, %configured_descriptor_capacity = index.assume %descriptor_capacity0, %configured_descriptor_capacity0 [range(%descriptor_capacity0, 1, 512), eq(%descriptor_capacity0, %configured_descriptor_capacity0)] : index, index + %bounded_descriptor_count, %queue_descriptor_capacity = index.assume %descriptor_count, %configured_descriptor_capacity [range(%descriptor_count, 1, 512), le(%descriptor_count, %configured_descriptor_capacity)] : index, index + %valid_descriptor = index.cmp ult, %descriptor_ordinal0, %bounded_descriptor_count : index + %safe_descriptor_ordinal0 = scf.select %valid_descriptor, %descriptor_ordinal0, %c0 : index + %descriptor_ordinal, %launch_descriptor_count = index.assume %safe_descriptor_ordinal0, %bounded_descriptor_count [lt(%safe_descriptor_ordinal0, %bounded_descriptor_count)] : index, index + %q4_block_count = index.div %input_size, %c256 : index + %weight_row_bytes = index.mul %q4_block_count, %q4_block_bytes : index + %weight_expert_bytes = index.mul %bounded_output_size, %weight_row_bytes : index + %q8_group_count = index.div %input_size, %c128 : index + %q8_bytes_per_token = index.mul %q8_group_count, %q4_block_bytes : index + %queue_descriptors_noalias, %assignment_ordinals_noalias, %q8_noalias, %gate_noalias, %up_noalias, %output_noalias = buffer.assume.noalias %queue_descriptors, %assignment_ordinals, %q8_input, %gate_weight, %up_weight, %output : buffer, buffer, buffer, buffer, buffer, buffer + %queue_descriptor_view = buffer.view %queue_descriptors_noalias[%c0_offset] : buffer -> view<4x[%configured_expert_count]x4xi32> + %assignment_view = buffer.view %assignment_ordinals_noalias[%c0_offset] : buffer -> view<[%assignment_count]xi32> + %output_view = buffer.view %output_noalias[%c0_offset] : buffer -> view<[%assignment_count]x[%bounded_output_size]xf32> + %descriptor = vector.load %queue_descriptor_view[%bounded_queue_ordinal, %descriptor_ordinal, 0] : view<4x[%configured_expert_count]x4xi32> -> vector<4xi32> + %expert, %assignment_base0, %row_count0 = func.call @qwen3_moe_unpack_batched_decode_expert_descriptor(%descriptor) : (vector<4xi32>) -> (index, index, index) + %assignment_base, %consumer_assignment_count = index.assume %assignment_base0, %assignment_count [lt(%assignment_base0, %assignment_count)] : index, index + %row_count = index.assume %row_count0 [range(%row_count0, 2, 2)] : index + %assignment1_ordinal0 = index.add %assignment_base, %c1 : index + %assignment1_ordinal, %pair_assignment_count = index.assume %assignment1_ordinal0, %assignment_count [lt(%assignment1_ordinal0, %assignment_count)] : index, index + %assignment0_i32 = view.load %assignment_view[%assignment_base] : view<[%assignment_count]xi32> -> i32 + %assignment1_i32 = view.load %assignment_view[%assignment1_ordinal] : view<[%assignment_count]xi32> -> i32 + %assignment0_index0 = index.cast %assignment0_i32 : i32 to index + %assignment1_index0 = index.cast %assignment1_i32 : i32 to index + %assignment0, %assignment1, %routed_assignment_count = index.assume %assignment0_index0, %assignment1_index0, %assignment_count [range(%assignment0_index0, 0, 65535), range(%assignment1_index0, 0, 65535), lt(%assignment0_index0, %assignment_count), lt(%assignment1_index0, %assignment_count)] : index, index, index + %token0_index0 = index.div %assignment0, %configured_route_count : index + %token1_index0 = index.div %assignment1, %configured_route_count : index + %token0, %token1, %q8_token_count = index.assume %token0_index0, %token1_index0, %bounded_token_count [lt(%token0_index0, %bounded_token_count), lt(%token1_index0, %bounded_token_count)] : index, index, index + %channel_base = index.mul %channel_workgroup, %c4 : index + %channel0 = index.add %channel_base, %subgroup : index + %valid_channel = index.cmp ult, %channel0, %bounded_output_size : index + %safe_channel0 = scf.select %valid_channel, %channel0, %c0 : index + %channel, %weight_output_size = index.assume %safe_channel0, %bounded_output_size [lt(%safe_channel0, %bounded_output_size)] : index, index + %expert_byte_base = index.mul %expert, %weight_expert_bytes : index + %channel_byte_add = index.mul %channel, %weight_row_bytes : index + %row_byte_index = index.add %expert_byte_base, %channel_byte_add : index + %row_byte_base = index.scale %row_byte_index, %c1_byte : index, offset -> offset + %q8_token0_byte_index = index.mul %token0, %q8_bytes_per_token : index + %q8_token1_byte_index = index.mul %token1, %q8_bytes_per_token : index + %q8_token0_byte_base = index.scale %q8_token0_byte_index, %c1_byte : index, offset -> offset + %q8_token1_byte_base = index.scale %q8_token1_byte_index, %c1_byte : index, offset -> offset + %gate0_acc, %gate1_acc = func.call @qwen3_moe_batched_decode_q4k_rows2_lane(%input_size, %gate_noalias, %q8_noalias, %row_byte_base, %q8_token0_byte_base, %q8_token1_byte_base, %lane) : (index, buffer, buffer, offset, offset, offset, index) -> (f32, f32) + %up0_acc, %up1_acc = func.call @qwen3_moe_batched_decode_q4k_rows2_lane(%input_size, %up_noalias, %q8_noalias, %row_byte_base, %q8_token0_byte_base, %q8_token1_byte_base, %lane) : (index, buffer, buffer, offset, offset, offset, index) -> (f32, f32) + %gate0_dot = kernel.subgroup.reduce %gate0_acc : f32 + %up0_dot = kernel.subgroup.reduce %up0_acc : f32 + %gate1_dot = kernel.subgroup.reduce %gate1_acc : f32 + %up1_dot = kernel.subgroup.reduce %up1_acc : f32 + %lane_i32 = index.cast %lane : index to i32 + %is_lane_zero = scalar.cmpi eq, %lane_i32, %c0_i32 : i32 + %valid_channel_descriptor = scalar.andi %valid_channel, %valid_descriptor : i1 + %writes_output = scalar.andi %valid_channel_descriptor, %is_lane_zero : i1 + scf.if %writes_output { + %gate0_silu = scalar.siluf %gate0_dot : f32 + %gate1_silu = scalar.siluf %gate1_dot : f32 + %result0 = scalar.mulf %gate0_silu, %up0_dot : f32 + %result1 = scalar.mulf %gate1_silu, %up1_dot : f32 + view.store %result0, %output_view[%assignment0, %channel] : f32, view<[%assignment_count]x[%bounded_output_size]xf32> + view.store %result1, %output_view[%assignment1, %channel] : f32, view<[%assignment_count]x[%bounded_output_size]xf32> + } + kernel.return +} + +// Two identical top-8 route rows produce eight exact-two-row descriptors. +// The direct route-centric provider remains the numerical oracle. +check.case public @qwen3_moe_batched_decode_gate_up_q4k_rows2_differential_case { + %descriptor_count = check.literal value(8) : index + %queue_ordinal = check.literal value(1) : index + %token_count = check.literal value(2) : index + %input_size = check.literal value(512) : index + %route_count = check.literal value(8) : index + %route_stride = check.literal value(8) : index + %expert_count = check.literal value(32) : index + %output_size = check.literal value(32) : index + %input = check.generate.iota offset(-0.5) step(0.001953125) period(1024) : tensor<2x512xf32> + %q8_input = check.generate.fill value(0) : tensor<2x576xi8> + %route_ids = check.generate.iota offset(0) step(1) period(8) : tensor<2x8xi32> + %assignment_ordinals = check.generate.fill value(-1) : tensor<16xi32> + %queue_counts = check.generate.fill value(-1) : tensor<4xi32> + %queue_descriptors = check.generate.fill value(-1) : tensor<4x32x4xi32> + %gate_weight = check.generate.iota offset(-72) step(1) period(144) : tensor<32x32x2x144xi8> + %up_weight = check.generate.iota offset(-71) step(1) period(144) : tensor<32x32x2x144xi8> + %expected = check.generate.fill value(0.0) : tensor<2x8x32xf32> + %actual = check.generate.fill value(1.0) : tensor<2x8x32xf32> + kernel.launch @ggml_quantize_q8_1_x4_f32[%token_count, %input_size](%token_count, %input_size, %input, %q8_input) : [index, index](index, index, tensor<2x512xf32>, tensor<2x576xi8>) + kernel.launch @qwen3_moe_build_batched_decode_expert_dispatch[%token_count, %route_count, %route_stride, %expert_count](%token_count, %route_count, %route_stride, %expert_count, %route_ids, %assignment_ordinals, %queue_counts, %queue_descriptors) : [index, index, index, index](index, index, index, index, tensor<2x8xi32>, tensor<16xi32>, tensor<4xi32>, tensor<4x32x4xi32>) + kernel.launch @qwen3_moe_routed_gate_up_swiglu_q4k_q8[%token_count, %route_count, %route_stride, %expert_count, %output_size](%token_count, %route_count, %route_stride, %expert_count, %output_size, %q8_input, %route_ids, %gate_weight, %up_weight, %expected) : [index, index, index, index, index](index, index, index, index, index, tensor<2x576xi8>, tensor<2x8xi32>, tensor<32x32x2x144xi8>, tensor<32x32x2x144xi8>, tensor<2x8x32xf32>) + kernel.launch @qwen3_moe_batched_decode_gate_up_q4k_rows2[%descriptor_count, %queue_ordinal, %token_count, %route_count, %expert_count, %output_size](%descriptor_count, %queue_ordinal, %token_count, %route_count, %expert_count, %output_size, %queue_descriptors, %assignment_ordinals, %q8_input, %gate_weight, %up_weight, %actual) : [index, index, index, index, index, index](index, index, index, index, index, index, tensor<4x32x4xi32>, tensor<16xi32>, tensor<2x576xi8>, tensor<32x32x2x144xi8>, tensor<32x32x2x144xi8>, tensor<2x8x32xf32>) + check.expect.close actual(%actual) expected(%expected) atol(0.25) rtol(9.9999999999999995e-07) nan(same) : tensor<2x8x32xf32> + check.return +} + +// Production dimensions with a uniform two-row-per-expert distribution. +// Timing includes the assignment producer so the row schedule cannot hide its +// routing preparation cost. +check.case public @qwen3_moe_batched_decode_gate_up_q4k_rows2_benchmark_case { + %descriptor_count = check.literal value(64) : index + %queue_ordinal = check.literal value(1) : index + %token_count = check.literal value(16) : index + %route_count = check.literal value(8) : index + %route_stride = check.literal value(8) : index + %expert_count = check.literal value(128) : index + %output_size = check.literal value(768) : index + %q8_input = check.generate.fill value(0) : tensor<16x2304xi8> + %route_ids = check.generate.iota offset(0) step(1) period(64) : tensor<16x8xi32> + %assignment_ordinals = check.generate.fill value(-1) : tensor<128xi32> + %queue_counts = check.generate.fill value(-1) : tensor<4xi32> + %queue_descriptors = check.generate.fill value(-1) : tensor<4x128x4xi32> + %gate_weight = check.generate.fill value(0) : tensor<128x768x8x144xi8> + %up_weight = check.generate.fill value(0) : tensor<128x768x8x144xi8> + %output = check.generate.fill value(1.0) : tensor<16x8x768xf32> + kernel.launch @qwen3_moe_build_batched_decode_expert_dispatch[%token_count, %route_count, %route_stride, %expert_count](%token_count, %route_count, %route_stride, %expert_count, %route_ids, %assignment_ordinals, %queue_counts, %queue_descriptors) : [index, index, index, index](index, index, index, index, tensor<16x8xi32>, tensor<128xi32>, tensor<4xi32>, tensor<4x128x4xi32>) + kernel.launch @qwen3_moe_batched_decode_gate_up_q4k_rows2[%descriptor_count, %queue_ordinal, %token_count, %route_count, %expert_count, %output_size](%descriptor_count, %queue_ordinal, %token_count, %route_count, %expert_count, %output_size, %queue_descriptors, %assignment_ordinals, %q8_input, %gate_weight, %up_weight, %output) : [index, index, index, index, index, index](index, index, index, index, index, index, tensor<4x128x4xi32>, tensor<128xi32>, tensor<16x2304xi8>, tensor<128x768x8x144xi8>, tensor<128x768x8x144xi8>, tensor<16x8x768xf32>) + check.return +} + +check.benchmark<@qwen3_moe_batched_decode_gate_up_q4k_rows2_benchmark_case> @qwen3_moe_batched_decode_gate_up_q4k_rows2_benchmark diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/dense_linear_quantized_f16_wmma.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/dense_linear_quantized_f16_wmma.loom new file mode 100644 index 000000000000..69b4f701f0c0 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/dense_linear_quantized_f16_wmma.loom @@ -0,0 +1,1349 @@ +// Dense raw-quantized projection for the Qwen attention path. +// +// Each two-wave workgroup computes 64 output channels for 32 contiguous +// tokens. Raw GGUF Q4_K or Q6_K rows are decoded directly into a padded FP16 +// LDS tile, while the same workgroup converts the corresponding F32 activation +// tile to FP16. Four wave64 WMMA accumulators cover the 32x32 result owned by +// each wave. A wave-private LDS slice transposes each accumulator for coalesced +// F32 publication. +// +// Both weight contracts use the unmodified GGUF layout: +// Q4_K: [output channel][input size / 256][144 bytes] +// Q6_K: [output channel][input size / 256][210 bytes] +// The entry point selects the format before the shared device template is +// instantiated, so inactive decode logic is absent from the emitted kernel. No +// persistent repacking or expanded-weight allocation is required. +config.decl @qwen3_moe.model.hidden_size : %value: index where [range(%value, 128, 32768), mul(%value, 128)] +config.decl @qwen3_moe.model.rms_epsilon : f32 + +func.def inline @ggml_q6k_f16_vector4(%weight: buffer, %weight_row_byte_base: offset, %q6_block: index, %q6_group: index, %packet: index) -> (vector<4xf16>) { + %c2 = index.constant 2 : index + %c4 = index.constant 4 : index + %c8 = index.constant 8 : index + %c16 = index.constant 16 : index + %block_bytes = index.constant 210 : offset + %qh_byte_add = index.constant 128 : offset + %scale_byte_add = index.constant 192 : offset + %d_byte_add = index.constant 208 : offset + %c4_i32v = vector.constant 4 : vector<1xi32> + %nibble_mask = vector.constant 252645135 : vector<1xi32> + %high_mask = vector.constant 50529027 : vector<1xi32> + %c32_f32v = vector.constant 32.0 : vector<4xf32> + %bounded_group = index.assume %q6_group [range(%q6_group, 0, 7)] : index + %bounded_packet = index.assume %packet [range(%packet, 0, 7)] : index + %block_byte_add = index.scale %q6_block, %block_bytes : index, offset -> offset + %block_byte_base = index.add %weight_row_byte_base, %block_byte_add : offset + %qh_byte_base = index.add %block_byte_base, %qh_byte_add : offset + %scale_byte_base = index.add %block_byte_base, %scale_byte_add : offset + %d_byte_base = index.add %block_byte_base, %d_byte_add : offset + %ql_view = buffer.view %weight[%block_byte_base] : buffer -> view<32xi32> + %qh_view = buffer.view %weight[%qh_byte_base] : buffer -> view<16xi32> + %scale_view = buffer.view %weight[%scale_byte_base] : buffer -> view<16xi8> + %d_view = buffer.view %weight[%d_byte_base] : buffer -> view<1xf16> + %group_in_half = index.rem %bounded_group, %c4 : index + %half = index.div %bounded_group, %c4 : index + %ql_side = index.rem %group_in_half, %c2 : index + %ql_half_word_base = index.mul %half, %c16 : index + %ql_side_word_add = index.mul %ql_side, %c8 : index + %ql_word_base = index.add %ql_half_word_base, %ql_side_word_add : index + %ql_word_index = index.add %ql_word_base, %bounded_packet : index + %qh_half_word_base = index.mul %half, %c8 : index + %qh_word_index = index.add %qh_half_word_base, %bounded_packet : index + %nibble = index.div %group_in_half, %c2 : index + %nibble_shift_index = index.mul %nibble, %c4 : index + %nibble_shift_i32 = index.cast %nibble_shift_index : index to i32 + %nibble_shift = vector.splat %nibble_shift_i32 : vector<1xi32> + %qh_shift_index = index.mul %group_in_half, %c2 : index + %qh_shift_i32 = index.cast %qh_shift_index : index to i32 + %qh_shift = vector.splat %qh_shift_i32 : vector<1xi32> + %scale_packet_half = index.div %bounded_packet, %c4 : index + %scale_group_base = index.mul %bounded_group, %c2 : index + %scale_index = index.add %scale_group_base, %scale_packet_half : index + %ql_word = vector.load %ql_view[%ql_word_index] : view<32xi32> -> vector<1xi32> + %qh_word = vector.load %qh_view[%qh_word_index] : view<16xi32> -> vector<1xi32> + %ql_shifted = vector.shrui %ql_word, %nibble_shift : vector<1xi32> + %ql = vector.andi %ql_shifted, %nibble_mask : vector<1xi32> + %qh_shifted = vector.shrui %qh_word, %qh_shift : vector<1xi32> + %qh_low = vector.andi %qh_shifted, %high_mask : vector<1xi32> + %qh = vector.shli %qh_low, %c4_i32v : vector<1xi32> + %code = vector.ori %ql, %qh : vector<1xi32> + %code_i8 = vector.bitcast %code : vector<1xi32> to vector<4xi8> + %code_f32 = vector.uitofp %code_i8 : vector<4xi8> to vector<4xf32> + %centered = vector.subf %code_f32, %c32_f32v : vector<4xf32> + %scale_i8 = view.load %scale_view[%scale_index] : view<16xi8> -> i8 + %d_f16 = view.load %d_view[0] : view<1xf16> -> f16 + %scale = scalar.sitofp %scale_i8 : i8 to f32 + %d = scalar.extf %d_f16 : f16 to f32 + %combined_scale = scalar.mulf %scale, %d : f32 + %combined_scale_vector = vector.splat %combined_scale : vector<4xf32> + %values_f32 = vector.mulf %centered, %combined_scale_vector : vector<4xf32> + %values = vector.fptrunc %values_f32 : vector<4xf32> to vector<4xf16> + func.return %values : vector<4xf16> +} + +func.def inline @qwen3_moe_q4k_chunk_pair_global(%weight: buffer, %row_byte_base: offset, %q4_block: index, %q4_group_pair: index, %q4_half: index, %header_words: vector<4xi32>) -> (vector<16xi8>, f32, f32, vector<16xi8>, f32, f32) { + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c4 = index.constant 4 : index + %c8 = index.constant 8 : index + %block_bytes = index.constant 144 : offset + %code_byte_add = index.constant 16 : offset + %c4_i32 = scalar.constant 4 : i32 + %nibble_mask = vector.constant 252645135 : vector<4xi32> + %bounded_pair = index.assume %q4_group_pair [range(%q4_group_pair, 0, 3)] : index + %bounded_half = index.assume %q4_half [range(%q4_half, 0, 1)] : index + %block_byte_add = index.scale %q4_block, %block_bytes : index, offset -> offset + %block_byte_base = index.add %row_byte_base, %block_byte_add : offset + %code_byte_base = index.add %block_byte_base, %code_byte_add : offset + %code_view = buffer.view %weight[%code_byte_base] : buffer -> view<32xi32> + %header_halves = vector.bitcast %header_words : vector<4xi32> to vector<8xf16> + %d_f16 = vector.extract %header_halves[0] : vector<8xf16> -> f16 + %dmin_f16 = vector.extract %header_halves[1] : vector<8xf16> -> f16 + %scale0 = vector.extract %header_words[1] : vector<4xi32> -> i32 + %scale1 = vector.extract %header_words[2] : vector<4xi32> -> i32 + %scale2 = vector.extract %header_words[3] : vector<4xi32> -> i32 + %d = scalar.extf %d_f16 : f16 to f32 + %dmin = scalar.extf %dmin_f16 : f16 to f32 + %pair_code_base = index.mul %bounded_pair, %c8 : index + %half_code_add = index.mul %bounded_half, %c4 : index + %code_index0 = index.add %pair_code_base, %half_code_add : index + %code_index = index.assume %code_index0 [range(%code_index0, 0, 28)] : index + %packed_codes = vector.load %code_view[%code_index] : view<32xi32> -> vector<4xi32> + %low_codes = vector.andi %packed_codes, %nibble_mask : vector<4xi32> + %c4_i32v = vector.splat %c4_i32 : vector<4xi32> + %high_shifted = vector.shrui %packed_codes, %c4_i32v : vector<4xi32> + %high_codes = vector.andi %high_shifted, %nibble_mask : vector<4xi32> + %q4_low = vector.bitcast %low_codes : vector<4xi32> to vector<16xi8> + %q4_high = vector.bitcast %high_codes : vector<4xi32> to vector<16xi8> + %low_group = index.mul %bounded_pair, %c2 : index + %high_group = index.add %low_group, %c1 : index + %low_scale, %low_minimum = func.call @qwen3_moe_q4k_scale_from_header(%scale0, %scale1, %scale2, %low_group) : (i32, i32, i32, index) -> (i32, i32) + %high_scale, %high_minimum = func.call @qwen3_moe_q4k_scale_from_header(%scale0, %scale1, %scale2, %high_group) : (i32, i32, i32, index) -> (i32, i32) + %low_scale_f32 = scalar.uitofp %low_scale : i32 to f32 + %low_minimum_f32 = scalar.uitofp %low_minimum : i32 to f32 + %high_scale_f32 = scalar.uitofp %high_scale : i32 to f32 + %high_minimum_f32 = scalar.uitofp %high_minimum : i32 to f32 + %low_d_scale = scalar.mulf %d, %low_scale_f32 : f32 + %low_dmin_scale = scalar.mulf %dmin, %low_minimum_f32 : f32 + %high_d_scale = scalar.mulf %d, %high_scale_f32 : f32 + %high_dmin_scale = scalar.mulf %dmin, %high_minimum_f32 : f32 + func.return %q4_low, %low_d_scale, %low_dmin_scale, %q4_high, %high_d_scale, %high_dmin_scale : vector<16xi8>, f32, f32, vector<16xi8>, f32, f32 +} + +func.def inline @qwen3_moe_q4k_q8_1_dot(%q4_values: vector<16xi8>, %d_scale: f32, %dmin_scale: f32, %q8_values: vector<16xi8>, %q8_d: f32, %q8_s: f32) -> (f32) { + %c0_i32 = scalar.constant 0 : i32 + %c0_i32v = vector.constant 0 : vector<4xi32> + %half_f32 = scalar.constant 0.5 : f32 + %partial_dots = vector.dot4i %q4_values, %q8_values, %c0_i32v : vector<16xi8>, vector<16xi8>, vector<4xi32> + %q_sum = vector.reduce %partial_dots, %c0_i32 : vector<4xi32>, i32 + %q_sum_f32 = scalar.sitofp %q_sum : i32 to f32 + %scaled_dot0 = scalar.mulf %q8_d, %d_scale : f32 + %scaled_dot = scalar.mulf %scaled_dot0, %q_sum_f32 : f32 + %q8_half_sum = scalar.mulf %q8_s, %half_f32 : f32 + %minimum_correction = scalar.mulf %dmin_scale, %q8_half_sum : f32 + %contribution = scalar.subf %scaled_dot, %minimum_correction : f32 + func.return %contribution : f32 +} + +func.def inline @qwen3_moe_q4k_q8_1_x4_paired_block_lane(%input_size: index, %weight: buffer, %weight_row_byte_base: offset, %q8_input: buffer, %q8_row_byte_base: offset, %q4_block: index, %block_lane: index) -> (f32) { + %bounded_input_size = index.assume %input_size [range(%input_size, 256, 32768), mul(%input_size, 256)] : index + %bounded_q4_block0 = index.assume %q4_block [range(%q4_block, 0, 127)] : index + %bounded_block_lane = index.assume %block_lane [range(%block_lane, 0, 7)] : index + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c4 = index.constant 4 : index + %c8 = index.constant 8 : index + %c128 = index.constant 128 : index + %c256 = index.constant 256 : index + %q4_block_bytes = index.constant 144 : offset + %q8_group_bytes = index.constant 144 : offset + %q8_payload_byte_add = index.constant 16 : offset + %q4_block_count = index.div %bounded_input_size, %c256 : index + %q8_group_count = index.div %bounded_input_size, %c128 : index + %bounded_q4_block, %bounded_q4_block_count = index.assume %bounded_q4_block0, %q4_block_count [lt(%bounded_q4_block0, %q4_block_count)] : index, index + %q4_group_pair0 = index.div %bounded_block_lane, %c2 : index + %q4_group_pair = index.assume %q4_group_pair0 [range(%q4_group_pair0, 0, 3)] : index + %q4_half0 = index.rem %bounded_block_lane, %c2 : index + %q4_half = index.assume %q4_half0 [range(%q4_half0, 0, 1)] : index + %q8_group_in_block0 = index.div %q4_group_pair, %c2 : index + %q8_group_in_block = index.assume %q8_group_in_block0 [range(%q8_group_in_block0, 0, 1)] : index + %pair_in_q8_group0 = index.rem %q4_group_pair, %c2 : index + %pair_in_q8_group = index.assume %pair_in_q8_group0 [range(%pair_in_q8_group0, 0, 1)] : index + %q8_low_inner_block0 = index.mul %pair_in_q8_group, %c2 : index + %q8_low_inner_block = index.assume %q8_low_inner_block0 [range(%q8_low_inner_block0, 0, 2)] : index + %q8_high_inner_block0 = index.add %q8_low_inner_block, %c1 : index + %q8_high_inner_block = index.assume %q8_high_inner_block0 [range(%q8_high_inner_block0, 1, 3)] : index + %q8_half_word_add = index.mul %q4_half, %c4 : index + %q8_low_inner_word_base = index.mul %q8_low_inner_block, %c8 : index + %q8_low_word_index0 = index.add %q8_low_inner_word_base, %q8_half_word_add : index + %q8_low_word_index = index.assume %q8_low_word_index0 [range(%q8_low_word_index0, 0, 20)] : index + %q8_high_inner_word_base = index.mul %q8_high_inner_block, %c8 : index + %q8_high_word_index0 = index.add %q8_high_inner_word_base, %q8_half_word_add : index + %q8_high_word_index = index.assume %q8_high_word_index0 [range(%q8_high_word_index0, 8, 28)] : index + %q8_low_ds_index0 = index.mul %q8_low_inner_block, %c2 : index + %q8_low_ds_index = index.assume %q8_low_ds_index0 [range(%q8_low_ds_index0, 0, 4)] : index + %q8_block_group_base = index.mul %bounded_q4_block, %c2 : index + %q8_group0 = index.add %q8_block_group_base, %q8_group_in_block : index + %q8_group, %bounded_q8_group_count = index.assume %q8_group0, %q8_group_count [lt(%q8_group0, %q8_group_count)] : index, index + %q8_group_byte_add = index.scale %q8_group, %q8_group_bytes : index, offset -> offset + %q8_group_byte_base = index.add %q8_row_byte_base, %q8_group_byte_add : offset + %q8_payload_byte_base = index.add %q8_group_byte_base, %q8_payload_byte_add : offset + %q8_ds_view = buffer.view %q8_input[%q8_group_byte_base] : buffer -> view<8xf16> + %q8_words_view = buffer.view %q8_input[%q8_payload_byte_base] : buffer -> view<32xi32> + %q8_ds = vector.load %q8_ds_view[%q8_low_ds_index] : view<8xf16> -> vector<4xf16> + %q8_low_d_f16 = vector.extract %q8_ds[0] : vector<4xf16> -> f16 + %q8_low_s_f16 = vector.extract %q8_ds[1] : vector<4xf16> -> f16 + %q8_high_d_f16 = vector.extract %q8_ds[2] : vector<4xf16> -> f16 + %q8_high_s_f16 = vector.extract %q8_ds[3] : vector<4xf16> -> f16 + %q8_low_d = scalar.extf %q8_low_d_f16 : f16 to f32 + %q8_low_s = scalar.extf %q8_low_s_f16 : f16 to f32 + %q8_high_d = scalar.extf %q8_high_d_f16 : f16 to f32 + %q8_high_s = scalar.extf %q8_high_s_f16 : f16 to f32 + %q8_low_words = vector.load %q8_words_view[%q8_low_word_index] : view<32xi32> -> vector<4xi32> + %q8_high_words = vector.load %q8_words_view[%q8_high_word_index] : view<32xi32> -> vector<4xi32> + %q8_low_values = vector.bitcast %q8_low_words : vector<4xi32> to vector<16xi8> + %q8_high_values = vector.bitcast %q8_high_words : vector<4xi32> to vector<16xi8> + %q4_block_byte_add = index.scale %bounded_q4_block, %q4_block_bytes : index, offset -> offset + %q4_block_byte_base = index.add %weight_row_byte_base, %q4_block_byte_add : offset + %q4_header_view = buffer.view %weight[%q4_block_byte_base] : buffer -> view<4xi32> + %q4_header_words = vector.load %q4_header_view[0] : view<4xi32> -> vector<4xi32> + %q4_low, %low_d_scale, %low_dmin_scale, %q4_high, %high_d_scale, %high_dmin_scale = func.call @qwen3_moe_q4k_chunk_pair_global(%weight, %weight_row_byte_base, %bounded_q4_block, %q4_group_pair, %q4_half, %q4_header_words) : (buffer, offset, index, index, index, vector<4xi32>) -> (vector<16xi8>, f32, f32, vector<16xi8>, f32, f32) + %low = func.call @qwen3_moe_q4k_q8_1_dot(%q4_low, %low_d_scale, %low_dmin_scale, %q8_low_values, %q8_low_d, %q8_low_s) : (vector<16xi8>, f32, f32, vector<16xi8>, f32, f32) -> (f32) + %high = func.call @qwen3_moe_q4k_q8_1_dot(%q4_high, %high_d_scale, %high_dmin_scale, %q8_high_values, %q8_high_d, %q8_high_s) : (vector<16xi8>, f32, f32, vector<16xi8>, f32, f32) -> (f32) + %pair = scalar.addf %low, %high : f32 + func.return %pair : f32 +} + +func.def inline @qwen3_moe_q4k_q8_1_x4_paired_row_lane(%input_size: index, %weight: buffer, %weight_row_byte_base: offset, %q8_input: buffer, %q8_row_byte_base: offset, %lane: index) -> (f32) { + %bounded_input_size = index.assume %input_size [range(%input_size, 256, 32768), mul(%input_size, 256)] : index + %bounded_lane = index.assume %lane [range(%lane, 0, 31)] : index + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c4 = index.constant 4 : index + %c8 = index.constant 8 : index + %c256 = index.constant 256 : index + %c1023 = index.constant 1023 : index + %c1024 = index.constant 1024 : index + %c0_f32 = scalar.constant 0.0 : f32 + %q4_block_count = index.div %bounded_input_size, %c256 : index + %padded_input_size = index.add %bounded_input_size, %c1023 : index + %iteration_count = index.div %padded_input_size, %c1024 : index + %lane_q4_block = index.div %bounded_lane, %c8 : index + %block_lane0 = index.rem %bounded_lane, %c8 : index + %block_lane = index.assume %block_lane0 [range(%block_lane0, 0, 7)] : index + %sum = scf.for %iteration = [%c0 to %iteration_count step %c1](%iteration_acc = %c0_f32 : f32) -> (f32) unroll { + %iteration_q4_block = index.mul %iteration, %c4 : index + %q4_block0 = index.add %iteration_q4_block, %lane_q4_block : index + %valid_q4_block = index.cmp ult, %q4_block0, %q4_block_count : index + %contribution = scf.if %valid_q4_block -> (f32) { + %q4_block, %bounded_q4_block_count = index.assume %q4_block0, %q4_block_count [lt(%q4_block0, %q4_block_count)] : index, index + %pair = func.call @qwen3_moe_q4k_q8_1_x4_paired_block_lane(%bounded_input_size, %weight, %weight_row_byte_base, %q8_input, %q8_row_byte_base, %q4_block, %block_lane) : (index, buffer, offset, buffer, offset, index, index) -> (f32) + scf.yield %pair : f32 + } else { + scf.yield %c0_f32 : f32 + } + %next = scalar.addf %iteration_acc, %contribution : f32 + scf.yield %next : f32 + } + func.return %sum : f32 +} + +func.def inline @qwen3_moe_q4k_scale_from_header(%scale0: i32, %scale1: i32, %scale2: i32, %q4_group: index) -> (i32, i32) { + %c4 = index.constant 4 : index + %c8 = index.constant 8 : index + %c2_i32 = scalar.constant 2 : i32 + %c4_i32 = scalar.constant 4 : i32 + %c15_i32 = scalar.constant 15 : i32 + %c48_i32 = scalar.constant 48 : i32 + %bounded_group = index.assume %q4_group [range(%q4_group, 0, 7)] : index + %is_low_group = index.cmp ult, %bounded_group, %c4 : index + %scale_lane = index.rem %bounded_group, %c4 : index + %scale_shift_index = index.mul %scale_lane, %c8 : index + %scale_shift = index.cast %scale_shift_index : index to i32 + %high_shift = scalar.addi %scale_shift, %c2_i32 : i32 + %minimum_shift = scalar.addi %scale_shift, %c4_i32 : i32 + %selected_scale_source = scf.select %is_low_group, %scale0, %scale2 : i32 + %selected_minimum_source = scf.select %is_low_group, %scale1, %scale2 : i32 + %selected_scale_high_shift = scf.select %is_low_group, %scale_shift, %high_shift : i32 + %selected_minimum_low_shift = scf.select %is_low_group, %scale_shift, %minimum_shift : i32 + %scale_low0 = scalar.shrui %selected_scale_source, %scale_shift : i32 + %scale_low = scalar.andi %scale_low0, %c15_i32 : i32 + %scale_high0 = scalar.shrui %scale0, %selected_scale_high_shift : i32 + %scale_high = scalar.andi %scale_high0, %c48_i32 : i32 + %scale = scalar.ori %scale_low, %scale_high : i32 + %minimum_low0 = scalar.shrui %selected_minimum_source, %selected_minimum_low_shift : i32 + %minimum_low = scalar.andi %minimum_low0, %c15_i32 : i32 + %minimum_high0 = scalar.shrui %scale1, %selected_scale_high_shift : i32 + %minimum_high = scalar.andi %minimum_high0, %c48_i32 : i32 + %minimum = scalar.ori %minimum_low, %minimum_high : i32 + func.return %scale, %minimum : i32, i32 +} + +func.def inline @qwen3_moe_rmsnorm_quantize_q8_1_x4_body(%publish_normalized: i1, %reduction_subgroup_count0: index, %token_count: index, %token0: index, %input: buffer, %weight: buffer, %normalized_output: buffer, %q8_output: buffer) { + %hidden_size0 = config.get @qwen3_moe.model.hidden_size : index + %epsilon = config.get @qwen3_moe.model.rms_epsilon : f32 + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 2048)] : index + %reduction_subgroup_count = index.assume %reduction_subgroup_count0 [range(%reduction_subgroup_count0, 1, 8)] : index + %hidden_size = index.assume %hidden_size0 [range(%hidden_size0, 128, 32768), mul(%hidden_size0, 128)] : index + %workitem = kernel.workitem.id : index + %subgroup0 = kernel.subgroup.id : index + %subgroup = index.assume %subgroup0 [range(%subgroup0, 0, 7)] : index + %lane = kernel.subgroup.lane.id : index + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c4 = index.constant 4 : index + %c8 = index.constant 8 : index + %c16 = index.constant 16 : index + %c32 = index.constant 32 : index + %c128 = index.constant 128 : index + %c256 = index.constant 256 : index + %c1024 = index.constant 1024 : index + %group_bytes = index.constant 144 : offset + %payload_byte_add = index.constant 16 : offset + %scratch_d_byte_add = index.constant 1024 : offset + %scratch_bytes = index.constant 1152 : offset + %c0_offset = index.constant 0 : offset + %c0_f32 = scalar.constant 0.0 : f32 + %c1_f32 = scalar.constant 1.0 : f32 + %c127 = scalar.constant 127.0 : f32 + %c0_f32x4 = vector.constant 0.0 : vector<4xf32> + %valid_token = index.cmp ult, %token0, %bounded_token_count : index + %safe_token0 = scf.select %valid_token, %token0, %c0 : index + %token, %launch_token_count = index.assume %safe_token0, %bounded_token_count [lt(%safe_token0, %bounded_token_count)] : index, index + %hidden_size_i32 = index.cast %hidden_size : index to i32 + %hidden_size_f32 = scalar.sitofp %hidden_size_i32 : i32 to f32 + %physical_group_count = index.div %hidden_size, %c128 : index + %row_bytes = index.scale %physical_group_count, %group_bytes : index, offset -> offset + %token_output_byte_base = index.scale %token, %row_bytes : index, offset -> offset + %input_noalias, %weight_noalias, %q8_output_noalias = buffer.assume.noalias %input, %weight, %q8_output : buffer, buffer, buffer + %input_view = buffer.view %input_noalias[%c0_offset] : buffer -> view<[%launch_token_count]x[%hidden_size]xf32> + %weight_view = buffer.view %weight_noalias[%c0_offset] : buffer -> view<[%hidden_size]xf32> + %normalized_output_view = buffer.view %normalized_output[%c0_offset] : buffer -> view<[%launch_token_count]x[%hidden_size]xf32> + %scratch = buffer.alloca align(16) %scratch_bytes : buffer + %scratch_values = buffer.view %scratch[%c0_offset] : buffer -> view<256xf32> + %scratch_d = buffer.view %scratch[%scratch_d_byte_add] : buffer -> view<32xf32> + // Reduce the complete row before any block-local quantization. + %thread_sum = scf.for %channel = [%workitem to %hidden_size step %c256](%running_sum = %c0_f32 : f32) -> (f32) { + %value = view.load %input_view[%token, %channel] : view<[%launch_token_count]x[%hidden_size]xf32> -> f32 + %square = scalar.mulf %value, %value : f32 + %next_sum = scalar.addf %running_sum, %square : f32 + scf.yield %next_sum : f32 + } + %subgroup_sum = kernel.subgroup.reduce %thread_sum : f32 + %is_subgroup_leader = index.cmp eq, %lane, %c0 : index + scf.if %is_subgroup_leader { + view.store %subgroup_sum, %scratch_values[%subgroup] : f32, view<256xf32> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + %is_reduction_subgroup = index.cmp eq, %subgroup, %c0 : index + %is_reduction_lane = index.cmp ult, %lane, %reduction_subgroup_count : index + %loads_subgroup_sum = scalar.andi %is_reduction_subgroup, %is_reduction_lane : i1 + %subgroup_partial = scf.if %loads_subgroup_sum -> (f32) { + %value = view.load %scratch_values[%lane] : view<256xf32> -> f32 + scf.yield %value : f32 + } else { + scf.yield %c0_f32 : f32 + } + %row_sum = kernel.subgroup.reduce %subgroup_partial : f32 + %writes_scale = scalar.andi %is_reduction_subgroup, %is_subgroup_leader : i1 + scf.if %writes_scale { + %mean = scalar.divf %row_sum, %hidden_size_f32 : f32 + %biased_mean = scalar.addf %mean, %epsilon : f32 + %scale = scalar.rsqrtf %biased_mean : f32 + view.store %scale, %scratch_values[%c0] : f32, view<256xf32> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + %row_scale = view.load %scratch_values[%c0] : view<256xf32> -> f32 + %row_scale_vector = vector.splat %row_scale : vector<4xf32> + %publishes_normalized = scalar.andi %publish_normalized, %valid_token : i1 + // Reuse one scratch frame for each 1024-element stripe. Hidden sizes need + // only be divisible by 128; inactive workitems in the final stripe carry + // zeros and never publish. + scf.for %stripe_base = [%c0 to %hidden_size step %c1024] { + %word_element_add = index.mul %workitem, %c4 : index + %channel = index.add %stripe_base, %word_element_add : index + %valid_word = index.cmp ult, %channel, %hidden_size : index + %mask = vector.mask.range [%channel to %hidden_size step %c1] : index -> vector<4xi1> + %input_values = vector.load.mask %input_view[%token, %channel], %mask, %c0_f32x4 : view<[%launch_token_count]x[%hidden_size]xf32>, vector<4xi1>, vector<4xf32> + %learned_weights = vector.load.mask %weight_view[%channel], %mask, %c0_f32x4 : view<[%hidden_size]xf32>, vector<4xi1>, vector<4xf32> + %normalized0 = vector.mulf %input_values, %row_scale_vector : vector<4xf32> + %normalized = vector.mulf %normalized0, %learned_weights : vector<4xf32> + scf.if %publishes_normalized { + vector.store.mask %normalized, %normalized_output_view[%token, %channel], %mask : vector<4xf32>, view<[%launch_token_count]x[%hidden_size]xf32>, vector<4xi1> + } + %absolute_values = vector.absf %normalized : vector<4xf32> + %thread_max = vector.reduce %absolute_values, %c0_f32 : vector<4xf32>, f32 + view.store %thread_max, %scratch_values[%workitem] : f32, view<256xf32> + kernel.barrier scope(workgroup) ordering(acq_rel) + %word_in_block = index.rem %workitem, %c8 : index + %block_in_stripe = index.div %workitem, %c8 : index + %is_block_leader = index.cmp eq, %word_in_block, %c0 : index + %writes_block_d = scalar.andi %valid_word, %is_block_leader : i1 + scf.if %writes_block_d { + %cohort_base = index.mul %block_in_stripe, %c8 : index + %cohort_maxima = vector.load %scratch_values[%cohort_base] : view<256xf32> -> vector<8xf32> + %amax = vector.reduce %cohort_maxima, %c0_f32 : vector<8xf32>, f32 + %d = scalar.divf %amax, %c127 : f32 + view.store %d, %scratch_d[%block_in_stripe] : f32, view<32xf32> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + %d = scf.if %valid_word -> (f32) { + %block_d = view.load %scratch_d[%block_in_stripe] : view<32xf32> -> f32 + scf.yield %block_d : f32 + } else { + scf.yield %c0_f32 : f32 + } + %d_nonzero = scalar.cmpf one, %d, %c0_f32 : f32 + %d_inverse = scf.if %d_nonzero -> (f32) { + %inverse = scalar.divf %c1_f32, %d : f32 + scf.yield %inverse : f32 + } else { + scf.yield %c0_f32 : f32 + } + %d_inverse_vector = vector.splat %d_inverse : vector<4xf32> + %scaled_values = vector.mulf %normalized, %d_inverse_vector : vector<4xf32> + %rounded_values = vector.roundf %scaled_values : vector<4xf32> + %quantized_values = vector.fptosi %rounded_values : vector<4xf32> to vector<4xi8> + %packed_word = vector.bitcast %quantized_values : vector<4xi8> to vector<1xi32> + %publishes_q8_word = scalar.andi %valid_word, %valid_token : i1 + scf.if %publishes_q8_word { + %q8_block = index.div %channel, %c32 : index + %physical_group = index.div %q8_block, %c4 : index + %block_in_group = index.rem %q8_block, %c4 : index + %group_byte_add = index.scale %physical_group, %group_bytes : index, offset -> offset + %group_byte_offset = index.add %token_output_byte_base, %group_byte_add : offset + %payload_byte_offset = index.add %group_byte_offset, %payload_byte_add : offset + %group_ds = buffer.view %q8_output_noalias[%group_byte_offset] : buffer -> view<8xf16> + %group_qs = buffer.view %q8_output_noalias[%payload_byte_offset] : buffer -> view<32xi32> + %block_word_base = index.mul %block_in_group, %c8 : index + %packed_word_index0 = index.add %block_word_base, %word_in_block : index + %packed_word_index = index.assume %packed_word_index0 [range(%packed_word_index0, 0, 31)] : index + vector.store %packed_word, %group_qs[%packed_word_index] : vector<1xi32>, view<32xi32> + } + %thread_quantized_sum = vector.reduce %rounded_values, %c0_f32 : vector<4xf32>, f32 + view.store %thread_quantized_sum, %scratch_values[%workitem] : f32, view<256xf32> + kernel.barrier scope(workgroup) ordering(acq_rel) + %publishes_block_ds = scalar.andi %writes_block_d, %valid_token : i1 + scf.if %publishes_block_ds { + %cohort_base = index.mul %block_in_stripe, %c8 : index + %cohort_sums = vector.load %scratch_values[%cohort_base] : view<256xf32> -> vector<8xf32> + %quantized_sum = vector.reduce %cohort_sums, %c0_f32 : vector<8xf32>, f32 + %s = scalar.mulf %quantized_sum, %d : f32 + %q8_block = index.div %channel, %c32 : index + %physical_group = index.div %q8_block, %c4 : index + %block_in_group = index.rem %q8_block, %c4 : index + %group_byte_add = index.scale %physical_group, %group_bytes : index, offset -> offset + %group_byte_offset = index.add %token_output_byte_base, %group_byte_add : offset + %group_ds = buffer.view %q8_output_noalias[%group_byte_offset] : buffer -> view<8xf16> + %d_f16 = scalar.fptrunc %d : f32 to f16 + %s_f16 = scalar.fptrunc %s : f32 to f16 + %ds_index = index.mul %block_in_group, %c2 : index + %s_index = index.add %ds_index, %c1 : index + view.store %d_f16, %group_ds[%ds_index] : f16, view<8xf16> + view.store %s_f16, %group_ds[%s_index] : f16, view<8xf16> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + } + func.return +} + +amdgpu.target @qwen3_moe_dense_gfx11_wave64 {subgroup_size = 64} + +amdgpu.target @qwen3_moe_dense_gfx11_wave32 {subgroup_size = 32} + +target.decl @qwen3_moe_attention_prepare_gfx11_wave32 + +config.decl @qwen3_moe.dense_quantized.input_size : %value: index where [range(%value, 256, 32768), mul(%value, 256)] + +config.decl @qwen3_moe.dense_quantized.output_size : %value: index where [range(%value, 1, 262144)] + +// Selects whether publication overwrites the output (0) or accumulates the +// projection into a caller-provided residual (1). +config.decl @qwen3_moe.dense_quantized.output_accumulation : %value: index where [range(%value, 0, 1)] + +config.decl @qwen3_moe.workload.token_capacity : %value: index where [range(%value, 1, 2048)] + +// These declarations are used only by the differential correctness case. The +// dense production kernel has no routing input or dependency. +kernel.decl @qwen3_moe_build_expert_table(%token_count: index, %route_count: index, %route_stride: index, %expert_count: index) launch(%token_count: index, %route_count: index, %route_stride: index, %expert_count: index, %route_ids: buffer, %expert_table: buffer) + +kernel.decl @qwen3_moe_routed_linear_q4k_f16_wmma(%token_count: index) launch(%token_count: index, %input: buffer, %expert_table: buffer, %weight: buffer, %output: buffer) + +// Shared raw-layout Q4_K row contraction primitive. Each lane consumes both +// nibbles of one packed-code load. +func.decl @qwen3_moe_q4k_q8_1_x4_paired_row_lane(%input_size: index, %weight: buffer, %weight_row_byte_base: offset, %q8_input: buffer, %q8_row_byte_base: offset, %lane: index) -> (f32) + +// Shared raw Q6_K decoder linked from the GGML physical-layout module. +func.decl @ggml_q6k_f16_vector4(%weight: buffer, %weight_row_byte_base: offset, %q6_block: index, %q6_group: index, %packet: index) -> (vector<4xf16>) + +// These declarations are used only by the Q6_K differential correctness case. +kernel.decl @ggml_quantize_q8_1_x4_f32(%token_count: index, %input_size: index) launch(%token_count: index, %input_size: index, %input: buffer, %output: buffer) + +kernel.decl @ggml_linear_q6k_q8_1_x4(%token_count: index, %input_size: index, %output_size: index) launch(%token_count: index, %input_size: index, %output_size: index, %q8_input: buffer, %weight: buffer, %output: buffer) + +kernel.decl @qwen3_moe_rmsnorm_f32_quantize_q8_1_x4(%token_count: index) launch(%token_count: index, %input: buffer, %rms_weight: buffer, %normalized_f32_output: buffer, %q8_1_x4_output: buffer) + +// Decodes the four adjacent Q4_K values owned by one load packet. One packed +// header load supplies both half scales and all twelve six-bit group scales; +// one packed code load supplies four adjacent unsigned nibbles. The final +// affine conversion is formed in F32 before truncation to the FP16 WMMA input. +func.def inline @qwen3_moe_dense_q4k_wmma_vector4(%weight: buffer, %row_byte_base: offset, %q4_block: index, %q4_group: index, %packet: index) -> (vector<4xf16>) { + %c0 = index.constant 0 : index + %c2 = index.constant 2 : index + %c4 = index.constant 4 : index + %c8 = index.constant 8 : index + %block_bytes = index.constant 144 : offset + %scale_offset = index.constant 4 : offset + %code_offset = index.constant 16 : offset + %c2_i32 = scalar.constant 2 : i32 + %c4_i32 = scalar.constant 4 : i32 + %c15 = vector.constant 15 : vector<1xi32> + %c48 = vector.constant 48 : vector<1xi32> + %q4_mask = vector.constant 252645135 : vector<1xi32> + %bounded_group = index.assume %q4_group [range(%q4_group, 0, 7)] : index + %bounded_packet = index.assume %packet [range(%packet, 0, 7)] : index + %block_byte_add = index.scale %q4_block, %block_bytes : index, offset -> offset + %block_byte_base = index.add %row_byte_base, %block_byte_add : offset + %scale_byte_base = index.add %block_byte_base, %scale_offset : offset + %code_byte_base = index.add %block_byte_base, %code_offset : offset + %dm_view = buffer.view %weight[%block_byte_base] : buffer -> view<2xf16> + %scale_view = buffer.view %weight[%scale_byte_base] : buffer -> view<3xi32> + %code_view = buffer.view %weight[%code_byte_base] : buffer -> view<32xi32> + %dm = vector.load %dm_view[%c0] : view<2xf16> -> vector<2xf16> + %scales = vector.load %scale_view[%c0] : view<3xi32> -> vector<3xi32> + %d_f16 = vector.extract %dm[0] : vector<2xf16> -> f16 + %dmin_f16 = vector.extract %dm[1] : vector<2xf16> -> f16 + %d = scalar.extf %d_f16 : f16 to f32 + %dmin = scalar.extf %dmin_f16 : f16 to f32 + %q_page0 = index.div %bounded_group, %c2 : index + %q_page = index.mul %q_page0, %c8 : index + %q_word_index0 = index.add %q_page, %bounded_packet : index + %q_word_index = index.assume %q_word_index0 [range(%q_word_index0, 0, 31)] : index + %is_low = index.cmp ult, %bounded_group, %c4 : index + %scale_lane = index.rem %bounded_group, %c4 : index + %scale_shift_index = index.mul %scale_lane, %c8 : index + %scale_shift_i32 = index.cast %scale_shift_index : index to i32 + %scale_shift = vector.splat %scale_shift_i32 : vector<1xi32> + %scale0_i32 = vector.extract %scales[0] : vector<3xi32> -> i32 + %scale1_i32 = vector.extract %scales[1] : vector<3xi32> -> i32 + %scale2_i32 = vector.extract %scales[2] : vector<3xi32> -> i32 + %scale0 = vector.splat %scale0_i32 : vector<1xi32> + %scale1 = vector.splat %scale1_i32 : vector<1xi32> + %scale2 = vector.splat %scale2_i32 : vector<1xi32> + %high_shift_i32 = scalar.addi %scale_shift_i32, %c2_i32 : i32 + %minimum_shift_i32 = scalar.addi %scale_shift_i32, %c4_i32 : i32 + %selected_scale_source = scf.select %is_low, %scale0, %scale2 : vector<1xi32> + %selected_minimum_source = scf.select %is_low, %scale1, %scale2 : vector<1xi32> + %selected_scale_high_shift_i32 = scf.select %is_low, %scale_shift_i32, %high_shift_i32 : i32 + %selected_minimum_low_shift_i32 = scf.select %is_low, %scale_shift_i32, %minimum_shift_i32 : i32 + %selected_scale_high_shift = vector.splat %selected_scale_high_shift_i32 : vector<1xi32> + %selected_minimum_low_shift = vector.splat %selected_minimum_low_shift_i32 : vector<1xi32> + %scale_low0 = vector.shrui %selected_scale_source, %scale_shift : vector<1xi32> + %scale_low = vector.andi %scale_low0, %c15 : vector<1xi32> + %scale_high0 = vector.shrui %scale0, %selected_scale_high_shift : vector<1xi32> + %scale_high = vector.andi %scale_high0, %c48 : vector<1xi32> + %scale = vector.ori %scale_low, %scale_high : vector<1xi32> + %minimum_low0 = vector.shrui %selected_minimum_source, %selected_minimum_low_shift : vector<1xi32> + %minimum_low = vector.andi %minimum_low0, %c15 : vector<1xi32> + %minimum_high0 = vector.shrui %scale1, %selected_scale_high_shift : vector<1xi32> + %minimum_high = vector.andi %minimum_high0, %c48 : vector<1xi32> + %minimum = vector.ori %minimum_low, %minimum_high : vector<1xi32> + %scale_f32 = vector.uitofp %scale : vector<1xi32> to vector<1xf32> + %minimum_f32 = vector.uitofp %minimum : vector<1xi32> to vector<1xf32> + %d_vector1 = vector.splat %d : vector<1xf32> + %dmin_vector1 = vector.splat %dmin : vector<1xf32> + %d_scale_vector1 = vector.mulf %d_vector1, %scale_f32 : vector<1xf32> + %minimum_scale_vector1 = vector.mulf %dmin_vector1, %minimum_f32 : vector<1xf32> + %d_scale = vector.extract %d_scale_vector1[0] : vector<1xf32> -> f32 + %minimum_scale = vector.extract %minimum_scale_vector1[0] : vector<1xf32> -> f32 + %q_word = vector.load %code_view[%q_word_index] : view<32xi32> -> vector<1xi32> + %q_half = index.rem %bounded_group, %c2 : index + %q_shift_index = index.mul %q_half, %c4 : index + %q_shift_i32 = index.cast %q_shift_index : index to i32 + %q_shift = vector.splat %q_shift_i32 : vector<1xi32> + %shifted_q = vector.shrui %q_word, %q_shift : vector<1xi32> + %masked_q = vector.andi %shifted_q, %q4_mask : vector<1xi32> + %q_i8 = vector.bitcast %masked_q : vector<1xi32> to vector<4xi8> + %q_f32 = vector.uitofp %q_i8 : vector<4xi8> to vector<4xf32> + %negative_minimum_scale = scalar.negf %minimum_scale : f32 + %q0 = vector.extract %q_f32[0] : vector<4xf32> -> f32 + %q1 = vector.extract %q_f32[1] : vector<4xf32> -> f32 + %q2 = vector.extract %q_f32[2] : vector<4xf32> -> f32 + %q3 = vector.extract %q_f32[3] : vector<4xf32> -> f32 + %value0 = scalar.fmaf %q0, %d_scale, %negative_minimum_scale : f32 + %value1 = scalar.fmaf %q1, %d_scale, %negative_minimum_scale : f32 + %value2 = scalar.fmaf %q2, %d_scale, %negative_minimum_scale : f32 + %value3 = scalar.fmaf %q3, %d_scale, %negative_minimum_scale : f32 + %half0 = scalar.fptrunc %value0 : f32 to f16 + %half1 = scalar.fptrunc %value1 : f32 to f16 + %half2 = scalar.fptrunc %value2 : f32 to f16 + %half3 = scalar.fptrunc %value3 : f32 to f16 + %result = vector.from_elements %half0, %half1, %half2, %half3 : vector<4xf16> + func.return %result : vector<4xf16> +} + +// Low-fixed-cost Q4_K body shared by the ordinary and completion-fused decode +// exports. Each wave owns one output channel and contracts a Q8_1 activation +// row directly against the original GGUF weight row. +func.def inline @qwen3_moe_dense_linear_q4k_q8_1_x4_body(%publish_output: i1, %token_count: index, %token0: index, %q8_input: buffer, %weight: buffer, %output: buffer) { + %input_size = config.get @qwen3_moe.dense_quantized.input_size : index + %output_size = config.get @qwen3_moe.dense_quantized.output_size : index + %output_accumulation = config.get @qwen3_moe.dense_quantized.output_accumulation : index + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 2048)] : index + %bounded_input_size = index.assume %input_size [range(%input_size, 256, 32768), mul(%input_size, 256)] : index + %bounded_output_size = index.assume %output_size [range(%output_size, 1, 262144)] : index + %channel_tile = kernel.workgroup.id : index + %subgroup0 = kernel.subgroup.id : index + %lane = kernel.subgroup.lane.id : index + %c1 = index.constant 1 : index + %c8 = index.constant 8 : index + %c128 = index.constant 128 : index + %c144_bytes = index.constant 144 : offset + %c256 = index.constant 256 : index + %c0_i32 = scalar.constant 0 : i32 + %c0_f32 = scalar.constant 0.0 : f32 + %c0_offset = index.constant 0 : offset + %token, %launch_token_count = index.assume %token0, %bounded_token_count [lt(%token0, %bounded_token_count)] : index, index + %subgroup = index.assume %subgroup0 [range(%subgroup0, 0, 7)] : index + %channel_base = index.mul %channel_tile, %c8 : index + %channel = index.add %channel_base, %subgroup : index + %valid_channel = index.cmp ult, %channel, %bounded_output_size : index + %lane_i32 = index.cast %lane : index to i32 + %is_lane_zero = scalar.cmpi eq, %lane_i32, %c0_i32 : i32 + %q4_block_count = index.div %bounded_input_size, %c256 : index + %weight_row_bytes = index.scale %q4_block_count, %c144_bytes : index, offset -> offset + %weight_row_byte_base = index.scale %channel, %weight_row_bytes : index, offset -> offset + %q8_group_count = index.div %bounded_input_size, %c128 : index + %q8_row_bytes = index.scale %q8_group_count, %c144_bytes : index, offset -> offset + %q8_row_byte_base = index.scale %token, %q8_row_bytes : index, offset -> offset + %q8_noalias, %weight_noalias, %output_noalias = buffer.assume.noalias %q8_input, %weight, %output : buffer, buffer, buffer + %lane_sum = scf.if %publish_output -> (f32) { + %channel_sum = scf.if %valid_channel -> (f32) { + %sum = func.call @qwen3_moe_q4k_q8_1_x4_paired_row_lane(%bounded_input_size, %weight_noalias, %weight_row_byte_base, %q8_noalias, %q8_row_byte_base, %lane) : (index, buffer, offset, buffer, offset, index) -> (f32) + scf.yield %sum : f32 + } else { + scf.yield %c0_f32 : f32 + } + scf.yield %channel_sum : f32 + } else { + scf.yield %c0_f32 : f32 + } + %dot = kernel.subgroup.reduce %lane_sum : f32 + %accumulates_output = index.cmp eq, %output_accumulation, %c1 : index + scf.if %publish_output { + scf.if %valid_channel { + scf.if %is_lane_zero { + %output_view = buffer.view %output_noalias[%c0_offset] : buffer -> view<[%launch_token_count]x[%bounded_output_size]xf32> + %result = scf.if %accumulates_output -> (f32) { + %residual = view.load %output_view[%token, %channel] : view<[%launch_token_count]x[%bounded_output_size]xf32> -> f32 + %sum = scalar.addf %residual, %dot : f32 + scf.yield %sum : f32 + } else { + scf.yield %dot : f32 + } + view.store %result, %output_view[%token, %channel] : f32, view<[%launch_token_count]x[%bounded_output_size]xf32> + } + } + } + func.return +} + +// Direct decode entry point. This complements the grouped WMMA schedule below: +// it gives up cross-token weight reuse in exchange for doing no padded matrix +// work when fewer than 32 tokens are available. +kernel.def target(@qwen3_moe_dense_gfx11_wave32) @qwen3_moe_dense_linear_q4k_q8_1_x4(%token_count: index) { + %output_size = config.get @qwen3_moe.dense_quantized.output_size : index + %token_capacity = config.get @qwen3_moe.workload.token_capacity : index + %c1 = index.constant 1 : index + %c7 = index.constant 7 : index + %c8 = index.constant 8 : index + %c256 = index.constant 256 : index + %padded_output_size = index.add %output_size, %c7 : index + %output_tiles = index.div %padded_output_size, %c8 : index + kernel.launch.config workgroups(%output_tiles, %token_capacity, %c1) workgroup_size(%c256, %c1, %c1) : index +} launch(%token_count: index, %q8_input: buffer, %weight: buffer, %output: buffer) { + %token_capacity = config.get @qwen3_moe.workload.token_capacity : index + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 2048), le(%token_count, %token_capacity)] : index + %token0 = kernel.workgroup.id : index + %c0 = index.constant 0 : index + %valid_token = index.cmp ult, %token0, %bounded_token_count : index + %safe_token = scf.select %valid_token, %token0, %c0 : index + func.call @qwen3_moe_dense_linear_q4k_q8_1_x4_body(%valid_token, %bounded_token_count, %safe_token, %q8_input, %weight, %output) : (i1, index, index, buffer, buffer, buffer) + kernel.return +} + +// Decode-only route that publishes the normalized F32 and Q8_1 x4 rows +// consumed by the following feed-forward boundary. +kernel.def target(@qwen3_moe_attention_prepare_gfx11_wave32) @qwen3_moe_dense_linear_q4k_q8_1_x4_next_q8(%token_count: index) { + %output_size = config.get @qwen3_moe.dense_quantized.output_size : index + %c1 = index.constant 1 : index + %c8 = index.constant 8 : index + %c256 = index.constant 256 : index + %output_tiles = index.div %output_size, %c8 : index + kernel.launch.config workgroups(%output_tiles, %c1, %c1) workgroup_size(%c256, %c1, %c1) : index +} launch(%token_count: index, %q8_input: buffer, %weight: buffer, %output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counter: buffer, %next_q8_output: buffer) { + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 1)] : index + %publish_output = scalar.constant true : i1 + %projection_token = kernel.workgroup.id : index + func.call @qwen3_moe_dense_linear_q4k_q8_1_x4_body(%publish_output, %bounded_token_count, %projection_token, %q8_input, %weight, %output) : (i1, index, index, buffer, buffer, buffer) + %output_size0 = config.get @qwen3_moe.dense_quantized.output_size : index + %output_size = index.assume %output_size0 [range(%output_size0, 128, 32768), mul(%output_size0, 128)] : index + %token0 = kernel.workgroup.id : index + %token = index.assume %token0 [lt(%token0, %bounded_token_count)] : index + %c0 = index.constant 0 : index + %c8 = index.constant 8 : index + %workitem = kernel.workitem.id : index + %output_tile_count = index.div %output_size, %c8 : index + %is_arrival_workitem = index.cmp eq, %workitem, %c0 : index + %c0_i32 = scalar.constant 0 : i32 + %c1_i32 = scalar.constant 1 : i32 + %c0_offset = index.constant 0 : offset + %counter_scratch_bytes = index.constant 4 : offset + %output_noalias, %norm_weight_noalias, %normalized_output_noalias, %completion_counter_noalias, %next_q8_output_noalias = buffer.assume.noalias %output, %norm_weight, %normalized_output, %completion_counter, %next_q8_output : buffer, buffer, buffer, buffer, buffer + %completion_counter_aligned = buffer.assume.alignment %completion_counter_noalias {minimum_alignment = 16} : buffer + %completion_counter_view = buffer.view %completion_counter_aligned[%c0_offset] : buffer -> view<1xi32> + %counter_scratch = buffer.alloca align(4) %counter_scratch_bytes : buffer + %counter_scratch_view = buffer.view %counter_scratch[%c0_offset] : buffer -> view<1xi32> + // Publish every producer's residual stores before the leader advances one + // workgroup arrival. The last arrival then acquires the complete row. + kernel.barrier scope(workgroup) ordering(release) + scf.if %is_arrival_workitem { + %old_counter = view.atomic.rmw %c1_i32, %completion_counter_view[%c0] {ordering = acq_rel, scope = device} : i32, view<1xi32> -> i32 + view.store %old_counter, %counter_scratch_view[%c0] : i32, view<1xi32> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + %old_counter = view.load %counter_scratch_view[%c0] : view<1xi32> -> i32 + %output_tile_count_i32 = index.cast %output_tile_count : index to i32 + %last_output_tile_i32 = scalar.subi %output_tile_count_i32, %c1_i32 : i32 + %negative_output_tile_count_i32 = scalar.subi %c0_i32, %output_tile_count_i32 : i32 + %is_last_output_tile = scalar.cmpi eq, %old_counter, %last_output_tile_i32 : i32 + scf.if %is_last_output_tile { + kernel.barrier scope(workgroup) ordering(acquire) + %publish_normalized = scalar.constant true : i1 + func.call @qwen3_moe_rmsnorm_quantize_q8_1_x4_body(%publish_normalized, %c8, %bounded_token_count, %token, %output_noalias, %norm_weight_noalias, %normalized_output_noalias, %next_q8_output_noalias) : (i1, index, index, index, buffer, buffer, buffer, buffer) + // Reset only after every normalized F32 and Q8 store completes. + kernel.barrier scope(workgroup) ordering(release) + scf.if %is_arrival_workitem { + view.atomic.reduce %negative_output_tile_count_i32, %completion_counter_view[%c0] {ordering = release, scope = device} : i32, view<1xi32> + } + } + kernel.return +} + +// Shared matrix-tile schedule. Callers own the launch grid and pass one +// validated tile coordinate plus a literal storage kind, allowing the linker +// and JIT to erase the inactive packed decoder. +func.def inline @qwen3_moe_dense_linear_quantized_f16_wmma_body(%weight_format: index, %token_count: index, %input_size0: index, %output_size0: index, %output_accumulation: index, %channel_tile: index, %token_tile: index, %input: buffer, %weight: buffer, %output: buffer) { + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 2048)] : index + %bounded_input_size = index.assume %input_size0 [range(%input_size0, 256, 32768), mul(%input_size0, 256)] : index + %bounded_output_size = index.assume %output_size0 [range(%output_size0, 1, 262144)] : index + %workitem = kernel.workitem.id : index + %subgroup0 = kernel.subgroup.id : index + %subgroup = index.assume %subgroup0 [range(%subgroup0, 0, 1)] : index + %lane = kernel.subgroup.lane.id : index + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c4 = index.constant 4 : index + %c6 = index.constant 6 : index + %c8 = index.constant 8 : index + %c16 = index.constant 16 : index + %c32 = index.constant 32 : index + %c40 = index.constant 40 : index + %c64 = index.constant 64 : index + %c210_bytes = index.constant 210 : offset + %c256 = index.constant 256 : index + %c0_offset = index.constant 0 : offset + %q4_block_bytes = index.constant 144 : offset + %weight_stage_bytes = index.constant 5120 : offset + %activation_stage_bytes = index.constant 2560 : offset + %wave_result_stage_bytes = index.constant 512 : offset + %result_stage_bytes = index.constant 1024 : offset + %c0_f16x4 = vector.constant 0.0 : vector<4xf16> + %c0_f32x4 = vector.constant 0.0 : vector<4xf32> + %zero_accumulator = vector.constant 0.0 : vector<8xf16> + %m = index.constant 16 : index + %n = index.constant 16 : index + %k = index.constant 16 : index + %is_q6 = index.cmp eq, %weight_format, %c6 : index + %weight_block_bytes = scf.select %is_q6, %c210_bytes, %q4_block_bytes : offset + %quant_block_count = index.div %bounded_input_size, %c256 : index + %weight_row_bytes = index.scale %quant_block_count, %weight_block_bytes : index, offset -> offset + %input_noalias, %weight_noalias, %output_noalias = buffer.assume.noalias %input, %weight, %output : buffer, buffer, buffer + %input_view = buffer.view %input_noalias[%c0_offset] : buffer -> view<[%bounded_token_count]x[%bounded_input_size]xf32> + %output_view = buffer.view %output_noalias[%c0_offset] : buffer -> view<[%bounded_token_count]x[%bounded_output_size]xf32> + %weight_stage = buffer.alloca align(16) %weight_stage_bytes : buffer + %activation_stage = buffer.alloca align(16) %activation_stage_bytes : buffer + %result_stage = buffer.alloca align(16) %result_stage_bytes : buffer + %weight_stage_view = buffer.view %weight_stage[%c0_offset] : buffer -> view<64x40xf16> + %activation_stage_physical_view = buffer.view %activation_stage[%c0_offset] : buffer -> view<32x40xf16> + %activation_fragment_layout = encoding.layout.strided [1, %c40] : encoding + %activation_fragment_view = buffer.view %activation_stage[%c0_offset] : buffer -> view<32x32xf16, %activation_fragment_layout> + %wave_result_stage_offset = index.scale %subgroup, %wave_result_stage_bytes : index, offset -> offset + %result_fragment_layout = encoding.layout.strided [1, %c16] : encoding + %result_fragment_view = buffer.view %result_stage[%wave_result_stage_offset] : buffer -> view<16x16xf16, %result_fragment_layout> + %result_physical_view = buffer.view %result_stage[%wave_result_stage_offset] : buffer -> view<16x16xf16> + %channel_tile_base = index.mul %channel_tile, %c64 : index + %token_tile_base = index.mul %token_tile, %c32 : index + %load_packet = index.rem %workitem, %c8 : index + %load_k = index.mul %load_packet, %c4 : index + %load_row0 = index.div %workitem, %c8 : index + %load_row = index.assume %load_row0 [range(%load_row0, 0, 15)] : index + %subgroup_channel_add = index.mul %subgroup, %c32 : index + %subgroup_channel1 = index.add %subgroup_channel_add, %c16 : index + %init00 = vector.fragment %zero_accumulator shape [%m, %n] : vector<8xf16> + %init01 = vector.fragment %zero_accumulator shape [%m, %n] : vector<8xf16> + %init10 = vector.fragment %zero_accumulator shape [%m, %n] : vector<8xf16> + %init11 = vector.fragment %zero_accumulator shape [%m, %n] : vector<8xf16> + %result00, %result01, %result10, %result11 = scf.for %quant_block = [%c0 to %quant_block_count step %c1](%block_acc00 = %init00 : vector<8xf16>, %block_acc01 = %init01 : vector<8xf16>, %block_acc10 = %init10 : vector<8xf16>, %block_acc11 = %init11 : vector<8xf16>) -> (vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>) { + %block_result00, %block_result01, %block_result10, %block_result11 = scf.for %quant_group = [%c0 to %c8 step %c1](%acc00 = %block_acc00 : vector<8xf16>, %acc01 = %block_acc01 : vector<8xf16>, %acc10 = %block_acc10 : vector<8xf16>, %acc11 = %block_acc11 : vector<8xf16>) -> (vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>) { + %block_k_base = index.mul %quant_block, %c256 : index + %group_k_add = index.mul %quant_group, %c32 : index + %k_origin = index.add %block_k_base, %group_k_add : index + scf.for %row_offset = [%c0 to %c64 step %c16] unroll { + %local_row0 = index.add %load_row, %row_offset : index + %local_row = index.assume %local_row0 [range(%local_row0, 0, 63)] : index + %channel = index.add %channel_tile_base, %local_row : index + %valid_channel = index.cmp ult, %channel, %bounded_output_size : index + %weight_values = scf.if %valid_channel -> (vector<4xf16>) { + %row_byte_base = index.scale %channel, %weight_row_bytes : index, offset -> offset + %decoded = scf.if %is_q6 -> (vector<4xf16>) { + %q6_values = func.call @ggml_q6k_f16_vector4(%weight_noalias, %row_byte_base, %quant_block, %quant_group, %load_packet) : (buffer, offset, index, index, index) -> (vector<4xf16>) + scf.yield %q6_values : vector<4xf16> + } else { + %q4_values = func.call @qwen3_moe_dense_q4k_wmma_vector4(%weight_noalias, %row_byte_base, %quant_block, %quant_group, %load_packet) : (buffer, offset, index, index, index) -> (vector<4xf16>) + scf.yield %q4_values : vector<4xf16> + } + scf.yield %decoded : vector<4xf16> + } else { + scf.yield %c0_f16x4 : vector<4xf16> + } + vector.store %weight_values, %weight_stage_view[%local_row, %load_k] : vector<4xf16>, view<64x40xf16> + %is_activation_row = index.cmp ult, %local_row, %c32 : index + scf.if %is_activation_row { + %activation_row = index.assume %local_row [range(%local_row, 0, 31)] : index + %token = index.add %token_tile_base, %activation_row : index + %valid_token = index.cmp ult, %token, %bounded_token_count : index + %activation_values = scf.if %valid_token -> (vector<4xf16>) { + %bounded_token, %input_token_count = index.assume %token, %bounded_token_count [lt(%token, %bounded_token_count)] : index, index + %input_k = index.add %k_origin, %load_k : index + %loaded = vector.load %input_view[%bounded_token, %input_k] : view<[%bounded_token_count]x[%bounded_input_size]xf32> -> vector<4xf32> + %converted = vector.fptrunc %loaded : vector<4xf32> to vector<4xf16> + scf.yield %converted : vector<4xf16> + } else { + scf.yield %c0_f16x4 : vector<4xf16> + } + vector.store %activation_values, %activation_stage_physical_view[%activation_row, %load_k] : vector<4xf16>, view<32x40xf16> + } + } + kernel.barrier scope(workgroup) ordering(acq_rel) + %next00, %next01, %next10, %next11 = scf.for %k_half = [%c0 to %c32 step %c16](%half_acc00 = %acc00 : vector<8xf16>, %half_acc01 = %acc01 : vector<8xf16>, %half_acc10 = %acc10 : vector<8xf16>, %half_acc11 = %acc11 : vector<8xf16>) -> (vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>) unroll { + %lhs0 = vector.fragment.load %weight_stage_view[%subgroup_channel_add, %k_half] shape [%m, %k] : view<64x40xf16> -> vector<16xf16> + %lhs1 = vector.fragment.load %weight_stage_view[%subgroup_channel1, %k_half] shape [%m, %k] : view<64x40xf16> -> vector<16xf16> + %rhs0 = vector.fragment.load %activation_fragment_view[%k_half, %c0] shape [%k, %n] : view<32x32xf16, %activation_fragment_layout> -> vector<16xf16> + %rhs1 = vector.fragment.load %activation_fragment_view[%k_half, %c16] shape [%k, %n] : view<32x32xf16, %activation_fragment_layout> -> vector<16xf16> + %half_next00 = vector.mma %lhs0, %rhs0, %half_acc00 : vector<16xf16>, vector<16xf16>, vector<8xf16> + %half_next01 = vector.mma %lhs0, %rhs1, %half_acc01 : vector<16xf16>, vector<16xf16>, vector<8xf16> + %half_next10 = vector.mma %lhs1, %rhs0, %half_acc10 : vector<16xf16>, vector<16xf16>, vector<8xf16> + %half_next11 = vector.mma %lhs1, %rhs1, %half_acc11 : vector<16xf16>, vector<16xf16>, vector<8xf16> + scf.yield %half_next00, %half_next01, %half_next10, %half_next11 : vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + scf.yield %next00, %next01, %next10, %next11 : vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16> + } + scf.yield %block_result00, %block_result01, %block_result10, %block_result11 : vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16> + } + // WMMA owns [channel][token] fragments. Each wave transposes one fragment + // through its private LDS slice so every lane publishes four contiguous F32 + // channels for one logical token. + %publish_token0 = index.div %lane, %c4 : index + %publish_token = index.assume %publish_token0 [range(%publish_token0, 0, 15)] : index + %publish_packet0 = index.rem %lane, %c4 : index + %publish_packet = index.assume %publish_packet0 [range(%publish_packet0, 0, 3)] : index + %publish_channel_add = index.mul %publish_packet, %c4 : index + %token0 = index.add %token_tile_base, %publish_token : index + %publish_token1 = index.add %publish_token, %c16 : index + %token1 = index.add %token_tile_base, %publish_token1 : index + %subgroup_channel_base = index.add %channel_tile_base, %subgroup_channel_add : index + %channel0 = index.add %subgroup_channel_base, %publish_channel_add : index + %channel1_base = index.add %subgroup_channel_base, %c16 : index + %channel1 = index.add %channel1_base, %publish_channel_add : index + %valid_token0 = index.cmp ult, %token0, %bounded_token_count : index + %valid_token1 = index.cmp ult, %token1, %bounded_token_count : index + %valid_channel0 = index.cmp ult, %channel0, %bounded_output_size : index + %valid_channel1 = index.cmp ult, %channel1, %bounded_output_size : index + %writes00 = scalar.andi %valid_token0, %valid_channel0 : i1 + %writes01 = scalar.andi %valid_token1, %valid_channel0 : i1 + %writes10 = scalar.andi %valid_token0, %valid_channel1 : i1 + %writes11 = scalar.andi %valid_token1, %valid_channel1 : i1 + %accumulates_output = index.cmp eq, %output_accumulation, %c1 : index + vector.fragment.store %result00, %result_fragment_view[%c0, %c0] shape [%m, %n] : vector<8xf16>, view<16x16xf16, %result_fragment_layout> + kernel.barrier scope(subgroup) ordering(acq_rel) + scf.if %writes00 { + %bounded_token, %output_token_count = index.assume %token0, %bounded_token_count [lt(%token0, %bounded_token_count)] : index, index + %values = vector.load %result_physical_view[%publish_token, %publish_channel_add] : view<16x16xf16> -> vector<4xf16> + %wide = vector.extf %values : vector<4xf16> to vector<4xf32> + %mask = vector.mask.range [%channel0 to %bounded_output_size step %c1] : index -> vector<4xi1> + %published = scf.if %accumulates_output -> (vector<4xf32>) { + %residual = vector.load.mask %output_view[%bounded_token, %channel0], %mask, %c0_f32x4 : view<[%bounded_token_count]x[%bounded_output_size]xf32>, vector<4xi1>, vector<4xf32> + %sum = vector.addf %residual, %wide : vector<4xf32> + scf.yield %sum : vector<4xf32> + } else { + scf.yield %wide : vector<4xf32> + } + vector.store.mask %published, %output_view[%bounded_token, %channel0], %mask : vector<4xf32>, view<[%bounded_token_count]x[%bounded_output_size]xf32>, vector<4xi1> + } + kernel.barrier scope(subgroup) ordering(acq_rel) + vector.fragment.store %result01, %result_fragment_view[%c0, %c0] shape [%m, %n] : vector<8xf16>, view<16x16xf16, %result_fragment_layout> + kernel.barrier scope(subgroup) ordering(acq_rel) + scf.if %writes01 { + %bounded_token, %output_token_count = index.assume %token1, %bounded_token_count [lt(%token1, %bounded_token_count)] : index, index + %values = vector.load %result_physical_view[%publish_token, %publish_channel_add] : view<16x16xf16> -> vector<4xf16> + %wide = vector.extf %values : vector<4xf16> to vector<4xf32> + %mask = vector.mask.range [%channel0 to %bounded_output_size step %c1] : index -> vector<4xi1> + %published = scf.if %accumulates_output -> (vector<4xf32>) { + %residual = vector.load.mask %output_view[%bounded_token, %channel0], %mask, %c0_f32x4 : view<[%bounded_token_count]x[%bounded_output_size]xf32>, vector<4xi1>, vector<4xf32> + %sum = vector.addf %residual, %wide : vector<4xf32> + scf.yield %sum : vector<4xf32> + } else { + scf.yield %wide : vector<4xf32> + } + vector.store.mask %published, %output_view[%bounded_token, %channel0], %mask : vector<4xf32>, view<[%bounded_token_count]x[%bounded_output_size]xf32>, vector<4xi1> + } + kernel.barrier scope(subgroup) ordering(acq_rel) + vector.fragment.store %result10, %result_fragment_view[%c0, %c0] shape [%m, %n] : vector<8xf16>, view<16x16xf16, %result_fragment_layout> + kernel.barrier scope(subgroup) ordering(acq_rel) + scf.if %writes10 { + %bounded_token, %output_token_count = index.assume %token0, %bounded_token_count [lt(%token0, %bounded_token_count)] : index, index + %values = vector.load %result_physical_view[%publish_token, %publish_channel_add] : view<16x16xf16> -> vector<4xf16> + %wide = vector.extf %values : vector<4xf16> to vector<4xf32> + %mask = vector.mask.range [%channel1 to %bounded_output_size step %c1] : index -> vector<4xi1> + %published = scf.if %accumulates_output -> (vector<4xf32>) { + %residual = vector.load.mask %output_view[%bounded_token, %channel1], %mask, %c0_f32x4 : view<[%bounded_token_count]x[%bounded_output_size]xf32>, vector<4xi1>, vector<4xf32> + %sum = vector.addf %residual, %wide : vector<4xf32> + scf.yield %sum : vector<4xf32> + } else { + scf.yield %wide : vector<4xf32> + } + vector.store.mask %published, %output_view[%bounded_token, %channel1], %mask : vector<4xf32>, view<[%bounded_token_count]x[%bounded_output_size]xf32>, vector<4xi1> + } + kernel.barrier scope(subgroup) ordering(acq_rel) + vector.fragment.store %result11, %result_fragment_view[%c0, %c0] shape [%m, %n] : vector<8xf16>, view<16x16xf16, %result_fragment_layout> + kernel.barrier scope(subgroup) ordering(acq_rel) + scf.if %writes11 { + %bounded_token, %output_token_count = index.assume %token1, %bounded_token_count [lt(%token1, %bounded_token_count)] : index, index + %values = vector.load %result_physical_view[%publish_token, %publish_channel_add] : view<16x16xf16> -> vector<4xf16> + %wide = vector.extf %values : vector<4xf16> to vector<4xf32> + %mask = vector.mask.range [%channel1 to %bounded_output_size step %c1] : index -> vector<4xi1> + %published = scf.if %accumulates_output -> (vector<4xf32>) { + %residual = vector.load.mask %output_view[%bounded_token, %channel1], %mask, %c0_f32x4 : view<[%bounded_token_count]x[%bounded_output_size]xf32>, vector<4xi1>, vector<4xf32> + %sum = vector.addf %residual, %wide : vector<4xf32> + scf.yield %sum : vector<4xf32> + } else { + scf.yield %wide : vector<4xf32> + } + vector.store.mask %published, %output_view[%bounded_token, %channel1], %mask : vector<4xf32>, view<[%bounded_token_count]x[%bounded_output_size]xf32>, vector<4xi1> + } + func.return +} + +// Shared launch geometry for configured standalone kernels and parameterized +// command-program kernels. +template.decl @qwen3_moe_dense_quantized_launch(%token_capacity: index, %output_size: index) -> (index, index, index, index) +template.def<@qwen3_moe_dense_quantized_launch> @qwen3_moe_dense_linear_quantized_f16_wmma_launch(%token_capacity: index, %output_size: index) -> (index, index, index, index) { + %c1 = index.constant 1 : index + %c31 = index.constant 31 : index + %c32 = index.constant 32 : index + %c63 = index.constant 63 : index + %c64 = index.constant 64 : index + %c128 = index.constant 128 : index + %padded_output_size = index.add %output_size, %c63 : index + %output_tiles = index.div %padded_output_size, %c64 : index + %padded_token_count = index.add %token_capacity, %c31 : index + %token_tiles = index.div %padded_token_count, %c32 : index + template.return %output_tiles, %token_tiles, %c1, %c128 : index, index, index, index +} + +// Configured entry points retain the compact standalone kernel ABI used by +// library-style callers. Shape configs specialize the native code once while +// token count remains the only per-dispatch scalar. +kernel.def target(@qwen3_moe_dense_gfx11_wave64) @qwen3_moe_dense_linear_q4k_f16_wmma(%token_count: index) { + %output_size = config.get @qwen3_moe.dense_quantized.output_size : index + %token_capacity = config.get @qwen3_moe.workload.token_capacity : index + %output_tiles, %token_tiles, %c1, %c128 = template.apply<@qwen3_moe_dense_quantized_launch>(%token_capacity, %output_size) pure : (index, index) -> (index, index, index, index) + kernel.launch.config workgroups(%output_tiles, %token_tiles, %c1) workgroup_size(%c128, %c1, %c1) : index +} launch(%token_count: index, %input: buffer, %weight: buffer, %output: buffer) { + %input_size = config.get @qwen3_moe.dense_quantized.input_size : index + %output_size = config.get @qwen3_moe.dense_quantized.output_size : index + %token_capacity = config.get @qwen3_moe.workload.token_capacity : index + %output_accumulation = config.get @qwen3_moe.dense_quantized.output_accumulation : index + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 2048), le(%token_count, %token_capacity)] : index + %channel_tile = kernel.workgroup.id : index + %token_tile = kernel.workgroup.id : index + %q4 = index.constant 4 : index + func.call @qwen3_moe_dense_linear_quantized_f16_wmma_body(%q4, %bounded_token_count, %input_size, %output_size, %output_accumulation, %channel_tile, %token_tile, %input, %weight, %output) : (index, index, index, index, index, index, index, buffer, buffer, buffer) + kernel.return +} + +kernel.def target(@qwen3_moe_dense_gfx11_wave64) @qwen3_moe_dense_linear_q6k_f16_wmma(%token_count: index) { + %output_size = config.get @qwen3_moe.dense_quantized.output_size : index + %token_capacity = config.get @qwen3_moe.workload.token_capacity : index + %output_tiles, %token_tiles, %c1, %c128 = template.apply<@qwen3_moe_dense_quantized_launch>(%token_capacity, %output_size) pure : (index, index) -> (index, index, index, index) + kernel.launch.config workgroups(%output_tiles, %token_tiles, %c1) workgroup_size(%c128, %c1, %c1) : index +} launch(%token_count: index, %input: buffer, %weight: buffer, %output: buffer) { + %input_size = config.get @qwen3_moe.dense_quantized.input_size : index + %output_size = config.get @qwen3_moe.dense_quantized.output_size : index + %token_capacity = config.get @qwen3_moe.workload.token_capacity : index + %output_accumulation = config.get @qwen3_moe.dense_quantized.output_accumulation : index + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 2048), le(%token_count, %token_capacity)] : index + %channel_tile = kernel.workgroup.id : index + %token_tile = kernel.workgroup.id : index + %q6 = index.constant 6 : index + func.call @qwen3_moe_dense_linear_quantized_f16_wmma_body(%q6, %bounded_token_count, %input_size, %output_size, %output_accumulation, %channel_tile, %token_tile, %input, %weight, %output) : (index, index, index, index, index, index, index, buffer, buffer, buffer) + kernel.return +} + +// Parameterized entry points carry launch-varying shapes through command IR. +// Command planning materializes each exact fact environment in a private unit, +// so these scalar values do not survive in the native device ABI. +kernel.def target(@qwen3_moe_dense_gfx11_wave64) @qwen3_moe_dense_linear_q4k_f16_wmma_parameterized(%token_count: index, %input_size: index, %output_size: index, %output_accumulation: index) { + %output_tiles, %token_tiles, %c1, %c128 = template.apply<@qwen3_moe_dense_quantized_launch>(%token_count, %output_size) : (index, index) -> (index, index, index, index) + kernel.launch.config workgroups(%output_tiles, %token_tiles, %c1) workgroup_size(%c128, %c1, %c1) : index +} launch(%token_count: index, %input_size: index, %output_size: index, %output_accumulation: index, %input: buffer, %weight: buffer, %output: buffer) where [range(%token_count, 1, 2048), range(%input_size, 256, 32768), mul(%input_size, 256), range(%output_size, 1, 262144), range(%output_accumulation, 0, 1)] { + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 2048)] : index + %channel_tile = kernel.workgroup.id : index + %token_tile = kernel.workgroup.id : index + %q4 = index.constant 4 : index + func.call @qwen3_moe_dense_linear_quantized_f16_wmma_body(%q4, %bounded_token_count, %input_size, %output_size, %output_accumulation, %channel_tile, %token_tile, %input, %weight, %output) : (index, index, index, index, index, index, index, buffer, buffer, buffer) + kernel.return +} + +kernel.def target(@qwen3_moe_dense_gfx11_wave64) @qwen3_moe_dense_linear_q6k_f16_wmma_parameterized(%token_count: index, %input_size: index, %output_size: index, %output_accumulation: index) { + %output_tiles, %token_tiles, %c1, %c128 = template.apply<@qwen3_moe_dense_quantized_launch>(%token_count, %output_size) : (index, index) -> (index, index, index, index) + kernel.launch.config workgroups(%output_tiles, %token_tiles, %c1) workgroup_size(%c128, %c1, %c1) : index +} launch(%token_count: index, %input_size: index, %output_size: index, %output_accumulation: index, %input: buffer, %weight: buffer, %output: buffer) where [range(%token_count, 1, 2048), range(%input_size, 256, 32768), mul(%input_size, 256), range(%output_size, 1, 262144), range(%output_accumulation, 0, 1)] { + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 2048)] : index + %channel_tile = kernel.workgroup.id : index + %token_tile = kernel.workgroup.id : index + %q6 = index.constant 6 : index + func.call @qwen3_moe_dense_linear_quantized_f16_wmma_body(%q6, %bounded_token_count, %input_size, %output_size, %output_accumulation, %channel_tile, %token_tile, %input, %weight, %output) : (index, index, index, index, index, index, index, buffer, buffer, buffer) + kernel.return +} + +// The exact-representable activation and nonzero packed weight bytes compare +// dense addressing against the already-certified routed WMMA provider. The +// shape crosses both the 32-token and 64-channel tile boundaries. +check.case public @qwen3_moe_dense_linear_q4k_f16_wmma_differential_case { + %token_count = check.literal value(33) : index + %route_count = check.literal value(1) : index + %route_stride = check.literal value(1) : index + %expert_count = check.literal value(1) : index + %input = check.generate.fill value(0.00390625) : tensor<33x512xf32> + %route_ids = check.generate.fill value(0) : tensor<33x1xi32> + %expert_table = check.generate.fill value(-1) : tensor<34xi32> + %weight = check.generate.fill value(34) : tensor<1x65x2x144xi8> + %expected = check.generate.fill value(0.0) : tensor<33x65xf32> + %actual = check.generate.fill value(1.0) : tensor<33x65xf32> + kernel.launch @qwen3_moe_build_expert_table[%token_count, %route_count, %route_stride, %expert_count](%token_count, %route_count, %route_stride, %expert_count, %route_ids, %expert_table) : [index, index, index, index](index, index, index, index, tensor<33x1xi32>, tensor<34xi32>) + kernel.launch @qwen3_moe_routed_linear_q4k_f16_wmma[%token_count](%token_count, %input, %expert_table, %weight, %expected) : [index](index, tensor<33x512xf32>, tensor<34xi32>, tensor<1x65x2x144xi8>, tensor<33x65xf32>) + kernel.launch @qwen3_moe_dense_linear_q4k_f16_wmma[%token_count](%token_count, %input, %weight, %actual) : [index](index, tensor<33x512xf32>, tensor<1x65x2x144xi8>, tensor<33x65xf32>) + check.expect.close actual(%actual) expected(%expected) atol(0.01) rtol(0.01) nan(same) : tensor<33x65xf32> + check.return +} + +// Q6_K uses the same exact-representable activation across both the Q8_1 dot +// reference and FP16 WMMA provider. The comparison crosses both matrix tile +// boundaries and validates the format-specialized packed decoder. +check.case public @qwen3_moe_dense_linear_q6k_f16_wmma_differential_case { + %token_count = check.literal value(33) : index + %input_size = check.literal value(512) : index + %output_size = check.literal value(65) : index + %input = check.generate.fill value(0.00390625) : tensor<33x512xf32> + %q8_input = check.generate.fill value(0) : tensor<33x576xi8> + %weight = check.generate.fill value(-86) : tensor<65x2x210xi8> + %expected = check.generate.fill value(0.0) : tensor<33x65xf32> + %actual = check.generate.fill value(1.0) : tensor<33x65xf32> + kernel.launch @ggml_quantize_q8_1_x4_f32[%token_count, %input_size](%token_count, %input_size, %input, %q8_input) : [index, index](index, index, tensor<33x512xf32>, tensor<33x576xi8>) + kernel.launch @ggml_linear_q6k_q8_1_x4[%token_count, %input_size, %output_size](%token_count, %input_size, %output_size, %q8_input, %weight, %expected) : [index, index, index](index, index, index, tensor<33x576xi8>, tensor<65x2x210xi8>, tensor<33x65xf32>) + kernel.launch @qwen3_moe_dense_linear_q6k_f16_wmma[%token_count](%token_count, %input, %weight, %actual) : [index](index, tensor<33x512xf32>, tensor<65x2x210xi8>, tensor<33x65xf32>) + check.expect.close actual(%actual) expected(%expected) atol(0.25) rtol(0.01) nan(same) : tensor<33x65xf32> + check.return +} + +// The O-projection dimensions double the Q/K input depth. This nonzero +// production-size differential keeps that loop regime covered independently +// of the benchmark parameter sweep. +check.case public @qwen3_moe_dense_linear_q4k_f16_wmma_o_differential_case { + %token_count = check.literal value(1) : index + %input_size = check.literal value(4096) : index + %route_count = check.literal value(1) : index + %route_stride = check.literal value(1) : index + %expert_count = check.literal value(1) : index + %input = check.generate.fill value(0.00390625) : tensor<1x4096xf32> + %q8_input = check.generate.fill value(0) : tensor<1x4608xi8> + %route_ids = check.generate.fill value(0) : tensor<1x1xi32> + %expert_table = check.generate.fill value(-1) : tensor<2xi32> + %weight = check.generate.fill value(34) : tensor<1x2048x16x144xi8> + %expected = check.generate.fill value(0.0) : tensor<1x2048xf32> + %actual = check.generate.fill value(0.0) : tensor<1x2048xf32> + %direct = check.generate.fill value(0.0) : tensor<1x2048xf32> + kernel.launch @qwen3_moe_build_expert_table[%token_count, %route_count, %route_stride, %expert_count](%token_count, %route_count, %route_stride, %expert_count, %route_ids, %expert_table) : [index, index, index, index](index, index, index, index, tensor<1x1xi32>, tensor<2xi32>) + kernel.launch @ggml_quantize_q8_1_x4_f32[%token_count, %input_size](%token_count, %input_size, %input, %q8_input) : [index, index](index, index, tensor<1x4096xf32>, tensor<1x4608xi8>) + kernel.launch @qwen3_moe_routed_linear_q4k_f16_wmma[%token_count](%token_count, %input, %expert_table, %weight, %expected) : [index](index, tensor<1x4096xf32>, tensor<2xi32>, tensor<1x2048x16x144xi8>, tensor<1x2048xf32>) + kernel.launch @qwen3_moe_dense_linear_q4k_f16_wmma[%token_count](%token_count, %input, %weight, %actual) : [index](index, tensor<1x4096xf32>, tensor<1x2048x16x144xi8>, tensor<1x2048xf32>) + kernel.launch @qwen3_moe_dense_linear_q4k_q8_1_x4[%token_count](%token_count, %q8_input, %weight, %direct) : [index](index, tensor<1x4608xi8>, tensor<1x2048x16x144xi8>, tensor<1x2048xf32>) + check.expect.close actual(%actual) expected(%expected) atol(0.01) rtol(0.01) nan(same) : tensor<1x2048xf32> + check.expect.close actual(%direct) expected(%expected) atol(0.25) rtol(0.01) nan(same) : tensor<1x2048xf32> + check.return +} + +// Crosses the direct schedule's eight-wave output tile while comparing its +// Q8_1 dot path against the independently structured FP16 WMMA provider. +check.case public @qwen3_moe_dense_linear_q4k_q8_1_x4_differential_case { + %token_count = check.literal value(2) : index + %input_size = check.literal value(512) : index + %input = check.generate.fill value(0.00390625) : tensor<2x512xf32> + %q8_input = check.generate.fill value(0) : tensor<2x576xi8> + %weight = check.generate.iota offset(-72) step(1) period(144) : tensor<65x2x144xi8> + %expected = check.generate.fill value(0.0) : tensor<2x65xf32> + %actual = check.generate.fill value(0.0) : tensor<2x65xf32> + kernel.launch @ggml_quantize_q8_1_x4_f32[%token_count, %input_size](%token_count, %input_size, %input, %q8_input) : [index, index](index, index, tensor<2x512xf32>, tensor<2x576xi8>) + kernel.launch @qwen3_moe_dense_linear_q4k_f16_wmma[%token_count](%token_count, %input, %weight, %expected) : [index](index, tensor<2x512xf32>, tensor<65x2x144xi8>, tensor<2x65xf32>) + kernel.launch @qwen3_moe_dense_linear_q4k_q8_1_x4[%token_count](%token_count, %q8_input, %weight, %actual) : [index](index, tensor<2x576xi8>, tensor<65x2x144xi8>, tensor<2x65xf32>) + check.expect.close actual(%actual) expected(%expected) atol(0.25) rtol(0.01) nan(same) : tensor<2x65xf32> + check.return +} + +// A zero projection leaves the nonzero caller-provided residual unchanged. +// Overwrite publication would instead produce zero, making the accumulation +// contract observable without coupling this case to another projection kernel. +check.case public @qwen3_moe_dense_linear_q4k_q8_1_x4_accumulation_case { + %token_count = check.literal value(2) : index + %input_size = check.literal value(512) : index + %input = check.generate.fill value(0.0) : tensor<2x512xf32> + %q8_input = check.generate.fill value(1) : tensor<2x576xi8> + %weight = check.generate.fill value(34) : tensor<65x2x144xi8> + %actual = check.generate.fill value(1.25) : tensor<2x65xf32> + %expected = check.generate.fill value(1.25) : tensor<2x65xf32> + kernel.launch @ggml_quantize_q8_1_x4_f32[%token_count, %input_size](%token_count, %input_size, %input, %q8_input) : [index, index](index, index, tensor<2x512xf32>, tensor<2x576xi8>) + kernel.launch @qwen3_moe_dense_linear_q4k_q8_1_x4[%token_count](%token_count, %q8_input, %weight, %actual) : [index](index, tensor<2x576xi8>, tensor<65x2x144xi8>, tensor<2x65xf32>) + check.expect.close actual(%actual) expected(%expected) atol(0.0) rtol(0.0) nan(same) : tensor<2x65xf32> + check.return +} + +// The production output geometry requires all 256 tiles to publish the +// accumulated hidden row before the final workgroup can normalize and pack it. +// Reuse each oversized attention-input Q8 allocation in place, invoke the fused +// kernel twice through one completion word, and compare every semantic output +// with the ordinary two-dispatch composition. +check.case public @qwen3_moe_dense_linear_q4k_q8_1_x4_next_q8_differential_case { + %token_count = check.literal value(1) : index + %input_size = check.literal value(4096) : index + %input = check.generate.iota offset(-0.5) step(0.000244140625) period(4096) : tensor<1x4096xf32> + %expected_q8 = check.generate.fill value(0) : tensor<4608xi8> + %actual_q8_0 = check.generate.fill value(0) : tensor<4608xi8> + %actual_q8_1 = check.generate.fill value(0) : tensor<4608xi8> + %weight = check.generate.fill value(34) : tensor<2048x16x144xi8> + %norm_weight = check.generate.iota offset(-1.0) step(0.0009765625) : tensor<2048xf32> + %expected_output = check.generate.iota offset(-0.25) step(0.000244140625) period(2048) : tensor<1x2048xf32> + %actual_output_0 = check.generate.iota offset(-0.25) step(0.000244140625) period(2048) : tensor<1x2048xf32> + %actual_output_1 = check.generate.iota offset(-0.25) step(0.000244140625) period(2048) : tensor<1x2048xf32> + %expected_normalized = check.generate.fill value(1.0) : tensor<1x2048xf32> + %actual_normalized_0 = check.generate.fill value(1.0) : tensor<1x2048xf32> + %actual_normalized_1 = check.generate.fill value(1.0) : tensor<1x2048xf32> + %completion_counter = check.generate.fill value(0) : tensor<1xi32> + %expected_counter = check.generate.fill value(0) : tensor<1xi32> + kernel.launch @ggml_quantize_q8_1_x4_f32[%token_count, %input_size](%token_count, %input_size, %input, %expected_q8) : [index, index](index, index, tensor<1x4096xf32>, tensor<4608xi8>) + kernel.launch @ggml_quantize_q8_1_x4_f32[%token_count, %input_size](%token_count, %input_size, %input, %actual_q8_0) : [index, index](index, index, tensor<1x4096xf32>, tensor<4608xi8>) + kernel.launch @ggml_quantize_q8_1_x4_f32[%token_count, %input_size](%token_count, %input_size, %input, %actual_q8_1) : [index, index](index, index, tensor<1x4096xf32>, tensor<4608xi8>) + kernel.launch @qwen3_moe_dense_linear_q4k_q8_1_x4[%token_count](%token_count, %expected_q8, %weight, %expected_output) : [index](index, tensor<4608xi8>, tensor<2048x16x144xi8>, tensor<1x2048xf32>) + kernel.launch @qwen3_moe_rmsnorm_f32_quantize_q8_1_x4[%token_count](%token_count, %expected_output, %norm_weight, %expected_normalized, %expected_q8) : [index](index, tensor<1x2048xf32>, tensor<2048xf32>, tensor<1x2048xf32>, tensor<4608xi8>) + kernel.launch @qwen3_moe_dense_linear_q4k_q8_1_x4_next_q8[%token_count](%token_count, %actual_q8_0, %weight, %actual_output_0, %norm_weight, %actual_normalized_0, %completion_counter, %actual_q8_0) : [index](index, tensor<4608xi8>, tensor<2048x16x144xi8>, tensor<1x2048xf32>, tensor<2048xf32>, tensor<1x2048xf32>, tensor<1xi32>, tensor<4608xi8>) + kernel.launch @qwen3_moe_dense_linear_q4k_q8_1_x4_next_q8[%token_count](%token_count, %actual_q8_1, %weight, %actual_output_1, %norm_weight, %actual_normalized_1, %completion_counter, %actual_q8_1) : [index](index, tensor<4608xi8>, tensor<2048x16x144xi8>, tensor<1x2048xf32>, tensor<2048xf32>, tensor<1x2048xf32>, tensor<1xi32>, tensor<4608xi8>) + check.expect.close actual(%actual_output_0) expected(%expected_output) atol(0.0) rtol(0.0) nan(same) : tensor<1x2048xf32> + check.expect.close actual(%actual_output_1) expected(%expected_output) atol(0.0) rtol(0.0) nan(same) : tensor<1x2048xf32> + check.expect.close actual(%actual_normalized_0) expected(%expected_normalized) atol(0.0) rtol(0.0) nan(same) : tensor<1x2048xf32> + check.expect.close actual(%actual_normalized_1) expected(%expected_normalized) atol(0.0) rtol(0.0) nan(same) : tensor<1x2048xf32> + check.expect.equal actual(%actual_q8_0) expected(%expected_q8) : tensor<4608xi8> + check.expect.equal actual(%actual_q8_1) expected(%expected_q8) : tensor<4608xi8> + check.expect.equal actual(%completion_counter) expected(%expected_counter) : tensor<1xi32> + check.return +} + +check.case public @qwen3_moe_dense_linear_q4k_q8_1_x4_next_q8_benchmark_case { + %token_count = check.literal value(1) : index + %q8_input_and_output = check.generate.fill value(0) : tensor<4608xi8> + %weight = check.generate.fill value(0) : tensor<2048x16x144xi8> + %output = check.generate.fill value(1.0) : tensor<1x2048xf32> + %norm_weight = check.generate.fill value(1.0) : tensor<2048xf32> + %normalized_output = check.generate.fill value(0.0) : tensor<1x2048xf32> + %completion_counter = check.generate.fill value(0) : tensor<1xi32> + kernel.launch @qwen3_moe_dense_linear_q4k_q8_1_x4_next_q8[%token_count](%token_count, %q8_input_and_output, %weight, %output, %norm_weight, %normalized_output, %completion_counter, %q8_input_and_output) : [index](index, tensor<4608xi8>, tensor<2048x16x144xi8>, tensor<1x2048xf32>, tensor<2048xf32>, tensor<1x2048xf32>, tensor<1xi32>, tensor<4608xi8>) + check.return +} + +check.case public @qwen3_moe_dense_linear_q4k_q8_1_x4_next_q8_composed_benchmark_case { + %token_count = check.literal value(1) : index + %q8_input_and_output = check.generate.fill value(0) : tensor<4608xi8> + %weight = check.generate.fill value(0) : tensor<2048x16x144xi8> + %output = check.generate.fill value(1.0) : tensor<1x2048xf32> + %norm_weight = check.generate.fill value(1.0) : tensor<2048xf32> + %normalized_output = check.generate.fill value(0.0) : tensor<1x2048xf32> + kernel.launch @qwen3_moe_dense_linear_q4k_q8_1_x4[%token_count](%token_count, %q8_input_and_output, %weight, %output) : [index](index, tensor<4608xi8>, tensor<2048x16x144xi8>, tensor<1x2048xf32>) + kernel.launch @qwen3_moe_rmsnorm_f32_quantize_q8_1_x4[%token_count](%token_count, %output, %norm_weight, %normalized_output, %q8_input_and_output) : [index](index, tensor<1x2048xf32>, tensor<2048xf32>, tensor<1x2048xf32>, tensor<4608xi8>) + check.return +} + +check.case public @qwen3_moe_dense_linear_q4k_f16_wmma_q_projection_benchmark_case { + %token_count = check.param.choice values([1, 17, 32, 63, 128, 129, 512, 1024, 2048]) name("token_count") : index + %input = check.generate.fill value(0.0) : tensor<[%token_count]x2048xf32> + %weight = check.generate.fill value(0) : tensor<4096x8x144xi8> + %output = check.generate.fill value(1.0) : tensor<[%token_count]x4096xf32> + %expected = check.generate.fill value(0.0) : tensor<[%token_count]x4096xf32> + kernel.launch @qwen3_moe_dense_linear_q4k_f16_wmma[%token_count](%token_count, %input, %weight, %output) : [index](index, tensor<[%token_count]x2048xf32>, tensor<4096x8x144xi8>, tensor<[%token_count]x4096xf32>) + check.expect.close actual(%output) expected(%expected) atol(0.0) rtol(0.0) nan(same) : tensor<[%token_count]x4096xf32> + check.return +} + +check.case public @qwen3_moe_dense_linear_q4k_q8_1_x4_q_projection_benchmark_case { + %token_count = check.param.choice values([1, 17, 32, 63, 128, 129, 512, 1024, 2048]) name("token_count") : index + %input_size = check.literal value(2048) : index + %input = check.generate.fill value(0.0) : tensor<[%token_count]x2048xf32> + %q8_input = check.generate.fill value(1) : tensor<[%token_count]x2304xi8> + %weight = check.generate.fill value(0) : tensor<4096x8x144xi8> + %output = check.generate.fill value(1.0) : tensor<[%token_count]x4096xf32> + %expected = check.generate.fill value(0.0) : tensor<[%token_count]x4096xf32> + kernel.launch @ggml_quantize_q8_1_x4_f32[%token_count, %input_size](%token_count, %input_size, %input, %q8_input) : [index, index](index, index, tensor<[%token_count]x2048xf32>, tensor<[%token_count]x2304xi8>) + kernel.launch @qwen3_moe_dense_linear_q4k_q8_1_x4[%token_count](%token_count, %q8_input, %weight, %output) : [index](index, tensor<[%token_count]x2304xi8>, tensor<4096x8x144xi8>, tensor<[%token_count]x4096xf32>) + check.expect.close actual(%output) expected(%expected) atol(0.0) rtol(0.0) nan(same) : tensor<[%token_count]x4096xf32> + check.return +} + +check.case public @qwen3_moe_dense_linear_q4k_f16_wmma_k_projection_benchmark_case { + %token_count = check.param.choice values([1, 17, 32, 63, 128, 129, 512, 1024, 2048]) name("token_count") : index + %input = check.generate.fill value(0.0) : tensor<[%token_count]x2048xf32> + %weight = check.generate.fill value(0) : tensor<512x8x144xi8> + %output = check.generate.fill value(1.0) : tensor<[%token_count]x512xf32> + %expected = check.generate.fill value(0.0) : tensor<[%token_count]x512xf32> + kernel.launch @qwen3_moe_dense_linear_q4k_f16_wmma[%token_count](%token_count, %input, %weight, %output) : [index](index, tensor<[%token_count]x2048xf32>, tensor<512x8x144xi8>, tensor<[%token_count]x512xf32>) + check.expect.close actual(%output) expected(%expected) atol(0.0) rtol(0.0) nan(same) : tensor<[%token_count]x512xf32> + check.return +} + +check.case public @qwen3_moe_dense_linear_q4k_q8_1_x4_k_projection_benchmark_case { + %token_count = check.param.choice values([1, 17, 32, 63, 128, 129, 512, 1024, 2048]) name("token_count") : index + %input_size = check.literal value(2048) : index + %input = check.generate.fill value(0.0) : tensor<[%token_count]x2048xf32> + %q8_input = check.generate.fill value(1) : tensor<[%token_count]x2304xi8> + %weight = check.generate.fill value(0) : tensor<512x8x144xi8> + %output = check.generate.fill value(1.0) : tensor<[%token_count]x512xf32> + %expected = check.generate.fill value(0.0) : tensor<[%token_count]x512xf32> + kernel.launch @ggml_quantize_q8_1_x4_f32[%token_count, %input_size](%token_count, %input_size, %input, %q8_input) : [index, index](index, index, tensor<[%token_count]x2048xf32>, tensor<[%token_count]x2304xi8>) + kernel.launch @qwen3_moe_dense_linear_q4k_q8_1_x4[%token_count](%token_count, %q8_input, %weight, %output) : [index](index, tensor<[%token_count]x2304xi8>, tensor<512x8x144xi8>, tensor<[%token_count]x512xf32>) + check.expect.close actual(%output) expected(%expected) atol(0.0) rtol(0.0) nan(same) : tensor<[%token_count]x512xf32> + check.return +} + +check.case public @qwen3_moe_dense_linear_q6k_f16_wmma_v_projection_benchmark_case { + %token_count = check.param.choice values([1, 17, 32, 63, 128, 129, 512, 1024, 2048]) name("token_count") : index + %input = check.generate.fill value(0.0) : tensor<[%token_count]x2048xf32> + %weight = check.generate.fill value(0) : tensor<512x8x210xi8> + %output = check.generate.fill value(1.0) : tensor<[%token_count]x512xf32> + %expected = check.generate.fill value(0.0) : tensor<[%token_count]x512xf32> + kernel.launch @qwen3_moe_dense_linear_q6k_f16_wmma[%token_count](%token_count, %input, %weight, %output) : [index](index, tensor<[%token_count]x2048xf32>, tensor<512x8x210xi8>, tensor<[%token_count]x512xf32>) + check.expect.close actual(%output) expected(%expected) atol(0.0) rtol(0.0) nan(same) : tensor<[%token_count]x512xf32> + check.return +} + +check.case public @qwen3_moe_dense_linear_q4k_f16_wmma_o_projection_benchmark_case { + %token_count = check.param.choice values([1, 17, 32, 63, 128, 129, 512, 1024, 2048]) name("token_count") : index + %input = check.generate.fill value(0.0) : tensor<[%token_count]x4096xf32> + %weight = check.generate.fill value(0) : tensor<2048x16x144xi8> + %output = check.generate.fill value(1.0) : tensor<[%token_count]x2048xf32> + %expected = check.generate.fill value(1.0) : tensor<[%token_count]x2048xf32> + kernel.launch @qwen3_moe_dense_linear_q4k_f16_wmma[%token_count](%token_count, %input, %weight, %output) : [index](index, tensor<[%token_count]x4096xf32>, tensor<2048x16x144xi8>, tensor<[%token_count]x2048xf32>) + check.expect.close actual(%output) expected(%expected) atol(0.0) rtol(0.0) nan(same) : tensor<[%token_count]x2048xf32> + check.return +} + +check.case public @qwen3_moe_dense_linear_q4k_q8_1_x4_o_projection_benchmark_case { + %token_count = check.param.choice values([1, 17, 32, 63, 128, 129, 512, 1024, 2048]) name("token_count") : index + %input_size = check.literal value(4096) : index + %input = check.generate.fill value(0.0) : tensor<[%token_count]x4096xf32> + %q8_input = check.generate.fill value(1) : tensor<[%token_count]x4608xi8> + %weight = check.generate.fill value(0) : tensor<2048x16x144xi8> + %output = check.generate.fill value(1.0) : tensor<[%token_count]x2048xf32> + %expected = check.generate.fill value(1.0) : tensor<[%token_count]x2048xf32> + kernel.launch @ggml_quantize_q8_1_x4_f32[%token_count, %input_size](%token_count, %input_size, %input, %q8_input) : [index, index](index, index, tensor<[%token_count]x4096xf32>, tensor<[%token_count]x4608xi8>) + kernel.launch @qwen3_moe_dense_linear_q4k_q8_1_x4[%token_count](%token_count, %q8_input, %weight, %output) : [index](index, tensor<[%token_count]x4608xi8>, tensor<2048x16x144xi8>, tensor<[%token_count]x2048xf32>) + check.expect.close actual(%output) expected(%expected) atol(0.0) rtol(0.0) nan(same) : tensor<[%token_count]x2048xf32> + check.return +} + +check.benchmark<@qwen3_moe_dense_linear_q4k_f16_wmma_differential_case> @qwen3_moe_dense_linear_q4k_f16_wmma_differential + +check.benchmark<@qwen3_moe_dense_linear_q6k_f16_wmma_differential_case> @qwen3_moe_dense_linear_q6k_f16_wmma_differential + +check.benchmark<@qwen3_moe_dense_linear_q4k_f16_wmma_o_differential_case> @qwen3_moe_dense_linear_q4k_f16_wmma_o_differential + +check.benchmark<@qwen3_moe_dense_linear_q4k_q8_1_x4_differential_case> @qwen3_moe_dense_linear_q4k_q8_1_x4_differential + +check.benchmark<@qwen3_moe_dense_linear_q4k_q8_1_x4_next_q8_benchmark_case> @qwen3_moe_dense_linear_q4k_q8_1_x4_next_q8_decode + +check.benchmark<@qwen3_moe_dense_linear_q4k_q8_1_x4_next_q8_composed_benchmark_case> @qwen3_moe_dense_linear_q4k_q8_1_x4_next_q8_composed_decode + +check.benchmark<@qwen3_moe_dense_linear_q4k_f16_wmma_q_projection_benchmark_case> @qwen3_moe_dense_linear_q4k_f16_wmma_q_decode {token_count = 1} + +check.benchmark<@qwen3_moe_dense_linear_q4k_f16_wmma_q_projection_benchmark_case> @qwen3_moe_dense_linear_q4k_f16_wmma_q_prefill_32 {token_count = 32} + +check.benchmark<@qwen3_moe_dense_linear_q4k_f16_wmma_q_projection_benchmark_case> @qwen3_moe_dense_linear_q4k_f16_wmma_q_prefill_128 {token_count = 128} + +check.benchmark<@qwen3_moe_dense_linear_q4k_f16_wmma_q_projection_benchmark_case> @qwen3_moe_dense_linear_q4k_f16_wmma_q_prefill_512 {token_count = 512} + +check.benchmark<@qwen3_moe_dense_linear_q4k_q8_1_x4_q_projection_benchmark_case> @qwen3_moe_dense_linear_q4k_q8_1_x4_q_decode {token_count = 1} + +check.benchmark<@qwen3_moe_dense_linear_q4k_q8_1_x4_q_projection_benchmark_case> @qwen3_moe_dense_linear_q4k_q8_1_x4_q_prefill_32 {token_count = 32} + +check.benchmark<@qwen3_moe_dense_linear_q4k_q8_1_x4_q_projection_benchmark_case> @qwen3_moe_dense_linear_q4k_q8_1_x4_q_prefill_128 {token_count = 128} + +check.benchmark<@qwen3_moe_dense_linear_q4k_q8_1_x4_q_projection_benchmark_case> @qwen3_moe_dense_linear_q4k_q8_1_x4_q_prefill_512 {token_count = 512} + +check.benchmark<@qwen3_moe_dense_linear_q4k_f16_wmma_k_projection_benchmark_case> @qwen3_moe_dense_linear_q4k_f16_wmma_k_decode {token_count = 1} + +check.benchmark<@qwen3_moe_dense_linear_q4k_f16_wmma_k_projection_benchmark_case> @qwen3_moe_dense_linear_q4k_f16_wmma_k_prefill_32 {token_count = 32} + +check.benchmark<@qwen3_moe_dense_linear_q4k_f16_wmma_k_projection_benchmark_case> @qwen3_moe_dense_linear_q4k_f16_wmma_k_prefill_128 {token_count = 128} + +check.benchmark<@qwen3_moe_dense_linear_q4k_f16_wmma_k_projection_benchmark_case> @qwen3_moe_dense_linear_q4k_f16_wmma_k_prefill_512 {token_count = 512} + +check.benchmark<@qwen3_moe_dense_linear_q4k_q8_1_x4_k_projection_benchmark_case> @qwen3_moe_dense_linear_q4k_q8_1_x4_k_decode {token_count = 1} + +check.benchmark<@qwen3_moe_dense_linear_q4k_q8_1_x4_k_projection_benchmark_case> @qwen3_moe_dense_linear_q4k_q8_1_x4_k_prefill_32 {token_count = 32} + +check.benchmark<@qwen3_moe_dense_linear_q4k_q8_1_x4_k_projection_benchmark_case> @qwen3_moe_dense_linear_q4k_q8_1_x4_k_prefill_128 {token_count = 128} + +check.benchmark<@qwen3_moe_dense_linear_q4k_q8_1_x4_k_projection_benchmark_case> @qwen3_moe_dense_linear_q4k_q8_1_x4_k_prefill_512 {token_count = 512} + +check.benchmark<@qwen3_moe_dense_linear_q6k_f16_wmma_v_projection_benchmark_case> @qwen3_moe_dense_linear_q6k_f16_wmma_v_decode {token_count = 1} + +check.benchmark<@qwen3_moe_dense_linear_q6k_f16_wmma_v_projection_benchmark_case> @qwen3_moe_dense_linear_q6k_f16_wmma_v_prefill_32 {token_count = 32} + +check.benchmark<@qwen3_moe_dense_linear_q6k_f16_wmma_v_projection_benchmark_case> @qwen3_moe_dense_linear_q6k_f16_wmma_v_prefill_128 {token_count = 128} + +check.benchmark<@qwen3_moe_dense_linear_q6k_f16_wmma_v_projection_benchmark_case> @qwen3_moe_dense_linear_q6k_f16_wmma_v_prefill_512 {token_count = 512} + +check.benchmark<@qwen3_moe_dense_linear_q4k_f16_wmma_o_projection_benchmark_case> @qwen3_moe_dense_linear_q4k_f16_wmma_o_decode {token_count = 1} + +check.benchmark<@qwen3_moe_dense_linear_q4k_f16_wmma_o_projection_benchmark_case> @qwen3_moe_dense_linear_q4k_f16_wmma_o_prefill_32 {token_count = 32} + +check.benchmark<@qwen3_moe_dense_linear_q4k_f16_wmma_o_projection_benchmark_case> @qwen3_moe_dense_linear_q4k_f16_wmma_o_prefill_128 {token_count = 128} + +check.benchmark<@qwen3_moe_dense_linear_q4k_f16_wmma_o_projection_benchmark_case> @qwen3_moe_dense_linear_q4k_f16_wmma_o_prefill_512 {token_count = 512} + +check.benchmark<@qwen3_moe_dense_linear_q4k_q8_1_x4_o_projection_benchmark_case> @qwen3_moe_dense_linear_q4k_q8_1_x4_o_decode {token_count = 1} + +check.benchmark<@qwen3_moe_dense_linear_q4k_q8_1_x4_o_projection_benchmark_case> @qwen3_moe_dense_linear_q4k_q8_1_x4_o_prefill_32 {token_count = 32} + +check.benchmark<@qwen3_moe_dense_linear_q4k_q8_1_x4_o_projection_benchmark_case> @qwen3_moe_dense_linear_q4k_q8_1_x4_o_prefill_128 {token_count = 128} + +check.benchmark<@qwen3_moe_dense_linear_q4k_q8_1_x4_o_projection_benchmark_case> @qwen3_moe_dense_linear_q4k_q8_1_x4_o_prefill_512 {token_count = 512} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/expert_table_partition_fused.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/expert_table_partition_fused.loom new file mode 100644 index 000000000000..8c912bbfc954 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/expert_table_partition_fused.loom @@ -0,0 +1,229 @@ +// Fuses the Prefill-512 expert assignment table and partition descriptor +// construction used by grouped routed projections. +// +// Each of 128 workgroups retains ownership of one expert row. After publishing +// its assignment count, lane zero releases a device-scope completion arrival. +// The last workgroup acquires all counts, compacts them into deterministic +// 32-row descriptors, and resets the counter so reusable command buffers can +// issue the kernel again against the same storage. +amdgpu.target @qwen3_moe_expert_table_partition_gfx11_wave32 {subgroup_size = 32} + +kernel.decl @qwen3_moe_build_expert_table(%token_count: index, %route_count: index, %route_stride: index, %expert_count: index) launch(%token_count: index, %route_count: index, %route_stride: index, %expert_count: index, %route_ids: buffer, %expert_table: buffer) + +kernel.decl @qwen3_moe_build_expert_partition_table(%token_count: index, %route_count: index, %expert_count: index) launch(%token_count: index, %route_count: index, %expert_count: index, %expert_table: buffer, %partition_table: buffer) + +kernel.def target(@qwen3_moe_expert_table_partition_gfx11_wave32) @qwen3_moe_build_expert_table_partition_prefill_512(%token_count: index, %route_count: index, %route_stride: index, %expert_count: index) { + %c1 = index.constant 1 : index + %launch_expert_count = index.constant 128 : index + %workgroup_size = index.constant 256 : index + kernel.launch.config workgroups(%launch_expert_count, %c1, %c1) workgroup_size(%workgroup_size, %c1, %c1) : index +} launch(%token_count: index, %route_count: index, %route_stride: index, %expert_count: index, %route_ids: buffer, %expert_table: buffer, %partition_table: buffer, %completion_counter: buffer) { + %bounded_token_count = index.assume %token_count [range(%token_count, 512, 512)] : index + %bounded_route_count = index.assume %route_count [range(%route_count, 8, 8)] : index + %bounded_route_stride = index.assume %route_stride [range(%route_stride, 8, 8)] : index + %bounded_expert_count = index.assume %expert_count [range(%expert_count, 128, 128)] : index + %expert0 = kernel.workgroup.id : index + %expert, %launch_expert_count = index.assume %expert0, %bounded_expert_count [lt(%expert0, %bounded_expert_count)] : index, index + %lane0 = kernel.workitem.id : index + %lane = index.assume %lane0 [range(%lane0, 0, 255)] : index + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c31 = index.constant 31 : index + %c32 = index.constant 32 : index + %workgroup_size = index.constant 256 : index + %c0_i32 = scalar.constant 0 : i32 + %c1_i32 = scalar.constant 1 : i32 + %c7_i32 = scalar.constant 7 : i32 + %c13_i32 = scalar.constant 13 : i32 + %c0_offset = index.constant 0 : offset + %c4_bytes = index.constant 4 : offset + %assignment_count = index.mul %bounded_token_count, %bounded_route_count : index + %assignment_table_byte_base = index.scale %launch_expert_count, %c4_bytes : index, offset -> offset + %route_ids_noalias, %expert_table_noalias, %partition_table_noalias, %completion_counter_noalias = buffer.assume.noalias %route_ids, %expert_table, %partition_table, %completion_counter : buffer, buffer, buffer, buffer + %completion_counter_aligned = buffer.assume.alignment %completion_counter_noalias {minimum_alignment = 16} : buffer + %route_view = buffer.view %route_ids_noalias[%c0_offset] : buffer -> view<[%bounded_token_count]x[%bounded_route_stride]xi32> + %count_view = buffer.view %expert_table_noalias[%c0_offset] : buffer -> view<[%launch_expert_count]xi32> + %assignment_view = buffer.view %expert_table_noalias[%assignment_table_byte_base] : buffer -> view<[%launch_expert_count]x[%bounded_token_count]xi32> + %completion_counter_view = buffer.view %completion_counter_aligned[%c0_offset] : buffer -> view<1xi32> + + %expert_route_count = scf.for %block_base = [%c0 to %assignment_count step %workgroup_size](%matched_base = %c0_i32 : i32) -> (i32) { + %assignment = index.add %block_base, %lane : index + %in_range = index.cmp ult, %assignment, %assignment_count : index + %route_expert_i32 = scf.if %in_range -> (i32) { + %token0 = index.div %assignment, %bounded_route_count : index + %route0 = index.rem %assignment, %bounded_route_count : index + %token, %route_token_count = index.assume %token0, %bounded_token_count [lt(%token0, %bounded_token_count)] : index, index + %route, %route_row_stride = index.assume %route0, %bounded_route_stride [lt(%route0, %bounded_route_stride)] : index, index + %loaded = view.load %route_view[%token, %route] : view<[%bounded_token_count]x[%bounded_route_stride]xi32> -> i32 + scf.yield %loaded : i32 + } else { + %cn1_i32 = scalar.constant -1 : i32 + scf.yield %cn1_i32 : i32 + } + %route_expert0 = index.cast %route_expert_i32 : i32 to index + %route_expert = index.assume %route_expert0 [range(%route_expert0, -1, 127)] : index + %matches = index.cmp eq, %route_expert, %expert : index + %match_i32 = scf.if %matches -> (i32) { + scf.yield %c1_i32 : i32 + } else { + scf.yield %c0_i32 : i32 + } + %block_prefix = kernel.workgroup.scan %match_i32 {direction = forward, mode = exclusive} : i32 + %block_match_count_reduced = kernel.workgroup.reduce %match_i32 : i32 + %block_match_count = kernel.subgroup.broadcast.first %block_match_count_reduced : i32 + scf.if %matches { + %match_ordinal_i32 = scalar.addi %matched_base, %block_prefix : i32 + %match_ordinal0 = index.cast %match_ordinal_i32 : i32 to index + %match_ordinal = index.assume %match_ordinal0 [range(%match_ordinal0, 0, 511)] : index + %bounded_match_ordinal, %table_token_count = index.assume %match_ordinal, %bounded_token_count [lt(%match_ordinal, %bounded_token_count)] : index, index + %assignment_i32 = index.cast %assignment : index to i32 + view.store %assignment_i32, %assignment_view[%expert, %bounded_match_ordinal] : i32, view<[%launch_expert_count]x[%bounded_token_count]xi32> + } + %next_matched_base = scalar.addi %matched_base, %block_match_count : i32 + scf.yield %next_matched_base : i32 + } + + %is_lane_zero = index.cmp eq, %lane, %c0 : index + scf.if %is_lane_zero { + view.store %expert_route_count, %count_view[%expert] : i32, view<[%launch_expert_count]xi32> + } + + // Every count store precedes its workgroup's release. The last arrival + // acquires all preceding releases before any lane loads the complete table. + %local_old_counter = scf.if %is_lane_zero -> (i32) { + %old_counter = view.atomic.rmw %c1_i32, %completion_counter_view[%c0] {ordering = acq_rel, scope = device} : i32, view<1xi32> -> i32 + scf.yield %old_counter : i32 + } else { + scf.yield %c0_i32 : i32 + } + %old_counter = kernel.workgroup.reduce %local_old_counter : i32 + %expert_count_i32 = index.cast %launch_expert_count : index to i32 + %last_expert_i32 = scalar.subi %expert_count_i32, %c1_i32 : i32 + %negative_expert_count_i32 = scalar.subi %c0_i32, %expert_count_i32 : i32 + %is_last_expert = scalar.cmpi eq, %old_counter, %last_expert_i32 : i32 + scf.if %is_last_expert { + %rounded_assignment_count = index.add %assignment_count, %c31 : index + %assignment_partition_count = index.div %rounded_assignment_count, %c32 : index + %maximum_partition_count = index.add %assignment_partition_count, %launch_expert_count : index + %partition_count_view = buffer.view %partition_table_noalias[%c0_offset] : buffer -> view<1xi32> + %partition_descriptor_view = buffer.view %partition_table_noalias[%c4_bytes] : buffer -> view<[%maximum_partition_count]xi32> + %has_expert = index.cmp ult, %lane, %launch_expert_count : index + %expert_assignment_count_i32 = scf.if %has_expert -> (i32) { + %partition_expert, %table_expert_count = index.assume %lane, %launch_expert_count [lt(%lane, %launch_expert_count)] : index, index + %loaded = view.load %count_view[%partition_expert] : view<[%launch_expert_count]xi32> -> i32 + scf.yield %loaded : i32 + } else { + scf.yield %c0_i32 : i32 + } + %expert_assignment_count0 = index.cast %expert_assignment_count_i32 : i32 to index + %expert_assignment_count = index.assume %expert_assignment_count0 [range(%expert_assignment_count0, 0, 512)] : index + %rounded_expert_assignment_count = index.add %expert_assignment_count, %c31 : index + %expert_partition_count = index.div %rounded_expert_assignment_count, %c32 : index + %expert_partition_count_i32 = index.cast %expert_partition_count : index to i32 + %expert_partition_base_i32 = kernel.workgroup.scan %expert_partition_count_i32 {direction = forward, mode = exclusive} : i32 + %partition_count_i32 = kernel.workgroup.reduce %expert_partition_count_i32 : i32 + %partition_count = index.cast %partition_count_i32 : i32 to index + %bounded_partition_count, %table_partition_capacity = index.assume %partition_count, %maximum_partition_count [lt(%partition_count, %maximum_partition_count)] : index, index + %expert_partition_base0 = index.cast %expert_partition_base_i32 : i32 to index + %expert_partition_base = index.assume %expert_partition_base0 [range(%expert_partition_base0, 0, 255)] : index + scf.if %has_expert { + %partition_expert, %table_expert_count = index.assume %lane, %launch_expert_count [lt(%lane, %launch_expert_count)] : index, index + %expert_i32 = index.cast %partition_expert : index to i32 + scf.for %partition = [%c0 to %expert_partition_count step %c1] { + %descriptor_ordinal0 = index.add %expert_partition_base, %partition : index + %descriptor_ordinal, %descriptor_count = index.assume %descriptor_ordinal0, %bounded_partition_count [lt(%descriptor_ordinal0, %bounded_partition_count)] : index, index + %table_descriptor_ordinal, %table_descriptor_capacity = index.assume %descriptor_ordinal, %maximum_partition_count [lt(%descriptor_ordinal, %maximum_partition_count)] : index, index + %partition_remainder = index.rem %expert_assignment_count, %c32 : index + %has_partial_tail = index.cmp ne, %partition_remainder, %c0 : index + %partition_row_count = scf.if %has_partial_tail -> (index) { + %next_partition = index.add %partition, %c1 : index + %is_tail_partition = index.cmp eq, %next_partition, %expert_partition_count : index + %tail_row_count = scf.if %is_tail_partition -> (index) { + scf.yield %partition_remainder : index + } else { + scf.yield %c32 : index + } + scf.yield %tail_row_count : index + } else { + scf.yield %c32 : index + } + %partition_i32 = index.cast %partition : index to i32 + %partition_row_count_i32 = index.cast %partition_row_count : index to i32 + %packed_partition = scalar.shli %partition_i32, %c7_i32 : i32 + %partition_row_count_minus_one = scalar.subi %partition_row_count_i32, %c1_i32 : i32 + %packed_row_count = scalar.shli %partition_row_count_minus_one, %c13_i32 : i32 + %packed_expert_partition = scalar.ori %expert_i32, %packed_partition : i32 + %packed_descriptor = scalar.ori %packed_expert_partition, %packed_row_count : i32 + view.store %packed_descriptor, %partition_descriptor_view[%table_descriptor_ordinal] : i32, view<[%maximum_partition_count]xi32> + } + } + scf.if %is_lane_zero { + view.store %partition_count_i32, %partition_count_view[%c0] : i32, view<1xi32> + } + // The counter cannot become reusable until every descriptor store is + // complete and visible to the following dispatch. + kernel.barrier scope(workgroup) ordering(acq_rel) + scf.if %is_lane_zero { + view.atomic.reduce %negative_expert_count_i32, %completion_counter_view[%c0] {ordering = release, scope = device} : i32, view<1xi32> + } + } + kernel.return +} + +// Compare both fused outputs against the production composition, then invoke +// the fused route twice against one counter to make reset correctness visible. +check.case public @qwen3_moe_expert_table_partition_fused_differential_case { + %token_count = check.literal value(512) : index + %route_count = check.literal value(8) : index + %route_stride = check.literal value(8) : index + %expert_count = check.literal value(128) : index + %route_ids = check.generate.iota offset(0) step(1) period(128) : tensor<512x8xi32> + %expected_expert_table = check.generate.fill value(-1) : tensor<65664xi32> + %expected_partition_table = check.generate.fill value(-1) : tensor<257xi32> + %actual_expert_table0 = check.generate.fill value(-1) : tensor<65664xi32> + %actual_partition_table0 = check.generate.fill value(-1) : tensor<257xi32> + %actual_expert_table1 = check.generate.fill value(-1) : tensor<65664xi32> + %actual_partition_table1 = check.generate.fill value(-1) : tensor<257xi32> + %completion_counter = check.generate.fill value(0) : tensor<1xi32> + %expected_counter = check.generate.fill value(0) : tensor<1xi32> + kernel.launch @qwen3_moe_build_expert_table[%token_count, %route_count, %route_stride, %expert_count](%token_count, %route_count, %route_stride, %expert_count, %route_ids, %expected_expert_table) : [index, index, index, index](index, index, index, index, tensor<512x8xi32>, tensor<65664xi32>) + kernel.launch @qwen3_moe_build_expert_partition_table[%token_count, %route_count, %expert_count](%token_count, %route_count, %expert_count, %expected_expert_table, %expected_partition_table) : [index, index, index](index, index, index, tensor<65664xi32>, tensor<257xi32>) + kernel.launch @qwen3_moe_build_expert_table_partition_prefill_512[%token_count, %route_count, %route_stride, %expert_count](%token_count, %route_count, %route_stride, %expert_count, %route_ids, %actual_expert_table0, %actual_partition_table0, %completion_counter) : [index, index, index, index](index, index, index, index, tensor<512x8xi32>, tensor<65664xi32>, tensor<257xi32>, tensor<1xi32>) + kernel.launch @qwen3_moe_build_expert_table_partition_prefill_512[%token_count, %route_count, %route_stride, %expert_count](%token_count, %route_count, %route_stride, %expert_count, %route_ids, %actual_expert_table1, %actual_partition_table1, %completion_counter) : [index, index, index, index](index, index, index, index, tensor<512x8xi32>, tensor<65664xi32>, tensor<257xi32>, tensor<1xi32>) + check.expect.equal actual(%actual_expert_table0) expected(%expected_expert_table) : tensor<65664xi32> + check.expect.equal actual(%actual_partition_table0) expected(%expected_partition_table) : tensor<257xi32> + check.expect.equal actual(%actual_expert_table1) expected(%expected_expert_table) : tensor<65664xi32> + check.expect.equal actual(%actual_partition_table1) expected(%expected_partition_table) : tensor<257xi32> + check.expect.equal actual(%completion_counter) expected(%expected_counter) : tensor<1xi32> + check.return +} + +check.case public @qwen3_moe_expert_table_partition_fused_benchmark_case { + %token_count = check.literal value(512) : index + %route_count = check.literal value(8) : index + %route_stride = check.literal value(8) : index + %expert_count = check.literal value(128) : index + %route_ids = check.generate.iota offset(0) step(1) period(128) : tensor<512x8xi32> + %expert_table = check.generate.fill value(-1) : tensor<65664xi32> + %partition_table = check.generate.fill value(-1) : tensor<257xi32> + %completion_counter = check.generate.fill value(0) : tensor<1xi32> + kernel.launch @qwen3_moe_build_expert_table_partition_prefill_512[%token_count, %route_count, %route_stride, %expert_count](%token_count, %route_count, %route_stride, %expert_count, %route_ids, %expert_table, %partition_table, %completion_counter) : [index, index, index, index](index, index, index, index, tensor<512x8xi32>, tensor<65664xi32>, tensor<257xi32>, tensor<1xi32>) + check.return +} + +check.case public @qwen3_moe_expert_table_partition_composed_benchmark_case { + %token_count = check.literal value(512) : index + %route_count = check.literal value(8) : index + %route_stride = check.literal value(8) : index + %expert_count = check.literal value(128) : index + %route_ids = check.generate.iota offset(0) step(1) period(128) : tensor<512x8xi32> + %expert_table = check.generate.fill value(-1) : tensor<65664xi32> + %partition_table = check.generate.fill value(-1) : tensor<257xi32> + kernel.launch @qwen3_moe_build_expert_table[%token_count, %route_count, %route_stride, %expert_count](%token_count, %route_count, %route_stride, %expert_count, %route_ids, %expert_table) : [index, index, index, index](index, index, index, index, tensor<512x8xi32>, tensor<65664xi32>) + kernel.launch @qwen3_moe_build_expert_partition_table[%token_count, %route_count, %expert_count](%token_count, %route_count, %expert_count, %expert_table, %partition_table) : [index, index, index](index, index, index, tensor<65664xi32>, tensor<257xi32>) + check.return +} + +check.benchmark<@qwen3_moe_expert_table_partition_fused_benchmark_case> @qwen3_moe_expert_table_partition_fused_prefill_512 + +check.benchmark<@qwen3_moe_expert_table_partition_composed_benchmark_case> @qwen3_moe_expert_table_partition_composed_prefill_512 diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/flash_attention_decode_f32_f16_wmma.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/flash_attention_decode_f32_f16_wmma.loom new file mode 100644 index 000000000000..fc523520717a --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/flash_attention_decode_f32_f16_wmma.loom @@ -0,0 +1,690 @@ +// Qwen3 MoE grouped-query decode FlashAttention. +// +// One or two four-wave workgroups compute up to 16 query heads that share one +// KV head. Each workgroup walks the exact KV length in 64-row blocks while +// carrying online-softmax max, sum, and output state in registers. With two +// output partitions, each workgroup owns one disjoint 64-channel output half; +// this duplicates QK work but exposes more parallelism without publishing +// partial softmax state. The ownership changes mirror the cooperative-matrix +// schedule used by llama.cpp's Vulkan CM1 kernel: +// +// 1. All workitems stage a scaled 16x128 F16 GQA-head tile. +// 2. Each wave computes one 16x16 QK score slice. +// 3. Scores cross LDS so each wave can normalize four complete query heads. +// 4. F16 probabilities cross LDS for four P*V WMMA steps. +// 5. Each active lane retains one four-channel F16 packet for each of its +// four query heads across subsequent KV blocks. +// +// K and V remain in the row-major llama.cpp cache layout +// [KV token][KV head][128]. Their aligned F16 fragments load directly from +// global memory; there is no expanded or repacked persistent allocation. QK +// and the online-softmax statistics remain F32, while the P*V accumulation and +// carried output match the Vulkan oracle's F16 policy. Unlike the split-K +// fallback, this kernel never publishes partial tensors or completion atomics. +amdgpu.target @qwen3_moe_attention_decode_gfx11_wave64 {subgroup_size = 64} + +config.decl @qwen3_moe.attention.query_head_count : %value: index where [range(%value, 1, 64)] + +config.decl @qwen3_moe.attention.key_value_head_count : %value: index where [range(%value, 1, 64)] + +// Number of independent 64-channel output partitions per KV head. One avoids +// redundant QK work; two exposes more workgroups for short decode contexts. +config.decl @qwen3_moe.attention.decode.output_partition_count : %value: index where [range(%value, 1, 2)] + +kernel.def target(@qwen3_moe_attention_decode_gfx11_wave64) @qwen3_moe_flash_attention_decode_f32_f16_wmma(%key_value_token_count: index) { + %key_value_head_count = config.get @qwen3_moe.attention.key_value_head_count : index + %output_partition_count = config.get @qwen3_moe.attention.decode.output_partition_count : index + %c1 = index.constant 1 : index + %c256 = index.constant 256 : index + %workgroup_count = index.mul %key_value_head_count, %output_partition_count : index + kernel.launch.config workgroups(%workgroup_count, %c1, %c1) workgroup_size(%c256, %c1, %c1) : index +} launch(%key_value_token_count: index, %query: buffer, %key: buffer, %value: buffer, %mask: buffer, %output: buffer) { + %bounded_key_value_token_count = index.assume %key_value_token_count [range(%key_value_token_count, 1, 32768)] : index + %query_head_count = config.get @qwen3_moe.attention.query_head_count : index + %key_value_head_count = config.get @qwen3_moe.attention.key_value_head_count : index + %output_partition_count = config.get @qwen3_moe.attention.decode.output_partition_count : index + %workgroup_x0 = kernel.workgroup.id : index + %workgroup_count0 = index.mul %key_value_head_count, %output_partition_count : index + %workgroup_x, %workgroup_count = index.assume %workgroup_x0, %workgroup_count0 [lt(%workgroup_x0, %workgroup_count0)] : index, index + %key_value_head0 = index.div %workgroup_x, %output_partition_count : index + %key_value_head = index.assume %key_value_head0 [range(%key_value_head0, 0, 63)] : index + %output_partition = index.rem %workgroup_x, %output_partition_count : index + %workitem = kernel.workitem.id : index + %subgroup0 = kernel.subgroup.id : index + %subgroup = index.assume %subgroup0 [range(%subgroup0, 0, 3)] : index + %lane = kernel.subgroup.lane.id : index + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c3 = index.constant 3 : index + %c4 = index.constant 4 : index + %c8 = index.constant 8 : index + %c15 = index.constant 15 : index + %c16 = index.constant 16 : index + %c32 = index.constant 32 : index + %c64 = index.constant 64 : index + %c128 = index.constant 128 : index + %c256 = index.constant 256 : index + %c0_offset = index.constant 0 : offset + %query_stage_bytes = index.constant 4352 : offset + %score_stage_bytes = index.constant 6144 : offset + %probability_stage_bytes = index.constant 3072 : offset + %product_stage_bytes = index.constant 2048 : offset + %tail_key_value_stage_capacity = index.constant 8192 : offset + %c0_f32 = scalar.constant 0.0 : f32 + %negative_large = scalar.constant -1e+30 : f32 + %head_size_f32 = scalar.constant 128.0 : f32 + %attention_scale = scalar.rsqrtf %head_size_f32 : f32 + %c0_f16 = scalar.constant 0.0 : f16 + %c1_f32 = scalar.constant 1.0 : f32 + %output_zero0 = vector.constant 0.0 : vector<4xf16> + %output_zero1 = vector.constant 0.0 : vector<4xf16> + %output_zero2 = vector.constant 0.0 : vector<4xf16> + %output_zero3 = vector.constant 0.0 : vector<4xf16> + %c0_f16x8 = vector.constant 0.0 : vector<8xf16> + %c0_f32x4 = vector.constant 0.0 : vector<4xf32> + %negative_f32x4 = vector.constant -1e+30 : vector<4xf32> + %m = index.constant 16 : index + %n = index.constant 16 : index + %k = index.constant 16 : index + %query_heads_per_key_value_head0 = index.div %query_head_count, %key_value_head_count : index + %query_heads_per_key_value_head = index.assume %query_heads_per_key_value_head0 [range(%query_heads_per_key_value_head0, 1, 16)] : index + %query_head_base = index.mul %key_value_head, %query_heads_per_key_value_head : index + %key_value_width = index.mul %key_value_head_count, %c128 : index + %key_value_head_base = index.mul %key_value_head, %c128 : index + %full_key_value_block_count = index.div %bounded_key_value_token_count, %c64 : index + %full_key_value_token_count0 = index.mul %full_key_value_block_count, %c64 : index + %full_key_value_token_count = index.assume %full_key_value_token_count0 [range(%full_key_value_token_count0, 0, 32768), mul(%full_key_value_token_count0, 64)] : index + %tail_key_value_token_count = index.sub %bounded_key_value_token_count, %full_key_value_token_count : index + %has_key_value_tail = index.cmp ne, %tail_key_value_token_count, %c0 : index + %has_single_key_value_tail = index.cmp eq, %tail_key_value_token_count, %c1 : index + %tail_key_value_stage_bytes = scf.select %has_key_value_tail, %tail_key_value_stage_capacity, %c0_offset : offset + %subgroup_score_column = index.mul %subgroup, %c16 : index + %subgroup_query_row = index.mul %subgroup, %c4 : index + %query_row0 = index.add %subgroup_query_row, %c0 : index + %query_row1 = index.add %subgroup_query_row, %c1 : index + %query_row2 = index.add %subgroup_query_row, %c2 : index + %query_row3 = index.add %subgroup_query_row, %c3 : index + %query_head0 = index.add %query_head_base, %query_row0 : index + %query_head1 = index.add %query_head_base, %query_row1 : index + %query_head2 = index.add %query_head_base, %query_row2 : index + %query_head3 = index.add %query_head_base, %query_row3 : index + %query_valid0 = index.cmp ult, %query_head0, %query_head_count : index + %query_valid1 = index.cmp ult, %query_head1, %query_head_count : index + %query_valid2 = index.cmp ult, %query_head2, %query_head_count : index + %query_valid3 = index.cmp ult, %query_head3, %query_head_count : index + %query_valid = vector.from_elements %query_valid0, %query_valid1, %query_valid2, %query_valid3 : vector<4xi1> + %subgroup_product_channel = index.mul %subgroup, %c16 : index + %lane_output_tile = index.div %lane, %c16 : index + %output_tile_count = index.div %c2, %output_partition_count : index + %output_tile_end = index.add %output_partition, %output_tile_count : index + %lane_product_channel0 = index.rem %lane, %c16 : index + %lane_product_channel = index.mul %lane_product_channel0, %c4 : index + %lane_output_channel = index.mul %lane, %c4 : index + %lane_output_at_or_after_partition = index.cmp uge, %lane_output_tile, %output_partition : index + %lane_output_before_end = index.cmp ult, %lane_output_tile, %output_tile_end : index + %lane_has_output = scalar.andi %lane_output_at_or_after_partition, %lane_output_before_end : i1 + // The padded LDS rows mirror the Vulkan oracle's ownership changes. Q uses + // eight spare F16 columns after its 128 channels. Score and probability + // transpose to key-major rows with eight spare columns after 16 queries. + // These strides avoid the bank pattern produced by dense transposed rows. + %query_transposed_layout = encoding.layout.strided [1, 136] : encoding + %probability_transposed_layout = encoding.layout.strided [1, 24] : encoding + %query_noalias, %key_noalias, %value_noalias, %mask_noalias, %output_noalias = buffer.assume.noalias %query, %key, %value, %mask, %output : buffer, buffer, buffer, buffer, buffer + %query_aligned = buffer.assume.alignment %query_noalias {minimum_alignment = 16} : buffer + %key_aligned = buffer.assume.alignment %key_noalias {minimum_alignment = 16} : buffer + %value_aligned = buffer.assume.alignment %value_noalias {minimum_alignment = 16} : buffer + %mask_aligned = buffer.assume.alignment %mask_noalias {minimum_alignment = 16} : buffer + %output_aligned = buffer.assume.alignment %output_noalias {minimum_alignment = 16} : buffer + %query_view = buffer.view %query_aligned[%c0_offset] : buffer -> view<[%query_head_count]x128xf32> + %key_view = buffer.view %key_aligned[%c0_offset] : buffer -> view<[%bounded_key_value_token_count]x[%key_value_width]xf16> + %value_view = buffer.view %value_aligned[%c0_offset] : buffer -> view<[%bounded_key_value_token_count]x[%key_value_width]xf16> + %mask_view = buffer.view %mask_aligned[%c0_offset] : buffer -> view<[%bounded_key_value_token_count]xf16> + %output_view = buffer.view %output_aligned[%c0_offset] : buffer -> view<[%query_head_count]x128xf32> + %query_stage = buffer.alloca align(16) %query_stage_bytes : buffer + %score_stage = buffer.alloca align(16) %score_stage_bytes : buffer + %probability_stage = buffer.alloca align(16) %probability_stage_bytes : buffer + %product_stage = buffer.alloca align(16) %product_stage_bytes : buffer + %tail_key_value_stage = buffer.alloca align(16) %tail_key_value_stage_bytes : buffer + %query_stage_view = buffer.view %query_stage[%c0_offset] : buffer -> view<16x136xf16> + %query_transposed_view = buffer.view %query_stage[%c0_offset] : buffer -> view<128x16xf16, %query_transposed_layout> + %score_stage_view = buffer.view %score_stage[%c0_offset] : buffer -> view<64x24xf32> + %probability_stage_view = buffer.view %probability_stage[%c0_offset] : buffer -> view<16x64xf16, %probability_transposed_layout> + %product_stage_view = buffer.view %product_stage[%c0_offset] : buffer -> view<16x64xf16> + %tail_key_value_stage_view = buffer.view %tail_key_value_stage[%c0_offset] : buffer -> view<32x128xf16> + // Scale and truncate Q exactly once. The Vulkan reference does this before + // entering its KV loop, making QK a native F16 WMMA while retaining F32 + // accumulation. + scf.for %load_iteration = [%c0 to %c8 step %c1] unroll { + %linear = index.madd %load_iteration, %c256, %workitem : index + %local_query_row = index.div %linear, %c128 : index + %query_channel = index.rem %linear, %c128 : index + %local_query_head = index.add %query_head_base, %local_query_row : index + %query_valid_load = index.cmp ult, %local_query_head, %query_head_count : index + %query_value = scf.if %query_valid_load -> (f16) { + %loaded = view.load %query_view[%local_query_head, %query_channel] : view<[%query_head_count]x128xf32> -> f32 + %scaled = scalar.mulf %loaded, %attention_scale : f32 + %truncated = scalar.fptrunc %scaled : f32 to f16 + scf.yield %truncated : f16 + } else { + scf.yield %c0_f16 : f16 + } + view.store %query_value, %query_stage_view[%local_query_row, %query_channel] : f16, view<16x136xf16> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + %full_max, %full_sum, %full_output0, %full_output1, %full_output2, %full_output3 = scf.for %key_origin = [%c0 to %full_key_value_token_count step %c64](%current_max = %negative_f32x4 : vector<4xf32>, %current_sum = %c0_f32x4 : vector<4xf32>, %current_output0 = %output_zero0 : vector<4xf16>, %current_output1 = %output_zero1 : vector<4xf16>, %current_output2 = %output_zero2 : vector<4xf16>, %current_output3 = %output_zero3 : vector<4xf16>) -> (vector<4xf32>, vector<4xf32>, vector<4xf16>, vector<4xf16>, vector<4xf16>, vector<4xf16>) { + // Four independent wave-level WMMAs produce a 16x64 score tile. + %score_key_origin0 = index.add %key_origin, %subgroup_score_column : index + %last_full_key_tile_start = index.sub %bounded_key_value_token_count, %c15 : index + %score_key_origin = index.assume %score_key_origin0 [lt(%score_key_origin0, %last_full_key_tile_start)] : index + %score_init_values = vector.constant 0.0 : vector<4xf32> + %score_init = vector.fragment %score_init_values shape [%m, %n] : vector<4xf32> + %score_fragment = scf.for %head_tile = [%c0 to %c128 step %c16](%score_accumulator = %score_init : vector<4xf32>) -> (vector<4xf32>) unroll { + %key_channel = index.add %key_value_head_base, %head_tile : index + %key_fragment = vector.fragment.load %key_view[%score_key_origin, %key_channel] shape [%m, %k] : view<[%bounded_key_value_token_count]x[%key_value_width]xf16> -> vector<16xf16> + %query_fragment = vector.fragment.load %query_transposed_view[%head_tile, %c0] shape [%k, %n] : view<128x16xf16, %query_transposed_layout> -> vector<16xf16> + %next_score_accumulator = vector.mma %key_fragment, %query_fragment, %score_accumulator : vector<16xf16>, vector<16xf16>, vector<4xf32> + scf.yield %next_score_accumulator : vector<4xf32> + } + vector.fragment.store %score_fragment, %score_stage_view[%subgroup_score_column, %c0] shape [%m, %n] : vector<4xf32>, view<64x24xf32> + kernel.barrier scope(workgroup) ordering(acq_rel) + // LDS transposes ownership from one 16-column score slice per wave to + // four complete query rows per wave. Every lane contributes one key + // column to each of those rows. + %key_token0 = index.add %key_origin, %lane : index + %key_token = index.assume %key_token0 [lt(%key_token0, %bounded_key_value_token_count)] : index + %raw_score0 = view.load %score_stage_view[%lane, %query_row0] : view<64x24xf32> -> f32 + %raw_score1 = view.load %score_stage_view[%lane, %query_row1] : view<64x24xf32> -> f32 + %raw_score2 = view.load %score_stage_view[%lane, %query_row2] : view<64x24xf32> -> f32 + %raw_score3 = view.load %score_stage_view[%lane, %query_row3] : view<64x24xf32> -> f32 + %mask_f16 = view.load %mask_view[%key_token] : view<[%bounded_key_value_token_count]xf16> -> f16 + %mask_f32 = scalar.extf %mask_f16 : f16 to f32 + %masked_score0 = scf.if %query_valid0 -> (f32) { + %score = scalar.addf %raw_score0, %mask_f32 : f32 + scf.yield %score : f32 + } else { + scf.yield %negative_large : f32 + } + %masked_score1 = scf.if %query_valid1 -> (f32) { + %score = scalar.addf %raw_score1, %mask_f32 : f32 + scf.yield %score : f32 + } else { + scf.yield %negative_large : f32 + } + %masked_score2 = scf.if %query_valid2 -> (f32) { + %score = scalar.addf %raw_score2, %mask_f32 : f32 + scf.yield %score : f32 + } else { + scf.yield %negative_large : f32 + } + %masked_score3 = scf.if %query_valid3 -> (f32) { + %score = scalar.addf %raw_score3, %mask_f32 : f32 + scf.yield %score : f32 + } else { + scf.yield %negative_large : f32 + } + %masked_scores = vector.from_elements %masked_score0, %masked_score1, %masked_score2, %masked_score3 : vector<4xf32> + %block_max = kernel.subgroup.reduce %masked_scores : vector<4xf32> + %next_max = vector.maxnumf %current_max, %block_max : vector<4xf32> + %score_delta = vector.subf %masked_scores, %next_max : vector<4xf32> + %raw_probability = vector.expf %score_delta : vector<4xf32> + %probability = vector.select %query_valid, %raw_probability, %c0_f32x4 : vector<4xf32> + %block_sum = kernel.subgroup.reduce %probability : vector<4xf32> + %old_delta = vector.subf %current_max, %next_max : vector<4xf32> + %old_scale = vector.expf %old_delta : vector<4xf32> + %scaled_current_sum = vector.mulf %current_sum, %old_scale : vector<4xf32> + %next_sum = vector.addf %scaled_current_sum, %block_sum : vector<4xf32> + %probability_f16 = vector.fptrunc %probability : vector<4xf32> to vector<4xf16> + %probability0 = vector.extract %probability_f16[0] : vector<4xf16> -> f16 + %probability1 = vector.extract %probability_f16[1] : vector<4xf16> -> f16 + %probability2 = vector.extract %probability_f16[2] : vector<4xf16> -> f16 + %probability3 = vector.extract %probability_f16[3] : vector<4xf16> -> f16 + view.store %probability0, %probability_stage_view[%query_row0, %lane] : f16, view<16x64xf16, %probability_transposed_layout> + view.store %probability1, %probability_stage_view[%query_row1, %lane] : f16, view<16x64xf16, %probability_transposed_layout> + view.store %probability2, %probability_stage_view[%query_row2, %lane] : f16, view<16x64xf16, %probability_transposed_layout> + view.store %probability3, %probability_stage_view[%query_row3, %lane] : f16, view<16x64xf16, %probability_transposed_layout> + kernel.barrier scope(workgroup) ordering(acq_rel) + // Match the Vulkan CM1 ownership schedule by computing two sequential + // 64-channel output tiles. This keeps one P*V accumulator live per wave + // and reuses a 2 KiB exchange tile instead of retaining both halves. + %old_scale_f16 = vector.fptrunc %old_scale : vector<4xf32> to vector<4xf16> + %old_scale0_scalar = vector.extract %old_scale_f16[0] : vector<4xf16> -> f16 + %old_scale1_scalar = vector.extract %old_scale_f16[1] : vector<4xf16> -> f16 + %old_scale2_scalar = vector.extract %old_scale_f16[2] : vector<4xf16> -> f16 + %old_scale3_scalar = vector.extract %old_scale_f16[3] : vector<4xf16> -> f16 + %old_scale0 = vector.splat %old_scale0_scalar : vector<4xf16> + %old_scale1 = vector.splat %old_scale1_scalar : vector<4xf16> + %old_scale2 = vector.splat %old_scale2_scalar : vector<4xf16> + %old_scale3 = vector.splat %old_scale3_scalar : vector<4xf16> + %scaled_current_output0 = vector.mulf %current_output0, %old_scale0 : vector<4xf16> + %scaled_current_output1 = vector.mulf %current_output1, %old_scale1 : vector<4xf16> + %scaled_current_output2 = vector.mulf %current_output2, %old_scale2 : vector<4xf16> + %scaled_current_output3 = vector.mulf %current_output3, %old_scale3 : vector<4xf16> + %next_output0, %next_output1, %next_output2, %next_output3 = scf.for %output_tile = [%c0 to %c2 step %c1](%tile_output0 = %scaled_current_output0 : vector<4xf16>, %tile_output1 = %scaled_current_output1 : vector<4xf16>, %tile_output2 = %scaled_current_output2 : vector<4xf16>, %tile_output3 = %scaled_current_output3 : vector<4xf16>) -> (vector<4xf16>, vector<4xf16>, vector<4xf16>, vector<4xf16>) unroll { + %tile_at_or_after_partition = index.cmp uge, %output_tile, %output_partition : index + %tile_before_partition_end = index.cmp ult, %output_tile, %output_tile_end : index + %tile_selected = scalar.andi %tile_at_or_after_partition, %tile_before_partition_end : i1 + %updated_output0, %updated_output1, %updated_output2, %updated_output3 = scf.if %tile_selected -> (vector<4xf16>, vector<4xf16>, vector<4xf16>, vector<4xf16>) { + %output_tile_channel = index.mul %output_tile, %c64 : index + %value_channel0 = index.add %key_value_head_base, %output_tile_channel : index + %value_channel = index.add %value_channel0, %subgroup_product_channel : index + %product_init = vector.fragment %c0_f16x8 shape [%m, %n] : vector<8xf16> + %product_fragment = scf.for %key_tile = [%c0 to %c64 step %c16](%product_accumulator = %product_init : vector<8xf16>) -> (vector<8xf16>) unroll { + %value_token0 = index.add %key_origin, %key_tile : index + %value_token = index.assume %value_token0 [lt(%value_token0, %last_full_key_tile_start)] : index + %probability_fragment = vector.fragment.load %probability_stage_view[%c0, %key_tile] shape [%m, %k] : view<16x64xf16, %probability_transposed_layout> -> vector<16xf16> + %value_fragment = vector.fragment.load %value_view[%value_token, %value_channel] shape [%k, %n] : view<[%bounded_key_value_token_count]x[%key_value_width]xf16> -> vector<16xf16> + %next_product_accumulator = vector.mma %probability_fragment, %value_fragment, %product_accumulator : vector<16xf16>, vector<16xf16>, vector<8xf16> + scf.yield %next_product_accumulator : vector<8xf16> + } + vector.fragment.store %product_fragment, %product_stage_view[%c0, %subgroup_product_channel] shape [%m, %n] : vector<8xf16>, view<16x64xf16> + kernel.barrier scope(workgroup) ordering(acq_rel) + %owns_output_tile = index.cmp eq, %lane_output_tile, %output_tile : index + %next_tile_output0, %next_tile_output1, %next_tile_output2, %next_tile_output3 = scf.if %owns_output_tile -> (vector<4xf16>, vector<4xf16>, vector<4xf16>, vector<4xf16>) { + %block_output0 = vector.load %product_stage_view[%query_row0, %lane_product_channel] : view<16x64xf16> -> vector<4xf16> + %block_output1 = vector.load %product_stage_view[%query_row1, %lane_product_channel] : view<16x64xf16> -> vector<4xf16> + %block_output2 = vector.load %product_stage_view[%query_row2, %lane_product_channel] : view<16x64xf16> -> vector<4xf16> + %block_output3 = vector.load %product_stage_view[%query_row3, %lane_product_channel] : view<16x64xf16> -> vector<4xf16> + %next_tile_output0 = vector.addf %tile_output0, %block_output0 : vector<4xf16> + %next_tile_output1 = vector.addf %tile_output1, %block_output1 : vector<4xf16> + %next_tile_output2 = vector.addf %tile_output2, %block_output2 : vector<4xf16> + %next_tile_output3 = vector.addf %tile_output3, %block_output3 : vector<4xf16> + scf.yield %next_tile_output0, %next_tile_output1, %next_tile_output2, %next_tile_output3 : vector<4xf16>, vector<4xf16>, vector<4xf16>, vector<4xf16> + } else { + scf.yield %tile_output0, %tile_output1, %tile_output2, %tile_output3 : vector<4xf16>, vector<4xf16>, vector<4xf16>, vector<4xf16> + } + // Complete every read before another output tile overwrites LDS. + kernel.barrier scope(workgroup) ordering(acq_rel) + scf.yield %next_tile_output0, %next_tile_output1, %next_tile_output2, %next_tile_output3 : vector<4xf16>, vector<4xf16>, vector<4xf16>, vector<4xf16> + } else { + scf.yield %tile_output0, %tile_output1, %tile_output2, %tile_output3 : vector<4xf16>, vector<4xf16>, vector<4xf16>, vector<4xf16> + } + scf.yield %updated_output0, %updated_output1, %updated_output2, %updated_output3 : vector<4xf16>, vector<4xf16>, vector<4xf16>, vector<4xf16> + } + scf.yield %next_max, %next_sum, %next_output0, %next_output1, %next_output2, %next_output3 : vector<4xf32>, vector<4xf32>, vector<4xf16>, vector<4xf16>, vector<4xf16>, vector<4xf16> + } + // A single trailing KV row is cheaper as a native wave reduction and direct + // V update. Skip the general WMMA tail loop for that exact JIT-specialized + // shape; larger tails continue through the masked 32-row schedule below. + %wmma_tail_start = scf.select %has_single_key_value_tail, %bounded_key_value_token_count, %full_key_value_token_count : index + %tail_score_wave = index.cmp ult, %subgroup, %c2 : index + %wmma_tail_max, %wmma_tail_sum, %wmma_tail_output0, %wmma_tail_output1, %wmma_tail_output2, %wmma_tail_output3 = scf.for %tail_key_origin = [%wmma_tail_start to %bounded_key_value_token_count step %c32](%current_max = %full_max : vector<4xf32>, %current_sum = %full_sum : vector<4xf32>, %current_output0 = %full_output0 : vector<4xf16>, %current_output1 = %full_output1 : vector<4xf16>, %current_output2 = %full_output2 : vector<4xf16>, %current_output3 = %full_output3 : vector<4xf16>) -> (vector<4xf32>, vector<4xf32>, vector<4xf16>, vector<4xf16>, vector<4xf16>, vector<4xf16>) { + %tail_remaining = index.sub %bounded_key_value_token_count, %tail_key_origin : index + %tail_key_count = index.min %tail_remaining, %c32 : index + // Cooperatively stage one K tile, explicitly zeroing the padded rows. + scf.for %load_iteration = [%c0 to %c16 step %c1] unroll { + %linear = index.madd %load_iteration, %c256, %workitem : index + %tail_key_row = index.div %linear, %c128 : index + %tail_key_channel = index.rem %linear, %c128 : index + %tail_key_valid = index.cmp ult, %tail_key_row, %tail_key_count : index + %tail_key_value = scf.if %tail_key_valid -> (f16) { + %tail_key_token0 = index.add %tail_key_origin, %tail_key_row : index + %tail_key_token = index.assume %tail_key_token0 [lt(%tail_key_token0, %bounded_key_value_token_count)] : index + %global_key_channel = index.add %key_value_head_base, %tail_key_channel : index + %loaded = view.load %key_view[%tail_key_token, %global_key_channel] : view<[%bounded_key_value_token_count]x[%key_value_width]xf16> -> f16 + scf.yield %loaded : f16 + } else { + scf.yield %c0_f16 : f16 + } + view.store %tail_key_value, %tail_key_value_stage_view[%tail_key_row, %tail_key_channel] : f16, view<32x128xf16> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + // Waves zero and one compute the two 16-column QK fragments in this tile. + scf.if %tail_score_wave { + %tail_score_subgroup = index.assume %subgroup [range(%subgroup, 0, 1)] : index + %tail_score_column = index.mul %tail_score_subgroup, %c16 : index + %tail_score_init_values = vector.constant 0.0 : vector<4xf32> + %tail_score_init = vector.fragment %tail_score_init_values shape [%m, %n] : vector<4xf32> + %tail_score_fragment = scf.for %head_tile = [%c0 to %c128 step %c16](%score_accumulator = %tail_score_init : vector<4xf32>) -> (vector<4xf32>) unroll { + %key_fragment = vector.fragment.load %tail_key_value_stage_view[%tail_score_column, %head_tile] shape [%m, %k] : view<32x128xf16> -> vector<16xf16> + %query_fragment = vector.fragment.load %query_transposed_view[%head_tile, %c0] shape [%k, %n] : view<128x16xf16, %query_transposed_layout> -> vector<16xf16> + %next_score_accumulator = vector.mma %key_fragment, %query_fragment, %score_accumulator : vector<16xf16>, vector<16xf16>, vector<4xf32> + scf.yield %next_score_accumulator : vector<4xf32> + } + vector.fragment.store %tail_score_fragment, %score_stage_view[%tail_score_column, %c0] shape [%m, %n] : vector<4xf32>, view<64x24xf32> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + // LDS changes ownership from the QK wave to four query-row waves. Lanes + // beyond the logical tail never read the score or mask buffers. + %tail_lane_valid = index.cmp ult, %lane, %tail_key_count : index + %tail_valid0 = scalar.andi %tail_lane_valid, %query_valid0 : i1 + %tail_valid1 = scalar.andi %tail_lane_valid, %query_valid1 : i1 + %tail_valid2 = scalar.andi %tail_lane_valid, %query_valid2 : i1 + %tail_valid3 = scalar.andi %tail_lane_valid, %query_valid3 : i1 + %tail_valid = vector.from_elements %tail_valid0, %tail_valid1, %tail_valid2, %tail_valid3 : vector<4xi1> + %tail_key_token0 = index.add %tail_key_origin, %lane : index + %tail_mask_f32 = scf.if %tail_lane_valid -> (f32) { + %tail_key_token = index.assume %tail_key_token0 [lt(%tail_key_token0, %bounded_key_value_token_count)] : index + %mask_f16 = view.load %mask_view[%tail_key_token] : view<[%bounded_key_value_token_count]xf16> -> f16 + %mask_f32 = scalar.extf %mask_f16 : f16 to f32 + scf.yield %mask_f32 : f32 + } else { + scf.yield %c0_f32 : f32 + } + %masked_score0 = scf.if %tail_valid0 -> (f32) { + %raw_score = view.load %score_stage_view[%lane, %query_row0] : view<64x24xf32> -> f32 + %score = scalar.addf %raw_score, %tail_mask_f32 : f32 + scf.yield %score : f32 + } else { + scf.yield %negative_large : f32 + } + %masked_score1 = scf.if %tail_valid1 -> (f32) { + %raw_score = view.load %score_stage_view[%lane, %query_row1] : view<64x24xf32> -> f32 + %score = scalar.addf %raw_score, %tail_mask_f32 : f32 + scf.yield %score : f32 + } else { + scf.yield %negative_large : f32 + } + %masked_score2 = scf.if %tail_valid2 -> (f32) { + %raw_score = view.load %score_stage_view[%lane, %query_row2] : view<64x24xf32> -> f32 + %score = scalar.addf %raw_score, %tail_mask_f32 : f32 + scf.yield %score : f32 + } else { + scf.yield %negative_large : f32 + } + %masked_score3 = scf.if %tail_valid3 -> (f32) { + %raw_score = view.load %score_stage_view[%lane, %query_row3] : view<64x24xf32> -> f32 + %score = scalar.addf %raw_score, %tail_mask_f32 : f32 + scf.yield %score : f32 + } else { + scf.yield %negative_large : f32 + } + %masked_scores = vector.from_elements %masked_score0, %masked_score1, %masked_score2, %masked_score3 : vector<4xf32> + %block_max = kernel.subgroup.reduce %masked_scores : vector<4xf32> + %next_max = vector.maxnumf %current_max, %block_max : vector<4xf32> + %score_delta = vector.subf %masked_scores, %next_max : vector<4xf32> + %raw_probability = vector.expf %score_delta : vector<4xf32> + %probability = vector.select %tail_valid, %raw_probability, %c0_f32x4 : vector<4xf32> + %block_sum = kernel.subgroup.reduce %probability : vector<4xf32> + %old_delta = vector.subf %current_max, %next_max : vector<4xf32> + %old_scale = vector.expf %old_delta : vector<4xf32> + %scaled_current_sum = vector.mulf %current_sum, %old_scale : vector<4xf32> + %next_sum = vector.addf %scaled_current_sum, %block_sum : vector<4xf32> + %probability_f16 = vector.fptrunc %probability : vector<4xf32> to vector<4xf16> + %probability0 = vector.extract %probability_f16[0] : vector<4xf16> -> f16 + %probability1 = vector.extract %probability_f16[1] : vector<4xf16> -> f16 + %probability2 = vector.extract %probability_f16[2] : vector<4xf16> -> f16 + %probability3 = vector.extract %probability_f16[3] : vector<4xf16> -> f16 + view.store %probability0, %probability_stage_view[%query_row0, %lane] : f16, view<16x64xf16, %probability_transposed_layout> + view.store %probability1, %probability_stage_view[%query_row1, %lane] : f16, view<16x64xf16, %probability_transposed_layout> + view.store %probability2, %probability_stage_view[%query_row2, %lane] : f16, view<16x64xf16, %probability_transposed_layout> + view.store %probability3, %probability_stage_view[%query_row3, %lane] : f16, view<16x64xf16, %probability_transposed_layout> + kernel.barrier scope(workgroup) ordering(acq_rel) + // Reuse product scratch for V after every probability is resident in its + // disjoint LDS tile. + scf.for %load_iteration = [%c0 to %c16 step %c1] unroll { + %linear = index.madd %load_iteration, %c256, %workitem : index + %tail_value_row = index.div %linear, %c128 : index + %tail_value_channel = index.rem %linear, %c128 : index + %tail_value_valid = index.cmp ult, %tail_value_row, %tail_key_count : index + %tail_value = scf.if %tail_value_valid -> (f16) { + %tail_value_token0 = index.add %tail_key_origin, %tail_value_row : index + %tail_value_token = index.assume %tail_value_token0 [lt(%tail_value_token0, %bounded_key_value_token_count)] : index + %global_value_channel = index.add %key_value_head_base, %tail_value_channel : index + %loaded = view.load %value_view[%tail_value_token, %global_value_channel] : view<[%bounded_key_value_token_count]x[%key_value_width]xf16> -> f16 + scf.yield %loaded : f16 + } else { + scf.yield %c0_f16 : f16 + } + view.store %tail_value, %tail_key_value_stage_view[%tail_value_row, %tail_value_channel] : f16, view<32x128xf16> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + // Keep tail V staging separate from the 2 KiB product exchange, then use + // the same two 64-channel phases as the aligned path. + %old_scale_f16 = vector.fptrunc %old_scale : vector<4xf32> to vector<4xf16> + %old_scale0_scalar = vector.extract %old_scale_f16[0] : vector<4xf16> -> f16 + %old_scale1_scalar = vector.extract %old_scale_f16[1] : vector<4xf16> -> f16 + %old_scale2_scalar = vector.extract %old_scale_f16[2] : vector<4xf16> -> f16 + %old_scale3_scalar = vector.extract %old_scale_f16[3] : vector<4xf16> -> f16 + %old_scale0 = vector.splat %old_scale0_scalar : vector<4xf16> + %old_scale1 = vector.splat %old_scale1_scalar : vector<4xf16> + %old_scale2 = vector.splat %old_scale2_scalar : vector<4xf16> + %old_scale3 = vector.splat %old_scale3_scalar : vector<4xf16> + %scaled_current_output0 = vector.mulf %current_output0, %old_scale0 : vector<4xf16> + %scaled_current_output1 = vector.mulf %current_output1, %old_scale1 : vector<4xf16> + %scaled_current_output2 = vector.mulf %current_output2, %old_scale2 : vector<4xf16> + %scaled_current_output3 = vector.mulf %current_output3, %old_scale3 : vector<4xf16> + %next_output0, %next_output1, %next_output2, %next_output3 = scf.for %output_tile = [%c0 to %c2 step %c1](%tile_output0 = %scaled_current_output0 : vector<4xf16>, %tile_output1 = %scaled_current_output1 : vector<4xf16>, %tile_output2 = %scaled_current_output2 : vector<4xf16>, %tile_output3 = %scaled_current_output3 : vector<4xf16>) -> (vector<4xf16>, vector<4xf16>, vector<4xf16>, vector<4xf16>) unroll { + %tile_at_or_after_partition = index.cmp uge, %output_tile, %output_partition : index + %tile_before_partition_end = index.cmp ult, %output_tile, %output_tile_end : index + %tile_selected = scalar.andi %tile_at_or_after_partition, %tile_before_partition_end : i1 + %updated_output0, %updated_output1, %updated_output2, %updated_output3 = scf.if %tile_selected -> (vector<4xf16>, vector<4xf16>, vector<4xf16>, vector<4xf16>) { + %output_tile_channel = index.mul %output_tile, %c64 : index + %value_channel = index.add %output_tile_channel, %subgroup_product_channel : index + %tail_product_init = vector.fragment %c0_f16x8 shape [%m, %n] : vector<8xf16> + %tail_product_fragment = scf.for %key_tile = [%c0 to %c32 step %c16](%product_accumulator = %tail_product_init : vector<8xf16>) -> (vector<8xf16>) unroll { + %probability_fragment = vector.fragment.load %probability_stage_view[%c0, %key_tile] shape [%m, %k] : view<16x64xf16, %probability_transposed_layout> -> vector<16xf16> + %value_fragment = vector.fragment.load %tail_key_value_stage_view[%key_tile, %value_channel] shape [%k, %n] : view<32x128xf16> -> vector<16xf16> + %next_product_accumulator = vector.mma %probability_fragment, %value_fragment, %product_accumulator : vector<16xf16>, vector<16xf16>, vector<8xf16> + scf.yield %next_product_accumulator : vector<8xf16> + } + vector.fragment.store %tail_product_fragment, %product_stage_view[%c0, %subgroup_product_channel] shape [%m, %n] : vector<8xf16>, view<16x64xf16> + kernel.barrier scope(workgroup) ordering(acq_rel) + %owns_output_tile = index.cmp eq, %lane_output_tile, %output_tile : index + %next_tile_output0, %next_tile_output1, %next_tile_output2, %next_tile_output3 = scf.if %owns_output_tile -> (vector<4xf16>, vector<4xf16>, vector<4xf16>, vector<4xf16>) { + %block_output0 = vector.load %product_stage_view[%query_row0, %lane_product_channel] : view<16x64xf16> -> vector<4xf16> + %block_output1 = vector.load %product_stage_view[%query_row1, %lane_product_channel] : view<16x64xf16> -> vector<4xf16> + %block_output2 = vector.load %product_stage_view[%query_row2, %lane_product_channel] : view<16x64xf16> -> vector<4xf16> + %block_output3 = vector.load %product_stage_view[%query_row3, %lane_product_channel] : view<16x64xf16> -> vector<4xf16> + %next_tile_output0 = vector.addf %tile_output0, %block_output0 : vector<4xf16> + %next_tile_output1 = vector.addf %tile_output1, %block_output1 : vector<4xf16> + %next_tile_output2 = vector.addf %tile_output2, %block_output2 : vector<4xf16> + %next_tile_output3 = vector.addf %tile_output3, %block_output3 : vector<4xf16> + scf.yield %next_tile_output0, %next_tile_output1, %next_tile_output2, %next_tile_output3 : vector<4xf16>, vector<4xf16>, vector<4xf16>, vector<4xf16> + } else { + scf.yield %tile_output0, %tile_output1, %tile_output2, %tile_output3 : vector<4xf16>, vector<4xf16>, vector<4xf16>, vector<4xf16> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + scf.yield %next_tile_output0, %next_tile_output1, %next_tile_output2, %next_tile_output3 : vector<4xf16>, vector<4xf16>, vector<4xf16>, vector<4xf16> + } else { + scf.yield %tile_output0, %tile_output1, %tile_output2, %tile_output3 : vector<4xf16>, vector<4xf16>, vector<4xf16>, vector<4xf16> + } + scf.yield %updated_output0, %updated_output1, %updated_output2, %updated_output3 : vector<4xf16>, vector<4xf16>, vector<4xf16>, vector<4xf16> + } + scf.yield %next_max, %next_sum, %next_output0, %next_output1, %next_output2, %next_output3 : vector<4xf32>, vector<4xf32>, vector<4xf16>, vector<4xf16>, vector<4xf16>, vector<4xf16> + } + %final_max, %final_sum, %final_output0, %final_output1, %final_output2, %final_output3 = scf.if %has_single_key_value_tail -> (vector<4xf32>, vector<4xf32>, vector<4xf16>, vector<4xf16>, vector<4xf16>, vector<4xf16>) { + %single_tail_token = index.assume %full_key_value_token_count [lt(%full_key_value_token_count, %bounded_key_value_token_count)] : index + %single_tail_channel = index.mul %lane, %c2 : index + %single_tail_key_channel = index.add %key_value_head_base, %single_tail_channel : index + %key_pair_f16 = vector.load %key_view[%single_tail_token, %single_tail_key_channel] : view<[%bounded_key_value_token_count]x[%key_value_width]xf16> -> vector<2xf16> + %query_pair0_f16 = vector.load %query_stage_view[%query_row0, %single_tail_channel] : view<16x136xf16> -> vector<2xf16> + %query_pair1_f16 = vector.load %query_stage_view[%query_row1, %single_tail_channel] : view<16x136xf16> -> vector<2xf16> + %query_pair2_f16 = vector.load %query_stage_view[%query_row2, %single_tail_channel] : view<16x136xf16> -> vector<2xf16> + %query_pair3_f16 = vector.load %query_stage_view[%query_row3, %single_tail_channel] : view<16x136xf16> -> vector<2xf16> + %key_pair = vector.extf %key_pair_f16 : vector<2xf16> to vector<2xf32> + %query_pair0 = vector.extf %query_pair0_f16 : vector<2xf16> to vector<2xf32> + %query_pair1 = vector.extf %query_pair1_f16 : vector<2xf16> to vector<2xf32> + %query_pair2 = vector.extf %query_pair2_f16 : vector<2xf16> to vector<2xf32> + %query_pair3 = vector.extf %query_pair3_f16 : vector<2xf16> to vector<2xf32> + %product_pair0 = vector.mulf %query_pair0, %key_pair : vector<2xf32> + %product_pair1 = vector.mulf %query_pair1, %key_pair : vector<2xf32> + %product_pair2 = vector.mulf %query_pair2, %key_pair : vector<2xf32> + %product_pair3 = vector.mulf %query_pair3, %key_pair : vector<2xf32> + %partial_score0 = vector.reduce %product_pair0, %c0_f32 : vector<2xf32>, f32 + %partial_score1 = vector.reduce %product_pair1, %c0_f32 : vector<2xf32>, f32 + %partial_score2 = vector.reduce %product_pair2, %c0_f32 : vector<2xf32>, f32 + %partial_score3 = vector.reduce %product_pair3, %c0_f32 : vector<2xf32>, f32 + %partial_scores = vector.from_elements %partial_score0, %partial_score1, %partial_score2, %partial_score3 : vector<4xf32> + %reduced_scores = kernel.subgroup.reduce %partial_scores : vector<4xf32> + %mask_f16 = view.load %mask_view[%single_tail_token] : view<[%bounded_key_value_token_count]xf16> -> f16 + %mask_f32 = scalar.extf %mask_f16 : f16 to f32 + %mask_vector = vector.splat %mask_f32 : vector<4xf32> + %raw_scores = vector.addf %reduced_scores, %mask_vector : vector<4xf32> + %masked_scores = vector.select %query_valid, %raw_scores, %negative_f32x4 : vector<4xf32> + %next_max = vector.maxnumf %wmma_tail_max, %masked_scores : vector<4xf32> + %score_delta = vector.subf %masked_scores, %next_max : vector<4xf32> + %raw_probability = vector.expf %score_delta : vector<4xf32> + %probability = vector.select %query_valid, %raw_probability, %c0_f32x4 : vector<4xf32> + %old_delta = vector.subf %wmma_tail_max, %next_max : vector<4xf32> + %old_scale = vector.expf %old_delta : vector<4xf32> + %scaled_current_sum = vector.mulf %wmma_tail_sum, %old_scale : vector<4xf32> + %next_sum = vector.addf %scaled_current_sum, %probability : vector<4xf32> + %next_output0, %next_output1, %next_output2, %next_output3 = scf.if %lane_has_output -> (vector<4xf16>, vector<4xf16>, vector<4xf16>, vector<4xf16>) { + %output_packet_channel = index.assume %lane_output_channel [range(%lane_output_channel, 0, 124), mul(%lane_output_channel, 4)] : index + %value_channel = index.add %key_value_head_base, %output_packet_channel : index + %value_packet = vector.load %value_view[%single_tail_token, %value_channel] : view<[%bounded_key_value_token_count]x[%key_value_width]xf16> -> vector<4xf16> + %old_scale_f16 = vector.fptrunc %old_scale : vector<4xf32> to vector<4xf16> + %probability_f16 = vector.fptrunc %probability : vector<4xf32> to vector<4xf16> + %old_scale0_scalar = vector.extract %old_scale_f16[0] : vector<4xf16> -> f16 + %old_scale1_scalar = vector.extract %old_scale_f16[1] : vector<4xf16> -> f16 + %old_scale2_scalar = vector.extract %old_scale_f16[2] : vector<4xf16> -> f16 + %old_scale3_scalar = vector.extract %old_scale_f16[3] : vector<4xf16> -> f16 + %probability0_scalar = vector.extract %probability_f16[0] : vector<4xf16> -> f16 + %probability1_scalar = vector.extract %probability_f16[1] : vector<4xf16> -> f16 + %probability2_scalar = vector.extract %probability_f16[2] : vector<4xf16> -> f16 + %probability3_scalar = vector.extract %probability_f16[3] : vector<4xf16> -> f16 + %old_scale0 = vector.splat %old_scale0_scalar : vector<4xf16> + %old_scale1 = vector.splat %old_scale1_scalar : vector<4xf16> + %old_scale2 = vector.splat %old_scale2_scalar : vector<4xf16> + %old_scale3 = vector.splat %old_scale3_scalar : vector<4xf16> + %probability0 = vector.splat %probability0_scalar : vector<4xf16> + %probability1 = vector.splat %probability1_scalar : vector<4xf16> + %probability2 = vector.splat %probability2_scalar : vector<4xf16> + %probability3 = vector.splat %probability3_scalar : vector<4xf16> + %scaled_current_output0 = vector.mulf %wmma_tail_output0, %old_scale0 : vector<4xf16> + %scaled_current_output1 = vector.mulf %wmma_tail_output1, %old_scale1 : vector<4xf16> + %scaled_current_output2 = vector.mulf %wmma_tail_output2, %old_scale2 : vector<4xf16> + %scaled_current_output3 = vector.mulf %wmma_tail_output3, %old_scale3 : vector<4xf16> + %tail_output0 = vector.mulf %value_packet, %probability0 : vector<4xf16> + %tail_output1 = vector.mulf %value_packet, %probability1 : vector<4xf16> + %tail_output2 = vector.mulf %value_packet, %probability2 : vector<4xf16> + %tail_output3 = vector.mulf %value_packet, %probability3 : vector<4xf16> + %updated_output0 = vector.addf %scaled_current_output0, %tail_output0 : vector<4xf16> + %updated_output1 = vector.addf %scaled_current_output1, %tail_output1 : vector<4xf16> + %updated_output2 = vector.addf %scaled_current_output2, %tail_output2 : vector<4xf16> + %updated_output3 = vector.addf %scaled_current_output3, %tail_output3 : vector<4xf16> + scf.yield %updated_output0, %updated_output1, %updated_output2, %updated_output3 : vector<4xf16>, vector<4xf16>, vector<4xf16>, vector<4xf16> + } else { + scf.yield %wmma_tail_output0, %wmma_tail_output1, %wmma_tail_output2, %wmma_tail_output3 : vector<4xf16>, vector<4xf16>, vector<4xf16>, vector<4xf16> + } + scf.yield %next_max, %next_sum, %next_output0, %next_output1, %next_output2, %next_output3 : vector<4xf32>, vector<4xf32>, vector<4xf16>, vector<4xf16>, vector<4xf16>, vector<4xf16> + } else { + scf.yield %wmma_tail_max, %wmma_tail_sum, %wmma_tail_output0, %wmma_tail_output1, %wmma_tail_output2, %wmma_tail_output3 : vector<4xf32>, vector<4xf32>, vector<4xf16>, vector<4xf16>, vector<4xf16>, vector<4xf16> + } + // Normalize and publish the lane-owned packets. Query-tail rows never + // participate in the output store. + scf.if %lane_has_output { + %output_packet_channel = index.assume %lane_output_channel [range(%lane_output_channel, 0, 124), mul(%lane_output_channel, 4)] : index + %sum0_scalar = vector.extract %final_sum[0] : vector<4xf32> -> f32 + %sum1_scalar = vector.extract %final_sum[1] : vector<4xf32> -> f32 + %sum2_scalar = vector.extract %final_sum[2] : vector<4xf32> -> f32 + %sum3_scalar = vector.extract %final_sum[3] : vector<4xf32> -> f32 + %inverse_sum0_f32 = scalar.divf %c1_f32, %sum0_scalar : f32 + %inverse_sum1_f32 = scalar.divf %c1_f32, %sum1_scalar : f32 + %inverse_sum2_f32 = scalar.divf %c1_f32, %sum2_scalar : f32 + %inverse_sum3_f32 = scalar.divf %c1_f32, %sum3_scalar : f32 + %inverse_sum0_f16 = scalar.fptrunc %inverse_sum0_f32 : f32 to f16 + %inverse_sum1_f16 = scalar.fptrunc %inverse_sum1_f32 : f32 to f16 + %inverse_sum2_f16 = scalar.fptrunc %inverse_sum2_f32 : f32 to f16 + %inverse_sum3_f16 = scalar.fptrunc %inverse_sum3_f32 : f32 to f16 + %inverse_sum0 = vector.splat %inverse_sum0_f16 : vector<4xf16> + %inverse_sum1 = vector.splat %inverse_sum1_f16 : vector<4xf16> + %inverse_sum2 = vector.splat %inverse_sum2_f16 : vector<4xf16> + %inverse_sum3 = vector.splat %inverse_sum3_f16 : vector<4xf16> + %normalized0_f16 = vector.mulf %final_output0, %inverse_sum0 : vector<4xf16> + %normalized1_f16 = vector.mulf %final_output1, %inverse_sum1 : vector<4xf16> + %normalized2_f16 = vector.mulf %final_output2, %inverse_sum2 : vector<4xf16> + %normalized3_f16 = vector.mulf %final_output3, %inverse_sum3 : vector<4xf16> + %normalized0 = vector.extf %normalized0_f16 : vector<4xf16> to vector<4xf32> + %normalized1 = vector.extf %normalized1_f16 : vector<4xf16> to vector<4xf32> + %normalized2 = vector.extf %normalized2_f16 : vector<4xf16> to vector<4xf32> + %normalized3 = vector.extf %normalized3_f16 : vector<4xf16> to vector<4xf32> + scf.if %query_valid0 { + vector.store %normalized0, %output_view[%query_head0, %output_packet_channel] : vector<4xf32>, view<[%query_head_count]x128xf32> + } + scf.if %query_valid1 { + vector.store %normalized1, %output_view[%query_head1, %output_packet_channel] : vector<4xf32>, view<[%query_head_count]x128xf32> + } + scf.if %query_valid2 { + vector.store %normalized2, %output_view[%query_head2, %output_packet_channel] : vector<4xf32>, view<[%query_head_count]x128xf32> + } + scf.if %query_valid3 { + vector.store %normalized3, %output_view[%query_head3, %output_packet_channel] : vector<4xf32>, view<[%query_head_count]x128xf32> + } + } + kernel.return +} + +// The mask selects the first KV row exactly. QK, F16 probability conversion, +// and P*V all execute, while the expected result remains an auditable iota. +check.case public @qwen3_moe_flash_attention_decode_f32_f16_wmma_selected_row_case { + %key_value_token_count = check.literal value(64) : index + %query = check.generate.fill value(1.0) : tensor<1x128xf32> + %key = check.generate.fill value(1.0) : tensor<64x1x128xf16> + %value = check.generate.iota offset(0.0) step(0.125) : tensor<64x1x128xf16> + %mask = check.generate.iota offset(0.0) step(-10000.0) : tensor<64xf16> + %output = check.generate.fill value(-1.0) : tensor<1x128xf32> + %expected = check.generate.iota offset(0.0) step(0.125) : tensor<1x128xf32> + kernel.launch @qwen3_moe_flash_attention_decode_f32_f16_wmma[%key_value_token_count](%key_value_token_count, %query, %key, %value, %mask, %output) : [index](index, tensor<1x128xf32>, tensor<64x1x128xf16>, tensor<64x1x128xf16>, tensor<64xf16>, tensor<1x128xf32>) + check.expect.close actual(%output) expected(%expected) atol(0.001) rtol(0.001) nan(same) : tensor<1x128xf32> + check.return +} + +// Thirty-two query heads share four KV heads in the production GQA ratio. +// Equal scores and constant values make every output exactly two while all +// head-group addressing and output ownership execute. +check.case public @qwen3_moe_flash_attention_decode_f32_f16_wmma_gqa_case { + %key_value_token_count = check.literal value(128) : index + %query = check.generate.fill value(1.0) : tensor<32x128xf32> + %key = check.generate.fill value(1.0) : tensor<128x4x128xf16> + %value = check.generate.fill value(2.0) : tensor<128x4x128xf16> + %mask = check.generate.fill value(0.0) : tensor<128xf16> + %output = check.generate.fill value(-1.0) : tensor<32x128xf32> + %expected = check.generate.fill value(2.0) : tensor<32x128xf32> + kernel.launch @qwen3_moe_flash_attention_decode_f32_f16_wmma[%key_value_token_count](%key_value_token_count, %query, %key, %value, %mask, %output) : [index](index, tensor<32x128xf32>, tensor<128x4x128xf16>, tensor<128x4x128xf16>, tensor<128xf16>, tensor<32x128xf32>) + check.expect.close actual(%output) expected(%expected) atol(0.001) rtol(0.001) nan(same) : tensor<32x128xf32> + check.return +} + +// Sixty-five KV rows force one masked cleanup tile after a full WMMA block. +// The mask selects only that final row, whose iota values begin at 1024. +check.case public @qwen3_moe_flash_attention_decode_f32_f16_wmma_tail_case { + %key_value_token_count = check.literal value(65) : index + %query = check.generate.fill value(1.0) : tensor<1x128xf32> + %key = check.generate.fill value(1.0) : tensor<65x1x128xf16> + %value = check.generate.iota offset(0.0) step(0.125) : tensor<65x1x128xf16> + %mask = check.generate.iota offset(-64000.0) step(1000.0) : tensor<65xf16> + %output = check.generate.fill value(-1.0) : tensor<1x128xf32> + %expected = check.generate.iota offset(1024.0) step(0.125) : tensor<1x128xf32> + kernel.launch @qwen3_moe_flash_attention_decode_f32_f16_wmma[%key_value_token_count](%key_value_token_count, %query, %key, %value, %mask, %output) : [index](index, tensor<1x128xf32>, tensor<65x1x128xf16>, tensor<65x1x128xf16>, tensor<65xf16>, tensor<1x128xf32>) + check.expect.close actual(%output) expected(%expected) atol(0.001) rtol(0.001) nan(same) : tensor<1x128xf32> + check.return +} + +check.case public @qwen3_moe_flash_attention_decode_f32_f16_wmma_benchmark_case { + %key_value_token_count = check.param.choice values([64, 65, 128, 256, 512, 768, 1024, 1280, 2048]) name("key_value_token_count") : index + %query = check.generate.fill value(0.0) : tensor<32x128xf32> + %key = check.generate.fill value(0.0) : tensor<[%key_value_token_count]x4x128xf16> + %value = check.generate.fill value(0.0) : tensor<[%key_value_token_count]x4x128xf16> + %mask = check.generate.fill value(0.0) : tensor<[%key_value_token_count]xf16> + %output = check.generate.fill value(1.0) : tensor<32x128xf32> + %expected = check.generate.fill value(0.0) : tensor<32x128xf32> + kernel.launch @qwen3_moe_flash_attention_decode_f32_f16_wmma[%key_value_token_count](%key_value_token_count, %query, %key, %value, %mask, %output) : [index](index, tensor<32x128xf32>, tensor<[%key_value_token_count]x4x128xf16>, tensor<[%key_value_token_count]x4x128xf16>, tensor<[%key_value_token_count]xf16>, tensor<32x128xf32>) + check.expect.close actual(%output) expected(%expected) atol(0.0) rtol(0.0) nan(same) : tensor<32x128xf32> + check.return +} + +check.benchmark<@qwen3_moe_flash_attention_decode_f32_f16_wmma_selected_row_case> @qwen3_moe_flash_attention_decode_f32_f16_wmma_selected_row + +check.benchmark<@qwen3_moe_flash_attention_decode_f32_f16_wmma_gqa_case> @qwen3_moe_flash_attention_decode_f32_f16_wmma_gqa + +check.benchmark<@qwen3_moe_flash_attention_decode_f32_f16_wmma_tail_case> @qwen3_moe_flash_attention_decode_f32_f16_wmma_tail + +check.benchmark<@qwen3_moe_flash_attention_decode_f32_f16_wmma_benchmark_case> @qwen3_moe_flash_attention_decode_f32_f16_wmma_decode_64 {key_value_token_count = 64} + +check.benchmark<@qwen3_moe_flash_attention_decode_f32_f16_wmma_benchmark_case> @qwen3_moe_flash_attention_decode_f32_f16_wmma_decode_65 {key_value_token_count = 65} + +check.benchmark<@qwen3_moe_flash_attention_decode_f32_f16_wmma_benchmark_case> @qwen3_moe_flash_attention_decode_f32_f16_wmma_decode_128 {key_value_token_count = 128} + +check.benchmark<@qwen3_moe_flash_attention_decode_f32_f16_wmma_benchmark_case> @qwen3_moe_flash_attention_decode_f32_f16_wmma_decode_256 {key_value_token_count = 256} + +check.benchmark<@qwen3_moe_flash_attention_decode_f32_f16_wmma_benchmark_case> @qwen3_moe_flash_attention_decode_f32_f16_wmma_decode_512 {key_value_token_count = 512} + +check.benchmark<@qwen3_moe_flash_attention_decode_f32_f16_wmma_benchmark_case> @qwen3_moe_flash_attention_decode_f32_f16_wmma_decode_768 {key_value_token_count = 768} + +check.benchmark<@qwen3_moe_flash_attention_decode_f32_f16_wmma_benchmark_case> @qwen3_moe_flash_attention_decode_f32_f16_wmma_decode_1024 {key_value_token_count = 1024} + +check.benchmark<@qwen3_moe_flash_attention_decode_f32_f16_wmma_benchmark_case> @qwen3_moe_flash_attention_decode_f32_f16_wmma_decode_1280 {key_value_token_count = 1280} + +check.benchmark<@qwen3_moe_flash_attention_decode_f32_f16_wmma_benchmark_case> @qwen3_moe_flash_attention_decode_f32_f16_wmma_decode_2048 {key_value_token_count = 2048} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/flash_attention_decode_q128_f32_f16_wmma.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/flash_attention_decode_q128_f32_f16_wmma.loom new file mode 100644 index 000000000000..5a0c8037f856 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/flash_attention_decode_q128_f32_f16_wmma.loom @@ -0,0 +1,266 @@ +// Exact-q128 grouped-query decode FlashAttention. +// +// One eight-wave workgroup owns each KV head. The waves cooperatively cover +// all 128 KV rows for QK, normalize the complete score matrix in LDS, and then +// each own one 16-channel P*V tile. This preserves 32 active waves for the +// production 32Q/4KV shape while removing duplicate QK work, global partial +// tensors, completion atomics, and a separate resolve dispatch. +amdgpu.target @qwen3_moe_attention_decode_q128_gfx11_wave64 {subgroup_size = 64} + +config.decl @qwen3_moe.attention.query_head_count : %value: index where [range(%value, 1, 64)] + +config.decl @qwen3_moe.attention.key_value_head_count : %value: index where [range(%value, 1, 64)] + +kernel.def target(@qwen3_moe_attention_decode_q128_gfx11_wave64) @qwen3_moe_flash_attention_decode_q128_fused_f32_f16_wmma(%key_value_token_count: index) { + %key_value_head_count = config.get @qwen3_moe.attention.key_value_head_count : index + %c1 = index.constant 1 : index + %c512 = index.constant 512 : index + kernel.launch.config workgroups(%key_value_head_count, %c1, %c1) workgroup_size(%c512, %c1, %c1) : index +} launch(%key_value_token_count: index, %query: buffer, %key: buffer, %value: buffer, %mask: buffer, %output: buffer) { + %query_head_count = config.get @qwen3_moe.attention.query_head_count : index + %key_value_head_count = config.get @qwen3_moe.attention.key_value_head_count : index + %workgroup_x0 = kernel.workgroup.id : index + %key_value_head = index.assume %workgroup_x0 [range(%workgroup_x0, 0, 63)] : index + %workitem0 = kernel.workitem.id : index + %workitem = index.assume %workitem0 [range(%workitem0, 0, 511)] : index + %subgroup0 = kernel.subgroup.id : index + %subgroup = index.assume %subgroup0 [range(%subgroup0, 0, 7)] : index + %lane0 = kernel.subgroup.lane.id : index + %lane = index.assume %lane0 [range(%lane0, 0, 63)] : index + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c4 = index.constant 4 : index + %c16 = index.constant 16 : index + %c32 = index.constant 32 : index + %c64 = index.constant 64 : index + %c128 = index.constant 128 : index + %c512 = index.constant 512 : index + %c0_offset = index.constant 0 : offset + %query_stage_bytes = index.constant 4352 : offset + %score_stage_bytes = index.constant 12288 : offset + %probability_stage_bytes = index.constant 6144 : offset + %product_stage_bytes = index.constant 8192 : offset + %c0_f16 = scalar.constant 0.0 : f16 + %c0_f32 = scalar.constant 0.0 : f32 + %c1_f32 = scalar.constant 1.0 : f32 + %negative_large = scalar.constant -1e+30 : f32 + %c0_f32x2 = vector.constant 0.0 : vector<2xf32> + %m = index.constant 16 : index + %n = index.constant 16 : index + %k = index.constant 16 : index + %head_size_f32 = scalar.constant 128.0 : f32 + %attention_scale = scalar.rsqrtf %head_size_f32 : f32 + %query_heads_per_key_value_head0 = index.div %query_head_count, %key_value_head_count : index + %query_heads_per_key_value_head = index.assume %query_heads_per_key_value_head0 [range(%query_heads_per_key_value_head0, 1, 16)] : index + %query_head_base = index.mul %key_value_head, %query_heads_per_key_value_head : index + %key_value_width = index.mul %key_value_head_count, %c128 : index + %key_value_head_base = index.mul %key_value_head, %c128 : index + %score_key_origin = index.mul %subgroup, %c16 : index + %query_row_base = index.mul %subgroup, %c2 : index + %query_row0 = index.add %query_row_base, %c0 : index + %query_row1 = index.add %query_row_base, %c1 : index + %query_head0 = index.add %query_head_base, %query_row0 : index + %query_head1 = index.add %query_head_base, %query_row1 : index + %query_row_valid0 = index.cmp ult, %query_row0, %query_heads_per_key_value_head : index + %query_row_valid1 = index.cmp ult, %query_row1, %query_heads_per_key_value_head : index + %query_head_in_range0 = index.cmp ult, %query_head0, %query_head_count : index + %query_head_in_range1 = index.cmp ult, %query_head1, %query_head_count : index + %query_valid0 = scalar.andi %query_row_valid0, %query_head_in_range0 : i1 + %query_valid1 = scalar.andi %query_row_valid1, %query_head_in_range1 : i1 + %query_valid0x2 = vector.from_elements %query_valid0, %query_valid0 : vector<2xi1> + %query_valid1x2 = vector.from_elements %query_valid1, %query_valid1 : vector<2xi1> + %key_token0 = index.add %lane, %c0 : index + %key_token1 = index.add %lane, %c64 : index + %output_tile_channel = index.mul %subgroup, %c16 : index + %query_transposed_layout = encoding.layout.strided [1, 136] : encoding + %probability_transposed_layout = encoding.layout.strided [1, 24] : encoding + %query_noalias, %key_noalias, %value_noalias, %mask_noalias, %output_noalias = buffer.assume.noalias %query, %key, %value, %mask, %output : buffer, buffer, buffer, buffer, buffer + %query_aligned = buffer.assume.alignment %query_noalias {minimum_alignment = 16} : buffer + %key_aligned = buffer.assume.alignment %key_noalias {minimum_alignment = 16} : buffer + %value_aligned = buffer.assume.alignment %value_noalias {minimum_alignment = 16} : buffer + %mask_aligned = buffer.assume.alignment %mask_noalias {minimum_alignment = 16} : buffer + %output_aligned = buffer.assume.alignment %output_noalias {minimum_alignment = 16} : buffer + %query_view = buffer.view %query_aligned[%c0_offset] : buffer -> view<[%query_head_count]x128xf32> + %key_view = buffer.view %key_aligned[%c0_offset] : buffer -> view<128x[%key_value_width]xf16> + %value_view = buffer.view %value_aligned[%c0_offset] : buffer -> view<128x[%key_value_width]xf16> + %mask_view = buffer.view %mask_aligned[%c0_offset] : buffer -> view<128xf16> + %output_view = buffer.view %output_aligned[%c0_offset] : buffer -> view<[%query_head_count]x128xf32> + %query_stage = buffer.alloca align(16) %query_stage_bytes : buffer + %score_stage = buffer.alloca align(16) %score_stage_bytes : buffer + %probability_stage = buffer.alloca align(16) %probability_stage_bytes : buffer + %product_stage = buffer.alloca align(16) %product_stage_bytes : buffer + %query_stage_view = buffer.view %query_stage[%c0_offset] : buffer -> view<16x136xf16> + %query_transposed_view = buffer.view %query_stage[%c0_offset] : buffer -> view<128x16xf16, %query_transposed_layout> + %score_stage_view = buffer.view %score_stage[%c0_offset] : buffer -> view<128x24xf32> + %probability_stage_view = buffer.view %probability_stage[%c0_offset] : buffer -> view<16x128xf16, %probability_transposed_layout> + %product_stage_view = buffer.view %product_stage[%c0_offset] : buffer -> view<16x128xf32> + // Scale and transpose the 16-row Q tile once for all eight score waves. + scf.for %load_iteration = [%c0 to %c4 step %c1] unroll { + %linear = index.madd %load_iteration, %c512, %workitem : index + %local_query_row = index.div %linear, %c128 : index + %query_channel = index.rem %linear, %c128 : index + %local_query_head = index.add %query_head_base, %local_query_row : index + %local_query_row_valid = index.cmp ult, %local_query_row, %query_heads_per_key_value_head : index + %local_query_head_in_range = index.cmp ult, %local_query_head, %query_head_count : index + %local_query_valid = scalar.andi %local_query_row_valid, %local_query_head_in_range : i1 + %query_value = scf.if %local_query_valid -> (f16) { + %loaded = view.load %query_view[%local_query_head, %query_channel] : view<[%query_head_count]x128xf32> -> f32 + %scaled = scalar.mulf %loaded, %attention_scale : f32 + %truncated = scalar.fptrunc %scaled : f32 to f16 + scf.yield %truncated : f16 + } else { + scf.yield %c0_f16 : f16 + } + view.store %query_value, %query_stage_view[%local_query_row, %query_channel] : f16, view<16x136xf16> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + // Each wave produces one 16x16 score slice, covering all 128 KV rows. + %score_init_values = vector.constant 0.0 : vector<4xf32> + %score_init = vector.fragment %score_init_values shape [%m, %n] : vector<4xf32> + %score_fragment = scf.for %head_tile = [%c0 to %c128 step %c16](%score_accumulator = %score_init : vector<4xf32>) -> (vector<4xf32>) unroll { + %key_channel = index.add %key_value_head_base, %head_tile : index + %key_fragment = vector.fragment.load %key_view[%score_key_origin, %key_channel] shape [%m, %k] : view<128x[%key_value_width]xf16> -> vector<16xf16> + %query_fragment = vector.fragment.load %query_transposed_view[%head_tile, %c0] shape [%k, %n] : view<128x16xf16, %query_transposed_layout> -> vector<16xf16> + %next_score_accumulator = vector.mma %key_fragment, %query_fragment, %score_accumulator : vector<16xf16>, vector<16xf16>, vector<4xf32> + scf.yield %next_score_accumulator : vector<4xf32> + } + vector.fragment.store %score_fragment, %score_stage_view[%score_key_origin, %c0] shape [%m, %n] : vector<4xf32>, view<128x24xf32> + kernel.barrier scope(workgroup) ordering(acq_rel) + // Each wave normalizes two query rows. Every lane owns two KV columns, so a + // vector reduction first combines its pair and the subgroup reduction then + // spans the complete 128-token row. + %mask0_f16 = view.load %mask_view[%key_token0] : view<128xf16> -> f16 + %mask1_f16 = view.load %mask_view[%key_token1] : view<128xf16> -> f16 + %mask0 = scalar.extf %mask0_f16 : f16 to f32 + %mask1 = scalar.extf %mask1_f16 : f16 to f32 + %mask_pair = vector.from_elements %mask0, %mask1 : vector<2xf32> + %raw_score00 = view.load %score_stage_view[%key_token0, %query_row0] : view<128x24xf32> -> f32 + %raw_score01 = view.load %score_stage_view[%key_token1, %query_row0] : view<128x24xf32> -> f32 + %raw_score10 = view.load %score_stage_view[%key_token0, %query_row1] : view<128x24xf32> -> f32 + %raw_score11 = view.load %score_stage_view[%key_token1, %query_row1] : view<128x24xf32> -> f32 + %raw_scores0 = vector.from_elements %raw_score00, %raw_score01 : vector<2xf32> + %raw_scores1 = vector.from_elements %raw_score10, %raw_score11 : vector<2xf32> + %added_scores0 = vector.addf %raw_scores0, %mask_pair : vector<2xf32> + %added_scores1 = vector.addf %raw_scores1, %mask_pair : vector<2xf32> + %negative_pair = vector.constant -1e+30 : vector<2xf32> + %masked_scores0 = vector.select %query_valid0x2, %added_scores0, %negative_pair : vector<2xf32> + %masked_scores1 = vector.select %query_valid1x2, %added_scores1, %negative_pair : vector<2xf32> + %lane_max0 = vector.reduce %masked_scores0, %negative_large : vector<2xf32>, f32 + %lane_max1 = vector.reduce %masked_scores1, %negative_large : vector<2xf32>, f32 + %lane_maxima = vector.from_elements %lane_max0, %lane_max1 : vector<2xf32> + %row_maxima = kernel.subgroup.reduce %lane_maxima : vector<2xf32> + %row_maximum0 = vector.extract %row_maxima[0] : vector<2xf32> -> f32 + %row_maximum1 = vector.extract %row_maxima[1] : vector<2xf32> -> f32 + %row_maximum0x2 = vector.splat %row_maximum0 : vector<2xf32> + %row_maximum1x2 = vector.splat %row_maximum1 : vector<2xf32> + %score_delta0 = vector.subf %masked_scores0, %row_maximum0x2 : vector<2xf32> + %score_delta1 = vector.subf %masked_scores1, %row_maximum1x2 : vector<2xf32> + %raw_probability0 = vector.expf %score_delta0 : vector<2xf32> + %raw_probability1 = vector.expf %score_delta1 : vector<2xf32> + %probability0 = vector.select %query_valid0x2, %raw_probability0, %c0_f32x2 : vector<2xf32> + %probability1 = vector.select %query_valid1x2, %raw_probability1, %c0_f32x2 : vector<2xf32> + %lane_sum0 = vector.reduce %probability0, %c0_f32 : vector<2xf32>, f32 + %lane_sum1 = vector.reduce %probability1, %c0_f32 : vector<2xf32>, f32 + %lane_sums = vector.from_elements %lane_sum0, %lane_sum1 : vector<2xf32> + %row_sums = kernel.subgroup.reduce %lane_sums : vector<2xf32> + %row_sum0 = vector.extract %row_sums[0] : vector<2xf32> -> f32 + %row_sum1 = vector.extract %row_sums[1] : vector<2xf32> -> f32 + %reciprocal_sum0 = scf.if %query_valid0 -> (f32) { + %reciprocal = scalar.divf %c1_f32, %row_sum0 : f32 + scf.yield %reciprocal : f32 + } else { + scf.yield %c0_f32 : f32 + } + %reciprocal_sum1 = scf.if %query_valid1 -> (f32) { + %reciprocal = scalar.divf %c1_f32, %row_sum1 : f32 + scf.yield %reciprocal : f32 + } else { + scf.yield %c0_f32 : f32 + } + %reciprocal_sum0x2 = vector.splat %reciprocal_sum0 : vector<2xf32> + %reciprocal_sum1x2 = vector.splat %reciprocal_sum1 : vector<2xf32> + %normalized_probability0 = vector.mulf %probability0, %reciprocal_sum0x2 : vector<2xf32> + %normalized_probability1 = vector.mulf %probability1, %reciprocal_sum1x2 : vector<2xf32> + %probability0_f16 = vector.fptrunc %normalized_probability0 : vector<2xf32> to vector<2xf16> + %probability1_f16 = vector.fptrunc %normalized_probability1 : vector<2xf32> to vector<2xf16> + %probability00 = vector.extract %probability0_f16[0] : vector<2xf16> -> f16 + %probability01 = vector.extract %probability0_f16[1] : vector<2xf16> -> f16 + %probability10 = vector.extract %probability1_f16[0] : vector<2xf16> -> f16 + %probability11 = vector.extract %probability1_f16[1] : vector<2xf16> -> f16 + view.store %probability00, %probability_stage_view[%query_row0, %key_token0] : f16, view<16x128xf16, %probability_transposed_layout> + view.store %probability01, %probability_stage_view[%query_row0, %key_token1] : f16, view<16x128xf16, %probability_transposed_layout> + view.store %probability10, %probability_stage_view[%query_row1, %key_token0] : f16, view<16x128xf16, %probability_transposed_layout> + view.store %probability11, %probability_stage_view[%query_row1, %key_token1] : f16, view<16x128xf16, %probability_transposed_layout> + kernel.barrier scope(workgroup) ordering(acq_rel) + // Each wave owns one disjoint 16-channel output tile. + %product_init_values = vector.constant 0.0 : vector<4xf32> + %product_init = vector.fragment %product_init_values shape [%m, %n] : vector<4xf32> + %value_channel = index.add %key_value_head_base, %output_tile_channel : index + %product_fragment = scf.for %key_tile = [%c0 to %c128 step %c16](%product_accumulator = %product_init : vector<4xf32>) -> (vector<4xf32>) unroll { + %probability_fragment = vector.fragment.load %probability_stage_view[%c0, %key_tile] shape [%m, %k] : view<16x128xf16, %probability_transposed_layout> -> vector<16xf16> + %value_fragment = vector.fragment.load %value_view[%key_tile, %value_channel] shape [%k, %n] : view<128x[%key_value_width]xf16> -> vector<16xf16> + %next_product_accumulator = vector.mma %probability_fragment, %value_fragment, %product_accumulator : vector<16xf16>, vector<16xf16>, vector<4xf32> + scf.yield %next_product_accumulator : vector<4xf32> + } + vector.fragment.store %product_fragment, %product_stage_view[%c0, %output_tile_channel] shape [%m, %n] : vector<4xf32>, view<16x128xf32> + kernel.barrier scope(workgroup) ordering(acq_rel) + // One vector packet per workitem drains the 16x128 result tile while + // suppressing padded GQA rows. + %output_row = index.div %workitem, %c32 : index + %output_packet = index.rem %workitem, %c32 : index + %output_channel = index.mul %output_packet, %c4 : index + %output_head = index.add %query_head_base, %output_row : index + %output_row_valid = index.cmp ult, %output_row, %query_heads_per_key_value_head : index + %output_head_in_range = index.cmp ult, %output_head, %query_head_count : index + %output_valid = scalar.andi %output_row_valid, %output_head_in_range : i1 + scf.if %output_valid { + %output_packet_value = vector.load %product_stage_view[%output_row, %output_channel] : view<16x128xf32> -> vector<4xf32> + vector.store %output_packet_value, %output_view[%output_head, %output_channel] : vector<4xf32>, view<[%query_head_count]x128xf32> + } + kernel.return +} + +// A sharp mask selects the first KV row and makes the expected result an +// auditable iota while QK, global softmax, F16 probabilities, and P*V execute. +check.case public @qwen3_moe_flash_attention_decode_q128_fused_selected_row_case { + %key_value_token_count = check.literal value(128) : index + %query = check.generate.fill value(1.0) : tensor<1x128xf32> + %key = check.generate.fill value(1.0) : tensor<128x1x128xf16> + %value = check.generate.iota offset(0.0) step(0.125) : tensor<128x1x128xf16> + %mask = check.generate.iota offset(0.0) step(-10000.0) : tensor<128xf16> + %output = check.generate.fill value(-1.0) : tensor<1x128xf32> + %expected = check.generate.iota offset(0.0) step(0.125) : tensor<1x128xf32> + kernel.launch @qwen3_moe_flash_attention_decode_q128_fused_f32_f16_wmma[%key_value_token_count](%key_value_token_count, %query, %key, %value, %mask, %output) : [index](index, tensor<1x128xf32>, tensor<128x1x128xf16>, tensor<128x1x128xf16>, tensor<128xf16>, tensor<1x128xf32>) + check.expect.close actual(%output) expected(%expected) atol(0.001) rtol(0.001) nan(same) : tensor<1x128xf32> + check.return +} + +// The production 32Q/4KV GQA shape verifies all head groups and output tiles. +check.case public @qwen3_moe_flash_attention_decode_q128_fused_gqa_case { + %key_value_token_count = check.literal value(128) : index + %query = check.generate.fill value(1.0) : tensor<32x128xf32> + %key = check.generate.fill value(1.0) : tensor<128x4x128xf16> + %value = check.generate.fill value(2.0) : tensor<128x4x128xf16> + %mask = check.generate.fill value(0.0) : tensor<128xf16> + %output = check.generate.fill value(-1.0) : tensor<32x128xf32> + %expected = check.generate.fill value(2.0) : tensor<32x128xf32> + kernel.launch @qwen3_moe_flash_attention_decode_q128_fused_f32_f16_wmma[%key_value_token_count](%key_value_token_count, %query, %key, %value, %mask, %output) : [index](index, tensor<32x128xf32>, tensor<128x4x128xf16>, tensor<128x4x128xf16>, tensor<128xf16>, tensor<32x128xf32>) + check.expect.close actual(%output) expected(%expected) atol(0.001) rtol(0.001) nan(same) : tensor<32x128xf32> + check.return +} + +check.case public @qwen3_moe_flash_attention_decode_q128_fused_benchmark_case { + %key_value_token_count = check.literal value(128) : index + %query = check.generate.fill value(0.0) : tensor<32x128xf32> + %key = check.generate.fill value(0.0) : tensor<128x4x128xf16> + %value = check.generate.fill value(0.0) : tensor<128x4x128xf16> + %mask = check.generate.fill value(0.0) : tensor<128xf16> + %output = check.generate.fill value(1.0) : tensor<32x128xf32> + %expected = check.generate.fill value(0.0) : tensor<32x128xf32> + kernel.launch @qwen3_moe_flash_attention_decode_q128_fused_f32_f16_wmma[%key_value_token_count](%key_value_token_count, %query, %key, %value, %mask, %output) : [index](index, tensor<32x128xf32>, tensor<128x4x128xf16>, tensor<128x4x128xf16>, tensor<128xf16>, tensor<32x128xf32>) + check.expect.close actual(%output) expected(%expected) atol(0.0) rtol(0.0) nan(same) : tensor<32x128xf32> + check.return +} + +check.benchmark<@qwen3_moe_flash_attention_decode_q128_fused_benchmark_case> @qwen3_moe_flash_attention_decode_q128_fused diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/flash_attention_decode_split_f32_f16_wmma.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/flash_attention_decode_split_f32_f16_wmma.loom new file mode 100644 index 000000000000..22cb0b7fdaa3 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/flash_attention_decode_split_f32_f16_wmma.loom @@ -0,0 +1,1214 @@ +// Qwen3 MoE grouped-query decode FlashAttention. +// +// Each workgroup processes one 64-token KV block for all GQA query heads that +// share a KV head. The workgroups publish online-softmax state, then the last +// arrival folds every block and resets the per-KV-head completion counter for +// the next invocation. Packing GQA heads removes redundant K/V traffic while +// split-K preserves enough parallelism for decode without another dispatch. +amdgpu.target @qwen3_moe_decode_split_gfx11_wave64 {subgroup_size = 64} + +config.decl @qwen3_moe.attention.query_head_count : %value: index where [range(%value, 1, 64)] + +config.decl @qwen3_moe.attention.key_value_head_count : %value: index where [range(%value, 1, 64)] + +// Maximum K/V storage capacity available to the compiled kernel. +config.decl @qwen3_moe.attention.key_value_token_capacity : %value: index where [range(%value, 64, 32768)] + +// Computes one active online-softmax partial for a 64-row KV block. Both the +// fused short-context export and the two-dispatch long-context export reach +// this body through the block-classifying producer below, keeping their +// different binding contracts honest without duplicating the attention math. +// Partial storage retains its capacity-specialized block-axis stride while the +// launched block count bounds issue-time work. Keeping those values distinct +// preserves constant address arithmetic across changing visible prefixes. +func.def inline @qwen3_moe_flash_attention_decode_split_produce_active_partials_body_f32_f16_wmma(%key_value_token_count: index, %partial_block_capacity0: index, %launched_block_count0: index, %query: buffer, %key: buffer, %value: buffer, %lane_mask: f32, %partial_max: buffer, %partial_sum: buffer, %partial_output: buffer) { + %bounded_key_value_token_count = index.assume %key_value_token_count [range(%key_value_token_count, 1, 32768)] : index + %partial_block_capacity, %launched_block_count = index.assume %partial_block_capacity0, %launched_block_count0 [range(%partial_block_capacity0, 1, 512), range(%launched_block_count0, 1, 512), le(%launched_block_count0, %partial_block_capacity0)] : index, index + %query_head_count = config.get @qwen3_moe.attention.query_head_count : index + %key_value_head_count = config.get @qwen3_moe.attention.key_value_head_count : index + %workgroup_x0 = kernel.workgroup.id : index + %workgroup_y0 = kernel.workgroup.id : index + %workgroup_x_in_launch = index.assume %workgroup_x0 [range(%workgroup_x0, 0, 511)] : index + %workgroup_y = index.assume %workgroup_y0 [range(%workgroup_y0, 0, 63)] : index + %workitem = kernel.workitem.id : index + %subgroup0 = kernel.subgroup.id : index + %subgroup = index.assume %subgroup0 [range(%subgroup0, 0, 3)] : index + %lane = kernel.subgroup.lane.id : index + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c3 = index.constant 3 : index + %c4 = index.constant 4 : index + %c8 = index.constant 8 : index + %c16 = index.constant 16 : index + %c32 = index.constant 32 : index + %c63 = index.constant 63 : index + %c64 = index.constant 64 : index + %c128 = index.constant 128 : index + %c256 = index.constant 256 : index + %c0_offset = index.constant 0 : offset + %query_stage_bytes = index.constant 4352 : offset + %score_stage_bytes = index.constant 6144 : offset + %probability_stage_bytes = index.constant 3072 : offset + %product_stage_bytes = index.constant 4096 : offset + %c0_f32 = scalar.constant 0.0 : f32 + %negative_large = scalar.constant -1e+30 : f32 + %c0_i32 = scalar.constant 0 : i32 + %c1_i32 = scalar.constant 1 : i32 + %head_size_f32 = scalar.constant 128.0 : f32 + %attention_scale = scalar.rsqrtf %head_size_f32 : f32 + %c0_f16 = scalar.constant 0.0 : f16 + %c0_f16x8 = vector.constant 0.0 : vector<8xf16> + %c0_f32x4 = vector.constant 0.0 : vector<4xf32> + %m = index.constant 16 : index + %n = index.constant 16 : index + %k = index.constant 16 : index + %workgroup_x, %launch_partial_block_capacity, %launch_launched_block_count = index.assume %workgroup_x_in_launch, %partial_block_capacity, %launched_block_count [lt(%workgroup_x_in_launch, %launched_block_count)] : index, index, index + %tail_key_value_token_count = index.rem %bounded_key_value_token_count, %c64 : index + %has_no_tail = index.cmp eq, %tail_key_value_token_count, %c0 : index + %active_padded_key_value_token_count = index.add %bounded_key_value_token_count, %c63 : index + %active_key_value_block_count = index.div %active_padded_key_value_token_count, %c64 : index + %last_block_ordinal = index.sub %active_key_value_block_count, %c1 : index + %block_ordinal = index.add %workgroup_x, %c0 : index + %is_not_last_block = index.cmp ne, %block_ordinal, %last_block_ordinal : index + %is_full_block = scalar.ori %has_no_tail, %is_not_last_block : i1 + %key_value_head = index.add %workgroup_y, %c0 : index + %query_heads_per_key_value_head0 = index.div %query_head_count, %key_value_head_count : index + %query_heads_per_key_value_head = index.assume %query_heads_per_key_value_head0 [range(%query_heads_per_key_value_head0, 1, 16)] : index + %query_head_base = index.mul %key_value_head, %query_heads_per_key_value_head : index + %key_value_width = index.mul %key_value_head_count, %c128 : index + %key_value_head_base = index.mul %key_value_head, %c128 : index + %key_origin = index.mul %block_ordinal, %c64 : index + %subgroup_score_column = index.mul %subgroup, %c16 : index + %subgroup_query_row = index.mul %subgroup, %c4 : index + %query_row0 = index.add %subgroup_query_row, %c0 : index + %query_row1 = index.add %subgroup_query_row, %c1 : index + %query_row2 = index.add %subgroup_query_row, %c2 : index + %query_row3 = index.add %subgroup_query_row, %c3 : index + %query_head0 = index.add %query_head_base, %query_row0 : index + %query_head1 = index.add %query_head_base, %query_row1 : index + %query_head2 = index.add %query_head_base, %query_row2 : index + %query_head3 = index.add %query_head_base, %query_row3 : index + %query_head_valid0 = index.cmp ult, %query_head0, %query_head_count : index + %query_head_valid1 = index.cmp ult, %query_head1, %query_head_count : index + %query_head_valid2 = index.cmp ult, %query_head2, %query_head_count : index + %query_head_valid3 = index.cmp ult, %query_head3, %query_head_count : index + %query_valid = vector.from_elements %query_head_valid0, %query_head_valid1, %query_head_valid2, %query_head_valid3 : vector<4xi1> + %subgroup_output_channel = index.mul %subgroup, %c32 : index + %subgroup_output_channel1 = index.add %subgroup_output_channel, %c16 : index + %lane_output_channel = index.mul %lane, %c4 : index + %lane_has_output = index.cmp ult, %lane, %c32 : index + %lane_is_zero = index.cmp eq, %lane, %c0 : index + %query_transposed_layout = encoding.layout.strided [1, 136] : encoding + %probability_transposed_layout = encoding.layout.strided [1, 24] : encoding + %query_noalias, %key_noalias, %value_noalias, %partial_max_noalias, %partial_sum_noalias, %partial_output_noalias = buffer.assume.noalias %query, %key, %value, %partial_max, %partial_sum, %partial_output : buffer, buffer, buffer, buffer, buffer, buffer + %query_aligned = buffer.assume.alignment %query_noalias {minimum_alignment = 16} : buffer + %key_aligned = buffer.assume.alignment %key_noalias {minimum_alignment = 16} : buffer + %value_aligned = buffer.assume.alignment %value_noalias {minimum_alignment = 16} : buffer + %partial_max_aligned = buffer.assume.alignment %partial_max_noalias {minimum_alignment = 16} : buffer + %partial_sum_aligned = buffer.assume.alignment %partial_sum_noalias {minimum_alignment = 16} : buffer + %partial_output_aligned = buffer.assume.alignment %partial_output_noalias {minimum_alignment = 16} : buffer + %query_view = buffer.view %query_aligned[%c0_offset] : buffer -> view<[%query_head_count]x128xf32> + %key_view = buffer.view %key_aligned[%c0_offset] : buffer -> view<[%bounded_key_value_token_count]x[%key_value_width]xf16> + %value_view = buffer.view %value_aligned[%c0_offset] : buffer -> view<[%bounded_key_value_token_count]x[%key_value_width]xf16> + %partial_max_view = buffer.view %partial_max_aligned[%c0_offset] : buffer -> view<[%key_value_head_count]x[%launch_partial_block_capacity]x16xf32> + %partial_sum_view = buffer.view %partial_sum_aligned[%c0_offset] : buffer -> view<[%key_value_head_count]x[%launch_partial_block_capacity]x16xf32> + %partial_output_view = buffer.view %partial_output_aligned[%c0_offset] : buffer -> view<[%key_value_head_count]x[%launch_partial_block_capacity]x16x128xf16> + %query_stage = buffer.alloca align(16) %query_stage_bytes : buffer + %score_stage = buffer.alloca align(16) %score_stage_bytes : buffer + %probability_stage = buffer.alloca align(16) %probability_stage_bytes : buffer + %product_stage = buffer.alloca align(16) %product_stage_bytes : buffer + %query_stage_view = buffer.view %query_stage[%c0_offset] : buffer -> view<16x136xf16> + %query_transposed_view = buffer.view %query_stage[%c0_offset] : buffer -> view<128x16xf16, %query_transposed_layout> + %score_stage_view = buffer.view %score_stage[%c0_offset] : buffer -> view<64x24xf32> + %probability_stage_view = buffer.view %probability_stage[%c0_offset] : buffer -> view<16x64xf16, %probability_transposed_layout> + %product_stage_view = buffer.view %product_stage[%c0_offset] : buffer -> view<16x128xf16> + %tail_key_stage_view = buffer.view %product_stage[%c0_offset] : buffer -> view<64x16xf16> + %tail_value_stage_view = buffer.view %score_stage[%c0_offset] : buffer -> view<16x128xf16> + scf.for %load_iteration = [%c0 to %c8 step %c1] unroll { + %linear = index.madd %load_iteration, %c256, %workitem : index + %local_query_row = index.div %linear, %c128 : index + %query_channel = index.rem %linear, %c128 : index + %local_query_head = index.add %query_head_base, %local_query_row : index + %local_query_head_valid = index.cmp ult, %local_query_head, %query_head_count : index + %query_value = scf.if %local_query_head_valid -> (f16) { + %loaded = view.load %query_view[%local_query_head, %query_channel] : view<[%query_head_count]x128xf32> -> f32 + %scaled = scalar.mulf %loaded, %attention_scale : f32 + %truncated = scalar.fptrunc %scaled : f32 to f16 + scf.yield %truncated : f16 + } else { + scf.yield %c0_f16 : f16 + } + view.store %query_value, %query_stage_view[%local_query_row, %query_channel] : f16, view<16x136xf16> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + %score_init_values = vector.constant 0.0 : vector<4xf32> + %score_init = vector.fragment %score_init_values shape [%m, %n] : vector<4xf32> + %score_fragment = scf.if %is_full_block -> (vector<4xf32>) { + %full_score_fragment = scf.for %head_tile = [%c0 to %c128 step %c16](%score_accumulator = %score_init : vector<4xf32>) -> (vector<4xf32>) unroll { + // Bound this view by the proven tile end so the full vector footprint is + // visible without treating physical tail padding as logical storage. + %score_key_origin0 = index.add %key_origin, %subgroup_score_column : index + %score_key_end0 = index.add %score_key_origin0, %c16 : index + %score_key_origin, %score_key_end = index.assume %score_key_origin0, %score_key_end0 [le(%score_key_end0, %bounded_key_value_token_count)] : index, index + %full_key_view = buffer.view %key_aligned[%c0_offset] : buffer -> view<[%score_key_end]x[%key_value_width]xf16> + %key_channel = index.add %key_value_head_base, %head_tile : index + %key_fragment = vector.fragment.load %full_key_view[%score_key_origin, %key_channel] shape [%m, %k] : view<[%score_key_end]x[%key_value_width]xf16> -> vector<16xf16> + %query_fragment = vector.fragment.load %query_transposed_view[%head_tile, %c0] shape [%k, %n] : view<128x16xf16, %query_transposed_layout> -> vector<16xf16> + %next_score_accumulator = vector.mma %key_fragment, %query_fragment, %score_accumulator : vector<16xf16>, vector<16xf16>, vector<4xf32> + scf.yield %next_score_accumulator : vector<4xf32> + } + scf.yield %full_score_fragment : vector<4xf32> + } else { + // Stage one 64x16 K panel at a time. Every physical load is guarded, so + // callers need no initialized padding beyond the logical KV length. + %tail_score_fragment = scf.for %head_tile = [%c0 to %c128 step %c16](%score_accumulator = %score_init : vector<4xf32>) -> (vector<4xf32>) unroll { + scf.for %load_iteration = [%c0 to %c4 step %c1] unroll { + %linear = index.madd %load_iteration, %c256, %workitem : index + %key_row = index.div %linear, %c16 : index + %head_channel = index.rem %linear, %c16 : index + %key_token = index.add %key_origin, %key_row : index + %key_valid = index.cmp ult, %key_token, %bounded_key_value_token_count : index + %key_value = scf.if %key_valid -> (f16) { + %head_channel0 = index.add %head_tile, %head_channel : index + %key_channel = index.add %key_value_head_base, %head_channel0 : index + %loaded = view.load %key_view[%key_token, %key_channel] : view<[%bounded_key_value_token_count]x[%key_value_width]xf16> -> f16 + scf.yield %loaded : f16 + } else { + scf.yield %c0_f16 : f16 + } + view.store %key_value, %tail_key_stage_view[%key_row, %head_channel] : f16, view<64x16xf16> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + %key_fragment = vector.fragment.load %tail_key_stage_view[%subgroup_score_column, %c0] shape [%m, %k] : view<64x16xf16> -> vector<16xf16> + %query_fragment = vector.fragment.load %query_transposed_view[%head_tile, %c0] shape [%k, %n] : view<128x16xf16, %query_transposed_layout> -> vector<16xf16> + %next_score_accumulator = vector.mma %key_fragment, %query_fragment, %score_accumulator : vector<16xf16>, vector<16xf16>, vector<4xf32> + kernel.barrier scope(workgroup) ordering(acq_rel) + scf.yield %next_score_accumulator : vector<4xf32> + } + scf.yield %tail_score_fragment : vector<4xf32> + } + vector.fragment.store %score_fragment, %score_stage_view[%subgroup_score_column, %c0] shape [%m, %n] : vector<4xf32>, view<64x24xf32> + kernel.barrier scope(workgroup) ordering(acq_rel) + %key_token = index.add %key_origin, %lane : index + %raw_score0 = view.load %score_stage_view[%lane, %query_row0] : view<64x24xf32> -> f32 + %raw_score1 = view.load %score_stage_view[%lane, %query_row1] : view<64x24xf32> -> f32 + %raw_score2 = view.load %score_stage_view[%lane, %query_row2] : view<64x24xf32> -> f32 + %raw_score3 = view.load %score_stage_view[%lane, %query_row3] : view<64x24xf32> -> f32 + %key_valid = index.cmp ult, %key_token, %bounded_key_value_token_count : index + %score_valid0 = scalar.andi %query_head_valid0, %key_valid : i1 + %score_valid1 = scalar.andi %query_head_valid1, %key_valid : i1 + %score_valid2 = scalar.andi %query_head_valid2, %key_valid : i1 + %score_valid3 = scalar.andi %query_head_valid3, %key_valid : i1 + %masked_score0 = scf.if %score_valid0 -> (f32) { + %score = scalar.addf %raw_score0, %lane_mask : f32 + scf.yield %score : f32 + } else { + scf.yield %negative_large : f32 + } + %masked_score1 = scf.if %score_valid1 -> (f32) { + %score = scalar.addf %raw_score1, %lane_mask : f32 + scf.yield %score : f32 + } else { + scf.yield %negative_large : f32 + } + %masked_score2 = scf.if %score_valid2 -> (f32) { + %score = scalar.addf %raw_score2, %lane_mask : f32 + scf.yield %score : f32 + } else { + scf.yield %negative_large : f32 + } + %masked_score3 = scf.if %score_valid3 -> (f32) { + %score = scalar.addf %raw_score3, %lane_mask : f32 + scf.yield %score : f32 + } else { + scf.yield %negative_large : f32 + } + %masked_scores = vector.from_elements %masked_score0, %masked_score1, %masked_score2, %masked_score3 : vector<4xf32> + %block_max = kernel.subgroup.reduce %masked_scores : vector<4xf32> + %score_delta = vector.subf %masked_scores, %block_max : vector<4xf32> + %raw_probability = vector.expf %score_delta : vector<4xf32> + %probability = vector.select %query_valid, %raw_probability, %c0_f32x4 : vector<4xf32> + %block_sum = kernel.subgroup.reduce %probability : vector<4xf32> + %probability_f16 = vector.fptrunc %probability : vector<4xf32> to vector<4xf16> + %probability0 = vector.extract %probability_f16[0] : vector<4xf16> -> f16 + %probability1 = vector.extract %probability_f16[1] : vector<4xf16> -> f16 + %probability2 = vector.extract %probability_f16[2] : vector<4xf16> -> f16 + %probability3 = vector.extract %probability_f16[3] : vector<4xf16> -> f16 + view.store %probability0, %probability_stage_view[%query_row0, %lane] : f16, view<16x64xf16, %probability_transposed_layout> + view.store %probability1, %probability_stage_view[%query_row1, %lane] : f16, view<16x64xf16, %probability_transposed_layout> + view.store %probability2, %probability_stage_view[%query_row2, %lane] : f16, view<16x64xf16, %probability_transposed_layout> + view.store %probability3, %probability_stage_view[%query_row3, %lane] : f16, view<16x64xf16, %probability_transposed_layout> + kernel.barrier scope(workgroup) ordering(acq_rel) + // Keep both P*V halves live so the scheduler can interleave their independent + // matrix chains. F16 accumulation and exchange preserve the Vulkan CM1 + // contract while halving the product LDS footprint. + %product_init0 = vector.fragment %c0_f16x8 shape [%m, %n] : vector<8xf16> + %product_init1 = vector.fragment %c0_f16x8 shape [%m, %n] : vector<8xf16> + %value_channel0 = index.add %key_value_head_base, %subgroup_output_channel : index + %value_channel1 = index.add %key_value_head_base, %subgroup_output_channel1 : index + %product_fragment0, %product_fragment1 = scf.if %is_full_block -> (vector<8xf16>, vector<8xf16>) { + %full_product_fragment0, %full_product_fragment1 = scf.for %key_tile = [%c0 to %c64 step %c16](%product_accumulator0 = %product_init0 : vector<8xf16>, %product_accumulator1 = %product_init1 : vector<8xf16>) -> (vector<8xf16>, vector<8xf16>) unroll { + // Retain the same logical-extent contract for the full P*V tile. + %value_token0 = index.add %key_origin, %key_tile : index + %value_token_end0 = index.add %value_token0, %c16 : index + %value_token, %value_token_end = index.assume %value_token0, %value_token_end0 [le(%value_token_end0, %bounded_key_value_token_count)] : index, index + %full_value_view = buffer.view %value_aligned[%c0_offset] : buffer -> view<[%value_token_end]x[%key_value_width]xf16> + %probability_fragment = vector.fragment.load %probability_stage_view[%c0, %key_tile] shape [%m, %k] : view<16x64xf16, %probability_transposed_layout> -> vector<16xf16> + %value_fragment0 = vector.fragment.load %full_value_view[%value_token, %value_channel0] shape [%k, %n] : view<[%value_token_end]x[%key_value_width]xf16> -> vector<16xf16> + %value_fragment1 = vector.fragment.load %full_value_view[%value_token, %value_channel1] shape [%k, %n] : view<[%value_token_end]x[%key_value_width]xf16> -> vector<16xf16> + %next_product_accumulator0 = vector.mma %probability_fragment, %value_fragment0, %product_accumulator0 : vector<16xf16>, vector<16xf16>, vector<8xf16> + %next_product_accumulator1 = vector.mma %probability_fragment, %value_fragment1, %product_accumulator1 : vector<16xf16>, vector<16xf16>, vector<8xf16> + scf.yield %next_product_accumulator0, %next_product_accumulator1 : vector<8xf16>, vector<8xf16> + } + scf.yield %full_product_fragment0, %full_product_fragment1 : vector<8xf16>, vector<8xf16> + } else { + // Stage one 16x128 V panel at a time and zero every lane beyond the + // logical KV length before the panel participates in P*V. + %tail_product_fragment0, %tail_product_fragment1 = scf.for %key_tile = [%c0 to %c64 step %c16](%product_accumulator0 = %product_init0 : vector<8xf16>, %product_accumulator1 = %product_init1 : vector<8xf16>) -> (vector<8xf16>, vector<8xf16>) unroll { + scf.for %load_iteration = [%c0 to %c8 step %c1] unroll { + %linear = index.madd %load_iteration, %c256, %workitem : index + %key_row = index.div %linear, %c128 : index + %value_channel = index.rem %linear, %c128 : index + %value_token0 = index.add %key_origin, %key_tile : index + %value_token = index.add %value_token0, %key_row : index + %value_valid = index.cmp ult, %value_token, %bounded_key_value_token_count : index + %value_element = scf.if %value_valid -> (f16) { + %global_value_channel = index.add %key_value_head_base, %value_channel : index + %loaded = view.load %value_view[%value_token, %global_value_channel] : view<[%bounded_key_value_token_count]x[%key_value_width]xf16> -> f16 + scf.yield %loaded : f16 + } else { + scf.yield %c0_f16 : f16 + } + view.store %value_element, %tail_value_stage_view[%key_row, %value_channel] : f16, view<16x128xf16> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + %probability_fragment = vector.fragment.load %probability_stage_view[%c0, %key_tile] shape [%m, %k] : view<16x64xf16, %probability_transposed_layout> -> vector<16xf16> + %value_fragment0 = vector.fragment.load %tail_value_stage_view[%c0, %subgroup_output_channel] shape [%k, %n] : view<16x128xf16> -> vector<16xf16> + %value_fragment1 = vector.fragment.load %tail_value_stage_view[%c0, %subgroup_output_channel1] shape [%k, %n] : view<16x128xf16> -> vector<16xf16> + %next_product_accumulator0 = vector.mma %probability_fragment, %value_fragment0, %product_accumulator0 : vector<16xf16>, vector<16xf16>, vector<8xf16> + %next_product_accumulator1 = vector.mma %probability_fragment, %value_fragment1, %product_accumulator1 : vector<16xf16>, vector<16xf16>, vector<8xf16> + kernel.barrier scope(workgroup) ordering(acq_rel) + scf.yield %next_product_accumulator0, %next_product_accumulator1 : vector<8xf16>, vector<8xf16> + } + scf.yield %tail_product_fragment0, %tail_product_fragment1 : vector<8xf16>, vector<8xf16> + } + vector.fragment.store %product_fragment0, %product_stage_view[%c0, %subgroup_output_channel] shape [%m, %n] : vector<8xf16>, view<16x128xf16> + vector.fragment.store %product_fragment1, %product_stage_view[%c0, %subgroup_output_channel1] shape [%m, %n] : vector<8xf16>, view<16x128xf16> + kernel.barrier scope(workgroup) ordering(acq_rel) + scf.if %lane_has_output { + %block_output0 = vector.load %product_stage_view[%query_row0, %lane_output_channel] : view<16x128xf16> -> vector<4xf16> + %block_output1 = vector.load %product_stage_view[%query_row1, %lane_output_channel] : view<16x128xf16> -> vector<4xf16> + %block_output2 = vector.load %product_stage_view[%query_row2, %lane_output_channel] : view<16x128xf16> -> vector<4xf16> + %block_output3 = vector.load %product_stage_view[%query_row3, %lane_output_channel] : view<16x128xf16> -> vector<4xf16> + scf.if %query_head_valid0 { + vector.store %block_output0, %partial_output_view[%key_value_head, %block_ordinal, %query_row0, %lane_output_channel] : vector<4xf16>, view<[%key_value_head_count]x[%launch_partial_block_capacity]x16x128xf16> + } + scf.if %query_head_valid1 { + vector.store %block_output1, %partial_output_view[%key_value_head, %block_ordinal, %query_row1, %lane_output_channel] : vector<4xf16>, view<[%key_value_head_count]x[%launch_partial_block_capacity]x16x128xf16> + } + scf.if %query_head_valid2 { + vector.store %block_output2, %partial_output_view[%key_value_head, %block_ordinal, %query_row2, %lane_output_channel] : vector<4xf16>, view<[%key_value_head_count]x[%launch_partial_block_capacity]x16x128xf16> + } + scf.if %query_head_valid3 { + vector.store %block_output3, %partial_output_view[%key_value_head, %block_ordinal, %query_row3, %lane_output_channel] : vector<4xf16>, view<[%key_value_head_count]x[%launch_partial_block_capacity]x16x128xf16> + } + } + scf.if %lane_is_zero { + scf.if %query_head_valid0 { + %maximum = vector.extract %block_max[0] : vector<4xf32> -> f32 + %sum = vector.extract %block_sum[0] : vector<4xf32> -> f32 + view.store %maximum, %partial_max_view[%key_value_head, %block_ordinal, %query_row0] : f32, view<[%key_value_head_count]x[%launch_partial_block_capacity]x16xf32> + view.store %sum, %partial_sum_view[%key_value_head, %block_ordinal, %query_row0] : f32, view<[%key_value_head_count]x[%launch_partial_block_capacity]x16xf32> + } + scf.if %query_head_valid1 { + %maximum = vector.extract %block_max[1] : vector<4xf32> -> f32 + %sum = vector.extract %block_sum[1] : vector<4xf32> -> f32 + view.store %maximum, %partial_max_view[%key_value_head, %block_ordinal, %query_row1] : f32, view<[%key_value_head_count]x[%launch_partial_block_capacity]x16xf32> + view.store %sum, %partial_sum_view[%key_value_head, %block_ordinal, %query_row1] : f32, view<[%key_value_head_count]x[%launch_partial_block_capacity]x16xf32> + } + scf.if %query_head_valid2 { + %maximum = vector.extract %block_max[2] : vector<4xf32> -> f32 + %sum = vector.extract %block_sum[2] : vector<4xf32> -> f32 + view.store %maximum, %partial_max_view[%key_value_head, %block_ordinal, %query_row2] : f32, view<[%key_value_head_count]x[%launch_partial_block_capacity]x16xf32> + view.store %sum, %partial_sum_view[%key_value_head, %block_ordinal, %query_row2] : f32, view<[%key_value_head_count]x[%launch_partial_block_capacity]x16xf32> + } + scf.if %query_head_valid3 { + %maximum = vector.extract %block_max[3] : vector<4xf32> -> f32 + %sum = vector.extract %block_sum[3] : vector<4xf32> -> f32 + view.store %maximum, %partial_max_view[%key_value_head, %block_ordinal, %query_row3] : f32, view<[%key_value_head_count]x[%launch_partial_block_capacity]x16xf32> + view.store %sum, %partial_sum_view[%key_value_head, %block_ordinal, %query_row3] : f32, view<[%key_value_head_count]x[%launch_partial_block_capacity]x16xf32> + } + } + func.return +} + +// Fully masked splits publish the online-softmax identity without running QK +// or P*V. Valid additive masks contain finite F16 values or negative infinity, +// and every finite F16 value is greater than -1e30. Comparing after extension +// therefore distinguishes active rows from masked rows exactly. +func.def inline @qwen3_moe_flash_attention_decode_split_produce_partials_body_f32_f16_wmma(%key_value_token_count: index, %query: buffer, %key: buffer, %value: buffer, %mask: buffer, %partial_max: buffer, %partial_sum: buffer, %partial_output: buffer) { + %bounded_key_value_token_count = index.assume %key_value_token_count [range(%key_value_token_count, 1, 32768)] : index + %query_head_count = config.get @qwen3_moe.attention.query_head_count : index + %key_value_head_count = config.get @qwen3_moe.attention.key_value_head_count : index + %workgroup_x0 = kernel.workgroup.id : index + %workgroup_y0 = kernel.workgroup.id : index + %workgroup_x_in_launch = index.assume %workgroup_x0 [range(%workgroup_x0, 0, 511)] : index + %workgroup_y = index.assume %workgroup_y0 [range(%workgroup_y0, 0, 63)] : index + %subgroup0 = kernel.subgroup.id : index + %subgroup = index.assume %subgroup0 [range(%subgroup0, 0, 3)] : index + %lane = kernel.subgroup.lane.id : index + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c3 = index.constant 3 : index + %c4 = index.constant 4 : index + %c32 = index.constant 32 : index + %c64 = index.constant 64 : index + %c0_offset = index.constant 0 : offset + %c0_f32 = scalar.constant 0.0 : f32 + %negative_large = scalar.constant -1e+30 : f32 + %c0_f16x4 = vector.constant 0.0 : vector<4xf16> + %key_value_token_capacity = config.get @qwen3_moe.attention.key_value_token_capacity : index + %c63 = index.constant 63 : index + %padded_key_value_token_capacity = index.add %key_value_token_capacity, %c63 : index + %key_value_block_count0 = index.div %padded_key_value_token_capacity, %c64 : index + %key_value_block_count = index.assume %key_value_block_count0 [range(%key_value_block_count0, 1, 512)] : index + %workgroup_x, %launch_key_value_block_count = index.assume %workgroup_x_in_launch, %key_value_block_count [lt(%workgroup_x_in_launch, %key_value_block_count)] : index, index + %block_ordinal = index.add %workgroup_x, %c0 : index + %key_origin = index.mul %block_ordinal, %c64 : index + %key_token = index.add %key_origin, %lane : index + %key_valid = index.cmp ult, %key_token, %bounded_key_value_token_count : index + %mask_aligned = buffer.assume.alignment %mask {minimum_alignment = 16} : buffer + %mask_view = buffer.view %mask_aligned[%c0_offset] : buffer -> view<[%bounded_key_value_token_count]xf16> + %lane_mask = scf.if %key_valid -> (f32) { + %mask_f16 = view.load %mask_view[%key_token] : view<[%bounded_key_value_token_count]xf16> -> f16 + %mask_f32 = scalar.extf %mask_f16 : f16 to f32 + scf.yield %mask_f32 : f32 + } else { + scf.yield %negative_large : f32 + } + %block_mask_maximum = kernel.workgroup.reduce %lane_mask : f32 + %block_has_attention = scalar.cmpf ogt, %block_mask_maximum, %negative_large : f32 + scf.if %block_has_attention { + func.call inline @qwen3_moe_flash_attention_decode_split_produce_active_partials_body_f32_f16_wmma(%bounded_key_value_token_count, %launch_key_value_block_count, %launch_key_value_block_count, %query, %key, %value, %lane_mask, %partial_max, %partial_sum, %partial_output) : (index, index, index, buffer, buffer, buffer, f32, buffer, buffer, buffer) + } else { + %query_heads_per_key_value_head0 = index.div %query_head_count, %key_value_head_count : index + %query_heads_per_key_value_head = index.assume %query_heads_per_key_value_head0 [range(%query_heads_per_key_value_head0, 1, 16)] : index + %key_value_head = index.add %workgroup_y, %c0 : index + %query_head_base = index.mul %key_value_head, %query_heads_per_key_value_head : index + %subgroup_query_row = index.mul %subgroup, %c4 : index + %query_row0 = index.add %subgroup_query_row, %c0 : index + %query_row1 = index.add %subgroup_query_row, %c1 : index + %query_row2 = index.add %subgroup_query_row, %c2 : index + %query_row3 = index.add %subgroup_query_row, %c3 : index + %query_head0 = index.add %query_head_base, %query_row0 : index + %query_head1 = index.add %query_head_base, %query_row1 : index + %query_head2 = index.add %query_head_base, %query_row2 : index + %query_head3 = index.add %query_head_base, %query_row3 : index + %query_head_valid0 = index.cmp ult, %query_head0, %query_head_count : index + %query_head_valid1 = index.cmp ult, %query_head1, %query_head_count : index + %query_head_valid2 = index.cmp ult, %query_head2, %query_head_count : index + %query_head_valid3 = index.cmp ult, %query_head3, %query_head_count : index + %lane_output_channel = index.mul %lane, %c4 : index + %lane_has_output = index.cmp ult, %lane, %c32 : index + %lane_is_zero = index.cmp eq, %lane, %c0 : index + %partial_max_noalias, %partial_sum_noalias, %partial_output_noalias = buffer.assume.noalias %partial_max, %partial_sum, %partial_output : buffer, buffer, buffer + %partial_max_aligned = buffer.assume.alignment %partial_max_noalias {minimum_alignment = 16} : buffer + %partial_sum_aligned = buffer.assume.alignment %partial_sum_noalias {minimum_alignment = 16} : buffer + %partial_output_aligned = buffer.assume.alignment %partial_output_noalias {minimum_alignment = 16} : buffer + %partial_max_view = buffer.view %partial_max_aligned[%c0_offset] : buffer -> view<[%key_value_head_count]x[%key_value_block_count]x16xf32> + %partial_sum_view = buffer.view %partial_sum_aligned[%c0_offset] : buffer -> view<[%key_value_head_count]x[%key_value_block_count]x16xf32> + %partial_output_view = buffer.view %partial_output_aligned[%c0_offset] : buffer -> view<[%key_value_head_count]x[%key_value_block_count]x16x128xf16> + scf.if %lane_has_output { + scf.if %query_head_valid0 { + vector.store %c0_f16x4, %partial_output_view[%key_value_head, %block_ordinal, %query_row0, %lane_output_channel] : vector<4xf16>, view<[%key_value_head_count]x[%key_value_block_count]x16x128xf16> + } + scf.if %query_head_valid1 { + vector.store %c0_f16x4, %partial_output_view[%key_value_head, %block_ordinal, %query_row1, %lane_output_channel] : vector<4xf16>, view<[%key_value_head_count]x[%key_value_block_count]x16x128xf16> + } + scf.if %query_head_valid2 { + vector.store %c0_f16x4, %partial_output_view[%key_value_head, %block_ordinal, %query_row2, %lane_output_channel] : vector<4xf16>, view<[%key_value_head_count]x[%key_value_block_count]x16x128xf16> + } + scf.if %query_head_valid3 { + vector.store %c0_f16x4, %partial_output_view[%key_value_head, %block_ordinal, %query_row3, %lane_output_channel] : vector<4xf16>, view<[%key_value_head_count]x[%key_value_block_count]x16x128xf16> + } + } + scf.if %lane_is_zero { + scf.if %query_head_valid0 { + view.store %negative_large, %partial_max_view[%key_value_head, %block_ordinal, %query_row0] : f32, view<[%key_value_head_count]x[%key_value_block_count]x16xf32> + view.store %c0_f32, %partial_sum_view[%key_value_head, %block_ordinal, %query_row0] : f32, view<[%key_value_head_count]x[%key_value_block_count]x16xf32> + } + scf.if %query_head_valid1 { + view.store %negative_large, %partial_max_view[%key_value_head, %block_ordinal, %query_row1] : f32, view<[%key_value_head_count]x[%key_value_block_count]x16xf32> + view.store %c0_f32, %partial_sum_view[%key_value_head, %block_ordinal, %query_row1] : f32, view<[%key_value_head_count]x[%key_value_block_count]x16xf32> + } + scf.if %query_head_valid2 { + view.store %negative_large, %partial_max_view[%key_value_head, %block_ordinal, %query_row2] : f32, view<[%key_value_head_count]x[%key_value_block_count]x16xf32> + view.store %c0_f32, %partial_sum_view[%key_value_head, %block_ordinal, %query_row2] : f32, view<[%key_value_head_count]x[%key_value_block_count]x16xf32> + } + scf.if %query_head_valid3 { + view.store %negative_large, %partial_max_view[%key_value_head, %block_ordinal, %query_row3] : f32, view<[%key_value_head_count]x[%key_value_block_count]x16xf32> + view.store %c0_f32, %partial_sum_view[%key_value_head, %block_ordinal, %query_row3] : f32, view<[%key_value_head_count]x[%key_value_block_count]x16xf32> + } + } + } + func.return +} + +// Up to four split-K blocks are cheapest to fold directly. Each workitem owns +// one output element, so the reducer needs no LDS or subgroup synchronization. +func.def inline @qwen3_moe_flash_attention_decode_split_reduce_completed_direct_f32(%partial_block_capacity0: index, %active_block_count0: index, %partial_max: buffer, %partial_sum: buffer, %partial_output: buffer, %output: buffer) where [range(%partial_block_capacity0, 1, 4)] { + %partial_block_capacity, %active_block_count = index.assume %partial_block_capacity0, %active_block_count0 [range(%partial_block_capacity0, 1, 4), range(%active_block_count0, 1, 4), le(%active_block_count0, %partial_block_capacity0)] : index, index + %query_head_count = config.get @qwen3_moe.attention.query_head_count : index + %key_value_head_count = config.get @qwen3_moe.attention.key_value_head_count : index + %workgroup_y0 = kernel.workgroup.id : index + %key_value_head = index.assume %workgroup_y0 [range(%workgroup_y0, 0, 63)] : index + %workitem = kernel.workitem.id : index + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c128 = index.constant 128 : index + %c256 = index.constant 256 : index + %c0_offset = index.constant 0 : offset + %c0_f32 = scalar.constant 0.0 : f32 + %negative_large = scalar.constant -1e+30 : f32 + %query_heads_per_key_value_head0 = index.div %query_head_count, %key_value_head_count : index + %query_heads_per_key_value_head = index.assume %query_heads_per_key_value_head0 [range(%query_heads_per_key_value_head0, 1, 16)] : index + %query_head_base = index.mul %key_value_head, %query_heads_per_key_value_head : index + %partial_max_aligned = buffer.assume.alignment %partial_max {minimum_alignment = 16} : buffer + %partial_sum_aligned = buffer.assume.alignment %partial_sum {minimum_alignment = 16} : buffer + %partial_output_aligned = buffer.assume.alignment %partial_output {minimum_alignment = 16} : buffer + %output_aligned = buffer.assume.alignment %output {minimum_alignment = 16} : buffer + %partial_max_view = buffer.view %partial_max_aligned[%c0_offset] : buffer -> view<[%key_value_head_count]x[%partial_block_capacity]x16xf32> + %partial_sum_view = buffer.view %partial_sum_aligned[%c0_offset] : buffer -> view<[%key_value_head_count]x[%partial_block_capacity]x16xf32> + %partial_output_view = buffer.view %partial_output_aligned[%c0_offset] : buffer -> view<[%key_value_head_count]x[%partial_block_capacity]x16x128xf16> + %output_view = buffer.view %output_aligned[%c0_offset] : buffer -> view<[%query_head_count]x128xf32> + %partial_element_count = index.mul %query_heads_per_key_value_head, %c128 : index + scf.for %linear = [%workitem to %partial_element_count step %c256] { + %query_row = index.div %linear, %c128 : index + %output_channel = index.rem %linear, %c128 : index + %query_head = index.add %query_head_base, %query_row : index + %query_head_valid = index.cmp ult, %query_head, %query_head_count : index + scf.if %query_head_valid { + %maximum = scf.for %block = [%c0 to %active_block_count step %c1](%running_maximum = %negative_large : f32) -> (f32) unroll schedule(interleaved) { + %block_maximum = view.load %partial_max_view[%key_value_head, %block, %query_row] : view<[%key_value_head_count]x[%partial_block_capacity]x16xf32> -> f32 + %next_maximum = scalar.maxnumf %running_maximum, %block_maximum : f32 + scf.yield %next_maximum : f32 + } + %sum, %unnormalized_output = scf.for %block = [%c0 to %active_block_count step %c1](%running_sum = %c0_f32 : f32, %running_output = %c0_f32 : f32) -> (f32, f32) unroll schedule(interleaved) { + %block_maximum = view.load %partial_max_view[%key_value_head, %block, %query_row] : view<[%key_value_head_count]x[%partial_block_capacity]x16xf32> -> f32 + %delta = scalar.subf %block_maximum, %maximum : f32 + %scale = scalar.expf %delta : f32 + %block_sum = view.load %partial_sum_view[%key_value_head, %block, %query_row] : view<[%key_value_head_count]x[%partial_block_capacity]x16xf32> -> f32 + %scaled_sum = scalar.mulf %block_sum, %scale : f32 + %next_sum = scalar.addf %running_sum, %scaled_sum : f32 + %block_output_f16 = view.load %partial_output_view[%key_value_head, %block, %query_row, %output_channel] : view<[%key_value_head_count]x[%partial_block_capacity]x16x128xf16> -> f16 + %block_output = scalar.extf %block_output_f16 : f16 to f32 + %scaled_output = scalar.mulf %block_output, %scale : f32 + %next_output = scalar.addf %running_output, %scaled_output : f32 + scf.yield %next_sum, %next_output : f32, f32 + } + %normalized_output = scalar.divf %unnormalized_output, %sum : f32 + view.store %normalized_output, %output_view[%query_head, %output_channel] : f32, view<[%query_head_count]x128xf32> + } + } + func.return +} + +// Longer bounded contexts amortize a cooperative reducer. Each wave folds two +// query rows, stages the per-block scales in LDS, and writes two output channels +// per lane. +func.def inline @qwen3_moe_flash_attention_decode_split_reduce_completed_cooperative_f32(%partial_block_capacity0: index, %active_block_count0: index, %partial_max: buffer, %partial_sum: buffer, %partial_output: buffer, %output: buffer) { + %partial_block_capacity, %active_block_count = index.assume %partial_block_capacity0, %active_block_count0 [range(%partial_block_capacity0, 1, 32), range(%active_block_count0, 1, 32), le(%active_block_count0, %partial_block_capacity0)] : index, index + %query_head_count = config.get @qwen3_moe.attention.query_head_count : index + %key_value_head_count = config.get @qwen3_moe.attention.key_value_head_count : index + %workgroup_y0 = kernel.workgroup.id : index + %key_value_head = index.assume %workgroup_y0 [range(%workgroup_y0, 0, 63)] : index + %subgroup0 = kernel.subgroup.id : index + %subgroup = index.assume %subgroup0 [range(%subgroup0, 0, 3)] : index + %lane = kernel.subgroup.lane.id : index + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c4 = index.constant 4 : index + %c7 = index.constant 7 : index + %c8 = index.constant 8 : index + %c0_offset = index.constant 0 : offset + %scale_stage_bytes = index.constant 1024 : offset + %c0_f32 = scalar.constant 0.0 : f32 + %negative_large = scalar.constant -1e+30 : f32 + %c0_f32x2 = vector.constant 0.0 : vector<2xf32> + %query_heads_per_key_value_head0 = index.div %query_head_count, %key_value_head_count : index + %query_heads_per_key_value_head = index.assume %query_heads_per_key_value_head0 [range(%query_heads_per_key_value_head0, 1, 16)] : index + %query_head_base = index.mul %key_value_head, %query_heads_per_key_value_head : index + %lane_output_channel0 = index.mul %lane, %c2 : index + %lane_output_channel = index.assume %lane_output_channel0 [range(%lane_output_channel0, 0, 126)] : index + %lane_has_block = index.cmp ult, %lane, %active_block_count : index + %partial_max_aligned = buffer.assume.alignment %partial_max {minimum_alignment = 16} : buffer + %partial_sum_aligned = buffer.assume.alignment %partial_sum {minimum_alignment = 16} : buffer + %partial_output_aligned = buffer.assume.alignment %partial_output {minimum_alignment = 16} : buffer + %output_aligned = buffer.assume.alignment %output {minimum_alignment = 16} : buffer + %partial_max_view = buffer.view %partial_max_aligned[%c0_offset] : buffer -> view<[%key_value_head_count]x[%partial_block_capacity]x16xf32> + %partial_sum_view = buffer.view %partial_sum_aligned[%c0_offset] : buffer -> view<[%key_value_head_count]x[%partial_block_capacity]x16xf32> + %partial_output_view = buffer.view %partial_output_aligned[%c0_offset] : buffer -> view<[%key_value_head_count]x[%partial_block_capacity]x16x128xf16> + %output_view = buffer.view %output_aligned[%c0_offset] : buffer -> view<[%query_head_count]x128xf32> + %scale_stage = buffer.alloca align(16) %scale_stage_bytes : buffer + %scale_stage_view = buffer.view %scale_stage[%c0_offset] : buffer -> view<4x32x2xf32> + %padded_reduction_row_count = index.add %query_heads_per_key_value_head, %c7 : index + %reduction_phase_count = index.div %padded_reduction_row_count, %c8 : index + scf.for %phase = [%c0 to %reduction_phase_count step %c1] unroll { + %phase_row_base = index.mul %phase, %c8 : index + %subgroup_row_base = index.mul %subgroup, %c2 : index + %query_row0 = index.add %phase_row_base, %subgroup_row_base : index + %query_row1 = index.add %query_row0, %c1 : index + %query_head0 = index.add %query_head_base, %query_row0 : index + %query_head1 = index.add %query_head0, %c1 : index + %query_head_valid0 = index.cmp ult, %query_head0, %query_head_count : index + %query_head_valid1 = index.cmp ult, %query_head1, %query_head_count : index + %safe_query_row0 = scf.select %query_head_valid0, %query_row0, %c0 : index + %safe_query_row1 = scf.select %query_head_valid1, %query_row1, %c0 : index + %reducer_active0 = scalar.andi %lane_has_block, %query_head_valid0 : i1 + %reducer_active1 = scalar.andi %lane_has_block, %query_head_valid1 : i1 + %local_maximum0 = scf.if %reducer_active0 -> (f32) { + %block_maximum = view.load %partial_max_view[%key_value_head, %lane, %query_row0] : view<[%key_value_head_count]x[%partial_block_capacity]x16xf32> -> f32 + scf.yield %block_maximum : f32 + } else { + scf.yield %negative_large : f32 + } + %local_maximum1 = scf.if %reducer_active1 -> (f32) { + %block_maximum = view.load %partial_max_view[%key_value_head, %lane, %query_row1] : view<[%key_value_head_count]x[%partial_block_capacity]x16xf32> -> f32 + scf.yield %block_maximum : f32 + } else { + scf.yield %negative_large : f32 + } + %local_maximums = vector.from_elements %local_maximum0, %local_maximum1 : vector<2xf32> + %maximums = kernel.subgroup.reduce %local_maximums : vector<2xf32> + %maximum0 = vector.extract %maximums[0] : vector<2xf32> -> f32 + %maximum1 = vector.extract %maximums[1] : vector<2xf32> -> f32 + %local_sum0, %scale0 = scf.if %reducer_active0 -> (f32, f32) { + %block_maximum = view.load %partial_max_view[%key_value_head, %lane, %query_row0] : view<[%key_value_head_count]x[%partial_block_capacity]x16xf32> -> f32 + %delta = scalar.subf %block_maximum, %maximum0 : f32 + %scale = scalar.expf %delta : f32 + %block_sum = view.load %partial_sum_view[%key_value_head, %lane, %query_row0] : view<[%key_value_head_count]x[%partial_block_capacity]x16xf32> -> f32 + %scaled_sum = scalar.mulf %block_sum, %scale : f32 + scf.yield %scaled_sum, %scale : f32, f32 + } else { + scf.yield %c0_f32, %c0_f32 : f32, f32 + } + %local_sum1, %scale1 = scf.if %reducer_active1 -> (f32, f32) { + %block_maximum = view.load %partial_max_view[%key_value_head, %lane, %query_row1] : view<[%key_value_head_count]x[%partial_block_capacity]x16xf32> -> f32 + %delta = scalar.subf %block_maximum, %maximum1 : f32 + %scale = scalar.expf %delta : f32 + %block_sum = view.load %partial_sum_view[%key_value_head, %lane, %query_row1] : view<[%key_value_head_count]x[%partial_block_capacity]x16xf32> -> f32 + %scaled_sum = scalar.mulf %block_sum, %scale : f32 + scf.yield %scaled_sum, %scale : f32, f32 + } else { + scf.yield %c0_f32, %c0_f32 : f32, f32 + } + scf.if %lane_has_block { + %bounded_block_lane = index.assume %lane [range(%lane, 0, 31)] : index + %scales = vector.from_elements %scale0, %scale1 : vector<2xf32> + vector.store %scales, %scale_stage_view[%subgroup, %bounded_block_lane, %c0] : vector<2xf32>, view<4x32x2xf32> + } + %local_sums = vector.from_elements %local_sum0, %local_sum1 : vector<2xf32> + %sums = kernel.subgroup.reduce %local_sums : vector<2xf32> + %sum0 = vector.extract %sums[0] : vector<2xf32> -> f32 + %sum1 = vector.extract %sums[1] : vector<2xf32> -> f32 + kernel.barrier scope(workgroup) ordering(acq_rel) + %unnormalized_output0, %unnormalized_output1 = scf.for %block = [%c0 to %active_block_count step %c1](%running_output0 = %c0_f32x2 : vector<2xf32>, %running_output1 = %c0_f32x2 : vector<2xf32>) -> (vector<2xf32>, vector<2xf32>) unroll(%c4) schedule(interleaved) { + %scales = vector.load %scale_stage_view[%subgroup, %block, %c0] : view<4x32x2xf32> -> vector<2xf32> + %scale0 = vector.extract %scales[0] : vector<2xf32> -> f32 + %scale1 = vector.extract %scales[1] : vector<2xf32> -> f32 + %block_output0_f16 = vector.load %partial_output_view[%key_value_head, %block, %safe_query_row0, %lane_output_channel] : view<[%key_value_head_count]x[%partial_block_capacity]x16x128xf16> -> vector<2xf16> + %block_output0 = vector.extf %block_output0_f16 : vector<2xf16> to vector<2xf32> + %scale_vector0 = vector.splat %scale0 : vector<2xf32> + %scaled_output0 = vector.mulf %block_output0, %scale_vector0 : vector<2xf32> + %next_output0 = vector.addf %running_output0, %scaled_output0 : vector<2xf32> + %block_output1_f16 = vector.load %partial_output_view[%key_value_head, %block, %safe_query_row1, %lane_output_channel] : view<[%key_value_head_count]x[%partial_block_capacity]x16x128xf16> -> vector<2xf16> + %block_output1 = vector.extf %block_output1_f16 : vector<2xf16> to vector<2xf32> + %scale_vector1 = vector.splat %scale1 : vector<2xf32> + %scaled_output1 = vector.mulf %block_output1, %scale_vector1 : vector<2xf32> + %next_output1 = vector.addf %running_output1, %scaled_output1 : vector<2xf32> + scf.yield %next_output0, %next_output1 : vector<2xf32>, vector<2xf32> + } + %sum_vector0 = vector.splat %sum0 : vector<2xf32> + %sum_vector1 = vector.splat %sum1 : vector<2xf32> + %normalized_output0 = vector.divf %unnormalized_output0, %sum_vector0 : vector<2xf32> + %normalized_output1 = vector.divf %unnormalized_output1, %sum_vector1 : vector<2xf32> + scf.if %query_head_valid0 { + vector.store %normalized_output0, %output_view[%query_head0, %lane_output_channel] : vector<2xf32>, view<[%query_head_count]x128xf32> + } + scf.if %query_head_valid1 { + vector.store %normalized_output1, %output_view[%query_head1, %lane_output_channel] : vector<2xf32>, view<[%query_head_count]x128xf32> + } + } + func.return +} + +// Packs the query heads owned by one completed KV head into GGML's Q8_1 x4 +// layout. Each 32-workitem cohort owns one contiguous 128-element query row; +// the production 8:1 GQA ratio therefore fills all 256 workitems. Smaller or +// larger valid ratios use the same phase loop without making barriers +// conditional. +func.def inline @qwen3_moe_flash_attention_decode_pack_completed_key_value_head_q8_1_x4(%key_value_head: index, %output: buffer, %q8_output: buffer) { + %query_head_count = config.get @qwen3_moe.attention.query_head_count : index + %key_value_head_count = config.get @qwen3_moe.attention.key_value_head_count : index + %query_heads_per_key_value_head0 = index.div %query_head_count, %key_value_head_count : index + %query_heads_per_key_value_head = index.assume %query_heads_per_key_value_head0 [range(%query_heads_per_key_value_head0, 1, 16)] : index + %workitem0 = kernel.workitem.id : index + %workitem = index.assume %workitem0 [range(%workitem0, 0, 255)] : index + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c4 = index.constant 4 : index + %c7 = index.constant 7 : index + %c8 = index.constant 8 : index + %c32 = index.constant 32 : index + %group_bytes = index.constant 144 : offset + %payload_byte_add = index.constant 16 : offset + %scratch_d_byte_add = index.constant 1024 : offset + %scratch_bytes = index.constant 1152 : offset + %c0_offset = index.constant 0 : offset + %c0_f32 = scalar.constant 0.0 : f32 + %c1_f32 = scalar.constant 1.0 : f32 + %c127 = scalar.constant 127.0 : f32 + %c0_f32x4 = vector.constant 0.0 : vector<4xf32> + %output_noalias, %q8_output_noalias = buffer.assume.noalias %output, %q8_output : buffer, buffer + %output_aligned = buffer.assume.alignment %output_noalias {minimum_alignment = 16} : buffer + %q8_output_aligned = buffer.assume.alignment %q8_output_noalias {minimum_alignment = 16} : buffer + %output_view = buffer.view %output_aligned[%c0_offset] : buffer -> view<[%query_head_count]x128xf32> + %scratch = buffer.alloca align(16) %scratch_bytes : buffer + %scratch_values = buffer.view %scratch[%c0_offset] : buffer -> view<256xf32> + %scratch_d = buffer.view %scratch[%scratch_d_byte_add] : buffer -> view<32xf32> + %group_in_phase = index.div %workitem, %c32 : index + %lane_in_group0 = index.rem %workitem, %c32 : index + %lane_in_group = index.assume %lane_in_group0 [range(%lane_in_group0, 0, 31)] : index + %block_in_group0 = index.div %lane_in_group, %c8 : index + %block_in_group = index.assume %block_in_group0 [range(%block_in_group0, 0, 3)] : index + %word_in_block0 = index.rem %lane_in_group, %c8 : index + %word_in_block = index.assume %word_in_block0 [range(%word_in_block0, 0, 7)] : index + %padded_phase_count = index.add %query_heads_per_key_value_head, %c7 : index + %phase_count = index.div %padded_phase_count, %c8 : index + scf.for %phase = [%c0 to %phase_count step %c1] unroll { + %phase_group_base = index.mul %phase, %c8 : index + %group_in_key_value_head = index.add %phase_group_base, %group_in_phase : index + %valid_group = index.cmp ult, %group_in_key_value_head, %query_heads_per_key_value_head : index + %key_value_query_head_base = index.mul %key_value_head, %query_heads_per_key_value_head : index + %query_head0 = index.add %key_value_query_head_base, %group_in_key_value_head : index + %safe_query_head = scf.select %valid_group, %query_head0, %c0 : index + %block_element_add = index.mul %block_in_group, %c32 : index + %word_element_add = index.mul %word_in_block, %c4 : index + %input_channel0 = index.add %block_element_add, %word_element_add : index + %input_channel = index.assume %input_channel0 [range(%input_channel0, 0, 124)] : index + %input_values = scf.if %valid_group -> (vector<4xf32>) { + %values = vector.load %output_view[%safe_query_head, %input_channel] : view<[%query_head_count]x128xf32> -> vector<4xf32> + scf.yield %values : vector<4xf32> + } else { + scf.yield %c0_f32x4 : vector<4xf32> + } + %absolute_values = vector.absf %input_values : vector<4xf32> + %thread_max = vector.reduce %absolute_values, %c0_f32 : vector<4xf32>, f32 + view.store %thread_max, %scratch_values[%workitem] : f32, view<256xf32> + kernel.barrier scope(workgroup) ordering(acq_rel) + %is_cohort_leader = index.cmp eq, %word_in_block, %c0 : index + %d_index0 = index.mul %group_in_phase, %c4 : index + %d_index1 = index.add %d_index0, %block_in_group : index + %d_index = index.assume %d_index1 [range(%d_index1, 0, 31)] : index + scf.if %is_cohort_leader { + %group_thread_base = index.mul %group_in_phase, %c32 : index + %block_thread_add = index.mul %block_in_group, %c8 : index + %cohort_base = index.add %group_thread_base, %block_thread_add : index + %cohort_maxima = vector.load %scratch_values[%cohort_base] : view<256xf32> -> vector<8xf32> + %amax = vector.reduce %cohort_maxima, %c0_f32 : vector<8xf32>, f32 + %d = scalar.divf %amax, %c127 : f32 + view.store %d, %scratch_d[%d_index] : f32, view<32xf32> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + %d = view.load %scratch_d[%d_index] : view<32xf32> -> f32 + %d_nonzero = scalar.cmpf one, %d, %c0_f32 : f32 + %d_inverse = scf.if %d_nonzero -> (f32) { + %inverse = scalar.divf %c1_f32, %d : f32 + scf.yield %inverse : f32 + } else { + scf.yield %c0_f32 : f32 + } + %d_inverse_vector = vector.splat %d_inverse : vector<4xf32> + %scaled_values = vector.mulf %input_values, %d_inverse_vector : vector<4xf32> + %rounded_values = vector.roundf %scaled_values : vector<4xf32> + %quantized_values = vector.fptosi %rounded_values : vector<4xf32> to vector<4xi8> + %packed_word = vector.bitcast %quantized_values : vector<4xi8> to vector<1xi32> + %group_byte_offset = index.scale %safe_query_head, %group_bytes : index, offset -> offset + %payload_byte_offset = index.add %group_byte_offset, %payload_byte_add : offset + %group_ds = buffer.view %q8_output_aligned[%group_byte_offset] : buffer -> view<8xf16> + %group_qs = buffer.view %q8_output_aligned[%payload_byte_offset] : buffer -> view<32xi32> + %block_word_add = index.mul %block_in_group, %c8 : index + %packed_word_index0 = index.add %block_word_add, %word_in_block : index + %packed_word_index = index.assume %packed_word_index0 [range(%packed_word_index0, 0, 31)] : index + scf.if %valid_group { + vector.store %packed_word, %group_qs[%packed_word_index] : vector<1xi32>, view<32xi32> + } + %thread_sum = vector.reduce %rounded_values, %c0_f32 : vector<4xf32>, f32 + view.store %thread_sum, %scratch_values[%workitem] : f32, view<256xf32> + kernel.barrier scope(workgroup) ordering(acq_rel) + scf.if %is_cohort_leader { + %group_thread_base = index.mul %group_in_phase, %c32 : index + %block_thread_add = index.mul %block_in_group, %c8 : index + %cohort_base = index.add %group_thread_base, %block_thread_add : index + %cohort_sums = vector.load %scratch_values[%cohort_base] : view<256xf32> -> vector<8xf32> + %quantized_sum = vector.reduce %cohort_sums, %c0_f32 : vector<8xf32>, f32 + %s = scalar.mulf %quantized_sum, %d : f32 + %d_f16 = scalar.fptrunc %d : f32 to f16 + %s_f16 = scalar.fptrunc %s : f32 to f16 + %ds_index = index.mul %block_in_group, %c2 : index + %s_index = index.add %ds_index, %c1 : index + scf.if %valid_group { + view.store %d_f16, %group_ds[%ds_index] : f16, view<8xf16> + view.store %s_f16, %group_ds[%s_index] : f16, view<8xf16> + } + } + kernel.barrier scope(workgroup) ordering(acq_rel) + } + func.return +} + +// Completes bounded decode contexts inside the last arriving producer +// workgroup. Four KV-head workgroups reduce their own query heads while all +// other producers retire, erasing a second dispatch and its execution barrier. +// Capacity selects the algorithm and partial layout; producer count owns the +// issue-time completion threshold and active reduction prefix. +template.decl @qwen3_moe_attention_decode_split_reduce_fused(%key_value_token_capacity: index, %partial_block_capacity0: index, %producer_block_count0: index, %publish_q8: i1, %partial_max: buffer, %partial_sum: buffer, %partial_output: buffer, %completion_counter: buffer, %output: buffer, %q8_output: buffer) -> () +template.def<@qwen3_moe_attention_decode_split_reduce_fused> priority(20) @qwen3_moe_flash_attention_decode_split_reduce_fused_direct_f32(%key_value_token_capacity: index, %partial_block_capacity0: index, %producer_block_count0: index, %publish_q8: i1, %partial_max: buffer, %partial_sum: buffer, %partial_output: buffer, %completion_counter: buffer, %output: buffer, %q8_output: buffer) where [range(%key_value_token_capacity, 64, 256)] { + %partial_block_capacity, %producer_block_count = index.assume %partial_block_capacity0, %producer_block_count0 [range(%partial_block_capacity0, 1, 4), range(%producer_block_count0, 1, 4), le(%producer_block_count0, %partial_block_capacity0)] : index, index + %key_value_head_count = config.get @qwen3_moe.attention.key_value_head_count : index + %workgroup_y0 = kernel.workgroup.id : index + %key_value_head = index.assume %workgroup_y0 [range(%workgroup_y0, 0, 63)] : index + %workitem = kernel.workitem.id : index + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c0_offset = index.constant 0 : offset + %counter_scratch_bytes = index.constant 4 : offset + %c0_i32 = scalar.constant 0 : i32 + %c1_i32 = scalar.constant 1 : i32 + %partial_max_noalias, %partial_sum_noalias, %partial_output_noalias, %completion_counter_noalias, %output_noalias = buffer.assume.noalias %partial_max, %partial_sum, %partial_output, %completion_counter, %output : buffer, buffer, buffer, buffer, buffer + %completion_counter_aligned = buffer.assume.alignment %completion_counter_noalias {minimum_alignment = 16} : buffer + %completion_counter_view = buffer.view %completion_counter_aligned[%c0_offset] : buffer -> view<[%key_value_head_count]xi32> + %counter_scratch = buffer.alloca align(4) %counter_scratch_bytes : buffer + %counter_scratch_view = buffer.view %counter_scratch[%c0_offset] : buffer -> view<1xi32> + %workitem_is_zero = index.cmp eq, %workitem, %c0 : index + // Publish every producer's partial stores before the leader advances one + // workgroup arrival. The last arrival then acquires every partial. + kernel.barrier scope(workgroup) ordering(release) + scf.if %workitem_is_zero { + %old_counter = view.atomic.rmw %c1_i32, %completion_counter_view[%key_value_head] {ordering = acq_rel, scope = device} : i32, view<[%key_value_head_count]xi32> -> i32 + view.store %old_counter, %counter_scratch_view[%c0] : i32, view<1xi32> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + %old_counter = view.load %counter_scratch_view[%c0] : view<1xi32> -> i32 + %key_value_block_count_i32 = index.cast %producer_block_count : index to i32 + %last_block_ordinal_i32 = scalar.subi %key_value_block_count_i32, %c1_i32 : i32 + %negative_key_value_block_count_i32 = scalar.subi %c0_i32, %key_value_block_count_i32 : i32 + %is_last_partition = scalar.cmpi eq, %old_counter, %last_block_ordinal_i32 : i32 + scf.if %is_last_partition { + kernel.barrier scope(workgroup) ordering(acquire) + func.call inline @qwen3_moe_flash_attention_decode_split_reduce_completed_direct_f32(%partial_block_capacity, %producer_block_count, %partial_max_noalias, %partial_sum_noalias, %partial_output_noalias, %output_noalias) : (index, index, buffer, buffer, buffer, buffer) + scf.if %publish_q8 { + kernel.barrier scope(workgroup) ordering(acq_rel) + func.call @qwen3_moe_flash_attention_decode_pack_completed_key_value_head_q8_1_x4(%key_value_head, %output_noalias, %q8_output) : (index, buffer, buffer) + } + // Do not expose the reset until every final F32 and Q8 store completes. + kernel.barrier scope(workgroup) ordering(release) + scf.if %workitem_is_zero { + view.atomic.reduce %negative_key_value_block_count_i32, %completion_counter_view[%key_value_head] {ordering = release, scope = device} : i32, view<[%key_value_head_count]xi32> + } + } + template.return +} + +// Contexts without a proven short bound use the cooperative completion path. +template.def<@qwen3_moe_attention_decode_split_reduce_fused> priority(10) @qwen3_moe_flash_attention_decode_split_reduce_fused_cooperative_f32(%key_value_token_capacity: index, %partial_block_capacity0: index, %producer_block_count0: index, %publish_q8: i1, %partial_max: buffer, %partial_sum: buffer, %partial_output: buffer, %completion_counter: buffer, %output: buffer, %q8_output: buffer) where [range(%key_value_token_capacity, 257, 2048)] { + %partial_block_capacity, %producer_block_count = index.assume %partial_block_capacity0, %producer_block_count0 [range(%partial_block_capacity0, 1, 32), range(%producer_block_count0, 1, 32), le(%producer_block_count0, %partial_block_capacity0)] : index, index + %key_value_head_count = config.get @qwen3_moe.attention.key_value_head_count : index + %workgroup_y0 = kernel.workgroup.id : index + %key_value_head = index.assume %workgroup_y0 [range(%workgroup_y0, 0, 63)] : index + %workitem = kernel.workitem.id : index + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c0_offset = index.constant 0 : offset + %counter_scratch_bytes = index.constant 4 : offset + %c0_i32 = scalar.constant 0 : i32 + %c1_i32 = scalar.constant 1 : i32 + %partial_max_noalias, %partial_sum_noalias, %partial_output_noalias, %completion_counter_noalias, %output_noalias = buffer.assume.noalias %partial_max, %partial_sum, %partial_output, %completion_counter, %output : buffer, buffer, buffer, buffer, buffer + %completion_counter_aligned = buffer.assume.alignment %completion_counter_noalias {minimum_alignment = 16} : buffer + %completion_counter_view = buffer.view %completion_counter_aligned[%c0_offset] : buffer -> view<[%key_value_head_count]xi32> + %counter_scratch = buffer.alloca align(4) %counter_scratch_bytes : buffer + %counter_scratch_view = buffer.view %counter_scratch[%c0_offset] : buffer -> view<1xi32> + %workitem_is_zero = index.cmp eq, %workitem, %c0 : index + // Publish every producer's partial stores before the leader advances one + // workgroup arrival. The last arrival then acquires every partial. + kernel.barrier scope(workgroup) ordering(release) + scf.if %workitem_is_zero { + %old_counter = view.atomic.rmw %c1_i32, %completion_counter_view[%key_value_head] {ordering = acq_rel, scope = device} : i32, view<[%key_value_head_count]xi32> -> i32 + view.store %old_counter, %counter_scratch_view[%c0] : i32, view<1xi32> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + %old_counter = view.load %counter_scratch_view[%c0] : view<1xi32> -> i32 + %key_value_block_count_i32 = index.cast %producer_block_count : index to i32 + %last_block_ordinal_i32 = scalar.subi %key_value_block_count_i32, %c1_i32 : i32 + %negative_key_value_block_count_i32 = scalar.subi %c0_i32, %key_value_block_count_i32 : i32 + %is_last_partition = scalar.cmpi eq, %old_counter, %last_block_ordinal_i32 : i32 + scf.if %is_last_partition { + kernel.barrier scope(workgroup) ordering(acquire) + func.call inline @qwen3_moe_flash_attention_decode_split_reduce_completed_cooperative_f32(%partial_block_capacity, %producer_block_count, %partial_max_noalias, %partial_sum_noalias, %partial_output_noalias, %output_noalias) : (index, index, buffer, buffer, buffer, buffer) + scf.if %publish_q8 { + kernel.barrier scope(workgroup) ordering(acq_rel) + func.call @qwen3_moe_flash_attention_decode_pack_completed_key_value_head_q8_1_x4(%key_value_head, %output_noalias, %q8_output) : (index, buffer, buffer) + } + // Do not expose the reset until every final F32 and Q8 store completes. + kernel.barrier scope(workgroup) ordering(release) + scf.if %workitem_is_zero { + view.atomic.reduce %negative_key_value_block_count_i32, %completion_counter_view[%key_value_head] {ordering = release, scope = device} : i32, view<[%key_value_head_count]xi32> + } + } + template.return +} + +// Short-context export: produce and reduce in one dispatch. +kernel.def target(@qwen3_moe_decode_split_gfx11_wave64) @qwen3_moe_flash_attention_decode_split_f32_f16_wmma(%key_value_token_count: index) { + %key_value_head_count = config.get @qwen3_moe.attention.key_value_head_count : index + %key_value_token_capacity = config.get @qwen3_moe.attention.key_value_token_capacity : index + %c1 = index.constant 1 : index + %c63 = index.constant 63 : index + %c64 = index.constant 64 : index + %c256 = index.constant 256 : index + %padded_key_value_token_capacity = index.add %key_value_token_capacity, %c63 : index + %key_value_block_count = index.div %padded_key_value_token_capacity, %c64 : index + kernel.launch.config workgroups(%key_value_block_count, %key_value_head_count, %c1) workgroup_size(%c256, %c1, %c1) : index +} launch(%key_value_token_count: index, %query: buffer, %key: buffer, %value: buffer, %mask: buffer, %partial_max: buffer, %partial_sum: buffer, %partial_output: buffer, %completion_counter: buffer, %output: buffer) { + %key_value_token_capacity = config.get @qwen3_moe.attention.key_value_token_capacity : index + %c63 = index.constant 63 : index + %c64 = index.constant 64 : index + %key_value_token_count_in_range = index.assume %key_value_token_count [range(%key_value_token_count, 1, 2048)] : index + %bounded_key_value_token_count, %launch_key_value_token_capacity = index.assume %key_value_token_count_in_range, %key_value_token_capacity [le(%key_value_token_count_in_range, %key_value_token_capacity)] : index, index + %padded_key_value_token_capacity = index.add %launch_key_value_token_capacity, %c63 : index + %producer_block_count = index.div %padded_key_value_token_capacity, %c64 : index + %publish_q8 = scalar.constant false : i1 + func.call @qwen3_moe_flash_attention_decode_split_produce_partials_body_f32_f16_wmma(%bounded_key_value_token_count, %query, %key, %value, %mask, %partial_max, %partial_sum, %partial_output) : (index, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + template.apply<@qwen3_moe_attention_decode_split_reduce_fused>(%launch_key_value_token_capacity, %producer_block_count, %producer_block_count, %publish_q8, %partial_max, %partial_sum, %partial_output, %completion_counter, %output, %partial_output) : (index, index, index, i1, buffer, buffer, buffer, buffer, buffer, buffer) + kernel.return +} + +// Short-context export that publishes the F32 attention result and the Q8_1 +// representation consumed by the following output projection. +kernel.def target(@qwen3_moe_decode_split_gfx11_wave64) @qwen3_moe_flash_attention_decode_split_f32_f16_wmma_next_q8(%key_value_token_count: index) { + %key_value_head_count = config.get @qwen3_moe.attention.key_value_head_count : index + %key_value_token_capacity = config.get @qwen3_moe.attention.key_value_token_capacity : index + %c1 = index.constant 1 : index + %c63 = index.constant 63 : index + %c64 = index.constant 64 : index + %c256 = index.constant 256 : index + %padded_key_value_token_capacity = index.add %key_value_token_capacity, %c63 : index + %key_value_block_count = index.div %padded_key_value_token_capacity, %c64 : index + kernel.launch.config workgroups(%key_value_block_count, %key_value_head_count, %c1) workgroup_size(%c256, %c1, %c1) : index +} launch(%key_value_token_count: index, %query: buffer, %key: buffer, %value: buffer, %mask: buffer, %partial_max: buffer, %partial_sum: buffer, %partial_output: buffer, %completion_counter: buffer, %output: buffer, %next_q8_output: buffer) { + %key_value_token_capacity = config.get @qwen3_moe.attention.key_value_token_capacity : index + %c63 = index.constant 63 : index + %c64 = index.constant 64 : index + %key_value_token_count_in_range = index.assume %key_value_token_count [range(%key_value_token_count, 1, 2048)] : index + %bounded_key_value_token_count, %launch_key_value_token_capacity = index.assume %key_value_token_count_in_range, %key_value_token_capacity [le(%key_value_token_count_in_range, %key_value_token_capacity)] : index, index + %padded_key_value_token_capacity = index.add %launch_key_value_token_capacity, %c63 : index + %producer_block_count = index.div %padded_key_value_token_capacity, %c64 : index + %publish_q8 = scalar.constant true : i1 + func.call @qwen3_moe_flash_attention_decode_split_produce_partials_body_f32_f16_wmma(%bounded_key_value_token_count, %query, %key, %value, %mask, %partial_max, %partial_sum, %partial_output) : (index, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + template.apply<@qwen3_moe_attention_decode_split_reduce_fused>(%launch_key_value_token_capacity, %producer_block_count, %producer_block_count, %publish_q8, %partial_max, %partial_sum, %partial_output, %completion_counter, %output, %next_q8_output) : (index, index, index, i1, buffer, buffer, buffer, buffer, buffer, buffer) + kernel.return +} + +// Long-context producer export: publish partials for a following parallel +// reducer without carrying short-context synchronization bindings. +kernel.def target(@qwen3_moe_decode_split_gfx11_wave64) @qwen3_moe_flash_attention_decode_split_produce_partials_f32_f16_wmma(%key_value_token_count: index) { + %key_value_head_count = config.get @qwen3_moe.attention.key_value_head_count : index + %key_value_token_capacity = config.get @qwen3_moe.attention.key_value_token_capacity : index + %c1 = index.constant 1 : index + %c63 = index.constant 63 : index + %c64 = index.constant 64 : index + %c256 = index.constant 256 : index + %padded_key_value_token_capacity = index.add %key_value_token_capacity, %c63 : index + %key_value_block_count = index.div %padded_key_value_token_capacity, %c64 : index + kernel.launch.config workgroups(%key_value_block_count, %key_value_head_count, %c1) workgroup_size(%c256, %c1, %c1) : index +} launch(%key_value_token_count: index, %query: buffer, %key: buffer, %value: buffer, %mask: buffer, %partial_max: buffer, %partial_sum: buffer, %partial_output: buffer) { + %key_value_token_capacity = config.get @qwen3_moe.attention.key_value_token_capacity : index + %key_value_token_count_in_range = index.assume %key_value_token_count [range(%key_value_token_count, 1, 32768)] : index + %bounded_key_value_token_count = index.assume %key_value_token_count_in_range [le(%key_value_token_count_in_range, %key_value_token_capacity)] : index + func.call @qwen3_moe_flash_attention_decode_split_produce_partials_body_f32_f16_wmma(%bounded_key_value_token_count, %query, %key, %value, %mask, %partial_max, %partial_sum, %partial_output) : (index, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + kernel.return +} + +// Long contexts have enough split-K partials that assigning the reduction to +// one last-arriving producer serializes useful work. This reducer launches one +// two-wave workgroup per query head after an execution barrier from the +// producer. The first wave computes normalization once and both waves consume +// it, avoiding the duplicate work of independent 64-channel output slices. +kernel.def target(@qwen3_moe_decode_split_gfx11_wave64) @qwen3_moe_flash_attention_decode_split_reduce_f32(%key_value_token_count: index) { + %query_head_count = config.get @qwen3_moe.attention.query_head_count : index + %c1 = index.constant 1 : index + %c128 = index.constant 128 : index + kernel.launch.config workgroups(%c1, %query_head_count, %c1) workgroup_size(%c128, %c1, %c1) : index +} launch(%key_value_token_count: index, %partial_max: buffer, %partial_sum: buffer, %partial_output: buffer, %output: buffer) { + %bounded_key_value_token_count = index.assume %key_value_token_count [range(%key_value_token_count, 1, 32768)] : index + %query_head_count = config.get @qwen3_moe.attention.query_head_count : index + %key_value_head_count = config.get @qwen3_moe.attention.key_value_head_count : index + %query_head0 = kernel.workgroup.id : index + %query_head = index.assume %query_head0 [range(%query_head0, 0, 63)] : index + %workitem = kernel.workitem.id : index + %subgroup = kernel.subgroup.id : index + %lane = kernel.subgroup.lane.id : index + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c4 = index.constant 4 : index + %c64 = index.constant 64 : index + %c0_offset = index.constant 0 : offset + %reduction_stage_bytes = index.constant 8 : offset + %c0_f32 = scalar.constant 0.0 : f32 + %negative_large = scalar.constant -1e+30 : f32 + %key_value_token_capacity = config.get @qwen3_moe.attention.key_value_token_capacity : index + %c63 = index.constant 63 : index + %padded_key_value_token_capacity = index.add %key_value_token_capacity, %c63 : index + %key_value_block_count0 = index.div %padded_key_value_token_capacity, %c64 : index + %key_value_block_count = index.assume %key_value_block_count0 [range(%key_value_block_count0, 1, 512)] : index + %query_heads_per_key_value_head0 = index.div %query_head_count, %key_value_head_count : index + %query_heads_per_key_value_head = index.assume %query_heads_per_key_value_head0 [range(%query_heads_per_key_value_head0, 1, 16)] : index + %key_value_head = index.div %query_head, %query_heads_per_key_value_head : index + %query_row = index.rem %query_head, %query_heads_per_key_value_head : index + %is_first_subgroup = index.cmp eq, %subgroup, %c0 : index + %workitem_is_zero = index.cmp eq, %workitem, %c0 : index + %output_channel = index.add %workitem, %c0 : index + %partial_max_noalias, %partial_sum_noalias, %partial_output_noalias, %output_noalias = buffer.assume.noalias %partial_max, %partial_sum, %partial_output, %output : buffer, buffer, buffer, buffer + %partial_max_aligned = buffer.assume.alignment %partial_max_noalias {minimum_alignment = 16} : buffer + %partial_sum_aligned = buffer.assume.alignment %partial_sum_noalias {minimum_alignment = 16} : buffer + %partial_output_aligned = buffer.assume.alignment %partial_output_noalias {minimum_alignment = 16} : buffer + %output_aligned = buffer.assume.alignment %output_noalias {minimum_alignment = 16} : buffer + %partial_max_view = buffer.view %partial_max_aligned[%c0_offset] : buffer -> view<[%key_value_head_count]x[%key_value_block_count]x16xf32> + %partial_sum_view = buffer.view %partial_sum_aligned[%c0_offset] : buffer -> view<[%key_value_head_count]x[%key_value_block_count]x16xf32> + %partial_output_view = buffer.view %partial_output_aligned[%c0_offset] : buffer -> view<[%key_value_head_count]x[%key_value_block_count]x16x128xf16> + %output_view = buffer.view %output_aligned[%c0_offset] : buffer -> view<[%query_head_count]x128xf32> + %reduction_stage = buffer.alloca align(8) %reduction_stage_bytes : buffer + %reduction_stage_view = buffer.view %reduction_stage[%c0_offset] : buffer -> view<2xf32> + %lane_maximum = scf.if %is_first_subgroup -> (f32) { + %maximum = scf.for %block = [%lane to %key_value_block_count step %c64](%running_maximum = %negative_large : f32) -> (f32) { + %block_maximum = view.load %partial_max_view[%key_value_head, %block, %query_row] : view<[%key_value_head_count]x[%key_value_block_count]x16xf32> -> f32 + %next_maximum = scalar.maxnumf %running_maximum, %block_maximum : f32 + scf.yield %next_maximum : f32 + } + scf.yield %maximum : f32 + } else { + scf.yield %negative_large : f32 + } + %subgroup_maximum = kernel.subgroup.reduce %lane_maximum : f32 + scf.if %workitem_is_zero { + view.store %subgroup_maximum, %reduction_stage_view[%c0] : f32, view<2xf32> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + %maximum = view.load %reduction_stage_view[%c0] : view<2xf32> -> f32 + %lane_sum = scf.if %is_first_subgroup -> (f32) { + %sum = scf.for %block = [%lane to %key_value_block_count step %c64](%running_sum = %c0_f32 : f32) -> (f32) { + %block_maximum = view.load %partial_max_view[%key_value_head, %block, %query_row] : view<[%key_value_head_count]x[%key_value_block_count]x16xf32> -> f32 + %delta = scalar.subf %block_maximum, %maximum : f32 + %scale = scalar.expf %delta : f32 + view.store %scale, %partial_max_view[%key_value_head, %block, %query_row] : f32, view<[%key_value_head_count]x[%key_value_block_count]x16xf32> + %block_sum = view.load %partial_sum_view[%key_value_head, %block, %query_row] : view<[%key_value_head_count]x[%key_value_block_count]x16xf32> -> f32 + %scaled_sum = scalar.mulf %block_sum, %scale : f32 + %next_sum = scalar.addf %running_sum, %scaled_sum : f32 + scf.yield %next_sum : f32 + } + scf.yield %sum : f32 + } else { + scf.yield %c0_f32 : f32 + } + %subgroup_sum = kernel.subgroup.reduce %lane_sum : f32 + scf.if %workitem_is_zero { + view.store %subgroup_sum, %reduction_stage_view[%c1] : f32, view<2xf32> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + %sum = view.load %reduction_stage_view[%c1] : view<2xf32> -> f32 + %unnormalized_output = scf.for %block = [%c0 to %key_value_block_count step %c1](%running_output = %c0_f32 : f32) -> (f32) unroll(%c4) schedule(interleaved) { + %scale = view.load %partial_max_view[%key_value_head, %block, %query_row] : view<[%key_value_head_count]x[%key_value_block_count]x16xf32> -> f32 + %block_output_f16 = view.load %partial_output_view[%key_value_head, %block, %query_row, %output_channel] : view<[%key_value_head_count]x[%key_value_block_count]x16x128xf16> -> f16 + %block_output = scalar.extf %block_output_f16 : f16 to f32 + %scaled_output = scalar.mulf %block_output, %scale : f32 + %next_output = scalar.addf %running_output, %scaled_output : f32 + scf.yield %next_output : f32 + } + %normalized_output = scalar.divf %unnormalized_output, %sum : f32 + view.store %normalized_output, %output_view[%query_head, %output_channel] : f32, view<[%query_head_count]x128xf32> + kernel.return +} + +check.case public @qwen3_moe_flash_attention_decode_split_f32_f16_wmma_case { + %key_value_token_count = check.literal value(256) : index + %query = check.generate.fill value(1.0) : tensor<32x128xf32> + %key = check.generate.fill value(1.0) : tensor<256x4x128xf16> + %value = check.generate.fill value(2.0) : tensor<256x4x128xf16> + %mask = check.generate.fill value(0.0) : tensor<256xf16> + %partial_max = check.generate.fill value(-1.0) : tensor<4x4x16xf32> + %partial_sum = check.generate.fill value(-1.0) : tensor<4x4x16xf32> + %partial_output = check.generate.fill value(-1.0) : tensor<4x4x16x128xf16> + %completion_counter = check.generate.fill value(0) : tensor<4xi32> + %output = check.generate.fill value(-1.0) : tensor<32x128xf32> + %expected = check.generate.fill value(2.0) : tensor<32x128xf32> + kernel.launch @qwen3_moe_flash_attention_decode_split_f32_f16_wmma[%key_value_token_count](%key_value_token_count, %query, %key, %value, %mask, %partial_max, %partial_sum, %partial_output, %completion_counter, %output) : [index](index, tensor<32x128xf32>, tensor<256x4x128xf16>, tensor<256x4x128xf16>, tensor<256xf16>, tensor<4x4x16xf32>, tensor<4x4x16xf32>, tensor<4x4x16x128xf16>, tensor<4xi32>, tensor<32x128xf32>) + check.expect.close actual(%output) expected(%expected) atol(0.001) rtol(0.001) nan(same) : tensor<32x128xf32> + check.return +} + +// Only KV row zero participates. The finite F32 iota step overflows to F16 +// negative infinity at every later row, leaving blocks one through three +// entirely masked. Each empty split must publish the online-softmax identity +// instead of evaluating -inf - -inf and contaminating the final reduction. +check.case public @qwen3_moe_flash_attention_decode_split_f32_f16_wmma_masked_blocks_case { + %key_value_token_count = check.literal value(256) : index + %query = check.generate.fill value(1.0) : tensor<32x128xf32> + %key = check.generate.fill value(1.0) : tensor<256x4x128xf16> + %value = check.generate.fill value(2.0) : tensor<256x4x128xf16> + %mask = check.generate.iota offset(0.0) step(-1e+30) : tensor<256xf16> + %partial_max = check.generate.fill value(-1.0) : tensor<4x4x16xf32> + %partial_sum = check.generate.fill value(-1.0) : tensor<4x4x16xf32> + %partial_output = check.generate.fill value(-1.0) : tensor<4x4x16x128xf16> + %completion_counter = check.generate.fill value(0) : tensor<4xi32> + %output = check.generate.fill value(-1.0) : tensor<32x128xf32> + %expected = check.generate.fill value(2.0) : tensor<32x128xf32> + kernel.launch @qwen3_moe_flash_attention_decode_split_f32_f16_wmma[%key_value_token_count](%key_value_token_count, %query, %key, %value, %mask, %partial_max, %partial_sum, %partial_output, %completion_counter, %output) : [index](index, tensor<32x128xf32>, tensor<256x4x128xf16>, tensor<256x4x128xf16>, tensor<256xf16>, tensor<4x4x16xf32>, tensor<4x4x16xf32>, tensor<4x4x16x128xf16>, tensor<4xi32>, tensor<32x128xf32>) + check.expect.close actual(%output) expected(%expected) atol(0.001) rtol(0.001) nan(same) : tensor<32x128xf32> + check.return +} + +// The production Decode-513 shape combines eight full blocks with a one-row +// tail and selects the cooperative fused reducer. Constant nonzero V keeps the +// expected result exact while all four KV heads, 32 GQA heads, completion +// counters, partial tensors, and tail guards participate. A second invocation +// reuses the partial and counter storage with a different V tensor, making the +// completion-counter reset observable. +check.case public @qwen3_moe_flash_attention_decode_split_f32_f16_wmma_decode_513_case { + %key_value_token_count = check.literal value(513) : index + %query = check.generate.fill value(1.0) : tensor<32x128xf32> + %key = check.generate.fill value(1.0) : tensor<513x4x128xf16> + %value0 = check.generate.fill value(2.0) : tensor<513x4x128xf16> + %value1 = check.generate.fill value(3.0) : tensor<513x4x128xf16> + %mask = check.generate.fill value(0.0) : tensor<513xf16> + %partial_max = check.generate.fill value(-1.0) : tensor<4x9x16xf32> + %partial_sum = check.generate.fill value(-1.0) : tensor<4x9x16xf32> + %partial_output = check.generate.fill value(-1.0) : tensor<4x9x16x128xf16> + %completion_counter = check.generate.fill value(0) : tensor<4xi32> + %output0 = check.generate.fill value(-1.0) : tensor<32x128xf32> + %output1 = check.generate.fill value(-1.0) : tensor<32x128xf32> + %expected0 = check.generate.fill value(2.0) : tensor<32x128xf32> + %expected1 = check.generate.fill value(3.0) : tensor<32x128xf32> + kernel.launch @qwen3_moe_flash_attention_decode_split_f32_f16_wmma[%key_value_token_count](%key_value_token_count, %query, %key, %value0, %mask, %partial_max, %partial_sum, %partial_output, %completion_counter, %output0) : [index](index, tensor<32x128xf32>, tensor<513x4x128xf16>, tensor<513x4x128xf16>, tensor<513xf16>, tensor<4x9x16xf32>, tensor<4x9x16xf32>, tensor<4x9x16x128xf16>, tensor<4xi32>, tensor<32x128xf32>) + kernel.launch @qwen3_moe_flash_attention_decode_split_f32_f16_wmma[%key_value_token_count](%key_value_token_count, %query, %key, %value1, %mask, %partial_max, %partial_sum, %partial_output, %completion_counter, %output1) : [index](index, tensor<32x128xf32>, tensor<513x4x128xf16>, tensor<513x4x128xf16>, tensor<513xf16>, tensor<4x9x16xf32>, tensor<4x9x16xf32>, tensor<4x9x16x128xf16>, tensor<4xi32>, tensor<32x128xf32>) + check.expect.close actual(%output0) expected(%expected0) atol(0.001) rtol(0.001) nan(same) : tensor<32x128xf32> + check.expect.close actual(%output1) expected(%expected1) atol(0.001) rtol(0.001) nan(same) : tensor<32x128xf32> + check.return +} + +// Sixty-five KV rows force a second split containing one valid row. The mask +// selects that final row, whose iota values begin at 1024, so an omitted or +// uninitialized tail cannot accidentally satisfy the check. +check.case public @qwen3_moe_flash_attention_decode_split_f32_f16_wmma_tail_case { + %key_value_token_count = check.literal value(65) : index + %query = check.generate.fill value(1.0) : tensor<1x128xf32> + %key = check.generate.fill value(1.0) : tensor<65x1x128xf16> + %value = check.generate.iota offset(0.0) step(0.125) : tensor<65x1x128xf16> + %mask = check.generate.iota offset(-64000.0) step(1000.0) : tensor<65xf16> + %partial_max = check.generate.fill value(-1.0) : tensor<1x2x16xf32> + %partial_sum = check.generate.fill value(-1.0) : tensor<1x2x16xf32> + %partial_output = check.generate.fill value(-1.0) : tensor<1x2x16x128xf16> + %completion_counter = check.generate.fill value(0) : tensor<1xi32> + %output = check.generate.fill value(-1.0) : tensor<1x128xf32> + %expected = check.generate.iota offset(1024.0) step(0.125) : tensor<1x128xf32> + kernel.launch @qwen3_moe_flash_attention_decode_split_f32_f16_wmma[%key_value_token_count](%key_value_token_count, %query, %key, %value, %mask, %partial_max, %partial_sum, %partial_output, %completion_counter, %output) : [index](index, tensor<1x128xf32>, tensor<65x1x128xf16>, tensor<65x1x128xf16>, tensor<65xf16>, tensor<1x2x16xf32>, tensor<1x2x16xf32>, tensor<1x2x16x128xf16>, tensor<1xi32>, tensor<1x128xf32>) + check.expect.close actual(%output) expected(%expected) atol(0.001) rtol(0.001) nan(same) : tensor<1x128xf32> + check.return +} + +// Thirty-two split-K blocks exercise the separate parallel reducer and the +// execution barrier between its producer and consumer dispatches. +check.case public @qwen3_moe_flash_attention_decode_split_f32_f16_wmma_long_case { + %key_value_token_count = check.literal value(2048) : index + %query = check.generate.fill value(0.0) : tensor<32x128xf32> + %key = check.generate.fill value(0.0) : tensor<2048x4x128xf16> + %value = check.generate.fill value(0.0) : tensor<2048x4x128xf16> + %mask = check.generate.fill value(0.0) : tensor<2048xf16> + %partial_max = check.generate.fill value(-1.0) : tensor<4x32x16xf32> + %partial_sum = check.generate.fill value(-1.0) : tensor<4x32x16xf32> + %partial_output = check.generate.fill value(-1.0) : tensor<4x32x16x128xf16> + %output = check.generate.fill value(1.0) : tensor<32x128xf32> + %expected = check.generate.fill value(0.0) : tensor<32x128xf32> + kernel.launch @qwen3_moe_flash_attention_decode_split_produce_partials_f32_f16_wmma[%key_value_token_count](%key_value_token_count, %query, %key, %value, %mask, %partial_max, %partial_sum, %partial_output) : [index](index, tensor<32x128xf32>, tensor<2048x4x128xf16>, tensor<2048x4x128xf16>, tensor<2048xf16>, tensor<4x32x16xf32>, tensor<4x32x16xf32>, tensor<4x32x16x128xf16>) + kernel.launch @qwen3_moe_flash_attention_decode_split_reduce_f32[%key_value_token_count](%key_value_token_count, %partial_max, %partial_sum, %partial_output, %output) : [index](index, tensor<4x32x16xf32>, tensor<4x32x16xf32>, tensor<4x32x16x128xf16>, tensor<32x128xf32>) + check.expect.close actual(%output) expected(%expected) atol(0.0) rtol(0.0) nan(same) : tensor<32x128xf32> + check.return +} + +check.case public @qwen3_moe_flash_attention_decode_split_f32_f16_wmma_benchmark_case { + %key_value_token_count = check.param.choice values([64, 65, 128, 256, 512, 513, 768, 1024, 1280, 2048]) name("key_value_token_count") : index + %query = check.generate.fill value(0.0) : tensor<32x128xf32> + %key = check.generate.fill value(0.0) : tensor<[%key_value_token_count]x4x128xf16> + %value = check.generate.fill value(0.0) : tensor<[%key_value_token_count]x4x128xf16> + %mask = check.generate.fill value(0.0) : tensor<[%key_value_token_count]xf16> + // Reserve the bounded scratch capacity once; each specialization addresses + // only ceildiv(key_value_token_count, 64) blocks. + %partial_max = check.generate.fill value(-1.0) : tensor<4x512x16xf32> + %partial_sum = check.generate.fill value(-1.0) : tensor<4x512x16xf32> + %partial_output = check.generate.fill value(-1.0) : tensor<4x512x16x128xf16> + %completion_counter = check.generate.fill value(0) : tensor<4xi32> + %output = check.generate.fill value(1.0) : tensor<32x128xf32> + %expected = check.generate.fill value(0.0) : tensor<32x128xf32> + kernel.launch @qwen3_moe_flash_attention_decode_split_f32_f16_wmma[%key_value_token_count](%key_value_token_count, %query, %key, %value, %mask, %partial_max, %partial_sum, %partial_output, %completion_counter, %output) : [index](index, tensor<32x128xf32>, tensor<[%key_value_token_count]x4x128xf16>, tensor<[%key_value_token_count]x4x128xf16>, tensor<[%key_value_token_count]xf16>, tensor<4x512x16xf32>, tensor<4x512x16xf32>, tensor<4x512x16x128xf16>, tensor<4xi32>, tensor<32x128xf32>) + check.expect.close actual(%output) expected(%expected) atol(0.0) rtol(0.0) nan(same) : tensor<32x128xf32> + check.return +} + +// Long-context execution is one reusable producer/reducer command buffer. The +// harness records an explicit dispatch execution barrier between these calls +// and profiles their complete device-side span as one semantic operation. +check.case public @qwen3_moe_flash_attention_decode_split_f32_f16_wmma_long_benchmark_case { + %key_value_token_count = check.param.choice values([2048, 32768]) name("key_value_token_count") : index + %query = check.generate.fill value(0.0) : tensor<32x128xf32> + %key = check.generate.fill value(0.0) : tensor<[%key_value_token_count]x4x128xf16> + %value = check.generate.fill value(0.0) : tensor<[%key_value_token_count]x4x128xf16> + %mask = check.generate.fill value(0.0) : tensor<[%key_value_token_count]xf16> + %partial_max = check.generate.fill value(-1.0) : tensor<4x512x16xf32> + %partial_sum = check.generate.fill value(-1.0) : tensor<4x512x16xf32> + %partial_output = check.generate.fill value(-1.0) : tensor<4x512x16x128xf16> + %output = check.generate.fill value(1.0) : tensor<32x128xf32> + %expected = check.generate.fill value(0.0) : tensor<32x128xf32> + kernel.launch @qwen3_moe_flash_attention_decode_split_produce_partials_f32_f16_wmma[%key_value_token_count](%key_value_token_count, %query, %key, %value, %mask, %partial_max, %partial_sum, %partial_output) : [index](index, tensor<32x128xf32>, tensor<[%key_value_token_count]x4x128xf16>, tensor<[%key_value_token_count]x4x128xf16>, tensor<[%key_value_token_count]xf16>, tensor<4x512x16xf32>, tensor<4x512x16xf32>, tensor<4x512x16x128xf16>) + kernel.launch @qwen3_moe_flash_attention_decode_split_reduce_f32[%key_value_token_count](%key_value_token_count, %partial_max, %partial_sum, %partial_output, %output) : [index](index, tensor<4x512x16xf32>, tensor<4x512x16xf32>, tensor<4x512x16x128xf16>, tensor<32x128xf32>) + check.expect.close actual(%output) expected(%expected) atol(0.0) rtol(0.0) nan(same) : tensor<32x128xf32> + check.return +} + +check.benchmark<@qwen3_moe_flash_attention_decode_split_f32_f16_wmma_benchmark_case> @qwen3_moe_flash_attention_decode_split_f32_f16_wmma_decode_64 {key_value_token_count = 64} + +check.benchmark<@qwen3_moe_flash_attention_decode_split_f32_f16_wmma_benchmark_case> @qwen3_moe_flash_attention_decode_split_f32_f16_wmma_decode_65 {key_value_token_count = 65} + +check.benchmark<@qwen3_moe_flash_attention_decode_split_f32_f16_wmma_benchmark_case> @qwen3_moe_flash_attention_decode_split_f32_f16_wmma_decode_128 {key_value_token_count = 128} + +check.benchmark<@qwen3_moe_flash_attention_decode_split_f32_f16_wmma_benchmark_case> @qwen3_moe_flash_attention_decode_split_f32_f16_wmma_decode_256 {key_value_token_count = 256} + +check.benchmark<@qwen3_moe_flash_attention_decode_split_f32_f16_wmma_masked_blocks_case> @qwen3_moe_flash_attention_decode_split_f32_f16_wmma_decode_256_masked_blocks + +check.benchmark<@qwen3_moe_flash_attention_decode_split_f32_f16_wmma_benchmark_case> @qwen3_moe_flash_attention_decode_split_f32_f16_wmma_decode_512 {key_value_token_count = 512} + +check.benchmark<@qwen3_moe_flash_attention_decode_split_f32_f16_wmma_benchmark_case> @qwen3_moe_flash_attention_decode_split_f32_f16_wmma_decode_513 {key_value_token_count = 513} + +check.benchmark<@qwen3_moe_flash_attention_decode_split_f32_f16_wmma_benchmark_case> @qwen3_moe_flash_attention_decode_split_f32_f16_wmma_decode_768 {key_value_token_count = 768} + +check.benchmark<@qwen3_moe_flash_attention_decode_split_f32_f16_wmma_benchmark_case> @qwen3_moe_flash_attention_decode_split_f32_f16_wmma_decode_1024 {key_value_token_count = 1024} + +check.benchmark<@qwen3_moe_flash_attention_decode_split_f32_f16_wmma_benchmark_case> @qwen3_moe_flash_attention_decode_split_f32_f16_wmma_decode_1280 {key_value_token_count = 1280} + +check.benchmark<@qwen3_moe_flash_attention_decode_split_f32_f16_wmma_benchmark_case> @qwen3_moe_flash_attention_decode_split_f32_f16_wmma_decode_2048_fused {key_value_token_count = 2048} + +check.benchmark<@qwen3_moe_flash_attention_decode_split_f32_f16_wmma_long_benchmark_case> @qwen3_moe_flash_attention_decode_split_f32_f16_wmma_decode_2048 {key_value_token_count = 2048} + +check.benchmark<@qwen3_moe_flash_attention_decode_split_f32_f16_wmma_long_benchmark_case> @qwen3_moe_flash_attention_decode_split_f32_f16_wmma_decode_32768 {key_value_token_count = 32768} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/flash_attention_decode_split_next_q8_test.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/flash_attention_decode_split_next_q8_test.loom new file mode 100644 index 000000000000..42203d79310c --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/flash_attention_decode_split_next_q8_test.loom @@ -0,0 +1,204 @@ +// Test-only differentials for fused decode-attention Q8 publication. The +// production kernels remain the actual providers; this module contributes only +// reference packing, masked metadata setup, and comparison cases. +func.def inline @ggml_quantize_q8_1_x4_group_body(%publish_output: i1, %group_count0: index, %group0: index, %input: buffer, %output: buffer) { + %group_count, %group = index.assume %group_count0, %group0 [range(%group_count0, 1, 524288), lt(%group0, %group_count0)] : index, index + %lane = kernel.workitem.id : index + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c4 = index.constant 4 : index + %c8 = index.constant 8 : index + %c32 = index.constant 32 : index + %c128 = index.constant 128 : index + %group_bytes = index.constant 144 : offset + %payload_byte_add = index.constant 16 : offset + %scratch_d_byte_add = index.constant 128 : offset + %scratch_bytes = index.constant 144 : offset + %c0_f32 = scalar.constant 0.0 : f32 + %c1_f32 = scalar.constant 1.0 : f32 + %c127 = scalar.constant 127.0 : f32 + %c0_offset = index.constant 0 : offset + %launched_element_count = index.mul %group_count, %c128 : index + %block_in_group0 = index.div %lane, %c8 : index + %block_in_group = index.assume %block_in_group0 [range(%block_in_group0, 0, 3)] : index + %word_in_block0 = index.rem %lane, %c8 : index + %word_in_block = index.assume %word_in_block0 [range(%word_in_block0, 0, 7)] : index + %group_element_base = index.mul %group, %c128 : index + %block_element_add = index.mul %block_in_group, %c32 : index + %block_word_add = index.mul %block_in_group, %c8 : index + %word_element_add = index.mul %word_in_block, %c4 : index + %input_block_base = index.add %group_element_base, %block_element_add : index + %input_index = index.add %input_block_base, %word_element_add : index + %input_noalias, %output_noalias = buffer.assume.noalias %input, %output : buffer, buffer + %input_view = buffer.view %input_noalias[%c0_offset] : buffer -> view<[%launched_element_count]xf32> + %input_values = vector.load %input_view[%input_index] : view<[%launched_element_count]xf32> -> vector<4xf32> + %absolute_values = vector.absf %input_values : vector<4xf32> + %thread_max = vector.reduce %absolute_values, %c0_f32 : vector<4xf32>, f32 + %scratch = buffer.alloca align(16) %scratch_bytes : buffer + %scratch_values = buffer.view %scratch[%c0_offset] : buffer -> view<32xf32> + %scratch_d = buffer.view %scratch[%scratch_d_byte_add] : buffer -> view<4xf32> + view.store %thread_max, %scratch_values[%lane] : f32, view<32xf32> + kernel.barrier scope(workgroup) ordering(acq_rel) + %is_cohort_leader = index.cmp eq, %word_in_block, %c0 : index + scf.if %is_cohort_leader { + %cohort_base = index.mul %block_in_group, %c8 : index + %cohort_maxima = vector.load %scratch_values[%cohort_base] : view<32xf32> -> vector<8xf32> + %amax = vector.reduce %cohort_maxima, %c0_f32 : vector<8xf32>, f32 + %d = scalar.divf %amax, %c127 : f32 + view.store %d, %scratch_d[%block_in_group] : f32, view<4xf32> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + %d = view.load %scratch_d[%block_in_group] : view<4xf32> -> f32 + %d_nonzero = scalar.cmpf one, %d, %c0_f32 : f32 + %d_inverse = scf.if %d_nonzero -> (f32) { + %inverse = scalar.divf %c1_f32, %d : f32 + scf.yield %inverse : f32 + } else { + scf.yield %c0_f32 : f32 + } + %d_inverse_vector = vector.splat %d_inverse : vector<4xf32> + %scaled_values = vector.mulf %input_values, %d_inverse_vector : vector<4xf32> + %rounded_values = vector.roundf %scaled_values : vector<4xf32> + %quantized_values = vector.fptosi %rounded_values : vector<4xf32> to vector<4xi8> + %packed_word = vector.bitcast %quantized_values : vector<4xi8> to vector<1xi32> + %group_byte_offset = index.scale %group, %group_bytes : index, offset -> offset + %payload_byte_offset = index.add %group_byte_offset, %payload_byte_add : offset + %group_ds = buffer.view %output_noalias[%group_byte_offset] : buffer -> view<8xf16> + %group_qs = buffer.view %output_noalias[%payload_byte_offset] : buffer -> view<32xi32> + %packed_word_index0 = index.add %block_word_add, %word_in_block : index + %packed_word_index = index.assume %packed_word_index0 [range(%packed_word_index0, 0, 31)] : index + scf.if %publish_output { + vector.store %packed_word, %group_qs[%packed_word_index] : vector<1xi32>, view<32xi32> + } + %thread_sum = vector.reduce %rounded_values, %c0_f32 : vector<4xf32>, f32 + view.store %thread_sum, %scratch_values[%lane] : f32, view<32xf32> + kernel.barrier scope(workgroup) ordering(acq_rel) + %publishes_metadata = scalar.andi %is_cohort_leader, %publish_output : i1 + scf.if %publishes_metadata { + %cohort_base = index.mul %block_in_group, %c8 : index + %cohort_sums = vector.load %scratch_values[%cohort_base] : view<32xf32> -> vector<8xf32> + %quantized_sum = vector.reduce %cohort_sums, %c0_f32 : vector<8xf32>, f32 + %s = scalar.mulf %quantized_sum, %d : f32 + %d_f16 = scalar.fptrunc %d : f32 to f16 + %s_f16 = scalar.fptrunc %s : f32 to f16 + %ds_index = index.mul %block_in_group, %c2 : index + view.store %d_f16, %group_ds[%ds_index] : f16, view<8xf16> + %s_index = index.add %ds_index, %c1 : index + view.store %s_f16, %group_ds[%s_index] : f16, view<8xf16> + } + func.return +} + +target.decl @qwen3_moe_decode_split_gfx11_wave64 + +kernel.decl target(@qwen3_moe_decode_split_gfx11_wave64) @qwen3_moe_flash_attention_decode_split_f32_f16_wmma(%key_value_token_count: index) launch(%key_value_token_count: index, %query: buffer, %key: buffer, %value: buffer, %mask: buffer, %partial_max: buffer, %partial_sum: buffer, %partial_output: buffer, %completion_counter: buffer, %output: buffer) + +kernel.decl target(@qwen3_moe_decode_split_gfx11_wave64) @qwen3_moe_flash_attention_decode_split_f32_f16_wmma_next_q8(%key_value_token_count: index) launch(%key_value_token_count: index, %query: buffer, %key: buffer, %value: buffer, %mask: buffer, %partial_max: buffer, %partial_sum: buffer, %partial_output: buffer, %completion_counter: buffer, %output: buffer, %next_q8_output: buffer) + +// Publishes the canonical reference Q8_1 x4 layout for one 4096-element row. +func.decl @qwen3_moe_flash_attention_decode_pack_completed_key_value_head_q8_1_x4(%key_value_head: index, %output: buffer, %q8_output: buffer) + +kernel.def @qwen3_moe_flash_attention_decode_split_quantize_reference_4096() { + %c1 = index.constant 1 : index + %c32 = index.constant 32 : index + kernel.launch.config workgroups(%c32, %c1, %c1) workgroup_size(%c32, %c1, %c1) : index +} launch(%input: buffer, %output: buffer) { + %publish_output = scalar.constant true : i1 + %group_count = index.constant 32 : index + %group = kernel.workgroup.id : index + func.call @ggml_quantize_q8_1_x4_group_body(%publish_output, %group_count, %group, %input, %output) : (i1, index, index, buffer, buffer) + kernel.return +} + +// Isolates the new publication phase for access-sanitized coverage without +// inflating the complete fused attention kernel beyond a short branch's range. +kernel.def target(@qwen3_moe_decode_split_gfx11_wave64) @qwen3_moe_flash_attention_decode_split_pack_completed_q8_test() { + %c1 = index.constant 1 : index + %c4 = index.constant 4 : index + %c256 = index.constant 256 : index + kernel.launch.config workgroups(%c1, %c4, %c1) workgroup_size(%c256, %c1, %c1) : index +} launch(%input: buffer, %output: buffer) { + %key_value_head = kernel.workgroup.id : index + func.call @qwen3_moe_flash_attention_decode_pack_completed_key_value_head_q8_1_x4(%key_value_head, %input, %output) : (index, buffer, buffer) + kernel.return +} + +// Models the request-owned Decode-513 mask inside its 576-row capacity class. +kernel.def @qwen3_moe_flash_attention_decode_split_mask_513_of_576() { + %c1 = index.constant 1 : index + %c64 = index.constant 64 : index + kernel.launch.config workgroups(%c1, %c1, %c1) workgroup_size(%c64, %c1, %c1) : index +} launch(%mask: buffer) { + %workitem = kernel.workitem.id : index + %c63 = index.constant 63 : index + %c513 = index.constant 513 : index + %c0 = index.constant 0 : index + %c0_offset = index.constant 0 : offset + %negative_large = scalar.constant -1e+30 : f32 + %negative_infinity = scalar.fptrunc %negative_large : f32 to f16 + %valid = index.cmp ult, %workitem, %c63 : index + %row0 = index.add %c513, %workitem : index + %row = scf.select %valid, %row0, %c0 : index + %mask_view = buffer.view %mask[%c0_offset] : buffer -> view<576xf16> + scf.if %valid { + view.store %negative_infinity, %mask_view[%row] : f16, view<576xf16> + } + kernel.return +} + +// The fused producer must match the ordinary attention export followed by the +// canonical GGML packer for every F32 value and every packed byte. Two value +// tensors reuse both paths' partials and completion counters, making counter +// reset and stale-publication failures observable. +check.case public @qwen3_moe_flash_attention_decode_split_next_q8_capacity_576_differential_case { + %key_value_token_count = check.literal value(576) : index + %query_seed = check.param.seed base(0x514d4f4551465138) count(1) : i64 + %key_seed = check.param.seed base(0x514d4f454b465138) count(1) : i64 + %value_seed0 = check.param.seed base(0x514d4f4556465130) count(1) : i64 + %value_seed1 = check.param.seed base(0x514d4f4556465131) count(1) : i64 + %query = check.generate.random.uniform seed(%query_seed) range(-1.0 to 1.0) : tensor<32x128xf32> + %key = check.generate.random.uniform seed(%key_seed) range(-1.0 to 1.0) : tensor<576x4x128xf16> + %value0 = check.generate.random.uniform seed(%value_seed0) range(-1.0 to 1.0) : tensor<576x4x128xf16> + %value1 = check.generate.random.uniform seed(%value_seed1) range(-1.0 to 1.0) : tensor<576x4x128xf16> + %mask = check.generate.fill value(0.0) : tensor<576xf16> + %reference_partial_max = check.generate.fill value(-1.0) : tensor<4x9x16xf32> + %reference_partial_sum = check.generate.fill value(-1.0) : tensor<4x9x16xf32> + %reference_partial_output = check.generate.fill value(-1.0) : tensor<4x9x16x128xf16> + %reference_counter = check.generate.fill value(0) : tensor<4xi32> + %actual_partial_max = check.generate.fill value(-1.0) : tensor<4x9x16xf32> + %actual_partial_sum = check.generate.fill value(-1.0) : tensor<4x9x16xf32> + %actual_partial_output = check.generate.fill value(-1.0) : tensor<4x9x16x128xf16> + %actual_counter = check.generate.fill value(0) : tensor<4xi32> + %reference_output0 = check.generate.fill value(-1.0) : tensor<32x128xf32> + %reference_output1 = check.generate.fill value(-1.0) : tensor<32x128xf32> + %actual_output0 = check.generate.fill value(-2.0) : tensor<32x128xf32> + %actual_output1 = check.generate.fill value(-2.0) : tensor<32x128xf32> + %reference_q8_0 = check.generate.fill value(0) : tensor<4608xi8> + %reference_q8_1 = check.generate.fill value(0) : tensor<4608xi8> + %actual_q8_0 = check.generate.fill value(1) : tensor<4608xi8> + %actual_q8_1 = check.generate.fill value(1) : tensor<4608xi8> + kernel.launch @qwen3_moe_flash_attention_decode_split_mask_513_of_576(%mask) : (tensor<576xf16>) + kernel.launch @qwen3_moe_flash_attention_decode_split_f32_f16_wmma[%key_value_token_count](%key_value_token_count, %query, %key, %value0, %mask, %reference_partial_max, %reference_partial_sum, %reference_partial_output, %reference_counter, %reference_output0) : [index](index, tensor<32x128xf32>, tensor<576x4x128xf16>, tensor<576x4x128xf16>, tensor<576xf16>, tensor<4x9x16xf32>, tensor<4x9x16xf32>, tensor<4x9x16x128xf16>, tensor<4xi32>, tensor<32x128xf32>) + kernel.launch @qwen3_moe_flash_attention_decode_split_quantize_reference_4096(%reference_output0, %reference_q8_0) : (tensor<32x128xf32>, tensor<4608xi8>) + kernel.launch @qwen3_moe_flash_attention_decode_split_f32_f16_wmma_next_q8[%key_value_token_count](%key_value_token_count, %query, %key, %value0, %mask, %actual_partial_max, %actual_partial_sum, %actual_partial_output, %actual_counter, %actual_output0, %actual_q8_0) : [index](index, tensor<32x128xf32>, tensor<576x4x128xf16>, tensor<576x4x128xf16>, tensor<576xf16>, tensor<4x9x16xf32>, tensor<4x9x16xf32>, tensor<4x9x16x128xf16>, tensor<4xi32>, tensor<32x128xf32>, tensor<4608xi8>) + kernel.launch @qwen3_moe_flash_attention_decode_split_f32_f16_wmma[%key_value_token_count](%key_value_token_count, %query, %key, %value1, %mask, %reference_partial_max, %reference_partial_sum, %reference_partial_output, %reference_counter, %reference_output1) : [index](index, tensor<32x128xf32>, tensor<576x4x128xf16>, tensor<576x4x128xf16>, tensor<576xf16>, tensor<4x9x16xf32>, tensor<4x9x16xf32>, tensor<4x9x16x128xf16>, tensor<4xi32>, tensor<32x128xf32>) + kernel.launch @qwen3_moe_flash_attention_decode_split_quantize_reference_4096(%reference_output1, %reference_q8_1) : (tensor<32x128xf32>, tensor<4608xi8>) + kernel.launch @qwen3_moe_flash_attention_decode_split_f32_f16_wmma_next_q8[%key_value_token_count](%key_value_token_count, %query, %key, %value1, %mask, %actual_partial_max, %actual_partial_sum, %actual_partial_output, %actual_counter, %actual_output1, %actual_q8_1) : [index](index, tensor<32x128xf32>, tensor<576x4x128xf16>, tensor<576x4x128xf16>, tensor<576xf16>, tensor<4x9x16xf32>, tensor<4x9x16xf32>, tensor<4x9x16x128xf16>, tensor<4xi32>, tensor<32x128xf32>, tensor<4608xi8>) + check.expect.equal actual(%actual_output0) expected(%reference_output0) : tensor<32x128xf32> + check.expect.equal actual(%actual_q8_0) expected(%reference_q8_0) : tensor<4608xi8> + check.expect.equal actual(%actual_output1) expected(%reference_output1) : tensor<32x128xf32> + check.expect.equal actual(%actual_q8_1) expected(%reference_q8_1) : tensor<4608xi8> + check.return +} + +check.case public @qwen3_moe_flash_attention_decode_split_pack_completed_q8_differential_case { + %input_seed = check.param.seed base(0x514d4f4550385138) count(1) : i64 + %input = check.generate.random.uniform seed(%input_seed) range(-1.0 to 1.0) : tensor<32x128xf32> + %reference = check.generate.fill value(0) : tensor<4608xi8> + %actual = check.generate.fill value(1) : tensor<4608xi8> + kernel.launch @qwen3_moe_flash_attention_decode_split_quantize_reference_4096(%input, %reference) : (tensor<32x128xf32>, tensor<4608xi8>) + kernel.launch @qwen3_moe_flash_attention_decode_split_pack_completed_q8_test(%input, %actual) : (tensor<32x128xf32>, tensor<4608xi8>) + check.expect.equal actual(%actual) expected(%reference) : tensor<4608xi8> + check.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/flash_attention_f32_f16_wmma.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/flash_attention_f32_f16_wmma.loom new file mode 100644 index 000000000000..d0b97ce7ee0f --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/flash_attention_f32_f16_wmma.loom @@ -0,0 +1,870 @@ +// Qwen3 MoE grouped-query FlashAttention for the F32-query/F16-cache path. +// +// One four-wave workgroup computes 16 query rows for one query head against +// 64 KV rows at a time. The ownership changes mirror the cooperative-matrix +// schedule used by llama.cpp's Vulkan CM1 kernel: +// +// 1. All workitems stage a scaled 16x128 F16 query tile. +// 2. Each wave computes one 16x16 QK score slice. +// 3. Scores cross LDS so each wave can normalize four complete query rows. +// 4. F16 probabilities cross LDS for four P*V WMMA steps. +// 5. Each active lane retains one four-channel F16 packet for each of its +// four query rows across subsequent 64-row KV blocks. +// +// K and V remain in the row-major llama.cpp cache layout +// [KV token][KV head][128]. Their aligned F16 fragments load directly from +// global memory; there is no expanded or repacked persistent allocation. QK +// and the online-softmax statistics remain F32, while the P*V accumulation and +// carried output match the Vulkan oracle's F16 policy. +amdgpu.target @qwen3_moe_attention_gfx11_wave64 {subgroup_size = 64} + +config.decl @qwen3_moe.attention.query_head_count : %value: index where [range(%value, 1, 64)] + +config.decl @qwen3_moe.attention.key_value_head_count : %value: index where [range(%value, 1, 64)] + +// Bounds the context copied by the test-only row-extraction kernel below. +config.decl @qwen3_moe.attention.test.context_capacity : %value: index where [range(%value, 1, 32768)] + +kernel.def target(@qwen3_moe_attention_gfx11_wave64) @qwen3_moe_flash_attention_f32_f16_wmma(%query_token_count: index, %key_value_token_count: index) { + %query_head_count = config.get @qwen3_moe.attention.query_head_count : index + %c1 = index.constant 1 : index + %c15 = index.constant 15 : index + %c16 = index.constant 16 : index + %c256 = index.constant 256 : index + %padded_query_token_count = index.add %query_token_count, %c15 : index + %query_tile_count = index.div %padded_query_token_count, %c16 : index + kernel.launch.config workgroups(%query_tile_count, %query_head_count, %c1) workgroup_size(%c256, %c1, %c1) : index +} launch(%query_token_count: index, %key_value_token_count: index, %query: buffer, %key: buffer, %value: buffer, %mask: buffer, %output: buffer) where [range(%query_token_count, 1, 2048)] { + %bounded_key_value_token_count = index.assume %key_value_token_count [range(%key_value_token_count, 1, 32768)] : index + %query_head_count = config.get @qwen3_moe.attention.query_head_count : index + %key_value_head_count = config.get @qwen3_moe.attention.key_value_head_count : index + %query_tile = kernel.workgroup.id : index + %query_head0 = kernel.workgroup.id : index + %query_head, %launch_query_head_count = index.assume %query_head0, %query_head_count [lt(%query_head0, %query_head_count)] : index, index + %workitem = kernel.workitem.id : index + %subgroup0 = kernel.subgroup.id : index + %subgroup = index.assume %subgroup0 [range(%subgroup0, 0, 3)] : index + %lane = kernel.subgroup.lane.id : index + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c3 = index.constant 3 : index + %c4 = index.constant 4 : index + %c8 = index.constant 8 : index + %c15 = index.constant 15 : index + %c16 = index.constant 16 : index + %c32 = index.constant 32 : index + %c64 = index.constant 64 : index + %c128 = index.constant 128 : index + %c256 = index.constant 256 : index + %c0_offset = index.constant 0 : offset + %query_stage_bytes = index.constant 4352 : offset + %score_stage_bytes = index.constant 6144 : offset + %probability_stage_bytes = index.constant 3072 : offset + %product_stage_bytes = index.constant 2048 : offset + %tail_key_value_stage_capacity = index.constant 8192 : offset + %c0_f32 = scalar.constant 0.0 : f32 + %negative_large = scalar.constant -1e+30 : f32 + %head_size_f32 = scalar.constant 128.0 : f32 + %attention_scale = scalar.rsqrtf %head_size_f32 : f32 + %c0_f16 = scalar.constant 0.0 : f16 + %c1_f32 = scalar.constant 1.0 : f32 + %output_zero0 = vector.constant 0.0 : vector<4xf16> + %output_zero1 = vector.constant 0.0 : vector<4xf16> + %output_zero2 = vector.constant 0.0 : vector<4xf16> + %output_zero3 = vector.constant 0.0 : vector<4xf16> + %c0_f16x8 = vector.constant 0.0 : vector<8xf16> + %c0_f32x4 = vector.constant 0.0 : vector<4xf32> + %negative_f32x4 = vector.constant -1e+30 : vector<4xf32> + %m = index.constant 16 : index + %n = index.constant 16 : index + %k = index.constant 16 : index + %query_heads_per_key_value_head = index.div %query_head_count, %key_value_head_count : index + %key_value_head = index.div %query_head, %query_heads_per_key_value_head : index + %key_value_width = index.mul %key_value_head_count, %c128 : index + %key_value_head_base = index.mul %key_value_head, %c128 : index + %query_origin = index.mul %query_tile, %c16 : index + %full_key_value_block_count = index.div %bounded_key_value_token_count, %c64 : index + %full_key_value_token_count0 = index.mul %full_key_value_block_count, %c64 : index + %full_key_value_token_count = index.assume %full_key_value_token_count0 [range(%full_key_value_token_count0, 0, 32768), mul(%full_key_value_token_count0, 64)] : index + %tail_key_value_token_count = index.sub %bounded_key_value_token_count, %full_key_value_token_count : index + %has_key_value_tail = index.cmp ne, %tail_key_value_token_count, %c0 : index + %tail_key_value_stage_bytes = scf.select %has_key_value_tail, %tail_key_value_stage_capacity, %c0_offset : offset + %subgroup_score_column = index.mul %subgroup, %c16 : index + %subgroup_query_row = index.mul %subgroup, %c4 : index + %query_row0 = index.add %subgroup_query_row, %c0 : index + %query_row1 = index.add %subgroup_query_row, %c1 : index + %query_row2 = index.add %subgroup_query_row, %c2 : index + %query_row3 = index.add %subgroup_query_row, %c3 : index + %query_token0 = index.add %query_origin, %query_row0 : index + %query_token1 = index.add %query_origin, %query_row1 : index + %query_token2 = index.add %query_origin, %query_row2 : index + %query_token3 = index.add %query_origin, %query_row3 : index + %query_valid0 = index.cmp ult, %query_token0, %query_token_count : index + %query_valid1 = index.cmp ult, %query_token1, %query_token_count : index + %query_valid2 = index.cmp ult, %query_token2, %query_token_count : index + %query_valid3 = index.cmp ult, %query_token3, %query_token_count : index + %query_valid = vector.from_elements %query_valid0, %query_valid1, %query_valid2, %query_valid3 : vector<4xi1> + %subgroup_product_channel = index.mul %subgroup, %c16 : index + %lane_output_tile = index.div %lane, %c16 : index + %lane_product_channel0 = index.rem %lane, %c16 : index + %lane_product_channel = index.mul %lane_product_channel0, %c4 : index + %lane_output_channel = index.mul %lane, %c4 : index + %lane_has_output = index.cmp ult, %lane, %c32 : index + // The padded LDS rows mirror the Vulkan oracle's ownership changes. Q uses + // eight spare F16 columns after its 128 channels. Score and probability + // transpose to key-major rows with eight spare columns after 16 queries. + // These strides avoid the bank pattern produced by dense transposed rows. + %query_transposed_layout = encoding.layout.strided [1, 136] : encoding + %probability_transposed_layout = encoding.layout.strided [1, 24] : encoding + %query_noalias, %key_noalias, %value_noalias, %mask_noalias, %output_noalias = buffer.assume.noalias %query, %key, %value, %mask, %output : buffer, buffer, buffer, buffer, buffer + %query_aligned = buffer.assume.alignment %query_noalias {minimum_alignment = 16} : buffer + %key_aligned = buffer.assume.alignment %key_noalias {minimum_alignment = 16} : buffer + %value_aligned = buffer.assume.alignment %value_noalias {minimum_alignment = 16} : buffer + %mask_aligned = buffer.assume.alignment %mask_noalias {minimum_alignment = 16} : buffer + %output_aligned = buffer.assume.alignment %output_noalias {minimum_alignment = 16} : buffer + %query_view = buffer.view %query_aligned[%c0_offset] : buffer -> view<[%query_token_count]x[%query_head_count]x128xf32> + %key_view = buffer.view %key_aligned[%c0_offset] : buffer -> view<[%bounded_key_value_token_count]x[%key_value_width]xf16> + %value_view = buffer.view %value_aligned[%c0_offset] : buffer -> view<[%bounded_key_value_token_count]x[%key_value_width]xf16> + %mask_view = buffer.view %mask_aligned[%c0_offset] : buffer -> view<[%query_token_count]x[%bounded_key_value_token_count]xf16> + %output_view = buffer.view %output_aligned[%c0_offset] : buffer -> view<[%query_token_count]x[%query_head_count]x128xf32> + %query_stage = buffer.alloca align(16) %query_stage_bytes : buffer + %score_stage = buffer.alloca align(16) %score_stage_bytes : buffer + %probability_stage = buffer.alloca align(16) %probability_stage_bytes : buffer + %product_stage = buffer.alloca align(16) %product_stage_bytes : buffer + %tail_key_value_stage = buffer.alloca align(16) %tail_key_value_stage_bytes : buffer + %query_stage_view = buffer.view %query_stage[%c0_offset] : buffer -> view<16x136xf16> + %query_transposed_view = buffer.view %query_stage[%c0_offset] : buffer -> view<128x16xf16, %query_transposed_layout> + %score_stage_view = buffer.view %score_stage[%c0_offset] : buffer -> view<64x24xf32> + %mask_summary_view = buffer.view %score_stage[%c0_offset] : buffer -> view<4xf32> + %probability_stage_view = buffer.view %probability_stage[%c0_offset] : buffer -> view<16x64xf16, %probability_transposed_layout> + %product_stage_view = buffer.view %product_stage[%c0_offset] : buffer -> view<16x64xf16> + %tail_key_value_stage_view = buffer.view %tail_key_value_stage[%c0_offset] : buffer -> view<32x128xf16> + // Scale and truncate Q exactly once. The Vulkan reference does this before + // entering its KV loop, making QK a native F16 WMMA while retaining F32 + // accumulation. + scf.for %load_iteration = [%c0 to %c8 step %c1] unroll { + %linear = index.madd %load_iteration, %c256, %workitem : index + %local_query_row = index.div %linear, %c128 : index + %query_channel = index.rem %linear, %c128 : index + %query_token = index.add %query_origin, %local_query_row : index + %query_valid_load = index.cmp ult, %query_token, %query_token_count : index + %query_value = scf.if %query_valid_load -> (f16) { + %loaded = view.load %query_view[%query_token, %query_head, %query_channel] : view<[%query_token_count]x[%query_head_count]x128xf32> -> f32 + %scaled = scalar.mulf %loaded, %attention_scale : f32 + %truncated = scalar.fptrunc %scaled : f32 to f16 + scf.yield %truncated : f16 + } else { + scf.yield %c0_f16 : f16 + } + view.store %query_value, %query_stage_view[%local_query_row, %query_channel] : f16, view<16x136xf16> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + %full_max, %full_sum, %full_output0, %full_output1, %full_output2, %full_output3 = scf.for %key_origin = [%c0 to %full_key_value_token_count step %c64](%current_max = %negative_f32x4 : vector<4xf32>, %current_sum = %c0_f32x4 : vector<4xf32>, %current_output0 = %output_zero0 : vector<4xf16>, %current_output1 = %output_zero1 : vector<4xf16>, %current_output2 = %output_zero2 : vector<4xf16>, %current_output3 = %output_zero3 : vector<4xf16>) -> (vector<4xf32>, vector<4xf32>, vector<4xf16>, vector<4xf16>, vector<4xf16>, vector<4xf16>) { + // Cache the complete 16x64 mask tile before QK. Causal masks leave future + // KV blocks entirely at negative infinity; reducing the cached tile lets + // the workgroup skip QK, softmax, and P*V for those blocks. This mirrors + // the Vulkan oracle without requiring its auxiliary compact mask buffer. + %key_token0 = index.add %key_origin, %lane : index + %key_token = index.assume %key_token0 [lt(%key_token0, %bounded_key_value_token_count)] : index + %mask0 = scf.if %query_valid0 -> (f16) { + %value = view.load %mask_view[%query_token0, %key_token] : view<[%query_token_count]x[%bounded_key_value_token_count]xf16> -> f16 + scf.yield %value : f16 + } else { + scf.yield %c0_f16 : f16 + } + %mask1 = scf.if %query_valid1 -> (f16) { + %value = view.load %mask_view[%query_token1, %key_token] : view<[%query_token_count]x[%bounded_key_value_token_count]xf16> -> f16 + scf.yield %value : f16 + } else { + scf.yield %c0_f16 : f16 + } + %mask2 = scf.if %query_valid2 -> (f16) { + %value = view.load %mask_view[%query_token2, %key_token] : view<[%query_token_count]x[%bounded_key_value_token_count]xf16> -> f16 + scf.yield %value : f16 + } else { + scf.yield %c0_f16 : f16 + } + %mask3 = scf.if %query_valid3 -> (f16) { + %value = view.load %mask_view[%query_token3, %key_token] : view<[%query_token_count]x[%bounded_key_value_token_count]xf16> -> f16 + scf.yield %value : f16 + } else { + scf.yield %c0_f16 : f16 + } + %mask_f16 = vector.from_elements %mask0, %mask1, %mask2, %mask3 : vector<4xf16> + %mask_summary_f32 = vector.extf %mask_f16 : vector<4xf16> to vector<4xf32> + %effective_mask = vector.select %query_valid, %mask_summary_f32, %negative_f32x4 : vector<4xf32> + %lane_mask_maximum = vector.reduce %effective_mask, %negative_large : vector<4xf32>, f32 + %subgroup_mask_maximum = kernel.subgroup.reduce %lane_mask_maximum : f32 + %is_subgroup_leader = index.cmp eq, %lane, %c0 : index + scf.if %is_subgroup_leader { + view.store %subgroup_mask_maximum, %mask_summary_view[%subgroup] : f32, view<4xf32> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + %subgroup_mask_maxima = vector.load %mask_summary_view[%c0] : view<4xf32> -> vector<4xf32> + %workgroup_mask_maximum = vector.reduce %subgroup_mask_maxima, %negative_large : vector<4xf32>, f32 + %block_has_attention = scalar.cmpf ogt, %workgroup_mask_maximum, %negative_large : f32 + %next_block_max, %next_block_sum, %next_block_output0, %next_block_output1, %next_block_output2, %next_block_output3 = scf.if %block_has_attention -> (vector<4xf32>, vector<4xf32>, vector<4xf16>, vector<4xf16>, vector<4xf16>, vector<4xf16>) { + // Four independent wave-level WMMAs produce a 16x64 score tile. + %score_key_origin0 = index.add %key_origin, %subgroup_score_column : index + %last_full_key_tile_start = index.sub %bounded_key_value_token_count, %c15 : index + %score_key_origin = index.assume %score_key_origin0 [lt(%score_key_origin0, %last_full_key_tile_start)] : index + %score_init_values = vector.constant 0.0 : vector<4xf32> + %score_init = vector.fragment %score_init_values shape [%m, %n] : vector<4xf32> + %score_fragment = scf.for %head_tile = [%c0 to %c128 step %c16](%score_accumulator = %score_init : vector<4xf32>) -> (vector<4xf32>) unroll schedule(recurrence) { + %key_channel = index.add %key_value_head_base, %head_tile : index + %key_fragment = vector.fragment.load %key_view[%score_key_origin, %key_channel] shape [%m, %k] : view<[%bounded_key_value_token_count]x[%key_value_width]xf16> -> vector<16xf16> + %query_fragment = vector.fragment.load %query_transposed_view[%head_tile, %c0] shape [%k, %n] : view<128x16xf16, %query_transposed_layout> -> vector<16xf16> + %next_score_accumulator = vector.mma %key_fragment, %query_fragment, %score_accumulator : vector<16xf16>, vector<16xf16>, vector<4xf32> + scf.yield %next_score_accumulator : vector<4xf32> + } + vector.fragment.store %score_fragment, %score_stage_view[%subgroup_score_column, %c0] shape [%m, %n] : vector<4xf32>, view<64x24xf32> + kernel.barrier scope(workgroup) ordering(acq_rel) + // LDS transposes ownership from one 16-column score slice per wave to + // four complete query rows per wave. Every lane contributes one key + // column to each of those rows. + %raw_score0 = view.load %score_stage_view[%lane, %query_row0] : view<64x24xf32> -> f32 + %raw_score1 = view.load %score_stage_view[%lane, %query_row1] : view<64x24xf32> -> f32 + %raw_score2 = view.load %score_stage_view[%lane, %query_row2] : view<64x24xf32> -> f32 + %raw_score3 = view.load %score_stage_view[%lane, %query_row3] : view<64x24xf32> -> f32 + %mask_f32 = vector.extf %mask_f16 : vector<4xf16> to vector<4xf32> + %raw_scores = vector.from_elements %raw_score0, %raw_score1, %raw_score2, %raw_score3 : vector<4xf32> + %masked_scores0 = vector.addf %raw_scores, %mask_f32 : vector<4xf32> + %masked_scores = vector.select %query_valid, %masked_scores0, %negative_f32x4 : vector<4xf32> + %block_max = kernel.subgroup.reduce %masked_scores : vector<4xf32> + %next_max = vector.maxnumf %current_max, %block_max : vector<4xf32> + %score_delta = vector.subf %masked_scores, %next_max : vector<4xf32> + %raw_probability = vector.expf %score_delta : vector<4xf32> + %probability = vector.select %query_valid, %raw_probability, %c0_f32x4 : vector<4xf32> + %block_sum = kernel.subgroup.reduce %probability : vector<4xf32> + %old_delta = vector.subf %current_max, %next_max : vector<4xf32> + %old_scale = vector.expf %old_delta : vector<4xf32> + %scaled_current_sum = vector.mulf %current_sum, %old_scale : vector<4xf32> + %next_sum = vector.addf %scaled_current_sum, %block_sum : vector<4xf32> + %probability_f16 = vector.fptrunc %probability : vector<4xf32> to vector<4xf16> + %probability0 = vector.extract %probability_f16[0] : vector<4xf16> -> f16 + %probability1 = vector.extract %probability_f16[1] : vector<4xf16> -> f16 + %probability2 = vector.extract %probability_f16[2] : vector<4xf16> -> f16 + %probability3 = vector.extract %probability_f16[3] : vector<4xf16> -> f16 + view.store %probability0, %probability_stage_view[%query_row0, %lane] : f16, view<16x64xf16, %probability_transposed_layout> + view.store %probability1, %probability_stage_view[%query_row1, %lane] : f16, view<16x64xf16, %probability_transposed_layout> + view.store %probability2, %probability_stage_view[%query_row2, %lane] : f16, view<16x64xf16, %probability_transposed_layout> + view.store %probability3, %probability_stage_view[%query_row3, %lane] : f16, view<16x64xf16, %probability_transposed_layout> + kernel.barrier scope(workgroup) ordering(acq_rel) + // Match the Vulkan CM1 ownership schedule by computing two sequential + // 64-channel output tiles. This keeps one P*V accumulator live per wave + // and reuses a 2 KiB exchange tile instead of retaining both halves. + %old_scale_f16 = vector.fptrunc %old_scale : vector<4xf32> to vector<4xf16> + %old_scale0_scalar = vector.extract %old_scale_f16[0] : vector<4xf16> -> f16 + %old_scale1_scalar = vector.extract %old_scale_f16[1] : vector<4xf16> -> f16 + %old_scale2_scalar = vector.extract %old_scale_f16[2] : vector<4xf16> -> f16 + %old_scale3_scalar = vector.extract %old_scale_f16[3] : vector<4xf16> -> f16 + %old_scale0 = vector.splat %old_scale0_scalar : vector<4xf16> + %old_scale1 = vector.splat %old_scale1_scalar : vector<4xf16> + %old_scale2 = vector.splat %old_scale2_scalar : vector<4xf16> + %old_scale3 = vector.splat %old_scale3_scalar : vector<4xf16> + %scaled_current_output0 = vector.mulf %current_output0, %old_scale0 : vector<4xf16> + %scaled_current_output1 = vector.mulf %current_output1, %old_scale1 : vector<4xf16> + %scaled_current_output2 = vector.mulf %current_output2, %old_scale2 : vector<4xf16> + %scaled_current_output3 = vector.mulf %current_output3, %old_scale3 : vector<4xf16> + %next_output0, %next_output1, %next_output2, %next_output3 = scf.for %output_tile = [%c0 to %c2 step %c1](%tile_output0 = %scaled_current_output0 : vector<4xf16>, %tile_output1 = %scaled_current_output1 : vector<4xf16>, %tile_output2 = %scaled_current_output2 : vector<4xf16>, %tile_output3 = %scaled_current_output3 : vector<4xf16>) -> (vector<4xf16>, vector<4xf16>, vector<4xf16>, vector<4xf16>) unroll { + %output_tile_channel = index.mul %output_tile, %c64 : index + %value_channel0 = index.add %key_value_head_base, %output_tile_channel : index + %value_channel = index.add %value_channel0, %subgroup_product_channel : index + %product_init = vector.fragment %c0_f16x8 shape [%m, %n] : vector<8xf16> + %product_fragment = scf.for %key_tile = [%c0 to %c64 step %c16](%product_accumulator = %product_init : vector<8xf16>) -> (vector<8xf16>) unroll schedule(recurrence) { + %value_token0 = index.add %key_origin, %key_tile : index + %value_token = index.assume %value_token0 [lt(%value_token0, %last_full_key_tile_start)] : index + %probability_fragment = vector.fragment.load %probability_stage_view[%c0, %key_tile] shape [%m, %k] : view<16x64xf16, %probability_transposed_layout> -> vector<16xf16> + %value_fragment = vector.fragment.load %value_view[%value_token, %value_channel] shape [%k, %n] : view<[%bounded_key_value_token_count]x[%key_value_width]xf16> -> vector<16xf16> + %next_product_accumulator = vector.mma %probability_fragment, %value_fragment, %product_accumulator : vector<16xf16>, vector<16xf16>, vector<8xf16> + scf.yield %next_product_accumulator : vector<8xf16> + } + vector.fragment.store %product_fragment, %product_stage_view[%c0, %subgroup_product_channel] shape [%m, %n] : vector<8xf16>, view<16x64xf16> + kernel.barrier scope(workgroup) ordering(acq_rel) + %owns_output_tile = index.cmp eq, %lane_output_tile, %output_tile : index + %updated_output0, %updated_output1, %updated_output2, %updated_output3 = scf.if %owns_output_tile -> (vector<4xf16>, vector<4xf16>, vector<4xf16>, vector<4xf16>) { + %block_output0 = vector.load %product_stage_view[%query_row0, %lane_product_channel] : view<16x64xf16> -> vector<4xf16> + %block_output1 = vector.load %product_stage_view[%query_row1, %lane_product_channel] : view<16x64xf16> -> vector<4xf16> + %block_output2 = vector.load %product_stage_view[%query_row2, %lane_product_channel] : view<16x64xf16> -> vector<4xf16> + %block_output3 = vector.load %product_stage_view[%query_row3, %lane_product_channel] : view<16x64xf16> -> vector<4xf16> + %updated_tile_output0 = vector.addf %tile_output0, %block_output0 : vector<4xf16> + %updated_tile_output1 = vector.addf %tile_output1, %block_output1 : vector<4xf16> + %updated_tile_output2 = vector.addf %tile_output2, %block_output2 : vector<4xf16> + %updated_tile_output3 = vector.addf %tile_output3, %block_output3 : vector<4xf16> + scf.yield %updated_tile_output0, %updated_tile_output1, %updated_tile_output2, %updated_tile_output3 : vector<4xf16>, vector<4xf16>, vector<4xf16>, vector<4xf16> + } else { + scf.yield %tile_output0, %tile_output1, %tile_output2, %tile_output3 : vector<4xf16>, vector<4xf16>, vector<4xf16>, vector<4xf16> + } + // Complete every read before the next output tile overwrites LDS. + kernel.barrier scope(workgroup) ordering(acq_rel) + scf.yield %updated_output0, %updated_output1, %updated_output2, %updated_output3 : vector<4xf16>, vector<4xf16>, vector<4xf16>, vector<4xf16> + } + scf.yield %next_max, %next_sum, %next_output0, %next_output1, %next_output2, %next_output3 : vector<4xf32>, vector<4xf32>, vector<4xf16>, vector<4xf16>, vector<4xf16>, vector<4xf16> + } else { + scf.yield %current_max, %current_sum, %current_output0, %current_output1, %current_output2, %current_output3 : vector<4xf32>, vector<4xf32>, vector<4xf16>, vector<4xf16>, vector<4xf16>, vector<4xf16> + } + scf.yield %next_block_max, %next_block_sum, %next_block_output0, %next_block_output1, %next_block_output2, %next_block_output3 : vector<4xf32>, vector<4xf32>, vector<4xf16>, vector<4xf16>, vector<4xf16>, vector<4xf16> + } + // A masked 32-row tile handles the final 1-63 KV rows with the same WMMA + // fragments as the aligned path. K and V are staged separately into product + // scratch, so every padded element is initialized and no physical padding is + // required of the caller. Each tile rounds probabilities to F16 before P*V. + %tail_score_wave = index.cmp ult, %subgroup, %c2 : index + %final_max, %final_sum, %final_output0, %final_output1, %final_output2, %final_output3 = scf.for %tail_key_origin = [%full_key_value_token_count to %bounded_key_value_token_count step %c32](%current_max = %full_max : vector<4xf32>, %current_sum = %full_sum : vector<4xf32>, %current_output0 = %full_output0 : vector<4xf16>, %current_output1 = %full_output1 : vector<4xf16>, %current_output2 = %full_output2 : vector<4xf16>, %current_output3 = %full_output3 : vector<4xf16>) -> (vector<4xf32>, vector<4xf32>, vector<4xf16>, vector<4xf16>, vector<4xf16>, vector<4xf16>) { + %tail_remaining = index.sub %bounded_key_value_token_count, %tail_key_origin : index + %tail_key_count = index.min %tail_remaining, %c32 : index + // Cooperatively stage one K tile, explicitly zeroing the padded rows. + scf.for %load_iteration = [%c0 to %c16 step %c1] unroll { + %linear = index.madd %load_iteration, %c256, %workitem : index + %tail_key_row = index.div %linear, %c128 : index + %tail_key_channel = index.rem %linear, %c128 : index + %tail_key_valid = index.cmp ult, %tail_key_row, %tail_key_count : index + %tail_key_value = scf.if %tail_key_valid -> (f16) { + %tail_key_token0 = index.add %tail_key_origin, %tail_key_row : index + %tail_key_token = index.assume %tail_key_token0 [lt(%tail_key_token0, %bounded_key_value_token_count)] : index + %global_key_channel = index.add %key_value_head_base, %tail_key_channel : index + %loaded = view.load %key_view[%tail_key_token, %global_key_channel] : view<[%bounded_key_value_token_count]x[%key_value_width]xf16> -> f16 + scf.yield %loaded : f16 + } else { + scf.yield %c0_f16 : f16 + } + view.store %tail_key_value, %tail_key_value_stage_view[%tail_key_row, %tail_key_channel] : f16, view<32x128xf16> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + // Waves zero and one compute the two 16-column QK fragments in this tile. + scf.if %tail_score_wave { + %tail_score_subgroup = index.assume %subgroup [range(%subgroup, 0, 1)] : index + %tail_score_column = index.mul %tail_score_subgroup, %c16 : index + %tail_score_init_values = vector.constant 0.0 : vector<4xf32> + %tail_score_init = vector.fragment %tail_score_init_values shape [%m, %n] : vector<4xf32> + %tail_score_fragment = scf.for %head_tile = [%c0 to %c128 step %c16](%score_accumulator = %tail_score_init : vector<4xf32>) -> (vector<4xf32>) unroll { + %key_fragment = vector.fragment.load %tail_key_value_stage_view[%tail_score_column, %head_tile] shape [%m, %k] : view<32x128xf16> -> vector<16xf16> + %query_fragment = vector.fragment.load %query_transposed_view[%head_tile, %c0] shape [%k, %n] : view<128x16xf16, %query_transposed_layout> -> vector<16xf16> + %next_score_accumulator = vector.mma %key_fragment, %query_fragment, %score_accumulator : vector<16xf16>, vector<16xf16>, vector<4xf32> + scf.yield %next_score_accumulator : vector<4xf32> + } + vector.fragment.store %tail_score_fragment, %score_stage_view[%tail_score_column, %c0] shape [%m, %n] : vector<4xf32>, view<64x24xf32> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + // LDS changes ownership from the QK wave to four query-row waves. Lanes + // beyond the logical tail never read the score or mask buffers. + %tail_lane_valid = index.cmp ult, %lane, %tail_key_count : index + %tail_valid0 = scalar.andi %tail_lane_valid, %query_valid0 : i1 + %tail_valid1 = scalar.andi %tail_lane_valid, %query_valid1 : i1 + %tail_valid2 = scalar.andi %tail_lane_valid, %query_valid2 : i1 + %tail_valid3 = scalar.andi %tail_lane_valid, %query_valid3 : i1 + %tail_valid = vector.from_elements %tail_valid0, %tail_valid1, %tail_valid2, %tail_valid3 : vector<4xi1> + %tail_key_token0 = index.add %tail_key_origin, %lane : index + %masked_score0 = scf.if %tail_valid0 -> (f32) { + %tail_key_token = index.assume %tail_key_token0 [lt(%tail_key_token0, %bounded_key_value_token_count)] : index + %raw_score = view.load %score_stage_view[%lane, %query_row0] : view<64x24xf32> -> f32 + %mask_f16 = view.load %mask_view[%query_token0, %tail_key_token] : view<[%query_token_count]x[%bounded_key_value_token_count]xf16> -> f16 + %mask_f32 = scalar.extf %mask_f16 : f16 to f32 + %score = scalar.addf %raw_score, %mask_f32 : f32 + scf.yield %score : f32 + } else { + scf.yield %negative_large : f32 + } + %masked_score1 = scf.if %tail_valid1 -> (f32) { + %tail_key_token = index.assume %tail_key_token0 [lt(%tail_key_token0, %bounded_key_value_token_count)] : index + %raw_score = view.load %score_stage_view[%lane, %query_row1] : view<64x24xf32> -> f32 + %mask_f16 = view.load %mask_view[%query_token1, %tail_key_token] : view<[%query_token_count]x[%bounded_key_value_token_count]xf16> -> f16 + %mask_f32 = scalar.extf %mask_f16 : f16 to f32 + %score = scalar.addf %raw_score, %mask_f32 : f32 + scf.yield %score : f32 + } else { + scf.yield %negative_large : f32 + } + %masked_score2 = scf.if %tail_valid2 -> (f32) { + %tail_key_token = index.assume %tail_key_token0 [lt(%tail_key_token0, %bounded_key_value_token_count)] : index + %raw_score = view.load %score_stage_view[%lane, %query_row2] : view<64x24xf32> -> f32 + %mask_f16 = view.load %mask_view[%query_token2, %tail_key_token] : view<[%query_token_count]x[%bounded_key_value_token_count]xf16> -> f16 + %mask_f32 = scalar.extf %mask_f16 : f16 to f32 + %score = scalar.addf %raw_score, %mask_f32 : f32 + scf.yield %score : f32 + } else { + scf.yield %negative_large : f32 + } + %masked_score3 = scf.if %tail_valid3 -> (f32) { + %tail_key_token = index.assume %tail_key_token0 [lt(%tail_key_token0, %bounded_key_value_token_count)] : index + %raw_score = view.load %score_stage_view[%lane, %query_row3] : view<64x24xf32> -> f32 + %mask_f16 = view.load %mask_view[%query_token3, %tail_key_token] : view<[%query_token_count]x[%bounded_key_value_token_count]xf16> -> f16 + %mask_f32 = scalar.extf %mask_f16 : f16 to f32 + %score = scalar.addf %raw_score, %mask_f32 : f32 + scf.yield %score : f32 + } else { + scf.yield %negative_large : f32 + } + %masked_scores = vector.from_elements %masked_score0, %masked_score1, %masked_score2, %masked_score3 : vector<4xf32> + %block_max = kernel.subgroup.reduce %masked_scores : vector<4xf32> + %next_max = vector.maxnumf %current_max, %block_max : vector<4xf32> + %score_delta = vector.subf %masked_scores, %next_max : vector<4xf32> + %raw_probability = vector.expf %score_delta : vector<4xf32> + %probability = vector.select %tail_valid, %raw_probability, %c0_f32x4 : vector<4xf32> + %block_sum = kernel.subgroup.reduce %probability : vector<4xf32> + %old_delta = vector.subf %current_max, %next_max : vector<4xf32> + %old_scale = vector.expf %old_delta : vector<4xf32> + %scaled_current_sum = vector.mulf %current_sum, %old_scale : vector<4xf32> + %next_sum = vector.addf %scaled_current_sum, %block_sum : vector<4xf32> + %probability_f16 = vector.fptrunc %probability : vector<4xf32> to vector<4xf16> + %probability0 = vector.extract %probability_f16[0] : vector<4xf16> -> f16 + %probability1 = vector.extract %probability_f16[1] : vector<4xf16> -> f16 + %probability2 = vector.extract %probability_f16[2] : vector<4xf16> -> f16 + %probability3 = vector.extract %probability_f16[3] : vector<4xf16> -> f16 + view.store %probability0, %probability_stage_view[%query_row0, %lane] : f16, view<16x64xf16, %probability_transposed_layout> + view.store %probability1, %probability_stage_view[%query_row1, %lane] : f16, view<16x64xf16, %probability_transposed_layout> + view.store %probability2, %probability_stage_view[%query_row2, %lane] : f16, view<16x64xf16, %probability_transposed_layout> + view.store %probability3, %probability_stage_view[%query_row3, %lane] : f16, view<16x64xf16, %probability_transposed_layout> + kernel.barrier scope(workgroup) ordering(acq_rel) + // Reuse product scratch for V after every probability is resident in its + // disjoint LDS tile. + scf.for %load_iteration = [%c0 to %c16 step %c1] unroll { + %linear = index.madd %load_iteration, %c256, %workitem : index + %tail_value_row = index.div %linear, %c128 : index + %tail_value_channel = index.rem %linear, %c128 : index + %tail_value_valid = index.cmp ult, %tail_value_row, %tail_key_count : index + %tail_value = scf.if %tail_value_valid -> (f16) { + %tail_value_token0 = index.add %tail_key_origin, %tail_value_row : index + %tail_value_token = index.assume %tail_value_token0 [lt(%tail_value_token0, %bounded_key_value_token_count)] : index + %global_value_channel = index.add %key_value_head_base, %tail_value_channel : index + %loaded = view.load %value_view[%tail_value_token, %global_value_channel] : view<[%bounded_key_value_token_count]x[%key_value_width]xf16> -> f16 + scf.yield %loaded : f16 + } else { + scf.yield %c0_f16 : f16 + } + view.store %tail_value, %tail_key_value_stage_view[%tail_value_row, %tail_value_channel] : f16, view<32x128xf16> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + // Keep tail V staging separate from the 2 KiB product exchange, then use + // the same two 64-channel phases as the aligned path. + %old_scale_f16 = vector.fptrunc %old_scale : vector<4xf32> to vector<4xf16> + %old_scale0_scalar = vector.extract %old_scale_f16[0] : vector<4xf16> -> f16 + %old_scale1_scalar = vector.extract %old_scale_f16[1] : vector<4xf16> -> f16 + %old_scale2_scalar = vector.extract %old_scale_f16[2] : vector<4xf16> -> f16 + %old_scale3_scalar = vector.extract %old_scale_f16[3] : vector<4xf16> -> f16 + %old_scale0 = vector.splat %old_scale0_scalar : vector<4xf16> + %old_scale1 = vector.splat %old_scale1_scalar : vector<4xf16> + %old_scale2 = vector.splat %old_scale2_scalar : vector<4xf16> + %old_scale3 = vector.splat %old_scale3_scalar : vector<4xf16> + %scaled_current_output0 = vector.mulf %current_output0, %old_scale0 : vector<4xf16> + %scaled_current_output1 = vector.mulf %current_output1, %old_scale1 : vector<4xf16> + %scaled_current_output2 = vector.mulf %current_output2, %old_scale2 : vector<4xf16> + %scaled_current_output3 = vector.mulf %current_output3, %old_scale3 : vector<4xf16> + %next_output0, %next_output1, %next_output2, %next_output3 = scf.for %output_tile = [%c0 to %c2 step %c1](%tile_output0 = %scaled_current_output0 : vector<4xf16>, %tile_output1 = %scaled_current_output1 : vector<4xf16>, %tile_output2 = %scaled_current_output2 : vector<4xf16>, %tile_output3 = %scaled_current_output3 : vector<4xf16>) -> (vector<4xf16>, vector<4xf16>, vector<4xf16>, vector<4xf16>) unroll { + %output_tile_channel = index.mul %output_tile, %c64 : index + %value_channel = index.add %output_tile_channel, %subgroup_product_channel : index + %tail_product_init = vector.fragment %c0_f16x8 shape [%m, %n] : vector<8xf16> + %tail_product_fragment = scf.for %key_tile = [%c0 to %c32 step %c16](%product_accumulator = %tail_product_init : vector<8xf16>) -> (vector<8xf16>) unroll { + %probability_fragment = vector.fragment.load %probability_stage_view[%c0, %key_tile] shape [%m, %k] : view<16x64xf16, %probability_transposed_layout> -> vector<16xf16> + %value_fragment = vector.fragment.load %tail_key_value_stage_view[%key_tile, %value_channel] shape [%k, %n] : view<32x128xf16> -> vector<16xf16> + %next_product_accumulator = vector.mma %probability_fragment, %value_fragment, %product_accumulator : vector<16xf16>, vector<16xf16>, vector<8xf16> + scf.yield %next_product_accumulator : vector<8xf16> + } + vector.fragment.store %tail_product_fragment, %product_stage_view[%c0, %subgroup_product_channel] shape [%m, %n] : vector<8xf16>, view<16x64xf16> + kernel.barrier scope(workgroup) ordering(acq_rel) + %owns_output_tile = index.cmp eq, %lane_output_tile, %output_tile : index + %updated_output0, %updated_output1, %updated_output2, %updated_output3 = scf.if %owns_output_tile -> (vector<4xf16>, vector<4xf16>, vector<4xf16>, vector<4xf16>) { + %block_output0 = vector.load %product_stage_view[%query_row0, %lane_product_channel] : view<16x64xf16> -> vector<4xf16> + %block_output1 = vector.load %product_stage_view[%query_row1, %lane_product_channel] : view<16x64xf16> -> vector<4xf16> + %block_output2 = vector.load %product_stage_view[%query_row2, %lane_product_channel] : view<16x64xf16> -> vector<4xf16> + %block_output3 = vector.load %product_stage_view[%query_row3, %lane_product_channel] : view<16x64xf16> -> vector<4xf16> + %updated_tile_output0 = vector.addf %tile_output0, %block_output0 : vector<4xf16> + %updated_tile_output1 = vector.addf %tile_output1, %block_output1 : vector<4xf16> + %updated_tile_output2 = vector.addf %tile_output2, %block_output2 : vector<4xf16> + %updated_tile_output3 = vector.addf %tile_output3, %block_output3 : vector<4xf16> + scf.yield %updated_tile_output0, %updated_tile_output1, %updated_tile_output2, %updated_tile_output3 : vector<4xf16>, vector<4xf16>, vector<4xf16>, vector<4xf16> + } else { + scf.yield %tile_output0, %tile_output1, %tile_output2, %tile_output3 : vector<4xf16>, vector<4xf16>, vector<4xf16>, vector<4xf16> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + scf.yield %updated_output0, %updated_output1, %updated_output2, %updated_output3 : vector<4xf16>, vector<4xf16>, vector<4xf16>, vector<4xf16> + } + scf.yield %next_max, %next_sum, %next_output0, %next_output1, %next_output2, %next_output3 : vector<4xf32>, vector<4xf32>, vector<4xf16>, vector<4xf16>, vector<4xf16>, vector<4xf16> + } + // Normalize and publish the lane-owned packets. Query-tail rows never + // participate in the output store. + scf.if %lane_has_output { + %sum0_scalar = vector.extract %final_sum[0] : vector<4xf32> -> f32 + %sum1_scalar = vector.extract %final_sum[1] : vector<4xf32> -> f32 + %sum2_scalar = vector.extract %final_sum[2] : vector<4xf32> -> f32 + %sum3_scalar = vector.extract %final_sum[3] : vector<4xf32> -> f32 + %inverse_sum0_f32 = scalar.divf %c1_f32, %sum0_scalar : f32 + %inverse_sum1_f32 = scalar.divf %c1_f32, %sum1_scalar : f32 + %inverse_sum2_f32 = scalar.divf %c1_f32, %sum2_scalar : f32 + %inverse_sum3_f32 = scalar.divf %c1_f32, %sum3_scalar : f32 + %inverse_sum0_f16 = scalar.fptrunc %inverse_sum0_f32 : f32 to f16 + %inverse_sum1_f16 = scalar.fptrunc %inverse_sum1_f32 : f32 to f16 + %inverse_sum2_f16 = scalar.fptrunc %inverse_sum2_f32 : f32 to f16 + %inverse_sum3_f16 = scalar.fptrunc %inverse_sum3_f32 : f32 to f16 + %inverse_sum0 = vector.splat %inverse_sum0_f16 : vector<4xf16> + %inverse_sum1 = vector.splat %inverse_sum1_f16 : vector<4xf16> + %inverse_sum2 = vector.splat %inverse_sum2_f16 : vector<4xf16> + %inverse_sum3 = vector.splat %inverse_sum3_f16 : vector<4xf16> + %normalized0_f16 = vector.mulf %final_output0, %inverse_sum0 : vector<4xf16> + %normalized1_f16 = vector.mulf %final_output1, %inverse_sum1 : vector<4xf16> + %normalized2_f16 = vector.mulf %final_output2, %inverse_sum2 : vector<4xf16> + %normalized3_f16 = vector.mulf %final_output3, %inverse_sum3 : vector<4xf16> + %normalized0 = vector.extf %normalized0_f16 : vector<4xf16> to vector<4xf32> + %normalized1 = vector.extf %normalized1_f16 : vector<4xf16> to vector<4xf32> + %normalized2 = vector.extf %normalized2_f16 : vector<4xf16> to vector<4xf32> + %normalized3 = vector.extf %normalized3_f16 : vector<4xf16> to vector<4xf32> + scf.if %query_valid0 { + vector.store %normalized0, %output_view[%query_token0, %query_head, %lane_output_channel] : vector<4xf32>, view<[%query_token_count]x[%query_head_count]x128xf32> + } + scf.if %query_valid1 { + vector.store %normalized1, %output_view[%query_token1, %query_head, %lane_output_channel] : vector<4xf32>, view<[%query_token_count]x[%query_head_count]x128xf32> + } + scf.if %query_valid2 { + vector.store %normalized2, %output_view[%query_token2, %query_head, %lane_output_channel] : vector<4xf32>, view<[%query_token_count]x[%query_head_count]x128xf32> + } + scf.if %query_valid3 { + vector.store %normalized3, %output_view[%query_token3, %query_head, %lane_output_channel] : vector<4xf32>, view<[%query_token_count]x[%query_head_count]x128xf32> + } + } + kernel.return +} + +// Test-only causal-mask construction for the production 14-token witness. +// Finite F16 minima retain exact zero probabilities without requiring infinity +// support from synthetic tensor generators. Selective linking drops this +// helper from production roots. +kernel.def target(@qwen3_moe_attention_gfx11_wave64) @qwen3_moe_flash_attention_test_make_causal_mask() { + %c1 = index.constant 1 : index + %c256 = index.constant 256 : index + kernel.launch.config workgroups(%c1, %c1, %c1) workgroup_size(%c256, %c1, %c1) : index +} launch(%mask: buffer) { + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c14 = index.constant 14 : index + %c196 = index.constant 196 : index + %c0_offset = index.constant 0 : offset + %c0_f16 = scalar.constant 0.0 : f16 + %negative_f16 = scalar.constant -65504.0 : f16 + %workitem = kernel.workitem.id : index + %in_bounds = index.cmp ult, %workitem, %c196 : index + %mask_noalias = buffer.assume.noalias %mask : buffer + %mask_view = buffer.view %mask_noalias[%c0_offset] : buffer -> view<14x14xf16> + scf.if %in_bounds { + %row = index.div %workitem, %c14 : index + %column = index.rem %workitem, %c14 : index + %row_limit = index.add %row, %c1 : index + %is_visible = index.cmp ult, %column, %row_limit : index + %value = scf.select %is_visible, %c0_f16, %negative_f16 : f16 + view.store %value, %mask_view[%row, %column] : f16, view<14x14xf16> + } + kernel.return +} + +// Test-only row extraction for comparing one multirow attention dispatch with +// independent one-row dispatches over identical data. Selective linking drops +// this helper from production roots. +kernel.def target(@qwen3_moe_attention_gfx11_wave64) @qwen3_moe_flash_attention_test_extract_row(%query_token_count: index, %context_count: index, %source_row: index) { + %query_head_count = config.get @qwen3_moe.attention.query_head_count : index + %context_capacity = config.get @qwen3_moe.attention.test.context_capacity : index + %c128 = index.constant 128 : index + %c255 = index.constant 255 : index + %c256 = index.constant 256 : index + %c1 = index.constant 1 : index + %query_element_count = index.mul %query_head_count, %c128 : index + %element_count = index.add %query_element_count, %context_capacity : index + %padded_element_count = index.add %element_count, %c255 : index + %workgroup_count = index.div %padded_element_count, %c256 : index + kernel.launch.config workgroups(%workgroup_count, %c1, %c1) workgroup_size(%c256, %c1, %c1) : index +} launch(%query_token_count: index, %context_count: index, %source_row: index, %source_query: buffer, %source_mask: buffer, %target_query: buffer, %target_mask: buffer) { + %c0 = index.constant 0 : index + %c0_offset = index.constant 0 : offset + %c128 = index.constant 128 : index + %c256 = index.constant 256 : index + %query_head_count = config.get @qwen3_moe.attention.query_head_count : index + %context_capacity = config.get @qwen3_moe.attention.test.context_capacity : index + %bounded_query_token_count = index.assume %query_token_count [range(%query_token_count, 1, 2048)] : index + %bounded_context_count = index.assume %context_count [range(%context_count, 1, 32768), le(%context_count, %context_capacity)] : index + %bounded_source_row, %source_query_token_count = index.assume %source_row, %bounded_query_token_count [range(%source_row, 0, 2047), lt(%source_row, %bounded_query_token_count)] : index, index + %workgroup = kernel.workgroup.id : index + %workitem = kernel.workitem.id : index + %linear = index.madd %workgroup, %c256, %workitem : index + %query_element_count = index.mul %query_head_count, %c128 : index + %element_count = index.add %query_element_count, %bounded_context_count : index + %in_bounds = index.cmp ult, %linear, %element_count : index + %source_query_noalias, %source_mask_noalias, %target_query_noalias, %target_mask_noalias = buffer.assume.noalias %source_query, %source_mask, %target_query, %target_mask : buffer, buffer, buffer, buffer + %source_query_view = buffer.view %source_query_noalias[%c0_offset] : buffer -> view<[%source_query_token_count]x[%query_head_count]x128xf32> + %source_mask_view = buffer.view %source_mask_noalias[%c0_offset] : buffer -> view<[%source_query_token_count]x[%bounded_context_count]xf16> + %target_query_view = buffer.view %target_query_noalias[%c0_offset] : buffer -> view<1x[%query_head_count]x128xf32> + %target_mask_view = buffer.view %target_mask_noalias[%c0_offset] : buffer -> view<1x[%bounded_context_count]xf16> + scf.if %in_bounds { + %is_query_element = index.cmp ult, %linear, %query_element_count : index + scf.if %is_query_element { + %query_head = index.div %linear, %c128 : index + %query_channel = index.rem %linear, %c128 : index + %value = view.load %source_query_view[%bounded_source_row, %query_head, %query_channel] : view<[%source_query_token_count]x[%query_head_count]x128xf32> -> f32 + view.store %value, %target_query_view[%c0, %query_head, %query_channel] : f32, view<1x[%query_head_count]x128xf32> + } else { + %mask_column = index.sub %linear, %query_element_count : index + %value = view.load %source_mask_view[%bounded_source_row, %mask_column] : view<[%source_query_token_count]x[%bounded_context_count]xf16> -> f16 + view.store %value, %target_mask_view[%c0, %mask_column] : f16, view<1x[%bounded_context_count]xf16> + } + } + kernel.return +} + +// The mask selects the first KV row exactly. QK, F16 probability conversion, +// GQA addressing, and the P*V path all execute, while the expected result is +// the first V row and remains auditable as an iota. +check.case public @qwen3_moe_flash_attention_f32_f16_wmma_selected_row_case { + %query_token_count = check.literal value(1) : index + %key_value_token_count = check.literal value(64) : index + %query = check.generate.fill value(1.0) : tensor<1x1x128xf32> + %key = check.generate.fill value(1.0) : tensor<64x1x128xf16> + %value = check.generate.iota offset(0.0) step(0.125) : tensor<64x1x128xf16> + %mask = check.generate.iota offset(0.0) step(-10000.0) : tensor<1x64xf16> + %output = check.generate.fill value(-1.0) : tensor<1x1x128xf32> + %expected = check.generate.iota offset(0.0) step(0.125) : tensor<1x1x128xf32> + kernel.launch @qwen3_moe_flash_attention_f32_f16_wmma[%query_token_count, %key_value_token_count](%query_token_count, %key_value_token_count, %query, %key, %value, %mask, %output) : [index, index](index, index, tensor<1x1x128xf32>, tensor<64x1x128xf16>, tensor<64x1x128xf16>, tensor<1x64xf16>, tensor<1x1x128xf32>) + check.expect.close actual(%output) expected(%expected) atol(0.001) rtol(0.001) nan(same) : tensor<1x1x128xf32> + check.return +} + +// Four query rows select four distinct KV rows. The 33-element mask period +// maps row-major index row*32+column to zero exactly when row == column for +// rows zero through three. The expected output is therefore the first four +// rows of V, preserving an auditable iota while distinguishing every component +// in the four-row per-wave ownership path. +check.case public @qwen3_moe_flash_attention_f32_f16_wmma_multirow_selected_case { + %query_token_count = check.literal value(4) : index + %key_value_token_count = check.literal value(32) : index + %query = check.generate.fill value(1.0) : tensor<4x1x128xf32> + %key = check.generate.fill value(1.0) : tensor<32x1x128xf16> + %value = check.generate.iota offset(0.0) step(0.125) : tensor<32x1x128xf16> + %mask = check.generate.iota offset(0.0) step(-10000.0) period(33) : tensor<4x32xf16> + %output = check.generate.fill value(-1.0) : tensor<4x1x128xf32> + %expected = check.generate.iota offset(0.0) step(0.125) : tensor<4x1x128xf32> + kernel.launch @qwen3_moe_flash_attention_f32_f16_wmma[%query_token_count, %key_value_token_count](%query_token_count, %key_value_token_count, %query, %key, %value, %mask, %output) : [index, index](index, index, tensor<4x1x128xf32>, tensor<32x1x128xf16>, tensor<32x1x128xf16>, tensor<4x32xf16>, tensor<4x1x128xf32>) + check.expect.close actual(%output) expected(%expected) atol(0.001) rtol(0.001) nan(same) : tensor<4x1x128xf32> + check.return +} + +// A unit mask step gives every row the same geometric softmax distribution, +// shifted to begin at the row-matched KV entry. Values advance by 1/4096, so +// each KV row adds exactly 1/32 and the infinite-series weighted row offset is +// 1/(32*(e-1)). Terms that wrap at period 33 are below F16 significance. This +// exercises four independent online-softmax and P*V states while retaining a +// compact closed-form expected iota. +check.case public @qwen3_moe_flash_attention_f32_f16_wmma_multirow_online_case { + %query_token_count = check.literal value(4) : index + %key_value_token_count = check.literal value(32) : index + %query = check.generate.fill value(1.0) : tensor<4x1x128xf32> + %key = check.generate.fill value(1.0) : tensor<32x1x128xf16> + %value = check.generate.iota offset(0.0) step(0.000244140625) : tensor<32x1x128xf16> + %mask = check.generate.iota offset(0.0) step(-1.0) period(33) : tensor<4x32xf16> + %output = check.generate.fill value(-1.0) : tensor<4x1x128xf32> + %expected = check.generate.iota offset(0.01818677224124075) step(0.000244140625) : tensor<4x1x128xf32> + kernel.launch @qwen3_moe_flash_attention_f32_f16_wmma[%query_token_count, %key_value_token_count](%query_token_count, %key_value_token_count, %query, %key, %value, %mask, %output) : [index, index](index, index, tensor<4x1x128xf32>, tensor<32x1x128xf16>, tensor<32x1x128xf16>, tensor<4x32xf16>, tensor<4x1x128xf32>) + check.expect.close actual(%output) expected(%expected) atol(0.001) rtol(0.001) nan(same) : tensor<4x1x128xf32> + check.return +} + +// Compare the first four rows of one fourteen-row launch with four one-row +// launches over identical, nonuniform data and the production causal mask. +// Fourteen rows retain the exact dynamic dense-row stride, ownership, and +// query-validity pressure, while the one-row path is the proven containment. +// Row extraction only reshapes bindings and performs no attention arithmetic. +check.case public @qwen3_moe_flash_attention_f32_f16_wmma_multirow_differential_case { + %query_token_count = check.literal value(14) : index + %c1 = check.literal value(1) : index + %key_value_token_count = check.literal value(14) : index + %row0 = check.literal value(0) : index + %row1 = check.literal value(1) : index + %row2 = check.literal value(2) : index + %row3 = check.literal value(3) : index + %query_seed = check.param.seed base(0x514d4f4551554552) count(1) : i64 + %key_seed = check.param.seed base(0x514d4f454b455931) count(1) : i64 + %value_seed = check.param.seed base(0x514d4f4556414c31) count(1) : i64 + %query = check.generate.random.uniform seed(%query_seed) range(-1.0 to 1.0) : tensor<14x32x128xf32> + %key = check.generate.random.uniform seed(%key_seed) range(-1.0 to 1.0) : tensor<14x4x128xf16> + %value = check.generate.random.uniform seed(%value_seed) range(-1.0 to 1.0) : tensor<14x4x128xf16> + %mask = check.generate.fill value(0.0) : tensor<14x14xf16> + %actual = check.generate.fill value(-1.0) : tensor<14x32x128xf32> + %row_query = check.generate.fill value(0.0) : tensor<1x32x128xf32> + %row_mask = check.generate.fill value(0.0) : tensor<1x14xf16> + %expected0 = check.generate.fill value(0.0) : tensor<1x32x128xf32> + %expected1 = check.generate.fill value(0.0) : tensor<1x32x128xf32> + %expected2 = check.generate.fill value(0.0) : tensor<1x32x128xf32> + %expected3 = check.generate.fill value(0.0) : tensor<1x32x128xf32> + %actual0 = check.generate.fill value(-1.0) : tensor<1x32x128xf32> + %actual1 = check.generate.fill value(-1.0) : tensor<1x32x128xf32> + %actual2 = check.generate.fill value(-1.0) : tensor<1x32x128xf32> + %actual3 = check.generate.fill value(-1.0) : tensor<1x32x128xf32> + kernel.launch @qwen3_moe_flash_attention_test_make_causal_mask(%mask) : (tensor<14x14xf16>) + kernel.launch @qwen3_moe_flash_attention_f32_f16_wmma[%query_token_count, %key_value_token_count](%query_token_count, %key_value_token_count, %query, %key, %value, %mask, %actual) : [index, index](index, index, tensor<14x32x128xf32>, tensor<14x4x128xf16>, tensor<14x4x128xf16>, tensor<14x14xf16>, tensor<14x32x128xf32>) + kernel.launch @qwen3_moe_flash_attention_test_extract_row[%query_token_count, %key_value_token_count, %row0](%query_token_count, %key_value_token_count, %row0, %query, %mask, %row_query, %row_mask) : [index, index, index](index, index, index, tensor<14x32x128xf32>, tensor<14x14xf16>, tensor<1x32x128xf32>, tensor<1x14xf16>) + kernel.launch @qwen3_moe_flash_attention_f32_f16_wmma[%c1, %key_value_token_count](%c1, %key_value_token_count, %row_query, %key, %value, %row_mask, %expected0) : [index, index](index, index, tensor<1x32x128xf32>, tensor<14x4x128xf16>, tensor<14x4x128xf16>, tensor<1x14xf16>, tensor<1x32x128xf32>) + kernel.launch @qwen3_moe_flash_attention_test_extract_row[%query_token_count, %key_value_token_count, %row0](%query_token_count, %key_value_token_count, %row0, %actual, %mask, %actual0, %row_mask) : [index, index, index](index, index, index, tensor<14x32x128xf32>, tensor<14x14xf16>, tensor<1x32x128xf32>, tensor<1x14xf16>) + kernel.launch @qwen3_moe_flash_attention_test_extract_row[%query_token_count, %key_value_token_count, %row1](%query_token_count, %key_value_token_count, %row1, %query, %mask, %row_query, %row_mask) : [index, index, index](index, index, index, tensor<14x32x128xf32>, tensor<14x14xf16>, tensor<1x32x128xf32>, tensor<1x14xf16>) + kernel.launch @qwen3_moe_flash_attention_f32_f16_wmma[%c1, %key_value_token_count](%c1, %key_value_token_count, %row_query, %key, %value, %row_mask, %expected1) : [index, index](index, index, tensor<1x32x128xf32>, tensor<14x4x128xf16>, tensor<14x4x128xf16>, tensor<1x14xf16>, tensor<1x32x128xf32>) + kernel.launch @qwen3_moe_flash_attention_test_extract_row[%query_token_count, %key_value_token_count, %row1](%query_token_count, %key_value_token_count, %row1, %actual, %mask, %actual1, %row_mask) : [index, index, index](index, index, index, tensor<14x32x128xf32>, tensor<14x14xf16>, tensor<1x32x128xf32>, tensor<1x14xf16>) + kernel.launch @qwen3_moe_flash_attention_test_extract_row[%query_token_count, %key_value_token_count, %row2](%query_token_count, %key_value_token_count, %row2, %query, %mask, %row_query, %row_mask) : [index, index, index](index, index, index, tensor<14x32x128xf32>, tensor<14x14xf16>, tensor<1x32x128xf32>, tensor<1x14xf16>) + kernel.launch @qwen3_moe_flash_attention_f32_f16_wmma[%c1, %key_value_token_count](%c1, %key_value_token_count, %row_query, %key, %value, %row_mask, %expected2) : [index, index](index, index, tensor<1x32x128xf32>, tensor<14x4x128xf16>, tensor<14x4x128xf16>, tensor<1x14xf16>, tensor<1x32x128xf32>) + kernel.launch @qwen3_moe_flash_attention_test_extract_row[%query_token_count, %key_value_token_count, %row2](%query_token_count, %key_value_token_count, %row2, %actual, %mask, %actual2, %row_mask) : [index, index, index](index, index, index, tensor<14x32x128xf32>, tensor<14x14xf16>, tensor<1x32x128xf32>, tensor<1x14xf16>) + kernel.launch @qwen3_moe_flash_attention_test_extract_row[%query_token_count, %key_value_token_count, %row3](%query_token_count, %key_value_token_count, %row3, %query, %mask, %row_query, %row_mask) : [index, index, index](index, index, index, tensor<14x32x128xf32>, tensor<14x14xf16>, tensor<1x32x128xf32>, tensor<1x14xf16>) + kernel.launch @qwen3_moe_flash_attention_f32_f16_wmma[%c1, %key_value_token_count](%c1, %key_value_token_count, %row_query, %key, %value, %row_mask, %expected3) : [index, index](index, index, tensor<1x32x128xf32>, tensor<14x4x128xf16>, tensor<14x4x128xf16>, tensor<1x14xf16>, tensor<1x32x128xf32>) + kernel.launch @qwen3_moe_flash_attention_test_extract_row[%query_token_count, %key_value_token_count, %row3](%query_token_count, %key_value_token_count, %row3, %actual, %mask, %actual3, %row_mask) : [index, index, index](index, index, index, tensor<14x32x128xf32>, tensor<14x14xf16>, tensor<1x32x128xf32>, tensor<1x14xf16>) + check.expect.close actual(%actual0) expected(%expected0) atol(0.001) rtol(0.001) nan(same) : tensor<1x32x128xf32> + check.expect.close actual(%actual1) expected(%expected1) atol(0.001) rtol(0.001) nan(same) : tensor<1x32x128xf32> + check.expect.close actual(%actual2) expected(%expected2) atol(0.001) rtol(0.001) nan(same) : tensor<1x32x128xf32> + check.expect.close actual(%actual3) expected(%expected3) atol(0.001) rtol(0.001) nan(same) : tensor<1x32x128xf32> + check.return +} + +// Seventeen rows force a partial second query tile. Equal scores and constant +// values make every valid output exactly two while still exercising online +// normalization and the query-tail guards. +check.case public @qwen3_moe_flash_attention_f32_f16_wmma_query_tail_case { + %query_token_count = check.literal value(17) : index + %key_value_token_count = check.literal value(64) : index + %query = check.generate.fill value(1.0) : tensor<17x1x128xf32> + %key = check.generate.fill value(1.0) : tensor<64x1x128xf16> + %value = check.generate.fill value(2.0) : tensor<64x1x128xf16> + %mask = check.generate.fill value(0.0) : tensor<17x64xf16> + %output = check.generate.fill value(-1.0) : tensor<17x1x128xf32> + %expected = check.generate.fill value(2.0) : tensor<17x1x128xf32> + kernel.launch @qwen3_moe_flash_attention_f32_f16_wmma[%query_token_count, %key_value_token_count](%query_token_count, %key_value_token_count, %query, %key, %value, %mask, %output) : [index, index](index, index, tensor<17x1x128xf32>, tensor<64x1x128xf16>, tensor<64x1x128xf16>, tensor<17x64xf16>, tensor<17x1x128xf32>) + check.expect.close actual(%output) expected(%expected) atol(0.001) rtol(0.001) nan(same) : tensor<17x1x128xf32> + check.return +} + +// Sixty-five KV rows force one masked cleanup tile after a full WMMA block. +// The mask selects only that final row, whose iota values begin at 1024, so +// omitting the cleanup path cannot accidentally satisfy the check. +check.case public @qwen3_moe_flash_attention_f32_f16_wmma_key_value_tail_case { + %query_token_count = check.literal value(1) : index + %key_value_token_count = check.literal value(65) : index + %query = check.generate.fill value(1.0) : tensor<1x1x128xf32> + %key = check.generate.fill value(1.0) : tensor<65x1x128xf16> + %value = check.generate.iota offset(0.0) step(0.125) : tensor<65x1x128xf16> + %mask = check.generate.iota offset(-64000.0) step(1000.0) : tensor<1x65xf16> + %output = check.generate.fill value(-1.0) : tensor<1x1x128xf32> + %expected = check.generate.iota offset(1024.0) step(0.125) : tensor<1x1x128xf32> + kernel.launch @qwen3_moe_flash_attention_f32_f16_wmma[%query_token_count, %key_value_token_count](%query_token_count, %key_value_token_count, %query, %key, %value, %mask, %output) : [index, index](index, index, tensor<1x1x128xf32>, tensor<65x1x128xf16>, tensor<65x1x128xf16>, tensor<1x65xf16>, tensor<1x1x128xf32>) + check.expect.close actual(%output) expected(%expected) atol(0.001) rtol(0.001) nan(same) : tensor<1x1x128xf32> + check.return +} + +// The first 64-row block contains the selected value while every mask in the +// second block rounds to negative infinity. This preserves a finite online +// softmax state while exercising the workgroup-wide block-pruning path. +check.case public @qwen3_moe_flash_attention_f32_f16_wmma_pruned_block_case { + %query_token_count = check.literal value(1) : index + %key_value_token_count = check.literal value(128) : index + %query = check.generate.fill value(1.0) : tensor<1x1x128xf32> + %key = check.generate.fill value(1.0) : tensor<128x1x128xf16> + %value = check.generate.iota offset(0.0) step(0.125) : tensor<128x1x128xf16> + %mask = check.generate.iota offset(0.0) step(-2000.0) : tensor<1x128xf16> + %output = check.generate.fill value(-1.0) : tensor<1x1x128xf32> + %expected = check.generate.iota offset(0.0) step(0.125) : tensor<1x1x128xf32> + kernel.launch @qwen3_moe_flash_attention_f32_f16_wmma[%query_token_count, %key_value_token_count](%query_token_count, %key_value_token_count, %query, %key, %value, %mask, %output) : [index, index](index, index, tensor<1x1x128xf32>, tensor<128x1x128xf16>, tensor<128x1x128xf16>, tensor<1x128xf16>, tensor<1x1x128xf32>) + check.expect.close actual(%output) expected(%expected) atol(0.001) rtol(0.001) nan(same) : tensor<1x1x128xf32> + check.return +} + +check.case public @qwen3_moe_flash_attention_f32_f16_wmma_benchmark_case { + %query_token_count = check.param.choice values([1, 32, 64, 128, 192, 255, 256, 257, 384, 511, 512, 513, 768, 1023, 1024, 1025, 1280, 1536, 1792, 2048]) name("query_token_count") : index + %key_value_token_count = check.param.choice values([64, 128, 192, 255, 256, 257, 384, 511, 512, 513, 768, 1023, 1024, 1025, 1280, 1536, 1792, 2048, 32768]) name("key_value_token_count") : index + %query = check.generate.fill value(0.0) : tensor<[%query_token_count]x32x128xf32> + %key = check.generate.fill value(0.0) : tensor<[%key_value_token_count]x4x128xf16> + %value = check.generate.fill value(0.0) : tensor<[%key_value_token_count]x4x128xf16> + %mask = check.generate.fill value(0.0) : tensor<[%query_token_count]x[%key_value_token_count]xf16> + %output = check.generate.fill value(1.0) : tensor<[%query_token_count]x32x128xf32> + %expected = check.generate.fill value(0.0) : tensor<[%query_token_count]x32x128xf32> + kernel.launch @qwen3_moe_flash_attention_f32_f16_wmma[%query_token_count, %key_value_token_count](%query_token_count, %key_value_token_count, %query, %key, %value, %mask, %output) : [index, index](index, index, tensor<[%query_token_count]x32x128xf32>, tensor<[%key_value_token_count]x4x128xf16>, tensor<[%key_value_token_count]x4x128xf16>, tensor<[%query_token_count]x[%key_value_token_count]xf16>, tensor<[%query_token_count]x32x128xf32>) + check.expect.close actual(%output) expected(%expected) atol(0.0) rtol(0.0) nan(same) : tensor<[%query_token_count]x32x128xf32> + check.return +} + +// Four finite 64-row blocks followed by four negative-infinity blocks model +// the aggregate computed/pruned work ratio of 512-token causal prefill while +// keeping every workgroup's path identical for a stable microbenchmark. +check.case public @qwen3_moe_flash_attention_f32_f16_wmma_pruned_half_benchmark_case { + %query_token_count = check.literal value(512) : index + %key_value_token_count = check.literal value(512) : index + %query = check.generate.fill value(0.0) : tensor<512x32x128xf32> + %key = check.generate.fill value(0.0) : tensor<512x4x128xf16> + %value = check.generate.fill value(0.0) : tensor<512x4x128xf16> + %mask = check.generate.iota offset(0.0) step(-300.0) period(512) : tensor<512x512xf16> + %output = check.generate.fill value(1.0) : tensor<512x32x128xf32> + %expected = check.generate.fill value(0.0) : tensor<512x32x128xf32> + kernel.launch @qwen3_moe_flash_attention_f32_f16_wmma[%query_token_count, %key_value_token_count](%query_token_count, %key_value_token_count, %query, %key, %value, %mask, %output) : [index, index](index, index, tensor<512x32x128xf32>, tensor<512x4x128xf16>, tensor<512x4x128xf16>, tensor<512x512xf16>, tensor<512x32x128xf32>) + check.expect.close actual(%output) expected(%expected) atol(0.0) rtol(0.0) nan(same) : tensor<512x32x128xf32> + check.return +} + +check.benchmark<@qwen3_moe_flash_attention_f32_f16_wmma_selected_row_case> @qwen3_moe_flash_attention_f32_f16_wmma_selected_row + +check.benchmark<@qwen3_moe_flash_attention_f32_f16_wmma_query_tail_case> @qwen3_moe_flash_attention_f32_f16_wmma_query_tail + +check.benchmark<@qwen3_moe_flash_attention_f32_f16_wmma_key_value_tail_case> @qwen3_moe_flash_attention_f32_f16_wmma_key_value_tail + +check.benchmark<@qwen3_moe_flash_attention_f32_f16_wmma_pruned_block_case> @qwen3_moe_flash_attention_f32_f16_wmma_pruned_block + +check.benchmark<@qwen3_moe_flash_attention_f32_f16_wmma_pruned_half_benchmark_case> @qwen3_moe_flash_attention_f32_f16_wmma_prefill_512_pruned_half + +check.benchmark<@qwen3_moe_flash_attention_f32_f16_wmma_benchmark_case> @qwen3_moe_flash_attention_f32_f16_wmma_decode_256 {key_value_token_count = 256, query_token_count = 1} + +check.benchmark<@qwen3_moe_flash_attention_f32_f16_wmma_benchmark_case> @qwen3_moe_flash_attention_f32_f16_wmma_decode_2048 {key_value_token_count = 2048, query_token_count = 1} + +check.benchmark<@qwen3_moe_flash_attention_f32_f16_wmma_benchmark_case> @qwen3_moe_flash_attention_f32_f16_wmma_decode_32768 {key_value_token_count = 32768, query_token_count = 1} + +check.benchmark<@qwen3_moe_flash_attention_f32_f16_wmma_benchmark_case> @qwen3_moe_flash_attention_f32_f16_wmma_prefill_32 {key_value_token_count = 256, query_token_count = 32} + +check.benchmark<@qwen3_moe_flash_attention_f32_f16_wmma_benchmark_case> @qwen3_moe_flash_attention_f32_f16_wmma_prefill_128 {key_value_token_count = 256, query_token_count = 128} + +check.benchmark<@qwen3_moe_flash_attention_f32_f16_wmma_benchmark_case> @qwen3_moe_flash_attention_f32_f16_wmma_prefill_512 {key_value_token_count = 512, query_token_count = 512} + +check.benchmark<@qwen3_moe_flash_attention_f32_f16_wmma_benchmark_case> @qwen3_moe_flash_attention_f32_f16_wmma_prefill_512_context_1024 {key_value_token_count = 1024, query_token_count = 512} + +check.benchmark<@qwen3_moe_flash_attention_f32_f16_wmma_benchmark_case> @qwen3_moe_flash_attention_f32_f16_wmma_prefill_512_context_1536 {key_value_token_count = 1536, query_token_count = 512} + +check.benchmark<@qwen3_moe_flash_attention_f32_f16_wmma_benchmark_case> @qwen3_moe_flash_attention_f32_f16_wmma_prefill_512_context_2048 {key_value_token_count = 2048, query_token_count = 512} + +check.benchmark<@qwen3_moe_flash_attention_f32_f16_wmma_benchmark_case> @qwen3_moe_flash_attention_f32_f16_wmma_prefill_1024 {key_value_token_count = 1024, query_token_count = 1024} + +check.benchmark<@qwen3_moe_flash_attention_f32_f16_wmma_benchmark_case> @qwen3_moe_flash_attention_f32_f16_wmma_prefill_2048 {key_value_token_count = 2048, query_token_count = 2048} + +// These aligned and boundary-adjacent self-attention shapes expose launch or +// tail cliffs that a powers-of-two-only benchmark would hide. They are +// measurement witnesses for one shape-specialized kernel, not routing buckets. +check.benchmark<@qwen3_moe_flash_attention_f32_f16_wmma_benchmark_case> @qwen3_moe_flash_attention_f32_f16_wmma_self_64 {key_value_token_count = 64, query_token_count = 64} + +check.benchmark<@qwen3_moe_flash_attention_f32_f16_wmma_benchmark_case> @qwen3_moe_flash_attention_f32_f16_wmma_self_128 {key_value_token_count = 128, query_token_count = 128} + +check.benchmark<@qwen3_moe_flash_attention_f32_f16_wmma_benchmark_case> @qwen3_moe_flash_attention_f32_f16_wmma_self_192 {key_value_token_count = 192, query_token_count = 192} + +check.benchmark<@qwen3_moe_flash_attention_f32_f16_wmma_benchmark_case> @qwen3_moe_flash_attention_f32_f16_wmma_self_255 {key_value_token_count = 255, query_token_count = 255} + +check.benchmark<@qwen3_moe_flash_attention_f32_f16_wmma_benchmark_case> @qwen3_moe_flash_attention_f32_f16_wmma_self_256 {key_value_token_count = 256, query_token_count = 256} + +check.benchmark<@qwen3_moe_flash_attention_f32_f16_wmma_benchmark_case> @qwen3_moe_flash_attention_f32_f16_wmma_self_257 {key_value_token_count = 257, query_token_count = 257} + +check.benchmark<@qwen3_moe_flash_attention_f32_f16_wmma_benchmark_case> @qwen3_moe_flash_attention_f32_f16_wmma_self_384 {key_value_token_count = 384, query_token_count = 384} + +check.benchmark<@qwen3_moe_flash_attention_f32_f16_wmma_benchmark_case> @qwen3_moe_flash_attention_f32_f16_wmma_self_511 {key_value_token_count = 511, query_token_count = 511} + +check.benchmark<@qwen3_moe_flash_attention_f32_f16_wmma_benchmark_case> @qwen3_moe_flash_attention_f32_f16_wmma_self_512 {key_value_token_count = 512, query_token_count = 512} + +check.benchmark<@qwen3_moe_flash_attention_f32_f16_wmma_benchmark_case> @qwen3_moe_flash_attention_f32_f16_wmma_self_513 {key_value_token_count = 513, query_token_count = 513} + +check.benchmark<@qwen3_moe_flash_attention_f32_f16_wmma_benchmark_case> @qwen3_moe_flash_attention_f32_f16_wmma_self_768 {key_value_token_count = 768, query_token_count = 768} + +check.benchmark<@qwen3_moe_flash_attention_f32_f16_wmma_benchmark_case> @qwen3_moe_flash_attention_f32_f16_wmma_self_1023 {key_value_token_count = 1023, query_token_count = 1023} + +check.benchmark<@qwen3_moe_flash_attention_f32_f16_wmma_benchmark_case> @qwen3_moe_flash_attention_f32_f16_wmma_self_1024 {key_value_token_count = 1024, query_token_count = 1024} + +check.benchmark<@qwen3_moe_flash_attention_f32_f16_wmma_benchmark_case> @qwen3_moe_flash_attention_f32_f16_wmma_self_1025 {key_value_token_count = 1025, query_token_count = 1025} + +check.benchmark<@qwen3_moe_flash_attention_f32_f16_wmma_benchmark_case> @qwen3_moe_flash_attention_f32_f16_wmma_self_1280 {key_value_token_count = 1280, query_token_count = 1280} + +check.benchmark<@qwen3_moe_flash_attention_f32_f16_wmma_benchmark_case> @qwen3_moe_flash_attention_f32_f16_wmma_self_1536 {key_value_token_count = 1536, query_token_count = 1536} + +check.benchmark<@qwen3_moe_flash_attention_f32_f16_wmma_benchmark_case> @qwen3_moe_flash_attention_f32_f16_wmma_self_1792 {key_value_token_count = 1792, query_token_count = 1792} + +check.benchmark<@qwen3_moe_flash_attention_f32_f16_wmma_benchmark_case> @qwen3_moe_flash_attention_f32_f16_wmma_self_2048 {key_value_token_count = 2048, query_token_count = 2048} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/model_config.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/model_config.loom new file mode 100644 index 000000000000..59f6b68a28cd --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/model_config.loom @@ -0,0 +1,18 @@ +// Qwen model-wide configuration shared across independently linked kernels. +// +// Layer-local storage formats and schedule choices remain with their owning +// kernels. These values describe the semantic model contract and therefore +// have one symbol definition regardless of how many kernels consume them. +config.decl @qwen3_moe.model.hidden_size : %value: index where [range(%value, 128, 32768), mul(%value, 128)] + +config.decl @qwen3_moe.model.rms_epsilon : f32 + +config.decl @qwen3_moe.attention.head_size : %value: index where [range(%value, 4, 1024), mul(%value, 4)] + +config.decl @qwen3_moe.attention.query_size : %value: index where [range(%value, 1, 262144)] + +config.decl @qwen3_moe.attention.key_value_size : %value: index where [range(%value, 1, 262144)] + +config.decl @qwen3_moe.router.expert_count : %value: index where [range(%value, 32, 512), mul(%value, 32)] + +config.decl @qwen3_moe.router.route_count : %value: index where [range(%value, 1, 32)] diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/routed_down_next_q8.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/routed_down_next_q8.loom new file mode 100644 index 000000000000..e569e268ba95 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/routed_down_next_q8.loom @@ -0,0 +1,227 @@ +// Shared decode completion protocol for routed-down projections that publish +// the normalized Q8_1 x4 row consumed by the next projection boundary. Each +// x workgroup must publish one output tile before applying this body and must +// provide both its tile width and the number of target subgroups participating +// in the row reduction. The last arrival acquires the complete residual row, +// normalizes and packs it, then resets the reusable completion word after every +// Q8 store. +// +// This outer device template deliberately applies the RMSNorm/Q8 device +// template. Keeping that composition authored here ensures the compiler +// preserves the eventual kernel ancestor through nested template selection. +config.decl @qwen3_moe.model.hidden_size : %value: index where [range(%value, 128, 32768), mul(%value, 128)] +config.decl @qwen3_moe.model.rms_epsilon : f32 + +func.def inline @qwen3_moe_rmsnorm_quantize_q8_1_x4_body(%publish_normalized: i1, %reduction_subgroup_count0: index, %token_count: index, %token0: index, %input: buffer, %weight: buffer, %normalized_output: buffer, %q8_output: buffer) { + %hidden_size0 = config.get @qwen3_moe.model.hidden_size : index + %epsilon = config.get @qwen3_moe.model.rms_epsilon : f32 + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 2048)] : index + %reduction_subgroup_count = index.assume %reduction_subgroup_count0 [range(%reduction_subgroup_count0, 1, 8)] : index + %hidden_size = index.assume %hidden_size0 [range(%hidden_size0, 128, 32768), mul(%hidden_size0, 128)] : index + %workitem = kernel.workitem.id : index + %subgroup0 = kernel.subgroup.id : index + %subgroup = index.assume %subgroup0 [range(%subgroup0, 0, 7)] : index + %lane = kernel.subgroup.lane.id : index + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c4 = index.constant 4 : index + %c8 = index.constant 8 : index + %c16 = index.constant 16 : index + %c32 = index.constant 32 : index + %c128 = index.constant 128 : index + %c256 = index.constant 256 : index + %c1024 = index.constant 1024 : index + %group_bytes = index.constant 144 : offset + %payload_byte_add = index.constant 16 : offset + %scratch_d_byte_add = index.constant 1024 : offset + %scratch_bytes = index.constant 1152 : offset + %c0_offset = index.constant 0 : offset + %c0_f32 = scalar.constant 0.0 : f32 + %c1_f32 = scalar.constant 1.0 : f32 + %c127 = scalar.constant 127.0 : f32 + %c0_f32x4 = vector.constant 0.0 : vector<4xf32> + %valid_token = index.cmp ult, %token0, %bounded_token_count : index + %safe_token0 = scf.select %valid_token, %token0, %c0 : index + %token, %launch_token_count = index.assume %safe_token0, %bounded_token_count [lt(%safe_token0, %bounded_token_count)] : index, index + %hidden_size_i32 = index.cast %hidden_size : index to i32 + %hidden_size_f32 = scalar.sitofp %hidden_size_i32 : i32 to f32 + %physical_group_count = index.div %hidden_size, %c128 : index + %row_bytes = index.scale %physical_group_count, %group_bytes : index, offset -> offset + %token_output_byte_base = index.scale %token, %row_bytes : index, offset -> offset + %input_noalias, %weight_noalias, %q8_output_noalias = buffer.assume.noalias %input, %weight, %q8_output : buffer, buffer, buffer + %input_view = buffer.view %input_noalias[%c0_offset] : buffer -> view<[%launch_token_count]x[%hidden_size]xf32> + %weight_view = buffer.view %weight_noalias[%c0_offset] : buffer -> view<[%hidden_size]xf32> + %normalized_output_view = buffer.view %normalized_output[%c0_offset] : buffer -> view<[%launch_token_count]x[%hidden_size]xf32> + %scratch = buffer.alloca align(16) %scratch_bytes : buffer + %scratch_values = buffer.view %scratch[%c0_offset] : buffer -> view<256xf32> + %scratch_d = buffer.view %scratch[%scratch_d_byte_add] : buffer -> view<32xf32> + // Reduce the complete row before any block-local quantization. + %thread_sum = scf.for %channel = [%workitem to %hidden_size step %c256](%running_sum = %c0_f32 : f32) -> (f32) { + %value = view.load %input_view[%token, %channel] : view<[%launch_token_count]x[%hidden_size]xf32> -> f32 + %square = scalar.mulf %value, %value : f32 + %next_sum = scalar.addf %running_sum, %square : f32 + scf.yield %next_sum : f32 + } + %subgroup_sum = kernel.subgroup.reduce %thread_sum : f32 + %is_subgroup_leader = index.cmp eq, %lane, %c0 : index + scf.if %is_subgroup_leader { + view.store %subgroup_sum, %scratch_values[%subgroup] : f32, view<256xf32> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + %is_reduction_subgroup = index.cmp eq, %subgroup, %c0 : index + %is_reduction_lane = index.cmp ult, %lane, %reduction_subgroup_count : index + %loads_subgroup_sum = scalar.andi %is_reduction_subgroup, %is_reduction_lane : i1 + %subgroup_partial = scf.if %loads_subgroup_sum -> (f32) { + %value = view.load %scratch_values[%lane] : view<256xf32> -> f32 + scf.yield %value : f32 + } else { + scf.yield %c0_f32 : f32 + } + %row_sum = kernel.subgroup.reduce %subgroup_partial : f32 + %writes_scale = scalar.andi %is_reduction_subgroup, %is_subgroup_leader : i1 + scf.if %writes_scale { + %mean = scalar.divf %row_sum, %hidden_size_f32 : f32 + %biased_mean = scalar.addf %mean, %epsilon : f32 + %scale = scalar.rsqrtf %biased_mean : f32 + view.store %scale, %scratch_values[%c0] : f32, view<256xf32> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + %row_scale = view.load %scratch_values[%c0] : view<256xf32> -> f32 + %row_scale_vector = vector.splat %row_scale : vector<4xf32> + %publishes_normalized = scalar.andi %publish_normalized, %valid_token : i1 + // Reuse one scratch frame for each 1024-element stripe. Hidden sizes need + // only be divisible by 128; inactive workitems in the final stripe carry + // zeros and never publish. + scf.for %stripe_base = [%c0 to %hidden_size step %c1024] { + %word_element_add = index.mul %workitem, %c4 : index + %channel = index.add %stripe_base, %word_element_add : index + %valid_word = index.cmp ult, %channel, %hidden_size : index + %mask = vector.mask.range [%channel to %hidden_size step %c1] : index -> vector<4xi1> + %input_values = vector.load.mask %input_view[%token, %channel], %mask, %c0_f32x4 : view<[%launch_token_count]x[%hidden_size]xf32>, vector<4xi1>, vector<4xf32> + %learned_weights = vector.load.mask %weight_view[%channel], %mask, %c0_f32x4 : view<[%hidden_size]xf32>, vector<4xi1>, vector<4xf32> + %normalized0 = vector.mulf %input_values, %row_scale_vector : vector<4xf32> + %normalized = vector.mulf %normalized0, %learned_weights : vector<4xf32> + scf.if %publishes_normalized { + vector.store.mask %normalized, %normalized_output_view[%token, %channel], %mask : vector<4xf32>, view<[%launch_token_count]x[%hidden_size]xf32>, vector<4xi1> + } + %absolute_values = vector.absf %normalized : vector<4xf32> + %thread_max = vector.reduce %absolute_values, %c0_f32 : vector<4xf32>, f32 + view.store %thread_max, %scratch_values[%workitem] : f32, view<256xf32> + kernel.barrier scope(workgroup) ordering(acq_rel) + %word_in_block = index.rem %workitem, %c8 : index + %block_in_stripe = index.div %workitem, %c8 : index + %is_block_leader = index.cmp eq, %word_in_block, %c0 : index + %writes_block_d = scalar.andi %valid_word, %is_block_leader : i1 + scf.if %writes_block_d { + %cohort_base = index.mul %block_in_stripe, %c8 : index + %cohort_maxima = vector.load %scratch_values[%cohort_base] : view<256xf32> -> vector<8xf32> + %amax = vector.reduce %cohort_maxima, %c0_f32 : vector<8xf32>, f32 + %d = scalar.divf %amax, %c127 : f32 + view.store %d, %scratch_d[%block_in_stripe] : f32, view<32xf32> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + %d = scf.if %valid_word -> (f32) { + %block_d = view.load %scratch_d[%block_in_stripe] : view<32xf32> -> f32 + scf.yield %block_d : f32 + } else { + scf.yield %c0_f32 : f32 + } + %d_nonzero = scalar.cmpf one, %d, %c0_f32 : f32 + %d_inverse = scf.if %d_nonzero -> (f32) { + %inverse = scalar.divf %c1_f32, %d : f32 + scf.yield %inverse : f32 + } else { + scf.yield %c0_f32 : f32 + } + %d_inverse_vector = vector.splat %d_inverse : vector<4xf32> + %scaled_values = vector.mulf %normalized, %d_inverse_vector : vector<4xf32> + %rounded_values = vector.roundf %scaled_values : vector<4xf32> + %quantized_values = vector.fptosi %rounded_values : vector<4xf32> to vector<4xi8> + %packed_word = vector.bitcast %quantized_values : vector<4xi8> to vector<1xi32> + %publishes_q8_word = scalar.andi %valid_word, %valid_token : i1 + scf.if %publishes_q8_word { + %q8_block = index.div %channel, %c32 : index + %physical_group = index.div %q8_block, %c4 : index + %block_in_group = index.rem %q8_block, %c4 : index + %group_byte_add = index.scale %physical_group, %group_bytes : index, offset -> offset + %group_byte_offset = index.add %token_output_byte_base, %group_byte_add : offset + %payload_byte_offset = index.add %group_byte_offset, %payload_byte_add : offset + %group_ds = buffer.view %q8_output_noalias[%group_byte_offset] : buffer -> view<8xf16> + %group_qs = buffer.view %q8_output_noalias[%payload_byte_offset] : buffer -> view<32xi32> + %block_word_base = index.mul %block_in_group, %c8 : index + %packed_word_index0 = index.add %block_word_base, %word_in_block : index + %packed_word_index = index.assume %packed_word_index0 [range(%packed_word_index0, 0, 31)] : index + vector.store %packed_word, %group_qs[%packed_word_index] : vector<1xi32>, view<32xi32> + } + %thread_quantized_sum = vector.reduce %rounded_values, %c0_f32 : vector<4xf32>, f32 + view.store %thread_quantized_sum, %scratch_values[%workitem] : f32, view<256xf32> + kernel.barrier scope(workgroup) ordering(acq_rel) + %publishes_block_ds = scalar.andi %writes_block_d, %valid_token : i1 + scf.if %publishes_block_ds { + %cohort_base = index.mul %block_in_stripe, %c8 : index + %cohort_sums = vector.load %scratch_values[%cohort_base] : view<256xf32> -> vector<8xf32> + %quantized_sum = vector.reduce %cohort_sums, %c0_f32 : vector<8xf32>, f32 + %s = scalar.mulf %quantized_sum, %d : f32 + %q8_block = index.div %channel, %c32 : index + %physical_group = index.div %q8_block, %c4 : index + %block_in_group = index.rem %q8_block, %c4 : index + %group_byte_add = index.scale %physical_group, %group_bytes : index, offset -> offset + %group_byte_offset = index.add %token_output_byte_base, %group_byte_add : offset + %group_ds = buffer.view %q8_output_noalias[%group_byte_offset] : buffer -> view<8xf16> + %d_f16 = scalar.fptrunc %d : f32 to f16 + %s_f16 = scalar.fptrunc %s : f32 to f16 + %ds_index = index.mul %block_in_group, %c2 : index + %s_index = index.add %ds_index, %c1 : index + view.store %d_f16, %group_ds[%ds_index] : f16, view<8xf16> + view.store %s_f16, %group_ds[%s_index] : f16, view<8xf16> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + } + func.return +} + +func.def inline @qwen3_moe_routed_down_next_q8_completion(%output_channels_per_workgroup0: index, %reduction_subgroup_count0: index, %token_count: index, %output_size: index, %output: buffer, %norm_weight: buffer, %completion_counter: buffer, %next_q8_output: buffer) { + %output_channels_per_workgroup = index.assume %output_channels_per_workgroup0 [range(%output_channels_per_workgroup0, 1, 8)] : index + %reduction_subgroup_count = index.assume %reduction_subgroup_count0 [range(%reduction_subgroup_count0, 1, 8)] : index + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 1)] : index + %bounded_output_size = index.assume %output_size [range(%output_size, 128, 32768), mul(%output_size, 128)] : index + %token0 = kernel.workgroup.id : index + %token = index.assume %token0 [lt(%token0, %bounded_token_count)] : index + %c0 = index.constant 0 : index + %workitem = kernel.workitem.id : index + %output_tile_count = index.div %bounded_output_size, %output_channels_per_workgroup : index + %is_arrival_workitem = index.cmp eq, %workitem, %c0 : index + %c0_i32 = scalar.constant 0 : i32 + %c1_i32 = scalar.constant 1 : i32 + %c0_offset = index.constant 0 : offset + %counter_scratch_bytes = index.constant 4 : offset + %output_noalias, %norm_weight_noalias, %completion_counter_noalias, %next_q8_output_noalias = buffer.assume.noalias %output, %norm_weight, %completion_counter, %next_q8_output : buffer, buffer, buffer, buffer + %completion_counter_aligned = buffer.assume.alignment %completion_counter_noalias {minimum_alignment = 16} : buffer + %completion_counter_view = buffer.view %completion_counter_aligned[%c0_offset] : buffer -> view<1xi32> + %counter_scratch = buffer.alloca align(4) %counter_scratch_bytes : buffer + %counter_scratch_view = buffer.view %counter_scratch[%c0_offset] : buffer -> view<1xi32> + // Publish every producer's residual stores before the leader advances one + // workgroup arrival. The last arrival then acquires the complete row. + kernel.barrier scope(workgroup) ordering(release) + scf.if %is_arrival_workitem { + %old_counter = view.atomic.rmw %c1_i32, %completion_counter_view[%c0] {ordering = acq_rel, scope = device} : i32, view<1xi32> -> i32 + view.store %old_counter, %counter_scratch_view[%c0] : i32, view<1xi32> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + %old_counter = view.load %counter_scratch_view[%c0] : view<1xi32> -> i32 + %output_tile_count_i32 = index.cast %output_tile_count : index to i32 + %last_output_tile_i32 = scalar.subi %output_tile_count_i32, %c1_i32 : i32 + %negative_output_tile_count_i32 = scalar.subi %c0_i32, %output_tile_count_i32 : i32 + %is_last_output_tile = scalar.cmpi eq, %old_counter, %last_output_tile_i32 : i32 + scf.if %is_last_output_tile { + kernel.barrier scope(workgroup) ordering(acquire) + %publish_normalized = scalar.constant false : i1 + func.call @qwen3_moe_rmsnorm_quantize_q8_1_x4_body(%publish_normalized, %reduction_subgroup_count, %bounded_token_count, %token, %output_noalias, %norm_weight_noalias, %next_q8_output_noalias, %next_q8_output_noalias) : (i1, index, index, index, buffer, buffer, buffer, buffer) + kernel.barrier scope(workgroup) ordering(release) + scf.if %is_arrival_workitem { + view.atomic.reduce %negative_output_tile_count_i32, %completion_counter_view[%c0] {ordering = release, scope = device} : i32, view<1xi32> + } + } + func.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/routed_down_q4k.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/routed_down_q4k.loom new file mode 100644 index 000000000000..3cdc8917b720 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/routed_down_q4k.loom @@ -0,0 +1,793 @@ +// Fuses the Qwen3 MoE Q4_K down projection with route weighting, top-8 +// reduction, and residual publication. The input contains one Q8_1 x4 row per +// [token, route], while weights remain in GGUF's raw +// [expert, output, K / 256, 144 bytes] layout. +// +// Route IDs retain an independent physical stride because llama.cpp selects a +// top-8 view from its 128-entry argsort storage. Normalized route weights are +// compact [token, route]. The output enters containing the residual and is +// updated in place, so the fused boundary never materializes an unweighted or +// route-indexed [token, route, hidden] down-projection tensor. +// +// One wave owns one output channel and contracts all selected routes in +// registers. Four eight-lane cohorts contract four independent routes while +// walking every Q4_K block. Qwen's eight-route, three-block decode shape uses +// all 32 lanes across two route batches without addressing synthetic blocks. +config.decl @qwen3_moe.model.hidden_size : %value: index where [range(%value, 128, 32768), mul(%value, 128)] +config.decl @qwen3_moe.model.rms_epsilon : f32 + +func.def inline @qwen3_moe_q4k_chunk_pair_global(%weight: buffer, %row_byte_base: offset, %q4_block: index, %q4_group_pair: index, %q4_half: index, %header_words: vector<4xi32>) -> (vector<16xi8>, f32, f32, vector<16xi8>, f32, f32) { + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c4 = index.constant 4 : index + %c8 = index.constant 8 : index + %block_bytes = index.constant 144 : offset + %code_byte_add = index.constant 16 : offset + %c4_i32 = scalar.constant 4 : i32 + %nibble_mask = vector.constant 252645135 : vector<4xi32> + %bounded_pair = index.assume %q4_group_pair [range(%q4_group_pair, 0, 3)] : index + %bounded_half = index.assume %q4_half [range(%q4_half, 0, 1)] : index + %block_byte_add = index.scale %q4_block, %block_bytes : index, offset -> offset + %block_byte_base = index.add %row_byte_base, %block_byte_add : offset + %code_byte_base = index.add %block_byte_base, %code_byte_add : offset + %code_view = buffer.view %weight[%code_byte_base] : buffer -> view<32xi32> + %header_halves = vector.bitcast %header_words : vector<4xi32> to vector<8xf16> + %d_f16 = vector.extract %header_halves[0] : vector<8xf16> -> f16 + %dmin_f16 = vector.extract %header_halves[1] : vector<8xf16> -> f16 + %scale0 = vector.extract %header_words[1] : vector<4xi32> -> i32 + %scale1 = vector.extract %header_words[2] : vector<4xi32> -> i32 + %scale2 = vector.extract %header_words[3] : vector<4xi32> -> i32 + %d = scalar.extf %d_f16 : f16 to f32 + %dmin = scalar.extf %dmin_f16 : f16 to f32 + %pair_code_base = index.mul %bounded_pair, %c8 : index + %half_code_add = index.mul %bounded_half, %c4 : index + %code_index0 = index.add %pair_code_base, %half_code_add : index + %code_index = index.assume %code_index0 [range(%code_index0, 0, 28)] : index + %packed_codes = vector.load %code_view[%code_index] : view<32xi32> -> vector<4xi32> + %low_codes = vector.andi %packed_codes, %nibble_mask : vector<4xi32> + %c4_i32v = vector.splat %c4_i32 : vector<4xi32> + %high_shifted = vector.shrui %packed_codes, %c4_i32v : vector<4xi32> + %high_codes = vector.andi %high_shifted, %nibble_mask : vector<4xi32> + %q4_low = vector.bitcast %low_codes : vector<4xi32> to vector<16xi8> + %q4_high = vector.bitcast %high_codes : vector<4xi32> to vector<16xi8> + %low_group = index.mul %bounded_pair, %c2 : index + %high_group = index.add %low_group, %c1 : index + %low_scale, %low_minimum = func.call @qwen3_moe_q4k_scale_from_header(%scale0, %scale1, %scale2, %low_group) : (i32, i32, i32, index) -> (i32, i32) + %high_scale, %high_minimum = func.call @qwen3_moe_q4k_scale_from_header(%scale0, %scale1, %scale2, %high_group) : (i32, i32, i32, index) -> (i32, i32) + %low_scale_f32 = scalar.uitofp %low_scale : i32 to f32 + %low_minimum_f32 = scalar.uitofp %low_minimum : i32 to f32 + %high_scale_f32 = scalar.uitofp %high_scale : i32 to f32 + %high_minimum_f32 = scalar.uitofp %high_minimum : i32 to f32 + %low_d_scale = scalar.mulf %d, %low_scale_f32 : f32 + %low_dmin_scale = scalar.mulf %dmin, %low_minimum_f32 : f32 + %high_d_scale = scalar.mulf %d, %high_scale_f32 : f32 + %high_dmin_scale = scalar.mulf %dmin, %high_minimum_f32 : f32 + func.return %q4_low, %low_d_scale, %low_dmin_scale, %q4_high, %high_d_scale, %high_dmin_scale : vector<16xi8>, f32, f32, vector<16xi8>, f32, f32 +} + +func.def inline @qwen3_moe_q4k_q8_1_dot(%q4_values: vector<16xi8>, %d_scale: f32, %dmin_scale: f32, %q8_values: vector<16xi8>, %q8_d: f32, %q8_s: f32) -> (f32) { + %c0_i32 = scalar.constant 0 : i32 + %c0_i32v = vector.constant 0 : vector<4xi32> + %half_f32 = scalar.constant 0.5 : f32 + %partial_dots = vector.dot4i %q4_values, %q8_values, %c0_i32v : vector<16xi8>, vector<16xi8>, vector<4xi32> + %q_sum = vector.reduce %partial_dots, %c0_i32 : vector<4xi32>, i32 + %q_sum_f32 = scalar.sitofp %q_sum : i32 to f32 + %scaled_dot0 = scalar.mulf %q8_d, %d_scale : f32 + %scaled_dot = scalar.mulf %scaled_dot0, %q_sum_f32 : f32 + %q8_half_sum = scalar.mulf %q8_s, %half_f32 : f32 + %minimum_correction = scalar.mulf %dmin_scale, %q8_half_sum : f32 + %contribution = scalar.subf %scaled_dot, %minimum_correction : f32 + func.return %contribution : f32 +} + +func.def inline @qwen3_moe_q4k_q8_1_x4_cohort_row_lane(%input_size: index, %weight: buffer, %weight_row_byte_base: offset, %q8_input: buffer, %q8_row_byte_base: offset, %cohort_lane: index) -> (f32) { + %bounded_input_size = index.assume %input_size [range(%input_size, 256, 32768), mul(%input_size, 256)] : index + %bounded_cohort_lane = index.assume %cohort_lane [range(%cohort_lane, 0, 7)] : index + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c256 = index.constant 256 : index + %c0_f32 = scalar.constant 0.0 : f32 + %q4_block_count = index.div %bounded_input_size, %c256 : index + %sum = scf.for %q4_block0 = [%c0 to %q4_block_count step %c1](%block_acc = %c0_f32 : f32) -> (f32) unroll { + %q4_block, %bounded_q4_block_count = index.assume %q4_block0, %q4_block_count [lt(%q4_block0, %q4_block_count)] : index, index + %pair = func.call @qwen3_moe_q4k_q8_1_x4_paired_block_lane(%bounded_input_size, %weight, %weight_row_byte_base, %q8_input, %q8_row_byte_base, %q4_block, %bounded_cohort_lane) : (index, buffer, offset, buffer, offset, index, index) -> (f32) + %next = scalar.addf %block_acc, %pair : f32 + scf.yield %next : f32 + } + func.return %sum : f32 +} + +func.def inline @qwen3_moe_q4k_q8_1_x4_paired_block_lane(%input_size: index, %weight: buffer, %weight_row_byte_base: offset, %q8_input: buffer, %q8_row_byte_base: offset, %q4_block: index, %block_lane: index) -> (f32) { + %bounded_input_size = index.assume %input_size [range(%input_size, 256, 32768), mul(%input_size, 256)] : index + %bounded_q4_block0 = index.assume %q4_block [range(%q4_block, 0, 127)] : index + %bounded_block_lane = index.assume %block_lane [range(%block_lane, 0, 7)] : index + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c4 = index.constant 4 : index + %c8 = index.constant 8 : index + %c128 = index.constant 128 : index + %c256 = index.constant 256 : index + %q4_block_bytes = index.constant 144 : offset + %q8_group_bytes = index.constant 144 : offset + %q8_payload_byte_add = index.constant 16 : offset + %q4_block_count = index.div %bounded_input_size, %c256 : index + %q8_group_count = index.div %bounded_input_size, %c128 : index + %bounded_q4_block, %bounded_q4_block_count = index.assume %bounded_q4_block0, %q4_block_count [lt(%bounded_q4_block0, %q4_block_count)] : index, index + %q4_group_pair0 = index.div %bounded_block_lane, %c2 : index + %q4_group_pair = index.assume %q4_group_pair0 [range(%q4_group_pair0, 0, 3)] : index + %q4_half0 = index.rem %bounded_block_lane, %c2 : index + %q4_half = index.assume %q4_half0 [range(%q4_half0, 0, 1)] : index + %q8_group_in_block0 = index.div %q4_group_pair, %c2 : index + %q8_group_in_block = index.assume %q8_group_in_block0 [range(%q8_group_in_block0, 0, 1)] : index + %pair_in_q8_group0 = index.rem %q4_group_pair, %c2 : index + %pair_in_q8_group = index.assume %pair_in_q8_group0 [range(%pair_in_q8_group0, 0, 1)] : index + %q8_low_inner_block0 = index.mul %pair_in_q8_group, %c2 : index + %q8_low_inner_block = index.assume %q8_low_inner_block0 [range(%q8_low_inner_block0, 0, 2)] : index + %q8_high_inner_block0 = index.add %q8_low_inner_block, %c1 : index + %q8_high_inner_block = index.assume %q8_high_inner_block0 [range(%q8_high_inner_block0, 1, 3)] : index + %q8_half_word_add = index.mul %q4_half, %c4 : index + %q8_low_inner_word_base = index.mul %q8_low_inner_block, %c8 : index + %q8_low_word_index0 = index.add %q8_low_inner_word_base, %q8_half_word_add : index + %q8_low_word_index = index.assume %q8_low_word_index0 [range(%q8_low_word_index0, 0, 20)] : index + %q8_high_inner_word_base = index.mul %q8_high_inner_block, %c8 : index + %q8_high_word_index0 = index.add %q8_high_inner_word_base, %q8_half_word_add : index + %q8_high_word_index = index.assume %q8_high_word_index0 [range(%q8_high_word_index0, 8, 28)] : index + %q8_low_ds_index0 = index.mul %q8_low_inner_block, %c2 : index + %q8_low_ds_index = index.assume %q8_low_ds_index0 [range(%q8_low_ds_index0, 0, 4)] : index + %q8_block_group_base = index.mul %bounded_q4_block, %c2 : index + %q8_group0 = index.add %q8_block_group_base, %q8_group_in_block : index + %q8_group, %bounded_q8_group_count = index.assume %q8_group0, %q8_group_count [lt(%q8_group0, %q8_group_count)] : index, index + %q8_group_byte_add = index.scale %q8_group, %q8_group_bytes : index, offset -> offset + %q8_group_byte_base = index.add %q8_row_byte_base, %q8_group_byte_add : offset + %q8_payload_byte_base = index.add %q8_group_byte_base, %q8_payload_byte_add : offset + %q8_ds_view = buffer.view %q8_input[%q8_group_byte_base] : buffer -> view<8xf16> + %q8_words_view = buffer.view %q8_input[%q8_payload_byte_base] : buffer -> view<32xi32> + %q8_ds = vector.load %q8_ds_view[%q8_low_ds_index] : view<8xf16> -> vector<4xf16> + %q8_low_d_f16 = vector.extract %q8_ds[0] : vector<4xf16> -> f16 + %q8_low_s_f16 = vector.extract %q8_ds[1] : vector<4xf16> -> f16 + %q8_high_d_f16 = vector.extract %q8_ds[2] : vector<4xf16> -> f16 + %q8_high_s_f16 = vector.extract %q8_ds[3] : vector<4xf16> -> f16 + %q8_low_d = scalar.extf %q8_low_d_f16 : f16 to f32 + %q8_low_s = scalar.extf %q8_low_s_f16 : f16 to f32 + %q8_high_d = scalar.extf %q8_high_d_f16 : f16 to f32 + %q8_high_s = scalar.extf %q8_high_s_f16 : f16 to f32 + %q8_low_words = vector.load %q8_words_view[%q8_low_word_index] : view<32xi32> -> vector<4xi32> + %q8_high_words = vector.load %q8_words_view[%q8_high_word_index] : view<32xi32> -> vector<4xi32> + %q8_low_values = vector.bitcast %q8_low_words : vector<4xi32> to vector<16xi8> + %q8_high_values = vector.bitcast %q8_high_words : vector<4xi32> to vector<16xi8> + %q4_block_byte_add = index.scale %bounded_q4_block, %q4_block_bytes : index, offset -> offset + %q4_block_byte_base = index.add %weight_row_byte_base, %q4_block_byte_add : offset + %q4_header_view = buffer.view %weight[%q4_block_byte_base] : buffer -> view<4xi32> + %q4_header_words = vector.load %q4_header_view[0] : view<4xi32> -> vector<4xi32> + %q4_low, %low_d_scale, %low_dmin_scale, %q4_high, %high_d_scale, %high_dmin_scale = func.call @qwen3_moe_q4k_chunk_pair_global(%weight, %weight_row_byte_base, %bounded_q4_block, %q4_group_pair, %q4_half, %q4_header_words) : (buffer, offset, index, index, index, vector<4xi32>) -> (vector<16xi8>, f32, f32, vector<16xi8>, f32, f32) + %low = func.call @qwen3_moe_q4k_q8_1_dot(%q4_low, %low_d_scale, %low_dmin_scale, %q8_low_values, %q8_low_d, %q8_low_s) : (vector<16xi8>, f32, f32, vector<16xi8>, f32, f32) -> (f32) + %high = func.call @qwen3_moe_q4k_q8_1_dot(%q4_high, %high_d_scale, %high_dmin_scale, %q8_high_values, %q8_high_d, %q8_high_s) : (vector<16xi8>, f32, f32, vector<16xi8>, f32, f32) -> (f32) + %pair = scalar.addf %low, %high : f32 + func.return %pair : f32 +} + +func.def inline @qwen3_moe_q4k_scale_from_header(%scale0: i32, %scale1: i32, %scale2: i32, %q4_group: index) -> (i32, i32) { + %c4 = index.constant 4 : index + %c8 = index.constant 8 : index + %c2_i32 = scalar.constant 2 : i32 + %c4_i32 = scalar.constant 4 : i32 + %c15_i32 = scalar.constant 15 : i32 + %c48_i32 = scalar.constant 48 : i32 + %bounded_group = index.assume %q4_group [range(%q4_group, 0, 7)] : index + %is_low_group = index.cmp ult, %bounded_group, %c4 : index + %scale_lane = index.rem %bounded_group, %c4 : index + %scale_shift_index = index.mul %scale_lane, %c8 : index + %scale_shift = index.cast %scale_shift_index : index to i32 + %high_shift = scalar.addi %scale_shift, %c2_i32 : i32 + %minimum_shift = scalar.addi %scale_shift, %c4_i32 : i32 + %selected_scale_source = scf.select %is_low_group, %scale0, %scale2 : i32 + %selected_minimum_source = scf.select %is_low_group, %scale1, %scale2 : i32 + %selected_scale_high_shift = scf.select %is_low_group, %scale_shift, %high_shift : i32 + %selected_minimum_low_shift = scf.select %is_low_group, %scale_shift, %minimum_shift : i32 + %scale_low0 = scalar.shrui %selected_scale_source, %scale_shift : i32 + %scale_low = scalar.andi %scale_low0, %c15_i32 : i32 + %scale_high0 = scalar.shrui %scale0, %selected_scale_high_shift : i32 + %scale_high = scalar.andi %scale_high0, %c48_i32 : i32 + %scale = scalar.ori %scale_low, %scale_high : i32 + %minimum_low0 = scalar.shrui %selected_minimum_source, %selected_minimum_low_shift : i32 + %minimum_low = scalar.andi %minimum_low0, %c15_i32 : i32 + %minimum_high0 = scalar.shrui %scale1, %selected_scale_high_shift : i32 + %minimum_high = scalar.andi %minimum_high0, %c48_i32 : i32 + %minimum = scalar.ori %minimum_low, %minimum_high : i32 + func.return %scale, %minimum : i32, i32 +} + +func.def inline @qwen3_moe_rmsnorm_quantize_q8_1_x4_body(%publish_normalized: i1, %reduction_subgroup_count0: index, %token_count: index, %token0: index, %input: buffer, %weight: buffer, %normalized_output: buffer, %q8_output: buffer) { + %hidden_size0 = config.get @qwen3_moe.model.hidden_size : index + %epsilon = config.get @qwen3_moe.model.rms_epsilon : f32 + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 2048)] : index + %reduction_subgroup_count = index.assume %reduction_subgroup_count0 [range(%reduction_subgroup_count0, 1, 8)] : index + %hidden_size = index.assume %hidden_size0 [range(%hidden_size0, 128, 32768), mul(%hidden_size0, 128)] : index + %workitem = kernel.workitem.id : index + %subgroup0 = kernel.subgroup.id : index + %subgroup = index.assume %subgroup0 [range(%subgroup0, 0, 7)] : index + %lane = kernel.subgroup.lane.id : index + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c4 = index.constant 4 : index + %c8 = index.constant 8 : index + %c16 = index.constant 16 : index + %c32 = index.constant 32 : index + %c128 = index.constant 128 : index + %c256 = index.constant 256 : index + %c1024 = index.constant 1024 : index + %group_bytes = index.constant 144 : offset + %payload_byte_add = index.constant 16 : offset + %scratch_d_byte_add = index.constant 1024 : offset + %scratch_bytes = index.constant 1152 : offset + %c0_offset = index.constant 0 : offset + %c0_f32 = scalar.constant 0.0 : f32 + %c1_f32 = scalar.constant 1.0 : f32 + %c127 = scalar.constant 127.0 : f32 + %c0_f32x4 = vector.constant 0.0 : vector<4xf32> + %valid_token = index.cmp ult, %token0, %bounded_token_count : index + %safe_token0 = scf.select %valid_token, %token0, %c0 : index + %token, %launch_token_count = index.assume %safe_token0, %bounded_token_count [lt(%safe_token0, %bounded_token_count)] : index, index + %hidden_size_i32 = index.cast %hidden_size : index to i32 + %hidden_size_f32 = scalar.sitofp %hidden_size_i32 : i32 to f32 + %physical_group_count = index.div %hidden_size, %c128 : index + %row_bytes = index.scale %physical_group_count, %group_bytes : index, offset -> offset + %token_output_byte_base = index.scale %token, %row_bytes : index, offset -> offset + %input_noalias, %weight_noalias, %q8_output_noalias = buffer.assume.noalias %input, %weight, %q8_output : buffer, buffer, buffer + %input_view = buffer.view %input_noalias[%c0_offset] : buffer -> view<[%launch_token_count]x[%hidden_size]xf32> + %weight_view = buffer.view %weight_noalias[%c0_offset] : buffer -> view<[%hidden_size]xf32> + %normalized_output_view = buffer.view %normalized_output[%c0_offset] : buffer -> view<[%launch_token_count]x[%hidden_size]xf32> + %scratch = buffer.alloca align(16) %scratch_bytes : buffer + %scratch_values = buffer.view %scratch[%c0_offset] : buffer -> view<256xf32> + %scratch_d = buffer.view %scratch[%scratch_d_byte_add] : buffer -> view<32xf32> + // Reduce the complete row before any block-local quantization. + %thread_sum = scf.for %channel = [%workitem to %hidden_size step %c256](%running_sum = %c0_f32 : f32) -> (f32) { + %value = view.load %input_view[%token, %channel] : view<[%launch_token_count]x[%hidden_size]xf32> -> f32 + %square = scalar.mulf %value, %value : f32 + %next_sum = scalar.addf %running_sum, %square : f32 + scf.yield %next_sum : f32 + } + %subgroup_sum = kernel.subgroup.reduce %thread_sum : f32 + %is_subgroup_leader = index.cmp eq, %lane, %c0 : index + scf.if %is_subgroup_leader { + view.store %subgroup_sum, %scratch_values[%subgroup] : f32, view<256xf32> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + %is_reduction_subgroup = index.cmp eq, %subgroup, %c0 : index + %is_reduction_lane = index.cmp ult, %lane, %reduction_subgroup_count : index + %loads_subgroup_sum = scalar.andi %is_reduction_subgroup, %is_reduction_lane : i1 + %subgroup_partial = scf.if %loads_subgroup_sum -> (f32) { + %value = view.load %scratch_values[%lane] : view<256xf32> -> f32 + scf.yield %value : f32 + } else { + scf.yield %c0_f32 : f32 + } + %row_sum = kernel.subgroup.reduce %subgroup_partial : f32 + %writes_scale = scalar.andi %is_reduction_subgroup, %is_subgroup_leader : i1 + scf.if %writes_scale { + %mean = scalar.divf %row_sum, %hidden_size_f32 : f32 + %biased_mean = scalar.addf %mean, %epsilon : f32 + %scale = scalar.rsqrtf %biased_mean : f32 + view.store %scale, %scratch_values[%c0] : f32, view<256xf32> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + %row_scale = view.load %scratch_values[%c0] : view<256xf32> -> f32 + %row_scale_vector = vector.splat %row_scale : vector<4xf32> + %publishes_normalized = scalar.andi %publish_normalized, %valid_token : i1 + // Reuse one scratch frame for each 1024-element stripe. Hidden sizes need + // only be divisible by 128; inactive workitems in the final stripe carry + // zeros and never publish. + scf.for %stripe_base = [%c0 to %hidden_size step %c1024] { + %word_element_add = index.mul %workitem, %c4 : index + %channel = index.add %stripe_base, %word_element_add : index + %valid_word = index.cmp ult, %channel, %hidden_size : index + %mask = vector.mask.range [%channel to %hidden_size step %c1] : index -> vector<4xi1> + %input_values = vector.load.mask %input_view[%token, %channel], %mask, %c0_f32x4 : view<[%launch_token_count]x[%hidden_size]xf32>, vector<4xi1>, vector<4xf32> + %learned_weights = vector.load.mask %weight_view[%channel], %mask, %c0_f32x4 : view<[%hidden_size]xf32>, vector<4xi1>, vector<4xf32> + %normalized0 = vector.mulf %input_values, %row_scale_vector : vector<4xf32> + %normalized = vector.mulf %normalized0, %learned_weights : vector<4xf32> + scf.if %publishes_normalized { + vector.store.mask %normalized, %normalized_output_view[%token, %channel], %mask : vector<4xf32>, view<[%launch_token_count]x[%hidden_size]xf32>, vector<4xi1> + } + %absolute_values = vector.absf %normalized : vector<4xf32> + %thread_max = vector.reduce %absolute_values, %c0_f32 : vector<4xf32>, f32 + view.store %thread_max, %scratch_values[%workitem] : f32, view<256xf32> + kernel.barrier scope(workgroup) ordering(acq_rel) + %word_in_block = index.rem %workitem, %c8 : index + %block_in_stripe = index.div %workitem, %c8 : index + %is_block_leader = index.cmp eq, %word_in_block, %c0 : index + %writes_block_d = scalar.andi %valid_word, %is_block_leader : i1 + scf.if %writes_block_d { + %cohort_base = index.mul %block_in_stripe, %c8 : index + %cohort_maxima = vector.load %scratch_values[%cohort_base] : view<256xf32> -> vector<8xf32> + %amax = vector.reduce %cohort_maxima, %c0_f32 : vector<8xf32>, f32 + %d = scalar.divf %amax, %c127 : f32 + view.store %d, %scratch_d[%block_in_stripe] : f32, view<32xf32> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + %d = scf.if %valid_word -> (f32) { + %block_d = view.load %scratch_d[%block_in_stripe] : view<32xf32> -> f32 + scf.yield %block_d : f32 + } else { + scf.yield %c0_f32 : f32 + } + %d_nonzero = scalar.cmpf one, %d, %c0_f32 : f32 + %d_inverse = scf.if %d_nonzero -> (f32) { + %inverse = scalar.divf %c1_f32, %d : f32 + scf.yield %inverse : f32 + } else { + scf.yield %c0_f32 : f32 + } + %d_inverse_vector = vector.splat %d_inverse : vector<4xf32> + %scaled_values = vector.mulf %normalized, %d_inverse_vector : vector<4xf32> + %rounded_values = vector.roundf %scaled_values : vector<4xf32> + %quantized_values = vector.fptosi %rounded_values : vector<4xf32> to vector<4xi8> + %packed_word = vector.bitcast %quantized_values : vector<4xi8> to vector<1xi32> + %publishes_q8_word = scalar.andi %valid_word, %valid_token : i1 + scf.if %publishes_q8_word { + %q8_block = index.div %channel, %c32 : index + %physical_group = index.div %q8_block, %c4 : index + %block_in_group = index.rem %q8_block, %c4 : index + %group_byte_add = index.scale %physical_group, %group_bytes : index, offset -> offset + %group_byte_offset = index.add %token_output_byte_base, %group_byte_add : offset + %payload_byte_offset = index.add %group_byte_offset, %payload_byte_add : offset + %group_ds = buffer.view %q8_output_noalias[%group_byte_offset] : buffer -> view<8xf16> + %group_qs = buffer.view %q8_output_noalias[%payload_byte_offset] : buffer -> view<32xi32> + %block_word_base = index.mul %block_in_group, %c8 : index + %packed_word_index0 = index.add %block_word_base, %word_in_block : index + %packed_word_index = index.assume %packed_word_index0 [range(%packed_word_index0, 0, 31)] : index + vector.store %packed_word, %group_qs[%packed_word_index] : vector<1xi32>, view<32xi32> + } + %thread_quantized_sum = vector.reduce %rounded_values, %c0_f32 : vector<4xf32>, f32 + view.store %thread_quantized_sum, %scratch_values[%workitem] : f32, view<256xf32> + kernel.barrier scope(workgroup) ordering(acq_rel) + %publishes_block_ds = scalar.andi %writes_block_d, %valid_token : i1 + scf.if %publishes_block_ds { + %cohort_base = index.mul %block_in_stripe, %c8 : index + %cohort_sums = vector.load %scratch_values[%cohort_base] : view<256xf32> -> vector<8xf32> + %quantized_sum = vector.reduce %cohort_sums, %c0_f32 : vector<8xf32>, f32 + %s = scalar.mulf %quantized_sum, %d : f32 + %q8_block = index.div %channel, %c32 : index + %physical_group = index.div %q8_block, %c4 : index + %block_in_group = index.rem %q8_block, %c4 : index + %group_byte_add = index.scale %physical_group, %group_bytes : index, offset -> offset + %group_byte_offset = index.add %token_output_byte_base, %group_byte_add : offset + %group_ds = buffer.view %q8_output_noalias[%group_byte_offset] : buffer -> view<8xf16> + %d_f16 = scalar.fptrunc %d : f32 to f16 + %s_f16 = scalar.fptrunc %s : f32 to f16 + %ds_index = index.mul %block_in_group, %c2 : index + %s_index = index.add %ds_index, %c1 : index + view.store %d_f16, %group_ds[%ds_index] : f16, view<8xf16> + view.store %s_f16, %group_ds[%s_index] : f16, view<8xf16> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + } + func.return +} + +func.def inline @qwen3_moe_routed_down_next_q8_completion(%output_channels_per_workgroup0: index, %reduction_subgroup_count0: index, %token_count: index, %output_size: index, %output: buffer, %norm_weight: buffer, %completion_counter: buffer, %next_q8_output: buffer) { + %output_channels_per_workgroup = index.assume %output_channels_per_workgroup0 [range(%output_channels_per_workgroup0, 1, 8)] : index + %reduction_subgroup_count = index.assume %reduction_subgroup_count0 [range(%reduction_subgroup_count0, 1, 8)] : index + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 1)] : index + %bounded_output_size = index.assume %output_size [range(%output_size, 128, 32768), mul(%output_size, 128)] : index + %token0 = kernel.workgroup.id : index + %token = index.assume %token0 [lt(%token0, %bounded_token_count)] : index + %c0 = index.constant 0 : index + %workitem = kernel.workitem.id : index + %output_tile_count = index.div %bounded_output_size, %output_channels_per_workgroup : index + %is_arrival_workitem = index.cmp eq, %workitem, %c0 : index + %c0_i32 = scalar.constant 0 : i32 + %c1_i32 = scalar.constant 1 : i32 + %c0_offset = index.constant 0 : offset + %counter_scratch_bytes = index.constant 4 : offset + %output_noalias, %norm_weight_noalias, %completion_counter_noalias, %next_q8_output_noalias = buffer.assume.noalias %output, %norm_weight, %completion_counter, %next_q8_output : buffer, buffer, buffer, buffer + %completion_counter_aligned = buffer.assume.alignment %completion_counter_noalias {minimum_alignment = 16} : buffer + %completion_counter_view = buffer.view %completion_counter_aligned[%c0_offset] : buffer -> view<1xi32> + %counter_scratch = buffer.alloca align(4) %counter_scratch_bytes : buffer + %counter_scratch_view = buffer.view %counter_scratch[%c0_offset] : buffer -> view<1xi32> + // Publish every producer's residual stores before the leader advances one + // workgroup arrival. The last arrival then acquires the complete row. + kernel.barrier scope(workgroup) ordering(release) + scf.if %is_arrival_workitem { + %old_counter = view.atomic.rmw %c1_i32, %completion_counter_view[%c0] {ordering = acq_rel, scope = device} : i32, view<1xi32> -> i32 + view.store %old_counter, %counter_scratch_view[%c0] : i32, view<1xi32> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + %old_counter = view.load %counter_scratch_view[%c0] : view<1xi32> -> i32 + %output_tile_count_i32 = index.cast %output_tile_count : index to i32 + %last_output_tile_i32 = scalar.subi %output_tile_count_i32, %c1_i32 : i32 + %negative_output_tile_count_i32 = scalar.subi %c0_i32, %output_tile_count_i32 : i32 + %is_last_output_tile = scalar.cmpi eq, %old_counter, %last_output_tile_i32 : i32 + scf.if %is_last_output_tile { + kernel.barrier scope(workgroup) ordering(acquire) + %publish_normalized = scalar.constant false : i1 + func.call @qwen3_moe_rmsnorm_quantize_q8_1_x4_body(%publish_normalized, %reduction_subgroup_count, %bounded_token_count, %token, %output_noalias, %norm_weight_noalias, %next_q8_output_noalias, %next_q8_output_noalias) : (i1, index, index, index, buffer, buffer, buffer, buffer) + kernel.barrier scope(workgroup) ordering(release) + scf.if %is_arrival_workitem { + view.atomic.reduce %negative_output_tile_count_i32, %completion_counter_view[%c0] {ordering = release, scope = device} : i32, view<1xi32> + } + } + func.return +} + +target.decl @qwen3_moe_attention_prepare_gfx11_wave32 + +config.decl @qwen3_moe.routed_down.input_size : %value: index where [range(%value, 256, 32768), mul(%value, 256)] + +config.decl @qwen3_moe.routed_down.route_count : %value: index where [range(%value, 1, 8)] + +config.decl @qwen3_moe.routed_down.expert_count : %value: index where [range(%value, 1, 512)] + +config.decl @qwen3_moe.routed_down.output_size : %value: index where [range(%value, 1, 4096)] + +config.decl @qwen3_moe.workload.token_capacity : %value: index where [range(%value, 1, 2048)] + +kernel.decl @ggml_quantize_q8_1_x4_f32(%token_count: index, %input_size: index) launch(%token_count: index, %input_size: index, %input: buffer, %output: buffer) + +func.decl @qwen3_moe_q4k_q8_1_x4_cohort_row_lane(%input_size: index, %weight: buffer, %weight_row_byte_base: offset, %q8_input: buffer, %q8_row_byte_base: offset, %cohort_lane: index) -> (f32) + +kernel.decl @qwen3_moe_attention_rmsnorm_quantize_q8_1_x4(%token_count: index) launch(%token_count: index, %input: buffer, %weight: buffer, %q8_output: buffer) + +// Shared Q4_K contraction and residual publication. Each lane consumes both +// nibbles of a packed code load before the subgroup reduces the row. +func.def inline @qwen3_moe_routed_down_q4k_q8_1_x4_body(%publish_output: i1, %token_count: index, %token: index, %input_size: index, %route_count: index, %route_id_stride: index, %expert_count: index, %output_size: index, %q8_input: buffer, %route_ids: buffer, %route_weights: buffer, %weight: buffer, %output: buffer) { + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 2048)] : index + %bounded_token, %body_token_count = index.assume %token, %bounded_token_count [lt(%token, %bounded_token_count)] : index, index + %bounded_input_size = index.assume %input_size [range(%input_size, 256, 32768), mul(%input_size, 256)] : index + %bounded_route_count = index.assume %route_count [range(%route_count, 1, 8)] : index + %bounded_route_id_stride = index.assume %route_id_stride [range(%route_id_stride, 1, 512)] : index + %bounded_expert_count = index.assume %expert_count [range(%expert_count, 1, 512)] : index + %bounded_output_size = index.assume %output_size [range(%output_size, 1, 262144)] : index + %channel_tile = kernel.workgroup.id : index + %subgroup0 = kernel.subgroup.id : index + %lane = kernel.subgroup.lane.id : index + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c3 = index.constant 3 : index + %c4 = index.constant 4 : index + %c8 = index.constant 8 : index + %c128 = index.constant 128 : index + %c144 = index.constant 144 : index + %c256 = index.constant 256 : index + %c1_byte = index.constant 1 : offset + %c0_i32 = scalar.constant 0 : i32 + %c0_f32 = scalar.constant 0.0 : f32 + %c0_offset = index.constant 0 : offset + %subgroup = index.assume %subgroup0 [range(%subgroup0, 0, 7)] : index + %cohort0 = index.div %lane, %c8 : index + %cohort = index.assume %cohort0 [range(%cohort0, 0, 3)] : index + %cohort_lane0 = index.rem %lane, %c8 : index + %cohort_lane = index.assume %cohort_lane0 [range(%cohort_lane0, 0, 7)] : index + %channel_base = index.mul %channel_tile, %c8 : index + %channel = index.add %channel_base, %subgroup : index + %valid_channel = index.cmp ult, %channel, %bounded_output_size : index + %lane_i32 = index.cast %lane : index to i32 + %is_lane_zero = scalar.cmpi eq, %lane_i32, %c0_i32 : i32 + %q4_block_count = index.div %bounded_input_size, %c256 : index + %weight_row_byte_count = index.mul %q4_block_count, %c144 : index + %q8_group_count = index.div %bounded_input_size, %c128 : index + %q8_row_byte_count = index.mul %q8_group_count, %c144 : index + %padded_route_count = index.add %bounded_route_count, %c3 : index + %route_batch_count = index.div %padded_route_count, %c4 : index + %q8_noalias, %route_id_noalias, %route_weight_noalias, %weight_noalias, %output_noalias = buffer.assume.noalias %q8_input, %route_ids, %route_weights, %weight, %output : buffer, buffer, buffer, buffer, buffer + %route_id_view = buffer.view %route_id_noalias[%c0_offset] : buffer -> view<[%body_token_count]x[%bounded_route_id_stride]xi32> + %route_weight_view = buffer.view %route_weight_noalias[%c0_offset] : buffer -> view<[%body_token_count]x[%bounded_route_count]xf32> + %output_view = buffer.view %output_noalias[%c0_offset] : buffer -> view<[%body_token_count]x[%bounded_output_size]xf32> + %lane_has_route = index.cmp ult, %lane, %bounded_route_count : index + %lane_within_route_stride = index.cmp ult, %lane, %bounded_route_id_stride : index + %loads_active_route = scalar.andi %publish_output, %lane_has_route : i1 + %loads_route_metadata = scalar.andi %loads_active_route, %lane_within_route_stride : i1 + %lane_expert_i32, %lane_route_weight = scf.if %loads_route_metadata -> (i32, f32) { + %route_lane0, %metadata_route_count = index.assume %lane, %bounded_route_count [lt(%lane, %bounded_route_count)] : index, index + %route_lane, %metadata_route_id_stride = index.assume %route_lane0, %bounded_route_id_stride [lt(%route_lane0, %bounded_route_id_stride)] : index, index + %expert_i32 = view.load %route_id_view[%bounded_token, %route_lane] : view<[%body_token_count]x[%bounded_route_id_stride]xi32> -> i32 + %route_weight = view.load %route_weight_view[%bounded_token, %route_lane0] : view<[%body_token_count]x[%bounded_route_count]xf32> -> f32 + scf.yield %expert_i32, %route_weight : i32, f32 + } else { + scf.yield %c0_i32, %c0_f32 : i32, f32 + } + %computes_channel = scalar.andi %publish_output, %valid_channel : i1 + %routed_lane_sum = scf.if %computes_channel -> (f32) { + %sum = scf.for %route_batch = [%c0 to %route_batch_count step %c1](%route_acc = %c0_f32 : f32) -> (f32) unroll { + %route_batch_base = index.mul %route_batch, %c4 : index + %route0 = index.add %route_batch_base, %cohort : index + %route = index.assume %route0 [range(%route0, 0, 7)] : index + %active_route = index.cmp ult, %route, %bounded_route_count : index + %weighted_lane = scf.if %active_route -> (f32) { + %bounded_route, %body_route_count = index.assume %route, %bounded_route_count [lt(%route, %bounded_route_count)] : index, index + %route_i32 = index.cast %bounded_route : index to i32 + %expert_i32 = kernel.subgroup.broadcast %lane_expert_i32 from %route_i32 : i32, i32 + %expert0 = index.cast %expert_i32 : i32 to index + %expert1 = index.assume %expert0 [range(%expert0, 0, 511)] : index + %expert, %weight_expert_count = index.assume %expert1, %bounded_expert_count [lt(%expert1, %bounded_expert_count)] : index, index + %expert_output_base = index.mul %expert, %bounded_output_size : index + %expert_channel = index.add %expert_output_base, %channel : index + %weight_row_byte_index = index.mul %expert_channel, %weight_row_byte_count : index + %weight_row_byte_base = index.scale %weight_row_byte_index, %c1_byte : index, offset -> offset + %q8_row_base0 = index.mul %bounded_token, %body_route_count : index + %q8_row = index.add %q8_row_base0, %bounded_route : index + %q8_row_byte_index = index.mul %q8_row, %q8_row_byte_count : index + %q8_row_byte_base = index.scale %q8_row_byte_index, %c1_byte : index, offset -> offset + %route_lane_sum = func.call @qwen3_moe_q4k_q8_1_x4_cohort_row_lane(%bounded_input_size, %weight_noalias, %weight_row_byte_base, %q8_noalias, %q8_row_byte_base, %cohort_lane) : (index, buffer, offset, buffer, offset, index) -> (f32) + %route_weight = kernel.subgroup.broadcast %lane_route_weight from %route_i32 : f32, i32 + %weighted = scalar.mulf %route_lane_sum, %route_weight : f32 + scf.yield %weighted : f32 + } else { + scf.yield %c0_f32 : f32 + } + %next = scalar.addf %route_acc, %weighted_lane : f32 + scf.yield %next : f32 + } + scf.yield %sum : f32 + } else { + scf.yield %c0_f32 : f32 + } + %routed_sum = kernel.subgroup.reduce %routed_lane_sum : f32 + %writes_active_channel = scalar.andi %publish_output, %valid_channel : i1 + %writes_output = scalar.andi %writes_active_channel, %is_lane_zero : i1 + scf.if %writes_output { + %bounded_channel, %body_output_size = index.assume %channel, %bounded_output_size [lt(%channel, %bounded_output_size)] : index, index + %residual = view.load %output_view[%bounded_token, %bounded_channel] : view<[%body_token_count]x[%bounded_output_size]xf32> -> f32 + %result = scalar.addf %residual, %routed_sum : f32 + view.store %result, %output_view[%bounded_token, %bounded_channel] : f32, view<[%body_token_count]x[%bounded_output_size]xf32> + } + func.return +} + +kernel.def @qwen3_moe_routed_down_q4k_q8_1_x4(%token_count: index, %input_size: index, %route_count: index, %route_id_stride: index, %expert_count: index, %output_size: index) { + %token_capacity = config.get @qwen3_moe.workload.token_capacity : index + %configured_output_size = config.get @qwen3_moe.routed_down.output_size : index + %c8 = index.constant 8 : index + %c7 = index.constant 7 : index + %c1 = index.constant 1 : index + %workgroup_size = index.constant 256 : index + %padded_output_size = index.add %configured_output_size, %c7 : index + %output_tiles = index.div %padded_output_size, %c8 : index + kernel.launch.config workgroups(%output_tiles, %token_capacity, %c1) workgroup_size(%workgroup_size, %c1, %c1) : index +} launch(%token_count: index, %input_size: index, %route_count: index, %route_id_stride: index, %expert_count: index, %output_size: index, %q8_input: buffer, %route_ids: buffer, %route_weights: buffer, %weight: buffer, %output: buffer) { + %token_capacity = config.get @qwen3_moe.workload.token_capacity : index + %configured_input_size0 = config.get @qwen3_moe.routed_down.input_size : index + %configured_route_count0 = config.get @qwen3_moe.routed_down.route_count : index + %configured_expert_count0 = config.get @qwen3_moe.routed_down.expert_count : index + %configured_output_size0 = config.get @qwen3_moe.routed_down.output_size : index + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 2048), le(%token_count, %token_capacity)] : index + %bounded_input_size, %configured_input_size = index.assume %input_size, %configured_input_size0 [range(%input_size, 256, 32768), mul(%input_size, 256), eq(%input_size, %configured_input_size0)] : index, index + %bounded_route_count, %configured_route_count = index.assume %route_count, %configured_route_count0 [range(%route_count, 1, 8), eq(%route_count, %configured_route_count0)] : index, index + %bounded_route_id_stride = index.assume %route_id_stride [range(%route_id_stride, 1, 512)] : index + %bounded_expert_count, %configured_expert_count = index.assume %expert_count, %configured_expert_count0 [range(%expert_count, 1, 512), eq(%expert_count, %configured_expert_count0)] : index, index + %bounded_output_size, %configured_output_size = index.assume %output_size, %configured_output_size0 [range(%output_size, 1, 4096), eq(%output_size, %configured_output_size0)] : index, index + %token0 = kernel.workgroup.id : index + %c0 = index.constant 0 : index + %valid_token = index.cmp ult, %token0, %bounded_token_count : index + %safe_token0 = scf.select %valid_token, %token0, %c0 : index + %safe_token, %body_token_count = index.assume %safe_token0, %bounded_token_count [lt(%safe_token0, %bounded_token_count)] : index, index + func.call @qwen3_moe_routed_down_q4k_q8_1_x4_body(%valid_token, %body_token_count, %safe_token, %configured_input_size, %configured_route_count, %bounded_route_id_stride, %configured_expert_count, %configured_output_size, %q8_input, %route_ids, %route_weights, %weight, %output) : (i1, index, index, index, index, index, index, index, buffer, buffer, buffer, buffer, buffer) + kernel.return +} + +// Decode-only route that also publishes the normalized Q8_1 x4 row consumed by +// the next projection boundary. +kernel.def target(@qwen3_moe_attention_prepare_gfx11_wave32) @qwen3_moe_routed_down_q4k_q8_1_x4_next_q8(%token_count: index, %input_size: index, %route_count: index, %route_id_stride: index, %expert_count: index, %output_size: index) { + %configured_output_size = config.get @qwen3_moe.routed_down.output_size : index + %c8 = index.constant 8 : index + %c7 = index.constant 7 : index + %c1 = index.constant 1 : index + %workgroup_size = index.constant 256 : index + %padded_output_size = index.add %configured_output_size, %c7 : index + %output_tiles = index.div %padded_output_size, %c8 : index + kernel.launch.config workgroups(%output_tiles, %c1, %c1) workgroup_size(%workgroup_size, %c1, %c1) : index +} launch(%token_count: index, %input_size: index, %route_count: index, %route_id_stride: index, %expert_count: index, %output_size: index, %q8_input: buffer, %route_ids: buffer, %route_weights: buffer, %weight: buffer, %output: buffer, %norm_weight: buffer, %completion_counter: buffer, %next_q8_output: buffer) { + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 1)] : index + %configured_input_size0 = config.get @qwen3_moe.routed_down.input_size : index + %configured_route_count0 = config.get @qwen3_moe.routed_down.route_count : index + %configured_expert_count0 = config.get @qwen3_moe.routed_down.expert_count : index + %configured_output_size0 = config.get @qwen3_moe.routed_down.output_size : index + %bounded_input_size, %configured_input_size = index.assume %input_size, %configured_input_size0 [range(%input_size, 256, 32768), mul(%input_size, 256), eq(%input_size, %configured_input_size0)] : index, index + %bounded_route_count, %configured_route_count = index.assume %route_count, %configured_route_count0 [range(%route_count, 1, 8), eq(%route_count, %configured_route_count0)] : index, index + %bounded_route_id_stride = index.assume %route_id_stride [range(%route_id_stride, 1, 512)] : index + %bounded_expert_count, %configured_expert_count = index.assume %expert_count, %configured_expert_count0 [range(%expert_count, 1, 512), eq(%expert_count, %configured_expert_count0)] : index, index + %bounded_output_size, %configured_output_size = index.assume %output_size, %configured_output_size0 [range(%output_size, 128, 4096), mul(%output_size, 128), eq(%output_size, %configured_output_size0)] : index, index + %c8 = index.constant 8 : index + %publishes_output = scalar.constant true : i1 + %body_token = index.constant 0 : index + func.call @qwen3_moe_routed_down_q4k_q8_1_x4_body(%publishes_output, %bounded_token_count, %body_token, %configured_input_size, %configured_route_count, %bounded_route_id_stride, %configured_expert_count, %configured_output_size, %q8_input, %route_ids, %route_weights, %weight, %output) : (i1, index, index, index, index, index, index, index, buffer, buffer, buffer, buffer, buffer) + func.call @qwen3_moe_routed_down_next_q8_completion(%c8, %c8, %bounded_token_count, %configured_output_size, %output, %norm_weight, %completion_counter, %next_q8_output) : (index, index, index, index, buffer, buffer, buffer, buffer) + kernel.return +} + +// The production hidden width requires all 256 output tiles to arrive before +// normalization. Compare both residual and packed Q8 publication with the +// ordinary two-dispatch composition, then reuse the same completion word. +check.case public @qwen3_moe_routed_down_q4k_q8_1_x4_next_q8_differential_case { + %token_count = check.literal value(1) : index + %input_size = check.literal value(768) : index + %route_count = check.literal value(2) : index + %route_id_stride = check.literal value(4) : index + %expert_count = check.literal value(2) : index + %output_size = check.literal value(2048) : index + %routed_row_count = check.literal value(2) : index + %routed_input = check.generate.fill value(0.00390625) : tensor<2x768xf32> + %q8_input = check.generate.fill value(0) : tensor<2x864xi8> + %route_ids = check.generate.iota offset(0) step(1) period(2) : tensor<1x4xi32> + %route_weights = check.generate.fill value(0.5) : tensor<1x2xf32> + %weight = check.generate.iota offset(-72) step(1) period(144) : tensor<2x2048x3x144xi8> + %norm_weight = check.generate.iota offset(-1.0) step(0.0009765625) : tensor<2048xf32> + %expected_output = check.generate.iota offset(-0.5) step(0.00048828125) period(2048) : tensor<1x2048xf32> + %actual_output0 = check.generate.iota offset(-0.5) step(0.00048828125) period(2048) : tensor<1x2048xf32> + %actual_output1 = check.generate.iota offset(-0.5) step(0.00048828125) period(2048) : tensor<1x2048xf32> + %expected_q8 = check.generate.fill value(0) : tensor<2304xi8> + %actual_q8_0 = check.generate.fill value(1) : tensor<2304xi8> + %actual_q8_1 = check.generate.fill value(1) : tensor<2304xi8> + %completion_counter = check.generate.fill value(0) : tensor<1xi32> + %expected_counter = check.generate.fill value(0) : tensor<1xi32> + kernel.launch @ggml_quantize_q8_1_x4_f32[%routed_row_count, %input_size](%routed_row_count, %input_size, %routed_input, %q8_input) : [index, index](index, index, tensor<2x768xf32>, tensor<2x864xi8>) + kernel.launch @qwen3_moe_routed_down_q4k_q8_1_x4[%token_count, %input_size, %route_count, %route_id_stride, %expert_count, %output_size](%token_count, %input_size, %route_count, %route_id_stride, %expert_count, %output_size, %q8_input, %route_ids, %route_weights, %weight, %expected_output) : [index, index, index, index, index, index](index, index, index, index, index, index, tensor<2x864xi8>, tensor<1x4xi32>, tensor<1x2xf32>, tensor<2x2048x3x144xi8>, tensor<1x2048xf32>) + kernel.launch @qwen3_moe_attention_rmsnorm_quantize_q8_1_x4[%token_count](%token_count, %expected_output, %norm_weight, %expected_q8) : [index](index, tensor<1x2048xf32>, tensor<2048xf32>, tensor<2304xi8>) + kernel.launch @qwen3_moe_routed_down_q4k_q8_1_x4_next_q8[%token_count, %input_size, %route_count, %route_id_stride, %expert_count, %output_size](%token_count, %input_size, %route_count, %route_id_stride, %expert_count, %output_size, %q8_input, %route_ids, %route_weights, %weight, %actual_output0, %norm_weight, %completion_counter, %actual_q8_0) : [index, index, index, index, index, index](index, index, index, index, index, index, tensor<2x864xi8>, tensor<1x4xi32>, tensor<1x2xf32>, tensor<2x2048x3x144xi8>, tensor<1x2048xf32>, tensor<2048xf32>, tensor<1xi32>, tensor<2304xi8>) + kernel.launch @qwen3_moe_routed_down_q4k_q8_1_x4_next_q8[%token_count, %input_size, %route_count, %route_id_stride, %expert_count, %output_size](%token_count, %input_size, %route_count, %route_id_stride, %expert_count, %output_size, %q8_input, %route_ids, %route_weights, %weight, %actual_output1, %norm_weight, %completion_counter, %actual_q8_1) : [index, index, index, index, index, index](index, index, index, index, index, index, tensor<2x864xi8>, tensor<1x4xi32>, tensor<1x2xf32>, tensor<2x2048x3x144xi8>, tensor<1x2048xf32>, tensor<2048xf32>, tensor<1xi32>, tensor<2304xi8>) + check.expect.close actual(%actual_output0) expected(%expected_output) atol(0.000001) rtol(0.000001) nan(same) : tensor<1x2048xf32> + check.expect.close actual(%actual_output1) expected(%expected_output) atol(0.000001) rtol(0.000001) nan(same) : tensor<1x2048xf32> + check.expect.equal actual(%actual_q8_0) expected(%expected_q8) : tensor<2304xi8> + check.expect.equal actual(%actual_q8_1) expected(%expected_q8) : tensor<2304xi8> + check.expect.equal actual(%completion_counter) expected(%expected_counter) : tensor<1xi32> + check.return +} + +check.case public @qwen3_moe_routed_down_q4k_q8_1_x4_next_q8_benchmark_case { + %token_count = check.literal value(1) : index + %input_size = check.literal value(768) : index + %route_count = check.literal value(8) : index + %route_id_stride = check.literal value(8) : index + %expert_count = check.literal value(128) : index + %output_size = check.literal value(2048) : index + %q8_input = check.generate.fill value(0) : tensor<1x8x864xi8> + %route_ids = check.generate.iota offset(0) step(1) period(128) : tensor<1x8xi32> + %route_weights = check.generate.fill value(0.125) : tensor<1x8xf32> + %weight = check.generate.fill value(0) : tensor<128x2048x3x144xi8> + %output = check.generate.fill value(1.0) : tensor<1x2048xf32> + %norm_weight = check.generate.fill value(1.0) : tensor<2048xf32> + %completion_counter = check.generate.fill value(0) : tensor<1xi32> + %next_q8_output = check.generate.fill value(0) : tensor<2304xi8> + kernel.launch @qwen3_moe_routed_down_q4k_q8_1_x4_next_q8[%token_count, %input_size, %route_count, %route_id_stride, %expert_count, %output_size](%token_count, %input_size, %route_count, %route_id_stride, %expert_count, %output_size, %q8_input, %route_ids, %route_weights, %weight, %output, %norm_weight, %completion_counter, %next_q8_output) : [index, index, index, index, index, index](index, index, index, index, index, index, tensor<1x8x864xi8>, tensor<1x8xi32>, tensor<1x8xf32>, tensor<128x2048x3x144xi8>, tensor<1x2048xf32>, tensor<2048xf32>, tensor<1xi32>, tensor<2304xi8>) + check.return +} + +check.case public @qwen3_moe_routed_down_q4k_q8_1_x4_next_q8_composed_benchmark_case { + %token_count = check.literal value(1) : index + %input_size = check.literal value(768) : index + %route_count = check.literal value(8) : index + %route_id_stride = check.literal value(8) : index + %expert_count = check.literal value(128) : index + %output_size = check.literal value(2048) : index + %q8_input = check.generate.fill value(0) : tensor<1x8x864xi8> + %route_ids = check.generate.iota offset(0) step(1) period(128) : tensor<1x8xi32> + %route_weights = check.generate.fill value(0.125) : tensor<1x8xf32> + %weight = check.generate.fill value(0) : tensor<128x2048x3x144xi8> + %output = check.generate.fill value(1.0) : tensor<1x2048xf32> + %norm_weight = check.generate.fill value(1.0) : tensor<2048xf32> + %next_q8_output = check.generate.fill value(0) : tensor<2304xi8> + kernel.launch @qwen3_moe_routed_down_q4k_q8_1_x4[%token_count, %input_size, %route_count, %route_id_stride, %expert_count, %output_size](%token_count, %input_size, %route_count, %route_id_stride, %expert_count, %output_size, %q8_input, %route_ids, %route_weights, %weight, %output) : [index, index, index, index, index, index](index, index, index, index, index, index, tensor<1x8x864xi8>, tensor<1x8xi32>, tensor<1x8xf32>, tensor<128x2048x3x144xi8>, tensor<1x2048xf32>) + kernel.launch @qwen3_moe_attention_rmsnorm_quantize_q8_1_x4[%token_count](%token_count, %output, %norm_weight, %next_q8_output) : [index](index, tensor<1x2048xf32>, tensor<2048xf32>, tensor<2304xi8>) + check.return +} + +// Two selected experts exercise the noncompact route-ID stride, normalized +// weighted reduction, in-place residual update, the odd Q4_K block tail at +// K=768, and the output tile tail. Uniform 0xaa bytes decode to a deterministic +// nonzero row. +check.case public @qwen3_moe_routed_down_q4k_q8_1_x4_nonzero_residual_case { + %token_count = check.literal value(1) : index + %input_size = check.literal value(768) : index + %route_count = check.literal value(2) : index + %route_id_stride = check.literal value(4) : index + %expert_count = check.literal value(2) : index + %output_size = check.literal value(9) : index + %routed_input = check.generate.fill value(0.00390625) : tensor<2x768xf32> + %q8_input = check.generate.fill value(0) : tensor<2x864xi8> + %route_ids = check.generate.iota offset(0) step(1) period(2) : tensor<1x4xi32> + %route_weights = check.generate.fill value(0.5) : tensor<1x2xf32> + %weight = check.generate.fill value(-86) : tensor<2x9x3x144xi8> + %output = check.generate.fill value(1.0) : tensor<1x9xf32> + %expected = check.generate.fill value(-58.0394287109375) : tensor<1x9xf32> + %routed_row_count = check.literal value(2) : index + kernel.launch @ggml_quantize_q8_1_x4_f32[%routed_row_count, %input_size](%routed_row_count, %input_size, %routed_input, %q8_input) : [index, index](index, index, tensor<2x768xf32>, tensor<2x864xi8>) + kernel.launch @qwen3_moe_routed_down_q4k_q8_1_x4[%token_count, %input_size, %route_count, %route_id_stride, %expert_count, %output_size](%token_count, %input_size, %route_count, %route_id_stride, %expert_count, %output_size, %q8_input, %route_ids, %route_weights, %weight, %output) : [index, index, index, index, index, index](index, index, index, index, index, index, tensor<2x864xi8>, tensor<1x4xi32>, tensor<1x2xf32>, tensor<2x9x3x144xi8>, tensor<1x9xf32>) + check.expect.close actual(%output) expected(%expected) atol(0.25) rtol(9.9999999999999995e-07) nan(same) : tensor<1x9xf32> + check.return +} + +// Two tokens use different route-weight sums so token, route, packed-row, and +// output addressing cannot collapse to the single-token case. +check.case public @qwen3_moe_routed_down_q4k_q8_1_x4_nonzero_prefill_case { + %token_count = check.literal value(2) : index + %input_size = check.literal value(768) : index + %route_count = check.literal value(2) : index + %route_id_stride = check.literal value(4) : index + %expert_count = check.literal value(2) : index + %output_size = check.literal value(1) : index + %routed_row_count = check.literal value(4) : index + %routed_input = check.generate.fill value(0.00390625) : tensor<2x2x768xf32> + %q8_input = check.generate.fill value(0) : tensor<2x2x864xi8> + %route_ids = check.generate.iota offset(0) step(1) period(2) : tensor<2x4xi32> + %route_weights = check.generate.iota offset(0.25) step(0.25) period(4) : tensor<2x2xf32> + %weight = check.generate.fill value(-86) : tensor<2x1x3x144xi8> + %output = check.generate.fill value(1.0) : tensor<2x1xf32> + %expected = check.generate.iota offset(-43.279571533203125) step(-59.0394287109375) : tensor<2x1xf32> + kernel.launch @ggml_quantize_q8_1_x4_f32[%routed_row_count, %input_size](%routed_row_count, %input_size, %routed_input, %q8_input) : [index, index](index, index, tensor<2x2x768xf32>, tensor<2x2x864xi8>) + kernel.launch @qwen3_moe_routed_down_q4k_q8_1_x4[%token_count, %input_size, %route_count, %route_id_stride, %expert_count, %output_size](%token_count, %input_size, %route_count, %route_id_stride, %expert_count, %output_size, %q8_input, %route_ids, %route_weights, %weight, %output) : [index, index, index, index, index, index](index, index, index, index, index, index, tensor<2x2x864xi8>, tensor<2x4xi32>, tensor<2x2xf32>, tensor<2x1x3x144xi8>, tensor<2x1xf32>) + check.expect.close actual(%output) expected(%expected) atol(0.25) rtol(9.9999999999999995e-07) nan(same) : tensor<2x1xf32> + check.return +} + +check.case public @qwen3_moe_routed_down_q4k_q8_1_x4_benchmark_case { + %token_count = check.param.choice values([1, 2, 4, 8, 16, 17, 32, 63, 128, 129, 512]) name("token_count") : index + %input_size = check.literal value(768) : index + %route_count = check.literal value(8) : index + %route_id_stride = check.literal value(128) : index + %expert_count = check.literal value(128) : index + %output_size = check.literal value(2048) : index + %q8_input = check.generate.fill value(0) : tensor<[%token_count]x8x864xi8> + %route_ids = check.generate.iota offset(0) step(1) period(127) : tensor<[%token_count]x128xi32> + %route_weights = check.generate.fill value(0.125) : tensor<[%token_count]x8xf32> + %weight = check.generate.fill value(0) : tensor<128x2048x3x144xi8> + %output = check.generate.fill value(1.0) : tensor<[%token_count]x2048xf32> + %expected = check.generate.fill value(1.0) : tensor<[%token_count]x2048xf32> + kernel.launch @qwen3_moe_routed_down_q4k_q8_1_x4[%token_count, %input_size, %route_count, %route_id_stride, %expert_count, %output_size](%token_count, %input_size, %route_count, %route_id_stride, %expert_count, %output_size, %q8_input, %route_ids, %route_weights, %weight, %output) : [index, index, index, index, index, index](index, index, index, index, index, index, tensor<[%token_count]x8x864xi8>, tensor<[%token_count]x128xi32>, tensor<[%token_count]x8xf32>, tensor<128x2048x3x144xi8>, tensor<[%token_count]x2048xf32>) + check.expect.close actual(%output) expected(%expected) atol(0.0) rtol(0.0) nan(same) : tensor<[%token_count]x2048xf32> + check.return +} + +check.case public @qwen3_moe_routed_down_q4k_q8_1_x4_pipeline_benchmark_case { + %token_count = check.param.choice values([1, 2, 4, 8, 16, 17, 32, 63, 128, 129, 512]) name("token_count") : index + %input_size = check.literal value(768) : index + %route_count = check.literal value(8) : index + %route_id_stride = check.literal value(128) : index + %expert_count = check.literal value(128) : index + %output_size = check.literal value(2048) : index + // Each 768-element route contains six complete Q8_1 x4 groups, so packing + // the eight contiguous routes as one 6,144-element token row preserves the + // exact per-route physical layout without a derived testbench scalar. + %routed_input_size = check.literal value(6144) : index + %routed_input = check.generate.fill value(0.0) : tensor<[%token_count]x8x768xf32> + %q8_input = check.generate.fill value(1) : tensor<[%token_count]x8x864xi8> + %route_ids = check.generate.iota offset(0) step(1) period(127) : tensor<[%token_count]x128xi32> + %route_weights = check.generate.fill value(0.125) : tensor<[%token_count]x8xf32> + %weight = check.generate.fill value(0) : tensor<128x2048x3x144xi8> + %output = check.generate.fill value(1.0) : tensor<[%token_count]x2048xf32> + %expected = check.generate.fill value(1.0) : tensor<[%token_count]x2048xf32> + kernel.launch @ggml_quantize_q8_1_x4_f32[%token_count, %routed_input_size](%token_count, %routed_input_size, %routed_input, %q8_input) : [index, index](index, index, tensor<[%token_count]x8x768xf32>, tensor<[%token_count]x8x864xi8>) + kernel.launch @qwen3_moe_routed_down_q4k_q8_1_x4[%token_count, %input_size, %route_count, %route_id_stride, %expert_count, %output_size](%token_count, %input_size, %route_count, %route_id_stride, %expert_count, %output_size, %q8_input, %route_ids, %route_weights, %weight, %output) : [index, index, index, index, index, index](index, index, index, index, index, index, tensor<[%token_count]x8x864xi8>, tensor<[%token_count]x128xi32>, tensor<[%token_count]x8xf32>, tensor<128x2048x3x144xi8>, tensor<[%token_count]x2048xf32>) + check.expect.close actual(%output) expected(%expected) atol(0.0) rtol(0.0) nan(same) : tensor<[%token_count]x2048xf32> + check.return +} + +check.benchmark<@qwen3_moe_routed_down_q4k_q8_1_x4_nonzero_residual_case> @qwen3_moe_routed_down_q4k_q8_1_x4_small + +check.benchmark<@qwen3_moe_routed_down_q4k_q8_1_x4_benchmark_case> @qwen3_moe_routed_down_q4k_q8_1_x4_decode {token_count = 1} + +check.benchmark<@qwen3_moe_routed_down_q4k_q8_1_x4_benchmark_case> @qwen3_moe_routed_down_q4k_q8_1_x4_small_batch_2 {token_count = 2} + +check.benchmark<@qwen3_moe_routed_down_q4k_q8_1_x4_benchmark_case> @qwen3_moe_routed_down_q4k_q8_1_x4_small_batch_4 {token_count = 4} + +check.benchmark<@qwen3_moe_routed_down_q4k_q8_1_x4_benchmark_case> @qwen3_moe_routed_down_q4k_q8_1_x4_small_batch_8 {token_count = 8} + +check.benchmark<@qwen3_moe_routed_down_q4k_q8_1_x4_benchmark_case> @qwen3_moe_routed_down_q4k_q8_1_x4_small_batch_16 {token_count = 16} + +check.benchmark<@qwen3_moe_routed_down_q4k_q8_1_x4_benchmark_case> @qwen3_moe_routed_down_q4k_q8_1_x4_prefill_32 {token_count = 32} + +check.benchmark<@qwen3_moe_routed_down_q4k_q8_1_x4_benchmark_case> @qwen3_moe_routed_down_q4k_q8_1_x4_prefill_128 {token_count = 128} + +check.benchmark<@qwen3_moe_routed_down_q4k_q8_1_x4_benchmark_case> @qwen3_moe_routed_down_q4k_q8_1_x4_prefill_512 {token_count = 512} + +check.benchmark<@qwen3_moe_routed_down_q4k_q8_1_x4_pipeline_benchmark_case> @qwen3_moe_routed_down_q4k_q8_1_x4_pipeline_decode {token_count = 1} + +check.benchmark<@qwen3_moe_routed_down_q4k_q8_1_x4_pipeline_benchmark_case> @qwen3_moe_routed_down_q4k_q8_1_x4_pipeline_prefill_32 {token_count = 32} + +check.benchmark<@qwen3_moe_routed_down_q4k_q8_1_x4_pipeline_benchmark_case> @qwen3_moe_routed_down_q4k_q8_1_x4_pipeline_prefill_128 {token_count = 128} + +check.benchmark<@qwen3_moe_routed_down_q4k_q8_1_x4_pipeline_benchmark_case> @qwen3_moe_routed_down_q4k_q8_1_x4_pipeline_prefill_512 {token_count = 512} + +check.benchmark<@qwen3_moe_routed_down_q4k_q8_1_x4_next_q8_benchmark_case> @qwen3_moe_routed_down_q4k_q8_1_x4_next_q8_decode + +check.benchmark<@qwen3_moe_routed_down_q4k_q8_1_x4_next_q8_composed_benchmark_case> @qwen3_moe_routed_down_q4k_q8_1_x4_next_q8_composed_decode diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/routed_down_q6k.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/routed_down_q6k.loom new file mode 100644 index 000000000000..175b8a1d253e --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/routed_down_q6k.loom @@ -0,0 +1,1291 @@ +// Fuses the Qwen3 MoE Q6_K down projection with route weighting, top-8 +// reduction, and residual publication. The input contains one Q8_1 x4 row per +// [token, route], while weights remain in GGUF's raw +// [expert, output, K / 256, 210 bytes] layout. +// +// Route IDs retain an independent physical stride because llama.cpp selects a +// top-8 view from its 128-entry argsort storage. Normalized route weights are +// compact [token, route]. The output enters containing the residual and is +// updated in place, so the fused boundary never materializes an unweighted or +// route-indexed [token, route, hidden] down-projection tensor. +// +// The Q8_1 provider assigns one output channel to a wave and reduces all routes +// in registers. The direct-F32 provider mirrors llama.cpp's gfx1151 Vulkan +// decode algorithm: four wave64 subgroups compute four channels, each subgroup +// contracts four raw Q6_K blocks at a time, and route weighting and residual +// publication remain fused. A grouped prefill provider can reuse the same raw +// Q6_K row primitive while staging each expert row across multiple tokens. +config.decl @qwen3_moe.model.hidden_size : %value: index where [range(%value, 128, 32768), mul(%value, 128)] +config.decl @qwen3_moe.model.rms_epsilon : f32 + +func.def inline @ggml_q6k_dot4_f32(%lhs: vector<4xf32>, %rhs: vector<4xf32>) -> (f32) { + %c0 = scalar.constant 0.0 : f32 + %lhs0 = vector.extract %lhs[0] : vector<4xf32> -> f32 + %lhs1 = vector.extract %lhs[1] : vector<4xf32> -> f32 + %lhs2 = vector.extract %lhs[2] : vector<4xf32> -> f32 + %lhs3 = vector.extract %lhs[3] : vector<4xf32> -> f32 + %rhs0 = vector.extract %rhs[0] : vector<4xf32> -> f32 + %rhs1 = vector.extract %rhs[1] : vector<4xf32> -> f32 + %rhs2 = vector.extract %rhs[2] : vector<4xf32> -> f32 + %rhs3 = vector.extract %rhs[3] : vector<4xf32> -> f32 + %sum0 = scalar.fmaf %lhs0, %rhs0, %c0 : f32 + %sum1 = scalar.fmaf %lhs1, %rhs1, %sum0 : f32 + %sum2 = scalar.fmaf %lhs2, %rhs2, %sum1 : f32 + %sum3 = scalar.fmaf %lhs3, %rhs3, %sum2 : f32 + func.return %sum3 : f32 +} + +func.def inline @ggml_q6k_f32_block_row(%input_size: index, %row: index, %block: index, %frame_count: index, %frame: index, %lane: index, %weight: buffer, %scale_stage: buffer, %input0: vector<4xf32>, %input1: vector<4xf32>, %input2: vector<4xf32>, %input3: vector<4xf32>) -> (f32) { + %c0 = index.constant 0 : index + %c2 = index.constant 2 : index + %c4 = index.constant 4 : index + %c6 = index.constant 6 : index + %c8 = index.constant 8 : index + %c16 = index.constant 16 : index + %c128 = index.constant 128 : offset + %c208 = index.constant 208 : offset + %c210_bytes = index.constant 210 : offset + %c256 = index.constant 256 : index + %c0_offset = index.constant 0 : offset + %c4_i32v = vector.constant 4 : vector<1xi32> + %c2_i32v = vector.constant 2 : vector<1xi32> + %nibble_mask = vector.constant 252645135 : vector<1xi32> + %high0_mask = vector.constant 50529027 : vector<1xi32> + %high2_mask = vector.constant 202116108 : vector<1xi32> + %high4_mask = vector.constant 808464432 : vector<1xi32> + %high6_mask = vector.constant -1061109568 : vector<1xi32> + %c32_f32v = vector.constant 32.0 : vector<4xf32> + %c0_f32 = scalar.constant 0.0 : f32 + %bounded_lane = index.assume %lane [range(%lane, 0, 63)] : index + %block_count = index.div %input_size, %c256 : index + %valid_block = index.cmp ult, %block, %block_count : index + %safe_block = scf.if %valid_block -> (index) { + scf.yield %block : index + } else { + scf.yield %c0 : index + } + %itid = index.rem %bounded_lane, %c16 : index + %cohort = index.div %bounded_lane, %c16 : index + %vector_half = index.div %itid, %c8 : index + %vector_index = index.rem %itid, %c8 : index + %vector_quarter = index.div %vector_index, %c4 : index + %weight_row_bytes = index.scale %block_count, %c210_bytes : index, offset -> offset + %weight_row_byte_base = index.scale %row, %weight_row_bytes : index, offset -> offset + %block_byte_add = index.scale %safe_block, %c210_bytes : index, offset -> offset + %weight_block_byte_base = index.add %weight_row_byte_base, %block_byte_add : offset + %qh_byte_base = index.add %weight_block_byte_base, %c128 : offset + %d_byte_base = index.add %weight_block_byte_base, %c208 : offset + %ql_view = buffer.view %weight[%weight_block_byte_base] : buffer -> view<32xi32> + %qh_view = buffer.view %weight[%qh_byte_base] : buffer -> view<16xi32> + %d_view = buffer.view %weight[%d_byte_base] : buffer -> view<1xf16> + %scale_stage_view = buffer.view %scale_stage[%c0_offset] : buffer -> view<[%frame_count]x64xf32> + %scale_cohort_base = index.mul %cohort, %c16 : index + %scale_half_base = index.mul %vector_half, %c8 : index + %scale_index0 = index.add %scale_half_base, %vector_quarter : index + %scale_index1 = index.add %scale_index0, %c2 : index + %scale_index2 = index.add %scale_index0, %c4 : index + %scale_index3 = index.add %scale_index0, %c6 : index + %stage_scale_index0 = index.add %scale_cohort_base, %scale_index0 : index + %stage_scale_index1 = index.add %scale_cohort_base, %scale_index1 : index + %stage_scale_index2 = index.add %scale_cohort_base, %scale_index2 : index + %stage_scale_index3 = index.add %scale_cohort_base, %scale_index3 : index + %scale0 = view.load %scale_stage_view[%frame, %stage_scale_index0] : view<[%frame_count]x64xf32> -> f32 + %scale1 = view.load %scale_stage_view[%frame, %stage_scale_index1] : view<[%frame_count]x64xf32> -> f32 + %scale2 = view.load %scale_stage_view[%frame, %stage_scale_index2] : view<[%frame_count]x64xf32> -> f32 + %scale3 = view.load %scale_stage_view[%frame, %stage_scale_index3] : view<[%frame_count]x64xf32> -> f32 + %ql_half_word_base = index.mul %vector_half, %c16 : index + %ql_word_index00 = index.add %ql_half_word_base, %vector_index : index + %ql_word_index10 = index.add %ql_word_index00, %c8 : index + %qh_half_word_base = index.mul %vector_half, %c8 : index + %qh_word_index0 = index.add %qh_half_word_base, %vector_index : index + %ql_word_index0, %ql_word_index1, %qh_word_index = index.assume %ql_word_index00, %ql_word_index10, %qh_word_index0 [range(%ql_word_index00, 0, 23), range(%ql_word_index10, 8, 31), range(%qh_word_index0, 0, 15)] : index, index, index + %ql_word0 = vector.load %ql_view[%ql_word_index0] : view<32xi32> -> vector<1xi32> + %ql_word1 = vector.load %ql_view[%ql_word_index1] : view<32xi32> -> vector<1xi32> + %qh_word = vector.load %qh_view[%qh_word_index] : view<16xi32> -> vector<1xi32> + %ql0 = vector.andi %ql_word0, %nibble_mask : vector<1xi32> + %ql1 = vector.andi %ql_word1, %nibble_mask : vector<1xi32> + %ql_word0_high = vector.shrui %ql_word0, %c4_i32v : vector<1xi32> + %ql_word1_high = vector.shrui %ql_word1, %c4_i32v : vector<1xi32> + %ql2 = vector.andi %ql_word0_high, %nibble_mask : vector<1xi32> + %ql3 = vector.andi %ql_word1_high, %nibble_mask : vector<1xi32> + %qh0_low = vector.andi %qh_word, %high0_mask : vector<1xi32> + %qh1_low = vector.andi %qh_word, %high2_mask : vector<1xi32> + %qh2 = vector.andi %qh_word, %high4_mask : vector<1xi32> + %qh3_high = vector.andi %qh_word, %high6_mask : vector<1xi32> + %qh0 = vector.shli %qh0_low, %c4_i32v : vector<1xi32> + %qh1 = vector.shli %qh1_low, %c2_i32v : vector<1xi32> + %qh3 = vector.shrui %qh3_high, %c2_i32v : vector<1xi32> + %code0 = vector.ori %ql0, %qh0 : vector<1xi32> + %code1 = vector.ori %ql1, %qh1 : vector<1xi32> + %code2 = vector.ori %ql2, %qh2 : vector<1xi32> + %code3 = vector.ori %ql3, %qh3 : vector<1xi32> + %code0_i8 = vector.bitcast %code0 : vector<1xi32> to vector<4xi8> + %code1_i8 = vector.bitcast %code1 : vector<1xi32> to vector<4xi8> + %code2_i8 = vector.bitcast %code2 : vector<1xi32> to vector<4xi8> + %code3_i8 = vector.bitcast %code3 : vector<1xi32> to vector<4xi8> + %code0_f32 = vector.uitofp %code0_i8 : vector<4xi8> to vector<4xf32> + %code1_f32 = vector.uitofp %code1_i8 : vector<4xi8> to vector<4xf32> + %code2_f32 = vector.uitofp %code2_i8 : vector<4xi8> to vector<4xf32> + %code3_f32 = vector.uitofp %code3_i8 : vector<4xi8> to vector<4xf32> + %q0 = vector.subf %code0_f32, %c32_f32v : vector<4xf32> + %q1 = vector.subf %code1_f32, %c32_f32v : vector<4xf32> + %q2 = vector.subf %code2_f32, %c32_f32v : vector<4xf32> + %q3 = vector.subf %code3_f32, %c32_f32v : vector<4xf32> + %dot0 = func.call @ggml_q6k_dot4_f32(%input0, %q0) : (vector<4xf32>, vector<4xf32>) -> (f32) + %dot1 = func.call @ggml_q6k_dot4_f32(%input1, %q1) : (vector<4xf32>, vector<4xf32>) -> (f32) + %dot2 = func.call @ggml_q6k_dot4_f32(%input2, %q2) : (vector<4xf32>, vector<4xf32>) -> (f32) + %dot3 = func.call @ggml_q6k_dot4_f32(%input3, %q3) : (vector<4xf32>, vector<4xf32>) -> (f32) + %scaled3 = scalar.mulf %dot3, %scale3 : f32 + %scaled2 = scalar.fmaf %dot2, %scale2, %scaled3 : f32 + %scaled1 = scalar.fmaf %dot1, %scale1, %scaled2 : f32 + %scaled0 = scalar.fmaf %dot0, %scale0, %scaled1 : f32 + %d_f16 = view.load %d_view[0] : view<1xf16> -> f16 + %d = scalar.extf %d_f16 : f16 to f32 + %contribution0 = scalar.mulf %scaled0, %d : f32 + %contribution = scf.if %valid_block -> (f32) { + scf.yield %contribution0 : f32 + } else { + scf.yield %c0_f32 : f32 + } + func.return %contribution : f32 +} + +func.def inline @ggml_q6k_load_f32_block(%token_count: index, %input_size: index, %token: index, %block: index, %lane: index, %input: buffer) -> (vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>) { + %c0 = index.constant 0 : index + %c4 = index.constant 4 : index + %c8 = index.constant 8 : index + %c16 = index.constant 16 : index + %c32 = index.constant 32 : index + %c64 = index.constant 64 : index + %c96 = index.constant 96 : index + %c128 = index.constant 128 : index + %c256 = index.constant 256 : index + %c0_offset = index.constant 0 : offset + %bounded_lane = index.assume %lane [range(%lane, 0, 63)] : index + %block_count = index.div %input_size, %c256 : index + %valid_block = index.cmp ult, %block, %block_count : index + %safe_block0 = scf.select %valid_block, %block, %c0 : index + %safe_block, %launch_block_count = index.assume %safe_block0, %block_count [lt(%safe_block0, %block_count)] : index, index + %itid = index.rem %bounded_lane, %c16 : index + %vector_half = index.div %itid, %c8 : index + %vector_index = index.rem %itid, %c8 : index + %input_view = buffer.view %input[%c0_offset] : buffer -> view<[%token_count]x[%launch_block_count]x256xf32> + %vector_half_base = index.mul %vector_half, %c128 : index + %vector_offset = index.mul %vector_index, %c4 : index + %input_index0 = index.add %vector_half_base, %vector_offset : index + %input_index1 = index.add %input_index0, %c32 : index + %input_index2 = index.add %input_index0, %c64 : index + %input_index3 = index.add %input_index0, %c96 : index + %input0 = vector.load %input_view[%token, %safe_block, %input_index0] : view<[%token_count]x[%launch_block_count]x256xf32> -> vector<4xf32> + %input1 = vector.load %input_view[%token, %safe_block, %input_index1] : view<[%token_count]x[%launch_block_count]x256xf32> -> vector<4xf32> + %input2 = vector.load %input_view[%token, %safe_block, %input_index2] : view<[%token_count]x[%launch_block_count]x256xf32> -> vector<4xf32> + %input3 = vector.load %input_view[%token, %safe_block, %input_index3] : view<[%token_count]x[%launch_block_count]x256xf32> -> vector<4xf32> + func.return %input0, %input1, %input2, %input3 : vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32> +} + +func.def inline @ggml_q6k_load_f32_scale(%input_size: index, %row: index, %block: index, %lane: index, %weight: buffer) -> (f32) { + %c0 = index.constant 0 : index + %c16 = index.constant 16 : index + %c192 = index.constant 192 : offset + %c210_bytes = index.constant 210 : offset + %c256 = index.constant 256 : index + %bounded_lane = index.assume %lane [range(%lane, 0, 63)] : index + %block_count = index.div %input_size, %c256 : index + %valid_block = index.cmp ult, %block, %block_count : index + %safe_block = scf.if %valid_block -> (index) { + scf.yield %block : index + } else { + scf.yield %c0 : index + } + %itid = index.rem %bounded_lane, %c16 : index + %weight_row_bytes = index.scale %block_count, %c210_bytes : index, offset -> offset + %weight_row_byte_base = index.scale %row, %weight_row_bytes : index, offset -> offset + %block_byte_add = index.scale %safe_block, %c210_bytes : index, offset -> offset + %weight_block_byte_base = index.add %weight_row_byte_base, %block_byte_add : offset + %scale_byte_base = index.add %weight_block_byte_base, %c192 : offset + %scale_view = buffer.view %weight[%scale_byte_base] : buffer -> view<16xi8> + %scale_i8 = view.load %scale_view[%itid] : view<16xi8> -> i8 + %scale = scalar.sitofp %scale_i8 : i8 to f32 + func.return %scale : f32 +} + +func.def inline @ggml_q6k_q8_1_x4_block_lane(%weight: buffer, %weight_row_byte_base: offset, %q8_input: buffer, %q8_row_byte_base: offset, %q6_block: index, %lane: index) -> (f32) { + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c4 = index.constant 4 : index + %c8 = index.constant 8 : index + %c16 = index.constant 16 : index + %block_bytes = index.constant 210 : offset + %qh_byte_add = index.constant 128 : offset + %scale_byte_add = index.constant 192 : offset + %d_byte_add = index.constant 208 : offset + %c4_i32v = vector.constant 4 : vector<1xi32> + %c0_i32v = vector.constant 0 : vector<1xi32> + %nibble_mask = vector.constant 252645135 : vector<1xi32> + %high_mask = vector.constant 808464432 : vector<1xi32> + %c0_f32 = scalar.constant 0.0 : f32 + %bounded_lane = index.assume %lane [range(%lane, 0, 31)] : index + %block_byte_add = index.scale %q6_block, %block_bytes : index, offset -> offset + %block_byte_base = index.add %weight_row_byte_base, %block_byte_add : offset + %qh_byte_base = index.add %block_byte_base, %qh_byte_add : offset + %scale_byte_base = index.add %block_byte_base, %scale_byte_add : offset + %d_byte_base = index.add %block_byte_base, %d_byte_add : offset + %ql_view = buffer.view %weight[%block_byte_base] : buffer -> view<32xi32> + %qh_view = buffer.view %weight[%qh_byte_base] : buffer -> view<16xi32> + %scale_view = buffer.view %weight[%scale_byte_base] : buffer -> view<16xi8> + %d_view = buffer.view %weight[%d_byte_base] : buffer -> view<1xf16> + %lane_mod8 = index.rem %bounded_lane, %c8 : index + %lane_mod16 = index.rem %bounded_lane, %c16 : index + %lane_div16 = index.div %bounded_lane, %c16 : index + %lane_div8_in_16 = index.div %lane_mod16, %c8 : index + %lane_div4_in_16 = index.div %lane_mod16, %c4 : index + %qh_high_base = index.mul %lane_div16, %c8 : index + %qh_index0 = index.add %qh_high_base, %lane_mod8 : index + %qh_index = index.assume %qh_index0 [range(%qh_index0, 0, 15)] : index + %ql_word = vector.load %ql_view[%bounded_lane] : view<32xi32> -> vector<1xi32> + %qh_word = vector.load %qh_view[%qh_index] : view<16xi32> -> vector<1xi32> + %qh_base_shift_index = index.mul %lane_div8_in_16, %c2 : index + %qh_base_shift_i32 = index.cast %qh_base_shift_index : index to i32 + %q8_block_base = index.mul %q6_block, %c8 : index + %q8_high_add = index.mul %lane_div16, %c4 : index + %q8_quadrant = index.add %q8_high_add, %lane_div8_in_16 : index + %scale_high_base = index.mul %lane_div16, %c8 : index + %scale_lane0 = index.add %scale_high_base, %lane_div4_in_16 : index + %d_f16 = view.load %d_view[0] : view<1xf16> -> f16 + %d = scalar.extf %d_f16 : f16 to f32 + %sum = scf.for %part = [%c0 to %c2 step %c1](%accumulator = %c0_f32 : f32) -> (f32) unroll { + %bounded_part = index.assume %part [range(%part, 0, 1)] : index + %part_shift_index = index.mul %bounded_part, %c4 : index + %part_shift_i32 = index.cast %part_shift_index : index to i32 + %part_shift = vector.splat %part_shift_i32 : vector<1xi32> + %ql_shifted = vector.shrui %ql_word, %part_shift : vector<1xi32> + %ql = vector.andi %ql_shifted, %nibble_mask : vector<1xi32> + %qh_shift_i32 = scalar.addi %qh_base_shift_i32, %part_shift_i32 : i32 + %qh_shift = vector.splat %qh_shift_i32 : vector<1xi32> + %qh_shifted = vector.shrui %qh_word, %qh_shift : vector<1xi32> + %qh_positioned = vector.shli %qh_shifted, %c4_i32v : vector<1xi32> + %qh = vector.andi %qh_positioned, %high_mask : vector<1xi32> + %code = vector.ori %ql, %qh : vector<1xi32> + %signed_weight = func.call @ggml_q6k_sign_extend_dot4(%code) : (vector<1xi32>) -> (vector<4xi8>) + %q8_part_add = index.mul %bounded_part, %c2 : index + %q8_block_part = index.add %q8_block_base, %q8_part_add : index + %q8_block = index.add %q8_block_part, %q8_quadrant : index + %q8_values, %q8_d = func.call @ggml_q8_1_x4_word(%q8_input, %q8_row_byte_base, %q8_block, %lane_mod8) : (buffer, offset, index, index) -> (vector<4xi8>, f32) + %scale_lane = index.add %scale_lane0, %part_shift_index : index + %scale_i8 = view.load %scale_view[%scale_lane] : view<16xi8> -> i8 + %scale = scalar.sitofp %scale_i8 : i8 to f32 + %dot = vector.dot4i %signed_weight, %q8_values, %c0_i32v : vector<4xi8>, vector<4xi8>, vector<1xi32> + %dot_i32 = vector.extract %dot[0] : vector<1xi32> -> i32 + %dot_f32 = scalar.sitofp %dot_i32 : i32 to f32 + %scaled0 = scalar.mulf %dot_f32, %scale : f32 + %scaled1 = scalar.mulf %scaled0, %d : f32 + %contribution = scalar.mulf %scaled1, %q8_d : f32 + %next = scalar.addf %accumulator, %contribution : f32 + scf.yield %next : f32 + } + func.return %sum : f32 +} + +func.def inline @ggml_q6k_q8_1_x4_row_lane(%input_size: index, %weight: buffer, %weight_row_byte_base: offset, %q8_input: buffer, %q8_row_byte_base: offset, %lane: index) -> (f32) { + %bounded_input_size = index.assume %input_size [range(%input_size, 256, 32768), mul(%input_size, 256)] : index + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c256 = index.constant 256 : index + %c0_f32 = scalar.constant 0.0 : f32 + %block_count = index.div %bounded_input_size, %c256 : index + %result = scf.for %block = [%c0 to %block_count step %c1](%block_acc = %c0_f32 : f32) -> (f32) { + %contribution = func.call @ggml_q6k_q8_1_x4_block_lane(%weight, %weight_row_byte_base, %q8_input, %q8_row_byte_base, %block, %lane) : (buffer, offset, buffer, offset, index, index) -> (f32) + %next = scalar.addf %block_acc, %contribution : f32 + scf.yield %next : f32 + } + func.return %result : f32 +} + +func.def inline @ggml_q6k_sign_extend_dot4(%code: vector<1xi32>) -> (vector<4xi8>) { + %c1_i32v = vector.constant 1 : vector<1xi32> + %c2_i32v = vector.constant 2 : vector<1xi32> + %low5_mask = vector.constant 522133279 : vector<1xi32> + %bit5_mask = vector.constant 538976288 : vector<1xi32> + %sign_mask = vector.constant -522133280 : vector<1xi32> + %low5 = vector.andi %code, %low5_mask : vector<1xi32> + %bit5 = vector.andi %code, %bit5_mask : vector<1xi32> + %bit6 = vector.shli %bit5, %c1_i32v : vector<1xi32> + %bit7 = vector.shli %bit5, %c2_i32v : vector<1xi32> + %high01 = vector.ori %bit5, %bit6 : vector<1xi32> + %high = vector.ori %high01, %bit7 : vector<1xi32> + %sign = vector.xori %high, %sign_mask : vector<1xi32> + %signed_i32 = vector.ori %low5, %sign : vector<1xi32> + %signed = vector.bitcast %signed_i32 : vector<1xi32> to vector<4xi8> + func.return %signed : vector<4xi8> +} + +func.def inline @ggml_q6k_stage_f32_scales(%input_size: index, %row: index, %block: index, %frame_count: index, %frame: index, %lane: index, %weight: buffer, %scale_stage: buffer) { + %c0_offset = index.constant 0 : offset + %bounded_lane = index.assume %lane [range(%lane, 0, 63)] : index + %scale_stage_view = buffer.view %scale_stage[%c0_offset] : buffer -> view<[%frame_count]x64xf32> + %scale = func.call @ggml_q6k_load_f32_scale(%input_size, %row, %block, %bounded_lane, %weight) : (index, index, index, index, buffer) -> (f32) + view.store %scale, %scale_stage_view[%frame, %bounded_lane] : f32, view<[%frame_count]x64xf32> + kernel.barrier scope(workgroup) ordering(acq_rel) + func.return +} + +func.def inline @ggml_q8_1_x4_word(%q8_input: buffer, %row_byte_base: offset, %q8_block: index, %word_in_block: index) -> (vector<4xi8>, f32) { + %c2 = index.constant 2 : index + %c4 = index.constant 4 : index + %c8 = index.constant 8 : index + %group_bytes = index.constant 144 : offset + %payload_byte_add = index.constant 16 : offset + %group = index.div %q8_block, %c4 : index + %inner0 = index.rem %q8_block, %c4 : index + %inner = index.assume %inner0 [range(%inner0, 0, 3)] : index + %word0 = index.assume %word_in_block [range(%word_in_block, 0, 7)] : index + %group_byte_add = index.scale %group, %group_bytes : index, offset -> offset + %group_byte_base = index.add %row_byte_base, %group_byte_add : offset + %payload_byte_base = index.add %group_byte_base, %payload_byte_add : offset + %ds_view = buffer.view %q8_input[%group_byte_base] : buffer -> view<8xf16> + %payload_view = buffer.view %q8_input[%payload_byte_base] : buffer -> view<32xi32> + %d_index = index.mul %inner, %c2 : index + %inner_word_base = index.mul %inner, %c8 : index + %word_index = index.add %inner_word_base, %word0 : index + %d_f16 = view.load %ds_view[%d_index] : view<8xf16> -> f16 + %packed = vector.load %payload_view[%word_index] : view<32xi32> -> vector<1xi32> + %values = vector.bitcast %packed : vector<1xi32> to vector<4xi8> + %d = scalar.extf %d_f16 : f16 to f32 + func.return %values, %d : vector<4xi8>, f32 +} + +func.def inline @qwen3_moe_rmsnorm_quantize_q8_1_x4_body(%publish_normalized: i1, %reduction_subgroup_count0: index, %token_count: index, %token0: index, %input: buffer, %weight: buffer, %normalized_output: buffer, %q8_output: buffer) { + %hidden_size0 = config.get @qwen3_moe.model.hidden_size : index + %epsilon = config.get @qwen3_moe.model.rms_epsilon : f32 + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 2048)] : index + %reduction_subgroup_count = index.assume %reduction_subgroup_count0 [range(%reduction_subgroup_count0, 1, 8)] : index + %hidden_size = index.assume %hidden_size0 [range(%hidden_size0, 128, 32768), mul(%hidden_size0, 128)] : index + %workitem = kernel.workitem.id : index + %subgroup0 = kernel.subgroup.id : index + %subgroup = index.assume %subgroup0 [range(%subgroup0, 0, 7)] : index + %lane = kernel.subgroup.lane.id : index + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c4 = index.constant 4 : index + %c8 = index.constant 8 : index + %c16 = index.constant 16 : index + %c32 = index.constant 32 : index + %c128 = index.constant 128 : index + %c256 = index.constant 256 : index + %c1024 = index.constant 1024 : index + %group_bytes = index.constant 144 : offset + %payload_byte_add = index.constant 16 : offset + %scratch_d_byte_add = index.constant 1024 : offset + %scratch_bytes = index.constant 1152 : offset + %c0_offset = index.constant 0 : offset + %c0_f32 = scalar.constant 0.0 : f32 + %c1_f32 = scalar.constant 1.0 : f32 + %c127 = scalar.constant 127.0 : f32 + %c0_f32x4 = vector.constant 0.0 : vector<4xf32> + %valid_token = index.cmp ult, %token0, %bounded_token_count : index + %safe_token0 = scf.select %valid_token, %token0, %c0 : index + %token, %launch_token_count = index.assume %safe_token0, %bounded_token_count [lt(%safe_token0, %bounded_token_count)] : index, index + %hidden_size_i32 = index.cast %hidden_size : index to i32 + %hidden_size_f32 = scalar.sitofp %hidden_size_i32 : i32 to f32 + %physical_group_count = index.div %hidden_size, %c128 : index + %row_bytes = index.scale %physical_group_count, %group_bytes : index, offset -> offset + %token_output_byte_base = index.scale %token, %row_bytes : index, offset -> offset + %input_noalias, %weight_noalias, %q8_output_noalias = buffer.assume.noalias %input, %weight, %q8_output : buffer, buffer, buffer + %input_view = buffer.view %input_noalias[%c0_offset] : buffer -> view<[%launch_token_count]x[%hidden_size]xf32> + %weight_view = buffer.view %weight_noalias[%c0_offset] : buffer -> view<[%hidden_size]xf32> + %normalized_output_view = buffer.view %normalized_output[%c0_offset] : buffer -> view<[%launch_token_count]x[%hidden_size]xf32> + %scratch = buffer.alloca align(16) %scratch_bytes : buffer + %scratch_values = buffer.view %scratch[%c0_offset] : buffer -> view<256xf32> + %scratch_d = buffer.view %scratch[%scratch_d_byte_add] : buffer -> view<32xf32> + // Reduce the complete row before any block-local quantization. + %thread_sum = scf.for %channel = [%workitem to %hidden_size step %c256](%running_sum = %c0_f32 : f32) -> (f32) { + %value = view.load %input_view[%token, %channel] : view<[%launch_token_count]x[%hidden_size]xf32> -> f32 + %square = scalar.mulf %value, %value : f32 + %next_sum = scalar.addf %running_sum, %square : f32 + scf.yield %next_sum : f32 + } + %subgroup_sum = kernel.subgroup.reduce %thread_sum : f32 + %is_subgroup_leader = index.cmp eq, %lane, %c0 : index + scf.if %is_subgroup_leader { + view.store %subgroup_sum, %scratch_values[%subgroup] : f32, view<256xf32> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + %is_reduction_subgroup = index.cmp eq, %subgroup, %c0 : index + %is_reduction_lane = index.cmp ult, %lane, %reduction_subgroup_count : index + %loads_subgroup_sum = scalar.andi %is_reduction_subgroup, %is_reduction_lane : i1 + %subgroup_partial = scf.if %loads_subgroup_sum -> (f32) { + %value = view.load %scratch_values[%lane] : view<256xf32> -> f32 + scf.yield %value : f32 + } else { + scf.yield %c0_f32 : f32 + } + %row_sum = kernel.subgroup.reduce %subgroup_partial : f32 + %writes_scale = scalar.andi %is_reduction_subgroup, %is_subgroup_leader : i1 + scf.if %writes_scale { + %mean = scalar.divf %row_sum, %hidden_size_f32 : f32 + %biased_mean = scalar.addf %mean, %epsilon : f32 + %scale = scalar.rsqrtf %biased_mean : f32 + view.store %scale, %scratch_values[%c0] : f32, view<256xf32> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + %row_scale = view.load %scratch_values[%c0] : view<256xf32> -> f32 + %row_scale_vector = vector.splat %row_scale : vector<4xf32> + %publishes_normalized = scalar.andi %publish_normalized, %valid_token : i1 + // Reuse one scratch frame for each 1024-element stripe. Hidden sizes need + // only be divisible by 128; inactive workitems in the final stripe carry + // zeros and never publish. + scf.for %stripe_base = [%c0 to %hidden_size step %c1024] { + %word_element_add = index.mul %workitem, %c4 : index + %channel = index.add %stripe_base, %word_element_add : index + %valid_word = index.cmp ult, %channel, %hidden_size : index + %mask = vector.mask.range [%channel to %hidden_size step %c1] : index -> vector<4xi1> + %input_values = vector.load.mask %input_view[%token, %channel], %mask, %c0_f32x4 : view<[%launch_token_count]x[%hidden_size]xf32>, vector<4xi1>, vector<4xf32> + %learned_weights = vector.load.mask %weight_view[%channel], %mask, %c0_f32x4 : view<[%hidden_size]xf32>, vector<4xi1>, vector<4xf32> + %normalized0 = vector.mulf %input_values, %row_scale_vector : vector<4xf32> + %normalized = vector.mulf %normalized0, %learned_weights : vector<4xf32> + scf.if %publishes_normalized { + vector.store.mask %normalized, %normalized_output_view[%token, %channel], %mask : vector<4xf32>, view<[%launch_token_count]x[%hidden_size]xf32>, vector<4xi1> + } + %absolute_values = vector.absf %normalized : vector<4xf32> + %thread_max = vector.reduce %absolute_values, %c0_f32 : vector<4xf32>, f32 + view.store %thread_max, %scratch_values[%workitem] : f32, view<256xf32> + kernel.barrier scope(workgroup) ordering(acq_rel) + %word_in_block = index.rem %workitem, %c8 : index + %block_in_stripe = index.div %workitem, %c8 : index + %is_block_leader = index.cmp eq, %word_in_block, %c0 : index + %writes_block_d = scalar.andi %valid_word, %is_block_leader : i1 + scf.if %writes_block_d { + %cohort_base = index.mul %block_in_stripe, %c8 : index + %cohort_maxima = vector.load %scratch_values[%cohort_base] : view<256xf32> -> vector<8xf32> + %amax = vector.reduce %cohort_maxima, %c0_f32 : vector<8xf32>, f32 + %d = scalar.divf %amax, %c127 : f32 + view.store %d, %scratch_d[%block_in_stripe] : f32, view<32xf32> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + %d = scf.if %valid_word -> (f32) { + %block_d = view.load %scratch_d[%block_in_stripe] : view<32xf32> -> f32 + scf.yield %block_d : f32 + } else { + scf.yield %c0_f32 : f32 + } + %d_nonzero = scalar.cmpf one, %d, %c0_f32 : f32 + %d_inverse = scf.if %d_nonzero -> (f32) { + %inverse = scalar.divf %c1_f32, %d : f32 + scf.yield %inverse : f32 + } else { + scf.yield %c0_f32 : f32 + } + %d_inverse_vector = vector.splat %d_inverse : vector<4xf32> + %scaled_values = vector.mulf %normalized, %d_inverse_vector : vector<4xf32> + %rounded_values = vector.roundf %scaled_values : vector<4xf32> + %quantized_values = vector.fptosi %rounded_values : vector<4xf32> to vector<4xi8> + %packed_word = vector.bitcast %quantized_values : vector<4xi8> to vector<1xi32> + %publishes_q8_word = scalar.andi %valid_word, %valid_token : i1 + scf.if %publishes_q8_word { + %q8_block = index.div %channel, %c32 : index + %physical_group = index.div %q8_block, %c4 : index + %block_in_group = index.rem %q8_block, %c4 : index + %group_byte_add = index.scale %physical_group, %group_bytes : index, offset -> offset + %group_byte_offset = index.add %token_output_byte_base, %group_byte_add : offset + %payload_byte_offset = index.add %group_byte_offset, %payload_byte_add : offset + %group_ds = buffer.view %q8_output_noalias[%group_byte_offset] : buffer -> view<8xf16> + %group_qs = buffer.view %q8_output_noalias[%payload_byte_offset] : buffer -> view<32xi32> + %block_word_base = index.mul %block_in_group, %c8 : index + %packed_word_index0 = index.add %block_word_base, %word_in_block : index + %packed_word_index = index.assume %packed_word_index0 [range(%packed_word_index0, 0, 31)] : index + vector.store %packed_word, %group_qs[%packed_word_index] : vector<1xi32>, view<32xi32> + } + %thread_quantized_sum = vector.reduce %rounded_values, %c0_f32 : vector<4xf32>, f32 + view.store %thread_quantized_sum, %scratch_values[%workitem] : f32, view<256xf32> + kernel.barrier scope(workgroup) ordering(acq_rel) + %publishes_block_ds = scalar.andi %writes_block_d, %valid_token : i1 + scf.if %publishes_block_ds { + %cohort_base = index.mul %block_in_stripe, %c8 : index + %cohort_sums = vector.load %scratch_values[%cohort_base] : view<256xf32> -> vector<8xf32> + %quantized_sum = vector.reduce %cohort_sums, %c0_f32 : vector<8xf32>, f32 + %s = scalar.mulf %quantized_sum, %d : f32 + %q8_block = index.div %channel, %c32 : index + %physical_group = index.div %q8_block, %c4 : index + %block_in_group = index.rem %q8_block, %c4 : index + %group_byte_add = index.scale %physical_group, %group_bytes : index, offset -> offset + %group_byte_offset = index.add %token_output_byte_base, %group_byte_add : offset + %group_ds = buffer.view %q8_output_noalias[%group_byte_offset] : buffer -> view<8xf16> + %d_f16 = scalar.fptrunc %d : f32 to f16 + %s_f16 = scalar.fptrunc %s : f32 to f16 + %ds_index = index.mul %block_in_group, %c2 : index + %s_index = index.add %ds_index, %c1 : index + view.store %d_f16, %group_ds[%ds_index] : f16, view<8xf16> + view.store %s_f16, %group_ds[%s_index] : f16, view<8xf16> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + } + func.return +} + +func.def inline @qwen3_moe_routed_down_next_q8_completion(%output_channels_per_workgroup0: index, %reduction_subgroup_count0: index, %token_count: index, %output_size: index, %output: buffer, %norm_weight: buffer, %completion_counter: buffer, %next_q8_output: buffer) { + %output_channels_per_workgroup = index.assume %output_channels_per_workgroup0 [range(%output_channels_per_workgroup0, 1, 8)] : index + %reduction_subgroup_count = index.assume %reduction_subgroup_count0 [range(%reduction_subgroup_count0, 1, 8)] : index + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 1)] : index + %bounded_output_size = index.assume %output_size [range(%output_size, 128, 32768), mul(%output_size, 128)] : index + %token0 = kernel.workgroup.id : index + %token = index.assume %token0 [lt(%token0, %bounded_token_count)] : index + %c0 = index.constant 0 : index + %workitem = kernel.workitem.id : index + %output_tile_count = index.div %bounded_output_size, %output_channels_per_workgroup : index + %is_arrival_workitem = index.cmp eq, %workitem, %c0 : index + %c0_i32 = scalar.constant 0 : i32 + %c1_i32 = scalar.constant 1 : i32 + %c0_offset = index.constant 0 : offset + %counter_scratch_bytes = index.constant 4 : offset + %output_noalias, %norm_weight_noalias, %completion_counter_noalias, %next_q8_output_noalias = buffer.assume.noalias %output, %norm_weight, %completion_counter, %next_q8_output : buffer, buffer, buffer, buffer + %completion_counter_aligned = buffer.assume.alignment %completion_counter_noalias {minimum_alignment = 16} : buffer + %completion_counter_view = buffer.view %completion_counter_aligned[%c0_offset] : buffer -> view<1xi32> + %counter_scratch = buffer.alloca align(4) %counter_scratch_bytes : buffer + %counter_scratch_view = buffer.view %counter_scratch[%c0_offset] : buffer -> view<1xi32> + // Publish every producer's residual stores before the leader advances one + // workgroup arrival. The last arrival then acquires the complete row. + kernel.barrier scope(workgroup) ordering(release) + scf.if %is_arrival_workitem { + %old_counter = view.atomic.rmw %c1_i32, %completion_counter_view[%c0] {ordering = acq_rel, scope = device} : i32, view<1xi32> -> i32 + view.store %old_counter, %counter_scratch_view[%c0] : i32, view<1xi32> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + %old_counter = view.load %counter_scratch_view[%c0] : view<1xi32> -> i32 + %output_tile_count_i32 = index.cast %output_tile_count : index to i32 + %last_output_tile_i32 = scalar.subi %output_tile_count_i32, %c1_i32 : i32 + %negative_output_tile_count_i32 = scalar.subi %c0_i32, %output_tile_count_i32 : i32 + %is_last_output_tile = scalar.cmpi eq, %old_counter, %last_output_tile_i32 : i32 + scf.if %is_last_output_tile { + kernel.barrier scope(workgroup) ordering(acquire) + %publish_normalized = scalar.constant false : i1 + func.call @qwen3_moe_rmsnorm_quantize_q8_1_x4_body(%publish_normalized, %reduction_subgroup_count, %bounded_token_count, %token, %output_noalias, %norm_weight_noalias, %next_q8_output_noalias, %next_q8_output_noalias) : (i1, index, index, index, buffer, buffer, buffer, buffer) + kernel.barrier scope(workgroup) ordering(release) + scf.if %is_arrival_workitem { + view.atomic.reduce %negative_output_tile_count_i32, %completion_counter_view[%c0] {ordering = release, scope = device} : i32, view<1xi32> + } + } + func.return +} + +amdgpu.target @qwen3_moe_routed_down_q6k_gfx11_wave64 {subgroup_size = 64} + +target.decl @qwen3_moe_attention_prepare_gfx11_wave32 + +config.decl @qwen3_moe.routed_down.input_size : %value: index where [range(%value, 256, 32768), mul(%value, 256)] + +config.decl @qwen3_moe.routed_down.route_count : %value: index where [range(%value, 1, 8)] + +config.decl @qwen3_moe.routed_down.expert_count : %value: index where [range(%value, 1, 512)] + +config.decl @qwen3_moe.routed_down.output_size : %value: index where [range(%value, 1, 4096)] + +config.decl @qwen3_moe.workload.token_capacity : %value: index where [range(%value, 1, 2048)] + +kernel.decl @ggml_quantize_q8_1_x4_f32(%token_count: index, %input_size: index) launch(%token_count: index, %input_size: index, %input: buffer, %output: buffer) + +func.decl @ggml_q6k_q8_1_x4_row_lane(%input_size: index, %weight: buffer, %weight_row_byte_base: offset, %q8_input: buffer, %q8_row_byte_base: offset, %lane: index) -> (f32) + +func.decl @ggml_q6k_stage_f32_scales(%input_size: index, %row: index, %block: index, %frame_count: index, %frame: index, %lane: index, %weight: buffer, %scale_stage: buffer) + +func.decl @ggml_q6k_load_f32_block(%token_count: index, %input_size: index, %token: index, %block: index, %lane: index, %input: buffer) -> (vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>) + +func.decl @ggml_q6k_f32_block_row(%input_size: index, %row: index, %block: index, %frame_count: index, %frame: index, %lane: index, %weight: buffer, %scale_stage: buffer, %input0: vector<4xf32>, %input1: vector<4xf32>, %input2: vector<4xf32>, %input3: vector<4xf32>) -> (f32) + +kernel.decl @qwen3_moe_attention_rmsnorm_quantize_q8_1_x4(%token_count: index) launch(%token_count: index, %input: buffer, %weight: buffer, %q8_output: buffer) + +// Shared Q6_K contraction and residual publication used by ordinary and +// completion-fused exports with identical launch geometry. +func.def inline @qwen3_moe_routed_down_q6k_q8_1_x4_body(%publish_output: i1, %token_count: index, %token: index, %input_size: index, %route_count: index, %route_id_stride: index, %expert_count: index, %output_size: index, %q8_input: buffer, %route_ids: buffer, %route_weights: buffer, %weight: buffer, %output: buffer) { + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 2048)] : index + %bounded_token, %body_token_count = index.assume %token, %bounded_token_count [lt(%token, %bounded_token_count)] : index, index + %bounded_input_size = index.assume %input_size [range(%input_size, 256, 32768), mul(%input_size, 256)] : index + %bounded_route_count = index.assume %route_count [range(%route_count, 1, 8)] : index + %bounded_route_id_stride = index.assume %route_id_stride [range(%route_id_stride, 1, 512)] : index + %bounded_expert_count = index.assume %expert_count [range(%expert_count, 1, 512)] : index + %bounded_output_size = index.assume %output_size [range(%output_size, 1, 262144)] : index + %channel_tile = kernel.workgroup.id : index + %subgroup0 = kernel.subgroup.id : index + %lane = kernel.subgroup.lane.id : index + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c8 = index.constant 8 : index + %c128 = index.constant 128 : index + %c144_bytes = index.constant 144 : offset + %c210_bytes = index.constant 210 : offset + %c256 = index.constant 256 : index + %c0_i32 = scalar.constant 0 : i32 + %c0_f32 = scalar.constant 0.0 : f32 + %c0_offset = index.constant 0 : offset + %subgroup = index.assume %subgroup0 [range(%subgroup0, 0, 7)] : index + %channel_base = index.mul %channel_tile, %c8 : index + %channel = index.add %channel_base, %subgroup : index + %valid_channel = index.cmp ult, %channel, %bounded_output_size : index + %lane_i32 = index.cast %lane : index to i32 + %is_lane_zero = scalar.cmpi eq, %lane_i32, %c0_i32 : i32 + %q6_block_count = index.div %bounded_input_size, %c256 : index + %weight_row_bytes = index.scale %q6_block_count, %c210_bytes : index, offset -> offset + %q8_group_count = index.div %bounded_input_size, %c128 : index + %q8_row_bytes = index.scale %q8_group_count, %c144_bytes : index, offset -> offset + %q8_noalias, %route_id_noalias, %route_weight_noalias, %weight_noalias, %output_noalias = buffer.assume.noalias %q8_input, %route_ids, %route_weights, %weight, %output : buffer, buffer, buffer, buffer, buffer + %route_id_view = buffer.view %route_id_noalias[%c0_offset] : buffer -> view<[%body_token_count]x[%bounded_route_id_stride]xi32> + %route_weight_view = buffer.view %route_weight_noalias[%c0_offset] : buffer -> view<[%body_token_count]x[%bounded_route_count]xf32> + %output_view = buffer.view %output_noalias[%c0_offset] : buffer -> view<[%body_token_count]x[%bounded_output_size]xf32> + %lane_has_route = index.cmp ult, %lane, %bounded_route_count : index + %lane_within_route_stride = index.cmp ult, %lane, %bounded_route_id_stride : index + %loads_active_route = scalar.andi %publish_output, %lane_has_route : i1 + %loads_route_metadata = scalar.andi %loads_active_route, %lane_within_route_stride : i1 + %lane_expert_i32, %lane_route_weight = scf.if %loads_route_metadata -> (i32, f32) { + %route_lane0, %metadata_route_count = index.assume %lane, %bounded_route_count [lt(%lane, %bounded_route_count)] : index, index + %route_lane, %metadata_route_id_stride = index.assume %route_lane0, %bounded_route_id_stride [lt(%route_lane0, %bounded_route_id_stride)] : index, index + %expert_i32 = view.load %route_id_view[%bounded_token, %route_lane] : view<[%body_token_count]x[%bounded_route_id_stride]xi32> -> i32 + %route_weight = view.load %route_weight_view[%bounded_token, %route_lane0] : view<[%body_token_count]x[%bounded_route_count]xf32> -> f32 + scf.yield %expert_i32, %route_weight : i32, f32 + } else { + scf.yield %c0_i32, %c0_f32 : i32, f32 + } + %computes_channel = scalar.andi %publish_output, %valid_channel : i1 + %routed_lane_sum = scf.if %computes_channel -> (f32) { + %sum = scf.for %route = [%c0 to %bounded_route_count step %c1](%route_acc = %c0_f32 : f32) -> (f32) unroll { + %route_i32 = index.cast %route : index to i32 + %expert_i32 = kernel.subgroup.broadcast %lane_expert_i32 from %route_i32 : i32, i32 + %expert0 = index.cast %expert_i32 : i32 to index + %expert1 = index.assume %expert0 [range(%expert0, 0, 511)] : index + %expert, %weight_expert_count = index.assume %expert1, %bounded_expert_count [lt(%expert1, %bounded_expert_count)] : index, index + %expert_output_base = index.mul %expert, %bounded_output_size : index + %expert_channel = index.add %expert_output_base, %channel : index + %weight_row_byte_base = index.scale %expert_channel, %weight_row_bytes : index, offset -> offset + %q8_row_base0 = index.mul %bounded_token, %bounded_route_count : index + %q8_row = index.add %q8_row_base0, %route : index + %q8_row_byte_base = index.scale %q8_row, %q8_row_bytes : index, offset -> offset + %lane_acc = func.call @ggml_q6k_q8_1_x4_row_lane(%bounded_input_size, %weight_noalias, %weight_row_byte_base, %q8_noalias, %q8_row_byte_base, %lane) : (index, buffer, offset, buffer, offset, index) -> (f32) + %route_weight = kernel.subgroup.broadcast %lane_route_weight from %route_i32 : f32, i32 + %weighted_lane = scalar.mulf %lane_acc, %route_weight : f32 + %next = scalar.addf %route_acc, %weighted_lane : f32 + scf.yield %next : f32 + } + scf.yield %sum : f32 + } else { + scf.yield %c0_f32 : f32 + } + %routed_sum = kernel.subgroup.reduce %routed_lane_sum : f32 + %writes_active_channel = scalar.andi %publish_output, %valid_channel : i1 + %writes_output = scalar.andi %writes_active_channel, %is_lane_zero : i1 + scf.if %writes_output { + %bounded_channel, %body_output_size = index.assume %channel, %bounded_output_size [lt(%channel, %bounded_output_size)] : index, index + %residual = view.load %output_view[%bounded_token, %bounded_channel] : view<[%body_token_count]x[%bounded_output_size]xf32> -> f32 + %result = scalar.addf %residual, %routed_sum : f32 + view.store %result, %output_view[%bounded_token, %bounded_channel] : f32, view<[%body_token_count]x[%bounded_output_size]xf32> + } + func.return +} + +kernel.def @qwen3_moe_routed_down_q6k_q8_1_x4(%token_count: index, %input_size: index, %route_count: index, %route_id_stride: index, %expert_count: index, %output_size: index) { + %token_capacity = config.get @qwen3_moe.workload.token_capacity : index + %configured_output_size = config.get @qwen3_moe.routed_down.output_size : index + %c8 = index.constant 8 : index + %c7 = index.constant 7 : index + %c1 = index.constant 1 : index + %workgroup_size = index.constant 256 : index + %padded_output_size = index.add %configured_output_size, %c7 : index + %output_tiles = index.div %padded_output_size, %c8 : index + kernel.launch.config workgroups(%output_tiles, %token_capacity, %c1) workgroup_size(%workgroup_size, %c1, %c1) : index +} launch(%token_count: index, %input_size: index, %route_count: index, %route_id_stride: index, %expert_count: index, %output_size: index, %q8_input: buffer, %route_ids: buffer, %route_weights: buffer, %weight: buffer, %output: buffer) { + %token_capacity = config.get @qwen3_moe.workload.token_capacity : index + %configured_input_size0 = config.get @qwen3_moe.routed_down.input_size : index + %configured_route_count0 = config.get @qwen3_moe.routed_down.route_count : index + %configured_expert_count0 = config.get @qwen3_moe.routed_down.expert_count : index + %configured_output_size0 = config.get @qwen3_moe.routed_down.output_size : index + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 2048), le(%token_count, %token_capacity)] : index + %bounded_input_size, %configured_input_size = index.assume %input_size, %configured_input_size0 [range(%input_size, 256, 32768), mul(%input_size, 256), eq(%input_size, %configured_input_size0)] : index, index + %bounded_route_count, %configured_route_count = index.assume %route_count, %configured_route_count0 [range(%route_count, 1, 8), eq(%route_count, %configured_route_count0)] : index, index + %bounded_route_id_stride = index.assume %route_id_stride [range(%route_id_stride, 1, 512)] : index + %bounded_expert_count, %configured_expert_count = index.assume %expert_count, %configured_expert_count0 [range(%expert_count, 1, 512), eq(%expert_count, %configured_expert_count0)] : index, index + %bounded_output_size, %configured_output_size = index.assume %output_size, %configured_output_size0 [range(%output_size, 1, 4096), eq(%output_size, %configured_output_size0)] : index, index + %token0 = kernel.workgroup.id : index + %c0 = index.constant 0 : index + %valid_token = index.cmp ult, %token0, %bounded_token_count : index + %safe_token0 = scf.select %valid_token, %token0, %c0 : index + %safe_token, %body_token_count = index.assume %safe_token0, %bounded_token_count [lt(%safe_token0, %bounded_token_count)] : index, index + func.call @qwen3_moe_routed_down_q6k_q8_1_x4_body(%valid_token, %body_token_count, %safe_token, %configured_input_size, %configured_route_count, %bounded_route_id_stride, %configured_expert_count, %configured_output_size, %q8_input, %route_ids, %route_weights, %weight, %output) : (i1, index, index, index, index, index, index, index, buffer, buffer, buffer, buffer, buffer) + kernel.return +} + +// Decode-only route that also publishes the normalized Q8_1 x4 row consumed by +// the next projection boundary. +kernel.def target(@qwen3_moe_attention_prepare_gfx11_wave32) @qwen3_moe_routed_down_q6k_q8_1_x4_next_q8(%token_count: index, %input_size: index, %route_count: index, %route_id_stride: index, %expert_count: index, %output_size: index) { + %configured_output_size = config.get @qwen3_moe.routed_down.output_size : index + %c8 = index.constant 8 : index + %c7 = index.constant 7 : index + %c1 = index.constant 1 : index + %workgroup_size = index.constant 256 : index + %padded_output_size = index.add %configured_output_size, %c7 : index + %output_tiles = index.div %padded_output_size, %c8 : index + kernel.launch.config workgroups(%output_tiles, %c1, %c1) workgroup_size(%workgroup_size, %c1, %c1) : index +} launch(%token_count: index, %input_size: index, %route_count: index, %route_id_stride: index, %expert_count: index, %output_size: index, %q8_input: buffer, %route_ids: buffer, %route_weights: buffer, %weight: buffer, %output: buffer, %norm_weight: buffer, %completion_counter: buffer, %next_q8_output: buffer) { + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 1)] : index + %configured_input_size0 = config.get @qwen3_moe.routed_down.input_size : index + %configured_route_count0 = config.get @qwen3_moe.routed_down.route_count : index + %configured_expert_count0 = config.get @qwen3_moe.routed_down.expert_count : index + %configured_output_size0 = config.get @qwen3_moe.routed_down.output_size : index + %bounded_input_size, %configured_input_size = index.assume %input_size, %configured_input_size0 [range(%input_size, 256, 32768), mul(%input_size, 256), eq(%input_size, %configured_input_size0)] : index, index + %bounded_route_count, %configured_route_count = index.assume %route_count, %configured_route_count0 [range(%route_count, 1, 8), eq(%route_count, %configured_route_count0)] : index, index + %bounded_route_id_stride = index.assume %route_id_stride [range(%route_id_stride, 1, 512)] : index + %bounded_expert_count, %configured_expert_count = index.assume %expert_count, %configured_expert_count0 [range(%expert_count, 1, 512), eq(%expert_count, %configured_expert_count0)] : index, index + %bounded_output_size, %configured_output_size = index.assume %output_size, %configured_output_size0 [range(%output_size, 128, 4096), mul(%output_size, 128), eq(%output_size, %configured_output_size0)] : index, index + %c8 = index.constant 8 : index + %publishes_output = scalar.constant true : i1 + %body_token = index.constant 0 : index + func.call @qwen3_moe_routed_down_q6k_q8_1_x4_body(%publishes_output, %bounded_token_count, %body_token, %configured_input_size, %configured_route_count, %bounded_route_id_stride, %configured_expert_count, %configured_output_size, %q8_input, %route_ids, %route_weights, %weight, %output) : (i1, index, index, index, index, index, index, index, buffer, buffer, buffer, buffer, buffer) + func.call @qwen3_moe_routed_down_next_q8_completion(%c8, %c8, %bounded_token_count, %configured_output_size, %output, %norm_weight, %completion_counter, %next_q8_output) : (index, index, index, index, buffer, buffer, buffer, buffer) + kernel.return +} + +// The production hidden width requires all 256 output tiles to arrive before +// normalization. Compare both residual and packed Q8 publication with the +// ordinary two-dispatch composition, then reuse the same completion word. +check.case public @qwen3_moe_routed_down_q6k_q8_1_x4_next_q8_differential_case { + %token_count = check.literal value(1) : index + %input_size = check.literal value(768) : index + %route_count = check.literal value(2) : index + %route_id_stride = check.literal value(4) : index + %expert_count = check.literal value(2) : index + %output_size = check.literal value(2048) : index + %routed_row_count = check.literal value(2) : index + %routed_input = check.generate.fill value(0.00390625) : tensor<2x768xf32> + %q8_input = check.generate.fill value(0) : tensor<2x864xi8> + %route_ids = check.generate.iota offset(0) step(1) period(2) : tensor<1x4xi32> + %route_weights = check.generate.fill value(0.5) : tensor<1x2xf32> + %weight = check.generate.fill value(-86) : tensor<2x2048x3x210xi8> + %norm_weight = check.generate.iota offset(-1.0) step(0.0009765625) : tensor<2048xf32> + %expected_output = check.generate.iota offset(-0.5) step(0.00048828125) period(2048) : tensor<1x2048xf32> + %actual_output0 = check.generate.iota offset(-0.5) step(0.00048828125) period(2048) : tensor<1x2048xf32> + %actual_output1 = check.generate.iota offset(-0.5) step(0.00048828125) period(2048) : tensor<1x2048xf32> + %expected_q8 = check.generate.fill value(0) : tensor<2304xi8> + %actual_q8_0 = check.generate.fill value(1) : tensor<2304xi8> + %actual_q8_1 = check.generate.fill value(1) : tensor<2304xi8> + %completion_counter = check.generate.fill value(0) : tensor<1xi32> + %expected_counter = check.generate.fill value(0) : tensor<1xi32> + kernel.launch @ggml_quantize_q8_1_x4_f32[%routed_row_count, %input_size](%routed_row_count, %input_size, %routed_input, %q8_input) : [index, index](index, index, tensor<2x768xf32>, tensor<2x864xi8>) + kernel.launch @qwen3_moe_routed_down_q6k_q8_1_x4[%token_count, %input_size, %route_count, %route_id_stride, %expert_count, %output_size](%token_count, %input_size, %route_count, %route_id_stride, %expert_count, %output_size, %q8_input, %route_ids, %route_weights, %weight, %expected_output) : [index, index, index, index, index, index](index, index, index, index, index, index, tensor<2x864xi8>, tensor<1x4xi32>, tensor<1x2xf32>, tensor<2x2048x3x210xi8>, tensor<1x2048xf32>) + kernel.launch @qwen3_moe_attention_rmsnorm_quantize_q8_1_x4[%token_count](%token_count, %expected_output, %norm_weight, %expected_q8) : [index](index, tensor<1x2048xf32>, tensor<2048xf32>, tensor<2304xi8>) + kernel.launch @qwen3_moe_routed_down_q6k_q8_1_x4_next_q8[%token_count, %input_size, %route_count, %route_id_stride, %expert_count, %output_size](%token_count, %input_size, %route_count, %route_id_stride, %expert_count, %output_size, %q8_input, %route_ids, %route_weights, %weight, %actual_output0, %norm_weight, %completion_counter, %actual_q8_0) : [index, index, index, index, index, index](index, index, index, index, index, index, tensor<2x864xi8>, tensor<1x4xi32>, tensor<1x2xf32>, tensor<2x2048x3x210xi8>, tensor<1x2048xf32>, tensor<2048xf32>, tensor<1xi32>, tensor<2304xi8>) + kernel.launch @qwen3_moe_routed_down_q6k_q8_1_x4_next_q8[%token_count, %input_size, %route_count, %route_id_stride, %expert_count, %output_size](%token_count, %input_size, %route_count, %route_id_stride, %expert_count, %output_size, %q8_input, %route_ids, %route_weights, %weight, %actual_output1, %norm_weight, %completion_counter, %actual_q8_1) : [index, index, index, index, index, index](index, index, index, index, index, index, tensor<2x864xi8>, tensor<1x4xi32>, tensor<1x2xf32>, tensor<2x2048x3x210xi8>, tensor<1x2048xf32>, tensor<2048xf32>, tensor<1xi32>, tensor<2304xi8>) + check.expect.close actual(%actual_output0) expected(%expected_output) atol(0.0) rtol(0.0) nan(same) : tensor<1x2048xf32> + check.expect.close actual(%actual_output1) expected(%expected_output) atol(0.0) rtol(0.0) nan(same) : tensor<1x2048xf32> + check.expect.equal actual(%actual_q8_0) expected(%expected_q8) : tensor<2304xi8> + check.expect.equal actual(%actual_q8_1) expected(%expected_q8) : tensor<2304xi8> + check.expect.equal actual(%completion_counter) expected(%expected_counter) : tensor<1xi32> + check.return +} + +check.case public @qwen3_moe_routed_down_q6k_q8_1_x4_next_q8_benchmark_case { + %token_count = check.literal value(1) : index + %input_size = check.literal value(768) : index + %route_count = check.literal value(8) : index + %route_id_stride = check.literal value(8) : index + %expert_count = check.literal value(128) : index + %output_size = check.literal value(2048) : index + %q8_input = check.generate.fill value(0) : tensor<1x8x864xi8> + %route_ids = check.generate.iota offset(0) step(1) period(128) : tensor<1x8xi32> + %route_weights = check.generate.fill value(0.125) : tensor<1x8xf32> + %weight = check.generate.fill value(0) : tensor<128x2048x3x210xi8> + %output = check.generate.fill value(1.0) : tensor<1x2048xf32> + %norm_weight = check.generate.fill value(1.0) : tensor<2048xf32> + %completion_counter = check.generate.fill value(0) : tensor<1xi32> + %next_q8_output = check.generate.fill value(0) : tensor<2304xi8> + kernel.launch @qwen3_moe_routed_down_q6k_q8_1_x4_next_q8[%token_count, %input_size, %route_count, %route_id_stride, %expert_count, %output_size](%token_count, %input_size, %route_count, %route_id_stride, %expert_count, %output_size, %q8_input, %route_ids, %route_weights, %weight, %output, %norm_weight, %completion_counter, %next_q8_output) : [index, index, index, index, index, index](index, index, index, index, index, index, tensor<1x8x864xi8>, tensor<1x8xi32>, tensor<1x8xf32>, tensor<128x2048x3x210xi8>, tensor<1x2048xf32>, tensor<2048xf32>, tensor<1xi32>, tensor<2304xi8>) + check.return +} + +check.case public @qwen3_moe_routed_down_q6k_q8_1_x4_next_q8_composed_benchmark_case { + %token_count = check.literal value(1) : index + %input_size = check.literal value(768) : index + %route_count = check.literal value(8) : index + %route_id_stride = check.literal value(8) : index + %expert_count = check.literal value(128) : index + %output_size = check.literal value(2048) : index + %q8_input = check.generate.fill value(0) : tensor<1x8x864xi8> + %route_ids = check.generate.iota offset(0) step(1) period(128) : tensor<1x8xi32> + %route_weights = check.generate.fill value(0.125) : tensor<1x8xf32> + %weight = check.generate.fill value(0) : tensor<128x2048x3x210xi8> + %output = check.generate.fill value(1.0) : tensor<1x2048xf32> + %norm_weight = check.generate.fill value(1.0) : tensor<2048xf32> + %next_q8_output = check.generate.fill value(0) : tensor<2304xi8> + kernel.launch @qwen3_moe_routed_down_q6k_q8_1_x4[%token_count, %input_size, %route_count, %route_id_stride, %expert_count, %output_size](%token_count, %input_size, %route_count, %route_id_stride, %expert_count, %output_size, %q8_input, %route_ids, %route_weights, %weight, %output) : [index, index, index, index, index, index](index, index, index, index, index, index, tensor<1x8x864xi8>, tensor<1x8xi32>, tensor<1x8xf32>, tensor<128x2048x3x210xi8>, tensor<1x2048xf32>) + kernel.launch @qwen3_moe_attention_rmsnorm_quantize_q8_1_x4[%token_count](%token_count, %output, %norm_weight, %next_q8_output) : [index](index, tensor<1x2048xf32>, tensor<2048xf32>, tensor<2304xi8>) + check.return +} + +// Shared direct-F32 decode schedule. Every subgroup owns one scale frame and +// one output channel. Providers specialize how many Q6_K blocks a subgroup +// covers per pass from the target subgroup width. +func.def inline @qwen3_moe_routed_down_q6k_f32_body(%publish_output: i1, %subgroup_count: index, %block_step: index, %scale_stage_bytes: offset, %token_count: index, %token: index, %input_size: index, %route_count: index, %route_id_stride: index, %expert_count: index, %output_size: index, %input: buffer, %route_ids: buffer, %route_weights: buffer, %weight: buffer, %output: buffer) { + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 2048)] : index + %bounded_token, %body_token_count = index.assume %token, %bounded_token_count [lt(%token, %bounded_token_count)] : index, index + %bounded_input_size = index.assume %input_size [range(%input_size, 256, 32768), mul(%input_size, 256)] : index + %bounded_route_count = index.assume %route_count [range(%route_count, 1, 8)] : index + %bounded_route_id_stride = index.assume %route_id_stride [range(%route_id_stride, 1, 512)] : index + %bounded_expert_count = index.assume %expert_count [range(%expert_count, 1, 512)] : index + %bounded_output_size = index.assume %output_size [range(%output_size, 1, 262144)] : index + %channel_tile = kernel.workgroup.id : index + %subgroup0 = kernel.subgroup.id : index + %lane0 = kernel.subgroup.lane.id : index + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c16 = index.constant 16 : index + %c210 = index.constant 210 : index + %c256 = index.constant 256 : index + %c0_i32 = scalar.constant 0 : i32 + %c0_f32 = scalar.constant 0.0 : f32 + %c0_offset = index.constant 0 : offset + %subgroup = index.assume %subgroup0 [range(%subgroup0, 0, 7)] : index + %lane = index.assume %lane0 [range(%lane0, 0, 63)] : index + %channel_tile_base = index.mul %channel_tile, %subgroup_count : index + %channel = index.add %channel_tile_base, %subgroup : index + %valid_channel = index.cmp ult, %channel, %bounded_output_size : index + %safe_channel = scf.select %valid_channel, %channel, %c0 : index + %lane_i32 = index.cast %lane : index to i32 + %is_lane_zero = scalar.cmpi eq, %lane_i32, %c0_i32 : i32 + %assignment_count = index.mul %body_token_count, %bounded_route_count : index + %q6_block_count = index.div %bounded_input_size, %c256 : index + %weight_row_byte_count = index.mul %q6_block_count, %c210 : index + %cohort = index.div %lane, %c16 : index + %input_noalias, %route_id_noalias, %route_weight_noalias, %weight_noalias, %output_noalias = buffer.assume.noalias %input, %route_ids, %route_weights, %weight, %output : buffer, buffer, buffer, buffer, buffer + %route_id_view = buffer.view %route_id_noalias[%c0_offset] : buffer -> view<[%body_token_count]x[%bounded_route_id_stride]xi32> + %route_weight_view = buffer.view %route_weight_noalias[%c0_offset] : buffer -> view<[%body_token_count]x[%bounded_route_count]xf32> + %output_view = buffer.view %output_noalias[%c0_offset] : buffer -> view<[%body_token_count]x[%bounded_output_size]xf32> + %scale_stage = buffer.alloca align(16) %scale_stage_bytes : buffer + %routed_sum0 = scf.if %publish_output -> (f32) { + %active_sum = scf.for %route = [%c0 to %bounded_route_count step %c1](%route_acc = %c0_f32 : f32) -> (f32) unroll { + // Route metadata is uniform across all channels. Keeping the access in + // the unrolled route body exposes that uniformity directly to target + // lowering instead of materializing a dynamic subgroup broadcast. + %bounded_route, %body_route_count = index.assume %route, %bounded_route_count [lt(%route, %bounded_route_count)] : index, index + %route_for_stride, %body_route_id_stride = index.assume %bounded_route, %bounded_route_id_stride [lt(%bounded_route, %bounded_route_id_stride)] : index, index + %expert_i32 = view.load %route_id_view[%bounded_token, %route_for_stride] : view<[%body_token_count]x[%bounded_route_id_stride]xi32> -> i32 + %route_weight = view.load %route_weight_view[%bounded_token, %bounded_route] : view<[%body_token_count]x[%bounded_route_count]xf32> -> f32 + %expert0 = index.cast %expert_i32 : i32 to index + %expert1 = index.assume %expert0 [range(%expert0, 0, 511)] : index + %expert, %weight_expert_count = index.assume %expert1, %bounded_expert_count [lt(%expert1, %bounded_expert_count)] : index, index + %expert_output_base = index.mul %expert, %bounded_output_size : index + %expert_channel = index.add %expert_output_base, %safe_channel : index + %input_row_base = index.mul %bounded_token, %bounded_route_count : index + %input_row = index.add %input_row_base, %bounded_route : index + %lane_sum = scf.for %block_base = [%c0 to %q6_block_count step %block_step](%block_acc = %c0_f32 : f32) -> (f32) { + %block = index.add %block_base, %cohort : index + func.call @ggml_q6k_stage_f32_scales(%bounded_input_size, %expert_channel, %block, %subgroup_count, %subgroup, %lane, %weight_noalias, %scale_stage) : (index, index, index, index, index, index, buffer, buffer) + %input0, %input1, %input2, %input3 = func.call @ggml_q6k_load_f32_block(%assignment_count, %bounded_input_size, %input_row, %block, %lane, %input_noalias) : (index, index, index, index, index, buffer) -> (vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>) + %contribution = func.call @ggml_q6k_f32_block_row(%bounded_input_size, %expert_channel, %block, %subgroup_count, %subgroup, %lane, %weight_noalias, %scale_stage, %input0, %input1, %input2, %input3) : (index, index, index, index, index, index, buffer, buffer, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>) -> (f32) + %next = scalar.addf %block_acc, %contribution : f32 + scf.yield %next : f32 + } + %route_sum = kernel.subgroup.reduce %lane_sum : f32 + %weighted = scalar.mulf %route_sum, %route_weight : f32 + %next = scalar.addf %route_acc, %weighted : f32 + scf.yield %next : f32 + } + scf.yield %active_sum : f32 + } else { + scf.yield %c0_f32 : f32 + } + %routed_sum = scf.select %valid_channel, %routed_sum0, %c0_f32 : f32 + %writes_active_channel = scalar.andi %publish_output, %valid_channel : i1 + %writes_output = scalar.andi %writes_active_channel, %is_lane_zero : i1 + scf.if %writes_output { + %bounded_channel, %body_output_size = index.assume %channel, %bounded_output_size [lt(%channel, %bounded_output_size)] : index, index + %residual = view.load %output_view[%bounded_token, %bounded_channel] : view<[%body_token_count]x[%bounded_output_size]xf32> -> f32 + %result = scalar.addf %residual, %routed_sum : f32 + view.store %result, %output_view[%bounded_token, %bounded_channel] : f32, view<[%body_token_count]x[%bounded_output_size]xf32> + } + func.return +} + +// Decode schedule matching the Vulkan oracle's two adjacent output rows per +// wave. Four wave64 subgroups publish eight channels per workgroup. Keeping the +// route loop rolled prevents its eight iterations from multiplying the two-row +// contraction's register pressure and instruction footprint. +func.def inline @qwen3_moe_routed_down_q6k_f32_pair_body(%token_count: index, %token: index, %input_size: index, %route_count: index, %route_id_stride: index, %expert_count: index, %output_size: index, %input: buffer, %route_ids: buffer, %route_weights: buffer, %weight: buffer, %output: buffer) { + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 1)] : index + %bounded_token, %body_token_count = index.assume %token, %bounded_token_count [lt(%token, %bounded_token_count)] : index, index + %bounded_input_size = index.assume %input_size [range(%input_size, 256, 32768), mul(%input_size, 256)] : index + %bounded_route_count = index.assume %route_count [range(%route_count, 1, 8)] : index + %bounded_route_id_stride = index.assume %route_id_stride [range(%route_id_stride, 1, 512)] : index + %bounded_expert_count = index.assume %expert_count [range(%expert_count, 1, 512)] : index + %bounded_output_size = index.assume %output_size [range(%output_size, 2, 4096)] : index + %pair_tile = kernel.workgroup.id : index + %subgroup0 = kernel.subgroup.id : index + %lane0 = kernel.subgroup.lane.id : index + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c4 = index.constant 4 : index + %c8 = index.constant 8 : index + %c16 = index.constant 16 : index + %c256 = index.constant 256 : index + %c0_f32 = scalar.constant 0.0 : f32 + %c0_offset = index.constant 0 : offset + %scale_stage_bytes = index.constant 2048 : offset + %subgroup = index.assume %subgroup0 [range(%subgroup0, 0, 3)] : index + %lane = index.assume %lane0 [range(%lane0, 0, 63)] : index + %pair_tile_base = index.mul %pair_tile, %c4 : index + %pair = index.add %pair_tile_base, %subgroup : index + %row00 = index.mul %pair, %c2 : index + %row0, %body_output_size = index.assume %row00, %bounded_output_size [lt(%row00, %bounded_output_size)] : index, index + %row1 = index.add %row0, %c1 : index + %row1_valid = index.cmp ult, %row1, %body_output_size : index + %safe_row1 = scf.select %row1_valid, %row1, %c0 : index + %assignment_count = index.mul %body_token_count, %bounded_route_count : index + %q6_block_count = index.div %bounded_input_size, %c256 : index + %cohort = index.div %lane, %c16 : index + %scale_frame0 = index.mul %subgroup, %c2 : index + %scale_frame1 = index.add %scale_frame0, %c1 : index + %input_noalias, %route_id_noalias, %route_weight_noalias, %weight_noalias, %output_noalias = buffer.assume.noalias %input, %route_ids, %route_weights, %weight, %output : buffer, buffer, buffer, buffer, buffer + %route_id_view = buffer.view %route_id_noalias[%c0_offset] : buffer -> view<[%body_token_count]x[%bounded_route_id_stride]xi32> + %route_weight_view = buffer.view %route_weight_noalias[%c0_offset] : buffer -> view<[%body_token_count]x[%bounded_route_count]xf32> + %output_view = buffer.view %output_noalias[%c0_offset] : buffer -> view<[%body_token_count]x[%body_output_size]xf32> + %scale_stage = buffer.alloca align(16) %scale_stage_bytes : buffer + %routed_sum0, %routed_sum1 = scf.for %route = [%c0 to %bounded_route_count step %c1](%route_acc0 = %c0_f32 : f32, %route_acc1 = %c0_f32 : f32) -> (f32, f32) { + %bounded_route, %body_route_count = index.assume %route, %bounded_route_count [lt(%route, %bounded_route_count)] : index, index + %route_for_stride, %body_route_id_stride = index.assume %bounded_route, %bounded_route_id_stride [lt(%bounded_route, %bounded_route_id_stride)] : index, index + %expert_i32 = view.load %route_id_view[%bounded_token, %route_for_stride] : view<[%body_token_count]x[%bounded_route_id_stride]xi32> -> i32 + %route_weight = view.load %route_weight_view[%bounded_token, %bounded_route] : view<[%body_token_count]x[%bounded_route_count]xf32> -> f32 + %expert0 = index.cast %expert_i32 : i32 to index + %expert1 = index.assume %expert0 [range(%expert0, 0, 511)] : index + %expert, %body_expert_count = index.assume %expert1, %bounded_expert_count [lt(%expert1, %bounded_expert_count)] : index, index + %expert_output_base = index.mul %expert, %body_output_size : index + %expert_row0 = index.add %expert_output_base, %row0 : index + %expert_row1 = index.add %expert_output_base, %safe_row1 : index + %input_row_base = index.mul %bounded_token, %body_route_count : index + %input_row = index.add %input_row_base, %bounded_route : index + %lane_sum0, %lane_sum1 = scf.for %block_base = [%c0 to %q6_block_count step %c4](%block_acc0 = %c0_f32 : f32, %block_acc1 = %c0_f32 : f32) -> (f32, f32) { + %block = index.add %block_base, %cohort : index + func.call @ggml_q6k_stage_f32_scales(%bounded_input_size, %expert_row0, %block, %c8, %scale_frame0, %lane, %weight_noalias, %scale_stage) : (index, index, index, index, index, index, buffer, buffer) + %input00, %input01, %input02, %input03 = func.call @ggml_q6k_load_f32_block(%assignment_count, %bounded_input_size, %input_row, %block, %lane, %input_noalias) : (index, index, index, index, index, buffer) -> (vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>) + %contribution0 = func.call @ggml_q6k_f32_block_row(%bounded_input_size, %expert_row0, %block, %c8, %scale_frame0, %lane, %weight_noalias, %scale_stage, %input00, %input01, %input02, %input03) : (index, index, index, index, index, index, buffer, buffer, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>) -> (f32) + %contribution1 = scf.if %row1_valid -> (f32) { + func.call @ggml_q6k_stage_f32_scales(%bounded_input_size, %expert_row1, %block, %c8, %scale_frame1, %lane, %weight_noalias, %scale_stage) : (index, index, index, index, index, index, buffer, buffer) + // Reloading the shared activations lets the first row's vectors die + // before the second scale stage. Carrying them across that stage + // increases register pressure and is slower on gfx1151. + %input10, %input11, %input12, %input13 = func.call @ggml_q6k_load_f32_block(%assignment_count, %bounded_input_size, %input_row, %block, %lane, %input_noalias) : (index, index, index, index, index, buffer) -> (vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>) + %row1_contribution = func.call @ggml_q6k_f32_block_row(%bounded_input_size, %expert_row1, %block, %c8, %scale_frame1, %lane, %weight_noalias, %scale_stage, %input10, %input11, %input12, %input13) : (index, index, index, index, index, index, buffer, buffer, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>) -> (f32) + scf.yield %row1_contribution : f32 + } else { + scf.yield %c0_f32 : f32 + } + %next0 = scalar.addf %block_acc0, %contribution0 : f32 + %next1 = scalar.addf %block_acc1, %contribution1 : f32 + scf.yield %next0, %next1 : f32, f32 + } + %route_sum0 = kernel.subgroup.reduce %lane_sum0 : f32 + %route_sum1 = kernel.subgroup.reduce %lane_sum1 : f32 + %weighted0 = scalar.mulf %route_sum0, %route_weight : f32 + %weighted1 = scalar.mulf %route_sum1, %route_weight : f32 + %next0 = scalar.addf %route_acc0, %weighted0 : f32 + %next1 = scalar.addf %route_acc1, %weighted1 : f32 + scf.yield %next0, %next1 : f32, f32 + } + %is_lane_zero = index.cmp eq, %lane, %c0 : index + scf.if %is_lane_zero { + %residual0 = view.load %output_view[%bounded_token, %row0] : view<[%body_token_count]x[%body_output_size]xf32> -> f32 + %result0 = scalar.addf %residual0, %routed_sum0 : f32 + view.store %result0, %output_view[%bounded_token, %row0] : f32, view<[%body_token_count]x[%body_output_size]xf32> + scf.if %row1_valid { + %residual1 = view.load %output_view[%bounded_token, %row1] : view<[%body_token_count]x[%body_output_size]xf32> -> f32 + %result1 = scalar.addf %residual1, %routed_sum1 : f32 + view.store %result1, %output_view[%bounded_token, %row1] : f32, view<[%body_token_count]x[%body_output_size]xf32> + } + } + func.return +} + +// Wave64 matches llama.cpp's Vulkan subgroup width and contracts four Q6_K +// blocks per subgroup pass. +kernel.def target(@qwen3_moe_routed_down_q6k_gfx11_wave64) @qwen3_moe_routed_down_q6k_f32_wave64(%token_count: index, %input_size: index, %route_count: index, %route_id_stride: index, %expert_count: index, %output_size: index) { + %token_capacity = config.get @qwen3_moe.workload.token_capacity : index + %configured_output_size = config.get @qwen3_moe.routed_down.output_size : index + %c1 = index.constant 1 : index + %c3 = index.constant 3 : index + %c4 = index.constant 4 : index + %c256 = index.constant 256 : index + %padded_output_size = index.add %configured_output_size, %c3 : index + %output_tiles = index.div %padded_output_size, %c4 : index + kernel.launch.config workgroups(%output_tiles, %token_capacity, %c1) workgroup_size(%c256, %c1, %c1) : index +} launch(%token_count: index, %input_size: index, %route_count: index, %route_id_stride: index, %expert_count: index, %output_size: index, %input: buffer, %route_ids: buffer, %route_weights: buffer, %weight: buffer, %output: buffer) { + %token_capacity = config.get @qwen3_moe.workload.token_capacity : index + %configured_input_size0 = config.get @qwen3_moe.routed_down.input_size : index + %configured_route_count0 = config.get @qwen3_moe.routed_down.route_count : index + %configured_expert_count0 = config.get @qwen3_moe.routed_down.expert_count : index + %configured_output_size0 = config.get @qwen3_moe.routed_down.output_size : index + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 2048), le(%token_count, %token_capacity)] : index + %bounded_input_size, %configured_input_size = index.assume %input_size, %configured_input_size0 [range(%input_size, 256, 32768), mul(%input_size, 256), eq(%input_size, %configured_input_size0)] : index, index + %bounded_route_count, %configured_route_count = index.assume %route_count, %configured_route_count0 [range(%route_count, 1, 8), eq(%route_count, %configured_route_count0)] : index, index + %bounded_route_id_stride = index.assume %route_id_stride [range(%route_id_stride, 1, 512)] : index + %bounded_expert_count, %configured_expert_count = index.assume %expert_count, %configured_expert_count0 [range(%expert_count, 1, 512), eq(%expert_count, %configured_expert_count0)] : index, index + %bounded_output_size, %configured_output_size = index.assume %output_size, %configured_output_size0 [range(%output_size, 1, 4096), eq(%output_size, %configured_output_size0)] : index, index + %token0 = kernel.workgroup.id : index + %c0 = index.constant 0 : index + %valid_token = index.cmp ult, %token0, %bounded_token_count : index + %safe_token0 = scf.select %valid_token, %token0, %c0 : index + %safe_token, %body_token_count = index.assume %safe_token0, %bounded_token_count [lt(%safe_token0, %bounded_token_count)] : index, index + %c4 = index.constant 4 : index + %scale_stage_bytes = index.constant 1024 : offset + func.call @qwen3_moe_routed_down_q6k_f32_body(%valid_token, %c4, %c4, %scale_stage_bytes, %body_token_count, %safe_token, %configured_input_size, %configured_route_count, %bounded_route_id_stride, %configured_expert_count, %configured_output_size, %input, %route_ids, %route_weights, %weight, %output) : (i1, index, index, offset, index, index, index, index, index, index, index, buffer, buffer, buffer, buffer, buffer) + kernel.return +} + +// Decode-only direct-F32 route that publishes the normalized Q8_1 x4 row +// consumed by the next projection boundary. Each wave64 subgroup owns two +// adjacent output channels, so four subgroups publish eight channels while the +// completion epilogue still reduces four subgroup sums. +kernel.def target(@qwen3_moe_routed_down_q6k_gfx11_wave64) @qwen3_moe_routed_down_q6k_f32_wave64_next_q8(%token_count: index, %input_size: index, %route_count: index, %route_id_stride: index, %expert_count: index, %output_size: index) { + %configured_output_size = config.get @qwen3_moe.routed_down.output_size : index + %c1 = index.constant 1 : index + %c7 = index.constant 7 : index + %c8 = index.constant 8 : index + %c256 = index.constant 256 : index + %padded_output_size = index.add %configured_output_size, %c7 : index + %output_tiles = index.div %padded_output_size, %c8 : index + kernel.launch.config workgroups(%output_tiles, %c1, %c1) workgroup_size(%c256, %c1, %c1) : index +} launch(%token_count: index, %input_size: index, %route_count: index, %route_id_stride: index, %expert_count: index, %output_size: index, %input: buffer, %route_ids: buffer, %route_weights: buffer, %weight: buffer, %output: buffer, %norm_weight: buffer, %completion_counter: buffer, %next_q8_output: buffer) { + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 1)] : index + %configured_input_size0 = config.get @qwen3_moe.routed_down.input_size : index + %configured_route_count0 = config.get @qwen3_moe.routed_down.route_count : index + %configured_expert_count0 = config.get @qwen3_moe.routed_down.expert_count : index + %configured_output_size0 = config.get @qwen3_moe.routed_down.output_size : index + %bounded_input_size, %configured_input_size = index.assume %input_size, %configured_input_size0 [range(%input_size, 256, 32768), mul(%input_size, 256), eq(%input_size, %configured_input_size0)] : index, index + %bounded_route_count, %configured_route_count = index.assume %route_count, %configured_route_count0 [range(%route_count, 1, 8), eq(%route_count, %configured_route_count0)] : index, index + %bounded_route_id_stride = index.assume %route_id_stride [range(%route_id_stride, 1, 512)] : index + %bounded_expert_count, %configured_expert_count = index.assume %expert_count, %configured_expert_count0 [range(%expert_count, 1, 512), eq(%expert_count, %configured_expert_count0)] : index, index + %bounded_output_size, %configured_output_size = index.assume %output_size, %configured_output_size0 [range(%output_size, 128, 4096), mul(%output_size, 128), eq(%output_size, %configured_output_size0)] : index, index + %c4 = index.constant 4 : index + %c8 = index.constant 8 : index + %body_token = index.constant 0 : index + func.call @qwen3_moe_routed_down_q6k_f32_pair_body(%bounded_token_count, %body_token, %configured_input_size, %configured_route_count, %bounded_route_id_stride, %configured_expert_count, %configured_output_size, %input, %route_ids, %route_weights, %weight, %output) : (index, index, index, index, index, index, index, buffer, buffer, buffer, buffer, buffer) + func.call @qwen3_moe_routed_down_next_q8_completion(%c8, %c4, %bounded_token_count, %configured_output_size, %output, %norm_weight, %completion_counter, %next_q8_output) : (index, index, index, index, buffer, buffer, buffer, buffer) + kernel.return +} + +// Two selected experts exercise the noncompact route-ID stride, normalized +// weighted reduction, in-place residual update, K=768, and the output tile +// tail. Uniform 0xaa Q6_K bytes decode to a deterministic nonzero row with +// negative signed group scales. +check.case public @qwen3_moe_routed_down_q6k_q8_1_x4_nonzero_residual_case { + %token_count = check.literal value(1) : index + %input_size = check.literal value(768) : index + %route_count = check.literal value(2) : index + %route_id_stride = check.literal value(4) : index + %expert_count = check.literal value(2) : index + %output_size = check.literal value(9) : index + %routed_input = check.generate.fill value(0.00390625) : tensor<2x768xf32> + %q8_input = check.generate.fill value(0) : tensor<2x864xi8> + %route_ids = check.generate.iota offset(0) step(1) period(2) : tensor<1x4xi32> + %route_weights = check.generate.fill value(0.5) : tensor<1x2xf32> + %weight = check.generate.fill value(-86) : tensor<2x9x3x210xi8> + %output = check.generate.fill value(1.0) : tensor<1x9xf32> + %expected = check.generate.fill value(135.31431579589844) : tensor<1x9xf32> + %routed_row_count = check.literal value(2) : index + kernel.launch @ggml_quantize_q8_1_x4_f32[%routed_row_count, %input_size](%routed_row_count, %input_size, %routed_input, %q8_input) : [index, index](index, index, tensor<2x768xf32>, tensor<2x864xi8>) + kernel.launch @qwen3_moe_routed_down_q6k_q8_1_x4[%token_count, %input_size, %route_count, %route_id_stride, %expert_count, %output_size](%token_count, %input_size, %route_count, %route_id_stride, %expert_count, %output_size, %q8_input, %route_ids, %route_weights, %weight, %output) : [index, index, index, index, index, index](index, index, index, index, index, index, tensor<2x864xi8>, tensor<1x4xi32>, tensor<1x2xf32>, tensor<2x9x3x210xi8>, tensor<1x9xf32>) + check.expect.close actual(%output) expected(%expected) atol(0.25) rtol(9.9999999999999995e-07) nan(same) : tensor<1x9xf32> + check.return +} + +// Two tokens use different route-weight sums so token, route, packed-row, and +// output addressing cannot collapse to the single-token case. +check.case public @qwen3_moe_routed_down_q6k_q8_1_x4_nonzero_prefill_case { + %token_count = check.literal value(2) : index + %input_size = check.literal value(768) : index + %route_count = check.literal value(2) : index + %route_id_stride = check.literal value(4) : index + %expert_count = check.literal value(2) : index + %output_size = check.literal value(1) : index + %routed_row_count = check.literal value(4) : index + %routed_input = check.generate.fill value(0.00390625) : tensor<2x2x768xf32> + %q8_input = check.generate.fill value(0) : tensor<2x2x864xi8> + %route_ids = check.generate.iota offset(0) step(1) period(2) : tensor<2x4xi32> + %route_weights = check.generate.iota offset(0.25) step(0.25) period(4) : tensor<2x2xf32> + %weight = check.generate.fill value(-86) : tensor<2x1x3x210xi8> + %output = check.generate.fill value(1.0) : tensor<2x1xf32> + %expected = check.generate.iota offset(101.73573684692383) step(134.31431579589844) : tensor<2x1xf32> + kernel.launch @ggml_quantize_q8_1_x4_f32[%routed_row_count, %input_size](%routed_row_count, %input_size, %routed_input, %q8_input) : [index, index](index, index, tensor<2x2x768xf32>, tensor<2x2x864xi8>) + kernel.launch @qwen3_moe_routed_down_q6k_q8_1_x4[%token_count, %input_size, %route_count, %route_id_stride, %expert_count, %output_size](%token_count, %input_size, %route_count, %route_id_stride, %expert_count, %output_size, %q8_input, %route_ids, %route_weights, %weight, %output) : [index, index, index, index, index, index](index, index, index, index, index, index, tensor<2x2x864xi8>, tensor<2x4xi32>, tensor<2x2xf32>, tensor<2x1x3x210xi8>, tensor<2x1xf32>) + check.expect.close actual(%output) expected(%expected) atol(0.25) rtol(9.9999999999999995e-07) nan(same) : tensor<2x1xf32> + check.return +} + +// The exact-representable routed inputs make direct-F32 and Q8_1 contraction +// comparable while preserving noncompact route IDs and an output-channel tail. +check.case public @qwen3_moe_routed_down_q6k_f32_wave64_differential_case { + %token_count = check.literal value(1) : index + %input_size = check.literal value(768) : index + %route_count = check.literal value(2) : index + %route_id_stride = check.literal value(4) : index + %expert_count = check.literal value(2) : index + %output_size = check.literal value(9) : index + %routed_row_count = check.literal value(2) : index + %routed_input = check.generate.fill value(0.00390625) : tensor<2x768xf32> + %q8_input = check.generate.fill value(0) : tensor<2x864xi8> + %route_ids = check.generate.iota offset(0) step(1) period(2) : tensor<1x4xi32> + %route_weights = check.generate.fill value(0.5) : tensor<1x2xf32> + %weight = check.generate.fill value(-86) : tensor<2x9x3x210xi8> + %expected = check.generate.fill value(1.0) : tensor<1x9xf32> + %actual_wave64 = check.generate.fill value(1.0) : tensor<1x9xf32> + kernel.launch @ggml_quantize_q8_1_x4_f32[%routed_row_count, %input_size](%routed_row_count, %input_size, %routed_input, %q8_input) : [index, index](index, index, tensor<2x768xf32>, tensor<2x864xi8>) + kernel.launch @qwen3_moe_routed_down_q6k_q8_1_x4[%token_count, %input_size, %route_count, %route_id_stride, %expert_count, %output_size](%token_count, %input_size, %route_count, %route_id_stride, %expert_count, %output_size, %q8_input, %route_ids, %route_weights, %weight, %expected) : [index, index, index, index, index, index](index, index, index, index, index, index, tensor<2x864xi8>, tensor<1x4xi32>, tensor<1x2xf32>, tensor<2x9x3x210xi8>, tensor<1x9xf32>) + kernel.launch @qwen3_moe_routed_down_q6k_f32_wave64[%token_count, %input_size, %route_count, %route_id_stride, %expert_count, %output_size](%token_count, %input_size, %route_count, %route_id_stride, %expert_count, %output_size, %routed_input, %route_ids, %route_weights, %weight, %actual_wave64) : [index, index, index, index, index, index](index, index, index, index, index, index, tensor<2x768xf32>, tensor<1x4xi32>, tensor<1x2xf32>, tensor<2x9x3x210xi8>, tensor<1x9xf32>) + check.expect.close actual(%actual_wave64) expected(%expected) atol(0.25) rtol(0.01) nan(same) : tensor<1x9xf32> + check.return +} + +// Wave64 reference for the row producer used by the fused completion path. +kernel.def target(@qwen3_moe_routed_down_q6k_gfx11_wave64) @qwen3_moe_rmsnorm_quantize_q8_1_x4_wave64_production_check() { + %c1 = index.constant 1 : index + %c256 = index.constant 256 : index + kernel.launch.config workgroups(%c1, %c1, %c1) workgroup_size(%c256, %c1, %c1) : index +} launch(%input: buffer, %norm_weight: buffer, %q8_output: buffer) { + %publish_normalized = scalar.constant false : i1 + %reduction_subgroup_count = index.constant 4 : index + %token_count = index.constant 1 : index + %token = index.constant 0 : index + func.call @qwen3_moe_rmsnorm_quantize_q8_1_x4_body(%publish_normalized, %reduction_subgroup_count, %token_count, %token, %input, %norm_weight, %q8_output, %q8_output) : (i1, index, index, index, buffer, buffer, buffer, buffer) + kernel.return +} + +// The production hidden width requires every direct-F32 output tile to arrive +// before normalization. Nonzero Q6_K data compares the paired-row producer +// against the ordinary wave64 producer and standalone normalization, then +// reuses the completion word so stale completion state is observable. +check.case public @qwen3_moe_routed_down_q6k_f32_wave64_next_q8_differential_case { + %token_count = check.literal value(1) : index + %input_size = check.literal value(768) : index + %route_count = check.literal value(8) : index + %route_id_stride = check.literal value(8) : index + %expert_count = check.literal value(8) : index + %output_size = check.literal value(2048) : index + %input = check.generate.iota offset(-0.5) step(0.00016276041666666666) period(6144) : tensor<8x768xf32> + %route_ids = check.generate.iota offset(0) step(1) period(8) : tensor<8xi32> + %route_weights = check.generate.iota offset(0.027777777777777776) step(0.027777777777777776) : tensor<8xf32> + %weight = check.generate.fill value(-86) : tensor<8x2048x3x210xi8> + %norm_weight = check.generate.iota offset(-1.0) step(0.0009765625) : tensor<2048xf32> + %expected_output = check.generate.iota offset(-0.5) step(0.00048828125) period(2048) : tensor<2048xf32> + %actual_output0 = check.generate.iota offset(-0.5) step(0.00048828125) period(2048) : tensor<2048xf32> + %actual_output1 = check.generate.iota offset(-0.5) step(0.00048828125) period(2048) : tensor<2048xf32> + %expected_q8 = check.generate.fill value(0) : tensor<2304xi8> + %actual_q8_0 = check.generate.fill value(1) : tensor<2304xi8> + %actual_q8_1 = check.generate.fill value(1) : tensor<2304xi8> + %completion_counter = check.generate.fill value(0) : tensor<1xi32> + %expected_counter = check.generate.fill value(0) : tensor<1xi32> + kernel.launch @qwen3_moe_routed_down_q6k_f32_wave64[%token_count, %input_size, %route_count, %route_id_stride, %expert_count, %output_size](%token_count, %input_size, %route_count, %route_id_stride, %expert_count, %output_size, %input, %route_ids, %route_weights, %weight, %expected_output) : [index, index, index, index, index, index](index, index, index, index, index, index, tensor<8x768xf32>, tensor<8xi32>, tensor<8xf32>, tensor<8x2048x3x210xi8>, tensor<2048xf32>) + kernel.launch @qwen3_moe_rmsnorm_quantize_q8_1_x4_wave64_production_check(%expected_output, %norm_weight, %expected_q8) : (tensor<2048xf32>, tensor<2048xf32>, tensor<2304xi8>) + kernel.launch @qwen3_moe_routed_down_q6k_f32_wave64_next_q8[%token_count, %input_size, %route_count, %route_id_stride, %expert_count, %output_size](%token_count, %input_size, %route_count, %route_id_stride, %expert_count, %output_size, %input, %route_ids, %route_weights, %weight, %actual_output0, %norm_weight, %completion_counter, %actual_q8_0) : [index, index, index, index, index, index](index, index, index, index, index, index, tensor<8x768xf32>, tensor<8xi32>, tensor<8xf32>, tensor<8x2048x3x210xi8>, tensor<2048xf32>, tensor<2048xf32>, tensor<1xi32>, tensor<2304xi8>) + kernel.launch @qwen3_moe_routed_down_q6k_f32_wave64_next_q8[%token_count, %input_size, %route_count, %route_id_stride, %expert_count, %output_size](%token_count, %input_size, %route_count, %route_id_stride, %expert_count, %output_size, %input, %route_ids, %route_weights, %weight, %actual_output1, %norm_weight, %completion_counter, %actual_q8_1) : [index, index, index, index, index, index](index, index, index, index, index, index, tensor<8x768xf32>, tensor<8xi32>, tensor<8xf32>, tensor<8x2048x3x210xi8>, tensor<2048xf32>, tensor<2048xf32>, tensor<1xi32>, tensor<2304xi8>) + check.expect.close actual(%actual_output0) expected(%expected_output) atol(0.0) rtol(0.0) nan(same) : tensor<2048xf32> + check.expect.close actual(%actual_output1) expected(%expected_output) atol(0.0) rtol(0.0) nan(same) : tensor<2048xf32> + check.expect.equal actual(%actual_q8_0) expected(%expected_q8) : tensor<2304xi8> + check.expect.equal actual(%actual_q8_1) expected(%expected_q8) : tensor<2304xi8> + check.expect.equal actual(%completion_counter) expected(%expected_counter) : tensor<1xi32> + check.return +} + +check.case public @qwen3_moe_routed_down_q6k_q8_1_x4_benchmark_case { + %token_count = check.param.choice values([1, 2, 4, 8, 16, 17, 32, 63, 128, 129, 512]) name("token_count") : index + %input_size = check.literal value(768) : index + %route_count = check.literal value(8) : index + %route_id_stride = check.literal value(128) : index + %expert_count = check.literal value(128) : index + %output_size = check.literal value(2048) : index + %q8_input = check.generate.fill value(0) : tensor<[%token_count]x8x864xi8> + %route_ids = check.generate.iota offset(0) step(1) period(127) : tensor<[%token_count]x128xi32> + %route_weights = check.generate.fill value(0.125) : tensor<[%token_count]x8xf32> + %weight = check.generate.fill value(0) : tensor<128x2048x3x210xi8> + %output = check.generate.fill value(1.0) : tensor<[%token_count]x2048xf32> + %expected = check.generate.fill value(1.0) : tensor<[%token_count]x2048xf32> + kernel.launch @qwen3_moe_routed_down_q6k_q8_1_x4[%token_count, %input_size, %route_count, %route_id_stride, %expert_count, %output_size](%token_count, %input_size, %route_count, %route_id_stride, %expert_count, %output_size, %q8_input, %route_ids, %route_weights, %weight, %output) : [index, index, index, index, index, index](index, index, index, index, index, index, tensor<[%token_count]x8x864xi8>, tensor<[%token_count]x128xi32>, tensor<[%token_count]x8xf32>, tensor<128x2048x3x210xi8>, tensor<[%token_count]x2048xf32>) + check.expect.close actual(%output) expected(%expected) atol(0.0) rtol(0.0) nan(same) : tensor<[%token_count]x2048xf32> + check.return +} + +check.case public @qwen3_moe_routed_down_q6k_q8_1_x4_pipeline_benchmark_case { + %token_count = check.param.choice values([1, 2, 4, 8, 16, 17, 32, 63, 128, 129, 512]) name("token_count") : index + %input_size = check.literal value(768) : index + %route_count = check.literal value(8) : index + %route_id_stride = check.literal value(128) : index + %expert_count = check.literal value(128) : index + %output_size = check.literal value(2048) : index + // Each 768-element route contains six complete Q8_1 x4 groups, so packing + // the eight contiguous routes as one 6,144-element token row preserves the + // exact per-route physical layout without a derived testbench scalar. + %routed_input_size = check.literal value(6144) : index + %routed_input = check.generate.fill value(0.0) : tensor<[%token_count]x8x768xf32> + %q8_input = check.generate.fill value(1) : tensor<[%token_count]x8x864xi8> + %route_ids = check.generate.iota offset(0) step(1) period(127) : tensor<[%token_count]x128xi32> + %route_weights = check.generate.fill value(0.125) : tensor<[%token_count]x8xf32> + %weight = check.generate.fill value(0) : tensor<128x2048x3x210xi8> + %output = check.generate.fill value(1.0) : tensor<[%token_count]x2048xf32> + %expected = check.generate.fill value(1.0) : tensor<[%token_count]x2048xf32> + kernel.launch @ggml_quantize_q8_1_x4_f32[%token_count, %routed_input_size](%token_count, %routed_input_size, %routed_input, %q8_input) : [index, index](index, index, tensor<[%token_count]x8x768xf32>, tensor<[%token_count]x8x864xi8>) + kernel.launch @qwen3_moe_routed_down_q6k_q8_1_x4[%token_count, %input_size, %route_count, %route_id_stride, %expert_count, %output_size](%token_count, %input_size, %route_count, %route_id_stride, %expert_count, %output_size, %q8_input, %route_ids, %route_weights, %weight, %output) : [index, index, index, index, index, index](index, index, index, index, index, index, tensor<[%token_count]x8x864xi8>, tensor<[%token_count]x128xi32>, tensor<[%token_count]x8xf32>, tensor<128x2048x3x210xi8>, tensor<[%token_count]x2048xf32>) + check.expect.close actual(%output) expected(%expected) atol(0.0) rtol(0.0) nan(same) : tensor<[%token_count]x2048xf32> + check.return +} + +check.case public @qwen3_moe_routed_down_q6k_f32_wave64_benchmark_case { + %token_count = check.param.choice values([1, 2, 4, 8, 16, 17, 32, 63, 128, 129, 512]) name("token_count") : index + %input_size = check.literal value(768) : index + %route_count = check.literal value(8) : index + %route_id_stride = check.literal value(128) : index + %expert_count = check.literal value(128) : index + %output_size = check.literal value(2048) : index + %routed_input = check.generate.fill value(0.0) : tensor<[%token_count]x8x768xf32> + %route_ids = check.generate.iota offset(0) step(1) period(127) : tensor<[%token_count]x128xi32> + %route_weights = check.generate.fill value(0.125) : tensor<[%token_count]x8xf32> + %weight = check.generate.fill value(0) : tensor<128x2048x3x210xi8> + %output = check.generate.fill value(1.0) : tensor<[%token_count]x2048xf32> + %expected = check.generate.fill value(1.0) : tensor<[%token_count]x2048xf32> + kernel.launch @qwen3_moe_routed_down_q6k_f32_wave64[%token_count, %input_size, %route_count, %route_id_stride, %expert_count, %output_size](%token_count, %input_size, %route_count, %route_id_stride, %expert_count, %output_size, %routed_input, %route_ids, %route_weights, %weight, %output) : [index, index, index, index, index, index](index, index, index, index, index, index, tensor<[%token_count]x8x768xf32>, tensor<[%token_count]x128xi32>, tensor<[%token_count]x8xf32>, tensor<128x2048x3x210xi8>, tensor<[%token_count]x2048xf32>) + check.expect.close actual(%output) expected(%expected) atol(0.0) rtol(0.0) nan(same) : tensor<[%token_count]x2048xf32> + check.return +} + +check.benchmark<@qwen3_moe_routed_down_q6k_q8_1_x4_nonzero_residual_case> @qwen3_moe_routed_down_q6k_q8_1_x4_small + +check.benchmark<@qwen3_moe_routed_down_q6k_q8_1_x4_benchmark_case> @qwen3_moe_routed_down_q6k_q8_1_x4_decode {token_count = 1} + +check.benchmark<@qwen3_moe_routed_down_q6k_q8_1_x4_benchmark_case> @qwen3_moe_routed_down_q6k_q8_1_x4_small_batch_2 {token_count = 2} + +check.benchmark<@qwen3_moe_routed_down_q6k_q8_1_x4_benchmark_case> @qwen3_moe_routed_down_q6k_q8_1_x4_small_batch_4 {token_count = 4} + +check.benchmark<@qwen3_moe_routed_down_q6k_q8_1_x4_benchmark_case> @qwen3_moe_routed_down_q6k_q8_1_x4_small_batch_8 {token_count = 8} + +check.benchmark<@qwen3_moe_routed_down_q6k_q8_1_x4_benchmark_case> @qwen3_moe_routed_down_q6k_q8_1_x4_small_batch_16 {token_count = 16} + +check.benchmark<@qwen3_moe_routed_down_q6k_q8_1_x4_benchmark_case> @qwen3_moe_routed_down_q6k_q8_1_x4_prefill_32 {token_count = 32} + +check.benchmark<@qwen3_moe_routed_down_q6k_q8_1_x4_benchmark_case> @qwen3_moe_routed_down_q6k_q8_1_x4_prefill_128 {token_count = 128} + +check.benchmark<@qwen3_moe_routed_down_q6k_q8_1_x4_benchmark_case> @qwen3_moe_routed_down_q6k_q8_1_x4_prefill_512 {token_count = 512} + +check.benchmark<@qwen3_moe_routed_down_q6k_q8_1_x4_pipeline_benchmark_case> @qwen3_moe_routed_down_q6k_q8_1_x4_pipeline_decode {token_count = 1} + +check.benchmark<@qwen3_moe_routed_down_q6k_q8_1_x4_pipeline_benchmark_case> @qwen3_moe_routed_down_q6k_q8_1_x4_pipeline_prefill_32 {token_count = 32} + +check.benchmark<@qwen3_moe_routed_down_q6k_q8_1_x4_pipeline_benchmark_case> @qwen3_moe_routed_down_q6k_q8_1_x4_pipeline_prefill_128 {token_count = 128} + +check.benchmark<@qwen3_moe_routed_down_q6k_q8_1_x4_pipeline_benchmark_case> @qwen3_moe_routed_down_q6k_q8_1_x4_pipeline_prefill_512 {token_count = 512} + +check.benchmark<@qwen3_moe_routed_down_q6k_f32_wave64_benchmark_case> @qwen3_moe_routed_down_q6k_f32_wave64_decode {token_count = 1} + +check.benchmark<@qwen3_moe_routed_down_q6k_f32_wave64_benchmark_case> @qwen3_moe_routed_down_q6k_f32_wave64_small_batch_2 {token_count = 2} + +check.benchmark<@qwen3_moe_routed_down_q6k_f32_wave64_benchmark_case> @qwen3_moe_routed_down_q6k_f32_wave64_small_batch_4 {token_count = 4} + +check.benchmark<@qwen3_moe_routed_down_q6k_f32_wave64_benchmark_case> @qwen3_moe_routed_down_q6k_f32_wave64_small_batch_8 {token_count = 8} + +check.benchmark<@qwen3_moe_routed_down_q6k_f32_wave64_benchmark_case> @qwen3_moe_routed_down_q6k_f32_wave64_small_batch_16 {token_count = 16} + +check.benchmark<@qwen3_moe_routed_down_q6k_q8_1_x4_next_q8_benchmark_case> @qwen3_moe_routed_down_q6k_q8_1_x4_next_q8_decode + +check.benchmark<@qwen3_moe_routed_down_q6k_q8_1_x4_next_q8_composed_benchmark_case> @qwen3_moe_routed_down_q6k_q8_1_x4_next_q8_composed_decode diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/routed_down_quantized_f16_wmma.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/routed_down_quantized_f16_wmma.loom new file mode 100644 index 000000000000..aca0ec90718b --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/routed_down_quantized_f16_wmma.loom @@ -0,0 +1,987 @@ +// Expert-grouped Q4_K and Q6_K down projections for gfx11 prefill. +// +// Each two-wave workgroup decodes 64 output channels for 32 routed rows of one +// expert. Raw GGUF weights and FP16 routed SwiGLU rows are staged as FP16 WMMA +// operands. Results remain FP16 and are scattered into compact +// [token, route, hidden] order without collisions. A following reduction +// widens each route, applies its normalized weight, and accumulates the +// residual in FP32. +// +// Top-k routing selects an expert at most once per token. The expert table +// therefore contains at most token_count assignments per expert, while each +// assignment ordinal still identifies the compact [token, route] activation +// and route-weight row. +func.def inline @qwen3_moe_q4k_scale_from_header(%scale0: i32, %scale1: i32, %scale2: i32, %q4_group: index) -> (i32, i32) { + %c4 = index.constant 4 : index + %c8 = index.constant 8 : index + %c2_i32 = scalar.constant 2 : i32 + %c4_i32 = scalar.constant 4 : i32 + %c15_i32 = scalar.constant 15 : i32 + %c48_i32 = scalar.constant 48 : i32 + %bounded_group = index.assume %q4_group [range(%q4_group, 0, 7)] : index + %is_low_group = index.cmp ult, %bounded_group, %c4 : index + %scale_lane = index.rem %bounded_group, %c4 : index + %scale_shift_index = index.mul %scale_lane, %c8 : index + %scale_shift = index.cast %scale_shift_index : index to i32 + %high_shift = scalar.addi %scale_shift, %c2_i32 : i32 + %minimum_shift = scalar.addi %scale_shift, %c4_i32 : i32 + %selected_scale_source = scf.select %is_low_group, %scale0, %scale2 : i32 + %selected_minimum_source = scf.select %is_low_group, %scale1, %scale2 : i32 + %selected_scale_high_shift = scf.select %is_low_group, %scale_shift, %high_shift : i32 + %selected_minimum_low_shift = scf.select %is_low_group, %scale_shift, %minimum_shift : i32 + %scale_low0 = scalar.shrui %selected_scale_source, %scale_shift : i32 + %scale_low = scalar.andi %scale_low0, %c15_i32 : i32 + %scale_high0 = scalar.shrui %scale0, %selected_scale_high_shift : i32 + %scale_high = scalar.andi %scale_high0, %c48_i32 : i32 + %scale = scalar.ori %scale_low, %scale_high : i32 + %minimum_low0 = scalar.shrui %selected_minimum_source, %selected_minimum_low_shift : i32 + %minimum_low = scalar.andi %minimum_low0, %c15_i32 : i32 + %minimum_high0 = scalar.shrui %scale1, %selected_scale_high_shift : i32 + %minimum_high = scalar.andi %minimum_high0, %c48_i32 : i32 + %minimum = scalar.ori %minimum_low, %minimum_high : i32 + func.return %scale, %minimum : i32, i32 +} + +func.def inline @qwen3_moe_q4k_wmma_load_code(%weight: buffer, %row_byte_base: offset, %q4_block: index, %q4_group_pair: index, %packet: index) -> (vector<1xi32>) { + %c8 = index.constant 8 : index + %block_bytes = index.constant 144 : offset + %code_offset = index.constant 16 : offset + %bounded_group_pair = index.assume %q4_group_pair [range(%q4_group_pair, 0, 3)] : index + %bounded_packet = index.assume %packet [range(%packet, 0, 7)] : index + %block_byte_add = index.scale %q4_block, %block_bytes : index, offset -> offset + %block_byte_base = index.add %row_byte_base, %block_byte_add : offset + %code_byte_base = index.add %block_byte_base, %code_offset : offset + %code_view = buffer.view %weight[%code_byte_base] : buffer -> view<32xi32> + %q_page = index.mul %bounded_group_pair, %c8 : index + %q_word_index0 = index.add %q_page, %bounded_packet : index + %q_word_index = index.assume %q_word_index0 [range(%q_word_index0, 0, 31)] : index + %q_word = vector.load %code_view[%q_word_index] : view<32xi32> -> vector<1xi32> + func.return %q_word : vector<1xi32> +} + +func.def inline @qwen3_moe_q4k_wmma_load_header(%weight: buffer, %row_byte_base: offset, %q4_block: index) -> (vector<4xi32>) { + %c0 = index.constant 0 : index + %block_bytes = index.constant 144 : offset + %block_byte_add = index.scale %q4_block, %block_bytes : index, offset -> offset + %block_byte_base = index.add %row_byte_base, %block_byte_add : offset + %header_view = buffer.view %weight[%block_byte_base] : buffer -> view<4xi32> + %header_words = vector.load %header_view[%c0] : view<4xi32> -> vector<4xi32> + func.return %header_words : vector<4xi32> +} + +func.def inline @qwen3_moe_q4k_wmma_vector4_from_header_code(%q4_group: index, %header_words: vector<4xi32>, %q_word: vector<1xi32>) -> (vector<4xf16>) { + %c2 = index.constant 2 : index + %c4 = index.constant 4 : index + %q4_mask = vector.constant 252645135 : vector<1xi32> + %bounded_group = index.assume %q4_group [range(%q4_group, 0, 7)] : index + %header_halves = vector.bitcast %header_words : vector<4xi32> to vector<8xf16> + %d_f16 = vector.extract %header_halves[0] : vector<8xf16> -> f16 + %dmin_f16 = vector.extract %header_halves[1] : vector<8xf16> -> f16 + %scale0 = vector.extract %header_words[1] : vector<4xi32> -> i32 + %scale1 = vector.extract %header_words[2] : vector<4xi32> -> i32 + %scale2 = vector.extract %header_words[3] : vector<4xi32> -> i32 + %d = scalar.extf %d_f16 : f16 to f32 + %dmin = scalar.extf %dmin_f16 : f16 to f32 + %scale, %minimum = func.call @qwen3_moe_q4k_scale_from_header(%scale0, %scale1, %scale2, %bounded_group) : (i32, i32, i32, index) -> (i32, i32) + %scale_f32 = scalar.uitofp %scale : i32 to f32 + %minimum_f32 = scalar.uitofp %minimum : i32 to f32 + %d_scale = scalar.mulf %d, %scale_f32 : f32 + %minimum_scale = scalar.mulf %dmin, %minimum_f32 : f32 + %q_half = index.rem %bounded_group, %c2 : index + %q_shift_index = index.mul %q_half, %c4 : index + %q_shift_i32 = index.cast %q_shift_index : index to i32 + %q_shift = vector.splat %q_shift_i32 : vector<1xi32> + %shifted_q = vector.shrui %q_word, %q_shift : vector<1xi32> + %masked_q = vector.andi %shifted_q, %q4_mask : vector<1xi32> + %q_i8 = vector.bitcast %masked_q : vector<1xi32> to vector<4xi8> + %q_f32 = vector.uitofp %q_i8 : vector<4xi8> to vector<4xf32> + // Form adjacent FP16 lanes from fused FP32 affine expressions. AMDGPU maps + // this natural shape to packed mixlo/mixhi instructions where available. + %negative_minimum_scale = scalar.negf %minimum_scale : f32 + %q0 = vector.extract %q_f32[0] : vector<4xf32> -> f32 + %q1 = vector.extract %q_f32[1] : vector<4xf32> -> f32 + %q2 = vector.extract %q_f32[2] : vector<4xf32> -> f32 + %q3 = vector.extract %q_f32[3] : vector<4xf32> -> f32 + %value0 = scalar.fmaf %q0, %d_scale, %negative_minimum_scale : f32 + %value1 = scalar.fmaf %q1, %d_scale, %negative_minimum_scale : f32 + %value2 = scalar.fmaf %q2, %d_scale, %negative_minimum_scale : f32 + %value3 = scalar.fmaf %q3, %d_scale, %negative_minimum_scale : f32 + %half0 = scalar.fptrunc %value0 : f32 to f16 + %half1 = scalar.fptrunc %value1 : f32 to f16 + %half2 = scalar.fptrunc %value2 : f32 to f16 + %half3 = scalar.fptrunc %value3 : f32 to f16 + %result = vector.from_elements %half0, %half1, %half2, %half3 : vector<4xf16> + func.return %result : vector<4xf16> +} + +amdgpu.target @qwen3_moe_routed_down_gfx11_wave64 {subgroup_size = 64} + +config.decl @qwen3_moe.routed_down.input_size : %value: index where [range(%value, 256, 32768), mul(%value, 256)] + +config.decl @qwen3_moe.routed_down.route_count : %value: index where [range(%value, 1, 8)] + +config.decl @qwen3_moe.routed_down.expert_count : %value: index where [range(%value, 1, 512)] + +config.decl @qwen3_moe.routed_down.output_size : %value: index where [range(%value, 1, 4096)] + +kernel.decl @ggml_quantize_q8_1_x4_f32(%token_count: index, %input_size: index) launch(%token_count: index, %input_size: index, %input: buffer, %output: buffer) + +kernel.decl @qwen3_moe_build_expert_table(%token_count: index, %route_count: index, %route_stride: index, %expert_count: index) launch(%token_count: index, %route_count: index, %route_stride: index, %expert_count: index, %route_ids: buffer, %expert_table: buffer) + +kernel.decl @qwen3_moe_build_expert_partition_table(%token_count: index, %route_count: index, %expert_count: index) launch(%token_count: index, %route_count: index, %expert_count: index, %expert_table: buffer, %partition_table: buffer) + +kernel.decl @qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma(%token_count: index) launch(%token_count: index, %input: buffer, %expert_table: buffer, %partition_table: buffer, %gate_weight: buffer, %up_weight: buffer, %output: buffer) + +kernel.decl @qwen3_moe_routed_gate_up_swiglu_q4k_q8(%token_count: index, %route_count: index, %route_stride: index, %expert_count: index, %output_size: index) launch(%token_count: index, %route_count: index, %route_stride: index, %expert_count: index, %output_size: index, %q8_input: buffer, %route_ids: buffer, %gate_weight: buffer, %up_weight: buffer, %output: buffer) + +func.decl @qwen3_moe_q4k_wmma_vector4(%weight: buffer, %row_byte_base: offset, %q4_block: index, %q4_group: index, %packet: index) -> (vector<4xf16>) + +func.decl @qwen3_moe_q4k_wmma_load_header(%weight: buffer, %row_byte_base: offset, %q4_block: index) -> (vector<4xi32>) + +func.decl @qwen3_moe_q4k_wmma_load_code(%weight: buffer, %row_byte_base: offset, %q4_block: index, %q4_group_pair: index, %packet: index) -> (vector<1xi32>) + +func.decl @qwen3_moe_q4k_wmma_vector4_from_header_code(%q4_group: index, %header_words: vector<4xi32>, %q_word: vector<1xi32>) -> (vector<4xf16>) + +func.decl @ggml_q6k_f16_vector4(%weight: buffer, %weight_row_byte_base: offset, %q6_block: index, %q6_group: index, %packet: index) -> (vector<4xf16>) + +kernel.decl @qwen3_moe_routed_down_q4k_q8_1_x4(%token_count: index, %input_size: index, %route_count: index, %route_id_stride: index, %expert_count: index, %output_size: index) launch(%token_count: index, %input_size: index, %route_count: index, %route_id_stride: index, %expert_count: index, %output_size: index, %q8_input: buffer, %route_ids: buffer, %route_weights: buffer, %weight: buffer, %output: buffer) + +kernel.decl @qwen3_moe_routed_down_q6k_q8_1_x4(%token_count: index, %input_size: index, %route_count: index, %route_id_stride: index, %expert_count: index, %output_size: index) launch(%token_count: index, %input_size: index, %route_count: index, %route_id_stride: index, %expert_count: index, %output_size: index, %q8_input: buffer, %route_ids: buffer, %route_weights: buffer, %weight: buffer, %output: buffer) + +// Acquires the packed words shared by one four-group half of a Q6_K block. +// Each QL word supplies two groups and the QH word supplies all four. +func.def inline @qwen3_moe_q6k_wmma_load_half_codes(%weight: buffer, %weight_row_byte_base: offset, %q6_block: index, %q6_half: index, %packet: index) -> (vector<1xi32>, vector<1xi32>, vector<1xi32>) { + %c8 = index.constant 8 : index + %c16 = index.constant 16 : index + %block_bytes = index.constant 210 : offset + %qh_byte_add = index.constant 128 : offset + %bounded_half = index.assume %q6_half [range(%q6_half, 0, 1)] : index + %bounded_packet = index.assume %packet [range(%packet, 0, 7)] : index + %block_byte_add = index.scale %q6_block, %block_bytes : index, offset -> offset + %block_byte_base = index.add %weight_row_byte_base, %block_byte_add : offset + %qh_byte_base = index.add %block_byte_base, %qh_byte_add : offset + %ql_view = buffer.view %weight[%block_byte_base] : buffer -> view<32xi32> + %qh_view = buffer.view %weight[%qh_byte_base] : buffer -> view<16xi32> + %ql_half_word_base = index.mul %bounded_half, %c16 : index + %ql0_word_index = index.add %ql_half_word_base, %bounded_packet : index + %ql1_word_base = index.add %ql_half_word_base, %c8 : index + %ql1_word_index = index.add %ql1_word_base, %bounded_packet : index + %qh_half_word_base = index.mul %bounded_half, %c8 : index + %qh_word_index = index.add %qh_half_word_base, %bounded_packet : index + %ql0_word = vector.load %ql_view[%ql0_word_index] : view<32xi32> -> vector<1xi32> + %ql1_word = vector.load %ql_view[%ql1_word_index] : view<32xi32> -> vector<1xi32> + %qh_word = vector.load %qh_view[%qh_word_index] : view<16xi32> -> vector<1xi32> + func.return %ql0_word, %ql1_word, %qh_word : vector<1xi32>, vector<1xi32>, vector<1xi32> +} + +// Decodes four adjacent values after the surrounding schedule has selected +// the scale and retained the packed code words at their natural lifetimes. +func.def inline @qwen3_moe_q6k_wmma_vector4_from_scale_codes(%q6_group: index, %scale_i8: i8, %d_f16: f16, %ql_word: vector<1xi32>, %qh_word: vector<1xi32>) -> (vector<4xf16>) { + %c2 = index.constant 2 : index + %c4 = index.constant 4 : index + %c4_i32v = vector.constant 4 : vector<1xi32> + %nibble_mask = vector.constant 252645135 : vector<1xi32> + %high_mask = vector.constant 50529027 : vector<1xi32> + %c32_f32v = vector.constant 32.0 : vector<4xf32> + %bounded_group = index.assume %q6_group [range(%q6_group, 0, 7)] : index + %group_in_half = index.rem %bounded_group, %c4 : index + %nibble = index.div %group_in_half, %c2 : index + %nibble_shift_index = index.mul %nibble, %c4 : index + %nibble_shift_i32 = index.cast %nibble_shift_index : index to i32 + %nibble_shift = vector.splat %nibble_shift_i32 : vector<1xi32> + %qh_shift_index = index.mul %group_in_half, %c2 : index + %qh_shift_i32 = index.cast %qh_shift_index : index to i32 + %qh_shift = vector.splat %qh_shift_i32 : vector<1xi32> + %ql_shifted = vector.shrui %ql_word, %nibble_shift : vector<1xi32> + %ql = vector.andi %ql_shifted, %nibble_mask : vector<1xi32> + %qh_shifted = vector.shrui %qh_word, %qh_shift : vector<1xi32> + %qh_low = vector.andi %qh_shifted, %high_mask : vector<1xi32> + %qh = vector.shli %qh_low, %c4_i32v : vector<1xi32> + %code = vector.ori %ql, %qh : vector<1xi32> + %code_i8 = vector.bitcast %code : vector<1xi32> to vector<4xi8> + %code_f32 = vector.uitofp %code_i8 : vector<4xi8> to vector<4xf32> + %centered = vector.subf %code_f32, %c32_f32v : vector<4xf32> + %scale = scalar.sitofp %scale_i8 : i8 to f32 + %d = scalar.extf %d_f16 : f16 to f32 + %combined_scale = scalar.mulf %scale, %d : f32 + %combined_scale_vector = vector.splat %combined_scale : vector<4xf32> + %values_f32 = vector.mulf %centered, %combined_scale_vector : vector<4xf32> + %values = vector.fptrunc %values_f32 : vector<4xf32> to vector<4xf16> + func.return %values : vector<4xf16> +} + +// Loads only the group scale and block multiplier while reusing packed words +// retained by an enclosing four-group half-block schedule. +func.def inline @qwen3_moe_q6k_wmma_vector4_from_half_codes(%weight: buffer, %weight_row_byte_base: offset, %q6_block: index, %q6_group: index, %packet: index, %ql0_word: vector<1xi32>, %ql1_word: vector<1xi32>, %qh_word: vector<1xi32>) -> (vector<4xf16>) { + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c4 = index.constant 4 : index + %block_bytes = index.constant 210 : offset + %scale_byte_add = index.constant 192 : offset + %d_byte_add = index.constant 208 : offset + %bounded_group = index.assume %q6_group [range(%q6_group, 0, 7)] : index + %bounded_packet = index.assume %packet [range(%packet, 0, 7)] : index + %block_byte_add = index.scale %q6_block, %block_bytes : index, offset -> offset + %block_byte_base = index.add %weight_row_byte_base, %block_byte_add : offset + %scale_byte_base = index.add %block_byte_base, %scale_byte_add : offset + %d_byte_base = index.add %block_byte_base, %d_byte_add : offset + %scale_view = buffer.view %weight[%scale_byte_base] : buffer -> view<16xi8> + %d_view = buffer.view %weight[%d_byte_base] : buffer -> view<1xf16> + %group_in_half = index.rem %bounded_group, %c4 : index + %ql_side = index.rem %group_in_half, %c2 : index + %uses_ql1 = index.cmp eq, %ql_side, %c1 : index + %ql_word = scf.select %uses_ql1, %ql1_word, %ql0_word : vector<1xi32> + %scale_packet_half = index.div %bounded_packet, %c4 : index + %scale_group_base = index.mul %bounded_group, %c2 : index + %scale_index = index.add %scale_group_base, %scale_packet_half : index + %scale_i8 = view.load %scale_view[%scale_index] : view<16xi8> -> i8 + %d_f16 = view.load %d_view[%c0] : view<1xf16> -> f16 + %values = func.call @qwen3_moe_q6k_wmma_vector4_from_scale_codes(%bounded_group, %scale_i8, %d_f16, %ql_word, %qh_word) : (index, i8, f16, vector<1xi32>, vector<1xi32>) -> (vector<4xf16>) + func.return %values : vector<4xf16> +} + +// Shared raw-quantized matrix schedule. Entry points pass a literal weight +// format so linking and JIT specialization erase the inactive packed decoder. +func.def inline @qwen3_moe_routed_down_quantized_f16_wmma_body(%weight_format: index, %token_count: index, %input: buffer, %expert_table: buffer, %weight: buffer, %output: buffer) { + %input_size = config.get @qwen3_moe.routed_down.input_size : index + %route_count = config.get @qwen3_moe.routed_down.route_count : index + %expert_count = config.get @qwen3_moe.routed_down.expert_count : index + %output_size = config.get @qwen3_moe.routed_down.output_size : index + %bounded_route_count = index.assume %route_count [range(%route_count, 1, 8)] : index + %bounded_expert_count = index.assume %expert_count [range(%expert_count, 1, 128)] : index + %bounded_output_size = index.assume %output_size [range(%output_size, 1, 4096)] : index + %channel_tile = kernel.workgroup.id : index + %route_tile = kernel.workgroup.id : index + %expert = kernel.workgroup.id : index + %workitem = kernel.workitem.id : index + %subgroup0 = kernel.subgroup.id : index + %subgroup = index.assume %subgroup0 [range(%subgroup0, 0, 1)] : index + %lane = kernel.subgroup.lane.id : index + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c4 = index.constant 4 : index + %c6 = index.constant 6 : index + %c8 = index.constant 8 : index + %c16 = index.constant 16 : index + %c32 = index.constant 32 : index + %c40 = index.constant 40 : index + %c48 = index.constant 48 : index + %c63 = index.constant 63 : index + %c64 = index.constant 64 : index + %c256 = index.constant 256 : index + %c0_offset = index.constant 0 : offset + %c4_bytes = index.constant 4 : offset + %q4_block_bytes = index.constant 144 : offset + %q6_block_bytes = index.constant 210 : offset + %weight_stage_bytes = index.constant 5120 : offset + %activation_stage_bytes = index.constant 2560 : offset + %route_stage_bytes = index.constant 128 : offset + %wave_result_stage_bytes = index.constant 512 : offset + %result_stage_bytes = index.constant 1024 : offset + %c0_i32 = scalar.constant 0 : i32 + %cn1_i32 = scalar.constant -1 : i32 + %c0_i32x1 = vector.constant 0 : vector<1xi32> + %c0_i32x4 = vector.constant 0 : vector<4xi32> + %c0_f16x4 = vector.constant 0.0 : vector<4xf16> + %zero_accumulator = vector.constant 0.0 : vector<8xf16> + %m = index.constant 16 : index + %n = index.constant 16 : index + %k = index.constant 16 : index + %assignment_count = index.mul %token_count, %bounded_route_count : index + %assignment_table_byte_base = index.scale %bounded_expert_count, %c4_bytes : index, offset -> offset + %is_q4 = index.cmp eq, %weight_format, %c4 : index + %is_q6 = index.cmp eq, %weight_format, %c6 : index + %quant_block_bytes = scf.select %is_q6, %q6_block_bytes, %q4_block_bytes : offset + %quant_block_count = index.div %input_size, %c256 : index + %weight_row_bytes = index.scale %quant_block_count, %quant_block_bytes : index, offset -> offset + %weight_expert_bytes = index.scale %bounded_output_size, %weight_row_bytes : index, offset -> offset + %input_noalias, %expert_table_noalias, %weight_noalias, %output_noalias = buffer.assume.noalias %input, %expert_table, %weight, %output : buffer, buffer, buffer, buffer + %input_view = buffer.view %input_noalias[%c0_offset] : buffer -> view<[%assignment_count]x[%input_size]xf16> + %count_view = buffer.view %expert_table_noalias[%c0_offset] : buffer -> view<[%bounded_expert_count]xi32> + %assignment_view = buffer.view %expert_table_noalias[%assignment_table_byte_base] : buffer -> view<[%bounded_expert_count]x[%token_count]xi32> + %output_view = buffer.view %output_noalias[%c0_offset] : buffer -> view<[%assignment_count]x[%bounded_output_size]xf16> + %weight_stage = buffer.alloca align(16) %weight_stage_bytes : buffer + %activation_stage = buffer.alloca align(16) %activation_stage_bytes : buffer + %route_stage = buffer.alloca align(16) %route_stage_bytes : buffer + %result_stage = buffer.alloca align(16) %result_stage_bytes : buffer + %weight_stage_view = buffer.view %weight_stage[%c0_offset] : buffer -> view<64x40xf16> + %activation_stage_physical_view = buffer.view %activation_stage[%c0_offset] : buffer -> view<32x40xf16> + %activation_fragment_layout = encoding.layout.strided [1, %c40] : encoding + %activation_fragment_view = buffer.view %activation_stage[%c0_offset] : buffer -> view<32x32xf16, %activation_fragment_layout> + %route_stage_view = buffer.view %route_stage[%c0_offset] : buffer -> view<32xi32> + %wave_result_stage_offset = index.scale %subgroup, %wave_result_stage_bytes : index, offset -> offset + %result_fragment_layout = encoding.layout.strided [1, %c16] : encoding + %result_fragment_view = buffer.view %result_stage[%wave_result_stage_offset] : buffer -> view<16x16xf16, %result_fragment_layout> + %result_physical_view = buffer.view %result_stage[%wave_result_stage_offset] : buffer -> view<16x16xf16> + %channel_tile_base = index.mul %channel_tile, %c64 : index + %initial_route_tile_base = index.mul %route_tile, %c32 : index + // The launch workload fixes the interleaved route partitions for this exact + // command-program specialization. + %padded_token_count = index.add %token_count, %c63 : index + %route_partition_count = index.div %padded_token_count, %c64 : index + %route_partition_step = index.mul %route_partition_count, %c32 : index + %bounded_expert, %table_expert_count = index.assume %expert, %bounded_expert_count [lt(%expert, %bounded_expert_count)] : index, index + %is_workitem_zero = index.cmp eq, %workitem, %c0 : index + %lane_expert_route_count = scf.if %is_workitem_zero -> (i32) { + %loaded = view.load %count_view[%bounded_expert] : view<[%bounded_expert_count]xi32> -> i32 + scf.yield %loaded : i32 + } else { + scf.yield %c0_i32 : i32 + } + %expert_route_count_reduced = kernel.workgroup.reduce %lane_expert_route_count : i32 + %expert_route_count_i32 = kernel.subgroup.broadcast.first %expert_route_count_reduced : i32 + %expert_route_count0 = index.cast %expert_route_count_i32 : i32 to index + %expert_route_count = index.assume %expert_route_count0 [range(%expert_route_count0, 0, 2048)] : index + scf.for %route_tile_base = [%initial_route_tile_base to %expert_route_count step %route_partition_step] { + // Snapshot this expert's compact assignment map once, then reuse it for + // every K group and output-channel tile. + %loads_route = index.cmp ult, %workitem, %c32 : index + scf.if %loads_route { + %local_route = index.assume %workitem [range(%workitem, 0, 31)] : index + %assignment_ordinal = index.add %route_tile_base, %local_route : index + %valid_row = index.cmp ult, %assignment_ordinal, %expert_route_count : index + %assignment_i32 = scf.if %valid_row -> (i32) { + %bounded_assignment_ordinal, %table_token_count = index.assume %assignment_ordinal, %token_count [lt(%assignment_ordinal, %token_count)] : index, index + %loaded = view.load %assignment_view[%bounded_expert, %bounded_assignment_ordinal] : view<[%bounded_expert_count]x[%token_count]xi32> -> i32 + scf.yield %loaded : i32 + } else { + scf.yield %cn1_i32 : i32 + } + view.store %assignment_i32, %route_stage_view[%local_route] : i32, view<32xi32> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + %load_packet = index.rem %workitem, %c8 : index + %load_k = index.mul %load_packet, %c4 : index + %load_row0 = index.div %workitem, %c8 : index + %load_row = index.assume %load_row0 [range(%load_row0, 0, 15)] : index + %expert_byte_base = index.scale %bounded_expert, %weight_expert_bytes : index, offset -> offset + %subgroup_channel_add = index.mul %subgroup, %c32 : index + %subgroup_channel1 = index.add %subgroup_channel_add, %c16 : index + %init00 = vector.fragment %zero_accumulator shape [%m, %n] : vector<8xf16> + %init01 = vector.fragment %zero_accumulator shape [%m, %n] : vector<8xf16> + %init10 = vector.fragment %zero_accumulator shape [%m, %n] : vector<8xf16> + %init11 = vector.fragment %zero_accumulator shape [%m, %n] : vector<8xf16> + %result00, %result01, %result10, %result11 = scf.for %quant_block = [%c0 to %quant_block_count step %c1](%block_acc00 = %init00 : vector<8xf16>, %block_acc01 = %init01 : vector<8xf16>, %block_acc10 = %init10 : vector<8xf16>, %block_acc11 = %init11 : vector<8xf16>) -> (vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>) { + // Four lanes cover the 64 output rows owned by one load packet. Snapshot + // each Q4_K block header once and retain it across all eight quant groups. + // The Q6 specialization erases this loop with %is_q4=false. + %header0, %header1, %header2, %header3 = scf.for %header_row_offset = [%c0 to %c64 step %c16](%prior_header0 = %c0_i32x4 : vector<4xi32>, %prior_header1 = %c0_i32x4 : vector<4xi32>, %prior_header2 = %c0_i32x4 : vector<4xi32>, %prior_header3 = %c0_i32x4 : vector<4xi32>) -> (vector<4xi32>, vector<4xi32>, vector<4xi32>, vector<4xi32>) unroll { + %header_local_row0 = index.add %load_row, %header_row_offset : index + %header_local_row = index.assume %header_local_row0 [range(%header_local_row0, 0, 63)] : index + %header_channel = index.add %channel_tile_base, %header_local_row : index + %valid_header_channel = index.cmp ult, %header_channel, %bounded_output_size : index + %loads_q4_header = scalar.andi %is_q4, %valid_header_channel : i1 + %loaded_header = scf.if %loads_q4_header -> (vector<4xi32>) { + %header_channel_byte_add = index.scale %header_channel, %weight_row_bytes : index, offset -> offset + %header_row_byte_base = index.add %expert_byte_base, %header_channel_byte_add : offset + %header_words = func.call @qwen3_moe_q4k_wmma_load_header(%weight_noalias, %header_row_byte_base, %quant_block) : (buffer, offset, index) -> (vector<4xi32>) + scf.yield %header_words : vector<4xi32> + } else { + scf.yield %c0_i32x4 : vector<4xi32> + } + %updates_header0 = index.cmp eq, %header_row_offset, %c0 : index + %updates_header1 = index.cmp eq, %header_row_offset, %c16 : index + %updates_header2 = index.cmp eq, %header_row_offset, %c32 : index + %updates_header3 = index.cmp eq, %header_row_offset, %c48 : index + %next_header0 = scf.select %updates_header0, %loaded_header, %prior_header0 : vector<4xi32> + %next_header1 = scf.select %updates_header1, %loaded_header, %prior_header1 : vector<4xi32> + %next_header2 = scf.select %updates_header2, %loaded_header, %prior_header2 : vector<4xi32> + %next_header3 = scf.select %updates_header3, %loaded_header, %prior_header3 : vector<4xi32> + scf.yield %next_header0, %next_header1, %next_header2, %next_header3 : vector<4xi32>, vector<4xi32>, vector<4xi32>, vector<4xi32> + } + %quant_group_outer_count = scf.select %is_q4, %c4, %c2 : index + %groups_per_outer = scf.select %is_q4, %c2, %c4 : index + %block_result00, %block_result01, %block_result10, %block_result11 = scf.for %quant_group_outer = [%c0 to %quant_group_outer_count step %c1](%acc00 = %block_acc00 : vector<8xf16>, %acc01 = %block_acc01 : vector<8xf16>, %acc10 = %block_acc10 : vector<8xf16>, %acc11 = %block_acc11 : vector<8xf16>) -> (vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>) { + // Q4_K retains one packed code word across its adjacent low/high group + // pair. The Q6 specialization erases this acquisition path. + %q_word0, %q_word1, %q_word2, %q_word3 = scf.for %code_row_offset = [%c0 to %c64 step %c16](%prior_q_word0 = %c0_i32x1 : vector<1xi32>, %prior_q_word1 = %c0_i32x1 : vector<1xi32>, %prior_q_word2 = %c0_i32x1 : vector<1xi32>, %prior_q_word3 = %c0_i32x1 : vector<1xi32>) -> (vector<1xi32>, vector<1xi32>, vector<1xi32>, vector<1xi32>) unroll { + %code_local_row0 = index.add %load_row, %code_row_offset : index + %code_local_row = index.assume %code_local_row0 [range(%code_local_row0, 0, 63)] : index + %code_channel = index.add %channel_tile_base, %code_local_row : index + %valid_code_channel = index.cmp ult, %code_channel, %bounded_output_size : index + %loads_q4_code = scalar.andi %is_q4, %valid_code_channel : i1 + %loaded_q_word = scf.if %loads_q4_code -> (vector<1xi32>) { + %q4_group_pair = index.assume %quant_group_outer [range(%quant_group_outer, 0, 3)] : index + %code_channel_byte_add = index.scale %code_channel, %weight_row_bytes : index, offset -> offset + %code_row_byte_base = index.add %expert_byte_base, %code_channel_byte_add : offset + %q_word = func.call @qwen3_moe_q4k_wmma_load_code(%weight_noalias, %code_row_byte_base, %quant_block, %q4_group_pair, %load_packet) : (buffer, offset, index, index, index) -> (vector<1xi32>) + scf.yield %q_word : vector<1xi32> + } else { + scf.yield %c0_i32x1 : vector<1xi32> + } + %updates_q_word0 = index.cmp eq, %code_row_offset, %c0 : index + %updates_q_word1 = index.cmp eq, %code_row_offset, %c16 : index + %updates_q_word2 = index.cmp eq, %code_row_offset, %c32 : index + %updates_q_word3 = index.cmp eq, %code_row_offset, %c48 : index + %next_q_word0 = scf.select %updates_q_word0, %loaded_q_word, %prior_q_word0 : vector<1xi32> + %next_q_word1 = scf.select %updates_q_word1, %loaded_q_word, %prior_q_word1 : vector<1xi32> + %next_q_word2 = scf.select %updates_q_word2, %loaded_q_word, %prior_q_word2 : vector<1xi32> + %next_q_word3 = scf.select %updates_q_word3, %loaded_q_word, %prior_q_word3 : vector<1xi32> + scf.yield %next_q_word0, %next_q_word1, %next_q_word2, %next_q_word3 : vector<1xi32>, vector<1xi32>, vector<1xi32>, vector<1xi32> + } + // Four Q6_K groups in one half-block share a QH word and consume two + // nibbles from each of two QL words. Retain the three packed words for + // all four groups. The Q4 specialization erases this acquisition path. + %q6_ql00, %q6_ql10, %q6_qh0, %q6_ql01, %q6_ql11, %q6_qh1, %q6_ql02, %q6_ql12, %q6_qh2, %q6_ql03, %q6_ql13, %q6_qh3 = scf.for %q6_code_row_offset = [%c0 to %c64 step %c16](%prior_q6_ql00 = %c0_i32x1 : vector<1xi32>, %prior_q6_ql10 = %c0_i32x1 : vector<1xi32>, %prior_q6_qh0 = %c0_i32x1 : vector<1xi32>, %prior_q6_ql01 = %c0_i32x1 : vector<1xi32>, %prior_q6_ql11 = %c0_i32x1 : vector<1xi32>, %prior_q6_qh1 = %c0_i32x1 : vector<1xi32>, %prior_q6_ql02 = %c0_i32x1 : vector<1xi32>, %prior_q6_ql12 = %c0_i32x1 : vector<1xi32>, %prior_q6_qh2 = %c0_i32x1 : vector<1xi32>, %prior_q6_ql03 = %c0_i32x1 : vector<1xi32>, %prior_q6_ql13 = %c0_i32x1 : vector<1xi32>, %prior_q6_qh3 = %c0_i32x1 : vector<1xi32>) -> (vector<1xi32>, vector<1xi32>, vector<1xi32>, vector<1xi32>, vector<1xi32>, vector<1xi32>, vector<1xi32>, vector<1xi32>, vector<1xi32>, vector<1xi32>, vector<1xi32>, vector<1xi32>) unroll { + %q6_code_local_row0 = index.add %load_row, %q6_code_row_offset : index + %q6_code_local_row = index.assume %q6_code_local_row0 [range(%q6_code_local_row0, 0, 63)] : index + %q6_code_channel = index.add %channel_tile_base, %q6_code_local_row : index + %valid_q6_code_channel = index.cmp ult, %q6_code_channel, %bounded_output_size : index + %loads_q6_code = scalar.andi %is_q6, %valid_q6_code_channel : i1 + %loaded_q6_ql0, %loaded_q6_ql1, %loaded_q6_qh = scf.if %loads_q6_code -> (vector<1xi32>, vector<1xi32>, vector<1xi32>) { + %q6_half = index.assume %quant_group_outer [range(%quant_group_outer, 0, 1)] : index + %q6_code_channel_byte_add = index.scale %q6_code_channel, %weight_row_bytes : index, offset -> offset + %q6_code_row_byte_base = index.add %expert_byte_base, %q6_code_channel_byte_add : offset + %ql0_word, %ql1_word, %qh_word = func.call @qwen3_moe_q6k_wmma_load_half_codes(%weight_noalias, %q6_code_row_byte_base, %quant_block, %q6_half, %load_packet) : (buffer, offset, index, index, index) -> (vector<1xi32>, vector<1xi32>, vector<1xi32>) + scf.yield %ql0_word, %ql1_word, %qh_word : vector<1xi32>, vector<1xi32>, vector<1xi32> + } else { + scf.yield %c0_i32x1, %c0_i32x1, %c0_i32x1 : vector<1xi32>, vector<1xi32>, vector<1xi32> + } + %updates_q6_code0 = index.cmp eq, %q6_code_row_offset, %c0 : index + %updates_q6_code1 = index.cmp eq, %q6_code_row_offset, %c16 : index + %updates_q6_code2 = index.cmp eq, %q6_code_row_offset, %c32 : index + %updates_q6_code3 = index.cmp eq, %q6_code_row_offset, %c48 : index + %next_q6_ql00 = scf.select %updates_q6_code0, %loaded_q6_ql0, %prior_q6_ql00 : vector<1xi32> + %next_q6_ql10 = scf.select %updates_q6_code0, %loaded_q6_ql1, %prior_q6_ql10 : vector<1xi32> + %next_q6_qh0 = scf.select %updates_q6_code0, %loaded_q6_qh, %prior_q6_qh0 : vector<1xi32> + %next_q6_ql01 = scf.select %updates_q6_code1, %loaded_q6_ql0, %prior_q6_ql01 : vector<1xi32> + %next_q6_ql11 = scf.select %updates_q6_code1, %loaded_q6_ql1, %prior_q6_ql11 : vector<1xi32> + %next_q6_qh1 = scf.select %updates_q6_code1, %loaded_q6_qh, %prior_q6_qh1 : vector<1xi32> + %next_q6_ql02 = scf.select %updates_q6_code2, %loaded_q6_ql0, %prior_q6_ql02 : vector<1xi32> + %next_q6_ql12 = scf.select %updates_q6_code2, %loaded_q6_ql1, %prior_q6_ql12 : vector<1xi32> + %next_q6_qh2 = scf.select %updates_q6_code2, %loaded_q6_qh, %prior_q6_qh2 : vector<1xi32> + %next_q6_ql03 = scf.select %updates_q6_code3, %loaded_q6_ql0, %prior_q6_ql03 : vector<1xi32> + %next_q6_ql13 = scf.select %updates_q6_code3, %loaded_q6_ql1, %prior_q6_ql13 : vector<1xi32> + %next_q6_qh3 = scf.select %updates_q6_code3, %loaded_q6_qh, %prior_q6_qh3 : vector<1xi32> + scf.yield %next_q6_ql00, %next_q6_ql10, %next_q6_qh0, %next_q6_ql01, %next_q6_ql11, %next_q6_qh1, %next_q6_ql02, %next_q6_ql12, %next_q6_qh2, %next_q6_ql03, %next_q6_ql13, %next_q6_qh3 : vector<1xi32>, vector<1xi32>, vector<1xi32>, vector<1xi32>, vector<1xi32>, vector<1xi32>, vector<1xi32>, vector<1xi32>, vector<1xi32>, vector<1xi32>, vector<1xi32>, vector<1xi32> + } + %outer_result00, %outer_result01, %outer_result10, %outer_result11 = scf.for %group_within_outer = [%c0 to %groups_per_outer step %c1](%group_acc00 = %acc00 : vector<8xf16>, %group_acc01 = %acc01 : vector<8xf16>, %group_acc10 = %acc10 : vector<8xf16>, %group_acc11 = %acc11 : vector<8xf16>) -> (vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>) { + %quant_group_base = index.mul %quant_group_outer, %groups_per_outer : index + %quant_group0 = index.add %quant_group_base, %group_within_outer : index + %quant_group = index.assume %quant_group0 [range(%quant_group0, 0, 7)] : index + %block_k_base = index.mul %quant_block, %c256 : index + %group_k_add = index.mul %quant_group, %c32 : index + %k_origin = index.add %block_k_base, %group_k_add : index + scf.for %row_offset = [%c0 to %c64 step %c16] unroll { + %local_row0 = index.add %load_row, %row_offset : index + %local_row = index.assume %local_row0 [range(%local_row0, 0, 63)] : index + %selects_row0 = index.cmp eq, %row_offset, %c0 : index + %selects_row2 = index.cmp eq, %row_offset, %c32 : index + %selects_low_row_pair = index.cmp ult, %row_offset, %c32 : index + %selected_header01 = scf.select %selects_row0, %header0, %header1 : vector<4xi32> + %selected_header23 = scf.select %selects_row2, %header2, %header3 : vector<4xi32> + %selected_header = scf.select %selects_low_row_pair, %selected_header01, %selected_header23 : vector<4xi32> + %selected_q_word01 = scf.select %selects_row0, %q_word0, %q_word1 : vector<1xi32> + %selected_q_word23 = scf.select %selects_row2, %q_word2, %q_word3 : vector<1xi32> + %selected_q_word = scf.select %selects_low_row_pair, %selected_q_word01, %selected_q_word23 : vector<1xi32> + %selected_q6_ql0_01 = scf.select %selects_row0, %q6_ql00, %q6_ql01 : vector<1xi32> + %selected_q6_ql0_23 = scf.select %selects_row2, %q6_ql02, %q6_ql03 : vector<1xi32> + %selected_q6_ql0 = scf.select %selects_low_row_pair, %selected_q6_ql0_01, %selected_q6_ql0_23 : vector<1xi32> + %selected_q6_ql1_01 = scf.select %selects_row0, %q6_ql10, %q6_ql11 : vector<1xi32> + %selected_q6_ql1_23 = scf.select %selects_row2, %q6_ql12, %q6_ql13 : vector<1xi32> + %selected_q6_ql1 = scf.select %selects_low_row_pair, %selected_q6_ql1_01, %selected_q6_ql1_23 : vector<1xi32> + %selected_q6_qh01 = scf.select %selects_row0, %q6_qh0, %q6_qh1 : vector<1xi32> + %selected_q6_qh23 = scf.select %selects_row2, %q6_qh2, %q6_qh3 : vector<1xi32> + %selected_q6_qh = scf.select %selects_low_row_pair, %selected_q6_qh01, %selected_q6_qh23 : vector<1xi32> + %channel = index.add %channel_tile_base, %local_row : index + %valid_channel = index.cmp ult, %channel, %bounded_output_size : index + %weight_values = scf.if %valid_channel -> (vector<4xf16>) { + %decoded = scf.if %is_q6 -> (vector<4xf16>) { + %channel_byte_add = index.scale %channel, %weight_row_bytes : index, offset -> offset + %row_byte_base = index.add %expert_byte_base, %channel_byte_add : offset + %q6_values = func.call @qwen3_moe_q6k_wmma_vector4_from_half_codes(%weight_noalias, %row_byte_base, %quant_block, %quant_group, %load_packet, %selected_q6_ql0, %selected_q6_ql1, %selected_q6_qh) : (buffer, offset, index, index, index, vector<1xi32>, vector<1xi32>, vector<1xi32>) -> (vector<4xf16>) + scf.yield %q6_values : vector<4xf16> + } else { + %q4_values = func.call @qwen3_moe_q4k_wmma_vector4_from_header_code(%quant_group, %selected_header, %selected_q_word) : (index, vector<4xi32>, vector<1xi32>) -> (vector<4xf16>) + scf.yield %q4_values : vector<4xf16> + } + scf.yield %decoded : vector<4xf16> + } else { + scf.yield %c0_f16x4 : vector<4xf16> + } + %is_activation_row = index.cmp ult, %local_row, %c32 : index + vector.store %weight_values, %weight_stage_view[%local_row, %load_k] : vector<4xf16>, view<64x40xf16> + scf.if %is_activation_row { + %activation_row = index.assume %local_row [range(%local_row, 0, 31)] : index + %assignment_i32 = view.load %route_stage_view[%activation_row] : view<32xi32> -> i32 + %valid_assignment = scalar.cmpi sge, %assignment_i32, %c0_i32 : i32 + %activation_values = scf.if %valid_assignment -> (vector<4xf16>) { + %assignment0 = index.cast %assignment_i32 : i32 to index + %assignment = index.assume %assignment0 [range(%assignment0, 0, 16383)] : index + %bounded_assignment, %bounded_assignment_count = index.assume %assignment, %assignment_count [lt(%assignment, %assignment_count)] : index, index + %input_k = index.add %k_origin, %load_k : index + %loaded = vector.load %input_view[%bounded_assignment, %input_k] : view<[%assignment_count]x[%input_size]xf16> -> vector<4xf16> + scf.yield %loaded : vector<4xf16> + } else { + scf.yield %c0_f16x4 : vector<4xf16> + } + vector.store %activation_values, %activation_stage_physical_view[%activation_row, %load_k] : vector<4xf16>, view<32x40xf16> + } + } + kernel.barrier scope(workgroup) ordering(acq_rel) + %next00, %next01, %next10, %next11 = scf.for %k_half = [%c0 to %c32 step %c16](%half_acc00 = %group_acc00 : vector<8xf16>, %half_acc01 = %group_acc01 : vector<8xf16>, %half_acc10 = %group_acc10 : vector<8xf16>, %half_acc11 = %group_acc11 : vector<8xf16>) -> (vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>) unroll { + %lhs0 = vector.fragment.load %weight_stage_view[%subgroup_channel_add, %k_half] shape [%m, %k] : view<64x40xf16> -> vector<16xf16> + %lhs1 = vector.fragment.load %weight_stage_view[%subgroup_channel1, %k_half] shape [%m, %k] : view<64x40xf16> -> vector<16xf16> + %rhs0 = vector.fragment.load %activation_fragment_view[%k_half, %c0] shape [%k, %n] : view<32x32xf16, %activation_fragment_layout> -> vector<16xf16> + %rhs1 = vector.fragment.load %activation_fragment_view[%k_half, %c16] shape [%k, %n] : view<32x32xf16, %activation_fragment_layout> -> vector<16xf16> + %half_next00 = vector.mma %lhs0, %rhs0, %half_acc00 : vector<16xf16>, vector<16xf16>, vector<8xf16> + %half_next01 = vector.mma %lhs0, %rhs1, %half_acc01 : vector<16xf16>, vector<16xf16>, vector<8xf16> + %half_next10 = vector.mma %lhs1, %rhs0, %half_acc10 : vector<16xf16>, vector<16xf16>, vector<8xf16> + %half_next11 = vector.mma %lhs1, %rhs1, %half_acc11 : vector<16xf16>, vector<16xf16>, vector<8xf16> + scf.yield %half_next00, %half_next01, %half_next10, %half_next11 : vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + scf.yield %next00, %next01, %next10, %next11 : vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16> + } + scf.yield %outer_result00, %outer_result01, %outer_result10, %outer_result11 : vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16> + } + scf.yield %block_result00, %block_result01, %block_result10, %block_result11 : vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16> + } + // WMMA produces [channel][route] fragments. Transpose through wave-private + // LDS so each lane publishes four adjacent channels for one assignment. + %publish_route0 = index.div %lane, %c4 : index + %publish_route = index.assume %publish_route0 [range(%publish_route0, 0, 15)] : index + %publish_packet0 = index.rem %lane, %c4 : index + %publish_packet = index.assume %publish_packet0 [range(%publish_packet0, 0, 3)] : index + %publish_channel_add = index.mul %publish_packet, %c4 : index + %local_route1 = index.add %c16, %publish_route : index + %assignment0_i32 = view.load %route_stage_view[%publish_route] : view<32xi32> -> i32 + %assignment1_i32 = view.load %route_stage_view[%local_route1] : view<32xi32> -> i32 + %assignment0_nonnegative = scalar.cmpi sge, %assignment0_i32, %c0_i32 : i32 + %assignment1_nonnegative = scalar.cmpi sge, %assignment1_i32, %c0_i32 : i32 + %safe_assignment0_i32 = scf.select %assignment0_nonnegative, %assignment0_i32, %c0_i32 : i32 + %safe_assignment1_i32 = scf.select %assignment1_nonnegative, %assignment1_i32, %c0_i32 : i32 + %safe_assignment0_0 = index.cast %safe_assignment0_i32 : i32 to index + %safe_assignment1_0 = index.cast %safe_assignment1_i32 : i32 to index + %safe_assignment0 = index.assume %safe_assignment0_0 [range(%safe_assignment0_0, 0, 16383)] : index + %safe_assignment1 = index.assume %safe_assignment1_0 [range(%safe_assignment1_0, 0, 16383)] : index + %bounded_assignment0, %bounded_assignment_count0 = index.assume %safe_assignment0, %assignment_count [lt(%safe_assignment0, %assignment_count)] : index, index + %bounded_assignment1, %bounded_assignment_count1 = index.assume %safe_assignment1, %assignment_count [lt(%safe_assignment1, %assignment_count)] : index, index + %subgroup_channel_base = index.add %channel_tile_base, %subgroup_channel_add : index + %channel0 = index.add %subgroup_channel_base, %publish_channel_add : index + %channel1_base = index.add %subgroup_channel_base, %c16 : index + %channel1 = index.add %channel1_base, %publish_channel_add : index + %valid_channel0 = index.cmp ult, %channel0, %bounded_output_size : index + %valid_channel1 = index.cmp ult, %channel1, %bounded_output_size : index + %writes00 = scalar.andi %assignment0_nonnegative, %valid_channel0 : i1 + %writes01 = scalar.andi %assignment1_nonnegative, %valid_channel0 : i1 + %writes10 = scalar.andi %assignment0_nonnegative, %valid_channel1 : i1 + %writes11 = scalar.andi %assignment1_nonnegative, %valid_channel1 : i1 + vector.fragment.store %result00, %result_fragment_view[%c0, %c0] shape [%m, %n] : vector<8xf16>, view<16x16xf16, %result_fragment_layout> + kernel.barrier scope(subgroup) ordering(acq_rel) + scf.if %writes00 { + %values = vector.load %result_physical_view[%publish_route, %publish_channel_add] : view<16x16xf16> -> vector<4xf16> + %mask = vector.mask.range [%channel0 to %bounded_output_size step %c1] : index -> vector<4xi1> + vector.store.mask %values, %output_view[%bounded_assignment0, %channel0], %mask : vector<4xf16>, view<[%assignment_count]x[%bounded_output_size]xf16>, vector<4xi1> + } + kernel.barrier scope(subgroup) ordering(acq_rel) + vector.fragment.store %result01, %result_fragment_view[%c0, %c0] shape [%m, %n] : vector<8xf16>, view<16x16xf16, %result_fragment_layout> + kernel.barrier scope(subgroup) ordering(acq_rel) + scf.if %writes01 { + %values = vector.load %result_physical_view[%publish_route, %publish_channel_add] : view<16x16xf16> -> vector<4xf16> + %mask = vector.mask.range [%channel0 to %bounded_output_size step %c1] : index -> vector<4xi1> + vector.store.mask %values, %output_view[%bounded_assignment1, %channel0], %mask : vector<4xf16>, view<[%assignment_count]x[%bounded_output_size]xf16>, vector<4xi1> + } + kernel.barrier scope(subgroup) ordering(acq_rel) + vector.fragment.store %result10, %result_fragment_view[%c0, %c0] shape [%m, %n] : vector<8xf16>, view<16x16xf16, %result_fragment_layout> + kernel.barrier scope(subgroup) ordering(acq_rel) + scf.if %writes10 { + %values = vector.load %result_physical_view[%publish_route, %publish_channel_add] : view<16x16xf16> -> vector<4xf16> + %mask = vector.mask.range [%channel1 to %bounded_output_size step %c1] : index -> vector<4xi1> + vector.store.mask %values, %output_view[%bounded_assignment0, %channel1], %mask : vector<4xf16>, view<[%assignment_count]x[%bounded_output_size]xf16>, vector<4xi1> + } + kernel.barrier scope(subgroup) ordering(acq_rel) + vector.fragment.store %result11, %result_fragment_view[%c0, %c0] shape [%m, %n] : vector<8xf16>, view<16x16xf16, %result_fragment_layout> + kernel.barrier scope(subgroup) ordering(acq_rel) + scf.if %writes11 { + %values = vector.load %result_physical_view[%publish_route, %publish_channel_add] : view<16x16xf16> -> vector<4xf16> + %mask = vector.mask.range [%channel1 to %bounded_output_size step %c1] : index -> vector<4xi1> + vector.store.mask %values, %output_view[%bounded_assignment1, %channel1], %mask : vector<4xf16>, view<[%assignment_count]x[%bounded_output_size]xf16>, vector<4xi1> + } + // Route, operand, and result stages are reused by the next concentrated + // routing partition. All waves must finish publication before reuse. + kernel.barrier scope(workgroup) ordering(acq_rel) + } + func.return +} + +// Q4_K and Q6_K retain separate entry points while sharing the complete matrix +// schedule. This keeps format routing outside the hot kernel. +kernel.def target(@qwen3_moe_routed_down_gfx11_wave64) @qwen3_moe_routed_down_q4k_f16_wmma_grouped(%token_count: index) { + %expert_count = config.get @qwen3_moe.routed_down.expert_count : index + %output_size = config.get @qwen3_moe.routed_down.output_size : index + %c1 = index.constant 1 : index + %c63 = index.constant 63 : index + %c64 = index.constant 64 : index + %c128 = index.constant 128 : index + %padded_output_size = index.add %output_size, %c63 : index + %output_tiles = index.div %padded_output_size, %c64 : index + %padded_token_count = index.add %token_count, %c63 : index + %route_tiles = index.div %padded_token_count, %c64 : index + kernel.launch.config workgroups(%output_tiles, %route_tiles, %expert_count) workgroup_size(%c128, %c1, %c1) : index +} launch(%token_count: index, %input: buffer, %expert_table: buffer, %weight: buffer, %output: buffer) where [range(%token_count, 1, 2048)] { + %q4 = index.constant 4 : index + func.call @qwen3_moe_routed_down_quantized_f16_wmma_body(%q4, %token_count, %input, %expert_table, %weight, %output) : (index, index, buffer, buffer, buffer, buffer) + kernel.return +} + +kernel.def target(@qwen3_moe_routed_down_gfx11_wave64) @qwen3_moe_routed_down_q6k_f16_wmma_grouped(%token_count: index) { + %expert_count = config.get @qwen3_moe.routed_down.expert_count : index + %output_size = config.get @qwen3_moe.routed_down.output_size : index + %c1 = index.constant 1 : index + %c63 = index.constant 63 : index + %c64 = index.constant 64 : index + %c128 = index.constant 128 : index + %padded_output_size = index.add %output_size, %c63 : index + %output_tiles = index.div %padded_output_size, %c64 : index + %padded_token_count = index.add %token_count, %c63 : index + %route_tiles = index.div %padded_token_count, %c64 : index + kernel.launch.config workgroups(%output_tiles, %route_tiles, %expert_count) workgroup_size(%c128, %c1, %c1) : index +} launch(%token_count: index, %input: buffer, %expert_table: buffer, %weight: buffer, %output: buffer) where [range(%token_count, 1, 2048)] { + %q6 = index.constant 6 : index + func.call @qwen3_moe_routed_down_quantized_f16_wmma_body(%q6, %token_count, %input, %expert_table, %weight, %output) : (index, index, buffer, buffer, buffer, buffer) + kernel.return +} + +// Reduces compact routed projections into the residual in logical token order. +// +// One workitem owns four adjacent output channels. All top-k routes are +// accumulated in FP32 before the residual is loaded and published once. The +// preceding grouped projection is the only producer of each FP16 route row, so +// this boundary requires no global atomics. +kernel.def target(@qwen3_moe_routed_down_gfx11_wave64) @qwen3_moe_routed_down_weighted_reduce_f16_f32(%token_count: index) { + %output_size = config.get @qwen3_moe.routed_down.output_size : index + %c1 = index.constant 1 : index + %c64 = index.constant 64 : index + %c256 = index.constant 256 : index + %padded_output_size = index.add %output_size, %c256 : index + %rounded_output_size = index.sub %padded_output_size, %c1 : index + %output_tiles = index.div %rounded_output_size, %c256 : index + kernel.launch.config workgroups(%output_tiles, %token_count, %c1) workgroup_size(%c64, %c1, %c1) : index +} launch(%token_count: index, %route_weights: buffer, %routed_output: buffer, %output: buffer) where [range(%token_count, 1, 2048)] { + %route_count = config.get @qwen3_moe.routed_down.route_count : index + %output_size = config.get @qwen3_moe.routed_down.output_size : index + %bounded_route_count = index.assume %route_count [range(%route_count, 1, 8)] : index + %bounded_output_size = index.assume %output_size [range(%output_size, 1, 4096)] : index + %channel_tile = kernel.workgroup.id : index + %token0 = kernel.workgroup.id : index + %lane = kernel.workitem.id : index + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c4 = index.constant 4 : index + %c256 = index.constant 256 : index + %c0_offset = index.constant 0 : offset + %c0_f16x4 = vector.constant 0.0 : vector<4xf16> + %c0_f32x4 = vector.constant 0.0 : vector<4xf32> + %assignment_count = index.mul %token_count, %bounded_route_count : index + %channel_tile_base = index.mul %channel_tile, %c256 : index + %lane_channel_add = index.mul %lane, %c4 : index + %channel = index.add %channel_tile_base, %lane_channel_add : index + %active_token = index.cmp ult, %token0, %token_count : index + %valid_channel = index.cmp ult, %channel, %bounded_output_size : index + %publishes_output = scalar.andi %active_token, %valid_channel : i1 + %route_weights_noalias, %routed_output_noalias, %output_noalias = buffer.assume.noalias %route_weights, %routed_output, %output : buffer, buffer, buffer + %route_weight_view = buffer.view %route_weights_noalias[%c0_offset] : buffer -> view<[%token_count]x[%bounded_route_count]xf32> + %routed_output_view = buffer.view %routed_output_noalias[%c0_offset] : buffer -> view<[%assignment_count]x[%bounded_output_size]xf16> + %output_view = buffer.view %output_noalias[%c0_offset] : buffer -> view<[%token_count]x[%bounded_output_size]xf32> + scf.if %publishes_output { + %token, %output_token_count = index.assume %token0, %token_count [lt(%token0, %token_count)] : index, index + %mask = vector.mask.range [%channel to %bounded_output_size step %c1] : index -> vector<4xi1> + %weighted_sum = scf.for %route = [%c0 to %bounded_route_count step %c1](%sum = %c0_f32x4 : vector<4xf32>) -> (vector<4xf32>) unroll { + %assignment = index.madd %token, %bounded_route_count, %route : index + %routed = vector.load.mask %routed_output_view[%assignment, %channel], %mask, %c0_f16x4 : view<[%assignment_count]x[%bounded_output_size]xf16>, vector<4xi1>, vector<4xf16> + %wide = vector.extf %routed : vector<4xf16> to vector<4xf32> + %route_weight = view.load %route_weight_view[%token, %route] : view<[%token_count]x[%bounded_route_count]xf32> -> f32 + %route_weight_x4 = vector.splat %route_weight : vector<4xf32> + %weighted = vector.mulf %wide, %route_weight_x4 : vector<4xf32> + %next = vector.addf %sum, %weighted : vector<4xf32> + scf.yield %next : vector<4xf32> + } + %residual = vector.load.mask %output_view[%token, %channel], %mask, %c0_f32x4 : view<[%token_count]x[%bounded_output_size]xf32>, vector<4xi1>, vector<4xf32> + %result = vector.addf %residual, %weighted_sum : vector<4xf32> + vector.store.mask %result, %output_view[%token, %channel], %mask : vector<4xf32>, view<[%token_count]x[%bounded_output_size]xf32>, vector<4xi1> + } + kernel.return +} + +// Complete expert-MLP coverage crosses route partitions and output tails. The +// established Q8_1 gate/up and Q4_K down providers form an independent +// reference while the production path keeps the routed SwiGLU handoff in F16. +// The complete-chain tolerance covers both independent quantized accumulation +// regimes; the focused gate/up and down cases constrain each boundary to 1%. +check.case public @qwen3_moe_expert_mlp_q4k_f16_handoff_differential_case { + %token_count = check.literal value(67) : index + %gate_input_size = check.literal value(512) : index + %gate_output_size = check.literal value(256) : index + %down_output_size = check.literal value(33) : index + %routed_row_count = check.literal value(134) : index + %route_count = check.literal value(2) : index + %route_stride = check.literal value(4) : index + %expert_count = check.literal value(4) : index + %input = check.generate.fill value(0.00390625) : tensor<67x512xf32> + %q8_input = check.generate.fill value(0) : tensor<67x576xi8> + %route_ids = check.generate.iota offset(0) step(1) period(4) : tensor<67x4xi32> + %route_weights = check.generate.fill value(0.5) : tensor<67x2xf32> + %expert_table = check.generate.fill value(-1) : tensor<272xi32> + %partition_table = check.generate.fill value(-1) : tensor<10xi32> + %gate_weight = check.generate.fill value(34) : tensor<4x256x2x144xi8> + %up_weight = check.generate.fill value(35) : tensor<4x256x2x144xi8> + %down_weight = check.generate.fill value(-86) : tensor<4x33x1x144xi8> + %reference_gate_output = check.generate.fill value(1.0) : tensor<67x2x256xf32> + %q8_down_input = check.generate.fill value(0) : tensor<67x2x288xi8> + %expected_output = check.generate.fill value(1.0) : tensor<67x33xf32> + %actual_gate_output = check.generate.fill value(1.0) : tensor<67x2x256xf16> + %actual_routed_output = check.generate.fill value(0.0) : tensor<67x2x33xf16> + %actual_output = check.generate.fill value(1.0) : tensor<67x33xf32> + kernel.launch @ggml_quantize_q8_1_x4_f32[%token_count, %gate_input_size](%token_count, %gate_input_size, %input, %q8_input) : [index, index](index, index, tensor<67x512xf32>, tensor<67x576xi8>) + kernel.launch @qwen3_moe_routed_gate_up_swiglu_q4k_q8[%token_count, %route_count, %route_stride, %expert_count, %gate_output_size](%token_count, %route_count, %route_stride, %expert_count, %gate_output_size, %q8_input, %route_ids, %gate_weight, %up_weight, %reference_gate_output) : [index, index, index, index, index](index, index, index, index, index, tensor<67x576xi8>, tensor<67x4xi32>, tensor<4x256x2x144xi8>, tensor<4x256x2x144xi8>, tensor<67x2x256xf32>) + kernel.launch @ggml_quantize_q8_1_x4_f32[%routed_row_count, %gate_output_size](%routed_row_count, %gate_output_size, %reference_gate_output, %q8_down_input) : [index, index](index, index, tensor<67x2x256xf32>, tensor<67x2x288xi8>) + kernel.launch @qwen3_moe_routed_down_q4k_q8_1_x4[%token_count, %gate_output_size, %route_count, %route_stride, %expert_count, %down_output_size](%token_count, %gate_output_size, %route_count, %route_stride, %expert_count, %down_output_size, %q8_down_input, %route_ids, %route_weights, %down_weight, %expected_output) : [index, index, index, index, index, index](index, index, index, index, index, index, tensor<67x2x288xi8>, tensor<67x4xi32>, tensor<67x2xf32>, tensor<4x33x1x144xi8>, tensor<67x33xf32>) + kernel.launch @qwen3_moe_build_expert_table[%token_count, %route_count, %route_stride, %expert_count](%token_count, %route_count, %route_stride, %expert_count, %route_ids, %expert_table) : [index, index, index, index](index, index, index, index, tensor<67x4xi32>, tensor<272xi32>) + kernel.launch @qwen3_moe_build_expert_partition_table[%token_count, %route_count, %expert_count](%token_count, %route_count, %expert_count, %expert_table, %partition_table) : [index, index, index](index, index, index, tensor<272xi32>, tensor<10xi32>) + kernel.launch @qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma[%token_count](%token_count, %input, %expert_table, %partition_table, %gate_weight, %up_weight, %actual_gate_output) : [index](index, tensor<67x512xf32>, tensor<272xi32>, tensor<10xi32>, tensor<4x256x2x144xi8>, tensor<4x256x2x144xi8>, tensor<67x2x256xf16>) + kernel.launch @qwen3_moe_routed_down_q4k_f16_wmma_grouped[%token_count](%token_count, %actual_gate_output, %expert_table, %down_weight, %actual_routed_output) : [index](index, tensor<67x2x256xf16>, tensor<272xi32>, tensor<4x33x1x144xi8>, tensor<67x2x33xf16>) + kernel.launch @qwen3_moe_routed_down_weighted_reduce_f16_f32[%token_count](%token_count, %route_weights, %actual_routed_output, %actual_output) : [index](index, tensor<67x2xf32>, tensor<67x2x33xf16>, tensor<67x33xf32>) + check.expect.close actual(%actual_output) expected(%expected_output) atol(0.5) rtol(0.05) nan(same) : tensor<67x33xf32> + check.return +} + +// Q6_K shares the production matrix schedule but has an independent packed +// decoder and Q8_1 reference provider. +check.case public @qwen3_moe_expert_mlp_q6k_f16_handoff_differential_case { + %token_count = check.literal value(67) : index + %gate_input_size = check.literal value(512) : index + %gate_output_size = check.literal value(256) : index + %down_output_size = check.literal value(33) : index + %routed_row_count = check.literal value(134) : index + %route_count = check.literal value(2) : index + %route_stride = check.literal value(4) : index + %expert_count = check.literal value(4) : index + %input = check.generate.fill value(0.00390625) : tensor<67x512xf32> + %q8_input = check.generate.fill value(0) : tensor<67x576xi8> + %route_ids = check.generate.iota offset(0) step(1) period(4) : tensor<67x4xi32> + %route_weights = check.generate.fill value(0.5) : tensor<67x2xf32> + %expert_table = check.generate.fill value(-1) : tensor<272xi32> + %partition_table = check.generate.fill value(-1) : tensor<10xi32> + %gate_weight = check.generate.fill value(34) : tensor<4x256x2x144xi8> + %up_weight = check.generate.fill value(35) : tensor<4x256x2x144xi8> + %down_weight = check.generate.fill value(-86) : tensor<4x33x1x210xi8> + %reference_gate_output = check.generate.fill value(1.0) : tensor<67x2x256xf32> + %q8_down_input = check.generate.fill value(0) : tensor<67x2x288xi8> + %expected_output = check.generate.fill value(1.0) : tensor<67x33xf32> + %actual_gate_output = check.generate.fill value(1.0) : tensor<67x2x256xf16> + %actual_routed_output = check.generate.fill value(0.0) : tensor<67x2x33xf16> + %actual_output = check.generate.fill value(1.0) : tensor<67x33xf32> + kernel.launch @ggml_quantize_q8_1_x4_f32[%token_count, %gate_input_size](%token_count, %gate_input_size, %input, %q8_input) : [index, index](index, index, tensor<67x512xf32>, tensor<67x576xi8>) + kernel.launch @qwen3_moe_routed_gate_up_swiglu_q4k_q8[%token_count, %route_count, %route_stride, %expert_count, %gate_output_size](%token_count, %route_count, %route_stride, %expert_count, %gate_output_size, %q8_input, %route_ids, %gate_weight, %up_weight, %reference_gate_output) : [index, index, index, index, index](index, index, index, index, index, tensor<67x576xi8>, tensor<67x4xi32>, tensor<4x256x2x144xi8>, tensor<4x256x2x144xi8>, tensor<67x2x256xf32>) + kernel.launch @ggml_quantize_q8_1_x4_f32[%routed_row_count, %gate_output_size](%routed_row_count, %gate_output_size, %reference_gate_output, %q8_down_input) : [index, index](index, index, tensor<67x2x256xf32>, tensor<67x2x288xi8>) + kernel.launch @qwen3_moe_routed_down_q6k_q8_1_x4[%token_count, %gate_output_size, %route_count, %route_stride, %expert_count, %down_output_size](%token_count, %gate_output_size, %route_count, %route_stride, %expert_count, %down_output_size, %q8_down_input, %route_ids, %route_weights, %down_weight, %expected_output) : [index, index, index, index, index, index](index, index, index, index, index, index, tensor<67x2x288xi8>, tensor<67x4xi32>, tensor<67x2xf32>, tensor<4x33x1x210xi8>, tensor<67x33xf32>) + kernel.launch @qwen3_moe_build_expert_table[%token_count, %route_count, %route_stride, %expert_count](%token_count, %route_count, %route_stride, %expert_count, %route_ids, %expert_table) : [index, index, index, index](index, index, index, index, tensor<67x4xi32>, tensor<272xi32>) + kernel.launch @qwen3_moe_build_expert_partition_table[%token_count, %route_count, %expert_count](%token_count, %route_count, %expert_count, %expert_table, %partition_table) : [index, index, index](index, index, index, tensor<272xi32>, tensor<10xi32>) + kernel.launch @qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma[%token_count](%token_count, %input, %expert_table, %partition_table, %gate_weight, %up_weight, %actual_gate_output) : [index](index, tensor<67x512xf32>, tensor<272xi32>, tensor<10xi32>, tensor<4x256x2x144xi8>, tensor<4x256x2x144xi8>, tensor<67x2x256xf16>) + kernel.launch @qwen3_moe_routed_down_q6k_f16_wmma_grouped[%token_count](%token_count, %actual_gate_output, %expert_table, %down_weight, %actual_routed_output) : [index](index, tensor<67x2x256xf16>, tensor<272xi32>, tensor<4x33x1x210xi8>, tensor<67x2x33xf16>) + kernel.launch @qwen3_moe_routed_down_weighted_reduce_f16_f32[%token_count](%token_count, %route_weights, %actual_routed_output, %actual_output) : [index](index, tensor<67x2xf32>, tensor<67x2x33xf16>, tensor<67x33xf32>) + check.expect.close actual(%actual_output) expected(%expected_output) atol(0.5) rtol(0.05) nan(same) : tensor<67x33xf32> + check.return +} + +// Differential coverage crosses token, route, expert, channel, and output-tail +// boundaries. The direct Q8_1 provider supplies an independent raw-Q4_K +// reference while the grouped provider exercises collision-free FP16 +// publication and a separate weighted FP32 residual reduction. +check.case public @qwen3_moe_routed_down_q4k_f16_wmma_differential_case { + %token_count = check.literal value(67) : index + %input_size = check.literal value(512) : index + %route_count = check.literal value(2) : index + %route_stride = check.literal value(4) : index + %expert_count = check.literal value(4) : index + %output_size = check.literal value(33) : index + %routed_row_count = check.literal value(134) : index + %input_f32 = check.generate.fill value(0.00390625) : tensor<67x2x512xf32> + %input_f16 = check.generate.fill value(0.00390625) : tensor<67x2x512xf16> + %q8_input = check.generate.fill value(0) : tensor<67x2x576xi8> + %route_ids = check.generate.iota offset(0) step(1) period(4) : tensor<67x4xi32> + %route_weights = check.generate.fill value(0.5) : tensor<67x2xf32> + %expert_table = check.generate.fill value(-1) : tensor<272xi32> + %weight = check.generate.fill value(-86) : tensor<4x33x2x144xi8> + %expected = check.generate.fill value(1.0) : tensor<67x33xf32> + %routed_output = check.generate.fill value(0.0) : tensor<67x2x33xf16> + %actual = check.generate.fill value(1.0) : tensor<67x33xf32> + kernel.launch @ggml_quantize_q8_1_x4_f32[%routed_row_count, %input_size](%routed_row_count, %input_size, %input_f32, %q8_input) : [index, index](index, index, tensor<67x2x512xf32>, tensor<67x2x576xi8>) + kernel.launch @qwen3_moe_routed_down_q4k_q8_1_x4[%token_count, %input_size, %route_count, %route_stride, %expert_count, %output_size](%token_count, %input_size, %route_count, %route_stride, %expert_count, %output_size, %q8_input, %route_ids, %route_weights, %weight, %expected) : [index, index, index, index, index, index](index, index, index, index, index, index, tensor<67x2x576xi8>, tensor<67x4xi32>, tensor<67x2xf32>, tensor<4x33x2x144xi8>, tensor<67x33xf32>) + kernel.launch @qwen3_moe_build_expert_table[%token_count, %route_count, %route_stride, %expert_count](%token_count, %route_count, %route_stride, %expert_count, %route_ids, %expert_table) : [index, index, index, index](index, index, index, index, tensor<67x4xi32>, tensor<272xi32>) + kernel.launch @qwen3_moe_routed_down_q4k_f16_wmma_grouped[%token_count](%token_count, %input_f16, %expert_table, %weight, %routed_output) : [index](index, tensor<67x2x512xf16>, tensor<272xi32>, tensor<4x33x2x144xi8>, tensor<67x2x33xf16>) + kernel.launch @qwen3_moe_routed_down_weighted_reduce_f16_f32[%token_count](%token_count, %route_weights, %routed_output, %actual) : [index](index, tensor<67x2xf32>, tensor<67x2x33xf16>, tensor<67x33xf32>) + check.expect.close actual(%actual) expected(%expected) atol(0.5) rtol(0.01) nan(same) : tensor<67x33xf32> + check.return +} + +// Q6_K follows the identical routed matrix schedule while selecting its +// independent packed-row decoder. The same irregular shape prevents format +// specialization from weakening routing or tail coverage. +check.case public @qwen3_moe_routed_down_q6k_f16_wmma_differential_case { + %token_count = check.literal value(67) : index + %input_size = check.literal value(512) : index + %route_count = check.literal value(2) : index + %route_stride = check.literal value(4) : index + %expert_count = check.literal value(4) : index + %output_size = check.literal value(33) : index + %routed_row_count = check.literal value(134) : index + %input_f32 = check.generate.fill value(0.00390625) : tensor<67x2x512xf32> + %input_f16 = check.generate.fill value(0.00390625) : tensor<67x2x512xf16> + %q8_input = check.generate.fill value(0) : tensor<67x2x576xi8> + %route_ids = check.generate.iota offset(0) step(1) period(4) : tensor<67x4xi32> + %route_weights = check.generate.fill value(0.5) : tensor<67x2xf32> + %expert_table = check.generate.fill value(-1) : tensor<272xi32> + %weight = check.generate.fill value(-86) : tensor<4x33x2x210xi8> + %expected = check.generate.fill value(1.0) : tensor<67x33xf32> + %routed_output = check.generate.fill value(0.0) : tensor<67x2x33xf16> + %actual = check.generate.fill value(1.0) : tensor<67x33xf32> + kernel.launch @ggml_quantize_q8_1_x4_f32[%routed_row_count, %input_size](%routed_row_count, %input_size, %input_f32, %q8_input) : [index, index](index, index, tensor<67x2x512xf32>, tensor<67x2x576xi8>) + kernel.launch @qwen3_moe_routed_down_q6k_q8_1_x4[%token_count, %input_size, %route_count, %route_stride, %expert_count, %output_size](%token_count, %input_size, %route_count, %route_stride, %expert_count, %output_size, %q8_input, %route_ids, %route_weights, %weight, %expected) : [index, index, index, index, index, index](index, index, index, index, index, index, tensor<67x2x576xi8>, tensor<67x4xi32>, tensor<67x2xf32>, tensor<4x33x2x210xi8>, tensor<67x33xf32>) + kernel.launch @qwen3_moe_build_expert_table[%token_count, %route_count, %route_stride, %expert_count](%token_count, %route_count, %route_stride, %expert_count, %route_ids, %expert_table) : [index, index, index, index](index, index, index, index, tensor<67x4xi32>, tensor<272xi32>) + kernel.launch @qwen3_moe_routed_down_q6k_f16_wmma_grouped[%token_count](%token_count, %input_f16, %expert_table, %weight, %routed_output) : [index](index, tensor<67x2x512xf16>, tensor<272xi32>, tensor<4x33x2x210xi8>, tensor<67x2x33xf16>) + kernel.launch @qwen3_moe_routed_down_weighted_reduce_f16_f32[%token_count](%token_count, %route_weights, %routed_output, %actual) : [index](index, tensor<67x2xf32>, tensor<67x2x33xf16>, tensor<67x33xf32>) + check.expect.close actual(%actual) expected(%expected) atol(0.5) rtol(0.01) nan(same) : tensor<67x33xf32> + check.return +} + +check.case public @qwen3_moe_routed_down_q4k_f16_wmma_benchmark_case { + %token_count = check.param.choice values([1, 2, 4, 8, 16, 17, 32, 63, 128, 129, 512]) name("token_count") : index + %route_count = check.literal value(8) : index + %route_stride = check.literal value(128) : index + %expert_count = check.literal value(128) : index + %route_ids = check.generate.iota offset(0) step(1) period(127) : tensor<[%token_count]x128xi32> + %route_weights = check.generate.fill value(0.125) : tensor<[%token_count]x8xf32> + %expert_table = check.generate.fill value(-1) : tensor<65664xi32> + %input = check.generate.fill value(0.0) : tensor<[%token_count]x8x768xf16> + %weight = check.generate.fill value(0) : tensor<128x2048x3x144xi8> + %routed_output = check.generate.fill value(0.0) : tensor<[%token_count]x8x2048xf16> + %output = check.generate.fill value(1.0) : tensor<[%token_count]x2048xf32> + %expected = check.generate.fill value(1.0) : tensor<[%token_count]x2048xf32> + kernel.launch @qwen3_moe_build_expert_table[%token_count, %route_count, %route_stride, %expert_count](%token_count, %route_count, %route_stride, %expert_count, %route_ids, %expert_table) : [index, index, index, index](index, index, index, index, tensor<[%token_count]x128xi32>, tensor<65664xi32>) + kernel.launch @qwen3_moe_routed_down_q4k_f16_wmma_grouped[%token_count](%token_count, %input, %expert_table, %weight, %routed_output) : [index](index, tensor<[%token_count]x8x768xf16>, tensor<65664xi32>, tensor<128x2048x3x144xi8>, tensor<[%token_count]x8x2048xf16>) + kernel.launch @qwen3_moe_routed_down_weighted_reduce_f16_f32[%token_count](%token_count, %route_weights, %routed_output, %output) : [index](index, tensor<[%token_count]x8xf32>, tensor<[%token_count]x8x2048xf16>, tensor<[%token_count]x2048xf32>) + check.expect.close actual(%output) expected(%expected) atol(0.0) rtol(0.0) nan(same) : tensor<[%token_count]x2048xf32> + check.return +} + +check.case public @qwen3_moe_routed_down_q6k_f16_wmma_benchmark_case { + %token_count = check.param.choice values([1, 2, 4, 8, 16, 17, 32, 63, 128, 129, 512]) name("token_count") : index + %route_count = check.literal value(8) : index + %route_stride = check.literal value(128) : index + %expert_count = check.literal value(128) : index + %route_ids = check.generate.iota offset(0) step(1) period(127) : tensor<[%token_count]x128xi32> + %route_weights = check.generate.fill value(0.125) : tensor<[%token_count]x8xf32> + %expert_table = check.generate.fill value(-1) : tensor<65664xi32> + %input = check.generate.fill value(0.0) : tensor<[%token_count]x8x768xf16> + %weight = check.generate.fill value(0) : tensor<128x2048x3x210xi8> + %routed_output = check.generate.fill value(0.0) : tensor<[%token_count]x8x2048xf16> + %output = check.generate.fill value(1.0) : tensor<[%token_count]x2048xf32> + %expected = check.generate.fill value(1.0) : tensor<[%token_count]x2048xf32> + kernel.launch @qwen3_moe_build_expert_table[%token_count, %route_count, %route_stride, %expert_count](%token_count, %route_count, %route_stride, %expert_count, %route_ids, %expert_table) : [index, index, index, index](index, index, index, index, tensor<[%token_count]x128xi32>, tensor<65664xi32>) + kernel.launch @qwen3_moe_routed_down_q6k_f16_wmma_grouped[%token_count](%token_count, %input, %expert_table, %weight, %routed_output) : [index](index, tensor<[%token_count]x8x768xf16>, tensor<65664xi32>, tensor<128x2048x3x210xi8>, tensor<[%token_count]x8x2048xf16>) + kernel.launch @qwen3_moe_routed_down_weighted_reduce_f16_f32[%token_count](%token_count, %route_weights, %routed_output, %output) : [index](index, tensor<[%token_count]x8xf32>, tensor<[%token_count]x8x2048xf16>, tensor<[%token_count]x2048xf32>) + check.expect.close actual(%output) expected(%expected) atol(0.0) rtol(0.0) nan(same) : tensor<[%token_count]x2048xf32> + check.return +} + +// Maximally diverse decode-batch controls. Compact route storage makes the +// flattened 0..127 iota assign every M=16 route to a different expert. +check.case public @qwen3_moe_routed_down_q4k_f16_wmma_diverse_benchmark_case { + %token_count = check.param.choice values([1, 2, 4, 8, 16]) name("token_count") : index + %route_count = check.literal value(8) : index + %route_stride = check.literal value(8) : index + %expert_count = check.literal value(128) : index + %route_ids = check.generate.iota offset(0) step(1) period(128) : tensor<[%token_count]x8xi32> + %route_weights = check.generate.fill value(0.125) : tensor<[%token_count]x8xf32> + %expert_table = check.generate.fill value(-1) : tensor<65664xi32> + %input = check.generate.fill value(0.0) : tensor<[%token_count]x8x768xf16> + %weight = check.generate.fill value(0) : tensor<128x2048x3x144xi8> + %routed_output = check.generate.fill value(0.0) : tensor<[%token_count]x8x2048xf16> + %output = check.generate.fill value(1.0) : tensor<[%token_count]x2048xf32> + %expected = check.generate.fill value(1.0) : tensor<[%token_count]x2048xf32> + kernel.launch @qwen3_moe_build_expert_table[%token_count, %route_count, %route_stride, %expert_count](%token_count, %route_count, %route_stride, %expert_count, %route_ids, %expert_table) : [index, index, index, index](index, index, index, index, tensor<[%token_count]x8xi32>, tensor<65664xi32>) + kernel.launch @qwen3_moe_routed_down_q4k_f16_wmma_grouped[%token_count](%token_count, %input, %expert_table, %weight, %routed_output) : [index](index, tensor<[%token_count]x8x768xf16>, tensor<65664xi32>, tensor<128x2048x3x144xi8>, tensor<[%token_count]x8x2048xf16>) + kernel.launch @qwen3_moe_routed_down_weighted_reduce_f16_f32[%token_count](%token_count, %route_weights, %routed_output, %output) : [index](index, tensor<[%token_count]x8xf32>, tensor<[%token_count]x8x2048xf16>, tensor<[%token_count]x2048xf32>) + check.expect.close actual(%output) expected(%expected) atol(0.0) rtol(0.0) nan(same) : tensor<[%token_count]x2048xf32> + check.return +} + +check.case public @qwen3_moe_routed_down_q6k_f16_wmma_diverse_benchmark_case { + %token_count = check.param.choice values([1, 2, 4, 8, 16]) name("token_count") : index + %route_count = check.literal value(8) : index + %route_stride = check.literal value(8) : index + %expert_count = check.literal value(128) : index + %route_ids = check.generate.iota offset(0) step(1) period(128) : tensor<[%token_count]x8xi32> + %route_weights = check.generate.fill value(0.125) : tensor<[%token_count]x8xf32> + %expert_table = check.generate.fill value(-1) : tensor<65664xi32> + %input = check.generate.fill value(0.0) : tensor<[%token_count]x8x768xf16> + %weight = check.generate.fill value(0) : tensor<128x2048x3x210xi8> + %routed_output = check.generate.fill value(0.0) : tensor<[%token_count]x8x2048xf16> + %output = check.generate.fill value(1.0) : tensor<[%token_count]x2048xf32> + %expected = check.generate.fill value(1.0) : tensor<[%token_count]x2048xf32> + kernel.launch @qwen3_moe_build_expert_table[%token_count, %route_count, %route_stride, %expert_count](%token_count, %route_count, %route_stride, %expert_count, %route_ids, %expert_table) : [index, index, index, index](index, index, index, index, tensor<[%token_count]x8xi32>, tensor<65664xi32>) + kernel.launch @qwen3_moe_routed_down_q6k_f16_wmma_grouped[%token_count](%token_count, %input, %expert_table, %weight, %routed_output) : [index](index, tensor<[%token_count]x8x768xf16>, tensor<65664xi32>, tensor<128x2048x3x210xi8>, tensor<[%token_count]x8x2048xf16>) + kernel.launch @qwen3_moe_routed_down_weighted_reduce_f16_f32[%token_count](%token_count, %route_weights, %routed_output, %output) : [index](index, tensor<[%token_count]x8xf32>, tensor<[%token_count]x8x2048xf16>, tensor<[%token_count]x2048xf32>) + check.expect.close actual(%output) expected(%expected) atol(0.0) rtol(0.0) nan(same) : tensor<[%token_count]x2048xf32> + check.return +} + +check.benchmark<@qwen3_moe_routed_down_q4k_f16_wmma_differential_case> @qwen3_moe_routed_down_q4k_f16_wmma_small + +check.benchmark<@qwen3_moe_routed_down_q4k_f16_wmma_benchmark_case> @qwen3_moe_routed_down_q4k_f16_wmma_decode {token_count = 1} + +check.benchmark<@qwen3_moe_routed_down_q4k_f16_wmma_benchmark_case> @qwen3_moe_routed_down_q4k_f16_wmma_small_batch_2 {token_count = 2} + +check.benchmark<@qwen3_moe_routed_down_q4k_f16_wmma_benchmark_case> @qwen3_moe_routed_down_q4k_f16_wmma_small_batch_4 {token_count = 4} + +check.benchmark<@qwen3_moe_routed_down_q4k_f16_wmma_benchmark_case> @qwen3_moe_routed_down_q4k_f16_wmma_small_batch_8 {token_count = 8} + +check.benchmark<@qwen3_moe_routed_down_q4k_f16_wmma_benchmark_case> @qwen3_moe_routed_down_q4k_f16_wmma_small_batch_16 {token_count = 16} + +check.benchmark<@qwen3_moe_routed_down_q4k_f16_wmma_diverse_benchmark_case> @qwen3_moe_routed_down_q4k_f16_wmma_diverse_decode {token_count = 1} + +check.benchmark<@qwen3_moe_routed_down_q4k_f16_wmma_diverse_benchmark_case> @qwen3_moe_routed_down_q4k_f16_wmma_diverse_small_batch_2 {token_count = 2} + +check.benchmark<@qwen3_moe_routed_down_q4k_f16_wmma_diverse_benchmark_case> @qwen3_moe_routed_down_q4k_f16_wmma_diverse_small_batch_4 {token_count = 4} + +check.benchmark<@qwen3_moe_routed_down_q4k_f16_wmma_diverse_benchmark_case> @qwen3_moe_routed_down_q4k_f16_wmma_diverse_small_batch_8 {token_count = 8} + +check.benchmark<@qwen3_moe_routed_down_q4k_f16_wmma_diverse_benchmark_case> @qwen3_moe_routed_down_q4k_f16_wmma_diverse_small_batch_16 {token_count = 16} + +check.benchmark<@qwen3_moe_routed_down_q4k_f16_wmma_benchmark_case> @qwen3_moe_routed_down_q4k_f16_wmma_prefill_32 {token_count = 32} + +check.benchmark<@qwen3_moe_routed_down_q4k_f16_wmma_benchmark_case> @qwen3_moe_routed_down_q4k_f16_wmma_prefill_128 {token_count = 128} + +check.benchmark<@qwen3_moe_routed_down_q4k_f16_wmma_benchmark_case> @qwen3_moe_routed_down_q4k_f16_wmma_prefill_512 {token_count = 512} + +check.benchmark<@qwen3_moe_routed_down_q6k_f16_wmma_differential_case> @qwen3_moe_routed_down_q6k_f16_wmma_small + +check.benchmark<@qwen3_moe_routed_down_q6k_f16_wmma_benchmark_case> @qwen3_moe_routed_down_q6k_f16_wmma_decode {token_count = 1} + +check.benchmark<@qwen3_moe_routed_down_q6k_f16_wmma_benchmark_case> @qwen3_moe_routed_down_q6k_f16_wmma_small_batch_2 {token_count = 2} + +check.benchmark<@qwen3_moe_routed_down_q6k_f16_wmma_benchmark_case> @qwen3_moe_routed_down_q6k_f16_wmma_small_batch_4 {token_count = 4} + +check.benchmark<@qwen3_moe_routed_down_q6k_f16_wmma_benchmark_case> @qwen3_moe_routed_down_q6k_f16_wmma_small_batch_8 {token_count = 8} + +check.benchmark<@qwen3_moe_routed_down_q6k_f16_wmma_benchmark_case> @qwen3_moe_routed_down_q6k_f16_wmma_small_batch_16 {token_count = 16} + +check.benchmark<@qwen3_moe_routed_down_q6k_f16_wmma_diverse_benchmark_case> @qwen3_moe_routed_down_q6k_f16_wmma_diverse_decode {token_count = 1} + +check.benchmark<@qwen3_moe_routed_down_q6k_f16_wmma_diverse_benchmark_case> @qwen3_moe_routed_down_q6k_f16_wmma_diverse_small_batch_2 {token_count = 2} + +check.benchmark<@qwen3_moe_routed_down_q6k_f16_wmma_diverse_benchmark_case> @qwen3_moe_routed_down_q6k_f16_wmma_diverse_small_batch_4 {token_count = 4} + +check.benchmark<@qwen3_moe_routed_down_q6k_f16_wmma_diverse_benchmark_case> @qwen3_moe_routed_down_q6k_f16_wmma_diverse_small_batch_8 {token_count = 8} + +check.benchmark<@qwen3_moe_routed_down_q6k_f16_wmma_diverse_benchmark_case> @qwen3_moe_routed_down_q6k_f16_wmma_diverse_small_batch_16 {token_count = 16} + +check.benchmark<@qwen3_moe_routed_down_q6k_f16_wmma_benchmark_case> @qwen3_moe_routed_down_q6k_f16_wmma_prefill_32 {token_count = 32} + +check.benchmark<@qwen3_moe_routed_down_q6k_f16_wmma_benchmark_case> @qwen3_moe_routed_down_q6k_f16_wmma_prefill_128 {token_count = 128} + +check.benchmark<@qwen3_moe_routed_down_q6k_f16_wmma_benchmark_case> @qwen3_moe_routed_down_q6k_f16_wmma_prefill_512 {token_count = 512} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/routed_down_weighted_reduce_next_rmsnorm_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/routed_down_weighted_reduce_next_rmsnorm_f32.loom new file mode 100644 index 000000000000..f08559f7e55f --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/routed_down_weighted_reduce_next_rmsnorm_f32.loom @@ -0,0 +1,160 @@ +// Fused grouped-prefill routed residual reduction and next-layer RMSNorm. +// +// One 256-workitem workgroup, comprising eight wave32 subgroups, owns one +// complete 2048-channel token row. Each workitem retains eight adjacent FP32 +// residual results across the workgroup RMS reduction, then publishes both +// the updated hidden state and the learned-weight normalized input consumed +// by the next layer. +amdgpu.target @qwen3_moe_routed_down_next_norm_gfx11_wave32 {subgroup_size = 32} + +config.decl @qwen3_moe.model.hidden_size : %value: index where [range(%value, 128, 32768), mul(%value, 128)] + +config.decl @qwen3_moe.model.rms_epsilon : f32 + +config.decl @qwen3_moe.routed_down.route_count : %value: index where [range(%value, 1, 8)] + +config.decl @qwen3_moe.routed_down.output_size : %value: index where [range(%value, 1, 4096)] + +kernel.decl @qwen3_moe_routed_down_weighted_reduce_f16_f32(%token_count: index) launch(%token_count: index, %route_weights: buffer, %routed_output: buffer, %hidden_state: buffer) + +kernel.decl @qwen3_moe_rmsnorm_f32(%token_count: index) launch(%token_count: index, %input: buffer, %weight: buffer, %output: buffer) + +kernel.def target(@qwen3_moe_routed_down_next_norm_gfx11_wave32) @qwen3_moe_routed_down_weighted_reduce_next_rmsnorm_f32(%token_count: index) { + %c1 = index.constant 1 : index + %c256 = index.constant 256 : index + kernel.launch.config workgroups(%token_count, %c1, %c1) workgroup_size(%c256, %c1, %c1) : index +} launch(%token_count: index, %route_weights: buffer, %routed_output: buffer, %hidden_state: buffer, %next_norm_weight: buffer, %next_projection_input: buffer) where [range(%token_count, 1, 2048)] { + %route_count0 = config.get @qwen3_moe.routed_down.route_count : index + %output_size0 = config.get @qwen3_moe.routed_down.output_size : index + %hidden_size0 = config.get @qwen3_moe.model.hidden_size : index + %epsilon = config.get @qwen3_moe.model.rms_epsilon : f32 + %route_count = index.assume %route_count0 [range(%route_count0, 8, 8)] : index + %output_size, %hidden_size = index.assume %output_size0, %hidden_size0 [range(%output_size0, 2048, 2048), range(%hidden_size0, 2048, 2048), eq(%output_size0, %hidden_size0)] : index, index + %token0 = kernel.workgroup.id : index + %workitem0 = kernel.workitem.id : index + %workitem = index.assume %workitem0 [range(%workitem0, 0, 255)] : index + %subgroup0 = kernel.subgroup.id : index + %subgroup = index.assume %subgroup0 [range(%subgroup0, 0, 7)] : index + %lane = kernel.subgroup.lane.id : index + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c8 = index.constant 8 : index + %channel = index.mul %workitem, %c8 : index + %scratch_bytes = index.constant 32 : offset + %c0_offset = index.constant 0 : offset + %c0_f32 = scalar.constant 0.0 : f32 + %c0_f32x8 = vector.constant 0.0 : vector<8xf32> + %active_token = index.cmp ult, %token0, %token_count : index + %safe_token0 = scf.select %active_token, %token0, %c0 : index + %token, %launch_token_count = index.assume %safe_token0, %token_count [lt(%safe_token0, %token_count)] : index, index + %hidden_size_i32 = index.cast %hidden_size : index to i32 + %hidden_size_f32 = scalar.sitofp %hidden_size_i32 : i32 to f32 + %assignment_count = index.mul %launch_token_count, %route_count : index + %route_weights_noalias, %routed_output_noalias, %hidden_state_noalias, %next_norm_weight_noalias, %next_projection_input_noalias = buffer.assume.noalias %route_weights, %routed_output, %hidden_state, %next_norm_weight, %next_projection_input : buffer, buffer, buffer, buffer, buffer + %route_weights_view = buffer.view %route_weights_noalias[%c0_offset] : buffer -> view<[%launch_token_count]x[%route_count]xf32> + %routed_output_view = buffer.view %routed_output_noalias[%c0_offset] : buffer -> view<[%assignment_count]x[%output_size]xf16> + %hidden_state_view = buffer.view %hidden_state_noalias[%c0_offset] : buffer -> view<[%launch_token_count]x[%hidden_size]xf32> + %next_norm_weight_view = buffer.view %next_norm_weight_noalias[%c0_offset] : buffer -> view<[%hidden_size]xf32> + %next_projection_input_view = buffer.view %next_projection_input_noalias[%c0_offset] : buffer -> view<[%launch_token_count]x[%hidden_size]xf32> + %scratch = buffer.alloca align(16) %scratch_bytes : buffer + %scratch_view = buffer.view %scratch[%c0_offset] : buffer -> view<8xf32> + + // The predicate depends only on the workgroup ID, so every workitem either + // executes the complete barrier-bearing row reduction or skips it together. + scf.if %active_token { + %weighted_sum = scf.for %route = [%c0 to %route_count step %c1](%sum = %c0_f32x8 : vector<8xf32>) -> (vector<8xf32>) unroll { + %assignment = index.madd %token, %route_count, %route : index + %routed_f16 = vector.load %routed_output_view[%assignment, %channel] : view<[%assignment_count]x[%output_size]xf16> -> vector<8xf16> + %routed = vector.extf %routed_f16 : vector<8xf16> to vector<8xf32> + %route_weight = view.load %route_weights_view[%token, %route] : view<[%launch_token_count]x[%route_count]xf32> -> f32 + %route_weight_x8 = vector.splat %route_weight : vector<8xf32> + %weighted = vector.mulf %routed, %route_weight_x8 : vector<8xf32> + %next = vector.addf %sum, %weighted : vector<8xf32> + scf.yield %next : vector<8xf32> + } + %residual = vector.load %hidden_state_view[%token, %channel] : view<[%launch_token_count]x[%hidden_size]xf32> -> vector<8xf32> + %result = vector.addf %residual, %weighted_sum : vector<8xf32> + %squares = vector.mulf %result, %result : vector<8xf32> + %thread_sum = vector.reduce %squares, %c0_f32 : vector<8xf32>, f32 + %subgroup_sum = kernel.subgroup.reduce %thread_sum : f32 + + %is_subgroup_leader = index.cmp eq, %lane, %c0 : index + scf.if %is_subgroup_leader { + view.store %subgroup_sum, %scratch_view[%subgroup] : f32, view<8xf32> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + %is_reduction_subgroup = index.cmp eq, %subgroup, %c0 : index + %is_reduction_lane = index.cmp ult, %lane, %c8 : index + %loads_subgroup_sum = scalar.andi %is_reduction_subgroup, %is_reduction_lane : i1 + %subgroup_partial = scf.if %loads_subgroup_sum -> (f32) { + %value = view.load %scratch_view[%lane] : view<8xf32> -> f32 + scf.yield %value : f32 + } else { + scf.yield %c0_f32 : f32 + } + %row_sum = kernel.subgroup.reduce %subgroup_partial : f32 + %writes_scale = scalar.andi %is_reduction_subgroup, %is_subgroup_leader : i1 + scf.if %writes_scale { + %mean = scalar.divf %row_sum, %hidden_size_f32 : f32 + %biased_mean = scalar.addf %mean, %epsilon : f32 + %scale = scalar.rsqrtf %biased_mean : f32 + view.store %scale, %scratch_view[%c0] : f32, view<8xf32> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + + %scale = view.load %scratch_view[%c0] : view<8xf32> -> f32 + %scale_x8 = vector.splat %scale : vector<8xf32> + %learned_weight = vector.load %next_norm_weight_view[%channel] : view<[%hidden_size]xf32> -> vector<8xf32> + %normalized = vector.mulf %result, %scale_x8 : vector<8xf32> + %next_projection = vector.mulf %normalized, %learned_weight : vector<8xf32> + vector.store %result, %hidden_state_view[%token, %channel] : vector<8xf32>, view<[%launch_token_count]x[%hidden_size]xf32> + vector.store %next_projection, %next_projection_input_view[%token, %channel] : vector<8xf32>, view<[%launch_token_count]x[%hidden_size]xf32> + } + kernel.return +} + +// The exact Prefill-512 shape compares both published tensors against the +// ordinary weighted reduction followed by the canonical RMSNorm. +check.case public @qwen3_moe_routed_down_weighted_reduce_next_rmsnorm_f32_differential_case { + %token_count = check.literal value(512) : index + %route_weights = check.generate.iota offset(0.03125) step(0.015625) period(8) : tensor<512x8xf32> + %routed_output = check.generate.iota offset(-0.25) step(0.015625) period(31) : tensor<512x8x2048xf16> + %expected_hidden_state = check.generate.iota offset(-0.5) step(0.0078125) period(17) : tensor<512x2048xf32> + %actual_hidden_state = check.generate.iota offset(-0.5) step(0.0078125) period(17) : tensor<512x2048xf32> + %next_norm_weight = check.generate.iota offset(0.5) step(0.0078125) period(29) : tensor<2048xf32> + %expected_next_projection_input = check.generate.fill value(0.0) : tensor<512x2048xf32> + %actual_next_projection_input = check.generate.fill value(1.0) : tensor<512x2048xf32> + kernel.launch @qwen3_moe_routed_down_weighted_reduce_f16_f32[%token_count](%token_count, %route_weights, %routed_output, %expected_hidden_state) : [index](index, tensor<512x8xf32>, tensor<512x8x2048xf16>, tensor<512x2048xf32>) + kernel.launch @qwen3_moe_rmsnorm_f32[%token_count](%token_count, %expected_hidden_state, %next_norm_weight, %expected_next_projection_input) : [index](index, tensor<512x2048xf32>, tensor<2048xf32>, tensor<512x2048xf32>) + kernel.launch @qwen3_moe_routed_down_weighted_reduce_next_rmsnorm_f32[%token_count](%token_count, %route_weights, %routed_output, %actual_hidden_state, %next_norm_weight, %actual_next_projection_input) : [index](index, tensor<512x8xf32>, tensor<512x8x2048xf16>, tensor<512x2048xf32>, tensor<2048xf32>, tensor<512x2048xf32>) + check.expect.close actual(%actual_hidden_state) expected(%expected_hidden_state) atol(0.0001) rtol(0.0001) nan(same) : tensor<512x2048xf32> + check.expect.close actual(%actual_next_projection_input) expected(%expected_next_projection_input) atol(0.001) rtol(0.001) nan(same) : tensor<512x2048xf32> + check.return +} + +check.case public @qwen3_moe_routed_down_weighted_reduce_next_rmsnorm_f32_fused_benchmark_case { + %token_count = check.literal value(512) : index + %route_weights = check.generate.fill value(0.125) : tensor<512x8xf32> + %routed_output = check.generate.fill value(0.5) : tensor<512x8x2048xf16> + %hidden_state = check.generate.fill value(1.0) : tensor<512x2048xf32> + %next_norm_weight = check.generate.fill value(1.0) : tensor<2048xf32> + %next_projection_input = check.generate.fill value(0.0) : tensor<512x2048xf32> + kernel.launch @qwen3_moe_routed_down_weighted_reduce_next_rmsnorm_f32[%token_count](%token_count, %route_weights, %routed_output, %hidden_state, %next_norm_weight, %next_projection_input) : [index](index, tensor<512x8xf32>, tensor<512x8x2048xf16>, tensor<512x2048xf32>, tensor<2048xf32>, tensor<512x2048xf32>) + check.return +} + +check.case public @qwen3_moe_routed_down_weighted_reduce_next_rmsnorm_f32_composed_benchmark_case { + %token_count = check.literal value(512) : index + %route_weights = check.generate.fill value(0.125) : tensor<512x8xf32> + %routed_output = check.generate.fill value(0.5) : tensor<512x8x2048xf16> + %hidden_state = check.generate.fill value(1.0) : tensor<512x2048xf32> + %next_norm_weight = check.generate.fill value(1.0) : tensor<2048xf32> + %next_projection_input = check.generate.fill value(0.0) : tensor<512x2048xf32> + kernel.launch @qwen3_moe_routed_down_weighted_reduce_f16_f32[%token_count](%token_count, %route_weights, %routed_output, %hidden_state) : [index](index, tensor<512x8xf32>, tensor<512x8x2048xf16>, tensor<512x2048xf32>) + kernel.launch @qwen3_moe_rmsnorm_f32[%token_count](%token_count, %hidden_state, %next_norm_weight, %next_projection_input) : [index](index, tensor<512x2048xf32>, tensor<2048xf32>, tensor<512x2048xf32>) + check.return +} + +check.benchmark<@qwen3_moe_routed_down_weighted_reduce_next_rmsnorm_f32_fused_benchmark_case> @qwen3_moe_routed_down_weighted_reduce_next_rmsnorm_f32_fused_prefill_512 + +check.benchmark<@qwen3_moe_routed_down_weighted_reduce_next_rmsnorm_f32_composed_benchmark_case> @qwen3_moe_routed_down_weighted_reduce_next_rmsnorm_f32_composed_prefill_512 diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/routed_down_weighted_reduce_next_rmsnorm_q8_1_x4.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/routed_down_weighted_reduce_next_rmsnorm_q8_1_x4.loom new file mode 100644 index 000000000000..00ff052de02c --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/routed_down_weighted_reduce_next_rmsnorm_q8_1_x4.loom @@ -0,0 +1,269 @@ +// Fused grouped-prefill routed residual reduction, next-layer RMSNorm, and +// GGML Q8_1 x4 publication. +// +// One 256-workitem workgroup, comprising eight wave32 subgroups, owns one +// complete 2048-channel token row. Each workitem retains eight adjacent +// residual results across the RMS reduction, updates hidden state, and +// quantizes the learned-weight normalized values directly into the physical +// row consumed by quantized QKV. This is the interlayer producer for the +// quantized attention schedule; no F32 normalized row is materialized. +// +// Four neighboring workitems own one logical 32-element Q8_1 block. Each +// workitem contributes two packed four-value words. LDS preserves the same +// eight-word max and sum reduction order as the standalone GGML packer so the +// fused and decomposed paths produce identical metadata and payload bytes. +amdgpu.target @qwen3_moe_routed_down_next_norm_q8_gfx11_wave32 {subgroup_size = 32} + +config.decl @qwen3_moe.model.hidden_size : %value: index where [range(%value, 128, 32768), mul(%value, 128)] + +config.decl @qwen3_moe.model.rms_epsilon : f32 + +config.decl @qwen3_moe.routed_down.route_count : %value: index where [range(%value, 1, 8)] + +config.decl @qwen3_moe.routed_down.output_size : %value: index where [range(%value, 1, 4096)] + +config.decl @qwen3_moe.workload.token_capacity : %value: index where [range(%value, 1, 2048)] + +kernel.decl @qwen3_moe_routed_down_weighted_reduce_f16_f32(%token_count: index) launch(%token_count: index, %route_weights: buffer, %routed_output: buffer, %hidden_state: buffer) + +kernel.decl @qwen3_moe_rmsnorm_f32(%token_count: index) launch(%token_count: index, %input: buffer, %weight: buffer, %output: buffer) + +kernel.decl @ggml_quantize_q8_1_x4_f32(%token_count: index, %input_size: index) launch(%token_count: index, %input_size: index, %input: buffer, %output: buffer) + +kernel.def target(@qwen3_moe_routed_down_next_norm_q8_gfx11_wave32) @qwen3_moe_routed_down_weighted_reduce_next_rmsnorm_q8_1_x4(%token_count: index) { + %token_capacity = config.get @qwen3_moe.workload.token_capacity : index + %c1 = index.constant 1 : index + %c256 = index.constant 256 : index + kernel.launch.config workgroups(%token_capacity, %c1, %c1) workgroup_size(%c256, %c1, %c1) : index +} launch(%token_count: index, %route_weights: buffer, %routed_output: buffer, %hidden_state: buffer, %next_norm_weight: buffer, %next_projection_input: buffer) { + %route_count0 = config.get @qwen3_moe.routed_down.route_count : index + %output_size0 = config.get @qwen3_moe.routed_down.output_size : index + %hidden_size0 = config.get @qwen3_moe.model.hidden_size : index + %epsilon = config.get @qwen3_moe.model.rms_epsilon : f32 + %token_capacity = config.get @qwen3_moe.workload.token_capacity : index + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 2048), le(%token_count, %token_capacity)] : index + %route_count = index.assume %route_count0 [range(%route_count0, 8, 8)] : index + %output_size, %hidden_size = index.assume %output_size0, %hidden_size0 [range(%output_size0, 2048, 2048), range(%hidden_size0, 2048, 2048), eq(%output_size0, %hidden_size0)] : index, index + %token0 = kernel.workgroup.id : index + %workitem0 = kernel.workitem.id : index + %workitem = index.assume %workitem0 [range(%workitem0, 0, 255)] : index + %subgroup0 = kernel.subgroup.id : index + %subgroup = index.assume %subgroup0 [range(%subgroup0, 0, 7)] : index + %lane = kernel.subgroup.lane.id : index + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c4 = index.constant 4 : index + %c8 = index.constant 8 : index + %c128 = index.constant 128 : index + %group_bytes = index.constant 144 : offset + %payload_byte_add = index.constant 16 : offset + %scratch_d_byte_add = index.constant 2048 : offset + %scratch_bytes = index.constant 2304 : offset + %c0_offset = index.constant 0 : offset + %c0_f32 = scalar.constant 0.0 : f32 + %c1_f32 = scalar.constant 1.0 : f32 + %c127 = scalar.constant 127.0 : f32 + %c0_f32x8 = vector.constant 0.0 : vector<8xf32> + %active_token = index.cmp ult, %token0, %bounded_token_count : index + %safe_token0 = scf.select %active_token, %token0, %c0 : index + %token, %launch_token_count = index.assume %safe_token0, %bounded_token_count [lt(%safe_token0, %bounded_token_count)] : index, index + %hidden_size_i32 = index.cast %hidden_size : index to i32 + %hidden_size_f32 = scalar.sitofp %hidden_size_i32 : i32 to f32 + %assignment_count = index.mul %launch_token_count, %route_count : index + %physical_group_count = index.div %hidden_size, %c128 : index + %row_bytes = index.scale %physical_group_count, %group_bytes : index, offset -> offset + %token_output_byte_base = index.scale %token, %row_bytes : index, offset -> offset + %route_weights_noalias, %routed_output_noalias, %hidden_state_noalias, %next_norm_weight_noalias, %next_projection_input_noalias = buffer.assume.noalias %route_weights, %routed_output, %hidden_state, %next_norm_weight, %next_projection_input : buffer, buffer, buffer, buffer, buffer + %route_weights_view = buffer.view %route_weights_noalias[%c0_offset] : buffer -> view<[%launch_token_count]x[%route_count]xf32> + %routed_output_view = buffer.view %routed_output_noalias[%c0_offset] : buffer -> view<[%assignment_count]x[%output_size]xf16> + %hidden_state_view = buffer.view %hidden_state_noalias[%c0_offset] : buffer -> view<[%launch_token_count]x[%hidden_size]xf32> + %next_norm_weight_view = buffer.view %next_norm_weight_noalias[%c0_offset] : buffer -> view<[%hidden_size]xf32> + %scratch = buffer.alloca align(16) %scratch_bytes : buffer + %scratch_values = buffer.view %scratch[%c0_offset] : buffer -> view<512xf32> + %scratch_d = buffer.view %scratch[%scratch_d_byte_add] : buffer -> view<64xf32> + // The token predicate is workgroup-uniform, keeping every RMSNorm and Q8_1 + // barrier either active for the whole workgroup or skipped by all workitems. + scf.if %active_token { + %channel = index.mul %workitem, %c8 : index + %weighted_sum = scf.for %route = [%c0 to %route_count step %c1](%sum = %c0_f32x8 : vector<8xf32>) -> (vector<8xf32>) unroll { + %assignment = index.madd %token, %route_count, %route : index + %routed_f16 = vector.load %routed_output_view[%assignment, %channel] : view<[%assignment_count]x[%output_size]xf16> -> vector<8xf16> + %routed = vector.extf %routed_f16 : vector<8xf16> to vector<8xf32> + %route_weight = view.load %route_weights_view[%token, %route] : view<[%launch_token_count]x[%route_count]xf32> -> f32 + %route_weight_x8 = vector.splat %route_weight : vector<8xf32> + %weighted = vector.mulf %routed, %route_weight_x8 : vector<8xf32> + %next = vector.addf %sum, %weighted : vector<8xf32> + scf.yield %next : vector<8xf32> + } + %residual = vector.load %hidden_state_view[%token, %channel] : view<[%launch_token_count]x[%hidden_size]xf32> -> vector<8xf32> + %result = vector.addf %residual, %weighted_sum : vector<8xf32> + %squares = vector.mulf %result, %result : vector<8xf32> + %thread_sum = vector.reduce %squares, %c0_f32 : vector<8xf32>, f32 + %subgroup_sum = kernel.subgroup.reduce %thread_sum : f32 + + %is_subgroup_leader = index.cmp eq, %lane, %c0 : index + scf.if %is_subgroup_leader { + view.store %subgroup_sum, %scratch_values[%subgroup] : f32, view<512xf32> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + %is_reduction_subgroup = index.cmp eq, %subgroup, %c0 : index + %is_reduction_lane = index.cmp ult, %lane, %c8 : index + %loads_subgroup_sum = scalar.andi %is_reduction_subgroup, %is_reduction_lane : i1 + %subgroup_partial = scf.if %loads_subgroup_sum -> (f32) { + %value = view.load %scratch_values[%lane] : view<512xf32> -> f32 + scf.yield %value : f32 + } else { + scf.yield %c0_f32 : f32 + } + %row_sum = kernel.subgroup.reduce %subgroup_partial : f32 + %writes_scale = scalar.andi %is_reduction_subgroup, %is_subgroup_leader : i1 + scf.if %writes_scale { + %mean = scalar.divf %row_sum, %hidden_size_f32 : f32 + %biased_mean = scalar.addf %mean, %epsilon : f32 + %scale = scalar.rsqrtf %biased_mean : f32 + view.store %scale, %scratch_values[%c0] : f32, view<512xf32> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + + %row_scale = view.load %scratch_values[%c0] : view<512xf32> -> f32 + // All waves must capture the row scale before scratch is reused for the + // independent Q8 block reductions. + kernel.barrier scope(workgroup) ordering(acq_rel) + %row_scale_vector = vector.splat %row_scale : vector<8xf32> + %learned_weight = vector.load %next_norm_weight_view[%channel] : view<[%hidden_size]xf32> -> vector<8xf32> + %normalized0 = vector.mulf %result, %row_scale_vector : vector<8xf32> + %normalized = vector.mulf %normalized0, %learned_weight : vector<8xf32> + vector.store %result, %hidden_state_view[%token, %channel] : vector<8xf32>, view<[%launch_token_count]x[%hidden_size]xf32> + + %low_values = vector.slice %normalized[0] : vector<8xf32> -> vector<4xf32> + %high_values = vector.slice %normalized[4] : vector<8xf32> -> vector<4xf32> + %low_absolute_values = vector.absf %low_values : vector<4xf32> + %high_absolute_values = vector.absf %high_values : vector<4xf32> + %low_max = vector.reduce %low_absolute_values, %c0_f32 : vector<4xf32>, f32 + %high_max = vector.reduce %high_absolute_values, %c0_f32 : vector<4xf32>, f32 + %partial_base = index.mul %workitem, %c2 : index + %partial_high = index.add %partial_base, %c1 : index + view.store %low_max, %scratch_values[%partial_base] : f32, view<512xf32> + view.store %high_max, %scratch_values[%partial_high] : f32, view<512xf32> + kernel.barrier scope(workgroup) ordering(acq_rel) + + %block_in_row = index.div %workitem, %c4 : index + %workitem_in_block = index.rem %workitem, %c4 : index + %is_block_leader = index.cmp eq, %workitem_in_block, %c0 : index + scf.if %is_block_leader { + %block_partial_base = index.mul %block_in_row, %c8 : index + %block_maxima = vector.load %scratch_values[%block_partial_base] : view<512xf32> -> vector<8xf32> + %amax = vector.reduce %block_maxima, %c0_f32 : vector<8xf32>, f32 + %d = scalar.divf %amax, %c127 : f32 + view.store %d, %scratch_d[%block_in_row] : f32, view<64xf32> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + + %d = view.load %scratch_d[%block_in_row] : view<64xf32> -> f32 + %d_nonzero = scalar.cmpf one, %d, %c0_f32 : f32 + %d_inverse = scf.if %d_nonzero -> (f32) { + %inverse = scalar.divf %c1_f32, %d : f32 + scf.yield %inverse : f32 + } else { + scf.yield %c0_f32 : f32 + } + %d_inverse_vector = vector.splat %d_inverse : vector<8xf32> + %scaled_values = vector.mulf %normalized, %d_inverse_vector : vector<8xf32> + %rounded_values = vector.roundf %scaled_values : vector<8xf32> + %quantized_values = vector.fptosi %rounded_values : vector<8xf32> to vector<8xi8> + %packed_words = vector.bitcast %quantized_values : vector<8xi8> to vector<2xi32> + %physical_group = index.div %block_in_row, %c4 : index + %block_in_group = index.rem %block_in_row, %c4 : index + %group_byte_add = index.scale %physical_group, %group_bytes : index, offset -> offset + %group_byte_offset = index.add %token_output_byte_base, %group_byte_add : offset + %payload_byte_offset = index.add %group_byte_offset, %payload_byte_add : offset + %group_ds = buffer.view %next_projection_input_noalias[%group_byte_offset] : buffer -> view<8xf16> + %group_qs = buffer.view %next_projection_input_noalias[%payload_byte_offset] : buffer -> view<32xi32> + %block_word_base = index.mul %block_in_group, %c8 : index + %workitem_word_add = index.mul %workitem_in_block, %c2 : index + %packed_word_index0 = index.add %block_word_base, %workitem_word_add : index + %packed_word_index = index.assume %packed_word_index0 [range(%packed_word_index0, 0, 31)] : index + vector.store %packed_words, %group_qs[%packed_word_index] : vector<2xi32>, view<32xi32> + + %low_rounded_values = vector.slice %rounded_values[0] : vector<8xf32> -> vector<4xf32> + %high_rounded_values = vector.slice %rounded_values[4] : vector<8xf32> -> vector<4xf32> + %low_quantized_sum = vector.reduce %low_rounded_values, %c0_f32 : vector<4xf32>, f32 + %high_quantized_sum = vector.reduce %high_rounded_values, %c0_f32 : vector<4xf32>, f32 + view.store %low_quantized_sum, %scratch_values[%partial_base] : f32, view<512xf32> + view.store %high_quantized_sum, %scratch_values[%partial_high] : f32, view<512xf32> + kernel.barrier scope(workgroup) ordering(acq_rel) + scf.if %is_block_leader { + %block_partial_base = index.mul %block_in_row, %c8 : index + %block_sums = vector.load %scratch_values[%block_partial_base] : view<512xf32> -> vector<8xf32> + %quantized_sum = vector.reduce %block_sums, %c0_f32 : vector<8xf32>, f32 + %s = scalar.mulf %quantized_sum, %d : f32 + %d_f16 = scalar.fptrunc %d : f32 to f16 + %s_f16 = scalar.fptrunc %s : f32 to f16 + %ds_index = index.mul %block_in_group, %c2 : index + %s_index = index.add %ds_index, %c1 : index + view.store %d_f16, %group_ds[%ds_index] : f16, view<8xf16> + view.store %s_f16, %group_ds[%s_index] : f16, view<8xf16> + } + } + kernel.return +} + +// Fourteen independent token rows lock the complete interlayer contract. The +// fused producer must match ordinary weighted reduction for hidden state and +// RMSNorm followed by the generic GGML packer for every Q8 metadata and +// payload byte. +check.case public @qwen3_moe_routed_down_weighted_reduce_next_rmsnorm_q8_1_x4_differential_case { + %token_count = check.literal value(14) : index + %hidden_size = check.literal value(2048) : index + %route_weights = check.generate.iota offset(0.03125) step(0.015625) period(8) : tensor<14x8xf32> + %routed_seed = check.param.seed base(5858425849430430008) count(1) : i64 + %routed_output = check.generate.random.uniform seed(%routed_seed) range(-0.5 to 0.5) : tensor<14x8x2048xf16> + %expected_hidden_state = check.generate.iota offset(-0.5) step(0.0078125) period(31) : tensor<14x2048xf32> + %actual_hidden_state = check.generate.iota offset(-0.5) step(0.0078125) period(31) : tensor<14x2048xf32> + %next_norm_weight = check.generate.iota offset(0.5) step(0.0078125) period(29) : tensor<2048xf32> + %normalized = check.generate.fill value(0.0) : tensor<14x2048xf32> + %expected_q8 = check.generate.fill value(0) : tensor<14x2304xi8> + %actual_q8 = check.generate.fill value(1) : tensor<14x2304xi8> + kernel.launch @qwen3_moe_routed_down_weighted_reduce_f16_f32[%token_count](%token_count, %route_weights, %routed_output, %expected_hidden_state) : [index](index, tensor<14x8xf32>, tensor<14x8x2048xf16>, tensor<14x2048xf32>) + kernel.launch @qwen3_moe_rmsnorm_f32[%token_count](%token_count, %expected_hidden_state, %next_norm_weight, %normalized) : [index](index, tensor<14x2048xf32>, tensor<2048xf32>, tensor<14x2048xf32>) + kernel.launch @ggml_quantize_q8_1_x4_f32[%token_count, %hidden_size](%token_count, %hidden_size, %normalized, %expected_q8) : [index, index](index, index, tensor<14x2048xf32>, tensor<14x2304xi8>) + kernel.launch @qwen3_moe_routed_down_weighted_reduce_next_rmsnorm_q8_1_x4[%token_count](%token_count, %route_weights, %routed_output, %actual_hidden_state, %next_norm_weight, %actual_q8) : [index](index, tensor<14x8xf32>, tensor<14x8x2048xf16>, tensor<14x2048xf32>, tensor<2048xf32>, tensor<14x2304xi8>) + check.expect.close actual(%actual_hidden_state) expected(%expected_hidden_state) atol(0.0001) rtol(0.0001) nan(same) : tensor<14x2048xf32> + check.expect.equal actual(%actual_q8) expected(%expected_q8) : tensor<14x2304xi8> + check.return +} + +check.case public @qwen3_moe_routed_down_weighted_reduce_next_rmsnorm_q8_1_x4_fused_benchmark_case { + %token_count = check.param.choice values([14, 32, 128, 512]) name("token_count") : index + %route_weights = check.generate.fill value(0.125) : tensor<[%token_count]x8xf32> + %routed_output = check.generate.fill value(0.0) : tensor<[%token_count]x8x2048xf16> + %hidden_state = check.generate.fill value(1.0) : tensor<[%token_count]x2048xf32> + %next_norm_weight = check.generate.fill value(1.0) : tensor<2048xf32> + %next_projection_input = check.generate.fill value(0) : tensor<[%token_count]x2304xi8> + kernel.launch @qwen3_moe_routed_down_weighted_reduce_next_rmsnorm_q8_1_x4[%token_count](%token_count, %route_weights, %routed_output, %hidden_state, %next_norm_weight, %next_projection_input) : [index](index, tensor<[%token_count]x8xf32>, tensor<[%token_count]x8x2048xf16>, tensor<[%token_count]x2048xf32>, tensor<2048xf32>, tensor<[%token_count]x2304xi8>) + check.return +} + +check.case public @qwen3_moe_routed_down_weighted_reduce_next_rmsnorm_q8_1_x4_composed_benchmark_case { + %token_count = check.param.choice values([14, 32, 128, 512]) name("token_count") : index + %hidden_size = check.literal value(2048) : index + %route_weights = check.generate.fill value(0.125) : tensor<[%token_count]x8xf32> + %routed_output = check.generate.fill value(0.0) : tensor<[%token_count]x8x2048xf16> + %hidden_state = check.generate.fill value(1.0) : tensor<[%token_count]x2048xf32> + %next_norm_weight = check.generate.fill value(1.0) : tensor<2048xf32> + %normalized = check.generate.fill value(0.0) : tensor<[%token_count]x2048xf32> + %next_projection_input = check.generate.fill value(0) : tensor<[%token_count]x2304xi8> + kernel.launch @qwen3_moe_routed_down_weighted_reduce_f16_f32[%token_count](%token_count, %route_weights, %routed_output, %hidden_state) : [index](index, tensor<[%token_count]x8xf32>, tensor<[%token_count]x8x2048xf16>, tensor<[%token_count]x2048xf32>) + kernel.launch @qwen3_moe_rmsnorm_f32[%token_count](%token_count, %hidden_state, %next_norm_weight, %normalized) : [index](index, tensor<[%token_count]x2048xf32>, tensor<2048xf32>, tensor<[%token_count]x2048xf32>) + kernel.launch @ggml_quantize_q8_1_x4_f32[%token_count, %hidden_size](%token_count, %hidden_size, %normalized, %next_projection_input) : [index, index](index, index, tensor<[%token_count]x2048xf32>, tensor<[%token_count]x2304xi8>) + check.return +} + +check.benchmark<@qwen3_moe_routed_down_weighted_reduce_next_rmsnorm_q8_1_x4_fused_benchmark_case> @qwen3_moe_routed_down_weighted_reduce_next_rmsnorm_q8_1_x4_fused_prefill_14 {token_count = 14} + +check.benchmark<@qwen3_moe_routed_down_weighted_reduce_next_rmsnorm_q8_1_x4_fused_benchmark_case> @qwen3_moe_routed_down_weighted_reduce_next_rmsnorm_q8_1_x4_fused_prefill_512 {token_count = 512} + +check.benchmark<@qwen3_moe_routed_down_weighted_reduce_next_rmsnorm_q8_1_x4_composed_benchmark_case> @qwen3_moe_routed_down_weighted_reduce_next_rmsnorm_q8_1_x4_composed_prefill_14 {token_count = 14} + +check.benchmark<@qwen3_moe_routed_down_weighted_reduce_next_rmsnorm_q8_1_x4_composed_benchmark_case> @qwen3_moe_routed_down_weighted_reduce_next_rmsnorm_q8_1_x4_composed_prefill_512 {token_count = 512} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/routed_gate_up_swiglu_q4k.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/routed_gate_up_swiglu_q4k.loom new file mode 100644 index 000000000000..f96abf03cc97 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/routed_gate_up_swiglu_q4k.loom @@ -0,0 +1,1688 @@ +// Fuses the Qwen3 MoE gate and up expert projections with the SwiGLU +// epilogue. Weights are consumed directly from GGUF's raw Q4_K layout: +// [expert][output channel][K / 256][144 bytes]. Activations use GGML's Q8_1 +// x4 layout produced by the shared quantization kernel. +// +// Route IDs have a logical route_count and an independent physical +// route_stride. This admits llama.cpp's [token][128] argsort storage while +// operating on only the selected top-8 entries, and also admits compact route +// buffers produced by a future fused router. +// +// Two schedules share this contract. Small token counts use one wave per +// [token, route, output channel]. Larger token counts first build a compact +// assignment table per expert, then group up to 32 selected rows with 32 +// output channels so raw weights and quantized activations can be reused +// through workgroup memory. +func.def inline @ggml_q8_1_x4_block(%q8_input: buffer, %row_byte_base: offset, %q8_block: index) -> (vector<32xi8>, f32, f32) { + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c4 = index.constant 4 : index + %c8 = index.constant 8 : index + %group_bytes = index.constant 144 : offset + %payload_byte_add = index.constant 16 : offset + %group = index.div %q8_block, %c4 : index + %inner0 = index.rem %q8_block, %c4 : index + %inner = index.assume %inner0 [range(%inner0, 0, 3)] : index + %group_byte_add = index.scale %group, %group_bytes : index, offset -> offset + %group_byte_base = index.add %row_byte_base, %group_byte_add : offset + %payload_byte_base = index.add %group_byte_base, %payload_byte_add : offset + %ds_view = buffer.view %q8_input[%group_byte_base] : buffer -> view<8xf16> + %payload_view = buffer.view %q8_input[%payload_byte_base] : buffer -> view<32xi32> + %d_index = index.mul %inner, %c2 : index + %s_index = index.add %d_index, %c1 : index + %word_index = index.mul %inner, %c8 : index + %d_f16 = view.load %ds_view[%d_index] : view<8xf16> -> f16 + %s_f16 = view.load %ds_view[%s_index] : view<8xf16> -> f16 + %words = vector.load %payload_view[%word_index] : view<32xi32> -> vector<8xi32> + %values = vector.bitcast %words : vector<8xi32> to vector<32xi8> + %d = scalar.extf %d_f16 : f16 to f32 + %s = scalar.extf %s_f16 : f16 to f32 + func.return %values, %d, %s : vector<32xi8>, f32, f32 +} + +config.decl @qwen3_moe.routed_gate_up.input_size : %value: index where [range(%value, 512, 32768), mul(%value, 512)] + +config.decl @qwen3_moe.routed_gate_up.expert_count : %value: index where [range(%value, 1, 512)] + +config.decl @qwen3_moe.routed_gate_up.route_count : %value: index where [range(%value, 1, 8)] + +config.decl @qwen3_moe.routed_gate_up.output_size : %value: index where [range(%value, 1, 4096)] + +config.decl @qwen3_moe.workload.token_capacity : %value: index where [range(%value, 1, 2048)] + +kernel.decl @ggml_quantize_q8_1_x4_f32(%token_count: index, %input_size: index) launch(%token_count: index, %input_size: index, %input: buffer, %output: buffer) + +func.decl @ggml_q8_1_x4_block(%q8_input: buffer, %row_byte_base: offset, %q8_block: index) -> (vector<32xi8>, f32, f32) + +// Decodes one 16-element Q4_K chunk while preserving its unsigned nibbles for +// AMDGPU's mixed u8*s8 dot4 instruction. The returned scale terms are shared +// by every routed activation row consuming the same weight chunk. +func.def inline @qwen3_moe_q4k_chunk_local(%weight: buffer, %row_byte_base: offset, %q4_block: index, %q4_group: index, %q4_half: index) -> (vector<16xi8>, f32, f32) { + %c4 = index.constant 4 : index + %c8 = index.constant 8 : index + %c16 = index.constant 16 : index + %block_bytes = index.constant 144 : offset + %header_byte_count = index.constant 4 : index + %code_byte_add = index.constant 16 : offset + %c4_i32 = scalar.constant 4 : i32 + %c6_i32 = scalar.constant 6 : i32 + %c15_i32 = scalar.constant 15 : i32 + %c63_i32 = scalar.constant 63 : i32 + %nibble_mask = vector.constant 252645135 : vector<4xi32> + %bounded_group = index.assume %q4_group [range(%q4_group, 0, 7)] : index + %bounded_half = index.assume %q4_half [range(%q4_half, 0, 1)] : index + %block_byte_add = index.scale %q4_block, %block_bytes : index, offset -> offset + %block_byte_base = index.add %row_byte_base, %block_byte_add : offset + %code_byte_base = index.add %block_byte_base, %code_byte_add : offset + %header_view = buffer.view %weight[%block_byte_base] : buffer -> view<4xi32> + %code_view = buffer.view %weight[%code_byte_base] : buffer -> view<32xi32> + %header_words = vector.load %header_view[0] : view<4xi32> -> vector<4xi32> + %header_halves = vector.bitcast %header_words : vector<4xi32> to vector<8xf16> + %d_f16 = vector.extract %header_halves[0] : vector<8xf16> -> f16 + %dmin_f16 = vector.extract %header_halves[1] : vector<8xf16> -> f16 + %d = scalar.extf %d_f16 : f16 to f32 + %dmin = scalar.extf %dmin_f16 : f16 to f32 + %header_bytes = vector.bitcast %header_words : vector<4xi32> to vector<16xi8> + %iqs_group = index.mul %bounded_group, %c8 : index + %iqs_half = index.mul %bounded_half, %c4 : index + %iqs = index.add %iqs_group, %iqs_half : index + %qs_page0 = index.div %iqs, %c16 : index + %qs_page = index.mul %qs_page0, %c8 : index + %qs_lane = index.rem %iqs, %c8 : index + %qs_index0 = index.add %qs_page, %qs_lane : index + %qs_index = index.assume %qs_index0 [range(%qs_index0, 0, 28)] : index + %iqs_mod16 = index.rem %iqs, %c16 : index + %nibble_page = index.div %iqs_mod16, %c8 : index + %nibble_shift_index = index.mul %nibble_page, %c4 : index + %nibble_shift_i32 = index.cast %nibble_shift_index : index to i32 + %nibble_shift = vector.splat %nibble_shift_i32 : vector<4xi32> + %packed_codes = vector.load %code_view[%qs_index] : view<32xi32> -> vector<4xi32> + %shifted_codes = vector.shrui %packed_codes, %nibble_shift : vector<4xi32> + %masked_codes = vector.andi %shifted_codes, %nibble_mask : vector<4xi32> + %q4_values = vector.bitcast %masked_codes : vector<4xi32> to vector<16xi8> + %is_low_group = index.cmp ult, %bounded_group, %c4 : index + %scale, %minimum = scf.if %is_low_group -> (i32, i32) { + %scale_index = index.add %bounded_group, %header_byte_count : index + %minimum_index0 = index.add %bounded_group, %c4 : index + %minimum_index = index.add %minimum_index0, %header_byte_count : index + %scale_i8 = vector.extract %header_bytes[%scale_index] : vector<16xi8> -> i8 + %minimum_i8 = vector.extract %header_bytes[%minimum_index] : vector<16xi8> -> i8 + %scale_u8 = scalar.extui %scale_i8 : i8 to i32 + %minimum_u8 = scalar.extui %minimum_i8 : i8 to i32 + %scale_low6 = scalar.andi %scale_u8, %c63_i32 : i32 + %minimum_low6 = scalar.andi %minimum_u8, %c63_i32 : i32 + scf.yield %scale_low6, %minimum_low6 : i32, i32 + } else { + %packed_index0 = index.add %bounded_group, %c4 : index + %packed_index = index.add %packed_index0, %header_byte_count : index + %scale_high_index = index.sub %bounded_group, %c4 : index + %scale_high_header_index = index.add %scale_high_index, %header_byte_count : index + %minimum_high_header_index = index.add %bounded_group, %header_byte_count : index + %packed_i8 = vector.extract %header_bytes[%packed_index] : vector<16xi8> -> i8 + %scale_high_i8 = vector.extract %header_bytes[%scale_high_header_index] : vector<16xi8> -> i8 + %minimum_high_i8 = vector.extract %header_bytes[%minimum_high_header_index] : vector<16xi8> -> i8 + %packed = scalar.extui %packed_i8 : i8 to i32 + %scale_high = scalar.extui %scale_high_i8 : i8 to i32 + %minimum_high = scalar.extui %minimum_high_i8 : i8 to i32 + %scale_low4 = scalar.andi %packed, %c15_i32 : i32 + %minimum_low4 = scalar.shrui %packed, %c4_i32 : i32 + %scale_high2_raw = scalar.shrui %scale_high, %c6_i32 : i32 + %minimum_high2_raw = scalar.shrui %minimum_high, %c6_i32 : i32 + %scale_high2 = scalar.shli %scale_high2_raw, %c4_i32 : i32 + %minimum_high2 = scalar.shli %minimum_high2_raw, %c4_i32 : i32 + %scale_value = scalar.ori %scale_low4, %scale_high2 : i32 + %minimum_value = scalar.ori %minimum_low4, %minimum_high2 : i32 + scf.yield %scale_value, %minimum_value : i32, i32 + } + %scale_f32 = scalar.uitofp %scale : i32 to f32 + %minimum_f32 = scalar.uitofp %minimum : i32 to f32 + %d_scale = scalar.mulf %d, %scale_f32 : f32 + %dmin_scale = scalar.mulf %dmin, %minimum_f32 : f32 + func.return %q4_values, %d_scale, %dmin_scale : vector<16xi8>, f32, f32 +} + +// Global-weight variant of the Q4_K chunk decoder. Scalar scale-byte loads are +// legal and cheaper for the one-wave schedule, while the LDS schedule uses the +// packed-header variant above because AMDGPU has no local i8 load descriptor. +func.def inline @qwen3_moe_q4k_chunk_global(%weight: buffer, %row_byte_base: offset, %q4_block: index, %q4_group: index, %q4_half: index) -> (vector<16xi8>, f32, f32) { + %c4 = index.constant 4 : index + %c8 = index.constant 8 : index + %c16 = index.constant 16 : index + %block_bytes = index.constant 144 : offset + %scale_byte_add = index.constant 4 : offset + %code_byte_add = index.constant 16 : offset + %c4_i32 = scalar.constant 4 : i32 + %c6_i32 = scalar.constant 6 : i32 + %c15_i32 = scalar.constant 15 : i32 + %c63_i32 = scalar.constant 63 : i32 + %nibble_mask = vector.constant 252645135 : vector<4xi32> + %bounded_group = index.assume %q4_group [range(%q4_group, 0, 7)] : index + %bounded_half = index.assume %q4_half [range(%q4_half, 0, 1)] : index + %block_byte_add = index.scale %q4_block, %block_bytes : index, offset -> offset + %block_byte_base = index.add %row_byte_base, %block_byte_add : offset + %scale_byte_base = index.add %block_byte_base, %scale_byte_add : offset + %code_byte_base = index.add %block_byte_base, %code_byte_add : offset + %half_view = buffer.view %weight[%block_byte_base] : buffer -> view<2xf16> + %scale_view = buffer.view %weight[%scale_byte_base] : buffer -> view<12xi8> + %code_view = buffer.view %weight[%code_byte_base] : buffer -> view<32xi32> + %d_f16 = view.load %half_view[0] : view<2xf16> -> f16 + %dmin_f16 = view.load %half_view[1] : view<2xf16> -> f16 + %d = scalar.extf %d_f16 : f16 to f32 + %dmin = scalar.extf %dmin_f16 : f16 to f32 + %iqs_group = index.mul %bounded_group, %c8 : index + %iqs_half = index.mul %bounded_half, %c4 : index + %iqs = index.add %iqs_group, %iqs_half : index + %qs_page0 = index.div %iqs, %c16 : index + %qs_page = index.mul %qs_page0, %c8 : index + %qs_lane = index.rem %iqs, %c8 : index + %qs_index0 = index.add %qs_page, %qs_lane : index + %qs_index = index.assume %qs_index0 [range(%qs_index0, 0, 28)] : index + %iqs_mod16 = index.rem %iqs, %c16 : index + %nibble_page = index.div %iqs_mod16, %c8 : index + %nibble_shift_index = index.mul %nibble_page, %c4 : index + %nibble_shift_i32 = index.cast %nibble_shift_index : index to i32 + %nibble_shift = vector.splat %nibble_shift_i32 : vector<4xi32> + %packed_codes = vector.load %code_view[%qs_index] : view<32xi32> -> vector<4xi32> + %shifted_codes = vector.shrui %packed_codes, %nibble_shift : vector<4xi32> + %masked_codes = vector.andi %shifted_codes, %nibble_mask : vector<4xi32> + %q4_values = vector.bitcast %masked_codes : vector<4xi32> to vector<16xi8> + %is_low_group = index.cmp ult, %bounded_group, %c4 : index + %scale, %minimum = scf.if %is_low_group -> (i32, i32) { + %minimum_index = index.add %bounded_group, %c4 : index + %scale_i8 = view.load %scale_view[%bounded_group] : view<12xi8> -> i8 + %minimum_i8 = view.load %scale_view[%minimum_index] : view<12xi8> -> i8 + %scale_u8 = scalar.extui %scale_i8 : i8 to i32 + %minimum_u8 = scalar.extui %minimum_i8 : i8 to i32 + %scale_low6 = scalar.andi %scale_u8, %c63_i32 : i32 + %minimum_low6 = scalar.andi %minimum_u8, %c63_i32 : i32 + scf.yield %scale_low6, %minimum_low6 : i32, i32 + } else { + %packed_index = index.add %bounded_group, %c4 : index + %scale_high_index = index.sub %bounded_group, %c4 : index + %packed_i8 = view.load %scale_view[%packed_index] : view<12xi8> -> i8 + %scale_high_i8 = view.load %scale_view[%scale_high_index] : view<12xi8> -> i8 + %minimum_high_i8 = view.load %scale_view[%bounded_group] : view<12xi8> -> i8 + %packed = scalar.extui %packed_i8 : i8 to i32 + %scale_high = scalar.extui %scale_high_i8 : i8 to i32 + %minimum_high = scalar.extui %minimum_high_i8 : i8 to i32 + %scale_low4 = scalar.andi %packed, %c15_i32 : i32 + %minimum_low4 = scalar.shrui %packed, %c4_i32 : i32 + %scale_high2_raw = scalar.shrui %scale_high, %c6_i32 : i32 + %minimum_high2_raw = scalar.shrui %minimum_high, %c6_i32 : i32 + %scale_high2 = scalar.shli %scale_high2_raw, %c4_i32 : i32 + %minimum_high2 = scalar.shli %minimum_high2_raw, %c4_i32 : i32 + %scale_value = scalar.ori %scale_low4, %scale_high2 : i32 + %minimum_value = scalar.ori %minimum_low4, %minimum_high2 : i32 + scf.yield %scale_value, %minimum_value : i32, i32 + } + %scale_f32 = scalar.uitofp %scale : i32 to f32 + %minimum_f32 = scalar.uitofp %minimum : i32 to f32 + %d_scale = scalar.mulf %d, %scale_f32 : f32 + %dmin_scale = scalar.mulf %dmin, %minimum_f32 : f32 + func.return %q4_values, %d_scale, %dmin_scale : vector<16xi8>, f32, f32 +} + +// Decodes one Q4_K scale/minimum pair from a header loaded as three packed +// scale words. Keeping the complete header in registers lets paired-group +// contractions share both the header and packed-code load. +func.def inline @qwen3_moe_q4k_scale_from_header(%scale0: i32, %scale1: i32, %scale2: i32, %q4_group: index) -> (i32, i32) { + %c4 = index.constant 4 : index + %c8 = index.constant 8 : index + %c2_i32 = scalar.constant 2 : i32 + %c4_i32 = scalar.constant 4 : i32 + %c15_i32 = scalar.constant 15 : i32 + %c48_i32 = scalar.constant 48 : i32 + %bounded_group = index.assume %q4_group [range(%q4_group, 0, 7)] : index + %is_low_group = index.cmp ult, %bounded_group, %c4 : index + %scale_lane = index.rem %bounded_group, %c4 : index + %scale_shift_index = index.mul %scale_lane, %c8 : index + %scale_shift = index.cast %scale_shift_index : index to i32 + %high_shift = scalar.addi %scale_shift, %c2_i32 : i32 + %minimum_shift = scalar.addi %scale_shift, %c4_i32 : i32 + %selected_scale_source = scf.select %is_low_group, %scale0, %scale2 : i32 + %selected_minimum_source = scf.select %is_low_group, %scale1, %scale2 : i32 + %selected_scale_high_shift = scf.select %is_low_group, %scale_shift, %high_shift : i32 + %selected_minimum_low_shift = scf.select %is_low_group, %scale_shift, %minimum_shift : i32 + %scale_low0 = scalar.shrui %selected_scale_source, %scale_shift : i32 + %scale_low = scalar.andi %scale_low0, %c15_i32 : i32 + %scale_high0 = scalar.shrui %scale0, %selected_scale_high_shift : i32 + %scale_high = scalar.andi %scale_high0, %c48_i32 : i32 + %scale = scalar.ori %scale_low, %scale_high : i32 + %minimum_low0 = scalar.shrui %selected_minimum_source, %selected_minimum_low_shift : i32 + %minimum_low = scalar.andi %minimum_low0, %c15_i32 : i32 + %minimum_high0 = scalar.shrui %scale1, %selected_scale_high_shift : i32 + %minimum_high = scalar.andi %minimum_high0, %c48_i32 : i32 + %minimum = scalar.ori %minimum_low, %minimum_high : i32 + func.return %scale, %minimum : i32, i32 +} + +// Decodes two adjacent Q4_K groups from the low and high nibbles of one +// packed-code load. The caller supplies the already-loaded 16-byte header. +func.def inline @qwen3_moe_q4k_chunk_pair_global(%weight: buffer, %row_byte_base: offset, %q4_block: index, %q4_group_pair: index, %q4_half: index, %header_words: vector<4xi32>) -> (vector<16xi8>, f32, f32, vector<16xi8>, f32, f32) { + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c4 = index.constant 4 : index + %c8 = index.constant 8 : index + %block_bytes = index.constant 144 : offset + %code_byte_add = index.constant 16 : offset + %c4_i32 = scalar.constant 4 : i32 + %nibble_mask = vector.constant 252645135 : vector<4xi32> + %bounded_pair = index.assume %q4_group_pair [range(%q4_group_pair, 0, 3)] : index + %bounded_half = index.assume %q4_half [range(%q4_half, 0, 1)] : index + %block_byte_add = index.scale %q4_block, %block_bytes : index, offset -> offset + %block_byte_base = index.add %row_byte_base, %block_byte_add : offset + %code_byte_base = index.add %block_byte_base, %code_byte_add : offset + %code_view = buffer.view %weight[%code_byte_base] : buffer -> view<32xi32> + %header_halves = vector.bitcast %header_words : vector<4xi32> to vector<8xf16> + %d_f16 = vector.extract %header_halves[0] : vector<8xf16> -> f16 + %dmin_f16 = vector.extract %header_halves[1] : vector<8xf16> -> f16 + %scale0 = vector.extract %header_words[1] : vector<4xi32> -> i32 + %scale1 = vector.extract %header_words[2] : vector<4xi32> -> i32 + %scale2 = vector.extract %header_words[3] : vector<4xi32> -> i32 + %d = scalar.extf %d_f16 : f16 to f32 + %dmin = scalar.extf %dmin_f16 : f16 to f32 + %pair_code_base = index.mul %bounded_pair, %c8 : index + %half_code_add = index.mul %bounded_half, %c4 : index + %code_index0 = index.add %pair_code_base, %half_code_add : index + %code_index = index.assume %code_index0 [range(%code_index0, 0, 28)] : index + %packed_codes = vector.load %code_view[%code_index] : view<32xi32> -> vector<4xi32> + %low_codes = vector.andi %packed_codes, %nibble_mask : vector<4xi32> + %c4_i32v = vector.splat %c4_i32 : vector<4xi32> + %high_shifted = vector.shrui %packed_codes, %c4_i32v : vector<4xi32> + %high_codes = vector.andi %high_shifted, %nibble_mask : vector<4xi32> + %q4_low = vector.bitcast %low_codes : vector<4xi32> to vector<16xi8> + %q4_high = vector.bitcast %high_codes : vector<4xi32> to vector<16xi8> + %low_group = index.mul %bounded_pair, %c2 : index + %high_group = index.add %low_group, %c1 : index + %low_scale, %low_minimum = func.call @qwen3_moe_q4k_scale_from_header(%scale0, %scale1, %scale2, %low_group) : (i32, i32, i32, index) -> (i32, i32) + %high_scale, %high_minimum = func.call @qwen3_moe_q4k_scale_from_header(%scale0, %scale1, %scale2, %high_group) : (i32, i32, i32, index) -> (i32, i32) + %low_scale_f32 = scalar.uitofp %low_scale : i32 to f32 + %low_minimum_f32 = scalar.uitofp %low_minimum : i32 to f32 + %high_scale_f32 = scalar.uitofp %high_scale : i32 to f32 + %high_minimum_f32 = scalar.uitofp %high_minimum : i32 to f32 + %low_d_scale = scalar.mulf %d, %low_scale_f32 : f32 + %low_dmin_scale = scalar.mulf %dmin, %low_minimum_f32 : f32 + %high_d_scale = scalar.mulf %d, %high_scale_f32 : f32 + %high_dmin_scale = scalar.mulf %dmin, %high_minimum_f32 : f32 + func.return %q4_low, %low_d_scale, %low_dmin_scale, %q4_high, %high_d_scale, %high_dmin_scale : vector<16xi8>, f32, f32, vector<16xi8>, f32, f32 +} + +// Contracts one decoded Q4_K half-group with a Q8_1 half-block. Both +// half-groups apply half of the Q8_1 block-sum correction. +func.def inline @qwen3_moe_q4k_q8_1_dot(%q4_values: vector<16xi8>, %d_scale: f32, %dmin_scale: f32, %q8_values: vector<16xi8>, %q8_d: f32, %q8_s: f32) -> (f32) { + %c0_i32 = scalar.constant 0 : i32 + %c0_i32v = vector.constant 0 : vector<4xi32> + %half_f32 = scalar.constant 0.5 : f32 + %partial_dots = vector.dot4i %q4_values, %q8_values, %c0_i32v : vector<16xi8>, vector<16xi8>, vector<4xi32> + %q_sum = vector.reduce %partial_dots, %c0_i32 : vector<4xi32>, i32 + %q_sum_f32 = scalar.sitofp %q_sum : i32 to f32 + %scaled_dot0 = scalar.mulf %q8_d, %d_scale : f32 + %scaled_dot = scalar.mulf %scaled_dot0, %q_sum_f32 : f32 + %q8_half_sum = scalar.mulf %q8_s, %half_f32 : f32 + %minimum_correction = scalar.mulf %dmin_scale, %q8_half_sum : f32 + %contribution = scalar.subf %scaled_dot, %minimum_correction : f32 + func.return %contribution : f32 +} + +// Contracts one decoded Q4_K half-group against four routed Q8_1 rows. The +// leading row axis stays explicit through dot4 and scale correction so the +// grouped schedule can reuse each weight decode without scalarizing the rows. +func.def inline @qwen3_moe_q4k_q8_1_dot4_rows(%q4_values: vector<16xi8>, %d_scale: f32, %dmin_scale: f32, %q8_values: vector<4x16xi8>, %q8_d: vector<4xf32>, %q8_s: vector<4xf32>) -> (vector<4xf32>) { + %q8_values0 = vector.extract %q8_values[0] : vector<4x16xi8> -> vector<16xi8> + %q8_values1 = vector.extract %q8_values[1] : vector<4x16xi8> -> vector<16xi8> + %q8_values2 = vector.extract %q8_values[2] : vector<4x16xi8> -> vector<16xi8> + %q8_values3 = vector.extract %q8_values[3] : vector<4x16xi8> -> vector<16xi8> + %q8_d0 = vector.extract %q8_d[0] : vector<4xf32> -> f32 + %q8_d1 = vector.extract %q8_d[1] : vector<4xf32> -> f32 + %q8_d2 = vector.extract %q8_d[2] : vector<4xf32> -> f32 + %q8_d3 = vector.extract %q8_d[3] : vector<4xf32> -> f32 + %q8_s0 = vector.extract %q8_s[0] : vector<4xf32> -> f32 + %q8_s1 = vector.extract %q8_s[1] : vector<4xf32> -> f32 + %q8_s2 = vector.extract %q8_s[2] : vector<4xf32> -> f32 + %q8_s3 = vector.extract %q8_s[3] : vector<4xf32> -> f32 + %contribution0 = func.call @qwen3_moe_q4k_q8_1_dot(%q4_values, %d_scale, %dmin_scale, %q8_values0, %q8_d0, %q8_s0) : (vector<16xi8>, f32, f32, vector<16xi8>, f32, f32) -> (f32) + %contribution1 = func.call @qwen3_moe_q4k_q8_1_dot(%q4_values, %d_scale, %dmin_scale, %q8_values1, %q8_d1, %q8_s1) : (vector<16xi8>, f32, f32, vector<16xi8>, f32, f32) -> (f32) + %contribution2 = func.call @qwen3_moe_q4k_q8_1_dot(%q4_values, %d_scale, %dmin_scale, %q8_values2, %q8_d2, %q8_s2) : (vector<16xi8>, f32, f32, vector<16xi8>, f32, f32) -> (f32) + %contribution3 = func.call @qwen3_moe_q4k_q8_1_dot(%q4_values, %d_scale, %dmin_scale, %q8_values3, %q8_d3, %q8_s3) : (vector<16xi8>, f32, f32, vector<16xi8>, f32, f32) -> (f32) + %contribution = vector.from_elements %contribution0, %contribution1, %contribution2, %contribution3 : vector<4xf32> + func.return %contribution : vector<4xf32> +} + +// Convenience wrapper used by the one-wave schedule. +func.def inline @qwen3_moe_q4k_q8_1_chunk(%weight: buffer, %row_byte_base: offset, %q4_block: index, %q4_group: index, %q4_half: index, %q8_values: vector<16xi8>, %q8_d: f32, %q8_s: f32) -> (f32) { + %q4_values, %d_scale, %dmin_scale = func.call @qwen3_moe_q4k_chunk_global(%weight, %row_byte_base, %q4_block, %q4_group, %q4_half) : (buffer, offset, index, index, index) -> (vector<16xi8>, f32, f32) + %contribution = func.call @qwen3_moe_q4k_q8_1_dot(%q4_values, %d_scale, %dmin_scale, %q8_values, %q8_d, %q8_s) : (vector<16xi8>, f32, f32, vector<16xi8>, f32, f32) -> (f32) + func.return %contribution : f32 +} + +// Computes one lane's partial for a complete Q4_K row. Dense and routed +// providers choose output ownership independently, then share this exact +// packed-row contraction before reducing across a wave32 subgroup. +func.def inline @qwen3_moe_q4k_q8_1_x4_row_lane(%input_size: index, %weight: buffer, %weight_row_byte_base: offset, %q8_input: buffer, %q8_row_byte_base: offset, %lane: index) -> (f32) { + %bounded_input_size = index.assume %input_size [range(%input_size, 256, 32768), mul(%input_size, 256)] : index + %bounded_lane = index.assume %lane [range(%lane, 0, 31)] : index + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c4 = index.constant 4 : index + %c8 = index.constant 8 : index + %c16 = index.constant 16 : index + %c128 = index.constant 128 : index + %c144 = index.constant 144 : index + %c256 = index.constant 256 : index + %c511 = index.constant 511 : index + %c512 = index.constant 512 : index + %q8_group_bytes = index.constant 144 : offset + %q8_payload_byte_add = index.constant 16 : offset + %c0_f32 = scalar.constant 0.0 : f32 + %q4_block_count = index.div %bounded_input_size, %c256 : index + %q8_group_count = index.div %bounded_input_size, %c128 : index + %padded_input_size = index.add %bounded_input_size, %c511 : index + %iteration_count = index.div %padded_input_size, %c512 : index + %q4_group0 = index.div %bounded_lane, %c2 : index + %q4_group1 = index.rem %q4_group0, %c8 : index + %q4_group = index.assume %q4_group1 [range(%q4_group1, 0, 7)] : index + %q4_half0 = index.rem %bounded_lane, %c2 : index + %q4_half = index.assume %q4_half0 [range(%q4_half0, 0, 1)] : index + %lane_q4_block = index.div %bounded_lane, %c16 : index + %lane_q8_group = index.div %bounded_lane, %c8 : index + %q8_inner_block0 = index.div %bounded_lane, %c2 : index + %q8_inner_block1 = index.rem %q8_inner_block0, %c4 : index + %q8_inner_block = index.assume %q8_inner_block1 [range(%q8_inner_block1, 0, 3)] : index + %q8_inner_word_base = index.mul %q8_inner_block, %c8 : index + %q8_half_word_add = index.mul %q4_half, %c4 : index + %q8_word_index0 = index.add %q8_inner_word_base, %q8_half_word_add : index + %q8_word_index = index.assume %q8_word_index0 [range(%q8_word_index0, 0, 28)] : index + %q8_ds_index = index.mul %q8_inner_block, %c2 : index + %q8_s_index = index.add %q8_ds_index, %c1 : index + %sum = scf.for %iteration = [%c0 to %iteration_count step %c1](%iteration_acc = %c0_f32 : f32) -> (f32) unroll { + %iteration_q4_block = index.mul %iteration, %c2 : index + %q4_block0 = index.add %iteration_q4_block, %lane_q4_block : index + %valid_q4_block = index.cmp ult, %q4_block0, %q4_block_count : index + %contribution = scf.if %valid_q4_block -> (f32) { + %q4_block, %bounded_q4_block_count = index.assume %q4_block0, %q4_block_count [lt(%q4_block0, %q4_block_count)] : index, index + %iteration_q8_group = index.mul %iteration, %c4 : index + %q8_group0 = index.add %iteration_q8_group, %lane_q8_group : index + %q8_group, %bounded_q8_group_count = index.assume %q8_group0, %q8_group_count [lt(%q8_group0, %q8_group_count)] : index, index + %q8_group_byte_add = index.scale %q8_group, %q8_group_bytes : index, offset -> offset + %q8_group_byte_base = index.add %q8_row_byte_base, %q8_group_byte_add : offset + %q8_payload_byte_base = index.add %q8_group_byte_base, %q8_payload_byte_add : offset + %q8_ds_view = buffer.view %q8_input[%q8_group_byte_base] : buffer -> view<8xf16> + %q8_words_view = buffer.view %q8_input[%q8_payload_byte_base] : buffer -> view<32xi32> + %q8_d_f16 = view.load %q8_ds_view[%q8_ds_index] : view<8xf16> -> f16 + %q8_s_f16 = view.load %q8_ds_view[%q8_s_index] : view<8xf16> -> f16 + %q8_d = scalar.extf %q8_d_f16 : f16 to f32 + %q8_s = scalar.extf %q8_s_f16 : f16 to f32 + %q8_words = vector.load %q8_words_view[%q8_word_index] : view<32xi32> -> vector<4xi32> + %q8_values = vector.bitcast %q8_words : vector<4xi32> to vector<16xi8> + %dot = func.call @qwen3_moe_q4k_q8_1_chunk(%weight, %weight_row_byte_base, %q4_block, %q4_group, %q4_half, %q8_values, %q8_d, %q8_s) : (buffer, offset, index, index, index, vector<16xi8>, f32, f32) -> (f32) + scf.yield %dot : f32 + } else { + scf.yield %c0_f32 : f32 + } + %next = scalar.addf %iteration_acc, %contribution : f32 + scf.yield %next : f32 + } + func.return %sum : f32 +} + +// Computes one lane's two adjacent-group contributions within one Q4_K block. +// Eight block lanes consume both nibbles from every packed code load and cover +// the block's full 256-element input extent. +func.def inline @qwen3_moe_q4k_q8_1_x4_paired_block_lane(%input_size: index, %weight: buffer, %weight_row_byte_base: offset, %q8_input: buffer, %q8_row_byte_base: offset, %q4_block: index, %block_lane: index) -> (f32) { + %bounded_input_size = index.assume %input_size [range(%input_size, 256, 32768), mul(%input_size, 256)] : index + %bounded_q4_block0 = index.assume %q4_block [range(%q4_block, 0, 127)] : index + %bounded_block_lane = index.assume %block_lane [range(%block_lane, 0, 7)] : index + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c4 = index.constant 4 : index + %c8 = index.constant 8 : index + %c128 = index.constant 128 : index + %c256 = index.constant 256 : index + %q4_block_bytes = index.constant 144 : offset + %q8_group_bytes = index.constant 144 : offset + %q8_payload_byte_add = index.constant 16 : offset + %q4_block_count = index.div %bounded_input_size, %c256 : index + %q8_group_count = index.div %bounded_input_size, %c128 : index + %bounded_q4_block, %bounded_q4_block_count = index.assume %bounded_q4_block0, %q4_block_count [lt(%bounded_q4_block0, %q4_block_count)] : index, index + %q4_group_pair0 = index.div %bounded_block_lane, %c2 : index + %q4_group_pair = index.assume %q4_group_pair0 [range(%q4_group_pair0, 0, 3)] : index + %q4_half0 = index.rem %bounded_block_lane, %c2 : index + %q4_half = index.assume %q4_half0 [range(%q4_half0, 0, 1)] : index + %q8_group_in_block0 = index.div %q4_group_pair, %c2 : index + %q8_group_in_block = index.assume %q8_group_in_block0 [range(%q8_group_in_block0, 0, 1)] : index + %pair_in_q8_group0 = index.rem %q4_group_pair, %c2 : index + %pair_in_q8_group = index.assume %pair_in_q8_group0 [range(%pair_in_q8_group0, 0, 1)] : index + %q8_low_inner_block0 = index.mul %pair_in_q8_group, %c2 : index + %q8_low_inner_block = index.assume %q8_low_inner_block0 [range(%q8_low_inner_block0, 0, 2)] : index + %q8_high_inner_block0 = index.add %q8_low_inner_block, %c1 : index + %q8_high_inner_block = index.assume %q8_high_inner_block0 [range(%q8_high_inner_block0, 1, 3)] : index + %q8_half_word_add = index.mul %q4_half, %c4 : index + %q8_low_inner_word_base = index.mul %q8_low_inner_block, %c8 : index + %q8_low_word_index0 = index.add %q8_low_inner_word_base, %q8_half_word_add : index + %q8_low_word_index = index.assume %q8_low_word_index0 [range(%q8_low_word_index0, 0, 20)] : index + %q8_high_inner_word_base = index.mul %q8_high_inner_block, %c8 : index + %q8_high_word_index0 = index.add %q8_high_inner_word_base, %q8_half_word_add : index + %q8_high_word_index = index.assume %q8_high_word_index0 [range(%q8_high_word_index0, 8, 28)] : index + %q8_low_ds_index0 = index.mul %q8_low_inner_block, %c2 : index + %q8_low_ds_index = index.assume %q8_low_ds_index0 [range(%q8_low_ds_index0, 0, 4)] : index + %q8_block_group_base = index.mul %bounded_q4_block, %c2 : index + %q8_group0 = index.add %q8_block_group_base, %q8_group_in_block : index + %q8_group, %bounded_q8_group_count = index.assume %q8_group0, %q8_group_count [lt(%q8_group0, %q8_group_count)] : index, index + %q8_group_byte_add = index.scale %q8_group, %q8_group_bytes : index, offset -> offset + %q8_group_byte_base = index.add %q8_row_byte_base, %q8_group_byte_add : offset + %q8_payload_byte_base = index.add %q8_group_byte_base, %q8_payload_byte_add : offset + %q8_ds_view = buffer.view %q8_input[%q8_group_byte_base] : buffer -> view<8xf16> + %q8_words_view = buffer.view %q8_input[%q8_payload_byte_base] : buffer -> view<32xi32> + %q8_ds = vector.load %q8_ds_view[%q8_low_ds_index] : view<8xf16> -> vector<4xf16> + %q8_low_d_f16 = vector.extract %q8_ds[0] : vector<4xf16> -> f16 + %q8_low_s_f16 = vector.extract %q8_ds[1] : vector<4xf16> -> f16 + %q8_high_d_f16 = vector.extract %q8_ds[2] : vector<4xf16> -> f16 + %q8_high_s_f16 = vector.extract %q8_ds[3] : vector<4xf16> -> f16 + %q8_low_d = scalar.extf %q8_low_d_f16 : f16 to f32 + %q8_low_s = scalar.extf %q8_low_s_f16 : f16 to f32 + %q8_high_d = scalar.extf %q8_high_d_f16 : f16 to f32 + %q8_high_s = scalar.extf %q8_high_s_f16 : f16 to f32 + %q8_low_words = vector.load %q8_words_view[%q8_low_word_index] : view<32xi32> -> vector<4xi32> + %q8_high_words = vector.load %q8_words_view[%q8_high_word_index] : view<32xi32> -> vector<4xi32> + %q8_low_values = vector.bitcast %q8_low_words : vector<4xi32> to vector<16xi8> + %q8_high_values = vector.bitcast %q8_high_words : vector<4xi32> to vector<16xi8> + %q4_block_byte_add = index.scale %bounded_q4_block, %q4_block_bytes : index, offset -> offset + %q4_block_byte_base = index.add %weight_row_byte_base, %q4_block_byte_add : offset + %q4_header_view = buffer.view %weight[%q4_block_byte_base] : buffer -> view<4xi32> + %q4_header_words = vector.load %q4_header_view[0] : view<4xi32> -> vector<4xi32> + %q4_low, %low_d_scale, %low_dmin_scale, %q4_high, %high_d_scale, %high_dmin_scale = func.call @qwen3_moe_q4k_chunk_pair_global(%weight, %weight_row_byte_base, %bounded_q4_block, %q4_group_pair, %q4_half, %q4_header_words) : (buffer, offset, index, index, index, vector<4xi32>) -> (vector<16xi8>, f32, f32, vector<16xi8>, f32, f32) + %low = func.call @qwen3_moe_q4k_q8_1_dot(%q4_low, %low_d_scale, %low_dmin_scale, %q8_low_values, %q8_low_d, %q8_low_s) : (vector<16xi8>, f32, f32, vector<16xi8>, f32, f32) -> (f32) + %high = func.call @qwen3_moe_q4k_q8_1_dot(%q4_high, %high_d_scale, %high_dmin_scale, %q8_high_values, %q8_high_d, %q8_high_s) : (vector<16xi8>, f32, f32, vector<16xi8>, f32, f32) -> (f32) + %pair = scalar.addf %low, %high : f32 + func.return %pair : f32 +} + +// Computes one lane's partial while consuming both nibbles of each packed +// Q4_K code load. Eight lanes cover one 256-element Q4_K block, so a wave32 +// advances through four blocks per iteration and reduces the adjacent-group +// contributions together. +func.def inline @qwen3_moe_q4k_q8_1_x4_paired_row_lane(%input_size: index, %weight: buffer, %weight_row_byte_base: offset, %q8_input: buffer, %q8_row_byte_base: offset, %lane: index) -> (f32) { + %bounded_input_size = index.assume %input_size [range(%input_size, 256, 32768), mul(%input_size, 256)] : index + %bounded_lane = index.assume %lane [range(%lane, 0, 31)] : index + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c4 = index.constant 4 : index + %c8 = index.constant 8 : index + %c256 = index.constant 256 : index + %c1023 = index.constant 1023 : index + %c1024 = index.constant 1024 : index + %c0_f32 = scalar.constant 0.0 : f32 + %q4_block_count = index.div %bounded_input_size, %c256 : index + %padded_input_size = index.add %bounded_input_size, %c1023 : index + %iteration_count = index.div %padded_input_size, %c1024 : index + %lane_q4_block = index.div %bounded_lane, %c8 : index + %block_lane0 = index.rem %bounded_lane, %c8 : index + %block_lane = index.assume %block_lane0 [range(%block_lane0, 0, 7)] : index + %sum = scf.for %iteration = [%c0 to %iteration_count step %c1](%iteration_acc = %c0_f32 : f32) -> (f32) unroll { + %iteration_q4_block = index.mul %iteration, %c4 : index + %q4_block0 = index.add %iteration_q4_block, %lane_q4_block : index + %valid_q4_block = index.cmp ult, %q4_block0, %q4_block_count : index + %contribution = scf.if %valid_q4_block -> (f32) { + %q4_block, %bounded_q4_block_count = index.assume %q4_block0, %q4_block_count [lt(%q4_block0, %q4_block_count)] : index, index + %pair = func.call @qwen3_moe_q4k_q8_1_x4_paired_block_lane(%bounded_input_size, %weight, %weight_row_byte_base, %q8_input, %q8_row_byte_base, %q4_block, %block_lane) : (index, buffer, offset, buffer, offset, index, index) -> (f32) + scf.yield %pair : f32 + } else { + scf.yield %c0_f32 : f32 + } + %next = scalar.addf %iteration_acc, %contribution : f32 + scf.yield %next : f32 + } + func.return %sum : f32 +} + +// Computes one lane's partial for a row owned by an eight-lane cohort. Every +// cohort lane consumes both nibbles for one adjacent Q4_K group pair while the +// cohort walks all blocks in the row. This schedule lets separate cohorts +// contract independent routed rows concurrently. +func.def inline @qwen3_moe_q4k_q8_1_x4_cohort_row_lane(%input_size: index, %weight: buffer, %weight_row_byte_base: offset, %q8_input: buffer, %q8_row_byte_base: offset, %cohort_lane: index) -> (f32) { + %bounded_input_size = index.assume %input_size [range(%input_size, 256, 32768), mul(%input_size, 256)] : index + %bounded_cohort_lane = index.assume %cohort_lane [range(%cohort_lane, 0, 7)] : index + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c256 = index.constant 256 : index + %c0_f32 = scalar.constant 0.0 : f32 + %q4_block_count = index.div %bounded_input_size, %c256 : index + %sum = scf.for %q4_block0 = [%c0 to %q4_block_count step %c1](%block_acc = %c0_f32 : f32) -> (f32) unroll { + %q4_block, %bounded_q4_block_count = index.assume %q4_block0, %q4_block_count [lt(%q4_block0, %q4_block_count)] : index, index + %pair = func.call @qwen3_moe_q4k_q8_1_x4_paired_block_lane(%bounded_input_size, %weight, %weight_row_byte_base, %q8_input, %q8_row_byte_base, %q4_block, %bounded_cohort_lane) : (index, buffer, offset, buffer, offset, index, index) -> (f32) + %next = scalar.addf %block_acc, %pair : f32 + scf.yield %next : f32 + } + func.return %sum : f32 +} + +// Decodes the packed queue descriptor shared by expert-grouped projections. +// The producer stores a 7-bit expert ordinal, a 6-bit 32-row partition +// ordinal, and a 5-bit row count minus one. +func.def inline @qwen3_moe_unpack_expert_partition_descriptor(%descriptor: i32) -> (index, index, index) { + %c1_i32 = scalar.constant 1 : i32 + %c5_i32 = scalar.constant 5 : i32 + %c7_i32 = scalar.constant 7 : i32 + %c13_i32 = scalar.constant 13 : i32 + %c31_i32 = scalar.constant 31 : i32 + %c63_i32 = scalar.constant 63 : i32 + %c127_i32 = scalar.constant 127 : i32 + %expert_i32 = scalar.andi %descriptor, %c127_i32 : i32 + %partition_shifted_i32 = scalar.shrui %descriptor, %c7_i32 : i32 + %partition_i32 = scalar.andi %partition_shifted_i32, %c63_i32 : i32 + %route_tile_base_i32 = scalar.shli %partition_i32, %c5_i32 : i32 + %row_count_shifted_i32 = scalar.shrui %descriptor, %c13_i32 : i32 + %row_count_minus_one_i32 = scalar.andi %row_count_shifted_i32, %c31_i32 : i32 + %partition_row_count_i32 = scalar.addi %row_count_minus_one_i32, %c1_i32 : i32 + %expert0 = index.cast %expert_i32 : i32 to index + %expert = index.assume %expert0 [range(%expert0, 0, 127)] : index + %route_tile_base0 = index.cast %route_tile_base_i32 : i32 to index + %route_tile_base = index.assume %route_tile_base0 [range(%route_tile_base0, 0, 2016)] : index + %partition_row_count0 = index.cast %partition_row_count_i32 : i32 to index + %partition_row_count = index.assume %partition_row_count0 [range(%partition_row_count0, 1, 32)] : index + func.return %expert, %route_tile_base, %partition_row_count : index, index, index +} + +// Builds the transient expert table consumed by the grouped projection. The +// table packs [expert_count] route counts followed by +// [expert_count][token_count] compact assignment ordinals. Top-k routing +// selects each expert at most once per token, so token_count entries are +// sufficient for every expert even though there are route_count assignments +// per token. +kernel.def @qwen3_moe_build_expert_table(%token_count: index, %route_count: index, %route_stride: index, %expert_count: index) { + %configured_expert_count = config.get @qwen3_moe.routed_gate_up.expert_count : index + %c1 = index.constant 1 : index + %workgroup_size = index.constant 256 : index + kernel.launch.config workgroups(%configured_expert_count, %c1, %c1) workgroup_size(%workgroup_size, %c1, %c1) : index +} launch(%token_count: index, %route_count: index, %route_stride: index, %expert_count: index, %route_ids: buffer, %expert_table: buffer) { + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 2048)] : index + %configured_route_count0 = config.get @qwen3_moe.routed_gate_up.route_count : index + %bounded_route_count, %configured_route_count = index.assume %route_count, %configured_route_count0 [range(%route_count, 1, 8), eq(%route_count, %configured_route_count0)] : index, index + %bounded_route_stride = index.assume %route_stride [range(%route_stride, 1, 512)] : index + %configured_expert_count0 = config.get @qwen3_moe.routed_gate_up.expert_count : index + %bounded_expert_count, %configured_expert_count = index.assume %expert_count, %configured_expert_count0 [range(%expert_count, 1, 512), eq(%expert_count, %configured_expert_count0)] : index, index + %expert0 = kernel.workgroup.id : index + %lane = kernel.workitem.id : index + %c0 = index.constant 0 : index + %workgroup_size = index.constant 256 : index + %c0_i32 = scalar.constant 0 : i32 + %c1_i32 = scalar.constant 1 : i32 + %c0_offset = index.constant 0 : offset + %c4_bytes = index.constant 4 : offset + %assignment_count = index.mul %bounded_token_count, %bounded_route_count : index + %expert, %table_expert_count = index.assume %expert0, %bounded_expert_count [lt(%expert0, %bounded_expert_count)] : index, index + %assignment_table_byte_base = index.scale %table_expert_count, %c4_bytes : index, offset -> offset + %route_view = buffer.view %route_ids[%c0_offset] : buffer -> view<[%bounded_token_count]x[%bounded_route_stride]xi32> + %count_view = buffer.view %expert_table[%c0_offset] : buffer -> view<[%table_expert_count]xi32> + %assignment_view = buffer.view %expert_table[%assignment_table_byte_base] : buffer -> view<[%table_expert_count]x[%bounded_token_count]xi32> + %expert_route_count = scf.for %block_base = [%c0 to %assignment_count step %workgroup_size](%matched_base = %c0_i32 : i32) -> (i32) { + %assignment = index.add %block_base, %lane : index + %in_range = index.cmp ult, %assignment, %assignment_count : index + %route_expert_i32 = scf.if %in_range -> (i32) { + %token0 = index.div %assignment, %configured_route_count : index + %route0 = index.rem %assignment, %configured_route_count : index + %token, %route_token_count = index.assume %token0, %bounded_token_count [lt(%token0, %bounded_token_count)] : index, index + // The route stride is the physical row width and must contain every + // logical top-k route. + %route, %route_row_stride = index.assume %route0, %bounded_route_stride [lt(%route0, %bounded_route_stride)] : index, index + %loaded = view.load %route_view[%token, %route] : view<[%bounded_token_count]x[%bounded_route_stride]xi32> -> i32 + scf.yield %loaded : i32 + } else { + %cn1_i32 = scalar.constant -1 : i32 + scf.yield %cn1_i32 : i32 + } + %route_expert0 = index.cast %route_expert_i32 : i32 to index + %route_expert = index.assume %route_expert0 [range(%route_expert0, -1, 511)] : index + %matches = index.cmp eq, %route_expert, %expert : index + %match_i32 = scf.if %matches -> (i32) { + scf.yield %c1_i32 : i32 + } else { + scf.yield %c0_i32 : i32 + } + %block_prefix = kernel.workgroup.scan %match_i32 {direction = forward, mode = exclusive} : i32 + %block_match_count_reduced = kernel.workgroup.reduce %match_i32 : i32 + %block_match_count = kernel.subgroup.broadcast.first %block_match_count_reduced : i32 + scf.if %matches { + %match_ordinal_i32 = scalar.addi %matched_base, %block_prefix : i32 + %match_ordinal0 = index.cast %match_ordinal_i32 : i32 to index + %match_ordinal = index.assume %match_ordinal0 [range(%match_ordinal0, 0, 2047)] : index + // Top-k route IDs are unique within a token, so one expert can own at + // most token_count assignments. + %bounded_match_ordinal, %table_token_count = index.assume %match_ordinal, %bounded_token_count [lt(%match_ordinal, %bounded_token_count)] : index, index + %assignment_i32 = index.cast %assignment : index to i32 + view.store %assignment_i32, %assignment_view[%expert, %bounded_match_ordinal] : i32, view<[%table_expert_count]x[%bounded_token_count]xi32> + } + %next_matched_base = scalar.addi %matched_base, %block_match_count : i32 + scf.yield %next_matched_base : i32 + } + %is_lane_zero = index.cmp eq, %lane, %c0 : index + scf.if %is_lane_zero { + view.store %expert_route_count, %count_view[%expert] : i32, view<[%table_expert_count]xi32> + } + kernel.return +} + +// Compacts expert assignment counts into exact 32-row projection partitions. +// +// One lane owns each expert. A workgroup scan assigns deterministic descriptor +// offsets, then each lane publishes an 18-bit descriptor containing the expert, +// 32-row partition ordinal, and tail row count. The consumer can launch a +// distribution-independent grid without serializing a concentrated expert +// inside one workgroup. +kernel.def @qwen3_moe_build_expert_partition_table(%token_count: index, %route_count: index, %expert_count: index) { + %c1 = index.constant 1 : index + %workgroup_size = index.constant 128 : index + kernel.launch.config workgroups(%c1, %c1, %c1) workgroup_size(%workgroup_size, %c1, %c1) : index +} launch(%token_count: index, %route_count: index, %expert_count: index, %expert_table: buffer, %partition_table: buffer) { + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 2048)] : index + %bounded_route_count = index.assume %route_count [range(%route_count, 1, 8)] : index + %bounded_expert_count = index.assume %expert_count [range(%expert_count, 1, 128)] : index + %lane = kernel.workitem.id : index + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c31 = index.constant 31 : index + %c32 = index.constant 32 : index + %c0_i32 = scalar.constant 0 : i32 + %c1_i32 = scalar.constant 1 : i32 + %c7_i32 = scalar.constant 7 : i32 + %c13_i32 = scalar.constant 13 : i32 + %c0_offset = index.constant 0 : offset + %c4_bytes = index.constant 4 : offset + %assignment_count = index.mul %bounded_token_count, %bounded_route_count : index + %rounded_assignment_count = index.add %assignment_count, %c31 : index + %assignment_partition_count = index.div %rounded_assignment_count, %c32 : index + %maximum_partition_count = index.add %assignment_partition_count, %bounded_expert_count : index + %count_view = buffer.view %expert_table[%c0_offset] : buffer -> view<[%bounded_expert_count]xi32> + %partition_count_view = buffer.view %partition_table[%c0_offset] : buffer -> view<1xi32> + %partition_descriptor_view = buffer.view %partition_table[%c4_bytes] : buffer -> view<[%maximum_partition_count]xi32> + %has_expert = index.cmp ult, %lane, %bounded_expert_count : index + %expert_assignment_count_i32 = scf.if %has_expert -> (i32) { + %expert, %table_expert_count = index.assume %lane, %bounded_expert_count [lt(%lane, %bounded_expert_count)] : index, index + %loaded = view.load %count_view[%expert] : view<[%bounded_expert_count]xi32> -> i32 + scf.yield %loaded : i32 + } else { + scf.yield %c0_i32 : i32 + } + %expert_assignment_count0 = index.cast %expert_assignment_count_i32 : i32 to index + %expert_assignment_count = index.assume %expert_assignment_count0 [range(%expert_assignment_count0, 0, 2048)] : index + %rounded_expert_assignment_count = index.add %expert_assignment_count, %c31 : index + %expert_partition_count = index.div %rounded_expert_assignment_count, %c32 : index + %expert_partition_count_i32 = index.cast %expert_partition_count : index to i32 + %expert_partition_base_i32 = kernel.workgroup.scan %expert_partition_count_i32 {direction = forward, mode = exclusive} : i32 + %partition_count_i32 = kernel.workgroup.reduce %expert_partition_count_i32 : i32 + %partition_count = index.cast %partition_count_i32 : i32 to index + %bounded_partition_count, %table_partition_capacity = index.assume %partition_count, %maximum_partition_count [lt(%partition_count, %maximum_partition_count)] : index, index + %expert_partition_base0 = index.cast %expert_partition_base_i32 : i32 to index + %expert_partition_base = index.assume %expert_partition_base0 [range(%expert_partition_base0, 0, 639)] : index + scf.if %has_expert { + %expert, %table_expert_count = index.assume %lane, %bounded_expert_count [lt(%lane, %bounded_expert_count)] : index, index + %expert_i32 = index.cast %expert : index to i32 + scf.for %partition = [%c0 to %expert_partition_count step %c1] { + %descriptor_ordinal0 = index.add %expert_partition_base, %partition : index + %descriptor_ordinal, %descriptor_count = index.assume %descriptor_ordinal0, %bounded_partition_count [lt(%descriptor_ordinal0, %bounded_partition_count)] : index, index + %table_descriptor_ordinal, %table_descriptor_capacity = index.assume %descriptor_ordinal, %maximum_partition_count [lt(%descriptor_ordinal, %maximum_partition_count)] : index, index + %partition_remainder = index.rem %expert_assignment_count, %c32 : index + %has_partial_tail = index.cmp ne, %partition_remainder, %c0 : index + %partition_row_count = scf.if %has_partial_tail -> (index) { + %next_partition = index.add %partition, %c1 : index + %is_tail_partition = index.cmp eq, %next_partition, %expert_partition_count : index + %tail_row_count = scf.if %is_tail_partition -> (index) { + scf.yield %partition_remainder : index + } else { + scf.yield %c32 : index + } + scf.yield %tail_row_count : index + } else { + scf.yield %c32 : index + } + %partition_i32 = index.cast %partition : index to i32 + %partition_row_count_i32 = index.cast %partition_row_count : index to i32 + %packed_partition = scalar.shli %partition_i32, %c7_i32 : i32 + %partition_row_count_minus_one = scalar.subi %partition_row_count_i32, %c1_i32 : i32 + %packed_row_count = scalar.shli %partition_row_count_minus_one, %c13_i32 : i32 + %packed_expert_partition = scalar.ori %expert_i32, %packed_partition : i32 + %packed_descriptor = scalar.ori %packed_expert_partition, %packed_row_count : i32 + view.store %packed_descriptor, %partition_descriptor_view[%table_descriptor_ordinal] : i32, view<[%maximum_partition_count]xi32> + } + } + %is_lane_zero = index.cmp eq, %lane, %c0 : index + scf.if %is_lane_zero { + view.store %partition_count_i32, %partition_count_view[%c0] : i32, view<1xi32> + } + kernel.return +} + +// Small-token body for routed Q4_K gate/up projections. One wave computes one +// [token, route, output channel]. Callers own launch geometry and may append a +// producer epilogue after the F32 SwiGLU value is visible. +func.def inline @qwen3_moe_routed_gate_up_swiglu_q4k_q8_body(%publish_output: i1, %token_count: index, %token: index, %route_count: index, %route: index, %route_stride: index, %expert_count: index, %output_size: index, %channel: index, %lane: index, %q8_input: buffer, %route_ids: buffer, %gate_weight: buffer, %up_weight: buffer, %output: buffer) { + %input_size = config.get @qwen3_moe.routed_gate_up.input_size : index + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 512)] : index + %bounded_route_count, %bounded_route_stride = index.assume %route_count, %route_stride [range(%route_count, 1, 8), range(%route_stride, 1, 128), le(%route_count, %route_stride)] : index, index + %bounded_expert_count = index.assume %expert_count [range(%expert_count, 1, 128)] : index + %bounded_output_size = index.assume %output_size [range(%output_size, 1, 4096)] : index + %bounded_token, %body_token_count = index.assume %token, %bounded_token_count [lt(%token, %bounded_token_count)] : index, index + %bounded_channel, %body_output_size = index.assume %channel, %bounded_output_size [lt(%channel, %bounded_output_size)] : index, index + %bounded_route, %body_route_count = index.assume %route, %bounded_route_count [lt(%route, %bounded_route_count)] : index, index + %bounded_lane = index.assume %lane [range(%lane, 0, 31)] : index + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c4 = index.constant 4 : index + %c8 = index.constant 8 : index + %c128 = index.constant 128 : index + %c256 = index.constant 256 : index + %c1023 = index.constant 1023 : index + %c1024 = index.constant 1024 : index + %q4_block_bytes = index.constant 144 : index + %q4_block_bytes_offset = index.constant 144 : offset + %q8_group_bytes = index.constant 144 : offset + %q8_payload_byte_add = index.constant 16 : offset + %c1_byte = index.constant 1 : offset + %c0_f32 = scalar.constant 0.0 : f32 + %c0_i32 = scalar.constant 0 : i32 + %c0_offset = index.constant 0 : offset + %q4_block_count = index.div %input_size, %c256 : index + %weight_row_bytes = index.mul %q4_block_count, %q4_block_bytes : index + %weight_expert_bytes = index.mul %bounded_output_size, %weight_row_bytes : index + %q8_group_count = index.div %input_size, %c128 : index + %q8_bytes_per_token = index.mul %q8_group_count, %q4_block_bytes : index + %padded_input_size = index.add %input_size, %c1023 : index + %iteration_count = index.div %padded_input_size, %c1024 : index + %q8_noalias, %route_ids_noalias, %gate_noalias, %up_noalias, %output_noalias = buffer.assume.noalias %q8_input, %route_ids, %gate_weight, %up_weight, %output : buffer, buffer, buffer, buffer, buffer + %route_ids_view = buffer.view %route_ids_noalias[%c0_offset] : buffer -> view<[%bounded_token_count]x[%bounded_route_stride]xi32> + %output_view = buffer.view %output_noalias[%c0_offset] : buffer -> view<[%bounded_token_count]x[%bounded_route_count]x[%bounded_output_size]xf32> + %route_index, %route_row_stride = index.assume %bounded_route, %bounded_route_stride [lt(%bounded_route, %bounded_route_stride)] : index, index + %expert_i32 = view.load %route_ids_view[%bounded_token, %route_index] : view<[%bounded_token_count]x[%bounded_route_stride]xi32> -> i32 + %expert0 = index.cast %expert_i32 : i32 to index + %expert = index.assume %expert0 [range(%expert0, 0, 127)] : index + %expert_byte_base = index.mul %expert, %weight_expert_bytes : index + %channel_byte_add = index.mul %bounded_channel, %weight_row_bytes : index + %row_byte_index = index.add %expert_byte_base, %channel_byte_add : index + %row_byte_base = index.scale %row_byte_index, %c1_byte : index, offset -> offset + %q8_token_byte_index = index.mul %bounded_token, %q8_bytes_per_token : index + %q8_token_byte_base = index.scale %q8_token_byte_index, %c1_byte : index, offset -> offset + %q4_group_pair0 = index.div %bounded_lane, %c2 : index + %q4_group_pair1 = index.rem %q4_group_pair0, %c4 : index + %q4_group_pair = index.assume %q4_group_pair1 [range(%q4_group_pair1, 0, 3)] : index + %q4_half0 = index.rem %bounded_lane, %c2 : index + %q4_half = index.assume %q4_half0 [range(%q4_half0, 0, 1)] : index + %lane_q4_block = index.div %bounded_lane, %c8 : index + %q8_group_in_block0 = index.div %q4_group_pair, %c2 : index + %q8_group_in_block = index.assume %q8_group_in_block0 [range(%q8_group_in_block0, 0, 1)] : index + %pair_in_q8_group0 = index.rem %q4_group_pair, %c2 : index + %pair_in_q8_group = index.assume %pair_in_q8_group0 [range(%pair_in_q8_group0, 0, 1)] : index + %q8_low_inner_block0 = index.mul %pair_in_q8_group, %c2 : index + %q8_low_inner_block = index.assume %q8_low_inner_block0 [range(%q8_low_inner_block0, 0, 2)] : index + %q8_high_inner_block0 = index.add %q8_low_inner_block, %c1 : index + %q8_high_inner_block = index.assume %q8_high_inner_block0 [range(%q8_high_inner_block0, 1, 3)] : index + %q8_half_word_add = index.mul %q4_half, %c4 : index + %q8_low_inner_word_base = index.mul %q8_low_inner_block, %c8 : index + %q8_low_word_index0 = index.add %q8_low_inner_word_base, %q8_half_word_add : index + %q8_low_word_index = index.assume %q8_low_word_index0 [range(%q8_low_word_index0, 0, 20)] : index + %q8_high_inner_word_base = index.mul %q8_high_inner_block, %c8 : index + %q8_high_word_index0 = index.add %q8_high_inner_word_base, %q8_half_word_add : index + %q8_high_word_index = index.assume %q8_high_word_index0 [range(%q8_high_word_index0, 8, 28)] : index + %q8_low_ds_index0 = index.mul %q8_low_inner_block, %c2 : index + %q8_low_ds_index = index.assume %q8_low_ds_index0 [range(%q8_low_ds_index0, 0, 4)] : index + %gate_acc, %up_acc = scf.for %iteration = [%c0 to %iteration_count step %c1](%gate_iter = %c0_f32 : f32, %up_iter = %c0_f32 : f32) -> (f32, f32) unroll { + %iteration_q4_block = index.mul %iteration, %c4 : index + %q4_block0 = index.add %iteration_q4_block, %lane_q4_block : index + %valid_q4_block = index.cmp ult, %q4_block0, %q4_block_count : index + %gate_contribution, %up_contribution = scf.if %valid_q4_block -> (f32, f32) { + %q4_block, %bounded_q4_block_count = index.assume %q4_block0, %q4_block_count [lt(%q4_block0, %q4_block_count)] : index, index + %q8_block_group_base = index.mul %q4_block, %c2 : index + %q8_group0 = index.add %q8_block_group_base, %q8_group_in_block : index + %q8_group, %bounded_q8_group_count = index.assume %q8_group0, %q8_group_count [lt(%q8_group0, %q8_group_count)] : index, index + %q8_group_byte_add = index.scale %q8_group, %q8_group_bytes : index, offset -> offset + %q8_group_byte_base = index.add %q8_token_byte_base, %q8_group_byte_add : offset + %q8_payload_byte_base = index.add %q8_group_byte_base, %q8_payload_byte_add : offset + %q8_ds_view = buffer.view %q8_noalias[%q8_group_byte_base] : buffer -> view<8xf16> + %q8_words_view = buffer.view %q8_noalias[%q8_payload_byte_base] : buffer -> view<32xi32> + %q8_ds = vector.load %q8_ds_view[%q8_low_ds_index] : view<8xf16> -> vector<4xf16> + %q8_low_d_f16 = vector.extract %q8_ds[0] : vector<4xf16> -> f16 + %q8_low_s_f16 = vector.extract %q8_ds[1] : vector<4xf16> -> f16 + %q8_high_d_f16 = vector.extract %q8_ds[2] : vector<4xf16> -> f16 + %q8_high_s_f16 = vector.extract %q8_ds[3] : vector<4xf16> -> f16 + %q8_low_d = scalar.extf %q8_low_d_f16 : f16 to f32 + %q8_low_s = scalar.extf %q8_low_s_f16 : f16 to f32 + %q8_high_d = scalar.extf %q8_high_d_f16 : f16 to f32 + %q8_high_s = scalar.extf %q8_high_s_f16 : f16 to f32 + %q8_low_words = vector.load %q8_words_view[%q8_low_word_index] : view<32xi32> -> vector<4xi32> + %q8_high_words = vector.load %q8_words_view[%q8_high_word_index] : view<32xi32> -> vector<4xi32> + %q8_low_values = vector.bitcast %q8_low_words : vector<4xi32> to vector<16xi8> + %q8_high_values = vector.bitcast %q8_high_words : vector<4xi32> to vector<16xi8> + %q4_block_byte_add = index.scale %q4_block, %q4_block_bytes_offset : index, offset -> offset + %q4_block_byte_base = index.add %row_byte_base, %q4_block_byte_add : offset + %gate_header_view = buffer.view %gate_noalias[%q4_block_byte_base] : buffer -> view<4xi32> + %up_header_view = buffer.view %up_noalias[%q4_block_byte_base] : buffer -> view<4xi32> + %gate_header_words = vector.load %gate_header_view[0] : view<4xi32> -> vector<4xi32> + %up_header_words = vector.load %up_header_view[0] : view<4xi32> -> vector<4xi32> + %gate_q4_low, %gate_low_d_scale, %gate_low_dmin_scale, %gate_q4_high, %gate_high_d_scale, %gate_high_dmin_scale = func.call @qwen3_moe_q4k_chunk_pair_global(%gate_noalias, %row_byte_base, %q4_block, %q4_group_pair, %q4_half, %gate_header_words) : (buffer, offset, index, index, index, vector<4xi32>) -> (vector<16xi8>, f32, f32, vector<16xi8>, f32, f32) + %up_q4_low, %up_low_d_scale, %up_low_dmin_scale, %up_q4_high, %up_high_d_scale, %up_high_dmin_scale = func.call @qwen3_moe_q4k_chunk_pair_global(%up_noalias, %row_byte_base, %q4_block, %q4_group_pair, %q4_half, %up_header_words) : (buffer, offset, index, index, index, vector<4xi32>) -> (vector<16xi8>, f32, f32, vector<16xi8>, f32, f32) + %gate_low = func.call @qwen3_moe_q4k_q8_1_dot(%gate_q4_low, %gate_low_d_scale, %gate_low_dmin_scale, %q8_low_values, %q8_low_d, %q8_low_s) : (vector<16xi8>, f32, f32, vector<16xi8>, f32, f32) -> (f32) + %gate_high = func.call @qwen3_moe_q4k_q8_1_dot(%gate_q4_high, %gate_high_d_scale, %gate_high_dmin_scale, %q8_high_values, %q8_high_d, %q8_high_s) : (vector<16xi8>, f32, f32, vector<16xi8>, f32, f32) -> (f32) + %up_low = func.call @qwen3_moe_q4k_q8_1_dot(%up_q4_low, %up_low_d_scale, %up_low_dmin_scale, %q8_low_values, %q8_low_d, %q8_low_s) : (vector<16xi8>, f32, f32, vector<16xi8>, f32, f32) -> (f32) + %up_high = func.call @qwen3_moe_q4k_q8_1_dot(%up_q4_high, %up_high_d_scale, %up_high_dmin_scale, %q8_high_values, %q8_high_d, %q8_high_s) : (vector<16xi8>, f32, f32, vector<16xi8>, f32, f32) -> (f32) + %gate_pair = scalar.addf %gate_low, %gate_high : f32 + %up_pair = scalar.addf %up_low, %up_high : f32 + scf.yield %gate_pair, %up_pair : f32, f32 + } else { + scf.yield %c0_f32, %c0_f32 : f32, f32 + } + %gate_next = scalar.addf %gate_iter, %gate_contribution : f32 + %up_next = scalar.addf %up_iter, %up_contribution : f32 + scf.yield %gate_next, %up_next : f32, f32 + } + %gate_dot = kernel.subgroup.reduce %gate_acc : f32 + %up_dot = kernel.subgroup.reduce %up_acc : f32 + %lane_i32 = index.cast %bounded_lane : index to i32 + %is_lane_zero = scalar.cmpi eq, %lane_i32, %c0_i32 : i32 + %writes_output = scalar.andi %publish_output, %is_lane_zero : i1 + scf.if %writes_output { + %gate_silu = scalar.siluf %gate_dot : f32 + %result = scalar.mulf %gate_silu, %up_dot : f32 + view.store %result, %output_view[%bounded_token, %bounded_route, %bounded_channel] : f32, view<[%bounded_token_count]x[%bounded_route_count]x[%bounded_output_size]xf32> + } + func.return +} + +// Small-token schedule for routed Q4_K gate/up projections. Each workgroup +// packs four independent channel waves for one [token, route] pair. Eight +// lanes consume both nibbles of the packed codes for one 256-element block, so +// each wave advances through K in 1024-element stripes. Gate and up share each +// paired Q8_1 load before independent raw-weight dot products and a fused +// SwiGLU epilogue. +kernel.def @qwen3_moe_routed_gate_up_swiglu_q4k_q8(%token_count: index, %route_count: index, %route_stride: index, %expert_count: index, %output_size: index) { + %configured_route_count = config.get @qwen3_moe.routed_gate_up.route_count : index + %configured_output_size = config.get @qwen3_moe.routed_gate_up.output_size : index + %unit = index.constant 1 : index + %c3 = index.constant 3 : index + %c4 = index.constant 4 : index + %workgroup_size = index.constant 128 : index + %padded_output_size = index.add %configured_output_size, %c3 : index + %channel_workgroup_count = index.div %padded_output_size, %c4 : index + kernel.launch.config workgroups(%channel_workgroup_count, %configured_route_count, %token_count) workgroup_size(%workgroup_size, %unit, %unit) : index +} launch(%token_count: index, %route_count: index, %route_stride: index, %expert_count: index, %output_size: index, %q8_input: buffer, %route_ids: buffer, %gate_weight: buffer, %up_weight: buffer, %output: buffer) { + %configured_route_count0 = config.get @qwen3_moe.routed_gate_up.route_count : index + %configured_expert_count0 = config.get @qwen3_moe.routed_gate_up.expert_count : index + %configured_output_size0 = config.get @qwen3_moe.routed_gate_up.output_size : index + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 512)] : index + %bounded_route_count, %configured_route_count = index.assume %route_count, %configured_route_count0 [range(%route_count, 1, 8), eq(%route_count, %configured_route_count0)] : index, index + %bounded_route_stride = index.assume %route_stride [range(%route_stride, 1, 128)] : index + %bounded_expert_count, %configured_expert_count = index.assume %expert_count, %configured_expert_count0 [range(%expert_count, 1, 128), eq(%expert_count, %configured_expert_count0)] : index, index + %bounded_output_size, %configured_output_size = index.assume %output_size, %configured_output_size0 [range(%output_size, 1, 4096), eq(%output_size, %configured_output_size0)] : index, index + %token0 = kernel.workgroup.id : index + %channel_workgroup = kernel.workgroup.id : index + %route0 = kernel.workgroup.id : index + %subgroup = kernel.subgroup.id : index + %lane = kernel.subgroup.lane.id : index + %c0 = index.constant 0 : index + %c4 = index.constant 4 : index + %valid_token = index.cmp ult, %token0, %bounded_token_count : index + %safe_token0 = scf.select %valid_token, %token0, %c0 : index + %safe_token, %body_token_count = index.assume %safe_token0, %bounded_token_count [lt(%safe_token0, %bounded_token_count)] : index, index + %channel_base = index.mul %channel_workgroup, %c4 : index + %channel0 = index.add %channel_base, %subgroup : index + %valid_channel = index.cmp ult, %channel0, %bounded_output_size : index + %safe_channel0 = scf.select %valid_channel, %channel0, %c0 : index + %safe_channel, %body_output_size = index.assume %safe_channel0, %bounded_output_size [lt(%safe_channel0, %bounded_output_size)] : index, index + %route, %body_route_count = index.assume %route0, %bounded_route_count [lt(%route0, %bounded_route_count)] : index, index + %publishes_output = scalar.andi %valid_token, %valid_channel : i1 + func.call @qwen3_moe_routed_gate_up_swiglu_q4k_q8_body(%publishes_output, %body_token_count, %safe_token, %body_route_count, %route, %bounded_route_stride, %bounded_expert_count, %body_output_size, %safe_channel, %lane, %q8_input, %route_ids, %gate_weight, %up_weight, %output) : (i1, index, index, index, index, index, index, index, index, index, buffer, buffer, buffer, buffer, buffer) + kernel.return +} + +// Packs one Q8_1 x4 physical group with one selected subgroup. Each lane owns +// four adjacent values and contributes its maximum and sum to one of four +// eight-lane logical blocks. One-hot vectors let ordinary subgroup reductions +// compute all four block aggregates without LDS or workgroup barriers. +func.def inline @qwen3_moe_routed_gate_up_quantize_q8_1_x4_subgroup_body(%publish_output: i1, %group_count0: index, %group0: index, %input: buffer, %output: buffer) { + %group_count, %group = index.assume %group_count0, %group0 [range(%group_count0, 1, 524288), lt(%group0, %group_count0)] : index, index + %lane0 = kernel.subgroup.lane.id : index + %lane = index.assume %lane0 [range(%lane0, 0, 31)] : index + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c4 = index.constant 4 : index + %c8 = index.constant 8 : index + %c128 = index.constant 128 : index + %group_bytes = index.constant 144 : offset + %payload_byte_add = index.constant 16 : offset + %c0_f32 = scalar.constant 0.0 : f32 + %c1_f32 = scalar.constant 1.0 : f32 + %c127_f32 = scalar.constant 127.0 : f32 + %c0_f32x4 = vector.constant 0.0 : vector<4xf32> + %c0_offset = index.constant 0 : offset + %launched_element_count = index.mul %group_count, %c128 : index + %group_element_base = index.mul %group, %c128 : index + %lane_element_add = index.mul %lane, %c4 : index + %input_index = index.add %group_element_base, %lane_element_add : index + %input_noalias, %output_noalias = buffer.assume.noalias %input, %output : buffer, buffer + %input_view = buffer.view %input_noalias[%c0_offset] : buffer -> view<[%launched_element_count]xf32> + %input_values = vector.load %input_view[%input_index] : view<[%launched_element_count]xf32> -> vector<4xf32> + %absolute_values = vector.absf %input_values : vector<4xf32> + %thread_maximum = vector.reduce %absolute_values, %c0_f32 : vector<4xf32>, f32 + %block0 = index.div %lane, %c8 : index + %block = index.assume %block0 [range(%block0, 0, 3)] : index + %block_maxima = vector.insert %thread_maximum into %c0_f32x4[%block] : f32, vector<4xf32> + %reduced_maxima = kernel.subgroup.reduce %block_maxima : vector<4xf32> + %amax = vector.extract %reduced_maxima[%block] : vector<4xf32> -> f32 + %d = scalar.divf %amax, %c127_f32 : f32 + %d_nonzero = scalar.cmpf one, %d, %c0_f32 : f32 + %d_inverse = scf.if %d_nonzero -> (f32) { + %inverse = scalar.divf %c1_f32, %d : f32 + scf.yield %inverse : f32 + } else { + scf.yield %c0_f32 : f32 + } + %d_inverse_vector = vector.splat %d_inverse : vector<4xf32> + %scaled_values = vector.mulf %input_values, %d_inverse_vector : vector<4xf32> + %rounded_values = vector.roundf %scaled_values : vector<4xf32> + %quantized_values = vector.fptosi %rounded_values : vector<4xf32> to vector<4xi8> + %packed_word = vector.bitcast %quantized_values : vector<4xi8> to vector<1xi32> + %group_byte_offset = index.scale %group, %group_bytes : index, offset -> offset + %payload_byte_offset = index.add %group_byte_offset, %payload_byte_add : offset + %group_ds = buffer.view %output_noalias[%group_byte_offset] : buffer -> view<8xf16> + %group_qs = buffer.view %output_noalias[%payload_byte_offset] : buffer -> view<32xi32> + scf.if %publish_output { + vector.store %packed_word, %group_qs[%lane] : vector<1xi32>, view<32xi32> + } + %thread_sum = vector.reduce %rounded_values, %c0_f32 : vector<4xf32>, f32 + %block_sums0 = vector.insert %thread_sum into %c0_f32x4[%block] : f32, vector<4xf32> + %reduced_sums = kernel.subgroup.reduce %block_sums0 : vector<4xf32> + %quantized_sum = vector.extract %reduced_sums[%block] : vector<4xf32> -> f32 + %s = scalar.mulf %quantized_sum, %d : f32 + %word_in_block = index.rem %lane, %c8 : index + %is_block_leader = index.cmp eq, %word_in_block, %c0 : index + %publishes_metadata = scalar.andi %publish_output, %is_block_leader : i1 + scf.if %publishes_metadata { + %d_f16 = scalar.fptrunc %d : f32 to f16 + %s_f16 = scalar.fptrunc %s : f32 to f16 + %ds_index = index.mul %block, %c2 : index + %s_index = index.add %ds_index, %c1 : index + view.store %d_f16, %group_ds[%ds_index] : f16, view<8xf16> + view.store %s_f16, %group_ds[%s_index] : f16, view<8xf16> + } + func.return +} + +// Decode producer that publishes both the ordinary F32 SwiGLU rows and their +// packed Q8_1 x4 representation. Four independent channel waves share one +// workgroup and publish one completion arrival. The last workgroup to complete +// a 128-channel physical group packs all four logical Q8_1 blocks with one +// selected subgroup, then resets the counter before returning. +kernel.def @qwen3_moe_routed_gate_up_swiglu_q4k_q8_1_x4_next_q8(%token_count: index, %route_count: index, %route_stride: index, %expert_count: index, %output_size: index) { + %configured_route_count = config.get @qwen3_moe.routed_gate_up.route_count : index + %configured_output_size = config.get @qwen3_moe.routed_gate_up.output_size : index + %unit = index.constant 1 : index + %c3 = index.constant 3 : index + %c4 = index.constant 4 : index + %workgroup_size = index.constant 128 : index + %padded_output_size = index.add %configured_output_size, %c3 : index + %channel_workgroup_count = index.div %padded_output_size, %c4 : index + kernel.launch.config workgroups(%channel_workgroup_count, %configured_route_count, %unit) workgroup_size(%workgroup_size, %unit, %unit) : index +} launch(%token_count: index, %route_count: index, %route_stride: index, %expert_count: index, %output_size: index, %q8_input: buffer, %route_ids: buffer, %gate_weight: buffer, %up_weight: buffer, %output: buffer, %completion_counters: buffer, %next_q8_output: buffer) { + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 1)] : index + %configured_route_count0 = config.get @qwen3_moe.routed_gate_up.route_count : index + %configured_expert_count0 = config.get @qwen3_moe.routed_gate_up.expert_count : index + %configured_output_size0 = config.get @qwen3_moe.routed_gate_up.output_size : index + %bounded_route_count, %configured_route_count = index.assume %route_count, %configured_route_count0 [range(%route_count, 1, 8), eq(%route_count, %configured_route_count0)] : index, index + %bounded_route_stride = index.assume %route_stride [range(%route_stride, 1, 128)] : index + %bounded_expert_count, %configured_expert_count = index.assume %expert_count, %configured_expert_count0 [range(%expert_count, 1, 128), eq(%expert_count, %configured_expert_count0)] : index, index + %bounded_output_size, %configured_output_size = index.assume %output_size, %configured_output_size0 [range(%output_size, 128, 4096), mul(%output_size, 128), eq(%output_size, %configured_output_size0)] : index, index + %q8_input_noalias, %route_ids_noalias, %gate_weight_noalias, %up_weight_noalias, %output_noalias, %completion_counters_noalias, %next_q8_output_noalias = buffer.assume.noalias %q8_input, %route_ids, %gate_weight, %up_weight, %output, %completion_counters, %next_q8_output : buffer, buffer, buffer, buffer, buffer, buffer, buffer + %publishes_swiglu = scalar.constant true : i1 + %body_token = index.constant 0 : index + %channel_workgroup = kernel.workgroup.id : index + %route = kernel.workgroup.id : index + %token = kernel.workgroup.id : index + %subgroup = kernel.subgroup.id : index + %lane = kernel.subgroup.lane.id : index + %workitem = kernel.workitem.id : index + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c4 = index.constant 4 : index + %c128 = index.constant 128 : index + %c0_i32 = scalar.constant 0 : i32 + %c4_i32 = scalar.constant 4 : i32 + %c0_offset = index.constant 0 : offset + %counter_scratch_bytes = index.constant 4 : offset + %channel_base = index.mul %channel_workgroup, %c4 : index + %channel = index.add %channel_base, %subgroup : index + func.call @qwen3_moe_routed_gate_up_swiglu_q4k_q8_body(%publishes_swiglu, %bounded_token_count, %body_token, %bounded_route_count, %route, %bounded_route_stride, %bounded_expert_count, %bounded_output_size, %channel, %lane, %q8_input_noalias, %route_ids_noalias, %gate_weight_noalias, %up_weight_noalias, %output_noalias) : (i1, index, index, index, index, index, index, index, index, index, buffer, buffer, buffer, buffer, buffer) + %physical_group_count = index.div %bounded_output_size, %c128 : index + %row_count = index.mul %bounded_token_count, %bounded_route_count : index + %completion_counter_count = index.mul %row_count, %physical_group_count : index + %token_row_base = index.mul %token, %bounded_route_count : index + %row = index.add %token_row_base, %route : index + %row_group_base = index.mul %row, %physical_group_count : index + %group_in_row = index.div %channel_base, %c128 : index + %counter_index0 = index.add %row_group_base, %group_in_row : index + %counter_index, %bounded_completion_counter_count = index.assume %counter_index0, %completion_counter_count [lt(%counter_index0, %completion_counter_count)] : index, index + %completion_counters_aligned = buffer.assume.alignment %completion_counters_noalias {minimum_alignment = 16} : buffer + %completion_counters_view = buffer.view %completion_counters_aligned[%c0_offset] : buffer -> view<[%bounded_completion_counter_count]xi32> + %counter_scratch = buffer.alloca align(4) %counter_scratch_bytes : buffer + %counter_scratch_view = buffer.view %counter_scratch[%c0_offset] : buffer -> view<1xi32> + %is_arrival_lane = index.cmp eq, %workitem, %c0 : index + // Publish every producer lane's SwiGLU store before the leader advances one + // workgroup arrival. The last arrival then acquires the physical group. + kernel.barrier scope(workgroup) ordering(release) + scf.if %is_arrival_lane { + %old_counter = view.atomic.rmw %c4_i32, %completion_counters_view[%counter_index] {ordering = acq_rel, scope = device} : i32, view<[%bounded_completion_counter_count]xi32> -> i32 + view.store %old_counter, %counter_scratch_view[%c0] : i32, view<1xi32> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + %old_counter = view.load %counter_scratch_view[%c0] : view<1xi32> -> i32 + %group_size_i32 = index.cast %c128 : index to i32 + %last_arrival_i32 = scalar.subi %group_size_i32, %c4_i32 : i32 + %negative_group_size_i32 = scalar.subi %c0_i32, %group_size_i32 : i32 + %is_last_arrival = scalar.cmpi eq, %old_counter, %last_arrival_i32 : i32 + scf.if %is_last_arrival { + kernel.barrier scope(workgroup) ordering(acquire) + %publish_output = scalar.constant true : i1 + %is_quantize_subgroup = index.cmp eq, %subgroup, %c0 : index + scf.if %is_quantize_subgroup { + func.call @qwen3_moe_routed_gate_up_quantize_q8_1_x4_subgroup_body(%publish_output, %bounded_completion_counter_count, %counter_index, %output_noalias, %next_q8_output_noalias) : (i1, index, index, buffer, buffer) + } + kernel.barrier scope(workgroup) ordering(release) + scf.if %is_arrival_lane { + view.atomic.reduce %negative_group_size_i32, %completion_counters_view[%counter_index] {ordering = release, scope = device} : i32, view<[%bounded_completion_counter_count]xi32> + } + } + kernel.return +} + +// Large-token schedule for routed Q4_K gate/up projections. Each workgroup +// owns one expert, 32 output channels, and up to 32 selected rows. It +// reads those rows from the compact expert table, stages one K=256 slice of +// both raw Q4_K projections and Q8_1 activations in LDS, and reuses each +// decoded weight chunk across four rows before a fused SwiGLU epilogue. +kernel.def @qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped(%token_count: index, %route_count: index, %expert_count: index, %output_size: index) { + %token_capacity = config.get @qwen3_moe.workload.token_capacity : index + %configured_expert_count = config.get @qwen3_moe.routed_gate_up.expert_count : index + %configured_output_size = config.get @qwen3_moe.routed_gate_up.output_size : index + %c1 = index.constant 1 : index + %c4 = index.constant 4 : index + %c31 = index.constant 31 : index + %c32 = index.constant 32 : index + %workgroup_size = index.constant 256 : index + %padded_output_size = index.add %configured_output_size, %c31 : index + %output_tiles = index.div %padded_output_size, %c32 : index + %padded_token_count = index.add %token_capacity, %c31 : index + %route_tiles = index.div %padded_token_count, %c32 : index + %route_partitions = index.min %route_tiles, %c4 : index + kernel.launch.config workgroups(%output_tiles, %route_partitions, %configured_expert_count) workgroup_size(%workgroup_size, %c1, %c1) : index +} launch(%token_count: index, %route_count: index, %expert_count: index, %output_size: index, %q8_input: buffer, %expert_table: buffer, %gate_weight: buffer, %up_weight: buffer, %output: buffer) { + %input_size = config.get @qwen3_moe.routed_gate_up.input_size : index + %token_capacity = config.get @qwen3_moe.workload.token_capacity : index + %configured_route_count0 = config.get @qwen3_moe.routed_gate_up.route_count : index + %configured_expert_count0 = config.get @qwen3_moe.routed_gate_up.expert_count : index + %configured_output_size0 = config.get @qwen3_moe.routed_gate_up.output_size : index + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 512), le(%token_count, %token_capacity)] : index + %bounded_route_count, %configured_route_count = index.assume %route_count, %configured_route_count0 [range(%route_count, 1, 8), eq(%route_count, %configured_route_count0)] : index, index + %bounded_expert_count, %configured_expert_count = index.assume %expert_count, %configured_expert_count0 [range(%expert_count, 1, 128), eq(%expert_count, %configured_expert_count0)] : index, index + %bounded_output_size, %configured_output_size = index.assume %output_size, %configured_output_size0 [range(%output_size, 1, 4096), eq(%output_size, %configured_output_size0)] : index, index + %channel_tile = kernel.workgroup.id : index + %route_tile = kernel.workgroup.id : index + %expert0 = kernel.workgroup.id : index + %lane = kernel.workitem.id : index + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c4 = index.constant 4 : index + %c8 = index.constant 8 : index + %c32 = index.constant 32 : index + %c36 = index.constant 36 : index + %c72 = index.constant 72 : index + %c128 = index.constant 128 : index + %c256 = index.constant 256 : index + %gate_stage_word_count = index.constant 1152 : index + %q8_stage_word_count = index.constant 2304 : index + %c0_i32 = scalar.constant 0 : i32 + %c0_f32_rows = vector.constant 0.0 : vector<4xf32> + %c0_offset = index.constant 0 : offset + %c4_bytes = index.constant 4 : offset + %q4_block_bytes = index.constant 144 : offset + %q8_row_bytes = index.constant 288 : offset + %row_id_bytes = index.constant 128 : offset + %weight_stage_bytes = index.constant 4608 : offset + %q8_stage_bytes = index.constant 9216 : offset + %assignment_count = index.mul %bounded_token_count, %bounded_route_count : index + %expert, %table_expert_count = index.assume %expert0, %bounded_expert_count [lt(%expert0, %bounded_expert_count)] : index, index + %assignment_table_byte_base = index.scale %table_expert_count, %c4_bytes : index, offset -> offset + %q4_block_count = index.div %input_size, %c256 : index + %output_route_count = index.mul %bounded_token_count, %bounded_route_count : index + %q8_noalias, %expert_table_noalias, %gate_noalias, %up_noalias, %output_noalias = buffer.assume.noalias %q8_input, %expert_table, %gate_weight, %up_weight, %output : buffer, buffer, buffer, buffer, buffer + %q8_words = buffer.view %q8_noalias[%c0_offset] : buffer -> view<[%bounded_token_count]x[%q4_block_count]x72xi32> + %count_view = buffer.view %expert_table_noalias[%c0_offset] : buffer -> view<[%table_expert_count]xi32> + %assignment_view = buffer.view %expert_table_noalias[%assignment_table_byte_base] : buffer -> view<[%table_expert_count]x[%bounded_token_count]xi32> + %gate_words = buffer.view %gate_noalias[%c0_offset] : buffer -> view<[%bounded_expert_count]x[%bounded_output_size]x[%q4_block_count]x36xi32> + %up_words = buffer.view %up_noalias[%c0_offset] : buffer -> view<[%bounded_expert_count]x[%bounded_output_size]x[%q4_block_count]x36xi32> + %output_view = buffer.view %output_noalias[%c0_offset] : buffer -> view<[%output_route_count]x[%bounded_output_size]xf32> + %row_ids = buffer.alloca align(16) %row_id_bytes : buffer + %gate_stage = buffer.alloca align(16) %weight_stage_bytes : buffer + %up_stage = buffer.alloca align(16) %weight_stage_bytes : buffer + %q8_stage = buffer.alloca align(16) %q8_stage_bytes : buffer + %row_ids_view = buffer.view %row_ids[%c0_offset] : buffer -> view<32xi32> + %gate_stage_words = buffer.view %gate_stage[%c0_offset] : buffer -> view<1152xi32> + %up_stage_words = buffer.view %up_stage[%c0_offset] : buffer -> view<1152xi32> + %q8_stage_words = buffer.view %q8_stage[%c0_offset] : buffer -> view<2304xi32> + %channel_base = index.mul %channel_tile, %c32 : index + %channel0 = index.rem %lane, %c32 : index + %channel = index.add %channel_base, %channel0 : index + %row_base0 = index.div %lane, %c32 : index + %row_base = index.assume %row_base0 [range(%row_base0, 0, 7)] : index + %initial_tile_base = index.mul %route_tile, %c32 : index + // Four interleaved route partitions cover the first 128 rows. Smaller + // token counts execute at most one iteration; larger expert populations + // continue in 128-row strides without increasing the launch grid. + %route_partition_step = index.constant 128 : index + %is_lane_zero = index.cmp eq, %lane, %c0 : index + %lane_expert_route_count = scf.if %is_lane_zero -> (i32) { + %loaded = view.load %count_view[%expert] : view<[%table_expert_count]xi32> -> i32 + scf.yield %loaded : i32 + } else { + scf.yield %c0_i32 : i32 + } + %expert_route_count_reduced = kernel.workgroup.reduce %lane_expert_route_count : i32 + %expert_route_count_i32 = kernel.subgroup.broadcast.first %expert_route_count_reduced : i32 + %expert_route_count0 = index.cast %expert_route_count_i32 : i32 to index + %expert_route_count = index.assume %expert_route_count0 [range(%expert_route_count0, 0, 4096)] : index + scf.for %tile_base = [%initial_tile_base to %expert_route_count step %route_partition_step] { + %tile_base_i32 = index.cast %tile_base : index to i32 + %remaining_rows_i32 = scalar.subi %expert_route_count_i32, %tile_base_i32 : i32 + %remaining_rows0 = index.cast %remaining_rows_i32 : i32 to index + %remaining_rows = index.assume %remaining_rows0 [range(%remaining_rows0, 1, 4096)] : index + %has_full_tile = index.cmp uge, %remaining_rows, %c32 : index + %tile_row_count = scf.if %has_full_tile -> (index) { + scf.yield %c32 : index + } else { + scf.yield %remaining_rows : index + } + %loads_row = index.cmp ult, %lane, %tile_row_count : index + scf.if %loads_row { + %local_row = index.assume %lane [range(%lane, 0, 31)] : index + %expert_assignment_ordinal0 = index.add %tile_base, %local_row : index + %expert_assignment_ordinal, %table_token_count = index.assume %expert_assignment_ordinal0, %bounded_token_count [lt(%expert_assignment_ordinal0, %bounded_token_count)] : index, index + %assignment_i32 = view.load %assignment_view[%expert, %expert_assignment_ordinal] : view<[%table_expert_count]x[%bounded_token_count]xi32> -> i32 + view.store %assignment_i32, %row_ids_view[%local_row] : i32, view<32xi32> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + %gate_acc, %up_acc = scf.for %q4_block = [%c0 to %q4_block_count step %c1](%gate_block_acc = %c0_f32_rows : vector<4xf32>, %up_block_acc = %c0_f32_rows : vector<4xf32>) -> (vector<4xf32>, vector<4xf32>) { + %bounded_q4_block, %body_q4_block_count = index.assume %q4_block, %q4_block_count [lt(%q4_block, %q4_block_count)] : index, index + scf.for %stage_word = [%lane to %gate_stage_word_count step %c256] { + %local_channel = index.div %stage_word, %c36 : index + %word_in_block0 = index.rem %stage_word, %c36 : index + %word_in_block = index.assume %word_in_block0 [range(%word_in_block0, 0, 35)] : index + %global_channel = index.add %channel_base, %local_channel : index + %valid_channel = index.cmp ult, %global_channel, %bounded_output_size : index + %gate_word, %up_word = scf.if %valid_channel -> (i32, i32) { + %bounded_global_channel, %table_output_size = index.assume %global_channel, %bounded_output_size [lt(%global_channel, %bounded_output_size)] : index, index + %gate_loaded = view.load %gate_words[%expert, %bounded_global_channel, %bounded_q4_block, %word_in_block] : view<[%bounded_expert_count]x[%bounded_output_size]x[%q4_block_count]x36xi32> -> i32 + %up_loaded = view.load %up_words[%expert, %bounded_global_channel, %bounded_q4_block, %word_in_block] : view<[%bounded_expert_count]x[%bounded_output_size]x[%q4_block_count]x36xi32> -> i32 + scf.yield %gate_loaded, %up_loaded : i32, i32 + } else { + scf.yield %c0_i32, %c0_i32 : i32, i32 + } + view.store %gate_word, %gate_stage_words[%stage_word] : i32, view<1152xi32> + view.store %up_word, %up_stage_words[%stage_word] : i32, view<1152xi32> + } + scf.for %stage_word = [%lane to %q8_stage_word_count step %c256] { + %local_row = index.div %stage_word, %c72 : index + %word_in_row0 = index.rem %stage_word, %c72 : index + %word_in_row = index.assume %word_in_row0 [range(%word_in_row0, 0, 71)] : index + %valid_row = index.cmp ult, %local_row, %tile_row_count : index + %q8_word = scf.if %valid_row -> (i32) { + %bounded_local_row = index.assume %local_row [range(%local_row, 0, 31)] : index + %assignment_i32 = view.load %row_ids_view[%bounded_local_row] : view<32xi32> -> i32 + %assignment0 = index.cast %assignment_i32 : i32 to index + %assignment = index.assume %assignment0 [range(%assignment0, 0, 4095)] : index + %bounded_assignment, %bounded_assignment_count = index.assume %assignment, %assignment_count [lt(%assignment, %assignment_count)] : index, index + %token0 = index.div %bounded_assignment, %configured_route_count : index + %token, %table_token_count = index.assume %token0, %bounded_token_count [lt(%token0, %bounded_token_count)] : index, index + %loaded = view.load %q8_words[%token, %bounded_q4_block, %word_in_row] : view<[%bounded_token_count]x[%q4_block_count]x72xi32> -> i32 + scf.yield %loaded : i32 + } else { + scf.yield %c0_i32 : i32 + } + view.store %q8_word, %q8_stage_words[%stage_word] : i32, view<2304xi32> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + %channel_byte_base = index.scale %channel0, %q4_block_bytes : index, offset -> offset + %after_groups_gate, %after_groups_up = scf.for %q4_group = [%c0 to %c8 step %c1](%gate_group_acc = %gate_block_acc : vector<4xf32>, %up_group_acc = %up_block_acc : vector<4xf32>) -> (vector<4xf32>, vector<4xf32>) unroll { + %row1 = index.add %row_base, %c8 : index + %row2_add = index.constant 16 : index + %row2 = index.add %row_base, %row2_add : index + %row3_add = index.constant 24 : index + %row3 = index.add %row_base, %row3_add : index + %row0_byte_base = index.scale %row_base, %q8_row_bytes : index, offset -> offset + %row1_byte_base = index.scale %row1, %q8_row_bytes : index, offset -> offset + %row2_byte_base = index.scale %row2, %q8_row_bytes : index, offset -> offset + %row3_byte_base = index.scale %row3, %q8_row_bytes : index, offset -> offset + %q8_values0, %q8_d0, %q8_s0 = func.call @ggml_q8_1_x4_block(%q8_stage, %row0_byte_base, %q4_group) : (buffer, offset, index) -> (vector<32xi8>, f32, f32) + %q8_values1, %q8_d1, %q8_s1 = func.call @ggml_q8_1_x4_block(%q8_stage, %row1_byte_base, %q4_group) : (buffer, offset, index) -> (vector<32xi8>, f32, f32) + %q8_values2, %q8_d2, %q8_s2 = func.call @ggml_q8_1_x4_block(%q8_stage, %row2_byte_base, %q4_group) : (buffer, offset, index) -> (vector<32xi8>, f32, f32) + %q8_values3, %q8_d3, %q8_s3 = func.call @ggml_q8_1_x4_block(%q8_stage, %row3_byte_base, %q4_group) : (buffer, offset, index) -> (vector<32xi8>, f32, f32) + %q8_low0 = vector.slice %q8_values0[0] : vector<32xi8> -> vector<16xi8> + %q8_low1 = vector.slice %q8_values1[0] : vector<32xi8> -> vector<16xi8> + %q8_low2 = vector.slice %q8_values2[0] : vector<32xi8> -> vector<16xi8> + %q8_low3 = vector.slice %q8_values3[0] : vector<32xi8> -> vector<16xi8> + %q8_high0 = vector.slice %q8_values0[16] : vector<32xi8> -> vector<16xi8> + %q8_high1 = vector.slice %q8_values1[16] : vector<32xi8> -> vector<16xi8> + %q8_high2 = vector.slice %q8_values2[16] : vector<32xi8> -> vector<16xi8> + %q8_high3 = vector.slice %q8_values3[16] : vector<32xi8> -> vector<16xi8> + %q8_rows_low0 = vector.constant 0 : vector<4x16xi8> + %q8_rows_low1 = vector.insert %q8_low0 into %q8_rows_low0[0] : vector<16xi8>, vector<4x16xi8> + %q8_rows_low2 = vector.insert %q8_low1 into %q8_rows_low1[1] : vector<16xi8>, vector<4x16xi8> + %q8_rows_low3 = vector.insert %q8_low2 into %q8_rows_low2[2] : vector<16xi8>, vector<4x16xi8> + %q8_rows_low = vector.insert %q8_low3 into %q8_rows_low3[3] : vector<16xi8>, vector<4x16xi8> + %q8_rows_high0 = vector.constant 0 : vector<4x16xi8> + %q8_rows_high1 = vector.insert %q8_high0 into %q8_rows_high0[0] : vector<16xi8>, vector<4x16xi8> + %q8_rows_high2 = vector.insert %q8_high1 into %q8_rows_high1[1] : vector<16xi8>, vector<4x16xi8> + %q8_rows_high3 = vector.insert %q8_high2 into %q8_rows_high2[2] : vector<16xi8>, vector<4x16xi8> + %q8_rows_high = vector.insert %q8_high3 into %q8_rows_high3[3] : vector<16xi8>, vector<4x16xi8> + %q8_d = vector.from_elements %q8_d0, %q8_d1, %q8_d2, %q8_d3 : vector<4xf32> + %q8_s = vector.from_elements %q8_s0, %q8_s1, %q8_s2, %q8_s3 : vector<4xf32> + %gate_q4_low, %gate_d_scale_low, %gate_dmin_scale_low = func.call @qwen3_moe_q4k_chunk_local(%gate_stage, %channel_byte_base, %c0, %q4_group, %c0) : (buffer, offset, index, index, index) -> (vector<16xi8>, f32, f32) + %gate_q4_high, %gate_d_scale_high, %gate_dmin_scale_high = func.call @qwen3_moe_q4k_chunk_local(%gate_stage, %channel_byte_base, %c0, %q4_group, %c1) : (buffer, offset, index, index, index) -> (vector<16xi8>, f32, f32) + %up_q4_low, %up_d_scale_low, %up_dmin_scale_low = func.call @qwen3_moe_q4k_chunk_local(%up_stage, %channel_byte_base, %c0, %q4_group, %c0) : (buffer, offset, index, index, index) -> (vector<16xi8>, f32, f32) + %up_q4_high, %up_d_scale_high, %up_dmin_scale_high = func.call @qwen3_moe_q4k_chunk_local(%up_stage, %channel_byte_base, %c0, %q4_group, %c1) : (buffer, offset, index, index, index) -> (vector<16xi8>, f32, f32) + %gate_low = func.call @qwen3_moe_q4k_q8_1_dot4_rows(%gate_q4_low, %gate_d_scale_low, %gate_dmin_scale_low, %q8_rows_low, %q8_d, %q8_s) : (vector<16xi8>, f32, f32, vector<4x16xi8>, vector<4xf32>, vector<4xf32>) -> (vector<4xf32>) + %gate_high = func.call @qwen3_moe_q4k_q8_1_dot4_rows(%gate_q4_high, %gate_d_scale_high, %gate_dmin_scale_high, %q8_rows_high, %q8_d, %q8_s) : (vector<16xi8>, f32, f32, vector<4x16xi8>, vector<4xf32>, vector<4xf32>) -> (vector<4xf32>) + %up_low = func.call @qwen3_moe_q4k_q8_1_dot4_rows(%up_q4_low, %up_d_scale_low, %up_dmin_scale_low, %q8_rows_low, %q8_d, %q8_s) : (vector<16xi8>, f32, f32, vector<4x16xi8>, vector<4xf32>, vector<4xf32>) -> (vector<4xf32>) + %up_high = func.call @qwen3_moe_q4k_q8_1_dot4_rows(%up_q4_high, %up_d_scale_high, %up_dmin_scale_high, %q8_rows_high, %q8_d, %q8_s) : (vector<16xi8>, f32, f32, vector<4x16xi8>, vector<4xf32>, vector<4xf32>) -> (vector<4xf32>) + %gate_pair = vector.addf %gate_low, %gate_high : vector<4xf32> + %up_pair = vector.addf %up_low, %up_high : vector<4xf32> + %gate_next = vector.addf %gate_group_acc, %gate_pair : vector<4xf32> + %up_next = vector.addf %up_group_acc, %up_pair : vector<4xf32> + scf.yield %gate_next, %up_next : vector<4xf32>, vector<4xf32> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + scf.yield %after_groups_gate, %after_groups_up : vector<4xf32>, vector<4xf32> + } + %gate_silu = vector.siluf %gate_acc : vector<4xf32> + %result_rows = vector.mulf %gate_silu, %up_acc : vector<4xf32> + %valid_channel = index.cmp ult, %channel, %bounded_output_size : index + scf.for %row_variant = [%c0 to %c4 step %c1] unroll { + %row_add = index.mul %row_variant, %c8 : index + %local_row = index.add %row_base, %row_add : index + %valid_row = index.cmp ult, %local_row, %tile_row_count : index + %writes_output = scalar.andi %valid_channel, %valid_row : i1 + scf.if %writes_output { + %bounded_local_row = index.assume %local_row [range(%local_row, 0, 31)] : index + %assignment_i32 = view.load %row_ids_view[%bounded_local_row] : view<32xi32> -> i32 + %assignment0 = index.cast %assignment_i32 : i32 to index + %assignment = index.assume %assignment0 [range(%assignment0, 0, 4095)] : index + %bounded_assignment, %bounded_assignment_count = index.assume %assignment, %assignment_count [lt(%assignment, %assignment_count)] : index, index + %result = vector.extract %result_rows[%row_variant] : vector<4xf32> -> f32 + view.store %result, %output_view[%bounded_assignment, %channel] : f32, view<[%output_route_count]x[%bounded_output_size]xf32> + } + } + } + kernel.return +} + +check.case public @qwen3_moe_routed_gate_up_swiglu_q4k_q8_nonzero_case { + %token_count = check.literal value(1) : index + %input_size = check.literal value(2048) : index + %route_count = check.literal value(1) : index + %route_stride = check.literal value(1) : index + %expert_count = check.literal value(1) : index + %output_size = check.literal value(1) : index + %input = check.generate.fill value(0.00390625) : tensor<1x2048xf32> + %q8_input = check.generate.fill value(0) : tensor<1x2304xi8> + %route_ids = check.generate.fill value(0) : tensor<1xi32> + %gate_weight = check.generate.fill value(85) : tensor<8x144xi8> + %up_weight = check.generate.fill value(-86) : tensor<8x144xi8> + %output = check.generate.fill value(0.0) : tensor<1xf32> + %expected = check.generate.fill value(-9024645.0) : tensor<1xf32> + kernel.launch @ggml_quantize_q8_1_x4_f32[%token_count, %input_size](%token_count, %input_size, %input, %q8_input) : [index, index](index, index, tensor<1x2048xf32>, tensor<1x2304xi8>) + kernel.launch @qwen3_moe_routed_gate_up_swiglu_q4k_q8[%token_count, %route_count, %route_stride, %expert_count, %output_size](%token_count, %route_count, %route_stride, %expert_count, %output_size, %q8_input, %route_ids, %gate_weight, %up_weight, %output) : [index, index, index, index, index](index, index, index, index, index, tensor<1x2304xi8>, tensor<1xi32>, tensor<8x144xi8>, tensor<8x144xi8>, tensor<1xf32>) + check.expect.close actual(%output) expected(%expected) atol(16.0) rtol(9.9999999999999995e-07) nan(same) : tensor<1xf32> + check.return +} + +// Exact decode topology with compact route IDs and every physical output +// group populated. Two fused invocations prove counter reuse in addition to +// comparing the F32 and packed outputs with the ordinary composition. +check.case public @qwen3_moe_routed_gate_up_swiglu_q4k_q8_1_x4_next_q8_differential_case { + %token_count = check.literal value(1) : index + %input_size = check.literal value(2048) : index + %route_count = check.literal value(8) : index + %route_stride = check.literal value(8) : index + %expert_count = check.literal value(128) : index + %output_size = check.literal value(768) : index + %routed_row_count = check.literal value(8) : index + %input = check.generate.fill value(0.00390625) : tensor<1x2048xf32> + %q8_input = check.generate.fill value(0) : tensor<2304xi8> + %route_ids = check.generate.iota offset(0) step(1) period(128) : tensor<8xi32> + %gate_weight = check.generate.iota offset(-72) step(1) period(144) : tensor<128x768x8x144xi8> + %up_weight = check.generate.iota offset(-71) step(1) period(144) : tensor<128x768x8x144xi8> + %expected_output = check.generate.fill value(0.0) : tensor<8x768xf32> + %expected_q8 = check.generate.fill value(0) : tensor<8x864xi8> + %actual_output0 = check.generate.fill value(1.0) : tensor<8x768xf32> + %actual_q8_0 = check.generate.fill value(1) : tensor<8x864xi8> + %actual_output1 = check.generate.fill value(2.0) : tensor<8x768xf32> + %actual_q8_1 = check.generate.fill value(2) : tensor<8x864xi8> + %completion_counters = check.generate.fill value(0) : tensor<48xi32> + %expected_counters = check.generate.fill value(0) : tensor<48xi32> + kernel.launch @ggml_quantize_q8_1_x4_f32[%token_count, %input_size](%token_count, %input_size, %input, %q8_input) : [index, index](index, index, tensor<1x2048xf32>, tensor<2304xi8>) + kernel.launch @qwen3_moe_routed_gate_up_swiglu_q4k_q8[%token_count, %route_count, %route_stride, %expert_count, %output_size](%token_count, %route_count, %route_stride, %expert_count, %output_size, %q8_input, %route_ids, %gate_weight, %up_weight, %expected_output) : [index, index, index, index, index](index, index, index, index, index, tensor<2304xi8>, tensor<8xi32>, tensor<128x768x8x144xi8>, tensor<128x768x8x144xi8>, tensor<8x768xf32>) + kernel.launch @ggml_quantize_q8_1_x4_f32[%routed_row_count, %output_size](%routed_row_count, %output_size, %expected_output, %expected_q8) : [index, index](index, index, tensor<8x768xf32>, tensor<8x864xi8>) + kernel.launch @qwen3_moe_routed_gate_up_swiglu_q4k_q8_1_x4_next_q8[%token_count, %route_count, %route_stride, %expert_count, %output_size](%token_count, %route_count, %route_stride, %expert_count, %output_size, %q8_input, %route_ids, %gate_weight, %up_weight, %actual_output0, %completion_counters, %actual_q8_0) : [index, index, index, index, index](index, index, index, index, index, tensor<2304xi8>, tensor<8xi32>, tensor<128x768x8x144xi8>, tensor<128x768x8x144xi8>, tensor<8x768xf32>, tensor<48xi32>, tensor<8x864xi8>) + kernel.launch @qwen3_moe_routed_gate_up_swiglu_q4k_q8_1_x4_next_q8[%token_count, %route_count, %route_stride, %expert_count, %output_size](%token_count, %route_count, %route_stride, %expert_count, %output_size, %q8_input, %route_ids, %gate_weight, %up_weight, %actual_output1, %completion_counters, %actual_q8_1) : [index, index, index, index, index](index, index, index, index, index, tensor<2304xi8>, tensor<8xi32>, tensor<128x768x8x144xi8>, tensor<128x768x8x144xi8>, tensor<8x768xf32>, tensor<48xi32>, tensor<8x864xi8>) + check.expect.close actual(%actual_output0) expected(%expected_output) atol(16.0) rtol(9.9999999999999995e-07) nan(same) : tensor<8x768xf32> + check.expect.close actual(%actual_output1) expected(%expected_output) atol(16.0) rtol(9.9999999999999995e-07) nan(same) : tensor<8x768xf32> + check.expect.equal actual(%actual_q8_0) expected(%expected_q8) : tensor<8x864xi8> + check.expect.equal actual(%actual_q8_1) expected(%expected_q8) : tensor<8x864xi8> + check.expect.equal actual(%completion_counters) expected(%expected_counters) : tensor<48xi32> + check.return +} + +check.case public @qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped_differential_case { + %token_count = check.literal value(64) : index + %input_size = check.literal value(512) : index + %route_count = check.literal value(2) : index + %route_stride = check.literal value(4) : index + %expert_count = check.literal value(4) : index + %output_size = check.literal value(32) : index + %input = check.generate.fill value(0.00390625) : tensor<64x512xf32> + %q8_input = check.generate.fill value(0) : tensor<64x576xi8> + // The physical row retains four argsort entries while the logical top-k view + // selects experts 0 and 1. Both experts span two 32-row grouped tiles. + %route_ids = check.generate.iota offset(0) step(1) period(4) : tensor<64x4xi32> + %expert_table = check.generate.fill value(-1) : tensor<260xi32> + %gate_weight = check.generate.iota offset(-72) step(1) period(144) : tensor<4x32x2x144xi8> + %up_weight = check.generate.iota offset(-71) step(1) period(144) : tensor<4x32x2x144xi8> + %expected = check.generate.fill value(0.0) : tensor<64x2x32xf32> + %actual = check.generate.fill value(1.0) : tensor<64x2x32xf32> + kernel.launch @ggml_quantize_q8_1_x4_f32[%token_count, %input_size](%token_count, %input_size, %input, %q8_input) : [index, index](index, index, tensor<64x512xf32>, tensor<64x576xi8>) + kernel.launch @qwen3_moe_build_expert_table[%token_count, %route_count, %route_stride, %expert_count](%token_count, %route_count, %route_stride, %expert_count, %route_ids, %expert_table) : [index, index, index, index](index, index, index, index, tensor<64x4xi32>, tensor<260xi32>) + kernel.launch @qwen3_moe_routed_gate_up_swiglu_q4k_q8[%token_count, %route_count, %route_stride, %expert_count, %output_size](%token_count, %route_count, %route_stride, %expert_count, %output_size, %q8_input, %route_ids, %gate_weight, %up_weight, %expected) : [index, index, index, index, index](index, index, index, index, index, tensor<64x576xi8>, tensor<64x4xi32>, tensor<4x32x2x144xi8>, tensor<4x32x2x144xi8>, tensor<64x2x32xf32>) + kernel.launch @qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped[%token_count, %route_count, %expert_count, %output_size](%token_count, %route_count, %expert_count, %output_size, %q8_input, %expert_table, %gate_weight, %up_weight, %actual) : [index, index, index, index](index, index, index, index, tensor<64x576xi8>, tensor<260xi32>, tensor<4x32x2x144xi8>, tensor<4x32x2x144xi8>, tensor<64x2x32xf32>) + check.expect.close actual(%actual) expected(%expected) atol(0.25) rtol(9.9999999999999995e-07) nan(same) : tensor<64x2x32xf32> + check.return +} + +check.case public @qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped_tail_case { + %token_count = check.literal value(37) : index + %input_size = check.literal value(512) : index + %route_count = check.literal value(3) : index + %route_stride = check.literal value(5) : index + %expert_count = check.literal value(4) : index + %output_size = check.literal value(33) : index + %input = check.generate.fill value(0.00390625) : tensor<37x512xf32> + %q8_input = check.generate.fill value(0) : tensor<37x576xi8> + // Three unique selected experts rotate through a five-entry physical row. + // This leaves the second route tile empty and the second channel tile with + // only one live output channel. + %route_ids = check.generate.iota offset(0) step(1) period(4) : tensor<37x5xi32> + %expert_table = check.generate.fill value(-1) : tensor<152xi32> + %gate_weight = check.generate.iota offset(-72) step(1) period(144) : tensor<4x33x2x144xi8> + %up_weight = check.generate.iota offset(-71) step(1) period(144) : tensor<4x33x2x144xi8> + %expected = check.generate.fill value(0.0) : tensor<37x3x33xf32> + %actual = check.generate.fill value(1.0) : tensor<37x3x33xf32> + kernel.launch @ggml_quantize_q8_1_x4_f32[%token_count, %input_size](%token_count, %input_size, %input, %q8_input) : [index, index](index, index, tensor<37x512xf32>, tensor<37x576xi8>) + kernel.launch @qwen3_moe_build_expert_table[%token_count, %route_count, %route_stride, %expert_count](%token_count, %route_count, %route_stride, %expert_count, %route_ids, %expert_table) : [index, index, index, index](index, index, index, index, tensor<37x5xi32>, tensor<152xi32>) + kernel.launch @qwen3_moe_routed_gate_up_swiglu_q4k_q8[%token_count, %route_count, %route_stride, %expert_count, %output_size](%token_count, %route_count, %route_stride, %expert_count, %output_size, %q8_input, %route_ids, %gate_weight, %up_weight, %expected) : [index, index, index, index, index](index, index, index, index, index, tensor<37x576xi8>, tensor<37x5xi32>, tensor<4x33x2x144xi8>, tensor<4x33x2x144xi8>, tensor<37x3x33xf32>) + kernel.launch @qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped[%token_count, %route_count, %expert_count, %output_size](%token_count, %route_count, %expert_count, %output_size, %q8_input, %expert_table, %gate_weight, %up_weight, %actual) : [index, index, index, index](index, index, index, index, tensor<37x576xi8>, tensor<152xi32>, tensor<4x33x2x144xi8>, tensor<4x33x2x144xi8>, tensor<37x3x33xf32>) + check.expect.close actual(%actual) expected(%expected) atol(0.25) rtol(9.9999999999999995e-07) nan(same) : tensor<37x3x33xf32> + check.return +} + +// Forces one route partition to process a second 128-row band. This models a +// maximally skewed router while retaining a noncompact physical route stride. +check.case public @qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped_partition_stride_case { + %token_count = check.literal value(129) : index + %input_size = check.literal value(512) : index + %route_count = check.literal value(1) : index + %route_stride = check.literal value(3) : index + %expert_count = check.literal value(2) : index + %output_size = check.literal value(1) : index + %input = check.generate.fill value(0.00390625) : tensor<129x512xf32> + %q8_input = check.generate.fill value(0) : tensor<129x576xi8> + %route_ids = check.generate.fill value(0) : tensor<129x3xi32> + %expert_table = check.generate.fill value(-1) : tensor<260xi32> + %gate_weight = check.generate.fill value(85) : tensor<2x1x2x144xi8> + %up_weight = check.generate.fill value(-86) : tensor<2x1x2x144xi8> + %expected = check.generate.fill value(0.0) : tensor<129x1x1xf32> + %actual = check.generate.fill value(1.0) : tensor<129x1x1xf32> + kernel.launch @ggml_quantize_q8_1_x4_f32[%token_count, %input_size](%token_count, %input_size, %input, %q8_input) : [index, index](index, index, tensor<129x512xf32>, tensor<129x576xi8>) + kernel.launch @qwen3_moe_build_expert_table[%token_count, %route_count, %route_stride, %expert_count](%token_count, %route_count, %route_stride, %expert_count, %route_ids, %expert_table) : [index, index, index, index](index, index, index, index, tensor<129x3xi32>, tensor<260xi32>) + kernel.launch @qwen3_moe_routed_gate_up_swiglu_q4k_q8[%token_count, %route_count, %route_stride, %expert_count, %output_size](%token_count, %route_count, %route_stride, %expert_count, %output_size, %q8_input, %route_ids, %gate_weight, %up_weight, %expected) : [index, index, index, index, index](index, index, index, index, index, tensor<129x576xi8>, tensor<129x3xi32>, tensor<2x1x2x144xi8>, tensor<2x1x2x144xi8>, tensor<129x1x1xf32>) + kernel.launch @qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped[%token_count, %route_count, %expert_count, %output_size](%token_count, %route_count, %expert_count, %output_size, %q8_input, %expert_table, %gate_weight, %up_weight, %actual) : [index, index, index, index](index, index, index, index, tensor<129x576xi8>, tensor<260xi32>, tensor<2x1x2x144xi8>, tensor<2x1x2x144xi8>, tensor<129x1x1xf32>) + check.expect.close actual(%actual) expected(%expected) atol(0.25) rtol(9.9999999999999995e-07) nan(same) : tensor<129x1x1xf32> + check.return +} + +check.case public @qwen3_moe_routed_gate_up_swiglu_q4k_q8_benchmark_case { + %token_count = check.param.choice values([1, 2, 4, 8, 16, 17, 32, 63, 128, 129, 512]) name("token_count") : index + %route_count = check.literal value(8) : index + %route_stride = check.literal value(128) : index + %expert_count = check.literal value(128) : index + %output_size = check.literal value(768) : index + %q8_input = check.generate.fill value(0) : tensor<[%token_count]x2304xi8> + %route_ids = check.generate.iota offset(0) step(1) period(127) : tensor<[%token_count]x128xi32> + %gate_weight = check.generate.fill value(0) : tensor<128x768x8x144xi8> + %up_weight = check.generate.fill value(0) : tensor<128x768x8x144xi8> + %output = check.generate.fill value(1.0) : tensor<[%token_count]x8x768xf32> + %expected = check.generate.fill value(0.0) : tensor<[%token_count]x8x768xf32> + kernel.launch @qwen3_moe_routed_gate_up_swiglu_q4k_q8[%token_count, %route_count, %route_stride, %expert_count, %output_size](%token_count, %route_count, %route_stride, %expert_count, %output_size, %q8_input, %route_ids, %gate_weight, %up_weight, %output) : [index, index, index, index, index](index, index, index, index, index, tensor<[%token_count]x2304xi8>, tensor<[%token_count]x128xi32>, tensor<128x768x8x144xi8>, tensor<128x768x8x144xi8>, tensor<[%token_count]x8x768xf32>) + check.expect.close actual(%output) expected(%expected) atol(0.0) rtol(0.0) nan(same) : tensor<[%token_count]x8x768xf32> + check.return +} + +check.case public @qwen3_moe_routed_gate_up_swiglu_q4k_q8_1_x4_next_q8_benchmark_case { + %token_count = check.literal value(1) : index + %route_count = check.literal value(8) : index + %route_stride = check.literal value(8) : index + %expert_count = check.literal value(128) : index + %output_size = check.literal value(768) : index + %q8_input = check.generate.fill value(0) : tensor<2304xi8> + %route_ids = check.generate.iota offset(0) step(1) period(128) : tensor<8xi32> + %gate_weight = check.generate.fill value(0) : tensor<128x768x8x144xi8> + %up_weight = check.generate.fill value(0) : tensor<128x768x8x144xi8> + %output = check.generate.fill value(1.0) : tensor<8x768xf32> + %completion_counters = check.generate.fill value(0) : tensor<48xi32> + %next_q8_output = check.generate.fill value(1) : tensor<8x864xi8> + %expected_output = check.generate.fill value(0.0) : tensor<8x768xf32> + %expected_q8 = check.generate.fill value(0) : tensor<8x864xi8> + %expected_counters = check.generate.fill value(0) : tensor<48xi32> + kernel.launch @qwen3_moe_routed_gate_up_swiglu_q4k_q8_1_x4_next_q8[%token_count, %route_count, %route_stride, %expert_count, %output_size](%token_count, %route_count, %route_stride, %expert_count, %output_size, %q8_input, %route_ids, %gate_weight, %up_weight, %output, %completion_counters, %next_q8_output) : [index, index, index, index, index](index, index, index, index, index, tensor<2304xi8>, tensor<8xi32>, tensor<128x768x8x144xi8>, tensor<128x768x8x144xi8>, tensor<8x768xf32>, tensor<48xi32>, tensor<8x864xi8>) + check.expect.close actual(%output) expected(%expected_output) atol(0.0) rtol(0.0) nan(same) : tensor<8x768xf32> + check.expect.equal actual(%next_q8_output) expected(%expected_q8) : tensor<8x864xi8> + check.expect.equal actual(%completion_counters) expected(%expected_counters) : tensor<48xi32> + check.return +} + +check.case public @qwen3_moe_routed_gate_up_swiglu_q4k_q8_1_x4_next_q8_composed_benchmark_case { + %token_count = check.literal value(1) : index + %route_count = check.literal value(8) : index + %route_stride = check.literal value(8) : index + %expert_count = check.literal value(128) : index + %output_size = check.literal value(768) : index + %routed_row_count = check.literal value(8) : index + %q8_input = check.generate.fill value(0) : tensor<2304xi8> + %route_ids = check.generate.iota offset(0) step(1) period(128) : tensor<8xi32> + %gate_weight = check.generate.fill value(0) : tensor<128x768x8x144xi8> + %up_weight = check.generate.fill value(0) : tensor<128x768x8x144xi8> + %output = check.generate.fill value(1.0) : tensor<8x768xf32> + %next_q8_output = check.generate.fill value(1) : tensor<8x864xi8> + %expected_output = check.generate.fill value(0.0) : tensor<8x768xf32> + %expected_q8 = check.generate.fill value(0) : tensor<8x864xi8> + kernel.launch @qwen3_moe_routed_gate_up_swiglu_q4k_q8[%token_count, %route_count, %route_stride, %expert_count, %output_size](%token_count, %route_count, %route_stride, %expert_count, %output_size, %q8_input, %route_ids, %gate_weight, %up_weight, %output) : [index, index, index, index, index](index, index, index, index, index, tensor<2304xi8>, tensor<8xi32>, tensor<128x768x8x144xi8>, tensor<128x768x8x144xi8>, tensor<8x768xf32>) + kernel.launch @ggml_quantize_q8_1_x4_f32[%routed_row_count, %output_size](%routed_row_count, %output_size, %output, %next_q8_output) : [index, index](index, index, tensor<8x768xf32>, tensor<8x864xi8>) + check.expect.close actual(%output) expected(%expected_output) atol(0.0) rtol(0.0) nan(same) : tensor<8x768xf32> + check.expect.equal actual(%next_q8_output) expected(%expected_q8) : tensor<8x864xi8> + check.return +} + +check.case public @qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped_benchmark_case { + %token_count = check.param.choice values([1, 2, 4, 8, 16, 17, 32, 63, 128, 129, 512]) name("token_count") : index + %route_count = check.literal value(8) : index + %route_stride = check.literal value(128) : index + %expert_count = check.literal value(128) : index + %output_size = check.literal value(768) : index + %q8_input = check.generate.fill value(0) : tensor<[%token_count]x2304xi8> + // A period of 127 rotates each physical 128-entry row by one expert while + // keeping the first eight logical routes distinct within every token. + %route_ids = check.generate.iota offset(0) step(1) period(127) : tensor<[%token_count]x128xi32> + // One packed transient buffer holds 128 counts followed by room for 512 + // assignments per expert. The production allocation uses the exact token + // count; this fixed test capacity permits one parameterized benchmark case. + %expert_table = check.generate.fill value(-1) : tensor<65664xi32> + %gate_weight = check.generate.fill value(0) : tensor<128x768x8x144xi8> + %up_weight = check.generate.fill value(0) : tensor<128x768x8x144xi8> + %output = check.generate.fill value(1.0) : tensor<[%token_count]x8x768xf32> + %expected = check.generate.fill value(0.0) : tensor<[%token_count]x8x768xf32> + kernel.launch @qwen3_moe_build_expert_table[%token_count, %route_count, %route_stride, %expert_count](%token_count, %route_count, %route_stride, %expert_count, %route_ids, %expert_table) : [index, index, index, index](index, index, index, index, tensor<[%token_count]x128xi32>, tensor<65664xi32>) + kernel.launch @qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped[%token_count, %route_count, %expert_count, %output_size](%token_count, %route_count, %expert_count, %output_size, %q8_input, %expert_table, %gate_weight, %up_weight, %output) : [index, index, index, index](index, index, index, index, tensor<[%token_count]x2304xi8>, tensor<65664xi32>, tensor<128x768x8x144xi8>, tensor<128x768x8x144xi8>, tensor<[%token_count]x8x768xf32>) + check.expect.close actual(%output) expected(%expected) atol(0.0) rtol(0.0) nan(same) : tensor<[%token_count]x8x768xf32> + check.return +} + +// Maximally diverse decode-batch control. Compact route storage makes the +// flattened 0..127 iota assign every M=16 route to a different expert. +check.case public @qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped_diverse_benchmark_case { + %token_count = check.param.choice values([1, 2, 4, 8, 16]) name("token_count") : index + %route_count = check.literal value(8) : index + %route_stride = check.literal value(8) : index + %expert_count = check.literal value(128) : index + %output_size = check.literal value(768) : index + %q8_input = check.generate.fill value(0) : tensor<[%token_count]x2304xi8> + %route_ids = check.generate.iota offset(0) step(1) period(128) : tensor<[%token_count]x8xi32> + %expert_table = check.generate.fill value(-1) : tensor<65664xi32> + %gate_weight = check.generate.fill value(0) : tensor<128x768x8x144xi8> + %up_weight = check.generate.fill value(0) : tensor<128x768x8x144xi8> + %output = check.generate.fill value(1.0) : tensor<[%token_count]x8x768xf32> + %expected = check.generate.fill value(0.0) : tensor<[%token_count]x8x768xf32> + kernel.launch @qwen3_moe_build_expert_table[%token_count, %route_count, %route_stride, %expert_count](%token_count, %route_count, %route_stride, %expert_count, %route_ids, %expert_table) : [index, index, index, index](index, index, index, index, tensor<[%token_count]x8xi32>, tensor<65664xi32>) + kernel.launch @qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped[%token_count, %route_count, %expert_count, %output_size](%token_count, %route_count, %expert_count, %output_size, %q8_input, %expert_table, %gate_weight, %up_weight, %output) : [index, index, index, index](index, index, index, index, tensor<[%token_count]x2304xi8>, tensor<65664xi32>, tensor<128x768x8x144xi8>, tensor<128x768x8x144xi8>, tensor<[%token_count]x8x768xf32>) + check.expect.close actual(%output) expected(%expected) atol(0.0) rtol(0.0) nan(same) : tensor<[%token_count]x8x768xf32> + check.return +} + +check.case public @qwen3_moe_routed_gate_up_swiglu_q4k_q8_pipeline_benchmark_case { + %token_count = check.param.choice values([1, 2, 4, 8, 16, 32, 128, 512]) name("token_count") : index + %input_size = check.literal value(2048) : index + %route_count = check.literal value(8) : index + %route_stride = check.literal value(128) : index + %expert_count = check.literal value(128) : index + %output_size = check.literal value(768) : index + %input = check.generate.fill value(0.0) : tensor<[%token_count]x2048xf32> + %q8_input = check.generate.fill value(1) : tensor<[%token_count]x2304xi8> + %route_ids = check.generate.iota offset(0) step(1) period(127) : tensor<[%token_count]x128xi32> + %gate_weight = check.generate.fill value(0) : tensor<128x768x8x144xi8> + %up_weight = check.generate.fill value(0) : tensor<128x768x8x144xi8> + %output = check.generate.fill value(1.0) : tensor<[%token_count]x8x768xf32> + %expected = check.generate.fill value(0.0) : tensor<[%token_count]x8x768xf32> + kernel.launch @ggml_quantize_q8_1_x4_f32[%token_count, %input_size](%token_count, %input_size, %input, %q8_input) : [index, index](index, index, tensor<[%token_count]x2048xf32>, tensor<[%token_count]x2304xi8>) + kernel.launch @qwen3_moe_routed_gate_up_swiglu_q4k_q8[%token_count, %route_count, %route_stride, %expert_count, %output_size](%token_count, %route_count, %route_stride, %expert_count, %output_size, %q8_input, %route_ids, %gate_weight, %up_weight, %output) : [index, index, index, index, index](index, index, index, index, index, tensor<[%token_count]x2304xi8>, tensor<[%token_count]x128xi32>, tensor<128x768x8x144xi8>, tensor<128x768x8x144xi8>, tensor<[%token_count]x8x768xf32>) + check.expect.close actual(%output) expected(%expected) atol(0.0) rtol(0.0) nan(same) : tensor<[%token_count]x8x768xf32> + check.return +} + +check.case public @qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped_pipeline_benchmark_case { + %token_count = check.param.choice values([1, 2, 4, 8, 16, 17, 32, 63, 128, 129, 512]) name("token_count") : index + %input_size = check.literal value(2048) : index + %route_count = check.literal value(8) : index + %route_stride = check.literal value(128) : index + %expert_count = check.literal value(128) : index + %output_size = check.literal value(768) : index + %input = check.generate.fill value(0.0) : tensor<[%token_count]x2048xf32> + %q8_input = check.generate.fill value(1) : tensor<[%token_count]x2304xi8> + %route_ids = check.generate.iota offset(0) step(1) period(127) : tensor<[%token_count]x128xi32> + %expert_table = check.generate.fill value(-1) : tensor<65664xi32> + %gate_weight = check.generate.fill value(0) : tensor<128x768x8x144xi8> + %up_weight = check.generate.fill value(0) : tensor<128x768x8x144xi8> + %output = check.generate.fill value(1.0) : tensor<[%token_count]x8x768xf32> + %expected = check.generate.fill value(0.0) : tensor<[%token_count]x8x768xf32> + kernel.launch @ggml_quantize_q8_1_x4_f32[%token_count, %input_size](%token_count, %input_size, %input, %q8_input) : [index, index](index, index, tensor<[%token_count]x2048xf32>, tensor<[%token_count]x2304xi8>) + kernel.launch @qwen3_moe_build_expert_table[%token_count, %route_count, %route_stride, %expert_count](%token_count, %route_count, %route_stride, %expert_count, %route_ids, %expert_table) : [index, index, index, index](index, index, index, index, tensor<[%token_count]x128xi32>, tensor<65664xi32>) + kernel.launch @qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped[%token_count, %route_count, %expert_count, %output_size](%token_count, %route_count, %expert_count, %output_size, %q8_input, %expert_table, %gate_weight, %up_weight, %output) : [index, index, index, index](index, index, index, index, tensor<[%token_count]x2304xi8>, tensor<65664xi32>, tensor<128x768x8x144xi8>, tensor<128x768x8x144xi8>, tensor<[%token_count]x8x768xf32>) + check.expect.close actual(%output) expected(%expected) atol(0.0) rtol(0.0) nan(same) : tensor<[%token_count]x8x768xf32> + check.return +} + +check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_q8_nonzero_case> @qwen3_moe_routed_gate_up_swiglu_q4k_q8_small + +check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_q8_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_q8_decode {token_count = 1} + +check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_q8_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_q8_small_batch_2 {token_count = 2} + +check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_q8_1_x4_next_q8_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_q8_1_x4_next_q8_decode + +check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_q8_1_x4_next_q8_composed_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_q8_1_x4_next_q8_composed_decode + +check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_q8_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_q8_small_batch_4 {token_count = 4} + +check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_q8_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_q8_small_batch_8 {token_count = 8} + +check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_q8_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_q8_small_batch_16 {token_count = 16} + +check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_q8_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_q8_prefill_17 {token_count = 17} + +check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_q8_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_q8_prefill_32 {token_count = 32} + +check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_q8_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_q8_prefill_63 {token_count = 63} + +check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_q8_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_q8_prefill_128 {token_count = 128} + +check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_q8_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_q8_prefill_129 {token_count = 129} + +check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_q8_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_q8_prefill_512 {token_count = 512} + +check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped_prefill_17 {token_count = 17} + +check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped_decode {token_count = 1} + +check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped_small_batch_2 {token_count = 2} + +check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped_small_batch_4 {token_count = 4} + +check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped_small_batch_8 {token_count = 8} + +check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped_small_batch_16 {token_count = 16} + +check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped_diverse_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped_diverse_decode {token_count = 1} + +check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped_diverse_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped_diverse_small_batch_2 {token_count = 2} + +check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped_diverse_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped_diverse_small_batch_4 {token_count = 4} + +check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped_diverse_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped_diverse_small_batch_8 {token_count = 8} + +check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped_diverse_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped_diverse_small_batch_16 {token_count = 16} + +check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped_prefill_32 {token_count = 32} + +check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped_prefill_63 {token_count = 63} + +check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped_prefill_128 {token_count = 128} + +check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped_prefill_129 {token_count = 129} + +check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped_prefill_512 {token_count = 512} + +check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_q8_pipeline_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_q8_pipeline_decode {token_count = 1} + +check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_q8_pipeline_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_q8_pipeline_small_batch_8 {token_count = 8} + +check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_q8_pipeline_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_q8_pipeline_prefill_32 {token_count = 32} + +check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_q8_pipeline_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_q8_pipeline_prefill_128 {token_count = 128} + +check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_q8_pipeline_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_q8_pipeline_prefill_512 {token_count = 512} + +check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped_pipeline_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped_pipeline_prefill_17 {token_count = 17} + +check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped_pipeline_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped_pipeline_prefill_32 {token_count = 32} + +check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped_pipeline_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped_pipeline_prefill_63 {token_count = 63} + +check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped_pipeline_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped_pipeline_prefill_128 {token_count = 128} + +check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped_pipeline_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped_pipeline_prefill_129 {token_count = 129} + +check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped_pipeline_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped_pipeline_prefill_512 {token_count = 512} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/routed_linear_q4k_f16_wmma.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/routed_linear_q4k_f16_wmma.loom new file mode 100644 index 000000000000..d24109910e3d --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/routed_linear_q4k_f16_wmma.loom @@ -0,0 +1,1004 @@ +// Gfx11 routed-expert projection matching llama.cpp Vulkan's matrix path. +// +// Each two-wave workgroup computes 64 output channels for 32 routed rows of one +// expert. The waves share padded FP16 operand stages, each owns 32 output +// channels and four WMMA accumulators, and scatters one transposed accumulator +// fragment at a time through a compact route map. Gate and up invoke this same +// projection independently before the separate SwiGLU epilogue. +// +// The raw weights remain [expert][output channel][K / 256][144 bytes]. +// No persistent repacking or expanded-weight allocation is required. +func.def inline @qwen3_moe_q4k_scale_from_header(%scale0: i32, %scale1: i32, %scale2: i32, %q4_group: index) -> (i32, i32) { + %c4 = index.constant 4 : index + %c8 = index.constant 8 : index + %c2_i32 = scalar.constant 2 : i32 + %c4_i32 = scalar.constant 4 : i32 + %c15_i32 = scalar.constant 15 : i32 + %c48_i32 = scalar.constant 48 : i32 + %bounded_group = index.assume %q4_group [range(%q4_group, 0, 7)] : index + %is_low_group = index.cmp ult, %bounded_group, %c4 : index + %scale_lane = index.rem %bounded_group, %c4 : index + %scale_shift_index = index.mul %scale_lane, %c8 : index + %scale_shift = index.cast %scale_shift_index : index to i32 + %high_shift = scalar.addi %scale_shift, %c2_i32 : i32 + %minimum_shift = scalar.addi %scale_shift, %c4_i32 : i32 + %selected_scale_source = scf.select %is_low_group, %scale0, %scale2 : i32 + %selected_minimum_source = scf.select %is_low_group, %scale1, %scale2 : i32 + %selected_scale_high_shift = scf.select %is_low_group, %scale_shift, %high_shift : i32 + %selected_minimum_low_shift = scf.select %is_low_group, %scale_shift, %minimum_shift : i32 + %scale_low0 = scalar.shrui %selected_scale_source, %scale_shift : i32 + %scale_low = scalar.andi %scale_low0, %c15_i32 : i32 + %scale_high0 = scalar.shrui %scale0, %selected_scale_high_shift : i32 + %scale_high = scalar.andi %scale_high0, %c48_i32 : i32 + %scale = scalar.ori %scale_low, %scale_high : i32 + %minimum_low0 = scalar.shrui %selected_minimum_source, %selected_minimum_low_shift : i32 + %minimum_low = scalar.andi %minimum_low0, %c15_i32 : i32 + %minimum_high0 = scalar.shrui %scale1, %selected_scale_high_shift : i32 + %minimum_high = scalar.andi %minimum_high0, %c48_i32 : i32 + %minimum = scalar.ori %minimum_low, %minimum_high : i32 + func.return %scale, %minimum : i32, i32 +} + +func.def inline @qwen3_moe_unpack_expert_partition_descriptor(%descriptor: i32) -> (index, index, index) { + %c1_i32 = scalar.constant 1 : i32 + %c5_i32 = scalar.constant 5 : i32 + %c7_i32 = scalar.constant 7 : i32 + %c13_i32 = scalar.constant 13 : i32 + %c31_i32 = scalar.constant 31 : i32 + %c63_i32 = scalar.constant 63 : i32 + %c127_i32 = scalar.constant 127 : i32 + %expert_i32 = scalar.andi %descriptor, %c127_i32 : i32 + %partition_shifted_i32 = scalar.shrui %descriptor, %c7_i32 : i32 + %partition_i32 = scalar.andi %partition_shifted_i32, %c63_i32 : i32 + %route_tile_base_i32 = scalar.shli %partition_i32, %c5_i32 : i32 + %row_count_shifted_i32 = scalar.shrui %descriptor, %c13_i32 : i32 + %row_count_minus_one_i32 = scalar.andi %row_count_shifted_i32, %c31_i32 : i32 + %partition_row_count_i32 = scalar.addi %row_count_minus_one_i32, %c1_i32 : i32 + %expert0 = index.cast %expert_i32 : i32 to index + %expert = index.assume %expert0 [range(%expert0, 0, 127)] : index + %route_tile_base0 = index.cast %route_tile_base_i32 : i32 to index + %route_tile_base = index.assume %route_tile_base0 [range(%route_tile_base0, 0, 2016)] : index + %partition_row_count0 = index.cast %partition_row_count_i32 : i32 to index + %partition_row_count = index.assume %partition_row_count0 [range(%partition_row_count0, 1, 32)] : index + func.return %expert, %route_tile_base, %partition_row_count : index, index, index +} + +amdgpu.target @qwen3_moe_gfx11_wave64 {subgroup_size = 64} + +amdgpu.target @qwen3_moe_gfx11_wave32 {subgroup_size = 32} + +config.decl @qwen3_moe.routed_gate_up.input_size : %value: index where [range(%value, 512, 32768), mul(%value, 512)] + +// Top-k is a model hyperparameter and is specialized with the kernel. Keeping +// it out of the dynamic workload lets address arithmetic and assignment decode +// fold to the exact model contract. +config.decl @qwen3_moe.routed_gate_up.route_count : %value: index where [range(%value, 1, 8)] + +config.decl @qwen3_moe.routed_gate_up.expert_count : %value: index where [range(%value, 1, 512)] + +config.decl @qwen3_moe.routed_gate_up.output_size : %value: index where [range(%value, 1, 4096)] + +kernel.decl @ggml_quantize_q8_1_x4_f32(%token_count: index, %input_size: index) launch(%token_count: index, %input_size: index, %input: buffer, %output: buffer) + +kernel.decl @qwen3_moe_build_expert_table(%token_count: index, %route_count: index, %route_stride: index, %expert_count: index) launch(%token_count: index, %route_count: index, %route_stride: index, %expert_count: index, %route_ids: buffer, %expert_table: buffer) + +kernel.decl @qwen3_moe_build_expert_partition_table(%token_count: index, %route_count: index, %expert_count: index) launch(%token_count: index, %route_count: index, %expert_count: index, %expert_table: buffer, %partition_table: buffer) + +func.decl @qwen3_moe_unpack_expert_partition_descriptor(%descriptor: i32) -> (index, index, index) + +kernel.decl @qwen3_moe_routed_gate_up_swiglu_q4k_q8(%token_count: index, %route_count: index, %route_stride: index, %expert_count: index, %output_size: index) launch(%token_count: index, %route_count: index, %route_stride: index, %expert_count: index, %output_size: index, %q8_input: buffer, %route_ids: buffer, %gate_weight: buffer, %up_weight: buffer, %output: buffer) + +func.decl @qwen3_moe_q4k_scale_from_header(%scale0: i32, %scale1: i32, %scale2: i32, %q4_group: index) -> (i32, i32) + +// Acquires the packed code word shared by one adjacent Q4_K group pair. +func.def inline @qwen3_moe_q4k_wmma_load_code(%weight: buffer, %row_byte_base: offset, %q4_block: index, %q4_group_pair: index, %packet: index) -> (vector<1xi32>) { + %c8 = index.constant 8 : index + %block_bytes = index.constant 144 : offset + %code_offset = index.constant 16 : offset + %bounded_group_pair = index.assume %q4_group_pair [range(%q4_group_pair, 0, 3)] : index + %bounded_packet = index.assume %packet [range(%packet, 0, 7)] : index + %block_byte_add = index.scale %q4_block, %block_bytes : index, offset -> offset + %block_byte_base = index.add %row_byte_base, %block_byte_add : offset + %code_byte_base = index.add %block_byte_base, %code_offset : offset + %code_view = buffer.view %weight[%code_byte_base] : buffer -> view<32xi32> + %q_page = index.mul %bounded_group_pair, %c8 : index + %q_word_index0 = index.add %q_page, %bounded_packet : index + %q_word_index = index.assume %q_word_index0 [range(%q_word_index0, 0, 31)] : index + %q_word = vector.load %code_view[%q_word_index] : view<32xi32> -> vector<1xi32> + func.return %q_word : vector<1xi32> +} + +// Decodes the four adjacent Q4_K values owned by one load packet from an +// already-loaded block header and packed code word. Matrix schedules choose +// the lifetime of both immutable packets. +func.def inline @qwen3_moe_q4k_wmma_vector4_from_header_code(%q4_group: index, %header_words: vector<4xi32>, %q_word: vector<1xi32>) -> (vector<4xf16>) { + %c2 = index.constant 2 : index + %c4 = index.constant 4 : index + %q4_mask = vector.constant 252645135 : vector<1xi32> + %bounded_group = index.assume %q4_group [range(%q4_group, 0, 7)] : index + %header_halves = vector.bitcast %header_words : vector<4xi32> to vector<8xf16> + %d_f16 = vector.extract %header_halves[0] : vector<8xf16> -> f16 + %dmin_f16 = vector.extract %header_halves[1] : vector<8xf16> -> f16 + %scale0 = vector.extract %header_words[1] : vector<4xi32> -> i32 + %scale1 = vector.extract %header_words[2] : vector<4xi32> -> i32 + %scale2 = vector.extract %header_words[3] : vector<4xi32> -> i32 + %d = scalar.extf %d_f16 : f16 to f32 + %dmin = scalar.extf %dmin_f16 : f16 to f32 + %scale, %minimum = func.call @qwen3_moe_q4k_scale_from_header(%scale0, %scale1, %scale2, %bounded_group) : (i32, i32, i32, index) -> (i32, i32) + %scale_f32 = scalar.uitofp %scale : i32 to f32 + %minimum_f32 = scalar.uitofp %minimum : i32 to f32 + %d_scale = scalar.mulf %d, %scale_f32 : f32 + %minimum_scale = scalar.mulf %dmin, %minimum_f32 : f32 + %q_half = index.rem %bounded_group, %c2 : index + %q_shift_index = index.mul %q_half, %c4 : index + %q_shift_i32 = index.cast %q_shift_index : index to i32 + %q_shift = vector.splat %q_shift_i32 : vector<1xi32> + %shifted_q = vector.shrui %q_word, %q_shift : vector<1xi32> + %masked_q = vector.andi %shifted_q, %q4_mask : vector<1xi32> + %q_i8 = vector.bitcast %masked_q : vector<1xi32> to vector<4xi8> + %q_f32 = vector.uitofp %q_i8 : vector<4xi8> to vector<4xf32> + // Form adjacent FP16 lanes from fused FP32 affine expressions. AMDGPU maps + // this natural shape to packed mixlo/mixhi instructions where available. + %negative_minimum_scale = scalar.negf %minimum_scale : f32 + %q0 = vector.extract %q_f32[0] : vector<4xf32> -> f32 + %q1 = vector.extract %q_f32[1] : vector<4xf32> -> f32 + %q2 = vector.extract %q_f32[2] : vector<4xf32> -> f32 + %q3 = vector.extract %q_f32[3] : vector<4xf32> -> f32 + %value0 = scalar.fmaf %q0, %d_scale, %negative_minimum_scale : f32 + %value1 = scalar.fmaf %q1, %d_scale, %negative_minimum_scale : f32 + %value2 = scalar.fmaf %q2, %d_scale, %negative_minimum_scale : f32 + %value3 = scalar.fmaf %q3, %d_scale, %negative_minimum_scale : f32 + %half0 = scalar.fptrunc %value0 : f32 to f16 + %half1 = scalar.fptrunc %value1 : f32 to f16 + %half2 = scalar.fptrunc %value2 : f32 to f16 + %half3 = scalar.fptrunc %value3 : f32 to f16 + %result = vector.from_elements %half0, %half1, %half2, %half3 : vector<4xf16> + func.return %result : vector<4xf16> +} + +// Decodes one group when its caller has retained only the Q4_K block header. +func.def inline @qwen3_moe_q4k_wmma_vector4_from_header(%weight: buffer, %row_byte_base: offset, %q4_block: index, %q4_group: index, %packet: index, %header_words: vector<4xi32>) -> (vector<4xf16>) { + %c2 = index.constant 2 : index + %bounded_group = index.assume %q4_group [range(%q4_group, 0, 7)] : index + %q4_group_pair = index.div %bounded_group, %c2 : index + %q_word = func.call @qwen3_moe_q4k_wmma_load_code(%weight, %row_byte_base, %q4_block, %q4_group_pair, %packet) : (buffer, offset, index, index, index) -> (vector<1xi32>) + %values = func.call @qwen3_moe_q4k_wmma_vector4_from_header_code(%bounded_group, %header_words, %q_word) : (index, vector<4xi32>, vector<1xi32>) -> (vector<4xf16>) + func.return %values : vector<4xf16> +} + +// Acquires one naturally aligned Q4_K block header as a single 16-byte packet. +func.def inline @qwen3_moe_q4k_wmma_load_header(%weight: buffer, %row_byte_base: offset, %q4_block: index) -> (vector<4xi32>) { + %c0 = index.constant 0 : index + %block_bytes = index.constant 144 : offset + %block_byte_add = index.scale %q4_block, %block_bytes : index, offset -> offset + %block_byte_base = index.add %row_byte_base, %block_byte_add : offset + %header_view = buffer.view %weight[%block_byte_base] : buffer -> view<4xi32> + %header_words = vector.load %header_view[%c0] : view<4xi32> -> vector<4xi32> + func.return %header_words : vector<4xi32> +} + +// Acquires one block header before decoding the selected four-value group +// packet. Matrix schedules that span several groups call the two operations +// separately so the header lifetime matches their complete block loop. +func.def inline @qwen3_moe_q4k_wmma_vector4(%weight: buffer, %row_byte_base: offset, %q4_block: index, %q4_group: index, %packet: index) -> (vector<4xf16>) { + %header_words = func.call @qwen3_moe_q4k_wmma_load_header(%weight, %row_byte_base, %q4_block) : (buffer, offset, index) -> (vector<4xi32>) + %values = func.call @qwen3_moe_q4k_wmma_vector4_from_header(%weight, %row_byte_base, %q4_block, %q4_group, %packet, %header_words) : (buffer, offset, index, index, index, vector<4xi32>) -> (vector<4xf16>) + func.return %values : vector<4xf16> +} + +kernel.def target(@qwen3_moe_gfx11_wave64) @qwen3_moe_routed_linear_q4k_f16_wmma(%token_count: index) { + %expert_count = config.get @qwen3_moe.routed_gate_up.expert_count : index + %output_size = config.get @qwen3_moe.routed_gate_up.output_size : index + %c1 = index.constant 1 : index + %c63 = index.constant 63 : index + %c64 = index.constant 64 : index + %c128 = index.constant 128 : index + %padded_output_size = index.add %output_size, %c63 : index + %output_tiles = index.div %padded_output_size, %c64 : index + %padded_token_count = index.add %token_count, %c63 : index + %route_tiles = index.div %padded_token_count, %c64 : index + kernel.launch.config workgroups(%output_tiles, %route_tiles, %expert_count) workgroup_size(%c128, %c1, %c1) : index +} launch(%token_count: index, %input: buffer, %expert_table: buffer, %weight: buffer, %output: buffer) where [range(%token_count, 1, 2048)] { + %input_size = config.get @qwen3_moe.routed_gate_up.input_size : index + %route_count = config.get @qwen3_moe.routed_gate_up.route_count : index + %expert_count = config.get @qwen3_moe.routed_gate_up.expert_count : index + %output_size = config.get @qwen3_moe.routed_gate_up.output_size : index + %bounded_route_count = index.assume %route_count [range(%route_count, 1, 8)] : index + %bounded_expert_count = index.assume %expert_count [range(%expert_count, 1, 128)] : index + %bounded_output_size = index.assume %output_size [range(%output_size, 1, 4096)] : index + %channel_tile = kernel.workgroup.id : index + %route_tile = kernel.workgroup.id : index + %expert = kernel.workgroup.id : index + %workitem = kernel.workitem.id : index + %subgroup0 = kernel.subgroup.id : index + %subgroup = index.assume %subgroup0 [range(%subgroup0, 0, 1)] : index + %lane = kernel.subgroup.lane.id : index + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c4 = index.constant 4 : index + %c8 = index.constant 8 : index + %c16 = index.constant 16 : index + %c32 = index.constant 32 : index + %c40 = index.constant 40 : index + %c63 = index.constant 63 : index + %c64 = index.constant 64 : index + %c256 = index.constant 256 : index + %c0_offset = index.constant 0 : offset + %c4_bytes = index.constant 4 : offset + %q4_block_bytes = index.constant 144 : offset + %weight_stage_bytes = index.constant 5120 : offset + %activation_stage_bytes = index.constant 2560 : offset + %route_stage_bytes = index.constant 128 : offset + %wave_result_stage_bytes = index.constant 512 : offset + %result_stage_bytes = index.constant 1024 : offset + %c0_i32 = scalar.constant 0 : i32 + %cn1_i32 = scalar.constant -1 : i32 + %c0_f16x4 = vector.constant 0.0 : vector<4xf16> + %zero_accumulator = vector.constant 0.0 : vector<8xf16> + %m = index.constant 16 : index + %n = index.constant 16 : index + %k = index.constant 16 : index + %assignment_count = index.mul %token_count, %bounded_route_count : index + %assignment_table_byte_base = index.scale %bounded_expert_count, %c4_bytes : index, offset -> offset + %q4_block_count = index.div %input_size, %c256 : index + %weight_row_bytes = index.scale %q4_block_count, %q4_block_bytes : index, offset -> offset + %weight_expert_bytes = index.scale %bounded_output_size, %weight_row_bytes : index, offset -> offset + %output_row_count = index.mul %token_count, %bounded_route_count : index + %input_noalias, %expert_table_noalias, %weight_noalias, %output_noalias = buffer.assume.noalias %input, %expert_table, %weight, %output : buffer, buffer, buffer, buffer + %input_view = buffer.view %input_noalias[%c0_offset] : buffer -> view<[%token_count]x[%input_size]xf32> + %count_view = buffer.view %expert_table_noalias[%c0_offset] : buffer -> view<[%bounded_expert_count]xi32> + %assignment_view = buffer.view %expert_table_noalias[%assignment_table_byte_base] : buffer -> view<[%bounded_expert_count]x[%token_count]xi32> + %output_view = buffer.view %output_noalias[%c0_offset] : buffer -> view<[%output_row_count]x[%bounded_output_size]xf32> + %weight_stage = buffer.alloca align(16) %weight_stage_bytes : buffer + %activation_stage = buffer.alloca align(16) %activation_stage_bytes : buffer + %route_stage = buffer.alloca align(16) %route_stage_bytes : buffer + %result_stage = buffer.alloca align(16) %result_stage_bytes : buffer + %weight_stage_view = buffer.view %weight_stage[%c0_offset] : buffer -> view<64x40xf16> + %activation_stage_physical_view = buffer.view %activation_stage[%c0_offset] : buffer -> view<32x40xf16> + %activation_fragment_layout = encoding.layout.strided [1, %c40] : encoding + %activation_fragment_view = buffer.view %activation_stage[%c0_offset] : buffer -> view<32x32xf16, %activation_fragment_layout> + %route_stage_view = buffer.view %route_stage[%c0_offset] : buffer -> view<32xi32> + %wave_result_stage_offset = index.scale %subgroup, %wave_result_stage_bytes : index, offset -> offset + %result_fragment_layout = encoding.layout.strided [1, %c16] : encoding + %result_fragment_view = buffer.view %result_stage[%wave_result_stage_offset] : buffer -> view<16x16xf16, %result_fragment_layout> + %result_physical_view = buffer.view %result_stage[%wave_result_stage_offset] : buffer -> view<16x16xf16> + %channel_tile_base = index.mul %channel_tile, %c64 : index + %initial_route_tile_base = index.mul %route_tile, %c32 : index + %padded_token_count = index.add %token_count, %c63 : index + %route_partition_count = index.div %padded_token_count, %c64 : index + %route_partition_step = index.mul %route_partition_count, %c32 : index + %bounded_expert, %table_expert_count = index.assume %expert, %bounded_expert_count [lt(%expert, %bounded_expert_count)] : index, index + %is_workitem_zero = index.cmp eq, %workitem, %c0 : index + %lane_expert_route_count = scf.if %is_workitem_zero -> (i32) { + %loaded = view.load %count_view[%bounded_expert] : view<[%bounded_expert_count]xi32> -> i32 + scf.yield %loaded : i32 + } else { + scf.yield %c0_i32 : i32 + } + %expert_route_count_reduced = kernel.workgroup.reduce %lane_expert_route_count : i32 + %expert_route_count_i32 = kernel.subgroup.broadcast.first %expert_route_count_reduced : i32 + %expert_route_count0 = index.cast %expert_route_count_i32 : i32 to index + %expert_route_count = index.assume %expert_route_count0 [range(%expert_route_count0, 0, 2048)] : index + // Route partitions are distributed across the launch grid and continue in + // uniform strides for concentrated routing. Balanced Qwen prefill gives each + // expert 32, 64, or 128 rows at 512, 1024, or 2048 tokens. Concentrated + // experts can consume the full token count without changing the launch + // geometry. + scf.for %route_tile_base = [%initial_route_tile_base to %expert_route_count step %route_partition_step] { + // The first wave snapshots the compact route map once. Every K tile then + // reuses these 32 entries while the full workgroup cooperatively fills the + // operand stages. + %loads_route = index.cmp ult, %workitem, %c32 : index + scf.if %loads_route { + %local_route = index.assume %workitem [range(%workitem, 0, 31)] : index + %assignment_ordinal = index.add %route_tile_base, %local_route : index + %valid_row = index.cmp ult, %assignment_ordinal, %expert_route_count : index + %assignment_i32 = scf.if %valid_row -> (i32) { + %bounded_assignment_ordinal, %table_token_count = index.assume %assignment_ordinal, %token_count [lt(%assignment_ordinal, %token_count)] : index, index + %loaded = view.load %assignment_view[%bounded_expert, %bounded_assignment_ordinal] : view<[%bounded_expert_count]x[%token_count]xi32> -> i32 + scf.yield %loaded : i32 + } else { + scf.yield %cn1_i32 : i32 + } + view.store %assignment_i32, %route_stage_view[%local_route] : i32, view<32xi32> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + %load_packet = index.rem %workitem, %c8 : index + %load_k = index.mul %load_packet, %c4 : index + %load_row0 = index.div %workitem, %c8 : index + %load_row = index.assume %load_row0 [range(%load_row0, 0, 15)] : index + %expert_byte_base = index.scale %bounded_expert, %weight_expert_bytes : index, offset -> offset + %subgroup_channel_add = index.mul %subgroup, %c32 : index + %subgroup_channel1 = index.add %subgroup_channel_add, %c16 : index + %init00 = vector.fragment %zero_accumulator shape [%m, %n] : vector<8xf16> + %init01 = vector.fragment %zero_accumulator shape [%m, %n] : vector<8xf16> + %init10 = vector.fragment %zero_accumulator shape [%m, %n] : vector<8xf16> + %init11 = vector.fragment %zero_accumulator shape [%m, %n] : vector<8xf16> + %result00, %result01, %result10, %result11 = scf.for %q4_block = [%c0 to %q4_block_count step %c1](%block_acc00 = %init00 : vector<8xf16>, %block_acc01 = %init01 : vector<8xf16>, %block_acc10 = %init10 : vector<8xf16>, %block_acc11 = %init11 : vector<8xf16>) -> (vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>) { + %block_result00, %block_result01, %block_result10, %block_result11 = scf.for %q4_group = [%c0 to %c8 step %c1](%acc00 = %block_acc00 : vector<8xf16>, %acc01 = %block_acc01 : vector<8xf16>, %acc10 = %block_acc10 : vector<8xf16>, %acc11 = %block_acc11 : vector<8xf16>) -> (vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>) { + %block_k_base = index.mul %q4_block, %c256 : index + %group_k_add = index.mul %q4_group, %c32 : index + %k_origin = index.add %block_k_base, %group_k_add : index + scf.for %row_offset = [%c0 to %c64 step %c16] unroll { + %local_row0 = index.add %load_row, %row_offset : index + %local_row = index.assume %local_row0 [range(%local_row0, 0, 63)] : index + %channel = index.add %channel_tile_base, %local_row : index + %valid_channel = index.cmp ult, %channel, %bounded_output_size : index + %weight_values = scf.if %valid_channel -> (vector<4xf16>) { + %channel_byte_add = index.scale %channel, %weight_row_bytes : index, offset -> offset + %row_byte_base = index.add %expert_byte_base, %channel_byte_add : offset + %decoded = func.call @qwen3_moe_q4k_wmma_vector4(%weight_noalias, %row_byte_base, %q4_block, %q4_group, %load_packet) : (buffer, offset, index, index, index) -> (vector<4xf16>) + scf.yield %decoded : vector<4xf16> + } else { + scf.yield %c0_f16x4 : vector<4xf16> + } + %is_activation_row = index.cmp ult, %local_row, %c32 : index + vector.store %weight_values, %weight_stage_view[%local_row, %load_k] : vector<4xf16>, view<64x40xf16> + scf.if %is_activation_row { + %activation_row = index.assume %local_row [range(%local_row, 0, 31)] : index + %assignment_i32 = view.load %route_stage_view[%activation_row] : view<32xi32> -> i32 + %valid_assignment = scalar.cmpi sge, %assignment_i32, %c0_i32 : i32 + %activation_values = scf.if %valid_assignment -> (vector<4xf16>) { + %assignment0 = index.cast %assignment_i32 : i32 to index + %assignment = index.assume %assignment0 [range(%assignment0, 0, 16383)] : index + %bounded_assignment, %bounded_assignment_count = index.assume %assignment, %assignment_count [lt(%assignment, %assignment_count)] : index, index + %token0 = index.div %bounded_assignment, %bounded_route_count : index + %token, %input_token_count = index.assume %token0, %token_count [lt(%token0, %token_count)] : index, index + %input_k = index.add %k_origin, %load_k : index + %loaded = vector.load %input_view[%token, %input_k] : view<[%token_count]x[%input_size]xf32> -> vector<4xf32> + %converted = vector.fptrunc %loaded : vector<4xf32> to vector<4xf16> + scf.yield %converted : vector<4xf16> + } else { + scf.yield %c0_f16x4 : vector<4xf16> + } + vector.store %activation_values, %activation_stage_physical_view[%activation_row, %load_k] : vector<4xf16>, view<32x40xf16> + } + } + kernel.barrier scope(workgroup) ordering(acq_rel) + %next00, %next01, %next10, %next11 = scf.for %k_half = [%c0 to %c32 step %c16](%half_acc00 = %acc00 : vector<8xf16>, %half_acc01 = %acc01 : vector<8xf16>, %half_acc10 = %acc10 : vector<8xf16>, %half_acc11 = %acc11 : vector<8xf16>) -> (vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>) unroll { + %lhs0 = vector.fragment.load %weight_stage_view[%subgroup_channel_add, %k_half] shape [%m, %k] : view<64x40xf16> -> vector<16xf16> + %lhs1 = vector.fragment.load %weight_stage_view[%subgroup_channel1, %k_half] shape [%m, %k] : view<64x40xf16> -> vector<16xf16> + %rhs0 = vector.fragment.load %activation_fragment_view[%k_half, %c0] shape [%k, %n] : view<32x32xf16, %activation_fragment_layout> -> vector<16xf16> + %rhs1 = vector.fragment.load %activation_fragment_view[%k_half, %c16] shape [%k, %n] : view<32x32xf16, %activation_fragment_layout> -> vector<16xf16> + %half_next00 = vector.mma %lhs0, %rhs0, %half_acc00 : vector<16xf16>, vector<16xf16>, vector<8xf16> + %half_next01 = vector.mma %lhs0, %rhs1, %half_acc01 : vector<16xf16>, vector<16xf16>, vector<8xf16> + %half_next10 = vector.mma %lhs1, %rhs0, %half_acc10 : vector<16xf16>, vector<16xf16>, vector<8xf16> + %half_next11 = vector.mma %lhs1, %rhs1, %half_acc11 : vector<16xf16>, vector<16xf16>, vector<8xf16> + scf.yield %half_next00, %half_next01, %half_next10, %half_next11 : vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + scf.yield %next00, %next01, %next10, %next11 : vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16> + } + scf.yield %block_result00, %block_result01, %block_result10, %block_result11 : vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16> + } + // WMMA produces [channel][route] fragments. Store each through a + // transposed, wave-private LDS slice so every lane can scatter four + // contiguous channels to one routed output row. This uses the full wave + // for conversion and publication instead of serializing 16 channels onto + // each of 16 lanes. Subgroup-scoped fences suffice because waves never + // access each other's result slice. + %publish_route0 = index.div %lane, %c4 : index + %publish_route = index.assume %publish_route0 [range(%publish_route0, 0, 15)] : index + %publish_packet0 = index.rem %lane, %c4 : index + %publish_packet = index.assume %publish_packet0 [range(%publish_packet0, 0, 3)] : index + %publish_channel_add = index.mul %publish_packet, %c4 : index + %local_route1 = index.add %c16, %publish_route : index + %assignment0_i32 = view.load %route_stage_view[%publish_route] : view<32xi32> -> i32 + %assignment1_i32 = view.load %route_stage_view[%local_route1] : view<32xi32> -> i32 + %assignment0_nonnegative = scalar.cmpi sge, %assignment0_i32, %c0_i32 : i32 + %assignment1_nonnegative = scalar.cmpi sge, %assignment1_i32, %c0_i32 : i32 + %safe_assignment0_i32 = scf.if %assignment0_nonnegative -> (i32) { + scf.yield %assignment0_i32 : i32 + } else { + scf.yield %c0_i32 : i32 + } + %safe_assignment1_i32 = scf.if %assignment1_nonnegative -> (i32) { + scf.yield %assignment1_i32 : i32 + } else { + scf.yield %c0_i32 : i32 + } + %safe_assignment0_0 = index.cast %safe_assignment0_i32 : i32 to index + %safe_assignment1_0 = index.cast %safe_assignment1_i32 : i32 to index + %safe_assignment0 = index.assume %safe_assignment0_0 [range(%safe_assignment0_0, 0, 16383)] : index + %safe_assignment1 = index.assume %safe_assignment1_0 [range(%safe_assignment1_0, 0, 16383)] : index + %bounded_assignment0, %bounded_assignment_count0 = index.assume %safe_assignment0, %assignment_count [lt(%safe_assignment0, %assignment_count)] : index, index + %bounded_assignment1, %bounded_assignment_count1 = index.assume %safe_assignment1, %assignment_count [lt(%safe_assignment1, %assignment_count)] : index, index + %subgroup_channel_base = index.add %channel_tile_base, %subgroup_channel_add : index + %channel0 = index.add %subgroup_channel_base, %publish_channel_add : index + %channel1_base = index.add %subgroup_channel_base, %c16 : index + %channel1 = index.add %channel1_base, %publish_channel_add : index + %valid_channel0 = index.cmp ult, %channel0, %bounded_output_size : index + %valid_channel1 = index.cmp ult, %channel1, %bounded_output_size : index + %writes00 = scalar.andi %assignment0_nonnegative, %valid_channel0 : i1 + %writes01 = scalar.andi %assignment1_nonnegative, %valid_channel0 : i1 + %writes10 = scalar.andi %assignment0_nonnegative, %valid_channel1 : i1 + %writes11 = scalar.andi %assignment1_nonnegative, %valid_channel1 : i1 + vector.fragment.store %result00, %result_fragment_view[%c0, %c0] shape [%m, %n] : vector<8xf16>, view<16x16xf16, %result_fragment_layout> + kernel.barrier scope(subgroup) ordering(acq_rel) + scf.if %writes00 { + %values = vector.load %result_physical_view[%publish_route, %publish_channel_add] : view<16x16xf16> -> vector<4xf16> + %wide = vector.extf %values : vector<4xf16> to vector<4xf32> + %mask = vector.mask.range [%channel0 to %bounded_output_size step %c1] : index -> vector<4xi1> + vector.store.mask %wide, %output_view[%bounded_assignment0, %channel0], %mask : vector<4xf32>, view<[%output_row_count]x[%bounded_output_size]xf32>, vector<4xi1> + } + kernel.barrier scope(subgroup) ordering(acq_rel) + vector.fragment.store %result01, %result_fragment_view[%c0, %c0] shape [%m, %n] : vector<8xf16>, view<16x16xf16, %result_fragment_layout> + kernel.barrier scope(subgroup) ordering(acq_rel) + scf.if %writes01 { + %values = vector.load %result_physical_view[%publish_route, %publish_channel_add] : view<16x16xf16> -> vector<4xf16> + %wide = vector.extf %values : vector<4xf16> to vector<4xf32> + %mask = vector.mask.range [%channel0 to %bounded_output_size step %c1] : index -> vector<4xi1> + vector.store.mask %wide, %output_view[%bounded_assignment1, %channel0], %mask : vector<4xf32>, view<[%output_row_count]x[%bounded_output_size]xf32>, vector<4xi1> + } + kernel.barrier scope(subgroup) ordering(acq_rel) + vector.fragment.store %result10, %result_fragment_view[%c0, %c0] shape [%m, %n] : vector<8xf16>, view<16x16xf16, %result_fragment_layout> + kernel.barrier scope(subgroup) ordering(acq_rel) + scf.if %writes10 { + %values = vector.load %result_physical_view[%publish_route, %publish_channel_add] : view<16x16xf16> -> vector<4xf16> + %wide = vector.extf %values : vector<4xf16> to vector<4xf32> + %mask = vector.mask.range [%channel1 to %bounded_output_size step %c1] : index -> vector<4xi1> + vector.store.mask %wide, %output_view[%bounded_assignment0, %channel1], %mask : vector<4xf32>, view<[%output_row_count]x[%bounded_output_size]xf32>, vector<4xi1> + } + kernel.barrier scope(subgroup) ordering(acq_rel) + vector.fragment.store %result11, %result_fragment_view[%c0, %c0] shape [%m, %n] : vector<8xf16>, view<16x16xf16, %result_fragment_layout> + kernel.barrier scope(subgroup) ordering(acq_rel) + scf.if %writes11 { + %values = vector.load %result_physical_view[%publish_route, %publish_channel_add] : view<16x16xf16> -> vector<4xf16> + %wide = vector.extf %values : vector<4xf16> to vector<4xf32> + %mask = vector.mask.range [%channel1 to %bounded_output_size step %c1] : index -> vector<4xi1> + vector.store.mask %wide, %output_view[%bounded_assignment1, %channel1], %mask : vector<4xf32>, view<[%output_row_count]x[%bounded_output_size]xf32>, vector<4xi1> + } + // Route and result stages are reused by the next concentrated-routing + // partition. All waves must finish publication before either stage changes. + kernel.barrier scope(workgroup) ordering(acq_rel) + } + kernel.return +} + +// Fuses gate and up projection through SwiGLU for one routed expert tile. +// +// Eight waves split the 64 output channels and 32 routed rows into 16x16 result +// tiles. Each wave carries one gate and one up fragment while the workgroup +// shares one routed activation tile across both contractions. The final +// fragments meet in a wave-private LDS slice and publish the activated product +// directly, avoiding both full-size projection intermediates. +// +// The SwiGLU product is rounded to FP16 at its sole publication point because +// the grouped-down contraction consumes that exact FP16 WMMA operand. Keeping +// the transient in FP16 avoids a widen-store-load-truncate round trip and +// halves its global-memory footprint. +kernel.def target(@qwen3_moe_gfx11_wave32) @qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma(%token_count: index) { + %route_count = config.get @qwen3_moe.routed_gate_up.route_count : index + %expert_count = config.get @qwen3_moe.routed_gate_up.expert_count : index + %output_size = config.get @qwen3_moe.routed_gate_up.output_size : index + %c1 = index.constant 1 : index + %c31 = index.constant 31 : index + %c32 = index.constant 32 : index + %c63 = index.constant 63 : index + %c64 = index.constant 64 : index + %c256 = index.constant 256 : index + %padded_output_size = index.add %output_size, %c63 : index + %output_tiles = index.div %padded_output_size, %c64 : index + // The exact partition count lives in device memory. The rounding bound keeps + // it below assignment_partitions + expert_count, so launching the larger + // term lets every workgroup consume at most two strided descriptors. + %assignment_count = index.mul %token_count, %route_count : index + %rounded_assignment_count = index.add %assignment_count, %c31 : index + %assignment_partition_count = index.div %rounded_assignment_count, %c32 : index + %has_more_assignment_partitions = index.cmp ugt, %assignment_partition_count, %expert_count : index + %launch_partition_count = scf.select %has_more_assignment_partitions, %assignment_partition_count, %expert_count : index + kernel.launch.config workgroups(%output_tiles, %launch_partition_count, %c1) workgroup_size(%c256, %c1, %c1) : index +} launch(%token_count: index, %input: buffer, %expert_table: buffer, %partition_table: buffer, %gate_weight: buffer, %up_weight: buffer, %output: buffer) where [range(%token_count, 1, 2048)] { + %input_size = config.get @qwen3_moe.routed_gate_up.input_size : index + %route_count = config.get @qwen3_moe.routed_gate_up.route_count : index + %expert_count = config.get @qwen3_moe.routed_gate_up.expert_count : index + %output_size = config.get @qwen3_moe.routed_gate_up.output_size : index + %bounded_route_count = index.assume %route_count [range(%route_count, 1, 8)] : index + %bounded_expert_count = index.assume %expert_count [range(%expert_count, 1, 128)] : index + %bounded_output_size = index.assume %output_size [range(%output_size, 1, 4096)] : index + %channel_tile = kernel.workgroup.id : index + %partition_ordinal = kernel.workgroup.id : index + %workitem = kernel.workitem.id : index + %subgroup0 = kernel.subgroup.id : index + %subgroup = index.assume %subgroup0 [range(%subgroup0, 0, 7)] : index + %lane = kernel.subgroup.lane.id : index + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c4 = index.constant 4 : index + %c8 = index.constant 8 : index + %c16 = index.constant 16 : index + %c31 = index.constant 31 : index + %c32 = index.constant 32 : index + %c40 = index.constant 40 : index + %c64 = index.constant 64 : index + %c256 = index.constant 256 : index + %c0_offset = index.constant 0 : offset + %c4_bytes = index.constant 4 : offset + %q4_block_bytes = index.constant 144 : offset + %weight_stage_bytes = index.constant 5120 : offset + %activation_stage_bytes = index.constant 2560 : offset + %route_stage_bytes = index.constant 128 : offset + %wave_result_stage_bytes = index.constant 512 : offset + %result_stage_bytes = index.constant 4096 : offset + %c0_i32 = scalar.constant 0 : i32 + %cn1_i32 = scalar.constant -1 : i32 + %c0_i32x1 = vector.constant 0 : vector<1xi32> + %c0_i32x4 = vector.constant 0 : vector<4xi32> + %c0_f16x4 = vector.constant 0.0 : vector<4xf16> + %zero_accumulator = vector.constant 0.0 : vector<16xf16> + %m = index.constant 16 : index + %n = index.constant 16 : index + %k = index.constant 16 : index + %assignment_count = index.mul %token_count, %bounded_route_count : index + %rounded_assignment_count = index.add %assignment_count, %c31 : index + %assignment_partition_count = index.div %rounded_assignment_count, %c32 : index + %maximum_partition_count = index.add %assignment_partition_count, %bounded_expert_count : index + %has_more_assignment_partitions = index.cmp ugt, %assignment_partition_count, %bounded_expert_count : index + %launch_partition_count = scf.select %has_more_assignment_partitions, %assignment_partition_count, %bounded_expert_count : index + %assignment_table_byte_base = index.scale %bounded_expert_count, %c4_bytes : index, offset -> offset + %q4_block_count = index.div %input_size, %c256 : index + %weight_row_bytes = index.scale %q4_block_count, %q4_block_bytes : index, offset -> offset + %weight_expert_bytes = index.scale %bounded_output_size, %weight_row_bytes : index, offset -> offset + %output_row_count = index.mul %token_count, %bounded_route_count : index + %input_noalias, %expert_table_noalias, %partition_table_noalias, %gate_weight_noalias, %up_weight_noalias, %output_noalias = buffer.assume.noalias %input, %expert_table, %partition_table, %gate_weight, %up_weight, %output : buffer, buffer, buffer, buffer, buffer, buffer + %input_view = buffer.view %input_noalias[%c0_offset] : buffer -> view<[%token_count]x[%input_size]xf32> + %assignment_view = buffer.view %expert_table_noalias[%assignment_table_byte_base] : buffer -> view<[%bounded_expert_count]x[%token_count]xi32> + %partition_count_view = buffer.view %partition_table_noalias[%c0_offset] : buffer -> view<1xi32> + %partition_descriptor_view = buffer.view %partition_table_noalias[%c4_bytes] : buffer -> view<[%maximum_partition_count]xi32> + %output_view = buffer.view %output_noalias[%c0_offset] : buffer -> view<[%output_row_count]x[%bounded_output_size]xf16> + %weight_stage = buffer.alloca align(16) %weight_stage_bytes : buffer + %activation_stage = buffer.alloca align(16) %activation_stage_bytes : buffer + %route_stage = buffer.alloca align(16) %route_stage_bytes : buffer + %result_stage = buffer.alloca align(16) %result_stage_bytes : buffer + %weight_stage_view = buffer.view %weight_stage[%c0_offset] : buffer -> view<64x40xf16> + %activation_stage_physical_view = buffer.view %activation_stage[%c0_offset] : buffer -> view<32x40xf16> + %activation_fragment_layout = encoding.layout.strided [1, %c40] : encoding + %activation_fragment_view = buffer.view %activation_stage[%c0_offset] : buffer -> view<32x32xf16, %activation_fragment_layout> + %route_stage_view = buffer.view %route_stage[%c0_offset] : buffer -> view<32xi32> + %wave_result_stage_offset = index.scale %subgroup, %wave_result_stage_bytes : index, offset -> offset + %result_fragment_layout = encoding.layout.strided [1, %c16] : encoding + %result_fragment_view = buffer.view %result_stage[%wave_result_stage_offset] : buffer -> view<16x16xf16, %result_fragment_layout> + %result_physical_view = buffer.view %result_stage[%wave_result_stage_offset] : buffer -> view<16x16xf16> + %channel_tile_base = index.mul %channel_tile, %c64 : index + %is_workitem_zero = index.cmp eq, %workitem, %c0 : index + %lane_partition_count_i32 = scf.if %is_workitem_zero -> (i32) { + %loaded = view.load %partition_count_view[%c0] : view<1xi32> -> i32 + scf.yield %loaded : i32 + } else { + scf.yield %c0_i32 : i32 + } + %partition_count_reduced = kernel.workgroup.reduce %lane_partition_count_i32 : i32 + %partition_count_i32 = kernel.subgroup.broadcast.first %partition_count_reduced : i32 + %partition_count0 = index.cast %partition_count_i32 : i32 to index + %partition_count, %partition_capacity = index.assume %partition_count0, %maximum_partition_count [lt(%partition_count0, %maximum_partition_count)] : index, index + scf.for %active_partition = [%partition_ordinal to %partition_count step %launch_partition_count] { + %descriptor_ordinal, %descriptor_count = index.assume %active_partition, %partition_count [lt(%active_partition, %partition_count)] : index, index + %table_descriptor_ordinal, %table_descriptor_capacity = index.assume %descriptor_ordinal, %maximum_partition_count [lt(%descriptor_ordinal, %maximum_partition_count)] : index, index + %lane_descriptor_i32 = scf.if %is_workitem_zero -> (i32) { + %loaded = view.load %partition_descriptor_view[%table_descriptor_ordinal] : view<[%maximum_partition_count]xi32> -> i32 + scf.yield %loaded : i32 + } else { + scf.yield %c0_i32 : i32 + } + %descriptor_reduced_i32 = kernel.workgroup.reduce %lane_descriptor_i32 : i32 + %descriptor_i32 = kernel.subgroup.broadcast.first %descriptor_reduced_i32 : i32 + %expert, %route_tile_base, %partition_row_count = func.call @qwen3_moe_unpack_expert_partition_descriptor(%descriptor_i32) : (i32) -> (index, index, index) + %bounded_expert, %table_expert_count = index.assume %expert, %bounded_expert_count [lt(%expert, %bounded_expert_count)] : index, index + %loads_route = index.cmp ult, %workitem, %c32 : index + scf.if %loads_route { + %local_route = index.assume %workitem [range(%workitem, 0, 31)] : index + %assignment_ordinal = index.add %route_tile_base, %local_route : index + %valid_row = index.cmp ult, %local_route, %partition_row_count : index + %assignment_i32 = scf.if %valid_row -> (i32) { + %bounded_assignment_ordinal, %table_token_count = index.assume %assignment_ordinal, %token_count [lt(%assignment_ordinal, %token_count)] : index, index + %loaded = view.load %assignment_view[%bounded_expert, %bounded_assignment_ordinal] : view<[%bounded_expert_count]x[%token_count]xi32> -> i32 + scf.yield %loaded : i32 + } else { + scf.yield %cn1_i32 : i32 + } + view.store %assignment_i32, %route_stage_view[%local_route] : i32, view<32xi32> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + %load_packet = index.rem %workitem, %c8 : index + %load_k = index.mul %load_packet, %c4 : index + %load_row0 = index.div %workitem, %c8 : index + %load_row = index.assume %load_row0 [range(%load_row0, 0, 31)] : index + %expert_byte_base = index.scale %bounded_expert, %weight_expert_bytes : index, offset -> offset + %channel_subgroup = index.rem %subgroup, %c4 : index + %route_subgroup = index.div %subgroup, %c4 : index + %subgroup_channel_add = index.mul %channel_subgroup, %c16 : index + %subgroup_route_add = index.mul %route_subgroup, %c16 : index + %init_gate = vector.fragment %zero_accumulator shape [%m, %n] : vector<16xf16> + %init_up = vector.fragment %zero_accumulator shape [%m, %n] : vector<16xf16> + %gate_result, %up_result = scf.for %q4_block = [%c0 to %q4_block_count step %c1](%block_gate = %init_gate : vector<16xf16>, %block_up = %init_up : vector<16xf16>) -> (vector<16xf16>, vector<16xf16>) { + // Two lanes cover the 64 output rows owned by one load packet. Retain + // both projection headers across the complete eight-group block. + %gate_header0, %gate_header1, %up_header0, %up_header1 = scf.for %header_row_offset = [%c0 to %c64 step %c32](%prior_gate_header0 = %c0_i32x4 : vector<4xi32>, %prior_gate_header1 = %c0_i32x4 : vector<4xi32>, %prior_up_header0 = %c0_i32x4 : vector<4xi32>, %prior_up_header1 = %c0_i32x4 : vector<4xi32>) -> (vector<4xi32>, vector<4xi32>, vector<4xi32>, vector<4xi32>) unroll { + %header_local_row0 = index.add %load_row, %header_row_offset : index + %header_local_row = index.assume %header_local_row0 [range(%header_local_row0, 0, 63)] : index + %header_channel = index.add %channel_tile_base, %header_local_row : index + %valid_header_channel = index.cmp ult, %header_channel, %bounded_output_size : index + %loaded_gate_header, %loaded_up_header = scf.if %valid_header_channel -> (vector<4xi32>, vector<4xi32>) { + %header_channel_byte_add = index.scale %header_channel, %weight_row_bytes : index, offset -> offset + %header_row_byte_base = index.add %expert_byte_base, %header_channel_byte_add : offset + %gate_header = func.call @qwen3_moe_q4k_wmma_load_header(%gate_weight_noalias, %header_row_byte_base, %q4_block) : (buffer, offset, index) -> (vector<4xi32>) + %up_header = func.call @qwen3_moe_q4k_wmma_load_header(%up_weight_noalias, %header_row_byte_base, %q4_block) : (buffer, offset, index) -> (vector<4xi32>) + scf.yield %gate_header, %up_header : vector<4xi32>, vector<4xi32> + } else { + scf.yield %c0_i32x4, %c0_i32x4 : vector<4xi32>, vector<4xi32> + } + %updates_header0 = index.cmp eq, %header_row_offset, %c0 : index + %next_gate_header0 = scf.select %updates_header0, %loaded_gate_header, %prior_gate_header0 : vector<4xi32> + %next_gate_header1 = scf.select %updates_header0, %prior_gate_header1, %loaded_gate_header : vector<4xi32> + %next_up_header0 = scf.select %updates_header0, %loaded_up_header, %prior_up_header0 : vector<4xi32> + %next_up_header1 = scf.select %updates_header0, %prior_up_header1, %loaded_up_header : vector<4xi32> + scf.yield %next_gate_header0, %next_gate_header1, %next_up_header0, %next_up_header1 : vector<4xi32>, vector<4xi32>, vector<4xi32>, vector<4xi32> + } + %next_block_gate, %next_block_up = scf.for %q4_group_pair = [%c0 to %c4 step %c1](%pair_gate_acc = %block_gate : vector<16xf16>, %pair_up_acc = %block_up : vector<16xf16>) -> (vector<16xf16>, vector<16xf16>) { + // Each code word supplies the low and high nibbles for one adjacent + // group pair. Retain both projections' words for exactly those uses. + %gate_q_word0, %gate_q_word1, %up_q_word0, %up_q_word1 = scf.for %code_row_offset = [%c0 to %c64 step %c32](%prior_gate_q_word0 = %c0_i32x1 : vector<1xi32>, %prior_gate_q_word1 = %c0_i32x1 : vector<1xi32>, %prior_up_q_word0 = %c0_i32x1 : vector<1xi32>, %prior_up_q_word1 = %c0_i32x1 : vector<1xi32>) -> (vector<1xi32>, vector<1xi32>, vector<1xi32>, vector<1xi32>) unroll { + %code_local_row0 = index.add %load_row, %code_row_offset : index + %code_local_row = index.assume %code_local_row0 [range(%code_local_row0, 0, 63)] : index + %code_channel = index.add %channel_tile_base, %code_local_row : index + %valid_code_channel = index.cmp ult, %code_channel, %bounded_output_size : index + %loaded_gate_q_word, %loaded_up_q_word = scf.if %valid_code_channel -> (vector<1xi32>, vector<1xi32>) { + %bounded_q4_group_pair = index.assume %q4_group_pair [range(%q4_group_pair, 0, 3)] : index + %code_channel_byte_add = index.scale %code_channel, %weight_row_bytes : index, offset -> offset + %code_row_byte_base = index.add %expert_byte_base, %code_channel_byte_add : offset + %gate_q_word = func.call @qwen3_moe_q4k_wmma_load_code(%gate_weight_noalias, %code_row_byte_base, %q4_block, %bounded_q4_group_pair, %load_packet) : (buffer, offset, index, index, index) -> (vector<1xi32>) + %up_q_word = func.call @qwen3_moe_q4k_wmma_load_code(%up_weight_noalias, %code_row_byte_base, %q4_block, %bounded_q4_group_pair, %load_packet) : (buffer, offset, index, index, index) -> (vector<1xi32>) + scf.yield %gate_q_word, %up_q_word : vector<1xi32>, vector<1xi32> + } else { + scf.yield %c0_i32x1, %c0_i32x1 : vector<1xi32>, vector<1xi32> + } + %updates_q_word0 = index.cmp eq, %code_row_offset, %c0 : index + %next_gate_q_word0 = scf.select %updates_q_word0, %loaded_gate_q_word, %prior_gate_q_word0 : vector<1xi32> + %next_gate_q_word1 = scf.select %updates_q_word0, %prior_gate_q_word1, %loaded_gate_q_word : vector<1xi32> + %next_up_q_word0 = scf.select %updates_q_word0, %loaded_up_q_word, %prior_up_q_word0 : vector<1xi32> + %next_up_q_word1 = scf.select %updates_q_word0, %prior_up_q_word1, %loaded_up_q_word : vector<1xi32> + scf.yield %next_gate_q_word0, %next_gate_q_word1, %next_up_q_word0, %next_up_q_word1 : vector<1xi32>, vector<1xi32>, vector<1xi32>, vector<1xi32> + } + %next_pair_gate, %next_pair_up = scf.for %group_within_pair = [%c0 to %c2 step %c1](%group_gate = %pair_gate_acc : vector<16xf16>, %group_up = %pair_up_acc : vector<16xf16>) -> (vector<16xf16>, vector<16xf16>) { + %q4_group_base = index.mul %q4_group_pair, %c2 : index + %q4_group0 = index.add %q4_group_base, %group_within_pair : index + %q4_group = index.assume %q4_group0 [range(%q4_group0, 0, 7)] : index + %block_k_base = index.mul %q4_block, %c256 : index + %group_k_add = index.mul %q4_group, %c32 : index + %k_origin = index.add %block_k_base, %group_k_add : index + scf.for %row_offset = [%c0 to %c64 step %c32] unroll { + %local_row0 = index.add %load_row, %row_offset : index + %local_row = index.assume %local_row0 [range(%local_row0, 0, 63)] : index + %selects_first_row = index.cmp eq, %row_offset, %c0 : index + %selected_gate_header = scf.select %selects_first_row, %gate_header0, %gate_header1 : vector<4xi32> + %selected_gate_q_word = scf.select %selects_first_row, %gate_q_word0, %gate_q_word1 : vector<1xi32> + %channel = index.add %channel_tile_base, %local_row : index + %valid_channel = index.cmp ult, %channel, %bounded_output_size : index + %gate_values = scf.if %valid_channel -> (vector<4xf16>) { + %decoded = func.call @qwen3_moe_q4k_wmma_vector4_from_header_code(%q4_group, %selected_gate_header, %selected_gate_q_word) : (index, vector<4xi32>, vector<1xi32>) -> (vector<4xf16>) + scf.yield %decoded : vector<4xf16> + } else { + scf.yield %c0_f16x4 : vector<4xf16> + } + vector.store %gate_values, %weight_stage_view[%local_row, %load_k] : vector<4xf16>, view<64x40xf16> + } + %assignment_i32 = view.load %route_stage_view[%load_row] : view<32xi32> -> i32 + %valid_assignment = scalar.cmpi sge, %assignment_i32, %c0_i32 : i32 + %activation_values = scf.if %valid_assignment -> (vector<4xf16>) { + %assignment0 = index.cast %assignment_i32 : i32 to index + %assignment = index.assume %assignment0 [range(%assignment0, 0, 16383)] : index + %bounded_assignment, %bounded_assignment_count = index.assume %assignment, %assignment_count [lt(%assignment, %assignment_count)] : index, index + %token0 = index.div %bounded_assignment, %bounded_route_count : index + %token, %input_token_count = index.assume %token0, %token_count [lt(%token0, %token_count)] : index, index + %input_k = index.add %k_origin, %load_k : index + %loaded = vector.load %input_view[%token, %input_k] : view<[%token_count]x[%input_size]xf32> -> vector<4xf32> + %converted = vector.fptrunc %loaded : vector<4xf32> to vector<4xf16> + scf.yield %converted : vector<4xf16> + } else { + scf.yield %c0_f16x4 : vector<4xf16> + } + vector.store %activation_values, %activation_stage_physical_view[%load_row, %load_k] : vector<4xf16>, view<32x40xf16> + kernel.barrier scope(workgroup) ordering(acq_rel) + %gate_next = scf.for %k_half = [%c0 to %c32 step %c16](%half_gate = %group_gate : vector<16xf16>) -> (vector<16xf16>) unroll { + %lhs = vector.fragment.load %weight_stage_view[%subgroup_channel_add, %k_half] shape [%m, %k] : view<64x40xf16> -> vector<16xf16> + %rhs = vector.fragment.load %activation_fragment_view[%k_half, %subgroup_route_add] shape [%k, %n] : view<32x32xf16, %activation_fragment_layout> -> vector<16xf16> + %next = vector.mma %lhs, %rhs, %half_gate : vector<16xf16>, vector<16xf16>, vector<16xf16> + scf.yield %next : vector<16xf16> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + scf.for %row_offset = [%c0 to %c64 step %c32] unroll { + %local_row0 = index.add %load_row, %row_offset : index + %local_row = index.assume %local_row0 [range(%local_row0, 0, 63)] : index + %selects_first_row = index.cmp eq, %row_offset, %c0 : index + %selected_up_header = scf.select %selects_first_row, %up_header0, %up_header1 : vector<4xi32> + %selected_up_q_word = scf.select %selects_first_row, %up_q_word0, %up_q_word1 : vector<1xi32> + %channel = index.add %channel_tile_base, %local_row : index + %valid_channel = index.cmp ult, %channel, %bounded_output_size : index + %up_values = scf.if %valid_channel -> (vector<4xf16>) { + %decoded = func.call @qwen3_moe_q4k_wmma_vector4_from_header_code(%q4_group, %selected_up_header, %selected_up_q_word) : (index, vector<4xi32>, vector<1xi32>) -> (vector<4xf16>) + scf.yield %decoded : vector<4xf16> + } else { + scf.yield %c0_f16x4 : vector<4xf16> + } + vector.store %up_values, %weight_stage_view[%local_row, %load_k] : vector<4xf16>, view<64x40xf16> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + %up_next = scf.for %k_half = [%c0 to %c32 step %c16](%half_up = %group_up : vector<16xf16>) -> (vector<16xf16>) unroll { + %lhs = vector.fragment.load %weight_stage_view[%subgroup_channel_add, %k_half] shape [%m, %k] : view<64x40xf16> -> vector<16xf16> + %rhs = vector.fragment.load %activation_fragment_view[%k_half, %subgroup_route_add] shape [%k, %n] : view<32x32xf16, %activation_fragment_layout> -> vector<16xf16> + %next = vector.mma %lhs, %rhs, %half_up : vector<16xf16>, vector<16xf16>, vector<16xf16> + scf.yield %next : vector<16xf16> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + scf.yield %gate_next, %up_next : vector<16xf16>, vector<16xf16> + } + scf.yield %next_pair_gate, %next_pair_up : vector<16xf16>, vector<16xf16> + } + scf.yield %next_block_gate, %next_block_up : vector<16xf16>, vector<16xf16> + } + %publish_route0 = index.div %lane, %c2 : index + %publish_route = index.assume %publish_route0 [range(%publish_route0, 0, 15)] : index + %publish_packet0 = index.rem %lane, %c2 : index + %publish_packet = index.assume %publish_packet0 [range(%publish_packet0, 0, 1)] : index + %publish_channel_add = index.mul %publish_packet, %c8 : index + %local_route = index.add %subgroup_route_add, %publish_route : index + %assignment_i32 = view.load %route_stage_view[%local_route] : view<32xi32> -> i32 + %assignment_nonnegative = scalar.cmpi sge, %assignment_i32, %c0_i32 : i32 + %safe_assignment_i32 = scf.if %assignment_nonnegative -> (i32) { + scf.yield %assignment_i32 : i32 + } else { + scf.yield %c0_i32 : i32 + } + %safe_assignment0 = index.cast %safe_assignment_i32 : i32 to index + %safe_assignment = index.assume %safe_assignment0 [range(%safe_assignment0, 0, 16383)] : index + %bounded_assignment, %bounded_assignment_count = index.assume %safe_assignment, %assignment_count [lt(%safe_assignment, %assignment_count)] : index, index + %subgroup_channel_base = index.add %channel_tile_base, %subgroup_channel_add : index + %channel = index.add %subgroup_channel_base, %publish_channel_add : index + %valid_channel = index.cmp ult, %channel, %bounded_output_size : index + %writes = scalar.andi %assignment_nonnegative, %valid_channel : i1 + vector.fragment.store %gate_result, %result_fragment_view[%c0, %c0] shape [%m, %n] : vector<16xf16>, view<16x16xf16, %result_fragment_layout> + kernel.barrier scope(subgroup) ordering(acq_rel) + %gate_values = vector.load %result_physical_view[%publish_route, %publish_channel_add] : view<16x16xf16> -> vector<8xf16> + %gate_wide = vector.extf %gate_values : vector<8xf16> to vector<8xf32> + %activated = vector.siluf %gate_wide : vector<8xf32> + kernel.barrier scope(subgroup) ordering(acq_rel) + vector.fragment.store %up_result, %result_fragment_view[%c0, %c0] shape [%m, %n] : vector<16xf16>, view<16x16xf16, %result_fragment_layout> + kernel.barrier scope(subgroup) ordering(acq_rel) + scf.if %writes { + %up_values = vector.load %result_physical_view[%publish_route, %publish_channel_add] : view<16x16xf16> -> vector<8xf16> + %up_wide = vector.extf %up_values : vector<8xf16> to vector<8xf32> + %wide_values = vector.mulf %activated, %up_wide : vector<8xf32> + %values = vector.fptrunc %wide_values : vector<8xf32> to vector<8xf16> + %mask = vector.mask.range [%channel to %bounded_output_size step %c1] : index -> vector<8xi1> + vector.store.mask %values, %output_view[%bounded_assignment, %channel], %mask : vector<8xf16>, view<[%output_row_count]x[%bounded_output_size]xf16>, vector<8xi1> + } + kernel.barrier scope(subgroup) ordering(acq_rel) + } + kernel.return +} + +// Applies the model's SwiGLU epilogue to gate and up projections already +// scattered into logical [token][route][output channel] order. +kernel.def target(@qwen3_moe_gfx11_wave64) @qwen3_moe_routed_swiglu_f32(%token_count: index) { + %c1 = index.constant 1 : index + %workgroup_size = index.constant 256 : index + %rounding = index.constant 255 : index + %output_size = config.get @qwen3_moe.routed_gate_up.output_size : index + %route_count = config.get @qwen3_moe.routed_gate_up.route_count : index + %row_count = index.mul %token_count, %route_count : index + %element_count = index.mul %row_count, %output_size : index + %rounded_count = index.add %element_count, %rounding : index + %workgroup_count = index.div %rounded_count, %workgroup_size : index + kernel.launch.config workgroups(%workgroup_count, %c1, %c1) workgroup_size(%workgroup_size, %c1, %c1) : index +} launch(%token_count: index, %gate: buffer, %up: buffer, %output: buffer) where [range(%token_count, 1, 2048)] { + %route_count = config.get @qwen3_moe.routed_gate_up.route_count : index + %output_size = config.get @qwen3_moe.routed_gate_up.output_size : index + %bounded_route_count = index.assume %route_count [range(%route_count, 1, 8)] : index + %bounded_output_size = index.assume %output_size [range(%output_size, 1, 4096)] : index + %workgroup = kernel.workgroup.id : index + %lane = kernel.workitem.id : index + %workgroup_size = index.constant 256 : index + %c0_offset = index.constant 0 : offset + %row_count = index.mul %token_count, %bounded_route_count : index + %element_count = index.mul %row_count, %bounded_output_size : index + %element = index.madd %workgroup, %workgroup_size, %lane : index + %in_bounds = index.cmp ult, %element, %element_count : index + %gate_noalias, %up_noalias, %output_noalias = buffer.assume.noalias %gate, %up, %output : buffer, buffer, buffer + scf.if %in_bounds { + %bounded_element, %view_element_count = index.assume %element, %element_count [lt(%element, %element_count)] : index, index + %gate_view = buffer.view %gate_noalias[%c0_offset] : buffer -> view<[%view_element_count]xf32> + %up_view = buffer.view %up_noalias[%c0_offset] : buffer -> view<[%view_element_count]xf32> + %output_view = buffer.view %output_noalias[%c0_offset] : buffer -> view<[%view_element_count]xf32> + %gate_value = view.load %gate_view[%bounded_element] : view<[%view_element_count]xf32> -> f32 + %up_value = view.load %up_view[%bounded_element] : view<[%view_element_count]xf32> -> f32 + %activated = scalar.siluf %gate_value : f32 + %result = scalar.mulf %activated, %up_value : f32 + view.store %result, %output_view[%bounded_element] : f32, view<[%view_element_count]xf32> + } + kernel.return +} + +// Reference FP16 publication path for checking the fused gate/up provider at +// its actual model boundary. Gate and up projections remain independently +// materialized so this schedule does not share the fused kernel's matrix or +// publication implementation. +kernel.def target(@qwen3_moe_gfx11_wave64) @qwen3_moe_routed_swiglu_f16(%token_count: index) { + %c1 = index.constant 1 : index + %workgroup_size = index.constant 256 : index + %rounding = index.constant 255 : index + %output_size = config.get @qwen3_moe.routed_gate_up.output_size : index + %route_count = config.get @qwen3_moe.routed_gate_up.route_count : index + %row_count = index.mul %token_count, %route_count : index + %element_count = index.mul %row_count, %output_size : index + %rounded_count = index.add %element_count, %rounding : index + %workgroup_count = index.div %rounded_count, %workgroup_size : index + kernel.launch.config workgroups(%workgroup_count, %c1, %c1) workgroup_size(%workgroup_size, %c1, %c1) : index +} launch(%token_count: index, %gate: buffer, %up: buffer, %output: buffer) where [range(%token_count, 1, 2048)] { + %route_count = config.get @qwen3_moe.routed_gate_up.route_count : index + %output_size = config.get @qwen3_moe.routed_gate_up.output_size : index + %bounded_route_count = index.assume %route_count [range(%route_count, 1, 8)] : index + %bounded_output_size = index.assume %output_size [range(%output_size, 1, 4096)] : index + %workgroup = kernel.workgroup.id : index + %lane = kernel.workitem.id : index + %workgroup_size = index.constant 256 : index + %c0_offset = index.constant 0 : offset + %row_count = index.mul %token_count, %bounded_route_count : index + %element_count = index.mul %row_count, %bounded_output_size : index + %element = index.madd %workgroup, %workgroup_size, %lane : index + %in_bounds = index.cmp ult, %element, %element_count : index + %gate_noalias, %up_noalias, %output_noalias = buffer.assume.noalias %gate, %up, %output : buffer, buffer, buffer + scf.if %in_bounds { + %bounded_element, %view_element_count = index.assume %element, %element_count [lt(%element, %element_count)] : index, index + %gate_view = buffer.view %gate_noalias[%c0_offset] : buffer -> view<[%view_element_count]xf32> + %up_view = buffer.view %up_noalias[%c0_offset] : buffer -> view<[%view_element_count]xf32> + %output_view = buffer.view %output_noalias[%c0_offset] : buffer -> view<[%view_element_count]xf16> + %gate_value = view.load %gate_view[%bounded_element] : view<[%view_element_count]xf32> -> f32 + %up_value = view.load %up_view[%bounded_element] : view<[%view_element_count]xf32> -> f32 + %activated = scalar.siluf %gate_value : f32 + %wide_result = scalar.mulf %activated, %up_value : f32 + %result = scalar.fptrunc %wide_result : f32 to f16 + view.store %result, %output_view[%bounded_element] : f16, view<[%view_element_count]xf16> + } + kernel.return +} + +// The constant input is exactly representable in Q8_1 and FP16. Comparing +// against the established integer-dot path therefore isolates Q4_K decode, +// routing, FP16 matrix accumulation, tail handling, and the SwiGLU boundary. +check.case public @qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma_differential_case { + %token_count = check.literal value(67) : index + %input_size = check.literal value(512) : index + %route_count = check.literal value(2) : index + %route_stride = check.literal value(4) : index + %expert_count = check.literal value(4) : index + %output_size = check.literal value(33) : index + %input = check.generate.fill value(0.00390625) : tensor<67x512xf32> + %q8_input = check.generate.fill value(0) : tensor<67x576xi8> + %route_ids = check.generate.iota offset(0) step(1) period(4) : tensor<67x4xi32> + %expert_table = check.generate.fill value(-1) : tensor<272xi32> + %partition_table = check.generate.fill value(-1) : tensor<10xi32> + %gate_weight = check.generate.fill value(34) : tensor<4x33x2x144xi8> + %up_weight = check.generate.fill value(35) : tensor<4x33x2x144xi8> + %gate_projection = check.generate.fill value(0.0) : tensor<67x2x33xf32> + %up_projection = check.generate.fill value(0.0) : tensor<67x2x33xf32> + %expected = check.generate.fill value(0.0) : tensor<67x2x33xf32> + %actual = check.generate.fill value(1.0) : tensor<67x2x33xf32> + %actual_f16 = check.generate.fill value(1.0) : tensor<67x2x33xf16> + %fused_actual = check.generate.fill value(1.0) : tensor<67x2x33xf16> + kernel.launch @ggml_quantize_q8_1_x4_f32[%token_count, %input_size](%token_count, %input_size, %input, %q8_input) : [index, index](index, index, tensor<67x512xf32>, tensor<67x576xi8>) + kernel.launch @qwen3_moe_build_expert_table[%token_count, %route_count, %route_stride, %expert_count](%token_count, %route_count, %route_stride, %expert_count, %route_ids, %expert_table) : [index, index, index, index](index, index, index, index, tensor<67x4xi32>, tensor<272xi32>) + kernel.launch @qwen3_moe_build_expert_partition_table[%token_count, %route_count, %expert_count](%token_count, %route_count, %expert_count, %expert_table, %partition_table) : [index, index, index](index, index, index, tensor<272xi32>, tensor<10xi32>) + kernel.launch @qwen3_moe_routed_gate_up_swiglu_q4k_q8[%token_count, %route_count, %route_stride, %expert_count, %output_size](%token_count, %route_count, %route_stride, %expert_count, %output_size, %q8_input, %route_ids, %gate_weight, %up_weight, %expected) : [index, index, index, index, index](index, index, index, index, index, tensor<67x576xi8>, tensor<67x4xi32>, tensor<4x33x2x144xi8>, tensor<4x33x2x144xi8>, tensor<67x2x33xf32>) + kernel.launch @qwen3_moe_routed_linear_q4k_f16_wmma[%token_count](%token_count, %input, %expert_table, %gate_weight, %gate_projection) : [index](index, tensor<67x512xf32>, tensor<272xi32>, tensor<4x33x2x144xi8>, tensor<67x2x33xf32>) + kernel.launch @qwen3_moe_routed_linear_q4k_f16_wmma[%token_count](%token_count, %input, %expert_table, %up_weight, %up_projection) : [index](index, tensor<67x512xf32>, tensor<272xi32>, tensor<4x33x2x144xi8>, tensor<67x2x33xf32>) + kernel.launch @qwen3_moe_routed_swiglu_f32[%token_count](%token_count, %gate_projection, %up_projection, %actual) : [index](index, tensor<67x2x33xf32>, tensor<67x2x33xf32>, tensor<67x2x33xf32>) + kernel.launch @qwen3_moe_routed_swiglu_f16[%token_count](%token_count, %gate_projection, %up_projection, %actual_f16) : [index](index, tensor<67x2x33xf32>, tensor<67x2x33xf32>, tensor<67x2x33xf16>) + kernel.launch @qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma[%token_count](%token_count, %input, %expert_table, %partition_table, %gate_weight, %up_weight, %fused_actual) : [index](index, tensor<67x512xf32>, tensor<272xi32>, tensor<10xi32>, tensor<4x33x2x144xi8>, tensor<4x33x2x144xi8>, tensor<67x2x33xf16>) + check.expect.close actual(%actual) expected(%expected) atol(0.01) rtol(0.01) nan(same) : tensor<67x2x33xf32> + check.expect.close actual(%fused_actual) expected(%actual_f16) atol(0.01) rtol(0.01) nan(same) : tensor<67x2x33xf16> + check.return +} + +check.case public @qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma_benchmark_case { + %token_count = check.param.choice values([1, 2, 4, 8, 16, 17, 32, 63, 128, 129, 512, 1024, 2048]) name("token_count") : index + %route_count = check.literal value(8) : index + %route_stride = check.literal value(8) : index + %expert_count = check.literal value(128) : index + %output_size = check.literal value(768) : index + %input = check.generate.fill value(0.0) : tensor<[%token_count]x2048xf32> + // The compact iota exactly matches the oracle's balanced ring: + // expert(token, route) = (8 * token + route) % 128. + %route_ids = check.generate.iota offset(0) step(1) period(128) : tensor<[%token_count]x8xi32> + // One count per expert followed by 2048 assignment slots per expert. + %expert_table = check.generate.fill value(-1) : tensor<262272xi32> + %gate_weight = check.generate.fill value(0) : tensor<128x768x8x144xi8> + %up_weight = check.generate.fill value(0) : tensor<128x768x8x144xi8> + %gate_projection = check.generate.fill value(1.0) : tensor<[%token_count]x8x768xf32> + %up_projection = check.generate.fill value(1.0) : tensor<[%token_count]x8x768xf32> + %output = check.generate.fill value(1.0) : tensor<[%token_count]x8x768xf32> + %expected = check.generate.fill value(0.0) : tensor<[%token_count]x8x768xf32> + kernel.launch @qwen3_moe_build_expert_table[%token_count, %route_count, %route_stride, %expert_count](%token_count, %route_count, %route_stride, %expert_count, %route_ids, %expert_table) : [index, index, index, index](index, index, index, index, tensor<[%token_count]x8xi32>, tensor<262272xi32>) + kernel.launch @qwen3_moe_routed_linear_q4k_f16_wmma[%token_count](%token_count, %input, %expert_table, %gate_weight, %gate_projection) : [index](index, tensor<[%token_count]x2048xf32>, tensor<262272xi32>, tensor<128x768x8x144xi8>, tensor<[%token_count]x8x768xf32>) + kernel.launch @qwen3_moe_routed_linear_q4k_f16_wmma[%token_count](%token_count, %input, %expert_table, %up_weight, %up_projection) : [index](index, tensor<[%token_count]x2048xf32>, tensor<262272xi32>, tensor<128x768x8x144xi8>, tensor<[%token_count]x8x768xf32>) + kernel.launch @qwen3_moe_routed_swiglu_f32[%token_count](%token_count, %gate_projection, %up_projection, %output) : [index](index, tensor<[%token_count]x8x768xf32>, tensor<[%token_count]x8x768xf32>, tensor<[%token_count]x8x768xf32>) + check.expect.close actual(%output) expected(%expected) atol(0.0) rtol(0.0) nan(same) : tensor<[%token_count]x8x768xf32> + check.return +} + +check.case public @qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma_fused_benchmark_case { + %token_count = check.param.choice values([1, 2, 4, 8, 16, 17, 32, 63, 128, 129, 512, 1024, 2048]) name("token_count") : index + %route_count = check.literal value(8) : index + %route_stride = check.literal value(8) : index + %expert_count = check.literal value(128) : index + %output_size = check.literal value(768) : index + %input = check.generate.fill value(0.0) : tensor<[%token_count]x2048xf32> + %route_ids = check.generate.iota offset(0) step(1) period(128) : tensor<[%token_count]x8xi32> + // One count per expert followed by 2048 assignment slots per expert. + %expert_table = check.generate.fill value(-1) : tensor<262272xi32> + // One exact count followed by at most 640 packed descriptors. + %partition_table = check.generate.fill value(-1) : tensor<641xi32> + %gate_weight = check.generate.fill value(0) : tensor<128x768x8x144xi8> + %up_weight = check.generate.fill value(0) : tensor<128x768x8x144xi8> + %output = check.generate.fill value(1.0) : tensor<[%token_count]x8x768xf16> + %expected = check.generate.fill value(0.0) : tensor<[%token_count]x8x768xf16> + kernel.launch @qwen3_moe_build_expert_table[%token_count, %route_count, %route_stride, %expert_count](%token_count, %route_count, %route_stride, %expert_count, %route_ids, %expert_table) : [index, index, index, index](index, index, index, index, tensor<[%token_count]x8xi32>, tensor<262272xi32>) + kernel.launch @qwen3_moe_build_expert_partition_table[%token_count, %route_count, %expert_count](%token_count, %route_count, %expert_count, %expert_table, %partition_table) : [index, index, index](index, index, index, tensor<262272xi32>, tensor<641xi32>) + kernel.launch @qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma[%token_count](%token_count, %input, %expert_table, %partition_table, %gate_weight, %up_weight, %output) : [index](index, tensor<[%token_count]x2048xf32>, tensor<262272xi32>, tensor<641xi32>, tensor<128x768x8x144xi8>, tensor<128x768x8x144xi8>, tensor<[%token_count]x8x768xf16>) + check.expect.close actual(%output) expected(%expected) atol(0.0) rtol(0.0) nan(same) : tensor<[%token_count]x8x768xf16> + check.return +} + +check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma_differential_case> @qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma_differential + +check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma_decode {token_count = 1} + +check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma_small_batch_2 {token_count = 2} + +check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma_small_batch_4 {token_count = 4} + +check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma_small_batch_8 {token_count = 8} + +check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma_small_batch_16 {token_count = 16} + +check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma_prefill_17 {token_count = 17} + +check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma_prefill_32 {token_count = 32} + +check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma_prefill_63 {token_count = 63} + +check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma_prefill_128 {token_count = 128} + +check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma_prefill_129 {token_count = 129} + +check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma_prefill_512 {token_count = 512} + +check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma_prefill_1024 {token_count = 1024} + +check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma_prefill_2048 {token_count = 2048} + +check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma_fused_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma_fused_decode {token_count = 1} + +check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma_fused_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma_fused_small_batch_2 {token_count = 2} + +check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma_fused_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma_fused_small_batch_4 {token_count = 4} + +check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma_fused_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma_fused_small_batch_8 {token_count = 8} + +check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma_fused_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma_fused_small_batch_16 {token_count = 16} + +check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma_fused_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma_fused_prefill_17 {token_count = 17} + +check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma_fused_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma_fused_prefill_32 {token_count = 32} + +check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma_fused_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma_fused_prefill_63 {token_count = 63} + +check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma_fused_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma_fused_prefill_128 {token_count = 128} + +check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma_fused_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma_fused_prefill_129 {token_count = 129} + +check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma_fused_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma_fused_prefill_512 {token_count = 512} + +check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma_fused_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma_fused_prefill_1024 {token_count = 1024} + +check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma_fused_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma_fused_prefill_2048 {token_count = 2048} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/router_projection_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/router_projection_f32.loom new file mode 100644 index 000000000000..271f5c73ad14 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/router_projection_f32.loom @@ -0,0 +1,265 @@ +// Dense F32 projection from hidden activations to MoE router logits. +// +// Two matrix-vector schedules cover the measured shape classes. Decode assigns +// one expert row to each wave64, matching the llama.cpp Vulkan geometry and +// maximizing independent waves for one token. Prefill assigns four adjacent +// expert rows to each wave32 and reuses every activation packet across their +// dot products. The latter reduces activation traffic once token parallelism +// already fills the device. +// +// gfx1151 uses descriptor-backed MUBUF accesses for the decode schedule. The +// generic gfx11 provider retains global addressing where descriptor setup costs +// more than it saves. +amdgpu.target @qwen3_moe_router_projection_gfx11_wave32 {subgroup_size = 32} + +amdgpu.target @qwen3_moe_router_projection_gfx11_wave64 {subgroup_size = 64} + +amdgpu.target @qwen3_moe_router_projection_gfx1151 + +config.decl @qwen3_moe.model.hidden_size : %value: index where [range(%value, 128, 32768), mul(%value, 128)] + +config.decl @qwen3_moe.router.expert_count : %value: index where [range(%value, 32, 512), mul(%value, 32)] + +// Selects the target's preferred storage contract without exposing addressing +// policy at the projection call site. +template.decl @qwen3_moe_router_projection_storage(%input: buffer, %weight: buffer, %output: buffer) -> (buffer, buffer, buffer) +template.def<@qwen3_moe_router_projection_storage> priority(20) @qwen3_moe_router_projection_descriptor_storage(%input: buffer, %weight: buffer, %output: buffer) -> (buffer, buffer, buffer) { + %input_descriptor = buffer.assume.memory_space %input : buffer + %weight_descriptor = buffer.assume.memory_space %weight : buffer + %output_descriptor = buffer.assume.memory_space %output : buffer + template.return %input_descriptor, %weight_descriptor, %output_descriptor : buffer, buffer, buffer +} + +template.def<@qwen3_moe_router_projection_storage> priority(1) @qwen3_moe_router_projection_global_storage(%input: buffer, %weight: buffer, %output: buffer) -> (buffer, buffer, buffer) { + template.return %input, %weight, %output : buffer, buffer, buffer +} + +// Scalar differential oracle. It is reachable only from check cases and keeps +// production validation independent of the packetized wave schedule. +kernel.def target(@qwen3_moe_router_projection_gfx11_wave32) @qwen3_moe_router_projection_f32_reference(%token_count: index) { + %expert_count = config.get @qwen3_moe.router.expert_count : index + %c1 = index.constant 1 : index + kernel.launch.config workgroups(%expert_count, %token_count, %c1) workgroup_size(%c1, %c1, %c1) : index +} launch(%token_count: index, %input: buffer, %weight: buffer, %output: buffer) where [range(%token_count, 1, 2048)] { + %hidden_size0 = config.get @qwen3_moe.model.hidden_size : index + %expert_count0 = config.get @qwen3_moe.router.expert_count : index + %hidden_size, %expert_count = index.assume %hidden_size0, %expert_count0 [range(%hidden_size0, 128, 32768), mul(%hidden_size0, 128), range(%expert_count0, 32, 512), mul(%expert_count0, 32)] : index, index + %expert0 = kernel.workgroup.id : index + %token0 = kernel.workgroup.id : index + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c0_f32 = scalar.constant 0.0 : f32 + %c0_offset = index.constant 0 : offset + %expert = index.assume %expert0 [lt(%expert0, %expert_count)] : index + %valid_token = index.cmp ult, %token0, %token_count : index + %safe_token0 = scf.select %valid_token, %token0, %c0 : index + %token, %launch_token_count = index.assume %safe_token0, %token_count [lt(%safe_token0, %token_count)] : index, index + %input_noalias, %weight_noalias, %output_noalias = buffer.assume.noalias %input, %weight, %output : buffer, buffer, buffer + %input_view = buffer.view %input_noalias[%c0_offset] : buffer -> view<[%launch_token_count]x[%hidden_size]xf32> + %weight_view = buffer.view %weight_noalias[%c0_offset] : buffer -> view<[%expert_count]x[%hidden_size]xf32> + %output_view = buffer.view %output_noalias[%c0_offset] : buffer -> view<[%launch_token_count]x[%expert_count]xf32> + %sum = scf.for %channel = [%c0 to %hidden_size step %c1](%accumulator = %c0_f32 : f32) -> (f32) { + %input_value = view.load %input_view[%token, %channel] : view<[%launch_token_count]x[%hidden_size]xf32> -> f32 + %weight_value = view.load %weight_view[%expert, %channel] : view<[%expert_count]x[%hidden_size]xf32> -> f32 + %next_accumulator = scalar.fmaf %input_value, %weight_value, %accumulator : f32 + scf.yield %next_accumulator : f32 + } + scf.if %valid_token { + view.store %sum, %output_view[%token, %expert] : f32, view<[%launch_token_count]x[%expert_count]xf32> + } + kernel.return +} + +// One wave64 owns one output row. This schedule is selected for decode. +kernel.def target(@qwen3_moe_router_projection_gfx11_wave64) @qwen3_moe_router_projection_f32_one_row_wave64(%token_count: index) { + %expert_count = config.get @qwen3_moe.router.expert_count : index + %c1 = index.constant 1 : index + %wave_size = target.subgroup.size : index + kernel.launch.config workgroups(%expert_count, %token_count, %c1) workgroup_size(%wave_size, %c1, %c1) : index +} launch(%token_count: index, %input: buffer, %weight: buffer, %output: buffer) where [range(%token_count, 1, 2048)] { + %hidden_size0 = config.get @qwen3_moe.model.hidden_size : index + %expert_count0 = config.get @qwen3_moe.router.expert_count : index + %hidden_size, %expert_count = index.assume %hidden_size0, %expert_count0 [range(%hidden_size0, 128, 32768), mul(%hidden_size0, 128), range(%expert_count0, 32, 512), mul(%expert_count0, 32)] : index, index + %expert0 = kernel.workgroup.id : index + %token0 = kernel.workgroup.id : index + %lane = kernel.subgroup.lane.id : index + %c0 = index.constant 0 : index + %c3 = index.constant 3 : index + %c4 = index.constant 4 : index + %c256 = index.constant 256 : index + %c1024 = index.constant 1024 : index + %c0_i32 = scalar.constant 0 : i32 + %c0_f32 = scalar.constant 0.0 : f32 + %c0_offset = index.constant 0 : offset + %valid_token = index.cmp ult, %token0, %token_count : index + %safe_token0 = scf.select %valid_token, %token0, %c0 : index + %token, %launch_token_count = index.assume %safe_token0, %token_count [lt(%safe_token0, %token_count)] : index, index + %expert = index.assume %expert0 [lt(%expert0, %expert_count)] : index + %lane_channel = index.mul %lane, %c4 : index + %input_storage, %weight_storage, %output_storage = template.apply<@qwen3_moe_router_projection_storage>(%input, %weight, %output) : (buffer, buffer, buffer) -> (buffer, buffer, buffer) + %input_noalias, %weight_noalias, %output_noalias = buffer.assume.noalias %input_storage, %weight_storage, %output_storage : buffer, buffer, buffer + %input_view = buffer.view %input_noalias[%c0_offset] : buffer -> view<[%launch_token_count]x[%hidden_size]xf32> + %weight_view = buffer.view %weight_noalias[%c0_offset] : buffer -> view<[%expert_count]x[%hidden_size]xf32> + %output_view = buffer.view %output_noalias[%c0_offset] : buffer -> view<[%launch_token_count]x[%expert_count]xf32> + %full_channel_limit = index.sub %hidden_size, %c3 : index + %unroll_remainder = index.rem %hidden_size, %c1024 : index + %uses_unrolled_schedule = index.cmp eq, %unroll_remainder, %c0 : index + %lane_sum = scf.if %uses_unrolled_schedule -> (f32) { + %unrolled_sum = scf.for %channel = [%lane_channel to %full_channel_limit step %c256](%accumulator = %c0_f32 : f32) -> (f32) unroll(%c4) schedule(interleaved) { + %input_values = vector.load %input_view[%token, %channel] : view<[%launch_token_count]x[%hidden_size]xf32> -> vector<4xf32> + %weight_values = vector.load %weight_view[%expert, %channel] : view<[%expert_count]x[%hidden_size]xf32> -> vector<4xf32> + %next_accumulator = vector.dotf %input_values, %weight_values, %accumulator : vector<4xf32>, vector<4xf32>, f32 + scf.yield %next_accumulator : f32 + } + scf.yield %unrolled_sum : f32 + } else { + %general_sum = scf.for %channel = [%lane_channel to %full_channel_limit step %c256](%accumulator = %c0_f32 : f32) -> (f32) { + %input_values = vector.load %input_view[%token, %channel] : view<[%launch_token_count]x[%hidden_size]xf32> -> vector<4xf32> + %weight_values = vector.load %weight_view[%expert, %channel] : view<[%expert_count]x[%hidden_size]xf32> -> vector<4xf32> + %next_accumulator = vector.dotf %input_values, %weight_values, %accumulator : vector<4xf32>, vector<4xf32>, f32 + scf.yield %next_accumulator : f32 + } + scf.yield %general_sum : f32 + } + %sum = kernel.subgroup.reduce %lane_sum : f32 + %lane_i32 = index.cast %lane : index to i32 + %writes_output = scalar.cmpi eq, %lane_i32, %c0_i32 : i32 + %publishes_output = scalar.andi %valid_token, %writes_output : i1 + scf.if %publishes_output { + view.store %sum, %output_view[%token, %expert] : f32, view<[%launch_token_count]x[%expert_count]xf32> + } + kernel.return +} + +// One wave32 owns four adjacent output rows and reuses each input packet across +// their contractions. This schedule is selected once token parallelism can +// populate the device independently. +kernel.def target(@qwen3_moe_router_projection_gfx11_wave32) @qwen3_moe_router_projection_f32_four_row_wave32(%token_count: index) { + %expert_count = config.get @qwen3_moe.router.expert_count : index + %c1 = index.constant 1 : index + %c4 = index.constant 4 : index + %wave_size = target.subgroup.size : index + %expert_tiles = index.div %expert_count, %c4 : index + kernel.launch.config workgroups(%expert_tiles, %token_count, %c1) workgroup_size(%wave_size, %c1, %c1) : index +} launch(%token_count: index, %input: buffer, %weight: buffer, %output: buffer) where [range(%token_count, 1, 2048)] { + %hidden_size0 = config.get @qwen3_moe.model.hidden_size : index + %expert_count0 = config.get @qwen3_moe.router.expert_count : index + %hidden_size, %expert_count = index.assume %hidden_size0, %expert_count0 [range(%hidden_size0, 128, 32768), mul(%hidden_size0, 128), range(%expert_count0, 32, 512), mul(%expert_count0, 32)] : index, index + %expert_tile = kernel.workgroup.id : index + %token0 = kernel.workgroup.id : index + %lane = kernel.subgroup.lane.id : index + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c3 = index.constant 3 : index + %c4 = index.constant 4 : index + %c128 = index.constant 128 : index + %c1024 = index.constant 1024 : index + %c0_i32 = scalar.constant 0 : i32 + %c0_f32 = scalar.constant 0.0 : f32 + %c0_offset = index.constant 0 : offset + %c0 = index.constant 0 : index + %valid_token = index.cmp ult, %token0, %token_count : index + %safe_token0 = scf.select %valid_token, %token0, %c0 : index + %token, %launch_token_count = index.assume %safe_token0, %token_count [lt(%safe_token0, %token_count)] : index, index + %expert_base = index.mul %expert_tile, %c4 : index + %expert1 = index.add %expert_base, %c1 : index + %expert2 = index.add %expert_base, %c2 : index + %expert3 = index.add %expert_base, %c3 : index + %lane_channel = index.mul %lane, %c4 : index + %c0_f32x4 = vector.splat %c0_f32 : vector<4xf32> + %input_noalias, %weight_noalias, %output_noalias = buffer.assume.noalias %input, %weight, %output : buffer, buffer, buffer + %input_view = buffer.view %input_noalias[%c0_offset] : buffer -> view<[%launch_token_count]x[%hidden_size]xf32> + %weight_view = buffer.view %weight_noalias[%c0_offset] : buffer -> view<[%expert_count]x[%hidden_size]xf32> + %output_view = buffer.view %output_noalias[%c0_offset] : buffer -> view<[%launch_token_count]x[%expert_count]xf32> + %full_channel_limit = index.sub %hidden_size, %c3 : index + %unroll_remainder = index.rem %hidden_size, %c1024 : index + %uses_unrolled_schedule = index.cmp eq, %unroll_remainder, %c0 : index + %unroll_count = scf.select %uses_unrolled_schedule, %c4, %c1 : index + %lane_sums = scf.for %channel = [%lane_channel to %full_channel_limit step %c128](%accumulators = %c0_f32x4 : vector<4xf32>) -> (vector<4xf32>) unroll(%unroll_count) schedule(interleaved) { + %input_values = vector.load %input_view[%token, %channel] : view<[%launch_token_count]x[%hidden_size]xf32> -> vector<4xf32> + %weight_values0 = vector.load %weight_view[%expert_base, %channel] : view<[%expert_count]x[%hidden_size]xf32> -> vector<4xf32> + %weight_values1 = vector.load %weight_view[%expert1, %channel] : view<[%expert_count]x[%hidden_size]xf32> -> vector<4xf32> + %weight_values2 = vector.load %weight_view[%expert2, %channel] : view<[%expert_count]x[%hidden_size]xf32> -> vector<4xf32> + %weight_values3 = vector.load %weight_view[%expert3, %channel] : view<[%expert_count]x[%hidden_size]xf32> -> vector<4xf32> + %accumulator0 = vector.extract %accumulators[0] : vector<4xf32> -> f32 + %accumulator1 = vector.extract %accumulators[1] : vector<4xf32> -> f32 + %accumulator2 = vector.extract %accumulators[2] : vector<4xf32> -> f32 + %accumulator3 = vector.extract %accumulators[3] : vector<4xf32> -> f32 + %next0 = vector.dotf %input_values, %weight_values0, %accumulator0 : vector<4xf32>, vector<4xf32>, f32 + %next1 = vector.dotf %input_values, %weight_values1, %accumulator1 : vector<4xf32>, vector<4xf32>, f32 + %next2 = vector.dotf %input_values, %weight_values2, %accumulator2 : vector<4xf32>, vector<4xf32>, f32 + %next3 = vector.dotf %input_values, %weight_values3, %accumulator3 : vector<4xf32>, vector<4xf32>, f32 + %next_accumulators = vector.from_elements %next0, %next1, %next2, %next3 : vector<4xf32> + scf.yield %next_accumulators : vector<4xf32> + } + %sums = kernel.subgroup.reduce %lane_sums : vector<4xf32> + %lane_i32 = index.cast %lane : index to i32 + %writes_output = scalar.cmpi eq, %lane_i32, %c0_i32 : i32 + %publishes_output = scalar.andi %valid_token, %writes_output : i1 + scf.if %publishes_output { + vector.store %sums, %output_view[%token, %expert_base] : vector<4xf32>, view<[%launch_token_count]x[%expert_count]xf32> + } + kernel.return +} + +// Nonuniform inputs and weights compare both production schedules against an +// independently structured scalar contraction. The shape crosses every lane +// and all eight production output tiles. +check.case public @qwen3_moe_router_projection_f32_differential_case { + %token_count = check.literal value(2) : index + %input = check.generate.iota offset(-0.25) step(0.0078125) period(17) : tensor<2x512xf32> + %weight = check.generate.iota offset(-0.5) step(0.015625) period(31) : tensor<32x512xf32> + %expected = check.generate.fill value(0.0) : tensor<2x32xf32> + %decode_actual = check.generate.fill value(1.0) : tensor<2x32xf32> + %prefill_actual = check.generate.fill value(1.0) : tensor<2x32xf32> + kernel.launch @qwen3_moe_router_projection_f32_reference[%token_count](%token_count, %input, %weight, %expected) : [index](index, tensor<2x512xf32>, tensor<32x512xf32>, tensor<2x32xf32>) + kernel.launch @qwen3_moe_router_projection_f32_one_row_wave64[%token_count](%token_count, %input, %weight, %decode_actual) : [index](index, tensor<2x512xf32>, tensor<32x512xf32>, tensor<2x32xf32>) + kernel.launch @qwen3_moe_router_projection_f32_four_row_wave32[%token_count](%token_count, %input, %weight, %prefill_actual) : [index](index, tensor<2x512xf32>, tensor<32x512xf32>, tensor<2x32xf32>) + check.expect.close actual(%decode_actual) expected(%expected) atol(0.0001) rtol(0.0001) nan(same) : tensor<2x32xf32> + check.expect.close actual(%prefill_actual) expected(%expected) atol(0.0001) rtol(0.0001) nan(same) : tensor<2x32xf32> + check.return +} + +// Binary-exact values give every production row the analytic result 0.125, +// retaining correctness checks in each measured runtime token bucket. Both +// schedule cases expose the full shape sweep so routing decisions remain +// directly measurable. +check.case public @qwen3_moe_router_projection_f32_decode_benchmark_case { + %token_count = check.param.choice values([1, 32, 128, 512]) name("token_count") : index + %input = check.generate.fill value(0.00390625) : tensor<[%token_count]x2048xf32> + %weight = check.generate.fill value(0.015625) : tensor<128x2048xf32> + %output = check.generate.fill value(1.0) : tensor<[%token_count]x128xf32> + %expected = check.generate.fill value(0.125) : tensor<[%token_count]x128xf32> + kernel.launch @qwen3_moe_router_projection_f32_one_row_wave64[%token_count](%token_count, %input, %weight, %output) : [index](index, tensor<[%token_count]x2048xf32>, tensor<128x2048xf32>, tensor<[%token_count]x128xf32>) + check.expect.close actual(%output) expected(%expected) atol(0.0001) rtol(0.0001) nan(same) : tensor<[%token_count]x128xf32> + check.return +} + +check.case public @qwen3_moe_router_projection_f32_prefill_benchmark_case { + %token_count = check.param.choice values([1, 32, 128, 512]) name("token_count") : index + %input = check.generate.fill value(0.00390625) : tensor<[%token_count]x2048xf32> + %weight = check.generate.fill value(0.015625) : tensor<128x2048xf32> + %output = check.generate.fill value(1.0) : tensor<[%token_count]x128xf32> + %expected = check.generate.fill value(0.125) : tensor<[%token_count]x128xf32> + kernel.launch @qwen3_moe_router_projection_f32_four_row_wave32[%token_count](%token_count, %input, %weight, %output) : [index](index, tensor<[%token_count]x2048xf32>, tensor<128x2048xf32>, tensor<[%token_count]x128xf32>) + check.expect.close actual(%output) expected(%expected) atol(0.0001) rtol(0.0001) nan(same) : tensor<[%token_count]x128xf32> + check.return +} + +check.benchmark<@qwen3_moe_router_projection_f32_differential_case> @qwen3_moe_router_projection_f32_differential + +check.benchmark<@qwen3_moe_router_projection_f32_decode_benchmark_case> @qwen3_moe_router_projection_f32_decode {token_count = 1} + +check.benchmark<@qwen3_moe_router_projection_f32_prefill_benchmark_case> @qwen3_moe_router_projection_f32_prefill_32 {token_count = 32} + +check.benchmark<@qwen3_moe_router_projection_f32_prefill_benchmark_case> @qwen3_moe_router_projection_f32_prefill_128 {token_count = 128} + +check.benchmark<@qwen3_moe_router_projection_f32_prefill_benchmark_case> @qwen3_moe_router_projection_f32_prefill_512 {token_count = 512} + +check.benchmark<@qwen3_moe_router_projection_f32_prefill_benchmark_case> @qwen3_moe_router_projection_f32_prefill_schedule_decode {token_count = 1} + +check.benchmark<@qwen3_moe_router_projection_f32_decode_benchmark_case> @qwen3_moe_router_projection_f32_decode_schedule_prefill_32 {token_count = 32} + +check.benchmark<@qwen3_moe_router_projection_f32_decode_benchmark_case> @qwen3_moe_router_projection_f32_decode_schedule_prefill_128 {token_count = 128} + +check.benchmark<@qwen3_moe_router_projection_f32_decode_benchmark_case> @qwen3_moe_router_projection_f32_decode_schedule_prefill_512 {token_count = 512} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/router_projection_top8_fused_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/router_projection_top8_fused_f32.loom new file mode 100644 index 000000000000..ede30fd40da4 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/router_projection_top8_fused_f32.loom @@ -0,0 +1,231 @@ +// Fused decode router projection and deterministic normalized top-8 routing. +// +// Each wave64 computes a target-selected adjacent expert tile. A device-scope completion +// sequence publishes the row across workgroups, and the last arrival invokes +// the same row-selection contract as the standalone router. The completion +// counter is reset after route publication so reusable command buffers can +// issue the kernel repeatedly against the same storage. +// +// This schedule is decode-only. Prefill has enough token parallelism to keep +// projection and routing independent without paying completion atomics. +amdgpu.target @qwen3_moe_router_fused_gfx11_wave64 {subgroup_size = 64} + +amdgpu.target @qwen3_moe_router_fused_gfx1151 + +config.decl @qwen3_moe.model.hidden_size : %value: index where [range(%value, 128, 32768), mul(%value, 128)] + +config.decl @qwen3_moe.router.expert_count : %value: index where [range(%value, 32, 512), mul(%value, 32)] + +kernel.decl @qwen3_moe_router_projection_f32_four_row_wave32(%token_count: index) launch(%token_count: index, %input: buffer, %weight: buffer, %output: buffer) + +kernel.decl @qwen3_moe_router_top8_f32(%token_count: index, %route_id_stride: index) launch(%token_count: index, %route_id_stride: index, %logits: buffer, %route_ids: buffer, %route_weights: buffer) + +template.decl @qwen3_moe_router_fused_experts_per_wave() -> (index) +template.def<@qwen3_moe_router_fused_experts_per_wave> requires [#target.subgroup.size<64>] priority(20) @qwen3_moe_router_fused_two_experts_per_wave() -> (index) { + %c2 = index.constant 2 : index + template.return %c2 : index +} + +template.def<@qwen3_moe_router_fused_experts_per_wave> requires [#target.subgroup.size<64>] priority(1) @qwen3_moe_router_fused_four_experts_per_wave() -> (index) { + %c4 = index.constant 4 : index + template.return %c4 : index +} + +template.decl @qwen3_moe_router_projection_storage(%input: buffer, %weight: buffer, %output: buffer) -> (buffer, buffer, buffer) +template.decl @qwen3_moe_router_top8_row(%valid_token: i1, %token: index, %token_count: index, %route_id_stride: index, %logits: buffer, %route_ids: buffer, %route_weights: buffer) -> () +template.decl @qwen3_moe_router_fused_projection(%hidden_size: index, %expert_count: index, %token_count: index, %token: index, %expert_tile: index, %lane: index, %writes_projection: i1, %input: buffer, %weight: buffer, %logits: buffer) -> () +template.def<@qwen3_moe_router_fused_projection> requires [#target.subgroup.size<64>] priority(20) @qwen3_moe_router_fused_two_expert_projection(%hidden_size: index, %expert_count: index, %token_count: index, %token: index, %expert_tile: index, %lane: index, %writes_projection: i1, %input: buffer, %weight: buffer, %logits: buffer) { + %c0_offset = index.constant 0 : offset + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c4 = index.constant 4 : index + %c256 = index.constant 256 : index + %c0_f32 = scalar.constant 0.0 : f32 + %expert_base = index.mul %expert_tile, %c2 : index + %expert1 = index.add %expert_base, %c1 : index + %lane_channel = index.mul %lane, %c4 : index + %input_view = buffer.view %input[%c0_offset] : buffer -> view<[%token_count]x[%hidden_size]xf32> + %weight_view = buffer.view %weight[%c0_offset] : buffer -> view<[%expert_count]x[%hidden_size]xf32> + %logits_view = buffer.view %logits[%c0_offset] : buffer -> view<[%token_count]x[%expert_count]xf32> + %c0_f32x2 = vector.splat %c0_f32 : vector<2xf32> + %lane_sums = scf.for %channel = [%lane_channel to %hidden_size step %c256](%accumulators = %c0_f32x2 : vector<2xf32>) -> (vector<2xf32>) unroll(%c4) schedule(interleaved) { + %input_values = vector.load %input_view[%token, %channel] : view<[%token_count]x[%hidden_size]xf32> -> vector<4xf32> + %weight_values0 = vector.load %weight_view[%expert_base, %channel] : view<[%expert_count]x[%hidden_size]xf32> -> vector<4xf32> + %weight_values1 = vector.load %weight_view[%expert1, %channel] : view<[%expert_count]x[%hidden_size]xf32> -> vector<4xf32> + %accumulator0 = vector.extract %accumulators[0] : vector<2xf32> -> f32 + %accumulator1 = vector.extract %accumulators[1] : vector<2xf32> -> f32 + %next0 = vector.dotf %input_values, %weight_values0, %accumulator0 : vector<4xf32>, vector<4xf32>, f32 + %next1 = vector.dotf %input_values, %weight_values1, %accumulator1 : vector<4xf32>, vector<4xf32>, f32 + %next_accumulators = vector.from_elements %next0, %next1 : vector<2xf32> + scf.yield %next_accumulators : vector<2xf32> + } + %sums = kernel.subgroup.reduce %lane_sums : vector<2xf32> + scf.if %writes_projection { + vector.store %sums, %logits_view[%token, %expert_base] : vector<2xf32>, view<[%token_count]x[%expert_count]xf32> + } + template.return +} + +template.def<@qwen3_moe_router_fused_projection> requires [#target.subgroup.size<64>] priority(1) @qwen3_moe_router_fused_four_expert_projection(%hidden_size: index, %expert_count: index, %token_count: index, %token: index, %expert_tile: index, %lane: index, %writes_projection: i1, %input: buffer, %weight: buffer, %logits: buffer) { + %c0_offset = index.constant 0 : offset + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c3 = index.constant 3 : index + %c4 = index.constant 4 : index + %c256 = index.constant 256 : index + %c0_f32 = scalar.constant 0.0 : f32 + %expert_base = index.mul %expert_tile, %c4 : index + %expert1 = index.add %expert_base, %c1 : index + %expert2 = index.add %expert_base, %c2 : index + %expert3 = index.add %expert_base, %c3 : index + %lane_channel = index.mul %lane, %c4 : index + %input_view = buffer.view %input[%c0_offset] : buffer -> view<[%token_count]x[%hidden_size]xf32> + %weight_view = buffer.view %weight[%c0_offset] : buffer -> view<[%expert_count]x[%hidden_size]xf32> + %logits_view = buffer.view %logits[%c0_offset] : buffer -> view<[%token_count]x[%expert_count]xf32> + %c0_f32x4 = vector.splat %c0_f32 : vector<4xf32> + %lane_sums = scf.for %channel = [%lane_channel to %hidden_size step %c256](%accumulators = %c0_f32x4 : vector<4xf32>) -> (vector<4xf32>) unroll(%c4) schedule(interleaved) { + %input_values = vector.load %input_view[%token, %channel] : view<[%token_count]x[%hidden_size]xf32> -> vector<4xf32> + %weight_values0 = vector.load %weight_view[%expert_base, %channel] : view<[%expert_count]x[%hidden_size]xf32> -> vector<4xf32> + %weight_values1 = vector.load %weight_view[%expert1, %channel] : view<[%expert_count]x[%hidden_size]xf32> -> vector<4xf32> + %weight_values2 = vector.load %weight_view[%expert2, %channel] : view<[%expert_count]x[%hidden_size]xf32> -> vector<4xf32> + %weight_values3 = vector.load %weight_view[%expert3, %channel] : view<[%expert_count]x[%hidden_size]xf32> -> vector<4xf32> + %accumulator0 = vector.extract %accumulators[0] : vector<4xf32> -> f32 + %accumulator1 = vector.extract %accumulators[1] : vector<4xf32> -> f32 + %accumulator2 = vector.extract %accumulators[2] : vector<4xf32> -> f32 + %accumulator3 = vector.extract %accumulators[3] : vector<4xf32> -> f32 + %next0 = vector.dotf %input_values, %weight_values0, %accumulator0 : vector<4xf32>, vector<4xf32>, f32 + %next1 = vector.dotf %input_values, %weight_values1, %accumulator1 : vector<4xf32>, vector<4xf32>, f32 + %next2 = vector.dotf %input_values, %weight_values2, %accumulator2 : vector<4xf32>, vector<4xf32>, f32 + %next3 = vector.dotf %input_values, %weight_values3, %accumulator3 : vector<4xf32>, vector<4xf32>, f32 + %next_accumulators = vector.from_elements %next0, %next1, %next2, %next3 : vector<4xf32> + scf.yield %next_accumulators : vector<4xf32> + } + %sums = kernel.subgroup.reduce %lane_sums : vector<4xf32> + scf.if %writes_projection { + vector.store %sums, %logits_view[%token, %expert_base] : vector<4xf32>, view<[%token_count]x[%expert_count]xf32> + } + template.return +} + +kernel.def target(@qwen3_moe_router_fused_gfx11_wave64) @qwen3_moe_router_projection_top8_fused_decode_f32(%token_count: index, %route_id_stride: index) { + %expert_count = config.get @qwen3_moe.router.expert_count : index + %c1 = index.constant 1 : index + %wave_size = target.subgroup.size : index + %experts_per_wave = template.apply<@qwen3_moe_router_fused_experts_per_wave>() pure : () -> (index) + %expert_tiles = index.div %expert_count, %experts_per_wave : index + kernel.launch.config workgroups(%expert_tiles, %c1, %c1) workgroup_size(%wave_size, %c1, %c1) : index +} launch(%token_count: index, %route_id_stride: index, %input: buffer, %weight: buffer, %logits: buffer, %completion_counter: buffer, %route_ids: buffer, %route_weights: buffer) { + %hidden_size0 = config.get @qwen3_moe.model.hidden_size : index + %expert_count0 = config.get @qwen3_moe.router.expert_count : index + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 1)] : index + %bounded_route_id_stride = index.assume %route_id_stride [range(%route_id_stride, 1, 512)] : index + %hidden_size, %expert_count = index.assume %hidden_size0, %expert_count0 [range(%hidden_size0, 128, 32768), mul(%hidden_size0, 128), range(%expert_count0, 64, 512), mul(%expert_count0, 64)] : index, index + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %expert_tile_count0 = kernel.workgroup.count : index + %expert_tile_count = index.assume %expert_tile_count0 [range(%expert_tile_count0, 16, 256)] : index + %expert_tile0 = kernel.workgroup.id : index + %expert_tile = index.assume %expert_tile0 [lt(%expert_tile0, %expert_tile_count)] : index + %lane = kernel.subgroup.lane.id : index + %c0_i32 = scalar.constant 0 : i32 + %c1_i32 = scalar.constant 1 : i32 + %publishes_row = scalar.constant true : i1 + %c0_offset = index.constant 0 : offset + %counter_scratch_bytes = index.constant 4 : offset + %token = index.assume %c0 [lt(%c0, %bounded_token_count)] : index + %input_storage, %weight_storage, %logits_storage = template.apply<@qwen3_moe_router_projection_storage>(%input, %weight, %logits) : (buffer, buffer, buffer) -> (buffer, buffer, buffer) + %input_noalias, %weight_noalias, %logits_noalias, %completion_counter_noalias, %route_ids_noalias, %route_weights_noalias = buffer.assume.noalias %input_storage, %weight_storage, %logits_storage, %completion_counter, %route_ids, %route_weights : buffer, buffer, buffer, buffer, buffer, buffer + %completion_counter_aligned = buffer.assume.alignment %completion_counter_noalias {minimum_alignment = 16} : buffer + %completion_counter_view = buffer.view %completion_counter_aligned[%c0_offset] : buffer -> view<1xi32> + %counter_scratch = buffer.alloca align(4) %counter_scratch_bytes : buffer + %counter_scratch_view = buffer.view %counter_scratch[%c0_offset] : buffer -> view<1xi32> + + %lane_i32 = index.cast %lane : index to i32 + %writes_projection = scalar.cmpi eq, %lane_i32, %c0_i32 : i32 + template.apply<@qwen3_moe_router_fused_projection>(%hidden_size, %expert_count, %bounded_token_count, %token, %expert_tile, %lane, %writes_projection, %input_noalias, %weight_noalias, %logits_noalias) : (index, index, index, index, index, index, i1, buffer, buffer, buffer) + + // Every projection store precedes its workgroup's release. The last arrival + // acquires all preceding releases before any lane loads the complete row. + kernel.barrier scope(workgroup) ordering(acq_rel) + scf.if %writes_projection { + %old_counter = view.atomic.rmw %c1_i32, %completion_counter_view[%c0] {ordering = acq_rel, scope = device} : i32, view<1xi32> -> i32 + view.store %old_counter, %counter_scratch_view[%c0] : i32, view<1xi32> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + %old_counter = view.load %counter_scratch_view[%c0] : view<1xi32> -> i32 + %expert_tile_count_i32 = index.cast %expert_tile_count : index to i32 + %last_expert_tile_i32 = scalar.subi %expert_tile_count_i32, %c1_i32 : i32 + %negative_expert_tile_count_i32 = scalar.subi %c0_i32, %expert_tile_count_i32 : i32 + %is_last_projection = scalar.cmpi eq, %old_counter, %last_expert_tile_i32 : i32 + scf.if %is_last_projection { + template.apply<@qwen3_moe_router_top8_row>(%publishes_row, %token, %bounded_token_count, %bounded_route_id_stride, %logits_noalias, %route_ids_noalias, %route_weights_noalias) : (i1, index, index, index, buffer, buffer, buffer) + // The counter cannot become reusable until all route stores complete. + kernel.barrier scope(workgroup) ordering(acq_rel) + scf.if %writes_projection { + view.atomic.reduce %negative_expert_tile_count_i32, %completion_counter_view[%c0] {ordering = release, scope = device} : i32, view<1xi32> + } + } + kernel.return +} + +// Compare the fused output against the production composition, then invoke the +// fused route again against the same counter to make reset correctness visible. +check.case public @qwen3_moe_router_projection_top8_fused_differential_case { + %token_count = check.literal value(1) : index + %route_id_stride = check.literal value(8) : index + %input = check.generate.iota offset(-0.25) step(0.0078125) period(17) : tensor<1x2048xf32> + %weight = check.generate.iota offset(-0.5) step(0.015625) period(31) : tensor<128x2048xf32> + %expected_logits = check.generate.fill value(0.0) : tensor<1x128xf32> + %expected_route_ids = check.generate.fill value(-1) : tensor<1x8xi32> + %expected_route_weights = check.generate.fill value(0.0) : tensor<1x8xf32> + %actual_logits0 = check.generate.fill value(1.0) : tensor<1x128xf32> + %actual_route_ids0 = check.generate.fill value(-1) : tensor<1x8xi32> + %actual_route_weights0 = check.generate.fill value(0.0) : tensor<1x8xf32> + %actual_logits1 = check.generate.fill value(1.0) : tensor<1x128xf32> + %actual_route_ids1 = check.generate.fill value(-1) : tensor<1x8xi32> + %actual_route_weights1 = check.generate.fill value(0.0) : tensor<1x8xf32> + %completion_counter = check.generate.fill value(0) : tensor<1xi32> + %expected_counter = check.generate.fill value(0) : tensor<1xi32> + kernel.launch @qwen3_moe_router_projection_f32_four_row_wave32[%token_count](%token_count, %input, %weight, %expected_logits) : [index](index, tensor<1x2048xf32>, tensor<128x2048xf32>, tensor<1x128xf32>) + kernel.launch @qwen3_moe_router_top8_f32[%token_count, %route_id_stride](%token_count, %route_id_stride, %expected_logits, %expected_route_ids, %expected_route_weights) : [index, index](index, index, tensor<1x128xf32>, tensor<1x8xi32>, tensor<1x8xf32>) + kernel.launch @qwen3_moe_router_projection_top8_fused_decode_f32[%token_count, %route_id_stride](%token_count, %route_id_stride, %input, %weight, %actual_logits0, %completion_counter, %actual_route_ids0, %actual_route_weights0) : [index, index](index, index, tensor<1x2048xf32>, tensor<128x2048xf32>, tensor<1x128xf32>, tensor<1xi32>, tensor<1x8xi32>, tensor<1x8xf32>) + kernel.launch @qwen3_moe_router_projection_top8_fused_decode_f32[%token_count, %route_id_stride](%token_count, %route_id_stride, %input, %weight, %actual_logits1, %completion_counter, %actual_route_ids1, %actual_route_weights1) : [index, index](index, index, tensor<1x2048xf32>, tensor<128x2048xf32>, tensor<1x128xf32>, tensor<1xi32>, tensor<1x8xi32>, tensor<1x8xf32>) + check.expect.close actual(%actual_logits0) expected(%expected_logits) atol(0.001) rtol(0.001) nan(same) : tensor<1x128xf32> + check.expect.close actual(%actual_logits1) expected(%expected_logits) atol(0.001) rtol(0.001) nan(same) : tensor<1x128xf32> + check.expect.equal actual(%actual_route_ids0) expected(%expected_route_ids) : tensor<1x8xi32> + check.expect.equal actual(%actual_route_ids1) expected(%expected_route_ids) : tensor<1x8xi32> + check.expect.close actual(%actual_route_weights0) expected(%expected_route_weights) atol(0.0001) rtol(0.0001) nan(same) : tensor<1x8xf32> + check.expect.close actual(%actual_route_weights1) expected(%expected_route_weights) atol(0.0001) rtol(0.0001) nan(same) : tensor<1x8xf32> + check.expect.equal actual(%completion_counter) expected(%expected_counter) : tensor<1xi32> + check.return +} + +check.case public @qwen3_moe_router_projection_top8_fused_benchmark_case { + %token_count = check.literal value(1) : index + %route_id_stride = check.literal value(8) : index + %input = check.generate.fill value(0.00390625) : tensor<1x2048xf32> + %weight = check.generate.fill value(0.015625) : tensor<128x2048xf32> + %logits = check.generate.fill value(0.0) : tensor<1x128xf32> + %completion_counter = check.generate.fill value(0) : tensor<1xi32> + %route_ids = check.generate.fill value(-1) : tensor<1x8xi32> + %route_weights = check.generate.fill value(0.0) : tensor<1x8xf32> + kernel.launch @qwen3_moe_router_projection_top8_fused_decode_f32[%token_count, %route_id_stride](%token_count, %route_id_stride, %input, %weight, %logits, %completion_counter, %route_ids, %route_weights) : [index, index](index, index, tensor<1x2048xf32>, tensor<128x2048xf32>, tensor<1x128xf32>, tensor<1xi32>, tensor<1x8xi32>, tensor<1x8xf32>) + check.return +} + +check.case public @qwen3_moe_router_projection_top8_composed_benchmark_case { + %token_count = check.literal value(1) : index + %route_id_stride = check.literal value(8) : index + %input = check.generate.fill value(0.00390625) : tensor<1x2048xf32> + %weight = check.generate.fill value(0.015625) : tensor<128x2048xf32> + %logits = check.generate.fill value(0.0) : tensor<1x128xf32> + %route_ids = check.generate.fill value(-1) : tensor<1x8xi32> + %route_weights = check.generate.fill value(0.0) : tensor<1x8xf32> + kernel.launch @qwen3_moe_router_projection_f32_four_row_wave32[%token_count](%token_count, %input, %weight, %logits) : [index](index, tensor<1x2048xf32>, tensor<128x2048xf32>, tensor<1x128xf32>) + kernel.launch @qwen3_moe_router_top8_f32[%token_count, %route_id_stride](%token_count, %route_id_stride, %logits, %route_ids, %route_weights) : [index, index](index, index, tensor<1x128xf32>, tensor<1x8xi32>, tensor<1x8xf32>) + check.return +} + +check.benchmark<@qwen3_moe_router_projection_top8_fused_benchmark_case> @qwen3_moe_router_projection_top8_fused_decode + +check.benchmark<@qwen3_moe_router_projection_top8_composed_benchmark_case> @qwen3_moe_router_projection_top8_composed_decode diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/router_top8_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/router_top8_f32.loom new file mode 100644 index 000000000000..fd6a2ea664d0 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/router_top8_f32.loom @@ -0,0 +1,213 @@ +// Deterministic Qwen MoE routing from F32 expert logits. +// +// One wave owns one token row. Lanes load adjacent expert packets, repeatedly +// select the largest remaining logit, and break equal-logit ties in favor of +// the lower expert ordinal. Route IDs use an explicit physical row stride so +// the same kernel can publish either compact `[token][route]` rows or the +// `[token][expert]` argsort storage exposed by GGML views. +// +// Qwen applies a full expert softmax, selects the top-k probabilities, and +// renormalizes those selected probabilities. The full-softmax denominator +// cancels during renormalization: +// +// (exp(x_i) / sum_all) / sum_topk(exp(x_j) / sum_all) +// = exp(x_i) / sum_topk(exp(x_j)) +// +// Selection therefore operates directly on logits and only the selected +// values reach the exponential. This preserves the model contract while +// avoiding exponentials for experts that cannot contribute to the result. +amdgpu.target @qwen3_moe_router_gfx11_wave64 {subgroup_size = 64} + +config.decl @qwen3_moe.router.expert_count : %value: index where [range(%value, 32, 512), mul(%value, 32)] + +config.decl @qwen3_moe.router.route_count : %value: index where [range(%value, 1, 32)] + +// Selects and normalizes one logical router row. The caller owns mapping a +// subgroup to a safe physical row and tells the helper whether that row should +// publish. Keeping row semantics independent of launch geometry lets fused +// producers consume this exact tie-breaking and normalization contract. +template.decl @qwen3_moe_router_top8_row(%valid_token: i1, %token: index, %token_count: index, %route_id_stride: index, %logits: buffer, %route_ids: buffer, %route_weights: buffer) -> () +template.def<@qwen3_moe_router_top8_row> requires [#target.subgroup.size<64>] @qwen3_moe_router_top8_row_f32(%valid_token: i1, %token: index, %token_count: index, %route_id_stride: index, %logits: buffer, %route_ids: buffer, %route_weights: buffer) { + %expert_count0 = config.get @qwen3_moe.router.expert_count : index + %route_count0 = config.get @qwen3_moe.router.route_count : index + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 2048)] : index + %bounded_route_id_stride0 = index.assume %route_id_stride [range(%route_id_stride, 1, 512)] : index + %expert_count, %route_count, %bounded_route_id_stride = index.assume %expert_count0, %route_count0, %bounded_route_id_stride0 [range(%expert_count0, 32, 512), mul(%expert_count0, 32), range(%route_count0, 1, 32), le(%route_count0, %expert_count0), le(%route_count0, %bounded_route_id_stride0)] : index, index, index + %safe_token, %launch_token_count = index.assume %token, %bounded_token_count [lt(%token, %bounded_token_count)] : index, index + %lane = kernel.subgroup.lane.id : index + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %wave_size = target.subgroup.size : index + %negative_large = scalar.constant -3.4028234663852885e+38 : f32 + %largest_i32 = scalar.constant 2147483647 : i32 + %c0_offset = index.constant 0 : offset + %experts_per_lane0 = index.div %expert_count, %wave_size : index + %experts_per_lane = index.assume %experts_per_lane0 [range(%experts_per_lane0, 1, 16)] : index + %lane_expert_base = index.mul %lane, %experts_per_lane : index + %route_id_storage_count = index.mul %launch_token_count, %bounded_route_id_stride : index + %route_weight_storage_count = index.mul %launch_token_count, %route_count : index + %logits_noalias, %route_ids_noalias, %route_weights_noalias = buffer.assume.noalias %logits, %route_ids, %route_weights : buffer, buffer, buffer + %logits_view = buffer.view %logits_noalias[%c0_offset] : buffer -> view<[%launch_token_count]x[%expert_count]xf32> + %route_ids_view = buffer.view %route_ids_noalias[%c0_offset] : buffer -> view<[%route_id_storage_count]xi32> + %route_weights_view = buffer.view %route_weights_noalias[%c0_offset] : buffer -> view<[%route_weight_storage_count]xf32> + %initial_logits = vector.load %logits_view[%safe_token, %lane_expert_base] : view<[%launch_token_count]x[%expert_count]xf32> -> vector<[%experts_per_lane]xf32> + %remaining_final, %selected_logit = scf.for %route = [%c0 to %route_count step %c1](%remaining_logits = %initial_logits : vector<[%experts_per_lane]xf32>, %lane_selected_logit = %negative_large : f32) -> (vector<[%experts_per_lane]xf32>, f32) { + %local_value, %local_id = scf.for %slot = [%c0 to %experts_per_lane step %c1](%best_value = %negative_large : f32, %best_id = %largest_i32 : i32) -> (f32, i32) unroll { + %candidate_value = vector.extract %remaining_logits[%slot] : vector<[%experts_per_lane]xf32> -> f32 + %candidate_expert = index.add %lane_expert_base, %slot : index + %candidate_id = index.cast %candidate_expert : index to i32 + %is_greater = scalar.cmpf ogt, %candidate_value, %best_value : f32 + %is_equal = scalar.cmpf oeq, %candidate_value, %best_value : f32 + %is_lower_id = scalar.cmpi ult, %candidate_id, %best_id : i32 + %is_lower_tie = scalar.andi %is_equal, %is_lower_id : i1 + %is_better = scalar.ori %is_greater, %is_lower_tie : i1 + %next_value, %next_id = scf.if %is_better -> (f32, i32) { + scf.yield %candidate_value, %candidate_id : f32, i32 + } else { + scf.yield %best_value, %best_id : f32, i32 + } + scf.yield %next_value, %next_id : f32, i32 + } + %winner_value = kernel.subgroup.reduce %local_value : f32 + %matches_winner = scalar.cmpf oeq, %local_value, %winner_value : f32 + %winner_lane_mask = kernel.subgroup.vote.ballot %matches_winner : i1 -> i64 + %nonzero_winner_lane_mask = scalar.assume %winner_lane_mask [ne(%winner_lane_mask, 0)] : i64 + %winner_lane_i64 = scalar.cttzi %nonzero_winner_lane_mask : i64 + %winner_lane0 = index.cast %winner_lane_i64 : i64 to index + %winner_lane = index.assume %winner_lane0 [range(%winner_lane0, 0, 63)] : index + %local_expert0 = index.cast %local_id : i32 to index + %local_expert = index.assume %local_expert0 [range(%local_expert0, 0, 511)] : index + %winner_slot = index.rem %local_expert, %experts_per_lane : index + %owns_winner = index.cmp eq, %lane, %winner_lane : index + %next_remaining_logits = scf.if %owns_winner -> (vector<[%experts_per_lane]xf32>) { + %removed = vector.insert %negative_large into %remaining_logits[%winner_slot] : f32, vector<[%experts_per_lane]xf32> + scf.yield %removed : vector<[%experts_per_lane]xf32> + } else { + scf.yield %remaining_logits : vector<[%experts_per_lane]xf32> + } + %publishes_route_id = scalar.andi %valid_token, %owns_winner : i1 + scf.if %publishes_route_id { + %route_id_token_base = index.mul %safe_token, %bounded_route_id_stride : index + %route_id_index = index.add %route_id_token_base, %route : index + view.store %local_id, %route_ids_view[%route_id_index] : i32, view<[%route_id_storage_count]xi32> + } + %lane_publishes_route = index.cmp eq, %lane, %route : index + %publishes_route = scalar.andi %valid_token, %lane_publishes_route : i1 + %next_lane_selected_logit = scf.if %publishes_route -> (f32) { + scf.yield %winner_value : f32 + } else { + scf.yield %lane_selected_logit : f32 + } + scf.yield %next_remaining_logits, %next_lane_selected_logit : vector<[%experts_per_lane]xf32>, f32 + } + %selected_max = kernel.subgroup.reduce %selected_logit : f32 + %selected_delta = scalar.subf %selected_logit, %selected_max : f32 + %unnormalized_weight = scalar.expf %selected_delta : f32 + %selected_sum = kernel.subgroup.reduce %unnormalized_weight : f32 + %route_weight = scalar.divf %unnormalized_weight, %selected_sum : f32 + %lane_publishes_weight = index.cmp ult, %lane, %route_count : index + %publishes_weight = scalar.andi %valid_token, %lane_publishes_weight : i1 + scf.if %publishes_weight { + %route_weight_token_base = index.mul %safe_token, %route_count : index + %route_weight_index = index.add %route_weight_token_base, %lane : index + view.store %route_weight, %route_weights_view[%route_weight_index] : f32, view<[%route_weight_storage_count]xf32> + } + template.return +} + +kernel.def target(@qwen3_moe_router_gfx11_wave64) @qwen3_moe_router_top8_f32(%token_count: index, %route_id_stride: index) { + %c1 = index.constant 1 : index + %c3 = index.constant 3 : index + %c4 = index.constant 4 : index + %subgroup_size = target.subgroup.size : index + %workgroup_size = index.mul %subgroup_size, %c4 : index + %padded_token_count = index.add %token_count, %c3 : index + %workgroup_count = index.div %padded_token_count, %c4 : index + kernel.launch.config workgroups(%workgroup_count, %c1, %c1) workgroup_size(%workgroup_size, %c1, %c1) : index +} launch(%token_count: index, %route_id_stride: index, %logits: buffer, %route_ids: buffer, %route_weights: buffer) where [range(%token_count, 1, 2048), range(%route_id_stride, 1, 512)] { + %token_workgroup = kernel.workgroup.id : index + %subgroup = kernel.subgroup.id : index + %c0 = index.constant 0 : index + %c4 = index.constant 4 : index + %token_base = index.mul %token_workgroup, %c4 : index + %token0 = index.add %token_base, %subgroup : index + %valid_token = index.cmp ult, %token0, %token_count : index + %safe_token = scf.select %valid_token, %token0, %c0 : index + template.apply<@qwen3_moe_router_top8_row>(%valid_token, %safe_token, %token_count, %route_id_stride, %logits, %route_ids, %route_weights) : (i1, index, index, index, buffer, buffer, buffer) + kernel.return +} + +// Builds the deterministic route-ID oracle for two physical rows with eight +// published IDs and eight untouched padding slots per row. +kernel.def @qwen3_moe_router_top8_wide_stride_reference() { + %c1 = index.constant 1 : index + %c32 = index.constant 32 : index + kernel.launch.config workgroups(%c1, %c1, %c1) workgroup_size(%c32, %c1, %c1) : index +} launch(%route_ids: buffer) { + %element = kernel.workitem.id : index + %c7 = index.constant 7 : index + %c8 = index.constant 8 : index + %c16 = index.constant 16 : index + %c0_offset = index.constant 0 : offset + %slot = index.rem %element, %c16 : index + %publishes_id = index.cmp ult, %slot, %c8 : index + %route_ids_view = buffer.view %route_ids[%c0_offset] : buffer -> view<32xi32> + scf.if %publishes_id { + %scaled_slot = index.mul %slot, %c8 : index + %route_id_index = index.add %scaled_slot, %c7 : index + %route_id = index.cast %route_id_index : index to i32 + view.store %route_id, %route_ids_view[%element] : i32, view<32xi32> + } + kernel.return +} + +// Period-eight logits place sixteen equal maxima in different lane-local +// slots. Repeated selection must remove each winner and retain the lower-ID +// half of that tie set. Two rows leave two inactive waves in the four-wave +// production workgroup and exercise guarded tail publication. A second call +// proves that a wider physical route-ID stride preserves its padding. +check.case public @qwen3_moe_router_top8_f32_repeated_maxima_case { + %token_count = check.literal value(2) : index + %route_id_stride = check.literal value(8) : index + %wide_route_id_stride = check.literal value(16) : index + %logits = check.generate.iota offset(0.0) step(1.0) period(8) : tensor<2x128xf32> + %route_ids = check.generate.fill value(-1) : tensor<2x8xi32> + %route_weights = check.generate.fill value(0.0) : tensor<2x8xf32> + %wide_route_ids = check.generate.fill value(-1) : tensor<2x16xi32> + %wide_route_weights = check.generate.fill value(0.0) : tensor<2x8xf32> + %expected_ids = check.generate.iota offset(7) step(8) period(8) : tensor<2x8xi32> + %expected_wide_ids = check.generate.fill value(-1) : tensor<2x16xi32> + %expected_weights = check.generate.fill value(0.125) : tensor<2x8xf32> + kernel.launch @qwen3_moe_router_top8_f32[%token_count, %route_id_stride](%token_count, %route_id_stride, %logits, %route_ids, %route_weights) : [index, index](index, index, tensor<2x128xf32>, tensor<2x8xi32>, tensor<2x8xf32>) + kernel.launch @qwen3_moe_router_top8_wide_stride_reference(%expected_wide_ids) : (tensor<2x16xi32>) + kernel.launch @qwen3_moe_router_top8_f32[%token_count, %wide_route_id_stride](%token_count, %wide_route_id_stride, %logits, %wide_route_ids, %wide_route_weights) : [index, index](index, index, tensor<2x128xf32>, tensor<2x16xi32>, tensor<2x8xf32>) + check.expect.equal actual(%route_ids) expected(%expected_ids) : tensor<2x8xi32> + check.expect.close actual(%route_weights) expected(%expected_weights) atol(1.0000000000000001e-05) rtol(1.0000000000000001e-05) nan(same) : tensor<2x8xf32> + check.expect.equal actual(%wide_route_ids) expected(%expected_wide_ids) : tensor<2x16xi32> + check.expect.close actual(%wide_route_weights) expected(%expected_weights) atol(1.0000000000000001e-05) rtol(1.0000000000000001e-05) nan(same) : tensor<2x8xf32> + check.return +} + +check.case public @qwen3_moe_router_top8_f32_benchmark_case { + %token_count = check.param.choice values([1, 32, 128, 512]) name("token_count") : index + %route_id_stride = check.literal value(8) : index + %logits = check.generate.iota offset(0.0) step(1.0) period(8) : tensor<[%token_count]x128xf32> + %route_ids = check.generate.fill value(-1) : tensor<[%token_count]x8xi32> + %route_weights = check.generate.fill value(0.0) : tensor<[%token_count]x8xf32> + %expected_ids = check.generate.iota offset(7) step(8) period(8) : tensor<[%token_count]x8xi32> + %expected_weights = check.generate.fill value(0.125) : tensor<[%token_count]x8xf32> + kernel.launch @qwen3_moe_router_top8_f32[%token_count, %route_id_stride](%token_count, %route_id_stride, %logits, %route_ids, %route_weights) : [index, index](index, index, tensor<[%token_count]x128xf32>, tensor<[%token_count]x8xi32>, tensor<[%token_count]x8xf32>) + check.expect.equal actual(%route_ids) expected(%expected_ids) : tensor<[%token_count]x8xi32> + check.expect.close actual(%route_weights) expected(%expected_weights) atol(1.0000000000000001e-05) rtol(1.0000000000000001e-05) nan(same) : tensor<[%token_count]x8xf32> + check.return +} + +check.benchmark<@qwen3_moe_router_top8_f32_repeated_maxima_case> @qwen3_moe_router_top8_f32_repeated_maxima + +check.benchmark<@qwen3_moe_router_top8_f32_benchmark_case> @qwen3_moe_router_top8_f32_decode {token_count = 1} + +check.benchmark<@qwen3_moe_router_top8_f32_benchmark_case> @qwen3_moe_router_top8_f32_prefill_32 {token_count = 32} + +check.benchmark<@qwen3_moe_router_top8_f32_benchmark_case> @qwen3_moe_router_top8_f32_prefill_128 {token_count = 128} + +check.benchmark<@qwen3_moe_router_top8_f32_benchmark_case> @qwen3_moe_router_top8_f32_prefill_512 {token_count = 512} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_owned/attention_metadata.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_owned/attention_metadata.loom new file mode 100644 index 000000000000..a75a3fa0059d --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_owned/attention_metadata.loom @@ -0,0 +1,184 @@ +// Copyright 2026 The IREE Authors +// +// Licensed under the Apache License v2.0 with LLVM Exceptions. +// See https://llvm.org/LICENSE.txt for license information. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +// Derives positions, separate K/V cache indices, and the dense causal-mask bit +// pattern on device from the owned runtime's compact context-base control word. +amdgpu.target @qwen_attention_metadata_gfx11_wave64 {subgroup_size = 64} + +kernel.def target(@qwen_attention_metadata_gfx11_wave64) export("qwen_attention_metadata") @qwen_attention_metadata(%token_count: index, %context_capacity: index) { + %bounded_context_capacity = index.assume %context_capacity [range(%context_capacity, 1, 32768)] : index + %c1 = index.constant 1 : index + %c255 = index.constant 255 : index + %c256 = index.constant 256 : index + %padded_context_capacity = index.add %bounded_context_capacity, %c255 : index + %key_workgroup_count = index.div %padded_context_capacity, %c256 : index + kernel.launch.config workgroups(%key_workgroup_count, %token_count, %c1) workgroup_size(%c256, %c1, %c1) : index +} launch(%token_count: index, %context_capacity: index, %control: buffer, %positions: buffer, %key_cache_indices: buffer, %value_cache_indices: buffer, %attention_mask: buffer) where [range(%token_count, 1, 2048)] { + %bounded_context_capacity = index.assume %context_capacity [range(%context_capacity, 1, 32768)] : index + %query0 = kernel.workgroup.id : index + %query, %launch_token_count = index.assume %query0, %token_count [lt(%query0, %token_count)] : index, index + %key_workgroup = kernel.workgroup.id : index + %workitem = kernel.workitem.id : index + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c256 = index.constant 256 : index + %c0_i16 = scalar.constant 0 : i16 + // 0xfc00 is the F16 negative-infinity bit pattern. + %negative_infinity_i16 = scalar.constant -1024 : i16 + %c0_offset = index.constant 0 : offset + %control_noalias, %positions_noalias, %key_cache_indices_noalias, %value_cache_indices_noalias, %attention_mask_noalias = buffer.assume.noalias %control, %positions, %key_cache_indices, %value_cache_indices, %attention_mask : buffer, buffer, buffer, buffer, buffer + %control_view = buffer.view %control_noalias[%c0_offset] : buffer -> view<1xi32> + %context_base_raw = view.load %control_view[%c0] : view<1xi32> -> i32 + %context_base_i32 = scalar.assume %context_base_raw [range(%context_base_raw, 0, 32767)] : i32 + %context_base0 = index.cast %context_base_i32 : i32 to index + %context_base = index.assume %context_base0 [range(%context_base0, 0, 32767)] : index + %visible_count0 = index.add %context_base, %token_count : index + %visible_count, %launch_context_capacity = index.assume %visible_count0, %bounded_context_capacity [le(%visible_count0, %bounded_context_capacity)] : index, index + %key_workgroup_base = index.mul %key_workgroup, %c256 : index + %key0 = index.add %key_workgroup_base, %workitem : index + %valid_key = index.cmp ult, %key0, %launch_context_capacity : index + %safe_key0 = scf.select %valid_key, %key0, %c0 : index + %key = index.assume %safe_key0 [lt(%safe_key0, %launch_context_capacity)] : index + %mask_element_count = index.mul %launch_token_count, %launch_context_capacity : index + %mask_row_base = index.mul %query, %launch_context_capacity : index + %mask_element0 = index.add %mask_row_base, %key : index + %mask_element = index.assume %mask_element0 [lt(%mask_element0, %mask_element_count)] : index + %positions_view = buffer.view %positions_noalias[%c0_offset] : buffer -> view<[%launch_token_count]xi32> + %key_cache_indices_view = buffer.view %key_cache_indices_noalias[%c0_offset] : buffer -> view<[%launch_token_count]xi64> + %value_cache_indices_view = buffer.view %value_cache_indices_noalias[%c0_offset] : buffer -> view<[%launch_token_count]xi64> + %attention_mask_view = buffer.view %attention_mask_noalias[%c0_offset] : buffer -> view<[%mask_element_count]xi16> + %query_position0 = index.add %context_base, %query : index + %query_position = index.assume %query_position0 [lt(%query_position0, %visible_count)] : index + %causal_end = index.add %query_position, %c1 : index + %before_causal_end = index.cmp ult, %key, %causal_end : index + %before_visible_count = index.cmp ult, %key, %visible_count : index + %is_visible = scalar.andi %before_causal_end, %before_visible_count : i1 + %mask_bits = scf.select %is_visible, %c0_i16, %negative_infinity_i16 : i16 + scf.if %valid_key { + view.store %mask_bits, %attention_mask_view[%mask_element] : i16, view<[%mask_element_count]xi16> + } + %is_first_key_workgroup = index.cmp eq, %key_workgroup, %c0 : index + %is_first_workitem = index.cmp eq, %workitem, %c0 : index + %publishes_metadata = scalar.andi %is_first_key_workgroup, %is_first_workitem : i1 + scf.if %publishes_metadata { + %absolute_position_i32 = index.cast %query_position : index to i32 + %cache_index_i64 = index.cast %query_position : index to i64 + // This fixed no-ring experiment maps each logical position directly to + // the same physical row in both cache planes. + view.store %absolute_position_i32, %positions_view[%query] : i32, view<[%launch_token_count]xi32> + view.store %cache_index_i64, %key_cache_indices_view[%query] : i64, view<[%launch_token_count]xi64> + view.store %cache_index_i64, %value_cache_indices_view[%query] : i64, view<[%launch_token_count]xi64> + } + kernel.return +} + +// Publishes the single position and cache row needed by one decode issue. +// Decode attention bounds itself with the same request control word and does +// not need a materialized causal mask: every prior row is visible. +kernel.def target(@qwen_attention_metadata_gfx11_wave64) export("qwen_decode_attention_metadata") @qwen_decode_attention_metadata() { + %c1 = index.constant 1 : index + kernel.launch.config workgroups(%c1, %c1, %c1) workgroup_size(%c1, %c1, %c1) : index +} launch(%control: buffer, %positions: buffer, %key_cache_indices: buffer, %value_cache_indices: buffer) { + %c0 = index.constant 0 : index + %c0_offset = index.constant 0 : offset + %control_noalias, %positions_noalias, %key_cache_indices_noalias, %value_cache_indices_noalias = buffer.assume.noalias %control, %positions, %key_cache_indices, %value_cache_indices : buffer, buffer, buffer, buffer + %control_view = buffer.view %control_noalias[%c0_offset] : buffer -> view<1xi32> + %positions_view = buffer.view %positions_noalias[%c0_offset] : buffer -> view<1xi32> + %key_cache_indices_view = buffer.view %key_cache_indices_noalias[%c0_offset] : buffer -> view<1xi64> + %value_cache_indices_view = buffer.view %value_cache_indices_noalias[%c0_offset] : buffer -> view<1xi64> + %context_base_raw = view.load %control_view[%c0] : view<1xi32> -> i32 + %context_base = scalar.assume %context_base_raw [range(%context_base_raw, 0, 32767)] : i32 + %cache_index = scalar.extsi %context_base : i32 to i64 + view.store %context_base, %positions_view[%c0] : i32, view<1xi32> + view.store %cache_index, %key_cache_indices_view[%c0] : i64, view<1xi64> + view.store %cache_index, %value_cache_indices_view[%c0] : i64, view<1xi64> + kernel.return +} + +// Position zero exposes the causal edge directly as `[0, -inf]`. A second +// invocation proves that a nonzero context base advances all three metadata +// streams and makes every prior cache row visible. +check.case public @qwen_attention_metadata_causal_and_nonzero_base_case { + %c1 = check.literal value(1) : index + %c2 = check.literal value(2) : index + %five = check.literal value(5) : index + %zero_control = check.generate.fill value(0) : tensor<1xi32> + %zero_positions = check.generate.fill value(-1) : tensor<1xi32> + %zero_key_indices = check.generate.fill value(-1) : tensor<1xi64> + %zero_value_indices = check.generate.fill value(-1) : tensor<1xi64> + %zero_mask = check.generate.fill value(1) : tensor<2xi16> + %expected_zero = check.generate.fill value(0) : tensor<1xi32> + %expected_zero_indices = check.generate.fill value(0) : tensor<1xi64> + %expected_zero_mask = check.generate.iota offset(0) step(-1024) : tensor<2xi16> + kernel.launch @qwen_attention_metadata[%c1, %c2](%c1, %c2, %zero_control, %zero_positions, %zero_key_indices, %zero_value_indices, %zero_mask) : [index, index](index, index, tensor<1xi32>, tensor<1xi32>, tensor<1xi64>, tensor<1xi64>, tensor<2xi16>) + check.expect.equal actual(%zero_positions) expected(%expected_zero) : tensor<1xi32> + check.expect.equal actual(%zero_key_indices) expected(%expected_zero_indices) : tensor<1xi64> + check.expect.equal actual(%zero_value_indices) expected(%expected_zero_indices) : tensor<1xi64> + check.expect.equal actual(%zero_mask) expected(%expected_zero_mask) : tensor<2xi16> + %nonzero_control = check.generate.fill value(4) : tensor<1xi32> + %nonzero_positions = check.generate.fill value(-1) : tensor<1xi32> + %nonzero_key_indices = check.generate.fill value(-1) : tensor<1xi64> + %nonzero_value_indices = check.generate.fill value(-1) : tensor<1xi64> + %nonzero_mask = check.generate.fill value(1) : tensor<5xi16> + %expected_nonzero = check.generate.fill value(4) : tensor<1xi32> + %expected_nonzero_indices = check.generate.fill value(4) : tensor<1xi64> + %expected_nonzero_mask = check.generate.fill value(0) : tensor<5xi16> + kernel.launch @qwen_attention_metadata[%c1, %five](%c1, %five, %nonzero_control, %nonzero_positions, %nonzero_key_indices, %nonzero_value_indices, %nonzero_mask) : [index, index](index, index, tensor<1xi32>, tensor<1xi32>, tensor<1xi64>, tensor<1xi64>, tensor<5xi16>) + check.expect.equal actual(%nonzero_positions) expected(%expected_nonzero) : tensor<1xi32> + check.expect.equal actual(%nonzero_key_indices) expected(%expected_nonzero_indices) : tensor<1xi64> + check.expect.equal actual(%nonzero_value_indices) expected(%expected_nonzero_indices) : tensor<1xi64> + check.expect.equal actual(%nonzero_mask) expected(%expected_nonzero_mask) : tensor<5xi16> + check.return +} + +check.case public @qwen_attention_metadata_benchmark_case { + %token_count = check.param.choice values([32, 128, 512]) name("token_count") : index + %control = check.generate.fill value(0) : tensor<1xi32> + %positions = check.generate.fill value(-1) : tensor<[%token_count]xi32> + %key_cache_indices = check.generate.fill value(-1) : tensor<[%token_count]xi64> + %value_cache_indices = check.generate.fill value(-1) : tensor<[%token_count]xi64> + %attention_mask = check.generate.fill value(1) : tensor<[%token_count]x[%token_count]xi16> + %expected_positions = check.generate.iota offset(0) step(1) : tensor<[%token_count]xi32> + %expected_cache_indices = check.generate.iota offset(0) step(1) : tensor<[%token_count]xi64> + kernel.launch @qwen_attention_metadata[%token_count, %token_count](%token_count, %token_count, %control, %positions, %key_cache_indices, %value_cache_indices, %attention_mask) : [index, index](index, index, tensor<1xi32>, tensor<[%token_count]xi32>, tensor<[%token_count]xi64>, tensor<[%token_count]xi64>, tensor<[%token_count]x[%token_count]xi16>) + check.expect.equal actual(%positions) expected(%expected_positions) : tensor<[%token_count]xi32> + check.expect.equal actual(%key_cache_indices) expected(%expected_cache_indices) : tensor<[%token_count]xi64> + check.expect.equal actual(%value_cache_indices) expected(%expected_cache_indices) : tensor<[%token_count]xi64> + check.return +} + +check.case public @qwen_decode_attention_metadata_sequence_case { + %control0 = check.generate.fill value(0) : tensor<1xi32> + %positions0 = check.generate.fill value(-1) : tensor<1xi32> + %key_indices0 = check.generate.fill value(-1) : tensor<1xi64> + %value_indices0 = check.generate.fill value(-1) : tensor<1xi64> + %expected_position0 = check.generate.fill value(0) : tensor<1xi32> + %expected_indices0 = check.generate.fill value(0) : tensor<1xi64> + kernel.launch @qwen_decode_attention_metadata(%control0, %positions0, %key_indices0, %value_indices0) : (tensor<1xi32>, tensor<1xi32>, tensor<1xi64>, tensor<1xi64>) + check.expect.equal actual(%positions0) expected(%expected_position0) : tensor<1xi32> + check.expect.equal actual(%key_indices0) expected(%expected_indices0) : tensor<1xi64> + check.expect.equal actual(%value_indices0) expected(%expected_indices0) : tensor<1xi64> + + %control575 = check.generate.fill value(575) : tensor<1xi32> + %positions575 = check.generate.fill value(-1) : tensor<1xi32> + %key_indices575 = check.generate.fill value(-1) : tensor<1xi64> + %value_indices575 = check.generate.fill value(-1) : tensor<1xi64> + %expected_position575 = check.generate.fill value(575) : tensor<1xi32> + %expected_indices575 = check.generate.fill value(575) : tensor<1xi64> + kernel.launch @qwen_decode_attention_metadata(%control575, %positions575, %key_indices575, %value_indices575) : (tensor<1xi32>, tensor<1xi32>, tensor<1xi64>, tensor<1xi64>) + check.expect.equal actual(%positions575) expected(%expected_position575) : tensor<1xi32> + check.expect.equal actual(%key_indices575) expected(%expected_indices575) : tensor<1xi64> + check.expect.equal actual(%value_indices575) expected(%expected_indices575) : tensor<1xi64> + check.return +} + +check.benchmark<@qwen_attention_metadata_causal_and_nonzero_base_case> @qwen_attention_metadata_causal_and_nonzero_base + +check.benchmark<@qwen_attention_metadata_benchmark_case> @qwen_attention_metadata_prefill_32 {token_count = 32} + +check.benchmark<@qwen_attention_metadata_benchmark_case> @qwen_attention_metadata_prefill_128 {token_count = 128} + +check.benchmark<@qwen_attention_metadata_benchmark_case> @qwen_attention_metadata_prefill_512 {token_count = 512} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_owned/attention_metadata_bringup_workaround.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_owned/attention_metadata_bringup_workaround.loom new file mode 100644 index 000000000000..079c1465fefb --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_owned/attention_metadata_bringup_workaround.loom @@ -0,0 +1,170 @@ +// Copyright 2026 The IREE Authors +// +// Licensed under the Apache License v2.0 with LLVM Exceptions. +// See https://llvm.org/LICENSE.txt for license information. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +// Temporary, non-sanctioned one-function kernel for Qwen bring-up. This is not +// a metadata-kernel framework or a second Loom authoring path. The owned +// runtime carries one compact context-base control word and needs positions, +// separate K/V cache indices, and the dense causal-mask bit pattern derived on +// device before attention begins. Delete this fork when the Qwen kernel corpus +// provides the canonical producer. +amdgpu.target @qwen_attention_metadata_gfx11_wave64 {subgroup_size = 64} + +kernel.def target(@qwen_attention_metadata_gfx11_wave64) export("qwen_attention_metadata_bringup_workaround") @qwen_attention_metadata_bringup_workaround(%token_count: index, %context_capacity: index) { + %one = index.constant 1 : index + %twofiftyfive = index.constant 255 : index + %twofiftysix = index.constant 256 : index + %padded_context_capacity = index.add %context_capacity, %twofiftyfive : index + %key_workgroup_count = index.div %padded_context_capacity, %twofiftysix : index + kernel.launch.config workgroups(%key_workgroup_count, %token_count, %one) workgroup_size(%twofiftysix, %one, %one) : index +} launch(%token_count: index, %context_capacity: index, %control: buffer, %positions: buffer, %key_cache_indices: buffer, %value_cache_indices: buffer, %attention_mask: buffer) { + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 2048)] : index + %bounded_context_capacity = index.assume %context_capacity [range(%context_capacity, 1, 32768)] : index + %query0 = kernel.workgroup.id : index + %query, %launch_token_count = index.assume %query0, %bounded_token_count [lt(%query0, %bounded_token_count)] : index, index + %key_workgroup = kernel.workgroup.id : index + %workitem = kernel.workitem.id : index + %zero = index.constant 0 : index + %one = index.constant 1 : index + %twofiftysix = index.constant 256 : index + %zero_i16 = scalar.constant 0 : i16 + // 0xfc00 is the F16 negative-infinity bit pattern. + %negative_infinity_i16 = scalar.constant -1024 : i16 + %zero_offset = index.constant 0 : offset + %control_noalias, %positions_noalias, %key_cache_indices_noalias, %value_cache_indices_noalias, %attention_mask_noalias = buffer.assume.noalias %control, %positions, %key_cache_indices, %value_cache_indices, %attention_mask : buffer, buffer, buffer, buffer, buffer + %control_view = buffer.view %control_noalias[%zero_offset] : buffer -> view<1xi32> + %context_base_raw = view.load %control_view[%zero] : view<1xi32> -> i32 + %context_base_i32 = scalar.assume %context_base_raw [range(%context_base_raw, 0, 32767)] : i32 + %context_base0 = index.cast %context_base_i32 : i32 to index + %context_base = index.assume %context_base0 [range(%context_base0, 0, 32767)] : index + %visible_count0 = index.add %context_base, %bounded_token_count : index + %visible_count, %launch_context_capacity = index.assume %visible_count0, %bounded_context_capacity [le(%visible_count0, %bounded_context_capacity)] : index, index + %key_workgroup_base = index.mul %key_workgroup, %twofiftysix : index + %key0 = index.add %key_workgroup_base, %workitem : index + %valid_key = index.cmp ult, %key0, %launch_context_capacity : index + %safe_key0 = scf.select %valid_key, %key0, %zero : index + %key = index.assume %safe_key0 [lt(%safe_key0, %launch_context_capacity)] : index + %mask_element_count = index.mul %launch_token_count, %launch_context_capacity : index + %mask_row_base = index.mul %query, %launch_context_capacity : index + %mask_element0 = index.add %mask_row_base, %key : index + %mask_element = index.assume %mask_element0 [lt(%mask_element0, %mask_element_count)] : index + %positions_view = buffer.view %positions_noalias[%zero_offset] : buffer -> view<[%launch_token_count]xi32> + %key_cache_indices_view = buffer.view %key_cache_indices_noalias[%zero_offset] : buffer -> view<[%launch_token_count]xi64> + %value_cache_indices_view = buffer.view %value_cache_indices_noalias[%zero_offset] : buffer -> view<[%launch_token_count]xi64> + %attention_mask_view = buffer.view %attention_mask_noalias[%zero_offset] : buffer -> view<[%mask_element_count]xi16> + %query_position0 = index.add %context_base, %query : index + %query_position = index.assume %query_position0 [lt(%query_position0, %visible_count)] : index + %causal_end = index.add %query_position, %one : index + %before_causal_end = index.cmp ult, %key, %causal_end : index + %before_visible_count = index.cmp ult, %key, %visible_count : index + %is_visible = scalar.andi %before_causal_end, %before_visible_count : i1 + %mask_bits = scf.select %is_visible, %zero_i16, %negative_infinity_i16 : i16 + scf.if %valid_key { + view.store %mask_bits, %attention_mask_view[%mask_element] : i16, view<[%mask_element_count]xi16> + } + %is_first_key_workgroup = index.cmp eq, %key_workgroup, %zero : index + %is_first_workitem = index.cmp eq, %workitem, %zero : index + %publishes_metadata = scalar.andi %is_first_key_workgroup, %is_first_workitem : i1 + scf.if %publishes_metadata { + %absolute_position_i32 = index.cast %query_position : index to i32 + %cache_index_i64 = index.cast %query_position : index to i64 + // This fixed no-ring experiment maps each logical position directly to + // the same physical row in both cache planes. + view.store %absolute_position_i32, %positions_view[%query] : i32, view<[%launch_token_count]xi32> + view.store %cache_index_i64, %key_cache_indices_view[%query] : i64, view<[%launch_token_count]xi64> + view.store %cache_index_i64, %value_cache_indices_view[%query] : i64, view<[%launch_token_count]xi64> + } + kernel.return +} + +// Position zero exposes the causal edge directly as `[0, -inf]`. A second +// invocation proves that a nonzero context base advances all three metadata +// streams and makes every prior cache row visible. +check.case public @qwen_attention_metadata_causal_and_nonzero_base_case { + %one = check.literal value(1) : index + %two = check.literal value(2) : index + %five = check.literal value(5) : index + %zero_control = check.generate.fill value(0) : tensor<1xi32> + %zero_positions = check.generate.fill value(-1) : tensor<1xi32> + %zero_key_indices = check.generate.fill value(-1) : tensor<1xi64> + %zero_value_indices = check.generate.fill value(-1) : tensor<1xi64> + %zero_mask = check.generate.fill value(1) : tensor<2xi16> + %expected_zero = check.generate.fill value(0) : tensor<1xi32> + %expected_zero_indices = check.generate.fill value(0) : tensor<1xi64> + %expected_zero_mask = check.generate.iota offset(0) step(-1024) : tensor<2xi16> + func.call @qwen_attention_metadata_bringup_workaround(%one, %two, %zero_control, %zero_positions, %zero_key_indices, %zero_value_indices, %zero_mask) : (index, index, tensor<1xi32>, tensor<1xi32>, tensor<1xi64>, tensor<1xi64>, tensor<2xi16>) + check.expect.equal actual(%zero_positions) expected(%expected_zero) : tensor<1xi32> + check.expect.equal actual(%zero_key_indices) expected(%expected_zero_indices) : tensor<1xi64> + check.expect.equal actual(%zero_value_indices) expected(%expected_zero_indices) : tensor<1xi64> + check.expect.equal actual(%zero_mask) expected(%expected_zero_mask) : tensor<2xi16> + %nonzero_control = check.generate.fill value(4) : tensor<1xi32> + %nonzero_positions = check.generate.fill value(-1) : tensor<1xi32> + %nonzero_key_indices = check.generate.fill value(-1) : tensor<1xi64> + %nonzero_value_indices = check.generate.fill value(-1) : tensor<1xi64> + %nonzero_mask = check.generate.fill value(1) : tensor<5xi16> + %expected_nonzero = check.generate.fill value(4) : tensor<1xi32> + %expected_nonzero_indices = check.generate.fill value(4) : tensor<1xi64> + %expected_nonzero_mask = check.generate.fill value(0) : tensor<5xi16> + func.call @qwen_attention_metadata_bringup_workaround(%one, %five, %nonzero_control, %nonzero_positions, %nonzero_key_indices, %nonzero_value_indices, %nonzero_mask) : (index, index, tensor<1xi32>, tensor<1xi32>, tensor<1xi64>, tensor<1xi64>, tensor<5xi16>) + check.expect.equal actual(%nonzero_positions) expected(%expected_nonzero) : tensor<1xi32> + check.expect.equal actual(%nonzero_key_indices) expected(%expected_nonzero_indices) : tensor<1xi64> + check.expect.equal actual(%nonzero_value_indices) expected(%expected_nonzero_indices) : tensor<1xi64> + check.expect.equal actual(%nonzero_mask) expected(%expected_nonzero_mask) : tensor<5xi16> + check.return +} + +check.case public @qwen_attention_metadata_benchmark_case { + %token_count = check.param.choice values([32, 128, 512]) name("token_count") : index + %control = check.generate.fill value(0) : tensor<1xi32> + %positions = check.generate.fill value(-1) : tensor<[%token_count]xi32> + %key_cache_indices = check.generate.fill value(-1) : tensor<[%token_count]xi64> + %value_cache_indices = check.generate.fill value(-1) : tensor<[%token_count]xi64> + %attention_mask = check.generate.fill value(1) : tensor<[%token_count]x[%token_count]xi16> + %expected_positions = check.generate.iota offset(0) step(1) : tensor<[%token_count]xi32> + %expected_cache_indices = check.generate.iota offset(0) step(1) : tensor<[%token_count]xi64> + func.call @qwen_attention_metadata_bringup_workaround(%token_count, %token_count, %control, %positions, %key_cache_indices, %value_cache_indices, %attention_mask) : (index, index, tensor<1xi32>, tensor<[%token_count]xi32>, tensor<[%token_count]xi64>, tensor<[%token_count]xi64>, tensor<[%token_count]x[%token_count]xi16>) + check.expect.equal actual(%positions) expected(%expected_positions) : tensor<[%token_count]xi32> + check.expect.equal actual(%key_cache_indices) expected(%expected_cache_indices) : tensor<[%token_count]xi64> + check.expect.equal actual(%value_cache_indices) expected(%expected_cache_indices) : tensor<[%token_count]xi64> + check.return +} + +// Locked decode geometry: one query row addressing the explicit 768-row KV +// bucket recovered from the live llama.cpp graph. +check.case public @qwen_attention_metadata_decode_768_case { + %token_count = check.literal value(1) : index + %context_capacity = check.literal value(768) : index + %control = check.generate.fill value(767) : tensor<1xi32> + %positions = check.generate.fill value(-1) : tensor<1xi32> + %key_cache_indices = check.generate.fill value(-1) : tensor<1xi64> + %value_cache_indices = check.generate.fill value(-1) : tensor<1xi64> + %attention_mask = check.generate.fill value(1) : tensor<768xi16> + func.call @qwen_attention_metadata_bringup_workaround(%token_count, %context_capacity, %control, %positions, %key_cache_indices, %value_cache_indices, %attention_mask) : (index, index, tensor<1xi32>, tensor<1xi32>, tensor<1xi64>, tensor<1xi64>, tensor<768xi16>) + check.return +} + +check.case public @qwen_attention_metadata_prefill_512_case { + %token_count = check.literal value(512) : index + %context_capacity = check.literal value(512) : index + %control = check.generate.fill value(0) : tensor<1xi32> + %positions = check.generate.fill value(-1) : tensor<512xi32> + %key_cache_indices = check.generate.fill value(-1) : tensor<512xi64> + %value_cache_indices = check.generate.fill value(-1) : tensor<512xi64> + %attention_mask = check.generate.fill value(1) : tensor<512x512xi16> + func.call @qwen_attention_metadata_bringup_workaround(%token_count, %context_capacity, %control, %positions, %key_cache_indices, %value_cache_indices, %attention_mask) : (index, index, tensor<1xi32>, tensor<512xi32>, tensor<512xi64>, tensor<512xi64>, tensor<512x512xi16>) + check.return +} + +check.benchmark<@qwen_attention_metadata_causal_and_nonzero_base_case> @qwen_attention_metadata_causal_and_nonzero_base + +check.benchmark<@qwen_attention_metadata_decode_768_case> @qwen_attention_metadata_decode_768 + +check.benchmark<@qwen_attention_metadata_prefill_512_case> @qwen_attention_metadata_model_prefill_512 + +check.benchmark<@qwen_attention_metadata_benchmark_case> @qwen_attention_metadata_prefill_32 {token_count = 32} + +check.benchmark<@qwen_attention_metadata_benchmark_case> @qwen_attention_metadata_prefill_128 {token_count = 128} + +check.benchmark<@qwen_attention_metadata_benchmark_case> @qwen_attention_metadata_prefill_512 {token_count = 512} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_owned/attention_state_initialize.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_owned/attention_state_initialize.loom new file mode 100644 index 000000000000..bade957f43d0 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_owned/attention_state_initialize.loom @@ -0,0 +1,112 @@ +// Copyright 2026 The IREE Authors +// +// Licensed under the Apache License v2.0 with LLVM Exceptions. +// See https://llvm.org/LICENSE.txt for license information. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +amdgpu.target @qwen_attention_state_gfx11_wave64 {subgroup_size = 64} + +// Captures the position before the metadata producer overwrites the positions +// buffer with the canonical sequence used by attention. +kernel.def target(@qwen_attention_state_gfx11_wave64) export("qwen_attention_context_base_capture") @qwen_attention_context_base_capture() { + %c1 = index.constant 1 : index + kernel.launch.config workgroups(%c1, %c1, %c1) workgroup_size(%c1, %c1, %c1) : index +} launch(%positions: buffer, %control: buffer) { + %c0 = index.constant 0 : index + %c0_offset = index.constant 0 : offset + %positions_noalias, %control_noalias = buffer.assume.noalias %positions, %control : buffer, buffer + %positions_view = buffer.view %positions_noalias[%c0_offset] : buffer -> view<1xi32> + %control_view = buffer.view %control_noalias[%c0_offset] : buffer -> view<1xi32> + %context_base = view.load %positions_view[%c0] : view<1xi32> -> i32 + view.store %context_base, %control_view[%c0] : i32, view<1xi32> + kernel.return +} + +// Decode has one query row, so the input position is already the absolute +// position consumed by attention. Publish its cache indices and causal mask +// while initializing every self-resetting completion counter in the replay. +kernel.def target(@qwen_attention_state_gfx11_wave64) export("qwen_attention_decode_state_initialize") @qwen_attention_decode_state_initialize(%context_capacity: index, %completion_counter_count: index) { + %c1 = index.constant 1 : index + %c256 = index.constant 256 : index + kernel.launch.config workgroups(%c1, %c1, %c1) workgroup_size(%c256, %c1, %c1) : index +} launch(%context_capacity: index, %completion_counter_count: index, %positions: buffer, %key_cache_indices: buffer, %value_cache_indices: buffer, %attention_mask: buffer, %completion_counters: buffer) { + %bounded_context_capacity = index.assume %context_capacity [range(%context_capacity, 1, 32768)] : index + %counter_count = index.assume %completion_counter_count [range(%completion_counter_count, 1, 64)] : index + %workitem0 = kernel.workitem.id : index + %workitem = index.assume %workitem0 [range(%workitem0, 0, 255)] : index + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c256 = index.constant 256 : index + %c0_i16 = scalar.constant 0 : i16 + %c0_i32 = scalar.constant 0 : i32 + // 0xfc00 is the F16 negative-infinity bit pattern. + %negative_infinity_i16 = scalar.constant -1024 : i16 + %c0_offset = index.constant 0 : offset + %positions_noalias, %key_indices_noalias, %value_indices_noalias, %mask_noalias, %counters_noalias = buffer.assume.noalias %positions, %key_cache_indices, %value_cache_indices, %attention_mask, %completion_counters : buffer, buffer, buffer, buffer, buffer + %positions_view = buffer.view %positions_noalias[%c0_offset] : buffer -> view<1xi32> + %key_indices_view = buffer.view %key_indices_noalias[%c0_offset] : buffer -> view<1xi64> + %value_indices_view = buffer.view %value_indices_noalias[%c0_offset] : buffer -> view<1xi64> + %mask_view = buffer.view %mask_noalias[%c0_offset] : buffer -> view<[%bounded_context_capacity]xi16> + %counters_view = buffer.view %counters_noalias[%c0_offset] : buffer -> view<[%counter_count]xi32> + %context_base_raw = view.load %positions_view[%c0] : view<1xi32> -> i32 + %context_base_i32 = scalar.assume %context_base_raw [range(%context_base_raw, 0, 32767)] : i32 + %context_base0 = index.cast %context_base_i32 : i32 to index + %context_base = index.assume %context_base0 [range(%context_base0, 0, 32767)] : index + %visible_count0 = index.add %context_base, %c1 : index + %visible_count, %launch_context_capacity = index.assume %visible_count0, %bounded_context_capacity [le(%visible_count0, %bounded_context_capacity)] : index, index + scf.for %key_base = [%c0 to %launch_context_capacity step %c256] { + %key0 = index.add %key_base, %workitem : index + %valid_key = index.cmp ult, %key0, %launch_context_capacity : index + %safe_key0 = scf.select %valid_key, %key0, %c0 : index + %key = index.assume %safe_key0 [lt(%safe_key0, %launch_context_capacity)] : index + %is_visible = index.cmp ult, %key, %visible_count : index + %mask_bits = scf.select %is_visible, %c0_i16, %negative_infinity_i16 : i16 + scf.if %valid_key { + view.store %mask_bits, %mask_view[%key] : i16, view<[%bounded_context_capacity]xi16> + } + } + %is_first = index.cmp eq, %workitem, %c0 : index + scf.if %is_first { + %cache_index = index.cast %context_base : index to i64 + view.store %cache_index, %key_indices_view[%c0] : i64, view<1xi64> + view.store %cache_index, %value_indices_view[%c0] : i64, view<1xi64> + } + %is_counter = index.cmp ult, %workitem, %counter_count : index + scf.if %is_counter { + %counter = index.assume %workitem [lt(%workitem, %counter_count)] : index + view.store %c0_i32, %counters_view[%counter] : i32, view<[%counter_count]xi32> + } + kernel.return +} + +check.case public @qwen_attention_context_base_capture_case { + %positions = check.generate.fill value(7) : tensor<1xi32> + %control = check.generate.fill value(-1) : tensor<1xi32> + %expected = check.generate.fill value(7) : tensor<1xi32> + kernel.launch @qwen_attention_context_base_capture[](%positions, %control) : [](tensor<1xi32>, tensor<1xi32>) + check.expect.equal actual(%control) expected(%expected) : tensor<1xi32> + check.return +} + +check.case public @qwen_attention_decode_state_initialize_case { + %context_capacity = check.literal value(768) : index + %counter_count = check.literal value(56) : index + %positions = check.generate.fill value(767) : tensor<1xi32> + %key_cache_indices = check.generate.fill value(-1) : tensor<1xi64> + %value_cache_indices = check.generate.fill value(-1) : tensor<1xi64> + %attention_mask = check.generate.fill value(1) : tensor<768xi16> + %counters = check.generate.fill value(-1) : tensor<56xi32> + %expected_cache_indices = check.generate.fill value(767) : tensor<1xi64> + %expected_mask = check.generate.fill value(0) : tensor<768xi16> + %expected_counters = check.generate.fill value(0) : tensor<56xi32> + kernel.launch @qwen_attention_decode_state_initialize[%context_capacity, %counter_count](%context_capacity, %counter_count, %positions, %key_cache_indices, %value_cache_indices, %attention_mask, %counters) : [index, index](index, index, tensor<1xi32>, tensor<1xi64>, tensor<1xi64>, tensor<768xi16>, tensor<56xi32>) + check.expect.equal actual(%key_cache_indices) expected(%expected_cache_indices) : tensor<1xi64> + check.expect.equal actual(%value_cache_indices) expected(%expected_cache_indices) : tensor<1xi64> + check.expect.equal actual(%attention_mask) expected(%expected_mask) : tensor<768xi16> + check.expect.equal actual(%counters) expected(%expected_counters) : tensor<56xi32> + check.return +} + +check.benchmark<@qwen_attention_context_base_capture_case> @qwen_attention_context_base_capture_benchmark + +check.benchmark<@qwen_attention_decode_state_initialize_case> @qwen_attention_decode_state_initialize_benchmark diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_owned/token_embedding_bringup_workaround.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_owned/token_embedding_bringup_workaround.loom new file mode 100644 index 000000000000..8977fee1060c --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_owned/token_embedding_bringup_workaround.loom @@ -0,0 +1,224 @@ +// Copyright 2026 The IREE Authors +// +// Licensed under the Apache License v2.0 with LLVM Exceptions. +// See https://llvm.org/LICENSE.txt for license information. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +// Temporary, non-sanctioned one-function kernel for Qwen bring-up. This is not +// an embedding-kernel framework, a generator, or a second Loom authoring path. +// It gathers token rows directly from the model's unmodified GGUF Q4_K payload +// and decodes them to the owned F32 hidden-state layout. Delete this file when +// the Qwen kernel corpus provides the canonical token-embedding producer. +// +// The owned model row contract is: +// Q4_K: [vocabulary row][hidden_size / 256 blocks][144 bytes] +// -> [hidden_size x f32] +amdgpu.target @qwen_token_embedding_gfx11_wave64 {subgroup_size = 64} + +kernel.def target(@qwen_token_embedding_gfx11_wave64) export("qwen_token_embedding_q4k_bringup_workaround") @qwen_token_embedding_q4k_bringup_workaround(%token_count: index, %vocabulary_count: index, %hidden_size: index) { + %one = index.constant 1 : index + %onethousandtwentyfour = index.constant 1024 : index + %workgroups_per_token = index.div %hidden_size, %onethousandtwentyfour : index + %workgroup_size = index.constant 256 : index + kernel.launch.config workgroups(%workgroups_per_token, %token_count, %one) workgroup_size(%workgroup_size, %one, %one) : index +} launch(%token_count: index, %vocabulary_count: index, %hidden_size: index, %token_ids: buffer, %weight: buffer, %output: buffer) { + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 2048)] : index + %bounded_vocabulary_count = index.assume %vocabulary_count [range(%vocabulary_count, 1, 262144)] : index + %bounded_hidden_size = index.assume %hidden_size [range(%hidden_size, 2048, 3072), mul(%hidden_size, 1024)] : index + %workgroup = kernel.workgroup.id : index + %token0 = kernel.workgroup.id : index + %workitem = kernel.workitem.id : index + %zero = index.constant 0 : index + %two = index.constant 2 : index + %three = index.constant 3 : index + %four = index.constant 4 : index + %eight = index.constant 8 : index + %sixtyfour = index.constant 64 : index + %twofiftysix = index.constant 256 : index + %packets_per_token0 = index.div %bounded_hidden_size, %four : index + %packets_per_token = index.assume %packets_per_token0 [range(%packets_per_token0, 512, 768)] : index + %q4_block_count0 = index.div %bounded_hidden_size, %twofiftysix : index + %q4_block_count = index.assume %q4_block_count0 [range(%q4_block_count0, 8, 12)] : index + %zero_offset = index.constant 0 : offset + %block_bytes = index.constant 144 : offset + %scale_offset = index.constant 4 : offset + %code_offset = index.constant 16 : offset + %row_bytes = index.scale %q4_block_count, %block_bytes : index, offset -> offset + %two_i32 = scalar.constant 2 : i32 + %four_i32 = scalar.constant 4 : i32 + %fifteen = vector.constant 15 : vector<1xi32> + %fortyeight = vector.constant 48 : vector<1xi32> + %q4_mask = vector.constant 252645135 : vector<1xi32> + %token, %launch_token_count = index.assume %token0, %bounded_token_count [lt(%token0, %bounded_token_count)] : index, index + %row_packet0 = index.madd %workgroup, %twofiftysix, %workitem : index + %row_packet, %launch_packets_per_token = index.assume %row_packet0, %packets_per_token [range(%row_packet0, 0, 767), lt(%row_packet0, %packets_per_token)] : index, index + %q4_block0 = index.div %row_packet, %sixtyfour : index + %q4_block, %weight_q4_block_count = index.assume %q4_block0, %q4_block_count [range(%q4_block0, 0, 11), lt(%q4_block0, %q4_block_count)] : index, index + %block_packet0 = index.rem %row_packet, %sixtyfour : index + %block_packet = index.assume %block_packet0 [range(%block_packet0, 0, 63)] : index + %q4_group0 = index.div %block_packet, %eight : index + %q4_group = index.assume %q4_group0 [range(%q4_group0, 0, 7)] : index + %load_packet0 = index.rem %block_packet, %eight : index + %load_packet = index.assume %load_packet0 [range(%load_packet0, 0, 7)] : index + %token_ids_noalias, %weight_noalias, %output_noalias = buffer.assume.noalias %token_ids, %weight, %output : buffer, buffer, buffer + %token_ids_view = buffer.view %token_ids_noalias[%zero_offset] : buffer -> view<[%launch_token_count]xi32> + %output_view = buffer.view %output_noalias[%zero_offset] : buffer -> view<[%launch_token_count]x[%bounded_hidden_size]xf32> + // Token IDs are validated at the request boundary. This trusted consumer + // carries that fact into row addressing instead of adding a fallback row. + %token_id_raw = view.load %token_ids_view[%token] : view<[%launch_token_count]xi32> -> i32 + %token_id_i32 = scalar.assume %token_id_raw [range(%token_id_raw, 0, 262143)] : i32 + %token_id0 = index.cast %token_id_i32 : i32 to index + %token_id, %weight_vocabulary_count = index.assume %token_id0, %bounded_vocabulary_count [lt(%token_id0, %bounded_vocabulary_count)] : index, index + %row_byte_base = index.scale %token_id, %row_bytes : index, offset -> offset + %block_byte_add = index.scale %q4_block, %block_bytes : index, offset -> offset + %block_byte_base = index.add %row_byte_base, %block_byte_add : offset + %scale_byte_base = index.add %block_byte_base, %scale_offset : offset + %code_byte_base = index.add %block_byte_base, %code_offset : offset + %dm_view = buffer.view %weight_noalias[%block_byte_base] : buffer -> view<2xf16> + %scale_view = buffer.view %weight_noalias[%scale_byte_base] : buffer -> view<3xi32> + %code_view = buffer.view %weight_noalias[%code_byte_base] : buffer -> view<32xi32> + %dm = vector.load %dm_view[%zero] : view<2xf16> -> vector<2xf16> + %scales = vector.load %scale_view[%zero] : view<3xi32> -> vector<3xi32> + %d_f16 = vector.extract %dm[0] : vector<2xf16> -> f16 + %dmin_f16 = vector.extract %dm[1] : vector<2xf16> -> f16 + %d = scalar.extf %d_f16 : f16 to f32 + %dmin = scalar.extf %dmin_f16 : f16 to f32 + %q_page0 = index.div %q4_group, %two : index + %q_page = index.mul %q_page0, %eight : index + %q_word_index0 = index.add %q_page, %load_packet : index + %q_word_index = index.assume %q_word_index0 [range(%q_word_index0, 0, 31)] : index + %is_low = index.cmp ult, %q4_group, %four : index + %scale_lane = index.rem %q4_group, %four : index + %scale_shift_index = index.mul %scale_lane, %eight : index + %scale_shift_i32 = index.cast %scale_shift_index : index to i32 + %scale_shift = vector.splat %scale_shift_i32 : vector<1xi32> + %scale0_i32 = vector.extract %scales[0] : vector<3xi32> -> i32 + %scale1_i32 = vector.extract %scales[1] : vector<3xi32> -> i32 + %scale2_i32 = vector.extract %scales[2] : vector<3xi32> -> i32 + %scale0 = vector.splat %scale0_i32 : vector<1xi32> + %scale1 = vector.splat %scale1_i32 : vector<1xi32> + %scale2 = vector.splat %scale2_i32 : vector<1xi32> + %high_shift_i32 = scalar.addi %scale_shift_i32, %two_i32 : i32 + %minimum_shift_i32 = scalar.addi %scale_shift_i32, %four_i32 : i32 + %selected_scale_source = scf.select %is_low, %scale0, %scale2 : vector<1xi32> + %selected_minimum_source = scf.select %is_low, %scale1, %scale2 : vector<1xi32> + %selected_scale_high_shift_i32 = scf.select %is_low, %scale_shift_i32, %high_shift_i32 : i32 + %selected_minimum_low_shift_i32 = scf.select %is_low, %scale_shift_i32, %minimum_shift_i32 : i32 + %selected_scale_high_shift = vector.splat %selected_scale_high_shift_i32 : vector<1xi32> + %selected_minimum_low_shift = vector.splat %selected_minimum_low_shift_i32 : vector<1xi32> + %scale_low0 = vector.shrui %selected_scale_source, %scale_shift : vector<1xi32> + %scale_low = vector.andi %scale_low0, %fifteen : vector<1xi32> + %scale_high0 = vector.shrui %scale0, %selected_scale_high_shift : vector<1xi32> + %scale_high = vector.andi %scale_high0, %fortyeight : vector<1xi32> + %scale = vector.ori %scale_low, %scale_high : vector<1xi32> + %minimum_low0 = vector.shrui %selected_minimum_source, %selected_minimum_low_shift : vector<1xi32> + %minimum_low = vector.andi %minimum_low0, %fifteen : vector<1xi32> + %minimum_high0 = vector.shrui %scale1, %selected_scale_high_shift : vector<1xi32> + %minimum_high = vector.andi %minimum_high0, %fortyeight : vector<1xi32> + %minimum = vector.ori %minimum_low, %minimum_high : vector<1xi32> + %scale_f32 = vector.uitofp %scale : vector<1xi32> to vector<1xf32> + %minimum_f32 = vector.uitofp %minimum : vector<1xi32> to vector<1xf32> + %d_vector1 = vector.splat %d : vector<1xf32> + %dmin_vector1 = vector.splat %dmin : vector<1xf32> + %d_scale_vector1 = vector.mulf %d_vector1, %scale_f32 : vector<1xf32> + %minimum_scale_vector1 = vector.mulf %dmin_vector1, %minimum_f32 : vector<1xf32> + %d_scale = vector.extract %d_scale_vector1[0] : vector<1xf32> -> f32 + %minimum_scale = vector.extract %minimum_scale_vector1[0] : vector<1xf32> -> f32 + %q_word = vector.load %code_view[%q_word_index] : view<32xi32> -> vector<1xi32> + %q_half = index.rem %q4_group, %two : index + %q_shift_index = index.mul %q_half, %four : index + %q_shift_i32 = index.cast %q_shift_index : index to i32 + %q_shift = vector.splat %q_shift_i32 : vector<1xi32> + %shifted_q = vector.shrui %q_word, %q_shift : vector<1xi32> + %masked_q = vector.andi %shifted_q, %q4_mask : vector<1xi32> + %q_i8 = vector.bitcast %masked_q : vector<1xi32> to vector<4xi8> + %q_f32 = vector.uitofp %q_i8 : vector<4xi8> to vector<4xf32> + %negative_minimum_scale = scalar.negf %minimum_scale : f32 + %q0 = vector.extract %q_f32[0] : vector<4xf32> -> f32 + %q1 = vector.extract %q_f32[1] : vector<4xf32> -> f32 + %q2 = vector.extract %q_f32[2] : vector<4xf32> -> f32 + %q3 = vector.extract %q_f32[3] : vector<4xf32> -> f32 + %value0 = scalar.fmaf %q0, %d_scale, %negative_minimum_scale : f32 + %value1 = scalar.fmaf %q1, %d_scale, %negative_minimum_scale : f32 + %value2 = scalar.fmaf %q2, %d_scale, %negative_minimum_scale : f32 + %value3 = scalar.fmaf %q3, %d_scale, %negative_minimum_scale : f32 + %values = vector.from_elements %value0, %value1, %value2, %value3 : vector<4xf32> + %output_channel0 = index.mul %row_packet, %four : index + %output_channel_end = index.sub %bounded_hidden_size, %three : index + %output_channel = index.assume %output_channel0 [range(%output_channel0, 0, 3068), lt(%output_channel0, %output_channel_end)] : index + vector.store %values, %output_view[%token, %output_channel] : vector<4xf32>, view<[%launch_token_count]x[%bounded_hidden_size]xf32> + kernel.return +} + +// Uniform 0x55 Q4_K bytes encode d=dmin=85.3125, scale=minimum=21, +// and q=5 in every group, producing 85.3125 * 21 * (5 - 1) = 7166.25. +check.case public @qwen_token_embedding_q4k_decode_case { + %token_count = check.literal value(1) : index + %vocabulary_count = check.literal value(1) : index + %hidden_size = check.literal value(2048) : index + %token_ids = check.generate.fill value(0) : tensor<1xi32> + %weight = check.generate.fill value(85) : tensor<1x8x144xi8> + %output = check.generate.fill value(0.0) : tensor<1x2048xf32> + %expected = check.generate.fill value(7166.25) : tensor<1x2048xf32> + func.call @qwen_token_embedding_q4k_bringup_workaround(%token_count, %vocabulary_count, %hidden_size, %token_ids, %weight, %output) : (index, index, index, tensor<1xi32>, tensor<1x8x144xi8>, tensor<1x2048xf32>) + check.expect.close actual(%output) expected(%expected) atol(0.0) rtol(0.0) nan(same) : tensor<1x2048xf32> + check.return +} + +// Reverse-order IDs exercise the first, interior, and final legal row while +// remaining suitable for access-sanitized execution. +check.case public @qwen_token_embedding_q4k_row_access_case { + %token_count = check.literal value(3) : index + %vocabulary_count = check.literal value(3) : index + %hidden_size = check.literal value(2048) : index + %token_ids = check.generate.iota offset(2) step(-1) : tensor<3xi32> + %weight = check.generate.fill value(0) : tensor<3x8x144xi8> + %output = check.generate.fill value(1.0) : tensor<3x2048xf32> + %expected = check.generate.fill value(0.0) : tensor<3x2048xf32> + func.call @qwen_token_embedding_q4k_bringup_workaround(%token_count, %vocabulary_count, %hidden_size, %token_ids, %weight, %output) : (index, index, index, tensor<3xi32>, tensor<3x8x144xi8>, tensor<3x2048xf32>) + check.expect.close actual(%output) expected(%expected) atol(0.0) rtol(0.0) nan(same) : tensor<3x2048xf32> + check.return +} + +check.case public @qwen_token_embedding_q4k_benchmark_case { + %token_count = check.param.choice values([1, 32, 128, 512]) name("token_count") : index + %vocabulary_count = check.literal value(151936) : index + %hidden_size = check.literal value(2048) : index + %token_ids = check.generate.iota offset(151424) step(1) period(512) : tensor<[%token_count]xi32> + %weight = check.generate.fill value(0) : tensor<151936x8x144xi8> + %output = check.generate.fill value(1.0) : tensor<[%token_count]x2048xf32> + %expected = check.generate.fill value(0.0) : tensor<[%token_count]x2048xf32> + func.call @qwen_token_embedding_q4k_bringup_workaround(%token_count, %vocabulary_count, %hidden_size, %token_ids, %weight, %output) : (index, index, index, tensor<[%token_count]xi32>, tensor<151936x8x144xi8>, tensor<[%token_count]x2048xf32>) + check.expect.close actual(%output) expected(%expected) atol(0.0) rtol(0.0) nan(same) : tensor<[%token_count]x2048xf32> + check.return +} + +// Exercises the final legal checkpoint row at the exact Qwen3.5 geometry. +check.case public @qwen35_token_embedding_q4k_row_access_case { + %token_count = check.literal value(1) : index + %vocabulary_count = check.literal value(248320) : index + %hidden_size = check.literal value(3072) : index + %token_ids = check.generate.fill value(248319) : tensor<1xi32> + %weight = check.generate.fill value(0) : tensor<248320x12x144xi8> + %output = check.generate.fill value(1.0) : tensor<1x3072xf32> + %expected = check.generate.fill value(0.0) : tensor<1x3072xf32> + func.call @qwen_token_embedding_q4k_bringup_workaround(%token_count, %vocabulary_count, %hidden_size, %token_ids, %weight, %output) : (index, index, index, tensor<1xi32>, tensor<248320x12x144xi8>, tensor<1x3072xf32>) + check.expect.close actual(%output) expected(%expected) atol(0.0) rtol(0.0) nan(same) : tensor<1x3072xf32> + check.return +} + +check.benchmark<@qwen_token_embedding_q4k_decode_case> @qwen_token_embedding_q4k_decode + +// Production decode specialization. Unlike the tiny differential case above, +// this carries the locked Q4_K_M vocabulary geometry into LoomC. +check.benchmark<@qwen_token_embedding_q4k_benchmark_case> @qwen_token_embedding_q4k_model_decode {token_count = 1} + +check.benchmark<@qwen_token_embedding_q4k_row_access_case> @qwen_token_embedding_q4k_row_access + +check.benchmark<@qwen35_token_embedding_q4k_row_access_case> @qwen35_token_embedding_q4k_decode + +check.benchmark<@qwen_token_embedding_q4k_benchmark_case> @qwen_token_embedding_q4k_prefill_32 {token_count = 32} + +check.benchmark<@qwen_token_embedding_q4k_benchmark_case> @qwen_token_embedding_q4k_prefill_128 {token_count = 128} + +check.benchmark<@qwen_token_embedding_q4k_benchmark_case> @qwen_token_embedding_q4k_prefill_512 {token_count = 512} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_owned/token_embedding_q4k.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_owned/token_embedding_q4k.loom new file mode 100644 index 000000000000..1bce11c86515 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_owned/token_embedding_q4k.loom @@ -0,0 +1,191 @@ +// Copyright 2026 The IREE Authors +// +// Licensed under the Apache License v2.0 with LLVM Exceptions. +// See https://llvm.org/LICENSE.txt for license information. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +// Gathers token rows directly from the model's unmodified GGUF Q4_K payload +// and decodes them to the owned F32 hidden-state layout. The fixed model row +// contract is: +// Q4_K: [vocabulary row][8 blocks][144 bytes] -> [2048xf32] +amdgpu.target @qwen_token_embedding_gfx11_wave64 {subgroup_size = 64} + +kernel.def target(@qwen_token_embedding_gfx11_wave64) export("qwen_token_embedding_q4k") @qwen_token_embedding_q4k(%token_count: index, %vocabulary_count: index) { + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %workgroup_count = index.mul %token_count, %c2 : index + %workgroup_size = index.constant 256 : index + kernel.launch.config workgroups(%workgroup_count, %c1, %c1) workgroup_size(%workgroup_size, %c1, %c1) : index +} launch(%token_count: index, %vocabulary_count: index, %token_ids: buffer, %weight: buffer, %output: buffer) where [range(%token_count, 1, 2048)] { + %bounded_vocabulary_count = index.assume %vocabulary_count [range(%vocabulary_count, 1, 262144)] : index + %workgroup = kernel.workgroup.id : index + %workitem = kernel.workitem.id : index + %c0 = index.constant 0 : index + %c2 = index.constant 2 : index + %c4 = index.constant 4 : index + %c8 = index.constant 8 : index + %c64 = index.constant 64 : index + %c256 = index.constant 256 : index + %packets_per_token = index.constant 512 : index + %c0_offset = index.constant 0 : offset + %block_bytes = index.constant 144 : offset + %scale_offset = index.constant 4 : offset + %code_offset = index.constant 16 : offset + %row_bytes = index.constant 1152 : offset + %c2_i32 = scalar.constant 2 : i32 + %c4_i32 = scalar.constant 4 : i32 + %c15 = vector.constant 15 : vector<1xi32> + %c48 = vector.constant 48 : vector<1xi32> + %q4_mask = vector.constant 252645135 : vector<1xi32> + %packet0 = index.madd %workgroup, %c256, %workitem : index + %packet_count = index.mul %token_count, %packets_per_token : index + %packet, %launch_packet_count = index.assume %packet0, %packet_count [lt(%packet0, %packet_count)] : index, index + %token0 = index.div %packet, %packets_per_token : index + %token, %launch_token_count = index.assume %token0, %token_count [lt(%token0, %token_count)] : index, index + %row_packet0 = index.rem %packet, %packets_per_token : index + %row_packet = index.assume %row_packet0 [range(%row_packet0, 0, 511)] : index + %q4_block0 = index.div %row_packet, %c64 : index + %q4_block = index.assume %q4_block0 [range(%q4_block0, 0, 7)] : index + %block_packet0 = index.rem %row_packet, %c64 : index + %block_packet = index.assume %block_packet0 [range(%block_packet0, 0, 63)] : index + %q4_group0 = index.div %block_packet, %c8 : index + %q4_group = index.assume %q4_group0 [range(%q4_group0, 0, 7)] : index + %load_packet0 = index.rem %block_packet, %c8 : index + %load_packet = index.assume %load_packet0 [range(%load_packet0, 0, 7)] : index + %token_ids_noalias, %weight_noalias, %output_noalias = buffer.assume.noalias %token_ids, %weight, %output : buffer, buffer, buffer + %token_ids_view = buffer.view %token_ids_noalias[%c0_offset] : buffer -> view<[%launch_token_count]xi32> + %output_view = buffer.view %output_noalias[%c0_offset] : buffer -> view<[%launch_token_count]x2048xf32> + // Token IDs are validated at the request boundary. This trusted consumer + // carries that fact into row addressing instead of adding a fallback row. + %token_id_raw = view.load %token_ids_view[%token] : view<[%launch_token_count]xi32> -> i32 + %token_id_i32 = scalar.assume %token_id_raw [range(%token_id_raw, 0, 262143)] : i32 + %token_id0 = index.cast %token_id_i32 : i32 to index + %token_id, %weight_vocabulary_count = index.assume %token_id0, %bounded_vocabulary_count [lt(%token_id0, %bounded_vocabulary_count)] : index, index + %row_byte_base = index.scale %token_id, %row_bytes : index, offset -> offset + %block_byte_add = index.scale %q4_block, %block_bytes : index, offset -> offset + %block_byte_base = index.add %row_byte_base, %block_byte_add : offset + %scale_byte_base = index.add %block_byte_base, %scale_offset : offset + %code_byte_base = index.add %block_byte_base, %code_offset : offset + %dm_view = buffer.view %weight_noalias[%block_byte_base] : buffer -> view<2xf16> + %scale_view = buffer.view %weight_noalias[%scale_byte_base] : buffer -> view<3xi32> + %code_view = buffer.view %weight_noalias[%code_byte_base] : buffer -> view<32xi32> + %dm = vector.load %dm_view[%c0] : view<2xf16> -> vector<2xf16> + %scales = vector.load %scale_view[%c0] : view<3xi32> -> vector<3xi32> + %d_f16 = vector.extract %dm[0] : vector<2xf16> -> f16 + %dmin_f16 = vector.extract %dm[1] : vector<2xf16> -> f16 + %d = scalar.extf %d_f16 : f16 to f32 + %dmin = scalar.extf %dmin_f16 : f16 to f32 + %q_page0 = index.div %q4_group, %c2 : index + %q_page = index.mul %q_page0, %c8 : index + %q_word_index0 = index.add %q_page, %load_packet : index + %q_word_index = index.assume %q_word_index0 [range(%q_word_index0, 0, 31)] : index + %is_low = index.cmp ult, %q4_group, %c4 : index + %scale_lane = index.rem %q4_group, %c4 : index + %scale_shift_index = index.mul %scale_lane, %c8 : index + %scale_shift_i32 = index.cast %scale_shift_index : index to i32 + %scale_shift = vector.splat %scale_shift_i32 : vector<1xi32> + %scale0_i32 = vector.extract %scales[0] : vector<3xi32> -> i32 + %scale1_i32 = vector.extract %scales[1] : vector<3xi32> -> i32 + %scale2_i32 = vector.extract %scales[2] : vector<3xi32> -> i32 + %scale0 = vector.splat %scale0_i32 : vector<1xi32> + %scale1 = vector.splat %scale1_i32 : vector<1xi32> + %scale2 = vector.splat %scale2_i32 : vector<1xi32> + %high_shift_i32 = scalar.addi %scale_shift_i32, %c2_i32 : i32 + %minimum_shift_i32 = scalar.addi %scale_shift_i32, %c4_i32 : i32 + %selected_scale_source = scf.select %is_low, %scale0, %scale2 : vector<1xi32> + %selected_minimum_source = scf.select %is_low, %scale1, %scale2 : vector<1xi32> + %selected_scale_high_shift_i32 = scf.select %is_low, %scale_shift_i32, %high_shift_i32 : i32 + %selected_minimum_low_shift_i32 = scf.select %is_low, %scale_shift_i32, %minimum_shift_i32 : i32 + %selected_scale_high_shift = vector.splat %selected_scale_high_shift_i32 : vector<1xi32> + %selected_minimum_low_shift = vector.splat %selected_minimum_low_shift_i32 : vector<1xi32> + %scale_low0 = vector.shrui %selected_scale_source, %scale_shift : vector<1xi32> + %scale_low = vector.andi %scale_low0, %c15 : vector<1xi32> + %scale_high0 = vector.shrui %scale0, %selected_scale_high_shift : vector<1xi32> + %scale_high = vector.andi %scale_high0, %c48 : vector<1xi32> + %scale = vector.ori %scale_low, %scale_high : vector<1xi32> + %minimum_low0 = vector.shrui %selected_minimum_source, %selected_minimum_low_shift : vector<1xi32> + %minimum_low = vector.andi %minimum_low0, %c15 : vector<1xi32> + %minimum_high0 = vector.shrui %scale1, %selected_scale_high_shift : vector<1xi32> + %minimum_high = vector.andi %minimum_high0, %c48 : vector<1xi32> + %minimum = vector.ori %minimum_low, %minimum_high : vector<1xi32> + %scale_f32 = vector.uitofp %scale : vector<1xi32> to vector<1xf32> + %minimum_f32 = vector.uitofp %minimum : vector<1xi32> to vector<1xf32> + %d_vector1 = vector.splat %d : vector<1xf32> + %dmin_vector1 = vector.splat %dmin : vector<1xf32> + %d_scale_vector1 = vector.mulf %d_vector1, %scale_f32 : vector<1xf32> + %minimum_scale_vector1 = vector.mulf %dmin_vector1, %minimum_f32 : vector<1xf32> + %d_scale = vector.extract %d_scale_vector1[0] : vector<1xf32> -> f32 + %minimum_scale = vector.extract %minimum_scale_vector1[0] : vector<1xf32> -> f32 + %q_word = vector.load %code_view[%q_word_index] : view<32xi32> -> vector<1xi32> + %q_half = index.rem %q4_group, %c2 : index + %q_shift_index = index.mul %q_half, %c4 : index + %q_shift_i32 = index.cast %q_shift_index : index to i32 + %q_shift = vector.splat %q_shift_i32 : vector<1xi32> + %shifted_q = vector.shrui %q_word, %q_shift : vector<1xi32> + %masked_q = vector.andi %shifted_q, %q4_mask : vector<1xi32> + %q_i8 = vector.bitcast %masked_q : vector<1xi32> to vector<4xi8> + %q_f32 = vector.uitofp %q_i8 : vector<4xi8> to vector<4xf32> + %negative_minimum_scale = scalar.negf %minimum_scale : f32 + %q0 = vector.extract %q_f32[0] : vector<4xf32> -> f32 + %q1 = vector.extract %q_f32[1] : vector<4xf32> -> f32 + %q2 = vector.extract %q_f32[2] : vector<4xf32> -> f32 + %q3 = vector.extract %q_f32[3] : vector<4xf32> -> f32 + %value0 = scalar.fmaf %q0, %d_scale, %negative_minimum_scale : f32 + %value1 = scalar.fmaf %q1, %d_scale, %negative_minimum_scale : f32 + %value2 = scalar.fmaf %q2, %d_scale, %negative_minimum_scale : f32 + %value3 = scalar.fmaf %q3, %d_scale, %negative_minimum_scale : f32 + %values = vector.from_elements %value0, %value1, %value2, %value3 : vector<4xf32> + %output_channel = index.mul %row_packet, %c4 : index + vector.store %values, %output_view[%token, %output_channel] : vector<4xf32>, view<[%launch_token_count]x2048xf32> + kernel.return +} + +// Uniform 0x55 Q4_K bytes encode d=dmin=85.3125, scale=minimum=21, +// and q=5 in every group, producing 85.3125 * 21 * (5 - 1) = 7166.25. +check.case public @qwen_token_embedding_q4k_decode_case { + %token_count = check.literal value(1) : index + %vocabulary_count = check.literal value(1) : index + %token_ids = check.generate.fill value(0) : tensor<1xi32> + %weight = check.generate.fill value(85) : tensor<1x8x144xi8> + %output = check.generate.fill value(0.0) : tensor<1x2048xf32> + %expected = check.generate.fill value(7166.25) : tensor<1x2048xf32> + kernel.launch @qwen_token_embedding_q4k[%token_count, %vocabulary_count](%token_count, %vocabulary_count, %token_ids, %weight, %output) : [index, index](index, index, tensor<1xi32>, tensor<1x8x144xi8>, tensor<1x2048xf32>) + check.expect.close actual(%output) expected(%expected) atol(0.0) rtol(0.0) nan(same) : tensor<1x2048xf32> + check.return +} + +// Reverse-order IDs exercise the first, interior, and final legal row while +// remaining suitable for access-sanitized execution. +check.case public @qwen_token_embedding_q4k_row_access_case { + %token_count = check.literal value(3) : index + %vocabulary_count = check.literal value(3) : index + %token_ids = check.generate.iota offset(2) step(-1) : tensor<3xi32> + %weight = check.generate.fill value(0) : tensor<3x8x144xi8> + %output = check.generate.fill value(1.0) : tensor<3x2048xf32> + %expected = check.generate.fill value(0.0) : tensor<3x2048xf32> + kernel.launch @qwen_token_embedding_q4k[%token_count, %vocabulary_count](%token_count, %vocabulary_count, %token_ids, %weight, %output) : [index, index](index, index, tensor<3xi32>, tensor<3x8x144xi8>, tensor<3x2048xf32>) + check.expect.close actual(%output) expected(%expected) atol(0.0) rtol(0.0) nan(same) : tensor<3x2048xf32> + check.return +} + +check.case public @qwen_token_embedding_q4k_benchmark_case { + %token_count = check.param.choice values([32, 128, 512]) name("token_count") : index + %vocabulary_count = check.literal value(151936) : index + %token_ids = check.generate.iota offset(151424) step(1) period(512) : tensor<[%token_count]xi32> + %weight = check.generate.fill value(0) : tensor<151936x8x144xi8> + %output = check.generate.fill value(1.0) : tensor<[%token_count]x2048xf32> + %expected = check.generate.fill value(0.0) : tensor<[%token_count]x2048xf32> + kernel.launch @qwen_token_embedding_q4k[%token_count, %vocabulary_count](%token_count, %vocabulary_count, %token_ids, %weight, %output) : [index, index](index, index, tensor<[%token_count]xi32>, tensor<151936x8x144xi8>, tensor<[%token_count]x2048xf32>) + check.expect.close actual(%output) expected(%expected) atol(0.0) rtol(0.0) nan(same) : tensor<[%token_count]x2048xf32> + check.return +} + +check.benchmark<@qwen_token_embedding_q4k_decode_case> @qwen_token_embedding_q4k_decode + +check.benchmark<@qwen_token_embedding_q4k_row_access_case> @qwen_token_embedding_q4k_row_access + +check.benchmark<@qwen_token_embedding_q4k_benchmark_case> @qwen_token_embedding_q4k_prefill_32 {token_count = 32} + +check.benchmark<@qwen_token_embedding_q4k_benchmark_case> @qwen_token_embedding_q4k_prefill_128 {token_count = 128} + +check.benchmark<@qwen_token_embedding_q4k_benchmark_case> @qwen_token_embedding_q4k_prefill_512 {token_count = 512} diff --git a/ggml/src/ggml-hrx/loom-jit.cpp b/ggml/src/ggml-hrx/loom-jit.cpp new file mode 100644 index 000000000000..fda979d3772c --- /dev/null +++ b/ggml/src/ggml-hrx/loom-jit.cpp @@ -0,0 +1,1048 @@ +// Copyright 2026 The HRX Authors +// SPDX-License-Identifier: Apache-2.0 + +#include "loom-jit.h" + +#include "loomc/launch_config.h" +#include "loomc/loomc.h" +#include "loomc/sanitizer.h" +#include "loomc/target/amdgpu.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +ggml_hrx_loom_jit_compile_result::~ggml_hrx_loom_jit_compile_result() { + reset(); +} + +ggml_hrx_loom_jit_compile_result::ggml_hrx_loom_jit_compile_result(ggml_hrx_loom_jit_compile_result && other) noexcept { + *this = std::move(other); +} + +ggml_hrx_loom_jit_compile_result & ggml_hrx_loom_jit_compile_result::operator=( + ggml_hrx_loom_jit_compile_result && other) noexcept { + if (this == &other) { + return *this; + } + reset(); + hsaco_data = std::exchange(other.hsaco_data, nullptr); + hsaco_size = std::exchange(other.hsaco_size, 0); + manifest_json = std::exchange(other.manifest_json, nullptr); + manifest_json_size = std::exchange(other.manifest_json_size, 0); + compile_report_json = std::exchange(other.compile_report_json, nullptr); + compile_report_json_size = std::exchange(other.compile_report_json_size, 0); + final_module_text = std::exchange(other.final_module_text, nullptr); + final_module_text_size = std::exchange(other.final_module_text_size, 0); + launch_config = std::exchange(other.launch_config, {}); + return *this; +} + +void ggml_hrx_loom_jit_compile_result::reset() { + hrx_host_allocator_t allocator = hrx_host_allocator_system(); + hrx_host_allocator_free(allocator, hsaco_data); + hrx_host_allocator_free(allocator, manifest_json); + hrx_host_allocator_free(allocator, compile_report_json); + hrx_host_allocator_free(allocator, final_module_text); + hsaco_data = nullptr; + hsaco_size = 0; + manifest_json = nullptr; + manifest_json_size = 0; + compile_report_json = nullptr; + compile_report_json_size = 0; + final_module_text = nullptr; + final_module_text_size = 0; + launch_config = {}; +} + +struct ggml_hrx_loom_jit_amdgpu { + loomc_target_environment_t * target_environment = nullptr; + loomc_context_t * context = nullptr; + loomc_target_profile_t * target_profile = nullptr; + loomc_compiler_t * compiler = nullptr; + loomc_pass_program_t * pass_program = nullptr; + loomc_amdgpu_runtime_global_flags_t runtime_globals = LOOMC_AMDGPU_RUNTIME_GLOBAL_NONE; +}; + +namespace { + +// LoomC currently accepts concrete workload values for launch evaluation but +// not for compilation. Until that becomes one operation in LoomC, specialize +// the pinned textual kernel source before parsing so body optimization sees +// the same exact facts as launch evaluation. The public kernel ABI is retained: +// scalar arguments remain present but their uses in both kernel regions are +// replaced by backend-authored constants. This is deliberately local to the +// HRX backend and its emitted module text is always available for inspection. +static bool ggml_hrx_loom_specialize_workload_text(const std::string & input, + const char * root_symbol, + const int64_t * workload_arguments, + size_t workload_argument_count, + std::string & output, + std::string & error) { + output = input; + if (workload_argument_count == 0) { + return true; + } + std::string symbol = root_symbol ? root_symbol : ""; + if (!symbol.empty() && symbol.front() == '@') { + symbol.erase(symbol.begin()); + } + const std::string marker = "@" + symbol + "("; + const size_t symbol_position = output.find(marker); + if (symbol_position == std::string::npos) { + error = "workload specialization cannot find root " + marker; + return false; + } + const size_t parameter_begin = symbol_position + marker.size(); + const size_t parameter_end = output.find(')', parameter_begin); + if (parameter_end == std::string::npos) { + error = "workload specialization cannot parse root parameters"; + return false; + } + const std::string parameters = output.substr(parameter_begin, parameter_end - parameter_begin); + std::vector names; + size_t cursor = 0; + while (names.size() < workload_argument_count) { + const size_t percent = parameters.find('%', cursor); + if (percent == std::string::npos) { + break; + } + size_t name_end = percent + 1; + while (name_end < parameters.size() && + (std::isalnum(static_cast(parameters[name_end])) || parameters[name_end] == '_')) { + ++name_end; + } + const size_t colon = parameters.find(':', name_end); + if (colon == std::string::npos) { + break; + } + const size_t type_begin = parameters.find_first_not_of(" \t", colon + 1); + if (type_begin != std::string::npos && parameters.compare(type_begin, 5, "index") == 0) { + names.push_back(parameters.substr(percent + 1, name_end - percent - 1)); + } + cursor = name_end; + } + if (names.size() != workload_argument_count) { + error = "workload specialization argument count does not match root index parameters"; + return false; + } + + auto matching_brace = [&](size_t open) -> size_t { + size_t depth = 0; + for (size_t i = open; i < output.size(); ++i) { + if (output[i] == '{') { + ++depth; + } + if (output[i] == '}' && --depth == 0) { + return i; + } + } + return std::string::npos; + }; + std::vector> regions; + const size_t config_open = output.find('{', parameter_end); + const size_t config_close = config_open == std::string::npos ? std::string::npos : matching_brace(config_open); + if (config_close == std::string::npos) { + error = "workload specialization cannot find kernel config region"; + return false; + } + regions.push_back({ config_open, config_close }); + const size_t launch = output.find("launch(", config_close + 1); + const size_t launch_open = launch == std::string::npos ? std::string::npos : output.find('{', launch); + const size_t launch_close = launch_open == std::string::npos ? std::string::npos : matching_brace(launch_open); + if (launch_close == std::string::npos) { + error = "workload specialization cannot find kernel launch region"; + return false; + } + regions.push_back({ launch_open, launch_close }); + + for (auto region = regions.rbegin(); region != regions.rend(); ++region) { + std::string body = output.substr(region->first + 1, region->second - region->first - 1); + std::string prefix; + for (size_t i = 0; i < names.size(); ++i) { + const std::string original = "%" + names[i]; + const std::string specialized = "%ggml_hrx_specialized_" + names[i]; + size_t use = 0; + while ((use = body.find(original, use)) != std::string::npos) { + const size_t after = use + original.size(); + if (after == body.size() || + (!std::isalnum(static_cast(body[after])) && body[after] != '_')) { + body.replace(use, original.size(), specialized); + use += specialized.size(); + } else { + use = after; + } + } + prefix += "\n " + specialized + " = index.constant " + std::to_string(workload_arguments[i]) + " : index"; + } + output.replace(region->first + 1, region->second - region->first - 1, prefix + body); + } + return true; +} + +template class LoomHandle { + public: + LoomHandle() = default; + LoomHandle(const LoomHandle &) = delete; + LoomHandle & operator=(const LoomHandle &) = delete; + + ~LoomHandle() { reset(); } + + T * get() const { return value_; } + + T ** out() { + reset(); + return &value_; + } + + void reset(T * value = nullptr) { + if (value_) { + Release(value_); + } + value_ = value; + } + + private: + T * value_ = nullptr; +}; + +using LoomWorkspace = LoomHandle; +using LoomSource = LoomHandle; +using LoomModule = LoomHandle; +using LoomResult = LoomHandle; +using LoomLinkIndexBuilder = LoomHandle; +using LoomLinkIndex = LoomHandle; +using LoomLinker = LoomHandle; +using LoomLaunchConfigProgram = LoomHandle; + +struct HrxLoomJitDeleter { + void operator()(ggml_hrx_loom_jit_amdgpu * jit) const { ggml_hrx_loom_jit_amdgpu_release(jit); } +}; + +hrx_status_t ggml_hrx_loom_jit_make_status(hrx_status_code_t code, const char * message) { + return hrx_make_status(code, message ? message : "GGML HRX Loom JIT failure"); +} + +hrx_status_code_t ggml_hrx_loom_jit_status_code_from_loom(loomc_status_code_t code) { + switch (code) { + case LOOMC_STATUS_OK: + return HRX_STATUS_OK; + case LOOMC_STATUS_CANCELLED: + return HRX_STATUS_CANCELLED; + case LOOMC_STATUS_UNKNOWN: + return HRX_STATUS_UNKNOWN; + case LOOMC_STATUS_INVALID_ARGUMENT: + return HRX_STATUS_INVALID_ARGUMENT; + case LOOMC_STATUS_DEADLINE_EXCEEDED: + return HRX_STATUS_DEADLINE_EXCEEDED; + case LOOMC_STATUS_NOT_FOUND: + return HRX_STATUS_NOT_FOUND; + case LOOMC_STATUS_ALREADY_EXISTS: + return HRX_STATUS_ALREADY_EXISTS; + case LOOMC_STATUS_PERMISSION_DENIED: + return HRX_STATUS_PERMISSION_DENIED; + case LOOMC_STATUS_RESOURCE_EXHAUSTED: + return HRX_STATUS_OUT_OF_MEMORY; + case LOOMC_STATUS_FAILED_PRECONDITION: + return HRX_STATUS_FAILED_PRECONDITION; + case LOOMC_STATUS_ABORTED: + return HRX_STATUS_ABORTED; + case LOOMC_STATUS_OUT_OF_RANGE: + return HRX_STATUS_OUT_OF_RANGE; + case LOOMC_STATUS_UNIMPLEMENTED: + return HRX_STATUS_UNIMPLEMENTED; + case LOOMC_STATUS_INTERNAL: + return HRX_STATUS_INTERNAL; + case LOOMC_STATUS_UNAVAILABLE: + return HRX_STATUS_UNAVAILABLE; + case LOOMC_STATUS_DATA_LOSS: + return HRX_STATUS_DATA_LOSS; + case LOOMC_STATUS_UNAUTHENTICATED: + return HRX_STATUS_PERMISSION_DENIED; + case LOOMC_STATUS_DEFERRED: + return HRX_STATUS_UNAVAILABLE; + case LOOMC_STATUS_INCOMPATIBLE: + return HRX_STATUS_FAILED_PRECONDITION; + case LOOMC_STATUS_CODE_MASK: + return HRX_STATUS_INTERNAL; + } + return HRX_STATUS_INTERNAL; +} + +std::string ggml_hrx_loom_jit_format_status(loomc_status_t status) { + if (loomc_status_is_ok(status)) { + return "OK"; + } + loomc_host_size_t length = 0; + loomc_status_format(status, 0, nullptr, &length); + std::unique_ptr buffer(new (std::nothrow) char[length + 1]()); + if (!buffer) { + char fallback[4096] = { 0 }; + loomc_host_size_t fallback_length = 0; + loomc_status_format(status, sizeof(fallback), fallback, &fallback_length); + if (fallback_length >= sizeof(fallback)) { + fallback_length = sizeof(fallback) - 1; + } + return std::string(fallback, fallback_length); + } + loomc_host_size_t actual_length = 0; + loomc_status_format(status, length + 1, buffer.get(), &actual_length); + return std::string(buffer.get(), actual_length); +} + +void ggml_hrx_loom_jit_spam_failure(const char * context, const std::string & message) { + std::fprintf(stderr, "HRX Loom JIT %s failed: %s\n", context ? context : "operation", message.c_str()); + std::fflush(stderr); +} + +hrx_status_t ggml_hrx_loom_jit_status_from_loom(loomc_status_t status, const char * context) { + if (loomc_status_is_ok(status)) { + return hrx_ok_status(); + } + const hrx_status_code_t hrx_code = ggml_hrx_loom_jit_status_code_from_loom(loomc_status_code(status)); + const std::string formatted_status = ggml_hrx_loom_jit_format_status(status); + loomc_status_free(status); + std::string message = std::string(context ? context : "loomc") + ": " + formatted_status; + ggml_hrx_loom_jit_spam_failure(context, message); + return ggml_hrx_loom_jit_make_status(hrx_code, message.c_str()); +} + +hrx_status_t ggml_hrx_loom_jit_status_from_result(const loomc_result_t * result, const char * context) { + if (result && loomc_result_succeeded(result)) { + return hrx_ok_status(); + } + std::string message = context ? context : "loomc result failed"; + if (result) { + const loomc_host_size_t diagnostic_count = loomc_result_diagnostic_count(result); + for (loomc_host_size_t i = 0; i < diagnostic_count; ++i) { + const loomc_diagnostic_t * diagnostic = loomc_result_diagnostic_at(result, i); + if (!diagnostic) { + continue; + } + message += "\n diagnostic["; + message += std::to_string(static_cast(i)); + message += "] "; + message.append(diagnostic->code.data, diagnostic->code.size); + message += ": "; + message.append(diagnostic->message.data, diagnostic->message.size); + if (diagnostic->range.start_line || diagnostic->range.start_column) { + message += " @ "; + message += std::to_string(diagnostic->range.start_line); + message += ":"; + message += std::to_string(diagnostic->range.start_column); + } + } + } + ggml_hrx_loom_jit_spam_failure(context, message); + return ggml_hrx_loom_jit_make_status(HRX_STATUS_FAILED_PRECONDITION, message.c_str()); +} + +void * ggml_hrx_loom_jit_malloc_copy(const void * data, size_t size, bool nul_terminate) { + if (!data || size == 0) { + return nullptr; + } + const size_t alloc_size = nul_terminate ? size + 1 : size; + void * result = nullptr; + hrx_status_t status = hrx_host_allocator_malloc_uninitialized(hrx_host_allocator_system(), alloc_size, &result); + if (!hrx_status_is_ok(status)) { + hrx_status_ignore(status); + return nullptr; + } + std::memcpy(result, data, size); + if (nul_terminate) { + static_cast(result)[size] = 0; + } + return result; +} + +const loomc_artifact_t * ggml_hrx_loom_jit_find_artifact(const loomc_result_t * result, + loomc_artifact_kind_t kind, + loomc_string_view_t format) { + for (loomc_host_size_t i = 0; i < loomc_result_artifact_count(result); ++i) { + const loomc_artifact_t * artifact = loomc_result_artifact_at(result, i); + if (!artifact) { + continue; + } + if (artifact->kind == kind && loomc_string_view_equal(artifact->format, format)) { + return artifact; + } + } + return nullptr; +} + +hrx_status_t ggml_hrx_loom_jit_copy_artifact_bytes(const loomc_artifact_t * artifact, + void ** out_data, + size_t * out_size, + bool nul_terminate) { + if (out_data) { + *out_data = nullptr; + } + if (out_size) { + *out_size = 0; + } + if (!artifact || !out_data || !out_size) { + return hrx_ok_status(); + } + // The new loomc API exposes artifact bytes as an opaque byte sequence; + // contiguous access is best-effort (true for the artifacts we produce). + loomc_byte_span_t span = {}; + if (!loomc_byte_sequence_try_get_contiguous_span(artifact->contents, &span)) { + return ggml_hrx_loom_jit_make_status(HRX_STATUS_INTERNAL, + "Loom artifact is not a contiguous byte sequence"); + } + void * copy = ggml_hrx_loom_jit_malloc_copy(span.data, span.data_length, nul_terminate); + if (!copy) { + return ggml_hrx_loom_jit_make_status(HRX_STATUS_OUT_OF_MEMORY, "failed to copy Loom artifact"); + } + *out_data = copy; + *out_size = span.data_length; + return hrx_ok_status(); +} + +hrx_status_t ggml_hrx_loom_jit_evaluate_launch_config(const loomc_artifact_t * artifact, + const char * root_symbol, + const int64_t * workload_arguments, + size_t workload_argument_count, + ggml_hrx_loom_jit_launch_config * out_launch_config) { + if (!out_launch_config) { + return ggml_hrx_loom_jit_make_status(HRX_STATUS_INVALID_ARGUMENT, "out_launch_config is required"); + } + if (!artifact) { + return ggml_hrx_loom_jit_make_status(HRX_STATUS_NOT_FOUND, "Loom did not return a launch-config artifact"); + } + + LoomLaunchConfigProgram program; + loomc_status_t status = + loomc_launch_config_program_load(artifact, loomc_allocator_system(), program.out()); + if (!loomc_status_is_ok(status)) { + return ggml_hrx_loom_jit_status_from_loom(status, "load Loom launch config program"); + } + + std::string export_name = root_symbol ? root_symbol : ""; + if (!export_name.empty() && export_name.front() == '@') { + export_name.erase(export_name.begin()); + } + loomc_launch_config_function_t function = loomc_launch_config_function_invalid(); + status = loomc_launch_config_program_lookup_function(program.get(), loomc_make_cstring_view(export_name.c_str()), + &function); + if (!loomc_status_is_ok(status)) { + return ggml_hrx_loom_jit_status_from_loom(status, "find Loom launch config function"); + } + + std::vector workload_bits; + workload_bits.reserve(workload_argument_count); + for (size_t i = 0; i < workload_argument_count; ++i) { + workload_bits.push_back(static_cast(workload_arguments[i])); + } + + loomc_launch_config_t launch_config = {}; + launch_config.type = LOOMC_STRUCTURE_TYPE_LAUNCH_CONFIG; + launch_config.structure_size = sizeof(launch_config); + + status = loomc_launch_config_program_invoke(program.get(), function, + workload_bits.empty() ? nullptr : workload_bits.data(), + workload_bits.size(), &launch_config); + if (!loomc_status_is_ok(status)) { + return ggml_hrx_loom_jit_status_from_loom(status, "invoke Loom launch config function"); + } + if (!launch_config.workgroup_count.x || !launch_config.workgroup_count.y || !launch_config.workgroup_count.z || + !launch_config.workgroup_size.x || !launch_config.workgroup_size.y || !launch_config.workgroup_size.z) { + return ggml_hrx_loom_jit_make_status(HRX_STATUS_FAILED_PRECONDITION, + "Loom launch config did not provide required workgroup count and size"); + } + + out_launch_config->fields = 0; + out_launch_config->workgroup_count[0] = launch_config.workgroup_count.x; + out_launch_config->workgroup_count[1] = launch_config.workgroup_count.y; + out_launch_config->workgroup_count[2] = launch_config.workgroup_count.z; + out_launch_config->workgroup_size[0] = launch_config.workgroup_size.x; + out_launch_config->workgroup_size[1] = launch_config.workgroup_size.y; + out_launch_config->workgroup_size[2] = launch_config.workgroup_size.z; + out_launch_config->subgroup_size = launch_config.subgroup_size; + out_launch_config->workgroup_storage_bytes = launch_config.workgroup_storage_bytes; + out_launch_config->workload_argument_count = workload_argument_count; + return hrx_ok_status(); +} + +hrx_status_t ggml_hrx_loom_jit_parse_sanitizer_checks(const char * value, loomc_sanitizer_checks_t * out_checks) { + *out_checks = 0; + if (!value || !value[0] || std::strcmp(value, "0") == 0 || std::strcmp(value, "none") == 0) { + return hrx_ok_status(); + } + if (std::strcmp(value, "access") == 0 || std::strcmp(value, "asan") == 0) { + *out_checks = LOOMC_SANITIZER_CHECKS_ASAN_LIKE; + return hrx_ok_status(); + } + if (std::strcmp(value, "value") == 0) { + *out_checks = LOOMC_SANITIZER_CHECK_VALUE; + return hrx_ok_status(); + } + if (std::strcmp(value, "operation") == 0) { + *out_checks = LOOMC_SANITIZER_CHECK_OPERATION; + return hrx_ok_status(); + } + if (std::strcmp(value, "ubsan") == 0) { + *out_checks = LOOMC_SANITIZER_CHECKS_UBSAN_LIKE; + return hrx_ok_status(); + } + if (std::strcmp(value, "all") == 0) { + *out_checks = LOOMC_SANITIZER_CHECK_ACCESS | LOOMC_SANITIZER_CHECK_VALUE | LOOMC_SANITIZER_CHECK_OPERATION; + return hrx_ok_status(); + } + char message[256] = { 0 }; + std::snprintf(message, sizeof(message), "unsupported GGML_HRX_LOOM_SANITIZER '%s'", value); + return ggml_hrx_loom_jit_make_status(HRX_STATUS_INVALID_ARGUMENT, message); +} + +hrx_status_t ggml_hrx_loom_jit_parse_sanitizer_reporting(const char * value, + loomc_sanitizer_reporting_mode_t * out_reporting_mode) { + *out_reporting_mode = LOOMC_SANITIZER_REPORTING_MODE_REPORT_ONLY; + if (!value || !value[0] || std::strcmp(value, "report-only") == 0 || std::strcmp(value, "report_only") == 0 || + std::strcmp(value, "report") == 0) { + return hrx_ok_status(); + } + if (std::strcmp(value, "default") == 0) { + *out_reporting_mode = LOOMC_SANITIZER_REPORTING_MODE_DEFAULT; + return hrx_ok_status(); + } + if (std::strcmp(value, "trap") == 0) { + *out_reporting_mode = LOOMC_SANITIZER_REPORTING_MODE_TRAP; + return hrx_ok_status(); + } + char message[256] = { 0 }; + std::snprintf(message, sizeof(message), "unsupported GGML_HRX_LOOM_SANITIZER_REPORTING '%s'", value); + return ggml_hrx_loom_jit_make_status(HRX_STATUS_INVALID_ARGUMENT, message); +} + +loomc_amdgpu_runtime_global_flags_t ggml_hrx_loom_jit_runtime_globals(loomc_sanitizer_checks_t sanitizer_checks) { + if (!sanitizer_checks) { + return LOOMC_AMDGPU_RUNTIME_GLOBAL_NONE; + } + loomc_amdgpu_runtime_global_flags_t runtime_globals = LOOMC_AMDGPU_RUNTIME_GLOBAL_FEEDBACK_CONFIG; + if (sanitizer_checks & LOOMC_SANITIZER_CHECK_ACCESS) { + runtime_globals |= LOOMC_AMDGPU_RUNTIME_GLOBAL_ASAN_CONFIG; + } + return runtime_globals; +} + +} // namespace + +hrx_status_t ggml_hrx_loom_jit_amdgpu_create(const ggml_hrx_loom_jit_amdgpu_options * options, + ggml_hrx_loom_jit_amdgpu ** out_jit) { + if (!out_jit) { + return ggml_hrx_loom_jit_make_status(HRX_STATUS_INVALID_ARGUMENT, "out_jit must not be NULL"); + } + *out_jit = nullptr; + if (!options || !options->processor || options->processor[0] == 0) { + return ggml_hrx_loom_jit_make_status(HRX_STATUS_INVALID_ARGUMENT, + "valid ggml_hrx_loom_jit_amdgpu_options_t with processor is required"); + } + + std::unique_ptr jit(new (std::nothrow) ggml_hrx_loom_jit_amdgpu()); + if (!jit) { + return ggml_hrx_loom_jit_make_status(HRX_STATUS_OUT_OF_MEMORY, "failed to allocate GGML HRX Loom JIT"); + } + + LoomResult result; + loomc_status_t status = loomc_target_environment_create_amdgpu(loomc_allocator_system(), &jit->target_environment); + if (!loomc_status_is_ok(status)) { + return ggml_hrx_loom_jit_status_from_loom(status, "create AMDGPU target environment"); + } + + loomc_context_target_options_t target_options = {}; + target_options.type = LOOMC_STRUCTURE_TYPE_CONTEXT_TARGET_OPTIONS; + target_options.structure_size = sizeof(target_options); + target_options.target_environment = jit->target_environment; + loomc_context_options_t context_options = {}; + context_options.type = LOOMC_STRUCTURE_TYPE_CONTEXT_OPTIONS; + context_options.structure_size = sizeof(context_options); + context_options.next = &target_options; + status = loomc_context_create(&context_options, loomc_allocator_system(), &jit->context); + if (!loomc_status_is_ok(status)) { + return ggml_hrx_loom_jit_status_from_loom(status, "create Loom context"); + } + + loomc_amdgpu_profile_options_t profile_options = {}; + profile_options.type = LOOMC_STRUCTURE_TYPE_AMDGPU_PROFILE_OPTIONS; + profile_options.structure_size = sizeof(profile_options); + profile_options.identifier = loomc_make_cstring_view(options->identifier); + profile_options.identity.target = loomc_make_cstring_view(options->processor); + status = loomc_target_profile_create_amdgpu(jit->target_environment, &profile_options, loomc_allocator_system(), + &jit->target_profile); + if (!loomc_status_is_ok(status)) { + return ggml_hrx_loom_jit_status_from_loom(status, "create AMDGPU target profile"); + } + status = loomc_compiler_create(jit->context, nullptr, loomc_allocator_system(), &jit->compiler); + if (!loomc_status_is_ok(status)) { + return ggml_hrx_loom_jit_status_from_loom(status, "create Loom compiler"); + } + + loomc_sanitizer_options_t sanitizer_options = {}; + sanitizer_options.type = LOOMC_STRUCTURE_TYPE_SANITIZER_OPTIONS; + sanitizer_options.structure_size = sizeof(sanitizer_options); + sanitizer_options.next = nullptr; + hrx_status_t sanitizer_status = + ggml_hrx_loom_jit_parse_sanitizer_checks(options->sanitizer, &sanitizer_options.checks); + if (!hrx_status_is_ok(sanitizer_status)) { + return sanitizer_status; + } + if (sanitizer_options.checks) { + hrx_status_t sanitizer_reporting_status = ggml_hrx_loom_jit_parse_sanitizer_reporting( + options->sanitizer_reporting, &sanitizer_options.reporting_mode); + if (!hrx_status_is_ok(sanitizer_reporting_status)) { + return sanitizer_reporting_status; + } + } + jit->runtime_globals = ggml_hrx_loom_jit_runtime_globals(sanitizer_options.checks); + loomc_target_pipeline_options_t pipeline_options = {}; + pipeline_options.type = LOOMC_STRUCTURE_TYPE_TARGET_PIPELINE_OPTIONS; + pipeline_options.structure_size = sizeof(pipeline_options); + pipeline_options.next = sanitizer_options.checks ? static_cast(&sanitizer_options) : nullptr; + pipeline_options.identifier = loomc_make_cstring_view("ggml-hrx-amdgpu-jit-prepared-low"); + pipeline_options.kind = LOOMC_TARGET_PIPELINE_KIND_PREPARED_LOW; + pipeline_options.control_flow_lowering = LOOMC_TARGET_CONTROL_FLOW_LOWERING_CFG; + pipeline_options.source_to_low_max_errors = 20; + status = loomc_pass_program_create_from_target_pipeline(jit->context, &pipeline_options, loomc_allocator_system(), + &jit->pass_program, result.out()); + if (!loomc_status_is_ok(status)) { + return ggml_hrx_loom_jit_status_from_loom(status, "create target pass program"); + } + if (!loomc_result_succeeded(result.get())) { + return ggml_hrx_loom_jit_status_from_result(result.get(), "target pass program failed"); + } + + *out_jit = jit.release(); + return hrx_ok_status(); +} + +void ggml_hrx_loom_jit_amdgpu_release(ggml_hrx_loom_jit_amdgpu * jit) { + if (!jit) { + return; + } + loomc_pass_program_release(jit->pass_program); + loomc_compiler_release(jit->compiler); + loomc_target_profile_release(jit->target_profile); + loomc_context_release(jit->context); + loomc_target_environment_release(jit->target_environment); + delete jit; +} + +hrx_status_t ggml_hrx_loom_jit_amdgpu_compile(ggml_hrx_loom_jit_amdgpu * jit, + const ggml_hrx_loom_jit_compile_options * options, + ggml_hrx_loom_jit_compile_result * out_result) { + if (!out_result) { + return ggml_hrx_loom_jit_make_status(HRX_STATUS_INVALID_ARGUMENT, "out_result must not be NULL"); + } + out_result->reset(); + if (!jit || !options || !options->source_data || options->source_size == 0 || !options->root_symbol || + options->root_symbol[0] == 0) { + return ggml_hrx_loom_jit_make_status( + HRX_STATUS_INVALID_ARGUMENT, "valid GGML HRX Loom JIT compile options with source and root are required"); + } + if (options->config_binding_count > 0 && !options->config_bindings) { + return ggml_hrx_loom_jit_make_status(HRX_STATUS_INVALID_ARGUMENT, + "GGML HRX Loom JIT config binding count requires config bindings"); + } + if (options->workload_argument_count > 0 && !options->workload_arguments) { + return ggml_hrx_loom_jit_make_status(HRX_STATUS_INVALID_ARGUMENT, + "GGML HRX Loom JIT workload argument count requires workload arguments"); + } + if (options->dependency_count > 0 && !options->dependencies) { + return ggml_hrx_loom_jit_make_status(HRX_STATUS_INVALID_ARGUMENT, + "GGML HRX Loom JIT dependency count requires dependencies"); + } + for (size_t i = 0; i < options->config_binding_count; ++i) { + if (!options->config_bindings[i].key || !options->config_bindings[i].value) { + return ggml_hrx_loom_jit_make_status(HRX_STATUS_INVALID_ARGUMENT, + "GGML HRX Loom JIT config binding keys and values must not be NULL"); + } + } + + std::unique_ptr config_bindings; + if (options->config_binding_count > 0) { + config_bindings.reset(new (std::nothrow) loomc_config_binding_t[options->config_binding_count]()); + if (!config_bindings) { + return ggml_hrx_loom_jit_make_status(HRX_STATUS_OUT_OF_MEMORY, + "failed to allocate GGML HRX Loom JIT config bindings"); + } + for (size_t i = 0; i < options->config_binding_count; ++i) { + config_bindings[i].key = loomc_make_cstring_view(options->config_bindings[i].key); + config_bindings[i].value = loomc_make_cstring_view(options->config_bindings[i].value); + } + } + + LoomWorkspace workspace; + LoomSource source; + LoomModule module; + LoomResult result; + std::string specialized_source; + if (options->source_format == GGML_HRX_LOOM_JIT_SOURCE_FORMAT_TEXT && options->dependency_count == 0 && + options->workload_argument_count > 0) { + std::string specialization_error; + const std::string source_text(static_cast(options->source_data), options->source_size); + if (!ggml_hrx_loom_specialize_workload_text(source_text, options->root_symbol, options->workload_arguments, + options->workload_argument_count, specialized_source, + specialization_error)) { + return ggml_hrx_loom_jit_make_status(HRX_STATUS_FAILED_PRECONDITION, specialization_error.c_str()); + } + } + + loomc_status_t status = loomc_workspace_create(nullptr, loomc_allocator_system(), workspace.out()); + if (!loomc_status_is_ok(status)) { + return ggml_hrx_loom_jit_status_from_loom(status, "create Loom workspace"); + } + + loomc_source_options_t source_options = {}; + source_options.type = LOOMC_STRUCTURE_TYPE_SOURCE_OPTIONS; + source_options.structure_size = sizeof(source_options); + source_options.format = options->source_format == GGML_HRX_LOOM_JIT_SOURCE_FORMAT_BYTECODE ? + LOOMC_SOURCE_FORMAT_BYTECODE : + LOOMC_SOURCE_FORMAT_TEXT; + source_options.identifier = loomc_make_cstring_view(options->source_identifier); + source_options.contents = specialized_source.empty() ? + loomc_make_byte_span(options->source_data, options->source_size) : + loomc_make_byte_span(specialized_source.data(), specialized_source.size()); + source_options.storage = LOOMC_SOURCE_STORAGE_BORROWED; + status = loomc_source_create(&source_options, loomc_allocator_system(), source.out()); + if (!loomc_status_is_ok(status)) { + return ggml_hrx_loom_jit_status_from_loom(status, "create Loom source"); + } + + LoomLinkIndexBuilder link_index_builder; + status = loomc_link_index_builder_create(jit->context, nullptr, loomc_allocator_system(), link_index_builder.out()); + if (!loomc_status_is_ok(status)) { + return ggml_hrx_loom_jit_status_from_loom(status, "create Loom link index builder"); + } + loomc_link_index_source_options_t link_source_options = {}; + link_source_options.provider_name = loomc_make_cstring_view(options->source_identifier); + link_source_options.role = LOOMC_LINK_PROVIDER_ROLE_INPUT; + status = loomc_link_index_builder_add_source(link_index_builder.get(), source.get(), &link_source_options, nullptr); + if (!loomc_status_is_ok(status)) { + return ggml_hrx_loom_jit_status_from_loom(status, "index Loom source"); + } + std::vector dependency_sources; + dependency_sources.reserve(options->dependency_count); + for (size_t i = 0; i < options->dependency_count; ++i) { + const ggml_hrx_loom_jit_source & dependency = options->dependencies[i]; + if (!dependency.source_data || dependency.source_size == 0 || !dependency.source_identifier) { + for (loomc_source_t * dependency_source : dependency_sources) { + loomc_source_release(dependency_source); + } + return ggml_hrx_loom_jit_make_status(HRX_STATUS_INVALID_ARGUMENT, "invalid Loom JIT dependency source"); + } + loomc_source_options_t dependency_options = {}; + dependency_options.type = LOOMC_STRUCTURE_TYPE_SOURCE_OPTIONS; + dependency_options.structure_size = sizeof(dependency_options); + dependency_options.format = dependency.source_format == GGML_HRX_LOOM_JIT_SOURCE_FORMAT_BYTECODE ? + LOOMC_SOURCE_FORMAT_BYTECODE : + LOOMC_SOURCE_FORMAT_TEXT; + dependency_options.identifier = loomc_make_cstring_view(dependency.source_identifier); + dependency_options.contents = loomc_make_byte_span(dependency.source_data, dependency.source_size); + dependency_options.storage = LOOMC_SOURCE_STORAGE_BORROWED; + loomc_source_t * dependency_source = nullptr; + status = loomc_source_create(&dependency_options, loomc_allocator_system(), &dependency_source); + if (!loomc_status_is_ok(status)) { + for (loomc_source_t * retained_source : dependency_sources) { + loomc_source_release(retained_source); + } + return ggml_hrx_loom_jit_status_from_loom(status, "create Loom dependency source"); + } + dependency_sources.push_back(dependency_source); + loomc_link_index_source_options_t dependency_link_options = {}; + dependency_link_options.provider_name = loomc_make_cstring_view(dependency.source_identifier); + if (dependency.source_format == GGML_HRX_LOOM_JIT_SOURCE_FORMAT_BYTECODE) { + dependency_link_options.role = LOOMC_LINK_PROVIDER_ROLE_INPUT; + } else { + const std::string dependency_text(static_cast(dependency.source_data), + dependency.source_size); + dependency_link_options.role = dependency_text.find("config.decl") != std::string::npos ? + LOOMC_LINK_PROVIDER_ROLE_INPUT : + LOOMC_LINK_PROVIDER_ROLE_LIBRARY; + } + status = loomc_link_index_builder_add_source(link_index_builder.get(), dependency_source, + &dependency_link_options, nullptr); + if (!loomc_status_is_ok(status)) { + for (loomc_source_t * retained_source : dependency_sources) { + loomc_source_release(retained_source); + } + return ggml_hrx_loom_jit_status_from_loom(status, "index Loom dependency source"); + } + } + LoomLinkIndex link_index; + status = loomc_link_index_builder_finish(link_index_builder.get(), link_index.out(), result.out()); + for (loomc_source_t * dependency_source : dependency_sources) { + loomc_source_release(dependency_source); + } + if (!loomc_status_is_ok(status)) { + return ggml_hrx_loom_jit_status_from_loom(status, "finish Loom link index"); + } + if (!loomc_result_succeeded(result.get())) { + return ggml_hrx_loom_jit_status_from_result(result.get(), "Loom source indexing failed"); + } + result.reset(); + + LoomLinker linker; + status = loomc_linker_create(jit->context, nullptr, loomc_allocator_system(), linker.out()); + if (!loomc_status_is_ok(status)) { + return ggml_hrx_loom_jit_status_from_loom(status, "create Loom linker"); + } + + // BUILD-authored kernel recipes first archive their primary source and + // libraries, then compile roots from that linked module. Preserve that + // composition exactly. In particular, config.decl operations are not + // callable dependency edges and disappear if raw source libraries are + // fed directly to a selective link. Bytecode sources with workload + // arguments also use this path so the linked module can be serialized + // back to text for the current workload specialization pass. + LoomSource archived_source; + LoomSource specialized_archive_source; + std::string specialized_archive_text; + const bool needs_archive_source = + options->dependency_count > 0 || + (options->source_format == GGML_HRX_LOOM_JIT_SOURCE_FORMAT_BYTECODE && options->workload_argument_count > 0); + if (needs_archive_source) { + LoomModule archive_module; + loomc_link_options_t archive_options = {}; + archive_options.type = LOOMC_STRUCTURE_TYPE_LINK_OPTIONS; + archive_options.structure_size = sizeof(archive_options); + archive_options.link_index = link_index.get(); + archive_options.module_name = loomc_make_cstring_view(options->module_name); + archive_options.flags = LOOMC_LINK_FLAG_STRIP_TEST_SYMBOLS; + status = loomc_link_module(linker.get(), workspace.get(), &archive_options, archive_module.out(), result.out()); + if (!loomc_status_is_ok(status)) { + return ggml_hrx_loom_jit_status_from_loom(status, "archive Loom kernel module"); + } + if (!loomc_result_succeeded(result.get())) { + return ggml_hrx_loom_jit_status_from_result(result.get(), "Loom kernel archive linking failed"); + } + result.reset(); + loomc_module_serialize_options_t serialize_options = {}; + serialize_options.type = LOOMC_STRUCTURE_TYPE_MODULE_SERIALIZE_OPTIONS; + serialize_options.structure_size = sizeof(serialize_options); + serialize_options.format = LOOMC_SOURCE_FORMAT_TEXT; + serialize_options.identifier = loomc_make_cstring_view(options->source_identifier); + status = loomc_module_serialize_to_source(archive_module.get(), &serialize_options, loomc_allocator_system(), + archived_source.out()); + if (!loomc_status_is_ok(status)) { + return ggml_hrx_loom_jit_status_from_loom(status, "serialize Loom kernel archive"); + } + loomc_source_t * archive_index_source = archived_source.get(); + if (options->workload_argument_count > 0) { + const loomc_byte_span_t archive_contents = loomc_source_contents(archived_source.get()); + const std::string archive_text(reinterpret_cast(archive_contents.data), + archive_contents.data_length); + std::string specialization_error; + if (!ggml_hrx_loom_specialize_workload_text(archive_text, options->root_symbol, options->workload_arguments, + options->workload_argument_count, specialized_archive_text, + specialization_error)) { + return ggml_hrx_loom_jit_make_status(HRX_STATUS_FAILED_PRECONDITION, specialization_error.c_str()); + } + loomc_source_options_t specialized_options = {}; + specialized_options.type = LOOMC_STRUCTURE_TYPE_SOURCE_OPTIONS; + specialized_options.structure_size = sizeof(specialized_options); + specialized_options.format = LOOMC_SOURCE_FORMAT_TEXT; + specialized_options.identifier = loomc_make_cstring_view(options->source_identifier); + specialized_options.contents = + loomc_make_byte_span(specialized_archive_text.data(), specialized_archive_text.size()); + specialized_options.storage = LOOMC_SOURCE_STORAGE_BORROWED; + status = + loomc_source_create(&specialized_options, loomc_allocator_system(), specialized_archive_source.out()); + if (!loomc_status_is_ok(status)) { + return ggml_hrx_loom_jit_status_from_loom(status, "create specialized Loom kernel archive"); + } + archive_index_source = specialized_archive_source.get(); + } + LoomLinkIndexBuilder archive_index_builder; + status = loomc_link_index_builder_create(jit->context, nullptr, loomc_allocator_system(), + archive_index_builder.out()); + if (!loomc_status_is_ok(status)) { + return ggml_hrx_loom_jit_status_from_loom(status, "create Loom kernel archive index"); + } + loomc_link_index_source_options_t archive_source_options = {}; + archive_source_options.provider_name = loomc_make_cstring_view(options->source_identifier); + archive_source_options.role = LOOMC_LINK_PROVIDER_ROLE_INPUT; + status = loomc_link_index_builder_add_source(archive_index_builder.get(), archive_index_source, + &archive_source_options, nullptr); + if (!loomc_status_is_ok(status)) { + return ggml_hrx_loom_jit_status_from_loom(status, "index Loom kernel archive"); + } + status = loomc_link_index_builder_finish(archive_index_builder.get(), link_index.out(), result.out()); + if (!loomc_status_is_ok(status)) { + return ggml_hrx_loom_jit_status_from_loom(status, "finish Loom kernel archive index"); + } + if (!loomc_result_succeeded(result.get())) { + return ggml_hrx_loom_jit_status_from_result(result.get(), "Loom kernel archive indexing failed"); + } + result.reset(); + } + const loomc_string_view_t root_symbols[] = { loomc_make_cstring_view(options->root_symbol) }; + loomc_link_options_t link_options = {}; + link_options.type = LOOMC_STRUCTURE_TYPE_LINK_OPTIONS; + link_options.structure_size = sizeof(link_options); + link_options.next = nullptr; + link_options.link_index = link_index.get(); + link_options.module_name = loomc_make_cstring_view(options->module_name); + link_options.mode = LOOMC_LINK_MODE_LINK; + link_options.root_symbols = root_symbols; + link_options.root_symbol_count = 1; + link_options.flags = LOOMC_LINK_FLAG_STRIP_TEST_SYMBOLS; + link_options.config.bindings = config_bindings.get(); + link_options.config.binding_count = options->config_binding_count; + link_options.config.flags = LOOMC_CONFIG_POLICY_FLAG_REQUIRE_RESOLVED; + status = loomc_link_module(linker.get(), workspace.get(), &link_options, module.out(), result.out()); + if (!loomc_status_is_ok(status)) { + return ggml_hrx_loom_jit_status_from_loom(status, "link Loom root"); + } + if (!loomc_result_succeeded(result.get())) { + return ggml_hrx_loom_jit_status_from_result(result.get(), "Loom root linking failed"); + } + result.reset(); + + const loomc_target_specialization_t specialization = { + loomc_make_cstring_view(options->root_symbol), + jit->target_profile, + }; + loomc_target_specialization_options_t compile_target_options = {}; + compile_target_options.type = LOOMC_STRUCTURE_TYPE_TARGET_SPECIALIZATION_OPTIONS; + compile_target_options.structure_size = sizeof(compile_target_options); + compile_target_options.specializations = &specialization; + compile_target_options.specialization_count = 1; + loomc_compile_options_t compile_options = {}; + compile_options.type = LOOMC_STRUCTURE_TYPE_COMPILE_OPTIONS; + compile_options.structure_size = sizeof(compile_options); + compile_options.next = &compile_target_options; + compile_options.module_name = loomc_make_cstring_view(options->module_name); + compile_options.artifact_flags = LOOMC_COMPILE_ARTIFACT_FLAG_MODULE_TEXT | LOOMC_COMPILE_ARTIFACT_FLAG_REPORT_JSON; + if (options->evaluate_launch_config) { + compile_options.artifact_flags |= LOOMC_COMPILE_ARTIFACT_FLAG_LAUNCH_CONFIG; + } + // Config bindings are materialized as typed `config.def` ops by the root + // link invocation above; the compiled module already carries exact values, + // so no compile-time config module is needed. REQUIRE_RESOLVED verifies + // every reachable config.get has an exact value. + compile_options.config_flags = LOOMC_CONFIG_POLICY_FLAG_REQUIRE_RESOLVED; + compile_options.config_module = nullptr; + status = loomc_compile_module(jit->compiler, workspace.get(), jit->pass_program, module.get(), &compile_options, + loomc_allocator_system(), result.out()); + if (!loomc_status_is_ok(status)) { + return ggml_hrx_loom_jit_status_from_loom(status, "compile Loom module"); + } + if (!loomc_result_succeeded(result.get())) { + return ggml_hrx_loom_jit_status_from_result(result.get(), "Loom compilation failed"); + } + + hrx_status_t hrx_status = hrx_ok_status(); + const loomc_artifact_t * compile_report = ggml_hrx_loom_jit_find_artifact( + result.get(), LOOMC_ARTIFACT_KIND_REPORT, loomc_make_cstring_view(LOOMC_ARTIFACT_FORMAT_JSON)); + hrx_status = ggml_hrx_loom_jit_copy_artifact_bytes(compile_report, + reinterpret_cast(&out_result->compile_report_json), + &out_result->compile_report_json_size, true); + if (!hrx_status_is_ok(hrx_status)) { + return hrx_status; + } + const loomc_artifact_t * final_module = ggml_hrx_loom_jit_find_artifact( + result.get(), LOOMC_ARTIFACT_KIND_MODULE, loomc_make_cstring_view(LOOMC_ARTIFACT_FORMAT_LOOM_TEXT)); + hrx_status = + ggml_hrx_loom_jit_copy_artifact_bytes(final_module, reinterpret_cast(&out_result->final_module_text), + &out_result->final_module_text_size, true); + if (!hrx_status_is_ok(hrx_status)) { + return hrx_status; + } + if (options->evaluate_launch_config) { + const char * launch_config_symbol = options->launch_config_symbol; + if (launch_config_symbol == nullptr || launch_config_symbol[0] == '\0') { + launch_config_symbol = options->root_symbol; + } + const loomc_artifact_t * launch_config = + ggml_hrx_loom_jit_find_artifact(result.get(), LOOMC_ARTIFACT_KIND_LAUNCH_CONFIG, + loomc_make_cstring_view(LOOMC_ARTIFACT_FORMAT_LOOM_BYTECODE)); + hrx_status = ggml_hrx_loom_jit_evaluate_launch_config( + launch_config, launch_config_symbol, + options->workload_argument_count == 0 ? nullptr : options->workload_arguments, + options->workload_argument_count, &out_result->launch_config); + if (!hrx_status_is_ok(hrx_status)) { + return hrx_status; + } + } + result.reset(); + + loomc_amdgpu_emit_options_t amdgpu_options = {}; + amdgpu_options.type = LOOMC_STRUCTURE_TYPE_AMDGPU_EMIT_OPTIONS; + amdgpu_options.structure_size = sizeof(amdgpu_options); + amdgpu_options.next = nullptr; + amdgpu_options.runtime_globals = jit->runtime_globals; + const loomc_option_entry_t emit_entries[] = { + { + loomc_make_cstring_view(LOOMC_EMIT_OPTION_KEY_IDENTIFIER), + loomc_make_cstring_view(options->artifact_identifier), + }, + }; + loomc_option_dict_t option_dict = {}; + option_dict.type = LOOMC_STRUCTURE_TYPE_OPTION_DICT; + option_dict.structure_size = sizeof(option_dict); + option_dict.next = &amdgpu_options; + option_dict.entries = emit_entries; + option_dict.entry_count = options->artifact_identifier ? 1 : 0; + loomc_artifact_manifest_options_t manifest_options = {}; + manifest_options.type = LOOMC_STRUCTURE_TYPE_ARTIFACT_MANIFEST_OPTIONS; + manifest_options.structure_size = sizeof(manifest_options); + manifest_options.next = &option_dict; + manifest_options.mode = LOOMC_ARTIFACT_MANIFEST_MODE_DETAILS; + loomc_compile_report_options_t report_options = {}; + report_options.type = LOOMC_STRUCTURE_TYPE_COMPILE_REPORT_OPTIONS; + report_options.structure_size = sizeof(report_options); + report_options.next = &manifest_options; + report_options.mode = LOOMC_COMPILE_REPORT_MODE_DETAILS; + loomc_emit_options_t emit_options = {}; + emit_options.type = LOOMC_STRUCTURE_TYPE_EMIT_OPTIONS; + emit_options.structure_size = sizeof(emit_options); + emit_options.next = &report_options; + emit_options.artifact_format = loomc_make_cstring_view(LOOMC_ARTIFACT_FORMAT_AMDGPU_HSACO); + emit_options.identifier = loomc_make_cstring_view(options->artifact_identifier); + emit_options.artifact_flags = LOOMC_EMIT_ARTIFACT_FLAG_PRIMARY; + status = loomc_emit_module(jit->target_environment, workspace.get(), module.get(), &emit_options, + loomc_allocator_system(), result.out()); + if (!loomc_status_is_ok(status)) { + out_result->reset(); + return ggml_hrx_loom_jit_status_from_loom(status, "emit AMDGPU HSACO"); + } + if (!loomc_result_succeeded(result.get())) { + hrx_status = ggml_hrx_loom_jit_status_from_result(result.get(), "AMDGPU HSACO emission failed"); + out_result->reset(); + return hrx_status; + } + + const loomc_artifact_t * hsaco = ggml_hrx_loom_jit_find_artifact( + result.get(), LOOMC_ARTIFACT_KIND_EXECUTABLE, loomc_make_cstring_view(LOOMC_ARTIFACT_FORMAT_AMDGPU_HSACO)); + if (!hsaco) { + out_result->reset(); + return ggml_hrx_loom_jit_make_status(HRX_STATUS_NOT_FOUND, "Loom did not return an AMDGPU HSACO artifact"); + } + hrx_status = ggml_hrx_loom_jit_copy_artifact_bytes(hsaco, &out_result->hsaco_data, &out_result->hsaco_size, false); + if (hrx_status_is_ok(hrx_status)) { + const loomc_artifact_t * report = + ggml_hrx_loom_jit_find_artifact(result.get(), LOOMC_ARTIFACT_KIND_REPORT, + loomc_make_cstring_view(LOOMC_ARTIFACT_FORMAT_COMPILE_REPORT_JSON)); + hrx_status = + ggml_hrx_loom_jit_copy_artifact_bytes(report, reinterpret_cast(&out_result->compile_report_json), + &out_result->compile_report_json_size, true); + } + if (hrx_status_is_ok(hrx_status)) { + const loomc_artifact_t * manifest = + ggml_hrx_loom_jit_find_artifact(result.get(), LOOMC_ARTIFACT_KIND_REPORT, + loomc_make_cstring_view(LOOMC_ARTIFACT_FORMAT_ARTIFACT_MANIFEST_JSON)); + hrx_status = ggml_hrx_loom_jit_copy_artifact_bytes( + manifest, reinterpret_cast(&out_result->manifest_json), &out_result->manifest_json_size, true); + } + + if (!hrx_status_is_ok(hrx_status)) { + out_result->reset(); + } + return hrx_status; +} diff --git a/ggml/src/ggml-hrx/loom-jit.h b/ggml/src/ggml-hrx/loom-jit.h new file mode 100644 index 000000000000..4dd05f73519c --- /dev/null +++ b/ggml/src/ggml-hrx/loom-jit.h @@ -0,0 +1,95 @@ +// Copyright 2026 The HRX Authors +// SPDX-License-Identifier: Apache-2.0 + +#pragma once + +#include "hrx_runtime.h" + +#include +#include +#include + +struct ggml_hrx_loom_jit_amdgpu; + +enum class ggml_hrx_loom_jit_source_format { + Text, + Bytecode, +}; +inline constexpr auto GGML_HRX_LOOM_JIT_SOURCE_FORMAT_TEXT = ggml_hrx_loom_jit_source_format::Text; +inline constexpr auto GGML_HRX_LOOM_JIT_SOURCE_FORMAT_BYTECODE = ggml_hrx_loom_jit_source_format::Bytecode; + +struct ggml_hrx_loom_jit_amdgpu_options { + const char * processor = nullptr; + const char * identifier = nullptr; + const char * sanitizer = nullptr; + const char * sanitizer_reporting = nullptr; +}; + +struct ggml_hrx_loom_jit_config_binding { + const char * key = nullptr; + const char * value = nullptr; +}; + +struct ggml_hrx_loom_jit_source { + const void * source_data = nullptr; + size_t source_size = 0; + ggml_hrx_loom_jit_source_format source_format = ggml_hrx_loom_jit_source_format::Text; + const char * source_identifier = nullptr; +}; + +struct ggml_hrx_loom_jit_launch_config { + std::array workgroup_count = {}; + std::array workgroup_size = {}; + uint32_t subgroup_size = 0; + uint64_t workgroup_storage_bytes = 0; + size_t workload_argument_count = 0; + uint32_t fields = 0; +}; + +struct ggml_hrx_loom_jit_compile_options { + const void * source_data = nullptr; + size_t source_size = 0; + ggml_hrx_loom_jit_source_format source_format = ggml_hrx_loom_jit_source_format::Text; + const char * source_identifier = nullptr; + const char * root_symbol = nullptr; + const char * launch_config_symbol = nullptr; + const char * module_name = nullptr; + const char * artifact_identifier = nullptr; + const ggml_hrx_loom_jit_source * dependencies = nullptr; + size_t dependency_count = 0; + const ggml_hrx_loom_jit_config_binding * config_bindings = nullptr; + size_t config_binding_count = 0; + const int64_t * workload_arguments = nullptr; + size_t workload_argument_count = 0; + bool evaluate_launch_config = false; +}; + +struct ggml_hrx_loom_jit_compile_result { + ggml_hrx_loom_jit_compile_result() = default; + ~ggml_hrx_loom_jit_compile_result(); + ggml_hrx_loom_jit_compile_result(const ggml_hrx_loom_jit_compile_result &) = delete; + ggml_hrx_loom_jit_compile_result & operator=(const ggml_hrx_loom_jit_compile_result &) = delete; + ggml_hrx_loom_jit_compile_result(ggml_hrx_loom_jit_compile_result && other) noexcept; + ggml_hrx_loom_jit_compile_result & operator=(ggml_hrx_loom_jit_compile_result && other) noexcept; + + void reset(); + + void * hsaco_data = nullptr; + size_t hsaco_size = 0; + char * manifest_json = nullptr; + size_t manifest_json_size = 0; + char * compile_report_json = nullptr; + size_t compile_report_json_size = 0; + char * final_module_text = nullptr; + size_t final_module_text_size = 0; + ggml_hrx_loom_jit_launch_config launch_config; +}; + +hrx_status_t ggml_hrx_loom_jit_amdgpu_create(const ggml_hrx_loom_jit_amdgpu_options * options, + ggml_hrx_loom_jit_amdgpu ** out_jit); + +void ggml_hrx_loom_jit_amdgpu_release(ggml_hrx_loom_jit_amdgpu * jit); + +hrx_status_t ggml_hrx_loom_jit_amdgpu_compile(ggml_hrx_loom_jit_amdgpu * jit, + const ggml_hrx_loom_jit_compile_options * options, + ggml_hrx_loom_jit_compile_result * out_result); diff --git a/ggml/src/ggml-hrx/runtime/command-program-executor.cpp b/ggml/src/ggml-hrx/runtime/command-program-executor.cpp new file mode 100644 index 000000000000..d33cd6bcc131 --- /dev/null +++ b/ggml/src/ggml-hrx/runtime/command-program-executor.cpp @@ -0,0 +1,1214 @@ +// Copyright 2026 The HRX Authors +// SPDX-License-Identifier: Apache-2.0 + +#include "command-program-executor.h" + +#include "dispatch/command-program-diagnostics.h" +#include "dispatch/command-program-resolver.h" +#include "ggml-impl.h" +#include "hrx-interop-utils.h" +#include "runtime/kernel-executable-cache.h" +#include "runtime/transient-arena.h" + +#include +#include +#include +#include + +namespace ggml::hrx { + +PreparedProgramConstantBuffer::~PreparedProgramConstantBuffer() { + if (buffer != nullptr) { + hrx_buffer_release(buffer); + } +} + +PreparedProgramConstantBuffer::PreparedProgramConstantBuffer(PreparedProgramConstantBuffer && other) noexcept : + value(other.value), + name(std::move(other.name)), + buffer(std::exchange(other.buffer, nullptr)), + size(other.size) { + other.size = 0; +} + +PreparedProgramConstantBuffer & +PreparedProgramConstantBuffer::operator=(PreparedProgramConstantBuffer && other) noexcept { + if (this != &other) { + if (buffer != nullptr) { + hrx_buffer_release(buffer); + } + value = other.value; + name = std::move(other.name); + buffer = std::exchange(other.buffer, nullptr); + size = other.size; + other.size = 0; + } + return *this; +} + +RecordedCommandGraph::~RecordedCommandGraph() { + if (exec != nullptr) { + hrx_graph_exec_release(exec); + } + if (graph != nullptr) { + hrx_graph_release(graph); + } +} + +RecordedCommandGraph::RecordedCommandGraph(RecordedCommandGraph && other) noexcept : + graph(std::exchange(other.graph, nullptr)), + exec(std::exchange(other.exec, nullptr)), + bound_transient_arena_allocation_id(other.bound_transient_arena_allocation_id), + dispatch_count(other.dispatch_count), + status(std::move(other.status)) { + other.bound_transient_arena_allocation_id = kInvalidTransientArenaAllocationId; + other.dispatch_count = 0; +} + +RecordedCommandGraph & RecordedCommandGraph::operator=(RecordedCommandGraph && other) noexcept { + if (this != &other) { + if (exec != nullptr) { + hrx_graph_exec_release(exec); + } + if (graph != nullptr) { + hrx_graph_release(graph); + } + graph = std::exchange(other.graph, nullptr); + exec = std::exchange(other.exec, nullptr); + bound_transient_arena_allocation_id = other.bound_transient_arena_allocation_id; + dispatch_count = other.dispatch_count; + status = std::move(other.status); + other.bound_transient_arena_allocation_id = kInvalidTransientArenaAllocationId; + other.dispatch_count = 0; + } + return *this; +} + +namespace { + +static const char * status_first_error(const Status & status) { + return status.errors().empty() ? "" : status.errors().front().c_str(); +} + +static bool resource_access_writes(ResourceAccess access) { + return access == ResourceAccess::Write || access == ResourceAccess::ReadWrite; +} + +static Status command_program_metadata_context_valid(const CommandProgramExecutionContext & context) { + Status status; + if (context.target == nullptr) { + status.log("missing HRX target"); + return status; + } + if (context.corpus == nullptr) { + status.log("missing HRX kernel corpus"); + return status; + } + return status; +} + +static Status command_program_preparation_context_valid(const CommandProgramExecutionContext & context) { + Status status; + if (context.device == nullptr) { + status.log("missing HRX device"); + return status; + } + if (context.kernel_executables == nullptr) { + status.log("missing HRX kernel executable cache"); + return status; + } + if (context.host_transfers == nullptr) { + status.log("missing HRX host transfer manager"); + return status; + } + if (context.host_weights == nullptr) { + status.log("missing HRX host weight cache"); + return status; + } + return status; +} + +static Status command_program_transient_context_valid(const CommandProgramExecutionContext & context, + const CommandProgram & commands) { + Status status; + if (commands.transients.arena_size == 0) { + return status; + } + if (context.transient_arena == nullptr) { + status.log("missing HRX transient arena"); + return status; + } + if (context.stream == nullptr) { + status.log("missing HRX stream for transient arena"); + return status; + } + return status; +} + +static bool prepared_program_has_constant(const PreparedCommandProgram & prepared, ValueId value) { + for (const PreparedProgramConstantBuffer & constant : prepared.program_constants) { + if (constant.value == value) { + return true; + } + } + return false; +} + +static bool prepared_execution_context_valid(const CommandProgramExecutionContext & context) { + if (context.stream == nullptr) { + GGML_LOG_ERROR("%s: missing HRX stream\n", __func__); + return false; + } + return true; +} + +static Status ensure_transient_arena(const CommandProgramExecutionContext & context, + const CommandProgram & commands, + TransientArenaAllocationRef & allocation) { + allocation = {}; + Status status = command_program_transient_context_valid(context, commands); + if (!status.success()) { + return status; + } + if (commands.transients.arena_size == 0) { + return status; + } + status = context.transient_arena->ensure_capacity(context.device, context.stream, commands.transients.arena_size); + if (!status.success()) { + return status; + } + allocation = context.transient_arena->current_allocation(); + return status; +} + +static Status initialize_command_program_constants(const CommandProgramExecutionContext & context, + const CommandProgram & commands, + const TransientArenaAllocationRef & allocation, + const PreparedCommandProgram & prepared) { + Status status; + if (commands.constant_initializations.empty()) { + return status; + } + for (const ConstantInitialization & initialization : commands.constant_initializations) { + if (prepared_program_has_constant(prepared, initialization.value)) { + continue; + } + if (allocation.buffer == nullptr) { + status.log("command program has constant initialization %s without a transient arena allocation", + initialization.name.c_str()); + continue; + } + const TransientAllocation * transient = find_transient_allocation(commands.transients, initialization.value); + if (transient == nullptr) { + status.log("constant initialization %s references missing transient value %d", initialization.name.c_str(), + initialization.value.value); + continue; + } + if (initialization.offset > transient->size || + initialization.data.size() > transient->size - initialization.offset) { + status.log("constant initialization %s is outside transient allocation length %zu", + initialization.name.c_str(), transient->size); + continue; + } + // TODO: Track initialized transient arena allocation ids so constants are not transferred every invocation. + if (context.host_transfers == nullptr) { + status.log("constant initialization %s requires a host transfer manager", initialization.name.c_str()); + continue; + } + Status upload_status = context.host_transfers->upload_synchronous( + context.stream, initialization.data.data(), allocation.buffer, + transient->arena_offset + initialization.offset, initialization.data.size()); + if (!upload_status.success()) { + status.log("failed to upload constant initialization %s", initialization.name.c_str()); + status.append(upload_status); + } + } + return status; +} + +static Status initialize_command_program_completion_counters(const CommandProgramExecutionContext & context, + const CommandProgram & commands, + const TransientArenaAllocationRef & allocation) { + Status status; + if (commands.completion_counters.byte_count == 0) { + return status; + } + if (allocation.buffer == nullptr) { + status.log("command program has completion counters without a transient arena allocation"); + return status; + } + if (commands.completion_counters.arena_offset > commands.transients.arena_size || + commands.completion_counters.byte_count > + commands.transients.arena_size - commands.completion_counters.arena_offset) { + status.log("completion counter initialization is outside transient arena length %zu", + commands.transients.arena_size); + return status; + } + const uint32_t zero_pattern = 0; + if (ErrorResult error = take_status( + hrx_stream_fill_buffer(context.stream, allocation.buffer, commands.completion_counters.arena_offset, + commands.completion_counters.byte_count, &zero_pattern, sizeof(zero_pattern)))) { + status.log("failed to initialize completion counters: %s", error->c_str()); + } + return status; +} + +static std::string format_resolved_command_context(const ResolvedCommand & command) { + std::ostringstream out; + out << "command " << command.ordinal << " kind=" << command_kind_name(command.kind) + << " kernel_id=" << command.kernel.kernel_id << " bindings=" << command.bindings.size(); + return out.str(); +} + +static std::string format_prepared_command_context(const PreparedCommand & command) { + std::ostringstream out; + out << "command " << command.ordinal << " kind=" << command_kind_name(command.kind); + if (command.kind == CommandKind::Kernel) { + out << " kernel_id=" << command.kernel.specialization.kernel_id + << " bindings=" << command.kernel.bindings.size(); + } + return out.str(); +} + +static Dispatch build_dispatch(const ResolvedCommand & command) { + Dispatch dispatch; + dispatch.kernel = command.kernel; + dispatch.bindings.reserve(command.bindings.size()); + for (const ResolvedCommandBinding & binding : command.bindings) { + dispatch.bindings.push_back({ binding.binding.value, binding.binding.offset, binding.binding.length }); + } + return dispatch; +} + +struct GraphValueAccess { + bool read = false; + bool write = false; +}; + +static std::unordered_map collect_graph_value_access(const CommandProgram & commands) { + std::unordered_map access_by_value; + auto append_command_list_access = [&](const std::vector & command_list) { + for (const Command & command : command_list) { + for (const CommandBinding & binding : command.bindings) { + if (binding.origin != CommandBindingOrigin::GraphValue) { + continue; + } + GraphValueAccess & access = access_by_value[binding.value.value]; + switch (binding.access) { + case ResourceAccess::Read: + access.read = true; + break; + case ResourceAccess::Write: + access.write = true; + break; + case ResourceAccess::ReadWrite: + access.read = true; + access.write = true; + break; + } + } + } + }; + append_command_list_access(commands.initialization_commands); + append_command_list_access(commands.commands); + return access_by_value; +} + +static CommandProgramBindings materialize_host_bindings(const CommandProgramExecutionContext & context, + const CommandProgram & commands, + const CommandProgramBindings & bindings, + PreparedCommandProgram & prepared) { + std::vector materialized; + Status status; + materialized.reserve(bindings.bindings().size()); + const std::unordered_map access_by_value = collect_graph_value_access(commands); + for (const CommandProgramBinding & binding : bindings.bindings()) { + if (!binding.requires_materialization()) { + materialized.push_back(binding); + continue; + } + const auto found_access = access_by_value.find(binding.value.value); + const GraphValueAccess access = + found_access != access_by_value.end() ? found_access->second : GraphValueAccess{}; + if (binding.weight && access.read && !access.write) { + HostWeightSource source; + source.host_data = binding.host_data; + source.identity = binding.identity; + source.generation = binding.generation; + source.capacity = binding.capacity; + source.offset = binding.offset; + source.length = binding.length; + HostWeightAcquireResult resident = + context.host_weights->acquire(context.device, context.stream, *context.host_transfers, source); + if (!resident.valid()) { + status.log("materialize host weight value %d failed", binding.value.value); + status.append(resident.status); + materialized.push_back(binding); + continue; + } + CommandProgramBinding device_binding = binding; + device_binding.buffer = resident.lease.buffer(); + device_binding.host_data = nullptr; + device_binding.offset = 0; + device_binding.capacity = binding.length; + materialized.push_back(device_binding); + prepared.resident_host_weights.push_back(std::move(resident.lease)); + continue; + } + + HostStagingBuffer staging; + Status allocation_status = allocate_host_staging_buffer(context.device, binding.length, staging); + if (!allocation_status.success()) { + status.log("allocate host staging for value %d failed", binding.value.value); + status.append(allocation_status); + materialized.push_back(binding); + continue; + } + staging.value = binding.value.value; + staging.host_data = static_cast(binding.host_data) + binding.offset; + staging.upload = access.read; + staging.download = access.write; + CommandProgramBinding device_binding = binding; + device_binding.buffer = staging.buffer; + device_binding.host_data = nullptr; + device_binding.offset = 0; + device_binding.capacity = binding.length; + materialized.push_back(device_binding); + prepared.host_staging.push_back(std::move(staging)); + } + return CommandProgramBindings::from_bindings(std::move(materialized), status); +} + +struct ProgramConstantImage { + ValueId value; + std::string name; + std::vector data; + bool read = false; +}; + +static Status collect_program_constant_images(const CommandProgram & commands, + std::vector & images) { + Status status; + std::unordered_map image_by_value; + for (const ConstantInitialization & initialization : commands.constant_initializations) { + const TransientAllocation * allocation = find_transient_allocation(commands.transients, initialization.value); + if (allocation == nullptr) { + status.log("constant initialization %s references missing transient value %d", initialization.name.c_str(), + initialization.value.value); + continue; + } + if (initialization.offset > allocation->size || + initialization.data.size() > allocation->size - initialization.offset) { + status.log("constant initialization %s is outside transient allocation length %zu", + initialization.name.c_str(), allocation->size); + continue; + } + + ProgramConstantImage * image = nullptr; + const auto found = image_by_value.find(initialization.value.value); + if (found == image_by_value.end()) { + ProgramConstantImage next; + next.value = initialization.value; + next.name = initialization.name; + next.data.resize(allocation->size); + image_by_value.emplace(initialization.value.value, images.size()); + images.push_back(std::move(next)); + image = &images.back(); + } else { + image = &images[found->second]; + } + + std::copy(initialization.data.begin(), initialization.data.end(), + image->data.begin() + initialization.offset); + } + return status; +} + +static Status validate_program_constant_access(const CommandProgram & commands, + std::vector & images) { + Status status; + std::unordered_map image_by_value; + for (size_t i = 0; i < images.size(); ++i) { + image_by_value.emplace(images[i].value.value, i); + } + + auto validate_command_list = [&](const std::vector & command_list) { + for (const Command & command : command_list) { + for (const CommandBinding & binding : command.bindings) { + const auto found = image_by_value.find(binding.value.value); + if (found == image_by_value.end()) { + continue; + } + ProgramConstantImage & image = images[found->second]; + if (resource_access_writes(binding.access)) { + status.log("constant initialization %s is written by command %u; prepared constant buffer copy " + "support is required", + image.name.c_str(), command.ordinal); + continue; + } + image.read = true; + } + } + }; + validate_command_list(commands.initialization_commands); + validate_command_list(commands.commands); + return status; +} + +static PreparedProgramConstantBuffer make_program_constant_buffer(ValueId value, + std::string name, + hrx_buffer_t buffer, + size_t size) { + PreparedProgramConstantBuffer result; + result.value = value; + result.name = std::move(name); + result.buffer = buffer; + result.size = size; + return result; +} + +static Status bind_prepared_command_list_program_constants( + const PreparedCommandProgram & prepared, + std::vector & prepared_commands) { + Status status; + for (PreparedCommand & command : prepared_commands) { + for (PreparedCommandBinding & binding : command.kernel.bindings) { + for (const PreparedProgramConstantBuffer & constant : prepared.program_constants) { + if (binding.binding.origin != CommandBindingOrigin::Transient || + binding.binding.value != constant.value) { + continue; + } + if (binding.binding.offset > constant.size || + binding.binding.length > constant.size - binding.binding.offset) { + status.log("%s is outside prepared constant %s length %zu", + format_command_binding(binding.binding).c_str(), constant.name.c_str(), constant.size); + continue; + } + binding.ref = { constant.buffer, binding.binding.offset, binding.binding.length }; + binding.binding.origin = CommandBindingOrigin::ProgramConstant; + } + } + } + return status; +} + +static Status prepare_program_constant_buffers(const CommandProgramExecutionContext & context, + const CommandProgram & commands, + PreparedCommandProgram & prepared) { + Status status; + if (commands.constant_initializations.empty()) { + return status; + } + + std::vector images; + status.append(collect_program_constant_images(commands, images)); + status.append(validate_program_constant_access(commands, images)); + if (!status.success()) { + return status; + } + if (images.empty()) { + return status; + } + if (context.device == nullptr) { + status.log("missing HRX device for prepared constants"); + return status; + } + if (context.stream == nullptr) { + status.log("missing HRX stream for prepared constants"); + return status; + } + + hrx_buffer_params_t params = { + HRX_MEMORY_TYPE_DEVICE_LOCAL, + HRX_MEMORY_ACCESS_ALL, + HRX_BUFFER_USAGE_DEFAULT, + 0, + }; + for (const ProgramConstantImage & image : images) { + if (!image.read) { + continue; + } + hrx_buffer_t buffer = nullptr; + if (ErrorResult error = take_status( + hrx_allocator_allocate_buffer(hrx_device_allocator(context.device), params, image.data.size(), + &buffer))) { + status.log("allocate prepared constant %s: %s", image.name.c_str(), error->c_str()); + continue; + } + if (context.host_transfers == nullptr) { + hrx_buffer_release(buffer); + status.log("prepared constant %s requires a host transfer manager", image.name.c_str()); + continue; + } + Status upload_status = + context.host_transfers->upload_synchronous(context.stream, image.data.data(), buffer, 0, image.data.size()); + if (!upload_status.success()) { + hrx_buffer_release(buffer); + status.log("upload prepared constant %s", image.name.c_str()); + status.append(upload_status); + continue; + } + prepared.program_constants.push_back( + make_program_constant_buffer(image.value, image.name, buffer, image.data.size())); + } + if (!status.success()) { + return status; + } + + status.append(bind_prepared_command_list_program_constants(prepared, prepared.initialization_commands)); + status.append(bind_prepared_command_list_program_constants(prepared, prepared.commands)); + return status; +} + +static Status rebind_prepared_host_staging(const CommandProgramBindings & bindings, PreparedCommandProgram & prepared) { + Status status; + for (HostStagingBuffer & staging : prepared.host_staging) { + const CommandProgramBinding * binding = bindings.find(ValueId(staging.value)); + if (binding == nullptr || binding->host_data == nullptr || binding->length != staging.length || + binding->offset > binding->capacity || binding->length > binding->capacity - binding->offset) { + status.log("live host binding does not match prepared value %d", staging.value); + continue; + } + staging.host_data = static_cast(binding->host_data) + binding->offset; + } + return status; +} + +static Status upload_prepared_host_staging(const CommandProgramExecutionContext & context, + const PreparedCommandProgram & prepared) { + Status status; + if (prepared.host_staging.empty()) { + return status; + } + if (context.host_transfers == nullptr) { + status.log("missing HRX host transfer manager"); + return status; + } + for (const HostStagingBuffer & staging : prepared.host_staging) { + if (!staging.upload) { + continue; + } + Status upload_status = + context.host_transfers->upload_async(context.stream, staging.host_data, staging.buffer, 0, staging.length); + status.append(upload_status); + } + return status; +} + +static Status download_prepared_host_staging(const CommandProgramExecutionContext & context, + const PreparedCommandProgram & prepared) { + Status status; + if (prepared.host_staging.empty()) { + return status; + } + if (context.host_transfers == nullptr) { + status.log("missing HRX host transfer manager"); + return status; + } + for (const HostStagingBuffer & staging : prepared.host_staging) { + if (!staging.download) { + continue; + } + Status download_status = context.host_transfers->download_synchronous( + context.stream, staging.buffer, 0, staging.host_data, staging.length); + status.append(download_status); + } + return status; +} + +static PreparedCommand make_prepared_command_shape(const ResolvedCommand & command) { + PreparedCommand prepared; + prepared.ordinal = command.ordinal; + prepared.kind = command.kind; + prepared.kernel.specialization = command.kernel; + prepared.kernel.bindings.reserve(command.bindings.size()); + for (const ResolvedCommandBinding & binding : command.bindings) { + prepared.kernel.bindings.push_back({ + binding.binding, + { binding.ref.buffer, binding.ref.offset, binding.ref.length }, + }); + } + return prepared; +} + +static Status prepare_kernel_command(const CommandProgramExecutionContext & context, + const ResolvedCommand & command, + PreparedCommand & prepared, + KernelExecutableRef & executable_ref) { + Status status; + const std::string command_context = format_resolved_command_context(command); + if (command.kind != CommandKind::Kernel) { + status.log("unsupported command kind in %s", command_context.c_str()); + return status; + } + Dispatch dispatch = build_dispatch(command); + + KernelResolveResult resolved = + resolve_kernel_definition(*context.corpus, context.target, dispatch.kernel.kernel_id); + if (!resolved.found()) { + status.log("%s: %s", command_context.c_str(), + format_kernel_resolve_error(resolved, dispatch.kernel.kernel_id).c_str()); + return status; + } + + prepared = make_prepared_command_shape(command); + executable_ref = context.kernel_executables->get_or_compile( + { context.device, context.target }, *resolved.definition, dispatch, prepared.kernel.constants); + if (!executable_ref.valid()) { + status.log("failed to prepare %s", command_context.c_str()); + return status; + } + return status; +} + +static bool execute_prepared_kernel_command(const CommandProgramExecutionContext & context, + const PreparedCommand & command) { + const std::string command_context = format_prepared_command_context(command); + if (command.kind != CommandKind::Kernel) { + GGML_LOG_ERROR("%s: unsupported command kind in %s\n", __func__, command_context.c_str()); + return false; + } + if (command.kernel.executable == nullptr) { + GGML_LOG_ERROR("%s: missing kernel executable for %s\n", __func__, command_context.c_str()); + return false; + } + + std::vector refs; + refs.reserve(command.kernel.bindings.size()); + for (const PreparedCommandBinding & binding : command.kernel.bindings) { + refs.push_back({ binding.ref.buffer, binding.ref.offset, binding.ref.length }); + } + + const KernelExecutable & executable = *command.kernel.executable; + hrx_dispatch_config_t config = { + { executable.launch.workgroup_count[0], executable.launch.workgroup_count[1], + executable.launch.workgroup_count[2] }, + { executable.launch.workgroup_size[0], executable.launch.workgroup_size[1], + executable.launch.workgroup_size[2] }, + executable.launch.subgroup_size, + }; + if (ErrorResult error = take_status(hrx_stream_dispatch( + context.stream, executable.executable, executable.export_ordinal, &config, command.kernel.constants.data(), + command.kernel.constants.size(), refs.data(), refs.size(), 0))) { + GGML_LOG_ERROR("%s: failed to execute %s: %s\n", __func__, command_context.c_str(), error->c_str()); + return false; + } + return true; +} + +static void prepare_command_list(const CommandProgramExecutionContext & context, + const std::vector & commands, + std::vector & prepared_commands, + std::vector & executable_refs, + Status & status) { + prepared_commands.reserve(commands.size()); + executable_refs.reserve(commands.size()); + for (const ResolvedCommand & command : commands) { + PreparedCommand prepared_command; + KernelExecutableRef executable_ref; + Status command_status = prepare_kernel_command(context, command, prepared_command, executable_ref); + if (command_status.success()) { + prepared_commands.push_back(std::move(prepared_command)); + executable_refs.push_back(std::move(executable_ref)); + } else { + status.append(command_status); + } + } +} + +static void materialize_command_list_executables(const CommandProgramExecutionContext & context, + std::vector & prepared_commands, + const std::vector & executable_refs, + Status & status) { + for (size_t i = 0; i < prepared_commands.size(); ++i) { + PreparedCommand & command = prepared_commands[i]; + command.kernel.executable = context.kernel_executables->materialize( + { context.device, context.target }, executable_refs[i], command.kernel.constants); + if (command.kernel.executable == nullptr) { + status.log("failed to prepare %s", format_prepared_command_context(command).c_str()); + } + } +} + +static bool bind_prepared_command_list_transients(const CommandProgram & commands, + const TransientArenaAllocationRef & transient_allocation, + std::vector & prepared_commands) { + for (PreparedCommand & command : prepared_commands) { + for (PreparedCommandBinding & binding : command.kernel.bindings) { + if (binding.binding.origin != CommandBindingOrigin::Transient) { + continue; + } + const TransientAllocation * allocation = + find_transient_allocation(commands.transients, binding.binding.value); + if (allocation == nullptr) { + GGML_LOG_ERROR("%s: %s has no transient allocation\n", __func__, + format_command_binding(binding.binding).c_str()); + return false; + } + if (binding.binding.offset > allocation->size || + binding.binding.length > allocation->size - binding.binding.offset || + commands.transients.arena_size > transient_allocation.capacity) { + GGML_LOG_ERROR("%s: %s is outside transient arena\n", __func__, + format_command_binding(binding.binding).c_str()); + return false; + } + binding.ref = { + transient_allocation.buffer, + allocation->arena_offset + binding.binding.offset, + binding.binding.length, + }; + } + } + return true; +} + +static bool execute_prepared_command_list(const CommandProgramExecutionContext & context, + const std::vector & commands) { + for (const PreparedCommand & command : commands) { + if (!execute_prepared_kernel_command(context, command)) { + return false; + } + } + return true; +} + +struct GraphDependencyChain { + hrx_graph_node_t last = nullptr; + + const hrx_graph_node_t * deps() const { return last == nullptr ? nullptr : &last; } + size_t dep_count() const { return last == nullptr ? 0 : 1; } + void update(hrx_graph_node_t node) { last = node; } +}; + +static Status record_completion_counter_fill(hrx_graph_t graph, + GraphDependencyChain & chain, + const CommandProgram & commands, + const TransientArenaAllocationRef & allocation) { + Status status; + if (commands.completion_counters.byte_count == 0) { + return status; + } + if (allocation.buffer == nullptr) { + status.log("command program has completion counters without a transient arena allocation"); + return status; + } + if (commands.completion_counters.arena_offset > commands.transients.arena_size || + commands.completion_counters.byte_count > + commands.transients.arena_size - commands.completion_counters.arena_offset) { + status.log("completion counter graph fill is outside transient arena length %zu", + commands.transients.arena_size); + return status; + } + + hrx_graph_fill_buffer_node_attrs_t attrs = { + { allocation.buffer, commands.completion_counters.arena_offset, commands.completion_counters.byte_count }, + 0, + sizeof(uint32_t), + }; + hrx_graph_node_t node = nullptr; + if (ErrorResult error = + take_status(hrx_graph_add_fill_buffer_node(graph, chain.deps(), chain.dep_count(), &attrs, &node))) { + status.log("record completion counter fill: %s", error->c_str()); + return status; + } + chain.update(node); + return status; +} + +static Status record_prepared_kernel_command(hrx_graph_t graph, + GraphDependencyChain & chain, + const PreparedCommand & command) { + Status status; + const std::string command_context = format_prepared_command_context(command); + if (command.kind != CommandKind::Kernel) { + status.log("unsupported command kind in %s", command_context.c_str()); + return status; + } + if (command.kernel.executable == nullptr) { + status.log("missing kernel executable for %s", command_context.c_str()); + return status; + } + + std::vector refs; + refs.reserve(command.kernel.bindings.size()); + for (const PreparedCommandBinding & binding : command.kernel.bindings) { + if (binding.ref.buffer == nullptr) { + status.log("%s has unbound buffer in %s", format_command_binding(binding.binding).c_str(), + command_context.c_str()); + continue; + } + refs.push_back({ binding.ref.buffer, binding.ref.offset, binding.ref.length }); + } + if (!status.success()) { + return status; + } + + const KernelExecutable & executable = *command.kernel.executable; + hrx_graph_kernel_node_attrs_t attrs = { + executable.executable, + executable.export_ordinal, + { + { executable.launch.workgroup_count[0], executable.launch.workgroup_count[1], + executable.launch.workgroup_count[2] }, + { executable.launch.workgroup_size[0], executable.launch.workgroup_size[1], + executable.launch.workgroup_size[2] }, + executable.launch.subgroup_size, + }, + command.kernel.constants.data(), + command.kernel.constants.size(), + refs.data(), + refs.size(), + 0, + }; + hrx_graph_node_t node = nullptr; + if (ErrorResult error = + take_status(hrx_graph_add_kernel_node(graph, chain.deps(), chain.dep_count(), &attrs, &node))) { + status.log("record %s: %s", command_context.c_str(), error->c_str()); + return status; + } + chain.update(node); + return status; +} + +static Status record_prepared_command_list(hrx_graph_t graph, + GraphDependencyChain & chain, + const std::vector & commands, + size_t & dispatch_count) { + Status status; + for (const PreparedCommand & command : commands) { + Status command_status = record_prepared_kernel_command(graph, chain, command); + if (!command_status.success()) { + status.append(command_status); + return status; + } + ++dispatch_count; + } + return status; +} + +static RecordedCommandGraph record_prepared_command_graph(const CommandProgramExecutionContext & context, + const CommandProgram & commands, + const PreparedCommandProgram & prepared, + const TransientArenaAllocationRef & allocation) { + RecordedCommandGraph recorded; + if (context.device == nullptr) { + recorded.status.log("missing HRX device for graph replay"); + return recorded; + } + + hrx_graph_t graph = nullptr; + if (ErrorResult error = take_status(hrx_graph_create(context.device, 0, &graph))) { + recorded.status.log("create HRX graph replay: %s", error->c_str()); + return recorded; + } + recorded.graph = graph; + + GraphDependencyChain chain; + recorded.status.append(record_completion_counter_fill(recorded.graph, chain, commands, allocation)); + if (!recorded.status.success()) { + return recorded; + } + recorded.status.append(record_prepared_command_list(recorded.graph, chain, prepared.initialization_commands, + recorded.dispatch_count)); + if (!recorded.status.success()) { + return recorded; + } + recorded.status.append(record_prepared_command_list(recorded.graph, chain, prepared.commands, + recorded.dispatch_count)); + if (!recorded.status.success()) { + return recorded; + } + + hrx_graph_exec_t exec = nullptr; + if (ErrorResult error = take_status(hrx_graph_instantiate(recorded.graph, 0, &exec))) { + recorded.status.log("instantiate HRX graph replay: %s", error->c_str()); + return recorded; + } + recorded.exec = exec; + recorded.bound_transient_arena_allocation_id = prepared.bound_transient_arena_allocation_id; + return recorded; +} + +} // namespace + +PreparedCommandProgram prepare_command_program(const CommandProgramExecutionContext & context, + const CommandProgram & commands, + const CommandProgramBindings & bindings) { + PreparedCommandProgram prepared; + prepared.status = command_program_metadata_context_valid(context); + if (!prepared.status.success()) { + return prepared; + } + + const VerificationResult verification = verify_command_program(commands, *context.corpus, context.target); + if (!verification.valid()) { + prepared.status.append(verification.status); + return prepared; + } + if (!bindings.valid()) { + prepared.status.append(bindings.status); + return prepared; + } + + TransientArenaAllocationRef transient_allocation; + prepared.status = ensure_transient_arena(context, commands, transient_allocation); + if (!prepared.status.success()) { + return prepared; + } + + prepared.status = command_program_preparation_context_valid(context); + if (!prepared.status.success()) { + return prepared; + } + + const CommandProgramBindings materialized_bindings = + materialize_host_bindings(context, commands, bindings, prepared); + if (!materialized_bindings.valid()) { + prepared.status.append(materialized_bindings.status); + return prepared; + } + + const TransientArenaAllocationRef * transient_allocation_ptr = + commands.transients.arena_size == 0 ? nullptr : &transient_allocation; + const ResolvedCommandProgram resolved = + resolve_command_program_bindings(commands, materialized_bindings, transient_allocation_ptr); + if (!resolved.valid()) { + prepared.status.append(resolved.status); + return prepared; + } + + std::vector initialization_executable_refs; + std::vector command_executable_refs; + prepare_command_list(context, resolved.initialization_commands, prepared.initialization_commands, + initialization_executable_refs, prepared.status); + prepare_command_list(context, resolved.commands, prepared.commands, command_executable_refs, prepared.status); + materialize_command_list_executables(context, prepared.initialization_commands, initialization_executable_refs, + prepared.status); + materialize_command_list_executables(context, prepared.commands, command_executable_refs, prepared.status); + prepared.bound_transient_arena_allocation_id = transient_allocation.allocation_id; + if (prepared.status.success()) { + prepared.status.append(prepare_program_constant_buffers(context, commands, prepared)); + } + return prepared; +} + +bool bind_prepared_command_program_transients(const CommandProgram & commands, + const TransientArenaAllocationRef & transient_allocation, + PreparedCommandProgram & prepared) { + if (!prepared.valid()) { + return false; + } + if (commands.transients.arena_size == 0) { + prepared.bound_transient_arena_allocation_id = kInvalidTransientArenaAllocationId; + return true; + } + if (transient_allocation.buffer == nullptr || + transient_allocation.allocation_id == kInvalidTransientArenaAllocationId) { + GGML_LOG_ERROR("%s: missing transient arena allocation\n", __func__); + return false; + } + if (prepared.bound_transient_arena_allocation_id == transient_allocation.allocation_id) { + return true; + } + if (!bind_prepared_command_list_transients(commands, transient_allocation, prepared.initialization_commands) || + !bind_prepared_command_list_transients(commands, transient_allocation, prepared.commands)) { + return false; + } + prepared.bound_transient_arena_allocation_id = transient_allocation.allocation_id; + return true; +} + +bool bind_and_execute_prepared_command_program(const CommandProgramExecutionContext & context, + const CommandProgram & commands, + const CommandProgramBindings & bindings, + PreparedCommandProgram & prepared) { + if (!prepared.valid()) { + return execute_prepared_command_program(context, prepared); + } + Status rebind_status = rebind_prepared_host_staging(bindings, prepared); + if (!rebind_status.success()) { + GGML_LOG_ERROR("%s: %s\n", __func__, status_first_error(rebind_status)); + return false; + } + if (commands.transients.arena_size == 0) { + Status status = initialize_command_program_constants(context, commands, {}, prepared); + if (!status.success()) { + GGML_LOG_ERROR("%s: %s\n", __func__, status_first_error(status)); + return false; + } + status = initialize_command_program_completion_counters(context, commands, {}); + if (!status.success()) { + GGML_LOG_ERROR("%s: %s\n", __func__, status_first_error(status)); + return false; + } + return bind_prepared_command_program_transients(commands, {}, prepared) && + execute_prepared_command_program(context, prepared); + } + + Status status = command_program_transient_context_valid(context, commands); + if (!status.success()) { + GGML_LOG_ERROR("%s: %s\n", __func__, status_first_error(status)); + return false; + } + + TransientArena::AllocationLease lease = context.transient_arena->acquire_allocation_lease(); + status = lease.ensure_capacity(context.device, context.stream, commands.transients.arena_size); + if (!status.success()) { + GGML_LOG_ERROR("%s: %s\n", __func__, status_first_error(status)); + return false; + } + if (!bind_prepared_command_program_transients(commands, lease.current_allocation(), prepared)) { + return false; + } + status = initialize_command_program_constants(context, commands, lease.current_allocation(), prepared); + if (!status.success()) { + GGML_LOG_ERROR("%s: %s\n", __func__, status_first_error(status)); + return false; + } + status = initialize_command_program_completion_counters(context, commands, lease.current_allocation()); + if (!status.success()) { + GGML_LOG_ERROR("%s: %s\n", __func__, status_first_error(status)); + return false; + } + return execute_prepared_command_program(context, prepared); +} + +RecordedCommandGraphExecutionResult bind_and_launch_recorded_command_graph( + const CommandProgramExecutionContext & context, + const CommandProgram & commands, + const CommandProgramBindings & bindings, + PreparedCommandProgram & prepared, + RecordedCommandGraph & recorded) { + RecordedCommandGraphExecutionResult result; + result.event = HrxGraphReplayEvent::Ineligible; + + if (!prepared.valid()) { + result.status.append(prepared.status); + if (result.status.success()) { + result.status.log("invalid prepared command program"); + } + return result; + } + if (!prepared_execution_context_valid(context)) { + result.status.log("missing HRX stream"); + result.event = HrxGraphReplayEvent::BuildFailed; + return result; + } + + Status rebind_status = rebind_prepared_host_staging(bindings, prepared); + if (!rebind_status.success()) { + result.status.append(rebind_status); + result.event = HrxGraphReplayEvent::BuildFailed; + return result; + } + + TransientArenaAllocationRef transient_allocation; + TransientArena::AllocationLease lease; + if (commands.transients.arena_size == 0) { + if (!bind_prepared_command_program_transients(commands, {}, prepared)) { + result.status.log("bind transient-free prepared command program failed"); + result.event = HrxGraphReplayEvent::BuildFailed; + return result; + } + } else { + Status status = command_program_transient_context_valid(context, commands); + if (!status.success()) { + result.status.append(status); + result.event = HrxGraphReplayEvent::BuildFailed; + return result; + } + lease = context.transient_arena->acquire_allocation_lease(); + status = lease.ensure_capacity(context.device, context.stream, commands.transients.arena_size); + if (!status.success()) { + result.status.append(status); + result.event = HrxGraphReplayEvent::BuildFailed; + return result; + } + transient_allocation = lease.current_allocation(); + if (!bind_prepared_command_program_transients(commands, transient_allocation, prepared)) { + result.status.log("bind prepared command program transients for graph replay failed"); + result.event = HrxGraphReplayEvent::BuildFailed; + return result; + } + } + + const bool had_recorded = recorded.valid(); + result.transient_allocation_changed = + had_recorded && recorded.bound_transient_arena_allocation_id != prepared.bound_transient_arena_allocation_id; + if (!had_recorded || result.transient_allocation_changed) { + result.event = + result.transient_allocation_changed ? HrxGraphReplayEvent::RebuildTransient : HrxGraphReplayEvent::MissBuild; + const uint64_t build_start_ns = hrx_graph_replay_now_ns(); + RecordedCommandGraph rebuilt = record_prepared_command_graph(context, commands, prepared, transient_allocation); + result.build_ns = hrx_graph_replay_now_ns() - build_start_ns; + if (!rebuilt.valid()) { + result.status.append(rebuilt.status); + result.event = HrxGraphReplayEvent::BuildFailed; + return result; + } + recorded = std::move(rebuilt); + } else { + result.event = HrxGraphReplayEvent::Hit; + } + + const uint64_t launch_start_ns = hrx_graph_replay_now_ns(); + Status upload_status = upload_prepared_host_staging(context, prepared); + if (!upload_status.success()) { + result.launch_ns = hrx_graph_replay_now_ns() - launch_start_ns; + result.status.append(upload_status); + result.event = HrxGraphReplayEvent::LaunchFailed; + return result; + } + if (ErrorResult error = take_status(hrx_graph_exec_launch(recorded.exec, context.stream))) { + result.launch_ns = hrx_graph_replay_now_ns() - launch_start_ns; + result.status.log("launch HRX graph replay: %s", error->c_str()); + result.event = HrxGraphReplayEvent::LaunchFailed; + return result; + } + Status download_status = download_prepared_host_staging(context, prepared); + if (!download_status.success()) { + result.launch_ns = hrx_graph_replay_now_ns() - launch_start_ns; + result.status.append(download_status); + result.event = HrxGraphReplayEvent::LaunchFailed; + return result; + } + result.launch_ns = hrx_graph_replay_now_ns() - launch_start_ns; + result.dispatch_count = recorded.dispatch_count; + result.success = true; + return result; +} + +bool execute_prepared_command_program(const CommandProgramExecutionContext & context, + const PreparedCommandProgram & commands) { + if (!commands.valid()) { + GGML_LOG_ERROR("%s: invalid HRX prepared command program: %s\n", __func__, status_first_error(commands.status)); + return false; + } + if (!prepared_execution_context_valid(context)) { + return false; + } + Status upload_status = upload_prepared_host_staging(context, commands); + if (!upload_status.success()) { + GGML_LOG_ERROR("%s: %s\n", __func__, status_first_error(upload_status)); + return false; + } + if (!execute_prepared_command_list(context, commands.initialization_commands) || + !execute_prepared_command_list(context, commands.commands)) { + return false; + } + Status download_status = download_prepared_host_staging(context, commands); + if (!download_status.success()) { + GGML_LOG_ERROR("%s: %s\n", __func__, status_first_error(download_status)); + return false; + } + return true; +} + +bool execute_command_program(const CommandProgramExecutionContext & context, + const CommandProgram & commands, + const CommandProgramBindings & bindings) { + PreparedCommandProgram prepared = prepare_command_program(context, commands, bindings); + return bind_and_execute_prepared_command_program(context, commands, bindings, prepared); +} + +} // namespace ggml::hrx diff --git a/ggml/src/ggml-hrx/runtime/command-program-executor.h b/ggml/src/ggml-hrx/runtime/command-program-executor.h new file mode 100644 index 000000000000..bfc8689b53b8 --- /dev/null +++ b/ggml/src/ggml-hrx/runtime/command-program-executor.h @@ -0,0 +1,144 @@ +#pragma once + +#include "dispatch/command-program-bindings.h" +#include "dispatch/command-program-resolver.h" +#include "dispatch/command-program.h" +#include "kernel-corpus/kernel-corpus.h" +#include "runtime/graph-replay.h" +#include "runtime/host-memory.h" + +#include +#include +#include +#include + +typedef struct hrx_device_s * hrx_device_t; +typedef struct hrx_stream_s * hrx_stream_t; +typedef struct hrx_buffer_s * hrx_buffer_t; +typedef struct hrx_graph_s * hrx_graph_t; +typedef struct hrx_graph_exec_s * hrx_graph_exec_t; +struct ggml_hrx_loom_jit_amdgpu; + +namespace ggml::hrx { + +class KernelExecutableCache; +struct KernelExecutable; +class TransientArena; + +struct CommandProgramExecutionContext { + hrx_device_t device = nullptr; + hrx_stream_t stream = nullptr; + const char * target = nullptr; + const KernelCorpus * corpus = nullptr; + KernelExecutableCache * kernel_executables = nullptr; + TransientArena * transient_arena = nullptr; + HostTransferManager * host_transfers = nullptr; + HostWeightCache * host_weights = nullptr; +}; + +struct PreparedCommandBinding { + CommandBinding binding; + ResolvedBufferRef ref; +}; + +struct PreparedKernelCommand { + KernelSpecialization specialization; + std::shared_ptr executable; + std::vector constants; + std::vector bindings; +}; + +struct PreparedCommand { + uint32_t ordinal = 0; + CommandKind kind = CommandKind::Invalid; + PreparedKernelCommand kernel; +}; + +struct PreparedProgramConstantBuffer { + ValueId value; + std::string name; + hrx_buffer_t buffer = nullptr; + size_t size = 0; + + PreparedProgramConstantBuffer() = default; + ~PreparedProgramConstantBuffer(); + + PreparedProgramConstantBuffer(PreparedProgramConstantBuffer && other) noexcept; + PreparedProgramConstantBuffer & operator=(PreparedProgramConstantBuffer && other) noexcept; + + PreparedProgramConstantBuffer(const PreparedProgramConstantBuffer &) = delete; + PreparedProgramConstantBuffer & operator=(const PreparedProgramConstantBuffer &) = delete; +}; + +struct PreparedCommandProgram { + std::vector initialization_commands; + std::vector commands; + std::vector host_staging; + std::vector resident_host_weights; + std::vector program_constants; + Status status; + uint64_t bound_transient_arena_allocation_id = kInvalidTransientArenaAllocationId; + + bool valid() const { return status.success(); } +}; + +struct RecordedCommandGraph { + hrx_graph_t graph = nullptr; + hrx_graph_exec_t exec = nullptr; + uint64_t bound_transient_arena_allocation_id = kInvalidTransientArenaAllocationId; + size_t dispatch_count = 0; + Status status; + + RecordedCommandGraph() = default; + ~RecordedCommandGraph(); + + RecordedCommandGraph(RecordedCommandGraph && other) noexcept; + RecordedCommandGraph & operator=(RecordedCommandGraph && other) noexcept; + + RecordedCommandGraph(const RecordedCommandGraph &) = delete; + RecordedCommandGraph & operator=(const RecordedCommandGraph &) = delete; + + bool valid() const { return status.success() && exec != nullptr; } +}; + +struct RecordedCommandGraphExecutionResult { + bool success = false; + Status status; + HrxGraphReplayEvent event = HrxGraphReplayEvent::Disabled; + std::string ineligible_reason; + size_t dispatch_count = 0; + bool transient_allocation_changed = false; + uint64_t build_ns = 0; + uint64_t launch_ns = 0; + + uint64_t total_ns() const { return build_ns + launch_ns; } +}; + +PreparedCommandProgram prepare_command_program(const CommandProgramExecutionContext & context, + const CommandProgram & commands, + const CommandProgramBindings & bindings); + +bool execute_prepared_command_program(const CommandProgramExecutionContext & context, + const PreparedCommandProgram & commands); + +bool bind_prepared_command_program_transients(const CommandProgram & commands, + const TransientArenaAllocationRef & transient_allocation, + PreparedCommandProgram & prepared); + +bool bind_and_execute_prepared_command_program(const CommandProgramExecutionContext & context, + const CommandProgram & commands, + const CommandProgramBindings & bindings, + PreparedCommandProgram & prepared); + +RecordedCommandGraphExecutionResult bind_and_launch_recorded_command_graph( + const CommandProgramExecutionContext & context, + const CommandProgram & commands, + const CommandProgramBindings & bindings, + PreparedCommandProgram & prepared, + RecordedCommandGraph & recorded); + +bool execute_command_program(const CommandProgramExecutionContext & context, + const CommandProgram & commands, + const CommandProgramBindings & bindings); + +} // namespace ggml::hrx diff --git a/ggml/src/ggml-hrx/runtime/graph-executor.cpp b/ggml/src/ggml-hrx/runtime/graph-executor.cpp new file mode 100644 index 000000000000..9ac434ef0a47 --- /dev/null +++ b/ggml/src/ggml-hrx/runtime/graph-executor.cpp @@ -0,0 +1,138 @@ +#include "graph-executor.h" + +#include "backend-buffer-binding.h" +#include "ggml-impl.h" +#include "runtime/kernel-executable-cache.h" +#include "runtime/prepared-command-program-cache.h" +#include "runtime/transient-arena.h" + +#include +#include + +namespace ggml::hrx { + +GraphExecutor::GraphExecutor(ggml_backend_hrx_context & context) : context_(context) {} + +Status GraphExecutor::context_valid_for_graph_programs() const { + Status status; + if (context_.device == nullptr) { + status.log("missing HRX device context"); + } else if (context_.device->architecture.empty()) { + status.log("missing HRX target"); + } + return status; +} + +Status GraphExecutor::context_valid_for_execution() const { + return context_valid_for_graph_programs(); +} + +GraphSupportResult GraphExecutor::can_execute(const ggml_cgraph & graph) const { + GraphSupportResult result; + if (graph.n_nodes == 0) { + result.supported = true; + return result; + } + result.status = context_valid_for_graph_programs(); + if (!result.status.success()) { + return result; + } + const KernelCorpus & corpus = get_qwen_kernel_corpus(); + const GraphProgramSupportResult support = + context_.graph_programs.check_support(graph, corpus, context_.device->architecture); + result.supported = support.supported; + result.status.append(support.status); + return result; +} + +CommandProgramBindings GraphExecutor::bind_external_value_buffers(const GraphProgramMatch & match) const { + std::vector bindings; + Status status; + bindings.reserve(match.external_bindings.size()); + for (const GraphProgramExternalBinding & external : match.external_bindings) { + ValueBufferBinding value_binding; + CommandProgramBinding binding; + binding.value = external.value; + if (ggml_backend_hrx_resolve_value_buffer(external.tensor, value_binding)) { + binding.buffer = value_binding.buffer; + binding.host_data = value_binding.host_data; + binding.offset = value_binding.offset; + binding.length = value_binding.length; + binding.identity = value_binding.identity; + binding.generation = value_binding.generation; + binding.capacity = value_binding.capacity; + binding.weight = value_binding.weight; + } else { + status.log("external value %d is not bound", external.value.value); + } + bindings.push_back(binding); + } + return CommandProgramBindings::from_bindings(std::move(bindings), status); +} + +GraphExecutionResult GraphExecutor::execute(const ggml_cgraph & graph) const { + GraphExecutionResult result; + if (graph.n_nodes == 0) { + result.code = GGML_STATUS_SUCCESS; + return result; + } + result.status = context_valid_for_execution(); + if (!result.status.success()) { + return result; + } + + const KernelCorpus & corpus = get_qwen_kernel_corpus(); + GraphProgramLookup lookup = context_.graph_programs.get_or_build(graph, corpus, context_.device->architecture); + if (!lookup.valid()) { + result.status.append(lookup.status); + result.status.append(lookup.match.status); + if (result.status.success()) { + result.status.log("build HRX graph program failed"); + } + return result; + } + + const bool use_graph_prepared = + !lookup.program->has_prepared_program() || lookup.program->can_use_prepared_fast_path(graph); + GraphProgramMatch binding_match = std::move(lookup.match); + if (use_graph_prepared && lookup.program->has_prepared_program()) { + binding_match = lookup.program->match_host_staging_graph(graph); + if (!binding_match.valid()) { + result.status.append(binding_match.status); + return result; + } + } + + CommandProgramBindings bindings = bind_external_value_buffers(binding_match); + if (!bindings.valid()) { + result.status.append(bindings.status); + return result; + } + const CommandProgramExecutionContext execution_context = { + context_.device->device, + context_.stream, + context_.device->architecture.c_str(), + &corpus, + &context_.kernel_executables, + &context_.transient_arena, + &context_.host_transfers, + &context_.host_weights, + }; + const PreparedCommandProgramCacheExecutionResult execution = + use_graph_prepared ? lookup.program->execute_with_result(execution_context, bindings) : + context_.prepared_programs.execute_with_result(execution_context, lookup.program->uid(), + lookup.program->command_shape(), + lookup.program->commands(), bindings); + if (!execution.success) { + result.status.append(execution.status); + if (result.status.success()) { + result.status.log("execute HRX command program failed"); + } + return result; + } + + result.code = GGML_STATUS_SUCCESS; + return result; +} + +} // namespace ggml::hrx diff --git a/ggml/src/ggml-hrx/runtime/graph-executor.h b/ggml/src/ggml-hrx/runtime/graph-executor.h new file mode 100644 index 000000000000..a2759ea71875 --- /dev/null +++ b/ggml/src/ggml-hrx/runtime/graph-executor.h @@ -0,0 +1,43 @@ +#pragma once + +#include "backend-context.h" +#include "dispatch/command-program-bindings.h" +#include "ggml.h" +#include "runtime/graph-program-cache.h" +#include "status.h" + +struct ggml_cgraph; + +namespace ggml::hrx { + +struct GraphSupportResult { + bool supported = false; + Status status; + + bool success() const { return supported && status.success(); } +}; + +struct GraphExecutionResult { + enum ggml_status code = GGML_STATUS_FAILED; + Status status; + + bool success() const { return code == GGML_STATUS_SUCCESS && status.success(); } +}; + +class GraphExecutor { + public: + explicit GraphExecutor(ggml_backend_hrx_context & context); + + GraphSupportResult can_execute(const ggml_cgraph & graph) const; + GraphExecutionResult execute(const ggml_cgraph & graph) const; + + private: + Status context_valid_for_graph_programs() const; + Status context_valid_for_execution() const; + + CommandProgramBindings bind_external_value_buffers(const GraphProgramMatch & match) const; + + ggml_backend_hrx_context & context_; +}; + +} // namespace ggml::hrx diff --git a/ggml/src/ggml-hrx/runtime/graph-program-cache.cpp b/ggml/src/ggml-hrx/runtime/graph-program-cache.cpp new file mode 100644 index 000000000000..b22f0818ceb1 --- /dev/null +++ b/ggml/src/ggml-hrx/runtime/graph-program-cache.cpp @@ -0,0 +1,610 @@ +#include "graph-program-cache.h" + +#include "dispatch/dispatch-scheduler.h" +#include "ggml-impl.h" +#include "ggml.h" + +#include +#include +#include +#include + +namespace ggml::hrx { +namespace { + +static bool tensor_metadata_matches(const Value & value, const ggml_tensor * tensor) { + if (tensor == nullptr || value.type != tensor->type || value.element_count != ggml_nelements(tensor) || + value.byte_count != ggml_nbytes(tensor) || value.contiguous != ggml_is_contiguous(tensor)) { + return false; + } + const bool tensor_alias = tensor->view_src != nullptr; + const bool value_alias = value.alias_source.value >= 0; + if (tensor_alias && (!value_alias || value.storage_offset != tensor->view_offs)) { + return false; + } + for (int i = 0; i < GGML_MAX_DIMS; ++i) { + if (value.ne[i] != tensor->ne[i] || value.nb[i] != tensor->nb[i]) { + return false; + } + } + return true; +} + +static bool graph_node_params_match(const GraphNode & cached_node, const ggml_tensor * current_node) { + if (current_node == nullptr) { + return false; + } + return op_params_equivalent(cached_node.op, cached_node.params, *current_node); +} + +static bool environment_flag_enabled(const char * name) { + const char * value = std::getenv(name); + return value != nullptr && value[0] != '\0' && value[0] != '0'; +} + +static void apply_graph_replay_result(PreparedCommandProgramCacheExecutionResult & result, + const RecordedCommandGraphExecutionResult & replay) { + result.graph_replay_event = replay.event; + result.graph_replay_ineligible_reason = replay.ineligible_reason; + result.graph_replay_build_ns = replay.build_ns; + result.graph_replay_launch_ns = replay.launch_ns; + result.graph_replay_total_ns = replay.total_ns(); + result.graph_replay_dispatches = replay.dispatch_count; + result.graph_replay_transient_allocation_changed = replay.transient_allocation_changed; +} + +static bool graph_replay_should_fallback(HrxGraphReplayEvent event) { + return event == HrxGraphReplayEvent::Ineligible || event == HrxGraphReplayEvent::BuildFailed; +} + +static Status bind_current_value(const ValueMap & values, + ValueId expected, + const ggml_tensor * tensor, + std::vector & tensor_by_value, + std::unordered_map & value_by_tensor, + const char * role, + size_t node_index) { + Status status; + const Value * value = values.find(expected); + if (value == nullptr || expected.value < 0 || static_cast(expected.value) >= tensor_by_value.size()) { + status.log("node %zu %s references missing cached value %d", node_index, role, expected.value); + return status; + } + if (tensor == nullptr) { + status.log("node %zu %s value %d maps to a null tensor", node_index, role, expected.value); + return status; + } + + const ggml_tensor * existing_tensor = tensor_by_value[static_cast(expected.value)]; + if (existing_tensor != nullptr && existing_tensor != tensor) { + status.log("node %zu %s value %d maps to multiple current tensors", node_index, role, expected.value); + return status; + } + + const auto existing_value = value_by_tensor.find(tensor); + if (existing_value != value_by_tensor.end() && existing_value->second != expected.value) { + status.log("node %zu %s tensor maps to cached values %d and %d", node_index, role, existing_value->second, + expected.value); + return status; + } + + if (existing_tensor == nullptr && existing_value == value_by_tensor.end() && + !tensor_metadata_matches(*value, tensor)) { + status.log("node %zu %s value %d metadata does not match current tensor", node_index, role, expected.value); + return status; + } + + tensor_by_value[static_cast(expected.value)] = tensor; + value_by_tensor.emplace(tensor, expected.value); + return status; +} + +static std::string command_program_shape_key(const CommandProgram & commands) { + std::ostringstream out; + out << "hrx-command-program-v1|commands=" << commands.commands.size(); + for (const Command & command : commands.commands) { + out << "|ordinal=" << command.ordinal << "|kind=" << static_cast(command.kind) + << "|kernel=" << command.kernel.kernel_id; + for (const auto & parameter : command.kernel.integer_parameters) { + out << "|ip:" << parameter.first << '=' << parameter.second; + } + for (const auto & parameter : command.kernel.compile_parameters) { + out << "|cp:" << parameter.first << '=' << parameter.second; + } + out << "|bindings=" << command.bindings.size(); + for (const CommandBinding & binding : command.bindings) { + out << "|b:" << binding.name << ':' << binding.value.value << ':' << static_cast(binding.origin) << ':' + << binding.offset << ':' << binding.length << ':' << static_cast(binding.access); + } + out << "|deps=" << command.dependencies.size(); + for (const uint32_t dependency : command.dependencies) { + out << ':' << dependency; + } + } + out << "|transients=" << commands.transients.allocations.size() << "|arena=" << commands.transients.arena_size + << "|arena_alignment=" << commands.transients.arena_alignment; + for (const TransientAllocation & allocation : commands.transients.allocations) { + out << "|t:" << allocation.value.value << ':' << allocation.arena_offset << ':' << allocation.size << ':' + << allocation.alignment; + } + return out.str(); +} + +} // namespace + +GraphProgram::GraphProgram(uint64_t uid, + std::string target, + std::unique_ptr graph, + std::unique_ptr commands, + std::string command_shape) : + uid_(uid), + target_(std::move(target)), + graph_(std::move(graph)), + commands_(std::move(commands)), + command_shape_(std::move(command_shape)) {} + +const GraphProgramExternalSlot * GraphProgram::find_external_slot(ValueId value) const { + const auto found = external_slot_by_value_.find(value.value); + if (found == external_slot_by_value_.end() || found->second >= external_slots_.size()) { + return nullptr; + } + return &external_slots_[found->second]; +} + +const ggml_tensor * GraphProgram::resolve_external_slot(const ggml_cgraph & graph, + const GraphProgramExternalSlot & slot, + Status & status) const { + if (slot.node_index >= static_cast(graph.n_nodes)) { + status.log("external value %d references node slot %zu but current graph has %d nodes", slot.value.value, + slot.node_index, graph.n_nodes); + return nullptr; + } + const ggml_tensor * node = graph.nodes[slot.node_index]; + if (node == nullptr) { + status.log("external value %d references null node slot %zu", slot.value.value, slot.node_index); + return nullptr; + } + if (slot.kind == GraphProgramExternalSlotKind::Node) { + return node; + } + if (slot.source_index < 0 || slot.source_index >= GGML_MAX_SRC) { + status.log("external value %d references invalid source slot %d", slot.value.value, slot.source_index); + return nullptr; + } + const ggml_tensor * source = node->src[slot.source_index]; + if (source == nullptr) { + status.log("external value %d references null source slot %zu:%d", slot.value.value, slot.node_index, + slot.source_index); + return nullptr; + } + return source; +} + +GraphProgramMatch GraphProgram::match_trusted_graph(const ggml_cgraph & current_graph, bool bind_external) const { + GraphProgramMatch result; + if (graph_ == nullptr || commands_ == nullptr) { + result.status.log("missing cached HRX graph program"); + return result; + } + if (graph_->nodes().size() != static_cast(current_graph.n_nodes)) { + result.status.log("cached graph has %zu nodes but current graph has %d", graph_->nodes().size(), + current_graph.n_nodes); + return result; + } + if (!graph_->nodes().empty()) { + const ggml_tensor * first = current_graph.nodes[0]; + const ggml_tensor * last = current_graph.nodes[current_graph.n_nodes - 1]; + if (first == nullptr || last == nullptr) { + result.status.log("current graph has null sentinel nodes"); + return result; + } + if (first->op != graph_->nodes().front().op || last->op != graph_->nodes().back().op) { + result.status.log("current graph sentinel ops do not match cached HRX graph"); + return result; + } + } + if (!bind_external) { + return result; + } + result.external_bindings.reserve(external_slots_.size()); + for (const GraphProgramExternalSlot & slot : external_slots_) { + const ggml_tensor * tensor = resolve_external_slot(current_graph, slot, result.status); + if (tensor == nullptr) { + return result; + } + result.external_bindings.push_back({ slot.value, tensor }); + } + return result; +} + +GraphProgramMatch GraphProgram::match_host_staging_graph(const ggml_cgraph & current_graph) const { + GraphProgramMatch result; + std::lock_guard lock(prepared_mutex_); + if (!has_prepared_) { + result.status.log("missing prepared HRX command program"); + return result; + } + result.external_bindings.reserve(prepared_.host_staging.size()); + for (const HostStagingBuffer & staging : prepared_.host_staging) { + const GraphProgramExternalSlot * slot = find_external_slot(ValueId(staging.value)); + if (slot == nullptr) { + result.status.log("prepared host staging value %d has no external graph slot", staging.value); + return result; + } + const ggml_tensor * tensor = resolve_external_slot(current_graph, *slot, result.status); + if (tensor == nullptr) { + return result; + } + result.external_bindings.push_back({ ValueId(staging.value), tensor }); + } + return result; +} + +Status GraphProgram::capture_external_slots(const ggml_cgraph & graph, const GraphProgramMatch & match) { + Status status; + external_slots_.clear(); + external_slot_by_value_.clear(); + fast_path_nodes_ = graph.nodes; + external_slots_.reserve(match.external_bindings.size()); + for (const GraphProgramExternalBinding & binding : match.external_bindings) { + GraphProgramExternalSlot slot; + slot.value = binding.value; + bool found = false; + for (int i = 0; i < graph.n_nodes && !found; ++i) { + const ggml_tensor * node = graph.nodes[i]; + if (node == nullptr) { + continue; + } + if (node == binding.tensor) { + slot.kind = GraphProgramExternalSlotKind::Node; + slot.node_index = static_cast(i); + found = true; + break; + } + for (int j = 0; j < GGML_MAX_SRC; ++j) { + if (node->src[j] == binding.tensor) { + slot.kind = GraphProgramExternalSlotKind::Source; + slot.node_index = static_cast(i); + slot.source_index = j; + found = true; + break; + } + } + } + if (!found) { + status.log("external value %d has no current graph slot", binding.value.value); + continue; + } + external_slot_by_value_[binding.value.value] = external_slots_.size(); + external_slots_.push_back(slot); + } + return status; +} + +bool GraphProgram::has_prepared_program() const { + std::lock_guard lock(prepared_mutex_); + return has_prepared_; +} + +bool GraphProgram::can_use_prepared_fast_path(const ggml_cgraph & graph) const { + return graph.nodes == fast_path_nodes_; +} + +PreparedCommandProgramCacheStats GraphProgram::prepared_stats() const { + std::lock_guard lock(prepared_mutex_); + return prepared_stats_; +} + +PreparedCommandProgramCacheExecutionResult GraphProgram::execute_with_result( + const CommandProgramExecutionContext & context, + const CommandProgramBindings & bindings) { + PreparedCommandProgramCacheExecutionResult result; + if (commands_ == nullptr || !commands_->valid() || !bindings.valid()) { + result.graph_replay_event = HrxGraphReplayEvent::Ineligible; + result.graph_replay_ineligible_reason = "invalid_graph_program"; + if (commands_ == nullptr) { + result.status.log("missing cached HRX command program"); + } + result.status.append(bindings.status); + return result; + } + + std::lock_guard lock(prepared_mutex_); + if (!has_prepared_) { + prepared_ = prepare_command_program(context, *commands_, bindings); + if (!prepared_.valid()) { + result.status.append(prepared_.status); + return result; + } + has_prepared_ = true; + ++prepared_stats_.builds; + } else { + ++prepared_stats_.hits; + } + + const RecordedCommandGraphExecutionResult replay = + bind_and_launch_recorded_command_graph(context, *commands_, bindings, prepared_, recorded_); + apply_graph_replay_result(result, replay); + if (replay.success) { + result.success = true; + return result; + } + if (!graph_replay_should_fallback(replay.event)) { + result.status.append(replay.status); + if (result.status.success()) { + result.status.log("execute cached HRX graph replay failed"); + } + return result; + } + result.success = bind_and_execute_prepared_command_program(context, *commands_, bindings, prepared_); + if (!result.success) { + result.status.log("execute cached HRX command program failed"); + } + return result; +} + +GraphProgramMatch GraphProgram::match_current_graph(const ggml_cgraph & current_graph) const { + GraphProgramMatch result; + if (graph_ == nullptr) { + result.status.log("missing cached HRX graph"); + return result; + } + if (commands_ == nullptr) { + result.status.log("missing cached HRX command program"); + return result; + } + if (graph_->nodes().size() != static_cast(current_graph.n_nodes)) { + result.status.log("cached graph has %zu nodes but current graph has %d", graph_->nodes().size(), + current_graph.n_nodes); + return result; + } + + const ValueMap & values = graph_->values(); + std::vector tensor_by_value(values.size(), nullptr); + std::unordered_map value_by_tensor; + + for (size_t node_index = 0; node_index < graph_->nodes().size(); ++node_index) { + const GraphNode & cached_node = graph_->nodes()[node_index]; + const ggml_tensor * current_node = current_graph.nodes[node_index]; + if (current_node == nullptr) { + result.status.log("current graph node %zu is null", node_index); + return result; + } + if (cached_node.op != current_node->op) { + result.status.log("node %zu cached op %s does not match current op %s", node_index, + ggml_op_name(cached_node.op), ggml_op_name(current_node->op)); + return result; + } + if (!graph_node_params_match(cached_node, current_node)) { + result.status.log("node %zu cached op params do not match current graph", node_index); + return result; + } + + size_t input_index = 0; + for (const ggml_tensor * source : current_node->src) { + if (source == nullptr) { + continue; + } + if (input_index >= cached_node.inputs.size()) { + result.status.log("node %zu has more inputs than the cached graph", node_index); + return result; + } + Status status = bind_current_value(values, cached_node.inputs[input_index], source, tensor_by_value, + value_by_tensor, "input", node_index); + if (!status.success()) { + result.status.append(status); + return result; + } + ++input_index; + } + if (input_index != cached_node.inputs.size()) { + result.status.log("node %zu has %zu inputs but cached graph has %zu", node_index, input_index, + cached_node.inputs.size()); + return result; + } + Status status = bind_current_value(values, cached_node.output, current_node, tensor_by_value, value_by_tensor, + "output", node_index); + if (!status.success()) { + result.status.append(status); + return result; + } + } + + for (const ValueId id : values.external_value_ids()) { + if (id.value < 0 || static_cast(id.value) >= tensor_by_value.size() || + tensor_by_value[static_cast(id.value)] == nullptr) { + result.status.log("external value %d is missing from the current graph", id.value); + return result; + } + result.external_bindings.push_back({ id, tensor_by_value[static_cast(id.value)] }); + } + return result; +} + +bool GraphProgramCache::can_execute(const ggml_cgraph & graph, + const KernelCorpus & corpus, + const std::string & target) const { + return check_support(graph, corpus, target).supported; +} + +GraphProgramSupportResult GraphProgramCache::check_support(const ggml_cgraph & graph, + const KernelCorpus & corpus, + const std::string & target) const { + GraphProgramSupportResult result; + if (graph.n_nodes == 0) { + result.supported = true; + return result; + } + GraphImportResult imported = import_ggml_graph(graph); + if (!imported.valid()) { + result.status.append(imported.status); + return result; + } + std::unique_ptr program = + build_program_from_imported(graph.uid, std::move(imported.graph), corpus, target, result.status); + result.supported = program != nullptr && result.status.success(); + return result; +} + +GraphProgramLookup GraphProgramCache::build_from_imported(const ggml_cgraph & graph, + Graph && imported_graph, + const KernelCorpus & corpus, + const std::string & target) { + GraphProgramLookup result; + std::unique_ptr program = + build_program_from_imported(graph.uid, std::move(imported_graph), corpus, target, result.status); + if (program == nullptr) { + return result; + } + + GraphProgramMatch match = program->match_current_graph(graph); + if (!match.valid()) { + result.status.append(match.status); + return result; + } + Status slot_status = program->capture_external_slots(graph, match); + if (!slot_status.success()) { + result.status.append(slot_status); + return result; + } + + if (graph.uid == 0) { + result.uncached_program = std::move(program); + result.program = result.uncached_program.get(); + result.match = std::move(match); + return result; + } + + GraphProgram * cached_program = program.get(); + { + std::lock_guard lock(mutex_); + programs_[graph.uid] = std::move(program); + cached_program = programs_[graph.uid].get(); + last_program_ = cached_program; + ++stats_.builds; + } + result.program = cached_program; + result.match = std::move(match); + return result; +} + +GraphProgramLookup GraphProgramCache::get_or_build(const ggml_cgraph & graph, + const KernelCorpus & corpus, + const std::string & target) { + GraphProgramLookup result; + if (graph.uid != 0) { + const bool disable_fast_path = environment_flag_enabled("GGML_HRX_DISABLE_GRAPH_UID_FAST_PATH"); + const bool validate_fast_path = environment_flag_enabled("GGML_HRX_VALIDATE_GRAPH_UID_CACHE"); + GraphProgram * cached_program = nullptr; + { + std::lock_guard lock(mutex_); + if (!disable_fast_path && last_program_ != nullptr && last_program_->uid() == graph.uid && + last_program_->target() == target) { + cached_program = last_program_; + } else { + const auto found = programs_.find(graph.uid); + if (found != programs_.end() && found->second->target() == target) { + cached_program = found->second.get(); + last_program_ = cached_program; + } + } + } + if (cached_program != nullptr) { + const bool bind_external = + !cached_program->has_prepared_program() || !cached_program->can_use_prepared_fast_path(graph); + GraphProgramMatch match = disable_fast_path || validate_fast_path ? + cached_program->match_current_graph(graph) : + cached_program->match_trusted_graph(graph, bind_external); + if (match.valid() && validate_fast_path && !disable_fast_path) { + GraphProgramMatch trusted_match = cached_program->match_trusted_graph(graph, bind_external); + if (!trusted_match.valid()) { + result.status.append(trusted_match.status); + return result; + } + } + if (match.valid()) { + result.program = cached_program; + result.match = std::move(match); + { + std::lock_guard lock(mutex_); + ++stats_.hits; + } + return result; + } + if (validate_fast_path) { + result.status.append(match.status); + return result; + } + } + } + + GraphImportResult imported = import_ggml_graph(graph); + if (!imported.valid()) { + result.status.append(imported.status); + return result; + } + return build_from_imported(graph, std::move(imported.graph), corpus, target); +} + +GraphProgramCacheStats GraphProgramCache::stats() const { + std::lock_guard lock(mutex_); + GraphProgramCacheStats stats = stats_; + for (const auto & entry : programs_) { + const PreparedCommandProgramCacheStats prepared = entry.second->prepared_stats(); + stats.prepared_program_builds += prepared.builds; + stats.prepared_program_hits += prepared.hits; + } + return stats; +} + +void GraphProgramCache::clear() { + std::lock_guard lock(mutex_); + programs_.clear(); + last_program_ = nullptr; +} + +std::unique_ptr GraphProgramCache::build_program_from_imported(uint64_t uid, + Graph && imported_graph, + const KernelCorpus & corpus, + const std::string & target, + Status & errors) const { + DispatchScheduler scheduler; + if (!scheduler.schedule_graph(imported_graph, { target })) { + errors.append(scheduler.plan().status); + return nullptr; + } + + CommandProgram commands = build_command_program(imported_graph, scheduler.plan(), corpus, target); + if (!commands.valid()) { + errors.append(commands.status); + return nullptr; + } + + std::string command_shape = command_program_shape_key(commands); + return std::make_unique(uid, target, std::make_unique(std::move(imported_graph)), + std::make_unique(std::move(commands)), + std::move(command_shape)); +} + +bool can_execute_standalone_op_as_graph(const ggml_tensor * op, const std::string & target) { + if (op == nullptr) { + return false; + } + Graph graph; + std::vector inputs; + for (const ggml_tensor * source : op->src) { + if (source == nullptr) { + continue; + } + inputs.push_back(graph.values().get_or_add_tensor_value(source, ValueKind::External)); + } + const ValueId output = graph.values().get_or_add_tensor_value(op, ValueKind::External); + GraphNode & node = graph.add_node(op->op, output, std::move(inputs)); + node.params = import_op_params(*op); + if (!graph.build_index().success()) { + return false; + } + return DispatchScheduler::supports_node(graph, &node, { target }); +} + +} // namespace ggml::hrx diff --git a/ggml/src/ggml-hrx/runtime/graph-program-cache.h b/ggml/src/ggml-hrx/runtime/graph-program-cache.h new file mode 100644 index 000000000000..0d2a804c295c --- /dev/null +++ b/ggml/src/ggml-hrx/runtime/graph-program-cache.h @@ -0,0 +1,160 @@ +#pragma once + +#include "dispatch/command-program.h" +#include "graph/graph.h" +#include "kernel-corpus/kernel-corpus.h" +#include "runtime/prepared-command-program-cache.h" +#include "status.h" + +#include +#include +#include +#include +#include +#include + +struct ggml_cgraph; +struct ggml_tensor; + +namespace ggml::hrx { + +struct GraphProgramExternalBinding { + ValueId value; + const ggml_tensor * tensor = nullptr; +}; + +enum class GraphProgramExternalSlotKind { + Node, + Source, +}; + +struct GraphProgramExternalSlot { + ValueId value; + GraphProgramExternalSlotKind kind = GraphProgramExternalSlotKind::Node; + size_t node_index = 0; + int32_t source_index = -1; +}; + +struct GraphProgramMatch { + std::vector external_bindings; + Status status; + + bool valid() const { return status.success(); } +}; + +class GraphProgram { + public: + GraphProgram(uint64_t uid, + std::string target, + std::unique_ptr graph, + std::unique_ptr commands, + std::string command_shape); + + uint64_t uid() const { return uid_; } + + const std::string & target() const { return target_; } + + const std::string & command_shape() const { return command_shape_; } + + const Graph & graph() const { return *graph_; } + + Graph & graph() { return *graph_; } + + const CommandProgram & commands() const { return *commands_; } + + CommandProgram & commands() { return *commands_; } + + GraphProgramMatch match_current_graph(const ggml_cgraph & graph) const; + GraphProgramMatch match_trusted_graph(const ggml_cgraph & graph, bool bind_external = true) const; + GraphProgramMatch match_host_staging_graph(const ggml_cgraph & graph) const; + + Status capture_external_slots(const ggml_cgraph & graph, const GraphProgramMatch & match); + + bool has_prepared_program() const; + bool can_use_prepared_fast_path(const ggml_cgraph & graph) const; + + PreparedCommandProgramCacheStats prepared_stats() const; + + PreparedCommandProgramCacheExecutionResult execute_with_result(const CommandProgramExecutionContext & context, + const CommandProgramBindings & bindings); + + private: + const GraphProgramExternalSlot * find_external_slot(ValueId value) const; + const ggml_tensor * resolve_external_slot(const ggml_cgraph & graph, + const GraphProgramExternalSlot & slot, + Status & status) const; + + uint64_t uid_ = 0; + std::string target_; + std::unique_ptr graph_; + std::unique_ptr commands_; + std::string command_shape_; + + std::vector external_slots_; + std::unordered_map external_slot_by_value_; + const ggml_tensor * const * fast_path_nodes_ = nullptr; + mutable std::mutex prepared_mutex_; + PreparedCommandProgram prepared_; + RecordedCommandGraph recorded_; + PreparedCommandProgramCacheStats prepared_stats_; + bool has_prepared_ = false; +}; + +struct GraphProgramCacheStats { + uint64_t builds = 0; + uint64_t hits = 0; + uint64_t prepared_program_builds = 0; + uint64_t prepared_program_hits = 0; +}; + +struct GraphProgramLookup { + GraphProgram * program = nullptr; + std::unique_ptr uncached_program; + GraphProgramMatch match; + Status status; + + bool valid() const { return program != nullptr && status.success() && match.valid(); } +}; + +struct GraphProgramSupportResult { + bool supported = false; + Status status; + + bool valid() const { return supported && status.success(); } +}; + +class GraphProgramCache { + public: + bool can_execute(const ggml_cgraph & graph, const KernelCorpus & corpus, const std::string & target) const; + + GraphProgramSupportResult check_support(const ggml_cgraph & graph, + const KernelCorpus & corpus, + const std::string & target) const; + + GraphProgramLookup get_or_build(const ggml_cgraph & graph, const KernelCorpus & corpus, const std::string & target); + + GraphProgramCacheStats stats() const; + + void clear(); + + private: + GraphProgramLookup build_from_imported(const ggml_cgraph & graph, + Graph && imported_graph, + const KernelCorpus & corpus, + const std::string & target); + + std::unique_ptr build_program_from_imported(uint64_t uid, + Graph && imported_graph, + const KernelCorpus & corpus, + const std::string & target, + Status & errors) const; + + mutable std::mutex mutex_; + std::unordered_map> programs_; + GraphProgram * last_program_ = nullptr; + GraphProgramCacheStats stats_; +}; + +bool can_execute_standalone_op_as_graph(const ggml_tensor * op, const std::string & target); + +} // namespace ggml::hrx diff --git a/ggml/src/ggml-hrx/runtime/graph-replay.h b/ggml/src/ggml-hrx/runtime/graph-replay.h new file mode 100644 index 000000000000..3bcd3abdcbbe --- /dev/null +++ b/ggml/src/ggml-hrx/runtime/graph-replay.h @@ -0,0 +1,44 @@ +#pragma once + +#include +#include + +namespace ggml::hrx { + +enum class HrxGraphReplayEvent { + Disabled, + Ineligible, + MissBuild, + Hit, + RebuildTransient, + BuildFailed, + LaunchFailed, +}; + +inline const char * hrx_graph_replay_event_name(HrxGraphReplayEvent event) { + switch (event) { + case HrxGraphReplayEvent::Disabled: + return "disabled"; + case HrxGraphReplayEvent::Ineligible: + return "ineligible"; + case HrxGraphReplayEvent::MissBuild: + return "miss_build"; + case HrxGraphReplayEvent::Hit: + return "hit"; + case HrxGraphReplayEvent::RebuildTransient: + return "rebuild_transient"; + case HrxGraphReplayEvent::BuildFailed: + return "build_failed"; + case HrxGraphReplayEvent::LaunchFailed: + return "launch_failed"; + } + return "unknown"; +} + +inline uint64_t hrx_graph_replay_now_ns() { + return static_cast( + std::chrono::duration_cast(std::chrono::steady_clock::now().time_since_epoch()) + .count()); +} + +} // namespace ggml::hrx diff --git a/ggml/src/ggml-hrx/runtime/host-buffer-registry.cpp b/ggml/src/ggml-hrx/runtime/host-buffer-registry.cpp new file mode 100644 index 000000000000..02b591fcfc9c --- /dev/null +++ b/ggml/src/ggml-hrx/runtime/host-buffer-registry.cpp @@ -0,0 +1,71 @@ +#include "host-buffer-registry.h" + +#include "hrx_runtime.h" + +#include +#include + +namespace ggml::hrx { + +HostBufferRef::HostBufferRef(hrx_buffer_t buffer, size_t offset) : buffer_(buffer), offset_(offset) {} + +HostBufferRef::~HostBufferRef() { + if (buffer_ != nullptr) { + hrx_buffer_release(buffer_); + } +} + +HostBufferRef::HostBufferRef(HostBufferRef && other) noexcept : + buffer_(std::exchange(other.buffer_, nullptr)), + offset_(other.offset_) { + other.offset_ = 0; +} + +HostBufferRef & HostBufferRef::operator=(HostBufferRef && other) noexcept { + if (this != &other) { + if (buffer_ != nullptr) { + hrx_buffer_release(buffer_); + } + buffer_ = std::exchange(other.buffer_, nullptr); + offset_ = other.offset_; + other.offset_ = 0; + } + return *this; +} + +void HostBufferRegistry::add(hrx_buffer_t buffer, void * base, size_t size) { + if (buffer == nullptr || base == nullptr || size == 0) { + return; + } + std::lock_guard lock(mutex_); + entries_.push_back({ buffer, reinterpret_cast(base), size }); +} + +void HostBufferRegistry::remove(hrx_buffer_t buffer) { + std::lock_guard lock(mutex_); + entries_.erase(std::remove_if(entries_.begin(), entries_.end(), + [buffer](const Entry & entry) { return entry.buffer == buffer; }), + entries_.end()); +} + +HostBufferRef HostBufferRegistry::find(const void * data, size_t size) const { + if (data == nullptr) { + return {}; + } + const uintptr_t address = reinterpret_cast(data); + std::lock_guard lock(mutex_); + for (const Entry & entry : entries_) { + if (address < entry.base) { + continue; + } + const size_t offset = static_cast(address - entry.base); + if (offset > entry.size || size > entry.size - offset) { + continue; + } + hrx_buffer_retain(entry.buffer); + return HostBufferRef(entry.buffer, offset); + } + return {}; +} + +} // namespace ggml::hrx diff --git a/ggml/src/ggml-hrx/runtime/host-buffer-registry.h b/ggml/src/ggml-hrx/runtime/host-buffer-registry.h new file mode 100644 index 000000000000..30ce0eba2ec9 --- /dev/null +++ b/ggml/src/ggml-hrx/runtime/host-buffer-registry.h @@ -0,0 +1,56 @@ +#pragma once + +#include +#include +#include +#include + +typedef struct hrx_buffer_s * hrx_buffer_t; + +namespace ggml::hrx { + +class HostBufferRef { + public: + HostBufferRef() = default; + ~HostBufferRef(); + + HostBufferRef(HostBufferRef && other) noexcept; + HostBufferRef & operator=(HostBufferRef && other) noexcept; + + HostBufferRef(const HostBufferRef &) = delete; + HostBufferRef & operator=(const HostBufferRef &) = delete; + + bool valid() const { return buffer_ != nullptr; } + + hrx_buffer_t buffer() const { return buffer_; } + + size_t offset() const { return offset_; } + + private: + hrx_buffer_t buffer_ = nullptr; + size_t offset_ = 0; + + HostBufferRef(hrx_buffer_t buffer, size_t offset); + friend class HostBufferRegistry; +}; + +// Tracks mapped HRX host buffers so pointer-based GGML transfers can remain stream ordered and handle based. +class HostBufferRegistry { + public: + void add(hrx_buffer_t buffer, void * base, size_t size); + void remove(hrx_buffer_t buffer); + + HostBufferRef find(const void * data, size_t size) const; + + private: + struct Entry { + hrx_buffer_t buffer = nullptr; + uintptr_t base = 0; + size_t size = 0; + }; + + mutable std::mutex mutex_; + std::vector entries_; +}; + +} // namespace ggml::hrx diff --git a/ggml/src/ggml-hrx/runtime/host-memory.cpp b/ggml/src/ggml-hrx/runtime/host-memory.cpp new file mode 100644 index 000000000000..e74ae16c98d3 --- /dev/null +++ b/ggml/src/ggml-hrx/runtime/host-memory.cpp @@ -0,0 +1,320 @@ +#include "host-memory.h" + +#include "hrx-interop-utils.h" +#include "hrx_runtime.h" + +#include +#include + +namespace ggml::hrx { +namespace { + +static constexpr size_t kMaxInlineUploadBytes = 63 * 1024; + +static Status allocate_device_buffer(hrx_device_t device, size_t size, hrx_buffer_t & buffer) { + Status status; + if (device == nullptr) { + status.log("missing HRX device for host memory allocation"); + return status; + } + if (size == 0) { + status.log("cannot allocate an empty host memory device buffer"); + return status; + } + hrx_buffer_params_t params = { + HRX_MEMORY_TYPE_DEVICE_LOCAL, + HRX_MEMORY_ACCESS_ALL, + HRX_BUFFER_USAGE_DEFAULT, + 0, + }; + if (ErrorResult error = + take_status(hrx_allocator_allocate_buffer(hrx_device_allocator(device), params, size, &buffer))) { + status.log("allocate host memory device buffer: %s", error->c_str()); + } + return status; +} + +} // namespace + +Status HostTransferManager::upload_synchronous(hrx_stream_t stream, + const void * host_source, + hrx_buffer_t destination, + size_t offset, + size_t size) { + Status status; + if (size == 0) { + return status; + } + if (stream == nullptr || host_source == nullptr || destination == nullptr) { + status.log("invalid HRX host upload"); + return status; + } + hrx_device_t device = nullptr; + if (ErrorResult error = take_status(hrx_stream_get_device(stream, &device))) { + status.log("query HRX upload device failed: %s", error->c_str()); + return status; + } + if (ErrorResult error = take_status(hrx_stream_synchronize(stream))) { + status.log("synchronize before HRX host upload failed: %s", error->c_str()); + return status; + } + if (ErrorResult error = take_status(hrx_synchronous_h2d(device, host_source, destination, offset, size))) { + status.log("synchronous HRX host upload failed: %s", error->c_str()); + return status; + } + std::lock_guard lock(mutex_); + ++stats_.uploads; + stats_.upload_bytes += size; + return status; +} + +Status HostTransferManager::upload_async(hrx_stream_t stream, + const void * host_source, + hrx_buffer_t destination, + size_t offset, + size_t size) { + Status status; + if (size == 0) { + return status; + } + if (stream == nullptr || host_source == nullptr || destination == nullptr) { + status.log("invalid HRX host upload"); + return status; + } + + const uint8_t * host_bytes = static_cast(host_source); + size_t uploaded = 0; + while (uploaded < size) { + const size_t remaining = size - uploaded; + const size_t chunk_size = remaining < kMaxInlineUploadBytes ? remaining : kMaxInlineUploadBytes; + if (ErrorResult error = take_status( + hrx_stream_update_buffer(stream, host_bytes + uploaded, chunk_size, destination, offset + uploaded))) { + status.log("HRX async host upload failed: %s", error->c_str()); + return status; + } + uploaded += chunk_size; + } + + std::lock_guard lock(mutex_); + ++stats_.uploads; + stats_.upload_bytes += size; + return status; +} + +Status HostTransferManager::download_synchronous(hrx_stream_t stream, + hrx_buffer_t source, + size_t offset, + void * host_destination, + size_t size) { + Status status; + if (size == 0) { + return status; + } + if (stream == nullptr || source == nullptr || host_destination == nullptr) { + status.log("invalid HRX host download"); + return status; + } + hrx_device_t device = nullptr; + if (ErrorResult error = take_status(hrx_stream_get_device(stream, &device))) { + status.log("query HRX download device failed: %s", error->c_str()); + return status; + } + if (ErrorResult error = take_status(hrx_stream_synchronize(stream))) { + status.log("synchronize before HRX host download failed: %s", error->c_str()); + return status; + } + if (ErrorResult error = take_status(hrx_synchronous_d2h(device, source, offset, host_destination, size))) { + status.log("synchronous HRX host download failed: %s", error->c_str()); + return status; + } + std::lock_guard lock(mutex_); + ++stats_.downloads; + stats_.download_bytes += size; + return status; +} + +HostTransferStats HostTransferManager::stats() const { + std::lock_guard lock(mutex_); + return stats_; +} + +void HostTransferManager::clear() { + std::lock_guard lock(mutex_); + stats_ = {}; +} + +struct HostWeightLease::Entry { + ~Entry() { + if (buffer != nullptr) { + hrx_buffer_release(buffer); + } + } + + hrx_buffer_t buffer = nullptr; + size_t length = 0; + std::string layout; +}; + +HostWeightLease::HostWeightLease(std::shared_ptr entry) : entry_(std::move(entry)) {} + +bool HostWeightLease::valid() const { + return entry_ != nullptr && entry_->buffer != nullptr; +} + +hrx_buffer_t HostWeightLease::buffer() const { + return valid() ? entry_->buffer : nullptr; +} + +size_t HostWeightLease::length() const { + return entry_ != nullptr ? entry_->length : 0; +} + +const std::string & HostWeightLease::layout() const { + static const std::string empty; + return entry_ != nullptr ? entry_->layout : empty; +} + +HostWeightCache::~HostWeightCache() { + clear(); +} + +size_t HostWeightCache::SourceKeyHash::operator()(const SourceKey & key) const { + uint64_t hash = UINT64_C(1469598103934665603); + auto mix = [&](uint64_t value) { + hash ^= value; + hash *= UINT64_C(1099511628211); + }; + mix(key.identity); + mix(key.generation); + mix(static_cast(key.capacity)); + mix(static_cast(key.offset)); + mix(static_cast(key.length)); + return static_cast(hash); +} + +HostWeightAcquireResult HostWeightCache::acquire(hrx_device_t device, + hrx_stream_t stream, + HostTransferManager & transfers, + const HostWeightSource & source) { + HostWeightAcquireResult result; + if (stream == nullptr) { + result.status.log("host weight residency requires an HRX stream"); + return result; + } + if (source.host_data == nullptr || source.identity == 0 || source.generation == 0 || source.length == 0 || + source.offset > source.capacity || source.length > source.capacity - source.offset) { + result.status.log("invalid host weight source"); + return result; + } + if (source.layout.empty()) { + result.status.log("host weight source has no layout"); + return result; + } + + const SourceKey key{ source.identity, source.generation, source.capacity, source.offset, source.length }; + { + std::lock_guard lock(mutex_); + const auto found = entries_.find(key); + if (found != entries_.end()) { + if (found->second->layout != source.layout) { + ++stats_.layout_conflicts; + result.status.log("host weight source already has resident layout %s, cannot also materialize %s", + found->second->layout.c_str(), source.layout.c_str()); + return result; + } + ++stats_.hits; + result.lease = HostWeightLease(found->second); + return result; + } + } + + auto entry = std::make_shared(); + entry->length = source.length; + entry->layout = source.layout; + result.status = allocate_device_buffer(device, source.length, entry->buffer); + if (!result.status.success()) { + return result; + } + result.status = transfers.upload_synchronous( + stream, static_cast(source.host_data) + source.offset, entry->buffer, 0, source.length); + if (!result.status.success()) { + return result; + } + + { + std::lock_guard lock(mutex_); + const auto inserted = entries_.emplace(key, entry); + if (!inserted.second) { + ++stats_.hits; + result.lease = HostWeightLease(inserted.first->second); + return result; + } + ++stats_.misses; + stats_.allocation_count = entries_.size(); + stats_.resident_bytes += source.length; + } + result.lease = HostWeightLease(std::move(entry)); + return result; +} + +HostWeightCacheStats HostWeightCache::stats() const { + std::lock_guard lock(mutex_); + return stats_; +} + +void HostWeightCache::clear() { + std::lock_guard lock(mutex_); + entries_.clear(); + stats_ = {}; +} + +HostStagingBuffer::~HostStagingBuffer() { + clear(); +} + +HostStagingBuffer::HostStagingBuffer(HostStagingBuffer && other) noexcept { + *this = std::move(other); +} + +HostStagingBuffer & HostStagingBuffer::operator=(HostStagingBuffer && other) noexcept { + if (this == &other) { + return *this; + } + clear(); + buffer = other.buffer; + host_data = other.host_data; + value = other.value; + length = other.length; + upload = other.upload; + download = other.download; + other.buffer = nullptr; + other.host_data = nullptr; + other.value = -1; + other.length = 0; + other.upload = false; + other.download = false; + return *this; +} + +void HostStagingBuffer::clear() { + if (buffer != nullptr) { + hrx_buffer_release(buffer); + buffer = nullptr; + } + host_data = nullptr; + value = -1; + length = 0; + upload = false; + download = false; +} + +Status allocate_host_staging_buffer(hrx_device_t device, size_t size, HostStagingBuffer & staging) { + staging.clear(); + Status status = allocate_device_buffer(device, size, staging.buffer); + if (status.success()) { + staging.length = size; + } + return status; +} + +} // namespace ggml::hrx diff --git a/ggml/src/ggml-hrx/runtime/host-memory.h b/ggml/src/ggml-hrx/runtime/host-memory.h new file mode 100644 index 000000000000..ee4dc45efec4 --- /dev/null +++ b/ggml/src/ggml-hrx/runtime/host-memory.h @@ -0,0 +1,147 @@ +#pragma once + +#include "status.h" + +#include +#include +#include +#include +#include +#include + +typedef struct hrx_buffer_s * hrx_buffer_t; +typedef struct hrx_device_s * hrx_device_t; +typedef struct hrx_stream_s * hrx_stream_t; + +namespace ggml::hrx { + +struct HostTransferStats { + uint64_t uploads = 0; + uint64_t downloads = 0; + size_t upload_bytes = 0; + size_t download_bytes = 0; +}; + +class HostTransferManager { + public: + Status upload_synchronous( + hrx_stream_t stream, const void * host_source, hrx_buffer_t destination, size_t offset, size_t size); + Status upload_async(hrx_stream_t stream, + const void * host_source, + hrx_buffer_t destination, + size_t offset, + size_t size); + Status download_synchronous( + hrx_stream_t stream, hrx_buffer_t source, size_t offset, void * host_destination, size_t size); + + HostTransferStats stats() const; + void clear(); + + private: + mutable std::mutex mutex_; + HostTransferStats stats_; +}; + +struct HostWeightSource { + const void * host_data = nullptr; + uint64_t identity = 0; + uint64_t generation = 0; + size_t capacity = 0; + size_t offset = 0; + size_t length = 0; + std::string layout = "ggml-native"; +}; + +struct HostWeightCacheStats { + uint64_t hits = 0; + uint64_t misses = 0; + uint64_t layout_conflicts = 0; + size_t allocation_count = 0; + size_t resident_bytes = 0; +}; + +class HostWeightLease { + public: + HostWeightLease() = default; + + bool valid() const; + hrx_buffer_t buffer() const; + size_t length() const; + const std::string & layout() const; + + private: + struct Entry; + std::shared_ptr entry_; + + explicit HostWeightLease(std::shared_ptr entry); + friend class HostWeightCache; +}; + +struct HostWeightAcquireResult { + HostWeightLease lease; + Status status; + + bool valid() const { return lease.valid() && status.success(); } +}; + +class HostWeightCache { + public: + HostWeightCache() = default; + ~HostWeightCache(); + + HostWeightCache(const HostWeightCache &) = delete; + HostWeightCache & operator=(const HostWeightCache &) = delete; + + HostWeightAcquireResult acquire(hrx_device_t device, + hrx_stream_t stream, + HostTransferManager & transfers, + const HostWeightSource & source); + HostWeightCacheStats stats() const; + void clear(); + + private: + struct SourceKey { + uint64_t identity = 0; + uint64_t generation = 0; + size_t capacity = 0; + size_t offset = 0; + size_t length = 0; + + bool operator==(const SourceKey & other) const { + return identity == other.identity && generation == other.generation && capacity == other.capacity && + offset == other.offset && length == other.length; + } + }; + + struct SourceKeyHash { + size_t operator()(const SourceKey & key) const; + }; + + mutable std::mutex mutex_; + std::unordered_map, SourceKeyHash> entries_; + HostWeightCacheStats stats_; +}; + +struct HostStagingBuffer { + HostStagingBuffer() = default; + ~HostStagingBuffer(); + + HostStagingBuffer(HostStagingBuffer && other) noexcept; + HostStagingBuffer & operator=(HostStagingBuffer && other) noexcept; + + HostStagingBuffer(const HostStagingBuffer &) = delete; + HostStagingBuffer & operator=(const HostStagingBuffer &) = delete; + + hrx_buffer_t buffer = nullptr; + void * host_data = nullptr; + int32_t value = -1; + size_t length = 0; + bool upload = false; + bool download = false; + + void clear(); +}; + +Status allocate_host_staging_buffer(hrx_device_t device, size_t size, HostStagingBuffer & staging); + +} // namespace ggml::hrx diff --git a/ggml/src/ggml-hrx/runtime/kernel-executable-cache.cpp b/ggml/src/ggml-hrx/runtime/kernel-executable-cache.cpp new file mode 100644 index 000000000000..516bc8247d02 --- /dev/null +++ b/ggml/src/ggml-hrx/runtime/kernel-executable-cache.cpp @@ -0,0 +1,412 @@ +#include "kernel-executable-cache.h" + +#include "ggml-impl.h" +#include "hrx-interop-utils.h" + +#include +#include +#include +#include +#include +#include +#include + +namespace ggml::hrx { +namespace { + +static ggml_hrx_loom_jit_source_format to_jit_source_format(KernelSourceFormat format) { + switch (format) { + case KERNEL_SOURCE_FORMAT_TEXT: + return GGML_HRX_LOOM_JIT_SOURCE_FORMAT_TEXT; + case KERNEL_SOURCE_FORMAT_BINARY: + return GGML_HRX_LOOM_JIT_SOURCE_FORMAT_BYTECODE; + } + return GGML_HRX_LOOM_JIT_SOURCE_FORMAT_TEXT; +} + +static void append_u32(std::vector & bytes, uint32_t value) { + const size_t offset = bytes.size(); + bytes.resize(offset + sizeof(value)); + std::memcpy(bytes.data() + offset, &value, sizeof(value)); +} + +static bool pack_kernel_constants(const KernelDefinition & definition, + const Dispatch & dispatch, + std::vector & constants) { + constants.clear(); + for (const KernelScalarDefinition & parameter : definition.launch_parameters) { + const char * name = parameter.name != nullptr ? parameter.name : ""; + const char * type = parameter.type != nullptr ? parameter.type : ""; + const auto item = dispatch.kernel.integer_parameters.find(name); + const bool supported_type = std::strcmp(type, "index") == 0 || std::strcmp(type, "i32") == 0; + if (item == dispatch.kernel.integer_parameters.end() || !supported_type || item->second < 0 || + static_cast(item->second) > std::numeric_limits::max()) { + constants.clear(); + GGML_LOG_ERROR("%s: invalid launch scalar %s for %s\n", __func__, name, + kernel_definition_name(definition).c_str()); + return false; + } + append_u32(constants, static_cast(item->second)); + } + return true; +} + +static std::string kernel_executable_key(const KernelDefinition & definition, + const Dispatch & dispatch, + const char * target) { + std::ostringstream out; + out << (target != nullptr ? target : "") << '|' << definition.source_digest << '|' << definition.symbol + << "|recipe=" << definition.compile_recipe.mode; + for (const KernelScalarDefinition & parameter : definition.workload_parameters) { + const char * name = parameter.name != nullptr ? parameter.name : ""; + const auto item = dispatch.kernel.integer_parameters.find(name); + out << '|' << name << '='; + if (item == dispatch.kernel.integer_parameters.end()) { + out << ""; + } else { + out << item->second; + } + } + for (const KernelCompileConfig & config : definition.compile_config) { + out << '|' << (config.key != nullptr ? config.key : "") << '=' << (config.value != nullptr ? config.value : ""); + } + for (const auto & config : dispatch.kernel.compile_parameters) { + out << '|' << config.first << '=' << config.second; + } + return out.str(); +} + +static bool build_compile_request(const KernelDefinition & definition, + const Dispatch & dispatch, + LoomKernelCompileRequest & request) { + if (definition.compile_recipe.primary_sources.empty()) { + GGML_LOG_ERROR("%s: kernel %s has no primary source\n", __func__, kernel_definition_name(definition).c_str()); + return false; + } + const KernelSourceRef & primary_source = definition.compile_recipe.primary_sources.front(); + const KernelSource * source = primary_source.contents; + if (source == nullptr) { + GGML_LOG_ERROR("%s: missing embedded source for %s\n", __func__, primary_source.path); + return false; + } + + request.source_data = source->source.data; + request.source_size = source->source.length; + request.source_format = to_jit_source_format(source->source.format); + request.source_identifier = primary_source.path != nullptr ? primary_source.path : ""; + request.symbol = definition.symbol != nullptr ? definition.symbol : ""; + request.launch_config_symbol = definition.name != nullptr ? definition.name : ""; + + request.dependencies.reserve(definition.compile_recipe.library_sources.size()); + for (const KernelSourceRef & dependency_ref : definition.compile_recipe.library_sources) { + const KernelSource * dependency = dependency_ref.contents; + if (dependency == nullptr) { + GGML_LOG_ERROR("%s: missing embedded dependency for %s\n", __func__, dependency_ref.path); + return false; + } + request.dependencies.push_back({ + dependency->source.data, + dependency->source.length, + to_jit_source_format(dependency->source.format), + dependency_ref.path, + }); + } + + std::map merged_configs; + for (const KernelCompileConfig & config : definition.compile_config) { + merged_configs[config.key != nullptr ? config.key : ""] = config.value != nullptr ? config.value : ""; + } + for (const auto & config : dispatch.kernel.compile_parameters) { + merged_configs[config.first] = config.second; + } + request.config_storage.reserve(merged_configs.size()); + for (const auto & config : merged_configs) { + request.config_storage.push_back(config); + } + + request.workload.reserve(definition.workload_parameters.size()); + for (const KernelScalarDefinition & parameter : definition.workload_parameters) { + const char * name = parameter.name != nullptr ? parameter.name : ""; + const char * type = parameter.type != nullptr ? parameter.type : ""; + const auto item = dispatch.kernel.integer_parameters.find(name); + const bool supported_type = std::strcmp(type, "index") == 0 || std::strcmp(type, "i32") == 0; + if (item == dispatch.kernel.integer_parameters.end() || !supported_type) { + GGML_LOG_ERROR("%s: invalid workload scalar %s for %s\n", __func__, name, + kernel_definition_name(definition).c_str()); + return false; + } + request.workload.push_back(item->second); + } + return true; +} + +static std::shared_ptr load_kernel_executable(const KernelExecutablePrepareContext & context, + const KernelDefinition & definition, + const Dispatch & dispatch, + const std::vector & constants, + const std::string & key, + ggml_hrx_loom_jit_compile_result & compiled, + std::string & error_message) { + if (context.device == nullptr) { + error_message = "missing HRX device"; + GGML_LOG_ERROR("%s: load %s: %s\n", __func__, key.c_str(), error_message.c_str()); + return nullptr; + } + if (context.target == nullptr) { + error_message = "missing HRX target"; + GGML_LOG_ERROR("%s: load %s: %s\n", __func__, key.c_str(), error_message.c_str()); + return nullptr; + } + + auto executable = std::make_shared(); + executable->launch = compiled.launch_config; + if (ErrorResult error = + take_status(hrx_executable_load_data(context.device, compiled.hsaco_data, compiled.hsaco_size, "amdgpu", + context.target, &executable->executable))) { + error_message = "load " + key + ": " + *error; + GGML_LOG_ERROR("%s: %s\n", __func__, error_message.c_str()); + return nullptr; + } + if (ErrorResult error = take_status(hrx_executable_lookup_export_by_name(executable->executable, definition.name, + &executable->export_ordinal))) { + error_message = "lookup " + key + ": " + *error; + GGML_LOG_ERROR("%s: %s\n", __func__, error_message.c_str()); + return nullptr; + } + if (ErrorResult error = take_status( + hrx_executable_export_info(executable->executable, executable->export_ordinal, &executable->export_info))) { + error_message = "inspect " + key + ": " + *error; + GGML_LOG_ERROR("%s: %s\n", __func__, error_message.c_str()); + return nullptr; + } + if (executable->export_info.binding_count != dispatch.bindings.size() || + executable->export_info.constant_byte_length != constants.size() || + executable->export_info.parameter_count != dispatch.bindings.size() + definition.launch_parameters.size()) { + error_message = "compiled ABI does not match manifest for " + key; + GGML_LOG_ERROR("%s: %s\n", __func__, error_message.c_str()); + return nullptr; + } + if (executable->launch.workgroup_count[0] == 0 || executable->launch.workgroup_size[0] == 0) { + error_message = "compiled launch geometry is empty for " + key; + GGML_LOG_ERROR("%s: %s\n", __func__, error_message.c_str()); + return nullptr; + } + return executable; +} + +} // namespace + +class KernelExecutableCacheEntry { + public: + KernelExecutableCacheEntry(std::string key, + const KernelDefinition & definition, + const Dispatch & dispatch, + LoomCompiledKernelRef compiled_ref) : + key(std::move(key)), + definition(&definition), + dispatch(dispatch), + compiled_ref(std::move(compiled_ref)) {} + + std::mutex mutex; + std::condition_variable complete; + std::string key; + const KernelDefinition * definition = nullptr; + Dispatch dispatch; + LoomCompiledKernelRef compiled_ref; + std::shared_ptr executable; + std::string error; + + enum class LoadState { + Unloaded, + Loading, + Loaded, + Failed, + }; + + std::atomic load_state = LoadState::Unloaded; +}; + +KernelExecutable::~KernelExecutable() { + if (executable != nullptr) { + hrx_executable_release(executable); + } +} + +KernelExecutableCache::KernelExecutableCache(LoomJitMode mode) : mode_(mode), mode_is_forced_(true) {} + +KernelExecutableCache::~KernelExecutableCache() { + clear(); +} + +bool KernelExecutableCache::ensure_jit_locked(const char * target, std::string & error_message) { + if (target == nullptr || target[0] == '\0') { + error_message = "missing HRX target"; + GGML_LOG_ERROR("%s: %s\n", __func__, error_message.c_str()); + return false; + } + if (jit_ != nullptr) { + if (target_ != target) { + error_message = + "HRX kernel executable cache target mismatch: existing " + target_ + ", requested " + target; + GGML_LOG_ERROR("%s: %s\n", __func__, error_message.c_str()); + return false; + } + return true; + } + + jit_ = mode_is_forced_ ? create_loom_jit(target, mode_, error_message) : create_loom_jit(target, error_message); + if (jit_ == nullptr) { + if (error_message.empty()) { + error_message = "create Loom JIT failed"; + } + return false; + } + target_ = target; + return true; +} + +KernelExecutableRef KernelExecutableCache::get_or_compile(const KernelExecutablePrepareContext & context, + const KernelDefinition & definition, + const Dispatch & dispatch, + std::vector & constants) { + KernelExecutableRef ref; + if (!pack_kernel_constants(definition, dispatch, constants)) { + return ref; + } + + const std::string key = kernel_executable_key(definition, dispatch, context.target); + LoomKernelCompileRequest request; + LoomCompiledKernelRef compiled_ref; + { + std::lock_guard lock(mutex_); + const auto found = cache_.find(key); + if (found != cache_.end()) { + ref.entry = found->second; + return ref; + } + std::string error_message; + if (!ensure_jit_locked(context.target, error_message)) { + return ref; + } + if (!build_compile_request(definition, dispatch, request)) { + return ref; + } + + compiled_ref = jit_->compile(key, std::move(request)); + auto entry = std::make_shared(key, definition, dispatch, compiled_ref); + cache_.emplace(key, entry); + ref.entry = std::move(entry); + } + return ref; +} + +std::shared_ptr KernelExecutableCache::materialize(const KernelExecutablePrepareContext & context, + const KernelExecutableRef & ref, + const std::vector & constants) { + if (!ref.valid()) { + return nullptr; + } + + KernelExecutableCacheEntry & entry = *ref.entry; + KernelExecutableCacheEntry::LoadState load_state = entry.load_state.load(std::memory_order_acquire); + if (load_state == KernelExecutableCacheEntry::LoadState::Loaded) { + return std::atomic_load_explicit(&entry.executable, std::memory_order_acquire); + } + if (load_state == KernelExecutableCacheEntry::LoadState::Failed) { + std::lock_guard entry_lock(entry.mutex); + GGML_LOG_ERROR("%s: %s\n", __func__, entry.error.c_str()); + return nullptr; + } + + KernelExecutableCacheEntry::LoadState expected = KernelExecutableCacheEntry::LoadState::Unloaded; + if (!entry.load_state.compare_exchange_strong(expected, KernelExecutableCacheEntry::LoadState::Loading, + std::memory_order_acq_rel, std::memory_order_acquire)) { + std::unique_lock entry_lock(entry.mutex); + entry.complete.wait(entry_lock, [&] { + const KernelExecutableCacheEntry::LoadState current = entry.load_state.load(std::memory_order_acquire); + return current == KernelExecutableCacheEntry::LoadState::Loaded || + current == KernelExecutableCacheEntry::LoadState::Failed; + }); + if (entry.load_state.load(std::memory_order_acquire) == KernelExecutableCacheEntry::LoadState::Failed) { + GGML_LOG_ERROR("%s: %s\n", __func__, entry.error.c_str()); + return nullptr; + } + return std::atomic_load_explicit(&entry.executable, std::memory_order_acquire); + } + + LoomCompiledKernelRef compiled_ref; + { + std::lock_guard entry_lock(entry.mutex); + compiled_ref = entry.compiled_ref; + } + + std::string error_message; + if (compiled_ref == nullptr || !compiled_ref->resolve()) { + error_message = compiled_ref ? compiled_ref->error_message() : "missing compiled kernel"; + GGML_LOG_ERROR("%s: %s\n", __func__, error_message.c_str()); + + { + std::lock_guard entry_lock(entry.mutex); + entry.error = error_message; + entry.load_state.store(KernelExecutableCacheEntry::LoadState::Failed, std::memory_order_release); + } + entry.complete.notify_all(); + + std::lock_guard lock(mutex_); + const auto found = cache_.find(entry.key); + if (found != cache_.end() && found->second == ref.entry) { + cache_.erase(found); + } + return nullptr; + } + + ggml_hrx_loom_jit_compile_result compiled = compiled_ref->take_result(); + std::shared_ptr executable = load_kernel_executable( + context, *entry.definition, entry.dispatch, constants, entry.key, compiled, error_message); + + if (executable != nullptr) { + compiled.reset(); + { + std::lock_guard entry_lock(entry.mutex); + entry.compiled_ref.reset(); + std::atomic_store_explicit(&entry.executable, executable, std::memory_order_release); + entry.load_state.store(KernelExecutableCacheEntry::LoadState::Loaded, std::memory_order_release); + } + } else { + { + std::lock_guard entry_lock(entry.mutex); + entry.error = std::move(error_message); + entry.load_state.store(KernelExecutableCacheEntry::LoadState::Failed, std::memory_order_release); + } + } + entry.complete.notify_all(); + + if (executable == nullptr) { + std::lock_guard lock(mutex_); + const auto found = cache_.find(entry.key); + if (found != cache_.end() && found->second == ref.entry) { + cache_.erase(found); + } + } + return executable; +} + +std::shared_ptr KernelExecutableCache::prepare(const KernelExecutablePrepareContext & context, + const KernelDefinition & definition, + const Dispatch & dispatch, + std::vector & constants) { + const KernelExecutableRef ref = get_or_compile(context, definition, dispatch, constants); + return materialize(context, ref, constants); +} + +void KernelExecutableCache::clear() { + if (jit_ != nullptr) { + jit_->clear(); + } + std::lock_guard lock(mutex_); + cache_.clear(); + jit_.reset(); + target_.clear(); +} + +} // namespace ggml::hrx diff --git a/ggml/src/ggml-hrx/runtime/kernel-executable-cache.h b/ggml/src/ggml-hrx/runtime/kernel-executable-cache.h new file mode 100644 index 000000000000..2ae130d9918e --- /dev/null +++ b/ggml/src/ggml-hrx/runtime/kernel-executable-cache.h @@ -0,0 +1,72 @@ +#pragma once + +#include "dispatch/dispatch.h" +#include "hrx_runtime.h" +#include "kernel-corpus/kernel-corpus.h" +#include "runtime/loom-kernel-jit.h" + +#include +#include +#include +#include +#include +#include + +namespace ggml::hrx { + +class KernelExecutableCacheEntry; + +struct KernelExecutable { + ~KernelExecutable(); + + hrx_executable_t executable = nullptr; + uint32_t export_ordinal = 0; + hrx_executable_export_info_t export_info = {}; + ggml_hrx_loom_jit_launch_config launch; +}; + +struct KernelExecutablePrepareContext { + hrx_device_t device = nullptr; + const char * target = nullptr; +}; + +struct KernelExecutableRef { + std::shared_ptr entry; + + bool valid() const { return entry != nullptr; } +}; + +class KernelExecutableCache { + public: + KernelExecutableCache() = default; + explicit KernelExecutableCache(LoomJitMode mode); + ~KernelExecutableCache(); + + KernelExecutableRef get_or_compile(const KernelExecutablePrepareContext & context, + const KernelDefinition & definition, + const Dispatch & dispatch, + std::vector & constants); + + std::shared_ptr materialize(const KernelExecutablePrepareContext & context, + const KernelExecutableRef & ref, + const std::vector & constants); + + std::shared_ptr prepare(const KernelExecutablePrepareContext & context, + const KernelDefinition & definition, + const Dispatch & dispatch, + std::vector & constants); + + void clear(); + + private: + bool ensure_jit_locked(const char * target, std::string & error_message); + + std::mutex mutex_; + std::unordered_map> cache_; + std::unique_ptr jit_; + std::string target_; + LoomJitMode mode_ = LoomJitMode::Async; + bool mode_is_forced_ = false; +}; + +} // namespace ggml::hrx diff --git a/ggml/src/ggml-hrx/runtime/loom-kernel-jit.cpp b/ggml/src/ggml-hrx/runtime/loom-kernel-jit.cpp new file mode 100644 index 000000000000..8d107144aa9b --- /dev/null +++ b/ggml/src/ggml-hrx/runtime/loom-kernel-jit.cpp @@ -0,0 +1,308 @@ +#include "loom-kernel-jit.h" + +#include "ggml-impl.h" +#include "hrx-interop-utils.h" + +#include +#include +#include +#include +#include +#include + +namespace ggml::hrx { +namespace { + +class LoomAmdgpuJit { + public: + LoomAmdgpuJit() = default; + LoomAmdgpuJit(const LoomAmdgpuJit &) = delete; + LoomAmdgpuJit & operator=(const LoomAmdgpuJit &) = delete; + + ~LoomAmdgpuJit() { reset(); } + + bool create(const char * target, std::string & error_message) { + reset(); + ggml_hrx_loom_jit_amdgpu_options options = {}; + options.processor = target; + options.identifier = target; + if (ErrorResult error = take_status(ggml_hrx_loom_jit_amdgpu_create(&options, &jit_))) { + error_message = "create Loom JIT: " + *error; + GGML_LOG_ERROR("%s: %s\n", __func__, error_message.c_str()); + return false; + } + return true; + } + + ggml_hrx_loom_jit_amdgpu * get() const { return jit_; } + + private: + void reset() { + if (jit_ != nullptr) { + ggml_hrx_loom_jit_amdgpu_release(jit_); + jit_ = nullptr; + } + } + + ggml_hrx_loom_jit_amdgpu * jit_ = nullptr; +}; + +static size_t default_worker_count() { + const unsigned int hardware_threads = std::thread::hardware_concurrency(); + if (hardware_threads == 0) { + return 1; + } + return std::min(hardware_threads, 4); +} + +static bool compile_kernel(ggml_hrx_loom_jit_amdgpu * jit, + const LoomKernelCompileRequest & request, + const std::string & key, + ggml_hrx_loom_jit_compile_result & compiled, + std::string & error_message) { + std::vector configs; + configs.reserve(request.config_storage.size()); + for (const auto & config : request.config_storage) { + configs.push_back({ config.first.c_str(), config.second.c_str() }); + } + + ggml_hrx_loom_jit_compile_options compile_options = {}; + compile_options.source_data = request.source_data; + compile_options.source_size = request.source_size; + compile_options.source_format = request.source_format; + compile_options.source_identifier = request.source_identifier.c_str(); + compile_options.root_symbol = request.symbol.c_str(); + compile_options.launch_config_symbol = request.launch_config_symbol.c_str(); + compile_options.module_name = request.symbol.c_str(); + compile_options.artifact_identifier = request.symbol.c_str(); + compile_options.dependencies = request.dependencies.data(); + compile_options.dependency_count = request.dependencies.size(); + compile_options.config_bindings = configs.data(); + compile_options.config_binding_count = configs.size(); + compile_options.workload_arguments = request.workload.data(); + compile_options.workload_argument_count = request.workload.size(); + compile_options.evaluate_launch_config = true; + + if (ErrorResult error = take_status(ggml_hrx_loom_jit_amdgpu_compile(jit, &compile_options, &compiled))) { + error_message = "compile " + key + ": " + *error; + GGML_LOG_ERROR("%s: %s\n", __func__, error_message.c_str()); + return false; + } + return true; +} + +static bool is_disabled_value(const char * value) { + if (value == nullptr) { + return false; + } + return std::strcmp(value, "0") == 0 || std::strcmp(value, "false") == 0 || std::strcmp(value, "FALSE") == 0 || + std::strcmp(value, "off") == 0 || std::strcmp(value, "OFF") == 0; +} + +} // namespace + +class LoomSyncJit final : public LoomJit { + public: + LoomSyncJit(const char * target, std::string & error_message) { valid_ = jit_.create(target, error_message); } + + LoomCompiledKernelRef compile(std::string key, LoomKernelCompileRequest request) override { + auto compiled_ref = std::make_shared(std::move(key), std::move(request)); + if (!valid_) { + compiled_ref->complete({}, false, "Loom JIT is not initialized"); + return compiled_ref; + } + compile_ref(*compiled_ref, jit_.get()); + return compiled_ref; + } + + bool async_enabled() const override { return false; } + + private: + static void compile_ref(LoomCompiledKernel & compiled_ref, ggml_hrx_loom_jit_amdgpu * jit) { + ggml_hrx_loom_jit_compile_result compiled; + std::string error_message; + const bool success = compile_kernel(jit, compiled_ref.request(), compiled_ref.key(), compiled, error_message); + compiled_ref.complete(std::move(compiled), success, std::move(error_message)); + } + + LoomAmdgpuJit jit_; + bool valid_ = false; +}; + +class LoomAsyncJit final : public LoomJit { + public: + LoomAsyncJit(const char * target, std::string & error_message) : target_(target != nullptr ? target : "") { + if (target_.empty()) { + error_message = "missing HRX target"; + return; + } + valid_ = true; + } + + ~LoomAsyncJit() override { clear(); } + + LoomCompiledKernelRef compile(std::string key, LoomKernelCompileRequest request) override { + auto compiled_ref = std::make_shared(std::move(key), std::move(request)); + if (!valid_) { + compiled_ref->complete({}, false, "Loom JIT is not initialized"); + return compiled_ref; + } + + { + std::lock_guard lock(mutex_); + start_workers_locked(); + pending_.push_back(compiled_ref); + } + work_available_.notify_one(); + return compiled_ref; + } + + void clear() override { stop_workers(); } + + bool async_enabled() const override { return true; } + + private: + void start_workers_locked() { + if (!workers_.empty()) { + return; + } + shutdown_ = false; + const size_t count = default_worker_count(); + workers_.reserve(count); + for (size_t i = 0; i < count; ++i) { + workers_.emplace_back([this] { worker_loop(); }); + } + } + + void stop_workers() { + { + std::lock_guard lock(mutex_); + if (workers_.empty()) { + return; + } + shutdown_ = true; + } + work_available_.notify_all(); + for (std::thread & worker : workers_) { + if (worker.joinable()) { + worker.join(); + } + } + { + std::lock_guard lock(mutex_); + pending_.clear(); + workers_.clear(); + shutdown_ = false; + } + } + + void worker_loop() { + LoomAmdgpuJit worker_jit; + std::string jit_error; + const bool jit_ready = worker_jit.create(target_.c_str(), jit_error); + for (;;) { + LoomCompiledKernelRef compiled_ref; + { + std::unique_lock lock(mutex_); + work_available_.wait(lock, [&] { return shutdown_ || !pending_.empty(); }); + if (shutdown_ && pending_.empty()) { + return; + } + compiled_ref = std::move(pending_.front()); + pending_.pop_front(); + } + + if (!jit_ready) { + compiled_ref->complete({}, false, jit_error); + continue; + } + compile_ref(*compiled_ref, worker_jit.get()); + } + } + + static void compile_ref(LoomCompiledKernel & compiled_ref, ggml_hrx_loom_jit_amdgpu * jit) { + ggml_hrx_loom_jit_compile_result compiled; + std::string error_message; + const bool success = compile_kernel(jit, compiled_ref.request(), compiled_ref.key(), compiled, error_message); + compiled_ref.complete(std::move(compiled), success, std::move(error_message)); + } + + const std::string target_; + bool valid_ = false; + std::mutex mutex_; + std::condition_variable work_available_; + std::deque pending_; + std::vector workers_; + bool shutdown_ = false; +}; + +bool loom_async_jit_enabled_from_environment() { + static const bool enabled = [] { + const char * value = std::getenv("GGML_HRX_ASYNC_JIT"); + return !is_disabled_value(value); + }(); + return enabled; +} + +std::unique_ptr create_loom_jit(const char * target, LoomJitMode mode, std::string & error_message) { + if (target == nullptr || target[0] == '\0') { + error_message = "missing HRX target"; + GGML_LOG_ERROR("%s: %s\n", __func__, error_message.c_str()); + return nullptr; + } + if (mode == LoomJitMode::Async) { + auto jit = std::make_unique(target, error_message); + if (!error_message.empty()) { + GGML_LOG_ERROR("%s: %s\n", __func__, error_message.c_str()); + return nullptr; + } + return jit; + } + auto jit = std::make_unique(target, error_message); + if (!error_message.empty()) { + return nullptr; + } + return jit; +} + +std::unique_ptr create_loom_jit(const char * target, std::string & error_message) { + return create_loom_jit(target, loom_async_jit_enabled_from_environment() ? LoomJitMode::Async : LoomJitMode::Sync, + error_message); +} + +LoomCompiledKernel::LoomCompiledKernel(std::string key, LoomKernelCompileRequest request) : + key_(std::move(key)), + request_(std::move(request)) {} + +bool LoomCompiledKernel::resolve() const { + const State current = state(); + if (current != State::Pending) { + return current == State::Succeeded; + } + + std::unique_lock lock(mutex_); + complete_.wait(lock, [&] { return state() != State::Pending; }); + return state() == State::Succeeded; +} + +std::string LoomCompiledKernel::error_message() const { + std::lock_guard lock(mutex_); + return error_; +} + +ggml_hrx_loom_jit_compile_result LoomCompiledKernel::take_result() { + std::lock_guard lock(mutex_); + return std::move(compiled_); +} + +void LoomCompiledKernel::complete(ggml_hrx_loom_jit_compile_result compiled, bool success, std::string error) { + { + std::lock_guard lock(mutex_); + compiled_ = std::move(compiled); + error_ = std::move(error); + state_.store(success ? State::Succeeded : State::Failed, std::memory_order_release); + } + complete_.notify_all(); +} + +} // namespace ggml::hrx diff --git a/ggml/src/ggml-hrx/runtime/loom-kernel-jit.h b/ggml/src/ggml-hrx/runtime/loom-kernel-jit.h new file mode 100644 index 000000000000..627c20d3f23a --- /dev/null +++ b/ggml/src/ggml-hrx/runtime/loom-kernel-jit.h @@ -0,0 +1,95 @@ +#pragma once + +#include "loom-jit.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace ggml::hrx { + +struct LoomKernelCompileRequest { + const void * source_data = nullptr; + size_t source_size = 0; + ggml_hrx_loom_jit_source_format source_format = GGML_HRX_LOOM_JIT_SOURCE_FORMAT_TEXT; + std::string source_identifier; + std::string symbol; + std::string launch_config_symbol; + std::vector dependencies; + std::vector> config_storage; + std::vector workload; +}; + +class LoomCompiledKernel { + public: + LoomCompiledKernel(std::string key, LoomKernelCompileRequest request); + + LoomCompiledKernel(const LoomCompiledKernel &) = delete; + LoomCompiledKernel & operator=(const LoomCompiledKernel &) = delete; + + const std::string & key() const { return key_; } + + bool resolve() const; + std::string error_message() const; + ggml_hrx_loom_jit_compile_result take_result(); + + private: + friend class LoomSyncJit; + friend class LoomAsyncJit; + + const LoomKernelCompileRequest & request() const { return request_; } + + void complete(ggml_hrx_loom_jit_compile_result compiled, bool success, std::string error); + + enum class State { + Pending, + Succeeded, + Failed, + }; + + State state() const { return state_.load(std::memory_order_acquire); } + + mutable std::mutex mutex_; + mutable std::condition_variable complete_; + std::string key_; + LoomKernelCompileRequest request_; + ggml_hrx_loom_jit_compile_result compiled_; + std::string error_; + std::atomic state_ = State::Pending; +}; + +using LoomCompiledKernelRef = std::shared_ptr; + +class LoomJit { + public: + virtual ~LoomJit() = default; + + LoomJit(const LoomJit &) = delete; + LoomJit & operator=(const LoomJit &) = delete; + + virtual LoomCompiledKernelRef compile(std::string key, LoomKernelCompileRequest request) = 0; + + virtual void clear() {} + + virtual bool async_enabled() const = 0; + + protected: + LoomJit() = default; +}; + +enum class LoomJitMode { + Sync, + Async, +}; + +std::unique_ptr create_loom_jit(const char * target, LoomJitMode mode, std::string & error_message); +std::unique_ptr create_loom_jit(const char * target, std::string & error_message); +bool loom_async_jit_enabled_from_environment(); + +} // namespace ggml::hrx diff --git a/ggml/src/ggml-hrx/runtime/prepared-command-program-cache.cpp b/ggml/src/ggml-hrx/runtime/prepared-command-program-cache.cpp new file mode 100644 index 000000000000..ba8c7db983fe --- /dev/null +++ b/ggml/src/ggml-hrx/runtime/prepared-command-program-cache.cpp @@ -0,0 +1,200 @@ +#include "prepared-command-program-cache.h" + +#include +#include +#include +#include + +namespace ggml::hrx { +namespace { + +static void mix_hash(uint64_t & hash, uint64_t value) { + hash ^= value; + hash *= UINT64_C(1099511628211); +} + +static uint64_t hash_text(const char * text) { + uint64_t hash = UINT64_C(1469598103934665603); + if (text == nullptr) { + return hash; + } + while (*text != 0) { + mix_hash(hash, static_cast(*text)); + ++text; + } + return hash; +} + +static uint64_t hash_string(const std::string & text) { + uint64_t hash = UINT64_C(1469598103934665603); + for (const char c : text) { + mix_hash(hash, static_cast(c)); + } + return hash; +} + +static void apply_graph_replay_result(PreparedCommandProgramCacheExecutionResult & result, + const RecordedCommandGraphExecutionResult & replay) { + result.graph_replay_event = replay.event; + result.graph_replay_ineligible_reason = replay.ineligible_reason; + result.graph_replay_build_ns = replay.build_ns; + result.graph_replay_launch_ns = replay.launch_ns; + result.graph_replay_total_ns = replay.total_ns(); + result.graph_replay_dispatches = replay.dispatch_count; + result.graph_replay_transient_allocation_changed = replay.transient_allocation_changed; +} + +static bool graph_replay_should_fallback(HrxGraphReplayEvent event) { + return event == HrxGraphReplayEvent::Ineligible || event == HrxGraphReplayEvent::BuildFailed; +} + +} // namespace + +size_t PreparedCommandProgramCache::KeyHash::operator()(const Key & key) const { + uint64_t hash = UINT64_C(1469598103934665603); + mix_hash(hash, key.graph_uid); + mix_hash(hash, key.target_hash); + mix_hash(hash, key.command_shape_hash); + mix_hash(hash, key.bindings_hash); + return static_cast(hash); +} + +PreparedCommandProgramCache::Key PreparedCommandProgramCache::cache_key(uint64_t graph_uid, + const CommandProgramExecutionContext & context, + const std::string & command_shape, + const CommandProgramBindings & bindings) const { + return { + graph_uid, + hash_text(context.target), + hash_string(command_shape), + command_program_bindings_hash(bindings).value, + }; +} + +bool PreparedCommandProgramCache::execute(const CommandProgramExecutionContext & context, + uint64_t graph_uid, + const std::string & command_shape, + const CommandProgram & commands, + const CommandProgramBindings & bindings) { + return execute_with_result(context, graph_uid, command_shape, commands, bindings).success; +} + +PreparedCommandProgramCacheExecutionResult PreparedCommandProgramCache::execute_with_result( + const CommandProgramExecutionContext & context, + uint64_t graph_uid, + const std::string & command_shape, + const CommandProgram & commands, + const CommandProgramBindings & bindings) { + PreparedCommandProgramCacheExecutionResult result; + if (graph_uid == 0 || !commands.valid() || !bindings.valid()) { + result.graph_replay_event = HrxGraphReplayEvent::Ineligible; + result.graph_replay_ineligible_reason = "uncached_graph"; + PreparedCommandProgram prepared = prepare_command_program(context, commands, bindings); + if (!prepared.valid()) { + result.status.append(prepared.status); + return result; + } + result.success = bind_and_execute_prepared_command_program(context, commands, bindings, prepared); + if (!result.success) { + result.status.log("execute uncached HRX command program failed"); + } + return result; + } + + const Key key = cache_key(graph_uid, context, command_shape, bindings); + std::shared_ptr entry; + bool created_entry = false; + { + std::lock_guard lock(mutex_); + auto found = programs_.find(key); + if (found == programs_.end()) { + entry = std::make_shared(); + programs_.emplace(key, entry); + created_entry = true; + } else { + entry = found->second; + } + } + + std::lock_guard entry_lock(entry->mutex); + if (entry->has_program && entry->program.valid()) { + record_hit(); + const RecordedCommandGraphExecutionResult replay = + bind_and_launch_recorded_command_graph(context, commands, bindings, entry->program, entry->recorded); + apply_graph_replay_result(result, replay); + if (replay.success) { + result.success = true; + return result; + } + if (!graph_replay_should_fallback(replay.event)) { + result.status.append(replay.status); + if (result.status.success()) { + result.status.log("execute cached HRX graph replay failed"); + } + return result; + } + result.success = bind_and_execute_prepared_command_program(context, commands, bindings, entry->program); + if (!result.success) { + result.status.log("execute cached HRX command program failed"); + } + return result; + } + + PreparedCommandProgram prepared = prepare_command_program(context, commands, bindings); + if (!prepared.valid()) { + if (created_entry) { + std::lock_guard lock(mutex_); + const auto found = programs_.find(key); + if (found != programs_.end() && found->second == entry && !entry->has_program) { + programs_.erase(found); + } + } + result.status.append(prepared.status); + return result; + } + entry->program = std::move(prepared); + entry->has_program = true; + record_build(); + + const RecordedCommandGraphExecutionResult replay = + bind_and_launch_recorded_command_graph(context, commands, bindings, entry->program, entry->recorded); + apply_graph_replay_result(result, replay); + if (replay.success) { + result.success = true; + return result; + } + if (!graph_replay_should_fallback(replay.event)) { + result.status.append(replay.status); + if (result.status.success()) { + result.status.log("execute prepared HRX graph replay failed"); + } + return result; + } + result.success = bind_and_execute_prepared_command_program(context, commands, bindings, entry->program); + if (!result.success) { + result.status.log("execute prepared HRX command program failed"); + } + return result; +} + +PreparedCommandProgramCacheStats PreparedCommandProgramCache::stats() const { + std::lock_guard lock(mutex_); + return stats_; +} + +void PreparedCommandProgramCache::clear() { + std::lock_guard lock(mutex_); + programs_.clear(); +} + +void PreparedCommandProgramCache::record_build() { + std::lock_guard lock(mutex_); + ++stats_.builds; +} + +void PreparedCommandProgramCache::record_hit() { + std::lock_guard lock(mutex_); + ++stats_.hits; +} + +} // namespace ggml::hrx diff --git a/ggml/src/ggml-hrx/runtime/prepared-command-program-cache.h b/ggml/src/ggml-hrx/runtime/prepared-command-program-cache.h new file mode 100644 index 000000000000..68b61b0ce8f2 --- /dev/null +++ b/ggml/src/ggml-hrx/runtime/prepared-command-program-cache.h @@ -0,0 +1,89 @@ +#pragma once + +#include "command-program-executor.h" +#include "dispatch/command-program-bindings.h" +#include "dispatch/command-program.h" +#include "runtime/graph-replay.h" + +#include +#include +#include +#include +#include +#include + +namespace ggml::hrx { + +struct PreparedCommandProgramCacheStats { + uint64_t builds = 0; + uint64_t hits = 0; +}; + +struct PreparedCommandProgramCacheExecutionResult { + bool success = false; + Status status; + HrxGraphReplayEvent graph_replay_event = HrxGraphReplayEvent::Disabled; + std::string graph_replay_ineligible_reason; + uint64_t graph_replay_build_ns = 0; + uint64_t graph_replay_launch_ns = 0; + uint64_t graph_replay_total_ns = 0; + size_t graph_replay_dispatches = 0; + bool graph_replay_transient_allocation_changed = false; +}; + +class PreparedCommandProgramCache { + public: + bool execute(const CommandProgramExecutionContext & context, + uint64_t graph_uid, + const std::string & command_shape, + const CommandProgram & commands, + const CommandProgramBindings & bindings); + + PreparedCommandProgramCacheExecutionResult execute_with_result(const CommandProgramExecutionContext & context, + uint64_t graph_uid, + const std::string & command_shape, + const CommandProgram & commands, + const CommandProgramBindings & bindings); + + PreparedCommandProgramCacheStats stats() const; + + void clear(); + + private: + struct Key { + uint64_t graph_uid = 0; + uint64_t target_hash = 0; + uint64_t command_shape_hash = 0; + uint64_t bindings_hash = 0; + + bool operator==(const Key & other) const { + return graph_uid == other.graph_uid && target_hash == other.target_hash && + command_shape_hash == other.command_shape_hash && bindings_hash == other.bindings_hash; + } + }; + + struct KeyHash { + size_t operator()(const Key & key) const; + }; + + Key cache_key(uint64_t graph_uid, + const CommandProgramExecutionContext & context, + const std::string & command_shape, + const CommandProgramBindings & bindings) const; + + struct Entry { + std::mutex mutex; + PreparedCommandProgram program; + RecordedCommandGraph recorded; + bool has_program = false; + }; + + void record_build(); + void record_hit(); + + mutable std::mutex mutex_; + std::unordered_map, KeyHash> programs_; + PreparedCommandProgramCacheStats stats_; +}; + +} // namespace ggml::hrx diff --git a/ggml/src/ggml-hrx/runtime/transient-arena.cpp b/ggml/src/ggml-hrx/runtime/transient-arena.cpp new file mode 100644 index 000000000000..a950f4ab8c76 --- /dev/null +++ b/ggml/src/ggml-hrx/runtime/transient-arena.cpp @@ -0,0 +1,125 @@ +#include "transient-arena.h" + +#include "hrx-interop-utils.h" +#include "hrx_runtime.h" + +#include + +namespace ggml::hrx { + +TransientArena::AllocationLease::AllocationLease(TransientArena & arena, std::unique_lock lock) : + arena_(&arena), + lock_(std::move(lock)) {} + +Status TransientArena::AllocationLease::ensure_capacity(hrx_device_t device, + hrx_stream_t stream, + size_t required_size) { + Status status; + if (arena_ == nullptr || !lock_.owns_lock()) { + status.log("missing transient arena allocation lease"); + return status; + } + return arena_->ensure_capacity_locked(device, stream, required_size); +} + +TransientArenaAllocationRef TransientArena::AllocationLease::current_allocation() const { + if (arena_ == nullptr || !lock_.owns_lock()) { + return {}; + } + return arena_->current_allocation_locked(); +} + +TransientArena::~TransientArena() { + clear(); +} + +void TransientArena::clear() { + std::lock_guard lock(mutex_); + if (buffer_ != nullptr) { + hrx_buffer_release(buffer_); + buffer_ = nullptr; + } + allocation_capacity_ = 0; + allocation_id_ = kInvalidTransientArenaAllocationId; +} + +uint64_t TransientArena::next_allocation_id() { + const uint64_t id = next_allocation_id_++; + if (next_allocation_id_ == kInvalidTransientArenaAllocationId) { + ++next_allocation_id_; + } + return id; +} + +Status TransientArena::ensure_capacity(hrx_device_t device, hrx_stream_t stream, size_t required_size) { + std::lock_guard lock(mutex_); + return ensure_capacity_locked(device, stream, required_size); +} + +TransientArena::AllocationLease TransientArena::acquire_allocation_lease() { + return AllocationLease(*this, std::unique_lock(mutex_)); +} + +Status TransientArena::ensure_capacity_locked(hrx_device_t device, hrx_stream_t stream, size_t required_size) { + Status status; + if (required_size == 0 || allocation_capacity_ >= required_size) { + return status; + } + if (device == nullptr) { + status.log("missing HRX device for transient arena allocation"); + return status; + } + if (stream == nullptr) { + status.log("missing HRX stream for transient arena allocation"); + return status; + } + if (buffer_ != nullptr) { + if (ErrorResult error = take_status(hrx_stream_synchronize(stream))) { + status.log("synchronize before growing transient arena: %s", error->c_str()); + return status; + } + hrx_buffer_release(buffer_); + buffer_ = nullptr; + allocation_capacity_ = 0; + allocation_id_ = kInvalidTransientArenaAllocationId; + } + + hrx_buffer_params_t params = { + HRX_MEMORY_TYPE_DEVICE_LOCAL, + HRX_MEMORY_ACCESS_ALL, + HRX_BUFFER_USAGE_DEFAULT, + 0, + }; + hrx_buffer_t allocation = nullptr; + if (ErrorResult error = take_status( + hrx_allocator_allocate_buffer(hrx_device_allocator(device), params, required_size, &allocation))) { + status.log("allocate transient arena: %s", error->c_str()); + return status; + } + + buffer_ = allocation; + allocation_capacity_ = required_size; + allocation_id_ = next_allocation_id(); + return status; +} + +TransientArenaAllocationRef TransientArena::current_allocation() const { + std::lock_guard lock(mutex_); + return current_allocation_locked(); +} + +size_t TransientArena::capacity() const { + std::lock_guard lock(mutex_); + return allocation_capacity_; +} + +uint64_t TransientArena::allocation_id() const { + std::lock_guard lock(mutex_); + return allocation_id_; +} + +TransientArenaAllocationRef TransientArena::current_allocation_locked() const { + return { buffer_, allocation_capacity_, allocation_id_ }; +} + +} // namespace ggml::hrx diff --git a/ggml/src/ggml-hrx/runtime/transient-arena.h b/ggml/src/ggml-hrx/runtime/transient-arena.h new file mode 100644 index 000000000000..a529bf938f10 --- /dev/null +++ b/ggml/src/ggml-hrx/runtime/transient-arena.h @@ -0,0 +1,66 @@ +#pragma once + +#include "dispatch/command-program-resolver.h" +#include "status.h" + +#include +#include +#include + +typedef struct hrx_device_s * hrx_device_t; +typedef struct hrx_stream_s * hrx_stream_t; + +namespace ggml::hrx { + +class TransientArena { + public: + class AllocationLease { + public: + AllocationLease() = default; + AllocationLease(AllocationLease &&) noexcept = default; + AllocationLease & operator=(AllocationLease &&) noexcept = default; + + AllocationLease(const AllocationLease &) = delete; + AllocationLease & operator=(const AllocationLease &) = delete; + + Status ensure_capacity(hrx_device_t device, hrx_stream_t stream, size_t required_size); + TransientArenaAllocationRef current_allocation() const; + + private: + friend class TransientArena; + + AllocationLease(TransientArena & arena, std::unique_lock lock); + + TransientArena * arena_ = nullptr; + std::unique_lock lock_; + }; + + TransientArena() = default; + ~TransientArena(); + + TransientArena(const TransientArena &) = delete; + TransientArena & operator=(const TransientArena &) = delete; + + Status ensure_capacity(hrx_device_t device, hrx_stream_t stream, size_t required_size); + AllocationLease acquire_allocation_lease(); + void clear(); + + TransientArenaAllocationRef current_allocation() const; + + size_t capacity() const; + + uint64_t allocation_id() const; + + private: + uint64_t next_allocation_id(); + Status ensure_capacity_locked(hrx_device_t device, hrx_stream_t stream, size_t required_size); + TransientArenaAllocationRef current_allocation_locked() const; + + mutable std::mutex mutex_; + hrx_buffer_t buffer_ = nullptr; + size_t allocation_capacity_ = 0; + uint64_t allocation_id_ = kInvalidTransientArenaAllocationId; + uint64_t next_allocation_id_ = 1; +}; + +} // namespace ggml::hrx diff --git a/ggml/src/ggml-hrx/status.h b/ggml/src/ggml-hrx/status.h new file mode 100644 index 000000000000..6d6aafc8ab6e --- /dev/null +++ b/ggml/src/ggml-hrx/status.h @@ -0,0 +1,72 @@ +#pragma once + +#include +#include +#include +#include +#include +#include + +namespace ggml::hrx { + +class [[nodiscard]] Status { + public: + Status() = default; + + Status(const Status &) = delete; + Status & operator=(const Status &) = delete; + + Status(Status &&) noexcept = default; + Status & operator=(Status &&) noexcept = default; + + bool success() const { return messages_.empty(); } + + const std::vector & errors() const { return messages_; } + + void append(const Status & other) { + messages_.insert(messages_.end(), other.messages_.begin(), other.messages_.end()); + } + + void log(const char * format, ...) { + va_list args; + va_start(args, format); + log_va(format, args); + va_end(args); + } + + private: + void push(std::string message) { messages_.push_back(std::move(message)); } + + void log_va(const char * format, va_list args) { + if (format == nullptr) { + push("failed to format error message"); + return; + } + + char stack[256]; + va_list args_copy; + va_copy(args_copy, args); + const int written = std::vsnprintf(stack, sizeof(stack), format, args_copy); + va_end(args_copy); + if (written < 0) { + push("failed to format error message"); + return; + } + if (static_cast(written) < sizeof(stack)) { + push(std::string(stack, static_cast(written))); + return; + } + + std::vector buffer(static_cast(written) + 1); + const int rewritten = std::vsnprintf(buffer.data(), buffer.size(), format, args); + if (rewritten < 0) { + push("failed to format error message"); + return; + } + push(std::string(buffer.data(), static_cast(rewritten))); + } + + std::vector messages_; +}; + +} // namespace ggml::hrx diff --git a/ggml/src/ggml-hrx/tools/analyze-graph.cpp b/ggml/src/ggml-hrx/tools/analyze-graph.cpp new file mode 100644 index 000000000000..8ae0b7febfb2 --- /dev/null +++ b/ggml/src/ggml-hrx/tools/analyze-graph.cpp @@ -0,0 +1,159 @@ +#include "dispatch/command-program-diagnostics.h" +#include "dispatch/command-program.h" +#include "dispatch/dispatch-scheduler.h" +#include "graph/graph-diagnostics.h" +#include "kernel-corpus/kernel-corpus.h" + +#include +#include +#include +#include +#include +#include + +namespace { + +struct Options { + std::filesystem::path graph; + std::filesystem::path output_directory; + std::string target; +}; + +static void print_usage(const char * program) { + std::cerr << "usage: " << program << " --graph [--target ] [--out ]\n"; +} + +static bool parse_options(int argc, char ** argv, Options & options) { + for (int i = 1; i < argc; ++i) { + const std::string arg = argv[i]; + if (arg == "--help" || arg == "-h") { + print_usage(argv[0]); + return false; + } + if (arg == "--graph" && i + 1 < argc) { + options.graph = argv[++i]; + continue; + } + if (arg == "--target" && i + 1 < argc) { + options.target = argv[++i]; + continue; + } + if (arg == "--out" && i + 1 < argc) { + options.output_directory = argv[++i]; + continue; + } + std::cerr << "unknown or incomplete argument: " << arg << '\n'; + print_usage(argv[0]); + return false; + } + if (options.graph.empty()) { + std::cerr << "--graph is required\n"; + print_usage(argv[0]); + return false; + } + return true; +} + +static std::string read_file(const std::filesystem::path & path) { + std::ifstream input(path, std::ios::binary); + if (!input) { + return {}; + } + return { std::istreambuf_iterator(input), std::istreambuf_iterator() }; +} + +static void write_file(const std::filesystem::path & path, const std::string & contents) { + std::filesystem::create_directories(path.parent_path()); + std::ofstream output(path, std::ios::binary | std::ios::trunc); + if (!output) { + throw std::runtime_error("cannot create " + path.string()); + } + output << contents; + if (contents.empty() || contents.back() != '\n') { + output << '\n'; + } +} + +static void write_output(const Options & options, const std::string & name, const std::string & contents) { + if (options.output_directory.empty()) { + std::cout << "== " << name << " ==\n" << contents; + if (contents.empty() || contents.back() != '\n') { + std::cout << '\n'; + } + return; + } + write_file(options.output_directory / name, contents); +} + +} // namespace + +int main(int argc, char ** argv) { + Options options; + if (!parse_options(argc, argv, options)) { + return 1; + } + + const std::string contents = read_file(options.graph); + if (contents.empty()) { + std::cerr << "cannot read graph snapshot: " << options.graph << '\n'; + return 1; + } + + ggml::hrx::GraphSnapshotLoadResult snapshot = ggml::hrx::load_graph_snapshot_json(contents); + if (!snapshot.valid()) { + for (const std::string & error : snapshot.status.errors()) { + std::cerr << error << '\n'; + } + return 1; + } + + if (options.target.empty()) { + options.target = snapshot.target; + } + if (options.target.empty()) { + std::cerr << "target is missing from both --target and graph snapshot\n"; + return 1; + } + + write_output(options, "graph.txt", + ggml::hrx::format_graph_snapshot_text(snapshot.graph, options.target, snapshot.uid)); + write_output(options, "graph.json", + ggml::hrx::serialize_graph_snapshot_json(snapshot.graph, options.target, snapshot.uid)); + + ggml::hrx::DispatchScheduler scheduler; + ggml::hrx::DispatchScheduleDiagnostics diagnostics; + const bool scheduled = scheduler.schedule_graph(snapshot.graph, { options.target }, &diagnostics); + write_output(options, "schedule.txt", + ggml::hrx::format_schedule_diagnostics_text(snapshot.graph, scheduler.plan(), diagnostics)); + write_output(options, "schedule.json", + ggml::hrx::serialize_schedule_diagnostics_json(snapshot.graph, scheduler.plan(), diagnostics)); + if (!scheduled) { + write_output(options, "unmatched.txt", + ggml::hrx::format_schedule_diagnostics_text(snapshot.graph, scheduler.plan(), diagnostics)); + write_output(options, "unmatched.json", + ggml::hrx::serialize_schedule_diagnostics_json(snapshot.graph, scheduler.plan(), diagnostics)); + return 2; + } + + const ggml::hrx::KernelCorpus & corpus = ggml::hrx::get_qwen_kernel_corpus(); + ggml::hrx::CommandProgram commands = + ggml::hrx::build_command_program(snapshot.graph, scheduler.plan(), corpus, options.target); + write_output(options, "commands.txt", ggml::hrx::format_command_program(commands)); + if (!commands.valid()) { + for (const std::string & error : commands.status.errors()) { + std::cerr << error << '\n'; + } + return 3; + } + + const ggml::hrx::VerificationResult verification = + ggml::hrx::verify_command_program(commands, corpus, options.target); + if (!verification.valid()) { + for (const std::string & error : verification.status.errors()) { + std::cerr << error << '\n'; + } + return 4; + } + + return 0; +} diff --git a/ggml/src/ggml-hrx/tools/compile-kernel.cpp b/ggml/src/ggml-hrx/tools/compile-kernel.cpp new file mode 100644 index 000000000000..390688eba0d9 --- /dev/null +++ b/ggml/src/ggml-hrx/tools/compile-kernel.cpp @@ -0,0 +1,110 @@ +#include "../loom-jit.h" +#include "tool-utils.h" + +#include +#include +#include +#include +#include + +using ggml::hrx::tool::report_status; +using ggml::hrx::tool::write_file; + +int main(int argc, char ** argv) { + std::string target; + std::string source_path; + std::string root; + std::string output_path; + std::vector config_storage; + std::vector workload; + for (int i = 1; i < argc; ++i) { + const std::string argument = argv[i]; + if (argument == "--target" && i + 1 < argc) { + target = argv[++i]; + } else if (argument == "--source" && i + 1 < argc) { + source_path = argv[++i]; + } else if (argument == "--root" && i + 1 < argc) { + root = argv[++i]; + } else if (argument == "--output" && i + 1 < argc) { + output_path = argv[++i]; + } else if (argument == "--config" && i + 1 < argc) { + config_storage.emplace_back(argv[++i]); + } else if (argument == "--workload" && i + 1 < argc) { + workload.push_back(std::stoll(argv[++i])); + } else { + std::cerr << "unknown or incomplete argument: " << argument << '\n'; + return 2; + } + } + if (target.empty() || source_path.empty() || root.empty() || output_path.empty()) { + std::cerr << "usage: ggml-hrx-compile-kernel --target gfx... --source linked.loom --root symbol " + "--output dir [--config key=value] [--workload value]\n"; + return 2; + } + const std::string source = ggml::hrx::tool::read_file(source_path); + if (source.empty()) { + std::cerr << "cannot read Loom source: " << source_path << '\n'; + return 2; + } + std::vector configs; + std::vector config_keys; + std::vector config_values; + for (const std::string & item : config_storage) { + const size_t equals = item.find('='); + if (equals == std::string::npos) { + std::cerr << "invalid config binding: " << item << '\n'; + return 2; + } + config_keys.push_back(item.substr(0, equals)); + config_values.push_back(item.substr(equals + 1)); + } + for (size_t i = 0; i < config_keys.size(); ++i) { + configs.push_back({ config_keys[i].c_str(), config_values[i].c_str() }); + } + + ggml_hrx_loom_jit_amdgpu_options jit_options; + jit_options.processor = target.c_str(); + jit_options.identifier = target.c_str(); + ggml_hrx_loom_jit_amdgpu * jit = nullptr; + hrx_status_t status = ggml_hrx_loom_jit_amdgpu_create(&jit_options, &jit); + if (!report_status(status, "create JIT")) { + return 1; + } + + ggml_hrx_loom_jit_compile_options options; + options.source_data = source.data(); + options.source_size = source.size(); + options.source_format = GGML_HRX_LOOM_JIT_SOURCE_FORMAT_TEXT; + options.source_identifier = source_path.c_str(); + options.root_symbol = root.c_str(); + options.module_name = root.c_str(); + options.artifact_identifier = root.c_str(); + options.config_bindings = configs.data(); + options.config_binding_count = configs.size(); + options.workload_arguments = workload.data(); + options.workload_argument_count = workload.size(); + options.evaluate_launch_config = !workload.empty(); + ggml_hrx_loom_jit_compile_result result; + status = ggml_hrx_loom_jit_amdgpu_compile(jit, &options, &result); + if (!hrx_status_is_ok(status)) { + ggml_hrx_loom_jit_amdgpu_release(jit); + report_status(status, "compile kernel"); + return 1; + } + std::filesystem::create_directories(output_path); + write_file(std::filesystem::path(output_path) / "kernel.hsaco", result.hsaco_data, result.hsaco_size); + write_file(std::filesystem::path(output_path) / "manifest.json", result.manifest_json, result.manifest_json_size); + write_file(std::filesystem::path(output_path) / "compile-report.json", result.compile_report_json, + result.compile_report_json_size); + write_file(std::filesystem::path(output_path) / "final.loom", result.final_module_text, + result.final_module_text_size); + std::ofstream launch(std::filesystem::path(output_path) / "launch.txt", std::ios::trunc); + launch << "target=" << target << '\n' + << "root=" << root << '\n' + << "workgroups=" << result.launch_config.workgroup_count[0] << ',' << result.launch_config.workgroup_count[1] + << ',' << result.launch_config.workgroup_count[2] << '\n' + << "workgroup_size=" << result.launch_config.workgroup_size[0] << ',' + << result.launch_config.workgroup_size[1] << ',' << result.launch_config.workgroup_size[2] << '\n'; + ggml_hrx_loom_jit_amdgpu_release(jit); + return 0; +} diff --git a/ggml/src/ggml-hrx/tools/compile_qwen_kernel_corpus.py b/ggml/src/ggml-hrx/tools/compile_qwen_kernel_corpus.py new file mode 100755 index 000000000000..31da06085261 --- /dev/null +++ b/ggml/src/ggml-hrx/tools/compile_qwen_kernel_corpus.py @@ -0,0 +1,256 @@ +#!/usr/bin/env python3 +"""Compiles the pinned Qwen corpus using only BUILD.bazel-authored recipes.""" + +from __future__ import annotations + +import argparse +import json +import os +import pathlib +import re +import subprocess +import sys + + +BENCHMARK_RE = re.compile( + r"check\.benchmark<@(?P[A-Za-z0-9_]+)>\s+@(?P[A-Za-z0-9_]+)" + r"(?:\s*\{(?P[^}]*)\})?" +) +ATTR_RE = re.compile(r"(?P[A-Za-z0-9_.]+)\s*=\s*(?P-?[0-9]+)") +CASE_RE = re.compile(r"check\.case(?:\s+public)?\s+@(?P[A-Za-z0-9_]+)\s*\{") +LITERAL_RE = re.compile(r"%(?P[A-Za-z0-9_]+)\s*=\s*check\.literal\s+value\((?P-?[0-9]+)\)\s*:\s*index") +FUNC_CALL_RE = re.compile(r"func\.call\s+@(?P[A-Za-z0-9_]+)\s*\(") + + +def run(command: list[str], *, env: dict[str, str] | None = None) -> subprocess.CompletedProcess[str]: + return subprocess.run(command, text=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE, env=env, check=False) + + +def require_run(command: list[str], *, env: dict[str, str] | None = None) -> subprocess.CompletedProcess[str]: + result = run(command, env=env) + if result.returncode: + raise RuntimeError(f"command failed ({result.returncode}): {' '.join(command)}\n{result.stderr}") + return result + + +def main() -> int: + parser = argparse.ArgumentParser() + parser.add_argument("--manifest", type=pathlib.Path, required=True) + parser.add_argument("--corpus-dir", type=pathlib.Path, required=True) + parser.add_argument("--hrx-source", type=pathlib.Path, required=True) + parser.add_argument("--loom-link", type=pathlib.Path, required=True) + parser.add_argument("--benchmark-tool", type=pathlib.Path, required=True) + parser.add_argument("--compiler", type=pathlib.Path, required=True) + parser.add_argument("--target", choices=("gfx1100", "gfx1151"), required=True) + parser.add_argument("--output-dir", type=pathlib.Path, required=True) + parser.add_argument("--schedule", type=pathlib.Path, action="append", default=[], + help="materialized program.json whose exact specializations must compile") + args = parser.parse_args() + + manifest = json.loads(args.manifest.read_text(encoding="utf-8")) + revision = require_run(["git", "-C", str(args.hrx_source), "rev-parse", "HEAD"]).stdout.strip() + if revision != manifest["upstream_revision"]: + raise RuntimeError(f"HRX checkout {revision} does not match corpus {manifest['upstream_revision']}") + if require_run(["git", "-C", str(args.hrx_source), "status", "--porcelain"]).stdout.strip(): + raise RuntimeError("refusing to compile against a dirty pinned HRX checkout") + override = os.environ.get("HSA_OVERRIDE_GFX_VERSION") + if args.target == "gfx1100" and override != "11.0.0": + raise RuntimeError("gfx1100 validation on this gfx1151 host requires HSA_OVERRIDE_GFX_VERSION=11.0.0") + if args.target == "gfx1151" and override: + raise RuntimeError("native gfx1151 validation must run without HSA_OVERRIDE_GFX_VERSION") + + args.output_dir.mkdir(parents=True, exist_ok=True) + modules = {item["name"]: item for item in manifest["link_modules"]} + linked_dir = args.output_dir / "linked" + linked_dir.mkdir(exist_ok=True) + linked_paths: dict[str, pathlib.Path] = {} + + def link_module(name: str) -> pathlib.Path: + if name in linked_paths: + return linked_paths[name] + recipe = modules[name] + output = linked_dir / f"{name}.loom" + command = [str(args.loom_link), "--mode=archive", f"--output={output}"] + command.extend(str(args.corpus_dir / source) for source in recipe["srcs"]) + for library in recipe["libraries"]: + path = link_module(library[1:]) if library.startswith(":") else args.corpus_dir / library + command.append(f"--library={path}") + require_run(command) + linked_paths[name] = output + return output + + exports = {item["symbol"]: item for item in manifest["exports"]} + + def source_for_root(root: str) -> pathlib.Path: + source_name = exports[root]["source"] + direct_modules = [name for name, recipe in modules.items() if source_name in recipe["srcs"]] + if direct_modules: + return link_module(direct_modules[0]) + library_modules = [name for name, recipe in modules.items() if source_name in recipe["libraries"]] + if library_modules: + return link_module(library_modules[0]) + return args.corpus_dir / source_name + + results: list[dict[str, object]] = [] + compiled_keys: set[tuple[str, tuple[int, ...], tuple[str, ...]]] = set() + root_sources: dict[str, pathlib.Path] = {} + planned_invocation_count = 0 + resolved_invocation_count = 0 + for case in manifest["plan_cases"]: + source = link_module(case["link_module"]) if "link_module" in case else args.corpus_dir / case["source"] + plan_command = [str(args.benchmark_tool)] + for item in case["args"]: + if item.startswith("$(location"): + plan_command.append(str(source)) + else: + plan_command.append(item) + plan = require_run(plan_command) + case_dir = args.output_dir / "recipes" / case["name"] + case_dir.mkdir(parents=True, exist_ok=True) + (case_dir / "plan.jsonl").write_text(plan.stdout, encoding="utf-8") + (case_dir / "plan.stderr.txt").write_text(plan.stderr, encoding="utf-8") + plan_rows = [json.loads(line) for line in plan.stdout.splitlines() if line.strip()] + plan_rows = [row for row in plan_rows if row.get("row") == "plan"] + if not plan_rows: + raise RuntimeError(f"BUILD recipe {case['name']} produced no planner rows") + source_text = source.read_text(encoding="utf-8") + benchmark_defs = { + match.group("name"): (match.group("case"), + {item.group("name"): int(item.group("value")) for item in ATTR_RE.finditer(match.group("attrs") or "")}) + for match in BENCHMARK_RE.finditer(source.read_text(encoding="utf-8")) + } + case_literals: dict[str, dict[str, int]] = {} + case_calls: dict[str, list[str]] = {} + for match in CASE_RE.finditer(source_text): + depth = 1 + cursor = match.end() + while cursor < len(source_text) and depth: + depth += source_text[cursor] == "{" + depth -= source_text[cursor] == "}" + cursor += 1 + body = source_text[match.end():cursor - 1] + case_literals[match.group("name")] = { + item.group("name"): int(item.group("value")) for item in LITERAL_RE.finditer(body) + } + case_calls[match.group("name")] = [item.group("name") for item in FUNC_CALL_RE.finditer(body)] + configs = [item.removeprefix("--config=") for item in case["args"] if item.startswith("--config=")] + selects_benchmark = any(item.startswith("--benchmark=") for item in case["args"]) + owned_sources = set(modules[case["link_module"]]["srcs"]) if "link_module" in case else {case["source"]} + for row in plan_rows: + benchmark_case, attrs = benchmark_defs.get(row["benchmark"], (row["case"], {})) + concrete_values = dict(case_literals.get(benchmark_case, {})) + concrete_values.update(attrs) + expected_count = int(row.get("actual_invocation_count", 1)) + planned_invocation_count += expected_count + if row.get("actual_entry"): + roots = [row["actual_entry"]] + else: + roots = [root for root in case_calls.get(benchmark_case, []) if root in exports] + if len(roots) != expected_count: + raise RuntimeError( + f"BUILD recipe {case['name']} planner reports {expected_count} invocations for " + f"{benchmark_case}, but source resolves {len(roots)} exported calls: {roots}" + ) + resolved_invocation_count += len(roots) + for root in roots: + if root not in exports: + raise RuntimeError(f"BUILD recipe {case['name']} selected unmanifested kernel {root}") + if not selects_benchmark and exports[root]["source"] not in owned_sources: + continue + root_sources.setdefault(root, source) + parameter_names = [item["name"] for item in exports[root]["workload_parameters"]] + missing = [name for name in parameter_names if name not in concrete_values] + if missing: + raise RuntimeError(f"{case['name']} recipe omits concrete {missing} for {root}") + workload = tuple(concrete_values[name] for name in parameter_names) + key = (root, workload, tuple(configs)) + if key in compiled_keys: + continue + compiled_keys.add(key) + compile_dir = case_dir / f"{root}-{len(compiled_keys):03d}" + command = [str(args.compiler), "--target", args.target, "--source", str(source), + "--root", root, "--output", str(compile_dir)] + for config in configs: + command.extend(("--config", config)) + for value in workload: + command.extend(("--workload", str(value))) + compile_result = run(command, env=os.environ.copy()) + (case_dir / f"{root}-{len(compiled_keys):03d}.stdout.txt").write_text(compile_result.stdout, encoding="utf-8") + (case_dir / f"{root}-{len(compiled_keys):03d}.stderr.txt").write_text(compile_result.stderr, encoding="utf-8") + result = {"case": case["name"], "root": root, "workload": workload, + "configs": configs, "artifact_dir": str(compile_dir), + "status": "ok" if compile_result.returncode == 0 else "failed"} + results.append(result) + + schedule_requirement_count = 0 + schedule_unique_requirement_count = 0 + for schedule_path in args.schedule: + schedule = json.loads(schedule_path.read_text(encoding="utf-8")) + schedule_dir = args.output_dir / "schedules" / schedule_path.parent.name + schedule_dir.mkdir(parents=True, exist_ok=True) + for invocation in schedule.get("invocations", []): + for dispatch in invocation.get("dispatches", []): + specialization = dispatch["kernel"] + if specialization.get("execution") != "native": + raise RuntimeError(f"schedule {schedule_path} contains non-native dispatch {specialization['variant']}") + root = specialization["variant"] + if root not in exports: + raise RuntimeError(f"schedule {schedule_path} selects unmanifested kernel {root}") + schedule_requirement_count += 1 + parameters = specialization.get("parameters", {}) + parameter_names = [item["name"] for item in exports[root]["workload_parameters"]] + missing = [name for name in parameter_names if name not in parameters] + if missing: + raise RuntimeError(f"schedule {schedule_path} omits concrete {missing} for {root}") + workload = tuple(int(parameters[name]) for name in parameter_names) + configs = tuple(f"{name}={value}" for name, value in + sorted(specialization.get("compile_parameters", {}).items())) + key = (root, workload, configs) + if key in compiled_keys: + continue + schedule_unique_requirement_count += 1 + compiled_keys.add(key) + source = root_sources[root] if root in root_sources else source_for_root(root) + compile_dir = schedule_dir / f"{root}-{schedule_unique_requirement_count:03d}" + command = [str(args.compiler), "--target", args.target, "--source", str(source), + "--root", root, "--output", str(compile_dir)] + for config in configs: + command.extend(("--config", config)) + for value in workload: + command.extend(("--workload", str(value))) + compile_result = run(command, env=os.environ.copy()) + (schedule_dir / f"{root}-{schedule_unique_requirement_count:03d}.stdout.txt").write_text( + compile_result.stdout, encoding="utf-8") + (schedule_dir / f"{root}-{schedule_unique_requirement_count:03d}.stderr.txt").write_text( + compile_result.stderr, encoding="utf-8") + results.append({"case": f"schedule:{schedule_path}", "root": root, "workload": workload, + "configs": configs, "artifact_dir": str(compile_dir), + "status": "ok" if compile_result.returncode == 0 else "failed"}) + + summary = { + "schema": "ggml-hrx-qwen-compile-report-v1", + "target": args.target, + "hsa_override_gfx_version": override, + "hrx_revision": revision, + "corpus_digest": manifest["corpus_sha256"], + "recipe_digest": manifest["build_bazel_sha256"], + "plan_case_count": len(manifest["plan_cases"]), + "planned_invocation_count": planned_invocation_count, + "resolved_invocation_count": resolved_invocation_count, + "schedule_requirement_count": schedule_requirement_count, + "schedule_unique_requirement_count": schedule_unique_requirement_count, + "compile_count": len(results), + "failed_count": sum(item["status"] != "ok" for item in results), + "results": results, + } + (args.output_dir / "summary.json").write_text(json.dumps(summary, indent=2) + "\n", encoding="utf-8") + print(json.dumps({key: summary[key] for key in ("target", "plan_case_count", "compile_count", "failed_count")})) + return 0 if summary["failed_count"] == 0 and planned_invocation_count == resolved_invocation_count else 1 + + +if __name__ == "__main__": + try: + raise SystemExit(main()) + except (OSError, RuntimeError, ValueError, json.JSONDecodeError) as error: + print(f"compile corpus: {error}", file=sys.stderr) + raise SystemExit(2) diff --git a/ggml/src/ggml-hrx/tools/generate_kernel_corpus.py b/ggml/src/ggml-hrx/tools/generate_kernel_corpus.py new file mode 100644 index 000000000000..b1074791401a --- /dev/null +++ b/ggml/src/ggml-hrx/tools/generate_kernel_corpus.py @@ -0,0 +1,536 @@ +#!/usr/bin/env python3 + +import argparse +import hashlib +import json +import pathlib +import re +import subprocess +import sys +import tempfile +from typing import Dict, Iterable, List, Tuple + + +DEFAULT_KERNEL_FAMILY = "qwen3_moe" + +SOURCE_ARRAY_TEMPLATE = """static const unsigned char {symbol}[] = {{ +{bytes} +}}; +static constexpr size_t {symbol}Size = {size}; +""" + +DEPENDENCY_TABLE_TEMPLATE = """static const KernelSourceSpan {dependency_table}[] = {{ +{dependencies} +}}; +""" + +SOURCE_RECORD_TEMPLATE = """static const KernelSource {record} = {{ + {{ reinterpret_cast({source_symbol}), {source_symbol}Size, {source_format} }}, + {dependency_table}, + {dependency_count}, +}}; +""" + +CORPUS_ARRAY_TEMPLATE = """static {type} {symbol}[] = {{ +{values} +}}; +""" + +KERNEL_RECORD_TEMPLATE = """ {{ + {family}, + {name}, + kernel_catalog_id({family}, {name}), + {source}, + {dependencies}, + {symbol}, + "amdgpu", + {target_selector}, + {{ nullptr, 0 }}, + {scalar_parameters}, + {bindings}, + {source_digest}, + {workload_parameters}, + {launch_parameters}, + {{ + {compile_mode}, + {link_module}, + {primary_sources}, + {library_sources}, + }}, + }},""" + +SOURCE_DATA_TEMPLATE = """{source_arrays} +{dependency_tables} +{source_records} +static const KernelSourceRecordEntry kKernelSourceRecords[] = {{ +{lookup_entries} +}}; +""" + +CORPUS_DATA_TEMPLATE = """{kernel_arrays} +static const KernelDefinition kQwenKernelDefinitions[] = {{ +{kernel_records} +}}; + +static const KernelCorpus kQwenKernelCorpus = {{ + "ggml-hrx-kernel-corpus-v2", + {upstream_revision}, + {corpus_digest}, + {recipe_digest}, + {plan_case_count}, + {{ kQwenKernelDefinitions, {kernel_count} }}, +}}; +""" + +CATALOG_DATA_TEMPLATE = """struct KernelCatalogEntry {{ + const char * family; + const char * name; +}}; + +static constexpr KernelCatalogEntry kKernelCatalogEntries[] = {{ +{kernel_entries} +}}; + +constexpr bool kernel_catalog_entry_exists(const char * family, const char * name) {{ + for (const KernelCatalogEntry & known : kKernelCatalogEntries) {{ + if (kernel_catalog_name_equal(family, known.family) && kernel_catalog_name_equal(name, known.name)) {{ + return true; + }} + }} + return false; +}} +""" + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description="Generate embedded kernel corpus data includes.") + parser.add_argument("--source-output", type=pathlib.Path, required=True) + parser.add_argument("--corpus-output", type=pathlib.Path, required=True) + parser.add_argument("--catalog-output", type=pathlib.Path, required=True) + parser.add_argument("--manifest", type=pathlib.Path, required=True) + parser.add_argument("--corpus-dir", type=pathlib.Path, required=True) + parser.add_argument("--source-format", choices=("text", "binary"), default="text") + parser.add_argument("--loom-link", type=pathlib.Path) + parser.add_argument("--loom-format", type=pathlib.Path) + parser.add_argument("--depfile", type=pathlib.Path) + return parser.parse_args() + + +def read_text(path: pathlib.Path) -> str: + try: + return path.read_text() + except OSError as exc: + raise RuntimeError(f"failed to read {path}: {exc}") from exc + + +def read_bytes(path: pathlib.Path) -> bytes: + try: + return path.read_bytes() + except OSError as exc: + raise RuntimeError(f"failed to read {path}: {exc}") from exc + + +def run_tool(command: List[str]) -> subprocess.CompletedProcess[str]: + return subprocess.run(command, text=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE, check=False) + + +def require_tool(command: List[str]) -> None: + result = run_tool(command) + if result.returncode: + raise RuntimeError(f"command failed ({result.returncode}): {' '.join(command)}\n{result.stderr}") + + +def convert_source_to_bytecode(source_path: pathlib.Path, loom_link: pathlib.Path, loom_format: pathlib.Path) -> bytes: + with tempfile.TemporaryDirectory(prefix="ggml-hrx-loom-") as temp_dir_name: + temp_dir = pathlib.Path(temp_dir_name) + stripped = temp_dir / "stripped.loom" + bytecode = temp_dir / "stripped.loombc" + require_tool([ + str(loom_link), + "--verify=false", + "--mode=merge", + "--strip-check", + "--to=text", + f"--output={stripped}", + str(source_path), + ]) + format_result = run_tool([ + str(loom_format), + "--from=text", + "--to=bc", + f"--output={bytecode}", + str(stripped), + ]) + if format_result.returncode: + require_tool([ + str(loom_link), + "--verify=false", + "--mode=merge", + "--strip-check", + "--to=bc", + f"--output={bytecode}", + str(source_path), + ]) + return read_bytes(bytecode) + + +def sanitize_symbol(path: str, index: int) -> str: + stem = re.sub(r"[^0-9A-Za-z_]", "_", path) + if not stem or stem[0].isdigit(): + stem = f"_{stem}" + return f"kernel_source_{index}_{stem}" + + +def format_byte_array(data: bytes) -> str: + if not data: + return "" + lines = [] + for offset in range(0, len(data), 16): + chunk = data[offset : offset + 16] + lines.append(" " + ", ".join(f"0x{byte:02x}" for byte in chunk) + ",") + return "\n".join(lines) + "\n" + + +def escape_cpp_string(text: str) -> str: + return text.replace("\\", "\\\\").replace('"', '\\"') + + +def cpp_string(text: str) -> str: + return f"\"{escape_cpp_string(text)}\"" + + +def resource_access_value(access: str) -> str: + if access == "read": + return "ResourceAccess::Read" + if access == "write": + return "ResourceAccess::Write" + if access == "read_write": + return "ResourceAccess::ReadWrite" + raise RuntimeError(f"invalid kernel binding access metadata: {access}") + + +def span_initializer(symbol: str, count: int) -> str: + if count == 0: + return "{ nullptr, 0 }" + return "{ " + symbol + ", " + str(count) + " }" + + +def typed_array(symbol: str, value_type: str, values: List[str]) -> Tuple[str, str]: + if not values: + return "", "{ nullptr, 0 }" + array = CORPUS_ARRAY_TEMPLATE.format( + type=value_type, + symbol=symbol, + values="\n".join(" " + value + "," for value in values), + ) + return array, span_initializer(symbol, len(values)) + + +def string_array(symbol: str, items: Iterable[str]) -> Tuple[str, str]: + return typed_array(symbol, "const char * const", [cpp_string(item) for item in items]) + + +def source_ref_array(symbol: str, items: Iterable[str], source_records: Dict[str, str]) -> Tuple[str, str]: + values = [ + "{ " + cpp_string(item) + ", &" + source_records[item] + " }" + for item in items + ] + return typed_array(symbol, "const KernelSourceRef", values) + + +def scalar_array(symbol: str, items: Iterable[dict]) -> Tuple[str, str]: + values = [ + "{ " + cpp_string(item["name"]) + ", " + cpp_string(item["type"]) + " }" + for item in items + ] + return typed_array(symbol, "const KernelScalarDefinition", values) + + +def binding_array(symbol: str, names: List[str], access: List[str]) -> Tuple[str, str]: + if len(names) != len(access): + raise RuntimeError("kernel binding access metadata has the wrong arity") + values = [ + "{ " + cpp_string(name) + ", " + resource_access_value(access_value) + " }" + for name, access_value in zip(names, access) + ] + return typed_array(symbol, "const KernelBindingDefinition", values) + + +def scalar_parameter_names(workload_parameters: List[dict], launch_parameters: List[dict]) -> List[str]: + names: List[str] = [] + for parameter in [*workload_parameters, *launch_parameters]: + name = parameter["name"] + if name not in names: + names.append(name) + return names + + +def depfile_escape(path: pathlib.Path) -> str: + return str(path).replace("\\", "\\\\").replace(" ", "\\ ") + + +def collect_sources(manifest: dict) -> Tuple[List[str], Dict[str, List[str]]]: + source_dependencies: Dict[str, List[str]] = {} + all_sources = set() + + for export in manifest.get("exports", []): + recipe = export.get("compile_recipe", {}) + primary_sources = recipe.get("primary_sources", []) + library_sources = recipe.get("library_sources", []) + if len(primary_sources) != 1: + raise RuntimeError("expected each export compile_recipe to have exactly one primary source") + + primary = primary_sources[0] + dependencies = list(library_sources) + previous = source_dependencies.get(primary) + if previous is not None and previous != dependencies: + raise RuntimeError(f"conflicting dependency list for {primary}") + source_dependencies[primary] = dependencies + all_sources.add(primary) + all_sources.update(dependencies) + + return sorted(all_sources), source_dependencies + + +def manifest_file_digests(manifest: dict) -> Dict[str, str]: + return {file["path"]: file["sha256"] for file in manifest.get("files", [])} + + +def sha256(data: bytes) -> str: + return hashlib.sha256(data).hexdigest() + + +def kernel_catalog_id(family: str, name: str) -> int: + hash_value = 1469598103934665603 + for byte in family.encode("utf-8"): + hash_value ^= byte + hash_value = (hash_value * 1099511628211) & 0xFFFFFFFFFFFFFFFF + hash_value ^= 0 + hash_value = (hash_value * 1099511628211) & 0xFFFFFFFFFFFFFFFF + for byte in name.encode("utf-8"): + hash_value ^= byte + hash_value = (hash_value * 1099511628211) & 0xFFFFFFFFFFFFFFFF + return hash_value + + +def generate_corpus_records(manifest: dict, source_records: Dict[str, str]) -> Tuple[str, str, int]: + digests = manifest_file_digests(manifest) + arrays = [] + records = [] + exports = manifest.get("exports", []) + for index, export in enumerate(exports): + recipe = export["compile_recipe"] + dependencies = list(export["compile_dependencies"]) + library_sources = list(recipe["library_sources"]) + if dependencies != library_sources: + raise RuntimeError("legacy dependency closure disagrees with compile recipe") + workload_parameters = list(export["workload_parameters"]) + launch_parameters = list(export["launch_parameters"]) + + dependencies_array, dependencies_span = string_array(f"kKernelDependencies{index}", dependencies) + scalar_array_text, scalar_span = string_array( + f"kKernelScalarParameters{index}", + scalar_parameter_names(workload_parameters, launch_parameters), + ) + bindings_array, bindings_span = binding_array( + f"kKernelBindings{index}", + list(export["bindings"]), + list(export.get("binding_access", [])), + ) + workload_array, workload_span = scalar_array(f"kKernelWorkloadParameters{index}", workload_parameters) + launch_array, launch_span = scalar_array(f"kKernelLaunchParameters{index}", launch_parameters) + primary_array, primary_span = source_ref_array(f"kKernelPrimarySources{index}", recipe["primary_sources"], source_records) + library_array, library_span = source_ref_array(f"kKernelLibrarySources{index}", library_sources, source_records) + arrays.extend( + item for item in [ + dependencies_array, + scalar_array_text, + bindings_array, + workload_array, + launch_array, + primary_array, + library_array, + ] if item + ) + records.append( + KERNEL_RECORD_TEMPLATE.format( + family=cpp_string(export.get("family", DEFAULT_KERNEL_FAMILY)), + name=cpp_string(export["name"]), + symbol=cpp_string(export["symbol"]), + target_selector=cpp_string(export.get("target_selector", "")), + source=cpp_string(export["source"]), + dependencies=dependencies_span, + source_digest=cpp_string(digests[export["source"]]), + scalar_parameters=scalar_span, + bindings=bindings_span, + workload_parameters=workload_span, + launch_parameters=launch_span, + compile_mode=cpp_string(recipe["mode"]), + link_module=cpp_string(recipe.get("link_module", "")), + primary_sources=primary_span, + library_sources=library_span, + ) + ) + return "\n".join(arrays), "\n".join(records), len(exports) + + +def generate_catalog_verifier(manifest: dict) -> str: + kernel_entries = sorted(set( + (export.get("family", DEFAULT_KERNEL_FAMILY), export["name"]) + for export in manifest.get("exports", []) + )) + ids: Dict[int, Tuple[str, str]] = {} + for family, kernel_name in kernel_entries: + catalog_id = kernel_catalog_id(family, kernel_name) + if catalog_id in ids: + previous_family, previous_name = ids[catalog_id] + raise RuntimeError( + "kernel catalog id collision: " + f"{previous_family}/{previous_name} and {family}/{kernel_name}") + ids[catalog_id] = (family, kernel_name) + return CATALOG_DATA_TEMPLATE.format( + kernel_entries="\n".join( + " { " + cpp_string(family) + ", " + cpp_string(kernel_name) + " }," + for family, kernel_name in kernel_entries + ), + ) + + +def generate_includes(args: argparse.Namespace, manifest: dict) -> Tuple[str, str, str, List[pathlib.Path], int]: + if args.source_format == "binary" and (args.loom_link is None or args.loom_format is None): + raise RuntimeError("binary source format requires --loom-link and --loom-format") + + corpus_dir = args.corpus_dir + sources, source_dependencies = collect_sources(manifest) + digests = manifest_file_digests(manifest) + source_bytes: Dict[str, bytes] = {} + source_format = "KERNEL_SOURCE_FORMAT_BINARY" if args.source_format == "binary" else "KERNEL_SOURCE_FORMAT_TEXT" + input_files: List[pathlib.Path] = [] + + for export in manifest.get("exports", []): + recipe = export.get("compile_recipe", {}) + for primary in recipe.get("primary_sources", []): + if primary not in sources: + raise RuntimeError(f"export primary source is not embedded: {primary}") + + for source in sources: + if source not in digests: + raise RuntimeError(f"embedded source is missing from manifest file table: {source}") + path = corpus_dir / source + input_files.append(path) + data = read_bytes(path) + digest = sha256(data) + if digest != digests[source]: + raise RuntimeError(f"manifest digest mismatch for {source}: got {digest}, expected {digests[source]}") + if args.source_format == "binary": + data = convert_source_to_bytecode(path, args.loom_link, args.loom_format) + source_bytes[source] = data + + source_symbols: Dict[str, str] = {} + source_arrays = [] + for index, source in enumerate(sources): + symbol = sanitize_symbol(source, index) + source_symbols[source] = symbol + data = source_bytes[source] + source_arrays.append( + SOURCE_ARRAY_TEMPLATE.format( + symbol=symbol, + bytes=format_byte_array(data), + size=len(data), + ) + ) + + dependency_tables = [] + source_record_definitions = [] + source_record_symbols = {} + lookup_entries = [] + for index, source in enumerate(sources): + dependencies = source_dependencies.get(source, []) + record = f"kernel_source_record_{index}" + source_record_symbols[source] = record + if dependencies: + dependency_table = f"kernel_source_dependencies_{index}" + entries = [] + for dependency in dependencies: + symbol = source_symbols[dependency] + entries.append( + f" {{ reinterpret_cast({symbol}), {symbol}Size, {source_format} }}," + ) + dependency_tables.append( + DEPENDENCY_TABLE_TEMPLATE.format( + dependency_table=dependency_table, + dependencies="\n".join(entries), + ) + ) + else: + dependency_table = "nullptr" + + source_record_definitions.append( + SOURCE_RECORD_TEMPLATE.format( + record=record, + source_symbol=source_symbols[source], + source_format=source_format, + dependency_table=dependency_table, + dependency_count=len(dependencies), + ) + ) + lookup_entries.append(f" {{ {cpp_string(source)}, &{record} }},") + + kernel_arrays, kernel_records, kernel_count = generate_corpus_records(manifest, source_record_symbols) + return ( + SOURCE_DATA_TEMPLATE.format( + source_arrays="\n".join(source_arrays), + dependency_tables="\n".join(dependency_tables), + source_records="\n".join(source_record_definitions), + lookup_entries="\n".join(lookup_entries), + ), + CORPUS_DATA_TEMPLATE.format( + kernel_arrays=kernel_arrays, + kernel_records=kernel_records, + upstream_revision=cpp_string(manifest["upstream_revision"]), + corpus_digest=cpp_string(manifest["corpus_sha256"]), + recipe_digest=cpp_string(manifest["build_bazel_sha256"]), + plan_case_count=len(manifest["plan_cases"]), + kernel_count=kernel_count, + ), + generate_catalog_verifier(manifest), + input_files, + sum(len(data) for data in source_bytes.values()), + ) + + +def write_depfile(path: pathlib.Path, outputs: List[pathlib.Path], inputs: List[pathlib.Path]) -> None: + targets = " ".join(depfile_escape(output_path) for output_path in outputs) + entries = [depfile_escape(input_path) for input_path in inputs] + path.write_text(f"{targets}: {' '.join(entries)}\n") + + +def main() -> int: + args = parse_args() + try: + manifest = json.loads(read_text(args.manifest)) + source_include, corpus_include, catalog_include, input_files, byte_count = generate_includes(args, manifest) + args.source_output.parent.mkdir(parents=True, exist_ok=True) + args.corpus_output.parent.mkdir(parents=True, exist_ok=True) + args.catalog_output.parent.mkdir(parents=True, exist_ok=True) + args.source_output.write_text(source_include) + args.corpus_output.write_text(corpus_include) + args.catalog_output.write_text(catalog_include) + if args.depfile is not None: + args.depfile.parent.mkdir(parents=True, exist_ok=True) + write_depfile(args.depfile, [args.source_output, args.corpus_output, args.catalog_output], + [args.manifest, *input_files]) + except Exception as exc: + print(f"generate_kernel_corpus.py: {exc}", file=sys.stderr) + return 1 + + print( + f"embedded kernel corpus: source_files={len(input_files)} source_bytes={byte_count} " + f"source_format={args.source_format}", + file=sys.stderr, + ) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/ggml/src/ggml-hrx/tools/sync_qwen_kernel_corpus.py b/ggml/src/ggml-hrx/tools/sync_qwen_kernel_corpus.py new file mode 100755 index 000000000000..5d4d9705c47d --- /dev/null +++ b/ggml/src/ggml-hrx/tools/sync_qwen_kernel_corpus.py @@ -0,0 +1,596 @@ +#!/usr/bin/env python3 +"""Mirrors the bounded Qwen MoE Loom corpus with reproducible provenance.""" + +from __future__ import annotations + +import argparse +import ast +import hashlib +import json +import pathlib +import re +import shutil +import subprocess +import sys +import tempfile + + +SOURCE_SUBDIR = pathlib.Path("experimental/qwen_moe/kernels") +QWEN_ENDPOINT_SOURCE_SUBDIR = pathlib.Path("experimental/qwen/kernels") +CORPUS_FILES = ( + "ggml/linear_q6k_f32.loom", + "ggml/linear_q6k_q8_1_x4.loom", + "ggml/quantize_q8_1_x4.loom", + "qwen3_moe/attention_postprocess_f32_f16.loom", + "qwen3_moe/attention_prepare_quantized.loom", + "qwen3_moe/attention_qkv_postprocess_fused.loom", + "qwen3_moe/attention_qkv_quantized.loom", + "qwen3_moe/attention_qkv_same_format_prefill.loom", + "qwen3_moe/batched_decode_expert_dispatch.loom", + "qwen3_moe/batched_decode_gate_up_q4k.loom", + "qwen3_moe/dense_linear_quantized_f16_wmma.loom", + "qwen3_moe/expert_table_partition_fused.loom", + "qwen3_moe/flash_attention_decode_f32_f16_wmma.loom", + "qwen3_moe/flash_attention_decode_q128_f32_f16_wmma.loom", + "qwen3_moe/flash_attention_decode_split_f32_f16_wmma.loom", + "qwen3_moe/flash_attention_decode_split_next_q8_test.loom", + "qwen3_moe/flash_attention_f32_f16_wmma.loom", + "qwen3_moe/model_config.loom", + "qwen3_moe/routed_down_q4k.loom", + "qwen3_moe/routed_down_q6k.loom", + "qwen3_moe/routed_down_next_q8.loom", + "qwen3_moe/routed_down_quantized_f16_wmma.loom", + "qwen3_moe/routed_down_weighted_reduce_next_rmsnorm_f32.loom", + "qwen3_moe/routed_down_weighted_reduce_next_rmsnorm_q8_1_x4.loom", + "qwen3_moe/routed_gate_up_swiglu_q4k.loom", + "qwen3_moe/routed_linear_q4k_f16_wmma.loom", + "qwen3_moe/router_projection_f32.loom", + "qwen3_moe/router_projection_top8_fused_f32.loom", + "qwen3_moe/router_top8_f32.loom", +) +QWEN_ENDPOINT_FILES = ( + ("token_embedding_q4k.loom", "qwen_owned/token_embedding_q4k.loom"), + ("attention_metadata.loom", "qwen_owned/attention_metadata.loom"), +) + +# These integration kernels are deliberately owned by the llama.cpp HRX +# backend. They are not attributed to the pinned qwen_moe corpus or its BUILD +# recipes. +OWNED_KERNEL_DIR = pathlib.Path(__file__).resolve().parent.parent / "kernel-corpus" / "kernels" +OWNED_FILES = ( + "qwen_owned/token_embedding_bringup_workaround.loom", + "qwen_owned/attention_state_initialize.loom", + "qwen_owned/attention_metadata_bringup_workaround.loom", + "hrx_owned/gather_add_f32.loom", + "hrx_owned/add_f32.loom", +) + +KERNEL_RE = re.compile( + r"kernel\.def(?P(?:\s+(?:target\([^)]*\)|export\(\"[^\"]+\"\)))*)" + r"\s+@(?P[A-Za-z0-9_]+)" + r"\((?P.*?)\)\s*\{.*?\}\s*launch\((?P.*?)\)" + r"(?:\s+where\s+\[[^\]]*\])?\s*\{", + re.DOTALL, +) +ARG_RE = re.compile(r"%(?P[A-Za-z0-9_]+)\s*:\s*(?P[A-Za-z0-9<>?]+)") +TARGET_MODIFIER_RE = re.compile(r"target\(@(?P[A-Za-z0-9_]+)\)") +EXPORT_MODIFIER_RE = re.compile(r"export\(\"(?P[^\"]+)\"\)") +AMDGPU_TARGET_RE = re.compile( + r"amdgpu\.target<(?P[A-Za-z0-9_.-]+)>\s+@(?P[A-Za-z0-9_]+)" +) + + +def binding_access(symbol: str, name: str) -> str: + """Authoritative launch ABI access contract; no name inference at runtime.""" + if symbol == "qwen_attention_context_base_capture": + return "read" if name == "positions" else "write" + if symbol == "qwen_attention_decode_state_initialize": + return "read" if name == "positions" else "write" + if symbol == "qwen_attention_metadata_bringup_workaround" and name != "control": + return "read_write" + if symbol == "qwen_attention_metadata" and name != "control": + return "read_write" + if symbol == "qwen_decode_attention_metadata" and name != "control": + return "read_write" + if symbol == "qwen3_moe_router_top8_f32" and name in ("route_ids", "route_weights"): + return "write" + if symbol == "qwen3_moe_router_projection_top8_fused_decode_f32" and name in ( + "logits", "completion_counter", "route_ids", "route_weights"): + return "read_write" + if symbol == "qwen3_moe_attention_qkv_postprocess_fused_decode" and name in ( + "query_output_raw", "key_output_raw", "value_output_raw", + "query_output", "key_cache", "value_cache", "completion_counters"): + return "read_write" + if symbol == "qwen3_moe_build_expert_table" and name == "expert_table": + return "write" + if symbol == "qwen3_moe_build_expert_partition_table" and name == "partition_table": + return "write" + if symbol == "qwen3_moe_build_expert_table_partition_prefill_512" and name in ("expert_table", "partition_table"): + return "write" + if symbol == "qwen3_moe_routed_down_weighted_reduce_next_rmsnorm_f32": + if name == "hidden_state": + return "read_write" + if name == "next_projection_input": + return "write" + if symbol == "ggml_q8_1_x4_inspect_one_group" and name != "packed": + return "write" + if name in ("output", "query_output", "key_output", "value_output", "normalized_output", "q8_output", + "next_q8_output", "key_cache", "value_cache", "partial_max", "partial_sum", "partial_output", + "completion_counter", "completion_counters"): + return "read_write" + return "read" + + +def starlark_calls(text: str, function: str) -> list[str]: + """Extracts the literal-only calls used by the pinned kernel BUILD file.""" + result: list[str] = [] + marker = function + "(" + cursor = 0 + while (start := text.find(marker, cursor)) != -1: + index = start + len(marker) + depth = 1 + quote: str | None = None + escaped = False + while index < len(text) and depth: + character = text[index] + if quote: + if escaped: + escaped = False + elif character == "\\": + escaped = True + elif character == quote: + quote = None + elif character in "\"'": + quote = character + elif character == "(": + depth += 1 + elif character == ")": + depth -= 1 + index += 1 + if depth: + raise RuntimeError(f"unterminated {function} call in BUILD.bazel") + result.append(text[start + len(marker) : index - 1]) + cursor = index + return result + + +def literal_assignment(call: str, name: str, default: object = None) -> object: + match = re.search(rf"(?:^|\n)\s*{re.escape(name)}\s*=\s*(\[[\s\S]*?\]|\"[^\"]*\")\s*,", call) + return default if match is None else ast.literal_eval(match.group(1)) + + +def parse_build_recipes(text: str) -> tuple[list[dict[str, object]], list[dict[str, object]]]: + modules: list[dict[str, object]] = [] + for call in starlark_calls(text, "loom_link_module"): + modules.append({ + "name": literal_assignment(call, "name"), + "srcs": literal_assignment(call, "srcs", []), + "libraries": literal_assignment(call, "libraries", []), + }) + module_names = {str(module["name"]) for module in modules} + cases: list[dict[str, object]] = [] + for call in starlark_calls(text, "iree_executable_test"): + name = literal_assignment(call, "name") + if not isinstance(name, str) or not name.endswith("_plan_test"): + continue + args = literal_assignment(call, "args", []) + data = literal_assignment(call, "data", []) + linked = [item[1:] for item in data if isinstance(item, str) and item.startswith(":") and item[1:] in module_names] + direct_sources = [item for item in data if isinstance(item, str) and item.endswith(".loom")] + if len(linked) + len(direct_sources) != 1: + raise RuntimeError(f"plan test {name} does not name exactly one linked module or direct Loom source") + case = { + "name": name, + "args": args, + } + if linked: + case["link_module"] = linked[0] + else: + case["source"] = direct_sources[0] + cases.append(case) + if not modules or not cases: + raise RuntimeError("BUILD.bazel contains no pinned Loom link/plan recipes") + return modules, cases + + +def git(repo: pathlib.Path, *args: str) -> str: + return subprocess.check_output(["git", "-C", str(repo), *args], text=True).strip() + + +def optional_git(repo: pathlib.Path, *args: str) -> str | None: + result = subprocess.run( + ["git", "-C", str(repo), *args], text=True, + stdout=subprocess.PIPE, stderr=subprocess.DEVNULL, check=False) + value = result.stdout.strip() + return value if result.returncode == 0 and value else None + + +def normalize_repository_url(url: str) -> str: + """Normalizes GitHub transport spelling without changing repository identity.""" + match = re.fullmatch(r"(?:ssh://)?git@github\.com[:/](?P.+)", url) + if match: + return f"https://github.com/{match.group('path')}" + return url + + +def upstream_repository(repo: pathlib.Path) -> str: + """Returns provenance without imposing a local Git remote name.""" + remotes: list[str] = [] + branch = optional_git(repo, "symbolic-ref", "--quiet", "--short", "HEAD") + if branch: + branch_remote = optional_git(repo, "config", "--get", f"branch.{branch}.remote") + if branch_remote and branch_remote != ".": + remotes.append(branch_remote) + remotes.append("origin") + remotes.extend(git(repo, "remote").splitlines()) + for remote in dict.fromkeys(remotes): + url = optional_git(repo, "config", "--get", f"remote.{remote}.url") + if url: + return normalize_repository_url(url) + raise RuntimeError("HRX source tree has no repository remote for provenance") + + +def sha256(data: bytes) -> str: + return hashlib.sha256(data).hexdigest() + + +def parse_exports(text: str, source: str) -> list[dict[str, object]]: + exports: list[dict[str, object]] = [] + for match in KERNEL_RE.finditer(text): + modifiers = match.group("modifiers") + target_match = TARGET_MODIFIER_RE.search(modifiers) + export_match = EXPORT_MODIFIER_RE.search(modifiers) + workload = [item.groupdict() for item in ARG_RE.finditer(match.group("workload"))] + launch = [item.groupdict() for item in ARG_RE.finditer(match.group("launch"))] + bindings = [item["name"] for item in launch if item["type"] == "buffer"] + exports.append( + { + "name": export_match.group("name") if export_match else match.group("symbol"), + "symbol": match.group("symbol"), + "target_symbol": target_match.group("symbol") if target_match else "", + "source": source, + "workload_parameters": workload, + "launch_parameters": [item for item in launch if item["type"] != "buffer"], + "bindings": bindings, + "binding_access": [binding_access(match.group("symbol"), name) for name in bindings], + } + ) + return exports + + +def parse_amdgpu_targets(text: str) -> dict[str, str]: + result: dict[str, str] = {} + for match in AMDGPU_TARGET_RE.finditer(text): + symbol = match.group("symbol") + selector = match.group("selector") + if symbol in result and result[symbol] != selector: + raise RuntimeError( + f"AMDGPU target @{symbol} is declared as both {result[symbol]} and {selector}") + result[symbol] = selector + return result + + +def merge_amdgpu_targets(target_selectors: dict[str, str], additions: dict[str, str]) -> None: + for symbol, selector in additions.items(): + if symbol in target_selectors and target_selectors[symbol] != selector: + raise RuntimeError( + f"AMDGPU target @{symbol} is declared as both {target_selectors[symbol]} and {selector}") + target_selectors[symbol] = selector + + +def resolve_export_variants(exports: list[dict[str, object]], target_selectors: dict[str, str]) -> None: + groups: dict[str, list[dict[str, object]]] = {} + for item in exports: + groups.setdefault(str(item["name"]), []).append(item) + + for name, variants in groups.items(): + selectors: set[str] = set() + for item in variants: + target_symbol = str(item.pop("target_symbol")) + selector = "" + if target_symbol: + if target_symbol not in target_selectors: + raise RuntimeError(f"kernel export {name} references unknown target @{target_symbol}") + selector = target_selectors[target_symbol] + if selector.endswith("-generic"): + selector = "" + if selector in selectors: + label = selector or "default" + raise RuntimeError(f"kernel export {name} repeats target variant {label}") + selectors.add(selector) + item["target_selector"] = selector + + +def construct(source_root: pathlib.Path, destination: pathlib.Path, expected_revision: str | None) -> None: + revision = git(source_root, "rev-parse", "HEAD") + if expected_revision and revision != expected_revision: + raise RuntimeError(f"HRX revision {revision} does not match expected {expected_revision}") + if git(source_root, "status", "--porcelain"): + raise RuntimeError("refusing to mirror a dirty HRX source tree") + + source_directory = source_root / SOURCE_SUBDIR + build_data = (source_directory / "BUILD.bazel").read_bytes() + all_link_modules, all_plan_cases = parse_build_recipes(build_data.decode("utf-8")) + + modules_by_name = {str(item["name"]): item for item in all_link_modules} + direct_plan_sources = { + str(item["source"]) + for item in all_plan_cases + if "source" in item + } + + def module_files(name: str) -> list[str]: + module = modules_by_name[name] + result = list(module["srcs"]) + for library in module["libraries"]: + if str(library).startswith(":"): + result.extend(module_files(str(library)[1:])) + else: + result.append(str(library)) + return list(dict.fromkeys(result)) + + def compile_recipe(source: str) -> dict[str, object]: + direct = [str(item["name"]) for item in all_link_modules if source in item["srcs"]] + indirect = [str(item["name"]) for item in all_link_modules if source in item["libraries"]] + if not direct and (source in direct_plan_sources or not indirect): + return {"mode": "direct", "primary_sources": [source], "library_sources": []} + module_name = (direct or indirect)[0] + module = modules_by_name[module_name] + files = module_files(module_name) + return { + "mode": "archive", + "link_module": module_name, + "primary_sources": list(module["srcs"]), + "library_sources": [item for item in files if item not in module["srcs"]], + } + file_rows: list[dict[str, object]] = [] + exports: list[dict[str, object]] = [] + target_selectors: dict[str, str] = {} + upstream_aggregate = hashlib.sha256() + for relative_text in CORPUS_FILES: + relative = pathlib.Path(relative_text) + source = source_directory / relative + if not source.is_file(): + raise RuntimeError(f"missing required corpus source: {source}") + data = source.read_bytes() + digest = sha256(data) + upstream_aggregate.update(relative_text.encode()) + upstream_aggregate.update(b"\0") + upstream_aggregate.update(bytes.fromhex(digest)) + target = destination / relative + target.parent.mkdir(parents=True, exist_ok=True) + target.write_bytes(data) + file_rows.append({"path": relative_text, "sha256": digest, "size": len(data)}) + source_text = data.decode("utf-8") + exports.extend(parse_exports(source_text, relative_text)) + merge_amdgpu_targets(target_selectors, parse_amdgpu_targets(source_text)) + + endpoint_source_directory = source_root / QWEN_ENDPOINT_SOURCE_SUBDIR + for source_text_name, local_text_name in QWEN_ENDPOINT_FILES: + source = endpoint_source_directory / source_text_name + if not source.is_file(): + raise RuntimeError(f"missing required Qwen endpoint source: {source}") + data = source.read_bytes() + digest = sha256(data) + provenance_path = f"{QWEN_ENDPOINT_SOURCE_SUBDIR.as_posix()}/{source_text_name}" + upstream_aggregate.update(provenance_path.encode()) + upstream_aggregate.update(b"\0") + upstream_aggregate.update(bytes.fromhex(digest)) + relative_text = f"../{local_text_name}" + target = destination.parent / local_text_name + target.parent.mkdir(parents=True, exist_ok=True) + target.write_bytes(data) + file_rows.append({ + "path": relative_text, + "sha256": digest, + "size": len(data), + "upstream_path": provenance_path, + }) + source_text = data.decode("utf-8") + exports.extend(parse_exports(source_text, relative_text)) + merge_amdgpu_targets(target_selectors, parse_amdgpu_targets(source_text)) + + owned_aggregate = hashlib.sha256() + for filename in OWNED_FILES: + source = OWNED_KERNEL_DIR / filename + if not source.is_file(): + raise RuntimeError(f"missing required backend-owned kernel source: {source}") + data = source.read_bytes() + digest = sha256(data) + relative_text = f"../{filename}" + owned_aggregate.update(filename.encode()) + owned_aggregate.update(b"\0") + owned_aggregate.update(bytes.fromhex(digest)) + file_rows.append({"path": relative_text, "sha256": digest, "size": len(data), "owner": "ggml-hrx"}) + source_text = data.decode("utf-8") + exports.extend(parse_exports(source_text, relative_text)) + merge_amdgpu_targets(target_selectors, parse_amdgpu_targets(source_text)) + + resolve_export_variants(exports, target_selectors) + + for item in exports: + recipe = compile_recipe(str(item["source"])) + item["compile_recipe"] = recipe + item["compile_dependencies"] = list(recipe["library_sources"]) + + required_modules: set[str] = set() + + def require_module(name: str) -> None: + if name in required_modules: + return + if name not in modules_by_name: + raise RuntimeError(f"selected kernel recipe references unknown link module {name}") + required_modules.add(name) + for library in modules_by_name[name]["libraries"]: + if str(library).startswith(":"): + require_module(str(library)[1:]) + + required_files = set(CORPUS_FILES) + for item in exports: + if str(item["source"]).startswith("../"): + continue + recipe = item["compile_recipe"] + link_module = str(recipe.get("link_module", "")) + if link_module: + require_module(link_module) + required_files.update(str(path) for path in recipe["primary_sources"]) + required_files.update(str(path) for path in recipe["library_sources"]) + + mirrored_files = set(CORPUS_FILES) + for relative_text in sorted(required_files - mirrored_files): + relative = pathlib.Path(relative_text) + if relative.is_absolute() or ".." in relative.parts: + raise RuntimeError(f"kernel recipe escapes the source corpus: {relative_text}") + source = source_directory / relative + if not source.is_file(): + raise RuntimeError(f"missing required kernel dependency: {source}") + data = source.read_bytes() + digest = sha256(data) + upstream_aggregate.update(relative_text.encode()) + upstream_aggregate.update(b"\0") + upstream_aggregate.update(bytes.fromhex(digest)) + target = destination / relative + target.parent.mkdir(parents=True, exist_ok=True) + target.write_bytes(data) + file_rows.append({"path": relative_text, "sha256": digest, "size": len(data)}) + + link_modules = [ + module for module in all_link_modules + if str(module["name"]) in required_modules + ] + plan_cases = [ + case for case in all_plan_cases + if (str(case.get("link_module", "")) in required_modules or + str(case.get("source", "")) in required_files) + ] + plan_cases.extend([ + { + "name": "owned_token_embedding_decode_plan_test", + "args": ["$(location ../qwen_owned/token_embedding_q4k.loom)", + "--benchmark=@qwen_token_embedding_q4k_decode", "--dry-run", + "--output-format=jsonl", "--sample-compilation=per_sample"], + "source": "../qwen_owned/token_embedding_q4k.loom", + }, + { + "name": "owned_token_embedding_prefill_plan_test", + "args": ["$(location ../qwen_owned/token_embedding_q4k.loom)", + "--benchmark=@qwen_token_embedding_q4k_prefill_512", "--dry-run", + "--output-format=jsonl", "--sample-compilation=per_sample"], + "source": "../qwen_owned/token_embedding_q4k.loom", + }, + { + "name": "owned_attention_context_base_capture_plan_test", + "args": ["$(location ../qwen_owned/attention_state_initialize.loom)", + "--benchmark=@qwen_attention_context_base_capture_benchmark", "--dry-run", + "--output-format=jsonl", "--sample-compilation=per_sample"], + "source": "../qwen_owned/attention_state_initialize.loom", + "owner": "ggml-hrx", + }, + { + "name": "owned_attention_decode_state_initialize_plan_test", + "args": ["$(location ../qwen_owned/attention_state_initialize.loom)", + "--benchmark=@qwen_attention_decode_state_initialize_benchmark", "--dry-run", + "--output-format=jsonl", "--sample-compilation=per_sample"], + "source": "../qwen_owned/attention_state_initialize.loom", + "owner": "ggml-hrx", + }, + { + "name": "owned_attention_metadata_prefill_plan_test", + "args": ["$(location ../qwen_owned/attention_metadata.loom)", + "--benchmark=@qwen_attention_metadata_prefill_512", "--dry-run", + "--output-format=jsonl", "--sample-compilation=per_sample"], + "source": "../qwen_owned/attention_metadata.loom", + }, + { + "name": "owned_gather_add_plan_test", + "args": ["$(location ../hrx_owned/gather_add_f32.loom)", + "--benchmark=@ggml_gather_add_noncontiguous", "--dry-run", + "--output-format=jsonl", "--sample-compilation=per_sample"], + "source": "../hrx_owned/gather_add_f32.loom", + "owner": "ggml-hrx", + }, + ]) + + upstream_digest = upstream_aggregate.hexdigest() + owned_digest = owned_aggregate.hexdigest() + combined_aggregate = hashlib.sha256() + combined_aggregate.update(bytes.fromhex(upstream_digest)) + combined_aggregate.update(bytes.fromhex(owned_digest)) + + manifest = { + "schema": "ggml-hrx-qwen-kernel-corpus-v2", + "upstream_repository": upstream_repository(source_root), + "upstream_revision": revision, + "source_subdirectory": SOURCE_SUBDIR.as_posix(), + "qwen_endpoint_source_subdirectory": QWEN_ENDPOINT_SOURCE_SUBDIR.as_posix(), + "corpus_sha256": combined_aggregate.hexdigest(), + "upstream_corpus_sha256": upstream_digest, + "owned_corpus_sha256": owned_digest, + "build_bazel_sha256": sha256(build_data), + "files": file_rows, + "exports": sorted(exports, key=lambda item: (str(item["symbol"]), str(item["source"]))), + "link_modules": link_modules, + "plan_cases": plan_cases, + } + (destination / "manifest.json").write_text(json.dumps(manifest, indent=2) + "\n", encoding="utf-8") + + +def trees_equal(lhs: pathlib.Path, rhs: pathlib.Path) -> bool: + lhs_files = sorted(path.relative_to(lhs) for path in lhs.rglob("*") if path.is_file()) + rhs_files = sorted(path.relative_to(rhs) for path in rhs.rglob("*") if path.is_file()) + return lhs_files == rhs_files and all((lhs / path).read_bytes() == (rhs / path).read_bytes() for path in lhs_files) + + +def endpoint_files_equal(generated_qwen_moe: pathlib.Path, destination_qwen_moe: pathlib.Path) -> bool: + generated_kernel_root = generated_qwen_moe.parent + destination_kernel_root = destination_qwen_moe.parent + for _, local_text_name in QWEN_ENDPOINT_FILES: + relative = pathlib.Path(local_text_name) + generated = generated_kernel_root / relative + destination = destination_kernel_root / relative + if not destination.is_file() or generated.read_bytes() != destination.read_bytes(): + return False + return True + + +def copy_endpoint_files(generated_qwen_moe: pathlib.Path, destination_qwen_moe: pathlib.Path) -> None: + generated_kernel_root = generated_qwen_moe.parent + destination_kernel_root = destination_qwen_moe.parent + for _, local_text_name in QWEN_ENDPOINT_FILES: + relative = pathlib.Path(local_text_name) + source = generated_kernel_root / relative + target = destination_kernel_root / relative + target.parent.mkdir(parents=True, exist_ok=True) + shutil.copy2(source, target) + + +def main() -> int: + parser = argparse.ArgumentParser() + parser.add_argument("--hrx-source", type=pathlib.Path, required=True) + parser.add_argument("--destination", type=pathlib.Path, required=True) + parser.add_argument("--expect-revision") + parser.add_argument("--check", action="store_true") + args = parser.parse_args() + + with tempfile.TemporaryDirectory(prefix="hrx-qwen-corpus-") as temporary: + generated = pathlib.Path(temporary) / args.destination.name + construct(args.hrx_source.resolve(), generated, args.expect_revision) + if args.check: + if (not args.destination.is_dir() or not trees_equal(generated, args.destination) or + not endpoint_files_equal(generated, args.destination)): + print("mirrored Qwen kernel corpus is stale", file=sys.stderr) + return 1 + return 0 + args.destination.mkdir(parents=True, exist_ok=True) + for child in args.destination.iterdir(): + if child.is_dir(): + shutil.rmtree(child) + else: + child.unlink() + shutil.copytree(generated, args.destination, dirs_exist_ok=True) + copy_endpoint_files(generated, args.destination) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/ggml/src/ggml-hrx/tools/tool-utils.h b/ggml/src/ggml-hrx/tools/tool-utils.h new file mode 100644 index 000000000000..c01e089ceab4 --- /dev/null +++ b/ggml/src/ggml-hrx/tools/tool-utils.h @@ -0,0 +1,61 @@ +#pragma once + +#include "../hrx-interop-utils.h" + +#include +#include +#include +#include +#include +#include + +namespace ggml::hrx::tool { + +inline std::string read_file(const std::filesystem::path & path) { + std::ifstream input(path, std::ios::binary); + if (!input) { + return {}; + } + return { std::istreambuf_iterator(input), std::istreambuf_iterator() }; +} + +inline void write_file(const std::filesystem::path & path, const void * data, size_t size) { + std::ofstream output(path, std::ios::binary | std::ios::trunc); + if (!output) { + throw std::runtime_error("cannot create " + path.string()); + } + output.write(static_cast(data), static_cast(size)); + if (!output) { + throw std::runtime_error("cannot write " + path.string()); + } +} + +inline void write_file(const std::filesystem::path & path, const std::string & contents) { + std::ofstream output(path, std::ios::binary | std::ios::trunc); + if (!output) { + throw std::runtime_error("cannot create " + path.string()); + } + output << contents; + if (contents.empty() || contents.back() != '\n') { + output << '\n'; + } + if (!output) { + throw std::runtime_error("cannot write " + path.string()); + } +} + +inline void check_status(hrx_status_t status, const std::string & operation) { + if (ErrorResult error = take_status(status)) { + throw std::runtime_error(operation + ": " + *error); + } +} + +inline bool report_status(hrx_status_t status, const std::string & operation) { + if (ErrorResult error = take_status(status)) { + std::cerr << operation << ": " << *error << '\n'; + return false; + } + return true; +} + +} // namespace ggml::hrx::tool diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 881e55c75a1d..b53be5e8a127 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -336,3 +336,25 @@ if (TARGET gguf-model-data) target_link_libraries(test-export-graph-ops PRIVATE gguf-model-data) target_compile_definitions(test-export-graph-ops PRIVATE LLAMA_HF_FETCH) endif() + +if (TARGET ggml-hrx) + add_executable(test-hrx-buffer test-hrx-buffer.cpp) + target_link_libraries(test-hrx-buffer PRIVATE ggml-hrx ggml hrx::hrx) + target_include_directories(test-hrx-buffer PRIVATE ../ggml/include ../ggml/src ../ggml/src/ggml-hrx) + add_test(NAME test-hrx-buffer COMMAND test-hrx-buffer) + + add_executable(hrx-backend-test hrx-backend-test.cpp) + target_link_libraries(hrx-backend-test PRIVATE ggml-hrx ggml hrx::hrx loomc::loomc) + target_include_directories(hrx-backend-test PRIVATE ../ggml/include ../ggml/src ../ggml/src/ggml-hrx) + add_test(NAME hrx-backend-test COMMAND hrx-backend-test) + + add_executable(test-hrx-loom-jit test-hrx-loom-jit.cpp) + target_link_libraries(test-hrx-loom-jit PRIVATE ggml-hrx ggml-hrx-kernel-corpus hrx::hrx loomc::loomc) + target_include_directories(test-hrx-loom-jit PRIVATE ../ggml/include ../ggml/src ../ggml/src/ggml-hrx) + add_test(NAME test-hrx-loom-jit COMMAND test-hrx-loom-jit) + + add_executable(test-hrx-ops test-hrx-ops.cpp) + target_link_libraries(test-hrx-ops PRIVATE ggml-hrx ggml hrx::hrx loomc::loomc) + target_include_directories(test-hrx-ops PRIVATE ../ggml/include ../ggml/src ../ggml/src/ggml-hrx) + add_test(NAME test-hrx-ops COMMAND test-hrx-ops) +endif() diff --git a/tests/hrx-backend-test.cpp b/tests/hrx-backend-test.cpp new file mode 100644 index 000000000000..1320f1255f66 --- /dev/null +++ b/tests/hrx-backend-test.cpp @@ -0,0 +1,4857 @@ +#include "backend-buffer-binding.h" +#include "backend-context.h" +#include "dispatch/command-program-bindings.h" +#include "dispatch/command-program-diagnostics.h" +#include "dispatch/command-program-resolver.h" +#include "dispatch/command-program.h" +#include "dispatch/dispatch-scheduler.h" +#include "dispatch_registration/dispatch-registry.h" +#include "ggml-alloc.h" +#include "ggml-backend.h" +#include "ggml-hrx.h" +#include "ggml-impl.h" +#include "ggml.h" +#include "graph/graph-diagnostics.h" +#include "graph/graph-matcher.h" +#include "graph/graph-traversal.h" +#include "graph/graph.h" +#include "hrx-interop-utils.h" +#include "kernel-corpus/kernel-corpus.h" +#include "runtime/command-program-executor.h" +#include "runtime/graph-executor.h" +#include "runtime/graph-program-cache.h" +#include "runtime/graph-replay.h" +#include "runtime/loom-kernel-jit.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#define REQUIRE(condition) \ + do { \ + if (!(condition)) { \ + std::fprintf(stderr, "%s:%d: requirement failed: %s\n", __FILE__, __LINE__, #condition); \ + std::abort(); \ + } \ + } while (false) + +static hrx_buffer_t dummy_hrx_buffer(uintptr_t value) { + return reinterpret_cast(value); +} + +static hrx_stream_t dummy_hrx_stream(uintptr_t value) { + return reinterpret_cast(value); +} + +static hrx_graph_exec_t dummy_hrx_graph_exec(uintptr_t value) { + return reinterpret_cast(value); +} + +static bool contains_value_id(const std::vector & ids, ggml::hrx::ValueId id) { + for (const ggml::hrx::ValueId candidate : ids) { + if (candidate == id) { + return true; + } + } + return false; +} + +static bool command_program_verifies(const ggml::hrx::CommandProgram & program) { + return ggml::hrx::verify_command_program(program, ggml::hrx::get_qwen_kernel_corpus(), "gfx1151").valid(); +} + +static bool async_jit_expected_from_environment() { + const char * value = std::getenv("GGML_HRX_ASYNC_JIT"); + return value == nullptr || + (std::strcmp(value, "0") != 0 && std::strcmp(value, "false") != 0 && std::strcmp(value, "FALSE") != 0 && + std::strcmp(value, "off") != 0 && std::strcmp(value, "OFF") != 0); +} + +static ggml::hrx::CommandProgram copy_command_program_shape(const ggml::hrx::CommandProgram & program) { + ggml::hrx::CommandProgram copy; + copy.initialization_commands = program.initialization_commands; + copy.commands = program.commands; + copy.transients = program.transients; + copy.completion_counters = program.completion_counters; + copy.constant_initializations = program.constant_initializations; + return copy; +} + +static bool status_contains(const ggml::hrx::Status & status, const char * text) { + for (const std::string & message : status.errors()) { + if (message.find(text) != std::string::npos) { + return true; + } + } + return false; +} + +static bool string_contains(const std::string & value, const char * text) { + return value.find(text) != std::string::npos; +} + +static void require_hrx_status(hrx_status_t status) { + if (ggml::hrx::ErrorResult error = ggml::hrx::take_status(status)) { + std::fprintf(stderr, "HRX status failed: %s\n", error->c_str()); + std::abort(); + } +} + +static ggml::hrx::DispatchTarget test_dispatch_target() { + return { "gfx1151" }; +} + +static const ggml::hrx::DispatchRegistry & test_dispatch_registry() { + const ggml::hrx::DispatchRegistry * registry = ggml::hrx::find_dispatch_registry(test_dispatch_target()); + REQUIRE(registry != nullptr); + return *registry; +} + +static std::string kernel_name_for_id(uint64_t kernel_id) { + const ggml::hrx::KernelResolveResult resolved = + ggml::hrx::resolve_kernel_definition(ggml::hrx::get_qwen_kernel_corpus(), "gfx1151", kernel_id); + REQUIRE(resolved.found()); + return ggml::hrx::kernel_definition_name(*resolved.definition); +} + +static void require_compile_parameter(const ggml::hrx::Dispatch & dispatch, + const char * name, + const std::string & value) { + const auto found = dispatch.kernel.compile_parameters.find(name); + REQUIRE(found != dispatch.kernel.compile_parameters.end()); + REQUIRE(found->second == value); +} + +static constexpr int64_t kQwenFlashHeadSize = 128; +static constexpr int64_t kQwenRouterExpertCount = 128; +static constexpr int64_t kQwenRouterRouteCount = 8; +static constexpr int64_t kQwenMoeHiddenSize = 2048; +static constexpr int64_t kQwenMoeIntermediateSize = 768; + +static size_t qwen_expert_table_size(int64_t token_count, int64_t expert_count = kQwenRouterExpertCount) { + return static_cast(expert_count + expert_count * token_count) * sizeof(int32_t); +} + +static size_t qwen_partition_table_size(int64_t token_count, + int64_t route_count = kQwenRouterRouteCount, + int64_t expert_count = kQwenRouterExpertCount) { + const int64_t assignment_count = token_count * route_count; + const int64_t assignment_partition_count = (assignment_count + 31) / 32; + return static_cast(1 + assignment_partition_count + expert_count) * sizeof(int32_t); +} + +static size_t qwen_q8_1_x4_size(int64_t token_count, int64_t hidden_size) { + return static_cast(token_count) * ggml_row_size(GGML_TYPE_Q8_1, hidden_size); +} + +static std::vector make_qwen_route_ids_iota(int64_t token_count, + int64_t route_count, + int64_t route_stride, + int64_t expert_count) { + std::vector route_ids(static_cast(token_count * route_stride), -1); + for (int64_t token = 0; token < token_count; ++token) { + for (int64_t route = 0; route < route_count; ++route) { + route_ids[static_cast(token * route_stride + route)] = + static_cast((token * route_count + route) % expert_count); + } + } + return route_ids; +} + +static std::vector make_qwen_expert_table_reference(const std::vector & route_ids, + int64_t token_count, + int64_t route_count, + int64_t route_stride, + int64_t expert_count) { + std::vector expert_table(qwen_expert_table_size(token_count, expert_count) / sizeof(int32_t), -1); + for (int64_t expert = 0; expert < expert_count; ++expert) { + expert_table[static_cast(expert)] = 0; + } + for (int64_t token = 0; token < token_count; ++token) { + for (int64_t route = 0; route < route_count; ++route) { + const int32_t expert = route_ids[static_cast(token * route_stride + route)]; + REQUIRE(expert >= 0); + REQUIRE(expert < expert_count); + int32_t & count = expert_table[static_cast(expert)]; + const size_t assignment_offset = + static_cast(expert_count + static_cast(expert) * token_count + count); + expert_table[assignment_offset] = static_cast(token * route_count + route); + ++count; + } + } + return expert_table; +} + +static std::vector make_qwen_partition_table_reference(const std::vector & expert_table, + int64_t token_count, + int64_t route_count, + int64_t expert_count) { + std::vector partition_table( + qwen_partition_table_size(token_count, route_count, expert_count) / sizeof(int32_t), -1); + int32_t partition_count = 0; + for (int64_t expert = 0; expert < expert_count; ++expert) { + const int32_t expert_assignment_count = expert_table[static_cast(expert)]; + REQUIRE(expert_assignment_count >= 0); + const int32_t expert_partition_count = (expert_assignment_count + 31) / 32; + for (int32_t partition = 0; partition < expert_partition_count; ++partition) { + int32_t row_count = expert_assignment_count - partition * 32; + if (row_count > 32) { + row_count = 32; + } + const int32_t descriptor = static_cast(expert) | (partition << 7) | ((row_count - 1) << 13); + partition_table[static_cast(1 + partition_count)] = descriptor; + ++partition_count; + } + } + partition_table[0] = partition_count; + return partition_table; +} + +static void require_qwen_expert_table_matches(const std::vector & actual, + const std::vector & expected, + int64_t token_count, + int64_t expert_count) { + REQUIRE(actual.size() == expected.size()); + for (int64_t expert = 0; expert < expert_count; ++expert) { + const size_t count_index = static_cast(expert); + REQUIRE(actual[count_index] == expected[count_index]); + for (int32_t ordinal = 0; ordinal < expected[count_index]; ++ordinal) { + const size_t assignment_index = + static_cast(expert_count + expert * token_count + static_cast(ordinal)); + REQUIRE(actual[assignment_index] == expected[assignment_index]); + } + } +} + +static void require_qwen_partition_table_matches(const std::vector & actual, + const std::vector & expected) { + REQUIRE(actual.size() == expected.size()); + REQUIRE(actual[0] == expected[0]); + for (int32_t i = 0; i < actual[0]; ++i) { + REQUIRE(actual[static_cast(1 + i)] == expected[static_cast(1 + i)]); + } +} + +static size_t qwen_routed_gate_up_f16_output_size(int64_t token_count) { + return static_cast(token_count * kQwenRouterRouteCount * kQwenMoeIntermediateSize) * sizeof(ggml_fp16_t); +} + +static size_t qwen_routed_down_f16_output_size(int64_t token_count) { + return static_cast(token_count * kQwenRouterRouteCount * kQwenMoeHiddenSize) * sizeof(ggml_fp16_t); +} + +static void set_qwen_flash_query_layout(ggml_tensor * tensor, int64_t head_count) { + REQUIRE(tensor != nullptr); + tensor->nb[0] = sizeof(float); + tensor->nb[1] = static_cast(head_count * kQwenFlashHeadSize) * sizeof(float); + tensor->nb[2] = static_cast(kQwenFlashHeadSize) * sizeof(float); +} + +static void set_qwen_flash_key_value_layout(ggml_tensor * tensor, int64_t head_count) { + REQUIRE(tensor != nullptr); + tensor->nb[0] = sizeof(ggml_fp16_t); + tensor->nb[1] = static_cast(head_count * kQwenFlashHeadSize) * sizeof(ggml_fp16_t); + tensor->nb[2] = static_cast(kQwenFlashHeadSize) * sizeof(ggml_fp16_t); +} + +static ggml_tensor * build_qwen_flash_attention_graph(ggml_context * ctx, + int64_t query_token_count, + int64_t key_value_token_count, + int64_t query_head_count, + int64_t key_value_head_count, + ggml_type query_type = GGML_TYPE_F32, + ggml_type key_value_type = GGML_TYPE_F16, + bool include_mask = true, + bool include_sinks = false, + int64_t head_size = kQwenFlashHeadSize, + float scale = 1.0f / std::sqrt(128.0f), + float max_bias = 0.0f, + float logit_softcap = 0.0f) { + ggml_tensor * query = ggml_new_tensor_3d(ctx, query_type, head_size, query_token_count, query_head_count); + ggml_tensor * key = ggml_new_tensor_3d(ctx, key_value_type, head_size, key_value_token_count, key_value_head_count); + ggml_tensor * value = + ggml_new_tensor_3d(ctx, key_value_type, head_size, key_value_token_count, key_value_head_count); + REQUIRE(query != nullptr); + REQUIRE(key != nullptr); + REQUIRE(value != nullptr); + if (query_type == GGML_TYPE_F32) { + set_qwen_flash_query_layout(query, query_head_count); + } + if (key_value_type == GGML_TYPE_F16) { + set_qwen_flash_key_value_layout(key, key_value_head_count); + set_qwen_flash_key_value_layout(value, key_value_head_count); + } + ggml_tensor * mask = nullptr; + if (include_mask) { + mask = ggml_new_tensor_2d(ctx, GGML_TYPE_F16, key_value_token_count, query_token_count); + REQUIRE(mask != nullptr); + } + ggml_tensor * output = ggml_flash_attn_ext(ctx, query, key, value, mask, scale, max_bias, logit_softcap); + REQUIRE(output != nullptr); + if (include_sinks) { + ggml_tensor * sinks = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, query_head_count); + REQUIRE(sinks != nullptr); + ggml_flash_attn_ext_add_sinks(output, sinks); + } + return output; +} + +static ggml_tensor * build_qwen_router_top8_graph(ggml_context * ctx, + ggml_tensor * logits, + ggml_tensor ** route_ids = nullptr, + ggml_sort_order order = GGML_SORT_ORDER_DESC, + int64_t route_count = kQwenRouterRouteCount, + float clamp_min = 1.0e-7f) { + ggml_tensor * probs = ggml_soft_max(ctx, logits); + REQUIRE(probs != nullptr); + ggml_tensor * probs_reshaped = ggml_reshape_3d(ctx, probs, 1, logits->ne[0], logits->ne[1]); + REQUIRE(probs_reshaped != nullptr); + ggml_tensor * argsort = ggml_argsort(ctx, probs, order); + REQUIRE(argsort != nullptr); + ggml_tensor * topk = ggml_view_2d(ctx, argsort, route_count, logits->ne[1], argsort->nb[1], 0); + REQUIRE(topk != nullptr); + if (route_ids != nullptr) { + *route_ids = topk; + } + ggml_tensor * selected = ggml_get_rows(ctx, probs_reshaped, topk); + REQUIRE(selected != nullptr); + ggml_tensor * selected_reshaped = ggml_reshape_2d(ctx, selected, route_count, logits->ne[1]); + REQUIRE(selected_reshaped != nullptr); + ggml_tensor * sum = ggml_sum_rows(ctx, selected_reshaped); + REQUIRE(sum != nullptr); + ggml_tensor * clamped_sum = ggml_clamp(ctx, sum, clamp_min, std::numeric_limits::infinity()); + REQUIRE(clamped_sum != nullptr); + ggml_tensor * normalized = ggml_div(ctx, selected_reshaped, clamped_sum); + REQUIRE(normalized != nullptr); + ggml_tensor * output = ggml_reshape_3d(ctx, normalized, 1, route_count, logits->ne[1]); + REQUIRE(output != nullptr); + return output; +} + +static std::vector traversal_indices(const ggml::hrx::Graph & graph) { + const ggml::hrx::GraphTraversalOrder order = ggml::hrx::GraphTraversalOrder::build(graph); + std::vector indices; + indices.reserve(order.nodes().size()); + for (const ggml::hrx::GraphNode * node : order.nodes()) { + size_t index = 0; + REQUIRE(node != nullptr); + REQUIRE(graph.index().node_index(node, index)); + indices.push_back(index); + } + return indices; +} + +static size_t find_position(const std::vector & indices, size_t node_index) { + const std::vector::const_iterator it = std::find(indices.begin(), indices.end(), node_index); + REQUIRE(it != indices.end()); + return static_cast(it - indices.begin()); +} + +static size_t producer_index_for_tensor(const ggml::hrx::Graph & graph, const ggml_tensor * tensor) { + const ggml::hrx::Value * value = graph.values().find_tensor(tensor); + REQUIRE(value != nullptr); + const ggml::hrx::GraphNode * producer = graph.index().producer(value->id); + REQUIRE(producer != nullptr); + size_t index = 0; + REQUIRE(graph.index().node_index(producer, index)); + return index; +} + +static bool match_dispatch_at_index(const ggml::hrx::Graph & graph, + const ggml::hrx::CommandPlan & plan, + const std::vector & covered_nodes, + size_t node_index, + ggml::hrx::DispatchMatch & match); + +static void append_match_to_plan(ggml::hrx::CommandPlan & plan, + ggml::hrx::DispatchMatch & match, + std::vector & covered_nodes, + ggml::hrx::Graph * graph = nullptr); + +static void run_status_checks() { + ggml::hrx::Status status; + REQUIRE(status.success()); + REQUIRE(status.errors().empty()); + + status.log("first"); + REQUIRE(!status.success()); + REQUIRE(!status.errors().empty()); + REQUIRE(status.errors().size() == 1); + REQUIRE(status.errors()[0] == "first"); + + status.log("value %d", 7); + REQUIRE(status.errors().size() == 2); + REQUIRE(status.errors()[1] == "value 7"); + + ggml::hrx::Status other; + other.log("third"); + status.append(other); + REQUIRE(status.errors().size() == 3); + REQUIRE(status.errors()[2] == "third"); +} + +static void run_command_plan_metadata_checks() { + const ggml::hrx::MoeRoutingResourceMetadata routing = { + 4, + 8, + 128, + 128, + }; + const ggml::hrx::CommandPlanResourceMetadata metadata = ggml::hrx::make_command_plan_resource_metadata(routing); + + REQUIRE(metadata.kind == ggml::hrx::CommandPlanResourceMetadataKind::MoeRoutingResource); + ggml::hrx::MoeRoutingResourceMetadata decoded; + REQUIRE(metadata.read(decoded)); + REQUIRE(decoded.token_count == routing.token_count); + REQUIRE(decoded.route_count == routing.route_count); + REQUIRE(decoded.route_stride == routing.route_stride); + REQUIRE(decoded.expert_count == routing.expert_count); + + const ggml::hrx::CommandPlanResourceMetadata empty; + REQUIRE(!empty.read(decoded)); + + ggml::hrx::CommandPlanMetadata metadata_plan; + ggml::hrx::Status status; + REQUIRE(metadata_plan.append_alternate_value( + { ggml::hrx::ValueId(1), ggml::hrx::ValueId(2), GGML_TYPE_F16, 16, "alternate" }, status)); + REQUIRE(metadata_plan.append_alternate_value( + { ggml::hrx::ValueId(1), ggml::hrx::ValueId(2), GGML_TYPE_F16, 16, "alternate" }, status)); + REQUIRE(metadata_plan.alternate_values().size() == 1); + REQUIRE(metadata_plan.find_alternate_value(ggml::hrx::ValueId(1), GGML_TYPE_F16, 16) != nullptr); + REQUIRE(metadata_plan.find_alternate_value(ggml::hrx::ValueId(1), GGML_TYPE_F32, 16) == nullptr); + REQUIRE(metadata_plan.find_alternate_value(ggml::hrx::ValueId(1), GGML_TYPE_F16, 32) == nullptr); + REQUIRE(!metadata_plan.append_alternate_value( + { ggml::hrx::ValueId(1), ggml::hrx::ValueId(3), GGML_TYPE_F16, 16, "alternate" }, status)); + REQUIRE(!status.success()); + + ggml::hrx::CommandPlanMetadata bundle_plan; + ggml::hrx::Status bundle_status; + const ggml::hrx::CommandPlanMoeRoutingBundle bundle = { + ggml::hrx::ValueId(10), + ggml::hrx::ValueId(11), + ggml::hrx::ValueId(12), + ggml::hrx::ValueId(13), + 128, + 64, + 4, + 8, + 128, + 128, + }; + REQUIRE(bundle_plan.append_moe_routing_bundle(bundle, bundle_status)); + REQUIRE(bundle_plan.append_moe_routing_bundle(bundle, bundle_status)); + REQUIRE(bundle_plan.moe_routing_bundles().size() == 1); + const ggml::hrx::CommandPlanMoeRoutingBundle * found_bundle = + bundle_plan.find_moe_routing_bundle(ggml::hrx::ValueId(10)); + REQUIRE(found_bundle != nullptr); + REQUIRE(found_bundle->route_weights == ggml::hrx::ValueId(11)); + REQUIRE(found_bundle->expert_table == ggml::hrx::ValueId(12)); + REQUIRE(found_bundle->partition_table == ggml::hrx::ValueId(13)); + ggml::hrx::CommandPlanMoeRoutingBundle conflicting_bundle = bundle; + conflicting_bundle.route_weights = ggml::hrx::ValueId(14); + REQUIRE(!bundle_plan.append_moe_routing_bundle(conflicting_bundle, bundle_status)); + REQUIRE(!bundle_status.success()); +} + +static bool has_dispatch_registration(const std::vector & registrations, + const char * name) { + for (const ggml::hrx::DispatchRegistration & registration : registrations) { + if (std::string(registration.name) == name) { + return true; + } + } + return false; +} + +static bool match_test_single_dispatch(const ggml::hrx::DispatchMatchContext & context, + ggml::hrx::DispatchMatch & match) { + ggml::hrx::Dispatch dispatch; + dispatch.kernel.integer_parameters.emplace("route", 1); + match.covered_nodes.push_back(context.root_index); + match.dispatches.push_back(std::move(dispatch)); + return true; +} + +static bool match_test_fused_dispatch(const ggml::hrx::DispatchMatchContext & context, + ggml::hrx::DispatchMatch & match) { + ggml::hrx::Dispatch dispatch; + dispatch.kernel.integer_parameters.emplace("route", 2); + match.covered_nodes.push_back(context.root_index); + match.dispatches.push_back(std::move(dispatch)); + return true; +} + +static bool match_test_wrong_root_dispatch(const ggml::hrx::DispatchMatchContext & context, + ggml::hrx::DispatchMatch & match) { + ggml::hrx::Dispatch dispatch; + dispatch.kernel.integer_parameters.emplace("route", 3); + match.covered_nodes.push_back(context.root_index); + match.dispatches.push_back(std::move(dispatch)); + return true; +} + +static void run_dispatch_registry_checks() { + const ggml::hrx::DispatchRegistry & registry = test_dispatch_registry(); + REQUIRE(ggml::hrx::find_dispatch_registry({ "gfx1100" }) != nullptr); + REQUIRE(ggml::hrx::find_dispatch_registry({ "gfx1151" }) != nullptr); + REQUIRE(ggml::hrx::find_dispatch_registry({ "gfx0000" }) == nullptr); + + REQUIRE(has_dispatch_registration(registry.registrations_for_root(GGML_OP_ADD), "common.add_f32")); + REQUIRE( + has_dispatch_registration(registry.registrations_for_root(GGML_OP_MUL_MAT), "llm.matmul.dense_q4k_f16_wmma")); + REQUIRE( + has_dispatch_registration(registry.registrations_for_root(GGML_OP_MUL_MAT), "llm.matmul.dense_q6k_f16_wmma")); + REQUIRE(has_dispatch_registration(registry.registrations_for_root(GGML_OP_MUL_MAT), "qwen.matmul.q6k_q8_1_x4")); + REQUIRE(has_dispatch_registration(registry.registrations_for_root(GGML_OP_MUL_MAT), + "llm.moe_router.projection_f32_four_row_wave32")); + REQUIRE( + has_dispatch_registration(registry.registrations_for_root(GGML_OP_RMS_NORM), "qwen.rmsnorm_f32.mul_weight")); + REQUIRE(has_dispatch_registration(registry.registrations_for_root(GGML_OP_RMS_NORM), + "qwen.rmsnorm_f32_quantize_q8_1_x4")); + REQUIRE(has_dispatch_registration(registry.registrations_for_root(GGML_OP_FLASH_ATTN_EXT), + "qwen.flash_attention_f32_f16_wmma")); + REQUIRE(has_dispatch_registration(registry.registrations_for_root(GGML_OP_RESHAPE), + "qwen.attention_postprocess_f32_f16")); + REQUIRE(has_dispatch_registration(registry.registrations_for_root(GGML_OP_GET_ROWS), + "qwen.preamble.token_embedding_q4k")); + REQUIRE(has_dispatch_registration(registry.registrations_for_root(GGML_OP_GET_ROWS), "common.gather_add_f32")); + REQUIRE(has_dispatch_registration(registry.registrations_for_root(GGML_OP_SOFT_MAX), "llm.moe_router.top8_f32")); + REQUIRE(has_dispatch_registration(registry.registrations_for_root(GGML_OP_MUL_MAT_ID), + "llm.routed_ffn.gate_up_swiglu_q4k_f16_wmma")); + REQUIRE(has_dispatch_registration(registry.registrations_for_root(GGML_OP_MUL_MAT_ID), + "llm.routed_ffn.down_q4k_f16_wmma_grouped")); + REQUIRE(has_dispatch_registration(registry.registrations_for_root(GGML_OP_MUL_MAT_ID), + "llm.routed_ffn.down_q6k_f16_wmma_grouped")); + REQUIRE(registry.single_op_registrations().size() >= 5); + + ggml::hrx::DispatchRegistryBuilder builder; + builder.add({ + "test.single_add", + GGML_OP_ADD, + ggml::hrx::DispatchMatchKind::SingleOp, + 1000, + ggml::hrx::DispatchSource::Common, + match_test_single_dispatch, + }); + builder.add({ + "test.fused_add", + GGML_OP_ADD, + ggml::hrx::DispatchMatchKind::Fused, + 0, + ggml::hrx::DispatchSource::Common, + match_test_fused_dispatch, + }); + builder.add({ + "test.wrong_root", + GGML_OP_MUL_MAT, + ggml::hrx::DispatchMatchKind::Fused, + 2000, + ggml::hrx::DispatchSource::Common, + match_test_wrong_root_dispatch, + }); + const ggml::hrx::DispatchRegistry ordering_registry = builder.build(); + + ggml::hrx::Graph graph; + graph.add_node(GGML_OP_ADD, ggml::hrx::ValueId(0), {}); + REQUIRE(graph.build_index().success()); + + const std::vector covered_nodes(graph.nodes().size(), false); + const ggml::hrx::CommandPlan plan; + const ggml::hrx::DispatchMatchContext context = { + graph, &graph.nodes().front(), + 0, covered_nodes, + plan, ggml::hrx::ValueId(static_cast(graph.values().size())), + }; + ggml::hrx::DispatchMatch match; + REQUIRE(ordering_registry.match(context, match)); + REQUIRE(match.dispatches.size() == 1); + REQUIRE(match.dispatches.front().kernel.integer_parameters.at("route") == 2); +} + +static void run_graph_import_checks() { + ggml_init_params params = {}; + params.mem_size = 256 * 1024; + params.no_alloc = true; + ggml_context * ctx = ggml_init(params); + REQUIRE(ctx != nullptr); + + ggml_tensor * a = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 8); + ggml_tensor * out = ggml_add(ctx, a, a); + REQUIRE(a != nullptr); + REQUIRE(out != nullptr); + + ggml_cgraph * graph = ggml_new_graph(ctx); + REQUIRE(graph != nullptr); + ggml_build_forward_expand(graph, out); + + ggml::hrx::GraphImportResult imported = ggml::hrx::import_ggml_graph(*graph); + REQUIRE(imported.valid()); + REQUIRE(imported.graph.nodes().size() == 1); + const ggml::hrx::GraphNode * node = &imported.graph.nodes().front(); + REQUIRE(node->op == GGML_OP_ADD); + REQUIRE(node->inputs.size() == 2); + REQUIRE(node->inputs[0] == node->inputs[1]); + REQUIRE(ggml::hrx::DispatchScheduler::supports_node(imported.graph, node, test_dispatch_target())); + + const ggml::hrx::Value * a_value = imported.graph.values().find_tensor(a); + const ggml::hrx::Value * out_value = imported.graph.values().find_tensor(out); + REQUIRE(a_value != nullptr); + REQUIRE(out_value != nullptr); + REQUIRE(a_value->kind == ggml::hrx::ValueKind::External); + REQUIRE(out_value->kind == ggml::hrx::ValueKind::External); + REQUIRE(!a_value->buffer.has_value()); + REQUIRE(!out_value->buffer.has_value()); + + const std::vector external_ids = imported.graph.values().external_value_ids(); + REQUIRE(external_ids.size() == 2); + REQUIRE(contains_value_id(external_ids, a_value->id)); + REQUIRE(contains_value_id(external_ids, out_value->id)); + + ggml::hrx::DispatchScheduler scheduler; + REQUIRE(scheduler.schedule_graph(imported.graph, test_dispatch_target())); + REQUIRE(scheduler.plan().valid()); + REQUIRE(scheduler.plan().dispatches.size() == 1); + REQUIRE(scheduler.plan().dispatches.front().bindings.size() == 3); + REQUIRE(scheduler.plan().dispatches.front().bindings[0].length == a_value->byte_count); + + ggml::hrx::CommandProgramBindings missing_bindings = + ggml::hrx::CommandProgramBindings::from_value_map(imported.graph.values()); + REQUIRE(!missing_bindings.valid()); + + REQUIRE(imported.graph.values().bind_buffer(a_value->id, { dummy_hrx_buffer(0x1000), 0, a_value->byte_count })); + ggml::hrx::CommandProgramBindings partial_bindings = + ggml::hrx::CommandProgramBindings::from_value_map(imported.graph.values()); + REQUIRE(!partial_bindings.valid()); + + REQUIRE(imported.graph.values().bind_buffer(out_value->id, { dummy_hrx_buffer(0x2000), 0, 0 })); + ggml::hrx::CommandProgramBindings empty_runtime_binding = + ggml::hrx::CommandProgramBindings::from_value_map(imported.graph.values()); + REQUIRE(!empty_runtime_binding.valid()); + + REQUIRE(imported.graph.values().bind_buffer(out_value->id, { dummy_hrx_buffer(0x2000), 0, out_value->byte_count })); + ggml::hrx::CommandProgramBindings runtime_bindings = + ggml::hrx::CommandProgramBindings::from_value_map(imported.graph.values()); + REQUIRE(runtime_bindings.valid()); + REQUIRE(runtime_bindings.bindings().size() == 2); + const ggml::hrx::CommandProgramBinding * a_binding = runtime_bindings.find(a_value->id); + REQUIRE(a_binding != nullptr); + REQUIRE(a_binding->buffer == dummy_hrx_buffer(0x1000)); + REQUIRE(a_binding->offset == 0); + REQUIRE(a_binding->length == a_value->byte_count); + REQUIRE(runtime_bindings.find(ggml::hrx::ValueId(123456)) == nullptr); + + const ggml::hrx::CommandProgramBindingsFingerprint runtime_fingerprint = + ggml::hrx::command_program_bindings_fingerprint(runtime_bindings); + REQUIRE(!runtime_fingerprint.value.empty()); + + ggml::hrx::ValueMap changed_identity = imported.graph.values(); + REQUIRE(changed_identity.bind_buffer( + a_value->id, { dummy_hrx_buffer(0x1000), 0, a_value->byte_count, 1, 0, a_value->byte_count })); + const ggml::hrx::CommandProgramBindings changed_identity_bindings = + ggml::hrx::CommandProgramBindings::from_value_map(changed_identity); + REQUIRE(changed_identity_bindings.valid()); + REQUIRE(ggml::hrx::command_program_bindings_fingerprint(changed_identity_bindings).value != + runtime_fingerprint.value); + + ggml::hrx::ValueMap changed_generation = imported.graph.values(); + REQUIRE(changed_generation.bind_buffer( + a_value->id, { dummy_hrx_buffer(0x1000), 0, a_value->byte_count, 0, 1, a_value->byte_count })); + const ggml::hrx::CommandProgramBindings changed_generation_bindings = + ggml::hrx::CommandProgramBindings::from_value_map(changed_generation); + REQUIRE(changed_generation_bindings.valid()); + REQUIRE(ggml::hrx::command_program_bindings_fingerprint(changed_generation_bindings).value != + runtime_fingerprint.value); + + ggml::hrx::ValueMap changed_capacity = imported.graph.values(); + REQUIRE(changed_capacity.bind_buffer( + a_value->id, { dummy_hrx_buffer(0x1000), 0, a_value->byte_count, 0, 0, a_value->byte_count + 256 })); + const ggml::hrx::CommandProgramBindings changed_capacity_bindings = + ggml::hrx::CommandProgramBindings::from_value_map(changed_capacity); + REQUIRE(changed_capacity_bindings.valid()); + REQUIRE(ggml::hrx::command_program_bindings_fingerprint(changed_capacity_bindings).value != + runtime_fingerprint.value); + + const ggml::hrx::CommandProgram commands = ggml::hrx::build_command_program( + imported.graph, scheduler.plan(), ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); + REQUIRE(commands.valid()); + REQUIRE(commands.commands.size() == 1); + const ggml::hrx::Command & command = commands.commands.front(); + REQUIRE(command.ordinal == 0); + REQUIRE(command.kind == ggml::hrx::CommandKind::Kernel); + REQUIRE(command.kernel.kernel_id != ggml::hrx::kUncatalogedKernelId); + REQUIRE(command.bindings.size() == 3); + REQUIRE(command.bindings[0].name == "a"); + REQUIRE(command.bindings[0].origin == ggml::hrx::CommandBindingOrigin::GraphValue); + REQUIRE(command.bindings[0].access == ggml::hrx::ResourceAccess::Read); + REQUIRE(command.bindings[1].name == "b"); + REQUIRE(command.bindings[1].origin == ggml::hrx::CommandBindingOrigin::GraphValue); + REQUIRE(command.bindings[1].access == ggml::hrx::ResourceAccess::Read); + REQUIRE(command.bindings[2].name == "output"); + REQUIRE(command.bindings[2].origin == ggml::hrx::CommandBindingOrigin::GraphValue); + REQUIRE(command.bindings[2].access == ggml::hrx::ResourceAccess::ReadWrite); + REQUIRE(command_program_verifies(commands)); + + const ggml::hrx::PreparedCommand default_prepared_command; + REQUIRE(default_prepared_command.kind == ggml::hrx::CommandKind::Invalid); + REQUIRE(ggml::hrx::command_kind_name(ggml::hrx::CommandKind::Invalid) == "Invalid"); + REQUIRE(ggml::hrx::command_kind_name(ggml::hrx::CommandKind::Kernel) == "Kernel"); + REQUIRE(ggml::hrx::command_kind_name(static_cast(255)) == "Unknown(255)"); + REQUIRE(ggml::hrx::command_binding_origin_name(ggml::hrx::CommandBindingOrigin::GraphValue) == "GraphValue"); + REQUIRE(ggml::hrx::command_binding_origin_name(ggml::hrx::CommandBindingOrigin::Transient) == "Transient"); + REQUIRE(ggml::hrx::command_binding_origin_name(ggml::hrx::CommandBindingOrigin::ProgramConstant) == + "ProgramConstant"); + REQUIRE(ggml::hrx::command_binding_origin_name(static_cast(255)) == + "Unknown(255)"); + REQUIRE(ggml::hrx::resource_access_name(ggml::hrx::ResourceAccess::Read) == "Read"); + REQUIRE(ggml::hrx::resource_access_name(ggml::hrx::ResourceAccess::Write) == "Write"); + REQUIRE(ggml::hrx::resource_access_name(ggml::hrx::ResourceAccess::ReadWrite) == "ReadWrite"); + REQUIRE(ggml::hrx::resource_access_name(static_cast(255)) == "Unknown(255)"); + + const std::string binding_text = ggml::hrx::format_command_binding(command.bindings[0]); + REQUIRE(string_contains(binding_text, "binding a")); + REQUIRE(string_contains(binding_text, "value=")); + REQUIRE(string_contains(binding_text, "origin=GraphValue")); + REQUIRE(string_contains(binding_text, "access=Read")); + REQUIRE(std::string(ggml::hrx::hrx_graph_replay_event_name(ggml::hrx::HrxGraphReplayEvent::Disabled)) == + "disabled"); + REQUIRE(std::string(ggml::hrx::hrx_graph_replay_event_name(ggml::hrx::HrxGraphReplayEvent::Hit)) == "hit"); + REQUIRE(string_contains(binding_text, "range=[0, ")); + REQUIRE(string_contains(binding_text, std::to_string(a_value->byte_count).c_str())); + + const std::string command_text = ggml::hrx::format_command(command); + REQUIRE(string_contains(command_text, "command 0")); + REQUIRE(string_contains(command_text, "kind=Kernel")); + REQUIRE(string_contains(command_text, "kernel_id=")); + REQUIRE(string_contains(command_text, "bindings=3")); + REQUIRE(string_contains(command_text, "deps=0")); + + const std::string program_text = ggml::hrx::format_command_program(commands); + REQUIRE(string_contains(program_text, "command_program commands=1")); + REQUIRE(string_contains(program_text, "command 0")); + REQUIRE(string_contains(program_text, "binding a")); + REQUIRE(string_contains(program_text, "binding b")); + REQUIRE(string_contains(program_text, "binding output")); + + ggml::hrx::ResolvedCommandProgram resolved = + ggml::hrx::resolve_command_program_bindings(commands, runtime_bindings); + REQUIRE(resolved.valid()); + REQUIRE(resolved.commands.size() == 1); + REQUIRE(resolved.commands.front().ordinal == command.ordinal); + REQUIRE(resolved.commands.front().kind == command.kind); + REQUIRE(resolved.commands.front().kernel.kernel_id == command.kernel.kernel_id); + REQUIRE(resolved.commands.front().bindings.size() == 3); + REQUIRE(resolved.commands.front().bindings[0].binding.name == "a"); + REQUIRE(resolved.commands.front().bindings[0].ref.buffer == dummy_hrx_buffer(0x1000)); + REQUIRE(resolved.commands.front().bindings[0].ref.offset == 0); + REQUIRE(resolved.commands.front().bindings[0].ref.length == a_value->byte_count); + REQUIRE(resolved.commands.front().bindings[1].binding.name == "b"); + REQUIRE(resolved.commands.front().bindings[1].ref.buffer == dummy_hrx_buffer(0x1000)); + REQUIRE(resolved.commands.front().bindings[1].ref.offset == 0); + REQUIRE(resolved.commands.front().bindings[1].ref.length == a_value->byte_count); + REQUIRE(resolved.commands.front().bindings[2].binding.name == "output"); + REQUIRE(resolved.commands.front().bindings[2].ref.buffer == dummy_hrx_buffer(0x2000)); + REQUIRE(resolved.commands.front().bindings[2].ref.offset == 0); + REQUIRE(resolved.commands.front().bindings[2].ref.length == out_value->byte_count); + + ggml::hrx::CommandProgram offset_command = copy_command_program_shape(commands); + offset_command.commands.front().bindings[0].offset = 4; + offset_command.commands.front().bindings[0].length = 8; + ggml::hrx::ValueMap offset_values = imported.graph.values(); + REQUIRE(offset_values.bind_buffer(a_value->id, { dummy_hrx_buffer(0x3000), 16, a_value->byte_count })); + REQUIRE(offset_values.bind_buffer(out_value->id, { dummy_hrx_buffer(0x4000), 32, out_value->byte_count })); + const ggml::hrx::CommandProgramBindings offset_bindings = + ggml::hrx::CommandProgramBindings::from_value_map(offset_values); + REQUIRE(offset_bindings.valid()); + resolved = ggml::hrx::resolve_command_program_bindings(offset_command, offset_bindings); + REQUIRE(resolved.valid()); + REQUIRE(resolved.commands.front().bindings[0].ref.buffer == dummy_hrx_buffer(0x3000)); + REQUIRE(resolved.commands.front().bindings[0].ref.offset == 20); + REQUIRE(resolved.commands.front().bindings[0].ref.length == 8); + + std::vector host_storage(a_value->byte_count + 64); + ggml::hrx::ValueBufferBinding host_value_binding; + host_value_binding.host_data = host_storage.data(); + host_value_binding.offset = 16; + host_value_binding.length = a_value->byte_count; + host_value_binding.identity = 42; + host_value_binding.generation = 1; + host_value_binding.capacity = host_storage.size(); + host_value_binding.weight = true; + ggml::hrx::ValueMap host_values = imported.graph.values(); + REQUIRE(host_values.bind_buffer(a_value->id, host_value_binding)); + REQUIRE(host_values.bind_buffer(out_value->id, { dummy_hrx_buffer(0x5000), 0, out_value->byte_count })); + const ggml::hrx::CommandProgramBindings host_bindings = + ggml::hrx::CommandProgramBindings::from_value_map(host_values); + REQUIRE(host_bindings.valid()); + const ggml::hrx::CommandProgramBinding * host_binding = host_bindings.find(a_value->id); + REQUIRE(host_binding != nullptr); + REQUIRE(host_binding->buffer == nullptr); + REQUIRE(host_binding->host_data == host_storage.data()); + REQUIRE(host_binding->offset == 16); + REQUIRE(host_binding->length == a_value->byte_count); + REQUIRE(host_binding->weight); + REQUIRE(ggml::hrx::command_program_bindings_fingerprint(host_bindings).value != runtime_fingerprint.value); + + resolved = ggml::hrx::resolve_command_program_bindings(commands, missing_bindings); + REQUIRE(!resolved.valid()); + REQUIRE(status_contains(resolved.status, "is not bound")); + REQUIRE(status_contains(resolved.status, "binding output")); + REQUIRE(status_contains(resolved.status, "value=")); + + resolved = ggml::hrx::resolve_command_program_bindings(commands, partial_bindings); + REQUIRE(!resolved.valid()); + REQUIRE(status_contains(resolved.status, "is not bound")); + + resolved = ggml::hrx::resolve_command_program_bindings(commands, empty_runtime_binding); + REQUIRE(!resolved.valid()); + REQUIRE(status_contains(resolved.status, "empty binding")); + + ggml::hrx::ValueMap null_values = imported.graph.values(); + REQUIRE(null_values.bind_buffer(a_value->id, { nullptr, 0, a_value->byte_count })); + REQUIRE(null_values.bind_buffer(out_value->id, { dummy_hrx_buffer(0x2000), 0, out_value->byte_count })); + const ggml::hrx::CommandProgramBindings null_bindings = + ggml::hrx::CommandProgramBindings::from_value_map(null_values); + REQUIRE(!null_bindings.valid()); + resolved = ggml::hrx::resolve_command_program_bindings(commands, null_bindings); + REQUIRE(!resolved.valid()); + REQUIRE(status_contains(resolved.status, "null buffer")); + REQUIRE(status_contains(resolved.status, "binding a")); + + ggml::hrx::CommandProgram empty_resolve_binding = copy_command_program_shape(commands); + empty_resolve_binding.commands.front().bindings[0].length = 0; + resolved = ggml::hrx::resolve_command_program_bindings(empty_resolve_binding, runtime_bindings); + REQUIRE(!resolved.valid()); + REQUIRE(status_contains(resolved.status, "empty range")); + REQUIRE(status_contains(resolved.status, "range=[0, 0)")); + + ggml::hrx::CommandProgram out_of_range_binding = copy_command_program_shape(commands); + out_of_range_binding.commands.front().bindings[0].offset = a_value->byte_count; + out_of_range_binding.commands.front().bindings[0].length = 4; + resolved = ggml::hrx::resolve_command_program_bindings(out_of_range_binding, runtime_bindings); + REQUIRE(!resolved.valid()); + REQUIRE(status_contains(resolved.status, "outside runtime binding length")); + REQUIRE(status_contains(resolved.status, "binding a")); + + ggml::hrx::CommandProgram unsupported_origin = copy_command_program_shape(commands); + unsupported_origin.commands.front().bindings[0].origin = static_cast(255); + resolved = ggml::hrx::resolve_command_program_bindings(unsupported_origin, runtime_bindings); + REQUIRE(!resolved.valid()); + REQUIRE(status_contains(resolved.status, "unsupported binding origin")); + REQUIRE(status_contains(resolved.status, "origin=Unknown(255)")); + ggml::hrx::VerificationResult verification = + ggml::hrx::verify_command_program(unsupported_origin, ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); + REQUIRE(!verification.valid()); + REQUIRE(status_contains(verification.status, "unsupported binding origin")); + REQUIRE(status_contains(verification.status, "origin=Unknown(255)")); + + ggml::hrx::CommandProgram invalid_kernel = copy_command_program_shape(commands); + invalid_kernel.commands.front().kernel.kernel_id = ggml::hrx::kUncatalogedKernelId; + verification = ggml::hrx::verify_command_program(invalid_kernel, ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); + REQUIRE(!verification.valid()); + REQUIRE(status_contains(verification.status, "command 0")); + REQUIRE(status_contains(verification.status, "kernel_id=")); + + ggml::hrx::CommandProgram empty_bindings = copy_command_program_shape(commands); + empty_bindings.commands.front().bindings.clear(); + verification = ggml::hrx::verify_command_program(empty_bindings, ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); + REQUIRE(!verification.valid()); + REQUIRE(status_contains(verification.status, "bindings=0")); + + ggml::hrx::CommandProgram empty_binding = copy_command_program_shape(commands); + empty_binding.commands.front().bindings[0].length = 0; + verification = ggml::hrx::verify_command_program(empty_binding, ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); + REQUIRE(!verification.valid()); + REQUIRE(status_contains(verification.status, "binding a")); + REQUIRE(status_contains(verification.status, "range=[0, 0)")); + + ggml::hrx::CommandProgram invalid_value = copy_command_program_shape(commands); + invalid_value.commands.front().bindings[0].value = ggml::hrx::ValueId(-1); + verification = ggml::hrx::verify_command_program(invalid_value, ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); + REQUIRE(!verification.valid()); + REQUIRE(status_contains(verification.status, "value=-1")); + + ggml::hrx::CommandProgram forward_dependency = copy_command_program_shape(commands); + forward_dependency.commands.front().dependencies.push_back(0); + verification = + ggml::hrx::verify_command_program(forward_dependency, ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); + REQUIRE(!verification.valid()); + REQUIRE(status_contains(verification.status, "forward dependency 0")); + + ggml::hrx::CommandProgram wrong_binding_name = copy_command_program_shape(commands); + wrong_binding_name.commands.front().bindings[0].name = "wrong"; + verification = + ggml::hrx::verify_command_program(wrong_binding_name, ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); + REQUIRE(!verification.valid()); + REQUIRE(status_contains(verification.status, "binding wrong")); + + ggml::hrx::CommandProgram wrong_binding_access = copy_command_program_shape(commands); + wrong_binding_access.commands.front().bindings[0].access = ggml::hrx::ResourceAccess::Write; + verification = + ggml::hrx::verify_command_program(wrong_binding_access, ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); + REQUIRE(!verification.valid()); + REQUIRE(status_contains(verification.status, "access=Write")); + + const ggml::hrx::KernelCorpus & corpus = ggml::hrx::get_qwen_kernel_corpus(); + const ggml::hrx::CommandProgramExecutionContext prepare_context = { + nullptr, nullptr, "gfx1151", &corpus, nullptr, nullptr, nullptr, nullptr, + }; + + ggml::hrx::PreparedCommandProgram prepared = + ggml::hrx::prepare_command_program(prepare_context, invalid_kernel, runtime_bindings); + REQUIRE(!prepared.valid()); + REQUIRE(status_contains(prepared.status, "kernel_id=")); + + prepared = ggml::hrx::prepare_command_program(prepare_context, commands, missing_bindings); + REQUIRE(!prepared.valid()); + REQUIRE(status_contains(prepared.status, "is not bound")); + + prepared = ggml::hrx::prepare_command_program(prepare_context, commands, runtime_bindings); + REQUIRE(!prepared.valid()); + REQUIRE(status_contains(prepared.status, "missing HRX device")); + + const ggml::hrx::CommandProgramExecutionContext missing_kernel_cache_context = { + reinterpret_cast(uintptr_t(1)), nullptr, "gfx1151", &corpus, nullptr, nullptr, nullptr, nullptr, + }; + prepared = ggml::hrx::prepare_command_program(missing_kernel_cache_context, commands, runtime_bindings); + REQUIRE(!prepared.valid()); + REQUIRE(status_contains(prepared.status, "missing HRX kernel executable cache")); + + ggml_free(ctx); +} + +static void run_graph_snapshot_diagnostics_checks() { + ggml_init_params params = {}; + params.mem_size = 256 * 1024; + params.no_alloc = true; + ggml_context * ctx = ggml_init(params); + REQUIRE(ctx != nullptr); + + ggml_tensor * a = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 8); + ggml_tensor * b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 8); + ggml_tensor * out = ggml_add(ctx, a, b); + REQUIRE(a != nullptr); + REQUIRE(b != nullptr); + REQUIRE(out != nullptr); + + ggml_cgraph * graph = ggml_new_graph(ctx); + REQUIRE(graph != nullptr); + ggml_build_forward_expand(graph, out); + + ggml::hrx::GraphImportResult imported = ggml::hrx::import_ggml_graph(*graph); + REQUIRE(imported.valid()); + + const std::string snapshot_json = ggml::hrx::serialize_graph_snapshot_json(imported.graph, "gfx1151", 42); + REQUIRE(string_contains(snapshot_json, "ggml-hrx-graph-snapshot-v1")); + REQUIRE(string_contains(snapshot_json, "ADD")); + + ggml::hrx::GraphSnapshotLoadResult loaded = ggml::hrx::load_graph_snapshot_json(snapshot_json); + REQUIRE(loaded.valid()); + REQUIRE(loaded.uid == 42); + REQUIRE(loaded.target == "gfx1151"); + REQUIRE(loaded.graph.nodes().size() == imported.graph.nodes().size()); + REQUIRE(loaded.graph.values().size() == imported.graph.values().size()); + + ggml::hrx::DispatchScheduler scheduler; + ggml::hrx::DispatchScheduleDiagnostics diagnostics; + REQUIRE(scheduler.schedule_graph(loaded.graph, { loaded.target }, &diagnostics)); + REQUIRE(scheduler.plan().valid()); + REQUIRE(scheduler.plan().dispatches.size() == 1); + + ggml::hrx::CommandProgram commands = ggml::hrx::build_command_program( + loaded.graph, scheduler.plan(), ggml::hrx::get_qwen_kernel_corpus(), loaded.target); + REQUIRE(commands.valid()); + REQUIRE(command_program_verifies(commands)); + + ggml_free(ctx); +} + +static void run_unmatched_graph_diagnostics_checks() { + ggml_init_params params = {}; + params.mem_size = 512 * 1024; + params.no_alloc = true; + ggml_context * ctx = ggml_init(params); + REQUIRE(ctx != nullptr); + + ggml_tensor * cache = ggml_new_tensor_2d(ctx, GGML_TYPE_F16, 512, 512); + ggml_tensor * rows = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 512, 2); + ggml_tensor * indices = ggml_new_tensor_1d(ctx, GGML_TYPE_I64, 2); + REQUIRE(cache != nullptr); + REQUIRE(rows != nullptr); + REQUIRE(indices != nullptr); + ggml_tensor * out = ggml_set_rows(ctx, cache, rows, indices); + REQUIRE(out != nullptr); + + ggml_cgraph * graph = ggml_new_graph(ctx); + REQUIRE(graph != nullptr); + ggml_build_forward_expand(graph, out); + + ggml::hrx::GraphImportResult imported = ggml::hrx::import_ggml_graph(*graph); + REQUIRE(imported.valid()); + REQUIRE(imported.graph.nodes().size() == 1); + REQUIRE(imported.graph.nodes().front().op == GGML_OP_SET_ROWS); + + ggml::hrx::DispatchScheduler scheduler; + ggml::hrx::DispatchScheduleDiagnostics diagnostics; + REQUIRE(!scheduler.schedule_graph(imported.graph, test_dispatch_target(), &diagnostics)); + REQUIRE(diagnostics.unsupported_node != nullptr); + REQUIRE(diagnostics.unsupported_node_index == 0); + REQUIRE(diagnostics.unsupported_node->op == GGML_OP_SET_ROWS); + REQUIRE(diagnostics.match.attempts.empty()); + REQUIRE(status_contains(scheduler.plan().status, "unsupported HRX node 0: SET_ROWS")); + + const std::string diagnostics_text = + ggml::hrx::format_schedule_diagnostics_text(imported.graph, scheduler.plan(), diagnostics); + REQUIRE(string_contains(diagnostics_text, "unsupported_node=0:SET_ROWS")); + REQUIRE(string_contains(diagnostics_text, "matcher_attempts=0")); + const std::string diagnostics_json = + ggml::hrx::serialize_schedule_diagnostics_json(imported.graph, scheduler.plan(), diagnostics); + REQUIRE(string_contains(diagnostics_json, "SET_ROWS")); + REQUIRE(string_contains(diagnostics_json, "matcher_attempts")); + + ggml_free(ctx); +} + +static void run_completion_counter_plan_checks() { + ggml::hrx::Graph graph; + ggml::hrx::CommandPlan plan; + const ggml::hrx::ValueId first_counter(1000); + const ggml::hrx::ValueId second_counter(1001); + plan.completion_counter_requests.push_back({ first_counter, "first_completion_counter", 1 }); + plan.completion_counter_requests.push_back({ second_counter, "second_completion_counters", 2 }); + + const ggml::hrx::CommandProgram commands = + ggml::hrx::build_command_program(graph, plan, ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); + REQUIRE(commands.valid()); + REQUIRE(commands.commands.empty()); + REQUIRE(commands.constant_initializations.empty()); + REQUIRE(commands.completion_counters.count == 3); + REQUIRE(commands.completion_counters.arena_offset == 0); + REQUIRE(commands.completion_counters.byte_count == 24); + REQUIRE(commands.transients.allocations.size() == 2); + REQUIRE(commands.transients.arena_size == 256); + REQUIRE(command_program_verifies(commands)); + + const ggml::hrx::TransientAllocation * first_allocation = + ggml::hrx::find_transient_allocation(commands.transients, first_counter); + const ggml::hrx::TransientAllocation * second_allocation = + ggml::hrx::find_transient_allocation(commands.transients, second_counter); + REQUIRE(first_allocation != nullptr); + REQUIRE(second_allocation != nullptr); + REQUIRE(first_allocation->size == sizeof(int32_t)); + REQUIRE(first_allocation->alignment == 16); + REQUIRE(first_allocation->arena_offset == commands.completion_counters.arena_offset); + REQUIRE(second_allocation->size == 2 * sizeof(int32_t)); + REQUIRE(second_allocation->alignment == 16); + REQUIRE(second_allocation->arena_offset == 16); + + ggml::hrx::CommandPlan bound_plan; + bound_plan.completion_counter_requests.push_back({ first_counter, "first_bound_completion_counter", 1 }); + bound_plan.completion_counter_requests.push_back({ second_counter, "second_bound_completion_counter", 1 }); + ggml::hrx::Dispatch first_dispatch; + first_dispatch.kernel = + ggml::hrx::make_kernel_specialization(ggml::hrx::kernel_catalog_ref("qwen3_moe", "ggml_add_f32")); + first_dispatch.kernel.integer_parameters.emplace("element_count", 1); + first_dispatch.bindings.push_back({ first_counter, 0, sizeof(int32_t) }); + first_dispatch.bindings.push_back({ second_counter, 0, sizeof(int32_t) }); + first_dispatch.bindings.push_back({ first_counter, 0, sizeof(int32_t) }); + bound_plan.dispatches.push_back(std::move(first_dispatch)); + ggml::hrx::Dispatch second_dispatch; + second_dispatch.kernel = + ggml::hrx::make_kernel_specialization(ggml::hrx::kernel_catalog_ref("qwen3_moe", "ggml_add_f32")); + second_dispatch.kernel.integer_parameters.emplace("element_count", 1); + second_dispatch.bindings.push_back({ second_counter, 0, sizeof(int32_t) }); + second_dispatch.bindings.push_back({ first_counter, 0, sizeof(int32_t) }); + second_dispatch.bindings.push_back({ second_counter, 0, sizeof(int32_t) }); + bound_plan.dispatches.push_back(std::move(second_dispatch)); + + const ggml::hrx::CommandProgram bound_commands = + ggml::hrx::build_command_program(graph, bound_plan, ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); + REQUIRE(bound_commands.valid()); + REQUIRE(bound_commands.commands.size() == 2); + REQUIRE(bound_commands.completion_counters.count == 2); + REQUIRE(bound_commands.completion_counters.arena_offset == 0); + REQUIRE(bound_commands.completion_counters.byte_count == 20); + REQUIRE(command_program_verifies(bound_commands)); + + const ggml::hrx::TransientArenaAllocationRef transient_arena = { + dummy_hrx_buffer(0x8000), + bound_commands.transients.arena_size, + 7, + }; + ggml::hrx::PreparedCommandProgram prepared_shape; + for (const ggml::hrx::Command & prepared_source : bound_commands.commands) { + ggml::hrx::PreparedCommand prepared_command; + prepared_command.ordinal = prepared_source.ordinal; + prepared_command.kind = prepared_source.kind; + prepared_command.kernel.specialization = prepared_source.kernel; + for (const ggml::hrx::CommandBinding & binding : prepared_source.bindings) { + prepared_command.kernel.bindings.push_back({ + binding, { dummy_hrx_buffer(0x4000), 123, binding.length } + }); + } + prepared_shape.commands.push_back(std::move(prepared_command)); + } + prepared_shape.bound_transient_arena_allocation_id = 1; + + REQUIRE(ggml::hrx::bind_prepared_command_program_transients(bound_commands, transient_arena, prepared_shape)); + REQUIRE(prepared_shape.commands[0].kernel.bindings[0].ref.buffer == dummy_hrx_buffer(0x8000)); + REQUIRE(prepared_shape.commands[0].kernel.bindings[0].ref.offset == 0); + REQUIRE(prepared_shape.commands[0].kernel.bindings[1].ref.buffer == dummy_hrx_buffer(0x8000)); + REQUIRE(prepared_shape.commands[0].kernel.bindings[1].ref.offset == 16); + REQUIRE(prepared_shape.commands[1].kernel.bindings[0].ref.offset == 16); + REQUIRE(prepared_shape.commands[1].kernel.bindings[1].ref.offset == 0); +} + +static void run_graph_index_checks() { + ggml_init_params params = {}; + params.mem_size = 256 * 1024; + params.no_alloc = true; + ggml_context * ctx = ggml_init(params); + REQUIRE(ctx != nullptr); + + ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 256, 1); + ggml_tensor * weight = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 256); + REQUIRE(input != nullptr); + REQUIRE(weight != nullptr); + ggml_tensor * rms = ggml_rms_norm(ctx, input, 0.000001f); + REQUIRE(rms != nullptr); + ggml_tensor * out = ggml_mul(ctx, rms, weight); + REQUIRE(out != nullptr); + + ggml_cgraph * graph = ggml_new_graph(ctx); + REQUIRE(graph != nullptr); + ggml_build_forward_expand(graph, out); + + ggml::hrx::GraphImportResult imported = ggml::hrx::import_ggml_graph(*graph); + REQUIRE(imported.valid()); + REQUIRE(imported.graph.has_index()); + REQUIRE(imported.graph.nodes().size() == 2); + const ggml::hrx::GraphNode * rms_node = &imported.graph.nodes()[0]; + const ggml::hrx::GraphNode * mul_node = &imported.graph.nodes()[1]; + REQUIRE(rms_node->op == GGML_OP_RMS_NORM); + REQUIRE(mul_node->op == GGML_OP_MUL); + const ggml::hrx::RmsNormParams * rms_params = ggml::hrx::op_params_as(rms_node->params); + REQUIRE(rms_params != nullptr); + REQUIRE(rms_params->eps == 0.000001f); + REQUIRE(imported.graph.index().producer(rms_node->output) == rms_node); + REQUIRE(imported.graph.index().producer(mul_node->output) == mul_node); + REQUIRE(imported.graph.index().has_single_consumer(rms_node->output)); + const std::vector & consumers = imported.graph.index().consumers(rms_node->output); + REQUIRE(consumers.size() == 1); + REQUIRE(consumers.front() == mul_node); + + ggml::hrx::DispatchScheduler scheduler; + REQUIRE(scheduler.schedule_graph(imported.graph, test_dispatch_target())); + REQUIRE(scheduler.plan().valid()); + REQUIRE(scheduler.plan().dispatches.size() == 1); + + ggml_free(ctx); +} + +static void run_graph_traversal_checks() { + { + ggml_init_params params = {}; + params.mem_size = 256 * 1024; + params.no_alloc = true; + ggml_context * ctx = ggml_init(params); + REQUIRE(ctx != nullptr); + + ggml_tensor * a = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 8); + ggml_tensor * b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 8); + ggml_tensor * c = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 8); + ggml_tensor * d = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 8); + ggml_tensor * out0 = ggml_add(ctx, a, b); + ggml_tensor * out1 = ggml_add(ctx, c, d); + REQUIRE(a != nullptr); + REQUIRE(b != nullptr); + REQUIRE(c != nullptr); + REQUIRE(d != nullptr); + REQUIRE(out0 != nullptr); + REQUIRE(out1 != nullptr); + + ggml_cgraph * graph = ggml_new_graph(ctx); + REQUIRE(graph != nullptr); + ggml_build_forward_expand(graph, out0); + ggml_build_forward_expand(graph, out1); + + ggml::hrx::GraphImportResult imported = ggml::hrx::import_ggml_graph(*graph); + REQUIRE(imported.valid()); + REQUIRE(imported.graph.nodes().size() == 2); + REQUIRE(imported.graph.nodes()[0].output == imported.graph.values().find_tensor(out0)->id); + REQUIRE(imported.graph.nodes()[1].output == imported.graph.values().find_tensor(out1)->id); + + const std::vector order = traversal_indices(imported.graph); + REQUIRE(order.size() == 2); + REQUIRE(order[0] == 0); + REQUIRE(order[1] == 1); + + ggml_free(ctx); + } + + { + ggml_init_params params = {}; + params.mem_size = 256 * 1024; + params.no_alloc = true; + ggml_context * ctx = ggml_init(params); + REQUIRE(ctx != nullptr); + + ggml_tensor * a = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 8); + ggml_tensor * b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 8); + ggml_tensor * add_out = ggml_add(ctx, a, b); + ggml_tensor * weight = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 4, 3); + ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 4, 2); + ggml_tensor * matmul = ggml_mul_mat(ctx, weight, input); + REQUIRE(a != nullptr); + REQUIRE(b != nullptr); + REQUIRE(add_out != nullptr); + REQUIRE(weight != nullptr); + REQUIRE(input != nullptr); + REQUIRE(matmul != nullptr); + + ggml_cgraph * graph = ggml_new_graph(ctx); + REQUIRE(graph != nullptr); + ggml_build_forward_expand(graph, add_out); + ggml_build_forward_expand(graph, matmul); + + ggml::hrx::GraphImportResult imported = ggml::hrx::import_ggml_graph(*graph); + REQUIRE(imported.valid()); + REQUIRE(imported.graph.nodes().size() == 2); + REQUIRE(imported.graph.nodes()[0].op == GGML_OP_ADD); + REQUIRE(imported.graph.nodes()[1].op == GGML_OP_MUL_MAT); + + const std::vector order = traversal_indices(imported.graph); + REQUIRE(order.size() == 2); + REQUIRE(order[0] == 1); + REQUIRE(order[1] == 0); + + ggml_free(ctx); + } + + { + ggml_init_params params = {}; + params.mem_size = 256 * 1024; + params.no_alloc = true; + ggml_context * ctx = ggml_init(params); + REQUIRE(ctx != nullptr); + + ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 256, 1); + ggml_tensor * weight = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 256); + ggml_tensor * rms = ggml_rms_norm(ctx, input, 0.000001f); + ggml_tensor * out = ggml_mul(ctx, rms, weight); + REQUIRE(input != nullptr); + REQUIRE(weight != nullptr); + REQUIRE(rms != nullptr); + REQUIRE(out != nullptr); + + ggml_cgraph * graph = ggml_new_graph(ctx); + REQUIRE(graph != nullptr); + ggml_build_forward_expand(graph, out); + + ggml::hrx::GraphImportResult imported = ggml::hrx::import_ggml_graph(*graph); + REQUIRE(imported.valid()); + REQUIRE(imported.graph.nodes().size() == 2); + REQUIRE(imported.graph.nodes()[0].op == GGML_OP_RMS_NORM); + REQUIRE(imported.graph.nodes()[1].op == GGML_OP_MUL); + + const std::vector order = traversal_indices(imported.graph); + REQUIRE(order.size() == 2); + REQUIRE(order[0] == 0); + REQUIRE(order[1] == 1); + + ggml_free(ctx); + } + + { + ggml_init_params params = {}; + params.mem_size = 256 * 1024; + params.no_alloc = true; + ggml_context * ctx = ggml_init(params); + REQUIRE(ctx != nullptr); + + ggml_tensor * a = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 8); + ggml_tensor * b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 8); + ggml_tensor * c = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 8); + ggml_tensor * left = ggml_add(ctx, a, b); + ggml_tensor * right = ggml_mul(ctx, a, c); + ggml_tensor * join = ggml_add(ctx, left, right); + REQUIRE(a != nullptr); + REQUIRE(b != nullptr); + REQUIRE(c != nullptr); + REQUIRE(left != nullptr); + REQUIRE(right != nullptr); + REQUIRE(join != nullptr); + + ggml_cgraph * graph = ggml_new_graph(ctx); + REQUIRE(graph != nullptr); + ggml_build_forward_expand(graph, join); + + ggml::hrx::GraphImportResult imported = ggml::hrx::import_ggml_graph(*graph); + REQUIRE(imported.valid()); + REQUIRE(imported.graph.nodes().size() == 3); + + const size_t left_index = producer_index_for_tensor(imported.graph, left); + const size_t right_index = producer_index_for_tensor(imported.graph, right); + const size_t join_index = producer_index_for_tensor(imported.graph, join); + + const std::vector order = traversal_indices(imported.graph); + REQUIRE(order.size() == 3); + REQUIRE(find_position(order, join_index) > find_position(order, left_index)); + REQUIRE(find_position(order, join_index) > find_position(order, right_index)); + + ggml_free(ctx); + } +} + +static void schedule_single_matmul_command(ggml_context * ctx, + ggml_tensor * output, + const char * expected_kernel_name, + int64_t expected_token_count, + int64_t expected_input_size, + int64_t expected_output_size) { + ggml_cgraph * graph = ggml_new_graph(ctx); + REQUIRE(graph != nullptr); + ggml_build_forward_expand(graph, output); + + ggml::hrx::GraphImportResult imported = ggml::hrx::import_ggml_graph(*graph); + REQUIRE(imported.valid()); + REQUIRE(imported.graph.nodes().size() == 1); + REQUIRE(imported.graph.nodes()[0].op == GGML_OP_MUL_MAT); + + ggml::hrx::DispatchScheduler scheduler; + REQUIRE(scheduler.schedule_graph(imported.graph, test_dispatch_target())); + REQUIRE(scheduler.plan().valid()); + REQUIRE(scheduler.plan().dispatches.size() == 1); + + const ggml::hrx::Dispatch & dispatch = scheduler.plan().dispatches.front(); + const std::string kernel_name = kernel_name_for_id(dispatch.kernel.kernel_id); + REQUIRE(kernel_name == expected_kernel_name); + REQUIRE(dispatch.kernel.integer_parameters.at("token_count") == expected_token_count); + REQUIRE(dispatch.bindings.size() == 3); + require_compile_parameter(dispatch, "qwen3_moe.workload.token_capacity", std::to_string(expected_token_count)); + if (string_contains(kernel_name, "dense_linear")) { + require_compile_parameter(dispatch, "qwen3_moe.dense_quantized.input_size", + std::to_string(expected_input_size)); + require_compile_parameter(dispatch, "qwen3_moe.dense_quantized.output_size", + std::to_string(expected_output_size)); + require_compile_parameter(dispatch, "qwen3_moe.dense_quantized.output_accumulation", "0"); + } else { + require_compile_parameter(dispatch, "qwen3_moe.model.hidden_size", std::to_string(expected_input_size)); + require_compile_parameter(dispatch, "qwen3_moe.router.expert_count", std::to_string(expected_output_size)); + } + + const ggml::hrx::CommandProgram commands = ggml::hrx::build_command_program( + imported.graph, scheduler.plan(), ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); + REQUIRE(commands.valid()); + REQUIRE(commands.commands.size() == 1); + REQUIRE(command_program_verifies(commands)); + REQUIRE(commands.commands.front().bindings.size() == 3); + REQUIRE(commands.commands.front().bindings[0].name == "input"); + REQUIRE(commands.commands.front().bindings[1].name == "weight"); + REQUIRE(commands.commands.front().bindings[2].name == "output"); +} + +static bool matmul_graph_is_supported(ggml_context * ctx, ggml_tensor * output) { + ggml_cgraph * graph = ggml_new_graph(ctx); + REQUIRE(graph != nullptr); + ggml_build_forward_expand(graph, output); + + ggml::hrx::GraphImportResult imported = ggml::hrx::import_ggml_graph(*graph); + REQUIRE(imported.valid()); + return ggml::hrx::DispatchScheduler::can_schedule_graph(imported.graph, test_dispatch_target()); +} + +static void schedule_qwen_terminal_q6k_q8_command(int64_t token_count) { + ggml_init_params params = {}; + params.mem_size = 4 * 1024 * 1024; + params.no_alloc = true; + ggml_context * ctx = ggml_init(params); + REQUIRE(ctx != nullptr); + + ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, token_count); + ggml_tensor * norm_weight = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 2048); + ggml_tensor * vocab_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_Q6_K, 2048, 151936); + REQUIRE(input != nullptr); + REQUIRE(norm_weight != nullptr); + REQUIRE(vocab_weight != nullptr); + ggml_tensor * rms = ggml_rms_norm(ctx, input, 0.000001f); + ggml_tensor * normalized = ggml_mul(ctx, rms, norm_weight); + ggml_tensor * logits = ggml_mul_mat(ctx, vocab_weight, normalized); + REQUIRE(rms != nullptr); + REQUIRE(normalized != nullptr); + REQUIRE(logits != nullptr); + + ggml_cgraph * graph = ggml_new_graph(ctx); + REQUIRE(graph != nullptr); + ggml_build_forward_expand(graph, logits); + + ggml::hrx::GraphImportResult imported = ggml::hrx::import_ggml_graph(*graph); + REQUIRE(imported.valid()); + REQUIRE(imported.graph.nodes().size() == 3); + + ggml::hrx::DispatchScheduler scheduler; + REQUIRE(scheduler.schedule_graph(imported.graph, test_dispatch_target())); + REQUIRE(scheduler.plan().valid()); + REQUIRE(scheduler.plan().dispatches.size() == 2); + REQUIRE(scheduler.plan().transients.size() == 1); + + const ggml::hrx::Dispatch & rms_dispatch = scheduler.plan().dispatches[0]; + REQUIRE(kernel_name_for_id(rms_dispatch.kernel.kernel_id) == "qwen3_moe:qwen3_moe_rmsnorm_f32_quantize_q8_1_x4"); + REQUIRE(rms_dispatch.kernel.integer_parameters.at("token_count") == token_count); + require_compile_parameter(rms_dispatch, "qwen3_moe.model.hidden_size", "2048"); + require_compile_parameter(rms_dispatch, "qwen3_moe.workload.token_capacity", std::to_string(token_count)); + REQUIRE(rms_dispatch.bindings.size() == 4); + + const ggml::hrx::CommandPlanTransient & q8_transient = scheduler.plan().transients.front(); + REQUIRE(q8_transient.size == qwen_q8_1_x4_size(token_count, 2048)); + REQUIRE(rms_dispatch.bindings[3].value == q8_transient.value); + REQUIRE(rms_dispatch.bindings[3].length == q8_transient.size); + + const ggml::hrx::Dispatch & vocab_dispatch = scheduler.plan().dispatches[1]; + REQUIRE(kernel_name_for_id(vocab_dispatch.kernel.kernel_id) == "qwen3_moe:ggml_linear_q6k_q8_1_x4"); + REQUIRE(vocab_dispatch.kernel.integer_parameters.at("token_count") == token_count); + REQUIRE(vocab_dispatch.kernel.integer_parameters.at("input_size") == 2048); + REQUIRE(vocab_dispatch.kernel.integer_parameters.at("output_size") == 151936); + require_compile_parameter(vocab_dispatch, "ggml.linear_q6k_q8_1_x4.token_capacity", std::to_string(token_count)); + require_compile_parameter(vocab_dispatch, "ggml.linear_q6k_q8_1_x4.output_capacity", "151936"); + REQUIRE(vocab_dispatch.bindings.size() == 3); + REQUIRE(vocab_dispatch.bindings[0].value == q8_transient.value); + REQUIRE(vocab_dispatch.bindings[0].length == q8_transient.size); + + const ggml::hrx::CommandProgram commands = ggml::hrx::build_command_program( + imported.graph, scheduler.plan(), ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); + REQUIRE(commands.valid()); + REQUIRE(commands.commands.size() == 2); + REQUIRE(command_program_verifies(commands)); + REQUIRE(commands.commands[0].bindings.size() == 4); + REQUIRE(commands.commands[0].bindings[3].name == "q8_output"); + REQUIRE(commands.commands[0].bindings[3].origin == ggml::hrx::CommandBindingOrigin::Transient); + REQUIRE(commands.commands[1].bindings.size() == 3); + REQUIRE(commands.commands[1].bindings[0].name == "q8_input"); + REQUIRE(commands.commands[1].bindings[0].origin == ggml::hrx::CommandBindingOrigin::Transient); + REQUIRE(commands.commands[1].bindings[0].value == q8_transient.value); + + ggml_free(ctx); +} + +static bool graph_is_supported(ggml_context * ctx, ggml_tensor * output) { + ggml_cgraph * graph = ggml_new_graph(ctx); + REQUIRE(graph != nullptr); + ggml_build_forward_expand(graph, output); + + ggml::hrx::GraphImportResult imported = ggml::hrx::import_ggml_graph(*graph); + REQUIRE(imported.valid()); + return ggml::hrx::DispatchScheduler::can_schedule_graph(imported.graph, test_dispatch_target()); +} + +struct ManualQwenRouterTop8Graph { + ggml::hrx::Graph graph; + ggml::hrx::ValueId route_ids; +}; + +static ggml::hrx::ValueId add_manual_graph_tensor(ggml::hrx::Graph & graph, + ggml_tensor * tensor, + ggml::hrx::ValueKind kind) { + REQUIRE(tensor != nullptr); + return graph.values().get_or_add_tensor_value(tensor, kind); +} + +static ManualQwenRouterTop8Graph build_manual_qwen_router_top8_graph(ggml_context * ctx, + int64_t expert_count, + int64_t route_count, + int64_t token_count) { + ManualQwenRouterTop8Graph manual; + ggml::hrx::Graph & graph = manual.graph; + + const ggml::hrx::ValueId logits = add_manual_graph_tensor( + graph, ggml_new_tensor_2d(ctx, GGML_TYPE_F32, expert_count, token_count), ggml::hrx::ValueKind::External); + const ggml::hrx::ValueId probs = add_manual_graph_tensor( + graph, ggml_new_tensor_2d(ctx, GGML_TYPE_F32, expert_count, token_count), ggml::hrx::ValueKind::Transient); + const ggml::hrx::ValueId probs_reshaped = add_manual_graph_tensor( + graph, ggml_new_tensor_3d(ctx, GGML_TYPE_F32, 1, expert_count, token_count), ggml::hrx::ValueKind::Transient); + const ggml::hrx::ValueId argsort = add_manual_graph_tensor( + graph, ggml_new_tensor_2d(ctx, GGML_TYPE_I32, expert_count, token_count), ggml::hrx::ValueKind::Transient); + manual.route_ids = add_manual_graph_tensor(graph, ggml_new_tensor_2d(ctx, GGML_TYPE_I32, route_count, token_count), + ggml::hrx::ValueKind::Transient); + const ggml::hrx::ValueId selected = add_manual_graph_tensor( + graph, ggml_new_tensor_3d(ctx, GGML_TYPE_F32, 1, route_count, token_count), ggml::hrx::ValueKind::Transient); + const ggml::hrx::ValueId weights_flat = add_manual_graph_tensor( + graph, ggml_new_tensor_2d(ctx, GGML_TYPE_F32, route_count, token_count), ggml::hrx::ValueKind::Transient); + const ggml::hrx::ValueId sum = add_manual_graph_tensor( + graph, ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 1, token_count), ggml::hrx::ValueKind::Transient); + const ggml::hrx::ValueId clamped = add_manual_graph_tensor( + graph, ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 1, token_count), ggml::hrx::ValueKind::Transient); + const ggml::hrx::ValueId normalized = add_manual_graph_tensor( + graph, ggml_new_tensor_2d(ctx, GGML_TYPE_F32, route_count, token_count), ggml::hrx::ValueKind::Transient); + const ggml::hrx::ValueId route_weights = add_manual_graph_tensor( + graph, ggml_new_tensor_3d(ctx, GGML_TYPE_F32, 1, route_count, token_count), ggml::hrx::ValueKind::External); + + ggml::hrx::GraphNode & softmax = graph.add_node(GGML_OP_SOFT_MAX, probs, { logits }); + softmax.params = ggml::hrx::SoftMaxParams{ 1.0f, 0.0f }; + graph.add_node(GGML_OP_RESHAPE, probs_reshaped, { probs }); + ggml::hrx::GraphNode & argsort_node = graph.add_node(GGML_OP_ARGSORT, argsort, { probs }); + argsort_node.params = ggml::hrx::ArgsortParams{ GGML_SORT_ORDER_DESC }; + graph.add_node(GGML_OP_VIEW, manual.route_ids, { argsort }); + graph.add_node(GGML_OP_GET_ROWS, selected, { probs_reshaped, manual.route_ids }); + graph.add_node(GGML_OP_RESHAPE, weights_flat, { selected }); + graph.add_node(GGML_OP_SUM_ROWS, sum, { weights_flat }); + ggml::hrx::GraphNode & clamp = graph.add_node(GGML_OP_CLAMP, clamped, { sum }); + clamp.params = ggml::hrx::ClampParams{ 0.00006103515625f, std::numeric_limits::infinity() }; + graph.add_node(GGML_OP_DIV, normalized, { weights_flat, clamped }); + graph.add_node(GGML_OP_RESHAPE, route_weights, { normalized }); + REQUIRE(graph.build_index().success()); + return manual; +} + +static ggml::hrx::Graph build_manual_token_embedding_graph(ggml_tensor * weight, + ggml_tensor * token_ids, + ggml_tensor * output) { + ggml::hrx::Graph graph; + ggml::hrx::ValueId weight_value = graph.values().get_or_add_tensor_value(weight, ggml::hrx::ValueKind::External); + ggml::hrx::ValueId token_ids_value = + graph.values().get_or_add_tensor_value(token_ids, ggml::hrx::ValueKind::External); + ggml::hrx::ValueId output_value = graph.values().get_or_add_tensor_value(output, ggml::hrx::ValueKind::External); + graph.add_node(GGML_OP_GET_ROWS, output_value, { weight_value, token_ids_value }); + REQUIRE(graph.build_index().success()); + return graph; +} + +static bool manual_token_embedding_graph_is_supported(ggml_context * ctx, + ggml_type weight_type, + ggml_type token_ids_type, + ggml_type output_type, + int64_t hidden_size, + int64_t vocabulary_count, + int64_t token_count, + int64_t output_hidden_size = -1, + int64_t output_token_count = -1) { + if (output_hidden_size < 0) { + output_hidden_size = hidden_size; + } + if (output_token_count < 0) { + output_token_count = token_count; + } + ggml_tensor * weight = ggml_new_tensor_2d(ctx, weight_type, hidden_size, vocabulary_count); + ggml_tensor * token_ids = ggml_new_tensor_1d(ctx, token_ids_type, token_count); + ggml_tensor * output = ggml_new_tensor_2d(ctx, output_type, output_hidden_size, output_token_count); + REQUIRE(weight != nullptr); + REQUIRE(token_ids != nullptr); + REQUIRE(output != nullptr); + + ggml::hrx::Graph graph = build_manual_token_embedding_graph(weight, token_ids, output); + return ggml::hrx::DispatchScheduler::can_schedule_graph(graph, test_dispatch_target()); +} + +static void schedule_qwen_token_embedding_command(ggml_context * ctx, + ggml_tensor * output, + int64_t expected_token_count, + int64_t expected_vocabulary_count) { + ggml_cgraph * graph = ggml_new_graph(ctx); + REQUIRE(graph != nullptr); + ggml_build_forward_expand(graph, output); + + ggml::hrx::GraphImportResult imported = ggml::hrx::import_ggml_graph(*graph); + REQUIRE(imported.valid()); + REQUIRE(imported.graph.nodes().size() == 1); + REQUIRE(imported.graph.nodes()[0].op == GGML_OP_GET_ROWS); + + ggml::hrx::DispatchScheduler scheduler; + REQUIRE(scheduler.schedule_graph(imported.graph, test_dispatch_target())); + REQUIRE(scheduler.plan().valid()); + REQUIRE(scheduler.plan().dispatches.size() == 1); + + const ggml::hrx::Dispatch & dispatch = scheduler.plan().dispatches.front(); + const std::string kernel_name = kernel_name_for_id(dispatch.kernel.kernel_id); + REQUIRE(kernel_name == "qwen3_moe:qwen_token_embedding_q4k"); + REQUIRE(dispatch.kernel.integer_parameters.at("token_count") == expected_token_count); + REQUIRE(dispatch.kernel.integer_parameters.at("vocabulary_count") == expected_vocabulary_count); + REQUIRE(dispatch.bindings.size() == 3); + + const ggml::hrx::CommandProgram commands = ggml::hrx::build_command_program( + imported.graph, scheduler.plan(), ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); + REQUIRE(commands.valid()); + REQUIRE(commands.commands.size() == 1); + REQUIRE(command_program_verifies(commands)); + REQUIRE(commands.commands.front().bindings.size() == 3); + REQUIRE(commands.commands.front().bindings[0].name == "token_ids"); + REQUIRE(commands.commands.front().bindings[1].name == "weight"); + REQUIRE(commands.commands.front().bindings[2].name == "output"); +} + +static void run_qwen_token_embedding_dispatch_checks() { + ggml_init_params params = {}; + params.mem_size = 2 * 1024 * 1024; + params.no_alloc = true; + ggml_context * ctx = ggml_init(params); + REQUIRE(ctx != nullptr); + + { + ggml_tensor * weight = ggml_new_tensor_2d(ctx, GGML_TYPE_Q4_K, 2048, 151936); + ggml_tensor * token_ids = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, 1); + REQUIRE(weight != nullptr); + REQUIRE(token_ids != nullptr); + ggml_tensor * output = ggml_get_rows(ctx, weight, token_ids); + REQUIRE(output != nullptr); + schedule_qwen_token_embedding_command(ctx, output, 1, 151936); + } + { + ggml_tensor * weight = ggml_new_tensor_2d(ctx, GGML_TYPE_Q4_K, 2048, 151936); + ggml_tensor * token_ids = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, 13); + REQUIRE(weight != nullptr); + REQUIRE(token_ids != nullptr); + ggml_tensor * output = ggml_get_rows(ctx, weight, token_ids); + REQUIRE(output != nullptr); + schedule_qwen_token_embedding_command(ctx, output, 13, 151936); + } + + REQUIRE( + !manual_token_embedding_graph_is_supported(ctx, GGML_TYPE_F32, GGML_TYPE_I32, GGML_TYPE_F32, 2048, 151936, 1)); + REQUIRE( + !manual_token_embedding_graph_is_supported(ctx, GGML_TYPE_Q4_K, GGML_TYPE_I64, GGML_TYPE_F32, 2048, 151936, 1)); + REQUIRE( + !manual_token_embedding_graph_is_supported(ctx, GGML_TYPE_Q4_K, GGML_TYPE_I32, GGML_TYPE_F16, 2048, 151936, 1)); + REQUIRE( + !manual_token_embedding_graph_is_supported(ctx, GGML_TYPE_Q4_K, GGML_TYPE_I32, GGML_TYPE_F32, 1024, 151936, 1)); + REQUIRE( + !manual_token_embedding_graph_is_supported(ctx, GGML_TYPE_Q4_K, GGML_TYPE_I32, GGML_TYPE_F32, 3072, 248320, 1)); + REQUIRE(!manual_token_embedding_graph_is_supported(ctx, GGML_TYPE_Q4_K, GGML_TYPE_I32, GGML_TYPE_F32, 2048, 151936, + 1, 2048, 2)); + + ggml_free(ctx); +} + +static ggml::hrx::Graph build_manual_gather_add_graph(ggml_context * ctx, + int64_t hidden_size, + int64_t source_token_count, + int64_t output_token_count, + bool shared_row_ids, + int64_t second_source_hidden_size = -1) { + if (second_source_hidden_size < 0) { + second_source_hidden_size = hidden_size; + } + + ggml::hrx::Graph graph; + + ggml_tensor * attention = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, hidden_size, source_token_count); + ggml_tensor * residual = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, second_source_hidden_size, source_token_count); + ggml_tensor * row_ids0 = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, output_token_count); + ggml_tensor * row_ids1 = shared_row_ids ? row_ids0 : ggml_new_tensor_1d(ctx, GGML_TYPE_I32, output_token_count); + ggml_tensor * selected0 = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, hidden_size, output_token_count); + ggml_tensor * selected1 = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, hidden_size, output_token_count); + ggml_tensor * output = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, hidden_size, output_token_count); + REQUIRE(attention != nullptr); + REQUIRE(residual != nullptr); + REQUIRE(row_ids0 != nullptr); + REQUIRE(row_ids1 != nullptr); + REQUIRE(selected0 != nullptr); + REQUIRE(selected1 != nullptr); + REQUIRE(output != nullptr); + + const ggml::hrx::ValueId attention_value = + graph.values().get_or_add_tensor_value(attention, ggml::hrx::ValueKind::External); + const ggml::hrx::ValueId residual_value = + graph.values().get_or_add_tensor_value(residual, ggml::hrx::ValueKind::External); + const ggml::hrx::ValueId row_ids0_value = + graph.values().get_or_add_tensor_value(row_ids0, ggml::hrx::ValueKind::External); + const ggml::hrx::ValueId row_ids1_value = + graph.values().get_or_add_tensor_value(row_ids1, ggml::hrx::ValueKind::External); + const ggml::hrx::ValueId selected0_value = + graph.values().get_or_add_tensor_value(selected0, ggml::hrx::ValueKind::Transient); + const ggml::hrx::ValueId selected1_value = + graph.values().get_or_add_tensor_value(selected1, ggml::hrx::ValueKind::Transient); + const ggml::hrx::ValueId output_value = + graph.values().get_or_add_tensor_value(output, ggml::hrx::ValueKind::External); + + graph.add_node(GGML_OP_GET_ROWS, selected0_value, { attention_value, row_ids0_value }); + graph.add_node(GGML_OP_GET_ROWS, selected1_value, { residual_value, row_ids1_value }); + graph.add_node(GGML_OP_ADD, output_value, { selected0_value, selected1_value }); + REQUIRE(graph.build_index().success()); + return graph; +} + +static bool manual_gather_add_graph_is_supported(ggml_context * ctx, + int64_t hidden_size, + int64_t source_token_count, + int64_t output_token_count, + bool shared_row_ids, + int64_t second_source_hidden_size = -1) { + ggml::hrx::Graph graph = build_manual_gather_add_graph(ctx, hidden_size, source_token_count, output_token_count, + shared_row_ids, second_source_hidden_size); + return ggml::hrx::DispatchScheduler::can_schedule_graph(graph, test_dispatch_target()); +} + +static bool partial_gather_add_graph_is_supported(ggml_context * ctx) { + ggml::hrx::Graph graph; + + ggml_tensor * attention = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, 13); + ggml_tensor * row_ids = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, 1); + ggml_tensor * selected = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, 1); + ggml_tensor * add_input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, 1); + ggml_tensor * add_output = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, 1); + REQUIRE(attention != nullptr); + REQUIRE(row_ids != nullptr); + REQUIRE(selected != nullptr); + REQUIRE(add_input != nullptr); + REQUIRE(add_output != nullptr); + + const ggml::hrx::ValueId attention_value = + graph.values().get_or_add_tensor_value(attention, ggml::hrx::ValueKind::External); + const ggml::hrx::ValueId row_ids_value = + graph.values().get_or_add_tensor_value(row_ids, ggml::hrx::ValueKind::External); + const ggml::hrx::ValueId selected_value = + graph.values().get_or_add_tensor_value(selected, ggml::hrx::ValueKind::Transient); + const ggml::hrx::ValueId add_input_value = + graph.values().get_or_add_tensor_value(add_input, ggml::hrx::ValueKind::External); + const ggml::hrx::ValueId add_output_value = + graph.values().get_or_add_tensor_value(add_output, ggml::hrx::ValueKind::External); + + graph.add_node(GGML_OP_GET_ROWS, selected_value, { attention_value, row_ids_value }); + graph.add_node(GGML_OP_ADD, add_output_value, { selected_value, add_input_value }); + REQUIRE(graph.build_index().success()); + return ggml::hrx::DispatchScheduler::can_schedule_graph(graph, test_dispatch_target()); +} + +static void schedule_gather_add_command(ggml::hrx::Graph & graph, + int64_t expected_hidden_size, + int64_t expected_source_token_count, + int64_t expected_output_token_count) { + ggml::hrx::DispatchScheduler scheduler; + REQUIRE(scheduler.schedule_graph(graph, test_dispatch_target())); + REQUIRE(scheduler.plan().valid()); + REQUIRE(scheduler.plan().dispatches.size() == 1); + + const ggml::hrx::Dispatch & dispatch = scheduler.plan().dispatches.front(); + REQUIRE(kernel_name_for_id(dispatch.kernel.kernel_id) == "qwen3_moe:ggml_gather_add_f32"); + REQUIRE(dispatch.kernel.integer_parameters.at("hidden_size") == expected_hidden_size); + REQUIRE(dispatch.kernel.integer_parameters.at("source_token_count") == expected_source_token_count); + REQUIRE(dispatch.kernel.integer_parameters.at("output_token_count") == expected_output_token_count); + REQUIRE(dispatch.bindings.size() == 4); + + const ggml::hrx::CommandProgram commands = + ggml::hrx::build_command_program(graph, scheduler.plan(), ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); + REQUIRE(commands.valid()); + REQUIRE(commands.commands.size() == 1); + REQUIRE(command_program_verifies(commands)); + REQUIRE(commands.commands.front().bindings.size() == 4); + REQUIRE(commands.commands.front().bindings[0].name == "attention"); + REQUIRE(commands.commands.front().bindings[1].name == "residual"); + REQUIRE(commands.commands.front().bindings[2].name == "output_ids"); + REQUIRE(commands.commands.front().bindings[3].name == "output"); +} + +static void run_gather_add_dispatch_checks() { + ggml_init_params params = {}; + params.mem_size = 4 * 1024 * 1024; + params.no_alloc = true; + ggml_context * ctx = ggml_init(params); + REQUIRE(ctx != nullptr); + + { + ggml::hrx::Graph graph = build_manual_gather_add_graph(ctx, 2048, 13, 1, true); + schedule_gather_add_command(graph, 2048, 13, 1); + } + { + ggml::hrx::Graph graph = build_manual_gather_add_graph(ctx, 2048, 128, 8, true); + schedule_gather_add_command(graph, 2048, 128, 8); + } + + REQUIRE(!manual_gather_add_graph_is_supported(ctx, 2048, 13, 1, false)); + REQUIRE(!manual_gather_add_graph_is_supported(ctx, 2048, 13, 1, true, 1024)); + REQUIRE(!manual_gather_add_graph_is_supported(ctx, 96, 13, 1, true)); + REQUIRE(!partial_gather_add_graph_is_supported(ctx)); + + ggml_free(ctx); +} + +static void schedule_qwen_flash_attention_command(ggml_context * ctx, ggml_tensor * output) { + ggml_cgraph * graph = ggml_new_graph(ctx); + REQUIRE(graph != nullptr); + ggml_build_forward_expand(graph, output); + + ggml::hrx::GraphImportResult imported = ggml::hrx::import_ggml_graph(*graph); + REQUIRE(imported.valid()); + REQUIRE(imported.graph.nodes().size() == 1); + REQUIRE(imported.graph.nodes()[0].op == GGML_OP_FLASH_ATTN_EXT); + + ggml::hrx::DispatchScheduler scheduler; + REQUIRE(scheduler.schedule_graph(imported.graph, test_dispatch_target())); + REQUIRE(scheduler.plan().valid()); + REQUIRE(scheduler.plan().dispatches.size() == 1); + + const ggml::hrx::Dispatch & dispatch = scheduler.plan().dispatches.front(); + const std::string kernel_name = kernel_name_for_id(dispatch.kernel.kernel_id); + REQUIRE(kernel_name == "qwen3_moe:qwen3_moe_flash_attention_f32_f16_wmma"); + REQUIRE(dispatch.kernel.integer_parameters.at("query_token_count") == 4); + REQUIRE(dispatch.kernel.integer_parameters.at("key_value_token_count") == 8); + REQUIRE(dispatch.bindings.size() == 5); + require_compile_parameter(dispatch, "qwen3_moe.attention.query_head_count", "4"); + require_compile_parameter(dispatch, "qwen3_moe.attention.key_value_head_count", "2"); + require_compile_parameter(dispatch, "qwen3_moe.workload.token_capacity", "4"); + + const ggml::hrx::CommandProgram commands = ggml::hrx::build_command_program( + imported.graph, scheduler.plan(), ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); + REQUIRE(commands.valid()); + REQUIRE(commands.commands.size() == 1); + REQUIRE(command_program_verifies(commands)); + REQUIRE(commands.commands.front().bindings.size() == 5); + REQUIRE(commands.commands.front().bindings[0].name == "query"); + REQUIRE(commands.commands.front().bindings[1].name == "key"); + REQUIRE(commands.commands.front().bindings[2].name == "value"); + REQUIRE(commands.commands.front().bindings[3].name == "mask"); + REQUIRE(commands.commands.front().bindings[4].name == "output"); +} + +static void run_qwen_flash_attention_dispatch_checks() { + ggml_init_params params = {}; + params.mem_size = 4 * 1024 * 1024; + params.no_alloc = true; + ggml_context * ctx = ggml_init(params); + REQUIRE(ctx != nullptr); + + { + ggml_tensor * output = build_qwen_flash_attention_graph(ctx, 4, 8, 4, 2); + schedule_qwen_flash_attention_command(ctx, output); + } + { + ggml_tensor * output = build_qwen_flash_attention_graph(ctx, 1, 8, 4, 2); + REQUIRE(!graph_is_supported(ctx, output)); + } + { + ggml_tensor * output = build_qwen_flash_attention_graph(ctx, 4, 8, 4, 2, GGML_TYPE_F32, GGML_TYPE_F32); + REQUIRE(!graph_is_supported(ctx, output)); + } + { + ggml_tensor * output = build_qwen_flash_attention_graph(ctx, 4, 8, 4, 2, GGML_TYPE_F32, GGML_TYPE_F16, false); + REQUIRE(!graph_is_supported(ctx, output)); + } + { + ggml_tensor * output = + build_qwen_flash_attention_graph(ctx, 4, 8, 4, 2, GGML_TYPE_F32, GGML_TYPE_F16, true, true); + REQUIRE(!graph_is_supported(ctx, output)); + } + { + ggml_tensor * output = build_qwen_flash_attention_graph(ctx, 4, 8, 4, 2, GGML_TYPE_F32, GGML_TYPE_F16, true, + false, 64, 1.0f / std::sqrt(64.0f)); + REQUIRE(!graph_is_supported(ctx, output)); + } + { + ggml_tensor * output = build_qwen_flash_attention_graph(ctx, 4, 8, 4, 2, GGML_TYPE_F32, GGML_TYPE_F16, true, + false, kQwenFlashHeadSize, 1.0f); + REQUIRE(!graph_is_supported(ctx, output)); + } + { + ggml_tensor * output = + build_qwen_flash_attention_graph(ctx, 4, 8, 4, 2, GGML_TYPE_F32, GGML_TYPE_F16, true, false, + kQwenFlashHeadSize, 1.0f / std::sqrt(128.0f), 1.0f); + REQUIRE(!graph_is_supported(ctx, output)); + } + + ggml_free(ctx); +} + +struct QwenAttentionPostprocessTensors { + ggml_tensor * query_raw = nullptr; + ggml_tensor * key_raw = nullptr; + ggml_tensor * value_raw = nullptr; + ggml_tensor * query_reshape = nullptr; + ggml_tensor * key_reshape = nullptr; + ggml_tensor * value_reshape = nullptr; + ggml_tensor * query_output = nullptr; + ggml_tensor * key_cache = nullptr; + ggml_tensor * value_cache = nullptr; + ggml_tensor * key_output = nullptr; + ggml_tensor * value_output = nullptr; + ggml_tensor * positions = nullptr; + ggml_tensor * key_cache_indices = nullptr; + ggml_tensor * value_cache_indices = nullptr; + ggml_tensor * mask = nullptr; + ggml_tensor * flash_output = nullptr; +}; + +static QwenAttentionPostprocessTensors build_qwen_attention_postprocess_graph(ggml_context * ctx, + int64_t token_count, + int64_t query_head_count, + int64_t key_value_head_count, + int64_t cache_row_count, + float rms_epsilon = 0.000001f, + bool include_inverse_frequencies = true) { + QwenAttentionPostprocessTensors tensors; + const int64_t query_size = query_head_count * kQwenFlashHeadSize; + const int64_t key_value_size = key_value_head_count * kQwenFlashHeadSize; + + ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, kQwenMoeHiddenSize, token_count); + ggml_tensor * query_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_Q4_K, kQwenMoeHiddenSize, query_size); + ggml_tensor * key_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_Q4_K, kQwenMoeHiddenSize, key_value_size); + ggml_tensor * value_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_Q6_K, kQwenMoeHiddenSize, key_value_size); + REQUIRE(input != nullptr); + REQUIRE(query_weight != nullptr); + REQUIRE(key_weight != nullptr); + REQUIRE(value_weight != nullptr); + + tensors.query_raw = ggml_mul_mat(ctx, query_weight, input); + tensors.key_raw = ggml_mul_mat(ctx, key_weight, input); + tensors.value_raw = ggml_mul_mat(ctx, value_weight, input); + REQUIRE(tensors.query_raw != nullptr); + REQUIRE(tensors.key_raw != nullptr); + REQUIRE(tensors.value_raw != nullptr); + + tensors.query_reshape = ggml_reshape_3d(ctx, tensors.query_raw, kQwenFlashHeadSize, query_head_count, token_count); + tensors.key_reshape = ggml_reshape_3d(ctx, tensors.key_raw, kQwenFlashHeadSize, key_value_head_count, token_count); + tensors.value_reshape = + ggml_reshape_3d(ctx, tensors.value_raw, kQwenFlashHeadSize, key_value_head_count, token_count); + REQUIRE(tensors.query_reshape != nullptr); + REQUIRE(tensors.key_reshape != nullptr); + REQUIRE(tensors.value_reshape != nullptr); + + ggml_tensor * query_norm_weight = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, kQwenFlashHeadSize); + ggml_tensor * key_norm_weight = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, kQwenFlashHeadSize); + tensors.positions = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, token_count); + ggml_tensor * inverse_frequencies = + include_inverse_frequencies ? ggml_new_tensor_1d(ctx, GGML_TYPE_F32, kQwenFlashHeadSize / 2) : nullptr; + REQUIRE(query_norm_weight != nullptr); + REQUIRE(key_norm_weight != nullptr); + REQUIRE(tensors.positions != nullptr); + REQUIRE(include_inverse_frequencies == (inverse_frequencies != nullptr)); + + ggml_tensor * query_norm = ggml_rms_norm(ctx, tensors.query_reshape, rms_epsilon); + ggml_tensor * query_mul = ggml_mul(ctx, query_norm, query_norm_weight); + REQUIRE(query_norm != nullptr); + REQUIRE(query_mul != nullptr); + tensors.query_output = include_inverse_frequencies ? + ggml_rope_ext(ctx, query_mul, tensors.positions, inverse_frequencies, kQwenFlashHeadSize, + GGML_ROPE_TYPE_NEOX, 0, 10000.0f, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f) : + ggml_rope(ctx, query_mul, tensors.positions, kQwenFlashHeadSize, GGML_ROPE_TYPE_NEOX); + REQUIRE(tensors.query_output != nullptr); + + ggml_tensor * key_norm = ggml_rms_norm(ctx, tensors.key_reshape, rms_epsilon); + ggml_tensor * key_mul = ggml_mul(ctx, key_norm, key_norm_weight); + REQUIRE(key_norm != nullptr); + REQUIRE(key_mul != nullptr); + ggml_tensor * key_rope = include_inverse_frequencies ? + ggml_rope_ext(ctx, key_mul, tensors.positions, inverse_frequencies, kQwenFlashHeadSize, + GGML_ROPE_TYPE_NEOX, 0, 10000.0f, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f) : + ggml_rope(ctx, key_mul, tensors.positions, kQwenFlashHeadSize, GGML_ROPE_TYPE_NEOX); + REQUIRE(key_rope != nullptr); + + ggml_tensor * key_cache_rows = + ggml_reshape_2d(ctx, key_rope, kQwenFlashHeadSize * key_value_head_count, token_count); + ggml_tensor * value_cache_rows = + ggml_reshape_2d(ctx, tensors.value_reshape, kQwenFlashHeadSize * key_value_head_count, token_count); + REQUIRE(key_cache_rows != nullptr); + REQUIRE(value_cache_rows != nullptr); + + tensors.key_cache = ggml_new_tensor_2d(ctx, GGML_TYPE_F16, key_value_size, cache_row_count); + tensors.value_cache = ggml_new_tensor_2d(ctx, GGML_TYPE_F16, key_value_size, cache_row_count); + tensors.key_cache_indices = ggml_new_tensor_1d(ctx, GGML_TYPE_I64, token_count); + tensors.value_cache_indices = ggml_new_tensor_1d(ctx, GGML_TYPE_I64, token_count); + REQUIRE(tensors.key_cache != nullptr); + REQUIRE(tensors.value_cache != nullptr); + REQUIRE(tensors.key_cache_indices != nullptr); + REQUIRE(tensors.value_cache_indices != nullptr); + + tensors.key_output = ggml_set_rows(ctx, tensors.key_cache, key_cache_rows, tensors.key_cache_indices); + tensors.value_output = ggml_set_rows(ctx, tensors.value_cache, value_cache_rows, tensors.value_cache_indices); + REQUIRE(tensors.key_output != nullptr); + REQUIRE(tensors.value_output != nullptr); + return tensors; +} + +static ggml::hrx::GraphImportResult import_qwen_attention_postprocess_graph( + ggml_context * ctx, + const QwenAttentionPostprocessTensors & tensors) { + ggml_cgraph * graph = ggml_new_graph(ctx); + REQUIRE(graph != nullptr); + ggml_build_forward_expand(graph, tensors.query_output); + ggml_build_forward_expand(graph, tensors.key_output); + ggml_build_forward_expand(graph, tensors.value_output); + + ggml::hrx::GraphImportResult imported = ggml::hrx::import_ggml_graph(*graph); + REQUIRE(imported.valid()); + return imported; +} + +static ggml_tensor * append_qwen_flash_attention_consumer(ggml_context * ctx, + QwenAttentionPostprocessTensors & tensors, + int64_t token_count, + int64_t query_head_count, + int64_t key_value_head_count, + int64_t cache_row_count) { + ggml_tensor * query_layout = + ggml_reshape_3d(ctx, tensors.query_output, kQwenFlashHeadSize, query_head_count, token_count); + ggml_tensor * query_permute = ggml_permute(ctx, query_layout, 0, 2, 1, 3); + ggml_tensor * key_cache_layout = + ggml_reshape_3d(ctx, tensors.key_cache, kQwenFlashHeadSize, key_value_head_count, cache_row_count); + ggml_tensor * key_permute = ggml_permute(ctx, key_cache_layout, 0, 2, 1, 3); + ggml_tensor * value_cache_layout = + ggml_reshape_3d(ctx, tensors.value_cache, kQwenFlashHeadSize, key_value_head_count, cache_row_count); + ggml_tensor * value_permute = ggml_permute(ctx, value_cache_layout, 0, 2, 1, 3); + tensors.mask = ggml_new_tensor_2d(ctx, GGML_TYPE_F16, cache_row_count, token_count); + REQUIRE(query_layout != nullptr); + REQUIRE(query_permute != nullptr); + REQUIRE(key_cache_layout != nullptr); + REQUIRE(key_permute != nullptr); + REQUIRE(value_cache_layout != nullptr); + REQUIRE(value_permute != nullptr); + REQUIRE(tensors.mask != nullptr); + + tensors.flash_output = ggml_flash_attn_ext(ctx, query_permute, key_permute, value_permute, tensors.mask, + 1.0f / std::sqrt(static_cast(kQwenFlashHeadSize)), 0.0f, 0.0f); + REQUIRE(tensors.flash_output != nullptr); + return tensors.flash_output; +} + +static void schedule_qwen_attention_postprocess_command(ggml_context * ctx, + const QwenAttentionPostprocessTensors & tensors, + int64_t token_count, + int64_t query_head_count, + int64_t key_value_head_count, + int64_t cache_row_count, + bool expect_synthetic_inverse_frequencies = false) { + ggml::hrx::GraphImportResult imported = import_qwen_attention_postprocess_graph(ctx, tensors); + + ggml::hrx::DispatchScheduler scheduler; + REQUIRE(scheduler.schedule_graph(imported.graph, test_dispatch_target())); + REQUIRE(scheduler.plan().valid()); + REQUIRE(scheduler.plan().initialization_dispatches.empty()); + REQUIRE(scheduler.plan().dispatches.size() == 4); + + const ggml::hrx::Dispatch & dispatch = scheduler.plan().dispatches.back(); + const std::string kernel_name = kernel_name_for_id(dispatch.kernel.kernel_id); + REQUIRE(kernel_name == "qwen3_moe:qwen3_moe_attention_postprocess_f32_f16"); + REQUIRE(dispatch.kernel.integer_parameters.at("token_count") == token_count); + REQUIRE(dispatch.kernel.integer_parameters.at("cache_row_count") == cache_row_count); + REQUIRE(dispatch.bindings.size() == 12); + require_compile_parameter(dispatch, "qwen3_moe.model.rms_epsilon", "0.000001"); + require_compile_parameter(dispatch, "qwen3_moe.attention.head_size", std::to_string(kQwenFlashHeadSize)); + require_compile_parameter(dispatch, "qwen3_moe.attention.query_size", + std::to_string(query_head_count * kQwenFlashHeadSize)); + require_compile_parameter(dispatch, "qwen3_moe.attention.key_value_size", + std::to_string(key_value_head_count * kQwenFlashHeadSize)); + require_compile_parameter(dispatch, "qwen3_moe.workload.token_capacity", std::to_string(token_count)); + + const ggml::hrx::Value * positions_value = imported.graph.values().find_tensor(tensors.positions); + const ggml::hrx::Value * key_indices_value = imported.graph.values().find_tensor(tensors.key_cache_indices); + const ggml::hrx::Value * value_indices_value = imported.graph.values().find_tensor(tensors.value_cache_indices); + const ggml::hrx::Value * query_raw_value = imported.graph.values().find_tensor(tensors.query_raw); + const ggml::hrx::Value * key_raw_value = imported.graph.values().find_tensor(tensors.key_raw); + const ggml::hrx::Value * value_raw_value = imported.graph.values().find_tensor(tensors.value_raw); + const ggml::hrx::Value * query_output_value = imported.graph.values().find_tensor(tensors.query_output); + const ggml::hrx::Value * key_cache_value = imported.graph.values().find_tensor(tensors.key_cache); + const ggml::hrx::Value * value_cache_value = imported.graph.values().find_tensor(tensors.value_cache); + REQUIRE(positions_value != nullptr); + REQUIRE(key_indices_value != nullptr); + REQUIRE(value_indices_value != nullptr); + REQUIRE(query_raw_value != nullptr); + REQUIRE(key_raw_value != nullptr); + REQUIRE(value_raw_value != nullptr); + REQUIRE(query_output_value != nullptr); + REQUIRE(key_cache_value != nullptr); + REQUIRE(value_cache_value != nullptr); + REQUIRE(dispatch.bindings[0].value == positions_value->id); + REQUIRE(dispatch.bindings[1].value == key_indices_value->id); + REQUIRE(dispatch.bindings[2].value == value_indices_value->id); + REQUIRE(dispatch.bindings[3].value == query_raw_value->id); + REQUIRE(dispatch.bindings[4].value == key_raw_value->id); + REQUIRE(dispatch.bindings[5].value == value_raw_value->id); + REQUIRE(dispatch.bindings[9].value == query_output_value->id); + REQUIRE(dispatch.bindings[10].value == key_cache_value->id); + REQUIRE(dispatch.bindings[11].value == value_cache_value->id); + if (expect_synthetic_inverse_frequencies) { + REQUIRE(scheduler.plan().transients.size() == 1); + REQUIRE(scheduler.plan().constant_initializations.size() == 1); + REQUIRE(dispatch.bindings[8].value == scheduler.plan().transients[0].value); + REQUIRE(scheduler.plan().constant_initializations[0].value == dispatch.bindings[8].value); + REQUIRE(scheduler.plan().constant_initializations[0].data.size() == + static_cast(kQwenFlashHeadSize / 2) * sizeof(float)); + } else { + REQUIRE(scheduler.plan().transients.empty()); + REQUIRE(scheduler.plan().constant_initializations.empty()); + } + + const ggml::hrx::CommandProgram commands = ggml::hrx::build_command_program( + imported.graph, scheduler.plan(), ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); + REQUIRE(commands.valid()); + REQUIRE(commands.initialization_commands.empty()); + REQUIRE(commands.commands.size() == 4); + REQUIRE(command_program_verifies(commands)); + REQUIRE(commands.commands.back().bindings.size() == 12); + REQUIRE(commands.commands.back().bindings[0].name == "positions"); + REQUIRE(commands.commands.back().bindings[1].name == "key_cache_indices"); + REQUIRE(commands.commands.back().bindings[2].name == "value_cache_indices"); + REQUIRE(commands.commands.back().bindings[3].name == "query_input"); + REQUIRE(commands.commands.back().bindings[4].name == "key_input"); + REQUIRE(commands.commands.back().bindings[5].name == "value_input"); + REQUIRE(commands.commands.back().bindings[6].name == "query_norm_weight"); + REQUIRE(commands.commands.back().bindings[7].name == "key_norm_weight"); + REQUIRE(commands.commands.back().bindings[8].name == "inverse_frequencies"); + REQUIRE(commands.commands.back().bindings[9].name == "query_output"); + REQUIRE(commands.commands.back().bindings[10].name == "key_cache"); + REQUIRE(commands.commands.back().bindings[11].name == "value_cache"); + REQUIRE(commands.constant_initializations.size() == (expect_synthetic_inverse_frequencies ? 1 : 0)); +} + +static void run_qwen_attention_postprocess_dispatch_checks() { + ggml_init_params params = {}; + params.mem_size = 8 * 1024 * 1024; + params.no_alloc = true; + ggml_context * ctx = ggml_init(params); + REQUIRE(ctx != nullptr); + + { + constexpr int64_t token_count = 4; + constexpr int64_t query_head_count = 4; + constexpr int64_t key_value_head_count = 2; + constexpr int64_t cache_row_count = 16; + const QwenAttentionPostprocessTensors tensors = build_qwen_attention_postprocess_graph( + ctx, token_count, query_head_count, key_value_head_count, cache_row_count); + schedule_qwen_attention_postprocess_command(ctx, tensors, token_count, query_head_count, key_value_head_count, + cache_row_count); + } + + { + const QwenAttentionPostprocessTensors tensors = + build_qwen_attention_postprocess_graph(ctx, 4, 4, 2, 16, 0.00001f); + ggml::hrx::GraphImportResult imported = import_qwen_attention_postprocess_graph(ctx, tensors); + REQUIRE(!ggml::hrx::DispatchScheduler::can_schedule_graph(imported.graph, test_dispatch_target())); + } + + { + const QwenAttentionPostprocessTensors tensors = + build_qwen_attention_postprocess_graph(ctx, 4, 4, 2, 16, 0.000001f, false); + schedule_qwen_attention_postprocess_command(ctx, tensors, 4, 4, 2, 16, true); + } + + { + constexpr int64_t token_count = 13; + constexpr int64_t query_head_count = 32; + constexpr int64_t key_value_head_count = 4; + constexpr int64_t cache_row_count = 512; + const QwenAttentionPostprocessTensors tensors = build_qwen_attention_postprocess_graph( + ctx, token_count, query_head_count, key_value_head_count, cache_row_count, 0.000001f, false); + schedule_qwen_attention_postprocess_command(ctx, tensors, token_count, query_head_count, key_value_head_count, + cache_row_count, true); + } + + { + constexpr int64_t token_count = 13; + constexpr int64_t query_head_count = 32; + constexpr int64_t key_value_head_count = 4; + constexpr int64_t cache_row_count = 512; + + const QwenAttentionPostprocessTensors first = build_qwen_attention_postprocess_graph( + ctx, token_count, query_head_count, key_value_head_count, cache_row_count, 0.000001f, false); + const QwenAttentionPostprocessTensors second = build_qwen_attention_postprocess_graph( + ctx, token_count, query_head_count, key_value_head_count, cache_row_count, 0.000001f, false); + + ggml_cgraph * graph = ggml_new_graph(ctx); + REQUIRE(graph != nullptr); + ggml_build_forward_expand(graph, first.query_output); + ggml_build_forward_expand(graph, first.key_output); + ggml_build_forward_expand(graph, first.value_output); + ggml_build_forward_expand(graph, second.query_output); + ggml_build_forward_expand(graph, second.key_output); + ggml_build_forward_expand(graph, second.value_output); + + ggml::hrx::GraphImportResult imported = ggml::hrx::import_ggml_graph(*graph); + REQUIRE(imported.valid()); + + ggml::hrx::DispatchScheduler scheduler; + REQUIRE(scheduler.schedule_graph(imported.graph, test_dispatch_target())); + REQUIRE(scheduler.plan().valid()); + REQUIRE(scheduler.plan().dispatches.size() == 8); + + const ggml::hrx::CommandProgram commands = ggml::hrx::build_command_program( + imported.graph, scheduler.plan(), ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); + REQUIRE(commands.valid()); + REQUIRE(commands.commands.size() == 8); + REQUIRE(command_program_verifies(commands)); + } + + { + constexpr int64_t token_count = 4; + constexpr int64_t query_head_count = 4; + constexpr int64_t key_value_head_count = 2; + constexpr int64_t cache_row_count = 16; + QwenAttentionPostprocessTensors tensors = build_qwen_attention_postprocess_graph( + ctx, token_count, query_head_count, key_value_head_count, cache_row_count); + ggml_tensor * flash_output = append_qwen_flash_attention_consumer(ctx, tensors, token_count, query_head_count, + key_value_head_count, cache_row_count); + + ggml_cgraph * graph = ggml_new_graph(ctx); + REQUIRE(graph != nullptr); + ggml_build_forward_expand(graph, tensors.key_output); + ggml_build_forward_expand(graph, tensors.value_output); + ggml_build_forward_expand(graph, flash_output); + + ggml::hrx::GraphImportResult imported = ggml::hrx::import_ggml_graph(*graph); + REQUIRE(imported.valid()); + + std::vector covered_nodes(imported.graph.nodes().size(), false); + ggml::hrx::CommandPlan plan; + ggml::hrx::DispatchMatch match; + REQUIRE(match_dispatch_at_index(imported.graph, plan, covered_nodes, + producer_index_for_tensor(imported.graph, tensors.query_reshape), match)); + append_match_to_plan(plan, match, covered_nodes); + REQUIRE(plan.valid()); + REQUIRE(plan.initialization_dispatches.size() == 2); + + const ggml::hrx::Dispatch & context_capture = plan.initialization_dispatches[0]; + const ggml::hrx::Dispatch & metadata = plan.initialization_dispatches[1]; + REQUIRE(kernel_name_for_id(context_capture.kernel.kernel_id) == + "qwen3_moe:qwen_attention_context_base_capture"); + REQUIRE(kernel_name_for_id(metadata.kernel.kernel_id) == "qwen3_moe:qwen_attention_metadata"); + REQUIRE(context_capture.bindings.size() == 2); + REQUIRE(metadata.bindings.size() == 5); + REQUIRE(context_capture.bindings[1].value == metadata.bindings[0].value); + REQUIRE(metadata.kernel.integer_parameters.at("token_count") == token_count); + REQUIRE(metadata.kernel.integer_parameters.at("context_capacity") == cache_row_count); + + const ggml::hrx::Value * positions_value = imported.graph.values().find_tensor(tensors.positions); + const ggml::hrx::Value * key_indices_value = imported.graph.values().find_tensor(tensors.key_cache_indices); + const ggml::hrx::Value * value_indices_value = imported.graph.values().find_tensor(tensors.value_cache_indices); + const ggml::hrx::Value * mask_value = imported.graph.values().find_tensor(tensors.mask); + REQUIRE(positions_value != nullptr); + REQUIRE(key_indices_value != nullptr); + REQUIRE(value_indices_value != nullptr); + REQUIRE(mask_value != nullptr); + REQUIRE(context_capture.bindings[0].value == positions_value->id); + REQUIRE(metadata.bindings[1].value == positions_value->id); + REQUIRE(metadata.bindings[2].value == key_indices_value->id); + REQUIRE(metadata.bindings[3].value == value_indices_value->id); + REQUIRE(metadata.bindings[4].value == mask_value->id); + + const ggml::hrx::CommandProgram commands = + ggml::hrx::build_command_program(imported.graph, plan, ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); + REQUIRE(commands.valid()); + REQUIRE(command_program_verifies(commands)); + REQUIRE(commands.initialization_commands.size() == 2); + REQUIRE(commands.initialization_commands[0].bindings[0].name == "positions"); + REQUIRE(commands.initialization_commands[0].bindings[1].name == "control"); + REQUIRE(commands.initialization_commands[1].bindings[0].name == "control"); + REQUIRE(commands.initialization_commands[1].bindings[4].name == "attention_mask"); + REQUIRE(ggml::hrx::find_transient_allocation(commands.transients, context_capture.bindings[1].value) != + nullptr); + } + + ggml_free(ctx); +} + +static void run_qwen_matmul_dispatch_checks() { + ggml_init_params params = {}; + params.mem_size = 2 * 1024 * 1024; + params.no_alloc = true; + ggml_context * ctx = ggml_init(params); + REQUIRE(ctx != nullptr); + + { + ggml_tensor * weight = ggml_new_tensor_2d(ctx, GGML_TYPE_Q4_K, 2048, 128); + ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, 4); + REQUIRE(weight != nullptr); + REQUIRE(input != nullptr); + ggml_tensor * output = ggml_mul_mat(ctx, weight, input); + REQUIRE(output != nullptr); + schedule_single_matmul_command(ctx, output, "qwen3_moe:qwen3_moe_dense_linear_q4k_f16_wmma", 4, 2048, 128); + } + + { + ggml_tensor * weight = ggml_new_tensor_2d(ctx, GGML_TYPE_Q4_K, 2048, 128); + ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, 1); + REQUIRE(weight != nullptr); + REQUIRE(input != nullptr); + ggml_tensor * output = ggml_mul_mat(ctx, weight, input); + REQUIRE(output != nullptr); + REQUIRE(!matmul_graph_is_supported(ctx, output)); + } + + { + ggml_tensor * weight = ggml_new_tensor_2d(ctx, GGML_TYPE_Q6_K, 2048, 128); + ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, 2); + REQUIRE(weight != nullptr); + REQUIRE(input != nullptr); + ggml_tensor * output = ggml_mul_mat(ctx, weight, input); + REQUIRE(output != nullptr); + schedule_single_matmul_command(ctx, output, "qwen3_moe:qwen3_moe_dense_linear_q6k_f16_wmma", 2, 2048, 128); + } + + { + ggml_tensor * weight = ggml_new_tensor_2d(ctx, GGML_TYPE_Q6_K, 2048, 128); + ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, 1); + REQUIRE(weight != nullptr); + REQUIRE(input != nullptr); + ggml_tensor * output = ggml_mul_mat(ctx, weight, input); + REQUIRE(output != nullptr); + REQUIRE(!matmul_graph_is_supported(ctx, output)); + } + + { + ggml_tensor * weight = ggml_new_tensor_2d(ctx, GGML_TYPE_Q6_K, 2048, 151936); + ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, 1); + REQUIRE(weight != nullptr); + REQUIRE(input != nullptr); + ggml_tensor * output = ggml_mul_mat(ctx, weight, input); + REQUIRE(output != nullptr); + schedule_single_matmul_command(ctx, output, "qwen3_moe:qwen3_moe_dense_linear_q6k_f16_wmma", 1, 2048, 151936); + } + + { + ggml_tensor * weight = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, 128); + ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, 4); + REQUIRE(weight != nullptr); + REQUIRE(input != nullptr); + ggml_tensor * output = ggml_mul_mat(ctx, weight, input); + REQUIRE(output != nullptr); + schedule_single_matmul_command(ctx, output, "qwen3_moe:qwen3_moe_router_projection_f32_four_row_wave32", 4, + 2048, 128); + } + + { + ggml_tensor * weight = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, 128); + ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, 1); + REQUIRE(weight != nullptr); + REQUIRE(input != nullptr); + ggml_tensor * output = ggml_mul_mat(ctx, weight, input); + REQUIRE(output != nullptr); + schedule_single_matmul_command(ctx, output, "qwen3_moe:qwen3_moe_router_projection_f32_four_row_wave32", 1, + 2048, 128); + } + + { + ggml_tensor * weight = ggml_new_tensor_2d(ctx, GGML_TYPE_F16, 2048, 128); + ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, 4); + REQUIRE(weight != nullptr); + REQUIRE(input != nullptr); + ggml_tensor * output = ggml_mul_mat(ctx, weight, input); + REQUIRE(output != nullptr); + REQUIRE(!matmul_graph_is_supported(ctx, output)); + } + + { + ggml_tensor * weight = ggml_new_tensor_2d(ctx, GGML_TYPE_Q4_K, 2048, 128); + ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F16, 2048, 4); + REQUIRE(weight != nullptr); + REQUIRE(input != nullptr); + ggml_tensor * output = ggml_mul_mat(ctx, weight, input); + REQUIRE(output != nullptr); + REQUIRE(!matmul_graph_is_supported(ctx, output)); + } + + { + ggml_tensor * weight = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, 2048, 128, 2); + ggml_tensor * input = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, 2048, 4, 2); + REQUIRE(weight != nullptr); + REQUIRE(input != nullptr); + ggml_tensor * output = ggml_mul_mat(ctx, weight, input); + REQUIRE(output != nullptr); + REQUIRE(!matmul_graph_is_supported(ctx, output)); + } + + { + ggml_tensor * weight = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 1024, 128); + ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 1024, 4); + REQUIRE(weight != nullptr); + REQUIRE(input != nullptr); + ggml_tensor * output = ggml_mul_mat(ctx, weight, input); + REQUIRE(output != nullptr); + REQUIRE(!matmul_graph_is_supported(ctx, output)); + } + + ggml_free(ctx); +} + +static void schedule_qwen_router_top8_command(ggml_context * ctx, + ggml_tensor * output, + ggml_tensor * route_ids, + int64_t expected_expert_count = kQwenRouterExpertCount, + int64_t expected_route_count = kQwenRouterRouteCount) { + ggml_cgraph * graph = ggml_new_graph(ctx); + REQUIRE(graph != nullptr); + ggml_build_forward_expand(graph, output); + + ggml::hrx::GraphImportResult imported = ggml::hrx::import_ggml_graph(*graph); + REQUIRE(imported.valid()); + REQUIRE(imported.graph.nodes().size() == 10); + + const ggml::hrx::Value * route_ids_value = imported.graph.values().find_tensor(route_ids); + const ggml::hrx::Value * output_value = imported.graph.values().find_tensor(output); + REQUIRE(route_ids_value != nullptr); + REQUIRE(output_value != nullptr); + REQUIRE(route_ids_value->kind == ggml::hrx::ValueKind::Transient); + REQUIRE(output_value->kind == ggml::hrx::ValueKind::External); + + const ggml::hrx::GraphNode * softmax_node = nullptr; + const ggml::hrx::GraphNode * get_rows_node = nullptr; + for (const ggml::hrx::GraphNode & node : imported.graph.nodes()) { + if (node.op == GGML_OP_SOFT_MAX) { + softmax_node = &node; + } else if (node.op == GGML_OP_GET_ROWS) { + get_rows_node = &node; + } + } + REQUIRE(softmax_node != nullptr); + REQUIRE(get_rows_node != nullptr); + const ggml::hrx::GraphNode * probs_reshape = + ggml::hrx::find_single_layout_alias_consumer_with_op(imported.graph, softmax_node->output, GGML_OP_RESHAPE); + REQUIRE(probs_reshape != nullptr); + REQUIRE(ggml::hrx::is_layout_alias_node(imported.graph, *probs_reshape)); + REQUIRE(ggml::hrx::find_single_consumer_with_op_through_layout_aliases(imported.graph, softmax_node->output, + GGML_OP_GET_ROWS) == get_rows_node); + + ggml::hrx::DispatchScheduler scheduler; + REQUIRE(scheduler.schedule_graph(imported.graph, test_dispatch_target())); + REQUIRE(scheduler.plan().valid()); + + const int64_t token_count = output->ne[2]; + const size_t route_id_length = static_cast(token_count) * route_ids->nb[1]; + const size_t expert_table_bytes = qwen_expert_table_size(token_count, expected_expert_count); + const size_t partition_table_bytes = + qwen_partition_table_size(token_count, expected_route_count, expected_expert_count); + const bool uses_fused_prefill_expert_table_partition = + token_count == 512 && expected_route_count == 8 && + route_ids->nb[1] / sizeof(int32_t) == static_cast(expected_route_count) && + expected_expert_count == 128; + REQUIRE(scheduler.plan().dispatches.size() == (uses_fused_prefill_expert_table_partition ? 2 : 3)); + REQUIRE(scheduler.plan().transients.size() == 2); + REQUIRE(scheduler.plan().constant_initializations.empty()); + REQUIRE(scheduler.plan().completion_counter_requests.size() == (uses_fused_prefill_expert_table_partition ? 1 : 0)); + + const ggml::hrx::CommandPlanTransient & expert_table_transient = scheduler.plan().transients[0]; + const ggml::hrx::CommandPlanTransient & partition_table_transient = scheduler.plan().transients[1]; + REQUIRE(expert_table_transient.value.value == static_cast(imported.graph.values().size())); + REQUIRE(expert_table_transient.name == "qwen.router.expert_table"); + REQUIRE(expert_table_transient.size == expert_table_bytes); + REQUIRE(partition_table_transient.value.value == expert_table_transient.value.value + 1); + REQUIRE(partition_table_transient.name == "qwen.router.partition_table"); + REQUIRE(partition_table_transient.size == partition_table_bytes); + if (uses_fused_prefill_expert_table_partition) { + const ggml::hrx::CommandPlanCompletionCounterRequest & completion_counter_request = + scheduler.plan().completion_counter_requests[0]; + REQUIRE(completion_counter_request.value.value == partition_table_transient.value.value + 1); + REQUIRE(completion_counter_request.name == "qwen.router.prefill_expert_table_partition_completion_counter"); + REQUIRE(completion_counter_request.count == 1); + } + + const ggml::hrx::Dispatch & dispatch = scheduler.plan().dispatches[0]; + const std::string kernel_name = kernel_name_for_id(dispatch.kernel.kernel_id); + REQUIRE(kernel_name == "qwen3_moe:qwen3_moe_router_top8_f32"); + REQUIRE(dispatch.kernel.integer_parameters.at("token_count") == token_count); + REQUIRE(dispatch.kernel.integer_parameters.at("route_id_stride") == route_ids->nb[1] / sizeof(int32_t)); + REQUIRE(dispatch.bindings.size() == 3); + REQUIRE(dispatch.bindings[1].value == route_ids_value->id); + REQUIRE(dispatch.bindings[1].length == route_id_length); + REQUIRE(dispatch.bindings[2].value == output_value->id); + require_compile_parameter(dispatch, "qwen3_moe.router.expert_count", std::to_string(expected_expert_count)); + require_compile_parameter(dispatch, "qwen3_moe.router.route_count", std::to_string(expected_route_count)); + require_compile_parameter(dispatch, "qwen3_moe.workload.token_capacity", std::to_string(token_count)); + + if (uses_fused_prefill_expert_table_partition) { + const ggml::hrx::CommandPlanCompletionCounterRequest & completion_counter_request = + scheduler.plan().completion_counter_requests[0]; + const ggml::hrx::Dispatch & expert_table_partition_dispatch = scheduler.plan().dispatches[1]; + REQUIRE(kernel_name_for_id(expert_table_partition_dispatch.kernel.kernel_id) == + "qwen3_moe:qwen3_moe_build_expert_table_partition_prefill_512"); + REQUIRE(expert_table_partition_dispatch.kernel.integer_parameters.at("token_count") == token_count); + REQUIRE(expert_table_partition_dispatch.kernel.integer_parameters.at("route_count") == expected_route_count); + REQUIRE(expert_table_partition_dispatch.kernel.integer_parameters.at("route_stride") == + route_ids->nb[1] / sizeof(int32_t)); + REQUIRE(expert_table_partition_dispatch.kernel.integer_parameters.at("expert_count") == expected_expert_count); + REQUIRE(expert_table_partition_dispatch.bindings.size() == 4); + REQUIRE(expert_table_partition_dispatch.bindings[0].value == route_ids_value->id); + REQUIRE(expert_table_partition_dispatch.bindings[0].length == route_id_length); + REQUIRE(expert_table_partition_dispatch.bindings[1].value == expert_table_transient.value); + REQUIRE(expert_table_partition_dispatch.bindings[1].length == expert_table_bytes); + REQUIRE(expert_table_partition_dispatch.bindings[2].value == partition_table_transient.value); + REQUIRE(expert_table_partition_dispatch.bindings[2].length == partition_table_bytes); + REQUIRE(expert_table_partition_dispatch.bindings[3].value == completion_counter_request.value); + REQUIRE(expert_table_partition_dispatch.bindings[3].length == sizeof(int32_t)); + } else { + const ggml::hrx::Dispatch & expert_table_dispatch = scheduler.plan().dispatches[1]; + REQUIRE(kernel_name_for_id(expert_table_dispatch.kernel.kernel_id) == "qwen3_moe:qwen3_moe_build_expert_table"); + REQUIRE(expert_table_dispatch.kernel.integer_parameters.at("token_count") == token_count); + REQUIRE(expert_table_dispatch.kernel.integer_parameters.at("route_count") == expected_route_count); + REQUIRE(expert_table_dispatch.kernel.integer_parameters.at("route_stride") == + route_ids->nb[1] / sizeof(int32_t)); + REQUIRE(expert_table_dispatch.kernel.integer_parameters.at("expert_count") == expected_expert_count); + REQUIRE(expert_table_dispatch.bindings.size() == 2); + REQUIRE(expert_table_dispatch.bindings[0].value == route_ids_value->id); + REQUIRE(expert_table_dispatch.bindings[0].length == route_id_length); + REQUIRE(expert_table_dispatch.bindings[1].value == expert_table_transient.value); + REQUIRE(expert_table_dispatch.bindings[1].length == expert_table_bytes); + require_compile_parameter(expert_table_dispatch, "qwen3_moe.routed_gate_up.expert_count", + std::to_string(expected_expert_count)); + require_compile_parameter(expert_table_dispatch, "qwen3_moe.routed_gate_up.route_count", + std::to_string(expected_route_count)); + require_compile_parameter(expert_table_dispatch, "qwen3_moe.workload.token_capacity", + std::to_string(token_count)); + + const ggml::hrx::Dispatch & partition_table_dispatch = scheduler.plan().dispatches[2]; + REQUIRE(kernel_name_for_id(partition_table_dispatch.kernel.kernel_id) == + "qwen3_moe:qwen3_moe_build_expert_partition_table"); + REQUIRE(partition_table_dispatch.kernel.integer_parameters.at("token_count") == token_count); + REQUIRE(partition_table_dispatch.kernel.integer_parameters.at("route_count") == expected_route_count); + REQUIRE(partition_table_dispatch.kernel.integer_parameters.at("expert_count") == expected_expert_count); + REQUIRE(partition_table_dispatch.bindings.size() == 2); + REQUIRE(partition_table_dispatch.bindings[0].value == expert_table_transient.value); + REQUIRE(partition_table_dispatch.bindings[0].length == expert_table_bytes); + REQUIRE(partition_table_dispatch.bindings[1].value == partition_table_transient.value); + REQUIRE(partition_table_dispatch.bindings[1].length == partition_table_bytes); + require_compile_parameter(partition_table_dispatch, "qwen3_moe.routed_gate_up.expert_count", + std::to_string(expected_expert_count)); + require_compile_parameter(partition_table_dispatch, "qwen3_moe.routed_gate_up.route_count", + std::to_string(expected_route_count)); + require_compile_parameter(partition_table_dispatch, "qwen3_moe.workload.token_capacity", + std::to_string(token_count)); + } + + const ggml::hrx::CommandProgram commands = ggml::hrx::build_command_program( + imported.graph, scheduler.plan(), ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); + REQUIRE(commands.valid()); + REQUIRE(commands.commands.size() == (uses_fused_prefill_expert_table_partition ? 2 : 3)); + REQUIRE(commands.constant_initializations.empty()); + REQUIRE(commands.completion_counters.count == (uses_fused_prefill_expert_table_partition ? 1 : 0)); + REQUIRE(commands.completion_counters.byte_count == + (uses_fused_prefill_expert_table_partition ? sizeof(int32_t) : 0)); + REQUIRE(command_program_verifies(commands)); + REQUIRE(commands.commands[0].bindings.size() == 3); + REQUIRE(commands.commands[0].bindings[0].name == "logits"); + REQUIRE(commands.commands[0].bindings[1].name == "route_ids"); + REQUIRE(commands.commands[0].bindings[1].origin == ggml::hrx::CommandBindingOrigin::Transient); + REQUIRE(commands.commands[0].bindings[1].value == route_ids_value->storage_root); + REQUIRE(commands.commands[0].bindings[1].offset == route_ids_value->storage_offset); + REQUIRE(commands.commands[0].bindings[1].length == dispatch.bindings[1].length); + REQUIRE(commands.commands[0].bindings[2].name == "route_weights"); + REQUIRE(commands.commands[1].dependencies.size() == 1); + REQUIRE(commands.commands[1].dependencies[0] == 0); + REQUIRE(commands.commands[1].bindings.size() == (uses_fused_prefill_expert_table_partition ? 4 : 2)); + REQUIRE(commands.commands[1].bindings[0].name == "route_ids"); + REQUIRE(commands.commands[1].bindings[0].origin == ggml::hrx::CommandBindingOrigin::Transient); + REQUIRE(commands.commands[1].bindings[0].value == route_ids_value->storage_root); + REQUIRE(commands.commands[1].bindings[0].offset == route_ids_value->storage_offset); + REQUIRE(commands.commands[1].bindings[1].name == "expert_table"); + REQUIRE(commands.commands[1].bindings[1].origin == ggml::hrx::CommandBindingOrigin::Transient); + REQUIRE(commands.commands[1].bindings[1].length == expert_table_bytes); + if (uses_fused_prefill_expert_table_partition) { + REQUIRE(commands.commands[1].bindings[2].name == "partition_table"); + REQUIRE(commands.commands[1].bindings[2].origin == ggml::hrx::CommandBindingOrigin::Transient); + REQUIRE(commands.commands[1].bindings[2].length == partition_table_bytes); + REQUIRE(commands.commands[1].bindings[3].name == "completion_counter"); + REQUIRE(commands.commands[1].bindings[3].origin == ggml::hrx::CommandBindingOrigin::Transient); + REQUIRE(commands.commands[1].bindings[3].length == sizeof(int32_t)); + } else { + REQUIRE(commands.commands[2].dependencies.size() == 1); + REQUIRE(commands.commands[2].dependencies[0] == 1); + REQUIRE(commands.commands[2].bindings.size() == 2); + REQUIRE(commands.commands[2].bindings[0].name == "expert_table"); + REQUIRE(commands.commands[2].bindings[0].origin == ggml::hrx::CommandBindingOrigin::Transient); + REQUIRE(commands.commands[2].bindings[1].name == "partition_table"); + REQUIRE(commands.commands[2].bindings[1].origin == ggml::hrx::CommandBindingOrigin::Transient); + REQUIRE(commands.commands[2].bindings[1].length == partition_table_bytes); + } + const ggml::hrx::TransientAllocation * route_ids_allocation = + ggml::hrx::find_transient_allocation(commands.transients, route_ids_value->storage_root); + REQUIRE(route_ids_allocation != nullptr); + REQUIRE(route_ids_allocation->size == dispatch.bindings[1].length); + const ggml::hrx::TransientAllocation * expert_table_allocation = + ggml::hrx::find_transient_allocation(commands.transients, expert_table_transient.value); + REQUIRE(expert_table_allocation != nullptr); + REQUIRE(expert_table_allocation->size == expert_table_bytes); + const ggml::hrx::TransientAllocation * partition_table_allocation = + ggml::hrx::find_transient_allocation(commands.transients, partition_table_transient.value); + REQUIRE(partition_table_allocation != nullptr); + REQUIRE(partition_table_allocation->size == partition_table_bytes); + if (uses_fused_prefill_expert_table_partition) { + const ggml::hrx::CommandPlanCompletionCounterRequest & completion_counter_request = + scheduler.plan().completion_counter_requests[0]; + const ggml::hrx::TransientAllocation * completion_counter_allocation = + ggml::hrx::find_transient_allocation(commands.transients, completion_counter_request.value); + REQUIRE(completion_counter_allocation != nullptr); + REQUIRE(completion_counter_allocation->size == sizeof(int32_t)); + REQUIRE(completion_counter_allocation->arena_offset == commands.completion_counters.arena_offset); + } +} + +static void schedule_manual_qwen_router_top8_command(ggml::hrx::Graph & graph, + ggml::hrx::ValueId route_ids, + int64_t token_count, + int64_t expert_count, + int64_t route_count) { + ggml::hrx::DispatchScheduler scheduler; + REQUIRE(scheduler.schedule_graph(graph, test_dispatch_target())); + REQUIRE(scheduler.plan().valid()); + REQUIRE(scheduler.plan().dispatches.size() == 3); + REQUIRE(scheduler.plan().transients.size() == 2); + + const size_t route_id_length = static_cast(token_count * route_count) * sizeof(int32_t); + const size_t expert_table_bytes = qwen_expert_table_size(token_count, expert_count); + const size_t partition_table_bytes = qwen_partition_table_size(token_count, route_count, expert_count); + + const ggml::hrx::Dispatch & dispatch = scheduler.plan().dispatches[0]; + REQUIRE(kernel_name_for_id(dispatch.kernel.kernel_id) == "qwen3_moe:qwen3_moe_router_top8_f32"); + REQUIRE(dispatch.kernel.integer_parameters.at("token_count") == token_count); + REQUIRE(dispatch.kernel.integer_parameters.at("route_id_stride") == route_count); + REQUIRE(dispatch.bindings.size() == 3); + REQUIRE(dispatch.bindings[1].value == route_ids); + REQUIRE(dispatch.bindings[1].length == route_id_length); + require_compile_parameter(dispatch, "qwen3_moe.router.expert_count", std::to_string(expert_count)); + require_compile_parameter(dispatch, "qwen3_moe.router.route_count", std::to_string(route_count)); + + const ggml::hrx::CommandPlanTransient & expert_table_transient = scheduler.plan().transients[0]; + const ggml::hrx::CommandPlanTransient & partition_table_transient = scheduler.plan().transients[1]; + REQUIRE(expert_table_transient.size == expert_table_bytes); + REQUIRE(partition_table_transient.size == partition_table_bytes); + + const ggml::hrx::CommandPlanMoeRoutingBundle * bundle = + scheduler.plan().metadata.find_moe_routing_bundle(route_ids); + REQUIRE(bundle != nullptr); + REQUIRE(bundle->token_count == token_count); + REQUIRE(bundle->route_count == route_count); + REQUIRE(bundle->route_stride == route_count); + REQUIRE(bundle->expert_count == expert_count); + REQUIRE(bundle->expert_table_byte_count == expert_table_bytes); + REQUIRE(bundle->partition_table_byte_count == partition_table_bytes); + + const ggml::hrx::CommandProgram commands = + ggml::hrx::build_command_program(graph, scheduler.plan(), ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); + REQUIRE(commands.valid()); + REQUIRE(commands.commands.size() == 3); + REQUIRE(command_program_verifies(commands)); + REQUIRE(commands.commands[0].bindings[1].value == route_ids); + REQUIRE(commands.commands[0].bindings[1].length == route_id_length); +} + +struct QwenRoutedGateUpTensors { + ggml_tensor * route_ids = nullptr; + ggml_tensor * route_weights = nullptr; + ggml_tensor * gate = nullptr; + ggml_tensor * up = nullptr; + ggml_tensor * glu = nullptr; + ggml_tensor * output = nullptr; + ggml_tensor * weighted = nullptr; + ggml_tensor * hidden_state = nullptr; + ggml_tensor * residual = nullptr; + ggml_tensor * next_rms = nullptr; + ggml_tensor * next_output = nullptr; + ggml_tensor * hidden_use = nullptr; + std::vector route_views; +}; + +static QwenRoutedGateUpTensors build_qwen_routed_gate_up_graph(ggml_context * ctx, + int64_t token_count, + ggml_glu_op glu_op = GGML_GLU_OP_SWIGLU, + ggml_type up_weight_type = GGML_TYPE_Q4_K, + bool include_down = true, + ggml_type down_weight_type = GGML_TYPE_Q6_K) { + QwenRoutedGateUpTensors tensors; + ggml_tensor * logits = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, kQwenRouterExpertCount, token_count); + REQUIRE(logits != nullptr); + tensors.route_weights = build_qwen_router_top8_graph(ctx, logits, &tensors.route_ids); + REQUIRE(tensors.route_weights != nullptr); + REQUIRE(tensors.route_ids != nullptr); + + ggml_tensor * input = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, kQwenMoeHiddenSize, 1, token_count); + ggml_tensor * gate_weight = + ggml_new_tensor_3d(ctx, GGML_TYPE_Q4_K, kQwenMoeHiddenSize, kQwenMoeIntermediateSize, kQwenRouterExpertCount); + ggml_tensor * up_weight = + ggml_new_tensor_3d(ctx, up_weight_type, kQwenMoeHiddenSize, kQwenMoeIntermediateSize, kQwenRouterExpertCount); + REQUIRE(input != nullptr); + REQUIRE(gate_weight != nullptr); + REQUIRE(up_weight != nullptr); + + tensors.gate = ggml_mul_mat_id(ctx, gate_weight, input, tensors.route_ids); + tensors.up = ggml_mul_mat_id(ctx, up_weight, input, tensors.route_ids); + REQUIRE(tensors.gate != nullptr); + REQUIRE(tensors.up != nullptr); + tensors.glu = ggml_glu_split(ctx, tensors.gate, tensors.up, glu_op); + REQUIRE(tensors.glu != nullptr); + + if (include_down) { + ggml_tensor * down_weight = ggml_new_tensor_3d(ctx, down_weight_type, kQwenMoeIntermediateSize, + kQwenMoeHiddenSize, kQwenRouterExpertCount); + REQUIRE(down_weight != nullptr); + tensors.output = ggml_mul_mat_id(ctx, down_weight, tensors.glu, tensors.route_ids); + REQUIRE(tensors.output != nullptr); + } else { + tensors.output = tensors.glu; + } + return tensors; +} + +static void append_qwen_weighted_reduce_tail(ggml_context * ctx, + QwenRoutedGateUpTensors & tensors, + bool include_next_rmsnorm = false) { + REQUIRE(tensors.output != nullptr); + REQUIRE(tensors.route_weights != nullptr); + tensors.weighted = ggml_mul(ctx, tensors.output, tensors.route_weights); + REQUIRE(tensors.weighted != nullptr); + tensors.route_views.clear(); + for (int64_t route = 0; route < kQwenRouterRouteCount; ++route) { + ggml_tensor * view = + ggml_view_2d(ctx, tensors.weighted, kQwenMoeHiddenSize, tensors.weighted->ne[2], tensors.weighted->nb[2], + static_cast(route) * tensors.weighted->nb[1]); + REQUIRE(view != nullptr); + tensors.route_views.push_back(view); + } + + ggml_tensor * reduced = tensors.route_views.front(); + for (size_t i = 1; i < tensors.route_views.size(); ++i) { + reduced = ggml_add(ctx, reduced, tensors.route_views[i]); + REQUIRE(reduced != nullptr); + } + + tensors.hidden_state = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, kQwenMoeHiddenSize, tensors.output->ne[2]); + REQUIRE(tensors.hidden_state != nullptr); + tensors.residual = ggml_add(ctx, tensors.hidden_state, reduced); + REQUIRE(tensors.residual != nullptr); + + if (include_next_rmsnorm) { + ggml_tensor * next_norm_weight = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, kQwenMoeHiddenSize); + REQUIRE(next_norm_weight != nullptr); + tensors.next_rms = ggml_rms_norm(ctx, tensors.residual, 0.000001f); + REQUIRE(tensors.next_rms != nullptr); + tensors.next_output = ggml_mul(ctx, tensors.next_rms, next_norm_weight); + REQUIRE(tensors.next_output != nullptr); + } +} + +static ggml::hrx::GraphImportResult import_qwen_routed_gate_up_graph(ggml_context * ctx, + const QwenRoutedGateUpTensors & tensors) { + ggml_cgraph * graph = ggml_new_graph(ctx); + REQUIRE(graph != nullptr); + ggml_build_forward_expand(graph, tensors.route_weights); + ggml_build_forward_expand(graph, tensors.output); + if (tensors.residual != nullptr) { + ggml_build_forward_expand(graph, tensors.residual); + } + if (tensors.next_output != nullptr) { + ggml_build_forward_expand(graph, tensors.next_output); + } + if (tensors.hidden_use != nullptr) { + ggml_build_forward_expand(graph, tensors.hidden_use); + } + + ggml::hrx::GraphImportResult imported = ggml::hrx::import_ggml_graph(*graph); + REQUIRE(imported.valid()); + return imported; +} + +static bool match_dispatch_at_index(const ggml::hrx::Graph & graph, + const ggml::hrx::CommandPlan & plan, + const std::vector & covered_nodes, + size_t node_index, + ggml::hrx::DispatchMatch & match) { + REQUIRE(node_index < graph.nodes().size()); + const ggml::hrx::DispatchMatchContext context = { + graph, + &graph.nodes()[node_index], + node_index, + covered_nodes, + plan, + ggml::hrx::ValueId(static_cast(graph.values().size() + plan.transients.size() + + plan.completion_counter_requests.size())), + }; + return test_dispatch_registry().match(context, match); +} + +static void append_match_to_plan(ggml::hrx::CommandPlan & plan, + ggml::hrx::DispatchMatch & match, + std::vector & covered_nodes, + ggml::hrx::Graph * graph) { + if (graph != nullptr) { + for (const ggml::hrx::DispatchValueAliasRequest & alias : match.value_aliases) { + ggml::hrx::Status status = graph->values().alias_storage(alias.target_value, alias.source_value); + REQUIRE(status.success()); + } + } + for (ggml::hrx::Dispatch & dispatch : match.initialization_dispatches) { + plan.initialization_dispatches.push_back(std::move(dispatch)); + } + for (ggml::hrx::Dispatch & dispatch : match.dispatches) { + plan.dispatches.push_back(std::move(dispatch)); + } + for (ggml::hrx::CommandPlanTransient & transient : match.transients) { + plan.transients.push_back(std::move(transient)); + } + for (ggml::hrx::CommandPlanConstantInitialization & initialization : match.constant_initializations) { + plan.constant_initializations.push_back(std::move(initialization)); + } + for (ggml::hrx::CommandPlanCompletionCounterRequest & request : match.completion_counter_requests) { + plan.completion_counter_requests.push_back(std::move(request)); + } + REQUIRE(plan.metadata.append(std::move(match.metadata), plan.status)); + for (const size_t covered_node : match.covered_nodes) { + REQUIRE(covered_node < covered_nodes.size()); + REQUIRE(!covered_nodes[covered_node]); + covered_nodes[covered_node] = true; + } +} + +static ggml::hrx::CommandPlan build_qwen_router_plan_for_graph(ggml::hrx::Graph & graph, + std::vector & covered_nodes) { + ggml::hrx::CommandPlan plan; + size_t softmax_index = graph.nodes().size(); + for (size_t i = 0; i < graph.nodes().size(); ++i) { + if (graph.nodes()[i].op == GGML_OP_SOFT_MAX) { + softmax_index = i; + break; + } + } + REQUIRE(softmax_index < graph.nodes().size()); + ggml::hrx::DispatchMatch router_match; + REQUIRE(match_dispatch_at_index(graph, plan, covered_nodes, softmax_index, router_match)); + append_match_to_plan(plan, router_match, covered_nodes, &graph); + return plan; +} + +static void append_qwen_routed_gate_up_for_graph(ggml::hrx::Graph & graph, + const QwenRoutedGateUpTensors & tensors, + std::vector & covered_nodes, + ggml::hrx::CommandPlan & plan) { + const size_t gate_index = producer_index_for_tensor(graph, tensors.gate); + ggml::hrx::DispatchMatch gate_up_match; + REQUIRE(match_dispatch_at_index(graph, plan, covered_nodes, gate_index, gate_up_match)); + append_match_to_plan(plan, gate_up_match, covered_nodes, &graph); +} + +static void append_qwen_routed_down_for_graph(ggml::hrx::Graph & graph, + const QwenRoutedGateUpTensors & tensors, + std::vector & covered_nodes, + ggml::hrx::CommandPlan & plan, + const char * expected_kernel_name) { + const size_t down_index = producer_index_for_tensor(graph, tensors.output); + ggml::hrx::DispatchMatch down_match; + REQUIRE(match_dispatch_at_index(graph, plan, covered_nodes, down_index, down_match)); + append_match_to_plan(plan, down_match, covered_nodes, &graph); + + REQUIRE(plan.dispatches.size() >= 1); + REQUIRE(plan.transients.size() >= 1); + const ggml::hrx::Value * route_ids_value = graph.values().find_tensor(tensors.route_ids); + const ggml::hrx::Value * glu_value = graph.values().find_tensor(tensors.glu); + const ggml::hrx::Value * output_value = graph.values().find_tensor(tensors.output); + REQUIRE(route_ids_value != nullptr); + REQUIRE(glu_value != nullptr); + REQUIRE(output_value != nullptr); + const ggml::hrx::CommandPlanMoeRoutingBundle * routing_bundle = + plan.metadata.find_moe_routing_bundle(route_ids_value->id); + const ggml::hrx::CommandPlanAlternateValue * gate_up_alternate = plan.metadata.find_alternate_value( + glu_value->id, GGML_TYPE_F16, qwen_routed_gate_up_f16_output_size(tensors.output->ne[2])); + const ggml::hrx::CommandPlanAlternateValue * routed_down_alternate = plan.metadata.find_alternate_value( + output_value->id, GGML_TYPE_F16, qwen_routed_down_f16_output_size(tensors.output->ne[2])); + REQUIRE(routing_bundle != nullptr); + REQUIRE(gate_up_alternate != nullptr); + REQUIRE(routed_down_alternate != nullptr); + const ggml::hrx::Dispatch & dispatch = plan.dispatches.back(); + const ggml::hrx::CommandPlanTransient & routed_down_transient = plan.transients.back(); + REQUIRE(kernel_name_for_id(dispatch.kernel.kernel_id) == expected_kernel_name); + REQUIRE(dispatch.kernel.integer_parameters.at("token_count") == tensors.output->ne[2]); + REQUIRE(dispatch.bindings.size() == 4); + REQUIRE(routed_down_transient.name == "qwen.moe.routed_down_f16"); + REQUIRE(routed_down_transient.size == qwen_routed_down_f16_output_size(tensors.output->ne[2])); + REQUIRE(dispatch.bindings[0].value == gate_up_alternate->alternate_value); + REQUIRE(dispatch.bindings[0].length == qwen_routed_gate_up_f16_output_size(tensors.output->ne[2])); + REQUIRE(dispatch.bindings[1].value == routing_bundle->expert_table); + REQUIRE(dispatch.bindings[1].length == routing_bundle->expert_table_byte_count); + REQUIRE(routed_down_alternate->alternate_value == routed_down_transient.value); + REQUIRE(routed_down_alternate->byte_count == routed_down_transient.size); + REQUIRE(dispatch.bindings[3].value == routed_down_transient.value); + REQUIRE(dispatch.bindings[3].length == routed_down_transient.size); + require_compile_parameter(dispatch, "qwen3_moe.routed_down.input_size", "768"); + require_compile_parameter(dispatch, "qwen3_moe.routed_down.route_count", "8"); + require_compile_parameter(dispatch, "qwen3_moe.routed_down.expert_count", "128"); + require_compile_parameter(dispatch, "qwen3_moe.routed_down.output_size", "2048"); + require_compile_parameter(dispatch, "qwen3_moe.workload.token_capacity", std::to_string(tensors.output->ne[2])); +} + +static void append_qwen_weighted_reduce_for_graph(ggml::hrx::Graph & graph, + const QwenRoutedGateUpTensors & tensors, + std::vector & covered_nodes, + ggml::hrx::CommandPlan & plan, + const char * expected_kernel_name) { + REQUIRE(tensors.weighted != nullptr); + REQUIRE(tensors.residual != nullptr); + const size_t weighted_index = producer_index_for_tensor(graph, tensors.weighted); + ggml::hrx::DispatchMatch weighted_match; + REQUIRE(match_dispatch_at_index(graph, plan, covered_nodes, weighted_index, weighted_match)); + append_match_to_plan(plan, weighted_match, covered_nodes, &graph); + + const ggml::hrx::Value * route_weights_value = graph.values().find_tensor(tensors.route_weights); + const ggml::hrx::Value * routed_output_value = graph.values().find_tensor(tensors.output); + const ggml::hrx::Value * residual_value = graph.values().find_tensor(tensors.residual); + REQUIRE(route_weights_value != nullptr); + REQUIRE(routed_output_value != nullptr); + REQUIRE(residual_value != nullptr); + + const ggml::hrx::Dispatch & dispatch = plan.dispatches.back(); + REQUIRE(kernel_name_for_id(dispatch.kernel.kernel_id) == expected_kernel_name); + REQUIRE(dispatch.kernel.integer_parameters.at("token_count") == tensors.output->ne[2]); + REQUIRE(dispatch.bindings.size() == (tensors.next_output != nullptr ? 5 : 3)); + REQUIRE(dispatch.bindings[0].value == route_weights_value->id); + REQUIRE(dispatch.bindings[0].length == route_weights_value->byte_count); + REQUIRE(dispatch.bindings[1].value == plan.metadata.alternate_values().back().alternate_value); + REQUIRE(dispatch.bindings[1].length == qwen_routed_down_f16_output_size(tensors.output->ne[2])); + REQUIRE(dispatch.bindings[2].value == residual_value->id); + REQUIRE(dispatch.bindings[2].length == residual_value->byte_count); + require_compile_parameter(dispatch, "qwen3_moe.routed_down.route_count", "8"); + require_compile_parameter(dispatch, "qwen3_moe.routed_down.output_size", "2048"); + require_compile_parameter(dispatch, "qwen3_moe.workload.token_capacity", std::to_string(tensors.output->ne[2])); + + if (tensors.next_output != nullptr) { + const ggml::hrx::Value * next_output_value = graph.values().find_tensor(tensors.next_output); + REQUIRE(next_output_value != nullptr); + REQUIRE(dispatch.bindings[4].value == next_output_value->id); + REQUIRE(dispatch.bindings[4].length == next_output_value->byte_count); + require_compile_parameter(dispatch, "qwen3_moe.model.hidden_size", "2048"); + require_compile_parameter(dispatch, "qwen3_moe.model.rms_epsilon", "0.000001"); + } +} + +static void run_qwen_routed_gate_up_dispatch_checks() { + ggml_init_params params = {}; + params.mem_size = 4 * 1024 * 1024; + params.no_alloc = true; + ggml_context * ctx = ggml_init(params); + REQUIRE(ctx != nullptr); + + { + constexpr int64_t token_count = 4; + const QwenRoutedGateUpTensors tensors = build_qwen_routed_gate_up_graph(ctx, token_count); + ggml::hrx::GraphImportResult imported = import_qwen_routed_gate_up_graph(ctx, tensors); + REQUIRE(imported.graph.nodes().size() == 14); + + const ggml::hrx::Value * route_ids_value = imported.graph.values().find_tensor(tensors.route_ids); + const ggml::hrx::Value * route_weights_value = imported.graph.values().find_tensor(tensors.route_weights); + const ggml::hrx::Value * glu_value = imported.graph.values().find_tensor(tensors.glu); + REQUIRE(route_ids_value != nullptr); + REQUIRE(route_weights_value != nullptr); + REQUIRE(glu_value != nullptr); + REQUIRE(route_ids_value->kind == ggml::hrx::ValueKind::Transient); + REQUIRE(route_weights_value->kind == ggml::hrx::ValueKind::External); + REQUIRE(glu_value->kind == ggml::hrx::ValueKind::Transient); + + std::vector covered_nodes(imported.graph.nodes().size(), false); + ggml::hrx::CommandPlan plan = build_qwen_router_plan_for_graph(imported.graph, covered_nodes); + const ggml::hrx::CommandPlanGeneratedResource * expert_table_resource = plan.metadata.find_generated_resource( + route_ids_value->id, ggml::hrx::GeneratedResourceRole::MoeExpertTable); + const ggml::hrx::CommandPlanGeneratedResource * partition_table_resource = + plan.metadata.find_generated_resource(route_ids_value->id, + ggml::hrx::GeneratedResourceRole::MoePartitionTable); + const ggml::hrx::CommandPlanMoeRoutingBundle * routing_bundle = + plan.metadata.find_moe_routing_bundle(route_ids_value->id); + REQUIRE(expert_table_resource != nullptr); + REQUIRE(partition_table_resource != nullptr); + REQUIRE(routing_bundle != nullptr); + REQUIRE(routing_bundle->route_ids == route_ids_value->id); + REQUIRE(routing_bundle->route_weights == route_weights_value->id); + REQUIRE(routing_bundle->expert_table == expert_table_resource->generated_value); + REQUIRE(routing_bundle->partition_table == partition_table_resource->generated_value); + REQUIRE(routing_bundle->expert_table_byte_count == qwen_expert_table_size(token_count)); + REQUIRE(routing_bundle->partition_table_byte_count == qwen_partition_table_size(token_count)); + REQUIRE(routing_bundle->route_count == kQwenRouterRouteCount); + REQUIRE(routing_bundle->expert_count == kQwenRouterExpertCount); + ggml::hrx::MoeRoutingResourceMetadata expert_metadata; + ggml::hrx::MoeRoutingResourceMetadata partition_metadata; + REQUIRE(expert_table_resource->metadata.read(expert_metadata)); + REQUIRE(partition_table_resource->metadata.read(partition_metadata)); + REQUIRE(expert_metadata.token_count == token_count); + REQUIRE(expert_metadata.route_count == kQwenRouterRouteCount); + REQUIRE(expert_metadata.expert_count == kQwenRouterExpertCount); + REQUIRE(partition_metadata.route_stride == expert_metadata.route_stride); + REQUIRE(expert_table_resource->byte_count == qwen_expert_table_size(token_count)); + REQUIRE(partition_table_resource->byte_count == qwen_partition_table_size(token_count)); + + const size_t gate_index = producer_index_for_tensor(imported.graph, tensors.gate); + ggml::hrx::DispatchMatch gate_up_match; + REQUIRE(match_dispatch_at_index(imported.graph, plan, covered_nodes, gate_index, gate_up_match)); + append_match_to_plan(plan, gate_up_match, covered_nodes); + + REQUIRE(plan.dispatches.size() == 4); + REQUIRE(plan.transients.size() == 3); + REQUIRE(plan.metadata.alternate_values().size() == 1); + const ggml::hrx::CommandPlanTransient & f16_output_transient = plan.transients.back(); + REQUIRE(f16_output_transient.name == "qwen.moe.gate_up_swiglu_f16"); + REQUIRE(f16_output_transient.size == qwen_routed_gate_up_f16_output_size(token_count)); + REQUIRE(plan.metadata.alternate_values().front().graph_value == glu_value->id); + REQUIRE(plan.metadata.alternate_values().front().alternate_value == f16_output_transient.value); + REQUIRE(plan.metadata.alternate_values().front().type == GGML_TYPE_F16); + REQUIRE(plan.metadata.alternate_values().front().byte_count == f16_output_transient.size); + + const ggml::hrx::Dispatch & dispatch = plan.dispatches.back(); + REQUIRE(kernel_name_for_id(dispatch.kernel.kernel_id) == + "qwen3_moe:qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma"); + REQUIRE(dispatch.kernel.integer_parameters.at("token_count") == token_count); + REQUIRE(dispatch.bindings.size() == 6); + REQUIRE(dispatch.bindings[1].value == routing_bundle->expert_table); + REQUIRE(dispatch.bindings[1].length == qwen_expert_table_size(token_count)); + REQUIRE(dispatch.bindings[2].value == routing_bundle->partition_table); + REQUIRE(dispatch.bindings[2].length == qwen_partition_table_size(token_count)); + REQUIRE(dispatch.bindings[5].value == f16_output_transient.value); + REQUIRE(dispatch.bindings[5].length == f16_output_transient.size); + require_compile_parameter(dispatch, "qwen3_moe.routed_gate_up.input_size", "2048"); + require_compile_parameter(dispatch, "qwen3_moe.routed_gate_up.expert_count", "128"); + require_compile_parameter(dispatch, "qwen3_moe.routed_gate_up.route_count", "8"); + require_compile_parameter(dispatch, "qwen3_moe.routed_gate_up.output_size", "768"); + require_compile_parameter(dispatch, "qwen3_moe.workload.token_capacity", std::to_string(token_count)); + + const ggml::hrx::CommandProgram commands = + ggml::hrx::build_command_program(imported.graph, plan, ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); + REQUIRE(commands.valid()); + REQUIRE(commands.commands.size() == 4); + REQUIRE(command_program_verifies(commands)); + REQUIRE(commands.commands[3].bindings.size() == 6); + REQUIRE(commands.commands[3].bindings[0].name == "input"); + REQUIRE(commands.commands[3].bindings[0].origin == ggml::hrx::CommandBindingOrigin::GraphValue); + REQUIRE(commands.commands[3].bindings[1].name == "expert_table"); + REQUIRE(commands.commands[3].bindings[1].origin == ggml::hrx::CommandBindingOrigin::Transient); + REQUIRE(commands.commands[3].bindings[2].name == "partition_table"); + REQUIRE(commands.commands[3].bindings[2].origin == ggml::hrx::CommandBindingOrigin::Transient); + REQUIRE(commands.commands[3].bindings[3].name == "gate_weight"); + REQUIRE(commands.commands[3].bindings[3].origin == ggml::hrx::CommandBindingOrigin::GraphValue); + REQUIRE(commands.commands[3].bindings[4].name == "up_weight"); + REQUIRE(commands.commands[3].bindings[4].origin == ggml::hrx::CommandBindingOrigin::GraphValue); + REQUIRE(commands.commands[3].bindings[5].name == "output"); + REQUIRE(commands.commands[3].bindings[5].origin == ggml::hrx::CommandBindingOrigin::Transient); + REQUIRE(ggml::hrx::find_transient_allocation(commands.transients, f16_output_transient.value) != nullptr); + } + + { + constexpr int64_t token_count = 1; + const QwenRoutedGateUpTensors tensors = build_qwen_routed_gate_up_graph(ctx, token_count); + ggml::hrx::GraphImportResult imported = import_qwen_routed_gate_up_graph(ctx, tensors); + std::vector covered_nodes(imported.graph.nodes().size(), false); + ggml::hrx::CommandPlan plan = build_qwen_router_plan_for_graph(imported.graph, covered_nodes); + const size_t gate_index = producer_index_for_tensor(imported.graph, tensors.gate); + ggml::hrx::DispatchMatch gate_up_match; + REQUIRE(match_dispatch_at_index(imported.graph, plan, covered_nodes, gate_index, gate_up_match)); + append_match_to_plan(plan, gate_up_match, covered_nodes); + + REQUIRE(plan.dispatches.size() == 4); + REQUIRE(plan.transients.size() == 3); + REQUIRE(kernel_name_for_id(plan.dispatches.back().kernel.kernel_id) == + "qwen3_moe:qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma"); + REQUIRE(plan.dispatches.back().kernel.integer_parameters.at("token_count") == 1); + + const ggml::hrx::CommandProgram commands = + ggml::hrx::build_command_program(imported.graph, plan, ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); + REQUIRE(commands.valid()); + REQUIRE(command_program_verifies(commands)); + } + + { + constexpr int64_t token_count = 4; + ggml_tensor * first_logits = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, kQwenRouterExpertCount, token_count); + ggml_tensor * first_route_ids = nullptr; + REQUIRE(first_logits != nullptr); + ggml_tensor * first_route_weights = build_qwen_router_top8_graph(ctx, first_logits, &first_route_ids); + REQUIRE(first_route_ids != nullptr); + REQUIRE(first_route_weights != nullptr); + + const QwenRoutedGateUpTensors tensors = build_qwen_routed_gate_up_graph(ctx, token_count); + ggml_cgraph * graph = ggml_new_graph(ctx); + REQUIRE(graph != nullptr); + ggml_build_forward_expand(graph, first_route_weights); + ggml_build_forward_expand(graph, tensors.route_weights); + ggml_build_forward_expand(graph, tensors.output); + + ggml::hrx::GraphImportResult imported = ggml::hrx::import_ggml_graph(*graph); + REQUIRE(imported.valid()); + + const ggml::hrx::Value * first_route_ids_value = imported.graph.values().find_tensor(first_route_ids); + const ggml::hrx::Value * second_route_ids_value = imported.graph.values().find_tensor(tensors.route_ids); + REQUIRE(first_route_ids_value != nullptr); + REQUIRE(second_route_ids_value != nullptr); + REQUIRE(first_route_ids_value->id != second_route_ids_value->id); + + std::vector covered_nodes(imported.graph.nodes().size(), false); + ggml::hrx::CommandPlan plan; + size_t router_matches = 0; + for (size_t i = 0; i < imported.graph.nodes().size(); ++i) { + if (imported.graph.nodes()[i].op != GGML_OP_SOFT_MAX) { + continue; + } + ggml::hrx::DispatchMatch router_match; + REQUIRE(match_dispatch_at_index(imported.graph, plan, covered_nodes, i, router_match)); + append_match_to_plan(plan, router_match, covered_nodes); + ++router_matches; + } + REQUIRE(router_matches == 2); + REQUIRE(plan.metadata.generated_resources().size() == 4); + REQUIRE(plan.metadata.moe_routing_bundles().size() == 2); + + const ggml::hrx::CommandPlanGeneratedResource * first_expert_table = plan.metadata.find_generated_resource( + first_route_ids_value->id, ggml::hrx::GeneratedResourceRole::MoeExpertTable); + const ggml::hrx::CommandPlanGeneratedResource * second_expert_table = plan.metadata.find_generated_resource( + second_route_ids_value->id, ggml::hrx::GeneratedResourceRole::MoeExpertTable); + const ggml::hrx::CommandPlanGeneratedResource * second_partition_table = plan.metadata.find_generated_resource( + second_route_ids_value->id, ggml::hrx::GeneratedResourceRole::MoePartitionTable); + const ggml::hrx::CommandPlanMoeRoutingBundle * second_routing_bundle = + plan.metadata.find_moe_routing_bundle(second_route_ids_value->id); + REQUIRE(first_expert_table != nullptr); + REQUIRE(second_expert_table != nullptr); + REQUIRE(second_partition_table != nullptr); + REQUIRE(second_routing_bundle != nullptr); + REQUIRE(second_routing_bundle->expert_table == second_expert_table->generated_value); + REQUIRE(second_routing_bundle->partition_table == second_partition_table->generated_value); + REQUIRE(first_expert_table->generated_value != second_expert_table->generated_value); + + const size_t gate_index = producer_index_for_tensor(imported.graph, tensors.gate); + ggml::hrx::DispatchMatch gate_up_match; + REQUIRE(match_dispatch_at_index(imported.graph, plan, covered_nodes, gate_index, gate_up_match)); + append_match_to_plan(plan, gate_up_match, covered_nodes); + + REQUIRE(plan.dispatches.size() == 7); + REQUIRE(plan.transients.size() == 5); + const ggml::hrx::Dispatch & dispatch = plan.dispatches.back(); + REQUIRE(dispatch.bindings.size() == 6); + REQUIRE(dispatch.bindings[1].value == second_routing_bundle->expert_table); + REQUIRE(dispatch.bindings[1].value != first_expert_table->generated_value); + REQUIRE(dispatch.bindings[2].value == second_routing_bundle->partition_table); + + const ggml::hrx::CommandProgram commands = + ggml::hrx::build_command_program(imported.graph, plan, ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); + REQUIRE(commands.valid()); + REQUIRE(commands.commands.size() == 7); + REQUIRE(command_program_verifies(commands)); + } + + { + constexpr int64_t token_count = 4; + const QwenRoutedGateUpTensors tensors = build_qwen_routed_gate_up_graph(ctx, token_count); + ggml::hrx::GraphImportResult imported = import_qwen_routed_gate_up_graph(ctx, tensors); + std::vector covered_nodes(imported.graph.nodes().size(), false); + ggml::hrx::CommandPlan plan = build_qwen_router_plan_for_graph(imported.graph, covered_nodes); + append_qwen_routed_gate_up_for_graph(imported.graph, tensors, covered_nodes, plan); + append_qwen_routed_down_for_graph(imported.graph, tensors, covered_nodes, plan, + "qwen3_moe:qwen3_moe_routed_down_q6k_f16_wmma_grouped"); + + REQUIRE(plan.dispatches.size() == 5); + REQUIRE(plan.transients.size() == 4); + const ggml::hrx::CommandProgram commands = + ggml::hrx::build_command_program(imported.graph, plan, ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); + REQUIRE(commands.valid()); + REQUIRE(commands.commands.size() == 5); + REQUIRE(command_program_verifies(commands)); + REQUIRE(ggml::hrx::find_transient_allocation(commands.transients, plan.transients.back().value) != nullptr); + } + + { + constexpr int64_t token_count = 4; + QwenRoutedGateUpTensors tensors = build_qwen_routed_gate_up_graph(ctx, token_count); + append_qwen_weighted_reduce_tail(ctx, tensors); + ggml::hrx::GraphImportResult imported = import_qwen_routed_gate_up_graph(ctx, tensors); + std::vector covered_nodes(imported.graph.nodes().size(), false); + ggml::hrx::CommandPlan plan = build_qwen_router_plan_for_graph(imported.graph, covered_nodes); + append_qwen_routed_gate_up_for_graph(imported.graph, tensors, covered_nodes, plan); + append_qwen_routed_down_for_graph(imported.graph, tensors, covered_nodes, plan, + "qwen3_moe:qwen3_moe_routed_down_q6k_f16_wmma_grouped"); + append_qwen_weighted_reduce_for_graph(imported.graph, tensors, covered_nodes, plan, + "qwen3_moe:qwen3_moe_routed_down_weighted_reduce_f16_f32"); + + REQUIRE(plan.dispatches.size() == 6); + REQUIRE(plan.transients.size() == 4); + const ggml::hrx::CommandProgram commands = + ggml::hrx::build_command_program(imported.graph, plan, ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); + REQUIRE(commands.valid()); + REQUIRE(commands.commands.size() == 6); + REQUIRE(command_program_verifies(commands)); + REQUIRE(commands.commands.back().bindings.size() == 3); + REQUIRE(commands.commands.back().bindings[0].name == "route_weights"); + REQUIRE(commands.commands.back().bindings[1].name == "routed_output"); + REQUIRE(commands.commands.back().bindings[1].origin == ggml::hrx::CommandBindingOrigin::Transient); + REQUIRE(commands.commands.back().bindings[2].name == "output"); + REQUIRE(commands.commands.back().bindings[2].access == ggml::hrx::ResourceAccess::ReadWrite); + } + + { + constexpr int64_t token_count = 1; + QwenRoutedGateUpTensors tensors = build_qwen_routed_gate_up_graph(ctx, token_count); + append_qwen_weighted_reduce_tail(ctx, tensors); + ggml::hrx::GraphImportResult imported = import_qwen_routed_gate_up_graph(ctx, tensors); + std::vector covered_nodes(imported.graph.nodes().size(), false); + ggml::hrx::CommandPlan plan = build_qwen_router_plan_for_graph(imported.graph, covered_nodes); + append_qwen_routed_gate_up_for_graph(imported.graph, tensors, covered_nodes, plan); + append_qwen_routed_down_for_graph(imported.graph, tensors, covered_nodes, plan, + "qwen3_moe:qwen3_moe_routed_down_q6k_f16_wmma_grouped"); + append_qwen_weighted_reduce_for_graph(imported.graph, tensors, covered_nodes, plan, + "qwen3_moe:qwen3_moe_routed_down_weighted_reduce_f16_f32"); + + REQUIRE(plan.dispatches.size() == 6); + REQUIRE(plan.dispatches.back().kernel.integer_parameters.at("token_count") == 1); + const ggml::hrx::CommandProgram commands = + ggml::hrx::build_command_program(imported.graph, plan, ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); + REQUIRE(commands.valid()); + REQUIRE(command_program_verifies(commands)); + } + + { + constexpr int64_t token_count = 4; + QwenRoutedGateUpTensors tensors = build_qwen_routed_gate_up_graph(ctx, token_count); + append_qwen_weighted_reduce_tail(ctx, tensors, true); + ggml::hrx::GraphImportResult imported = import_qwen_routed_gate_up_graph(ctx, tensors); + std::vector covered_nodes(imported.graph.nodes().size(), false); + ggml::hrx::CommandPlan plan = build_qwen_router_plan_for_graph(imported.graph, covered_nodes); + append_qwen_routed_gate_up_for_graph(imported.graph, tensors, covered_nodes, plan); + append_qwen_routed_down_for_graph(imported.graph, tensors, covered_nodes, plan, + "qwen3_moe:qwen3_moe_routed_down_q6k_f16_wmma_grouped"); + append_qwen_weighted_reduce_for_graph(imported.graph, tensors, covered_nodes, plan, + "qwen3_moe:qwen3_moe_routed_down_weighted_reduce_next_rmsnorm_f32"); + + REQUIRE(plan.dispatches.size() == 6); + REQUIRE(plan.transients.size() == 4); + const ggml::hrx::Value * hidden_state_value = imported.graph.values().find_tensor(tensors.hidden_state); + const ggml::hrx::Value * residual_value = imported.graph.values().find_tensor(tensors.residual); + REQUIRE(hidden_state_value != nullptr); + REQUIRE(residual_value != nullptr); + REQUIRE(residual_value->alias_source == hidden_state_value->id); + REQUIRE(imported.graph.values().same_storage(hidden_state_value->id, residual_value->id)); + REQUIRE(residual_value->storage_root == hidden_state_value->storage_root); + const ggml::hrx::CommandProgram commands = + ggml::hrx::build_command_program(imported.graph, plan, ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); + REQUIRE(commands.valid()); + REQUIRE(commands.commands.size() == 6); + REQUIRE(command_program_verifies(commands)); + REQUIRE(commands.commands.back().bindings.size() == 5); + REQUIRE(commands.commands.back().bindings[0].name == "route_weights"); + REQUIRE(commands.commands.back().bindings[1].name == "routed_output"); + REQUIRE(commands.commands.back().bindings[1].origin == ggml::hrx::CommandBindingOrigin::Transient); + REQUIRE(commands.commands.back().bindings[2].name == "hidden_state"); + REQUIRE(commands.commands.back().bindings[2].value == hidden_state_value->storage_root); + REQUIRE(commands.commands.back().bindings[2].access == ggml::hrx::ResourceAccess::ReadWrite); + REQUIRE(commands.commands.back().bindings[3].name == "next_norm_weight"); + REQUIRE(commands.commands.back().bindings[4].name == "next_projection_input"); + REQUIRE(commands.commands.back().bindings[4].access == ggml::hrx::ResourceAccess::Write); + } + + { + constexpr int64_t token_count = 4; + QwenRoutedGateUpTensors tensors = build_qwen_routed_gate_up_graph(ctx, token_count); + append_qwen_weighted_reduce_tail(ctx, tensors, true); + ggml_tensor * hidden_bias = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, kQwenMoeHiddenSize, token_count); + REQUIRE(hidden_bias != nullptr); + tensors.hidden_use = ggml_add(ctx, tensors.hidden_state, hidden_bias); + REQUIRE(tensors.hidden_use != nullptr); + ggml::hrx::GraphImportResult imported = import_qwen_routed_gate_up_graph(ctx, tensors); + std::vector covered_nodes(imported.graph.nodes().size(), false); + ggml::hrx::CommandPlan plan = build_qwen_router_plan_for_graph(imported.graph, covered_nodes); + append_qwen_routed_gate_up_for_graph(imported.graph, tensors, covered_nodes, plan); + append_qwen_routed_down_for_graph(imported.graph, tensors, covered_nodes, plan, + "qwen3_moe:qwen3_moe_routed_down_q6k_f16_wmma_grouped"); + + const size_t weighted_index = producer_index_for_tensor(imported.graph, tensors.weighted); + ggml::hrx::DispatchMatch weighted_match; + REQUIRE(match_dispatch_at_index(imported.graph, plan, covered_nodes, weighted_index, weighted_match)); + append_match_to_plan(plan, weighted_match, covered_nodes, &imported.graph); + + REQUIRE(plan.dispatches.size() == 6); + REQUIRE(kernel_name_for_id(plan.dispatches.back().kernel.kernel_id) == + "qwen3_moe:qwen3_moe_routed_down_weighted_reduce_f16_f32"); + REQUIRE(weighted_match.value_aliases.empty()); + REQUIRE(plan.dispatches.back().bindings.size() == 3); + const ggml::hrx::Value * hidden_state_value = imported.graph.values().find_tensor(tensors.hidden_state); + const ggml::hrx::Value * residual_value = imported.graph.values().find_tensor(tensors.residual); + REQUIRE(hidden_state_value != nullptr); + REQUIRE(residual_value != nullptr); + REQUIRE(residual_value->alias_source != hidden_state_value->id); + REQUIRE(!imported.graph.values().same_storage(hidden_state_value->id, residual_value->id)); + } + + { + constexpr int64_t token_count = 4; + const QwenRoutedGateUpTensors tensors = + build_qwen_routed_gate_up_graph(ctx, token_count, GGML_GLU_OP_SWIGLU, GGML_TYPE_Q4_K, true, GGML_TYPE_Q4_K); + ggml::hrx::GraphImportResult imported = import_qwen_routed_gate_up_graph(ctx, tensors); + std::vector covered_nodes(imported.graph.nodes().size(), false); + ggml::hrx::CommandPlan plan = build_qwen_router_plan_for_graph(imported.graph, covered_nodes); + append_qwen_routed_gate_up_for_graph(imported.graph, tensors, covered_nodes, plan); + append_qwen_routed_down_for_graph(imported.graph, tensors, covered_nodes, plan, + "qwen3_moe:qwen3_moe_routed_down_q4k_f16_wmma_grouped"); + + REQUIRE(plan.dispatches.size() == 5); + REQUIRE(plan.transients.size() == 4); + const ggml::hrx::CommandProgram commands = + ggml::hrx::build_command_program(imported.graph, plan, ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); + REQUIRE(commands.valid()); + REQUIRE(commands.commands.size() == 5); + REQUIRE(command_program_verifies(commands)); + } + + { + const QwenRoutedGateUpTensors tensors = build_qwen_routed_gate_up_graph(ctx, 4); + ggml::hrx::GraphImportResult imported = import_qwen_routed_gate_up_graph(ctx, tensors); + const size_t gate_index = producer_index_for_tensor(imported.graph, tensors.gate); + std::vector covered_nodes(imported.graph.nodes().size(), false); + const ggml::hrx::CommandPlan empty_plan; + ggml::hrx::DispatchMatch gate_up_match; + REQUIRE(!match_dispatch_at_index(imported.graph, empty_plan, covered_nodes, gate_index, gate_up_match)); + } + + { + const QwenRoutedGateUpTensors tensors = build_qwen_routed_gate_up_graph(ctx, 4, GGML_GLU_OP_GEGLU); + ggml::hrx::GraphImportResult imported = import_qwen_routed_gate_up_graph(ctx, tensors); + std::vector covered_nodes(imported.graph.nodes().size(), false); + ggml::hrx::CommandPlan plan = build_qwen_router_plan_for_graph(imported.graph, covered_nodes); + const size_t gate_index = producer_index_for_tensor(imported.graph, tensors.gate); + ggml::hrx::DispatchMatch gate_up_match; + REQUIRE(!match_dispatch_at_index(imported.graph, plan, covered_nodes, gate_index, gate_up_match)); + } + + { + const QwenRoutedGateUpTensors tensors = + build_qwen_routed_gate_up_graph(ctx, 4, GGML_GLU_OP_SWIGLU, GGML_TYPE_Q6_K); + ggml::hrx::GraphImportResult imported = import_qwen_routed_gate_up_graph(ctx, tensors); + std::vector covered_nodes(imported.graph.nodes().size(), false); + ggml::hrx::CommandPlan plan = build_qwen_router_plan_for_graph(imported.graph, covered_nodes); + const size_t gate_index = producer_index_for_tensor(imported.graph, tensors.gate); + ggml::hrx::DispatchMatch gate_up_match; + REQUIRE(!match_dispatch_at_index(imported.graph, plan, covered_nodes, gate_index, gate_up_match)); + } + + { + const QwenRoutedGateUpTensors tensors = + build_qwen_routed_gate_up_graph(ctx, 4, GGML_GLU_OP_SWIGLU, GGML_TYPE_Q4_K, false); + ggml::hrx::GraphImportResult imported = import_qwen_routed_gate_up_graph(ctx, tensors); + std::vector covered_nodes(imported.graph.nodes().size(), false); + ggml::hrx::CommandPlan plan = build_qwen_router_plan_for_graph(imported.graph, covered_nodes); + const size_t gate_index = producer_index_for_tensor(imported.graph, tensors.gate); + ggml::hrx::DispatchMatch gate_up_match; + REQUIRE(!match_dispatch_at_index(imported.graph, plan, covered_nodes, gate_index, gate_up_match)); + } + + { + const QwenRoutedGateUpTensors tensors = build_qwen_routed_gate_up_graph(ctx, 4); + ggml::hrx::GraphImportResult imported = import_qwen_routed_gate_up_graph(ctx, tensors); + std::vector covered_nodes(imported.graph.nodes().size(), false); + ggml::hrx::CommandPlan plan = build_qwen_router_plan_for_graph(imported.graph, covered_nodes); + const size_t down_index = producer_index_for_tensor(imported.graph, tensors.output); + ggml::hrx::DispatchMatch down_match; + REQUIRE(!match_dispatch_at_index(imported.graph, plan, covered_nodes, down_index, down_match)); + } + + { + const QwenRoutedGateUpTensors tensors = + build_qwen_routed_gate_up_graph(ctx, 4, GGML_GLU_OP_SWIGLU, GGML_TYPE_Q4_K, true, GGML_TYPE_Q5_K); + ggml::hrx::GraphImportResult imported = import_qwen_routed_gate_up_graph(ctx, tensors); + std::vector covered_nodes(imported.graph.nodes().size(), false); + ggml::hrx::CommandPlan plan = build_qwen_router_plan_for_graph(imported.graph, covered_nodes); + append_qwen_routed_gate_up_for_graph(imported.graph, tensors, covered_nodes, plan); + const size_t down_index = producer_index_for_tensor(imported.graph, tensors.output); + ggml::hrx::DispatchMatch down_match; + REQUIRE(!match_dispatch_at_index(imported.graph, plan, covered_nodes, down_index, down_match)); + } + + ggml_free(ctx); +} + +static void run_qwen_router_top8_dispatch_checks() { + ggml_init_params params = {}; + params.mem_size = 2 * 1024 * 1024; + params.no_alloc = true; + ggml_context * ctx = ggml_init(params); + REQUIRE(ctx != nullptr); + + { + ggml_tensor * logits = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, kQwenRouterExpertCount, 4); + ggml_tensor * route_ids = nullptr; + REQUIRE(logits != nullptr); + ggml_tensor * output = build_qwen_router_top8_graph(ctx, logits, &route_ids); + schedule_qwen_router_top8_command(ctx, output, route_ids); + } + { + ggml_tensor * logits = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, kQwenRouterExpertCount, 1); + ggml_tensor * route_ids = nullptr; + REQUIRE(logits != nullptr); + ggml_tensor * output = build_qwen_router_top8_graph(ctx, logits, &route_ids); + schedule_qwen_router_top8_command(ctx, output, route_ids); + } + { + ggml_tensor * logits = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, kQwenRouterExpertCount, 13); + ggml_tensor * route_ids = nullptr; + REQUIRE(logits != nullptr); + ggml_tensor * output = build_qwen_router_top8_graph(ctx, logits, &route_ids, GGML_SORT_ORDER_DESC, + kQwenRouterRouteCount, 0.00006103515625f); + schedule_qwen_router_top8_command(ctx, output, route_ids); + } + { + ggml_tensor * logits = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, kQwenRouterExpertCount, 512); + ggml_tensor * route_ids = nullptr; + REQUIRE(logits != nullptr); + ggml_tensor * output = build_qwen_router_top8_graph(ctx, logits, &route_ids); + schedule_qwen_router_top8_command(ctx, output, route_ids); + } + { + ggml_tensor * logits = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 64, 4); + ggml_tensor * route_ids = nullptr; + REQUIRE(logits != nullptr); + ggml_tensor * output = build_qwen_router_top8_graph(ctx, logits, &route_ids); + schedule_qwen_router_top8_command(ctx, output, route_ids, 64, kQwenRouterRouteCount); + } + { + ggml_tensor * logits = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, kQwenRouterExpertCount, 4); + ggml_tensor * route_ids = nullptr; + REQUIRE(logits != nullptr); + ggml_tensor * output = build_qwen_router_top8_graph(ctx, logits, &route_ids, GGML_SORT_ORDER_DESC, 4); + schedule_qwen_router_top8_command(ctx, output, route_ids, kQwenRouterExpertCount, 4); + } + { + ManualQwenRouterTop8Graph manual = build_manual_qwen_router_top8_graph(ctx, 64, 4, 5); + schedule_manual_qwen_router_top8_command(manual.graph, manual.route_ids, 5, 64, 4); + } + { + ggml_tensor * logits = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, kQwenRouterExpertCount, 4); + REQUIRE(logits != nullptr); + ggml_tensor * output = build_qwen_router_top8_graph(ctx, logits, nullptr, GGML_SORT_ORDER_ASC); + REQUIRE(!graph_is_supported(ctx, output)); + } + { + ggml_tensor * logits = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, kQwenRouterExpertCount, 4); + REQUIRE(logits != nullptr); + ggml_tensor * output = build_qwen_router_top8_graph(ctx, logits, nullptr, GGML_SORT_ORDER_DESC, 33); + REQUIRE(!graph_is_supported(ctx, output)); + } + { + ggml_tensor * logits = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 16, 4); + REQUIRE(logits != nullptr); + ggml_tensor * output = build_qwen_router_top8_graph(ctx, logits); + REQUIRE(!graph_is_supported(ctx, output)); + } + { + ggml_tensor * logits = ggml_new_tensor_2d(ctx, GGML_TYPE_F16, kQwenRouterExpertCount, 4); + REQUIRE(logits != nullptr); + ggml_tensor * output = build_qwen_router_top8_graph(ctx, logits); + REQUIRE(!graph_is_supported(ctx, output)); + } + + ggml_free(ctx); +} + +static void bind_external_values(ggml::hrx::ValueMap & values) { + uintptr_t buffer = 0x1000; + for (const ggml::hrx::ValueId id : values.external_value_ids()) { + const ggml::hrx::Value * value = values.find(id); + REQUIRE(value != nullptr); + REQUIRE(values.bind_buffer(id, { dummy_hrx_buffer(buffer), 0, value->byte_count })); + buffer += 0x1000; + } +} + +static void run_alias_value_import_checks() { + ggml_init_params params = {}; + params.mem_size = 256 * 1024; + params.no_alloc = true; + ggml_context * ctx = ggml_init(params); + REQUIRE(ctx != nullptr); + + ggml_tensor * source = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 8); + ggml_tensor * view = ggml_view_1d(ctx, source, 4, 2 * sizeof(float)); + REQUIRE(source != nullptr); + REQUIRE(view != nullptr); + + ggml_cgraph * graph = ggml_new_graph(ctx); + REQUIRE(graph != nullptr); + ggml_build_forward_expand(graph, view); + + ggml::hrx::GraphImportResult imported = ggml::hrx::import_ggml_graph(*graph); + REQUIRE(imported.valid()); + REQUIRE(imported.graph.nodes().size() == 1); + REQUIRE(imported.graph.nodes()[0].op == GGML_OP_VIEW); + + const ggml::hrx::Value * source_value = imported.graph.values().find_tensor(source); + const ggml::hrx::Value * view_value = imported.graph.values().find_tensor(view); + REQUIRE(source_value != nullptr); + REQUIRE(view_value != nullptr); + REQUIRE(source_value->kind == ggml::hrx::ValueKind::External); + REQUIRE(view_value->kind == ggml::hrx::ValueKind::External); + REQUIRE(imported.graph.values().same_storage(source_value->id, view_value->id)); + REQUIRE(view_value->storage_root == source_value->id); + REQUIRE(view_value->alias_source == source_value->id); + REQUIRE(view_value->storage_offset == 2 * sizeof(float)); + REQUIRE(view_value->byte_count == 4 * sizeof(float)); + REQUIRE(ggml::hrx::is_layout_alias_node(imported.graph, imported.graph.nodes()[0])); + + REQUIRE(imported.graph.values().bind_buffer(source_value->id, + { dummy_hrx_buffer(0x4000), 128, source_value->byte_count })); + const ggml::hrx::CommandProgramBindings bindings = + ggml::hrx::CommandProgramBindings::from_value_map(imported.graph.values()); + REQUIRE(bindings.valid()); + const ggml::hrx::CommandProgramBinding * view_binding = bindings.find(view_value->id); + REQUIRE(view_binding != nullptr); + REQUIRE(view_binding->buffer == dummy_hrx_buffer(0x4000)); + REQUIRE(view_binding->offset == 128 + 2 * sizeof(float)); + REQUIRE(view_binding->length == view_value->byte_count); + + ggml_tensor * internal_source = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 8, 2); + ggml_tensor * internal_bias = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 8, 1); + ggml_tensor * internal_view = + ggml_view_2d(ctx, internal_source, 8, 1, internal_source->nb[1], internal_source->nb[1]); + ggml_tensor * internal_out = ggml_add(ctx, internal_view, internal_bias); + REQUIRE(internal_source != nullptr); + REQUIRE(internal_bias != nullptr); + REQUIRE(internal_view != nullptr); + REQUIRE(internal_out != nullptr); + + ggml_cgraph * internal_graph = ggml_new_graph(ctx); + REQUIRE(internal_graph != nullptr); + ggml_build_forward_expand(internal_graph, internal_out); + + ggml::hrx::GraphImportResult internal_imported = ggml::hrx::import_ggml_graph(*internal_graph); + REQUIRE(internal_imported.valid()); + REQUIRE(internal_imported.graph.nodes().size() == 2); + REQUIRE(internal_imported.graph.nodes()[0].op == GGML_OP_VIEW); + REQUIRE(internal_imported.graph.nodes()[1].op == GGML_OP_ADD); + + const ggml::hrx::Value * internal_source_value = internal_imported.graph.values().find_tensor(internal_source); + const ggml::hrx::Value * internal_view_value = internal_imported.graph.values().find_tensor(internal_view); + const ggml::hrx::Value * internal_bias_value = internal_imported.graph.values().find_tensor(internal_bias); + const ggml::hrx::Value * internal_out_value = internal_imported.graph.values().find_tensor(internal_out); + REQUIRE(internal_source_value != nullptr); + REQUIRE(internal_view_value != nullptr); + REQUIRE(internal_bias_value != nullptr); + REQUIRE(internal_out_value != nullptr); + REQUIRE(internal_source_value->kind == ggml::hrx::ValueKind::External); + REQUIRE(internal_view_value->kind == ggml::hrx::ValueKind::External); + REQUIRE(internal_view_value->storage_root == internal_source_value->id); + REQUIRE(internal_view_value->storage_offset == internal_source->nb[1]); + + ggml::hrx::DispatchScheduler internal_scheduler; + REQUIRE(internal_scheduler.schedule_graph(internal_imported.graph, test_dispatch_target())); + REQUIRE(internal_scheduler.plan().valid()); + REQUIRE(internal_scheduler.plan().dispatches.size() == 1); + const ggml::hrx::CommandProgram internal_commands = ggml::hrx::build_command_program( + internal_imported.graph, internal_scheduler.plan(), ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); + REQUIRE(internal_commands.valid()); + REQUIRE(internal_commands.commands.size() == 1); + REQUIRE(internal_commands.commands[0].bindings[0].value == internal_view_value->id); + REQUIRE(internal_commands.commands[0].bindings[0].offset == 0); + REQUIRE(internal_commands.commands[0].bindings[0].origin == ggml::hrx::CommandBindingOrigin::GraphValue); + REQUIRE(ggml::hrx::find_transient_allocation(internal_commands.transients, internal_view_value->id) == nullptr); + + REQUIRE(internal_imported.graph.values().bind_buffer( + internal_source_value->id, { dummy_hrx_buffer(0x5000), 256, internal_source_value->byte_count })); + REQUIRE(internal_imported.graph.values().bind_buffer( + internal_bias_value->id, { dummy_hrx_buffer(0x6000), 0, internal_bias_value->byte_count })); + REQUIRE(internal_imported.graph.values().bind_buffer( + internal_out_value->id, { dummy_hrx_buffer(0x7000), 0, internal_out_value->byte_count })); + const ggml::hrx::CommandProgramBindings internal_bindings = + ggml::hrx::CommandProgramBindings::from_value_map(internal_imported.graph.values()); + REQUIRE(internal_bindings.valid()); + const ggml::hrx::CommandProgramBinding * internal_view_binding = internal_bindings.find(internal_view_value->id); + REQUIRE(internal_view_binding != nullptr); + REQUIRE(internal_view_binding->buffer == dummy_hrx_buffer(0x5000)); + REQUIRE(internal_view_binding->offset == 256 + internal_source->nb[1]); + REQUIRE(internal_view_binding->length == internal_view_value->byte_count); + const ggml::hrx::ResolvedCommandProgram internal_resolved = + ggml::hrx::resolve_command_program_bindings(internal_commands, internal_bindings); + REQUIRE(internal_resolved.valid()); + REQUIRE(internal_resolved.commands[0].bindings[0].ref.buffer == dummy_hrx_buffer(0x5000)); + REQUIRE(internal_resolved.commands[0].bindings[0].ref.offset == 256 + internal_source->nb[1]); + + ggml_tensor * cache = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 8, 4); + ggml_tensor * rows = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 8, 2); + ggml_tensor * row_indices = ggml_new_tensor_1d(ctx, GGML_TYPE_I64, 2); + ggml_tensor * updated = ggml_set_rows(ctx, cache, rows, row_indices); + REQUIRE(cache != nullptr); + REQUIRE(rows != nullptr); + REQUIRE(row_indices != nullptr); + REQUIRE(updated != nullptr); + ggml_tensor * updated_view = ggml_view_2d(ctx, updated, 8, 2, updated->nb[1], 0); + ggml_tensor * updated_bias = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 8, 2); + ggml_tensor * updated_out = ggml_add(ctx, updated_view, updated_bias); + REQUIRE(updated_view != nullptr); + REQUIRE(updated_bias != nullptr); + REQUIRE(updated_out != nullptr); + + ggml_cgraph * set_rows_graph = ggml_new_graph(ctx); + REQUIRE(set_rows_graph != nullptr); + ggml_build_forward_expand(set_rows_graph, updated_out); + + ggml::hrx::GraphImportResult set_rows_imported = ggml::hrx::import_ggml_graph(*set_rows_graph); + REQUIRE(set_rows_imported.valid()); + REQUIRE(set_rows_imported.graph.nodes().size() == 3); + REQUIRE(set_rows_imported.graph.nodes()[0].op == GGML_OP_SET_ROWS); + REQUIRE(set_rows_imported.graph.nodes()[1].op == GGML_OP_VIEW); + REQUIRE(set_rows_imported.graph.nodes()[2].op == GGML_OP_ADD); + + const ggml::hrx::Value * cache_value = set_rows_imported.graph.values().find_tensor(cache); + const ggml::hrx::Value * updated_value = set_rows_imported.graph.values().find_tensor(updated); + const ggml::hrx::Value * updated_view_value = set_rows_imported.graph.values().find_tensor(updated_view); + REQUIRE(cache_value != nullptr); + REQUIRE(updated_value != nullptr); + REQUIRE(updated_view_value != nullptr); + REQUIRE(cache_value->kind == ggml::hrx::ValueKind::External); + REQUIRE(updated_value->kind == ggml::hrx::ValueKind::External); + REQUIRE(updated_view_value->kind == ggml::hrx::ValueKind::External); + REQUIRE(updated_value->storage_root == cache_value->id); + REQUIRE(updated_view_value->storage_root == cache_value->id); + REQUIRE(updated_view_value->alias_source == cache_value->id); + + ggml::hrx::ValueMap value_map; + const std::array ne = { 8, 1, 1, 1 }; + const std::array nb = { sizeof(float), 8 * sizeof(float), 8 * sizeof(float), + 8 * sizeof(float) }; + REQUIRE(value_map.add_snapshot_storage({ ggml::hrx::ValueStorageId(0), ggml::hrx::ValueId(0), 8 * sizeof(float) }) + .success()); + REQUIRE( + value_map + .add_snapshot_value({ ggml::hrx::ValueId(0), ggml::hrx::ValueKind::External, ggml::hrx::ValueStorageId(0), + ggml::hrx::ValueId(0), ggml::hrx::ValueId(), 0, 8 * sizeof(float), GGML_TYPE_F32, ne, + nb, 8, 8 * sizeof(float), true, nullptr, std::nullopt }) + .success()); + REQUIRE(value_map.add_snapshot_storage({ ggml::hrx::ValueStorageId(1), ggml::hrx::ValueId(1), 8 * sizeof(float) }) + .success()); + REQUIRE( + value_map + .add_snapshot_value({ ggml::hrx::ValueId(1), ggml::hrx::ValueKind::Transient, ggml::hrx::ValueStorageId(1), + ggml::hrx::ValueId(1), ggml::hrx::ValueId(), 0, 8 * sizeof(float), GGML_TYPE_F32, ne, + nb, 8, 8 * sizeof(float), true, nullptr, std::nullopt }) + .success()); + REQUIRE(value_map.alias_storage(ggml::hrx::ValueId(1), ggml::hrx::ValueId(0)).success()); + const ggml::hrx::Value * aliased_value = value_map.find(ggml::hrx::ValueId(1)); + REQUIRE(aliased_value != nullptr); + REQUIRE(aliased_value->kind == ggml::hrx::ValueKind::Transient); + REQUIRE(aliased_value->alias_source == ggml::hrx::ValueId(0)); + REQUIRE(aliased_value->storage_root == ggml::hrx::ValueId(0)); + REQUIRE(value_map.same_storage(ggml::hrx::ValueId(0), ggml::hrx::ValueId(1))); + REQUIRE(value_map.bind_buffer(ggml::hrx::ValueId(0), { dummy_hrx_buffer(0x8000), 64, 8 * sizeof(float) })); + const std::optional aliased_binding = + value_map.resolve_buffer_binding(ggml::hrx::ValueId(1)); + REQUIRE(aliased_binding.has_value()); + REQUIRE(aliased_binding->buffer == dummy_hrx_buffer(0x8000)); + REQUIRE(aliased_binding->offset == 64); + REQUIRE(aliased_binding->length == 8 * sizeof(float)); + + ggml_free(ctx); +} + +static void run_multi_dispatch_checks() { + ggml_init_params params = {}; + params.mem_size = 256 * 1024; + params.no_alloc = true; + ggml_context * ctx = ggml_init(params); + REQUIRE(ctx != nullptr); + + ggml_tensor * a = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 8); + ggml_tensor * b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 8); + ggml_tensor * c = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 8); + ggml_tensor * d = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 8); + ggml_tensor * out0 = ggml_add(ctx, a, b); + ggml_tensor * out1 = ggml_add(ctx, c, d); + REQUIRE(a != nullptr); + REQUIRE(b != nullptr); + REQUIRE(c != nullptr); + REQUIRE(d != nullptr); + REQUIRE(out0 != nullptr); + REQUIRE(out1 != nullptr); + + ggml_cgraph * graph = ggml_new_graph(ctx); + REQUIRE(graph != nullptr); + ggml_build_forward_expand(graph, out0); + ggml_build_forward_expand(graph, out1); + graph->uid = 1001; + + ggml::hrx::GraphImportResult imported = ggml::hrx::import_ggml_graph(*graph); + REQUIRE(imported.valid()); + REQUIRE(imported.graph.nodes().size() == 2); + const ggml::hrx::Value * out0_value = imported.graph.values().find_tensor(out0); + const ggml::hrx::Value * out1_value = imported.graph.values().find_tensor(out1); + REQUIRE(out0_value != nullptr); + REQUIRE(out1_value != nullptr); + REQUIRE(imported.graph.nodes()[0].output == out0_value->id); + REQUIRE(imported.graph.nodes()[1].output == out1_value->id); + + ggml::hrx::DispatchScheduler scheduler; + REQUIRE(scheduler.schedule_graph(imported.graph, test_dispatch_target())); + REQUIRE(scheduler.plan().valid()); + REQUIRE(scheduler.plan().dispatches.size() == 2); + + const ggml::hrx::CommandProgram commands = ggml::hrx::build_command_program( + imported.graph, scheduler.plan(), ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); + REQUIRE(commands.valid()); + REQUIRE(commands.commands.size() == 2); + REQUIRE(commands.commands[0].ordinal == 0); + REQUIRE(commands.commands[0].dependencies.empty()); + REQUIRE(commands.commands[1].ordinal == 1); + REQUIRE(commands.commands[1].dependencies.size() == 1); + REQUIRE(commands.commands[1].dependencies[0] == 0); + REQUIRE(command_program_verifies(commands)); + + bind_external_values(imported.graph.values()); + const ggml::hrx::CommandProgramBindings bindings = + ggml::hrx::CommandProgramBindings::from_value_map(imported.graph.values()); + REQUIRE(bindings.valid()); + ggml::hrx::ResolvedCommandProgram resolved = ggml::hrx::resolve_command_program_bindings(commands, bindings); + REQUIRE(resolved.valid()); + REQUIRE(resolved.commands.size() == 2); + + ggml_free(ctx); +} + +static void run_layout_alias_scheduler_elision_checks() { + ggml_init_params params = {}; + params.mem_size = 256 * 1024; + params.no_alloc = true; + ggml_context * ctx = ggml_init(params); + REQUIRE(ctx != nullptr); + + ggml_tensor * a = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 8); + ggml_tensor * b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 8); + ggml_tensor * sum = ggml_add(ctx, a, b); + ggml_tensor * reshaped = ggml_reshape_2d(ctx, sum, 4, 2); + REQUIRE(a != nullptr); + REQUIRE(b != nullptr); + REQUIRE(sum != nullptr); + REQUIRE(reshaped != nullptr); + + ggml_cgraph * graph = ggml_new_graph(ctx); + REQUIRE(graph != nullptr); + ggml_build_forward_expand(graph, reshaped); + + ggml::hrx::GraphImportResult imported = ggml::hrx::import_ggml_graph(*graph); + REQUIRE(imported.valid()); + REQUIRE(imported.graph.nodes().size() == 2); + REQUIRE(imported.graph.nodes()[0].op == GGML_OP_ADD); + REQUIRE(imported.graph.nodes()[1].op == GGML_OP_RESHAPE); + + const ggml::hrx::Value * sum_value = imported.graph.values().find_tensor(sum); + const ggml::hrx::Value * reshaped_value = imported.graph.values().find_tensor(reshaped); + REQUIRE(sum_value != nullptr); + REQUIRE(reshaped_value != nullptr); + REQUIRE(sum_value->kind == ggml::hrx::ValueKind::Transient); + REQUIRE(reshaped_value->kind == ggml::hrx::ValueKind::External); + REQUIRE(imported.graph.values().same_storage(sum_value->id, reshaped_value->id)); + REQUIRE(reshaped_value->storage_root == sum_value->id); + REQUIRE(reshaped_value->alias_source == sum_value->id); + REQUIRE(ggml::hrx::is_layout_alias_node(imported.graph, imported.graph.nodes()[1])); + + ggml::hrx::DispatchScheduler scheduler; + REQUIRE(scheduler.schedule_graph(imported.graph, test_dispatch_target())); + REQUIRE(scheduler.plan().valid()); + REQUIRE(scheduler.plan().dispatches.size() == 1); + + const ggml::hrx::CommandProgram commands = ggml::hrx::build_command_program( + imported.graph, scheduler.plan(), ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); + REQUIRE(commands.valid()); + REQUIRE(commands.commands.size() == 1); + REQUIRE(commands.commands[0].bindings.size() == 3); + REQUIRE(commands.commands[0].bindings[2].value == sum_value->id); + REQUIRE(commands.commands[0].bindings[2].origin == ggml::hrx::CommandBindingOrigin::Transient); + REQUIRE(commands.transients.allocations.size() == 1); + REQUIRE(ggml::hrx::find_transient_allocation(commands.transients, sum_value->id) != nullptr); + REQUIRE(ggml::hrx::find_transient_allocation(commands.transients, reshaped_value->id) == nullptr); + REQUIRE(command_program_verifies(commands)); + + ggml_free(ctx); +} + +static void run_transient_import_checks() { + ggml_init_params params = {}; + params.mem_size = 256 * 1024; + params.no_alloc = true; + ggml_context * ctx = ggml_init(params); + REQUIRE(ctx != nullptr); + + ggml_tensor * a = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 8); + ggml_tensor * b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 8); + ggml_tensor * sum = ggml_add(ctx, a, b); + ggml_tensor * out = ggml_sqr(ctx, sum); + REQUIRE(a != nullptr); + REQUIRE(b != nullptr); + REQUIRE(sum != nullptr); + REQUIRE(out != nullptr); + + ggml_cgraph * graph = ggml_new_graph(ctx); + REQUIRE(graph != nullptr); + ggml_build_forward_expand(graph, out); + + ggml::hrx::GraphImportResult imported = ggml::hrx::import_ggml_graph(*graph); + REQUIRE(imported.valid()); + REQUIRE(imported.graph.nodes().size() == 2); + REQUIRE(imported.graph.nodes()[0].op == GGML_OP_ADD); + REQUIRE(imported.graph.nodes()[1].op == GGML_OP_SQR); + + const ggml::hrx::Value * a_value = imported.graph.values().find_tensor(a); + const ggml::hrx::Value * sum_value = imported.graph.values().find_tensor(sum); + const ggml::hrx::Value * out_value = imported.graph.values().find_tensor(out); + REQUIRE(a_value != nullptr); + REQUIRE(sum_value != nullptr); + REQUIRE(out_value != nullptr); + REQUIRE(a_value->kind == ggml::hrx::ValueKind::External); + REQUIRE(sum_value->kind == ggml::hrx::ValueKind::Transient); + REQUIRE(out_value->kind == ggml::hrx::ValueKind::External); + REQUIRE( + !imported.graph.values().bind_buffer(sum_value->id, { dummy_hrx_buffer(0x3000), 0, sum_value->byte_count })); + + ggml::hrx::DispatchScheduler scheduler; + REQUIRE(!scheduler.schedule_graph(imported.graph, test_dispatch_target())); + REQUIRE(!scheduler.plan().valid()); + REQUIRE(scheduler.plan().dispatches.empty()); + REQUIRE(status_contains(scheduler.plan().status, "unsupported HRX node 1")); + + ggml_free(ctx); +} + +static void run_chained_dispatch_requires_transients() { + ggml_init_params params = {}; + params.mem_size = 256 * 1024; + params.no_alloc = true; + ggml_context * ctx = ggml_init(params); + REQUIRE(ctx != nullptr); + + ggml_tensor * a = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 8); + ggml_tensor * b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 8); + ggml_tensor * c = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 8); + ggml_tensor * sum = ggml_add(ctx, a, b); + ggml_tensor * out = ggml_add(ctx, sum, c); + REQUIRE(a != nullptr); + REQUIRE(b != nullptr); + REQUIRE(c != nullptr); + REQUIRE(sum != nullptr); + REQUIRE(out != nullptr); + + ggml_cgraph * graph = ggml_new_graph(ctx); + REQUIRE(graph != nullptr); + ggml_build_forward_expand(graph, out); + + ggml::hrx::GraphImportResult imported = ggml::hrx::import_ggml_graph(*graph); + REQUIRE(imported.valid()); + REQUIRE(imported.graph.nodes().size() == 2); + REQUIRE(imported.graph.nodes()[0].op == GGML_OP_ADD); + REQUIRE(imported.graph.nodes()[1].op == GGML_OP_ADD); + + const ggml::hrx::Value * sum_value = imported.graph.values().find_tensor(sum); + REQUIRE(sum_value != nullptr); + REQUIRE(sum_value->kind == ggml::hrx::ValueKind::Transient); + + ggml::hrx::DispatchScheduler scheduler; + REQUIRE(scheduler.schedule_graph(imported.graph, test_dispatch_target())); + REQUIRE(scheduler.plan().valid()); + REQUIRE(scheduler.plan().dispatches.size() == 2); + + const ggml::hrx::CommandProgram commands = ggml::hrx::build_command_program( + imported.graph, scheduler.plan(), ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); + REQUIRE(commands.valid()); + REQUIRE(commands.commands.size() == 2); + REQUIRE(commands.commands[1].dependencies.size() == 1); + REQUIRE(commands.commands[1].dependencies[0] == 0); + REQUIRE(commands.commands[0].bindings.size() == 3); + REQUIRE(commands.commands[1].bindings.size() == 3); + REQUIRE(commands.commands[0].bindings[0].origin == ggml::hrx::CommandBindingOrigin::GraphValue); + REQUIRE(commands.commands[0].bindings[1].origin == ggml::hrx::CommandBindingOrigin::GraphValue); + REQUIRE(commands.commands[0].bindings[2].value == sum_value->id); + REQUIRE(commands.commands[0].bindings[2].origin == ggml::hrx::CommandBindingOrigin::Transient); + REQUIRE(commands.commands[1].bindings[0].value == sum_value->id); + REQUIRE(commands.commands[1].bindings[0].origin == ggml::hrx::CommandBindingOrigin::Transient); + REQUIRE(commands.commands[1].bindings[1].origin == ggml::hrx::CommandBindingOrigin::GraphValue); + REQUIRE(commands.commands[1].bindings[2].origin == ggml::hrx::CommandBindingOrigin::GraphValue); + REQUIRE(commands.transients.allocations.size() == 1); + const ggml::hrx::TransientAllocation * sum_allocation = + ggml::hrx::find_transient_allocation(commands.transients, sum_value->id); + REQUIRE(sum_allocation != nullptr); + REQUIRE(sum_allocation->value == sum_value->id); + REQUIRE(sum_allocation->size == sum_value->byte_count); + REQUIRE(sum_allocation->alignment == 256); + REQUIRE(sum_allocation->arena_offset == 0); + REQUIRE(commands.transients.arena_size == 256); + REQUIRE(command_program_verifies(commands)); + + const std::string transient_binding_text = ggml::hrx::format_command_binding(commands.commands[1].bindings[0]); + REQUIRE(string_contains(transient_binding_text, "origin=Transient")); + + bind_external_values(imported.graph.values()); + const ggml::hrx::CommandProgramBindings bindings = + ggml::hrx::CommandProgramBindings::from_value_map(imported.graph.values()); + REQUIRE(bindings.valid()); + REQUIRE(bindings.find(sum_value->id) == nullptr); + + const ggml::hrx::ResolvedCommandProgram resolved = ggml::hrx::resolve_command_program_bindings(commands, bindings); + REQUIRE(!resolved.valid()); + REQUIRE(status_contains(resolved.status, "no transient arena")); + REQUIRE(status_contains(resolved.status, "origin=Transient")); + REQUIRE(status_contains(resolved.status, "value=")); + + const ggml::hrx::TransientArenaAllocationRef transient_arena = { + dummy_hrx_buffer(0x8000), + commands.transients.arena_size, + 7, + }; + const ggml::hrx::ResolvedCommandProgram resolved_with_transients = + ggml::hrx::resolve_command_program_bindings(commands, bindings, &transient_arena); + REQUIRE(resolved_with_transients.valid()); + REQUIRE(resolved_with_transients.commands.size() == 2); + REQUIRE(resolved_with_transients.commands[0].bindings[2].ref.buffer == dummy_hrx_buffer(0x8000)); + REQUIRE(resolved_with_transients.commands[0].bindings[2].ref.offset == 0); + REQUIRE(resolved_with_transients.commands[0].bindings[2].ref.length == sum_value->byte_count); + REQUIRE(resolved_with_transients.commands[1].bindings[0].ref.buffer == dummy_hrx_buffer(0x8000)); + REQUIRE(resolved_with_transients.commands[1].bindings[0].ref.offset == 0); + REQUIRE(resolved_with_transients.commands[1].bindings[0].ref.length == sum_value->byte_count); + + ggml::hrx::PreparedCommandProgram prepared_shape; + for (const ggml::hrx::Command & prepared_source : commands.commands) { + ggml::hrx::PreparedCommand prepared_command; + prepared_command.ordinal = prepared_source.ordinal; + prepared_command.kind = prepared_source.kind; + prepared_command.kernel.specialization = prepared_source.kernel; + for (const ggml::hrx::CommandBinding & binding : prepared_source.bindings) { + prepared_command.kernel.bindings.push_back({ + binding, { dummy_hrx_buffer(0x4000), 123, binding.length } + }); + } + prepared_shape.commands.push_back(prepared_command); + } + prepared_shape.bound_transient_arena_allocation_id = 1; + + REQUIRE(ggml::hrx::bind_prepared_command_program_transients(commands, transient_arena, prepared_shape)); + REQUIRE(prepared_shape.bound_transient_arena_allocation_id == transient_arena.allocation_id); + REQUIRE(prepared_shape.commands[0].kernel.bindings[2].ref.buffer == dummy_hrx_buffer(0x8000)); + REQUIRE(prepared_shape.commands[0].kernel.bindings[2].ref.offset == 0); + REQUIRE(prepared_shape.commands[1].kernel.bindings[0].ref.buffer == dummy_hrx_buffer(0x8000)); + REQUIRE(prepared_shape.commands[1].kernel.bindings[0].ref.offset == 0); + + const ggml::hrx::TransientArenaAllocationRef grown_transient_arena = { + dummy_hrx_buffer(0x9000), + commands.transients.arena_size + 256, + 8, + }; + REQUIRE(ggml::hrx::bind_prepared_command_program_transients(commands, grown_transient_arena, prepared_shape)); + REQUIRE(prepared_shape.bound_transient_arena_allocation_id == grown_transient_arena.allocation_id); + REQUIRE(prepared_shape.commands[0].kernel.bindings[2].ref.buffer == dummy_hrx_buffer(0x9000)); + REQUIRE(prepared_shape.commands[1].kernel.bindings[0].ref.buffer == dummy_hrx_buffer(0x9000)); + + prepared_shape.commands[1].kernel.bindings[0].binding.origin = ggml::hrx::CommandBindingOrigin::ProgramConstant; + prepared_shape.commands[1].kernel.bindings[0].ref = { dummy_hrx_buffer(0xb000), 32, sum_value->byte_count }; + const ggml::hrx::TransientArenaAllocationRef rebinding_transient_arena = { + dummy_hrx_buffer(0xc000), + commands.transients.arena_size + 512, + 9, + }; + REQUIRE(ggml::hrx::bind_prepared_command_program_transients(commands, rebinding_transient_arena, prepared_shape)); + REQUIRE(prepared_shape.commands[0].kernel.bindings[2].ref.buffer == dummy_hrx_buffer(0xc000)); + REQUIRE(prepared_shape.commands[1].kernel.bindings[0].ref.buffer == dummy_hrx_buffer(0xb000)); + REQUIRE(prepared_shape.commands[1].kernel.bindings[0].ref.offset == 32); + + const ggml::hrx::TransientArenaAllocationRef invalid_transient_arena = { + dummy_hrx_buffer(0xa000), + commands.transients.arena_size, + ggml::hrx::kInvalidTransientArenaAllocationId, + }; + const ggml::hrx::ResolvedCommandProgram invalid_transient_resolved = + ggml::hrx::resolve_command_program_bindings(commands, bindings, &invalid_transient_arena); + REQUIRE(!invalid_transient_resolved.valid()); + REQUIRE(status_contains(invalid_transient_resolved.status, "no transient arena allocation id")); + + ggml::hrx::CommandProgram missing_allocation = copy_command_program_shape(commands); + missing_allocation.transients.allocations.clear(); + ggml::hrx::VerificationResult verification = + ggml::hrx::verify_command_program(missing_allocation, ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); + REQUIRE(!verification.valid()); + REQUIRE(status_contains(verification.status, "no transient allocation")); + + ggml::hrx::CommandProgram out_of_range = copy_command_program_shape(commands); + out_of_range.commands[0].bindings[2].length = sum_value->byte_count + 1; + verification = ggml::hrx::verify_command_program(out_of_range, ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); + REQUIRE(!verification.valid()); + REQUIRE(status_contains(verification.status, "outside transient allocation length")); + + ggml_free(ctx); +} + +static void run_graph_replay_host_staging_is_not_ineligible() { + ggml::hrx::CommandProgram commands; + + std::vector source0(64, 1); + std::vector source1(64, 2); + + ggml::hrx::PreparedCommandProgram prepared; + ggml::hrx::HostStagingBuffer staging; + staging.buffer = dummy_hrx_buffer(0x1000); + staging.host_data = source0.data(); + staging.value = 7; + staging.length = source0.size(); + staging.upload = true; + prepared.host_staging.push_back(std::move(staging)); + + ggml::hrx::CommandProgramBinding live_binding; + live_binding.value = ggml::hrx::ValueId(7); + live_binding.length = source1.size(); + live_binding.capacity = source1.size(); + live_binding.host_data = source1.data(); + const ggml::hrx::CommandProgramBindings bindings = + ggml::hrx::CommandProgramBindings::from_bindings({ live_binding }); + REQUIRE(bindings.valid()); + + ggml::hrx::RecordedCommandGraph recorded; + recorded.exec = dummy_hrx_graph_exec(0x2000); + recorded.bound_transient_arena_allocation_id = ggml::hrx::kInvalidTransientArenaAllocationId; + + ggml::hrx::CommandProgramExecutionContext context; + context.stream = dummy_hrx_stream(0x3000); + + const ggml::hrx::RecordedCommandGraphExecutionResult result = + ggml::hrx::bind_and_launch_recorded_command_graph(context, commands, bindings, prepared, recorded); + recorded.exec = nullptr; + + REQUIRE(!result.success); + REQUIRE(result.event == ggml::hrx::HrxGraphReplayEvent::LaunchFailed); + REQUIRE(result.ineligible_reason.empty()); + REQUIRE(status_contains(result.status, "missing HRX host transfer manager")); + REQUIRE(prepared.host_staging.size() == 1); + REQUIRE(prepared.host_staging[0].buffer == dummy_hrx_buffer(0x1000)); + REQUIRE(prepared.host_staging[0].host_data == source1.data()); + prepared.host_staging[0].buffer = nullptr; +} + +static void run_multiple_transient_plan_checks() { + ggml_init_params params = {}; + params.mem_size = 256 * 1024; + params.no_alloc = true; + ggml_context * ctx = ggml_init(params); + REQUIRE(ctx != nullptr); + + ggml_tensor * a = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 8); + ggml_tensor * b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 8); + ggml_tensor * c = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 8); + ggml_tensor * d = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 8); + ggml_tensor * sum0 = ggml_add(ctx, a, b); + ggml_tensor * sum1 = ggml_add(ctx, c, d); + ggml_tensor * out = ggml_add(ctx, sum0, sum1); + REQUIRE(a != nullptr); + REQUIRE(b != nullptr); + REQUIRE(c != nullptr); + REQUIRE(d != nullptr); + REQUIRE(sum0 != nullptr); + REQUIRE(sum1 != nullptr); + REQUIRE(out != nullptr); + + ggml_cgraph * graph = ggml_new_graph(ctx); + REQUIRE(graph != nullptr); + ggml_build_forward_expand(graph, out); + + ggml::hrx::GraphImportResult imported = ggml::hrx::import_ggml_graph(*graph); + REQUIRE(imported.valid()); + REQUIRE(imported.graph.nodes().size() == 3); + + const ggml::hrx::Value * sum0_value = imported.graph.values().find_tensor(sum0); + const ggml::hrx::Value * sum1_value = imported.graph.values().find_tensor(sum1); + REQUIRE(sum0_value != nullptr); + REQUIRE(sum1_value != nullptr); + REQUIRE(sum0_value->kind == ggml::hrx::ValueKind::Transient); + REQUIRE(sum1_value->kind == ggml::hrx::ValueKind::Transient); + + ggml::hrx::DispatchScheduler scheduler; + REQUIRE(scheduler.schedule_graph(imported.graph, test_dispatch_target())); + REQUIRE(scheduler.plan().valid()); + REQUIRE(scheduler.plan().dispatches.size() == 3); + + const ggml::hrx::CommandProgram commands = ggml::hrx::build_command_program( + imported.graph, scheduler.plan(), ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); + REQUIRE(commands.valid()); + REQUIRE(commands.transients.allocations.size() == 2); + REQUIRE(commands.transients.arena_size == 512); + const ggml::hrx::TransientAllocation * sum0_allocation = + ggml::hrx::find_transient_allocation(commands.transients, sum0_value->id); + const ggml::hrx::TransientAllocation * sum1_allocation = + ggml::hrx::find_transient_allocation(commands.transients, sum1_value->id); + REQUIRE(sum0_allocation != nullptr); + REQUIRE(sum1_allocation != nullptr); + REQUIRE(sum0_allocation->arena_offset != sum1_allocation->arena_offset); + REQUIRE(sum0_allocation->arena_offset % 256 == 0); + REQUIRE(sum1_allocation->arena_offset % 256 == 0); + + ggml::hrx::CommandProgram overlapping_live_transients = copy_command_program_shape(commands); + ggml::hrx::TransientAllocation * overlapping_sum0 = nullptr; + ggml::hrx::TransientAllocation * overlapping_sum1 = nullptr; + for (ggml::hrx::TransientAllocation & allocation : overlapping_live_transients.transients.allocations) { + if (allocation.value == sum0_value->id) { + overlapping_sum0 = &allocation; + } else if (allocation.value == sum1_value->id) { + overlapping_sum1 = &allocation; + } + } + REQUIRE(overlapping_sum0 != nullptr); + REQUIRE(overlapping_sum1 != nullptr); + overlapping_sum1->arena_offset = overlapping_sum0->arena_offset; + ggml::hrx::VerificationResult verification = + ggml::hrx::verify_command_program(overlapping_live_transients, ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); + REQUIRE(!verification.valid()); + REQUIRE(status_contains(verification.status, "transient allocations overlap")); + + ggml_free(ctx); +} + +static void run_disjoint_transient_plan_packing_checks() { + ggml_init_params params = {}; + params.mem_size = 256 * 1024; + params.no_alloc = true; + ggml_context * ctx = ggml_init(params); + REQUIRE(ctx != nullptr); + + ggml_tensor * a = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 8); + ggml_tensor * b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 8); + ggml_tensor * c = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 8); + ggml_tensor * d = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 8); + ggml_tensor * e = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 8); + ggml_tensor * f = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 8); + ggml_tensor * sum0 = ggml_add(ctx, a, b); + ggml_tensor * out0 = ggml_add(ctx, sum0, c); + ggml_tensor * sum1 = ggml_add(ctx, d, e); + ggml_tensor * out1 = ggml_add(ctx, sum1, f); + REQUIRE(a != nullptr); + REQUIRE(b != nullptr); + REQUIRE(c != nullptr); + REQUIRE(d != nullptr); + REQUIRE(e != nullptr); + REQUIRE(f != nullptr); + REQUIRE(sum0 != nullptr); + REQUIRE(out0 != nullptr); + REQUIRE(sum1 != nullptr); + REQUIRE(out1 != nullptr); + + ggml_cgraph * graph = ggml_new_graph(ctx); + REQUIRE(graph != nullptr); + ggml_build_forward_expand(graph, out0); + ggml_build_forward_expand(graph, out1); + + ggml::hrx::GraphImportResult imported = ggml::hrx::import_ggml_graph(*graph); + REQUIRE(imported.valid()); + REQUIRE(imported.graph.nodes().size() == 4); + + const ggml::hrx::Value * sum0_value = imported.graph.values().find_tensor(sum0); + const ggml::hrx::Value * sum1_value = imported.graph.values().find_tensor(sum1); + REQUIRE(sum0_value != nullptr); + REQUIRE(sum1_value != nullptr); + REQUIRE(sum0_value->kind == ggml::hrx::ValueKind::Transient); + REQUIRE(sum1_value->kind == ggml::hrx::ValueKind::Transient); + + ggml::hrx::DispatchScheduler scheduler; + REQUIRE(scheduler.schedule_graph(imported.graph, test_dispatch_target())); + REQUIRE(scheduler.plan().valid()); + REQUIRE(scheduler.plan().dispatches.size() == 4); + + const ggml::hrx::CommandProgram commands = ggml::hrx::build_command_program( + imported.graph, scheduler.plan(), ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); + REQUIRE(commands.valid()); + REQUIRE(commands.transients.allocations.size() == 2); + REQUIRE(commands.transients.arena_size == 256); + const ggml::hrx::TransientAllocation * sum0_allocation = + ggml::hrx::find_transient_allocation(commands.transients, sum0_value->id); + const ggml::hrx::TransientAllocation * sum1_allocation = + ggml::hrx::find_transient_allocation(commands.transients, sum1_value->id); + REQUIRE(sum0_allocation != nullptr); + REQUIRE(sum1_allocation != nullptr); + REQUIRE(sum0_allocation->arena_offset == sum1_allocation->arena_offset); + REQUIRE(sum0_allocation->arena_offset % 256 == 0); + REQUIRE(command_program_verifies(commands)); + + ggml_free(ctx); +} + +static void run_graph_program_cache_uid_mismatch_checks() { + ggml::hrx::GraphProgramCache cache; + const ggml::hrx::KernelCorpus & corpus = ggml::hrx::get_qwen_kernel_corpus(); + + ggml_init_params params = {}; + params.mem_size = 512 * 1024; + params.no_alloc = true; + ggml_context * ctx = ggml_init(params); + REQUIRE(ctx != nullptr); + + ggml_tensor * a = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 8); + ggml_tensor * b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 8); + ggml_tensor * c = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 8); + ggml_tensor * d = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 8); + ggml_tensor * out0 = ggml_add(ctx, a, b); + ggml_tensor * out1 = ggml_add(ctx, c, d); + REQUIRE(a != nullptr); + REQUIRE(b != nullptr); + REQUIRE(c != nullptr); + REQUIRE(d != nullptr); + REQUIRE(out0 != nullptr); + REQUIRE(out1 != nullptr); + + ggml_cgraph * graph0 = ggml_new_graph(ctx); + REQUIRE(graph0 != nullptr); + ggml_build_forward_expand(graph0, out0); + graph0->uid = 3001; + + ggml::hrx::GraphProgramLookup lookup = cache.get_or_build(*graph0, corpus, "gfx1151"); + REQUIRE(lookup.valid()); + ggml::hrx::GraphProgramCacheStats stats = cache.stats(); + REQUIRE(stats.builds == 1); + REQUIRE(stats.hits == 0); + + lookup = cache.get_or_build(*graph0, corpus, "gfx1151"); + REQUIRE(lookup.valid()); + stats = cache.stats(); + REQUIRE(stats.builds == 1); + REQUIRE(stats.hits == 1); + + ggml_cgraph * graph1 = ggml_new_graph(ctx); + REQUIRE(graph1 != nullptr); + ggml_build_forward_expand(graph1, out0); + ggml_build_forward_expand(graph1, out1); + graph1->uid = 3001; + + lookup = cache.get_or_build(*graph1, corpus, "gfx1151"); + REQUIRE(lookup.valid()); + stats = cache.stats(); + REQUIRE(stats.builds == 2); + REQUIRE(stats.hits == 1); + + ggml_tensor * unsupported = ggml_sqr(ctx, a); + REQUIRE(unsupported != nullptr); + ggml_cgraph * graph2 = ggml_new_graph(ctx); + REQUIRE(graph2 != nullptr); + ggml_build_forward_expand(graph2, unsupported); + graph2->uid = 3001; + + lookup = cache.get_or_build(*graph2, corpus, "gfx1151"); + REQUIRE(!lookup.valid()); + stats = cache.stats(); + REQUIRE(stats.builds == 2); + REQUIRE(stats.hits == 1); + + ggml::hrx::GraphProgramCache alias_cache; + ggml_tensor * alias_sum = ggml_add(ctx, a, b); + ggml_tensor * view0 = ggml_view_1d(ctx, alias_sum, 4, 0); + ggml_tensor * view1 = ggml_view_1d(ctx, alias_sum, 4, sizeof(float)); + REQUIRE(alias_sum != nullptr); + REQUIRE(view0 != nullptr); + REQUIRE(view1 != nullptr); + + ggml_cgraph * alias_graph0 = ggml_new_graph(ctx); + REQUIRE(alias_graph0 != nullptr); + ggml_build_forward_expand(alias_graph0, view0); + alias_graph0->uid = 3002; + lookup = alias_cache.get_or_build(*alias_graph0, corpus, "gfx1151"); + REQUIRE(lookup.valid()); + stats = alias_cache.stats(); + REQUIRE(stats.builds == 1); + REQUIRE(stats.hits == 0); + + ggml_cgraph * alias_graph1 = ggml_new_graph(ctx); + REQUIRE(alias_graph1 != nullptr); + ggml_build_forward_expand(alias_graph1, view1); + alias_graph1->uid = 3002; + lookup = alias_cache.get_or_build(*alias_graph1, corpus, "gfx1151"); + REQUIRE(lookup.valid()); + stats = alias_cache.stats(); + REQUIRE(stats.builds == 1); + REQUIRE(stats.hits == 1); + + ggml_free(ctx); +} + +static void run_graph_executor_contract_checks() { + ggml_backend_hrx_device_context device_context = {}; + ggml_backend_hrx_context backend_context = {}; + device_context.architecture = "gfx1151"; + backend_context.device = &device_context; + const ggml::hrx::GraphExecutor executor(backend_context); + + ggml_init_params params = {}; + params.mem_size = 256 * 1024; + params.no_alloc = true; + ggml_context * ctx = ggml_init(params); + REQUIRE(ctx != nullptr); + + ggml_tensor * a = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 8); + ggml_tensor * b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 8); + ggml_tensor * add_out = ggml_add(ctx, a, b); + ggml_tensor * sqr_out = ggml_sqr(ctx, a); + REQUIRE(a != nullptr); + REQUIRE(b != nullptr); + REQUIRE(add_out != nullptr); + REQUIRE(sqr_out != nullptr); + + ggml_cgraph * add_graph = ggml_new_graph(ctx); + REQUIRE(add_graph != nullptr); + ggml_build_forward_expand(add_graph, add_out); + const ggml::hrx::GraphSupportResult add_support = executor.can_execute(*add_graph); + REQUIRE(add_support.supported); + REQUIRE(add_support.status.success()); + + const ggml::hrx::GraphExecutionResult missing_binding = executor.execute(*add_graph); + REQUIRE(!missing_binding.success()); + REQUIRE(missing_binding.code == GGML_STATUS_FAILED); + REQUIRE(status_contains(missing_binding.status, "external value")); + REQUIRE(status_contains(missing_binding.status, "not bound")); + + ggml_cgraph * sqr_graph = ggml_new_graph(ctx); + REQUIRE(sqr_graph != nullptr); + ggml_build_forward_expand(sqr_graph, sqr_out); + const ggml::hrx::GraphSupportResult sqr_support = executor.can_execute(*sqr_graph); + REQUIRE(!sqr_support.supported); + REQUIRE(status_contains(sqr_support.status, "unsupported HRX node 0")); + REQUIRE(status_contains(sqr_support.status, "SQR")); + + ggml_free(ctx); +} + +static std::vector read_transient_i32(ggml_backend_hrx_context * context, + ggml::hrx::TransientArenaAllocationRef arena, + const ggml::hrx::TransientAllocation & allocation) { + REQUIRE(context != nullptr); + REQUIRE(arena.buffer != nullptr); + REQUIRE(allocation.size % sizeof(int32_t) == 0); + std::vector data(allocation.size / sizeof(int32_t)); + require_hrx_status(hrx_synchronous_d2h(context->device->device, arena.buffer, allocation.arena_offset, data.data(), + allocation.size)); + return data; +} + +static void write_transient_i32_value(ggml_backend_hrx_context * context, + ggml::hrx::TransientArenaAllocationRef arena, + const ggml::hrx::TransientAllocation & allocation, + int32_t value) { + REQUIRE(context != nullptr); + REQUIRE(arena.buffer != nullptr); + REQUIRE(allocation.size >= sizeof(value)); + require_hrx_status( + hrx_synchronous_h2d(context->device->device, &value, arena.buffer, allocation.arena_offset, sizeof(value))); +} + +static void run_qwen_expert_table_partition_prefill_512_execution() { + ggml_backend_t backend = ggml_backend_hrx_init(0); + REQUIRE(backend != nullptr); + auto * backend_context = static_cast(backend->context); + REQUIRE(backend_context != nullptr); + + ggml_init_params params = {}; + params.mem_size = 256 * 1024; + params.no_alloc = true; + ggml_context * ctx = ggml_init(params); + REQUIRE(ctx != nullptr); + + constexpr int64_t token_count = 512; + constexpr int64_t route_count = 8; + constexpr int64_t route_stride = 8; + constexpr int64_t expert_count = 128; + + ggml_tensor * route_ids_tensor = ggml_new_tensor_2d(ctx, GGML_TYPE_I32, route_stride, token_count); + REQUIRE(route_ids_tensor != nullptr); + ggml_backend_buffer_t buffer = ggml_backend_alloc_ctx_tensors(ctx, backend); + REQUIRE(buffer != nullptr); + + const std::vector route_ids = + make_qwen_route_ids_iota(token_count, route_count, route_stride, expert_count); + const std::vector expected_expert_table = + make_qwen_expert_table_reference(route_ids, token_count, route_count, route_stride, expert_count); + const std::vector expected_partition_table = + make_qwen_partition_table_reference(expected_expert_table, token_count, route_count, expert_count); + ggml_backend_tensor_set(route_ids_tensor, route_ids.data(), 0, route_ids.size() * sizeof(int32_t)); + ggml_backend_synchronize(backend); + + ggml::hrx::Graph graph; + const ggml::hrx::ValueId route_ids_value = + graph.values().get_or_add_tensor_value(route_ids_tensor, ggml::hrx::ValueKind::External); + ggml::hrx::ValueBufferBinding route_ids_binding; + REQUIRE(ggml_backend_hrx_resolve_value_buffer(route_ids_tensor, route_ids_binding)); + REQUIRE(graph.values().bind_buffer(route_ids_value, route_ids_binding)); + + const ggml::hrx::ValueId expert_table_value(static_cast(graph.values().size())); + const ggml::hrx::ValueId partition_table_value(expert_table_value.value + 1); + const ggml::hrx::ValueId completion_counter_value(expert_table_value.value + 2); + ggml::hrx::CommandPlan plan; + plan.transients.push_back( + { expert_table_value, "qwen.test.expert_table", qwen_expert_table_size(token_count, expert_count), 256 }); + plan.transients.push_back({ partition_table_value, "qwen.test.partition_table", + qwen_partition_table_size(token_count, route_count, expert_count), 256 }); + plan.completion_counter_requests.push_back({ completion_counter_value, "qwen.test.completion_counter", 1 }); + ggml::hrx::Dispatch dispatch; + dispatch.kernel = ggml::hrx::make_kernel_specialization( + ggml::hrx::kernel_catalog_ref("qwen3_moe", "qwen3_moe_build_expert_table_partition_prefill_512")); + dispatch.kernel.integer_parameters.emplace("token_count", token_count); + dispatch.kernel.integer_parameters.emplace("route_count", route_count); + dispatch.kernel.integer_parameters.emplace("route_stride", route_stride); + dispatch.kernel.integer_parameters.emplace("expert_count", expert_count); + dispatch.bindings.push_back({ route_ids_value, 0, route_ids.size() * sizeof(int32_t) }); + dispatch.bindings.push_back({ expert_table_value, 0, qwen_expert_table_size(token_count, expert_count) }); + dispatch.bindings.push_back( + { partition_table_value, 0, qwen_partition_table_size(token_count, route_count, expert_count) }); + dispatch.bindings.push_back({ completion_counter_value, 0, sizeof(int32_t) }); + plan.dispatches.push_back(std::move(dispatch)); + + const ggml::hrx::KernelCorpus & corpus = ggml::hrx::get_qwen_kernel_corpus(); + const char * target = backend_context->device->architecture.c_str(); + const ggml::hrx::CommandProgram commands = ggml::hrx::build_command_program(graph, plan, corpus, target); + REQUIRE(commands.valid()); + REQUIRE(commands.commands.size() == 1); + REQUIRE(commands.completion_counters.count == 1); + REQUIRE(commands.completion_counters.byte_count == sizeof(int32_t)); + REQUIRE(commands.transients.allocations.size() == 3); + REQUIRE(ggml::hrx::verify_command_program(commands, corpus, target).valid()); + + const ggml::hrx::TransientAllocation * expert_table_allocation = + ggml::hrx::find_transient_allocation(commands.transients, expert_table_value); + const ggml::hrx::TransientAllocation * partition_table_allocation = + ggml::hrx::find_transient_allocation(commands.transients, partition_table_value); + const ggml::hrx::TransientAllocation * completion_counter_allocation = + ggml::hrx::find_transient_allocation(commands.transients, completion_counter_value); + REQUIRE(expert_table_allocation != nullptr); + REQUIRE(partition_table_allocation != nullptr); + REQUIRE(completion_counter_allocation != nullptr); + REQUIRE(completion_counter_allocation->arena_offset == commands.completion_counters.arena_offset); + + const ggml::hrx::CommandProgramBindings bindings = + ggml::hrx::CommandProgramBindings::from_value_map(graph.values()); + REQUIRE(bindings.valid()); + const ggml::hrx::CommandProgramExecutionContext execution_context = { + backend_context->device->device, + backend_context->stream, + target, + &corpus, + &backend_context->kernel_executables, + &backend_context->transient_arena, + &backend_context->host_transfers, + &backend_context->host_weights, + }; + + REQUIRE(ggml::hrx::execute_command_program(execution_context, commands, bindings)); + ggml_backend_synchronize(backend); + ggml::hrx::TransientArenaAllocationRef arena = backend_context->transient_arena.current_allocation(); + require_qwen_expert_table_matches(read_transient_i32(backend_context, arena, *expert_table_allocation), + expected_expert_table, token_count, expert_count); + require_qwen_partition_table_matches(read_transient_i32(backend_context, arena, *partition_table_allocation), + expected_partition_table); + REQUIRE(read_transient_i32(backend_context, arena, *completion_counter_allocation)[0] == 0); + + write_transient_i32_value(backend_context, arena, *completion_counter_allocation, 17); + REQUIRE(ggml::hrx::execute_command_program(execution_context, commands, bindings)); + ggml_backend_synchronize(backend); + arena = backend_context->transient_arena.current_allocation(); + require_qwen_expert_table_matches(read_transient_i32(backend_context, arena, *expert_table_allocation), + expected_expert_table, token_count, expert_count); + require_qwen_partition_table_matches(read_transient_i32(backend_context, arena, *partition_table_allocation), + expected_partition_table); + REQUIRE(read_transient_i32(backend_context, arena, *completion_counter_allocation)[0] == 0); + + ggml_backend_buffer_free(buffer); + ggml_free(ctx); + ggml_backend_free(backend); +} + +static void run_add_f32() { + ggml_backend_t backend = ggml_backend_hrx_init(0); + REQUIRE(backend != nullptr); + + ggml_init_params params = {}; + params.mem_size = 256 * 1024; + params.no_alloc = true; + ggml_context * ctx = ggml_init(params); + REQUIRE(ctx != nullptr); + + constexpr int64_t element_count = 1024; + ggml_tensor * a = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, element_count); + ggml_tensor * b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, element_count); + ggml_tensor * out = ggml_add(ctx, a, b); + REQUIRE(a != nullptr); + REQUIRE(b != nullptr); + REQUIRE(out != nullptr); + + ggml_cgraph * graph = ggml_new_graph(ctx); + REQUIRE(graph != nullptr); + ggml_build_forward_expand(graph, out); + + ggml_backend_buffer_t buffer = ggml_backend_alloc_ctx_tensors(ctx, backend); + REQUIRE(buffer != nullptr); + + std::vector a_data(element_count); + std::vector b_data(element_count); + std::vector expected(element_count); + for (int64_t i = 0; i < element_count; ++i) { + a_data[i] = static_cast(i % 17) * 0.25f - 2.0f; + b_data[i] = static_cast(i % 13) * -0.5f + 3.0f; + expected[i] = a_data[i] + b_data[i]; + } + + ggml_backend_tensor_set(a, a_data.data(), 0, a_data.size() * sizeof(float)); + ggml_backend_tensor_set(b, b_data.data(), 0, b_data.size() * sizeof(float)); + + REQUIRE(ggml_backend_graph_compute(backend, graph) == GGML_STATUS_SUCCESS); + ggml_backend_synchronize(backend); + + std::vector actual(element_count); + ggml_backend_tensor_get(out, actual.data(), 0, actual.size() * sizeof(float)); + for (int64_t i = 0; i < element_count; ++i) { + REQUIRE(actual[i] == expected[i]); + } + + ggml_backend_buffer_free(buffer); + ggml_free(ctx); + ggml_backend_free(backend); +} + +static void run_two_independent_add_f32() { + ggml_backend_t backend = ggml_backend_hrx_init(0); + REQUIRE(backend != nullptr); + + ggml_init_params params = {}; + params.mem_size = 256 * 1024; + params.no_alloc = true; + ggml_context * ctx = ggml_init(params); + REQUIRE(ctx != nullptr); + + constexpr int64_t element_count = 1024; + ggml_tensor * a = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, element_count); + ggml_tensor * b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, element_count); + ggml_tensor * c = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, element_count); + ggml_tensor * d = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, element_count); + ggml_tensor * out0 = ggml_add(ctx, a, b); + ggml_tensor * out1 = ggml_add(ctx, c, d); + REQUIRE(a != nullptr); + REQUIRE(b != nullptr); + REQUIRE(c != nullptr); + REQUIRE(d != nullptr); + REQUIRE(out0 != nullptr); + REQUIRE(out1 != nullptr); + + ggml_cgraph * graph = ggml_new_graph(ctx); + REQUIRE(graph != nullptr); + ggml_build_forward_expand(graph, out0); + ggml_build_forward_expand(graph, out1); + graph->uid = 1002; + + ggml_backend_buffer_t buffer = ggml_backend_alloc_ctx_tensors(ctx, backend); + REQUIRE(buffer != nullptr); + + std::vector a_data(element_count); + std::vector b_data(element_count); + std::vector c_data(element_count); + std::vector d_data(element_count); + std::vector expected0(element_count); + std::vector expected1(element_count); + for (int64_t i = 0; i < element_count; ++i) { + a_data[i] = static_cast(i % 17) * 0.25f - 2.0f; + b_data[i] = static_cast(i % 13) * -0.5f + 3.0f; + c_data[i] = static_cast(i % 19) * 0.125f + 1.0f; + d_data[i] = static_cast(i % 11) * 0.75f - 4.0f; + expected0[i] = a_data[i] + b_data[i]; + expected1[i] = c_data[i] + d_data[i]; + } + + ggml_backend_tensor_set(a, a_data.data(), 0, a_data.size() * sizeof(float)); + ggml_backend_tensor_set(b, b_data.data(), 0, b_data.size() * sizeof(float)); + ggml_backend_tensor_set(c, c_data.data(), 0, c_data.size() * sizeof(float)); + ggml_backend_tensor_set(d, d_data.data(), 0, d_data.size() * sizeof(float)); + + ggml_backend_hrx_cache_stats cache_stats = {}; + REQUIRE(ggml_backend_hrx_get_cache_stats(backend, &cache_stats)); + REQUIRE(cache_stats.graph_program_builds == 0); + REQUIRE(cache_stats.graph_program_hits == 0); + REQUIRE(cache_stats.prepared_program_builds == 0); + REQUIRE(cache_stats.prepared_program_hits == 0); + + REQUIRE(ggml_backend_graph_compute(backend, graph) == GGML_STATUS_SUCCESS); + ggml_backend_synchronize(backend); + + REQUIRE(ggml_backend_hrx_get_cache_stats(backend, &cache_stats)); + REQUIRE(cache_stats.graph_program_builds == 1); + REQUIRE(cache_stats.graph_program_hits == 0); + REQUIRE(cache_stats.prepared_program_builds == 1); + REQUIRE(cache_stats.prepared_program_hits == 0); + + std::vector actual0(element_count); + std::vector actual1(element_count); + ggml_backend_tensor_get(out0, actual0.data(), 0, actual0.size() * sizeof(float)); + ggml_backend_tensor_get(out1, actual1.data(), 0, actual1.size() * sizeof(float)); + for (int64_t i = 0; i < element_count; ++i) { + REQUIRE(actual0[i] == expected0[i]); + REQUIRE(actual1[i] == expected1[i]); + } + + for (int64_t i = 0; i < element_count; ++i) { + a_data[i] = static_cast(i % 23) * -0.25f + 5.0f; + b_data[i] = static_cast(i % 7) * 0.5f - 1.0f; + c_data[i] = static_cast(i % 5) * -0.125f + 2.0f; + d_data[i] = static_cast(i % 29) * 0.75f - 6.0f; + expected0[i] = a_data[i] + b_data[i]; + expected1[i] = c_data[i] + d_data[i]; + } + ggml_backend_tensor_set(a, a_data.data(), 0, a_data.size() * sizeof(float)); + ggml_backend_tensor_set(b, b_data.data(), 0, b_data.size() * sizeof(float)); + ggml_backend_tensor_set(c, c_data.data(), 0, c_data.size() * sizeof(float)); + ggml_backend_tensor_set(d, d_data.data(), 0, d_data.size() * sizeof(float)); + + REQUIRE(ggml_backend_graph_compute(backend, graph) == GGML_STATUS_SUCCESS); + ggml_backend_synchronize(backend); + + REQUIRE(ggml_backend_hrx_get_cache_stats(backend, &cache_stats)); + REQUIRE(cache_stats.graph_program_builds == 1); + REQUIRE(cache_stats.graph_program_hits == 1); + REQUIRE(cache_stats.prepared_program_builds == 1); + REQUIRE(cache_stats.prepared_program_hits == 1); + + ggml_backend_tensor_get(out0, actual0.data(), 0, actual0.size() * sizeof(float)); + ggml_backend_tensor_get(out1, actual1.data(), 0, actual1.size() * sizeof(float)); + for (int64_t i = 0; i < element_count; ++i) { + REQUIRE(actual0[i] == expected0[i]); + REQUIRE(actual1[i] == expected1[i]); + } + + ggml_backend_buffer_free(buffer); + ggml_free(ctx); + ggml_backend_free(backend); +} + +static void run_chained_add_f32() { + ggml_backend_t backend = ggml_backend_hrx_init(0); + REQUIRE(backend != nullptr); + + ggml_init_params params = {}; + params.mem_size = 256 * 1024; + params.no_alloc = true; + ggml_context * ctx = ggml_init(params); + REQUIRE(ctx != nullptr); + + constexpr int64_t element_count = 1024; + ggml_tensor * a = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, element_count); + ggml_tensor * b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, element_count); + ggml_tensor * c = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, element_count); + ggml_tensor * sum = ggml_add(ctx, a, b); + ggml_tensor * out = ggml_add(ctx, sum, c); + REQUIRE(a != nullptr); + REQUIRE(b != nullptr); + REQUIRE(c != nullptr); + REQUIRE(sum != nullptr); + REQUIRE(out != nullptr); + + ggml_cgraph * graph = ggml_new_graph(ctx); + REQUIRE(graph != nullptr); + ggml_build_forward_expand(graph, out); + graph->uid = 1004; + + ggml_backend_buffer_t buffer = ggml_backend_alloc_ctx_tensors(ctx, backend); + REQUIRE(buffer != nullptr); + + std::vector a_data(element_count); + std::vector b_data(element_count); + std::vector c_data(element_count); + std::vector expected(element_count); + for (int64_t i = 0; i < element_count; ++i) { + a_data[i] = static_cast(i % 17) * 0.25f - 2.0f; + b_data[i] = static_cast(i % 13) * -0.5f + 3.0f; + c_data[i] = static_cast(i % 7) * 0.125f + 1.0f; + expected[i] = a_data[i] + b_data[i] + c_data[i]; + } + + ggml_backend_tensor_set(a, a_data.data(), 0, a_data.size() * sizeof(float)); + ggml_backend_tensor_set(b, b_data.data(), 0, b_data.size() * sizeof(float)); + ggml_backend_tensor_set(c, c_data.data(), 0, c_data.size() * sizeof(float)); + + REQUIRE(ggml_backend_graph_compute(backend, graph) == GGML_STATUS_SUCCESS); + ggml_backend_synchronize(backend); + + std::vector actual(element_count); + ggml_backend_tensor_get(out, actual.data(), 0, actual.size() * sizeof(float)); + for (int64_t i = 0; i < element_count; ++i) { + REQUIRE(actual[i] == expected[i]); + } + + ggml_backend_hrx_cache_stats cache_stats = {}; + REQUIRE(ggml_backend_hrx_get_cache_stats(backend, &cache_stats)); + REQUIRE(cache_stats.graph_program_builds == 1); + REQUIRE(cache_stats.prepared_program_builds == 1); + + for (int64_t i = 0; i < element_count; ++i) { + a_data[i] = static_cast(i % 11) * -0.25f + 5.0f; + b_data[i] = static_cast(i % 5) * 0.5f - 1.0f; + c_data[i] = static_cast(i % 19) * 0.75f - 6.0f; + expected[i] = a_data[i] + b_data[i] + c_data[i]; + } + ggml_backend_tensor_set(a, a_data.data(), 0, a_data.size() * sizeof(float)); + ggml_backend_tensor_set(b, b_data.data(), 0, b_data.size() * sizeof(float)); + ggml_backend_tensor_set(c, c_data.data(), 0, c_data.size() * sizeof(float)); + + REQUIRE(ggml_backend_graph_compute(backend, graph) == GGML_STATUS_SUCCESS); + ggml_backend_synchronize(backend); + + REQUIRE(ggml_backend_hrx_get_cache_stats(backend, &cache_stats)); + REQUIRE(cache_stats.graph_program_hits == 1); + REQUIRE(cache_stats.prepared_program_hits == 1); + + ggml_backend_tensor_get(out, actual.data(), 0, actual.size() * sizeof(float)); + for (int64_t i = 0; i < element_count; ++i) { + REQUIRE(actual[i] == expected[i]); + } + + ggml_backend_buffer_free(buffer); + ggml_free(ctx); + ggml_backend_free(backend); +} + +static void run_same_uid_distinct_graph_reuses_graph_program() { + ggml_backend_t backend = ggml_backend_hrx_init(0); + REQUIRE(backend != nullptr); + + constexpr int64_t element_count = 1024; + std::vector a_data(element_count); + std::vector b_data(element_count); + std::vector expected(element_count); + + ggml_init_params params0 = {}; + params0.mem_size = 256 * 1024; + params0.no_alloc = true; + ggml_context * ctx0 = ggml_init(params0); + REQUIRE(ctx0 != nullptr); + + ggml_tensor * a0 = ggml_new_tensor_1d(ctx0, GGML_TYPE_F32, element_count); + ggml_tensor * b0 = ggml_new_tensor_1d(ctx0, GGML_TYPE_F32, element_count); + ggml_tensor * out0 = ggml_add(ctx0, a0, b0); + REQUIRE(a0 != nullptr); + REQUIRE(b0 != nullptr); + REQUIRE(out0 != nullptr); + + ggml_cgraph * graph0 = ggml_new_graph(ctx0); + REQUIRE(graph0 != nullptr); + ggml_build_forward_expand(graph0, out0); + graph0->uid = 1003; + + ggml_backend_buffer_t buffer0 = ggml_backend_alloc_ctx_tensors(ctx0, backend); + REQUIRE(buffer0 != nullptr); + + for (int64_t i = 0; i < element_count; ++i) { + a_data[i] = static_cast(i % 17) * 0.25f - 2.0f; + b_data[i] = static_cast(i % 13) * -0.5f + 3.0f; + expected[i] = a_data[i] + b_data[i]; + } + ggml_backend_tensor_set(a0, a_data.data(), 0, a_data.size() * sizeof(float)); + ggml_backend_tensor_set(b0, b_data.data(), 0, b_data.size() * sizeof(float)); + REQUIRE(ggml_backend_graph_compute(backend, graph0) == GGML_STATUS_SUCCESS); + ggml_backend_synchronize(backend); + + ggml_backend_hrx_cache_stats cache_stats = {}; + REQUIRE(ggml_backend_hrx_get_cache_stats(backend, &cache_stats)); + REQUIRE(cache_stats.graph_program_builds == 1); + REQUIRE(cache_stats.graph_program_hits == 0); + REQUIRE(cache_stats.prepared_program_builds == 1); + REQUIRE(cache_stats.prepared_program_hits == 0); + + ggml_init_params params1 = {}; + params1.mem_size = 256 * 1024; + params1.no_alloc = true; + ggml_context * ctx1 = ggml_init(params1); + REQUIRE(ctx1 != nullptr); + + ggml_tensor * a1 = ggml_new_tensor_1d(ctx1, GGML_TYPE_F32, element_count); + ggml_tensor * b1 = ggml_new_tensor_1d(ctx1, GGML_TYPE_F32, element_count); + ggml_tensor * out1 = ggml_add(ctx1, a1, b1); + REQUIRE(a1 != nullptr); + REQUIRE(b1 != nullptr); + REQUIRE(out1 != nullptr); + + ggml_cgraph * graph1 = ggml_new_graph(ctx1); + REQUIRE(graph1 != nullptr); + ggml_build_forward_expand(graph1, out1); + graph1->uid = 1003; + + ggml_backend_buffer_t buffer1 = ggml_backend_alloc_ctx_tensors(ctx1, backend); + REQUIRE(buffer1 != nullptr); + + for (int64_t i = 0; i < element_count; ++i) { + a_data[i] = static_cast(i % 23) * -0.25f + 5.0f; + b_data[i] = static_cast(i % 7) * 0.5f - 1.0f; + expected[i] = a_data[i] + b_data[i]; + } + ggml_backend_tensor_set(a1, a_data.data(), 0, a_data.size() * sizeof(float)); + ggml_backend_tensor_set(b1, b_data.data(), 0, b_data.size() * sizeof(float)); + REQUIRE(ggml_backend_graph_compute(backend, graph1) == GGML_STATUS_SUCCESS); + ggml_backend_synchronize(backend); + + REQUIRE(ggml_backend_hrx_get_cache_stats(backend, &cache_stats)); + REQUIRE(cache_stats.graph_program_builds == 1); + REQUIRE(cache_stats.graph_program_hits == 1); + REQUIRE(cache_stats.prepared_program_builds == 2); + REQUIRE(cache_stats.prepared_program_hits == 0); + + std::vector actual(element_count); + ggml_backend_tensor_get(out1, actual.data(), 0, actual.size() * sizeof(float)); + for (int64_t i = 0; i < element_count; ++i) { + REQUIRE(actual[i] == expected[i]); + } + + ggml_backend_buffer_free(buffer1); + ggml_free(ctx1); + ggml_backend_buffer_free(buffer0); + ggml_free(ctx0); + ggml_backend_free(backend); +} + +static void run_unsupported_op_fails() { + ggml_backend_t backend = ggml_backend_hrx_init(0); + REQUIRE(backend != nullptr); + + ggml_init_params params = {}; + params.mem_size = 256 * 1024; + params.no_alloc = true; + ggml_context * ctx = ggml_init(params); + REQUIRE(ctx != nullptr); + + ggml_tensor * a = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 8); + ggml_tensor * out = ggml_sqr(ctx, a); + REQUIRE(a != nullptr); + REQUIRE(out != nullptr); + + ggml_cgraph * graph = ggml_new_graph(ctx); + REQUIRE(graph != nullptr); + ggml_build_forward_expand(graph, out); + + ggml_backend_buffer_t buffer = ggml_backend_alloc_ctx_tensors(ctx, backend); + REQUIRE(buffer != nullptr); + + std::vector input(8, 2.0f); + ggml_backend_tensor_set(a, input.data(), 0, input.size() * sizeof(float)); + REQUIRE(ggml_backend_graph_compute(backend, graph) == GGML_STATUS_FAILED); + + ggml_backend_buffer_free(buffer); + ggml_free(ctx); + ggml_backend_free(backend); +} + +int main() { + run_status_checks(); + run_command_plan_metadata_checks(); + run_dispatch_registry_checks(); + run_graph_import_checks(); + run_graph_snapshot_diagnostics_checks(); + run_unmatched_graph_diagnostics_checks(); + run_completion_counter_plan_checks(); + run_graph_index_checks(); + run_graph_traversal_checks(); + run_qwen_token_embedding_dispatch_checks(); + run_gather_add_dispatch_checks(); + run_qwen_flash_attention_dispatch_checks(); + run_qwen_attention_postprocess_dispatch_checks(); + run_qwen_matmul_dispatch_checks(); + schedule_qwen_terminal_q6k_q8_command(1); + schedule_qwen_terminal_q6k_q8_command(18); + run_qwen_router_top8_dispatch_checks(); + run_qwen_routed_gate_up_dispatch_checks(); + run_alias_value_import_checks(); + run_multi_dispatch_checks(); + run_layout_alias_scheduler_elision_checks(); + run_transient_import_checks(); + run_chained_dispatch_requires_transients(); + run_graph_replay_host_staging_is_not_ineligible(); + run_multiple_transient_plan_checks(); + run_disjoint_transient_plan_packing_checks(); + run_graph_program_cache_uid_mismatch_checks(); + run_graph_executor_contract_checks(); + REQUIRE(ggml::hrx::loom_async_jit_enabled_from_environment() == async_jit_expected_from_environment()); + + if (ggml_backend_hrx_get_device_count() == 0) { + std::fprintf(stderr, "test skipped: no HRX devices available\n"); + return 0; + } + + run_qwen_expert_table_partition_prefill_512_execution(); + run_add_f32(); + run_two_independent_add_f32(); + run_chained_add_f32(); + run_same_uid_distinct_graph_reuses_graph_program(); + run_unsupported_op_fails(); + return 0; +} diff --git a/tests/test-hrx-buffer.cpp b/tests/test-hrx-buffer.cpp new file mode 100644 index 000000000000..6d7d25703b03 --- /dev/null +++ b/tests/test-hrx-buffer.cpp @@ -0,0 +1,377 @@ +#include "backend-buffer-binding.h" +#include "backend-context.h" +#include "ggml-backend-impl.h" +#include "ggml-backend.h" +#include "ggml-hrx.h" +#include "ggml.h" +#include "hrx-interop-utils.h" +#include "runtime/host-memory.h" + +#include +#include +#include +#include +#include +#include + +#define REQUIRE(condition) \ + do { \ + if (!(condition)) { \ + std::fprintf(stderr, "%s:%d: requirement failed: %s\n", __FILE__, __LINE__, #condition); \ + std::abort(); \ + } \ + } while (false) + +static void require_hrx_status(hrx_status_t status) { + if (ggml::hrx::ErrorResult error = ggml::hrx::take_status(status)) { + std::fprintf(stderr, "HRX status failed: %s\n", error->c_str()); + std::abort(); + } +} + +static ggml_backend_hrx_context * backend_context(ggml_backend_t backend) { + auto * context = static_cast(backend->context); + REQUIRE(context != nullptr); + REQUIRE(context->device != nullptr); + REQUIRE(context->device->device != nullptr); + REQUIRE(context->stream != nullptr); + return context; +} + +static void run_backend_buffer_checks(ggml_backend_t backend) { + ggml_backend_hrx_context * hrx = backend_context(backend); + + ggml_init_params params = {}; + params.mem_size = 16 * 1024; + params.no_alloc = true; + ggml_context * context = ggml_init(params); + REQUIRE(context != nullptr); + ggml_tensor * tensor = ggml_new_tensor_1d(context, GGML_TYPE_I32, 64); + ggml_tensor * copy = ggml_new_tensor_1d(context, GGML_TYPE_I32, 64); + ggml_backend_buffer_t buffer = ggml_backend_alloc_buffer(backend, 4096); + ggml_backend_buffer_t copy_buffer = ggml_backend_alloc_buffer(backend, 4096); + REQUIRE(buffer != nullptr); + REQUIRE(copy_buffer != nullptr); + tensor->buffer = buffer; + tensor->data = ggml_backend_buffer_get_base(buffer); + copy->buffer = copy_buffer; + copy->data = ggml_backend_buffer_get_base(copy_buffer); + REQUIRE(ggml_backend_buffer_init_tensor(buffer, tensor) == GGML_STATUS_SUCCESS); + REQUIRE(ggml_backend_buffer_init_tensor(copy_buffer, copy) == GGML_STATUS_SUCCESS); + + std::array input = {}; + for (size_t i = 0; i < input.size(); ++i) { + input[i] = static_cast(i * 17 + 3); + } + ggml_backend_tensor_set(tensor, input.data(), 0, sizeof(input)); + std::array output = {}; + ggml_backend_tensor_get(tensor, output.data(), 0, sizeof(output)); + REQUIRE(output == input); + ggml_backend_tensor_copy(tensor, copy); + output.fill(0); + ggml_backend_tensor_get(copy, output.data(), 0, sizeof(output)); + REQUIRE(output == input); + + input[0] = 0x12345678; + ggml_backend_tensor_set_async(backend, tensor, input.data(), 0, sizeof(input)); + ggml_backend_synchronize(backend); + output.fill(0); + ggml_backend_tensor_get_async(backend, tensor, output.data(), 0, sizeof(output)); + ggml_backend_synchronize(backend); + REQUIRE(output == input); + REQUIRE(hrx->device->synchronous_upload_fallbacks.load(std::memory_order_relaxed) == 1); + REQUIRE(hrx->device->synchronous_download_fallbacks.load(std::memory_order_relaxed) == 1); + + ggml_backend_tensor_memset(tensor, 0x5a, 16, 32); + ggml_backend_tensor_get(tensor, output.data(), 0, sizeof(output)); + const uint8_t * bytes = reinterpret_cast(output.data()); + for (size_t i = 16; i < 48; ++i) { + REQUIRE(bytes[i] == 0x5a); + } + + ggml_backend_buffer_clear(buffer, 0); + ggml_backend_tensor_get(tensor, output.data(), 0, sizeof(output)); + for (uint32_t value : output) { + REQUIRE(value == 0); + } + + ggml_backend_buffer_free(buffer); + ggml_backend_buffer_free(copy_buffer); + ggml_free(context); + ggml_backend_synchronize(backend); + REQUIRE(hrx->device->device != nullptr); +} + +static void run_host_buffer_checks(ggml_backend_t backend) { + ggml_backend_hrx_context * context = backend_context(backend); + ggml_backend_buffer_type_t buft = ggml_backend_dev_host_buffer_type(ggml_backend_get_device(backend)); + REQUIRE(buft != nullptr); + REQUIRE(ggml_backend_buft_is_host(buft)); + + const bool original_direct_host_bindings = context->device->use_direct_host_bindings; + context->device->use_direct_host_bindings = false; + ggml_backend_buffer_t buffer = ggml_backend_buft_alloc_buffer(buft, 4096); + context->device->use_direct_host_bindings = original_direct_host_bindings; + REQUIRE(buffer != nullptr); + REQUIRE(ggml_backend_buffer_is_host(buffer)); + auto * buffer_context = ggml_backend_hrx_buffer_context_from_buffer(buffer); + REQUIRE(buffer_context != nullptr); + REQUIRE(buffer_context->buffer != nullptr); + REQUIRE(buffer_context->base == ggml_backend_buffer_get_base(buffer)); + REQUIRE(!buffer_context->direct_host_binding); + + const uint32_t pattern = 0x12345678; + require_hrx_status( + hrx_stream_fill_buffer(context->stream, buffer_context->buffer, 0, 4096, &pattern, sizeof(pattern))); + require_hrx_status(hrx_stream_synchronize(context->stream)); + const auto * words = static_cast(ggml_backend_buffer_get_base(buffer)); + for (size_t i = 0; i < 4096 / sizeof(uint32_t); ++i) { + REQUIRE(words[i] == pattern); + } + + ggml_init_params params = {}; + params.mem_size = 4096; + params.no_alloc = true; + ggml_context * ggml = ggml_init(params); + REQUIRE(ggml != nullptr); + ggml_tensor * host_tensor = ggml_new_tensor_1d(ggml, GGML_TYPE_I32, 64); + host_tensor->buffer = buffer; + host_tensor->data = ggml_backend_buffer_get_base(buffer); + REQUIRE(ggml_backend_buffer_init_tensor(buffer, host_tensor) == GGML_STATUS_SUCCESS); + ggml::hrx::ValueBufferBinding staged_binding; + REQUIRE(ggml_backend_hrx_resolve_value_buffer(host_tensor, staged_binding)); + REQUIRE(staged_binding.buffer == nullptr); + REQUIRE(staged_binding.host_data == ggml_backend_buffer_get_base(buffer)); + REQUIRE(staged_binding.offset == 0); + REQUIRE(staged_binding.length == ggml_nbytes(host_tensor)); + + context->device->use_direct_host_bindings = true; + ggml_backend_buffer_t direct_buffer = ggml_backend_buft_alloc_buffer(buft, 4096); + context->device->use_direct_host_bindings = original_direct_host_bindings; + REQUIRE(direct_buffer != nullptr); + auto * direct_buffer_context = ggml_backend_hrx_buffer_context_from_buffer(direct_buffer); + REQUIRE(direct_buffer_context->direct_host_binding); + ggml_tensor * direct_tensor = ggml_new_tensor_1d(ggml, GGML_TYPE_I32, 64); + direct_tensor->buffer = direct_buffer; + direct_tensor->data = ggml_backend_buffer_get_base(direct_buffer); + REQUIRE(ggml_backend_buffer_init_tensor(direct_buffer, direct_tensor) == GGML_STATUS_SUCCESS); + ggml::hrx::ValueBufferBinding direct_binding; + REQUIRE(ggml_backend_hrx_resolve_value_buffer(direct_tensor, direct_binding)); + REQUIRE(direct_binding.buffer == direct_buffer_context->buffer); + REQUIRE(direct_binding.host_data == nullptr); + REQUIRE(direct_binding.offset == 0); + REQUIRE(direct_binding.length == ggml_nbytes(direct_tensor)); + + ggml_tensor * tensor = ggml_new_tensor_1d(ggml, GGML_TYPE_I32, 64); + ggml_backend_buffer_t local = ggml_backend_alloc_buffer(backend, 4096); + REQUIRE(local != nullptr); + tensor->buffer = local; + tensor->data = ggml_backend_buffer_get_base(local); + REQUIRE(ggml_backend_buffer_init_tensor(local, tensor) == GGML_STATUS_SUCCESS); + + const uint64_t upload_fallbacks = + context->device->synchronous_upload_fallbacks.load(std::memory_order_relaxed); + const uint64_t download_fallbacks = + context->device->synchronous_download_fallbacks.load(std::memory_order_relaxed); + auto * host_words = static_cast(ggml_backend_buffer_get_base(buffer)); + for (size_t i = 0; i < 64; ++i) { + host_words[i] = static_cast(i * 13 + 7); + } + ggml_backend_tensor_set_async(backend, tensor, host_words, 0, 64 * sizeof(uint32_t)); + ggml_backend_synchronize(backend); + std::memset(host_words, 0, 64 * sizeof(uint32_t)); + ggml_backend_tensor_get_async(backend, tensor, host_words, 0, 64 * sizeof(uint32_t)); + ggml_backend_synchronize(backend); + for (size_t i = 0; i < 64; ++i) { + REQUIRE(host_words[i] == static_cast(i * 13 + 7)); + } + REQUIRE(context->device->synchronous_upload_fallbacks.load(std::memory_order_relaxed) == upload_fallbacks); + REQUIRE(context->device->synchronous_download_fallbacks.load(std::memory_order_relaxed) == download_fallbacks); + + ggml_backend_buffer_free(local); + ggml_backend_buffer_free(direct_buffer); + ggml_backend_buffer_free(buffer); + ggml_free(ggml); +} + +static void run_host_transfer_checks(ggml_backend_hrx_context * context) { + ggml::hrx::HostTransferManager transfers; + ggml::hrx::HostStagingBuffer staging; + REQUIRE(ggml::hrx::allocate_host_staging_buffer(context->device->device, 64, staging).success()); + + const std::array zero = {}; + require_hrx_status(hrx_synchronous_h2d(context->device->device, zero.data(), staging.buffer, 0, zero.size())); + + std::array host = {}; + for (size_t i = 0; i < host.size(); ++i) { + host[i] = static_cast(i + 1); + } + + REQUIRE(transfers.upload_synchronous(context->stream, host.data(), staging.buffer, 0, 0).success()); + REQUIRE(transfers.upload_synchronous(context->stream, host.data() + 8, staging.buffer, 16, 24).success()); + ggml::hrx::HostTransferStats stats = transfers.stats(); + REQUIRE(stats.uploads == 1); + REQUIRE(stats.upload_bytes == 24); + require_hrx_status(hrx_stream_synchronize(context->stream)); + + std::array upload_result = {}; + require_hrx_status( + hrx_synchronous_d2h(context->device->device, staging.buffer, 0, upload_result.data(), upload_result.size())); + for (size_t i = 0; i < upload_result.size(); ++i) { + const uint8_t expected = i >= 16 && i < 40 ? host[i - 8] : 0; + REQUIRE(upload_result[i] == expected); + } + + std::array device_values = {}; + for (size_t i = 0; i < device_values.size(); ++i) { + device_values[i] = static_cast(0xa0 + i); + } + require_hrx_status( + hrx_synchronous_h2d(context->device->device, device_values.data(), staging.buffer, 0, device_values.size())); + + std::array download_result = {}; + REQUIRE(transfers.download_synchronous( + context->stream, staging.buffer, 12, download_result.data() + 4, 20).success()); + stats = transfers.stats(); + REQUIRE(stats.downloads == 1); + REQUIRE(stats.download_bytes == 20); + require_hrx_status(hrx_stream_synchronize(context->stream)); + for (size_t i = 0; i < download_result.size(); ++i) { + const uint8_t expected = i >= 4 && i < 24 ? device_values[i + 8] : 0; + REQUIRE(download_result[i] == expected); + } + + REQUIRE(!transfers.upload_synchronous(nullptr, host.data(), staging.buffer, 0, 4).success()); + REQUIRE(!transfers.upload_synchronous(context->stream, nullptr, staging.buffer, 0, 4).success()); + REQUIRE(!transfers.upload_synchronous(context->stream, host.data(), nullptr, 0, 4).success()); + REQUIRE(!transfers.download_synchronous(nullptr, staging.buffer, 0, download_result.data(), 4).success()); + REQUIRE(!transfers.download_synchronous(context->stream, nullptr, 0, download_result.data(), 4).success()); + REQUIRE(!transfers.download_synchronous(context->stream, staging.buffer, 0, nullptr, 4).success()); + + transfers.clear(); + stats = transfers.stats(); + REQUIRE(stats.uploads == 0); + REQUIRE(stats.downloads == 0); + REQUIRE(stats.upload_bytes == 0); + REQUIRE(stats.download_bytes == 0); +} + +static void run_host_staging_checks(ggml_backend_hrx_context * context) { + ggml::hrx::HostStagingBuffer staging; + REQUIRE(ggml::hrx::allocate_host_staging_buffer(context->device->device, 32, staging).success()); + REQUIRE(staging.buffer != nullptr); + REQUIRE(staging.length == 32); + + hrx_buffer_t original = staging.buffer; + ggml::hrx::HostStagingBuffer moved(std::move(staging)); + REQUIRE(moved.buffer == original); + REQUIRE(moved.length == 32); + REQUIRE(staging.buffer == nullptr); + REQUIRE(staging.length == 0); + + ggml::hrx::HostStagingBuffer assigned; + assigned = std::move(moved); + REQUIRE(assigned.buffer == original); + REQUIRE(assigned.length == 32); + REQUIRE(moved.buffer == nullptr); + REQUIRE(moved.length == 0); + + assigned.clear(); + REQUIRE(assigned.buffer == nullptr); + REQUIRE(assigned.length == 0); + assigned.clear(); + REQUIRE(assigned.buffer == nullptr); +} + +static void run_host_weight_cache_checks(ggml_backend_hrx_context * context) { + ggml::hrx::HostTransferManager transfers; + ggml::hrx::HostWeightCache weights; + + std::array host = {}; + for (size_t i = 0; i < host.size(); ++i) { + host[i] = static_cast(i); + } + + ggml::hrx::HostWeightSource source; + source.host_data = host.data(); + source.identity = 0x1234; + source.generation = 1; + source.capacity = host.size(); + source.offset = 16; + source.length = 32; + source.layout = "ggml-native"; + + ggml::hrx::HostWeightAcquireResult first = + weights.acquire(context->device->device, context->stream, transfers, source); + REQUIRE(first.valid()); + + std::array first_bytes = {}; + require_hrx_status( + hrx_synchronous_d2h(context->device->device, first.lease.buffer(), 0, first_bytes.data(), first_bytes.size())); + for (size_t i = 0; i < first_bytes.size(); ++i) { + REQUIRE(first_bytes[i] == host[source.offset + i]); + } + + ggml::hrx::HostWeightAcquireResult second = + weights.acquire(context->device->device, context->stream, transfers, source); + REQUIRE(second.valid()); + REQUIRE(second.lease.buffer() == first.lease.buffer()); + + source.offset = 32; + ggml::hrx::HostWeightAcquireResult slice = + weights.acquire(context->device->device, context->stream, transfers, source); + REQUIRE(slice.valid()); + REQUIRE(slice.lease.buffer() != first.lease.buffer()); + + source.offset = 16; + source.generation = 2; + ggml::hrx::HostWeightAcquireResult next_generation = + weights.acquire(context->device->device, context->stream, transfers, source); + REQUIRE(next_generation.valid()); + REQUIRE(next_generation.lease.buffer() != first.lease.buffer()); + + source.generation = 1; + source.layout = "alternate-layout"; + ggml::hrx::HostWeightAcquireResult conflict = + weights.acquire(context->device->device, context->stream, transfers, source); + REQUIRE(!conflict.valid()); + + ggml::hrx::HostWeightCacheStats weight_stats = weights.stats(); + REQUIRE(weight_stats.hits == 1); + REQUIRE(weight_stats.misses == 3); + REQUIRE(weight_stats.layout_conflicts == 1); + REQUIRE(weight_stats.allocation_count == 3); + REQUIRE(weight_stats.resident_bytes == 96); + + const ggml::hrx::HostTransferStats transfer_stats = transfers.stats(); + REQUIRE(transfer_stats.uploads == 3); + REQUIRE(transfer_stats.upload_bytes == 96); + + weights.clear(); + weight_stats = weights.stats(); + REQUIRE(weight_stats.hits == 0); + REQUIRE(weight_stats.misses == 0); + REQUIRE(weight_stats.layout_conflicts == 0); + REQUIRE(weight_stats.allocation_count == 0); + REQUIRE(weight_stats.resident_bytes == 0); +} + +int main() { + if (ggml_backend_hrx_get_device_count() == 0) { + std::fprintf(stderr, "test skipped: no HRX devices available\n"); + return 0; + } + + ggml_backend_t backend = ggml_backend_hrx_init(0); + REQUIRE(backend != nullptr); + ggml_backend_hrx_context * context = backend_context(backend); + + run_backend_buffer_checks(backend); + run_host_buffer_checks(backend); + run_host_transfer_checks(context); + run_host_staging_checks(context); + run_host_weight_cache_checks(context); + + ggml_backend_free(backend); + return 0; +} diff --git a/tests/test-hrx-loom-jit.cpp b/tests/test-hrx-loom-jit.cpp new file mode 100644 index 000000000000..d8480a3492bf --- /dev/null +++ b/tests/test-hrx-loom-jit.cpp @@ -0,0 +1,431 @@ +#include "dispatch/dispatch.h" +#include "hrx-interop-utils.h" +#include "kernel-corpus/kernel-corpus.h" +#include "runtime/kernel-executable-cache.h" +#include "runtime/loom-kernel-jit.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#define REQUIRE(condition) \ + do { \ + if (!(condition)) { \ + std::fprintf(stderr, "%s:%d: requirement failed: %s\n", __FILE__, __LINE__, #condition); \ + std::abort(); \ + } \ + } while (false) + +namespace { + +class HrxTestDevice { + public: + ~HrxTestDevice() { + if (device != nullptr) { + hrx_device_release(device); + } + } + + bool open() { + hrx_status_t init_status = hrx_gpu_initialize(0); + if (!hrx_status_is_ok(init_status)) { + if (hrx_status_code(init_status) == HRX_STATUS_ALREADY_EXISTS) { + hrx_status_ignore(init_status); + } else { + hrx_status_ignore(init_status); + return false; + } + } + + int count = 0; + if (ggml::hrx::ErrorResult error = ggml::hrx::take_status(hrx_gpu_device_count(&count))) { + std::fprintf(stderr, "skip executable cache materialization: device count failed: %s\n", error->c_str()); + return false; + } + for (int i = 0; i < count; ++i) { + hrx_device_t candidate = nullptr; + if (ggml::hrx::ErrorResult error = ggml::hrx::take_status(hrx_gpu_device_get(i, &candidate))) { + std::fprintf(stderr, "skip HRX device %d: %s\n", i, error->c_str()); + continue; + } + if (candidate == nullptr) { + continue; + } + hrx_device_retain(candidate); + std::optional candidate_architecture = + read_string_property(candidate, HRX_DEVICE_PROPERTY_ARCHITECTURE); + if (!candidate_architecture) { + hrx_device_release(candidate); + continue; + } + device = candidate; + architecture = *candidate_architecture; + return true; + } + return false; + } + + hrx_device_t device = nullptr; + std::string architecture; + + private: + static std::optional read_string_property(hrx_device_t device, hrx_device_property_t property) { + std::vector buffer(64); + while (buffer.size() <= 4096) { + hrx_status_t status = hrx_device_get_property(device, property, buffer.data(), buffer.size()); + if (hrx_status_is_ok(status)) { + return std::string(buffer.data()); + } + if (hrx_status_code(status) != HRX_STATUS_OUT_OF_RANGE) { + if (ggml::hrx::ErrorResult error = ggml::hrx::take_status(status)) { + std::fprintf(stderr, "HRX property query failed: %s\n", error->c_str()); + } + return std::nullopt; + } + hrx_status_ignore(status); + buffer.resize(buffer.size() * 2); + } + return std::nullopt; + } +}; + +static ggml_hrx_loom_jit_source_format to_jit_source_format(ggml::hrx::KernelSourceFormat format) { + switch (format) { + case ggml::hrx::KERNEL_SOURCE_FORMAT_TEXT: + return GGML_HRX_LOOM_JIT_SOURCE_FORMAT_TEXT; + case ggml::hrx::KERNEL_SOURCE_FORMAT_BINARY: + return GGML_HRX_LOOM_JIT_SOURCE_FORMAT_BYTECODE; + } + return GGML_HRX_LOOM_JIT_SOURCE_FORMAT_TEXT; +} + +static const ggml::hrx::KernelDefinition & find_kernel(const char * name) { + const ggml::hrx::KernelCorpus & corpus = ggml::hrx::get_qwen_kernel_corpus(); + for (const ggml::hrx::KernelDefinition & kernel : corpus.kernels) { + if (std::strcmp(kernel.name, name) == 0) { + return kernel; + } + } + std::fprintf(stderr, "missing test kernel: %s\n", name); + std::abort(); +} + +static const ggml::hrx::KernelDefinition * find_targeted_kernel(const char * name, const char * target) { + const ggml::hrx::KernelCorpus & corpus = ggml::hrx::get_qwen_kernel_corpus(); + for (const ggml::hrx::KernelDefinition & kernel : corpus.kernels) { + if (std::strcmp(kernel.name, name) == 0 && std::strcmp(kernel.target_selector, target) == 0) { + return &kernel; + } + } + return nullptr; +} + +static ggml::hrx::Dispatch make_add_dispatch(const ggml::hrx::KernelDefinition & definition, int64_t element_count) { + ggml::hrx::Dispatch dispatch; + dispatch.kernel.kernel_id = definition.id; + dispatch.kernel.integer_parameters.emplace("element_count", element_count); + dispatch.bindings.reserve(definition.bindings.size()); + for (size_t i = 0; i < definition.bindings.size(); ++i) { + dispatch.bindings.push_back({ + ggml::hrx::ValueId(static_cast(i + 1)), + 0, + static_cast(element_count) * sizeof(float), + }); + } + return dispatch; +} + +static ggml::hrx::LoomKernelCompileRequest make_compile_request( + const ggml::hrx::KernelDefinition & definition, + const std::map & workload, + const std::map & compile_config = {}) { + ggml::hrx::LoomKernelCompileRequest request; + REQUIRE(!definition.compile_recipe.primary_sources.empty()); + + const ggml::hrx::KernelSourceRef & primary_source = definition.compile_recipe.primary_sources.front(); + REQUIRE(primary_source.contents != nullptr); + request.source_data = primary_source.contents->source.data; + request.source_size = primary_source.contents->source.length; + request.source_format = to_jit_source_format(primary_source.contents->source.format); + request.source_identifier = primary_source.path != nullptr ? primary_source.path : ""; + request.symbol = definition.symbol != nullptr ? definition.symbol : ""; + request.launch_config_symbol = definition.name != nullptr ? definition.name : ""; + + request.dependencies.reserve(definition.compile_recipe.library_sources.size()); + for (const ggml::hrx::KernelSourceRef & dependency_ref : definition.compile_recipe.library_sources) { + REQUIRE(dependency_ref.contents != nullptr); + request.dependencies.push_back({ + dependency_ref.contents->source.data, + dependency_ref.contents->source.length, + to_jit_source_format(dependency_ref.contents->source.format), + dependency_ref.path, + }); + } + + std::map merged_config; + for (const ggml::hrx::KernelCompileConfig & config : definition.compile_config) { + merged_config[config.key != nullptr ? config.key : ""] = config.value != nullptr ? config.value : ""; + } + for (const auto & item : compile_config) { + merged_config[item.first] = item.second; + } + request.config_storage.reserve(merged_config.size()); + for (const auto & item : merged_config) { + request.config_storage.push_back(item); + } + + request.workload.reserve(definition.workload_parameters.size()); + for (const ggml::hrx::KernelScalarDefinition & parameter : definition.workload_parameters) { + REQUIRE(parameter.type != nullptr); + REQUIRE(std::strcmp(parameter.type, "index") == 0); + const auto found = workload.find(parameter.name != nullptr ? parameter.name : ""); + REQUIRE(found != workload.end()); + request.workload.push_back(found->second); + } + return request; +} + +static ggml::hrx::LoomCompiledKernelRef compile_kernel(ggml::hrx::LoomJit & jit, + const std::string & key, + const ggml::hrx::KernelDefinition & definition, + const std::map & workload, + const std::map & compile_config = {}) { + return jit.compile(key, make_compile_request(definition, workload, compile_config)); +} + +static bool resolve_with_timeout(const ggml::hrx::LoomCompiledKernelRef & ref, std::chrono::seconds timeout) { + std::atomic done = false; + std::atomic success = false; + std::thread resolver([&] { + success.store(ref->resolve(), std::memory_order_release); + done.store(true, std::memory_order_release); + }); + + const auto deadline = std::chrono::steady_clock::now() + timeout; + while (!done.load(std::memory_order_acquire)) { + if (std::chrono::steady_clock::now() >= deadline) { + std::fprintf(stderr, "timed out waiting for async Loom compile: %s\n", ref->key().c_str()); + std::abort(); + } + std::this_thread::sleep_for(std::chrono::milliseconds(10)); + } + resolver.join(); + return success.load(std::memory_order_acquire); +} + +static void require_compiled_kernel(const ggml::hrx::LoomCompiledKernelRef & ref) { + REQUIRE(ref != nullptr); + REQUIRE(resolve_with_timeout(ref, std::chrono::seconds(120))); + ggml_hrx_loom_jit_compile_result compiled = ref->take_result(); + REQUIRE(compiled.hsaco_data != nullptr); + REQUIRE(compiled.hsaco_size > 0); + REQUIRE(compiled.launch_config.workgroup_count[0] > 0); + REQUIRE(compiled.launch_config.workgroup_size[0] > 0); + compiled.reset(); +} + +static int64_t elapsed_us(std::chrono::steady_clock::time_point begin, std::chrono::steady_clock::time_point end) { + return std::chrono::duration_cast(end - begin).count(); +} + +static void run_cache_materialize_case(HrxTestDevice & device, + ggml::hrx::LoomJitMode mode, + const ggml::hrx::KernelDefinition & add) { + ggml::hrx::KernelExecutablePrepareContext context = {}; + context.device = device.device; + context.target = device.architecture.c_str(); + + ggml::hrx::Dispatch dispatch = make_add_dispatch(add, 64); + std::vector constants; + ggml::hrx::KernelExecutableCache cache(mode); + + ggml::hrx::KernelExecutableRef ref = cache.get_or_compile(context, add, dispatch, constants); + REQUIRE(ref.valid()); + + std::shared_ptr executable = cache.materialize(context, ref, constants); + REQUIRE(executable != nullptr); + REQUIRE(executable->executable != nullptr); + REQUIRE(executable->export_info.binding_count == dispatch.bindings.size()); + REQUIRE(executable->export_info.constant_byte_length == constants.size()); + REQUIRE(executable->launch.workgroup_count[0] > 0); + REQUIRE(executable->launch.workgroup_size[0] > 0); + + std::shared_ptr loaded_hit = cache.materialize(context, ref, constants); + REQUIRE(loaded_hit == executable); + + std::vector constants_for_prepare; + std::shared_ptr prepared = + cache.prepare(context, add, dispatch, constants_for_prepare); + REQUIRE(prepared == executable); + + const char * mode_name = mode == ggml::hrx::LoomJitMode::Async ? "async" : "sync"; + std::printf("%s KernelExecutableCache materialized ggml_add_f32 for %s\n", mode_name, device.architecture.c_str()); +} + +static void run_targeted_export_materialize_case(HrxTestDevice & device) { + const ggml::hrx::KernelDefinition * definition = + find_targeted_kernel("ggml_linear_q6k_q8_1_x4", device.architecture.c_str()); + if (definition == nullptr) { + return; + } + + ggml::hrx::KernelExecutablePrepareContext context = {}; + context.device = device.device; + context.target = device.architecture.c_str(); + + ggml::hrx::Dispatch dispatch; + dispatch.kernel.kernel_id = definition->id; + dispatch.kernel.integer_parameters.emplace("token_count", 1); + dispatch.kernel.integer_parameters.emplace("input_size", 2048); + dispatch.kernel.integer_parameters.emplace("output_size", 151936); + dispatch.kernel.compile_parameters.emplace("ggml.linear_q6k_q8_1_x4.token_capacity", "1"); + dispatch.kernel.compile_parameters.emplace("ggml.linear_q6k_q8_1_x4.output_capacity", "151936"); + dispatch.bindings.resize(3); + + ggml::hrx::KernelExecutableCache cache(ggml::hrx::LoomJitMode::Sync); + std::vector constants; + std::shared_ptr executable = cache.prepare(context, *definition, dispatch, constants); + REQUIRE(executable != nullptr); + REQUIRE(executable->executable != nullptr); + REQUIRE(executable->launch.workgroup_count[0] == 151936); + std::printf("KernelExecutableCache materialized targeted %s as export %s for %s\n", definition->symbol, + definition->name, device.architecture.c_str()); +} + +} // namespace + +int main() { + static constexpr const char * kTarget = "gfx1100"; + + const ggml::hrx::KernelDefinition & add = find_kernel("ggml_add_f32"); + const ggml::hrx::KernelDefinition & gather_add = find_kernel("ggml_gather_add_f32"); + const ggml::hrx::KernelDefinition & rmsnorm = find_kernel("qwen3_moe_rmsnorm_f32"); + const ggml::hrx::KernelDefinition & router_top8 = find_kernel("qwen3_moe_router_top8_f32"); + const ggml::hrx::KernelDefinition & expert_table = find_kernel("qwen3_moe_build_expert_table"); + const ggml::hrx::KernelDefinition & partition_table = find_kernel("qwen3_moe_build_expert_partition_table"); + + std::string sync_error; + std::unique_ptr sync_jit = + ggml::hrx::create_loom_jit(kTarget, ggml::hrx::LoomJitMode::Sync, sync_error); + REQUIRE(sync_jit != nullptr); + REQUIRE(!sync_jit->async_enabled()); + + const auto sync_begin = std::chrono::steady_clock::now(); + ggml::hrx::LoomCompiledKernelRef sync_ref = compile_kernel(*sync_jit, "sync-add-64", add, + { + { "element_count", 64 } + }); + require_compiled_kernel(sync_ref); + const auto sync_end = std::chrono::steady_clock::now(); + const int64_t sync_compile_us = elapsed_us(sync_begin, sync_end); + std::printf("sync Loom compile completed in %ld us\n", static_cast(sync_compile_us)); + + std::string async_error; + std::unique_ptr async_jit = + ggml::hrx::create_loom_jit(kTarget, ggml::hrx::LoomJitMode::Async, async_error); + REQUIRE(async_jit != nullptr); + REQUIRE(async_jit->async_enabled()); + + std::vector refs; + refs.reserve(9); + const auto enqueue_begin = std::chrono::steady_clock::now(); + refs.push_back(compile_kernel(*async_jit, "async-add-64", add, + { + { "element_count", 64 } + })); + refs.push_back(compile_kernel(*async_jit, "async-add-128", add, + { + { "element_count", 128 } + })); + refs.push_back(compile_kernel(*async_jit, "async-add-256", add, + { + { "element_count", 256 } + })); + refs.push_back(compile_kernel(*async_jit, "async-add-512", add, + { + { "element_count", 512 } + })); + refs.push_back(compile_kernel(*async_jit, "async-rmsnorm-1", rmsnorm, + { + { "token_count", 1 } + }, + { + { "qwen3_moe.model.hidden_size", "2048" }, + { "qwen3_moe.model.rms_epsilon", "0.000001" }, + { "qwen3_moe.workload.token_capacity", "1" }, + { "ggml.quantize_q8_1_x4.group_capacity", "256" }, + })); + refs.push_back(compile_kernel(*async_jit, "async-router-top8-1", router_top8, + { + { "token_count", 1 }, + { "route_id_stride", 8 }, + }, + { + { "qwen3_moe.router.expert_count", "128" }, + { "qwen3_moe.router.route_count", "8" }, + { "qwen3_moe.workload.token_capacity", "1" }, + })); + refs.push_back(compile_kernel(*async_jit, "async-expert-table-1", expert_table, + { + { "token_count", 1 }, + { "route_count", 8 }, + { "route_stride", 8 }, + { "expert_count", 128 }, + }, + { + { "qwen3_moe.routed_gate_up.expert_count", "128" }, + { "qwen3_moe.routed_gate_up.route_count", "8" }, + { "qwen3_moe.workload.token_capacity", "1" }, + })); + refs.push_back(compile_kernel(*async_jit, "async-partition-table-1", partition_table, + { + { "token_count", 1 }, + { "route_count", 8 }, + { "expert_count", 128 }, + }, + { + { "qwen3_moe.routed_gate_up.expert_count", "128" }, + { "qwen3_moe.routed_gate_up.route_count", "8" }, + { "qwen3_moe.workload.token_capacity", "1" }, + })); + // Exercise gather-add coverage in the async JIT path. + refs.push_back(compile_kernel(*async_jit, "async-gather-add-2-to-1", gather_add, + { + { "source_token_count", 2 }, + { "output_token_count", 1 }, + { "hidden_size", 2048 }, + })); + const auto enqueue_end = std::chrono::steady_clock::now(); + const int64_t enqueue_us = elapsed_us(enqueue_begin, enqueue_end); + std::printf("async Loom enqueue completed in %ld us\n", static_cast(enqueue_us)); + + const int64_t max_expected_enqueue_us = std::max(100000, sync_compile_us / 2); + REQUIRE(enqueue_us < max_expected_enqueue_us); + + for (const ggml::hrx::LoomCompiledKernelRef & ref : refs) { + require_compiled_kernel(ref); + } + + std::printf("async Loom JIT compiled %zu kernels\n", refs.size()); + + HrxTestDevice device; + if (device.open()) { + run_cache_materialize_case(device, ggml::hrx::LoomJitMode::Sync, add); + run_cache_materialize_case(device, ggml::hrx::LoomJitMode::Async, add); + run_targeted_export_materialize_case(device); + } else { + std::printf("skipping KernelExecutableCache materialization checks: no HRX device available\n"); + } + + return 0; +} diff --git a/tests/test-hrx-ops.cpp b/tests/test-hrx-ops.cpp new file mode 100644 index 000000000000..0eb83eba8f20 --- /dev/null +++ b/tests/test-hrx-ops.cpp @@ -0,0 +1,2028 @@ +#include "backend-context.h" +#include "dispatch/dispatch-scheduler.h" +#include "ggml-alloc.h" +#include "ggml-backend.h" +#include "ggml-hrx.h" +#include "ggml.h" +#include "graph/graph.h" +#include "kernel-corpus/kernel-corpus.h" +#include "runtime/graph-executor.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#define REQUIRE(condition) \ + do { \ + if (!(condition)) { \ + std::fprintf(stderr, "%s:%d: requirement failed: %s\n", __FILE__, __LINE__, #condition); \ + std::abort(); \ + } \ + } while (false) + +static constexpr float kQwenRmsNormEps = 0.000001f; +static constexpr int64_t kQwenFlashHeadSize = 128; +static constexpr int64_t kQwenRouterExpertCount = 128; +static constexpr int64_t kQwenRouterRouteCount = 8; +static constexpr int64_t kQwenHiddenSize = 2048; +static constexpr int64_t kQwenMoeIntermediate = 768; +static constexpr int64_t kQwenVocabularyCount = 151936; + +static ggml::hrx::Value make_test_value(ggml::hrx::ValueId id, + ggml::hrx::ValueStorageId storage, + ggml::hrx::ValueId storage_root, + ggml::hrx::ValueId alias_source, + size_t storage_offset, + size_t storage_byte_count, + ggml_type type, + int64_t element_count) { + ggml::hrx::Value value = {}; + value.id = id; + value.kind = ggml::hrx::ValueKind::Transient; + value.storage = storage; + value.storage_root = storage_root; + value.alias_source = alias_source; + value.storage_offset = storage_offset; + value.storage_byte_count = storage_byte_count; + value.type = type; + value.ne = { element_count, 1, 1, 1 }; + value.nb = { ggml_type_size(type), ggml_type_size(type) * static_cast(element_count), + ggml_type_size(type) * static_cast(element_count), + ggml_type_size(type) * static_cast(element_count) }; + value.element_count = element_count; + value.byte_count = ggml_row_size(type, element_count); + value.contiguous = true; + return value; +} + +static std::vector make_input(int64_t hidden_size, int64_t token_count) { + std::vector data(hidden_size * token_count); + for (int64_t i = 0; i < static_cast(data.size()); ++i) { + data[i] = static_cast((i % 29) - 14) * 0.125f; + } + return data; +} + +static std::vector make_weight(int64_t hidden_size) { + std::vector data(hidden_size); + for (int64_t i = 0; i < hidden_size; ++i) { + data[i] = 0.5f + static_cast(i % 17) * 0.03125f; + } + return data; +} + +static std::vector make_router_input(int64_t hidden_size, int64_t token_count) { + std::vector data(hidden_size * token_count); + for (int64_t i = 0; i < static_cast(data.size()); ++i) { + data[i] = static_cast((i % 41) - 20) * 0.01f; + } + return data; +} + +static std::vector make_router_weight(int64_t hidden_size, int64_t expert_count) { + std::vector data(hidden_size * expert_count); + for (int64_t expert = 0; expert < expert_count; ++expert) { + for (int64_t column = 0; column < hidden_size; ++column) { + data[expert * hidden_size + column] = static_cast(((expert + column) % 31) - 15) * 0.0025f; + } + } + return data; +} + +static std::vector make_router_logits(int64_t token_count) { + std::vector data(kQwenRouterExpertCount * token_count); + for (int64_t i = 0; i < static_cast(data.size()); ++i) { + data[i] = static_cast(i % kQwenRouterRouteCount); + } + return data; +} + +static std::vector make_flash_query(int64_t token_count) { + std::vector data(kQwenFlashHeadSize * token_count); + for (int64_t i = 0; i < static_cast(data.size()); ++i) { + data[i] = static_cast((i % 37) - 18) * 0.01f; + } + return data; +} + +static std::vector make_flash_key_value(int64_t token_count, int offset) { + std::vector data(kQwenFlashHeadSize * token_count); + for (int64_t i = 0; i < static_cast(data.size()); ++i) { + const float value = static_cast(((i + offset) % 31) - 15) * 0.015f; + data[i] = ggml_fp32_to_fp16(value); + } + return data; +} + +static std::vector make_flash_mask(int64_t query_token_count, int64_t key_value_token_count) { + std::vector data(query_token_count * key_value_token_count); + for (int64_t query = 0; query < query_token_count; ++query) { + for (int64_t key = 0; key < key_value_token_count; ++key) { + const float value = key <= query + 1 ? 0.0f : -10000.0f; + data[query * key_value_token_count + key] = ggml_fp32_to_fp16(value); + } + } + return data; +} + +static std::vector rmsnorm_mul_reference(const std::vector & input, + const std::vector & weight, + int64_t hidden_size, + int64_t token_count) { + std::vector output(input.size()); + for (int64_t token = 0; token < token_count; ++token) { + float sum_squares = 0.0f; + for (int64_t column = 0; column < hidden_size; ++column) { + const float value = input[token * hidden_size + column]; + sum_squares += value * value; + } + const float scale = 1.0f / std::sqrt(sum_squares / static_cast(hidden_size) + kQwenRmsNormEps); + for (int64_t column = 0; column < hidden_size; ++column) { + output[token * hidden_size + column] = input[token * hidden_size + column] * scale * weight[column]; + } + } + return output; +} + +static std::vector router_projection_reference(const std::vector & input, + const std::vector & weight, + int64_t hidden_size, + int64_t expert_count, + int64_t token_count) { + std::vector output(expert_count * token_count); + for (int64_t token = 0; token < token_count; ++token) { + for (int64_t expert = 0; expert < expert_count; ++expert) { + float sum = 0.0f; + for (int64_t column = 0; column < hidden_size; ++column) { + sum += input[token * hidden_size + column] * weight[expert * hidden_size + column]; + } + output[token * expert_count + expert] = sum; + } + } + return output; +} + +static std::vector router_top8_weights_reference(const std::vector & logits, int64_t token_count) { + std::vector output(kQwenRouterRouteCount * token_count); + for (int64_t token = 0; token < token_count; ++token) { + bool used[kQwenRouterExpertCount] = {}; + int64_t selected[kQwenRouterRouteCount] = {}; + for (int64_t route = 0; route < kQwenRouterRouteCount; ++route) { + int64_t best_expert = -1; + float best_value = -std::numeric_limits::infinity(); + for (int64_t expert = 0; expert < kQwenRouterExpertCount; ++expert) { + const float value = logits[token * kQwenRouterExpertCount + expert]; + if (!used[expert] && + (best_expert < 0 || value > best_value || (value == best_value && expert < best_expert))) { + best_value = value; + best_expert = expert; + } + } + selected[route] = best_expert; + used[best_expert] = true; + } + + float max_selected = -std::numeric_limits::infinity(); + for (const int64_t expert : selected) { + max_selected = std::max(max_selected, logits[token * kQwenRouterExpertCount + expert]); + } + float sum = 0.0f; + for (int64_t route = 0; route < kQwenRouterRouteCount; ++route) { + const float value = std::exp(logits[token * kQwenRouterExpertCount + selected[route]] - max_selected); + output[token * kQwenRouterRouteCount + route] = value; + sum += value; + } + for (int64_t route = 0; route < kQwenRouterRouteCount; ++route) { + output[token * kQwenRouterRouteCount + route] /= sum; + } + } + return output; +} + +static std::vector flash_attention_reference(const std::vector & query, + const std::vector & key, + const std::vector & value, + const std::vector & mask, + int64_t query_token_count, + int64_t key_value_token_count) { + std::vector output(query_token_count * kQwenFlashHeadSize); + const float scale = 1.0f / std::sqrt(static_cast(kQwenFlashHeadSize)); + for (int64_t query_token = 0; query_token < query_token_count; ++query_token) { + std::vector scores(key_value_token_count); + float max_score = -std::numeric_limits::infinity(); + for (int64_t key_token = 0; key_token < key_value_token_count; ++key_token) { + float dot = 0.0f; + for (int64_t channel = 0; channel < kQwenFlashHeadSize; ++channel) { + dot += query[query_token * kQwenFlashHeadSize + channel] * + ggml_fp16_to_fp32(key[key_token * kQwenFlashHeadSize + channel]); + } + const float score = dot * scale + ggml_fp16_to_fp32(mask[query_token * key_value_token_count + key_token]); + scores[key_token] = score; + max_score = std::max(max_score, score); + } + + float sum = 0.0f; + for (float & score : scores) { + score = std::exp(score - max_score); + sum += score; + } + for (int64_t channel = 0; channel < kQwenFlashHeadSize; ++channel) { + float weighted_sum = 0.0f; + for (int64_t key_token = 0; key_token < key_value_token_count; ++key_token) { + const float probability = scores[key_token] / sum; + weighted_sum += probability * ggml_fp16_to_fp32(value[key_token * kQwenFlashHeadSize + channel]); + } + output[query_token * kQwenFlashHeadSize + channel] = weighted_sum; + } + } + return output; +} + +static ggml_tensor * build_rmsnorm_mul_graph(ggml_context * ctx, + ggml_tensor * input, + ggml_tensor * weight, + float eps = kQwenRmsNormEps) { + ggml_tensor * rms = ggml_rms_norm(ctx, input, eps); + REQUIRE(rms != nullptr); + ggml_tensor * output = ggml_mul(ctx, rms, weight); + REQUIRE(output != nullptr); + return output; +} + +static ggml_tensor * build_qwen_flash_attention_graph(ggml_context * ctx, + ggml_tensor * query, + ggml_tensor * key, + ggml_tensor * value, + ggml_tensor * mask) { + ggml_tensor * output = ggml_flash_attn_ext(ctx, query, key, value, mask, + 1.0f / std::sqrt(static_cast(kQwenFlashHeadSize)), 0.0f, 0.0f); + REQUIRE(output != nullptr); + return output; +} + +static ggml_tensor * build_qwen_router_top8_graph(ggml_context * ctx, + ggml_tensor * logits, + ggml_tensor ** route_ids = nullptr) { + ggml_tensor * probs = ggml_soft_max(ctx, logits); + REQUIRE(probs != nullptr); + ggml_tensor * probs_reshaped = ggml_reshape_3d(ctx, probs, 1, kQwenRouterExpertCount, logits->ne[1]); + REQUIRE(probs_reshaped != nullptr); + ggml_tensor * argsort = ggml_argsort(ctx, probs, GGML_SORT_ORDER_DESC); + REQUIRE(argsort != nullptr); + ggml_tensor * topk = ggml_view_2d(ctx, argsort, kQwenRouterRouteCount, logits->ne[1], argsort->nb[1], 0); + REQUIRE(topk != nullptr); + if (route_ids != nullptr) { + *route_ids = topk; + } + ggml_tensor * selected = ggml_get_rows(ctx, probs_reshaped, topk); + REQUIRE(selected != nullptr); + ggml_tensor * selected_reshaped = ggml_reshape_2d(ctx, selected, kQwenRouterRouteCount, logits->ne[1]); + REQUIRE(selected_reshaped != nullptr); + ggml_tensor * sum = ggml_sum_rows(ctx, selected_reshaped); + REQUIRE(sum != nullptr); + ggml_tensor * clamped_sum = ggml_clamp(ctx, sum, 1.0e-7f, std::numeric_limits::infinity()); + REQUIRE(clamped_sum != nullptr); + ggml_tensor * normalized = ggml_div(ctx, selected_reshaped, clamped_sum); + REQUIRE(normalized != nullptr); + ggml_tensor * output = ggml_reshape_3d(ctx, normalized, 1, kQwenRouterRouteCount, logits->ne[1]); + REQUIRE(output != nullptr); + return output; +} + +static std::string kernel_name_for_id(uint64_t kernel_id) { + const ggml::hrx::KernelResolveResult resolved = + ggml::hrx::resolve_kernel_definition(ggml::hrx::get_qwen_kernel_corpus(), "gfx1151", kernel_id); + REQUIRE(resolved.found()); + return ggml::hrx::kernel_definition_name(*resolved.definition); +} + +static std::vector scheduled_kernel_sequence(ggml_cgraph * graph) { + ggml::hrx::GraphImportResult imported = ggml::hrx::import_ggml_graph(*graph); + REQUIRE(imported.valid()); + ggml::hrx::DispatchScheduler scheduler; + ggml::hrx::DispatchScheduleDiagnostics diagnostics; + if (!scheduler.schedule_graph(imported.graph, { "gfx1151" }, &diagnostics)) { + for (const std::string & error : scheduler.plan().status.errors()) { + std::fprintf(stderr, "scheduler error: %s\n", error.c_str()); + } + std::fprintf(stderr, "unsupported: %s\n", diagnostics.unsupported_message.c_str()); + for (const ggml::hrx::DispatchRegistrationAttempt & attempt : diagnostics.match.attempts) { + std::fprintf(stderr, " attempt %s matched=%d\n", attempt.name.c_str(), attempt.matched ? 1 : 0); + if (!attempt.covered_nodes.empty()) { + std::fprintf(stderr, " covered:"); + for (size_t node : attempt.covered_nodes) { + std::fprintf(stderr, " %zu", node); + } + std::fprintf(stderr, "\n"); + } + for (const std::string & error : attempt.errors) { + std::fprintf(stderr, " %s\n", error.c_str()); + } + } + std::abort(); + } + REQUIRE(scheduler.plan().valid()); + + std::vector names; + names.reserve(scheduler.plan().dispatches.size()); + for (const ggml::hrx::Dispatch & dispatch : scheduler.plan().dispatches) { + names.push_back(kernel_name_for_id(dispatch.kernel.kernel_id)); + } + return names; +} + +static size_t producer_index_for_tensor(const ggml::hrx::Graph & graph, const ggml_tensor * tensor) { + const ggml::hrx::Value * value = graph.values().find_tensor(tensor); + REQUIRE(value != nullptr); + const ggml::hrx::GraphNode * producer = graph.index().producer(value->id); + REQUIRE(producer != nullptr); + size_t index = 0; + REQUIRE(graph.index().node_index(producer, index)); + return index; +} + +static ggml::hrx::ValueId next_plan_value(const ggml::hrx::Graph & graph, const ggml::hrx::CommandPlan & plan) { + return ggml::hrx::ValueId( + static_cast(graph.values().size() + plan.transients.size() + plan.completion_counter_requests.size())); +} + +static void append_match_to_plan(ggml::hrx::CommandPlan & plan, + ggml::hrx::DispatchMatch & match, + std::vector & covered_nodes) { + for (ggml::hrx::Dispatch & dispatch : match.initialization_dispatches) { + plan.initialization_dispatches.push_back(std::move(dispatch)); + } + for (ggml::hrx::Dispatch & dispatch : match.dispatches) { + plan.dispatches.push_back(std::move(dispatch)); + } + for (ggml::hrx::CommandPlanTransient & transient : match.transients) { + plan.transients.push_back(std::move(transient)); + } + for (ggml::hrx::CommandPlanConstantInitialization & initialization : match.constant_initializations) { + plan.constant_initializations.push_back(std::move(initialization)); + } + for (ggml::hrx::CommandPlanCompletionCounterRequest & request : match.completion_counter_requests) { + plan.completion_counter_requests.push_back(std::move(request)); + } + REQUIRE(plan.metadata.append(std::move(match.metadata), plan.status)); + plan.status.append(match.status); + for (size_t covered_node : match.covered_nodes) { + REQUIRE(covered_node < covered_nodes.size()); + REQUIRE(!covered_nodes[covered_node]); + covered_nodes[covered_node] = true; + } +} + +static void match_dispatch_at_index(const ggml::hrx::Graph & graph, + const ggml::hrx::DispatchRegistry & registry, + ggml::hrx::CommandPlan & plan, + std::vector & covered_nodes, + size_t root_index, + ggml::hrx::DispatchMatch & match) { + REQUIRE(root_index < graph.nodes().size()); + const ggml::hrx::DispatchMatchContext context = { + graph, &graph.nodes()[root_index], root_index, covered_nodes, plan, next_plan_value(graph, plan), + }; + ggml::hrx::DispatchMatchDiagnostics diagnostics; + if (!registry.match(context, match, &diagnostics)) { + std::fprintf(stderr, "manual matcher failed for node %zu %s\n", root_index, + ggml_op_name(graph.nodes()[root_index].op)); + for (const ggml::hrx::DispatchRegistrationAttempt & attempt : diagnostics.attempts) { + std::fprintf(stderr, " attempt %s matched=%d\n", attempt.name.c_str(), attempt.matched ? 1 : 0); + for (const std::string & error : attempt.errors) { + std::fprintf(stderr, " %s\n", error.c_str()); + } + } + std::abort(); + } + append_match_to_plan(plan, match, covered_nodes); + REQUIRE(plan.valid()); +} + +static void require_kernel_subsequence(const std::vector & sequence, + const std::vector & expected) { + size_t sequence_index = 0; + for (const std::string & name : expected) { + while (sequence_index < sequence.size() && sequence[sequence_index] != name) { + ++sequence_index; + } + if (sequence_index >= sequence.size()) { + std::fprintf(stderr, "missing expected kernel: %s\nscheduled kernels:\n", name.c_str()); + for (const std::string & scheduled : sequence) { + std::fprintf(stderr, " %s\n", scheduled.c_str()); + } + std::abort(); + } + ++sequence_index; + } +} + +static void run_alternate_value_alias_lookup_checks() { + constexpr int64_t element_count = 2048; + const size_t full_bytes = ggml_row_size(GGML_TYPE_F32, element_count); + const size_t q8_bytes = ggml_row_size(GGML_TYPE_Q8_1, element_count); + + ggml::hrx::Graph graph; + ggml::hrx::Status status; + ggml::hrx::CommandPlan plan; + + const ggml::hrx::ValueId root(0); + const ggml::hrx::ValueId full_alias(1); + const ggml::hrx::ValueId partial_alias(2); + const ggml::hrx::ValueId q8_alternate(100); + const ggml::hrx::ValueStorageId storage(0); + + status = graph.values().add_snapshot_storage({ storage, root, full_bytes }); + REQUIRE(status.success()); + status = graph.values().add_snapshot_value( + make_test_value(root, storage, root, ggml::hrx::ValueId(), 0, full_bytes, GGML_TYPE_F32, element_count)); + REQUIRE(status.success()); + status = graph.values().add_snapshot_value( + make_test_value(full_alias, storage, root, root, 0, full_bytes, GGML_TYPE_F32, element_count)); + REQUIRE(status.success()); + status = graph.values().add_snapshot_value( + make_test_value(partial_alias, storage, root, root, 0, full_bytes, GGML_TYPE_F32, element_count / 2)); + REQUIRE(status.success()); + + REQUIRE(plan.metadata.append_alternate_value({ root, q8_alternate, GGML_TYPE_Q8_1, q8_bytes, "q8" }, status)); + + const ggml::hrx::CommandPlanAlternateValue * exact = + ggml::hrx::find_alternate_value(graph, plan, root, GGML_TYPE_Q8_1, q8_bytes); + REQUIRE(exact != nullptr); + REQUIRE(exact->alternate_value == q8_alternate); + + const ggml::hrx::CommandPlanAlternateValue * through_full_alias = + ggml::hrx::find_alternate_value(graph, plan, full_alias, GGML_TYPE_Q8_1, q8_bytes); + REQUIRE(through_full_alias != nullptr); + REQUIRE(through_full_alias->alternate_value == q8_alternate); + + const ggml::hrx::CommandPlanAlternateValue * through_partial_alias = + ggml::hrx::find_alternate_value(graph, plan, partial_alias, GGML_TYPE_Q8_1, q8_bytes); + REQUIRE(through_partial_alias == nullptr); +} + +static std::vector make_pattern_f32(size_t element_count, int seed, float scale = 0.01f) { + std::vector data(element_count); + for (size_t i = 0; i < element_count; ++i) { + const int value = static_cast((i * 17 + static_cast(seed) * 29) % 97) - 48; + data[i] = static_cast(value) * scale; + } + return data; +} + +static std::vector make_i32_mod_data(size_t element_count, int32_t modulo) { + std::vector data(element_count); + for (size_t i = 0; i < element_count; ++i) { + data[i] = static_cast(i % static_cast(modulo)); + } + return data; +} + +static std::vector make_i64_mod_data(size_t element_count, int64_t modulo) { + std::vector data(element_count); + for (size_t i = 0; i < element_count; ++i) { + data[i] = static_cast(i % static_cast(modulo)); + } + return data; +} + +static std::vector make_pattern_f16(size_t element_count, int seed, float scale = 0.01f) { + const std::vector f32 = make_pattern_f32(element_count, seed, scale); + std::vector data(element_count); + for (size_t i = 0; i < element_count; ++i) { + data[i] = ggml_fp32_to_fp16(f32[i]); + } + return data; +} + +static std::vector make_quantized_rows(ggml_type type, int64_t row_length, int64_t row_count, int seed) { + const ggml_type_traits * traits = ggml_get_type_traits(type); + REQUIRE(traits != nullptr); + REQUIRE(traits->from_float_ref != nullptr); + const size_t row_size = ggml_row_size(type, row_length); + std::vector data(static_cast(row_count) * row_size); + std::vector row(static_cast(row_length)); + for (int64_t r = 0; r < row_count; ++r) { + for (int64_t c = 0; c < row_length; ++c) { + const int value = static_cast((r * 13 + c * 7 + seed * 31) % 101) - 50; + row[static_cast(c)] = static_cast(value) * 0.005f; + } + traits->from_float_ref(row.data(), data.data() + static_cast(r) * row_size, row_length); + } + return data; +} + +static void set_tensor_bytes(ggml_backend_t backend, ggml_tensor * tensor, const void * data, size_t byte_count) { + REQUIRE(tensor != nullptr); + REQUIRE(ggml_nbytes(tensor) == byte_count); + ggml_backend_tensor_set(tensor, data, 0, byte_count); + ggml_backend_synchronize(backend); +} + +static void set_tensor_pair_bytes(ggml_backend_t cpu_backend, + ggml_tensor * cpu_tensor, + ggml_backend_t hrx_backend, + ggml_tensor * hrx_tensor, + const void * data, + size_t byte_count) { + set_tensor_bytes(cpu_backend, cpu_tensor, data, byte_count); + set_tensor_bytes(hrx_backend, hrx_tensor, data, byte_count); +} + +static std::vector get_f32_tensor(ggml_backend_t backend, ggml_tensor * tensor) { + REQUIRE(tensor != nullptr); + const size_t element_count = static_cast(ggml_nelements(tensor)); + std::vector data(element_count); + if (tensor->type == GGML_TYPE_F32) { + ggml_backend_tensor_get(tensor, data.data(), 0, data.size() * sizeof(float)); + } else if (tensor->type == GGML_TYPE_F16) { + std::vector f16(element_count); + ggml_backend_tensor_get(tensor, f16.data(), 0, f16.size() * sizeof(ggml_fp16_t)); + for (size_t i = 0; i < element_count; ++i) { + data[i] = ggml_fp16_to_fp32(f16[i]); + } + } else { + REQUIRE(false); + } + ggml_backend_synchronize(backend); + return data; +} + +static void require_close(const std::vector & actual, + const std::vector & expected, + float abs_tolerance, + float rel_tolerance = 0.0f) { + REQUIRE(actual.size() == expected.size()); + for (size_t i = 0; i < actual.size(); ++i) { + const float diff = std::fabs(actual[i] - expected[i]); + const float allowed = abs_tolerance + rel_tolerance * std::fabs(expected[i]); + if (diff > allowed) { + std::fprintf(stderr, "value mismatch at %zu: actual=%g expected=%g diff=%g allowed=%g\n", i, actual[i], + expected[i], diff, allowed); + std::abort(); + } + } +} + +static ggml_backend_t init_cpu_backend() { + ggml_backend_t backend = ggml_backend_init_by_type(GGML_BACKEND_DEVICE_TYPE_CPU, nullptr); + REQUIRE(backend != nullptr); + return backend; +} + +struct AttentionPostprocessGraph { + ggml_tensor * input = nullptr; + ggml_tensor * query_weight = nullptr; + ggml_tensor * key_weight = nullptr; + ggml_tensor * value_weight = nullptr; + ggml_tensor * query_norm_weight = nullptr; + ggml_tensor * key_norm_weight = nullptr; + ggml_tensor * positions = nullptr; + ggml_tensor * inverse_frequencies = nullptr; + ggml_tensor * key_cache = nullptr; + ggml_tensor * value_cache = nullptr; + ggml_tensor * key_cache_indices = nullptr; + ggml_tensor * value_cache_indices = nullptr; + ggml_tensor * attention_mask = nullptr; + ggml_tensor * query_reshape = nullptr; + ggml_tensor * query_output = nullptr; + ggml_tensor * key_output = nullptr; + ggml_tensor * value_output = nullptr; +}; + +static AttentionPostprocessGraph build_attention_postprocess_graph(ggml_context * ctx, + int64_t token_count, + int64_t query_head_count, + int64_t key_value_head_count, + int64_t cache_row_count) { + AttentionPostprocessGraph graph; + const int64_t query_size = query_head_count * kQwenFlashHeadSize; + const int64_t key_value_size = key_value_head_count * kQwenFlashHeadSize; + + graph.input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, kQwenHiddenSize, token_count); + graph.query_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_Q4_K, kQwenHiddenSize, query_size); + graph.key_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_Q4_K, kQwenHiddenSize, key_value_size); + graph.value_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_Q6_K, kQwenHiddenSize, key_value_size); + REQUIRE(graph.input != nullptr); + REQUIRE(graph.query_weight != nullptr); + REQUIRE(graph.key_weight != nullptr); + REQUIRE(graph.value_weight != nullptr); + + ggml_tensor * query_raw = ggml_mul_mat(ctx, graph.query_weight, graph.input); + ggml_tensor * key_raw = ggml_mul_mat(ctx, graph.key_weight, graph.input); + ggml_tensor * value_raw = ggml_mul_mat(ctx, graph.value_weight, graph.input); + REQUIRE(query_raw != nullptr); + REQUIRE(key_raw != nullptr); + REQUIRE(value_raw != nullptr); + + ggml_tensor * query_reshape = ggml_reshape_3d(ctx, query_raw, kQwenFlashHeadSize, query_head_count, token_count); + ggml_tensor * key_reshape = ggml_reshape_3d(ctx, key_raw, kQwenFlashHeadSize, key_value_head_count, token_count); + ggml_tensor * value_reshape = + ggml_reshape_3d(ctx, value_raw, kQwenFlashHeadSize, key_value_head_count, token_count); + REQUIRE(query_reshape != nullptr); + REQUIRE(key_reshape != nullptr); + REQUIRE(value_reshape != nullptr); + graph.query_reshape = query_reshape; + + graph.query_norm_weight = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, kQwenFlashHeadSize); + graph.key_norm_weight = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, kQwenFlashHeadSize); + graph.positions = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, token_count); + graph.inverse_frequencies = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, kQwenFlashHeadSize / 2); + REQUIRE(graph.query_norm_weight != nullptr); + REQUIRE(graph.key_norm_weight != nullptr); + REQUIRE(graph.positions != nullptr); + REQUIRE(graph.inverse_frequencies != nullptr); + + ggml_tensor * query_norm = ggml_rms_norm(ctx, query_reshape, kQwenRmsNormEps); + ggml_tensor * query_mul = ggml_mul(ctx, query_norm, graph.query_norm_weight); + REQUIRE(query_norm != nullptr); + REQUIRE(query_mul != nullptr); + graph.query_output = ggml_rope_ext(ctx, query_mul, graph.positions, graph.inverse_frequencies, kQwenFlashHeadSize, + GGML_ROPE_TYPE_NEOX, 0, 10000.0f, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f); + REQUIRE(graph.query_output != nullptr); + + ggml_tensor * key_norm = ggml_rms_norm(ctx, key_reshape, kQwenRmsNormEps); + ggml_tensor * key_mul = ggml_mul(ctx, key_norm, graph.key_norm_weight); + REQUIRE(key_norm != nullptr); + REQUIRE(key_mul != nullptr); + ggml_tensor * key_rope = ggml_rope_ext(ctx, key_mul, graph.positions, graph.inverse_frequencies, kQwenFlashHeadSize, + GGML_ROPE_TYPE_NEOX, 0, 10000.0f, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f); + REQUIRE(key_rope != nullptr); + + ggml_tensor * key_cache_rows = ggml_reshape_2d(ctx, key_rope, key_value_size, token_count); + ggml_tensor * value_cache_rows = ggml_reshape_2d(ctx, value_reshape, key_value_size, token_count); + REQUIRE(key_cache_rows != nullptr); + REQUIRE(value_cache_rows != nullptr); + + graph.key_cache = ggml_new_tensor_2d(ctx, GGML_TYPE_F16, key_value_size, cache_row_count); + graph.value_cache = ggml_new_tensor_2d(ctx, GGML_TYPE_F16, key_value_size, cache_row_count); + graph.key_cache_indices = ggml_new_tensor_1d(ctx, GGML_TYPE_I64, token_count); + graph.value_cache_indices = ggml_new_tensor_1d(ctx, GGML_TYPE_I64, token_count); + REQUIRE(graph.key_cache != nullptr); + REQUIRE(graph.value_cache != nullptr); + REQUIRE(graph.key_cache_indices != nullptr); + REQUIRE(graph.value_cache_indices != nullptr); + + graph.key_output = ggml_set_rows(ctx, graph.key_cache, key_cache_rows, graph.key_cache_indices); + graph.value_output = ggml_set_rows(ctx, graph.value_cache, value_cache_rows, graph.value_cache_indices); + REQUIRE(graph.key_output != nullptr); + REQUIRE(graph.value_output != nullptr); + return graph; +} + +static ggml_tensor * append_qwen_full_cache_flash_attention_consumer(ggml_context * ctx, + AttentionPostprocessGraph & graph, + int64_t token_count, + int64_t query_head_count, + int64_t key_value_head_count, + int64_t cache_row_count) { + ggml_tensor * query_layout = + ggml_reshape_3d(ctx, graph.query_output, kQwenFlashHeadSize, query_head_count, token_count); + ggml_tensor * query_permute = ggml_permute(ctx, query_layout, 0, 2, 1, 3); + ggml_tensor * key_cache_layout = + ggml_reshape_3d(ctx, graph.key_cache, kQwenFlashHeadSize, key_value_head_count, cache_row_count); + ggml_tensor * key_permute = ggml_permute(ctx, key_cache_layout, 0, 2, 1, 3); + ggml_tensor * value_cache_layout = + ggml_reshape_3d(ctx, graph.value_cache, kQwenFlashHeadSize, key_value_head_count, cache_row_count); + ggml_tensor * value_permute = ggml_permute(ctx, value_cache_layout, 0, 2, 1, 3); + graph.attention_mask = ggml_new_tensor_2d(ctx, GGML_TYPE_F16, cache_row_count, token_count); + REQUIRE(query_layout != nullptr); + REQUIRE(query_permute != nullptr); + REQUIRE(key_cache_layout != nullptr); + REQUIRE(key_permute != nullptr); + REQUIRE(value_cache_layout != nullptr); + REQUIRE(value_permute != nullptr); + REQUIRE(graph.attention_mask != nullptr); + return build_qwen_flash_attention_graph(ctx, query_permute, key_permute, value_permute, graph.attention_mask); +} + +struct QwenFlashAttentionLayoutGraph { + ggml_tensor * query = nullptr; + ggml_tensor * key = nullptr; + ggml_tensor * value = nullptr; + ggml_tensor * mask = nullptr; + ggml_tensor * output = nullptr; +}; + +static QwenFlashAttentionLayoutGraph build_qwen_flash_attention_layout_graph(ggml_context * ctx, + int64_t query_token_count, + int64_t key_value_token_count) { + QwenFlashAttentionLayoutGraph graph; + constexpr int64_t query_head_count = 32; + constexpr int64_t key_value_head_count = 4; + + ggml_tensor * query_storage = + ggml_new_tensor_3d(ctx, GGML_TYPE_F32, kQwenFlashHeadSize, query_head_count, query_token_count); + ggml_tensor * key_storage = + ggml_new_tensor_3d(ctx, GGML_TYPE_F16, kQwenFlashHeadSize, key_value_head_count, key_value_token_count); + ggml_tensor * value_storage = + ggml_new_tensor_3d(ctx, GGML_TYPE_F16, kQwenFlashHeadSize, key_value_head_count, key_value_token_count); + REQUIRE(query_storage != nullptr); + REQUIRE(key_storage != nullptr); + REQUIRE(value_storage != nullptr); + + graph.query = ggml_permute(ctx, query_storage, 0, 2, 1, 3); + graph.key = ggml_permute(ctx, key_storage, 0, 2, 1, 3); + graph.value = ggml_permute(ctx, value_storage, 0, 2, 1, 3); + graph.mask = ggml_new_tensor_2d(ctx, GGML_TYPE_F16, key_value_token_count, query_token_count); + REQUIRE(graph.query != nullptr); + REQUIRE(graph.key != nullptr); + REQUIRE(graph.value != nullptr); + REQUIRE(graph.mask != nullptr); + + graph.output = build_qwen_flash_attention_graph(ctx, graph.query, graph.key, graph.value, graph.mask); + return graph; +} + +static AttentionPostprocessGraph build_decode_attention_qkv_graph(ggml_context * ctx) { + constexpr int64_t token_count = 1; + constexpr int64_t query_head_count = 32; + constexpr int64_t key_value_head_count = 4; + constexpr int64_t cache_row_count = 1024; + AttentionPostprocessGraph graph = + build_attention_postprocess_graph(ctx, token_count, query_head_count, key_value_head_count, cache_row_count); + + ggml_tensor * hidden_state = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, kQwenHiddenSize, token_count); + ggml_tensor * attention_weight = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, kQwenHiddenSize); + REQUIRE(hidden_state != nullptr); + REQUIRE(attention_weight != nullptr); + ggml_tensor * attention_rms = ggml_rms_norm(ctx, hidden_state, kQwenRmsNormEps); + REQUIRE(attention_rms != nullptr); + graph.input = ggml_mul(ctx, attention_rms, attention_weight); + REQUIRE(graph.input != nullptr); + + const int64_t key_value_size = key_value_head_count * kQwenFlashHeadSize; + ggml_tensor * query_raw = ggml_mul_mat(ctx, graph.query_weight, graph.input); + ggml_tensor * key_raw = ggml_mul_mat(ctx, graph.key_weight, graph.input); + ggml_tensor * value_raw = ggml_mul_mat(ctx, graph.value_weight, graph.input); + REQUIRE(query_raw != nullptr); + REQUIRE(key_raw != nullptr); + REQUIRE(value_raw != nullptr); + + ggml_tensor * query_reshape = ggml_reshape_3d(ctx, query_raw, kQwenFlashHeadSize, query_head_count, token_count); + ggml_tensor * key_reshape = ggml_reshape_3d(ctx, key_raw, kQwenFlashHeadSize, key_value_head_count, token_count); + ggml_tensor * value_reshape = + ggml_reshape_3d(ctx, value_raw, kQwenFlashHeadSize, key_value_head_count, token_count); + REQUIRE(query_reshape != nullptr); + REQUIRE(key_reshape != nullptr); + REQUIRE(value_reshape != nullptr); + + ggml_tensor * query_norm = ggml_rms_norm(ctx, query_reshape, kQwenRmsNormEps); + ggml_tensor * query_mul = ggml_mul(ctx, query_norm, graph.query_norm_weight); + REQUIRE(query_norm != nullptr); + REQUIRE(query_mul != nullptr); + graph.query_output = ggml_rope_ext(ctx, query_mul, graph.positions, graph.inverse_frequencies, kQwenFlashHeadSize, + GGML_ROPE_TYPE_NEOX, 0, 10000.0f, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f); + REQUIRE(graph.query_output != nullptr); + + ggml_tensor * key_norm = ggml_rms_norm(ctx, key_reshape, kQwenRmsNormEps); + ggml_tensor * key_mul = ggml_mul(ctx, key_norm, graph.key_norm_weight); + REQUIRE(key_norm != nullptr); + REQUIRE(key_mul != nullptr); + ggml_tensor * key_rope = ggml_rope_ext(ctx, key_mul, graph.positions, graph.inverse_frequencies, kQwenFlashHeadSize, + GGML_ROPE_TYPE_NEOX, 0, 10000.0f, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f); + REQUIRE(key_rope != nullptr); + + ggml_tensor * key_cache_rows = ggml_reshape_2d(ctx, key_rope, key_value_size, token_count); + ggml_tensor * value_cache_rows = ggml_reshape_2d(ctx, value_reshape, key_value_size, token_count); + REQUIRE(key_cache_rows != nullptr); + REQUIRE(value_cache_rows != nullptr); + graph.key_output = ggml_set_rows(ctx, graph.key_cache, key_cache_rows, graph.key_cache_indices); + graph.value_output = ggml_set_rows(ctx, graph.value_cache, value_cache_rows, graph.value_cache_indices); + REQUIRE(graph.key_output != nullptr); + REQUIRE(graph.value_output != nullptr); + return graph; +} + +struct RoutedMoeGraph { + ggml_tensor * logits = nullptr; + ggml_tensor * input = nullptr; + ggml_tensor * gate_weight = nullptr; + ggml_tensor * up_weight = nullptr; + ggml_tensor * down_weight = nullptr; + ggml_tensor * hidden_state = nullptr; + ggml_tensor * norm_weight = nullptr; + ggml_tensor * output = nullptr; +}; + +static RoutedMoeGraph build_routed_moe_graph(ggml_context * ctx, + ggml_type down_weight_type, + bool include_next_rmsnorm) { + RoutedMoeGraph graph; + constexpr int64_t token_count = 1; + graph.logits = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, kQwenRouterExpertCount, token_count); + REQUIRE(graph.logits != nullptr); + ggml_tensor * route_ids = nullptr; + ggml_tensor * route_weights = build_qwen_router_top8_graph(ctx, graph.logits, &route_ids); + REQUIRE(route_weights != nullptr); + REQUIRE(route_ids != nullptr); + + graph.input = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, kQwenHiddenSize, 1, token_count); + graph.gate_weight = + ggml_new_tensor_3d(ctx, GGML_TYPE_Q4_K, kQwenHiddenSize, kQwenMoeIntermediate, kQwenRouterExpertCount); + graph.up_weight = + ggml_new_tensor_3d(ctx, GGML_TYPE_Q4_K, kQwenHiddenSize, kQwenMoeIntermediate, kQwenRouterExpertCount); + graph.down_weight = + ggml_new_tensor_3d(ctx, down_weight_type, kQwenMoeIntermediate, kQwenHiddenSize, kQwenRouterExpertCount); + graph.hidden_state = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, kQwenHiddenSize, token_count); + REQUIRE(graph.input != nullptr); + REQUIRE(graph.gate_weight != nullptr); + REQUIRE(graph.up_weight != nullptr); + REQUIRE(graph.down_weight != nullptr); + REQUIRE(graph.hidden_state != nullptr); + + ggml_tensor * gate = ggml_mul_mat_id(ctx, graph.gate_weight, graph.input, route_ids); + ggml_tensor * up = ggml_mul_mat_id(ctx, graph.up_weight, graph.input, route_ids); + REQUIRE(gate != nullptr); + REQUIRE(up != nullptr); + ggml_tensor * glu = ggml_glu_split(ctx, gate, up, GGML_GLU_OP_SWIGLU); + REQUIRE(glu != nullptr); + ggml_tensor * down = ggml_mul_mat_id(ctx, graph.down_weight, glu, route_ids); + REQUIRE(down != nullptr); + ggml_tensor * weighted = ggml_mul(ctx, down, route_weights); + REQUIRE(weighted != nullptr); + + std::vector route_views; + route_views.reserve(kQwenRouterRouteCount); + for (int64_t route = 0; route < kQwenRouterRouteCount; ++route) { + ggml_tensor * view = ggml_view_2d(ctx, weighted, kQwenHiddenSize, token_count, weighted->nb[2], + static_cast(route) * weighted->nb[1]); + REQUIRE(view != nullptr); + route_views.push_back(view); + } + + ggml_tensor * reduced = route_views.front(); + for (size_t i = 1; i < route_views.size(); ++i) { + reduced = ggml_add(ctx, reduced, route_views[i]); + REQUIRE(reduced != nullptr); + } + + ggml_tensor * residual = ggml_add(ctx, graph.hidden_state, reduced); + REQUIRE(residual != nullptr); + graph.output = residual; + + if (include_next_rmsnorm) { + graph.norm_weight = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, kQwenHiddenSize); + REQUIRE(graph.norm_weight != nullptr); + ggml_tensor * rms = ggml_rms_norm(ctx, residual, kQwenRmsNormEps); + REQUIRE(rms != nullptr); + graph.output = ggml_mul(ctx, rms, graph.norm_weight); + REQUIRE(graph.output != nullptr); + } + + return graph; +} + +static void run_rmsnorm_support_checks() { + ggml_backend_hrx_device_context device_context = {}; + ggml_backend_hrx_context backend_context = {}; + device_context.architecture = "gfx1151"; + backend_context.device = &device_context; + const ggml::hrx::GraphExecutor executor(backend_context); + + ggml_init_params params = {}; + params.mem_size = 1024 * 1024; + params.no_alloc = true; + ggml_context * ctx = ggml_init(params); + REQUIRE(ctx != nullptr); + + ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 256, 1); + ggml_tensor * weight = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 256); + REQUIRE(input != nullptr); + REQUIRE(weight != nullptr); + ggml_tensor * output = build_rmsnorm_mul_graph(ctx, input, weight); + ggml_cgraph * graph = ggml_new_graph(ctx); + REQUIRE(graph != nullptr); + ggml_build_forward_expand(graph, output); + const ggml::hrx::GraphSupportResult support = executor.can_execute(*graph); + REQUIRE(support.supported); + REQUIRE(support.status.success()); + + ggml_tensor * wrong_eps_output = build_rmsnorm_mul_graph(ctx, input, weight, 1.0e-5f); + ggml_cgraph * wrong_eps_graph = ggml_new_graph(ctx); + REQUIRE(wrong_eps_graph != nullptr); + ggml_build_forward_expand(wrong_eps_graph, wrong_eps_output); + const ggml::hrx::GraphSupportResult wrong_eps_support = executor.can_execute(*wrong_eps_graph); + REQUIRE(!wrong_eps_support.supported); + + ggml_tensor * wrong_type_input = ggml_new_tensor_2d(ctx, GGML_TYPE_F16, 256, 1); + ggml_tensor * wrong_type_weight = ggml_new_tensor_1d(ctx, GGML_TYPE_F16, 256); + REQUIRE(wrong_type_input != nullptr); + REQUIRE(wrong_type_weight != nullptr); + ggml_tensor * wrong_type_output = build_rmsnorm_mul_graph(ctx, wrong_type_input, wrong_type_weight); + ggml_cgraph * wrong_type_graph = ggml_new_graph(ctx); + REQUIRE(wrong_type_graph != nullptr); + ggml_build_forward_expand(wrong_type_graph, wrong_type_output); + const ggml::hrx::GraphSupportResult wrong_type_support = executor.can_execute(*wrong_type_graph); + REQUIRE(!wrong_type_support.supported); + + ggml_tensor * wrong_weight = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 128); + REQUIRE(wrong_weight != nullptr); + ggml_tensor * wrong_weight_output = build_rmsnorm_mul_graph(ctx, input, wrong_weight); + ggml_cgraph * wrong_weight_graph = ggml_new_graph(ctx); + REQUIRE(wrong_weight_graph != nullptr); + ggml_build_forward_expand(wrong_weight_graph, wrong_weight_output); + const ggml::hrx::GraphSupportResult wrong_weight_support = executor.can_execute(*wrong_weight_graph); + REQUIRE(!wrong_weight_support.supported); + + ggml_free(ctx); +} + +static void run_rmsnorm_mul_case(int64_t hidden_size, int64_t token_count) { + ggml_backend_t backend = ggml_backend_hrx_init(0); + REQUIRE(backend != nullptr); + + ggml_init_params params = {}; + params.mem_size = static_cast(hidden_size * token_count * sizeof(float) * 8 + 1024 * 1024); + params.no_alloc = true; + ggml_context * ctx = ggml_init(params); + REQUIRE(ctx != nullptr); + + ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, hidden_size, token_count); + ggml_tensor * weight = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, hidden_size); + REQUIRE(input != nullptr); + REQUIRE(weight != nullptr); + ggml_tensor * output = build_rmsnorm_mul_graph(ctx, input, weight); + + ggml_cgraph * graph = ggml_new_graph(ctx); + REQUIRE(graph != nullptr); + ggml_build_forward_expand(graph, output); + + require_kernel_subsequence(scheduled_kernel_sequence(graph), { "qwen3_moe:qwen3_moe_rmsnorm_f32" }); + + ggml_backend_buffer_t buffer = ggml_backend_alloc_ctx_tensors(ctx, backend); + REQUIRE(buffer != nullptr); + + const std::vector input_data = make_input(hidden_size, token_count); + const std::vector weight_data = make_weight(hidden_size); + const std::vector expected = rmsnorm_mul_reference(input_data, weight_data, hidden_size, token_count); + + ggml_backend_tensor_set(input, input_data.data(), 0, input_data.size() * sizeof(float)); + ggml_backend_tensor_set(weight, weight_data.data(), 0, weight_data.size() * sizeof(float)); + + REQUIRE(ggml_backend_graph_compute(backend, graph) == GGML_STATUS_SUCCESS); + ggml_backend_synchronize(backend); + + std::vector actual(expected.size()); + ggml_backend_tensor_get(output, actual.data(), 0, actual.size() * sizeof(float)); + for (size_t i = 0; i < actual.size(); ++i) { + const float diff = std::fabs(actual[i] - expected[i]); + REQUIRE(diff <= 5.0e-4f); + } + + ggml_backend_buffer_free(buffer); + ggml_free(ctx); + ggml_backend_free(backend); +} + +static void run_router_projection_case(int64_t token_count) { + static constexpr int64_t kHiddenSize = 2048; + static constexpr int64_t kExpertCount = 128; + + ggml_backend_t backend = ggml_backend_hrx_init(0); + REQUIRE(backend != nullptr); + + ggml_init_params params = {}; + params.mem_size = static_cast( + (kHiddenSize * token_count + kHiddenSize * kExpertCount + kExpertCount * token_count) * sizeof(float) * 4 + + 1024 * 1024); + params.no_alloc = true; + ggml_context * ctx = ggml_init(params); + REQUIRE(ctx != nullptr); + + ggml_tensor * weight = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, kHiddenSize, kExpertCount); + ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, kHiddenSize, token_count); + REQUIRE(weight != nullptr); + REQUIRE(input != nullptr); + ggml_tensor * output = ggml_mul_mat(ctx, weight, input); + REQUIRE(output != nullptr); + + ggml_cgraph * graph = ggml_new_graph(ctx); + REQUIRE(graph != nullptr); + ggml_build_forward_expand(graph, output); + + require_kernel_subsequence(scheduled_kernel_sequence(graph), + { "qwen3_moe:qwen3_moe_router_projection_f32_four_row_wave32" }); + + ggml_backend_buffer_t buffer = ggml_backend_alloc_ctx_tensors(ctx, backend); + REQUIRE(buffer != nullptr); + + const std::vector input_data = make_router_input(kHiddenSize, token_count); + const std::vector weight_data = make_router_weight(kHiddenSize, kExpertCount); + const std::vector expected = + router_projection_reference(input_data, weight_data, kHiddenSize, kExpertCount, token_count); + + ggml_backend_tensor_set(input, input_data.data(), 0, input_data.size() * sizeof(float)); + ggml_backend_tensor_set(weight, weight_data.data(), 0, weight_data.size() * sizeof(float)); + + REQUIRE(ggml_backend_graph_compute(backend, graph) == GGML_STATUS_SUCCESS); + ggml_backend_synchronize(backend); + + std::vector actual(expected.size()); + ggml_backend_tensor_get(output, actual.data(), 0, actual.size() * sizeof(float)); + for (size_t i = 0; i < actual.size(); ++i) { + const float diff = std::fabs(actual[i] - expected[i]); + REQUIRE(diff <= 1.0e-2f); + } + + ggml_backend_buffer_free(buffer); + ggml_free(ctx); + ggml_backend_free(backend); +} + +static void run_router_top8_case(int64_t token_count) { + ggml_backend_t backend = ggml_backend_hrx_init(0); + REQUIRE(backend != nullptr); + + ggml_init_params params = {}; + params.mem_size = static_cast(kQwenRouterExpertCount * token_count * sizeof(float) * 16 + 1024 * 1024); + params.no_alloc = true; + ggml_context * ctx = ggml_init(params); + REQUIRE(ctx != nullptr); + + ggml_tensor * logits = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, kQwenRouterExpertCount, token_count); + REQUIRE(logits != nullptr); + ggml_tensor * output = build_qwen_router_top8_graph(ctx, logits); + + ggml_cgraph * graph = ggml_new_graph(ctx); + REQUIRE(graph != nullptr); + ggml_build_forward_expand(graph, output); + + require_kernel_subsequence(scheduled_kernel_sequence(graph), { "qwen3_moe:qwen3_moe_router_top8_f32" }); + + ggml_backend_buffer_t buffer = ggml_backend_alloc_ctx_tensors(ctx, backend); + REQUIRE(buffer != nullptr); + + const std::vector logits_data = make_router_logits(token_count); + const std::vector expected = router_top8_weights_reference(logits_data, token_count); + + ggml_backend_tensor_set(logits, logits_data.data(), 0, logits_data.size() * sizeof(float)); + + REQUIRE(ggml_backend_graph_compute(backend, graph) == GGML_STATUS_SUCCESS); + ggml_backend_synchronize(backend); + + std::vector actual(expected.size()); + ggml_backend_tensor_get(output, actual.data(), 0, actual.size() * sizeof(float)); + for (size_t i = 0; i < actual.size(); ++i) { + const float diff = std::fabs(actual[i] - expected[i]); + REQUIRE(diff <= 1.0e-5f); + } + + ggml_backend_buffer_free(buffer); + ggml_free(ctx); + ggml_backend_free(backend); +} + +static void run_qwen_flash_attention_case() { + static constexpr int64_t kQueryTokenCount = 2; + static constexpr int64_t kKeyValueTokenCount = 4; + + ggml_backend_t backend = ggml_backend_hrx_init(0); + REQUIRE(backend != nullptr); + + ggml_init_params params = {}; + params.mem_size = 2 * 1024 * 1024; + params.no_alloc = true; + ggml_context * ctx = ggml_init(params); + REQUIRE(ctx != nullptr); + + ggml_tensor * query = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, kQwenFlashHeadSize, kQueryTokenCount, 1); + ggml_tensor * key = ggml_new_tensor_3d(ctx, GGML_TYPE_F16, kQwenFlashHeadSize, kKeyValueTokenCount, 1); + ggml_tensor * value = ggml_new_tensor_3d(ctx, GGML_TYPE_F16, kQwenFlashHeadSize, kKeyValueTokenCount, 1); + ggml_tensor * mask = ggml_new_tensor_2d(ctx, GGML_TYPE_F16, kKeyValueTokenCount, kQueryTokenCount); + REQUIRE(query != nullptr); + REQUIRE(key != nullptr); + REQUIRE(value != nullptr); + REQUIRE(mask != nullptr); + ggml_tensor * output = build_qwen_flash_attention_graph(ctx, query, key, value, mask); + + ggml_cgraph * graph = ggml_new_graph(ctx); + REQUIRE(graph != nullptr); + ggml_build_forward_expand(graph, output); + + require_kernel_subsequence(scheduled_kernel_sequence(graph), + { "qwen3_moe:qwen3_moe_flash_attention_f32_f16_wmma" }); + + ggml_backend_buffer_t buffer = ggml_backend_alloc_ctx_tensors(ctx, backend); + REQUIRE(buffer != nullptr); + + const std::vector query_data = make_flash_query(kQueryTokenCount); + const std::vector key_data = make_flash_key_value(kKeyValueTokenCount, 3); + const std::vector value_data = make_flash_key_value(kKeyValueTokenCount, 11); + const std::vector mask_data = make_flash_mask(kQueryTokenCount, kKeyValueTokenCount); + const std::vector expected = + flash_attention_reference(query_data, key_data, value_data, mask_data, kQueryTokenCount, kKeyValueTokenCount); + + ggml_backend_tensor_set(query, query_data.data(), 0, query_data.size() * sizeof(float)); + ggml_backend_tensor_set(key, key_data.data(), 0, key_data.size() * sizeof(ggml_fp16_t)); + ggml_backend_tensor_set(value, value_data.data(), 0, value_data.size() * sizeof(ggml_fp16_t)); + ggml_backend_tensor_set(mask, mask_data.data(), 0, mask_data.size() * sizeof(ggml_fp16_t)); + + REQUIRE(ggml_backend_graph_compute(backend, graph) == GGML_STATUS_SUCCESS); + ggml_backend_synchronize(backend); + + std::vector actual(expected.size()); + ggml_backend_tensor_get(output, actual.data(), 0, actual.size() * sizeof(float)); + for (size_t i = 0; i < actual.size(); ++i) { + const float diff = std::fabs(actual[i] - expected[i]); + REQUIRE(diff <= 5.0e-2f); + } + + ggml_backend_buffer_free(buffer); + ggml_free(ctx); + ggml_backend_free(backend); +} + +static void run_qwen_full_cache_prefill_flash_attention_scheduling_case(int64_t token_count) { + constexpr int64_t kQueryHeadCount = 32; + constexpr int64_t kKeyValueHeadCount = 4; + constexpr int64_t kFullCacheRowCount = 40960; + const size_t active_mask_byte_count = + static_cast(token_count) * static_cast(token_count) * sizeof(ggml_fp16_t); + + ggml_init_params params = {}; + params.mem_size = 128 * 1024 * 1024; + params.no_alloc = true; + ggml_context * ctx = ggml_init(params); + REQUIRE(ctx != nullptr); + + AttentionPostprocessGraph attention = + build_attention_postprocess_graph(ctx, token_count, kQueryHeadCount, kKeyValueHeadCount, kFullCacheRowCount); + ggml_tensor * flash_output = append_qwen_full_cache_flash_attention_consumer( + ctx, attention, token_count, kQueryHeadCount, kKeyValueHeadCount, kFullCacheRowCount); + + ggml_cgraph * graph = ggml_new_graph(ctx); + REQUIRE(graph != nullptr); + ggml_build_forward_expand(graph, attention.key_output); + ggml_build_forward_expand(graph, attention.value_output); + ggml_build_forward_expand(graph, flash_output); + + ggml::hrx::GraphImportResult imported = ggml::hrx::import_ggml_graph(*graph); + REQUIRE(imported.valid()); + const ggml::hrx::DispatchRegistry * registry = ggml::hrx::find_dispatch_registry({ "gfx1151" }); + REQUIRE(registry != nullptr); + + std::vector covered_nodes(imported.graph.nodes().size(), false); + ggml::hrx::CommandPlan plan; + ggml::hrx::DispatchMatch postprocess_match; + match_dispatch_at_index(imported.graph, *registry, plan, covered_nodes, + producer_index_for_tensor(imported.graph, attention.query_reshape), postprocess_match); + ggml::hrx::DispatchMatch flash_match; + match_dispatch_at_index(imported.graph, *registry, plan, covered_nodes, + producer_index_for_tensor(imported.graph, flash_output), flash_match); + + const ggml::hrx::Value * mask_value = imported.graph.values().find_tensor(attention.attention_mask); + REQUIRE(mask_value != nullptr); + const ggml::hrx::CommandPlanAlternateValue * compact_mask = + ggml::hrx::find_alternate_value(plan, mask_value->id, GGML_TYPE_F16, active_mask_byte_count); + REQUIRE(compact_mask != nullptr); + + const ggml::hrx::Dispatch * metadata_dispatch = nullptr; + for (const ggml::hrx::Dispatch & dispatch : plan.initialization_dispatches) { + if (kernel_name_for_id(dispatch.kernel.kernel_id) == "qwen3_moe:qwen_attention_metadata") { + metadata_dispatch = &dispatch; + } + } + REQUIRE(metadata_dispatch != nullptr); + REQUIRE(metadata_dispatch->kernel.integer_parameters.at("token_count") == token_count); + REQUIRE(metadata_dispatch->kernel.integer_parameters.at("context_capacity") == token_count); + REQUIRE(metadata_dispatch->bindings.size() == 5); + REQUIRE(metadata_dispatch->bindings[4].value == compact_mask->alternate_value); + REQUIRE(metadata_dispatch->bindings[4].length == active_mask_byte_count); + + const ggml::hrx::Dispatch * flash_dispatch = nullptr; + for (const ggml::hrx::Dispatch & dispatch : plan.dispatches) { + if (kernel_name_for_id(dispatch.kernel.kernel_id) == "qwen3_moe:qwen3_moe_flash_attention_f32_f16_wmma") { + flash_dispatch = &dispatch; + } + } + REQUIRE(flash_dispatch != nullptr); + REQUIRE(flash_dispatch->kernel.integer_parameters.at("query_token_count") == token_count); + REQUIRE(flash_dispatch->kernel.integer_parameters.at("key_value_token_count") == token_count); + REQUIRE(flash_dispatch->bindings.size() == 5); + REQUIRE(flash_dispatch->bindings[3].value == compact_mask->alternate_value); + REQUIRE(flash_dispatch->bindings[3].length == active_mask_byte_count); + + ggml_free(ctx); +} + +static void run_qwen_decode_split_flash_attention_scheduling_case(int64_t query_token_count, + int64_t key_value_token_count) { + constexpr int64_t query_head_count = 32; + constexpr int64_t key_value_head_count = 4; + constexpr int64_t hidden_size = query_head_count * kQwenFlashHeadSize; + const size_t q8_row_bytes = ggml_row_size(GGML_TYPE_Q8_1, hidden_size); + const size_t q8_output_bytes = static_cast(query_token_count) * q8_row_bytes; + const int64_t key_value_capacity = (key_value_token_count + 63) / 64 * 64; + const int64_t key_value_blocks = key_value_capacity / 64; + const size_t partial_scalar_bytes = + static_cast(key_value_head_count * key_value_blocks * 16) * sizeof(float); + const size_t partial_output_bytes = + static_cast(key_value_head_count * key_value_blocks * 16 * kQwenFlashHeadSize) * sizeof(ggml_fp16_t); + + ggml_init_params params = {}; + params.mem_size = 128 * 1024 * 1024; + params.no_alloc = true; + ggml_context * ctx = ggml_init(params); + REQUIRE(ctx != nullptr); + + QwenFlashAttentionLayoutGraph attention = + build_qwen_flash_attention_layout_graph(ctx, query_token_count, key_value_token_count); + + ggml_cgraph * cgraph = ggml_new_graph(ctx); + REQUIRE(cgraph != nullptr); + ggml_build_forward_expand(cgraph, attention.output); + + ggml::hrx::GraphImportResult imported = ggml::hrx::import_ggml_graph(*cgraph); + REQUIRE(imported.valid()); + ggml::hrx::DispatchScheduler scheduler; + ggml::hrx::DispatchScheduleDiagnostics diagnostics; + REQUIRE(scheduler.schedule_graph(imported.graph, { "gfx1151" }, &diagnostics)); + const ggml::hrx::CommandPlan & plan = scheduler.plan(); + REQUIRE(plan.valid()); + REQUIRE(plan.dispatches.size() == static_cast(query_token_count)); + REQUIRE(plan.transients.size() == 4); + REQUIRE(plan.completion_counter_requests.size() == 1); + + REQUIRE(plan.transients[0].size == partial_scalar_bytes); + REQUIRE(plan.transients[1].size == partial_scalar_bytes); + REQUIRE(plan.transients[2].size == partial_output_bytes); + REQUIRE(plan.transients[3].size == q8_output_bytes); + REQUIRE(plan.completion_counter_requests[0].count == key_value_head_count); + + const ggml::hrx::Value * query_value = imported.graph.values().find_tensor(attention.query); + const ggml::hrx::Value * mask_value = imported.graph.values().find_tensor(attention.mask); + const ggml::hrx::Value * output_value = imported.graph.values().find_tensor(attention.output); + REQUIRE(query_value != nullptr); + REQUIRE(mask_value != nullptr); + REQUIRE(output_value != nullptr); + const ggml::hrx::CommandPlanAlternateValue * q8_alternate = + plan.metadata.find_alternate_value(output_value->id, GGML_TYPE_Q8_1, q8_output_bytes); + REQUIRE(q8_alternate != nullptr); + REQUIRE(q8_alternate->alternate_value == plan.transients[3].value); + + const size_t query_row_bytes = static_cast(hidden_size) * sizeof(float); + const size_t mask_row_bytes = static_cast(key_value_token_count) * sizeof(ggml_fp16_t); + const size_t output_row_bytes = query_row_bytes; + for (int64_t row = 0; row < query_token_count; ++row) { + const ggml::hrx::Dispatch & dispatch = plan.dispatches[static_cast(row)]; + REQUIRE(kernel_name_for_id(dispatch.kernel.kernel_id) == + "qwen3_moe:qwen3_moe_flash_attention_decode_split_f32_f16_wmma_next_q8"); + REQUIRE(dispatch.kernel.integer_parameters.at("key_value_token_count") == key_value_token_count); + REQUIRE(dispatch.kernel.compile_parameters.at("qwen3_moe.attention.key_value_token_capacity") == + std::to_string(key_value_capacity)); + REQUIRE(dispatch.bindings.size() == 10); + REQUIRE(dispatch.bindings[0].value == query_value->id); + REQUIRE(dispatch.bindings[0].offset == static_cast(row) * query_value->nb[1]); + REQUIRE(dispatch.bindings[0].length == query_row_bytes); + REQUIRE(dispatch.bindings[3].value == mask_value->id); + REQUIRE(dispatch.bindings[3].offset == static_cast(row) * mask_value->nb[1]); + REQUIRE(dispatch.bindings[3].length == mask_row_bytes); + REQUIRE(dispatch.bindings[8].value == output_value->id); + REQUIRE(dispatch.bindings[8].offset == static_cast(row) * output_value->nb[2]); + REQUIRE(dispatch.bindings[8].length == output_row_bytes); + REQUIRE(dispatch.bindings[9].value == q8_alternate->alternate_value); + REQUIRE(dispatch.bindings[9].offset == static_cast(row) * q8_row_bytes); + REQUIRE(dispatch.bindings[9].length == q8_row_bytes); + } + + ggml_free(ctx); +} + +static void run_qwen_decode_attention_output_next_q8_scheduling_case(bool include_get_rows_selectors) { + constexpr int64_t query_token_count = 1; + constexpr int64_t key_value_token_count = 512; + constexpr int64_t attention_hidden_size = 4096; + const size_t attention_q8_bytes = ggml_row_size(GGML_TYPE_Q8_1, attention_hidden_size); + const size_t next_q8_bytes = ggml_row_size(GGML_TYPE_Q8_1, kQwenHiddenSize); + + ggml_init_params params = {}; + params.mem_size = 128 * 1024 * 1024; + params.no_alloc = true; + ggml_context * ctx = ggml_init(params); + REQUIRE(ctx != nullptr); + + QwenFlashAttentionLayoutGraph attention = + build_qwen_flash_attention_layout_graph(ctx, query_token_count, key_value_token_count); + ggml_tensor * attention_output = ggml_reshape_2d(ctx, attention.output, attention_hidden_size, query_token_count); + ggml_tensor * output_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_Q4_K, attention_hidden_size, kQwenHiddenSize); + ggml_tensor * projection = ggml_mul_mat(ctx, output_weight, attention_output); + ggml_tensor * residual_input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, kQwenHiddenSize, query_token_count); + ggml_tensor * selected_projection = projection; + ggml_tensor * selected_residual = residual_input; + ggml_tensor * row_indices = nullptr; + if (include_get_rows_selectors) { + row_indices = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, 1); + selected_projection = ggml_get_rows(ctx, projection, row_indices); + selected_residual = ggml_get_rows(ctx, residual_input, row_indices); + REQUIRE(row_indices != nullptr); + REQUIRE(selected_projection != nullptr); + REQUIRE(selected_residual != nullptr); + } + ggml_tensor * residual = ggml_add(ctx, selected_projection, selected_residual); + ggml_tensor * norm_weight = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, kQwenHiddenSize); + ggml_tensor * rms = ggml_rms_norm(ctx, residual, kQwenRmsNormEps); + ggml_tensor * normalized = ggml_mul(ctx, rms, norm_weight); + REQUIRE(attention_output != nullptr); + REQUIRE(output_weight != nullptr); + REQUIRE(projection != nullptr); + REQUIRE(residual_input != nullptr); + REQUIRE(residual != nullptr); + REQUIRE(norm_weight != nullptr); + REQUIRE(rms != nullptr); + REQUIRE(normalized != nullptr); + + ggml_cgraph * cgraph = ggml_new_graph(ctx); + REQUIRE(cgraph != nullptr); + ggml_build_forward_expand(cgraph, normalized); + + ggml::hrx::GraphImportResult imported = ggml::hrx::import_ggml_graph(*cgraph); + REQUIRE(imported.valid()); + ggml::hrx::DispatchScheduler scheduler; + ggml::hrx::DispatchScheduleDiagnostics diagnostics; + REQUIRE(scheduler.schedule_graph(imported.graph, { "gfx1151" }, &diagnostics)); + const ggml::hrx::CommandPlan & plan = scheduler.plan(); + REQUIRE(plan.valid()); + REQUIRE(plan.dispatches.size() == 2); + REQUIRE(plan.completion_counter_requests.size() == 2); + + const ggml::hrx::Dispatch & projection_dispatch = plan.dispatches.back(); + REQUIRE(kernel_name_for_id(projection_dispatch.kernel.kernel_id) == + "qwen3_moe:qwen3_moe_dense_linear_q4k_q8_1_x4_next_q8"); + REQUIRE(projection_dispatch.kernel.integer_parameters.at("token_count") == query_token_count); + REQUIRE(projection_dispatch.kernel.compile_parameters.at("qwen3_moe.dense_quantized.input_size") == + std::to_string(attention_hidden_size)); + REQUIRE(projection_dispatch.kernel.compile_parameters.at("qwen3_moe.dense_quantized.output_size") == + std::to_string(kQwenHiddenSize)); + REQUIRE(projection_dispatch.kernel.compile_parameters.at("qwen3_moe.dense_quantized.output_accumulation") == "1"); + REQUIRE(projection_dispatch.bindings.size() == 7); + + const ggml::hrx::Value * flash_output_value = imported.graph.values().find_tensor(attention.output); + const ggml::hrx::Value * residual_input_value = imported.graph.values().find_tensor(residual_input); + const ggml::hrx::Value * residual_value = imported.graph.values().find_tensor(residual); + const ggml::hrx::Value * normalized_value = imported.graph.values().find_tensor(normalized); + REQUIRE(flash_output_value != nullptr); + REQUIRE(residual_input_value != nullptr); + REQUIRE(residual_value != nullptr); + REQUIRE(normalized_value != nullptr); + const ggml::hrx::CommandPlanAlternateValue * attention_q8 = + ggml::hrx::find_alternate_value(plan, flash_output_value->id, GGML_TYPE_Q8_1, attention_q8_bytes); + const ggml::hrx::CommandPlanAlternateValue * next_q8 = + ggml::hrx::find_alternate_value(plan, normalized_value->id, GGML_TYPE_Q8_1, next_q8_bytes); + REQUIRE(attention_q8 != nullptr); + REQUIRE(next_q8 != nullptr); + REQUIRE(projection_dispatch.bindings[0].value == attention_q8->alternate_value); + REQUIRE(projection_dispatch.bindings[2].value == residual_value->id); + REQUIRE(projection_dispatch.bindings[4].value == normalized_value->id); + REQUIRE(projection_dispatch.bindings[6].value == next_q8->alternate_value); + + const ggml::hrx::Value * aliased_residual = imported.graph.values().find(residual_value->id); + REQUIRE(aliased_residual != nullptr); + REQUIRE(aliased_residual->alias_source == residual_input_value->id); + + ggml_free(ctx); +} + +static void run_add_f32_cpu_reference_case() { + ggml_backend_t cpu_backend = init_cpu_backend(); + ggml_backend_t hrx_backend = ggml_backend_hrx_init(0); + REQUIRE(hrx_backend != nullptr); + + ggml_init_params cpu_params = {}; + cpu_params.mem_size = 256 * 1024; + cpu_params.no_alloc = true; + ggml_init_params hrx_params = cpu_params; + ggml_context * cpu_ctx = ggml_init(cpu_params); + ggml_context * hrx_ctx = ggml_init(hrx_params); + REQUIRE(cpu_ctx != nullptr); + REQUIRE(hrx_ctx != nullptr); + + constexpr int64_t element_count = 257; + ggml_tensor * cpu_a = ggml_new_tensor_1d(cpu_ctx, GGML_TYPE_F32, element_count); + ggml_tensor * cpu_b = ggml_new_tensor_1d(cpu_ctx, GGML_TYPE_F32, element_count); + ggml_tensor * cpu_output = ggml_add(cpu_ctx, cpu_a, cpu_b); + ggml_tensor * hrx_a = ggml_new_tensor_1d(hrx_ctx, GGML_TYPE_F32, element_count); + ggml_tensor * hrx_b = ggml_new_tensor_1d(hrx_ctx, GGML_TYPE_F32, element_count); + ggml_tensor * hrx_output = ggml_add(hrx_ctx, hrx_a, hrx_b); + REQUIRE(cpu_output != nullptr); + REQUIRE(hrx_output != nullptr); + + ggml_cgraph * cpu_graph = ggml_new_graph(cpu_ctx); + ggml_cgraph * hrx_graph = ggml_new_graph(hrx_ctx); + REQUIRE(cpu_graph != nullptr); + REQUIRE(hrx_graph != nullptr); + ggml_build_forward_expand(cpu_graph, cpu_output); + ggml_build_forward_expand(hrx_graph, hrx_output); + + require_kernel_subsequence(scheduled_kernel_sequence(hrx_graph), { "qwen3_moe:ggml_add_f32" }); + + ggml_backend_buffer_t cpu_buffer = ggml_backend_alloc_ctx_tensors(cpu_ctx, cpu_backend); + ggml_backend_buffer_t hrx_buffer = ggml_backend_alloc_ctx_tensors(hrx_ctx, hrx_backend); + REQUIRE(cpu_buffer != nullptr); + REQUIRE(hrx_buffer != nullptr); + + const std::vector a = make_pattern_f32(element_count, 1, 0.125f); + const std::vector b = make_pattern_f32(element_count, 2, 0.25f); + set_tensor_pair_bytes(cpu_backend, cpu_a, hrx_backend, hrx_a, a.data(), a.size() * sizeof(float)); + set_tensor_pair_bytes(cpu_backend, cpu_b, hrx_backend, hrx_b, b.data(), b.size() * sizeof(float)); + + REQUIRE(ggml_backend_graph_compute(cpu_backend, cpu_graph) == GGML_STATUS_SUCCESS); + REQUIRE(ggml_backend_graph_compute(hrx_backend, hrx_graph) == GGML_STATUS_SUCCESS); + ggml_backend_synchronize(cpu_backend); + ggml_backend_synchronize(hrx_backend); + require_close(get_f32_tensor(hrx_backend, hrx_output), get_f32_tensor(cpu_backend, cpu_output), 0.0f); + + ggml_backend_buffer_free(cpu_buffer); + ggml_backend_buffer_free(hrx_buffer); + ggml_free(cpu_ctx); + ggml_free(hrx_ctx); + ggml_backend_free(cpu_backend); + ggml_backend_free(hrx_backend); +} + +static void run_gather_add_f32_cpu_reference_case() { + ggml_backend_t cpu_backend = init_cpu_backend(); + ggml_backend_t hrx_backend = ggml_backend_hrx_init(0); + REQUIRE(hrx_backend != nullptr); + + ggml_init_params params = {}; + params.mem_size = 512 * 1024; + params.no_alloc = true; + ggml_context * cpu_ctx = ggml_init(params); + ggml_context * hrx_ctx = ggml_init(params); + REQUIRE(cpu_ctx != nullptr); + REQUIRE(hrx_ctx != nullptr); + + constexpr int64_t hidden_size = 256; + constexpr int64_t source_token_count = 11; + constexpr int64_t output_token_count = 5; + ggml_tensor * cpu_a = ggml_new_tensor_2d(cpu_ctx, GGML_TYPE_F32, hidden_size, source_token_count); + ggml_tensor * cpu_b = ggml_new_tensor_2d(cpu_ctx, GGML_TYPE_F32, hidden_size, source_token_count); + ggml_tensor * cpu_ids = ggml_new_tensor_1d(cpu_ctx, GGML_TYPE_I32, output_token_count); + ggml_tensor * cpu_rows_a = ggml_get_rows(cpu_ctx, cpu_a, cpu_ids); + ggml_tensor * cpu_rows_b = ggml_get_rows(cpu_ctx, cpu_b, cpu_ids); + ggml_tensor * cpu_output = ggml_add(cpu_ctx, cpu_rows_a, cpu_rows_b); + + ggml_tensor * hrx_a = ggml_new_tensor_2d(hrx_ctx, GGML_TYPE_F32, hidden_size, source_token_count); + ggml_tensor * hrx_b = ggml_new_tensor_2d(hrx_ctx, GGML_TYPE_F32, hidden_size, source_token_count); + ggml_tensor * hrx_ids = ggml_new_tensor_1d(hrx_ctx, GGML_TYPE_I32, output_token_count); + ggml_tensor * hrx_rows_a = ggml_get_rows(hrx_ctx, hrx_a, hrx_ids); + ggml_tensor * hrx_rows_b = ggml_get_rows(hrx_ctx, hrx_b, hrx_ids); + ggml_tensor * hrx_output = ggml_add(hrx_ctx, hrx_rows_a, hrx_rows_b); + REQUIRE(cpu_output != nullptr); + REQUIRE(hrx_output != nullptr); + + ggml_cgraph * cpu_graph = ggml_new_graph(cpu_ctx); + ggml_cgraph * hrx_graph = ggml_new_graph(hrx_ctx); + REQUIRE(cpu_graph != nullptr); + REQUIRE(hrx_graph != nullptr); + ggml_build_forward_expand(cpu_graph, cpu_output); + ggml_build_forward_expand(hrx_graph, hrx_output); + + require_kernel_subsequence(scheduled_kernel_sequence(hrx_graph), { "qwen3_moe:ggml_gather_add_f32" }); + + ggml_backend_buffer_t cpu_buffer = ggml_backend_alloc_ctx_tensors(cpu_ctx, cpu_backend); + ggml_backend_buffer_t hrx_buffer = ggml_backend_alloc_ctx_tensors(hrx_ctx, hrx_backend); + REQUIRE(cpu_buffer != nullptr); + REQUIRE(hrx_buffer != nullptr); + + const std::vector a = make_pattern_f32(hidden_size * source_token_count, 3, 0.05f); + const std::vector b = make_pattern_f32(hidden_size * source_token_count, 4, 0.075f); + const std::vector ids = { 9, 3, 7, 1, 5 }; + set_tensor_pair_bytes(cpu_backend, cpu_a, hrx_backend, hrx_a, a.data(), a.size() * sizeof(float)); + set_tensor_pair_bytes(cpu_backend, cpu_b, hrx_backend, hrx_b, b.data(), b.size() * sizeof(float)); + set_tensor_pair_bytes(cpu_backend, cpu_ids, hrx_backend, hrx_ids, ids.data(), ids.size() * sizeof(int32_t)); + + REQUIRE(ggml_backend_graph_compute(cpu_backend, cpu_graph) == GGML_STATUS_SUCCESS); + REQUIRE(ggml_backend_graph_compute(hrx_backend, hrx_graph) == GGML_STATUS_SUCCESS); + ggml_backend_synchronize(cpu_backend); + ggml_backend_synchronize(hrx_backend); + require_close(get_f32_tensor(hrx_backend, hrx_output), get_f32_tensor(cpu_backend, cpu_output), 0.0f); + + ggml_backend_buffer_free(cpu_buffer); + ggml_backend_buffer_free(hrx_buffer); + ggml_free(cpu_ctx); + ggml_free(hrx_ctx); + ggml_backend_free(cpu_backend); + ggml_backend_free(hrx_backend); +} + +static void run_token_embedding_q4k_cpu_reference_case() { + ggml_backend_t cpu_backend = init_cpu_backend(); + ggml_backend_t hrx_backend = ggml_backend_hrx_init(0); + REQUIRE(hrx_backend != nullptr); + + ggml_init_params params = {}; + params.mem_size = 4 * 1024 * 1024; + params.no_alloc = true; + ggml_context * cpu_ctx = ggml_init(params); + ggml_context * hrx_ctx = ggml_init(params); + REQUIRE(cpu_ctx != nullptr); + REQUIRE(hrx_ctx != nullptr); + + constexpr int64_t vocabulary_count = 64; + constexpr int64_t hidden_size = kQwenHiddenSize; + constexpr int64_t token_count = 7; + ggml_tensor * cpu_weight = ggml_new_tensor_2d(cpu_ctx, GGML_TYPE_Q4_K, hidden_size, vocabulary_count); + ggml_tensor * cpu_ids = ggml_new_tensor_1d(cpu_ctx, GGML_TYPE_I32, token_count); + ggml_tensor * cpu_output = ggml_get_rows(cpu_ctx, cpu_weight, cpu_ids); + ggml_tensor * hrx_weight = ggml_new_tensor_2d(hrx_ctx, GGML_TYPE_Q4_K, hidden_size, vocabulary_count); + ggml_tensor * hrx_ids = ggml_new_tensor_1d(hrx_ctx, GGML_TYPE_I32, token_count); + ggml_tensor * hrx_output = ggml_get_rows(hrx_ctx, hrx_weight, hrx_ids); + REQUIRE(cpu_output != nullptr); + REQUIRE(hrx_output != nullptr); + + ggml_cgraph * cpu_graph = ggml_new_graph(cpu_ctx); + ggml_cgraph * hrx_graph = ggml_new_graph(hrx_ctx); + REQUIRE(cpu_graph != nullptr); + REQUIRE(hrx_graph != nullptr); + ggml_build_forward_expand(cpu_graph, cpu_output); + ggml_build_forward_expand(hrx_graph, hrx_output); + + require_kernel_subsequence(scheduled_kernel_sequence(hrx_graph), { "qwen3_moe:qwen_token_embedding_q4k" }); + + ggml_backend_buffer_t cpu_buffer = ggml_backend_alloc_ctx_tensors(cpu_ctx, cpu_backend); + ggml_backend_buffer_t hrx_buffer = ggml_backend_alloc_ctx_tensors(hrx_ctx, hrx_backend); + REQUIRE(cpu_buffer != nullptr); + REQUIRE(hrx_buffer != nullptr); + + const std::vector weight = make_quantized_rows(GGML_TYPE_Q4_K, hidden_size, vocabulary_count, 5); + const std::vector ids = { 3, 17, 29, 41, 53, 7, 19 }; + set_tensor_pair_bytes(cpu_backend, cpu_weight, hrx_backend, hrx_weight, weight.data(), weight.size()); + set_tensor_pair_bytes(cpu_backend, cpu_ids, hrx_backend, hrx_ids, ids.data(), ids.size() * sizeof(int32_t)); + + REQUIRE(ggml_backend_graph_compute(cpu_backend, cpu_graph) == GGML_STATUS_SUCCESS); + REQUIRE(ggml_backend_graph_compute(hrx_backend, hrx_graph) == GGML_STATUS_SUCCESS); + ggml_backend_synchronize(cpu_backend); + ggml_backend_synchronize(hrx_backend); + require_close(get_f32_tensor(hrx_backend, hrx_output), get_f32_tensor(cpu_backend, cpu_output), 5.0e-4f); + + ggml_backend_buffer_free(cpu_buffer); + ggml_backend_buffer_free(hrx_buffer); + ggml_free(cpu_ctx); + ggml_free(hrx_ctx); + ggml_backend_free(cpu_backend); + ggml_backend_free(hrx_backend); +} + +static void run_dense_matmul_cpu_reference_case(ggml_type weight_type, + const char * expected_kernel, + int64_t token_count, + int64_t output_size) { + ggml_backend_t cpu_backend = init_cpu_backend(); + ggml_backend_t hrx_backend = ggml_backend_hrx_init(0); + REQUIRE(hrx_backend != nullptr); + + ggml_init_params params = {}; + params.mem_size = static_cast(32 * 1024 * 1024); + params.no_alloc = true; + ggml_context * cpu_ctx = ggml_init(params); + ggml_context * hrx_ctx = ggml_init(params); + REQUIRE(cpu_ctx != nullptr); + REQUIRE(hrx_ctx != nullptr); + + constexpr int64_t input_size = kQwenHiddenSize; + ggml_tensor * cpu_weight = ggml_new_tensor_2d(cpu_ctx, weight_type, input_size, output_size); + ggml_tensor * cpu_input = ggml_new_tensor_2d(cpu_ctx, GGML_TYPE_F32, input_size, token_count); + ggml_tensor * cpu_output = ggml_mul_mat(cpu_ctx, cpu_weight, cpu_input); + ggml_tensor * hrx_weight = ggml_new_tensor_2d(hrx_ctx, weight_type, input_size, output_size); + ggml_tensor * hrx_input = ggml_new_tensor_2d(hrx_ctx, GGML_TYPE_F32, input_size, token_count); + ggml_tensor * hrx_output = ggml_mul_mat(hrx_ctx, hrx_weight, hrx_input); + REQUIRE(cpu_output != nullptr); + REQUIRE(hrx_output != nullptr); + + ggml_cgraph * cpu_graph = ggml_new_graph(cpu_ctx); + ggml_cgraph * hrx_graph = ggml_new_graph(hrx_ctx); + REQUIRE(cpu_graph != nullptr); + REQUIRE(hrx_graph != nullptr); + ggml_build_forward_expand(cpu_graph, cpu_output); + ggml_build_forward_expand(hrx_graph, hrx_output); + + require_kernel_subsequence(scheduled_kernel_sequence(hrx_graph), { expected_kernel }); + + ggml_backend_buffer_t cpu_buffer = ggml_backend_alloc_ctx_tensors(cpu_ctx, cpu_backend); + ggml_backend_buffer_t hrx_buffer = ggml_backend_alloc_ctx_tensors(hrx_ctx, hrx_backend); + REQUIRE(cpu_buffer != nullptr); + REQUIRE(hrx_buffer != nullptr); + + const std::vector weight = make_quantized_rows(weight_type, input_size, output_size, 6); + const std::vector input = make_pattern_f32(input_size * token_count, 7, 0.01f); + set_tensor_pair_bytes(cpu_backend, cpu_weight, hrx_backend, hrx_weight, weight.data(), weight.size()); + set_tensor_pair_bytes(cpu_backend, cpu_input, hrx_backend, hrx_input, input.data(), input.size() * sizeof(float)); + + REQUIRE(ggml_backend_graph_compute(cpu_backend, cpu_graph) == GGML_STATUS_SUCCESS); + REQUIRE(ggml_backend_graph_compute(hrx_backend, hrx_graph) == GGML_STATUS_SUCCESS); + ggml_backend_synchronize(cpu_backend); + ggml_backend_synchronize(hrx_backend); + require_close(get_f32_tensor(hrx_backend, hrx_output), get_f32_tensor(cpu_backend, cpu_output), 3.0e-1f, 2.0e-2f); + + ggml_backend_buffer_free(cpu_buffer); + ggml_backend_buffer_free(hrx_buffer); + ggml_free(cpu_ctx); + ggml_free(hrx_ctx); + ggml_backend_free(cpu_backend); + ggml_backend_free(hrx_backend); +} + +static void run_endpoint_rmsnorm_q6k_q8_cpu_reference_case() { + ggml_backend_t cpu_backend = init_cpu_backend(); + ggml_backend_t hrx_backend = ggml_backend_hrx_init(0); + REQUIRE(hrx_backend != nullptr); + + ggml_init_params params = {}; + params.mem_size = 32 * 1024 * 1024; + params.no_alloc = true; + ggml_context * cpu_ctx = ggml_init(params); + ggml_context * hrx_ctx = ggml_init(params); + REQUIRE(cpu_ctx != nullptr); + REQUIRE(hrx_ctx != nullptr); + + constexpr int64_t token_count = 1; + ggml_tensor * cpu_input = ggml_new_tensor_2d(cpu_ctx, GGML_TYPE_F32, kQwenHiddenSize, token_count); + ggml_tensor * cpu_norm_w = ggml_new_tensor_1d(cpu_ctx, GGML_TYPE_F32, kQwenHiddenSize); + ggml_tensor * cpu_weight = ggml_new_tensor_2d(cpu_ctx, GGML_TYPE_Q6_K, kQwenHiddenSize, kQwenVocabularyCount); + ggml_tensor * cpu_norm = build_rmsnorm_mul_graph(cpu_ctx, cpu_input, cpu_norm_w); + ggml_tensor * cpu_output = ggml_mul_mat(cpu_ctx, cpu_weight, cpu_norm); + + ggml_tensor * hrx_input = ggml_new_tensor_2d(hrx_ctx, GGML_TYPE_F32, kQwenHiddenSize, token_count); + ggml_tensor * hrx_norm_w = ggml_new_tensor_1d(hrx_ctx, GGML_TYPE_F32, kQwenHiddenSize); + ggml_tensor * hrx_weight = ggml_new_tensor_2d(hrx_ctx, GGML_TYPE_Q6_K, kQwenHiddenSize, kQwenVocabularyCount); + ggml_tensor * hrx_norm = build_rmsnorm_mul_graph(hrx_ctx, hrx_input, hrx_norm_w); + ggml_tensor * hrx_output = ggml_mul_mat(hrx_ctx, hrx_weight, hrx_norm); + REQUIRE(cpu_output != nullptr); + REQUIRE(hrx_output != nullptr); + + ggml_cgraph * cpu_graph = ggml_new_graph(cpu_ctx); + ggml_cgraph * hrx_graph = ggml_new_graph(hrx_ctx); + REQUIRE(cpu_graph != nullptr); + REQUIRE(hrx_graph != nullptr); + ggml_build_forward_expand(cpu_graph, cpu_output); + ggml_build_forward_expand(hrx_graph, hrx_output); + + require_kernel_subsequence( + scheduled_kernel_sequence(hrx_graph), + { "qwen3_moe:qwen3_moe_rmsnorm_f32_quantize_q8_1_x4", "qwen3_moe:ggml_linear_q6k_q8_1_x4" }); + + ggml_backend_buffer_t cpu_buffer = ggml_backend_alloc_ctx_tensors(cpu_ctx, cpu_backend); + ggml_backend_buffer_t hrx_buffer = ggml_backend_alloc_ctx_tensors(hrx_ctx, hrx_backend); + REQUIRE(cpu_buffer != nullptr); + REQUIRE(hrx_buffer != nullptr); + + const std::vector input = make_pattern_f32(kQwenHiddenSize * token_count, 8, 0.01f); + const std::vector norm_w = make_weight(kQwenHiddenSize); + const std::vector weight = make_quantized_rows(GGML_TYPE_Q6_K, kQwenHiddenSize, kQwenVocabularyCount, 9); + set_tensor_pair_bytes(cpu_backend, cpu_input, hrx_backend, hrx_input, input.data(), input.size() * sizeof(float)); + set_tensor_pair_bytes(cpu_backend, cpu_norm_w, hrx_backend, hrx_norm_w, norm_w.data(), + norm_w.size() * sizeof(float)); + set_tensor_pair_bytes(cpu_backend, cpu_weight, hrx_backend, hrx_weight, weight.data(), weight.size()); + + REQUIRE(ggml_backend_graph_compute(cpu_backend, cpu_graph) == GGML_STATUS_SUCCESS); + REQUIRE(ggml_backend_graph_compute(hrx_backend, hrx_graph) == GGML_STATUS_SUCCESS); + ggml_backend_synchronize(cpu_backend); + ggml_backend_synchronize(hrx_backend); + require_close(get_f32_tensor(hrx_backend, hrx_output), get_f32_tensor(cpu_backend, cpu_output), 6.0e-1f, 3.0e-2f); + + ggml_backend_buffer_free(cpu_buffer); + ggml_backend_buffer_free(hrx_buffer); + ggml_free(cpu_ctx); + ggml_free(hrx_ctx); + ggml_backend_free(cpu_backend); + ggml_backend_free(hrx_backend); +} + +static void run_attention_postprocess_cpu_reference_case() { + ggml_backend_t cpu_backend = init_cpu_backend(); + ggml_backend_t hrx_backend = ggml_backend_hrx_init(0); + REQUIRE(hrx_backend != nullptr); + + ggml_init_params params = {}; + params.mem_size = 32 * 1024 * 1024; + params.no_alloc = true; + ggml_context * cpu_ctx = ggml_init(params); + ggml_context * hrx_ctx = ggml_init(params); + REQUIRE(cpu_ctx != nullptr); + REQUIRE(hrx_ctx != nullptr); + + constexpr int64_t token_count = 2; + constexpr int64_t query_head_count = 4; + constexpr int64_t key_value_head_count = 2; + constexpr int64_t cache_row_count = 8; + const int64_t query_size = query_head_count * kQwenFlashHeadSize; + const int64_t key_value_size = key_value_head_count * kQwenFlashHeadSize; + AttentionPostprocessGraph cpu = build_attention_postprocess_graph(cpu_ctx, token_count, query_head_count, + key_value_head_count, cache_row_count); + AttentionPostprocessGraph hrx = build_attention_postprocess_graph(hrx_ctx, token_count, query_head_count, + key_value_head_count, cache_row_count); + + ggml_cgraph * cpu_graph = ggml_new_graph(cpu_ctx); + ggml_cgraph * hrx_graph = ggml_new_graph(hrx_ctx); + REQUIRE(cpu_graph != nullptr); + REQUIRE(hrx_graph != nullptr); + ggml_build_forward_expand(cpu_graph, cpu.query_output); + ggml_build_forward_expand(cpu_graph, cpu.key_output); + ggml_build_forward_expand(cpu_graph, cpu.value_output); + ggml_build_forward_expand(hrx_graph, hrx.query_output); + ggml_build_forward_expand(hrx_graph, hrx.key_output); + ggml_build_forward_expand(hrx_graph, hrx.value_output); + + require_kernel_subsequence(scheduled_kernel_sequence(hrx_graph), + { "qwen3_moe:qwen3_moe_attention_postprocess_f32_f16" }); + + ggml_backend_buffer_t cpu_buffer = ggml_backend_alloc_ctx_tensors(cpu_ctx, cpu_backend); + ggml_backend_buffer_t hrx_buffer = ggml_backend_alloc_ctx_tensors(hrx_ctx, hrx_backend); + REQUIRE(cpu_buffer != nullptr); + REQUIRE(hrx_buffer != nullptr); + + const std::vector input = make_pattern_f32(kQwenHiddenSize * token_count, 10, 0.01f); + const std::vector query_w = make_quantized_rows(GGML_TYPE_Q4_K, kQwenHiddenSize, query_size, 11); + const std::vector key_w = make_quantized_rows(GGML_TYPE_Q4_K, kQwenHiddenSize, key_value_size, 12); + const std::vector value_w = make_quantized_rows(GGML_TYPE_Q6_K, kQwenHiddenSize, key_value_size, 13); + const std::vector query_nw = make_weight(kQwenFlashHeadSize); + const std::vector key_nw = make_pattern_f32(kQwenFlashHeadSize, 14, 0.02f); + const std::vector positions = make_i32_mod_data(token_count, 1024); + std::vector inv_freq(static_cast(kQwenFlashHeadSize / 2)); + for (size_t i = 0; i < inv_freq.size(); ++i) { + inv_freq[i] = 1.0f / std::pow(10000.0f, static_cast(2 * i) / static_cast(kQwenFlashHeadSize)); + } + const std::vector key_cache(static_cast(key_value_size * cache_row_count), + ggml_fp32_to_fp16(0.0f)); + const std::vector value_cache(static_cast(key_value_size * cache_row_count), + ggml_fp32_to_fp16(0.0f)); + const std::vector cache_indices = make_i64_mod_data(token_count, cache_row_count); + + set_tensor_pair_bytes(cpu_backend, cpu.input, hrx_backend, hrx.input, input.data(), input.size() * sizeof(float)); + set_tensor_pair_bytes(cpu_backend, cpu.query_weight, hrx_backend, hrx.query_weight, query_w.data(), query_w.size()); + set_tensor_pair_bytes(cpu_backend, cpu.key_weight, hrx_backend, hrx.key_weight, key_w.data(), key_w.size()); + set_tensor_pair_bytes(cpu_backend, cpu.value_weight, hrx_backend, hrx.value_weight, value_w.data(), value_w.size()); + set_tensor_pair_bytes(cpu_backend, cpu.query_norm_weight, hrx_backend, hrx.query_norm_weight, query_nw.data(), + query_nw.size() * sizeof(float)); + set_tensor_pair_bytes(cpu_backend, cpu.key_norm_weight, hrx_backend, hrx.key_norm_weight, key_nw.data(), + key_nw.size() * sizeof(float)); + set_tensor_pair_bytes(cpu_backend, cpu.positions, hrx_backend, hrx.positions, positions.data(), + positions.size() * sizeof(int32_t)); + set_tensor_pair_bytes(cpu_backend, cpu.inverse_frequencies, hrx_backend, hrx.inverse_frequencies, inv_freq.data(), + inv_freq.size() * sizeof(float)); + set_tensor_pair_bytes(cpu_backend, cpu.key_cache, hrx_backend, hrx.key_cache, key_cache.data(), + key_cache.size() * sizeof(ggml_fp16_t)); + set_tensor_pair_bytes(cpu_backend, cpu.value_cache, hrx_backend, hrx.value_cache, value_cache.data(), + value_cache.size() * sizeof(ggml_fp16_t)); + set_tensor_pair_bytes(cpu_backend, cpu.key_cache_indices, hrx_backend, hrx.key_cache_indices, cache_indices.data(), + cache_indices.size() * sizeof(int64_t)); + set_tensor_pair_bytes(cpu_backend, cpu.value_cache_indices, hrx_backend, hrx.value_cache_indices, + cache_indices.data(), cache_indices.size() * sizeof(int64_t)); + + REQUIRE(ggml_backend_graph_compute(cpu_backend, cpu_graph) == GGML_STATUS_SUCCESS); + REQUIRE(ggml_backend_graph_compute(hrx_backend, hrx_graph) == GGML_STATUS_SUCCESS); + ggml_backend_synchronize(cpu_backend); + ggml_backend_synchronize(hrx_backend); + require_close(get_f32_tensor(hrx_backend, hrx.query_output), get_f32_tensor(cpu_backend, cpu.query_output), 2.0f, + 5.0e-2f); + require_close(get_f32_tensor(hrx_backend, hrx.key_output), get_f32_tensor(cpu_backend, cpu.key_output), 2.0f, + 5.0e-2f); + require_close(get_f32_tensor(hrx_backend, hrx.value_output), get_f32_tensor(cpu_backend, cpu.value_output), 2.0f, + 5.0e-2f); + + ggml_backend_buffer_free(cpu_buffer); + ggml_backend_buffer_free(hrx_buffer); + ggml_free(cpu_ctx); + ggml_free(hrx_ctx); + ggml_backend_free(cpu_backend); + ggml_backend_free(hrx_backend); +} + +static void run_routed_moe_cpu_reference_case(ggml_type down_weight_type, bool include_next_rmsnorm) { + ggml_backend_t cpu_backend = init_cpu_backend(); + ggml_backend_t hrx_backend = ggml_backend_hrx_init(0); + REQUIRE(hrx_backend != nullptr); + + ggml_init_params params = {}; + params.mem_size = 64 * 1024 * 1024; + params.no_alloc = true; + ggml_context * cpu_ctx = ggml_init(params); + ggml_context * hrx_ctx = ggml_init(params); + REQUIRE(cpu_ctx != nullptr); + REQUIRE(hrx_ctx != nullptr); + + RoutedMoeGraph cpu = build_routed_moe_graph(cpu_ctx, down_weight_type, include_next_rmsnorm); + RoutedMoeGraph hrx = build_routed_moe_graph(hrx_ctx, down_weight_type, include_next_rmsnorm); + + ggml_cgraph * cpu_graph = ggml_new_graph(cpu_ctx); + ggml_cgraph * hrx_graph = ggml_new_graph(hrx_ctx); + REQUIRE(cpu_graph != nullptr); + REQUIRE(hrx_graph != nullptr); + ggml_build_forward_expand(cpu_graph, cpu.output); + ggml_build_forward_expand(hrx_graph, hrx.output); + + std::vector expected = { + "qwen3_moe:qwen3_moe_router_top8_f32", + "qwen3_moe:qwen3_moe_build_expert_table", + "qwen3_moe:qwen3_moe_build_expert_partition_table", + "qwen3_moe:qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma", + down_weight_type == GGML_TYPE_Q4_K ? "qwen3_moe:qwen3_moe_routed_down_q4k_f16_wmma_grouped" : + "qwen3_moe:qwen3_moe_routed_down_q6k_f16_wmma_grouped", + include_next_rmsnorm ? "qwen3_moe:qwen3_moe_routed_down_weighted_reduce_next_rmsnorm_f32" : + "qwen3_moe:qwen3_moe_routed_down_weighted_reduce_f16_f32", + }; + require_kernel_subsequence(scheduled_kernel_sequence(hrx_graph), expected); + + ggml_backend_buffer_t cpu_buffer = ggml_backend_alloc_ctx_tensors(cpu_ctx, cpu_backend); + ggml_backend_buffer_t hrx_buffer = ggml_backend_alloc_ctx_tensors(hrx_ctx, hrx_backend); + REQUIRE(cpu_buffer != nullptr); + REQUIRE(hrx_buffer != nullptr); + + const std::vector logits = make_router_logits(1); + const std::vector input = make_pattern_f32(kQwenHiddenSize, 15, 0.01f); + const std::vector hidden = make_pattern_f32(kQwenHiddenSize, 16, 0.02f); + set_tensor_pair_bytes(cpu_backend, cpu.logits, hrx_backend, hrx.logits, logits.data(), + logits.size() * sizeof(float)); + set_tensor_pair_bytes(cpu_backend, cpu.input, hrx_backend, hrx.input, input.data(), input.size() * sizeof(float)); + set_tensor_pair_bytes(cpu_backend, cpu.hidden_state, hrx_backend, hrx.hidden_state, hidden.data(), + hidden.size() * sizeof(float)); + if (include_next_rmsnorm) { + const std::vector norm = make_weight(kQwenHiddenSize); + set_tensor_pair_bytes(cpu_backend, cpu.norm_weight, hrx_backend, hrx.norm_weight, norm.data(), + norm.size() * sizeof(float)); + } + + { + const std::vector gate = + make_quantized_rows(GGML_TYPE_Q4_K, kQwenHiddenSize, kQwenMoeIntermediate * kQwenRouterExpertCount, 17); + set_tensor_pair_bytes(cpu_backend, cpu.gate_weight, hrx_backend, hrx.gate_weight, gate.data(), gate.size()); + } + { + const std::vector up = + make_quantized_rows(GGML_TYPE_Q4_K, kQwenHiddenSize, kQwenMoeIntermediate * kQwenRouterExpertCount, 18); + set_tensor_pair_bytes(cpu_backend, cpu.up_weight, hrx_backend, hrx.up_weight, up.data(), up.size()); + } + { + const std::vector down = + make_quantized_rows(down_weight_type, kQwenMoeIntermediate, kQwenHiddenSize * kQwenRouterExpertCount, 19); + set_tensor_pair_bytes(cpu_backend, cpu.down_weight, hrx_backend, hrx.down_weight, down.data(), down.size()); + } + + REQUIRE(ggml_backend_graph_compute(cpu_backend, cpu_graph) == GGML_STATUS_SUCCESS); + REQUIRE(ggml_backend_graph_compute(hrx_backend, hrx_graph) == GGML_STATUS_SUCCESS); + ggml_backend_synchronize(cpu_backend); + ggml_backend_synchronize(hrx_backend); + require_close(get_f32_tensor(hrx_backend, hrx.output), get_f32_tensor(cpu_backend, cpu.output), 2.5f, 1.0e-1f); + + ggml_backend_buffer_free(cpu_buffer); + ggml_backend_buffer_free(hrx_buffer); + ggml_free(cpu_ctx); + ggml_free(hrx_ctx); + ggml_backend_free(cpu_backend); + ggml_backend_free(hrx_backend); +} + +static void run_decode_routed_moe_scheduling_case(ggml_type down_weight_type, bool alias_gate_input = false) { + ggml_init_params params = {}; + params.mem_size = 128 * 1024 * 1024; + params.no_alloc = true; + ggml_context * ctx = ggml_init(params); + REQUIRE(ctx != nullptr); + + ggml_tensor * hidden_state = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, kQwenHiddenSize, 1); + ggml_tensor * attention_weight = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, kQwenHiddenSize); + REQUIRE(hidden_state != nullptr); + REQUIRE(attention_weight != nullptr); + ggml_tensor * attention_rms = ggml_rms_norm(ctx, hidden_state, kQwenRmsNormEps); + REQUIRE(attention_rms != nullptr); + ggml_tensor * attention_prepared = ggml_mul(ctx, attention_rms, attention_weight); + REQUIRE(attention_prepared != nullptr); + ggml_tensor * moe_input = attention_prepared; + if (alias_gate_input) { + moe_input = ggml_reshape_2d(ctx, attention_prepared, kQwenHiddenSize, 1); + REQUIRE(moe_input != nullptr); + } + + ggml_tensor * router_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, kQwenHiddenSize, kQwenRouterExpertCount); + REQUIRE(router_weight != nullptr); + ggml_tensor * logits = ggml_mul_mat(ctx, router_weight, attention_prepared); + REQUIRE(logits != nullptr); + ggml_tensor * route_ids = nullptr; + ggml_tensor * route_weights = build_qwen_router_top8_graph(ctx, logits, &route_ids); + REQUIRE(route_ids != nullptr); + REQUIRE(route_weights != nullptr); + + ggml_tensor * gate_weight = + ggml_new_tensor_3d(ctx, GGML_TYPE_Q4_K, kQwenHiddenSize, kQwenMoeIntermediate, kQwenRouterExpertCount); + ggml_tensor * up_weight = + ggml_new_tensor_3d(ctx, GGML_TYPE_Q4_K, kQwenHiddenSize, kQwenMoeIntermediate, kQwenRouterExpertCount); + ggml_tensor * down_weight = + ggml_new_tensor_3d(ctx, down_weight_type, kQwenMoeIntermediate, kQwenHiddenSize, kQwenRouterExpertCount); + REQUIRE(gate_weight != nullptr); + REQUIRE(up_weight != nullptr); + REQUIRE(down_weight != nullptr); + + ggml_tensor * gate = ggml_mul_mat_id(ctx, gate_weight, moe_input, route_ids); + ggml_tensor * up = ggml_mul_mat_id(ctx, up_weight, moe_input, route_ids); + REQUIRE(gate != nullptr); + REQUIRE(up != nullptr); + ggml_tensor * glu = ggml_glu_split(ctx, gate, up, GGML_GLU_OP_SWIGLU); + REQUIRE(glu != nullptr); + ggml_tensor * down = ggml_mul_mat_id(ctx, down_weight, glu, route_ids); + REQUIRE(down != nullptr); + ggml_tensor * weighted = ggml_mul(ctx, down, route_weights); + REQUIRE(weighted != nullptr); + + std::vector route_views; + route_views.reserve(kQwenRouterRouteCount); + for (int64_t route = 0; route < kQwenRouterRouteCount; ++route) { + ggml_tensor * view = ggml_view_2d(ctx, weighted, kQwenHiddenSize, 1, weighted->nb[2], + static_cast(route) * weighted->nb[1]); + REQUIRE(view != nullptr); + route_views.push_back(view); + } + + ggml_tensor * reduced = route_views.front(); + for (size_t i = 1; i < route_views.size(); ++i) { + reduced = ggml_add(ctx, reduced, route_views[i]); + REQUIRE(reduced != nullptr); + } + ggml_tensor * residual = ggml_add(ctx, hidden_state, reduced); + REQUIRE(residual != nullptr); + ggml_tensor * next_norm_weight = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, kQwenHiddenSize); + REQUIRE(next_norm_weight != nullptr); + ggml_tensor * next_rms = ggml_rms_norm(ctx, residual, kQwenRmsNormEps); + REQUIRE(next_rms != nullptr); + ggml_tensor * output = ggml_mul(ctx, next_rms, next_norm_weight); + REQUIRE(output != nullptr); + + ggml_cgraph * graph = ggml_new_graph(ctx); + REQUIRE(graph != nullptr); + ggml_build_forward_expand(graph, output); + + std::vector expected = { + "qwen3_moe:qwen3_moe_rmsnorm_f32_quantize_q8_1_x4", + "qwen3_moe:qwen3_moe_router_projection_top8_fused_decode_f32", + down_weight_type == GGML_TYPE_Q4_K ? "qwen3_moe:qwen3_moe_routed_gate_up_swiglu_q4k_q8_1_x4_next_q8" : + "qwen3_moe:qwen3_moe_routed_gate_up_swiglu_q4k_q8", + down_weight_type == GGML_TYPE_Q4_K ? "qwen3_moe:qwen3_moe_routed_down_q4k_q8_1_x4_next_q8" : + "qwen3_moe:qwen3_moe_routed_down_q6k_f32_wave64_next_q8", + }; + require_kernel_subsequence(scheduled_kernel_sequence(graph), expected); + ggml_free(ctx); +} + +static void run_decode_attention_qkv_scheduling_case() { + ggml_init_params params = {}; + params.mem_size = 128 * 1024 * 1024; + params.no_alloc = true; + ggml_context * ctx = ggml_init(params); + REQUIRE(ctx != nullptr); + + AttentionPostprocessGraph graph = build_decode_attention_qkv_graph(ctx); + ggml_cgraph * cgraph = ggml_new_graph(ctx); + REQUIRE(cgraph != nullptr); + ggml_build_forward_expand(cgraph, graph.query_output); + ggml_build_forward_expand(cgraph, graph.key_output); + ggml_build_forward_expand(cgraph, graph.value_output); + + require_kernel_subsequence(scheduled_kernel_sequence(cgraph), + { "qwen3_moe:qwen3_moe_rmsnorm_f32_quantize_q8_1_x4", + "qwen3_moe:qwen3_moe_attention_qkv_postprocess_fused_decode" }); + ggml_free(ctx); +} + +int main() { + run_rmsnorm_support_checks(); + run_alternate_value_alias_lookup_checks(); + + if (ggml_backend_hrx_get_device_count() == 0) { + std::fprintf(stderr, "test skipped: no HRX devices available\n"); + return 0; + } + + run_add_f32_cpu_reference_case(); + run_gather_add_f32_cpu_reference_case(); + run_token_embedding_q4k_cpu_reference_case(); + run_dense_matmul_cpu_reference_case(GGML_TYPE_Q4_K, "qwen3_moe:qwen3_moe_dense_linear_q4k_f16_wmma", 2, 128); + run_dense_matmul_cpu_reference_case(GGML_TYPE_Q6_K, "qwen3_moe:qwen3_moe_dense_linear_q6k_f16_wmma", 2, 128); + run_endpoint_rmsnorm_q6k_q8_cpu_reference_case(); + run_attention_postprocess_cpu_reference_case(); + run_routed_moe_cpu_reference_case(GGML_TYPE_Q4_K, true); + run_routed_moe_cpu_reference_case(GGML_TYPE_Q6_K, false); + run_decode_attention_qkv_scheduling_case(); + run_decode_routed_moe_scheduling_case(GGML_TYPE_Q4_K); + run_decode_routed_moe_scheduling_case(GGML_TYPE_Q6_K); + run_decode_routed_moe_scheduling_case(GGML_TYPE_Q6_K, true); + run_rmsnorm_mul_case(256, 1); + run_rmsnorm_mul_case(256, 4); + run_rmsnorm_mul_case(2048, 1); + run_router_projection_case(4); + run_router_top8_case(4); + run_qwen_flash_attention_case(); + run_qwen_decode_split_flash_attention_scheduling_case(1, 512); + run_qwen_decode_split_flash_attention_scheduling_case(4, 513); + run_qwen_decode_attention_output_next_q8_scheduling_case(false); + run_qwen_decode_attention_output_next_q8_scheduling_case(true); + run_qwen_full_cache_prefill_flash_attention_scheduling_case(16); + run_qwen_full_cache_prefill_flash_attention_scheduling_case(512); + return 0; +}