decode.hpp

Go to the documentation of this file.

#pragma once
 
// SPDX-License-Identifier: LicenseRef-FSL-1.1-Apache-2.0
// Copyright 2026 Lloyal Labs
 
 
#include "common.hpp"
#include <algorithm>
#include <array>
#include <cstdint>
#include <llama/llama.h>
#include <memory>
#include <span>
#include <stdexcept>
#include <string>
#include <vector>
 
namespace lloyal::decode {
 
// Scratch::as_batch and as_embd_batch write llama_batch field by field; a new
// upstream field would arrive as zero without a word. Fail the build instead:
// the size on 64-bit, and the member count, which padding cannot hide.
#if INTPTR_MAX == INT64_MAX
static_assert(sizeof(llama_batch) == 56, "llama_batch changed shape: audit Scratch::as_batch and as_embd_batch");
#endif
namespace shape_guard {
inline void llama_batch_has_seven_members() {
  [[maybe_unused]] auto [n_tokens, token, embd, pos, n_seq_id, seq_id, logits] = llama_batch{};
}
}  // namespace shape_guard
 
namespace detail {
inline void batch_clear(llama_batch& batch) { batch.n_tokens = 0; }
inline void batch_add(llama_batch& batch, llama_token id, llama_pos pos,
                      llama_seq_id seq_id, bool logits) {
  const int32_t i = batch.n_tokens;
  batch.token[i] = id;
  batch.pos[i] = pos;
  batch.n_seq_id[i] = 1;
  batch.seq_id[i][0] = seq_id;
  batch.logits[i] = logits ? int8_t{1} : int8_t{0};
  batch.n_tokens++;
}
} // namespace detail
 
inline bool leaves_branch_intact(int32_t rc, bool partial) {
  return (rc == 1 || rc == -1) && !partial;
}
 
struct DecodeError : std::runtime_error {
  int32_t rc;
  bool partial;
  DecodeError(int32_t rc_, bool partial_, const std::string& msg)
      : std::runtime_error(msg + " (rc=" + std::to_string(rc_) +
                           (partial_ ? ", partial)" : ")")),
        rc(rc_), partial(partial_), msg_(msg) {}
 
  DecodeError as_partial() const { return DecodeError(rc, true, msg_); }
 
private:
  std::string msg_;
};
 
inline void require_fits_batch(llama_context* ctx, int32_t n_tokens, const char* who) {
  const int32_t limit = static_cast<int32_t>(llama_n_batch(ctx));
  if (n_tokens > limit) {
    throw std::runtime_error(std::string(who) + " - a batch of " + std::to_string(n_tokens) +
                             " tokens exceeds the context's n_batch (" + std::to_string(limit) + ")");
  }
}
 
[[nodiscard]] inline int many(llama_context *ctx, const llama_token *tokens,
                               int32_t n_tokens, int32_t n_past, int32_t n_batch,
                               llama_seq_id seq_id = 0,
                               int32_t* n_committed = nullptr) {
  LLOYAL_LOG_DEBUG(
      "[decode::many] Processing %d tokens at position %d", n_tokens,
      n_past);
 
  if (!ctx) {
    LLOYAL_LOG_DEBUG("[decode::many] ERROR: NULL context");
    throw std::runtime_error("decode::many - NULL context");
  }
 
  if (!tokens || n_tokens <= 0) {
    LLOYAL_LOG_DEBUG("[decode::many] ERROR: Invalid token array");
    throw std::runtime_error("decode::many - Invalid token array");
  }
 
  if (n_batch <= 0) {
    throw std::runtime_error("decode::many - n_batch must be positive");
  }
  require_fits_batch(ctx, std::min(n_batch, n_tokens), "decode::many");
 
  // Thread-local batch avoids per-call allocation. Grows if needed, never shrinks.
  struct ThreadLocalBatch {
    llama_batch batch{};
    int32_t capacity = 0;
 
    void ensure(int32_t n) {
      if (n <= capacity) return;
      if (capacity > 0) llama_batch_free(batch);
      batch = llama_batch_init(n, 0, 1);
      capacity = n;
    }
 
    ~ThreadLocalBatch() {
      if (capacity > 0) llama_batch_free(batch);
    }
  };
  thread_local ThreadLocalBatch tl;
  tl.ensure(n_batch);
  llama_batch& batch = tl.batch;
 
  // Process tokens in chunks
  int32_t processed = 0;
  while (processed < n_tokens) {
    const int32_t n_eval = std::min(n_tokens - processed, n_batch);
 
    detail::batch_clear(batch);
 
    // Add tokens one by one, mark logits=true only on the final chunk's last token
    const bool is_last_chunk = (processed + n_eval >= n_tokens);
    for (int32_t i = 0; i < n_eval; ++i) {
      const int32_t pos = n_past + i;
      const bool want_logits = is_last_chunk && (i == n_eval - 1);
 
      detail::batch_add(batch, tokens[processed + i], pos, seq_id, want_logits);
    }
 
    // Decode chunk (updates KV cache)
    const int rc = llama_decode(ctx, batch);
    if (rc != 0) {
      LLOYAL_LOG_DEBUG(
          "[decode::many] ERROR: llama_decode failed at position %d (rc=%d)",
          n_past, rc);
      if (n_committed) *n_committed = processed;
      return rc;
    }
 
    n_past += n_eval;
    processed += n_eval;
 
    LLOYAL_LOG_DEBUG("[decode::many] Processed %d/%d tokens",
                     processed, n_tokens);
  }
 
  LLOYAL_LOG_DEBUG("[decode::many] Decode complete");
  if (n_committed) *n_committed = n_tokens;
  return 0;
}
 
[[nodiscard]] inline int many(llama_context *ctx,
                               const std::vector<llama_token> &tokens,
                               int32_t n_past, int32_t n_batch,
                               llama_seq_id seq_id = 0) {
  return many(ctx, tokens.data(), static_cast<int32_t>(tokens.size()), n_past,
              n_batch, seq_id);
}
 
[[nodiscard]] inline int one(llama_context *ctx, llama_token tok, llama_pos pos,
                              llama_seq_id seq_id = 0, bool want_logits = true) {
  if (!ctx) {
    throw std::runtime_error("decode::one - NULL context");
  }
 
  struct ThreadLocalBatch {
    llama_batch batch = llama_batch_init(1, 0, 1);
    ~ThreadLocalBatch() { llama_batch_free(batch); }
  };
  thread_local ThreadLocalBatch tl;
 
  detail::batch_clear(tl.batch);
  detail::batch_add(tl.batch, tok, pos, seq_id, want_logits);
 
  return llama_decode(ctx, tl.batch);
}
 
// ============================================================================
// Multi-Sequence Decode
// ============================================================================
 
struct EachItem {
  llama_token token;            
  llama_pos pos;                
  llama_seq_id seq_id;          
  bool output_logits = false;   
};
 
struct ScatterItem {
  std::span<const llama_token> tokens;    
  llama_pos start_pos;                    
  llama_seq_id seq_id;                    
  bool output_logits = false;             
};
 
struct Scratch {
  std::vector<llama_token> tokens_;
  std::vector<llama_pos> pos_;
  std::vector<int32_t> n_seq_id_;
  std::vector<llama_seq_id> seq_id_single_;
  std::vector<llama_seq_id*> seq_id_ptrs_;
  std::vector<int8_t> logits_;
 
  void resize(int32_t n, int32_t n_pos_per_embd = 1) {
    tokens_.resize(n);
    pos_.resize(static_cast<size_t>(n) * n_pos_per_embd);
    n_seq_id_.resize(n);
    seq_id_single_.resize(n);
    seq_id_ptrs_.resize(n);
    logits_.resize(n);
  }
 
  llama_batch as_batch(int32_t n_tokens) {
    llama_batch batch{};
    batch.n_tokens = n_tokens;
    batch.token = tokens_.data();
    batch.embd = nullptr;
    batch.pos = pos_.data();
    batch.n_seq_id = n_seq_id_.data();
    batch.seq_id = seq_id_ptrs_.data();
    batch.logits = logits_.data();
    return batch;
  }
 
  llama_batch as_embd_batch(int32_t n_tokens, const float* rows) {
    llama_batch batch{};
    batch.n_tokens = n_tokens;
    batch.token = nullptr;
    batch.embd = const_cast<float*>(rows);
    batch.pos = pos_.data();
    batch.n_seq_id = n_seq_id_.data();
    batch.seq_id = seq_id_ptrs_.data();
    batch.logits = logits_.data();
    return batch;
  }
};
 
inline void require_row_width(llama_context* ctx, int32_t n_embd_inp, const char* who) {
  if (const llama_model* m = llama_get_model(ctx)) {
    const int32_t expected = llama_model_n_embd_inp(m);
    if (expected > 0 && n_embd_inp != expected) {
      throw std::runtime_error(std::string(who) + " - n_embd_inp " + std::to_string(n_embd_inp) +
                               " does not match the model's input width " + std::to_string(expected));
    }
  }
}
 
[[nodiscard]] inline int each(llama_context* ctx,
                               const EachItem* items,
                               int32_t n,
                               Scratch& scratch) {
  if (!ctx) {
    throw std::runtime_error("decode::each - NULL context");
  }
  if (n < 0) {
    throw std::runtime_error("decode::each - negative item count");
  }
  if (n == 0) return 0;
  require_fits_batch(ctx, n, "decode::each");
 
  scratch.resize(n);
 
  for (int32_t i = 0; i < n; ++i) {
    scratch.tokens_[i] = items[i].token;
    scratch.pos_[i] = items[i].pos;
    scratch.n_seq_id_[i] = 1;
    scratch.seq_id_single_[i] = items[i].seq_id;
    scratch.seq_id_ptrs_[i] = &scratch.seq_id_single_[i];
    scratch.logits_[i] = items[i].output_logits ? int8_t{1} : int8_t{0};
  }
 
  llama_batch batch = scratch.as_batch(n);
 
  LLOYAL_LOG_DEBUG("[decode::each] Submitting %d tokens across %d sequences", n, n);
 
  return llama_decode(ctx, batch);
}
 
[[nodiscard]] inline int each(llama_context* ctx,
                               const std::vector<EachItem>& items,
                               Scratch& scratch) {
  return each(ctx, items.data(), static_cast<int32_t>(items.size()), scratch);
}
 
[[nodiscard]] inline int scatter(llama_context* ctx,
                                        const ScatterItem* items,
                                        int32_t n,
                                        Scratch& scratch) {
  if (!ctx) {
    throw std::runtime_error("decode::scatter - NULL context");
  }
  if (n < 0) {
    throw std::runtime_error("decode::scatter - negative item count");
  }
 
  int32_t total = 0;
  for (int32_t i = 0; i < n; ++i) {
    total += static_cast<int32_t>(items[i].tokens.size());
  }
  if (total == 0) return 0;
  require_fits_batch(ctx, total, "decode::scatter");
 
  scratch.resize(total);
 
  int32_t cursor = 0;
  for (int32_t i = 0; i < n; ++i) {
    const auto& item = items[i];
    const llama_pos base_pos = item.start_pos;
    const int32_t item_n = static_cast<int32_t>(item.tokens.size());
 
    for (int32_t j = 0; j < item_n; ++j) {
      scratch.tokens_[cursor] = item.tokens[j];
      scratch.pos_[cursor] = base_pos + j;
      scratch.n_seq_id_[cursor] = 1;
      scratch.seq_id_single_[cursor] = item.seq_id;
      scratch.seq_id_ptrs_[cursor] = &scratch.seq_id_single_[cursor];
 
      const bool want_logits =
          item.output_logits ? (j == item_n - 1) : false;
      scratch.logits_[cursor] = want_logits ? int8_t{1} : int8_t{0};
 
      ++cursor;
    }
  }
 
  llama_batch batch = scratch.as_batch(total);
 
  LLOYAL_LOG_DEBUG("[decode::scatter] Submitting %d total tokens across %d sequences", total, n);
 
  return llama_decode(ctx, batch);
}
 
[[nodiscard]] inline int scatter(llama_context* ctx,
                                        const std::vector<ScatterItem>& items,
                                        Scratch& scratch) {
  return scatter(ctx, items.data(), static_cast<int32_t>(items.size()), scratch);
}
 
// ============================================================================
// Embedding-Row Decode (multimodal ingress)
// ============================================================================
 
struct EmbdItem {
  const float* rows = nullptr;
  int32_t n_rows = 0;
  int32_t n_embd_inp = 0;
  const llama_pos* pos = nullptr;
  int32_t n_pos_per_embd = 1;
  llama_seq_id seq_id = 0;
  bool non_causal = false;
  bool output_logits = false;
};
 
struct Segment {
  enum class Kind { Text, Embd };
  Kind kind = Kind::Text;
 
  std::span<const llama_token> tokens;
 
  const float* rows = nullptr;
  int32_t n_rows = 0;
  int32_t n_embd_inp = 0;
  llama_pos n_pos = 0;
  int32_t n_pos_per_embd = 1;
  bool non_causal = false;
};
 
struct SegmentSource {
  virtual ~SegmentSource() = default;
 
  virtual size_t size() = 0;
 
  virtual size_t cells() const = 0;
 
  virtual Segment at(size_t i) = 0;
 
  virtual void positions(size_t i, llama_pos base, llama_pos* out) = 0;
 
  virtual int32_t input_of(size_t /*i*/) const { return -1; }
 
  virtual size_t total_of(int32_t /*input*/) const { return 0; }
};
 
class TokenSource final : public SegmentSource {
public:
  explicit TokenSource(std::span<const llama_token> tokens) : tokens_(tokens) {}
  size_t size() override { return 1; }
  size_t cells() const override { return tokens_.size(); }
  Segment at(size_t) override {
    Segment seg;
    seg.tokens = tokens_;
    return seg;
  }
  void positions(size_t, llama_pos, llama_pos*) override {}
 
private:
  std::span<const llama_token> tokens_;
};
 
struct SegmentProgress {
  int32_t input = -1;
  size_t landed = 0;
  size_t total = 0;
};
 
[[nodiscard]] inline int embd(llama_context* ctx,
                              const EmbdItem& item,
                              int32_t n_batch,
                              Scratch& scratch,
                              int32_t* n_committed = nullptr) {
  if (!ctx) {
    throw std::runtime_error("decode::embd - NULL context");
  }
  if (!item.rows || item.n_rows <= 0 || item.n_embd_inp <= 0) {
    throw std::runtime_error("decode::embd - invalid rows");
  }
  if (!item.pos || (item.n_pos_per_embd != 1 && item.n_pos_per_embd != 4)) {
    throw std::runtime_error("decode::embd - invalid positions");
  }
  if (n_batch <= 0) {
    throw std::runtime_error("decode::embd - n_batch must be positive");
  }
  // The row width must match the RESIDENT model, not merely be positive.
  // llama_batch carries no width metadata: llama_decode consumes rows at the
  // model's own input width while the chunk loop below strides by this one.
  // A wrong-but-positive width therefore starts later chunks mid-row and reads
  // past the caller's allocation — corrupt vision state, or an out-of-bounds
  // native read, with nothing to signal it.
  require_row_width(ctx, item.n_embd_inp, "decode::embd");
 
  const int32_t n    = item.n_rows;
  const int32_t nppe = item.n_pos_per_embd;
  require_fits_batch(ctx, std::min(n_batch, n), "decode::embd");
 
  // A non-causal block is bidirectional: every row must be able to attend to
  // every other row, which only holds if they share one forward pass. The
  // chunk loop below would split an oversized block across separate
  // llama_decode calls, and rows in an earlier call cannot see later ones —
  // the block silently stops being bidirectional and the vision state is
  // wrong with no error anywhere. Refuse the configuration instead.
  if (item.non_causal) {
    const int32_t n_ubatch = static_cast<int32_t>(llama_n_ubatch(ctx));
    if (n > n_batch || n > n_ubatch) {
      throw std::runtime_error(
          "decode::embd - non-causal block of " + std::to_string(n) +
          " rows exceeds n_batch (" + std::to_string(n_batch) +
          ") or n_ubatch (" + std::to_string(n_ubatch) +
          "); a bidirectional image must decode in a single dispatch");
    }
  }
 
  struct CausalGuard {
    llama_context* ctx = nullptr;
    bool engaged = false;
    ~CausalGuard() { if (engaged) llama_set_causal_attn(ctx, true); }
  } causal_guard{ctx, item.non_causal};
 
  if (item.non_causal) llama_set_causal_attn(ctx, false);
 
  int32_t processed = 0;
  while (processed < n) {
    const int32_t n_view = std::min(n - processed, n_batch);
    scratch.resize(n_view, nppe);
 
    // Section-major re-pack for this view: source section s spans
    // [s*n + processed, +n_view); destination is contiguous per section.
    for (int32_t s = 0; s < nppe; ++s) {
      const llama_pos* src = item.pos + static_cast<size_t>(s) * n + processed;
      std::copy(src, src + n_view,
                scratch.pos_.begin() + static_cast<size_t>(s) * n_view);
    }
 
    for (int32_t i = 0; i < n_view; ++i) {
      scratch.n_seq_id_[i]     = 1;
      scratch.seq_id_single_[i] = item.seq_id;
      scratch.seq_id_ptrs_[i]  = &scratch.seq_id_single_[i];
      scratch.logits_[i]       = 0;
    }
    const bool is_last_view = (processed + n_view >= n);
    if (item.output_logits && is_last_view) {
      scratch.logits_[n_view - 1] = 1;
    }
 
    llama_batch batch = scratch.as_embd_batch(
        n_view, item.rows + static_cast<size_t>(processed) * item.n_embd_inp);
 
    LLOYAL_LOG_DEBUG("[decode::embd] Submitting %d/%d rows (seq %d)",
                     processed + n_view, n, item.seq_id);
 
    const int rc = llama_decode(ctx, batch);
    if (rc != 0) {
      LLOYAL_LOG_DEBUG("[decode::embd] ERROR: llama_decode failed (rc=%d)", rc);
      if (n_committed) *n_committed = processed;
      return rc;
    }
 
    processed += n_view;
  }
 
  LLOYAL_LOG_DEBUG("[decode::embd] Decode complete (%d rows)", n);
  if (n_committed) *n_committed = n;
  return 0;
}
 
// ============================================================================
// The segments rail
// ============================================================================
 
struct LandedSegment {
  size_t index = 0;
  Segment::Kind kind = Segment::Kind::Text;
  int32_t cells = 0;
  llama_pos n_pos = 0;
  bool has_logits = false;
};
 
struct StreamReading {
  bool keeps_memory = true;
  bool causal = true;
  bool pooled = false;
 
  bool whole() const { return !keeps_memory || !causal || pooled; }
 
  bool takes_bidirectional() const { return !keeps_memory || !causal; }
};
 
inline StreamReading reading_of(llama_context* ctx) {
  return StreamReading{llama_get_memory(ctx) != nullptr, llama_get_causal_attn(ctx),
                       llama_pooling_type(ctx) != LLAMA_POOLING_TYPE_NONE};
}
 
namespace detail {
 
inline void require_segment(const Segment& seg, size_t i) {
  if (seg.kind == Segment::Kind::Text) {
    // Terminality comes from the index, so an empty trailing segment would make
    // the real final one non-terminal: [EMBD, empty TEXT] would decode the image
    // without logits and hand back nothing to sample.
    if (seg.tokens.empty()) {
      throw std::runtime_error("decode::segments - empty TEXT segment at " + std::to_string(i) +
                               "; a source must not yield empty segments");
    }
    return;
  }
  if (!seg.rows || seg.n_rows <= 0) {
    throw std::runtime_error("decode::segments - empty embedding segment at " + std::to_string(i));
  }
  // n_pos_per_embd of -1 wraps a size_t multiply into an enormous allocation,
  // and 0 hands positions() a zero-length buffer to write into.
  if (seg.n_pos_per_embd != 1 && seg.n_pos_per_embd != 4) {
    throw std::runtime_error("decode::segments - segment " + std::to_string(i) + " has n_pos_per_embd " +
                             std::to_string(seg.n_pos_per_embd) + " (expected 1 or 4)");
  }
  if (seg.n_embd_inp <= 0) {
    throw std::runtime_error("decode::segments - segment " + std::to_string(i) + " has non-positive n_embd_inp");
  }
  if (seg.n_pos <= 0 || seg.n_pos > seg.n_rows) {
    throw std::runtime_error("decode::segments - segment " + std::to_string(i) + " has n_pos " +
                             std::to_string(seg.n_pos) + " outside (0, n_rows=" + std::to_string(seg.n_rows) + "]");
  }
}
 
template <typename OnLanded>
void segments_each(llama_context* ctx, llama_seq_id seq, llama_pos base, SegmentSource& source,
                   int32_t rows_n_batch, Scratch& scratch, std::vector<llama_pos>& positions,
                   OnLanded&& landed) {
  const size_t n = source.size();
  const int32_t text_n_batch = static_cast<int32_t>(llama_n_batch(ctx));
  llama_pos at = base;
  bool any_landed = false;
 
  const auto failed = [&](int rc, int32_t committed, const std::string& what) {
    return DecodeError(rc, any_landed || committed > 0, "decode::segments - " + what);
  };
 
  for (size_t i = 0; i < n; ++i) {
    const Segment seg = source.at(i);
    const bool is_last = (i + 1 == n);
    require_segment(seg, i);
 
    if (seg.kind == Segment::Kind::Text) {
      const int32_t n_tokens = static_cast<int32_t>(seg.tokens.size());
      int32_t committed = 0;
      if (const int rc = many(ctx, seg.tokens.data(), n_tokens, at, text_n_batch, seq, &committed); rc != 0) {
        throw failed(rc, committed, "TEXT segment " + std::to_string(i) + " failed");
      }
      at += n_tokens;
      any_landed = true;
      landed(LandedSegment{i, Segment::Kind::Text, n_tokens, n_tokens, true});
      continue;
    }
 
    // The base is handed to the source, which applies its own model's rules.
    positions.assign(static_cast<size_t>(seg.n_rows) * seg.n_pos_per_embd, 0);
    source.positions(i, at, positions.data());
 
    EmbdItem item;
    item.rows           = seg.rows;
    item.n_rows         = seg.n_rows;
    item.n_embd_inp     = seg.n_embd_inp;
    item.pos            = positions.data();
    item.n_pos_per_embd = seg.n_pos_per_embd;
    item.seq_id         = seq;
    item.non_causal     = seg.non_causal;
    item.output_logits  = is_last;
    int32_t committed = 0;
    if (const int rc = embd(ctx, item, rows_n_batch, scratch, &committed); rc != 0) {
      throw failed(rc, committed, "EMBD segment " + std::to_string(i) + " failed");
    }
    at += seg.n_pos;
    any_landed = true;
    landed(LandedSegment{i, Segment::Kind::Embd, seg.n_rows, seg.n_pos, is_last});
  }
}
 
template <typename OnLanded>
void segments_whole(llama_context* ctx, const StreamReading& reading, llama_seq_id seq, llama_pos base,
                    SegmentSource& source, std::vector<llama_pos>& positions, OnLanded&& landed) {
  // An oversized dispatch on this context aborts the process inside llama.cpp
  // rather than returning an error. The source's quote refuses one early, before
  // anything is encoded; the entries actually built are counted against the same
  // bound below, since the source is untrusted.
  const size_t n_fit = std::min(llama_n_batch(ctx), llama_n_ubatch(ctx));
  const auto too_large = [&](size_t n_entries) {
    return std::runtime_error("decode::segments - a stream of " + std::to_string(n_entries) +
                              " entries must go in one dispatch on this context, which holds " +
                              std::to_string(n_fit));
  };
  if (source.cells() > n_fit) throw too_large(source.cells());
 
  std::unique_ptr<llama_batch_ext, decltype(&llama_batch_ext_free)> batch(llama_batch_ext_init(ctx),
                                                                          llama_batch_ext_free);
  if (!batch) {
    throw std::runtime_error("decode::segments - llama_batch_ext_init failed");
  }
  // A stream placed whole is read whole: every entry's output is on, pooled or
  // per entry, as llama.cpp's encoder computes them for such a context anyway.
  const auto entry = [&](int32_t idx, const llama_pos* pos) {
    if (idx >= 0 && static_cast<size_t>(idx) >= n_fit) throw too_large(static_cast<size_t>(idx) + 1);
    if (idx < 0 || !llama_batch_ext_set_pos(batch.get(), idx, pos) ||
        !llama_batch_ext_set_output_embd(batch.get(), idx, true)) {
      throw std::runtime_error("decode::segments - the batch refused an entry (" + std::to_string(idx) + ")");
    }
    return idx;
  };
 
  const size_t n = source.size();
  std::vector<LandedSegment> placed;
  placed.reserve(n);
  llama_pos at = base;
  int32_t last = -1;
 
  for (size_t i = 0; i < n; ++i) {
    const Segment seg = source.at(i);
    require_segment(seg, i);
 
    if (seg.kind == Segment::Kind::Text) {
      for (size_t k = 0; k < seg.tokens.size(); ++k) {
        const llama_pos pos = at + static_cast<llama_pos>(k);
        last = entry(llama_batch_ext_add_token(batch.get(), seq, seg.tokens[k]), &pos);
      }
      const int32_t n_tokens = static_cast<int32_t>(seg.tokens.size());
      placed.push_back(LandedSegment{i, Segment::Kind::Text, n_tokens, n_tokens, false});
      at += n_tokens;
      continue;
    }
 
    require_row_width(ctx, seg.n_embd_inp, "decode::segments");
    if (seg.non_causal && !reading.takes_bidirectional()) {
      throw std::runtime_error("decode::segments - segment " + std::to_string(i) +
                               " is bidirectional, and this context attends causally");
    }
    positions.assign(static_cast<size_t>(seg.n_rows) * seg.n_pos_per_embd, 0);
    source.positions(i, at, positions.data());
    // positions() is section-major; an entry takes its own [t, y, x, z].
    std::array<llama_pos, 4> pos{};
    for (int32_t r = 0; r < seg.n_rows; ++r) {
      for (int32_t k = 0; k < seg.n_pos_per_embd; ++k) {
        pos[k] = positions[static_cast<size_t>(k) * seg.n_rows + r];
      }
      const llama_embd row{seg.rows + static_cast<size_t>(r) * seg.n_embd_inp, 1,
                           static_cast<size_t>(seg.n_embd_inp)};
      last = entry(llama_batch_ext_add_embd(batch.get(), seq, row), pos.data());
    }
    placed.push_back(LandedSegment{i, Segment::Kind::Embd, seg.n_rows, seg.n_pos, false});
    at += seg.n_pos;
  }
 
  if (last < 0) return;
  // The stream's last entry carries its logits, as a per-segment placement's does.
  llama_batch_ext_set_output_logits(batch.get(), last, true);
  placed.back().has_logits = true;
 
  // A model that cannot mix tokens and rows in one batch answers -1 here.
  if (const int rc = llama_process(ctx, LLAMA_PROCESS_TYPE_DECODE, batch.get()); rc != 0) {
    throw DecodeError(rc, false, "decode::segments - the stream's one dispatch failed");
  }
  for (const auto& seg : placed) landed(seg);
}
 
}  // namespace detail
 
 
template <typename OnLanded>
inline void segments(llama_context* ctx, llama_seq_id seq, llama_pos base,
                     SegmentSource& source, int32_t rows_n_batch,
                     Scratch& scratch, std::vector<llama_pos>& positions,
                     OnLanded&& landed) {
  if (!ctx) {
    throw std::runtime_error("decode::segments - NULL context");
  }
  if (const StreamReading reading = reading_of(ctx); reading.whole()) {
    detail::segments_whole(ctx, reading, seq, base, source, positions, landed);
    return;
  }
  detail::segments_each(ctx, seq, base, source, rows_n_batch, scratch, positions, landed);
}
 
// ============================================================================
// Bin-Packing Utility
// ============================================================================
 
struct PackedChunk {
  std::vector<int32_t> indices;   
  bool oversized = false;         
};
 
inline std::vector<PackedChunk> bin_pack(
    const std::span<const llama_token>* items,
    int32_t n,
    int32_t n_batch) {
 
  std::vector<PackedChunk> chunks;
  int32_t chunk_total = 0;
 
  for (int32_t i = 0; i < n; ++i) {
    int32_t tc = static_cast<int32_t>(items[i].size());
    if (tc == 0) continue;
 
    if (tc > n_batch) {
      chunks.push_back({{i}, true});
      continue;
    }
 
    if (chunks.empty() || chunks.back().oversized ||
        chunk_total + tc > n_batch) {
      chunks.push_back({{i}, false});
      chunk_total = tc;
    } else {
      chunks.back().indices.push_back(i);
      chunk_total += tc;
    }
  }
 
  return chunks;
}
 
} // namespace lloyal::decode