minke.yu/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
python/sglang/jit_kernel/csrc/elementwise/clamp_position.cuh
+1 -1
View File
@@ -29,7 +29,7 @@ struct ClampPosition {
SymbolicSize N = {"num_elements"}; SymbolicSize N = {"num_elements"};
SymbolicDevice device_; SymbolicDevice device_;
device_.set_options<kDLCUDA>(); device_.set_options<kDLCUDA, kDLROCM>();
TensorMatcher({N}) // TensorMatcher({N}) //
.with_dtype<T>() .with_dtype<T>()
python/sglang/jit_kernel/csrc/elementwise/resolve_future_token_ids.cuh
+1 -1
View File
@@ -34,7 +34,7 @@ struct ResolveFutureTokenIds {
SymbolicSize N = {"num_tokens"}; SymbolicSize N = {"num_tokens"};
SymbolicSize M = {"map_size"}; SymbolicSize M = {"map_size"};
SymbolicDevice device_; SymbolicDevice device_;
device_.set_options<kDLCUDA>(); device_.set_options<kDLCUDA, kDLROCM>();
TensorMatcher({N}).with_dtype<T>().with_device(device_).verify(input_ids); TensorMatcher({N}).with_dtype<T>().with_device(device_).verify(input_ids);
python/sglang/srt/managers/overlap_utils.py
+2 -8
View File
@@ -6,7 +6,7 @@ from typing import TYPE_CHECKING, Optional
import torch import torch
from sglang.srt.speculative.spec_utils import spec_need_hidden_states 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: if TYPE_CHECKING:
from sglang.srt.managers.schedule_batch import ModelWorkerBatch 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 ( from sglang.jit_kernel.resolve_future_token_ids import (
resolve_future_token_ids_cuda, resolve_future_token_ids_cuda,
) )
_resolve_future_token_ids = 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: else:
_resolve_future_token_ids = _resolve_future_token_ids_native _resolve_future_token_ids = _resolve_future_token_ids_native
python/sglang/srt/model_executor/forward_batch_info.py
+1 -6
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.server_args import get_global_server_args
from sglang.srt.utils import ( from sglang.srt.utils import (
get_compiler_backend,
is_cuda, is_cuda,
is_hip, is_hip,
is_npu, is_npu,
@@ -1187,13 +1186,9 @@ def _clamp_position_native(seq_lens):
return torch.clamp((seq_lens - 1), min=0).to(torch.int64) 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 from sglang.jit_kernel.clamp_position import clamp_position_cuda
clamp_position = 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: else:
clamp_position = _clamp_position_native clamp_position = _clamp_position_native