wmantly/llama.cpp-xrip
1
0 Fork 0
Code Issues Pull Requests Actions 32 Packages Projects Releases Wiki Activity

Merge upstream/master into master

Browse Source 
Resolve llama-context.cpp: keep compute lock on the new llama_batch_ext
encode/decode, and reset both gf_res_prev slots in the shared compute path.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
This commit is contained in:
xripandClaude Opus 5.5 committed 2026-09-25 11:33:07 +03:00
2 parents 80034be9a8 d028c697b5
commit 74facb828e
57 files changed
+2947 -983
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

.github/workflows/build-cuda-ubuntu.yml
+4 -4
View File
@@ -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' }}
.github/workflows/server-sanitize.yml
+3 -3
View File
@@ -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
.pi/gg/SYSTEM.md
+1
View File
@@ -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
ci/README-MUSA.md
+1 -1
View File
@@ -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:
ci/run.sh
+2 -2
View File
@@ -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
common/common.cpp
+30 -7
View File
@@ -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<llama_token> & 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<llama_token*>(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<llama_token*>(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;
}
common/common.h
+5 -1
View File
@@ -1023,6 +1023,10 @@ void common_batch_add(
const std::vector<llama_seq_id> & 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<llama_token> & all_tokens,
const llama_tokens & all_tokens,
int n_new,
int & n_past,
int n_batch,
conversion/mimo.py
+31 -3
View File
@@ -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
docs/backend/SYCL.md
+3
View File
@@ -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.<br>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. |
docs/build.md
+2 -2
View File
@@ -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
ggml/CMakeLists.txt
+1 -1
View File
@@ -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/")
ggml/src/ggml-cuda/common.cuh
+6 -15
View File
@@ -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<int width = WARP_SIZE>
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) || \
ggml/src/ggml-cuda/conv2d-dw.cu
+27 -13
View File
@@ -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 <typename T, typename Layout>
__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<float>(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<float, whcn_layout><<<blocks, CUDA_CONV2D_DW_BLOCK_SIZE, 0, st>>>(
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<half, whcn_layout><<<blocks, CUDA_CONV2D_DW_BLOCK_SIZE, 0, st>>>(
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<float, whcn_layout><<<blocks, CUDA_CONV2D_DW_BLOCK_SIZE, 0, st>>>(
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<float, cwhn_layout><<<blocks, CUDA_CONV2D_DW_BLOCK_SIZE, 0, st>>>(
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<half, cwhn_layout><<<blocks, CUDA_CONV2D_DW_BLOCK_SIZE, 0, st>>>(
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<float, cwhn_layout><<<blocks, CUDA_CONV2D_DW_BLOCK_SIZE, 0, st>>>(
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");
}
ggml/src/ggml-cuda/fattn-mma-f16.cuh
-4
View File
@@ -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<DKQ, DV, ncols1, ncols2, use_logit_softcap, V_is_K_view, use_sparse_kernel>;
#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_ptr_t>(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<DKQ, DV, ncols1, ncols2, use_logit_softcap, V_is_K_view, use_sparse_kernel>;
#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_ptr_t>(fattn_kernel), cudaFuncAttributeMaxDynamicSharedMemorySize, nbytes_shared_total));
shared_memory_limit_raised[id] = true;
}
#endif // !defined(GGML_USE_MUSA)
}
launch_fattn<DV, ncols1, ncols2>
ggml/src/ggml-cuda/ggml-cuda.cu
+20 -12
View File
@@ -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;
ggml/src/ggml-cuda/mmq.cu
+5
View File
@@ -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;
}
ggml/src/ggml-cuda/norm.cu
+44 -11
View File
@@ -73,7 +73,7 @@ static __global__ void group_norm_f32(const float * x, float * dst, const int gr
}
}
template <int block_size, bool do_multiply = false, bool do_add = false>
template <int block_size, bool do_multiply = false, bool do_add = false, bool do_scale = false>
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 <bool do_scale = false>
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<true>(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;
ggml/src/ggml-cuda/norm.cuh
+2
View File
@@ -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,
ggml/src/ggml-cuda/ssm-scan.cu
+4 -4
View File
@@ -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 <cub/cub.cuh>
@@ -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));
ggml/src/ggml-cuda/topk-moe.cu
+5
View File
@@ -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;
}
ggml/src/ggml-cuda/vendors/hip.h
Vendored
+2 -2
View File
@@ -248,11 +248,11 @@
typedef __hip_bfloat16 nv_bfloat16;
typedef __hip_bfloat162 nv_bfloat162;
#if HIP_VERSION >= 60200000
#if HIP_VERSION >= 60300000
#include <hip/hip_fp8.h>
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)));
ggml/src/ggml-cuda/vendors/musa.h
Vendored
+5
View File
@@ -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;
ggml/src/ggml-cuda/wkv.cu
-2
View File
@@ -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];
}
ggml/src/ggml-hexagon/ggml-hexagon.cpp
+233 -13
View File
@@ -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");
ggml/src/ggml-hexagon/htp/concat-ops.c
+132 -55
View File
@@ -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;
ggml/src/ggml-hexagon/htp/cpy-ops.c
+17 -2
View File
@@ -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;
ggml/src/ggml-hexagon/htp/hmx-mm-kernels-tiled.h
+105 -92
View File
@@ -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,
ggml/src/ggml-hexagon/htp/htp-ops.h
+1
View File
@@ -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,
ggml/src/ggml-hexagon/htp/hvx-mm-kernels-tiled.h
+212 -47
View File
@@ -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) {
ggml/src/ggml-hexagon/htp/matmul-ops.c
+340 -310
View File
File diff suppressed because it is too large. Load diff
ggml/src/ggml-hexagon/htp/matmul-ops.h
+88 -53
View File
@@ -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)) {
ggml/src/ggml-metal/ggml-metal-context.m
+6 -4
View File
@@ -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) {
ggml/src/ggml-metal/ggml-metal-ops.cpp
+23 -36
View File
@@ -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;
ggml/src/ggml-metal/kernels/fa.metal
+27 -10
View File
@@ -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<kd4_t, k4_t>::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<vd4_t, v4_t>::value) {
device const v4_t * pv4 = (device const v4_t *) (v + i11*args.nb21);
ggml/src/ggml-rpc/ggml-rpc.cpp
+9 -2
View File
@@ -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<std::mutex> lock(cache_mutex);
auto it = cache.find(cache_hash);
if (it != cache.end()) {
return it->second;
return std::max<size_t>(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<size_t>(response.alloc_size, min_size);
}
return ggml_nbytes(tensor);
ggml/src/ggml-sycl/fattn-sparse.cpp
+267
View File
@@ -0,0 +1,267 @@
#include "fattn.hpp"
#include "fattn-sparse.hpp"
#include <cstdint>
#include <cstdio>
#include <cstring>
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<int32_t,
sycl::memory_order::relaxed,
sycl::memory_scope::device,
sycl::access::address_space::global_space> 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<int32_t> idx_alloc(ctx.pool(), (size_t) n_kv_g);
ggml_sycl_pool_alloc<int32_t> cnt_alloc(ctx.pool(), 1);
ggml_sycl_pool_alloc<uint8_t> k_alloc(ctx.pool(), (size_t) n_head_k * n_kv_g * k_row);
ggml_sycl_pool_alloc<uint8_t> v_alloc(ctx.pool(), (size_t) n_head_k * n_kv_g * v_row);
ggml_sycl_pool_alloc<sycl::half> 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;
}
ggml/src/ggml-sycl/fattn-sparse.hpp
+10
View File
@@ -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
ggml/src/ggml-sycl/fattn.cpp
+6 -1
View File
@@ -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);
ggml/src/ggml-sycl/ggml-sycl.cpp
+12
View File
@@ -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__);
ggml/src/ggml.c
+2 -2
View File
@@ -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) {
include/llama-cpp.h
+5
View File
@@ -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, llama_model_deleter> llama_model_ptr;
typedef std::unique_ptr<llama_context, llama_context_deleter> llama_context_ptr;
typedef std::unique_ptr<llama_sampler, llama_sampler_deleter> llama_sampler_ptr;
typedef std::unique_ptr<llama_adapter_lora, llama_adapter_lora_deleter> llama_adapter_lora_ptr;
typedef std::unique_ptr<llama_batch_ext, llama_batch_ext_deleter> llama_batch_ext_ptr;
include/llama.h
+90
View File
@@ -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)
scripts/sync-ggml.last
+1 -1
View File
@@ -1 +1 @@
e565a8f4ce2e462c4973a24c51098dd3c81c0256
353b63b439f27ab2cc19dac97ab1681ba6d2d084
src/llama-batch.cpp
+440 -103
View File
@@ -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 <cassert>
#include <cstring>
@@ -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<llama_pos> 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;
}
src/llama-batch.h
+76 -7
View File
@@ -2,6 +2,7 @@
#include "llama.h"
#include "llama-arch.h"
#include "llama-cparams.h"
#include <array>
@@ -10,6 +11,7 @@
#include <bitset>
#include <memory>
#include <unordered_map>
#include <unordered_set>
// keep this struct lightweight
struct llama_ubatch {
@@ -68,19 +70,71 @@ struct llama_ubatch {
std::shared_ptr<data_t> 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<llama_seq_id> seq_ids;
std::array<llama_pos, GGML_MROPE_SECTIONS> pos = {0, 0, 0, 0};
};
std::vector<token> tokens;
std::vector<float> 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<llama_seq_id, 1> seq_id_0 = {{ 0 }}; // default sequence id
std::vector<llama_token> token_vec; // owned token IDs built from llama_batch_ext
std::vector<float> embd_vec; // owned embeddings built from llama_batch_ext
std::vector<llama_seq_id> seq_id_data; // flat storage for seq_id pointers below
std::vector<llama_pos> pos;
std::vector<int32_t> 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);
};
src/llama-context.cpp
+59 -34
View File
@@ -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<llama_s
return false; // all sequences use backend sampling
}
int llama_context::decode(const llama_batch & batch_inp) {
int llama_context::decode(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 (!memory) {
LLAMA_LOG_DEBUG("%s: cannot decode batches with this context (calling encode() instead)\n", __func__);
return encode(batch_inp);
}
if (batch_inp.n_tokens == 0) {
if (batch_inp.tokens.empty()) {
LLAMA_LOG_ERROR("%s: n_tokens == 0\n", __func__);
return -1;
}
if (batch_inp.n_embd > 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<int32_t> 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
//
src/llama-context.h
+4
View File
@@ -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);
src/llama-grammar.cpp
+5 -1
View File
@@ -194,7 +194,11 @@ static std::pair<uint32_t, const char *> 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<uint32_t>::max()) {
throw std::runtime_error(std::string("parsed token id is too big at ") + pos);
}
uint32_t token_id = static_cast<uint32_t>(id);
pos = int_end;
if (*pos != ']') {
throw std::runtime_error(std::string("expecting ']' at ") + pos);
src/llama-hparams.cpp
+2 -1
View File
@@ -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 {
src/llama-model.cpp
+22 -8
View File
@@ -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};
}
}
src/models/mimo2.cpp
+3 -1
View File
@@ -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);
tests/test-backend-ops.cpp
+4
View File
@@ -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 {
tests/test-batch-alloc.cpp
+506 -111
View File
@@ -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<float> embd;
std::vector<llama_pos> pos;
std::vector<int32_t> n_seq_id;
std::vector<int8_t> logits;
llama_batch_ext b;
std::vector<std::vector<llama_seq_id>> seq;
std::vector<llama_seq_id *> 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<llama_seq_id> 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<float> row(int32_t i, uint32_t width) const {
std::vector<float> 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<llama_seq_id> 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<llama_seq_id> 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<llama_pos> 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<llama_pos> 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();
}
tests/test-rpc-multi-server.cpp
+23 -1
View File
@@ -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);
tools/server/tests/tests.sh
+1 -1
View File
@@ -6,7 +6,7 @@ cd $SCRIPT_DIR
set -eu
WORKERS="${PYTEST_WORKERS:-auto}"
WORKERS="${PYTEST_WORKERS:-4}"
if [ $# -lt 1 ]
then
tools/ui/src/lib/constants/svg-blocks.constants.ts
+2
View File
@@ -41,6 +41,8 @@ export const SVG = {
* so an author <style> stays scoped to that root and can not reach the page.
*/
SANITIZE_CONFIG: {
ADD_ATTR: ['calcMode', 'from', 'to'],
ADD_TAGS: ['animate', 'set', 'use'],
FORBID_TAGS: ['foreignObject', 'script'],
USE_PROFILES: { svg: true, svgFilters: true }
},
tools/ui/src/lib/utils/sanitize-svg.ts
+11
View File
@@ -1,6 +1,17 @@
import { SVG } from '$lib/constants';
import DOMPurify from 'dompurify';
/**
* animate and set can retarget href or xlink:href to a javascript: uri through
* to, from, by or values, none of which DOMPurify checks as a uri. Dropping
* attributeName in that case leaves the animation inert.
*/
DOMPurify.addHook('uponSanitizeAttribute', (_node, data) => {
if (data.attrName === 'attributename' && /href$/i.test(data.attrValue.trim())) {
data.keepAttr = false;
}
});
/**
* Sanitizes a raw svg string for safe inline rendering.
* Returns the cleaned svg markup, or an empty string when the input is not a