Slot penalties and logit bias in batch windows: shared slot_row_draw for CUDA/HIP/Metal, server accepts logit_bias with --batch, docs updated

docs/BATCHING.md

@@ -631,7 +631,6 @@ counter-based draw (Philox(seed, position)).
   with `--batch-groups` above 1 (pipelined slot groups run no drafts, #1413) unless `STRATA_PIPED_MTP` is set with
   one slot per group ([above](#mtp-drafts-in-the-pipelined-groups-strata_piped_mtp-opt-in)), and has not been run
   with helper expert caches.
-- Repetition / frequency / presence penalties are not applied in batch windows.
 - A prompt shorter than one chunk is read in one piece (the slots wait for it); a read gives way only at a chunk
   boundary, and not for pictures.
 - Admissions are one at a time: two new long prompts are read one after the other.

docs/DETAILS.md

@@ -536,7 +536,7 @@ applies at the next start. Network, key, MCP and program keys are not editable t
   only after the strict parse failed, the server doubles the backslashes that are not valid escapes and checks the answer again.
   Valid output is never touched. `\b`, `\f`, `\n`, `\r`, `\t` and `\uXXXX` are valid JSON, so they stay: a `\beta` beside a `\alpha`
   comes back with a backspace. Off by default.
-- `logit_bias` and token/bias pair lists work on CUDA and HIP engines without continuous batching ([LOGIT_BIAS.md](LOGIT_BIAS.md)).
+- `logit_bias` and token/bias pair lists work on CUDA, HIP and Metal engines, also with continuous batching (`--batch` / `--parallel`: each slot keeps its own bias) ([LOGIT_BIAS.md](LOGIT_BIAS.md)).
 - Stop strings (`stop`, `stop_sequences`), `temperature`, `top_p`, `top_k`, `min_p` and `seed` are honored per request.
   `top_k` keeps at most 64 candidates. Penalties (`presence_penalty`, `frequency_penalty`, `repetition_penalty`,
   `penalty_last_n`; llama.cpp's `repeat_penalty` and `repeat_last_n` work too) apply to every token the speculative check

docs/LOGIT_BIAS.md

@@ -1,7 +1,7 @@
 # Per-request token biases
 
-`POST /v1/chat/completions` accepts `logit_bias` with a supporting CUDA or HIP
-engine, without continuous batching. The same bias applies to every generated
+`POST /v1/chat/completions` accepts `logit_bias` with a supporting CUDA, HIP or Metal
+engine, also with continuous batching (each slot keeps its own bias). The same bias applies to every generated
 token in the request, including tokens checked together during speculative
 decoding. It does not force single-token verification windows.
 
@@ -65,8 +65,8 @@ vector is uploaded when the bias changes and reused within the request.
 Malformed entries, duplicate IDs in pair lists, IDs outside the vocabulary,
 out-of-range biases, or banning the entire vocabulary return HTTP 400. Validation
 on the chat route occurs before SSE streaming headers are sent. Nonempty biases
-also return HTTP 400 with older engines, the current SYCL engine, or continuous
-batching (`--batch`), rather than being silently ignored. The engine advertises
+also return HTTP 400 with older engines or the current SYCL engine, rather than
+being silently ignored. The engine advertises
 `logit_bias=1` in its `INFO` response when this implementation is present.
 
 MTP target verification applies the bias; the draft model is not biased, so a

include/strata/core/verify.hpp

@@ -462,7 +462,8 @@ private:
     std::vector<std::vector<int32_t>> slot_tail_;        ///< per slot: its reply tail for the penalty window (set_slot_history)
     int32_t* slot_hist_d_ = nullptr;                     ///< the penalty-history rows of the current batch window (device)
     size_t slot_hist_cap_ = 0;
-    bool sample_rows(int S, std::string& err);   ///< the sampled slots' rows of the last batch window
+    bool sample_rows(int S, std::string& err);
+    int slot_row_draw(int S, int t, strata::kernels::SamplerParams& sp, const int32_t*& hist, int& hlen, std::vector<int32_t>& hrows, std::string& err);   ///< the sampled slots' rows of the last batch window
     int32_t* h_commitb_ = nullptr; int32_t* m_commitb_ = nullptr;   // per slot [1, 0, pos, -1 ..], stride 2 + max_t
     int32_t* commitb_ = nullptr;
     float* tail_snap_b_ = nullptr;         ///< per (slot, QSA layer) indexer tail snapshot

metal/src/engine/window.cpp

@@ -194,7 +194,7 @@ public:
     // history, Tn rows of hist_len tokens (the last penalty_last_n of a row count, -1 = unused), read when the window's rows are sampled.
     void set_sampling(const strata::kernels::SamplerParams* sp, const int32_t* hist, int hist_len);
     bool sample_rows(int Tn, int32_t* out, std::string& err);
-    bool sample_slot_row(int row, const strata::kernels::SamplerParams& sp, int32_t* tok, std::string& err);   // one slot's row of a batch window, its own params
+    bool sample_slot_row(int row, const strata::kernels::SamplerParams& sp, const int32_t* hist, int hist_len, int32_t* tok, std::string& err);   // one slot's row of a batch window, its own params
     bool sample_run(int row0, int n, const strata::kernels::SamplerParams& sp, const int32_t* hist, int hl, std::string& err);   // after run(): replaces out[0..Tn) (and tokout_) by the sampled tokens
     bool copy_logits(int t, float* host) const;
     int64_t vocab() const { return n_vocab_; }
@@ -3097,9 +3097,9 @@ bool Window::sample_rows(int Tn, int32_t* out, std::string& err) {
     return true;
 }
 
-bool Window::sample_slot_row(int row, const strata::kernels::SamplerParams& sp, int32_t* tok, std::string& err) {
+bool Window::sample_slot_row(int row, const strata::kernels::SamplerParams& sp, const int32_t* hist, int hist_len, int32_t* tok, std::string& err) {
     if (row < 0 || row >= MAXT) return true;
-    if (!sample_run(row, 1, sp, nullptr, 0, err)) return false;
+    if (!sample_run(row, 1, sp, hist, hist_len, err)) return false;
     *tok = tokout_.as<int32_t>()[row];
     return true;
 }
@@ -3219,8 +3219,8 @@ bool window_ensure(const strata::core::ModelGeometry& g, strata::core::SessionSt
     int64_t nv = 0;
     return window_init(window_create(), g, ss, MAXT, max_cells, &nv, err);
 }
-bool window_sample_slot_row(void* w, int row, const strata::kernels::SamplerParams& sp, int32_t* tok, std::string& err) {
-    return ((Window*) w)->sample_slot_row(row, sp, tok, err);
+bool window_sample_slot_row(void* w, int row, const strata::kernels::SamplerParams& sp, const int32_t* hist, int hist_len, int32_t* tok, std::string& err) {
+    return ((Window*) w)->sample_slot_row(row, sp, hist, hist_len, tok, err);
 }
 void window_set_sampling(void* w, const strata::kernels::SamplerParams* sp, const int32_t* hist, int hist_len) {
     ((Window*) w)->set_sampling(sp, hist, hist_len);

serve/logit_bias.py

@@ -103,11 +103,9 @@ def engine_key(value):
 
 
 def validate_request(req, engine, vocab_size):
-    """Fail closed for old engines, unsupported backends and continuous batching."""
+    """Fail closed for old engines and unsupported backends (continuous batching carries a slot's bias since the slot-penalty change)."""
     normalized = normalize(req.get("logit_bias"), vocab_size)
     if normalized:
         if getattr(engine, "info", {}).get("logit_bias") != 1:
             raise ValueError("logit_bias: this engine does not support token biases; use a supporting CUDA/HIP build")
-        if getattr(engine, "batch", 0):
-            raise ValueError("logit_bias: continuous batching is not supported; run without --batch")
     return normalized

serve/test_logit_bias.py

@@ -34,10 +34,12 @@ class LogitBias(unittest.TestCase):
         self.assertEqual(engine_key(entries).count(':-100'), 103215)
 
     def test_capability_and_batch(self):
-        for engine in [SimpleNamespace(info={}), SimpleNamespace(info={'logit_bias': 1}, batch=2)]:
+        for engine in [SimpleNamespace(info={})]:
             with self.assertRaises(ValueError):
                 validate_request({'logit_bias': {'1': -100}}, engine, 10)
             self.assertEqual(validate_request({}, engine, 10), {})
+        # --batch: a slot keeps its request's own bias (Verifier::set_slot_logit_bias)
+        self.assertEqual(validate_request({'logit_bias': {'1': -100}}, SimpleNamespace(info={'logit_bias': 1}, batch=2), 10), {1: -100})
         self.assertEqual(validate_request({'logit_bias': {'1': -100}}, SimpleNamespace(info={'logit_bias': 1}), 10), {1: -100})
 
     def test_http_rejects_before_stream(self):

src/core/verify.cpp

@@ -98,7 +98,7 @@ bool window_run(void* w, int T, const int32_t* tokens, int64_t pos0, bool head,
 bool window_commit(void* w, int n_keep, std::string& err);
 bool window_copy_logits(void* w, int t, float* host);
 void window_set_sampling(void* w, const strata::kernels::SamplerParams* sp, const int32_t* hist, int hist_len);
-bool window_sample_slot_row(void* w, int row, const strata::kernels::SamplerParams& sp, int32_t* tok, std::string& err);
+bool window_sample_slot_row(void* w, int row, const strata::kernels::SamplerParams& sp, const int32_t* hist, int hist_len, int32_t* tok, std::string& err);
 bool window_init_slots(void* w, const std::vector<strata::core::SessionState*>& slots, std::string& err);
 bool window_run_slots(void* w, const int* rows, int S, const int32_t* tokens, const int64_t* pos, int32_t* out, std::string& err);
 }  // namespace strata::metal_engine
@@ -4562,18 +4562,23 @@ bool Verifier::run_slot_rows(const int* rows, int S, const int32_t* tokens, cons
     if (S < 1 || S > (int) strata::kernels::kVerifyMaxT) { err = "verify: batch window size out of range"; return false; }
     (void) pool; (void) user;
     if (!strata::metal_engine::window_run_slots(metal_, rows, S, tokens, pos, out, err)) return false;
-    // a slot with its own sampling (set_slot_sampling: temperature / top_p / top_k / min_p / seed; penalties are not applied in batch windows): its row is
-    // drawn again from the window's logits with Philox(seed, position), as a solo window draws it
-    for (int t = 0; t < S; ++t) {
-        strata::kernels::SamplerParams sp = slot_sp_[(size_t) rows[t]];
-        if (sp.greedy || sp.temperature <= 0.0f) continue;
-        sp.counter = (uint64_t) pos[t];
-        sp.penalty_last_n = 0;
-        if (!strata::metal_engine::window_sample_slot_row(metal_, t, sp, &out[t], err)) return false;
-    }
     last_t_ = S;
     for (int t = 0; t < S; ++t) { last_rows_[t] = rows[t]; last_tokens_[t] = tokens[t]; last_pos_b_[t] = pos[t]; }
     last_batch_ = true;
+    // a slot with its own sampling, penalties or logit bias (set_slot_sampling / set_slot_history / set_slot_logit_bias): its row is drawn again from the
+    // window's logits with Philox(seed, position), as a solo window draws it; a plain greedy slot keeps the window's argmax token
+    {
+        std::vector<int32_t> hrows;
+        for (int t = 0; t < S; ++t) {
+            strata::kernels::SamplerParams sp;
+            const int32_t* hist = nullptr;
+            int hlen = 0;
+            const int r = slot_row_draw(S, t, sp, hist, hlen, hrows, err);
+            if (r < 0) return false;
+            if (r == 0) continue;
+            if (!strata::metal_engine::window_sample_slot_row(metal_, t, sp, hist, hlen, &out[t], err)) return false;
+        }
+    }
     return true;
 #endif
     using namespace strata::kernels;
@@ -4751,46 +4756,66 @@ bool Verifier::set_slot_logit_bias(int slot, const std::vector<float>& bias, std
     return true;
 }
 
-bool Verifier::sample_rows(int S, std::string& err) {
-    bool any = false;
-    // the slots' penalty-history rows of this window (host), built once for the rows that need one
+// One batch row's draw inputs: 0 = the row keeps the window's argmax token (a plain greedy slot), 1 = draw it with `sp` (+ `hist` / `hlen`), -1 = error.
+// Row t of the last window feeds last_tokens_[t] into slot last_rows_[t]; the penalty-history row is the slot's reply tail plus its tokens in this window up to t
+// (slot_penalty_rows), staged on the device like a solo window's.
+int Verifier::slot_row_draw(int S, int t, strata::kernels::SamplerParams& sp, const int32_t*& hist, int& hlen, std::vector<int32_t>& hrows,
+                            std::string& err) {
     constexpr int kSlotHistMax = 4096;   // serve's kPenaltyWindowCap
-    std::vector<int32_t> hrows;
+    sp = slot_sp_[(size_t) last_rows_[t]];
+    const bool pen = sp.penalty_last_n > 0 && (sp.penalty_repeat != 1.0f || sp.penalty_freq != 0.0f || sp.penalty_present != 0.0f);
+    const bool bias = sp.logit_bias != nullptr;
+    hist = nullptr;
+    hlen = 0;
+    if ((sp.greedy || sp.temperature <= 0.0f) && !pen && !bias) return 0;   // a plain greedy slot keeps the window's argmax token untouched
+    sp.counter = (uint64_t) last_pos_b_[t];   // Philox(seed, position): the solo window's draw for this position
+    if (!pen) { sp.penalty_last_n = 0; return 1; }
     std::vector<const int32_t*> tails(slots_.size(), nullptr);
     std::vector<int64_t> n_tails(slots_.size(), 0);
     for (size_t b = 0; b < slots_.size() && b < slot_tail_.size(); ++b) { tails[b] = slot_tail_[b].data(); n_tails[b] = (int64_t) slot_tail_[b].size(); }
+    hlen = std::min(sp.penalty_last_n, kSlotHistMax);
+    const size_t need = (size_t) S * (size_t) hlen;
+    hrows.assign(need, -1);
+    strata::kernels::slot_penalty_rows(last_rows_, last_tokens_, S, tails.data(), n_tails.data(), hlen, hrows.data());
+    const OnDevice on_hist(device_);
+    if (slot_hist_cap_ < need) {
+        if (slot_hist_d_ != nullptr) cudaFree(slot_hist_d_);
+        slot_hist_d_ = nullptr;
+        if (cudaMalloc((void**) &slot_hist_d_, need * sizeof(int32_t)) != cudaSuccess) { err = "verify batch: the penalty-history allocation failed"; return -1; }
+        slot_hist_cap_ = need;
+    }
+#if defined(STRATA_GPU_METAL)
+    if (cudaMemcpy(slot_hist_d_ + (size_t) t * (size_t) hlen, hrows.data() + (size_t) t * (size_t) hlen, (size_t) hlen * sizeof(int32_t), cudaMemcpyHostToDevice) != cudaSuccess)
+#else
+    if (cudaMemcpyAsync(slot_hist_d_ + (size_t) t * (size_t) hlen, hrows.data() + (size_t) t * (size_t) hlen, (size_t) hlen * sizeof(int32_t),
+                        cudaMemcpyHostToDevice, cs_) != cudaSuccess)
+#endif
+    { err = "verify batch: the penalty-history copy failed"; return -1; }
+    hist = slot_hist_d_ + (size_t) t * (size_t) hlen;
+    return 1;
+}
+
+bool Verifier::sample_rows(int S, std::string& err) {
+#if defined(STRATA_GPU_METAL)
+    (void) S;
+    err = "verify batch: sample_rows is the CUDA / HIP path";
+    return false;
+#else
+    bool any = false;
+    std::vector<int32_t> hrows;
     for (int t = 0; t < S; ++t) {
-        strata::kernels::SamplerParams sp = slot_sp_[(size_t) last_rows_[t]];
-        const bool pen = sp.penalty_last_n > 0 && (sp.penalty_repeat != 1.0f || sp.penalty_freq != 0.0f || sp.penalty_present != 0.0f);
-        const bool bias = sp.logit_bias != nullptr;
-        if ((sp.greedy || sp.temperature <= 0.0f) && !pen && !bias) continue;   // a plain greedy slot keeps the window's argmax token untouched
-        sp.counter = (uint64_t) last_pos_b_[t];   // Philox(seed, position): the solo window's draw for this position
+        strata::kernels::SamplerParams sp;
         const int32_t* hist = nullptr;
         int hlen = 0;
-        if (pen) {
-            hlen = std::min(sp.penalty_last_n, kSlotHistMax);
-            hrows.assign((size_t) S * (size_t) hlen, -1);
-            strata::kernels::slot_penalty_rows(last_rows_, last_tokens_, S, tails.data(), n_tails.data(), hlen, hrows.data());
-            const size_t need = (size_t) S * (size_t) hlen;
-            const OnDevice on_hist(device_);
-            if (slot_hist_cap_ < need) {
-                if (slot_hist_d_ != nullptr) cudaFree(slot_hist_d_);
-                slot_hist_d_ = nullptr;
-                if (cudaMalloc((void**) &slot_hist_d_, need * sizeof(int32_t)) != cudaSuccess) { err = "verify batch: the penalty-history allocation failed"; return false; }
-                slot_hist_cap_ = need;
-            }
-            if (cudaMemcpyAsync(slot_hist_d_ + (size_t) t * (size_t) hlen, hrows.data() + (size_t) t * (size_t) hlen, (size_t) hlen * sizeof(int32_t),
-                                cudaMemcpyHostToDevice, cs_) != cudaSuccess) { err = "verify batch: the penalty-history copy failed"; return false; }
-            hist = slot_hist_d_ + (size_t) t * (size_t) hlen;
-        } else {
-            sp.penalty_last_n = 0;
-        }
-        strata::kernels::sample_tokens(head_logits_ + (size_t) t * (size_t) n_vocab_, 1, (int) n_vocab_, hist, hlen, sp,
-                                       m_out_ + t, cs_);
+        const int r = slot_row_draw(S, t, sp, hist, hlen, hrows, err);
+        if (r < 0) return false;
+        if (r == 0) continue;
+        strata::kernels::sample_tokens(head_logits_ + (size_t) t * (size_t) n_vocab_, 1, (int) n_vocab_, hist, hlen, sp, m_out_ + t, cs_);
         any = true;
     }
     if (any && cudaStreamSynchronize(cs_) != cudaSuccess) { err = "verify batch: the row sampling failed"; return false; }
     return true;
+#endif
 }
 
 bool Verifier::batch_launch(int base, int S, const int32_t* tokens, const int64_t* pos, std::string& err) {