Skip to content

[ExecuTorch][WebGPU] Add Llama K16 online causal attention #21132

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
JCNTH wants to merge 4 commits into
base: gh/JCNTH/108/base
Choose a base branch
from
Open

[ExecuTorch][WebGPU] Add Llama K16 online causal attention #21132

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
5f3c4f6
Update
JCNTH Jul 22, 2026
b325261
Update
JCNTH Jul 22, 2026
c94e353
Update
JCNTH Jul 22, 2026
684fd2d
Update
JCNTH Jul 22, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
209 changes: 181 additions & 28 deletions backends/webgpu/runtime/ops/sdpa/Sdpa.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@
* LICENSE file in the root directory of this source tree.
*/

#include <executorch/backends/webgpu/runtime/WebGPUDevice.h>
#include <executorch/backends/webgpu/runtime/WebGPUGraph.h>
#include <executorch/backends/webgpu/runtime/WebGPUUtils.h>
#include <executorch/backends/webgpu/runtime/ops/OperatorRegistry.h>
Expand All @@ -14,6 +15,7 @@
#include <executorch/backends/webgpu/runtime/ops/sdpa/sdpa_compute_out_half_wgsl.h>
#include <executorch/backends/webgpu/runtime/ops/sdpa/sdpa_compute_out_wgsl.h>
#include <executorch/backends/webgpu/runtime/ops/sdpa/sdpa_softmax_wgsl.h>
#include <executorch/backends/webgpu/runtime/ops/sdpa/streaming_attention_k16_causal_bound_wgsl.h>
#include <executorch/backends/webgpu/runtime/ops/sdpa_fd_decode/SdpaFdDecode.h>
#include <executorch/backends/webgpu/runtime/ops/update_cache/update_cache_half_wgsl.h>
#include <executorch/backends/webgpu/runtime/ops/update_cache/update_cache_wgsl.h>
Expand Down Expand Up @@ -74,6 +76,24 @@ struct ComputeOutParams {
};
static_assert(sizeof(ComputeOutParams) == 32, "ComputeOutParams must be 32B");

struct StreamingAttentionK16Params {
uint32_t S;
uint32_t context_len;
uint32_t input_pos;
uint32_t q_token_stride4;
uint32_t q_head_stride4;
uint32_t kv_token_stride4;
uint32_t kv_head_stride4;
uint32_t o_token_stride4;
uint32_t o_head_stride4;
uint32_t _pad0;
uint32_t _pad1;
uint32_t _pad2;
};
static_assert(
sizeof(StreamingAttentionK16Params) == 48,
"StreamingAttentionK16Params must be 48B");

struct SdpaLiveState {
int64_t s;
int64_t pos;
Expand All @@ -82,10 +102,12 @@ struct SdpaLiveState {
AttnWeightsParams attn_weights;
SoftmaxParams softmax;
ComputeOutParams compute_out;
StreamingAttentionK16Params streaming_k16;
utils::WgCount update_cache_grid;
utils::WgCount qk_grid;
utils::WgCount softmax_grid;
utils::WgCount av_grid;
utils::WgCount streaming_k16_grid;
bool use_fd;
SdpaFdDecodeState fd;
};
Expand Down Expand Up @@ -147,6 +169,55 @@ static ComputeOutParams make_compute_out_params(
return p;
}

static StreamingAttentionK16Params make_streaming_attention_k16_params(
int64_t S,
int64_t context_len,
int64_t input_pos,
int64_t Hq,
int64_t Hkv,
int64_t D) {
StreamingAttentionK16Params p = {};
p.S = static_cast<uint32_t>(S);
p.context_len = static_cast<uint32_t>(context_len);
p.input_pos = static_cast<uint32_t>(input_pos);
p.q_token_stride4 = static_cast<uint32_t>(Hq * D / 4);
p.q_head_stride4 = static_cast<uint32_t>(D / 4);
p.kv_token_stride4 = static_cast<uint32_t>(Hkv * D / 4);
p.kv_head_stride4 = static_cast<uint32_t>(D / 4);
p.o_token_stride4 = static_cast<uint32_t>(Hq * D / 4);
p.o_head_stride4 = static_cast<uint32_t>(D / 4);
return p;
}

static bool streaming_attention_k16_device_supported(WGPUDevice device) {
WGPULimits limits = {};
const WebGPUContext* context = get_default_webgpu_context();
return context != nullptr && context->shader_f16_supported &&
wgpuDeviceGetLimits(device, &limits) == WGPUStatus_Success &&
limits.maxComputeInvocationsPerWorkgroup >= 128u &&
limits.maxComputeWorkgroupSizeX >= 32u &&
limits.maxComputeWorkgroupSizeY >= 4u &&
limits.maxComputeWorkgroupStorageSize >= 14720u;
}

static utils::WgCount streaming_attention_k16_grid(
WGPUDevice device,
int64_t S,
int64_t Hkv,
int64_t g) {
const uint64_t groups_per_kv =
(static_cast<uint64_t>(S) * static_cast<uint64_t>(g) + 31u) / 32u;
const uint64_t workgroups = static_cast<uint64_t>(Hkv) * groups_per_kv;
if (workgroups == 0u || workgroups > UINT32_MAX) {
throw std::runtime_error("WebGPU sdpa: K16 workgroup count exceeds uint32");
}
if (workgroups > utils::queried_max_workgroups(device)) {
throw std::runtime_error(
"WebGPU sdpa: K16 workgroup count exceeds the 1D dispatch limit");
}
return {static_cast<uint32_t>(workgroups), 1u};
}

// A buffer + its byte size, for binding.
struct BufferBinding {
WGPUBuffer buffer;
Expand Down Expand Up @@ -436,6 +507,18 @@ void sdpa_with_kv_cache_impl(WebGPUGraph& graph, const std::vector<int>& args) {
}

const WGPUDevice device = graph.device();
const WGPUBuffer k16_buffers[] = {
q.buffer, k.buffer, v.buffer, k_cache.buffer, v_cache.buffer, out.buffer};
bool k16_buffers_distinct = true;
for (size_t i = 0; i < 6; i++) {
for (size_t j = i + 1; j < 6; j++) {
k16_buffers_distinct =
k16_buffers_distinct && k16_buffers[i] != k16_buffers[j];
}
}
const bool k16_eligible = graph.kv_f16() && Hq == 32 && Hkv == 8 && g == 4 &&
D == 64 && scale == 0.125f && out.dims == q.dims &&
k16_buffers_distinct && streaming_attention_k16_device_supported(device);
const uint32_t uc_wg =
utils::clamp_workgroup_size(device, kUpdateCacheWorkgroupSizeX);
const uint32_t qk_wg = utils::clamp_workgroup_size(
Expand Down Expand Up @@ -466,7 +549,8 @@ void sdpa_with_kv_cache_impl(WebGPUGraph& graph, const std::vector<int>& args) {
uc_wg,
qk_wg,
av_wg,
fd_eligible](WebGPUGraph& gr) {
fd_eligible,
k16_eligible](WebGPUGraph& gr) {
SdpaLiveState state = {};
const auto& q_live_dims = gr.cur_dims(q_id);
state.s = q_live_dims[qn - 3];
Expand Down Expand Up @@ -503,21 +587,8 @@ void sdpa_with_kv_cache_impl(WebGPUGraph& graph, const std::vector<int>& args) {
static_cast<uint64_t>(Hkv) * static_cast<uint64_t>(D);
const uint64_t cache_numel = static_cast<uint64_t>(Cmax) *
static_cast<uint64_t>(Hkv) * static_cast<uint64_t>(D);
const uint64_t aw_floats = static_cast<uint64_t>(Hq) *
static_cast<uint64_t>(state.s) *
static_cast<uint64_t>(state.context_len);
const uint64_t qk_tiles = static_cast<uint64_t>(Hq) *
static_cast<uint64_t>(utils::div_up(state.s, kSdpaTileM)) *
static_cast<uint64_t>(utils::div_up(state.context_len, kSdpaTileN));
const uint64_t softmax_rows =
static_cast<uint64_t>(Hq) * static_cast<uint64_t>(state.s);
const uint64_t av_tiles = static_cast<uint64_t>(Hq) *
static_cast<uint64_t>(utils::div_up(state.s, kSdpaTileM)) *
static_cast<uint64_t>(utils::div_up(D, kSdpaTileN));
if (kv_numel > UINT32_MAX || kv_offset > UINT32_MAX ||
cache_numel > UINT32_MAX || aw_floats > UINT32_MAX ||
qk_tiles > UINT32_MAX || softmax_rows > UINT32_MAX ||
av_tiles > UINT32_MAX) {
cache_numel > UINT32_MAX) {
throw std::runtime_error("WebGPU sdpa: live workload exceeds uint32");
}

Expand All @@ -528,16 +599,43 @@ void sdpa_with_kv_cache_impl(WebGPUGraph& graph, const std::vector<int>& args) {
state.softmax = make_softmax_params(Hq, state.s, state.context_len);
state.compute_out =
make_compute_out_params(state.s, Hq, Hkv, D, state.context_len, g);
if (k16_eligible) {
state.streaming_k16 = make_streaming_attention_k16_params(
state.s, state.context_len, state.pos, Hq, Hkv, D);
state.streaming_k16_grid =
streaming_attention_k16_grid(gr.device(), state.s, Hkv, g);
}
state.update_cache_grid = {
utils::compute_1d_workgroup_count(
gr.device(), static_cast<uint32_t>(kv_numel), uc_wg, "uc(resize)"),
1u};
state.qk_grid = utils::compute_2d_workgroup_count(
gr.device(), static_cast<uint32_t>(qk_tiles), qk_wg, "QK(resize)");
state.softmax_grid = utils::compute_2d_workgroup_count(
gr.device(), static_cast<uint32_t>(softmax_rows), 1, "softmax(resize)");
state.av_grid = utils::compute_2d_workgroup_count(
gr.device(), static_cast<uint32_t>(av_tiles), av_wg, "AV(resize)");
if (!k16_eligible) {
const uint64_t aw_floats = static_cast<uint64_t>(Hq) *
static_cast<uint64_t>(state.s) *
static_cast<uint64_t>(state.context_len);
const uint64_t qk_tiles = static_cast<uint64_t>(Hq) *
static_cast<uint64_t>(utils::div_up(state.s, kSdpaTileM)) *
static_cast<uint64_t>(utils::div_up(state.context_len, kSdpaTileN));
const uint64_t softmax_rows =
static_cast<uint64_t>(Hq) * static_cast<uint64_t>(state.s);
const uint64_t av_tiles = static_cast<uint64_t>(Hq) *
static_cast<uint64_t>(utils::div_up(state.s, kSdpaTileM)) *
static_cast<uint64_t>(utils::div_up(D, kSdpaTileN));
if (aw_floats > UINT32_MAX || qk_tiles > UINT32_MAX ||
softmax_rows > UINT32_MAX || av_tiles > UINT32_MAX) {
throw std::runtime_error(
"WebGPU sdpa: materialized workload exceeds uint32");
}
state.qk_grid = utils::compute_2d_workgroup_count(
gr.device(), static_cast<uint32_t>(qk_tiles), qk_wg, "QK(resize)");
state.softmax_grid = utils::compute_2d_workgroup_count(
gr.device(),
static_cast<uint32_t>(softmax_rows),
1,
"softmax(resize)");
state.av_grid = utils::compute_2d_workgroup_count(
gr.device(), static_cast<uint32_t>(av_tiles), av_wg, "AV(resize)");
}
state.use_fd = fd_eligible && state.s == 1;
// make_sdpa_fd_decode_state requires D % 4 == 0; the op-level guard above
// ("head_dim (D) must be a multiple of 4") rejects any other D before this
Expand All @@ -550,9 +648,10 @@ void sdpa_with_kv_cache_impl(WebGPUGraph& graph, const std::vector<int>& args) {
};

const SdpaLiveState initial_state = compute_live_state(graph);
const uint64_t aw_cap_floats = static_cast<uint64_t>(Hq) *
static_cast<uint64_t>(S) *
static_cast<uint64_t>(dynamic_pos ? Cmax : context_len);
const uint64_t aw_cap_floats = k16_eligible
? 0u
: static_cast<uint64_t>(Hq) * static_cast<uint64_t>(S) *
static_cast<uint64_t>(dynamic_pos ? Cmax : context_len);
const uint64_t aw_bytes = aw_cap_floats * sizeof(float);

WGPUBuffer uc_k_buf = record_update_cache_dispatch(
Expand Down Expand Up @@ -584,7 +683,9 @@ void sdpa_with_kv_cache_impl(WebGPUGraph& graph, const std::vector<int>& args) {
graph.tensor_has_dynamic_dims(v_id);
const bool dual_route =
utils::should_record_sdpa_dual_route(fd_eligible, dynamic_sequence);
const bool record_materialized = dual_route || !initial_state.use_fd;
const bool record_k16 = k16_eligible && (dual_route || !initial_state.use_fd);
const bool record_materialized =
!k16_eligible && (dual_route || !initial_state.use_fd);
const bool record_fd = dual_route || initial_state.use_fd;

WGPUBuffer qk_buf = nullptr;
Expand Down Expand Up @@ -666,21 +767,55 @@ void sdpa_with_kv_cache_impl(WebGPUGraph& graph, const std::vector<int>& args) {
materialized_range.end = graph.num_dispatches();
}

WGPUBuffer k16_buf = nullptr;
size_t k16_idx = 0;
utils::DispatchRange k16_range = {};
if (record_k16) {
k16_range.begin = graph.num_dispatches();
k16_buf = graph.make_uniform_buffer(
&initial_state.streaming_k16, sizeof(StreamingAttentionK16Params));
BufferBinding k16_bindings[4] = {
{out.buffer, out.nbytes},
{q.buffer, q.nbytes},
{k_cache.buffer, k_cache.nbytes},
{v_cache.buffer, v_cache.nbytes}};
const utils::WgCount initial_grid = initial_state.use_fd
? utils::WgCount{0u, 0u}
: initial_state.streaming_k16_grid;
build_dispatch(
graph,
kStreamingAttentionK16CausalBoundWGSL,
k16_bindings,
4,
k16_buf,
sizeof(StreamingAttentionK16Params),
initial_grid.x,
initial_grid.y,
0,
true,
"sdpa_streaming_attention_k16_causal_bound");
k16_idx = graph.num_dispatches() - 1;
k16_range.end = graph.num_dispatches();
}

SdpaFdDecodeResources fd_resources = {};
size_t route_group = 0;
if (record_fd) {
fd_resources = record_sdpa_fd_decode_dispatches(
graph, q, k_cache, v_cache, out, initial_state.fd);
}
if (dual_route) {
const utils::DispatchRange prefill_range =
record_k16 ? k16_range : materialized_range;
route_group = graph.register_dispatch_route_group(
{materialized_range, fd_resources.dispatch_range});
{prefill_range, fd_resources.dispatch_range});
}

auto refresh_state = [compute_live_state,
q_id,
out_id,
dual_route,
record_k16,
record_materialized,
record_fd,
fixed_use_fd = initial_state.use_fd,
Expand All @@ -690,11 +825,13 @@ void sdpa_with_kv_cache_impl(WebGPUGraph& graph, const std::vector<int>& args) {
qk_idx,
softmax_idx,
av_idx,
k16_idx,
uc_k_buf,
uc_v_buf,
qk_buf,
softmax_buf,
av_buf,
k16_buf,
fd_resources](WebGPUGraph& gr) {
const SdpaLiveState state = compute_live_state(gr);

Expand Down Expand Up @@ -722,6 +859,14 @@ void sdpa_with_kv_cache_impl(WebGPUGraph& graph, const std::vector<int>& args) {
wgpuQueueWriteBuffer(
gr.queue(), av_buf, 0, &state.compute_out, sizeof(state.compute_out));
}
if (record_k16) {
wgpuQueueWriteBuffer(
gr.queue(),
k16_buf,
0,
&state.streaming_k16,
sizeof(state.streaming_k16));
}
if (record_fd) {
write_sdpa_fd_decode_uniforms(gr.queue(), fd_resources, state.fd);
}
Expand All @@ -735,8 +880,10 @@ void sdpa_with_kv_cache_impl(WebGPUGraph& graph, const std::vector<int>& args) {
const std::vector<utils::WgCount> active_grids = state.use_fd
? std::vector<
utils::WgCount>{state.fd.split_grid, state.fd.reduce_grid}
: std::vector<utils::WgCount>{
state.qk_grid, state.softmax_grid, state.av_grid};
: (record_k16
? std::vector<utils::WgCount>{state.streaming_k16_grid}
: std::vector<utils::WgCount>{
state.qk_grid, state.softmax_grid, state.av_grid});
gr.select_dispatch_route(route_group, active_route, active_grids);
} else if (state.use_fd) {
if (!fixed_use_fd) {
Expand All @@ -750,6 +897,12 @@ void sdpa_with_kv_cache_impl(WebGPUGraph& graph, const std::vector<int>& args) {
state.fd.reduce_grid.x;
gr.dispatch_at(fd_resources.dispatch_range.begin + 1).workgroup_count_y =
state.fd.reduce_grid.y;
} else if (record_k16) {
if (fixed_use_fd) {
throw std::runtime_error("WebGPU sdpa: static route changed");
}
gr.dispatch_at(k16_idx).workgroup_count_x = state.streaming_k16_grid.x;
gr.dispatch_at(k16_idx).workgroup_count_y = state.streaming_k16_grid.y;
} else {
if (fixed_use_fd) {
throw std::runtime_error("WebGPU sdpa: static route changed");
Expand Down
Loading
Loading