Skip to content

fix(quantize): add NVFP4 default type mapping and scale tensors #22897

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

Open
t-timms wants to merge 2 commits into
base: master
Choose a base branch
from
+94 −0
Open

fix(quantize): add NVFP4 default type mapping and scale tensors #22897

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
56f62c4
fix(quantize): add NVFP4 default type mapping and scale tensors
t-timms May 10, 2026
f5df518
perf(quantize): compute optimal NVFP4 .scale via MSE minimization
t-timms May 10, 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
93 changes: 93 additions & 0 deletions src/llama-quant.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -203,6 +203,7 @@ struct tensor_metadata {
std::string remapped_imatrix_name;
bool allows_quantization;
bool requires_imatrix;
bool write_nvfp4_scales = false;
};

//
Expand Down Expand Up @@ -802,6 +803,7 @@ ggml_type llama_ftype_get_default_type(llama_ftype ftype) {
case LLAMA_FTYPE_MOSTLY_Q1_0: return GGML_TYPE_Q1_0;

case LLAMA_FTYPE_MOSTLY_MXFP4_MOE: return GGML_TYPE_MXFP4;
case LLAMA_FTYPE_MOSTLY_NVFP4: return GGML_TYPE_NVFP4;

// K-quants
case LLAMA_FTYPE_MOSTLY_Q2_K_S:
Expand Down Expand Up @@ -850,6 +852,42 @@ static void init_quantize_state_counters(quantize_state_impl & qs, std::vector<t
qs.n_ffn_down = qs.n_ffn_gate = qs.n_ffn_up = (int)qs.model.hparams.n_layer;
}

//
// NVFP4 scale tensor helpers
//

static void add_f32_scalar_tensor(struct gguf_context * ctx, const char * name) {
struct ggml_init_params iparams = {
/*.mem_size =*/ ggml_tensor_overhead(),
/*.mem_buffer =*/ NULL,
/*.no_alloc =*/ true,
};
struct ggml_context * gctx = ggml_init(iparams);
struct ggml_tensor * t = ggml_new_tensor_1d(gctx, GGML_TYPE_F32, 1);
ggml_set_name(t, name);
gguf_add_tensor(ctx, t);
ggml_free(gctx);
}

static void add_nvfp4_scale_tensors(struct gguf_context * ctx, const char * weight_name) {
std::string base(weight_name);

auto pos = base.rfind(".weight");
std::string scale_name = (pos != std::string::npos)
? base.substr(0, pos) + ".scale"
: base + ".scale";
std::string input_scale_name = (pos != std::string::npos)
? base.substr(0, pos) + ".input_scale"
: base + ".input_scale";

if (gguf_find_tensor(ctx, scale_name.c_str()) == -1) {
add_f32_scalar_tensor(ctx, scale_name.c_str());
}
if (gguf_find_tensor(ctx, input_scale_name.c_str()) == -1) {
add_f32_scalar_tensor(ctx, input_scale_name.c_str());
}
}

//
// main quantization driver
//
Expand Down Expand Up @@ -1038,6 +1076,12 @@ static void llama_model_quantize_impl(const std::string & fname_inp, const std::

metadata[i].requires_imatrix = tensor_requires_imatrix(tensor->name, metadata[i].target_type, ftype);

if (metadata[i].target_type == GGML_TYPE_NVFP4 &&
strstr(tensor->name, ".experts.") == nullptr) {
add_nvfp4_scale_tensors(ctx_outs[i_split].get(), tensor->name);
metadata[i].write_nvfp4_scales = true;
}

if (params->imatrix) {
metadata[i].remapped_imatrix_name = remap_imatrix(tensor->name, mapped);
} else if (metadata[i].allows_quantization && metadata[i].requires_imatrix) {
Expand Down Expand Up @@ -1149,6 +1193,7 @@ static void llama_model_quantize_impl(const std::string & fname_inp, const std::

void * new_data;
size_t new_size;
float nvfp4_scale_val = 1.0f;

if (params->dry_run) {
// the --dry-run option calculates the final quantization size without quantizing
Expand Down Expand Up @@ -1247,6 +1292,32 @@ static void llama_model_quantize_impl(const std::string & fname_inp, const std::
new_size += llama_tensor_quantize_impl(new_type, f32_data_03, new_data_03, chunk_size, nrows, n_per_row, imatrix_03, workers, nthread_use);
}
LLAMA_LOG_INFO("size = %8.2f MiB -> %8.2f MiB\n", tensor_size/1024.0/1024.0, new_size/1024.0/1024.0);

if (new_type == GGML_TYPE_NVFP4) {
const int64_t nvfp4_row_size = ggml_row_size(GGML_TYPE_NVFP4, n_per_row);
const ggml_type_traits * nvfp4_traits = ggml_get_type_traits(GGML_TYPE_NVFP4);

std::vector<float> dequant_buf(nelements);

for (int64_t i03 = 0; i03 < tensor->ne[2]; ++i03) {
for (int64_t i02 = 0; i02 < nrows; ++i02) {
const void * row = (const char *)new_data
+ i03 * nrows * nvfp4_row_size
+ i02 * nvfp4_row_size;
nvfp4_traits->to_float(row,
dequant_buf.data() + i03 * nelements_matrix + i02 * n_per_row,
n_per_row);
}
}

double sum_prod = 0.0;
double sum_sq = 0.0;
for (int64_t i = 0; i < nelements; ++i) {
sum_prod += (double)f32_data[i] * (double)dequant_buf[i];
sum_sq += (double)dequant_buf[i] * (double)dequant_buf[i];
}
nvfp4_scale_val = (sum_sq > 1e-10) ? (float)(sum_prod / sum_sq) : 1.0f;
}
}
total_size_org += tensor_size;
total_size_new += new_size;
Expand All @@ -1260,6 +1331,28 @@ static void llama_model_quantize_impl(const std::string & fname_inp, const std::
fout.write((const char *) new_data, new_size);
zeros(fout, GGML_PAD(new_size, align) - new_size);
} // no --dry-run

if (!params->dry_run && tm.write_nvfp4_scales) {
const float scale_val = nvfp4_scale_val;
const size_t scale_size = sizeof(float);

std::string base(tm.name);
auto pos = base.rfind(".weight");
std::string scale_name = (pos != std::string::npos)
? base.substr(0, pos) + ".scale"
: base + ".scale";
std::string input_scale_name = (pos != std::string::npos)
? base.substr(0, pos) + ".input_scale"
: base + ".input_scale";

gguf_set_tensor_data(ctx_outs[cur_split].get(), scale_name.c_str(), &scale_val);
fout.write((const char *) &scale_val, scale_size);
zeros(fout, GGML_PAD(scale_size, align) - scale_size);

gguf_set_tensor_data(ctx_outs[cur_split].get(), input_scale_name.c_str(), &scale_val);
fout.write((const char *) &scale_val, scale_size);
zeros(fout, GGML_PAD(scale_size, align) - scale_size);
}
} // main loop

if (!params->dry_run) {
Expand Down
1 change: 1 addition & 0 deletions tools/quantize/quantize.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -36,6 +36,7 @@ static const std::vector<quant_option> QUANT_OPTIONS = {
{ "Q4_0", LLAMA_FTYPE_MOSTLY_Q4_0, " 4.34G, +0.4685 ppl @ Llama-3-8B", },
{ "Q4_1", LLAMA_FTYPE_MOSTLY_Q4_1, " 4.78G, +0.4511 ppl @ Llama-3-8B", },
{ "MXFP4_MOE",LLAMA_FTYPE_MOSTLY_MXFP4_MOE," MXFP4 MoE", },
{ "NVFP4", LLAMA_FTYPE_MOSTLY_NVFP4, " NVFP4 4.0 bpw quantization", },
{ "Q5_0", LLAMA_FTYPE_MOSTLY_Q5_0, " 5.21G, +0.1316 ppl @ Llama-3-8B", },
{ "Q5_1", LLAMA_FTYPE_MOSTLY_Q5_1, " 5.65G, +0.1062 ppl @ Llama-3-8B", },
{ "IQ2_XXS", LLAMA_FTYPE_MOSTLY_IQ2_XXS, " 2.06 bpw quantization", },
Expand Down