kanshan/TensorRT-LLMs
1
0
Fork 0
mirror of https://github.com/NVIDIA/TensorRT-LLM.git synced 2026-02-05 02:31:33 +08:00
Code Issues Actions 1 Packages Projects Releases Wiki Activity

[None][fix] AutoDeploy: remove mem check for a log unit test (#11120)

Browse Source 
Signed-off-by: Lucas Liebenwein <11156568+lucaslie@users.noreply.github.com>
This commit is contained in:
Lucas Liebenwein 2026-01-29 15:41:51 -05:00 committed by GitHub
parent 4345636b04
commit a4880ffdbb
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 0 additions and 37 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

37
tests/unittest/_torch/auto_deploy/unit/singlegpu/transformations/library/test_fuse_mamba_a_log.py
View File

@ -13,7 +13,6 @@
# See the License for the specific language governing permissions and
# limitations under the License.
import pytest
import torch
import torch.nn as nn
from _graph_test_helpers import run_test_transformed_gm
@ -89,39 +88,3 @@ def test_fuse_mamba_a_log_creates_fused_param():
node.target in {torch.exp, torch.ops.aten.exp.default}
for node in gm_transformed.graph.nodes
), "exp node should be removed after fusion."
def test_fuse_mamba_a_log_memory_usage():
torch.manual_seed(42)
torch.cuda.manual_seed(42)
torch.cuda.empty_cache()
device = "cuda"
dtype = torch.float32
num_features = 1024 * 1024
model = DummyMambaALogModule(num_features=num_features, dtype=dtype, device=device).to(
device=device, dtype=dtype
)
x = model.get_input(device=device, dtype=dtype)
gm = torch_export_to_gm(model, args=(x,), clone=True)
torch.cuda.synchronize()
torch.cuda.empty_cache()
mem_before = torch.cuda.memory_allocated()
gm_transformed = _apply_fuse_mamba_a_log(gm)
torch.cuda.synchronize()
torch.cuda.empty_cache()
mem_after = torch.cuda.memory_allocated()
diff = mem_after - mem_before
tolerance = 5 * 1024 # 5_KiB tolerance for allocator variance
assert abs(diff) <= tolerance, (
f"Unexpected memory delta after fusion. Expected no additional memory, got {diff} bytes."
)
with pytest.raises(AttributeError):
gm_transformed.get_parameter("A_log")