Skip to content

Commit

Permalink
#0: Enable Mamba end-to-end performance test
Browse files Browse the repository at this point in the history
This change also disables the persistent program cache since it triggers
a hang on certain CI machines (see #8606).
  • Loading branch information
esmalTT committed Jun 12, 2024
1 parent 0d70c31 commit c1c3f6b
Showing 1 changed file with 9 additions and 3 deletions.
12 changes: 9 additions & 3 deletions models/demos/mamba/tests/test_mamba_perf.py
Original file line number Diff line number Diff line change
Expand Up @@ -15,11 +15,16 @@

from models.perf.perf_utils import prep_perf_report
from models.perf.device_perf_utils import run_device_perf, check_device_perf, prep_device_perf_report
from models.utility_functions import profiler, enable_persistent_kernel_cache, skip_for_grayskull, skip_for_wormhole_b0
from models.utility_functions import (
profiler,
enable_persistent_kernel_cache,
disable_persistent_kernel_cache,
skip_for_grayskull,
skip_for_wormhole_b0,
)
from tt_metal.tools.profiler.process_model_log import get_samples_per_s


@skip_for_wormhole_b0("Non-deterministic hang on CI (#8606)")
@skip_for_grayskull("Requires eth connected devices to run")
@pytest.mark.models_performance_bare_metal
@pytest.mark.parametrize(
Expand Down Expand Up @@ -58,7 +63,8 @@ def test_mamba_e2e_perf(

sequences: torch.Tensor = tokenizer(prompts, return_tensors="pt", padding=True).input_ids

enable_persistent_kernel_cache()
# Required due to non-deterministic hang on CI (#8606)
disable_persistent_kernel_cache()

# prefill
prefill_iterations = sequences.shape[1] - 1
Expand Down

0 comments on commit c1c3f6b

Please sign in to comment.