diff --git a/python/sglang/jit_kernel/benchmark/bench_resolve_future_token_ids.py b/python/sglang/jit_kernel/benchmark/bench_resolve_future_token_ids.py new file mode 100644 index 000000000..31430442c --- /dev/null +++ b/python/sglang/jit_kernel/benchmark/bench_resolve_future_token_ids.py @@ -0,0 +1,69 @@ +import itertools + +import torch +import triton +import triton.testing + +from sglang.jit_kernel.benchmark.utils import ( + DEFAULT_DEVICE, + get_benchmark_range, + run_benchmark, +) +from sglang.jit_kernel.resolve_future_token_ids import resolve_future_token_ids_cuda +from sglang.srt.utils import get_compiler_backend + +SIZE_LIST = get_benchmark_range( + full_range=[2**n for n in range(4, 16)], # 16 … 32K elements + ci_range=[256, 4096], +) + +configs = list(itertools.product(SIZE_LIST)) + + +def _torch_resolve(input_ids, future_map): + input_ids[:] = torch.where( + input_ids < 0, + future_map[torch.clamp(-input_ids, min=0)], + input_ids, + ) + + +_compiled_resolve = torch.compile( + _torch_resolve, dynamic=True, backend=get_compiler_backend() +) + + +@triton.testing.perf_report( + triton.testing.Benchmark( + x_names=["size"], + x_vals=configs, + line_arg="provider", + line_vals=["jit", "torch_compile", "torch"], + line_names=["SGL JIT Kernel", "torch.compile", "PyTorch"], + styles=[("blue", "-"), ("green", "-."), ("red", "--")], + ylabel="us", + plot_name="resolve-future-token-ids-performance", + args={}, + ) +) +def benchmark(size: int, provider: str): + map_size = 8192 + future_map = torch.randint( + 0, 50000, (map_size,), dtype=torch.int64, device=DEFAULT_DEVICE + ) + input_ids = torch.randint( + -map_size + 1, 50000, (size,), dtype=torch.int64, device=DEFAULT_DEVICE + ) + + if provider == "jit": + fn = lambda: resolve_future_token_ids_cuda(input_ids.clone(), future_map) + elif provider == "torch_compile": + fn = lambda: _compiled_resolve(input_ids.clone(), future_map) + else: + fn = lambda: _torch_resolve(input_ids.clone(), future_map) + + return run_benchmark(fn) + + +if __name__ == "__main__": + benchmark.run(print_data=True) diff --git a/python/sglang/jit_kernel/csrc/elementwise/resolve_future_token_ids.cuh b/python/sglang/jit_kernel/csrc/elementwise/resolve_future_token_ids.cuh new file mode 100644 index 000000000..0ec5d60d7 --- /dev/null +++ b/python/sglang/jit_kernel/csrc/elementwise/resolve_future_token_ids.cuh @@ -0,0 +1,57 @@ +#include // For TensorMatcher, SymbolicSize, SymbolicDevice +#include // For RuntimeCheck, div_ceil + +#include // For LaunchKernel + +#include +#include + +#include +#include + +namespace { + +template +__global__ void resolve_future_token_ids_kernel(T* __restrict__ input_ids, const T* __restrict__ future_map, size_t n) { + size_t idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx < n) { + T val = input_ids[idx]; + if (val < 0) { + T key = -val; + if (key < 0) key = 0; // clamp for overflow + input_ids[idx] = future_map[key]; + } + } +} + +constexpr size_t kBlockSize = 256; + +template +struct ResolveFutureTokenIds { + static void run(tvm::ffi::TensorView input_ids, tvm::ffi::TensorView future_map) { + using namespace host; + + SymbolicSize N = {"num_tokens"}; + SymbolicSize M = {"map_size"}; + SymbolicDevice device_; + device_.set_options(); + + TensorMatcher({N}).with_dtype().with_device(device_).verify(input_ids); + + TensorMatcher({M}).with_dtype().with_device(device_).verify(future_map); + + const size_t num_tokens = N.unwrap(); + if (num_tokens == 0) return; + + const size_t grid_size = div_ceil(num_tokens, kBlockSize); + const DLDevice device = device_.unwrap(); + + LaunchKernel(grid_size, kBlockSize, device)( + resolve_future_token_ids_kernel, + static_cast(input_ids.data_ptr()), + static_cast(future_map.data_ptr()), + num_tokens); + } +}; + +} // namespace diff --git a/python/sglang/jit_kernel/resolve_future_token_ids.py b/python/sglang/jit_kernel/resolve_future_token_ids.py new file mode 100644 index 000000000..d0068a4ad --- /dev/null +++ b/python/sglang/jit_kernel/resolve_future_token_ids.py @@ -0,0 +1,41 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING + +import torch + +from sglang.jit_kernel.utils import cache_once, load_jit, make_cpp_args + +if TYPE_CHECKING: + from tvm_ffi.module import Module + + +@cache_once +def _jit_resolve_future_token_ids_module(dtype: torch.dtype) -> Module: + """Compile and cache the JIT module for a given dtype.""" + args = make_cpp_args(dtype) + return load_jit( + "resolve_future_token_ids", + *args, + cuda_files=["elementwise/resolve_future_token_ids.cuh"], + cuda_wrappers=[ + ( + "resolve_future_token_ids", + f"ResolveFutureTokenIds<{args}>::run", + ) + ], + ) + + +def resolve_future_token_ids_cuda( + input_ids: torch.Tensor, future_token_ids_map: torch.Tensor +) -> None: + """Resolve future token IDs in-place on CUDA. + + For each negative value in input_ids, replaces it with + future_token_ids_map[-value]. Non-negative values are unchanged. + + Supported dtypes: torch.int32, torch.int64. + """ + module = _jit_resolve_future_token_ids_module(input_ids.dtype) + module.resolve_future_token_ids(input_ids, future_token_ids_map) diff --git a/python/sglang/jit_kernel/tests/test_resolve_future_token_ids.py b/python/sglang/jit_kernel/tests/test_resolve_future_token_ids.py new file mode 100644 index 000000000..1e10ceb72 --- /dev/null +++ b/python/sglang/jit_kernel/tests/test_resolve_future_token_ids.py @@ -0,0 +1,63 @@ +import pytest +import torch + +from sglang.jit_kernel.resolve_future_token_ids import resolve_future_token_ids_cuda + + +def _reference_resolve(input_ids, future_map): + """Reference implementation using plain torch.""" + result = input_ids.clone() + result[:] = torch.where( + result < 0, + future_map[torch.clamp(-result, min=0)], + result, + ) + return result + + +@pytest.mark.parametrize("size", [1, 2, 127, 128, 255, 256, 1024, 4097]) +@pytest.mark.parametrize("dtype", [torch.int32, torch.int64]) +class TestResolveFutureTokenIds: + def test_all_negative(self, size: int, dtype: torch.dtype) -> None: + map_size = 8192 + future_map = torch.randint(0, 50000, (map_size,), dtype=dtype, device="cuda") + # Negative indices in range [-map_size+1, -1] + input_ids = -torch.randint(1, map_size, (size,), dtype=dtype, device="cuda") + + expected = _reference_resolve(input_ids, future_map) + resolve_future_token_ids_cuda(input_ids, future_map) + assert torch.equal(input_ids, expected) + + def test_all_non_negative(self, size: int, dtype: torch.dtype) -> None: + map_size = 16 + future_map = torch.randint(0, 50000, (map_size,), dtype=dtype, device="cuda") + input_ids = torch.randint(0, 50000, (size,), dtype=dtype, device="cuda") + + expected = input_ids.clone() + resolve_future_token_ids_cuda(input_ids, future_map) + assert torch.equal(input_ids, expected) + + def test_mixed(self, size: int, dtype: torch.dtype) -> None: + map_size = 8192 + future_map = torch.randint(0, 50000, (map_size,), dtype=dtype, device="cuda") + # Mix of negative and non-negative + input_ids = torch.randint( + -map_size + 1, 50000, (size,), dtype=dtype, device="cuda" + ) + + expected = _reference_resolve(input_ids, future_map) + resolve_future_token_ids_cuda(input_ids, future_map) + assert torch.equal(input_ids, expected) + + def test_zeros(self, size: int, dtype: torch.dtype) -> None: + map_size = 16 + future_map = torch.randint(0, 50000, (map_size,), dtype=dtype, device="cuda") + input_ids = torch.zeros(size, dtype=dtype, device="cuda") + + expected = input_ids.clone() + resolve_future_token_ids_cuda(input_ids, future_map) + assert torch.equal(input_ids, expected) + + +if __name__ == "__main__": + pytest.main([__file__, "-v", "-s"]) diff --git a/python/sglang/srt/managers/overlap_utils.py b/python/sglang/srt/managers/overlap_utils.py index 9246c2098..ee3886616 100644 --- a/python/sglang/srt/managers/overlap_utils.py +++ b/python/sglang/srt/managers/overlap_utils.py @@ -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_npu +from sglang.srt.utils import get_compiler_backend, is_cuda, is_hip if TYPE_CHECKING: from sglang.srt.managers.schedule_batch import ModelWorkerBatch @@ -14,11 +14,11 @@ if TYPE_CHECKING: from sglang.srt.speculative.eagle_info import EagleDraftInput from sglang.srt.speculative.spec_info import SpeculativeAlgorithm -_is_npu = is_npu() +_is_cuda = is_cuda() +_is_hip = is_hip() -@torch.compile(dynamic=True, backend=get_compiler_backend(), disable=_is_npu) -def _resolve_future_token_ids(input_ids, future_token_ids_map): +def _resolve_future_token_ids_native(input_ids, future_token_ids_map): input_ids[:] = torch.where( input_ids < 0, future_token_ids_map[torch.clamp(-input_ids, min=0)], @@ -26,6 +26,22 @@ def _resolve_future_token_ids(input_ids, future_token_ids_map): ) +if _is_cuda: + 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 + + @dataclass class FutureIndices: indices: torch.Tensor