[#10013][feat] AutoDeploy: native cache manager integration - #10635
Conversation
📝 WalkthroughWalkthroughThis pull request refactors the AutoDeploy KV cache management system to use native KVCacheManager with abstract cache resource handlers. It replaces page-based configuration ( Changes
Estimated code review effort🎯 4 (Complex) | ⏱️ ~65 minutes 🚥 Pre-merge checks | ✅ 2 | ❌ 3❌ Failed checks (1 warning, 2 inconclusive)
✅ Passed checks (2 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing touches
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 7
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (6)
tests/unittest/_torch/auto_deploy/_utils_test/_graph_test_helpers.py (1)
1-10: Missing NVIDIA copyright header.Per coding guidelines, all TensorRT-LLM source files (including
.py) should contain an NVIDIA copyright header with the year of latest meaningful modification.Proposed fix
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import copy from typing import Callable, Dict, List, Optionaltests/unittest/_torch/auto_deploy/unit/singlegpu/test_ad_build_small_single.py (1)
1-2: Missing NVIDIA copyright header.As per the coding guidelines, all TensorRT-LLM source files should contain an NVIDIA copyright header with the year of latest meaningful modification. This file is missing the required copyright header at the top.
Proposed fix to add copyright header
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. """Testing build_and_run_ad end2end."""tensorrt_llm/_torch/auto_deploy/transform/library/build_model.py (1)
71-91: Missingupdate_kv_cache_configcall for consistency.
BuildAndLoadFactoryModeloverrides_apply_to_full_model()but doesn't callcm.update_kv_cache_config(). Since this class explicitly useshf.AutoModelFactory(which returns dtype updates fromget_cache_config_updates()), the KV cache config updates may not be applied for HF models loaded via this path.Suggested fix
# build and load the model model = factory.build_and_load_model(cm.device) + # update the kv cache config + cm.update_kv_cache_config(**factory.get_cache_config_updates()) + # we set the standard example sequence WITHOUT extra_args to set them to None so that # only the text portion of the model gets called. cm.info.set_example_sequence()tests/unittest/_torch/auto_deploy/unit/singlegpu/shim/test_engine.py (1)
1-4: Missing NVIDIA copyright header.As per coding guidelines, all TensorRT-LLM source files should contain an NVIDIA copyright header with the year of latest meaningful modification.
Suggested fix
+# SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + from types import SimpleNamespace from typing import List, Optional, Typetests/integration/defs/accuracy/test_llm_api_autodeploy.py (2)
107-119: Duplicatekv_cache_configkey overwritesenable_block_reusesetting.The dictionary has
"kv_cache_config"defined twice (lines 107-109 and 117-119). Python dictionaries use the last occurrence, soenable_block_reuse: Falseis silently discarded. Since SSMs do not support cache reuse (as noted in the comment), this bug could cause test failures or incorrect behavior.Proposed fix: merge into single kv_cache_config
# SSMs do not support cache reuse. "kv_cache_config": { - "enable_block_reuse": False - }, - # Keep max_batch_size as in the PyTorch test to avoid OOM - "max_batch_size": 128, - # Model context length is 8K - "max_seq_len": 8192, - # Set explicitly to match default build_config behavior - "max_num_tokens": 8192, - "skip_loading_weights": False, - "kv_cache_config": { - "free_gpu_memory_fraction": 0.7 + "enable_block_reuse": False, + "free_gpu_memory_fraction": 0.7, }, + # Keep max_batch_size as in the PyTorch test to avoid OOM + "max_batch_size": 128, + # Model context length is 8K + "max_seq_len": 8192, + # Set explicitly to match default build_config behavior + "max_num_tokens": 8192, + "skip_loading_weights": False,
165-181: Duplicatekv_cache_configkey overwritesenable_block_reuseand mamba config.Same issue:
"kv_cache_config"is defined at lines 165-169 and again at lines 179-181. The second definition overwrites the first, losingenable_block_reuse: Falseand the mamba cache dtype comments. This could cause incorrect behavior for the MOE model tests.Proposed fix: merge into single kv_cache_config
# SSMs do not support cache reuse. "kv_cache_config": { - "enable_block_reuse": False + "enable_block_reuse": False, + "free_gpu_memory_fraction": 0.7, # NOTE: some accuracy benchmarks may require fp32 precision for mamba cache # "mamba_ssm_cache_dtype": "float32", }, # Keep max_batch_size as in the PyTorch test to avoid OOM "max_batch_size": 128, # Model context length is 8K "enable_chunked_prefill": True, "max_seq_len": 8192, # Set explicitly to match default build_config behavior "max_num_tokens": 8192, "skip_loading_weights": False, "compile_backend": "torch-cudagraph", - "kv_cache_config": { - "free_gpu_memory_fraction": 0.7 - }, "cuda_graph_batch_sizes": [1, 2, 4, 8, 16, 32, 64, 128],
🤖 Fix all issues with AI agents
In @docs/source/features/auto_deploy/advanced/benchmarking_with_trtllm_bench.md:
- Line 86: The docs show a mismatch: the table documents
kv_cache_config.free_gpu_memory_fraction default as 0.9 but the example uses
0.8; update the example value to 0.9 to match the documented default (or
alternatively keep 0.8 and add an inline comment in the example explaining why
0.8 was chosen) so the example and the documented default are consistent; edit
the example block that sets kv_cache_config.free_gpu_memory_fraction to reflect
this change and ensure the table and example no longer disagree.
- Around line 58-63: Remove the incorrect duplicated header "# transform
options" that precedes the kv_cache_config block; ensure the block begins with
the correct KV cache comment (e.g., the line "# KV cache configuration") and
leave the kv_cache_config: and free_gpu_memory_fraction: 0.8 entries unchanged
so the separate "# transform options" comment remains only above the actual
transforms section later.
In @docs/source/torch/auto_deploy/advanced/serving_with_trtllm_serve.md:
- Around line 57-63: Remove the duplicated and misleading comment header by
deleting or replacing the top "# transform options" line that precedes the
kv_cache_config block; ensure the kv_cache_config stanza remains and, if a
header is desired, use a correct comment like "# KV cache configuration" so only
the accurate header (for kv_cache_config) appears and the original "# transform
options" remains only once before the actual transforms section (refer to the
kv_cache_config block and the later transform options header).
In @examples/models/core/nemotron/README_nemotron_nano_v3.md:
- Around line 42-47: The README guidance is inconsistent: the YAML example
leaves mamba_ssm_cache_dtype commented while the text recommends float32; either
make the example match the recommendation by uncommenting mamba_ssm_cache_dtype:
float32 in the YAML snippet (so users see the recommended default), or change
the textual guidance that references mamba-cache-dtype to state explicitly that
float32 is optional and explain the accuracy vs. speed tradeoff and point to the
commented example; update references to the setting name mamba_ssm_cache_dtype
and ensure the nano v3 doc aligns with nano-v2-vl behavior.
In @tensorrt_llm/_torch/auto_deploy/shim/demollm.py:
- Around line 81-82: The calculation for num_extra_pages over-allocates when
extra_tokens divides tokens_per_block; replace the current expression that adds
a boolean (extra_tokens // tokens_per_block) + (extra_tokens > 0) with a proper
ceiling division and ensure non-positive extra_tokens yield zero pages: compute
num_extra_pages using ceiling division of extra_tokens by tokens_per_block
(e.g., (extra_tokens + tokens_per_block - 1) // tokens_per_block or
-(-extra_tokens // tokens_per_block)) and clamp to 0 when extra_tokens <= 0,
updating the assignment that currently sets num_extra_pages.
In @tensorrt_llm/_torch/auto_deploy/shim/interface.py:
- Around line 101-166: The byte-level abstraction relies on Mapping().tp_size ==
1 but this is not validated; update _create_kv_cache_manager to either (A)
validate and raise a clear error when Mapping().tp_size != 1 (e.g., assert or
raise ValueError referencing Mapping.tp_size and total_bytes_per_token), or (B)
compute num_kv_heads as total_bytes_per_token // mapping.tp_size (with proper
ceil/overflow handling) before assigning into kv_cache_kwargs; target the
_create_kv_cache_manager function, the kv_cache_kwargs["num_kv_heads"]
assignment, and the Mapping() instance to implement the chosen fix so the
implicit division in resource_manager.py no longer produces incorrect buffer
sizes.
In @tests/unittest/_torch/auto_deploy/unit/singlegpu/shim/test_engine.py:
- Around line 45-46: The pytest parametrization for test_engine contains a
duplicated value ([0, 2, 0]) so update the pytest.mark.parametrize call for
test_engine to remove the duplicate 0 (e.g., change to [0, 2]) so
tokens_per_block only gets unique cases; locate the pytest.mark.parametrize
annotation above the test_engine function and modify the list accordingly.
🧹 Nitpick comments (12)
tests/unittest/_torch/auto_deploy/unit/singlegpu/test_ad_build_small_single.py (1)
48-50: Test coverage gap forresize_kv_cachewith new configuration.The
resize_kv_cachetransform is now commented out, leaving no test coverage for triggering the resize behavior through the newkv_cache_config.free_gpu_memory_fractionpath. Consider adding a test case that configureskv_cache_configappropriately to validate the resize functionality still works as expected.Would you like me to help draft a test case that uses the new
kv_cache_config.free_gpu_memory_fractionapproach? Alternatively, should this TODO be tracked as a follow-up issue?tensorrt_llm/_torch/auto_deploy/models/hf.py (1)
263-272: Add return type hint and use explicit exception instead ofassert.Two issues with this method:
Missing return type hint: The base class
ModelFactory.get_cache_config_updatesspecifies-> Dict[str, Any]. This override should include the same type hint for consistency and clarity.Using
assertfor input validation:assertstatements are stripped when Python runs with the-O(optimize) flag, which would silently allow unsupported dtype values in production. Use a proper exception for validation that must always execute.♻️ Suggested improvement
- def get_cache_config_updates(self): - """Return kv cache dtype updates.""" + def get_cache_config_updates(self) -> Dict[str, Any]: + """Return updates for the KVCacheConfig for the model. + + Returns: + A dictionary with "dtype" key for kv cache dtype configuration. + """ if not self._quant_config_reader: return {} kv_cache_dtype = self._quant_config_reader.get_config().get("kv_cache_dtype", "auto") - assert kv_cache_dtype in ("fp8", "auto"), ( - f"Unsupported dtype: {kv_cache_dtype}. Only fp8 and auto are supported." - ) + if kv_cache_dtype not in ("fp8", "auto"): + raise ValueError( + f"Unsupported kv_cache_dtype: {kv_cache_dtype}. Only 'fp8' and 'auto' are supported." + ) return {"dtype": kv_cache_dtype}tensorrt_llm/_torch/auto_deploy/models/quant_config_reader.py (1)
1-7: Missing NVIDIA copyright header.As per coding guidelines, all TensorRT-LLM source files should contain an NVIDIA copyright header with the year of latest meaningful modification. Please add the appropriate copyright header at the top of the file.
📄 Suggested addition
+# SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. """ Quantization Config Reader Registry.tensorrt_llm/llmapi/llm_args.py (1)
1681-1690: Consider adding type constraints consistent withmamba_ssm_cache_dtype.The new
causal_conv_cache_dtypeanddelta_net_cache_dtypefields usestrtype, while the existingmamba_ssm_cache_dtypefield (lines 1674-1679) usesLiteral["auto", "float16", "bfloat16", "float32"]with explicit validation. If these new cache dtype fields support the same set of values, consider using the same Literal type for consistency and input validation.♻️ Suggested improvement
- # This is a pure python field, not a pybind field. It is only for the AutoDeploy backend. - causal_conv_cache_dtype: str = Field( - default="auto", - description="The data type to use for the causal conv cache.") - - # This is a pure python field, not a pybind field. It is only for the AutoDeploy backend. - delta_net_cache_dtype: str = Field( - default="auto", - description="The data type to use for the delta net cache.") + # This is a pure python field, not a pybind field. It is only for the AutoDeploy backend. + causal_conv_cache_dtype: Literal[ + "auto", "float16", "bfloat16", "float32"] = Field( + default="auto", + description="The data type to use for the causal conv cache.") + + # This is a pure python field, not a pybind field. It is only for the AutoDeploy backend. + delta_net_cache_dtype: Literal[ + "auto", "float16", "bfloat16", "float32"] = Field( + default="auto", + description="The data type to use for the delta net cache.")tensorrt_llm/_torch/auto_deploy/shim/demollm.py (1)
80-80: Consider addingstrict=Truetozip()for safety.Adding
strict=Trueensurestotal_lensandpage_assignmentshave matching lengths, failing fast if there's a mismatch rather than silently truncating.- for t_l, pages in zip(total_lens, page_assignments): + for t_l, pages in zip(total_lens, page_assignments, strict=True):tests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/test_flashinfer_attention_op.py (1)
1-6: Missing copyright header.This test file is missing the NVIDIA copyright header that should be present in all TensorRT-LLM source files. As per coding guidelines, all source files should contain an NVIDIA copyright header with the year of the latest meaningful modification.
Suggested header to add at the top of the file
+# SPDX-FileCopyrightText: Copyright (c) 2022-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import flashinfer import pytesttensorrt_llm/_torch/auto_deploy/custom_ops/triton_attention.py (1)
1-4: Missing copyright header.This source file is missing the NVIDIA copyright header. As per coding guidelines, all TensorRT-LLM source files should contain an NVIDIA copyright header with the year of the latest meaningful modification.
Suggested header to add at the top of the file
+# SPDX-FileCopyrightText: Copyright (c) 2022-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + """Custom ops for MHA/XQA attention."""tensorrt_llm/_torch/auto_deploy/transform/library/hidden_states.py (1)
230-236: Consider prefixing unused parameter with underscore.The
cache_configparameter is intentionally unused since the hidden states dtype is derived from the source node's metadata. This is correct behavior, but the static analysis hint (ARG003) could be silenced by prefixing the parameter with an underscore to indicate it's intentionally unused while maintaining interface consistency.Optional fix
@classmethod def get_cache_initializers( - cls, source_attn_node: Node, cache_config: KvCacheConfig + cls, source_attn_node: Node, _cache_config: KvCacheConfig ) -> ResourceHandlerDict:tensorrt_llm/_torch/auto_deploy/custom_ops/fla/fla_backend_delta.py (1)
1-6: Missing copyright header.This source file has a docstring but is missing the required NVIDIA copyright header. As per coding guidelines, all TensorRT-LLM source files should contain an NVIDIA copyright header.
Suggested header to add at the top of the file
+# SPDX-FileCopyrightText: Copyright (c) 2022-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + """Cached attention op for delta rule using the fla kernel library.tensorrt_llm/_torch/auto_deploy/custom_ops/mla.py (1)
23-36: Hardcodedbfloat16dtype may cause precision issues for float16 models.The
_precompute_inv_freqfunction hardcodestorch.bfloat16for the cos/sin stacked tensor:cos_sin_stacked = torch.stack([emb.cos().to(torch.bfloat16), emb.sin().to(torch.bfloat16)])For models using
float16, this could introduce unnecessary precision mismatches or type casting overhead. Consider accepting a dtype parameter or inferring from the input tensors:💡 Suggested improvement
def _precompute_inv_freq( - max_seq_len: int, head_dim: int, rope_theta: float, device: torch.device + max_seq_len: int, head_dim: int, rope_theta: float, device: torch.device, dtype: torch.dtype = torch.bfloat16 ) -> torch.Tensor: inv_freq = 1.0 / ( rope_theta ** (torch.arange(0, head_dim, 2, device=device).float() / head_dim) ) t = torch.arange(max_seq_len, device=inv_freq.device, dtype=inv_freq.dtype) freqs = torch.outer(t, inv_freq.to(t.device)) # Different from paper, but it uses a different permutation in order to obtain the same calculation emb = torch.cat((freqs, freqs), dim=-1) - cos_sin_stacked = torch.stack([emb.cos().to(torch.bfloat16), emb.sin().to(torch.bfloat16)]) + cos_sin_stacked = torch.stack([emb.cos().to(dtype), emb.sin().to(dtype)]) return cos_sin_stackedtensorrt_llm/_torch/auto_deploy/shim/interface.py (1)
89-95: Consider using a more specific exception type.The method correctly validates against
model_fields, but per the static analysis hint (TRY003), consider defining a custom exception or usingKeyError:💡 Optional improvement
def update_kv_cache_config(self, **kwargs) -> None: """Update the KVCacheConfig with the given kwargs.""" for k, v in kwargs.items(): if k in type(self.kv_cache_config).model_fields: setattr(self.kv_cache_config, k, v) else: - raise ValueError(f"Invalid KVCacheConfig field: {k}") + raise KeyError(f"Invalid KVCacheConfig field: {k!r}")This is a minor suggestion; the current implementation is functionally correct.
tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py (1)
273-274: TODO acknowledged - dummy request handling for MambaHybridCacheManager.The TODO notes that
_generate_dummy_requestmay need updates forMambaHybridCacheManagerand thatmax_state_slotsshouldn't be used anymore. Consider creating a tracking issue if this needs to be addressed before the feature is fully stable.Do you want me to open an issue to track this TODO for proper MambaHybridCacheManager support in dummy request generation?
eb50af0 to
0d9f9c9
Compare
0d9f9c9 to
f45b082
Compare
|
/bot run --add-multi-gpu-test --disable-fail-fast |
|
PR_Github #32059 [ run ] triggered by Bot. Commit: |
|
PR_Github #32059 [ run ] completed with state
|
f45b082 to
f9eeead
Compare
|
/bot run --add-multi-gpu-test --disable-fail-fast |
|
PR_Github #32158 [ run ] triggered by Bot. Commit: |
|
PR_Github #32158 [ run ] completed with state
|
nvchenghaoz
left a comment
nvchenghaoz
left a comment
There was a problem hiding this comment.
review all files except interface.py and the kvcache.py. will get back to this tomorrow. Consider split this into small PRs next time for such big changes...
|
/bot run --extra-stage "DGX_B200-4_GPUs-AutoDeploy-1, DGX_H100-4_GPUs-AutoDeploy-1" --disable-fail-fast |
|
PR_Github #33235 [ run ] triggered by Bot. Commit: |
|
PR_Github #33235 [ run ] completed with state |
|
/bot run --extra-stage "DGX_B200-4_GPUs-AutoDeploy-1, DGX_H100-4_GPUs-AutoDeploy-1" --disable-fail-fast |
|
PR_Github #33301 [ run ] triggered by Bot. Commit: |
|
The mem tracker is neat! |
|
PR_Github #33301 [ run ] completed with state
|
147dc88 to
6fe5b23
Compare
|
/bot run --extra-stage "DGX_B200-4_GPUs-AutoDeploy-1, DGX_H100-4_GPUs-AutoDeploy-1" --disable-fail-fast |
|
PR_Github #33416 [ run ] triggered by Bot. Commit: |
|
PR_Github #33416 [ run ] completed with state
|
Signed-off-by: Lucas Liebenwein <11156568+lucaslie@users.noreply.github.com>
6fe5b23 to
f6ec4d8
Compare
|
/bot run --extra-stage "DGX_B200-4_GPUs-AutoDeploy-1, DGX_H100-4_GPUs-AutoDeploy-1" --disable-fail-fast |
|
PR_Github #33638 [ run ] triggered by Bot. Commit: |
nv-guomingz
left a comment
nv-guomingz
left a comment
There was a problem hiding this comment.
LGTM on doc changes.
|
PR_Github #33638 [ run ] completed with state |
Performance overview before and after. No visible regression
|
8 participants
Description
fixes #10013
New Features
Backwards breaking changes (ACTION REQUIRED)
Deprecation of arguments
transforms.resize_kv_cache.free_mem_ratio-->kv_cache_config.free_gpu_memory_fractionattn_page_size-->kv_cache_config.tokens_per_blocktransforms.insert_cached_ssm_attention.cache_config.mamba_dtype-->kv_cache_config.mamba_ssm_cache_dtypeNew features and behavior
expect_mem_change: False. If you expect a transform to change mem allocation, make sure to setexpect_mem_change: Truebuffersin the AttentionDescriptor has been deprecated without replacement as it is not needed anymore.