tpu-inference 0.12.0.dev20251222__py3-none-any.whl → 0.12.0.dev20251224__py3-none-any.whl

tests/layers/vllm/test_unquantized.py

@@ -118,12 +118,16 @@ def test_loading_model(model, mesh):
 
 @pytest.mark.parametrize("model", MODELS)
 @pytest.mark.parametrize("bias", [False, True])
-@pytest.mark.parametrize("mesh", [
-    test_utils.get_spmd_mesh(1),
-    test_utils.get_spmd_mesh(jax.local_device_count())
-])
+@pytest.mark.parametrize("num_devices", [1, jax.local_device_count()])
 @pytest.mark.parametrize("enable_sp", [False, True])
-def test_row_parallel_linear(model, bias, mesh, enable_sp):
+@pytest.mark.parametrize("enable_attn_dp", [False, True])
+def test_row_parallel_linear(model, bias, num_devices, enable_sp,
+                             enable_attn_dp):
+    # Skip if enable_attn_dp is True but we don't have enough devices
+    if enable_attn_dp and num_devices < 2:
+        pytest.skip("enable_attn_dp requires at least 2 devices")
+
+    mesh = test_utils.get_spmd_mesh(num_devices, enable_attn_dp)
     dtype = torch.bfloat16
 
     engine_args = EngineArgs(
@@ -191,12 +195,16 @@ def test_row_parallel_linear(model, bias, mesh, enable_sp):
 
 @pytest.mark.parametrize("model", MODELS)
 @pytest.mark.parametrize("bias", [False, True])
-@pytest.mark.parametrize("mesh", [
-    test_utils.get_spmd_mesh(1),
-    test_utils.get_spmd_mesh(jax.local_device_count())
-])
+@pytest.mark.parametrize("num_devices", [1, jax.local_device_count()])
 @pytest.mark.parametrize("enable_sp", [False, True])
-def test_column_parallel_linear(model, bias, mesh, enable_sp):
+@pytest.mark.parametrize("enable_attn_dp", [False, True])
+def test_column_parallel_linear(model, bias, num_devices, enable_sp,
+                                enable_attn_dp):
+    # Skip if enable_attn_dp is True but we don't have enough devices
+    if enable_attn_dp and num_devices < 2:
+        pytest.skip("enable_attn_dp requires at least 2 devices")
+
+    mesh = test_utils.get_spmd_mesh(num_devices, enable_attn_dp)
     dtype = torch.bfloat16
 
     engine_args = EngineArgs(
@@ -263,13 +271,17 @@ def test_column_parallel_linear(model, bias, mesh, enable_sp):
 
 @pytest.mark.parametrize("model", MODELS)
 @pytest.mark.parametrize("bias", [False, True])
-@pytest.mark.parametrize("mesh", [
-    test_utils.get_spmd_mesh(1),
-    test_utils.get_spmd_mesh(jax.local_device_count())
-])
+@pytest.mark.parametrize("num_devices", [1, jax.local_device_count()])
 @pytest.mark.parametrize("enable_sp", [False, True])
 @pytest.mark.parametrize("fuse_matmuls", [False, True])
-def test_qkv_parallel_linear(model, bias, mesh, enable_sp, fuse_matmuls):
+@pytest.mark.parametrize("enable_attn_dp", [False, True])
+def test_qkv_parallel_linear(model, bias, num_devices, enable_sp, fuse_matmuls,
+                             enable_attn_dp):
+    # Skip if enable_attn_dp is True but we don't have enough devices
+    if enable_attn_dp and num_devices < 2:
+        pytest.skip("enable_attn_dp requires at least 2 devices")
+
+    mesh = test_utils.get_spmd_mesh(num_devices, enable_attn_dp)
     dtype = torch.bfloat16
 
     engine_args = EngineArgs(
@@ -341,14 +353,17 @@ def test_qkv_parallel_linear(model, bias, mesh, enable_sp, fuse_matmuls):
 
 @pytest.mark.parametrize("model", MODELS)
 @pytest.mark.parametrize("bias", [False, True])
-@pytest.mark.parametrize("mesh", [
-    test_utils.get_spmd_mesh(1),
-    test_utils.get_spmd_mesh(jax.local_device_count())
-])
+@pytest.mark.parametrize("num_devices", [1, jax.local_device_count()])
 @pytest.mark.parametrize("fuse_matmuls", [False, True])
 @pytest.mark.parametrize("enable_sp", [False, True])
-def test_merged_column_parallel_linear(model, bias, mesh, fuse_matmuls,
-                                       enable_sp):
+@pytest.mark.parametrize("enable_attn_dp", [False, True])
+def test_merged_column_parallel_linear(model, bias, num_devices, fuse_matmuls,
+                                       enable_sp, enable_attn_dp):
+    # Skip if enable_attn_dp is True but we don't have enough devices
+    if enable_attn_dp and num_devices < 2:
+        pytest.skip("enable_attn_dp requires at least 2 devices")
+
+    mesh = test_utils.get_spmd_mesh(num_devices, enable_attn_dp)
     dtype = torch.bfloat16
 
     engine_args = EngineArgs(
@@ -418,10 +433,7 @@ def test_merged_column_parallel_linear(model, bias, mesh, fuse_matmuls,
 
 
 @pytest.mark.parametrize("use_ep", [True, False])
-@pytest.mark.parametrize("mesh", [
-    test_utils.get_spmd_mesh(1),
-    test_utils.get_spmd_mesh(jax.local_device_count())
-])
+@pytest.mark.parametrize("num_devices", [1, jax.local_device_count()])
 @pytest.mark.parametrize("num_tokens", [8])
 @pytest.mark.parametrize("intermediate_size", [1024, 2048])
 @pytest.mark.parametrize("hidden_size", [128, 512])
@@ -429,8 +441,15 @@ def test_merged_column_parallel_linear(model, bias, mesh, fuse_matmuls,
 @pytest.mark.parametrize("topk", [2])
 @pytest.mark.parametrize("has_bias", [False, True])
 @pytest.mark.parametrize("activation", ["silu", "swigluoai"])
-def test_fused_moe(use_ep, mesh, num_tokens, intermediate_size, hidden_size,
-                   num_experts, topk, has_bias, activation):
+@pytest.mark.parametrize("enable_attn_dp", [False, True])
+def test_fused_moe(use_ep, num_devices, num_tokens, intermediate_size,
+                   hidden_size, num_experts, topk, has_bias, activation,
+                   enable_attn_dp):
+    # Skip if enable_attn_dp is True but we don't have enough devices
+    if enable_attn_dp and num_devices < 2:
+        pytest.skip("enable_attn_dp requires at least 2 devices")
+
+    mesh = test_utils.get_spmd_mesh(num_devices, enable_attn_dp)
 
     torch.manual_seed(42)
     dtype = torch.bfloat16
@@ -502,16 +521,27 @@ def test_fused_moe(use_ep, mesh, num_tokens, intermediate_size, hidden_size,
                                    rtol=1e-1)
 
 
-@pytest.mark.parametrize("mesh",
-                         [test_utils.get_spmd_mesh(jax.local_device_count())])
+@pytest.mark.parametrize("num_devices", [jax.local_device_count()])
 @pytest.mark.parametrize("num_tokens", [128, 512])
 @pytest.mark.parametrize("intermediate_size", [512])
 @pytest.mark.parametrize("hidden_size", [512])
 @pytest.mark.parametrize("num_experts", [32])
 @pytest.mark.parametrize("topk", [8])
 @pytest.mark.parametrize("has_bias", [False, True])
-def test_fused_moe_use_kernel(mesh, num_tokens, intermediate_size, hidden_size,
-                              num_experts, topk, has_bias):
+@pytest.mark.parametrize("enable_attn_dp", [False, True])
+def test_fused_moe_use_kernel(num_devices, num_tokens, intermediate_size,
+                              hidden_size, num_experts, topk, has_bias,
+                              enable_attn_dp):
+    # Skip if enable_attn_dp is True but we don't have enough devices
+    if enable_attn_dp and num_devices < 2:
+        pytest.skip("enable_attn_dp requires at least 2 devices")
+
+    # Skip attn_dp tests for fused_moe_use_kernel since the kernel only supports 2D mesh
+    if enable_attn_dp:
+        pytest.skip(
+            "fused_moe kernel does not support attn_dp (requires 2D mesh)")
+
+    mesh = test_utils.get_spmd_mesh(num_devices, enable_attn_dp)
 
     # TODO(Qiliang Cui): Remove when issue is resolved.
     if not jtu.is_device_tpu_at_least(version=7):

tests/layers/vllm/utils.py

@@ -16,12 +16,27 @@ import jax
 import torch
 import torch.nn.functional as F
 
+from tpu_inference.layers.common.sharding import (MESH_AXIS_NAMES,
+                                                  MESH_AXIS_NAMES_2D)
 
-def get_spmd_mesh(num_devices: int = 1):
-    axis_names = ("data", "model")
+
+def get_spmd_mesh(num_devices: int = 1, enable_attn_dp: bool = False):
     devices = sorted(jax.devices(), key=lambda d: d.id)[0:num_devices]
-    mesh_shape = (1, len(devices))
-    return jax.make_mesh(mesh_shape, axis_names, devices=devices)
+
+    if enable_attn_dp:
+        if num_devices < 2:
+            raise ValueError(
+                f"enable_attn_dp requires at least 2 devices, got {num_devices}"
+            )
+        axis_names = MESH_AXIS_NAMES
+        attn_dp_size = 2
+        model_size = num_devices // attn_dp_size
+        mesh_shape = (1, attn_dp_size, 1, model_size)
+        return jax.make_mesh(mesh_shape, axis_names, devices=devices)
+    else:
+        axis_names = MESH_AXIS_NAMES_2D
+        mesh_shape = (1, len(devices))
+        return jax.make_mesh(mesh_shape, axis_names, devices=devices)
 
 
 def find_all_layer_type(module: torch.nn.Module, layer_type: torch.nn.Module):

tests/models/common/test_model_loader.py

@@ -218,9 +218,9 @@ def test_register_model_vllm_wrapper_methods():
     with pytest.raises(NotImplementedError, match="JAX model"):
         instance.forward(input_ids=None, positions=None)
 
-    # `get_input_embeddings` should be unimplemented.
+    # `embed_input_ids` should be unimplemented.
     with pytest.raises(NotImplementedError, match="JAX model"):
-        instance.get_input_embeddings(input_ids=None, positions=None)
+        instance.embed_input_ids(input_ids=None, positions=None)
 
     # `load_weights` should be a no-op that returns None.
     assert instance.load_weights() is None

tests/models/jax/test_qwen2_5_vl.py

@@ -491,8 +491,7 @@ class TestQwen2_5_VLForConditionalGeneration:
         assert embeddings[1].shape == (tokens_per_image, vc.out_hidden_size)
         assert model.visual.call_count == 2
 
-    def test_get_multimodal_embeddings(
-            self, model: Qwen2_5_VLForConditionalGeneration):
+    def test_embed_multimodal(self, model: Qwen2_5_VLForConditionalGeneration):
         grid_thw = ((2, 28, 28), )
         vc = model.config.vision_config
         patch_dim = vc.in_channels * vc.temporal_patch_size * vc.patch_size * vc.patch_size
@@ -503,20 +502,20 @@ class TestQwen2_5_VLForConditionalGeneration:
         with patch.object(model,
                           '_process_image_input',
                           return_value=(mock_vision_output, )) as mock_process:
-            mm_embeds = model.get_multimodal_embeddings(
-                grid_thw, pixel_values=pixel_values)
+            mm_embeds = model.embed_multimodal(grid_thw,
+                                               pixel_values=pixel_values)
             mock_process.assert_called_once()
             assert isinstance(mm_embeds, tuple)
             assert len(mm_embeds) == 1
             assert mm_embeds[0].shape == (tokens_per_image, vc.out_hidden_size)
 
-        mm_embeds_none = model.get_multimodal_embeddings(grid_thw)
+        mm_embeds_none = model.embed_multimodal(grid_thw)
         assert len(mm_embeds_none) == 0
 
     @patch('tpu_inference.models.jax.qwen2_5_vl.merge_multimodal_embeddings')
-    def test_get_input_embeddings(self, mock_merge_embeddings: MagicMock,
-                                  model: Qwen2_5_VLForConditionalGeneration,
-                                  rng: PRNGKey):
+    def test_embed_input_ids(self, mock_merge_embeddings: MagicMock,
+                             model: Qwen2_5_VLForConditionalGeneration,
+                             rng: PRNGKey):
         input_ids = jax.random.randint(rng, (1, 10), 0,
                                        model.config.vocab_size)
         mock_text_embeds = jnp.ones((1, 10, model.config.hidden_size))
@@ -524,12 +523,12 @@ class TestQwen2_5_VLForConditionalGeneration:
         model.language_model.model.embed = MagicMock(
             return_value=mock_text_embeds)
 
-        embeds = model.get_input_embeddings(input_ids, None)
+        embeds = model.embed_input_ids(input_ids, None)
         np.testing.assert_array_equal(embeds, mock_text_embeds)
         mock_merge_embeddings.assert_not_called()
 
         empty_mm = jnp.ones((0, model.config.hidden_size), )
-        embeds_empty_mm = model.get_input_embeddings(input_ids, empty_mm)
+        embeds_empty_mm = model.embed_input_ids(input_ids, empty_mm)
         np.testing.assert_array_equal(embeds_empty_mm, mock_text_embeds)
         mock_merge_embeddings.assert_not_called()
 
@@ -537,7 +536,7 @@ class TestQwen2_5_VLForConditionalGeneration:
         mock_merged = jnp.ones((1, 15, model.config.hidden_size))
         mock_merge_embeddings.return_value = mock_merged
 
-        embeds_mm = model.get_input_embeddings(input_ids, mm_embeds)
+        embeds_mm = model.embed_input_ids(input_ids, mm_embeds)
         np.testing.assert_array_equal(embeds_mm, mock_merged)
         mock_merge_embeddings.assert_called_once_with(
             input_ids, mock_text_embeds, mm_embeds,

tests/runner/test_multimodal_manager.py

@@ -88,7 +88,7 @@ class TestMultiModalManager:
         # 1. ===== Setup =====
         self.runner.is_multimodal_model = True
         self.mock_get_mm_embed_fn = MagicMock()
-        self.runner.get_multimodal_embeddings_fn = self.mock_get_mm_embed_fn
+        self.runner.embed_multimodal_fn = self.mock_get_mm_embed_fn
 
         self.runner.state = MagicMock()
         # Mock scheduler output
@@ -139,7 +139,7 @@ class TestMultiModalManager:
         np.testing.assert_array_equal(np.asarray(cached_embedding),
                                       np.asarray(dummy_embedding))
 
-        # Check if get_multimodal_embeddings_fn was called with correct args
+        # Check if embed_multimodal_fn was called with correct args
         self.mock_get_mm_embed_fn.assert_called_once()
         call_args = self.mock_get_mm_embed_fn.call_args
 
@@ -169,7 +169,7 @@ class TestMultiModalManager:
         # 1. ===== Setup =====
         self.runner.is_multimodal_model = True
         self.mock_get_mm_embed_fn = MagicMock()
-        self.runner.get_multimodal_embeddings_fn = self.mock_get_mm_embed_fn
+        self.runner.embed_multimodal_fn = self.mock_get_mm_embed_fn
 
         self.runner.state = MagicMock()
         # Mock scheduler output for two requests

tests/runner/test_tpu_runner.py

@@ -88,7 +88,7 @@ class TestTPUJaxRunner:
 
         # Mock the embedding function
         self.mock_get_input_embed_fn = MagicMock()
-        self.runner.get_input_embeddings_fn = self.mock_get_input_embed_fn
+        self.runner.embed_input_ids_fn = self.mock_get_input_embed_fn
         self.mock_get_input_embed_fn.return_value = dummy_final_embeds
         self.runner.state = MagicMock()
 
@@ -116,6 +116,65 @@ class TestTPUJaxRunner:
                                       np.asarray(dummy_input_ids))
         self.mock_get_input_embed_fn.assert_not_called()
 
+    @patch('tpu_inference.runner.tpu_runner.TPUSupportedSamplingMetadata')
+    def test_prepare_inputs_hybrid_kvcache(self, mock_sampling_metadata):
+        # create hybrid kv cache config
+        # 20 layers, 10 full attn + 10 sw attn
+        self._create_mock_hybrid_kv_cache_config()
+
+        # Mock scheduler output.
+        scheduler_output = MagicMock()
+        scheduler_output.total_num_scheduled_tokens = 10
+        scheduler_output.num_scheduled_tokens = {'req1': 10}
+        scheduler_output.scheduled_spec_decode_tokens = {}
+        scheduler_output.grammar_bitmask = None
+
+        # Mock input_batch
+        self.runner.input_batch = MagicMock()
+        self.runner.input_batch.num_reqs = 1
+        self.runner.input_batch.req_ids = ['req1']
+        self.runner.input_batch.req_id_to_index = {'req1': 0}
+        self.runner.input_batch.num_computed_tokens_cpu = np.array([10])
+        self.runner.input_batch.token_ids_cpu = np.random.randint(
+            0, 1000, (8, 64), dtype=np.int32)
+
+        # Mock block tables
+        # there will be 2 block tables since there are 2 kv cache groups
+        mock_block_table = MagicMock()
+        mock_block_table.get_cpu_tensor.return_value = np.zeros(
+            self.runner.block_tables_cpu[0].shape)
+        self.runner.input_batch.block_table = [
+            mock_block_table, mock_block_table
+        ]
+        self.runner.block_tables_cpu = [
+            np.zeros(self.runner.block_tables_cpu[0].shape, dtype=np.int32),
+            np.zeros(self.runner.block_tables_cpu[0].shape, dtype=np.int32)
+        ]
+
+        mock_sampling_instance = MagicMock()
+        mock_sampling_metadata.from_input_batch.return_value = mock_sampling_instance
+
+        output = self.runner._prepare_inputs_non_dp(scheduler_output)
+        assert len(output) == 8
+        input_ids, positions, attention_metadata, sampling_metadata, logits_indices, spec_decode_metadata, logits_indices_selector, padded_num_reqs = output
+        # assert it will create attention metadata for each layer.
+        assert isinstance(attention_metadata, dict)
+        assert len(attention_metadata) == 20
+
+    def _create_mock_hybrid_kv_cache_config(self):
+        mock_kv_cache_config = MagicMock()
+        mock_kv_cache_group1 = MagicMock()
+        mock_kv_cache_group1.layer_names = [f'layer.{i}' for i in range(10)]
+        mock_kv_cache_group2 = MagicMock()
+        mock_kv_cache_group2.layer_names = [
+            f'layer.{i}' for i in range(10, 20)
+        ]
+        mock_kv_cache_config.kv_cache_groups = [
+            mock_kv_cache_group1, mock_kv_cache_group2
+        ]
+        self.runner.kv_cache_config = mock_kv_cache_config
+        self.runner.use_hybrid_kvcache = True
+
 
 class TestTPUJaxRunnerMultimodalModelLoadedForTextOnly:
 
@@ -126,7 +185,7 @@ class TestTPUJaxRunnerMultimodalModelLoadedForTextOnly:
         device_array = np.array(jax.devices()[:1]).reshape(1, 1, 1, -1)
         self.mock_mesh = jax.make_mesh(device_array.shape,
                                        ('data', 'attn_dp', 'expert', 'model'))
-        # Setup the runner with the model_config.is_multimodal_model set to True but get_model returning None for get_multimodal_embeddings_fn and get_input_embeddings_fn.
+        # Setup the runner with the model_config.is_multimodal_model set to True but get_model returning None for embed_multimodal_fn and embed_input_ids_fn.
         with patch('jax.devices', return_value=self.mock_devices), \
              patch('jax.make_mesh', return_value=self.mock_mesh), \
              patch('jax.random.key', return_value=self.mock_rng_key), \
@@ -172,8 +231,8 @@ class TestTPUJaxRunnerMultimodalModelLoadedForTextOnly:
     def _model_get_model(self):
         mock_multimodal_fns = {
             "precompile_vision_encoder_fn": None,
-            "get_multimodal_embeddings_fn": None,
-            "get_input_embeddings_fn": None,
+            "embed_multimodal_fn": None,
+            "embed_input_ids_fn": None,
             "get_mrope_input_positions_fn": None
         }
         return (
@@ -190,13 +249,13 @@ class TestTPUJaxRunnerMultimodalModelLoadedForTextOnly:
         # Precondition: make sure the model_config claims the model supports MM.
         assert self.runner.model_config.is_multimodal_model
 
-        # Precondition: load the model and returns get_multimodal_embeddings_fn as None.
-        assert self.runner.get_multimodal_embeddings_fn is None
+        # Precondition: load the model and returns embed_multimodal_fn as None.
+        assert self.runner.embed_multimodal_fn is None
 
         assert not self.runner.is_multimodal_model
 
-        self.runner.get_input_embeddings_fn = MagicMock()
+        self.runner.embed_input_ids_fn = MagicMock()
         dummy_input_ids = jnp.array([1, 2, 3])
         dummy_mm_embeds = jnp.ones((10, 128))
         _ = self.runner._get_input_ids_embeds(dummy_input_ids, dummy_mm_embeds)
-        self.runner.get_input_embeddings_fn.assert_not_called()
+        self.runner.embed_input_ids_fn.assert_not_called()

tests/runner/test_tpu_runner_dp.py

@@ -97,6 +97,20 @@ class TestTPUJaxRunnerDPInputsLightweight:
         mock_output.grammar_bitmask = None
         return mock_output
 
+    def _create_mock_hybrid_kv_cache_config(self):
+        mock_kv_cache_config = MagicMock()
+        mock_kv_cache_group1 = MagicMock()
+        mock_kv_cache_group1.layer_names = [f'layer.{i}' for i in range(10)]
+        mock_kv_cache_group2 = MagicMock()
+        mock_kv_cache_group2.layer_names = [
+            f'layer.{i}' for i in range(10, 20)
+        ]
+        mock_kv_cache_config.kv_cache_groups = [
+            mock_kv_cache_group1, mock_kv_cache_group2
+        ]
+        self.runner.kv_cache_config = mock_kv_cache_config
+        self.runner.use_hybrid_kvcache = True
+
     @patch('tpu_inference.runner.tpu_runner.NamedSharding')
     @patch('tpu_inference.runner.tpu_runner.runner_utils')
     @patch('tpu_inference.runner.tpu_runner.device_array',
@@ -146,6 +160,58 @@ class TestTPUJaxRunnerDPInputsLightweight:
         with pytest.raises(AssertionError):
             self.runner._prepare_inputs_dp(scheduler_output)
 
+    @patch('tpu_inference.runner.tpu_runner.NamedSharding')
+    @patch('tpu_inference.runner.tpu_runner.runner_utils')
+    @patch('tpu_inference.runner.tpu_runner.device_array',
+           side_effect=lambda mesh, tensors, **kwargs: tensors)
+    @patch('tpu_inference.runner.tpu_runner.TPUSupportedSamplingMetadata')
+    def test_prepare_inputs_dp_hybrid_kvcache(self, mock_sampling_metadata,
+                                              mock_device_array,
+                                              mock_runner_utils,
+                                              mock_named_sharding):
+        """Test basic functionality of _prepare_inputs_dp."""
+        # Mock utility functions
+        mock_runner_utils.get_padded_token_len.return_value = 16
+        mock_sampling_metadata.from_input_batch.return_value = MagicMock()
+        mock_named_sharding.return_value = MagicMock()
+
+        # Create test data - only use req1 and req2 to match num_reqs=2
+        num_scheduled_tokens = {"req1": 5, "req2": 3}
+        assigned_dp_ranks = {"req1": 0, "req2": 1}
+        scheduler_output = self._create_mock_scheduler_output(
+            num_scheduled_tokens, assigned_dp_ranks)
+
+        # Create hybrid kv cache config with 10 full attn layers, 10 sw attn layers
+        self._create_mock_hybrid_kv_cache_config()
+
+        # update input_batch's block_table
+        mock_block_table = MagicMock()
+        mock_block_table.get_cpu_tensor.return_value = np.arange(32).reshape(
+            4, 8)
+        self.runner.input_batch.block_table = [
+            mock_block_table, mock_block_table
+        ]
+
+        # update model runner's block_tables_cpu:
+        self.runner.block_tables_cpu = [
+            np.zeros((8, 8), dtype=np.int32),
+            np.zeros((8, 8), dtype=np.int32)
+        ]
+
+        # Execute the method
+        result = self.runner._prepare_inputs_dp(scheduler_output)
+
+        # Basic assertions
+        assert len(result) == 8
+        input_ids, positions, attention_metadata, sampling_metadata, logits_indices, spec_decode_metadata, logits_indices_selector, padded_num_reqs = result
+
+        # Verify utility functions were called
+        mock_runner_utils.get_padded_token_len.assert_called()
+
+        # Verify there's attention_metadata for each layer
+        assert isinstance(attention_metadata, dict)
+        assert len(attention_metadata) == 20
+
     def test_prepare_dp_input_metadata(self):
         num_scheduled_tokens = {"req1": 10, "req2": 5, "req3": 8, "req4": 3}
         assigned_dp_ranks = {"req1": 0, "req2": 0, "req3": 1, "req4": 1}