Update GPUGenerationModelRunner to clarify KV connector metadata handling during draft model runs

- Changed variable name to better reflect its purpose in managing KV connector metadata.
- Updated comments for improved clarity regarding the handling of speculative configurations.

Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
This commit is contained in:
tzhouam committed 2026-03-13 03:16:33 +00:00
1 parent d911475b22
commit 62c8695dfd
1 file changed
+4 -4
@@ -265,9 +265,9 @@ class GPUGenerationModelRunner(OmniGPUModelRunner):
# Run the model.
# Use persistent buffers for CUDA graphs.
# When spec decode is enabled, defer connector finalization
# (wait_for_save + clear metadata) until after draft model runs.
defer_kv_connector_finalize = self.speculative_config is not None
# When spec decode is enabled, defer clearing KV connector metadata
# until after draft model runs in sample_tokens.
clear_kv_metadata = self.speculative_config is None
with (
set_forward_context(
attn_metadata,
@@ -282,7 +282,7 @@ class GPUGenerationModelRunner(OmniGPUModelRunner):
record_function_or_nullcontext("Forward"),
self.maybe_get_kv_connector_output(
scheduler_output,
defer_finalize=defer_kv_connector_finalize,
clear_metadata=clear_kv_metadata,
) as kv_connector_output,
):
outputs = self._run_generation_model(