@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
cpp 172 lines 6.3 KB
Raw
1 #include <ATen/ATen.h>
2 #include <ATen/Parallel.h>
3 #include <torch/all.h>
4
5 // _dyn_quant_matmul_4bit is only available on AArch64.
6 #if defined(__aarch64__)
7 #include <ATen/ops/_dyn_quant_matmul_4bit.h>
8 #endif
9
10 inline torch::Tensor mm(const torch::Tensor& a, const torch::Tensor& packed_w,
11 int64_t group_size_eff, int64_t in_features,
12 int64_t out_features) {
13 #if defined(__aarch64__)
14 return at::_ops::_dyn_quant_matmul_4bit::call(a, packed_w, group_size_eff,
15 in_features, out_features);
16 #else
17 TORCH_CHECK(false,
18 "dynamic 4-bit int MoE path requires AArch64 (ARM64); "
19 "_dyn_quant_matmul_4bit is unavailable on this architecture");
20 return {};
21 #endif
22 }
23
24 enum ActivationKind : int64_t {
25 SwiGLU_Gu = 0, // act = SiLU(g) * u
26 SwiGLUOAI = 1, // act = SiLU(u) * g
27 SiLU = 2 // SiLU
28 };
29
30 torch::Tensor dynamic_4bit_int_moe_cpu(
31 torch::Tensor x, torch::Tensor topk_ids, torch::Tensor topk_weights,
32 torch::Tensor w13_packed, torch::Tensor w2_packed, int64_t hidden_size,
33 int64_t intermediate_size, int64_t group_size,
34 bool apply_router_weight_on_input, int64_t activation_kind) {
35 TORCH_CHECK(x.dim() == 2, "x must be 2D");
36 TORCH_CHECK(topk_ids.dim() == 2 && topk_weights.dim() == 2,
37 "topk tensors must be [T, K]");
38 TORCH_CHECK(
39 w13_packed.size(0) == w2_packed.size(0),
40 "w13_packed and w2_packed must have same number of experts in dim 0");
41
42 const int64_t T = x.size(0);
43 const int64_t K = topk_ids.size(1);
44 const int64_t E = w13_packed.size(0);
45 const int64_t N = T * K;
46 const int64_t w13_out_features = 2 * intermediate_size;
47
48 auto x_c = x.contiguous();
49 // _dyn_quant_matmul_4bit kernel natively supports these pre-quant activation
50 // dtypes:
51 // - fp32: with channelwise and groupwise
52 // - bf16: with channelwise -> upcast to fp32 for groupwise
53 // - fp16: not supported -> upcast to fp32 for groupwise & channelwise
54 const auto output_dtype = x_c.scalar_type();
55 const bool should_cast_input =
56 ((group_size != -1) && output_dtype == at::kBFloat16) ||
57 output_dtype == at::kHalf;
58 if (should_cast_input) {
59 x_c = x_c.to(at::kFloat);
60 }
61 auto ids_c = topk_ids.contiguous();
62 auto gates_c = topk_weights.to(x_c.scalar_type()).contiguous();
63
64 // bucketing tokens -> experts
65 c10::SmallVector<int64_t, 64> counts(
66 E, 0); // Small vector uses stack allocation
67 {
68 const auto* ids_ptr = ids_c.data_ptr<int64_t>();
69 for (int64_t i = 0; i < N; ++i) {
70 const int64_t e_id = ids_ptr[i];
71 TORCH_CHECK(0 <= e_id && e_id < E, "expert id out of range");
72 counts[e_id]++;
73 }
74 }
75 c10::SmallVector<int64_t, 65> offsets(E + 1, 0); // ( E +1 )
76 for (int64_t e = 0; e < E; ++e) offsets[e + 1] = offsets[e] + counts[e];
77
78 // expert_tokens = [tokens indices for expert 0, ...]
79 // expert_gates = [router weights for tokens assigned to expert 0, ...]
80 auto expert_tokens = at::empty({offsets[E]}, ids_c.options());
81 auto expert_gates = at::empty({offsets[E]}, gates_c.options());
82 {
83 c10::SmallVector<int64_t, 64> cursor(E, 0);
84 AT_DISPATCH_FLOATING_TYPES_AND2(
85 at::ScalarType::BFloat16, at::ScalarType::Half, gates_c.scalar_type(),
86 "bucket_expert_tokens_and_gates", [&] {
87 const auto* ids_ptr = ids_c.data_ptr<int64_t>();
88 const auto* gts_ptr = gates_c.data_ptr<scalar_t>();
89 auto* tok_ptr = expert_tokens.data_ptr<int64_t>();
90 auto* gate_ptr = expert_gates.data_ptr<scalar_t>();
91
92 for (int64_t t = 0; t < T; ++t) {
93 const int64_t base = t * K;
94 for (int64_t k = 0; k < K; ++k) {
95 const int64_t idx = base + k;
96 const int64_t e = ids_ptr[idx];
97 const int64_t p = offsets[e] + (cursor[e]++);
98 tok_ptr[p] = t;
99 gate_ptr[p] = gts_ptr[idx];
100 }
101 }
102 });
103 }
104
105 const int64_t g_eff_13 = (group_size != -1) ? group_size : hidden_size;
106 const int64_t g_eff_2 = (group_size != -1) ? group_size : intermediate_size;
107
108 // X_all [num_tokens * K, hidden_size]
109 auto X_all = x_c.index_select(/*dim=*/0, expert_tokens);
110 if (apply_router_weight_on_input) {
111 X_all = X_all.mul(expert_gates.unsqueeze(1));
112 }
113 auto Y_all = at::empty({offsets[E], hidden_size}, x_c.options());
114
115 at::parallel_for(0, offsets[E], 0, [&](int64_t idx_begin, int64_t idx_end) {
116 c10::InferenceMode guard;
117 for (int64_t e = 0; e < E; ++e) {
118 int64_t start = std::max(offsets[e], idx_begin);
119 int64_t end = std::min(offsets[e + 1], idx_end);
120 int64_t te = end - start;
121 if (te <= 0) {
122 continue;
123 }
124
125 auto x_e = X_all.narrow(/*dim=*/0, /*start=*/start, /*length=*/te);
126
127 auto w13_e = w13_packed.select(/*dim=*/0, e);
128 auto w2_e = w2_packed.select(/*dim=*/0, e);
129
130 // W13
131 auto y13 = mm(x_e, w13_e, g_eff_13, /*in_features=*/hidden_size,
132 /*out_features=*/w13_out_features);
133
134 auto g_part =
135 y13.narrow(/*dim=*/1, /*start=*/0, /*length=*/intermediate_size);
136 auto u_part = y13.narrow(/*dim=*/1, /*start=*/intermediate_size,
137 /*length=*/intermediate_size);
138
139 torch::Tensor act;
140 if (activation_kind == ActivationKind::SwiGLUOAI) { // SwiGLUOAI
141 constexpr double kAlpha = 1.702; // GPT-OSS default
142 constexpr double kLimit = 7.0; // GPT-OSS default
143 auto gate_c = at::clamp_max(g_part, kLimit);
144 auto up_c = at::clamp(u_part, -kLimit, kLimit);
145 auto glu = gate_c.mul(at::sigmoid(gate_c.mul(kAlpha)));
146 act = up_c.add(1.0).mul(glu);
147 } else { // SiLU , SwiGLU_GU, vLLM maps silu to SiluAndMul()
148 act = at::silu(g_part).mul(u_part);
149 }
150
151 // W2
152 auto y = mm(act, w2_e, g_eff_2, /*in_features=*/intermediate_size,
153 /*out_features=*/hidden_size);
154
155 // Store per-expert result
156 Y_all.narrow(/*dim=*/0, /*start=*/start, /*length=*/te).copy_(y);
157 }
158 });
159
160 if (!apply_router_weight_on_input) {
161 Y_all = Y_all.mul(expert_gates.unsqueeze(1));
162 }
163 if (Y_all.scalar_type() != output_dtype) {
164 Y_all = Y_all.to(output_dtype);
165 }
166
167 auto out = at::zeros({T, hidden_size}, x.options());
168 out =
169 at::index_add(out, /*dim=*/0, /*index=*/expert_tokens, /*source=*/Y_all);
170
171 return out;
172 }