chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

Enable JIT clamp_position and resolve_future_token_ids on ROCm (#21116)

Browse Source
This commit is contained in:
Lianmin Zheng
2026-03-22 22:33:54 -07:00
committed by GitHub
parent 32a85ef128
commit 7050011dee
4 changed files with 5 additions and 16 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
python/sglang/jit_kernel/csrc/elementwise/clamp_position.cuh
View File

@@ -29,7 +29,7 @@ struct ClampPosition {
SymbolicSize N = {"num_elements"};
SymbolicDevice device_;
device_.set_options<kDLCUDA>();
device_.set_options<kDLCUDA, kDLROCM>();
TensorMatcher({N}) //
.with_dtype<T>()

2
python/sglang/jit_kernel/csrc/elementwise/resolve_future_token_ids.cuh
View File

@@ -34,7 +34,7 @@ struct ResolveFutureTokenIds {
SymbolicSize N = {"num_tokens"};
SymbolicSize M = {"map_size"};
SymbolicDevice device_;
device_.set_options<kDLCUDA>();
device_.set_options<kDLCUDA, kDLROCM>();
TensorMatcher({N}).with_dtype<T>().with_device(device_).verify(input_ids);

10
python/sglang/srt/managers/overlap_utils.py
View File

@@ -6,7 +6,7 @@ from typing import TYPE_CHECKING, Optional
import torch
from sglang.srt.speculative.spec_utils import spec_need_hidden_states
from sglang.srt.utils import get_compiler_backend, is_cuda, is_hip
from sglang.srt.utils import is_cuda, is_hip
if TYPE_CHECKING:
from sglang.srt.managers.schedule_batch import ModelWorkerBatch
@@ -26,18 +26,12 @@ def _resolve_future_token_ids_native(input_ids, future_token_ids_map):
)
if _is_cuda:
if _is_cuda or _is_hip:
from sglang.jit_kernel.resolve_future_token_ids import (
resolve_future_token_ids_cuda,
)
_resolve_future_token_ids = resolve_future_token_ids_cuda
elif _is_hip:
_resolve_future_token_ids = torch.compile(
_resolve_future_token_ids_native,
dynamic=True,
backend=get_compiler_backend(),
)
else:
_resolve_future_token_ids = _resolve_future_token_ids_native

7
python/sglang/srt/model_executor/forward_batch_info.py
View File

@@ -58,7 +58,6 @@ from sglang.srt.model_executor.forward_batch_deepseek_mha_mixin import (
)
from sglang.srt.server_args import get_global_server_args
from sglang.srt.utils import (
get_compiler_backend,
is_cuda,
is_hip,
is_npu,
@@ -1187,13 +1186,9 @@ def _clamp_position_native(seq_lens):
return torch.clamp((seq_lens - 1), min=0).to(torch.int64)
if is_cuda():
if is_cuda() or is_hip():
from sglang.jit_kernel.clamp_position import clamp_position_cuda
clamp_position = clamp_position_cuda
elif is_hip():
clamp_position = torch.compile(
_clamp_position_native, dynamic=True, backend=get_compiler_backend()
)
else:
clamp_position = _clamp_position_native