@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
cu 175 lines 7.49 KB
Raw
1 /*
2 * Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
3 *
4 * Licensed under the Apache License, Version 2.0 (the "License");
5 * you may not use this file except in compliance with the License.
6 * You may obtain a copy of the License at
7 *
8 * http://www.apache.org/licenses/LICENSE-2.0
9 *
10 * Unless required by applicable law or agreed to in writing, software
11 * distributed under the License is distributed on an "AS IS" BASIS,
12 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13 * See the License for the specific language governing permissions and
14 * limitations under the License.
15 */
16
17 #include <torch/csrc/stable/tensor.h>
18
19 #include "libtorch_stable/torch_utils.h"
20
21 #include "libtorch_stable/cutlass_extensions/common.hpp"
22 #include "nvfp4_utils.cuh"
23
24 #if (defined(ENABLE_NVFP4_SM100) && ENABLE_NVFP4_SM100) || \
25 (defined(ENABLE_NVFP4_SM120) && ENABLE_NVFP4_SM120)
26 void scaled_fp4_quant_sm1xxa(torch::stable::Tensor const& output,
27 torch::stable::Tensor const& input,
28 torch::stable::Tensor const& output_sf,
29 torch::stable::Tensor const& input_sf,
30 bool is_sf_swizzled_layout);
31 #endif
32
33 #if (defined(ENABLE_NVFP4_SM100) && ENABLE_NVFP4_SM100) || \
34 (defined(ENABLE_NVFP4_SM120) && ENABLE_NVFP4_SM120)
35 void scaled_fp4_experts_quant_sm1xxa(
36 torch::stable::Tensor& output, torch::stable::Tensor& output_scale,
37 torch::stable::Tensor const& input,
38 torch::stable::Tensor const& input_global_scale,
39 torch::stable::Tensor const& input_offset_by_experts,
40 torch::stable::Tensor const& output_scale_offset_by_experts);
41 #endif
42
43 #if (defined(ENABLE_NVFP4_SM100) && ENABLE_NVFP4_SM100) || \
44 (defined(ENABLE_NVFP4_SM120) && ENABLE_NVFP4_SM120)
45 void silu_and_mul_nvfp4_quant_sm1xxa(torch::stable::Tensor& output,
46 torch::stable::Tensor& output_sf,
47 torch::stable::Tensor& input,
48 torch::stable::Tensor& input_sf);
49 #endif
50
51 #if (defined(ENABLE_NVFP4_SM100) && ENABLE_NVFP4_SM100) || \
52 (defined(ENABLE_NVFP4_SM120) && ENABLE_NVFP4_SM120)
53 void silu_and_mul_scaled_fp4_experts_quant_sm1xxa(
54 torch::stable::Tensor& output, torch::stable::Tensor& output_scale,
55 torch::stable::Tensor const& input,
56 torch::stable::Tensor const& input_global_scale,
57 torch::stable::Tensor const& input_offset_by_experts,
58 torch::stable::Tensor const& output_scale_offset_by_experts);
59 #endif
60
61 #if (defined(ENABLE_NVFP4_SM100) && ENABLE_NVFP4_SM100) || \
62 (defined(ENABLE_NVFP4_SM120) && ENABLE_NVFP4_SM120)
63 static bool nvfp4_quant_sm_supported() {
64 const int32_t sm = get_sm_version_num();
65 #if defined(ENABLE_NVFP4_SM100) && ENABLE_NVFP4_SM100
66 if (sm >= 100 && sm < 120) return true;
67 #endif
68 #if defined(ENABLE_NVFP4_SM120) && ENABLE_NVFP4_SM120
69 if (sm >= 120 && sm < 130) return true;
70 #endif
71 return false;
72 }
73 #endif
74
75 void scaled_fp4_quant_out(torch::stable::Tensor const& input,
76 torch::stable::Tensor const& input_sf,
77 bool is_sf_swizzled_layout,
78 torch::stable::Tensor& output,
79 torch::stable::Tensor& output_sf) {
80 #if (defined(ENABLE_NVFP4_SM100) && ENABLE_NVFP4_SM100) || \
81 (defined(ENABLE_NVFP4_SM120) && ENABLE_NVFP4_SM120)
82 STD_TORCH_CHECK(nvfp4_quant_sm_supported(),
83 "No compiled nvfp4 quantization kernel for SM ",
84 get_sm_version_num(),
85 ". Recompile with the appropriate CUDA arch.");
86 return scaled_fp4_quant_sm1xxa(output, input, output_sf, input_sf,
87 is_sf_swizzled_layout);
88 #endif
89 STD_TORCH_CHECK_NOT_IMPLEMENTED(false,
90 "No compiled nvfp4 quantization kernel");
91 }
92
93 std::tuple<torch::stable::Tensor, torch::stable::Tensor> scaled_fp4_quant_func(
94 torch::stable::Tensor const& input, torch::stable::Tensor const& input_sf,
95 bool is_sf_swizzled_layout) {
96 int64_t n = input.size(-1);
97 int64_t m = input.numel() / n;
98 auto device = input.device();
99
100 // Two fp4 values packed into a uint8
101 auto output = torch::stable::empty(
102 {m, n / 2}, torch::headeronly::ScalarType::Byte, std::nullopt, device);
103
104 torch::stable::Tensor output_sf;
105 if (is_sf_swizzled_layout) {
106 auto [sf_m, sf_n] = vllm::computeSwizzledSFShape(m, n);
107 output_sf = torch::stable::empty(
108 {sf_m, sf_n}, torch::headeronly::ScalarType::Int, std::nullopt, device);
109 } else {
110 output_sf = torch::stable::empty({m, n / CVT_FP4_SF_VEC_SIZE},
111 torch::headeronly::ScalarType::Byte,
112 std::nullopt, device);
113 }
114
115 scaled_fp4_quant_out(input, input_sf, is_sf_swizzled_layout, output,
116 output_sf);
117 return {output, output_sf};
118 }
119
120 void scaled_fp4_experts_quant(
121 torch::stable::Tensor& output, torch::stable::Tensor& output_scale,
122 torch::stable::Tensor const& input,
123 torch::stable::Tensor const& input_global_scale,
124 torch::stable::Tensor const& input_offset_by_experts,
125 torch::stable::Tensor const& output_scale_offset_by_experts) {
126 #if (defined(ENABLE_NVFP4_SM100) && ENABLE_NVFP4_SM100) || \
127 (defined(ENABLE_NVFP4_SM120) && ENABLE_NVFP4_SM120)
128 STD_TORCH_CHECK(nvfp4_quant_sm_supported(),
129 "No compiled nvfp4 experts quantization kernel for SM ",
130 get_sm_version_num(),
131 ". Recompile with the appropriate CUDA arch.");
132 return scaled_fp4_experts_quant_sm1xxa(
133 output, output_scale, input, input_global_scale, input_offset_by_experts,
134 output_scale_offset_by_experts);
135 #endif
136 STD_TORCH_CHECK_NOT_IMPLEMENTED(
137 false, "No compiled nvfp4 experts quantization kernel");
138 }
139
140 void silu_and_mul_nvfp4_quant(torch::stable::Tensor& output,
141 torch::stable::Tensor& output_sf,
142 torch::stable::Tensor& input,
143 torch::stable::Tensor& input_sf) {
144 #if (defined(ENABLE_NVFP4_SM100) && ENABLE_NVFP4_SM100) || \
145 (defined(ENABLE_NVFP4_SM120) && ENABLE_NVFP4_SM120)
146 STD_TORCH_CHECK(nvfp4_quant_sm_supported(),
147 "No compiled silu_and_mul nvfp4 quantization kernel for SM ",
148 get_sm_version_num(),
149 ". Recompile with the appropriate CUDA arch.");
150 return silu_and_mul_nvfp4_quant_sm1xxa(output, output_sf, input, input_sf);
151 #endif
152 STD_TORCH_CHECK_NOT_IMPLEMENTED(
153 false, "No compiled silu_and_mul nvfp4 quantization kernel");
154 }
155
156 void silu_and_mul_scaled_fp4_experts_quant(
157 torch::stable::Tensor& output, torch::stable::Tensor& output_scale,
158 torch::stable::Tensor const& input,
159 torch::stable::Tensor const& input_global_scale,
160 torch::stable::Tensor const& input_offset_by_experts,
161 torch::stable::Tensor const& output_scale_offset_by_experts) {
162 #if (defined(ENABLE_NVFP4_SM100) && ENABLE_NVFP4_SM100) || \
163 (defined(ENABLE_NVFP4_SM120) && ENABLE_NVFP4_SM120)
164 STD_TORCH_CHECK(nvfp4_quant_sm_supported(),
165 "No compiled silu_and_mul nvfp4 experts quantization kernel "
166 "for SM ",
167 get_sm_version_num(),
168 ". Recompile with the appropriate CUDA arch.");
169 return silu_and_mul_scaled_fp4_experts_quant_sm1xxa(
170 output, output_scale, input, input_global_scale, input_offset_by_experts,
171 output_scale_offset_by_experts);
172 #endif
173 STD_TORCH_CHECK_NOT_IMPLEMENTED(
174 false, "No compiled silu_and_mul nvfp4 experts quantization kernel");
175 }