diff --git a/python/sglang/kernels/ops/attention/clamp_position.py b/python/sglang/kernels/ops/attention/clamp_position.py index d7156a66b..8753bfcaa 100644 --- a/python/sglang/kernels/ops/attention/clamp_position.py +++ b/python/sglang/kernels/ops/attention/clamp_position.py @@ -4,7 +4,13 @@ from typing import TYPE_CHECKING import torch +from sglang.kernels.fused_op import BaseFusedOp, register_fused_op from sglang.kernels.jit.utils import cache_once, load_jit, make_cpp_args +from sglang.kernels.spec import ( + CapabilityRequirement, + FormatSignature, + KernelBackend, +) if TYPE_CHECKING: from tvm_ffi.module import Module @@ -33,3 +39,38 @@ def clamp_position_cuda(seq_lens: torch.Tensor) -> torch.Tensor: module = _jit_clamp_position_module(seq_lens.dtype) module.clamp_position(dst, seq_lens) return dst + + +class ClampPositionOp(BaseFusedOp): + """Compute non-negative, zero-based decode positions.""" + + op = "attention.clamp_position" + priority = (KernelBackend.JIT, KernelBackend.TORCH) + capabilities = { + KernelBackend.JIT: frozenset( + {CapabilityRequirement.CUDA, CapabilityRequirement.HIP} + ) + } + format_signature = FormatSignature( + supported_dtypes=("int32", "int64"), + description="clamp(seq_lens - 1, min=0); returns int64 positions", + ) + descriptions = { + KernelBackend.JIT: "Fused clamp-position kernel (sglang.kernels.jit).", + } + + def forward_native(self, seq_lens: torch.Tensor) -> torch.Tensor: + return torch.clamp(seq_lens - 1, min=0).to(torch.int64) + + def forward_jit(self, seq_lens: torch.Tensor) -> torch.Tensor: + return clamp_position_cuda(seq_lens).to(torch.int64) + + +_CLAMP_POSITION = register_fused_op(ClampPositionOp(), __name__, "_CLAMP_POSITION") + + +def clamp_position(seq_lens: torch.Tensor) -> torch.Tensor: + return _CLAMP_POSITION(seq_lens) + + +__all__ = ["ClampPositionOp", "clamp_position", "clamp_position_cuda"] diff --git a/python/sglang/srt/model_executor/forward_batch_info.py b/python/sglang/srt/model_executor/forward_batch_info.py index 42318607a..04b28b8ac 100644 --- a/python/sglang/srt/model_executor/forward_batch_info.py +++ b/python/sglang/srt/model_executor/forward_batch_info.py @@ -36,6 +36,7 @@ from typing import TYPE_CHECKING, Dict, List, Optional, Set, Tuple, Union import torch +from sglang.kernels.ops.attention.clamp_position import clamp_position from sglang.kernels.ops.attention.position import compute_position_triton from sglang.srt.configs.hybrid_arch import mambaish_config from sglang.srt.environ import envs @@ -61,7 +62,6 @@ from sglang.srt.runtime_context import ( from sglang.srt.speculative.spec_info import SpecInputType from sglang.srt.utils import ( is_cpu, - is_cuda, is_hip, is_npu, support_triton, @@ -1940,18 +1940,6 @@ def compute_position_torch( return positions.to(torch.int64), extend_start_loc -def _clamp_position_native(seq_lens): - return torch.clamp((seq_lens - 1), min=0).to(torch.int64) - - -if is_cuda() or is_hip(): - from sglang.kernels.ops.attention.clamp_position import clamp_position_cuda - - clamp_position = clamp_position_cuda -else: - clamp_position = _clamp_position_native - - def _hash_rids_to_tensor(*, rids: List[str], device: torch.device) -> torch.Tensor: values: List[int] = [_stable_hash_str_to_i64(rid) for rid in rids] return torch.tensor(values, dtype=torch.int64, device=device) diff --git a/test/registered/kernels/ops/attention/test_clamp_position.py b/test/registered/kernels/ops/attention/test_clamp_position.py index c33d3b844..579e5b4d7 100644 --- a/test/registered/kernels/ops/attention/test_clamp_position.py +++ b/test/registered/kernels/ops/attention/test_clamp_position.py @@ -3,7 +3,11 @@ import sys import pytest import torch -from sglang.kernels.ops.attention.clamp_position import clamp_position_cuda +from sglang.kernels.ops.attention.clamp_position import ( + ClampPositionOp, + clamp_position, + clamp_position_cuda, +) from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci register_cuda_ci(est_time=12, stage="base-b-kernel-unit", runner_config="1-gpu-large") @@ -14,6 +18,10 @@ def _reference_clamp_position(seq_lens): return torch.clamp(seq_lens - 1, min=0).to(seq_lens.dtype) +def _reference_public_clamp_position(seq_lens): + return torch.clamp(seq_lens - 1, min=0).to(torch.int64) + + @pytest.mark.parametrize("size", [1, 2, 127, 128, 255, 256, 1024, 4097]) @pytest.mark.parametrize("dtype", [torch.int32, torch.int64]) class TestClampPosition: @@ -42,5 +50,23 @@ class TestClampPosition: assert torch.equal(result, expected) +@pytest.mark.parametrize("dtype", [torch.int32, torch.int64]) +def test_native(dtype: torch.dtype) -> None: + seq_lens = torch.tensor([0, 1, 5, 128], dtype=dtype) + expected = _reference_public_clamp_position(seq_lens) + result = ClampPositionOp().forward_native(seq_lens) + assert result.dtype == torch.int64 + assert torch.equal(result, expected) + + +@pytest.mark.parametrize("dtype", [torch.int32, torch.int64]) +def test_public_dispatch(dtype: torch.dtype) -> None: + seq_lens = torch.tensor([0, 1, 5, 128], dtype=dtype, device="cuda") + expected = _reference_public_clamp_position(seq_lens) + result = clamp_position(seq_lens) + assert result.dtype == torch.int64 + assert torch.equal(result, expected) + + if __name__ == "__main__": sys.exit(pytest.main([__file__, "-v", "-s"]))