[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)