summary
AOTAutograd's synthetic-base handling (merge_view_inputs) supports functionalizing mutations on inputs that alias each other, but it bails out when the aliased views have different dtypes:
# torch/_functorch/_aot_autograd/runtime_wrappers.py:2299
# Regenerating views when reinterpreting complex / real tensors seems non-trivial,
# not handling for now
if not _same_dtype_views(view1, view2):
raise AssertionError(
"aot_autograd() does not yet handle input mutations on views with different dtypes."
)
The comment suggests this was banned for the complex/real reinterpretation case, but it also blocks the general mixed-dtype case, which appears to be legitimate and implementable.
Motivating use case
Tracing vLLM on a non-CUDA backend :
The vLLM engine allocates cache buffers as aliases into a single base tensor — GDN conv + SSM states and the regular softmax-attention KV cache all share one storage.
These tensors are mutated in-place in the forward pass to update cache state.
The aliased views have different dtypes, so merge_view_inputs raises the assertion above and tracing fails.
The desired behavior is to represent this in the traced graph via functionalization → defunctionalization, rather than working around it.
Repro sketch
base = torch.empty(1024, dtype=torch.uint8, device=dev)
a = base[:512].view(torch.float16) # cache A
b = base[512:].view(torch.bfloat16) # cache B, different dtype, same storage
def f(a, b, x):
a.copy_(x.to(a.dtype)) # in-place cache write
b.mul_(2)
return a.sum() + b.sum()
torch.compile(f)(a, b, x) # AssertionError: ... input mutations on views with different dtypes
(needs to be turned into a real minimal repro / test case)
Proposal
Extend merge_view_inputs / synthetic-base construction to handle aliased mutated inputs with differing dtypes. The core piece is regenerating each view off the synthetic base when a view(dtype) reinterpretation is involved (sizes/strides/offsets are in element units of the view dtype, so they need to be translated through the base's element size). Suggested scope:
Handle the mixed-dtype case for is_inference=True first (fewer restrictions than the autograd path).
Keep the existing bail-out only for cases that are genuinely unsound (e.g. complex ↔ real where regeneration is ill-defined), with a clearer error message.
Add tests to test/functorch/test_aotdispatch.py covering: mixed-dtype aliased inputs with data mutation, disjoint vs. overlapping views, and the false-aliasing path (compute_overlapping_inputs).
Alternatives considered
Tag mutating ops as "does not affect aliases." For KV-cache writes the mutated region is disjoint from other views, so the model can tag those ops and make them look functional, bypassing the alias-mutation logic entirely, then defunctionalize afterwards. This is the workaround currently being used for vLLM.
Downside: it is unsound in general. Viewing a mutated region of an alias through another op is a legitimate pattern; today's GenAI models mostly don't do it, but any model could, depending on how the cache is viewed. "Pretending these aren't views" would then silently hide the required aliasing resolution and produce wrong outputs.
Given that, proper support in merge_view_inputs seems preferable to a global "pretend it's not a view" escape hatch.
cc @chauhang @penguinwu @bdhirsh @bobrenjc93 @aorenste
summary
AOTAutograd's synthetic-base handling (merge_view_inputs) supports functionalizing mutations on inputs that alias each other, but it bails out when the aliased views have different dtypes:
The comment suggests this was banned for the complex/real reinterpretation case, but it also blocks the general mixed-dtype case, which appears to be legitimate and implementable.
Motivating use case
Tracing vLLM on a non-CUDA backend :
The vLLM engine allocates cache buffers as aliases into a single base tensor — GDN conv + SSM states and the regular softmax-attention KV cache all share one storage.
These tensors are mutated in-place in the forward pass to update cache state.
The aliased views have different dtypes, so merge_view_inputs raises the assertion above and tracing fails.
The desired behavior is to represent this in the traced graph via functionalization → defunctionalization, rather than working around it.
Repro sketch
Proposal
Extend merge_view_inputs / synthetic-base construction to handle aliased mutated inputs with differing dtypes. The core piece is regenerating each view off the synthetic base when a view(dtype) reinterpretation is involved (sizes/strides/offsets are in element units of the view dtype, so they need to be translated through the base's element size). Suggested scope:
Handle the mixed-dtype case for is_inference=True first (fewer restrictions than the autograd path).
Keep the existing bail-out only for cases that are genuinely unsound (e.g. complex ↔ real where regeneration is ill-defined), with a clearer error message.
Add tests to test/functorch/test_aotdispatch.py covering: mixed-dtype aliased inputs with data mutation, disjoint vs. overlapping views, and the false-aliasing path (compute_overlapping_inputs).
Alternatives considered
Tag mutating ops as "does not affect aliases." For KV-cache writes the mutated region is disjoint from other views, so the model can tag those ops and make them look functional, bypassing the alias-mutation logic entirely, then defunctionalize afterwards. This is the workaround currently being used for vLLM.
Downside: it is unsound in general. Viewing a mutated region of an alias through another op is a legitimate pattern; today's GenAI models mostly don't do it, but any model could, depending on how the cache is viewed. "Pretending these aren't views" would then silently hide the required aliasing resolution and produce wrong outputs.
Given that, proper support in merge_view_inputs seems preferable to a global "pretend it's not a view" escape hatch.
cc @chauhang @penguinwu @bdhirsh @bobrenjc93 @aorenste