address comments

Signed-off-by: qizixi <qizixi@meta.com>

Author: zixi-qi

vllm/v1/spec_decode/metadata.py

@@ -20,7 +20,6 @@ class SpecDecodeMetadata:
     bonus_logits_indices: torch.Tensor
     # [num_tokens + batch_size]
     logits_indices: torch.Tensor
-    total_num_scheduled_tokens: int
 
     def __post_init__(self):
         self.max_spec_len = max(self.num_draft_tokens)
@@ -59,5 +58,4 @@ def make_dummy(
             target_logits_indices=target_logits_indices,
             bonus_logits_indices=bonus_logits_indices,
             logits_indices=logits_indices,
-            total_num_scheduled_tokens=num_tokens,
         )

vllm/v1/worker/gpu_model_runner.py

@@ -34,8 +34,9 @@
 from vllm.sampling_params import SamplingType
 from vllm.sequence import IntermediateTensors
 from vllm.utils import (STR_DTYPE_TO_TORCH_DTYPE, DeviceMemoryProfiler,
-                        GiB_bytes, LazyLoader, async_tensor_h2d, cdiv,
-                        check_use_alibi, is_pin_memory_available)
+                        GiB_bytes, LayerBlockType, LazyLoader,
+                        async_tensor_h2d, cdiv, check_use_alibi,
+                        is_pin_memory_available)
 from vllm.v1.attention.backends.flash_attn import FlashAttentionMetadata
 from vllm.v1.attention.backends.utils import CommonAttentionMetadata
 from vllm.v1.core.encoder_cache_manager import compute_encoder_budget
@@ -898,7 +899,6 @@ def _calc_spec_decode_metadata(
             target_logits_indices=target_logits_indices,
             bonus_logits_indices=bonus_logits_indices,
             logits_indices=logits_indices,
-            total_num_scheduled_tokens=cu_num_scheduled_tokens[-1],
         )
         return metadata
 
@@ -1397,8 +1397,7 @@ def execute_model(
                     dtype=torch.int32,
                     target_device=self.device,
                     pin_memory=True)
-                num_tokens = spec_decode_metadata.total_num_scheduled_tokens - \
-                    sum(num_rejected_tokens)
+                num_tokens = num_scheduled_tokens - sum(num_rejected_tokens)
                 cu_num_tokens, token_indices = self.drafter.prepare_inputs(
                     eagle_attn_metadata.query_start_loc,
                     num_rejected_tokens_tensor,