From 73393caeea25dde75f09a4079ee7895a87d408ea Mon Sep 17 00:00:00 2001 From: Tianmu Li Date: Mon, 7 Sep 2026 10:14:01 +0000 Subject: [PATCH 1/4] CPU: add vectorized top-k and top-p sampling kernels (#43452) (cherry picked from commit 9ffd571a1f351c36bc95bc0de0e6cd1ce50d21c4) Signed-off-by: Tianmu Li --- .../kernels/cpu/benchmark_cpu_topk_topp.py | 421 ++++++++++ cmake/cpu_extension.cmake | 6 +- csrc/cpu/cpu_arch_macros.h | 58 ++ csrc/cpu/topk_topp_sampler.cpp | 775 ++++++++++++++++++ csrc/cpu/torch_bindings.cpp | 19 +- tests/v1/sample/test_topk_topp_sampler.py | 525 +++++++++++- vllm/_custom_ops.py | 17 + vllm/v1/sample/ops/topk_topp_sampler.py | 61 +- 8 files changed, 1851 insertions(+), 31 deletions(-) create mode 100644 benchmarks/kernels/cpu/benchmark_cpu_topk_topp.py create mode 100644 csrc/cpu/topk_topp_sampler.cpp diff --git a/benchmarks/kernels/cpu/benchmark_cpu_topk_topp.py b/benchmarks/kernels/cpu/benchmark_cpu_topk_topp.py new file mode 100644 index 000000000000..f7dfc45e5f3b --- /dev/null +++ b/benchmarks/kernels/cpu/benchmark_cpu_topk_topp.py @@ -0,0 +1,421 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +""" +Benchmark comparing CPU C++ kernels vs PyTorch sort-based top-k/top-p. + +Compares: +- apply_top_k_top_p_cpu (C++ kernel — ISA determined at build time) +- apply_top_k_top_p_pytorch (PyTorch sort-based) + +Scenarios: +- top_k only (whole batch, partial batch) +- top_p only (whole batch, partial batch) +- mix of top_k and top_p +""" + +import argparse +import time +from dataclasses import dataclass + +import torch + +from vllm.v1.sample.ops.topk_topp_sampler import ( + _HAS_CPU_SAMPLING_OPS, + apply_top_k_top_p_cpu, + apply_top_k_top_p_pytorch, +) + + +def _pytorch_reference( + logits: torch.Tensor, + k: torch.Tensor | None, + p: torch.Tensor | None, +) -> torch.Tensor: + """Wrapper that enables CPU sync (required outside GPU context).""" + return apply_top_k_top_p_pytorch(logits, k, p, allow_cpu_sync=True) + + +@dataclass +class BenchmarkConfig: + """Configuration for a benchmark run.""" + + name: str + batch_size: int + vocab_size: int + # k and p can be tensors or None + k_values: torch.Tensor | None # [batch_size] or None + p_values: torch.Tensor | None # [batch_size] or None + description: str + ops_pct: float = 0.0 # Percentage of ops relative to batch size + + +def calculate_ops_pct( + k_values: torch.Tensor | None, + p_values: torch.Tensor | None, + vocab_size: int, + batch_size: int, +) -> float: + """ + Calculate the percentage of active top-k and top-p operations. + + Returns percentage where 100% = batch_size ops. + E.g., if all rows have both top-k and top-p active, returns 200%. + """ + active_ops = 0 + + if k_values is not None: + # Count rows where k < vocab_size (active top-k filtering) + active_ops += (k_values < vocab_size).sum().item() + + if p_values is not None: + # Count rows where p < 1.0 (active top-p filtering) + active_ops += (p_values < 1.0).sum().item() + + return (active_ops / batch_size) * 100 if batch_size > 0 else 0.0 + + +def create_logits( + batch_size: int, vocab_size: int, device: str = "cpu" +) -> torch.Tensor: + """Create random logits mimicking a realistic LLM distribution. + + Uses a Zipf-like probability distribution (rank^-1.1) converted to logits + via log, then randomly permuted per row. This produces a peaked distribution + where a small number of tokens capture most probability mass, similar to + real model outputs. + """ + # Create Zipf-like probabilities: p(rank) ~ rank^(-alpha) + ranks = torch.arange(1, vocab_size + 1, dtype=torch.float32, device=device) + probs = ranks.pow(-1.1) + probs = probs / probs.sum() + + # Convert to logits (log-probabilities, unnormalized is fine) + base_logits = probs.log() + + # Broadcast to batch and randomly permute each row + logits = base_logits.unsqueeze(0).expand(batch_size, -1).clone() + for i in range(batch_size): + logits[i] = logits[i, torch.randperm(vocab_size, device=device)] + + return logits + + +def benchmark_function( + func, + logits: torch.Tensor, + k: torch.Tensor | None, + p: torch.Tensor | None, + warmup_iters: int = 5, + benchmark_iters: int = 20, +) -> float: + """ + Benchmark a function and return avg_time_ms. + """ + # Warmup + for _ in range(warmup_iters): + logits_copy = logits.clone() + func(logits_copy, k, p) + + # Benchmark + times = [] + for _ in range(benchmark_iters): + logits_copy = logits.clone() + t0 = time.perf_counter_ns() + func(logits_copy, k, p) + times.append((time.perf_counter_ns() - t0) / 1e6) + + return sum(times) / len(times) + + +def create_benchmark_configs( + batch_sizes: list[int], + vocab_sizes: list[int], + device: str = "cpu", +) -> list[BenchmarkConfig]: + """Create all benchmark configurations.""" + configs = [] + + for vocab_size in vocab_sizes: + for batch_size in batch_sizes: + # 1. Top-k only - whole batch (all rows have k < vocab_size) + k_all = torch.full((batch_size,), 50, dtype=torch.int32, device=device) + configs.append( + BenchmarkConfig( + name=f"topk_whole_b{batch_size}_v{vocab_size // 1000}k", + batch_size=batch_size, + vocab_size=vocab_size, + k_values=k_all, + p_values=None, + description=f"Top-k only (whole batch, k=50), " + f"batch={batch_size}, vocab={vocab_size}", + ops_pct=calculate_ops_pct(k_all, None, vocab_size, batch_size), + ) + ) + + # 2. Top-k only - partial batch (half have k=50, half have k=vocab_size) + k_partial = torch.full((batch_size,), 50, dtype=torch.int32, device=device) + k_partial[batch_size // 2 :] = vocab_size # No filtering for second half + configs.append( + BenchmarkConfig( + name=f"topk_partial_b{batch_size}_v{vocab_size // 1000}k", + batch_size=batch_size, + vocab_size=vocab_size, + k_values=k_partial, + p_values=None, + description=f"Top-k only (partial batch, 50% k=50, 50% k=vocab), " + f"batch={batch_size}, vocab={vocab_size}", + ops_pct=calculate_ops_pct(k_partial, None, vocab_size, batch_size), + ) + ) + + # 3. Top-p only - whole batch (all rows have p < 1.0) + p_all = torch.full((batch_size,), 0.9, dtype=torch.float32, device=device) + configs.append( + BenchmarkConfig( + name=f"topp_whole_b{batch_size}_v{vocab_size // 1000}k", + batch_size=batch_size, + vocab_size=vocab_size, + k_values=None, + p_values=p_all, + description=f"Top-p only (whole batch, p=0.9), " + f"batch={batch_size}, vocab={vocab_size}", + ops_pct=calculate_ops_pct(None, p_all, vocab_size, batch_size), + ) + ) + + # 4. Top-p only - partial batch (half have p=0.9, half have p=1.0) + p_partial = torch.full( + (batch_size,), 0.9, dtype=torch.float32, device=device + ) + p_partial[batch_size // 2 :] = 1.0 # No filtering for second half + configs.append( + BenchmarkConfig( + name=f"topp_partial_b{batch_size}_v{vocab_size // 1000}k", + batch_size=batch_size, + vocab_size=vocab_size, + k_values=None, + p_values=p_partial, + description=f"Top-p only (partial batch, 50% p=0.9, 50% p=1.0), " + f"batch={batch_size}, vocab={vocab_size}", + ops_pct=calculate_ops_pct(None, p_partial, vocab_size, batch_size), + ) + ) + + # 5. Mix of top-k and top-p (both applied to whole batch) + k_mix = torch.full((batch_size,), 100, dtype=torch.int32, device=device) + p_mix = torch.full((batch_size,), 0.9, dtype=torch.float32, device=device) + configs.append( + BenchmarkConfig( + name=f"topk_topp_whole_b{batch_size}_v{vocab_size // 1000}k", + batch_size=batch_size, + vocab_size=vocab_size, + k_values=k_mix, + p_values=p_mix, + description=f"Top-k + Top-p (whole batch, k=100, p=0.9), " + f"batch={batch_size}, vocab={vocab_size}", + ops_pct=calculate_ops_pct(k_mix, p_mix, vocab_size, batch_size), + ) + ) + + # 6. Mix with partial application (some rows k only, some p only, some both) + k_mixed = torch.full( + (batch_size,), vocab_size, dtype=torch.int32, device=device + ) + p_mixed = torch.full((batch_size,), 1.0, dtype=torch.float32, device=device) + # First third: k only + third = batch_size // 3 + k_mixed[:third] = 50 + # Second third: p only + p_mixed[third : 2 * third] = 0.5 + # Last third: both k and p + k_mixed[2 * third :] = 100 + p_mixed[2 * third :] = 0.9 + configs.append( + BenchmarkConfig( + name=f"mixed_partial_b{batch_size}_v{vocab_size // 1000}k", + batch_size=batch_size, + vocab_size=vocab_size, + k_values=k_mixed, + p_values=p_mixed, + description=f"Mixed partial (1/3 k=50, 1/3 p=0.9, 1/3 both), " + f"batch={batch_size}, vocab={vocab_size}", + ops_pct=calculate_ops_pct(k_mixed, p_mixed, vocab_size, batch_size), + ) + ) + + return configs + + +def run_benchmark( + configs: list[BenchmarkConfig], + warmup_iters: int = 5, + benchmark_iters: int = 20, + verbose: bool = True, +): + """Run all benchmarks and print results.""" + results = [] + + print("=" * 100) + print("Top-k/Top-p Benchmark: CPU C++ vs PyTorch Sort-based") + print("=" * 100) + print() + + for config in configs: + if verbose: + print(f"Running: {config.description}") + + # Create fresh logits for this config + logits = create_logits(config.batch_size, config.vocab_size) + + # Benchmark CPU + cpu_time = benchmark_function( + apply_top_k_top_p_cpu, + logits, + config.k_values, + config.p_values, + warmup_iters, + benchmark_iters, + ) + + # Benchmark PyTorch + pytorch_time = benchmark_function( + _pytorch_reference, + logits, + config.k_values, + config.p_values, + warmup_iters, + benchmark_iters, + ) + + speedup = pytorch_time / cpu_time if cpu_time > 0 else float("inf") + + result = { + "config": config, + "cpu_time_ms": cpu_time, + "pytorch_time_ms": pytorch_time, + "speedup": speedup, + } + results.append(result) + + if verbose: + print(f" CPU: {cpu_time:.3f} ms") + print(f" PyTorch: {pytorch_time:.3f} ms") + print(f" Speedup: {speedup:.2f}x") + print() + + del logits + + return results + + +def print_summary_table(results: list[dict]): + """Print a summary table of results.""" + print() + print("=" * 100) + print("SUMMARY TABLE") + print("=" * 100) + print() + + # Header + header = ( + f"{'Scenario':<40} {'Batch':>6} {'Vocab':>7} {'Ops%':>6} " + f"{'CPU (ms)':>10} {'PyTorch (ms)':>13} {'Speedup':>8}" + ) + print(header) + print("-" * 100) + + # Group by scenario type + current_vocab = None + for result in results: + config = result["config"] + + # Add separator between vocab sizes + if current_vocab != config.vocab_size: + if current_vocab is not None: + print("-" * 100) + current_vocab = config.vocab_size + + scenario = config.name.split("_b")[0] # Extract scenario name + print( + f"{scenario:<40} {config.batch_size:>6} {config.vocab_size:>7} " + f"{config.ops_pct:>5.0f}% " + f"{result['cpu_time_ms']:>10.3f} {result['pytorch_time_ms']:>13.3f} " + f"{result['speedup']:>7.2f}x" + ) + + print("=" * 100) + + +def main(): + parser = argparse.ArgumentParser( + description="Benchmark CPU C++ vs PyTorch sort-based top-k/top-p" + ) + parser.add_argument( + "--batch-sizes", + type=int, + nargs="+", + default=[1, 4, 16, 64, 128, 512, 1024, 2048], + help="Batch sizes to test (default: 1 4 16 64 128 512 1024 2048)", + ) + parser.add_argument( + "--vocab-sizes", + type=int, + nargs="+", + default=[32768, 131072], # 32k, 128k + help="Vocabulary sizes to test (default: 32768 131072)", + ) + parser.add_argument( + "--warmup-iters", + type=int, + default=5, + help="Number of warmup iterations (default: 5)", + ) + parser.add_argument( + "--benchmark-iters", + type=int, + default=20, + help="Number of benchmark iterations (default: 20)", + ) + parser.add_argument( + "--quiet", + action="store_true", + help="Only print summary table", + ) + + args = parser.parse_args() + + if not _HAS_CPU_SAMPLING_OPS: + print( + "ERROR: CPU sampling ops not available. " + "Build vLLM with VLLM_TARGET_DEVICE=cpu." + ) + return + + # Print configuration + print(f"Batch sizes: {args.batch_sizes}") + print(f"Vocab sizes: {args.vocab_sizes}") + print(f"Warmup iterations: {args.warmup_iters}") + print(f"Benchmark iterations: {args.benchmark_iters}") + print() + + # Create configs + configs = create_benchmark_configs( + args.batch_sizes, + args.vocab_sizes, + ) + + # Run benchmarks + results = run_benchmark( + configs, + warmup_iters=args.warmup_iters, + benchmark_iters=args.benchmark_iters, + verbose=not args.quiet, + ) + + # Print summary + print_summary_table(results) + + +if __name__ == "__main__": + main() diff --git a/cmake/cpu_extension.cmake b/cmake/cpu_extension.cmake index d92ad58a9cf3..42ed9e3919c5 100644 --- a/cmake/cpu_extension.cmake +++ b/cmake/cpu_extension.cmake @@ -186,7 +186,8 @@ if (CMAKE_SYSTEM_PROCESSOR MATCHES "x86_64|amd64" OR ENABLE_X86_ISA) endif() list(APPEND CXX_COMPILE_FLAGS_AVX2 - "-mavx2") + "-mavx2" + "-mfma") elseif (POWER9_FOUND OR POWER10_FOUND OR POWER11_FOUND) message(STATUS "PowerPC detected") if (POWER9_FOUND) @@ -497,6 +498,7 @@ set(VLLM_EXT_SRC "csrc/cpu/mla_decode.cpp" "csrc/cpu/pos_encoding.cpp" "csrc/cpu/mamba_cpu.cpp" + "csrc/cpu/topk_topp_sampler.cpp" "csrc/cpu/dynamic_4bit_int_moe_cpu.cpp" "csrc/cpu/cpu_fused_moe.cpp" "csrc/cpu/cpu_attn.cpp" @@ -595,6 +597,7 @@ if (ENABLE_X86_ISA) "csrc/cpu/mamba_cpu.cpp" "csrc/cpu/sampling_kernels.cpp" "csrc/cpu/fp8_quant.cpp" + "csrc/cpu/topk_topp_sampler.cpp" "csrc/cpu/torch_bindings.cpp" # TODO: Remove these files "csrc/cpu/activation.cpp" @@ -614,6 +617,7 @@ if (ENABLE_X86_ISA) "csrc/cpu/dnnl_kernels.cpp" "csrc/cpu/sampling_kernels.cpp" "csrc/cpu/fp8_quant.cpp" + "csrc/cpu/topk_topp_sampler.cpp" "csrc/cpu/torch_bindings.cpp" # TODO: Remove these files "csrc/cpu/activation.cpp" diff --git a/csrc/cpu/cpu_arch_macros.h b/csrc/cpu/cpu_arch_macros.h index ccec4c2e87c7..4d7cd42f089d 100644 --- a/csrc/cpu/cpu_arch_macros.h +++ b/csrc/cpu/cpu_arch_macros.h @@ -58,6 +58,64 @@ }; #endif + #if !defined(__AVX512F__) && defined(__AVX2__) + #define DEFINE_FAST_EXP \ + const __m256 vec_factorial_1 = _mm256_set1_ps(0.999999701f); \ + const __m256 vec_factorial_2 = _mm256_set1_ps(0.499991506f); \ + const __m256 vec_factorial_3 = _mm256_set1_ps(0.166676521f); \ + const __m256 vec_factorial_4 = _mm256_set1_ps(0.0418978221f); \ + const __m256 vec_factorial_5 = _mm256_set1_ps(0.00828929059f); \ + const __m256 vec_exp_log2ef = \ + _mm256_castsi256_ps(_mm256_set1_epi32(0x3fb8aa3b)); \ + const __m256 vec_half = _mm256_set1_ps(0.5f); \ + const __m256 vec_one = _mm256_set1_ps(1.f); \ + const __m256 vec_zero = _mm256_set1_ps(0.f); \ + const __m256 vec_two = _mm256_set1_ps(2.f); \ + const __m256 vec_ln2f = \ + _mm256_castsi256_ps(_mm256_set1_epi32(0x3f317218)); \ + const __m256 vec_ln_flt_min = \ + _mm256_castsi256_ps(_mm256_set1_epi32(0xc2aeac50)); \ + const __m256 vec_ln_flt_max = \ + _mm256_castsi256_ps(_mm256_set1_epi32(0x42b17218)); \ + const __m256i vec_127 = _mm256_set1_epi32(0x0000007f); \ + const int n_mantissa_bits = 23; \ + auto avx2_expf = [&](__m256 values) __attribute__((always_inline)) { \ + auto less_ln_flt_min_mask = \ + _mm256_cmp_ps(values, vec_ln_flt_min, 1 /*_CMP_LT_OS*/); \ + auto vec_src = _mm256_min_ps(values, vec_ln_flt_max); \ + vec_src = _mm256_max_ps(vec_src, vec_ln_flt_min); \ + auto vec_fx = _mm256_fmadd_ps(vec_src, vec_exp_log2ef, vec_half); \ + /* AVX2 lacks cvt_roundps_epi32; use round_ps + cvtps_epi32 */ \ + auto vec_fx_rnd = _mm256_round_ps( \ + vec_fx, _MM_FROUND_TO_NEG_INF | _MM_FROUND_NO_EXC); \ + auto vec_fx_i = _mm256_cvtps_epi32(vec_fx_rnd); \ + vec_fx = _mm256_cvtepi32_ps(vec_fx_i); \ + auto vec_exp_poly = _mm256_fnmadd_ps(vec_fx, vec_ln2f, vec_src); \ + auto vec_res = \ + _mm256_fmadd_ps(vec_exp_poly, vec_factorial_5, vec_factorial_4); \ + vec_res = _mm256_fmadd_ps(vec_exp_poly, vec_res, vec_factorial_3); \ + vec_res = _mm256_fmadd_ps(vec_exp_poly, vec_res, vec_factorial_2); \ + vec_res = _mm256_fmadd_ps(vec_exp_poly, vec_res, vec_factorial_1); \ + vec_res = _mm256_fmadd_ps(vec_exp_poly, vec_res, vec_one); \ + auto vec_exp_number = _mm256_sub_ps(vec_fx, vec_one); \ + auto vec_exp_number_i = _mm256_cvtps_epi32(vec_exp_number); \ + auto vec_two_pow_n_i = _mm256_add_epi32(vec_exp_number_i, vec_127); \ + vec_two_pow_n_i = _mm256_slli_epi32(vec_two_pow_n_i, n_mantissa_bits); \ + auto vec_two_pow_n = _mm256_castsi256_ps(vec_two_pow_n_i); \ + /* _mm256_blendv_ps: mask sign-bit=1 selects b (zero), else a */ \ + vec_two_pow_n = \ + _mm256_blendv_ps(vec_two_pow_n, vec_zero, less_ln_flt_min_mask); \ + vec_res = _mm256_mul_ps(vec_res, vec_two_pow_n); \ + vec_res = _mm256_mul_ps(vec_res, vec_two); \ + return vec_res; \ + }; \ + auto fast_exp = [&](const vec_op::FP32Vec16& vec) \ + __attribute__((always_inline)) { \ + return vec_op::FP32Vec16(avx2_expf(vec.reg_low), \ + avx2_expf(vec.reg_high)); \ + }; + #endif + #endif #ifdef __aarch64__ diff --git a/csrc/cpu/topk_topp_sampler.cpp b/csrc/cpu/topk_topp_sampler.cpp new file mode 100644 index 000000000000..8a3baa2f8748 --- /dev/null +++ b/csrc/cpu/topk_topp_sampler.cpp @@ -0,0 +1,775 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project +#include "cpu_arch_macros.h" +#include "cpu_types.hpp" +#include +#include +#include + +namespace { +constexpr float kPadSentinel = -1e30f; // logits below this are PAD +constexpr float kBoundaryTol = 1e-5f; // tie tolerance for top-p cut +constexpr int kMeanStdSampleN = 8192; // mean/std sample cap +constexpr int kBsearchMaxIters = 32; // binary-search iteration cap +constexpr int kTernarySearchMaxIters = 18; // matches Triton +constexpr float kKDuplicateTol = 1e-9f; // tie tolerance for top-k (Triton) +} // namespace + +// Matches _NORMAL_CDF_TO_SIGMA_TABLE in topk_topp_triton.py (200 entries). +static const float NORMAL_CDF_TO_SIGMA_TABLE[200] = { + 3.656f, 3.650f, 3.650f, 3.650f, 3.626f, 3.626f, 3.626f, 3.514f, 3.514f, + 3.503f, 3.503f, 3.434f, 3.434f, 3.428f, 3.428f, 3.387f, 3.380f, 3.380f, + 3.376f, 3.373f, 3.373f, 3.356f, 3.354f, 3.354f, 3.291f, 3.249f, 3.234f, + 3.214f, 3.198f, 3.198f, 3.185f, 3.177f, 3.177f, 3.165f, 3.164f, 3.161f, + 3.138f, 3.120f, 3.115f, 3.113f, 3.093f, 3.066f, 3.054f, 3.043f, 3.037f, + 3.023f, 2.993f, 2.991f, 2.976f, 2.970f, 2.952f, 2.946f, 2.932f, 2.908f, + 2.902f, 2.895f, 2.886f, 2.874f, 2.861f, 2.844f, 2.836f, 2.810f, 2.801f, + 2.790f, 2.784f, 2.779f, 2.767f, 2.757f, 2.745f, 2.733f, 2.723f, 2.716f, + 2.693f, 2.678f, 2.671f, 2.656f, 2.649f, 2.629f, 2.611f, 2.595f, 2.592f, + 2.585f, 2.574f, 2.550f, 2.543f, 2.534f, 2.521f, 2.518f, 2.497f, 2.485f, + 2.468f, 2.450f, 2.441f, 2.430f, 2.412f, 2.402f, 2.389f, 2.383f, 2.377f, + 2.364f, 2.349f, 2.338f, 2.332f, 2.319f, 2.310f, 2.301f, 2.282f, 2.274f, + 2.266f, 2.250f, 2.242f, 2.236f, 2.226f, 2.215f, 2.207f, 2.196f, 2.179f, + 2.171f, 2.162f, 2.147f, 2.135f, 2.121f, 2.109f, 2.095f, 2.085f, 2.073f, + 2.063f, 2.045f, 2.030f, 2.016f, 2.003f, 1.992f, 1.983f, 1.972f, 1.960f, + 1.949f, 1.940f, 1.928f, 1.912f, 1.897f, 1.881f, 1.869f, 1.854f, 1.838f, + 1.824f, 1.807f, 1.792f, 1.779f, 1.764f, 1.751f, 1.739f, 1.726f, 1.711f, + 1.697f, 1.685f, 1.668f, 1.652f, 1.636f, 1.622f, 1.603f, 1.585f, 1.568f, + 1.551f, 1.534f, 1.513f, 1.499f, 1.480f, 1.464f, 1.441f, 1.422f, 1.394f, + 1.373f, 1.347f, 1.320f, 1.296f, 1.270f, 1.246f, 1.219f, 1.190f, 1.163f, + 1.135f, 1.104f, 1.073f, 1.041f, 1.006f, 0.969f, 0.931f, 0.894f, 0.851f, + 0.806f, 0.757f, 0.702f, 0.643f, 0.574f, 0.498f, 0.405f, 0.288f, 0.134f, + -0.110f, -3.813f, +}; + +// Matches _PERCENTILE_TO_STD_TABLE in topk_topp_triton.py (200 entries). +static const float PERCENTILE_TO_STD_TABLE[200] = { + 2.576f, 2.319f, 2.178f, 2.064f, 1.968f, 1.892f, 1.819f, 1.757f, + 1.708f, 1.659f, 1.616f, 1.568f, 1.526f, 1.492f, 1.456f, 1.420f, + 1.382f, 1.342f, 1.309f, 1.280f, 1.249f, 1.221f, 1.193f, 1.169f, + 1.145f, 1.121f, 1.095f, 1.073f, 1.050f, 1.030f, 1.008f, 0.987f, + 0.966f, 0.945f, 0.926f, 0.910f, 0.891f, 0.871f, 0.854f, 0.837f, + 0.819f, 0.803f, 0.784f, 0.767f, 0.753f, 0.734f, 0.719f, 0.702f, + 0.690f, 0.675f, 0.658f, 0.640f, 0.625f, 0.609f, 0.595f, 0.578f, + 0.564f, 0.550f, 0.537f, 0.521f, 0.509f, 0.495f, 0.481f, 0.466f, + 0.453f, 0.439f, 0.424f, 0.410f, 0.397f, 0.383f, 0.370f, 0.356f, + 0.343f, 0.330f, 0.316f, 0.302f, 0.289f, 0.274f, 0.261f, 0.247f, + 0.235f, 0.223f, 0.209f, 0.196f, 0.184f, 0.172f, 0.159f, 0.149f, + 0.137f, 0.124f, 0.112f, 0.100f, 0.086f, 0.074f, 0.062f, 0.050f, + 0.035f, 0.023f, 0.009f, -0.003f, -0.015f, -0.027f, -0.039f, -0.052f, + -0.063f, -0.074f, -0.085f, -0.097f, -0.109f, -0.122f, -0.134f, -0.147f, + -0.158f, -0.171f, -0.184f, -0.196f, -0.210f, -0.223f, -0.235f, -0.248f, + -0.261f, -0.275f, -0.289f, -0.302f, -0.317f, -0.328f, -0.341f, -0.353f, + -0.368f, -0.382f, -0.396f, -0.410f, -0.426f, -0.439f, -0.452f, -0.465f, + -0.480f, -0.493f, -0.507f, -0.521f, -0.537f, -0.551f, -0.568f, -0.582f, + -0.597f, -0.614f, -0.628f, -0.643f, -0.658f, -0.673f, -0.691f, -0.706f, + -0.721f, -0.738f, -0.754f, -0.769f, -0.789f, -0.808f, -0.824f, -0.838f, + -0.857f, -0.877f, -0.893f, -0.912f, -0.929f, -0.947f, -0.965f, -0.983f, + -1.003f, -1.027f, -1.050f, -1.070f, -1.092f, -1.117f, -1.139f, -1.162f, + -1.189f, -1.216f, -1.241f, -1.272f, -1.300f, -1.330f, -1.367f, -1.404f, + -1.441f, -1.485f, -1.523f, -1.564f, -1.607f, -1.658f, -1.710f, -1.778f, + -1.832f, -1.901f, -1.978f, -2.068f, -2.174f, -2.325f, -2.577f, -3.813f, +}; + +namespace { + +// FP64 accumulator for softmax sum over FP32Vec16 chunks. +// FP32 accumulation over 128K elements loses ~3 ULP; FP64 avoids this. +#if defined(__AVX512F__) +struct Fp64Acc16 { + __m512d acc0 = _mm512_setzero_pd(); + __m512d acc1 = _mm512_setzero_pd(); + inline void add(const vec_op::FP32Vec16& v) { + __m256 lo8 = _mm512_castps512_ps256(v.reg); + __m256 hi8 = _mm512_extractf32x8_ps(v.reg, 1); + acc0 = _mm512_add_pd(acc0, _mm512_cvtps_pd(lo8)); + acc1 = _mm512_add_pd(acc1, _mm512_cvtps_pd(hi8)); + } + inline double reduce() const { + return _mm512_reduce_add_pd(acc0) + _mm512_reduce_add_pd(acc1); + } +}; +#elif defined(__AVX2__) +struct Fp64Acc16 { + __m256d acc0 = _mm256_setzero_pd(); + __m256d acc1 = _mm256_setzero_pd(); + __m256d acc2 = _mm256_setzero_pd(); + __m256d acc3 = _mm256_setzero_pd(); + inline void add(const vec_op::FP32Vec16& v) { + __m128 e_ll = _mm256_extractf128_ps(v.reg_low, 0); + __m128 e_lh = _mm256_extractf128_ps(v.reg_low, 1); + __m128 e_hl = _mm256_extractf128_ps(v.reg_high, 0); + __m128 e_hh = _mm256_extractf128_ps(v.reg_high, 1); + acc0 = _mm256_add_pd(acc0, _mm256_cvtps_pd(e_ll)); + acc1 = _mm256_add_pd(acc1, _mm256_cvtps_pd(e_lh)); + acc2 = _mm256_add_pd(acc2, _mm256_cvtps_pd(e_hl)); + acc3 = _mm256_add_pd(acc3, _mm256_cvtps_pd(e_hh)); + } + inline double reduce() const { + __m256d sum_a = + _mm256_add_pd(_mm256_add_pd(acc0, acc1), _mm256_add_pd(acc2, acc3)); + __m128d hi128 = _mm256_extractf128_pd(sum_a, 1); + __m128d lo128 = _mm256_castpd256_pd128(sum_a); + __m128d s2 = _mm_add_pd(lo128, hi128); + return _mm_cvtsd_f64(_mm_hadd_pd(s2, s2)); + } +}; +#endif + +__attribute__((always_inline)) static inline double sum_gt_to_double( + const float* buf, int n, float threshold) { + double s = 0.0; + int j = 0; +#if defined(__AVX512F__) + __m512d acc0 = _mm512_setzero_pd(), acc1 = _mm512_setzero_pd(); + const __m512 thr16 = _mm512_set1_ps(threshold); + for (; j + 16 <= n; j += 16) { + __m512 b16 = _mm512_loadu_ps(buf + j); + __mmask16 mask = _mm512_cmp_ps_mask(b16, thr16, _CMP_GT_OQ); + __m256 lo8 = _mm512_castps512_ps256(b16); + __m256 hi8 = _mm512_extractf32x8_ps(b16, 1); + acc0 = _mm512_add_pd(acc0, + _mm512_maskz_cvtps_pd((__mmask8)(mask & 0xFF), lo8)); + acc1 = + _mm512_add_pd(acc1, _mm512_maskz_cvtps_pd((__mmask8)(mask >> 8), hi8)); + } + s = _mm512_reduce_add_pd(acc0) + _mm512_reduce_add_pd(acc1); +#elif defined(__AVX2__) + // and_ps zeros excluded lanes; blendv with -inf would corrupt the sum. + const __m256 thr = _mm256_set1_ps(threshold); + __m256d acc0 = _mm256_setzero_pd(), acc1 = _mm256_setzero_pd(); + __m256d acc2 = _mm256_setzero_pd(), acc3 = _mm256_setzero_pd(); + for (; j + 16 <= n; j += 16) { + __m256 b_lo = _mm256_loadu_ps(buf + j); + __m256 b_hi = _mm256_loadu_ps(buf + j + 8); + __m256 masked_lo = + _mm256_and_ps(b_lo, _mm256_cmp_ps(b_lo, thr, _CMP_GT_OQ)); + __m256 masked_hi = + _mm256_and_ps(b_hi, _mm256_cmp_ps(b_hi, thr, _CMP_GT_OQ)); + acc0 = _mm256_add_pd(acc0, + _mm256_cvtps_pd(_mm256_extractf128_ps(masked_lo, 0))); + acc1 = _mm256_add_pd(acc1, + _mm256_cvtps_pd(_mm256_extractf128_ps(masked_lo, 1))); + acc2 = _mm256_add_pd(acc2, + _mm256_cvtps_pd(_mm256_extractf128_ps(masked_hi, 0))); + acc3 = _mm256_add_pd(acc3, + _mm256_cvtps_pd(_mm256_extractf128_ps(masked_hi, 1))); + } + __m256d sum_a = + _mm256_add_pd(_mm256_add_pd(acc0, acc1), _mm256_add_pd(acc2, acc3)); + __m128d hi128 = _mm256_extractf128_pd(sum_a, 1); + __m128d lo128 = _mm256_castpd256_pd128(sum_a); + __m128d s2 = _mm_add_pd(lo128, hi128); + s = _mm_cvtsd_f64(_mm_hadd_pd(s2, s2)); +#endif + for (; j < n; ++j) + if (buf[j] > threshold) s += (double)buf[j]; + return s; +} + +__attribute__((always_inline)) static inline int count_within_tol( + const float* row, int V, float center, float tol) { + int n = 0; + int i = 0; +#if defined(__AVX512F__) + const vec_op::FP32Vec16 c16(center); + const vec_op::FP32Vec16 t16(tol); + for (; i + 16 <= V; i += 16) { + vec_op::FP32Vec16 absdiff = (vec_op::FP32Vec16(row + i) - c16).abs(); + __mmask16 mask = _mm512_cmp_ps_mask(absdiff.reg, t16.reg, _CMP_LT_OQ); + n += __builtin_popcount((unsigned)mask); + } +#elif defined(__AVX2__) + const vec_op::FP32Vec16 c16(center); + const vec_op::FP32Vec16 t16(tol); + for (; i + 16 <= V; i += 16) { + vec_op::FP32Vec16 absdiff = (vec_op::FP32Vec16(row + i) - c16).abs(); + __m256 clo = _mm256_cmp_ps(absdiff.reg_low, t16.reg_low, _CMP_LT_OQ); + __m256 chi = _mm256_cmp_ps(absdiff.reg_high, t16.reg_high, _CMP_LT_OQ); + n += __builtin_popcount((unsigned)_mm256_movemask_ps(clo)); + n += __builtin_popcount((unsigned)_mm256_movemask_ps(chi)); + } +#endif + for (; i < V; ++i) + if (fabsf(row[i] - center) < tol) ++n; + return n; +} + +__attribute__((always_inline)) static inline void mask_write_below( + float* row, int V, float threshold, float fill) { + int i = 0; +#if defined(__AVX512F__) + const vec_op::FP32Vec16 thr16(threshold); + const vec_op::FP32Vec16 f16(fill); + for (; i + 16 <= V; i += 16) { + vec_op::FP32Vec16 v16(row + i); + __mmask16 keep = _mm512_cmp_ps_mask(v16.reg, thr16.reg, _CMP_GT_OQ); + vec_op::FP32Vec16(_mm512_mask_blend_ps(keep, f16.reg, v16.reg)) + .save(row + i); + } +#elif defined(__AVX2__) + const vec_op::FP32Vec16 thr16(threshold); + const vec_op::FP32Vec16 f16(fill); + for (; i + 16 <= V; i += 16) { + vec_op::FP32Vec16 v16(row + i); + __m256 klo = _mm256_cmp_ps(v16.reg_low, thr16.reg_low, _CMP_GT_OQ); + __m256 khi = _mm256_cmp_ps(v16.reg_high, thr16.reg_high, _CMP_GT_OQ); + vec_op::FP32Vec16(_mm256_blendv_ps(f16.reg_low, v16.reg_low, klo), + _mm256_blendv_ps(f16.reg_high, v16.reg_high, khi)) + .save(row + i); + } +#endif + for (; i < V; ++i) + if (!(row[i] > threshold)) row[i] = fill; +} + +// Vectorised max/min/finite-count over row[0..V). PAD lanes (<= pad_sentinel) +// are blended to +inf so they don't corrupt reduce_min. Returns SIMD tail +// start in *tail_i_out; caller handles remaining elements and clamps min<=max. +__attribute__((always_inline)) static inline void vec_max_min_with_pad_blend( + const float* row, int V, float pad_sentinel, float* max_out, float* min_out, + int* n_finite_out, int* tail_i_out) { + int i = 0; + float max_l = -1e38f, min_l = 1e38f; + int n_finite = 0; +#if defined(__AVX512F__) || defined(__AVX2__) + { + const float pos_inf_val = std::numeric_limits::infinity(); + vec_op::FP32Vec16 maxv(row[0]); + vec_op::FP32Vec16 minv(pos_inf_val); + const vec_op::FP32Vec16 sentinel16(pad_sentinel); + const vec_op::FP32Vec16 pos_inf16(pos_inf_val); + for (; i + 16 <= V; i += 16) { + vec_op::FP32Vec16 v16(row + i); + maxv = maxv.max(v16); + #if defined(__AVX512F__) + __mmask16 is_pad = + _mm512_cmp_ps_mask(v16.reg, sentinel16.reg, _CMP_LE_OS); + vec_op::FP32Vec16 safe( + _mm512_mask_blend_ps(is_pad, v16.reg, pos_inf16.reg)); + __mmask16 fin = _mm512_cmp_ps_mask(v16.reg, sentinel16.reg, _CMP_GT_OQ); + n_finite += __builtin_popcount((unsigned)fin); + #else + __m256 lt_lo = _mm256_cmp_ps(v16.reg_low, sentinel16.reg_low, _CMP_LE_OS); + __m256 lt_hi = + _mm256_cmp_ps(v16.reg_high, sentinel16.reg_high, _CMP_LE_OS); + vec_op::FP32Vec16 safe( + _mm256_blendv_ps(v16.reg_low, pos_inf16.reg_low, lt_lo), + _mm256_blendv_ps(v16.reg_high, pos_inf16.reg_high, lt_hi)); + __m256 fin_lo = + _mm256_cmp_ps(v16.reg_low, sentinel16.reg_low, _CMP_GT_OQ); + __m256 fin_hi = + _mm256_cmp_ps(v16.reg_high, sentinel16.reg_high, _CMP_GT_OQ); + n_finite += __builtin_popcount((unsigned)_mm256_movemask_ps(fin_lo)); + n_finite += __builtin_popcount((unsigned)_mm256_movemask_ps(fin_hi)); + #endif + minv = minv.min(safe); + } + max_l = maxv.reduce_max(); + min_l = minv.reduce_min(); + } +#endif + *max_out = max_l; + *min_out = min_l; + *n_finite_out = n_finite; + *tail_i_out = i; +} + +__attribute__((always_inline)) static inline int gather_outliers( + const float* buf, int n, float threshold, float* out, double* sum_out) { + int m = 0; + double s = 0.0; + for (int j = 0; j < n; ++j) { + float v = buf[j]; + if (v > threshold) { + out[m++] = v; + if (sum_out) s += (double)v; + } + } + if (sum_out) *sum_out = s; + return m; +} + +static inline void compute_mean_std(const float* row, int V, float* avg_out, + float* std_out) { + float sum_s = 0.f, sum_sq = 0.f; + int nf = 0; + for (int i = 0; i < V && nf < kMeanStdSampleN; ++i) { + float v = row[i]; + if (v > kPadSentinel) { + sum_s += v; + sum_sq += v * v; + ++nf; + } + } + float avg = nf > 0 ? sum_s / nf : 0.f; + // E[X²] - (E[X])² is numerically unstable for large logit magnitudes, but + // this is only a heuristic pivot initializer; the ternary/binary search + // converges correctly regardless of pivot quality. + float var = nf > 0 ? sum_sq / nf - avg * avg : 1.f; + *avg_out = avg; + *std_out = sqrtf(var > 0.f ? var : 0.f); +} + +// Apply boundary-tie resolution in row order. Keeps all v > pivot, plus up to +// n_keep_at_boundary values within tol of dup_value. Rare path: called only +// when binary/ternary search lands exactly on a duplicate boundary value. +static inline void apply_boundary_tie_loop(float* row, int V, float pivot, + float dup_value, float tol, + int n_keep_at_boundary) { + const float neg_inf = -std::numeric_limits::infinity(); + int kept_boundary = 0; + for (int i = 0; i < V; ++i) { + float v = row[i]; + bool keep = v > pivot; + if (!keep && fabsf(v - dup_value) < tol && + kept_boundary < n_keep_at_boundary) { + keep = true; + ++kept_boundary; + } + if (!keep) row[i] = neg_inf; + } +} + +} // namespace + +// Pre-computed per-row statistics shared between top-k and top-p passes. +struct RowStats { + float avg, std_v, max_l, min_l; + int n_finite; +}; + +// Compute RowStats for row[0..V). min_l is clamped to max_l if no finite min +// was found (degenerate all-PAD row). +static inline void compute_row_stats(const float* row, int V, RowStats* st) { + compute_mean_std(row, V, &st->avg, &st->std_v); + int tail_i; + vec_max_min_with_pad_blend(row, V, kPadSentinel, &st->max_l, &st->min_l, + &st->n_finite, &tail_i); + for (int i = tail_i; i < V; ++i) { + float v = row[i]; + if (v > st->max_l) st->max_l = v; + if (v > kPadSentinel) { + ++st->n_finite; + if (v < st->min_l) st->min_l = v; + } + } + if (st->min_l > st->max_l) st->min_l = st->max_l; +} + +// Entry-point validation helpers shared by all three public functions. +static void check_logits(const torch::Tensor& logits) { + TORCH_CHECK(logits.dim() == 2, "logits must be 2D"); + TORCH_CHECK(logits.is_cpu(), "logits must be CPU"); + TORCH_CHECK(logits.dtype() == torch::kFloat32, "logits must be float32"); + TORCH_CHECK(logits.is_contiguous(), "logits must be contiguous"); +} + +static void check_batch_float(const torch::Tensor& t, + const torch::Tensor& logits, const char* name) { + TORCH_CHECK(t.is_cpu(), name, " must be CPU"); + TORCH_CHECK(t.dtype() == torch::kFloat32, name, " must be float32"); + TORCH_CHECK(t.dim() == 1 && t.size(0) == logits.size(0), name, + " must be 1D with size == batch size"); +} + +static void check_batch_int(const torch::Tensor& t, const torch::Tensor& logits, + const char* name) { + TORCH_CHECK(t.is_cpu(), name, " must be CPU"); + TORCH_CHECK(t.dtype() == torch::kInt32, name, " must be int32"); + TORCH_CHECK(t.dim() == 1 && t.size(0) == logits.size(0), name, + " must be 1D with size == batch size"); +} + +static float binary_search_buffer(const float* buf, int n, float p_val, + float lo, float hi, double* sum_above_out) { + double s_hi = 0.0; + for (int iter = 0; iter < kBsearchMaxIters; ++iter) { + float mid = lo + (hi - lo) * 0.5f; + if (mid == lo || mid == hi) break; + double s = sum_gt_to_double(buf, n, mid); + if (s >= (double)p_val) { + lo = mid; + } else { + hi = mid; + s_hi = s; + } + } + *sum_above_out = s_hi; + return hi; +} + +static void scan_pivot_stats(const float* buf, int n, float pivot, float tol, + int* num_above, float* min_above, + int* num_at_min) { + float ma = std::numeric_limits::infinity(); + int na = 0; + int i = 0; +#if defined(__AVX512F__) || defined(__AVX2__) + { + const float pos_inf_val = std::numeric_limits::infinity(); + const vec_op::FP32Vec16 vpivot16(pivot); + const vec_op::FP32Vec16 pos_inf16(pos_inf_val); + vec_op::FP32Vec16 vmin16(pos_inf_val); + for (; i + 16 <= n; i += 16) { + vec_op::FP32Vec16 b16(buf + i); + #if defined(__AVX512F__) + __mmask16 above = _mm512_cmp_ps_mask(b16.reg, vpivot16.reg, _CMP_GT_OQ); + na += __builtin_popcount((unsigned)above); + vmin16 = vmin16.min(vec_op::FP32Vec16( + _mm512_mask_blend_ps(above, pos_inf16.reg, b16.reg))); + #else + __m256 above_lo = + _mm256_cmp_ps(b16.reg_low, vpivot16.reg_low, _CMP_GT_OQ); + __m256 above_hi = + _mm256_cmp_ps(b16.reg_high, vpivot16.reg_high, _CMP_GT_OQ); + na += __builtin_popcount((unsigned)_mm256_movemask_ps(above_lo)); + na += __builtin_popcount((unsigned)_mm256_movemask_ps(above_hi)); + vmin16 = vmin16.min(vec_op::FP32Vec16( + _mm256_blendv_ps(pos_inf16.reg_low, b16.reg_low, above_lo), + _mm256_blendv_ps(pos_inf16.reg_high, b16.reg_high, above_hi))); + #endif + } + ma = vmin16.reduce_min(); + } +#endif + for (; i < n; ++i) { + if (buf[i] > pivot) { + ++na; + if (buf[i] < ma) ma = buf[i]; + } + } + *num_above = na; + *min_above = ma; + // count_within_tol runs over the full buf; elements below pivot are all + // < min_above, so only ULP-equal values fall within kKDuplicateTol and + // are legitimately counted as boundary duplicates. + *num_at_min = count_within_tol(buf, n, ma, tol); +} + +static float ternary_search_topk(const float* buf, int n, int k, float lo, + float hi, float* dup_value_out, + int* num_dup_out, int* num_above_out) { + float final_pivot = (lo + hi) * 0.5f; + float dup_value = std::numeric_limits::infinity(); + int num_dup = 0, num_above = 0; + for (int iter = 0; iter < kTernarySearchMaxIters; ++iter) { + float p0 = lo + (hi - lo) * (1.0f / 3.0f); + float p1 = lo + (hi - lo) * (2.0f / 3.0f); + int na0, na1, nd0, nd1; + float mn0, mn1; + scan_pivot_stats(buf, n, p0, kKDuplicateTol, &na0, &mn0, &nd0); + scan_pivot_stats(buf, n, p1, kKDuplicateTol, &na1, &mn1, &nd1); + + bool found0 = (na0 >= k) && (na0 - nd0 < k); + bool found1 = (na1 >= k) && (na1 - nd1 < k); + if (found1) { + final_pivot = p1; + dup_value = mn1; + num_dup = nd1; + num_above = na1; + break; + } + if (found0) { + final_pivot = p0; + dup_value = mn0; + num_dup = nd0; + num_above = na0; + break; + } + if (na1 > k) + lo = p1; + else if (na0 > k) + lo = p0; + if (na0 < k) + hi = p0; + else if (na1 < k) + hi = p1; + + if (hi - lo < kKDuplicateTol) { + final_pivot = (lo + hi) * 0.5f; + scan_pivot_stats(buf, n, final_pivot, kKDuplicateTol, &num_above, + &dup_value, &num_dup); + break; + } + } + *dup_value_out = dup_value; + *num_dup_out = num_dup; + *num_above_out = num_above; + return final_pivot; +} + +// Core top-k logic. Receives pre-computed RowStats; does not re-scan the row +// for mean/std or max/min. +static void _top_k_row_core(float* __restrict__ row, int V, int k_val, + const RowStats& st, + float* __restrict__ outlier_buf) { + int tidx = (int)((double)k_val / V * 200); + if (tidx < 0) tidx = 0; + if (tidx > 199) tidx = 199; + float sigma = PERCENTILE_TO_STD_TABLE[tidx]; + sigma = sigma + fabsf(sigma) * -0.15f; + float outlier_logit = st.avg + st.std_v * sigma; + + int n_outliers = gather_outliers(row, V, outlier_logit, outlier_buf, nullptr); + + if (st.n_finite <= k_val) return; + + const float neg_inf = -std::numeric_limits::infinity(); + + // Degenerate-flat: ternary search is unstable when all values are equal. + if (st.max_l - st.min_l < kKDuplicateTol) { + int kept = 0; + for (int i = 0; i < V; ++i) { + if (row[i] > kPadSentinel) { + if (kept < k_val) + ++kept; + else + row[i] = neg_inf; + } + } + return; + } + + float final_pivot, dup_value; + int n_dup, n_above; + if (n_outliers > k_val) { + final_pivot = + ternary_search_topk(outlier_buf, n_outliers, k_val, outlier_logit, + st.max_l, &dup_value, &n_dup, &n_above); + } else { + final_pivot = ternary_search_topk(row, V, k_val, st.min_l, st.max_l, + &dup_value, &n_dup, &n_above); + } + + int n_keep_at_boundary = n_dup - (n_above - k_val); + if (n_keep_at_boundary < 0) n_keep_at_boundary = 0; + if (n_keep_at_boundary > n_dup) n_keep_at_boundary = n_dup; + + if (n_keep_at_boundary == 0) { + mask_write_below(row, V, final_pivot, neg_inf); + } else { + apply_boundary_tie_loop(row, V, dup_value, dup_value, kKDuplicateTol, + n_keep_at_boundary); + } +} + +// Core top-p logic. Receives pre-computed RowStats; does not re-scan the row. +static void _top_p_row_core(float* __restrict__ row, int V, float p_val, + const RowStats& st, float* __restrict__ scratch, + float* __restrict__ outlier_buf) { + if (p_val >= 1.0f) return; + + const float neg_inf = -std::numeric_limits::infinity(); + if (p_val <= 0.0f) { + mask_write_below(row, V, st.max_l - kBoundaryTol, neg_inf); + return; + } + + int tidx = (int)(p_val * 200); + if (tidx < 0) tidx = 0; + if (tidx > 199) tidx = 199; + float sigma = NORMAL_CDF_TO_SIGMA_TABLE[tidx]; + sigma = sigma + fabsf(sigma) * -0.25f; + float outlier_logit = st.avg + st.std_v * sigma; + +#if defined(__AVX512F__) || defined(__AVX2__) + DEFINE_FAST_EXP + vec_op::FP32Vec16 base16(st.max_l); +#endif + double sum_exp_d = 0.0; + { + int i = 0; +#if defined(__AVX512F__) || defined(__AVX2__) + Fp64Acc16 acc; + for (; i + 16 <= V; i += 16) { + vec_op::FP32Vec16 v16(row + i); + vec_op::FP32Vec16 e16 = fast_exp(v16 - base16); + acc.add(e16); + } + sum_exp_d += acc.reduce(); +#endif + for (; i < V; ++i) sum_exp_d += (double)expf(row[i] - st.max_l); + } + float sum_exp = (float)sum_exp_d; + + float outlier_prob = expf(outlier_logit - st.max_l) / sum_exp; + float max_prob = 1.0f / sum_exp; + float min_prob = expf(st.min_l - st.max_l) / sum_exp; + + double sum_out_d = 0.0; + { + const float inv_sum = 1.0f / sum_exp; + int i = 0; +#if defined(__AVX512F__) || defined(__AVX2__) + vec_op::FP32Vec16 isum16(inv_sum); + for (; i + 16 <= V; i += 16) { + vec_op::FP32Vec16 v16(row + i); + vec_op::FP32Vec16 p16 = fast_exp(v16 - base16) * isum16; + p16.save(scratch + i); + } +#endif + for (; i < V; ++i) scratch[i] = expf(row[i] - st.max_l) * inv_sum; + } + int n_out = + gather_outliers(scratch, V, outlier_prob, outlier_buf, &sum_out_d); + + float* sbuf; + int sn; + float buf_lo, buf_hi; + if (sum_out_d >= (double)p_val) { + sbuf = outlier_buf; + sn = n_out; + buf_hi = max_prob; + buf_lo = outlier_prob; + } else { + sbuf = scratch; + sn = V; + buf_hi = max_prob; + buf_lo = min_prob; + } + + double sum_above; + float boundary_prob = + binary_search_buffer(sbuf, sn, p_val, buf_lo, buf_hi, &sum_above); + + float dup_logit = logf(boundary_prob * sum_exp) + st.max_l; + + // NaN or overflow above max_l means binary search degenerated; skip masking. + if (!(dup_logit <= st.max_l)) { + return; + } + + int n_at_boundary = count_within_tol(row, V, dup_logit, kBoundaryTol); + + int n_keep_at_boundary = 0; + if (n_at_boundary > 0) { + double remaining = (double)p_val - sum_above; + if (remaining > 0.0 && (double)boundary_prob > 0.0) { + double needed_d = ceil(remaining / (double)boundary_prob); + if (needed_d >= (double)n_at_boundary) + n_keep_at_boundary = n_at_boundary; + else if (needed_d > 0.0) + n_keep_at_boundary = (int)needed_d; + } + } + + if (n_keep_at_boundary == 0) { + mask_write_below(row, V, dup_logit, neg_inf); + } else { + apply_boundary_tie_loop(row, V, dup_logit, dup_logit, kBoundaryTol, + n_keep_at_boundary); + } +} + +static void top_k_row(float* __restrict__ row, int V, int k_val, + float* __restrict__ outlier_buf) { + RowStats st; + compute_row_stats(row, V, &st); + _top_k_row_core(row, V, k_val, st, outlier_buf); +} + +static void top_p_row(float* __restrict__ row, int V, float p_val, + float* __restrict__ scratch, + float* __restrict__ outlier_buf) { + RowStats st; + compute_row_stats(row, V, &st); + _top_p_row_core(row, V, p_val, st, scratch, outlier_buf); +} + +// Fused top-k + top-p: shared mean/std and max/min passes save ~2 row reads +// vs sequential top_k_row + top_p_row. +static void top_k_p_row(float* __restrict__ row, int V, int k_val, float p_val, + float* __restrict__ scratch, + float* __restrict__ outlier_buf) { + RowStats st; + compute_row_stats(row, V, &st); + // k_val > 0 && k_val < V is pre-checked by the caller. + _top_k_row_core(row, V, k_val, st, outlier_buf); + // Top-k filtered lanes hold neg_inf; exp(-inf)=0 so they don't affect sum. + _top_p_row_core(row, V, p_val, st, scratch, outlier_buf); +} + +void cpu_topp_sampling(torch::Tensor& logits, const torch::Tensor& p) { + check_logits(logits); + check_batch_float(p, logits, "p"); + + const int64_t B = logits.size(0); + const int64_t V = logits.size(1); + const int V_int = static_cast(V); + if (V == 0) return; + float* lp = logits.data_ptr(); + const float* pp = p.data_ptr(); + + // thread_local avoids reallocation across calls; resize skips zero-fill. +#pragma omp parallel + { + thread_local std::vector scratch_tls; + thread_local std::vector outlier_tls; + if ((int64_t)scratch_tls.size() < V) scratch_tls.resize(V); + if ((int64_t)outlier_tls.size() < V) outlier_tls.resize(V); +#pragma omp for schedule(dynamic, 1) + for (int64_t b = 0; b < B; ++b) { + top_p_row(lp + b * V, V_int, pp[b], scratch_tls.data(), + outlier_tls.data()); + } + } +} + +void cpu_topk_sampling(torch::Tensor& logits, const torch::Tensor& k) { + check_logits(logits); + check_batch_int(k, logits, "k"); + + const int64_t B = logits.size(0); + const int64_t V = logits.size(1); + const int V_int = static_cast(V); + if (V == 0) return; + float* lp = logits.data_ptr(); + const int* kp = k.data_ptr(); + +#pragma omp parallel + { + thread_local std::vector outlier_tls; + if ((int64_t)outlier_tls.size() < V) outlier_tls.resize(V); +#pragma omp for schedule(dynamic, 1) + for (int64_t b = 0; b < B; ++b) { + int kv = kp[b]; + if (kv > 0 && kv < V_int) + top_k_row(lp + b * V, V_int, kv, outlier_tls.data()); + } + } +} + +void cpu_topk_topp_sampling(torch::Tensor& logits, const torch::Tensor& k, + const torch::Tensor& p) { + check_logits(logits); + check_batch_int(k, logits, "k"); + check_batch_float(p, logits, "p"); + + const int64_t B = logits.size(0); + const int64_t V = logits.size(1); + const int V_int = static_cast(V); + if (V == 0) return; + float* lp = logits.data_ptr(); + const int* kp = k.data_ptr(); + const float* pp = p.data_ptr(); + +#pragma omp parallel + { + thread_local std::vector scratch_tls; + thread_local std::vector outlier_tls; + if ((int64_t)scratch_tls.size() < V) scratch_tls.resize(V); + if ((int64_t)outlier_tls.size() < V) outlier_tls.resize(V); +#pragma omp for schedule(dynamic, 1) + for (int64_t b = 0; b < B; ++b) { + float* row = lp + b * V; + int kv = kp[b]; + float pv = pp[b]; + if (kv > 0 && kv < V_int) { + top_k_p_row(row, V_int, kv, pv, scratch_tls.data(), outlier_tls.data()); + } else { + top_p_row(row, V_int, pv, scratch_tls.data(), outlier_tls.data()); + } + } + } +} diff --git a/csrc/cpu/torch_bindings.cpp b/csrc/cpu/torch_bindings.cpp index c68d5c6317df..d74260b849c2 100644 --- a/csrc/cpu/torch_bindings.cpp +++ b/csrc/cpu/torch_bindings.cpp @@ -4,8 +4,9 @@ #include -// Note: overwrite the external definition for sharing same name between -// libraries use different ISAs. +// ISA-specific libs (_C_AVX512, _C_AVX2) must register under the same _C +// namespace so torch.ops._C.* dispatch works uniformly from Python. +#undef TORCH_EXTENSION_NAME #define TORCH_EXTENSION_NAME _C void release_dnnl_matmul_handler(int64_t handler); @@ -408,6 +409,11 @@ void dynamic_per_token_scaled_fp8_quant( void init_cpu_memory_env(std::vector node_ids); +void cpu_topp_sampling(torch::Tensor& logits, const torch::Tensor& p); +void cpu_topk_sampling(torch::Tensor& logits, const torch::Tensor& k); +void cpu_topk_topp_sampling(torch::Tensor& logits, const torch::Tensor& k, + const torch::Tensor& p); + namespace cpu_utils { void eagle_prepare_inputs_padded_kernel_impl( const torch::Tensor& cu_num_draft_tokens, @@ -1006,6 +1012,15 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) { ops.def("greedy_argmax(Tensor logits) -> Tensor"); ops.impl("greedy_argmax", torch::kCPU, &greedy_argmax); + ops.def("cpu_topp_sampling(Tensor! logits, Tensor p) -> ()"); + ops.impl("cpu_topp_sampling", torch::kCPU, &cpu_topp_sampling); + + ops.def("cpu_topk_sampling(Tensor! logits, Tensor k) -> ()"); + ops.impl("cpu_topk_sampling", torch::kCPU, &cpu_topk_sampling); + + ops.def("cpu_topk_topp_sampling(Tensor! logits, Tensor k, Tensor p) -> ()"); + ops.impl("cpu_topk_topp_sampling", torch::kCPU, &cpu_topk_topp_sampling); + // Speculative decoding kernels ops.def( "eagle_prepare_inputs_padded_kernel_impl(Tensor cu_num_draft_tokens, " diff --git a/tests/v1/sample/test_topk_topp_sampler.py b/tests/v1/sample/test_topk_topp_sampler.py index 924f0703d3e5..de5e88b10705 100644 --- a/tests/v1/sample/test_topk_topp_sampler.py +++ b/tests/v1/sample/test_topk_topp_sampler.py @@ -9,6 +9,8 @@ from vllm.triton_utils import HAS_TRITON from vllm.utils.torch_utils import set_random_seed from vllm.v1.sample.ops.topk_topp_sampler import ( + _HAS_CPU_SAMPLING_OPS, + apply_top_k_top_p_cpu, apply_top_k_top_p_pytorch, random_sample, ) @@ -20,6 +22,27 @@ VOCAB_SIZE = 128 * 1024 +def _chi2_check(empirical, expected, chisquare_fn, *, label, alpha): + import numpy as np + + outside = (expected == 0) & (empirical > 0) + assert not outside.any(), ( + f"{label}: sampled out-of-support tokens " + f"(zero expected prob): indices={outside.nonzero()[0].tolist()}" + ) + in_support = expected > 0 + if int(in_support.sum()) <= 1: + return + emp = empirical[in_support].astype(np.float64) + exp = expected[in_support].astype(np.float64) + exp = exp * (emp.sum() / exp.sum()) + chi2, p_value = chisquare_fn(emp, exp) + assert p_value > alpha, ( + f"{label}: distribution differs from theoretical: " + f"chi2={chi2:.2f} p_value={p_value:.2e} alpha={alpha}" + ) + + def _flashinfer_topk_topp_supported() -> bool: """True iff the FlashInfer top-k/top-p sampler is usable on this host. @@ -1124,7 +1147,6 @@ def test_distribution_matches_theoretical(self, topk, topp): from vllm.v1.sample.ops.topk_topp_sampler import ( apply_top_k_top_p, flashinfer_sample, - random_sample, ) torch.set_default_device(DEVICE_TYPE) @@ -1182,30 +1204,479 @@ def test_distribution_matches_theoretical(self, topk, topp): ) def _chi2_check(self, empirical, expected, chisquare_fn, *, label): - import numpy as np - - # Hard check: the sampler must never produce a token outside the - # expected support (zero theoretical probability). - outside = (expected == 0) & (empirical > 0) - assert not outside.any(), ( - f"{label}: sampled out-of-support tokens " - f"(zero expected prob): indices={outside.nonzero()[0].tolist()}" - ) - # Skip chi-square in the degenerate case where the support - # collapses to a single token (e.g. very restrictive joint - # top-k + top-p): all samples must land there and the hard - # check above already verified they do. - in_support = expected > 0 - if int(in_support.sum()) <= 1: - return - # Soft check: chi-square goodness-of-fit on in-support tokens. - # Cast to float64 so the rescaling step below stays within - # scipy.chisquare's strict 1.5e-8 sum-equality tolerance. - emp = empirical[in_support].astype(np.float64) - exp = expected[in_support].astype(np.float64) - exp = exp * (emp.sum() / exp.sum()) - chi2, p_value = chisquare_fn(emp, exp) - assert p_value > self.ALPHA, ( - f"{label}: distribution differs from theoretical: " - f"chi2={chi2:.2f} p_value={p_value:.2e} alpha={self.ALPHA}" + _chi2_check(empirical, expected, chisquare_fn, label=label, alpha=self.ALPHA) + + +# ============================================================================= +# CPU kernel tests +# ============================================================================= + + +@pytest.mark.skipif( + not _HAS_CPU_SAMPLING_OPS, + reason="CPU top-k/top-p C++ kernels not available on this build", +) +class TestCpuTopkTopp: + """Tests for the CPU top-k/top-p kernel. + + Structurally mirrors `TestTritonTopkTopp`: same `_compare_results` + contract (parity vs. PyTorch reference, with the same survivor-count + tolerance for top-p), same per-mode and edge-case test names. The + CPU kernel runs on `cpu` device regardless of `DEVICE_TYPE`. + """ + + @pytest.fixture(autouse=True) + def setup(self): + """Set up test fixtures.""" + torch.set_default_device("cpu") + self.generator = Generator(device="cpu").manual_seed(42) + + def _compare_results( + self, + logits: torch.Tensor, + k: torch.Tensor | None, + p: torch.Tensor | None, + ): + """Compare CPU kernel results with PyTorch sorting implementation. + + For top-k only, we expect exact match. + For top-p (with or without top-k), we allow small differences due to + floating-point precision in probability sum calculations. + """ + logits_pytorch = logits.clone() + logits_cpu = logits.clone().to(torch.float32) + + result_pytorch = apply_top_k_top_p_pytorch( + logits_pytorch, k, p, allow_cpu_sync=True + ) + k_i32 = k.to(torch.int32) if k is not None else None + p_f32 = p.to(torch.float32) if p is not None else None + result_cpu = apply_top_k_top_p_cpu(logits_cpu, k_i32, p_f32) + + assert not result_cpu.isnan().any(), "NaN found in CPU kernel result" + + pytorch_kept = (result_pytorch != float("-inf")).sum(dim=-1) + cpu_kept = (result_cpu != float("-inf")).sum(dim=-1) + + if p is None: + # Top-k only: expect exact match + assert torch.equal(pytorch_kept, cpu_kept), ( + f"Top-k mask mismatch: PyTorch kept {pytorch_kept.tolist()}, " + f"CPU kept {cpu_kept.tolist()}" + ) + else: + # Top-p involved: allow small differences + # Either < 1% of kept values OR < 5 values absolute + max_diff = (pytorch_kept - cpu_kept).abs().max().item() + max_kept = pytorch_kept.max().item() + if max_kept > 0 and max_diff > 3: + diff_pct = max_diff / max_kept * 100 + assert diff_pct < 0.5, ( + f"Top-p mask difference too large: {diff_pct:.2f}% " + f"(max diff {max_diff} values out of {max_kept})" + ) + + @pytest.mark.parametrize("batch_size", [1, 8, 32, 128, 512, 1024]) + @pytest.mark.parametrize("vocab_size", [1024, 32000, 128256]) + def test_topk_only(self, batch_size: int, vocab_size: int): + """Test top-k only (p=None).""" + logits = torch.randn( + batch_size, vocab_size, generator=self.generator, dtype=torch.float32 + ) + k = torch.randint( + 1, min(100, vocab_size), (batch_size,), generator=self.generator + ) + # Randomly disable top-k for some rows (~25%) + disable_mask = torch.randint(0, 4, (batch_size,), generator=self.generator) == 0 + k.masked_fill_(disable_mask, vocab_size) + + self._compare_results(logits, k, p=None) + + @pytest.mark.parametrize("batch_size", [1, 8, 32, 128, 512, 1024]) + @pytest.mark.parametrize("vocab_size", [1024, 32000, 128256]) + def test_topp_only(self, batch_size: int, vocab_size: int): + """Test top-p only (k=None).""" + logits = torch.randn( + batch_size, vocab_size, generator=self.generator, dtype=torch.float32 + ) + p = torch.rand(batch_size, generator=self.generator) * 0.9 + 0.1 # [0.1, 1.0] + # Randomly disable top-p for some rows (~25%) + disable_mask = torch.randint(0, 4, (batch_size,), generator=self.generator) == 0 + p.masked_fill_(disable_mask, 1.0) + + self._compare_results(logits, k=None, p=p) + + @pytest.mark.parametrize("batch_size", [1, 8, 32, 128, 512, 1024]) + @pytest.mark.parametrize("vocab_size", [1024, 32000, 128256]) + def test_topk_and_topp(self, batch_size: int, vocab_size: int): + """Test combined top-k and top-p.""" + logits = torch.randn( + batch_size, vocab_size, generator=self.generator, dtype=torch.float32 + ) + k = torch.randint( + 1, min(100, vocab_size), (batch_size,), generator=self.generator + ) + p = torch.rand(batch_size, generator=self.generator) * 0.9 + 0.1 # [0.1, 1.0] + + # Randomly disable top-k for some rows (~25%) + disable_k = torch.randint(0, 4, (batch_size,), generator=self.generator) == 0 + k.masked_fill_(disable_k, vocab_size) + # Randomly disable top-p for some rows (~25%) + disable_p = torch.randint(0, 4, (batch_size,), generator=self.generator) == 0 + p.masked_fill_(disable_p, 1.0) + + self._compare_results(logits, k, p) + + def test_both_disabled(self): + """Test when both k and p are None (should be no-op).""" + logits = torch.randn(32, 1024, generator=self.generator, dtype=torch.float32) + logits_clone = logits.clone() + + result = apply_top_k_top_p_cpu(logits_clone, k=None, p=None) + + assert torch.equal(result, logits), "Should be no-op when both k and p are None" + + def test_extreme_k_values(self): + """Test edge cases for k values.""" + batch_size, vocab_size = 16, 1024 + logits = torch.randn( + batch_size, vocab_size, generator=self.generator, dtype=torch.float32 + ) + + # k=1 (keep only top 1) + k = torch.ones(batch_size, dtype=torch.int32) + self._compare_results(logits.clone(), k, p=None) + + # k=vocab_size (keep all) + k = torch.full((batch_size,), vocab_size, dtype=torch.int32) + self._compare_results(logits.clone(), k, p=None) + + # Mixed extreme values + k = torch.tensor([1, vocab_size, 2, vocab_size - 1] * 4, dtype=torch.int32) + self._compare_results(logits.clone(), k, p=None) + + def test_extreme_p_values(self): + """Test edge cases for p values.""" + batch_size, vocab_size = 16, 1024 + logits = torch.randn( + batch_size, vocab_size, generator=self.generator, dtype=torch.float32 + ) + + # p=1.0 (keep all) + p = torch.ones(batch_size, dtype=torch.float32) + self._compare_results(logits.clone(), k=None, p=p) + + # Mixed values + p = torch.tensor([0.1, 0.5, 0.9, 1.0] * 4, dtype=torch.float32) + self._compare_results(logits.clone(), k=None, p=p) + + # p close to 0: kernel diff vs. sort-based ref is ≤2 tokens at + # vocab=1024 with no PAD inputs (standard binary-search precision). + p = torch.full((batch_size,), 0.01, dtype=torch.float32) + self._compare_results(logits.clone(), k=None, p=p) + + def test_large_batch(self): + """Test with a large batch size.""" + batch_size, vocab_size = 512, 32000 + logits = torch.randn( + batch_size, vocab_size, generator=self.generator, dtype=torch.float32 + ) + k = torch.randint(1, 50, (batch_size,), generator=self.generator) + p = torch.rand(batch_size, generator=self.generator) * 0.5 + 0.5 + + self._compare_results(logits, k, p) + + # ----------------------------------------------------------------- + # Tests for -inf logits (e.g. from grammar / structured output masks) + # ----------------------------------------------------------------- + + @pytest.mark.parametrize("inf_fraction", [0.5, 0.9, 0.99]) + def test_topk_with_neginf_logits(self, inf_fraction: float): + """Top-k with many -inf logits (simulating grammar bitmask).""" + batch_size, vocab_size = 32, 128256 + logits = torch.randn( + batch_size, vocab_size, generator=self.generator, dtype=torch.float32 + ) + mask = ( + torch.rand(batch_size, vocab_size, generator=self.generator) < inf_fraction + ) + logits[mask] = float("-inf") + + k = torch.randint( + 1, 50, (batch_size,), generator=self.generator, dtype=torch.int32 + ) + result = apply_top_k_top_p_cpu(logits.clone(), k, None) + + assert not result.isnan().any(), "NaN found in top-k result with -inf logits" + for i in range(batch_size): + kept = (result[i] > float("-inf")).sum().item() + assert kept <= k[i].item(), f"Row {i}: kept {kept} > k={k[i].item()}" + finite_in = (logits[i] > float("-inf")).sum().item() + if finite_in > 0: + assert kept > 0, f"Row {i}: no tokens kept despite finite input" + + @pytest.mark.parametrize("inf_fraction", [0.5, 0.9, 0.99]) + def test_topp_with_neginf_logits(self, inf_fraction: float): + """Top-p with many -inf logits. + + Exercises the PAD-sentinel blend in Pass 1 of the kernel. + """ + batch_size, vocab_size = 32, 128256 + logits = torch.randn( + batch_size, vocab_size, generator=self.generator, dtype=torch.float32 + ) + mask = ( + torch.rand(batch_size, vocab_size, generator=self.generator) < inf_fraction + ) + logits[mask] = float("-inf") + + p = ( + torch.rand(batch_size, generator=self.generator, dtype=torch.float32) * 0.9 + + 0.1 + ) + result = apply_top_k_top_p_cpu(logits.clone(), None, p) + + assert not result.isnan().any(), "NaN found in top-p result with -inf logits" + for i in range(batch_size): + finite_in = (logits[i] > float("-inf")).sum().item() + kept = (result[i] > float("-inf")).sum().item() + if finite_in > 0: + assert kept > 0, f"Row {i}: no tokens kept despite finite input" + + @pytest.mark.parametrize("inf_fraction", [0.5, 0.9, 0.99]) + def test_topk_topp_with_neginf_logits(self, inf_fraction: float): + """Combined top-k + top-p with many -inf logits.""" + batch_size, vocab_size = 32, 128256 + logits = torch.randn( + batch_size, vocab_size, generator=self.generator, dtype=torch.float32 + ) + mask = ( + torch.rand(batch_size, vocab_size, generator=self.generator) < inf_fraction + ) + logits[mask] = float("-inf") + + k = torch.randint( + 1, 50, (batch_size,), generator=self.generator, dtype=torch.int32 + ) + p = ( + torch.rand(batch_size, generator=self.generator, dtype=torch.float32) * 0.9 + + 0.1 + ) + result = apply_top_k_top_p_cpu(logits.clone(), k, p) + + assert not result.isnan().any(), ( + "NaN found in top-k+top-p result with -inf logits" + ) + for i in range(batch_size): + kept = (result[i] > float("-inf")).sum().item() + assert kept <= k[i].item(), f"Row {i}: kept {kept} > k={k[i].item()}" + + def test_all_neginf_logits(self): + """All logits are -inf (fully masked). Kernel should be a no-op.""" + batch_size, vocab_size = 16, 128256 + logits = torch.full( + (batch_size, vocab_size), float("-inf"), dtype=torch.float32 + ) + + k = torch.randint( + 1, 50, (batch_size,), generator=self.generator, dtype=torch.int32 + ) + p = torch.full((batch_size,), 0.9, dtype=torch.float32) + + # top-k only + result = apply_top_k_top_p_cpu(logits.clone(), k, None) + assert not result.isnan().any(), "NaN from all-inf top-k" + assert (result == float("-inf")).all(), "Expected all -inf unchanged" + + # top-p only + result = apply_top_k_top_p_cpu(logits.clone(), None, p) + assert not result.isnan().any(), "NaN from all-inf top-p" + assert (result == float("-inf")).all(), "Expected all -inf unchanged" + + # top-k + top-p + result = apply_top_k_top_p_cpu(logits.clone(), k, p) + assert not result.isnan().any(), "NaN from all-inf top-k+top-p" + assert (result == float("-inf")).all(), "Expected all -inf unchanged" + + def test_few_valid_tokens_with_neginf(self): + """Only a handful of tokens are finite per row (strict grammar).""" + batch_size, vocab_size = 32, 128256 + logits = torch.full( + (batch_size, vocab_size), float("-inf"), dtype=torch.float32 + ) + for i in range(batch_size): + indices = torch.randperm(vocab_size, generator=self.generator)[:5] + logits[i, indices] = torch.randn( + 5, generator=self.generator, dtype=torch.float32 + ) + + k = torch.full((batch_size,), 50, dtype=torch.int32) + p = torch.full((batch_size,), 0.9, dtype=torch.float32) + + # top-k only (k=50 but only 5 finite → keep all 5) + result = apply_top_k_top_p_cpu(logits.clone(), k, None) + assert not result.isnan().any() + for i in range(batch_size): + kept = (result[i] > float("-inf")).sum().item() + assert kept == 5, f"Row {i}: expected 5 kept, got {kept}" + + # top-k with k < num_finite + k_small = torch.full((batch_size,), 3, dtype=torch.int32) + result = apply_top_k_top_p_cpu(logits.clone(), k_small, None) + assert not result.isnan().any() + for i in range(batch_size): + kept = (result[i] > float("-inf")).sum().item() + assert kept <= 3, f"Row {i}: expected <=3 kept, got {kept}" + + # top-p only + result = apply_top_k_top_p_cpu(logits.clone(), None, p) + assert not result.isnan().any() + for i in range(batch_size): + kept = (result[i] > float("-inf")).sum().item() + assert kept > 0, f"Row {i}: no tokens kept" + + @pytest.mark.parametrize("num_valid", [1, 2, 5, 10, 50]) + @pytest.mark.parametrize( + "mode", + ["topk_only", "topp_only", "topk_and_topp"], + ) + def test_equal_logits_few_valid(self, num_valid: int, mode: str): + """Few valid tokens all sharing the same logit value. + + Stresses the boundary-tie duplicate-keep logic. With all-equal + finite logits the binary-search pivot can't differentiate — the + kernel's safe-fallback path must keep at least one token rather + than masking everything. + """ + batch_size, vocab_size = 32, 128256 + logits = torch.full( + (batch_size, vocab_size), float("-inf"), dtype=torch.float32 + ) + for i in range(batch_size): + indices = torch.randperm(vocab_size, generator=self.generator)[:num_valid] + logits[i, indices] = 1.0 # all equal + + k: torch.Tensor | None = None + p: torch.Tensor | None = None + if mode in ("topk_only", "topk_and_topp"): + k = torch.full((batch_size,), max(1, num_valid - 1), dtype=torch.int32) + if mode in ("topp_only", "topk_and_topp"): + p = torch.full((batch_size,), 0.95, dtype=torch.float32) + + result = apply_top_k_top_p_cpu(logits.clone(), k, p) + + assert not result.isnan().any(), "NaN in equal-logit result" + for i in range(batch_size): + kept = (result[i] > float("-inf")).sum().item() + assert kept > 0, ( + f"Row {i}: all tokens masked with {num_valid} equal-valued " + f"finite logits ({mode})" + ) + + @pytest.mark.parametrize("num_valid", [2, 5, 10]) + def test_nearly_equal_logits_topp(self, num_valid: int): + """Few valid tokens with very similar (but not identical) logits.""" + batch_size, vocab_size = 32, 128256 + logits = torch.full( + (batch_size, vocab_size), float("-inf"), dtype=torch.float32 + ) + for i in range(batch_size): + indices = torch.randperm(vocab_size, generator=self.generator)[:num_valid] + # Tiny spread: values in [1.0, 1.0 + 1e-6] + logits[i, indices] = ( + 1.0 + + torch.rand(num_valid, generator=self.generator, dtype=torch.float32) + * 1e-6 + ) + + p = torch.full((batch_size,), 0.95, dtype=torch.float32) + result = apply_top_k_top_p_cpu(logits.clone(), None, p) + + assert not result.isnan().any(), "NaN in nearly-equal-logit result" + for i in range(batch_size): + kept = (result[i] > float("-inf")).sum().item() + assert kept > 0, ( + f"Row {i}: all tokens masked with {num_valid} " + f"nearly-equal finite logits" + ) + + def test_mixed_neginf_and_normal_rows(self): + """Batch with a mix of normal rows and heavily-masked rows.""" + batch_size, vocab_size = 32, 32000 + logits = torch.randn( + batch_size, vocab_size, generator=self.generator, dtype=torch.float32 + ) + # Mask even rows heavily (99% -inf), leave odd rows normal. + for i in range(0, batch_size, 2): + mask = torch.rand(vocab_size, generator=self.generator) < 0.99 + logits[i][mask] = float("-inf") + + k = torch.randint( + 1, 50, (batch_size,), generator=self.generator, dtype=torch.int32 + ) + p = ( + torch.rand(batch_size, generator=self.generator, dtype=torch.float32) * 0.9 + + 0.1 + ) + + result = apply_top_k_top_p_cpu(logits.clone(), k, p) + assert not result.isnan().any(), "NaN in mixed normal/-inf batch" + for i in range(batch_size): + kept = (result[i] > float("-inf")).sum().item() + assert kept <= k[i].item() + finite_in = (logits[i] > float("-inf")).sum().item() + if finite_in > 0: + assert kept > 0, f"Row {i}: no tokens kept" + + @pytest.mark.skipif( + not _HAS_CPU_SAMPLING_OPS, reason="CPU top-p/top-k kernel unavailable" + ) + @pytest.mark.parametrize("vocab_size", [17, 128, 1024, 32000]) + @pytest.mark.parametrize("batch_size", [1, 8, 32]) + def test_fused_equals_sequential(self, batch_size: int, vocab_size: int): + """Fused top_k_p_row must produce the same output as sequential + top_k_row + top_p_row on non-tied rows (no boundary duplicates). + + Uses torch.allclose with rtol=1e-5, atol=1e-7 because FP64 + reduction order may differ between the two code paths. + """ + from vllm._custom_ops import ( + cpu_topk_sampling as _cpu_topk_sampling_op, + ) + from vllm._custom_ops import ( + cpu_topk_topp_sampling as _cpu_topk_topp_sampling_op, + ) + from vllm._custom_ops import ( + cpu_topp_sampling as _cpu_topp_sampling_op, + ) + + # Random non-tied logits; avoid uniform rows so boundary-tie is unlikely. + logits = torch.randn( + batch_size, vocab_size, generator=self.generator, dtype=torch.float32 + ) + k = torch.randint( + 1, + min(50, vocab_size - 1), + (batch_size,), + generator=self.generator, + dtype=torch.int32, + ) + p = ( + torch.rand(batch_size, generator=self.generator, dtype=torch.float32) * 0.7 + + 0.2 + ) # [0.2, 0.9] + + # Sequential path: apply top-k then top-p independently. + seq = logits.clone() + _cpu_topk_sampling_op(seq, k) + _cpu_topp_sampling_op(seq, p) + + # Fused path. + fused = logits.clone() + _cpu_topk_topp_sampling_op(fused, k, p) + + assert torch.allclose(fused, seq, rtol=1e-5, atol=1e-7), ( + f"Fused and sequential outputs differ at " + f"batch_size={batch_size}, vocab_size={vocab_size}.\n" + f"Max abs diff: {(fused - seq).abs().max().item():.2e}" ) diff --git a/vllm/_custom_ops.py b/vllm/_custom_ops.py index 850917f0efa9..93d2d25b8bf9 100644 --- a/vllm/_custom_ops.py +++ b/vllm/_custom_ops.py @@ -236,6 +236,23 @@ def rotary_embedding( ) +def cpu_topp_sampling(logits: torch.Tensor, p: torch.Tensor) -> None: + """In-place top-p masking using the C++ ternary-search kernel (CPU only).""" + torch.ops._C.cpu_topp_sampling(logits, p) + + +def cpu_topk_sampling(logits: torch.Tensor, k: torch.Tensor) -> None: + """In-place top-k masking using the C++ ternary-search kernel (CPU only).""" + torch.ops._C.cpu_topk_sampling(logits, k) + + +def cpu_topk_topp_sampling( + logits: torch.Tensor, k: torch.Tensor, p: torch.Tensor +) -> None: + """In-place fused top-k+top-p masking (CPU only).""" + torch.ops._C.cpu_topk_topp_sampling(logits, k, p) + + # layer norm ops def rms_norm( out: torch.Tensor, diff --git a/vllm/v1/sample/ops/topk_topp_sampler.py b/vllm/v1/sample/ops/topk_topp_sampler.py index 14ee9b7064f4..218609eb170f 100644 --- a/vllm/v1/sample/ops/topk_topp_sampler.py +++ b/vllm/v1/sample/ops/topk_topp_sampler.py @@ -23,6 +23,21 @@ apply_top_k_top_p_triton, ) +try: + from vllm._custom_ops import cpu_topk_sampling as _cpu_topk_sampling_op + from vllm._custom_ops import ( + cpu_topk_topp_sampling as _cpu_topk_topp_sampling_op, + ) + from vllm._custom_ops import cpu_topp_sampling as _cpu_topp_sampling_op + + # Probe the dispatcher: raises AttributeError on non-CPU builds where + # the op is absent (the Python wrappers import cleanly regardless). + torch.ops._C.cpu_topk_sampling # noqa: B018 + + _HAS_CPU_SAMPLING_OPS = True +except (ImportError, AttributeError): + _HAS_CPU_SAMPLING_OPS = False + logger = init_logger(__name__) @@ -354,7 +369,7 @@ def forward_cpu( (#59786). Falls back to native sampling when fp64 Gumbel noise is requested. """ - logits = apply_top_k_top_p(logits, k, p) + logits = apply_top_k_top_p_cpu(logits, k, p) logits_to_return = None if self.logprobs_mode == "processed_logits": logits_to_return = logits @@ -587,6 +602,50 @@ def apply_top_k_only(logits: torch.Tensor, k: torch.Tensor) -> torch.Tensor: return logits.masked_fill_(logits < top_k_mask, -float("inf")) +def _apply_top_k_only_cpu(logits: torch.Tensor, k: torch.Tensor) -> torch.Tensor: + if not _HAS_CPU_SAMPLING_OPS: + return apply_top_k_only(logits, k) + logits = logits.contiguous() + _cpu_topk_sampling_op(logits, k.to(torch.int32)) + return logits + + +def apply_top_k_top_p_cpu( + logits: torch.Tensor, + k: torch.Tensor | None, + p: torch.Tensor | None, +) -> torch.Tensor: + """CPU-optimized top-k + top-p. Uses C++ ternary-search kernels.""" + if p is None and k is None: + return logits + + if k is not None: + # k=0 is not valid; callers use k=vocab_size to disable top-k. + assert (k > 0).all(), ( + "k values must be > 0; use k=vocab_size to disable top-k filtering" + ) + + if p is None: + return _apply_top_k_only_cpu(logits, k) + + if k is None: + if _HAS_CPU_SAMPLING_OPS: + logits = logits.contiguous() + _cpu_topp_sampling_op(logits, p.float()) + return logits + return apply_top_k_top_p_pytorch(logits, None, p, allow_cpu_sync=True) + + # Joint k+p path. + if _HAS_CPU_SAMPLING_OPS: + logits = logits.contiguous() + _cpu_topk_topp_sampling_op(logits, k.to(torch.int32), p.float()) + return logits + + # Fallback: sequential top-k then top-p. + logits = _apply_top_k_only_cpu(logits, k) + return apply_top_k_top_p_pytorch(logits, None, p, allow_cpu_sync=True) + + def empty_exponential_noise_like( probs: torch.Tensor, use_fp64_gumbel: bool ) -> torch.Tensor: From f091d01f11d96fc82185d6d7bec9cd06904db236 Mon Sep 17 00:00:00 2001 From: Fadi Arafeh Date: Mon, 7 Sep 2026 10:14:01 +0000 Subject: [PATCH 2/4] CPU: vectorize top-k and top-p sampling - AArch64 (cherry picked from commit 267cbdf810e5f7542f0c442c091ae48af63cd0d0) Signed-off-by: Fadi Arafeh --- csrc/cpu/topk_topp_sampler.cpp | 90 +++++++++++++++++++++++++++++++--- 1 file changed, 84 insertions(+), 6 deletions(-) diff --git a/csrc/cpu/topk_topp_sampler.cpp b/csrc/cpu/topk_topp_sampler.cpp index 8a3baa2f8748..48045cb3a89c 100644 --- a/csrc/cpu/topk_topp_sampler.cpp +++ b/csrc/cpu/topk_topp_sampler.cpp @@ -114,6 +114,26 @@ struct Fp64Acc16 { return _mm_cvtsd_f64(_mm_hadd_pd(s2, s2)); } }; +#elif defined(__aarch64__) +struct Fp64Acc16 { + float64x2_t acc[8] = {vdupq_n_f64(0.0), vdupq_n_f64(0.0), + vdupq_n_f64(0.0), vdupq_n_f64(0.0), + vdupq_n_f64(0.0), vdupq_n_f64(0.0), + vdupq_n_f64(0.0), vdupq_n_f64(0.0)}; + inline void add(const vec_op::FP32Vec16& v) { + for (int i = 0; i < 4; ++i) { + const float32x4_t x = v.reg.val[i]; + acc[2 * i] = vaddq_f64(acc[2 * i], vcvt_f64_f32(vget_low_f32(x))); + acc[2 * i + 1] = + vaddq_f64(acc[2 * i + 1], vcvt_f64_f32(vget_high_f32(x))); + } + } + inline double reduce() const { + float64x2_t sum = acc[0]; + for (int i = 1; i < 8; ++i) sum = vaddq_f64(sum, acc[i]); + return vaddvq_f64(sum); + } +}; #endif __attribute__((always_inline)) static inline double sum_gt_to_double( @@ -161,6 +181,19 @@ __attribute__((always_inline)) static inline double sum_gt_to_double( __m128d lo128 = _mm256_castpd256_pd128(sum_a); __m128d s2 = _mm_add_pd(lo128, hi128); s = _mm_cvtsd_f64(_mm_hadd_pd(s2, s2)); +#elif defined(__aarch64__) + Fp64Acc16 acc; + const float32x4_t thr = vdupq_n_f32(threshold); + const float32x4_t zero = vdupq_n_f32(0.0f); + for (; j + 16 <= n; j += 16) { + float32x4x4_t masked; + for (int r = 0; r < 4; ++r) { + const float32x4_t v = vld1q_f32(buf + j + 4 * r); + masked.val[r] = vbslq_f32(vcgtq_f32(v, thr), v, zero); + } + acc.add(vec_op::FP32Vec16(masked)); + } + s = acc.reduce(); #endif for (; j < n; ++j) if (buf[j] > threshold) s += (double)buf[j]; @@ -189,6 +222,16 @@ __attribute__((always_inline)) static inline int count_within_tol( n += __builtin_popcount((unsigned)_mm256_movemask_ps(clo)); n += __builtin_popcount((unsigned)_mm256_movemask_ps(chi)); } +#elif defined(__aarch64__) + const float32x4_t c = vdupq_n_f32(center); + const float32x4_t t = vdupq_n_f32(tol); + for (; i + 16 <= V; i += 16) { + for (int r = 0; r < 4; ++r) { + const float32x4_t v = vld1q_f32(row + i + 4 * r); + const uint32x4_t mask = vcltq_f32(vabsq_f32(vsubq_f32(v, c)), t); + n += (int)vaddvq_u32(vshrq_n_u32(mask, 31)); + } + } #endif for (; i < V; ++i) if (fabsf(row[i] - center) < tol) ++n; @@ -218,6 +261,16 @@ __attribute__((always_inline)) static inline void mask_write_below( _mm256_blendv_ps(f16.reg_high, v16.reg_high, khi)) .save(row + i); } +#elif defined(__aarch64__) + const float32x4_t thr = vdupq_n_f32(threshold); + const float32x4_t fv = vdupq_n_f32(fill); + for (; i + 16 <= V; i += 16) { + for (int r = 0; r < 4; ++r) { + float32x4_t v = vld1q_f32(row + i + 4 * r); + v = vbslq_f32(vcgtq_f32(v, thr), v, fv); + vst1q_f32(row + i + 4 * r, v); + } + } #endif for (; i < V; ++i) if (!(row[i] > threshold)) row[i] = fill; @@ -232,7 +285,7 @@ __attribute__((always_inline)) static inline void vec_max_min_with_pad_blend( int i = 0; float max_l = -1e38f, min_l = 1e38f; int n_finite = 0; -#if defined(__AVX512F__) || defined(__AVX2__) +#if defined(__AVX512F__) || defined(__AVX2__) || defined(__aarch64__) { const float pos_inf_val = std::numeric_limits::infinity(); vec_op::FP32Vec16 maxv(row[0]); @@ -249,7 +302,7 @@ __attribute__((always_inline)) static inline void vec_max_min_with_pad_blend( _mm512_mask_blend_ps(is_pad, v16.reg, pos_inf16.reg)); __mmask16 fin = _mm512_cmp_ps_mask(v16.reg, sentinel16.reg, _CMP_GT_OQ); n_finite += __builtin_popcount((unsigned)fin); - #else + #elif defined(__AVX2__) __m256 lt_lo = _mm256_cmp_ps(v16.reg_low, sentinel16.reg_low, _CMP_LE_OS); __m256 lt_hi = _mm256_cmp_ps(v16.reg_high, sentinel16.reg_high, _CMP_LE_OS); @@ -262,6 +315,15 @@ __attribute__((always_inline)) static inline void vec_max_min_with_pad_blend( _mm256_cmp_ps(v16.reg_high, sentinel16.reg_high, _CMP_GT_OQ); n_finite += __builtin_popcount((unsigned)_mm256_movemask_ps(fin_lo)); n_finite += __builtin_popcount((unsigned)_mm256_movemask_ps(fin_hi)); + #else + float32x4x4_t safe_raw; + for (int r = 0; r < 4; ++r) { + const float32x4_t x = v16.reg.val[r]; + const uint32x4_t fin = vcgtq_f32(x, vdupq_n_f32(pad_sentinel)); + n_finite += (int)vaddvq_u32(vshrq_n_u32(fin, 31)); + safe_raw.val[r] = vbslq_f32(fin, x, vdupq_n_f32(pos_inf_val)); + } + vec_op::FP32Vec16 safe(safe_raw); #endif minv = minv.min(safe); } @@ -405,8 +467,23 @@ static void scan_pivot_stats(const float* buf, int n, float pivot, float tol, float ma = std::numeric_limits::infinity(); int na = 0; int i = 0; -#if defined(__AVX512F__) || defined(__AVX2__) +#if defined(__AVX512F__) || defined(__AVX2__) || defined(__aarch64__) { +#if defined(__aarch64__) + const float32x4_t vpivot = vdupq_n_f32(pivot); + const float32x4_t pos_inf = + vdupq_n_f32(std::numeric_limits::infinity()); + float32x4_t vmin = pos_inf; + for (; i + 16 <= n; i += 16) { + for (int r = 0; r < 4; ++r) { + const float32x4_t values = vld1q_f32(buf + i + 4 * r); + const uint32x4_t above = vcgtq_f32(values, vpivot); + na += (int)vaddvq_u32(vshrq_n_u32(above, 31)); + vmin = vminq_f32(vmin, vbslq_f32(above, values, pos_inf)); + } + } + ma = vminvq_f32(vmin); +#else const float pos_inf_val = std::numeric_limits::infinity(); const vec_op::FP32Vec16 vpivot16(pivot); const vec_op::FP32Vec16 pos_inf16(pos_inf_val); @@ -431,6 +508,7 @@ static void scan_pivot_stats(const float* buf, int n, float pivot, float tol, #endif } ma = vmin16.reduce_min(); +#endif } #endif for (; i < n; ++i) { @@ -573,14 +651,14 @@ static void _top_p_row_core(float* __restrict__ row, int V, float p_val, sigma = sigma + fabsf(sigma) * -0.25f; float outlier_logit = st.avg + st.std_v * sigma; -#if defined(__AVX512F__) || defined(__AVX2__) +#if defined(__AVX512F__) || defined(__AVX2__) || defined(__aarch64__) DEFINE_FAST_EXP vec_op::FP32Vec16 base16(st.max_l); #endif double sum_exp_d = 0.0; { int i = 0; -#if defined(__AVX512F__) || defined(__AVX2__) +#if defined(__AVX512F__) || defined(__AVX2__) || defined(__aarch64__) Fp64Acc16 acc; for (; i + 16 <= V; i += 16) { vec_op::FP32Vec16 v16(row + i); @@ -601,7 +679,7 @@ static void _top_p_row_core(float* __restrict__ row, int V, float p_val, { const float inv_sum = 1.0f / sum_exp; int i = 0; -#if defined(__AVX512F__) || defined(__AVX2__) +#if defined(__AVX512F__) || defined(__AVX2__) || defined(__aarch64__) vec_op::FP32Vec16 isum16(inv_sum); for (; i + 16 <= V; i += 16) { vec_op::FP32Vec16 v16(row + i); From bf6cd7fb269c487e5f1a38c0818c96c8b7e28a22 Mon Sep 17 00:00:00 2001 From: Michael Yang Date: Fri, 16 Oct 2026 14:07:02 +0100 Subject: [PATCH 3/4] CPU: clean up sampling kernels for AArch64 Remove the AVX512 and AVX2 sampling paths introduced by the base implementation while retaining the AArch64 implementations and scalar fallbacks. Remove the added AVX2 fast-exp macro and -mfma flag, apply repository formatting, and run the CPU top-k/top-p tests in the Arm CI kernel shard. Assisted-by: OpenAI Signed-off-by: Michael Yang --- .../scripts/hardware_ci/run-cpu-test-arm.sh | 1 + .../kernels/cpu/benchmark_cpu_topk_topp.py | 10 +- cmake/cpu_extension.cmake | 3 +- csrc/cpu/cpu_arch_macros.h | 58 ------ csrc/cpu/topk_topp_sampler.cpp | 192 +----------------- 5 files changed, 16 insertions(+), 248 deletions(-) diff --git a/.buildkite/scripts/hardware_ci/run-cpu-test-arm.sh b/.buildkite/scripts/hardware_ci/run-cpu-test-arm.sh index 1f598b277bb2..6f214e2cb582 100755 --- a/.buildkite/scripts/hardware_ci/run-cpu-test-arm.sh +++ b/.buildkite/scripts/hardware_ci/run-cpu-test-arm.sh @@ -44,6 +44,7 @@ kernel_tests() { set -e pytest -x -v -s tests/kernels/test_onednn.py pytest -x -v -s tests/kernels/test_cpu_fused_sampling.py + pytest -x -v -s tests/v1/sample/test_topk_topp_sampler.py::TestCpuTopkTopp pytest -x -v -s tests/kernels/attention/test_cpu_attn.py pytest -x -v -s tests/kernels/core/test_cpu_activation.py pytest -x -v -s tests/kernels/moe/test_cpu_fused_moe.py diff --git a/benchmarks/kernels/cpu/benchmark_cpu_topk_topp.py b/benchmarks/kernels/cpu/benchmark_cpu_topk_topp.py index f7dfc45e5f3b..e3322bc3c3a7 100644 --- a/benchmarks/kernels/cpu/benchmark_cpu_topk_topp.py +++ b/benchmarks/kernels/cpu/benchmark_cpu_topk_topp.py @@ -1,8 +1,7 @@ #!/usr/bin/env python3 # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Benchmark comparing CPU C++ kernels vs PyTorch sort-based top-k/top-p. +"""Benchmark comparing CPU C++ kernels vs PyTorch sort-based top-k/top-p. Compares: - apply_top_k_top_p_cpu (C++ kernel — ISA determined at build time) @@ -56,8 +55,7 @@ def calculate_ops_pct( vocab_size: int, batch_size: int, ) -> float: - """ - Calculate the percentage of active top-k and top-p operations. + """Calculate the percentage of active top-k and top-p operations. Returns percentage where 100% = batch_size ops. E.g., if all rows have both top-k and top-p active, returns 200%. @@ -109,9 +107,7 @@ def benchmark_function( warmup_iters: int = 5, benchmark_iters: int = 20, ) -> float: - """ - Benchmark a function and return avg_time_ms. - """ + """Benchmark a function and return avg_time_ms.""" # Warmup for _ in range(warmup_iters): logits_copy = logits.clone() diff --git a/cmake/cpu_extension.cmake b/cmake/cpu_extension.cmake index 42ed9e3919c5..a42ee4af8005 100644 --- a/cmake/cpu_extension.cmake +++ b/cmake/cpu_extension.cmake @@ -186,8 +186,7 @@ if (CMAKE_SYSTEM_PROCESSOR MATCHES "x86_64|amd64" OR ENABLE_X86_ISA) endif() list(APPEND CXX_COMPILE_FLAGS_AVX2 - "-mavx2" - "-mfma") + "-mavx2") elseif (POWER9_FOUND OR POWER10_FOUND OR POWER11_FOUND) message(STATUS "PowerPC detected") if (POWER9_FOUND) diff --git a/csrc/cpu/cpu_arch_macros.h b/csrc/cpu/cpu_arch_macros.h index 4d7cd42f089d..ccec4c2e87c7 100644 --- a/csrc/cpu/cpu_arch_macros.h +++ b/csrc/cpu/cpu_arch_macros.h @@ -58,64 +58,6 @@ }; #endif - #if !defined(__AVX512F__) && defined(__AVX2__) - #define DEFINE_FAST_EXP \ - const __m256 vec_factorial_1 = _mm256_set1_ps(0.999999701f); \ - const __m256 vec_factorial_2 = _mm256_set1_ps(0.499991506f); \ - const __m256 vec_factorial_3 = _mm256_set1_ps(0.166676521f); \ - const __m256 vec_factorial_4 = _mm256_set1_ps(0.0418978221f); \ - const __m256 vec_factorial_5 = _mm256_set1_ps(0.00828929059f); \ - const __m256 vec_exp_log2ef = \ - _mm256_castsi256_ps(_mm256_set1_epi32(0x3fb8aa3b)); \ - const __m256 vec_half = _mm256_set1_ps(0.5f); \ - const __m256 vec_one = _mm256_set1_ps(1.f); \ - const __m256 vec_zero = _mm256_set1_ps(0.f); \ - const __m256 vec_two = _mm256_set1_ps(2.f); \ - const __m256 vec_ln2f = \ - _mm256_castsi256_ps(_mm256_set1_epi32(0x3f317218)); \ - const __m256 vec_ln_flt_min = \ - _mm256_castsi256_ps(_mm256_set1_epi32(0xc2aeac50)); \ - const __m256 vec_ln_flt_max = \ - _mm256_castsi256_ps(_mm256_set1_epi32(0x42b17218)); \ - const __m256i vec_127 = _mm256_set1_epi32(0x0000007f); \ - const int n_mantissa_bits = 23; \ - auto avx2_expf = [&](__m256 values) __attribute__((always_inline)) { \ - auto less_ln_flt_min_mask = \ - _mm256_cmp_ps(values, vec_ln_flt_min, 1 /*_CMP_LT_OS*/); \ - auto vec_src = _mm256_min_ps(values, vec_ln_flt_max); \ - vec_src = _mm256_max_ps(vec_src, vec_ln_flt_min); \ - auto vec_fx = _mm256_fmadd_ps(vec_src, vec_exp_log2ef, vec_half); \ - /* AVX2 lacks cvt_roundps_epi32; use round_ps + cvtps_epi32 */ \ - auto vec_fx_rnd = _mm256_round_ps( \ - vec_fx, _MM_FROUND_TO_NEG_INF | _MM_FROUND_NO_EXC); \ - auto vec_fx_i = _mm256_cvtps_epi32(vec_fx_rnd); \ - vec_fx = _mm256_cvtepi32_ps(vec_fx_i); \ - auto vec_exp_poly = _mm256_fnmadd_ps(vec_fx, vec_ln2f, vec_src); \ - auto vec_res = \ - _mm256_fmadd_ps(vec_exp_poly, vec_factorial_5, vec_factorial_4); \ - vec_res = _mm256_fmadd_ps(vec_exp_poly, vec_res, vec_factorial_3); \ - vec_res = _mm256_fmadd_ps(vec_exp_poly, vec_res, vec_factorial_2); \ - vec_res = _mm256_fmadd_ps(vec_exp_poly, vec_res, vec_factorial_1); \ - vec_res = _mm256_fmadd_ps(vec_exp_poly, vec_res, vec_one); \ - auto vec_exp_number = _mm256_sub_ps(vec_fx, vec_one); \ - auto vec_exp_number_i = _mm256_cvtps_epi32(vec_exp_number); \ - auto vec_two_pow_n_i = _mm256_add_epi32(vec_exp_number_i, vec_127); \ - vec_two_pow_n_i = _mm256_slli_epi32(vec_two_pow_n_i, n_mantissa_bits); \ - auto vec_two_pow_n = _mm256_castsi256_ps(vec_two_pow_n_i); \ - /* _mm256_blendv_ps: mask sign-bit=1 selects b (zero), else a */ \ - vec_two_pow_n = \ - _mm256_blendv_ps(vec_two_pow_n, vec_zero, less_ln_flt_min_mask); \ - vec_res = _mm256_mul_ps(vec_res, vec_two_pow_n); \ - vec_res = _mm256_mul_ps(vec_res, vec_two); \ - return vec_res; \ - }; \ - auto fast_exp = [&](const vec_op::FP32Vec16& vec) \ - __attribute__((always_inline)) { \ - return vec_op::FP32Vec16(avx2_expf(vec.reg_low), \ - avx2_expf(vec.reg_high)); \ - }; - #endif - #endif #ifdef __aarch64__ diff --git a/csrc/cpu/topk_topp_sampler.cpp b/csrc/cpu/topk_topp_sampler.cpp index 48045cb3a89c..4cb2e6593b71 100644 --- a/csrc/cpu/topk_topp_sampler.cpp +++ b/csrc/cpu/topk_topp_sampler.cpp @@ -75,50 +75,10 @@ namespace { // FP64 accumulator for softmax sum over FP32Vec16 chunks. // FP32 accumulation over 128K elements loses ~3 ULP; FP64 avoids this. -#if defined(__AVX512F__) -struct Fp64Acc16 { - __m512d acc0 = _mm512_setzero_pd(); - __m512d acc1 = _mm512_setzero_pd(); - inline void add(const vec_op::FP32Vec16& v) { - __m256 lo8 = _mm512_castps512_ps256(v.reg); - __m256 hi8 = _mm512_extractf32x8_ps(v.reg, 1); - acc0 = _mm512_add_pd(acc0, _mm512_cvtps_pd(lo8)); - acc1 = _mm512_add_pd(acc1, _mm512_cvtps_pd(hi8)); - } - inline double reduce() const { - return _mm512_reduce_add_pd(acc0) + _mm512_reduce_add_pd(acc1); - } -}; -#elif defined(__AVX2__) -struct Fp64Acc16 { - __m256d acc0 = _mm256_setzero_pd(); - __m256d acc1 = _mm256_setzero_pd(); - __m256d acc2 = _mm256_setzero_pd(); - __m256d acc3 = _mm256_setzero_pd(); - inline void add(const vec_op::FP32Vec16& v) { - __m128 e_ll = _mm256_extractf128_ps(v.reg_low, 0); - __m128 e_lh = _mm256_extractf128_ps(v.reg_low, 1); - __m128 e_hl = _mm256_extractf128_ps(v.reg_high, 0); - __m128 e_hh = _mm256_extractf128_ps(v.reg_high, 1); - acc0 = _mm256_add_pd(acc0, _mm256_cvtps_pd(e_ll)); - acc1 = _mm256_add_pd(acc1, _mm256_cvtps_pd(e_lh)); - acc2 = _mm256_add_pd(acc2, _mm256_cvtps_pd(e_hl)); - acc3 = _mm256_add_pd(acc3, _mm256_cvtps_pd(e_hh)); - } - inline double reduce() const { - __m256d sum_a = - _mm256_add_pd(_mm256_add_pd(acc0, acc1), _mm256_add_pd(acc2, acc3)); - __m128d hi128 = _mm256_extractf128_pd(sum_a, 1); - __m128d lo128 = _mm256_castpd256_pd128(sum_a); - __m128d s2 = _mm_add_pd(lo128, hi128); - return _mm_cvtsd_f64(_mm_hadd_pd(s2, s2)); - } -}; -#elif defined(__aarch64__) +#if defined(__aarch64__) struct Fp64Acc16 { - float64x2_t acc[8] = {vdupq_n_f64(0.0), vdupq_n_f64(0.0), - vdupq_n_f64(0.0), vdupq_n_f64(0.0), - vdupq_n_f64(0.0), vdupq_n_f64(0.0), + float64x2_t acc[8] = {vdupq_n_f64(0.0), vdupq_n_f64(0.0), vdupq_n_f64(0.0), + vdupq_n_f64(0.0), vdupq_n_f64(0.0), vdupq_n_f64(0.0), vdupq_n_f64(0.0), vdupq_n_f64(0.0)}; inline void add(const vec_op::FP32Vec16& v) { for (int i = 0; i < 4; ++i) { @@ -140,48 +100,7 @@ __attribute__((always_inline)) static inline double sum_gt_to_double( const float* buf, int n, float threshold) { double s = 0.0; int j = 0; -#if defined(__AVX512F__) - __m512d acc0 = _mm512_setzero_pd(), acc1 = _mm512_setzero_pd(); - const __m512 thr16 = _mm512_set1_ps(threshold); - for (; j + 16 <= n; j += 16) { - __m512 b16 = _mm512_loadu_ps(buf + j); - __mmask16 mask = _mm512_cmp_ps_mask(b16, thr16, _CMP_GT_OQ); - __m256 lo8 = _mm512_castps512_ps256(b16); - __m256 hi8 = _mm512_extractf32x8_ps(b16, 1); - acc0 = _mm512_add_pd(acc0, - _mm512_maskz_cvtps_pd((__mmask8)(mask & 0xFF), lo8)); - acc1 = - _mm512_add_pd(acc1, _mm512_maskz_cvtps_pd((__mmask8)(mask >> 8), hi8)); - } - s = _mm512_reduce_add_pd(acc0) + _mm512_reduce_add_pd(acc1); -#elif defined(__AVX2__) - // and_ps zeros excluded lanes; blendv with -inf would corrupt the sum. - const __m256 thr = _mm256_set1_ps(threshold); - __m256d acc0 = _mm256_setzero_pd(), acc1 = _mm256_setzero_pd(); - __m256d acc2 = _mm256_setzero_pd(), acc3 = _mm256_setzero_pd(); - for (; j + 16 <= n; j += 16) { - __m256 b_lo = _mm256_loadu_ps(buf + j); - __m256 b_hi = _mm256_loadu_ps(buf + j + 8); - __m256 masked_lo = - _mm256_and_ps(b_lo, _mm256_cmp_ps(b_lo, thr, _CMP_GT_OQ)); - __m256 masked_hi = - _mm256_and_ps(b_hi, _mm256_cmp_ps(b_hi, thr, _CMP_GT_OQ)); - acc0 = _mm256_add_pd(acc0, - _mm256_cvtps_pd(_mm256_extractf128_ps(masked_lo, 0))); - acc1 = _mm256_add_pd(acc1, - _mm256_cvtps_pd(_mm256_extractf128_ps(masked_lo, 1))); - acc2 = _mm256_add_pd(acc2, - _mm256_cvtps_pd(_mm256_extractf128_ps(masked_hi, 0))); - acc3 = _mm256_add_pd(acc3, - _mm256_cvtps_pd(_mm256_extractf128_ps(masked_hi, 1))); - } - __m256d sum_a = - _mm256_add_pd(_mm256_add_pd(acc0, acc1), _mm256_add_pd(acc2, acc3)); - __m128d hi128 = _mm256_extractf128_pd(sum_a, 1); - __m128d lo128 = _mm256_castpd256_pd128(sum_a); - __m128d s2 = _mm_add_pd(lo128, hi128); - s = _mm_cvtsd_f64(_mm_hadd_pd(s2, s2)); -#elif defined(__aarch64__) +#if defined(__aarch64__) Fp64Acc16 acc; const float32x4_t thr = vdupq_n_f32(threshold); const float32x4_t zero = vdupq_n_f32(0.0f); @@ -204,25 +123,7 @@ __attribute__((always_inline)) static inline int count_within_tol( const float* row, int V, float center, float tol) { int n = 0; int i = 0; -#if defined(__AVX512F__) - const vec_op::FP32Vec16 c16(center); - const vec_op::FP32Vec16 t16(tol); - for (; i + 16 <= V; i += 16) { - vec_op::FP32Vec16 absdiff = (vec_op::FP32Vec16(row + i) - c16).abs(); - __mmask16 mask = _mm512_cmp_ps_mask(absdiff.reg, t16.reg, _CMP_LT_OQ); - n += __builtin_popcount((unsigned)mask); - } -#elif defined(__AVX2__) - const vec_op::FP32Vec16 c16(center); - const vec_op::FP32Vec16 t16(tol); - for (; i + 16 <= V; i += 16) { - vec_op::FP32Vec16 absdiff = (vec_op::FP32Vec16(row + i) - c16).abs(); - __m256 clo = _mm256_cmp_ps(absdiff.reg_low, t16.reg_low, _CMP_LT_OQ); - __m256 chi = _mm256_cmp_ps(absdiff.reg_high, t16.reg_high, _CMP_LT_OQ); - n += __builtin_popcount((unsigned)_mm256_movemask_ps(clo)); - n += __builtin_popcount((unsigned)_mm256_movemask_ps(chi)); - } -#elif defined(__aarch64__) +#if defined(__aarch64__) const float32x4_t c = vdupq_n_f32(center); const float32x4_t t = vdupq_n_f32(tol); for (; i + 16 <= V; i += 16) { @@ -241,27 +142,7 @@ __attribute__((always_inline)) static inline int count_within_tol( __attribute__((always_inline)) static inline void mask_write_below( float* row, int V, float threshold, float fill) { int i = 0; -#if defined(__AVX512F__) - const vec_op::FP32Vec16 thr16(threshold); - const vec_op::FP32Vec16 f16(fill); - for (; i + 16 <= V; i += 16) { - vec_op::FP32Vec16 v16(row + i); - __mmask16 keep = _mm512_cmp_ps_mask(v16.reg, thr16.reg, _CMP_GT_OQ); - vec_op::FP32Vec16(_mm512_mask_blend_ps(keep, f16.reg, v16.reg)) - .save(row + i); - } -#elif defined(__AVX2__) - const vec_op::FP32Vec16 thr16(threshold); - const vec_op::FP32Vec16 f16(fill); - for (; i + 16 <= V; i += 16) { - vec_op::FP32Vec16 v16(row + i); - __m256 klo = _mm256_cmp_ps(v16.reg_low, thr16.reg_low, _CMP_GT_OQ); - __m256 khi = _mm256_cmp_ps(v16.reg_high, thr16.reg_high, _CMP_GT_OQ); - vec_op::FP32Vec16(_mm256_blendv_ps(f16.reg_low, v16.reg_low, klo), - _mm256_blendv_ps(f16.reg_high, v16.reg_high, khi)) - .save(row + i); - } -#elif defined(__aarch64__) +#if defined(__aarch64__) const float32x4_t thr = vdupq_n_f32(threshold); const float32x4_t fv = vdupq_n_f32(fill); for (; i + 16 <= V; i += 16) { @@ -285,37 +166,14 @@ __attribute__((always_inline)) static inline void vec_max_min_with_pad_blend( int i = 0; float max_l = -1e38f, min_l = 1e38f; int n_finite = 0; -#if defined(__AVX512F__) || defined(__AVX2__) || defined(__aarch64__) +#if defined(__aarch64__) { const float pos_inf_val = std::numeric_limits::infinity(); vec_op::FP32Vec16 maxv(row[0]); vec_op::FP32Vec16 minv(pos_inf_val); - const vec_op::FP32Vec16 sentinel16(pad_sentinel); - const vec_op::FP32Vec16 pos_inf16(pos_inf_val); for (; i + 16 <= V; i += 16) { vec_op::FP32Vec16 v16(row + i); maxv = maxv.max(v16); - #if defined(__AVX512F__) - __mmask16 is_pad = - _mm512_cmp_ps_mask(v16.reg, sentinel16.reg, _CMP_LE_OS); - vec_op::FP32Vec16 safe( - _mm512_mask_blend_ps(is_pad, v16.reg, pos_inf16.reg)); - __mmask16 fin = _mm512_cmp_ps_mask(v16.reg, sentinel16.reg, _CMP_GT_OQ); - n_finite += __builtin_popcount((unsigned)fin); - #elif defined(__AVX2__) - __m256 lt_lo = _mm256_cmp_ps(v16.reg_low, sentinel16.reg_low, _CMP_LE_OS); - __m256 lt_hi = - _mm256_cmp_ps(v16.reg_high, sentinel16.reg_high, _CMP_LE_OS); - vec_op::FP32Vec16 safe( - _mm256_blendv_ps(v16.reg_low, pos_inf16.reg_low, lt_lo), - _mm256_blendv_ps(v16.reg_high, pos_inf16.reg_high, lt_hi)); - __m256 fin_lo = - _mm256_cmp_ps(v16.reg_low, sentinel16.reg_low, _CMP_GT_OQ); - __m256 fin_hi = - _mm256_cmp_ps(v16.reg_high, sentinel16.reg_high, _CMP_GT_OQ); - n_finite += __builtin_popcount((unsigned)_mm256_movemask_ps(fin_lo)); - n_finite += __builtin_popcount((unsigned)_mm256_movemask_ps(fin_hi)); - #else float32x4x4_t safe_raw; for (int r = 0; r < 4; ++r) { const float32x4_t x = v16.reg.val[r]; @@ -324,7 +182,6 @@ __attribute__((always_inline)) static inline void vec_max_min_with_pad_blend( safe_raw.val[r] = vbslq_f32(fin, x, vdupq_n_f32(pos_inf_val)); } vec_op::FP32Vec16 safe(safe_raw); - #endif minv = minv.min(safe); } max_l = maxv.reduce_max(); @@ -467,9 +324,8 @@ static void scan_pivot_stats(const float* buf, int n, float pivot, float tol, float ma = std::numeric_limits::infinity(); int na = 0; int i = 0; -#if defined(__AVX512F__) || defined(__AVX2__) || defined(__aarch64__) - { #if defined(__aarch64__) + { const float32x4_t vpivot = vdupq_n_f32(pivot); const float32x4_t pos_inf = vdupq_n_f32(std::numeric_limits::infinity()); @@ -483,32 +339,6 @@ static void scan_pivot_stats(const float* buf, int n, float pivot, float tol, } } ma = vminvq_f32(vmin); -#else - const float pos_inf_val = std::numeric_limits::infinity(); - const vec_op::FP32Vec16 vpivot16(pivot); - const vec_op::FP32Vec16 pos_inf16(pos_inf_val); - vec_op::FP32Vec16 vmin16(pos_inf_val); - for (; i + 16 <= n; i += 16) { - vec_op::FP32Vec16 b16(buf + i); - #if defined(__AVX512F__) - __mmask16 above = _mm512_cmp_ps_mask(b16.reg, vpivot16.reg, _CMP_GT_OQ); - na += __builtin_popcount((unsigned)above); - vmin16 = vmin16.min(vec_op::FP32Vec16( - _mm512_mask_blend_ps(above, pos_inf16.reg, b16.reg))); - #else - __m256 above_lo = - _mm256_cmp_ps(b16.reg_low, vpivot16.reg_low, _CMP_GT_OQ); - __m256 above_hi = - _mm256_cmp_ps(b16.reg_high, vpivot16.reg_high, _CMP_GT_OQ); - na += __builtin_popcount((unsigned)_mm256_movemask_ps(above_lo)); - na += __builtin_popcount((unsigned)_mm256_movemask_ps(above_hi)); - vmin16 = vmin16.min(vec_op::FP32Vec16( - _mm256_blendv_ps(pos_inf16.reg_low, b16.reg_low, above_lo), - _mm256_blendv_ps(pos_inf16.reg_high, b16.reg_high, above_hi))); - #endif - } - ma = vmin16.reduce_min(); -#endif } #endif for (; i < n; ++i) { @@ -651,14 +481,14 @@ static void _top_p_row_core(float* __restrict__ row, int V, float p_val, sigma = sigma + fabsf(sigma) * -0.25f; float outlier_logit = st.avg + st.std_v * sigma; -#if defined(__AVX512F__) || defined(__AVX2__) || defined(__aarch64__) +#if defined(__aarch64__) DEFINE_FAST_EXP vec_op::FP32Vec16 base16(st.max_l); #endif double sum_exp_d = 0.0; { int i = 0; -#if defined(__AVX512F__) || defined(__AVX2__) || defined(__aarch64__) +#if defined(__aarch64__) Fp64Acc16 acc; for (; i + 16 <= V; i += 16) { vec_op::FP32Vec16 v16(row + i); @@ -679,7 +509,7 @@ static void _top_p_row_core(float* __restrict__ row, int V, float p_val, { const float inv_sum = 1.0f / sum_exp; int i = 0; -#if defined(__AVX512F__) || defined(__AVX2__) || defined(__aarch64__) +#if defined(__aarch64__) vec_op::FP32Vec16 isum16(inv_sum); for (; i + 16 <= V; i += 16) { vec_op::FP32Vec16 v16(row + i); From 73f15af9a8cd73398f8e1b62646e1681e03f710b Mon Sep 17 00:00:00 2001 From: Rehan Khan Date: Mon, 2 Nov 2026 16:27:08 +0530 Subject: [PATCH 4/4] [CPU][S390x] Vectorise topk topp for s390x Signed-off-by: Rehan Khan --- csrc/cpu/topk_topp_sampler.cpp | 125 ++++++++++++++++++++++++++++++++- 1 file changed, 122 insertions(+), 3 deletions(-) diff --git a/csrc/cpu/topk_topp_sampler.cpp b/csrc/cpu/topk_topp_sampler.cpp index 4cb2e6593b71..d3f630e0131d 100644 --- a/csrc/cpu/topk_topp_sampler.cpp +++ b/csrc/cpu/topk_topp_sampler.cpp @@ -2,6 +2,7 @@ // SPDX-FileCopyrightText: Copyright contributors to the vLLM project #include "cpu_arch_macros.h" #include "cpu_types.hpp" +#include #include #include #include @@ -94,6 +95,44 @@ struct Fp64Acc16 { return vaddvq_f64(sum); } }; +#elif defined(__s390x__) +using vf32 = __vector float; +using vi32 = __vector signed int; + +static inline int hsum_i32(vi32 v) { return v[0] + v[1] + v[2] + v[3]; } + +static inline float hmin_f32(vf32 v) { + return std::min(std::min(v[0], v[1]), std::min(v[2], v[3])); +} + +struct Fp64Acc16 { + __vector double acc[8] = {}; + inline void add(const vec_op::FP32Vec16& v) { + for (int i = 0; i < 4; ++i) { + const vf32 x = v.reg.val[i]; + // vec_doublee widens lanes 0 and 2; rotate by one lane for 1 and 3. + acc[2 * i] += vec_doublee(x); + acc[2 * i + 1] += vec_doublee(vec_sld(x, x, 4)); + } + } + inline double reduce() const { + __vector double sum = acc[0]; + for (int i = 1; i < 8; ++i) sum += acc[i]; + return sum[0] + sum[1]; + } +}; + +// VXE exp() clamps its input, so exp(-inf) yields ~1e-38 instead of 0. +// Zero PAD lanes to match the scalar and NEON paths. +static inline vec_op::FP32Vec16 zero_pad_lanes(const vec_op::FP32Vec16& v, + const vec_op::FP32Vec16& e) { + const vf32 pad = vec_splats(kPadSentinel); + const vf32 zero = vec_splats(0.0f); + vec_op::f32x4x4_t out; + for (int r = 0; r < 4; ++r) + out.val[r] = vec_sel(zero, e.reg.val[r], vec_cmpgt(v.reg.val[r], pad)); + return vec_op::FP32Vec16(out); +} #endif __attribute__((always_inline)) static inline double sum_gt_to_double( @@ -113,6 +152,19 @@ __attribute__((always_inline)) static inline double sum_gt_to_double( acc.add(vec_op::FP32Vec16(masked)); } s = acc.reduce(); +#elif defined(__s390x__) + Fp64Acc16 acc; + const vf32 thr = vec_splats(threshold); + const vf32 zero = vec_splats(0.0f); + for (; j + 16 <= n; j += 16) { + vec_op::f32x4x4_t masked; + for (int r = 0; r < 4; ++r) { + const vf32 v = vec_xl(0, buf + j + 4 * r); + masked.val[r] = vec_sel(zero, v, vec_cmpgt(v, thr)); + } + acc.add(vec_op::FP32Vec16(masked)); + } + s = acc.reduce(); #endif for (; j < n; ++j) if (buf[j] > threshold) s += (double)buf[j]; @@ -133,6 +185,18 @@ __attribute__((always_inline)) static inline int count_within_tol( n += (int)vaddvq_u32(vshrq_n_u32(mask, 31)); } } +#elif defined(__s390x__) + const vf32 c = vec_splats(center); + const vf32 t = vec_splats(tol); + // Mask lanes are all-ones (-1), so subtracting counts matches. + vi32 cnt = vec_splats(0); + for (; i + 16 <= V; i += 16) { + for (int r = 0; r < 4; ++r) { + const vf32 v = vec_xl(0, row + i + 4 * r); + cnt -= (vi32)vec_cmplt(vec_abs(v - c), t); + } + } + n += hsum_i32(cnt); #endif for (; i < V; ++i) if (fabsf(row[i] - center) < tol) ++n; @@ -152,6 +216,15 @@ __attribute__((always_inline)) static inline void mask_write_below( vst1q_f32(row + i + 4 * r, v); } } +#elif defined(__s390x__) + const vf32 thr = vec_splats(threshold); + const vf32 fv = vec_splats(fill); + for (; i + 16 <= V; i += 16) { + for (int r = 0; r < 4; ++r) { + const vf32 v = vec_xl(0, row + i + 4 * r); + vec_xst(vec_sel(fv, v, vec_cmpgt(v, thr)), 0, row + i + 4 * r); + } + } #endif for (; i < V; ++i) if (!(row[i] > threshold)) row[i] = fill; @@ -187,6 +260,29 @@ __attribute__((always_inline)) static inline void vec_max_min_with_pad_blend( max_l = maxv.reduce_max(); min_l = minv.reduce_min(); } +#elif defined(__s390x__) + { + const vf32 pad = vec_splats(pad_sentinel); + const vf32 pos_inf = vec_splats(std::numeric_limits::infinity()); + vec_op::FP32Vec16 maxv(row[0]); + vec_op::FP32Vec16 minv(std::numeric_limits::infinity()); + vi32 cnt = vec_splats(0); + for (; i + 16 <= V; i += 16) { + vec_op::FP32Vec16 v16(row + i); + maxv = maxv.max(v16); + vec_op::f32x4x4_t safe_raw; + for (int r = 0; r < 4; ++r) { + const vf32 x = v16.reg.val[r]; + const auto fin = vec_cmpgt(x, pad); + cnt -= (vi32)fin; + safe_raw.val[r] = vec_sel(pos_inf, x, fin); + } + minv = minv.min(vec_op::FP32Vec16(safe_raw)); + } + n_finite = hsum_i32(cnt); + max_l = maxv.reduce_max(); + min_l = minv.reduce_min(); + } #endif *max_out = max_l; *min_out = min_l; @@ -340,6 +436,23 @@ static void scan_pivot_stats(const float* buf, int n, float pivot, float tol, } ma = vminvq_f32(vmin); } +#elif defined(__s390x__) + { + const vf32 vpivot = vec_splats(pivot); + const vf32 pos_inf = vec_splats(std::numeric_limits::infinity()); + vf32 vmin = pos_inf; + vi32 cnt = vec_splats(0); + for (; i + 16 <= n; i += 16) { + for (int r = 0; r < 4; ++r) { + const vf32 values = vec_xl(0, buf + i + 4 * r); + const auto above = vec_cmpgt(values, vpivot); + cnt -= (vi32)above; + vmin = vec_min(vmin, vec_sel(pos_inf, values, above)); + } + } + na = hsum_i32(cnt); + ma = hmin_f32(vmin); + } #endif for (; i < n; ++i) { if (buf[i] > pivot) { @@ -481,18 +594,21 @@ static void _top_p_row_core(float* __restrict__ row, int V, float p_val, sigma = sigma + fabsf(sigma) * -0.25f; float outlier_logit = st.avg + st.std_v * sigma; -#if defined(__aarch64__) +#if defined(__aarch64__) || defined(__s390x__) DEFINE_FAST_EXP vec_op::FP32Vec16 base16(st.max_l); #endif double sum_exp_d = 0.0; { int i = 0; -#if defined(__aarch64__) +#if defined(__aarch64__) || defined(__s390x__) Fp64Acc16 acc; for (; i + 16 <= V; i += 16) { vec_op::FP32Vec16 v16(row + i); vec_op::FP32Vec16 e16 = fast_exp(v16 - base16); + #if defined(__s390x__) + e16 = zero_pad_lanes(v16, e16); + #endif acc.add(e16); } sum_exp_d += acc.reduce(); @@ -509,11 +625,14 @@ static void _top_p_row_core(float* __restrict__ row, int V, float p_val, { const float inv_sum = 1.0f / sum_exp; int i = 0; -#if defined(__aarch64__) +#if defined(__aarch64__) || defined(__s390x__) vec_op::FP32Vec16 isum16(inv_sum); for (; i + 16 <= V; i += 16) { vec_op::FP32Vec16 v16(row + i); vec_op::FP32Vec16 p16 = fast_exp(v16 - base16) * isum16; + #if defined(__s390x__) + p16 = zero_pad_lanes(v16, p16); + #endif p16.save(scratch + i); } #endif