@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 118 lines 3.87 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.model_executor.layers.fused_qk_norm_rope import fused_qk_rmsnorm_rope_gate
8 from vllm.model_executor.layers.layernorm import GemmaRMSNorm
9 from vllm.model_executor.layers.rotary_embedding import (
10 MRotaryEmbedding,
11 RotaryEmbedding,
12 )
13 from vllm.platforms import current_platform
14 from vllm.utils.torch_utils import set_random_seed
15
16 # Qwen3.6 TP=1 attention geometry.
17 HEAD_DIM = 256
18 ROTARY_DIM = 64
19 RMS_NORM_EPS = 1e-6
20 MAX_POSITION_EMBEDDINGS = 262144
21 ROPE_THETA = 10000000.0
22 DTYPE = torch.bfloat16
23 SEED = 13
24 MROPE_SECTION = (11, 11, 10)
25 ROPE_CASES = [
26 pytest.param(24, 4, None, id="rope"),
27 pytest.param(16, 2, MROPE_SECTION, id="interleaved-mrope"),
28 ]
29
30
31 @pytest.mark.skipif(
32 not current_platform.is_cuda_alike(),
33 reason="fused_qk_rmsnorm_rope_gate Triton kernel requires CUDA/ROCm",
34 )
35 @pytest.mark.parametrize("num_q_heads,num_kv_heads,mrope_section", ROPE_CASES)
36 @pytest.mark.parametrize("num_tokens", [1, 4, 37])
37 @torch.inference_mode()
38 def test_fused_qk_norm_rope_gate_matches_reference(
39 default_vllm_config,
40 num_tokens: int,
41 num_q_heads: int,
42 num_kv_heads: int,
43 mrope_section: tuple[int, int, int] | None,
44 ) -> None:
45 device = torch.device("cuda", torch.accelerator.current_device_index())
46 torch.set_default_device(device)
47 set_random_seed(SEED)
48
49 q_gate = torch.randn(
50 num_tokens, num_q_heads * 2 * HEAD_DIM, dtype=DTYPE, device=device
51 )
52 k = torch.randn(num_tokens, num_kv_heads * HEAD_DIM, dtype=DTYPE, device=device)
53
54 q_norm = GemmaRMSNorm(HEAD_DIM, eps=RMS_NORM_EPS).to(device, dtype=DTYPE)
55 k_norm = GemmaRMSNorm(HEAD_DIM, eps=RMS_NORM_EPS).to(device, dtype=DTYPE)
56 q_norm.weight.normal_(std=0.1)
57 k_norm.weight.normal_(std=0.1)
58
59 q_gate_heads = q_gate.view(num_tokens, num_q_heads, 2 * HEAD_DIM)
60 q = q_gate_heads[..., :HEAD_DIM]
61 gate_ref = q_gate_heads[..., HEAD_DIM:].reshape(num_tokens, num_q_heads * HEAD_DIM)
62 q_ref = q_norm.forward_native(q)
63 k_ref = k_norm.forward_native(k.view(num_tokens, num_kv_heads, HEAD_DIM))
64 assert isinstance(q_ref, torch.Tensor)
65 assert isinstance(k_ref, torch.Tensor)
66 q_ref = q_ref.reshape(num_tokens, num_q_heads * HEAD_DIM)
67 k_ref = k_ref.reshape(num_tokens, num_kv_heads * HEAD_DIM)
68
69 if mrope_section is None:
70 rope = RotaryEmbedding(
71 HEAD_DIM,
72 ROTARY_DIM,
73 MAX_POSITION_EMBEDDINGS,
74 ROPE_THETA,
75 True,
76 DTYPE,
77 ).to(device)
78 positions = torch.arange(num_tokens, dtype=torch.long, device=device)
79 else:
80 rope = MRotaryEmbedding(
81 HEAD_DIM,
82 ROTARY_DIM,
83 MAX_POSITION_EMBEDDINGS,
84 ROPE_THETA,
85 True,
86 DTYPE,
87 mrope_section=list(mrope_section),
88 mrope_interleaved=True,
89 ).to(device)
90 positions = torch.arange(3 * num_tokens, dtype=torch.long, device=device).view(
91 3, num_tokens
92 )
93 assert torch.unique(positions[:, 0]).numel() == 3
94
95 q_ref, k_ref = rope.forward_native(positions, q_ref, k_ref)
96 assert k_ref is not None
97
98 q_out, k_out, gate_out = fused_qk_rmsnorm_rope_gate(
99 q_gate,
100 k,
101 q_norm.weight,
102 k_norm.weight,
103 rope.cos_sin_cache,
104 positions,
105 RMS_NORM_EPS,
106 num_q_heads,
107 num_kv_heads,
108 HEAD_DIM,
109 ROTARY_DIM,
110 mrope_section=mrope_section,
111 norm_beta=1.0,
112 )
113
114 # The built-in reference performs RoPE in BF16, while the fused kernel
115 # promotes the BF16-normalized values to FP32 for RoPE before storing BF16.
116 torch.testing.assert_close(q_out, q_ref, atol=1e-2, rtol=1e-2)
117 torch.testing.assert_close(k_out, k_ref, atol=1e-2, rtol=1e-2)
118 torch.testing.assert_close(gate_out, gate_ref, atol=0, rtol=0)