From 8255516eb479630991f4b4a2a4c42f732b5d97aa Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mateusz=20S=C5=82uszniak?= Date: Tue, 1 Sep 2026 11:16:32 +0200 Subject: [PATCH] Vulkan: size texture-vs-buffer choice by the bound, not the trace hint filter_invalid_reprs() decides whether a tensor can live in a texture by comparing required_image_extents(tensor_val.shape) against the device limits. For a symbolic dimension that comparison resolves at the dim's *hint* -- the size of the example the model happened to be traced with -- rather than the maximum its exported range allows. So a model traced with a small example picks a texture that is only valid at that size, and silently produces wrong results once it runs larger. Nothing raises. Repro: the Supertonic TTS text encoder, dynamic T in [8, 512], always executed at T=512 on a Galaxy S26 Ultra (Adreno 840). Only the trace-time example length differs; the edge graphs are identical (552 nodes, same ops, same literal args, same symbolic shapes): traced at T= 64 cosine 0.537966 .pte 18147970 bytes traced at T=128 cosine 0.537966 .pte 18147970 bytes (identical file) traced at T=192 cosine 0.999414 .pte 18149890 bytes traced at T=512 cosine 0.999414 .pte 18149890 bytes (identical file) The emitted binaries cluster exactly on the correctness boundary, which is what identifies this as a lowering decision rather than a bad kernel. With this change the T=64 export emits the 18149890-byte artifact and scores 0.999414. Evaluating the bound instead also covers the unbounded case: a dim with no finite upper bound cannot be shown to fit any texture, so it falls back to buffer storage. Measured on Supertonic (CPU reference reproduced by the XNNPACK delegate at cosine 1.000000): text_encoder 0.406303 -> 0.999414 vector_estimator 0.994432 -> 0.999994 vocoder 0.016757 -> 0.999977 --- backends/vulkan/utils.py | 41 ++++++++++++++++++++++++++++++++++++---- 1 file changed, 37 insertions(+), 4 deletions(-) diff --git a/backends/vulkan/utils.py b/backends/vulkan/utils.py index 84b901b6b6e..19aaebdf144 100644 --- a/backends/vulkan/utils.py +++ b/backends/vulkan/utils.py @@ -1181,6 +1181,33 @@ def make_tensor_repset(tensor_repr: TensorRepr) -> TensorRepSet: raise RuntimeError(f"Unsupported storage type {tensor_repr.storage_type}") +def upper_bound_size(dim: Union[int, torch.SymInt]) -> Optional[int]: + """Largest value a (possibly symbolic) tensor dimension can take. + + Returns None if no finite bound is known. + + A symbolic dim compares against a limit using its *hint* -- the size of the + example input the model happened to be traced with -- not the maximum the + exported range allows. Sizing decisions must use the bound instead, or a + model traced with a small example will make a choice that is invalid once it + runs at a larger size. + """ + if not isinstance(dim, torch.SymInt): + return int(dim) + if not dim.node.expr.free_symbols: + return int(dim.node.expr) + shape_env = dim.node.shape_env + if shape_env is None: + return None + try: + upper = shape_env.bound_sympy(dim.node.expr).upper + except Exception: + return None + if upper is None or not upper.is_finite: + return None + return int(upper) + + def filter_invalid_reprs( tensor_val: FakeTensor, tensor_repset: TensorRepSet, @@ -1198,11 +1225,17 @@ def filter_invalid_reprs( can be used to produce a valid image texture for the given tensor (i.e. fits within texture limits). """ + # Size the texture by what the dimension CAN be, not by the example the + # model was traced with. An unbounded dim cannot be shown to fit, so it + # falls back to buffer storage. + bounds = [upper_bound_size(d) for d in tensor_val.shape] valid_texture_layouts = set() - for memory_layout in tensor_repset.valid_texture_layouts: - extents = required_image_extents(tensor_val.shape, memory_layout) - if extents_are_valid(extents, texture_limits): - valid_texture_layouts.add(memory_layout) + if all(b is not None for b in bounds): + max_shape = torch.Size(bounds) + for memory_layout in tensor_repset.valid_texture_layouts: + extents = required_image_extents(max_shape, memory_layout) + if extents_are_valid(extents, texture_limits): + valid_texture_layouts.add(memory_layout) # High dimensional tensors require buffer storage if len(tensor_val.shape) > 4: