diff --git a/ggml/src/ggml-cuda/mmq-load-tiles.cuh b/ggml/src/ggml-cuda/mmq-load-tiles.cuh index 7f00bad94..e19f4f24d 100644 --- a/ggml/src/ggml-cuda/mmq-load-tiles.cuh +++ b/ggml/src/ggml-cuda/mmq-load-tiles.cuh @@ -1631,16 +1631,16 @@ template static __device__ __forceinline_ template static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_mxfp4_fp4( const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) { constexpr int warp_size = ggml_cuda_get_physical_warp_size(); - constexpr int nwarps = ggml_cuda_mmq_get_nthreads(type, J, fallback) / warp_size; - constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback); - constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback); + constexpr int nwarps = ggml_cuda_mmq_get_nthreads(type, J, fallback, GGML_PREC_Q4) / warp_size; + constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback, GGML_PREC_Q4); + constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback, GGML_PREC_Q4); int * x_qs = (int *) x_tile; uint32_t * x_sc = (uint32_t *) (x_qs + 2 * MMQ_TILE_NE_K); const int txi = threadIdx.x; - constexpr int iter_k = ggml_cuda_mmq_get_K_vram(type, J, fallback); + constexpr int iter_k = ggml_cuda_mmq_get_K_vram(type, J, fallback, GGML_PREC_Q4); constexpr int threads_per_row = iter_k / QK_MXFP4; // each thread processes 1 block constexpr int rows_per_warp = warp_size / threads_per_row; @@ -1670,12 +1670,12 @@ template static __device__ __forceinline_ } } -template static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_nvfp4( +template static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_nvfp4( const char * __restrict__ x, int * __restrict__ x_tile, const int kb0, const int i_max, const int stride) { constexpr int warp_size = ggml_cuda_get_physical_warp_size(); - constexpr int nwarps = ggml_cuda_mmq_get_nthreads(type, J, fallback) / warp_size; - constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback); - constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback); + constexpr int nwarps = ggml_cuda_mmq_get_nthreads(type, J, fallback, prec_src1) / warp_size; + constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback, prec_src1); + constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback, prec_src1); #if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) int * x_qs = (int *) x_tile; @@ -1729,12 +1729,12 @@ template static __device__ __forceinline_ template static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_nvfp4_nvfp4( const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) { constexpr int warp_size = ggml_cuda_get_physical_warp_size(); - constexpr int nwarps = ggml_cuda_mmq_get_nthreads(type, J, fallback) / warp_size; - constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback); - constexpr int iter_k = ggml_cuda_mmq_get_K_vram(type, J, fallback); + constexpr int nwarps = ggml_cuda_mmq_get_nthreads(type, J, fallback, GGML_PREC_Q4) / warp_size; + constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback, GGML_PREC_Q4); + constexpr int iter_k = ggml_cuda_mmq_get_K_vram(type, J, fallback, GGML_PREC_Q4); constexpr int threads_per_row = iter_k / QK_NVFP4; // each thread processes 1 block constexpr int rows_per_warp = warp_size / threads_per_row; - constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback); + constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback, GGML_PREC_Q4); uint32_t * x_u32 = (uint32_t *) x_tile; diff --git a/ggml/src/ggml-cuda/mmq-vec-dot.cuh b/ggml/src/ggml-cuda/mmq-vec-dot.cuh index 4d1c398fc..4ca6542d8 100644 --- a/ggml/src/ggml-cuda/mmq-vec-dot.cuh +++ b/ggml/src/ggml-cuda/mmq-vec-dot.cuh @@ -474,7 +474,7 @@ template static __device__ __forceinline_ } // Used for Q3_K, IQ2_S, and IQ2_XS: -template static __device__ __forceinline__ void ggml_cuda_mmq_vec_dot_q8_0_16_q8_1_mma( +template static __device__ __forceinline__ void ggml_cuda_mmq_vec_dot_q8_0_16_q8_1_mma( const int * __restrict__ x, const int * __restrict__ y, float * __restrict__ sum, const int k00) { #if defined(AMD_MFMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) constexpr data_layout input_layout = get_input_data_layout(); @@ -482,7 +482,7 @@ template static __device__ __forceinline_ typedef tile<16, 4, int, input_layout> tile_B; typedef tile<16, 16, int, DATA_LAYOUT_J_MAJOR> tile_C; - constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback); + constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback, prec_src1); constexpr int rows_per_warp = ggml_cuda_mmq_get_rows_per_warp(type, J, fallback); constexpr int ntx = rows_per_warp/tile_C::I; // Number of x minitiles per warp. @@ -532,7 +532,7 @@ template static __device__ __forceinline_ typedef tile< 8, 4, int> tile_B; typedef tile<16, 8, int> tile_C; - constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback); + constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback, prec_src1); constexpr int rows_per_warp = ggml_cuda_mmq_get_rows_per_warp(type, J, fallback); constexpr int ntx = rows_per_warp/tile_C::I; // Number of x minitiles per warp. @@ -1180,7 +1180,7 @@ template static __device__ __forceinline_ typedef tile<8, 8, int> tile_B; typedef tile<16, 8, float> tile_C; - constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback); + constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback, GGML_PREC_Q4); constexpr int rows_per_warp = ggml_cuda_mmq_get_rows_per_warp(type, J, fallback); constexpr int ntx = rows_per_warp / tile_C::I; constexpr int nfrags = MMQ_TILE_NE_K / tile_A::J; diff --git a/ggml/src/ggml-cuda/mmq.cu b/ggml/src/ggml-cuda/mmq.cu index f68b3df18..3e1721bbe 100644 --- a/ggml/src/ggml-cuda/mmq.cu +++ b/ggml/src/ggml-cuda/mmq.cu @@ -5,7 +5,7 @@ #include -static void ggml_cuda_mul_mat_q_switch_type(ggml_backend_cuda_context & ctx, const mmq_args & args, cudaStream_t stream) { +static void ggml_cuda_mul_mat_q_switch_type(ggml_backend_cuda_context & ctx, const mmq_args & args, cudaStream_t stream, const ggml_prec prec_src1) { switch (args.type_x) { case GGML_TYPE_Q1_0: mul_mat_q_case(ctx, args, stream); @@ -71,9 +71,18 @@ static void ggml_cuda_mul_mat_q_switch_type(ggml_backend_cuda_context & ctx, con break; // ----------------------------------------------------------------------- case GGML_TYPE_MXFP4: + // src1 at Q4 uses the native FP4 instructions, which are Blackwell-only + if (prec_src1 == GGML_PREC_Q4) { + mul_mat_q_case(ctx, args, stream); + break; + } mul_mat_q_case(ctx, args, stream); break; case GGML_TYPE_NVFP4: + if (prec_src1 == GGML_PREC_Q4) { + mul_mat_q_case(ctx, args, stream); + break; + } mul_mat_q_case(ctx, args, stream); break; default: @@ -82,6 +91,47 @@ static void ggml_cuda_mul_mat_q_switch_type(ggml_backend_cuda_context & ctx, con } } +// overrides the src1 precision requested by the graph, "auto" keeps the requested one +static ggml_prec ggml_cuda_mmq_get_prec_env() { + const char * env_c = getenv("GGML_CUDA_MMQ_PREC"); + if (env_c == nullptr) { + return GGML_PREC_UNDEFINED; + } + std::string env_cpp = env_c; + for (char & c : env_cpp) { + c = std::tolower(c); + } + if (env_cpp == "q4") { + return GGML_PREC_Q4; + } + if (env_cpp == "q8") { + return GGML_PREC_Q8; + } + if (env_cpp != "auto") { + GGML_LOG_WARN("%s: Unknown value for GGML_CUDA_MMQ_PREC: '%s'. Available: 'q4', 'q8', 'auto'.\n", __func__, env_cpp.c_str()); + } + return GGML_PREC_UNDEFINED; +} + +// src1 is quantized to Q8_1 unless the FP4 types can use 4-bit activations, in which case they +// default to the native W4A4 instructions on Blackwell. +static ggml_prec ggml_cuda_mmq_get_prec_src1(const ggml_tensor * src0, const ggml_tensor * dst, const int cc) { + static const ggml_prec prec_env = ggml_cuda_mmq_get_prec_env(); + + ggml_prec prec = prec_env; + if (prec == GGML_PREC_UNDEFINED) { + prec = (ggml_prec) ggml_get_op_params_i32(dst, 3); + } + + // Q4 only for the FP4 types on Blackwell + GGML_ASSERT(prec == GGML_PREC_UNDEFINED || prec == GGML_PREC_Q8 || prec == GGML_PREC_Q4); + const bool can_use_q4 = (src0->type == GGML_TYPE_NVFP4 || src0->type == GGML_TYPE_MXFP4) && blackwell_mma_available(cc); + if (prec == GGML_PREC_Q8 || !can_use_q4) { + return GGML_PREC_Q8; + } + return GGML_PREC_Q4; +} + void ggml_cuda_mul_mat_q( ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * ids, ggml_tensor * dst) { GGML_ASSERT( src1->type == GGML_TYPE_F32); @@ -128,7 +178,9 @@ void ggml_cuda_mul_mat_q( const bool fallback = ne01 % 128 != 0; - const bool use_native_fp4 = blackwell_mma_available(cc) && (src0->type == GGML_TYPE_MXFP4 || src0->type == GGML_TYPE_NVFP4); + const ggml_prec prec_src1 = ggml_cuda_mmq_get_prec_src1(src0, dst, cc); + + const bool use_native_fp4 = prec_src1 == GGML_PREC_Q4; const size_t y_block_size = use_native_fp4 ? sizeof(block_fp4_mmq) : sizeof(block_q8_1_mmq); const size_t y_values_per_block = use_native_fp4 ? QK_FP4_MMQ : QK8_1_MMQ; @@ -172,7 +224,7 @@ void ggml_cuda_mul_mat_q( ne02, ne12, s02, s12, s2, ne03, ne13, s03, s13, s3, ne1, ne1}; - ggml_cuda_mul_mat_q_switch_type(ctx, args, stream); + ggml_cuda_mul_mat_q_switch_type(ctx, args, stream, prec_src1); return; } @@ -260,7 +312,7 @@ void ggml_cuda_mul_mat_q( ne03, ne13, s03, s13, s3, ne12, ncols_opt}; - ggml_cuda_mul_mat_q_switch_type(ctx, args, stream); + ggml_cuda_mul_mat_q_switch_type(ctx, args, stream, prec_src1); } bool ggml_cuda_should_use_mmq(enum ggml_type type, int cc, int64_t ne11, int64_t n_experts) { diff --git a/ggml/src/ggml-cuda/mmq.cuh b/ggml/src/ggml-cuda/mmq.cuh index 6923f3510..4b50d1dc3 100644 --- a/ggml/src/ggml-cuda/mmq.cuh +++ b/ggml/src/ggml-cuda/mmq.cuh @@ -227,7 +227,7 @@ struct ggml_cuda_mmq_config { #undef CASE -static __host__ ggml_cuda_mmq_config ggml_cuda_mmq_get_config(const ggml_type type, const int J, const bool fallback, const int cc) { +static __host__ ggml_cuda_mmq_config ggml_cuda_mmq_get_config(const ggml_type type, const int J, const bool fallback, const int cc, const ggml_prec prec_src1 = GGML_PREC_Q8) { if (GGML_CUDA_CC_IS_AMD(cc)) { if (GGML_CUDA_CC_IS_GCN(cc)) { return ggml_cuda_mmq_get_config_gcn(type, J, fallback); @@ -247,6 +247,10 @@ static __host__ ggml_cuda_mmq_config ggml_cuda_mmq_get_config(const ggml_type ty return ggml_cuda_mmq_get_config_rdna2(type, J, fallback); } if (blackwell_mma_available(cc)) { + // only src1 at Q4 uses the native FP4 config, higher precisions keep src1 at Q8_1 + if (prec_src1 != GGML_PREC_Q4 && (type == GGML_TYPE_NVFP4 || type == GGML_TYPE_MXFP4)) { + return ggml_cuda_mmq_get_config_ampere(type, J, fallback); + } return ggml_cuda_mmq_get_config_blackwell(type, J, fallback); } if (ggml_cuda_highest_compiled_arch(cc) >= GGML_CUDA_CC_VOLTA) { @@ -258,7 +262,7 @@ static __host__ ggml_cuda_mmq_config ggml_cuda_mmq_get_config(const ggml_type ty return ggml_cuda_mmq_get_config_pascal_older(type, J, fallback); } -static constexpr __device__ ggml_cuda_mmq_config ggml_cuda_mmq_get_config(ggml_type type, int J, bool fallback) { +static constexpr __device__ ggml_cuda_mmq_config ggml_cuda_mmq_get_config(ggml_type type, int J, bool fallback, ggml_prec prec_src1 = GGML_PREC_Q8) { #ifdef GGML_USE_HIP #ifdef GCN return ggml_cuda_mmq_get_config_gcn(type, J, fallback); @@ -275,6 +279,10 @@ static constexpr __device__ ggml_cuda_mmq_config ggml_cuda_mmq_get_config(ggml_t #endif // CDNA #else #ifdef BLACKWELL_MMA_AVAILABLE + // only src1 at Q4 uses the native FP4 config, higher precisions keep src1 at Q8_1 + if (prec_src1 != GGML_PREC_Q4 && (type == GGML_TYPE_NVFP4 || type == GGML_TYPE_MXFP4)) { + return ggml_cuda_mmq_get_config_ampere(type, J, fallback); + } return ggml_cuda_mmq_get_config_blackwell(type, J, fallback); #elif __CUDA_ARCH__ >= GGML_CUDA_CC_VOLTA return ggml_cuda_mmq_get_config_ampere(type, J, fallback); @@ -284,79 +292,71 @@ static constexpr __device__ ggml_cuda_mmq_config ggml_cuda_mmq_get_config(ggml_t return ggml_cuda_mmq_get_config_pascal_older(type, J, fallback); #endif // BLACKWELL_MMA_AVAILABLE #endif // GGML_USE_HIP - GGML_UNUSED_VARS(type, J, fallback); + GGML_UNUSED_VARS(type, J, fallback, prec_src1); } static __host__ int ggml_cuda_mmq_get_type(const ggml_type type, const int J, const bool fallback, const int cc) { return ggml_cuda_mmq_get_config(type, J, fallback, cc).type; } -static constexpr __device__ int ggml_cuda_mmq_get_type(ggml_type type, int J, bool fallback) { - return ggml_cuda_mmq_get_config(type, J, fallback).type; +static constexpr __device__ int ggml_cuda_mmq_get_type(ggml_type type, int J, bool fallback, ggml_prec prec_src1 = GGML_PREC_Q8) { + return ggml_cuda_mmq_get_config(type, J, fallback, prec_src1).type; } -static __host__ int ggml_cuda_mmq_get_nthreads(const ggml_type type, const int J, const bool fallback, const int cc) { - return ggml_cuda_mmq_get_config(type, J, fallback, cc).nthreads; +static constexpr __device__ int ggml_cuda_mmq_get_nthreads(ggml_type type, int J, bool fallback, ggml_prec prec_src1 = GGML_PREC_Q8) { + return ggml_cuda_mmq_get_config(type, J, fallback, prec_src1).nthreads; } -static constexpr __device__ int ggml_cuda_mmq_get_nthreads(ggml_type type, int J, bool fallback) { - return ggml_cuda_mmq_get_config(type, J, fallback).nthreads; -} - -static __host__ int ggml_cuda_mmq_get_occupancy(const ggml_type type, const int J, const bool fallback, const int cc) { - return ggml_cuda_mmq_get_config(type, J, fallback, cc).occupancy; -} - -static constexpr __device__ int ggml_cuda_mmq_get_occupancy(ggml_type type, int J, bool fallback) { - return ggml_cuda_mmq_get_config(type, J, fallback).occupancy; +static constexpr __device__ int ggml_cuda_mmq_get_occupancy(ggml_type type, int J, bool fallback, ggml_prec prec_src1 = GGML_PREC_Q8) { + return ggml_cuda_mmq_get_config(type, J, fallback, prec_src1).occupancy; } static __host__ int ggml_cuda_mmq_get_I(const ggml_type type, const int J, const bool fallback, const int cc) { return ggml_cuda_mmq_get_config(type, J, fallback, cc).I; } -static constexpr __device__ int ggml_cuda_mmq_get_I(ggml_type type, int J, bool fallback) { - return ggml_cuda_mmq_get_config(type, J, fallback).I; +static constexpr __device__ int ggml_cuda_mmq_get_I(ggml_type type, int J, bool fallback, ggml_prec prec_src1 = GGML_PREC_Q8) { + return ggml_cuda_mmq_get_config(type, J, fallback, prec_src1).I; } static __host__ int ggml_cuda_mmq_get_J(const ggml_type type, const int J, const bool fallback, const int cc) { return ggml_cuda_mmq_get_config(type, J, fallback, cc).J; } -static constexpr __device__ int ggml_cuda_mmq_get_J(ggml_type type, int J, bool fallback) { - return ggml_cuda_mmq_get_config(type, J, fallback).J; +static constexpr __device__ int ggml_cuda_mmq_get_J(ggml_type type, int J, bool fallback, ggml_prec prec_src1 = GGML_PREC_Q8) { + return ggml_cuda_mmq_get_config(type, J, fallback, prec_src1).J; } static __host__ ggml_cuda_mmq_sram_layout ggml_cuda_mmq_get_sram_layout(const ggml_type type, const int J, const bool fallback, const int cc) { return ggml_cuda_mmq_get_config(type, J, fallback, cc).sram_layout; } -static constexpr __device__ ggml_cuda_mmq_sram_layout ggml_cuda_mmq_get_sram_layout(ggml_type type, int J, bool fallback) { - return ggml_cuda_mmq_get_config(type, J, fallback).sram_layout; +static constexpr __device__ ggml_cuda_mmq_sram_layout ggml_cuda_mmq_get_sram_layout(ggml_type type, int J, bool fallback, ggml_prec prec_src1 = GGML_PREC_Q8) { + return ggml_cuda_mmq_get_config(type, J, fallback, prec_src1).sram_layout; } static __host__ int ggml_cuda_mmq_get_K_vram(const ggml_type type, const int J, const bool fallback, const int cc) { return ggml_cuda_mmq_get_config(type, J, fallback, cc).K_vram; } -static constexpr __device__ int ggml_cuda_mmq_get_K_vram(ggml_type type, int J, bool fallback) { - return ggml_cuda_mmq_get_config(type, J, fallback).K_vram; +static constexpr __device__ int ggml_cuda_mmq_get_K_vram(ggml_type type, int J, bool fallback, ggml_prec prec_src1 = GGML_PREC_Q8) { + return ggml_cuda_mmq_get_config(type, J, fallback, prec_src1).K_vram; } static __host__ bool ggml_cuda_mmq_get_stream_k(const ggml_type type, const int J, const bool fallback, const int cc) { return ggml_cuda_mmq_get_config(type, J, fallback, cc).stream_k; } -static constexpr __device__ bool ggml_cuda_mmq_get_stream_k(ggml_type type, int J, bool fallback) { - return ggml_cuda_mmq_get_config(type, J, fallback).stream_k; +static constexpr __device__ bool ggml_cuda_mmq_get_stream_k(ggml_type type, int J, bool fallback, ggml_prec prec_src1 = GGML_PREC_Q8) { + return ggml_cuda_mmq_get_config(type, J, fallback, prec_src1).stream_k; } static __host__ int ggml_cuda_mmq_get_fallback(const ggml_type type, const int J, const bool fallback, const int cc) { return ggml_cuda_mmq_get_config(type, J, fallback, cc).fallback; } -static constexpr __device__ int ggml_cuda_mmq_get_fallback(ggml_type type, int J, bool fallback) { - return ggml_cuda_mmq_get_config(type, J, fallback).fallback; +static constexpr __device__ int ggml_cuda_mmq_get_fallback(ggml_type type, int J, bool fallback, ggml_prec prec_src1 = GGML_PREC_Q8) { + return ggml_cuda_mmq_get_config(type, J, fallback, prec_src1).fallback; } // --------------------------------------------------------------------------------------------- @@ -365,8 +365,8 @@ static __host__ int ggml_cuda_mmq_get_sram_stride(const ggml_type type, const in return ggml_cuda_mmq_get_sram_stride(ggml_cuda_mmq_get_sram_layout(type, J, fallback, cc)); } -static constexpr __device__ int ggml_cuda_mmq_get_sram_stride(ggml_type type, int J, bool fallback) { - return ggml_cuda_mmq_get_sram_stride(ggml_cuda_mmq_get_sram_layout(type, J, fallback)); +static constexpr __device__ int ggml_cuda_mmq_get_sram_stride(ggml_type type, int J, bool fallback, ggml_prec prec_src1 = GGML_PREC_Q8) { + return ggml_cuda_mmq_get_sram_stride(ggml_cuda_mmq_get_sram_layout(type, J, fallback, prec_src1)); } static __host__ int ggml_cuda_mmq_get_J_max(const ggml_type type, const bool fallback, const int cc, const int64_t ne11) { @@ -541,9 +541,9 @@ struct ggml_cuda_mmq_util_funcs { vdr(vdr), load_tiles(load_tiles), vec_dot(vec_dot), write_back(write_back) {} }; -template +template static constexpr __device__ ggml_cuda_mmq_util_funcs ggml_cuda_mmq_get_util_funcs() { - if (!ggml_cuda_mmq_get_config(type, J, fallback).use_mma_data_layout()) { + if (!ggml_cuda_mmq_get_config(type, J, fallback, prec_src1).use_mma_data_layout()) { switch (type) { case GGML_TYPE_Q1_0: return ggml_cuda_mmq_util_funcs( @@ -690,17 +690,23 @@ static constexpr __device__ ggml_cuda_mmq_util_funcs ggml_cuda_mmq_get_util_func #ifdef BLACKWELL_MMA_AVAILABLE switch (type) { case GGML_TYPE_MXFP4: - return ggml_cuda_mmq_util_funcs( - -1, - ggml_cuda_mmq_load_tiles_mxfp4_fp4, - ggml_cuda_mmq_vec_dot_fp4_fp4_mma, - ggml_cuda_mmq_write_back_mma); + if (prec_src1 == GGML_PREC_Q4) { + return ggml_cuda_mmq_util_funcs( + -1, + ggml_cuda_mmq_load_tiles_mxfp4_fp4, + ggml_cuda_mmq_vec_dot_fp4_fp4_mma, + ggml_cuda_mmq_write_back_mma); + } + break; case GGML_TYPE_NVFP4: - return ggml_cuda_mmq_util_funcs( - -1, - ggml_cuda_mmq_load_tiles_nvfp4_nvfp4, - ggml_cuda_mmq_vec_dot_fp4_fp4_mma, - ggml_cuda_mmq_write_back_mma); + if (prec_src1 == GGML_PREC_Q4) { + return ggml_cuda_mmq_util_funcs( + -1, + ggml_cuda_mmq_load_tiles_nvfp4_nvfp4, + ggml_cuda_mmq_vec_dot_fp4_fp4_mma, + ggml_cuda_mmq_write_back_mma); + } + break; default: break; } @@ -841,37 +847,37 @@ static constexpr __device__ ggml_cuda_mmq_util_funcs ggml_cuda_mmq_get_util_func case GGML_TYPE_NVFP4: return ggml_cuda_mmq_util_funcs( -1, - ggml_cuda_mmq_load_tiles_nvfp4, - ggml_cuda_mmq_vec_dot_q8_0_16_q8_1_mma, + ggml_cuda_mmq_load_tiles_nvfp4, + ggml_cuda_mmq_vec_dot_q8_0_16_q8_1_mma, ggml_cuda_mmq_write_back_mma); default: return ggml_cuda_mmq_util_funcs(1, nullptr, nullptr, nullptr); } } -template +template static constexpr __device__ int ggml_cuda_mmq_get_vdr() { - return ggml_cuda_mmq_get_util_funcs().vdr; + return ggml_cuda_mmq_get_util_funcs().vdr; } -template +template static constexpr __device__ ggml_cuda_mmq_load_tiles_t ggml_cuda_mmq_get_load_tiles() { - return ggml_cuda_mmq_get_util_funcs().load_tiles; + return ggml_cuda_mmq_get_util_funcs().load_tiles; } -template +template static constexpr __device__ ggml_cuda_mmq_vec_dot_t ggml_cuda_mmq_get_vec_dot() { - return ggml_cuda_mmq_get_util_funcs().vec_dot; + return ggml_cuda_mmq_get_util_funcs().vec_dot; } -template +template static constexpr __device__ ggml_cuda_mmq_write_back_t ggml_cuda_mmq_get_write_back() { - return ggml_cuda_mmq_get_util_funcs().write_back; + return ggml_cuda_mmq_get_util_funcs().write_back; } // --------------------------------------------------------------------------------------------- -template +template static __device__ __forceinline__ void mul_mat_q_process_tile( const char * __restrict__ x, const int offset_x, const int * __restrict__ y, const int * __restrict__ ids_dst, float * __restrict__ dst, float * __restrict__ tmp_fixup, @@ -880,25 +886,27 @@ static __device__ __forceinline__ void mul_mat_q_process_tile( const int tile_x_max_i, const int tile_y_max_j, const int kb0_start, const int kb0_stop) { constexpr int warp_size = ggml_cuda_get_physical_warp_size(); - constexpr int nwarps = ggml_cuda_mmq_get_nthreads(type, J, fallback) / warp_size; + constexpr int nwarps = ggml_cuda_mmq_get_nthreads(type, J, fallback, prec_src1) / warp_size; constexpr int qk = ggml_cuda_type_traits::qk; - constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback); - constexpr ggml_cuda_mmq_load_tiles_t load_tiles = ggml_cuda_mmq_get_load_tiles(); - constexpr ggml_cuda_mmq_vec_dot_t vec_dot = ggml_cuda_mmq_get_vec_dot(); - constexpr ggml_cuda_mmq_write_back_t write_back = ggml_cuda_mmq_get_write_back(); + constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback, prec_src1); + constexpr ggml_cuda_mmq_load_tiles_t load_tiles = ggml_cuda_mmq_get_load_tiles(); + constexpr ggml_cuda_mmq_vec_dot_t vec_dot = ggml_cuda_mmq_get_vec_dot(); + constexpr ggml_cuda_mmq_write_back_t write_back = ggml_cuda_mmq_get_write_back(); extern __shared__ int data_mul_mat_q[]; int * tile_y = data_mul_mat_q + J; int * tile_x = tile_y + GGML_PAD(J*MMQ_TILE_Y_K, nwarps*warp_size); #if defined(BLACKWELL_MMA_AVAILABLE) - // FP4 tile stores 8 blocks - constexpr int ne_block = (type == GGML_TYPE_MXFP4 || type == GGML_TYPE_NVFP4) ? QK_FP4_MMQ : QK8_1_MMQ; + // FP4 tile stores 8 blocks. src1 above Q4 uses the generic + // Q8_1 tile layout instead of the packed FP4 tile. + constexpr int ne_block = ((type == GGML_TYPE_MXFP4 || type == GGML_TYPE_NVFP4) && prec_src1 == GGML_PREC_Q4) ? + QK_FP4_MMQ : QK8_1_MMQ; #else constexpr int ne_block = QK8_1_MMQ; #endif // defined(BLACKWELL_MMA_AVAILABLE) - constexpr int ITER_K = ggml_cuda_mmq_get_K_vram(type, J, fallback); + constexpr int ITER_K = ggml_cuda_mmq_get_K_vram(type, J, fallback, prec_src1); constexpr int blocks_per_iter = ITER_K / qk; float sum[J*I / (nwarps*warp_size)] = {0.0f}; @@ -950,8 +958,8 @@ static __device__ __forceinline__ void mul_mat_q_process_tile( // The mul_mat_q kernel implements "stream-k" work partitioning as described in https://arxiv.org/abs/2301.03598 -template -__launch_bounds__(ggml_cuda_mmq_get_nthreads(type, J, fallback), ggml_cuda_mmq_get_occupancy(type, J, fallback)) +template +__launch_bounds__(ggml_cuda_mmq_get_nthreads(type, J, fallback, prec_src1), ggml_cuda_mmq_get_occupancy(type, J, fallback, prec_src1)) static __global__ void mul_mat_q( const char * __restrict__ x, const int * __restrict__ y, const int32_t * __restrict__ ids_dst, const int32_t * __restrict__ expert_bounds, float * __restrict__ dst, float * __restrict__ tmp_fixup, @@ -962,15 +970,15 @@ static __global__ void mul_mat_q( const uint3 ntx) { // Skip unused template specializations for faster compilation: - if (ggml_cuda_mmq_get_config(type, J, fallback).type == GGML_TYPE_COUNT) { + if (ggml_cuda_mmq_get_config(type, J, fallback, prec_src1).type == GGML_TYPE_COUNT) { NO_DEVICE_CODE; return; } constexpr int warp_size = ggml_cuda_get_physical_warp_size(); - constexpr int nwarps = ggml_cuda_mmq_get_nthreads(type, J, fallback) / warp_size; + constexpr int nwarps = ggml_cuda_mmq_get_nthreads(type, J, fallback, prec_src1) / warp_size; constexpr int qk = ggml_cuda_type_traits::qk; - constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback); + constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback, prec_src1); const uint32_t nty = (nrows_x + I - 1) / I; // Number of tiles y @@ -990,7 +998,7 @@ static __global__ void mul_mat_q( } __syncthreads(); - if constexpr (!ggml_cuda_mmq_get_stream_k(type, J, fallback)) { + if constexpr (!ggml_cuda_mmq_get_stream_k(type, J, fallback, prec_src1)) { const uint2 tmp2 = fast_div_modulo(blockIdx.z, nchannels_y); const int wt = tmp2.x; const int zt = tmp2.y; @@ -1053,14 +1061,14 @@ static __global__ void mul_mat_q( const int offset_x = fastdiv(wt, sample_ratio)*stride_sample_x + fastdiv(zt, channel_ratio)*stride_channel_x + it*I*stride_row_x; constexpr bool fixup = false; - mul_mat_q_process_tile + mul_mat_q_process_tile (x, offset_x, y + offset_y, ids_dst_shared, dst + offset_dst, tmp_fixup, y_scale_tile, stride_row_x, ncols_y, stride_col_dst, tile_x_max_i, tile_y_max_j, 0, blocks_per_ne00.z); return; } - constexpr int ITER_K = ggml_cuda_mmq_get_K_vram(type, J, fallback); + constexpr int ITER_K = ggml_cuda_mmq_get_K_vram(type, J, fallback, prec_src1); constexpr int blocks_per_iter = ITER_K / qk; // kbc == k block continuous, current index in continuous ijk space. @@ -1147,7 +1155,7 @@ static __global__ void mul_mat_q( const int offset_x = fastdiv(wt, sample_ratio)*stride_sample_x + fastdiv(zt, channel_ratio)*stride_channel_x + it*I*stride_row_x; constexpr bool fixup = false; // All but (potentially) the last iterations write their data to dst rather than the fixup buffer. - mul_mat_q_process_tile + mul_mat_q_process_tile (x, offset_x, y + offset_y, ids_dst_shared, dst + offset_dst, tmp_fixup, y_scale_tile, stride_row_x, ncols_y, stride_col_dst, tile_x_max_i, tile_y_max_j, kb0_start, kb0_stop); @@ -1231,24 +1239,24 @@ static __global__ void mul_mat_q( const int offset_x = fastdiv(wt, sample_ratio)*stride_sample_x + fastdiv(zt, channel_ratio)*stride_channel_x + it*I*stride_row_x; constexpr bool fixup = true; // Last index writes its data to fixup buffer to avoid data races with other blocks. - mul_mat_q_process_tile + mul_mat_q_process_tile (x, offset_x, y + offset_y, ids_dst_shared, dst + offset_dst, tmp_fixup, y_scale_tile, stride_row_x, ncols_y, stride_col_dst, tile_x_max_i, tile_y_max_j, kb0_start, kb0_stop); } -template -__launch_bounds__(ggml_cuda_mmq_get_nthreads(type, J, fallback)/2, 1) +template +__launch_bounds__(ggml_cuda_mmq_get_nthreads(type, J, fallback, prec_src1)/2, 1) static __global__ void mul_mat_q_stream_k_fixup( const int32_t * __restrict__ ids_dst, const int32_t * __restrict__ expert_bounds, float * __restrict__ dst, float * __restrict__ tmp_last_tile, const uint3 blocks_per_ne00, const int nrows_x, const int ncols_dst, const int stride_col_dst, const uint3 nchannels_y, const int stride_channel_dst, const uint3 nsamples_y, const int stride_sample_dst, const uint3 ntx) { constexpr int warp_size = ggml_cuda_get_physical_warp_size(); - constexpr int nwarps = (ggml_cuda_mmq_get_nthreads(type, J, fallback) / 2) / warp_size; - constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback); + constexpr int nwarps = (ggml_cuda_mmq_get_nthreads(type, J, fallback, prec_src1) / 2) / warp_size; + constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback, prec_src1); constexpr int qk = ggml_cuda_type_traits::qk; - constexpr int ITER_K = ggml_cuda_mmq_get_K_vram(type, J, fallback); + constexpr int ITER_K = ggml_cuda_mmq_get_K_vram(type, J, fallback, prec_src1); constexpr int blocks_per_iter = ITER_K / qk; float sum[J / nwarps] = {0.0f}; @@ -1392,22 +1400,22 @@ static size_t mmq_get_nbytes_shared(const ggml_cuda_mmq_config & config, const i return nbs_ids + nbs_x + GGML_PAD(nbs_y, config.nthreads*sizeof(int)); } -template +template static void launch_mul_mat_q(ggml_backend_cuda_context & ctx, const mmq_args & args, cudaStream_t stream) { const int id = ggml_cuda_get_device(); const int cc = ggml_cuda_info().devices[id].cc; const int nsm = ggml_cuda_info().devices[id].nsm; const int warp_size = ggml_cuda_info().devices[id].warp_size; - const ggml_cuda_mmq_config config = ggml_cuda_mmq_get_config(type, J, fallback, cc); + const ggml_cuda_mmq_config config = ggml_cuda_mmq_get_config(type, J, fallback, cc, prec_src1); GGML_ASSERT(config.nthreads % warp_size == 0); const int nwarps = config.nthreads / warp_size; const int nbytes_shared = mmq_get_nbytes_shared(config, cc); const dim3 block_dims(warp_size, nwarps, 1); - CUDA_SET_SHARED_MEMORY_LIMIT((mul_mat_q), nbytes_shared); - CUDA_SET_SHARED_MEMORY_LIMIT((mul_mat_q), nbytes_shared); + CUDA_SET_SHARED_MEMORY_LIMIT((mul_mat_q), nbytes_shared); + CUDA_SET_SHARED_MEMORY_LIMIT((mul_mat_q), nbytes_shared); const int nty = (args.nrows_x + config.I - 1) / config.I; const int ntx = (args.ncols_max + config.J - 1) / config.J; @@ -1426,8 +1434,8 @@ static void launch_mul_mat_q(ggml_backend_cuda_context & ctx, const mmq_args & a const uint3 channel_ratio_fd = init_fastdiv_values(channel_ratio); const uint3 sample_ratio_fd = init_fastdiv_values(sample_ratio); - if (!ggml_cuda_mmq_get_stream_k(type, J, fallback, cc)) { - mul_mat_q<<>> + if (!config.stream_k) { + mul_mat_q<<>> (args.x, args.y, args.ids_dst, args.expert_bounds, args.dst, nullptr, args.y_scale, blocks_per_ne00_fd, args.nrows_x, args.ncols_dst, args.stride_row_x, args.ncols_y, args.nrows_dst, channel_ratio_fd, nchannels_y_fd, args.stride_channel_x, args.stride_channel_y, args.stride_channel_dst, @@ -1456,7 +1464,7 @@ static void launch_mul_mat_q(ggml_backend_cuda_context & ctx, const mmq_args & a const dim3 block_nums_fixup(block_nums_stream_k.x, config.I/warp_size, 1); const dim3 block_dims_fixup(block_dims.x, block_dims.y/2, block_dims.z); - mul_mat_q<<>> + mul_mat_q<<>> (args.x, args.y, args.ids_dst, args.expert_bounds, args.dst, tmp_fixup.ptr, args.y_scale, blocks_per_ne00_fd, args.nrows_x, args.ncols_dst, args.stride_row_x, args.ncols_y, args.nrows_dst, channel_ratio_fd, nchannels_y_fd, args.stride_channel_x, args.stride_channel_y, args.stride_channel_dst, @@ -1468,13 +1476,13 @@ static void launch_mul_mat_q(ggml_backend_cuda_context & ctx, const mmq_args & a } CUDA_CHECK(cudaGetLastError()); - mul_mat_q_stream_k_fixup<<>> + mul_mat_q_stream_k_fixup<<>> (args.ids_dst, args.expert_bounds, args.dst, tmp_fixup.ptr, blocks_per_ne00_fd, args.nrows_x, args.ncols_dst, args.nrows_dst, nchannels_y_fd, args.stride_channel_dst, nsamples_y_fd, args.stride_sample_dst, ntx_fd); } -template +template void mul_mat_q_switch_J(ggml_backend_cuda_context & ctx, const mmq_args & args, cudaStream_t stream) { const int id = ggml_cuda_get_device(); const int cc = ggml_cuda_info().devices[id].cc; @@ -1484,7 +1492,7 @@ void mul_mat_q_switch_J(ggml_backend_cuda_context & ctx, const mmq_args & args, int ntiles_J_best = INT_MAX; for (int J = 8; J <= 128 && ntiles_J_best > 1; J += 8) { - const ggml_cuda_mmq_config config = ggml_cuda_mmq_get_config(type, J, fallback, cc); + const ggml_cuda_mmq_config config = ggml_cuda_mmq_get_config(type, J, fallback, cc, prec_src1); if (config.type == GGML_TYPE_COUNT) { continue; } @@ -1503,52 +1511,52 @@ void mul_mat_q_switch_J(ggml_backend_cuda_context & ctx, const mmq_args & args, switch (J_best) { case 8: - launch_mul_mat_q(ctx, args, stream); + launch_mul_mat_q(ctx, args, stream); break; case 16: - launch_mul_mat_q(ctx, args, stream); + launch_mul_mat_q(ctx, args, stream); break; case 24: - launch_mul_mat_q(ctx, args, stream); + launch_mul_mat_q(ctx, args, stream); break; case 32: - launch_mul_mat_q(ctx, args, stream); + launch_mul_mat_q(ctx, args, stream); break; case 40: - launch_mul_mat_q(ctx, args, stream); + launch_mul_mat_q(ctx, args, stream); break; case 48: - launch_mul_mat_q(ctx, args, stream); + launch_mul_mat_q(ctx, args, stream); break; case 56: - launch_mul_mat_q(ctx, args, stream); + launch_mul_mat_q(ctx, args, stream); break; case 64: - launch_mul_mat_q(ctx, args, stream); + launch_mul_mat_q(ctx, args, stream); break; case 72: - launch_mul_mat_q(ctx, args, stream); + launch_mul_mat_q(ctx, args, stream); break; case 80: - launch_mul_mat_q(ctx, args, stream); + launch_mul_mat_q(ctx, args, stream); break; case 88: - launch_mul_mat_q(ctx, args, stream); + launch_mul_mat_q(ctx, args, stream); break; case 96: - launch_mul_mat_q(ctx, args, stream); + launch_mul_mat_q(ctx, args, stream); break; case 104: - launch_mul_mat_q(ctx, args, stream); + launch_mul_mat_q(ctx, args, stream); break; case 112: - launch_mul_mat_q(ctx, args, stream); + launch_mul_mat_q(ctx, args, stream); break; case 120: - launch_mul_mat_q(ctx, args, stream); + launch_mul_mat_q(ctx, args, stream); break; case 128: - launch_mul_mat_q(ctx, args, stream); + launch_mul_mat_q(ctx, args, stream); break; default: fprintf(stderr, "J_best=%d\n", J_best); @@ -1557,20 +1565,24 @@ void mul_mat_q_switch_J(ggml_backend_cuda_context & ctx, const mmq_args & args, } } -template +template void mul_mat_q_case(ggml_backend_cuda_context & ctx, const mmq_args & args, cudaStream_t stream) { if (args.nrows_x % 128 == 0) { constexpr bool fallback = false; - mul_mat_q_switch_J(ctx, args, stream); + mul_mat_q_switch_J(ctx, args, stream); } else { constexpr bool fallback = true; - mul_mat_q_switch_J(ctx, args, stream); + mul_mat_q_switch_J(ctx, args, stream); } } #define DECL_MMQ_CASE(type) \ template void mul_mat_q_case(ggml_backend_cuda_context & ctx, const mmq_args & args, cudaStream_t stream) \ +// FP4 variant: uses native FP4 MMA instead of keeping src1 at Q8_1. +#define DECL_MMQ_CASE_W4A4(type) \ + template void mul_mat_q_case(ggml_backend_cuda_context & ctx, const mmq_args & args, cudaStream_t stream) \ + extern DECL_MMQ_CASE(GGML_TYPE_Q1_0); extern DECL_MMQ_CASE(GGML_TYPE_Q2_0); extern DECL_MMQ_CASE(GGML_TYPE_Q4_0); @@ -1596,6 +1608,8 @@ extern DECL_MMQ_CASE(GGML_TYPE_IQ4_XS); // ----------------------------------------- extern DECL_MMQ_CASE(GGML_TYPE_MXFP4); extern DECL_MMQ_CASE(GGML_TYPE_NVFP4); +extern DECL_MMQ_CASE_W4A4(GGML_TYPE_MXFP4); +extern DECL_MMQ_CASE_W4A4(GGML_TYPE_NVFP4); // ------------------------------------------------------------------------------------------------------------------------- diff --git a/ggml/src/ggml-cuda/template-instances/generate_cu_files.py b/ggml/src/ggml-cuda/template-instances/generate_cu_files.py index d7cd27167..7be409bd3 100755 --- a/ggml/src/ggml-cuda/template-instances/generate_cu_files.py +++ b/ggml/src/ggml-cuda/template-instances/generate_cu_files.py @@ -50,6 +50,12 @@ SOURCE_MMQ = """// This file has been autogenerated by generate_cu_files.py, do DECL_MMQ_CASE({type}); """ +TYPES_MMQ_W4A4 = ["GGML_TYPE_MXFP4", "GGML_TYPE_NVFP4"] + +SOURCE_MMQ_W4A4 = """ +DECL_MMQ_CASE_W4A4({type}); +""" + SOURCE_MMF = """// This file has been autogenerated by generate_cu_files.py, do not edit manually. #include "../mmf.cuh" @@ -105,6 +111,8 @@ for ncols in [8, 16, 32, 64]: for type in TYPES_MMQ: with open(f"mmq-instance-{get_short_name(type)}.cu", "w") as f: f.write(SOURCE_MMQ.format(type=type)) + if type in TYPES_MMQ_W4A4: + f.write(SOURCE_MMQ_W4A4.format(type=type)) for type in range(1, 17): with open(f"mmf-instance-ncols_{type}.cu", "w") as f: diff --git a/ggml/src/ggml-cuda/template-instances/mmq-instance-mxfp4.cu b/ggml/src/ggml-cuda/template-instances/mmq-instance-mxfp4.cu index c14624c52..0c1c7f291 100644 --- a/ggml/src/ggml-cuda/template-instances/mmq-instance-mxfp4.cu +++ b/ggml/src/ggml-cuda/template-instances/mmq-instance-mxfp4.cu @@ -3,3 +3,5 @@ #include "../mmq.cuh" DECL_MMQ_CASE(GGML_TYPE_MXFP4); + +DECL_MMQ_CASE_W4A4(GGML_TYPE_MXFP4); diff --git a/ggml/src/ggml-cuda/template-instances/mmq-instance-nvfp4.cu b/ggml/src/ggml-cuda/template-instances/mmq-instance-nvfp4.cu index 2cb140d35..bed8c8e2e 100644 --- a/ggml/src/ggml-cuda/template-instances/mmq-instance-nvfp4.cu +++ b/ggml/src/ggml-cuda/template-instances/mmq-instance-nvfp4.cu @@ -3,3 +3,5 @@ #include "../mmq.cuh" DECL_MMQ_CASE(GGML_TYPE_NVFP4); + +DECL_MMQ_CASE_W4A4(GGML_TYPE_NVFP4);