@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 185 lines 7.08 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 import pytest
5 import torch
6
7 from vllm import _custom_ops as ops
8
9 NUM_TOKENS = [1, 4, 16, 64, 128]
10 NUM_HEADS = [128]
11 NOPE_DIM = [512]
12 ROPE_DIM = [64]
13 DTYPES = [torch.bfloat16, torch.float16]
14
15
16 @pytest.mark.parametrize("num_tokens", NUM_TOKENS)
17 @pytest.mark.parametrize("num_heads", NUM_HEADS)
18 @pytest.mark.parametrize("nope_dim", NOPE_DIM)
19 @pytest.mark.parametrize("rope_dim", ROPE_DIM)
20 @pytest.mark.parametrize("dtype", DTYPES)
21 def test_concat_mla_q_contiguous(num_tokens, num_heads, nope_dim, rope_dim, dtype):
22 """Test with contiguous inputs (standard layout)."""
23 torch.manual_seed(42)
24 ql_nope = torch.randn(num_tokens, num_heads, nope_dim, dtype=dtype, device="cuda")
25 q_pe = torch.randn(num_tokens, num_heads, rope_dim, dtype=dtype, device="cuda")
26
27 ref = torch.cat((ql_nope, q_pe), dim=-1)
28
29 q_out = torch.empty(
30 num_tokens, num_heads, nope_dim + rope_dim, dtype=dtype, device="cuda"
31 )
32 ops.concat_mla_q(ql_nope, q_pe, q_out)
33
34 torch.testing.assert_close(q_out, ref, atol=0, rtol=0)
35
36
37 def test_concat_mla_q_padded_output_heads():
38 """The fused concat can write directly into a head-padded workspace."""
39 ql_nope = torch.randn(4, 8, 512, dtype=torch.bfloat16, device="cuda")
40 q_pe = torch.randn(4, 8, 64, dtype=torch.bfloat16, device="cuda")
41 q_out = torch.empty(4, 64, 576, dtype=torch.bfloat16, device="cuda")
42
43 ops.concat_mla_q(ql_nope, q_pe, q_out)
44
45 torch.testing.assert_close(q_out[:, :8], torch.cat((ql_nope, q_pe), dim=-1))
46
47
48 @pytest.mark.parametrize("num_tokens", [t for t in NUM_TOKENS if t > 1])
49 @pytest.mark.parametrize("num_heads", NUM_HEADS)
50 @pytest.mark.parametrize("nope_dim", NOPE_DIM)
51 @pytest.mark.parametrize("rope_dim", ROPE_DIM)
52 @pytest.mark.parametrize("dtype", DTYPES)
53 def test_concat_mla_q_transposed_nope(num_tokens, num_heads, nope_dim, rope_dim, dtype):
54 """Test with transposed nope input (simulates BMM output after transpose).
55
56 In the real code path, mqa_ql_nope is the result of:
57 torch.bmm(q_nope, W_UK_T) # [N, B, L]
58 .transpose(0, 1) # [B, N, L] — non-contiguous!
59 """
60 torch.manual_seed(42)
61 nope_raw = torch.randn(num_heads, num_tokens, nope_dim, dtype=dtype, device="cuda")
62 ql_nope = nope_raw.transpose(0, 1) # [B, N, L], non-contiguous
63 assert not ql_nope.is_contiguous()
64
65 q_pe = torch.randn(num_tokens, num_heads, rope_dim, dtype=dtype, device="cuda")
66
67 ref = torch.cat((ql_nope, q_pe), dim=-1)
68
69 q_out = torch.empty(
70 num_tokens, num_heads, nope_dim + rope_dim, dtype=dtype, device="cuda"
71 )
72 ops.concat_mla_q(ql_nope, q_pe, q_out)
73
74 torch.testing.assert_close(q_out, ref, atol=0, rtol=0)
75
76
77 @pytest.mark.parametrize("num_tokens", NUM_TOKENS)
78 @pytest.mark.parametrize("num_heads", NUM_HEADS)
79 @pytest.mark.parametrize("dtype", DTYPES)
80 def test_concat_mla_q_split_rope(num_tokens, num_heads, dtype):
81 """Test with rope from a split (simulates the actual code path).
82
83 In the real code path, q_pe comes from:
84 mqa_q.split([qk_nope_head_dim, qk_rope_head_dim], dim=-1)
85 which creates a non-contiguous view with stride(1) != rope_dim.
86 """
87 torch.manual_seed(42)
88 nope_dim = 512
89 rope_dim = 64
90 orig_dim = 128 + 64 # original q before absorption: [B, N, 192]
91
92 # Simulate split from original q tensor
93 q_orig = torch.randn(num_tokens, num_heads, orig_dim, dtype=dtype, device="cuda")
94 q_nope_orig, q_pe = q_orig.split([128, 64], dim=-1)
95
96 # q_pe is non-contiguous: stride(1) = 192, not 64
97 assert q_pe.stride(1) == orig_dim
98 assert q_pe.stride(2) == 1 # but innermost is fine
99
100 # Simulate absorbed nope (contiguous, different size)
101 ql_nope = torch.randn(num_tokens, num_heads, nope_dim, dtype=dtype, device="cuda")
102
103 ref = torch.cat((ql_nope, q_pe), dim=-1)
104
105 q_out = torch.empty(
106 num_tokens, num_heads, nope_dim + rope_dim, dtype=dtype, device="cuda"
107 )
108 ops.concat_mla_q(ql_nope, q_pe, q_out)
109
110 torch.testing.assert_close(q_out, ref, atol=0, rtol=0)
111
112
113 def test_concat_mla_q_zero_tokens():
114 """Test with zero tokens (edge case)."""
115 ql_nope = torch.empty(0, 128, 512, dtype=torch.bfloat16, device="cuda")
116 q_pe = torch.empty(0, 128, 64, dtype=torch.bfloat16, device="cuda")
117 q_out = torch.empty(0, 128, 576, dtype=torch.bfloat16, device="cuda")
118
119 ops.concat_mla_q(ql_nope, q_pe, q_out)
120
121
122 @pytest.mark.parametrize("num_tokens", [1, 64])
123 def test_concat_mla_q_values_preserved(num_tokens):
124 """Verify exact bit-level preservation (no computation, pure copy).
125
126 Compares raw int16 bits to avoid NaN != NaN issues from IEEE 754.
127 """
128 nope_dim, rope_dim = 512, 64
129
130 # Use specific bit patterns (stay in int16 for bit-exact comparison)
131 ql_nope_bits = torch.arange(
132 num_tokens * 128 * nope_dim, dtype=torch.int16, device="cuda"
133 ).view(num_tokens, 128, nope_dim)
134 q_pe_bits = torch.arange(
135 num_tokens * 128 * rope_dim, dtype=torch.int16, device="cuda"
136 ).view(num_tokens, 128, rope_dim)
137
138 ql_nope = ql_nope_bits.view(torch.bfloat16)
139 q_pe = q_pe_bits.view(torch.bfloat16)
140
141 q_out = torch.empty(
142 num_tokens, 128, nope_dim + rope_dim, dtype=torch.bfloat16, device="cuda"
143 )
144 ops.concat_mla_q(ql_nope, q_pe, q_out)
145
146 out_bits = q_out.view(torch.int16)
147
148 assert torch.equal(out_bits[..., :nope_dim], ql_nope_bits)
149
150 assert torch.equal(out_bits[..., nope_dim:], q_pe_bits)
151
152
153 @pytest.mark.parametrize("rope_dim,transposed", [(0, False), (0, True), (64, False)])
154 @pytest.mark.parametrize("enabled", [False, True])
155 def test_concat_mla_q_fp8_nope_and_rope(rope_dim, transposed, enabled):
156 """FP8 query packing preserves values with empty RoPE and strided BMM output."""
157 from vllm.config import CompilationConfig, VllmConfig, set_current_vllm_config
158 from vllm.model_executor.layers.attention.mla_attention import _DecodeConcatQuantFP8
159 from vllm.model_executor.layers.quantization.utils.quant_utils import (
160 GroupShape,
161 get_fp8_min_max,
162 )
163 from vllm.platforms import current_platform
164
165 torch.manual_seed(42)
166 shape = (8, 17, 512) if transposed else (17, 8, 512)
167 ql = torch.randn(shape, dtype=torch.bfloat16, device="cuda") * 128
168 if transposed:
169 ql = ql.transpose(0, 1)
170 qp = torch.randn(17, 8, rope_dim, dtype=torch.bfloat16, device="cuda") * 128
171 scale = torch.tensor(0.3, device="cuda", dtype=torch.float32)
172 joined = torch.cat((ql, qp), dim=-1)
173 fp8_min, fp8_max = get_fp8_min_max()
174 expected = (joined.float() * scale.reciprocal()).clamp(fp8_min, fp8_max)
175 expected = expected.to(current_platform.fp8_dtype())
176 custom_ops = ["none"] + (["+mla_decode_concat_quant_fp8"] if enabled else [])
177 config = VllmConfig(
178 compilation_config=CompilationConfig(mode=0, custom_ops=custom_ops)
179 )
180 with set_current_vllm_config(config):
181 quant = _DecodeConcatQuantFP8(static=True, group_shape=GroupShape.PER_TENSOR)
182 actual = quant(ql, qp, scale)
183 assert actual.dtype == current_platform.fp8_dtype()
184 torch.testing.assert_close(actual.float(), expected.float(), atol=0, rtol=0)
185 assert actual.is_contiguous()