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, 22:15 UTC
[Neuron] Add custom_ops for neuron backend (#13246)
Signed-off-by: Liangfu Chen <liangfc@amazon.com> Co-authored-by: George Novack <gnovack@amazon.com> Co-authored-by: Aoyu Zhang <aoyuzhan@amazon.com>
Liangfu Chen committed
Feb 25, 2025 at 11:47 UTC
f75aa727320f0edabe5d31410ef86ac8ae27d86b
9 files changed
+346
-3
tests/neuron/test_activation.py
new
+42
@@ -0,0 +1,42 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+
3
+import pytest
4
+import torch
5
+import torch.nn.functional as F
6
+
7
+from vllm.model_executor.layers.activation import FastGELU, SiluAndMul
8
+from vllm.platforms import current_platform
9
+
10
+
11
+@pytest.mark.parametrize("activation", ["silu_and_mul", "gelu_fast"])
12
+@pytest.mark.parametrize("num_tokens,d,dtype", [
13
+ (7, 512, torch.half),
14
+ (7, 512, torch.float),
15
+ (83, 512, torch.half),
16
+])
17
+@torch.inference_mode()
18
+def test_act_and_mul(
19
+ activation: str,
20
+ num_tokens: int,
21
+ d: int,
22
+ dtype: torch.dtype,
23
+) -> None:
24
+ import torch_xla.core.xla_model as xm
25
+
26
+ device = xm.xla_device()
27
+ current_platform.seed_everything(0)
28
+ torch.set_default_device("cpu")
29
+ x = torch.randn(num_tokens, 2 * d, dtype=dtype).to(device=device)
30
+ if activation == "silu_and_mul":
31
+ layer = SiluAndMul()
32
+ fn = layer.forward_native
33
+ elif activation == "gelu_fast":
34
+ layer = FastGELU()
35
+ fn = F.gelu
36
+ else:
37
+ raise NotImplementedError(
38
+ f"activation {activation} is not implemented.")
39
+ assert x.is_xla, "input tensor under testing is expected to be XLA tensor."
40
+ out = layer.to(device=device).forward_neuron(x)
41
+ ref_out = fn(x.cpu())
42
+ torch.testing.assert_close(out.cpu(), ref_out, atol=0.01, rtol=0.0)
tests/neuron/test_layernorm.py
new
+56
@@ -0,0 +1,56 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+
3
+import pytest
4
+import torch
5
+
6
+from vllm.model_executor.layers.layernorm import RMSNorm
7
+from vllm.platforms import current_platform
8
+
9
+
10
+@pytest.mark.parametrize("num_tokens,hidden_size,add_residual,dtype", [
11
+ (7, 8, False, torch.half),
12
+ (83, 768, False, torch.half),
13
+ (83, 768, True, torch.half),
14
+ (83, 768, True, torch.bfloat16),
15
+ (83, 768, True, torch.float32),
16
+])
17
+@torch.inference_mode()
18
+def test_rms_norm(
19
+ num_tokens: int,
20
+ hidden_size: int,
21
+ add_residual: bool,
22
+ dtype: torch.dtype,
23
+) -> None:
24
+ import torch_xla.core.xla_model as xm
25
+
26
+ device = xm.xla_device()
27
+ current_platform.seed_everything(0)
28
+ torch.set_default_device("cpu")
29
+ layer = RMSNorm(hidden_size).to(dtype=dtype)
30
+ layer.weight.data.normal_(mean=1.0, std=0.1)
31
+ scale = 1 / (2 * hidden_size)
32
+ x = torch.randn(num_tokens, hidden_size, dtype=dtype).to(device=device)
33
+ x *= scale
34
+ residual = torch.randn_like(x) * scale if add_residual else None
35
+
36
+ residual_cpu = residual.cpu() if add_residual else None
37
+ ref_out = layer.to(device="cpu").forward_native(x.cpu(), residual_cpu)
38
+ assert x.is_xla, "input tensor under testing is expected to be XLA tensor."
39
+ out = layer.to(device=device)(x, residual)
40
+
41
+ # NOTE(woosuk): LayerNorm operators (including RMS) typically have larger
42
+ # numerical errors than other operators because they involve reductions.
43
+ # Therefore, we use a larger tolerance.
44
+ if add_residual:
45
+ assert out[0].is_xla, "output tensor is expected to be XLA tensor"
46
+ torch.testing.assert_close(out[0].cpu(),
47
+ ref_out[0],
48
+ atol=1e-2,
49
+ rtol=1e-2)
50
+ torch.testing.assert_close(out[1].cpu(),
51
+ ref_out[1],
52
+ atol=1e-2,
53
+ rtol=1e-2)
54
+ else:
55
+ assert out.is_xla, "output tensor is expected to be XLA tensor"
56
+ torch.testing.assert_close(out.cpu(), ref_out, atol=1e-2, rtol=1e-2)
tests/neuron/test_logits_processor.py
new
+95
@@ -0,0 +1,95 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+
3
+import random
4
+from typing import Tuple
5
+from unittest.mock import patch
6
+
7
+import pytest
8
+import torch
9
+
10
+from vllm.model_executor.layers.logits_processor import LogitsProcessor
11
+from vllm.model_executor.sampling_metadata import SamplingMetadata
12
+from vllm.model_executor.utils import set_random_seed
13
+from vllm.sequence import SamplingParams, SequenceData, SequenceGroupMetadata
14
+from vllm.utils import is_pin_memory_available
15
+
16
+
17
+class MockLogitsProcessor(LogitsProcessor):
18
+
19
+ def __init__(self, vocab_size: int, scale: float,
20
+ fake_logits: torch.Tensor):
21
+ super().__init__(vocab_size=vocab_size, scale=scale)
22
+ self.fake_logits = fake_logits.clone()
23
+
24
+ def forward(self, *args, **kwargs):
25
+ with patch(
26
+ "vllm.model_executor.layers.logits_processor._prune_hidden_states",
27
+ lambda x, y: x
28
+ ), patch(
29
+ "vllm.model_executor.layers.logits_processor.LogitsProcessor._get_logits",
30
+ lambda *args, **kwargs: self.fake_logits):
31
+ return super().forward(*args, **kwargs)
32
+
33
+
34
+def _prepare_test(
35
+ batch_size: int
36
+) -> Tuple[torch.Tensor, torch.Tensor, MockLogitsProcessor]:
37
+ vocab_size = 32000
38
+ input_tensor = torch.rand((batch_size, 1024), dtype=torch.float16)
39
+ fake_logits = torch.full((batch_size, vocab_size),
40
+ 1e-2,
41
+ dtype=input_tensor.dtype)
42
+ logits_processor = MockLogitsProcessor(32000, 0.5, fake_logits)
43
+ return input_tensor, fake_logits, logits_processor
44
+
45
+
46
+RANDOM_SEEDS = list(range(8))
47
+
48
+
49
+@pytest.mark.parametrize("seed", RANDOM_SEEDS)
50
+def test_logits_processors(seed: int):
51
+ import torch_xla.core.xla_model as xm
52
+
53
+ device = xm.xla_device()
54
+ set_random_seed(seed)
55
+ torch.set_default_device("cpu")
56
+ batch_size = random.randint(1, 256)
57
+ input_tensor, fake_logits, logits_processor = _prepare_test(batch_size)
58
+
59
+ # This sample logits processor gives infinite score to the i-th token,
60
+ # where i is the length of the input sequence.
61
+ # We therefore expect the output token sequence to be [0, 1, 2, ...]
62
+ def pick_ith(token_ids, logits):
63
+ logits[len(token_ids)] = float("inf")
64
+ return logits
65
+
66
+ seq_group_metadata_list = []
67
+ seq_lens = []
68
+ for i in range(batch_size):
69
+ seq_group_metadata_list.append(
70
+ SequenceGroupMetadata(
71
+ request_id=f"test_{i}",
72
+ is_prompt=True,
73
+ seq_data={0: SequenceData.from_seqs([1, 2, 3])},
74
+ sampling_params=SamplingParams(temperature=0,
75
+ logits_processors=[pick_ith]),
76
+ block_tables={0: [1]},
77
+ ))
78
+ seq_lens.append(seq_group_metadata_list[-1].seq_data[0].get_len())
79
+
80
+ sampling_metadata = SamplingMetadata.prepare(
81
+ seq_group_metadata_list,
82
+ seq_lens,
83
+ query_lens=seq_lens,
84
+ device=device,
85
+ pin_memory=is_pin_memory_available())
86
+ logits_processor_output = logits_processor(
87
+ lm_head=None,
88
+ hidden_states=input_tensor,
89
+ sampling_metadata=sampling_metadata)
90
+
91
+ fake_logits *= logits_processor.scale
92
+ torch.testing.assert_close(logits_processor_output[:, 1],
93
+ fake_logits[:, 1],
94
+ rtol=1e-4,
95
+ atol=0.0)
tests/neuron/test_prefix_prefill.py
+4
-3
@@ -345,6 +345,7 @@ def test_contexted_kv_attention(
345
346
torch.manual_seed(0)
347
torch.set_printoptions(sci_mode=False)
348
+ torch.set_default_device("cpu")
349
dtype = torch.float32
350
351
min_ctx_len = 32
@@ -438,9 +439,9 @@ def test_contexted_kv_attention(
439
440
# transform block table
441
active_block_table = get_active_block_tables(
441
- block_table,
442
- torch.tensor(query_lens),
443
- torch.tensor(seq_lens),
442
+ block_table.cpu(),
443
+ torch.tensor(query_lens).cpu(),
444
+ torch.tensor(seq_lens).cpu(),
445
block_size,
446
num_active_blocks,
447
)
tests/neuron/test_rotary_embedding.py
new
+58
@@ -0,0 +1,58 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+"""
3
+Tests for miscellaneous utilities
4
+"""
5
+
6
+import pytest
7
+import torch
8
+
9
+from vllm.model_executor.layers.rotary_embedding import RotaryEmbedding
10
+from vllm.platforms import current_platform
11
+
12
+
13
+@pytest.mark.parametrize(
14
+ "max_position,is_neox_style,rotary_dim,head_size,seq_len", [
15
+ (16, False, 32, 32, 1024),
16
+ (16, False, 32, 128, 1024),
17
+ (16, True, 32, 32, 1024),
18
+ (16, True, 32, 128, 1024),
19
+ ])
20
+def test_rotary_embedding_opcheck(max_position, is_neox_style, rotary_dim,
21
+ head_size, seq_len):
22
+ import torch_xla.core.xla_model as xm
23
+
24
+ device = xm.xla_device()
25
+ current_platform.seed_everything(0)
26
+ torch.set_default_device("cpu")
27
+
28
+ batch_size = 1
29
+ base = 10000
30
+ num_heads = 8
31
+
32
+ rot = RotaryEmbedding(head_size, rotary_dim, max_position, base,
33
+ is_neox_style, torch.float32)
34
+
35
+ positions = torch.randint(0,
36
+ max_position, (batch_size, seq_len),
37
+ device="cpu")
38
+ query = torch.randn(batch_size,
39
+ seq_len,
40
+ num_heads * head_size,
41
+ dtype=torch.float32,
42
+ device="cpu")
43
+ key = torch.randn_like(query)
44
+
45
+ assert positions.is_cpu, \
46
+ "reference input tensor is expected to be CPU tensor."
47
+ ref_query, ref_key = rot.to(device="cpu").forward_native(
48
+ positions, query, key)
49
+ out_query, out_key = rot.to(device=device).forward_neuron(
50
+ positions.to(device=device), query.to(device=device),
51
+ key.to(device=device))
52
+ assert out_query.is_xla and out_key.is_xla, \
53
+ "output tensor is expected to be XLA tensor"
54
+ torch.testing.assert_close(out_query.cpu(),
55
+ ref_query,
56
+ atol=1e-2,
57
+ rtol=1e-2)
58
+ torch.testing.assert_close(out_key.cpu(), ref_key, atol=1e-2, rtol=1e-2)
vllm/model_executor/custom_op.py
+7
@@ -59,6 +59,11 @@ class CustomOp(nn.Module):
59
# PyTorch-native implementation.
60
return self.forward_native(*args, **kwargs)
61
62
+ def forward_neuron(self, *args, **kwargs):
63
+ # By default, we assume that Neuron ops are compatible with the
64
+ # PyTorch-native implementation.
65
+ return self.forward_native(*args, **kwargs)
66
+
67
def forward_oot(self, *args, **kwargs):
68
# By default, we assume that OOT ops are compatible with the
69
# PyTorch-native implementation.
@@ -88,6 +93,8 @@ class CustomOp(nn.Module):
93
return self.forward_tpu
94
elif current_platform.is_xpu():
95
return self.forward_xpu
96
+ elif current_platform.is_neuron():
97
+ return self.forward_neuron
98
elif current_platform.is_out_of_tree():
99
return self.forward_oot
100
else:
vllm/model_executor/layers/activation.py
+7
@@ -89,6 +89,13 @@ class SiluAndMul(CustomOp):
89
self.op(out, x)
90
return out
91
92
+ def forward_neuron(self, x: torch.Tensor) -> torch.Tensor:
93
+ d = x.shape[-1] // 2
94
+ x_reshaped = x.view(-1, x.shape[-1])
95
+ s = x_reshaped[:, :d] * F.sigmoid(x_reshaped[:, :d])
96
+ result = s * x_reshaped[:, d:]
97
+ return result.view(*x.shape[:-1], d)
98
+
99
100
@CustomOp.register("mul_and_silu")
101
class MulAndSilu(CustomOp):
vllm/model_executor/layers/logits_processor.py
+1
@@ -53,6 +53,7 @@ class LogitsProcessor(nn.Module):
53
# Whether to use gather or all-gather to gather the logits.
54
parallel_config = get_current_vllm_config().parallel_config
55
self.use_all_gather = current_platform.is_tpu() \
56
+ or current_platform.is_neuron() \
57
or envs.VLLM_USE_V1 \
58
or parallel_config.distributed_executor_backend == "external_launcher" # noqa
59
vllm/model_executor/layers/rotary_embedding.py
+76
@@ -254,6 +254,82 @@ class RotaryEmbedding(CustomOp):
254
key = torch.cat((key_rot, key_pass), dim=-1).reshape(key_shape)
255
return query, key
256
257
+ def forward_neuron(
258
+ self,
259
+ positions: torch.Tensor,
260
+ query: torch.Tensor,
261
+ key: torch.Tensor,
262
+ offsets: Optional[torch.Tensor] = None,
263
+ ) -> Tuple[torch.Tensor, torch.Tensor]:
264
+
265
+ def _apply_rotary_emb_neuron(
266
+ x: torch.Tensor,
267
+ cos: torch.Tensor,
268
+ sin: torch.Tensor,
269
+ is_neox_style: bool,
270
+ ) -> torch.Tensor:
271
+ cos = cos.unsqueeze(-2).to(x.dtype)
272
+ sin = sin.unsqueeze(-2).to(x.dtype)
273
+ if is_neox_style:
274
+ x1, x2 = torch.chunk(x, 2, dim=-1)
275
+ else:
276
+ # x1 = x[..., ::2]
277
+
278
+ # x2 = x[..., 1::2]
279
+ d = x.shape[-1] // 2
280
+ x_reshaped = x.view(-1, x.shape[-1])
281
+ x1 = x_reshaped[:, ::2].view(*x.shape[:-1], d)
282
+ x2 = x_reshaped[:, 1::2].view(*x.shape[:-1], d)
283
+ o1 = x1 * cos - x2 * sin
284
+ o2 = x2 * cos + x1 * sin
285
+ if is_neox_style:
286
+ return torch.cat((o1, o2), dim=-1)
287
+ else:
288
+ return torch.stack((o1, o2), dim=-1).flatten(-2)
289
+
290
+ if offsets is not None:
291
+ positions = positions + offsets
292
+
293
+ self.cos_sin_cache = self.cos_sin_cache.to(query.device,
294
+ dtype=query.dtype)
295
+
296
+ positions = positions.flatten()
297
+ num_tokens = positions.shape[0]
298
+ cos_sin = self.cos_sin_cache.index_select(0, positions)
299
+ cos, sin = cos_sin.chunk(2, dim=-1)
300
+
301
+ query_shape = query.shape
302
+ query = query.view(num_tokens, -1, self.head_size)
303
+ key_shape = key.shape
304
+ key = key.view(num_tokens, -1, self.head_size)
305
+
306
+ if self.rotary_dim == self.head_size:
307
+ query = _apply_rotary_emb(query, cos, sin, self.is_neox_style)
308
+ query = query.reshape(query_shape)
309
+ key = _apply_rotary_emb(key, cos, sin, self.is_neox_style)
310
+ key = key.reshape(key_shape)
311
+ else:
312
+ head_size = query.shape[-1]
313
+ query_reshaped = query.view(-1, head_size)
314
+ query_pass = query_reshaped[:, self.rotary_dim:].view(
315
+ *query.shape[:-1], head_size - self.rotary_dim)
316
+ query_rot = query_reshaped[:, :self.rotary_dim].view(
317
+ *query.shape[:-1], self.rotary_dim)
318
+ query_rot = _apply_rotary_emb_neuron(query_rot, cos, sin,
319
+ self.is_neox_style)
320
+ query = torch.cat((query_rot, query_pass),
321
+ dim=-1).reshape(query_shape)
322
+
323
+ key_reshaped = key.view(-1, head_size)
324
+ key_pass = key_reshaped[:, self.rotary_dim:].view(
325
+ *key.shape[:-1], head_size - self.rotary_dim)
326
+ key_rot = key_reshaped[:, :self.rotary_dim].view(
327
+ *key.shape[:-1], self.rotary_dim)
328
+ key_rot = _apply_rotary_emb_neuron(key_rot, cos, sin,
329
+ self.is_neox_style)
330
+ key = torch.cat((key_rot, key_pass), dim=-1).reshape(key_shape)
331
+ return query, key
332
+
333
def extra_repr(self) -> str:
334
s = f"head_size={self.head_size}, rotary_dim={self.rotary_dim}"
335
s += f", max_position_embeddings={self.max_position_embeddings}"