@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
cu 153 lines 6.26 KB
Raw
1 // SPDX-License-Identifier: Apache-2.0
2 // SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 #include <torch/csrc/stable/library.h>
5 #include <torch/csrc/stable/tensor.h>
6 #include <torch/headeronly/core/ScalarType.h>
7
8 #include "core/registration.h"
9 #include "libtorch_stable/torch_utils.h"
10
11 #include <cuda_bf16.h>
12 #include <cuda_runtime.h>
13
14 #include <stdexcept>
15
16 namespace {
17
18 inline int getSMVersion() {
19 auto* props = get_device_prop();
20 return props->major * 10 + props->minor;
21 }
22
23 } // namespace
24
25 static constexpr int FP32_MAX_TOKENS = 32;
26
27 // Supported (hidden_dim, num_experts) pairs (must match the instantiations in
28 // fp32_router_gemm.cu): (3072, 256) for MiniMax-M2/M2.5, (6144, 128) for M3.
29 static inline bool fp32_router_gemm_supported(int hidden_dim, int num_experts) {
30 return (hidden_dim == 3072 && num_experts == 256) ||
31 (hidden_dim == 6144 && num_experts == 128);
32 }
33
34 // Forward declarations — 4 template params must match fp32_router_gemm.cu
35 template <typename InputT, int kNumTokens, int kNumExperts, int kHiddenDim>
36 void invokeFp32RouterGemm(float* output, InputT const* mat_a,
37 float const* mat_b, cudaStream_t stream);
38
39 // LoopUnroller templated on InputT, kNumExperts and kHiddenDim
40 template <typename InputT, int kNumExperts, int kHiddenDim, int kBegin,
41 int kEnd>
42 struct Fp32LoopUnroller {
43 static void unroll(int num_tokens, float* output, InputT const* mat_a,
44 float const* mat_b, cudaStream_t stream) {
45 if (num_tokens == kBegin) {
46 invokeFp32RouterGemm<InputT, kBegin, kNumExperts, kHiddenDim>(
47 output, mat_a, mat_b, stream);
48 } else {
49 Fp32LoopUnroller<InputT, kNumExperts, kHiddenDim, kBegin + 1,
50 kEnd>::unroll(num_tokens, output, mat_a, mat_b, stream);
51 }
52 }
53 };
54
55 template <typename InputT, int kNumExperts, int kHiddenDim, int kEnd>
56 struct Fp32LoopUnroller<InputT, kNumExperts, kHiddenDim, kEnd, kEnd> {
57 static void unroll(int num_tokens, float* output, InputT const* mat_a,
58 float const* mat_b, cudaStream_t stream) {
59 if (num_tokens == kEnd) {
60 invokeFp32RouterGemm<InputT, kEnd, kNumExperts, kHiddenDim>(
61 output, mat_a, mat_b, stream);
62 } else {
63 throw std::invalid_argument(
64 "fp32_router_gemm: num_tokens must be in [1, 32]");
65 }
66 }
67 };
68
69 // Dispatch over the supported (num_experts, hidden_dim) pairs.
70 template <typename InputT>
71 void dispatchFp32RouterGemm(int num_experts, int hidden_dim, int num_tokens,
72 float* output, InputT const* mat_a,
73 float const* mat_b, cudaStream_t stream) {
74 if (num_experts == 256 && hidden_dim == 3072) {
75 Fp32LoopUnroller<InputT, 256, 3072, 1, FP32_MAX_TOKENS>::unroll(
76 num_tokens, output, mat_a, mat_b, stream);
77 } else if (num_experts == 128 && hidden_dim == 6144) {
78 Fp32LoopUnroller<InputT, 128, 6144, 1, FP32_MAX_TOKENS>::unroll(
79 num_tokens, output, mat_a, mat_b, stream);
80 } else {
81 throw std::invalid_argument(
82 "fp32_router_gemm: unsupported (hidden_dim, num_experts) pair");
83 }
84 }
85
86 void fp32_router_gemm(
87 torch::stable::Tensor& output, // [num_tokens, num_experts]
88 torch::stable::Tensor const& mat_a, // [num_tokens, hidden_dim]
89 torch::stable::Tensor const& mat_b // [num_experts, hidden_dim]
90 ) {
91 STD_TORCH_CHECK(output.dim() == 2 && mat_a.dim() == 2 && mat_b.dim() == 2);
92 STD_TORCH_CHECK(output.is_cuda() && mat_a.is_cuda() && mat_b.is_cuda(),
93 "fp32_router_gemm: all tensors must be CUDA tensors");
94 STD_TORCH_CHECK(output.get_device_index() == mat_a.get_device_index() &&
95 output.get_device_index() == mat_b.get_device_index(),
96 "fp32_router_gemm: all tensors must be on the same device");
97 STD_TORCH_CHECK(
98 output.is_contiguous() && mat_a.is_contiguous() && mat_b.is_contiguous(),
99 "fp32_router_gemm: all tensors must be contiguous");
100
101 const int num_tokens = mat_a.size(0);
102 const int num_experts = mat_b.size(0);
103 const int hidden_dim = mat_a.size(1);
104
105 STD_TORCH_CHECK(output.size(0) == num_tokens && output.size(1) == num_experts,
106 "fp32_router_gemm: output must have shape [num_tokens, "
107 "num_experts]");
108 STD_TORCH_CHECK(
109 mat_a.size(1) == mat_b.size(1),
110 "fp32_router_gemm: mat_a and mat_b must have the same hidden_dim");
111 STD_TORCH_CHECK(
112 fp32_router_gemm_supported(hidden_dim, num_experts),
113 "fp32_router_gemm: supported (hidden_dim, num_experts) pairs are "
114 "(3072, 256) and (6144, 128)");
115 STD_TORCH_CHECK(num_tokens <= FP32_MAX_TOKENS,
116 "fp32_router_gemm: num_tokens must be in [0, 32]");
117 STD_TORCH_CHECK(
118 mat_a.scalar_type() == torch::headeronly::ScalarType::Float ||
119 mat_a.scalar_type() == torch::headeronly::ScalarType::BFloat16,
120 "fp32_router_gemm: mat_a must be float32 or bfloat16");
121 STD_TORCH_CHECK(mat_b.scalar_type() == torch::headeronly::ScalarType::Float,
122 "fp32_router_gemm: mat_b (weight) must be float32");
123 STD_TORCH_CHECK(output.scalar_type() == torch::headeronly::ScalarType::Float,
124 "fp32_router_gemm: output must be float32");
125
126 if (num_tokens == 0) {
127 return;
128 }
129
130 const torch::stable::accelerator::DeviceGuard device_guard(
131 mat_a.get_device_index());
132 STD_TORCH_CHECK(getSMVersion() >= 90, "fp32_router_gemm: requires SM90+");
133
134 auto stream = get_current_cuda_stream(mat_a.get_device_index());
135 float* out_ptr = reinterpret_cast<float*>(output.mutable_data_ptr());
136 float const* mat_b_ptr = reinterpret_cast<float const*>(mat_b.data_ptr());
137
138 if (mat_a.scalar_type() == torch::headeronly::ScalarType::BFloat16) {
139 auto const* mat_a_ptr =
140 reinterpret_cast<__nv_bfloat16 const*>(mat_a.data_ptr());
141 dispatchFp32RouterGemm<__nv_bfloat16>(num_experts, hidden_dim, num_tokens,
142 out_ptr, mat_a_ptr, mat_b_ptr,
143 stream);
144 } else {
145 auto const* mat_a_ptr = reinterpret_cast<float const*>(mat_a.data_ptr());
146 dispatchFp32RouterGemm<float>(num_experts, hidden_dim, num_tokens, out_ptr,
147 mat_a_ptr, mat_b_ptr, stream);
148 }
149 }
150
151 STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, m) {
152 m.impl("fp32_router_gemm", TORCH_BOX(&fp32_router_gemm));
153 }