[Paged KV] Add paged attention deterministic smoke test

[WindChimeRan]

Summary

• Add test_paged_deterministic.py: 5-prompt smoke test using vLLM offline inference (temp=0, greedy) against hardcoded golden token IDs from Qwen3-0.6B
• Golden values generated on main from both MLX inline cache and HF paged KV cache paths
• Add tools/gen_golden.py helper to regenerate golden values

Motivation

Prerequisite for the native Metal kernel PR (#136). After inlining the vendored Metal shaders, paged attention output must remain identical to the current HF kernel baseline. This test anchors that.

Test

• python -m pytest tests/test_paged_deterministic.py -v -s (paged path by default)
• Passes on main with HF kernel: 5/5

Relevant Issue & PR

• Issue Metal paged-attention parity mismatch vs standard path #119
• PR [Paged KV] Inline metal kernel. deprecate hf pytorch kernel #136 : This inline metal kernel need to either pass this test, or explain the possible non-deterministics from the kernel.

upstream batch invariant feature

• blog: https://thinkingmachines.ai/blog/defeating-nondeterminism-in-llm-inference/
• main feature: [Feature]: Batch Invariant Feature and Performance Optimization vllm#27433
• vllm upstream batch invariant feature is only compatible with H / B series NVIDIA GPU. A100 not working. See my exp results https://github.com/WindChimeRan/spec_deterministic
• community work: [Feature] Batch-Invariant Support for FA2 and LoRA vllm#30018

Batch invariant is hardware & kernel dependent. Supporting this feature is non-trivial on metal.

output example:

[WindChimeRan]

@LxYuan0420 request for review

[LxYuan0420]

Nice direction and the tools/ helper is a good addition.

To recap the main issue I ran into locally: on macOS 15 (HF compat kernel 8968951), this test can flake under the current batched llm.generate(PROMPTS, ...) setup. Specifically, greedy output for One plus one equals sometimes diverges from both golden sets.

For CI stability, I think we should make the execution mode explicit. Would you be open to forcing single-seq execution (LLM(..., max_num_seqs=1)) for this test?

If your intent is to also cover batched behavior, maybe we can add a separate batched smoke test with a weaker assertion (since exact token IDs aren’t batch-invariant on Metal today)

[LxYuan0420]

On my machine, I need to set max_seq_nums=1 to pass the test

EDIT Typo: it is max_num_seqs=1 (vLLM LLM(..., max_num_seqs=1)), not max_seq_nums

[WindChimeRan]

On my machine, I need to set max_seq_nums=1 to pass the test

EDIT Typo: it is max_num_seqs=1 `(vLLM LLM(..., max_num_seqs=1))`, not max_seq_nums

This is really intriguing... Maybe it's due to Metal 3 vs. Metal 4, or the 0.2 × RAM auto-calculated max sequence length, or differences in our chips.

[LxYuan0420]

Thanks for iterating on this. Forcing single-seq execution (LLM(..., max_num_seqs=1)) resolves the macOS 15 flake I hit with batched scheduling.

I also pushed a small maintainer tweak on top to (1) fix the generator script name in comments/usage and (2) make the autouse env fixture respect any user-provided env overrides so the MLX run instructions match actual behavior.

LGTM.

Re batch-invariant determinism: I’m not expecting that in this PR (I only wanted a CI-stable single-seq golden smoke test). Let’s defer any batch-invariant feature/test until we support real continuous batching on Metal.