Skip to content

Refactor batch size handling for offloading host operations to device #1214

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
DocShotgun wants to merge 2 commits into from
Closed

Refactor batch size handling for offloading host operations to device #1214

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
a510272
Enhance op offload batch size handling
DocShotgun Feb 2, 2026
1abb68a
Read env var once and pull from device context
DocShotgun Feb 2, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 2 additions & 2 deletions ggml/CMakeLists.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -127,8 +127,8 @@ set (GGML_CUDA_KQUANTS_ITER "2" CACHE STRING
"ggml: iters./thread per block for Q2_K/Q6_K")
set (GGML_CUDA_PEER_MAX_BATCH_SIZE "128" CACHE STRING
"ggml: min batch size for GPU offload")
set (GGML_CUDA_MIN_BATCH_OFFLOAD "32" CACHE STRING
"ggml: max. batch size for using peer access")
set (GGML_OP_OFFLOAD_HEURISTIC_MIN "32" CACHE STRING
"ggml: min. batch size for offload heuristic")
option(GGML_CUDA_NO_PEER_COPY "ggml: do not use peer to peer copies" OFF)
option(GGML_CUDA_NO_VMM "ggml: do not try to use CUDA VMM" OFF)
option(GGML_CUDA_FA_ALL_QUANTS "ggml: compile all quants for FlashAttention" OFF)
Expand Down
2 changes: 1 addition & 1 deletion ggml/src/CMakeLists.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -379,7 +379,7 @@ if (GGML_CUDA)
add_compile_definitions(GGML_CUDA_MMV_Y=${GGML_CUDA_MMV_Y})
add_compile_definitions(K_QUANTS_PER_ITERATION=${GGML_CUDA_KQUANTS_ITER})
add_compile_definitions(GGML_CUDA_PEER_MAX_BATCH_SIZE=${GGML_CUDA_PEER_MAX_BATCH_SIZE})
add_compile_definitions(GGML_CUDA_MIN_BATCH_OFFLOAD=${GGML_CUDA_MIN_BATCH_OFFLOAD})
add_compile_definitions(GGML_OP_OFFLOAD_HEURISTIC_MIN=${GGML_OP_OFFLOAD_HEURISTIC_MIN})
add_compile_definitions(GGML_CUDA_FUSION=${GGML_CUDA_FUSION})

if (GGML_CUDA_USE_GRAPHS)
Expand Down
25 changes: 23 additions & 2 deletions ggml/src/ggml-cann.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1792,10 +1792,26 @@ GGML_CALL static bool ggml_backend_cann_supports_buft(
*/
GGML_CALL static bool ggml_backend_cann_offload_op(ggml_backend_t backend,
const ggml_tensor* op) {
const int min_batch_size = 32;
GGML_UNUSED(backend);
ggml_backend_cann_context* ctx = (ggml_backend_cann_context*)backend->context;

if (ctx->op_offload_min_batch_size >= 0) {
return op->ne[1] >= ctx->op_offload_min_batch_size && op->op != GGML_OP_GET_ROWS;
}

// No env var: use MoE-aware heuristic
int min_batch_size = GGML_OP_OFFLOAD_HEURISTIC_MIN;

if (op->op == GGML_OP_MUL_MAT_ID || op->op == GGML_OP_MOE_FUSED_UP_GATE) {
auto ids = op->op == GGML_OP_MUL_MAT_ID ? op->src[2] : op->src[3];
int64_t batch_size = op->ne[2];
if (batch_size < min_batch_size) return false;
int64_t n_experts_tot = op->src[0]->ne[2];
int64_t n_experts_active = ids->ne[0];
return batch_size * n_experts_active >= min_batch_size * n_experts_tot;
}

return op->ne[1] >= min_batch_size && op->op != GGML_OP_GET_ROWS;
GGML_UNUSED(backend);
}

/**
Expand Down Expand Up @@ -1944,6 +1960,11 @@ GGML_CALL ggml_backend_t ggml_backend_cann_init(int32_t device) {
return nullptr;
}

const char * env_value = getenv("GGML_OP_OFFLOAD_MIN_BATCH");
if (env_value) {
ctx->op_offload_min_batch_size = atoi(env_value);
}

ggml_backend_t cann_backend =
new ggml_backend{/* .guid = */ ggml_backend_cann_guid(),
/* .interface = */ ggml_backend_cann_interface,
Expand Down
1 change: 1 addition & 0 deletions ggml/src/ggml-cann/common.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -212,6 +212,7 @@ struct ggml_backend_cann_context {
int32_t device; /**< Device ID. */
std::string name; /**< Name of the device. */
aclrtEvent copy_event = nullptr; /**< Event for managing copy operations. */
int op_offload_min_batch_size = -1;

aclrtStream streams[GGML_CANN_MAX_STREAMS] = {
{nullptr}}; /**< Array of streams for the device. */
Expand Down
19 changes: 13 additions & 6 deletions ggml/src/ggml-cuda.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4409,8 +4409,12 @@ GGML_CALL static bool ggml_backend_cuda_supports_buft(ggml_backend_t backend, gg

GGML_CALL static bool ggml_backend_cuda_offload_op(ggml_backend_t backend, const ggml_tensor * op) {
auto ctx = (const ggml_backend_cuda_context *)backend->context;
int min_batch_size = ctx->offload_batch_size; //originally: GGML_CUDA_MIN_BATCH_OFFLOAD;

if (ctx->op_offload_min_batch_size >= 0) {

@ikawrakow ikawrakow Feb 2, 2026

Copy link
Copy Markdown
Owner

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

If, for whatever reason, somebody decided to use this new way of determining minimum offload batch size, it overrides the existing functionality for MoE models below.

Also, the batch size for MoE related ops (GGML_OP_MUL_MAT_ID and GGML_OP_MOE_FUSED_UP_GATE is not given by op->ne[1], it is op->ne[2] instead.

@DocShotgun DocShotgun Feb 2, 2026 •
edited
Loading

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The intention is to allow the user to override the heuristic with a raw minimum batch size value, so we skip the MoE calculation entirely if the user chooses to do so. If this is misguided in your view, then feel free to close this PR, as that is the main purpose of it.

Correct me if I'm misunderstanding here, but I do believe that the value that should be exposed for the user to adjust is the final batch size to be compared with the size of the input, and not the min_batch_size that subsequently undergoes the MoE adjustment with min_batch_size * total_experts / active_experts. If we are to bench speeds with op offload forced on/off to find the real break-even point on particular model/hardware, the value that is empirically determined should not be adjusted any further based on MoE. offload-batch-size could be adjusted to allow for this as well, without having to introduce a new env var.

return op->ne[1] >= ctx->op_offload_min_batch_size && op->op != GGML_OP_GET_ROWS;
}

// No env var: use MoE-aware heuristic
// Why do we want to do this? The heuristics that the batch must have more than min_batch_size tokens to be worth it
// offloading the required model weights comes from dense models. For MoE models, the average number of tokens
// each expert deals with in a batch is (active_experts / total_experts) * batch_size. Hence, according to the
Expand All @@ -4421,6 +4425,8 @@ GGML_CALL static bool ggml_backend_cuda_offload_op(ggml_backend_t backend, const
//
// as the condition for offloading model weights resinding in RAM to the GPU.
// In this case, the number of tokens is not as usual in op->ne[1] but rather in op->ne[2].
int min_batch_size = ctx->offload_batch_size;

if (op->op == GGML_OP_MUL_MAT_ID || op->op == GGML_OP_MOE_FUSED_UP_GATE) {
auto ids = op->op == GGML_OP_MUL_MAT_ID ? op->src[2] : op->src[3];
int64_t batch_size = op->ne[2];
Expand All @@ -4434,10 +4440,6 @@ GGML_CALL static bool ggml_backend_cuda_offload_op(ggml_backend_t backend, const

return op->ne[1] >= min_batch_size && op->op != GGML_OP_GET_ROWS;

// Original:
//return (op->ne[1] >= min_batch_size && op->op != GGML_OP_GET_ROWS) ||
// (op->ne[2] >= min_batch_size && (op->op == GGML_OP_MUL_MAT_ID || op->op == GGML_OP_MOE_FUSED_UP_GATE));

GGML_UNUSED(backend);
}

Expand Down Expand Up @@ -4524,7 +4526,7 @@ static ggml_guid_t ggml_backend_cuda_guid() {

struct cuda_params {
int fusion = GGML_CUDA_FUSION;
int offload_batch_size = GGML_CUDA_MIN_BATCH_OFFLOAD;
int offload_batch_size = GGML_OP_OFFLOAD_HEURISTIC_MIN;
int mmq_id_thresh = 32;
float fa_offset = 0;
#ifdef USE_CUDA_GRAPH
Expand Down Expand Up @@ -4618,6 +4620,11 @@ GGML_CALL ggml_backend_t ggml_backend_cuda_init(int device, [[maybe_unused]] con
return nullptr;
}

const char * env_value = getenv("GGML_OP_OFFLOAD_MIN_BATCH");
if (env_value) {
ctx->op_offload_min_batch_size = atoi(env_value);
}

ggml_backend_t cuda_backend = new ggml_backend {
/* .guid = */ ggml_backend_cuda_guid(),
/* .interface = */ ggml_backend_cuda_interface,
Expand Down
3 changes: 2 additions & 1 deletion ggml/src/ggml-cuda/common.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -851,7 +851,8 @@ struct ggml_backend_cuda_context {
cublasHandle_t cublas_handles[GGML_CUDA_MAX_DEVICES] = {nullptr};

int fusion = GGML_CUDA_FUSION;
int offload_batch_size = GGML_CUDA_MIN_BATCH_OFFLOAD;
int offload_batch_size = GGML_OP_OFFLOAD_HEURISTIC_MIN;
int op_offload_min_batch_size = -1;
int mmq_id_thresh = 32;
float fa_offset = 0.0f;
#ifdef USE_CUDA_GRAPH
Expand Down
26 changes: 24 additions & 2 deletions ggml/src/ggml-sycl.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5294,8 +5294,25 @@ GGML_CALL static bool ggml_backend_sycl_supports_op(ggml_backend_t backend, cons
}

GGML_CALL static bool ggml_backend_sycl_offload_op(ggml_backend_t backend, const ggml_tensor * op) {
const int min_batch_size = 32;
return op->ne[1] >= min_batch_size && op->op != GGML_OP_GET_ROWS && op->op != GGML_OP_MUL_MAT_ID;
auto ctx = (ggml_backend_sycl_context *)backend->context;

if (ctx->op_offload_min_batch_size >= 0) {
return op->ne[1] >= ctx->op_offload_min_batch_size && op->op != GGML_OP_GET_ROWS;
}

// No env var: use MoE-aware heuristic
int min_batch_size = GGML_OP_OFFLOAD_HEURISTIC_MIN;

if (op->op == GGML_OP_MUL_MAT_ID || op->op == GGML_OP_MOE_FUSED_UP_GATE) {
auto ids = op->op == GGML_OP_MUL_MAT_ID ? op->src[2] : op->src[3];
int64_t batch_size = op->ne[2];
if (batch_size < min_batch_size) return false;
int64_t n_experts_tot = op->src[0]->ne[2];
int64_t n_experts_active = ids->ne[0];
return batch_size * n_experts_active >= min_batch_size * n_experts_tot;
}

return op->ne[1] >= min_batch_size && op->op != GGML_OP_GET_ROWS;
GGML_UNUSED(backend);
}

Expand Down Expand Up @@ -5348,6 +5365,11 @@ GGML_CALL ggml_backend_t ggml_backend_sycl_init(int device) {
return nullptr;
};

const char * env_value = getenv("GGML_OP_OFFLOAD_MIN_BATCH");
if (env_value) {
ctx->op_offload_min_batch_size = atoi(env_value);
}

ggml_backend_t sycl_backend = new ggml_backend {
/* .guid = */ ggml_backend_sycl_guid(),
/* .interface = */ ggml_backend_sycl_interface,
Expand Down
1 change: 1 addition & 0 deletions ggml/src/ggml-sycl/common.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -257,6 +257,7 @@ struct ggml_tensor_extra_gpu {
struct ggml_backend_sycl_context {
int device;
std::string name;
int op_offload_min_batch_size = -1;

queue_ptr qptrs[GGML_SYCL_MAX_DEVICES][GGML_SYCL_MAX_STREAMS] = { { nullptr } };

Expand Down
27 changes: 24 additions & 3 deletions ggml/src/ggml-vulkan.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1051,6 +1051,7 @@ struct ggml_backend_vk_context {
vk_buffer prealloc_x, prealloc_y, prealloc_split_k;
vk::Fence fence, almost_ready_fence;
bool almost_ready_fence_pending {};
int op_offload_min_batch_size = -1;

vk_buffer buffer_pool[MAX_VK_BUFFERS];

Expand Down Expand Up @@ -10689,10 +10690,25 @@ static bool ggml_backend_vk_supports_buft(ggml_backend_t backend, ggml_backend_b
}

static bool ggml_backend_vk_offload_op(ggml_backend_t backend, const ggml_tensor * op) {
const int min_batch_size = 32;
ggml_backend_vk_context * ctx = (ggml_backend_vk_context *)backend->context;

if (ctx->op_offload_min_batch_size >= 0) {
return op->ne[1] >= ctx->op_offload_min_batch_size && op->op != GGML_OP_GET_ROWS;
}

// No env var: use MoE-aware heuristic
int min_batch_size = GGML_OP_OFFLOAD_HEURISTIC_MIN;

return (op->ne[1] >= min_batch_size && op->op != GGML_OP_GET_ROWS) ||
(op->ne[2] >= min_batch_size && op->op == GGML_OP_MUL_MAT_ID);
if (op->op == GGML_OP_MUL_MAT_ID || op->op == GGML_OP_MOE_FUSED_UP_GATE) {
auto ids = op->op == GGML_OP_MUL_MAT_ID ? op->src[2] : op->src[3];
int64_t batch_size = op->ne[2];
if (batch_size < min_batch_size) return false;
int64_t n_experts_tot = op->src[0]->ne[2];
int64_t n_experts_active = ids->ne[0];
return batch_size * n_experts_active >= min_batch_size * n_experts_tot;
}

return (op->ne[1] >= min_batch_size && op->op != GGML_OP_GET_ROWS);

UNUSED(backend);
}
Expand Down Expand Up @@ -10732,6 +10748,11 @@ ggml_backend_t ggml_backend_vk_init(size_t dev_num) {
ggml_backend_vk_context * ctx = new ggml_backend_vk_context;
ggml_vk_init(ctx, dev_num);

const char * env_value = getenv("GGML_OP_OFFLOAD_MIN_BATCH");
if (env_value) {
ctx->op_offload_min_batch_size = atoi(env_value);
}

ggml_backend_t vk_backend = new ggml_backend {
/* .guid = */ ggml_backend_vk_guid(),
/* .interface = */ ggml_backend_vk_interface,
Expand Down