[Config] Remove diffusion_batch_size and use max_num_seqs instead. (#6484)

Signed-off-by: LyxWxj <yuxuanlou621@foxmail.com>
Signed-off-by: Yuxuan Rho <104440472+LyxWxj@users.noreply.github.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>

tests/diffusion/batching/test_diffusion_batching.py

@@ -7,7 +7,7 @@ End-to-end test for diffusion batching via AsyncOmni.
 This test fires multiple concurrent ``AsyncOmni.generate()`` calls for a
 diffusion model and validates that every caller receives its correct
 individual result.  When the underlying diffusion stage is configured with
-``batch_size > 1`` (via stage config or ``StageDiffusionClient``), the
+``max_num_seqs > 1``, the
 requests will be batched internally.
 
 Even without explicit batching config this test is useful for verifying
@@ -260,7 +260,7 @@ async def compare_single_vs_parallel(
 ) -> None:
     """Run the same prompts sequentially then in parallel and print a comparison."""
 
-    omni = AsyncOmni(model=model, max_num_seqs=batch_size, diffusion_batch_size=batch_size)
+    omni = AsyncOmni(model=model, max_num_seqs=batch_size)
     try:
         await warmup(omni, WARMUP_PROMPTS)
         single_time = await run_single(omni, prompts)
@@ -288,7 +288,7 @@ async def main(model: str, num_prompts: int, mode: str, batch_size: int = 1) ->
         await compare_single_vs_parallel(model, prompts, batch_size=batch_size)
         return
 
-    omni = AsyncOmni(model=model, max_num_seqs=batch_size, diffusion_batch_size=batch_size)
+    omni = AsyncOmni(model=model, max_num_seqs=batch_size)
     try:
         await warmup(omni, WARMUP_PROMPTS)
 
@@ -389,7 +389,7 @@ def test_diffusion_batching_async_concurrent(model_name: str):
     fired via asyncio.gather. Each request_id must appear in the results."""
 
     async def _inner():
-        omni = AsyncOmni(model=model_name, diffusion_batch_size=1)
+        omni = AsyncOmni(model=model_name, max_num_seqs=1)
         try:
             prompts = TEST_PROMPTS[:4]
             sp = _default_sampling_params()
@@ -430,7 +430,7 @@ def test_diffusion_batching_list_prompt_rejected(model_name: str):
     """
 
     async def _inner():
-        omni = AsyncOmni(model=model_name, diffusion_batch_size=4)
+        omni = AsyncOmni(model=model_name, max_num_seqs=4)
         try:
             prompts = TEST_PROMPTS[:4]
             sp = _default_sampling_params()
@@ -528,7 +528,7 @@ if __name__ == "__main__":
         "--batch-size",
         type=int,
         default=1,
-        help="Scheduler in-flight width (--max-num-seqs) and client diffusion_batch_size",
+        help="Scheduler in-flight width (--max-num-seqs)",
     )
     parser.add_argument(
         "--mode",

tests/diffusion/test_diffusion_streaming_output.py

@@ -441,7 +441,6 @@ class TestPipelineStreamingOutputToEntrypoint:
                     model="test_model",
                     od_config=od_config,
                     metadata=metadata,
-                    batch_size=1,
                 )
 
     @staticmethod
@@ -526,6 +525,7 @@ class TestSupportedPipelines:
         runner = object.__new__(DiffusionModelRunner)
         runner.vllm_config = _make_vllm_config()
         runner.od_config = SimpleNamespace(
+            max_num_seqs=1,
             enable_cpu_offload=False,
             enable_layerwise_offload=False,
             enforce_eager=True,

tests/diffusion/test_inline_stage_diffusion_client.py

@@ -47,7 +47,7 @@ def client(mock_engine):
     with patch.object(InlineStageDiffusionClient, "_enrich_config"):
         od_config = MagicMock(spec=OmniDiffusionConfig)
         od_config.streaming_output = False
-        c = InlineStageDiffusionClient(model="test_model", od_config=od_config, metadata=metadata, batch_size=1)
+        c = InlineStageDiffusionClient(model="test_model", od_config=od_config, metadata=metadata)
         yield c
         c.shutdown()
 
@@ -234,5 +234,4 @@ def test_inline_client_requires_replica_id(mock_engine):
                 model="test_model",
                 od_config=od_config,
                 metadata=metadata,
-                batch_size=1,
             )

tests/diffusion/test_multiproc_engine_concurrency.py

@@ -1171,7 +1171,6 @@ class TestStageDiffusionClientErrorPropagation:
                 metadata,
                 "tcp://req",
                 "tcp://resp",
-                batch_size=1,
             )
 
     @pytest.mark.asyncio

tests/diffusion/test_prompt_update.py

@@ -387,7 +387,6 @@ class TestPromptUpdateIntegration:
                     model="test_model",
                     od_config=od_config,
                     metadata=metadata,
-                    batch_size=1,
                 )
 
     @staticmethod