Feat/precomputed noise pool

[GOavi101]

Description

• Port Holtz async Exp(1) log-noise ring buffer from sendnn-inference#1046: background thread fills noise; host sampling consumes from the pool.
• Sample on TP rank 0 and broadcast token ids; use log-space Gumbel (argmax(logits - log_noise)) so the hot path avoids softmax.
• Wire via build_spyre_sampler in the Spyre model runner; document host-sampler config; add unit tests for the ring buffer, sampler, and top-k/top-p path.

docs in docs/user_guide/configuration.md.

Related Issues

Relates to host sampling performance / sendnn-style Exp(1) pool reuse.

Test Plan

• Unit tests: pytest tests/test_spyre_noise_pool.py -q (pool fill, slice, seeded path, opt-in wiring)
• Manual (Spyre host), compare sampler timing with pool on vs off:

  SPYRE_USE_NOISE_POOL=1 SPYRE_SAMPLER_TIMING=50 \
    python examples/offline_inference/torch_spyre_inference.py
  # use temperature > 0; greedy does not exercise the pool

Checklist

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

[github-actions]

👋 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.

[github-actions]

📖 This PR includes doc changes. Download the preview (built from 2f8165a), or build locally:

uv pip install -r docs/requirements-docs.txt
mkdocs serve

[joerunde]

a few thoughts:

1. SPYRE_SAMPLER_TIMING smells like something that we should be using the profiler for instead of injecting into a production code path
2. We need to do config properly here- at the minimum the envs.py approach that puts all the config levers in one place with documentation, caching, and override semantics is required.
3. This feels like running before we crawl, it's an optimization for noise generation for sampling on a specific cpu platform but we haven't even done the work yet to put the rest of the sampling into a torch.compiled path on the spyre devices. The model / sampler split on sendnn-inference was a limitation that we couldn't work around that no longer exists here

[GOavi101]

Agree on sequencing. this is only an opt-in host stopgap for slow Exp(1) on s390x while sampling is still on CPU — it doesn’t move sampling onto Spyre.

for that few pain point i have analyzed:

1. argmax— CPU fallback
2. topk— native but unusable (k≤4, fp16 indices wrong past 2048)
3. cumsum— CPU fallback (top-p)
4. sort— unsupported (top-p / joint top-k+top-p)
5. scatter / scatter_add / gather / index_put_— incomplete/broken (top-p, penalties, logprobs)
6. exponential/ RNG— CPU fallback (temp / random sample)
   CC: @dilipgb @rishikakedia

[tdoublep]

I am concerned that this approach will change the sampling statistics in strange and ways that are hard to interpret. I would prefer to hold off on merging.

[dmholtz]

As the scope of this PR changed towards the direction of torch-spyre/sendnn-inference#1046:

I agreed with @tdoublep to provide an analogous solution to sendnn-inference here; however I think PR 1046 should not be mirrored 1:1 here. For example, the softmax function could be properly computed on device with this stack. There is currently no rush to address the sampler before batching support lands, so I will provide a draft of that idea in about two weeks.