A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 2, 11:00 UTC
[CPU Backend] CPU top-k and top-p sampling kernels using Triton (#43633)
Signed-off-by: Li, Tianmu <tianmu.li@intel.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Tianmu Li committed
May 29, 2026 at 00:02 UTC
94d3f4d205fdfe44212a7958c6af647ecba15556
5 files changed
+23
-6
.buildkite/hardware_tests/cpu.yaml
+6
-1
@@ -62,11 +62,16 @@ steps:
62
source_file_dependencies:
63
- vllm/v1/worker/cpu/
64
- vllm/v1/worker/gpu/
65
+ - vllm/v1/sample/ops/topk_topp_triton.py
66
+ - vllm/v1/sample/ops/topk_topp_sampler.py
67
+ - tests/v1/sample/test_topk_topp_sampler.py
68
commands:
69
- |
70
bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 45m "
71
uv pip install git+https://github.com/triton-lang/triton-cpu.git@270e696d
69
- VLLM_USE_V2_MODEL_RUNNER=1 pytest -x -v -s tests/models/language/generation/test_granite.py -m cpu_model"
72
+ VLLM_USE_V2_MODEL_RUNNER=1 pytest -x -v -s tests/models/language/generation/test_granite.py -m cpu_model
73
+ # TODO: move to CPU-Kernel Tests once triton-cpu has a pre-built wheel
74
+ pytest -x -v -s tests/v1/sample/test_topk_topp_sampler.py::TestTritonTopkTopp"
75
76
- label: CPU-Quantization Model Tests
77
depends_on: []
tests/v1/sample/test_topk_topp_sampler.py
+2
-1
@@ -5,6 +5,7 @@ import torch
5
from torch import Generator
6
7
from vllm.platforms import current_platform
8
+from vllm.triton_utils import HAS_TRITON
9
from vllm.v1.sample.ops.topk_topp_sampler import apply_top_k_top_p_pytorch
10
11
DEVICE_TYPE = current_platform.device_type
@@ -151,7 +152,7 @@ def test_flashinfer_sampler():
152
# =============================================================================
153
154
154
-@pytest.mark.skipif("cpu" in DEVICE_TYPE, reason="CUDA/XPU not available")
155
+@pytest.mark.skipif(not HAS_TRITON, reason="Triton not available on this platform")
156
class TestTritonTopkTopp:
157
"""Tests for the Triton top-k/top-p kernel."""
158
vllm/platforms/cpu.py
+4
@@ -408,6 +408,10 @@ class CpuPlatform(Platform):
408
def support_hybrid_kv_cache(cls) -> bool:
409
return True
410
411
+ @classmethod
412
+ def num_compute_units(cls, device_id: int = 0) -> int:
413
+ return torch.get_num_threads()
414
+
415
@classmethod
416
def import_kernels(cls) -> None:
417
if Platform.get_cpu_architecture() in (CpuArchEnum.X86,):
vllm/v1/sample/ops/topk_topp_sampler.py
+3
-2
@@ -168,7 +168,7 @@ class TopKTopPSampler(nn.Module):
168
169
The logits tensor may be updated in-place.
170
"""
171
- logits = apply_top_k_top_p_pytorch(logits, k, p, allow_cpu_sync=True)
171
+ logits = apply_top_k_top_p(logits, k, p)
172
logits_to_return = None
173
if self.logprobs_mode == "processed_logits":
174
logits_to_return = logits
@@ -310,8 +310,9 @@ def apply_top_k_top_p(
310
if p is None and k is None:
311
return logits
312
313
- # Keep CPU logits on the PyTorch path to avoid invoking Triton kernels.
313
if current_platform.is_cpu():
314
+ if HAS_TRITON:
315
+ return apply_top_k_top_p_triton(logits, k, p)
316
return apply_top_k_top_p_pytorch(logits, k, p, allow_cpu_sync=True)
317
318
if HAS_TRITON and logits.shape[0] >= 8:
vllm/v1/sample/ops/topk_topp_triton.py
+8
-2
@@ -1038,6 +1038,12 @@ def apply_top_k_top_p_triton(
1038
else:
1039
normal_cdf_to_sigma_table, percentile_to_std_table = tables
1040
1041
+ # Smaller tiles compile and run faster on CPU; GPU benefits from larger tiles.
1042
+ if logits.device.type == "cpu":
1043
+ block_size, block_size_trunc = 256, 128
1044
+ else:
1045
+ block_size, block_size_trunc = 8192, 4096
1046
+
1047
_topk_topp_kernel[(NUM_PROGRAMS,)](
1048
logits,
1049
logits.stride(0),
@@ -1049,8 +1055,8 @@ def apply_top_k_top_p_triton(
1055
BATCH_SIZE=batch_size,
1056
MASK_VALUE=mask_value,
1057
VOCAB_SIZE=vocab_size,
1052
- BLOCK_SIZE=8192,
1053
- BLOCK_SIZE_TRUNC=4096,
1058
+ BLOCK_SIZE=block_size,
1059
+ BLOCK_SIZE_TRUNC=block_size_trunc,
1060
TOPK_ENABLED=topk_enabled,
1061
TOPP_ENABLED=topp_enabled,
1062
)