Skip to content

Upstream the Gemma 4 Vulkan support we carry on our ExecuTorch fork #1361

Closed

Description

@msluszniak

Our ExecuTorch fork (software-mansion-labs/executorch) carries Gemma 4 Vulkan
support that upstream does not have. It should go upstream so we stop carrying
it, but it is a re-port rather than a patch, so it needs planning.

Context

Commit 24f0f69 "Add Gemma4 Vulkan support (Android GPU)" is 45 files,
+4247/-1188. It was already a port ("Ports Mateusz Kopciński's Gemma4 Vulkan
work onto 1.3"), and the upstream files it rewrites have moved again since
1.3.1. Upstream now has its own RMSNorm.cpp and IndexSelect.cpp, and its
SDPA.cpp has evolved independently, so our version collides rather than
applies.

Everything else SWM carries has now been checked against upstream and is either
already there or has been sent:

So Gemma 4 Vulkan is the only substantial piece left on the fork.

Suggested split

Four independent pieces, easiest first:

  1. Shader precision/perf helpers — linear_q4gsw / linear_dq8ca_q4gsw
    changes plus sdpa_fp32_acc.glslh and linear_fp32_scale_dq_acc.glslh.
    Additive, likely stand alone.
  2. The aten.rms_norm lowering (replacing the legacy et_vk.rms_norm).
    Upstream already has an RMSNorm.cpp; this needs a diff against theirs first
    to establish whether ours is an improvement, a duplicate, or a regression.
  3. GemmaSDPA.cpp + the gemma SDPA shaders — the actual feature, ~738 lines
    plus ~10 shader files. New files, so it can land as an addition, but it
    depends on the reworked SDPA tile-load helpers, which is exactly where it
    collides with upstream's evolved SDPA.cpp. This is the bulk of the work.
  4. Do not upstream — the gen_vulkan_spv.py backward-compat shim (lets
    ported shaders call define_required_extensions with one arg; a porting
    artifact) and the DispatchNode.cpp tracer, see below.

Separate cleanup, worth doing regardless

The same commit added an env-gated dispatch tracer to
backends/vulkan/runtime/graph/ops/DispatchNode.cpp — ETVK_DISPATCH_TRACE,
with <cstdio>/fopen/setvbuf, commented "Iter 21 instrumentation". It is
still in the fork HEAD, so it ships in every runtime we build, including the one
RNE distributes. Inert unless the env var is set, but it is leftover debugging
in a shipped runtime and should probably come out of the fork.

Activity

  1. self-assigned this
    on Aug 14, 2026
  2. added
    platform: androidIssues and tasks related to Android
    performanceRelated to all issues and tasks focused on improving performance
    on Aug 14, 2026
  3. added theissue type on Aug 14, 2026
  4. msluszniak commented on Oct 20, 2026

    @msluszniak
    MemberAuthor

    Upstream PRs: pytorch/executorch#23536 (merged), pytorch/executorch#23539, pytorch/executorch#24023.

    Exporter changes to export Gemma 4 with upstream ET:

    • Replace et_vk::gemma_sdpa_with_kv_cache with llama.sdpa_with_kv_cache (scale=1.0, sliding mask squeezed to [S, max_context_len]).
    • Keep weight scales in fp16 (lossless for Gemma 4) instead of fp32.
    • Keep the _edge_adapters dtype promotion and _to_dim_order_copy rewrite; they work around upstream _to_copy and embedding_q4gsw fp16 bugs.

    The fp32 scale and accumulator shaders are not needed: no measurable accuracy gain on S26.

  5. msluszniak commented on Oct 20, 2026

    @msluszniak
    MemberAuthor

    Since all needed changes are now as upstream PRs and only changes are in exports I'm closing this issue as comleted.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

Labels

performanceRelated to all issues and tasks focused on improving performanceplatform: androidIssues and tasks related to Android

Type

Projects

No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions