Skip to content

skip HPU graphs for long prefills #780

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

Merged
michalkuligowski merged 2 commits into from
Jan 8, 2026
Merged

skip HPU graphs for long prefills #780

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
30c9ac9
skip HPU graphs for long prefills
yangulei Jan 7, 2026
ad91776
use max_num_batched_tokens instead of fixed 8192
yangulei Jan 7, 2026
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
10 changes: 6 additions & 4 deletions vllm_gaudi/v1/worker/hpu_model_runner.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -783,14 +783,14 @@ def __init__(
self.use_hpu_graph = not self.model_config.enforce_eager
self.max_batch_size = self.scheduler_config.max_num_seqs
self.max_num_seqs = self.scheduler_config.max_num_seqs
Copy link

Copilot AI Jan 7, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The variable is renamed from max_cudagraph_capture_size to max_graph_capture_tokens, but the source still references max_cudagraph_capture_size. This creates inconsistent terminology between HPU-specific naming and the CUDA-derived config field name. Consider either keeping the original name for clarity or adding a comment explaining the name change from size to tokens.

Suggested change
self.max_num_seqs = self.scheduler_config.max_num_seqs
self.max_num_seqs = self.scheduler_config.max_num_seqs
# NOTE: `max_graph_capture_tokens` is derived from the CUDA-derived
# config field `max_cudagraph_capture_size`. The HPU path uses
# "tokens" terminology, but the underlying config name is kept for
# compatibility.

Copilot uses AI. Check for mistakes.
self.max_cudagraph_capture_size = self.vllm_config.compilation_config.max_cudagraph_capture_size
if prompt_profile_cfg:
self.max_prefill_batch_size = prompt_profile_cfg[0]
else:
self.max_prefill_batch_size = with_default(get_config().VLLM_PROMPT_BS_BUCKET_MAX, 1)
self.seen_configs: set = set()
self.max_num_batched_tokens = \
self.scheduler_config.max_num_batched_tokens
self.max_num_batched_tokens = self.scheduler_config.max_num_batched_tokens
self.max_graph_capture_tokens = self.vllm_config.compilation_config.max_cudagraph_capture_size if \
self.vllm_config.compilation_config.max_cudagraph_capture_size is not None else self.max_num_batched_tokens
self.use_prefix_caching = (self.vllm_config.cache_config.enable_prefix_caching)
self.bucketing_manager = HPUBucketingManager()
max_num_prefill_seqs = self.max_num_seqs if self.use_merged_prefill \
Expand Down Expand Up @@ -2575,7 +2575,9 @@ def _execute_model_generic(self,
additional_kwargs = {}
if htorch.utils.internal.is_lazy():
use_graphs = self._use_graphs()
if self.max_cudagraph_capture_size is not None and batch_size * seq_len > self.max_cudagraph_capture_size:
# skip HPU graphs for long prefills
Copy link

Copilot AI Jan 7, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The formula batch_size * (seq_len + num_blocks * self.block_size) calculates total tokens including context, but the logic assumes num_blocks * self.block_size represents context tokens. This assumption should be documented in a comment to clarify what num_blocks represents and why this calculation correctly accounts for context tokens.

Suggested change
# skip HPU graphs for long prefills
# skip HPU graphs for long prefills
# NOTE: num_blocks is the number of KV/cache blocks per sequence and
# self.block_size is tokens per block, so num_blocks * self.block_size
# represents the number of cached context tokens per sequence. Adding
# seq_len gives total tokens per sequence (context + current), which
# is then scaled by batch_size to compare against max_graph_capture_tokens.

Copilot uses AI. Check for mistakes.
if seq_len > 1 and \
batch_size * (seq_len + num_blocks * self.block_size) > self.max_graph_capture_tokens:
use_graphs = False
additional_kwargs.update({"bypass_hpu_graphs": not use_graphs})
else:
Expand Down