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:
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(
|
||||
|
||||
Reference in new issue
Block a user