@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 170 lines 5.83 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.v1.attention.ops.xpu_mla_sparse import triton_bf16_mla_sparse_interface
8
9
10 # https://github.com/deepseek-ai/FlashMLA/blob/main/tests/ref.py#L7
11 def _merge_two_lse(
12 lse0: torch.Tensor, lse1: torch.Tensor | None, s_q: int, h_q: int
13 ) -> torch.Tensor:
14 if lse1 is None:
15 return lse0
16 else:
17 return torch.logsumexp(
18 torch.stack([lse0.view(s_q, h_q), lse1.broadcast_to(s_q, h_q)], dim=0),
19 dim=0,
20 )
21
22
23 # Adapted from https://github.com/deepseek-ai/FlashMLA/blob/main/tests/ref.py#L19
24 def reference_mla_sparse_prefill(
25 q: torch.Tensor,
26 kv: torch.Tensor,
27 indices: torch.Tensor,
28 sm_scale: float,
29 d_v: int,
30 topk_length: torch.Tensor | None = None,
31 attn_sink: torch.Tensor | None = None,
32 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
33 """Returns:
34 - o: [s_q, h_q, dv]
35 - o_fp32: [s_q, h_q, dv]
36 - max_logits: [s_q, h_q]
37 - lse: [s_q, h_q]
38
39 """
40 s_q, h_q, d_qk = q.shape
41 s_kv, _, _ = kv.shape
42 _, _, topk = indices.shape
43
44 indices = indices.clone().squeeze(1)
45 if topk_length is not None:
46 mask = torch.arange(topk, device=topk_length.device).unsqueeze(0).broadcast_to(
47 s_q, topk
48 ) >= topk_length.unsqueeze(1) # [s_q, topk]
49 indices[mask] = -1
50 invalid_mask = (indices < 0) | (indices >= s_kv) # [s_q, topk]
51 indices[invalid_mask] = 0
52
53 q = q.float()
54 gathered_kv = (
55 kv.index_select(dim=0, index=indices.flatten()).reshape(s_q, topk, d_qk).float()
56 ) # [s_q, topk, d_qk]
57 P = q @ gathered_kv.transpose(1, 2) # [s_q, h_q, topk]
58 P *= sm_scale
59 P[invalid_mask.unsqueeze(1).broadcast_to(P.shape)] = float("-inf")
60
61 orig_lse = torch.logsumexp(P, dim=-1) # [s_q, h_q]
62 max_logits = P.max(dim=-1).values # [s_q, h_q]
63
64 lse_for_o = _merge_two_lse(orig_lse, attn_sink, s_q, h_q)
65 if not torch.is_inference_mode_enabled():
66 lse_for_o = lse_for_o.clone()
67 lse_for_o[lse_for_o == float("-inf")] = float(
68 "+inf"
69 ) # So that corresponding O will be 0
70 s_for_o = torch.exp(P - lse_for_o.unsqueeze(-1))
71 out = s_for_o @ gathered_kv[..., :d_v] # [s_q, h_q, dv]
72
73 lonely_q_mask = orig_lse == float("-inf") # [s_q, h_q]
74 orig_lse[lonely_q_mask] = float("+inf")
75 return (out.to(kv.dtype), out, max_logits, orig_lse)
76
77
78 @pytest.mark.parametrize("device_str", ["xpu"])
79 @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
80 @pytest.mark.skipif(
81 not torch.xpu.is_available(),
82 reason="XPU is required",
83 )
84 def test_bf16_triton_sparse_mla(device_str, dtype):
85 device = torch.device(device_str)
86 s_q = 1
87 s_kv = 256
88 h_q = 64 # kernel expects multiple of 64
89 h_kv = 1
90 d_qk = 576
91 d_v = 512
92 topk = 128
93
94 torch.random.manual_seed(1234)
95
96 q = torch.randn((s_q, h_q, d_qk), dtype=dtype, device=device)
97 kv = torch.randn((s_kv, h_kv, d_qk), dtype=dtype, device=device)
98 indices = torch.full((s_q, h_kv, topk), -1, dtype=torch.int32, device=device)
99 for t in range(s_q):
100 for h in range(h_kv):
101 i_i = torch.randperm(max(1, t))[:topk]
102 indices[t, h, : len(i_i)] = i_i
103
104 sm_scale = d_qk**-0.5
105
106 out, max_logits, lse = triton_bf16_mla_sparse_interface(
107 q, kv, indices, sm_scale, d_v
108 )
109 assert out.shape == (s_q, h_q, d_v)
110 assert max_logits.shape == (s_q, h_q)
111 assert lse.shape == (s_q, h_q)
112
113 ref_out, ref_out_fp32, ref_max_logits, ref_lse = reference_mla_sparse_prefill(
114 q, kv, indices, sm_scale, d_v
115 )
116 assert torch.allclose(out, ref_out, atol=1e-2, rtol=1e-2)
117 assert torch.allclose(max_logits, ref_max_logits, atol=1e-3, rtol=1e-3)
118 assert torch.allclose(lse, ref_lse, atol=1e-3, rtol=1e-3)
119
120
121 @pytest.mark.parametrize("device_str", ["xpu"])
122 @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
123 @pytest.mark.skipif(
124 not torch.xpu.is_available(),
125 reason="XPU is required",
126 )
127 def test_bf16_triton_sparse_mla_masked_chunks(device_str, dtype):
128 """Rows whose leading BLOCK_N index entries are all masked must not NaN.
129
130 Regression test: with an -inf running max, a fully-masked leading chunk
131 produced re_scale = exp2(-inf - -inf) = NaN, permanently poisoning the
132 accumulator even though valid keys followed in later chunks.
133 """
134 device = torch.device(device_str)
135 s_q = 3
136 s_kv = 256
137 h_q = 64
138 h_kv = 1
139 d_qk = 576
140 d_v = 512
141 topk = 128 # 8 chunks of BLOCK_N=16
142
143 torch.random.manual_seed(1234)
144
145 q = torch.randn((s_q, h_q, d_qk), dtype=dtype, device=device)
146 kv = torch.randn((s_kv, h_kv, d_qk), dtype=dtype, device=device)
147 indices = torch.full((s_q, h_kv, topk), -1, dtype=torch.int32, device=device)
148 # row 0: valid keys only in chunks 1-2 -> leading AND trailing masked chunks
149 indices[0, 0, 16:48] = torch.arange(32, dtype=torch.int32, device=device)
150 # row 1: fully valid
151 indices[1, 0, :] = torch.arange(topk, dtype=torch.int32, device=device)
152 # row 2: no valid key at all
153
154 sm_scale = d_qk**-0.5
155
156 out, max_logits, lse = triton_bf16_mla_sparse_interface(
157 q, kv, indices, sm_scale, d_v
158 )
159 assert out.isfinite().all()
160
161 ref_out, _, ref_max_logits, ref_lse = reference_mla_sparse_prefill(
162 q, kv, indices, sm_scale, d_v
163 )
164 assert torch.allclose(out[:2], ref_out[:2], atol=1e-2, rtol=1e-2)
165 assert torch.allclose(max_logits[:2], ref_max_logits[:2], atol=1e-3, rtol=1e-3)
166 assert torch.allclose(lse[:2], ref_lse[:2], atol=1e-3, rtol=1e-3)
167 # A row with no valid key yields zeros (the reference's convention); its
168 # lse/max_logits are large-negative finite rather than the reference's
169 # +inf/-inf placeholders, so only the output is compared here.
170 assert torch.allclose(out[2], torch.zeros_like(out[2]))