Skip to content

vllm.worker.multi_step_neuron_model_runner

MultiStepNeuronModelRunner

Bases: NeuronModelRunner

A model runner for multi step decoding using the transformers_neuronx framework

Source code in vllm/worker/multi_step_neuron_model_runner.py 
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
class MultiStepNeuronModelRunner(NeuronModelRunner):
    """A model runner for multi step decoding using the transformers_neuronx
    framework"""

    def __init__(
        self,
        vllm_config: VllmConfig,
    ):
        super().__init__(vllm_config)
        self.speculation_config = self.speculative_config
        from transformers_neuronx.config import GenerationConfig
        self.speculation_config.draft_model_config.neuron_sampling_params = (
            GenerationConfig(
            max_length=self.scheduler_config.max_model_len,
            do_sample=True,
            per_batch_line=True,
            top_k=[self._MAX_NEURON_SAMPLING_TOP_K] \
                  * self.scheduler_config.max_num_seqs,
            top_p=[1.0] * self.scheduler_config.max_num_seqs,
            temperature=[1.0] * self.scheduler_config.max_num_seqs,
            dynamic=True,
            global_top_k=self._MAX_NEURON_SAMPLING_TOP_K
        ))

    def load_model(self) -> None:
        if find_spec("transformers_neuronx") is not None:
            from vllm.model_executor.model_loader.neuron import (
                get_neuron_eagle_speculation_model,
                get_neuron_speculation_model)
            if self.speculation_config.speculative_token_tree is not None:
                self.model = get_neuron_eagle_speculation_model(
                    self.model_config,
                    parallel_config=self.parallel_config,
                    scheduler_config=self.scheduler_config,
                    speculation_config=self.speculation_config)
            else:
                self.model = get_neuron_speculation_model(
                    self.model_config,
                    parallel_config=self.parallel_config,
                    scheduler_config=self.scheduler_config,
                    speculation_config=self.speculation_config)
        else:
            raise NotImplementedError(
                "Supports only Transformer-NeuronX based models.")

    @torch.inference_mode()
    def execute_model(
        self,
        model_input: ModelInputForNeuron,
        kv_caches: Optional[List[torch.Tensor]] = None,
        intermediate_tensors: Optional[IntermediateTensors] = None,
        num_steps: int = 1,
    ) -> Optional[List[SamplerOutput]]:
        logits = self.model(
            input_ids=model_input.input_tokens,
            positions=model_input.input_positions,
            input_block_ids=model_input.input_block_ids,
            **MultiModalKwargs.as_kwargs(
                model_input.multi_modal_kwargs or {},
                device=self.device,
            ),
        )

        output = self.model.sample(
            logits=logits,
            sampling_metadata=model_input.sampling_metadata,
        )
        return output

speculation_config instance-attribute 

speculation_config = speculative_config

__init__ 

__init__(vllm_config: VllmConfig)
Source code in vllm/worker/multi_step_neuron_model_runner.py 
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
def __init__(
    self,
    vllm_config: VllmConfig,
):
    super().__init__(vllm_config)
    self.speculation_config = self.speculative_config
    from transformers_neuronx.config import GenerationConfig
    self.speculation_config.draft_model_config.neuron_sampling_params = (
        GenerationConfig(
        max_length=self.scheduler_config.max_model_len,
        do_sample=True,
        per_batch_line=True,
        top_k=[self._MAX_NEURON_SAMPLING_TOP_K] \
              * self.scheduler_config.max_num_seqs,
        top_p=[1.0] * self.scheduler_config.max_num_seqs,
        temperature=[1.0] * self.scheduler_config.max_num_seqs,
        dynamic=True,
        global_top_k=self._MAX_NEURON_SAMPLING_TOP_K
    ))

execute_model 

execute_model(
    model_input: ModelInputForNeuron,
    kv_caches: Optional[List[Tensor]] = None,
    intermediate_tensors: Optional[
        IntermediateTensors
    ] = None,
    num_steps: int = 1,
) -> Optional[List[SamplerOutput]]
Source code in vllm/worker/multi_step_neuron_model_runner.py 
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
@torch.inference_mode()
def execute_model(
    self,
    model_input: ModelInputForNeuron,
    kv_caches: Optional[List[torch.Tensor]] = None,
    intermediate_tensors: Optional[IntermediateTensors] = None,
    num_steps: int = 1,
) -> Optional[List[SamplerOutput]]:
    logits = self.model(
        input_ids=model_input.input_tokens,
        positions=model_input.input_positions,
        input_block_ids=model_input.input_block_ids,
        **MultiModalKwargs.as_kwargs(
            model_input.multi_modal_kwargs or {},
            device=self.device,
        ),
    )

    output = self.model.sample(
        logits=logits,
        sampling_metadata=model_input.sampling_metadata,
    )
    return output

load_model 

load_model() -> None
Source code in vllm/worker/multi_step_neuron_model_runner.py 
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
def load_model(self) -> None:
    if find_spec("transformers_neuronx") is not None:
        from vllm.model_executor.model_loader.neuron import (
            get_neuron_eagle_speculation_model,
            get_neuron_speculation_model)
        if self.speculation_config.speculative_token_tree is not None:
            self.model = get_neuron_eagle_speculation_model(
                self.model_config,
                parallel_config=self.parallel_config,
                scheduler_config=self.scheduler_config,
                speculation_config=self.speculation_config)
        else:
            self.model = get_neuron_speculation_model(
                self.model_config,
                parallel_config=self.parallel_config,
                scheduler_config=self.scheduler_config,
                speculation_config=self.speculation_config)
    else:
        raise NotImplementedError(
            "Supports only Transformer-NeuronX based models.")