Replace _resolve_future_token_ids with JIT kernel + platform dispatch (#20976)

Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

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)

python/sglang/jit_kernel/csrc/elementwise/resolve_future_token_ids.cuh

@@ -0,0 +1,57 @@
+#include <sgl_kernel/tensor.h>  // For TensorMatcher, SymbolicSize, SymbolicDevice
+#include <sgl_kernel/utils.h>   // For RuntimeCheck, div_ceil
+
+#include <sgl_kernel/utils.cuh>  // For LaunchKernel
+
+#include <dlpack/dlpack.h>
+#include <tvm/ffi/container/tensor.h>
+
+#include <cstddef>
+#include <cstdint>
+
+namespace {
+
+template <typename T>
+__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 <typename T>
+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<kDLCUDA>();
+
+    TensorMatcher({N}).with_dtype<T>().with_device(device_).verify(input_ids);
+
+    TensorMatcher({M}).with_dtype<T>().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<T>,
+        static_cast<T*>(input_ids.data_ptr()),
+        static_cast<const T*>(future_map.data_ptr()),
+        num_tokens);
+  }
+};
+
+}  // namespace

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)

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"])

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