[ROCm][CI] Fix Kimi K3 KDA on ROCm (#50262)

Signed-off-by: Stefan Koncarevic <stefan.koncarevic@amd.com>

.buildkite/test-amd.yaml

@@ -1576,10 +1576,14 @@ steps:
   - vllm/third_party/flash_linear_attention/ops/kda.py
   - vllm/third_party/flash_linear_attention/ops/chunk_delta_h.py
   - vllm/third_party/flash_linear_attention/ops/l2norm.py
-  - tests/kernels/test_kda.py
+  - vllm/models/kimi_k3/nvidia/kda.py
+  - vllm/models/kimi_k3/nvidia/kda_metadata.py
+  - vllm/models/kimi_k3/nvidia/ops/third_party/kda/
+  - tests/models/kimi_k3/test_kda.py
+  - tests/models/kimi_k3/test_kda_metadata.py
   - vllm/platforms/rocm.py
   commands:
-  - pytest -v -s kernels/test_kda.py
+  - pytest -v -s models/kimi_k3/test_kda.py models/kimi_k3/test_kda_metadata.py
 
 - label: Kernels Mamba Test # TBD
   timeout_in_minutes: 180

tests/v1/sample/test_head_dtype.py

@@ -14,6 +14,7 @@ import torch
 from vllm import LLM, SamplingParams
 from vllm.model_executor.layers.logits_processor import LogitsProcessor
 from vllm.model_executor.layers.vocab_parallel_embedding import (
+    ParallelLMHead,
     UnquantizedEmbeddingMethod,
 )
 
@@ -28,6 +29,7 @@ class _FakeLmHead:
         self.weight = weight
         self.quant_method = object() if quantized else UnquantizedEmbeddingMethod()
         self.shard_indices = shard_indices
+        self.tp_size = 1
 
 
 def _build_processor(vocab_size: int) -> LogitsProcessor:
@@ -135,11 +137,59 @@ def test_fp32_head_rejects_quantized_lm_head(default_vllm_config):
         lp._get_logits(torch.randn(4, 16, dtype=torch.bfloat16), lm_head, None)
 
 
+def test_replicated_lm_head_skips_tp_communication_and_preserves_processing(
+    default_vllm_config,
+):
+    from unittest import mock
+
+    vocab_size, hidden_size = 12, 8
+    soft_cap, scale = 2.0, 0.5
+    lp = LogitsProcessor(
+        vocab_size,
+        soft_cap=soft_cap,
+        scale=scale,
+    )
+    lp.head_dtype = torch.float32
+
+    hidden_states = torch.randn(4, hidden_size, dtype=torch.bfloat16)
+    weight = torch.randn(vocab_size, hidden_size, dtype=torch.bfloat16)
+    world_size_getter = (
+        "vllm.model_executor.layers.vocab_parallel_embedding."
+        "get_tensor_model_parallel_world_size"
+    )
+    with mock.patch(world_size_getter, return_value=2):
+        lm_head = ParallelLMHead(
+            vocab_size,
+            hidden_size,
+            params_dtype=torch.bfloat16,
+            disable_tp=True,
+        )
+    lm_head.weight_loader(lm_head.weight, weight)
+    assert lm_head.tp_size == 1
+
+    with mock.patch.object(lp, "_gather_logits") as gather_mock:
+        logits = lp(lm_head, hidden_states)
+
+    gather_mock.assert_not_called()
+
+    expected = torch.nn.functional.linear(hidden_states.float(), weight.float())
+    expected = torch.tanh(expected / soft_cap) * soft_cap * scale
+    torch.testing.assert_close(logits, expected)
+
+    all_gather_path = (
+        "vllm.model_executor.layers.logits_processor.tensor_model_parallel_all_gather"
+    )
+    with mock.patch(all_gather_path) as all_gather:
+        top = lp.get_top_tokens(lm_head, hidden_states)
+
+    all_gather.assert_not_called()
+    assert torch.equal(top, expected.argmax(dim=-1))
+
+
 def test_get_top_tokens_honors_head_dtype(default_vllm_config):
     # The spec-decode local-argmax path (get_top_tokens) must run the lm_head
     # in head_dtype too, not just _get_logits.
     import types
-    from unittest import mock
 
     vocab_size, hidden_size = 64, 16
     lp = _build_processor(vocab_size)
@@ -154,13 +204,7 @@ def test_get_top_tokens_honors_head_dtype(default_vllm_config):
         ),
     )
 
-    with mock.patch(
-        "vllm.model_executor.layers.logits_processor."
-        "get_tensor_model_parallel_world_size",
-        return_value=1,
-    ):
-        top = lp.get_top_tokens(lm_head, hidden_states, None)
-
+    top = lp.get_top_tokens(lm_head, hidden_states, None)
     expected = torch.nn.functional.linear(hidden_states.float(), weight.float()).argmax(
         dim=-1
     )

vllm/compilation/passes/fx_utils.py

@@ -14,6 +14,10 @@ def is_func(node: fx.Node, target: Target) -> bool:
     return bool(node.op == "call_function" and node.target == target)
 
 
+def is_auto_func(node: fx.Node, op: OpOverload) -> bool:
+    return is_func(node, auto_functionalized) and node.args[0] == op
+
+
 # Returns the first auto_functionalized node with the given op (if it exists)
 def find_auto_fn_maybe(nodes: Iterable[fx.Node], op: OpOverload) -> fx.Node | None:
     for node in nodes:
@@ -38,6 +42,13 @@ def find_getitem_maybe(node: fx.Node, idx: int) -> fx.Node | None:
     return None
 
 
+# Returns the getitem node that extracts the idx-th element from node
+def find_getitem(node: fx.Node, idx: int) -> fx.Node:
+    ret = find_getitem_maybe(node, idx)
+    assert ret is not None, f"Could not find getitem {idx} in node {node}"
+    return ret
+
+
 # An auto-functionalization-aware utility for finding nodes with a specific op
 # Also handles op overload packets and finds all overloads
 def find_op_nodes(
@@ -56,3 +67,11 @@ def find_op_nodes(
     for n in graph.find_nodes(op="call_function", target=auto_functionalized):
         if n.args[0] == op:
             yield n
+
+
+# Asserts that the node only has one user and returns it
+# Even if a node has only 1 user, it might share storage with another node,
+# which might need to be taken into account.
+def get_only_user(node: fx.Node) -> fx.Node:
+    assert len(node.users) == 1
+    return next(iter(node.users))

vllm/compilation/passes/vllm_inductor_pass.py

@@ -333,3 +333,12 @@ class VllmFusionPatternMatcherPass(VllmPatternMatcherPass):
     def __call__(self, graph: torch.fx.Graph) -> None:
         self.matched_count = self.pm_pass.apply(graph)
         VllmPatternMatcherPass.match_table[self.pass_name] += self.matched_count
+
+
+class PrinterInductorPass(VllmInductorPass):
+    def __init__(self, name: str, config: VllmConfig) -> None:
+        super().__init__(config)
+        self.name = name
+
+    def __call__(self, graph: torch.fx.Graph) -> None:
+        self.dump_graph(graph, self.name)

vllm/config/utils.py

@@ -415,6 +415,31 @@ class Range:
         return self.__str__()
 
 
+def handle_deprecated(
+    config: ConfigT,
+    old_name: str,
+    new_name_or_names: str | list[str],
+    removal_version: str,
+) -> None:
+    old_val = getattr(config, old_name)
+    if old_val is None:
+        return
+
+    if isinstance(new_name_or_names, str):
+        new_names = [new_name_or_names]
+    else:
+        new_names = new_name_or_names
+
+    msg = (
+        f"{old_name} is deprecated and will be removed in {removal_version}. "
+        f"Use {', '.join(new_names)} instead."
+    )
+    logger.warning(msg)
+
+    for new_name in new_names:
+        setattr(config, new_name, old_val)
+
+
 def get_from_deprecated_env_if_set(
     env_name: str,
     removal_version: str,

vllm/distributed/parallel_state.py

@@ -2081,6 +2081,9 @@ def model_parallel_is_initialized():
     return _TP is not None and _PP is not None
 
 
+_TP_STATE_PATCHED = False
+
+
 def get_tensor_model_parallel_world_size() -> int:
     """Return world size for the tensor model parallel group."""
     return get_tp_group().world_size

vllm/entrypoints/chat_utils.py

@@ -2015,6 +2015,15 @@ async def parse_chat_messages_async(
     return conversation, mm_data, mm_uuids
 
 
+def get_history_tool_calls_cnt(conversation: list[ConversationMessage]):
+    idx = 0
+    for msg in conversation:
+        if msg["role"] == "assistant":
+            tool_calls = msg.get("tool_calls")
+            idx += len(list(tool_calls)) if tool_calls is not None else 0  # noqa
+    return idx
+
+
 _KIMI_MODEL_TYPES = ("kimi_k2", "kimi_k25", "kimi_k3")
 
 

vllm/model_executor/layers/activation.py

@@ -8,8 +8,14 @@ import torch
 import torch.nn as nn
 import torch.nn.functional as F
 
+from vllm.distributed import (
+    divide,
+    get_tensor_model_parallel_rank,
+    get_tensor_model_parallel_world_size,
+)
 from vllm.logger import init_logger
 from vllm.model_executor.custom_op import CustomOp
+from vllm.model_executor.utils import set_weight_attrs
 from vllm.platforms import CpuArchEnum, current_platform
 from vllm.triton_utils import tl, triton
 from vllm.utils.collection_utils import LazyDict
@@ -782,6 +788,48 @@ class XIELU(CustomOp):
         return self.forward_native(input)
 
 
+class ScaledActivation(nn.Module):
+    """An activation function with post-scale parameters.
+
+    This is used for some quantization methods like AWQ.
+    """
+
+    def __init__(
+        self,
+        act_module: nn.Module,
+        intermediate_size: int,
+        input_is_parallel: bool = True,
+        params_dtype: torch.dtype | None = None,
+    ):
+        super().__init__()
+        self.act = act_module
+        self.input_is_parallel = input_is_parallel
+        if input_is_parallel:
+            tp_size = get_tensor_model_parallel_world_size()
+            intermediate_size_per_partition = divide(intermediate_size, tp_size)
+        else:
+            intermediate_size_per_partition = intermediate_size
+        if params_dtype is None:
+            params_dtype = torch.get_default_dtype()
+        self.scales = nn.Parameter(
+            torch.empty(intermediate_size_per_partition, dtype=params_dtype)
+        )
+        set_weight_attrs(self.scales, {"weight_loader": self.weight_loader})
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.act(x) / self.scales
+
+    def weight_loader(self, param: nn.Parameter, loaded_weight: torch.Tensor):
+        param_data = param.data
+        if self.input_is_parallel:
+            tp_rank = get_tensor_model_parallel_rank()
+            shard_size = param_data.shape[0]
+            start_idx = tp_rank * shard_size
+            loaded_weight = loaded_weight.narrow(0, start_idx, shard_size)
+        assert param_data.shape == loaded_weight.shape
+        param_data.copy_(loaded_weight)
+
+
 _ACTIVATION_REGISTRY = LazyDict(
     {
         "gelu": lambda: GELU(),

vllm/model_executor/layers/conv.py

@@ -155,6 +155,58 @@ class Conv2dLayer(ConvLayerBase):
         return self._forward_conv(x)
 
 
+class CausalConv2dLayer(Conv2dLayer):
+    """
+    A causal version of nn.Conv2d where each location in the 2D matrix would
+    have no access to locations on its right or down
+    All arguments are the same as nn.Conv2d except padding which should be
+    set as None
+    """
+
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        kernel_size: int,
+        stride: int,
+        padding: int = 0,
+        dilation: int = 1,
+        groups: int = 1,
+        bias: bool = True,
+        padding_mode: str = "zeros",
+        *,
+        params_dtype: torch.dtype | None = None,
+    ) -> None:
+        if padding is not None:
+            raise ValueError(
+                "Argument padding should be set to None for CausalConv2dLayer."
+            )
+        self._left_padding: int = kernel_size - 1
+        self._right_padding: int = stride - 1
+        padding = 0
+
+        super().__init__(
+            in_channels,
+            out_channels,
+            kernel_size,
+            stride,
+            padding,
+            dilation,
+            groups,
+            bias,
+            padding_mode,
+            params_dtype=params_dtype,
+        )
+
+    def forward(
+        self,
+        x: torch.Tensor,
+    ) -> torch.Tensor:
+        x = F.pad(x, pad=(self._left_padding, self._right_padding, 0, 0))
+        x = super().forward(x)
+        return x
+
+
 # --8<-- [start:conv3d]
 @CustomOp.register("conv3d")
 class Conv3dLayer(ConvLayerBase):

vllm/model_executor/layers/logits_processor.py

@@ -7,7 +7,6 @@ import torch.nn.functional as F
 
 from vllm.config import get_current_vllm_config
 from vllm.distributed import (
-    get_tensor_model_parallel_world_size,
     tensor_model_parallel_all_gather,
     tensor_model_parallel_gather,
 )
@@ -145,7 +144,8 @@ class LogitsProcessor(PluggableLayer):
         logits = self._apply_head(lm_head, hidden_states, embedding_bias)
 
         # Gather logits for TP
-        logits = self._gather_logits(logits)
+        if lm_head.tp_size > 1:
+            logits = self._gather_logits(logits)
 
         # Remove paddings in vocab (if any).
         if logits is not None:
@@ -169,7 +169,7 @@ class LogitsProcessor(PluggableLayer):
                 "The local argmax reduction optimization is not supported for "
                 "non-positive logit scaling factors."
             )
-        tp_size = get_tensor_model_parallel_world_size()
+        tp_size = lm_head.tp_size
 
         logits = self._apply_head(lm_head, hidden_states, embedding_bias)
         if self.soft_cap is not None:

vllm/model_executor/layers/quantization/fp8.py

@@ -4,6 +4,7 @@
 from typing import TYPE_CHECKING, Any
 
 import torch
+from torch.utils._python_dispatch import TorchDispatchMode
 
 import vllm.envs as envs
 import vllm.model_executor.layers.fused_moe.modular_kernel as mk
@@ -233,6 +234,26 @@ class Fp8Config(QuantizationConfig):
         return cache_scale_mapper | QuantizationConfig.get_cache_scale_mapper()
 
 
+class CopyNumelCounter(TorchDispatchMode):
+    """
+    Tracks total number of elements modified with `copy_`. Useful for keeping
+    track of weight loading where underlying weights can be arbitrarily
+    transformed (such as with `narrow`) before calling copy.
+    """
+
+    def __init__(self):
+        super().__init__()
+        self.copied_numel = 0
+
+    def __torch_dispatch__(self, func, types, args=(), kwargs=None):
+        if kwargs is None:
+            kwargs = {}
+        out = func(*args, **kwargs)
+        if func == torch.ops.aten.copy_.default:
+            self.copied_numel += args[0].numel()
+        return out
+
+
 class Fp8LinearMethod(LinearMethodBase):
     """Linear method for FP8.
     Supports loading FP8 checkpoints with static weight scale and

vllm/model_executor/layers/utils.py

@@ -25,6 +25,12 @@ MOE_LAYER_ROUTER_GATE_SUFFIXES = {
 }
 
 
+def is_layer_moe_router_gate(prefix: str) -> bool:
+    if not prefix:
+        return False
+    return prefix.rsplit(".", 1)[-1] in MOE_LAYER_ROUTER_GATE_SUFFIXES
+
+
 def get_token_bin_counts_and_mask(
     tokens: torch.Tensor,
     vocab_size: int,

vllm/model_executor/layers/vocab_parallel_embedding.py

@@ -232,6 +232,7 @@ class VocabParallelEmbedding(PluggableLayer):
         padding_size: padding size for the vocabulary.
         quant_config: quant config for the layer
         prefix: full name of the layer in the state dict
+        disable_tp: If true, tensor parallelism will be disabled for this layer.
     """  # noqa: E501
 
     # --8<-- [end:vocab_parallel_embedding]
@@ -245,12 +246,19 @@ class VocabParallelEmbedding(PluggableLayer):
         padding_size: int = DEFAULT_VOCAB_PADDING_SIZE,
         quant_config: QuantizationConfig | None = None,
         prefix: str = "",
+        *,
+        disable_tp: bool = False,
     ):
         super().__init__()
 
         # Keep the input dimensions.
-        tp_rank = get_tensor_model_parallel_rank()
-        self.tp_size = get_tensor_model_parallel_world_size()
+        self.disable_tp = disable_tp
+        if disable_tp:
+            tp_rank, self.tp_size = 0, 1
+        else:
+            tp_rank = get_tensor_model_parallel_rank()
+            self.tp_size = get_tensor_model_parallel_world_size()
+        self.tp_rank = tp_rank
         self.num_embeddings = num_embeddings
         self.padding_size = padding_size
         self.org_vocab_size = org_num_embeddings or num_embeddings
@@ -323,6 +331,13 @@ class VocabParallelEmbedding(PluggableLayer):
             params_dtype=params_dtype,
             weight_loader=self.weight_loader,
         )
+        self.update_param_tp_status()
+
+    def update_param_tp_status(self):
+        for param in self.parameters():
+            if isinstance(param, BasevLLMParameter):
+                param.tp_rank = self.tp_rank
+                param.tp_size = self.tp_size
 
     @classmethod
     def _get_indices(
@@ -487,9 +502,9 @@ class VocabParallelEmbedding(PluggableLayer):
         # Mask the output embedding.
         if self.tp_size > 1:
             output_parallel.masked_fill_(input_mask.unsqueeze(-1), 0)
-        # Reduce across all the model parallel GPUs.
-        output = tensor_model_parallel_all_reduce(output_parallel)
-        return output
+            # Reduce across all the model parallel GPUs.
+            return tensor_model_parallel_all_reduce(output_parallel)
+        return output_parallel
 
     def extra_repr(self) -> str:
         s = f"num_embeddings={self.num_embeddings_per_partition}"
@@ -516,6 +531,7 @@ class ParallelLMHead(VocabParallelEmbedding):
         params_dtype: type of the parameters.
         org_num_embeddings: original vocabulary size (without LoRA).
         padding_size: padding size for the vocabulary.
+        disable_tp: If true, tensor parallelism will be disabled for this layer.
     """
 
     # --8<-- [end:parallel_lm_head]
@@ -530,6 +546,8 @@ class ParallelLMHead(VocabParallelEmbedding):
         padding_size: int = DEFAULT_VOCAB_PADDING_SIZE,
         quant_config: QuantizationConfig | None = None,
         prefix: str = "",
+        *,
+        disable_tp: bool = False,
     ):
         super().__init__(
             num_embeddings,
@@ -539,6 +557,7 @@ class ParallelLMHead(VocabParallelEmbedding):
             padding_size,
             quant_config,
             prefix,
+            disable_tp=disable_tp,
         )
         self.quant_config = quant_config
         if bias:

vllm/model_executor/models/glm_ocr.py

@@ -242,6 +242,10 @@ class GlmOcrVisionBlock(Glm4vVisionBlock):
         )
 
 
+class GlmOcrVisionPatchEmbed(Glm4vVisionPatchEmbed):
+    pass
+
+
 class GlmOcrPatchMerger(Glm4vPatchMerger):
     pass
 

vllm/model_executor/models/idefics3.py

@@ -103,6 +103,41 @@ class Idefics3ProcessingInfo(BaseProcessingInfo):
     def get_supported_mm_limits(self) -> Mapping[str, int | None]:
         return {"image": None}
 
+    def _resize_output_size(
+        self,
+        *,
+        height: int,
+        width: int,
+        max_len: int | None = None,
+        min_len: int = 1,
+        max_size: int | None = None,
+    ) -> tuple[int, int]:
+        # Set default value for max_len if not provided
+        max_len = max(height, width) if max_len is None else max_len
+        aspect_ratio = width / height
+
+        # Handle the maximum size constraint
+        if max_size is not None:
+            max_len = min(max_len, max_size)
+
+        # Adjust dimensions according to the aspect ratio
+        if width >= height:
+            width = max_len
+            height = int(width / aspect_ratio)
+        else:
+            height = max_len
+            width = int(height * aspect_ratio)
+
+        # Ensure both width and height are even (if needed)
+        height += height % 2
+        width += width % 2
+
+        # Ensure dimensions are not smaller than the minimum length
+        height = max(height, min_len)
+        width = max(width, min_len)
+
+        return height, width
+
     def _get_image_feature_grid_size(
         self,
         *,

vllm/model_executor/models/llava_onevision2.py

@@ -1531,6 +1531,13 @@ class LlavaOnevision2MultiModalDataParser(MultiModalDataParser):
 class LlavaOnevision2MultiModalProcessor(
     BaseMultiModalProcessor[LlavaOnevision2ProcessingInfo]
 ):
+    def _get_data_parser(self) -> MultiModalDataParser:
+        # Retained for symmetry; vLLM actually fetches the parser via
+        # info.get_data_parser() (see ProcessingInfo override above).
+        return LlavaOnevision2MultiModalDataParser(
+            self.info.get_hf_config().vision_config.spatial_merge_size
+        )
+
     def _call_hf_processor(
         self,
         prompt: str,

vllm/model_executor/models/nemotron_parse.py

@@ -77,6 +77,24 @@ class BartScaledWordEmbedding(VocabParallelEmbedding):
         return super().forward(input_ids) * self.embed_scale
 
 
+class BartParallelLMHead(ParallelLMHead):
+    """
+    This module overrides ParallelLMHead's
+    forward by dividing by embeddings scale,
+    yielding effectively the inverse of
+    BartScaledWordEmbedding
+    """
+
+    def __init__(
+        self, num_embeddings: int, embedding_dim: int, embed_scale: float = 1.0
+    ):
+        super().__init__(num_embeddings, embedding_dim)
+        self.embed_scale = embed_scale
+
+    def forward(self, input_ids: torch.Tensor) -> torch.Tensor:
+        return super().forward(input_ids) / self.embed_scale
+
+
 class BartDecoderLayer(nn.Module):
     def __init__(
         self,

vllm/model_executor/models/phi4mm.py

@@ -521,6 +521,32 @@ class Phi4MMAudioEmbeddingInputs(TensorSchema):
 Phi4MMAudioInputs: TypeAlias = Phi4MMAudioFeatureInputs | Phi4MMAudioEmbeddingInputs
 
 
+def cat_with_pad(tensors, dim, padding_value=0):
+    """
+    cat along dim, while pad to max for all other dims
+    """
+    ndim = tensors[0].dim()
+    assert all(t.dim() == ndim for t in tensors[1:]), (
+        "All tensors must have the same number of dimensions"
+    )
+
+    out_size = [max(t.shape[i] for t in tensors) for i in range(ndim)]
+    out_size[dim] = sum(t.shape[dim] for t in tensors)
+    output = tensors[0].new_full(out_size, padding_value)
+
+    index = 0
+    for t in tensors:
+        # Create a slice list where every dimension except dim is full slice
+        slices = [slice(0, t.shape[d]) for d in range(ndim)]
+        # Update only the concat dimension slice
+        slices[dim] = slice(index, index + t.shape[dim])
+
+        output[slices] = t
+        index += t.shape[dim]
+
+    return output
+
+
 def stack_with_pad(
     tensors: torch.Tensor | list[torch.Tensor],
     padding_value: int | float = 0,

vllm/model_executor/models/phi4mm_utils.py

@@ -1595,6 +1595,14 @@ class AttModule(nn.Module):
         return x, memory, pos_emb, att_mask
 
 
+class AttBlock(BlockBase, AttModule):
+    """Attention Block module to support both Attention and Block module."""
+
+    def memory_dims(self, max_len: bool = False) -> tuple[int, int]:
+        """memory dimensions"""
+        return (1, self.input_size)
+
+
 def masked_softmax(
     scores: Tensor,
     mask: Tensor | None,

vllm/model_executor/models/qwen3_dspark.py

@@ -24,7 +24,6 @@ from vllm.logger import init_logger
 from vllm.model_executor.layers.logits_processor import LogitsProcessor
 from vllm.model_executor.layers.vocab_parallel_embedding import (
     ParallelLMHead,
-    VocabParallelEmbedding,
 )
 
 from .qwen3_dflash import DFlashQwen3ForCausalLM, DFlashQwen3Model
@@ -40,6 +39,10 @@ class DSparkMarkovHead(nn.Module):
     ``vocab_size``); ``markov_w2`` projects it to a draft-vocab bias
     (``draft_vocab_size``) added to the base draft logits. The two sizes
     coincide for full-vocab drafts.
+
+    Both weights are replicated because the head runs sequentially for every
+    draft position. Sharding them would add an all-reduce and a full-vocab
+    gather to each position.
     """
 
     def __init__(
@@ -50,19 +53,24 @@ class DSparkMarkovHead(nn.Module):
         prefix: str,
     ) -> None:
         super().__init__()
-        # TODO(ben): profile for which (if any) it makes sense to replicate or TP-shard
-        self.markov_w1 = VocabParallelEmbedding(
-            vocab_size, markov_rank, prefix=maybe_prefix(prefix, "markov_w1")
-        )
+        self.markov_w1 = nn.Embedding(vocab_size, markov_rank)
         self.markov_w2 = ParallelLMHead(
-            draft_vocab_size, markov_rank, prefix=maybe_prefix(prefix, "markov_w2")
+            draft_vocab_size,
+            markov_rank,
+            bias=False,
+            prefix=maybe_prefix(prefix, "markov_w2"),
+            disable_tp=True,
         )
 
     def embed(self, token_ids: torch.Tensor) -> torch.Tensor:
         """r-dim Markov embedding of ``token_ids`` ([B] -> [B, r])."""
         return self.markov_w1(token_ids)
 
-    def bias(self, markov_embed: torch.Tensor, logits_processor) -> torch.Tensor:
+    def bias(
+        self,
+        markov_embed: torch.Tensor,
+        logits_processor: LogitsProcessor,
+    ) -> torch.Tensor:
         """Vocab-size transition bias from a Markov embedding ([B, r] -> [B, V])."""
         return logits_processor(self.markov_w2, markov_embed)
 

vllm/model_executor/models/skyworkr1v.py

@@ -7,7 +7,7 @@
 # Copyright (c) 2025 Skywork
 # Licensed under The MIT License [see LICENSE for details]
 # --------------------------------------------------------
-from collections.abc import Iterable
+from collections.abc import Iterable, Mapping
 from typing import Annotated, Literal, TypeAlias
 
 import torch
@@ -15,6 +15,8 @@ import torch.nn as nn
 from transformers import PretrainedConfig
 
 from vllm.config import VllmConfig
+from vllm.config.multimodal import BaseDummyOptions
+from vllm.inputs import MultiModalDataDict
 from vllm.model_executor.layers.linear import ReplicatedLinear
 from vllm.model_executor.layers.quantization import QuantizationConfig
 from vllm.model_executor.layers.quantization.auto_awq import AutoAWQConfig
@@ -22,6 +24,7 @@ from vllm.model_executor.models.intern_vit import (
     InternVisionModel,
 )
 from vllm.multimodal import MULTIMODAL_REGISTRY
+from vllm.multimodal.processing import BaseDummyInputsBuilder
 from vllm.sequence import IntermediateTensors
 from vllm.transformers_utils.processors.internvl import (
     InternVLImageProcessor,
@@ -114,6 +117,33 @@ class SkyworkR1VProcessingInfo(BaseInternVLProcessingInfo):
         )
 
 
+class SkyworkR1VDummyInputsBuilder(BaseDummyInputsBuilder[SkyworkR1VProcessingInfo]):
+    def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str:
+        num_images = mm_counts.get("image", 0)
+
+        return "<image>" * num_images
+
+    def get_dummy_mm_data(
+        self,
+        seq_len: int,
+        mm_counts: Mapping[str, int],
+        mm_options: Mapping[str, BaseDummyOptions],
+    ) -> MultiModalDataDict:
+        target_width, target_height = self.info.get_image_size_with_most_features()
+        num_images = mm_counts.get("image", 0)
+
+        image_overrides = mm_options.get("image")
+
+        return {
+            "image": self._get_dummy_images(
+                width=target_width,
+                height=target_height,
+                num_images=num_images,
+                overrides=image_overrides,
+            )
+        }
+
+
 @MULTIMODAL_REGISTRY.register_processor(
     BaseInternVLMultiModalProcessor,
     info=SkyworkR1VProcessingInfo,

vllm/model_executor/models/utils.py

@@ -27,6 +27,7 @@ from vllm.multimodal import NestedTensors
 from vllm.sequence import IntermediateTensors
 from vllm.utils.math_utils import cdiv
 from vllm.utils.torch_utils import (
+    async_tensor_h2d,
     direct_register_custom_op,
 )
 
@@ -672,6 +673,17 @@ def _merge_multimodal_embeddings(
     return inputs_embeds
 
 
+def isin_list(
+    elements: torch.Tensor,
+    test_elements_list: list[int],
+) -> torch.Tensor:
+    test_elements = async_tensor_h2d(
+        test_elements_list, dtype=torch.int64, device=elements.device
+    )
+
+    return torch.isin(elements, test_elements)
+
+
 class StageMissingLayer(nn.Module):
     def __init__(self, stage_name: str, module: nn.Module | None = None) -> None:
         super().__init__()

vllm/model_executor/models/vision.py

@@ -580,3 +580,40 @@ def run_dp_sharded_mrope_vision_model(
         "Found unassigned embeddings"
     )
     return out_embeddings
+
+
+def get_llm_pos_ids_for_vision(
+    start_idx: int,
+    vision_idx: int,
+    spatial_merge_size: int,
+    t_index: list[int],
+    grid_hs: torch.Tensor,
+    grid_ws: torch.Tensor,
+) -> torch.Tensor:
+    llm_pos_ids_list = []
+    llm_grid_h = grid_hs[vision_idx] // spatial_merge_size
+    llm_grid_w = grid_ws[vision_idx] // spatial_merge_size
+    h_index = (
+        torch.arange(llm_grid_h)
+        .view(1, -1, 1)
+        .expand(len(t_index), -1, llm_grid_w)
+        .flatten()
+    )
+    w_index = (
+        torch.arange(llm_grid_w)
+        .view(1, 1, -1)
+        .expand(len(t_index), llm_grid_h, -1)
+        .flatten()
+    )
+    t_index_tensor = (
+        torch.Tensor(t_index)
+        .to(llm_grid_h.device)
+        .view(-1, 1)
+        .expand(-1, llm_grid_h * llm_grid_w)
+        .long()
+        .flatten()
+    )
+    _llm_pos_ids = torch.stack([t_index_tensor, h_index, w_index])
+    llm_pos_ids_list.append(_llm_pos_ids + start_idx)
+    llm_pos_ids = torch.cat(llm_pos_ids_list, dim=1)
+    return llm_pos_ids

vllm/models/kimi_k3/nvidia/kda.py

@@ -164,6 +164,8 @@ def is_flashkda_supported(
     dtype: torch.dtype,
     lower_bound: float | None,
 ) -> bool:
+    if not current_platform.is_cuda():
+        return False
     capability = current_platform.get_device_capability()
     return (
         capability is not None

vllm/profiler/utils.py

@@ -12,6 +12,15 @@ from torch._C._profiler import _EventType, _ProfilerEvent, _TensorMetadata
 #
 
 
+def trim_string_front(string: str, width: int) -> str:
+    if len(string) > width:
+        offset = len(string) - width + 3
+        string = string[offset:]
+        if len(string) > 3:
+            string = "..." + string[3:]
+    return string
+
+
 def trim_string_back(string: str, width: int) -> str:
     if len(string) > width:
         offset = len(string) - width + 3

vllm/utils/cache.py

@@ -182,6 +182,10 @@ class LRUCache(cachetools.LRUCache[_K, _V]):
 
         self.popitem(remove_pinned=remove_pinned)
 
+    def _remove_old_if_needed(self) -> None:
+        while self.currsize > self.capacity:
+            self.remove_oldest()
+
     def popitem(self, remove_pinned: bool = False):
         """Remove and return the `(key, value)` pair least recently used."""
         if not remove_pinned:

vllm/utils/multi_stream_utils.py

@@ -8,6 +8,10 @@ from typing import Any
 import torch
 
 
+class AuxStreamType(Enum):
+    Attention = 1
+
+
 class EventType(Enum):
     Main = 0
     Attention = 1

vllm/v1/core/kv_cache_utils.py

@@ -427,6 +427,26 @@ class FreeKVCacheBlockQueue:
             curr_block = curr_block.next_free_block
 
 
+def need_extra_keys(request: Request) -> bool:
+    """Check whether the blocks allocated to this request need extra hash keys.
+
+    Args:
+        request (Request): The request.
+
+    Returns:
+        bool: Whether blocks allocated to this request need extra hash keys.
+    """
+
+    # Multimodal requests need to include the MM hash.
+    # LoRA requests need to include the LoRA name.
+    # Request with provided cache salt need to include the salt.
+    return (
+        bool(request.mm_features)
+        or (request.lora_request is not None)
+        or (request.cache_salt is not None)
+    )
+
+
 def _gen_mm_extra_hash_keys(
     request: Request, start_token_idx: int, end_token_idx: int, start_mm_idx: int
 ) -> tuple[list[Any], int]:
@@ -1033,6 +1053,20 @@ def _get_kv_cache_groups_uniform_type(
     return [KVCacheGroupSpec(list(spec.kv_cache_specs.keys()), spec)]
 
 
+def is_kv_cache_page_size_uniform(kv_cache_spec: dict[str, KVCacheSpec]) -> bool:
+    """
+    Whether all layers in the given KVCacheSpec have the same page size.
+    Args:
+        kv_cache_spec: The KVCacheSpec of each attention layer in the model
+
+    Returns:
+        True if all layers have the same page size, False otherwise.
+    """
+
+    page_sizes = {layer.page_size_bytes for layer in kv_cache_spec.values()}
+    return len(page_sizes) == 1
+
+
 def unify_kv_cache_spec_page_size(
     kv_cache_spec: dict[str, KVCacheSpec],
 ) -> dict[str, KVCacheSpec]: