Skip to content

[https://nvbugs/6223556][fix] Correct cached_tokens accounting in dis-agg gen-first overlap path #14910

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
tensorrt-cicd wants to merge 1 commit into from
Closed

[https://nvbugs/6223556][fix] Correct cached_tokens accounting in dis-agg gen-first overlap path #14910

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
11 changes: 10 additions & 1 deletion tensorrt_llm/serve/openai_disagg_service.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
8 changes: 4 additions & 4 deletions tests/integration/defs/disaggregated/test_configs/disagg_config_overlap_gen_first.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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
Expand Down
8 changes: 4 additions & 4 deletions tests/integration/defs/disaggregated/test_configs/disagg_config_overlap_gen_first_pp4.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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
Expand Down
2 changes: 0 additions & 2 deletions tests/integration/test_lists/waives.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
Loading