From b313fd3ce8c151a0ddc18c85cf40b76c1bec4958 Mon Sep 17 00:00:00 2001 From: Masashi Yoshimura Date: Tue, 30 Jun 2026 17:20:04 +0900 Subject: [PATCH] ggml-webgpu: add support for NVFP4 (llama/25143) --- .../ggml-webgpu/ggml-webgpu-shader-lib.hpp | 7 +++ ggml/src/ggml-webgpu/ggml-webgpu.cpp | 3 ++ .../wgsl-shaders/common_decls.tmpl | 18 +++++++- .../ggml-webgpu/wgsl-shaders/get_rows.wgsl | 21 +++++++++ .../wgsl-shaders/mul_mat_decls.tmpl | 43 ++++++++++++++++- .../wgsl-shaders/mul_mat_vec_acc.tmpl | 46 +++++++++++++++++++ 6 files changed, 135 insertions(+), 3 deletions(-) diff --git a/ggml/src/ggml-webgpu/ggml-webgpu-shader-lib.hpp b/ggml/src/ggml-webgpu/ggml-webgpu-shader-lib.hpp index c00a2e9ee..80a16d16d 100644 --- a/ggml/src/ggml-webgpu/ggml-webgpu-shader-lib.hpp +++ b/ggml/src/ggml-webgpu/ggml-webgpu-shader-lib.hpp @@ -1563,6 +1563,7 @@ class ggml_webgpu_shader_lib { case GGML_TYPE_IQ1_S: case GGML_TYPE_IQ4_NL: case GGML_TYPE_MXFP4: + case GGML_TYPE_NVFP4: { // Quantized types using u32 buffers for portability. defines.push_back("SRC_TYPE=u32"); @@ -1593,6 +1594,8 @@ class ggml_webgpu_shader_lib { } else if ((key.src_type >= GGML_TYPE_Q4_0 && key.src_type <= GGML_TYPE_Q8_1) || key.src_type == GGML_TYPE_IQ4_NL || key.src_type == GGML_TYPE_MXFP4) { defines.push_back("BLOCK_SIZE=32u"); + } else if (key.src_type == GGML_TYPE_NVFP4) { + defines.push_back("BLOCK_SIZE=64u"); } else if (key.src_type >= GGML_TYPE_Q2_K) { defines.push_back("BLOCK_SIZE=256u"); } else { @@ -1960,6 +1963,7 @@ class ggml_webgpu_shader_lib { defines.push_back(type_upper + "_TABLES"); break; case GGML_TYPE_MXFP4: + case GGML_TYPE_NVFP4: defines.push_back(type_upper + "_LUT"); break; default: @@ -2103,6 +2107,7 @@ class ggml_webgpu_shader_lib { defines.push_back(type_upper + "_TABLES"); break; case GGML_TYPE_MXFP4: + case GGML_TYPE_NVFP4: defines.push_back(type_upper + "_LUT"); break; default: @@ -2274,6 +2279,7 @@ class ggml_webgpu_shader_lib { defines.push_back(type_upper + "_TABLES"); break; case GGML_TYPE_MXFP4: + case GGML_TYPE_NVFP4: defines.push_back(type_upper + "_LUT"); break; default: @@ -2394,6 +2400,7 @@ class ggml_webgpu_shader_lib { defines.push_back(type_upper + "_TABLES"); break; case GGML_TYPE_MXFP4: + case GGML_TYPE_NVFP4: defines.push_back(type_upper + "_LUT"); break; default: diff --git a/ggml/src/ggml-webgpu/ggml-webgpu.cpp b/ggml/src/ggml-webgpu/ggml-webgpu.cpp index f0ec18abd..29025e9ba 100644 --- a/ggml/src/ggml-webgpu/ggml-webgpu.cpp +++ b/ggml/src/ggml-webgpu/ggml-webgpu.cpp @@ -4056,6 +4056,7 @@ static bool ggml_webgpu_supported_qtype(ggml_type type) { case GGML_TYPE_IQ4_NL: case GGML_TYPE_IQ4_XS: case GGML_TYPE_MXFP4: + case GGML_TYPE_NVFP4: return true; default: return false; @@ -4156,6 +4157,7 @@ static bool ggml_backend_webgpu_device_supports_op(ggml_backend_dev_t dev, const case GGML_TYPE_IQ4_NL: case GGML_TYPE_IQ4_XS: case GGML_TYPE_MXFP4: + case GGML_TYPE_NVFP4: supports_op = true; break; default: @@ -4196,6 +4198,7 @@ static bool ggml_backend_webgpu_device_supports_op(ggml_backend_dev_t dev, const case GGML_TYPE_IQ4_NL: case GGML_TYPE_IQ4_XS: case GGML_TYPE_MXFP4: + case GGML_TYPE_NVFP4: supports_op = true; break; default: diff --git a/ggml/src/ggml-webgpu/wgsl-shaders/common_decls.tmpl b/ggml/src/ggml-webgpu/wgsl-shaders/common_decls.tmpl index 758efa17d..6634fbd65 100644 --- a/ggml/src/ggml-webgpu/wgsl-shaders/common_decls.tmpl +++ b/ggml/src/ggml-webgpu/wgsl-shaders/common_decls.tmpl @@ -896,9 +896,23 @@ const kvalues_iq4nl = array( #endif -#ifdef MXFP4_LUT +#if defined(MXFP4_LUT) || defined(NVFP4_LUT) const kvalues_mxfp4 = array( 0, 1, 2, 3, 4, 6, 8, 12, 0, -1, -2, -3, -4, -6, -8, -12 ); -#endif +#endif // MXFP4_LUT || NVFP4_LUT +#ifdef NVFP4_LUT +fn ue4m3_to_fp32(u: u32) -> f32 { + if (u == 0u || u == 127u) { + return 0.0; + } + let exp = (u >> 3u) & 15u; + let man = u & 7u; + if (exp == 0u) { + return f32(man) * (1.0 / 512.0); + } + let bits = ((exp + 120u) << 23u) | (man << 20u); + return bitcast(bits); +} +#endif // NVFP4_LUT diff --git a/ggml/src/ggml-webgpu/wgsl-shaders/get_rows.wgsl b/ggml/src/ggml-webgpu/wgsl-shaders/get_rows.wgsl index 78d61a93d..487edb327 100644 --- a/ggml/src/ggml-webgpu/wgsl-shaders/get_rows.wgsl +++ b/ggml/src/ggml-webgpu/wgsl-shaders/get_rows.wgsl @@ -672,6 +672,27 @@ fn copy_elements(src_base: u32, dst_base: u32, offset: u32) { } #endif +#ifdef NVFP4 +fn copy_elements(src_base: u32, dst_base: u32, offset: u32) { + let block_byte_base = (src_base + offset) * 36; + let d_word = load_u32_at_src(block_byte_base); + for (var sub: u32 = 0u; sub < 4; sub++) { + let d = ue4m3_to_fp32(get_byte(d_word, sub)) * 0.5; + for (var j: u32 = 0u; j < 2; j++) { + let q_packed = load_u32_at_src(block_byte_base + 4 + sub * 8 + j * 4); + for (var k: u32 = 0; k < 4; k++) { + let q_byte = get_byte(q_packed, k); + let q_lo = f32(kvalues_mxfp4[q_byte & 0xFu]) * d; + let q_hi = f32(kvalues_mxfp4[(q_byte >> 4) & 0xF]) * d; + let dst_offset = dst_base + offset * 64 + sub * 16 + j * 4 + k; + dst[dst_offset] = q_lo; + dst[dst_offset + 8u] = q_hi; + } + } + } +} +#endif + @group(0) @binding(0) var src: array; diff --git a/ggml/src/ggml-webgpu/wgsl-shaders/mul_mat_decls.tmpl b/ggml/src/ggml-webgpu/wgsl-shaders/mul_mat_decls.tmpl index 6a2eb8c82..13996ab51 100644 --- a/ggml/src/ggml-webgpu/wgsl-shaders/mul_mat_decls.tmpl +++ b/ggml/src/ggml-webgpu/wgsl-shaders/mul_mat_decls.tmpl @@ -241,7 +241,7 @@ fn init_shmem_src0(thread_id: u32, batch_offset: u32, offset_m: u32, k_outer: u3 #endif // INIT_SRC0_SHMEM_Q8_1 #if defined(INIT_SRC0_SHMEM_MXFP4) - let block_byte_base = src0_idx * 17u; + let block_byte_base = src0_idx * 17u; // BLOCK_SIZE_BYTES = 17u; let eu8 = get_byte(load_u32_at_src0_aligned(block_byte_base), block_byte_base & 3u); let e = ldexp(1.0, i32(eu8) - 128); @@ -263,6 +263,47 @@ fn init_shmem_src0(thread_id: u32, batch_offset: u32, offset_m: u32, k_outer: u3 } #endif // legacy-quants +#if defined(INIT_SRC0_SHMEM_NVFP4) +const BLOCK_SIZE = 64u; +const BLOCK_SIZE_BYTES = 36u; +const SUB_BLOCK_SIZE = 16u; // elements sharing one UE4M3 scale +const NQ = 16u; +const BYTES_PER_THREAD = 8u; +const BYTES_PER_INNER_LOOP = 4u; + +fn init_shmem_src0(thread_id: u32, batch_offset: u32, offset_m: u32, k_outer: u32) { + for (var i = thread_id * NQ; i < TILE_SRC0_SHMEM; i += TOTAL_WORKGROUP_SIZE * NQ) { + let tile_m = i / TILE_K; + let tile_k_start = i % TILE_K; + let global_m = offset_m + tile_m; + let global_k_start = k_outer + tile_k_start; + + if (global_m >= params.m) { + break; + } + + let block_k = global_k_start / BLOCK_SIZE; + let sub_block = (global_k_start % BLOCK_SIZE) / SUB_BLOCK_SIZE; + let src0_idx = batch_offset + global_m * params.stride_01 + block_k; + + let block_byte_base = src0_idx * BLOCK_SIZE_BYTES; + let d_byte_base = block_byte_base; + let qs_byte_base = block_byte_base + 4u; + + let d = ue4m3_to_fp32(get_byte(load_u32_at_src0_aligned(d_byte_base), sub_block)) * 0.5; + + for (var j = 0u; j < BYTES_PER_THREAD / BYTES_PER_INNER_LOOP; j++) { + let q_packed = load_u32_at_src0_aligned(qs_byte_base + sub_block * 8u + j * 4u); + for (var k = 0u; k < BYTES_PER_INNER_LOOP; k++) { + let q_byte = get_byte(q_packed, k); + shmem[i + j * BYTES_PER_INNER_LOOP + k] = f16(f32(kvalues_mxfp4[q_byte & 0xF]) * d); + shmem[i + j * BYTES_PER_INNER_LOOP + k + 8u] = f16(f32(kvalues_mxfp4[(q_byte >> 4) & 0xF]) * d); + } + } + } +} +#endif // INIT_SRC0_SHMEM_NVFP4 + // k-quants #if defined(INIT_SRC0_SHMEM_Q2_K) || defined(INIT_SRC0_SHMEM_Q3_K) || defined(INIT_SRC0_SHMEM_Q4_K) || defined(INIT_SRC0_SHMEM_Q5_K) || defined(INIT_SRC0_SHMEM_Q6_K) const BLOCK_SIZE = 256u; diff --git a/ggml/src/ggml-webgpu/wgsl-shaders/mul_mat_vec_acc.tmpl b/ggml/src/ggml-webgpu/wgsl-shaders/mul_mat_vec_acc.tmpl index b0703fe90..8fd0d1907 100644 --- a/ggml/src/ggml-webgpu/wgsl-shaders/mul_mat_vec_acc.tmpl +++ b/ggml/src/ggml-webgpu/wgsl-shaders/mul_mat_vec_acc.tmpl @@ -1505,3 +1505,49 @@ fn accumulate_vec_dot(thread_id: u32, row_base: u32, src0_batch_offset: u32, src return acc; } #endif + +#ifdef MUL_ACC_NVFP4 +#define BLOCK_SIZE 64 +#define BLOCK_SIZE_BYTES 36 +#define THREADS_PER_BLOCK 4 +#define ELEMS_PER_THREAD (BLOCK_SIZE/THREADS_PER_BLOCK) +fn accumulate_vec_dot(thread_id: u32, row_base: u32, src0_batch_offset: u32, src1_idx_base: u32) -> array, NUM_COLS> { + var acc: array, NUM_COLS>; + + let num_blocks = params.k / BLOCK_SIZE; + let sub = thread_id % THREADS_PER_BLOCK; + for (var block = thread_id/THREADS_PER_BLOCK; block < num_blocks; block += WG_SIZE/THREADS_PER_BLOCK) { + let x_base = src1_idx_base + block * BLOCK_SIZE + sub * ELEMS_PER_THREAD; + var x_block: array, NUM_COLS>; + for (var col = 0u; col < NUM_COLS;col += 1) { + for (var i = 0u; i < ELEMS_PER_THREAD / 2; i++) { + x_block[col][i] = f32(src1[x_base + col * params.stride_11 + i]); + x_block[col][i + 8] = f32(src1[x_base + col * params.stride_11 + i + 8]); + } + } + for (var row = 0u; row < OUTPUTS_PER_WG; row++) { + let output_row = row_base + row; + if (output_row < params.m) { + let block_byte_base = (src0_batch_offset + output_row * params.stride_01 + block) * BLOCK_SIZE_BYTES; + let d = ue4m3_to_fp32(get_byte(load_u32_at_src0_aligned(block_byte_base), sub)) * 0.5; + let q_w0 = load_u32_at_src0_aligned(block_byte_base + 4u + 8u * sub); + let q_w1 = load_u32_at_src0_aligned(block_byte_base + 8u + 8u * sub); + for (var col = 0u;col < NUM_COLS;col += 1) { + var row_sum = 0.0; + for (var l = 0u; l < 8u; l++) { + let q_word = select(q_w0, q_w1, l >= 4u); + let q_byte = get_byte(q_word, l % 4u); + let q_lo = f32(kvalues_mxfp4[q_byte & 0xFu]) * d; + let q_hi = f32(kvalues_mxfp4[(q_byte >> 4u) & 0xFu]) * d; + row_sum += q_lo * x_block[col][l]; + row_sum += q_hi * x_block[col][l + 8u]; + } + acc[col][row] += row_sum; + } + } + } + } + + return acc; +} +#endif