Skip to content

MTP: option to use re-quantized output tensor for better TG performance #1809

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

Merged
ikawrakow merged 3 commits into from
May 16, 2026
Merged

MTP: option to use re-quantized output tensor for better TG performance #1809

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
ece2adf
Option to use re-quantized output tensor for MTP
ikawrakow May 16, 2026
c7668b7
Remove quantize extra output option
ikawrakow May 16, 2026
b5da743
Handle interleaved types
ikawrakow May 16, 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
20 changes: 20 additions & 0 deletions common/common.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1727,6 +1727,11 @@ bool gpt_params_find_arg(int argc, char ** argv, const std::string & arg, gpt_pa
}
return true;
}
if (arg == "--mtp-requantized-output-tensor" || arg == "-mtprot") {

@ahoplock ahoplock May 16, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

I know it's already merged at this point but you have this as --mtp-requantize-output-tensor (missing 'd') in line 3036 and in the PR description.

CHECK_ARG
params.extra_output_type = argv[i];
return true;
}
if (arg == "-ctkd" || arg == "--cache-type-k-draft") {
params.speculative.cache_type_k = argv[++i];
return true;
Expand Down Expand Up @@ -3028,6 +3033,7 @@ void gpt_params_print_usage(int /*argc*/, char ** argv, const gpt_params & param
options.push_back({ "*", "-ctv-last, --cache-type-k-last TYPE,N", "KV cache data type for the last N layers of K (default: %s,-1)", params.type_k_last.c_str() });
options.push_back({ "*", "-ctv-first, --cache-type-v-first TYPE,N", "KV cache data type for the first N layers of V (default: %s,-1)", params.type_v_first.c_str() });
options.push_back({ "*", "-ctk-last, --cache-type-v-last TYPE,N", "KV cache data type for the last N layers of V (default: %s,-1)", params.type_v_last.c_str() });
options.push_back({ "*", "-mtprot, --mtp-requantize-output-tensor type", "Use output requantized to type for MTP (default: %s)", params.extra_output_type.c_str() });
options.push_back({ "*", "-ctkd, --cache-type-k-draft TYPE", "KV cache data type for K for the draft model" });
options.push_back({ "*", "-ctvd, --cache-type-v-draft TYPE", "KV cache data type for V for the draft model" });

Expand Down Expand Up @@ -3926,6 +3932,17 @@ static std::pair<int, int> get_batch_ubatch(const gpt_params & params) {
return {n_batch, n_ubatch};
}

static ggml_type parse_ggml_type(const char * arg) {
for (int j = 0; j < GGML_TYPE_COUNT; ++j) {
auto type = ggml_type(j);
const auto * name = ggml_type_name(type);
if (name && strcmp(arg, name) == 0) {
return type;
}
}
return GGML_TYPE_COUNT;
}

struct llama_model_params common_model_params_to_llama(const gpt_params & params) {
auto mparams = llama_model_default_params();
mparams.devices = params.devices.c_str();
Expand All @@ -3948,6 +3965,9 @@ struct llama_model_params common_model_params_to_llama(const gpt_params & params
mparams.type_k_last = kv_cache_type_from_str(params.type_k_last );
mparams.type_v_first = kv_cache_type_from_str(params.type_v_first);
mparams.type_v_last = kv_cache_type_from_str(params.type_v_last );
if (!params.extra_output_type.empty()) {
mparams.extra_output_type = parse_ggml_type(params.extra_output_type.c_str());
}
mparams.n_k_first = params.n_k_first;
mparams.n_k_last = params.n_k_last;
mparams.n_v_first = params.n_v_first;
Expand Down
2 changes: 2 additions & 0 deletions common/common.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -443,6 +443,8 @@ struct gpt_params {
int32_t n_v_first = -1;
int32_t n_v_last = -1;

std::string extra_output_type = "";

// multimodal models (see examples/mtmd)
common_params_model mmproj;
bool mmproj_use_gpu = true; // use GPU for multimodal model
Expand Down
98 changes: 90 additions & 8 deletions ggml/src/iqk/iqk_quantize.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -240,6 +240,81 @@ void IQ1BNQuantizer::quantize_one_row_2bn(const float * src, block_iq2_bn * y, i
}
}

static inline int num_rows([[maybe_unused]] ggml_type type) {
#ifdef HAVE_FANCY_SIMD
switch (type) {
case GGML_TYPE_Q2_K_R4:
case GGML_TYPE_Q3_K_R4:
case GGML_TYPE_Q6_K_R4:
case GGML_TYPE_IQ2_K_R4:
case GGML_TYPE_IQ3_K_R4:
case GGML_TYPE_IQ4_K_R4:
case GGML_TYPE_IQ5_K_R4:
case GGML_TYPE_IQ4_KS_R4:
case GGML_TYPE_IQ5_KS_R4:
case GGML_TYPE_IQ2_XXS_R4:
case GGML_TYPE_IQ2_XS_R4:
case GGML_TYPE_IQ2_S_R4:
case GGML_TYPE_IQ3_XXS_R4:
case GGML_TYPE_IQ1_S_R4:
case GGML_TYPE_IQ1_M_R4:
case GGML_TYPE_IQ3_S_R4: return 4;
case GGML_TYPE_IQ4_NL_R4:
case GGML_TYPE_Q5_0_R4:
case GGML_TYPE_Q6_0_R4:
case GGML_TYPE_IQ2_BN_R4:
case GGML_TYPE_IQ4_XS_R8:
case GGML_TYPE_Q4_K_R4:
case GGML_TYPE_Q5_K_R4:
case GGML_TYPE_Q8_KV:
case GGML_TYPE_Q8_KV_R8:
case GGML_TYPE_Q8_K_R8: return 8;
case GGML_TYPE_Q4_0_R8:
case GGML_TYPE_Q8_0_R8:
case GGML_TYPE_Q8_1:
case GGML_TYPE_Q8_K_R16:
case GGML_TYPE_BF16_R16: return 16;
default: return 1;
}
#else
switch (type) {
case GGML_TYPE_Q2_K_R4:
case GGML_TYPE_Q3_K_R4:
case GGML_TYPE_Q4_K_R4:
case GGML_TYPE_Q5_K_R4:
case GGML_TYPE_Q6_K_R4:
case GGML_TYPE_Q5_0_R4:
case GGML_TYPE_Q6_0_R4:
case GGML_TYPE_IQ4_NL_R4:
case GGML_TYPE_IQ2_K_R4:
case GGML_TYPE_IQ3_K_R4:
case GGML_TYPE_IQ4_K_R4:
case GGML_TYPE_IQ5_K_R4:
case GGML_TYPE_IQ4_KS_R4:
case GGML_TYPE_IQ5_KS_R4:
case GGML_TYPE_IQ2_XXS_R4:
case GGML_TYPE_IQ2_XS_R4:
case GGML_TYPE_IQ2_S_R4:
case GGML_TYPE_IQ3_XXS_R4:
case GGML_TYPE_IQ3_S_R4:
case GGML_TYPE_IQ1_S_R4:
case GGML_TYPE_IQ1_M_R4:
case GGML_TYPE_IQ2_BN_R4: return 4;
case GGML_TYPE_IQ4_XS_R8:
case GGML_TYPE_Q4_0_R8:
case GGML_TYPE_Q8_0_R8:
case GGML_TYPE_Q8_KV:
case GGML_TYPE_Q8_KV_R8:
case GGML_TYPE_Q8_1:
case GGML_TYPE_Q8_K_R8: return 8;
case GGML_TYPE_Q8_K_R16:
case GGML_TYPE_BF16_R16: return 16;
default: return 1;
}
#endif
}


}

void iqk_quantize_any(int from_type, int to_type,
Expand All @@ -251,21 +326,28 @@ void iqk_quantize_any(int from_type, int to_type,
GGML_ASSERT(ggml_type_size(type_x) == nb0);
auto type_y = ggml_type(to_type);
auto row_size_y = ggml_row_size(type_y, ne0);
int64_t nrows = ne1*ne2*ne3;
auto n_interleaved = num_rows(type_y);
GGML_ASSERT(ne1 % n_interleaved == 0);
int64_t ne1i = ne1/n_interleaved;
int64_t nrows = ne1i*ne2*ne3;
int64_t nrows_per_thread = (nrows + nth - 1)/nth;
int64_t first_row = nrows_per_thread*ith;
if (first_row >= nrows) return;
int64_t last_row = std::min(first_row + nrows_per_thread, nrows);
for (int64_t row = first_row; row < last_row; ++row) {
int64_t i3 = row/(ne1*ne2);
int64_t i2 = (row - i3*ne1*ne2)/ne1;
int64_t i1 = row - i3*ne1*ne2 - i2*ne1;
const char * cx = (const char *)x + i1*nb1 + i2*nb2 + i3*nb3;
int64_t i3 = row/(ne1i*ne2);
int64_t i2 = (row - i3*ne1i*ne2)/ne1i;
int64_t i1 = row - i3*ne1i*ne2 - i2*ne1i;
auto cx = (const char *)x + i1*n_interleaved*nb1 + i2*nb2 + i3*nb3;
auto cy = (char *)y + (i3*ne1*ne2 + i2*ne1 + i1*n_interleaved)*row_size_y;
// TODO: special case common types such as f16, q8_0
// (although the performance gains may be too small to justify the added complexity)
to_float((const void *)cx, (float *)work_buffer, ne0);
auto cy = (char *)y + (i3*ne1*ne2 + i2*ne1 + i1)*row_size_y;
from_float((const float *)work_buffer, (void *)cy, ne0);
if (type_x != GGML_TYPE_F32) {
to_float((const void *)cx, (float *)work_buffer, ne0*n_interleaved);
from_float((const float *)work_buffer, (void *)cy, ne0*n_interleaved);
} else {
from_float((const float *)cx, (void *)cy, ne0*n_interleaved);
}
}
}

Expand Down
2 changes: 2 additions & 0 deletions include/llama.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -394,6 +394,8 @@ extern "C" {
int32_t n_v_first;
int32_t n_v_last;

enum ggml_type extra_output_type;

// proportion of the model (layers or rows) to offload to each GPU, size: llama_max_devices()
const float * tensor_split;

Expand Down
4 changes: 2 additions & 2 deletions src/graphs/build_qwen35.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -238,7 +238,7 @@ struct ggml_tensor * llm_build_context::build_qwen35moe_mtp(

cb(cur, "result_norm", -1);

cur = build_output(lctx, ctx0, cur, model.output, mtp_layer.nextn.shared_head_norm, cb);
cur = build_output(lctx, ctx0, cur, model.output_mtp, mtp_layer.nextn.shared_head_norm, cb);
cb(cur, "result_output", -1);

return cur;
Expand Down Expand Up @@ -317,7 +317,7 @@ struct ggml_tensor * llm_build_context::build_qwen35_mtp(
cb(cur, "result_norm", -1);

//cur = build_output(lctx, ctx0, cur, model.output, nullptr, cb);
cur = build_output(lctx, ctx0, cur, model.output, mtp_layer.nextn.shared_head_norm, cb);
cur = build_output(lctx, ctx0, cur, model.output_mtp, mtp_layer.nextn.shared_head_norm, cb);
cb(cur, "result_output", -1);

return cur;
Expand Down
20 changes: 18 additions & 2 deletions src/llama-load-tensors.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1520,6 +1520,13 @@ bool create_tensors_helper::create_qwen35moe_tensors(const LLM_TN & tn) {
if (model.output == NULL) {
model.output = create_tensor(ctx_output, tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, llama_model_loader::TENSOR_DUPLICATED);
}
if (model.mtp) {
model.output_mtp = create_tensor(ctx_output, "output_extra.weight", {n_embd, n_vocab},
llama_model_loader::TENSOR_NOT_REQUIRED);
if (!model.output_mtp) {
model.output_mtp = model.output;
}
}
}

const int64_t n_ff_exp = hparams.n_ff_exp ? hparams.n_ff_exp : n_ff / n_expert_used;
Expand Down Expand Up @@ -1614,9 +1621,18 @@ bool create_tensors_helper::create_qwen35_tensors(const LLM_TN & tn) {
// output
{
model.output_norm = create_tensor(ctx_output, tn(LLM_TENSOR_OUTPUT_NORM, "weight"), {n_embd});
model.output = create_tensor(ctx_output, tn(LLM_TENSOR_OUTPUT, "weight"), {n_embd, n_vocab}, llama_model_loader::TENSOR_NOT_REQUIRED);
model.output = create_tensor(ctx_output, tn(LLM_TENSOR_OUTPUT, "weight"), {n_embd, n_vocab},
llama_model_loader::TENSOR_NOT_REQUIRED);
if (model.output == NULL) {
model.output = create_tensor(ctx_output, tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, llama_model_loader::TENSOR_DUPLICATED);
model.output = create_tensor(ctx_output, tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab},
llama_model_loader::TENSOR_DUPLICATED);
}
if (model.mtp) {
model.output_mtp = create_tensor(ctx_output, "output_extra.weight", {n_embd, n_vocab},
llama_model_loader::TENSOR_NOT_REQUIRED);
if (!model.output_mtp) {
model.output_mtp = model.output;
}
}
}

Expand Down
4 changes: 4 additions & 0 deletions src/llama-model-loader.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -70,6 +70,10 @@ struct llama_model_loader {
offs = gguf_get_data_offset(gguf_ctx) + gguf_get_tensor_offset(gguf_ctx, tensor_idx);

if (offs + ggml_nbytes(tensor) < offs || offs + ggml_nbytes(tensor) > file->size()) {
auto data_offset = gguf_get_data_offset(gguf_ctx);
auto tensor_offset = gguf_get_tensor_offset(gguf_ctx, tensor_idx);
fprintf(stderr, "Error while loading tensor %s: offs = %zu (%zu, %zu), size: %zu, file size: %zu\n", name,
offs, data_offset, tensor_offset, ggml_nbytes(tensor), file->size());
throw std::runtime_error(format("tensor '%s' data is not within the file bounds, model is corrupted or incomplete", name));
}
}
Expand Down
3 changes: 3 additions & 0 deletions src/llama-model.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -415,6 +415,9 @@ struct llama_model {
struct ggml_tensor * output;
struct ggml_tensor * output_b;
struct ggml_tensor * output_norm_enc;
struct ggml_tensor * output_mtp = nullptr;

std::unique_ptr<ggml_tensor> output_mtp_ptr;

llama_split_tensor split_output;
llama_split_tensor split_output_norm;
Expand Down
Loading