chore: import upstream snapshot with attribution

python/sglang/kernels/ops/grammar/__init__.py

@@ -0,0 +1,26 @@
+"""Constrained-decoding / grammar kernels (Triton).
+
+The Triton kernels migrated here live in this package
+(``sglang.kernels.ops.grammar.<module>``); import them from there. Their
+``KernelSpec`` metadata is registered below for inventory (backend = Triton).
+"""
+
+from sglang.kernels.registry import register_kernel
+from sglang.kernels.spec import KernelBackend, KernelSpec
+
+# (module, public_fn) migrated from constrained/triton_ops.
+_TRITON_KERNELS = [
+    ("bitmask_ops", "apply_token_bitmask_inplace_triton"),
+    ("token_filter_ops", "set_token_filter_triton"),
+]
+for _mod, _fn in _TRITON_KERNELS:
+    register_kernel(
+        KernelSpec(
+            op=f"grammar.{_fn}",
+            backend=KernelBackend.TRITON,
+            target=f"sglang.kernels.ops.grammar.{_mod}:{_fn}",
+        )
+    )
+del _mod, _fn
+
+__all__ = []

python/sglang/kernels/ops/grammar/bitmask_ops.py

@@ -0,0 +1,141 @@
+# Adapt from
+# https://github.com/mlc-ai/xgrammar/blob/v0.1.17/python/xgrammar/kernels/apply_token_bitmask_inplace_triton.py
+
+from typing import List, Optional, Union
+
+import torch
+import triton
+import triton.language as tl
+
+from sglang.srt.utils import get_device_core_count
+
+
+@triton.jit
+def apply_token_bitmask_inplace_kernel(
+    logits_ptr,
+    bitmask_ptr,
+    indices_ptr,
+    num_rows,
+    vocab_size,
+    logits_strides,
+    bitmask_strides,
+    NUM_SMS: tl.constexpr,
+    BLOCK_SIZE: tl.constexpr,
+):
+    """Apply a bitmask to logits in-place using Triton. The bitmask is a 01 bitwise compressed tensor,
+    where 0 means the token is masked and 1 means the token is not masked. After applying the bitmask,
+    the masked logits will be set to -inf.
+
+    Parameters
+    ----------
+    logits_ptr : tl.tensor
+        Pointer to the logits tensor to apply the bitmask to.
+
+    bitmask_ptr : tl.tensor
+        Pointer to the bitmask tensor to apply.
+
+    indices_ptr : Optional[tl.tensor]
+        Optional pointer to indices tensor specifying which rows to apply the mask to.
+
+    num_rows : int
+        Number of rows to process. If indices_ptr is provided, this is the number of unique indices.
+
+    vocab_size : int
+        Size of the vocabulary dimension. If the logits does not have a vocab padding, this is the
+        same as the logits's second dimension. Otherwise, this is the actual size of the vocabulary.
+
+    logits_strides : int
+        Stride between rows in the logits tensor.
+
+    bitmask_strides : int
+        Stride between rows in the bitmask tensor.
+
+    NUM_SMS : int
+        Number of streaming multiprocessors to use.
+
+    BLOCK_SIZE : int
+        Size of processing blocks.
+    """
+
+    pid = tl.program_id(0)
+    num_blocks = tl.cdiv(vocab_size, BLOCK_SIZE)
+    for work_id in tl.range(pid, num_rows * num_blocks, NUM_SMS):
+        row_id = work_id // num_blocks
+        block_offset = (work_id % num_blocks) * BLOCK_SIZE
+        batch_id = row_id if indices_ptr is None else tl.load(indices_ptr + row_id)
+        offsets = block_offset + tl.arange(0, BLOCK_SIZE)
+        bitmask_offsets = block_offset // 32 + tl.arange(0, BLOCK_SIZE // 32)
+        vocab_mask = offsets < vocab_size
+        packed_bitmask_mask = bitmask_offsets < bitmask_strides
+        packed_bitmask = tl.load(
+            bitmask_ptr + batch_id * bitmask_strides + bitmask_offsets,
+            packed_bitmask_mask,
+        )
+        bitmask = ((packed_bitmask[:, None] >> (tl.arange(0, 32)[None, :])) & 1) == 0
+        bitmask = bitmask.reshape(BLOCK_SIZE)
+
+        tl.store(
+            logits_ptr + batch_id * logits_strides + offsets,
+            -float("inf"),
+            vocab_mask & bitmask,
+        )
+
+
+def apply_token_bitmask_inplace_triton(
+    logits: torch.Tensor,
+    bitmask: torch.Tensor,
+    indices: Optional[Union[List[int], torch.Tensor]] = None,
+):
+    NUM_SMS = get_device_core_count()
+    BLOCK_SIZE = 4096
+    BITS_PER_BLOCK = 32
+
+    # Check input dtype
+    assert bitmask.dtype == torch.int32, "bitmask must be of type int32"
+
+    # Check input tensor shapes.
+    logits_shape = logits.shape
+    bitmask_shape = bitmask.shape
+    if logits.ndim == 1:
+        logits_shape = (1, logits_shape[0])
+    if bitmask.ndim == 1:
+        bitmask_shape = (1, bitmask_shape[0])
+
+    required_bitmask_width = (logits_shape[1] + BITS_PER_BLOCK - 1) // BITS_PER_BLOCK
+    assert required_bitmask_width >= bitmask_shape[1], (
+        f"Bitmask width too large: allow at most {required_bitmask_width} int32s for "
+        f"logits' width {logits_shape[1]}, but got {bitmask_shape[1]}"
+    )
+
+    vocab_size = min(logits_shape[1], bitmask_shape[1] * BITS_PER_BLOCK)
+
+    num_rows = None
+    if isinstance(indices, list) or isinstance(indices, torch.Tensor):
+        indices = torch.tensor(indices, dtype=torch.int32, device=logits.device)
+        num_rows = indices.shape[0]
+    else:
+        assert (
+            logits_shape[0] == bitmask_shape[0]
+        ), f"batch size mismatch: logits {logits_shape[0]} vs bitmask {bitmask_shape[0]}"
+        num_rows = logits_shape[0]
+
+    if NUM_SMS > 0:
+        grid = (NUM_SMS,)
+    else:
+        num_blocks = triton.cdiv(vocab_size, BLOCK_SIZE)
+        grid = (num_rows * num_blocks,)
+        NUM_SMS = triton.next_power_of_2(grid[0])
+
+    apply_token_bitmask_inplace_kernel[grid](
+        logits,
+        bitmask,
+        indices,
+        num_rows,
+        vocab_size,
+        logits_shape[1],
+        bitmask_shape[1],
+        NUM_SMS,
+        BLOCK_SIZE,
+        num_warps=BLOCK_SIZE // 32 // (16 // logits.element_size()),
+        num_stages=3,
+    )

python/sglang/kernels/ops/grammar/token_filter_ops.py

@@ -0,0 +1,175 @@
+# Copyright 2026 SGLang Team
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Triton kernels for token filter operations."""
+
+from collections import OrderedDict
+from typing import List
+
+import torch
+import triton
+import triton.language as tl
+
+from sglang.srt.utils import get_device_core_count
+
+
+@triton.jit
+def reset_vocab_mask_kernel(
+    vocab_mask_ptr,
+    batch_idx: int,
+    num_elements: int,
+    reset_value: tl.constexpr,
+):
+    """Reset the vocab mask for a specific batch index to a given value.
+
+    Parameters
+    ----------
+    vocab_mask_ptr : tl.tensor
+        Pointer to the vocab mask tensor.
+
+    batch_idx : int
+        The batch index to reset.
+
+    num_elements : int
+        Number of int32 elements in the vocab mask for each batch.
+
+    reset_value : int
+        The value to reset the vocab mask to (typically -1 or 0).
+    """
+    pid = tl.program_id(0)
+    num_threads = tl.num_programs(0)
+
+    for i in tl.range(pid, num_elements, num_threads):
+        offset = batch_idx * num_elements + i
+        tl.store(vocab_mask_ptr + offset, reset_value)
+
+
+@triton.jit
+def set_token_filter_batch_kernel(
+    vocab_mask_ptr,
+    token_ids_ptr,
+    batch_idx: int,
+    num_tokens: int,
+    num_elements: int,
+    is_allowed: tl.constexpr,
+):
+    """Set or clear specific tokens in the vocab mask for a batch.
+
+    Each token ID maps to a specific bit in the int32 bitmask array.
+    The kernel sets or clears those bits using atomic operations.
+
+    Parameters
+    ----------
+    vocab_mask_ptr : tl.tensor
+        Pointer to the vocab mask tensor.
+
+    token_ids_ptr : tl.tensor
+        Pointer to the token IDs to set/clear.
+
+    batch_idx : int
+        The batch index to modify.
+
+    num_tokens : int
+        Number of tokens to process.
+
+    num_elements : int
+        Number of int32 elements in the vocab mask for each batch.
+
+    is_allowed : bool
+        If True, set the bit to 1 (allow token).
+        If False, clear the bit to 0 (block token).
+    """
+    pid = tl.program_id(0)
+    num_threads = tl.num_programs(0)
+
+    for i in tl.range(pid, num_tokens, num_threads):
+        token_id = tl.load(token_ids_ptr + i)
+        element_idx = token_id // 32
+        bit_idx = token_id % 32
+
+        offset = batch_idx * num_elements + element_idx
+
+        if is_allowed:
+            tl.atomic_or(vocab_mask_ptr + offset, 1 << bit_idx)
+        else:
+            tl.atomic_and(vocab_mask_ptr + offset, ~(1 << bit_idx))
+
+
+_cached_num_sms = None
+_cached_token_id_tensors: OrderedDict[tuple[int, tuple[int, ...]], torch.Tensor] = (
+    OrderedDict()
+)
+_MAX_TOKEN_ID_TENSOR_CACHE_SIZE = 32
+
+
+def _compute_grid(work_items: int):
+    global _cached_num_sms
+    if _cached_num_sms is None:
+        _cached_num_sms = get_device_core_count()
+    if _cached_num_sms > 0:
+        return (min(_cached_num_sms, work_items),)
+    return (work_items,)
+
+
+def _get_cached_token_ids_tensor(
+    token_ids: List[int], device: torch.device
+) -> torch.Tensor:
+    key = (device.index or 0, tuple(token_ids))
+    cached = _cached_token_id_tensors.get(key)
+    if cached is not None:
+        _cached_token_id_tensors.move_to_end(key)
+        return cached
+
+    token_ids_tensor = torch.tensor(token_ids, dtype=torch.int32, device=device)
+    _cached_token_id_tensors[key] = token_ids_tensor
+    if len(_cached_token_id_tensors) > _MAX_TOKEN_ID_TENSOR_CACHE_SIZE:
+        _cached_token_id_tensors.popitem(last=False)
+    return token_ids_tensor
+
+
+def set_token_filter_triton(
+    vocab_mask: torch.Tensor,
+    token_ids: List[int],
+    batch_idx: int,
+    is_allowed: bool = True,
+    reset_vocab_mask: bool = True,
+):
+    """Set or clear specific tokens in the vocab mask using Triton."""
+    assert vocab_mask.device.type == "cuda"
+
+    num_elements = vocab_mask.shape[1]
+
+    if reset_vocab_mask:
+        reset_value = 0 if is_allowed else -1
+        reset_vocab_mask_kernel[_compute_grid(num_elements)](
+            vocab_mask,
+            batch_idx,
+            num_elements,
+            reset_value,
+            num_warps=4,
+        )
+
+    if not token_ids:
+        return
+
+    num_tokens = len(token_ids)
+    token_ids_tensor = _get_cached_token_ids_tensor(token_ids, vocab_mask.device)
+    set_token_filter_batch_kernel[_compute_grid(num_tokens)](
+        vocab_mask,
+        token_ids_tensor,
+        batch_idx,
+        num_tokens,
+        num_elements,
+        is_allowed,
+        num_warps=4,
+    )