From 67c4d9913d53d44e9b4f0291a1070c9d14e3844a Mon Sep 17 00:00:00 2001 From: apbose Date: Fri, 21 Aug 2026 16:26:08 -0700 Subject: [PATCH] Fix strided slice_scatter KV cache updates --- .../dynamo/conversion/impl/slice_scatter.py | 9 ++++---- tests/py/dynamo/runtime/test_aliased_io.py | 23 +++++++++++++++++++ 2 files changed, 28 insertions(+), 4 deletions(-) diff --git a/py/torch_tensorrt/dynamo/conversion/impl/slice_scatter.py b/py/torch_tensorrt/dynamo/conversion/impl/slice_scatter.py index 13c19f5e8b..41a4c50e40 100644 --- a/py/torch_tensorrt/dynamo/conversion/impl/slice_scatter.py +++ b/py/torch_tensorrt/dynamo/conversion/impl/slice_scatter.py @@ -206,10 +206,11 @@ def slice_scatter( update_len = end - start - # KV fast path. - kv_out = try_emit_kv_cache_update(ctx, name, input, src, dim, start, update_len) - if kv_out is not None: - return kv_out + # IKVCacheUpdateLayer only supports contiguous updates. + if step == 1: + kv_out = try_emit_kv_cache_update(ctx, name, input, src, dim, start, update_len) + if kv_out is not None: + return kv_out # Fallback: build broadcast indices and scatter. indices_np: np.ndarray = np.arange(start, end, step, dtype=np.int64) diff --git a/tests/py/dynamo/runtime/test_aliased_io.py b/tests/py/dynamo/runtime/test_aliased_io.py index 86fe8dfdf7..aa704b1ca9 100644 --- a/tests/py/dynamo/runtime/test_aliased_io.py +++ b/tests/py/dynamo/runtime/test_aliased_io.py @@ -54,6 +54,29 @@ def _find_aliased_io(compiled): return {} +class TestSliceScatterFallback(TestCase): + def test_strided_slice_uses_scatter_fallback(self): + class M(torch.nn.Module): + def forward(self, cache, update): + out = torch.ops.aten.slice_scatter.default(cache, update, 2, 0, 16, 2) + return out + 0 + + cache = torch.zeros(2, 4, 16, 8, device="cuda") + update = torch.ones(2, 4, 8, 8, device="cuda") + compiled = _compile_cpp(M().cuda(), (cache.clone(), update.clone())) + + # A strided update is not eligible for the contiguous KV-cache layer. + self.assertEqual(_find_aliased_io(compiled), {}) + + cache_run = cache.clone() + actual = compiled(cache_run, update) + actual = actual[0] if isinstance(actual, tuple) else actual + expected = torch.ops.aten.slice_scatter.default(cache, update, 2, 0, 16, 2) + + self.assertTrue(torch.allclose(actual, expected)) + self.assertTrue(torch.equal(cache_run, cache)) + + class TestUserInputKVCache(TestCase): """User passes the cache tensor on every call; engine mutates in place."""