Merge branch 'v0.9.2-dev' of http://10.16.6.30/dcutoolkit/deeplearing/vllm into v0.9.2-dev

vllm/attention/backends/flashmla.py

@@ -211,8 +211,9 @@ class FlashMLAImpl(MLACommonImpl[FlashMLAMetadata]):
                                       "FlashMLAImpl")
 
         if is_quantized_kv_cache(self.kv_cache_dtype):
-            raise NotImplementedError(
-                "FlashMLA with FP8 KV cache not yet supported")
+            if self.kv_cache_dtype != "fp8":
+                raise NotImplementedError(
+                    "FlashMLA with other KV cache not yet supported")
 
     def _forward_decode(
         self,
@@ -220,6 +221,8 @@ class FlashMLAImpl(MLACommonImpl[FlashMLAMetadata]):
         q_pe: torch.Tensor,
         kv_c_and_k_pe_cache: torch.Tensor,
         attn_metadata: FlashMLAMetadata,
+        k_scale = None,
+        kv_cache_dtype = "auto",
     ) -> torch.Tensor:
         assert kv_c_and_k_pe_cache.numel() > 0
 
@@ -239,6 +242,8 @@ class FlashMLAImpl(MLACommonImpl[FlashMLAMetadata]):
             num_splits=decode_meta.decode_num_splits,
             softmax_scale=self.scale,
             causal=True,
+            k_scale = k_scale,
+            kv_cache_dtype = kv_cache_dtype,
         )
 
         return self._v_up_proj(o)

vllm/attention/backends/mla/common.py

@@ -1397,6 +1397,6 @@ class MLACommonImpl(MLAAttentionImpl[T], Generic[T]):
             decode_ql_nope = decode_ql_nope.transpose(0, 1)
 
             output[num_prefill_tokens:] = self._forward_decode(
-                decode_ql_nope, decode_q_pe, kv_cache, attn_metadata)
+                decode_ql_nope, decode_q_pe, kv_cache, attn_metadata, layer._k_scale, self.kv_cache_dtype)
 
         return output
\ No newline at end of file

vllm/attention/ops/flashmla.py

@@ -75,6 +75,8 @@ def flash_mla_with_kvcache(
     num_splits: torch.Tensor,
     softmax_scale: Optional[float] = None,
     causal: bool = False,
+    k_scale = None,
+    kv_cache_dtype = "auto",
 ) -> Tuple[torch.Tensor, torch.Tensor]:
     """
     Arguments:
@@ -97,6 +99,22 @@ def flash_mla_with_kvcache(
     if softmax_scale is None:
         softmax_scale = q.shape[-1]**(-0.5)
     if current_platform.is_rocm():
+        if kv_cache_dtype == "fp8":         
+            out, softmax_lse = flash_mla_cuda.fwd_kvcache_mla(
+                q,
+                k_cache,
+                None,
+                head_dim_v,
+                cache_seqlens,
+                block_table,
+                softmax_scale,
+                causal,
+                tile_scheduler_metadata,
+                num_splits,
+                k_scale,
+                "fp8_e4m3",
+            )
+            return out, softmax_lse
         out, softmax_lse = flash_mla_cuda.fwd_kvcache_mla(
             q,
             k_cache,

vllm/compilation/decorators.py

@@ -11,6 +11,7 @@ from torch._dynamo.symbolic_convert import InliningInstructionTranslator
 
 from vllm.compilation.counter import compilation_counter
 from vllm.compilation.wrapper import TorchCompileWrapperWithCustomDispatcher
+from vllm.forward_context import get_profilling
 from vllm.config import CompilationLevel, VllmConfig
 from vllm.logger import init_logger
 from vllm.sequence import IntermediateTensors
@@ -169,7 +170,7 @@ def _support_torch_compile(
         # torch.compiler.is_compiling() means we are inside the compilation
         # e.g. TPU has the compilation logic in model runner, so we don't
         # need to compile the model inside.
-        if self.do_not_compile or torch.compiler.is_compiling():
+        if self.do_not_compile or torch.compiler.is_compiling() or get_profilling():
             return self.forward(*args, **kwargs)
 
         # the first compilation needs to have dynamic shapes marked

vllm/envs.py

@@ -1087,7 +1087,7 @@ environment_variables: dict[str, Callable[[], Any]] = {
              ("true", "1")),
     # vLLM will use global cache for moe
     "VLLM_USE_GLOBAL_CACHE13":
-        lambda: (os.environ.get("VLLM_USE_GLOBAL_CACHE13", "True").lower() in
+        lambda: (os.environ.get("VLLM_USE_GLOBAL_CACHE13", "False").lower() in
                  ("true", "1")),
 
 }
@@ -1162,4 +1162,4 @@ def compute_hash() -> str:
     hash_str = hashlib.md5(str(factors).encode(),
                            usedforsecurity=False).hexdigest()
 
-    return hash_str
\ No newline at end of file
+    return hash_str

vllm/forward_context.py

@@ -196,3 +196,16 @@ def set_forward_context(
         _forward_context = prev_context
         if envs.VLLM_ENABLE_TBO:
             set_tbo_forward_context(_forward_context)
+
+
+_profiling: bool = False
+
+@contextmanager
+def set_profilling(profiling):
+    global _profiling
+    _profiling = profiling
+
+
+def get_profilling() -> bool:
+    global _profiling
+    return _profiling
\ No newline at end of file

vllm/model_executor/layers/fused_moe/fused_marlin_moe.py

@@ -18,7 +18,7 @@ from vllm.model_executor.layers.quantization.utils.marlin_utils import (
     marlin_make_workspace_new, maybe_warn_marlin_atomic_add)
 from vllm.scalar_type import ScalarType, scalar_types
 from vllm.utils import direct_register_custom_op
-
+from vllm.model_executor.layers.fused_moe.fused_moe import get_moe_cache
 def get_scalar_type(num_bits: int, has_zp: bool):
     if has_zp:
         return scalar_types.uint4 if num_bits == 4 else scalar_types.uint8
@@ -104,8 +104,8 @@ def fused_marlin_moe(
     topk = topk_ids.shape[1] # 8
     
     #暂时固定为16384
-    CHUNK_SIZE = 16384
-
+    #CHUNK_SIZE = 16384
+    CHUNK_SIZE = envs.VLLM_FUSED_MOE_CHUNK_SIZE
     M = min(num_tokens, CHUNK_SIZE)
 
     if workspace is None:
@@ -120,18 +120,21 @@ def fused_marlin_moe(
     if global_num_experts == -1:
         global_num_experts = E
     intermediate_cache2 = torch.empty( 
-        (M * topk_ids.shape[1], N),
-        device=hidden_states.device,
-        dtype=hidden_states.dtype,
-    )
-    intermediate_cache13 = torch.empty(
-        (M * topk_ids.shape[1] * max(2 * N, K), ),
+        (M * topk, N),
         device=hidden_states.device,
         dtype=hidden_states.dtype,
     )
-    intermediate_cache1 = intermediate_cache13[:M * topk_ids.shape[1] * 2 * N]
+    if envs.VLLM_USE_GLOBAL_CACHE13:
+        intermediate_cache13 = get_moe_cache(topk, N, K, device=hidden_states.device, dtype=hidden_states.dtype)
+    else:    
+        intermediate_cache13 = torch.empty(
+            (M * topk * max(2 * N, K), ),
+            device=hidden_states.device,
+            dtype=hidden_states.dtype,
+        )
+    intermediate_cache1 = intermediate_cache13[:M * topk * 2 * N]
     intermediate_cache1 = intermediate_cache1.view(-1, 2 * N)
-    intermediate_cache3 = intermediate_cache13[:M * topk_ids.shape[1] * K] 
+    intermediate_cache3 = intermediate_cache13[:M * topk * K] 
     intermediate_cache3 = intermediate_cache3.view(-1, K) 
 
     use_atomic_add = hidden_states.dtype == torch.half or \

vllm/model_executor/models/deepseek_mtp.py

@@ -58,6 +58,11 @@ class DeepSeekMultiTokenPredictorLayer(nn.Module):
         quant_config: Optional[QuantizationConfig] = None,
     ) -> None:
         super().__init__()
+        self.embed_tokens = VocabParallelEmbedding(
+            config.vocab_size,
+            config.hidden_size,
+        )
+
         self.enorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
         self.hnorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
         self.eh_proj = nn.Linear(config.hidden_size * 2,
@@ -75,6 +80,8 @@ class DeepSeekMultiTokenPredictorLayer(nn.Module):
         inputs_embeds: Optional[torch.Tensor] = None,
         spec_step_index: int = 0,
     ) -> torch.Tensor:
+        if inputs_embeds is None:
+            inputs_embeds = self.embed_tokens(input_ids)
         assert inputs_embeds is not None
         # masking inputs at position 0, as not needed by MTP
         inputs_embeds[positions == 0] = 0
@@ -111,10 +118,7 @@ class DeepSeekMultiTokenPredictor(nn.Module):
             for idx in range(self.mtp_start_layer_idx,
                              self.mtp_start_layer_idx + self.num_mtp_layers)
         })
-        self.embed_tokens = VocabParallelEmbedding(
-            config.vocab_size,
-            config.hidden_size,
-        )
+
         self.logits_processor = LogitsProcessor(config.vocab_size)
 
     def forward(
@@ -125,8 +129,6 @@ class DeepSeekMultiTokenPredictor(nn.Module):
         inputs_embeds: Optional[torch.Tensor] = None,
         spec_step_idx: int = 0,
     ) -> torch.Tensor:
-        if inputs_embeds is None:
-            inputs_embeds = self.embed_tokens(input_ids)
         current_step_idx = (spec_step_idx % self.num_mtp_layers)
         return self.layers[str(self.mtp_start_layer_idx + current_step_idx)](
             input_ids,
@@ -308,25 +310,353 @@ class DeepSeekMTP(nn.Module, SupportsPP):
         """
         Rewrite the weight name to match the format of the original model.
         Add .mtp_block for modules in transformer layer block for spec layer
-        and rename shared layer weights to be top level.
         """
         spec_layer_weight_names = [
             "embed_tokens", "enorm", "hnorm", "eh_proj", "shared_head"
         ]
-        shared_weight_names = ["embed_tokens"]
         spec_layer_weight = False
-        shared_weight = False
         for weight_name in spec_layer_weight_names:
             if weight_name in name:
                 spec_layer_weight = True
-                if weight_name in shared_weight_names:
-                    shared_weight = True
                 break
         if not spec_layer_weight:
             # treat rest weights as weights for transformer layer block
             name = name.replace(f"model.layers.{spec_layer}.",
                                 f"model.layers.{spec_layer}.mtp_block.")
-        elif shared_weight:
-            # treat shared weights as top level weights
-            name = name.replace(f"model.layers.{spec_layer}.", "model.")
         return name
+    
+
+
+# # SPDX-License-Identifier: Apache-2.0
+# # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# import os
+# import re
+
+# from collections.abc import Iterable
+# from typing import Iterable, Optional
+
+
+# import torch
+# import torch.nn as nn
+# from transformers import PretrainedConfig
+
+# from vllm.config import CacheConfig, ModelConfig, VllmConfig
+# from vllm.model_executor.layers.fused_moe import FusedMoE
+# from vllm.model_executor.layers.layernorm import RMSNorm
+# from vllm.model_executor.layers.logits_processor import LogitsProcessor
+# from vllm.model_executor.layers.quantization import QuantizationConfig
+# from vllm.model_executor.layers.vocab_parallel_embedding import (
+#     ParallelLMHead, VocabParallelEmbedding)
+# from vllm.model_executor.model_loader.weight_utils import default_weight_loader
+# from vllm.model_executor.sampling_metadata import SamplingMetadata
+# from vllm.sequence import IntermediateTensors
+# from vllm.compilation.decorators import support_torch_compile
+# from .deepseek_v2 import (DeepseekV2DecoderLayer,
+#                           get_spec_layer_idx_from_weight_name)
+# from .interfaces import SupportsPP
+# from .utils import maybe_prefix
+# from vllm import _custom_ops as ops
+# from vllm.model_executor.layers.quantization.blockwise_int8 import BlockInt8Config
+
+
+# class SharedHead(nn.Module):
+
+#     def __init__(
+#         self,
+#         config: PretrainedConfig,
+#         quant_config: Optional[QuantizationConfig] = None,
+#     ) -> None:
+#         super().__init__()
+#         self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+#         self.head = ParallelLMHead(config.vocab_size,
+#                                    config.hidden_size,
+#                                    quant_config=quant_config)
+
+#     def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+#         return self.norm(hidden_states)
+
+
+# class DeepSeekMultiTokenPredictorLayer(nn.Module):
+
+#     def __init__(
+#         self,
+#         config: PretrainedConfig,
+#         prefix: str,
+#         model_config: ModelConfig,
+#         cache_config: Optional[CacheConfig] = None,
+#         quant_config: Optional[QuantizationConfig] = None,
+#     ) -> None:
+#         super().__init__()
+#         self.enorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+#         self.hnorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+#         self.eh_proj = nn.Linear(config.hidden_size * 2,
+#                                  config.hidden_size,
+#                                  bias=False)
+#         self.shared_head = SharedHead(config=config, quant_config=quant_config)
+#         self.mtp_block = DeepseekV2DecoderLayer(config, prefix, model_config,
+#                                                 cache_config, quant_config)
+
+#     def forward(
+#         self,
+#         input_ids: torch.Tensor,
+#         positions: torch.Tensor,
+#         previous_hidden_states: torch.Tensor,
+#         inputs_embeds: Optional[torch.Tensor] = None,
+#         spec_step_index: int = 0,
+#     ) -> torch.Tensor:
+#         assert inputs_embeds is not None
+#         # masking inputs at position 0, as not needed by MTP
+#         inputs_embeds[positions == 0] = 0
+#         inputs_embeds = self.enorm(inputs_embeds)
+#         previous_hidden_states = self.hnorm(previous_hidden_states)
+
+#         hidden_states = self.eh_proj(
+#             torch.cat([inputs_embeds, previous_hidden_states], dim=-1))
+
+#         hidden_states, residual = self.mtp_block(positions=positions,
+#                                                  hidden_states=hidden_states,
+#                                                  residual=None)
+#         hidden_states = residual + hidden_states
+#         return hidden_states
+
+
+# class DeepSeekMultiTokenPredictor(nn.Module):
+
+#     def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
+#         super().__init__()
+#         config = vllm_config.model_config.hf_config
+#         self.mtp_start_layer_idx = config.num_hidden_layers
+#         self.num_mtp_layers = config.num_nextn_predict_layers
+#         # to map the exact layer index from weights
+#         self.layers = torch.nn.ModuleDict({
+#             str(idx):
+#             DeepSeekMultiTokenPredictorLayer(
+#                 config,
+#                 f"{prefix}.layers.{idx}",
+#                 model_config=vllm_config.model_config,
+#                 cache_config=vllm_config.cache_config,
+#                 quant_config=vllm_config.quant_config,
+#             )
+#             for idx in range(self.mtp_start_layer_idx,
+#                              self.mtp_start_layer_idx + self.num_mtp_layers)
+#         })
+#         self.embed_tokens = VocabParallelEmbedding(
+#             config.vocab_size,
+#             config.hidden_size,
+#         )
+#         self.logits_processor = LogitsProcessor(config.vocab_size)
+
+#     def forward(
+#         self,
+#         input_ids: torch.Tensor,
+#         positions: torch.Tensor,
+#         previous_hidden_states: torch.Tensor,
+#         inputs_embeds: Optional[torch.Tensor] = None,
+#         spec_step_idx: int = 0,
+#     ) -> torch.Tensor:
+#         if inputs_embeds is None:
+#             inputs_embeds = self.embed_tokens(input_ids)
+#         current_step_idx = (spec_step_idx % self.num_mtp_layers)
+#         return self.layers[str(self.mtp_start_layer_idx + current_step_idx)](
+#             input_ids,
+#             positions,
+#             previous_hidden_states,
+#             inputs_embeds,
+#             current_step_idx,
+#         )
+
+#     def compute_logits(
+#         self,
+#         hidden_states: torch.Tensor,
+#         sampling_metadata: SamplingMetadata,
+#         spec_step_idx: int = 0,
+#     ) -> torch.Tensor:
+#         current_step_idx = (spec_step_idx % self.num_mtp_layers)
+#         mtp_layer = self.layers[str(self.mtp_start_layer_idx +
+#                                     current_step_idx)]
+#         logits = self.logits_processor(mtp_layer.shared_head.head,
+#                                        mtp_layer.shared_head(hidden_states),
+#                                        sampling_metadata)
+#         return logits
+
+# @support_torch_compile
+# class DeepSeekMTP(nn.Module, SupportsPP):
+
+#     def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
+#         super().__init__()
+#         self.config = vllm_config.model_config.hf_config
+#         quant_config = vllm_config.quant_config
+
+#         self.quant_method = None
+#         if quant_config is not None:
+#             self.quant_method = quant_config.get_name()
+#             os.environ['LLAMA_NN'] = '0'
+#             os.environ['LM_NN'] = '0'
+#             # The AWQ layer of MTP uses BlockInt8W8A8.
+#             if self.quant_method == "moe_wna16" or self.quant_method == "awq_marlin":
+#                 vllm_config.quant_config = BlockInt8Config(is_checkpoint_int8_serialized=True, weight_block_size=[128,128])
+
+#         self.model = DeepSeekMultiTokenPredictor(vllm_config=vllm_config,
+#                                                  prefix=maybe_prefix(
+#                                                      prefix, "model"))
+#         self.use_llama_nn = os.environ.get('LLAMA_NN') == '1'
+
+
+#     def forward(
+#         self,
+#         input_ids: torch.Tensor,
+#         positions: torch.Tensor,
+#         previous_hidden_states: torch.Tensor,
+#         intermediate_tensors: Optional[IntermediateTensors] = None,
+#         inputs_embeds: Optional[torch.Tensor] = None,
+#         spec_step_idx: int = 0,
+#     ) -> torch.Tensor:
+#         hidden_states = self.model(input_ids, positions,
+#                                    previous_hidden_states, inputs_embeds,
+#                                    spec_step_idx)
+#         return hidden_states
+
+#     def compute_logits(
+#         self,
+#         hidden_states: torch.Tensor,
+#         sampling_metadata: SamplingMetadata,
+#         spec_step_idx: int = 0,
+#     ) -> Optional[torch.Tensor]:
+#         return self.model.compute_logits(hidden_states, sampling_metadata,
+#                                          spec_step_idx)
+
+#     def load_weights(self, weights: Iterable[tuple[str,
+#                                                    torch.Tensor]]) -> set[str]:
+#         stacked_params_mapping = [
+#             ("gate_up_proj", "gate_proj", 0),
+#             ("gate_up_proj", "up_proj", 1),
+#         ]
+
+#         expert_params_mapping = FusedMoE.make_expert_params_mapping(
+#             ckpt_gate_proj_name="gate_proj",
+#             ckpt_down_proj_name="down_proj",
+#             ckpt_up_proj_name="up_proj",
+#             num_experts=self.config.n_routed_experts)
+
+#         params_dict = dict(self.named_parameters())
+#         loaded_params: set[str] = set()
+#         for name, loaded_weight in weights:
+#             if "rotary_emb.inv_freq" in name:
+#                 continue
+#             spec_layer = get_spec_layer_idx_from_weight_name(self.config, name)
+#             if spec_layer is None:
+#                 continue
+#             name = self._rewrite_spec_layer_name(spec_layer, name)
+#             for (param_name, weight_name, shard_id) in stacked_params_mapping:
+#                 # Skip non-stacked layers and experts (experts handled below).
+#                 if weight_name not in name:
+#                     continue
+#                 # We have mlp.experts[0].gate_proj in the checkpoint.
+#                 # Since we handle the experts below in expert_params_mapping,
+#                 # we need to skip here BEFORE we update the name, otherwise
+#                 # name will be updated to mlp.experts[0].gate_up_proj, which
+#                 # will then be updated below in expert_params_mapping
+#                 # for mlp.experts[0].gate_gate_up_proj, which breaks load.
+#                 if (("mlp.experts." in name) and name not in params_dict):
+#                     continue
+#                 name = name.replace(weight_name, param_name)
+#                 # Skip loading extra bias for GPTQ models.
+#                 if name.endswith(".bias") and name not in params_dict:
+#                     continue
+
+#                 param = params_dict[name]
+#                 weight_loader = param.weight_loader
+#                 weight_loader(param, loaded_weight, shard_id)
+#                 break
+#             else:
+#                 for mapping in expert_params_mapping:
+#                     param_name, weight_name, expert_id, shard_id = mapping
+#                     if weight_name not in name:
+#                         continue
+#                     name = name.replace(weight_name, param_name)
+
+#                     param = params_dict[name]
+#                     weight_loader = param.weight_loader
+#                     weight_loader(param,
+#                                   loaded_weight,
+#                                   name,
+#                                   shard_id=shard_id,
+#                                   expert_id=expert_id)
+#                     break
+#                 else:
+#                     # Skip loading extra bias for GPTQ models.
+#                     if name.endswith(".bias") and name not in params_dict:
+#                         continue
+
+#                     # According to DeepSeek-V3 Technical Report, MTP modules
+#                     # shares embedding layer. We only load the first weights.
+#                     if (spec_layer != self.model.mtp_start_layer_idx
+#                             and ".layers" not in name):
+#                         continue
+
+#                     param = params_dict[name]
+#                     weight_loader = getattr(param, "weight_loader",
+#                                             default_weight_loader)
+#                     weight_loader(param, loaded_weight)
+#             loaded_params.add(name)
+            
+#         if self.use_llama_nn and self.quant_method is None:
+#             lay_key_words = [
+#                 "self_attn.eh_proj.weight",
+#                 "self_attn.q_proj.weight",
+#                 "self_attn.q_a_proj.weight",
+#                 "self_attn.q_b_proj.weight",
+#                 "self_attn.kv_a_proj_with_mqa.weight",
+#                 "self_attn.kv_b_proj.weight",
+#                 "self_attn.o_proj.weight",
+#                 "mlp.gate_up_proj.weight",
+#                 "mlp.down_proj.weight",
+#                 "mlp.gate.weight",
+#                 "shared_experts.gate_up_proj.weight",
+#                 "shared_experts.down_proj.weight",
+#                 "shared_head.head.weight",
+#             ]
+
+#             combined_words = "|".join(lay_key_words)
+            
+#             for layername in loaded_params:
+#                 weight = params_dict[layername]
+#                 matches = re.findall(combined_words, layername)
+#                 if matches:
+#                     _weight = torch.zeros_like(weight.data)
+#                     ori_shape =_weight.shape
+                    
+#                     ops.trans_w16_gemm(_weight, weight.data, _weight.shape[0], _weight.shape[1])
+#                     weight.data.copy_(_weight)
+                    
+#                     weight.data=weight.data.reshape(ori_shape[1],-1)
+
+#         return loaded_params
+
+#     def _rewrite_spec_layer_name(self, spec_layer: int, name: str) -> str:
+#         """
+#         Rewrite the weight name to match the format of the original model.
+#         Add .mtp_block for modules in transformer layer block for spec layer
+#         and rename shared layer weights to be top level.
+#         """
+#         spec_layer_weight_names = [
+#             "embed_tokens", "enorm", "hnorm", "eh_proj", "shared_head"
+#         ]
+#         shared_weight_names = ["embed_tokens"]
+#         spec_layer_weight = False
+#         shared_weight = False
+#         for weight_name in spec_layer_weight_names:
+#             if weight_name in name:
+#                 spec_layer_weight = True
+#                 if weight_name in shared_weight_names:
+#                     shared_weight = True
+#                 break
+#         if not spec_layer_weight:
+#             # treat rest weights as weights for transformer layer block
+#             name = name.replace(f"model.layers.{spec_layer}.",
+#                                 f"model.layers.{spec_layer}.mtp_block.")
+#         elif shared_weight:
+#             # treat shared weights as top level weights
+#             name = name.replace(f"model.layers.{spec_layer}.", "model.")
+#         return name
+

vllm/v1/attention/backends/mla/common.py

@@ -647,10 +647,22 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):
 
                 repeats = torch.from_numpy(query_lens).pin_memory().to(
                     block_table_tensor.device, non_blocking=True).contiguous()
-                decode_block_table_tensor = torch.repeat_interleave(
-                    block_table_tensor[:self._num_decodes, ...],
-                    repeats, dim=0).contiguous()
-                decode_seq_lens = torch.repeat_interleave(seq_lens[:self._num_decodes], repeats, dim=0).contiguous()
+                
+                if envs.VLLM_ZERO_OVERHEAD:
+                    decode_block_table_tensor = torch.empty((self._num_decode_tokens, block_table_tensor.shape[1]), 
+                                                            device=block_table_tensor.device)
+                    arange_np = np.arange(self._num_decodes)
+                    indices_np = np.repeat(arange_np, query_lens)
+                    
+                    indices = torch.from_numpy(indices_np).pin_memory().to(
+                        block_table_tensor.device, non_blocking=True)
+                    decode_block_table_tensor = block_table_tensor[indices].contiguous()
+                    decode_seq_lens = seq_lens[indices].contiguous()
+                else:
+                    decode_block_table_tensor = torch.repeat_interleave(
+                        block_table_tensor[:self._num_decodes, ...],
+                        repeats, dim=0).contiguous()
+                    decode_seq_lens = torch.repeat_interleave(seq_lens[:self._num_decodes], repeats, dim=0).contiguous()
                 seq_lens_minus = torch.from_numpy(rarange).to(torch.int32).pin_memory().to(
                     seq_lens.device, non_blocking=True).contiguous()
                 decode_seq_lens = decode_seq_lens - seq_lens_minus
@@ -1086,6 +1098,6 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):
             decode_ql_nope = decode_ql_nope.transpose(0, 1)
 
             output[:num_decode_tokens] = self._forward_decode(
-                decode_ql_nope, decode_q_pe, kv_cache, attn_metadata)
+                decode_ql_nope, decode_q_pe, kv_cache, attn_metadata, layer._k_scale, self.kv_cache_dtype)
 
         return output_padded
\ No newline at end of file

vllm/v1/attention/backends/mla/flashmla.py

@@ -148,8 +148,9 @@ class FlashMLAImpl(MLACommonImpl[FlashMLAMetadata]):
                                       "FlashMLAImpl")
 
         if is_quantized_kv_cache(self.kv_cache_dtype):
-            raise NotImplementedError(
-                "FlashMLA V1 with FP8 KV cache not yet supported")
+            if self.kv_cache_dtype != "fp8":
+                raise NotImplementedError(
+                    "FlashMLA with other KV cache not yet supported")
 
     def _forward_decode(
         self,
@@ -157,6 +158,8 @@ class FlashMLAImpl(MLACommonImpl[FlashMLAMetadata]):
         q_pe: torch.Tensor,
         kv_c_and_k_pe_cache: torch.Tensor,
         attn_metadata: FlashMLAMetadata,
+        k_scale = None,
+        kv_cache_dtype = "auto",        
     ) -> torch.Tensor:
         assert kv_c_and_k_pe_cache.numel() > 0
         assert attn_metadata.decode is not None
@@ -175,6 +178,8 @@ class FlashMLAImpl(MLACommonImpl[FlashMLAMetadata]):
             num_splits=attn_metadata.decode.num_splits,
             softmax_scale=self.scale,
             causal=True,
+            k_scale = k_scale,
+            kv_cache_dtype = kv_cache_dtype,            
         )
 
         return self._v_up_proj(o)

vllm/v1/worker/gpu_model_runner.py

@@ -29,7 +29,7 @@ from vllm.distributed.parallel_state import (
     get_pp_group, get_tp_group, graph_capture, is_global_first_rank,
     prepare_communication_buffer_for_model)
 from vllm.forward_context import (DPMetadata, get_forward_context,
-                                  set_forward_context)
+                                  set_forward_context, set_profilling)
 from vllm.logger import init_logger
 from vllm.model_executor.layers.mamba.mamba_mixer2 import MambaMixer2
 from vllm.model_executor.layers.rotary_embedding import MRotaryEmbedding
@@ -69,7 +69,6 @@ from vllm.v1.worker.gpu_input_batch import CachedRequestState, InputBatch
 from vllm.v1.worker.lora_model_runner_mixin import LoRAModelRunnerMixin
 from vllm.platforms import current_platform
 from vllm.two_batch_overlap.v1.model_input_split_v1 import tbo_split_and_execute_model
-from vllm.zero_overhead.v1.gpu_model_runner import execute_model_sampled, zero_prepare_inputs
 
 from ..sample.logits_processor import LogitsProcessorManager
 from .utils import (gather_mm_placeholders, initialize_kv_cache_for_kv_sharing,
@@ -955,15 +954,25 @@ class GPUModelRunner(LoRAModelRunnerMixin):
         # [0, 1, 2, 5, 6, 9]
         target_logits_indices += arange
 
-        # TODO: Optimize the CPU -> GPU copy.
-        cu_num_draft_tokens = torch.from_numpy(cu_num_draft_tokens).to(
-            self.device, non_blocking=True)
-        logits_indices = torch.from_numpy(logits_indices).to(self.device,
-                                                             non_blocking=True)
-        target_logits_indices = torch.from_numpy(target_logits_indices).to(
-            self.device, non_blocking=True)
-        bonus_logits_indices = torch.from_numpy(bonus_logits_indices).to(
-            self.device, non_blocking=True)
+        if envs.VLLM_ZERO_OVERHEAD:
+            cu_num_draft_tokens = torch.from_numpy(cu_num_draft_tokens).pin_memory().to(
+                self.device, non_blocking=True)
+            logits_indices = torch.from_numpy(logits_indices).pin_memory().to(self.device,
+                                                                non_blocking=True)
+            target_logits_indices = torch.from_numpy(target_logits_indices).pin_memory().to(
+                self.device, non_blocking=True)
+            bonus_logits_indices = torch.from_numpy(bonus_logits_indices).pin_memory().to(
+                self.device, non_blocking=True)
+        else:
+            # TODO: Optimize the CPU -> GPU copy.
+            cu_num_draft_tokens = torch.from_numpy(cu_num_draft_tokens).to(
+                self.device, non_blocking=True)
+            logits_indices = torch.from_numpy(logits_indices).to(self.device,
+                                                                non_blocking=True)
+            target_logits_indices = torch.from_numpy(target_logits_indices).to(
+                self.device, non_blocking=True)
+            bonus_logits_indices = torch.from_numpy(bonus_logits_indices).to(
+                self.device, non_blocking=True)
 
         # Compute the draft token ids.
         # draft_token_indices:      [  1,   2,   3, 105, 106, 208]
@@ -1364,8 +1373,6 @@ class GPUModelRunner(LoRAModelRunnerMixin):
         # compiled with full CUDA graphs, we have to skip them entirely.
         skip_cuda_graphs = self.full_cuda_graph and not attention_cuda_graphs
 
-        if envs.VLLM_ZERO_OVERHEAD:
-            zero_prepare_inputs(self, scheduler_output, input_ids)
         if envs.VLLM_ENABLE_TBO and not self.use_cuda_graph:
             model_output, finished_sending, finished_recving = \
                  tbo_split_and_execute_model(self, attn_metadata, num_input_tokens,
@@ -1507,21 +1514,6 @@ class GPUModelRunner(LoRAModelRunnerMixin):
         sampled_token_ids = sampler_output.sampled_token_ids
         max_gen_len = sampled_token_ids.shape[-1]
 
-        if envs.VLLM_ZERO_OVERHEAD:
-            return execute_model_sampled(self, max_gen_len, sampled_token_ids, 
-                          discard_sampled_tokens_req_indices, scheduler_output,
-                          sampling_metadata,
-                          hidden_states,
-                          sample_hidden_states,
-                          aux_hidden_states,
-                          spec_decode_metadata,
-                          attn_metadata,
-                          logprobs_lists,
-                          prompt_logprobs_dict,
-                          finished_sending,
-                          finished_recving,
-                          num_nans_in_logits)
-
         if max_gen_len == 1:
             # No spec decode tokens.
             valid_sampled_token_ids = sampled_token_ids.tolist()
@@ -2095,7 +2087,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
             else:
                 hidden_states = outputs
 
-            if self.speculative_config and self.speculative_config.use_eagle():
+            if self.speculative_config and self.speculative_config.use_eagle() and not is_profile:
                 assert isinstance(self.drafter, EagleProposer)
                 self.drafter.dummy_run(num_tokens, attn_metadata)
 
@@ -2230,6 +2222,10 @@ class GPUModelRunner(LoRAModelRunnerMixin):
         return pooler_output
 
     def profile_run(self) -> None:
+        # set profiling flag to avoid torch compile
+        set_profilling(True)
+        self._sync_device()
+
         # Profile with multimodal encoder & encoder cache.
         # TODO: handle encoder-decoder models once we support them.
         if (self.is_multimodal_model and self.max_num_encoder_input_tokens > 0
@@ -2313,6 +2309,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
         del hidden_states, output
         self.encoder_cache.clear()
         gc.collect()
+        set_profilling(False)
 
     def capture_model(self) -> None:
         if not self.use_cuda_graph:

vllm/v1/worker/gpu_worker.py

@@ -29,6 +29,7 @@ from vllm.v1.utils import report_usage_stats
 from vllm.v1.worker.gpu_model_runner import GPUModelRunner
 from vllm.v1.worker.worker_base import WorkerBase
 from vllm.zero_overhead.utils import zero_overhead_stream
+from vllm.zero_overhead.v1.gpu_model_runner import V1ZeroModelRunner
 
 logger = init_logger(__name__)
 
@@ -163,8 +164,13 @@ class Worker(WorkerBase):
         set_random_seed(self.model_config.seed)
 
         # Construct the model runner
-        self.model_runner: GPUModelRunner = GPUModelRunner(
-            self.vllm_config, self.device)
+        if envs.VLLM_ZERO_OVERHEAD:
+            logger.info('use zero overhead model_runner')
+            self.model_runner: GPUModelRunner = V1ZeroModelRunner(
+                self.vllm_config, self.device)
+        else:
+            self.model_runner: GPUModelRunner = GPUModelRunner(
+                self.vllm_config, self.device)
 
         if self.rank == 0:
             # If usage stat is enabled, collect relevant info.

vllm/zero_overhead/v1/core.py

@@ -14,11 +14,15 @@ requsets_valid_token_len = {}
 
 def check_stop(request: Request,
                max_model_len: int,
-               pooler_output: Optional[torch.Tensor] = None) -> bool:
-    if request.request_id not in requsets_valid_token_len:
-        requsets_valid_token_len[request.request_id] = 0
-        return False
-    valid_output_len = requsets_valid_token_len[request.request_id]
+               pooler_output: Optional[torch.Tensor] = None,
+               use_valid_token_len:bool = False) -> bool:
+    if use_valid_token_len:
+        if request.request_id not in requsets_valid_token_len:
+            requsets_valid_token_len[request.request_id] = 0
+            return False
+        valid_output_len = requsets_valid_token_len[request.request_id]
+    else:
+        valid_output_len = request.num_output_tokens
     valid_num_tokens = request.num_prompt_tokens + valid_output_len
     if (valid_num_tokens >= max_model_len
             or valid_output_len >= request.max_tokens):
@@ -62,110 +66,121 @@ def zero_overhead_update_from_output(scheduler:Scheduler,
     spec_decoding_stats: Optional[SpecDecodingStats] = None
 
     # fix last model out in zero overhead
-    for req_idx, req_id in enumerate(model_runner_output.fix_req_ids):
-        if req_id not in scheduler.requests:
-            continue
-        request = scheduler.requests[req_id]
-        generated_token_ids = model_runner_output.fix_sampled_token_ids[req_idx]
-        if req_id not in requsets_valid_token_len:
-            requsets_valid_token_len[req_id] = 0
-        valid_output_len = requsets_valid_token_len[req_id]
-        fix_offset = valid_output_len - request.num_output_tokens
-        if isinstance(generated_token_ids, int):
-            request._output_token_ids[fix_offset] = generated_token_ids
-            request._all_token_ids[fix_offset] = generated_token_ids
-            requsets_valid_token_len[req_id] += 1
-        else:
-            valid_output_end = valid_output_len + len(generated_token_ids) - request.num_output_tokens
-            if valid_output_end == 0:
-                request._output_token_ids[fix_offset : ] = generated_token_ids
-                request._all_token_ids[fix_offset : ] = generated_token_ids
+    if model_runner_output.fix_req_ids is not None:
+        for req_idx, req_id in enumerate(model_runner_output.fix_req_ids):
+            if req_id not in scheduler.requests:
+                continue
+            request = scheduler.requests[req_id]
+            generated_token_ids = model_runner_output.fix_sampled_token_ids[req_idx]
+            if req_id not in requsets_valid_token_len:
+                requsets_valid_token_len[req_id] = 0
+            valid_output_len = requsets_valid_token_len[req_id]
+            fix_offset = valid_output_len - request.num_output_tokens
+            if isinstance(generated_token_ids, int):
+                request._output_token_ids[fix_offset] = generated_token_ids
+                request._all_token_ids[fix_offset] = generated_token_ids
+                requsets_valid_token_len[req_id] += 1
             else:
-                request._output_token_ids[fix_offset : valid_output_end] = generated_token_ids
-                request._all_token_ids[fix_offset : valid_output_end] = generated_token_ids
-            requsets_valid_token_len[req_id] += len(generated_token_ids) 
-
+                valid_output_end = valid_output_len + len(generated_token_ids) - request.num_output_tokens
+                if valid_output_end == 0:
+                    request._output_token_ids[fix_offset : ] = generated_token_ids
+                    request._all_token_ids[fix_offset : ] = generated_token_ids
+                else:
+                    request._output_token_ids[fix_offset : valid_output_end] = generated_token_ids
+                    request._all_token_ids[fix_offset : valid_output_end] = generated_token_ids
+                requsets_valid_token_len[req_id] += len(generated_token_ids) 
+
+                
+            stopped = False
+            new_logprobs = None
+            new_token_ids = generated_token_ids
+            kv_transfer_params = None
             
-        stopped = False
-        new_logprobs = None
-        new_token_ids = generated_token_ids
-        kv_transfer_params = None
-        
-        # Check for stop and update request state.
-        # This must be called before we make the EngineCoreOutput.
-        for num_new, output_token_id in enumerate(new_token_ids, 1):
-            stopped = check_stop(request, scheduler.max_model_len)
-            if stopped:
-                kv_transfer_params = scheduler._free_request(request)
-                del new_token_ids[num_new:]  # Trim new tokens if needed.
-                break
-
-        pooler_output = None
-        if pooler_outputs:
-            pooler_output = pooler_outputs[req_index]
-            stopped = check_stop(request, scheduler.max_model_len,
-                                    pooler_output)
-            if stopped:
-                kv_transfer_params = scheduler._free_request(request)
-
-        # Extract sample logprobs if needed.
-        if request.sampling_params is not None \
-            and request.sampling_params.logprobs is not None and logprobs:
-            # NOTE: once we support N tokens per step (spec decode),
-            # the outer lists can be of length > 1.
-            new_logprobs = logprobs.slice(req_index, req_index + 1)
-
-        if new_token_ids and scheduler.structured_output_manager.should_advance(
-                request):
-            # NOTE: structured_output_request
-            # should not be None if use_structured_output, we have
-            # check above, so safe to ignore type warning
-            request.structured_output_request.grammar.accept_tokens(  # type: ignore[union-attr]
-                req_id, new_token_ids)
-
-        # spec_token_ids comes from the model runner output
-        if num_nans_in_logits is not None and req_id in num_nans_in_logits:
-            request.num_nans_in_logits = num_nans_in_logits[req_id]
+            # Check for stop and update request state.
+            # This must be called before we make the EngineCoreOutput.
+            for num_new, output_token_id in enumerate(new_token_ids, 1):
+                stopped = check_stop(request, scheduler.max_model_len, True)
+                if stopped:
+                    kv_transfer_params = scheduler._free_request(request)
+                    del new_token_ids[num_new:]  # Trim new tokens if needed.
+                    break
+
+            pooler_output = None
+            if pooler_outputs:
+                pooler_output = pooler_outputs[req_index]
+                stopped = check_stop(request, scheduler.max_model_len,
+                                        pooler_output, True)
+                if stopped:
+                    kv_transfer_params = scheduler._free_request(request)
+
+            # Extract sample logprobs if needed.
+            if request.sampling_params is not None \
+                and request.sampling_params.logprobs is not None and logprobs:
+                # NOTE: once we support N tokens per step (spec decode),
+                # the outer lists can be of length > 1.
+                new_logprobs = logprobs.slice(req_index, req_index + 1)
+
+            if new_token_ids and scheduler.structured_output_manager.should_advance(
+                    request):
+                # NOTE: structured_output_request
+                # should not be None if use_structured_output, we have
+                # check above, so safe to ignore type warning
+                request.structured_output_request.grammar.accept_tokens(  # type: ignore[union-attr]
+                    req_id, new_token_ids)
+
+            # spec_token_ids comes from the model runner output
+            if num_nans_in_logits is not None and req_id in num_nans_in_logits:
+                request.num_nans_in_logits = num_nans_in_logits[req_id]
+
+            # Get prompt logprobs for this request.
+            prompt_logprobs_tensors = prompt_logprobs_dict.get(req_id)
+            if new_token_ids or pooler_output is not None \
+                or kv_transfer_params:
+
+                # Add EngineCoreOutput for this Request.
+                outputs[request.client_index].append(
+                    EngineCoreOutput(
+                        request_id=req_id,
+                        new_token_ids=new_token_ids,
+                        finish_reason=request.get_finished_reason(),
+                        new_logprobs=new_logprobs,
+                        new_prompt_logprobs_tensors=prompt_logprobs_tensors,
+                        pooling_output=pooler_output,
+                        stop_reason=request.stop_reason,
+                        events=request.take_events(),
+                        kv_transfer_params=kv_transfer_params,
+                        num_cached_tokens=request.num_cached_tokens,
+                    ))
 
-        # Add newly generated spec token ids to the request.
-        if spec_token_ids is not None:
-            if scheduler.structured_output_manager.should_advance(request):
-                metadata = request.structured_output_request
-                # Needs to happen after new_token_ids are accepted.
-                request.spec_token_ids = metadata.grammar.validate_tokens(  # type: ignore[union-attr]
-                    spec_token_ids[req_index])
             else:
-                request.spec_token_ids = spec_token_ids[req_index]
-
-        # Get prompt logprobs for this request.
-        prompt_logprobs_tensors = prompt_logprobs_dict.get(req_id)
-        if new_token_ids or pooler_output is not None \
-            or kv_transfer_params:
-
-            # Add EngineCoreOutput for this Request.
-            outputs[request.client_index].append(
-                EngineCoreOutput(
-                    request_id=req_id,
-                    new_token_ids=new_token_ids,
-                    finish_reason=request.get_finished_reason(),
-                    new_logprobs=new_logprobs,
-                    new_prompt_logprobs_tensors=prompt_logprobs_tensors,
-                    pooling_output=pooler_output,
-                    stop_reason=request.stop_reason,
-                    events=request.take_events(),
-                    kv_transfer_params=kv_transfer_params,
-                    num_cached_tokens=request.num_cached_tokens,
-                ))
-
-        else:
-            # Invariant: EngineCore returns no partial prefill outputs.
-            assert not prompt_logprobs_tensors
+                # Invariant: EngineCore returns no partial prefill outputs.
+                assert not prompt_logprobs_tensors
 
+    # fix last model out in zero overhead
+    if model_runner_output.fix_draft_req_ids is not None:
+        for req_idx, req_id in enumerate(model_runner_output.fix_draft_req_ids):
+            if req_id not in scheduler.requests:
+                continue
+            request = scheduler.requests[req_id]
+
+            # Add newly generated spec token ids to the request.
+            if model_runner_output.fix_draft_tokens_ids is not None:
+                if scheduler.structured_output_manager.should_advance(request):
+                    metadata = request.structured_output_request
+                    # Needs to happen after new_token_ids are accepted.
+                    request.spec_token_ids = metadata.grammar.validate_tokens(  # type: ignore[union-attr]
+                        model_runner_output.fix_draft_tokens_ids[req_idx])
+                else:
+                    request.spec_token_ids = model_runner_output.fix_draft_tokens_ids[req_idx]
 
     # NOTE(woosuk): As len(self.running) can be up to 1K or more, the below
     # loop can be a performance bottleneck. We should do our best to avoid
     # expensive operations inside the loop.
     for request in scheduler.running:
+        if request.is_finished():
+            if req_id in requsets_valid_token_len:
+                requsets_valid_token_len.pop(req_id)
+            continue
         req_id = request.request_id
         num_tokens_scheduled = num_scheduled_tokens.get(req_id, 0)
         if num_tokens_scheduled == 0:
@@ -212,19 +227,24 @@ def zero_overhead_update_from_output(scheduler:Scheduler,
 
             # Check for stop and update request state.
             # This must be called before we make the EngineCoreOutput.
-            stopped = check_stop(request, scheduler.max_model_len)
-            # if stopped:
-            #     kv_transfer_params = scheduler._free_request(request)
-            #     del new_token_ids[num_new:]  # Trim new tokens if needed.
-            #     break
+            
+            if model_runner_output.is_output_valid:
+                stopped = check_stop(request, scheduler.max_model_len,
+                                     False)
+                if stopped:
+                    kv_transfer_params = scheduler._free_request(request)
+                    del new_token_ids[num_new:]  # Trim new tokens if needed.
+                    break
 
         pooler_output = None
         if pooler_outputs:
-            pooler_output = pooler_outputs[req_index]
-            stopped = check_stop(request, scheduler.max_model_len,
-                                    pooler_output)
-            # if stopped:
-            #     kv_transfer_params = scheduler._free_request(request)
+            if model_runner_output.is_output_valid:
+                pooler_output = pooler_outputs[req_index]
+                stopped = check_stop(request, scheduler.max_model_len,
+                                        pooler_output,
+                                        False)
+                if stopped:
+                    kv_transfer_params = scheduler._free_request(request)
 
         # Extract sample logprobs if needed.
         if request.sampling_params is not None \
@@ -255,7 +275,30 @@ def zero_overhead_update_from_output(scheduler:Scheduler,
             else:
                 request.spec_token_ids = spec_token_ids[req_index]
 
-        if not stopped:
+        if model_runner_output.is_output_valid:
+            # # Get prompt logprobs for this request.
+            prompt_logprobs_tensors = prompt_logprobs_dict.get(req_id)
+            if new_token_ids or pooler_output is not None \
+                or kv_transfer_params:
+
+                # Add EngineCoreOutput for this Request.
+                outputs[request.client_index].append(
+                    EngineCoreOutput(
+                        request_id=req_id,
+                        new_token_ids=new_token_ids,
+                        finish_reason=request.get_finished_reason(),
+                        new_logprobs=new_logprobs,
+                        new_prompt_logprobs_tensors=prompt_logprobs_tensors,
+                        pooling_output=pooler_output,
+                        stop_reason=request.stop_reason,
+                        events=request.take_events(),
+                        kv_transfer_params=kv_transfer_params,
+                        num_cached_tokens=request.num_cached_tokens,
+                    ))
+        if stopped:
+            if req_id in requsets_valid_token_len:
+                requsets_valid_token_len.pop(req_id)
+        else:
             new_running.append(request)
 
     scheduler.running = new_running

vllm/zero_overhead/v1/gpu_model_runner.py

 
+from typing import Any, Optional, Union
 import torch
 import numpy as np
+from vllm import envs
 from vllm.distributed.kv_transfer.kv_transfer_state import get_kv_transfer_group, has_kv_transfer_group
-from vllm.distributed.parallel_state import get_tp_group
-from vllm.utils import async_tensor_h2d
+from vllm.distributed.parallel_state import get_pp_group, get_tp_group
+from vllm.forward_context import set_forward_context
+from vllm.sequence import IntermediateTensors
+from vllm.utils import async_tensor_h2d, round_up
+from vllm.v1.attention.backends.utils import CommonAttentionMetadata
+from vllm.v1.core.sched.output import SchedulerOutput
+from vllm.v1.outputs import EMPTY_MODEL_RUNNER_OUTPUT, ModelRunnerOutput
+from vllm.v1.sample.metadata import SamplingMetadata
+from vllm.v1.spec_decode.eagle import EagleProposer
+from vllm.v1.spec_decode.medusa import MedusaProposer
+from vllm.v1.spec_decode.metadata import SpecDecodeMetadata
+from vllm.v1.spec_decode.ngram_proposer import NgramProposer
+from vllm.v1.worker.block_table import BlockTable
+from vllm.v1.worker.gpu_model_runner import GPUModelRunner
 from vllm.zero_overhead.v1.outputs import ZeroV1ModelRunnerOutput
 from vllm.profiler.prof import profile
+from vllm.two_batch_overlap.v1.model_input_split_v1 import tbo_split_and_execute_model
 
 
-class V1ZeroModelRunner():
-    def __init__(self):
+class V1ZeroModelRunner(GPUModelRunner):
+    def __init__(self, vllm_config, device):
+        super().__init__(vllm_config, device)
         self.last_sampled_token_ids = None
         self.last_sampled_req_ids = []
         self.last_sampled_token_lens = []
         self.last_sampler_event = torch.cuda.Event(enable_timing=False)
         self.last_sampler_host_tokens = None
         self.token_ids_cpu_fix_recode = []
+        self.last_draft_token_ids = None
+        self.last_draft_host_tokens = None
+        self.last_draft_event = torch.cuda.Event(enable_timing=False)
     
-    def set_last_sampled_token_ids(self, sampled_token_ids):
-        self.last_sampled_token_ids = sampled_token_ids
-        self.last_sampled_req_ids = []
-        self.last_sampled_token_lens = []
+    def _prepare_inputs(
+        self,
+        scheduler_output: "SchedulerOutput",
+    ) -> tuple[dict[str, Any], bool, torch.Tensor,
+               Optional[SpecDecodeMetadata], np.ndarray]:
+        """
+        :return: tuple[
+            attn_metadata: layer-to-attention_metadata mapping,
+            attention_cuda_graphs: whether attention can run in cudagraph
+            logits_indices, spec_decode_metadata
+        ]
+        """
+        total_num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens
+        assert total_num_scheduled_tokens > 0
+        num_reqs = self.input_batch.num_reqs
+        assert num_reqs > 0
+
+        # OPTIMIZATION: Start copying the block table first.
+        # This way, we can overlap the copy with the following CPU operations.
+        self.input_batch.block_table.commit(num_reqs)
+
+        # Get the number of scheduled tokens for each request.
+        req_ids = self.input_batch.req_ids
+        tokens = [scheduler_output.num_scheduled_tokens[i] for i in req_ids]
+        num_scheduled_tokens = np.array(tokens, dtype=np.int32)
+        max_num_scheduled_tokens = max(tokens)
+
+        # Get request indices.
+        # E.g., [2, 5, 3] -> [0, 0, 1, 1, 1, 1, 1, 2, 2, 2]
+        req_indices = np.repeat(self.arange_np[:num_reqs],
+                                num_scheduled_tokens)
+
+        # cu_num_tokens: [2, 5, 3] -> [2, 7, 10]
+        # arange: [0, 1, 0, 1, 2, 3, 4, 0, 1, 2]
+        cu_num_tokens, arange = self._get_cumsum_and_arange(
+            num_scheduled_tokens)
+
+        # Get positions.
+        positions_np = self.positions_np[:total_num_scheduled_tokens]
+        np.add(self.input_batch.num_computed_tokens_cpu[req_indices],
+               arange,
+               out=positions_np)
+
+        # Calculate M-RoPE positions.
+        # Only relevant for models using M-RoPE (e.g, Qwen2-VL)
+        if self.uses_mrope:
+            self._calc_mrope_positions(scheduler_output)
+
+        # Get token indices.
+        # E.g., [0, 1, 0, 1, 2, 3, 4, 0, 1, 2]
+        # -> [0, 1, M, M + 1, M + 2, M + 3, M + 4, 2 * M, 2 * M + 1, 2 * M + 2]
+        # where M is the max_model_len.
+        token_indices = (positions_np +
+                         req_indices * self.input_batch.token_ids_cpu.shape[1])
+
+        # NOTE(woosuk): We use torch.index_select instead of np.take here
+        # because torch.index_select is much faster than np.take for large
+        # tensors.
+        torch.index_select(self.input_batch.token_ids_cpu_tensor.flatten(),
+                           0,
+                           torch.from_numpy(token_indices),
+                           out=self.input_ids_cpu[:total_num_scheduled_tokens])
+
+        # Calculate the slot mapping for each KV cache group.
+        for kv_cache_group_id, kv_cache_group_spec in enumerate(
+                self.kv_cache_config.kv_cache_groups):
+            block_size = kv_cache_group_spec.kv_cache_spec.block_size
+            block_table: BlockTable = self.input_batch.block_table[
+                kv_cache_group_id]
+            # E.g., [0, 1, 0, 1, 2, 3, 4, 0, 1, 2]
+            # -> [0, 0, K, K, K + 1, K + 1, K + 2, 2 * K, 2 * K, 2 * K + 1]
+            # where K is the max_num_blocks_per_req and the block size is 2.
+            # NOTE(woosuk): We can't simply use `token_indices // block_size`
+            # here because M (max_model_len) is not necessarily divisible by
+            # block_size.
+            block_table_indices = (
+                req_indices * block_table.max_num_blocks_per_req +
+                positions_np // block_size)
+            block_table_cpu = block_table.get_cpu_tensor()
+            block_numbers = block_table_cpu.flatten(
+            )[block_table_indices].numpy()
+            block_offsets = positions_np % block_size
+            np.add(
+                block_numbers * block_size,
+                block_offsets,
+                out=block_table.slot_mapping_np[:total_num_scheduled_tokens])
+
+        # Prepare the attention metadata.
+        self.query_start_loc_np[0] = 0
+        self.query_start_loc_np[1:num_reqs + 1] = cu_num_tokens
+
+        self.seq_lens_np[:num_reqs] = (
+            self.input_batch.num_computed_tokens_cpu[:num_reqs] +
+            num_scheduled_tokens)
+
+        # Copy the tensors to the GPU.
+        self.input_ids[:total_num_scheduled_tokens].copy_(
+            self.input_ids_cpu[:total_num_scheduled_tokens], non_blocking=True)
+        
+        self.zero_prepare_inputs(scheduler_output, self.input_ids)
+
+        if self.uses_mrope:
+            # Only relevant for models using M-RoPE (e.g, Qwen2-VL)
+            self.mrope_positions[:, :total_num_scheduled_tokens].copy_(
+                self.mrope_positions_cpu[:, :total_num_scheduled_tokens],
+                non_blocking=True)
+        else:
+            # Common case (1D positions)
+            self.positions[:total_num_scheduled_tokens].copy_(
+                self.positions_cpu[:total_num_scheduled_tokens],
+                non_blocking=True)
+
+        self.query_start_loc[:num_reqs + 1].copy_(
+            self.query_start_loc_cpu[:num_reqs + 1], non_blocking=True)
+        self.seq_lens[:num_reqs].copy_(self.seq_lens_cpu[:num_reqs],
+                                       non_blocking=True)
 
-v1_zero_overhead = V1ZeroModelRunner()
-
-def zero_prepare_inputs(runner, scheduler_output, input_ids):
-    req_ids = runner.input_batch.req_ids
-    update_req_indices = []
-    input_ids_indices = []
-    token_idx = 0
-    if v1_zero_overhead.last_sampled_token_ids is None:
-        return
-    sampled_tokens_num = v1_zero_overhead.last_sampled_token_ids.shape[1]
-    for req_id in req_ids:
-        if req_id in v1_zero_overhead.last_sampled_req_ids:
-            req_idx = v1_zero_overhead.last_sampled_req_ids.index(req_id) * sampled_tokens_num
-            update_req_indices.append(req_idx)
-            input_ids_indices.append(token_idx)
-        token_idx += scheduler_output.num_scheduled_tokens[req_id]
-    if len(update_req_indices) > 0:
-        update_req_indices_tensor = async_tensor_h2d(update_req_indices, torch.int32,
-                                                    runner.device,
-                                                    True)
-        input_ids_indices_tensor = async_tensor_h2d(input_ids_indices, torch.int32,
-                                                    runner.device,
-                                                    True)
-        last_sampled_token_ids = v1_zero_overhead.last_sampled_token_ids.flatten()
-        for i in range(sampled_tokens_num):
-            input_ids[input_ids_indices_tensor + i] = last_sampled_token_ids[update_req_indices_tensor + i]
-
-
-def execute_model_sampled(runner, max_gen_len, sampled_token_ids, 
-                          discard_sampled_tokens_req_indices, scheduler_output,
-                          sampling_metadata,
-                          hidden_states,
-                          sample_hidden_states,
-                          aux_hidden_states,
-                          spec_decode_metadata,
-                          attn_metadata,
-                          logprobs_lists,
-                          prompt_logprobs_dict,
-                          finished_sending,
-                          finished_recving,
-                          num_nans_in_logits
-                          ):
-    fix_req_ids = None
-    fix_sampled_token_ids = None
-    if max_gen_len == 1:
-        # No spec decode tokens.
-        if v1_zero_overhead.last_sampler_host_tokens != None:
-            v1_zero_overhead.last_sampler_event.synchronize()
-            fix_sampled_token_ids = v1_zero_overhead.last_sampler_host_tokens.tolist()
-            for req_idx, start_idx, end_idx in v1_zero_overhead.token_ids_cpu_fix_recode:
-                runner.input_batch.token_ids_cpu[req_idx, start_idx:end_idx] = fix_sampled_token_ids[req_idx]
-        fix_req_ids = v1_zero_overhead.last_sampled_req_ids
-        for req_idx, req_id in enumerate(fix_req_ids):
-            if req_id in runner.requests:
-                req_state = runner.requests[req_id]
-                token_idx = v1_zero_overhead.last_sampled_token_lens[req_idx]
-                req_state.output_token_ids[token_idx] = fix_sampled_token_ids[req_idx][0]
-        v1_zero_overhead.last_sampler_host_tokens = sampled_token_ids.to('cpu', non_blocking=True)
-        v1_zero_overhead.last_sampler_event.record()
-        v1_zero_overhead.set_last_sampled_token_ids(sampled_token_ids)
-        valid_sampled_token_ids = np.ones(sampled_token_ids.shape, dtype=int).tolist()
-    else:
-        # Includes spec decode tokens.
-        valid_sampled_token_ids = runner.rejection_sampler.parse_output(
-            sampled_token_ids,
-            runner.input_batch.vocab_size,
+        # Fill unused with -1. Needed for reshape_and_cache
+        self.seq_lens[num_reqs:].fill_(0)
+        # Note: pad query_start_loc to be non-decreasing, as kernels
+        # like FlashAttention requires that
+        self.query_start_loc[num_reqs + 1:].fill_(
+            self.query_start_loc_cpu[num_reqs].item())
+
+        query_start_loc = self.query_start_loc[:num_reqs + 1]
+        seq_lens = self.seq_lens[:num_reqs]
+
+        common_attn_metadata = CommonAttentionMetadata(
+            query_start_loc=query_start_loc,
+            seq_lens=seq_lens,
+            num_reqs=num_reqs,
+            num_actual_tokens=total_num_scheduled_tokens,
+            max_query_len=max_num_scheduled_tokens,
         )
-    # Mask out the sampled tokens that should not be sampled.
-    for i in discard_sampled_tokens_req_indices:
-        valid_sampled_token_ids[i].clear()
-
-    # Cache the sampled tokens in the model runner, so that the scheduler
-    # doesn't need to send them back.
-    # NOTE(woosuk): As an exception, when using PP, the scheduler sends
-    # the sampled tokens back, because there's no direct communication
-    # between the first-stage worker and the last-stage worker.
-    v1_zero_overhead.token_ids_cpu_fix_recode.clear()
-    for req_idx, sampled_ids in enumerate(valid_sampled_token_ids):
-        if not sampled_ids:
-            continue
-
-        start_idx = runner.input_batch.num_tokens_no_spec[req_idx]
-        end_idx = start_idx + len(sampled_ids)
-        assert end_idx <= runner.max_model_len, (
-            "Sampled token IDs exceed the max model length. "
-            f"Total number of tokens: {end_idx} > max_model_len: "
-            f"{runner.max_model_len}")
-
-        runner.input_batch.token_ids_cpu[req_idx,
-                                        start_idx:end_idx] = sampled_ids
-        v1_zero_overhead.token_ids_cpu_fix_recode.append([req_idx, start_idx, end_idx])
-        runner.input_batch.num_tokens_no_spec[req_idx] = end_idx
-        runner.input_batch.num_tokens[req_idx] = end_idx
-        req_id = runner.input_batch.req_ids[req_idx]
-        if req_id in runner.requests:
-            req_state = runner.requests[req_id]
-            v1_zero_overhead.last_sampled_req_ids.append(req_id)
-            v1_zero_overhead.last_sampled_token_lens.append(len(req_state.output_token_ids))
-            req_state.output_token_ids.extend(sampled_ids)
-    if not runner.speculative_config:
-        # Speculative decoding is not enabled.
-        spec_token_ids = None
-    else:
-        spec_token_ids = runner.propose_draft_token_ids(
+
+        attn_metadata: dict[str, Any] = {}
+        # Prepare the attention metadata for each KV cache group and make layers
+        # in the same group share the same metadata.
+        for kv_cache_group_id, kv_cache_group_spec in enumerate(
+                self.kv_cache_config.kv_cache_groups):
+
+            # Prepare for cascade attention if enabled & beneficial.
+            common_prefix_len = 0
+            builder = self.attn_metadata_builders[kv_cache_group_id]
+            if self.cascade_attn_enabled:
+                common_prefix_len = self._compute_cascade_attn_prefix_len(
+                    num_scheduled_tokens,
+                    scheduler_output.
+                    num_common_prefix_blocks[kv_cache_group_id],
+                    kv_cache_group_spec.kv_cache_spec,
+                    builder,
+                )
+
+            attn_metadata_i = (builder.build(
+                common_prefix_len=common_prefix_len,
+                common_attn_metadata=common_attn_metadata,
+            ))
+
+            for layer_name in kv_cache_group_spec.layer_names:
+                attn_metadata[layer_name] = attn_metadata_i
+
+        attention_cuda_graphs = all(
+            b.can_run_in_cudagraph(common_attn_metadata)
+            for b in self.attn_metadata_builders)
+
+        use_spec_decode = len(
+            scheduler_output.scheduled_spec_decode_tokens) > 0
+        if not use_spec_decode:
+            # NOTE(woosuk): Due to chunked prefills, the batch may contain
+            # partial requests. While we should not sample any token
+            # from these partial requests, we do so for simplicity.
+            # We will ignore the sampled tokens from the partial requests.
+            # TODO: Support prompt logprobs.
+            logits_indices = query_start_loc[1:] - 1
+            spec_decode_metadata = None
+        else:
+            # Get the number of draft tokens for each request.
+            # Iterate over the dictionary rather than all requests since not all
+            # requests have draft tokens.
+            num_draft_tokens = np.zeros(num_reqs, dtype=np.int32)
+            for req_id, draft_token_ids in (
+                    scheduler_output.scheduled_spec_decode_tokens.items()):
+                req_idx = self.input_batch.req_id_to_index[req_id]
+                num_draft_tokens[req_idx] = len(draft_token_ids)
+
+            spec_decode_metadata = self._calc_spec_decode_metadata(
+                num_draft_tokens, cu_num_tokens)
+            logits_indices = spec_decode_metadata.logits_indices
+
+        # Hot-Swap lora model
+        if self.lora_config:
+            self.set_active_loras(self.input_batch, num_scheduled_tokens)
+
+        return (attn_metadata, attention_cuda_graphs, logits_indices,
+                spec_decode_metadata, num_scheduled_tokens)
+
+    def zero_prepare_inputs(self, scheduler_output, input_ids):
+        req_ids = self.input_batch.req_ids
+        update_req_indices = []
+        input_ids_indices = []
+        token_idx = 0
+        if self.last_draft_token_ids is not None:
+            draft_tokens_num = self.last_draft_token_ids.shape[1]
+            for req_id in req_ids:
+                if req_id in self.last_sampled_req_ids:
+                    req_idx = self.last_sampled_req_ids.index(req_id) * draft_tokens_num
+                    for num_idx in range(draft_tokens_num):
+                        update_req_indices.append(req_idx + num_idx)
+                        input_ids_indices.append(token_idx + num_idx + 1)
+                token_idx += draft_tokens_num + 1
+            if len(update_req_indices) > 0:
+                update_req_indices_tensor = async_tensor_h2d(update_req_indices, torch.int32,
+                                                            self.device,
+                                                            True)
+                input_ids_indices_tensor = async_tensor_h2d(input_ids_indices, torch.int32,
+                                                            self.device,
+                                                            True)
+                last_draft_token_ids = self.last_draft_token_ids.flatten().to(torch.int)
+                input_ids[input_ids_indices_tensor] = last_draft_token_ids[update_req_indices_tensor]
+
+        update_req_indices = []
+        input_ids_indices = []
+        token_idx = 0
+        if self.last_sampled_token_ids is not None:
+            sampled_tokens_num = self.last_sampled_token_ids.shape[1]
+            for req_id in req_ids:
+                if req_id in self.last_sampled_req_ids:
+                    req_idx = self.last_sampled_req_ids.index(req_id) * sampled_tokens_num
+                    update_req_indices.append(req_idx)
+                    input_ids_indices.append(token_idx)
+                token_idx += scheduler_output.num_scheduled_tokens[req_id]
+            if len(update_req_indices) > 0:
+                update_req_indices_tensor = async_tensor_h2d(update_req_indices, torch.int32,
+                                                            self.device,
+                                                            True)
+                input_ids_indices_tensor = async_tensor_h2d(input_ids_indices, torch.int32,
+                                                            self.device,
+                                                            True)
+                last_sampled_token_ids = self.last_sampled_token_ids.flatten()
+                for i in range(sampled_tokens_num):
+                    input_ids[input_ids_indices_tensor + i] = last_sampled_token_ids[update_req_indices_tensor + i]
+
+    def propose_draft_token_ids(
+        self,
+        scheduler_output: "SchedulerOutput",
+        sampled_token_ids: list[list[int]],
+        sampling_metadata: SamplingMetadata,
+        hidden_states: torch.Tensor,
+        sample_hidden_states: torch.Tensor,
+        aux_hidden_states: Optional[torch.Tensor],
+        spec_decode_metadata: Optional[SpecDecodeMetadata],
+        attn_metadata: dict[str, Any],
+    ) -> list[list[int]]:
+        num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens
+        if self.speculative_config.method == "ngram":
+            assert isinstance(self.drafter, NgramProposer)
+            spec_token_ids = self.propose_ngram_draft_token_ids(
+                sampled_token_ids)
+        elif self.speculative_config.method == "medusa":
+            assert isinstance(self.drafter, MedusaProposer)
+            if sample_hidden_states.shape[0] == len(sampled_token_ids):
+                # The input to the target model does not include draft tokens.
+                hidden_states = sample_hidden_states
+            else:
+                indices = []
+                offset = 0
+                for num_draft, tokens in zip(
+                        spec_decode_metadata.num_draft_tokens,
+                        sampled_token_ids):
+                    indices.append(offset + len(tokens) - 1)
+                    offset += num_draft + 1
+                indices = torch.tensor(indices, device=self.device)
+                hidden_states = sample_hidden_states[indices]
+
+            spec_token_ids = self.drafter.propose(
+                target_hidden_states=hidden_states,
+                sampling_metadata=sampling_metadata,
+            )
+        elif self.speculative_config.use_eagle():
+            assert isinstance(self.drafter, EagleProposer)
+            # TODO(woosuk): Refactor the loop.
+            if self.last_sampled_token_ids is not None:
+                next_token_ids = self.last_sampled_token_ids.flatten()
+            else:
+                next_token_ids: list[int] = []
+                for i, token_ids in enumerate(sampled_token_ids):
+                    if token_ids:
+                        # Common case.
+                        next_token_id = token_ids[-1]
+                    else:
+                        # Partial prefill (rare case).
+                        # Get the next token id from the request state.
+                        req_id = self.input_batch.req_ids[i]
+                        req_state = self.requests[req_id]
+                        seq_len = (req_state.num_computed_tokens +
+                                scheduler_output.num_scheduled_tokens[req_id])
+                        next_token_id = req_state.get_token_id(seq_len)
+                    next_token_ids.append(next_token_id)
+                next_token_ids = torch.tensor(next_token_ids,
+                                            dtype=torch.int32,
+                                            device=self.device)
+            # At this moment, we assume all eagle layers belong to the same KV
+            # cache group, thus using the same attention metadata.
+            eagle_attn_metadata = attn_metadata[
+                self.drafter.attn_layer_names[0]]
+
+            # NOTE: deepseek_mtp uses MLA which does not have `block_table`
+            if hasattr(eagle_attn_metadata, "block_table"):
+                block_table = eagle_attn_metadata.block_table
+            else:
+                block_table = None
+
+            if spec_decode_metadata is None:
+                # input_ids can be None for multimodal models.
+                target_token_ids = self.input_ids[:num_scheduled_tokens]
+                # TODO(woosuk): Support M-RoPE.
+                target_positions = self.positions[:num_scheduled_tokens]
+                if self.use_aux_hidden_state_outputs:
+                    target_hidden_states = torch.cat(
+                        [h[:num_scheduled_tokens] for h in aux_hidden_states],
+                        dim=-1)
+                else:
+                    target_hidden_states = hidden_states[:num_scheduled_tokens]
+                target_slot_mapping = eagle_attn_metadata.slot_mapping
+                cu_num_tokens = eagle_attn_metadata.query_start_loc
+            else:
+                # TODO(woosuk): Refactor this.
+                num_accepted_tokens = [len(s) - 1 for s in sampled_token_ids]
+                num_accepted_tokens_tensor = async_tensor_h2d(
+                    num_accepted_tokens,
+                    dtype=torch.int32,
+                    target_device=self.device,
+                    pin_memory=True)
+                cu_num_tokens, token_indices = self.drafter.prepare_inputs(
+                    eagle_attn_metadata.query_start_loc,
+                    num_accepted_tokens_tensor,
+                )
+                target_token_ids = self.input_ids[token_indices]
+                # TODO(woosuk): Support M-RoPE.
+                target_positions = self.positions[token_indices]
+                if self.use_aux_hidden_state_outputs:
+                    target_hidden_states = torch.cat(
+                        [h[token_indices] for h in aux_hidden_states], dim=-1)
+                else:
+                    target_hidden_states = hidden_states[token_indices]
+                target_slot_mapping = eagle_attn_metadata.slot_mapping[
+                    token_indices]
+            draft_token_ids = self.drafter.propose(
+                target_token_ids=target_token_ids,
+                target_positions=target_positions,
+                target_hidden_states=target_hidden_states,
+                target_slot_mapping=target_slot_mapping,
+                next_token_ids=next_token_ids,
+                cu_num_tokens=cu_num_tokens,
+                block_table=block_table,
+                sampling_metadata=sampling_metadata,
+                decoding=spec_decode_metadata is not None
+            )
+            spec_token_ids = np.ones(draft_token_ids.shape, dtype=int).tolist()
+            self.last_draft_token_ids = draft_token_ids
+            self.last_draft_host_tokens = draft_token_ids.to('cpu', non_blocking=True)
+            self.last_draft_event.record()
+        return spec_token_ids
+
+    @torch.inference_mode()
+    def execute_model(
+        self,
+        scheduler_output: "SchedulerOutput",
+        intermediate_tensors: Optional[IntermediateTensors] = None,
+    ) -> Union[ModelRunnerOutput, IntermediateTensors]:
+        self._update_states(scheduler_output)
+        if not scheduler_output.total_num_scheduled_tokens:
+            if not has_kv_transfer_group():
+                # Return empty ModelRunnerOutput if there's no work to do.
+                return EMPTY_MODEL_RUNNER_OUTPUT
+
+            return self.kv_connector_no_forward(scheduler_output)
+
+        # Prepare the decoder inputs.
+        (attn_metadata, attention_cuda_graphs, logits_indices,
+         spec_decode_metadata,
+         num_scheduled_tokens_np) = (self._prepare_inputs(scheduler_output))
+        num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens
+        if (self.use_cuda_graph
+                and num_scheduled_tokens <= self.cudagraph_batch_sizes[-1]):
+            # Use piecewise CUDA graphs.
+            # Add padding to the batch size.
+            num_input_tokens = self.vllm_config.pad_for_cudagraph(
+                num_scheduled_tokens)
+        else:
+            # Eager mode.
+            # Pad tokens to multiple of tensor_parallel_size when
+            # enabled collective fusion for SP
+            tp_size = self.vllm_config.parallel_config.tensor_parallel_size
+            if self.compilation_config.pass_config. \
+                enable_sequence_parallelism and tp_size > 1:
+                num_input_tokens = round_up(num_scheduled_tokens, tp_size)
+            else:
+                num_input_tokens = num_scheduled_tokens
+
+        # Padding for DP
+        num_pad, num_tokens_across_dp = self.get_dp_padding(num_input_tokens)
+        num_input_tokens += num_pad
+
+        # _prepare_inputs may reorder the batch, so we must gather multi
+        # modal outputs after that to ensure the correct order
+        if self.is_multimodal_model:
+            # Run the multimodal encoder if any.
+            self._execute_mm_encoder(scheduler_output)
+            mm_embeds = self._gather_mm_embeddings(scheduler_output)
+        else:
+            mm_embeds = []
+
+        if self.is_multimodal_model and get_pp_group().is_first_rank:
+            # NOTE(woosuk): To unify token ids and soft tokens (vision
+            # embeddings), we always use embeddings (rather than token ids)
+            # as input to the multimodal model, even when the input is text.
+            input_ids = self.input_ids[:num_scheduled_tokens]
+            if mm_embeds:
+                inputs_embeds = self.model.get_input_embeddings(
+                    input_ids, mm_embeds)
+            else:
+                inputs_embeds = self.model.get_input_embeddings(input_ids)
+            # TODO(woosuk): Avoid the copy. Optimize.
+            self.inputs_embeds[:num_scheduled_tokens].copy_(inputs_embeds)
+            inputs_embeds = self.inputs_embeds[:num_input_tokens]
+            input_ids = None
+        else:
+            # For text-only models, we use token ids as input.
+            # While it is possible to use embeddings as input just like the
+            # multimodal models, it is not desirable for performance since
+            # then the embedding layer is not included in the CUDA graph.
+            input_ids = self.input_ids[:num_input_tokens]
+            inputs_embeds = None
+        if self.uses_mrope:
+            positions = self.mrope_positions[:, :num_input_tokens]
+        else:
+            positions = self.positions[:num_input_tokens]
+
+        if get_pp_group().is_first_rank:
+            intermediate_tensors = None
+        else:
+            intermediate_tensors = self.sync_and_slice_intermediate_tensors(
+                num_input_tokens, intermediate_tensors, True)
+
+        # Some attention backends only support CUDA Graphs in pure decode.
+        # If attention doesn't support CUDA Graphs for this batch, but we
+        # compiled with full CUDA graphs, we have to skip them entirely.
+        skip_cuda_graphs = self.full_cuda_graph and not attention_cuda_graphs
+
+        if envs.VLLM_ENABLE_TBO and not self.use_cuda_graph:
+            model_output, finished_sending, finished_recving = \
+                 tbo_split_and_execute_model(self, attn_metadata, num_input_tokens,
+                                             num_tokens_across_dp, input_ids, positions,
+                                             inputs_embeds, scheduler_output, intermediate_tensors)
+        else:
+            # Run the model.
+            # Use persistent buffers for CUDA graphs.
+            with set_forward_context(
+                    attn_metadata,
+                    self.vllm_config,
+                    num_tokens=num_input_tokens,
+                    num_tokens_across_dp=num_tokens_across_dp,
+                    skip_cuda_graphs=skip_cuda_graphs,
+            ):
+                self.maybe_setup_kv_connector(scheduler_output)
+
+                model_output = self.model(
+                    input_ids=input_ids,
+                    positions=positions,
+                    intermediate_tensors=intermediate_tensors,
+                    inputs_embeds=inputs_embeds,
+                )
+
+                self.maybe_wait_for_kv_save()
+                finished_sending, finished_recving = (
+                    self.get_finished_kv_transfers(scheduler_output))
+        if self.use_aux_hidden_state_outputs:
+            hidden_states, aux_hidden_states = model_output
+        else:
+            hidden_states = model_output
+            aux_hidden_states = None
+
+        # Broadcast PP output for external_launcher (torchrun)
+        # to make sure we are synced across pp ranks
+        # TODO: Support overlapping mirco-batches
+        # https://github.com/vllm-project/vllm/issues/18019
+        broadcast_pp_output = \
+            self.parallel_config.distributed_executor_backend \
+            == "external_launcher" and len(get_pp_group().ranks) > 0
+        if not get_pp_group().is_last_rank:
+            # For mid-pipeline stages, return the hidden states.
+            if not broadcast_pp_output:
+                return hidden_states
+            assert isinstance(hidden_states, IntermediateTensors)
+            get_pp_group().send_tensor_dict(hidden_states.tensors,
+                                            all_gather_group=get_tp_group())
+            logits = None
+        else:
+            if self.input_batch.pooling_params:
+                return self._pool(hidden_states, num_scheduled_tokens,
+                                  num_scheduled_tokens_np, finished_sending,
+                                  finished_recving)
+
+            sample_hidden_states = hidden_states[logits_indices]
+            logits = self.model.compute_logits(sample_hidden_states, None)
+        if broadcast_pp_output:
+            model_output_broadcast_data = {
+                "logits": logits.contiguous(),
+            } if logits is not None else {}
+            model_output_broadcast_data = get_pp_group().broadcast_tensor_dict(
+                model_output_broadcast_data, src=len(get_pp_group().ranks) - 1)
+            assert model_output_broadcast_data is not None
+            logits = model_output_broadcast_data["logits"]
+
+        # Apply structured output bitmasks if present
+        if scheduler_output.grammar_bitmask is not None:
+            self.apply_grammar_bitmask(scheduler_output, logits)
+
+        # Sample the next token and get logprobs if needed.
+        sampling_metadata = self.input_batch.sampling_metadata
+        if spec_decode_metadata is None:
+            sampler_output = self.sampler(
+                logits=logits,
+                sampling_metadata=sampling_metadata,
+            )
+        else:
+            # When indexing with a tensor (bonus_logits_indices), PyTorch
+            # creates a new tensor with separate storage from the original
+            # logits tensor. This means any in-place operations on bonus_logits
+            # won't affect the original logits tensor.
+            assert logits is not None
+            bonus_logits = logits[spec_decode_metadata.bonus_logits_indices]
+            sampler_output = self.sampler(
+                logits=bonus_logits,
+                sampling_metadata=sampling_metadata,
+            )
+            bonus_token_ids = sampler_output.sampled_token_ids
+
+            # Just like `bonus_logits`, `target_logits` is a new tensor with
+            # separate storage from the original `logits` tensor. Therefore,
+            # it is safe to update `target_logits` in place.
+            target_logits = logits[spec_decode_metadata.target_logits_indices]
+            output_token_ids = self.rejection_sampler(
+                spec_decode_metadata,
+                None,  # draft_probs
+                target_logits,
+                bonus_token_ids,
+                sampling_metadata,
+            )
+            sampler_output.sampled_token_ids = output_token_ids
+
+        num_nans_in_logits = {}
+        if envs.VLLM_COMPUTE_NANS_IN_LOGITS:
+            num_nans_in_logits = self._get_nans_in_logits(logits)
+
+        # TODO(woosuk): The following loop can be slow since it iterates over
+        # the requests one by one. Optimize.
+        discard_sampled_tokens_req_indices = []
+        for i, req_id in enumerate(self.input_batch.req_ids):
+            req_state = self.requests[req_id]
+            seq_len = (req_state.num_computed_tokens +
+                       scheduler_output.num_scheduled_tokens[req_id])
+            if seq_len < req_state.num_tokens:
+                # Ignore the sampled token for partial prefills.
+                # Rewind the generator state as if the token was not sampled.
+                # This relies on cuda-specific torch-internal impl details
+                generator = self.input_batch.generators.get(i)
+                if generator is not None:
+                    generator.set_offset(generator.get_offset() - 4)
+                # Record the index of the request that should not be sampled,
+                # so that we could clear the sampled tokens before returning.
+                discard_sampled_tokens_req_indices.append(i)
+
+        # NOTE: GPU -> CPU Sync happens here.
+        # Move as many CPU operations as possible before this sync point.
+        logprobs_tensors = sampler_output.logprobs_tensors
+        logprobs_lists = logprobs_tensors.tolists() \
+            if logprobs_tensors is not None else None
+
+        # Compute prompt logprobs if needed.
+        prompt_logprobs_dict = self._get_prompt_logprobs_dict(
+            hidden_states[:num_scheduled_tokens],
             scheduler_output,
-            valid_sampled_token_ids,
-            sampling_metadata,
-            hidden_states,
-            sample_hidden_states,
-            aux_hidden_states,
-            spec_decode_metadata,
-            attn_metadata,
         )
 
-    # Clear KVConnector state after all KVs are generated.
-    if has_kv_transfer_group():
-        get_kv_transfer_group().clear_connector_metadata()
-
-    runner.eplb_step()
-
-    model_output = ZeroV1ModelRunnerOutput(
-        req_ids=runner.input_batch.req_ids,
-        req_id_to_index=runner.input_batch.req_id_to_index,
-        sampled_token_ids=valid_sampled_token_ids,
-        spec_token_ids=spec_token_ids,
-        logprobs=logprobs_lists,
-        prompt_logprobs_dict=prompt_logprobs_dict,
-        pooler_output=[],
-        finished_sending=finished_sending,
-        finished_recving=finished_recving,
-        num_nans_in_logits=num_nans_in_logits,
-        fix_req_ids = fix_req_ids,
-        fix_sampled_token_ids = fix_sampled_token_ids
-    )
-    return model_output
\ No newline at end of file
+        # Get the valid generated tokens.
+        sampled_token_ids = sampler_output.sampled_token_ids
+        max_gen_len = sampled_token_ids.shape[-1]
+
+        fix_req_ids = None
+        fix_sampled_token_ids = None
+        fix_draft_token_ids = None
+        fix_draft_req_ids = self.last_sampled_req_ids
+        is_output_valid = False
+        if self.speculative_config:
+            if max_gen_len == 1:
+                valid_sampled_token_ids = sampled_token_ids.tolist()
+            else:
+                # Includes spec decode tokens.
+                valid_sampled_token_ids = self.rejection_sampler.parse_output(
+                    sampled_token_ids,
+                    self.input_batch.vocab_size,
+                )
+                self.last_sampler_host_tokens = None
+                self.last_sampled_token_ids = None
+            is_output_valid = True
+        else:
+            # No spec decode tokens.
+            fix_req_ids = self.last_sampled_req_ids
+            if self.last_sampler_host_tokens != None:
+                self.last_sampler_event.synchronize()
+                fix_sampled_token_ids = self.last_sampler_host_tokens.tolist()
+                for req_idx, start_idx, end_idx in self.token_ids_cpu_fix_recode:
+                    self.input_batch.token_ids_cpu[req_idx, start_idx:end_idx] = fix_sampled_token_ids[req_idx]
+            for req_idx, req_id in enumerate(fix_req_ids):
+                if req_id in self.requests:
+                    req_state = self.requests[req_id]
+                    token_idx = self.last_sampled_token_lens[req_idx]
+                    req_state.output_token_ids[token_idx] = fix_sampled_token_ids[req_idx][0]
+            self.last_sampler_host_tokens = sampled_token_ids.to('cpu', non_blocking=True)
+            self.last_sampler_event.record()
+            self.last_sampled_token_ids = sampled_token_ids
+            valid_sampled_token_ids = np.ones(sampled_token_ids.shape, dtype=int).tolist()
+            
+        # Mask out the sampled tokens that should not be sampled.
+        for i in discard_sampled_tokens_req_indices:
+            valid_sampled_token_ids[i].clear()
+
+        # Cache the sampled tokens in the model runner, so that the scheduler
+        # doesn't need to send them back.
+        # NOTE(woosuk): As an exception, when using PP, the scheduler sends
+        # the sampled tokens back, because there's no direct communication
+        # between the first-stage worker and the last-stage worker.
+        self.token_ids_cpu_fix_recode.clear()
+        self.last_sampled_req_ids = []
+        self.last_sampled_token_lens = []
+        for req_idx, sampled_ids in enumerate(valid_sampled_token_ids):
+            if not sampled_ids:
+                continue
+
+            start_idx = self.input_batch.num_tokens_no_spec[req_idx]
+            end_idx = start_idx + len(sampled_ids)
+            assert end_idx <= self.max_model_len, (
+                "Sampled token IDs exceed the max model length. "
+                f"Total number of tokens: {end_idx} > max_model_len: "
+                f"{self.max_model_len}")
+
+            self.input_batch.token_ids_cpu[req_idx,
+                                            start_idx:end_idx] = sampled_ids
+            self.token_ids_cpu_fix_recode.append([req_idx, start_idx, end_idx])
+            self.input_batch.num_tokens_no_spec[req_idx] = end_idx
+            self.input_batch.num_tokens[req_idx] = end_idx
+            req_id = self.input_batch.req_ids[req_idx]
+            if req_id in self.requests:
+                req_state = self.requests[req_id]
+                self.last_sampled_req_ids.append(req_id)
+                self.last_sampled_token_lens.append(len(req_state.output_token_ids))
+                req_state.output_token_ids.extend(sampled_ids)
+                
+        if not self.speculative_config:
+            # Speculative decoding is not enabled.
+            spec_token_ids = None
+            fix_draft_req_ids = None
+        else:
+            if self.last_draft_host_tokens is not None:
+                self.last_draft_event.synchronize()
+                fix_draft_token_ids = self.last_draft_host_tokens.tolist()
+            spec_token_ids = self.propose_draft_token_ids(
+                scheduler_output,
+                valid_sampled_token_ids,
+                sampling_metadata,
+                hidden_states,
+                sample_hidden_states,
+                aux_hidden_states,
+                spec_decode_metadata,
+                attn_metadata,
+            )
+
+        # Clear KVConnector state after all KVs are generated.
+        if has_kv_transfer_group():
+            get_kv_transfer_group().clear_connector_metadata()
+
+        self.eplb_step()
+
+        model_output = ZeroV1ModelRunnerOutput(
+            req_ids=self.input_batch.req_ids,
+            req_id_to_index=self.input_batch.req_id_to_index,
+            sampled_token_ids=valid_sampled_token_ids,
+            spec_token_ids=spec_token_ids,
+            logprobs=logprobs_lists,
+            prompt_logprobs_dict=prompt_logprobs_dict,
+            pooler_output=[],
+            finished_sending=finished_sending,
+            finished_recving=finished_recving,
+            num_nans_in_logits=num_nans_in_logits,
+            fix_req_ids = fix_req_ids,
+            fix_sampled_token_ids = fix_sampled_token_ids,
+            fix_draft_tokens_ids = fix_draft_token_ids, 
+            fix_draft_req_ids = fix_draft_req_ids,
+            is_output_valid=is_output_valid
+        )
+        return model_output
\ No newline at end of file

vllm/zero_overhead/v1/outputs.py

@@ -8,4 +8,7 @@ from vllm.v1.outputs import ModelRunnerOutput
 class ZeroV1ModelRunnerOutput(ModelRunnerOutput):
     # [num_reqs]
     fix_req_ids: list[str] = None
-    fix_sampled_token_ids:list[list[int]] = None
\ No newline at end of file
+    fix_sampled_token_ids:list[list[int]] = None
+    fix_draft_req_ids:list[list[int]] = None
+    fix_draft_tokens_ids:list[list[int]] = None
+    is_output_valid:bool = True
\ No newline at end of file