A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 12, 09:45 UTC
[Perf] Optimize EAGLE prepare_inputs_padded with triton kernels (#28597)
Signed-off-by: Benjamin Chislett <bchislett@nvidia.com> Signed-off-by: Benjamin Chislett <chislett.ben@gmail.com>
Benjamin Chislett committed
Nov 28, 2025 at 17:25 UTC
1986de137502d0d767cb4c1d3cad23dedbd22397
4 files changed
+198
-107
tests/v1/spec_decode/test_eagle.py
+11
-19
@@ -103,16 +103,23 @@ def test_prepare_next_token_ids():
103
mock_request.num_computed_tokens = 0
104
mock_requests[req_id] = mock_request
105
106
+ # explicitly discard the last request
107
+ discarded_req_mask = torch.tensor(
108
+ [False, False, False, True], dtype=torch.bool, device=device
109
+ )
110
sampled_token_ids = [
111
[0, 1, -1, -1, -1], # 1 accepted, 3 rejected, "1" sampled
112
[0, 1, 2, 3, 4], # all accepted, "4" sampled
113
[-1, -1, -1, -1, -1], # sampling skipped, use backup token "30"
110
- [-1, -1, -1, -1, -1], # this request will be discarded
114
+ [0, 1, 2, -1, -1], # explicitly discarded, sampling should be ignored
115
]
116
sampled_token_ids_tensor = torch.tensor(
117
sampled_token_ids, dtype=torch.int32, device=device
118
)
119
sampled_token_ids_cpu = [[i for i in seq if i != -1] for seq in sampled_token_ids]
120
+ for i in range(len(sampled_token_ids_cpu)):
121
+ if discarded_req_mask[i]:
122
+ sampled_token_ids_cpu[i] = []
123
124
expected_next_token_ids_cpu = [1, 4, 30, 40]
125
expected_next_token_ids_tensor = torch.tensor(
@@ -136,9 +143,6 @@ def test_prepare_next_token_ids():
143
device=device,
144
)
145
139
- discarded_req_indices = torch.tensor([3], dtype=torch.int64, device=device)
140
- num_discarded_reqs = 1
141
-
146
expected_valid_sampled_tokens_count = torch.tensor(
147
[2, 5, 0, 0], dtype=torch.int32, device=device
148
)
@@ -149,8 +153,7 @@ def test_prepare_next_token_ids():
153
sampled_token_ids_tensor,
154
mock_requests,
155
mock_input_batch,
152
- discarded_req_indices,
153
- num_discarded_reqs,
156
+ discarded_req_mask,
157
)
158
)
159
@@ -256,11 +259,6 @@ def test_prepare_inputs_padded():
259
- Request 3: query_len = 3, rejected = 2
260
261
Expected outputs:
259
- token_indices: [0, 1, 2,
260
- 3, 4, 5,
261
- 6, 7, 8]
262
- Reason: Deferred computation should not disturb the original indices.
263
-
262
token_indices_to_sample: [1, 5, 6]
263
Reason: After accounting for rejections, these are the valid token positions
264
from the original indices to sample from.
@@ -268,9 +266,6 @@ def test_prepare_inputs_padded():
266
267
device = torch.device(current_platform.device_type)
268
271
- expected_token_indices = torch.tensor(
272
- [0, 1, 2, 3, 4, 5, 6, 7, 8], dtype=torch.int32, device=device
273
- )
269
expected_token_indices_to_sample = torch.tensor(
270
[1, 5, 6], dtype=torch.int32, device=device
271
)
@@ -305,15 +300,12 @@ def test_prepare_inputs_padded():
300
301
proposer = _create_proposer("eagle", num_speculative_tokens)
302
308
- output_metadata, token_indices, token_indices_to_sample = (
309
- proposer.prepare_inputs_padded(
310
- common_attn_metadata, spec_decode_metadata, valid_sampled_tokens_count
311
- )
303
+ output_metadata, token_indices_to_sample = proposer.prepare_inputs_padded(
304
+ common_attn_metadata, spec_decode_metadata, valid_sampled_tokens_count
305
)
306
307
assert output_metadata.max_query_len == 3
308
assert torch.equal(output_metadata.query_start_loc, expected_query_start_loc)
316
- assert torch.equal(token_indices, expected_token_indices)
309
assert torch.equal(token_indices_to_sample, expected_token_indices_to_sample)
310
311
vllm/v1/spec_decode/eagle.py
+49
-58
@@ -25,6 +25,7 @@ from vllm.model_executor.models.deepseek_v2 import DeepseekV32IndexerCache
25
from vllm.model_executor.models.llama_eagle3 import Eagle3LlamaForCausalLM
26
from vllm.multimodal import MULTIMODAL_REGISTRY
27
from vllm.platforms import current_platform
28
+from vllm.triton_utils import triton
29
from vllm.utils.platform_utils import is_pin_memory_available
30
from vllm.v1.attention.backends.flash_attn import FlashAttentionMetadata
31
from vllm.v1.attention.backends.tree_attn import (
@@ -40,6 +41,10 @@ from vllm.v1.kv_cache_interface import KVCacheConfig
41
from vllm.v1.sample.metadata import SamplingMetadata
42
from vllm.v1.sample.sampler import _SAMPLING_EPS
43
from vllm.v1.spec_decode.metadata import SpecDecodeMetadata
44
+from vllm.v1.spec_decode.utils import (
45
+ eagle_prepare_inputs_padded_kernel,
46
+ eagle_prepare_next_token_padded_kernel,
47
+)
48
from vllm.v1.utils import CpuGpuBuffer
49
from vllm.v1.worker.dp_utils import coordinate_batch_across_dp
50
from vllm.v1.worker.gpu_input_batch import CachedRequestState, InputBatch
@@ -555,20 +560,15 @@ class EagleProposer:
560
sampled_token_ids: torch.Tensor,
561
requests: dict[str, CachedRequestState],
562
gpu_input_batch: InputBatch,
558
- discard_request_indices: torch.Tensor,
559
- num_discarded_requests: int,
563
+ discard_request_mask: torch.Tensor,
564
) -> tuple[torch.Tensor, torch.Tensor]:
565
"""
566
This function is used to prepare the inputs for speculative decoding.
567
It calculates the next token ids and the number of valid sampled tokens
568
for each request, considering the "discarded" requests whose next token
565
- is not sampled and comes from `request.get_token_id()` instead.
566
- It also accounts for the rejected tokens in `sampled_token_ids`.
567
- This function must use device functions to operate on the inputs, and
568
- should not introduce any blocking CPU-GPU synchronization.
569
+ is not sampled and comes from `request.get_token_id()` instead. This is denoted
570
+ the "backup" token id. It also counts rejected tokens via `sampled_token_ids`.
571
"""
570
- # TODO(Ben): Combine this into a custom fused kernel
571
-
572
# Precompute get_token_id for when there is no valid next token
573
num_reqs = gpu_input_batch.num_reqs
574
self.backup_next_token_ids.np[:num_reqs] = np.array(
@@ -577,44 +577,39 @@ class EagleProposer:
577
common_attn_metadata.seq_lens_cpu[i].item()
578
)
579
for i in range(num_reqs)
580
- ]
580
+ ],
581
+ dtype=np.int32,
582
)
583
self.backup_next_token_ids.copy_to_gpu(num_reqs)
584
+ backup_tokens_gpu = self.backup_next_token_ids.gpu
585
584
- # Mask out the sampled tokens indices that should not be sampled.
585
- discard_sampled_tokens_req_indices = discard_request_indices[
586
- :num_discarded_requests
587
- ]
586
+ batch_size, num_tokens = sampled_token_ids.shape
587
+ device = sampled_token_ids.device
588
589
- valid_sampled_token_ids_gpu = sampled_token_ids.clone()
590
- valid_sampled_token_ids_gpu.index_fill_(
591
- 0, discard_sampled_tokens_req_indices, -1
592
- )
589
+ assert discard_request_mask.dtype == torch.bool
590
+ assert backup_tokens_gpu.dtype == torch.int32
591
594
- # Generate a mask for all valid tokens within those requests
595
- valid_mask = (valid_sampled_token_ids_gpu != -1) & (
596
- valid_sampled_token_ids_gpu < gpu_input_batch.vocab_size
592
+ next_token_ids = torch.empty((batch_size,), dtype=torch.int32, device=device)
593
+ valid_sampled_tokens_count = torch.empty(
594
+ (batch_size,), dtype=torch.int32, device=device
595
)
596
599
- # Count the number of valid tokens in each request
600
- valid_sampled_tokens_count = valid_mask.sum(dim=1)
597
+ # Kernel grid: one program per request (row)
598
+ grid = (batch_size,)
599
602
- # Get the rightmost valid index per row
603
- last_valid_indices = valid_sampled_tokens_count - 1
604
- last_valid_indices_safe = torch.clamp(last_valid_indices, min=0)
605
-
606
- # Get last valid token from each row
607
- # (assume undefined state where there is no valid token)
608
- selected_tokens = torch.gather(
609
- valid_sampled_token_ids_gpu, 1, last_valid_indices_safe.unsqueeze(1)
610
- ).squeeze(1)
611
-
612
- # Use last token if valid, pre-computed backup if not
613
- batch_size = valid_sampled_token_ids_gpu.shape[0]
614
- next_token_ids = torch.where(
615
- last_valid_indices != -1,
616
- selected_tokens,
617
- self.backup_next_token_ids.gpu[:batch_size],
600
+ # Find the next power of 2 for block sizes
601
+ BLOCK_SIZE_TOKENS = triton.next_power_of_2(num_tokens)
602
+ eagle_prepare_next_token_padded_kernel[grid](
603
+ sampled_token_ids,
604
+ discard_request_mask,
605
+ backup_tokens_gpu,
606
+ next_token_ids,
607
+ valid_sampled_tokens_count,
608
+ gpu_input_batch.vocab_size,
609
+ num_tokens,
610
+ batch_size,
611
+ sampled_token_ids.stride(0),
612
+ BLOCK_SIZE_TOKENS=BLOCK_SIZE_TOKENS,
613
)
614
615
return next_token_ids, valid_sampled_tokens_count
@@ -624,35 +619,35 @@ class EagleProposer:
619
common_attn_metadata: CommonAttentionMetadata,
620
spec_decode_metadata: SpecDecodeMetadata,
621
valid_sampled_tokens_count: torch.Tensor,
627
- ) -> tuple[CommonAttentionMetadata, torch.Tensor, torch.Tensor]:
622
+ ) -> tuple[CommonAttentionMetadata, torch.Tensor]:
623
"""
624
This function is used to prepare the inputs for speculative decoding
625
It updates the common_attn_metadata for speculative decoding,
626
but does not consider the rejected tokens. Instead, all tokens
627
are included as inputs to the speculator, with the rejected tokens
628
used as padding and filtered out later by `token_indices_to_sample`.
634
- No blocking CPU operations should be introduced in this function.
629
"""
636
- num_draft_tokens_gpu = torch.cat(
637
- [
638
- spec_decode_metadata.cu_num_draft_tokens[0:1],
639
- spec_decode_metadata.cu_num_draft_tokens[1:]
640
- - spec_decode_metadata.cu_num_draft_tokens[:-1],
641
- ]
630
+ num_reqs = common_attn_metadata.num_reqs
631
+ device = valid_sampled_tokens_count.device
632
+
633
+ token_indices_to_sample = torch.empty(
634
+ (num_reqs,), dtype=torch.int32, device=device
635
)
636
644
- num_rejected_tokens_gpu = torch.where(
645
- num_draft_tokens_gpu > 0,
646
- num_draft_tokens_gpu + 1 - valid_sampled_tokens_count,
647
- torch.zeros_like(num_draft_tokens_gpu),
637
+ # Kernel grid: one program per request (row)
638
+ grid = (num_reqs,)
639
+ eagle_prepare_inputs_padded_kernel[grid](
640
+ spec_decode_metadata.cu_num_draft_tokens,
641
+ valid_sampled_tokens_count,
642
+ common_attn_metadata.query_start_loc,
643
+ token_indices_to_sample,
644
+ num_reqs,
645
)
646
647
query_start_loc_cpu = common_attn_metadata.query_start_loc_cpu
651
-
648
new_query_len_per_req = query_start_loc_cpu[1:] - query_start_loc_cpu[:-1]
649
650
total_num_tokens = query_start_loc_cpu[-1].item()
655
- token_indices = self.arange[:total_num_tokens]
651
652
spec_common_attn_metadata = CommonAttentionMetadata(
653
query_start_loc=common_attn_metadata.query_start_loc,
@@ -665,16 +660,12 @@ class EagleProposer:
660
max_query_len=new_query_len_per_req.max().item(),
661
max_seq_len=common_attn_metadata.seq_lens_cpu.max().item(),
662
block_table_tensor=common_attn_metadata.block_table_tensor,
668
- slot_mapping=common_attn_metadata.slot_mapping[token_indices],
663
+ slot_mapping=common_attn_metadata.slot_mapping[:total_num_tokens],
664
causal=True,
665
dcp_local_seq_lens=common_attn_metadata.dcp_local_seq_lens,
666
)
667
673
- token_indices_to_sample = (
674
- common_attn_metadata.query_start_loc[1:] - 1 - num_rejected_tokens_gpu
675
- )
676
-
677
- return spec_common_attn_metadata, token_indices, token_indices_to_sample
668
+ return spec_common_attn_metadata, token_indices_to_sample
669
670
def propose_tree(
671
self,
vllm/v1/spec_decode/utils.py
+105
@@ -1,6 +1,7 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
from vllm.sampling_params import SamplingParams
4
+from vllm.triton_utils import tl, triton
5
6
_SAMPLING_EPS = 1e-5
7
@@ -14,3 +15,107 @@ def is_spec_decode_unsupported(sampling_params: SamplingParams) -> bool:
15
or sampling_params.min_p > _SAMPLING_EPS
16
or sampling_params.logprobs is not None
17
)
18
+
19
+
20
+@triton.jit
21
+def eagle_prepare_inputs_padded_kernel(
22
+ cu_num_draft_tokens_ptr, # [num_reqs]
23
+ valid_sampled_tokens_count_ptr, # [num_reqs]
24
+ query_start_loc_gpu_ptr, # [num_reqs + 1]
25
+ token_indices_to_sample_ptr, # [num_reqs] (output)
26
+ num_reqs, # tl.int32
27
+):
28
+ """
29
+ Fused kernel for Eagle prepare_input_padded. This kernel computes the
30
+ token index to sample for each request, taking into account the number
31
+ of draft tokens and the number of valid sampled tokens (which is one more than
32
+ the number of accepted tokens).
33
+ """
34
+ req_idx = tl.program_id(axis=0)
35
+ if req_idx >= num_reqs:
36
+ return
37
+
38
+ # Calculate num_draft_tokens from cu_num_draft_tokens, which is an inclusive
39
+ # cumulative sum (first entry is the first value, not zero).
40
+ cu_draft_curr = tl.load(cu_num_draft_tokens_ptr + req_idx)
41
+
42
+ num_draft_tokens = 0
43
+ if req_idx == 0:
44
+ num_draft_tokens = cu_draft_curr
45
+ else:
46
+ cu_draft_prev = tl.load(cu_num_draft_tokens_ptr + req_idx - 1)
47
+ num_draft_tokens = cu_draft_curr - cu_draft_prev
48
+
49
+ valid_count = tl.load(valid_sampled_tokens_count_ptr + req_idx)
50
+ num_rejected_tokens = num_draft_tokens + 1 - valid_count
51
+ num_rejected_tokens = tl.where(num_draft_tokens > 0, num_rejected_tokens, 0)
52
+
53
+ # query_start_loc[req_idx + 1] is the start position of the next request,
54
+ # which is one past the last token of this request.
55
+ q_last_tok_idx = tl.load(query_start_loc_gpu_ptr + req_idx + 1) - 1
56
+
57
+ index_to_sample = q_last_tok_idx - num_rejected_tokens
58
+ tl.store(token_indices_to_sample_ptr + req_idx, index_to_sample)
59
+
60
+
61
+@triton.jit
62
+def eagle_prepare_next_token_padded_kernel(
63
+ sampled_token_ids_ptr, # [num_reqs, num_sampled_tokens_per_req]
64
+ discard_request_mask_ptr, # [num_reqs]
65
+ backup_next_token_ids_ptr, # [num_reqs]
66
+ next_token_ids_ptr, # [num_reqs] (output)
67
+ valid_sampled_tokens_count_ptr, # [num_reqs] (output)
68
+ vocab_size, # tl.int32
69
+ num_sampled_tokens_per_req, # tl.int32 (num_spec_tokens + 1)
70
+ num_reqs, # tl.int32
71
+ stride_sampled_token_ids, # tl.int32 (stride for dim 0)
72
+ BLOCK_SIZE_TOKENS: tl.constexpr, # Power-of-2 >= num_sampled_tokens_per_req
73
+):
74
+ """
75
+ Fused kernel for Eagle prepare_next_token_ids_padded. This kernel computes the
76
+ number of valid (1 + accepted) tokens for each request, and the corresponding
77
+ "next" token id to sample from during speculative decoding. This is the
78
+ "last accepted token" from the sampled tokens, or the backup token if no
79
+ tokens were accepted or if the request is marked as discarded.
80
+ """
81
+ req_idx = tl.program_id(axis=0)
82
+ if req_idx >= num_reqs:
83
+ return
84
+
85
+ # Check if this request is discarded.
86
+ is_discarded = tl.load(discard_request_mask_ptr + req_idx)
87
+
88
+ if is_discarded:
89
+ backup_token = tl.load(backup_next_token_ids_ptr + req_idx)
90
+ valid_count = tl.full((), 0, dtype=tl.uint32)
91
+ tl.store(next_token_ids_ptr + req_idx, backup_token)
92
+ tl.store(valid_sampled_tokens_count_ptr + req_idx, valid_count)
93
+ else:
94
+ # Count the number of valid tokens among the sampled tokens.
95
+ token_offs = tl.arange(0, BLOCK_SIZE_TOKENS)
96
+ token_mask = token_offs < num_sampled_tokens_per_req
97
+
98
+ row_ptr = sampled_token_ids_ptr + req_idx * stride_sampled_token_ids
99
+ token_ids = tl.load(row_ptr + token_offs, mask=token_mask, other=-1)
100
+
101
+ # Rejected tokens are -1, valid tokens are in [0, vocab_size)
102
+ is_valid_mask = (token_ids != -1) & (token_ids < vocab_size) & token_mask
103
+ valid_count = tl.sum(is_valid_mask)
104
+
105
+ if valid_count > 0:
106
+ # Guaranteed to be well-defined since
107
+ # valid_count > 0 implies is_valid_mask is not empty
108
+ last_valid_index = tl.max(tl.where(is_valid_mask, token_offs, -1))
109
+
110
+ # Select the token at that index, using a sum trick since
111
+ # we don't want to load again to access token_ids[last_valid_index].
112
+ last_valid_token = tl.sum(
113
+ tl.where(token_offs == last_valid_index, token_ids, 0)
114
+ )
115
+ tl.store(next_token_ids_ptr + req_idx, last_valid_token)
116
+ else:
117
+ # No valid tokens found, use backup token
118
+ backup_token = tl.load(backup_next_token_ids_ptr + req_idx)
119
+ tl.store(next_token_ids_ptr + req_idx, backup_token)
120
+
121
+ tl.store(valid_sampled_tokens_count_ptr + req_idx, valid_count)
vllm/v1/worker/gpu_model_runner.py
+33
-30
@@ -488,11 +488,9 @@ class GPUModelRunner(
488
self.max_num_tokens, self.hidden_size, dtype=self.dtype, numpy=False
489
)
490
self.is_token_ids = self._make_buffer(self.max_num_tokens, dtype=torch.bool)
491
- self.discard_request_indices = self._make_buffer(
492
- self.max_num_reqs, dtype=torch.int64
491
+ self.discard_request_mask = self._make_buffer(
492
+ self.max_num_reqs, dtype=torch.bool
493
)
494
- self.num_discarded_requests = 0
495
-
494
self.num_decode_draft_tokens = self._make_buffer(
495
self.max_num_reqs, dtype=torch.int32
496
)
@@ -1369,16 +1367,12 @@ class GPUModelRunner(
1367
num_tokens = [self.requests[r].num_tokens for r in self.input_batch.req_ids]
1368
num_tokens_np = np.array(num_tokens, dtype=np.int32)
1369
1372
- # Record the index of requests that should not be sampled,
1370
+ # Record which requests should not be sampled,
1371
# so that we could clear the sampled tokens before returning
1374
- discard_requests_mask = self.seq_lens.np[:num_reqs] < num_tokens_np
1375
- discard_request_indices = np.nonzero(discard_requests_mask)[0]
1376
- self.num_discarded_requests = len(discard_request_indices)
1377
- self.discard_request_indices.np[: self.num_discarded_requests] = (
1378
- discard_request_indices
1372
+ self.discard_request_mask.np[:num_reqs] = (
1373
+ self.seq_lens.np[:num_reqs] < num_tokens_np
1374
)
1380
-
1381
- self.discard_request_indices.copy_to_gpu(self.num_discarded_requests)
1375
+ self.discard_request_mask.copy_to_gpu(num_reqs)
1376
1377
# Copy the tensors to the GPU.
1378
self._prepare_input_ids(
@@ -2548,9 +2542,10 @@ class GPUModelRunner(
2542
if envs.VLLM_COMPUTE_NANS_IN_LOGITS:
2543
num_nans_in_logits = self._get_nans_in_logits(logits)
2544
2551
- discard_sampled_tokens_req_indices = self.discard_request_indices.np[
2552
- : self.num_discarded_requests
2553
- ]
2545
+ num_reqs = self.input_batch.num_reqs
2546
+ discard_sampled_tokens_req_indices = np.nonzero(
2547
+ self.discard_request_mask.np[:num_reqs]
2548
+ )[0]
2549
for i in discard_sampled_tokens_req_indices:
2550
gen = self.input_batch.generators.get(int(i))
2551
if gen is not None:
@@ -3131,8 +3126,7 @@ class GPUModelRunner(
3126
sampled_token_ids,
3127
self.requests,
3128
self.input_batch,
3134
- self.discard_request_indices.gpu,
3135
- self.num_discarded_requests,
3129
+ self.discard_request_mask.gpu,
3130
)
3131
)
3132
self._copy_valid_sampled_token_count(
@@ -3335,8 +3329,7 @@ class GPUModelRunner(
3329
sampled_token_ids,
3330
self.requests,
3331
self.input_batch,
3338
- self.discard_request_indices.gpu,
3339
- self.num_discarded_requests,
3332
+ self.discard_request_mask.gpu,
3333
)
3334
)
3335
self._copy_valid_sampled_token_count(
@@ -3363,24 +3356,34 @@ class GPUModelRunner(
3356
sampled_token_ids,
3357
spec_decode_metadata.num_draft_tokens,
3358
)
3359
+ target_token_ids = self.input_ids.gpu[token_indices]
3360
+ target_positions = self._get_positions(token_indices)
3361
+ if self.use_aux_hidden_state_outputs:
3362
+ assert aux_hidden_states is not None
3363
+ target_hidden_states = torch.cat(
3364
+ [h[token_indices] for h in aux_hidden_states], dim=-1
3365
+ )
3366
+ else:
3367
+ target_hidden_states = hidden_states[token_indices]
3368
else:
3367
- common_attn_metadata, token_indices, token_indices_to_sample = (
3369
+ common_attn_metadata, token_indices_to_sample = (
3370
self.drafter.prepare_inputs_padded(
3371
common_attn_metadata,
3372
spec_decode_metadata,
3373
valid_sampled_tokens_count,
3374
)
3375
)
3374
-
3375
- target_token_ids = self.input_ids.gpu[token_indices]
3376
- target_positions = self._get_positions(token_indices)
3377
- if self.use_aux_hidden_state_outputs:
3378
- assert aux_hidden_states is not None
3379
- target_hidden_states = torch.cat(
3380
- [h[token_indices] for h in aux_hidden_states], dim=-1
3381
- )
3382
- else:
3383
- target_hidden_states = hidden_states[token_indices]
3376
+ total_num_tokens = common_attn_metadata.num_actual_tokens
3377
+ # When padding the batch, token_indices is just a range
3378
+ target_token_ids = self.input_ids.gpu[:total_num_tokens]
3379
+ target_positions = self._get_positions(total_num_tokens)
3380
+ if self.use_aux_hidden_state_outputs:
3381
+ assert aux_hidden_states is not None
3382
+ target_hidden_states = torch.cat(
3383
+ [h[:total_num_tokens] for h in aux_hidden_states], dim=-1
3384
+ )
3385
+ else:
3386
+ target_hidden_states = hidden_states[:total_num_tokens]
3387
3388
if self.supports_mm_inputs:
3389
mm_embed_inputs = self._gather_mm_embeddings(