diff --git a/conversion/base.py b/conversion/base.py index 56547ace009..29d513feb17 100644 --- a/conversion/base.py +++ b/conversion/base.py @@ -160,6 +160,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_all_w4a16 = False # checkpoint quantized entirely as W4A16_NVFP4 + self._prec_a8_w4a4_override: set[str] = set() # gguf names explicitly kept at plain NVFP4 (W4A4) under a global W4A16 default + self._allow_prec_a8: dict[str, bool] = {} # gguf tensor name -> wants = 8-bit (A8) activations self._fp8_as_q8 = fp8_as_q8 self._fp8_dequantized: set[str] = set() @@ -617,6 +620,47 @@ def map_tensor_name(self, name: str, try_suffixes: Sequence[str] = (".weight", " raise ValueError(f"Can not map tensor {name!r}") return new_name + def _gguf_weight_name(self, name: str) -> str: + if name.endswith((".weight", ".bias")): + return name + return name + ".weight" + + def _tag_prec_a8(self, new_name: str) -> None: + # Tag a W4A16_NVFP4 weight as needing 8-bit activation. + if self._nvfp4_all_w4a16: + gguf_name = self._gguf_weight_name(new_name) + if gguf_name not in self._prec_a8_w4a4_override: + self._allow_prec_a8[gguf_name] = True + + def _hf_quant_tensors_to_gguf(self, hf_name: str) -> list[str]: + # Map an HF quantized-layer name to its GGUF tensor name(s). + if hf_name == "lm_head" or hf_name.endswith(".lm_head"): + return ["output.weight"] + + name = hf_name + if name.startswith("model.language_model."): + name = "model." + name[len("model.language_model."):] + + m = re.fullmatch(r"model\.layers\.(\d+)\.mlp\.experts", name) + if m: + bid = int(m.group(1)) + return [ + self._gguf_weight_name(self.format_tensor_name(gguf.MODEL_TENSOR.FFN_GATE_EXP, bid)), + self._gguf_weight_name(self.format_tensor_name(gguf.MODEL_TENSOR.FFN_UP_EXP, bid)), + self._gguf_weight_name(self.format_tensor_name(gguf.MODEL_TENSOR.FFN_DOWN_EXP, bid)), + ] + + candidates = [name] + if not name.endswith((".weight", ".bias")): + candidates.append(name + ".weight") + + for cand in candidates: + try: + return [self._gguf_weight_name(self.map_tensor_name(cand))] + except ValueError: + continue + return [] + def set_gguf_parameters(self): raise NotImplementedError("set_gguf_parameters() must be implemented in subclasses") @@ -726,6 +770,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_a8(new_name) self._write_scale_tensor(new_name.replace(".weight", ".scale"), scale2) self._write_scale_tensor(new_name.replace(".weight", ".input_scale"), input_scale) @@ -818,6 +863,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_a8(new_name) scales.sort(key=lambda x: x[0]) self._write_scales_tensor(new_name.replace(".weight", ".scale"), [s[1] for s in scales]) @@ -860,6 +906,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)) ) + # A checkpoint quantized entirely as W4A16_NVFP4 uses a single global quant_algo + self._nvfp4_all_w4a16 = quant_algo == "W4A16_NVFP4" + if quant_algo != "NVFP4": if nvfp4_compressed_tensors: quant_algo = "NVFP4" @@ -869,6 +918,21 @@ def prepare_tensors(self): self._is_nvfp4 = quant_algo in ("NVFP4", "W4A16_NVFP4") self._is_mxfp4 = quant_method == "mxfp4" + # Collect per-tensor W4A16_NVFP4 metadata (activations were not quantized to 4-bit). + if self._is_nvfp4 and not self._nvfp4_all_w4a16: + # Per-layer map tag only the W4A16_NVFP4 layers. + for tensor_name, entry in quant_layers.items(): + if not isinstance(entry, dict) or entry.get("quant_algo") != "W4A16_NVFP4": + continue + for gguf_name in self._hf_quant_tensors_to_gguf(tensor_name): + self._allow_prec_a8[gguf_name] = True + elif self._nvfp4_all_w4a16: + # Record any per-layer entries that override back to plain NVFP4 (W4A4) + for tensor_name, entry in quant_layers.items(): + algo = entry.get("quant_algo") if isinstance(entry, dict) else None + if isinstance(algo, str) and algo.endswith("NVFP4") and algo != "W4A16_NVFP4": + self._prec_a8_w4a4_override.update(self._hf_quant_tensors_to_gguf(tensor_name)) + # 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. @@ -1061,6 +1125,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._allow_prec_a8: + names = sorted(self._allow_prec_a8.keys()) + values = [self._allow_prec_a8[n] for n in names] + logger.info(f"Set allow_prec_a8 metadata for {len(names)} tensor(s)") + self.gguf_writer.add_tensor_extra_allow_prec_a8(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 ed48e7a05ec..2f7cad1799a 100644 --- a/docs/build.md +++ b/docs/build.md @@ -281,6 +281,10 @@ 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_FORCE_W4A4 + +NVFP4 models that carry W4A16 layers request higher-precision activations (W4A8), so on Blackwell those layers run through the W4A8 path instead of the native W4A4 path. Set `GGML_CUDA_FORCE_W4A4=1` to override that request and keep the native W4A4 path for faster prompt processing at the cost of accuracy. + ### 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/include/ggml.h b/ggml/include/ggml.h index 5f6774a630c..44d026c0201 100644 --- a/ggml/include/ggml.h +++ b/ggml/include/ggml.h @@ -434,9 +434,20 @@ extern "C" { }; // precision + // this enum is used to declare the allowed floating-point types that can be used during the compute of an op + // the declared types can be: + // - result accumulation type + // - source tensor data representation type + // - etc. + // the precision parameters are stored as ggml_tensor.op_params to the respective ops enum ggml_prec { - GGML_PREC_DEFAULT = 0, // stored as ggml_tensor.op_params, 0 by default - GGML_PREC_F32 = 10, + GGML_PREC_UNDEFINED = 0, + GGML_PREC_DEFAULT = 0, // note: deprecated, use GGML_PREC_UNDEFINED + GGML_PREC_F32 = 10, + GGML_PREC_BF16 = 15, + GGML_PREC_F16 = 20, + GGML_PREC_Q8 = 30, + GGML_PREC_Q4 = 40, }; // op hint @@ -1422,6 +1433,37 @@ extern "C" { struct ggml_tensor * b, float eps); + // set the minimum required accumulator type for the implementation to use during the compute + // for example: + // - GGML_PREC_F32 - requires accumulation of the results in F32 + // - GGML_PREC_BF16 - can accumulate the results in BF16, F32 + // - GGML_PREC_F16 - can accumulate the results in F16, F32 + // - GGML_PREC_Q8 - not allowed + // - GGML_PREC_Q4 - not allowed + GGML_API void ggml_prec_set_acc( + struct ggml_tensor * a, + enum ggml_prec prec); + + // set the smallest rank that the implementation can use to internally convert the src[idx] data to + // ranks in decreasing order: + // - GGML_PREC_F32 - GGML_TYPE_F32 + // - GGML_PREC_BF16 - GGML_TYPE_BF16 + // - GGML_PREC_F16 - GGML_TYPE_F16, + // - GGML_PREC_Q8 - GGML_TYPE_Q8_0, GGML_TYPE_Q8_1, GGML_TYPE_Q8_K, etc. + // - GGML_PREC_Q4 - GGML_TYPE_Q4_0, GGML_TYPE_Q4_1, GGML_TYPE_Q4_K, GGML_TYPE_NVFP4, GGML_TYPE_MXFP4, etc. + // + // for example: + // - ggml_prec_set_src(a, GGML_PREC_Q8, 1): + // - allows the implementation to quantize F32, BF16, F16 data of src[1] down to GGML_TYPE_Q8_0 + // - cannot quantize it down to GGML_TYPE_Q4_0 or GGML_TYPE_NVFP4 + // - ggml_prec_set_src(a, GGML_PREC_Q4, 1): + // - allows the implementation to quantize F32, BF16, F16 data of src[1] down to GGML_TYPE_Q8_0 or GGML_TYPE_NVFP4 + // + GGML_API void ggml_prec_set_src( + struct ggml_tensor * a, + enum ggml_prec prec, + int idx); + // A: k columns, n rows => [ne03, ne02, n, k] // B: k columns, m rows (i.e. we transpose it internally) => [ne03 * x, ne02 * y, m, k] // result is n columns, m rows => [ne03 * x, ne02 * y, m, n] @@ -1432,9 +1474,10 @@ extern "C" { // change the precision of a matrix multiplication // set to GGML_PREC_F32 for higher precision (useful for phi-2) - GGML_API void ggml_mul_mat_set_prec( + GGML_DEPRECATED(GGML_API void ggml_mul_mat_set_prec( struct ggml_tensor * a, - enum ggml_prec prec); + enum ggml_prec prec), + "use ggml_prec_set_acc() instead"); // change the hint of a matrix multiplication GGML_API void ggml_mul_mat_set_hint( @@ -2439,9 +2482,10 @@ extern "C" { float max_bias, float logit_softcap); - GGML_API void ggml_flash_attn_ext_set_prec( + GGML_DEPRECATED(GGML_API void ggml_flash_attn_ext_set_prec( struct ggml_tensor * a, - enum ggml_prec prec); + enum ggml_prec prec), + "use ggml_prec_set_acc() instead"); GGML_API enum ggml_prec ggml_flash_attn_ext_get_prec( const struct ggml_tensor * a); diff --git a/ggml/src/ggml-cuda/CMakeLists.txt b/ggml/src/ggml-cuda/CMakeLists.txt index d3953eee962..f690d6938ab 100644 --- a/ggml/src/ggml-cuda/CMakeLists.txt +++ b/ggml/src/ggml-cuda/CMakeLists.txt @@ -131,6 +131,10 @@ if (CUDAToolkit_FOUND) add_compile_definitions(GGML_CUDA_PEER_MAX_BATCH_SIZE=${GGML_CUDA_PEER_MAX_BATCH_SIZE}) + if (CMAKE_CUDA_ARCHITECTURES MATCHES "(^|;)12[0-9]a(-real|-virtual)?($|;)") + add_compile_definitions(GGML_CUDA_HAS_BLACKWELL_TARGET) + endif() + if (GGML_CUDA_GRAPHS) add_compile_definitions(GGML_CUDA_USE_GRAPHS) endif() diff --git a/ggml/src/ggml-cuda/mmq-load-tiles.cuh b/ggml/src/ggml-cuda/mmq-load-tiles.cuh index 8ed704c281a..8ae0eb0bb62 100644 --- a/ggml/src/ggml-cuda/mmq-load-tiles.cuh +++ b/ggml/src/ggml-cuda/mmq-load-tiles.cuh @@ -1662,12 +1662,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, force_w4a8) / warp_size; + constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback, force_w4a8); + constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback, force_w4a8); #if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) int * x_qs = (int *) x_tile; diff --git a/ggml/src/ggml-cuda/mmq-vec-dot.cuh b/ggml/src/ggml-cuda/mmq-vec-dot.cuh index d573433865f..17f26f81e1c 100644 --- a/ggml/src/ggml-cuda/mmq-vec-dot.cuh +++ b/ggml/src/ggml-cuda/mmq-vec-dot.cuh @@ -478,7 +478,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(); @@ -486,8 +486,8 @@ 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 I = ggml_cuda_mmq_get_I(type, J, fallback); - constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback); + constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback, force_w4a8); + constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback, force_w4a8); 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. @@ -537,8 +537,8 @@ template static __device__ __forceinline_ typedef tile< 8, 4, int> tile_B; typedef tile<16, 8, int> tile_C; - 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 I = ggml_cuda_mmq_get_I(type, J, fallback, force_w4a8); + constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback, force_w4a8); 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. diff --git a/ggml/src/ggml-cuda/mmq.cu b/ggml/src/ggml-cuda/mmq.cu index 707437ea3e5..48d06567b43 100644 --- a/ggml/src/ggml-cuda/mmq.cu +++ b/ggml/src/ggml-cuda/mmq.cu @@ -5,7 +5,8 @@ #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, + [[maybe_unused]] bool force_w4a8 = false) { switch (args.type_x) { case GGML_TYPE_Q1_0: mul_mat_q_case(ctx, args, stream); @@ -74,6 +75,13 @@ static void ggml_cuda_mul_mat_q_switch_type(ggml_backend_cuda_context & ctx, con mul_mat_q_case(ctx, args, stream); break; case GGML_TYPE_NVFP4: +#ifdef GGML_CUDA_HAS_BLACKWELL_TARGET + // W4A16 NVFP4: dispatch the W4A8 instantiation so activations stay at higher precision even on Blackwell. + if (force_w4a8) { + mul_mat_q_case(ctx, args, stream); + break; + } +#endif // GGML_CUDA_HAS_BLACKWELL_TARGET mul_mat_q_case(ctx, args, stream); break; default: @@ -82,6 +90,19 @@ static void ggml_cuda_mul_mat_q_switch_type(ggml_backend_cuda_context & ctx, con } } +// NVFP4 defaults to native W4A4 on Blackwell. A src1 precision of GGML_PREC_Q8 selects the W4A8 +// path, unless GGML_CUDA_FORCE_W4A4 overrides it. +static inline bool ggml_cuda_mmq_force_w4a8(const ggml_tensor * src0, const ggml_tensor * dst) { + static const bool force_w4a4 = []() { + const char * env = getenv("GGML_CUDA_FORCE_W4A4"); + return env != nullptr && std::atoi(env) != 0; + }(); + if (force_w4a4 || src0->type != GGML_TYPE_NVFP4) { + return false; + } + return ggml_get_op_params_i32(dst, 3) == GGML_PREC_Q8; +} + 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 +149,10 @@ 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 bool force_w4a8 = ggml_cuda_mmq_force_w4a8(src0, dst); + + const bool use_native_fp4 = !force_w4a8 && blackwell_mma_available(cc) && + (src0->type == GGML_TYPE_MXFP4 || src0->type == GGML_TYPE_NVFP4); 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 +196,7 @@ void ggml_cuda_mul_mat_q( ne02, ne12, s02, s12, s2, ne03, ne13, s03, s13, s3, ne1}; - ggml_cuda_mul_mat_q_switch_type(ctx, args, stream); + ggml_cuda_mul_mat_q_switch_type(ctx, args, stream, force_w4a8); return; } @@ -253,7 +277,7 @@ void ggml_cuda_mul_mat_q( ne03, ne13, s03, s13, s3, ne12}; - ggml_cuda_mul_mat_q_switch_type(ctx, args, stream); + ggml_cuda_mul_mat_q_switch_type(ctx, args, stream, force_w4a8); } 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 2eb15fdfad9..c04ddc790b8 100644 --- a/ggml/src/ggml-cuda/mmq.cuh +++ b/ggml/src/ggml-cuda/mmq.cuh @@ -225,7 +225,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 bool force_w4a8 = false) { if (GGML_CUDA_CC_IS_AMD(cc)) { if (GGML_CUDA_CC_IS_CDNA(cc)) { return ggml_cuda_mmq_get_config_cdna(type, J, fallback); @@ -242,6 +242,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)) { + // W4A16 NVFP4: keep src1 at Q8_1 (W4A8) by using the generic NVFP4 config even on Blackwell. + if (force_w4a8 && type == GGML_TYPE_NVFP4) { + 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) { @@ -250,7 +254,7 @@ static __host__ ggml_cuda_mmq_config ggml_cuda_mmq_get_config(const ggml_type ty return ggml_cuda_mmq_get_config_pascal(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, bool force_w4a8 = false) { #ifdef GGML_USE_HIP #ifdef CDNA return ggml_cuda_mmq_get_config_cdna(type, J, fallback); @@ -265,6 +269,10 @@ static constexpr __device__ ggml_cuda_mmq_config ggml_cuda_mmq_get_config(ggml_t #endif // CDNA #else #ifdef BLACKWELL_MMA_AVAILABLE + // W4A16 NVFP4: keep src1 at Q8_1 (W4A8) by using the generic NVFP4 config even on Blackwell. + if (force_w4a8 && type == GGML_TYPE_NVFP4) { + 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); @@ -272,79 +280,79 @@ static constexpr __device__ ggml_cuda_mmq_config ggml_cuda_mmq_get_config(ggml_t return ggml_cuda_mmq_get_config_pascal(type, J, fallback); #endif // BLACKWELL_MMA_AVAILABLE #endif // GGML_USE_HIP - GGML_UNUSED_VARS(type, J, fallback); + GGML_UNUSED_VARS(type, J, fallback, force_w4a8); } 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 constexpr __device__ int ggml_cuda_mmq_get_type(ggml_type type, int J, bool fallback, bool force_w4a8 = false) { + return ggml_cuda_mmq_get_config(type, J, fallback, force_w4a8).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_nthreads(ggml_type type, int J, bool fallback, bool force_w4a8 = false) { + return ggml_cuda_mmq_get_config(type, J, fallback, force_w4a8).nthreads; } 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_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, bool force_w4a8 = false) { + return ggml_cuda_mmq_get_config(type, J, fallback, force_w4a8).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, bool force_w4a8 = false) { + return ggml_cuda_mmq_get_config(type, J, fallback, force_w4a8).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, bool force_w4a8 = false) { + return ggml_cuda_mmq_get_config(type, J, fallback, force_w4a8).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, bool force_w4a8 = false) { + return ggml_cuda_mmq_get_config(type, J, fallback, force_w4a8).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, bool force_w4a8 = false) { + return ggml_cuda_mmq_get_config(type, J, fallback, force_w4a8).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, bool force_w4a8 = false) { + return ggml_cuda_mmq_get_config(type, J, fallback, force_w4a8).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, bool force_w4a8 = false) { + return ggml_cuda_mmq_get_config(type, J, fallback, force_w4a8).fallback; } // --------------------------------------------------------------------------------------------- @@ -353,8 +361,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, bool force_w4a8 = false) { + return ggml_cuda_mmq_get_sram_stride(ggml_cuda_mmq_get_sram_layout(type, J, fallback, force_w4a8)); } static __host__ int ggml_cuda_mmq_get_J_max(const ggml_type type, const bool fallback, const int cc, const int64_t ne11) { @@ -532,11 +540,11 @@ 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() { - constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback); + constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback, force_w4a8); - if (!ggml_cuda_mmq_get_config(type, J, fallback).use_mma_data_layout()) { + if (!ggml_cuda_mmq_get_config(type, J, fallback, force_w4a8).use_mma_data_layout()) { switch (type) { case GGML_TYPE_Q1_0: return ggml_cuda_mmq_util_funcs( @@ -689,11 +697,14 @@ static constexpr __device__ ggml_cuda_mmq_util_funcs ggml_cuda_mmq_get_util_func ggml_cuda_mmq_vec_dot_fp4_fp4_mma, ggml_cuda_mmq_write_back_mma); 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 (!force_w4a8) { + 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; } @@ -834,37 +845,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, @@ -873,25 +884,26 @@ 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, force_w4a8) / 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, force_w4a8); + 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. The NVFP4 W4A8 (force_w4a8) path 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) && !(type == GGML_TYPE_NVFP4 && force_w4a8)) ? 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, force_w4a8); constexpr int blocks_per_iter = ITER_K / qk; float sum[J*I / (nwarps*warp_size)] = {0.0f}; @@ -943,8 +955,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, force_w4a8), ggml_cuda_mmq_get_occupancy(type, J, fallback, force_w4a8)) 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, @@ -955,15 +967,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, force_w4a8).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, force_w4a8) / 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, force_w4a8); const uint32_t nty = (nrows_x + I - 1) / I; // Number of tiles y @@ -983,7 +995,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, force_w4a8)) { const uint2 tmp2 = fast_div_modulo(blockIdx.z, nchannels_y); const int wt = tmp2.x; const int zt = tmp2.y; @@ -1046,14 +1058,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, force_w4a8); constexpr int blocks_per_iter = ITER_K / qk; // kbc == k block continuous, current index in continuous ijk space. @@ -1140,7 +1152,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); @@ -1224,24 +1236,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, force_w4a8)/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, force_w4a8) / 2) / warp_size; + constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback, force_w4a8); 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, force_w4a8); constexpr int blocks_per_iter = ITER_K / qk; float sum[J / nwarps] = {0.0f}; @@ -1384,22 +1396,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, force_w4a8); 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; @@ -1419,7 +1431,7 @@ static void launch_mul_mat_q(ggml_backend_cuda_context & ctx, const mmq_args & a const uint3 sample_ratio_fd = init_fastdiv_values(sample_ratio); if (!ggml_cuda_mmq_get_stream_k(type, J, fallback, cc)) { - mul_mat_q<<>> + 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, @@ -1448,7 +1460,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, @@ -1460,13 +1472,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; @@ -1476,7 +1488,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, force_w4a8); if (config.type == GGML_TYPE_COUNT) { continue; } @@ -1495,52 +1507,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); @@ -1549,20 +1561,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) \ +// W4A16 NVFP4 variant: keeps src1 at Q8_1 (W4A8) instead of native FP4 MMA on Blackwell. +#define DECL_MMQ_CASE_W4A8(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); @@ -1588,6 +1604,9 @@ extern DECL_MMQ_CASE(GGML_TYPE_IQ4_XS); // ----------------------------------------- extern DECL_MMQ_CASE(GGML_TYPE_MXFP4); extern DECL_MMQ_CASE(GGML_TYPE_NVFP4); +#ifdef GGML_CUDA_HAS_BLACKWELL_TARGET +extern DECL_MMQ_CASE_W4A8(GGML_TYPE_NVFP4); // W4A8 path only differs on Blackwell +#endif // GGML_CUDA_HAS_BLACKWELL_TARGET // ------------------------------------------------------------------------------------------------------------------------- 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 d7cd271675e..673448f0c17 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,10 @@ DECL_MMQ_CASE({type}); """ +SOURCE_MMQ_EXTRA = { + "GGML_TYPE_NVFP4": "#ifdef GGML_CUDA_HAS_BLACKWELL_TARGET\nDECL_MMQ_CASE_W4A8(GGML_TYPE_NVFP4);\n#endif // GGML_CUDA_HAS_BLACKWELL_TARGET\n", +} + SOURCE_MMF = """// This file has been autogenerated by generate_cu_files.py, do not edit manually. #include "../mmf.cuh" @@ -105,6 +109,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 SOURCE_MMQ_EXTRA: + f.write(SOURCE_MMQ_EXTRA[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-nvfp4.cu b/ggml/src/ggml-cuda/template-instances/mmq-instance-nvfp4.cu index 2cb140d35a3..78d59c01ba1 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,6 @@ #include "../mmq.cuh" DECL_MMQ_CASE(GGML_TYPE_NVFP4); +#ifdef GGML_CUDA_HAS_BLACKWELL_TARGET +DECL_MMQ_CASE_W4A8(GGML_TYPE_NVFP4); +#endif // GGML_CUDA_HAS_BLACKWELL_TARGET diff --git a/ggml/src/ggml-impl.h b/ggml/src/ggml-impl.h index 62b76abbcec..f63529b90e6 100644 --- a/ggml/src/ggml-impl.h +++ b/ggml/src/ggml-impl.h @@ -160,6 +160,17 @@ static float ggml_get_op_params_f32(const struct ggml_tensor * tensor, uint32_t return ((const float *)(tensor->op_params))[i]; } +// - GGML_OP_MUL_MAT +// 0 - acc +// 1 - hint +// 2 - src0 precision +// 3 - src1 precision +// +// - GGML_OP_MUL_MAT_ID +// 0 - acc +// 1 - hint +// 2 - src0 precision +// 3 - src1 precision static void ggml_set_op_params_i32(struct ggml_tensor * tensor, uint32_t i, int32_t value) { assert(i < GGML_MAX_OP_PARAMS / sizeof(int32_t)); ((int32_t *)(tensor->op_params))[i] = value; diff --git a/ggml/src/ggml.c b/ggml/src/ggml.c index e0b615c07ed..0bc01c1f655 100644 --- a/ggml/src/ggml.c +++ b/ggml/src/ggml.c @@ -3265,6 +3265,54 @@ struct ggml_tensor * ggml_l2_norm_inplace( return ggml_l2_norm_impl(ctx, a, eps, true); } +// ggml_prec + +void ggml_prec_set_acc( + struct ggml_tensor * a, + enum ggml_prec prec) { + switch (a->op) { + case GGML_OP_MUL_MAT: + case GGML_OP_MUL_MAT_ID: + { + const int32_t prec_i32 = (int32_t) prec; + + ggml_set_op_params_i32(a, 0, prec_i32); + } + break; + case GGML_OP_FLASH_ATTN_EXT: + { + const int32_t prec_i32 = (int32_t) prec; + + ggml_set_op_params_i32(a, 3, prec_i32); + } + break; + default: + GGML_ABORT("not implemented"); + }; +} + +void ggml_prec_set_src( + struct ggml_tensor * a, + enum ggml_prec prec, + int idx) { + GGML_ASSERT(idx >= 0 && idx < GGML_MAX_SRC); + + switch (a->op) { + case GGML_OP_MUL_MAT: + case GGML_OP_MUL_MAT_ID: + { + GGML_ASSERT(idx == 1); + + const int32_t prec_i32 = (int32_t) prec; + + ggml_set_op_params_i32(a, 2 + idx, prec_i32); + } + break; + default: + GGML_ABORT("not implemented"); + }; +} + // ggml_mul_mat static inline bool ggml_can_mul_mat(const struct ggml_tensor * t0, const struct ggml_tensor * t1) { diff --git a/gguf-py/gguf/constants.py b/gguf-py/gguf/constants.py index f236a5d2c98..2f78d3a02f0 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. allow_prec_a8). + TENSOR_EXTRA_NAME = "general.tensor_extra.name" + TENSOR_EXTRA_ALLOW_PREC_A8 = "general.tensor_extra.allow_prec_a8" + # 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 d8a96a27bdd..3cf94d4523c 100644 --- a/gguf-py/gguf/gguf_writer.py +++ b/gguf-py/gguf/gguf_writer.py @@ -515,6 +515,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_allow_prec_a8(self, tensor_names: Sequence[str], values: Sequence[bool]) -> None: + if len(tensor_names) != len(values): + raise ValueError("tensor_extra allow_prec_a8 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_ALLOW_PREC_A8, 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 eecf444fcf3..de201f4cd94 100644 --- a/src/llama-arch.cpp +++ b/src/llama-arch.cpp @@ -157,31 +157,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_ALLOW_PREC_A8, "general.tensor_extra.allow_prec_a8" }, { 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 7159e23bf7a..84254d2927b 100644 --- a/src/llama-arch.h +++ b/src/llama-arch.h @@ -187,6 +187,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_ALLOW_PREC_A8, LLM_KV_VOCAB_SIZE, LLM_KV_CONTEXT_LENGTH, diff --git a/src/llama-context.cpp b/src/llama-context.cpp index 0402044da6b..91d234f9fc4 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -2469,6 +2469,7 @@ llm_graph_params llama_context::graph_params( /*.n_outputs =*/ n_outputs, /*.cb =*/ graph_get_cb(), /*.res =*/ res, + /*.act_policy =*/ &model.act_policy, }; } diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp index 8fca8e1bc0e..6f500ed1263 100644 --- a/src/llama-graph.cpp +++ b/src/llama-graph.cpp @@ -18,6 +18,7 @@ #include #include +#include #include #include #include @@ -1489,6 +1490,7 @@ llm_graph_context::llm_graph_context(const llm_graph_params & params) : samplers (params.samplers), cb_func (params.cb), res (params.res), + act_policy (params.act_policy), ctx0 (res->get_ctx()), gf (res->get_gf()) { res->set_params(params); @@ -1514,6 +1516,10 @@ ggml_tensor * llm_graph_context::build_lora_mm( ggml_tensor * w_s) const { ggml_tensor * res = ggml_mul_mat(ctx0, w, cur); + if (llama_act_policy_prec_a8(act_policy, w)) { + ggml_prec_set_src(res, GGML_PREC_Q8, 1); + } + if (w_s) { res = ggml_mul(ctx0, res, w_s); } @@ -1546,6 +1552,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 (llama_act_policy_prec_a8(act_policy, w)) { + ggml_prec_set_src(res, GGML_PREC_Q8, 1); + } + if (w_s) { const int64_t n_expert = w_s->ne[0]; const int64_t n_tokens = cur->ne[2]; @@ -1554,6 +1564,7 @@ ggml_tensor * llm_graph_context::build_lora_mm_id( s = ggml_get_rows(ctx0, s, ids); res = ggml_mul(ctx0, res, s); } + for (const auto & lora : *loras) { llama_adapter_lora_weight * lw = lora.first->get_weight(w); if (lw == nullptr) { @@ -1874,7 +1885,7 @@ ggml_tensor * llm_graph_context::build_ffn( cur = build_lora_mm(down, cur); if (arch == LLM_ARCH_GLM4 || arch == LLM_ARCH_GLM4_MOE || arch == LLM_ARCH_JAIS2) { // GLM4, GLM4_MOE, and JAIS2 seem to have numerical issues with half-precision accumulators - ggml_mul_mat_set_prec(cur, GGML_PREC_F32); + ggml_prec_set_acc(cur, GGML_PREC_F32); } } @@ -1972,7 +1983,7 @@ ggml_tensor * llm_graph_context::build_moe_ffn( if (probs_in == nullptr) { logits = build_lora_mm(gate_inp, cur); // [n_expert, n_tokens] if (gating_op == LLAMA_EXPERT_GATING_FUNC_TYPE_SQRT_SOFTPLUS) { - ggml_mul_mat_set_prec(logits, GGML_PREC_F32); + ggml_prec_set_acc(logits, GGML_PREC_F32); } cb(logits, "ffn_moe_logits", il); } else { @@ -2583,7 +2594,7 @@ ggml_tensor * llm_graph_context::build_attn_mha( res->add_fused_node({LLM_FUSED_OP_FLASH_ATTN, cur, il}); ggml_flash_attn_ext_add_sinks(cur, sinks); - ggml_flash_attn_ext_set_prec (cur, GGML_PREC_F32); + ggml_prec_set_acc(cur, GGML_PREC_F32); if (v_mla) { #if 0 @@ -2609,7 +2620,7 @@ ggml_tensor * llm_graph_context::build_attn_mha( // note: this op tends to require high floating point range // while for some models F16 is enough, for others it is not, so we default to F32 here - ggml_mul_mat_set_prec(kq, GGML_PREC_F32); + ggml_prec_set_acc(kq, GGML_PREC_F32); if (arch == LLM_ARCH_GROK) { // need to do the following: @@ -2842,7 +2853,7 @@ ggml_tensor * llm_graph_context::build_attn( if (arch == LLM_ARCH_GLM4 || arch == LLM_ARCH_GLM4_MOE || arch == LLM_ARCH_JAIS2) { // GLM4, GLM4_MOE, and JAIS2 seem to have numerical issues with half-precision accumulators cur = build_lora_mm(wo, cur); - ggml_mul_mat_set_prec(cur, GGML_PREC_F32); + ggml_prec_set_acc(cur, GGML_PREC_F32); if (wo_s) { cur = ggml_mul(ctx0, cur, wo_s); } @@ -2929,7 +2940,7 @@ ggml_tensor * llm_graph_context::build_attn( if (arch == LLM_ARCH_GLM4 || arch == LLM_ARCH_GLM4_MOE) { // GLM4 and GLM4_MOE seem to have numerical issues with half-precision accumulators cur = build_lora_mm(wo, cur); - ggml_mul_mat_set_prec(cur, GGML_PREC_F32); + ggml_prec_set_acc(cur, GGML_PREC_F32); if (wo_s) { cur = ggml_mul(ctx0, cur, wo_s); } diff --git a/src/llama-graph.h b/src/llama-graph.h index b388e028cb5..77c23264320 100644 --- a/src/llama-graph.h +++ b/src/llama-graph.h @@ -18,6 +18,7 @@ struct ggml_tensor; struct llama_cparams; struct llama_layer; +struct llama_act_policy; struct llama_memory_context_i; @@ -809,6 +810,8 @@ struct llm_graph_params { llm_graph_result * res; + const llama_act_policy * act_policy = nullptr; + // return true if the "other" params would result in a graph with the same topology as with the current params // having the same topology allows us to reuse the graph in some cases bool allow_reuse(const llm_graph_params & other) const { @@ -1032,6 +1035,8 @@ struct llm_graph_context { llm_graph_result * res; + const llama_act_policy * act_policy; + ggml_context * ctx0 = nullptr; ggml_cgraph * gf = nullptr; diff --git a/src/llama-model-loader.cpp b/src/llama-model-loader.cpp index 9b22cb05f29..1536d5c2e49 100644 --- a/src/llama-model-loader.cpp +++ b/src/llama-model-loader.cpp @@ -410,6 +410,7 @@ namespace GGUFMeta { template bool llama_model_loader::get_arr>(enum llm_kv kid, std::array & result, bool required); template bool llama_model_loader::get_arr>(enum llm_kv kid, std::vector & result, bool required); template bool llama_model_loader::get_arr>(enum llm_kv kid, std::array & result, bool required); + template bool llama_model_loader::get_arr>(enum llm_kv kid, std::vector & result, bool required); template bool llama_model_loader::get_key(const std::string & key, T & result, bool required) { diff --git a/src/llama-model-saver.cpp b/src/llama-model-saver.cpp index 9adaa93f62e..e854c957d35 100644 --- a/src/llama-model-saver.cpp +++ b/src/llama-model-saver.cpp @@ -198,6 +198,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->act_policy.per_tensor.empty()) { + std::vector tensor_names; + std::vector values; + tensor_names.reserve(model->act_policy.per_tensor.size()); + values.reserve(model->act_policy.per_tensor.size()); + for (const auto & [name, prec_a8] : model->act_policy.per_tensor) { + tensor_names.push_back(name); + values.push_back(prec_a8 ? 1 : 0); + } + add_kv(LLM_KV_GENERAL_TENSOR_EXTRA_NAME, tensor_names); + gguf_set_arr_data(gguf_ctx, llm_kv(LLM_KV_GENERAL_TENSOR_EXTRA_ALLOW_PREC_A8).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 c34700ff563..bd5c0ce5b31 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -1154,6 +1154,60 @@ struct llama_model::impl { std::vector tensor_split_owned; }; +bool llama_act_policy::wants_prec_a8(const ggml_tensor * w) const { + if (!w) { + return false; + } + + const auto it = per_tensor.find(w->name); + if (it != per_tensor.end()) { + return it->second; + } + + return false; +} + +static bool load_act_policy_arr( + llama_model_loader & ml, + const std::string & key_tensor, + const std::string & key_value, + llama_act_policy & policy) { + std::vector tensor_names; + if (!ml.get_arr(key_tensor, tensor_names, false)) { + return false; + } + + const gguf_context * ctx = ml.metadata; + const int kid = gguf_find_key(ctx, key_value.c_str()); + if (kid < 0 || gguf_get_kv_type(ctx, kid) != GGUF_TYPE_ARRAY) { + throw std::runtime_error(format("missing %s array", key_value.c_str())); + } + if (gguf_get_arr_type(ctx, kid) != GGUF_TYPE_BOOL) { + throw std::runtime_error(format("%s must be a bool array", key_value.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_tensor.c_str(), tensor_names.size(), n_values)); + } + + const int8_t * values = (const int8_t *) gguf_get_arr_data(ctx, kid); + for (size_t i = 0; i < n_values; ++i) { + policy.per_tensor.emplace(tensor_names[i], values[i] != 0); + } + + return true; +} + +static void load_act_policy(llama_model_loader & ml, llama_act_policy & policy) { + load_act_policy_arr(ml, + ml.llm_kv(LLM_KV_GENERAL_TENSOR_EXTRA_NAME), + ml.llm_kv(LLM_KV_GENERAL_TENSOR_EXTRA_ALLOW_PREC_A8), + policy); +} + 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 @@ -1205,6 +1259,10 @@ void llama_model_base::load_hparams(llama_model_loader & ml) { ml.get_key(LLM_KV_POOLING_TYPE, hparams.pooling_type, false); ml.get_key(LLM_KV_BLOCK_COUNT, hparams.n_layer_all); GGML_ASSERT(hparams.n_layer_all > 0 && hparams.n_layer_all <= LLAMA_MAX_LAYERS); + + // per-tensor activation precision policy + load_act_policy(ml, act_policy); + ml.get_key(LLM_KV_EXPERT_COUNT, hparams.n_expert, false); ml.get_key(LLM_KV_EXPERT_USED_COUNT, hparams.n_expert_used, false); ml.get_key(LLM_KV_EXPERT_GROUP_COUNT, hparams.n_expert_groups, false); diff --git a/src/llama-model.h b/src/llama-model.h index 44bd9675754..7a22600cdec 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 ggml_tensor; // available models enum llm_type { @@ -582,6 +583,17 @@ 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); +// Per-tensor activation precision from GGUF, unlisted tensors default to native (4-bit) activations. +struct llama_act_policy { + std::unordered_map per_tensor; + + bool wants_prec_a8(const ggml_tensor * w) const; +}; + +inline bool llama_act_policy_prec_a8(const llama_act_policy * policy, const ggml_tensor * w) { + return policy && policy->wants_prec_a8(w); +} + struct llama_model { llm_type type = LLM_TYPE_UNKNOWN; llm_arch arch = LLM_ARCH_UNKNOWN; @@ -591,6 +603,9 @@ struct llama_model { llama_hparams hparams = {}; llama_vocab vocab; + // per-tensor activation precision policy + llama_act_policy act_policy; + // for classifier models std::vector classifier_labels; diff --git a/src/models/minimax-m3.cpp b/src/models/minimax-m3.cpp index 1ba699d0166..8b9543cb2f4 100644 --- a/src/models/minimax-m3.cpp +++ b/src/models/minimax-m3.cpp @@ -191,7 +191,7 @@ ggml_tensor * llama_model_minimax_m3::graph::build_attn_msa_fa( ggml_tensor * o = ggml_flash_attn_ext(ctx0, q, k, v, mask, kq_scale, hparams.f_max_alibi_bias, 0.0f); - ggml_flash_attn_ext_set_prec(o, GGML_PREC_F32); + ggml_prec_set_acc(o, GGML_PREC_F32); cb(o, "msa_fattn", il); // [D, Gp, R, C] -> [D, Gp, C, R] -> [n_embd, T] @@ -389,7 +389,7 @@ llama_model_minimax_m3::graph::graph(const llama_model & model, const llm_graph_ ggml_tensor * iq4 = ggml_reshape_4d(ctx0, iq, n_idx_dim, Hd, 1, ns); ggml_tensor * sc = ggml_mul_mat(ctx0, ggml_reshape_4d(ctx0, ikp, n_idx_dim, n_ps, 1, ns), iq4); - ggml_mul_mat_set_prec(sc, GGML_PREC_F32); + ggml_prec_set_acc(sc, GGML_PREC_F32); // unmapped positions come out -inf, so they can never rank into the top-k sc = ggml_add_inplace(ctx0, sc, ggml_reshape_4d(ctx0, msa->pos_mask, n_ps, 1, 1, ns)); @@ -471,7 +471,7 @@ llama_model_minimax_m3::graph::graph(const llama_model & model, const llm_graph_ ggml_tensor * sc = ggml_mul_mat(ctx0, ikp, ggml_reshape_2d(ctx0, iq_s, n_idx_dim, Hd*n_tps)); // indexer scores run in F32 - ggml_mul_mat_set_prec(sc, GGML_PREC_F32); + ggml_prec_set_acc(sc, GGML_PREC_F32); sc = ggml_reshape_3d(ctx0, sc, n_ps, Hd, n_tps); // unmapped positions (holes, padding, empty cells) come out -inf sc = ggml_add_inplace(ctx0, sc, pm_s); diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index 99014df09cb..2d1ed89ef13 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -4466,6 +4466,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; @@ -4490,7 +4511,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(); @@ -4649,6 +4672,41 @@ struct test_mul_mat_hadamard : public test_mul_mat { } }; +// NVFP4 W4A8 path (GGML_PREC_Q8 on src1 disallows 4-bit activations) +struct test_mul_mat_nvfp4_w4a8 : public test_mul_mat { + test_mul_mat_nvfp4_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_NVFP4_W4A8"; + } +}; + +// NVFP4 native W4A4 path (default precision) +struct test_mul_mat_nvfp4_w4a4 : public test_mul_mat { + test_mul_mat_nvfp4_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_NVFP4_W4A4"; + } +}; + static void init_mul_mat_id_tensors(ggml_context * ctx, int n_mats) { std::random_device rd; std::default_random_engine rng(rd()); @@ -4691,7 +4749,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(); @@ -4736,6 +4796,39 @@ struct test_mul_mat_id : public test_case { } }; +// NVFP4 W4A8 path on the MoE path (GGML_PREC_Q8 on src1 disallows 4-bit activations) +struct test_mul_mat_id_nvfp4_w4a8 : public test_mul_mat_id { + test_mul_mat_id_nvfp4_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_NVFP4_W4A8"; + } +}; + +// NVFP4 native W4A4 path on the MoE path (default precision) +struct test_mul_mat_id_nvfp4_w4a4 : public test_mul_mat_id { + test_mul_mat_id_nvfp4_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_NVFP4_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; @@ -7167,7 +7260,7 @@ struct test_flash_attn_ext : public test_case { ggml_tensor * out = ggml_flash_attn_ext(ctx, q, k, v, m, 1.0f/sqrtf(hsk), max_bias, logit_softcap); ggml_flash_attn_ext_add_sinks(out, s); - ggml_flash_attn_ext_set_prec (out, prec); + ggml_prec_set_acc(out, prec); ggml_set_name(out, "out"); return out; @@ -9139,6 +9232,17 @@ static std::vector> make_test_cases_eval() { test_cases.emplace_back(new test_mul_mat_hadamard(GGML_TYPE_F32, GGML_TYPE_F32, 32, 1, 32)); // too small (N<64) test_cases.emplace_back(new test_mul_mat_hadamard(GGML_TYPE_F32, GGML_TYPE_F32, 1024, 1, 1024)); // too big (N>512) + // NVFP4 activation precision (default = native W4A4, src1 GGML_PREC_Q8 = W4A8) + test_cases.emplace_back(new test_mul_mat_nvfp4_w4a8(GGML_TYPE_NVFP4, GGML_TYPE_F32, 32, 1, 256)); + test_cases.emplace_back(new test_mul_mat_nvfp4_w4a8(GGML_TYPE_NVFP4, GGML_TYPE_F32, 32, 32, 256)); + test_cases.emplace_back(new test_mul_mat_nvfp4_w4a8(GGML_TYPE_NVFP4, GGML_TYPE_F32, 64, 16, 512)); + test_cases.emplace_back(new test_mul_mat_nvfp4_w4a4(GGML_TYPE_NVFP4, GGML_TYPE_F32, 32, 1, 256)); + test_cases.emplace_back(new test_mul_mat_nvfp4_w4a4(GGML_TYPE_NVFP4, GGML_TYPE_F32, 32, 32, 256)); + test_cases.emplace_back(new test_mul_mat_id_nvfp4_w4a8(GGML_TYPE_NVFP4, GGML_TYPE_F32, 8, 2, false, 32, 32, 256)); + test_cases.emplace_back(new test_mul_mat_id_nvfp4_w4a8(GGML_TYPE_NVFP4, GGML_TYPE_F32, 4, 2, true, 64, 16, 256)); + test_cases.emplace_back(new test_mul_mat_id_nvfp4_w4a4(GGML_TYPE_NVFP4, GGML_TYPE_F32, 8, 2, false, 32, 32, 256)); + test_cases.emplace_back(new test_mul_mat_id_nvfp4_w4a4(GGML_TYPE_NVFP4, GGML_TYPE_F32, 4, 2, true, 64, 16, 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})); diff --git a/tools/mtmd/clip.cpp b/tools/mtmd/clip.cpp index 90de1957586..76adae87276 100644 --- a/tools/mtmd/clip.cpp +++ b/tools/mtmd/clip.cpp @@ -780,7 +780,7 @@ ggml_tensor * clip_graph::build_attn( } cur = ggml_flash_attn_ext(ctx0, q, k, v, kq_mask, kq_scale, 0.0f, 0.0f); - ggml_flash_attn_ext_set_prec(cur, GGML_PREC_F32); + ggml_prec_set_acc(cur, GGML_PREC_F32); if (sinks != nullptr) { ggml_flash_attn_ext_add_sinks(cur, sinks); } @@ -793,7 +793,7 @@ ggml_tensor * clip_graph::build_attn( ggml_tensor * kq = ggml_mul_mat(ctx0, k, q); // F32 may not needed for vision encoders? - // ggml_mul_mat_set_prec(kq, GGML_PREC_F32); + // ggml_prec_set_acc(kq, GGML_PREC_F32); kq = ggml_soft_max_ext(ctx0, kq, kq_mask, kq_scale, 0.0f); if (sinks != nullptr) { diff --git a/tools/mtmd/models/mimovl.cpp b/tools/mtmd/models/mimovl.cpp index 6ff1124a02f..e1fbe2671dc 100644 --- a/tools/mtmd/models/mimovl.cpp +++ b/tools/mtmd/models/mimovl.cpp @@ -2,7 +2,7 @@ ggml_tensor * clip_graph_mimovl::build_mm(ggml_tensor * w, ggml_tensor * x) const { ggml_tensor * cur = ggml_mul_mat(ctx0, w, x); - ggml_mul_mat_set_prec(cur, GGML_PREC_F32); + ggml_prec_set_acc(cur, GGML_PREC_F32); return cur; } diff --git a/tools/mtmd/models/qwen3tts-spkenc.cpp b/tools/mtmd/models/qwen3tts-spkenc.cpp index d4659fd63dc..405fbb9cbc2 100644 --- a/tools/mtmd/models/qwen3tts-spkenc.cpp +++ b/tools/mtmd/models/qwen3tts-spkenc.cpp @@ -27,7 +27,7 @@ ggml_tensor * clip_graph_qwen3tts_spkenc::conv1d_same(ggml_tensor * x, ggml_tens ggml_tensor * w2d = ggml_reshape_2d(ctx0, w, (int64_t) K * IC, OC); ggml_tensor * y = ggml_mul_mat(ctx0, w2d, col); // [OC, T_out] - ggml_mul_mat_set_prec(y, GGML_PREC_F32); + ggml_prec_set_acc(y, GGML_PREC_F32); ggml_tensor * b2d = ggml_reshape_2d(ctx0, b, OC, 1); y = ggml_add(ctx0, y, b2d); diff --git a/tools/tuning/fa-vec.cpp b/tools/tuning/fa-vec.cpp index f904379695e..3d6cbeb2c1e 100644 --- a/tools/tuning/fa-vec.cpp +++ b/tools/tuning/fa-vec.cpp @@ -56,7 +56,7 @@ static ggml_tensor * fa_build_graph(ggml_context * ctx, const fa_shape & s) { ggml_set_name(m, "m"); ggml_tensor * out = ggml_flash_attn_ext(ctx, q, k, v, m, 1.0f / sqrtf((float) s.dk), 0.0f, 0.0f); - ggml_flash_attn_ext_set_prec(out, GGML_PREC_F32); + ggml_prec_set_acc(out, GGML_PREC_F32); ggml_set_name(out, "out"); return out;