Commit 126e3a6
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: 6777b70cab1b2a3fa0cf117d3399a1d8149e46321 parent 3320dbb commit 126e3a6
1 file changed
Lines changed: 897 additions & 0 deletions
0 commit comments