lightseekorg--tokenspeed/tokenspeed-kernel/python/tokenspeed_kernel/ops/sampling/triton/penalties.py

# SPDX-License-Identifier: MIT AND Apache-2.0
# SPDX-FileCopyrightText: Copyright (c) 2026 LightSeek Foundation
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
#
# Copyright (c) 2026 LightSeek Foundation
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

# TokenSpeed-owned penalty, logit-bias, and count kernels.

from __future__ import annotations

import torch
from tokenspeed_kernel._triton import tl, triton


@triton.jit
def _apply_penalties_logit_bias_inplace_kernel(
    logits_ptr,
    req_pool_indices_ptr,
    counts_ptr,
    logit_bias_ptr,
    freq_pen_pool_ptr,
    pres_pen_pool_ptr,
    rep_pen_pool_ptr,
    vocab_size: tl.constexpr,
    logits_row_stride: tl.constexpr,
    counts_row_stride: tl.constexpr,
    bias_row_stride: tl.constexpr,
    BLOCK_SIZE: tl.constexpr,
    NUM_TOKENS_PER_REQ: tl.constexpr,
):
    row = tl.program_id(0)
    block = tl.program_id(1)
    req_row = row // NUM_TOKENS_PER_REQ
    pool_idx = tl.load(req_pool_indices_ptr + req_row)
    cols = block * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
    mask = cols < vocab_size

    logits_offsets = row * logits_row_stride + cols
    state_offsets = pool_idx * counts_row_stride + cols
    bias_offsets = pool_idx * bias_row_stride + cols

    vals = tl.load(logits_ptr + logits_offsets, mask=mask, other=0.0).to(tl.float32)
    counts = tl.load(counts_ptr + state_offsets, mask=mask, other=0).to(tl.float32)
    active = counts > 0.0

    rep = tl.load(rep_pen_pool_ptr + pool_idx).to(tl.float32)
    freq = tl.load(freq_pen_pool_ptr + pool_idx).to(tl.float32)
    presence = tl.load(pres_pen_pool_ptr + pool_idx).to(tl.float32)

    rep_vals = tl.where(vals > 0.0, vals / rep, vals * rep)
    vals = tl.where(active, rep_vals, vals)
    vals = vals - freq * counts - presence * active.to(tl.float32)
    vals += tl.load(logit_bias_ptr + bias_offsets, mask=mask, other=0.0).to(tl.float32)

    tl.store(logits_ptr + logits_offsets, vals, mask=mask)


def apply_penalties_logit_bias_inplace(
    logits: torch.Tensor,
    req_pool_indices: torch.Tensor,
    counts: torch.Tensor,
    logit_bias: torch.Tensor,
    freq_pen_pool: torch.Tensor,
    pres_pen_pool: torch.Tensor,
    rep_pen_pool: torch.Tensor,
    *,
    num_tokens_per_req: int = 1,
) -> torch.Tensor:
    """Apply repetition/frequency/presence penalties and logit_bias in-place."""
    if logits.ndim != 2:
        raise ValueError(f"logits must be 2D, got {logits.ndim}D")
    if counts.ndim != 2 or logit_bias.ndim != 2:
        raise ValueError("counts and logit_bias must be 2D")
    if logits.device.type != "cuda":
        raise ValueError("apply_penalties_logit_bias_inplace requires CUDA logits")
    if logits.stride(-1) != 1:
        raise ValueError(
            "apply_penalties_logit_bias_inplace requires stride-1 vocab dimension, "
            f"got stride={logits.stride()}"
        )
    if req_pool_indices.dtype != torch.int32:
        raise ValueError(
            f"req_pool_indices must be int32, got {req_pool_indices.dtype}"
        )
    if counts.dtype != torch.int32:
        raise ValueError(f"counts must be int32, got {counts.dtype}")
    if num_tokens_per_req <= 0:
        raise ValueError("num_tokens_per_req must be positive")

    rows, vocab_size = logits.shape
    if rows % num_tokens_per_req != 0:
        raise ValueError(
            "logits rows must be divisible by num_tokens_per_req, "
            f"got rows={rows}, num_tokens_per_req={num_tokens_per_req}"
        )
    request_rows = rows // num_tokens_per_req
    if req_pool_indices.shape[0] != request_rows:
        raise ValueError(
            "req_pool_indices length must match request rows, "
            f"got {req_pool_indices.shape[0]} and {request_rows}"
        )
    if counts.shape[1] < vocab_size or logit_bias.shape[1] < vocab_size:
        raise ValueError(
            "counts/logit_bias vocab dimension must cover logits vocab, "
            f"got counts={counts.shape}, logit_bias={logit_bias.shape}, logits={logits.shape}"
        )
    if rows == 0:
        return logits

    num_blocks = triton.cdiv(vocab_size, 1024)
    _apply_penalties_logit_bias_inplace_kernel[(rows, num_blocks)](
        logits,
        req_pool_indices,
        counts,
        logit_bias,
        freq_pen_pool,
        pres_pen_pool,
        rep_pen_pool,
        vocab_size=vocab_size,
        logits_row_stride=logits.stride(0),
        counts_row_stride=counts.stride(0),
        bias_row_stride=logit_bias.stride(0),
        BLOCK_SIZE=1024,
        NUM_TOKENS_PER_REQ=num_tokens_per_req,
        num_warps=4,
        num_stages=3,
    )
    return logits


@triton.jit
def _accumulate_counts_inplace_kernel(
    counts_ptr,
    pool_idx_ptr,
    tokens_ptr,
    weights_ptr,
    total: tl.constexpr,
    counts_row_stride: tl.constexpr,
    vocab_size: tl.constexpr,
    BLOCK_SIZE: tl.constexpr,
):
    offs = tl.program_id(0) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
    mask = offs < total
    weights = tl.load(weights_ptr + offs, mask=mask, other=0).to(tl.int32)
    pool_idx = tl.load(pool_idx_ptr + offs, mask=mask, other=0).to(tl.int64)
    tokens = tl.load(tokens_ptr + offs, mask=mask, other=0).to(tl.int64)
    valid = mask & (weights != 0) & (tokens >= 0) & (tokens < vocab_size)
    tl.atomic_add(
        counts_ptr + pool_idx * counts_row_stride + tokens,
        weights,
        sem="relaxed",
        mask=valid,
    )


def accumulate_counts_inplace(
    counts: torch.Tensor,
    pool_idx: torch.Tensor,
    tokens: torch.Tensor,
    weights: torch.Tensor,
) -> None:
    """Graph-safe ``counts[pool_idx, tokens] += weights``."""
    if counts.ndim != 2:
        raise ValueError(f"counts must be 2D, got {counts.ndim}D")
    if counts.device.type != "cuda":
        raise ValueError("accumulate_counts_inplace requires CUDA counts")
    if counts.dtype != torch.int32:
        raise ValueError(f"counts must be int32, got {counts.dtype}")
    if pool_idx.dtype != torch.int32:
        raise ValueError(f"pool_idx must be int32, got {pool_idx.dtype}")
    if weights.dtype != torch.int32:
        raise ValueError(f"weights must be int32, got {weights.dtype}")
    if tokens.dtype not in (torch.int32, torch.int64):
        raise ValueError(f"tokens must be int32 or int64, got {tokens.dtype}")
    total = int(tokens.numel())
    if pool_idx.numel() != total or weights.numel() != total:
        raise ValueError(
            "pool_idx, tokens, and weights must have the same number of elements"
        )
    if total == 0:
        return

    _accumulate_counts_inplace_kernel[(triton.cdiv(total, 256),)](
        counts,
        pool_idx.reshape(-1),
        tokens.reshape(-1),
        weights.reshape(-1),
        total=total,
        counts_row_stride=counts.stride(0),
        vocab_size=counts.shape[1],
        BLOCK_SIZE=256,
        num_warps=4,
    )