feat(sampling): Add async ring buffer and Spyre-optimized samplers by dmholtz · Pull Request #1046 · torch-spyre/sendnn-inference

dmholtz · 2026-08-17T11:43:32Z

Description

This PR optimizes the Spyre sampling path in three stages to remove the main latency bottlenecks from token generation while preserving correctness.
To ease the review, each stage is a separate commit.

Problem

The default sampling path spends too much time in the critical decoding loop on CPU-heavy work:

Exponential noise generation is computed synchronously in the sampling path.
Softmax normalization remains expensive for large vocabularies.
Tensor-parallel CPU ranks redundantly sample and diverge from one another.

Solution

Stage 1: Async noise pre-sampling
- Introduce an asynchronous ring buffer that pre-generates exponential noise on a background thread.
- Replace synchronous exponential_() calls in the critical path with zero-copy buffer access.
- Add SpyreTopKTopPSampler and SpyreSampler integration for the Spyre path.
Stage 2: Single-rank sampling on TP
- Only TP rank 0 performs sampling and broadcasts the selected token IDs to the remaining ranks.
- This removes redundant CPU work and keeps all ranks aligned on the same tokens.
Stage 3: Log-space Gumbel sampling
- Reformulate sampling as argmax(logits - log_noise) instead of argmax(log(softmax(logits)) + log_noise).
- This preserves sampling behavior while removing softmax from the critical path.

Key Insight

The three optimizations are complementary: async noise generation removes expensive background work from the hot loop, TP rank gating removes redundant computation and divergence, and log-space Gumbel sampling removes softmax entirely. Together they reduce sampling latency without changing output distribution.
Performance tests show that OTPS increase by >30% for Granite 4.1 while staying in a fixed, thight compute budget.

Related Issues

n/a

Test Plan

Unit tests to verify ring-buffer correctness and sampling correctnes
Performance tests

Checklist

I have read the contributing guidelines
My code follows the project's code style (bash format.sh)
I have added tests for my changes
I have updated the documentation where needed
My commits include a Signed-off-by: line (DCO compliance)

github-actions · 2026-08-17T11:43:45Z

👋 Hi! Thank you for contributing.
Just a reminder: Make sure that your code passes all the linting checks, otherwise your PR won't be able to be merged. To do so, run ./format.sh.
Now you are good to go 🚀.

We also recommend installing prek and configuring it to check your code before every local commit.

dilipgb · 2026-08-28T14:16:22Z

Otherwise looks good to me. Pele tests passed fine. Maybe we wanted to run e2e tests on card once? @gkumbhat can you please take a look at this once?

dmholtz · 2026-08-28T18:27:11Z

Thank you for your review @dilipgb, let's run another round of tests to get this PR green

dilipgb · 2026-08-31T12:47:38Z

@joerunde would take look at this PR once?

tdoublep · 2026-08-31T14:08:37Z

Can we rebase the PR against the new version of vLLM to check there is no interaction?

Implement asynchronous exponential noise pre-sampling to optimize the latency-critical sampling path for Spyre hardware. Signed-off-by: David Holtz <56723830+dmholtz@users.noreply.github.com>

…ll ranks. Previously, the full sampling path was redundantly executed on every rank and the results of non-zero ranks were discarded. In the context of CPU bound sampling, this wastes CPU resources and can be streamlined by sampling only on TP rank 0, broadcasting results to all ranks to initialize the next step. Signed-off-by: David Holtz <56723830+dmholtz@users.noreply.github.com>

…rick. Currently, the softmax runs on CPU critical path, adding few milliseconds of runtime to every decode. This change rewrites the score computation by transfering the computation to log-space without changing the relative order of tokens. Co-authored-by: Jan Hofmeier <jan.hofmeier@de.ibm.com> Signed-off-by: David Holtz <56723830+dmholtz@users.noreply.github.com>

dmholtz · 2026-09-01T05:10:01Z

@tdoublep I rebased this PR to main, it went smooth without any conflicts. Could you please re-run the tests to confirm the behavior with vLLM v0.27.1?

tdoublep

LGTM - thanks for the awesome work!

tdoublep · 2026-09-01T09:39:16Z

@dilipgb Please take a final look (you requested changes, which is blocking merge). Tests are passing against latests main

R3hankhan123

1-2 nits rest it looks good to me

Signed-off-by: dmholtz <56723830+dmholtz@users.noreply.github.com>

dmholtz · 2026-09-09T05:46:50Z

@R3hankhan123 Thanks for your review, I addressed both NITs.

R3hankhan123

LGTM

dilipgb

LGTM

dmholtz force-pushed the stage1 branch from b45b7ba to 5bf03a4 Compare August 17, 2026 12:19

dmholtz marked this pull request as ready for review August 17, 2026 12:21

dmholtz requested review from prashantgupta24, rafvasq, sducouedic, tdoublep and yannicks1 as code owners August 17, 2026 12:21

dmholtz force-pushed the stage1 branch 2 times, most recently from 3bdd387 to 91e99a1 Compare August 18, 2026 05:54

dilipgb reviewed Aug 19, 2026

View reviewed changes

Comment thread tests/v1/sample/test_spyre_topk_topp_sampler.py

GOavi101 mentioned this pull request Aug 20, 2026

Feat/precomputed noise pool torch-spyre/spyre-inference#518

Draft

5 tasks

dilipgb requested changes Aug 28, 2026

View reviewed changes

dmholtz and others added 3 commits September 1, 2026 07:05

feat(sampling): Add async ring buffer and Spyre-optimized samplers

439d8c7

Implement asynchronous exponential noise pre-sampling to optimize the latency-critical sampling path for Spyre hardware. Signed-off-by: David Holtz <56723830+dmholtz@users.noreply.github.com>

dmholtz force-pushed the stage1 branch from 9127b6d to 60a2486 Compare September 1, 2026 05:07

tdoublep approved these changes Sep 1, 2026

View reviewed changes

R3hankhan123 reviewed Sep 4, 2026

View reviewed changes

Comment thread sendnn_inference/model_executor/model_loader/spyre.py Outdated

Comment thread sendnn_inference/v1/sample/spyre_topk_topp_sampler.py Outdated

dmholtz added 2 commits September 6, 2026 06:01

Replace warnings.warn with logger.warning in SpyreCausalLM

75f5288

Signed-off-by: dmholtz <56723830+dmholtz@users.noreply.github.com>

Address nit: Replace warnings.warn with logger.warning_once call

7c7f724

Signed-off-by: dmholtz <56723830+dmholtz@users.noreply.github.com>

R3hankhan123 approved these changes Sep 9, 2026

View reviewed changes

dilipgb approved these changes Sep 9, 2026

View reviewed changes

dilipgb merged commit f488d54 into torch-spyre:main Sep 9, 2026
16 checks passed

gkumbhat mentioned this pull request Sep 9, 2026

✨ Add precomputed noise pool for sampler #1045

Closed

5 tasks

Uh oh!

Conversation

dmholtz commented Aug 17, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Problem

Solution

Key Insight

Related Issues

Test Plan

Checklist

Uh oh!

github-actions Bot commented Aug 17, 2026

Uh oh!

Uh oh!

dilipgb commented Aug 28, 2026

Uh oh!

dmholtz commented Aug 28, 2026

Uh oh!

dilipgb commented Aug 31, 2026

Uh oh!

tdoublep commented Aug 31, 2026

Uh oh!

dmholtz commented Sep 1, 2026

Uh oh!

tdoublep left a comment

Choose a reason for hiding this comment

Uh oh!

tdoublep commented Sep 1, 2026

Uh oh!

R3hankhan123 left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

dmholtz commented Sep 9, 2026

Uh oh!

R3hankhan123 left a comment

Choose a reason for hiding this comment

Uh oh!

dilipgb left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

dmholtz commented Aug 17, 2026 •

edited

Loading