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 9, 14:45 UTC
[fix] CPUDNNLGEMMHandler pointer baked into inductor artifact (#32913)
Signed-off-by: dolpm <34420038+dolpm@users.noreply.github.com>
dolpm committed
Jan 26, 2026 at 13:59 UTC
58a05b0ca184048e11ea1a54759732ad634ce34a
4 files changed
+38
-23
csrc/cpu/dnnl_kernels.cpp
+6
-4
@@ -360,13 +360,14 @@ void onednn_scaled_mm(
360
const std::optional<torch::Tensor>& azp, // [M] or [1]
361
const std::optional<torch::Tensor>& azp_adj, // [M] or [1]
362
const std::optional<torch::Tensor>& bias, // [N]
363
- int64_t handler) {
363
+ const torch::Tensor& handler_tensor) {
364
CPU_KERNEL_GUARD_IN(onednn_scaled_mm)
365
TORCH_CHECK(a.dim() == 2);
366
TORCH_CHECK(a.is_contiguous());
367
TORCH_CHECK(c.is_contiguous());
368
W8A8MatMulPrimitiveHandler* ptr =
369
- reinterpret_cast<W8A8MatMulPrimitiveHandler*>(handler);
369
+ reinterpret_cast<W8A8MatMulPrimitiveHandler*>(
370
+ handler_tensor.item<int64_t>());
371
const int32_t* azp_ptr = nullptr;
372
if (azp.has_value()) {
373
azp_ptr = azp->data_ptr<int32_t>();
@@ -519,13 +520,14 @@ int64_t create_onednn_mm_handler(const torch::Tensor& b,
520
521
void onednn_mm(torch::Tensor& c, // [M, OC], row-major
522
const torch::Tensor& a, // [M, IC], row-major
522
- const std::optional<torch::Tensor>& bias, int64_t handler) {
523
+ const std::optional<torch::Tensor>& bias,
524
+ const torch::Tensor& handler_tensor) {
525
CPU_KERNEL_GUARD_IN(onednn_mm)
526
TORCH_CHECK(a.dim() == 2);
527
TORCH_CHECK(a.stride(-1) == 1);
528
TORCH_CHECK(c.stride(-1) == 1);
529
MatMulPrimitiveHandler* ptr =
528
- reinterpret_cast<MatMulPrimitiveHandler*>(handler);
530
+ reinterpret_cast<MatMulPrimitiveHandler*>(handler_tensor.item<int64_t>());
531
532
// ACL matmuls expect contiguous source tensors
533
#ifdef VLLM_USE_ACL
csrc/cpu/torch_bindings.cpp
+5
-4
@@ -19,13 +19,14 @@ void onednn_scaled_mm(torch::Tensor& c, const torch::Tensor& a,
19
const std::optional<torch::Tensor>& azp,
20
const std::optional<torch::Tensor>& azp_adj,
21
const std::optional<torch::Tensor>& bias,
22
- int64_t handler);
22
+ const torch::Tensor& handler_tensor);
23
24
int64_t create_onednn_mm_handler(const torch::Tensor& b,
25
int64_t primitive_cache_size);
26
27
void onednn_mm(torch::Tensor& c, const torch::Tensor& a,
28
- const std::optional<torch::Tensor>& bias, int64_t handler);
28
+ const std::optional<torch::Tensor>& bias,
29
+ const torch::Tensor& handler_tensor);
30
31
bool is_onednn_acl_supported();
32
@@ -196,7 +197,7 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
197
// oneDNN GEMM
198
ops.def(
199
"onednn_mm(Tensor! c, Tensor a, Tensor? bias, "
199
- "int handler) -> ()");
200
+ "Tensor handler_tensor) -> ()");
201
ops.impl("onednn_mm", torch::kCPU, &onednn_mm);
202
203
// Check if oneDNN was built with ACL backend
@@ -212,7 +213,7 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
213
// oneDNN scaled_mm for W8A8 with static per-tensor activation quantization
214
ops.def(
215
"onednn_scaled_mm(Tensor! c, Tensor a, Tensor a_scales, Tensor? azp, "
215
- "Tensor? azp_adj, Tensor? bias, int handler) -> ()");
216
+ "Tensor? azp_adj, Tensor? bias, Tensor handler_tensor) -> ()");
217
ops.impl("onednn_scaled_mm", torch::kCPU, &onednn_scaled_mm);
218
219
// Compute int8 quantized tensor for given scaling factor.
vllm/_custom_ops.py
+26
-9
@@ -2845,13 +2845,13 @@ if hasattr(torch.ops._C, "int8_scaled_mm_with_quant"):
2845
2846
class CPUDNNLGEMMHandler:
2847
def __init__(self) -> None:
2848
- self.handler: int | None = None
2848
+ self.handler_tensor: torch.Tensor | None = None
2849
self.n = -1
2850
self.k = -1
2851
2852
def __del__(self):
2853
- if self.handler is not None:
2854
- torch.ops._C.release_dnnl_matmul_handler(self.handler)
2853
+ if self.handler_tensor is not None:
2854
+ torch.ops._C.release_dnnl_matmul_handler(self.handler_tensor.item())
2855
2856
2857
_supports_onednn = bool(hasattr(torch.ops._C, "create_onednn_mm_handler"))
@@ -2867,8 +2867,10 @@ def create_onednn_mm(
2867
) -> CPUDNNLGEMMHandler:
2868
handler = CPUDNNLGEMMHandler()
2869
handler.k, handler.n = weight.size()
2870
- handler.handler = torch.ops._C.create_onednn_mm_handler(
2871
- weight, primitive_cache_size
2870
+ # store the handler pointer in a tensor it doesn't get inlined
2871
+ handler.handler_tensor = torch.tensor(
2872
+ torch.ops._C.create_onednn_mm_handler(weight, primitive_cache_size),
2873
+ dtype=torch.int64,
2874
)
2875
return handler
2876
@@ -2880,7 +2882,7 @@ def onednn_mm(
2882
) -> torch.Tensor:
2883
output = torch.empty((*x.shape[0:-1], dnnl_handler.n), dtype=x.dtype)
2884
torch.ops._C.onednn_mm(
2883
- output, x.reshape(-1, dnnl_handler.k), bias, dnnl_handler.handler
2885
+ output, x.reshape(-1, dnnl_handler.k), bias, dnnl_handler.handler_tensor
2886
)
2887
2888
return output
@@ -2896,8 +2898,17 @@ def create_onednn_scaled_mm(
2898
) -> CPUDNNLGEMMHandler:
2899
handler = CPUDNNLGEMMHandler()
2900
handler.k, handler.n = weight.size()
2899
- handler.handler = torch.ops._C.create_onednn_scaled_mm_handler(
2900
- weight, weight_scales, output_type, dynamic_quant, use_azp, primitive_cache_size
2901
+ # store the handler pointer in a tensor so it doesn't get inlined
2902
+ handler.handler_tensor = torch.tensor(
2903
+ torch.ops._C.create_onednn_scaled_mm_handler(
2904
+ weight,
2905
+ weight_scales,
2906
+ output_type,
2907
+ dynamic_quant,
2908
+ use_azp,
2909
+ primitive_cache_size,
2910
+ ),
2911
+ dtype=torch.int64,
2912
)
2913
return handler
2914
@@ -2950,7 +2961,13 @@ def onednn_scaled_mm(
2961
bias: torch.Tensor | None,
2962
) -> torch.Tensor:
2963
torch.ops._C.onednn_scaled_mm(
2953
- output, x, input_scale, input_zp, input_zp_adj, bias, dnnl_handler.handler
2964
+ output,
2965
+ x,
2966
+ input_scale,
2967
+ input_zp,
2968
+ input_zp_adj,
2969
+ bias,
2970
+ dnnl_handler.handler_tensor,
2971
)
2972
2973
return output
vllm/envs.py
+1
-6
@@ -289,16 +289,11 @@ def use_aot_compile() -> bool:
289
from vllm.model_executor.layers.batch_invariant import (
290
vllm_is_batch_invariant,
291
)
292
- from vllm.platforms import current_platform
292
from vllm.utils.torch_utils import is_torch_equal_or_newer
293
294
default_value = (
295
"1"
297
- if is_torch_equal_or_newer("2.10.0.dev")
298
- and not disable_compile_cache()
299
- # Disabling AOT_COMPILE for CPU
300
- # See: https://github.com/vllm-project/vllm/issues/32033
301
- and not current_platform.is_cpu()
296
+ if is_torch_equal_or_newer("2.10.0.dev") and not disable_compile_cache()
297
else "0"
298
)
299