[AMD][bugfix] Place TBO cuda-graph num_token_non_padded buffer on model devices (#28337)

2026-06-20 18:06:22 -07:00
parent 5b3eeaf504
commit 2552b860a3
2 changed files with 85 additions and 1 deletions
@@ -316,7 +316,9 @@ def compute_split_indices_for_cuda_graph_replay(

 class TboCudaGraphRunnerPlugin:
    def __init__(self):
-        self._tbo_children_num_token_non_padded = torch.zeros((2,), dtype=torch.int32)
+        self._tbo_children_num_token_non_padded = torch.zeros(
+            (2,), dtype=torch.int32, device=get_global_server_args().device
+        )

    def capture_one_batch_size(self, batch: ForwardBatch, num_tokens: int):
        if not is_tbo_enabled():