From 38c233fd041f134f96ae23065bb763456881269f Mon Sep 17 00:00:00 2001 From: Binyao Jiang Date: Tue, 20 Jan 2026 14:05:32 -0800 Subject: [PATCH] [Piecewise] Support PCG weak_ref_tensor cuda kernel on AMD (#17291) --- python/sglang/srt/compilation/weak_ref_tensor.py | 4 ++-- sgl-kernel/csrc/common_extension_rocm.cc | 6 ++++++ sgl-kernel/setup_rocm.py | 1 + 3 files changed, 9 insertions(+), 2 deletions(-) diff --git a/python/sglang/srt/compilation/weak_ref_tensor.py b/python/sglang/srt/compilation/weak_ref_tensor.py index 83cb7e64b..30f59d46f 100644 --- a/python/sglang/srt/compilation/weak_ref_tensor.py +++ b/python/sglang/srt/compilation/weak_ref_tensor.py @@ -2,9 +2,9 @@ from typing import Any, Union import torch -from sglang.srt.utils.common import is_cuda, is_npu +from sglang.srt.utils.common import is_cuda, is_hip, is_npu -if is_cuda(): +if is_cuda() or is_hip(): from sgl_kernel import weak_ref_tensor elif is_npu(): from torch_npu._C import _weak_ref_tensor as weak_ref_tensor diff --git a/sgl-kernel/csrc/common_extension_rocm.cc b/sgl-kernel/csrc/common_extension_rocm.cc index d0147e5f3..add827787 100644 --- a/sgl-kernel/csrc/common_extension_rocm.cc +++ b/sgl-kernel/csrc/common_extension_rocm.cc @@ -231,6 +231,12 @@ TORCH_LIBRARY_EXPAND(sgl_kernel, m) { "float scale," "int max_period) -> Tensor"); m.impl("timestep_embedding", torch::kCUDA, ×tep_embedding); + + /* + * From csrc/memory + */ + m.def("weak_ref_tensor(Tensor tensor) -> Tensor"); + m.impl("weak_ref_tensor", torch::kCUDA, &weak_ref_tensor); } REGISTER_EXTENSION(common_ops) diff --git a/sgl-kernel/setup_rocm.py b/sgl-kernel/setup_rocm.py index a7336a103..16fdbbd2f 100644 --- a/sgl-kernel/setup_rocm.py +++ b/sgl-kernel/setup_rocm.py @@ -53,6 +53,7 @@ sources = [ "csrc/moe/moe_topk_sigmoid_kernels.cu", "csrc/speculative/eagle_utils.cu", "csrc/kvcacheio/transfer.cu", + "csrc/memory/weak_ref_tensor.cpp", "csrc/elementwise/pos_enc.cu", "csrc/sgl_diffusion/elementwise/timestep_embedding.cu", ]