diff --git a/.gitignore b/.gitignore
index fd95f7d8d..28f9b76a9 100644
--- a/.gitignore
+++ b/.gitignore
@@ -141,5 +141,4 @@ examples/data
 .pixi

 /imod/tests/mydask.png
-/imod/tests/unittest_report.xml
-/imod/tests/examples_report.xml
+/imod/tests/*_report.xml
diff --git a/docs/api/changelog.rst b/docs/api/changelog.rst
index 95985f29f..525f330ef 100644
--- a/docs/api/changelog.rst
+++ b/docs/api/changelog.rst
@@ -17,6 +17,8 @@ Fixed
   flow barrier for MODFLOW 6
 - Bug where error would be thrown when barriers in a ``HorizontalFlowBarrier``
   would be snapped to the same cell edge. These are now summed.
+- Improve performance validation upon Package initialization
+- Improve performance writing ``HorizontalFlowBarrier`` objects

 Changed
 ~~~~~~~
diff --git a/imod/mf6/simulation.py b/imod/mf6/simulation.py
index 810a42a1a..54613eb3c 100644
--- a/imod/mf6/simulation.py
+++ b/imod/mf6/simulation.py
@@ -1378,8 +1378,5 @@ def from_imod5_data(
         )
         simulation["ims"] = solution

-        # cleanup packages for validation
-        idomain = groundwaterFlowModel.domain
-        simulation.mask_all_models(idomain)
         simulation.create_time_discretization(additional_times=times)
         return simulation
diff --git a/imod/schemata.py b/imod/schemata.py
index 95b5d6a3a..fd9d8d737 100644
--- a/imod/schemata.py
+++ b/imod/schemata.py
@@ -74,7 +74,7 @@ def scalar_None(obj):
     if not isinstance(obj, (xr.DataArray, xu.UgridDataArray)):
         return False
     else:
-        return (len(obj.shape) == 0) & (~obj.notnull()).all()
+        return (len(obj.shape) == 0) and (obj.isnull()).all()

 def align_other_obj_with_coords(
diff --git a/imod/typing/grid.py b/imod/typing/grid.py
index 3381b32ae..509da6aef 100644
--- a/imod/typing/grid.py
+++ b/imod/typing/grid.py
@@ -435,10 +435,58 @@ def is_transient_data_grid(
     return False

+class GridCache:
+    """
+    Cache grids in this object for a specific function, lookup grids based on
+    unique geometry hash.
+    """
+
+    def __init__(self, func: Callable, max_cache_size=5):
+        self.max_cache_size = max_cache_size
+        self.grid_cache: dict[int, GridDataArray] = {}
+        self.func = func
+
+    def get_grid(self, grid: GridDataArray):
+        geom_hash = get_grid_geometry_hash(grid)
+        if geom_hash not in self.grid_cache.keys():
+            if len(self.grid_cache.keys()) >= self.max_cache_size:
+                self.remove_first()
+            self.grid_cache[geom_hash] = self.func(grid)
+        return self.grid_cache[geom_hash]
+
+    def remove_first(self):
+        keys = list(self.grid_cache.keys())
+        self.grid_cache.pop(keys[0])
+
+    def clear(self):
+        self.grid_cache = {}
+
+
+UGRID2D_FROM_STRUCTURED_CACHE = GridCache(xu.Ugrid2d.from_structured)
+
+
 @typedispatch
 def as_ugrid_dataarray(grid: xr.DataArray) -> xu.UgridDataArray:
-    """Enforce GridDataArray to UgridDataArray"""
-    return xu.UgridDataArray.from_structured(grid)
+    """
+    Enforce GridDataArray to UgridDataArray, calls
+    xu.UgridDataArray.from_structured, which is a costly operation. Therefore
+    cache results.
+    """
+
+    topology = UGRID2D_FROM_STRUCTURED_CACHE.get_grid(grid)
+
+    # Copied from:
+    # https://github.com/Deltares/xugrid/blob/3dee693763da1c4c0859a4f53ac38d4b99613a33/xugrid/core/wrap.py#L236
+    # Note that "da" is renamed to "grid" and "grid" to "topology"
+    dims = grid.dims[:-2]
+    coords = {k: grid.coords[k] for k in dims}
+    face_da = xr.DataArray(
+        grid.data.reshape(*grid.shape[:-2], -1),
+        coords=coords,
+        dims=[*dims, topology.face_dimension],
+        name=grid.name,
+    )
+    return xu.UgridDataArray(face_da, topology)

 @typedispatch  # type: ignore[no-redef]
diff --git a/pixi.toml b/pixi.toml
index eac3c0f4d..e7c867b60 100644
--- a/pixi.toml
+++ b/pixi.toml
@@ -19,11 +19,12 @@ install_with_deps = "python -m pip install --editable ."
 format = "ruff check --fix .; ruff format ."
 lint = "ruff check . ; ruff format --check ."
 tests = { depends_on = ["unittests", "examples"] }
-unittests = { cmd = [
+unittests = { depends_on = ["unittests_njit", "unittests_jit"] }
+unittests_njit = { cmd = [
     "NUMBA_DISABLE_JIT=1",
     "pytest",
     "-n", "auto",
-    "-m", "not example and not user_acceptance",
+    "-m", "not example and not user_acceptance and not unittest_jit",
     "--cache-clear",
     "--verbose",
     "--junitxml=unittest_report.xml",
@@ -32,6 +33,14 @@ unittests = { cmd = [
     "--cov-report=html:coverage",
     "--cov-config=.coveragerc"
 ], depends_on = ["install"], cwd = "imod/tests" }
+unittests_jit = { cmd = [
+    "pytest",
+    "-n", "auto",
+    "-m", "unittest_jit",
+    "--cache-clear",
+    "--verbose",
+    "--junitxml=unittest_jit_report.xml",
+], depends_on = ["install"], cwd = "imod/tests" }
 examples = { cmd = [
     "NUMBA_DISABLE_JIT=1",
     "pytest",
diff --git a/pyproject.toml b/pyproject.toml
index dcaaaa437..c265cabe8 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -158,6 +158,7 @@ ignore_missing_imports = true
 markers = [
     "example: marks test as example (deselect with '-m \"not example\"')",
     "user_acceptance: marks user acceptance tests (deselect with '-m \"not user_acceptance\"')",
+    "unittest_jit: marks unit tests that should be jitted (deselect with '-m \"not unittest_jit\"')"
 ]

 [tool.hatch.version]

diff --git a/imod/tests/test_mf6/test_mf6_chd.py b/imod/tests/test_mf6/test_mf6_chd.py
index ef0a04108..c3b96b982 100644
--- a/imod/tests/test_mf6/test_mf6_chd.py
+++ b/imod/tests/test_mf6/test_mf6_chd.py
@@ -238,6 +238,7 @@ def test_from_imod5_shd(imod5_dataset, tmp_path):
     chd_shd.write("chd_shd", [1], write_context)

+@pytest.mark.unittest_jit
 @pytest.mark.parametrize("remove_merged_packages", [True, False])
 @pytest.mark.usefixtures("imod5_dataset")
 def test_concatenate_chd(imod5_dataset, tmp_path, remove_merged_packages):
diff --git a/imod/tests/test_mf6/test_mf6_simulation.py b/imod/tests/test_mf6/test_mf6_simulation.py
index b60643793..b9b78fdcb 100644
--- a/imod/tests/test_mf6/test_mf6_simulation.py
+++ b/imod/tests/test_mf6/test_mf6_simulation.py
@@ -476,6 +476,7 @@ def compare_submodel_partition_info(first: PartitionInfo, second: PartitionInfo)
     )

+@pytest.mark.unittest_jit
 @pytest.mark.usefixtures("imod5_dataset")
 def test_import_from_imod5(imod5_dataset, tmp_path):
     imod5_data = imod5_dataset[0]
@@ -495,18 +496,20 @@ def test_import_from_imod5(imod5_dataset, tmp_path):
     simulation["imported_model"]["oc"] = OutputControl(
         save_head="last", save_budget="last"
     )
-
     simulation.create_time_discretization(["01-01-2003", "02-01-2003"])
-
+    # Cleanup
     # Remove HFB packages outside domain
     # TODO: Build in support for hfb packages outside domain
     for hfb_outside in ["hfb-24", "hfb-26"]:
         simulation["imported_model"].pop(hfb_outside)
-
+    # Align NoData to domain
+    idomain = simulation["imported_model"].domain
+    simulation.mask_all_models(idomain)
     # write and validate the simulation.
     simulation.write(tmp_path, binary=False, validate=True)

+@pytest.mark.unittest_jit
 @pytest.mark.usefixtures("imod5_dataset")
 def test_import_from_imod5__correct_well_type(imod5_dataset):
     # Unpack
@@ -537,6 +540,7 @@ def test_import_from_imod5__correct_well_type(imod5_dataset):
     assert isinstance(simulation["imported_model"]["wel-WELLS_L5"], LayeredWell)

+@pytest.mark.unittest_jit
 @pytest.mark.usefixtures("imod5_dataset")
 def test_import_from_imod5__nonstandard_regridding(imod5_dataset, tmp_path):
     imod5_data = imod5_dataset[0]
@@ -558,22 +562,23 @@ def test_import_from_imod5__nonstandard_regridding(imod5_dataset, tmp_path):
         times,
         regridding_option,
     )
-
     simulation["imported_model"]["oc"] = OutputControl(
         save_head="last", save_budget="last"
     )
-
     simulation.create_time_discretization(["01-01-2003", "02-01-2003"])
-
+    # Cleanup
     # Remove HFB packages outside domain
     # TODO: Build in support for hfb packages outside domain
     for hfb_outside in ["hfb-24", "hfb-26"]:
         simulation["imported_model"].pop(hfb_outside)
-
+    # Align NoData to domain
+    idomain = simulation["imported_model"].domain
+    simulation.mask_all_models(idomain)
     # write and validate the simulation.
     simulation.write(tmp_path, binary=False, validate=True)

+@pytest.mark.unittest_jit
 @pytest.mark.usefixtures("imod5_dataset")
 def test_import_from_imod5_no_storage_no_recharge(imod5_dataset, tmp_path):
     # this test imports an imod5 simulation, but it has no recharge and no storage package.
@@ -594,23 +599,22 @@ def test_import_from_imod5_no_storage_no_recharge(imod5_dataset, tmp_path):
         default_simulation_distributing_options,
         times,
     )
-
     simulation["imported_model"]["oc"] = OutputControl(
         save_head="last", save_budget="last"
     )
-
     simulation.create_time_discretization(["01-01-2003", "02-01-2003"])
-
+    # Cleanup
     # Remove HFB packages outside domain
     # TODO: Build in support for hfb packages outside domain
     for hfb_outside in ["hfb-24", "hfb-26"]:
         simulation["imported_model"].pop(hfb_outside)
-
     # check storage is present and rch is absent
     assert not simulation["imported_model"]["sto"].dataset["transient"].values[()]
     package_keys = simulation["imported_model"].keys()
     for key in package_keys:
         assert key[0:3] != "rch"
-
+    # Align NoData to domain
+    idomain = simulation["imported_model"].domain
+    simulation.mask_all_models(idomain)
     # write and validate the simulation.
     simulation.write(tmp_path, binary=False, validate=True)
diff --git a/imod/tests/test_typing/test_typing_grid.py b/imod/tests/test_typing/test_typing_grid.py
index a1129940b..ff7f398b0 100644
--- a/imod/tests/test_typing/test_typing_grid.py
+++ b/imod/tests/test_typing/test_typing_grid.py
@@ -1,7 +1,11 @@
+import numpy as np
 import xarray as xr
 import xugrid as xu

 from imod.typing.grid import (
+    UGRID2D_FROM_STRUCTURED_CACHE,
+    GridCache,
+    as_ugrid_dataarray,
     enforce_dim_order,
     is_planar_grid,
     is_spatial_grid,
@@ -145,3 +149,73 @@ def test_merge_dictionary__unstructured(basic_unstructured_dis):
     assert isinstance(uds["bottom"], xr.DataArray)
     assert uds["ibound"].dims == ("layer", "mesh2d_nFaces")
     assert uds["bottom"].dims == ("layer",)
+
+
+def test_as_ugrid_dataarray__structured(basic_dis):
+    # Arrange
+    ibound, top, bottom = basic_dis
+    top_3d = top * ibound
+    bottom_3d = bottom * ibound
+    # Clear cache
+    UGRID2D_FROM_STRUCTURED_CACHE.clear()
+    # Act
+    ibound_disv = as_ugrid_dataarray(ibound)
+    top_disv = as_ugrid_dataarray(top_3d)
+    bottom_disv = as_ugrid_dataarray(bottom_3d)
+    # Assert
+    # Test types
+    assert isinstance(ibound_disv, xu.UgridDataArray)
+    assert isinstance(top_disv, xu.UgridDataArray)
+    assert isinstance(bottom_disv, xu.UgridDataArray)
+    # Test cache proper size
+    assert len(UGRID2D_FROM_STRUCTURED_CACHE.grid_cache) == 1
+    # Test that data is different
+    assert np.all(ibound_disv != top_disv)
+    assert np.all(top_disv != bottom_disv)
+    # Test that grid is equal
+    assert np.all(ibound_disv.grid == top_disv.grid)
+    assert np.all(top_disv.grid == bottom_disv.grid)
+
+
+def test_as_ugrid_dataarray__unstructured(basic_unstructured_dis):
+    # Arrange
+    ibound, top, bottom = basic_unstructured_dis
+    top_3d = enforce_dim_order(ibound * top)
+    bottom_3d = enforce_dim_order(ibound * bottom)
+    # Clear cache
+    UGRID2D_FROM_STRUCTURED_CACHE.clear()
+    # Act
+    ibound_disv = as_ugrid_dataarray(ibound)
+    top_disv = as_ugrid_dataarray(top_3d)
+    bottom_disv = as_ugrid_dataarray(bottom_3d)
+    # Assert
+    # Test types
+    assert isinstance(ibound_disv, xu.UgridDataArray)
+    assert isinstance(top_disv, xu.UgridDataArray)
+    assert isinstance(bottom_disv, xu.UgridDataArray)
+    assert len(UGRID2D_FROM_STRUCTURED_CACHE.grid_cache) == 0
+
+
+def test_ugrid2d_cache(basic_dis):
+    # Arrange
+    ibound, _, _ = basic_dis
+    # Act
+    cache = GridCache(xu.Ugrid2d.from_structured, max_cache_size=3)
+    for i in range(5):
+        ugrid2d = cache.get_grid(ibound[:, i:, :])
+    # Assert
+    # Test types
+    assert isinstance(ugrid2d, xu.Ugrid2d)
+    # Test cache proper size
+    assert cache.max_cache_size == 3
+    assert len(cache.grid_cache) == 3
+    # Check if smallest grid in last cache list by checking if amount of faces
+    # correct
+    expected_size = ibound[0, i:, :].size
+    keys = list(cache.grid_cache.keys())
+    last_ugrid = cache.grid_cache[keys[-1]]
+    actual_size = last_ugrid.n_face
+    assert expected_size == actual_size
+    # Test clear cache
+    cache.clear()
+    assert len(cache.grid_cache) == 0

Bottlenecks writing HFBs LHM
https://github.com/Deltares/imod-python/pull/1157 considerably improves the speed with which the LHM is written: From 34 minutes to 12.5 minutes. This is still slower than expected however. Based on profiling, I've identified 3 main bottlenecks:

- 4.5 minutes: ``xu.DataArray.from_structured`` is called 70 times. This is quite a computationally intensive operation. It can probably be reworked to be called only once.
- 3 minutes are spent in ``DTypeSchema.validate``. This make as call to the ``scalar_None`` function, which is supposed to shortcut if len(da.shape)  != 0, but because I used the ``&`` operator accidentily instead of ``and``, this is not happening.
- 3 minutes: Spent in ``mask_all_models``. This is called because the automatic regridding might introduce inconsistencies in where idomain is located and where data is specified.  I think we can remove this function, because:
     1. in this case no regridding takes place and the masking is unnecessary, besides I see that ``_regrid_like(model: IModel, *args)`` also does masking?
     2. automatic masking can hide issues already present in the original iMOD5 model.
     3. Regridding might introduce other issues as well: For example, model tops and bottoms might be regridded in such a manner that the bottom exceeds top. Automatically cleaning up all these potentially introduced artifacts is excessive. Probably better is to create cleanup utilities and guide the user to these in the ValidationError messages.

So tasks:

- [x] Ensure  ``xu.DataArray.from_structured`` is only called once
- [x] Make sure ``scalar_None`` shortcuts appropriately
- [x] Remove call to mask_all_models


diff --git a/monai/networks/nets/segresnet_ds.py b/monai/networks/nets/segresnet_ds.py
index 1ac5a79ee3..098e490511 100644
--- a/monai/networks/nets/segresnet_ds.py
+++ b/monai/networks/nets/segresnet_ds.py
@@ -508,8 +508,10 @@ def forward(  # type: ignore

         outputs: list[torch.Tensor] = []
         outputs_auto: list[torch.Tensor] = []
-        x_ = x.clone()
+        x_ = x
         if with_point:
+            if with_label:
+                x_ = x.clone()
             i = 0
             for level in self.up_layers:
                 x = level["upsample"](x)
diff --git a/monai/networks/nets/vista3d.py b/monai/networks/nets/vista3d.py
index 4215a9a594..6313b7812d 100644
--- a/monai/networks/nets/vista3d.py
+++ b/monai/networks/nets/vista3d.py
@@ -639,12 +639,10 @@ def forward(self, src: torch.Tensor, class_vector: torch.Tensor):
         if self.use_mlp:
             class_embedding = self.mlp(class_embedding)
         # [b,1,feat] @ [1,feat,dim], batch dimension become class_embedding batch dimension.
-        masks = []
-        for i in range(b):
-            mask = class_embedding @ src[[i]].view(1, c, h * w * d)
-            masks.append(mask.view(-1, 1, h, w, d))
+        masks_embedding = class_embedding.squeeze() @ src.view(b, c, h * w * d)
+        masks_embedding = masks_embedding.view(b, -1, h, w, d).transpose(0, 1)

-        return torch.cat(masks, 1), class_embedding
+        return masks_embedding, class_embedding

 class TwoWayTransformer(nn.Module):

Optimize the VISTA3D latency
**Is your feature request related to a problem? Please describe.**
The VISTA3D is a very brilliant model for everything segmentation in a medical image. However it suffers a slowdown [issue](https://github.com/Project-MONAI/model-zoo/tree/dev/models/vista3d/docs#inference-gpu-benchmarks) when running on a GPU platform, which takes minutes to run a single inference. Therefore, a lower latency version of the VISTA3D model should be a user-friendly model and gain more usage.

**Describe the solution you'd like**
1. Update the VISTA3D bundle parameters to achieve a better latency.
2. Update the VISTA3D network to solve some latency bottlenecks.
3. Try to accelerate the bundle with acceleration tools like TensorRT.

**Additional context**
Some analyzation has already done in [this PR](https://github.com/Project-MONAI/model-zoo/pull/671).
The corresponding MONAI core code change is included [here](https://github.com/binliunls/MONAI/pull/1)


diff --git a/benchmarks/eko/benchmark_inverse_matching.py b/benchmarks/eko/benchmark_inverse_matching.py
index d63b776f0..1004a5481 100644
--- a/benchmarks/eko/benchmark_inverse_matching.py
+++ b/benchmarks/eko/benchmark_inverse_matching.py
@@ -90,14 +90,10 @@ def benchmark_inverse_matching():
     with pytest.raises(AssertionError):
         np.testing.assert_allclose(op1_nf3.operator, op2_nf3.operator)

-    pdf1 = apply.apply_pdf(eko_output1, toy.mkPDF("ToyLH", 0))
-    pdf2 = apply.apply_pdf(eko_output2, toy.mkPDF("ToyLH", 0))
+    pdf1, _ = apply.apply_pdf(eko_output1, toy.mkPDF("ToyLH", 0))
+    pdf2, _ = apply.apply_pdf(eko_output2, toy.mkPDF("ToyLH", 0))

     # test that different PTO matching is applied correctly
-    np.testing.assert_allclose(
-        pdf1[(MC**2, 4)]["pdfs"][C_PID], pdf2[(MC**2, 4)]["pdfs"][C_PID]
-    )
+    np.testing.assert_allclose(pdf1[(MC**2, 4)][C_PID], pdf2[(MC**2, 4)][C_PID])
     with pytest.raises(AssertionError):
-        np.testing.assert_allclose(
-            pdf1[(MC**2, 3)]["pdfs"][C_PID], pdf2[(MC**2, 3)]["pdfs"][C_PID]
-        )
+        np.testing.assert_allclose(pdf1[(MC**2, 3)][C_PID], pdf2[(MC**2, 3)][C_PID])
diff --git a/src/ekobox/apply.py b/src/ekobox/apply.py
index 8d92f9a7b..09dad2d0f 100644
--- a/src/ekobox/apply.py
+++ b/src/ekobox/apply.py
@@ -1,131 +1,181 @@
-"""Apply operator evolution to PDF set."""
+"""Apply evolution operator to a PDF."""

-from dataclasses import dataclass
-from typing import Dict, Optional, Union
+from collections.abc import Sequence
+from typing import Optional

 import numpy as np
+import numpy.typing as npt

 from eko import basis_rotation as br
 from eko import interpolation
 from eko.io import EKO
 from eko.io.types import EvolutionPoint

+RawPdfResult = dict[EvolutionPoint, npt.ArrayLike]
+"""Evolved PDFs as raw grids.
+
+The key is given by the associated evolution point. The values are
+tensors sorted by (replica, flavor, xgrid).
+"""
+
+RawErrorResult = dict[EvolutionPoint, Optional[npt.ArrayLike]]
+"""Integration errors for evolved PDFs as raw grids.
+
+The key is given by the associated evolution point. The values are
+tensors sorted by (replica, flavor, xgrid).
+"""
+
+
+LabeledPdfResult = dict[EvolutionPoint, dict[int, npt.ArrayLike]]
+"""Evolved PDFs labeled by their PDF identifier.
+
+The outer key is given by the associated evolution point. The inner key
+is the |PID|. The inner values are the values for along the xgrid.
+"""
+LabeledErrorResult = dict[EvolutionPoint, dict[int, Optional[npt.ArrayLike]]]
+"""Integration errors for evolved PDFs labeled by their PDF identifier.
+
+The outer key is given by the associated evolution point. The inner key
+is the |PID|. The inner values are the values for along the xgrid.
+"""
+

 def apply_pdf(
     eko: EKO,
     lhapdf_like,
-    targetgrid=None,
-    rotate_to_evolution_basis=False,
-):
-    """Apply all available operators to the input PDFs.
+    targetgrid: npt.ArrayLike = None,
+    rotate_to_evolution_basis: bool = False,
+) -> tuple[LabeledPdfResult, LabeledErrorResult]:
+    """Apply all available operators to the input PDF.

     Parameters
     ----------
-        output : eko.output.EKO
-            eko output object containing all informations
-        lhapdf_like : object
-            object that provides an xfxQ2 callable (as `lhapdf <https://lhapdf.hepforge.org/>`_
-            and :class:`ekomark.toyLH.toyPDF` do) (and thus is in flavor basis)
-        targetgrid : list
-            if given, interpolates to the pdfs given at targetgrid (instead of xgrid)
-        rotate_to_evolution_basis : bool
-            if True rotate to evoluton basis
+    eko :
+        eko output object containing all informations
+    lhapdf_like : Any
+        object that provides an `xfxQ2` callable (as `lhapdf <https://lhapdf.hepforge.org/>`_
+        and :class:`ekomark.toyLH.toyPDF` do) (and thus is in flavor basis)
+    targetgrid :
+        if given, interpolates to the targetgrid (instead of xgrid)
+    rotate_to_evolution_basis :
+        if True rotate to evoluton basis

     Returns
     -------
-        out_grid : dict
-            output PDFs and their associated errors for the computed mu2grid
+    pdfs :
+        PDFs for the computed evolution points
+    errors :
+        Integration errors for PDFs for the computed evolution points
     """
+    # prepare post-process
     qed = eko.theory_card.order[1] > 0
+    flavor_rotation = None
+    labels = br.flavor_basis_pids
     if rotate_to_evolution_basis:
         if not qed:
-            rotate_flavor_to_evolution = br.rotate_flavor_to_evolution
-            labels = br.evol_basis
+            flavor_rotation = br.rotate_flavor_to_evolution
+            labels = br.evol_basis_pids
         else:
-            rotate_flavor_to_evolution = br.rotate_flavor_to_unified_evolution
-            labels = br.unified_evol_basis
-        return apply_pdf_flavor(
-            eko, lhapdf_like, targetgrid, rotate_flavor_to_evolution, labels=labels
-        )
-    return apply_pdf_flavor(eko, lhapdf_like, targetgrid)
-
-
-CONTRACTION = "ajbk,bk"
-
-_PdfLabel = Union[int, str]
-"""PDF identifier: either PID or label."""
-
-
-@dataclass
-class PdfResult:
-    """Helper class to collect PDF results."""
-
-    pdfs: Dict[_PdfLabel, float]
-    errors: Optional[Dict[_PdfLabel, float]] = None
+            flavor_rotation = br.rotate_flavor_to_unified_evolution
+            labels = br.unified_evol_basis_pids
+    return apply_pdf_flavor(eko, lhapdf_like, labels, targetgrid, flavor_rotation)

 def apply_pdf_flavor(
-    eko: EKO, lhapdf_like, targetgrid=None, flavor_rotation=None, labels=None
-):
-    """Apply all available operators to the input PDFs.
+    eko: EKO,
+    lhapdf_like,
+    flavor_labels: Sequence[int],
+    targetgrid: npt.ArrayLike = None,
+    flavor_rotation: npt.ArrayLike = None,
+) -> tuple[LabeledPdfResult, LabeledErrorResult]:
+    """Apply all available operators to the input PDF.

     Parameters
     ----------
-        output : eko.output.EKO
-            eko output object containing all informations
-        lhapdf_like : object
-            object that provides an xfxQ2 callable (as `lhapdf <https://lhapdf.hepforge.org/>`_
-            and :class:`ekomark.toyLH.toyPDF` do) (and thus is in flavor basis)
-        targetgrid : list
-            if given, interpolates to the pdfs given at targetgrid (instead of xgrid)
-        flavor_rotation : np.ndarray
-            Rotation matrix in flavor space
-        labels : list
-            list of labels
+    eko :
+        eko output object containing all informations
+    lhapdf_like : Any
+        object that provides an `xfxQ2` callable (as `lhapdf <https://lhapdf.hepforge.org/>`_
+        and :class:`ekomark.toyLH.toyPDF` do) (and thus is in flavor basis)
+    flavor_labels :
+        flavor names
+    targetgrid :
+        if given, interpolates to the targetgrid (instead of xgrid)
+    flavor_rotation :
+        if give, rotate in flavor space

     Returns
     -------
-        out_grid : dict
-            output PDFs and their associated errors for the computed mu2grid
+    pdfs :
+        PDFs for the computed evolution points
+    errors :
+        Integration errors for PDFs for the computed evolution points
     """
     # create pdfs
-    pdfs = np.zeros((len(br.flavor_basis_pids), len(eko.xgrid)))
+    input_pdfs = np.zeros((len(br.flavor_basis_pids), len(eko.xgrid)))
     for j, pid in enumerate(br.flavor_basis_pids):
         if not lhapdf_like.hasFlavor(pid):
             continue
-        pdfs[j] = np.array(
+        input_pdfs[j] = np.array(
             [lhapdf_like.xfxQ2(pid, x, eko.mu20) / x for x in eko.xgrid.raw]
         )
+    # apply
+    grids, grid_errors = apply_grids(eko, input_pdfs[None, :])
+    new_grids = rotate_result(eko, grids, flavor_labels, targetgrid, flavor_rotation)
+    new_errors = rotate_result(
+        eko, grid_errors, flavor_labels, targetgrid, flavor_rotation
+    )
+    # unwrap the replica axis again
+    pdfs: LabeledPdfResult = {}
+    errors: LabeledErrorResult = {}
+    for ep, pdf in new_grids.items():
+        pdfs[ep] = {
+            lab: grid[0] if grid is not None else None for lab, grid in pdf.items()
+        }
+        errors[ep] = {
+            lab: (grid[0] if grid is not None else None)
+            for lab, grid in new_errors[ep].items()
+        }
+    return pdfs, errors

-    # build output
-    out_grid: Dict[EvolutionPoint, PdfResult] = {}
-    for ep, elem in eko.items():
-        pdf_final = np.einsum(CONTRACTION, elem.operator, pdfs, optimize="optimal")
-        if elem.error is not None:
-            error_final = np.einsum(CONTRACTION, elem.error, pdfs, optimize="optimal")
-        else:
-            error_final = None
-        out_grid[ep] = PdfResult(dict(zip(br.flavor_basis_pids, pdf_final)))
-        if error_final is not None:
-            out_grid[ep].errors = dict(zip(br.flavor_basis_pids, error_final))

-    qed = eko.theory_card.order[1] > 0
+def rotate_result(
+    eko: EKO,
+    grids: RawErrorResult,
+    flavor_labels: Sequence[int],
+    targetgrid: Optional[npt.ArrayLike] = None,
+    flavor_rotation: Optional[npt.ArrayLike] = None,
+) -> LabeledErrorResult:
+    """Rotate and relabel PDFs.
+
+    Parameters
+    ----------
+    eko :
+        eko output object containing all informations
+    grids :
+        Raw grids coming from evolution
+    flavor_labels :
+        flavors names
+    targetgrid :
+        if given, interpolates to the targetgrid (instead of xgrid)
+    flavor_rotation :
+        if given, rotates in flavor space
+
+    Returns
+    -------
+    pdfs :
+        relabeled and, if requested rotated, PDFs
+    """
     # rotate to evolution basis
     if flavor_rotation is not None:
-        for q2, op in out_grid.items():
-            pdf = flavor_rotation @ np.array(
-                [op.pdfs[pid] for pid in br.flavor_basis_pids]
+        new_grids = {}
+        for ep, pdf_grid in grids.items():
+            new_grids[ep] = (
+                np.einsum("ab,rbk->rak", flavor_rotation, pdf_grid, optimize="optimal")
+                if pdf_grid is not None
+                else None
             )
-            if not qed:
-                evol_basis = br.evol_basis
-            else:
-                evol_basis = br.unified_evol_basis
-            op.pdfs = dict(zip(evol_basis, pdf))
-            if op.errors is not None:
-                errors = flavor_rotation @ np.array(
-                    [op.errors[pid] for pid in br.flavor_basis_pids]
-                )
-                op.errors = dict(zip(evol_basis, errors))
+        grids = new_grids

     # rotate/interpolate to target grid
     if targetgrid is not None:
@@ -135,14 +185,74 @@ def apply_pdf_flavor(
             mode_N=False,
         )

-        rot = b.get_interpolation(targetgrid)
-        for evpdf in out_grid.values():
-            for pdf_label in evpdf.pdfs:
-                evpdf.pdfs[pdf_label] = np.matmul(rot, evpdf.pdfs[pdf_label])
-                if evpdf.errors is not None:
-                    evpdf.errors[pdf_label] = np.matmul(rot, evpdf.errors[pdf_label])
-    # cast back to be backward compatible
-    real_out_grid = {}
-    for ep, res in out_grid.items():
-        real_out_grid[ep] = {"pdfs": res.pdfs, "errors": res.errors}
-    return real_out_grid
+        x_rotation = b.get_interpolation(targetgrid)
+        new_grids = {}
+        for ep, pdf_grid in grids.items():
+            new_grids[ep] = (
+                np.einsum("jk,rbk->rbj", x_rotation, pdf_grid, optimize="optimal")
+                if pdf_grid is not None
+                else None
+            )
+        grids = new_grids
+
+    # relabel
+    new_grids = {}
+    for ep, pdf_grid in grids.items():
+        new_grids[ep] = dict(
+            zip(
+                flavor_labels,
+                np.swapaxes(pdf_grid, 0, 1)
+                if pdf_grid is not None
+                else [None] * len(flavor_labels),
+            )
+        )
+    grids = new_grids
+
+    return grids
+
+
+_EKO_CONTRACTION = "ajbk,rbk->raj"
+"""Contract eko for all replicas."""
+
+
+def apply_grids(
+    eko: EKO, input_grids: npt.ArrayLike
+) -> tuple[RawPdfResult, RawErrorResult]:
+    """Apply all available operators to the input grids.
+
+    Parameters
+    ----------
+    eko :
+        eko output object
+    input_grids :
+        3D PDF grids evaluated at the inital scale. The axis have to be (replica, flavor, xgrid)
+
+    Returns
+    -------
+    pdfs :
+        output PDFs for the computed evolution points
+    pdfs :
+        associated integration errors for the computed evolution points
+    """
+    # sanity check
+    if len(input_grids.shape) != 3 or input_grids.shape[1:] != (
+        len(br.flavor_basis_pids),
+        len(eko.xgrid),
+    ):
+        raise ValueError(
+            "input grids have to be sorted by replica, flavor, xgrid!"
+            f"The shape has to be (r,{len(br.flavor_basis_pids)},{len(eko.xgrid)})"
+        )
+    # iterate
+    pdfs: RawPdfResult = {}
+    errors: RawErrorResult = {}
+    for ep, elem in eko.items():
+        pdfs[ep] = np.einsum(
+            _EKO_CONTRACTION, elem.operator, input_grids, optimize="optimal"
+        )
+        errors[ep] = (
+            np.einsum(_EKO_CONTRACTION, elem.error, input_grids, optimize="optimal")
+            if elem.error is not None
+            else None
+        )
+    return pdfs, errors
diff --git a/src/ekobox/evol_pdf.py b/src/ekobox/evol_pdf.py
index 15e837736..b629dd249 100644
--- a/src/ekobox/evol_pdf.py
+++ b/src/ekobox/evol_pdf.py
@@ -78,11 +78,11 @@ def evolve_pdfs(
     with EKO.read(eko_path) as eko_output:
         for initial_PDF in initial_PDF_list:
             evolved_PDF_list.append(
-                apply.apply_pdf(eko_output, initial_PDF, targetgrid)
+                apply.apply_pdf(eko_output, initial_PDF, targetgrid)[0]
             )

         # separate by nf the evolgrid (and order per nf/q)
-        q2block_per_nf = regroup_evolgrid(eko_output.evolgrid)  # pylint: disable=E1101
+        q2block_per_nf = regroup_evolgrid(eko_output.evolgrid)

     # update info file
     if targetgrid is None:
@@ -128,7 +128,7 @@ def collect_blocks(evolved_PDF: dict, q2block_per_nf: dict, xgrid: list):
     # fake xfxQ2
     def pdf_xq2(pid, x, Q2):
         x_idx = xgrid.index(x)
-        return x * evolved_PDF[(Q2, nf)]["pdfs"][pid][x_idx]
+        return x * evolved_PDF[(Q2, nf)][pid][x_idx]

     # loop on nf patches
     for nf, q2grid in q2block_per_nf.items():
diff --git a/src/ekomark/benchmark/runner.py b/src/ekomark/benchmark/runner.py
index 4b8e94c4d..a1d79c67a 100644
--- a/src/ekomark/benchmark/runner.py
+++ b/src/ekomark/benchmark/runner.py
@@ -228,12 +228,12 @@ def log(self, theory, _, pdf, me, ext):
                 rotate_to_evolution[3, :] = [0, 0, 0, 0, 0, -1, -1, 0, 1, 1, 0, 0, 0, 0]

         with EKO.read(me) as eko:
-            pdf_grid = apply.apply_pdf_flavor(
+            pdf_grid, errors = apply.apply_pdf_flavor(
                 eko,
                 pdf,
+                labels,
                 xgrid,
                 flavor_rotation=rotate_to_evolution,
-                labels=labels,
             )
         for q2, ref_pdfs in ext["values"].items():
             log_tab = dfdict.DFdict()
@@ -243,9 +243,8 @@ def log(self, theory, _, pdf, me, ext):
                 if len(eps) == 0:
                     raise KeyError(f"PDF at Q2={q2} not found")
                 raise KeyError(f"More than one evolution point found for Q2={q2}")
-            res = pdf_grid[eps[0]]
-            my_pdfs = res["pdfs"]
-            my_pdf_errs = res["errors"]
+            my_pdfs = pdf_grid[eps[0]]
+            my_pdf_errs = errors[eps[0]]

             for key in my_pdfs:
                 if key in self.skip_pdfs(theory):

diff --git a/tests/ekobox/test_apply.py b/tests/ekobox/test_apply.py
index 70eb74adc..522c4dc6d 100644
--- a/tests/ekobox/test_apply.py
+++ b/tests/ekobox/test_apply.py
@@ -5,34 +5,31 @@
 from tests.conftest import EKOFactory

-class TestApply:
-    def test_apply(self, eko_factory: EKOFactory, fake_pdf):
-        eko = eko_factory.get()
-        ep_out = eko.evolgrid[0]
-        # fake pdfs
-        pdf_grid = apply.apply_pdf(eko, fake_pdf)
-        assert len(pdf_grid) == len(eko.evolgrid)
-        pdfs = pdf_grid[ep_out]["pdfs"]
-        assert list(pdfs.keys()) == list(br.flavor_basis_pids)
-        # rotate to target_grid
-        target_grid = [0.75]
-        pdf_grid = apply.apply_pdf(eko, fake_pdf, target_grid)
-        assert len(pdf_grid) == 1
-        pdfs = pdf_grid[ep_out]["pdfs"]
-        assert list(pdfs.keys()) == list(br.flavor_basis_pids)
+def test_apply(eko_factory: EKOFactory, fake_pdf):
+    eko = eko_factory.get()
+    ep_out = eko.evolgrid[0]
+    # base application
+    pdfs, errors = apply.apply_pdf(eko, fake_pdf)
+    assert len(pdfs) == len(eko.evolgrid)
+    assert len(pdfs) == len(errors)
+    ep_pdfs = pdfs[ep_out]
+    assert list(ep_pdfs.keys()) == list(br.flavor_basis_pids)
+    # rotate to target_grid
+    for target_grid in ([0.75], [0.5, 0.6, 0.7]):
+        pdfs, _errors = apply.apply_pdf(eko, fake_pdf, target_grid)
+        ep_pdfs = pdfs[ep_out]
+        assert list(ep_pdfs.keys()) == list(br.flavor_basis_pids)
+        assert len(ep_pdfs[21]) == len(target_grid)
+    # rotate flavor
+    pdfs, _errors = apply.apply_pdf(eko, fake_pdf, rotate_to_evolution_basis=True)
+    ep_pdfs = pdfs[ep_out]
+    assert list(ep_pdfs.keys()) == list(br.evol_basis_pids)

-    def test_apply_flavor(self, eko_factory: EKOFactory, fake_pdf, monkeypatch):
-        eko = eko_factory.get()
-        ep_out = eko.evolgrid[0]
-        # fake pdfs
-        monkeypatch.setattr(
-            "eko.basis_rotation.rotate_flavor_to_evolution", np.ones((14, 14))
-        )
-        fake_evol_basis = tuple(
-            ["a", "b"] + [str(x) for x in range(len(br.flavor_basis_pids) - 2)]
-        )
-        monkeypatch.setattr("eko.basis_rotation.evol_basis", fake_evol_basis)
-        pdf_grid = apply.apply_pdf(eko, fake_pdf, rotate_to_evolution_basis=True)
-        assert len(pdf_grid) == len(eko.evolgrid)
-        pdfs = pdf_grid[ep_out]["pdfs"]
-        assert list(pdfs.keys()) == list(fake_evol_basis)
+
+def test_apply_grids(eko_factory: EKOFactory):
+    eko = eko_factory.get()
+    # since everything is random, we can only test the tensor shapes here
+    input_grids = np.random.rand(3, len(br.flavor_basis_pids), len(eko.xgrid))
+    pdfs, errors = apply.apply_grids(eko, input_grids)
+    for ep, res in pdfs.items():
+        assert res.shape == input_grids.shape
diff --git a/tests/ekobox/test_evol_pdf.py b/tests/ekobox/test_evol_pdf.py
index 7d788277b..baf70726d 100644
--- a/tests/ekobox/test_evol_pdf.py
+++ b/tests/ekobox/test_evol_pdf.py
@@ -113,11 +113,7 @@ def test_collect_blocks():
     def mk(eps):
         f = {}
         for ep in eps:
-            f[ep] = {
-                "pdfs": {
-                    pid: np.random.rand(len(xgrid)) for pid in br.flavor_basis_pids
-                }
-            }
+            f[ep] = {pid: np.random.rand(len(xgrid)) for pid in br.flavor_basis_pids}
         return f

     # basic

Allow for parallelization in `evolven3fit_new`
Currently `evolven3fit_new` evolves the PDFs sequentially both in Q2 and replicas. For many replicas this means that evolution can take up to several hours while even for the usual 100 replica fits it generally will take over 20 minutes for the evolution to complete.

Whenever resources are available (which they are on any clusters) it would be nice to have the option to parallelize parts of the the evolution to speed up the computation.

As per the discussion w @felixhekhorn, the necessary change here is for eko to allow for more than one replica at once. To first approximation it should be trivial (the EKO x PDF contraction should have one extra axis for the number of replicas).

Then of course the changes will need to be propagated to evolven3fit and there might also be collateral damage somewhere.


diff --git a/crates/accelerate/src/filter_op_nodes.rs b/crates/accelerate/src/filter_op_nodes.rs
new file mode 100644
index 000000000000..7c41391f3788
--- /dev/null
+++ b/crates/accelerate/src/filter_op_nodes.rs
@@ -0,0 +1,63 @@
+// This code is part of Qiskit.
+//
+// (C) Copyright IBM 2024
+//
+// This code is licensed under the Apache License, Version 2.0. You may
+// obtain a copy of this license in the LICENSE.txt file in the root directory
+// of this source tree or at http://www.apache.org/licenses/LICENSE-2.0.
+//
+// Any modifications or derivative works of this code must retain this
+// copyright notice, and modified files need to carry a notice indicating
+// that they have been altered from the originals.
+
+use pyo3::prelude::*;
+use pyo3::wrap_pyfunction;
+
+use qiskit_circuit::dag_circuit::DAGCircuit;
+use qiskit_circuit::packed_instruction::PackedInstruction;
+use rustworkx_core::petgraph::stable_graph::NodeIndex;
+
+#[pyfunction]
+#[pyo3(name = "filter_op_nodes")]
+pub fn py_filter_op_nodes(
+    py: Python,
+    dag: &mut DAGCircuit,
+    predicate: &Bound<PyAny>,
+) -> PyResult<()> {
+    let callable = |node: NodeIndex| -> PyResult<bool> {
+        let dag_op_node = dag.get_node(py, node)?;
+        predicate.call1((dag_op_node,))?.extract()
+    };
+    let mut remove_nodes: Vec<NodeIndex> = Vec::new();
+    for node in dag.op_nodes(true) {
+        if !callable(node)? {
+            remove_nodes.push(node);
+        }
+    }
+    for node in remove_nodes {
+        dag.remove_op_node(node);
+    }
+    Ok(())
+}
+
+/// Remove any nodes that have the provided label set
+///
+/// Args:
+///     dag (DAGCircuit): The dag circuit to filter the ops from
+///     label (str): The label to filter nodes on
+#[pyfunction]
+pub fn filter_labeled_op(dag: &mut DAGCircuit, label: String) {
+    let predicate = |node: &PackedInstruction| -> bool {
+        match node.label() {
+            Some(inst_label) => inst_label != label,
+            None => false,
+        }
+    };
+    dag.filter_op_nodes(predicate);
+}
+
+pub fn filter_op_nodes_mod(m: &Bound<PyModule>) -> PyResult<()> {
+    m.add_wrapped(wrap_pyfunction!(py_filter_op_nodes))?;
+    m.add_wrapped(wrap_pyfunction!(filter_labeled_op))?;
+    Ok(())
+}
diff --git a/crates/accelerate/src/lib.rs b/crates/accelerate/src/lib.rs
index e8760ee2c616..78eea97faad0 100644
--- a/crates/accelerate/src/lib.rs
+++ b/crates/accelerate/src/lib.rs
@@ -22,6 +22,7 @@ pub mod dense_layout;
 pub mod edge_collections;
 pub mod error_map;
 pub mod euler_one_qubit_decomposer;
+pub mod filter_op_nodes;
 pub mod isometry;
 pub mod nlayout;
 pub mod optimize_1q_gates;
diff --git a/crates/accelerate/src/remove_diagonal_gates_before_measure.rs b/crates/accelerate/src/remove_diagonal_gates_before_measure.rs
index 10916a77fca8..cf2c738f131a 100644
--- a/crates/accelerate/src/remove_diagonal_gates_before_measure.rs
+++ b/crates/accelerate/src/remove_diagonal_gates_before_measure.rs
@@ -49,7 +49,9 @@ fn run_remove_diagonal_before_measure(dag: &mut DAGCircuit) -> PyResult<()> {
     let mut nodes_to_remove = Vec::new();
     for index in dag.op_nodes(true) {
         let node = &dag.dag[index];
-        let NodeType::Operation(inst) = node else {panic!()};
+        let NodeType::Operation(inst) = node else {
+            panic!()
+        };

         if inst.op.name() == "measure" {
             let predecessor = (dag.quantum_predecessors(index))
diff --git a/crates/circuit/src/dag_circuit.rs b/crates/circuit/src/dag_circuit.rs
index 381ef25b7a7e..32b3a77ed24c 100644
--- a/crates/circuit/src/dag_circuit.rs
+++ b/crates/circuit/src/dag_circuit.rs
@@ -5794,6 +5794,25 @@ impl DAGCircuit {
         }
     }

+    // Filter any nodes that don't match a given predicate function
+    pub fn filter_op_nodes<F>(&mut self, mut predicate: F)
+    where
+        F: FnMut(&PackedInstruction) -> bool,
+    {
+        let mut remove_nodes: Vec<NodeIndex> = Vec::new();
+        for node in self.op_nodes(true) {
+            let NodeType::Operation(op) = &self.dag[node] else {
+                unreachable!()
+            };
+            if !predicate(op) {
+                remove_nodes.push(node);
+            }
+        }
+        for node in remove_nodes {
+            self.remove_op_node(node);
+        }
+    }
+
     pub fn op_nodes_by_py_type<'a>(
         &'a self,
         op: &'a Bound<PyType>,
diff --git a/crates/circuit/src/packed_instruction.rs b/crates/circuit/src/packed_instruction.rs
index 77ca0c6c02dd..df8f9801314a 100644
--- a/crates/circuit/src/packed_instruction.rs
+++ b/crates/circuit/src/packed_instruction.rs
@@ -553,6 +553,13 @@ impl PackedInstruction {
             .and_then(|extra| extra.condition.as_ref())
     }

+    #[inline]
+    pub fn label(&self) -> Option<&str> {
+        self.extra_attrs
+            .as_ref()
+            .and_then(|extra| extra.label.as_deref())
+    }
+
     /// Build a reference to the Python-space operation object (the `Gate`, etc) packed into this
     /// instruction.  This may construct the reference if the `PackedInstruction` is a standard
     /// gate with no already stored operation.
diff --git a/crates/pyext/src/lib.rs b/crates/pyext/src/lib.rs
index 1478fb367a13..49e44bffa2ec 100644
--- a/crates/pyext/src/lib.rs
+++ b/crates/pyext/src/lib.rs
@@ -16,8 +16,9 @@ use qiskit_accelerate::{
     circuit_library::circuit_library, commutation_analysis::commutation_analysis,
     commutation_checker::commutation_checker, convert_2q_block_matrix::convert_2q_block_matrix,
     dense_layout::dense_layout, error_map::error_map,
-    euler_one_qubit_decomposer::euler_one_qubit_decomposer, isometry::isometry, nlayout::nlayout,
-    optimize_1q_gates::optimize_1q_gates, pauli_exp_val::pauli_expval,
+    euler_one_qubit_decomposer::euler_one_qubit_decomposer, filter_op_nodes::filter_op_nodes_mod,
+    isometry::isometry, nlayout::nlayout, optimize_1q_gates::optimize_1q_gates,
+    pauli_exp_val::pauli_expval,
     remove_diagonal_gates_before_measure::remove_diagonal_gates_before_measure, results::results,
     sabre::sabre, sampled_exp_val::sampled_exp_val, sparse_pauli_op::sparse_pauli_op,
     star_prerouting::star_prerouting, stochastic_swap::stochastic_swap, synthesis::synthesis,
@@ -46,6 +47,7 @@ fn _accelerate(m: &Bound<PyModule>) -> PyResult<()> {
     add_submodule(m, dense_layout, "dense_layout")?;
     add_submodule(m, error_map, "error_map")?;
     add_submodule(m, euler_one_qubit_decomposer, "euler_one_qubit_decomposer")?;
+    add_submodule(m, filter_op_nodes_mod, "filter_op_nodes")?;
     add_submodule(m, isometry, "isometry")?;
     add_submodule(m, nlayout, "nlayout")?;
     add_submodule(m, optimize_1q_gates, "optimize_1q_gates")?;
diff --git a/qiskit/__init__.py b/qiskit/__init__.py
index d9979c9d4d92..3cc10bf96a31 100644
--- a/qiskit/__init__.py
+++ b/qiskit/__init__.py
@@ -92,6 +92,7 @@
 sys.modules["qiskit._accelerate.commutation_checker"] = _accelerate.commutation_checker
 sys.modules["qiskit._accelerate.commutation_analysis"] = _accelerate.commutation_analysis
 sys.modules["qiskit._accelerate.synthesis.linear_phase"] = _accelerate.synthesis.linear_phase
+sys.modules["qiskit._accelerate.filter_op_nodes"] = _accelerate.filter_op_nodes

 from qiskit.exceptions import QiskitError, MissingOptionalLibraryError

diff --git a/qiskit/transpiler/passes/utils/filter_op_nodes.py b/qiskit/transpiler/passes/utils/filter_op_nodes.py
index 344d2280e3f4..75b824332aee 100644
--- a/qiskit/transpiler/passes/utils/filter_op_nodes.py
+++ b/qiskit/transpiler/passes/utils/filter_op_nodes.py
@@ -18,6 +18,8 @@
 from qiskit.transpiler.basepasses import TransformationPass
 from qiskit.transpiler.passes.utils import control_flow

+from qiskit._accelerate.filter_op_nodes import filter_op_nodes
+

 class FilterOpNodes(TransformationPass):
     """Remove all operations that match a filter function
@@ -59,7 +61,5 @@ def __init__(self, predicate: Callable[[DAGOpNode], bool]):
     @control_flow.trivial_recurse
     def run(self, dag: DAGCircuit) -> DAGCircuit:
         """Run the RemoveBarriers pass on `dag`."""
-        for node in dag.op_nodes():
-            if not self.predicate(node):
-                dag.remove_op_node(node)
+        filter_op_nodes(dag, self.predicate)
         return dag

Port `FilterOpNodes` to Rust

Port `FilterOpNodes` to Rust

performance
Changelog: None
Rust
mod: transpiler


diff --git a/docs/source/configuration_file.rst b/docs/source/configuration_file.rst
index 8e9169961..ab31be8a5 100644
--- a/docs/source/configuration_file.rst
+++ b/docs/source/configuration_file.rst
@@ -2278,8 +2278,28 @@ Postprocessing

 Evaluation Parameters
 ---------------------
-Dict. Parameters to get object detection metrics (true positive and false detection rates), and this, for defined
-object sizes.
+Dict. Parameters to get object detection metrics (lesions true positive rate, lesions false detection rate
+and Hausdorff score), and this, for defined object sizes.
+
+.. jsonschema::
+
+    {
+        "$schema": "http://json-schema.org/draft-04/schema#",
+        "title": "object_detection_metrics",
+        "$$description": [
+            "Indicate if object detection metrics (lesions true positive rate, lesions false detection rate\n",
+            "and Hausdorff score) are computed or not at evaluation time. Default: ``true``",
+        ],
+        "type": "boolean"
+    }
+
+.. code-block:: JSON
+
+    {
+        "evaluation_parameters": {
+            "object_detection_metrics": true
+        }
+    }

 .. jsonschema::
@@ -2295,7 +2315,8 @@ object sizes.
                     "These values will create several consecutive target size bins. For instance\n",
                     "with a list of two values, we will have three target size bins: minimal size\n",
                     "to first list element, first list element to second list element, and second\n",
-                    "list element to infinity. Default: ``[20, 100]``."
+                    "list element to infinity. Default: ``[20, 100]``.\n",
+                    "``object_detection_metrics`` must be ``true`` for the target_size to apply."
                 ]
             },
             "unit": {
@@ -2329,7 +2350,10 @@ object sizes.
         "options": {
             "thr": {
                 "type": "int",
-                "description": "Minimal object size overlapping to be considered a TP, FP, or FN. Default: ``3``."
+                "$$description": [
+                    "Minimal object size overlapping to be considered a TP, FP, or FN. Default: ``3``.\n",
+                    "``object_detection_metrics`` must be ``true`` for the overlap to apply."
+                ]
             },
             "unit": {
                 "type": "string",
diff --git a/docs/source/tutorials/two_class_microscopy_seg_2d_unet.rst b/docs/source/tutorials/two_class_microscopy_seg_2d_unet.rst
index d01bade35..555a1f1bf 100644
--- a/docs/source/tutorials/two_class_microscopy_seg_2d_unet.rst
+++ b/docs/source/tutorials/two_class_microscopy_seg_2d_unet.rst
@@ -114,9 +114,9 @@ microscopy segmentation training.
       "train_fraction": 0.6
       "test_fraction": 0.1

-- ``training_parameters:training_time:num_epochs``: The maximum number of epochs that will be run during training. Each epoch is composed
-  of a training part and a validation part. It should be a strictly positive integer. In our case, we will use
-  50 epochs.
+- ``training_parameters:training_time:num_epochs``: The maximum number of epochs that will be run during training.
+  Each epoch is composed of a training part and a validation part. It should be a strictly positive integer.
+  In our case, we will use 50 epochs.

   .. code-block:: xml

@@ -132,12 +132,21 @@ microscopy segmentation training.
      "length_2D": [256, 256]
      "stride_2D": [244, 244]

-- ``postprocessing:binarize_maxpooling``: Used to binarize predictions across all classes in multiclass models. For each pixel, the class, including the background class, with the highest output probability will be segmented.
+- ``postprocessing:binarize_maxpooling``: Used to binarize predictions across all classes in multiclass models.
+  For each pixel, the class, including the background class, with the highest output probability will be segmented.

   .. code-block:: xml

       "binarize_maxpooling": {}

+- ``evaluation_parameters:object_detection_metrics``: Used to indicate if object detection metrics
+  (lesions true positive rate, lesions false detection rate and Hausdorff score) are computed or
+  not at evaluation time. For the axons and myelin segmentation task, we set this parameter to ``false``.
+
+  .. code-block:: xml
+
+      "object_detection_metrics": false
+
 - ``transformation:Resample``: Used to resample images to a common resolution (in mm) before splitting into patches,
   according to each image real pixel size. In our case, we resample the images to a common resolution of 0.0001 mm
   (0.1 μm) in both dimensions.
diff --git a/ivadomed/config/config_default.json b/ivadomed/config/config_default.json
index 70aa7c830..70c971f05 100644
--- a/ivadomed/config/config_default.json
+++ b/ivadomed/config/config_default.json
@@ -95,6 +95,7 @@
     },
     "postprocessing": {},
     "evaluation_parameters": {
+        "object_detection_metrics": true
     },
     "transformation": {
     }
diff --git a/ivadomed/config/config_microscopy.json b/ivadomed/config/config_microscopy.json
index c80556b06..d17dab35d 100644
--- a/ivadomed/config/config_microscopy.json
+++ b/ivadomed/config/config_microscopy.json
@@ -91,6 +91,9 @@
     "postprocessing": {
         "binarize_maxpooling": {}
     },
+    "evaluation_parameters": {
+        "object_detection_metrics": false
+    },
     "transformation": {
         "Resample":
         {
diff --git a/ivadomed/evaluation.py b/ivadomed/evaluation.py
index 8e7c4826d..241ac5383 100644
--- a/ivadomed/evaluation.py
+++ b/ivadomed/evaluation.py
@@ -86,24 +86,25 @@ def evaluate(bids_df, path_output, target_suffix, eval_params):
                                    params=eval_params)
         results_pred, data_painted = eval.run_eval()

-        # SAVE PAINTED DATA, TP FP FN
-        fname_paint = str(fname_pred).split('.nii.gz')[0] + '_painted.nii.gz'
-        nib_painted = nib.Nifti1Image(
-            dataobj=data_painted,
-            affine=nib_pred.header.get_best_affine(),
-            header=nib_pred.header.copy()
-        )
-        nib.save(nib_painted, fname_paint)
-
-        # For Microscopy PNG/TIF files (TODO: implement OMETIFF behavior)
-        if "nii" not in extension:
-            painted_list = imed_inference.split_classes(nib_painted)
-            # Reformat target list to include class index and be compatible with multiple raters
-            target_list = ["_class-%d" % i for i in range(len(target_suffix))]
-            imed_inference.pred_to_png(painted_list,
-                                       target_list,
-                                       str(path_preds.joinpath(subj_acq)),
-                                       suffix="_pred_painted.png")
+        if eval_params['object_detection_metrics']:
+            # SAVE PAINTED DATA, TP FP FN
+            fname_paint = str(fname_pred).split('.nii.gz')[0] + '_painted.nii.gz'
+            nib_painted = nib.Nifti1Image(
+                dataobj=data_painted,
+                affine=nib_pred.header.get_best_affine(),
+                header=nib_pred.header.copy()
+            )
+            nib.save(nib_painted, fname_paint)
+
+            # For Microscopy PNG/TIF files (TODO: implement OMETIFF behavior)
+            if "nii" not in extension:
+                painted_list = imed_inference.split_classes(nib_painted)
+                # Reformat target list to include class index and be compatible with multiple raters
+                target_list = ["_class-%d" % i for i in range(len(target_suffix))]
+                imed_inference.pred_to_png(painted_list,
+                                           target_list,
+                                           str(path_preds.joinpath(subj_acq)),
+                                           suffix="_pred_painted.png")

         # SAVE RESULTS FOR THIS PRED
         results_pred['image_id'] = subj_acq
@@ -135,6 +136,8 @@ class Evaluation3DMetrics(object):
         bin_struct (ndarray): Binary structure.
         size_min (int): Minimum size of objects. Objects that are smaller than this limit can be removed if
             "removeSmall" is in params.
+        object_detection_metrics (bool): Indicate if object detection metrics (lesions true positive and false detection
+            rates) are computed or not.
         overlap_vox (int): A prediction and ground-truth are considered as overlapping if they overlap for at least this
             amount of voxels.
         overlap_ratio (float): A prediction and ground-truth are considered as overlapping if they overlap for at least
@@ -167,6 +170,8 @@ def __init__(self, data_pred, data_gt, dim_lst, params=None):
         self.postprocessing_dict = {}
         self.size_min = 0

+        self.object_detection_metrics = params["object_detection_metrics"]
+
         if "target_size" in params:
             self.size_rng_lst, self.size_suffix_lst = \
                 self._get_size_ranges(thr_lst=params["target_size"]["thr"],
@@ -389,8 +394,12 @@ def get_ltpr(self, label_size=None, class_idx=0):
             label_size (int): Size of label.
             class_idx (int): Label index. If monolabel 0, else ranges from 0 to number of output channels - 1.

-        Note: computed only if n_obj >= 1.
+        Note: computed only if n_obj >= 1 and "object_detection_metrics" evaluation parameter is True.
         """
+        if not self.object_detection_metrics:
+            n_obj = 0
+            return np.nan, n_obj
+
         ltp, lfn, n_obj = self._get_ltp_lfn(label_size, class_idx)

         denom = ltp + lfn
@@ -406,8 +415,12 @@ def get_lfdr(self, label_size=None, class_idx=0):
             label_size (int): Size of label.
             class_idx (int): Label index. If monolabel 0, else ranges from 0 to number of output channels - 1.

-        Note: computed only if n_obj >= 1.
+        Note: computed only if n_obj >= 1 and "object_detection_metrics" evaluation parameter is True.
         """
+
+        if not self.object_detection_metrics:
+            return np.nan
+
         ltp, _, n_obj = self._get_ltp_lfn(label_size, class_idx)
         lfp = self._get_lfp(label_size, class_idx)

diff --git a/ivadomed/main.py b/ivadomed/main.py
index 2ea68a458..375f65aa1 100644
--- a/ivadomed/main.py
+++ b/ivadomed/main.py
@@ -517,7 +517,8 @@ def run_command(context, n_gif=0, thr_increment=None, resume_training=False, no_
                                                                            'requires_undo': True}}, device=device,
                                            cuda_available=cuda_available)

-        metric_fns = imed_metrics.get_metric_fns(ds_test.task)
+        eval_params = context[ConfigKW.EVALUATION_PARAMETERS]
+        metric_fns = imed_metrics.get_metric_fns(ds_test.task, eval_params)

         if ModelParamsKW.FILM_LAYERS in model_params and any(model_params[ModelParamsKW.FILM_LAYERS]):
             ds_test, model_params = update_film_model_params(context, ds_test, model_params, path_output)
@@ -535,7 +536,7 @@ def run_command(context, n_gif=0, thr_increment=None, resume_training=False, no_
         # RUN EVALUATION
         df_results = imed_evaluation.evaluate(bids_df, path_output=path_output,
                                               target_suffix=loader_params[LoaderParamsKW.TARGET_SUFFIX],
-                                              eval_params=context[ConfigKW.EVALUATION_PARAMETERS])
+                                              eval_params=eval_params)
         return df_results, pred_metrics

diff --git a/ivadomed/metrics.py b/ivadomed/metrics.py
index a0ba1f466..f8a49ab73 100644
--- a/ivadomed/metrics.py
+++ b/ivadomed/metrics.py
@@ -6,7 +6,7 @@

 # METRICS
-def get_metric_fns(task):
+def get_metric_fns(task, eval_params=None):
     metric_fns = [dice_score,
                   multi_class_dice_score,
                   precision_score,
@@ -14,7 +14,11 @@ def get_metric_fns(task):
                   specificity_score,
                   intersection_over_union,
                   accuracy_score]
-    if task == "segmentation":
+    if eval_params:
+        object_detection_metrics = eval_params['object_detection_metrics']
+    else:
+        object_detection_metrics = True
+    if task == "segmentation" and object_detection_metrics:
         metric_fns = metric_fns + [hausdorff_score]

     return metric_fns

Time issue when running `--test` on large microscopy images
**Disclaimer**
This issue is based on preliminary observations and does not contain a lot of details/examples at the moment.

## Issue description

I tried running the `--test` command on large microscopy images on `rosenberg`, and it takes a very long time.

I'll reference the following screenshot:
![test_on_om](https://user-images.githubusercontent.com/54086142/138940644-3f9618f9-d75e-41fb-9b3a-f10e8dc7a782.png)

From the `test` function called in main [here](https://github.com/ivadomed/ivadomed/blob/5207e7a5f0c882ae66248322ffbc518c0d2d6ce7/ivadomed/main.py#L470), the following happen:

At 19h20m56s, the segmentation of the first image starts (8557x7115 pixels).
At 19h29m59s, the segmentation of the second image starts  (5895x5266 pixels).
At 19h32m25s, the segmentations are completed, those timings are about normal (maybe a bit long, as describe in #972)

At this point, the test function is at line 85  in `testing.py` [here](https://github.com/ivadomed/ivadomed/blob/5207e7a5f0c882ae66248322ffbc518c0d2d6ce7/ivadomed/testing.py#L85), so it seems that the `metric_mgr.get_results()` function takes 40 minutes to compute the evaluation metrics (until 20h12m09s), which seems excessive.

After that, the `evaluate` function called in main [here](https://github.com/ivadomed/ivadomed/blob/5207e7a5f0c882ae66248322ffbc518c0d2d6ce7/ivadomed/main.py#L480) runs on the first image (8557x7115 pixels).
I killed the process after 2.5 hours without progress.
In that case, I suspect that the bottleneck may be in [run_eval](https://github.com/ivadomed/ivadomed/blob/5207e7a5f0c882ae66248322ffbc518c0d2d6ce7/ivadomed/evaluation.py#L87).

Further investigations are needed.
I would be happy to make my model and images available if someone is interested to look into this.
Thanks!

Update: I ran the same test again to have a little bit more details.
Turns out, the `evaluate` function took ~5h30min to compute metrics on my 2 images 😬.

![eval_om_02](https://user-images.githubusercontent.com/54086142/139498300-1c13bd3f-72af-4ce2-b0d7-daa66aa4f9f5.png)

I also ran the command with a cProfile (below) so we may have an idea where to look first.

<details>
<summary> cProfile output </summary>

```
         5766383 function calls (5632045 primitive calls) in 22538.294 seconds

   Ordered by: cumulative time

   ncalls  tottime  percall  cumtime  percall filename:lineno(function)
   2852/1    0.044    0.000 22538.307 22538.307 {built-in method builtins.exec}
        1    0.000    0.000 22538.307 22538.307 profile_test.py:1(<module>)
        1    0.032    0.032 22536.006 22536.006 main.py:295(run_command)
        1    1.535    1.535 19570.459 19570.459 evaluation.py:20(evaluate)
        2    0.001    0.001 19544.642 9772.321 evaluation.py:414(run_eval)
        4    0.019    0.005 14119.131 3529.783 evaluation.py:396(get_lfdr)
        8 4776.952  597.119 10804.197 1350.525 evaluation.py:305(_get_ltp_lfn)
        4 4290.862 1072.716 8717.269 2179.317 evaluation.py:344(_get_lfp)
    65904 7953.343    0.121 7953.343    0.121 {method 'astype' of 'numpy.ndarray' objects}
        4    0.008    0.002 5402.362 1350.591 evaluation.py:379(get_ltpr)
        1    0.018    0.018 2953.772 2953.772 testing.py:26(test)
69279/68504    1.798    0.000 2733.607    0.040 {built-in method numpy.core._multiarray_umath.implement_array_function}
        1    0.000    0.000 2381.068 2381.068 metrics.py:40(__call__)
        2    0.114    0.057 2357.215 1178.608 metrics.py:144(hausdorff_score)
    28944    0.160    0.000 2357.101    0.081 distance.py:365(directed_hausdorff)
    28944 2351.782    0.081 2355.295    0.081 {scipy.spatial._hausdorff.directed_hausdorff}
    27641    0.146    0.000 2327.396    0.084 <__array_function__ internals>:2(count_nonzero)
    27641    0.081    0.000 2327.107    0.084 numeric.py:424(count_nonzero)
    27641 2327.026    0.084 2327.026    0.084 {built-in method numpy.core._multiarray_umath.count_nonzero}
        1    0.061    0.061  570.711  570.711 testing.py:91(run_inference)
      141    0.003    0.000  531.454    3.769 std.py:1157(__iter__)
      135    0.001    0.000  531.343    3.936 dataloader.py:344(__next__)
      135    0.431    0.003  531.341    3.936 dataloader.py:383(_next_data)
      134    0.002    0.000  530.769    3.961 fetch.py:42(fetch)
      134    0.003    0.000  530.618    3.960 fetch.py:44(<listcomp>)
      402    1.184    0.003  530.616    1.320 mri2d_segmentation_dataset.py:169(__getitem__)
    12293  362.383    0.029  362.383    0.029 {method 'reduce' of 'numpy.ufunc' objects}
     9989    0.091    0.000  335.476    0.034 fromnumeric.py:70(_wrapreduction)
     9407    0.059    0.000  330.356    0.035 <__array_function__ internals>:2(amax)
     9407    0.083    0.000  330.240    0.035 fromnumeric.py:2617(amax)
      406    0.527    0.001  287.663    0.709 postprocessing.py:24(wrapper)
      406  107.077    0.264  287.136    0.707 postprocessing.py:75(threshold_predictions)
     1206    0.228    0.000  204.129    0.169 transforms.py:147(__call__)
127600/127428  168.456    0.001  173.261    0.001 {built-in method numpy.array}
      402    0.004    0.000  107.236    0.267 transforms.py:31(wrapper)
      402   41.435    0.103  107.231    0.267 transforms.py:298(__call__)
      804    0.012    0.000   96.664    0.120 transforms.py:212(__call__)
      409    0.002    0.000   65.816    0.161 <__array_function__ internals>:2(copy)
      409    0.001    0.000   65.809    0.161 function_base.py:715(copy)
      404    0.004    0.000   57.760    0.143 {method 'std' of 'numpy.ndarray' objects}
      404    0.234    0.001   57.756    0.143 _methods.py:259(_std)
      404   41.825    0.104   57.521    0.142 _methods.py:194(_var)
       22   33.902    1.541   37.671    1.712 metrics.py:62(numeric_score)
        4    0.000    0.000   21.153    5.288 inference.py:211(pred_to_png)
        8    0.000    0.000   20.135    2.517 functions.py:270(imwrite)
        8    0.007    0.001   20.129    2.516 format.py:477(append_data)
        8    0.013    0.002   20.122    2.515 pillow.py:359(_append_data)
        8    0.001    0.000   16.736    2.092 pillow.py:202(_append_data)
    39030   15.270    0.000   15.270    0.000 {method 'encode' of 'ImagingEncoder' objects}
      134   14.787    0.110   14.787    0.110 {method 'cpu' of 'torch._C._TensorBase' objects}
        8    0.000    0.000   14.650    1.831 Image.py:2095(save)
        8    0.000    0.000   14.470    1.809 PngImagePlugin.py:1178(_save)
        8    0.000    0.000   14.470    1.809 ImageFile.py:478(_save)
        1    0.004    0.004    9.926    9.926 loader.py:12(load_dataset)
        1    0.001    0.001    9.909    9.909 mri2d_segmentation_dataset.py:90(load_filenames)
        6    0.000    0.000    9.660    1.610 metrics.py:215(specificity_score)
        6    0.000    0.000    9.653    1.609 metrics.py:194(recall_score)
        6    0.000    0.000    9.651    1.608 metrics.py:172(precision_score)
        4    0.000    0.000    9.160    2.290 segmentation_pair.py:38(__init__)
        8    0.000    0.000    9.156    1.144 segmentation_pair.py:272(read_file)
        8    0.720    0.090    9.156    1.144 segmentation_pair.py:296(convert_file_to_nifti)
        2    1.160    0.580    8.364    4.182 inference.py:121(pred_to_nib)
      402    0.003    0.000    8.293    0.021 {method 'mean' of 'numpy.ndarray' objects}
      402    0.022    0.000    8.290    0.021 _methods.py:161(_mean)
        4    0.000    0.000    7.991    1.998 loadsave.py:81(save)
        4    0.000    0.000    7.991    1.998 filebasedimages.py:318(to_filename)
        4    0.000    0.000    7.991    1.998 analyze.py:996(to_file_map)
        4    0.000    0.000    7.863    1.966 arraywriters.py:539(to_fileobj)
        4    0.002    0.000    7.863    1.966 volumeutils.py:538(array_to_file)
        4    0.627    0.157    7.860    1.965 volumeutils.py:763(_write_data)
        8    0.000    0.000    6.200    0.775 functions.py:237(imread)
        2    0.391    0.196    5.694    2.847 evaluation.py:145(__init__)
      174    0.001    0.000    5.258    0.030 <__array_function__ internals>:2(sum)
      174    0.001    0.000    5.257    0.030 fromnumeric.py:2111(sum)
1928/1896    0.003    0.000    5.087    0.003 _asarray.py:110(asanyarray)
       34    0.000    0.000    5.080    0.149 dataobj_images.py:212(get_fdata)
        8    0.000    0.000    4.726    0.591 measurements.py:44(label)
        8    4.725    0.591    4.726    0.591 {built-in method scipy.ndimage._ni_label._label}
       16    0.000    0.000    4.460    0.279 openers.py:164(write)
       20    0.000    0.000    4.460    0.223 gzip.py:263(write)
        2    0.000    0.000    4.356    2.178 metrics.py:253(accuracy_score)
        2    0.000    0.000    4.353    2.176 metrics.py:236(intersection_over_union)
      209    0.008    0.000    4.020    0.019 __init__.py:1(<module>)
       10    2.341    0.234    3.618    0.362 metrics.py:91(dice_score)
       72    3.577    0.050    3.577    0.050 {method 'tobytes' of 'numpy.ndarray' objects}
        2    0.009    0.004    3.478    1.739 postprocessing.py:271(apply)
        2    0.540    0.270    3.470    1.735 postprocessing.py:284(binarize_maxpooling)
        8    0.698    0.087    3.372    0.422 util.py:45(image_as_uint)
        8    0.000    0.000    3.216    0.402 format.py:333(get_data)
        8    0.000    0.000    3.216    0.402 pillow.py:300(_get_data)
        8    0.000    0.000    3.216    0.402 pillow.py:167(_get_data)
        8    0.043    0.005    3.215    0.402 pillow.py:698(pil_get_frame)
       10    0.000    0.000    3.065    0.306 arrayproxy.py:370(__array__)
       10    0.022    0.002    3.065    0.306 arrayproxy.py:347(_get_scaled)
    28945    0.078    0.000    3.061    0.000 contextlib.py:72(inner)
       16    2.991    0.187    2.991    0.187 {method 'compress' of 'zlib.Compress' objects}
        8    0.000    0.000    2.983    0.373 functions.py:148(get_reader)
       16    0.000    0.000    2.978    0.186 format.py:215(__init__)
        8    0.000    0.000    2.977    0.372 format.py:157(get_reader)
        8    0.000    0.000    2.977    0.372 pillow.py:297(_open)
        8    0.000    0.000    2.977    0.372 pillow.py:118(_open)
        8    0.000    0.000    2.976    0.372 pillow.py:666(pil_try_read)
       16    0.000    0.000    2.976    0.186 Image.py:1274(getdata)
       24    0.003    0.000    2.975    0.124 ImageFile.py:154(load)
     1412    2.937    0.002    2.937    0.002 {method 'decode' of 'ImagingDecoder' objects}
    28945    2.643    0.000    2.643    0.000 {function SeedSequence.generate_state at 0x7f6986d91160}
10050/134    0.051    0.000    2.415    0.018 module.py:540(__call__)
      134    0.002    0.000    2.410    0.018 models.py:509(forward)
      539    2.342    0.004    2.342    0.004 {method 'copy' of 'numpy.ndarray' objects}
  2674/13    0.014    0.000    2.312    0.178 <frozen importlib._bootstrap>:986(_find_and_load)
  2668/13    0.010    0.000    2.311    0.178 <frozen importlib._bootstrap>:956(_find_and_load_unlocked)
  2223/13    0.011    0.000    2.310    0.178 <frozen importlib._bootstrap>:650(_load_unlocked)
  1934/13    0.006    0.000    2.309    0.178 <frozen importlib._bootstrap_external>:777(exec_module)
  3094/13    0.002    0.000    2.309    0.178 <frozen importlib._bootstrap>:211(_call_with_frames_removed)
5733/2516    0.008    0.000    2.227    0.001 <frozen importlib._bootstrap>:1017(_handle_fromlist)
   952/26    0.003    0.000    2.224    0.086 {built-in method builtins.__import__}
        2    0.000    0.000    2.190    1.095 main.py:1(<module>)
     1206    0.034    0.000    2.150    0.002 models.py:255(forward)
      120    0.000    0.000    2.112    0.018 fromnumeric.py:52(_wrapfunc)
120752/120617    0.057    0.000    2.106    0.000 _asarray.py:23(asarray)
        8    0.001    0.000    2.084    0.261 pillow.py:795(ndarray_to_pil)
        2    0.000    0.000    2.013    1.006 testing.py:412(get_gt)
        4    0.476    0.119    2.000    0.500 metrics.py:120(mse)
        1    0.000    0.000    1.970    1.970 serialization.py:499(load)
        1    0.000    0.000    1.970    1.970 serialization.py:613(_legacy_load)
        1    0.002    0.002    1.941    1.941 {method 'load' of '_pickle.Unpickler' objects}
      139    0.001    0.000    1.935    0.014 serialization.py:713(persistent_load)
     2546    0.006    0.000    1.925    0.001 conv.py:348(forward)
     1200    1.921    0.002    1.921    0.002 {built-in method zlib.crc32}
     2546    0.009    0.000    1.917    0.001 conv.py:340(_conv_forward)
    29058    1.912    0.000    1.912    0.000 {built-in method __new__ of type object at 0xab0800}
     2546    1.906    0.001    1.906    0.001 {built-in method conv2d}
      128    0.001    0.000    1.903    0.015 serialization.py:810(restore_location)
      128    0.000    0.000    1.902    0.015 serialization.py:176(default_restore_location)
      128    0.001    0.000    1.902    0.015 serialization.py:152(_cuda_deserialize)
      128    0.001    0.000    1.899    0.015 __init__.py:428(_lazy_new)
        8    0.046    0.006    1.736    0.217 Image.py:679(__array_interface__)
       76    0.000    0.000    1.693    0.022 <__array_function__ internals>:2(argmax)
       76    0.000    0.000    1.693    0.022 fromnumeric.py:1119(argmax)
       76    1.693    0.022    1.693    0.022 {method 'argmax' of 'numpy.ndarray' objects}
        8    0.013    0.002    1.690    0.211 Image.py:710(tobytes)
    65183    0.538    0.000    1.658    0.000 __init__.py:266(__getattr__)
       10    0.000    0.000    1.527    0.153 arrayproxy.py:328(_get_unscaled)
       10    0.288    0.029    1.525    0.153 volumeutils.py:449(array_from_file)
       24    0.000    0.000    1.370    0.057 evaluation.py:270(get_vol)
        8    0.000    0.000    1.338    0.167 Image.py:2736(fromarray)
      134    0.008    0.000    1.290    0.010 models.py:434(forward)
       39    0.000    0.000    1.277    0.033 {method 'sum' of 'numpy.ndarray' objects}
       39    0.000    0.000    1.277    0.033 _methods.py:45(_sum)
        1    0.000    0.000    1.245    1.245 bids_dataframe.py:31(__init__)
        1    0.001    0.001    1.241    1.241 bids_dataframe.py:69(create_bids_dataframe)
       10    0.140    0.014    1.236    0.124 {method 'readinto' of '_io._BufferedIOBase' objects}
    29558    0.035    0.000    1.217    0.000 <__array_function__ internals>:2(concatenate)
        1    0.000    0.000    1.195    1.195 evaluation.py:1(<module>)
      536    0.008    0.000    1.176    0.002 models.py:286(forward)
20458/18466    0.078    0.000    1.170    0.000 {method 'read' of '_io.BufferedReader' objects}
    65393    1.121    0.000    1.121    0.000 {built-in method _warnings.warn}
        2    0.000    0.000    1.115    0.558 inference.py:1(<module>)
      134    0.008    0.000    1.115    0.008 models.py:340(forward)
       46    0.000    0.000    1.099    0.024 gzip.py:287(read)
     1012    0.290    0.000    1.098    0.001 _compression.py:66(readinto)
        2    0.000    0.000    1.052    0.526 metrics.py:274(multi_class_dice_score)
```
</details>

Ouf yeah! The number of false positive lesions and number of true positive and false negative lesions (`_get_ltp_lfn` and `get_lfp`) don't really makes sense here and are probably mostly responsible for this huge computing time! It takes every unconnected objects et verifies overlaps, but in the case of axon and myelin segmentation, there are far from only 1 or 2 objects segmented. A quick workaround would be to put these metrics as options so they can be used only for relevant tasks like MS lesion segmentation.


diff --git a/zerver/lib/users.py b/zerver/lib/users.py
index 3b00d4187ff2d..cb24813fb1a6a 100644
--- a/zerver/lib/users.py
+++ b/zerver/lib/users.py
@@ -10,6 +10,7 @@
 from django.conf import settings
 from django.core.exceptions import ValidationError
 from django.db.models import Q, QuerySet
+from django.db.models.functions import Upper
 from django.utils.translation import gettext as _
 from django_otp.middleware import is_verified
 from typing_extensions import NotRequired
@@ -347,6 +348,74 @@ def access_user_by_email(
     return access_user_common(target, user_profile, allow_deactivated, allow_bots, for_admin)

+def bulk_access_users_by_email(
+    emails: list[str],
+    *,
+    acting_user: UserProfile,
+    allow_deactivated: bool = False,
+    allow_bots: bool = False,
+    for_admin: bool,
+) -> set[UserProfile]:
+    # We upper-case the email addresses ourselves here, because
+    # `email__iexact__in=emails` is not supported by Django.
+    target_emails_upper = [email.strip().upper() for email in emails]
+    users = (
+        UserProfile.objects.annotate(email_upper=Upper("email"))
+        .select_related(
+            "realm",
+            "realm__can_access_all_users_group",
+            "realm__can_access_all_users_group__named_user_group",
+            "realm__direct_message_initiator_group",
+            "realm__direct_message_initiator_group__named_user_group",
+            "realm__direct_message_permission_group",
+            "realm__direct_message_permission_group__named_user_group",
+            "bot_owner",
+        )
+        .filter(email_upper__in=target_emails_upper, realm=acting_user.realm)
+    )
+    valid_emails_upper = {user_profile.email_upper for user_profile in users}
+    all_users_exist = all(email in valid_emails_upper for email in target_emails_upper)
+
+    if not all_users_exist:
+        raise JsonableError(_("No such user"))
+
+    return {
+        access_user_common(user_profile, acting_user, allow_deactivated, allow_bots, for_admin)
+        for user_profile in users
+    }
+
+
+def bulk_access_users_by_id(
+    user_ids: list[int],
+    *,
+    acting_user: UserProfile,
+    allow_deactivated: bool = False,
+    allow_bots: bool = False,
+    for_admin: bool,
+) -> set[UserProfile]:
+    users = UserProfile.objects.select_related(
+        "realm",
+        "realm__can_access_all_users_group",
+        "realm__can_access_all_users_group__named_user_group",
+        "realm__direct_message_initiator_group",
+        "realm__direct_message_initiator_group__named_user_group",
+        "realm__direct_message_permission_group",
+        "realm__direct_message_permission_group__named_user_group",
+        "bot_owner",
+    ).filter(id__in=user_ids, realm=acting_user.realm)
+
+    valid_user_ids = {user_profile.id for user_profile in users}
+    all_users_exist = all(user_id in valid_user_ids for user_id in user_ids)
+
+    if not all_users_exist:
+        raise JsonableError(_("No such user"))
+
+    return {
+        access_user_common(user_profile, acting_user, allow_deactivated, allow_bots, for_admin)
+        for user_profile in users
+    }
+
+
 class Account(TypedDict):
     realm_name: str
     realm_id: int
diff --git a/zerver/views/streams.py b/zerver/views/streams.py
index 0625369640378..9d83a8ef28618 100644
--- a/zerver/views/streams.py
+++ b/zerver/views/streams.py
@@ -75,20 +75,37 @@
 from zerver.lib.typed_endpoint import ApiParamConfig, PathOnly, typed_endpoint
 from zerver.lib.typed_endpoint_validators import check_color, check_int_in_validator
 from zerver.lib.user_groups import access_user_group_for_setting
-from zerver.lib.users import access_user_by_email, access_user_by_id
+from zerver.lib.users import bulk_access_users_by_email, bulk_access_users_by_id
 from zerver.lib.utils import assert_is_not_none
 from zerver.models import NamedUserGroup, Realm, Stream, UserProfile
 from zerver.models.users import get_system_bot

-def principal_to_user_profile(agent: UserProfile, principal: str | int) -> UserProfile:
-    if isinstance(principal, str):
-        return access_user_by_email(
-            agent, principal, allow_deactivated=False, allow_bots=True, for_admin=False
+def bulk_principals_to_user_profiles(
+    principals: list[str] | list[int],
+    acting_user: UserProfile,
+) -> set[UserProfile]:
+    # Since principals is guaranteed to be non-empty and to have the same type of elements,
+    # the following if/else is safe and enough.
+
+    # principals are user emails.
+    if isinstance(principals[0], str):
+        return bulk_access_users_by_email(
+            principals,  # type: ignore[arg-type] # principals guaranteed to be list[str] only.
+            acting_user=acting_user,
+            allow_deactivated=False,
+            allow_bots=True,
+            for_admin=False,
         )
+
+    # principals are user ids.
     else:
-        return access_user_by_id(
-            agent, principal, allow_deactivated=False, allow_bots=True, for_admin=False
+        return bulk_access_users_by_id(
+            principals,  # type: ignore[arg-type] # principals guaranteed to be list[int] only.
+            acting_user=acting_user,
+            allow_deactivated=False,
+            allow_bots=True,
+            for_admin=False,
         )

@@ -473,12 +490,11 @@ def remove_subscriptions_backend(

     unsubscribing_others = False
     if principals:
-        people_to_unsub = set()
-        for principal in principals:
-            target_user = principal_to_user_profile(user_profile, principal)
-            people_to_unsub.add(target_user)
-            if not user_directly_controls_user(user_profile, target_user):
-                unsubscribing_others = True
+        people_to_unsub = bulk_principals_to_user_profiles(principals, user_profile)
+        unsubscribing_others = any(
+            not user_directly_controls_user(user_profile, target) for target in people_to_unsub
+        )
+
     else:
         people_to_unsub = {user_profile}

@@ -551,6 +567,7 @@ def add_subscriptions_backend(
     realm = user_profile.realm
     stream_dicts = []
     color_map = {}
+    # UserProfile ids or emails.
     if principals is None:
         principals = []

@@ -607,9 +624,8 @@ def add_subscriptions_backend(
             # Guest users case will not be handled here as it will
             # be handled by the decorator above.
             raise JsonableError(_("Insufficient permission"))
-        subscribers = {
-            principal_to_user_profile(user_profile, principal) for principal in principals
-        }
+        subscribers = bulk_principals_to_user_profiles(principals, user_profile)
+
     else:
         subscribers = {user_profile}

diff --git a/zerver/tests/test_subs.py b/zerver/tests/test_subs.py
index f4a3d0b3be9bc..da7dfe34e1660 100644
--- a/zerver/tests/test_subs.py
+++ b/zerver/tests/test_subs.py
@@ -2684,7 +2684,7 @@ def test_realm_admin_remove_multiple_users_from_stream(self) -> None:
             for name in ["cordelia", "prospero", "iago", "hamlet", "outgoing_webhook_bot"]
         ]
         result = self.attempt_unsubscribe_of_principal(
-            query_count=28,
+            query_count=24,
             cache_count=8,
             target_users=target_users,
             is_realm_admin=True,
@@ -2759,7 +2759,7 @@ def test_admin_remove_others_from_stream_legacy_emails(self) -> None:

     def test_admin_remove_multiple_users_from_stream_legacy_emails(self) -> None:
         result = self.attempt_unsubscribe_of_principal(
-            query_count=20,
+            query_count=19,
             target_users=[self.example_user("cordelia"), self.example_user("prospero")],
             is_realm_admin=True,
             is_subbed=True,
@@ -2773,7 +2773,7 @@ def test_admin_remove_multiple_users_from_stream_legacy_emails(self) -> None:

     def test_remove_unsubbed_user_along_with_subbed(self) -> None:
         result = self.attempt_unsubscribe_of_principal(
-            query_count=17,
+            query_count=16,
             target_users=[self.example_user("cordelia"), self.example_user("iago")],
             is_realm_admin=True,
             is_subbed=True,
@@ -2909,6 +2909,43 @@ def test_remove_invalid_user(self) -> None:
         )
         self.assert_json_error(result, "No such user", status_code=400)

+    def test_user_unsubscribe_theirself(self) -> None:
+        """
+        User trying to unsubscribe theirself from the stream, where
+        principals has the id of the acting_user performing the
+        unsubscribe action.
+        """
+        admin = self.example_user("iago")
+        self.login_user(admin)
+        self.assertTrue(admin.is_realm_admin)
+
+        stream_name = "hümbüǵ"
+        self.make_stream(stream_name)
+        self.subscribe(admin, stream_name)
+
+        # unsubscribing when subscribed.
+        result = self.client_delete(
+            "/json/users/me/subscriptions",
+            {
+                "subscriptions": orjson.dumps([stream_name]).decode(),
+                "principals": orjson.dumps([admin.id]).decode(),
+            },
+        )
+        json = self.assert_json_success(result)
+        self.assert_length(json["removed"], 1)
+
+        # unsubscribing after already being unsubscribed.
+        result = self.client_delete(
+            "/json/users/me/subscriptions",
+            {
+                "subscriptions": orjson.dumps([stream_name]).decode(),
+                "principals": orjson.dumps([admin.id]).decode(),
+            },
+        )
+
+        json = self.assert_json_success(result)
+        self.assert_length(json["not_removed"], 1)
+

 class DefaultStreamTest(ZulipTestCase):
     def get_default_stream_names(self, realm: Realm) -> set[str]:
@@ -4686,7 +4723,7 @@ def test_multi_user_subscription(self) -> None:
         streams_to_sub = ["multi_user_stream"]
         with (
             self.capture_send_event_calls(expected_num_events=5) as events,
-            self.assert_database_query_count(38),
+            self.assert_database_query_count(37),
         ):
             self.common_subscribe_to_streams(
                 self.test_user,
@@ -5157,11 +5194,8 @@ def test_bulk_subscribe_many(self) -> None:

         test_user_ids = [user.id for user in test_users]

-        # The only known O(N) behavior here is that we call
-        # principal_to_user_profile for each of our users, but it
-        # should be cached.
         with (
-            self.assert_database_query_count(21),
+            self.assert_database_query_count(16),
             self.assert_memcached_count(3),
             mock.patch("zerver.views.streams.send_messages_for_new_subscribers"),
         ):
@@ -5517,7 +5551,7 @@ def test_subscriptions_query_count(self) -> None:
         ]

         # Test creating a public stream when realm does not have a notification stream.
-        with self.assert_database_query_count(38):
+        with self.assert_database_query_count(37):
             self.common_subscribe_to_streams(
                 self.test_user,
                 [new_streams[0]],
@@ -5525,7 +5559,7 @@ def test_subscriptions_query_count(self) -> None:
             )

         # Test creating private stream.
-        with self.assert_database_query_count(40):
+        with self.assert_database_query_count(39):
             self.common_subscribe_to_streams(
                 self.test_user,
                 [new_streams[1]],
@@ -5537,7 +5571,7 @@ def test_subscriptions_query_count(self) -> None:
         new_stream_announcements_stream = get_stream(self.streams[0], self.test_realm)
         self.test_realm.new_stream_announcements_stream_id = new_stream_announcements_stream.id
         self.test_realm.save()
-        with self.assert_database_query_count(49):
+        with self.assert_database_query_count(48):
             self.common_subscribe_to_streams(
                 self.test_user,
                 [new_streams[2]],
@@ -6015,7 +6049,7 @@ def test_gather_subscriptions(self) -> None:
             polonius.id,
         ]

-        with self.assert_database_query_count(46):
+        with self.assert_database_query_count(43):
             self.common_subscribe_to_streams(
                 self.user_profile,
                 streams,

Make creating streams faster in large organizations
In large organizations (like chat.zulip.org), creating a stream can be very slow (seconds to tens of seconds). We should fix this.

It should be possible to reproduce this by creating 20K users with `manage.py populate_db --extra-users` in a development environment.
Make creating streams faster in large organizations
In large organizations (like chat.zulip.org), creating a stream can be very slow (seconds to tens of seconds). We should fix this.

It should be possible to reproduce this by creating 20K users with `manage.py populate_db --extra-users` in a development environment.

Hello @zulip/server-streams members, this issue was labeled with the "area: stream settings" label, so you may want to check it out!

<!-- areaLabelAddition -->

Hello @zulip/server-streams members, this issue was labeled with the "area: stream settings" label, so you may want to check it out!

<!-- areaLabelAddition -->

area: channel settings
priority: high
size: XL
post release
area: performance


diff --git a/vllm/core/evictor.py b/vllm/core/evictor.py
index ed7e06cab2996..44adc4158abec 100644
--- a/vllm/core/evictor.py
+++ b/vllm/core/evictor.py
@@ -1,6 +1,7 @@
 import enum
+import heapq
 from abc import ABC, abstractmethod
-from typing import OrderedDict, Tuple
+from typing import Dict, List, Tuple

 class EvictionPolicy(enum.Enum):
@@ -75,8 +76,14 @@ class LRUEvictor(Evictor):
     highest num_hashed_tokens value, then one will be chose arbitrarily
     """

+    # CLEANUP_THRESHOLD determines the maximum allowable size of the priority
+    # queue relative to the free table size. When this threshold is exceeded,
+    # a cleanup operation is triggered to reduce memory usage.
+    CLEANUP_THRESHOLD = 50
+
     def __init__(self):
-        self.free_table: OrderedDict[int, BlockMetaData] = OrderedDict()
+        self.free_table: Dict[int, BlockMetaData] = {}
+        self.priority_queue = []

     def __contains__(self, block_id: int) -> bool:
         return block_id in self.free_table
@@ -85,34 +92,50 @@ def evict(self) -> Tuple[int, int]:
         if len(self.free_table) == 0:
             raise ValueError("No usable cache memory left")

-        evicted_block, evicted_block_id = None, None
-        # The blocks with the lowest timestamps should be placed consecutively
-        # at the start of OrderedDict. Loop through all these blocks to
-        # find the one with maximum number of hashed tokens.
-        for _id, block in self.free_table.items():
-            if evicted_block is None:
-                evicted_block, evicted_block_id = block, _id
-                continue
-            if evicted_block.last_accessed < block.last_accessed:
-                break
-            if evicted_block.num_hashed_tokens < block.num_hashed_tokens:
-                evicted_block, evicted_block_id = block, _id
-
-        assert evicted_block is not None
-        assert evicted_block_id is not None
-        self.free_table.pop(evicted_block_id)
-
-        return evicted_block_id, evicted_block.content_hash
+        while self.priority_queue:
+            # We do not remove outdated entries from the priority queue at the
+            # time of updating the last_accessed timestamp. Instead, outdated
+            # entries are filtered out here during eviction. Outdated entries
+            # would either not in the free table, or have older last accessed
+            # time.
+            last_accessed, _, block_id, content_hash = heapq.heappop(
+                self.priority_queue)
+            if (block_id in self.free_table and
+                    self.free_table[block_id].last_accessed == last_accessed):
+                self.free_table.pop(block_id)
+                return block_id, content_hash
+
+        raise ValueError("No usable cache memory left")

     def add(self, block_id: int, content_hash: int, num_hashed_tokens: int,
             last_accessed: float):
         self.free_table[block_id] = BlockMetaData(content_hash,
                                                   num_hashed_tokens,
                                                   last_accessed)
+        heapq.heappush(
+            self.priority_queue,
+            (last_accessed, -num_hashed_tokens, block_id, content_hash))
+        self._cleanup_if_necessary()

     def update(self, block_id: int, last_accessed: float):
         self.free_table[block_id].last_accessed = last_accessed

+    def _cleanup_if_necessary(self):
+        if len(self.priority_queue) > LRUEvictor.CLEANUP_THRESHOLD * len(
+                self.free_table):
+            self._cleanup()
+
+    def _cleanup(self):
+        new_priority_queue: List[Tuple[float, int, int, int]] = []
+
+        for block_id, block in self.free_table.items():
+            new_priority_queue.append(
+                (block.last_accessed, -block.num_hashed_tokens, block_id,
+                 block.content_hash))
+        heapq.heapify(new_priority_queue)
+
+        self.priority_queue = new_priority_queue
+
     def remove(self, block_id: int):
         if block_id not in self.free_table:
             raise ValueError(

[Bug]: Prefix Caching in BlockSpaceManagerV1 and BlockSpaceManagerV2 Increases Time to First Token(TTFT) and Slows Down System
### Your current environment

```text
PyTorch version: 2.3.1+cu121
Is debug build: False
CUDA used to build PyTorch: 12.1
ROCM used to build PyTorch: N/A

OS: Ubuntu 22.04.3 LTS (x86_64)
GCC version: (Ubuntu 11.4.0-1ubuntu1~22.04) 11.4.0
Clang version: Could not collect
CMake version: version 3.27.6
Libc version: glibc-2.35

Python version: 3.10.12 (main, Nov 20 2023, 15:14:05) [GCC 11.4.0] (64-bit runtime)
Python platform: Linux-3.10.0-1160.114.2.el7.x86_64-x86_64-with-glibc2.35
Is CUDA available: True
CUDA runtime version: 12.2.140
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration:
GPU 0: NVIDIA A100-SXM4-80GB
GPU 1: NVIDIA A100-SXM4-80GB
GPU 2: NVIDIA A100-SXM4-80GB
GPU 3: NVIDIA A100-SXM4-80GB
GPU 4: NVIDIA A100-SXM4-80GB
GPU 5: NVIDIA A100-SXM4-80GB
GPU 6: NVIDIA A100-SXM4-80GB
GPU 7: NVIDIA A100-SXM4-80GB

Nvidia driver version: 535.129.03

Versions of relevant libraries:
[pip3] mypy-extensions==1.0.0
[pip3] mypy-protobuf==3.6.0
[pip3] numpy==1.22.0
[pip3] nvidia-nccl-cu12==2.20.5
[pip3] onnx==1.14.0
[pip3] pytorch-quantization==2.1.2
[pip3] torch==2.3.1
[pip3] torch-tensorrt==0.0.0
[pip3] torchdata==0.7.0a0
[pip3] torchtext==0.16.0a0
[pip3] torchvision==0.18.1
[pip3] transformers==4.43.2
[pip3] triton==2.3.1
[pip3] vllm-nccl-cu12==2.18.1.0.4.0
[conda] Could not collect
ROCM Version: Could not collect
Neuron SDK Version: N/A
vLLM Version: v0.5.3
vLLM Build Flags:
CUDA Archs: 5.2 6.0 6.1 7.0 7.2 7.5 8.0 8.6 8.7 9.0+PTX; ROCm: Disabled; Neuron: Disabled
```

### 🐛 Describe the bug

When enabling prefix caching in BlockSpaceManagerV2, the system experiences a slowdown, particularly with an increased time to first token(ttft). The following are the results of my experiments with the Llama3 8B and 70B models.
Initially, I suspected the tensor parallel logic, but since the 8B model, which I ran on a single GPU, also showed a slowdown, I believe the cause lies elsewhere.
(By the way, prefix caching in BlockSpaceManagerV1 showed a 1.3-1.5x speedup in the time to first token metric with the same experiment settings.)

### Experiment Settings
- I tested with [llmperf](https://github.com/ray-project/llmperf), which can vary the number of clients, affecting the maximum batch size in the vLLM worker.
- The total input length was set to 1536, with a prefix prompt length of 512.

### Llama 70B result
Model | Num GPUs | Num Clients | Block Manager | Prefix Caching | TTFT (mean)
-- | -- | -- | -- | -- | --
Llama3-70B | 4 | 16 | v2 | X | 2415 ms
Llama3-70B | 4 | 32 | v2 | X | 4179 ms
Llama3-70B | 4 | 64 | v2 | X | 7670 ms
Llama3-70B | 4 | 128 | v2 | X | 12883 ms
Llama3-70B | 4 | 16 | v2 | O | 2755 ms
Llama3-70B | 4 | 32 | v2 | O | 4652 ms
Llama3-70B | 4 | 64 | v2 | O | 14344 ms
Llama3-70B | 4 | 128 | v2 | O | 25500 ms

### Llama 8B result
Model | Num GPUs | Num Clients | Block Manager | Prefix Caching | TTFT (mean)
-- | -- | -- | -- | -- | --
Llama3-8B | 1 | 16 | v2 | X | 841 ms
Llama3-8B | 1 | 32 | v2 | X | 1441 ms
Llama3-8B | 1 | 64 | v2 | X | 2619 ms
Llama3-8B | 1 | 128 | v2 | X | 4729 ms
Llama3-8B | 1 | 16 | v2 | O | 1962 ms
Llama3-8B | 1 | 32 | v2 | O | 8382 ms
Llama3-8B | 1 | 64 | v2 | O | 12665 ms
Llama3-8B | 1 | 128 | v2 | O | 22439 ms

I meet the problem when enabling prefix caching in BlockSpaceManagerV2 ...