Qwen2.5-VL eagle3 infer (#8801)

This commit is contained in:
Lzhang-hub
2025-09-08 11:44:34 +08:00
committed by GitHub
parent 7802586cab
commit 37d83c6e6d
9 changed files with 114 additions and 5 deletions

View File

@@ -317,7 +317,9 @@ class CudaGraphRunner:
(self.max_num_token,), dtype=self._cache_loc_dtype()
)
self.positions = torch.zeros((self.max_num_token,), dtype=torch.int64)
self.mrope_positions = torch.zeros((3, self.max_bs), dtype=torch.int64)
self.mrope_positions = torch.zeros(
(3, self.max_num_token), dtype=torch.int64
)
self.num_token_non_padded = torch.zeros((1,), dtype=torch.int32)
self.tbo_plugin = TboCudaGraphRunnerPlugin()
@@ -532,7 +534,7 @@ class CudaGraphRunner:
encoder_lens = self.encoder_lens[:bs]
else:
encoder_lens = None
mrope_positions = self.mrope_positions[:, :bs]
mrope_positions = self.mrope_positions[:, :num_tokens]
next_token_logits_buffer = self.next_token_logits_buffer[:num_tokens]
self.num_token_non_padded[...] = num_tokens
@@ -751,7 +753,7 @@ class CudaGraphRunner:
if self.is_encoder_decoder:
self.encoder_lens[:raw_bs].copy_(forward_batch.encoder_lens)
if forward_batch.mrope_positions is not None:
self.mrope_positions[:, :raw_bs].copy_(forward_batch.mrope_positions)
self.mrope_positions[:, :raw_num_token].copy_(forward_batch.mrope_positions)
if self.require_gathered_buffer:
self.global_num_tokens_gpu.fill_(bs * self.num_tokens_per_bs)
self.global_num_tokens_for_logprob_gpu.fill_(bs * self.num_tokens_per_bs)