[Kernel] Add OOT dispatch for clamp position (#38687)

Co-authored-by: Xiaoyu Zhang <1182563586@qq.com>

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

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)

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