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, 09:45 UTC
Order sequence ids + config update to support specifying custom quantization layers (#18279)
Signed-off-by: Elaine Zhao <elaineyz@amazon.com> Co-authored-by: Tailin Pan <tailinpa@amazon.com> Co-authored-by: Rishabh Rajesh <rishyraj@amazon.com> Co-authored-by: Yishan McNabb <yishanm@amazon.com> Co-authored-by: Patrick Lange <patlange@amazon.com> Co-authored-by: Maxwell Goldberg <mgld@amazon.com> Co-authored-by: Aakash Shetty <sheaak@amazon.com>
aws-elaineyz committed
May 22, 2025 at 02:20 UTC
fa72f9a8126051abb9d00144f11aeb7615f36d21
2 files changed
+73
-10
tests/neuron/2_core/test_mistral.py
+35
-5
@@ -7,7 +7,7 @@ def test_mistral():
7
llm = LLM(model="mistralai/Mistral-7B-v0.1",
8
tensor_parallel_size=2,
9
max_num_seqs=4,
10
- max_model_len=512,
10
+ max_model_len=128,
11
use_v2_block_manager=True,
12
override_neuron_config={
13
"sequence_parallel_enabled": False,
@@ -15,16 +15,46 @@ def test_mistral():
15
},
16
device="neuron")
17
18
+ # Send more prompts than the compiled batch size (4) and request
19
+ # varying generation lengths to test accuracy related to Neuron
20
+ # specific sequence id sorting.
21
prompts = [
22
"The president of the United States is",
23
"The capital of France is",
24
+ "What is Annapurna labs?",
25
+ "I believe the meaning of life is",
26
+ "Tell me a story about a brave knight",
27
+ "Hello, my name is Llama",
28
]
22
- outputs = llm.generate(prompts, SamplingParams(top_k=1))
29
+
30
+ sampling_params = [
31
+ SamplingParams(top_k=1, max_tokens=10),
32
+ SamplingParams(top_k=1, max_tokens=20),
33
+ SamplingParams(top_k=1, max_tokens=30),
34
+ SamplingParams(top_k=1, max_tokens=40),
35
+ SamplingParams(top_k=1, max_tokens=50),
36
+ SamplingParams(top_k=1, max_tokens=60)
37
+ ]
38
+
39
+ outputs = llm.generate(prompts, sampling_params)
40
41
expected_outputs = [
25
- " the most powerful person in the world. He is the head of state "
26
- "and head",
27
- " a city of many faces. It is a city of history, culture, art"
42
+ " the most powerful person in the world. He is",
43
+ " a city of many faces. It is a city of history, culture, art, "
44
+ "fashion, and",
45
+ "\n\nAnnapurna Labs is a semiconductor company that was founded "
46
+ "in 2013 by Amazon. The company is",
47
+ " to be happy.\n\nI believe that happiness is a choice.\n\nI "
48
+ "believe that happiness is a state of mind.\n\nI believe that "
49
+ "happiness is a journey.\n\nI believe",
50
+ " who rescued a princess from a dragon.\n\nTell me a story about"
51
+ " a princess who rescued herself from a dragon.\n\nTell me a "
52
+ "story about a princess who rescued herself from a dragon and "
53
+ "then rescued a knight from",
54
+ " and I am a 10 year old male. I am a very friendly and "
55
+ "affectionate boy who loves to be around people. I am a very "
56
+ "active boy who loves to play and run around. I am a very smart "
57
+ "boy who loves to learn new things. I am a very loyal boy"
58
]
59
60
for expected_output, output in zip(expected_outputs, outputs):
vllm/model_executor/model_loader/neuronx_distributed.py
+38
-5
@@ -87,16 +87,29 @@ class NeuronCausalLM(nn.Module):
87
input_block_ids: torch.Tensor,
88
sampling_params: torch.Tensor,
89
) -> torch.Tensor:
90
+ # sort block ids sequentially for perf/neuron support reasons
91
+ sorted_input_block_ids, sorted_indices = torch.sort(input_block_ids)
92
+ input_ids = torch.index_select(input_ids, 0, sorted_indices)
93
+ positions = torch.index_select(positions, 0, sorted_indices)
94
+ sampling_params = torch.index_select(sampling_params, 0,
95
+ sorted_indices)
96
+
97
output = self.model(input_ids,
98
attention_mask=None,
99
position_ids=positions,
93
- seq_ids=input_block_ids,
100
+ seq_ids=sorted_input_block_ids,
101
sampling_params=sampling_params)
102
# on-device sampling
103
if self.config.neuron_config.on_device_sampling_config:
97
- return output.hidden_states
104
+ output = output.hidden_states
105
else:
99
- return output.logits[:, -1, :]
106
+ output = output.logits[:, -1, :]
107
+
108
+ restored_indices = torch.argsort(sorted_indices)
109
+ if input_block_ids.shape[0] != 1:
110
+ output = torch.index_select(output, 0, restored_indices)
111
+
112
+ return output
113
114
def compute_logits(self, hidden_states: torch.Tensor,
115
sampling_metadata: SamplingMetadata) -> torch.Tensor:
@@ -340,14 +353,26 @@ class NeuronSpeculationCausalLM(nn.Module):
353
input_block_ids: torch.Tensor,
354
sampling_params: torch.Tensor,
355
) -> torch.Tensor:
356
+ # sort block ids sequentially for perf/neuron support reasons
357
+ sorted_input_block_ids, sorted_indices = torch.sort(input_block_ids)
358
+ input_ids = torch.index_select(input_ids, 0, sorted_indices)
359
+ positions = torch.index_select(positions, 0, sorted_indices)
360
+ sampling_params = torch.index_select(sampling_params, 0,
361
+ sorted_indices)
362
+
363
output = self.model(input_ids,
364
attention_mask=None,
365
position_ids=positions,
346
- seq_ids=input_block_ids,
366
+ seq_ids=sorted_input_block_ids,
367
sampling_params=sampling_params)
368
+ restored_indices = torch.argsort(sorted_indices)
369
+
370
# CTX encoding
371
if (positions[:, 0]).sum().item() == 0:
350
- return output.fused_outputs[0][:, 0:1]
372
+ output = output.fused_outputs[0][:, 0:1]
373
+ if input_block_ids.shape[0] != 1:
374
+ output = torch.index_select(output, 0, restored_indices)
375
+ return output
376
377
# Fused Spec (Generation)
378
accepted_tokens_with_padding = output.fused_outputs[0]
@@ -362,6 +387,10 @@ class NeuronSpeculationCausalLM(nn.Module):
387
-1) >= generated_token_counts
388
accepted_tokens_with_padding[mask] = -1
389
390
+ if input_block_ids.shape[0] != 1:
391
+ accepted_tokens_with_padding = torch.index_select(
392
+ accepted_tokens_with_padding, 0, restored_indices)
393
+
394
return accepted_tokens_with_padding
395
396
def sample(
@@ -416,6 +445,10 @@ class NeuronSpeculationCausalLM(nn.Module):
445
draft_neuron_config.speculation_length = 0
446
draft_neuron_config.trace_tokengen_model = True
447
draft_neuron_config.enable_fused_speculation = False
448
+ if getattr(config.neuron_config, "draft_model_modules_to_not_convert",
449
+ None):
450
+ draft_neuron_config.modules_to_not_convert = (
451
+ draft_neuron_config.draft_model_modules_to_not_convert)
452
if config.neuron_config.enable_eagle_speculation:
453
draft_neuron_config.is_eagle_draft = True
454
draft_neuron_config.sequence_parallel_enabled = False