diff --git a/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py b/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py index e7634dfa66..559cd1543c 100644 --- a/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py +++ b/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py @@ -1224,7 +1224,8 @@ def _index_copy_kv_eligible( if len(node.args) < 4: return False if input_node is None: - input_node = node.args[0] # type: ignore[assignment] + cache_arg = node.args[0] + input_node = cache_arg if isinstance(cache_arg, Node) else None dim, _index_node, src_node = node.args[1:4] if not isinstance(input_node, Node) or input_node.op != "placeholder": @@ -1309,8 +1310,47 @@ def aten_ops_index_copy_fallback( ) +def slice_scatter_validator( + node: Node, settings: Optional[CompilationSettings] = None +) -> bool: + """Reject a write whose bounds need the size of the dim it writes -- a negative + index, or the open end torch.export writes for ``x[..., start:]`` -- when that dim + is dynamic. The converter cannot resolve those (see ``resolve_slice_scatter_write``) + and raises, so they run in PyTorch until it gains dynamic bounds. + + Missing metadata is passed, not rejected: the KV-cache classifier in + ``lowering/_buffer_lifting.py`` reads the same metadata, and vetoing a write it + classified as engine-aliased fails ``assert_predicted_kv_aliased``. + """ + input_meta = getattr(node.args[0], "meta", {}) + input_val = input_meta.get("val", input_meta.get("tensor_meta")) + if input_val is None: + _LOGGER.debug( + f"slice_scatter node {node.name} has no shape metadata; leaving its bounds " + "for the converter to resolve against the TensorRT shape." + ) + return True + + _start, _end, _step, status = impl.slice_scatter.resolve_slice_scatter_write( + tuple(input_val.shape), + args_bounds_check(node.args, 2, 0), + args_bounds_check(node.args, 3), + args_bounds_check(node.args, 4), + args_bounds_check(node.args, 5), + ) + if status is impl.slice_scatter.KVWriteStatus.DYNAMIC_DIM_SIZE: + _LOGGER.debug( + f"slice_scatter node {node.name} needs the size of a dynamic dim to " + "resolve its bounds; falling back to PyTorch operation." + ) + return False + return True + + @dynamo_tensorrt_converter( - torch.ops.aten.slice_scatter.default, supports_dynamic_shapes=True + torch.ops.aten.slice_scatter.default, + capability_validator=slice_scatter_validator, + supports_dynamic_shapes=True, ) @enforce_tensor_types({0: (TRTTensor,), 1: (TRTTensor,)}) def aten_ops_slice_scatter( diff --git a/py/torch_tensorrt/dynamo/conversion/impl/slice_scatter.py b/py/torch_tensorrt/dynamo/conversion/impl/slice_scatter.py index a6922dbfdf..154ba61414 100644 --- a/py/torch_tensorrt/dynamo/conversion/impl/slice_scatter.py +++ b/py/torch_tensorrt/dynamo/conversion/impl/slice_scatter.py @@ -19,6 +19,7 @@ from __future__ import annotations import logging +import sys from enum import Enum, auto from typing import Any, Optional, Tuple @@ -79,9 +80,25 @@ class KVWriteStatus(Enum): OK = auto() FULL_OVERWRITE = auto() DYNAMIC_BOUNDS = auto() + DYNAMIC_DIM_SIZE = auto() BAD_DIM = auto() +# torch.export uses INT64_MAX as "open end". `sys.maxsize` is the same value and is +# what the aten.slice converter matches on, so both are accepted. +_OPEN_END = (sys.maxsize, 2**63 - 1) + + +def _needs_dim_size(bound: Any) -> bool: + """Whether ``bound`` has to be resolved against the dim being written: ``None`` and + an open end run to it, a negative index counts back from it. A non-int bound is + reported as ``DYNAMIC_BOUNDS`` instead, so it passes here. + """ + if bound is None: + return True + return isinstance(bound, int) and (bound < 0 or bound in _OPEN_END) + + def resolve_slice_scatter_write( cache_shape: Tuple[Any, ...], dim: Any, @@ -92,9 +109,10 @@ def resolve_slice_scatter_write( """Resolve ``slice_scatter``'s slice bounds into the form ``_kv_eligible`` takes. Returns ``(start, end, step, status)``. Under ``OK`` and ``FULL_OVERWRITE`` the - three bounds are Python ``int``s, with the op's defaults filled in and negative - indices counted from the end; under ``DYNAMIC_BOUNDS`` and ``BAD_DIM`` all three - are ``None``, since nothing resolved. ``status`` is one of: + three bounds are Python ``int``s, with the op's defaults filled in, negative + indices counted from the end and the whole slice clamped to the dim; under the + other statuses all three are ``None``, since nothing resolved. ``status`` is one + of: * ``OK`` — the bounds are concrete, and the caller goes on to ``_kv_eligible(cache_shape, dim, start, end - start)``. @@ -105,6 +123,10 @@ def resolve_slice_scatter_write( under this status. * ``DYNAMIC_BOUNDS`` — a bound is not a Python ``int``, so the converter raises ``NotImplementedError``. + * ``DYNAMIC_DIM_SIZE`` — a bound needs the size of the dim it writes (see + :func:`_needs_dim_size`) and that dim is dynamic. ``aten_ops_slice_scatter``'s + validator rejects these so the partitioner runs them in PyTorch; the converter + raises if one reaches it anyway. * ``BAD_DIM`` — ``dim`` is either not a Python ``int`` or does not index ``cache_shape``; the converter raises ``IndexError`` for both. A ``numpy.int64`` is rejected on the type check even when its value is in range. @@ -123,25 +145,40 @@ def resolve_slice_scatter_write( if not isinstance(dim, int) or not -len(cache_shape) <= dim < len(cache_shape): return None, None, None, KVWriteStatus.BAD_DIM dim_size = cache_shape[dim] + # A dynamic dim reaches the converter as DYNAMIC_DIM (-1) and the predictor as a + # SymInt; neither is a size, and folding the -1 into an index is how an open-ended + # write on a dynamic dim used to resolve to an empty index range. + dim_is_static = isinstance(dim_size, int) and dim_size >= 0 if start is None: start = 0 - if isinstance(start, int) and start < 0 and isinstance(dim_size, int): - start = dim_size + start - if end is None: - end = dim_size - if isinstance(end, int) and end < 0 and isinstance(dim_size, int): - end = dim_size + end if step is None: step = 1 + if not dim_is_static and (_needs_dim_size(start) or _needs_dim_size(end)): + return None, None, None, KVWriteStatus.DYNAMIC_DIM_SIZE + + if end is None: + end = dim_size + if dim_is_static: + if isinstance(start, int) and start < 0: + start = dim_size + start + if isinstance(end, int) and end < 0: + end = dim_size + end + # Aten clamps a slice to its dim, and so must this: unclamped, the open end + # reaches the fallback's np.arange as an INT64_MAX-long index range. + if isinstance(start, int): + start = min(max(start, 0), dim_size) + if isinstance(end, int): + end = min(max(end, 0), dim_size) + # A slice covering the whole dim is a plain copy of the source whatever `step` is # made of, so it is settled before the bounds are required to be concrete: `step` # only has to compare equal to 1, which a symbolic step can do. if ( isinstance(start, int) and isinstance(end, int) - and isinstance(dim_size, int) + and dim_is_static and start == 0 and end == dim_size and step == 1 @@ -285,6 +322,14 @@ def slice_scatter( raise NotImplementedError( "slice_scatter with dynamic start/end/step is not yet supported" ) + + if status is KVWriteStatus.DYNAMIC_DIM_SIZE: + # The validator keeps these out of TensorRT, so this is only reachable for a + # node it could not read a shape from. + raise NotImplementedError( + f"slice_scatter: dim {dim} of the input is dynamic, so this write's bounds " + "cannot be resolved without its size" + ) # OK is the only status left, and it resolves all three bounds to Python ints. assert start is not None and end is not None and step is not None diff --git a/py/torch_tensorrt/dynamo/lowering/_buffer_lifting.py b/py/torch_tensorrt/dynamo/lowering/_buffer_lifting.py index c66b0674c2..90b7af47b3 100644 --- a/py/torch_tensorrt/dynamo/lowering/_buffer_lifting.py +++ b/py/torch_tensorrt/dynamo/lowering/_buffer_lifting.py @@ -150,14 +150,16 @@ def _kv_write_will_alias( direct network input, so this reuses the converters' own eligibility predicates -- and, for ``slice_scatter``, the converter's own derivation of the arguments those predicates take (:func:`resolve_slice_scatter_write`), since - a divergence there mis-predicts just as effectively as a divergence in the + a divergence there mispredicts just as effectively as a divergence in the predicate. Returning ``False`` routes the write to copy-back, which is right whatever the converter goes on to do with it: it lowers most ineligible writes to a non-aliasing scatter, returns the source outright for a full overwrite, and raises for the rest (a ``slice_scatter`` with dynamic bounds or a bad ``dim``, an ``index_copy`` its fallback cannot express). The first two both have a write-back to preserve -- for a full overwrite the source *is* the buffer's new contents -- - and the ones that raise never get as far as needing one. + and the ones that raise never get as far as needing one. One more never reaches a + converter: a ``slice_scatter`` whose bounds need the size of a dynamic dim, which + its validator runs in PyTorch, and which needs the copy-back like any other. Imports are local to avoid a lowering<->conversion import cycle. """ if not (isinstance(value_node, torch.fx.Node) and value_node.op == "call_function"): @@ -201,8 +203,9 @@ def _kv_write_will_alias( # Returning False routes the write to copy-back -- the copy_ is erased # either way, and the new value is re-attached as a graph output instead. A # full overwrite needs that: it returns the source, emits no KV layer, and - # its write still has to be recorded. The other two statuses raise out of the - # converter, so the compile aborts and the output is never reached. + # its write still has to be recorded. Two of the rest raise out of the + # converter, so the compile aborts and the output is never reached; the + # dynamic-dim one is validated away to PyTorch and keeps its copy-back. dim = args[2] if len(args) > 2 else 0 start, end, _step, status = resolve_slice_scatter_write( tuple(cache_shape), @@ -484,7 +487,7 @@ def lift_mutated_buffers( copyback: List[Tuple[torch.fx.Node, str, str]] = [] # Input-binding names of writes predicted to alias (KV). compile() asserts each # actually appears in the engine's aliased_io, and that no copy-back binding - # does, turning either mis-prediction into a loud error instead of a silently + # does, turning either misprediction into a loud error instead of a silently # dropped write-back or a module that raises on every call. The binding name # (buf_*) is the stable key: it survives the buffer renaming inline does later, # and it is exactly what aliased_io records on the input side. diff --git a/tests/py/dynamo/conversion/test_slice_scatter_aten.py b/tests/py/dynamo/conversion/test_slice_scatter_aten.py index ed4bacf050..0a38693f42 100644 --- a/tests/py/dynamo/conversion/test_slice_scatter_aten.py +++ b/tests/py/dynamo/conversion/test_slice_scatter_aten.py @@ -23,15 +23,20 @@ import numpy as np import torch +import torch_tensorrt from parameterized import parameterized -from torch.testing._internal.common_utils import run_tests +from torch.testing._internal.common_utils import TestCase, run_tests from torch_tensorrt import Input +from torch_tensorrt.dynamo.conversion.aten_ops_converters import slice_scatter_validator from torch_tensorrt.dynamo.conversion.impl.slice_scatter import ( slice_scatter as slice_scatter_impl, ) from .harness import DispatchTestCase +# What torch.export writes in place of the dim size for an open-ended slice. +OPEN_END = 2**63 - 1 + class _SliceScatterNotInputModule(torch.nn.Module): """Helper: forces the fallback path by making `cache` not a direct @@ -102,6 +107,43 @@ def test_fallback_dynamic_shape(self): ] self.run_test_with_dynamic_shape(module, input_specs) + def test_fallback_open_end_step_two(self): + """``cache[:, :, ::2, :]``: the open end has to be clamped to the dim, since the + index range is built with ``np.arange``.""" + module = _SliceScatterNotInputModule(2, 0, OPEN_END, step=2) + cache = torch.randn(2, 4, 16, 8) + update = torch.randn(2, 4, 8, 8) + self.run_test(module, [cache, update]) + + def test_fallback_open_end_interior_start(self): + """``cache[:, :, 3:, :]``: a start that rules out the full-overwrite shortcut, + so the clamp is what makes the write 13 slots wide and not INT64_MAX - 3.""" + module = _SliceScatterNotInputModule(2, 3, OPEN_END, step=1) + cache = torch.randn(2, 4, 16, 8) + update = torch.randn(2, 4, 13, 8) + self.run_test(module, [cache, update]) + + def test_fallback_dynamic_sliced_dim(self): + """The dim being *written* varies here; ``test_fallback_dynamic_shape`` leaves + it fixed at 64. Bounds that stand on their own index a dynamic dim as they are; + ones that need its size are validated away (``TestSliceScatterValidator``).""" + module = _SliceScatterNotInputModule(2, 1, 5, step=1) + input_specs = [ + Input( + min_shape=(2, 4, 20, 8), + opt_shape=(2, 4, 32, 8), + max_shape=(2, 4, 64, 8), + dtype=torch.float32, + ), + Input( + min_shape=(2, 4, 4, 8), + opt_shape=(2, 4, 4, 8), + max_shape=(2, 4, 4, 8), + dtype=torch.float32, + ), + ] + self.run_test_with_dynamic_shape(module, input_specs) + def test_full_overwrite_is_identity(self): """When start=0, end=dim_size, step=1, the converter short-circuits and returns ``src`` directly. Wrap the returned tensor in a small op @@ -121,12 +163,12 @@ def forward(self, cache_in, update): class TestSliceScatterEarlyExits(unittest.TestCase): - """The converter's two raising exits, driven through the converter itself. + """The converter's three raising exits, driven through the converter itself. - Both are reached before the converter touches anything but ``input.shape``, so + All are reached before the converter touches anything but ``input.shape``, so ``_call`` passes ``None`` for ``ctx``, ``target``, ``source_ir`` and ``src``, and an object carrying only a shape for the cache. Those five stand-ins are what - breaks if either exit is ever moved below a line that reads one of them. + breaks if any exit is ever moved below a line that reads one of them. ``run_test`` reaches neither exit, for a different reason per test. A bound that is not a Python int has no concrete ``aten.slice_scatter`` to be traced into. An @@ -139,8 +181,8 @@ class TestSliceScatterEarlyExits(unittest.TestCase): _CACHE_SHAPE = (2, 4, 16, 8) - def _call(self, dim, start, end, step): - cache = SimpleNamespace(shape=self._CACHE_SHAPE) + def _call(self, dim, start, end, step, cache_shape=None): + cache = SimpleNamespace(shape=cache_shape or self._CACHE_SHAPE) return slice_scatter_impl( None, None, None, "test_slice_scatter", cache, None, dim, start, end, step ) @@ -179,6 +221,135 @@ def test_non_int_dim_is_an_index_error(self): ): self._call(np.int64(2), 0, 4, 1) + def test_open_end_on_a_dynamic_dim_is_not_implemented(self): + """The -1 TensorRT reports for a dynamic dim is no size to clamp an open end + against. ``slice_scatter_validator`` keeps these out of the engine, so this exit + is the backstop for a node it could not read a shape from. Matching on the dim + is what separates it from the dynamic-bounds exit above, which raises the same + type for a slice whose bounds are symbolic on a perfectly static dim.""" + with self.assertRaises(NotImplementedError) as raised: + self._call(2, 3, OPEN_END, 1, cache_shape=(2, 4, -1, 8)) + self.assertIn("dim 2 of the input is dynamic", str(raised.exception)) + + +class TestSliceScatterValidator(TestCase): + """What the validator keeps out of TensorRT, checked directly rather than through a + compile that happens to succeed. Nodes are built by hand so the metadata under test + is chosen here, the dynamic ones taking their ``SymInt`` dim from a real export.""" + + # Bounds that need the size of the dim being written. + _BOUNDS_NEEDING_THE_DIM = ( + (3, OPEN_END), + (3, None), + (None, None), + (-4, None), + (-4, 12), + ) + + @staticmethod + def _static_node(*slice_args): + graph = torch.fx.Graph() + cache = graph.placeholder("cache") + cache.meta["val"] = torch.empty((2, 4, 16, 8), device="meta") + src = graph.placeholder("src") + src.meta["val"] = torch.empty((2, 4, 13, 8), device="meta") + return graph.call_function( + torch.ops.aten.slice_scatter.default, args=(cache, src, *slice_args) + ) + + @staticmethod + def _dynamic_seq_node(*slice_args): + """A ``slice_scatter`` spliced into an exported graph, so its cache placeholder + carries the ``SymInt`` dim export gives it.""" + + class Passthrough(torch.nn.Module): + def forward(self, cache, update): + return cache + 0 + + seq = torch.export.Dim("seq", min=8, max=32) + ep = torch.export.export( + Passthrough(), + (torch.randn(2, 4, 16, 8), torch.randn(2, 4, 13, 8)), + dynamic_shapes={"cache": {2: seq}, "update": None}, + ) + gm = ep.module() + cache, update = [n for n in gm.graph.nodes if n.op == "placeholder"][:2] + output = next(n for n in gm.graph.nodes if n.op == "output") + with gm.graph.inserting_before(output): + return gm.graph.call_function( + torch.ops.aten.slice_scatter.default, args=(cache, update, *slice_args) + ) + + def test_bounds_relative_to_a_dynamic_dim_are_rejected(self): + """Left in the engine, the open end reaches ``np.arange`` as a request for + INT64_MAX entries.""" + for start, end in self._BOUNDS_NEEDING_THE_DIM: + self.assertFalse( + slice_scatter_validator(self._dynamic_seq_node(2, start, end)) + ) + + def test_self_contained_bounds_on_a_dynamic_dim_are_accepted(self): + """A non-negative concrete bound means the same thing whatever the dim turns out + to be, so the converter indexes with it as given.""" + self.assertTrue(slice_scatter_validator(self._dynamic_seq_node(2, 1, 5))) + + def test_a_static_dim_resolves_every_bound(self): + """The same bounds on a static dim all resolve -- clamped, or counted from the + end -- so none of them is the validator's business.""" + for start, end in self._BOUNDS_NEEDING_THE_DIM: + self.assertTrue(slice_scatter_validator(self._static_node(2, start, end))) + + def test_a_node_without_shape_metadata_is_passed(self): + """Rejecting is the more damaging guess: the KV-cache classifier reads the same + metadata, and vetoing a write it classified as engine-aliased fails + ``assert_predicted_kv_aliased``. The converter resolves against the TensorRT + shape instead, and raises if that turns out to be dynamic.""" + graph = torch.fx.Graph() + cache = graph.placeholder("cache") + src = graph.placeholder("src") + node = graph.call_function( + torch.ops.aten.slice_scatter.default, args=(cache, src, 2, 3, OPEN_END) + ) + self.assertEqual(cache.meta, {}) + self.assertTrue(slice_scatter_validator(node)) + + +@unittest.skipIf(not torch.cuda.is_available(), "Skip because CUDA is not available") +class TestSliceScatterDynamicDimEndToEnd(TestCase): + """``cache[:, :, 3:, :] = update`` on a dynamic dim is the write with no lowering, + and the point of the validator is that the model compiles anyway. The ``+ 1`` gives + the engine something to take, so the write has to be partitioned out to PyTorch + rather than the whole graph falling back.""" + + def test_open_end_on_a_dynamic_dim_matches_eager(self): + class Write(torch.nn.Module): + def forward(self, cache, update): + # The op export emits for `cache[:, :, 3:, :] = update`, written out so + # the open end reaches the validator as the sentinel rather than as a + # symbolic bound the decomposition pass would rewrite first. + written = torch.ops.aten.slice_scatter.default( + cache, update, 2, 3, OPEN_END + ) + return written + 1.0 + + mod = Write().eval().cuda() + seq = torch.export.Dim("seq", min=8, max=32) + cache = torch.randn(2, 4, 16, 8).cuda() + update = torch.randn(2, 4, 13, 8).cuda() + ep = torch.export.export( + mod, + (cache, update), + dynamic_shapes={"cache": {2: seq}, "update": {2: seq - 3}}, + ) + trt_mod = torch_tensorrt.dynamo.compile(ep, [cache, update], min_block_size=1) + torch.testing.assert_close(trt_mod(cache, update), mod(cache, update)) + + longer_cache = torch.randn(2, 4, 24, 8).cuda() + longer_update = torch.randn(2, 4, 21, 8).cuda() + torch.testing.assert_close( + trt_mod(longer_cache, longer_update), mod(longer_cache, longer_update) + ) + if __name__ == "__main__": run_tests() diff --git a/tests/py/dynamo/lowering/test_buffer_lifting.py b/tests/py/dynamo/lowering/test_buffer_lifting.py index 71ebe05f09..888bae0feb 100644 --- a/tests/py/dynamo/lowering/test_buffer_lifting.py +++ b/tests/py/dynamo/lowering/test_buffer_lifting.py @@ -25,6 +25,7 @@ """ import inspect +import sys import unittest from unittest import mock @@ -559,7 +560,7 @@ class TestSliceScatterDerivationIsShared(TestCase): ``_kv_write_will_alias`` reuses the converter's eligibility predicate, but a predicate is only as good as what it is handed: computing ``start`` / - ``update_len`` differently on the two sides mis-predicts just as effectively + ``update_len`` differently on the two sides mispredicts just as effectively as a different predicate would. ``resolve_slice_scatter_write`` is the single derivation both sides call, and these pin the corners where an independent derivation would part company with the converter. @@ -587,23 +588,31 @@ def test_full_overwrite_is_not_kv(self): # Same slice written implicitly (no start/end args). self.assertFalse(self._classify((2, 4, 16, 8), (2, 4, 16, 8), 2)) - def test_open_ended_slice_is_not_kv(self): - """``cache[:, :, 3:, :]`` lowers with ``end == INT64_MAX``. The converter - takes its write length from ``end - start``, which fails the ``start + - update_len <= s_max`` bound, so the predictor has to read the length the - same way rather than from the source's extent along the dim.""" - self.assertFalse( - self._classify((2, 4, 16, 8), (2, 4, 13, 8), 2, 3, 9223372036854775807) - ) + def test_open_ended_slice_is_clamped_to_the_dim(self): + """``cache[:, :, 3:, :]`` lowers with ``end == INT64_MAX``, which the shared + derivation clamps to the dim as aten does, leaving the 13 slots from 3 to s_max + -- a KV-eligible write. Both sides have to clamp: the predictor takes the write + length from ``end - start``, so an unclamped end fails the ``start + update_len + <= s_max`` bound and files copy-back for a write the converter goes on to alias, + which is what ``assert_predicted_kv_aliased`` raises on.""" + for open_end in (sys.maxsize, 9223372036854775807): + self.assertTrue( + self._classify((2, 4, 16, 8), (2, 4, 13, 8), 2, 3, open_end) + ) def test_negative_start_is_normalised(self): - """A negative ``start`` counts from the end for the converter, so the - predictor must normalise it before applying the ``start + update_len <= - s_max`` bound rather than passing the raw negative through.""" - # -4 normalises to 12; 12 + 4 == s_max, so this is eligible. + """A negative ``start`` counts from the end for the converter, so the predictor + has to normalise it the same way -- ``resolve_slice_scatter_write`` is what pins + the resulting 12 -- before applying the ``start + update_len <= s_max`` bound. + + Both of these now pass that bound, and no slice can fail it: ``end`` is clamped + to the dim, so ``start + update_len == end <= s_max`` holds by construction. The + bound still guards ``index_copy``, whose write position comes from an index + tensor rather than a slice.""" + # -4 normalises to 12; the write is the 4 slots from 12 to s_max. self.assertTrue(self._classify((2, 4, 16, 8), (2, 4, 4, 8), 2, -4, 16)) - # -4 normalises to 12 but the write runs past s_max, so it is not. - self.assertFalse(self._classify((2, 4, 16, 8), (2, 4, 8, 8), 2, -4, 20)) + # An end past the dim is clamped back to it, leaving the same 4-slot write. + self.assertTrue(self._classify((2, 4, 16, 8), (2, 4, 4, 8), 2, -4, 20)) def test_non_int_start_is_not_kv(self): """A non-constant bound makes the converter raise rather than emit a KV @@ -654,6 +663,11 @@ def test_resolve_reports_the_converter_early_exits(self): resolve_slice_scatter_write(shape, 2, -4, None, None), (12, 16, 1, KVWriteStatus.OK), ) + # An end past the dim is clamped to it, as aten does. + self.assertEqual( + resolve_slice_scatter_write(shape, 2, -4, 20, 1), + (12, 16, 1, KVWriteStatus.OK), + ) self.assertEqual( resolve_slice_scatter_write(shape, 2, None, None, None), (0, 16, 1, KVWriteStatus.FULL_OVERWRITE), @@ -680,6 +694,36 @@ def __eq__(self, other): (None, None, None, KVWriteStatus.BAD_DIM), ) + def test_a_dynamic_dim_leaves_relative_bounds_unresolved(self): + """A bound needing the size of a dynamic dim is reported rather than resolved + against a stand-in: reading TensorRT's -1 as a size is how ``cache[:, :, 3:]`` + used to resolve to ``arange(3, -2)``, an empty write. Both shapes that dim + arrives in are checked -- -1 from TensorRT, a non-int for the fx graph's + ``SymInt`` -- since the two callers have to read them the same way.""" + from torch_tensorrt.dynamo.conversion.impl.slice_scatter import ( + KVWriteStatus, + resolve_slice_scatter_write, + ) + + unresolved = (None, None, None, KVWriteStatus.DYNAMIC_DIM_SIZE) + for dynamic_size in (-1, "sym"): + shape = (2, 4, dynamic_size, 8) + for start, end in ( + (3, sys.maxsize), + (3, 9223372036854775807), + (3, None), + (None, None), + (-4, None), + (-4, 12), + ): + self.assertEqual( + resolve_slice_scatter_write(shape, 2, start, end, 1), unresolved + ) + self.assertEqual( + resolve_slice_scatter_write(shape, 2, 1, 5, 1), + (1, 5, 1, KVWriteStatus.OK), + ) + class TestCacheMustReachTheConverterAsANetworkInput(TestCase): """``emit_kv_cache_update_layer`` aliases the cache only when it is handed a @@ -1274,7 +1318,7 @@ def forward(self, x): def test_compile_runs_the_predicted_kv_cross_check(self): """``compile()`` must actually call ``assert_predicted_kv_aliased`` with the - predictions lift made; without that call a mis-classified write silently + predictions lift made; without that call a misclassified write silently loses its write-back.""" from torch_tensorrt.dynamo import _compiler as C