Skip to content

ggml-webgpu: simplify flash_attn shaders and refactor several WGSL shaders - #26134

Merged
ggerganov merged 1 commit into
ggml-org:masterfrom
yomaytk:refactor-wgsl
Aug 10, 2026
Merged

ggml-webgpu: simplify flash_attn shaders and refactor several WGSL shaders#26134
ggerganov merged 1 commit into
ggml-org:masterfrom
yomaytk:refactor-wgsl

Conversation

@yomaytk

@yomaytk yomaytk commented Jul 26, 2026

Copy link
Copy Markdown
Member

Overview

This PR refactors several WGSL files and simplifies the flash attention WGSL files. The changes to each WGSL shader are as follows:

  • concat.wgsl, rms_norm_mul.wgsl, row_norm.wgsl, solve_tri.wgsl, ssm_scan.wgsl: Remove unused params.
  • conv2d.wgsl, conv2d_dw.wgsl, im2col.wgsl: Define the macros as the actual type for weights, input, and output buffer and remove some unnecessary functions.
  • soft_max.wgsl: Remove the unused param ne. Also, binding numbers for each tensor buffer are computed sequentially instead of defining buffers with nested macros, since the nested macro structures aren't easy to read.
  • flash_attn_decls.tmpl: This file contains the common definitions of params and types shared among flash attention WGSL shaders. Binding numbers for the tensor buffers are computed in the same way as in soft_max.wgsl.
  • flash_attn_staging.tmpl: flash_attn_quant_staging.tmpl was renamed to this file, and load_(k|v)_tile_block for the float type is extracted from the three flash attention shaders.
  • flash_attn.wgsl, flash_attn_tile.wgsl, flash_attn_vec_split.wgsl: Common code is extracted into flash_attn_decls.tmpl and flash_attn_staging.tmpl.

Requirements

  • I have read and agree with the contributing guidelines
  • AI usage disclosure: YES, I used AI to find the redundant code for the all wgsl shaders, and to refactor the several parts. Those changes are reviewed and the other parts (mainly for FA) were written by me.

@github-actions github-actions Bot added ggml changes relating to the ggml tensor library for machine learning WebGPU labels Jul 26, 2026
@yomaytk
yomaytk force-pushed the refactor-wgsl branch 3 times, most recently from 99bf54f to 9dd0ec4 Compare August 1, 2026 09:16
@yomaytk
yomaytk marked this pull request as ready for review August 1, 2026 09:24
@yomaytk
yomaytk requested a review from a team as a code owner August 1, 2026 09:24
@yomaytk yomaytk added the merge ready A maintainer can use this label to indicate that they consider the changes final and ready to merge. label Aug 8, 2026
@ggerganov
ggerganov merged commit f401bb1 into ggml-org:master Aug 10, 2026
27 of 29 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

ggml changes relating to the ggml tensor library for machine learning merge ready A maintainer can use this label to indicate that they consider the changes final and ready to merge. WebGPU

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants