Repository navigation
fix(conversion): handle slice_scatter edge cases and write a contiguous slice_scatter as a concatenation - #4898
Open
shoumikhin wants to merge 8 commits into
Open
shoumikhin wants to merge 8 commits into
shoumikhin wants to merge 8 commits into
Conversation
github-actions
Bot
added
documentation
shoumikhin
force-pushed
the
slice-scatter-concat
branch
5 times, most recently
from
October 24, 2026 22:56
05ec618 to
99e031c
Compare
shoumikhin
changed the title
shoumikhin
force-pushed
the
slice-scatter-concat
branch
from
October 25, 2026 23:19
99e031c to
2466df3
Compare
shoumikhin
changed the title
lanluo-nvidia
left a comment
lanluo-nvidia
left a comment
Collaborator
There was a problem hiding this comment.
lgtm, i requested @apbose for a final review and approve
apbose
reviewed
Oct 27, 2026
| else try_emit_kv_cache_update( | ||
| ctx, name, input, src, dim, start, update_len, step | ||
| ) | ||
| ) |
Collaborator
There was a problem hiding this comment.
Could we cast src to input.dtype before try_emit_kv_cache_update(), since TensorRT requires the KV cache and update tensors to have the same dtype? This is in the static bounds kv-cache case
Contributor
Author
There was a problem hiding this comment.
Good catch, thanks. It can happen: torch.slice_scatter accepts a src of another dtype, and a float16 update written into a float32 cache failed to build with "IKVCacheUpdateLayer update and output tensors must have identical types".
The converter now casts src to the input's dtype once, before both the KV cache path and the full overwrite, in "fix(conversion): write slice_scatter's src in the input's dtype". test_update_of_another_dtype in test_aliased_io.py covers it and fails without the cast.
| # only part is src. The scatter keeps the input's shape. | ||
| if step == 1 and not full_overwrite and src.shape[dim] != DYNAMIC_DIM: | ||
| return _concat_write(ctx, target, source_ir, name, input, src, dim, start, end) | ||
|
|
Collaborator
There was a problem hiding this comment.
Given the reported 1.35×–1.96× H100 regression for the rotary-write microbenchmark, do we have an H100 end-to-end result confirming that selecting concatenation unconditionally is still beneficial? Since we are replacing this with concat in ecvery case
Contributor
Author
There was a problem hiding this comment.
Yes. I measured end to end on an H100, with one build and only the concat commit switched on and off:
- SmolVLA through ExecuTorch + TensorRT: 14.9 ms per call with the scatter, 13.5 ms with the concat. Both match eager to the same error.
- A 16-layer Llama-style decoder with the same rotary, 64 to 2048 tokens: the concat takes 16% to 22% less time in FP16, and up to 14% less in FP32 (equal at 2048 tokens).
- pi0.5 has no such writes, so it does not change.
The 1.35x to 1.96x slowdown came from a chain of rotary writes that nothing else reads. None of these models was slower, so I kept the concat unconditional.
…atter lowering When a slice_scatter has a symbolic start, end or step, a lowering pass rewrites it into arange, expand and scatter. torch.slice_scatter accepts a src of another dtype and writes it in the input's dtype, but aten.scatter requires both dtypes to match. So a float32 cache written with a float16 src failed to compile with "scatter(): Expected self.dtype to be equal to src.dtype". The pass now casts src to the input's dtype first when the two differ.
`x[:, -1] = y` exports as select_scatter with index -1. Torch-TensorRT decomposes select_scatter into a one-wide slice_scatter from index to index + 1. For index -1 that is the slice -1:0, which is empty, so the write did not happen as intended: the converter failed with a numpy broadcast error. The decomposition now adds the dim size to a negative int index first, so the slice is -1 + size to size. An int index outside a static dim raises IndexError, as torch does, instead of landing in range. On a dynamic dim the index is not checked. A non-negative int index and an index that is only known at run time (a SymInt, like `x.shape[1] - 3` or `idx.item()`) are used as before. `x[:, -1] = y`, `x[..., -1] = y` and `x[-1] = y` now compile and match PyTorch, also when the written dim is dynamic.
TensorRT refuses to build an engine where one tensor is both an input and an output, and a tensor can be marked as an output only once. Both happen when a graph output is a tensor that is already an engine input or already an output: - a converter returns one of its inputs as is (slice_scatter does that for a write that covers the whole dim, since the result is just src); - the graph returns an input next to a computed output, like `return x + y, y`, with require_full_compilation=True; - the graph returns the same tensor twice, like `z = x + y; return z, z`. On main the first two fail. When the input is the only output, TensorRT refuses to build the engine. Next to another output, the engine builds, and then fails to load or run because that output has no binding. The interpreter skipped a tensor it had already marked, so the third builds an engine with one output fewer than the graph, and the module returns a single tensor instead of a tuple. The repair_input_as_output lowering pass is meant to catch a returned input by inserting a clone, but it does not fire on exported graphs. It looks for placeholders whose node.type is a Tensor, and torch.export leaves node.type unset, so the pass finds no placeholders and the graph still ends in `return (add, y)`. The interpreter now marks a copy of such a tensor at every output position, so the engine always has one output per position. The copy is an identity layer, which TensorRT allows for every dtype, while a cast is refused on FP8 tensors. An engine loaded from the engine cache now counts its outputs the same way; it counted distinct tensors, which made a cached `return x + y, y, y` return two outputs where a fresh build returns three. A converter cannot make the copy itself: it does not know whether its result will be an output, and a copy layer inside the engine is not always allowed. TensorRT takes a uint8 tensor only as an engine input or output, so `slice_scatter(x, src, 1, 0, 8).to(torch.float32)` on uint8 tensors would fail to build if the converter always copied src.
torch.slice_scatter writes src in the input's dtype. Two paths of the converter did not: - For a write that covers the whole dim, it returned src as is, so with an int32 input and a float32 src the result stayed float32. For example, with `src = torch.full((4, 8), 2.7)`, `slice_scatter(x, src, 1, 0, 8) * 3` gave 8 where PyTorch gives 6. - The KV cache path passed src to IKVCacheUpdateLayer as is. That layer requires the cache and the update to have the same dtype, so a float16 src written into a float32 cache failed to build the engine. The converter now casts src to the input's dtype once, before both paths. When the dtypes already match, the cast adds no layer, and a full overwrite still returns src itself. The scatter fallback already cast.
torch raises when src does not have the shape of the slice it writes. The converter did not check. Most such writes failed with a numpy broadcast error that did not say what was wrong. A strided write over a one-element slice built without error and wrote the wrong values, and a write over the whole dim returned src at its own shape, also when src had another rank or another size on a dim that is not written. The converter now checks src against the shape of the slice, for every kind of write: the same rank, and the same size on every dim whose size is known. It raises ValueError if they differ.
…rite A write that covers the whole dim returns src. When TensorRT sees src as dynamic on a dim where the input is static, for example after a reshape with -1, the result is then dynamic there too, although the graph records the input's static size. A later slice_scatter on that dim cannot resolve its bounds and raises NotImplementedError, even though the validator accepted it. This happens whether the dynamic dim is the written dim or another one. Such a write now reshapes src to the input's sizes, which keeps the static ones static. A scatter would have kept the shape too, but TensorRT has no int8 scatter, while an int8 reshape builds. TensorRT cannot reshape uint8, so the reshape runs in whichever of the two dtypes is not uint8: after the cast for a uint8 src, before it for a uint8 input. When both are uint8, src keeps its own shape, as before. The reshape reads a static size of 0 as 0 rather than as "copy this dim".
…nput
The KV cache layer updates its cache in place, so the slice_scatter
converter uses it only when the cache is an engine input. It decided that
from the TensorRT tensor. A converter can hand back its input as is, clone
for example, so a write on such a result saw the engine input itself and
updated the caller's tensor, while PyTorch leaves it unchanged:
out = cache.clone()
out[:, :, 3:4] = update
return out # the caller's cache changed as well
The converter now also requires the write's input in the graph to be a
graph input, as the index_copy KV path already does. A clone that a
lowering pass removes before conversion is not covered: the write is then
on the graph input itself.
When a slice_scatter cannot use the KV cache layer, the converter builds an index tensor as large as the written values and emits a scatter. In the engines we profiled, TensorRT ran each scatter as its own kernel, apart from the math that produces the values. The rotary embedding in SmolVLA does `res[..., :half] = a; res[..., half:] = b` in every attention layer, which became more than 500 scatter kernels per call. A write with step 1 is the input before the slice, the written values, and the input after it, joined along the written dim. The converter now emits that concatenation when src has a static size on that dim. An empty slice returns the input itself. That relies on the interpreter giving every output position its own tensor, and on the KV cache path taking only a write on a graph input. Strided writes and writes whose src is dynamic on the written dim keep the scatter, which keeps the input's static shape. A concatenation takes its size from its parts, so it relies on the earlier size check to reject a src that does not fit the slice. With the change, an ExecuTorch export of SmolVLA ran 130.5 ms to 125.3 ms per call on a Jetson Orin Nano, 28.7 ms to 27.0 ms on a Jetson AGX Thor, and 14.9 ms to 13.5 ms on an H100, with the same outputs. On an H100, a 16-layer Llama-style decoder with this rotary on q and k ran 16% to 22% faster in FP16 and up to 14% faster in FP32 (the same at 2048 tokens), from 64 to 2048 tokens. pi0.5 has no such writes and does not change. Only a chain of rotary writes that nothing else reads got slower: up to about 10% on an A100, and 1.35x to 1.96x on an H100 at 64 to 256 tokens.
shoumikhin
force-pushed
the
slice-scatter-concat
branch
from
October 27, 2026 09:33
2466df3 to
dc2c2c2
Compare
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What is wrong today
aten.slice_scatterwritessrcinto a slice of a tensor, likey[..., a:b] = src;select_scatter(y[:, i] = src) becomes one. Some writes give a wrong result or do not compile, likex[:, -1] = yor asrcof another data type (dtype).Also, outside the key/value (KV) cache case, every partial write is a scatter, its own TensorRT kernel: over 500 per call in SmolVLA's rotary embedding.
What this change does
Each fix is one commit (item 1 has two), with a test that fails without it:
srcof another dtype is written in the input's dtype, as in PyTorch, also with symbolic bounds, for a whole-dim write, and on the KV cache path.select_scatterindex is counted from the end, sox[:, -1] = yno longer becomes the empty slice-1:0.return z, zreturned one tensor.srcthat does not match the slice's shape raises a clearValueError, not a numpy error or a wrong result.srcis dynamic, so later writes compile.x.clone(), which changed the caller'sx.The last commit is the speedup: a step-1 write whose
srcis static on the written dim becomescat(input[:start], src, input[end:]). Other writes keep the scatter.End to end:
Only a microbenchmark chain of rotary writes that nothing else reads got slower (1.35x to 1.96x on an H100 at 64 to 256 tokens).
What was tested
H100, TensorRT 11.4, PyTorch 2.16 nightly: the slice_scatter, lowering, decomposition, buffer lifting and runtime tests pass (483, without hardware compatibility, cross-compile and weight streaming). Each commit passes its own tests.