[Vulkan] Missing/partial op support blocks full delegation of a dynamic-shape transformer encoder

[mergennachin]

🐛 Describe the bug

I'm adding a Vulkan option to examples/models/nemotron3_diarization, which has a 31-layer transformer encoder with a dynamic sequence length and runs in FP32. The encoder can't be delegated as a single Vulkan partition. The current workarounds are:

• eager attention with an additive mask;
• RemoveRedundantOpsTransform before partitioning;
• aten.gelu in operator_blocklist, so it falls back to XNNPACK.

With these, the encoder lowers to 32 Vulkan + 31 XNNPACK partitions and the process peaks at about 3.5 GiB. When every GELU stays on Vulkan, peak is about 0.7 GiB.

Checked at 00e5030. backends/vulkan/op_registry.py is unchanged on current main.

Repro

Partitioning only; no GPU needed.

import collections

import torch
import torch.nn.functional as F
from executorch.backends.vulkan.partitioner import vulkan_partitioner as vp
from executorch.backends.vulkan.partitioner.vulkan_partitioner import VulkanPartitioner
from executorch.exir import EdgeCompileConfig, to_edge_transform_and_lower
from torch.export import Dim, export

skips = collections.Counter()
_log_skip = vp.VulkanSupportedOperators.log_skip


def log_skip(self, node, reason):
    if node.op == "call_function":
        name = vp.utils.node_io_str(node).split(" = ")[1].split("(")[0]
        skips[(name, reason)] += 1
    _log_skip(self, node, reason)


vp.VulkanSupportedOperators.log_skip = log_skip


class Block(torch.nn.Module):
    def __init__(self, sdpa):
        super().__init__()
        self.sdpa = sdpa
        self.qkv = torch.nn.Linear(64, 192)
        self.ff = torch.nn.Linear(64, 64)

    def forward(self, x, lengths):
        b, s, _ = x.shape
        mask = (torch.arange(s)[None, :] < lengths[:, None])[:, None, None, :]
        q, k, v = self.qkv(x).view(b, s, 3, 2, 32).permute(2, 0, 3, 1, 4)
        if self.sdpa:
            y = F.scaled_dot_product_attention(q, k, v, attn_mask=mask)
        else:
            bias = torch.where(mask, 0.0, -torch.inf)
            y = torch.softmax(q @ k.transpose(-1, -2) * 32**-0.5 + bias, -1) @ v
        return F.gelu(self.ff(y.transpose(1, 2).reshape(b, s, 64)))


for sdpa in (False, True):
    skips.clear()
    ep = export(
        Block(sdpa).eval(),
        (torch.randn(1, 16, 64), torch.tensor([16])),
        dynamic_shapes=({1: Dim("s", min=2, max=1000)}, {}),
        strict=False,
    )
    edge = to_edge_transform_and_lower(
        ep,
        partitioner=[VulkanPartitioner({"require_dynamic_shapes": True})],
        compile_config=EdgeCompileConfig(_check_ir_validity=False),
    )
    graph = edge.exported_program().graph_module.graph
    delegates = sum("call_delegate" in str(n.target) for n in graph.nodes)
    print(f"\n== {'sdpa' if sdpa else 'eager'}: {delegates} Vulkan delegate(s)")
    for (name, reason), count in sorted(skips.items()):
        print(f"  {count}x {name}: {reason}")

Output:

== eager: 3 Vulkan delegate(s)
  4x aten.expand_copy.default: no dynamic shape support
  2x aten::scalar_tensor: no operator implementation

== sdpa: 4 Vulkan delegate(s)
  1x aten.any.dim: no operator implementation
  4x aten.expand_copy.default: no dynamic shape support
  1x aten.full_like.default: no dynamic shape support
  2x aten.mul.Scalar: no operator implementation
  2x aten::scalar_tensor: no operator implementation

Exact GELU is missing from this output because Vulkan accepts it (see item 1).

Gaps

1. aten.gelu with approximate="none" is computed with the tanh approximation.

   • The registry accepts every GELU: op_registry.py#L216.
   • gelu() ignores approximate: UnaryOp.cpp#L179-L185.
   • The shader always uses tanh: unary_op.yaml#L33-L34.

   This changes numerics without any warning. With every GELU on Vulkan, the encoder's output probabilities differed from the reference by up to 6.3e-3 (tested on MoltenVK), against a 1e-4 FP32 tolerance.

   Suggested fix: add an erf-based variant selected from args[1]. GLSL has no erf, so it needs a polynomial approximation, e.g. Abramowitz–Stegun 7.1.26 (max error ≈1.5e-7). Until then, the partitioner could reject approximate != "tanh".

2. aten.scalar_tensor is never partitioned.

   • EXIR keeps scalar_tensor as an ATen op in the edge graph: replace_aten_with_edge_pass.py#L17-L21.
   • Vulkan only registers exir_ops.edge.aten.scalar_tensor.default: op_registry.py#L1569-L1575.

   The partitioner therefore reports "no operator implementation", even though a kernel exists: ScalarTensor.cpp#L51.

   Registering torch.ops.aten.scalar_tensor.default is not enough by itself. The graph builder serializes node.target.__name__ (vulkan_graph_builder.py#L486). For an ATen overload that is scalar_tensor.default, but the runtime looks up aten.scalar_tensor.default.

   It shows up from torch.where(mask, 0.0, -inf) and from the SDPA decomposition.

3. aten.expand_copy has supports_resize=False: op_registry.py#L1315.

   • With require_dynamic_shapes=True, every expand is rejected. That includes static-shape expands, such as HF RoPE's inv_freq[None, :, None].expand(...), and the same-shape expands from matmul decomposition.
   • The runtime already implements resizing: Expand.cpp#L19-L49.
   • This may only need the flag plus a dynamic-shape test.

4. aten.full / full_like / zeros* / ones* don't set supports_resize: op_registry.py#L1547-L1562.

   • The runtime already implements resizing: Full.cpp#L19-L32.
   • Same fix as 3.

5. Decomposed SDPA needs two ops that aren't registered:

   • aten.mul.Scalar: the decomposition scales both q and k by sqrt(scale).
   • aten.any.dim on bool: comes from _safe_softmax, along with full_like, eq.Scalar and logical_not.

   Instead of adding these one by one, a better option may be to lower aten.scaled_dot_product_attention to the existing fused et_vk.sdpa. It takes q, k, v, an additive mask and a scale, and supports resizing: SDPA.cpp#L890-L999.

   Nothing produces et_vk.sdpa today. SDPA is not in ops_not_to_decompose (vulkan_partitioner.py#L48-L52), and no pass rewrites it. Bool masks would need converting to additive masks.

Expected impact

I tested a 2-layer model built from the real Nemotron config, with the same per-layer ops, using eager attention and RemoveRedundantOpsTransform. Leaving GELU on Vulkan gives a single encoder delegate; only scalar_tensor stays outside. So item 1 alone restores single-partition delegation for this model. Items 2–5 would remove the need for eager attention and the extra pass.

This issue was drafted with Claude Code. The repro output above is from an actual run.

Versions

ExecuTorch at 00e5030 (Vulkan registry identical on current main), macOS arm64. Runtime numbers from MoltenVK on Apple M1 Pro.

cc @SS-JIA @manuelcandales @digantdesai @cbilgin