@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 141 lines 5.22 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 import torch
5
6 import vllm.model_executor.offloader as offloader
7 from vllm.models.deepseek_v4.amd.rocm import (
8 DeepseekV4ROCMAiterMLAAttention,
9 )
10
11
12 class _WeightOnlyLinear(torch.nn.Module):
13 def __init__(self, weight: torch.Tensor) -> None:
14 super().__init__()
15 self.weight = torch.nn.Parameter(weight, requires_grad=False)
16
17
18 class _Compressor(torch.nn.Module):
19 def __init__(self, weight: torch.Tensor) -> None:
20 super().__init__()
21 self.fused_wkv_wgate = _WeightOnlyLinear(weight)
22
23
24 class _IndexerWeightsProjection(torch.nn.Module):
25 def forward(self, hidden_states: torch.Tensor) -> tuple[torch.Tensor, None]:
26 return hidden_states.sum(dim=-1, keepdim=True), None
27
28
29 class _Indexer(torch.nn.Module):
30 def __init__(self, compressor_weight: torch.Tensor) -> None:
31 super().__init__()
32 self.compressor = _Compressor(compressor_weight)
33 self.weights_proj = _IndexerWeightsProjection()
34
35
36 def _make_attention(
37 main_weight: torch.Tensor,
38 indexer_weight: torch.Tensor,
39 ) -> DeepseekV4ROCMAiterMLAAttention:
40 attention = DeepseekV4ROCMAiterMLAAttention.__new__(DeepseekV4ROCMAiterMLAAttention)
41 torch.nn.Module.__init__(attention)
42 attention.compressor = _Compressor(main_weight)
43 attention.indexer = _Indexer(indexer_weight)
44 attention.register_buffer("_fused_compressor_weight", None, persistent=False)
45 attention._fused_compressor_split_sizes = None
46 return attention
47
48
49 def test_prepare_compressor_gemm_fusion_aliases_original_weights():
50 main_value = torch.arange(12, dtype=torch.float32).reshape(4, 3)
51 indexer_value = torch.arange(6, dtype=torch.float32).reshape(2, 3) + 100
52 attention = _make_attention(main_value, indexer_value)
53 main_weight = attention.compressor.fused_wkv_wgate.weight
54 indexer_weight = attention.indexer.compressor.fused_wkv_wgate.weight
55
56 assert attention.prepare_compressor_gemm_fusion() is True
57
58 fused_weight = attention._fused_compressor_weight
59 assert fused_weight is not None
60 assert attention._fused_compressor_split_sizes == (4, 2)
61 torch.testing.assert_close(fused_weight[:4], main_value)
62 torch.testing.assert_close(fused_weight[4:], indexer_value)
63
64 fused_storage = fused_weight.untyped_storage().data_ptr()
65 assert main_weight.untyped_storage().data_ptr() == fused_storage
66 assert indexer_weight.untyped_storage().data_ptr() == fused_storage
67 assert main_weight.storage_offset() == 0
68 assert indexer_weight.storage_offset() == main_weight.numel()
69
70 state_dict = attention.state_dict()
71 assert "_fused_compressor_weight" not in state_dict
72 torch.testing.assert_close(
73 state_dict["compressor.fused_wkv_wgate.weight"], main_value
74 )
75 torch.testing.assert_close(
76 state_dict["indexer.compressor.fused_wkv_wgate.weight"],
77 indexer_value,
78 )
79
80 fused_data_ptr = fused_weight.data_ptr()
81 assert attention.prepare_compressor_gemm_fusion() is False
82 assert attention._fused_compressor_weight.data_ptr() == fused_data_ptr
83
84
85 def test_prepare_compressor_gemm_fusion_skips_weight_offloading(monkeypatch):
86 attention = _make_attention(torch.ones(4, 3), torch.ones(2, 3))
87 main_weight = attention.compressor.fused_wkv_wgate.weight
88 indexer_weight = attention.indexer.compressor.fused_wkv_wgate.weight
89
90 monkeypatch.setattr(offloader, "get_offloader", object)
91
92 assert attention.prepare_compressor_gemm_fusion() is False
93 assert attention._fused_compressor_weight is None
94 assert main_weight.untyped_storage().data_ptr() != (
95 indexer_weight.untyped_storage().data_ptr()
96 )
97
98
99 def test_fused_input_projection_uses_one_mm_and_returns_expected_tuple(
100 monkeypatch,
101 ):
102 main_weight = torch.tensor([[1.0, 0.0, 0.0], [0.0, 1.0, 0.0]], dtype=torch.float32)
103 indexer_weight = torch.tensor([[0.0, 0.0, 1.0]], dtype=torch.float32)
104 attention = _make_attention(main_weight, indexer_weight)
105 assert attention.prepare_compressor_gemm_fusion() is True
106
107 hidden_states = torch.tensor(
108 [[2.0, 3.0, 5.0], [7.0, 11.0, 13.0]], dtype=torch.float32
109 )
110 qr_kv = torch.full((2, 4), 17.0)
111 monkeypatch.setattr(
112 DeepseekV4ROCMAiterMLAAttention,
113 "_fused_wqa_wkv_gemm",
114 lambda self, hidden: qr_kv,
115 )
116
117 mm_calls: list[tuple[torch.Tensor, torch.Tensor, torch.dtype | None]] = []
118
119 def cpu_mm_with_out_dtype(
120 left: torch.Tensor,
121 right: torch.Tensor,
122 *,
123 out_dtype: torch.dtype | None = None,
124 ) -> torch.Tensor:
125 mm_calls.append((left, right, out_dtype))
126 return torch.matmul(left, right).to(out_dtype)
127
128 monkeypatch.setattr(torch, "mm", cpu_mm_with_out_dtype)
129
130 qr_out, kv_score, indexer_kv_score, indexer_weights = (
131 attention._run_parallel_input_projections(hidden_states)
132 )
133
134 assert qr_out is qr_kv
135 torch.testing.assert_close(kv_score, hidden_states[:, :2])
136 torch.testing.assert_close(indexer_kv_score, hidden_states[:, 2:])
137 torch.testing.assert_close(indexer_weights, hidden_states.sum(dim=-1, keepdim=True))
138 assert len(mm_calls) == 1
139 assert mm_calls[0][0] is hidden_states
140 assert mm_calls[0][2] is torch.float32
141 assert mm_calls[0][1].shape == (3, 3)