A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 13, 07:00 UTC
Rewrite C++ meta funcs to Python (#28595)
Signed-off-by: Jane Xu <janeyx@meta.com>
Jane (Yuan) Xu committed
Nov 13, 2025 at 11:52 UTC
06c4873d959feb0d4cb062ef17cdd0dd09dbf10f
3 files changed
+38
-33
csrc/quantization/gptq_marlin/awq_marlin_repack.cu
-16
@@ -247,22 +247,6 @@ torch::Tensor awq_marlin_repack(torch::Tensor& b_q_weight, int64_t size_k,
247
return out;
248
}
249
250
-torch::Tensor awq_marlin_repack_meta(torch::Tensor& b_q_weight,
251
- c10::SymInt size_k, c10::SymInt size_n,
252
- int64_t num_bits) {
253
- int const pack_factor = 32 / num_bits;
254
- auto options = torch::TensorOptions()
255
- .dtype(b_q_weight.dtype())
256
- .device(b_q_weight.device());
257
- return torch::empty_symint(
258
- {size_k / marlin::tile_size, size_n * marlin::tile_size / pack_factor},
259
- options);
260
-}
261
-
250
TORCH_LIBRARY_IMPL_EXPAND(TORCH_EXTENSION_NAME, CUDA, m) {
251
m.impl("awq_marlin_repack", &awq_marlin_repack);
252
}
265
-
266
-TORCH_LIBRARY_IMPL_EXPAND(TORCH_EXTENSION_NAME, Meta, m) {
267
- m.impl("awq_marlin_repack", &awq_marlin_repack_meta);
268
-}
csrc/quantization/gptq_marlin/gptq_marlin_repack.cu
-16
@@ -321,22 +321,6 @@ torch::Tensor gptq_marlin_repack(torch::Tensor& b_q_weight, torch::Tensor& perm,
321
return out;
322
}
323
324
-torch::Tensor gptq_marlin_repack_meta(torch::Tensor& b_q_weight,
325
- torch::Tensor& perm, c10::SymInt size_k,
326
- c10::SymInt size_n, int64_t num_bits) {
327
- int const pack_factor = 32 / num_bits;
328
- auto options = torch::TensorOptions()
329
- .dtype(b_q_weight.dtype())
330
- .device(b_q_weight.device());
331
- return torch::empty_symint(
332
- {size_k / marlin::tile_size, size_n * marlin::tile_size / pack_factor},
333
- options);
334
-}
335
-
324
TORCH_LIBRARY_IMPL_EXPAND(TORCH_EXTENSION_NAME, CUDA, m) {
325
m.impl("gptq_marlin_repack", &gptq_marlin_repack);
326
}
339
-
340
-TORCH_LIBRARY_IMPL_EXPAND(TORCH_EXTENSION_NAME, Meta, m) {
341
- m.impl("gptq_marlin_repack", &gptq_marlin_repack_meta);
342
-}
vllm/_custom_ops.py
+38
-1
@@ -1174,13 +1174,50 @@ def gptq_marlin_repack(
1174
return torch.ops._C.gptq_marlin_repack(b_q_weight, perm, size_k, size_n, num_bits)
1175
1176
1177
-# gptq_marlin
1177
+if hasattr(torch.ops._C, "gptq_marlin_repack"):
1178
+
1179
+ @register_fake("_C::gptq_marlin_repack")
1180
+ def _gptq_marlin_repack_fake(
1181
+ b_q_weight: torch.Tensor,
1182
+ perm: torch.Tensor,
1183
+ size_k: torch.SymInt,
1184
+ size_n: torch.SymInt,
1185
+ num_bits: int,
1186
+ ) -> torch.Tensor:
1187
+ pack_factor = 32 // num_bits
1188
+ marlin_tile_size = 16
1189
+ return torch.empty(
1190
+ (size_k // marlin_tile_size, size_n * marlin_tile_size // pack_factor),
1191
+ dtype=b_q_weight.dtype,
1192
+ device=b_q_weight.device,
1193
+ )
1194
+
1195
+
1196
+# awq_marlin
1197
def awq_marlin_repack(
1198
b_q_weight: torch.Tensor, size_k: int, size_n: int, num_bits: int
1199
) -> torch.Tensor:
1200
return torch.ops._C.awq_marlin_repack(b_q_weight, size_k, size_n, num_bits)
1201
1202
1203
+if hasattr(torch.ops._C, "awq_marlin_repack"):
1204
+
1205
+ @register_fake("_C::awq_marlin_repack")
1206
+ def _awq_marlin_repack_fake(
1207
+ b_q_weight: torch.Tensor,
1208
+ size_k: torch.SymInt,
1209
+ size_n: torch.SymInt,
1210
+ num_bits: int,
1211
+ ) -> torch.Tensor:
1212
+ pack_factor = 32 // num_bits
1213
+ marlin_tile_size = 16
1214
+ return torch.empty(
1215
+ (size_k // marlin_tile_size, size_n * marlin_tile_size // pack_factor),
1216
+ dtype=b_q_weight.dtype,
1217
+ device=b_q_weight.device,
1218
+ )
1219
+
1220
+
1221
def gptq_marlin_moe_repack(
1222
b_q_weight: torch.Tensor,
1223
perm: torch.Tensor,