Skip to content

remove spurious cpu->gpu and gpu->cpu transfers #123

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

Open
wants to merge 17 commits into
base: feat/ad-2025-07-22
Choose a base branch
from
Open

remove spurious cpu->gpu and gpu->cpu transfers #123

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
17 commits
Select commit Hold shift + click to select a range
42a507e
avoid gpu->cpu transfer when using overlap scheduler
suyoggupta Jul 27, 2025
a420975
prealloc
suyoggupta Jul 28, 2025
4b6b5b2
optimize prepare input
suyoggupta Jul 29, 2025
67c0957
more changes
suyoggupta Jul 29, 2025
4a6f1c2
remove more syncs
suyoggupta Jul 29, 2025
e3d4415
clean up
suyoggupta Aug 1, 2025
f2ba5da
clean up position_id handling
suyoggupta Aug 1, 2025
c309859
revert spurious changes
suyoggupta Aug 1, 2025
dfa1d1b
revert spurious print
suyoggupta Aug 1, 2025
63fe801
Update export_to_gm.py
suyoggupta Aug 1, 2025
e6b5a08
revert spurious print
suyoggupta Aug 1, 2025
8fef393
revert nvtx debug
suyoggupta Aug 1, 2025
e4aa045
revert nvtx debug
suyoggupta Aug 1, 2025
6483f0f
update_position_ids is the bottleneck - add nvtx markers and small li…
galagam Aug 3, 2025
e2ad155
avoid double call to update_position_ids
galagam Aug 3, 2025
7d5a102
optimize prepare_list for generation mode
galagam Aug 4, 2025
b4df6b4
bugfix in _update_position_ids optimized code
galagam Aug 4, 2025
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
12 changes: 9 additions & 3 deletions tensorrt_llm/_torch/auto_deploy/custom_ops/attention_interface.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -514,7 +514,12 @@ def unnest_sequences(self, t_nested: torch.Tensor) -> List[torch.Tensor]:
return list(torch.split(t_squeezed, self.sequence_lengths))

@nvtx_range("ad_update_pos")
def update_pos(self, seq_len: Union[torch.Tensor, List[int], int], reset: bool = False) -> None:
def update_pos(
self,
seq_len: Union[torch.Tensor, List[int], int],
reset: bool = False,
update_position_ids: bool = True,
) -> None:
"""Update the starting position for each sequence in the cache.

If ``reset=True`, ``input_pos`` will be reset to zero before updating.
Expand All @@ -528,8 +533,9 @@ def update_pos(self, seq_len: Union[torch.Tensor, List[int], int], reset: bool =
else:
self.input_pos_host[:bs] += seq_len.to(self.device)
Copy link

Copilot AI Aug 1, 2025

Choose a reason for hiding this comment

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

This operation moves seq_len to device before adding to host tensor, which defeats the purpose of keeping calculations on host. Consider converting seq_len to CPU first: self.input_pos_host[:bs] += seq_len.cpu()

Suggested change
self.input_pos_host[:bs] += seq_len.to(self.device)
self.input_pos_host[:bs] += seq_len.cpu()

Copilot uses AI. Check for mistakes.

# update position_ids
self._update_position_ids()
# In ad_executor context, this is done later in nest_sequences, so no need to do it here
if update_position_ids:
self._update_position_ids()
self.input_pos[:bs].copy_(self.input_pos_host[:bs], non_blocking=True)

@nvtx_range("ad_assign_cache_loc")
Expand Down
3 changes: 2 additions & 1 deletion tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -209,7 +209,8 @@ def _prepare_inputs(

# update the sequence info object now
si = self.cache_seq_interface.info
si.update_pos(input_pos, reset=True)
# skip calling _update_position_ids() here, as it will be called in nest_sequences
si.update_pos(input_pos, reset=True, update_position_ids=False)
Copy link
Author

@suyoggupta suyoggupta Aug 4, 2025

Choose a reason for hiding this comment

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

maybe it's better to not call update_pos here at all and introduce a different method that does what update_pos (update_position_ids=False) does? As-is, it is bit confusing to call update_pos without updating positions.
@galagam

Copy link

@galagam galagam Aug 4, 2025

Choose a reason for hiding this comment

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

Updating the position ids requires both the input positions and the sequence lengths, so it makes sense to update it whenever either is updated, but it's a bit wasteful.
A possible alternative would be to require the user to call it explicitly.
That is

si.update_input_pos()  #  rename update_pos
si.nest_sequences()
si.update_position_ids()

In any case, due to my recent changes, run time of update_position_ids decreased by x30, so it's not as critical to add this specific optimization as I initially believed. I'll run a more exhaustive check and consider to keep this optimization out of this PR for code simplicity.
@suyoggupta

si.assign_cache_loc(page_assignments)
si.nest_sequences(input_ids)

Expand Down