@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 104 lines 3.72 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """DecoderReplayLayers runs the replay layers on the forward's replay batch."""
4
5 import pytest
6 import torch
7
8 from vllm.forward_context import (
9 ForwardContext,
10 get_forward_context,
11 override_forward_context,
12 )
13 from vllm.models.deepseek_v41.decoder_replay_layers import DecoderReplayLayers
14
15 pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA")
16
17 DEVICE = torch.device("cuda")
18 NUM_TOKENS = 401
19 HC, HIDDEN = 2, 3
20 # decode (1 token), trimmed prefill (300 of 300), untrimmed prefill (100 of 300)
21 WINDOW = 128
22 REPLAY_ROWS = [0, *range(301 - WINDOW, 301), *range(301, 401)]
23
24
25 def _context(metadata):
26 return ForwardContext(
27 no_compile_layers={},
28 attn_metadata=metadata,
29 slot_mapping={},
30 is_padding=torch.zeros(NUM_TOKENS, dtype=torch.bool, device=DEVICE),
31 )
32
33
34 def _set_replay_batch(layers, rows, metadata):
35 layers.rows = torch.tensor(rows, device=DEVICE)
36 layers.forward_context = ForwardContext(
37 no_compile_layers={}, attn_metadata=metadata, slot_mapping={}
38 )
39
40
41 def _states(num_tokens: int):
42 """(hidden_states, positions, input_ids, pre_mix, post_mix, res_mix, residual),
43 shaped as the model's: hidden states [T, H], mixes [T, HC], residual [T, HC, H]."""
44 hidden = torch.arange(num_tokens, device=DEVICE, dtype=torch.float32)
45 hidden = hidden[:, None].expand(num_tokens, HIDDEN).contiguous()
46 mix = hidden[:, :HC].contiguous()
47 residual = hidden[:, None, :].expand(num_tokens, HC, HIDDEN).contiguous()
48 return (
49 hidden,
50 torch.arange(num_tokens, device=DEVICE),
51 None,
52 mix + 1,
53 mix + 2,
54 mix + 3,
55 residual,
56 )
57
58
59 def test_run_gathers_states_and_realigns_shared_indexer_buffers():
60 topk = torch.arange(NUM_TOKENS * 4, device=DEVICE).view(NUM_TOKENS, 4).int()
61 candidates = torch.arange(NUM_TOKENS * 3, device=DEVICE).view(NUM_TOKENS, 3).int()
62 topk_before, candidates_before = topk.clone(), candidates.clone()
63 seen = {}
64
65 def run_layers(hidden_states, positions, input_ids, pre_mix, post, res, residual):
66 seen["hidden_states"] = hidden_states.clone()
67 replay_context = get_forward_context()
68 seen["attn_metadata"] = replay_context.attn_metadata
69 seen["is_padding"] = replay_context.is_padding
70 return residual, pre_mix
71
72 layers = DecoderReplayLayers(WINDOW, run_layers, [topk, candidates])
73 states = _states(NUM_TOKENS)
74 hidden, residual = states[0], states[-1]
75 full, sub = object(), object()
76 _set_replay_batch(layers, REPLAY_ROWS, sub)
77 context = _context(full)
78 with override_forward_context(context):
79 outputs = layers(*states)
80 assert get_forward_context() is context
81
82 rows = torch.tensor(REPLAY_ROWS, device=DEVICE)
83 assert torch.equal(seen["hidden_states"], hidden[rows])
84 assert seen["attn_metadata"] is sub and seen["is_padding"] is None
85 assert torch.equal(topk[: len(REPLAY_ROWS)], topk_before[rows])
86 assert torch.equal(candidates[: len(REPLAY_ROWS)], candidates_before[rows])
87 assert len(outputs) == 2 and outputs[0].shape == residual.shape
88 assert torch.equal(outputs[0][rows], residual[rows])
89 assert outputs[0][1:173].abs().sum() == 0
90
91
92 def test_no_replay_batch_runs_the_whole_batch():
93 seen = {}
94
95 def run_layers(hidden_states, *rest):
96 seen["attn_metadata"] = get_forward_context().attn_metadata
97 return (hidden_states,)
98
99 layers = DecoderReplayLayers(WINDOW, run_layers, [])
100 states = _states(NUM_TOKENS)
101 full = object()
102 with override_forward_context(_context(full)):
103 (output,) = layers(*states)
104 assert output is states[0] and seen["attn_metadata"] is full