diff --git a/.github/workflows/build-cuda-ubuntu.yml b/.github/workflows/build-cuda-ubuntu.yml index 30029887c..3dc9255be 100644 --- a/.github/workflows/build-cuda-ubuntu.yml +++ b/.github/workflows/build-cuda-ubuntu.yml @@ -145,7 +145,7 @@ jobs: musa: runs-on: ubuntu-22.04 - container: mthreads/musa:rc4.3.0-devel-ubuntu22.04-amd64 + container: registry.mthreads.com/mcconline/inference/pytorch:2.9.1.post1-py3.10-musa5.2.0-mp31-devel-ubuntu22.04-amd64 steps: - name: Clone @@ -156,7 +156,7 @@ jobs: id: depends run: | apt-get update - apt-get install -y build-essential git cmake libssl-dev jq + apt-get install -y build-essential git cmake libssl-dev jq python3-venv - name: ccache uses: ggml-org/ccache-action@v1.2.24 @@ -178,8 +178,8 @@ jobs: run: | cmake -B build -S . \ -DGGML_MUSA=ON \ - -DMUSA_ARCHITECTURES=21 - time cmake --build build --config Release -j $(nproc) + -DMUSA_ARCHITECTURES=31 + cmake --build build --config Release -j $(nproc) - name: ccache-buckets-save if: ${{ github.event_name == 'push' && github.ref == 'refs/heads/master' }} diff --git a/.github/workflows/server-sanitize.yml b/.github/workflows/server-sanitize.yml index 8ff285a42..69777b6a2 100644 --- a/.github/workflows/server-sanitize.yml +++ b/.github/workflows/server-sanitize.yml @@ -45,7 +45,7 @@ concurrency: jobs: server: - runs-on: hf-jobs-cpu-xl + runs-on: hf-jobs-cpu-performance strategy: matrix: @@ -116,7 +116,7 @@ jobs: run: | source .venv/bin/activate cd tools/server/tests - PYTEST_WORKERS=1 ./tests.sh + PYTEST_WORKERS=4 ./tests.sh - name: Slow tests id: server_integration_tests_slow @@ -124,4 +124,4 @@ jobs: run: | source .venv/bin/activate cd tools/server/tests - PYTEST_WORKERS=1 SLOW_TESTS=1 ./tests.sh + PYTEST_WORKERS=4 SLOW_TESTS=1 ./tests.sh diff --git a/.pi/gg/SYSTEM.md b/.pi/gg/SYSTEM.md index bd308ea96..8b8b8a158 100644 --- a/.pi/gg/SYSTEM.md +++ b/.pi/gg/SYSTEM.md @@ -7,6 +7,7 @@ General: - Don't try to build or run the code unless you are explicitly asked to do so - Use the `gh` CLI tool when querying PRs, issues, or other GitHub resources - When [MODEL] is needed, first try to get it from the `PI_MODEL_NAME` env var before asking the user +- Never read the `AGENTS.md` file Coding: - When in doubt, always refer to the CONTRIBUTING.md file of the project diff --git a/ci/README-MUSA.md b/ci/README-MUSA.md index c5e24c5d9..2101dfa2c 100644 --- a/ci/README-MUSA.md +++ b/ci/README-MUSA.md @@ -21,7 +21,7 @@ docker run --privileged -it \ -v $HOME/llama.cpp/ci-cache:/ci-cache \ -v $HOME/llama.cpp/ci-results:/ci-results \ -v $PWD:/ws -w /ws \ - mthreads/musa:rc4.3.0-devel-ubuntu22.04-amd64 + registry.mthreads.com/mcconline/inference/pytorch:2.9.1.post1-py3.10-musa5.2.0-mp31-devel-ubuntu22.04-amd64 ``` Inside the container, execute the following commands: diff --git a/ci/run.sh b/ci/run.sh index 0595fac5a..fc661f976 100755 --- a/ci/run.sh +++ b/ci/run.sh @@ -158,8 +158,8 @@ if [ ! -z ${GG_BUILD_WEBGPU} ]; then fi if [ ! -z ${GG_BUILD_MUSA} ]; then - # Use qy1 by default (MTT S80) - MUSA_ARCH=${MUSA_ARCH:-21} + # Use ph1 by default (MTT S5000) + MUSA_ARCH=${MUSA_ARCH:-31} CMAKE_EXTRA="${CMAKE_EXTRA} -DGGML_MUSA=ON -DMUSA_ARCHITECTURES=${MUSA_ARCH}" fi diff --git a/common/common.cpp b/common/common.cpp index 750defd3c..e2b9a8765 100644 --- a/common/common.cpp +++ b/common/common.cpp @@ -2220,9 +2220,28 @@ bool common_replay_last_token(struct llama_context * ctx, llama_token last_token return true; } +llama_batch_ext_ptr common_batch_ext_get_one(llama_context * ctx, const llama_tokens & tokens) { + llama_batch_ext_ptr batch(llama_batch_ext_init(ctx)); + + auto mem = llama_get_memory(ctx); + llama_pos pos = mem ? llama_memory_seq_pos_max(mem, 0) + 1 : 0; + + for (size_t i = 0; i < tokens.size(); ++i) { + const int32_t idx = llama_batch_ext_add_token(batch.get(), 0, tokens[i]); + llama_batch_ext_set_pos(batch.get(), idx, &pos); + pos++; + } + + if (!tokens.empty()) { + llama_batch_ext_set_output_logits(batch.get(), (int32_t) tokens.size() - 1, true); + } + + return batch; +} + bool common_prompt_batch_decode( struct llama_context * ctx, - const std::vector & all_tokens, + const llama_tokens & all_tokens, int n_new, int & n_past, int n_batch, @@ -2243,7 +2262,9 @@ bool common_prompt_batch_decode( // Memory implementations in recurrent/hybrid models don't support removing tokens from their // memory, so we can't just remove the last token from the memory and replay the last token which // is the reason for this logic. - if (llama_decode(ctx, llama_batch_get_one(const_cast(all_tokens.data() + offset), n_tokens_before_last))) { + llama_tokens prefix_tokens(all_tokens.begin() + offset, all_tokens.begin() + offset + n_tokens_before_last); + llama_batch_ext_ptr batch_prefix = common_batch_ext_get_one(ctx, prefix_tokens); + if (llama_process(ctx, LLAMA_PROCESS_TYPE_DECODE, batch_prefix.get())) { COM_ERR("%s", "failed to eval\n"); return false; } @@ -2253,17 +2274,19 @@ bool common_prompt_batch_decode( COM_INF("saved session before last token to %s, n_new = %zu\n", state_path.data(), all_tokens.size()); llama_token last_token = all_tokens.back(); - llama_batch batch = llama_batch_get_one(&last_token, 1); - int32_t pos = n_past; - batch.pos = &pos; + llama_batch_ext_ptr batch_last = common_batch_ext_get_one(ctx, { last_token }); + llama_pos pos = n_past; + llama_batch_ext_set_pos(batch_last.get(), 0, &pos); - if (llama_decode(ctx, batch)) { + if (llama_process(ctx, LLAMA_PROCESS_TYPE_DECODE, batch_last.get())) { COM_ERR("%s", "failed to eval last token\n"); return false; } n_past++; } else { - if (llama_decode(ctx, llama_batch_get_one(const_cast(all_tokens.data() + offset), n_new))) { + llama_tokens new_tokens(all_tokens.begin() + offset, all_tokens.begin() + offset + n_new); + llama_batch_ext_ptr batch = common_batch_ext_get_one(ctx, new_tokens); + if (llama_process(ctx, LLAMA_PROCESS_TYPE_DECODE, batch.get())) { COM_ERR("%s", "failed to eval\n"); return false; } diff --git a/common/common.h b/common/common.h index 730f08d5e..20274f8cf 100644 --- a/common/common.h +++ b/common/common.h @@ -1023,6 +1023,10 @@ void common_batch_add( const std::vector & seq_ids, bool logits); +// create a single-sequence batch from a list of tokens +// last token always have output_logits set to true +llama_batch_ext_ptr common_batch_ext_get_one(struct llama_context * ctx, const llama_tokens & tokens); + // decodes a single batch of tokens for a prompt and manages session tokens // // Note: We save state before the last token so that we can replay it to ensure @@ -1030,7 +1034,7 @@ void common_batch_add( // tokens from memory, so this approach works across all model architectures. bool common_prompt_batch_decode( struct llama_context * ctx, - const std::vector & all_tokens, + const llama_tokens & all_tokens, int n_new, int & n_past, int n_batch, diff --git a/conversion/mimo.py b/conversion/mimo.py index 8a2689b96..982594e85 100644 --- a/conversion/mimo.py +++ b/conversion/mimo.py @@ -17,6 +17,7 @@ from .base import MmprojModel, ModelBase, TextModel, gguf, logger @ModelBase.example("XiaomiMiMo/MiMo-V2.5") class MimoV2Model(TextModel): model_arch = gguf.MODEL_ARCH.MIMO2 + supports_mtp_export = True # MiMo V2-Flash, V2.5 and V2.5-Pro all ship 3 trained MTP layers under model.mtp.layers.{0,1,2}. # The HF config does not expose the count, so it's hardcoded to match the count found in the safetensors. @@ -25,6 +26,8 @@ class MimoV2Model(TextModel): def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) + if self.no_mtp: + self._n_nextn = 0 self.block_count = self.hparams["num_hidden_layers"] + self._n_nextn self.tensor_map = gguf.get_tensor_name_map(self.model_arch, self.block_count) @@ -101,7 +104,7 @@ class MimoV2Model(TextModel): qkv_overrides: dict[str, tuple[Callable, Callable, int]] = {} qc = self.hparams.get("quantization_config") if isinstance(qc, dict) and qc.get("quant_method") == "fp8": - pat = re.compile(r"^model\.layers\.(\d+)\.self_attn\.qkv_proj\.weight_scale_inv$") + pat = re.compile(r"^model\.(mtp\.)?layers\.(\d+)\.self_attn\.qkv_proj\.weight_scale_inv$") for name in list(self.model_tensors.keys()): m = pat.match(name) if not m: @@ -109,10 +112,13 @@ class MimoV2Model(TextModel): weight_name = name.removesuffix("_scale_inv") if weight_name not in self.model_tensors: continue + bid = int(m.group(2)) + if m.group(1) is not None: + bid += self.hparams["num_hidden_layers"] qkv_overrides[weight_name] = ( self.model_tensors[weight_name], self.model_tensors[name], - int(m.group(1)), + bid, ) super().dequant_model() @@ -165,7 +171,8 @@ class MimoV2Model(TextModel): if v_scale is not None: self.gguf_writer.add_attn_value_scale(float(v_scale)) - self.gguf_writer.add_nextn_predict_layers(self._n_nextn) + if self._n_nextn > 0: + self.gguf_writer.add_nextn_predict_layers(self._n_nextn) _MXFP4_EXPERT_RE = re.compile( r"^model\.layers\.(\d+)\.mlp\.experts\.(\d+)\.(gate|up|down)_proj\.weight$" @@ -251,11 +258,32 @@ class MimoV2Model(TextModel): def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None: name, gen = item + is_mtp = name.startswith("model.mtp.layers.") + if is_mtp and cls.no_mtp: + return None + if cls.mtp_only and not is_mtp and name not in ( + "model.embed_tokens.weight", "model.norm.weight", "lm_head.weight", + ): + return None + if "attention_sink" in name and not name.endswith(".weight"): name += ".weight" return super().filter_tensors((name, gen)) + def prepare_metadata(self, vocab_only: bool): + from_dir = self.fname_out.is_dir() + super().prepare_metadata(vocab_only=vocab_only) + + if not self.mtp_only or not from_dir: + return + + output_type: str = self.ftype.name.partition("_")[2] + fname_default: str = gguf.naming_convention( + self.metadata.name, self.metadata.basename, self.metadata.finetune, + self.metadata.version, size_label=None, output_type=output_type, model_type=None) + self.fname_out = self.fname_out.parent / f"mtp-{fname_default}.gguf" + def modify_tensors(self, data_torch, name, bid): # Remap MTP/NextN tensors to additional layer slots so the standard tensor map handles them. # HF: model.mtp.layers.{i}.foo -> model.layers.{n_layer_text + i}.foo diff --git a/docs/backend/SYCL.md b/docs/backend/SYCL.md index 91b209741..f9da1ab3e 100644 --- a/docs/backend/SYCL.md +++ b/docs/backend/SYCL.md @@ -811,6 +811,9 @@ User can use the device management in [docs/multi-gpu.md](https://github.com/ggm | GGML_SYCL_MKL_FA_DIAG | 0 (default) or 1 | Enable output fingerprinting for MKL flash attention. Dumps the first 64 float output values for the first 6 FA calls with n_kv ≥ 1024, labeled with kernel type (MKL/TILE/VEC) for cross-kernel comparison. | | GGML_SYCL_ENABLE_FUSION | 0 or 1 (default) | Enable fused-kernel dispatch in graph compute. Unsupported types and layouts fall back to the standalone op kernels. See `ggml_sycl_can_fuse()`. | | GGML_SYCL_ENABLE_ESIMD | 0 or 1 (default)| Enable ESIMD kernels when available. | +| GGML_SYCL_SPARSE_FA | 0 (default) or 1 | Enable Sparse Flash-attention.| +| GGML_SYCL_SPARSE_FA_DEBUG | 0 (default) or 1 | Enable to debug for Sparse Flash-attention.| +| GGML_SYCL_SPARSE_FA_MARGIN | [0,..] default:256 | Set the margin value for Sparse Flash-attention.| | ZES_ENABLE_SYSMAN | 0 (default) or 1 | Support to get free memory of GPU by sycl::aspect::ext_intel_free_memory.
Recommended to use when --split-mode = layer | | UR_L0_ENABLE_RELAXED_ALLOCATION_LIMITS | 0 (default) or 1 | Allow SYCL/Unified Runtime Level Zero device allocations larger than 4 GiB. llama.cpp's direct Level Zero allocation path requests the relaxed maximum-size limit itself when GGML_SYCL_ENABLE_LEVEL_ZERO=1. | | GGML_SYCL_USM_SYSTEM | 0 (default) or 1 | Enable experimental support for [USM system allocations](https://github.khronos.org/SYCL_Reference/iface/usm_basic_concept.html#system-allocations) for large GPU buffers. This requires enough host memory for model weights and caches, an Intel Xe2+ GPU such as BMG or newer and supported on Linux only, with CONFIG_DRM_XE_GPUSVM enabled. | diff --git a/docs/build.md b/docs/build.md index 70fc17af2..bd666c145 100644 --- a/docs/build.md +++ b/docs/build.md @@ -323,11 +323,11 @@ cmake --build build --config Release By default, all supported compute capabilities are enabled. To customize this behavior, you can specify the `MUSA_ARCHITECTURES` option in the CMake command: ```bash -cmake -B build -DGGML_MUSA=ON -DMUSA_ARCHITECTURES="21" +cmake -B build -DGGML_MUSA=ON -DMUSA_ARCHITECTURES="31" cmake --build build --config Release ``` -This configuration enables only compute capability `2.1` (MTT S80) during compilation, which can help reduce compilation time. +This configuration enables only compute capability `3.1` (MTT S5000) during compilation, which can help reduce compilation time. #### Compilation options diff --git a/ggml/CMakeLists.txt b/ggml/CMakeLists.txt index 1df78fddf..752dabbb6 100644 --- a/ggml/CMakeLists.txt +++ b/ggml/CMakeLists.txt @@ -5,7 +5,7 @@ project("ggml" C CXX ASM) ### GGML Version set(GGML_VERSION_MAJOR 0) set(GGML_VERSION_MINOR 25) -set(GGML_VERSION_PATCH 1) +set(GGML_VERSION_PATCH 3) set(GGML_VERSION_BASE "${GGML_VERSION_MAJOR}.${GGML_VERSION_MINOR}.${GGML_VERSION_PATCH}") list(APPEND CMAKE_MODULE_PATH "${CMAKE_CURRENT_SOURCE_DIR}/cmake/") diff --git a/ggml/src/ggml-cuda/common.cuh b/ggml/src/ggml-cuda/common.cuh index 2e78ae4fa..5363dd001 100644 --- a/ggml/src/ggml-cuda/common.cuh +++ b/ggml/src/ggml-cuda/common.cuh @@ -111,9 +111,9 @@ #define GGML_CUDA_CC_IS_QY2(cc) (cc >= GGML_CUDA_CC_QY2 && cc < GGML_CUDA_CC_PH1) #define GGML_CUDA_CC_IS_PH1(cc) (cc >= GGML_CUDA_CC_PH1) -#if !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) && CUDART_VERSION >= 11070 +#if !defined(GGML_USE_HIP) && (defined(GGML_USE_MUSA) || CUDART_VERSION >= 11070) # define GGML_CUDA_USE_CUB -#endif // !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) && CUDART_VERSION >= 11070 +#endif // !defined(GGML_USE_HIP) && (defined(GGML_USE_MUSA) || CUDART_VERSION >= 11070) // PDL host-side support (cudaLaunchKernelEx) requires CUDART >= 11.8. // However, this has been bugged in CTK < 12.3 for MSVC builds, see @@ -237,7 +237,7 @@ static const char * cu_get_error_str(CUresult err) { #define CU_CHECK(err) CUDA_CHECK_GEN(err, CUDA_SUCCESS, cu_get_error_str) #endif -#if !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) +#if !defined(GGML_USE_HIP) # define CUDA_SET_SHARED_MEMORY_LIMIT(kernel, nbytes) \ do { \ static bool shared_memory_limit_raised[GGML_CUDA_MAX_DEVICES] = { false }; \ @@ -252,7 +252,7 @@ static const char * cu_get_error_str(CUresult err) { do { \ GGML_UNUSED(nbytes); \ } while (0) -#endif // !(defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) +#endif // !defined(GGML_USE_HIP) #if CUDART_VERSION >= 11010 || defined(GGML_USE_MUSA) #define GGML_CUDA_ASSUME(x) __builtin_assume(x) @@ -397,7 +397,7 @@ static constexpr __device__ int ggml_cuda_get_physical_warp_size() { // Maximum number of bytes that can be copied in a single instruction. static constexpr __device__ int ggml_cuda_get_max_cpy_bytes() { -#ifdef GGML_USE_HIP +#if defined(GGML_USE_HIP) || defined(GGML_USE_MUSA) return 16; #else #if __CUDA_ARCH__ >= GGML_CUDA_CC_VOLTA @@ -405,7 +405,7 @@ static constexpr __device__ int ggml_cuda_get_max_cpy_bytes() { #else return 8; #endif // __CUDA_ARCH__ >= GGML_CUDA_CC_VOLTA -#endif // GGML_USE_HIP +#endif // defined(GGML_USE_HIP) || defined(GGML_USE_MUSA) } @@ -424,10 +424,6 @@ static __device__ void no_device_code( __trap(); GGML_UNUSED(no_device_code); // suppress unused function warning - -#if defined(GGML_USE_MUSA) - __builtin_unreachable(); -#endif // defined(GGML_USE_MUSA) } #ifdef __CUDA_ARCH__ @@ -696,16 +692,11 @@ static __device__ __forceinline__ half2 ggml_cuda_hmax2(const half2 a, const hal template static __device__ __forceinline__ half2 warp_reduce_max(half2 x) { -#if !defined(GGML_USE_HIP) && __CUDA_ARCH__ >= GGML_CUDA_CC_PASCAL || defined(GGML_USE_HIP) #pragma unroll for (int offset = width/2; offset > 0; offset >>= 1) { x = ggml_cuda_hmax2(x, __shfl_xor_sync(0xffffffff, x, offset, width)); } return x; -#else - GGML_UNUSED(x); - NO_DEVICE_CODE; -#endif // !defined(GGML_USE_HIP) && __CUDA_ARCH__ >= GGML_CUDA_CC_PASCAL || defined(GGML_USE_HIP) } #if (defined(CUDART_VERSION) && CUDART_VERSION < CUDART_HMASK) || defined(GGML_USE_HIP) || \ diff --git a/ggml/src/ggml-cuda/conv2d-dw.cu b/ggml/src/ggml-cuda/conv2d-dw.cu index 7583233b1..6371d0a87 100644 --- a/ggml/src/ggml-cuda/conv2d-dw.cu +++ b/ggml/src/ggml-cuda/conv2d-dw.cu @@ -1,4 +1,5 @@ #include "conv2d-dw.cuh" +#include "convert.cuh" struct conv_params { int in_w, in_h; @@ -79,7 +80,7 @@ struct cwhn_layout { }; template -__global__ void conv2d_dw_kernel(const T * __restrict__ input, const T * __restrict__ kernel, T * __restrict__ output, +__global__ void conv2d_dw_kernel(const float * __restrict__ input, const T * __restrict__ kernel, float * __restrict__ output, const int in_w, const int in_h, const int out_w, const int out_h, const int kernel_w, const int kernel_h, const int stride_x, const int stride_y, const int padding_x, const int padding_y, const int dilation_x, const int dilation_y, @@ -97,7 +98,7 @@ __global__ void conv2d_dw_kernel(const T * __restrict__ input, const T * __restr int batch_idx, channel_idx, out_y_idx, out_x_idx; Layout::unpack_indices(global_idx, params, batch_idx, channel_idx, out_y_idx, out_x_idx); - T accumulator = 0; + float accumulator = 0.0f; kernel_bounds bounds = calculate_kernel_bounds(out_x_idx, out_y_idx, params); for (int kern_y = bounds.y_min; kern_y < bounds.y_max; ++kern_y) { @@ -106,10 +107,10 @@ __global__ void conv2d_dw_kernel(const T * __restrict__ input, const T * __restr for (int kern_x = bounds.x_min; kern_x < bounds.x_max; ++kern_x) { int in_x_idx = calculate_input_coord(out_x_idx, kern_x, params.stride_x, params.dilation_x, params.padding_x); - const T input_val = input[Layout::input_index(batch_idx, channel_idx, in_y_idx, in_x_idx, params)]; - const T kernel_val = kernel[Layout::kernel_index(channel_idx, kern_y, kern_x, params)]; + const float input_val = input[Layout::input_index(batch_idx, channel_idx, in_y_idx, in_x_idx, params)]; + const T kernel_val = kernel[Layout::kernel_index(channel_idx, kern_y, kern_x, params)]; - accumulator += input_val * kernel_val; + accumulator += input_val * ggml_cuda_cast(kernel_val); } } @@ -120,8 +121,9 @@ void ggml_cuda_op_conv2d_dw(ggml_backend_cuda_context & ctx, ggml_tensor * dst) const ggml_tensor * kernel = dst->src[0]; const ggml_tensor * input = dst->src[1]; - GGML_ASSERT(kernel->type == GGML_TYPE_F32 && input->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32); - const float * w_d = (const float *) kernel->data; + GGML_ASSERT(kernel->type == GGML_TYPE_F16 || kernel->type == GGML_TYPE_F32); + GGML_ASSERT(input->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32); + const void * w_d = kernel->data; const float * x_d = (const float *) input->data; float * y_d = (float *) dst->data; @@ -148,13 +150,25 @@ void ggml_cuda_op_conv2d_dw(ggml_backend_cuda_context & ctx, ggml_tensor * dst) const int blocks = (total + CUDA_CONV2D_DW_BLOCK_SIZE - 1) / CUDA_CONV2D_DW_BLOCK_SIZE; if (ggml_is_contiguous(input)) { - conv2d_dw_kernel<<>>( - x_d, w_d, y_d, in_w, in_h, out_w, out_h, kernel_w, kernel_h, stride_x, stride_y, padding_x, padding_y, - dilation_x, dilation_y, channels, batches); + if (kernel->type == GGML_TYPE_F16) { + conv2d_dw_kernel<<>>( + x_d, (const half *) w_d, y_d, in_w, in_h, out_w, out_h, kernel_w, kernel_h, stride_x, stride_y, + padding_x, padding_y, dilation_x, dilation_y, channels, batches); + } else { + conv2d_dw_kernel<<>>( + x_d, (const float *) w_d, y_d, in_w, in_h, out_w, out_h, kernel_w, kernel_h, stride_x, stride_y, + padding_x, padding_y, dilation_x, dilation_y, channels, batches); + } } else if (ggml_is_contiguous_channels(input)) { - conv2d_dw_kernel<<>>( - x_d, w_d, y_d, in_w, in_h, out_w, out_h, kernel_w, kernel_h, stride_x, stride_y, padding_x, padding_y, - dilation_x, dilation_y, channels, batches); + if (kernel->type == GGML_TYPE_F16) { + conv2d_dw_kernel<<>>( + x_d, (const half *) w_d, y_d, in_w, in_h, out_w, out_h, kernel_w, kernel_h, stride_x, stride_y, + padding_x, padding_y, dilation_x, dilation_y, channels, batches); + } else { + conv2d_dw_kernel<<>>( + x_d, (const float *) w_d, y_d, in_w, in_h, out_w, out_h, kernel_w, kernel_h, stride_x, stride_y, + padding_x, padding_y, dilation_x, dilation_y, channels, batches); + } } else { GGML_ABORT("Unsupported memory layout for conv_2d_dw"); } diff --git a/ggml/src/ggml-cuda/fattn-mma-f16.cuh b/ggml/src/ggml-cuda/fattn-mma-f16.cuh index 449a77c5b..abb99a354 100644 --- a/ggml/src/ggml-cuda/fattn-mma-f16.cuh +++ b/ggml/src/ggml-cuda/fattn-mma-f16.cuh @@ -2091,26 +2091,22 @@ void ggml_cuda_flash_attn_ext_mma_f16_case(ggml_backend_cuda_context & ctx, ggml constexpr bool use_sparse_kernel = false; fattn_kernel = flash_attn_ext_f16; -#if !defined(GGML_USE_MUSA) static bool shared_memory_limit_raised[GGML_CUDA_MAX_DEVICES] = {false}; if (!shared_memory_limit_raised[id]) { CUDA_CHECK(cudaFuncSetAttribute(reinterpret_cast(fattn_kernel), cudaFuncAttributeMaxDynamicSharedMemorySize, nbytes_shared_total)); shared_memory_limit_raised[id] = true; } -#endif // !defined(GGML_USE_MUSA) } } else { constexpr bool use_logit_softcap = true; constexpr bool use_sparse_kernel = false; fattn_kernel = flash_attn_ext_f16; -#if !defined(GGML_USE_MUSA) static bool shared_memory_limit_raised[GGML_CUDA_MAX_DEVICES] = {false}; if (!shared_memory_limit_raised[id]) { CUDA_CHECK(cudaFuncSetAttribute(reinterpret_cast(fattn_kernel), cudaFuncAttributeMaxDynamicSharedMemorySize, nbytes_shared_total)); shared_memory_limit_raised[id] = true; } -#endif // !defined(GGML_USE_MUSA) } launch_fattn diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu index 1783a23c3..faace54e3 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu @@ -311,13 +311,9 @@ static ggml_cuda_device_info ggml_cuda_init() { info.devices[id].smpb = prop.sharedMemPerBlock; info.devices[id].warp_size = prop.warpSize; -#ifndef GGML_USE_MUSA int supports_coop_launch = 0; CUDA_CHECK(cudaDeviceGetAttribute(&supports_coop_launch, cudaDevAttrCooperativeLaunch, physical_id)); info.devices[id].supports_cooperative_launch = !!supports_coop_launch; -#else - info.devices[id].supports_cooperative_launch = false; -#endif // !(GGML_USE_MUSA) #if defined(GGML_USE_HIP) info.devices[id].smpbo = prop.sharedMemPerBlock; @@ -338,8 +334,6 @@ static ggml_cuda_device_info ggml_cuda_init() { device_vmm ? "yes" : "no", prop.warpSize, device_vram_mib); #elif defined(GGML_USE_MUSA) - // FIXME: Ensure compatibility with varying warp sizes across different MUSA archs. - info.devices[id].warp_size = 32; info.devices[id].smpbo = prop.sharedMemPerBlockOptin; info.devices[id].cc = GGML_CUDA_CC_OFFSET_MTHREADS + prop.major * 0x100; info.devices[id].cc += prop.minor * 0x10; @@ -3399,6 +3393,19 @@ static bool ggml_cuda_can_fuse(const struct ggml_cgraph * cgraph, return true; } + if (ops.size() == 2 && ops.begin()[0] == GGML_OP_RMS_NORM && ops.begin()[1] == GGML_OP_SCALE) { + const ggml_tensor * rms_norm = cgraph->nodes[node_idx]; + const ggml_tensor * scale = cgraph->nodes[node_idx+1]; + + GGML_ASSERT(rms_norm->src[0]->type == GGML_TYPE_F32); + GGML_ASSERT(rms_norm->type == GGML_TYPE_F32); + + float bias; + memcpy(&bias, (const float *) scale->op_params + 1, sizeof(float)); + + return bias == 0.0f && scale->type == GGML_TYPE_F32; + } + if (ops.size() == 2 && ops.begin()[0] == GGML_OP_SSM_CONV && ops.begin()[1] == GGML_OP_UNARY && unary_ops.size() == 1 && unary_ops.begin()[0] == GGML_UNARY_OP_SILU) { const ggml_tensor * ssm_conv = cgraph->nodes[node_idx]; @@ -4270,6 +4277,11 @@ static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph return 1; } + if (ggml_cuda_can_fuse(cgraph, i, { GGML_OP_RMS_NORM, GGML_OP_SCALE }, {})) { + ggml_cuda_op_rms_norm_scale_fused(*cuda_ctx, node, cgraph->nodes[i + 1]); + return 1; + } + if (ggml_cuda_can_fuse(cgraph, i, { GGML_OP_SSM_CONV, GGML_OP_ADD, GGML_OP_UNARY }, { GGML_UNARY_OP_SILU })) { ggml_cuda_op_ssm_conv(*cuda_ctx, node, cgraph->nodes[i + 1], cgraph->nodes[i + 2]); return 2; @@ -5634,7 +5646,8 @@ static bool ggml_backend_cuda_device_supports_op(ggml_backend_dev_t dev, const g op->src[1]->type == GGML_TYPE_F32 && op->type == GGML_TYPE_F32 && ggml_is_contiguous(op->src[0]) && ggml_is_contiguous(op->src[1]) && ggml_is_contiguous(op); case GGML_OP_CONV_2D_DW: - return op->src[0]->type == GGML_TYPE_F32; + return (op->src[0]->type == GGML_TYPE_F16 || op->src[0]->type == GGML_TYPE_F32) && + op->src[1]->type == GGML_TYPE_F32 && op->type == GGML_TYPE_F32; case GGML_OP_CONV_TRANSPOSE_2D: case GGML_OP_POOL_1D: case GGML_OP_POOL_2D: @@ -5675,12 +5688,7 @@ static bool ggml_backend_cuda_device_supports_op(ggml_backend_dev_t dev, const g case GGML_OP_RWKV_WKV7: return true; case GGML_OP_GATED_DELTA_NET: - //TODO: enable once MUSA compiler is solved https://github.com/ggml-org/llama.cpp/pull/19504#issuecomment-4018634327 -#ifdef GGML_USE_MUSA - return false; -#else return true; -#endif // GGML_USE_MUSA case GGML_OP_DSV4_HC_COMB: return op->src[0]->type == GGML_TYPE_F32 && op->src[1]->type == GGML_TYPE_F32 && op->src[2]->type == GGML_TYPE_F32 && op->type == GGML_TYPE_F32; diff --git a/ggml/src/ggml-cuda/mmq.cu b/ggml/src/ggml-cuda/mmq.cu index f0a7ea7bf..040fb946e 100644 --- a/ggml/src/ggml-cuda/mmq.cu +++ b/ggml/src/ggml-cuda/mmq.cu @@ -393,5 +393,10 @@ bool ggml_cuda_should_use_mmq(enum ggml_type type, int cc, int64_t ne11, int64_t return n_experts > 0; } + // MUSA: the MMQ kernels compute wrong values on PH1 (MTT S5000). + if (cc == GGML_CUDA_CC_PH1) { + return false; + } + return (!GGML_CUDA_CC_IS_CDNA(cc)) || ne11 < MMQ_DP4A_MAX_BATCH_SIZE; } diff --git a/ggml/src/ggml-cuda/norm.cu b/ggml/src/ggml-cuda/norm.cu index c3758cd50..5543307be 100644 --- a/ggml/src/ggml-cuda/norm.cu +++ b/ggml/src/ggml-cuda/norm.cu @@ -73,7 +73,7 @@ static __global__ void group_norm_f32(const float * x, float * dst, const int gr } } -template +template static __global__ void rms_norm_f32(const float * x, float * dst, const int ncols, @@ -96,7 +96,8 @@ static __global__ void rms_norm_f32(const float * x, const uint3 add_ncols_packed = make_uint3(0, 0, 0), const uint3 add_nrows_packed = make_uint3(0, 0, 0), const uint3 add_nchannels_packed = make_uint3(0, 0, 0), - const uint3 add_nsamples_packed = make_uint3(0, 0, 0)) { + const uint3 add_nsamples_packed = make_uint3(0, 0, 0), + const float scale_out = 1.0f) { ggml_cuda_pdl_lc(); const int nrows = gridDim.x; const int nchannels = gridDim.y; @@ -107,6 +108,7 @@ static __global__ void rms_norm_f32(const float * x, const int tid = threadIdx.x; static_assert(!do_add || do_multiply, "fusing add is not supported without multiplying"); + static_assert(!do_scale || !do_multiply, "fusing scale is not supported with multiplying"); x += sample*stride_sample + channel*stride_channel + row*stride_row; dst += ((sample*nchannels + channel)*nrows + row)*ncols; @@ -148,6 +150,8 @@ static __global__ void rms_norm_f32(const float * x, } else if constexpr (do_multiply) { const int mul_col = fastmodulo(col, mul_ncols_packed); dst[col] = scale * x[col] * mul[mul_col]; + } else if constexpr (do_scale) { + dst[col] = scale_out * (scale * x[col]); } else { dst[col] = scale * x[col]; } @@ -301,25 +305,27 @@ static void group_norm_f32_cuda( } } +template static void rms_norm_f32_cuda( const float * x, float * dst, const int ncols, const int nrows, const int nchannels, const int nsamples, - const int64_t stride_row, const int64_t stride_channel, const int64_t stride_sample, const float eps, cudaStream_t stream) { + const int64_t stride_row, const int64_t stride_channel, const int64_t stride_sample, const float eps, cudaStream_t stream, + const float scale_out = 1.0f) { const dim3 blocks_num(nrows, nchannels, nsamples); if (ncols < 1024) { const dim3 block_dims(256, 1, 1); const ggml_cuda_kernel_launch_params launch_params = {blocks_num, block_dims, block_dims.x > WARP_SIZE ? 32 * sizeof(float): 0, stream}; - ggml_cuda_kernel_launch(rms_norm_f32<256, false>, launch_params, + ggml_cuda_kernel_launch(rms_norm_f32<256, false, false, do_scale>, launch_params, x, dst, ncols, stride_row, stride_channel, stride_sample, eps, // underlying cudaLaunchKernelEx does not support default params nullptr, 0, 0, 0, make_uint3(0, 0, 0), make_uint3(0, 0, 0), make_uint3(0, 0, 0), make_uint3(0, 0, 0), - nullptr, 0, 0, 0, make_uint3(0, 0, 0), make_uint3(0, 0, 0), make_uint3(0, 0, 0), make_uint3(0, 0, 0)); + nullptr, 0, 0, 0, make_uint3(0, 0, 0), make_uint3(0, 0, 0), make_uint3(0, 0, 0), make_uint3(0, 0, 0), scale_out); } else { const dim3 block_dims(1024, 1, 1); const ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params{blocks_num, block_dims, block_dims.x > WARP_SIZE ? 32 * sizeof(float): 0, stream}; - ggml_cuda_kernel_launch(rms_norm_f32<1024, false>, launch_params, x, dst, ncols, stride_row, stride_channel, stride_sample, eps, + ggml_cuda_kernel_launch(rms_norm_f32<1024, false, false, do_scale>, launch_params, x, dst, ncols, stride_row, stride_channel, stride_sample, eps, // underlying cudaLaunchKernelEx does not support default params nullptr, 0, 0, 0, make_uint3(0, 0, 0), make_uint3(0, 0, 0), make_uint3(0, 0, 0), make_uint3(0, 0, 0), - nullptr, 0, 0, 0, make_uint3(0, 0, 0), make_uint3(0, 0, 0), make_uint3(0, 0, 0), make_uint3(0, 0, 0)); + nullptr, 0, 0, 0, make_uint3(0, 0, 0), make_uint3(0, 0, 0), make_uint3(0, 0, 0), make_uint3(0, 0, 0), scale_out); } } @@ -367,7 +373,7 @@ static void rms_norm_mul_f32_cuda(const float * x, x, dst, ncols, stride_row, stride_channel, stride_sample, eps, mul, mul_stride_row, mul_stride_channel, mul_stride_sample, mul_ncols_packed, mul_nrows_packed, mul_nchannels_packed, mul_nsamples_packed, // underlying cudaLaunchKernelEx does not support default params - nullptr, 0, 0, 0, make_uint3(0, 0, 0), make_uint3(0, 0, 0), make_uint3(0, 0, 0), make_uint3(0, 0, 0)); + nullptr, 0, 0, 0, make_uint3(0, 0, 0), make_uint3(0, 0, 0), make_uint3(0, 0, 0), make_uint3(0, 0, 0), 1.0f); } else { const dim3 block_dims(1024, 1, 1); const ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params{blocks_num, block_dims, block_dims.x > WARP_SIZE ? 32 * sizeof(float): 0, stream}; @@ -375,7 +381,7 @@ static void rms_norm_mul_f32_cuda(const float * x, x, dst, ncols, stride_row, stride_channel, stride_sample, eps, mul, mul_stride_row, mul_stride_channel, mul_stride_sample, mul_ncols_packed, mul_nrows_packed, mul_nchannels_packed, mul_nsamples_packed, // underlying cudaLaunchKernelEx does not support default params - nullptr, 0, 0, 0, make_uint3(0, 0, 0), make_uint3(0, 0, 0), make_uint3(0, 0, 0), make_uint3(0, 0, 0)); + nullptr, 0, 0, 0, make_uint3(0, 0, 0), make_uint3(0, 0, 0), make_uint3(0, 0, 0), make_uint3(0, 0, 0), 1.0f); } } else { const uint3 mul_ncols_packed = init_fastdiv_values(mul_ncols); @@ -394,7 +400,7 @@ static void rms_norm_mul_f32_cuda(const float * x, x, dst, ncols, stride_row, stride_channel, stride_sample, eps, mul, mul_stride_row, mul_stride_channel, mul_stride_sample, mul_ncols_packed, mul_nrows_packed, mul_nchannels_packed, mul_nsamples_packed, add, add_stride_row, add_stride_channel, add_stride_sample, add_ncols_packed, add_nrows_packed, - add_nchannels_packed, add_nsamples_packed); + add_nchannels_packed, add_nsamples_packed, 1.0f); } else { const dim3 block_dims(1024, 1, 1); const ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params{blocks_num, block_dims, block_dims.x > WARP_SIZE ? 32 * sizeof(float): 0, stream}; @@ -402,7 +408,7 @@ static void rms_norm_mul_f32_cuda(const float * x, x, dst, ncols, stride_row, stride_channel, stride_sample, eps, mul, mul_stride_row, mul_stride_channel, mul_stride_sample, mul_ncols_packed, mul_nrows_packed, mul_nchannels_packed, mul_nsamples_packed, add, add_stride_row, add_stride_channel, add_stride_sample, add_ncols_packed, add_nrows_packed, - add_nchannels_packed, add_nsamples_packed); + add_nchannels_packed, add_nsamples_packed, 1.0f); } } } @@ -499,6 +505,33 @@ void ggml_cuda_op_rms_norm(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { rms_norm_f32_cuda(src0_d, dst_d, ne00, ne01, ne02, ne03, s01, s02, s03, eps, stream); } +void ggml_cuda_op_rms_norm_scale_fused(ggml_backend_cuda_context & ctx, ggml_tensor * dst, ggml_tensor * scale_tensor) { + const ggml_tensor * src0 = dst->src[0]; + const float * src0_d = (const float *) src0->data; + float * dst_d = (float *) scale_tensor->data; + cudaStream_t stream = ctx.stream(); + + GGML_ASSERT(src0->type == GGML_TYPE_F32); + GGML_ASSERT(scale_tensor->type == GGML_TYPE_F32); + + GGML_TENSOR_UNARY_OP_LOCALS; + + float eps; + memcpy(&eps, dst->op_params, sizeof(float)); + GGML_ASSERT(eps >= 0.0f); + + float scale; + memcpy(&scale, (const float *) scale_tensor->op_params + 0, sizeof(float)); + + const size_t ts0 = ggml_type_size(src0->type); + GGML_ASSERT(nb00 == ts0); + const int64_t s01 = nb01 / ts0; + const int64_t s02 = nb02 / ts0; + const int64_t s03 = nb03 / ts0; + + rms_norm_f32_cuda(src0_d, dst_d, ne00, ne01, ne02, ne03, s01, s02, s03, eps, stream, scale); +} + void ggml_cuda_op_rms_norm_fused(ggml_backend_cuda_context & ctx, ggml_tensor * dst, ggml_tensor * mul_tensor) { const ggml_tensor * rms_norm_src = (ggml_tensor *) dst->src[0]; float eps = 0.0f; diff --git a/ggml/src/ggml-cuda/norm.cuh b/ggml/src/ggml-cuda/norm.cuh index a74f63767..95618df5f 100644 --- a/ggml/src/ggml-cuda/norm.cuh +++ b/ggml/src/ggml-cuda/norm.cuh @@ -8,6 +8,8 @@ void ggml_cuda_op_rms_norm(ggml_backend_cuda_context & ctx, ggml_tensor * dst); void ggml_cuda_op_rms_norm_fused(ggml_backend_cuda_context & ctx, ggml_tensor * dst, ggml_tensor * mul_tensor); +void ggml_cuda_op_rms_norm_scale_fused(ggml_backend_cuda_context & ctx, ggml_tensor * dst, ggml_tensor * scale_tensor); + void ggml_cuda_op_rms_norm_fused_add(ggml_backend_cuda_context & ctx, ggml_tensor * dst, ggml_tensor * mul_tensor, diff --git a/ggml/src/ggml-cuda/ssm-scan.cu b/ggml/src/ggml-cuda/ssm-scan.cu index 40cb38dee..e9a1043fb 100644 --- a/ggml/src/ggml-cuda/ssm-scan.cu +++ b/ggml/src/ggml-cuda/ssm-scan.cu @@ -1,6 +1,6 @@ -#if !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) && CUDART_VERSION >= 11070 +#if !defined(GGML_USE_HIP) && (defined(GGML_USE_MUSA) || CUDART_VERSION >= 11070) #define USE_CUB -#endif // !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) && CUDART_VERSION >= 11070 +#endif // !defined(GGML_USE_HIP) && (defined(GGML_USE_MUSA) || CUDART_VERSION >= 11070) #ifdef USE_CUB #include @@ -342,7 +342,7 @@ static void ssm_scan_f32_cuda(const float * src0, const float * src1, const floa } } -#if !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) +#if !defined(GGML_USE_HIP) // ============================================================================ // SSD (State Space Duality) kernels for Mamba-2 prefill (n_tok > SSM_SSD_MIN_TOKENS) // @@ -821,7 +821,7 @@ void ggml_cuda_op_ssm_scan(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { GGML_ASSERT(src5->nb[2] <= (size_t)INT_MAX); GGML_ASSERT(src5->nb[3] <= (size_t)INT_MAX); -#if !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) +#if !defined(GGML_USE_HIP) // Mamba-2 with scalar A per head: use SSD matmul path for long sequences. // Requires NVIDIA Turing+ otherwise fallback to scan. const bool is_mamba2 = (src3->nb[1] == sizeof(float)); diff --git a/ggml/src/ggml-cuda/topk-moe.cu b/ggml/src/ggml-cuda/topk-moe.cu index dadcd601c..ee903bf21 100644 --- a/ggml/src/ggml-cuda/topk-moe.cu +++ b/ggml/src/ggml-cuda/topk-moe.cu @@ -98,7 +98,12 @@ __global__ void topk_moe_cuda(const float * logits, const float clamp_val, const float scale_val, const topk_moe_config config) { +#if defined(GGML_USE_MUSA) + // MUSA: every warp of a partially filled block must reach the barrier below. + const int row = MIN(blockIdx.x * blockDim.y + threadIdx.y, n_rows - 1); +#else const int row = blockIdx.x * blockDim.y + threadIdx.y; +#endif // defined(GGML_USE_MUSA) if (row >= n_rows) { return; } diff --git a/ggml/src/ggml-cuda/vendors/hip.h b/ggml/src/ggml-cuda/vendors/hip.h index 0a2f2829e..da7e3da0b 100644 --- a/ggml/src/ggml-cuda/vendors/hip.h +++ b/ggml/src/ggml-cuda/vendors/hip.h @@ -248,11 +248,11 @@ typedef __hip_bfloat16 nv_bfloat16; typedef __hip_bfloat162 nv_bfloat162; -#if HIP_VERSION >= 60200000 +#if HIP_VERSION >= 60300000 #include typedef __hip_fp8_e4m3 __nv_fp8_e4m3; #define FP8_AVAILABLE -#endif // HIP_VERSION >= 60200000 +#endif // HIP_VERSION >= 60300000 typedef int8_t int8x4_t __attribute__((ext_vector_type(4))); typedef uint8_t uint8x4_t __attribute__((ext_vector_type(4))); diff --git a/ggml/src/ggml-cuda/vendors/musa.h b/ggml/src/ggml-cuda/vendors/musa.h index 6d725c7ec..4243caab4 100644 --- a/ggml/src/ggml-cuda/vendors/musa.h +++ b/ggml/src/ggml-cuda/vendors/musa.h @@ -44,6 +44,7 @@ #define cudaDeviceGetPCIBusId musaDeviceGetPCIBusId #define cudaDeviceProp musaDeviceProp #define cudaDeviceSynchronize musaDeviceSynchronize +#define cudaDeviceGetAttribute musaDeviceGetAttribute #define cudaError_t musaError_t #define cudaErrorMemoryAllocation musaErrorMemoryAllocation #define cudaErrorPeerAccessAlreadyEnabled musaErrorPeerAccessAlreadyEnabled @@ -114,6 +115,7 @@ #define cuMemRelease muMemRelease #define cuMemSetAccess muMemSetAccess #define cuMemUnmap muMemUnmap +#define cudaDevAttrCooperativeLaunch musaDevAttrCooperativeLaunch #define cudaFuncAttributeMaxDynamicSharedMemorySize musaFuncAttributeMaxDynamicSharedMemorySize #define cudaFuncSetAttribute musaFuncSetAttribute #define cudaMemcpy3DPeerParms musaMemcpy3DPeerParms @@ -145,6 +147,9 @@ #define cudaStreamCaptureModeRelaxed musaStreamCaptureModeRelaxed #define cudaStreamBeginCapture musaStreamBeginCapture #define cudaStreamEndCapture musaStreamEndCapture +#define cudaStreamCaptureStatus musaStreamCaptureStatus +#define cudaStreamCaptureStatusNone musaStreamCaptureStatusNone +#define cudaStreamIsCapturing musaStreamIsCapturing #define cudaOccupancyMaxActiveBlocksPerMultiprocessor musaOccupancyMaxActiveBlocksPerMultiprocessor typedef __mt_bfloat16 nv_bfloat16; diff --git a/ggml/src/ggml-cuda/wkv.cu b/ggml/src/ggml-cuda/wkv.cu index 236111212..0bf997760 100644 --- a/ggml/src/ggml-cuda/wkv.cu +++ b/ggml/src/ggml-cuda/wkv.cu @@ -79,9 +79,7 @@ static __global__ void rwkv_wkv7_f32(const int B, const int T, const int C, cons float state[head_size]; __shared__ float _r[head_size], _w[head_size], _k[head_size], _a[head_size], _b[head_size]; -#ifndef GGML_USE_MUSA #pragma unroll -#endif for (int i = 0; i < head_size; i++) { state[i] = s[batch_i * state_size + head_i * head_size * head_size + tid * head_size + i]; } diff --git a/ggml/src/ggml-hexagon/ggml-hexagon.cpp b/ggml/src/ggml-hexagon/ggml-hexagon.cpp index 9a3500250..581f07aa2 100644 --- a/ggml/src/ggml-hexagon/ggml-hexagon.cpp +++ b/ggml/src/ggml-hexagon/ggml-hexagon.cpp @@ -267,10 +267,10 @@ static inline bool ggml_hexagon_is_repack_type(enum ggml_type type) { return type == GGML_TYPE_Q4_0 || type == GGML_TYPE_Q4_1 || type == GGML_TYPE_Q8_0 || type == GGML_TYPE_IQ4_NL || type == GGML_TYPE_MXFP4 || type == GGML_TYPE_Q6_K || - type == GGML_TYPE_Q4_K; + type == GGML_TYPE_Q4_K || type == GGML_TYPE_Q5_K; } -// Size of one repacked row in the DSP tiled layout. The Q6_K and Q4_K tiles store uncompressed scales/mins, +// Size of one repacked row in the DSP tiled layout. The Q6_K, Q5_K and Q4_K tiles store uncompressed scales/mins, // so they are larger than the ggml blocks. For the other repack types the tile has the same size as the ggml blocks. static inline size_t ggml_hexagon_tiled_row_size(enum ggml_type type, int64_t ne0) { if (type == GGML_TYPE_Q6_K) { @@ -279,6 +279,9 @@ static inline size_t ggml_hexagon_tiled_row_size(enum ggml_type type, int64_t ne if (type == GGML_TYPE_Q4_K) { return (size_t) (ne0 / 32) * (HTP_MM_WEIGHT_TILE_SIZE_Q4_1 / 32); } + if (type == GGML_TYPE_Q5_K) { + return (size_t) (ne0 / 32) * (HTP_MM_WEIGHT_TILE_SIZE_Q5_K / 32); + } return ggml_row_size(type, ne0); } @@ -1742,6 +1745,202 @@ static void repack_tiled_q4_K(void * data, const ggml_tensor * t, size_t offset, GGML_UNUSED(size); } +// tile layout: see HTP_MM_WEIGHT_TILE_SIZE_Q5_K in htp/matmul-ops.h +static void repack_q5_K_tiled(ggml_tensor * t, const void * data, size_t offset, size_t size) { + GGML_ASSERT(offset == 0); + + const block_q5_K * src_matrix = (const block_q5_K *) data; + int64_t ne0 = t->ne[0]; + int64_t ne1 = t->ne[1]; + int64_t ne2 = t->ne[2]; + int64_t ne3 = t->ne[3]; + int64_t ne0_padded = hex_round_up(ne0, 32); + int64_t ne1_padded = hex_round_up(ne1, 32); + + GGML_ASSERT(ne0 % QK_K == 0); + + const int n_col_tiles = ne1_padded / 32; + const int n_k_tiles = ne0_padded / 32; + const size_t tile_size = HTP_MM_WEIGHT_TILE_SIZE_Q5_K; + const size_t matrix_size = (size_t) n_col_tiles * n_k_tiles * tile_size; + + const int64_t sb_per_row = ne0 / QK_K; + + for (int i3 = 0; i3 < ne3; i3++) { + for (int i2 = 0; i2 < ne2; i2++) { + const block_q5_K * src_slice = src_matrix + (i3 * ne2 + i2) * (ne1 * sb_per_row); + uint8_t * matrix_dst = (uint8_t *) t->data + (i3 * ne2 + i2) * matrix_size; + + memset(matrix_dst, 0, matrix_size); + + for (int64_t r = 0; r < ne1; r++) { + const int ct = (int) (r / 32); + const int row = (int) (r % 32); + const block_q5_K * src_row = src_slice + r * sb_per_row; + + for (int kt = 0; kt < n_k_tiles; kt++) { + const int kt_local = kt % 8; + const block_q5_K * b = &src_row[kt / 8]; + const float d = GGML_FP16_TO_FP32(b->d); + const float dmin = GGML_FP16_TO_FP32(b->dmin); + + uint8_t * tile_dst = matrix_dst + ((size_t) ct * n_k_tiles + kt) * tile_size; + uint8_t * plane = tile_dst + 640; + + uint8_t sc, m; + get_scale_min_k4(kt_local, b->scales, &sc, &m); + + const float D = d * (float) sc; + const float M = -dmin * (float) m; + + const uint8_t * qs_sub = b->qs + (kt_local / 2) * 32; + const int shift = (kt_local & 1) ? 4 : 0; + const uint8_t hbit = (uint8_t) (1 << kt_local); + + for (int cp = 0; cp < 16; cp++) { + const uint8_t q0 = (qs_sub[2 * cp + 0] >> shift) & 0x0F; + const uint8_t q1 = (qs_sub[2 * cp + 1] >> shift) & 0x0F; + tile_dst[cp * 32 + row] = (uint8_t) ((q1 << 4) | q0); + + const int i = cp / 4; + const int lane = (cp % 4) * 32 + row; + if (b->qh[2 * cp + 0] & hbit) { + plane[lane] |= (uint8_t) (1 << (2 * i)); + } + if (b->qh[2 * cp + 1] & hbit) { + plane[lane] |= (uint8_t) (1 << (2 * i + 1)); + } + } + + ggml_half * scale_dst = (ggml_half *) (tile_dst + 512); + scale_dst[2 * row + 0] = GGML_FP32_TO_FP16(D); + scale_dst[2 * row + 1] = GGML_FP32_TO_FP16(M); + } + } + } + } + + GGML_UNUSED(size); +} + +// Reverse of repack_q5_K_tiled. Unpacks quants losslessly and normalizes scales/mins. Read-back only. +static void repack_tiled_q5_K(void * data, const ggml_tensor * t, size_t offset, size_t size) { + GGML_ASSERT(offset == 0); + + block_q5_K * dst_matrix = (block_q5_K *) data; + int64_t ne0 = t->ne[0]; + int64_t ne1 = t->ne[1]; + int64_t ne2 = t->ne[2]; + int64_t ne3 = t->ne[3]; + int64_t ne0_padded = hex_round_up(ne0, 32); + int64_t ne1_padded = hex_round_up(ne1, 32); + + GGML_ASSERT(ne0 % QK_K == 0); + + const int n_col_tiles = ne1_padded / 32; + const int n_k_tiles = ne0_padded / 32; + const size_t tile_size = HTP_MM_WEIGHT_TILE_SIZE_Q5_K; + const size_t matrix_size = (size_t) n_col_tiles * n_k_tiles * tile_size; + + const int64_t sb_per_row = ne0 / QK_K; + + for (int i3 = 0; i3 < ne3; i3++) { + for (int i2 = 0; i2 < ne2; i2++) { + block_q5_K * dst_slice = dst_matrix + (i3 * ne2 + i2) * (ne1 * sb_per_row); + const uint8_t * matrix_src = (const uint8_t *) t->data + (i3 * ne2 + i2) * matrix_size; + + for (int64_t r = 0; r < ne1; r++) { + const int ct = (int) (r / 32); + const int row = (int) (r % 32); + block_q5_K * dst_row = dst_slice + r * sb_per_row; + + for (int64_t sb = 0; sb < sb_per_row; sb++) { + block_q5_K * b = &dst_row[sb]; + memset(b, 0, sizeof(block_q5_K)); + + float sub_scales[8]; + float sub_mins[8]; + + for (int kt_local = 0; kt_local < 8; kt_local++) { + const int kt = sb * 8 + kt_local; + const uint8_t * tile_src = matrix_src + ((size_t) ct * n_k_tiles + kt) * tile_size; + const uint8_t * plane = tile_src + 640; + const ggml_half * scale_src = (const ggml_half *) (tile_src + 512); + + uint8_t * qs_sub = b->qs + (kt_local / 2) * 32; + const int shift = (kt_local & 1) ? 4 : 0; + const uint8_t hbit = (uint8_t) (1 << kt_local); + + for (int cp = 0; cp < 16; cp++) { + const uint8_t val = tile_src[cp * 32 + row]; + const uint8_t q0 = val & 0x0F; + const uint8_t q1 = val >> 4; + qs_sub[2 * cp + 0] |= (uint8_t) (q0 << shift); + qs_sub[2 * cp + 1] |= (uint8_t) (q1 << shift); + + const int i = cp / 4; + const int lane = (cp % 4) * 32 + row; + if (plane[lane] & (1 << (2 * i))) { + b->qh[2 * cp + 0] |= hbit; + } + if (plane[lane] & (1 << (2 * i + 1))) { + b->qh[2 * cp + 1] |= hbit; + } + } + + const float D = GGML_FP16_TO_FP32(scale_src[2 * row + 0]); + const float M = GGML_FP16_TO_FP32(scale_src[2 * row + 1]); + sub_scales[kt_local] = (D > 0.0f) ? D : 0.0f; + sub_mins[kt_local] = (-M > 0.0f) ? -M : 0.0f; + } + + float max_scale = 0.0f; + float max_min = 0.0f; + for (int j = 0; j < 8; j++) { + if (sub_scales[j] > max_scale) max_scale = sub_scales[j]; + if (sub_mins[j] > max_min) max_min = sub_mins[j]; + } + + float inv_scale = 0.0f; + if (max_scale > 0.0f) { + b->d = GGML_FP32_TO_FP16(max_scale / 63.0f); + const float d_actual = GGML_FP16_TO_FP32(b->d); + inv_scale = (d_actual > 0.0f) ? (1.0f / d_actual) : 0.0f; + } else { + b->d = GGML_FP32_TO_FP16(0.0f); + } + + float inv_min = 0.0f; + if (max_min > 0.0f) { + b->dmin = GGML_FP32_TO_FP16(max_min / 63.0f); + const float dmin_actual = GGML_FP16_TO_FP32(b->dmin); + inv_min = (dmin_actual > 0.0f) ? (1.0f / dmin_actual) : 0.0f; + } else { + b->dmin = GGML_FP32_TO_FP16(0.0f); + } + + for (int j = 0; j < 8; j++) { + uint8_t ls = (uint8_t) roundf(inv_scale * sub_scales[j]); + uint8_t lm = (uint8_t) roundf(inv_min * sub_mins[j]); + ls = (std::min)((uint8_t) 63, ls); + lm = (std::min)((uint8_t) 63, lm); + if (j < 4) { + b->scales[j] = ls; + b->scales[j + 4] = lm; + } else { + b->scales[j + 4] = (ls & 0xF) | ((lm & 0xF) << 4); + b->scales[j - 4] |= ((ls >> 4) << 6); + b->scales[j - 0] |= ((lm >> 4) << 6); + } + } + } + } + } + } + + GGML_UNUSED(size); +} + static void repack_tensor_tiled(ggml_tensor * tensor, const void * data, size_t size) { switch (tensor->type) { case GGML_TYPE_Q4_0: @@ -1768,6 +1967,10 @@ static void repack_tensor_tiled(ggml_tensor * tensor, const void * data, size_t repack_mxfp4_tiled(tensor, data, 0, size); break; + case GGML_TYPE_Q5_K: + repack_q5_K_tiled(tensor, data, 0, size); + break; + case GGML_TYPE_Q6_K: repack_q6_K_tiled(tensor, data, 0, size); break; @@ -1856,6 +2059,12 @@ static void ggml_backend_hexagon_buffer_get_tensor(ggml_backend_buffer_t buffer, repack_tiled_q4_K(data, tensor, offset, size); break; + case GGML_TYPE_Q5_K: + GGML_ASSERT(offset == 0); + GGML_ASSERT(offset + size <= ggml_nbytes(tensor)); + repack_tiled_q5_K(data, tensor, offset, size); + break; + case GGML_TYPE_Q8_0: GGML_ASSERT(offset == 0); GGML_ASSERT(offset + size <= ggml_nbytes(tensor)); @@ -1989,6 +2198,10 @@ static void ggml_backend_hexagon_buffer_get_tensor_2d(ggml_backend_buffer_t buff repack_tiled_q4_K(temp_buf.data(), tensor, offset, temp_size); break; + case GGML_TYPE_Q5_K: + repack_tiled_q5_K(temp_buf.data(), tensor, offset, temp_size); + break; + case GGML_TYPE_Q8_0: repack_tiled_q8_0(temp_buf.data(), tensor, offset, temp_size); break; @@ -4494,8 +4707,7 @@ static bool ggml_hexagon_precompute_hmx_mm_params( if (is_batched_val && wtype == GGML_TYPE_F16 && group_size > 1) { // Try grouped path first - const bool use_dma_activation = (src1->nb[1]/sizeof(float) > (size_t)ne00_padded); - if (htp_mm_hmx_solve_batched_params(wtype, ne00_padded, ne01_padded, ne11, group_size, use_dma_activation, n_threads, pipeline, src2_size, vtcm_budget, &m_chunk, &n_chunk, &act_threads_selected, &vtcm_size)) { + if (htp_mm_hmx_solve_batched_params(wtype, ne00_padded, ne01_padded, ne11, group_size, n_threads, pipeline, src2_size, vtcm_budget, &m_chunk, &n_chunk, &act_threads_selected, &vtcm_size)) { use_grouped = true; } } @@ -5404,12 +5616,13 @@ static bool ggml_hexagon_supported_mul_mat(const struct ggml_hexagon_session * s case GGML_TYPE_IQ4_NL: case GGML_TYPE_MXFP4: case GGML_TYPE_Q4_K: + case GGML_TYPE_Q5_K: case GGML_TYPE_Q6_K: if (!ggml_is_contiguous(src0) || ggml_is_permuted(src0)) { return false; } - if (src0->ne[0] % ((src0->type == GGML_TYPE_Q6_K || src0->type == GGML_TYPE_Q4_K) ? QK_K : 32)) { + if (src0->ne[0] % ((src0->type == GGML_TYPE_Q6_K || src0->type == GGML_TYPE_Q5_K || src0->type == GGML_TYPE_Q4_K) ? QK_K : 32)) { return false; } @@ -5487,12 +5700,13 @@ static bool ggml_hexagon_supported_mul_mat_id(const struct ggml_hexagon_session case GGML_TYPE_IQ4_NL: case GGML_TYPE_MXFP4: case GGML_TYPE_Q4_K: + case GGML_TYPE_Q5_K: case GGML_TYPE_Q6_K: if (!ggml_is_contiguous(src0) || ggml_is_permuted(src0)) { return false; } - if (src0->ne[0] % ((src0->type == GGML_TYPE_Q6_K || src0->type == GGML_TYPE_Q4_K) ? QK_K : 32)) { + if (src0->ne[0] % ((src0->type == GGML_TYPE_Q6_K || src0->type == GGML_TYPE_Q5_K || src0->type == GGML_TYPE_Q4_K) ? QK_K : 32)) { return false; } @@ -7089,18 +7303,21 @@ static bool ggml_hexagon_supported_cpy(const struct ggml_hexagon_session * sess, const struct ggml_tensor * src0 = op->src[0]; const struct ggml_tensor * dst = op; - // for now we can do f32 -> f16 and f16 -> f32 (without reshaping) - if (src0->type != GGML_TYPE_F32 && src0->type != GGML_TYPE_F16) return false; - if ( dst->type != GGML_TYPE_F32 && dst->type != GGML_TYPE_F16) return false; + if (src0->type != GGML_TYPE_F32 && src0->type != GGML_TYPE_F16 && + src0->type != GGML_TYPE_I32) return false; + if (dst->type != GGML_TYPE_F32 && dst->type != GGML_TYPE_F16 && + dst->type != GGML_TYPE_I32) return false; const bool sametype = (src0->type == dst->type); const bool transposed = ggml_is_transposed(src0) || ggml_is_transposed(dst); const bool sameshape = !transposed && ggml_are_same_shape(src0, dst); - // can handle any shape and any same-type (pretty slow if reshaping is required) + // Same-type copies also support I32. if (sametype) return true; - // cannot handle re-shaping and type conversion at the same time + // Type conversion is only supported between F32 and F16. + if (src0->type == GGML_TYPE_I32 || dst->type == GGML_TYPE_I32) return false; + if (!sameshape) return false; return true; @@ -7110,8 +7327,9 @@ static bool ggml_hexagon_supported_cont(const struct ggml_hexagon_session * sess GGML_UNUSED(sess); const struct ggml_tensor * src0 = op->src[0]; - // CONT is same-type only, supports f32 and f16 - if (src0->type != GGML_TYPE_F32 && src0->type != GGML_TYPE_F16) return false; + // CONT is same-type only and supports F32, F16, and I32. + if (src0->type != GGML_TYPE_F32 && src0->type != GGML_TYPE_F16 && + src0->type != GGML_TYPE_I32) return false; return true; } @@ -7754,6 +7972,8 @@ static void ggml_hexagon_init(ggml_backend_reg * reg) { "please update hexagon_type to match ggml_type"); static_assert((unsigned int) HTP_TYPE_Q4_K == (unsigned int) GGML_TYPE_Q4_K, "please update hexagon_type to match ggml_type"); + static_assert((unsigned int) HTP_TYPE_Q5_K == (unsigned int) GGML_TYPE_Q5_K, + "please update hexagon_type to match ggml_type"); static_assert((unsigned int) HTP_TYPE_Q6_K == (unsigned int) GGML_TYPE_Q6_K, "please update hexagon_type to match ggml_type"); diff --git a/ggml/src/ggml-hexagon/htp/concat-ops.c b/ggml/src/ggml-hexagon/htp/concat-ops.c index 1fa6ec1bd..382a94660 100644 --- a/ggml/src/ggml-hexagon/htp/concat-ops.c +++ b/ggml/src/ggml-hexagon/htp/concat-ops.c @@ -20,6 +20,7 @@ struct htp_concat_context { uint32_t nrows; uint32_t elem_start; uint32_t nelems; + uint32_t nplanes; struct fastdiv_values div_ne0; struct fastdiv_values div_ne1; struct fastdiv_values div_ne2; @@ -60,39 +61,47 @@ static void concat_2d_f32_transposed(unsigned int nth, unsigned int ith, void * struct htp_thread_trace * tr = &octx->ctx->trace[ith]; - for (uint32_t i = start_i; i < end_i; i += block_i) { - uint32_t current_block_i = (end_i - i < block_i) ? (end_i - i) : block_i; + for (uint32_t p = 0; p < cctx->nplanes; p++) { + const uint32_t i3 = p / dst->ne[2]; + const uint32_t i2 = p - i3 * dst->ne[2]; + const dma_addr_t src0_plane = src0->data + i2 * src0->nb[2] + i3 * src0->nb[3]; + const dma_addr_t src1_plane = src1->data + i2 * src1->nb[2] + i3 * src1->nb[3]; + const dma_addr_t dst_plane = dst->data + i2 * dst->nb[2] + i3 * dst->nb[3]; - uint32_t src1_width_bytes = current_block_i * sizeof(float); - const dma_addr_t src1_addr = src1->data + i * src1->nb[1]; - dma_queue_push(dma_q, dma_make_data(spad1_base, src1_addr), spad1_stride, src1->nb[0], src1_width_bytes, src1_ne0); + for (uint32_t i = start_i; i < end_i; i += block_i) { + uint32_t current_block_i = (end_i - i < block_i) ? (end_i - i) : block_i; - uint32_t src0_row_bytes = src0_ne0 * sizeof(float); - const dma_addr_t src0_addr = src0->data + i * src0->nb[1]; - dma_queue_push(dma_q, dma_make_data(spad0_base, src0_addr), spad0_row_bytes, src0->nb[1], src0_row_bytes, current_block_i); + uint32_t src1_width_bytes = current_block_i * sizeof(float); + const dma_addr_t src1_addr = src1_plane + i * src1->nb[1]; + dma_queue_push(dma_q, dma_make_data(spad1_base, src1_addr), spad1_stride, src1->nb[0], src1_width_bytes, src1_ne0); - dma_queue_pop(dma_q); // src1 + uint32_t src0_row_bytes = src0_ne0 * sizeof(float); + const dma_addr_t src0_addr = src0_plane + i * src0->nb[1]; + dma_queue_push(dma_q, dma_make_data(spad0_base, src0_addr), spad0_row_bytes, src0->nb[1], src0_row_bytes, current_block_i); - HVX_Vector * vtcm_tmp = (HVX_Vector *)(spad1_base + src1_ne0_padded * spad1_stride); + dma_queue_pop(dma_q); // src1 - htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) i); - for (uint32_t j = 0; j < src1_ne0_padded; j += 32) { - #pragma unroll(4) - for (uint32_t ii = 0; ii < current_block_i; ii++) { - size_t rt = (size_t)(spad1_base + j * spad1_stride + ii * sizeof(float)); - Q6_vgather_ARMVw(&vtcm_tmp[ii], rt, mu, vv); - uint8_t * dst_ptr = spad0_base + ii * spad0_row_bytes + (src0_ne0 + j) * sizeof(float); - hvx_vmemu(dst_ptr) = vtcm_tmp[ii]; + HVX_Vector * vtcm_tmp = (HVX_Vector *)(spad1_base + src1_ne0_padded * spad1_stride); + + htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) i); + for (uint32_t j = 0; j < src1_ne0_padded; j += 32) { + #pragma unroll(4) + for (uint32_t ii = 0; ii < current_block_i; ii++) { + size_t rt = (size_t)(spad1_base + j * spad1_stride + ii * sizeof(float)); + Q6_vgather_ARMVw(&vtcm_tmp[ii], rt, mu, vv); + uint8_t * dst_ptr = spad0_base + ii * spad0_row_bytes + (src0_ne0 + j) * sizeof(float); + hvx_vmemu(dst_ptr) = vtcm_tmp[ii]; + } } + htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) i); + + dma_queue_pop(dma_q); // src0 + + const dma_addr_t dst_addr = dst_plane + i * dst->nb[1]; + dma_queue_push(dma_q, dma_make_data(dst_addr, spad0_base), dst->nb[1], spad0_row_bytes, (src0_ne0 + src1_ne0) * sizeof(float), current_block_i); + + dma_queue_pop(dma_q); } - htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) i); - - dma_queue_pop(dma_q); // src0 - - const dma_addr_t dst_addr = dst->data + i * dst->nb[1]; - dma_queue_push(dma_q, dma_make_data(dst_addr, spad0_base), dst->nb[1], spad0_row_bytes, (src0_ne0 + src1_ne0) * sizeof(float), current_block_i); - - dma_queue_pop(dma_q); } } @@ -131,39 +140,47 @@ static void concat_2d_f16_transposed(unsigned int nth, unsigned int ith, void * struct htp_thread_trace * tr = &octx->ctx->trace[ith]; - for (uint32_t i = start_i; i < end_i; i += block_i) { - uint32_t current_block_i = (end_i - i < block_i) ? (end_i - i) : block_i; + for (uint32_t p = 0; p < cctx->nplanes; p++) { + const uint32_t i3 = p / dst->ne[2]; + const uint32_t i2 = p - i3 * dst->ne[2]; + const dma_addr_t src0_plane = src0->data + i2 * src0->nb[2] + i3 * src0->nb[3]; + const dma_addr_t src1_plane = src1->data + i2 * src1->nb[2] + i3 * src1->nb[3]; + const dma_addr_t dst_plane = dst->data + i2 * dst->nb[2] + i3 * dst->nb[3]; - uint32_t src1_width_bytes = current_block_i * sizeof(__fp16); - const dma_addr_t src1_addr = src1->data + i * src1->nb[1]; - dma_queue_push(dma_q, dma_make_data(spad1_base, src1_addr), spad1_stride, src1->nb[0], src1_width_bytes, src1_ne0); + for (uint32_t i = start_i; i < end_i; i += block_i) { + uint32_t current_block_i = (end_i - i < block_i) ? (end_i - i) : block_i; - uint32_t src0_row_bytes = src0_ne0 * sizeof(__fp16); - const dma_addr_t src0_addr = src0->data + i * src0->nb[1]; - dma_queue_push(dma_q, dma_make_data(spad0_base, src0_addr), spad0_row_bytes, src0->nb[1], src0_row_bytes, current_block_i); + uint32_t src1_width_bytes = current_block_i * sizeof(__fp16); + const dma_addr_t src1_addr = src1_plane + i * src1->nb[1]; + dma_queue_push(dma_q, dma_make_data(spad1_base, src1_addr), spad1_stride, src1->nb[0], src1_width_bytes, src1_ne0); - dma_queue_pop(dma_q); // src1 + uint32_t src0_row_bytes = src0_ne0 * sizeof(__fp16); + const dma_addr_t src0_addr = src0_plane + i * src0->nb[1]; + dma_queue_push(dma_q, dma_make_data(spad0_base, src0_addr), spad0_row_bytes, src0->nb[1], src0_row_bytes, current_block_i); - HVX_Vector * vtcm_tmp = (HVX_Vector *)(spad1_base + src1_ne0_padded * spad1_stride); + dma_queue_pop(dma_q); // src1 - htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) i); - for (uint32_t j = 0; j < src1_ne0_padded; j += 64) { - #pragma unroll(4) - for (uint32_t ii = 0; ii < current_block_i; ii++) { - size_t rt = (size_t)(spad1_base + j * spad1_stride + ii * sizeof(__fp16)); - Q6_vgather_ARMVh(&vtcm_tmp[ii], rt, mu, vv); - uint8_t * dst_ptr = spad0_base + ii * spad0_row_bytes + (src0_ne0 + j) * sizeof(__fp16); - hvx_vmemu(dst_ptr) = vtcm_tmp[ii]; + HVX_Vector * vtcm_tmp = (HVX_Vector *)(spad1_base + src1_ne0_padded * spad1_stride); + + htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) i); + for (uint32_t j = 0; j < src1_ne0_padded; j += 64) { + #pragma unroll(4) + for (uint32_t ii = 0; ii < current_block_i; ii++) { + size_t rt = (size_t)(spad1_base + j * spad1_stride + ii * sizeof(__fp16)); + Q6_vgather_ARMVh(&vtcm_tmp[ii], rt, mu, vv); + uint8_t * dst_ptr = spad0_base + ii * spad0_row_bytes + (src0_ne0 + j) * sizeof(__fp16); + hvx_vmemu(dst_ptr) = vtcm_tmp[ii]; + } } + htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) i); + + dma_queue_pop(dma_q); // src0 + + const dma_addr_t dst_addr = dst_plane + i * dst->nb[1]; + dma_queue_push(dma_q, dma_make_data(dst_addr, spad0_base), dst->nb[1], spad0_row_bytes, (src0_ne0 + src1_ne0) * sizeof(__fp16), current_block_i); + + dma_queue_pop(dma_q); } - htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) i); - - dma_queue_pop(dma_q); // src0 - - const dma_addr_t dst_addr = dst->data + i * dst->nb[1]; - dma_queue_push(dma_q, dma_make_data(dst_addr, spad0_base), dst->nb[1], spad0_row_bytes, (src0_ne0 + src1_ne0) * sizeof(__fp16), current_block_i); - - dma_queue_pop(dma_q); } } @@ -230,6 +247,61 @@ static void concat_generic(unsigned int nth, unsigned int ith, void * data) { } } +static bool concat_dim1_contiguous_dma(struct htp_ops_context * octx, int dim, uint32_t type_size) { + const struct htp_tensor * src0 = octx->src[0]; + const struct htp_tensor * src1 = octx->src[1]; + const struct htp_tensor * dst = octx->dst; + + if (dim != 1 || octx->ctx->mdev.count > 1 || + (dst->type != HTP_TYPE_F32 && dst->type != HTP_TYPE_F16 && dst->type != HTP_TYPE_I32) || + src0->type != dst->type || src1->type != dst->type || + src0->ne[0] != dst->ne[0] || src1->ne[0] != dst->ne[0] || + src0->ne[2] != dst->ne[2] || src1->ne[2] != dst->ne[2] || + src0->ne[3] != dst->ne[3] || src1->ne[3] != dst->ne[3] || + dst->ne[1] != src0->ne[1] + src1->ne[1] || + !htp_tensor_is_contiguous(src0, type_size) || + !htp_tensor_is_contiguous(src1, type_size) || + !htp_tensor_is_contiguous(dst, type_size)) { + return false; + } + + const uint32_t src0_row_size = src0->ne[0] * type_size; + const uint32_t src1_row_size = src1->ne[0] * type_size; + + // v75+ dma_queue_push() writes a 2D descriptor directly and does not split overflow. +#if __HVX_ARCH__ >= 75 + if (src0_row_size > 0xffffffu || src1_row_size > 0xffffffu || + src0->nb[1] > 0xffffffu || src1->nb[1] > 0xffffffu || dst->nb[1] > 0xffffffu || + src0->ne[1] > UINT16_MAX || src1->ne[1] > UINT16_MAX) { + return false; + } +#endif + + dma_queue * q = octx->ctx->dma[0]; + + for (uint32_t i3 = 0; i3 < dst->ne[3]; ++i3) { + for (uint32_t i2 = 0; i2 < dst->ne[2]; ++i2) { + dma_addr_t dst_addr = dst->data + i3 * dst->nb[3] + i2 * dst->nb[2]; + dma_addr_t src0_addr = src0->data + i3 * src0->nb[3] + i2 * src0->nb[2]; + dma_addr_t src1_addr = src1->data + i3 * src1->nb[3] + i2 * src1->nb[2]; + + if (!dma_queue_push(q, dma_make_data(dst_addr, src0_addr), dst->nb[1], src0->nb[1], src0_row_size, src0->ne[1])) { + dma_queue_flush(q); + dma_queue_push(q, dma_make_data(dst_addr, src0_addr), dst->nb[1], src0->nb[1], src0_row_size, src0->ne[1]); + } + + dst_addr += src0->ne[1] * dst->nb[1]; + if (!dma_queue_push(q, dma_make_data(dst_addr, src1_addr), dst->nb[1], src1->nb[1], src1_row_size, src1->ne[1])) { + dma_queue_flush(q); + dma_queue_push(q, dma_make_data(dst_addr, src1_addr), dst->nb[1], src1->nb[1], src1_row_size, src1->ne[1]); + } + } + } + + dma_queue_flush(q); + return true; +} + int op_concat(struct htp_ops_context * octx) { const struct htp_tensor * src0 = octx->src[0]; const struct htp_tensor * src1 = octx->src[1]; @@ -237,12 +309,14 @@ int op_concat(struct htp_ops_context * octx) { int dim = octx->op_params[0]; - bool is_2d = dst->ne[2] == 1 && dst->ne[3] == 1; - const uint32_t type_size = (dst->type == HTP_TYPE_F32 || dst->type == HTP_TYPE_I32) ? 4 : 2; bool is_src1_transposed = (src1->nb[0] > src1->nb[1]); bool is_src0_transposed = (src0->nb[0] > src0->nb[1]); + if (concat_dim1_contiguous_dma(octx, dim, type_size)) { + return HTP_STATUS_OK; + } + uint32_t n_threads = octx->n_threads; struct htp_concat_context cctx; cctx.octx = octx; @@ -253,7 +327,9 @@ int op_concat(struct htp_ops_context * octx) { void (*worker_func)(unsigned int, unsigned int, void *) = concat_generic; - if (dim == 0 && is_2d && is_src1_transposed && !is_src0_transposed) { + const bool rows_ok = src0->nb[0] == type_size && src1->nb[1] == type_size && dst->nb[0] == type_size; + + if (dim == 0 && is_src1_transposed && !is_src0_transposed && rows_ok) { const uint32_t total_rows = dst->ne[1]; const size_t dst_data_row_size = dst->ne[0] * type_size; uint32_t row_start = 0; @@ -272,6 +348,7 @@ int op_concat(struct htp_ops_context * octx) { cctx.row_start = row_start; cctx.nrows = nrows; + cctx.nplanes = dst->ne[2] * dst->ne[3]; uint32_t block_i = (type_size == 4) ? 32 : 64; diff --git a/ggml/src/ggml-hexagon/htp/cpy-ops.c b/ggml/src/ggml-hexagon/htp/cpy-ops.c index 4453dda3a..785e3e49d 100644 --- a/ggml/src/ggml-hexagon/htp/cpy-ops.c +++ b/ggml/src/ggml-hexagon/htp/cpy-ops.c @@ -140,6 +140,7 @@ static void cpy_thread_##NAME##_sameshape(unsigned int nth, unsigned int ith, vo DEFINE_CPY_SAMESHAPE(f32, float, 4) DEFINE_CPY_SAMESHAPE(f16, __fp16, 2) +DEFINE_CPY_SAMESHAPE(i32, int32_t, 4) #define DEFINE_CPY_RESHAPE(NAME, ELEM_TYPE, ELEM_SIZE) \ static void cpy_thread_##NAME##_reshape(unsigned int nth, unsigned int ith, void * data) { \ @@ -226,6 +227,7 @@ static void cpy_thread_##NAME##_reshape(unsigned int nth, unsigned int ith, void DEFINE_CPY_RESHAPE(f32, float, 4) DEFINE_CPY_RESHAPE(f16, __fp16, 2) +DEFINE_CPY_RESHAPE(i32, int32_t, 4) static void cpy_thread_f16_f32_sameshape(unsigned int nth, unsigned int ith, void * data) { struct htp_copy_context * ct = (struct htp_copy_context *) data; @@ -370,6 +372,7 @@ static int exec_cpy(struct htp_ops_context * octx, bool * use_dma) { switch (src0->type) { case HTP_TYPE_F32: ct.src0_type_size = 4; ct.src0_block_size = 1; ct.src0_blocks_per_row = ne00 / 1; break; case HTP_TYPE_F16: ct.src0_type_size = 2; ct.src0_block_size = 1; ct.src0_blocks_per_row = ne00 / 1; break; + case HTP_TYPE_I32: ct.src0_type_size = 4; ct.src0_block_size = 1; ct.src0_blocks_per_row = ne00 / 1; break; default: return HTP_STATUS_NO_SUPPORT; } @@ -377,6 +380,7 @@ static int exec_cpy(struct htp_ops_context * octx, bool * use_dma) { switch (dst->type) { case HTP_TYPE_F32: ct.dst_type_size = 4; ct.dst_block_size = 1; ct.dst_blocks_per_row = ne0 / 1; break; case HTP_TYPE_F16: ct.dst_type_size = 2; ct.dst_block_size = 1; ct.dst_blocks_per_row = ne0 / 1; break; + case HTP_TYPE_I32: ct.dst_type_size = 4; ct.dst_block_size = 1; ct.dst_blocks_per_row = ne0 / 1; break; default: return HTP_STATUS_NO_SUPPORT; } @@ -436,7 +440,12 @@ static int exec_cpy(struct htp_ops_context * octx, bool * use_dma) { } else { work_queue_func_t copy_fun = NULL; if (sametype) { - copy_fun = (src0->type == HTP_TYPE_F32) ? cpy_thread_f32_sameshape : cpy_thread_f16_sameshape; + switch (src0->type) { + case HTP_TYPE_F32: copy_fun = cpy_thread_f32_sameshape; break; + case HTP_TYPE_F16: copy_fun = cpy_thread_f16_sameshape; break; + case HTP_TYPE_I32: copy_fun = cpy_thread_i32_sameshape; break; + default: return HTP_STATUS_NO_SUPPORT; + } } else if (dst->type == HTP_TYPE_F16 && src0->type == HTP_TYPE_F32) { copy_fun = cpy_thread_f16_f32_sameshape; } else if (dst->type == HTP_TYPE_F32 && src0->type == HTP_TYPE_F16) { @@ -482,7 +491,13 @@ static int exec_cpy(struct htp_ops_context * octx, bool * use_dma) { ct.nelem = nelem; ct.elem_per_thread = fastdiv(nelem + n_threads - 1, &octx->n_threads_div); - work_queue_func_t copy_fun = (src0->type == HTP_TYPE_F32) ? cpy_thread_f32_reshape : cpy_thread_f16_reshape; + work_queue_func_t copy_fun = NULL; + switch (src0->type) { + case HTP_TYPE_F32: copy_fun = cpy_thread_f32_reshape; break; + case HTP_TYPE_F16: copy_fun = cpy_thread_f16_reshape; break; + case HTP_TYPE_I32: copy_fun = cpy_thread_i32_reshape; break; + default: return HTP_STATUS_NO_SUPPORT; + } work_queue_run(octx->ctx->work_queue, copy_fun, &ct, n_threads); } else { return HTP_STATUS_NO_SUPPORT; diff --git a/ggml/src/ggml-hexagon/htp/hmx-mm-kernels-tiled.h b/ggml/src/ggml-hexagon/htp/hmx-mm-kernels-tiled.h index d6d40586c..5b7f34023 100644 --- a/ggml/src/ggml-hexagon/htp/hmx-mm-kernels-tiled.h +++ b/ggml/src/ggml-hexagon/htp/hmx-mm-kernels-tiled.h @@ -506,6 +506,111 @@ static void dequantize_tiled_weight_to_fp16_task_q8_0( } } +static void dequantize_tiled_weight_to_fp16_task_q5_k( + const tiled_dequantize_state_t *state, + uint32_t start_tile, uint32_t end_tile) { + + const HVX_Vector mask_h4 = Q6_Vb_vsplat_R(0x0F); + + for (uint32_t t = start_tile; t < end_tile; t++) { + const uint8_t * tile_src = state->src + t * state->aligned_tile_size; + __fp16 * dst_ptr = state->dst + t * HTP_MM_HMX_TILE_N_ELMS; + + HVX_Vector vscale_offset = hvx_vmem(tile_src + 512); + HVX_VectorPair dm_deal = Q6_W_vdeal_VVR(vscale_offset, vscale_offset, -2); + HVX_Vector vd = Q6_V_lo_W(dm_deal); + HVX_Vector vm = Q6_V_hi_W(dm_deal); + + HVX_Vector v_scale_duplicated = Q6_V_lo_W(Q6_W_vshuff_VVR(vd, vd, -2)); + HVX_Vector v_offset_duplicated = Q6_V_lo_W(Q6_W_vshuff_VVR(vm, vm, -2)); + + // Load all 4 groups in parallel + HVX_Vector vq0 = hvx_vmem(tile_src + 0 * 128); + HVX_Vector vq1 = hvx_vmem(tile_src + 1 * 128); + HVX_Vector vq2 = hvx_vmem(tile_src + 2 * 128); + HVX_Vector vq3 = hvx_vmem(tile_src + 3 * 128); + + // Nibble extraction + HVX_Vector v_lo0 = Q6_V_vand_VV(vq0, mask_h4); + HVX_Vector v_hi0 = Q6_Vub_vlsr_VubR(vq0, 4); + HVX_Vector v_lo1 = Q6_V_vand_VV(vq1, mask_h4); + HVX_Vector v_hi1 = Q6_Vub_vlsr_VubR(vq1, 4); + HVX_Vector v_lo2 = Q6_V_vand_VV(vq2, mask_h4); + HVX_Vector v_hi2 = Q6_Vub_vlsr_VubR(vq2, 4); + HVX_Vector v_lo3 = Q6_V_vand_VV(vq3, mask_h4); + HVX_Vector v_hi3 = Q6_Vub_vlsr_VubR(vq3, 4); + + // Q5_K: OR in the 5th bit from the plane + HVX_Vector v_plane = hvx_vmem(tile_src + 640); + v_lo0 = hvx_q5k_or_hibit(v_lo0, v_plane, 0); + v_hi0 = hvx_q5k_or_hibit(v_hi0, v_plane, 1); + v_lo1 = hvx_q5k_or_hibit(v_lo1, v_plane, 2); + v_hi1 = hvx_q5k_or_hibit(v_hi1, v_plane, 3); + v_lo2 = hvx_q5k_or_hibit(v_lo2, v_plane, 4); + v_hi2 = hvx_q5k_or_hibit(v_hi2, v_plane, 5); + v_lo3 = hvx_q5k_or_hibit(v_lo3, v_plane, 6); + v_hi3 = hvx_q5k_or_hibit(v_hi3, v_plane, 7); + + // Shuffling + HVX_VectorPair vp_shuf0 = Q6_W_vshuff_VVR(v_hi0, v_lo0, -1); + HVX_VectorPair vp_shuf1 = Q6_W_vshuff_VVR(v_hi1, v_lo1, -1); + HVX_VectorPair vp_shuf2 = Q6_W_vshuff_VVR(v_hi2, v_lo2, -1); + HVX_VectorPair vp_shuf3 = Q6_W_vshuff_VVR(v_hi3, v_lo3, -1); + + // Unpack to 16-bit + HVX_VectorPair vp_int16_lo0 = Q6_Wh_vunpack_Vb(Q6_V_lo_W(vp_shuf0)); + HVX_VectorPair vp_int16_hi0 = Q6_Wh_vunpack_Vb(Q6_V_hi_W(vp_shuf0)); + HVX_VectorPair vp_int16_lo1 = Q6_Wh_vunpack_Vb(Q6_V_lo_W(vp_shuf1)); + HVX_VectorPair vp_int16_hi1 = Q6_Wh_vunpack_Vb(Q6_V_hi_W(vp_shuf1)); + HVX_VectorPair vp_int16_lo2 = Q6_Wh_vunpack_Vb(Q6_V_lo_W(vp_shuf2)); + HVX_VectorPair vp_int16_hi2 = Q6_Wh_vunpack_Vb(Q6_V_hi_W(vp_shuf2)); + HVX_VectorPair vp_int16_lo3 = Q6_Wh_vunpack_Vb(Q6_V_lo_W(vp_shuf3)); + HVX_VectorPair vp_int16_hi3 = Q6_Wh_vunpack_Vb(Q6_V_hi_W(vp_shuf3)); + + // Convert, multiply, add offset + HVX_Vector v_grp0_0 = Q6_Vhf_equals_Vqf16(Q6_Vqf16_vadd_Vqf16Vhf(Q6_Vqf16_vmpy_VhfVhf(Q6_Vhf_equals_Vh(Q6_V_lo_W(vp_int16_lo0)), v_scale_duplicated), v_offset_duplicated)); + HVX_Vector v_grp0_1 = Q6_Vhf_equals_Vqf16(Q6_Vqf16_vadd_Vqf16Vhf(Q6_Vqf16_vmpy_VhfVhf(Q6_Vhf_equals_Vh(Q6_V_hi_W(vp_int16_lo0)), v_scale_duplicated), v_offset_duplicated)); + HVX_Vector v_grp0_2 = Q6_Vhf_equals_Vqf16(Q6_Vqf16_vadd_Vqf16Vhf(Q6_Vqf16_vmpy_VhfVhf(Q6_Vhf_equals_Vh(Q6_V_lo_W(vp_int16_hi0)), v_scale_duplicated), v_offset_duplicated)); + HVX_Vector v_grp0_3 = Q6_Vhf_equals_Vqf16(Q6_Vqf16_vadd_Vqf16Vhf(Q6_Vqf16_vmpy_VhfVhf(Q6_Vhf_equals_Vh(Q6_V_hi_W(vp_int16_hi0)), v_scale_duplicated), v_offset_duplicated)); + + HVX_Vector v_grp1_0 = Q6_Vhf_equals_Vqf16(Q6_Vqf16_vadd_Vqf16Vhf(Q6_Vqf16_vmpy_VhfVhf(Q6_Vhf_equals_Vh(Q6_V_lo_W(vp_int16_lo1)), v_scale_duplicated), v_offset_duplicated)); + HVX_Vector v_grp1_1 = Q6_Vhf_equals_Vqf16(Q6_Vqf16_vadd_Vqf16Vhf(Q6_Vqf16_vmpy_VhfVhf(Q6_Vhf_equals_Vh(Q6_V_hi_W(vp_int16_lo1)), v_scale_duplicated), v_offset_duplicated)); + HVX_Vector v_grp1_2 = Q6_Vhf_equals_Vqf16(Q6_Vqf16_vadd_Vqf16Vhf(Q6_Vqf16_vmpy_VhfVhf(Q6_Vhf_equals_Vh(Q6_V_lo_W(vp_int16_hi1)), v_scale_duplicated), v_offset_duplicated)); + HVX_Vector v_grp1_3 = Q6_Vhf_equals_Vqf16(Q6_Vqf16_vadd_Vqf16Vhf(Q6_Vqf16_vmpy_VhfVhf(Q6_Vhf_equals_Vh(Q6_V_hi_W(vp_int16_hi1)), v_scale_duplicated), v_offset_duplicated)); + + HVX_Vector v_grp2_0 = Q6_Vhf_equals_Vqf16(Q6_Vqf16_vadd_Vqf16Vhf(Q6_Vqf16_vmpy_VhfVhf(Q6_Vhf_equals_Vh(Q6_V_lo_W(vp_int16_lo2)), v_scale_duplicated), v_offset_duplicated)); + HVX_Vector v_grp2_1 = Q6_Vhf_equals_Vqf16(Q6_Vqf16_vadd_Vqf16Vhf(Q6_Vqf16_vmpy_VhfVhf(Q6_Vhf_equals_Vh(Q6_V_hi_W(vp_int16_lo2)), v_scale_duplicated), v_offset_duplicated)); + HVX_Vector v_grp2_2 = Q6_Vhf_equals_Vqf16(Q6_Vqf16_vadd_Vqf16Vhf(Q6_Vqf16_vmpy_VhfVhf(Q6_Vhf_equals_Vh(Q6_V_lo_W(vp_int16_hi2)), v_scale_duplicated), v_offset_duplicated)); + HVX_Vector v_grp2_3 = Q6_Vhf_equals_Vqf16(Q6_Vqf16_vadd_Vqf16Vhf(Q6_Vqf16_vmpy_VhfVhf(Q6_Vhf_equals_Vh(Q6_V_hi_W(vp_int16_hi2)), v_scale_duplicated), v_offset_duplicated)); + + HVX_Vector v_grp3_0 = Q6_Vhf_equals_Vqf16(Q6_Vqf16_vadd_Vqf16Vhf(Q6_Vqf16_vmpy_VhfVhf(Q6_Vhf_equals_Vh(Q6_V_lo_W(vp_int16_lo3)), v_scale_duplicated), v_offset_duplicated)); + HVX_Vector v_grp3_1 = Q6_Vhf_equals_Vqf16(Q6_Vqf16_vadd_Vqf16Vhf(Q6_Vqf16_vmpy_VhfVhf(Q6_Vhf_equals_Vh(Q6_V_hi_W(vp_int16_lo3)), v_scale_duplicated), v_offset_duplicated)); + HVX_Vector v_grp3_2 = Q6_Vhf_equals_Vqf16(Q6_Vqf16_vadd_Vqf16Vhf(Q6_Vqf16_vmpy_VhfVhf(Q6_Vhf_equals_Vh(Q6_V_lo_W(vp_int16_hi3)), v_scale_duplicated), v_offset_duplicated)); + HVX_Vector v_grp3_3 = Q6_Vhf_equals_Vqf16(Q6_Vqf16_vadd_Vqf16Vhf(Q6_Vqf16_vmpy_VhfVhf(Q6_Vhf_equals_Vh(Q6_V_hi_W(vp_int16_hi3)), v_scale_duplicated), v_offset_duplicated)); + + // Parallel Stores + hvx_vmem(dst_ptr + 0 * 64) = v_grp0_0; + hvx_vmem(dst_ptr + 1 * 64) = v_grp0_1; + hvx_vmem(dst_ptr + 2 * 64) = v_grp0_2; + hvx_vmem(dst_ptr + 3 * 64) = v_grp0_3; + + hvx_vmem(dst_ptr + 4 * 64) = v_grp1_0; + hvx_vmem(dst_ptr + 5 * 64) = v_grp1_1; + hvx_vmem(dst_ptr + 6 * 64) = v_grp1_2; + hvx_vmem(dst_ptr + 7 * 64) = v_grp1_3; + + hvx_vmem(dst_ptr + 8 * 64) = v_grp2_0; + hvx_vmem(dst_ptr + 9 * 64) = v_grp2_1; + hvx_vmem(dst_ptr + 10 * 64) = v_grp2_2; + hvx_vmem(dst_ptr + 11 * 64) = v_grp2_3; + + hvx_vmem(dst_ptr + 12 * 64) = v_grp3_0; + hvx_vmem(dst_ptr + 13 * 64) = v_grp3_1; + hvx_vmem(dst_ptr + 14 * 64) = v_grp3_2; + hvx_vmem(dst_ptr + 15 * 64) = v_grp3_3; + } +} + // Q6_K stores 6-bit weights and one fp16 scale per 16 k, see HTP_MM_WEIGHT_TILE_SIZE_Q6_K. // A k-group holds 4 k per row, the HMX tile holds 2, so each group is dealt into two tiles. static void dequantize_tiled_weight_to_fp16_task_q6_k( @@ -914,98 +1019,6 @@ typedef struct { // activations : fp32 -> fp16 -static void transfer_activation_chunk_fp32_to_fp16(__fp16 *restrict vtcm_dst, const float *restrict src, uint32_t n_rows, uint32_t k_block, uint32_t k_stride, uint32_t k_valid) { - const uint32_t n_rows_padded = hex_align_up(n_rows, HTP_MM_HMX_TILE_N_ROWS); - const uint32_t n_rows_tiled = (n_rows / HTP_MM_HMX_TILE_N_ROWS) * HTP_MM_HMX_TILE_N_ROWS; - - uint32_t r = 0; - - #pragma unroll(2) - for (r = 0; r < n_rows_tiled; r += 2) { - uint32_t r0 = r / HTP_MM_HMX_TILE_N_ROWS; // tile row index - uint32_t r1 = r % HTP_MM_HMX_TILE_N_ROWS; // intra-tile row idx - - const float *ptr_in0 = src + (r + 0) * k_stride; - const float *ptr_in1 = src + (r + 1) * k_stride; - - uint32_t c = 0; - for (; c + 32 <= k_valid; c += 32) { - HVX_Vector v0 = *(const HVX_Vector *)(ptr_in0 + c); - HVX_Vector v1 = *(const HVX_Vector *)(ptr_in1 + c); - HVX_Vector v_out = hvx_vec_f32_to_f16_shuff(v0, v1); - - uint32_t c0 = c / HTP_MM_HMX_TILE_N_COLS; // tile column index - uint32_t tile_idx = r0 * (k_block / HTP_MM_HMX_TILE_N_COLS) + c0; - - HVX_Vector *tile = (HVX_Vector *) (vtcm_dst + tile_idx * HTP_MM_HMX_TILE_N_ELMS); - tile[r1 / 2] = v_out; - } - if (c < k_block) { - HVX_Vector v0 = *(const HVX_Vector *)(ptr_in0 + c); - HVX_Vector v1 = *(const HVX_Vector *)(ptr_in1 + c); - - uint32_t rem = k_valid - c; - HVX_VectorPred mask = Q6_Q_vsetq2_R(rem > 0 ? rem * sizeof(float) : 0); - v0 = Q6_V_vmux_QVV(mask, v0, Q6_V_vzero()); - v1 = Q6_V_vmux_QVV(mask, v1, Q6_V_vzero()); - - HVX_Vector v_out = hvx_vec_f32_to_f16_shuff(v0, v1); - - uint32_t c0 = c / HTP_MM_HMX_TILE_N_COLS; // tile column index - uint32_t tile_idx = r0 * (k_block / HTP_MM_HMX_TILE_N_COLS) + c0; - - HVX_Vector *tile = (HVX_Vector *) (vtcm_dst + tile_idx * HTP_MM_HMX_TILE_N_ELMS); - tile[r1 / 2] = v_out; - } - } - - for (; r < n_rows_padded; r += 2) { - uint32_t r0 = r / HTP_MM_HMX_TILE_N_ROWS; // tile row index - uint32_t r1 = r % HTP_MM_HMX_TILE_N_ROWS; // intra-tile row idx - - const bool row0_valid = r < n_rows; - const bool row1_valid = (r + 1) < n_rows; - - const float *ptr_in0 = row0_valid ? (src + (r + 0) * k_stride) : NULL; - const float *ptr_in1 = row1_valid ? (src + (r + 1) * k_stride) : NULL; - - uint32_t c = 0; - for (; c + 32 <= k_valid; c += 32) { - HVX_Vector v0 = Q6_V_vzero(); - HVX_Vector v1 = Q6_V_vzero(); - if (row0_valid) v0 = *(const HVX_Vector *)(ptr_in0 + c); - if (row1_valid) v1 = *(const HVX_Vector *)(ptr_in1 + c); - - HVX_Vector v_out = hvx_vec_f32_to_f16_shuff(v0, v1); - - uint32_t c0 = c / HTP_MM_HMX_TILE_N_COLS; // tile column index - uint32_t tile_idx = r0 * (k_block / HTP_MM_HMX_TILE_N_COLS) + c0; - - HVX_Vector *tile = (HVX_Vector *) (vtcm_dst + tile_idx * HTP_MM_HMX_TILE_N_ELMS); - tile[r1 / 2] = v_out; - } - if (c < k_block) { - HVX_Vector v0 = Q6_V_vzero(); - HVX_Vector v1 = Q6_V_vzero(); - if (row0_valid) v0 = *(const HVX_Vector *)(ptr_in0 + c); - if (row1_valid) v1 = *(const HVX_Vector *)(ptr_in1 + c); - - uint32_t rem = k_valid - c; - HVX_VectorPred mask = Q6_Q_vsetq2_R(rem > 0 ? rem * sizeof(float) : 0); - v0 = Q6_V_vmux_QVV(mask, v0, Q6_V_vzero()); - v1 = Q6_V_vmux_QVV(mask, v1, Q6_V_vzero()); - - HVX_Vector v_out = hvx_vec_f32_to_f16_shuff(v0, v1); - - uint32_t c0 = c / HTP_MM_HMX_TILE_N_COLS; // tile column index - uint32_t tile_idx = r0 * (k_block / HTP_MM_HMX_TILE_N_COLS) + c0; - - HVX_Vector *tile = (HVX_Vector *) (vtcm_dst + tile_idx * HTP_MM_HMX_TILE_N_ELMS); - tile[r1 / 2] = v_out; - } - } -} - static void transfer_activation_row_pair_fp32_to_fp16( __fp16 *restrict vtcm_dst, const float *restrict row0, diff --git a/ggml/src/ggml-hexagon/htp/htp-ops.h b/ggml/src/ggml-hexagon/htp/htp-ops.h index ee5b92441..a03c8550f 100644 --- a/ggml/src/ggml-hexagon/htp/htp-ops.h +++ b/ggml/src/ggml-hexagon/htp/htp-ops.h @@ -23,6 +23,7 @@ enum htp_data_type { HTP_TYPE_Q4_1 = 3, HTP_TYPE_Q8_0 = 8, HTP_TYPE_Q4_K = 12, + HTP_TYPE_Q5_K = 13, HTP_TYPE_Q6_K = 14, HTP_TYPE_IQ4_NL = 20, HTP_TYPE_I32 = 26, diff --git a/ggml/src/ggml-hexagon/htp/hvx-mm-kernels-tiled.h b/ggml/src/ggml-hexagon/htp/hvx-mm-kernels-tiled.h index 4d6110ffa..4564b1524 100644 --- a/ggml/src/ggml-hexagon/htp/hvx-mm-kernels-tiled.h +++ b/ggml/src/ggml-hexagon/htp/hvx-mm-kernels-tiled.h @@ -121,29 +121,43 @@ static inline void quantize_block_f32_q8_0_tiled(float * restrict x, uint8_t * r HVX_Vector * vx = (HVX_Vector *) x; HVX_Vector zero = Q6_V_vzero(); + HVX_Vector vmax0_sf = hvx_vec_reduce_max_f32(hvx_vec_abs_f32(vx[0])); + HVX_Vector vmax1_sf = hvx_vec_reduce_max_f32(hvx_vec_abs_f32(vx[1])); + HVX_Vector vmax2_sf = hvx_vec_reduce_max_f32(hvx_vec_abs_f32(vx[2])); + HVX_Vector vmax3_sf = hvx_vec_reduce_max_f32(hvx_vec_abs_f32(vx[3])); + HVX_Vector vx0_qf = Q6_Vqf32_vsub_VsfVsf(vx[0], zero); HVX_Vector vx1_qf = Q6_Vqf32_vsub_VsfVsf(vx[1], zero); HVX_Vector vx2_qf = Q6_Vqf32_vsub_VsfVsf(vx[2], zero); HVX_Vector vx3_qf = Q6_Vqf32_vsub_VsfVsf(vx[3], zero); + HVX_Vector vmax0_qf = Q6_Vqf32_vsub_VsfVsf(vmax0_sf, zero); + HVX_Vector vmax1_qf = Q6_Vqf32_vsub_VsfVsf(vmax1_sf, zero); + HVX_Vector vmax2_qf = Q6_Vqf32_vsub_VsfVsf(vmax2_sf, zero); + HVX_Vector vmax3_qf = Q6_Vqf32_vsub_VsfVsf(vmax3_sf, zero); + + HVX_Vector vmax01_hf = Q6_Vh_vdeal_Vh(Q6_Vhf_equals_Wqf32(Q6_W_vcombine_VV(vmax1_qf, vmax0_qf))); + HVX_Vector vmax23_hf = Q6_Vh_vdeal_Vh(Q6_Vhf_equals_Wqf32(Q6_W_vcombine_VV(vmax3_qf, vmax2_qf))); + HVX_Vector vx01_hf = Q6_Vh_vdeal_Vh(Q6_Vhf_equals_Wqf32(Q6_W_vcombine_VV(vx1_qf, vx0_qf))); HVX_Vector vx23_hf = Q6_Vh_vdeal_Vh(Q6_Vhf_equals_Wqf32(Q6_W_vcombine_VV(vx3_qf, vx2_qf))); - HVX_Vector vmax_hf = hvx_vec_reduce_max_f16(hvx_vec_abs_f16(vx01_hf)); - vmax_hf = hvx_vec_reduce_max2_f16(hvx_vec_abs_f16(vx23_hf), vmax_hf); + HVX_Vector vd01_qf16 = Q6_Vqf16_vmpy_VhfVhf(vmax01_hf, Q6_Vh_vsplat_R(0x2008)); + HVX_Vector vd23_qf16 = Q6_Vqf16_vmpy_VhfVhf(vmax23_hf, Q6_Vh_vsplat_R(0x2008)); + HVX_Vector vd01_hf = Q6_Vhf_equals_Vqf16(vd01_qf16); + HVX_Vector vd23_hf = Q6_Vhf_equals_Vqf16(vd23_qf16); - HVX_Vector vd_qf16 = Q6_Vqf16_vmpy_VhfVhf(vmax_hf, Q6_Vh_vsplat_R(0x2008)); - HVX_Vector vd_hf = Q6_Vhf_equals_Vqf16(vd_qf16); - - HVX_Vector vd_inv_hf = hvx_vec_inverse_f16(vd_hf); - vx01_hf = Q6_Vhf_equals_Vqf16(Q6_Vqf16_vmpy_VhfVhf(vx01_hf, vd_inv_hf)); - vx23_hf = Q6_Vhf_equals_Vqf16(Q6_Vqf16_vmpy_VhfVhf(vx23_hf, vd_inv_hf)); + HVX_Vector vd01_inv_hf = hvx_vec_inverse_f16(vd01_hf); + HVX_Vector vd23_inv_hf = hvx_vec_inverse_f16(vd23_hf); + vx01_hf = Q6_Vhf_equals_Vqf16(Q6_Vqf16_vmpy_VhfVhf(vx01_hf, vd01_inv_hf)); + vx23_hf = Q6_Vhf_equals_Vqf16(Q6_Vqf16_vmpy_VhfVhf(vx23_hf, vd23_inv_hf)); HVX_Vector vx01_i16 = hvx_vec_i16_from_hf_rnd_sat(vx01_hf); HVX_Vector vx23_i16 = hvx_vec_i16_from_hf_rnd_sat(vx23_hf); HVX_Vector vx_i8 = Q6_Vb_vpack_VhVh_sat(vx23_i16, vx01_i16); - HVX_Vector r_scale = hvx_vec_repl_f16(vd_hf); + HVX_VectorPair vp01 = Q6_W_vshuff_VVR(vd01_hf, vd01_hf, -64); + HVX_VectorPair vp23 = Q6_W_vshuff_VVR(vd23_hf, vd23_hf, -64); static const uint8_t __attribute__((aligned(128))) repl[128] = { 0x00, 0x00, 0x00, 0x00, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, @@ -157,8 +171,19 @@ static inline void quantize_block_f32_q8_0_tiled(float * restrict x, uint8_t * r }; HVX_Vector v_repl_ctrl = * (const HVX_Vector *) repl; + #pragma unroll for (int b = 0; b < 4; b++) { HVX_Vector v_act = Q6_V_vror_VR(vx_i8, b * 32); + HVX_Vector r_scale; + if (b == 0) { + r_scale = Q6_V_lo_W(vp01); + } else if (b == 1) { + r_scale = Q6_V_hi_W(vp01); + } else if (b == 2) { + r_scale = Q6_V_lo_W(vp23); + } else { + r_scale = Q6_V_hi_W(vp23); + } HVX_Vector r0 = Q6_V_vdelta_VV(v_act, v_repl_ctrl); HVX_Vector r1 = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act, 4), v_repl_ctrl); @@ -371,6 +396,78 @@ static inline HVX_VectorPair accum_q8_0_32x2( return Q6_W_vcombine_VV(v_sum1, v_sum0); } +// Q5_K: OR 0x10 into every lane of v whose flag j is set in the plane (see HTP_MM_WEIGHT_TILE_SIZE_Q5_K) +static inline HVX_Vector hvx_q5k_or_hibit(HVX_Vector v, HVX_Vector v_plane, int j) { + HVX_VectorPred q = Q6_Q_vand_VR(v_plane, 0x01010101u << j); + return Q6_V_vandor_VQR(v, q, 0x10101010); +} + +// 5-bit variant: the high bit comes from the plane, see hvx_q5k_or_hibit +static inline HVX_VectorPair unpack_and_interleave_5bit_x2(HVX_Vector v_src, HVX_Vector v_plane, int i, HVX_Vector mask_h4) { + HVX_Vector v_lo = hvx_q5k_or_hibit(Q6_V_vand_VV(v_src, mask_h4), v_plane, 2 * i); + HVX_Vector v_hi = hvx_q5k_or_hibit(Q6_Vub_vlsr_VubR(v_src, 4), v_plane, 2 * i + 1); + HVX_VectorPair v01_pair = Q6_W_vshuff_VVR(v_hi, v_lo, -1); + HVX_Vector v01_lo = Q6_V_lo_W(v01_pair); + HVX_Vector v01_hi = Q6_V_hi_W(v01_pair); + + HVX_Vector v23_lo = Q6_V_valign_VVR(v01_hi, v01_lo, 64); + HVX_Vector v_W0 = Q6_V_lo_W(Q6_W_vshuff_VVR(v23_lo, v01_lo, -2)); + + HVX_Vector v67_lo = Q6_V_valign_VVR(v01_lo, v01_hi, 64); + HVX_Vector v_W1 = Q6_V_lo_W(Q6_W_vshuff_VVR(v67_lo, v01_hi, -2)); + + return Q6_W_vcombine_VV(v_W1, v_W0); +} + +static inline HVX_Vector accum_5bit_32x1( + const HVX_Vector * restrict vptr, + const HVX_Vector * restrict v_act, + HVX_Vector i8 +) { + HVX_Vector v_sum0 = Q6_V_vzero(); + HVX_Vector v_sum1 = Q6_V_vzero(); + HVX_Vector mask_h4 = Q6_Vb_vsplat_R(0x0F); + HVX_Vector v_plane = vptr[5]; + + #pragma unroll + for (int i = 0; i < 4; i++) { + HVX_VectorPair v_W_pair = unpack_and_interleave_5bit_x2(vptr[i], v_plane, i, mask_h4); + HVX_Vector v_W0 = Q6_Vb_vsub_VbVb(Q6_V_lo_W(v_W_pair), i8); + HVX_Vector v_W1 = Q6_Vb_vsub_VbVb(Q6_V_hi_W(v_W_pair), i8); + v_sum0 = Q6_Vw_vrmpyacc_VwVbVb(v_sum0, v_W0, v_act[i * 2 + 0]); + v_sum1 = Q6_Vw_vrmpyacc_VwVbVb(v_sum1, v_W1, v_act[i * 2 + 1]); + } + + return Q6_Vw_vadd_VwVw(v_sum0, v_sum1); +} + +static inline HVX_VectorPair accum_5bit_32x2( + const HVX_Vector * restrict vptr, + const HVX_Vector * restrict v_act0, + const HVX_Vector * restrict v_act1, + HVX_Vector i8 +) { + HVX_Vector v_sum0 = Q6_V_vzero(); + HVX_Vector v_sum1 = Q6_V_vzero(); + HVX_Vector mask_h4 = Q6_Vb_vsplat_R(0x0F); + HVX_Vector v_plane = vptr[5]; + + #pragma unroll + for (int i = 0; i < 4; i++) { + HVX_VectorPair v_W_pair = unpack_and_interleave_5bit_x2(vptr[i], v_plane, i, mask_h4); + HVX_Vector v_W0 = Q6_Vb_vsub_VbVb(Q6_V_lo_W(v_W_pair), i8); + HVX_Vector v_W1 = Q6_Vb_vsub_VbVb(Q6_V_hi_W(v_W_pair), i8); + + v_sum0 = Q6_Vw_vrmpyacc_VwVbVb(v_sum0, v_W0, v_act0[i * 2 + 0]); + v_sum0 = Q6_Vw_vrmpyacc_VwVbVb(v_sum0, v_W1, v_act0[i * 2 + 1]); + + v_sum1 = Q6_Vw_vrmpyacc_VwVbVb(v_sum1, v_W0, v_act1[i * 2 + 0]); + v_sum1 = Q6_Vw_vrmpyacc_VwVbVb(v_sum1, v_W1, v_act1[i * 2 + 1]); + } + + return Q6_W_vcombine_VV(v_sum1, v_sum0); +} + // Q6_K weights are stored unsigned (0..63), see HTP_MM_WEIGHT_TILE_SIZE_Q6_K. Unpack k-group g of a tile to signed bytes (q - 32) static inline HVX_Vector unpack_q6_k_group(const HVX_Vector * restrict vptr, int g, HVX_Vector mask_0f, HVX_Vector mask_03, HVX_Vector i32) { HVX_Vector v_lo = (g & 1) ? Q6_Vub_vlsr_VubR(vptr[g >> 1], 4) : Q6_V_vand_VV(vptr[g >> 1], mask_0f); @@ -689,6 +786,102 @@ static void tiled_vec_dot_q8_0_32x2(const uint32_t n, float * restrict s0, float } } +static void tiled_vec_dot_q5_k_32x1(const uint32_t n, float * restrict s, const void * restrict vx, const void * restrict vy, uint32_t valid_rows, const float * restrict sz) { + const uint8_t * restrict tile_ptr = vx; + const uint8_t * restrict y_q = vy; + + HVX_Vector v_sum_float = Q6_V_vzero(); + + uint32_t n_k_tiles = n / 32; + for (uint32_t kt = 0; kt < n_k_tiles; kt++) { + const HVX_Vector * restrict vptr = (const HVX_Vector *) (tile_ptr + kt * 768); + const HVX_Vector * restrict v_act = (const HVX_Vector *) (y_q + kt * 1280); + + HVX_Vector v_sum = accum_5bit_32x1(vptr, v_act, Q6_V_vzero()); + HVX_Vector v_sum_sf = Q6_Vsf_equals_Vw(v_sum); + + HVX_Vector v_scale_offset = vptr[4]; + HVX_VectorPair p_deal = Q6_W_vdeal_VVR(v_scale_offset, v_scale_offset, -2); + HVX_Vector v_scale = Q6_V_lo_W(p_deal); + HVX_Vector v_offset = Q6_V_hi_W(p_deal); + + HVX_Vector v_scale_a = v_act[8]; + HVX_Vector v_sum_a = v_act[9]; + + HVX_Vector v_scale_comb = hvx_vec_mul_f16_f16_to_f32_lower32(v_scale, v_scale_a); + HVX_Vector v_offset_comb = hvx_vec_mul_f16_f16_to_f32_lower32(v_offset, v_sum_a); + + HVX_Vector v_scaled_dot = hvx_vec_mul_f32_f32(v_sum_sf, v_scale_comb); + HVX_Vector v_sum_scaled = hvx_vec_add_f32_f32(v_scaled_dot, v_offset_comb); + + v_sum_float = hvx_vec_add_f32_f32(v_sum_float, v_sum_scaled); + } + + if (sz) { + hvx_vec_store_u(s, valid_rows * sizeof(float), hvx_vec_add_f32_f32(v_sum_float, hvx_vmemu(sz))); + } else { + hvx_vec_store_u(s, valid_rows * sizeof(float), v_sum_float); + } +} + +static void tiled_vec_dot_q5_k_32x2(const uint32_t n, float * restrict s0, float * restrict s1, const void * restrict vx, const void * restrict vy0, const void * restrict vy1, uint32_t valid_rows, const float * restrict sz0, const float * restrict sz1) { + const uint8_t * restrict tile_ptr = vx; + const uint8_t * restrict y0_q = vy0; + const uint8_t * restrict y1_q = vy1; + + HVX_Vector v_sum_float_c0 = Q6_V_vzero(); + HVX_Vector v_sum_float_c1 = Q6_V_vzero(); + + uint32_t n_k_tiles = n / 32; + for (uint32_t kt = 0; kt < n_k_tiles; kt++) { + const HVX_Vector * restrict vptr = (const HVX_Vector *) (tile_ptr + kt * 768); + const HVX_Vector * restrict v_act0 = (const HVX_Vector *) (y0_q + kt * 1280); + const HVX_Vector * restrict v_act1 = (const HVX_Vector *) (y1_q + kt * 1280); + + HVX_VectorPair v_sums = accum_5bit_32x2(vptr, v_act0, v_act1, Q6_V_vzero()); + HVX_Vector v_sum_c0 = Q6_V_lo_W(v_sums); + HVX_Vector v_sum_c1 = Q6_V_hi_W(v_sums); + + HVX_Vector v_sum_sf_c0 = Q6_Vsf_equals_Vw(v_sum_c0); + HVX_Vector v_sum_sf_c1 = Q6_Vsf_equals_Vw(v_sum_c1); + + HVX_Vector v_scale_offset = vptr[4]; + HVX_VectorPair p_deal = Q6_W_vdeal_VVR(v_scale_offset, v_scale_offset, -2); + HVX_Vector v_scale = Q6_V_lo_W(p_deal); + HVX_Vector v_offset = Q6_V_hi_W(p_deal); + + HVX_Vector v_scale_a_c0 = v_act0[8]; + HVX_Vector v_sum_a_c0 = v_act0[9]; + HVX_Vector v_scale_a_c1 = v_act1[8]; + HVX_Vector v_sum_a_c1 = v_act1[9]; + + HVX_Vector v_scale_comb_c0 = hvx_vec_mul_f16_f16_to_f32_lower32(v_scale, v_scale_a_c0); + HVX_Vector v_offset_comb_c0 = hvx_vec_mul_f16_f16_to_f32_lower32(v_offset, v_sum_a_c0); + HVX_Vector v_scale_comb_c1 = hvx_vec_mul_f16_f16_to_f32_lower32(v_scale, v_scale_a_c1); + HVX_Vector v_offset_comb_c1 = hvx_vec_mul_f16_f16_to_f32_lower32(v_offset, v_sum_a_c1); + + HVX_Vector v_scaled_dot_c0 = hvx_vec_mul_f32_f32(v_sum_sf_c0, v_scale_comb_c0); + HVX_Vector v_sum_scaled_c0 = hvx_vec_add_f32_f32(v_scaled_dot_c0, v_offset_comb_c0); + + HVX_Vector v_scaled_dot_c1 = hvx_vec_mul_f32_f32(v_sum_sf_c1, v_scale_comb_c1); + HVX_Vector v_sum_scaled_c1 = hvx_vec_add_f32_f32(v_scaled_dot_c1, v_offset_comb_c1); + + v_sum_float_c0 = hvx_vec_add_f32_f32(v_sum_float_c0, v_sum_scaled_c0); + v_sum_float_c1 = hvx_vec_add_f32_f32(v_sum_float_c1, v_sum_scaled_c1); + } + + if (sz0) { + hvx_vec_store_u(s0, valid_rows * sizeof(float), hvx_vec_add_f32_f32(v_sum_float_c0, hvx_vmemu(sz0))); + } else { + hvx_vec_store_u(s0, valid_rows * sizeof(float), v_sum_float_c0); + } + if (sz1) { + hvx_vec_store_u(s1, valid_rows * sizeof(float), hvx_vec_add_f32_f32(v_sum_float_c1, hvx_vmemu(sz1))); + } else { + hvx_vec_store_u(s1, valid_rows * sizeof(float), v_sum_float_c1); + } +} + static void tiled_vec_dot_q6_k_32x1(const uint32_t n, float * restrict s, const void * restrict vx, const void * restrict vy, uint32_t valid_rows, const float * restrict sz) { const uint8_t * restrict tile_ptr = vx; const uint8_t * restrict y_q = vy; @@ -941,14 +1134,9 @@ static inline void quantize_f32_q8_0_tiled_kernel( size_t src_row_size, size_t dst_row_size ) { - const size_t src_row_size_padded = hex_round_up(src_row_size, QK_Q8_0_TILED * sizeof(float)); - hvx_splat_f32_a(tmp_data, 0.0f, src_row_size_padded / sizeof(float)); - + (void) tmp_data; for (uint32_t i = 0; i < nrows; ++i) { - hex_l2fetch(src_data, src_row_size, src_row_size, 2); - hvx_copy_f32_aa(tmp_data, src_data, ne0); - - quantize_row_f32_q8_0_tiled((float *) tmp_data, dst_data, ne0); + quantize_row_f32_q8_0_tiled((float *) src_data, dst_data, ne0); dst_data += dst_row_size; src_data += src_row_size; } @@ -963,14 +1151,9 @@ static inline void quantize_f32_q8_1_tiled_kernel( size_t src_row_size, size_t dst_row_size ) { - const size_t src_row_size_padded = hex_round_up(src_row_size, QK_Q8_0_TILED * sizeof(float)); - hvx_splat_f32_a(tmp_data, 0.0f, src_row_size_padded / sizeof(float)); - + (void) tmp_data; for (uint32_t i = 0; i < nrows; ++i) { - hex_l2fetch(src_data, src_row_size, src_row_size, 2); - hvx_copy_f32_aa(tmp_data, src_data, ne0); - - quantize_row_f32_q8_1_tiled((float *) tmp_data, dst_data, ne0); + quantize_row_f32_q8_1_tiled((float *) src_data, dst_data, ne0); dst_data += dst_row_size; src_data += src_row_size; } @@ -988,24 +1171,15 @@ static inline void quantize_f32_q8_0_tiled_block_kernel( uint32_t r, uint32_t c ) { + (void) tmp_data; const uint32_t qk = QK_Q8_0_TILED; const uint32_t nb = (ne0 + qk - 1) / qk; for (uint32_t ib = ib_first; ib < ib_last; ++ib) { - const uint8_t * restrict src_ptr = (const uint8_t *) src + r * src_row_size + c * qk * sizeof(float); + const float * restrict src_ptr = (const float *) ((const uint8_t *) src + r * src_row_size + c * qk * sizeof(float)); uint8_t * restrict dst_ptr = dst + r * dst_row_size + c * 4 * 1152; - hex_l2fetch(src_ptr, qk * sizeof(float), qk * sizeof(float), 1); - - if (c == nb - 1) { - uint32_t active_elements = ne0 - c * qk; - hvx_splat_f32_a(tmp_data, 0.0f, qk); - hvx_copy_f32_aa(tmp_data, src_ptr, active_elements); - } else { - hvx_copy_f32_aa(tmp_data, src_ptr, qk); - } - - quantize_block_f32_q8_0_tiled((float *) tmp_data, dst_ptr); + quantize_block_f32_q8_0_tiled((float *) src_ptr, dst_ptr); c++; if (c == nb) { @@ -1027,24 +1201,15 @@ static inline void quantize_f32_q8_1_tiled_block_kernel( uint32_t r, uint32_t c ) { + (void) tmp_data; const uint32_t qk = QK_Q8_0_TILED; const uint32_t nb = (ne0 + qk - 1) / qk; for (uint32_t ib = ib_first; ib < ib_last; ++ib) { - const uint8_t * restrict src_ptr = (const uint8_t *) src + r * src_row_size + c * qk * sizeof(float); + const float * restrict src_ptr = (const float *) ((const uint8_t *) src + r * src_row_size + c * qk * sizeof(float)); uint8_t * restrict dst_ptr = dst + r * dst_row_size + c * 4 * 1280; - hex_l2fetch(src_ptr, qk * sizeof(float), qk * sizeof(float), 1); - - if (c == nb - 1) { - uint32_t active_elements = ne0 - c * qk; - hvx_splat_f32_a(tmp_data, 0.0f, qk); - hvx_copy_f32_aa(tmp_data, src_ptr, active_elements); - } else { - hvx_copy_f32_aa(tmp_data, src_ptr, qk); - } - - quantize_block_f32_q8_1_tiled((float *) tmp_data, dst_ptr); + quantize_block_f32_q8_1_tiled((float *) src_ptr, dst_ptr); c++; if (c == nb) { diff --git a/ggml/src/ggml-hexagon/htp/matmul-ops.c b/ggml/src/ggml-hexagon/htp/matmul-ops.c index a09bc7a28..727dc281f 100644 --- a/ggml/src/ggml-hexagon/htp/matmul-ops.c +++ b/ggml/src/ggml-hexagon/htp/matmul-ops.c @@ -29,7 +29,7 @@ typedef struct { float *dst; dma_addr_t src2_addr; size_t src2_bytes; - const float *activation; + dma_addr_t act_dma_addr; dma_addr_t weight; dma_queue * weight_dma; int m; @@ -45,8 +45,8 @@ typedef struct { int ne13; size_t src0_nb2; size_t src0_nb3; - size_t src1_nb2; - size_t src1_nb3; + size_t act_nb2; + size_t act_nb3; size_t src2_nb2; size_t src2_nb3; size_t dst_nb2; @@ -93,7 +93,7 @@ struct htp_mm_context { uint32_t src0_row_start; uint32_t src0_row_end; uint32_t src0_row_size_padded; - uint32_t src1_nrows; + uint32_t act_nrows; uint32_t cur_m_start; uint32_t cur_m_rows; @@ -103,16 +103,13 @@ struct htp_mm_context { struct fastdiv_values mm_div_r3; struct fastdiv_values mm_div_ne11; - // Per thread quant tasks // Precomputed block-parallel quantization values - worker_callback_t quant_task_func; uint32_t quant_ib_first[WORK_QUEUE_MAX_N_THREADS]; uint32_t quant_ib_last[WORK_QUEUE_MAX_N_THREADS]; uint32_t quant_r[WORK_QUEUE_MAX_N_THREADS]; uint32_t quant_c[WORK_QUEUE_MAX_N_THREADS]; uint32_t n_quant_tasks; uint32_t n_quant_rows_per_thread; - atomic_uint quant_barrier; // Fields for scattered mapping & HMX support in MUL_MAT_ID const uint32_t * matrix_row_counts; @@ -125,12 +122,14 @@ struct htp_mm_context { uint8_t * vtcm_src2; uint8_t * vtcm_src3; uint8_t * vtcm_dst; + uint8_t * vtcm_act_raw; // Cached strides uint32_t vtcm_src0_stride; uint32_t vtcm_src1_stride; uint32_t vtcm_src2_stride; uint32_t vtcm_src3_stride; + uint32_t vtcm_act_raw_stride; // Cached thread offsets/sizes uint32_t vtcm_src0_size_per_thread; @@ -234,17 +233,6 @@ static const uint8_t __attribute__((aligned(VLEN))) kvalues_mxfp4_lut[] = { uint32_t src0_nrows_per_thread = mmctx->src0_nrows_per_thread; \ htp_matmul_tensors_preamble; -static inline void hvx_mm_run_quant_task(struct htp_mm_context * mmctx, unsigned int ith) { - if (mmctx->quant_task_func) { - if (ith < mmctx->n_quant_tasks) { - mmctx->quant_task_func(mmctx->n_quant_tasks, ith, mmctx); - atomic_fetch_sub(&mmctx->quant_barrier, 1); - } - while (atomic_load(&mmctx->quant_barrier) > 0) { - // spin - } - } -} @@ -301,8 +289,6 @@ static void hvx_mm_2d_repacked_##SUFFIX(unsigned int nth, unsigned int ith, void } \ } \ \ - hvx_mm_run_quant_task(mmctx, ith); \ - \ if (src0_start_row >= src0_end_row) { \ return; \ } \ @@ -416,8 +402,6 @@ static void hvx_mv_2d_repacked_##SUFFIX(unsigned int nth, unsigned int ith, void } \ } \ \ - hvx_mm_run_quant_task(mmctx, ith); \ - \ if (src0_start_row >= src0_end_row) { \ return; \ } \ @@ -479,8 +463,6 @@ static void hvx_mm_nx_2d_repacked_##SUFFIX(unsigned int nth, unsigned int ith, v uint32_t n_k_tiles_a = ne10 / 32; \ uint32_t tile_row_transfer_size_aligned = n_k_tiles_a * aligned_tile_size; \ \ - hvx_mm_run_quant_task(mmctx, ith); \ - \ for (uint32_t widx = 0; widx < n_weights; widx++) { \ const struct htp_tensor * restrict src_w = octx->src[widx]; \ const struct htp_tensor * restrict dst = octx->dsts[widx]; \ @@ -562,63 +544,124 @@ MATMUL_2D_REPACKED_IMPL(q4_0, 576, tiled_vec_dot_q4_0_32x2, tiled_vec_do MATMUL_2D_REPACKED_IMPL(q4_1, 640, tiled_vec_dot_q4_1_32x2, tiled_vec_dot_q4_1_32x1) MATMUL_2D_REPACKED_IMPL(q8_0, 1088, tiled_vec_dot_q8_0_32x2, tiled_vec_dot_q8_0_32x1) MATMUL_2D_REPACKED_IMPL(q6_k, 896, tiled_vec_dot_q6_k_32x2, tiled_vec_dot_q6_k_32x1) +MATMUL_2D_REPACKED_IMPL(q5_k, 768, tiled_vec_dot_q5_k_32x2, tiled_vec_dot_q5_k_32x1) MATMUL_2D_REPACKED_IMPL(iq4nl, 576, tiled_vec_dot_iq4nl_32x2, tiled_vec_dot_iq4nl_32x1) MATMUL_2D_REPACKED_IMPL(mxfp4, 544, tiled_vec_dot_mxfp4_32x2, tiled_vec_dot_mxfp4_32x1) -#define QUANTIZE_IMPL(name, log_name, kernel_fn, dst_row_size_expr) \ -static void name(unsigned int nth, unsigned int ith, void * data) { \ - struct htp_mm_context * mmctx = data; \ - struct htp_ops_context * octx = mmctx->octx; \ - const struct htp_mm_kernel_params * kparams = (const struct htp_mm_kernel_params *) octx->kernel_params; \ - const struct htp_tensor * src = mmctx->act; \ - const uint32_t ne0 = src->ne[0]; \ - const uint32_t nrows = mmctx->cur_m_rows ? mmctx->cur_m_rows : mmctx->src1_nrows; \ - const uint32_t nrows_per_thread = mmctx->n_quant_rows_per_thread; \ - \ - const uint32_t ir_first = nrows_per_thread * ith; \ - if (ir_first >= nrows) { \ - return; \ - } \ - \ - struct htp_thread_trace * tr = &octx->ctx->trace[ith]; \ - htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_A_QUANT, ir_first); \ - \ - uint8_t * restrict dst = mmctx->vtcm_src1; \ - const uint32_t ir_last = MIN(ir_first + nrows_per_thread, nrows); \ - const size_t src_row_size = src->nb[1]; \ - const size_t dst_row_size = (dst_row_size_expr); \ - uint8_t * restrict tmp_data = (uint8_t *) mmctx->vtcm_dst + (mmctx->vtcm_dst_size_per_thread * ith); \ - \ - const bool is_contiguous = (src->nb[2] == src->ne[1] * src->nb[1]) && (src->nb[3] == src->ne[2] * src->nb[2]); \ - if (is_contiguous) { \ - const uint8_t * restrict src_data = (const uint8_t *) src->data + (src_row_size * (mmctx->cur_m_start + ir_first)); \ - uint8_t * restrict dst_data = (uint8_t *) dst + (dst_row_size * ir_first); \ - kernel_fn(src_data, dst_data, tmp_data, ne0, ir_last - ir_first, src_row_size, dst_row_size); \ - } else { \ - const uint32_t ne12_ne1 = src->ne[2] * src->ne[1]; \ - for (uint32_t ir = ir_first; ir < ir_last; ++ir) { \ - const uint32_t ir1 = mmctx->cur_m_start + ir; \ - const uint32_t i13 = fastdiv(ir1, &kparams->div_ne12_ne1); \ - const uint32_t rem = ir1 - i13 * ne12_ne1; \ - const uint32_t i12 = fastdiv(rem, &kparams->div_ne1); \ - const uint32_t i11 = rem - i12 * src->ne[1]; \ - const uint8_t * restrict row_src = (const uint8_t *) src->data + ((size_t) i11 * src->nb[1] + (size_t) i12 * src->nb[2] + (size_t) i13 * src->nb[3]); \ - uint8_t * restrict row_dst = dst + (dst_row_size * ir); \ - kernel_fn(row_src, row_dst, tmp_data, ne0, 1, src_row_size, dst_row_size); \ - } \ - } \ - \ - htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_A_QUANT, ir_first); \ +static void hvx_mm_transfer_src1_dma( + struct htp_ops_context * octx, + const struct htp_mm_kernel_params * kparams, + const struct htp_tensor * src1, + uint8_t * dst_base, + size_t dst_row_size, + uint32_t m_start, + uint32_t m_rows +) { + if (m_rows == 0) { + return; + } + + dma_queue * dma_q = octx->ctx->dma[0]; + const uint32_t ne0 = src1->ne[0]; + const size_t elem_size = (src1->type == HTP_TYPE_F16) ? sizeof(__fp16) : sizeof(float); + const size_t row_bytes = ne0 * elem_size; + const size_t src1_nb1 = src1->nb[1]; + const dma_addr_t src_base = src1->data; + + const bool is_contiguous = (src1->nb[2] == src1->ne[1] * src1_nb1) && + (src1->nb[3] == src1->ne[2] * src1->nb[2]); + + if (is_contiguous) { + const dma_addr_t src_addr = src_base + m_start * src1_nb1; + dma_queue_push(dma_q, dma_make_data(dst_base, src_addr), + dst_row_size, src1_nb1, row_bytes, m_rows); + dma_queue_pop(dma_q); + } else { + const uint32_t ne12_ne1 = src1->ne[2] * src1->ne[1]; + const bool use_fastdiv = kparams->div_ne12_ne1.mp != 0; + for (uint32_t ir = 0; ir < m_rows; ++ir) { + const uint32_t ir1 = m_start + ir; + uint32_t i11, i12, i13; + if (use_fastdiv) { + i13 = fastdiv(ir1, &kparams->div_ne12_ne1); + const uint32_t rem = ir1 - i13 * ne12_ne1; + i12 = fastdiv(rem, &kparams->div_ne1); + i11 = rem - i12 * src1->ne[1]; + } else { + i13 = ne12_ne1 ? ir1 / ne12_ne1 : 0; + const uint32_t rem = ir1 - i13 * ne12_ne1; + i12 = src1->ne[1] ? rem / src1->ne[1] : 0; + i11 = rem - i12 * src1->ne[1]; + } + const dma_addr_t row_src = src_base + (i11 * src1->nb[1] + + i12 * src1->nb[2] + + i13 * src1->nb[3]); + uint8_t * row_dst = dst_base + ir * dst_row_size; + dma_queue_push(dma_q, dma_make_data(row_dst, row_src), + dst_row_size, src1_nb1, row_bytes, 1); + dma_queue_pop(dma_q); + } + } + + if (dst_row_size > row_bytes) { + struct htp_thread_trace * tr = &octx->ctx->trace[0]; + htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_A_PREP, (uint16_t) m_start); + const uint32_t pad_elems = (dst_row_size - row_bytes) / elem_size; + if (elem_size == sizeof(float)) { + for (uint32_t ir = 0; ir < m_rows; ++ir) { + hvx_splat_f32_u(dst_base + ir * dst_row_size + row_bytes, 0.0f, pad_elems); + } + } else { + for (uint32_t ir = 0; ir < m_rows; ++ir) { + hvx_splat_f16_u(dst_base + ir * dst_row_size + row_bytes, (_Float16) 0.0f, pad_elems); + } + } + htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_A_PREP, (uint16_t) m_start); + } +} + +#define QUANTIZE_IMPL(name, log_name, kernel_fn, dst_row_size_expr) \ +static void name(unsigned int nth, unsigned int ith, void * data) { \ + (void) nth; \ + struct htp_mm_context * mmctx = data; \ + struct htp_ops_context * octx = mmctx->octx; \ + const struct htp_tensor * src = mmctx->act; \ + const uint32_t ne0 = src->ne[0]; \ + const uint32_t nrows = mmctx->cur_m_rows ? mmctx->cur_m_rows : mmctx->act_nrows; \ + const uint32_t nrows_per_thread = mmctx->n_quant_rows_per_thread; \ + \ + const uint32_t ir_first = nrows_per_thread * ith; \ + if (ir_first >= nrows) { \ + return; \ + } \ + \ + struct htp_thread_trace * tr = &octx->ctx->trace[ith]; \ + htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_A_QUANT, ir_first); \ + \ + uint8_t * restrict dst = mmctx->vtcm_src1; \ + const uint32_t ir_last = MIN(ir_first + nrows_per_thread, nrows); \ + const size_t raw_row_size = mmctx->vtcm_act_raw_stride; \ + const size_t dst_row_size = (dst_row_size_expr); \ + \ + const uint8_t * restrict src_data = (const uint8_t *) mmctx->vtcm_act_raw + (raw_row_size * ir_first); \ + uint8_t * restrict dst_data = (uint8_t *) dst + (dst_row_size * ir_first); \ + kernel_fn(src_data, dst_data, NULL, ne0, ir_last - ir_first, raw_row_size, dst_row_size); \ + \ + htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_A_QUANT, ir_first); \ } QUANTIZE_IMPL(quantize_f32_q8_0_tiled, "quantize-f32-q8_0_tiled", quantize_f32_q8_0_tiled_kernel, htp_mm_q8_0_tiled_row_size(ne0)) QUANTIZE_IMPL(quantize_f32_q8_1_tiled, "quantize-f32-q8_1_tiled", quantize_f32_q8_1_tiled_kernel, htp_mm_q8_1_tiled_row_size(ne0)) -QUANTIZE_IMPL(quantize_f32_f32, "quantize-f32-f32", quantize_f32_f32_kernel, mmctx->vtcm_src1_stride) -QUANTIZE_IMPL(quantize_f32_f16, "quantize-f32-f16", quantize_f32_f16_kernel, mmctx->vtcm_src1_stride) -QUANTIZE_IMPL(quantize_f16_f16, "quantize-f16-f16", quantize_f16_f16_kernel, mmctx->vtcm_src1_stride) +QUANTIZE_IMPL(quantize_f32_f32, "quantize-f32-f32", quantize_f32_f32_kernel, mmctx->vtcm_src1_stride) +QUANTIZE_IMPL(quantize_f32_f16, "quantize-f32-f16", quantize_f32_f16_kernel, mmctx->vtcm_src1_stride) +QUANTIZE_IMPL(quantize_f16_f16, "quantize-f16-f16", quantize_f16_f16_kernel, mmctx->vtcm_src1_stride) static void quantize_f32_q8_0_tiled_block(unsigned int nth, unsigned int ith, void * data) { + (void) nth; struct htp_mm_context * mmctx = data; + if (mmctx->quant_ib_first[ith] >= mmctx->quant_ib_last[ith]) { + return; + } struct htp_ops_context * octx = mmctx->octx; struct htp_thread_trace * tr = &octx->ctx->trace[ith]; htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_A_QUANT, mmctx->quant_ib_first[ith]); @@ -626,13 +669,13 @@ static void quantize_f32_q8_0_tiled_block(unsigned int nth, unsigned int ith, vo const struct htp_tensor * src = mmctx->act; quantize_f32_q8_0_tiled_block_kernel( - (const float *) src->data, + (const float *) mmctx->vtcm_act_raw, mmctx->vtcm_src1, - (uint8_t *) mmctx->vtcm_dst + (mmctx->vtcm_dst_size_per_thread * ith), + NULL, src->ne[0], mmctx->quant_ib_first[ith], mmctx->quant_ib_last[ith], - src->nb[1], + mmctx->vtcm_act_raw_stride, htp_mm_q8_0_tiled_row_size(src->ne[0]), mmctx->quant_r[ith], mmctx->quant_c[ith] @@ -642,7 +685,11 @@ static void quantize_f32_q8_0_tiled_block(unsigned int nth, unsigned int ith, vo } static void quantize_f32_q8_1_tiled_block(unsigned int nth, unsigned int ith, void * data) { + (void) nth; struct htp_mm_context * mmctx = data; + if (mmctx->quant_ib_first[ith] >= mmctx->quant_ib_last[ith]) { + return; + } struct htp_ops_context * octx = mmctx->octx; struct htp_thread_trace * tr = &octx->ctx->trace[ith]; htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_A_QUANT, mmctx->quant_ib_first[ith]); @@ -650,13 +697,13 @@ static void quantize_f32_q8_1_tiled_block(unsigned int nth, unsigned int ith, vo const struct htp_tensor * src = mmctx->act; quantize_f32_q8_1_tiled_block_kernel( - (const float *) src->data, + (const float *) mmctx->vtcm_act_raw, mmctx->vtcm_src1, - (uint8_t *) mmctx->vtcm_dst + (mmctx->vtcm_dst_size_per_thread * ith), + NULL, src->ne[0], mmctx->quant_ib_first[ith], mmctx->quant_ib_last[ith], - src->nb[1], + mmctx->vtcm_act_raw_stride, htp_mm_q8_1_tiled_row_size(src->ne[0]), mmctx->quant_r[ith], mmctx->quant_c[ith] @@ -668,15 +715,17 @@ static void quantize_f32_q8_1_tiled_block(unsigned int nth, unsigned int ith, vo MATVEC_2D_REPACKED_IMPL(q4_0, 576, tiled_vec_dot_q4_0_32x1) MATVEC_2D_REPACKED_IMPL(q4_1, 640, tiled_vec_dot_q4_1_32x1) MATVEC_2D_REPACKED_IMPL(q8_0, 1088, tiled_vec_dot_q8_0_32x1) +MATVEC_2D_REPACKED_IMPL(q5_k, 768, tiled_vec_dot_q5_k_32x1) MATVEC_2D_REPACKED_IMPL(q6_k, 896, tiled_vec_dot_q6_k_32x1) MATVEC_2D_REPACKED_IMPL(iq4nl, 576, tiled_vec_dot_iq4nl_32x1) MATVEC_2D_REPACKED_IMPL(mxfp4, 544, tiled_vec_dot_mxfp4_32x1) -MATMUL_NX_2D_REPACKED_IMPL(q4_0, 576, tiled_vec_dot_q4_0_32x2, tiled_vec_dot_q4_0_32x1) -MATMUL_NX_2D_REPACKED_IMPL(q4_1, 640, tiled_vec_dot_q4_1_32x2, tiled_vec_dot_q4_1_32x1) -MATMUL_NX_2D_REPACKED_IMPL(q8_0, 1088, tiled_vec_dot_q8_0_32x2, tiled_vec_dot_q8_0_32x1) -MATMUL_NX_2D_REPACKED_IMPL(iq4nl, 576, tiled_vec_dot_iq4nl_32x2, tiled_vec_dot_iq4nl_32x1) -MATMUL_NX_2D_REPACKED_IMPL(mxfp4, 544, tiled_vec_dot_mxfp4_32x2, tiled_vec_dot_mxfp4_32x1) +MATMUL_NX_2D_REPACKED_IMPL(q4_0, 576, tiled_vec_dot_q4_0_32x2, tiled_vec_dot_q4_0_32x1) +MATMUL_NX_2D_REPACKED_IMPL(q4_1, 640, tiled_vec_dot_q4_1_32x2, tiled_vec_dot_q4_1_32x1) +MATMUL_NX_2D_REPACKED_IMPL(q8_0, 1088, tiled_vec_dot_q8_0_32x2, tiled_vec_dot_q8_0_32x1) +MATMUL_NX_2D_REPACKED_IMPL(iq4nl, 576, tiled_vec_dot_iq4nl_32x2, tiled_vec_dot_iq4nl_32x1) +MATMUL_NX_2D_REPACKED_IMPL(mxfp4, 544, tiled_vec_dot_mxfp4_32x2, tiled_vec_dot_mxfp4_32x1) +MATMUL_NX_2D_REPACKED_IMPL(q5_k, 768, tiled_vec_dot_q5_k_32x2, tiled_vec_dot_q5_k_32x1) #define MATMUL_4D_REPACKED_IMPL(SUFFIX, TILE_SIZE, DOT_2X2, DOT_2X1) \ static void hvx_mm_4d_repacked_##SUFFIX(unsigned int nth, unsigned int ith, void * data) { \ @@ -713,7 +762,6 @@ static void hvx_mm_4d_repacked_##SUFFIX(unsigned int nth, unsigned int ith, void const uint32_t ct_start = src0_start_row / 32; \ const uint32_t ct_end = (src0_end_row + 31) / 32; \ \ - hvx_mm_run_quant_task(mmctx, ith); \ \ if (src0_start_row >= src0_end_row || cur_m_rows == 0) { \ return; \ @@ -810,6 +858,7 @@ MATMUL_4D_REPACKED_IMPL(q4_0, 576, tiled_vec_dot_q4_0_32x2, tiled_vec_do MATMUL_4D_REPACKED_IMPL(q4_1, 640, tiled_vec_dot_q4_1_32x2, tiled_vec_dot_q4_1_32x1) MATMUL_4D_REPACKED_IMPL(q8_0, 1088, tiled_vec_dot_q8_0_32x2, tiled_vec_dot_q8_0_32x1) MATMUL_4D_REPACKED_IMPL(q6_k, 896, tiled_vec_dot_q6_k_32x2, tiled_vec_dot_q6_k_32x1) +MATMUL_4D_REPACKED_IMPL(q5_k, 768, tiled_vec_dot_q5_k_32x2, tiled_vec_dot_q5_k_32x1) MATMUL_4D_REPACKED_IMPL(iq4nl, 576, tiled_vec_dot_iq4nl_32x2, tiled_vec_dot_iq4nl_32x1) MATMUL_4D_REPACKED_IMPL(mxfp4, 544, tiled_vec_dot_mxfp4_32x2, tiled_vec_dot_mxfp4_32x1) @@ -822,7 +871,7 @@ static void hvx_mm_2d(unsigned int nth, unsigned int ith, void * data) { const uint32_t prefetch_mask = n_prefetch - 1; const uint32_t src0_nrows = mmctx->src0_row_end - mmctx->src0_row_start; // src0 rows - const uint32_t src1_nrows = mmctx->cur_m_rows ? mmctx->cur_m_rows : mmctx->src1_nrows; // src1 rows + const uint32_t src1_nrows = mmctx->cur_m_rows ? mmctx->cur_m_rows : mmctx->act_nrows; // src1 rows const uint32_t cur_m_start = mmctx->cur_m_start; const uint32_t src0_start_row = mmctx->src0_row_start + src0_nrows_per_thread * ith; @@ -845,6 +894,7 @@ static void hvx_mm_2d(unsigned int nth, unsigned int ith, void * data) { const dma_addr_t src0_row = src0->data; + // Prefill vtcm with src0 rows if (src0_start_row < src0_end_row) { for (uint32_t ir0 = src0_start_row; ir0 < src0_end_row_x2; ir0 += 2) { @@ -857,8 +907,6 @@ static void hvx_mm_2d(unsigned int nth, unsigned int ith, void * data) { } } - hvx_mm_run_quant_task(mmctx, ith); - if (src0_start_row >= src0_end_row) { return; } @@ -974,7 +1022,6 @@ static void hvx_mv_2d(unsigned int nth, unsigned int ith, void * data) { } } - hvx_mm_run_quant_task(mmctx, ith); if (src0_start_row >= src0_end_row) { return; @@ -1049,7 +1096,6 @@ static void hvx_mm_4d(unsigned int nth, unsigned int ith, void * data) { uint8_t * restrict vtcm_src0_ptr = mmctx->vtcm_src0 + mmctx->vtcm_src0_size_per_thread * ith; uint8_t * restrict src1_data = mmctx->vtcm_src1; - hvx_mm_run_quant_task(mmctx, ith); if (src0_start_row >= src0_end_row || cur_m_rows == 0) { return; @@ -1183,7 +1229,6 @@ static void hvx_mm_id(unsigned int nth, unsigned int ith, void * data) { const uint32_t src0_start_row = mmctx->src0_row_start + src0_nrows_per_thread * ith; const uint32_t src0_end_row = MIN(src0_start_row + src0_nrows_per_thread, mmctx->src0_row_end); - hvx_mm_run_quant_task(mmctx, ith); if (src0_start_row >= src0_end_row) { return; @@ -1272,7 +1317,6 @@ static void hvx_mv_id(unsigned int nth, unsigned int ith, void * data) { const uint32_t src0_start_row = mmctx->src0_row_start + src0_nrows_per_thread * ith; const uint32_t src0_end_row = MIN(src0_start_row + src0_nrows_per_thread, mmctx->src0_row_end); - hvx_mm_run_quant_task(mmctx, ith); if (src0_start_row >= src0_end_row) { return; @@ -1351,7 +1395,6 @@ static void hvx_mv_id_nx(unsigned int nth, unsigned int ith, void * data) { const struct htp_tensor * restrict act = octx->src[n_weights]; const struct htp_tensor * restrict ids = octx->src[n_weights + 1]; - hvx_mm_run_quant_task(mmctx, ith); struct htp_thread_trace * tr = &octx->ctx->trace[ith]; @@ -1442,7 +1485,6 @@ static void hvx_mm_id_nx(unsigned int nth, unsigned int ith, void * data) { const struct htp_tensor * restrict act = octx->src[n_weights]; const struct htp_tensor * restrict ids = octx->src[n_weights + 1]; - hvx_mm_run_quant_task(mmctx, ith); struct htp_thread_trace * tr = &octx->ctx->trace[ith]; @@ -1550,6 +1592,10 @@ static int hvx_mm_init_vec_dot(struct htp_mm_context * mmctx, enum htp_data_type mmctx->type = "q8_0_tiled-f32"; mmctx->vec_dot_32x1 = tiled_vec_dot_q8_0_32x1; return 0; + case HTP_TYPE_Q5_K: + mmctx->type = "q5_k_tiled-f32"; + mmctx->vec_dot_32x1 = tiled_vec_dot_q5_k_32x1; + return 0; case HTP_TYPE_Q6_K: mmctx->type = "q6_k_tiled-f32"; mmctx->vec_dot_32x1 = tiled_vec_dot_q6_k_32x1; @@ -1582,7 +1628,7 @@ static int hvx_mm_matmul(struct htp_ops_context * octx) { const uint32_t src0_nrows = ne01; const uint32_t src1_nrows = ne11 * ne12 * ne13; - mmctx->src1_nrows = src1_nrows; + mmctx->act_nrows = src1_nrows; uint32_t src0_row_start = 0; uint32_t src0_row_end = src0_nrows; @@ -1605,7 +1651,7 @@ static int hvx_mm_matmul(struct htp_ops_context * octx) { bool is_repacked = (src0->type == HTP_TYPE_Q4_0 || src0->type == HTP_TYPE_Q4_1 || src0->type == HTP_TYPE_Q8_0 || src0->type == HTP_TYPE_IQ4_NL || src0->type == HTP_TYPE_MXFP4 || src0->type == HTP_TYPE_Q6_K || - src0->type == HTP_TYPE_Q4_K); + src0->type == HTP_TYPE_Q4_K || src0->type == HTP_TYPE_Q5_K); // Compute src0_nrows_per_thread mmctx->src0_nrows_per_thread = fastdiv(nrows + octx->n_threads - 1, &octx->n_threads_div); @@ -1634,6 +1680,7 @@ static int hvx_mm_matmul(struct htp_ops_context * octx) { case HTP_TYPE_Q4_K: matmul_job_func = hvx_mm_4d_repacked_q4_1; break; case HTP_TYPE_Q8_0: matmul_job_func = hvx_mm_4d_repacked_q8_0; break; case HTP_TYPE_Q6_K: matmul_job_func = hvx_mm_4d_repacked_q6_k; break; + case HTP_TYPE_Q5_K: matmul_job_func = hvx_mm_4d_repacked_q5_k; break; case HTP_TYPE_IQ4_NL: matmul_job_func = hvx_mm_4d_repacked_iq4nl; break; case HTP_TYPE_MXFP4: matmul_job_func = hvx_mm_4d_repacked_mxfp4; break; default: return HTP_STATUS_NO_SUPPORT; @@ -1649,6 +1696,7 @@ static int hvx_mm_matmul(struct htp_ops_context * octx) { case HTP_TYPE_Q4_K: matmul_job_func = hvx_mm_2d_repacked_q4_1; break; case HTP_TYPE_Q8_0: matmul_job_func = hvx_mm_2d_repacked_q8_0; break; case HTP_TYPE_Q6_K: matmul_job_func = hvx_mm_2d_repacked_q6_k; break; + case HTP_TYPE_Q5_K: matmul_job_func = hvx_mm_2d_repacked_q5_k; break; case HTP_TYPE_IQ4_NL: matmul_job_func = hvx_mm_2d_repacked_iq4nl; break; case HTP_TYPE_MXFP4: matmul_job_func = hvx_mm_2d_repacked_mxfp4; break; default: return HTP_STATUS_NO_SUPPORT; @@ -1663,6 +1711,7 @@ static int hvx_mm_matmul(struct htp_ops_context * octx) { case HTP_TYPE_Q4_1: case HTP_TYPE_Q4_K: matmul_job_func = hvx_mv_2d_repacked_q4_1; break; case HTP_TYPE_Q8_0: matmul_job_func = hvx_mv_2d_repacked_q8_0; break; + case HTP_TYPE_Q5_K: matmul_job_func = hvx_mv_2d_repacked_q5_k; break; case HTP_TYPE_Q6_K: matmul_job_func = hvx_mv_2d_repacked_q6_k; break; case HTP_TYPE_IQ4_NL: matmul_job_func = hvx_mv_2d_repacked_iq4nl; break; case HTP_TYPE_MXFP4: matmul_job_func = hvx_mv_2d_repacked_mxfp4; break; @@ -1678,7 +1727,8 @@ static int hvx_mm_matmul(struct htp_ops_context * octx) { switch (kparams->kernel_type) { case HTP_MM_KERNEL_HVX_F16_F16_VTCM: quant_task_func = (src1->type == HTP_TYPE_F32) ? quantize_f32_f16 : quantize_f16_f16; - mmctx->type = "f16-f16"; + need_quant = (src1->type == HTP_TYPE_F32); + mmctx->type = (src1->type == HTP_TYPE_F32) ? "f32-f16" : "f16-f16"; mmctx->vec_dot_1x1 = vec_dot_f16_f16_aa_1x1; mmctx->vec_dot_2x1 = vec_dot_f16_f16_aa_2x1; mmctx->vec_dot_2x2 = vec_dot_f16_f16_aa_2x2; @@ -1686,7 +1736,8 @@ static int hvx_mm_matmul(struct htp_ops_context * octx) { break; case HTP_MM_KERNEL_HVX_F32_F32_VTCM: - quant_task_func = quantize_f32_f32; + quant_task_func = NULL; + need_quant = false; mmctx->type = "f32-f32"; mmctx->vec_dot_1x1 = vec_dot_f32_f32_aa_1x1; mmctx->vec_dot_2x1 = vec_dot_f32_f32_aa_2x1; @@ -1707,7 +1758,7 @@ static int hvx_mm_matmul(struct htp_ops_context * octx) { if (src1_nrows < octx->n_threads && !is_batched) { n_quant_tasks = MIN(total_nb, octx->n_threads); - quant_task_func = (src0->type == HTP_TYPE_Q4_1 || src0->type == HTP_TYPE_Q4_K) ? quantize_f32_q8_1_tiled_block : quantize_f32_q8_0_tiled_block; + quant_task_func = htp_mm_weight_has_offset(src0->type) ? quantize_f32_q8_1_tiled_block : quantize_f32_q8_0_tiled_block; for (uint32_t ith = 0; ith < n_quant_tasks; ++ith) { uint32_t ib_first = (total_nb * ith) / n_quant_tasks; uint32_t ib_last = (total_nb * (ith + 1)) / n_quant_tasks; @@ -1718,9 +1769,9 @@ static int hvx_mm_matmul(struct htp_ops_context * octx) { } } else { n_quant_tasks = MIN(src1_nrows, octx->n_threads); - quant_task_func = (src0->type == HTP_TYPE_Q4_1 || src0->type == HTP_TYPE_Q4_K) ? quantize_f32_q8_1_tiled : quantize_f32_q8_0_tiled; + quant_task_func = htp_mm_weight_has_offset(src0->type) ? quantize_f32_q8_1_tiled : quantize_f32_q8_0_tiled; } - src1_row_size = (src0->type == HTP_TYPE_Q4_1 || src0->type == HTP_TYPE_Q4_K) ? htp_mm_q8_1_tiled_row_size(ne10) : htp_mm_q8_0_tiled_row_size(ne10); + src1_row_size = htp_mm_weight_has_offset(src0->type) ? htp_mm_q8_1_tiled_row_size(ne10) : htp_mm_q8_0_tiled_row_size(ne10); break; } @@ -1760,10 +1811,11 @@ static int hvx_mm_matmul(struct htp_ops_context * octx) { } uint8_t * const base = (uint8_t *) octx->ctx->vtcm_base; - mmctx->vtcm_src1 = VTCM_LAYOUT_PTR(uint8_t, base, L.off_src1); - mmctx->vtcm_src0 = VTCM_LAYOUT_PTR(uint8_t, base, L.off_src0); - mmctx->vtcm_src2 = VTCM_LAYOUT_PTR(uint8_t, base, L.off_src2); - mmctx->vtcm_dst = VTCM_LAYOUT_PTR(uint8_t, base, L.off_dst); + mmctx->vtcm_src1 = VTCM_LAYOUT_PTR(uint8_t, base, L.off_src1); + mmctx->vtcm_src0 = VTCM_LAYOUT_PTR(uint8_t, base, L.off_src0); + mmctx->vtcm_src2 = VTCM_LAYOUT_PTR(uint8_t, base, L.off_src2); + mmctx->vtcm_dst = VTCM_LAYOUT_PTR(uint8_t, base, L.off_dst); + mmctx->vtcm_act_raw = VTCM_LAYOUT_PTR(uint8_t, base, L.off_act_raw); octx->src1_spad.src = NULL; octx->src0_spad.src = NULL; @@ -1771,46 +1823,68 @@ static int hvx_mm_matmul(struct htp_ops_context * octx) { mmctx->vtcm_src0_stride = src0_row_size_padded; mmctx->vtcm_src1_stride = src1_row_size; + if (kparams->kernel_type == HTP_MM_KERNEL_HVX_QUANT_BLOCK || kparams->kernel_type == HTP_MM_KERNEL_HVX_QUANT_ROW) { + mmctx->vtcm_act_raw_stride = hex_round_up(ne10 * sizeof(float), QK_Q8_0_TILED * sizeof(float)); + } else if (kparams->kernel_type == HTP_MM_KERNEL_HVX_F16_F16_VTCM) { + mmctx->vtcm_act_raw_stride = hex_round_up(ne10 * sizeof(float), 128); + } else { + mmctx->vtcm_act_raw_stride = 0; + } + + htp_trace_event_stop(tr, HTP_TRACE_EVT_INIT, 0); if (kparams->m_chunk > 0 && (uint32_t) kparams->m_chunk < src1_nrows) { - atomic_init(&mmctx->quant_barrier, 0); - htp_trace_event_stop(tr, HTP_TRACE_EVT_INIT, 0); - for (uint32_t m_start = 0; m_start < src1_nrows; m_start += m_chunk) { const uint32_t cur_m_rows = MIN(src1_nrows - m_start, m_chunk); mmctx->cur_m_start = m_start; mmctx->cur_m_rows = cur_m_rows; if (need_quant) { - const uint32_t quant_tasks = MIN(cur_m_rows, octx->n_threads); - mmctx->n_quant_rows_per_thread = (cur_m_rows + quant_tasks - 1) / quant_tasks; + hvx_mm_transfer_src1_dma(octx, kparams, src1, mmctx->vtcm_act_raw, mmctx->vtcm_act_raw_stride, m_start, cur_m_rows); + + const uint32_t qk = QK_Q8_0_TILED; + const uint32_t nb = (ne10 + qk - 1) / qk; + const uint32_t total_nb = cur_m_rows * nb; + uint32_t quant_tasks; + work_queue_func_t q_func; + if (cur_m_rows < octx->n_threads && (kparams->kernel_type == HTP_MM_KERNEL_HVX_QUANT_BLOCK || kparams->kernel_type == HTP_MM_KERNEL_HVX_QUANT_ROW)) { + quant_tasks = MIN(total_nb, octx->n_threads); + q_func = htp_mm_weight_has_offset(src0->type) ? quantize_f32_q8_1_tiled_block : quantize_f32_q8_0_tiled_block; + for (uint32_t ith = 0; ith < quant_tasks; ++ith) { + uint32_t ib_first = (total_nb * ith) / quant_tasks; + uint32_t ib_last = (total_nb * (ith + 1)) / quant_tasks; + mmctx->quant_ib_first[ith] = ib_first; + mmctx->quant_ib_last[ith] = ib_last; + mmctx->quant_r[ith] = ib_first / nb; + mmctx->quant_c[ith] = ib_first % nb; + } + } else { + quant_tasks = MIN(cur_m_rows, octx->n_threads); + q_func = quant_task_func; + mmctx->n_quant_rows_per_thread = (cur_m_rows + quant_tasks - 1) / quant_tasks; + } mmctx->n_quant_tasks = quant_tasks; - atomic_store(&mmctx->quant_barrier, quant_tasks); - mmctx->quant_task_func = quant_task_func; + work_queue_run(octx->ctx->work_queue, q_func, mmctx, quant_tasks); } else { - mmctx->quant_task_func = NULL; - mmctx->n_quant_tasks = 0; + hvx_mm_transfer_src1_dma(octx, kparams, src1, mmctx->vtcm_src1, mmctx->vtcm_src1_stride, m_start, cur_m_rows); } - worker_pool_run_func(octx->ctx->worker_pool, matmul_job_func, mmctx, octx->n_threads); + work_queue_run(octx->ctx->work_queue, matmul_job_func, mmctx, octx->n_threads); } } else { mmctx->cur_m_start = 0; mmctx->cur_m_rows = src1_nrows; if (need_quant) { + hvx_mm_transfer_src1_dma(octx, kparams, src1, mmctx->vtcm_act_raw, mmctx->vtcm_act_raw_stride, 0, src1_nrows); mmctx->n_quant_rows_per_thread = (src1_nrows + n_quant_tasks - 1) / n_quant_tasks; - mmctx->quant_task_func = quant_task_func; mmctx->n_quant_tasks = n_quant_tasks; - atomic_init(&mmctx->quant_barrier, n_quant_tasks); + work_queue_run(octx->ctx->work_queue, quant_task_func, mmctx, n_quant_tasks); } else { - mmctx->quant_task_func = NULL; - mmctx->n_quant_tasks = 0; + hvx_mm_transfer_src1_dma(octx, kparams, src1, mmctx->vtcm_src1, mmctx->vtcm_src1_stride, 0, src1_nrows); } - htp_trace_event_stop(tr, HTP_TRACE_EVT_INIT, 0); - - worker_pool_run_func(octx->ctx->worker_pool, matmul_job_func, mmctx, octx->n_threads); + work_queue_run(octx->ctx->work_queue, matmul_job_func, mmctx, octx->n_threads); } return HTP_STATUS_OK; @@ -1835,7 +1909,6 @@ static void hvx_mm_nx_2d(unsigned int nth, unsigned int ith, void * data) { struct htp_thread_trace * tr = &octx->ctx->trace[ith]; - hvx_mm_run_quant_task(mmctx, ith); for (uint32_t widx = 0; widx < n_weights; widx++) { const struct htp_tensor * restrict src_w = octx->src[widx]; @@ -1938,6 +2011,7 @@ DEQUANTIZE_WORKER_LOOP_IMPL(iq4_nl) DEQUANTIZE_WORKER_LOOP_IMPL(mxfp4) DEQUANTIZE_WORKER_LOOP_IMPL(q8_0) DEQUANTIZE_WORKER_LOOP_IMPL(q6_k) +DEQUANTIZE_WORKER_LOOP_IMPL(q5_k) static void convert_f16_worker_loop(unsigned int n, unsigned int i, void *data) { tiled_dequantize_state_t *state = (tiled_dequantize_state_t *)data; @@ -1990,8 +2064,7 @@ typedef struct { struct htp_context * ctx; struct htp_thread_trace * traces; __fp16 * dst; - const float * src; - const struct mmid_row_mapping * matrix_rows; + dma_addr_t act_dma_addr; float * vtcm_f32_act; uint32_t n_tasks; uint32_t n_tot_chunks; @@ -2008,7 +2081,7 @@ typedef struct { struct htp_context * ctx; struct htp_thread_trace * traces; __fp16 * dst; - const float * src; + dma_addr_t act_dma_addr; float * vtcm_f32_act; uint32_t n_rows; uint32_t k_block; @@ -2024,7 +2097,7 @@ typedef struct { static void transfer_activation_chunk_fp32_to_fp16_dma_pipelined_col_chunk( dma_queue *dma_q, __fp16 *restrict vtcm_dst, - const float *restrict src, + dma_addr_t act_dma_addr, uint32_t n_rows, uint32_t k_block, uint32_t k_stride, @@ -2044,7 +2117,7 @@ static void transfer_activation_chunk_fp32_to_fp16_dma_pipelined_col_chunk( // Push step 0 if (n_steps > 0 && n_rows > 0) { uint32_t nrows_to_fetch = hex_smin(n_rows, R); - dma_queue_push(dma_q, dma_make_data(thread_f32_act, src + c_first), + dma_queue_push(dma_q, dma_make_data(thread_f32_act, act_dma_addr + (size_t) c_first * sizeof(float)), c_len * sizeof(float), k_stride * sizeof(float), k_chunk_valid * sizeof(float), nrows_to_fetch); } // Push step 1 @@ -2052,9 +2125,8 @@ static void transfer_activation_chunk_fp32_to_fp16_dma_pipelined_col_chunk( uint32_t next_r = R * 1; if (next_r < n_rows) { uint32_t nrows_to_fetch = hex_smin(n_rows - next_r, R); - const float *next_src = src + next_r * k_stride + c_first; float *next_buf = thread_f32_act + 1 * R * c_len; - dma_queue_push(dma_q, dma_make_data(next_buf, next_src), + dma_queue_push(dma_q, dma_make_data(next_buf, act_dma_addr + ((size_t) next_r * k_stride + c_first) * sizeof(float)), c_len * sizeof(float), k_stride * sizeof(float), k_chunk_valid * sizeof(float), nrows_to_fetch); } } @@ -2084,50 +2156,12 @@ static void transfer_activation_chunk_fp32_to_fp16_dma_pipelined_col_chunk( uint32_t next_r = next_s << dma_step_rows_shift; if (next_r < n_rows) { uint32_t nrows_to_fetch = hex_smin(n_rows - next_r, R); - const float *next_src = src + next_r * k_stride + c_first; - dma_queue_push(dma_q, dma_make_data(curr_buf, next_src), + dma_queue_push(dma_q, dma_make_data(curr_buf, act_dma_addr + ((size_t) next_r * k_stride + c_first) * sizeof(float)), c_len * sizeof(float), k_stride * sizeof(float), k_chunk_valid * sizeof(float), nrows_to_fetch); } } } -static void transfer_activation_chunk_fp32_to_fp16_col_chunk( - __fp16 *restrict vtcm_dst, - const float *restrict src, - uint32_t n_rows, - uint32_t k_block, - uint32_t k_stride, - uint32_t c_first, - uint32_t c_len, - uint32_t k_chunk_valid) { - const uint32_t n_rows_padded = hex_align_up(n_rows, HTP_MM_HMX_TILE_N_ROWS); - const uint32_t n_rows_tiled = (n_rows / HTP_MM_HMX_TILE_N_ROWS) * HTP_MM_HMX_TILE_N_ROWS; - - uint32_t r = 0; - - #pragma unroll(2) - for (r = 0; r < n_rows_tiled; r += 2) { - const float *ptr_in0 = src + (r + 0) * k_stride + c_first; - const float *ptr_in1 = src + (r + 1) * k_stride + c_first; - - transfer_activation_row_pair_fp32_to_fp16_col_chunk( - vtcm_dst, ptr_in0, ptr_in1, r, k_block, c_first, c_len, k_chunk_valid, true, true - ); - } - - for (; r < n_rows_padded; r += 2) { - const bool row0_valid = r < n_rows; - const bool row1_valid = (r + 1) < n_rows; - - const float *ptr_in0 = row0_valid ? (src + (r + 0) * k_stride + c_first) : NULL; - const float *ptr_in1 = row1_valid ? (src + (r + 1) * k_stride + c_first) : NULL; - - transfer_activation_row_pair_fp32_to_fp16_col_chunk( - vtcm_dst, ptr_in0, ptr_in1, r, k_block, c_first, c_len, k_chunk_valid, row0_valid, row1_valid - ); - } -} - static void transfer_activation_chunk_col_chunk_worker_fn(unsigned int n, unsigned int i, void *data) { activation_transfer_col_chunk_state_t *st = (activation_transfer_col_chunk_state_t *) data; struct htp_thread_trace * tr = &st->traces[i]; @@ -2149,29 +2183,20 @@ static void transfer_activation_chunk_col_chunk_worker_fn(unsigned int n, unsign } __fp16 *dst = st->dst; - const float *src = st->src; - if (st->vtcm_f32_act) { - size_t thread_scratch_bytes = hex_align_down(fastdiv(st->vtcm_f32_act_bytes, &st->n_threads_div), 128); - float *thread_f32_act = (float *)((char *)st->vtcm_f32_act + i * thread_scratch_bytes); + size_t thread_scratch_bytes = hex_align_down(fastdiv(st->vtcm_f32_act_bytes, &st->n_threads_div), 128); + float *thread_f32_act = (float *)((char *)st->vtcm_f32_act + i * thread_scratch_bytes); - transfer_activation_chunk_fp32_to_fp16_dma_pipelined_col_chunk( - st->ctx->dma[i], dst, src, st->n_rows, st->k_block, st->k_stride, k_chunk_valid, - c_first, c_len, thread_f32_act, tr, st->dma_step_rows, st->dma_step_rows_shift - ); - } else { - htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_A_PREP, c_first); - transfer_activation_chunk_fp32_to_fp16_col_chunk( - dst, src, st->n_rows, st->k_block, st->k_stride, c_first, c_len, k_chunk_valid - ); - htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_A_PREP, c_first); - } + transfer_activation_chunk_fp32_to_fp16_dma_pipelined_col_chunk( + st->ctx->dma[i], dst, st->act_dma_addr, st->n_rows, st->k_block, st->k_stride, k_chunk_valid, + c_first, c_len, thread_f32_act, tr, st->dma_step_rows, st->dma_step_rows_shift + ); } static void transfer_activation_chunk_fp32_to_fp16_dma_pipelined( dma_queue *dma_q, __fp16 *restrict vtcm_dst, - const float *restrict src, + dma_addr_t act_dma_addr, uint32_t n_rows, uint32_t k_block, uint32_t k_stride, @@ -2189,7 +2214,7 @@ static void transfer_activation_chunk_fp32_to_fp16_dma_pipelined( // Push step 0 if (n_steps > 0 && n_rows > 0) { uint32_t nrows_to_fetch = hex_smin(n_rows, R); - dma_queue_push(dma_q, dma_make_data(thread_f32_act, src), + dma_queue_push(dma_q, dma_make_data(thread_f32_act, act_dma_addr), k_block * sizeof(float), k_stride * sizeof(float), k_valid * sizeof(float), nrows_to_fetch); } // Push step 1 (if valid) @@ -2197,9 +2222,8 @@ static void transfer_activation_chunk_fp32_to_fp16_dma_pipelined( uint32_t next_r = R * 1; if (next_r < n_rows) { uint32_t nrows_to_fetch = hex_smin(n_rows - next_r, R); - const float *next_src = src + next_r * k_stride; float *next_buf = thread_f32_act + 1 * R * k_block; - dma_queue_push(dma_q, dma_make_data(next_buf, next_src), + dma_queue_push(dma_q, dma_make_data(next_buf, act_dma_addr + (size_t) next_r * k_stride * sizeof(float)), k_block * sizeof(float), k_stride * sizeof(float), k_valid * sizeof(float), nrows_to_fetch); } } @@ -2227,8 +2251,7 @@ static void transfer_activation_chunk_fp32_to_fp16_dma_pipelined( uint32_t next_r = next_s << dma_step_rows_shift; if (next_r < n_rows) { uint32_t nrows_to_fetch = hex_smin(n_rows - next_r, R); - const float *next_src = src + next_r * k_stride; - dma_queue_push(dma_q, dma_make_data(curr_buf, next_src), + dma_queue_push(dma_q, dma_make_data(curr_buf, act_dma_addr + (size_t) next_r * k_stride * sizeof(float)), k_block * sizeof(float), k_stride * sizeof(float), k_valid * sizeof(float), nrows_to_fetch); } } @@ -2244,18 +2267,12 @@ static void transfer_activation_chunk_worker_fn(unsigned int n, unsigned int i, size_t chunk_size = hex_smin(st->n_tot_chunks - chunk_idx, st->n_chunks_per_task); __fp16 *dst = st->dst + chunk_idx * st->k_block; - const float *src = st->src + chunk_idx * st->k_stride; + const dma_addr_t act_dma_addr = st->act_dma_addr + (size_t) chunk_idx * st->k_stride * sizeof(float); - if (st->vtcm_f32_act) { - float *thread_f32_act = (float *)((char *)st->vtcm_f32_act + i * st->vtcm_f32_act_bytes_per_thread); - transfer_activation_chunk_fp32_to_fp16_dma_pipelined( - st->ctx->dma[i], dst, src, chunk_size, st->k_block, st->k_stride, st->k_valid, thread_f32_act, tr, st->dma_step_rows, st->dma_step_rows_shift - ); - } else { - htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_A_PREP, chunk_idx); - transfer_activation_chunk_fp32_to_fp16(dst, src, chunk_size, st->k_block, st->k_stride, st->k_valid); - htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_A_PREP, chunk_idx); - } + float *thread_f32_act = (float *)((char *)st->vtcm_f32_act + i * st->vtcm_f32_act_bytes_per_thread); + transfer_activation_chunk_fp32_to_fp16_dma_pipelined( + st->ctx->dma[i], dst, act_dma_addr, chunk_size, st->k_block, st->k_stride, st->k_valid, thread_f32_act, tr, st->dma_step_rows, st->dma_step_rows_shift + ); } } @@ -2494,7 +2511,7 @@ static void transfer_output_chunk_threaded(struct htp_context *ctx, float *dst, struct activation_transfer_params { struct htp_context * ctx; __fp16 * dst; - const float * src; + dma_addr_t act_dma_addr; int n_rows; int k_block; int k_stride; @@ -2509,7 +2526,7 @@ struct activation_transfer_params { static void transfer_activation_chunk_threaded(const struct activation_transfer_params * params) { struct htp_context * ctx = params->ctx; __fp16 * dst = params->dst; - const float * src = params->src; + const dma_addr_t act_dma_addr = params->act_dma_addr; int n_rows = params->n_rows; int k_block = params->k_block; int k_stride = params->k_stride; @@ -2529,7 +2546,7 @@ static void transfer_activation_chunk_threaded(const struct activation_transfer_ // Calculate step rows parameters for column-chunked dma pipelining uint32_t dma_step_rows = 2; uint32_t dma_step_rows_shift = 1; - if (vtcm_f32_act && vtcm_f32_act_bytes > 0 && k_block > 0) { + if (vtcm_f32_act_bytes > 0) { size_t thread_scratch_bytes = hex_align_down(fastdiv(vtcm_f32_act_bytes, act_threads_div), 128); size_t thread_scratch_elements = thread_scratch_bytes / sizeof(float); size_t dma_step_rows_max = fastdiv(thread_scratch_elements / 2, k_div); @@ -2541,7 +2558,7 @@ static void transfer_activation_chunk_threaded(const struct activation_transfer_ activation_transfer_col_chunk_state_t col_state; col_state.dst = dst; - col_state.src = src; + col_state.act_dma_addr = act_dma_addr; col_state.n_rows = n_rows; col_state.k_block = k_block; col_state.k_stride = k_stride; @@ -2569,7 +2586,7 @@ static void transfer_activation_chunk_threaded(const struct activation_transfer_ state.n_tot_chunks = n_tot_chunks; state.n_chunks_per_task = n_chunks_per_task; state.dst = dst; - state.src = src; + state.act_dma_addr = act_dma_addr; state.k_block = k_block; state.k_stride = k_stride; state.k_valid = k_valid; @@ -2581,7 +2598,7 @@ static void transfer_activation_chunk_threaded(const struct activation_transfer_ uint32_t dma_step_rows = 2; uint32_t dma_step_rows_shift = 1; - if (vtcm_f32_act && state.vtcm_f32_act_bytes_per_thread > 0 && k_block > 0) { + if (state.vtcm_f32_act_bytes_per_thread > 0) { size_t thread_scratch_elements = state.vtcm_f32_act_bytes_per_thread / sizeof(float); size_t dma_step_rows_max = fastdiv(thread_scratch_elements / 2, k_div); if (dma_step_rows_max >= 4) { @@ -2639,7 +2656,7 @@ static int hmx_mm_2d_f32(struct htp_context *ctx, float *restrict dst, dma_addr_t src2_addr, size_t src2_bytes, - const float *activation, + dma_addr_t act_dma_addr, dma_addr_t weight, int m, int k, int n, int act_stride, @@ -2663,7 +2680,7 @@ static int hmx_mm_2d_f32(struct htp_context *ctx, htp_trace_event_start(tr, HTP_TRACE_EVT_INIT, 0); if (k % 32 != 0 || n % 32 != 0) { return -1; } - if (!hex_is_aligned(dst, VLEN) || !hex_is_aligned(activation, VLEN)) { return -1; } + if (!hex_is_aligned(dst, VLEN) || (act_dma_addr & (VLEN - 1)) != 0) { return -1; } size_t row_stride = htp_mm_get_tiled_row_stride(weight_type, k); if (row_stride == 0) { @@ -2678,6 +2695,7 @@ static int hmx_mm_2d_f32(struct htp_context *ctx, case HTP_TYPE_Q4_K: dequant_worker_fn = dequantize_tiled_worker_loop_q4_1; break; case HTP_TYPE_MXFP4: dequant_worker_fn = dequantize_tiled_worker_loop_mxfp4; break; case HTP_TYPE_Q8_0: dequant_worker_fn = dequantize_tiled_worker_loop_q8_0; break; + case HTP_TYPE_Q5_K: dequant_worker_fn = dequantize_tiled_worker_loop_q5_k; break; case HTP_TYPE_Q6_K: dequant_worker_fn = dequantize_tiled_worker_loop_q6_k; break; case HTP_TYPE_F16: dequant_worker_fn = convert_f16_worker_loop; break; case HTP_TYPE_F32: dequant_worker_fn = quantize_f32_worker_loop; break; @@ -2703,7 +2721,7 @@ static int hmx_mm_2d_f32(struct htp_context *ctx, const size_t qweight_row_stride = is_quant ? (size_t)(n_k_tiles * aligned_tile_size) / 32 : 0; struct htp_mm_hmx_vtcm_layout L; - htp_mm_hmx_vtcm_layout_build(&L, HTP_MM_KERNEL_HMX_2D, weight_type, k, m_chunk_n_rows, n_chunk_n_cols, 1, false, pipeline, act_threads, aligned_tile_size, src2_bytes); + htp_mm_hmx_vtcm_layout_build(&L, HTP_MM_KERNEL_HMX_2D, weight_type, k, m_chunk_n_rows, n_chunk_n_cols, 1, pipeline, act_threads, aligned_tile_size, src2_bytes); vtcm_used = L.total_bytes; if (vtcm_used > vtcm_budget) { @@ -2755,7 +2773,7 @@ static int hmx_mm_2d_f32(struct htp_context *ctx, struct activation_transfer_params act_params = { .ctx = ctx, .dst = vtcm_f16_act, - .src = activation + mr * act_stride, + .act_dma_addr = act_dma_addr + mr * act_stride * sizeof(float), .n_rows = (int) n_rows, .k_block = k, .k_stride = act_stride, @@ -2846,7 +2864,7 @@ static int hmx_mm_2d_f32(struct htp_context *ctx, struct activation_transfer_params act_params = { .ctx = ctx, .dst = vtcm_f16_act, - .src = activation + mr * act_stride, + .act_dma_addr = act_dma_addr + mr * act_stride * sizeof(float), .n_rows = (int) n_rows, .k_block = k, .k_stride = act_stride, @@ -2925,10 +2943,10 @@ static int hmx_mm_nx_2d_f32(struct htp_ops_context * octx, const struct htp_mm_k const int k_valid = (int) act->ne[0]; const int m = (int) (act->ne[1] * act->ne[2] * act->ne[3]); const int act_stride = (int) (act->nb[1] / sizeof(float)); - const float * activation = (const float *) act->data; + const dma_addr_t act_dma_addr = act->data; if (k % 32 != 0) { return HTP_STATUS_NO_SUPPORT; } - if (!hex_is_aligned(activation, VLEN)) { return HTP_STATUS_NO_SUPPORT; } + if ((act_dma_addr & (VLEN - 1)) != 0) { return HTP_STATUS_NO_SUPPORT; } size_t row_stride = htp_mm_get_tiled_row_stride(weight_type, k); if (row_stride == 0) { @@ -2943,6 +2961,7 @@ static int hmx_mm_nx_2d_f32(struct htp_ops_context * octx, const struct htp_mm_k case HTP_TYPE_Q4_K: dequant_worker_fn = dequantize_tiled_worker_loop_q4_1; break; case HTP_TYPE_MXFP4: dequant_worker_fn = dequantize_tiled_worker_loop_mxfp4; break; case HTP_TYPE_Q8_0: dequant_worker_fn = dequantize_tiled_worker_loop_q8_0; break; + case HTP_TYPE_Q5_K: dequant_worker_fn = dequantize_tiled_worker_loop_q5_k; break; case HTP_TYPE_Q6_K: dequant_worker_fn = dequantize_tiled_worker_loop_q6_k; break; case HTP_TYPE_F16: dequant_worker_fn = convert_f16_worker_loop; break; case HTP_TYPE_F32: dequant_worker_fn = quantize_f32_worker_loop; break; @@ -2970,7 +2989,7 @@ static int hmx_mm_nx_2d_f32(struct htp_ops_context * octx, const struct htp_mm_k const uint32_t dma_width_bytes = is_quant ? tile_size : row_stride; struct htp_mm_hmx_vtcm_layout L; - htp_mm_hmx_vtcm_layout_build(&L, HTP_MM_KERNEL_HMX_2D, weight_type, k, m_chunk_n_rows, n_chunk_n_cols, 1, false, pipeline, act_threads, aligned_tile_size, 0); + htp_mm_hmx_vtcm_layout_build(&L, HTP_MM_KERNEL_HMX_2D, weight_type, k, m_chunk_n_rows, n_chunk_n_cols, 1, pipeline, act_threads, aligned_tile_size, 0); if (L.total_bytes > vtcm_budget) { FARF(ERROR, "hmx-mm-nx-2d: VTCM overflow: used %zu budget %zu, m %d k %d mc %d nc %d", @@ -3026,7 +3045,7 @@ static int hmx_mm_nx_2d_f32(struct htp_ops_context * octx, const struct htp_mm_k struct activation_transfer_params act_params = { .ctx = ctx, .dst = vtcm_f16_act, - .src = activation + mr * act_stride, + .act_dma_addr = act_dma_addr + mr * act_stride * sizeof(float), .n_rows = (int) n_rows, .k_block = k, .k_stride = act_stride, @@ -3124,7 +3143,7 @@ static int hmx_mm_nx_2d_f32(struct htp_ops_context * octx, const struct htp_mm_k struct activation_transfer_params act_params = { .ctx = ctx, .dst = vtcm_f16_act, - .src = activation + mr * act_stride, + .act_dma_addr = act_dma_addr + mr * act_stride * sizeof(float), .n_rows = (int) n_rows, .k_block = k, .k_stride = act_stride, @@ -3202,11 +3221,10 @@ static inline dma_addr_t hmx_mm_weight_batch_data(const hmx_mm_f16_f32_batched_p return params->weight + b2_idx * params->src0_nb2 + b3_idx * params->src0_nb3; } -static inline const float *hmx_mm_activation_batch_ptr(const hmx_mm_f16_f32_batched_params_t *params, - int dst_b2, int dst_b3) { - return (const float *) ((const uint8_t *) params->activation + - (size_t) dst_b2 * params->src1_nb2 + - (size_t) dst_b3 * params->src1_nb3); +static inline dma_addr_t hmx_mm_act_batch_addr(const hmx_mm_f16_f32_batched_params_t *params, + int dst_b2, int dst_b3) { + return params->act_dma_addr + dst_b2 * params->act_nb2 + + dst_b3 * params->act_nb3; } static inline float *hmx_mm_dst_batch_ptr(const hmx_mm_f16_f32_batched_params_t *params, @@ -3224,11 +3242,11 @@ static int hmx_mm_f16_f32_batched_simple(struct htp_context *ctx, for (int b3 = 0; b3 < params->ne13 && ret == 0; ++b3) { for (int b2 = 0; b2 < params->ne12 && ret == 0; ++b2) { dma_addr_t cur_src2_addr = params->src2_addr ? (params->src2_addr + - (dma_addr_t) b2 * params->src2_nb2 + - (dma_addr_t) b3 * params->src2_nb3) : 0; + b2 * params->src2_nb2 + + b3 * params->src2_nb3) : 0; ret = hmx_mm_2d_f32(ctx, params->weight_dma, hmx_mm_dst_batch_ptr(params, b2, b3), cur_src2_addr, params->src2_bytes, - hmx_mm_activation_batch_ptr(params, b2, b3), + hmx_mm_act_batch_addr(params, b2, b3), hmx_mm_weight_batch_data(params, b2, b3), params->m, params->k, params->n, params->act_stride, params->weight_stride * (int)sizeof(__fp16), @@ -3249,7 +3267,7 @@ static int hmx_mm_f16_f32_batched(struct htp_context *ctx, const hmx_mm_f16_f32_ if (params->ne02 <= 0 || params->ne03 <= 0 || params->ne12 <= 0 || params->ne13 <= 0) { return -1; } if (params->ne12 % params->ne02 != 0 || params->ne13 % params->ne03 != 0) { return -1; } if (params->k % 32 != 0 || params->n % 32 != 0) { return -1; } - if (!hex_is_aligned(params->dst, VLEN) || !hex_is_aligned(params->activation, VLEN)) { return -1; } + if (!hex_is_aligned(params->dst, VLEN) || (params->act_dma_addr & (VLEN - 1)) != 0) { return -1; } const int group_size = params->r2; const size_t vtcm_budget = ctx->vtcm_size; @@ -3267,16 +3285,12 @@ static int hmx_mm_f16_f32_batched(struct htp_context *ctx, const hmx_mm_f16_f32_ const size_t vec_dot_size = params->k * sizeof(__fp16); - const bool use_dma_activation = (params->act_stride > params->k); - const size_t f32_scratch_size = use_dma_activation - ? hex_align_up((size_t)act_threads * HTP_MM_DMA_ACT_MULTIPLIER * (size_t) params->k * sizeof(float), HTP_MM_HMX_TILE_SIZE) : 0; - size_t m_chunk_n_rows = m_chunk; size_t n_chunk_n_cols = n_chunk; size_t vtcm_used = vtcm_size; struct htp_mm_hmx_vtcm_layout L; - htp_mm_hmx_vtcm_layout_build(&L, HTP_MM_KERNEL_HMX_F16_BATCHED, HTP_TYPE_F16, params->k, m_chunk_n_rows, n_chunk_n_cols, group_size, use_dma_activation, false, act_threads, 0, params->src2_bytes); + htp_mm_hmx_vtcm_layout_build(&L, HTP_MM_KERNEL_HMX_F16_BATCHED, HTP_TYPE_F16, params->k, m_chunk_n_rows, n_chunk_n_cols, group_size, false, act_threads, 0, params->src2_bytes); if (L.total_bytes > vtcm_budget) { FARF(HIGH, "%s: grouped layout overflowed VTCM, falling back to simple batched loop", __func__); @@ -3291,7 +3305,7 @@ static int hmx_mm_f16_f32_batched(struct htp_context *ctx, const hmx_mm_f16_f32_ void *vtcm_scratch0 = VTCM_LAYOUT_PTR(void, base, L.off_scratch[0]); void *vtcm_scratch1 = VTCM_LAYOUT_PTR(void, base, L.off_scratch[1]); __fp16 *vtcm_scales = VTCM_LAYOUT_PTR(__fp16, base, L.off_scales); - float *vtcm_f32_act = VTCM_LAYOUT_PTR_OPTIONAL(float, base, L.off_act_f32, use_dma_activation); + float *vtcm_f32_act = VTCM_LAYOUT_PTR(float, base, L.off_act_f32); const bool has_src2 = (params->src2_bytes > 0 && params->src2_addr != 0); float *vtcm_src2 = VTCM_LAYOUT_PTR_OPTIONAL(float, base, L.off_src2, has_src2); @@ -3329,12 +3343,13 @@ static int hmx_mm_f16_f32_batched(struct htp_context *ctx, const hmx_mm_f16_f32_ // converts from the contiguous VTCM buffer. This avoids L2 cache // thrashing from HVX loads at large strides. for (int g = 0; g < group_size; ++g) { - const float *activation_chunk = hmx_mm_activation_batch_ptr(params, b2_base + g, b3) + mr * params->act_stride; + const dma_addr_t act_dma_addr = hmx_mm_act_batch_addr(params, b2_base + g, b3) + + mr * params->act_stride * sizeof(float); __fp16 *vtcm_act_g = vtcm_f16_act + (size_t) g * L.act_head_stride; struct activation_transfer_params act_params = { .ctx = ctx, .dst = vtcm_act_g, - .src = activation_chunk, + .act_dma_addr = act_dma_addr, .n_rows = (int) n_rows, .k_block = params->k, .k_stride = params->act_stride, @@ -3542,6 +3557,7 @@ static int hmx_mm_id_2d_f32(struct htp_context *ctx, case HTP_TYPE_Q4_K: dequant_worker_fn = dequantize_tiled_worker_loop_q4_1; break; case HTP_TYPE_MXFP4: dequant_worker_fn = dequantize_tiled_worker_loop_mxfp4; break; case HTP_TYPE_Q8_0: dequant_worker_fn = dequantize_tiled_worker_loop_q8_0; break; + case HTP_TYPE_Q5_K: dequant_worker_fn = dequantize_tiled_worker_loop_q5_k; break; case HTP_TYPE_Q6_K: dequant_worker_fn = dequantize_tiled_worker_loop_q6_k; break; case HTP_TYPE_F16: dequant_worker_fn = convert_f16_worker_loop; break; case HTP_TYPE_F32: dequant_worker_fn = quantize_f32_worker_loop; break; @@ -3692,15 +3708,15 @@ static int hmx_mm_op_matmul(struct htp_ops_context * octx, const struct htp_mm_k size_t src2_nb3 = 0; if (src2) { src2_stride = (src2->ne[1] == 1) ? 0 : (uint32_t) (src2->nb[1] / sizeof(float)); - src2_addr = src2->data + (dma_addr_t) m_start * src2_stride * sizeof(float); + src2_addr = src2->data + m_start * src2_stride * sizeof(float); src2_bytes = (size_t) kparams->vtcm_src2_size; src2_nb2 = (src2->ne[2] == 1) ? 0 : src2->nb[2]; src2_nb3 = (src2->ne[3] == 1) ? 0 : src2->nb[3]; } const int dst_stride = (int)(dst->nb[1] / sizeof(float)); - float * dst_ptr = (float *) dst->data + m_start * dst_stride; - const float * act_ptr = (const float *) src1->data + m_start * act_stride; + float * dst_ptr = (float *) dst->data + m_start * dst_stride; + const dma_addr_t act_addr = src1->data + m_start * act_stride * sizeof(float); int ret = -1; const int n_threads = kparams->n_threads; @@ -3709,7 +3725,7 @@ static int hmx_mm_op_matmul(struct htp_ops_context * octx, const struct htp_mm_k .dst = dst_ptr, .src2_addr = src2_addr, .src2_bytes = src2_bytes, - .activation = act_ptr, + .act_dma_addr = act_addr, .weight = src0->data, .weight_dma = octx->ctx->dma[0], .m = m_rows, @@ -3725,8 +3741,8 @@ static int hmx_mm_op_matmul(struct htp_ops_context * octx, const struct htp_mm_k .ne13 = ne13, .src0_nb2 = src0->nb[2], .src0_nb3 = src0->nb[3], - .src1_nb2 = src1->nb[2], - .src1_nb3 = src1->nb[3], + .act_nb2 = src1->nb[2], + .act_nb3 = src1->nb[3], .dst_nb2 = dst->nb[2], .dst_nb3 = dst->nb[3], .src2_nb2 = src2_nb2, @@ -3745,7 +3761,8 @@ static int hmx_mm_op_matmul(struct htp_ops_context * octx, const struct htp_mm_k kparams->vtcm_size); } else { ret = hmx_mm_2d_f32( - octx->ctx, octx->ctx->dma[0], dst_ptr, src2_addr, src2_bytes, act_ptr, src0->data, + octx->ctx, octx->ctx->dma[0], dst_ptr, src2_addr, src2_bytes, + act_addr, src0->data, m_rows, k, n, act_stride, (int) src0->nb[1], (int) src0->type, (int) src1->ne[0], dst_stride, src2_stride, (int)dst->ne[0], kparams->m_chunk, kparams->n_chunk, kparams->pipeline, n_threads, @@ -3829,7 +3846,7 @@ static int hvx_mm_matmul_id( ) { htp_matmul_tensors_preamble; const uint32_t src0_row_size_padded = mmctx->src0_row_size_padded; - const uint32_t src1_nrows = mmctx->src1_nrows; + const uint32_t act_nrows = mmctx->act_nrows; struct htp_thread_trace * tr = &octx->ctx->trace[0]; htp_trace_event_start(tr, HTP_TRACE_EVT_INIT, 0); @@ -3840,13 +3857,13 @@ static int hvx_mm_matmul_id( const uint32_t qk = QK_Q8_0_TILED; const uint32_t nb = (ne10 + qk - 1) / qk; - const uint32_t total_nb = src1_nrows * nb; + const uint32_t total_nb = act_nrows * nb; work_queue_func_t quant_task_func; uint32_t n_quant_tasks = 1; - if (src1_nrows < octx->n_threads) { + if (act_nrows < octx->n_threads) { n_quant_tasks = MIN(total_nb, octx->n_threads); - quant_task_func = (src0->type == HTP_TYPE_Q4_1 || src0->type == HTP_TYPE_Q4_K) ? quantize_f32_q8_1_tiled_block : quantize_f32_q8_0_tiled_block; + quant_task_func = htp_mm_weight_has_offset(src0->type) ? quantize_f32_q8_1_tiled_block : quantize_f32_q8_0_tiled_block; for (uint32_t ith = 0; ith < n_quant_tasks; ++ith) { uint32_t ib_first = (total_nb * ith) / n_quant_tasks; uint32_t ib_last = (total_nb * (ith + 1)) / n_quant_tasks; @@ -3856,13 +3873,13 @@ static int hvx_mm_matmul_id( mmctx->quant_c[ith] = ib_first % nb; } } else { - n_quant_tasks = MIN(src1_nrows, octx->n_threads); - quant_task_func = (src0->type == HTP_TYPE_Q4_1 || src0->type == HTP_TYPE_Q4_K) ? quantize_f32_q8_1_tiled : quantize_f32_q8_0_tiled; + n_quant_tasks = MIN(act_nrows, octx->n_threads); + quant_task_func = htp_mm_weight_has_offset(src0->type) ? quantize_f32_q8_1_tiled : quantize_f32_q8_0_tiled; } - size_t src1_row_size = (src0->type == HTP_TYPE_Q4_1 || src0->type == HTP_TYPE_Q4_K) ? htp_mm_q8_1_tiled_row_size(ne10) : htp_mm_q8_0_tiled_row_size(ne10); + size_t src1_row_size = htp_mm_weight_has_offset(src0->type) ? htp_mm_q8_1_tiled_row_size(ne10) : htp_mm_q8_0_tiled_row_size(ne10); struct htp_mm_hvx_vtcm_layout L; - htp_mm_hvx_vtcm_layout_build(&L, kparams->kernel_type, src0->type, ne10, src1_nrows, octx->n_threads, + htp_mm_hvx_vtcm_layout_build(&L, kparams->kernel_type, src0->type, ne10, act_nrows, octx->n_threads, 0, src0_row_size, src1_row_size, 0, kparams->n_prefetch, true, false); const size_t vtcm_size = L.total_bytes; @@ -3882,18 +3899,20 @@ static int hvx_mm_matmul_id( } uint8_t * const base = (uint8_t *) octx->ctx->vtcm_base; - mmctx->vtcm_src1 = VTCM_LAYOUT_PTR(uint8_t, base, L.off_src1); - mmctx->vtcm_src0 = VTCM_LAYOUT_PTR(uint8_t, base, L.off_src0); - mmctx->vtcm_src2 = NULL; - mmctx->vtcm_dst = VTCM_LAYOUT_PTR(uint8_t, base, L.off_dst); + mmctx->vtcm_src1 = VTCM_LAYOUT_PTR(uint8_t, base, L.off_src1); + mmctx->vtcm_src0 = VTCM_LAYOUT_PTR(uint8_t, base, L.off_src0); + mmctx->vtcm_src2 = NULL; + mmctx->vtcm_dst = VTCM_LAYOUT_PTR(uint8_t, base, L.off_dst); + mmctx->vtcm_act_raw = VTCM_LAYOUT_PTR(uint8_t, base, L.off_act_raw); octx->src1_spad.src = NULL; octx->src0_spad.src = NULL; octx->src2_spad.src = NULL; octx->dst_spad.src = NULL; - mmctx->vtcm_src0_stride = src0_row_size_padded; - mmctx->vtcm_src1_stride = src1_row_size; + mmctx->vtcm_src0_stride = src0_row_size_padded; + mmctx->vtcm_src1_stride = src1_row_size; + mmctx->vtcm_act_raw_stride = hex_round_up(ne10 * sizeof(float), QK_Q8_0_TILED * sizeof(float)); mmctx->vtcm_src0_size_per_thread = fastdiv(L.src0_bytes, &octx->n_threads_div); mmctx->vtcm_src1_size_per_thread = L.src1_bytes; @@ -3901,16 +3920,17 @@ static int hvx_mm_matmul_id( mmctx->vtcm_dst_size_per_thread = fastdiv(L.dst_bytes, &octx->n_threads_div); mmctx->cur_m_start = 0; - mmctx->cur_m_rows = src1_nrows; - - mmctx->n_quant_rows_per_thread = (src1_nrows + n_quant_tasks - 1) / n_quant_tasks; - mmctx->quant_task_func = quant_task_func; - mmctx->n_quant_tasks = n_quant_tasks; - atomic_init(&mmctx->quant_barrier, n_quant_tasks); + mmctx->cur_m_rows = act_nrows; htp_trace_event_stop(tr, HTP_TRACE_EVT_INIT, 0); - worker_pool_run_func(octx->ctx->worker_pool, hvx_mmid_task_func, mmctx, octx->n_threads); + hvx_mm_transfer_src1_dma(octx, kparams, src1, mmctx->vtcm_act_raw, mmctx->vtcm_act_raw_stride, 0, act_nrows); + + mmctx->n_quant_rows_per_thread = (act_nrows + n_quant_tasks - 1) / n_quant_tasks; + mmctx->n_quant_tasks = n_quant_tasks; + work_queue_run(octx->ctx->work_queue, quant_task_func, mmctx, n_quant_tasks); + + work_queue_run(octx->ctx->work_queue, hvx_mmid_task_func, mmctx, octx->n_threads); return HTP_STATUS_OK; } @@ -3976,7 +3996,7 @@ static int hvx_mm_matmul_id_nx( work_queue_func_t hvx_mmid_task_func ) { const uint32_t src0_row_size_padded = mmctx->src0_row_size_padded; - const uint32_t src1_nrows = mmctx->src1_nrows; + const uint32_t act_nrows = mmctx->act_nrows; struct htp_thread_trace * tr = &octx->ctx->trace[0]; htp_trace_event_start(tr, HTP_TRACE_EVT_INIT, 0); @@ -3990,13 +4010,13 @@ static int hvx_mm_matmul_id_nx( const uint32_t qk = QK_Q8_0_TILED; const uint32_t nb = (act->ne[0] + qk - 1) / qk; - const uint32_t total_nb = src1_nrows * nb; + const uint32_t total_nb = act_nrows * nb; work_queue_func_t quant_task_func; uint32_t n_quant_tasks = 1; - if (src1_nrows < octx->n_threads) { + if (act_nrows < octx->n_threads) { n_quant_tasks = MIN(total_nb, octx->n_threads); - quant_task_func = (src0->type == HTP_TYPE_Q4_1 || src0->type == HTP_TYPE_Q4_K) ? quantize_f32_q8_1_tiled_block : quantize_f32_q8_0_tiled_block; + quant_task_func = htp_mm_weight_has_offset(src0->type) ? quantize_f32_q8_1_tiled_block : quantize_f32_q8_0_tiled_block; for (uint32_t ith = 0; ith < n_quant_tasks; ++ith) { uint32_t ib_first = (total_nb * ith) / n_quant_tasks; uint32_t ib_last = (total_nb * (ith + 1)) / n_quant_tasks; @@ -4006,13 +4026,13 @@ static int hvx_mm_matmul_id_nx( mmctx->quant_c[ith] = ib_first % nb; } } else { - n_quant_tasks = MIN(src1_nrows, octx->n_threads); - quant_task_func = (src0->type == HTP_TYPE_Q4_1 || src0->type == HTP_TYPE_Q4_K) ? quantize_f32_q8_1_tiled : quantize_f32_q8_0_tiled; + n_quant_tasks = MIN(act_nrows, octx->n_threads); + quant_task_func = htp_mm_weight_has_offset(src0->type) ? quantize_f32_q8_1_tiled : quantize_f32_q8_0_tiled; } - size_t src1_row_size = (src0->type == HTP_TYPE_Q4_1 || src0->type == HTP_TYPE_Q4_K) ? htp_mm_q8_1_tiled_row_size(act->ne[0]) : htp_mm_q8_0_tiled_row_size(act->ne[0]); + size_t src1_row_size = htp_mm_weight_has_offset(src0->type) ? htp_mm_q8_1_tiled_row_size(act->ne[0]) : htp_mm_q8_0_tiled_row_size(act->ne[0]); struct htp_mm_hvx_vtcm_layout L; - htp_mm_hvx_vtcm_layout_build(&L, kparams->kernel_type, src0->type, act->ne[0], src1_nrows, octx->n_threads, + htp_mm_hvx_vtcm_layout_build(&L, kparams->kernel_type, src0->type, act->ne[0], act_nrows, octx->n_threads, 0, src0_row_size, src1_row_size, 0, kparams->n_prefetch, true, false); const size_t vtcm_size = L.total_bytes; @@ -4024,9 +4044,10 @@ static int hvx_mm_matmul_id_nx( } uint8_t * const base = (uint8_t *) octx->ctx->vtcm_base; - mmctx->vtcm_src0 = VTCM_LAYOUT_PTR(uint8_t, base, L.off_src0); - mmctx->vtcm_src1 = VTCM_LAYOUT_PTR(uint8_t, base, L.off_src1); - mmctx->vtcm_dst = VTCM_LAYOUT_PTR(uint8_t, base, L.off_dst); + mmctx->vtcm_src0 = VTCM_LAYOUT_PTR(uint8_t, base, L.off_src0); + mmctx->vtcm_src1 = VTCM_LAYOUT_PTR(uint8_t, base, L.off_src1); + mmctx->vtcm_dst = VTCM_LAYOUT_PTR(uint8_t, base, L.off_dst); + mmctx->vtcm_act_raw = VTCM_LAYOUT_PTR(uint8_t, base, L.off_act_raw); octx->src0_spad.src = NULL; octx->src1_spad.src = NULL; @@ -4034,29 +4055,31 @@ static int hvx_mm_matmul_id_nx( octx->src3_spad.src = NULL; octx->dst_spad.src = NULL; - mmctx->vtcm_src0_stride = 0; - mmctx->vtcm_src1_stride = src1_row_size; + mmctx->vtcm_src0_stride = 0; + mmctx->vtcm_src1_stride = src1_row_size; + mmctx->vtcm_act_raw_stride = hex_round_up(act->ne[0] * sizeof(float), QK_Q8_0_TILED * sizeof(float)); mmctx->vtcm_src0_size_per_thread = fastdiv(L.src0_bytes, &octx->n_threads_div); mmctx->vtcm_src1_size_per_thread = L.src1_bytes; mmctx->vtcm_dst_size_per_thread = fastdiv(L.dst_bytes, &octx->n_threads_div); mmctx->cur_m_start = 0; - mmctx->cur_m_rows = src1_nrows; - - mmctx->n_quant_rows_per_thread = (src1_nrows + n_quant_tasks - 1) / n_quant_tasks; - mmctx->quant_task_func = quant_task_func; - mmctx->n_quant_tasks = n_quant_tasks; - atomic_init(&mmctx->quant_barrier, n_quant_tasks); + mmctx->cur_m_rows = act_nrows; FARF(HIGH, "matmul-id-nx: src0 %d:%d:%d type %s nrows %u, src1 %d:%d:%d nrows %u, vtcm %zu/%zu, threads %d\n", src0->ne[0], src0->ne[1], src0->ne[2], mmctx->type, src0->ne[1], - act->ne[0], act->ne[1], act->ne[2], src1_nrows, + act->ne[0], act->ne[1], act->ne[2], act_nrows, L.total_bytes, octx->ctx->vtcm_size, octx->n_threads); htp_trace_event_stop(tr, HTP_TRACE_EVT_INIT, 0); - worker_pool_run_func(octx->ctx->worker_pool, hvx_mmid_task_func, mmctx, octx->n_threads); + hvx_mm_transfer_src1_dma(octx, kparams, act, mmctx->vtcm_act_raw, mmctx->vtcm_act_raw_stride, 0, act_nrows); + + mmctx->n_quant_rows_per_thread = (act_nrows + n_quant_tasks - 1) / n_quant_tasks; + mmctx->n_quant_tasks = n_quant_tasks; + work_queue_run(octx->ctx->work_queue, quant_task_func, mmctx, n_quant_tasks); + + work_queue_run(octx->ctx->work_queue, hvx_mmid_task_func, mmctx, octx->n_threads); return HTP_STATUS_OK; } @@ -4202,7 +4225,7 @@ int op_matmul_id(struct htp_ops_context * octx) { mmctx->mapping_stride = mapping_stride; mmctx->mm_div_ne11 = kparams->div_ne1; mmctx->src0_row_size_padded = src0_row_size_padded; - mmctx->src1_nrows = src1_nrows; + mmctx->act_nrows = src1_nrows; mmctx->cur_m_start = 0; mmctx->cur_m_rows = src1_nrows; @@ -4281,7 +4304,7 @@ int op_matmul_id_nx(struct htp_ops_context * octx) { const size_t src0_row_size = src0->nb[1]; const size_t src0_row_size_padded = hex_round_up(src0_row_size, 128); - const uint32_t src1_nrows = act->ne[1] * act->ne[2] * act->ne[3]; + const uint32_t act_nrows = act->ne[1] * act->ne[2] * act->ne[3]; const int n_ids = ids->ne[0]; const int n_as = src0->ne[2]; @@ -4291,7 +4314,7 @@ int op_matmul_id_nx(struct htp_ops_context * octx) { uint32_t * matrix_row_counts = (uint32_t *) mapping_buf; struct mmid_row_mapping * matrix_rows = NULL; - if (src1_nrows > 1) { + if (act_nrows > 1) { const size_t matrix_row_counts_size = n_as * sizeof(uint32_t); assert(octx->ctx->ddr_spad_size >= matrix_row_counts_size); @@ -4325,9 +4348,9 @@ int op_matmul_id_nx(struct htp_ops_context * octx) { mmctx->mapping_stride = mapping_stride; mmctx->mm_div_ne11 = kparams->div_ne1; mmctx->src0_row_size_padded = src0_row_size_padded; - mmctx->src1_nrows = src1_nrows; + mmctx->act_nrows = act_nrows; mmctx->cur_m_start = 0; - mmctx->cur_m_rows = src1_nrows; + mmctx->cur_m_rows = act_nrows; htp_trace_event_stop(tr, HTP_TRACE_EVT_INIT, 0); @@ -4336,7 +4359,7 @@ int op_matmul_id_nx(struct htp_ops_context * octx) { s = hmx_mm_op_matmul_id_nx(octx, mmctx); } else { if (hvx_mm_init_vec_dot(mmctx, src0->type) == 0) { - s = hvx_mm_matmul_id_nx(octx, mmctx, src1_nrows > 1 ? hvx_mm_id_nx : hvx_mv_id_nx); + s = hvx_mm_matmul_id_nx(octx, mmctx, act_nrows > 1 ? hvx_mm_id_nx : hvx_mv_id_nx); } else { s = HTP_STATUS_NO_SUPPORT; } @@ -4370,17 +4393,18 @@ int op_matmul_nx(struct htp_ops_context * octx) { bool is_repacked = (src0->type == HTP_TYPE_Q4_0 || src0->type == HTP_TYPE_Q4_1 || src0->type == HTP_TYPE_Q8_0 || src0->type == HTP_TYPE_IQ4_NL || - src0->type == HTP_TYPE_MXFP4 || src0->type == HTP_TYPE_Q4_K); + src0->type == HTP_TYPE_MXFP4 || src0->type == HTP_TYPE_Q4_K || + src0->type == HTP_TYPE_Q5_K); struct htp_mm_context mmctx_struct = {0}; struct htp_mm_context * mmctx = &mmctx_struct; mmctx->octx = octx; mmctx->act = act; - const uint32_t src1_nrows = act->ne[1] * act->ne[2] * act->ne[3]; - mmctx->src1_nrows = src1_nrows; + const uint32_t act_nrows = act->ne[1] * act->ne[2] * act->ne[3]; + mmctx->act_nrows = act_nrows; mmctx->cur_m_start = 0; - mmctx->cur_m_rows = src1_nrows; + mmctx->cur_m_rows = act_nrows; const size_t src0_row_size = src0->nb[1]; const size_t src0_row_size_padded = hex_round_up(src0_row_size, 128); @@ -4391,13 +4415,13 @@ int op_matmul_nx(struct htp_ops_context * octx) { const uint32_t qk = QK_Q8_0_TILED; const uint32_t nb = (act->ne[0] + qk - 1) / qk; - const uint32_t total_nb = src1_nrows * nb; + const uint32_t total_nb = act_nrows * nb; worker_callback_t quant_task_func; uint32_t n_quant_tasks = 1; - if (src1_nrows < octx->n_threads) { + if (act_nrows < octx->n_threads) { n_quant_tasks = MIN(total_nb, octx->n_threads); - quant_task_func = (src0->type == HTP_TYPE_Q4_1 || src0->type == HTP_TYPE_Q4_K) ? quantize_f32_q8_1_tiled_block : quantize_f32_q8_0_tiled_block; + quant_task_func = htp_mm_weight_has_offset(src0->type) ? quantize_f32_q8_1_tiled_block : quantize_f32_q8_0_tiled_block; for (uint32_t ith = 0; ith < n_quant_tasks; ++ith) { uint32_t ib_first = (total_nb * ith) / n_quant_tasks; uint32_t ib_last = (total_nb * (ith + 1)) / n_quant_tasks; @@ -4407,16 +4431,16 @@ int op_matmul_nx(struct htp_ops_context * octx) { mmctx->quant_c[ith] = ib_first % nb; } } else { - n_quant_tasks = MIN(src1_nrows, octx->n_threads); - quant_task_func = (src0->type == HTP_TYPE_Q4_1 || src0->type == HTP_TYPE_Q4_K) ? quantize_f32_q8_1_tiled : quantize_f32_q8_0_tiled; + n_quant_tasks = MIN(act_nrows, octx->n_threads); + quant_task_func = htp_mm_weight_has_offset(src0->type) ? quantize_f32_q8_1_tiled : quantize_f32_q8_0_tiled; } - const size_t src1_row_size = (src0->type == HTP_TYPE_Q4_1 || src0->type == HTP_TYPE_Q4_K) + const size_t src1_row_size = htp_mm_weight_has_offset(src0->type) ? htp_mm_q8_1_tiled_row_size(act->ne[0]) : htp_mm_q8_0_tiled_row_size(act->ne[0]); struct htp_mm_hvx_vtcm_layout L; - htp_mm_hvx_vtcm_layout_build(&L, kparams->kernel_type, src0->type, act->ne[0], src1_nrows, octx->n_threads, + htp_mm_hvx_vtcm_layout_build(&L, kparams->kernel_type, src0->type, act->ne[0], act_nrows, octx->n_threads, 0, src0_row_size, src1_row_size, 0, kparams->n_prefetch, false, true); const size_t vtcm_size = L.total_bytes; @@ -4428,9 +4452,10 @@ int op_matmul_nx(struct htp_ops_context * octx) { } uint8_t * const base = (uint8_t *) octx->ctx->vtcm_base; - mmctx->vtcm_src0 = VTCM_LAYOUT_PTR(uint8_t, base, L.off_src0); - mmctx->vtcm_src1 = VTCM_LAYOUT_PTR(uint8_t, base, L.off_src1); - mmctx->vtcm_dst = VTCM_LAYOUT_PTR(uint8_t, base, L.off_dst); + mmctx->vtcm_src0 = VTCM_LAYOUT_PTR(uint8_t, base, L.off_src0); + mmctx->vtcm_src1 = VTCM_LAYOUT_PTR(uint8_t, base, L.off_src1); + mmctx->vtcm_dst = VTCM_LAYOUT_PTR(uint8_t, base, L.off_dst); + mmctx->vtcm_act_raw = VTCM_LAYOUT_PTR(uint8_t, base, L.off_act_raw); octx->src0_spad.src = NULL; octx->src1_spad.src = NULL; @@ -4438,18 +4463,14 @@ int op_matmul_nx(struct htp_ops_context * octx) { octx->src3_spad.src = NULL; octx->dst_spad.src = NULL; - mmctx->vtcm_src0_stride = is_repacked ? 0 : src0_row_size_padded; - mmctx->vtcm_src1_stride = src1_row_size; + mmctx->vtcm_src0_stride = is_repacked ? 0 : src0_row_size_padded; + mmctx->vtcm_src1_stride = src1_row_size; + mmctx->vtcm_act_raw_stride = hex_round_up(act->ne[0] * sizeof(float), QK_Q8_0_TILED * sizeof(float)); mmctx->vtcm_src0_size_per_thread = fastdiv(L.src0_bytes, &octx->n_threads_div); mmctx->vtcm_src1_size_per_thread = L.src1_bytes; mmctx->vtcm_dst_size_per_thread = fastdiv(L.dst_bytes, &octx->n_threads_div); - mmctx->n_quant_rows_per_thread = (src1_nrows + n_quant_tasks - 1) / n_quant_tasks; - mmctx->quant_task_func = quant_task_func; - mmctx->n_quant_tasks = n_quant_tasks; - atomic_init(&mmctx->quant_barrier, n_quant_tasks); - // Run fused matmul const uint32_t n_matmul_jobs = octx->n_threads; worker_callback_t matmul_job_func; @@ -4458,10 +4479,13 @@ int op_matmul_nx(struct htp_ops_context * octx) { case HTP_TYPE_Q4_0: matmul_job_func = hvx_mm_nx_2d_repacked_q4_0; break; case HTP_TYPE_Q4_1: case HTP_TYPE_Q4_K: matmul_job_func = hvx_mm_nx_2d_repacked_q4_1; break; + case HTP_TYPE_Q5_K: matmul_job_func = hvx_mm_nx_2d_repacked_q5_k; break; case HTP_TYPE_Q8_0: matmul_job_func = hvx_mm_nx_2d_repacked_q8_0; break; case HTP_TYPE_IQ4_NL: matmul_job_func = hvx_mm_nx_2d_repacked_iq4nl; break; case HTP_TYPE_MXFP4: matmul_job_func = hvx_mm_nx_2d_repacked_mxfp4; break; - default: return HTP_STATUS_NO_SUPPORT; + default: + htp_trace_event_stop(tr, HTP_TRACE_EVT_INIT, 0); + return HTP_STATUS_NO_SUPPORT; } } else { matmul_job_func = hvx_mm_nx_2d; @@ -4469,7 +4493,13 @@ int op_matmul_nx(struct htp_ops_context * octx) { htp_trace_event_stop(tr, HTP_TRACE_EVT_INIT, 0); - worker_pool_run_func(octx->ctx->worker_pool, matmul_job_func, mmctx, n_matmul_jobs); + hvx_mm_transfer_src1_dma(octx, kparams, act, mmctx->vtcm_act_raw, mmctx->vtcm_act_raw_stride, 0, act_nrows); + + mmctx->n_quant_rows_per_thread = (act_nrows + n_quant_tasks - 1) / n_quant_tasks; + mmctx->n_quant_tasks = n_quant_tasks; + work_queue_run(octx->ctx->work_queue, quant_task_func, mmctx, n_quant_tasks); + + work_queue_run(octx->ctx->work_queue, matmul_job_func, mmctx, n_matmul_jobs); return HTP_STATUS_OK; } diff --git a/ggml/src/ggml-hexagon/htp/matmul-ops.h b/ggml/src/ggml-hexagon/htp/matmul-ops.h index fe9dbb61c..cfb3bfbbd 100644 --- a/ggml/src/ggml-hexagon/htp/matmul-ops.h +++ b/ggml/src/ggml-hexagon/htp/matmul-ops.h @@ -25,6 +25,9 @@ extern "C" { #define HTP_MM_WEIGHT_TILE_SIZE_Q8_0 1088 #define HTP_MM_WEIGHT_TILE_SIZE_IQ4_NL 576 #define HTP_MM_WEIGHT_TILE_SIZE_MXFP4 544 +// Q5_K: the Q4_1 tile (640) followed by a 128-byte plane with the 5th bit of every quant, transposed so that +// plane byte l holds the eight flags of lane l: bit 2i = low nibble of nibble vector i, bit 2i+1 = high nibble +#define HTP_MM_WEIGHT_TILE_SIZE_Q5_K 768 // Q6_K native 6-bit tile (32 rows x 32 k), vrmpy-ready: byte 4*row+b of a vector holds k = 4*group+b // vectors 0..3: low nibbles, vector i holds group 2i (low nibble) and group 2i+1 (high nibble) // vectors 4..5: high 2 bits, vector m holds groups 4m..4m+3 at bit offsets 0,2,4,6 @@ -37,6 +40,7 @@ extern "C" { #define HTP_MM_WEIGHT_ALIGNED_TILE_SIZE_Q8_0 1152 #define HTP_MM_WEIGHT_ALIGNED_TILE_SIZE_IQ4_NL 640 #define HTP_MM_WEIGHT_ALIGNED_TILE_SIZE_MXFP4 640 +#define HTP_MM_WEIGHT_ALIGNED_TILE_SIZE_Q5_K 768 #define HTP_MM_WEIGHT_ALIGNED_TILE_SIZE_Q6_K 896 // --- Activation Tiled Block Sizes (including padding) --- @@ -199,6 +203,8 @@ static inline uint32_t htp_mm_get_weight_tile_size(int weight_type) { return HTP_MM_WEIGHT_TILE_SIZE_Q4_1; case HTP_TYPE_Q8_0: return HTP_MM_WEIGHT_TILE_SIZE_Q8_0; + case HTP_TYPE_Q5_K: + return HTP_MM_WEIGHT_TILE_SIZE_Q5_K; case HTP_TYPE_Q6_K: return HTP_MM_WEIGHT_TILE_SIZE_Q6_K; case HTP_TYPE_MXFP4: @@ -218,6 +224,8 @@ static inline uint32_t htp_mm_get_weight_aligned_tile_size(int weight_type) { return HTP_MM_WEIGHT_ALIGNED_TILE_SIZE_Q4_1; case HTP_TYPE_Q8_0: return HTP_MM_WEIGHT_ALIGNED_TILE_SIZE_Q8_0; + case HTP_TYPE_Q5_K: + return HTP_MM_WEIGHT_ALIGNED_TILE_SIZE_Q5_K; case HTP_TYPE_Q6_K: return HTP_MM_WEIGHT_ALIGNED_TILE_SIZE_Q6_K; case HTP_TYPE_MXFP4: @@ -227,6 +235,11 @@ static inline uint32_t htp_mm_get_weight_aligned_tile_size(int weight_type) { } } +// weight types whose tiles carry a per-block offset (x = d * q + m): the activations need block sums (q8_1) +static inline bool htp_mm_weight_has_offset(int weight_type) { + return weight_type == HTP_TYPE_Q4_1 || weight_type == HTP_TYPE_Q4_K || weight_type == HTP_TYPE_Q5_K; +} + // --- Activation/Row Size Helpers --- static inline size_t htp_mm_q8_0_tiled_row_size(uint32_t ne) { const uint32_t ne_padded = ((ne + 127) / 128) * 128; @@ -248,6 +261,7 @@ static inline size_t htp_mm_get_tiled_row_stride(int weight_type, uint32_t k) { case HTP_TYPE_Q4_1: case HTP_TYPE_Q4_K: case HTP_TYPE_Q8_0: + case HTP_TYPE_Q5_K: case HTP_TYPE_Q6_K: case HTP_TYPE_MXFP4: return (size_t) nb * htp_mm_get_weight_tile_size(weight_type); @@ -332,6 +346,7 @@ struct htp_mm_hvx_vtcm_layout { size_t off_src2; // vtcm_src2 (Wq / fused only) size_t off_src3; // vtcm_src3 (Wv / fused only) size_t off_dst; // vtcm_dst (output scratch) + size_t off_act_raw; // vtcm_act_raw (raw activation DMA staging) // Cached sizes size_t src0_bytes; @@ -339,6 +354,7 @@ struct htp_mm_hvx_vtcm_layout { size_t src2_bytes; size_t src3_bytes; size_t dst_bytes; + size_t act_raw_bytes; size_t total_bytes; }; @@ -351,7 +367,6 @@ static inline void htp_mm_hmx_vtcm_layout_build( size_t mc, size_t nc, uint32_t group_size, - bool use_dma_activation, bool pipeline, uint32_t act_threads, uint32_t aligned_tile_size, @@ -366,8 +381,7 @@ static inline void htp_mm_hmx_vtcm_layout_build( const size_t activation_area_size = hex_align_up(group_size * act_head_stride * sizeof(uint16_t), HTP_MM_HMX_TILE_SIZE); const size_t output_area_size = hex_align_up(group_size * mc * nc * sizeof(uint16_t), HTP_MM_HMX_TILE_SIZE); const size_t scratch_area_size = hex_align_up(nc * vec_dot_size, HTP_MM_HMX_TILE_SIZE); - const size_t min_f32_size = use_dma_activation - ? hex_align_up(act_threads * HTP_MM_DMA_ACT_MULTIPLIER * k * sizeof(float), 128) : 0; + const size_t min_f32_size = hex_align_up(act_threads * HTP_MM_DMA_ACT_MULTIPLIER * k * sizeof(float), 128); // Group A: Permanent activation tiles and scales size_t off_group_a = 0; @@ -388,10 +402,9 @@ static inline void htp_mm_hmx_vtcm_layout_build( // Group C: Activation prep temporary buffer (overlaps Group B, starting at off_group_a) const size_t max_f32_size = act_threads * 64 * k * sizeof(float); - const size_t act_f32_size = use_dma_activation - ? hex_align_up(hex_smin(max_f32_size, hex_smax(min_f32_size, group_b_size)), 128) : 0; + const size_t act_f32_size = hex_align_up(hex_smin(max_f32_size, hex_smax(min_f32_size, group_b_size)), 128); size_t off_group_c = off_group_a; - VTCM_LAYOUT_ALLOC_OPTIONAL(off_group_c, off_act_f32, act_f32_size, use_dma_activation); + VTCM_LAYOUT_ALLOC(off_group_c, off_act_f32, act_f32_size); const size_t group_c_size = off_group_c - off_group_a; @@ -478,20 +491,20 @@ static inline void htp_mm_hvx_vtcm_layout_build( bool is_fused_nx ) { (void)src1_row_size; - size_t src0_sz = 0; - size_t src1_sz = 0; - size_t src2_sz = src2_row_size > 0 ? htp_mm_round_up(src2_row_size, 128) : 0; - size_t src3_sz = 0; - size_t dst_sz = 0; + size_t src0_sz = 0; + size_t src1_sz = 0; + size_t src2_sz = src2_row_size > 0 ? htp_mm_round_up(src2_row_size, 128) : 0; + size_t src3_sz = 0; + size_t dst_sz = 0; + size_t act_raw_sz = 0; const bool is_repack = (wtype == HTP_TYPE_Q4_0 || wtype == HTP_TYPE_Q4_1 || wtype == HTP_TYPE_Q8_0 || wtype == HTP_TYPE_IQ4_NL || wtype == HTP_TYPE_MXFP4 || wtype == HTP_TYPE_Q6_K || - wtype == HTP_TYPE_Q4_K); + wtype == HTP_TYPE_Q4_K || wtype == HTP_TYPE_Q5_K); if (is_fused_nx) { const size_t src0_row_size_padded = hex_round_up(src0_row_size, 128); - const size_t quant_scratch_size = hex_round_up(ne10 * sizeof(float), QK_Q8_0_TILED * sizeof(float)) * n_threads; size_t weight_sz_per_thread = 0; @@ -505,17 +518,19 @@ static inline void htp_mm_hvx_vtcm_layout_build( weight_sz_per_thread = hex_round_up(n_prefetch * src0_row_size_padded, 128); } - size_t tiled_act_row_size = (wtype == HTP_TYPE_Q4_1 || wtype == HTP_TYPE_Q4_K) ? htp_mm_q8_1_tiled_row_size(ne10) : htp_mm_q8_0_tiled_row_size(ne10); + size_t tiled_act_row_size = htp_mm_weight_has_offset(wtype) ? htp_mm_q8_1_tiled_row_size(ne10) : htp_mm_q8_0_tiled_row_size(ne10); size_t act_sz = hex_round_up(tiled_act_row_size * src1_nrows, 128); + size_t raw_row_size = hex_round_up(ne10 * sizeof(float), QK_Q8_0_TILED * sizeof(float)); - src0_sz = weight_sz_per_thread * n_threads; // shared single-weight prefetch buffer - src1_sz = act_sz; // quantized activation buffer - src2_sz = 0; - src3_sz = 0; - dst_sz = quant_scratch_size; + src0_sz = weight_sz_per_thread * n_threads; // shared single-weight prefetch buffer + src1_sz = act_sz; // quantized activation buffer + src2_sz = 0; + src3_sz = 0; + dst_sz = 0; + act_raw_sz = hex_round_up(raw_row_size * src1_nrows, 128); } else if (is_matmul_id) { const size_t src0_row_size_padded = htp_mm_round_up(src0_row_size, 128); - const size_t src1_row_size_tiled = (wtype == HTP_TYPE_Q4_1 || wtype == HTP_TYPE_Q4_K) ? htp_mm_q8_1_tiled_row_size(ne10) + const size_t src1_row_size_tiled = htp_mm_weight_has_offset(wtype) ? htp_mm_q8_1_tiled_row_size(ne10) : htp_mm_q8_0_tiled_row_size(ne10); size_t src0_sz_per_thread = htp_mm_round_up(n_prefetch * src0_row_size_padded, 256); @@ -529,10 +544,13 @@ static inline void htp_mm_hvx_vtcm_layout_build( src0_sz_per_thread = repacked_vtcm_size; } - src0_sz = src0_sz_per_thread * n_threads; - dst_sz = htp_mm_round_up(ne10 * sizeof(float), QK_Q8_0_TILED * sizeof(float)) * n_threads; - src2_sz = 0; - src3_sz = 0; + size_t raw_row_size = hex_round_up(ne10 * sizeof(float), QK_Q8_0_TILED * sizeof(float)); + + src0_sz = src0_sz_per_thread * n_threads; + dst_sz = 0; + src2_sz = 0; + src3_sz = 0; + act_raw_sz = hex_round_up(raw_row_size * src1_nrows, 128); } else { const size_t src0_row_size_padded = htp_mm_round_up(src0_row_size, 128); const size_t dst_nrows = (src1_nrows > 1) ? 0 : 1; @@ -540,21 +558,23 @@ static inline void htp_mm_hvx_vtcm_layout_build( switch (kernel_type) { case HTP_MM_KERNEL_HVX_F16_F16_VTCM: { size_t f16_src1_row_size = htp_mm_round_up(ne10 * 2, 128); - src1_sz = htp_mm_round_up(f16_src1_row_size * src1_nrows, 256); - src0_sz = htp_mm_round_up(n_prefetch * src0_row_size_padded, 256) * n_threads; - dst_sz = dst_nrows > 0 ? htp_mm_round_up(dst_row_size, 128) * n_threads : 0; + src1_sz = htp_mm_round_up(f16_src1_row_size * src1_nrows, 256); + src0_sz = htp_mm_round_up(n_prefetch * src0_row_size_padded, 256) * n_threads; + dst_sz = dst_nrows > 0 ? htp_mm_round_up(dst_row_size, 128) * n_threads : 0; + act_raw_sz = hex_round_up(hex_round_up(ne10 * sizeof(float), 128) * src1_nrows, 128); break; } case HTP_MM_KERNEL_HVX_F32_F32_VTCM: { size_t f32_src1_row_size = htp_mm_round_up(ne10 * 4, 128); - src1_sz = htp_mm_round_up(f32_src1_row_size * src1_nrows, 256); - src0_sz = htp_mm_round_up(n_prefetch * src0_row_size_padded, 256) * n_threads; - dst_sz = dst_nrows > 0 ? htp_mm_round_up(dst_row_size, 128) * n_threads : 0; + src1_sz = htp_mm_round_up(f32_src1_row_size * src1_nrows, 256); + src0_sz = htp_mm_round_up(n_prefetch * src0_row_size_padded, 256) * n_threads; + dst_sz = dst_nrows > 0 ? htp_mm_round_up(dst_row_size, 128) * n_threads : 0; + act_raw_sz = 0; break; } case HTP_MM_KERNEL_HVX_QUANT_BLOCK: case HTP_MM_KERNEL_HVX_QUANT_ROW: { - size_t q_src1_row_size = (wtype == HTP_TYPE_Q4_1 || wtype == HTP_TYPE_Q4_K) ? htp_mm_q8_1_tiled_row_size(ne10) : htp_mm_q8_0_tiled_row_size(ne10); + size_t q_src1_row_size = htp_mm_weight_has_offset(wtype) ? htp_mm_q8_1_tiled_row_size(ne10) : htp_mm_q8_0_tiled_row_size(ne10); src0_sz = htp_mm_round_up(n_prefetch * src0_row_size_padded, 256); src1_sz = htp_mm_round_up(q_src1_row_size * src1_nrows, 256); @@ -569,10 +589,10 @@ static inline void htp_mm_hvx_vtcm_layout_build( src0_sz = repacked_vtcm_size * n_threads; } - size_t quant_scratch_size_per_thread = htp_mm_round_up(ne10 * sizeof(float), QK_Q8_0_TILED * sizeof(float)); size_t dst_slice_per_thread = (dst_nrows > 0 && src1_nrows == 1) ? htp_mm_round_up((dst_row_size + n_threads - 1) / n_threads, 128) : 0; - size_t dst_size_per_thread = (dst_slice_per_thread > quant_scratch_size_per_thread) ? dst_slice_per_thread : quant_scratch_size_per_thread; - dst_sz = dst_size_per_thread * n_threads; + dst_sz = dst_slice_per_thread * n_threads; + size_t raw_row_size = hex_round_up(ne10 * sizeof(float), QK_Q8_0_TILED * sizeof(float)); + act_raw_sz = hex_round_up(raw_row_size * src1_nrows, 128); break; } default: @@ -580,19 +600,30 @@ static inline void htp_mm_hvx_vtcm_layout_build( } } - size_t off = 0; - VTCM_LAYOUT_ALLOC(off, off_src0, src0_sz); - VTCM_LAYOUT_ALLOC(off, off_src1, src1_sz); - VTCM_LAYOUT_ALLOC(off, off_src2, src2_sz); - VTCM_LAYOUT_ALLOC(off, off_src3, src3_sz); - VTCM_LAYOUT_ALLOC(off, off_dst, dst_sz); + // Group A: Persistent buffers across chunk compute + size_t off_group_a = 0; + VTCM_LAYOUT_ALLOC(off_group_a, off_src1, src1_sz); + VTCM_LAYOUT_ALLOC_OPTIONAL(off_group_a, off_src2, src2_sz, src2_sz > 0); + VTCM_LAYOUT_ALLOC_OPTIONAL(off_group_a, off_src3, src3_sz, src3_sz > 0); - L->src0_bytes = src0_sz; - L->src1_bytes = src1_sz; - L->src2_bytes = src2_sz; - L->src3_bytes = src3_sz; - L->dst_bytes = dst_sz; - L->total_bytes = off; + // Group B: Compute-only buffers (starts at off_group_a) + size_t off_group_b = off_group_a; + VTCM_LAYOUT_ALLOC(off_group_b, off_src0, src0_sz); + VTCM_LAYOUT_ALLOC_OPTIONAL(off_group_b, off_dst, dst_sz, dst_sz > 0); + const size_t group_b_size = off_group_b - off_group_a; + + // Group C: Raw activation staging buffer (overlaps Group B, starts at off_group_a) + size_t off_group_c = off_group_a; + VTCM_LAYOUT_ALLOC_OPTIONAL(off_group_c, off_act_raw, act_raw_sz, act_raw_sz > 0); + const size_t group_c_size = off_group_c - off_group_a; + + L->src0_bytes = src0_sz; + L->src1_bytes = src1_sz; + L->src2_bytes = src2_sz; + L->src3_bytes = src3_sz; + L->dst_bytes = dst_sz; + L->act_raw_bytes = act_raw_sz; + L->total_bytes = off_group_a + hex_smax(group_b_size, group_c_size); } static inline bool htp_mm_hvx_solve_vtcm_params( @@ -630,7 +661,7 @@ static inline bool htp_mm_hvx_solve_vtcm_params( const size_t avail_act = vtcm_budget - fixed_bytes; size_t row_size = 0; if (kernel_type == HTP_MM_KERNEL_HVX_QUANT_ROW || kernel_type == HTP_MM_KERNEL_HVX_QUANT_BLOCK) { - row_size = (wtype == HTP_TYPE_Q4_1 || wtype == HTP_TYPE_Q4_K) + row_size = htp_mm_weight_has_offset(wtype) ? htp_mm_q8_1_tiled_row_size(ne10) : htp_mm_q8_0_tiled_row_size(ne10); } else if (kernel_type == HTP_MM_KERNEL_HVX_F16_F16_VTCM) { @@ -642,7 +673,12 @@ static inline bool htp_mm_hvx_solve_vtcm_params( return false; } - uint32_t m_chunk = (uint32_t) (avail_act / row_size); + size_t eff_row_size = row_size; + if (kernel_type == HTP_MM_KERNEL_HVX_QUANT_ROW || kernel_type == HTP_MM_KERNEL_HVX_QUANT_BLOCK) { + eff_row_size += hex_round_up(ne10 * sizeof(float), QK_Q8_0_TILED * sizeof(float)); + } + + uint32_t m_chunk = (uint32_t) (avail_act / eff_row_size); if (m_chunk > 1) { m_chunk &= ~1U; } @@ -679,15 +715,15 @@ static inline size_t htp_mm_hmx_get_2d_vtcm_size( int wtype, uint32_t k, size_t mc, size_t nc, bool pipeline, uint32_t act_threads, uint32_t aligned_tile_size, size_t src2_size ) { struct htp_mm_hmx_vtcm_layout L; - htp_mm_hmx_vtcm_layout_build(&L, HTP_MM_KERNEL_HMX_2D, wtype, k, mc, nc, 1, false, pipeline, act_threads, aligned_tile_size, src2_size); + htp_mm_hmx_vtcm_layout_build(&L, HTP_MM_KERNEL_HMX_2D, wtype, k, mc, nc, 1, pipeline, act_threads, aligned_tile_size, src2_size); return L.total_bytes; } static inline size_t htp_mm_hmx_get_batched_vtcm_size( - int wtype, uint32_t k, size_t mc, size_t nc, uint32_t group_size, bool use_dma_activation, bool pipeline, uint32_t act_threads, size_t src2_size) { + int wtype, uint32_t k, size_t mc, size_t nc, uint32_t group_size, bool pipeline, uint32_t act_threads, size_t src2_size) { (void)pipeline; struct htp_mm_hmx_vtcm_layout L; - htp_mm_hmx_vtcm_layout_build(&L, HTP_MM_KERNEL_HMX_F16_BATCHED, wtype, k, mc, nc, group_size, use_dma_activation, false, act_threads, 0, src2_size); + htp_mm_hmx_vtcm_layout_build(&L, HTP_MM_KERNEL_HMX_F16_BATCHED, wtype, k, mc, nc, group_size, false, act_threads, 0, src2_size); return L.total_bytes; } @@ -697,7 +733,6 @@ static inline bool htp_mm_hmx_solve_batched_params( uint32_t ne01_padded, uint32_t ne11, uint32_t group_size, - bool use_dma_activation, int n_threads, bool pipeline, size_t src2_size, @@ -726,7 +761,7 @@ static inline bool htp_mm_hmx_solve_batched_params( if (htp_mm_hmx_compute_chunks(vtcm_budget, group_overhead, group_size_per_n, group_size_per_m, group_size_per_mn, hex_align_up(ne11, 32), ne01_padded, (size_t) ne01_padded * HTP_MM_HMX_COST_W_DEQUANT, (size_t) ne11 * HTP_MM_HMX_COST_A_CONVERT, &m_chunk_candidate, &n_chunk_candidate, &vtcm_size_candidate) == 0) { - size_t exact_size = htp_mm_hmx_get_batched_vtcm_size(wtype, k, m_chunk_candidate, n_chunk_candidate, group_size, use_dma_activation, pipeline, act_threads, src2_size); + size_t exact_size = htp_mm_hmx_get_batched_vtcm_size(wtype, k, m_chunk_candidate, n_chunk_candidate, group_size, pipeline, act_threads, src2_size); if (exact_size <= vtcm_budget) { size_t mblocks = ((size_t) ne11 + m_chunk_candidate - 1) / m_chunk_candidate; if (mblocks < best_mblocks || (mblocks == best_mblocks && act_threads > best_act_threads)) { diff --git a/ggml/src/ggml-metal/ggml-metal-context.m b/ggml/src/ggml-metal/ggml-metal-context.m index 442ed2a07..e9064666f 100644 --- a/ggml/src/ggml-metal/ggml-metal-context.m +++ b/ggml/src/ggml-metal/ggml-metal-context.m @@ -477,6 +477,10 @@ enum ggml_status ggml_metal_graph_compute(ggml_metal_t ctx, struct ggml_cgraph * return GGML_STATUS_FAILED; } + if (gf->n_nodes == 0) { + return GGML_STATUS_SUCCESS; + } + // number of nodes encoded by the main thread (empirically determined) const int n_main = MAX(64, 0.1*gf->n_nodes); @@ -514,8 +518,6 @@ enum ggml_status ggml_metal_graph_compute(ggml_metal_t ctx, struct ggml_cgraph * const bool use_capture = ctx->capture_compute == 0; if (use_capture) { - ctx->capture_compute = -1; - // make sure all previous computations have finished before starting the capture if (ctx->cmd_buf_last) { [ctx->cmd_buf_last waitUntilCompleted]; @@ -538,7 +540,7 @@ enum ggml_status ggml_metal_graph_compute(ggml_metal_t ctx, struct ggml_cgraph * NSError * error = nil; if (![[MTLCaptureManager sharedCaptureManager] startCaptureWithDescriptor:descriptor error:&error]) { - GGML_LOG_ERROR("%s: error: unable to start capture '%s'\n", __func__, [[error localizedDescription] UTF8String]); + GGML_LOG_ERROR("%s: error: unable to start capture '%s' (did you set METAL_CAPTURE_ENABLED=1 ?)\n", __func__, [[error localizedDescription] UTF8String]); } else { [ctx->capture_scope beginScope]; ctx->capture_started = true; @@ -749,7 +751,7 @@ void ggml_metal_set_n_cb(ggml_metal_t ctx, int n_cb) { idx_start, idx_end, ctx->use_concurrency, - ctx->capture_compute, + ctx->capture_compute == 0, ctx->debug_graph); for (int idx = 0; idx < ggml_metal_op_n_nodes(ctx_op); ++idx) { diff --git a/ggml/src/ggml-metal/ggml-metal-ops.cpp b/ggml/src/ggml-metal/ggml-metal-ops.cpp index 708703e51..54d5b9808 100644 --- a/ggml/src/ggml-metal/ggml-metal-ops.cpp +++ b/ggml/src/ggml-metal/ggml-metal-ops.cpp @@ -3494,34 +3494,20 @@ int ggml_metal_op_flash_attn_ext(ggml_metal_op_t ctx, int idx) { const int is_q = !use_kv_f16 && ggml_is_quantized(op->src[1]->type) ? 1 : 0; - // 2*(2*ncpsg) - // ncpsg soft_max values + ncpsg mask values - // - // 16*32*(nsg) - // the shared memory needed for the simdgroups to load the KV cache - // each thread loads (dequantizes) 16 head elements, there are 32 threads in th SG - // -#define FATTN_SMEM(nsg) (GGML_PAD((nqptg*(ne00 + 2*GGML_PAD(ne20, 64) + 2*(2*ncpsg)) + is_q*(16*32*(nsg)))*(sizeof(float)/2), 16)) + // shared memory layout (halfs unless noted): + // queries/attn/result: Q*(DK + 2*PAD2(DV,64) + 4*C) + // quantized KV scratch: 16*32*NSG (only when is_q) + const int64_t dv_pad = GGML_PAD(ne20, 64); - //int64_t nsgmax = 4; - // - //if (is_q) { - // nsgmax = 2; - // while (true) { - // const size_t smem = FATTN_SMEM(nsgmax); - // if (smem > props_dev->max_theadgroup_memory_size) { - // break; - // } - // nsgmax *= 2; - // } - // nsgmax /= 2; - //} + auto fa_smem = [&](int32_t nsg) -> size_t { + const size_t smem_half = nqptg*(ne00 + 2*dv_pad + 4*ncpsg) + is_q*(16*32*nsg); + return GGML_PAD(smem_half*sizeof(ggml_fp16_t), 16); + }; // simdgroups per threadgroup (a.k.a. warps) - //nsg = ne01 <= nqptg ? MAX(4, MIN(nsgmax, MIN(ne11/ncpsg, (int64_t) pipeline.maxTotalThreadsPerThreadgroup/32))) : 4; int32_t nsg = ne00 >= 512 ? 8 : 4; - const size_t smem = FATTN_SMEM(nsg); + const size_t smem = fa_smem(nsg); const int32_t ns10 = nb11_attn/nb10_attn; const int32_t ns20 = nb21_attn/nb20_attn; @@ -3577,7 +3563,6 @@ int ggml_metal_op_flash_attn_ext(ggml_metal_op_t ctx, int idx) { ggml_metal_encoder_set_threadgroup_memory_size(enc, smem, 0); ggml_metal_encoder_dispatch_threadgroups(enc, (ne01 + nqptg - 1)/nqptg, ne02, ne03, 32, nsg, 1); -#undef FATTN_SMEM } else { // half4x4 kernel // sparse: the index lists are per query row, so a threadgroup can share KV with Q == 1 only @@ -3679,14 +3664,18 @@ int ggml_metal_op_flash_attn_ext(ggml_metal_op_t ctx, int idx) { // note: for simplicity assume the K is larger or equal than V GGML_ASSERT(ne10 >= ne20); - // ne00 + 2*ncpsg*(nsg) - // for each query, we load it as f16 in shared memory (ne00) - // and store the soft_max values and the mask - // - // ne20*(nsg) - // each simdgroup has a full f32 head vector in shared mem to accumulate results - // -#define FATTN_SMEM(nsg) (GGML_PAD(((GGML_PAD(ne00, 128) + 4*ncpsg + 2*GGML_PAD(ne20, 128))*(nsg)*nqptg)*(sizeof(float)/2), 16)) + // shared memory layout (halfs unless noted): + // queries: Q*NSG*PAD2(ne00, 128) + // attn + mask: NSG*4*Q*C + // results: 2*NSG*Q*PAD2(ne20, 128) + // sparse idx: NSG*C ints (only when use_sparse) + const int64_t dk_pad = GGML_PAD(ne00, 128); + const int64_t dv_pad = GGML_PAD(ne20, 128); + + auto fa_vec_smem = [&](int64_t nsg, int32_t nqptg) -> size_t { + const size_t smem_half = (size_t) (dk_pad + 4*ncpsg + 2*dv_pad)*nqptg*nsg; + return GGML_PAD(smem_half*sizeof(ggml_fp16_t) + (use_sparse ? (size_t) nsg*ncpsg*sizeof(int) : 0), 16); + }; int64_t nsg = 1; @@ -3722,7 +3711,7 @@ int ggml_metal_op_flash_attn_ext(ggml_metal_op_t ctx, int idx) { } // fall back to baseline (Q=1) if the tuned config exceeds threadgroup memory - if ((size_t) FATTN_SMEM(nsg) > props_dev->max_theadgroup_memory_size) { + if (fa_vec_smem(nsg, nqptg) > props_dev->max_theadgroup_memory_size) { cfg = ggml_metal_tuning::fa_vec_baseline_cfg((int) ne00, (int) ne20); nqptg = cfg.Q; // = 1 } @@ -3779,9 +3768,8 @@ int ggml_metal_op_flash_attn_ext(ggml_metal_op_t ctx, int idx) { ggml_metal_encoder_set_buffer (enc, bid_src4, 5); ggml_metal_encoder_set_buffer (enc, use_sparse ? bid_idx : bid_src0, 8); - const size_t smem = FATTN_SMEM(nsg); + const size_t smem = fa_vec_smem(nsg, nqptg); - //printf("smem: %zu, max: %zu, nsg = %d, nsgmax = %d\n", smem, props_dev->max_theadgroup_memory_size, (int) nsg, (int) nsgmax); GGML_ASSERT(smem <= props_dev->max_theadgroup_memory_size); if (nwg == 1) { @@ -3827,7 +3815,6 @@ int ggml_metal_op_flash_attn_ext(ggml_metal_op_t ctx, int idx) { ggml_metal_encoder_dispatch_threadgroups(enc, nrows, 1, 1, 32*nwg, 1, 1); } } -#undef FATTN_SMEM } return 1; diff --git a/ggml/src/ggml-metal/kernels/fa.metal b/ggml/src/ggml-metal/kernels/fa.metal index f26d493d5..cf23ef04f 100644 --- a/ggml/src/ggml-metal/kernels/fa.metal +++ b/ggml/src/ggml-metal/kernels/fa.metal @@ -1248,17 +1248,24 @@ kernel void kernel_flash_attn_ext_vec( constexpr short NL = NW/NE; // note: this can be adjusted to support different head sizes and simdgroup work loads constexpr short SH = 4*Q*C; // shared memory per simdgroup + const int SMEM_Q = Q*NSG*PK; + const int SMEM_S = NSG*SH; + const int SMEM_O = 2*NSG*Q*PV; + const int SMEM = SMEM_Q + SMEM_S + SMEM_O; + static_assert(DK4 % NL == 0, "DK4 must be divisible by NL"); static_assert(DV4 % NL == 0, "DV4 must be divisible by NL"); - //const short T = PK + NSG*SH; // shared memory size per query in (half) + threadgroup q4_t * sq4 = (threadgroup q4_t *) shmem_f16; // holds the query data + threadgroup s_t * ss = (threadgroup s_t *) (shmem_f16 + SMEM_Q + sgitg*SH); // scratch buffer for attention + threadgroup s4_t * ss4 = (threadgroup s4_t *) (shmem_f16 + SMEM_Q + sgitg*SH); // same as above but in s4_t + threadgroup half * sm = (threadgroup half *) (shmem_f16 + SMEM_Q + sgitg*SH + 2*Q*C); // scratch buffer for mask + threadgroup o4_t * so4 = (threadgroup o4_t *) (shmem_f16 + SMEM_Q + SMEM_S + 2*sgitg*Q*PV); // scratch buffer for the results - //threadgroup q_t * sq = (threadgroup q_t *) (shmem_f16 + 0*PK); // holds the query data - threadgroup q4_t * sq4 = (threadgroup q4_t *) (shmem_f16 + 0*PK); // same as above but in q4_t - threadgroup s_t * ss = (threadgroup s_t *) (shmem_f16 + sgitg*SH + Q*NSG*PK); // scratch buffer for attention - threadgroup s4_t * ss4 = (threadgroup s4_t *) (shmem_f16 + sgitg*SH + Q*NSG*PK); // same as above but in s4_t - threadgroup half * sm = (threadgroup half *) (shmem_f16 + sgitg*SH + 2*Q*C + Q*NSG*PK); // scratch buffer for mask - threadgroup o4_t * so4 = (threadgroup o4_t *) (shmem_f16 + 2*sgitg*Q*PV + Q*NSG*PK + NSG*SH); // scratch buffer for the results + // sparse indices for the current block + threadgroup int * spidx = FC_flash_attn_ext_vec_has_sparse + ? (threadgroup int *) (shmem_f16 + SMEM) + sgitg*C + : nullptr; // store the result for all queries in shared memory (the O matrix from the paper) so4 += tiisg; @@ -1388,10 +1395,20 @@ kernel void kernel_flash_attn_ext_vec( ic = 0; } + // load the sparse KV indices for the current block into shared memory + if (FC_flash_attn_ext_vec_has_sparse) { + FOR_UNROLL (short ii = 0; ii < C/NW; ++ii) { + const short i = ii*NW + tiisg; + + spidx[i] = pidx[ic + i]; + } + simdgroup_barrier(mem_flags::mem_threadgroup); + } + if (FC_flash_attn_ext_vec_has_mask) { if (FC_flash_attn_ext_vec_has_sparse) { FOR_UNROLL (short qq = 0; qq < Q; ++qq) { - const int i11 = pidx[ic + tiisg]; + const int i11 = spidx[tiisg]; if ((iq1*Q + qq) < args.ne01 && i11 >= 0) { sm[qq*C + tiisg] = pm[qq][i11]; } else { @@ -1449,7 +1466,7 @@ kernel void kernel_flash_attn_ext_vec( FOR_UNROLL (short cc = 0; cc < C/NE; ++cc) { if (FC_flash_attn_ext_vec_has_sparse) { // the KV rows are gathered from the index list; -1 entries are padding - const int i11 = pidx[ic + NE*cc + ty]; + const int i11 = spidx[NE*cc + ty]; if (i11 >= 0) { if (is_same::value) { device const k4_t * pk4s = (device const k4_t *) (k + i11*args.nb11) + tx; @@ -1593,7 +1610,7 @@ kernel void kernel_flash_attn_ext_vec( if (FC_flash_attn_ext_vec_has_sparse) { FOR_UNROLL (short cc = 0; cc < C/NE; ++cc) { // the KV rows are gathered from the index list; -1 entries are padding - const int i11 = pidx[ic + NE*cc + ty]; + const int i11 = spidx[NE*cc + ty]; if (i11 >= 0) { if (is_same::value) { device const v4_t * pv4 = (device const v4_t *) (v + i11*args.nb21); diff --git a/ggml/src/ggml-rpc/ggml-rpc.cpp b/ggml/src/ggml-rpc/ggml-rpc.cpp index c24caad77..353b79b07 100644 --- a/ggml/src/ggml-rpc/ggml-rpc.cpp +++ b/ggml/src/ggml-rpc/ggml-rpc.cpp @@ -858,6 +858,11 @@ static size_t ggml_backend_rpc_buffer_type_get_alloc_size(ggml_backend_buffer_ty if (rpc_get) { ggml_backend_rpc_buffer_type_context * buft_ctx = (ggml_backend_rpc_buffer_type_context *)buft->context; + // the reported size must never be below ggml_nbytes: rpc_tensor stores nb[] as uint32_t, + // so a stride over 4 GiB is truncated on the wire and the remote size comes back too small + // TODO: change rpc_tensor nb to 64-bit int + const size_t min_size = ggml_nbytes(tensor); + // Cache key for calls to read the alloc_size. // We deliberately exclude src tensor dimensions from the key because: // 1. For CPU backends, alloc_size = ggml_nbytes(output) regardless of src shapes @@ -871,6 +876,7 @@ static size_t ggml_backend_rpc_buffer_type_get_alloc_size(ggml_backend_buffer_ty uint32_t op; int32_t op_params[GGML_MAX_OP_PARAMS / sizeof(int32_t)]; uint32_t ne[GGML_MAX_DIMS]; + uint64_t nb[GGML_MAX_DIMS]; }; alloc_size_cache_key key = {}; @@ -880,6 +886,7 @@ static size_t ggml_backend_rpc_buffer_type_get_alloc_size(ggml_backend_buffer_ty memcpy(key.op_params, tensor->op_params, sizeof(key.op_params)); for (int i = 0; i < GGML_MAX_DIMS; i++) { key.ne[i] = (uint32_t)tensor->ne[i]; + key.nb[i] = (uint64_t)tensor->nb[i]; } uint64_t cache_hash = fnv_hash((const uint8_t *)&key, sizeof(key)); @@ -893,7 +900,7 @@ static size_t ggml_backend_rpc_buffer_type_get_alloc_size(ggml_backend_buffer_ty std::lock_guard lock(cache_mutex); auto it = cache.find(cache_hash); if (it != cache.end()) { - return it->second; + return std::max(it->second, min_size); } } @@ -915,7 +922,7 @@ static size_t ggml_backend_rpc_buffer_type_get_alloc_size(ggml_backend_buffer_ty cache[cache_hash] = response.alloc_size; } - return response.alloc_size; + return std::max(response.alloc_size, min_size); } return ggml_nbytes(tensor); diff --git a/ggml/src/ggml-sycl/fattn-sparse.cpp b/ggml/src/ggml-sycl/fattn-sparse.cpp new file mode 100644 index 000000000..df2837094 --- /dev/null +++ b/ggml/src/ggml-sycl/fattn-sparse.cpp @@ -0,0 +1,267 @@ +#include "fattn.hpp" +#include "fattn-sparse.hpp" + +#include +#include +#include + +static constexpr int64_t SPARSE_FA_PAD = 256; +static constexpr int64_t SPARSE_FA_MIN_RATIO = 2; + +extern int g_ggml_sycl_enable_sparse_fa; +extern int g_ggml_sycl_debug_sparse_fa; +extern int g_ggml_sycl_sparse_fa_margin; + +static int sparse_fa_enabled(void) { + return g_ggml_sycl_enable_sparse_fa; +} + +static int sparse_fa_debug(void) { + return g_ggml_sycl_debug_sparse_fa; +} + +// slack above n_kv_max; callers may exceed the hint by a few always-attended positions +static int sparse_fa_margin(void) { + return g_ggml_sycl_sparse_fa_margin; +} + +// Unordered output is fine: softmax over the selected set is permutation invariant. +static void sparse_fa_compact_mask(sycl::queue * stream, + const sycl::half * __restrict__ mask, + int32_t * __restrict__ indices, + int32_t * __restrict__ count, + const int64_t n_kv, + const int64_t n_kv_g) { + constexpr size_t WG = 256; + const size_t global = (size_t) GGML_PAD(n_kv, (int64_t) WG); + + stream->parallel_for( + sycl::nd_range<1>(sycl::range<1>(global), sycl::range<1>(WG)), + [=](sycl::nd_item<1> item) { + const int64_t i = (int64_t) item.get_global_id(0); + if (i >= n_kv || !sycl::isfinite((float) mask[i])) { + return; + } + + sycl::atomic_ref ctr(*count); + + const int32_t pos = ctr.fetch_add(1); + if (pos < (int32_t) n_kv_g) { + indices[pos] = (int32_t) i; + } + }); +} + +// Rows along ne[0] are contiguous for every type used as a KV cache, so this is +// a plain byte copy and needs no per-type code. Padding slots are zeroed. +static void sparse_fa_gather_rows(sycl::queue * stream, + const uint8_t * __restrict__ src, + uint8_t * __restrict__ dst, + const int32_t * __restrict__ indices, + const int32_t * __restrict__ count, + const size_t row_size, + const size_t src_nb1, + const size_t src_nb2, + const int64_t n_kv_g, + const int64_t n_head) { + GGML_ASSERT(row_size % sizeof(uint32_t) == 0); + const size_t words = row_size / sizeof(uint32_t); + + stream->parallel_for( + sycl::range<3>((size_t) n_head, (size_t) n_kv_g, words), + [=](sycl::id<3> id) { + const int64_t h = (int64_t) id[0]; + const int64_t slot = (int64_t) id[1]; + const size_t w = id[2]; + + uint32_t * dst_row = + (uint32_t *) (dst + ((size_t) (h * n_kv_g + slot)) * row_size); + + if (slot >= (int64_t) *count) { + dst_row[w] = 0; + return; + } + + const uint32_t * src_row = + (const uint32_t *) (src + (size_t) indices[slot] * src_nb1 + + (size_t) h * src_nb2); + dst_row[w] = src_row[w]; + }); +} + +static void sparse_fa_gather_mask(sycl::queue * stream, + const sycl::half * __restrict__ mask, + sycl::half * __restrict__ mask_g, + const int32_t * __restrict__ indices, + const int32_t * __restrict__ count, + const int64_t n_kv_g, + const int64_t n_rows, + const size_t mask_s1) { + stream->parallel_for( + sycl::range<2>((size_t) n_rows, (size_t) n_kv_g), + [=](sycl::id<2> id) { + const int64_t r = (int64_t) id[0]; + const int64_t slot = (int64_t) id[1]; + + sycl::half v = sycl::half(-INFINITY); + if (slot < (int64_t) *count) { + v = mask[(size_t) r * mask_s1 + (size_t) indices[slot]]; + } + mask_g[(size_t) r * n_kv_g + slot] = v; + }); +} + +static bool sparse_fa_applicable(const ggml_tensor * dst, int64_t & n_kv_g_out) { + const ggml_tensor * Q = dst->src[0]; + const ggml_tensor * K = dst->src[1]; + const ggml_tensor * V = dst->src[2]; + const ggml_tensor * mask = dst->src[3]; + + if (!Q || !K || !V || !mask) { + return false; + } + + const int32_t n_kv_max = ggml_get_op_params_i32(dst, 4); + if (n_kv_max <= 0) { + return false; + } + + float max_bias = 0.0f; + float logit_softcap = 0.0f; + memcpy(&max_bias, (const float *) dst->op_params + 1, sizeof(float)); + memcpy(&logit_softcap, (const float *) dst->op_params + 2, sizeof(float)); + if (max_bias != 0.0f || logit_softcap != 0.0f) { + return false; + } + + // single-token decode only; prefill amortises the scan already + if (Q->ne[1] != 1) { + return false; + } + if (K->ne[3] != 1 || V->ne[3] != 1 || mask->ne[2] != 1 || mask->ne[3] != 1) { + return false; + } + if (mask->type != GGML_TYPE_F16 || mask->ne[0] < K->ne[1]) { + return false; + } + if (K->ne[2] != V->ne[2]) { + return false; + } + + // nb[1] may stride over heads (interleaved cache); only ne[0] must be contiguous + if (K->nb[0] != ggml_type_size(K->type) || V->nb[0] != ggml_type_size(V->type)) { + return false; + } + + const size_t k_row = ggml_row_size(K->type, K->ne[0]); + const size_t v_row = ggml_row_size(V->type, V->ne[0]); + if (k_row % sizeof(uint32_t) || v_row % sizeof(uint32_t)) { + return false; + } + + const int64_t n_kv_g = GGML_PAD((int64_t) n_kv_max + sparse_fa_margin(), SPARSE_FA_PAD); + if (n_kv_g * SPARSE_FA_MIN_RATIO > K->ne[1]) { + return false; + } + + n_kv_g_out = n_kv_g; + return true; +} + +bool ggml_sycl_flash_attn_ext_sparse(ggml_backend_sycl_context & ctx, ggml_tensor * dst) { + int64_t n_kv_g = 0; + if (!sparse_fa_enabled() || !sparse_fa_applicable(dst, n_kv_g)) { + return false; + } + + ggml_tensor * K = dst->src[1]; + ggml_tensor * V = dst->src[2]; + ggml_tensor * mask = dst->src[3]; + + const int64_t n_kv = K->ne[1]; + const int64_t n_head_k = K->ne[2]; + const int64_t n_rows_m = mask->ne[1]; + + const size_t k_row = ggml_row_size(K->type, K->ne[0]); + const size_t v_row = ggml_row_size(V->type, V->ne[0]); + + dpct::queue_ptr stream = ctx.stream(); + + ggml_sycl_pool_alloc idx_alloc(ctx.pool(), (size_t) n_kv_g); + ggml_sycl_pool_alloc cnt_alloc(ctx.pool(), 1); + ggml_sycl_pool_alloc k_alloc(ctx.pool(), (size_t) n_head_k * n_kv_g * k_row); + ggml_sycl_pool_alloc v_alloc(ctx.pool(), (size_t) n_head_k * n_kv_g * v_row); + ggml_sycl_pool_alloc m_alloc(ctx.pool(), (size_t) n_rows_m * n_kv_g); + + int32_t * d_idx = idx_alloc.get(); + int32_t * d_cnt = cnt_alloc.get(); + uint8_t * d_K = k_alloc.get(); + uint8_t * d_V = v_alloc.get(); + sycl::half * d_mask = m_alloc.get(); + + SYCL_CHECK(CHECK_TRY_ERROR(stream->memset(d_cnt, 0, sizeof(int32_t)))); + + sparse_fa_compact_mask(stream, (const sycl::half *) mask->data, + d_idx, d_cnt, n_kv, n_kv_g); + + sparse_fa_gather_rows(stream, (const uint8_t *) K->data, d_K, d_idx, d_cnt, + k_row, K->nb[1], K->nb[2], n_kv_g, n_head_k); + + sparse_fa_gather_rows(stream, (const uint8_t *) V->data, d_V, d_idx, d_cnt, + v_row, V->nb[1], V->nb[2], n_kv_g, n_head_k); + + sparse_fa_gather_mask(stream, (const sycl::half *) mask->data, d_mask, + d_idx, d_cnt, n_kv_g, n_rows_m, + mask->nb[1] / sizeof(sycl::half)); + + if (sparse_fa_debug()) { + int32_t h_cnt = 0; + SYCL_CHECK(CHECK_TRY_ERROR(stream->memcpy(&h_cnt, d_cnt, sizeof(int32_t)))); + SYCL_CHECK(CHECK_TRY_ERROR(stream->wait())); + fprintf(stderr, "[FA-SPARSE] n_kv=%lld n_kv_max=%d n_kv_g=%lld finite=%d%s\n", + (long long) n_kv, ggml_get_op_params_i32(dst, 4), + (long long) n_kv_g, (int) h_cnt, + h_cnt > (int32_t) n_kv_g ? " OVERFLOW" : ""); + } + + // shallow copies retargeted at the gathered buffers; kernels are unchanged + ggml_tensor K_g = *K; + K_g.data = d_K; + K_g.ne[1] = n_kv_g; + K_g.nb[1] = k_row; + K_g.nb[2] = (size_t) n_kv_g * k_row; + K_g.nb[3] = (size_t) n_head_k * n_kv_g * k_row; + K_g.view_src = nullptr; + K_g.view_offs = 0; + + ggml_tensor V_g = *V; + V_g.data = d_V; + V_g.ne[1] = n_kv_g; + V_g.nb[1] = v_row; + V_g.nb[2] = (size_t) n_kv_g * v_row; + V_g.nb[3] = (size_t) V->ne[2] * n_kv_g * v_row; + V_g.view_src = nullptr; + V_g.view_offs = 0; + + ggml_tensor M_g = *mask; + M_g.data = d_mask; + M_g.ne[0] = n_kv_g; + M_g.nb[1] = (size_t) n_kv_g * sizeof(sycl::half); + M_g.nb[2] = M_g.nb[1] * mask->ne[1]; + M_g.nb[3] = M_g.nb[2]; + M_g.view_src = nullptr; + M_g.view_offs = 0; + + ggml_tensor dst_g = *dst; + dst_g.src[1] = &K_g; + dst_g.src[2] = &V_g; + dst_g.src[3] = &M_g; + dst_g.op_params[4] = 0; // avoid re-entering this path + + ggml_sycl_flash_attn_ext(ctx, &dst_g); + + return true; +} diff --git a/ggml/src/ggml-sycl/fattn-sparse.hpp b/ggml/src/ggml-sycl/fattn-sparse.hpp new file mode 100644 index 000000000..98b06bddc --- /dev/null +++ b/ggml/src/ggml-sycl/fattn-sparse.hpp @@ -0,0 +1,10 @@ +#ifndef GGML_SYCL_FATTN_SPARSE_HPP +#define GGML_SYCL_FATTN_SPARSE_HPP + +#include "common.hpp" + +// Gather the K/V rows selected by a sparse mask and re-dispatch the dense +// kernels onto them. Returns false if the caller should use the dense path. +bool ggml_sycl_flash_attn_ext_sparse(ggml_backend_sycl_context & ctx, ggml_tensor * dst); + +#endif // GGML_SYCL_FATTN_SPARSE_HPP diff --git a/ggml/src/ggml-sycl/fattn.cpp b/ggml/src/ggml-sycl/fattn.cpp index 394cda593..541ae8a82 100644 --- a/ggml/src/ggml-sycl/fattn.cpp +++ b/ggml/src/ggml-sycl/fattn.cpp @@ -19,7 +19,7 @@ #include "fattn-vec.hpp" #include "fattn.hpp" #include "fattn-onednn.hpp" - +#include "fattn-sparse.hpp" #define FATTN_VEC_CASE(D, type_K, type_V) \ { \ @@ -276,6 +276,11 @@ static best_fattn_kernel ggml_sycl_get_best_fattn_kernel(const int device, const void ggml_sycl_flash_attn_ext(ggml_backend_sycl_context & ctx, ggml_tensor * dst) { ggml_sycl_set_device(ctx.device); + // sparse nodes are gathered down to n_kv_max rows and re-dispatched here + if (ggml_sycl_flash_attn_ext_sparse(ctx, dst)) { + return; + } + // n_kv watchdog: log when n_kv differs from the last FA call with // the same D — helps detect cache-truncation issues. static int nkv_debug = ggml_sycl_get_env("GGML_SYCL_MKL_FA_DEBUG", 0); diff --git a/ggml/src/ggml-sycl/ggml-sycl.cpp b/ggml/src/ggml-sycl/ggml-sycl.cpp index e13ec8525..99ebdb3e1 100644 --- a/ggml/src/ggml-sycl/ggml-sycl.cpp +++ b/ggml/src/ggml-sycl/ggml-sycl.cpp @@ -113,6 +113,9 @@ int g_ggml_sycl_usm_system = 0; int g_ggml_sycl_enable_host_pinned_mem = 1; int g_ggml_sycl_host_pinned_mem_2g = 0; int g_ggml_sycl_get_mem_api = MEMORY_API_TYPE_LEVEL_ZERO; +int g_ggml_sycl_enable_sparse_fa = 0; +int g_ggml_sycl_debug_sparse_fa = 0; +int g_ggml_sycl_sparse_fa_margin = 256; static ggml_sycl_device_info ggml_sycl_init() { GGML_SYCL_DEBUG("[SYCL] call ggml_sycl_init\n"); @@ -384,6 +387,10 @@ static void ggml_check_sycl() try { g_ggml_sycl_host_pinned_mem_2g = ggml_sycl_get_env("GGML_SYCL_HOST_PINNED_MEM_2G", 0) & g_ggml_sycl_enable_host_pinned_mem; + g_ggml_sycl_enable_sparse_fa = ggml_sycl_get_env("GGML_SYCL_SPARSE_FA", 0); + g_ggml_sycl_debug_sparse_fa = ggml_sycl_get_env("GGML_SYCL_SPARSE_FA_DEBUG", 0); + g_ggml_sycl_sparse_fa_margin = ggml_sycl_get_env("GGML_SYCL_SPARSE_FA_MARGIN", 256); + GGML_SYCL_DEBUG("[SYCL] call ggml_check_sycl\n"); GGML_LOG_INFO("Build with Macros:\n"); @@ -422,6 +429,7 @@ static void ggml_check_sycl() try { GGML_LOG_INFO(" GGML_SYCL_SUPPORT_VMM: no\n"); #endif + //Print the running environment variables for SYCL backend GGML_LOG_INFO("Running with Environment Variables:\n"); GGML_LOG_INFO(" GGML_SYCL_DEBUG: %d\n", g_ggml_sycl_debug); GGML_LOG_INFO(" GGML_SYCL_DEV_DEBUG: %d\n", g_ggml_sycl_dev_debug); @@ -491,6 +499,10 @@ static void ggml_check_sycl() try { GGML_LOG_INFO(" GGML_SYCL_ENABLE_HOST_PINNED_MEM: %d\n", g_ggml_sycl_enable_host_pinned_mem); GGML_LOG_INFO(" GGML_SYCL_HOST_PINNED_MEM_2G: %d\n", g_ggml_sycl_host_pinned_mem_2g); + GGML_LOG_INFO(" GGML_SYCL_SPARSE_FA: %d\n", g_ggml_sycl_enable_sparse_fa); + GGML_LOG_INFO(" GGML_SYCL_SPARSE_FA_DEBUG: %d\n", g_ggml_sycl_debug_sparse_fa); + GGML_LOG_INFO(" GGML_SYCL_SPARSE_FA_MARGIN: %d\n", g_ggml_sycl_sparse_fa_margin); + /* NOT REMOVE, keep it for next optimize for XMX. #if defined(SYCL_USE_XMX) fprintf(stderr, "%s: SYCL_USE_XMX: yes\n", __func__); diff --git a/ggml/src/ggml.c b/ggml/src/ggml.c index 6bcaee4cd..a085952f6 100644 --- a/ggml/src/ggml.c +++ b/ggml/src/ggml.c @@ -7465,7 +7465,7 @@ static void * incr_ptr_aligned(void ** p, size_t size, size_t align) { static size_t ggml_graph_nbytes(size_t size, bool grads) { size_t hash_size = ggml_hash_size(size * 2); - void * p = 0; + void * p = (char *) 1024; // workaround for ubsan error "applying non-zero offset X to null pointer" incr_ptr_aligned(&p, sizeof(struct ggml_cgraph), 1); incr_ptr_aligned(&p, size * sizeof(struct ggml_tensor *), sizeof(struct ggml_tensor *)); // nodes incr_ptr_aligned(&p, size * sizeof(struct ggml_tensor *), sizeof(struct ggml_tensor *)); // leafs @@ -7478,7 +7478,7 @@ static size_t ggml_graph_nbytes(size_t size, bool grads) { incr_ptr_aligned(&p, ggml_bitset_size(hash_size) * sizeof(ggml_bitset_t), sizeof(ggml_bitset_t)); size_t nbytes = (size_t) p; - return nbytes; + return nbytes - 1024; } size_t ggml_graph_overhead_custom(size_t size, bool grads) { diff --git a/include/llama-cpp.h b/include/llama-cpp.h index 8f6368177..880a6a5fa 100644 --- a/include/llama-cpp.h +++ b/include/llama-cpp.h @@ -24,7 +24,12 @@ struct llama_adapter_lora_deleter { void operator()(llama_adapter_lora * adapter) { llama_adapter_lora_free(adapter); } }; +struct llama_batch_ext_deleter { + void operator()(llama_batch_ext * batch) { llama_batch_ext_free(batch); } +}; + typedef std::unique_ptr llama_model_ptr; typedef std::unique_ptr llama_context_ptr; typedef std::unique_ptr llama_sampler_ptr; typedef std::unique_ptr llama_adapter_lora_ptr; +typedef std::unique_ptr llama_batch_ext_ptr; diff --git a/include/llama.h b/include/llama.h index cb177ece5..9758122f8 100644 --- a/include/llama.h +++ b/include/llama.h @@ -294,6 +294,11 @@ extern "C" { LLAMA_MODEL_META_KEY_SAMPLING_MIROSTAT_ETA, }; + enum llama_process_type { + LLAMA_PROCESS_TYPE_ENCODE, + LLAMA_PROCESS_TYPE_DECODE, + }; + struct llama_model_kv_override { enum llama_model_kv_override_type tag; @@ -1006,6 +1011,91 @@ extern "C" { struct llama_context * ctx, struct llama_batch batch); + // + // Extended batch API + // + + struct llama_batch_ext; + + struct llama_embd { + const float * data; + size_t n_rows; // number of embedding rows in data + size_t n_embd; // size of one row + }; + + LLAMA_API struct llama_batch_ext * llama_batch_ext_init (struct llama_context * ctx); + LLAMA_API void llama_batch_ext_free (struct llama_batch_ext * batch); + LLAMA_API void llama_batch_ext_clear(struct llama_batch_ext * batch); + + // Add an input token to the batch, with default values: + // id = LLAMA_TOKEN_NULL + // embd = nullptr + // pos = not set, the caller must set it with llama_batch_ext_set_pos() + // Returns the batch index (>= 0) + // On error: + // -1: batch is full + // -2: token is invalid (id == LLAMA_TOKEN_NULL or invalid embd) + // -3: invalid sequence id + LLAMA_API int32_t llama_batch_ext_add (struct llama_batch_ext * batch, llama_seq_id seq_id); + + // Add an input token to the batch, with a specified token ID or token embedding + LLAMA_API int32_t llama_batch_ext_add_token(struct llama_batch_ext * batch, llama_seq_id seq_id, llama_token id); + LLAMA_API int32_t llama_batch_ext_add_embd (struct llama_batch_ext * batch, llama_seq_id seq_id, struct llama_embd embd); + + // Add the token at index idx in the batch to another sequence id. The position will stays the same. + // Note: this should be called before other _set() functions + LLAMA_API bool llama_batch_ext_add_seq( + struct llama_batch_ext * batch, + int32_t idx, + llama_seq_id seq_id); + + // Set the token embedding for the token at index idx in the batch + // use it after llama_batch_ext_add_token() to have an entry with both a token id and an embedding + LLAMA_API bool llama_batch_ext_set_embd_token( + struct llama_batch_ext * batch, + int32_t idx, + struct llama_embd embd); + + // Set the "state" embedding for the token at index idx in the batch + // "state" here means extra hidden state carried over from a previous stage, e.g.: + // - MTP: state from N layers of the target model + // - Qwen3 VL (deepstack): state from N layers of the vision encoder + LLAMA_API bool llama_batch_ext_set_embd_state( + struct llama_batch_ext * batch, + int32_t idx, + struct llama_embd embd); + + // Set if output embeddings should be available for the token at index idx in the batch + // Note: for now, this is equivalent to setting the output logits + LLAMA_API bool llama_batch_ext_set_output_embd( + struct llama_batch_ext * batch, + int32_t idx, + bool value); + + // Set output logits for the token at index idx in the batch + // Note: for now, this is equivalent to setting the output embd + LLAMA_API bool llama_batch_ext_set_output_logits( + struct llama_batch_ext * batch, + int32_t idx, + bool value); + + // Set custom position for the token at index idx in the batch + // For M-RoPE models: + // - Embedding tokens must have multiple positions per token + // - Text token only requires one single position per token + LLAMA_API bool llama_batch_ext_set_pos( + struct llama_batch_ext * batch, + int32_t idx, + const llama_pos * pos); + + // TODO: implement get_embeddings() and get_logits() for llama_batch_ext + + // Return values are the same as llama_decode() + LLAMA_API int32_t llama_process( + struct llama_context * ctx, + enum llama_process_type type, + struct llama_batch_ext * batch); + // Set the number of threads used for decoding // n_threads is the number of threads used for generation (single token) // n_threads_batch is the number of threads used for prompt and batch processing (multiple tokens) diff --git a/scripts/sync-ggml.last b/scripts/sync-ggml.last index 8f4afc721..35cd376db 100644 --- a/scripts/sync-ggml.last +++ b/scripts/sync-ggml.last @@ -1 +1 @@ -e565a8f4ce2e462c4973a24c51098dd3c81c0256 +353b63b439f27ab2cc19dac97ab1681ba6d2d084 diff --git a/src/llama-batch.cpp b/src/llama-batch.cpp index 2b98a552f..89a1f3f37 100644 --- a/src/llama-batch.cpp +++ b/src/llama-batch.cpp @@ -3,6 +3,9 @@ #include "llama-impl.h" #include "llama-vocab.h" #include "llama-memory.h" +#include "llama-hparams.h" +#include "llama-model.h" +#include "llama-context.h" #include #include @@ -23,43 +26,120 @@ llama_batch_allocr::llama_batch_allocr(uint32_t n_pos_per_embd) : n_pos_per_embd } bool llama_batch_allocr::init( - const llama_batch & batch_inp, + const llama_batch_ext & batch_inp, const llama_vocab & vocab, - const llama_memory_i * memory, - uint32_t n_embd, - uint32_t n_seq_max, bool output_all) { clear(); - batch = batch_inp; + this->vocab = &vocab; + this->n_embd = batch_inp.n_embd > 0 ? batch_inp.n_embd : batch_inp.n_embd_inp; + this->n_seq_max = batch_inp.n_seq_max; - this->vocab = &vocab; + const int32_t n_tok = (int32_t) batch_inp.tokens.size(); - GGML_ASSERT(batch.n_tokens > 0); + GGML_ASSERT(n_tok > 0); - // - // validate input batch - // - - if (n_seq_max > LLAMA_MAX_SEQ) { + if ((uint32_t) n_seq_max > LLAMA_MAX_SEQ) { LLAMA_LOG_ERROR("%s: n_seq_max = %d > %d\n", __func__, n_seq_max, LLAMA_MAX_SEQ); return false; } - if (batch.token) { - for (int32_t i = 0; i < batch.n_tokens; ++i) { - if (batch.token[i] < 0 || (uint32_t) batch.token[i] >= vocab.n_tokens()) { - LLAMA_LOG_ERROR("%s: invalid token[%d] = %d\n", __func__, i, batch.token[i]); + const llama_memory_i * mem = batch_inp.mem; + + // + // determine the content types of the batch + // an entry can carry a token id, a token embedding, or both (e.g. MTP hook batches) + // all entries must carry the same combination + // + + const bool has_token = batch_inp.tokens[0].id != LLAMA_TOKEN_NULL; + const bool has_embd = batch_inp.tokens[0].has_embd; + + for (int32_t i = 1; i < n_tok; ++i) { + if ((batch_inp.tokens[i].id != LLAMA_TOKEN_NULL) != has_token || + batch_inp.tokens[i].has_embd != has_embd) { + LLAMA_LOG_ERROR("%s: all entries in the batch must have the same content types\n", __func__); + return false; + } + } + + if (!has_token && !has_embd) { + LLAMA_LOG_ERROR("%s: batch has neither token ids nor embeddings\n", __func__); + return false; + } + + // + // build flat token/embd array + // + + if (has_token) { + token_vec.resize(n_tok); + for (int32_t i = 0; i < n_tok; ++i) { + const llama_token id = batch_inp.tokens[i].id; + if (id < 0 || id >= batch_inp.n_vocab) { + LLAMA_LOG_ERROR("%s: invalid token[%d] = %d\n", __func__, i, id); return false; } + token_vec[i] = id; + } + } + + if (has_embd) { + embd_vec = batch_inp.embd; + } + + // + // build flat pos array + // token batch: pos[i] = tokens[i].pos[0] + // embedding batch: pos[j*n_tok + i] = tokens[i].pos[j] (section-major) + // + + { + const int32_t n_pos_total = has_token ? n_tok : n_tok * (int32_t) n_pos_per_embd; + pos.resize(n_pos_total); + if (has_token) { + for (int32_t i = 0; i < n_tok; ++i) { + pos[i] = batch_inp.tokens[i].pos[0]; + } + } else { + for (int32_t i = 0; i < n_tok; ++i) { + for (uint32_t j = 0; j < n_pos_per_embd; ++j) { + pos[(int32_t) j * n_tok + i] = batch_inp.tokens[i].pos[j]; + } + } } } - if (batch.seq_id) { - for (int32_t i = 0; i < batch.n_tokens; ++i) { - for (int32_t s = 0; s < batch.n_seq_id[i]; ++s) { - if (batch.seq_id && (batch.seq_id[i][s] < 0 || batch.seq_id[i][s] >= (llama_seq_id) n_seq_max)) { - LLAMA_LOG_ERROR("%s: invalid seq_id[%d][%d] = %d >= %d\n", __func__, i, s, batch.seq_id[i][s], (llama_seq_id) n_seq_max); + // + // build n_seq_id / seq_id arrays + // + + n_seq_id.resize(n_tok); + seq_id.resize(n_tok + 1); + seq_id[n_tok] = nullptr; + + { + size_t total = 0; + for (int32_t i = 0; i < n_tok; ++i) { + total += batch_inp.tokens[i].seq_ids.size(); + } + seq_id_data.reserve(total); + + for (int32_t i = 0; i < n_tok; ++i) { + for (auto sid : batch_inp.tokens[i].seq_ids) { + seq_id_data.push_back(sid); + } + } + + size_t off = 0; + for (int32_t i = 0; i < n_tok; ++i) { + n_seq_id[i] = (int32_t) batch_inp.tokens[i].seq_ids.size(); + seq_id[i] = seq_id_data.data() + off; + off += n_seq_id[i]; + + for (int32_t s = 0; s < n_seq_id[i]; ++s) { + if (seq_id[i][s] < 0 || seq_id[i][s] >= (llama_seq_id) n_seq_max) { + LLAMA_LOG_ERROR("%s: invalid seq_id[%d][%d] = %d >= %d\n", __func__, i, s, seq_id[i][s], (llama_seq_id) n_seq_max); return false; } } @@ -67,91 +147,43 @@ bool llama_batch_allocr::init( } // - // auto-generate missing fields + // build output/logits array // - if (!batch.n_seq_id) { - n_seq_id.resize(batch.n_tokens); - for (int32_t i = 0; i < batch.n_tokens; i++) { - n_seq_id[i] = seq_id_0.size(); - } - batch.n_seq_id = n_seq_id.data(); - } - - if (!batch.seq_id) { - seq_id.resize(batch.n_tokens + 1); - seq_id[batch.n_tokens] = NULL; - for (int32_t i = 0; i < batch.n_tokens; i++) { - seq_id[i] = seq_id_0.data(); - } - batch.seq_id = seq_id.data(); - } - - if (!batch.pos) { - pos.resize(batch.n_tokens); - - // initialize the starting position for each sequence based on the positions in the memory - llama_pos p0[LLAMA_MAX_SEQ]; - for (uint32_t s = 0; s < n_seq_max; ++s) { - if (!memory) { - // if no memory -> start from 0 - p0[s] = 0; - } else { - p0[s] = memory->seq_pos_max(s) + 1; - } + { + output.resize(n_tok, 0); + for (int32_t i = 0; i < n_tok; ++i) { + output[i] = batch_inp.tokens[i].output ? 1 : 0; } - for (int32_t i = 0; i < batch.n_tokens; i++) { - const llama_seq_id seq_id = batch.seq_id[i][0]; - - pos[i] = p0[seq_id]; - - // update the starting position for all sequences that are assigned to the this token - for (int32_t s = 0; s < batch.n_seq_id[i]; ++s) { - const llama_seq_id seq_id = batch.seq_id[i][s]; - - p0[seq_id] = pos[i] + 1; - } - } - - batch.pos = pos.data(); - } - - if (!batch.logits) { if (output_all) { - // return the output for all tokens - output.resize(batch.n_tokens, true); - } else { - // return the output only for the last token - output.resize(batch.n_tokens, false); - output[output.size() - 1] = true; - } - - batch.logits = output.data(); - } else if (output_all) { - bool warn = false; - - for (int32_t i = 0; i < batch.n_tokens; ++i) { - if (batch.logits[i] == 0) { - warn = true; + bool warn = false; + for (int32_t i = 0; i < n_tok; ++i) { + if (!output[i]) { warn = true; break; } + } + if (warn) { + LLAMA_LOG_WARN("%s: embeddings required but some input tokens were not marked as outputs -> overriding\n", __func__); + std::fill(output.begin(), output.end(), 1); } } - - if (warn) { - LLAMA_LOG_WARN("%s: embeddings required but some input tokens were not marked as outputs -> overriding\n", __func__); - - output.resize(batch.n_tokens, true); - batch.logits = output.data(); - } } + // + // set up the internal llama_batch to point to our owned arrays + // + + batch.n_tokens = n_tok; + batch.token = has_token ? token_vec.data() : nullptr; + batch.embd = has_embd ? embd_vec.data() : nullptr; + batch.pos = pos.data(); + batch.n_seq_id = n_seq_id.data(); + batch.seq_id = seq_id.data(); + batch.logits = output.data(); + // // compute stats // - this->n_embd = n_embd; - this->n_seq_max = n_seq_max; - // count the outputs in this batch for (int32_t i = 0; i < batch.n_tokens; ++i) { n_outputs += batch.logits[i] != 0; @@ -259,7 +291,7 @@ bool llama_batch_allocr::init( continue; } - const llama_pos p0 = memory ? memory->seq_pos_max(s) : -1; + const llama_pos p0 = mem ? mem->seq_pos_max(s) : -1; if (batch.token) { if (p0 >= 0 && p0 >= seq_pos_min(s)) { @@ -292,7 +324,7 @@ bool llama_batch_allocr::init( continue; } - const llama_pos p0 = memory ? memory->seq_pos_max(s) : -1; + const llama_pos p0 = mem ? mem->seq_pos_max(s) : -1; if (p0 >= 0) { bool ok = true; @@ -320,12 +352,12 @@ bool llama_batch_allocr::init( } } - if (memory) { + if (mem) { for (uint32_t s0 = 0; s0 < n_seq_max; ++s0) { for (uint32_t s1 = 0; s1 < n_seq_max; ++s1) { if (seq_cpl[s0][s1]) { - if (memory->seq_pos_min(s0) != memory->seq_pos_min(s1) || - memory->seq_pos_max(s0) != memory->seq_pos_max(s1)) { + if (mem->seq_pos_min(s0) != mem->seq_pos_min(s1) || + mem->seq_pos_max(s0) != mem->seq_pos_max(s1)) { LLAMA_LOG_ERROR("%s: sequence %d is coupled to %d in the input batch, but have divereged\n", __func__, s0, s1); return false; } @@ -725,11 +757,14 @@ void llama_batch_allocr::clear() { batch = {}; - pos .clear(); - n_seq_id .clear(); - seq_id .clear(); - seq_id_unq.clear(); - output .clear(); + token_vec .clear(); + embd_vec .clear(); + seq_id_data .clear(); + pos .clear(); + n_seq_id .clear(); + seq_id .clear(); + seq_id_unq .clear(); + output .clear(); for (auto & cur : seq_pos) { cur.clear(); @@ -985,3 +1020,305 @@ void llama_batch_free(struct llama_batch batch) { } if (batch.logits) free(batch.logits); } + + +// llama_batch_ext + +size_t llama_batch_ext_select_n_embd_inp(llama_context_type ctx_type, llm_arch arch, const llama_hparams & hparams) { + if (ctx_type == LLAMA_CONTEXT_TYPE_MTP) { + return hparams.n_embd_out(); + } + if (arch == LLM_ARCH_DFLASH) { + return hparams.n_embd_inp_enc(); + } + return hparams.n_embd_inp(); +} + +llama_batch_ext::llama_batch_ext(llama_context * ctx) : + n_tokens_max(llama_n_batch(ctx)), + n_embd_inp(llama_batch_ext_select_n_embd_inp(ctx->get_cparams().ctx_type, llama_get_model(ctx)->arch, llama_get_model(ctx)->hparams)), + n_embd_inp_enc(llama_get_model(ctx)->hparams.n_embd_inp_enc()), + n_seq_max(llama_n_seq_max(ctx)), + mem(llama_get_memory(ctx)), + n_vocab(llama_vocab_n_tokens(llama_model_get_vocab(llama_get_model(ctx)))), + n_pos_per_embd(llama_get_model(ctx)->hparams.n_pos_per_embd()) { + clear(); +} + +llama_batch_ext::llama_batch_ext( + size_t n_tokens_max, + size_t n_embd_inp, + size_t n_embd_inp_enc, + llama_seq_id n_seq_max, + llama_memory_i * mem, + llama_token n_vocab, + size_t n_pos_per_embd) : + n_tokens_max(n_tokens_max), + n_embd_inp(n_embd_inp), + n_embd_inp_enc(n_embd_inp_enc), + n_seq_max(n_seq_max), + mem(mem), + n_vocab(n_vocab), + n_pos_per_embd(n_pos_per_embd) { + clear(); +} + +void llama_batch_ext::clear() { + tokens.clear(); + embd .clear(); + n_embd = 0; +} + +int32_t llama_batch_ext::add_token(llama_seq_id seq_id) { + if (tokens.size() >= n_tokens_max) { + return -1; // size limit reached + } + if (seq_id < 0 || seq_id >= n_seq_max) { + return -3; // invalid sequence id + } + + // position is left undefined; call set_token_pos() before decoding + token t; + t.seq_ids.insert(seq_id); + + tokens.push_back(t); + + return (int32_t)(tokens.size() - 1); +} + +bool llama_batch_ext::add_seq(int32_t idx, llama_seq_id seq_id) { + if (idx < 0 || idx >= (int32_t) tokens.size()) { + return false; + } + if (seq_id < 0 || seq_id >= n_seq_max) { + return false; + } + + token & t = tokens[idx]; + + t.seq_ids.insert(seq_id); + + return true; +} + +bool llama_batch_ext::set_token_id(int32_t idx, llama_token id) { + if (idx < 0 || idx >= (int32_t) tokens.size()) { + return false; + } + if (id < 0 || id >= n_vocab) { + return false; + } + tokens[idx].id = id; + return true; +} + +bool llama_batch_ext::set_token_embd(int32_t idx, llama_embd embd_in) { + if (idx < 0 || idx >= (int32_t) tokens.size()) { + return false; + } + if (!embd_in.data) { + return false; + } + + const size_t n_total = embd_in.n_rows * embd_in.n_embd; + if (n_embd == 0) { + if (n_total != n_embd_inp && n_total != n_embd_inp_enc) { + LLAMA_LOG_ERROR("%s: embedding size mismatch, got %zu rows x %zu = %zu, expected %zu or %zu\n", + __func__, embd_in.n_rows, embd_in.n_embd, n_total, n_embd_inp, n_embd_inp_enc); + return false; + } + n_embd = n_total; + } else if (n_total != n_embd) { + LLAMA_LOG_ERROR("%s: embedding size mismatch, got %zu rows x %zu = %zu, expected %zu\n", + __func__, embd_in.n_rows, embd_in.n_embd, n_total, n_embd); + return false; + } + + token & t = tokens[idx]; + + if (t.has_embd) { + LLAMA_LOG_ERROR("%s: embedding for token %d is already set\n", __func__, idx); + return false; + } + + t.has_embd = true; + t.embd_off = embd.size(); + embd.insert(embd.end(), embd_in.data, embd_in.data + n_total); + + return true; +} + +bool llama_batch_ext::set_token_pos(int32_t idx, const llama_pos * pos_in) { + if (idx < 0 || idx >= (int32_t) tokens.size()) { + return false; + } + if (!pos_in) { + return false; + } + + token & t = tokens[idx]; + + size_t n_pos = t.id != LLAMA_TOKEN_NULL ? 1 : n_pos_per_embd; + for (size_t i = 0; i < n_pos; ++i) { + t.pos[i] = pos_in[i]; + } + + return true; +} + +bool llama_batch_ext::set_output(int32_t idx, bool output_last) { + if (idx < 0 || idx >= (int32_t) tokens.size()) { + return false; + } + tokens[idx].output = output_last; + return true; +} + +// llama_batch_ext C API + +llama_batch_ext * llama_batch_ext_init(llama_context * ctx) { + return new llama_batch_ext(ctx); +} + +void llama_batch_ext_free(llama_batch_ext * batch) { + delete batch; +} + +void llama_batch_ext_clear(llama_batch_ext * batch) { + batch->clear(); +} + +int32_t llama_batch_ext_add(llama_batch_ext * batch, llama_seq_id seq_id) { + return batch->add_token(seq_id); +} + +int32_t llama_batch_ext_add_token(llama_batch_ext * batch, llama_seq_id seq_id, llama_token id) { + int32_t idx = batch->add_token(seq_id); + if (idx < 0) { + return idx; + } + if (!batch->set_token_id(idx, id)) { + return -2; + } + return idx; +} + +int32_t llama_batch_ext_add_embd(llama_batch_ext * batch, llama_seq_id seq_id, llama_embd embd) { + int32_t idx = batch->add_token(seq_id); + if (idx < 0) { + return idx; + } + if (!batch->set_token_embd(idx, embd)) { + return -2; + } + return idx; +} + +bool llama_batch_ext_add_seq(llama_batch_ext * batch, int32_t idx, llama_seq_id seq_id) { + return batch->add_seq(idx, seq_id); +} + +bool llama_batch_ext_set_pos(llama_batch_ext * batch, int32_t idx, const llama_pos * pos) { + return batch->set_token_pos(idx, pos); +} + +bool llama_batch_ext_set_embd_token(llama_batch_ext * batch, int32_t idx, llama_embd embd) { + return batch->set_token_embd(idx, embd); +} + +bool llama_batch_ext_set_embd_state(llama_batch_ext * batch, int32_t idx, llama_embd embd) { + // TODO + GGML_UNUSED(batch); + GGML_UNUSED(idx); + GGML_UNUSED(embd); + return false; +} + +bool llama_batch_ext_set_output_embd(llama_batch_ext * batch, int32_t idx, bool value) { + return batch->set_output(idx, value); +} + +bool llama_batch_ext_set_output_logits(llama_batch_ext * batch, int32_t idx, bool value) { + return batch->set_output(idx, value); +} + +// llama_batch_compat + +void llama_batch_compat::init(llama_batch_ext & dst, const llama_batch & batch_inp, size_t n_embd_row) { + llama_batch_ext * batch_ext = &dst; + + if (n_embd_row == 0) { + n_embd_row = batch_ext->n_embd_inp; + } + + // a batch can carry both, for example the MTP hook batches + const bool has_token = batch_inp.token != nullptr; + const bool has_embd = batch_inp.embd != nullptr; + + static const llama_seq_id default_seq_id = 0; + static const int32_t default_n_seq_id = 1; + + // auto-generates positions locally when batch_inp.pos is null, continuing from memory + std::vector pos_next(batch_ext->n_seq_max); + for (llama_seq_id s = 0; s < (llama_seq_id) batch_ext->n_seq_max; ++s) { + pos_next[s] = llama_memory_seq_pos_max(batch_ext->mem, s) + 1; // assume next pos + } + + for (int32_t i = 0; i < batch_inp.n_tokens; ++i) { + const int32_t n_sid = batch_inp.n_seq_id ? batch_inp.n_seq_id[i] : default_n_seq_id; + const llama_seq_id * sids = batch_inp.seq_id ? batch_inp.seq_id[i] : &default_seq_id; + + llama_batch_ext::token t; + + // seq_ids + for (int32_t s = 0; s < n_sid; ++s) { + t.seq_ids.insert(sids[s]); + } + + // position(s) + if (batch_inp.pos) { + if (has_token) { + // token batch: one position per token + t.pos[0] = batch_inp.pos[i]; + } else { + // embedding batch (M-RoPE): section-major layout pos[j*n_tokens + i] + for (uint32_t j = 0; j < batch_ext->n_pos_per_embd; ++j) { + t.pos[j] = batch_inp.pos[(int32_t) j * batch_inp.n_tokens + i]; + } + } + } else { + // auto-generate position from the first seq_id + t.pos[0] = pos_next[sids[0]]++; + } + + // token id and/or embeddings + if (has_token) { + t.id = batch_inp.token[i]; + } + + if (has_embd) { + t.has_embd = true; + t.embd_off = batch_ext->embd.size(); + const float * src = batch_inp.embd + (size_t) i * n_embd_row; + batch_ext->embd.insert(batch_ext->embd.end(), src, src + n_embd_row); + batch_ext->n_embd = n_embd_row; + } + + // output flag + // if no logits array is given, default to only the last token being an output + t.output = batch_inp.logits + ? (batch_inp.logits[i] != 0) + : (i == batch_inp.n_tokens - 1); + + batch_ext->tokens.push_back(t); + } +} + +llama_batch_compat::llama_batch_compat(llama_context * ctx, const llama_batch & batch_inp, size_t n_embd_row) { + batch_ext = new llama_batch_ext(ctx); + init(*batch_ext, batch_inp, n_embd_row); +} + +llama_batch_compat::~llama_batch_compat() { + delete batch_ext; +} diff --git a/src/llama-batch.h b/src/llama-batch.h index a3d1889d4..201d48cce 100644 --- a/src/llama-batch.h +++ b/src/llama-batch.h @@ -2,6 +2,7 @@ #include "llama.h" +#include "llama-arch.h" #include "llama-cparams.h" #include @@ -10,6 +11,7 @@ #include #include #include +#include // keep this struct lightweight struct llama_ubatch { @@ -68,19 +70,71 @@ struct llama_ubatch { std::shared_ptr data; }; +struct llama_hparams; + +// MTP hook batches carry the target model's hidden state (n_embd_out size). +// DFlash batches carry the fused target features at the encoder input width (n_embd_inp_enc size). +// Normal batches carry token embeddings (n_embd_inp size). +size_t llama_batch_ext_select_n_embd_inp(llama_context_type ctx_type, llm_arch arch, const llama_hparams & hparams); + +struct llama_batch_ext { + const size_t n_tokens_max; // max number of tokens that can be stored in the batch + const size_t n_embd_inp; // decoder embd row width + const size_t n_embd_inp_enc; // encoder embd row width (e.g. eagle3/dflash extracted features) + const llama_seq_id n_seq_max; // max number of sequences + llama_memory_i * mem; // memory for position inference + const llama_token n_vocab; // max token ID that we accept + const size_t n_pos_per_embd; + + // actual embd row width of this batch, set by the first set_token_embd() + // must be either n_embd_inp or n_embd_inp_enc; encode/decode verify it against the graph input + size_t n_embd = 0; + + struct token { + llama_token id = LLAMA_TOKEN_NULL; + bool has_embd = false; // whether embd_off is set + size_t embd_off = 0; // index offset in the embd array + bool output = false; // TODO: have dedicated output flags + std::unordered_set seq_ids; + std::array pos = {0, 0, 0, 0}; + }; + std::vector tokens; + std::vector embd; + + llama_batch_ext(llama_context * ctx); + + // build without a llama_context, used by tests + llama_batch_ext( + size_t n_tokens_max, + size_t n_embd_inp, + size_t n_embd_inp_enc, + llama_seq_id n_seq_max, + llama_memory_i * mem, + llama_token n_vocab, + size_t n_pos_per_embd); + + void clear(); + + // add an entry with an undefined position + // the caller must set it explicitly via set_token_pos() + int32_t add_token(llama_seq_id seq_id); + + bool add_seq(int32_t idx, llama_seq_id seq_id); + bool set_token_id(int32_t idx, llama_token id); + bool set_token_embd(int32_t idx, llama_embd embd_in); + bool set_token_pos(int32_t idx, const llama_pos * pos_in); + bool set_output(int32_t idx, bool output_last); +}; + // a helper for sanitizing, fulfilling and splitting a batch class llama_batch_allocr { public: llama_batch_allocr(uint32_t n_pos_per_embd); - // sanitize and auto-gen missing data in the input batch - // memory is optional. if provided will be used to check for sequence continuity and to determine the positions + // convert a llama_batch_ext to internal llama_batch and sanitize it bool init( - const llama_batch & batch_inp, + const llama_batch_ext & batch_inp, const llama_vocab & vocab, - const llama_memory_i * memory, - uint32_t n_embd, - uint32_t n_seq_max, bool output_all); const llama_batch & get_batch() const; @@ -137,7 +191,9 @@ private: uint32_t n_seq_max; uint32_t n_outputs; - std::array seq_id_0 = {{ 0 }}; // default sequence id + std::vector token_vec; // owned token IDs built from llama_batch_ext + std::vector embd_vec; // owned embeddings built from llama_batch_ext + std::vector seq_id_data; // flat storage for seq_id pointers below std::vector pos; std::vector n_seq_id; @@ -172,3 +228,16 @@ private: int debug; }; + +// RAII translation layer: converts a llama_batch (old API) into a llama_batch_ext +struct llama_batch_compat { + llama_batch_ext * batch_ext; + + // n_embd_row is the embd row width of batch_inp, 0 = use the decoder width + llama_batch_compat(llama_context * ctx, const llama_batch & batch_inp, size_t n_embd_row = 0); + ~llama_batch_compat(); + + // fill an existing llama_batch_ext from a llama_batch (old API) + // note: this is called directly by the tests, skipping llama_context creation + static void init(llama_batch_ext & batch_ext, const llama_batch & batch_inp, size_t n_embd_row = 0); +}; diff --git a/src/llama-context.cpp b/src/llama-context.cpp index 37e68d626..128233914 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -779,10 +779,13 @@ void llama_context::sched_reserve() { // Both graphs must be discarded before either reserved buffer can be replaced. synchronize(); compute_source->synchronize(); - gf_res_prev->reset(); - gf_res_reserve->reset(); - compute_source->gf_res_prev->reset(); - compute_source->gf_res_reserve->reset(); + for (auto * ctx : {this, compute_source}) { + for (auto & res : ctx->gf_res_prev) { + res.reset(); + } + ctx->gf_res_prev_active = nullptr; + ctx->gf_res_reserve->reset(); + } ggml_backend_sched_reset(sched.get()); ggml_backend_sched_reset(compute_source->sched.get()); if (supported && devices == 1 && ggml_backend_sched_share_compute_buffers(sched.get(), compute_source->sched.get())) { @@ -1538,25 +1541,27 @@ llm_graph_result * llama_context::process_ubatch(const llama_ubatch & ubatch, ll return res; } -int llama_context::encode(const llama_batch & batch_inp) { +int llama_context::encode(const llama_batch_ext & batch_inp) { auto compute_lock = lock_compute(); - // MTP hook batches carry both token (next-token id) and embd (h_nextn row), - // so accept either present rather than requiring exactly one. - GGML_ASSERT(batch_inp.token || batch_inp.embd); - if (batch_inp.n_tokens == 0) { + if (batch_inp.tokens.empty()) { LLAMA_LOG_ERROR("%s: n_tokens == 0\n", __func__); return -1; } const auto & hparams = model.hparams; + if (batch_inp.n_embd > 0 && batch_inp.n_embd != hparams.n_embd_inp_enc()) { + LLAMA_LOG_ERROR("%s: embd row width %zu does not match the encoder input %u\n", + __func__, batch_inp.n_embd, hparams.n_embd_inp_enc()); + return -1; + } + // eagle3/DFlash: features as encoder input, and non-draft paths fall back to model's input dim - const int64_t n_embd = hparams.n_embd_inp_enc(); const int64_t n_vocab = model.vocab.n_tokens(); - // note: during encode, we always pass the full sequence starting from pos = 0 - if (!balloc->init(batch_inp, model.vocab, nullptr, n_embd, cparams.kv_unified ? LLAMA_MAX_SEQ : cparams.n_seq_max, true)) { + // note: during encode, we always output all tokens and skip position continuity checks (output_all=true) + if (!balloc->init(batch_inp, model.vocab, true)) { LLAMA_LOG_ERROR("%s: failed to initialize batch\n", __func__); return -1; } @@ -1777,30 +1782,29 @@ static bool needs_raw_logits(const llama_ubatch & ubatch, const std::map 0 && batch_inp.n_embd != batch_inp.n_embd_inp) { + LLAMA_LOG_ERROR("%s: embd row width %zu does not match the decoder input %zu\n", + __func__, batch_inp.n_embd, batch_inp.n_embd_inp); + return -1; + } + const auto & vocab = model.vocab; const auto & hparams = model.hparams; const int64_t n_vocab = vocab.n_tokens(); - const bool mtp_embd = cparams.ctx_type == LLAMA_CONTEXT_TYPE_MTP && batch_inp.embd; - // DFlash embd batches carry the fused target features at the encoder input width - const bool dflash_embd = model.arch == LLM_ARCH_DFLASH && batch_inp.embd; - const int64_t n_embd = mtp_embd ? hparams.n_embd_out() : dflash_embd ? hparams.n_embd_inp_enc() : hparams.n_embd_inp(); // when computing embeddings, all tokens are output const bool output_all = cparams.embeddings; @@ -1808,20 +1812,17 @@ int llama_context::decode(const llama_batch & batch_inp) { const uint32_t n_seq_max = cparams.kv_unified ? LLAMA_MAX_SEQ : cparams.n_seq_max; - // embedding contexts output every token even when batch.logits is not set - if (has_samplers && (output_all || batch_inp.logits)) { + // TODO: avoid this workaround in the future + // embedding contexts output every token even when no token is explicitly marked as output + if (has_samplers) { std::vector seq_output_count(n_seq_max, 0); - for (int32_t i = 0; i < batch_inp.n_tokens; ++i) { - if (!output_all && batch_inp.logits[i] == 0) { + for (const auto & tok : batch_inp.tokens) { + if (!output_all && !tok.output) { continue; } - const int ns = batch_inp.n_seq_id ? batch_inp.n_seq_id[i] : 1; - - for (int32_t s = 0; s < ns; ++s) { - const llama_seq_id seq_id = batch_inp.seq_id ? batch_inp.seq_id[i][s] : 0; - + for (auto seq_id : tok.seq_ids) { if (seq_id < 0 || (uint32_t) seq_id >= n_seq_max) { continue; } @@ -1839,7 +1840,7 @@ int llama_context::decode(const llama_batch & batch_inp) { } } - if (!balloc->init(batch_inp, vocab, memory.get(), n_embd, n_seq_max, output_all)) { + if (!balloc->init(batch_inp, vocab, output_all)) { LLAMA_LOG_ERROR("%s: failed to initialize batch\n", __func__); return -1; } @@ -3637,9 +3638,13 @@ void llama_context::opt_epoch_iter( batch.logits [pos_batch] = true; } - if (!balloc->init(batch, model.vocab, nullptr, model.hparams.n_embd_inp(), cparams.kv_unified ? LLAMA_MAX_SEQ : cparams.n_seq_max, true)) { - LLAMA_LOG_ERROR("%s: failed to initialize batch\n", __func__); - return; + // TODO: use llama_batch_ext here + { + llama_batch_compat compat(this, batch); + if (!balloc->init(*compat.batch_ext, model.vocab, true)) { + LLAMA_LOG_ERROR("%s: failed to initialize batch\n", __func__); + return; + } } const uint32_t n_tokens_all = balloc->get_n_tokens(); @@ -4392,6 +4397,18 @@ size_t llama_state_seq_load_file(llama_context * ctx, const char * filepath, lla } } +// compat: llama_batch -> llama_batch_ext -> encode/decode + +int llama_context::encode(const llama_batch & batch_inp) { + llama_batch_compat compat(this, batch_inp, model.hparams.n_embd_inp_enc()); + return encode(*compat.batch_ext); +} + +int llama_context::decode(const llama_batch & batch_inp) { + llama_batch_compat compat(this, batch_inp); + return decode(*compat.batch_ext); +} + /// int32_t llama_encode( @@ -4481,6 +4498,14 @@ void llama_opt_epoch( callback_eval); } +int32_t llama_process(llama_context * ctx, llama_process_type type, llama_batch_ext * batch) { + switch (type) { + case LLAMA_PROCESS_TYPE_ENCODE: return ctx->encode(*batch); + case LLAMA_PROCESS_TYPE_DECODE: return ctx->decode(*batch); + } + return -1; +} + // // ext // diff --git a/src/llama-context.h b/src/llama-context.h index 08d0f7fce..ff015eb19 100644 --- a/src/llama-context.h +++ b/src/llama-context.h @@ -148,6 +148,10 @@ struct llama_context { llama_memory_context_i * mctx, ggml_status & ret); + int encode(const llama_batch_ext & batch_inp); + int decode(const llama_batch_ext & batch_inp); + + // compat version int encode(const llama_batch & batch_inp); int decode(const llama_batch & batch_inp); diff --git a/src/llama-grammar.cpp b/src/llama-grammar.cpp index deffec8c0..0a8a61e7e 100644 --- a/src/llama-grammar.cpp +++ b/src/llama-grammar.cpp @@ -194,7 +194,11 @@ static std::pair parse_token(const llama_vocab * vocab, if (*pos == '[') { pos++; const char * int_end = parse_int(pos); - uint32_t token_id = std::stoul(std::string(pos, int_end - pos)); + unsigned long id = std::stoul(std::string(pos, int_end - pos)); + if (id > std::numeric_limits::max()) { + throw std::runtime_error(std::string("parsed token id is too big at ") + pos); + } + uint32_t token_id = static_cast(id); pos = int_end; if (*pos != ']') { throw std::runtime_error(std::string("expecting ']' at ") + pos); diff --git a/src/llama-hparams.cpp b/src/llama-hparams.cpp index 34b3c6880..b83f45ba9 100644 --- a/src/llama-hparams.cpp +++ b/src/llama-hparams.cpp @@ -283,7 +283,8 @@ bool llama_hparams::is_ple(uint32_t il) const { } uint32_t llama_hparams::n_pos_per_embd() const { - return rope_type == LLAMA_ROPE_TYPE_MROPE || rope_type == LLAMA_ROPE_TYPE_IMROPE ? 4 : 1; + return (rope_type == LLAMA_ROPE_TYPE_MROPE || rope_type == LLAMA_ROPE_TYPE_IMROPE) + ? GGML_MROPE_SECTIONS : 1; } bool llama_hparams::is_swa(uint32_t il) const { diff --git a/src/llama-model.cpp b/src/llama-model.cpp index 8e8136641..201c75f88 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -671,11 +671,16 @@ struct ggml_backend_meta_split_state llama_meta_device_get_split_state(const str } if (std::regex_match(tensor_name, pattern_qkv_weight) || std::regex_match(tensor_name, pattern_qkv_bias)) { - const int64_t n_embd = hparams.n_head(il) * hparams.n_embd_head_k(il); - const int64_t n_embd_gqa = hparams.n_embd_v_gqa(il); - GGML_ASSERT(hparams.n_embd_k_gqa(il) == n_embd_gqa); - GGML_ASSERT(tensor->ne[axis] == n_embd + 2*n_embd_gqa); - return {{n_embd, 1}, {n_embd_gqa, 2}}; + const int64_t n_embd_q = hparams.n_head(il) * hparams.n_embd_head_k(il); + const int64_t n_embd_k = hparams.n_embd_k_gqa(il); + const int64_t n_embd_v = hparams.n_embd_v_gqa(il); + GGML_ASSERT(tensor->ne[axis] == n_embd_q + n_embd_k + n_embd_v); + if (n_embd_k == n_embd_v) { + return {{n_embd_q, 1}, {n_embd_k, 2}}; + } + // uneven K/V head sizes (e.g. MiMo d_k=192 d_v=128): split K and V as separate + // segments so each device gets whole heads of both + return {{n_embd_q, 1}, {n_embd_k, 1}, {n_embd_v, 1}}; } if (std::regex_match(tensor_name, pattern_ffn_up_weight) || std::regex_match(tensor_name, pattern_ffn_up_bias)) { const int64_t n_ff = hparams.n_ff(il); @@ -768,7 +773,7 @@ struct ggml_backend_meta_split_state llama_meta_device_get_split_state(const str } if (std::regex_match(tensor_name, pattern_attn_out_weight)) { GGML_ASSERT(segments.size() == 1); - return {granularity_q}; + return {granularity_head * hparams.n_embd_head_v(il)}; } if (std::regex_match(tensor_name, pattern_attn_gate_weight)) { GGML_ASSERT(segments.size() == 1); @@ -779,20 +784,29 @@ struct ggml_backend_meta_split_state llama_meta_device_get_split_state(const str } const int64_t granularity_kv = granularity_q / n_gqa; + // the V head size can differ from the K head size (e.g. MiMo d_k=192 d_v=128): + // align V tensors to whole V heads at the same head-index scale as Q and K so all + // three stay in lockstep per device + const int64_t granularity_v = (granularity_kv / hparams.n_embd_head_k(il)) * hparams.n_embd_head_v(il); if (std::regex_match(tensor_name, pattern_kv_weight) || std::regex_match(tensor_name, pattern_kv_bias) || std::regex_match(tensor_name, pattern_kv_cache)) { GGML_ASSERT(segments.size() == 1); - return {granularity_kv}; + const bool is_v = tensor_name.find("attn_v") != std::string::npos || tensor_name.find("cache_v") != std::string::npos; + return {is_v ? granularity_v : granularity_kv}; } if (std::regex_match(tensor_name, pattern_qkv_weight) || std::regex_match(tensor_name, pattern_qkv_bias)) { - GGML_ASSERT(segments.size() == 2); // fused full attention layers need Q gate tensors handled like above: // TODO: deduplicate condition [TAG_SPLIT_QGATE_QWEN] if (ud->model->arch == LLM_ARCH_QWEN3NEXT || ud->model->arch == LLM_ARCH_QWEN35 || ud->model->arch == LLM_ARCH_QWEN35MOE || ud->model->arch == LLM_ARCH_QWEN4EXP) { return {std::lcm(2*n_embd_q, blck_size_perf), granularity_kv}; } + if (segments.size() == 3) { + // uneven K/V head sizes: per-segment granularity, V aligned to whole V heads + return {granularity_q, granularity_kv, granularity_v}; + } + GGML_ASSERT(segments.size() == 2); return {granularity_q, granularity_kv}; } } diff --git a/src/models/mimo2.cpp b/src/models/mimo2.cpp index a466984af..b6d7aceda 100644 --- a/src/models/mimo2.cpp +++ b/src/models/mimo2.cpp @@ -25,8 +25,10 @@ void llama_model_mimo2::load_arch_hparams(llama_model_loader & ml) { void llama_model_mimo2::load_arch_tensors(llama_model_loader & ml) { LLAMA_LOAD_LOCALS; + const bool mtp_only = (hparams.n_layer_nextn > 0) && (ml.get_weight("blk.0.attn_norm.weight") == nullptr); const std::string mtp_probe = "blk." + std::to_string(n_layer) + ".nextn.eh_proj.weight"; const bool trunk_only = (hparams.n_layer_nextn > 0) && (ml.get_weight(mtp_probe.c_str()) == nullptr); + const int trunk_flags = mtp_only ? TENSOR_NOT_REQUIRED : 0; int mtp_flags = trunk_only ? TENSOR_NOT_REQUIRED : 0; if (!ml.load_mtp) { @@ -46,7 +48,7 @@ void llama_model_mimo2::load_arch_tensors(llama_model_loader & ml) { uint32_t n_head = hparams.n_head(i); const bool is_nextn = i >= n_layer; - const int flags = is_nextn ? mtp_flags : 0; + const int flags = is_nextn ? mtp_flags : trunk_flags; create_tensor_qkv(layer, i, n_embd, n_embd_head_k * n_head, n_embd_k_gqa, n_embd_v_gqa, flags); layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", i), { n_embd_head_v * n_head, n_embd }, flags); diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index 5d3f83213..d5fa5104a 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -845,6 +845,7 @@ struct console_printer : public printer { } else if (result.test_mode == "support") { print_support_console(result); } + fflush(stdout); } void print_operation(const test_operation_info & info) override { @@ -854,6 +855,7 @@ struct console_printer : public printer { // Handle large tensor skip first if (info.is_large_tensor_skip) { printf("skipping large tensors for speed \n"); + fflush(stdout); return; } @@ -864,6 +866,7 @@ struct console_printer : public printer { } else { printf("not supported [%s]\n", info.backend_name.c_str()); } + fflush(stdout); return; } @@ -900,6 +903,7 @@ struct console_printer : public printer { } else { printf("\033[1;31mFAIL\033[0m\n"); } + fflush(stdout); } void print_summary(const test_summary_info & info) override { diff --git a/tests/test-batch-alloc.cpp b/tests/test-batch-alloc.cpp index 66d29d6f5..ad186c693 100644 --- a/tests/test-batch-alloc.cpp +++ b/tests/test-batch-alloc.cpp @@ -3,6 +3,8 @@ #include "llama.h" #include "../src/llama-batch.h" +#include "../src/llama-arch.h" +#include "../src/llama-hparams.h" #include "../src/llama-memory.h" #include "../src/llama-vocab.h" @@ -47,49 +49,55 @@ struct mock_memory : public llama_memory_i { void state_read (llama_io_read_i &, llama_seq_id, llama_state_seq_flags) override { GGML_ASSERT(false && "not implemented"); } }; -// builds embedding batches - an empty llama_vocab rejects all token ids, so -// the tests use embeddings everywhere except the token validation tests +// builds a llama_batch_ext without a llama_context +// n_vocab = 0 by default, so every token id is invalid and the tests use embeddings unless stated otherwise struct batch_builder { - uint32_t n_embd; + const uint32_t n_embd; - std::vector embd; - std::vector pos; - std::vector n_seq_id; - std::vector logits; + llama_batch_ext b; - std::vector> seq; - std::vector seq_ptr; + batch_builder( + uint32_t n_embd = 2, + llama_memory_i * mem = nullptr, + llama_seq_id n_seq_max = 4, + uint32_t n_pos_per_embd = 1, + llama_token n_vocab = 0, + uint32_t n_embd_inp_enc = 0) + : n_embd(n_embd), + b(/*n_tokens_max*/ 64, n_embd, n_embd_inp_enc > 0 ? n_embd_inp_enc : n_embd, n_seq_max, mem, n_vocab, n_pos_per_embd) {} - batch_builder(uint32_t n_embd = 2) : n_embd(n_embd) {} - - // embd values are 100*i + k so that ubatch contents can be traced back to batch indices - void add(llama_pos p, std::initializer_list seq_ids, bool output) { - const int32_t i = (int32_t) seq.size(); - for (uint32_t k = 0; k < n_embd; ++k) { - embd.push_back(100.0f*i + k); + // one embedding row for batch index i, values 100*i + k so ubatch contents can be traced back + std::vector row(int32_t i, uint32_t width) const { + std::vector r(width); + for (uint32_t k = 0; k < width; ++k) { + r[k] = 100.0f*i + k; } - pos.push_back(p); - n_seq_id.push_back((int32_t) seq_ids.size()); - seq.emplace_back(seq_ids); - logits.push_back(output ? 1 : 0); + return r; } - llama_batch make(bool with_pos = true, bool with_seq = true, bool with_logits = true) { - seq_ptr.clear(); - for (auto & s : seq) { - seq_ptr.push_back(s.data()); + // embedding entry with full M-RoPE positions + int32_t add_embd(const llama_pos * pos, std::initializer_list seq_ids, bool output, uint32_t width = 0) { + width = width > 0 ? width : n_embd; + + auto it = seq_ids.begin(); + const int32_t idx = b.add_token(*it); + GGML_ASSERT(idx >= 0); + for (++it; it != seq_ids.end(); ++it) { + GGML_ASSERT(b.add_seq(idx, *it)); } - seq_ptr.push_back(nullptr); - llama_batch res = {}; - res.n_tokens = (int32_t) seq.size(); - res.embd = embd.data(); - res.pos = with_pos ? pos.data() : nullptr; - res.n_seq_id = with_seq ? n_seq_id.data() : nullptr; - res.seq_id = with_seq ? seq_ptr.data() : nullptr; - res.logits = with_logits ? logits.data() : nullptr; + const auto r = row(idx, width); + GGML_ASSERT(b.set_token_embd(idx, { r.data(), 1, width })); + GGML_ASSERT(b.set_token_pos(idx, pos)); + GGML_ASSERT(b.set_output(idx, output)); - return res; + return idx; + } + + // embedding entry with a single sequential position + int32_t add(llama_pos p, std::initializer_list seq_ids, bool output) { + const llama_pos pos[GGML_MROPE_SECTIONS] = { p, 0, 0, 0 }; + return add_embd(pos, seq_ids, output); } }; @@ -97,22 +105,31 @@ static void test_init(testing & t) { llama_vocab vocab; t.test("rejects_n_seq_max_too_large", [&](testing & t) { - batch_builder bb; + batch_builder bb(2, nullptr, LLAMA_MAX_SEQ + 1); bb.add(0, {0}, true); llama_batch_allocr ba(1); - t.assert_true(!ba.init(bb.make(), vocab, nullptr, bb.n_embd, LLAMA_MAX_SEQ + 1, false)); + t.assert_true(!ba.init(bb.b, vocab, false)); }); t.test("rejects_invalid_token", [&](testing & t) { - llama_token tok = 0; // empty vocab -> every token id is out of range - llama_batch batch = llama_batch_get_one(&tok, 1); + // n_vocab = 0 -> every token id is out of range + // set_token_id() refuses such ids, so the token is poked directly to reach the init() check + batch_builder bb; + const int32_t idx = bb.b.add_token(0); + const llama_pos pos = 0; + bb.b.set_token_pos(idx, &pos); + bb.b.set_output(idx, true); llama_batch_allocr ba(1); - t.assert_true("token id >= n_tokens", !ba.init(batch, vocab, nullptr, 0, 1, false)); - tok = -1; - t.assert_true("negative token id", !ba.init(batch, vocab, nullptr, 0, 1, false)); + t.assert_true("set_token_id refuses out of range id", !bb.b.set_token_id(idx, 0)); + + bb.b.tokens[idx].id = 0; + t.assert_true("token id >= n_vocab", !ba.init(bb.b, vocab, false)); + + bb.b.tokens[idx].id = -1; + t.assert_true("negative token id", !ba.init(bb.b, vocab, false)); }); t.test("rejects_invalid_seq_id", [&](testing & t) { @@ -120,33 +137,44 @@ static void test_init(testing & t) { { batch_builder bb; - bb.add(0, {4}, true); - t.assert_true("seq_id >= n_seq_max", !ba.init(bb.make(), vocab, nullptr, bb.n_embd, 4, false)); + t.assert_true("add_token refuses seq_id >= n_seq_max", bb.b.add_token(4) == -3); + t.assert_true("add_token refuses negative seq_id", bb.b.add_token(-1) == -3); + } + { + // poke the seq_ids directly to reach the init() check + batch_builder bb; + const int32_t idx = bb.add(0, {0}, true); + bb.b.tokens[idx].seq_ids = { 4 }; + t.assert_true("seq_id >= n_seq_max", !ba.init(bb.b, vocab, false)); } { batch_builder bb; - bb.add(0, {-1}, true); - t.assert_true("negative seq_id", !ba.init(bb.make(), vocab, nullptr, bb.n_embd, 4, false)); + const int32_t idx = bb.add(0, {0}, true); + bb.b.tokens[idx].seq_ids = { -1 }; + t.assert_true("negative seq_id", !ba.init(bb.b, vocab, false)); } }); - t.test("autofill_defaults", [&](testing & t) { + t.test("copies_pos_seq_output", [&](testing & t) { batch_builder bb; for (int i = 0; i < 4; ++i) { - bb.add(0, {0}, false); + bb.add(i, {0}, i == 3); } llama_batch_allocr ba(1); - t.assert_true(ba.init(bb.make(false, false, false), vocab, nullptr, bb.n_embd, 4, false)); + t.assert_true(ba.init(bb.b, vocab, false)); const llama_batch & batch = ba.get_batch(); t.assert_equal(4u, ba.get_n_tokens()); + t.assert_true("embedding batch", batch.embd != nullptr); + t.assert_true("no token ids", batch.token == nullptr); for (int i = 0; i < 4; ++i) { - t.assert_equal("pos defaults to 0..n-1", i, batch.pos[i]); - t.assert_equal("n_seq_id defaults to 1", 1, batch.n_seq_id[i]); - t.assert_equal("seq_id defaults to 0", 0, batch.seq_id[i][0]); + t.assert_equal(i, batch.pos[i]); + t.assert_equal(1, batch.n_seq_id[i]); + t.assert_equal(0, batch.seq_id[i][0]); + t.assert_equal(100.0f*i, batch.embd[i*bb.n_embd]); } t.assert_equal("only the last token is an output", 1u, ba.get_n_outputs()); @@ -165,7 +193,7 @@ static void test_init(testing & t) { } llama_batch_allocr ba(1); - t.assert_true(ba.init(bb.make(true, true, false), vocab, nullptr, bb.n_embd, 4, true)); + t.assert_true(ba.init(bb.b, vocab, true)); t.assert_equal(4u, ba.get_n_outputs()); }); @@ -176,7 +204,7 @@ static void test_init(testing & t) { bb.add(2, {0}, true); llama_batch_allocr ba(1); - t.assert_true(ba.init(bb.make(), vocab, nullptr, bb.n_embd, 4, false)); + t.assert_true(ba.init(bb.b, vocab, false)); t.assert_equal(2u, ba.get_n_outputs()); llama_ubatch ub = ba.split_simple(10); @@ -191,17 +219,17 @@ static void test_init(testing & t) { t.assert_equal(2, out_ids[1]); }); - t.test("pos_from_memory", [&](testing & t) { + t.test("pos_after_memory", [&](testing & t) { mock_memory mem; mem.ranges[0] = {0, 9}; - batch_builder bb; + batch_builder bb(2, &mem); for (int i = 0; i < 3; ++i) { - bb.add(0, {0}, false); + bb.add(10 + i, {0}, false); } llama_batch_allocr ba(1); - t.assert_true(ba.init(bb.make(false, true, false), vocab, &mem, bb.n_embd, 4, false)); + t.assert_true(ba.init(bb.b, vocab, false)); t.assert_equal("pos continues after memory", 10, ba.seq_pos_min(0)); t.assert_equal(12, ba.seq_pos_max(0)); @@ -214,22 +242,22 @@ static void test_init(testing & t) { llama_batch_allocr ba(1); { - batch_builder bb; + batch_builder bb(2, &mem); bb.add(10, {0}, false); bb.add(11, {0}, true); - t.assert_true("pos_max + 1 is accepted", ba.init(bb.make(), vocab, &mem, bb.n_embd, 4, false)); + t.assert_true("pos_max + 1 is accepted", ba.init(bb.b, vocab, false)); } { - batch_builder bb; + batch_builder bb(2, &mem); bb.add(11, {0}, false); bb.add(12, {0}, true); - t.assert_true("gap after memory is rejected", !ba.init(bb.make(), vocab, &mem, bb.n_embd, 4, false)); + t.assert_true("gap after memory is rejected", !ba.init(bb.b, vocab, false)); } { - batch_builder bb; + batch_builder bb(2, &mem); bb.add(9, {0}, false); bb.add(10, {0}, true); - t.assert_true("overlap with memory is rejected", !ba.init(bb.make(), vocab, &mem, bb.n_embd, 4, false)); + t.assert_true("overlap with memory is rejected", !ba.init(bb.b, vocab, false)); } }); @@ -240,7 +268,7 @@ static void test_init(testing & t) { bb.add(3, {0}, true); llama_batch_allocr ba(1); - t.assert_true(!ba.init(bb.make(), vocab, nullptr, bb.n_embd, 4, false)); + t.assert_true(!ba.init(bb.b, vocab, false)); }); t.test("rejects_decreasing_positions", [&](testing & t) { @@ -253,7 +281,7 @@ static void test_init(testing & t) { // seq 0 sees positions 4,5,6,3 in batch order -> the trailing 3 decreases llama_batch_allocr ba(1); - t.assert_true(!ba.init(bb.make(true, true, false), vocab, nullptr, bb.n_embd, 4, false)); + t.assert_true(!ba.init(bb.b, vocab, false)); }); t.test("allows_equal_positions_in_seq", [&](testing & t) { @@ -263,23 +291,143 @@ static void test_init(testing & t) { bb.add(1, {0}, true); llama_batch_allocr ba(1); - t.assert_true(ba.init(bb.make(true, true, false), vocab, nullptr, bb.n_embd, 4, false)); + t.assert_true(ba.init(bb.b, vocab, false)); }); - t.test("rejects_coupled_diverged_seqs", [&](testing & t) { - batch_builder bb; - bb.add(6, {0, 1}, true); - llama_batch_allocr ba(1); mock_memory mem; mem.ranges[0] = {0, 5}; mem.ranges[1] = {2, 5}; // same pos_max, different pos_min -> diverged - t.assert_true(!ba.init(bb.make(), vocab, &mem, bb.n_embd, 4, false)); + { + batch_builder bb(2, &mem); + bb.add(6, {0, 1}, true); + t.assert_true(!ba.init(bb.b, vocab, false)); + } mem.ranges[1] = {0, 5}; - t.assert_true(ba.init(bb.make(), vocab, &mem, bb.n_embd, 4, false)); + { + batch_builder bb(2, &mem); + bb.add(6, {0, 1}, true); + t.assert_true(ba.init(bb.b, vocab, false)); + } + }); +} + +static void test_content_types(testing & t) { + llama_vocab vocab; + + t.test("token_and_embd_together", [&](testing & t) { + // e.g. MTP hook batches: a token id and its embedding on the same entry + batch_builder bb(2, nullptr, 4, 1, /*n_vocab*/ 10); + + const int32_t idx = bb.b.add_token(0); + t.assert_true(bb.b.set_token_id(idx, 3)); + const auto r = bb.row(idx, bb.n_embd); + t.assert_true(bb.b.set_token_embd(idx, { r.data(), 1, bb.n_embd })); + const llama_pos pos = 0; + bb.b.set_token_pos(idx, &pos); + bb.b.set_output(idx, true); + + llama_batch_allocr ba(1); + t.assert_true(ba.init(bb.b, vocab, false)); + + const llama_batch & batch = ba.get_batch(); + t.assert_true("token ids are kept", batch.token != nullptr); + t.assert_true("embeddings are kept", batch.embd != nullptr); + t.assert_equal(3, batch.token[0]); + t.assert_equal(0.0f, batch.embd[0]); + t.assert_equal(1.0f, batch.embd[1]); + + llama_ubatch ub = ba.split_simple(1); + t.assert_true(ub.token != nullptr && ub.embd != nullptr); + t.assert_equal(3, ub.token[0]); + }); + + t.test("rejects_mixed_content_types", [&](testing & t) { + batch_builder bb(2, nullptr, 4, 1, /*n_vocab*/ 10); + + // entry 0: token only, entry 1: token + embd + const llama_pos p0 = 0; + const llama_pos p1 = 1; + + int32_t i0 = bb.b.add_token(0); + bb.b.set_token_id(i0, 1); + bb.b.set_token_pos(i0, &p0); + + int32_t i1 = bb.b.add_token(0); + bb.b.set_token_id(i1, 2); + const auto r = bb.row(i1, bb.n_embd); + bb.b.set_token_embd(i1, { r.data(), 1, bb.n_embd }); + bb.b.set_token_pos(i1, &p1); + bb.b.set_output(i1, true); + + llama_batch_allocr ba(1); + t.assert_true(!ba.init(bb.b, vocab, false)); + }); + + t.test("rejects_neither_token_nor_embd", [&](testing & t) { + batch_builder bb; + const int32_t idx = bb.b.add_token(0); + const llama_pos pos = 0; + bb.b.set_token_pos(idx, &pos); + bb.b.set_output(idx, true); + + llama_batch_allocr ba(1); + t.assert_true(!ba.init(bb.b, vocab, false)); + }); + + t.test("rejects_embd_size_mismatch", [&](testing & t) { + batch_builder bb; // n_embd = 2, n_embd_inp_enc = 2 + const int32_t idx = bb.b.add_token(0); + const auto r = bb.row(idx, 8); + + t.assert_true("too small", !bb.b.set_token_embd(idx, { r.data(), 1, 1 })); + t.assert_true("too large", !bb.b.set_token_embd(idx, { r.data(), 1, 3 })); + t.assert_true("zero rows", !bb.b.set_token_embd(idx, { r.data(), 0, 2 })); + t.assert_true("null data", !bb.b.set_token_embd(idx, { nullptr, 1, 2 })); + t.assert_true("same total via a different split is accepted", bb.b.set_token_embd(idx, { r.data(), 2, 1 })); + }); + + t.test("rejects_double_embd", [&](testing & t) { + batch_builder bb; + const int32_t idx = bb.add(0, {0}, true); + const auto r = bb.row(idx, bb.n_embd); + t.assert_true(!bb.b.set_token_embd(idx, { r.data(), 1, bb.n_embd })); + }); + + t.test("encoder_width", [&](testing & t) { + // e.g. eagle3/dflash: extracted features are wider than the decoder input + const uint32_t n_embd_enc = 6; + batch_builder bb(2, nullptr, 4, 1, 0, n_embd_enc); + + const llama_pos p0 = 0; + const llama_pos p1 = 1; + bb.add_embd(&p0, {0}, false, n_embd_enc); + bb.add_embd(&p1, {0}, true, n_embd_enc); + + t.assert_equal("batch width follows the first embedding", (size_t) n_embd_enc, bb.b.n_embd); + + llama_batch_allocr ba(1); + t.assert_true(ba.init(bb.b, vocab, false)); + + // the ubatch uses the encoder stride: token 1 starts at offset n_embd_enc + llama_ubatch ub = ba.split_simple(2); + t.assert_equal(2u, ub.n_tokens); + t.assert_equal(100.0f, ub.embd[n_embd_enc]); + t.assert_equal(105.0f, ub.embd[n_embd_enc + 5]); + }); + + t.test("rejects_mixing_widths", [&](testing & t) { + batch_builder bb(2, nullptr, 4, 1, 0, /*n_embd_inp_enc*/ 6); + + const llama_pos p0 = 0; + bb.add_embd(&p0, {0}, false, 2); // first entry fixes the batch width to 2 + + const int32_t idx = bb.b.add_token(0); + const auto r = bb.row(idx, 6); + t.assert_true(!bb.b.set_token_embd(idx, { r.data(), 1, 6 })); }); } @@ -293,7 +441,7 @@ static void test_split(testing & t) { } llama_batch_allocr ba(1); - t.assert_true(ba.init(bb.make(), vocab, nullptr, bb.n_embd, 4, false)); + t.assert_true(ba.init(bb.b, vocab, false)); llama_ubatch ub = ba.split_simple(2); t.assert_equal(2u, ub.n_tokens); @@ -336,7 +484,7 @@ static void test_split(testing & t) { } llama_batch_allocr ba(1); - t.assert_true(ba.init(bb.make(), vocab, nullptr, bb.n_embd, 4, false)); + t.assert_true(ba.init(bb.b, vocab, false)); while (ba.split_simple(1).n_tokens > 0) { } @@ -359,7 +507,7 @@ static void test_split(testing & t) { } llama_batch_allocr ba(1); - t.assert_true(ba.init(bb.make(), vocab, nullptr, bb.n_embd, 4, false)); + t.assert_true(ba.init(bb.b, vocab, false)); llama_ubatch ub = ba.split_equal(8, false, 0); t.assert_true(ub.equal_seqs()); @@ -395,7 +543,7 @@ static void test_split(testing & t) { bb.add(1, {0, 1}, true); llama_batch_allocr ba(1); - t.assert_true(ba.init(bb.make(), vocab, nullptr, bb.n_embd, 4, false)); + t.assert_true(ba.init(bb.b, vocab, false)); llama_ubatch ub = ba.split_equal(4, true, 0); t.assert_equal("sequential split rejects coupled seqs", 0u, ub.n_tokens); @@ -417,7 +565,7 @@ static void test_split(testing & t) { } llama_batch_allocr ba(1); - t.assert_true(ba.init(bb.make(), vocab, nullptr, bb.n_embd, 4, false)); + t.assert_true(ba.init(bb.b, vocab, false)); for (llama_seq_id s = 0; s < 3; ++s) { llama_ubatch ub = ba.split_seq(8); @@ -459,14 +607,14 @@ static void test_keep_tail(testing & t) { } ++s; } - return bb.make(); }; t.test("noop_when_seqs_complete", [&](testing & t) { batch_builder bb; + make_batch(bb, {2, 2}); llama_batch_allocr ba(1); - t.assert_true(ba.init(make_batch(bb, {2, 2}), vocab, nullptr, bb.n_embd, 4, false)); + t.assert_true(ba.init(bb.b, vocab, false)); llama_ubatch ub = ba.split_equal(4, false, 2); t.assert_equal("both seqs fit whole", 4u, ub.n_tokens); @@ -478,9 +626,10 @@ static void test_keep_tail(testing & t) { t.test("defers_seq_with_short_remainder", [&](testing & t) { batch_builder bb; + make_batch(bb, {2, 3}); llama_batch_allocr ba(1); - t.assert_true(ba.init(make_batch(bb, {2, 3}), vocab, nullptr, bb.n_embd, 4, false)); + t.assert_true(ba.init(bb.b, vocab, false)); // expansion stops at 2 tokens per seq: seq 0 completes, seq 1 would be left // with 1 < n_keep_tail remaining, so it is deferred entirely @@ -504,9 +653,10 @@ static void test_keep_tail(testing & t) { t.test("completes_first_seq_when_all_violate", [&](testing & t) { batch_builder bb; + make_batch(bb, {3, 3}); llama_batch_allocr ba(1); - t.assert_true(ba.init(make_batch(bb, {3, 3}), vocab, nullptr, bb.n_embd, 4, false)); + t.assert_true(ba.init(bb.b, vocab, false)); // expansion stops at 2 tokens per seq, leaving both with 1 < n_keep_tail remaining; // seq 0 still fits in n_ubatch, so it is extended to completion and emitted alone @@ -528,9 +678,10 @@ static void test_keep_tail(testing & t) { t.test("truncates_to_preserve_tail", [&](testing & t) { batch_builder bb; + make_batch(bb, {5}); llama_batch_allocr ba(1); - t.assert_true(ba.init(make_batch(bb, {5}), vocab, nullptr, bb.n_embd, 4, false)); + t.assert_true(ba.init(bb.b, vocab, false)); // 4 tokens would leave a remainder of 1, and the seq does not fit in n_ubatch, // so the ubatch is truncated until n_keep_tail tokens remain @@ -551,9 +702,10 @@ static void test_keep_tail(testing & t) { t.test("keeps_full_ubatch_with_sufficient_remainder", [&](testing & t) { batch_builder bb; + make_batch(bb, {6}); llama_batch_allocr ba(1); - t.assert_true(ba.init(make_batch(bb, {6}), vocab, nullptr, bb.n_embd, 4, false)); + t.assert_true(ba.init(bb.b, vocab, false)); llama_ubatch ub = ba.split_equal(4, false, 2); t.assert_equal("remainder >= n_keep_tail, no truncation", 4u, ub.n_tokens); @@ -567,10 +719,11 @@ static void test_keep_tail(testing & t) { }); t.test("multi_seq_prefix_kept", [&](testing & t) { - batch_builder bb; + batch_builder bb(2, nullptr, 6); + make_batch(bb, {3, 4}); llama_batch_allocr ba(1); - t.assert_true(ba.init(make_batch(bb, {3, 4}), vocab, nullptr, bb.n_embd, 6, false)); + t.assert_true(ba.init(bb.b, vocab, false)); // expansion stops at 3 tokens per seq: seq 0 completes, seq 1 has 1 < n_keep_tail // remaining and is deferred even though its tokens were already gathered @@ -591,32 +744,26 @@ static void test_mrope(testing & t) { llama_vocab vocab; t.test("pos_layout_and_split", [&](testing & t) { - const uint32_t n_pos = 4; + const uint32_t n_pos = 4; const uint32_t n_embd = 2; - batch_builder bb(n_embd); - bb.add(10, {0}, false); - bb.add(11, {0}, true); + batch_builder bb(n_embd, nullptr, 4, n_pos); - // M-RoPE positions for embeddings are laid out [n_pos][n_tokens] - std::vector pos = { - 10, 11, // temporal - 5, 6, // y - 7, 8, // x - 0, 0, - }; - - llama_batch batch = bb.make(false, true, true); - batch.pos = pos.data(); + // M-RoPE positions per embedding: [temporal, y, x, other] + const llama_pos pos0[n_pos] = { 10, 5, 7, 0 }; + const llama_pos pos1[n_pos] = { 11, 6, 8, 0 }; + bb.add_embd(pos0, {0}, false); + bb.add_embd(pos1, {0}, true); llama_batch_allocr ba(n_pos); - t.assert_true(ba.init(batch, vocab, nullptr, n_embd, 4, false)); + t.assert_true(ba.init(bb.b, vocab, false)); llama_ubatch ub = ba.split_simple(2); t.assert_equal(2u, ub.n_tokens); t.assert_equal(n_pos, ub.n_pos); t.assert_true(ub.is_pos_2d()); + // the ubatch stores positions section-major: [n_pos][n_tokens] const llama_pos expected[8] = {10, 11, 5, 6, 7, 8, 0, 0}; for (int i = 0; i < 8; ++i) { t.assert_equal(expected[i], ub.pos[i]); @@ -624,7 +771,7 @@ static void test_mrope(testing & t) { }); t.test("pos_jump_allowed", [&](testing & t) { - const uint32_t n_pos = 4; + const uint32_t n_pos = 4; const uint32_t n_embd = 2; mock_memory mem; @@ -633,15 +780,12 @@ static void test_mrope(testing & t) { llama_batch_allocr ba(n_pos); auto try_pos = [&](llama_pos p0) { - batch_builder bb(n_embd); - bb.add(p0, {0}, true); + batch_builder bb(n_embd, &mem, 4, n_pos); - std::vector pos = {p0, 1, 1, 0}; + const llama_pos pos[n_pos] = { p0, 1, 1, 0 }; + bb.add_embd(pos, {0}, true); - llama_batch batch = bb.make(false, true, true); - batch.pos = pos.data(); - - return ba.init(batch, vocab, &mem, n_embd, 4, false); + return ba.init(bb.b, vocab, false); }; t.assert_true("gap after memory is allowed", try_pos(15)); @@ -650,6 +794,254 @@ static void test_mrope(testing & t) { }); } +// conversion from the old llama_batch API (llama_batch_compat::init) +static void test_compat(testing & t) { + llama_vocab vocab; + + t.test("token_batch_explicit_fields", [&](testing & t) { + llama_token token[3] = { 5, 6, 7 }; + llama_pos pos[3] = { 3, 4, 5 }; + int32_t n_seq_id[3] = { 1, 1, 2 }; + llama_seq_id s0[1] = { 1 }; + llama_seq_id s1[1] = { 1 }; + llama_seq_id s2[2] = { 1, 2 }; + llama_seq_id * seq_id[4] = { s0, s1, s2, nullptr }; + int8_t logits[3] = { 0, 1, 0 }; + + llama_batch lb = {}; + lb.n_tokens = 3; + lb.token = token; + lb.pos = pos; + lb.n_seq_id = n_seq_id; + lb.seq_id = seq_id; + lb.logits = logits; + + batch_builder bb(2, nullptr, 4, 1, /*n_vocab*/ 100); + llama_batch_compat::init(bb.b, lb); + + t.assert_equal((size_t) 3, bb.b.tokens.size()); + t.assert_true("no embeddings", bb.b.embd.empty() && bb.b.n_embd == 0); + for (int i = 0; i < 3; ++i) { + t.assert_equal(token[i], bb.b.tokens[i].id); + t.assert_equal(pos[i], bb.b.tokens[i].pos[0]); + t.assert_true(!bb.b.tokens[i].has_embd); + t.assert_equal(logits[i] != 0, bb.b.tokens[i].output); + } + t.assert_equal((size_t) 1, bb.b.tokens[0].seq_ids.size()); + t.assert_true(bb.b.tokens[0].seq_ids.count(1) == 1); + t.assert_equal((size_t) 2, bb.b.tokens[2].seq_ids.size()); + t.assert_true(bb.b.tokens[2].seq_ids.count(1) == 1 && bb.b.tokens[2].seq_ids.count(2) == 1); + + // round trip through the allocator + llama_batch_allocr ba(1); + t.assert_true(ba.init(bb.b, vocab, false)); + const llama_batch & batch = ba.get_batch(); + t.assert_true(batch.token != nullptr && batch.embd == nullptr); + for (int i = 0; i < 3; ++i) { + t.assert_equal(token[i], batch.token[i]); + t.assert_equal(pos[i], batch.pos[i]); + } + t.assert_equal(1u, ba.get_n_outputs()); + }); + + t.test("defaults_for_null_fields", [&](testing & t) { + // llama_batch_get_one: only token and n_tokens are set + mock_memory mem; + mem.ranges[0] = {0, 9}; + + llama_token token[3] = { 5, 6, 7 }; + llama_batch lb = llama_batch_get_one(token, 3); + + batch_builder bb(2, &mem, 4, 1, /*n_vocab*/ 100); + llama_batch_compat::init(bb.b, lb); + + t.assert_equal((size_t) 3, bb.b.tokens.size()); + for (int i = 0; i < 3; ++i) { + t.assert_equal("pos continues after memory", 10 + i, bb.b.tokens[i].pos[0]); + t.assert_equal("seq_id defaults to 0", (size_t) 1, bb.b.tokens[i].seq_ids.size()); + t.assert_true(bb.b.tokens[i].seq_ids.count(0) == 1); + } + t.assert_true("only the last token is an output", !bb.b.tokens[0].output && !bb.b.tokens[1].output && bb.b.tokens[2].output); + + llama_batch_allocr ba(1); + t.assert_true(ba.init(bb.b, vocab, false)); + t.assert_equal(10, ba.seq_pos_min(0)); + t.assert_equal(12, ba.seq_pos_max(0)); + }); + + t.test("auto_pos_starts_at_zero_without_memory", [&](testing & t) { + llama_token token[2] = { 5, 6 }; + llama_batch lb = llama_batch_get_one(token, 2); + + batch_builder bb(2, nullptr, 4, 1, /*n_vocab*/ 100); + llama_batch_compat::init(bb.b, lb); + + t.assert_equal(0, bb.b.tokens[0].pos[0]); + t.assert_equal(1, bb.b.tokens[1].pos[0]); + }); + + t.test("auto_pos_is_tracked_per_seq", [&](testing & t) { + mock_memory mem; + mem.ranges[0] = {0, 9}; // seq 1 is empty + + llama_token token[4] = { 5, 6, 7, 8 }; + int32_t n_seq_id[4] = { 1, 1, 1, 1 }; + llama_seq_id s0[1] = { 0 }; + llama_seq_id s1[1] = { 1 }; + llama_seq_id * seq_id[5] = { s0, s1, s0, s1, nullptr }; + + llama_batch lb = {}; + lb.n_tokens = 4; + lb.token = token; + lb.n_seq_id = n_seq_id; + lb.seq_id = seq_id; + + batch_builder bb(2, &mem, 4, 1, /*n_vocab*/ 100); + llama_batch_compat::init(bb.b, lb); + + t.assert_equal("seq 0 continues after memory", 10, bb.b.tokens[0].pos[0]); + t.assert_equal("seq 1 starts from 0", 0, bb.b.tokens[1].pos[0]); + t.assert_equal(11, bb.b.tokens[2].pos[0]); + t.assert_equal( 1, bb.b.tokens[3].pos[0]); + }); + + t.test("embd_batch_with_mrope_positions", [&](testing & t) { + const uint32_t n_pos = 4; + const uint32_t n_embd = 2; + + float embd[2*n_embd] = { 0, 1, 100, 101 }; + // section-major layout: pos[j*n_tokens + i] + llama_pos pos[n_pos*2] = { + 10, 11, // temporal + 5, 6, // y + 7, 8, // x + 0, 0, + }; + + llama_batch lb = {}; + lb.n_tokens = 2; + lb.embd = embd; + lb.pos = pos; + + batch_builder bb(n_embd, nullptr, 4, n_pos); + llama_batch_compat::init(bb.b, lb); + + t.assert_equal((size_t) 2, bb.b.tokens.size()); + t.assert_equal("batch width", (size_t) n_embd, bb.b.n_embd); + for (int i = 0; i < 2; ++i) { + t.assert_true(bb.b.tokens[i].has_embd); + t.assert_equal(LLAMA_TOKEN_NULL, bb.b.tokens[i].id); + t.assert_equal((size_t) i*n_embd, bb.b.tokens[i].embd_off); + for (uint32_t j = 0; j < n_pos; ++j) { + t.assert_equal(pos[j*2 + i], bb.b.tokens[i].pos[j]); + } + } + t.assert_equal(100.0f, bb.b.embd[2]); + t.assert_equal(101.0f, bb.b.embd[3]); + + llama_batch_allocr ba(n_pos); + t.assert_true(ba.init(bb.b, vocab, false)); + llama_ubatch ub = ba.split_simple(2); + const llama_pos expected[8] = {10, 11, 5, 6, 7, 8, 0, 0}; + for (int i = 0; i < 8; ++i) { + t.assert_equal(expected[i], ub.pos[i]); + } + }); + + t.test("token_and_embd_both_set", [&](testing & t) { + // e.g. MTP hook batches + llama_token token[2] = { 5, 6 }; + float embd[4] = { 0, 1, 100, 101 }; + llama_pos pos[2] = { 3, 4 }; + + llama_batch lb = {}; + lb.n_tokens = 2; + lb.token = token; + lb.embd = embd; + lb.pos = pos; + + batch_builder bb(2, nullptr, 4, 1, /*n_vocab*/ 100); + llama_batch_compat::init(bb.b, lb); + + for (int i = 0; i < 2; ++i) { + t.assert_equal(token[i], bb.b.tokens[i].id); + t.assert_true(bb.b.tokens[i].has_embd); + t.assert_equal("one position per token", pos[i], bb.b.tokens[i].pos[0]); + } + t.assert_equal(100.0f, bb.b.embd[2]); + + llama_batch_allocr ba(1); + t.assert_true(ba.init(bb.b, vocab, false)); + const llama_batch & batch = ba.get_batch(); + t.assert_true("both kept", batch.token != nullptr && batch.embd != nullptr); + }); + + t.test("embd_row_width_override", [&](testing & t) { + // encoder input (e.g. eagle3/dflash) is wider than the decoder input + const uint32_t n_embd_enc = 6; + float embd[2*n_embd_enc]; + for (int i = 0; i < 2*6; ++i) { + embd[i] = (float) i; + } + + llama_batch lb = {}; + lb.n_tokens = 2; + lb.embd = embd; + + batch_builder bb(2, nullptr, 4, 1, 0, n_embd_enc); + llama_batch_compat::init(bb.b, lb, n_embd_enc); + + t.assert_equal((size_t) n_embd_enc, bb.b.n_embd); + t.assert_equal((size_t) 2*n_embd_enc, bb.b.embd.size()); + t.assert_equal((size_t) n_embd_enc, bb.b.tokens[1].embd_off); + t.assert_equal(6.0f, bb.b.embd[n_embd_enc]); + + llama_batch_allocr ba(1); + t.assert_true(ba.init(bb.b, vocab, false)); + llama_ubatch ub = ba.split_simple(2); + t.assert_equal("ubatch uses the encoder stride", 6.0f, ub.embd[n_embd_enc]); + }); +} + +static void test_mtp_embd_width(testing & t) { + t.test("mtp_uses_n_embd_out", [&](testing & t) { + llama_hparams hparams = {}; + hparams.n_embd = 64; + hparams.n_deepstack_layers = 2; // makes n_embd_inp() = 64 + 64*2 = 192 + hparams.n_embd_out_impl = 96; // makes n_embd_out() = 96 + + t.assert_equal("default context uses n_embd_inp (deepstack-aware)", + (size_t) 192, llama_batch_ext_select_n_embd_inp(LLAMA_CONTEXT_TYPE_DEFAULT, LLM_ARCH_LLAMA, hparams)); + + t.assert_equal("MTP context uses n_embd_out instead (target-model hidden state width)", + (size_t) 96, llama_batch_ext_select_n_embd_inp(LLAMA_CONTEXT_TYPE_MTP, LLM_ARCH_LLAMA, hparams)); + }); + + t.test("mtp_falls_back_to_n_embd_when_no_override", [&](testing & t) { + llama_hparams hparams = {}; + hparams.n_embd = 64; // no deepstack, no n_embd_out_impl override + + t.assert_equal((size_t) 64, llama_batch_ext_select_n_embd_inp(LLAMA_CONTEXT_TYPE_DEFAULT, LLM_ARCH_LLAMA, hparams)); + t.assert_equal((size_t) 64, llama_batch_ext_select_n_embd_inp(LLAMA_CONTEXT_TYPE_MTP, LLM_ARCH_LLAMA, hparams)); + }); + + t.test("dflash_uses_n_embd_inp_enc", [&](testing & t) { + llama_hparams hparams = {}; + hparams.n_embd = 64; + hparams.n_embd_inp_enc_impl = 128; // makes n_embd_inp_enc() = 128 + hparams.n_embd_out_impl = 96; // makes n_embd_out() = 96 + + t.assert_equal("DFlash uses the encoder input width", + (size_t) 128, llama_batch_ext_select_n_embd_inp(LLAMA_CONTEXT_TYPE_DEFAULT, LLM_ARCH_DFLASH, hparams)); + + t.assert_equal("other archs ignore n_embd_inp_enc", + (size_t) 64, llama_batch_ext_select_n_embd_inp(LLAMA_CONTEXT_TYPE_DEFAULT, LLM_ARCH_LLAMA, hparams)); + + t.assert_equal("MTP takes precedence over DFlash", + (size_t) 96, llama_batch_ext_select_n_embd_inp(LLAMA_CONTEXT_TYPE_MTP, LLM_ARCH_DFLASH, hparams)); + }); +} + int main(int argc, char ** argv) { testing t; @@ -665,10 +1057,13 @@ int main(int argc, char ** argv) { t.set_filter(argv[1]); } - t.test("init", test_init); - t.test("split", test_split); - t.test("keep_tail", test_keep_tail); - t.test("mrope", test_mrope); + t.test("init", test_init); + t.test("content_types", test_content_types); + t.test("compat", test_compat); + t.test("split", test_split); + t.test("keep_tail", test_keep_tail); + t.test("mrope", test_mrope); + t.test("mtp_embd_width", test_mtp_embd_width); return t.summary(); } diff --git a/tests/test-rpc-multi-server.cpp b/tests/test-rpc-multi-server.cpp index 4502e2ce7..f253b6e14 100644 --- a/tests/test-rpc-multi-server.cpp +++ b/tests/test-rpc-multi-server.cpp @@ -17,7 +17,7 @@ int main(int argc, char ** argv) { GGML_ASSERT(backend_b != nullptr); ggml_init_params params = { - /* .mem_size = */ ggml_tensor_overhead() + ggml_graph_overhead_custom(1, false), + /* .mem_size = */ 3*ggml_tensor_overhead() + ggml_graph_overhead_custom(1, false), /* .mem_buffer = */ nullptr, /* .no_alloc = */ true, }; @@ -40,6 +40,28 @@ int main(int argc, char ** argv) { ggml_backend_rpc_get_device_memory(endpoint_b, 0, &free_mem, &total_mem); GGML_ASSERT(total_mem > 0); ggml_backend_buffer_free(buffer); + + // Two tensors with the same ne[] but different nb[] must not share a cached alloc size. + // ref: https://github.com/ggml-org/llama.cpp/issues/28360 + ggml_backend_buffer_type_t buft = ggml_backend_rpc_buffer_type(endpoint_a, 0); + GGML_ASSERT(buft != nullptr); + + // MUL_MAT may need extra memory, so the size is read from the server [TAG_ALLOC_SIZE_EXPAND] + ggml_tensor * packed = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 64, 64); + packed->op = GGML_OP_MUL_MAT; + + // same ne[], twice the row stride + ggml_tensor * strided = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 64, 64); + strided->op = GGML_OP_MUL_MAT; + strided->nb[1] = 2*strided->nb[1]; + strided->nb[2] = strided->ne[1]*strided->nb[1]; + strided->nb[3] = strided->nb[2]; + GGML_ASSERT(ggml_nbytes(strided) > ggml_nbytes(packed)); + + // ask for the packed tensor first, so a cache keyed without nb[] holds the smaller size + GGML_ASSERT(ggml_backend_buft_get_alloc_size(buft, packed) >= ggml_nbytes(packed)); + GGML_ASSERT(ggml_backend_buft_get_alloc_size(buft, strided) >= ggml_nbytes(strided)); + ggml_free(ctx); ggml_backend_free(backend_b); ggml_backend_free(backend_a); diff --git a/tools/server/tests/tests.sh b/tools/server/tests/tests.sh index 1d1415663..c7d5d7805 100755 --- a/tools/server/tests/tests.sh +++ b/tools/server/tests/tests.sh @@ -6,7 +6,7 @@ cd $SCRIPT_DIR set -eu -WORKERS="${PYTEST_WORKERS:-auto}" +WORKERS="${PYTEST_WORKERS:-4}" if [ $# -lt 1 ] then diff --git a/tools/ui/src/lib/constants/svg-blocks.constants.ts b/tools/ui/src/lib/constants/svg-blocks.constants.ts index 705800c26..785e5a57f 100644 --- a/tools/ui/src/lib/constants/svg-blocks.constants.ts +++ b/tools/ui/src/lib/constants/svg-blocks.constants.ts @@ -41,6 +41,8 @@ export const SVG = { * so an author