diff --git a/conversion/base.py b/conversion/base.py index 9fba5a86b054..5561481e7716 100644 --- a/conversion/base.py +++ b/conversion/base.py @@ -170,6 +170,9 @@ def __init__(self, dir_model: Path, ftype: gguf.LlamaFileType, fname_out: Path, 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 @@ def prepare_qkv_fusion(self) -> None: 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 @@ def _repack_nvfp4(self, name: str, weight: Tensor, scale: Tensor, scale2: Tensor 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 @@ def _flush_nvfp4_experts(self, key, expert_blocks, expert_scales, expert_input_s 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 @@ def prepare_tensors(self): 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 @@ def prepare_tensors(self): 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 @@ def prepare_metadata(self, vocab_only: bool): 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 70fc17af2402..f3fdbbc79e2f 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 7f00bad943e2..e19f4f24d78b 100644 --- a/ggml/src/ggml-cuda/mmq-load-tiles.cuh +++ b/ggml/src/ggml-cuda/mmq-load-tiles.cuh @@ -1631,16 +1631,16 @@ template static __device__ __forceinline_ template 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 static __device__ __forceinline_ } } -template static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_nvfp4( +template 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 static __device__ __forceinline_ template 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 4d1c398fc545..4ca6542d80bf 100644 --- a/ggml/src/ggml-cuda/mmq-vec-dot.cuh +++ b/ggml/src/ggml-cuda/mmq-vec-dot.cuh @@ -474,7 +474,7 @@ template static __device__ __forceinline_ } // Used for Q3_K, IQ2_S, and IQ2_XS: -template static __device__ __forceinline__ void ggml_cuda_mmq_vec_dot_q8_0_16_q8_1_mma( +template 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 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 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 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 b13b34ee9cf0..27708035eb8a 100644 --- a/ggml/src/ggml-cuda/mmq.cu +++ b/ggml/src/ggml-cuda/mmq.cu @@ -5,7 +5,7 @@ #include -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(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(ctx, args, stream); + break; + } mul_mat_q_case(ctx, args, stream); break; case GGML_TYPE_NVFP4: + if (prec_src1 == GGML_PREC_Q4) { + mul_mat_q_case(ctx, args, stream); + break; + } mul_mat_q_case(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 6923f3510c03..4b50d1dc333b 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 +template 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, - ggml_cuda_mmq_vec_dot_fp4_fp4_mma, - ggml_cuda_mmq_write_back_mma); + if (prec_src1 == GGML_PREC_Q4) { + return ggml_cuda_mmq_util_funcs( + -1, + ggml_cuda_mmq_load_tiles_mxfp4_fp4, + ggml_cuda_mmq_vec_dot_fp4_fp4_mma, + ggml_cuda_mmq_write_back_mma); + } + break; case GGML_TYPE_NVFP4: - return ggml_cuda_mmq_util_funcs( - -1, - ggml_cuda_mmq_load_tiles_nvfp4_nvfp4, - ggml_cuda_mmq_vec_dot_fp4_fp4_mma, - ggml_cuda_mmq_write_back_mma); + if (prec_src1 == GGML_PREC_Q4) { + return ggml_cuda_mmq_util_funcs( + -1, + ggml_cuda_mmq_load_tiles_nvfp4_nvfp4, + ggml_cuda_mmq_vec_dot_fp4_fp4_mma, + ggml_cuda_mmq_write_back_mma); + } + 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, - ggml_cuda_mmq_vec_dot_q8_0_16_q8_1_mma, + ggml_cuda_mmq_load_tiles_nvfp4, + ggml_cuda_mmq_vec_dot_q8_0_16_q8_1_mma, ggml_cuda_mmq_write_back_mma); default: return ggml_cuda_mmq_util_funcs(1, nullptr, nullptr, nullptr); } } -template +template static constexpr __device__ int ggml_cuda_mmq_get_vdr() { - return ggml_cuda_mmq_get_util_funcs().vdr; + return ggml_cuda_mmq_get_util_funcs().vdr; } -template +template static constexpr __device__ ggml_cuda_mmq_load_tiles_t ggml_cuda_mmq_get_load_tiles() { - return ggml_cuda_mmq_get_util_funcs().load_tiles; + return ggml_cuda_mmq_get_util_funcs().load_tiles; } -template +template static constexpr __device__ ggml_cuda_mmq_vec_dot_t ggml_cuda_mmq_get_vec_dot() { - return ggml_cuda_mmq_get_util_funcs().vec_dot; + return ggml_cuda_mmq_get_util_funcs().vec_dot; } -template +template static constexpr __device__ ggml_cuda_mmq_write_back_t ggml_cuda_mmq_get_write_back() { - return ggml_cuda_mmq_get_util_funcs().write_back; + return ggml_cuda_mmq_get_util_funcs().write_back; } // --------------------------------------------------------------------------------------------- -template +template 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::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(); - constexpr ggml_cuda_mmq_vec_dot_t vec_dot = ggml_cuda_mmq_get_vec_dot(); - constexpr ggml_cuda_mmq_write_back_t write_back = ggml_cuda_mmq_get_write_back(); + 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(); + constexpr ggml_cuda_mmq_vec_dot_t vec_dot = ggml_cuda_mmq_get_vec_dot(); + constexpr ggml_cuda_mmq_write_back_t write_back = ggml_cuda_mmq_get_write_back(); 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 -__launch_bounds__(ggml_cuda_mmq_get_nthreads(type, J, fallback), ggml_cuda_mmq_get_occupancy(type, J, fallback)) +template +__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::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 + mul_mat_q_process_tile (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 + mul_mat_q_process_tile (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 + mul_mat_q_process_tile (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 -__launch_bounds__(ggml_cuda_mmq_get_nthreads(type, J, fallback)/2, 1) +template +__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::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 +template 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), nbytes_shared); - CUDA_SET_SHARED_MEMORY_LIMIT((mul_mat_q), nbytes_shared); + CUDA_SET_SHARED_MEMORY_LIMIT((mul_mat_q), nbytes_shared); + CUDA_SET_SHARED_MEMORY_LIMIT((mul_mat_q), 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<<>> + if (!config.stream_k) { + mul_mat_q<<>> (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<<>> + mul_mat_q<<>> (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<<>> + mul_mat_q_stream_k_fixup<<>> (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 +template 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(ctx, args, stream); + launch_mul_mat_q(ctx, args, stream); break; case 16: - launch_mul_mat_q(ctx, args, stream); + launch_mul_mat_q(ctx, args, stream); break; case 24: - launch_mul_mat_q(ctx, args, stream); + launch_mul_mat_q(ctx, args, stream); break; case 32: - launch_mul_mat_q(ctx, args, stream); + launch_mul_mat_q(ctx, args, stream); break; case 40: - launch_mul_mat_q(ctx, args, stream); + launch_mul_mat_q(ctx, args, stream); break; case 48: - launch_mul_mat_q(ctx, args, stream); + launch_mul_mat_q(ctx, args, stream); break; case 56: - launch_mul_mat_q(ctx, args, stream); + launch_mul_mat_q(ctx, args, stream); break; case 64: - launch_mul_mat_q(ctx, args, stream); + launch_mul_mat_q(ctx, args, stream); break; case 72: - launch_mul_mat_q(ctx, args, stream); + launch_mul_mat_q(ctx, args, stream); break; case 80: - launch_mul_mat_q(ctx, args, stream); + launch_mul_mat_q(ctx, args, stream); break; case 88: - launch_mul_mat_q(ctx, args, stream); + launch_mul_mat_q(ctx, args, stream); break; case 96: - launch_mul_mat_q(ctx, args, stream); + launch_mul_mat_q(ctx, args, stream); break; case 104: - launch_mul_mat_q(ctx, args, stream); + launch_mul_mat_q(ctx, args, stream); break; case 112: - launch_mul_mat_q(ctx, args, stream); + launch_mul_mat_q(ctx, args, stream); break; case 120: - launch_mul_mat_q(ctx, args, stream); + launch_mul_mat_q(ctx, args, stream); break; case 128: - launch_mul_mat_q(ctx, args, stream); + launch_mul_mat_q(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 +template 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(ctx, args, stream); + mul_mat_q_switch_J(ctx, args, stream); } else { constexpr bool fallback = true; - mul_mat_q_switch_J(ctx, args, stream); + mul_mat_q_switch_J(ctx, args, stream); } } #define DECL_MMQ_CASE(type) \ template void mul_mat_q_case(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(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 d7cd271675e0..7be409bd36ab 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 @@ 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 @@ def get_short_name(long_quant_name): 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 c14624c52cad..0c1c7f291b03 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 2cb140d35a3c..bed8c8e2e4d6 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 f585adcee34a..9a7a5e5bfaf6 100644 --- a/gguf-py/gguf/constants.py +++ b/gguf-py/gguf/constants.py @@ -26,6 +26,10 @@ class General: 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 56cc65a70a0d..cf7b367e520e 100644 --- a/gguf-py/gguf/gguf_writer.py +++ b/gguf-py/gguf/gguf_writer.py @@ -520,6 +520,12 @@ def add_custom_alignment(self, alignment: int) -> None: 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 03b7951c3cc1..9d2331c2fefc 100644 --- a/src/llama-arch.cpp +++ b/src/llama-arch.cpp @@ -162,31 +162,33 @@ static const std::map LLM_ARCH_NAMES = { }; static const std::map 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 af64f4ec57a9..23b6b38100c0 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 da9c558d288c..8675f6087336 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 07ca49ad09d1..0b3bab612375 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 cc4110639d4b..3daa425bc07b 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 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 samplers; const llm_graph_cb & cb_func; diff --git a/src/llama-model-saver.cpp b/src/llama-model-saver.cpp index 39160a417e6a..6f58dd15092c 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 tensor_names; + std::vector 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 6fdde70ce0db..ab5e744b5d10 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -1221,6 +1221,56 @@ struct llama_model::impl { std::vector 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 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 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()) { 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 984b2cf38837..a0f9f11423e3 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 prec_acc; + + // the key is the weight tensor `res->src[0]`, stores the recommended activation precision type + std::unordered_map 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 classifier_labels; diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index 917faabfaa60..89c18a2d5269 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 bs = {1, 1}, + std::array 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 bs = {1, 1}, + std::array 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> 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}));