[Engine] OpenAI logprobs / top_logprobs / allowed_token_ids for the CUDA/HIP engine

CMakeLists.txt

@@ -527,6 +527,10 @@ if(STRATA_ENABLE_CUDA OR STRATA_ENABLE_HIP)
   endif()
   # plan v0.3 P6: llama.cpp's block layouts and codebook grids for the i-quant formats (MIT, third_party/ggml)
   target_include_directories(strata_kernels PRIVATE third_party/ggml)
+  # gfx906 uses 64-lane waves; this reduction's fixed 32-lane groups are for CUDA and wave32 HIP.
+  if(NOT STRATA_HIP_GFX906)
+    target_sources(strata_kernels PRIVATE src/kernels/cuda/logprobs.cu)
+  endif()
   # Match the pinned llama.cpp CUDA arithmetic contract for this adapter only.
   if(STRATA_ENABLE_CUDA)
     set_source_files_properties(src/kernels/cuda/native_mmvq.cu src/kernels/cuda/native_bf16.cu
@@ -1261,6 +1265,13 @@ if(STRATA_ENABLE_CUDA OR STRATA_ENABLE_HIP)
   target_link_libraries(sampler_parity PRIVATE strata_kernels ${_strata_gpu_runtime_target})
   add_test(NAME sampler_parity COMMAND sampler_parity --selftest)
   set_tests_properties(sampler_parity PROPERTIES ENVIRONMENT "STRATA_OLD_SAMPLER=0;STRATA_SAMPLER_ONE_BLOCK=0")
+  if(NOT STRATA_HIP_GFX906)
+    add_executable(logprobs_parity src/kernels/logprobs_parity.cpp)
+    target_include_directories(logprobs_parity PRIVATE ${_strata_gpu_include_directories})
+    target_link_libraries(logprobs_parity PRIVATE strata_kernels ${_strata_gpu_runtime_target})
+    add_test(NAME logprobs_parity COMMAND logprobs_parity)
+    set_tests_properties(logprobs_parity PROPERTIES SKIP_RETURN_CODE 77)
+  endif()
   # the GPU verifier of probabilistic draft acceptance against the host reference (STRATA_SPEC_PROB)
   add_executable(spec_verify_parity src/kernels/spec_verify_parity.cpp)
   target_include_directories(spec_verify_parity PRIVATE ${_strata_gpu_include_directories})

docs/DETAILS.md

@@ -505,6 +505,24 @@ applies at the next start. Network, key, MCP and program keys are not editable t
   `STRATA_SPEC_GUMBEL=1` raise the accepted-draft rate on some machines; `STRATA_SPEC_PROB=1` is speculative rejection
   sampling. `STRATA_SPEC_DEPTH=1` prints acceptance by draft depth and `STRATA_SPEC_MIN_TEMP=0.9` limits sampled drafting
   to hotter requests. Try them on your own card before relying on them.
+- **logprobs: how likely each token was (CUDA/HIP).** `"logprobs": true` with `"top_logprobs": 0..20` on
+  `/v1/chat/completions`, streamed or not, adds `choices[0].logprobs.content[]` to the reply (`token`, `logprob`,
+  `bytes`, `top_logprobs`); with thinking on, the thinking's entries come under `reasoning_content` instead. The
+  tokens, joined, spell the reply exactly (a character split across two tokens lands in the one that finishes it).
+  A reply is bounded by `max_tokens` -- or by the rest of the context when it is not sent -- and by no cap of our
+  own, as on OpenAI and vLLM, neither of which ties a length limit to logprobs. A reply that is not streamed, or
+  that answers a `json_schema` (buffered to be validated), does hold every entry until it ends, so a long one holds
+  a lot of them; a streamed reply hands each entry out with its chunk and keeps none. `"allowed_token_ids": [ids]`
+  keeps only those ids in `top_logprobs`, most likely first, renormalised to sum to 1; the chosen token is one of
+  them when it was named (vLLM's field, but vLLM masks generation with it, while here it filters what is reported,
+  so a token outside the named ids keeps its own `logprob` and an empty list). An id past the model's vocabulary is
+  a 400. A token whose probability could not be read gets no entry, and the reply goes on. On a server started with
+  `--parallel`, only a `max_tokens` 1 request gets them: it is the one a batch slot decodes through the single-token
+  path, which writes an LP line. A longer reply there is a 400, since the slot's `BT` lines carry none -- entries for
+  a prefix only, which a client cannot tell from the model's own answer.
+  In serial CUDA and wave32 HIP decode, the engine scores accepted tokens and top IDs on the GPU and transfers
+  those scores. Equal-logit top IDs are ordered by token ID; score rounding can differ from the CPU path. Named
+  `allowed_token_ids` and pipelined windows keep the full-logit CPU scoring path.
 - **Non-finite logits end the request and drop the cached state (on by default).** If a verify window's logits are NaN, the
   window emits no tokens, the request ends, and the next request drops the live session, checkpoints and parked
   conversations and reads its prompt from token 0. The server retries a request that had sent nothing once, transparently.

include/strata/core/verify.hpp

@@ -27,6 +27,7 @@
 #include "strata/core/layer.hpp"
 #include "strata/core/session.hpp"
 #include "strata/kernels/sampler.hpp"
+#include "strata/kernels/logprobs.hpp"
 
 #include <cuda_runtime.h>
 
@@ -105,6 +106,12 @@ class Verifier {
     bool run(int T, const int32_t* tokens, int64_t pos0, PoolMultiFn pool, void* user, int32_t* out, std::string& err);
     /// Diagnostics: row `t` of the last window's head logits (n_vocab floats) to the host. Valid after run().
     bool copy_logits(int t, float* host) const;
+#if !defined(STRATA_GPU_METAL) && !defined(STRATA_HIP_GFX906)
+    /// Raw head-softmax scores for accepted tokens in the first `rows` of the last window.
+    /// Returns one chosen token plus `top_k` alternatives per row, before the next window overwrites logits.
+    bool logprobs(int rows, const int32_t* chosen, int top_k, strata::kernels::LogprobItem* out,
+                  std::string& err);
+#endif
     int64_t vocab() const { return next_ ? next_->vocab() : n_vocab_; }
     /// #879 GUARD: the first row of the last finished window (run, run_slot_rows, batch_poll, pl_finish) whose head
     /// logits held a NaN or an infinity, else -1; nonfinite_pos() is that row's position.  Every window graph
@@ -645,6 +652,11 @@ class Verifier {
     float* hit_xs_ = nullptr;
     void* hit_scratch_ = nullptr;
     float *head_mixed_ = nullptr, *head_inj_ = nullptr, *head_logits_ = nullptr;
+#if !defined(STRATA_GPU_METAL) && !defined(STRATA_HIP_GFX906)
+    void* logprobs_scratch_ = nullptr;
+    int32_t* logprobs_chosen_ = nullptr;
+    strata::kernels::LogprobItem* logprobs_out_ = nullptr;
+#endif
     uint16_t* sh_bf16_ = nullptr;
     uint8_t* arg_scratch_ = nullptr;   ///< argmax_rows' partials and counters
     int32_t* one_ = nullptr;           ///< device {1}: the n_keep of a one-token window, which commits itself

include/strata/kernels/logprobs.hpp

@@ -0,0 +1,29 @@
+#pragma once
+
+#include <cstddef>
+#include <cstdint>
+
+namespace strata::kernels {
+
+// One chosen token followed by K highest-scoring alternatives for each row.
+// Scores are natural log probabilities under the unmodified, full-vocabulary
+// softmax of the head logits. The chosen token need not be in the top K.
+struct LogprobItem {
+    int32_t id;
+    float logprob;
+};
+
+constexpr int kLogprobsMaxTop = 20;
+
+// Device scratch for logprobs_extract. The head vocabulary must contain at most
+// 262144 entries, as in the split sampler. Returns zero for unsupported sizes.
+size_t logprobs_scratch_bytes(int rows, int vocab, int top_k);
+
+// All pointers, including chosen and out, are device pointers. Scratch is at
+// least logprobs_scratch_bytes(rows, vocab, top_k) bytes. The caller preserves
+// logits until the stream has finished, and transfers only rows*(K+1) items.
+// Returns false without launching for invalid dimensions.
+bool logprobs_extract(const float* logits, int rows, int vocab, const int32_t* chosen,
+                      int top_k, void* scratch, LogprobItem* out, void* stream);
+
+}  // namespace strata::kernels