Skip to content

Commit 126e3a6

Browse files
alexmalyshevmeta-codesync[bot]
authored andcommitted
Add inference_pipeline CPU serving-stage benchmark
Summary: Add `inference_pipeline`, a benchmark of CPU-bound stages from real PyTorch inference pipelines -- the work done on the host CPU, around the model, to transform its inputs and outputs. Each workload is a deliberately pure-Python implementation of a stage that production often runs in native code; the pure-Python form still occurs in real systems (CPU-only or dependency-constrained serving, on-device, prototypes, and the reference implementations the native libraries were derived from), and isolating it is what lets the benchmark exercise the interpreter/JIT. These implementations are influenced by code using PyTorch, but they themselves do not use PyTorch. Workloads: - `nms`: object-detection post-processing -- anchor box decode + greedy non-max suppression (vs the C++/CUDA `torchvision.ops.nms` kernel). - `bpe`: byte-level byte-pair-encoding tokenization, GPT-2-style with a per-request cache and real token ids (vs Rust `tokenizers`/`tiktoken`). - `beam`: beam-search decode bookkeeping mirroring HuggingFace `BeamSearchScorer` (vs decode fused inside vLLM/TensorRT-LLM). - `ngram`: no-repeat-ngram blocking + repetition penalty mirroring HuggingFace logits processors (vs fairseq's fused CUDA `NGramRepeatBlock`). - `ctc`: CTC greedy decoding mirroring torchaudio (vs the C++ flashlight beam decoder). - `rec`: recommendation feature preprocessing -- Sigrid-style hashing, bucketization, and KeyedJaggedTensor assembly mirroring the TorchRec transforms (host-side Python glue with no single fused native op). By default the binary runs all workloads and prints a per-workload summary; pass `--workload <name>` to run a single stage. The harness mirrors `torchbench.py` (`--cinderx`, warmup/repeat timing, and a `--compare` mode that re-execs an interpreter baseline against the JIT in separate subprocesses); under `--compare` it prints a per-workload speedup table and the geomean. Measured CinderX JIT speedups (`--compare`, vs the `CINDERX_DISABLE=1` interpreter baseline): workload baseline jit speedup nms 17.318 13.276 1.30x bpe 1.635 1.817 0.90x beam 2.398 2.079 1.15x ngram 1.173 1.125 1.04x ctc 0.546 0.465 1.17x rec 7.340 5.237 1.40x geomean 1.15x The spread is the point: `rec` (no native fast path) is the biggest win and `bpe` regresses, with the rest modest -- genuine CinderX signals (the bpe merge loop's exception-driven control flow, and the tuple/dict hashing + heapq bookkeeping in beam/ngram, are not strong JIT wins). Reviewed By: yoney Differential Revision: D109588862 fbshipit-source-id: 6777b70cab1b2a3fa0cf117d3399a1d8149e4632
1 parent 3320dbb commit 126e3a6

1 file changed

Lines changed: 897 additions & 0 deletions

File tree

0 commit comments

Comments
 (0)