[Advanced Compiler]Add Vllm::apply_repetition_penalties_kernel (flagos-ai#1239)

* Add apply_repetition_penalties operator: implement operation, update tests and benchmarks

* Corrected the registration logic

* Corrected the registration logic

* Corrected the registration logic

* Corrected the registration logic

---------

Co-authored-by: yy33min <2811552420@qq.com>
Co-authored-by: henghengxiedaima <1149963331@qq.com>

Author: 3 people

benchmark/test_unary_pointwise_perf.py

@@ -305,3 +305,65 @@ def test_bitwise_right_shift_perf():
         dtypes=INT_DTYPES,
     )
     bench.run()
+
+
+class RepetitionPenaltyBenchmark(Benchmark):
+    def __init__(self, op_name, torch_op, dtypes):
+        super().__init__(op_name, torch_op, dtypes)
+        self.gems_op = None
+
+    def set_shapes(self, shape_file_path=None):
+        self.shapes = [
+            (1, 1024),
+            (1, 4096),
+            (1, 8192),
+            (8, 4096),
+            (16, 4096),
+            (32, 1024),
+            (8, 8192),
+            (64, 32000),
+        ]
+
+    def get_input_iter(self, cur_dtype):
+        for shape in self.shapes:
+            num_seqs, vocab_size = shape
+            yield (
+                torch.randn(shape, dtype=cur_dtype, device=self.device),
+                torch.randint(0, 2, shape, dtype=torch.bool, device=self.device),
+                torch.randint(0, 2, shape, dtype=torch.bool, device=self.device),
+                torch.empty(num_seqs, dtype=cur_dtype, device=self.device).uniform_(
+                    1.0, 2.0
+                ),
+            )
+
+    def set_gems(self, gems_op):
+        self.gems_op = gems_op
+
+
+UNSUPPORTED_VENDORS = {
+    "metax",
+    "kunlunxin",
+    "iluvatar",
+    "mthreads",
+    "hygon",
+    "cambricon",
+}
+
+
+@pytest.mark.skipif(SkipVersion("vllm", "<0.4"), reason="vLLM <0.4 not supported")
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
+@pytest.mark.skipif(
+    flag_gems.vendor_name in UNSUPPORTED_VENDORS, reason="Vendor not supported"
+)
+@pytest.mark.apply_repetition_penalties
+@pytest.mark.performance
+def test_perf_repetition_penalty():
+    vllm_ops = pytest.importorskip("vllm._custom_ops")
+
+    bench = RepetitionPenaltyBenchmark(
+        op_name="apply_repetition_penalties",
+        torch_op=vllm_ops.apply_repetition_penalties,
+        dtypes=FLOAT_DTYPES,
+    )
+    bench.set_gems(flag_gems.apply_repetition_penalties)
+    bench.run()

src/flag_gems/fused/__init__.py

@@ -1,3 +1,4 @@
+from flag_gems.fused.apply_repetition_penalties import apply_repetition_penalties
 from flag_gems.fused.concat_and_cache_mla import concat_and_cache_mla
 from flag_gems.fused.cross_entropy_loss import cross_entropy_loss
 from flag_gems.fused.flash_mla import flash_mla
@@ -23,6 +24,7 @@
 from flag_gems.fused.weight_norm import weight_norm
 
 __all__ = [
+    "apply_repetition_penalties",
     "apply_rotary_pos_emb",
     "skip_layer_norm",
     "fused_add_rms_norm",

src/flag_gems/fused/apply_repetition_penalties.py

@@ -0,0 +1,81 @@
+import torch
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def _repetition_penalty_kernel(
+    logits_ptr,
+    prompt_mask_ptr,
+    output_mask_ptr,
+    penalties_ptr,
+    num_seqs,
+    vocab_size,
+    BLOCK_SIZE: tl.constexpr,
+):
+    seq_idx = tl.program_id(0)
+    vocab_offset = tl.program_id(1) * BLOCK_SIZE
+
+    if seq_idx >= num_seqs:
+        return
+
+    penalty = tl.load(penalties_ptr + seq_idx)
+
+    vocab_idx = vocab_offset + tl.arange(0, BLOCK_SIZE)
+
+    valid_vocab = vocab_idx < vocab_size
+
+    logits_idx = seq_idx * vocab_size + vocab_idx
+    mask_idx = logits_idx
+
+    prompt_mask = tl.load(prompt_mask_ptr + mask_idx, mask=valid_vocab, other=False)
+    output_mask = tl.load(output_mask_ptr + mask_idx, mask=valid_vocab, other=False)
+    logits = tl.load(logits_ptr + logits_idx, mask=valid_vocab, other=0.0)
+
+    is_repeated = prompt_mask | output_mask
+
+    logits = tl.where(is_repeated & (logits > 0), logits / penalty, logits)
+    logits = tl.where(is_repeated & (logits <= 0), logits * penalty, logits)
+
+    tl.store(logits_ptr + logits_idx, logits, mask=valid_vocab)
+
+
+def apply_repetition_penalties(logits, prompt_mask, output_mask, repetition_penalties):
+    assert logits.is_contiguous(), "logits must be contiguous"
+    assert (
+        prompt_mask.is_contiguous() and prompt_mask.dtype == torch.bool
+    ), "prompt_mask must be contiguous bool tensor"
+    assert (
+        output_mask.is_contiguous() and output_mask.dtype == torch.bool
+    ), "output_mask must be contiguous bool tensor"
+    assert (
+        repetition_penalties.is_contiguous()
+    ), "repetition_penalties must be contiguous"
+    assert logits.dim() == 2, f"logits must be 2D, got {logits.dim()}D"
+    assert (
+        logits.shape == prompt_mask.shape == output_mask.shape
+    ), "shape mismatch between logits and masks"
+    assert (
+        repetition_penalties.dim() == 1
+        and repetition_penalties.numel() == logits.shape[0]
+    ), "repetition_penalties must be 1D with length equal to num_seqs"
+
+    num_seqs, vocab_size = logits.shape
+
+    BLOCK_SIZE = 1024
+
+    grid = (
+        num_seqs,
+        triton.cdiv(vocab_size, BLOCK_SIZE),
+    )
+
+    _repetition_penalty_kernel[grid](
+        logits,
+        prompt_mask,
+        output_mask,
+        repetition_penalties,
+        num_seqs,
+        vocab_size,
+        BLOCK_SIZE=BLOCK_SIZE,
+    )
+    return None

src/flag_gems/patches/patch_vllm_all.py

@@ -286,6 +286,17 @@ def custom_topk_softmax(
     )
 
 
+def custom_apply_repetition_penalties(
+    logits: torch.Tensor,
+    prompt_mask: torch.Tensor,
+    output_mask: torch.Tensor,
+    repetition_penalties: torch.Tensor,
+):
+    return flag_gems.apply_repetition_penalties(
+        logits, prompt_mask, output_mask, repetition_penalties
+    )
+
+
 def custom_get_scheduler_metadata(
     batch_size: int,
     max_seqlen_q: int,
@@ -409,3 +420,10 @@ def apply_gems_patches_to_vllm(verbose=True):
         "CUDA",
         verbose,
     )
+    patch_vllm_lib(
+        "_C",
+        "apply_repetition_penalties_",
+        custom_apply_repetition_penalties,
+        "CUDA",
+        verbose,
+    )

tests/test_unary_pointwise_ops.py

@@ -1458,3 +1458,88 @@ def test_accuracy_reglu(shape, dtype):
         res_out = flag_gems.reglu(input_tensor)
 
     gems_assert_close(res_out, ref_out, dtype)
+
+
+def _init_vllm():
+    if not torch.cuda.is_available():
+        return None, False
+    try:
+        from vllm._custom_ops import apply_repetition_penalties as fn
+
+        t, m = torch.randn(2, 1024, device="cuda"), torch.zeros(
+            2, 1024, dtype=torch.bool, device="cuda"
+        )
+        fn(t, m, m, torch.full((2,), 1.2, device="cuda"))
+        return fn, True
+    except (ImportError, RuntimeError):
+
+        def fallback(logits, pm, om, pens):
+            for i in range(logits.shape[0]):
+                m = pm[i] | om[i]
+                logits[i][m] = torch.where(
+                    logits[i][m] > 0, logits[i][m] / pens[i], logits[i][m] * pens[i]
+                )
+
+        return fallback, True
+
+
+_vllm_fn, _VLLM_OK = _init_vllm()
+
+_REP_PENALTY_CFG = {
+    "shapes": [
+        (1, 1024),
+        (1, 4096),
+        (1, 8192),
+        (8, 4096),
+        (16, 4096),
+        (32, 1024),
+        (8, 8192),
+    ],
+    "penalties": [1.0, 1.2, 1.5],
+    "device": torch.device("cuda:0"),
+}
+
+
+@pytest.mark.apply_repetition_penalties
+@pytest.mark.skipif(
+    not _VLLM_OK or not torch.cuda.is_available(), reason="need VLLM+CUDA"
+)
+@pytest.mark.parametrize("shape", _REP_PENALTY_CFG["shapes"])
+@pytest.mark.parametrize("penalty", _REP_PENALTY_CFG["penalties"])
+@pytest.mark.parametrize("dtype", FLOAT_DTYPES)
+@pytest.mark.parametrize("mask_mode", ["random", "empty"])
+def test_repetition_penalty(shape, penalty, dtype, mask_mode):
+    device = _REP_PENALTY_CFG["device"]
+
+    logits = torch.randn(shape, dtype=dtype, device=device).contiguous()
+    logits_ori = logits.clone()
+
+    if mask_mode == "random":
+        prompt_mask = torch.randint(0, 2, shape, dtype=torch.bool, device=device)
+        output_mask = torch.randint(0, 2, shape, dtype=torch.bool, device=device)
+    else:
+        prompt_mask = torch.zeros(shape, dtype=torch.bool, device=device)
+        output_mask = torch.zeros(shape, dtype=torch.bool, device=device)
+
+    penalties = torch.full((shape[0],), penalty, dtype=dtype, device=device)
+
+    logits_vllm = logits.clone()
+    _vllm_fn(logits_vllm, prompt_mask.clone(), output_mask.clone(), penalties.clone())
+    ref = to_reference(logits_vllm, True).to(dtype)
+
+    with flag_gems.use_gems():
+        flag_gems.apply_repetition_penalties(
+            logits, prompt_mask, output_mask, penalties
+        )
+    res = to_reference(logits, True).to(dtype)
+
+    gems_assert_close(res, ref, dtype)
+
+    has_mask = (prompt_mask | output_mask).any().item()
+    should_modify = has_mask and penalty != 1.0
+    if should_modify:
+        assert not torch.equal(
+            to_reference(logits, True), to_reference(logits_ori, True)
+        ), "In-place未生效"
+    elif mask_mode == "empty":
+        gems_assert_close(res, to_reference(logits_ori, True).to(dtype), dtype)