Commit e9f824d8c for llama.cpp
commit e9f824d8c0f011662a742c9d15d4aa18a41e32c0
Author: ynankani <ynankani@nvidia.com>
Date: Fri Sep 25 08:36:35 2026 +0000
llama : add `llama_prec_policy` + model-driven W4A4 path (#24364)
* Rebase and update based on #26675
Signed-off-by: ynankani <ynankani@nvidia.com>
* CI failure fix(launh_bounds overload on HIP) and cleanup
Signed-off-by: ynankani <ynankani@nvidia.com>
* Address review comments
Signed-off-by: ynankani <ynankani@nvidia.com>
* Use ggml tensor instead of name in act policy map
Signed-off-by: ynankani <ynankani@nvidia.com>
* Address review comments and cleanup
Signed-off-by: ynankani <ynankani@nvidia.com>
* Address review comments
Signed-off-by: ynankani <ynankani@nvidia.com>
* Rename changes
Signed-off-by: ynankani <ynankani@nvidia.com>
* Update ggml/src/ggml-cuda/mmq.cu
Co-authored-by: Georgi Gerganov <ggerganov@gmail.com>
* MXFP4 dispatch changes for higher src prec
Signed-off-by: ynankani <ynankani@nvidia.com>
* Refactor and address review comments
Signed-off-by: ynankani <ynankani@nvidia.com>
* Updates based on review comments
Signed-off-by: ynankani <ynankani@nvidia.com>
* Apply batched suggestions from code review
Co-authored-by: Johannes Gäßler <johannesg@5d6.de>
* Address review comments
Signed-off-by: ynankani <ynankani@nvidia.com>
* Apply patch from review
Signed-off-by: ynankani <ynankani@nvidia.com>
---------
Signed-off-by: ynankani <ynankani@nvidia.com>
Co-authored-by: Georgi Gerganov <ggerganov@gmail.com>
Co-authored-by: Johannes Gäßler <johannesg@5d6.de>
diff --git a/conversion/base.py b/conversion/base.py
index 9fba5a86b..5561481e7 100644
--- a/conversion/base.py
+++ b/conversion/base.py
@@ -170,6 +170,9 @@ class ModelBase:
self.dir_model_card = dir_model # overridden in convert_lora_to_gguf.py
self._is_nvfp4 = False
self._is_mxfp4 = False
+ self._nvfp4_global_algo: str | None = None # checkpoint-wide NVFP4 quant_algo
+ self._nvfp4_layer_algo: dict[str, str | None] = {} # per-layer quant_algo, keyed by HF module path
+ self._prec_a4: dict[str, bool] = {} # gguf tensor name -> can use 4-bit (A4) activations
self._fp8_as_q8 = fp8_as_q8
self._fp8_dequantized: set[str] = set()
@@ -664,6 +667,18 @@ class ModelBase:
if bias_types:
self._fusable_qkv_bias_layers.add(bid)
+ def _tag_prec_a4(self, hf_name: str, gguf_name: str) -> None:
+ # W4A16_NVFP4 should not use 4-bit activations
+ name = hf_name.removesuffix(".weight").removesuffix(".bias")
+ algo = self._nvfp4_global_algo
+ while name:
+ if name in self._nvfp4_layer_algo:
+ algo = self._nvfp4_layer_algo[name]
+ break
+ name = name.rpartition(".")[0]
+ if algo == "W4A16_NVFP4":
+ self._prec_a4[gguf_name] = False
+
def set_gguf_parameters(self):
raise NotImplementedError("set_gguf_parameters() must be implemented in subclasses")
@@ -837,6 +852,7 @@ class ModelBase:
raw, shape = self._nvfp4_pack(weight, scale)
logger.info(f"Repacked {new_name} with shape {shape} and quantization NVFP4")
self.gguf_writer.add_tensor(new_name, raw, raw_dtype=gguf.GGMLQuantizationType.NVFP4)
+ self._tag_prec_a4(name, new_name)
self._write_scale_tensor(new_name.replace(".weight", ".scale"), scale2)
self._write_scale_tensor(new_name.replace(".weight", ".input_scale"), input_scale)
@@ -929,6 +945,7 @@ class ModelBase:
new_name = self.map_tensor_name(merged_name)
logger.info(f"Repacked {new_name} with shape [{len(experts)}, {shape[0]}, {shape[1]}] and quantization NVFP4")
self.gguf_writer.add_tensor(new_name, merged, raw_dtype=gguf.GGMLQuantizationType.NVFP4)
+ self._tag_prec_a4(merged_name, new_name)
scales.sort(key=lambda x: x[0])
self._write_scales_tensor(new_name.replace(".weight", ".scale"), [s[1] for s in scales])
@@ -971,6 +988,9 @@ class ModelBase:
and bool(quant_groups)
and all(g.get("format") == "nvfp4-pack-quantized" for g in quant_groups.values() if isinstance(g, dict))
)
+
+ self._nvfp4_global_algo = quant_algo
+
if quant_algo != "NVFP4":
if nvfp4_compressed_tensors:
quant_algo = "NVFP4"
@@ -980,6 +1000,22 @@ class ModelBase:
self._is_nvfp4 = quant_algo in ("NVFP4", "W4A16_NVFP4")
self._is_mxfp4 = quant_method == "mxfp4"
+ # Per-tensor NVFP4 precision.
+ self._nvfp4_layer_algo = {}
+ if quant_layers:
+ # store all possible module paths and assert if a quantized layer is not in the model
+ modules: set[str] = set()
+ for name in self.model_tensors:
+ while name := name.rpartition(".")[0]:
+ modules.add(name)
+
+ for layer_name, entry in quant_layers.items():
+ if not isinstance(entry, dict):
+ continue
+ if titem := self.filter_tensors((layer_name, lambda: torch.empty(0))):
+ assert titem[0] in modules, f"quantized_layers entry {layer_name!r} is not in the model tensors"
+ self._nvfp4_layer_algo[titem[0]] = entry.get("quant_algo")
+
# NVFP4 weights are repacked and written directly to gguf_writer.
# This must run before dequant_model so NVFP4 tensors are removed
# from model_tensors, leaving only non-NVFP4 (e.g. FP8) for dequant.
@@ -1185,6 +1221,12 @@ class ModelBase:
logger.info("Set model quantization version")
self.gguf_writer.add_quantization_version(gguf.GGML_QUANT_VERSION)
+ if self._prec_a4:
+ names = sorted(self._prec_a4.keys())
+ values = [self._prec_a4[n] for n in names]
+ logger.info(f"Set prec_a4 metadata for {len(names)} tensor(s)")
+ self.gguf_writer.add_tensor_extra_prec_a4(names, values)
+
def write_vocab(self):
raise NotImplementedError("write_vocab() must be implemented in subclasses")
diff --git a/docs/build.md b/docs/build.md
index bd666c145..08424940c 100644
--- a/docs/build.md
+++ b/docs/build.md
@@ -282,6 +282,13 @@ Consider setting `CUDA_SCALE_LAUNCH_QUEUES=4x`, which increases the CUDA command
Override default, speed-optimized compute types for cuBLAS matrix multiplications.
Legal values: `auto`, `f16`, `fp16`, `bf16`, `f32`, `fp32`.
+#### GGML_CUDA_MMQ_PREC
+
+Override the activation precision that the model requests for NVFP4 and MXFP4 matrix multiplications.
+Currently supported values: `auto`, `q8`, `q4`.
+
+NVFP4 and MXFP4 layers marked as W4A16 request 8-bit activations, so on Blackwell those layers run through the W4A8 path instead of the native W4A4 path. Set `q4` to keep the native W4A4 path for faster prompt processing at the cost of accuracy, or `q8` to use the W4A8 path for every layer, `auto` uses per-tensor prec metadata (this is the same behavior as when the environment variable is not set).
+
### Unified Memory
The environment variable `GGML_CUDA_ENABLE_UNIFIED_MEMORY=1` can be used to enable unified memory in Linux. This allows swapping to system RAM instead of crashing when the GPU VRAM is exhausted. In Windows this setting is available in the NVIDIA control panel as `System Memory Fallback`.
diff --git a/ggml/src/ggml-cuda/mmq-load-tiles.cuh b/ggml/src/ggml-cuda/mmq-load-tiles.cuh
index 7f00bad94..e19f4f24d 100644
--- a/ggml/src/ggml-cuda/mmq-load-tiles.cuh
+++ b/ggml/src/ggml-cuda/mmq-load-tiles.cuh
@@ -1631,16 +1631,16 @@ template <ggml_type type, int J, bool fallback> static __device__ __forceinline_
template <ggml_type type, int J, bool fallback> static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_mxfp4_fp4(
const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
constexpr int warp_size = ggml_cuda_get_physical_warp_size();
- constexpr int nwarps = ggml_cuda_mmq_get_nthreads(type, J, fallback) / warp_size;
- constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback);
- constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback);
+ constexpr int nwarps = ggml_cuda_mmq_get_nthreads(type, J, fallback, GGML_PREC_Q4) / warp_size;
+ constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback, GGML_PREC_Q4);
+ constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback, GGML_PREC_Q4);
int * x_qs = (int *) x_tile;
uint32_t * x_sc = (uint32_t *) (x_qs + 2 * MMQ_TILE_NE_K);
const int txi = threadIdx.x;
- constexpr int iter_k = ggml_cuda_mmq_get_K_vram(type, J, fallback);
+ constexpr int iter_k = ggml_cuda_mmq_get_K_vram(type, J, fallback, GGML_PREC_Q4);
constexpr int threads_per_row = iter_k / QK_MXFP4; // each thread processes 1 block
constexpr int rows_per_warp = warp_size / threads_per_row;
@@ -1670,12 +1670,12 @@ template <ggml_type type, int J, bool fallback> static __device__ __forceinline_
}
}
-template <ggml_type type, int J, bool fallback> static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_nvfp4(
+template <ggml_type type, int J, bool fallback, ggml_prec prec_src1 = GGML_PREC_Q8> static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_nvfp4(
const char * __restrict__ x, int * __restrict__ x_tile, const int kb0, const int i_max, const int stride) {
constexpr int warp_size = ggml_cuda_get_physical_warp_size();
- constexpr int nwarps = ggml_cuda_mmq_get_nthreads(type, J, fallback) / warp_size;
- constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback);
- constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback);
+ constexpr int nwarps = ggml_cuda_mmq_get_nthreads(type, J, fallback, prec_src1) / warp_size;
+ constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback, prec_src1);
+ constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback, prec_src1);
#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
int * x_qs = (int *) x_tile;
@@ -1729,12 +1729,12 @@ template <ggml_type type, int J, bool fallback> static __device__ __forceinline_
template <ggml_type type, int J, bool fallback> static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_nvfp4_nvfp4(
const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
constexpr int warp_size = ggml_cuda_get_physical_warp_size();
- constexpr int nwarps = ggml_cuda_mmq_get_nthreads(type, J, fallback) / warp_size;
- constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback);
- constexpr int iter_k = ggml_cuda_mmq_get_K_vram(type, J, fallback);
+ constexpr int nwarps = ggml_cuda_mmq_get_nthreads(type, J, fallback, GGML_PREC_Q4) / warp_size;
+ constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback, GGML_PREC_Q4);
+ constexpr int iter_k = ggml_cuda_mmq_get_K_vram(type, J, fallback, GGML_PREC_Q4);
constexpr int threads_per_row = iter_k / QK_NVFP4; // each thread processes 1 block
constexpr int rows_per_warp = warp_size / threads_per_row;
- constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback);
+ constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback, GGML_PREC_Q4);
uint32_t * x_u32 = (uint32_t *) x_tile;
diff --git a/ggml/src/ggml-cuda/mmq-vec-dot.cuh b/ggml/src/ggml-cuda/mmq-vec-dot.cuh
index 4d1c398fc..4ca6542d8 100644
--- a/ggml/src/ggml-cuda/mmq-vec-dot.cuh
+++ b/ggml/src/ggml-cuda/mmq-vec-dot.cuh
@@ -474,7 +474,7 @@ template <ggml_type type, int J, bool fallback> static __device__ __forceinline_
}
// Used for Q3_K, IQ2_S, and IQ2_XS:
-template <ggml_type type, int J, bool fallback> static __device__ __forceinline__ void ggml_cuda_mmq_vec_dot_q8_0_16_q8_1_mma(
+template <ggml_type type, int J, bool fallback, ggml_prec prec_src1 = GGML_PREC_Q8> static __device__ __forceinline__ void ggml_cuda_mmq_vec_dot_q8_0_16_q8_1_mma(
const int * __restrict__ x, const int * __restrict__ y, float * __restrict__ sum, const int k00) {
#if defined(AMD_MFMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
constexpr data_layout input_layout = get_input_data_layout();
@@ -482,7 +482,7 @@ template <ggml_type type, int J, bool fallback> static __device__ __forceinline_
typedef tile<16, 4, int, input_layout> tile_B;
typedef tile<16, 16, int, DATA_LAYOUT_J_MAJOR> tile_C;
- constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback);
+ constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback, prec_src1);
constexpr int rows_per_warp = ggml_cuda_mmq_get_rows_per_warp(type, J, fallback);
constexpr int ntx = rows_per_warp/tile_C::I; // Number of x minitiles per warp.
@@ -532,7 +532,7 @@ template <ggml_type type, int J, bool fallback> static __device__ __forceinline_
typedef tile< 8, 4, int> tile_B;
typedef tile<16, 8, int> tile_C;
- constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback);
+ constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback, prec_src1);
constexpr int rows_per_warp = ggml_cuda_mmq_get_rows_per_warp(type, J, fallback);
constexpr int ntx = rows_per_warp/tile_C::I; // Number of x minitiles per warp.
@@ -1180,7 +1180,7 @@ template <ggml_type type, int J, bool fallback> static __device__ __forceinline_
typedef tile<8, 8, int> tile_B;
typedef tile<16, 8, float> tile_C;
- constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback);
+ constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback, GGML_PREC_Q4);
constexpr int rows_per_warp = ggml_cuda_mmq_get_rows_per_warp(type, J, fallback);
constexpr int ntx = rows_per_warp / tile_C::I;
constexpr int nfrags = MMQ_TILE_NE_K / tile_A::J;
diff --git a/ggml/src/ggml-cuda/mmq.cu b/ggml/src/ggml-cuda/mmq.cu
index f68b3df18..3e1721bbe 100644
--- a/ggml/src/ggml-cuda/mmq.cu
+++ b/ggml/src/ggml-cuda/mmq.cu
@@ -5,7 +5,7 @@
#include <cstdint>
-static void ggml_cuda_mul_mat_q_switch_type(ggml_backend_cuda_context & ctx, const mmq_args & args, cudaStream_t stream) {
+static void ggml_cuda_mul_mat_q_switch_type(ggml_backend_cuda_context & ctx, const mmq_args & args, cudaStream_t stream, const ggml_prec prec_src1) {
switch (args.type_x) {
case GGML_TYPE_Q1_0:
mul_mat_q_case<GGML_TYPE_Q1_0>(ctx, args, stream);
@@ -71,9 +71,18 @@ static void ggml_cuda_mul_mat_q_switch_type(ggml_backend_cuda_context & ctx, con
break;
// -----------------------------------------------------------------------
case GGML_TYPE_MXFP4:
+ // src1 at Q4 uses the native FP4 instructions, which are Blackwell-only
+ if (prec_src1 == GGML_PREC_Q4) {
+ mul_mat_q_case<GGML_TYPE_MXFP4, GGML_PREC_Q4>(ctx, args, stream);
+ break;
+ }
mul_mat_q_case<GGML_TYPE_MXFP4>(ctx, args, stream);
break;
case GGML_TYPE_NVFP4:
+ if (prec_src1 == GGML_PREC_Q4) {
+ mul_mat_q_case<GGML_TYPE_NVFP4, GGML_PREC_Q4>(ctx, args, stream);
+ break;
+ }
mul_mat_q_case<GGML_TYPE_NVFP4>(ctx, args, stream);
break;
default:
@@ -82,6 +91,47 @@ static void ggml_cuda_mul_mat_q_switch_type(ggml_backend_cuda_context & ctx, con
}
}
+// overrides the src1 precision requested by the graph, "auto" keeps the requested one
+static ggml_prec ggml_cuda_mmq_get_prec_env() {
+ const char * env_c = getenv("GGML_CUDA_MMQ_PREC");
+ if (env_c == nullptr) {
+ return GGML_PREC_UNDEFINED;
+ }
+ std::string env_cpp = env_c;
+ for (char & c : env_cpp) {
+ c = std::tolower(c);
+ }
+ if (env_cpp == "q4") {
+ return GGML_PREC_Q4;
+ }
+ if (env_cpp == "q8") {
+ return GGML_PREC_Q8;
+ }
+ if (env_cpp != "auto") {
+ GGML_LOG_WARN("%s: Unknown value for GGML_CUDA_MMQ_PREC: '%s'. Available: 'q4', 'q8', 'auto'.\n", __func__, env_cpp.c_str());
+ }
+ return GGML_PREC_UNDEFINED;
+}
+
+// src1 is quantized to Q8_1 unless the FP4 types can use 4-bit activations, in which case they
+// default to the native W4A4 instructions on Blackwell.
+static ggml_prec ggml_cuda_mmq_get_prec_src1(const ggml_tensor * src0, const ggml_tensor * dst, const int cc) {
+ static const ggml_prec prec_env = ggml_cuda_mmq_get_prec_env();
+
+ ggml_prec prec = prec_env;
+ if (prec == GGML_PREC_UNDEFINED) {
+ prec = (ggml_prec) ggml_get_op_params_i32(dst, 3);
+ }
+
+ // Q4 only for the FP4 types on Blackwell
+ GGML_ASSERT(prec == GGML_PREC_UNDEFINED || prec == GGML_PREC_Q8 || prec == GGML_PREC_Q4);
+ const bool can_use_q4 = (src0->type == GGML_TYPE_NVFP4 || src0->type == GGML_TYPE_MXFP4) && blackwell_mma_available(cc);
+ if (prec == GGML_PREC_Q8 || !can_use_q4) {
+ return GGML_PREC_Q8;
+ }
+ return GGML_PREC_Q4;
+}
+
void ggml_cuda_mul_mat_q(
ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * ids, ggml_tensor * dst) {
GGML_ASSERT( src1->type == GGML_TYPE_F32);
@@ -128,7 +178,9 @@ void ggml_cuda_mul_mat_q(
const bool fallback = ne01 % 128 != 0;
- const bool use_native_fp4 = blackwell_mma_available(cc) && (src0->type == GGML_TYPE_MXFP4 || src0->type == GGML_TYPE_NVFP4);
+ const ggml_prec prec_src1 = ggml_cuda_mmq_get_prec_src1(src0, dst, cc);
+
+ const bool use_native_fp4 = prec_src1 == GGML_PREC_Q4;
const size_t y_block_size = use_native_fp4 ? sizeof(block_fp4_mmq) : sizeof(block_q8_1_mmq);
const size_t y_values_per_block = use_native_fp4 ? QK_FP4_MMQ : QK8_1_MMQ;
@@ -172,7 +224,7 @@ void ggml_cuda_mul_mat_q(
ne02, ne12, s02, s12, s2,
ne03, ne13, s03, s13, s3,
ne1, ne1};
- ggml_cuda_mul_mat_q_switch_type(ctx, args, stream);
+ ggml_cuda_mul_mat_q_switch_type(ctx, args, stream, prec_src1);
return;
}
@@ -260,7 +312,7 @@ void ggml_cuda_mul_mat_q(
ne03, ne13, s03, s13, s3,
ne12, ncols_opt};
- ggml_cuda_mul_mat_q_switch_type(ctx, args, stream);
+ ggml_cuda_mul_mat_q_switch_type(ctx, args, stream, prec_src1);
}
bool ggml_cuda_should_use_mmq(enum ggml_type type, int cc, int64_t ne11, int64_t n_experts) {
diff --git a/ggml/src/ggml-cuda/mmq.cuh b/ggml/src/ggml-cuda/mmq.cuh
index 6923f3510..4b50d1dc3 100644
--- a/ggml/src/ggml-cuda/mmq.cuh
+++ b/ggml/src/ggml-cuda/mmq.cuh
@@ -227,7 +227,7 @@ struct ggml_cuda_mmq_config {
#undef CASE
-static __host__ ggml_cuda_mmq_config ggml_cuda_mmq_get_config(const ggml_type type, const int J, const bool fallback, const int cc) {
+static __host__ ggml_cuda_mmq_config ggml_cuda_mmq_get_config(const ggml_type type, const int J, const bool fallback, const int cc, const ggml_prec prec_src1 = GGML_PREC_Q8) {
if (GGML_CUDA_CC_IS_AMD(cc)) {
if (GGML_CUDA_CC_IS_GCN(cc)) {
return ggml_cuda_mmq_get_config_gcn(type, J, fallback);
@@ -247,6 +247,10 @@ static __host__ ggml_cuda_mmq_config ggml_cuda_mmq_get_config(const ggml_type ty
return ggml_cuda_mmq_get_config_rdna2(type, J, fallback);
}
if (blackwell_mma_available(cc)) {
+ // only src1 at Q4 uses the native FP4 config, higher precisions keep src1 at Q8_1
+ if (prec_src1 != GGML_PREC_Q4 && (type == GGML_TYPE_NVFP4 || type == GGML_TYPE_MXFP4)) {
+ return ggml_cuda_mmq_get_config_ampere(type, J, fallback);
+ }
return ggml_cuda_mmq_get_config_blackwell(type, J, fallback);
}
if (ggml_cuda_highest_compiled_arch(cc) >= GGML_CUDA_CC_VOLTA) {
@@ -258,7 +262,7 @@ static __host__ ggml_cuda_mmq_config ggml_cuda_mmq_get_config(const ggml_type ty
return ggml_cuda_mmq_get_config_pascal_older(type, J, fallback);
}
-static constexpr __device__ ggml_cuda_mmq_config ggml_cuda_mmq_get_config(ggml_type type, int J, bool fallback) {
+static constexpr __device__ ggml_cuda_mmq_config ggml_cuda_mmq_get_config(ggml_type type, int J, bool fallback, ggml_prec prec_src1 = GGML_PREC_Q8) {
#ifdef GGML_USE_HIP
#ifdef GCN
return ggml_cuda_mmq_get_config_gcn(type, J, fallback);
@@ -275,6 +279,10 @@ static constexpr __device__ ggml_cuda_mmq_config ggml_cuda_mmq_get_config(ggml_t
#endif // CDNA
#else
#ifdef BLACKWELL_MMA_AVAILABLE
+ // only src1 at Q4 uses the native FP4 config, higher precisions keep src1 at Q8_1
+ if (prec_src1 != GGML_PREC_Q4 && (type == GGML_TYPE_NVFP4 || type == GGML_TYPE_MXFP4)) {
+ return ggml_cuda_mmq_get_config_ampere(type, J, fallback);
+ }
return ggml_cuda_mmq_get_config_blackwell(type, J, fallback);
#elif __CUDA_ARCH__ >= GGML_CUDA_CC_VOLTA
return ggml_cuda_mmq_get_config_ampere(type, J, fallback);
@@ -284,79 +292,71 @@ static constexpr __device__ ggml_cuda_mmq_config ggml_cuda_mmq_get_config(ggml_t
return ggml_cuda_mmq_get_config_pascal_older(type, J, fallback);
#endif // BLACKWELL_MMA_AVAILABLE
#endif // GGML_USE_HIP
- GGML_UNUSED_VARS(type, J, fallback);
+ GGML_UNUSED_VARS(type, J, fallback, prec_src1);
}
static __host__ int ggml_cuda_mmq_get_type(const ggml_type type, const int J, const bool fallback, const int cc) {
return ggml_cuda_mmq_get_config(type, J, fallback, cc).type;
}
-static constexpr __device__ int ggml_cuda_mmq_get_type(ggml_type type, int J, bool fallback) {
- return ggml_cuda_mmq_get_config(type, J, fallback).type;
-}
-
-static __host__ int ggml_cuda_mmq_get_nthreads(const ggml_type type, const int J, const bool fallback, const int cc) {
- return ggml_cuda_mmq_get_config(type, J, fallback, cc).nthreads;
-}
-
-static constexpr __device__ int ggml_cuda_mmq_get_nthreads(ggml_type type, int J, bool fallback) {
- return ggml_cuda_mmq_get_config(type, J, fallback).nthreads;
+static constexpr __device__ int ggml_cuda_mmq_get_type(ggml_type type, int J, bool fallback, ggml_prec prec_src1 = GGML_PREC_Q8) {
+ return ggml_cuda_mmq_get_config(type, J, fallback, prec_src1).type;
}
-static __host__ int ggml_cuda_mmq_get_occupancy(const ggml_type type, const int J, const bool fallback, const int cc) {
- return ggml_cuda_mmq_get_config(type, J, fallback, cc).occupancy;
+static constexpr __device__ int ggml_cuda_mmq_get_nthreads(ggml_type type, int J, bool fallback, ggml_prec prec_src1 = GGML_PREC_Q8) {
+ return ggml_cuda_mmq_get_config(type, J, fallback, prec_src1).nthreads;
}
-static constexpr __device__ int ggml_cuda_mmq_get_occupancy(ggml_type type, int J, bool fallback) {
- return ggml_cuda_mmq_get_config(type, J, fallback).occupancy;
+static constexpr __device__ int ggml_cuda_mmq_get_occupancy(ggml_type type, int J, bool fallback, ggml_prec prec_src1 = GGML_PREC_Q8) {
+ return ggml_cuda_mmq_get_config(type, J, fallback, prec_src1).occupancy;
}
static __host__ int ggml_cuda_mmq_get_I(const ggml_type type, const int J, const bool fallback, const int cc) {
return ggml_cuda_mmq_get_config(type, J, fallback, cc).I;
}
-static constexpr __device__ int ggml_cuda_mmq_get_I(ggml_type type, int J, bool fallback) {
- return ggml_cuda_mmq_get_config(type, J, fallback).I;
+static constexpr __device__ int ggml_cuda_mmq_get_I(ggml_type type, int J, bool fallback, ggml_prec prec_src1 = GGML_PREC_Q8) {
+ return ggml_cuda_mmq_get_config(type, J, fallback, prec_src1).I;
}
static __host__ int ggml_cuda_mmq_get_J(const ggml_type type, const int J, const bool fallback, const int cc) {
return ggml_cuda_mmq_get_config(type, J, fallback, cc).J;
}
-static constexpr __device__ int ggml_cuda_mmq_get_J(ggml_type type, int J, bool fallback) {
- return ggml_cuda_mmq_get_config(type, J, fallback).J;
+static constexpr __device__ int ggml_cuda_mmq_get_J(ggml_type type, int J, bool fallback, ggml_prec prec_src1 = GGML_PREC_Q8) {
+ return ggml_cuda_mmq_get_config(type, J, fallback, prec_src1).J;
}
static __host__ ggml_cuda_mmq_sram_layout ggml_cuda_mmq_get_sram_layout(const ggml_type type, const int J, const bool fallback, const int cc) {
return ggml_cuda_mmq_get_config(type, J, fallback, cc).sram_layout;
}
-static constexpr __device__ ggml_cuda_mmq_sram_layout ggml_cuda_mmq_get_sram_layout(ggml_type type, int J, bool fallback) {
- return ggml_cuda_mmq_get_config(type, J, fallback).sram_layout;
+static constexpr __device__ ggml_cuda_mmq_sram_layout ggml_cuda_mmq_get_sram_layout(ggml_type type, int J, bool fallback, ggml_prec prec_src1 = GGML_PREC_Q8) {
+ return ggml_cuda_mmq_get_config(type, J, fallback, prec_src1).sram_layout;
}
static __host__ int ggml_cuda_mmq_get_K_vram(const ggml_type type, const int J, const bool fallback, const int cc) {
return ggml_cuda_mmq_get_config(type, J, fallback, cc).K_vram;
}
-static constexpr __device__ int ggml_cuda_mmq_get_K_vram(ggml_type type, int J, bool fallback) {
- return ggml_cuda_mmq_get_config(type, J, fallback).K_vram;
+static constexpr __device__ int ggml_cuda_mmq_get_K_vram(ggml_type type, int J, bool fallback, ggml_prec prec_src1 = GGML_PREC_Q8) {
+ return ggml_cuda_mmq_get_config(type, J, fallback, prec_src1).K_vram;
}
static __host__ bool ggml_cuda_mmq_get_stream_k(const ggml_type type, const int J, const bool fallback, const int cc) {
return ggml_cuda_mmq_get_config(type, J, fallback, cc).stream_k;
}
-static constexpr __device__ bool ggml_cuda_mmq_get_stream_k(ggml_type type, int J, bool fallback) {
- return ggml_cuda_mmq_get_config(type, J, fallback).stream_k;
+static constexpr __device__ bool ggml_cuda_mmq_get_stream_k(ggml_type type, int J, bool fallback, ggml_prec prec_src1 = GGML_PREC_Q8) {
+ return ggml_cuda_mmq_get_config(type, J, fallback, prec_src1).stream_k;
}
static __host__ int ggml_cuda_mmq_get_fallback(const ggml_type type, const int J, const bool fallback, const int cc) {
return ggml_cuda_mmq_get_config(type, J, fallback, cc).fallback;
}
-static constexpr __device__ int ggml_cuda_mmq_get_fallback(ggml_type type, int J, bool fallback) {
- return ggml_cuda_mmq_get_config(type, J, fallback).fallback;
+static constexpr __device__ int ggml_cuda_mmq_get_fallback(ggml_type type, int J, bool fallback, ggml_prec prec_src1 = GGML_PREC_Q8) {
+ return ggml_cuda_mmq_get_config(type, J, fallback, prec_src1).fallback;
}
// ---------------------------------------------------------------------------------------------
@@ -365,8 +365,8 @@ static __host__ int ggml_cuda_mmq_get_sram_stride(const ggml_type type, const in
return ggml_cuda_mmq_get_sram_stride(ggml_cuda_mmq_get_sram_layout(type, J, fallback, cc));
}
-static constexpr __device__ int ggml_cuda_mmq_get_sram_stride(ggml_type type, int J, bool fallback) {
- return ggml_cuda_mmq_get_sram_stride(ggml_cuda_mmq_get_sram_layout(type, J, fallback));
+static constexpr __device__ int ggml_cuda_mmq_get_sram_stride(ggml_type type, int J, bool fallback, ggml_prec prec_src1 = GGML_PREC_Q8) {
+ return ggml_cuda_mmq_get_sram_stride(ggml_cuda_mmq_get_sram_layout(type, J, fallback, prec_src1));
}
static __host__ int ggml_cuda_mmq_get_J_max(const ggml_type type, const bool fallback, const int cc, const int64_t ne11) {
@@ -541,9 +541,9 @@ struct ggml_cuda_mmq_util_funcs {
vdr(vdr), load_tiles(load_tiles), vec_dot(vec_dot), write_back(write_back) {}
};
-template <ggml_type type, int J, bool fallback>
+template <ggml_type type, int J, bool fallback, ggml_prec prec_src1 = GGML_PREC_Q8>
static constexpr __device__ ggml_cuda_mmq_util_funcs ggml_cuda_mmq_get_util_funcs() {
- if (!ggml_cuda_mmq_get_config(type, J, fallback).use_mma_data_layout()) {
+ if (!ggml_cuda_mmq_get_config(type, J, fallback, prec_src1).use_mma_data_layout()) {
switch (type) {
case GGML_TYPE_Q1_0:
return ggml_cuda_mmq_util_funcs(
@@ -690,17 +690,23 @@ static constexpr __device__ ggml_cuda_mmq_util_funcs ggml_cuda_mmq_get_util_func
#ifdef BLACKWELL_MMA_AVAILABLE
switch (type) {
case GGML_TYPE_MXFP4:
- return ggml_cuda_mmq_util_funcs(
- -1,
- ggml_cuda_mmq_load_tiles_mxfp4_fp4<type, J, fallback>,
- ggml_cuda_mmq_vec_dot_fp4_fp4_mma<type, J, fallback>,
- ggml_cuda_mmq_write_back_mma<type, J, fallback>);
+ if (prec_src1 == GGML_PREC_Q4) {
+ return ggml_cuda_mmq_util_funcs(
+ -1,
+ ggml_cuda_mmq_load_tiles_mxfp4_fp4<type, J, fallback>,
+ ggml_cuda_mmq_vec_dot_fp4_fp4_mma<type, J, fallback>,
+ ggml_cuda_mmq_write_back_mma<type, J, fallback>);
+ }
+ break;
case GGML_TYPE_NVFP4:
- return ggml_cuda_mmq_util_funcs(
- -1,
- ggml_cuda_mmq_load_tiles_nvfp4_nvfp4<type, J, fallback>,
- ggml_cuda_mmq_vec_dot_fp4_fp4_mma<type, J, fallback>,
- ggml_cuda_mmq_write_back_mma<type, J, fallback>);
+ if (prec_src1 == GGML_PREC_Q4) {
+ return ggml_cuda_mmq_util_funcs(
+ -1,
+ ggml_cuda_mmq_load_tiles_nvfp4_nvfp4<type, J, fallback>,
+ ggml_cuda_mmq_vec_dot_fp4_fp4_mma<type, J, fallback>,
+ ggml_cuda_mmq_write_back_mma<type, J, fallback>);
+ }
+ break;
default:
break;
}
@@ -841,37 +847,37 @@ static constexpr __device__ ggml_cuda_mmq_util_funcs ggml_cuda_mmq_get_util_func
case GGML_TYPE_NVFP4:
return ggml_cuda_mmq_util_funcs(
-1,
- ggml_cuda_mmq_load_tiles_nvfp4<type, J, fallback>,
- ggml_cuda_mmq_vec_dot_q8_0_16_q8_1_mma<type, J, fallback>,
+ ggml_cuda_mmq_load_tiles_nvfp4<type, J, fallback, prec_src1>,
+ ggml_cuda_mmq_vec_dot_q8_0_16_q8_1_mma<type, J, fallback, prec_src1>,
ggml_cuda_mmq_write_back_mma<type, J, fallback>);
default:
return ggml_cuda_mmq_util_funcs(1, nullptr, nullptr, nullptr);
}
}
-template <ggml_type type, int J, bool fallback>
+template <ggml_type type, int J, bool fallback, ggml_prec prec_src1 = GGML_PREC_Q8>
static constexpr __device__ int ggml_cuda_mmq_get_vdr() {
- return ggml_cuda_mmq_get_util_funcs<type, J, fallback>().vdr;
+ return ggml_cuda_mmq_get_util_funcs<type, J, fallback, prec_src1>().vdr;
}
-template <ggml_type type, int J, bool fallback>
+template <ggml_type type, int J, bool fallback, ggml_prec prec_src1 = GGML_PREC_Q8>
static constexpr __device__ ggml_cuda_mmq_load_tiles_t ggml_cuda_mmq_get_load_tiles() {
- return ggml_cuda_mmq_get_util_funcs<type, J, fallback>().load_tiles;
+ return ggml_cuda_mmq_get_util_funcs<type, J, fallback, prec_src1>().load_tiles;
}
-template <ggml_type type, int J, bool fallback>
+template <ggml_type type, int J, bool fallback, ggml_prec prec_src1 = GGML_PREC_Q8>
static constexpr __device__ ggml_cuda_mmq_vec_dot_t ggml_cuda_mmq_get_vec_dot() {
- return ggml_cuda_mmq_get_util_funcs<type, J, fallback>().vec_dot;
+ return ggml_cuda_mmq_get_util_funcs<type, J, fallback, prec_src1>().vec_dot;
}
-template <ggml_type type, int J, bool fallback>
+template <ggml_type type, int J, bool fallback, ggml_prec prec_src1 = GGML_PREC_Q8>
static constexpr __device__ ggml_cuda_mmq_write_back_t ggml_cuda_mmq_get_write_back() {
- return ggml_cuda_mmq_get_util_funcs<type, J, fallback>().write_back;
+ return ggml_cuda_mmq_get_util_funcs<type, J, fallback, prec_src1>().write_back;
}
// ---------------------------------------------------------------------------------------------
-template <ggml_type type, int J, bool fallback, bool fixup>
+template <ggml_type type, int J, bool fallback, bool fixup, ggml_prec prec_src1 = GGML_PREC_Q8>
static __device__ __forceinline__ void mul_mat_q_process_tile(
const char * __restrict__ x, const int offset_x, const int * __restrict__ y,
const int * __restrict__ ids_dst, float * __restrict__ dst, float * __restrict__ tmp_fixup,
@@ -880,25 +886,27 @@ static __device__ __forceinline__ void mul_mat_q_process_tile(
const int tile_x_max_i, const int tile_y_max_j, const int kb0_start, const int kb0_stop) {
constexpr int warp_size = ggml_cuda_get_physical_warp_size();
- constexpr int nwarps = ggml_cuda_mmq_get_nthreads(type, J, fallback) / warp_size;
+ constexpr int nwarps = ggml_cuda_mmq_get_nthreads(type, J, fallback, prec_src1) / warp_size;
constexpr int qk = ggml_cuda_type_traits<type>::qk;
- constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback);
- constexpr ggml_cuda_mmq_load_tiles_t load_tiles = ggml_cuda_mmq_get_load_tiles<type, J, fallback>();
- constexpr ggml_cuda_mmq_vec_dot_t vec_dot = ggml_cuda_mmq_get_vec_dot<type, J, fallback>();
- constexpr ggml_cuda_mmq_write_back_t write_back = ggml_cuda_mmq_get_write_back<type, J, fallback>();
+ constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback, prec_src1);
+ constexpr ggml_cuda_mmq_load_tiles_t load_tiles = ggml_cuda_mmq_get_load_tiles<type, J, fallback, prec_src1>();
+ constexpr ggml_cuda_mmq_vec_dot_t vec_dot = ggml_cuda_mmq_get_vec_dot<type, J, fallback, prec_src1>();
+ constexpr ggml_cuda_mmq_write_back_t write_back = ggml_cuda_mmq_get_write_back<type, J, fallback, prec_src1>();
extern __shared__ int data_mul_mat_q[];
int * tile_y = data_mul_mat_q + J;
int * tile_x = tile_y + GGML_PAD(J*MMQ_TILE_Y_K, nwarps*warp_size);
#if defined(BLACKWELL_MMA_AVAILABLE)
- // FP4 tile stores 8 blocks
- constexpr int ne_block = (type == GGML_TYPE_MXFP4 || type == GGML_TYPE_NVFP4) ? QK_FP4_MMQ : QK8_1_MMQ;
+ // FP4 tile stores 8 blocks. src1 above Q4 uses the generic
+ // Q8_1 tile layout instead of the packed FP4 tile.
+ constexpr int ne_block = ((type == GGML_TYPE_MXFP4 || type == GGML_TYPE_NVFP4) && prec_src1 == GGML_PREC_Q4) ?
+ QK_FP4_MMQ : QK8_1_MMQ;
#else
constexpr int ne_block = QK8_1_MMQ;
#endif // defined(BLACKWELL_MMA_AVAILABLE)
- constexpr int ITER_K = ggml_cuda_mmq_get_K_vram(type, J, fallback);
+ constexpr int ITER_K = ggml_cuda_mmq_get_K_vram(type, J, fallback, prec_src1);
constexpr int blocks_per_iter = ITER_K / qk;
float sum[J*I / (nwarps*warp_size)] = {0.0f};
@@ -950,8 +958,8 @@ static __device__ __forceinline__ void mul_mat_q_process_tile(
// The mul_mat_q kernel implements "stream-k" work partitioning as described in https://arxiv.org/abs/2301.03598
-template <ggml_type type, int J, bool fallback>
-__launch_bounds__(ggml_cuda_mmq_get_nthreads(type, J, fallback), ggml_cuda_mmq_get_occupancy(type, J, fallback))
+template <ggml_type type, int J, bool fallback, ggml_prec prec_src1 = GGML_PREC_Q8>
+__launch_bounds__(ggml_cuda_mmq_get_nthreads(type, J, fallback, prec_src1), ggml_cuda_mmq_get_occupancy(type, J, fallback, prec_src1))
static __global__ void mul_mat_q(
const char * __restrict__ x, const int * __restrict__ y, const int32_t * __restrict__ ids_dst,
const int32_t * __restrict__ expert_bounds, float * __restrict__ dst, float * __restrict__ tmp_fixup,
@@ -962,15 +970,15 @@ static __global__ void mul_mat_q(
const uint3 ntx) {
// Skip unused template specializations for faster compilation:
- if (ggml_cuda_mmq_get_config(type, J, fallback).type == GGML_TYPE_COUNT) {
+ if (ggml_cuda_mmq_get_config(type, J, fallback, prec_src1).type == GGML_TYPE_COUNT) {
NO_DEVICE_CODE;
return;
}
constexpr int warp_size = ggml_cuda_get_physical_warp_size();
- constexpr int nwarps = ggml_cuda_mmq_get_nthreads(type, J, fallback) / warp_size;
+ constexpr int nwarps = ggml_cuda_mmq_get_nthreads(type, J, fallback, prec_src1) / warp_size;
constexpr int qk = ggml_cuda_type_traits<type>::qk;
- constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback);
+ constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback, prec_src1);
const uint32_t nty = (nrows_x + I - 1) / I; // Number of tiles y
@@ -990,7 +998,7 @@ static __global__ void mul_mat_q(
}
__syncthreads();
- if constexpr (!ggml_cuda_mmq_get_stream_k(type, J, fallback)) {
+ if constexpr (!ggml_cuda_mmq_get_stream_k(type, J, fallback, prec_src1)) {
const uint2 tmp2 = fast_div_modulo(blockIdx.z, nchannels_y);
const int wt = tmp2.x;
const int zt = tmp2.y;
@@ -1053,14 +1061,14 @@ static __global__ void mul_mat_q(
const int offset_x = fastdiv(wt, sample_ratio)*stride_sample_x + fastdiv(zt, channel_ratio)*stride_channel_x + it*I*stride_row_x;
constexpr bool fixup = false;
- mul_mat_q_process_tile<type, J, fallback, fixup>
+ mul_mat_q_process_tile<type, J, fallback, fixup, prec_src1>
(x, offset_x, y + offset_y, ids_dst_shared, dst + offset_dst, tmp_fixup, y_scale_tile,
stride_row_x, ncols_y, stride_col_dst,
tile_x_max_i, tile_y_max_j, 0, blocks_per_ne00.z);
return;
}
- constexpr int ITER_K = ggml_cuda_mmq_get_K_vram(type, J, fallback);
+ constexpr int ITER_K = ggml_cuda_mmq_get_K_vram(type, J, fallback, prec_src1);
constexpr int blocks_per_iter = ITER_K / qk;
// kbc == k block continuous, current index in continuous ijk space.
@@ -1147,7 +1155,7 @@ static __global__ void mul_mat_q(
const int offset_x = fastdiv(wt, sample_ratio)*stride_sample_x + fastdiv(zt, channel_ratio)*stride_channel_x + it*I*stride_row_x;
constexpr bool fixup = false; // All but (potentially) the last iterations write their data to dst rather than the fixup buffer.
- mul_mat_q_process_tile<type, J, fallback, fixup>
+ mul_mat_q_process_tile<type, J, fallback, fixup, prec_src1>
(x, offset_x, y + offset_y, ids_dst_shared, dst + offset_dst, tmp_fixup, y_scale_tile,
stride_row_x, ncols_y, stride_col_dst,
tile_x_max_i, tile_y_max_j, kb0_start, kb0_stop);
@@ -1231,24 +1239,24 @@ static __global__ void mul_mat_q(
const int offset_x = fastdiv(wt, sample_ratio)*stride_sample_x + fastdiv(zt, channel_ratio)*stride_channel_x + it*I*stride_row_x;
constexpr bool fixup = true; // Last index writes its data to fixup buffer to avoid data races with other blocks.
- mul_mat_q_process_tile<type, J, fallback, fixup>
+ mul_mat_q_process_tile<type, J, fallback, fixup, prec_src1>
(x, offset_x, y + offset_y, ids_dst_shared, dst + offset_dst, tmp_fixup, y_scale_tile,
stride_row_x, ncols_y, stride_col_dst,
tile_x_max_i, tile_y_max_j, kb0_start, kb0_stop);
}
-template <ggml_type type, int J, bool fallback>
-__launch_bounds__(ggml_cuda_mmq_get_nthreads(type, J, fallback)/2, 1)
+template <ggml_type type, int J, bool fallback, ggml_prec prec_src1 = GGML_PREC_Q8>
+__launch_bounds__(ggml_cuda_mmq_get_nthreads(type, J, fallback, prec_src1)/2, 1)
static __global__ void mul_mat_q_stream_k_fixup(
const int32_t * __restrict__ ids_dst, const int32_t * __restrict__ expert_bounds, float * __restrict__ dst,
float * __restrict__ tmp_last_tile, const uint3 blocks_per_ne00, const int nrows_x, const int ncols_dst,
const int stride_col_dst, const uint3 nchannels_y, const int stride_channel_dst, const uint3 nsamples_y,
const int stride_sample_dst, const uint3 ntx) {
constexpr int warp_size = ggml_cuda_get_physical_warp_size();
- constexpr int nwarps = (ggml_cuda_mmq_get_nthreads(type, J, fallback) / 2) / warp_size;
- constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback);
+ constexpr int nwarps = (ggml_cuda_mmq_get_nthreads(type, J, fallback, prec_src1) / 2) / warp_size;
+ constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback, prec_src1);
constexpr int qk = ggml_cuda_type_traits<type>::qk;
- constexpr int ITER_K = ggml_cuda_mmq_get_K_vram(type, J, fallback);
+ constexpr int ITER_K = ggml_cuda_mmq_get_K_vram(type, J, fallback, prec_src1);
constexpr int blocks_per_iter = ITER_K / qk;
float sum[J / nwarps] = {0.0f};
@@ -1392,22 +1400,22 @@ static size_t mmq_get_nbytes_shared(const ggml_cuda_mmq_config & config, const i
return nbs_ids + nbs_x + GGML_PAD(nbs_y, config.nthreads*sizeof(int));
}
-template <ggml_type type, int J, bool fallback>
+template <ggml_type type, int J, bool fallback, ggml_prec prec_src1 = GGML_PREC_Q8>
static void launch_mul_mat_q(ggml_backend_cuda_context & ctx, const mmq_args & args, cudaStream_t stream) {
const int id = ggml_cuda_get_device();
const int cc = ggml_cuda_info().devices[id].cc;
const int nsm = ggml_cuda_info().devices[id].nsm;
const int warp_size = ggml_cuda_info().devices[id].warp_size;
- const ggml_cuda_mmq_config config = ggml_cuda_mmq_get_config(type, J, fallback, cc);
+ const ggml_cuda_mmq_config config = ggml_cuda_mmq_get_config(type, J, fallback, cc, prec_src1);
GGML_ASSERT(config.nthreads % warp_size == 0);
const int nwarps = config.nthreads / warp_size;
const int nbytes_shared = mmq_get_nbytes_shared(config, cc);
const dim3 block_dims(warp_size, nwarps, 1);
- CUDA_SET_SHARED_MEMORY_LIMIT((mul_mat_q<type, J, false>), nbytes_shared);
- CUDA_SET_SHARED_MEMORY_LIMIT((mul_mat_q<type, J, true>), nbytes_shared);
+ CUDA_SET_SHARED_MEMORY_LIMIT((mul_mat_q<type, J, false, prec_src1>), nbytes_shared);
+ CUDA_SET_SHARED_MEMORY_LIMIT((mul_mat_q<type, J, true, prec_src1>), nbytes_shared);
const int nty = (args.nrows_x + config.I - 1) / config.I;
const int ntx = (args.ncols_max + config.J - 1) / config.J;
@@ -1426,8 +1434,8 @@ static void launch_mul_mat_q(ggml_backend_cuda_context & ctx, const mmq_args & a
const uint3 channel_ratio_fd = init_fastdiv_values(channel_ratio);
const uint3 sample_ratio_fd = init_fastdiv_values(sample_ratio);
- if (!ggml_cuda_mmq_get_stream_k(type, J, fallback, cc)) {
- mul_mat_q<type, J, fallback><<<block_nums_xy_tiling, block_dims, nbytes_shared, stream>>>
+ if (!config.stream_k) {
+ mul_mat_q<type, J, fallback, prec_src1><<<block_nums_xy_tiling, block_dims, nbytes_shared, stream>>>
(args.x, args.y, args.ids_dst, args.expert_bounds, args.dst, nullptr, args.y_scale,
blocks_per_ne00_fd, args.nrows_x, args.ncols_dst, args.stride_row_x, args.ncols_y, args.nrows_dst,
channel_ratio_fd, nchannels_y_fd, args.stride_channel_x, args.stride_channel_y, args.stride_channel_dst,
@@ -1456,7 +1464,7 @@ static void launch_mul_mat_q(ggml_backend_cuda_context & ctx, const mmq_args & a
const dim3 block_nums_fixup(block_nums_stream_k.x, config.I/warp_size, 1);
const dim3 block_dims_fixup(block_dims.x, block_dims.y/2, block_dims.z);
- mul_mat_q<type, J, fallback><<<block_nums_stream_k, block_dims, nbytes_shared, stream>>>
+ mul_mat_q<type, J, fallback, prec_src1><<<block_nums_stream_k, block_dims, nbytes_shared, stream>>>
(args.x, args.y, args.ids_dst, args.expert_bounds, args.dst, tmp_fixup.ptr, args.y_scale,
blocks_per_ne00_fd, args.nrows_x, args.ncols_dst, args.stride_row_x, args.ncols_y, args.nrows_dst,
channel_ratio_fd, nchannels_y_fd, args.stride_channel_x, args.stride_channel_y, args.stride_channel_dst,
@@ -1468,13 +1476,13 @@ static void launch_mul_mat_q(ggml_backend_cuda_context & ctx, const mmq_args & a
}
CUDA_CHECK(cudaGetLastError());
- mul_mat_q_stream_k_fixup<type, J, fallback><<<block_nums_fixup, block_dims_fixup, 0, stream>>>
+ mul_mat_q_stream_k_fixup<type, J, fallback, prec_src1><<<block_nums_fixup, block_dims_fixup, 0, stream>>>
(args.ids_dst, args.expert_bounds, args.dst, tmp_fixup.ptr, blocks_per_ne00_fd, args.nrows_x, args.ncols_dst,
args.nrows_dst, nchannels_y_fd, args.stride_channel_dst, nsamples_y_fd, args.stride_sample_dst,
ntx_fd);
}
-template <ggml_type type, bool fallback>
+template <ggml_type type, bool fallback, ggml_prec prec_src1 = GGML_PREC_Q8>
void mul_mat_q_switch_J(ggml_backend_cuda_context & ctx, const mmq_args & args, cudaStream_t stream) {
const int id = ggml_cuda_get_device();
const int cc = ggml_cuda_info().devices[id].cc;
@@ -1484,7 +1492,7 @@ void mul_mat_q_switch_J(ggml_backend_cuda_context & ctx, const mmq_args & args,
int ntiles_J_best = INT_MAX;
for (int J = 8; J <= 128 && ntiles_J_best > 1; J += 8) {
- const ggml_cuda_mmq_config config = ggml_cuda_mmq_get_config(type, J, fallback, cc);
+ const ggml_cuda_mmq_config config = ggml_cuda_mmq_get_config(type, J, fallback, cc, prec_src1);
if (config.type == GGML_TYPE_COUNT) {
continue;
}
@@ -1503,52 +1511,52 @@ void mul_mat_q_switch_J(ggml_backend_cuda_context & ctx, const mmq_args & args,
switch (J_best) {
case 8:
- launch_mul_mat_q<type, 8, fallback>(ctx, args, stream);
+ launch_mul_mat_q<type, 8, fallback, prec_src1>(ctx, args, stream);
break;
case 16:
- launch_mul_mat_q<type, 16, fallback>(ctx, args, stream);
+ launch_mul_mat_q<type, 16, fallback, prec_src1>(ctx, args, stream);
break;
case 24:
- launch_mul_mat_q<type, 24, fallback>(ctx, args, stream);
+ launch_mul_mat_q<type, 24, fallback, prec_src1>(ctx, args, stream);
break;
case 32:
- launch_mul_mat_q<type, 32, fallback>(ctx, args, stream);
+ launch_mul_mat_q<type, 32, fallback, prec_src1>(ctx, args, stream);
break;
case 40:
- launch_mul_mat_q<type, 40, fallback>(ctx, args, stream);
+ launch_mul_mat_q<type, 40, fallback, prec_src1>(ctx, args, stream);
break;
case 48:
- launch_mul_mat_q<type, 48, fallback>(ctx, args, stream);
+ launch_mul_mat_q<type, 48, fallback, prec_src1>(ctx, args, stream);
break;
case 56:
- launch_mul_mat_q<type, 56, fallback>(ctx, args, stream);
+ launch_mul_mat_q<type, 56, fallback, prec_src1>(ctx, args, stream);
break;
case 64:
- launch_mul_mat_q<type, 64, fallback>(ctx, args, stream);
+ launch_mul_mat_q<type, 64, fallback, prec_src1>(ctx, args, stream);
break;
case 72:
- launch_mul_mat_q<type, 72, fallback>(ctx, args, stream);
+ launch_mul_mat_q<type, 72, fallback, prec_src1>(ctx, args, stream);
break;
case 80:
- launch_mul_mat_q<type, 80, fallback>(ctx, args, stream);
+ launch_mul_mat_q<type, 80, fallback, prec_src1>(ctx, args, stream);
break;
case 88:
- launch_mul_mat_q<type, 88, fallback>(ctx, args, stream);
+ launch_mul_mat_q<type, 88, fallback, prec_src1>(ctx, args, stream);
break;
case 96:
- launch_mul_mat_q<type, 96, fallback>(ctx, args, stream);
+ launch_mul_mat_q<type, 96, fallback, prec_src1>(ctx, args, stream);
break;
case 104:
- launch_mul_mat_q<type, 104, fallback>(ctx, args, stream);
+ launch_mul_mat_q<type, 104, fallback, prec_src1>(ctx, args, stream);
break;
case 112:
- launch_mul_mat_q<type, 112, fallback>(ctx, args, stream);
+ launch_mul_mat_q<type, 112, fallback, prec_src1>(ctx, args, stream);
break;
case 120:
- launch_mul_mat_q<type, 120, fallback>(ctx, args, stream);
+ launch_mul_mat_q<type, 120, fallback, prec_src1>(ctx, args, stream);
break;
case 128:
- launch_mul_mat_q<type, 128, fallback>(ctx, args, stream);
+ launch_mul_mat_q<type, 128, fallback, prec_src1>(ctx, args, stream);
break;
default:
fprintf(stderr, "J_best=%d\n", J_best);
@@ -1557,20 +1565,24 @@ void mul_mat_q_switch_J(ggml_backend_cuda_context & ctx, const mmq_args & args,
}
}
-template <ggml_type type>
+template <ggml_type type, ggml_prec prec_src1 = GGML_PREC_Q8>
void mul_mat_q_case(ggml_backend_cuda_context & ctx, const mmq_args & args, cudaStream_t stream) {
if (args.nrows_x % 128 == 0) {
constexpr bool fallback = false;
- mul_mat_q_switch_J<type, fallback>(ctx, args, stream);
+ mul_mat_q_switch_J<type, fallback, prec_src1>(ctx, args, stream);
} else {
constexpr bool fallback = true;
- mul_mat_q_switch_J<type, fallback>(ctx, args, stream);
+ mul_mat_q_switch_J<type, fallback, prec_src1>(ctx, args, stream);
}
}
#define DECL_MMQ_CASE(type) \
template void mul_mat_q_case<type>(ggml_backend_cuda_context & ctx, const mmq_args & args, cudaStream_t stream) \
+// FP4 variant: uses native FP4 MMA instead of keeping src1 at Q8_1.
+#define DECL_MMQ_CASE_W4A4(type) \
+ template void mul_mat_q_case<type, GGML_PREC_Q4>(ggml_backend_cuda_context & ctx, const mmq_args & args, cudaStream_t stream) \
+
extern DECL_MMQ_CASE(GGML_TYPE_Q1_0);
extern DECL_MMQ_CASE(GGML_TYPE_Q2_0);
extern DECL_MMQ_CASE(GGML_TYPE_Q4_0);
@@ -1596,6 +1608,8 @@ extern DECL_MMQ_CASE(GGML_TYPE_IQ4_XS);
// -----------------------------------------
extern DECL_MMQ_CASE(GGML_TYPE_MXFP4);
extern DECL_MMQ_CASE(GGML_TYPE_NVFP4);
+extern DECL_MMQ_CASE_W4A4(GGML_TYPE_MXFP4);
+extern DECL_MMQ_CASE_W4A4(GGML_TYPE_NVFP4);
// -------------------------------------------------------------------------------------------------------------------------
diff --git a/ggml/src/ggml-cuda/template-instances/generate_cu_files.py b/ggml/src/ggml-cuda/template-instances/generate_cu_files.py
index d7cd27167..7be409bd3 100755
--- a/ggml/src/ggml-cuda/template-instances/generate_cu_files.py
+++ b/ggml/src/ggml-cuda/template-instances/generate_cu_files.py
@@ -50,6 +50,12 @@ SOURCE_MMQ = """// This file has been autogenerated by generate_cu_files.py, do
DECL_MMQ_CASE({type});
"""
+TYPES_MMQ_W4A4 = ["GGML_TYPE_MXFP4", "GGML_TYPE_NVFP4"]
+
+SOURCE_MMQ_W4A4 = """
+DECL_MMQ_CASE_W4A4({type});
+"""
+
SOURCE_MMF = """// This file has been autogenerated by generate_cu_files.py, do not edit manually.
#include "../mmf.cuh"
@@ -105,6 +111,8 @@ for ncols in [8, 16, 32, 64]:
for type in TYPES_MMQ:
with open(f"mmq-instance-{get_short_name(type)}.cu", "w") as f:
f.write(SOURCE_MMQ.format(type=type))
+ if type in TYPES_MMQ_W4A4:
+ f.write(SOURCE_MMQ_W4A4.format(type=type))
for type in range(1, 17):
with open(f"mmf-instance-ncols_{type}.cu", "w") as f:
diff --git a/ggml/src/ggml-cuda/template-instances/mmq-instance-mxfp4.cu b/ggml/src/ggml-cuda/template-instances/mmq-instance-mxfp4.cu
index c14624c52..0c1c7f291 100644
--- a/ggml/src/ggml-cuda/template-instances/mmq-instance-mxfp4.cu
+++ b/ggml/src/ggml-cuda/template-instances/mmq-instance-mxfp4.cu
@@ -3,3 +3,5 @@
#include "../mmq.cuh"
DECL_MMQ_CASE(GGML_TYPE_MXFP4);
+
+DECL_MMQ_CASE_W4A4(GGML_TYPE_MXFP4);
diff --git a/ggml/src/ggml-cuda/template-instances/mmq-instance-nvfp4.cu b/ggml/src/ggml-cuda/template-instances/mmq-instance-nvfp4.cu
index 2cb140d35..bed8c8e2e 100644
--- a/ggml/src/ggml-cuda/template-instances/mmq-instance-nvfp4.cu
+++ b/ggml/src/ggml-cuda/template-instances/mmq-instance-nvfp4.cu
@@ -3,3 +3,5 @@
#include "../mmq.cuh"
DECL_MMQ_CASE(GGML_TYPE_NVFP4);
+
+DECL_MMQ_CASE_W4A4(GGML_TYPE_NVFP4);
diff --git a/gguf-py/gguf/constants.py b/gguf-py/gguf/constants.py
index f585adcee..9a7a5e5bf 100644
--- a/gguf-py/gguf/constants.py
+++ b/gguf-py/gguf/constants.py
@@ -26,6 +26,10 @@ class Keys:
ALIGNMENT = "general.alignment"
FILE_TYPE = "general.file_type"
+ # Per-tensor extra options (tensor name array + parallel option arrays, e.g. prec_a4).
+ TENSOR_EXTRA_NAME = "general.tensor_extra.name"
+ TENSOR_EXTRA_PREC_A4 = "general.tensor_extra.prec_a4"
+
# Recommended Sampler Parameters
SAMPLING_SEQUENCE = "general.sampling.sequence"
SAMPLING_TOP_K = "general.sampling.top_k"
diff --git a/gguf-py/gguf/gguf_writer.py b/gguf-py/gguf/gguf_writer.py
index 56cc65a70..cf7b367e5 100644
--- a/gguf-py/gguf/gguf_writer.py
+++ b/gguf-py/gguf/gguf_writer.py
@@ -520,6 +520,12 @@ class GGUFWriter:
def add_file_type(self, ftype: int) -> None:
self.add_uint32(Keys.General.FILE_TYPE, ftype)
+ def add_tensor_extra_prec_a4(self, tensor_names: Sequence[str], values: Sequence[bool]) -> None:
+ if len(tensor_names) != len(values):
+ raise ValueError("tensor_extra prec_a4 names and values must have the same length")
+ self.add_array(Keys.General.TENSOR_EXTRA_NAME, list(tensor_names))
+ self.add_array(Keys.General.TENSOR_EXTRA_PREC_A4, list(values))
+
def add_sampling_sequence(self, sequence: str) -> None:
self.add_string(Keys.General.SAMPLING_SEQUENCE, sequence)
diff --git a/src/llama-arch.cpp b/src/llama-arch.cpp
index 03b7951c3..9d2331c2f 100644
--- a/src/llama-arch.cpp
+++ b/src/llama-arch.cpp
@@ -162,31 +162,33 @@ static const std::map<llm_arch, const char *> LLM_ARCH_NAMES = {
};
static const std::map<llm_kv, const char *> LLM_KV_NAMES = {
- { LLM_KV_GENERAL_TYPE, "general.type" },
- { LLM_KV_GENERAL_ARCHITECTURE, "general.architecture" },
- { LLM_KV_GENERAL_QUANTIZATION_VERSION, "general.quantization_version" },
- { LLM_KV_GENERAL_ALIGNMENT, "general.alignment" },
- { LLM_KV_GENERAL_FILE_TYPE, "general.file_type" },
- { LLM_KV_GENERAL_SAMPLING_SEQUENCE, "general.sampling.sequence" },
- { LLM_KV_GENERAL_SAMPLING_TOP_K, "general.sampling.top_k" },
- { LLM_KV_GENERAL_SAMPLING_TOP_P, "general.sampling.top_p" },
- { LLM_KV_GENERAL_SAMPLING_MIN_P, "general.sampling.min_p" },
- { LLM_KV_GENERAL_SAMPLING_XTC_PROBABILITY, "general.sampling.xtc_probability" },
- { LLM_KV_GENERAL_SAMPLING_XTC_THRESHOLD, "general.sampling.xtc_threshold" },
- { LLM_KV_GENERAL_SAMPLING_TEMP, "general.sampling.temp" },
- { LLM_KV_GENERAL_SAMPLING_PENALTY_LAST_N, "general.sampling.penalty_last_n" },
- { LLM_KV_GENERAL_SAMPLING_PENALTY_REPEAT, "general.sampling.penalty_repeat" },
- { LLM_KV_GENERAL_SAMPLING_MIROSTAT, "general.sampling.mirostat" },
- { LLM_KV_GENERAL_SAMPLING_MIROSTAT_TAU, "general.sampling.mirostat_tau" },
- { LLM_KV_GENERAL_SAMPLING_MIROSTAT_ETA, "general.sampling.mirostat_eta" },
- { LLM_KV_GENERAL_NAME, "general.name" },
- { LLM_KV_GENERAL_AUTHOR, "general.author" },
- { LLM_KV_GENERAL_VERSION, "general.version" },
- { LLM_KV_GENERAL_URL, "general.url" },
- { LLM_KV_GENERAL_DESCRIPTION, "general.description" },
- { LLM_KV_GENERAL_LICENSE, "general.license" },
- { LLM_KV_GENERAL_SOURCE_URL, "general.source.url" },
- { LLM_KV_GENERAL_SOURCE_HF_REPO, "general.source.huggingface.repository" },
+ { LLM_KV_GENERAL_TYPE, "general.type" },
+ { LLM_KV_GENERAL_ARCHITECTURE, "general.architecture" },
+ { LLM_KV_GENERAL_QUANTIZATION_VERSION, "general.quantization_version" },
+ { LLM_KV_GENERAL_ALIGNMENT, "general.alignment" },
+ { LLM_KV_GENERAL_FILE_TYPE, "general.file_type" },
+ { LLM_KV_GENERAL_SAMPLING_SEQUENCE, "general.sampling.sequence" },
+ { LLM_KV_GENERAL_SAMPLING_TOP_K, "general.sampling.top_k" },
+ { LLM_KV_GENERAL_SAMPLING_TOP_P, "general.sampling.top_p" },
+ { LLM_KV_GENERAL_SAMPLING_MIN_P, "general.sampling.min_p" },
+ { LLM_KV_GENERAL_SAMPLING_XTC_PROBABILITY, "general.sampling.xtc_probability" },
+ { LLM_KV_GENERAL_SAMPLING_XTC_THRESHOLD, "general.sampling.xtc_threshold" },
+ { LLM_KV_GENERAL_SAMPLING_TEMP, "general.sampling.temp" },
+ { LLM_KV_GENERAL_SAMPLING_PENALTY_LAST_N, "general.sampling.penalty_last_n" },
+ { LLM_KV_GENERAL_SAMPLING_PENALTY_REPEAT, "general.sampling.penalty_repeat" },
+ { LLM_KV_GENERAL_SAMPLING_MIROSTAT, "general.sampling.mirostat" },
+ { LLM_KV_GENERAL_SAMPLING_MIROSTAT_TAU, "general.sampling.mirostat_tau" },
+ { LLM_KV_GENERAL_SAMPLING_MIROSTAT_ETA, "general.sampling.mirostat_eta" },
+ { LLM_KV_GENERAL_NAME, "general.name" },
+ { LLM_KV_GENERAL_AUTHOR, "general.author" },
+ { LLM_KV_GENERAL_VERSION, "general.version" },
+ { LLM_KV_GENERAL_URL, "general.url" },
+ { LLM_KV_GENERAL_DESCRIPTION, "general.description" },
+ { LLM_KV_GENERAL_LICENSE, "general.license" },
+ { LLM_KV_GENERAL_SOURCE_URL, "general.source.url" },
+ { LLM_KV_GENERAL_SOURCE_HF_REPO, "general.source.huggingface.repository" },
+ { LLM_KV_GENERAL_TENSOR_EXTRA_NAME, "general.tensor_extra.name" },
+ { LLM_KV_GENERAL_TENSOR_EXTRA_PREC_A4, "general.tensor_extra.prec_a4" },
{ LLM_KV_VOCAB_SIZE, "%s.vocab_size" },
{ LLM_KV_CONTEXT_LENGTH, "%s.context_length" },
diff --git a/src/llama-arch.h b/src/llama-arch.h
index af64f4ec5..23b6b3810 100644
--- a/src/llama-arch.h
+++ b/src/llama-arch.h
@@ -192,6 +192,8 @@ enum llm_kv {
LLM_KV_GENERAL_LICENSE,
LLM_KV_GENERAL_SOURCE_URL,
LLM_KV_GENERAL_SOURCE_HF_REPO,
+ LLM_KV_GENERAL_TENSOR_EXTRA_NAME,
+ LLM_KV_GENERAL_TENSOR_EXTRA_PREC_A4,
LLM_KV_VOCAB_SIZE,
LLM_KV_CONTEXT_LENGTH,
diff --git a/src/llama-context.cpp b/src/llama-context.cpp
index da9c558d2..8675f6087 100644
--- a/src/llama-context.cpp
+++ b/src/llama-context.cpp
@@ -2555,6 +2555,7 @@ llm_graph_params llama_context::graph_params(
/*.loras =*/ loras.get(),
/*.mctx =*/ mctx,
/*.cross =*/ &cross,
+ /*.prec_policy =*/ &model.prec_policy,
/*.samplers =*/ sampling.samplers,
/*.n_outputs =*/ n_outputs,
/*.cb =*/ graph_get_cb(),
diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp
index 07ca49ad0..0b3bab612 100644
--- a/src/llama-graph.cpp
+++ b/src/llama-graph.cpp
@@ -1489,6 +1489,7 @@ llm_graph_context::llm_graph_context(const llm_graph_params & params) :
loras (params.loras),
mctx (params.mctx),
cross (params.cross),
+ prec_policy (params.prec_policy),
samplers (params.samplers),
cb_func (params.cb),
res (params.res),
@@ -1517,6 +1518,10 @@ ggml_tensor * llm_graph_context::build_lora_mm(
ggml_tensor * w_s) const {
ggml_tensor * res = ggml_mul_mat(ctx0, w, cur);
+ if (prec_policy) {
+ prec_policy->apply(res);
+ }
+
if (w_s) {
res = ggml_mul(ctx0, res, w_s);
}
@@ -1549,6 +1554,10 @@ ggml_tensor * llm_graph_context::build_lora_mm_id(
ggml_tensor * w_s) const {
ggml_tensor * res = ggml_mul_mat_id(ctx0, w, cur, ids);
+ if (prec_policy) {
+ prec_policy->apply(res);
+ }
+
if (w_s) {
const int64_t n_expert = w_s->ne[0];
const int64_t n_tokens = cur->ne[2];
diff --git a/src/llama-graph.h b/src/llama-graph.h
index cc4110639..3daa425bc 100644
--- a/src/llama-graph.h
+++ b/src/llama-graph.h
@@ -19,6 +19,7 @@ struct ggml_tensor;
struct llama_cparams;
struct llama_layer;
+struct llama_prec_policy;
struct llama_memory_context_i;
@@ -787,6 +788,8 @@ struct llm_graph_params {
const llama_memory_context_i * mctx;
const llama_cross * cross;
+ const llama_prec_policy * prec_policy = nullptr;
+
std::map<llama_seq_id, llama_sampler *> samplers;
static bool samplers_equal(
@@ -1027,6 +1030,8 @@ struct llm_graph_context {
const llama_memory_context_i * mctx;
const llama_cross * cross;
+ const llama_prec_policy * prec_policy;
+
std::map<llama_seq_id, llama_sampler *> samplers;
const llm_graph_cb & cb_func;
diff --git a/src/llama-model-saver.cpp b/src/llama-model-saver.cpp
index 39160a417..6f58dd150 100644
--- a/src/llama-model-saver.cpp
+++ b/src/llama-model-saver.cpp
@@ -193,6 +193,19 @@ void llama_model_saver::add_kv_from_model() {
// add_kv(LLM_KV_GENERAL_SAMPLING_MIROSTAT_TAU, ???);
// add_kv(LLM_KV_GENERAL_SAMPLING_MIROSTAT_ETA, ???);
add_kv(LLM_KV_GENERAL_NAME, model->name);
+
+ if (!model->prec_policy.prec_src1.empty()) {
+ std::vector<std::string> tensor_names;
+ std::vector<int8_t> values;
+ tensor_names.reserve(model->prec_policy.prec_src1.size());
+ values.reserve(model->prec_policy.prec_src1.size());
+ for (const auto & [w, prec] : model->prec_policy.prec_src1) {
+ tensor_names.push_back(ggml_get_name(w));
+ values.push_back(prec == GGML_PREC_Q8 ? 0 : 1);
+ }
+ add_kv(LLM_KV_GENERAL_TENSOR_EXTRA_NAME, tensor_names);
+ gguf_set_arr_data(gguf_ctx, llm_kv(LLM_KV_GENERAL_TENSOR_EXTRA_PREC_A4).c_str(), GGUF_TYPE_BOOL, values.data(), values.size());
+ }
// add_kv(LLM_KV_GENERAL_AUTHOR, ???);
// add_kv(LLM_KV_GENERAL_VERSION, ???);
// add_kv(LLM_KV_GENERAL_URL, ???);
diff --git a/src/llama-model.cpp b/src/llama-model.cpp
index 6fdde70ce..ab5e744b5 100644
--- a/src/llama-model.cpp
+++ b/src/llama-model.cpp
@@ -1221,6 +1221,56 @@ struct llama_model::impl {
std::vector<float> tensor_split_owned;
};
+bool llama_prec_policy::apply(ggml_tensor * res) const {
+ if (!res || !res->src[0]) {
+ return false;
+ }
+
+ const auto it = prec_src1.find(res->src[0]);
+ if (it == prec_src1.end()) {
+ return false;
+ }
+
+ return ggml_prec_set_src(res, it->second, 1);
+}
+
+void llama_prec_policy::load(llama_model_loader & ml, const llama_model & model) {
+ std::vector<std::string> tensor_names;
+ if (!ml.get_arr(LLM_KV_GENERAL_TENSOR_EXTRA_NAME, tensor_names, false)) {
+ return;
+ }
+
+ const gguf_context * ctx = ml.metadata;
+ const std::string key = ml.llm_kv(LLM_KV_GENERAL_TENSOR_EXTRA_PREC_A4);
+ const int kid = gguf_find_key(ctx, key.c_str());
+ if (kid < 0 || gguf_get_kv_type(ctx, kid) != GGUF_TYPE_ARRAY || gguf_get_arr_type(ctx, kid) != GGUF_TYPE_BOOL) {
+ throw std::runtime_error(format("%s must be a bool array", key.c_str()));
+ }
+
+ const size_t n_values = gguf_get_arr_n(ctx, kid);
+ if (n_values != tensor_names.size()) {
+ throw std::runtime_error(format(
+ "%s tensor/value length mismatch (%zu vs %zu)",
+ key.c_str(), tensor_names.size(), n_values));
+ }
+
+ // tensors that can not use 4-bit activations, keep src1 at higher precision
+ const int8_t * values = (const int8_t *) gguf_get_arr_data(ctx, kid);
+ std::unordered_set<std::string> want;
+ for (size_t i = 0; i < n_values; ++i) {
+ if (values[i] == 0) {
+ want.insert(tensor_names[i]);
+ }
+ }
+
+ // resolve names to tensor pointers
+ for (const auto & [name, w] : model.tensors_by_name) {
+ if (want.count(name)) {
+ prec_src1.emplace(w, GGML_PREC_Q8);
+ }
+ }
+}
+
llama_model::llama_model(const llama_model_params & params) : params(params), pimpl(std::make_unique<impl>()) {
if (params.tensor_split != nullptr) {
// llama_model_params stores tensor_split as a borrowed pointer, but the model
@@ -1745,6 +1795,9 @@ bool llama_model_base::load_tensors(llama_model_loader & ml) {
}
}
+ // per-tensor activation precision policy
+ prec_policy.load(ml, *this);
+
ml.init_mappings(true, use_mlock ? &pimpl->mlock_mmaps : nullptr);
pimpl->mappings.reserve(ml.mappings.size());
diff --git a/src/llama-model.h b/src/llama-model.h
index 984b2cf38..a0f9f1142 100644
--- a/src/llama-model.h
+++ b/src/llama-model.h
@@ -17,6 +17,7 @@
struct llama_cparams;
struct llama_ubatch;
struct llama_model_loader;
+struct llama_model;
// available models
enum llm_type {
@@ -609,6 +610,19 @@ struct llama_meta_device_get_split_state_userdata {
struct ggml_backend_meta_split_state llama_meta_device_get_split_state(const struct ggml_tensor * tensor, void * userdata);
+struct llama_prec_policy {
+ // the key is the weight tensor `res->src[0]`, stores the recommended accumulation type of the op (unused for now)
+ // TODO: migrate ad-hoc ggml_prec_set_acc() calls to this container + update apply() to use it
+ std::unordered_map<const ggml_tensor *, ggml_prec> prec_acc;
+
+ // the key is the weight tensor `res->src[0]`, stores the recommended activation precision type
+ std::unordered_map<const ggml_tensor *, ggml_prec> prec_src1;
+
+ bool apply(ggml_tensor * res) const;
+
+ void load(llama_model_loader & ml, const llama_model & model);
+};
+
struct llama_model {
llm_type type = LLM_TYPE_UNKNOWN;
llm_arch arch = LLM_ARCH_UNKNOWN;
@@ -618,6 +632,9 @@ struct llama_model {
llama_hparams hparams = {};
llama_vocab vocab;
+ // per-tensor activation precision policy
+ llama_prec_policy prec_policy;
+
// for classifier models
std::vector<std::string> classifier_labels;
diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp
index 917faabfa..89c18a2d5 100644
--- a/tests/test-backend-ops.cpp
+++ b/tests/test-backend-ops.cpp
@@ -4920,6 +4920,27 @@ struct test_rwkv_wkv7 : public test_case {
}
};
+static int32_t test_get_op_params_i32(const ggml_tensor * tensor, uint32_t i) {
+ GGML_ASSERT(i < GGML_MAX_OP_PARAMS / sizeof(int32_t));
+ return tensor->op_params[i];
+}
+
+// true if any node of the given op requests 8-bit src1 (GGML_PREC_Q8)
+static bool graph_mul_mat_hi_prec_act(ggml_cgraph * gf, ggml_op op) {
+ if (gf == nullptr) {
+ return false;
+ }
+
+ ggml_tensor ** nodes = ggml_graph_nodes(gf);
+ for (int i = 0; i < ggml_graph_n_nodes(gf); ++i) {
+ if (nodes[i]->op == op && test_get_op_params_i32(nodes[i], 3) == GGML_PREC_Q8) {
+ return true;
+ }
+ }
+
+ return false;
+}
+
// GGML_OP_MUL_MAT
struct test_mul_mat : public test_case {
const ggml_type type_a;
@@ -4944,7 +4965,9 @@ struct test_mul_mat : public test_case {
double max_nmse_err(ggml_backend_t backend) override {
// for blackwell we quantize activations to mxfp4 instead of q8_1 so we add higher tolerance
- if ((type_a == GGML_TYPE_MXFP4 || type_a == GGML_TYPE_NVFP4) && backend_has_feature(backend, "BLACKWELL_NATIVE_FP4")) {
+ if ((type_a == GGML_TYPE_MXFP4 || type_a == GGML_TYPE_NVFP4) &&
+ !graph_mul_mat_hi_prec_act(gf, GGML_OP_MUL_MAT) &&
+ backend_has_feature(backend, "BLACKWELL_NATIVE_FP4")) {
return 2e-2;
}
return max_nmse_err();
@@ -5103,6 +5126,41 @@ struct test_mul_mat_hadamard : public test_mul_mat {
}
};
+// FP4 W4A8 path (GGML_PREC_Q8 on src1 disallows 4-bit activations)
+struct test_mul_mat_w4a8 : public test_mul_mat {
+ test_mul_mat_w4a8(ggml_type type_a = GGML_TYPE_NVFP4, ggml_type type_b = GGML_TYPE_F32,
+ int64_t m = 32, int64_t n = 32, int64_t k = 256,
+ std::array<int64_t, 2> bs = {1, 1},
+ std::array<int64_t, 2> nr = {1, 1})
+ : test_mul_mat(type_a, type_b, m, n, k, bs, nr) {}
+ ggml_tensor * build_graph(ggml_context * ctx) override {
+ ggml_tensor * out = test_mul_mat::build_graph(ctx);
+ for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != NULL; t = ggml_get_next_tensor(ctx, t)) {
+ if (t->op == GGML_OP_MUL_MAT) {
+ ggml_prec_set_src(t, GGML_PREC_Q8, 1);
+ }
+ }
+ return out;
+ }
+ std::string op_desc(ggml_tensor * t) override {
+ GGML_UNUSED(t);
+ return "MUL_MAT_W4A8";
+ }
+};
+
+// FP4 native W4A4 path (default precision)
+struct test_mul_mat_w4a4 : public test_mul_mat {
+ test_mul_mat_w4a4(ggml_type type_a = GGML_TYPE_NVFP4, ggml_type type_b = GGML_TYPE_F32,
+ int64_t m = 32, int64_t n = 32, int64_t k = 256,
+ std::array<int64_t, 2> bs = {1, 1},
+ std::array<int64_t, 2> nr = {1, 1})
+ : test_mul_mat(type_a, type_b, m, n, k, bs, nr) {}
+ std::string op_desc(ggml_tensor * t) override {
+ GGML_UNUSED(t);
+ return "MUL_MAT_W4A4";
+ }
+};
+
static void init_mul_mat_id_ids(ggml_context * ctx, int n_mats) {
std::random_device rd;
std::default_random_engine rng(rd());
@@ -5156,7 +5214,9 @@ struct test_mul_mat_id : public test_case {
double max_nmse_err(ggml_backend_t backend) override {
// for blackwell we quantize activations to mxfp4 instead of q8_1 so we add higher tolerance
- if ((type_a == GGML_TYPE_MXFP4 || type_a == GGML_TYPE_NVFP4) && backend_has_feature(backend, "BLACKWELL_NATIVE_FP4")) {
+ if ((type_a == GGML_TYPE_MXFP4 || type_a == GGML_TYPE_NVFP4) &&
+ !graph_mul_mat_hi_prec_act(gf, GGML_OP_MUL_MAT_ID) &&
+ backend_has_feature(backend, "BLACKWELL_NATIVE_FP4")) {
return 2e-2;
}
return max_nmse_err();
@@ -5211,6 +5271,39 @@ struct test_mul_mat_id : public test_case {
}
};
+// FP4 W4A8 path on the MoE path (GGML_PREC_Q8 on src1 disallows 4-bit activations)
+struct test_mul_mat_id_w4a8 : public test_mul_mat_id {
+ test_mul_mat_id_w4a8(ggml_type type_a = GGML_TYPE_NVFP4, ggml_type type_b = GGML_TYPE_F32,
+ int n_mats = 8, int n_used = 2, bool b = false,
+ int64_t m = 32, int64_t n = 32, int64_t k = 256)
+ : test_mul_mat_id(type_a, type_b, n_mats, n_used, b, m, n, k) {}
+ ggml_tensor * build_graph(ggml_context * ctx) override {
+ ggml_tensor * out = test_mul_mat_id::build_graph(ctx);
+ for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != NULL; t = ggml_get_next_tensor(ctx, t)) {
+ if (t->op == GGML_OP_MUL_MAT_ID) {
+ ggml_prec_set_src(t, GGML_PREC_Q8, 1);
+ }
+ }
+ return out;
+ }
+ std::string op_desc(ggml_tensor * t) override {
+ GGML_UNUSED(t);
+ return "MUL_MAT_ID_W4A8";
+ }
+};
+
+// FP4 native W4A4 path on the MoE path (default precision)
+struct test_mul_mat_id_w4a4 : public test_mul_mat_id {
+ test_mul_mat_id_w4a4(ggml_type type_a = GGML_TYPE_NVFP4, ggml_type type_b = GGML_TYPE_F32,
+ int n_mats = 8, int n_used = 2, bool b = false,
+ int64_t m = 32, int64_t n = 32, int64_t k = 256)
+ : test_mul_mat_id(type_a, type_b, n_mats, n_used, b, m, n, k) {}
+ std::string op_desc(ggml_tensor * t) override {
+ GGML_UNUSED(t);
+ return "MUL_MAT_ID_W4A4";
+ }
+};
+
// GGML_OP_MUL_MAT_ID + GGML_OP_ADD or GGML_OP_MUL
struct test_mul_mat_id_fusion : public test_case {
const ggml_type type_a;
@@ -9934,6 +10027,22 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
test_cases.emplace_back(new test_mul_mat_hadamard(GGML_TYPE_F32, GGML_TYPE_F16, 128, 4, 128, {2, 3}));
test_cases.emplace_back(new test_mul_mat_hadamard(GGML_TYPE_F32, GGML_TYPE_F16, 256, 512, 256)); // many rows
+ // FP4 activation precision (default = native W4A4, src1 GGML_PREC_Q8 = W4A8)
+ test_cases.emplace_back(new test_mul_mat_w4a8(GGML_TYPE_NVFP4, GGML_TYPE_F32, 32, 1, 256));
+ test_cases.emplace_back(new test_mul_mat_w4a8(GGML_TYPE_NVFP4, GGML_TYPE_F32, 32, 32, 256));
+ test_cases.emplace_back(new test_mul_mat_w4a8(GGML_TYPE_NVFP4, GGML_TYPE_F32, 64, 16, 512));
+ test_cases.emplace_back(new test_mul_mat_w4a4(GGML_TYPE_NVFP4, GGML_TYPE_F32, 32, 1, 256));
+ test_cases.emplace_back(new test_mul_mat_w4a4(GGML_TYPE_NVFP4, GGML_TYPE_F32, 32, 32, 256));
+ test_cases.emplace_back(new test_mul_mat_id_w4a8(GGML_TYPE_NVFP4, GGML_TYPE_F32, 8, 2, false, 32, 32, 256));
+ test_cases.emplace_back(new test_mul_mat_id_w4a8(GGML_TYPE_NVFP4, GGML_TYPE_F32, 4, 2, true, 64, 16, 256));
+ test_cases.emplace_back(new test_mul_mat_id_w4a4(GGML_TYPE_NVFP4, GGML_TYPE_F32, 8, 2, false, 32, 32, 256));
+ test_cases.emplace_back(new test_mul_mat_id_w4a4(GGML_TYPE_NVFP4, GGML_TYPE_F32, 4, 2, true, 64, 16, 256));
+ test_cases.emplace_back(new test_mul_mat_w4a8(GGML_TYPE_MXFP4, GGML_TYPE_F32, 32, 32, 256));
+ test_cases.emplace_back(new test_mul_mat_w4a8(GGML_TYPE_MXFP4, GGML_TYPE_F32, 64, 16, 512));
+ test_cases.emplace_back(new test_mul_mat_w4a4(GGML_TYPE_MXFP4, GGML_TYPE_F32, 32, 32, 256));
+ test_cases.emplace_back(new test_mul_mat_id_w4a8(GGML_TYPE_MXFP4, GGML_TYPE_F32, 8, 2, false, 32, 32, 256));
+ test_cases.emplace_back(new test_mul_mat_id_w4a4(GGML_TYPE_MXFP4, GGML_TYPE_F32, 8, 2, false, 32, 32, 256));
+
#if 0
// > 4GB A matrix. Too slow to be enabled by default.
test_cases.emplace_back(new test_mul_mat(GGML_TYPE_F16, GGML_TYPE_F16, 900000, 3, 2592, {1, 1}, {1, 1}));