From 8e99703d75df3bb5319bf24d3bde9718880e684f Mon Sep 17 00:00:00 2001 From: Ruben Ortlam Date: Mon, 14 Sep 2026 13:55:45 +0200 Subject: [PATCH] ggml: add tensor scaling (per-tensor, per-channel, per-expert, per-expert-per-channel) support --- conversion/base.py | 22 +- ggml/include/ggml.h | 17 ++ ggml/src/ggml-cpu/ggml-cpu.c | 66 ++++++ ggml/src/ggml-cuda/CMakeLists.txt | 11 +- ggml/src/ggml-cuda/common.cuh | 92 ++++++++ ggml/src/ggml-cuda/convert.cu | 12 + ggml/src/ggml-cuda/dequantize.cuh | 8 + ggml/src/ggml-cuda/fattn-common.cuh | 51 ++++ ggml/src/ggml-cuda/fattn-vec.cuh | 4 + ggml/src/ggml-cuda/fattn.cu | 27 ++- ggml/src/ggml-cuda/fp8.cu | 217 ++++++++++++++++++ ggml/src/ggml-cuda/fp8.cuh | 9 + ggml/src/ggml-cuda/getrows.cu | 4 + ggml/src/ggml-cuda/ggml-cuda.cu | 191 ++++++++++++--- ggml/src/ggml-cuda/mmvq.cu | 48 ++-- ggml/src/ggml-cuda/set-rows.cu | 16 +- .../fattn-vec-instance-bf16-f8_e4m3.cu | 7 + .../fattn-vec-instance-f16-f8_e4m3.cu | 7 + .../fattn-vec-instance-f8_e4m3-bf16.cu | 7 + .../fattn-vec-instance-f8_e4m3-f16.cu | 7 + .../fattn-vec-instance-f8_e4m3-f8_e4m3.cu | 7 + .../fattn-vec-instance-f8_e4m3-q4_0.cu | 7 + .../fattn-vec-instance-f8_e4m3-q4_1.cu | 7 + .../fattn-vec-instance-f8_e4m3-q5_0.cu | 7 + .../fattn-vec-instance-f8_e4m3-q5_1.cu | 7 + .../fattn-vec-instance-f8_e4m3-q8_0.cu | 7 + .../fattn-vec-instance-q4_0-f8_e4m3.cu | 7 + .../fattn-vec-instance-q4_1-f8_e4m3.cu | 7 + .../fattn-vec-instance-q5_0-f8_e4m3.cu | 7 + .../fattn-vec-instance-q5_1-f8_e4m3.cu | 7 + .../fattn-vec-instance-q8_0-f8_e4m3.cu | 7 + .../template-instances/generate_cu_files.py | 2 +- ggml/src/ggml-cuda/vecdotq.cuh | 19 ++ ggml/src/ggml-cuda/vendors/cuda.h | 1 + ggml/src/ggml.c | 61 +++++ src/llama-graph.cpp | 72 ++---- src/llama-model-loader.cpp | 22 +- src/llama-model.cpp | 60 +++-- src/models/granite-hybrid.cpp | 18 +- src/models/granite.cpp | 18 +- tests/test-backend-ops.cpp | 101 ++++---- 41 files changed, 1072 insertions(+), 202 deletions(-) create mode 100644 ggml/src/ggml-cuda/fp8.cu create mode 100644 ggml/src/ggml-cuda/fp8.cuh create mode 100644 ggml/src/ggml-cuda/template-instances/fattn-vec-instance-bf16-f8_e4m3.cu create mode 100644 ggml/src/ggml-cuda/template-instances/fattn-vec-instance-f16-f8_e4m3.cu create mode 100644 ggml/src/ggml-cuda/template-instances/fattn-vec-instance-f8_e4m3-bf16.cu create mode 100644 ggml/src/ggml-cuda/template-instances/fattn-vec-instance-f8_e4m3-f16.cu create mode 100644 ggml/src/ggml-cuda/template-instances/fattn-vec-instance-f8_e4m3-f8_e4m3.cu create mode 100644 ggml/src/ggml-cuda/template-instances/fattn-vec-instance-f8_e4m3-q4_0.cu create mode 100644 ggml/src/ggml-cuda/template-instances/fattn-vec-instance-f8_e4m3-q4_1.cu create mode 100644 ggml/src/ggml-cuda/template-instances/fattn-vec-instance-f8_e4m3-q5_0.cu create mode 100644 ggml/src/ggml-cuda/template-instances/fattn-vec-instance-f8_e4m3-q5_1.cu create mode 100644 ggml/src/ggml-cuda/template-instances/fattn-vec-instance-f8_e4m3-q8_0.cu create mode 100644 ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q4_0-f8_e4m3.cu create mode 100644 ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q4_1-f8_e4m3.cu create mode 100644 ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q5_0-f8_e4m3.cu create mode 100644 ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q5_1-f8_e4m3.cu create mode 100644 ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q8_0-f8_e4m3.cu diff --git a/conversion/base.py b/conversion/base.py index a88e8cc8de..1ba41efec9 100644 --- a/conversion/base.py +++ b/conversion/base.py @@ -607,8 +607,15 @@ class ModelBase: continue scale = LazyTorchTensor.to_eager(self.model_tensors[scale_name]()).float().flatten() - if scale.numel() != 1: - continue + + expert_match = re.search(r"\.layers\.(\d+)\.mlp\.experts\.(\d+)\.(gate_proj|up_proj|down_proj)\.weight$", weight_name) + # accept per-tensor (scalar) or per-output-channel scales + if scale.numel() != 1 and scale.numel() != weight.shape[0]: + raise ValueError( + f"FP8 weight {weight_name!r} has an unsupported block/group scale " + f"(scale numel {scale.numel()}, weight shape {list(weight.shape)}); " + f"only per-tensor or per-output-channel scales can be preserved. " + f"Re-run with --fp8-as-q8 to dequantize instead.") weight_prefix = weight_name.removesuffix(".weight") # Transformers fine-grained FP8 uses activation_scale while ModelOpt uses input_scale. @@ -629,7 +636,6 @@ class ModelBase: self._fp8_e4m3_preserved.add(weight_name) consumed.append(scale_name) - expert_match = re.search(r"\.layers\.(\d+)\.mlp\.experts\.(\d+)\.(gate_proj|up_proj|down_proj)\.weight$", weight_name) if expert_match: bid = int(expert_match.group(1)) expert_id = int(expert_match.group(2)) @@ -639,7 +645,7 @@ class ModelBase: target_name = new_name.replace(".weight", ".scale") entries = scale_tensors.setdefault(target_name, []) assert isinstance(entries, list) - cast(list[tuple[int, float]], entries).append((expert_id, float(scale[0]))) + cast(list, entries).append((expert_id, scale.numpy().astype(np.float32))) if input_scale is not None: target_name = new_name.replace(".weight", ".input_scale") entries = input_scale_tensors.setdefault(target_name, []) @@ -657,11 +663,15 @@ class ModelBase: for name, values in chain(scale_tensors.items(), input_scale_tensors.items()): if isinstance(values, list): values.sort(key=lambda item: item[0]) - scale = np.array([item[1] for item in values], dtype=np.float32) + arrays = [np.asarray(item[1], dtype=np.float32).reshape(-1) for item in values] + if all(a.size == 1 for a in arrays): + scale = np.array([a[0] for a in arrays], dtype=np.float32) + else: + scale = np.stack(arrays, axis=0) else: scale = values.astype(np.float32) if not np.allclose(scale, 1.0, atol=1e-6): - logger.info(f" + {name} (FP8 scale, shape [{scale.size}])") + logger.info(f" + {name} (FP8 scale, shape {list(scale.shape)})") self.gguf_writer.add_tensor(name, scale) @classmethod diff --git a/ggml/include/ggml.h b/ggml/include/ggml.h index 6e677d60ae..169efe7a36 100644 --- a/ggml/include/ggml.h +++ b/ggml/include/ggml.h @@ -768,6 +768,7 @@ extern "C" { GGML_API size_t ggml_element_size(const struct ggml_tensor * tensor); GGML_API bool ggml_is_quantized(enum ggml_type type); + GGML_API bool ggml_needs_scale_quantized(enum ggml_type type); // TODO: temporary until model loading of ggml examples is refactored GGML_API enum ggml_type ggml_ftype_to_ggml_type(enum ggml_ftype ftype); @@ -1441,6 +1442,13 @@ extern "C" { struct ggml_tensor * a, struct ggml_tensor * b); + GGML_API struct ggml_tensor * ggml_mul_mat_ext( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + struct ggml_tensor * scale_weight, + struct ggml_tensor * scale_activations); + // 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( @@ -1459,6 +1467,14 @@ extern "C" { struct ggml_tensor * b, struct ggml_tensor * ids); + GGML_API struct ggml_tensor * ggml_mul_mat_id_ext( + struct ggml_context * ctx, + struct ggml_tensor * as, + struct ggml_tensor * b, + struct ggml_tensor * ids, + struct ggml_tensor * scale_weight, + struct ggml_tensor * scale_activations); + // A: m columns, n rows, // B: p columns, n rows, // result is m columns, p rows @@ -2924,6 +2940,7 @@ extern "C" { int64_t blck_size_interleave; // interleave elements in blocks size_t type_size; bool is_quantized; + bool needs_scale; // whether the quantization type needs a scale factor for valid dequantization ggml_to_float_t to_float; ggml_from_float_t from_float_ref; }; diff --git a/ggml/src/ggml-cpu/ggml-cpu.c b/ggml/src/ggml-cpu/ggml-cpu.c index 9777d8dca0..cc46434f4c 100644 --- a/ggml/src/ggml-cpu/ggml-cpu.c +++ b/ggml/src/ggml-cpu/ggml-cpu.c @@ -1258,6 +1258,38 @@ static void ggml_compute_forward_mul_mat_one_chunk( } } +// apply the weight scale (src[2]) carried by ggml_mul_mat_ext, broadcast over dst +static void ggml_compute_forward_mul_mat_scale( + const struct ggml_compute_params * params, + struct ggml_tensor * dst) { + const struct ggml_tensor * scale = dst->src[2]; + if (!scale) { + return; + } + + const int ith = params->ith; + const int nth = params->nth; + + const int64_t ne0 = dst->ne[0], ne1 = dst->ne[1], ne2 = dst->ne[2], ne3 = dst->ne[3]; + const size_t nb1 = dst->nb[1], nb2 = dst->nb[2], nb3 = dst->nb[3]; + + ggml_barrier(params->threadpool); + + const float * sd = (const float *) scale->data; + const int64_t sne0 = scale->ne[0], sne1 = scale->ne[1], sne2 = scale->ne[2], sne3 = scale->ne[3]; + const int64_t nr = ne1 * ne2 * ne3; + for (int64_t ir = ith; ir < nr; ir += nth) { + const int64_t i1 = ir % ne1; + const int64_t i2 = (ir / ne1) % ne2; + const int64_t i3 = ir / (ne1 * ne2); + float * dp = (float *) ((char *) dst->data + i1*nb1 + i2*nb2 + i3*nb3); + const int64_t so = (i1 % sne1)*sne0 + (i2 % sne2)*sne0*sne1 + (i3 % sne3)*sne0*sne1*sne2; + for (int64_t i0 = 0; i0 < ne0; i0++) { + dp[i0] *= sd[(i0 % sne0) + so]; + } + } +} + void ggml_compute_forward_mul_mat( const struct ggml_compute_params * params, struct ggml_tensor * dst) { @@ -1321,6 +1353,7 @@ void ggml_compute_forward_mul_mat( src1->type, dst->type)) goto UseGgmlGemm1; + ggml_compute_forward_mul_mat_scale(params, dst); return; } UseGgmlGemm1:; @@ -1374,6 +1407,7 @@ UseGgmlGemm1:; // of src1 from the work buffer if (ggml_cpu_iqp_supports_mul_mat(dst) && !params->use_ref) { ggml_compute_forward_mul_mat_iqp(params, dst); + ggml_compute_forward_mul_mat_scale(params, dst); return; } @@ -1396,6 +1430,7 @@ UseGgmlGemm1:; vec_dot_type, dst->type)) goto UseGgmlGemm2; + ggml_compute_forward_mul_mat_scale(params, dst); return; } UseGgmlGemm2:; @@ -1463,6 +1498,8 @@ UseGgmlGemm2:; current_chunk = atomic_fetch_add_explicit(¶ms->threadpool->current_chunk, 1, memory_order_relaxed); } + + ggml_compute_forward_mul_mat_scale(params, dst); } // ggml_compute_forward_mul_mat_id @@ -1735,6 +1772,35 @@ static void ggml_compute_forward_mul_mat_id( current_chunk = atomic_fetch_add_explicit(current_chunk_ctr, 1, memory_order_relaxed); } } + + // apply the weight scale (src[3]) carried by ggml_mul_mat_id_ext: per-tensor [1], + // per-expert [n_expert] or per-channel-per-expert [n_out, n_expert], indexed via ids + const struct ggml_tensor * scale = dst->src[3]; + if (scale) { + ggml_barrier(params->threadpool); + + const float * sd = (const float *) scale->data; + const bool per_ch_exp = scale->ne[1] == ne02; + const bool per_expert = ggml_nelements(scale) == ne02; + const int64_t nr = ne1 * ne2; + for (int64_t ir = ith; ir < nr; ir += nth) { + const int64_t i1 = ir % ne1; + const int64_t i2 = ir / ne1; + const int32_t expert = *(const int32_t *) ((const char *) ids->data + i1*ids->nb[0] + i2*ids->nb[1]); + float * dp = (float *) ((char *) dst->data + i1*nb1 + i2*nb2); + if (per_ch_exp) { + const float * sp = sd + (int64_t) expert*ne0; + for (int64_t i0 = 0; i0 < ne0; i0++) { + dp[i0] *= sp[i0]; + } + } else { + const float s = sd[per_expert ? expert : 0]; + for (int64_t i0 = 0; i0 < ne0; i0++) { + dp[i0] *= s; + } + } + } + } } ///////////////////////////////// diff --git a/ggml/src/ggml-cuda/CMakeLists.txt b/ggml/src/ggml-cuda/CMakeLists.txt index 10828ad817..37ab8c4b09 100644 --- a/ggml/src/ggml-cuda/CMakeLists.txt +++ b/ggml/src/ggml-cuda/CMakeLists.txt @@ -121,7 +121,12 @@ if (CUDAToolkit_FOUND) template-instances/fattn-vec-instance-f16-f16.cu template-instances/fattn-vec-instance-q4_0-q4_0.cu template-instances/fattn-vec-instance-q8_0-q8_0.cu - template-instances/fattn-vec-instance-bf16-bf16.cu) + template-instances/fattn-vec-instance-bf16-bf16.cu + template-instances/fattn-vec-instance-f16-f8_e4m3.cu + template-instances/fattn-vec-instance-f8_e4m3-f16.cu + template-instances/fattn-vec-instance-bf16-f8_e4m3.cu + template-instances/fattn-vec-instance-f8_e4m3-bf16.cu + template-instances/fattn-vec-instance-f8_e4m3-f8_e4m3.cu) endif() ggml_add_backend_library(ggml-cuda @@ -156,7 +161,7 @@ if (CUDAToolkit_FOUND) if (GGML_STATIC) if (WIN32) # As of 12.3.1 CUDA Toolkit for Windows does not offer a static cublas library - target_link_libraries(ggml-cuda PRIVATE CUDA::cudart_static CUDA::cublas) + target_link_libraries(ggml-cuda PRIVATE CUDA::cudart_static CUDA::cublas CUDA::cublasLt) else () if (GGML_CUDA_CUB_3DOT2) target_link_libraries(ggml-cuda PRIVATE CCCL::CCCL) @@ -171,7 +176,7 @@ if (CUDAToolkit_FOUND) if (GGML_CUDA_CUB_3DOT2) target_link_libraries(ggml-cuda PRIVATE CCCL::CCCL) endif() - target_link_libraries(ggml-cuda PRIVATE CUDA::cudart CUDA::cublas) + target_link_libraries(ggml-cuda PRIVATE CUDA::cudart CUDA::cublas CUDA::cublasLt) endif() if (GGML_CUDA_NO_VMM) diff --git a/ggml/src/ggml-cuda/common.cuh b/ggml/src/ggml-cuda/common.cuh index 9918c03947..22dd4bac37 100644 --- a/ggml/src/ggml-cuda/common.cuh +++ b/ggml/src/ggml-cuda/common.cuh @@ -363,6 +363,11 @@ static bool blackwell_mma_available(const int cc) { ggml_cuda_highest_compiled_arch(cc) < GGML_CUDA_CC_RUBIN; } +static bool fp8_mma_hardware_available(const int cc) { + return GGML_CUDA_CC_IS_NVIDIA(cc) && (cc == GGML_CUDA_CC_ADA_LOVELACE || + (cc >= GGML_CUDA_CC_BLACKWELL && cc < GGML_CUDA_CC_RUBIN)); +} + // Checks whether the tensor's base data pointer and higher-dimensional strides are byte-aligned to `alignment` bytes. static bool ggml_cuda_is_aligned(const ggml_tensor * tensor, const size_t alignment) { GGML_ASSERT(tensor != nullptr); @@ -867,6 +872,73 @@ static __device__ __forceinline__ float ggml_cuda_ue4m3_to_fp32(uint8_t x) { #endif // defined(GGML_USE_HIP) && defined(CDNA3) && defined(FP8_AVAILABLE) && HIP_VERSION >= 60200000 } +static __device__ __forceinline__ float ggml_cuda_f8_e4m3_to_fp32(uint8_t x) { +#if defined(FP8_AVAILABLE) && !defined(GGML_USE_HIP) + __nv_fp8_e4m3 xf; + xf.__x = x; + return static_cast(xf); +#else + const uint8_t magnitude = x & 0x7F; + if (magnitude == 0x7F) { + return NAN; + } + + const int exp = (magnitude >> 3) & 0x0F; + const int man = magnitude & 0x07; + float value; + if (exp == 0) { + value = ldexpf((float) man, -9); + } else { + value = ldexpf(1.0f + (float) man / 8.0f, exp - 7); + } + return x & 0x80 ? -value : value; +#endif // defined(FP8_AVAILABLE) && !defined(GGML_USE_HIP) +} + +#if !defined(FP8_AVAILABLE) || defined(GGML_USE_HIP) +static __device__ __forceinline__ int ggml_cuda_round_to_nearest_even(float x) { + const int value = (int) floorf(x); + const float fraction = x - value; + return fraction > 0.5f || (fraction == 0.5f && (value & 1)) ? value + 1 : value; +} +#endif // !defined(FP8_AVAILABLE) || defined(GGML_USE_HIP) + +static __device__ __forceinline__ uint8_t ggml_cuda_fp32_to_f8_e4m3(float x) { +#if defined(FP8_AVAILABLE) && !defined(GGML_USE_HIP) + // TODO: Check how incoming NaNs are treated (i.e. is sign-bit preserved)? + return __nv_cvt_float_to_fp8(x, __NV_SATFINITE, __NV_E4M3); +#else + const uint8_t sign = signbit(x) ? 0x80 : 0; + x = fabsf(x); + + if (isnan(x)) { + return sign | 0x7F; + } + if (x == 0.0f) { + return sign; + } + if (isinf(x) || x >= 448.0f) { + return sign | 0x7E; + } + if (x < 0.015625f) { + return sign | (uint8_t) ggml_cuda_round_to_nearest_even(x * 512.0f); + } + + int exp; + const float mantissa = frexpf(x, &exp) * 2.0f; + int encoded_exp = exp + 6; + int encoded_man = ggml_cuda_round_to_nearest_even((mantissa - 1.0f) * 8.0f); + if (encoded_man == 8) { + encoded_man = 0; + encoded_exp++; + } + if (encoded_exp > 15 || (encoded_exp == 15 && encoded_man > 6)) { + return sign | 0x7E; + } + return sign | (uint8_t) (encoded_exp << 3) | (uint8_t) encoded_man; +#endif // defined(FP8_AVAILABLE) && !defined(GGML_USE_HIP) +} + static __device__ __forceinline__ uint8_t ggml_cuda_fp32_to_ue4m3(float x) { #if defined(BLACKWELL_MMA_AVAILABLE) // This is used for NVFP4 subblock scale quantizations only if (!(x > 0.0f)) { @@ -1035,6 +1107,13 @@ struct ggml_cuda_type_traits { static constexpr int qi = QI_NVFP4; }; +template<> +struct ggml_cuda_type_traits { + static constexpr int qk = QK8_1; + static constexpr int qr = QR8_1; + static constexpr int qi = QI8_1; +}; + template<> struct ggml_cuda_type_traits { static constexpr int qk = QK_K; @@ -1422,6 +1501,9 @@ struct ggml_backend_cuda_context { cublasHandle_t cublas_handles[GGML_CUDA_MAX_DEVICES][GGML_CUDA_MAX_STREAMS] = {nullptr}; void * cublas_workspaces[GGML_CUDA_MAX_DEVICES][GGML_CUDA_MAX_STREAMS] = {nullptr}; size_t cublas_workspace_sizes[GGML_CUDA_MAX_DEVICES] = {0}; +#if !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) && CUDART_VERSION >= 11080 + cublasLtHandle_t cublaslt_handles[GGML_CUDA_MAX_DEVICES] = {nullptr}; +#endif int curr_stream_no = 0; @@ -1516,6 +1598,16 @@ struct ggml_backend_cuda_context { return cublas_handles[device][curr_stream_no]; } +#if !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) && CUDART_VERSION >= 11080 + cublasLtHandle_t cublaslt_handle() { + if (cublaslt_handles[device] == nullptr) { + ggml_cuda_set_device(device); + CUBLAS_CHECK(cublasLtCreate(&cublaslt_handles[device])); + } + return cublaslt_handles[device]; + } +#endif + // pool std::unique_ptr pools[GGML_CUDA_MAX_DEVICES][GGML_CUDA_MAX_STREAMS]; diff --git a/ggml/src/ggml-cuda/convert.cu b/ggml/src/ggml-cuda/convert.cu index 360c614a44..ec601ff7d2 100644 --- a/ggml/src/ggml-cuda/convert.cu +++ b/ggml/src/ggml-cuda/convert.cu @@ -503,6 +503,8 @@ to_bf16_cuda_t ggml_get_to_bf16_cuda(ggml_type type) { return dequantize_row_mxfp4_cuda; case GGML_TYPE_NVFP4: return dequantize_row_nvfp4_cuda; + case GGML_TYPE_F8_E4M3: + return dequantize_block_cont_cuda<1, 1, dequantize_f8_e4m3>; case GGML_TYPE_F32: return convert_unary_cont_cuda; case GGML_TYPE_F16: @@ -563,6 +565,8 @@ to_fp16_cuda_t ggml_get_to_fp16_cuda(ggml_type type) { return dequantize_row_mxfp4_cuda; case GGML_TYPE_NVFP4: return dequantize_row_nvfp4_cuda; + case GGML_TYPE_F8_E4M3: + return dequantize_block_cont_cuda<1, 1, dequantize_f8_e4m3>; case GGML_TYPE_F32: return convert_unary_cont_cuda; case GGML_TYPE_BF16: @@ -620,6 +624,8 @@ to_fp32_cuda_t ggml_get_to_fp32_cuda(ggml_type type) { return dequantize_row_mxfp4_cuda; case GGML_TYPE_NVFP4: return dequantize_row_nvfp4_cuda; + case GGML_TYPE_F8_E4M3: + return dequantize_block_cont_cuda<1, 1, dequantize_f8_e4m3>; case GGML_TYPE_F16: return convert_unary_cont_cuda; case GGML_TYPE_BF16: @@ -647,6 +653,8 @@ to_fp16_nc_cuda_t ggml_get_to_fp16_nc_cuda(ggml_type type) { return dequantize_block_cuda; case GGML_TYPE_Q8_0: return dequantize_block_cuda; + case GGML_TYPE_F8_E4M3: + return dequantize_block_cuda<1, 1, dequantize_f8_e4m3>; case GGML_TYPE_BF16: return convert_unary_cuda; default: @@ -672,6 +680,8 @@ to_bf16_nc_cuda_t ggml_get_to_bf16_nc_cuda(ggml_type type) { return dequantize_block_cuda; case GGML_TYPE_Q8_0: return dequantize_block_cuda; + case GGML_TYPE_F8_E4M3: + return dequantize_block_cuda<1, 1, dequantize_f8_e4m3>; case GGML_TYPE_F16: return convert_unary_cuda; default: @@ -697,6 +707,8 @@ to_fp32_nc_cuda_t ggml_get_to_fp32_nc_cuda(ggml_type type) { return dequantize_block_cuda; case GGML_TYPE_Q8_0: return dequantize_block_cuda; + case GGML_TYPE_F8_E4M3: + return dequantize_block_cuda<1, 1, dequantize_f8_e4m3>; case GGML_TYPE_BF16: return convert_unary_cuda; default: diff --git a/ggml/src/ggml-cuda/dequantize.cuh b/ggml/src/ggml-cuda/dequantize.cuh index 651524be7b..8ccb2514ea 100644 --- a/ggml/src/ggml-cuda/dequantize.cuh +++ b/ggml/src/ggml-cuda/dequantize.cuh @@ -119,6 +119,14 @@ static __device__ __forceinline__ void dequantize_q8_0(const void * vx, const in v.y *= d; } +static __device__ __forceinline__ void dequantize_f8_e4m3(const void * vx, const int64_t ib, const int iqs, float2 & v) { + const ggml_fp8_e4m3_t * x = (const ggml_fp8_e4m3_t *) vx; + + v.x = ggml_cuda_f8_e4m3_to_fp32(x[ib + 0].bits); + v.y = ggml_cuda_f8_e4m3_to_fp32(x[ib + 1].bits); + GGML_UNUSED(iqs); +} + //================================== k-quants // Each call dequantizes one super-block of QK_K values into y using the diff --git a/ggml/src/ggml-cuda/fattn-common.cuh b/ggml/src/ggml-cuda/fattn-common.cuh index 7442bc22af..64924d5d4e 100644 --- a/ggml/src/ggml-cuda/fattn-common.cuh +++ b/ggml/src/ggml-cuda/fattn-common.cuh @@ -145,6 +145,34 @@ static __device__ __forceinline__ float vec_dot_fattn_vec_KQ_bf16( return sum; } +template +static __device__ __forceinline__ float vec_dot_fattn_vec_KQ_f8_e4m3( + const char * __restrict__ K_c, const void * __restrict__ Q_v, const int * __restrict__ Q_q8, const void * __restrict__ Q_ds_v) { + + const uint8_t * K_f8 = (const uint8_t *) K_c; + GGML_UNUSED(Q_q8); + GGML_UNUSED(Q_ds_v); + + float sum = 0.0f; + +#pragma unroll + for (int k_KQ_0 = 0; k_KQ_0 < D; k_KQ_0 += nthreads*4) { + uint8_t tmp[4]; + ggml_cuda_memcpy_1(tmp, K_f8 + k_KQ_0 + (threadIdx.x % nthreads)*4); +#pragma unroll + for (int k_KQ_1 = 0; k_KQ_1 < 4; ++k_KQ_1) { + const float k = ggml_cuda_f8_e4m3_to_fp32(tmp[k_KQ_1]); +#ifdef V_DOT2_F32_F16_AVAILABLE + sum += k * __half2float(((const half *) Q_v)[k_KQ_0/nthreads + k_KQ_1]); +#else + sum += k * ((const float *) Q_v)[k_KQ_0/nthreads + k_KQ_1]; +#endif // V_DOT2_F32_F16_AVAILABLE + } + } + + return sum; +} + template static __device__ __forceinline__ float vec_dot_fattn_vec_KQ_q4_0( const char * __restrict__ K_c, const void * __restrict__ Q_v, const int * __restrict__ Q_q8, const void * __restrict__ Q_ds_v) { @@ -405,6 +433,25 @@ static __device__ __forceinline__ void dequantize_V_bf16(const void * __restrict } } +template +static __device__ __forceinline__ void dequantize_V_f8_e4m3(const void * __restrict__ vx, void * __restrict__ dst, const int64_t i0) { + static_assert(ne == 2 || ne == 4, "bad ne"); + uint8_t tmp[ne]; + ggml_cuda_memcpy_1(tmp, (const uint8_t *) vx + i0); + +#pragma unroll + for (int l = 0; l < ne; ++l) { + const float value = ggml_cuda_f8_e4m3_to_fp32(tmp[l]); + if constexpr (std::is_same_v) { + ((half *) dst)[l] = __float2half(value); + } else if constexpr (std::is_same_v) { + ((float *) dst)[l] = value; + } else { + static_assert(std::is_same_v, "unsupported type"); + } + } +} + template static __device__ __forceinline__ void dequantize_V_q4_0(const void * __restrict__ vx, void * __restrict__ dst, const int64_t i0) { const block_q4_0 * x = (const block_q4_0 *) vx; @@ -633,6 +680,8 @@ constexpr __device__ vec_dot_KQ_t get_vec_dot_KQ() { return vec_dot_fattn_vec_KQ_q8_0; } else if constexpr (type_K == GGML_TYPE_BF16) { return vec_dot_fattn_vec_KQ_bf16; + } else if constexpr (type_K == GGML_TYPE_F8_E4M3) { + return vec_dot_fattn_vec_KQ_f8_e4m3; } else { static_assert(type_K == -1, "bad type"); return nullptr; @@ -655,6 +704,8 @@ constexpr __device__ dequantize_V_t get_dequantize_V() { return dequantize_V_q8_0; } else if constexpr (type_V == GGML_TYPE_BF16) { return dequantize_V_bf16; + } else if constexpr (type_V == GGML_TYPE_F8_E4M3) { + return dequantize_V_f8_e4m3; } else { static_assert(type_V == -1, "bad type"); return nullptr; diff --git a/ggml/src/ggml-cuda/fattn-vec.cuh b/ggml/src/ggml-cuda/fattn-vec.cuh index 519b36b9ff..34a07ad55a 100644 --- a/ggml/src/ggml-cuda/fattn-vec.cuh +++ b/ggml/src/ggml-cuda/fattn-vec.cuh @@ -585,6 +585,7 @@ void ggml_cuda_flash_attn_ext_vec_case(ggml_backend_cuda_context & ctx, ggml_ten extern DECL_FATTN_VEC_CASE(D, type_K, GGML_TYPE_Q5_1); \ extern DECL_FATTN_VEC_CASE(D, type_K, GGML_TYPE_Q8_0); \ extern DECL_FATTN_VEC_CASE(D, type_K, GGML_TYPE_BF16); \ + extern DECL_FATTN_VEC_CASE(D, type_K, GGML_TYPE_F8_E4M3); \ EXTERN_DECL_FATTN_VEC_CASES( 64, GGML_TYPE_F16) EXTERN_DECL_FATTN_VEC_CASES( 64, GGML_TYPE_Q4_0) @@ -593,6 +594,7 @@ EXTERN_DECL_FATTN_VEC_CASES( 64, GGML_TYPE_Q5_0) EXTERN_DECL_FATTN_VEC_CASES( 64, GGML_TYPE_Q5_1) EXTERN_DECL_FATTN_VEC_CASES( 64, GGML_TYPE_Q8_0) EXTERN_DECL_FATTN_VEC_CASES( 64, GGML_TYPE_BF16) +EXTERN_DECL_FATTN_VEC_CASES( 64, GGML_TYPE_F8_E4M3) EXTERN_DECL_FATTN_VEC_CASES(128, GGML_TYPE_F16) EXTERN_DECL_FATTN_VEC_CASES(128, GGML_TYPE_Q4_0) @@ -601,6 +603,7 @@ EXTERN_DECL_FATTN_VEC_CASES(128, GGML_TYPE_Q5_0) EXTERN_DECL_FATTN_VEC_CASES(128, GGML_TYPE_Q5_1) EXTERN_DECL_FATTN_VEC_CASES(128, GGML_TYPE_Q8_0) EXTERN_DECL_FATTN_VEC_CASES(128, GGML_TYPE_BF16) +EXTERN_DECL_FATTN_VEC_CASES(128, GGML_TYPE_F8_E4M3) EXTERN_DECL_FATTN_VEC_CASES(256, GGML_TYPE_F16) EXTERN_DECL_FATTN_VEC_CASES(256, GGML_TYPE_Q4_0) @@ -609,3 +612,4 @@ EXTERN_DECL_FATTN_VEC_CASES(256, GGML_TYPE_Q5_0) EXTERN_DECL_FATTN_VEC_CASES(256, GGML_TYPE_Q5_1) EXTERN_DECL_FATTN_VEC_CASES(256, GGML_TYPE_Q8_0) EXTERN_DECL_FATTN_VEC_CASES(256, GGML_TYPE_BF16) +EXTERN_DECL_FATTN_VEC_CASES(256, GGML_TYPE_F8_E4M3) diff --git a/ggml/src/ggml-cuda/fattn.cu b/ggml/src/ggml-cuda/fattn.cu index ae217fbd9d..7d34c5cc21 100644 --- a/ggml/src/ggml-cuda/fattn.cu +++ b/ggml/src/ggml-cuda/fattn.cu @@ -402,6 +402,7 @@ static void ggml_cuda_flash_attn_ext_vec(ggml_backend_cuda_context & ctx, ggml_t FATTN_VEC_CASES_ALL_D(GGML_TYPE_Q5_1, GGML_TYPE_F16) FATTN_VEC_CASES_ALL_D(GGML_TYPE_Q8_0, GGML_TYPE_F16) FATTN_VEC_CASES_ALL_D(GGML_TYPE_BF16, GGML_TYPE_F16) + FATTN_VEC_CASES_ALL_D(GGML_TYPE_F8_E4M3, GGML_TYPE_F16) FATTN_VEC_CASES_ALL_D(GGML_TYPE_F16, GGML_TYPE_Q4_0) FATTN_VEC_CASES_ALL_D(GGML_TYPE_Q4_0, GGML_TYPE_Q4_0) @@ -410,6 +411,7 @@ static void ggml_cuda_flash_attn_ext_vec(ggml_backend_cuda_context & ctx, ggml_t FATTN_VEC_CASES_ALL_D(GGML_TYPE_Q5_1, GGML_TYPE_Q4_0) FATTN_VEC_CASES_ALL_D(GGML_TYPE_Q8_0, GGML_TYPE_Q4_0) FATTN_VEC_CASES_ALL_D(GGML_TYPE_BF16, GGML_TYPE_Q4_0) + FATTN_VEC_CASES_ALL_D(GGML_TYPE_F8_E4M3, GGML_TYPE_Q4_0) FATTN_VEC_CASES_ALL_D(GGML_TYPE_F16, GGML_TYPE_Q4_1) FATTN_VEC_CASES_ALL_D(GGML_TYPE_Q4_0, GGML_TYPE_Q4_1) @@ -418,6 +420,7 @@ static void ggml_cuda_flash_attn_ext_vec(ggml_backend_cuda_context & ctx, ggml_t FATTN_VEC_CASES_ALL_D(GGML_TYPE_Q5_1, GGML_TYPE_Q4_1) FATTN_VEC_CASES_ALL_D(GGML_TYPE_Q8_0, GGML_TYPE_Q4_1) FATTN_VEC_CASES_ALL_D(GGML_TYPE_BF16, GGML_TYPE_Q4_1) + FATTN_VEC_CASES_ALL_D(GGML_TYPE_F8_E4M3, GGML_TYPE_Q4_1) FATTN_VEC_CASES_ALL_D(GGML_TYPE_F16, GGML_TYPE_Q5_0) FATTN_VEC_CASES_ALL_D(GGML_TYPE_Q4_0, GGML_TYPE_Q5_0) @@ -426,6 +429,7 @@ static void ggml_cuda_flash_attn_ext_vec(ggml_backend_cuda_context & ctx, ggml_t FATTN_VEC_CASES_ALL_D(GGML_TYPE_Q5_1, GGML_TYPE_Q5_0) FATTN_VEC_CASES_ALL_D(GGML_TYPE_Q8_0, GGML_TYPE_Q5_0) FATTN_VEC_CASES_ALL_D(GGML_TYPE_BF16, GGML_TYPE_Q5_0) + FATTN_VEC_CASES_ALL_D(GGML_TYPE_F8_E4M3, GGML_TYPE_Q5_0) FATTN_VEC_CASES_ALL_D(GGML_TYPE_F16, GGML_TYPE_Q5_1) FATTN_VEC_CASES_ALL_D(GGML_TYPE_Q4_0, GGML_TYPE_Q5_1) @@ -434,6 +438,7 @@ static void ggml_cuda_flash_attn_ext_vec(ggml_backend_cuda_context & ctx, ggml_t FATTN_VEC_CASES_ALL_D(GGML_TYPE_Q5_1, GGML_TYPE_Q5_1) FATTN_VEC_CASES_ALL_D(GGML_TYPE_Q8_0, GGML_TYPE_Q5_1) FATTN_VEC_CASES_ALL_D(GGML_TYPE_BF16, GGML_TYPE_Q5_1) + FATTN_VEC_CASES_ALL_D(GGML_TYPE_F8_E4M3, GGML_TYPE_Q5_1) FATTN_VEC_CASES_ALL_D(GGML_TYPE_F16, GGML_TYPE_Q8_0) FATTN_VEC_CASES_ALL_D(GGML_TYPE_Q4_0, GGML_TYPE_Q8_0) @@ -442,6 +447,7 @@ static void ggml_cuda_flash_attn_ext_vec(ggml_backend_cuda_context & ctx, ggml_t FATTN_VEC_CASES_ALL_D(GGML_TYPE_Q5_1, GGML_TYPE_Q8_0) FATTN_VEC_CASES_ALL_D(GGML_TYPE_Q8_0, GGML_TYPE_Q8_0) FATTN_VEC_CASES_ALL_D(GGML_TYPE_BF16, GGML_TYPE_Q8_0) + FATTN_VEC_CASES_ALL_D(GGML_TYPE_F8_E4M3, GGML_TYPE_Q8_0) FATTN_VEC_CASES_ALL_D(GGML_TYPE_F16, GGML_TYPE_BF16) FATTN_VEC_CASES_ALL_D(GGML_TYPE_Q4_0, GGML_TYPE_BF16) @@ -450,11 +456,26 @@ static void ggml_cuda_flash_attn_ext_vec(ggml_backend_cuda_context & ctx, ggml_t FATTN_VEC_CASES_ALL_D(GGML_TYPE_Q5_1, GGML_TYPE_BF16) FATTN_VEC_CASES_ALL_D(GGML_TYPE_Q8_0, GGML_TYPE_BF16) FATTN_VEC_CASES_ALL_D(GGML_TYPE_BF16, GGML_TYPE_BF16) + FATTN_VEC_CASES_ALL_D(GGML_TYPE_F8_E4M3, GGML_TYPE_BF16) + + FATTN_VEC_CASES_ALL_D(GGML_TYPE_F16, GGML_TYPE_F8_E4M3) + FATTN_VEC_CASES_ALL_D(GGML_TYPE_Q4_0, GGML_TYPE_F8_E4M3) + FATTN_VEC_CASES_ALL_D(GGML_TYPE_Q4_1, GGML_TYPE_F8_E4M3) + FATTN_VEC_CASES_ALL_D(GGML_TYPE_Q5_0, GGML_TYPE_F8_E4M3) + FATTN_VEC_CASES_ALL_D(GGML_TYPE_Q5_1, GGML_TYPE_F8_E4M3) + FATTN_VEC_CASES_ALL_D(GGML_TYPE_Q8_0, GGML_TYPE_F8_E4M3) + FATTN_VEC_CASES_ALL_D(GGML_TYPE_BF16, GGML_TYPE_F8_E4M3) + FATTN_VEC_CASES_ALL_D(GGML_TYPE_F8_E4M3, GGML_TYPE_F8_E4M3) #else FATTN_VEC_CASES_ALL_D(GGML_TYPE_F16, GGML_TYPE_F16) FATTN_VEC_CASES_ALL_D(GGML_TYPE_Q4_0, GGML_TYPE_Q4_0) FATTN_VEC_CASES_ALL_D(GGML_TYPE_Q8_0, GGML_TYPE_Q8_0) FATTN_VEC_CASES_ALL_D(GGML_TYPE_BF16, GGML_TYPE_BF16) + FATTN_VEC_CASES_ALL_D(GGML_TYPE_F16, GGML_TYPE_F8_E4M3) + FATTN_VEC_CASES_ALL_D(GGML_TYPE_F8_E4M3, GGML_TYPE_F16) + FATTN_VEC_CASES_ALL_D(GGML_TYPE_BF16, GGML_TYPE_F8_E4M3) + FATTN_VEC_CASES_ALL_D(GGML_TYPE_F8_E4M3, GGML_TYPE_BF16) + FATTN_VEC_CASES_ALL_D(GGML_TYPE_F8_E4M3, GGML_TYPE_F8_E4M3) #endif // GGML_CUDA_FA_ALL_QUANTS GGML_ABORT("fatal error"); @@ -482,6 +503,7 @@ static bool ggml_cuda_fattn_kv_type_supported(ggml_type type) { case GGML_TYPE_Q4_0: case GGML_TYPE_Q8_0: case GGML_TYPE_BF16: + case GGML_TYPE_F8_E4M3: return true; default: return false; @@ -573,7 +595,10 @@ static best_fattn_kernel ggml_cuda_get_best_fattn_kernel(const int device, const } #ifndef GGML_CUDA_FA_ALL_QUANTS - if (K->type != V->type) { + const bool mixed_fp8 = + (K->type == GGML_TYPE_F8_E4M3 && (V->type == GGML_TYPE_F32 || V->type == GGML_TYPE_F16 || V->type == GGML_TYPE_BF16)) || + (V->type == GGML_TYPE_F8_E4M3 && (K->type == GGML_TYPE_F32 || K->type == GGML_TYPE_F16 || K->type == GGML_TYPE_BF16)); + if (K->type != V->type && !mixed_fp8) { return BEST_FATTN_KERNEL_NONE; } #endif // GGML_CUDA_FA_ALL_QUANTS diff --git a/ggml/src/ggml-cuda/fp8.cu b/ggml/src/ggml-cuda/fp8.cu new file mode 100644 index 0000000000..019432d1e6 --- /dev/null +++ b/ggml/src/ggml-cuda/fp8.cu @@ -0,0 +1,217 @@ +#include "fp8.cuh" + +static __global__ void mul_mat_fp8_fallback( + const char * src0, const char * src1, char * dst, int64_t ne00, int64_t ne01, int64_t ne11, + int64_t ne12, int64_t ne13, int64_t r2, int64_t r3, int64_t nb01, int64_t nb02, int64_t nb03, + int64_t nb11, int64_t nb12, int64_t nb13, int64_t nb1, int64_t nb2, int64_t nb3, int64_t ne_dst) { + for (int64_t id = blockIdx.x; id < ne_dst; id += gridDim.x) { + int64_t tmp = id / ne01; + const int64_t i0 = id - tmp*ne01; + const int64_t i1 = tmp % ne11; + tmp /= ne11; + const int64_t i2 = tmp % ne12; + const int64_t i3 = tmp / ne12; + + const ggml_fp8_e4m3_t * x = (const ggml_fp8_e4m3_t *) (src0 + i0*nb01 + (i2/r2)*nb02 + (i3/r3)*nb03); + const float * y = (const float *) (src1 + i1*nb11 + i2*nb12 + i3*nb13); + float sum = 0.0f; + for (int64_t k = threadIdx.x; k < ne00; k += blockDim.x) { + sum = fmaf(ggml_cuda_f8_e4m3_to_fp32(x[k].bits), y[k], sum); + } + + __shared__ float shared[WARP_SIZE]; + sum = block_reduce(sum, shared); + if (threadIdx.x == 0) { + *(float *) (dst + i0*sizeof(float) + i1*nb1 + i2*nb2 + i3*nb3) = sum; + } + } +} + +void ggml_cuda_mul_mat_fp8_fallback( + ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst) { + GGML_ASSERT(src0->type == GGML_TYPE_F8_E4M3); + GGML_ASSERT(src1->type == GGML_TYPE_F32); + GGML_ASSERT(dst->type == GGML_TYPE_F32); + GGML_ASSERT(src0->nb[0] == sizeof(ggml_fp8_e4m3_t)); + GGML_ASSERT(src1->nb[0] == sizeof(float)); + + const int64_t r2 = src1->ne[2] / src0->ne[2]; + const int64_t r3 = src1->ne[3] / src0->ne[3]; + const int64_t ne_dst = ggml_nelements(dst); + const int blocks = std::min(ne_dst, 65535); + mul_mat_fp8_fallback<<>>( + (const char *) src0->data, (const char *) src1->data, (char *) dst->data, + src0->ne[0], src0->ne[1], src1->ne[1], src1->ne[2], src1->ne[3], r2, r3, + src0->nb[1], src0->nb[2], src0->nb[3], src1->nb[1], src1->nb[2], src1->nb[3], + dst->nb[1], dst->nb[2], dst->nb[3], ne_dst); + CUDA_CHECK(cudaGetLastError()); +} + +#if !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) && CUDART_VERSION >= 11080 + +struct fp8_abs_src { + const float * x; + int64_t ne0; + int64_t ne1; + int64_t ne2; + int64_t s1; + int64_t s2; + int64_t s3; + + __device__ float operator()(int64_t i) const { + const int64_t i0 = i % ne0; + i /= ne0; + const int64_t i1 = i % ne1; + i /= ne1; + const int64_t i2 = i % ne2; + const int64_t i3 = i / ne2; + const float value = fabsf(x[i0 + i1*s1 + i2*s2 + i3*s3]); + return isfinite(value) ? value : 448.0f; + } +}; + +static __global__ void fp8_amax_partials(fp8_abs_src src, int64_t ne, float * partials) { + float amax = 0.0f; + for (int64_t i = (int64_t) blockIdx.x*blockDim.x + threadIdx.x; i < ne; i += (int64_t) blockDim.x*gridDim.x) { + amax = fmaxf(amax, src(i)); + } + + __shared__ float shared[WARP_SIZE]; + amax = block_reduce(amax, shared); + if (threadIdx.x == 0) { + partials[blockIdx.x] = amax; + } +} + +static __global__ void fp8_amax_final(const float * partials, int n, float * amax) { + float value = 0.0f; + for (int i = threadIdx.x; i < n; i += blockDim.x) { + value = fmaxf(value, partials[i]); + } + + __shared__ float shared[WARP_SIZE]; + value = block_reduce(value, shared); + if (threadIdx.x == 0) { + *amax = value; + } +} + +static __global__ void quantize_fp8_e4m3( + const float * __restrict__ x, uint8_t * __restrict__ y, const float * __restrict__ amax, + float * __restrict__ scale, int64_t ne0, int64_t ne1, int64_t ne2, int64_t ne, int64_t s1, int64_t s2, int64_t s3) { + const int64_t i = (int64_t) blockIdx.x*blockDim.x + threadIdx.x; + if (i >= ne) { + return; + } + + int64_t tmp = i / ne0; + const int64_t i0 = i - tmp*ne0; + const int64_t i1 = tmp % ne1; + tmp /= ne1; + const int64_t i2 = tmp % ne2; + const int64_t i3 = tmp / ne2; + + const float d = *amax > 0.0f ? *amax / 448.0f : 1.0f; + const __nv_fp8_e4m3 q(x[i0 + i1*s1 + i2*s2 + i3*s3] / d); + y[i] = q.__x; + if (i == 0) { + *scale = d; + } +} + +static void fp8_destroy_matmul( + cublasLtMatmulDesc_t op_desc, cublasLtMatrixLayout_t a_desc, cublasLtMatrixLayout_t b_desc, + cublasLtMatrixLayout_t d_desc) { + CUBLAS_CHECK(cublasLtMatrixLayoutDestroy(d_desc)); + CUBLAS_CHECK(cublasLtMatrixLayoutDestroy(b_desc)); + CUBLAS_CHECK(cublasLtMatrixLayoutDestroy(a_desc)); + CUBLAS_CHECK(cublasLtMatmulDescDestroy(op_desc)); +} + +bool ggml_cuda_mul_mat_fp8( + ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst) { + const int cc = ggml_cuda_info().devices[ctx.device].cc; + if (!fp8_mma_hardware_available(cc) || src0->type != GGML_TYPE_F8_E4M3 || src1->type != GGML_TYPE_F32 || + dst->type != GGML_TYPE_F32 || !ggml_is_contiguous(dst) || src0->ne[0] % 16 != 0 || src0->ne[1] % 16 != 0 || + src0->nb[0] != sizeof(uint8_t) || src0->nb[1] != (size_t) src0->ne[0] || src1->nb[0] != sizeof(float)) { + return false; + } + + GGML_TENSOR_BINARY_OP_LOCALS + GGML_ASSERT(ne10 == ne00); + GGML_ASSERT(ne0 == ne01); + GGML_ASSERT(ne12 % ne02 == 0); + GGML_ASSERT(ne13 % ne03 == 0); + + cudaStream_t stream = ctx.stream(); + const int64_t ne_src1 = ggml_nelements(src1); + ggml_cuda_pool_alloc src1_fp8(ctx.pool(), ne_src1); + ggml_cuda_pool_alloc src1_scale(ctx.pool(), 1); + ggml_cuda_pool_alloc src1_amax(ctx.pool(), 1); + + const fp8_abs_src abs_src = { + (const float *) src1->data, ne10, ne11, ne12, + (int64_t) (nb11 / sizeof(float)), (int64_t) (nb12 / sizeof(float)), (int64_t) (nb13 / sizeof(float)) + }; + const int reduce_blocks = std::min((ne_src1 + 255)/256, 1024); + ggml_cuda_pool_alloc reduce_tmp(ctx.pool(), reduce_blocks); + fp8_amax_partials<<>>(abs_src, ne_src1, reduce_tmp.ptr); + fp8_amax_final<<<1, 256, 0, stream>>>(reduce_tmp.ptr, reduce_blocks, src1_amax.ptr); + + quantize_fp8_e4m3<<<(ne_src1 + 255)/256, 256, 0, stream>>>( + (const float *) src1->data, src1_fp8.ptr, src1_amax.ptr, src1_scale.ptr, ne10, ne11, ne12, ne_src1, + nb11 / sizeof(float), nb12 / sizeof(float), nb13 / sizeof(float)); + CUDA_CHECK(cudaGetLastError()); + + cublasLtMatmulDesc_t op_desc; + cublasLtMatrixLayout_t a_desc; + cublasLtMatrixLayout_t b_desc; + cublasLtMatrixLayout_t d_desc; + CUBLAS_CHECK(cublasLtMatmulDescCreate(&op_desc, CUBLAS_COMPUTE_32F, CUDA_R_32F)); + const cublasOperation_t trans_a = CUBLAS_OP_T; + CUBLAS_CHECK(cublasLtMatmulDescSetAttribute(op_desc, CUBLASLT_MATMUL_DESC_TRANSA, &trans_a, sizeof(trans_a))); + CUBLAS_CHECK(cublasLtMatmulDescSetAttribute( + op_desc, CUBLASLT_MATMUL_DESC_B_SCALE_POINTER, &src1_scale.ptr, sizeof(src1_scale.ptr))); + CUBLAS_CHECK(cublasLtMatrixLayoutCreate(&a_desc, CUDA_R_8F_E4M3, ne00, ne01, ne00)); + CUBLAS_CHECK(cublasLtMatrixLayoutCreate(&b_desc, CUDA_R_8F_E4M3, ne10, ne11, ne10)); + CUBLAS_CHECK(cublasLtMatrixLayoutCreate(&d_desc, CUDA_R_32F, ne0, ne1, ne0)); + + cublasLtMatmulPreference_t preference; + CUBLAS_CHECK(cublasLtMatmulPreferenceCreate(&preference)); + cublasLtMatmulHeuristicResult_t heuristic; + int returned = 0; + const cublasStatus_t status = cublasLtMatmulAlgoGetHeuristic( + ctx.cublaslt_handle(), op_desc, a_desc, b_desc, d_desc, d_desc, preference, 1, &heuristic, &returned); + CUBLAS_CHECK(cublasLtMatmulPreferenceDestroy(preference)); + if (status != CUBLAS_STATUS_SUCCESS || returned == 0) { + fp8_destroy_matmul(op_desc, a_desc, b_desc, d_desc); + return false; + } + + const float alpha = 1.0f; + const float beta = 0.0f; + const int64_t r2 = ne12 / ne02; + const int64_t r3 = ne13 / ne03; + for (int64_t i3 = 0; i3 < ne13; ++i3) { + for (int64_t i2 = 0; i2 < ne12; ++i2) { + const char * a = (const char *) src0->data + (i2/r2)*nb02 + (i3/r3)*nb03; + const uint8_t * b = src1_fp8.ptr + (i3*ne12 + i2)*ne11*ne10; + float * d = (float *) ((char *) dst->data + i2*dst->nb[2] + i3*dst->nb[3]); + CUBLAS_CHECK(cublasLtMatmul(ctx.cublaslt_handle(), op_desc, &alpha, a, a_desc, b, b_desc, + &beta, d, d_desc, d, d_desc, &heuristic.algo, nullptr, 0, stream)); + } + } + + fp8_destroy_matmul(op_desc, a_desc, b_desc, d_desc); + return true; +} + +#else + +bool ggml_cuda_mul_mat_fp8( + ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst) { + GGML_UNUSED_VARS(ctx, src0, src1, dst); + return false; +} + +#endif diff --git a/ggml/src/ggml-cuda/fp8.cuh b/ggml/src/ggml-cuda/fp8.cuh new file mode 100644 index 0000000000..bcaf1d46a7 --- /dev/null +++ b/ggml/src/ggml-cuda/fp8.cuh @@ -0,0 +1,9 @@ +#pragma once + +#include "common.cuh" + +bool ggml_cuda_mul_mat_fp8( + ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst); + +void ggml_cuda_mul_mat_fp8_fallback( + ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst); diff --git a/ggml/src/ggml-cuda/getrows.cu b/ggml/src/ggml-cuda/getrows.cu index 6b36b8fbc1..71c3321d25 100644 --- a/ggml/src/ggml-cuda/getrows.cu +++ b/ggml/src/ggml-cuda/getrows.cu @@ -344,6 +344,10 @@ static void ggml_cuda_get_rows_switch_src0_type( get_rows_cuda_q(src0_d, src1_d, dst_d, ne00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb1, nb2, nb3, stream); break; + case GGML_TYPE_F8_E4M3: + get_rows_cuda_q<1, 1, dequantize_f8_e4m3>(src0_d, src1_d, dst_d, + ne00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb1, nb2, nb3, stream); + break; case GGML_TYPE_Q2_K: get_rows_cuda_kq<64, dst_t, dequantize_q2_K>(src0_d, src1_d, dst_d, ne00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb1, nb2, nb3, stream); diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu index 45e9537f0e..a74fcb41ae 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu @@ -25,6 +25,7 @@ #include "ggml-cuda/diagmask.cuh" #include "ggml-cuda/diag.cuh" #include "ggml-cuda/fattn.cuh" +#include "ggml-cuda/fp8.cuh" #include "ggml-cuda/fwht.cuh" #include "ggml-cuda/getrows.cuh" #include "ggml-cuda/im2col.cuh" @@ -720,6 +721,11 @@ ggml_backend_cuda_context::~ggml_backend_cuda_context() { CUDA_CHECK(cudaFree(cublas_workspaces[i][j])); } } +#if !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) && CUDART_VERSION >= 11080 + if (cublaslt_handles[i] != nullptr) { + CUBLAS_CHECK(cublasLtDestroy(cublaslt_handles[i])); + } +#endif } } @@ -1795,11 +1801,12 @@ static bool ggml_cuda_should_fuse_mul_mat_vec_q(const ggml_tensor * tensor) { ggml_nbytes(src0) != ggml_backend_buffer_get_alloc_size(src0->buffer, src0) && src0->view_src; - bool use_mul_mat_vec_q = ggml_is_quantized(src0->type) && !bad_padding_clear && src1->type == GGML_TYPE_F32 && - dst->type == GGML_TYPE_F32 && src1->ne[1] <= MMVQ_MAX_BATCH_SIZE; + const int cc = ggml_cuda_info().devices[ggml_cuda_get_device()].cc; + bool use_mul_mat_vec_q = ggml_cuda_should_use_mmvq(src0->type, cc, src1->ne[1]) && !bad_padding_clear && + src1->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32; + use_mul_mat_vec_q = use_mul_mat_vec_q && (src0->type != GGML_TYPE_F8_E4M3 || src0->ne[0] % QK8_1 == 0); // fusion is not universally faster on Pascal - const int cc = ggml_cuda_info().devices[ggml_cuda_get_device()].cc; if (cc <= GGML_CUDA_CC_PASCAL) { return false; } @@ -1815,12 +1822,23 @@ static bool ggml_cuda_should_fuse_mul_mat_vec_q(const ggml_tensor * tensor) { return use_mul_mat_vec_q; } -static void ggml_cuda_mul_mat(ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst) { +// returns true if a scaled matmul applied its weight scale in-kernel (mmvq epilogue fold), +// so the caller can skip the separate ggml_cuda_op_mul_mat_scale pass +static bool ggml_cuda_mul_mat(ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst) { GGML_TENSOR_BINARY_OP_LOCALS const int32_t hint = ggml_get_op_params_i32(dst, 1); if (hint == GGML_HINT_SRC0_IS_HADAMARD && ggml_cuda_op_fwht(ctx, src1, dst)) { - return; + return false; + } + + // mmvq can fold the weight scale (src[2]) into its epilogue for NVFP4 per-tensor scales + const bool scale_fold = dst->src[2] != NULL && + src0->type == GGML_TYPE_NVFP4 && ggml_nelements(dst->src[2]) == 1 && + ggml_cuda_should_fuse_mul_mat_vec_q(dst); + ggml_cuda_mm_fusion_args_host fusion = {}; + if (scale_fold) { + fusion.x_scale = dst->src[2]; } // If src0 is a temporary compute buffer it may have some padding that needs to be cleared for mul_mat_vec_q or mul_mat_q. @@ -1830,7 +1848,7 @@ static void ggml_cuda_mul_mat(ggml_backend_cuda_context & ctx, const ggml_tensor && ggml_nbytes(src0) != ggml_backend_buffer_get_alloc_size(src0->buffer, src0) && src0->view_src; if (bad_padding_clear || src1->type != GGML_TYPE_F32 || dst->type != GGML_TYPE_F32) { ggml_cuda_mul_mat_cublas(ctx, src0, src1, dst); - return; + return false; } const int cc = ggml_cuda_info().devices[ctx.device].cc; @@ -1840,7 +1858,7 @@ static void ggml_cuda_mul_mat(ggml_backend_cuda_context & ctx, const ggml_tensor // The custom F16 vector kernel can be used over batched cuBLAS GEMM. // But this is only faster for GPUs without tensor cores or with a thin src0 matrix (particularly KQV in attention) ggml_cuda_mul_mat_vec_f(ctx, src0, src1, nullptr, dst); - return; + return false; } // A transposed vector can still use MMVQ (i.e. ne01 == 1) if (ne01 == 1 && ne11 > MMVF_MAX_BATCH_SIZE && ne2 == 1 && ne3 == 1 @@ -1854,21 +1872,30 @@ static void ggml_cuda_mul_mat(ggml_backend_cuda_context & ctx, const ggml_tensor dst_vec.nb[2] = dst_vec.nb[1]; dst_vec.nb[3] = dst_vec.nb[1]; ggml_cuda_mul_mat_vec_f(ctx, src1, src0, nullptr, &dst_vec); - return; + return false; } if (ggml_cuda_should_use_mmf(src0->type, cc, warp_size, src0->ne, src0->nb, ne11, /*mul_mat_id =*/ false)) { ggml_cuda_mul_mat_f(ctx, src0, src1, nullptr, dst); - return; + return false; } - if (ggml_cuda_should_use_mmvq(src0->type, cc, ne11)) { - ggml_cuda_mul_mat_vec_q(ctx, src0, src1, nullptr, dst); - return; + if (ggml_cuda_should_use_mmvq(src0->type, cc, ne11) && + (src0->type != GGML_TYPE_F8_E4M3 || ne00 % QK8_1 == 0)) { + ggml_cuda_mul_mat_vec_q(ctx, src0, src1, nullptr, dst, scale_fold ? &fusion : nullptr); + return scale_fold; } if (ggml_cuda_should_use_mmq(src0->type, cc, ne11, /*n_experts =*/ 0)) { ggml_cuda_mul_mat_q(ctx, src0, src1, nullptr, dst); - return; + return false; + } + if (src0->type == GGML_TYPE_F8_E4M3) { + // F8 kernels apply no weight scale; the trailing ggml_cuda_op_mul_mat_scale epilogue does (src[2]) + if (!ggml_cuda_mul_mat_fp8(ctx, src0, src1, dst)) { + ggml_cuda_mul_mat_fp8_fallback(ctx, src0, src1, dst); + } + return false; } ggml_cuda_mul_mat_cublas(ctx, src0, src1, dst); + return false; } // returns true when ggml_cuda_mul_mat_id takes the fallback path that requires stream synchronization @@ -1882,8 +1909,9 @@ static bool ggml_cuda_mul_mat_id_needs_sync(const ggml_tensor * dst, const int c } if (dst->ne[2] <= MMVQ_MAX_BATCH_SIZE) { - if (ggml_is_quantized(src0->type)) { - if (dst->ne[2] <= get_mmvq_mmid_max_batch(src0->type, cc)) { + if (ggml_is_quantized(src0->type) || src0->type == GGML_TYPE_F8_E4M3) { + if (dst->ne[2] <= get_mmvq_mmid_max_batch(src0->type, cc) && + (src0->type != GGML_TYPE_F8_E4M3 || src0->ne[0] % QK8_1 == 0)) { return false; } } else if (GGML_CUDA_CC_IS_AMD(cc)) { @@ -1902,7 +1930,8 @@ static bool ggml_cuda_mul_mat_id_needs_sync(const ggml_tensor * dst, const int c return true; } -static void ggml_cuda_mul_mat_id(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { +// returns true if a scaled matmul applied its weight scale in-kernel (mmvq epilogue fold) +static bool ggml_cuda_mul_mat_id(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { const ggml_tensor * src0 = dst->src[0]; const ggml_tensor * src1 = dst->src[1]; const ggml_tensor * ids = dst->src[2]; @@ -1914,32 +1943,42 @@ static void ggml_cuda_mul_mat_id(ggml_backend_cuda_context & ctx, ggml_tensor * const int cc = ggml_cuda_info().devices[ggml_cuda_get_device()].cc; + // mmvq can fold the per-expert weight scale (src[3]) into its epilogue for NVFP4 + const bool scale_fold = dst->src[3] != NULL && + src0->type == GGML_TYPE_NVFP4 && ggml_nelements(dst->src[3]) == src0->ne[2] && + ggml_cuda_should_fuse_mul_mat_vec_q(dst); + ggml_cuda_mm_fusion_args_host fusion = {}; + if (scale_fold) { + fusion.x_scale = dst->src[3]; + } + // [TAG_MUL_MAT_ID_CUDA_GRAPHS] if (src1->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) { static_assert(MMVQ_MAX_BATCH_SIZE == MMVF_MAX_BATCH_SIZE); if (ne2 <= MMVQ_MAX_BATCH_SIZE) { - if (ggml_is_quantized(src0->type)) { + if (ggml_is_quantized(src0->type) || src0->type == GGML_TYPE_F8_E4M3) { const int mmvq_mmid_max = get_mmvq_mmid_max_batch(src0->type, cc); - if (ne2 <= mmvq_mmid_max) { - ggml_cuda_mul_mat_vec_q(ctx, src0, src1, ids, dst); - return; + if (ne2 <= mmvq_mmid_max && + (src0->type != GGML_TYPE_F8_E4M3 || ne00 % QK8_1 == 0)) { + ggml_cuda_mul_mat_vec_q(ctx, src0, src1, ids, dst, scale_fold ? &fusion : nullptr); + return scale_fold; } } else { if (GGML_CUDA_CC_IS_AMD(cc)) { ggml_cuda_mul_mat_vec_f(ctx, src0, src1, ids, dst); - return; + return false; } } } if (ggml_cuda_should_use_mmq(src0->type, cc, ne12, /*n_experts=*/ne02)) { ggml_cuda_mul_mat_q(ctx, src0, src1, ids, dst); - return; + return false; } if (ggml_cuda_should_use_mmf(src0->type, cc, WARP_SIZE, src0->ne, src0->nb, src1->ne[2], /*mul_mat_id=*/true)) { ggml_cuda_mul_mat_f(ctx, src0, src1, ids, dst); - return; + return false; } } @@ -1951,7 +1990,7 @@ static void ggml_cuda_mul_mat_id(ggml_backend_cuda_context & ctx, ggml_tensor * GGML_ASSERT(nb2 % nb1 == 0); const ggml_type type_src1_sorted = (src0->type == GGML_TYPE_F16 && !fast_fp16_hardware_available(cc)) - || ggml_is_quantized(src0->type) ? GGML_TYPE_F32 : src0->type; + || ggml_is_quantized(src0->type) || src0->type == GGML_TYPE_F8_E4M3 ? GGML_TYPE_F32 : src0->type; const ggml_type type_dst_sorted = GGML_TYPE_F32; const size_t ts_src1_sorted = ggml_type_size(type_src1_sorted); const size_t ts_dst_sorted = ggml_type_size(type_dst_sorted); @@ -2057,6 +2096,72 @@ static void ggml_cuda_mul_mat_id(ggml_backend_cuda_context & ctx, ggml_tensor * ne0, ne0*ts_dst_sorted, ne_get_rows*ne0*ts_dst_sorted, ne_get_rows*ne0*ts_dst_sorted, ne_get_rows, 1, 1, sizeof(int32_t), ne_get_rows*sizeof(int32_t), ne_get_rows*sizeof(int32_t), nb1, nb2, nb3, stream); + return false; +} + +// epilogue for a scaled MUL_MAT/MUL_MAT_ID (weight scale in src[2] dense / src[3] id) +#define CUDA_MUL_MAT_SCALE_BLOCK_SIZE 256 +static __global__ void mul_mat_scale_dense(float * dst, const float * scale, + int64_t ne0, int64_t ne1, int64_t ne2, int64_t ne3, + int64_t sne0, int64_t sne1, int64_t sne2, int64_t sne3) { + const int64_t i = (int64_t) blockIdx.x * blockDim.x + threadIdx.x; + if (i >= ne0*ne1*ne2*ne3) { + return; + } + const int64_t i0 = i % ne0; + const int64_t i1 = (i / ne0) % ne1; + const int64_t i2 = (i / (ne0*ne1)) % ne2; + const int64_t i3 = i / (ne0*ne1*ne2); + const int64_t si = (i0 % sne0) + (i1 % sne1)*sne0 + (i2 % sne2)*sne0*sne1 + (i3 % sne3)*sne0*sne1*sne2; + dst[i] *= scale[si]; +} + +// sne0 == 0: scalar per-expert scale[expert]; sne0 > 0: per-channel-per-expert scale[i0 + expert*sne0] +static __global__ void mul_mat_scale_id(float * dst, const float * scale, const int32_t * ids, + int64_t ne0, int64_t ne1, int64_t ne2, int64_t ids_s0, int64_t ids_s1, int64_t sne0) { + const int64_t i = (int64_t) blockIdx.x * blockDim.x + threadIdx.x; + if (i >= ne0*ne1*ne2) { + return; + } + const int64_t i0 = i % ne0; + const int64_t i1 = (i / ne0) % ne1; + const int64_t i2 = i / (ne0*ne1); + const int64_t expert = ids[i1*ids_s0 + i2*ids_s1]; + dst[i] *= scale[sne0 > 0 ? i0 + expert*sne0 : expert]; +} + +static void ggml_cuda_op_mul_mat_scale(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { + const bool is_id = dst->op == GGML_OP_MUL_MAT_ID; + const ggml_tensor * scale = is_id ? dst->src[3] : dst->src[2]; + GGML_ASSERT(scale != NULL && scale->type == GGML_TYPE_F32 && ggml_is_contiguous(scale)); + GGML_ASSERT(ggml_is_contiguous(dst)); + + float * dst_d = (float *) dst->data; + const float * scale_d = (const float *) scale->data; + cudaStream_t stream = ctx.stream(); + + const int64_t ne0 = dst->ne[0], ne1 = dst->ne[1], ne2 = dst->ne[2], ne3 = dst->ne[3]; + + // per-expert scale is indexed through ids (scalar [n_expert] or 2D per-channel [n_out, n_expert]); + // per-tensor / dense scale is a plain broadcast + const bool per_channel_expert = is_id && scale->ne[1] == dst->src[0]->ne[2]; + const bool per_expert = is_id && (ggml_nelements(scale) == dst->src[0]->ne[2] || per_channel_expert); + + if (per_expert) { + const ggml_tensor * ids = dst->src[2]; + const int64_t total = ne0*ne1*ne2; + const int64_t blocks = (total + CUDA_MUL_MAT_SCALE_BLOCK_SIZE - 1) / CUDA_MUL_MAT_SCALE_BLOCK_SIZE; + mul_mat_scale_id<<>>( + dst_d, scale_d, (const int32_t *) ids->data, ne0, ne1, ne2, + ids->nb[0]/sizeof(int32_t), ids->nb[1]/sizeof(int32_t), per_channel_expert ? scale->ne[0] : 0); + } else { + const int64_t total = ne0*ne1*ne2*ne3; + const int64_t blocks = (total + CUDA_MUL_MAT_SCALE_BLOCK_SIZE - 1) / CUDA_MUL_MAT_SCALE_BLOCK_SIZE; + mul_mat_scale_dense<<>>( + dst_d, scale_d, ne0, ne1, ne2, ne3, + scale->ne[0], scale->ne[1], scale->ne[2], scale->ne[3]); + } + CUDA_CHECK(cudaGetLastError()); } static bool ggml_cuda_compute_forward(ggml_backend_cuda_context & ctx, struct ggml_tensor * dst) { @@ -2251,12 +2356,18 @@ static bool ggml_cuda_compute_forward(ggml_backend_cuda_context & ctx, struct gg case GGML_OP_RMS_NORM_BACK: ggml_cuda_op_rms_norm_back(ctx, dst); break; - case GGML_OP_MUL_MAT: - ggml_cuda_mul_mat(ctx, dst->src[0], dst->src[1], dst); - break; - case GGML_OP_MUL_MAT_ID: - ggml_cuda_mul_mat_id(ctx, dst); - break; + case GGML_OP_MUL_MAT: { + const bool scale_folded = ggml_cuda_mul_mat(ctx, dst->src[0], dst->src[1], dst); + if (dst->src[2] != NULL && !scale_folded) { + ggml_cuda_op_mul_mat_scale(ctx, dst); + } + } break; + case GGML_OP_MUL_MAT_ID: { + const bool scale_folded = ggml_cuda_mul_mat_id(ctx, dst); + if (dst->src[3] != NULL && !scale_folded) { + ggml_cuda_op_mul_mat_scale(ctx, dst); + } + } break; case GGML_OP_OUT_PROD: ggml_cuda_out_prod(ctx, dst); break; @@ -3433,6 +3544,12 @@ static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph ggml_tensor * node = cgraph->nodes[i]; + // a scaled matmul applies its weight scale in-kernel or via epilogue; keep it out of multi-op fusion + if ((node->op == GGML_OP_MUL_MAT && node->src[2] != NULL) || + (node->op == GGML_OP_MUL_MAT_ID && node->src[3] != NULL)) { + return 0; + } + if (node->op == GGML_OP_MUL) { ggml_cuda_moe_weighted_reduction_match match; if (ggml_cuda_match_moe_weighted_reduction(cgraph, i, match)) { @@ -3638,7 +3755,7 @@ static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph } const ggml_tensor * scale = scale_lhs_mm ? scale_node->src[1] : scale_node->src[0]; - if (mm_node->src[0]->type != GGML_TYPE_NVFP4 || scale_node->type != GGML_TYPE_F32 || + if ((mm_node->src[0]->type != GGML_TYPE_NVFP4 && mm_node->src[0]->type != GGML_TYPE_F8_E4M3) || scale_node->type != GGML_TYPE_F32 || scale->type != GGML_TYPE_F32 || !ggml_is_contiguous(scale) || ggml_nelements(scale) != 1 || !ggml_are_same_shape(scale_node, mm_node)) { return nullptr; @@ -3658,7 +3775,7 @@ static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph } const ggml_tensor * scale = reshape->src[0]; - if (mm_node->src[0]->type != GGML_TYPE_NVFP4 || scale_node->type != GGML_TYPE_F32 || + if ((mm_node->src[0]->type != GGML_TYPE_NVFP4 && mm_node->src[0]->type != GGML_TYPE_F8_E4M3) || scale_node->type != GGML_TYPE_F32 || scale->type != GGML_TYPE_F32 || !ggml_is_contiguous(scale) || ggml_nelements(scale) != mm_node->src[0]->ne[2] || !ggml_are_same_shape(scale_node, mm_node)) { return nullptr; @@ -5151,6 +5268,7 @@ static bool ggml_backend_cuda_device_supports_op(ggml_backend_dev_t dev, const g case GGML_TYPE_Q8_0: case GGML_TYPE_MXFP4: case GGML_TYPE_NVFP4: + case GGML_TYPE_F8_E4M3: case GGML_TYPE_Q2_K: case GGML_TYPE_Q3_K: case GGML_TYPE_Q4_K: @@ -5188,6 +5306,7 @@ static bool ggml_backend_cuda_device_supports_op(ggml_backend_dev_t dev, const g case GGML_TYPE_Q5_0: case GGML_TYPE_Q5_1: case GGML_TYPE_Q8_0: + case GGML_TYPE_F8_E4M3: case GGML_TYPE_Q2_K: case GGML_TYPE_Q3_K: case GGML_TYPE_Q4_K: @@ -5219,7 +5338,7 @@ static bool ggml_backend_cuda_device_supports_op(ggml_backend_dev_t dev, const g { return ( ( - (op->type == GGML_TYPE_F32 || op->type == GGML_TYPE_F16 || op->type == GGML_TYPE_BF16 || + (op->type == GGML_TYPE_F32 || op->type == GGML_TYPE_F16 || op->type == GGML_TYPE_BF16 || op->type == GGML_TYPE_F8_E4M3 || op->type == GGML_TYPE_Q4_0 || op->type == GGML_TYPE_Q4_1 || op->type == GGML_TYPE_Q5_0 || op->type == GGML_TYPE_Q5_1 || op->type == GGML_TYPE_Q8_0 || op->type == GGML_TYPE_IQ4_NL) && op->src[0]->type == GGML_TYPE_F32 @@ -5649,6 +5768,12 @@ static ggml_backend_feature * ggml_backend_cuda_get_features(ggml_backend_reg_t { const auto & info = ggml_cuda_info(); + for (int id = 0; id < info.device_count; ++id) { + if (fp8_mma_hardware_available(info.devices[id].cc)) { + features.push_back({ "NATIVE_FP8", "1"}); + break; + } + } for (int id = 0; id < info.device_count; ++id) { if (blackwell_mma_available(info.devices[id].cc)) { features.push_back({ "BLACKWELL_NATIVE_FP4", "1"}); diff --git a/ggml/src/ggml-cuda/mmvq.cu b/ggml/src/ggml-cuda/mmvq.cu index f65e0fbcd7..831e287f7a 100644 --- a/ggml/src/ggml-cuda/mmvq.cu +++ b/ggml/src/ggml-cuda/mmvq.cu @@ -8,6 +8,10 @@ typedef float (*vec_dot_q_cuda_t)(const void * __restrict__ vbq, const block_q8_1 * __restrict__ bq8_1, const int & kbx, const int & iqs); +static constexpr __host__ __device__ bool is_scaled_low_precision_type(ggml_type type) { + return type == GGML_TYPE_NVFP4 || type == GGML_TYPE_F8_E4M3; +} + static constexpr __device__ vec_dot_q_cuda_t get_vec_dot_q_cuda(ggml_type type) { switch (type) { case GGML_TYPE_Q1_0: return vec_dot_q1_0_q8_1; @@ -19,6 +23,7 @@ static constexpr __device__ vec_dot_q_cuda_t get_vec_dot_q_cuda(ggml_type type) case GGML_TYPE_Q8_0: return vec_dot_q8_0_q8_1; case GGML_TYPE_MXFP4: return vec_dot_mxfp4_q8_1; case GGML_TYPE_NVFP4: return vec_dot_nvfp4_q8_1; + case GGML_TYPE_F8_E4M3: return vec_dot_f8_e4m3_q8_1; case GGML_TYPE_Q2_K: return vec_dot_q2_K_q8_1; case GGML_TYPE_Q3_K: return vec_dot_q3_K_q8_1; case GGML_TYPE_Q4_K: return vec_dot_q4_K_q8_1; @@ -48,6 +53,7 @@ static constexpr __host__ __device__ int get_vdr_mmvq(ggml_type type) { case GGML_TYPE_Q8_0: return VDR_Q8_0_Q8_1_MMVQ; case GGML_TYPE_MXFP4: return VDR_MXFP4_Q8_1_MMVQ; case GGML_TYPE_NVFP4: return VDR_NVFP4_Q8_1_MMVQ; + case GGML_TYPE_F8_E4M3: return VDR_F8_E4M3_Q8_1_MMVQ; case GGML_TYPE_Q2_K: return VDR_Q2_K_Q8_1_MMVQ; case GGML_TYPE_Q3_K: return VDR_Q3_K_Q8_1_MMVQ; case GGML_TYPE_Q4_K: return VDR_Q4_K_Q8_1_MMVQ; @@ -287,7 +293,7 @@ int get_mmvq_mmid_max_batch(ggml_type type, int cc) { } bool ggml_cuda_should_use_mmvq(enum ggml_type type, int cc, int64_t ne11) { - if (!ggml_is_quantized(type)) { + if (!ggml_is_quantized(type) && type != GGML_TYPE_F8_E4M3) { return false; } // k-quants cost more to decode and mvq redoes that per column, so MMQ wins sooner. @@ -570,6 +576,7 @@ static __global__ void mul_mat_vec_q( constexpr int qk = ggml_cuda_type_traits::qk; constexpr int qi = ggml_cuda_type_traits::qi; constexpr int vdr = get_vdr_mmvq(type); + constexpr int kbx_stride = type == GGML_TYPE_F8_E4M3 ? QK8_1 : 1; constexpr mmvq_parameter_table_id table_id = get_device_table_id(); constexpr int nwarps = calc_nwarps(type, ncols_dst, table_id, small_k, halve_iters); constexpr int rows_per_cuda_block = calc_rows_per_block(ncols_dst, table_id, small_k, nwarps); @@ -618,7 +625,7 @@ static __global__ void mul_mat_vec_q( gate_bias = (const float *) fusion.gate_bias; active_glu = fusion.glu_op; glu_limit = fusion.glu_limit; - if constexpr (type == GGML_TYPE_NVFP4) { + if constexpr (is_scaled_low_precision_type(type)) { use_scale = fusion.x_scale != nullptr; use_gate_scale = fusion.gate_scale != nullptr && use_gate; x_scale = (const float *) fusion.x_scale; @@ -651,7 +658,7 @@ static __global__ void mul_mat_vec_q( gate_biases[j] = gate_bias[j * stride_col_dst + threadIdx.x]; } } - if constexpr (type == GGML_TYPE_NVFP4) { + if constexpr (is_scaled_low_precision_type(type)) { if (use_scale) { x_scales = x_scale[ids ? channel_x : 0]; } @@ -680,11 +687,11 @@ static __global__ void mul_mat_vec_q( #pragma unroll for (int i = 0; i < rows_per_cuda_block; ++i) { tmp[j][i] += vec_dot_q_cuda( - vx, &y[j*stride_col_y + kby], kbx_offset + i*stride_row_x + kbx, kqs); + vx, &y[j*stride_col_y + kby], kbx_offset + i*stride_row_x + kbx_stride*kbx, kqs); if constexpr (has_fusion) { if (use_gate) { tmp_gate[j][i] += vec_dot_q_cuda( - vgate, &y[j*stride_col_y + kby], kbx_offset + i*stride_row_x + kbx, kqs); + vgate, &y[j*stride_col_y + kby], kbx_offset + i*stride_row_x + kbx_stride*kbx, kqs); } } } @@ -739,13 +746,13 @@ static __global__ void mul_mat_vec_q( if (threadIdx.x == i && (rows_per_cuda_block == 1 || uint32_t(row0 + i) < stride_col_dst)) { float result = tmp[j][i]; if constexpr (has_fusion) { - if constexpr (type == GGML_TYPE_NVFP4) { + if constexpr (is_scaled_low_precision_type(type)) { result *= x_scales; } result += x_biases[j]; if (use_gate) { float gate_value = tmp_gate[j][i]; - if constexpr (type == GGML_TYPE_NVFP4) { + if constexpr (is_scaled_low_precision_type(type)) { gate_value *= gate_scales; } gate_value += gate_biases[j]; @@ -776,7 +783,7 @@ static __global__ void mul_mat_vec_q( if constexpr (!has_fusion) { GGML_UNUSED_VARS(use_gate, use_bias, use_gate_bias, use_scale, use_gate_scale, active_glu, glu_limit, gate_bias, x_bias, x_scale, gate_scale, tmp_gate); } - if constexpr (type != GGML_TYPE_NVFP4) { + if constexpr (!is_scaled_low_precision_type(type)) { GGML_UNUSED_VARS(use_scale, use_gate_scale, x_scale, gate_scale, x_scales, gate_scales); } } @@ -802,6 +809,7 @@ static __global__ void mul_mat_vec_q_moe( constexpr int qk = ggml_cuda_type_traits::qk; constexpr int qi = ggml_cuda_type_traits::qi; constexpr int vdr = get_vdr_mmvq(type); + constexpr int kbx_stride = type == GGML_TYPE_F8_E4M3 ? QK8_1 : 1; constexpr int warp_size = ggml_cuda_get_physical_warp_size(); constexpr vec_dot_q_cuda_t vec_dot_q_cuda = get_vec_dot_q_cuda(type); @@ -823,7 +831,7 @@ static __global__ void mul_mat_vec_q_moe( gate_bias = (const float *) fusion.gate_bias; active_glu = fusion.glu_op; glu_limit = fusion.glu_limit; - if constexpr (type == GGML_TYPE_NVFP4) { + if constexpr (is_scaled_low_precision_type(type)) { x_scale = (const float *) fusion.x_scale; gate_scale = (const float *) fusion.gate_scale; } @@ -857,10 +865,10 @@ static __global__ void mul_mat_vec_q_moe( #pragma unroll for (int i = 0; i < c_rows_per_block; ++i) { - tmp[i] += vec_dot_q_cuda(vx, &y[kby], kbx_offset + i*stride_row_x + kbx, kqs); + tmp[i] += vec_dot_q_cuda(vx, &y[kby], kbx_offset + i*stride_row_x + kbx_stride*kbx, kqs); if constexpr (has_fusion) { if (use_gate) { - tmp_gate[i] += vec_dot_q_cuda(vgate, &y[kby], kbx_offset + i*stride_row_x + kbx, kqs); + tmp_gate[i] += vec_dot_q_cuda(vgate, &y[kby], kbx_offset + i*stride_row_x + kbx_stride*kbx, kqs); } } } @@ -885,7 +893,7 @@ static __global__ void mul_mat_vec_q_moe( if constexpr (has_fusion) { const uint32_t bias_idx = channel_x*stride_channel_dst + row0 + threadIdx.x; - if constexpr (type == GGML_TYPE_NVFP4) { + if constexpr (is_scaled_low_precision_type(type)) { if (x_scale) { result *= x_scale[channel_x]; } @@ -895,7 +903,7 @@ static __global__ void mul_mat_vec_q_moe( } if (use_gate) { float gate_value = tmp_gate[threadIdx.x]; - if constexpr (type == GGML_TYPE_NVFP4) { + if constexpr (is_scaled_low_precision_type(type)) { if (gate_scale) { gate_value *= gate_scale[channel_x]; } @@ -927,7 +935,7 @@ static __global__ void mul_mat_vec_q_moe( if constexpr (!has_fusion) { GGML_UNUSED_VARS(use_gate, tmp_gate, vgate, x_bias, gate_bias, active_glu, glu_limit, x_scale, gate_scale); - } else if constexpr (type != GGML_TYPE_NVFP4) { + } else if constexpr (!is_scaled_low_precision_type(type)) { GGML_UNUSED_VARS(x_scale, gate_scale); } } @@ -1019,7 +1027,7 @@ static void mul_mat_vec_q_switch_ncols_dst( const int nsamples_x, const int nsamples_dst, const int stride_sample_x, const int stride_sample_y, const int stride_sample_dst, const int ids_stride, cudaStream_t stream) { - GGML_ASSERT(ncols_x % ggml_blck_size(type) == 0); + GGML_ASSERT(ncols_x % ggml_cuda_type_traits::qk == 0); GGML_ASSERT(ncols_dst <= MMVQ_MAX_BATCH_SIZE); const uint3 nchannels_y_fd = ids ? init_fastdiv_values(nchannels_y) : make_uint3(0, 0, 0); @@ -1265,6 +1273,12 @@ static void mul_mat_vec_q_switch_type( nchannels_x, nchannels_y, nchannels_dst, stride_channel_x, stride_channel_y, stride_channel_dst, nsamples_x, nsamples_dst, stride_sample_x, stride_sample_y, stride_sample_dst, ids_stride, stream); break; + case GGML_TYPE_F8_E4M3: + mul_mat_vec_q_switch_ncols_dst + (vx, vy, ids, fusion, dst, ncols_x, nrows_x, ncols_dst, stride_row_x, stride_col_y, stride_col_dst, + nchannels_x, nchannels_y, nchannels_dst, stride_channel_x, stride_channel_y, stride_channel_dst, + nsamples_x, nsamples_dst, stride_sample_x, stride_sample_y, stride_sample_dst, ids_stride, stream); + break; case GGML_TYPE_Q2_K: mul_mat_vec_q_switch_ncols_dst (vx, vy, ids, fusion, dst, ncols_x, nrows_x, ncols_dst, stride_row_x, stride_col_y, stride_col_dst, @@ -1387,9 +1401,9 @@ void ggml_cuda_mul_mat_vec_q( const int cc = ggml_cuda_info().devices[ggml_cuda_get_device()].cc; GGML_ASSERT( !ids || dst->ne[2] <= get_mmvq_mmid_max_batch(src0->type, cc)); GGML_ASSERT( ids || dst->ne[1] == 1); - // Scale fusion is only allowed for NVFP4 currently as the cost of checking this at run-time in the prologue is + // Scale fusion is only allowed for scaled low-precision types as the cost of checking this at run-time in the prologue is // non-negligible for some models such as gpt-oss-20b - GGML_ASSERT((fusion->x_scale == nullptr && fusion->gate_scale == nullptr) || src0->type == GGML_TYPE_NVFP4); + GGML_ASSERT((fusion->x_scale == nullptr && fusion->gate_scale == nullptr) || is_scaled_low_precision_type(src0->type)); if (fusion->x_bias) { GGML_ASSERT(fusion->x_bias->type == GGML_TYPE_F32); diff --git a/ggml/src/ggml-cuda/set-rows.cu b/ggml/src/ggml-cuda/set-rows.cu index 4659970651..166c6748d4 100644 --- a/ggml/src/ggml-cuda/set-rows.cu +++ b/ggml/src/ggml-cuda/set-rows.cu @@ -168,7 +168,11 @@ static __global__ void k_set_rows(const src_t * src0_ptr, const src_t * src0_row = src0 + i01*s01 + i02*s02 + i03*s03; dst_t * dst_row_ptr = dst + dst_row*s1 + i02*s2 + i03*s3; - dst_row_ptr[i00] = ggml_cuda_cast(src0_row[i00]); + if constexpr (std::is_same_v) { + dst_row_ptr[i00].bits = ggml_cuda_fp32_to_f8_e4m3(src0_row[i00]); + } else { + dst_row_ptr[i00] = ggml_cuda_cast(src0_row[i00]); + } GGML_UNUSED(ne10); GGML_UNUSED(ne11); @@ -257,6 +261,16 @@ static void set_rows_cuda(ggml_backend_cuda_context & ctx, const ggml_tensor * s nb1, nb2, nb3, stream ); + } else if (dst->type == GGML_TYPE_F8_E4M3) { + set_rows_cuda( + src0_d, src1_d, (ggml_fp8_e4m3_t *) dst->data, + ne00, ne01, ne02, ne03, + ne10, ne11, ne12, ne13, + nb01, nb02, nb03, + nb10, nb11, nb12, + nb1, nb2, nb3, + stream + ); } else if (dst->type == GGML_TYPE_Q4_0) { set_rows_cuda_quant( src0_d, src1_d, (block_q4_0*)dst->data, diff --git a/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-bf16-f8_e4m3.cu b/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-bf16-f8_e4m3.cu new file mode 100644 index 0000000000..ae9672a9e5 --- /dev/null +++ b/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-bf16-f8_e4m3.cu @@ -0,0 +1,7 @@ +// This file has been autogenerated by generate_cu_files.py, do not edit manually. + +#include "../fattn-vec.cuh" + +DECL_FATTN_VEC_CASE( 64, GGML_TYPE_BF16, GGML_TYPE_F8_E4M3); +DECL_FATTN_VEC_CASE(128, GGML_TYPE_BF16, GGML_TYPE_F8_E4M3); +DECL_FATTN_VEC_CASE(256, GGML_TYPE_BF16, GGML_TYPE_F8_E4M3); diff --git a/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-f16-f8_e4m3.cu b/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-f16-f8_e4m3.cu new file mode 100644 index 0000000000..2afaefa838 --- /dev/null +++ b/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-f16-f8_e4m3.cu @@ -0,0 +1,7 @@ +// This file has been autogenerated by generate_cu_files.py, do not edit manually. + +#include "../fattn-vec.cuh" + +DECL_FATTN_VEC_CASE( 64, GGML_TYPE_F16, GGML_TYPE_F8_E4M3); +DECL_FATTN_VEC_CASE(128, GGML_TYPE_F16, GGML_TYPE_F8_E4M3); +DECL_FATTN_VEC_CASE(256, GGML_TYPE_F16, GGML_TYPE_F8_E4M3); diff --git a/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-f8_e4m3-bf16.cu b/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-f8_e4m3-bf16.cu new file mode 100644 index 0000000000..9dfb9c45e7 --- /dev/null +++ b/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-f8_e4m3-bf16.cu @@ -0,0 +1,7 @@ +// This file has been autogenerated by generate_cu_files.py, do not edit manually. + +#include "../fattn-vec.cuh" + +DECL_FATTN_VEC_CASE( 64, GGML_TYPE_F8_E4M3, GGML_TYPE_BF16); +DECL_FATTN_VEC_CASE(128, GGML_TYPE_F8_E4M3, GGML_TYPE_BF16); +DECL_FATTN_VEC_CASE(256, GGML_TYPE_F8_E4M3, GGML_TYPE_BF16); diff --git a/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-f8_e4m3-f16.cu b/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-f8_e4m3-f16.cu new file mode 100644 index 0000000000..9e7570d36f --- /dev/null +++ b/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-f8_e4m3-f16.cu @@ -0,0 +1,7 @@ +// This file has been autogenerated by generate_cu_files.py, do not edit manually. + +#include "../fattn-vec.cuh" + +DECL_FATTN_VEC_CASE( 64, GGML_TYPE_F8_E4M3, GGML_TYPE_F16); +DECL_FATTN_VEC_CASE(128, GGML_TYPE_F8_E4M3, GGML_TYPE_F16); +DECL_FATTN_VEC_CASE(256, GGML_TYPE_F8_E4M3, GGML_TYPE_F16); diff --git a/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-f8_e4m3-f8_e4m3.cu b/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-f8_e4m3-f8_e4m3.cu new file mode 100644 index 0000000000..6c380e54c8 --- /dev/null +++ b/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-f8_e4m3-f8_e4m3.cu @@ -0,0 +1,7 @@ +// This file has been autogenerated by generate_cu_files.py, do not edit manually. + +#include "../fattn-vec.cuh" + +DECL_FATTN_VEC_CASE( 64, GGML_TYPE_F8_E4M3, GGML_TYPE_F8_E4M3); +DECL_FATTN_VEC_CASE(128, GGML_TYPE_F8_E4M3, GGML_TYPE_F8_E4M3); +DECL_FATTN_VEC_CASE(256, GGML_TYPE_F8_E4M3, GGML_TYPE_F8_E4M3); diff --git a/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-f8_e4m3-q4_0.cu b/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-f8_e4m3-q4_0.cu new file mode 100644 index 0000000000..fe22bb9642 --- /dev/null +++ b/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-f8_e4m3-q4_0.cu @@ -0,0 +1,7 @@ +// This file has been autogenerated by generate_cu_files.py, do not edit manually. + +#include "../fattn-vec.cuh" + +DECL_FATTN_VEC_CASE( 64, GGML_TYPE_F8_E4M3, GGML_TYPE_Q4_0); +DECL_FATTN_VEC_CASE(128, GGML_TYPE_F8_E4M3, GGML_TYPE_Q4_0); +DECL_FATTN_VEC_CASE(256, GGML_TYPE_F8_E4M3, GGML_TYPE_Q4_0); diff --git a/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-f8_e4m3-q4_1.cu b/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-f8_e4m3-q4_1.cu new file mode 100644 index 0000000000..317273afc2 --- /dev/null +++ b/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-f8_e4m3-q4_1.cu @@ -0,0 +1,7 @@ +// This file has been autogenerated by generate_cu_files.py, do not edit manually. + +#include "../fattn-vec.cuh" + +DECL_FATTN_VEC_CASE( 64, GGML_TYPE_F8_E4M3, GGML_TYPE_Q4_1); +DECL_FATTN_VEC_CASE(128, GGML_TYPE_F8_E4M3, GGML_TYPE_Q4_1); +DECL_FATTN_VEC_CASE(256, GGML_TYPE_F8_E4M3, GGML_TYPE_Q4_1); diff --git a/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-f8_e4m3-q5_0.cu b/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-f8_e4m3-q5_0.cu new file mode 100644 index 0000000000..e05b1ee381 --- /dev/null +++ b/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-f8_e4m3-q5_0.cu @@ -0,0 +1,7 @@ +// This file has been autogenerated by generate_cu_files.py, do not edit manually. + +#include "../fattn-vec.cuh" + +DECL_FATTN_VEC_CASE( 64, GGML_TYPE_F8_E4M3, GGML_TYPE_Q5_0); +DECL_FATTN_VEC_CASE(128, GGML_TYPE_F8_E4M3, GGML_TYPE_Q5_0); +DECL_FATTN_VEC_CASE(256, GGML_TYPE_F8_E4M3, GGML_TYPE_Q5_0); diff --git a/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-f8_e4m3-q5_1.cu b/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-f8_e4m3-q5_1.cu new file mode 100644 index 0000000000..6cc6386873 --- /dev/null +++ b/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-f8_e4m3-q5_1.cu @@ -0,0 +1,7 @@ +// This file has been autogenerated by generate_cu_files.py, do not edit manually. + +#include "../fattn-vec.cuh" + +DECL_FATTN_VEC_CASE( 64, GGML_TYPE_F8_E4M3, GGML_TYPE_Q5_1); +DECL_FATTN_VEC_CASE(128, GGML_TYPE_F8_E4M3, GGML_TYPE_Q5_1); +DECL_FATTN_VEC_CASE(256, GGML_TYPE_F8_E4M3, GGML_TYPE_Q5_1); diff --git a/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-f8_e4m3-q8_0.cu b/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-f8_e4m3-q8_0.cu new file mode 100644 index 0000000000..0a88e240b0 --- /dev/null +++ b/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-f8_e4m3-q8_0.cu @@ -0,0 +1,7 @@ +// This file has been autogenerated by generate_cu_files.py, do not edit manually. + +#include "../fattn-vec.cuh" + +DECL_FATTN_VEC_CASE( 64, GGML_TYPE_F8_E4M3, GGML_TYPE_Q8_0); +DECL_FATTN_VEC_CASE(128, GGML_TYPE_F8_E4M3, GGML_TYPE_Q8_0); +DECL_FATTN_VEC_CASE(256, GGML_TYPE_F8_E4M3, GGML_TYPE_Q8_0); diff --git a/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q4_0-f8_e4m3.cu b/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q4_0-f8_e4m3.cu new file mode 100644 index 0000000000..44b4d4fc9d --- /dev/null +++ b/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q4_0-f8_e4m3.cu @@ -0,0 +1,7 @@ +// This file has been autogenerated by generate_cu_files.py, do not edit manually. + +#include "../fattn-vec.cuh" + +DECL_FATTN_VEC_CASE( 64, GGML_TYPE_Q4_0, GGML_TYPE_F8_E4M3); +DECL_FATTN_VEC_CASE(128, GGML_TYPE_Q4_0, GGML_TYPE_F8_E4M3); +DECL_FATTN_VEC_CASE(256, GGML_TYPE_Q4_0, GGML_TYPE_F8_E4M3); diff --git a/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q4_1-f8_e4m3.cu b/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q4_1-f8_e4m3.cu new file mode 100644 index 0000000000..f16490953e --- /dev/null +++ b/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q4_1-f8_e4m3.cu @@ -0,0 +1,7 @@ +// This file has been autogenerated by generate_cu_files.py, do not edit manually. + +#include "../fattn-vec.cuh" + +DECL_FATTN_VEC_CASE( 64, GGML_TYPE_Q4_1, GGML_TYPE_F8_E4M3); +DECL_FATTN_VEC_CASE(128, GGML_TYPE_Q4_1, GGML_TYPE_F8_E4M3); +DECL_FATTN_VEC_CASE(256, GGML_TYPE_Q4_1, GGML_TYPE_F8_E4M3); diff --git a/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q5_0-f8_e4m3.cu b/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q5_0-f8_e4m3.cu new file mode 100644 index 0000000000..6dab7b408a --- /dev/null +++ b/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q5_0-f8_e4m3.cu @@ -0,0 +1,7 @@ +// This file has been autogenerated by generate_cu_files.py, do not edit manually. + +#include "../fattn-vec.cuh" + +DECL_FATTN_VEC_CASE( 64, GGML_TYPE_Q5_0, GGML_TYPE_F8_E4M3); +DECL_FATTN_VEC_CASE(128, GGML_TYPE_Q5_0, GGML_TYPE_F8_E4M3); +DECL_FATTN_VEC_CASE(256, GGML_TYPE_Q5_0, GGML_TYPE_F8_E4M3); diff --git a/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q5_1-f8_e4m3.cu b/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q5_1-f8_e4m3.cu new file mode 100644 index 0000000000..2cdaa289f3 --- /dev/null +++ b/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q5_1-f8_e4m3.cu @@ -0,0 +1,7 @@ +// This file has been autogenerated by generate_cu_files.py, do not edit manually. + +#include "../fattn-vec.cuh" + +DECL_FATTN_VEC_CASE( 64, GGML_TYPE_Q5_1, GGML_TYPE_F8_E4M3); +DECL_FATTN_VEC_CASE(128, GGML_TYPE_Q5_1, GGML_TYPE_F8_E4M3); +DECL_FATTN_VEC_CASE(256, GGML_TYPE_Q5_1, GGML_TYPE_F8_E4M3); diff --git a/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q8_0-f8_e4m3.cu b/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q8_0-f8_e4m3.cu new file mode 100644 index 0000000000..2556c8c82a --- /dev/null +++ b/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q8_0-f8_e4m3.cu @@ -0,0 +1,7 @@ +// This file has been autogenerated by generate_cu_files.py, do not edit manually. + +#include "../fattn-vec.cuh" + +DECL_FATTN_VEC_CASE( 64, GGML_TYPE_Q8_0, GGML_TYPE_F8_E4M3); +DECL_FATTN_VEC_CASE(128, GGML_TYPE_Q8_0, GGML_TYPE_F8_E4M3); +DECL_FATTN_VEC_CASE(256, GGML_TYPE_Q8_0, GGML_TYPE_F8_E4M3); 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 d7cd271675..509351d6d5 100755 --- a/ggml/src/ggml-cuda/template-instances/generate_cu_files.py +++ b/ggml/src/ggml-cuda/template-instances/generate_cu_files.py @@ -8,7 +8,7 @@ HEAD_SIZES_KQ = [40, 64, 72, 80, 96, 112, 128, 192, 256, 320, 512, 576] # DKQ -> DV override for asymmetric head dims. HEAD_SIZES_V_OVERRIDE = {576: 512, 320: 256, 192: 128} -TYPES_KV = ["GGML_TYPE_F16", "GGML_TYPE_Q4_0", "GGML_TYPE_Q4_1", "GGML_TYPE_Q5_0", "GGML_TYPE_Q5_1", "GGML_TYPE_Q8_0", "GGML_TYPE_BF16"] +TYPES_KV = ["GGML_TYPE_F16", "GGML_TYPE_Q4_0", "GGML_TYPE_Q4_1", "GGML_TYPE_Q5_0", "GGML_TYPE_Q5_1", "GGML_TYPE_Q8_0", "GGML_TYPE_BF16", "GGML_TYPE_F8_E4M3"] SOURCE_FATTN_TILE = """// This file has been autogenerated by generate_cu_files.py, do not edit manually. diff --git a/ggml/src/ggml-cuda/vecdotq.cuh b/ggml/src/ggml-cuda/vecdotq.cuh index ec117c57df..38a18f7e6f 100644 --- a/ggml/src/ggml-cuda/vecdotq.cuh +++ b/ggml/src/ggml-cuda/vecdotq.cuh @@ -360,6 +360,25 @@ static __device__ __forceinline__ float vec_dot_nvfp4_q8_1( return sum; } + +#define VDR_F8_E4M3_Q8_1_MMVQ 1 + +static __device__ __forceinline__ float vec_dot_f8_e4m3_q8_1( + const void * __restrict__ vbq, + const block_q8_1 * __restrict__ bq8_1, + const int32_t & kbx, + const int32_t & iqs) { + const ggml_fp8_e4m3_t * bq8 = (const ggml_fp8_e4m3_t *) vbq + kbx; + const int8_t * q8 = bq8_1->qs + 4*iqs; + + float sum = 0.0f; +#pragma unroll + for (int i = 0; i < 4; ++i) { + sum += ggml_cuda_f8_e4m3_to_fp32(bq8[4*iqs + i].bits) * q8[i]; + } + return __low2float(bq8_1->ds) * sum; +} + #define VDR_Q2_K_Q8_1_MMVQ 1 #define VDR_Q2_K_Q8_1_MMQ 4 diff --git a/ggml/src/ggml-cuda/vendors/cuda.h b/ggml/src/ggml-cuda/vendors/cuda.h index 323c980193..86c1ed9ce9 100644 --- a/ggml/src/ggml-cuda/vendors/cuda.h +++ b/ggml/src/ggml-cuda/vendors/cuda.h @@ -3,6 +3,7 @@ #include #include #include +#include #include #include diff --git a/ggml/src/ggml.c b/ggml/src/ggml.c index 6bae112d6d..4bbe3d9477 100644 --- a/ggml/src/ggml.c +++ b/ggml/src/ggml.c @@ -762,6 +762,7 @@ static const struct ggml_type_traits type_traits[GGML_TYPE_COUNT] = { .blck_size = QK_NVFP4, .type_size = sizeof(block_nvfp4), .is_quantized = true, + .needs_scale = true, .to_float = (ggml_to_float_t) dequantize_row_nvfp4, .from_float_ref = (ggml_from_float_t)quantize_row_nvfp4_ref, }, @@ -1369,6 +1370,13 @@ bool ggml_is_quantized(enum ggml_type type) { return type_traits[type].is_quantized; } +bool ggml_needs_scale_quantized(enum ggml_type type) { + assert(type >= 0); + assert(type < GGML_TYPE_COUNT); + assert(!type_traits[type].needs_scale || type_traits[type].is_quantized); + return type_traits[type].needs_scale; +} + const char * ggml_op_name(enum ggml_op op) { return GGML_OP_NAME[op]; } @@ -3300,8 +3308,30 @@ struct ggml_tensor * ggml_mul_mat( struct ggml_context * ctx, struct ggml_tensor * a, struct ggml_tensor * b) { + if (ggml_needs_scale_quantized(a->type) || ggml_needs_scale_quantized(b->type)) { + GGML_LOG_ERROR("%s: tensor types %s or %s requires explicit dequantization scales; use ggml_mul_mat_ext instead\n", + __func__, ggml_type_name(a->type), ggml_type_name(b->type)); + GGML_ABORT("fatal error"); + } + + return ggml_mul_mat_ext(ctx, a, b, NULL, NULL); +} + +struct ggml_tensor * ggml_mul_mat_ext( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + struct ggml_tensor * scale_weight, + struct ggml_tensor * scale_activations) { GGML_ASSERT(ggml_can_mul_mat(a, b)); GGML_ASSERT(!ggml_is_transposed(a)); + if (ggml_needs_scale_quantized(b->type)) { + GGML_LOG_ERROR("%s: scaled tensor type %s currently cannot be used as the activation tensor\n", + __func__, ggml_type_name(b->type)); + GGML_ABORT("fatal error"); + } + GGML_ASSERT(scale_weight == NULL || scale_weight->type == GGML_TYPE_F32); + GGML_ASSERT(scale_activations == NULL || scale_activations->type == GGML_TYPE_F32); const int64_t ne[4] = { a->ne[1], b->ne[1], b->ne[2], b->ne[3] }; struct ggml_tensor * result = ggml_new_tensor(ctx, GGML_TYPE_F32, 4, ne); @@ -3310,6 +3340,12 @@ struct ggml_tensor * ggml_mul_mat( result->src[0] = a; result->src[1] = b; + if (scale_weight) { + GGML_ASSERT(ggml_can_repeat(scale_weight, result)); + result->src[2] = scale_weight; + } + + GGML_UNUSED(scale_activations); return result; } @@ -3352,8 +3388,23 @@ struct ggml_tensor * ggml_mul_mat_id( struct ggml_tensor * as, struct ggml_tensor * b, struct ggml_tensor * ids) { + GGML_ASSERT(!ggml_needs_scale_quantized(as->type) && !ggml_needs_scale_quantized(b->type)); + + return ggml_mul_mat_id_ext(ctx, as, b, ids, NULL, NULL); +} + +struct ggml_tensor * ggml_mul_mat_id_ext( + struct ggml_context * ctx, + struct ggml_tensor * as, + struct ggml_tensor * b, + struct ggml_tensor * ids, + struct ggml_tensor * scale_weight, + struct ggml_tensor * scale_activations) { GGML_ASSERT(!ggml_is_transposed(as)); GGML_ASSERT(ids->type == GGML_TYPE_I32); + GGML_ASSERT(!ggml_needs_scale_quantized(b->type)); + GGML_ASSERT(scale_weight == NULL || scale_weight->type == GGML_TYPE_F32); + GGML_ASSERT(scale_activations == NULL || scale_activations->type == GGML_TYPE_F32); GGML_ASSERT(as->ne[3] == 1); // as is 3d (one matrix per expert) GGML_ASSERT(b->ne[3] == 1); // b is 3d @@ -3370,6 +3421,16 @@ struct ggml_tensor * ggml_mul_mat_id( result->src[1] = b; result->src[2] = ids; + if (scale_weight) { + struct ggml_tensor * s = scale_weight; + const bool per_tensor = s->ne[0] == 1 && s->ne[1] == 1 && s->ne[2] == 1 && s->ne[3] == 1; + const bool per_expert = s->ne[0] == as->ne[2] && s->ne[1] == 1 && s->ne[2] == 1 && s->ne[3] == 1; + const bool per_ch_exp = s->ne[0] == as->ne[1] && s->ne[1] == as->ne[2] && s->ne[2] == 1 && s->ne[3] == 1; + GGML_ASSERT(per_tensor || per_expert || per_ch_exp); + result->src[3] = s; + } + + GGML_UNUSED(scale_activations); return result; } diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp index a75f632b8c..f84455a882 100644 --- a/src/llama-graph.cpp +++ b/src/llama-graph.cpp @@ -1515,11 +1515,7 @@ ggml_tensor * llm_graph_context::build_lora_mm( ggml_tensor * w, ggml_tensor * cur, ggml_tensor * w_s) const { - ggml_tensor * res = ggml_mul_mat(ctx0, w, cur); - - if (w_s) { - res = ggml_mul(ctx0, res, w_s); - } + ggml_tensor * res = ggml_mul_mat_ext(ctx0, w, cur, w_s, nullptr); for (const auto & lora : *loras) { llama_adapter_lora_weight * lw = lora.first->get_weight(w); @@ -1547,16 +1543,8 @@ ggml_tensor * llm_graph_context::build_lora_mm_id( ggml_tensor * cur, // ggml_tensor * b ggml_tensor * ids, ggml_tensor * w_s) const { - ggml_tensor * res = ggml_mul_mat_id(ctx0, w, cur, ids); + ggml_tensor * res = ggml_mul_mat_id_ext(ctx0, w, cur, ids, w_s, nullptr); - if (w_s) { - const int64_t n_expert = w_s->ne[0]; - const int64_t n_tokens = cur->ne[2]; - ggml_tensor * s = ggml_reshape_3d(ctx0, w_s, 1, n_expert, 1); - s = ggml_repeat_4d(ctx0, s, 1, n_expert, n_tokens, 1); - 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) { @@ -1708,34 +1696,18 @@ ggml_tensor * llm_graph_context::build_ffn( llm_ffn_op_type type_op, llm_ffn_gate_type type_gate, int il) const { - // NVFP4 support is currently restricted to - // 1) LORA absence (*_s would be applied after LORA residual, which is incorrect) - // 2) bias absense (*_s would be applied after bias addition, which is incorrect) - // TODO: disambiguate LLM-architectural scales (which use *_s) from NVFP4 scale_2 (which also uses *_s currently) - auto has_lora = [this](ggml_tensor * w) { - if (!w) { - return false; - } - for (const auto & lora : *loras) { - if (lora.first->get_weight(w) != nullptr) { - return true; - } - } - return false; - }; - + // Low-precision weights (NVFP4, F8_E4M3) carry a per-tensor dequant scale that must be applied + // in dequant space. Route it into the matmul via the ext API so it lands before bias/LoRA. + // Architectural *_s (on full-precision weights) stays a post-matmul multiply. auto is_scaled_low_precision = [](ggml_tensor * w) { return w && (w->type == GGML_TYPE_NVFP4 || w->type == GGML_TYPE_F8_E4M3); }; - GGML_ASSERT(!up_s || !up_b || !is_scaled_low_precision(up)); - GGML_ASSERT(!gate_s || !gate_b || !is_scaled_low_precision(gate)); - GGML_ASSERT(!down_s || !down_b || !is_scaled_low_precision(down)); - GGML_ASSERT(!up_s || !is_scaled_low_precision(up) || !has_lora(up)); - GGML_ASSERT(!gate_s || !is_scaled_low_precision(gate) || !has_lora(gate)); - GGML_ASSERT(!down_s || !is_scaled_low_precision(down) || !has_lora(down)); + const bool up_derived = is_scaled_low_precision(up); + const bool gate_derived = is_scaled_low_precision(gate); + const bool down_derived = is_scaled_low_precision(down); - ggml_tensor * tmp = up ? build_lora_mm(up, cur) : cur; + ggml_tensor * tmp = up ? build_lora_mm(up, cur, up_derived ? up_s : nullptr) : cur; cb(tmp, "ffn_up", il); if (up_b) { @@ -1743,7 +1715,7 @@ ggml_tensor * llm_graph_context::build_ffn( cb(tmp, "ffn_up_b", il); } - if (up_s) { + if (up_s && !up_derived) { tmp = ggml_mul(ctx0, tmp, up_s); cb(tmp, "ffn_up_s", il); } @@ -1752,12 +1724,12 @@ ggml_tensor * llm_graph_context::build_ffn( switch (type_gate) { case LLM_FFN_SEQ: { - cur = build_lora_mm(gate, tmp); + cur = build_lora_mm(gate, tmp, gate_derived ? gate_s : nullptr); cb(cur, "ffn_gate", il); } break; case LLM_FFN_PAR: { - cur = build_lora_mm(gate, cur); + cur = build_lora_mm(gate, cur, gate_derived ? gate_s : nullptr); cb(cur, "ffn_gate", il); } break; } @@ -1767,7 +1739,7 @@ ggml_tensor * llm_graph_context::build_ffn( cb(cur, "ffn_gate_b", il); } - if (gate_s) { + if (gate_s && !gate_derived) { cur = ggml_mul(ctx0, cur, gate_s); cb(cur, "ffn_gate_s", il); } @@ -1875,7 +1847,7 @@ ggml_tensor * llm_graph_context::build_ffn( } if (down) { - cur = build_lora_mm(down, cur); + cur = build_lora_mm(down, cur, down_derived ? down_s : nullptr); 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); @@ -1890,7 +1862,7 @@ ggml_tensor * llm_graph_context::build_ffn( cur = ggml_add(ctx0, cur, down_b); } - if (down_s) { + if (down_s && !down_derived) { cur = ggml_mul(ctx0, cur, down_s); cb(cur, "ffn_down_s", il); } @@ -2844,15 +2816,10 @@ ggml_tensor * llm_graph_context::build_attn( } if (wo) { + cur = build_lora_mm(wo, cur, wo_s); 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); - if (wo_s) { - cur = ggml_mul(ctx0, cur, wo_s); - } - } else { - cur = build_lora_mm(wo, cur, wo_s); } } @@ -2931,15 +2898,10 @@ ggml_tensor * llm_graph_context::build_attn( cb(cur, "kqv_out", il); if (wo) { + cur = build_lora_mm(wo, cur, wo_s); 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); - if (wo_s) { - cur = ggml_mul(ctx0, cur, wo_s); - } - } else { - cur = build_lora_mm(wo, cur, wo_s); } } diff --git a/src/llama-model-loader.cpp b/src/llama-model-loader.cpp index 161f36c7a3..99461d1ade 100644 --- a/src/llama-model-loader.cpp +++ b/src/llama-model-loader.cpp @@ -919,7 +919,7 @@ const struct ggml_tensor * llama_model_loader::check_tensor_dims( } // checks if the weight tensor can be used with the specified buffer type and device -static bool weight_buft_supported(const llama_hparams & hparams, ggml_tensor * w, ggml_op op, ggml_backend_buffer_type_t buft, ggml_backend_dev_t dev) { +static bool weight_buft_supported(const llama_hparams & hparams, ggml_tensor * w, const ggml_tensor * scale_meta, ggml_op op, ggml_backend_buffer_type_t buft, ggml_backend_dev_t dev) { GGML_ASSERT(w != nullptr); if (op == GGML_OP_NONE) { @@ -948,7 +948,8 @@ static bool weight_buft_supported(const llama_hparams & hparams, ggml_tensor * w case GGML_OP_MUL_MAT: { ggml_tensor * b = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, w->ne[0], 512, w->ne[2], w->ne[3]); - op_tensor = ggml_mul_mat(ctx, w, b); + ggml_tensor * s = scale_meta ? ggml_new_tensor(ctx, GGML_TYPE_F32, GGML_MAX_DIMS, scale_meta->ne) : nullptr; + op_tensor = ggml_mul_mat_ext(ctx, w, b, s, nullptr); } break; case GGML_OP_MUL_MAT_ID: { @@ -957,7 +958,8 @@ static bool weight_buft_supported(const llama_hparams & hparams, ggml_tensor * w GGML_ASSERT(n_ids_used > 0); ggml_tensor * b = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, w->ne[0], n_ids_used, 512); ggml_tensor * ids = ggml_new_tensor_2d(ctx, GGML_TYPE_I32, n_ids_used, 512); - op_tensor = ggml_mul_mat_id(ctx, w, b, ids); + ggml_tensor * s = scale_meta ? ggml_new_tensor(ctx, GGML_TYPE_F32, GGML_MAX_DIMS, scale_meta->ne) : nullptr; + op_tensor = ggml_mul_mat_id_ext(ctx, w, b, ids, s, nullptr); } break; case GGML_OP_ADD: { @@ -1059,12 +1061,12 @@ static bool weight_buft_supported(const llama_hparams & hparams, ggml_tensor * w } // find the first buffer type in the list that can use the tensor -static ggml_backend_buffer_type_t select_weight_buft(const llama_hparams & hparams, ggml_tensor * tensor, ggml_op op, const buft_list_t * buft_list) { +static ggml_backend_buffer_type_t select_weight_buft(const llama_hparams & hparams, ggml_tensor * tensor, const ggml_tensor * scale_meta, ggml_op op, const buft_list_t * buft_list) { GGML_ASSERT(!buft_list->empty()); for (const auto & cur : *buft_list) { ggml_backend_dev_t cur_dev = cur.first; ggml_backend_buffer_type_t cur_buft = cur.second; - if (weight_buft_supported(hparams, tensor, op, cur_buft, cur_dev)) { + if (weight_buft_supported(hparams, tensor, scale_meta, op, cur_buft, cur_dev)) { return cur_buft; } } @@ -1207,6 +1209,12 @@ struct ggml_tensor * llama_model_loader::create_tensor( return lazy_read::buft(); } + const ggml_tensor * scale_meta = nullptr; + if ((op == GGML_OP_MUL_MAT || op == GGML_OP_MUL_MAT_ID) && ggml_needs_scale_quantized(t_meta->type)) { + const std::string scale_name = LLM_TN_IMPL(tn.arch, tn.tensor, "scale", tn.bid, tn.xid).str(); + scale_meta = get_tensor_meta(scale_name.c_str()); + } + // select the buffer type for this tensor const buft_list_t * buft_list; switch (info.layer) { @@ -1234,7 +1242,7 @@ struct ggml_tensor * llama_model_loader::create_tensor( if (std::regex_search(tensor_name, pattern)) { if (overrides->buft == ggml_backend_cpu_buffer_type()) { // when overriding to a CPU buffer, consider the extra buffer types - buft = select_weight_buft(hparams, t_meta, op, buft_list_cpu); + buft = select_weight_buft(hparams, t_meta, scale_meta, op, buft_list_cpu); if (use_mmap) { static std::once_flag once; std::call_once(once, [] { @@ -1255,7 +1263,7 @@ struct ggml_tensor * llama_model_loader::create_tensor( } if (!buft) { - buft = select_weight_buft(hparams, t_meta, op, buft_list); + buft = select_weight_buft(hparams, t_meta, scale_meta, op, buft_list); if (!buft) { throw std::runtime_error(format("failed to find a compatible buffer type for tensor %s", tn.str().c_str())); } diff --git a/src/llama-model.cpp b/src/llama-model.cpp index 14078fde52..9738a39bc6 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -1519,80 +1519,92 @@ bool llama_model_base::load_tensors(llama_model_loader & ml) { // call the per-model loading function load_arch_tensors(ml); - // generic pass: load optional per-tensor/per-expert ".scale" tensors (e.g. NVFP4 scale2) - // this avoids having to add scale loading to every architecture + // generic pass: load optional ".scale" tensors at their stored shape + // (per-tensor {1}, per-channel {n_out}, per-expert {n_expert}, or per-channel-per-expert {n_out, n_expert}) + auto create_scale = [&](const LLM_TN_IMPL & tn_scale) -> ggml_tensor * { + const std::string name = tn_scale.str(); + const ggml_tensor * meta = ml.get_tensor_meta(name.c_str()); + if (meta == nullptr) { + return nullptr; + } + if (meta->ne[1] > 1) { + return create_tensor(tn_scale, {meta->ne[0], meta->ne[1]}, TENSOR_NOT_REQUIRED); + } + return create_tensor(tn_scale, {ggml_nelements(meta)}, TENSOR_NOT_REQUIRED); + }; + for (int i = 0; i < n_layer_all; ++i) { auto & layer = layers[i]; // attention weight scales (per-tensor, shape {1}) if (!layer.wq_s && layer.wq) { - layer.wq_s = create_tensor(tn(LLM_TENSOR_ATTN_Q, "scale", i), {1}, TENSOR_NOT_REQUIRED); + layer.wq_s = create_scale(tn(LLM_TENSOR_ATTN_Q, "scale", i)); } if (!layer.wk_s && layer.wk) { - layer.wk_s = create_tensor(tn(LLM_TENSOR_ATTN_K, "scale", i), {1}, TENSOR_NOT_REQUIRED); + layer.wk_s = create_scale(tn(LLM_TENSOR_ATTN_K, "scale", i)); } if (!layer.wv_s && layer.wv) { - layer.wv_s = create_tensor(tn(LLM_TENSOR_ATTN_V, "scale", i), {1}, TENSOR_NOT_REQUIRED); + layer.wv_s = create_scale(tn(LLM_TENSOR_ATTN_V, "scale", i)); } if (!layer.wo_s && layer.wo) { - layer.wo_s = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "scale", i), {1}, TENSOR_NOT_REQUIRED); + layer.wo_s = create_scale(tn(LLM_TENSOR_ATTN_OUT, "scale", i)); } if (!layer.wqkv_s && layer.wqkv) { - layer.wqkv_s = create_tensor(tn(LLM_TENSOR_ATTN_QKV, "scale", i), {1}, TENSOR_NOT_REQUIRED); + layer.wqkv_s = create_scale(tn(LLM_TENSOR_ATTN_QKV, "scale", i)); } if (!layer.wqkv_gate_s && layer.wqkv_gate) { - layer.wqkv_gate_s = create_tensor(tn(LLM_TENSOR_ATTN_GATE, "scale", i), {1}, TENSOR_NOT_REQUIRED); + layer.wqkv_gate_s = create_scale(tn(LLM_TENSOR_ATTN_GATE, "scale", i)); } // dense FFN weight scales (per-tensor, shape {1}) if (!layer.ffn_gate_s && layer.ffn_gate) { - layer.ffn_gate_s = create_tensor(tn(LLM_TENSOR_FFN_GATE, "scale", i), {1}, TENSOR_NOT_REQUIRED); + layer.ffn_gate_s = create_scale(tn(LLM_TENSOR_FFN_GATE, "scale", i)); } if (!layer.ffn_down_s && layer.ffn_down) { - layer.ffn_down_s = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "scale", i), {1}, TENSOR_NOT_REQUIRED); + layer.ffn_down_s = create_scale(tn(LLM_TENSOR_FFN_DOWN, "scale", i)); } if (!layer.ffn_up_s && layer.ffn_up) { - layer.ffn_up_s = create_tensor(tn(LLM_TENSOR_FFN_UP, "scale", i), {1}, TENSOR_NOT_REQUIRED); + layer.ffn_up_s = create_scale(tn(LLM_TENSOR_FFN_UP, "scale", i)); } if (!layer.ffn_gate_shexp_s && layer.ffn_gate_shexp) { - layer.ffn_gate_shexp_s = create_tensor(tn(LLM_TENSOR_FFN_GATE_SHEXP, "scale", i), {1}, TENSOR_NOT_REQUIRED); + layer.ffn_gate_shexp_s = create_scale(tn(LLM_TENSOR_FFN_GATE_SHEXP, "scale", i)); } if (!layer.ffn_down_shexp_s && layer.ffn_down_shexp) { - layer.ffn_down_shexp_s = create_tensor(tn(LLM_TENSOR_FFN_DOWN_SHEXP, "scale", i), {1}, TENSOR_NOT_REQUIRED); + layer.ffn_down_shexp_s = create_scale(tn(LLM_TENSOR_FFN_DOWN_SHEXP, "scale", i)); } if (!layer.ffn_up_shexp_s && layer.ffn_up_shexp) { - layer.ffn_up_shexp_s = create_tensor(tn(LLM_TENSOR_FFN_UP_SHEXP, "scale", i), {1}, TENSOR_NOT_REQUIRED); + layer.ffn_up_shexp_s = create_scale(tn(LLM_TENSOR_FFN_UP_SHEXP, "scale", i)); } // MoE expert weight scales (per-expert, shape {n_expert}) if (!layer.ffn_gate_exps_s && layer.ffn_gate_exps) { - layer.ffn_gate_exps_s = create_tensor(tn(LLM_TENSOR_FFN_GATE_EXPS, "scale", i), {n_expert}, TENSOR_NOT_REQUIRED); + layer.ffn_gate_exps_s = create_scale(tn(LLM_TENSOR_FFN_GATE_EXPS, "scale", i)); } if (!layer.ffn_down_exps_s && layer.ffn_down_exps) { - layer.ffn_down_exps_s = create_tensor(tn(LLM_TENSOR_FFN_DOWN_EXPS, "scale", i), {n_expert}, TENSOR_NOT_REQUIRED); + layer.ffn_down_exps_s = create_scale(tn(LLM_TENSOR_FFN_DOWN_EXPS, "scale", i)); } if (!layer.ffn_up_exps_s && layer.ffn_up_exps) { - layer.ffn_up_exps_s = create_tensor(tn(LLM_TENSOR_FFN_UP_EXPS, "scale", i), {n_expert}, TENSOR_NOT_REQUIRED); + layer.ffn_up_exps_s = create_scale(tn(LLM_TENSOR_FFN_UP_EXPS, "scale", i)); } // recurrent / linear-attention weight scales (per-tensor, shape {1}) if (!layer.ssm_in_s && layer.ssm_in) { - layer.ssm_in_s = create_tensor(tn(LLM_TENSOR_SSM_IN, "scale", i), {1}, TENSOR_NOT_REQUIRED); + layer.ssm_in_s = create_scale(tn(LLM_TENSOR_SSM_IN, "scale", i)); } if (!layer.ssm_out_s && layer.ssm_out) { - layer.ssm_out_s = create_tensor(tn(LLM_TENSOR_SSM_OUT, "scale", i), {1}, TENSOR_NOT_REQUIRED); + layer.ssm_out_s = create_scale(tn(LLM_TENSOR_SSM_OUT, "scale", i)); } if (!layer.ssm_alpha_s && layer.ssm_alpha) { - layer.ssm_alpha_s = create_tensor(tn(LLM_TENSOR_SSM_ALPHA, "scale", i), {1}, TENSOR_NOT_REQUIRED); + layer.ssm_alpha_s = create_scale(tn(LLM_TENSOR_SSM_ALPHA, "scale", i)); } if (!layer.ssm_beta_s && layer.ssm_beta) { - layer.ssm_beta_s = create_tensor(tn(LLM_TENSOR_SSM_BETA, "scale", i), {1}, TENSOR_NOT_REQUIRED); + layer.ssm_beta_s = create_scale(tn(LLM_TENSOR_SSM_BETA, "scale", i)); } if (!layer.nextn.eh_proj_s && layer.nextn.eh_proj) { - layer.nextn.eh_proj_s = create_tensor(tn(LLM_TENSOR_NEXTN_EH_PROJ, "scale", i), {1}, TENSOR_NOT_REQUIRED); + layer.nextn.eh_proj_s = create_scale(tn(LLM_TENSOR_NEXTN_EH_PROJ, "scale", i)); } if (!layer.nextn.shared_head_head_s && layer.nextn.shared_head_head) { - layer.nextn.shared_head_head_s = create_tensor(tn(LLM_TENSOR_NEXTN_SHARED_HEAD_HEAD, "scale", i), {1}, TENSOR_NOT_REQUIRED); + layer.nextn.shared_head_head_s = create_scale(tn(LLM_TENSOR_NEXTN_SHARED_HEAD_HEAD, "scale", i)); } // input scales @@ -1664,7 +1676,7 @@ bool llama_model_base::load_tensors(llama_model_loader & ml) { if (output && (output->type == GGML_TYPE_NVFP4 || output->type == GGML_TYPE_F8_E4M3)) { // weight scale if (!output_s) { - output_s = create_tensor(tn(LLM_TENSOR_OUTPUT, "scale"), {1}, TENSOR_NOT_REQUIRED); + output_s = create_scale(tn(LLM_TENSOR_OUTPUT, "scale")); } // input scale if (!output_in_s) { diff --git a/src/models/granite-hybrid.cpp b/src/models/granite-hybrid.cpp index 8a8f7e19ff..b1840127e1 100644 --- a/src/models/granite-hybrid.cpp +++ b/src/models/granite-hybrid.cpp @@ -246,9 +246,9 @@ ggml_tensor * llama_model_granite_hybrid::graph::build_layer_ffn(ggml_tensor * cb(cur, "ffn_norm", il); cur = build_ffn(cur, - model.layers[il].ffn_up, model.layers[il].ffn_up_b, NULL, - model.layers[il].ffn_gate, model.layers[il].ffn_gate_b, NULL, - model.layers[il].ffn_down, model.layers[il].ffn_down_b, NULL, + model.layers[il].ffn_up, model.layers[il].ffn_up_b, model.layers[il].ffn_up_s, + model.layers[il].ffn_gate, model.layers[il].ffn_gate_b, model.layers[il].ffn_gate_s, + model.layers[il].ffn_down, model.layers[il].ffn_down_b, model.layers[il].ffn_down_s, NULL, LLM_FFN_SILU, LLM_FFN_PAR, il); cb(cur, "ffn_out", il); @@ -268,16 +268,20 @@ ggml_tensor * llama_model_granite_hybrid::graph::build_layer_ffn(ggml_tensor * LLM_FFN_SILU, true, hparams.expert_weights_scale, LLAMA_EXPERT_GATING_FUNC_TYPE_SOFTMAX, - il); + il, + nullptr, nullptr, + model.layers[il].ffn_up_exps_s, + model.layers[il].ffn_gate_exps_s, + model.layers[il].ffn_down_exps_s); cb(moe_out, "ffn_moe_out", il); // For Granite MoE Shared if (hparams.n_ff_shexp > 0) { ggml_tensor * ffn_shexp = build_ffn(cur, - model.layers[il].ffn_up_shexp, NULL, NULL, - model.layers[il].ffn_gate_shexp, NULL, NULL, - model.layers[il].ffn_down_shexp, NULL, NULL, + model.layers[il].ffn_up_shexp, NULL, model.layers[il].ffn_up_shexp_s, + model.layers[il].ffn_gate_shexp, NULL, model.layers[il].ffn_gate_shexp_s, + model.layers[il].ffn_down_shexp, NULL, model.layers[il].ffn_down_shexp_s, NULL, LLM_FFN_SILU, LLM_FFN_PAR, il); cb(ffn_shexp, "ffn_shexp", il); diff --git a/src/models/granite.cpp b/src/models/granite.cpp index 9e9f97e94d..35b4215b14 100644 --- a/src/models/granite.cpp +++ b/src/models/granite.cpp @@ -253,9 +253,9 @@ ggml_tensor * llama_model_granite::graph::build_layer_ffn( cb(cur, "ffn_norm", il); cur = build_ffn(cur, - model.layers[il].ffn_up, model.layers[il].ffn_up_b, NULL, - model.layers[il].ffn_gate, model.layers[il].ffn_gate_b, NULL, - model.layers[il].ffn_down, model.layers[il].ffn_down_b, NULL, + model.layers[il].ffn_up, model.layers[il].ffn_up_b, model.layers[il].ffn_up_s, + model.layers[il].ffn_gate, model.layers[il].ffn_gate_b, model.layers[il].ffn_gate_s, + model.layers[il].ffn_down, model.layers[il].ffn_down_b, model.layers[il].ffn_down_s, NULL, LLM_FFN_SILU, LLM_FFN_PAR, il); cb(cur, "ffn_out", il); @@ -277,15 +277,19 @@ ggml_tensor * llama_model_granite::graph::build_layer_ffn( LLM_FFN_SILU, true, hparams.expert_weights_scale, LLAMA_EXPERT_GATING_FUNC_TYPE_SOFTMAX, - il); + il, + nullptr, nullptr, + model.layers[il].ffn_up_exps_s, + model.layers[il].ffn_gate_exps_s, + model.layers[il].ffn_down_exps_s); cb(moe_out, "ffn_moe_out", il); // For Granite MoE Shared if (hparams.n_ff_shexp > 0) { ggml_tensor * ffn_shexp = build_ffn(cur, - model.layers[il].ffn_up_shexp, NULL, NULL, - model.layers[il].ffn_gate_shexp, NULL, NULL, - model.layers[il].ffn_down_shexp, NULL, NULL, + model.layers[il].ffn_up_shexp, NULL, model.layers[il].ffn_up_shexp_s, + model.layers[il].ffn_gate_shexp, NULL, model.layers[il].ffn_gate_shexp_s, + model.layers[il].ffn_down_shexp, NULL, model.layers[il].ffn_down_shexp_s, NULL, LLM_FFN_SILU, LLM_FFN_PAR, il); cb(ffn_shexp, "ffn_shexp", il); diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index be6ee77161..5efd5a7e6f 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -4625,9 +4625,10 @@ struct test_mul_mat : public test_case { const int64_t k_v; // size of k in memory, resulting in a non-contiguous view for k_v > k, no view for k_v == 0 const uint32_t o; // number of outputs const bool src_overlap; // a and b are overlapping views of the same tensor + const int scale_gran; // weight scale granularity for non-needs_scale types: 0=none, 1=per-tensor, 2=per-output-channel std::string vars() override { - return VARS_TO_STR11(type_a, type_b, m, n, k, bs, nr, per, k_v, o, src_overlap); + return VARS_TO_STR12(type_a, type_b, m, n, k, bs, nr, per, k_v, o, src_overlap, scale_gran); } double max_nmse_err() override { @@ -4659,8 +4660,8 @@ struct test_mul_mat : public test_case { std::array bs = {10, 10}, std::array nr = {2, 2}, std::array per = {0, 1, 2, 3}, - int64_t k_v = 0, uint32_t o = 1, bool src_overlap = false) - : type_a(type_a), type_b(type_b), m(m), n(n), k(k), bs(bs), nr(nr), per(per), k_v(k_v), o(o), src_overlap(src_overlap) {} + int64_t k_v = 0, uint32_t o = 1, bool src_overlap = false, int scale_gran = 0) + : type_a(type_a), type_b(type_b), m(m), n(n), k(k), bs(bs), nr(nr), per(per), k_v(k_v), o(o), src_overlap(src_overlap), scale_gran(scale_gran) {} ggml_tensor * build_graph(ggml_context * ctx) override { // C^T = A * B^T: (k, m) * (k, n) => (m, n) @@ -4726,10 +4727,20 @@ struct test_mul_mat : public test_case { ggml_set_name(b, "b"); } - ggml_tensor * out = ggml_mul_mat(ctx, a, b); + ggml_tensor * s = ggml_needs_scale_quantized(type_a) ? ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 1) : nullptr; + if (!s && scale_gran == 1) { + s = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 1); + } else if (!s && scale_gran == 2) { + s = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, m); + } + if (s) { + ggml_set_name(s, "s"); + } + + ggml_tensor * out = ggml_mul_mat_ext(ctx, a, b, s, nullptr); ggml_set_name(out, "out"); for (uint32_t i = 1; i < o; ++i) { - ggml_tensor * out2 = ggml_mul_mat(ctx, a, b); + ggml_tensor * out2 = ggml_mul_mat_ext(ctx, a, b, s, nullptr); ggml_set_name(out2, "out2"); out = ggml_add(ctx, out, out2); } @@ -4829,9 +4840,10 @@ struct test_mul_mat_id : public test_case { const int64_t m; const int64_t n; const int64_t k; + const int scale_gran; // weight scale granularity for non-needs_scale types: 0=none, 1=per-expert scalar, 2=per-channel-per-expert std::string vars() override { - return VARS_TO_STR8(type_a, type_b, n_mats, n_used, b, m, n, k); + return VARS_TO_STR9(type_a, type_b, n_mats, n_used, b, m, n, k, scale_gran); } double max_nmse_err() override { @@ -4854,11 +4866,16 @@ struct test_mul_mat_id : public test_case { return 2 * m * k * n * n_used; } + std::string op_desc(ggml_tensor * t) override { + GGML_UNUSED(t); + return ggml_op_name(GGML_OP_MUL_MAT_ID); + } + test_mul_mat_id(ggml_type type_a = GGML_TYPE_F32, 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 = 32) + int64_t m = 32, int64_t n = 32, int64_t k = 32, int scale_gran = 0) : type_a(type_a), type_b(type_b), n_mats(n_mats), n_used(n_used), b(b), - m(m), n(n), k(k) { + m(m), n(n), k(k), scale_gran(scale_gran) { GGML_ASSERT(n_used <= n_mats); } @@ -4877,7 +4894,17 @@ struct test_mul_mat_id : public test_case { ggml_tensor * b = ggml_new_tensor_3d(ctx, type_b, k, this->b ? 1 : n_used, n); ggml_set_name(b, "b"); - ggml_tensor * out = ggml_mul_mat_id(ctx, as, b, ids); + ggml_tensor * s = ggml_needs_scale_quantized(type_a) ? ggml_new_tensor_1d(ctx, GGML_TYPE_F32, n_mats) : nullptr; + if (!s && scale_gran == 1) { + s = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, n_mats); + } else if (!s && scale_gran == 2) { + s = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, m, n_mats); + } + if (s) { + ggml_set_name(s, "s"); + } + + ggml_tensor * out = ggml_mul_mat_id_ext(ctx, as, b, ids, s, nullptr); ggml_set_name(out, "out"); return out; @@ -6670,7 +6697,7 @@ struct test_mul_mat_vec_fusion : public test_case { } bool run_whole_graph() override { return true; } - bool use_weight_context() override { return use_id && with_lane_scale; } + bool use_weight_context() override { return use_id && (with_lane_scale || ggml_needs_scale_quantized(type)); } ggml_tensor * build_gate(ggml_context * ctx, ggml_tensor * ffn_gate, ggml_tensor * ffn_up) { ggml_tensor * out = nullptr; @@ -6689,20 +6716,6 @@ struct test_mul_mat_vec_fusion : public test_case { return out; } - ggml_tensor * build_lane_scale_dense(ggml_context * ctx, ggml_tensor * out) { - ggml_tensor * scale = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 1); - return ggml_mul(ctx, out, scale); - } - - ggml_tensor * build_lane_scale_id(ggml_context * ctx, ggml_context * ctx_weights, ggml_tensor * out, ggml_tensor * ids) { - GGML_ASSERT(ctx_weights); - ggml_tensor * scale = ggml_new_tensor_1d(ctx_weights, GGML_TYPE_F32, n_mats); - ggml_tensor * s = ggml_reshape_3d(ctx, scale, 1, n_mats, 1); - s = ggml_repeat_4d(ctx, s, 1, n_mats, m, 1); - s = ggml_get_rows(ctx, s, ids); - return ggml_mul(ctx, out, s); - } - ggml_tensor * build_graph(ggml_context * ctx) override { GGML_ASSERT(!use_weight_context()); return build_graph(ctx, nullptr); @@ -6720,10 +6733,9 @@ struct test_mul_mat_vec_fusion : public test_case { ggml_tensor * up = ggml_new_tensor(ctx, type, 4, ne0.data()); auto build_lane_up = [&]() { - ggml_tensor * ffn_up = ggml_mul_mat(ctx, up, cur); - if (with_lane_scale) { - ffn_up = build_lane_scale_dense(ctx, ffn_up); - } + ggml_tensor * scale = (with_lane_scale || ggml_needs_scale_quantized(up->type)) + ? ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 1) : nullptr; + ggml_tensor * ffn_up = ggml_mul_mat_ext(ctx, up, cur, scale, nullptr); if (with_bias) { std::array bias_ne = { ffn_up->ne[0], 1, channels, samples }; ggml_tensor * up_bias = ggml_new_tensor(ctx, GGML_TYPE_F32, 4, bias_ne.data()); @@ -6733,10 +6745,9 @@ struct test_mul_mat_vec_fusion : public test_case { }; auto build_lane_gate = [&]() { - ggml_tensor * ffn_gate = ggml_mul_mat(ctx, gate, cur); - if (with_lane_scale) { - ffn_gate = build_lane_scale_dense(ctx, ffn_gate); - } + ggml_tensor * scale = (with_lane_scale || ggml_needs_scale_quantized(gate->type)) + ? ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 1) : nullptr; + ggml_tensor * ffn_gate = ggml_mul_mat_ext(ctx, gate, cur, scale, nullptr); if (with_bias) { std::array bias_ne = { ffn_gate->ne[0], 1, channels, samples }; ggml_tensor * gate_bias = ggml_new_tensor(ctx, GGML_TYPE_F32, 4, bias_ne.data()); @@ -6769,10 +6780,9 @@ struct test_mul_mat_vec_fusion : public test_case { ggml_set_name(cur, "cur"); auto build_lane_up = [&]() { - ggml_tensor * ffn_up = ggml_mul_mat_id(ctx, ups, cur, ids); - if (with_lane_scale) { - ffn_up = build_lane_scale_id(ctx, ctx_weights, ffn_up, ids); - } + ggml_tensor * scale = (with_lane_scale || ggml_needs_scale_quantized(ups->type)) + ? ggml_new_tensor_1d(ctx_weights, GGML_TYPE_F32, n_mats) : nullptr; + ggml_tensor * ffn_up = ggml_mul_mat_id_ext(ctx, ups, cur, ids, scale, nullptr); if (with_bias) { ggml_tensor * up_bias_param = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, ffn_up->ne[0], n_mats); ffn_up = ggml_add_id(ctx, ffn_up, up_bias_param, ids); @@ -6781,10 +6791,9 @@ struct test_mul_mat_vec_fusion : public test_case { }; auto build_lane_gate = [&]() { - ggml_tensor * ffn_gate = ggml_mul_mat_id(ctx, gates, cur, ids); - if (with_lane_scale) { - ffn_gate = build_lane_scale_id(ctx, ctx_weights, ffn_gate, ids); - } + ggml_tensor * scale = (with_lane_scale || ggml_needs_scale_quantized(gates->type)) + ? ggml_new_tensor_1d(ctx_weights, GGML_TYPE_F32, n_mats) : nullptr; + ggml_tensor * ffn_gate = ggml_mul_mat_id_ext(ctx, gates, cur, ids, scale, nullptr); if (with_bias) { ggml_tensor * gate_bias_param = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, ffn_gate->ne[0], n_mats); ffn_gate = ggml_add_id(ctx, ffn_gate, gate_bias_param, ids); @@ -9478,6 +9487,18 @@ 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) + // FP8 weight-scale coverage (Qwen3.5-35B-A3B-FP8 shapes): dense per-channel + MoE per-channel-per-expert + for (int scale_gran : {1, 2}) { + // dense attention/shexp: (k, m) + test_cases.emplace_back(new test_mul_mat(GGML_TYPE_F8_E4M3, GGML_TYPE_F32, 8192, 1, 2048, {1, 1}, {1, 1}, {0,1,2,3}, 0, 1, false, scale_gran)); + test_cases.emplace_back(new test_mul_mat(GGML_TYPE_F8_E4M3, GGML_TYPE_F32, 8192, 16, 2048, {1, 1}, {1, 1}, {0,1,2,3}, 0, 1, false, scale_gran)); + // MoE experts: (type, n_mats, n_used, b, m, n, k) + test_cases.emplace_back(new test_mul_mat_id(GGML_TYPE_F8_E4M3, GGML_TYPE_F32, 256, 8, false, 512, 1, 2048, scale_gran)); // gate/up decode + test_cases.emplace_back(new test_mul_mat_id(GGML_TYPE_F8_E4M3, GGML_TYPE_F32, 256, 8, false, 512, 16, 2048, scale_gran)); // gate/up prompt + test_cases.emplace_back(new test_mul_mat_id(GGML_TYPE_F8_E4M3, GGML_TYPE_F32, 256, 8, false, 2048, 1, 512, scale_gran)); // down decode + test_cases.emplace_back(new test_mul_mat_id(GGML_TYPE_F8_E4M3, GGML_TYPE_F32, 256, 8, false, 2048, 16, 512, scale_gran)); // down prompt + } + #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}));