fix(sampling): treat top_k above the vocabulary as no top-k in the C++ chain #2247

Problem / Background

The stock sampling chain's top_k_filter (src/lib/mlxcel-core/cpp/mlx_cxx_bridge.cpp:6082-6097, origin/main 1b4e3657) calls mlx::core::argpartition(neg_x, top_k - 1, -1) without comparing top_k to the vocabulary size. Found while reviewing PR #2218. This is from reading the code; it has not been reproduced yet.

Current Behavior

• argpartition throws std::invalid_argument("[argpartition] Received invalid kth ...") when kth >= shape(axis) (src/lib/mlx-cpp/patches-rocm/mlx/ops.cpp:2873-2878, same check upstream). So top_k > vocab throws while the graph is being built. top_k == vocab is valid and keeps every token.
• fused_sample, fused_sample_xtc, fused_sample_categorical and fused_sample_probs are not Result bridge functions (src/lib/mlxcel-core/src/lib.rs:2503-2509, :2536, :2557), so the exception reaches cxx's noexcept boundary and calls std::terminate: the process aborts.
• The default GPU path reaches it. A sampled request with only top-k active (temperature > 0, top_k > vocab, top_p 1.0) skips Gumbel (top_k > 0). sampling_rejection_routes (:5858-5871) returns false when top-p is inactive, so fused_sample_impl (:6644-6762) falls through to fused_sample_argpartition_chain. XTC-active configs, fused_sample_categorical, MLXCEL_SAMPLING_REJECTION=0, the CPU device and backends without a rejection port take the chain unconditionally. With top-p active, top-k at or above vocab counts as inactive (:5869) and routes to the kernel, which is why this rarely shows.
• Rust does not validate first. The server passes the request's top_k through unchanged (src/server/request_options.rs:293), and fused_sample_dispatch (src/lib/mlxcel-core/src/sampling.rs:560-584) forwards config.top_k. The Rust-side filters already guard (sampling.rs:1235-1238, :1280-1284, :1557-1561 apply top_k_filter only when top_k < vocab). The C++ chain is the only unguarded call site.

Proposed Solution

In fused_sample_filter_logits, define use_top_k as top_k > 0 && top_k < x.shape(-1). That is the same "top-k at or above vocab is inactive" rule sampling_rejection_routes uses, and llama.cpp treats top_k >= n_vocab as no filtering. Leave the Rust reference top_k_filter (sampling.rs:2518) unchanged, since its callers guard, but add a one-line doc comment that k must be below the vocabulary. Request validation is not needed: an out-of-range top_k is a valid "no top-k" request.

Acceptance Criteria

• (PR fix(sampling): treat top_k above the vocabulary as no top-k in the chain #2254, in src/lib/mlxcel-core/src/sampling_top_k_vocab_tests.rs to keep sampling.rs from growing) A test in src/lib/mlxcel-core/src/sampling.rs runs ffi::fused_sample, ffi::fused_sample_categorical and ffi::fused_sample_probs on a [1, 64] row with top_k = 65 and top_k = 1_000_000 (temperature 0.7, top-p 1.0). None aborts, and the reported distribution equals the top_k = 0 distribution within 1e-6.
• (PR fix(sampling): treat top_k above the vocabulary as no top-k in the chain #2254) The same test with XTC active through fused_sample_xtc.
• (PR fix(sampling): treat top_k above the vocabulary as no top-k in the chain #2254, also a server-level tests/sampling_top_k_above_vocab_e2e.rs) Before the fix, the new test aborts the test binary (state this in the PR).

Verification

cargo test -p mlxcel-core sampling -- --test-threads=1
make verify

Related: PR #2218, #901 (rejection routing), #1379 (chain order).