@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 572ddf83ce

Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 14, 05:30 UTC

[Chore] Remove unused sampler in models (#25324)

Signed-off-by: Woosuk Kwon <woosuk.kwon@berkeley.edu>

Woosuk Kwon committed Sep 20, 2025 at 19:53 UTC 572ddf83ce1fe8d52670e01a7a6cc0d8a99fa90c

5 files changed -49

tests/lora/conftest.py

-3

@@ -17,7 +17,6 @@ from vllm.model_executor.layers.linear import (ColumnParallelLinear,

17 MergedColumnParallelLinear,

18 RowParallelLinear)

19 from vllm.model_executor.layers.logits_processor import LogitsProcessor

20 -from vllm.model_executor.layers.sampler import Sampler

20 from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead

21 from vllm.model_executor.models.interfaces import SupportsLoRA

22 from vllm.platforms import current_platform

@@ -97,7 +96,6 @@ def dummy_model() -> nn.Module:

96 # Special handling for lm_head & sampler

97 ("lm_head", ParallelLMHead(512, 10)),

98 ("logits_processor", LogitsProcessor(512)),

100 - ("sampler", Sampler())

99 ]))

100 model.config = MagicMock()

101 model.embedding_modules = {"lm_head": "lm_head"}

@@ -125,7 +123,6 @@ def dummy_model_gate_up() -> nn.Module:

123 # Special handling for lm_head & sampler

124 ("lm_head", ParallelLMHead(512, 10)),

125 ("logits_processor", LogitsProcessor(512)),

128 - ("sampler", Sampler())

126 ]))

127 model.config = MagicMock()

128 model.packed_modules_mapping = {

vllm/model_executor/models/ernie_mtp.py

-10

@@ -33,7 +33,6 @@ from vllm.config import CacheConfig, ModelConfig, VllmConfig

33 from vllm.model_executor.layers.layernorm import RMSNorm

34 from vllm.model_executor.layers.logits_processor import LogitsProcessor

35 from vllm.model_executor.layers.quantization import QuantizationConfig

36 -from vllm.model_executor.layers.sampler import SamplerOutput, get_sampler

36 from vllm.model_executor.layers.vocab_parallel_embedding import (

37 ParallelLMHead, VocabParallelEmbedding)

38 from vllm.model_executor.model_loader.weight_utils import default_weight_loader

@@ -160,7 +159,6 @@ class ErnieMTP(nn.Module, SupportsPP):

159 self.lm_head = ParallelLMHead(self.config.vocab_size,

160 self.config.hidden_size,

161 prefix=maybe_prefix(prefix, "lm_head"))

163 - self.sampler = get_sampler()

162

163 if self.config.tie_word_embeddings:

164 self.lm_head.weight = self.model.embed_tokens.weight

@@ -188,14 +186,6 @@ class ErnieMTP(nn.Module, SupportsPP):

186 return self.model.compute_logits(hidden_states, self.lm_head,

187 sampling_metadata, spec_step_idx)

188

191 - def sample(

192 - self,

193 - logits: torch.Tensor,

194 - sampling_metadata: SamplingMetadata,

195 - ) -> Optional[SamplerOutput]:

196 - next_tokens = self.sampler(logits, sampling_metadata)

197 - return next_tokens

198 -

189 def load_weights(self, weights: Iterable[tuple[str,

190 torch.Tensor]]) -> set[str]:

191 stacked_params_mapping = [

vllm/model_executor/models/plamo2.py

-10

@@ -41,7 +41,6 @@ from vllm.model_executor.layers.mamba.ops.ssd_combined import (

41 mamba_chunk_scan_combined)

42 from vllm.model_executor.layers.quantization import QuantizationConfig

43 from vllm.model_executor.layers.rotary_embedding import get_rope

44 -from vllm.model_executor.layers.sampler import SamplerOutput, get_sampler

44 from vllm.model_executor.layers.vocab_parallel_embedding import (

45 DEFAULT_VOCAB_PADDING_SIZE, ParallelLMHead, VocabParallelEmbedding)

46 from vllm.model_executor.model_loader.weight_utils import (

@@ -932,7 +931,6 @@ class Plamo2ForCausalLM(torch.nn.Module, HasInnerState, SupportsPP, IsHybrid):

931

932 self.logits_processor = LogitsProcessor(self.unpadded_vocab_size,

933 self.config.vocab_size)

935 - self.sampler = get_sampler()

934 self.make_empty_intermediate_tensors = (

935 self.model.make_empty_intermediate_tensors)

936

@@ -1030,14 +1028,6 @@ class Plamo2ForCausalLM(torch.nn.Module, HasInnerState, SupportsPP, IsHybrid):

1028 sampling_metadata)

1029 return logits

1030

1033 - def sample(

1034 - self,

1035 - logits: Optional[torch.Tensor],

1036 - sampling_metadata: SamplingMetadata,

1037 - ) -> Optional[SamplerOutput]:

1038 - next_tokens = self.sampler(logits, sampling_metadata)

1039 - return next_tokens

1040 -

1031 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):

1032 params_dict = dict(self.named_parameters())

1033 for name, loaded_weight in weights:

vllm/model_executor/models/step3_text.py

-10

@@ -26,7 +26,6 @@ from vllm.model_executor.layers.logits_processor import LogitsProcessor

26 from vllm.model_executor.layers.quantization.base_config import (

27 QuantizationConfig)

28 from vllm.model_executor.layers.rotary_embedding import get_rope

29 -from vllm.model_executor.layers.sampler import SamplerOutput, get_sampler

29 from vllm.model_executor.layers.vocab_parallel_embedding import (

30 DEFAULT_VOCAB_PADDING_SIZE, ParallelLMHead, VocabParallelEmbedding)

31 from vllm.model_executor.model_loader.weight_utils import default_weight_loader

@@ -391,7 +390,6 @@ class Step3TextForCausalLM(nn.Module, SupportsPP):

390 )

391 self.logits_processor = LogitsProcessor(self.unpadded_vocab_size,

392 config.vocab_size)

394 - self.sampler = get_sampler()

393 else:

394 self.lm_head = PPMissingLayer()

395

@@ -413,14 +411,6 @@ class Step3TextForCausalLM(nn.Module, SupportsPP):

411 sampling_metadata)

412 return logits

413

416 - def sample(

417 - self,

418 - logits: Optional[torch.Tensor],

419 - sampling_metadata: SamplingMetadata,

420 - ) -> Optional[SamplerOutput]:

421 - next_tokens = self.sampler(logits, sampling_metadata)

422 - return next_tokens

423 -

414 def load_weights(self, weights: Iterable[tuple[str,

415 torch.Tensor]]) -> set[str]:

416 qkv_params_mapping = [

vllm/model_executor/models/step3_vl.py

-16

@@ -2,7 +2,6 @@

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 import math

4 from collections.abc import Iterable, Mapping, Sequence

5 -from functools import cached_property

5 from itertools import product

6 from math import ceil, sqrt

7 from typing import Any, Literal, Optional, TypedDict, Union

@@ -24,7 +23,6 @@ from vllm.model_executor.layers.linear import (ColumnParallelLinear,

23 QKVParallelLinear,

24 RowParallelLinear)

25 from vllm.model_executor.layers.quantization import QuantizationConfig

27 -from vllm.model_executor.layers.sampler import SamplerOutput, get_sampler

26 from vllm.model_executor.sampling_metadata import SamplingMetadata

27 from vllm.multimodal import MULTIMODAL_REGISTRY

28 from vllm.multimodal.inputs import (MultiModalDataDict, MultiModalFieldConfig,

@@ -897,13 +895,6 @@ class Step3VLForConditionalGeneration(nn.Module, SupportsMultiModal,

895 self.make_empty_intermediate_tensors = (

896 self.language_model.make_empty_intermediate_tensors)

897

900 - @cached_property

901 - def sampler(self):

902 - if hasattr(self.language_model, "sampler"):

903 - return self.language_model.sampler

904 -

905 - return get_sampler()

906 -

898 @property

899 def device(self):

900 return next(self.parameters()).device

@@ -1069,13 +1060,6 @@ class Step3VLForConditionalGeneration(nn.Module, SupportsMultiModal,

1060 return self.language_model.compute_logits(hidden_states,

1061 sampling_metadata)

1062

1072 - def sample(

1073 - self,

1074 - logits: torch.Tensor,

1075 - sampling_metadata: SamplingMetadata,

1076 - ) -> Optional[SamplerOutput]:

1077 - return self.language_model.sample(logits, sampling_metadata)

1078 -

1063 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):

1064

1065 skip_prefixes = []