@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 157 lines 5.74 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.models.qwen4_exp.nvidia.ops.cute_dsl.hc_down_silu import hc_down_silu
8 from vllm.models.qwen4_exp.nvidia.ops.hc import (
9 grouped_gemma_rmsnorm,
10 hc_combine,
11 hc_combine_norm,
12 hc_gate_mix,
13 hc_silu,
14 )
15 from vllm.platforms import current_platform
16 from vllm.triton_utils import HAS_TRITON
17
18 pytestmark = pytest.mark.skipif(
19 not current_platform.is_cuda() or not HAS_TRITON,
20 reason="HC kernels require CUDA and Triton",
21 )
22
23 HC = 4
24 HIDDEN_SIZE = 2560
25 HYPER_HIDDEN_SIZE = HC * HIDDEN_SIZE
26 EPS = 1e-6
27 LORA_RANK = 320
28 DOWN_N = LORA_RANK + HC + 12 # merged down+inject weight, 16-row padded
29
30 requires_sm90 = pytest.mark.skipif(
31 not current_platform.has_device_capability(90),
32 reason="fused HC down+SiLU requires SM90+",
33 )
34
35
36 def test_grouped_gemma_rmsnorm() -> None:
37 torch.manual_seed(0)
38 x = torch.randn(2, HYPER_HIDDEN_SIZE, dtype=torch.bfloat16, device="cuda")
39 weight = torch.randn(HYPER_HIDDEN_SIZE, dtype=torch.bfloat16, device="cuda")
40
41 actual = grouped_gemma_rmsnorm(x, weight, EPS, HC)
42
43 grouped = x.float().unflatten(-1, (HC, HIDDEN_SIZE))
44 variance = grouped.square().mean(-1, keepdim=True)
45 expected = grouped * torch.rsqrt(variance + EPS)
46 expected = expected.flatten(-2) * (1.0 + weight.float())
47 torch.testing.assert_close(actual, expected.to(torch.bfloat16))
48
49
50 def test_hc_gate_mix() -> None:
51 torch.manual_seed(0)
52 x = torch.randn(2, HYPER_HIDDEN_SIZE, dtype=torch.bfloat16, device="cuda")
53 gate = torch.randn(2, HYPER_HIDDEN_SIZE, dtype=torch.bfloat16, device="cuda")
54
55 actual = hc_gate_mix(x, gate, HC)
56 expected = (
57 torch.sigmoid(gate.float().unflatten(-1, (HC, HIDDEN_SIZE)))
58 * x.float().unflatten(-1, (HC, HIDDEN_SIZE))
59 ).mean(-2)
60
61 torch.testing.assert_close(actual, expected.to(torch.bfloat16))
62
63
64 def test_hc_combine() -> None:
65 torch.manual_seed(0)
66 block_output = torch.randn(2, HIDDEN_SIZE, dtype=torch.bfloat16, device="cuda")
67 residual = torch.randn(2, HYPER_HIDDEN_SIZE, dtype=torch.bfloat16, device="cuda")
68 injection = torch.randn(2, HC, dtype=torch.bfloat16, device="cuda")
69
70 actual = hc_combine(residual, block_output, injection, HC)
71 injection_weight = 2.0 * torch.sigmoid(injection.float() / HC)
72 expected = residual.float().unflatten(-1, (HC, HIDDEN_SIZE))
73 expected = expected + block_output.float().unsqueeze(
74 -2
75 ) * injection_weight.unsqueeze(-1)
76
77 torch.testing.assert_close(actual, expected.flatten(-2).to(torch.bfloat16))
78
79
80 def test_hc_combine_unit_injection() -> None:
81 torch.manual_seed(0)
82 block_output = torch.randn(2, HIDDEN_SIZE, dtype=torch.bfloat16, device="cuda")
83 residual = torch.randn(2, HYPER_HIDDEN_SIZE, dtype=torch.bfloat16, device="cuda")
84
85 actual = hc_combine(residual, block_output, None, HC)
86 expected = residual.unflatten(-1, (HC, HIDDEN_SIZE))
87 expected = expected + block_output.unsqueeze(-2)
88
89 assert torch.equal(actual, expected.flatten(-2))
90
91
92 def test_hc_combine_norm() -> None:
93 torch.manual_seed(0)
94 block_output = torch.randn(2, HIDDEN_SIZE, dtype=torch.bfloat16, device="cuda")
95 residual = torch.randn(2, HYPER_HIDDEN_SIZE, dtype=torch.bfloat16, device="cuda")
96 injection = torch.randn(2, HC, dtype=torch.bfloat16, device="cuda")
97 weight = torch.randn(HYPER_HIDDEN_SIZE, dtype=torch.bfloat16, device="cuda")
98
99 actual, actual_norm = hc_combine_norm(
100 residual, block_output, injection, weight, EPS, HC
101 )
102
103 injection_weight = 2.0 * torch.sigmoid(injection.float() / HC)
104 expected = residual.float().unflatten(-1, (HC, HIDDEN_SIZE))
105 expected = expected + block_output.float().unsqueeze(
106 -2
107 ) * injection_weight.unsqueeze(-1)
108 expected = expected.flatten(-2).to(residual.dtype)
109 grouped = expected.float().unflatten(-1, (HC, HIDDEN_SIZE))
110 variance = grouped.square().mean(-1, keepdim=True)
111 expected_norm = grouped * torch.rsqrt(variance + EPS)
112 expected_norm = expected_norm.flatten(-2) * (1.0 + weight.float())
113
114 torch.testing.assert_close(actual, expected)
115 torch.testing.assert_close(actual_norm, expected_norm.to(torch.bfloat16))
116
117
118 @pytest.mark.parametrize("num_tokens", [1, 17, 2048])
119 def test_hc_combine_norm_unit_injection(num_tokens: int) -> None:
120 torch.manual_seed(0)
121 embedding = torch.randn(
122 num_tokens, HIDDEN_SIZE, dtype=torch.bfloat16, device="cuda"
123 )
124 hidden = torch.randn(
125 num_tokens, HC, HIDDEN_SIZE, dtype=torch.bfloat16, device="cuda"
126 )
127 weight = torch.randn(HYPER_HIDDEN_SIZE, dtype=torch.bfloat16, device="cuda")
128
129 actual, actual_norm = hc_combine_norm(
130 hidden.flatten(1), embedding, None, weight, EPS, HC
131 )
132
133 expected = (hidden + embedding.unsqueeze(1)).flatten(1)
134 expected_norm = grouped_gemma_rmsnorm(expected, weight, EPS, HC)
135 assert torch.equal(actual, expected)
136 torch.testing.assert_close(actual_norm, expected_norm)
137
138
139 @requires_sm90
140 @pytest.mark.parametrize("num_tokens", [1, 3, 5, 17, 48])
141 def test_hc_down_silu_fused(num_tokens: int) -> None:
142 # Compare computed columns with the unfused ll_bf16 + hc_silu reference.
143 from vllm.model_executor.kernels.linear.cute_dsl.ll_bf16 import ll_bf16_gemm
144
145 torch.manual_seed(0)
146 x = torch.randn(num_tokens, HYPER_HIDDEN_SIZE, dtype=torch.bfloat16, device="cuda")
147 weight = torch.randn(DOWN_N, HYPER_HIDDEN_SIZE, dtype=torch.bfloat16, device="cuda")
148
149 lora, injection = hc_down_silu(x, weight, LORA_RANK, HC)
150
151 down = ll_bf16_gemm(x, weight).to(torch.bfloat16)
152 torch.testing.assert_close(
153 lora, hc_silu(down[:, :LORA_RANK], HC), rtol=0.01, atol=0.01
154 )
155 torch.testing.assert_close(
156 injection, down[:, LORA_RANK : LORA_RANK + HC], rtol=0.01, atol=0.01
157 )