diff --git a/tensorrt_llm/serve/openai_disagg_service.py b/tensorrt_llm/serve/openai_disagg_service.py index a4c100f6458c..d8870fad0668 100644 --- a/tensorrt_llm/serve/openai_disagg_service.py +++ b/tensorrt_llm/serve/openai_disagg_service.py @@ -23,6 +23,7 @@ MetadataServerConfig, ServerRole, get_global_disagg_request_id, + rewrite_usage_response_from_ctx, ) from tensorrt_llm.logger import logger from tensorrt_llm.serve.cluster_storage import ClusterStorage, WatchEventType @@ -499,4 +500,12 @@ async def _yield_from_queue(): ) ) responses = await asyncio.gather(*tasks) - return responses[-1] + gen_response = responses[-1] + if need_ctx: + # The generation worker treats the whole transferred prompt as + # cached, so its usage accounting over-reports cached_tokens. + # Adopt the context worker's accounting, mirroring the + # context-first path which propagates ctx_usage to the gen + # worker (not possible here since both requests run concurrently). + rewrite_usage_response_from_ctx(gen_response, responses[0].usage) + return gen_response diff --git a/tests/integration/defs/disaggregated/test_configs/disagg_config_overlap_gen_first.yaml b/tests/integration/defs/disaggregated/test_configs/disagg_config_overlap_gen_first.yaml index 72a316ad86bc..50e8f172101b 100644 --- a/tests/integration/defs/disaggregated/test_configs/disagg_config_overlap_gen_first.yaml +++ b/tests/integration/defs/disaggregated/test_configs/disagg_config_overlap_gen_first.yaml @@ -12,9 +12,9 @@ context_servers: tensor_parallel_size: 1 pipeline_parallel_size: 1 kv_cache_config: - enable_block_reuse: False + enable_block_reuse: True free_gpu_memory_fraction: 0.2 - enable_partial_reuse: False + enable_partial_reuse: True cache_transceiver_config: backend: DEFAULT @@ -29,9 +29,9 @@ generation_servers: max_num_tokens: 4096 max_seq_len: 4096 kv_cache_config: - enable_block_reuse: False + enable_block_reuse: True free_gpu_memory_fraction: 0.2 - enable_partial_reuse: False + enable_partial_reuse: True cache_transceiver_config: backend: DEFAULT transceiver_runtime: PYTHON diff --git a/tests/integration/defs/disaggregated/test_configs/disagg_config_overlap_gen_first_pp4.yaml b/tests/integration/defs/disaggregated/test_configs/disagg_config_overlap_gen_first_pp4.yaml index e924722aef1d..4f50e9b57150 100644 --- a/tests/integration/defs/disaggregated/test_configs/disagg_config_overlap_gen_first_pp4.yaml +++ b/tests/integration/defs/disaggregated/test_configs/disagg_config_overlap_gen_first_pp4.yaml @@ -12,9 +12,9 @@ context_servers: tensor_parallel_size: 1 pipeline_parallel_size: 4 kv_cache_config: - enable_block_reuse: False + enable_block_reuse: True free_gpu_memory_fraction: 0.2 - enable_partial_reuse: False + enable_partial_reuse: True cache_transceiver_config: backend: DEFAULT @@ -29,9 +29,9 @@ generation_servers: max_num_tokens: 4096 max_seq_len: 4096 kv_cache_config: - enable_block_reuse: False + enable_block_reuse: True free_gpu_memory_fraction: 0.2 - enable_partial_reuse: False + enable_partial_reuse: True cache_transceiver_config: backend: DEFAULT transceiver_runtime: PYTHON diff --git a/tests/integration/test_lists/waives.txt b/tests/integration/test_lists/waives.txt index a1fd69684188..ba8cad14c36b 100644 --- a/tests/integration/test_lists/waives.txt +++ b/tests/integration/test_lists/waives.txt @@ -178,8 +178,6 @@ disaggregated/test_disaggregated.py::test_disaggregated_deepseek_v3_lite_fp8_tp1 disaggregated/test_disaggregated.py::test_disaggregated_genbs1[TinyLlama-1.1B-Chat-v1.0] SKIP (https://nvbugs/6162322) disaggregated/test_disaggregated.py::test_disaggregated_gpt_oss_120b_harmony[gpt_oss/gpt-oss-120b] SKIP (https://nvbugs/6245317) disaggregated/test_disaggregated.py::test_disaggregated_kv_cache_time_output[TinyLlama-1.1B-Chat-v1.0] SKIP (https://nvbugs/6184906) -disaggregated/test_disaggregated.py::test_disaggregated_overlap_gen_first[ctx_pp1-TinyLlama-1.1B-Chat-v1.0] SKIP (https://nvbugs/6223556) -disaggregated/test_disaggregated.py::test_disaggregated_overlap_gen_first[ctx_pp4-TinyLlama-1.1B-Chat-v1.0] SKIP (https://nvbugs/6223556) disaggregated/test_disaggregated.py::test_disaggregated_single_gpu[TinyLlama-1.1B-Chat-v1.0] SKIP (https://nvbugs/6184906) disaggregated/test_disaggregated_single_gpu.py::test_disaggregated_llama_context_capacity[False-False-DeepSeek-V3-Lite-fp8/fp8] SKIP (https://nvbugs/6266302) disaggregated/test_workers.py::test_workers_conversation_router[TinyLlama-1.1B-Chat-v1.0] SKIP (https://nvbugs/6162322)