Optimize ngram decode id computation (#24757)

Co-authored-by: Codex <codex@example.com> Co-authored-by: BBuf <xiaoyu.zhang@radixark.net>
2026-06-02 17:37:34 +08:00
parent f651b48764
commit 84e1108312
5 changed files with 417 additions and 22 deletions
@@ -0,0 +1,97 @@
+import sys
+
+import pytest
+import torch
+
+from sglang.jit_kernel.ngram_embedding import (
+    compute_n_gram_ids,
+    compute_n_gram_ids_decode,
+)
+from sglang.test.ci.ci_register import register_cuda_ci
+
+register_cuda_ci(est_time=30, suite="base-b-kernel-unit-1-gpu-large")
+
+
+def _make_ngram_params(ne_n: int, ne_k: int, vocab_size: int):
+    ne_weights = torch.zeros([ne_n - 1, ne_k, ne_n], dtype=torch.int32)
+    ne_mods = torch.zeros([ne_n - 1, ne_k], dtype=torch.int32)
+    exclusive_sums = torch.zeros([(ne_n - 1) * ne_k + 1], dtype=torch.int32)
+
+    for n in range(2, ne_n + 1):
+        for k in range(ne_k):
+            config_id = (n - 2) * ne_k + k
+            mod = 65537 + 2 * config_id
+            ne_mods[n - 2][k] = mod
+            exclusive_sums[config_id + 1] = exclusive_sums[config_id] + mod
+            for delta in range(ne_n):
+                ne_weights[n - 2][k][delta] = pow(vocab_size, delta, mod)
+
+    return (
+        ne_weights.cuda(),
+        ne_mods.cuda(),
+        exclusive_sums.cuda(),
+    )
+
+
+@pytest.mark.parametrize("batch_size", [1, 2, 17, 128, 1024])
+def test_compute_n_gram_ids_decode_matches_general(batch_size: int) -> None:
+    ne_n = 8
+    ne_k = 2
+    vocab_size = 32000
+    max_context_len = 1024
+    max_running_reqs = batch_size + 8
+    num_configs = (ne_n - 1) * ne_k
+
+    ne_weights, ne_mods, exclusive_sums = _make_ngram_params(ne_n, ne_k, vocab_size)
+    ne_token_table = torch.randint(
+        0,
+        vocab_size,
+        (max_running_reqs, max_context_len),
+        dtype=torch.int32,
+        device="cuda",
+    )
+    row_indices = torch.randperm(max_running_reqs, device="cuda")[:batch_size].to(
+        torch.int64
+    )
+    column_starts = torch.randint(
+        0, max_context_len, (batch_size,), dtype=torch.int32, device="cuda"
+    )
+    tokens = torch.empty(batch_size, dtype=torch.int32, device="cuda")
+    exclusive_req_len_sums = torch.arange(
+        batch_size + 1, dtype=torch.int32, device="cuda"
+    )
+    n_gram_ids_general = torch.empty(
+        (batch_size, num_configs), dtype=torch.int32, device="cuda"
+    )
+    n_gram_ids_decode = torch.empty_like(n_gram_ids_general)
+
+    compute_n_gram_ids(
+        ne_n=ne_n,
+        ne_k=ne_k,
+        ne_weights=ne_weights,
+        ne_mods=ne_mods,
+        exclusive_ne_embedder_size_sums=exclusive_sums,
+        tokens=tokens,
+        exclusive_req_len_sums=exclusive_req_len_sums,
+        ne_token_table=ne_token_table,
+        row_indices=row_indices,
+        column_starts=column_starts,
+        n_gram_ids=n_gram_ids_general,
+    )
+    compute_n_gram_ids_decode(
+        ne_n=ne_n,
+        ne_k=ne_k,
+        ne_weights=ne_weights,
+        ne_mods=ne_mods,
+        exclusive_ne_embedder_size_sums=exclusive_sums,
+        ne_token_table=ne_token_table,
+        row_indices=row_indices,
+        column_starts=column_starts,
+        n_gram_ids=n_gram_ids_decode,
+    )
+
+    torch.testing.assert_close(n_gram_ids_decode, n_gram_ids_general, atol=0, rtol=0)
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main([__file__, "-v", "-s"]))