Repository navigation
[ModelRunner V2] Don't pin reused flashinfer tensors - #32799
Conversation
Signed-off-by: Nick Hill <nickhill123@gmail.com>
There was a problem hiding this comment.
Code Review
The pull request correctly addresses a potential race condition in ModelRunnerV2 by conditionally disabling pin_memory for FlashInfer tensors. The added comment clearly explains the rationale behind this change, which is crucial for preventing asynchronous copy issues between CPU and GPU buffers. The implementation aligns with the stated objective of avoiding pinning when ModelRunnerV2 is active.
| self.pin_memory = ( | ||
| not envs.VLLM_USE_V2_MODEL_RUNNER and is_pin_memory_available() |
There was a problem hiding this comment.
The decision to disable pin_memory is critical for ModelRunnerV2 due to its lack of explicit synchronization, as noted in the comment. However, relying on the global envs.VLLM_USE_V2_MODEL_RUNNER environment variable introduces a tight coupling between FlashInferMetadataBuilder and the global state. This approach can make the system harder to reason about, test, and debug, especially if the activation of ModelRunnerV2 is not perfectly synchronized with the environment variable's state. For improved modularity and explicitness, consider passing a use_v2_model_runner boolean parameter directly to the FlashInferMetadataBuilder constructor. This would ensure that the pin_memory behavior is directly controlled by the active ModelRunner instance rather than an implicit global flag.
Since we do not have explicit synchronization in ModelRunnerV2, we do not pin reused CPU buffers to avoid a race condition between step N async copies to GPU and step N+1 buffer updates.