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

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 15, 01:00 UTC

[Model] Remove model sampler (#21059)

Signed-off-by: DarkLight1337 <tlleungac@connect.ust.hk>

Cyrus Leung committed Jul 17, 2025 at 03:03 UTC ac2bf41e5390e4cec61c7d87eb36699a882aceae

6 files changed -45

vllm/model_executor/models/bailing_moe.py

-10

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

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

48 QuantizationConfig)

49 from vllm.model_executor.layers.rotary_embedding import get_rope

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

50 from vllm.model_executor.layers.vocab_parallel_embedding import (

51 ParallelLMHead, VocabParallelEmbedding)

52 from vllm.model_executor.model_loader.weight_utils import default_weight_loader

@@ -485,7 +484,6 @@ class BailingMoeForCausalLM(nn.Module, SupportsPP):

484 else:

485 self.lm_head = PPMissingLayer()

486

488 - self.sampler = get_sampler()

487 self.make_empty_intermediate_tensors = (

488 self.model.make_empty_intermediate_tensors)

489

@@ -512,14 +510,6 @@ class BailingMoeForCausalLM(nn.Module, SupportsPP):

510 sampling_metadata)

511 return logits

512

515 - def sample(

516 - self,

517 - logits: torch.Tensor,

518 - sampling_metadata: SamplingMetadata,

519 - ) -> Optional[SamplerOutput]:

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

521 - return next_tokens

522 -

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

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

515 loader = AutoWeightsLoader(

vllm/model_executor/models/granite_speech.py

-2

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

36 from vllm.model_executor.layers.linear import (ColumnParallelLinear,

37 RowParallelLinear)

38 from vllm.model_executor.layers.quantization import QuantizationConfig

39 -from vllm.model_executor.layers.sampler import get_sampler

39 from vllm.model_executor.models.module_mapping import MultiModelKeys

40 from vllm.model_executor.sampling_metadata import SamplingMetadata

41 from vllm.multimodal import MULTIMODAL_REGISTRY

@@ -549,7 +548,6 @@ class GraniteSpeechForConditionalGeneration(

548 self.config = config

549 self.quant_config = quant_config

550 self.cache_config = cache_config

552 - self.sampler = get_sampler()

551

552 # The language model is typically a Granite LLM

553 self.language_model = init_vllm_registered_model(

vllm/model_executor/models/hunyuan_v1_moe.py

-10

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

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

50 QuantizationConfig)

51 from vllm.model_executor.layers.rotary_embedding import get_rope

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

52 from vllm.model_executor.layers.vocab_parallel_embedding import (

53 DEFAULT_VOCAB_PADDING_SIZE, ParallelLMHead, VocabParallelEmbedding)

54 from vllm.model_executor.model_loader.weight_utils import (

@@ -661,7 +660,6 @@ class HunYuanMoEV1ForCausalLM(nn.Module):

660 self.logits_processor = LogitsProcessor(self.unpadded_vocab_size,

661 config.vocab_size,

662 logit_scale)

664 - self.sampler = get_sampler()

663 else:

664 self.lm_head = PPMissingLayer()

665

@@ -685,14 +683,6 @@ class HunYuanMoEV1ForCausalLM(nn.Module):

683 sampling_metadata)

684 return logits

685

688 - def sample(

689 - self,

690 - logits: torch.Tensor,

691 - sampling_metadata: SamplingMetadata,

692 - ) -> Optional[SamplerOutput]:

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

694 - return next_tokens

695 -

686 def make_empty_intermediate_tensors(

687 self, batch_size: int, dtype: torch.dtype,

688 device: torch.device) -> IntermediateTensors:

vllm/model_executor/models/mimo.py

-2

@@ -36,7 +36,6 @@ from vllm.config import VllmConfig

36 from vllm.distributed import get_pp_group

37 from vllm.logger import init_logger

38 from vllm.model_executor.layers.logits_processor import LogitsProcessor

39 -from vllm.model_executor.layers.sampler import get_sampler

39 from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead

40 from vllm.model_executor.model_loader.weight_utils import (

41 default_weight_loader, maybe_remap_kv_scale_name)

@@ -176,7 +175,6 @@ class MiMoForCausalLM(Qwen2ForCausalLM, nn.Module):

175 self.lm_head = PPMissingLayer()

176

177 self.logits_processor = LogitsProcessor(config.vocab_size)

179 - self.sampler = get_sampler()

178

179 self.make_empty_intermediate_tensors = (

180 self.model.make_empty_intermediate_tensors)

vllm/model_executor/models/mimo_mtp.py

-11

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

30 from vllm.model_executor.layers.layernorm import RMSNorm

31 from vllm.model_executor.layers.logits_processor import LogitsProcessor

32 from vllm.model_executor.layers.quantization import QuantizationConfig

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

33 from vllm.model_executor.layers.vocab_parallel_embedding import (

34 ParallelLMHead, VocabParallelEmbedding)

35 from vllm.model_executor.model_loader.weight_utils import default_weight_loader

@@ -161,8 +160,6 @@ class MiMoMTP(nn.Module):

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

161 self.config.hidden_size)

162

164 - self.sampler = get_sampler()

165 -

163 def forward(

164 self,

165 input_ids: torch.Tensor,

@@ -187,14 +184,6 @@ class MiMoMTP(nn.Module):

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

185 sampling_metadata, spec_step_idx)

186

190 - def sample(

191 - self,

192 - logits: torch.Tensor,

193 - sampling_metadata: SamplingMetadata,

194 - ) -> Optional[SamplerOutput]:

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

196 - return next_tokens

197 -

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

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

189 stacked_params_mapping = [

vllm/model_executor/models/phi4flash.py

-10

@@ -23,7 +23,6 @@ from vllm.model_executor.layers.mamba.ops.causal_conv1d import (

23 causal_conv1d_fn, causal_conv1d_update)

24 from vllm.model_executor.layers.mamba.ops.mamba_ssm import (

25 selective_scan_fn, selective_state_update)

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

26 from vllm.model_executor.layers.vocab_parallel_embedding import (

27 DEFAULT_VOCAB_PADDING_SIZE, ParallelLMHead, VocabParallelEmbedding)

28 from vllm.model_executor.models.interfaces import (HasInnerState, IsHybrid,

@@ -641,7 +640,6 @@ class Phi4FlashForCausalLM(nn.Module, HasInnerState, IsHybrid, SupportsV0Only):

640 self.logits_processor = LogitsProcessor(self.unpadded_vocab_size,

641 config.vocab_size,

642 logits_as_input=False)

644 - self.sampler = get_sampler()

643

644 def forward(

645 self,

@@ -709,14 +707,6 @@ class Phi4FlashForCausalLM(nn.Module, HasInnerState, IsHybrid, SupportsV0Only):

707 prune_hidden_states=prune_hidden_states)

708 return processed_logits

709

712 - def sample(

713 - self,

714 - logits: torch.Tensor,

715 - sampling_metadata: SamplingMetadata,

716 - ) -> Optional[SamplerOutput]:

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

718 - return next_tokens

719 -

710 def load_weights(

711 self,

712 weights: Iterable[tuple[str, torch.Tensor]],