#include "ds4.h"

#include "linenoise.h"

/* ds4 CLI.

*

* One-shot mode builds a single DeepSeek chat prompt and exits. Interactive

* mode keeps a rendered token transcript plus one ds4_session, so follow-up

* turns reuse the live Metal KV checkpoint just like the server does. The CLI

* deliberately keeps policy here and leaves graph/cache mechanics inside the

* engine API. */

#include <ctype.h>

#include <errno.h>

#include <limits.h>

#include <math.h>

#include <signal.h>

#include <stdbool.h>

#include <stdint.h>

#include <stdio.h>

#include <stdlib.h>

#include <string.h>

#include <stdarg.h>

#include <time.h>

#include <unistd.h>

typedef struct {

const char *prompt;

const char *system;

int n_predict;

int ctx_size;

float temperature;

float top_p;

uint64_t seed;

bool dump_tokens;

const char *dump_logprobs_path;

int dump_logprobs_top_k;

ds4_think_mode think_mode;

bool head_test;

bool first_token_test;

bool metal_graph_test;

bool metal_graph_full_test;

bool metal_graph_prompt_test;

} cli_generation_options;

typedef struct {

ds4_engine_options engine;

cli_generation_options gen;

char *prompt_owned;

bool inspect;

} cli_config;

static volatile sig_atomic_t cli_interrupted;

static void cli_sigint_handler(int sig) {

(void)sig;

cli_interrupted = 1;

}

static bool cli_interrupt_requested(void) {

return cli_interrupted != 0;

}

static void cli_interrupt_clear(void) {

cli_interrupted = 0;

}

static void usage(FILE *fp) {

fprintf(fp,

"Usage: ds4 [(-p PROMPT | --prompt-file FILE)] [options]\n"

"\n"

"Invocation modes:\n"

" ds4\n"

" Start the interactive chat prompt: ds4>\n"

" ds4 -p TEXT\n"

" Run one prompt and exit.\n"

" ds4 --prompt-file FILE\n"

" Run one prompt read from FILE and exit. Useful for long prompts.\n"

"\n"

"Model and runtime:\n"

" -m, --model FILE\n"

" GGUF model path. Default: ds4flash.gguf\n"

" --mtp FILE\n"

" Optional MTP support GGUF used for draft-token probes.\n"

" --mtp-draft N\n"

" Maximum autoregressive MTP draft tokens per speculative step. Default: 1\n"

" --mtp-margin F\n"

" Minimum recursive-draft confidence for the fast N=2 verifier. Default: 3\n"

" -c, --ctx N\n"

" Context size allocated for the session. Default: 32768\n"

" --metal\n"

" Use the Metal graph backend. This is the normal fast path and the default.\n"

" --cpu\n"

" Use the CPU reference/debug backend. Not recommended for normal inference.\n"

" --backend NAME\n"

" Select backend explicitly: metal or cpu. Default: metal\n"

" -t, --threads N\n"

" CPU helper threads for host-side or reference work.\n"

" --quality\n"

" Prefer exact kernels where faster approximate paths exist; MTP uses strict verification.\n"

" --warm-weights\n"

" Touch mapped tensor pages before generation. Slower startup, fewer first-use stalls.\n"

" --disk-offload-weights\n"

" Lazily map GGUF tensor pages for Metal weights instead of eagerly registering all weights.\n"

"\n"

"Prompt and generation:\n"

" -p, --prompt TEXT\n"

" Prompt to generate from.\n"

" --prompt-file FILE\n"

" Read the prompt text from FILE.\n"

" -sys, --system TEXT\n"

" System prompt. Empty string disables the default. Default: You are a helpful assistant\n"

" -n, --tokens N\n"

" Maximum tokens to generate. Default: 50000\n"

" --temp F\n"

" Sampling temperature. 0 is greedy/deterministic. Default: 1\n"

" --top-p F\n"

" Nucleus sampling probability. Default: 1\n"

" --seed N\n"

" Sampling seed for reproducible non-greedy runs. Default: time-based\n"

" --think\n"

" Use normal thinking mode. This is the default.\n"

" --think-max\n"

" Use Think Max when --ctx is at least 393216 tokens; otherwise normal thinking.\n"

" --nothink\n"

" Start assistant turns with </think> for direct non-thinking replies.\n"

"\n"

"Interactive commands:\n"

" /help\n"

" Show interactive commands.\n"

" /think, /think-max, /nothink\n"

" Select normal thinking, context-gated Think Max, or non-thinking mode.\n"

" /ctx N\n"

" Recreate the interactive session with a new context size.\n"

" /read FILE\n"

" Read a prompt from FILE and run it as the next user message.\n"

" /quit, /exit\n"

" Leave the interactive prompt.\n"

" Ctrl+C\n"

" Stop the current generation and return to ds4> without exiting.\n"

"\n"

"Diagnostics:\n"

" --inspect\n"

" Load the model and print a summary only.\n"

" --dump-tokens\n"

" Print the encoded chat prompt tokens.\n"

" --dump-logprobs FILE\n"

" Write greedy continuation top-logprobs as JSON without printing text.\n"

" --logprobs-top-k N\n"

" Number of local alternatives stored by --dump-logprobs. Default: 20\n"

" --head-test\n"

" Run the output HC/logits head after the native slice.\n"

" --first-token-test\n"

" Run an exact CPU whole-model pass for the first prompt token.\n"

" --metal-graph-test\n"

" Compare first GPU-resident graph stages with CPU.\n"

" --metal-graph-full-test\n"

" Run the GPU-resident self-token graph across all layers.\n"

" --metal-graph-prompt-test\n"

" Compare CPU and GPU graph logits for the full prompt.\n"

"\n"

"Normal CLI commands:\n"

" ./ds4\n"

" ./ds4 -p \"Scrivi una storia su una papera scansafatiche\"\n"

" ./ds4 --think-max --prompt-file prompt.txt --ctx 393216\n"

"\n"

"Notes:\n"

" The CLI keeps KV cache state across interactive turns on the Metal backend.\n"

" Long added input is processed with batched prefill; short continuations use decode.\n"

" Startup prints the extra context-buffer memory for the selected context size.\n"

"\n"

" -h, --help\n"

" Show this help.\n");

}

static int parse_int(const char *s, const char *opt) {

char *end = NULL;

long v = strtol(s, &end, 10);

if (s[0] == '\0' || *end != '\0' || v <= 0 || v > INT32_MAX) {

fprintf(stderr, "ds4: invalid value for %s: %s\n", opt, s);

exit(2);

}

return (int)v;

}

static uint64_t parse_u64(const char *s, const char *opt) {

char *end = NULL;

unsigned long long v = strtoull(s, &end, 10);

if (s[0] == '\0' || *end != '\0' || v == 0) {

fprintf(stderr, "ds4: invalid value for %s: %s\n", opt, s);

exit(2);

}

return (uint64_t)v;

}

static float parse_float_range(const char *s, const char *opt, float min, float max) {

char *end = NULL;

float v = strtof(s, &end);

if (s[0] == '\0' || *end != '\0' || !isfinite(v) || v < min || v > max) {

fprintf(stderr, "ds4: invalid value for %s: %s\n", opt, s);

exit(2);

}

return v;

}

static ds4_backend parse_backend(const char *s) {

if (!strcmp(s, "metal")) return DS4_BACKEND_METAL;

if (!strcmp(s, "cpu")) return DS4_BACKEND_CPU;

fprintf(stderr, "ds4: invalid backend: %s\n", s);

fprintf(stderr, "ds4: valid backends are: metal, cpu\n");

exit(2);

}

static void log_context_memory(ds4_backend backend, int ctx_size) {

ds4_context_memory m = ds4_context_memory_estimate(backend, ctx_size);

fprintf(stderr,

"ds4: context buffers %.2f MiB (ctx=%d, backend=%s, prefill_chunk=%u, raw_kv_rows=%u, compressed_kv_rows=%u)\n",

(double)m.total_bytes / (1024.0 * 1024.0),

ctx_size,

ds4_backend_name(backend),

m.prefill_cap,

m.raw_cap,

m.comp_cap);

}

static ds4_think_mode cli_effective_think_mode(const cli_generation_options *gen) {

return ds4_think_mode_for_context(gen->think_mode, gen->ctx_size);

}

static bool cli_think_max_downgraded(const cli_generation_options *gen) {

return gen->think_mode == DS4_THINK_MAX &&

cli_effective_think_mode(gen) != DS4_THINK_MAX;

}

static void cli_warn_think_max_downgraded(const cli_generation_options *gen, const char *name) {

if (!cli_think_max_downgraded(gen)) return;

ds4_log(stderr,

DS4_LOG_WARNING,

"ds4: warning: %s needs --ctx >= %u; ctx=%d uses normal thinking instead\n",

name,

ds4_think_max_min_context(),

gen->ctx_size);

}

static double cli_now_sec(void) {

struct timespec ts;

clock_gettime(CLOCK_MONOTONIC, &ts);

return (double)ts.tv_sec + (double)ts.tv_nsec * 1.0e-9;

}

static char *read_prompt_file(const char *path, bool fatal);

typedef struct {

int base_tokens;

int input_tokens;

bool use_color;

} cli_prefill_progress;

static void cli_prefill_progress_cb(void *ud, const char *event, int current, int total) {

(void)total;

cli_prefill_progress *p = ud;

if (!p || !event || strcmp(event, "prefill_chunk") || p->input_tokens <= 0) return;

int processed = current - p->base_tokens;

if (processed < 0) processed = 0;

if (processed > p->input_tokens) processed = p->input_tokens;

double pct = 100.0 * (double)processed / (double)p->input_tokens;

if (pct > 100.0) pct = 100.0;

if (p->use_color) {

fputc('\r', stderr);

ds4_log(stderr,

DS4_LOG_PREFILL,

"processing %d input tokens: %d/%d (%.1f%%)",

p->input_tokens,

processed,

p->input_tokens,

pct);

fputs("\x1b[K", stderr);

if (processed >= p->input_tokens) fputc('\n', stderr);

} else {

fprintf(stderr,

"processing %d input tokens: %d/%d (%.1f%%)\n",

p->input_tokens,

processed,

p->input_tokens,

pct);

}

fflush(stderr);

}

static bool is_rendered_chat_prompt(const char *prompt) {

const char *bos = "<｜begin▁of▁sentence｜>";

return prompt && strncmp(prompt, bos, strlen(bos)) == 0;

}

typedef struct {

ds4_engine *engine;

FILE *fp;

bool format_thinking;

bool in_think;

bool color_open;

bool use_color;

bool last_output_newline;

char pending[16];

size_t pending_len;

} token_printer;

static bool bytes_has_prefix(const char *p, size_t n, const char *prefix) {

size_t plen = strlen(prefix);

return n >= plen && memcmp(p, prefix, plen) == 0;

}

static bool bytes_is_partial_prefix(const char *p, size_t n, const char *prefix) {

size_t plen = strlen(prefix);

return n < plen && memcmp(prefix, p, n) == 0;

}

static void token_printer_set_grey(token_printer *p) {

if (p->use_color && !p->color_open) {

fputs("\x1b[90m", p->fp);

p->color_open = true;

}

}

static void token_printer_reset_color(token_printer *p) {

if (p->use_color && p->color_open) {

fputs("\x1b[0m", p->fp);

p->color_open = false;

}

}

static void token_printer_write_char(token_printer *p, char c) {

if (p->in_think) token_printer_set_grey(p);

fputc((unsigned char)c, p->fp);

p->last_output_newline = c == '\n';

}

static void token_printer_process(token_printer *p, const char *text, size_t len, bool finish) {

const char *think_open = "<think>";

const char *think_close = "</think>";

size_t total = p->pending_len + len;

char *buf = malloc(total ? total : 1);

if (!buf) return;

if (p->pending_len) memcpy(buf, p->pending, p->pending_len);

if (len) memcpy(buf + p->pending_len, text, len);

p->pending_len = 0;

size_t i = 0;

while (i < total) {

const char *cur = buf + i;

const size_t rem = total - i;

if (bytes_has_prefix(cur, rem, think_open)) {

p->in_think = true;

i += strlen(think_open);

continue;

}

if (bytes_has_prefix(cur, rem, think_close)) {

p->in_think = false;

token_printer_reset_color(p);

if (!p->last_output_newline) {

fputc('\n', p->fp);

p->last_output_newline = true;

}

i += strlen(think_close);

continue;

}

if (!finish && cur[0] == '<' &&

(bytes_is_partial_prefix(cur, rem, think_open) ||

bytes_is_partial_prefix(cur, rem, think_close)))

{

if (rem < sizeof(p->pending)) {

memcpy(p->pending, cur, rem);

p->pending_len = rem;

}

break;

}

token_printer_write_char(p, cur[0]);

i++;

}

free(buf);

}

static void token_printer_finish(token_printer *p) {

if (p->format_thinking) {

token_printer_process(p, NULL, 0, true);

token_printer_reset_color(p);

}

fflush(p->fp);

}

static void generation_done(void *ud) {

token_printer *p = ud;

token_printer_finish(p);

if (!p->last_output_newline) {

fputc('\n', p->fp);

p->last_output_newline = true;

}

fflush(p->fp);

}

static void token_printer_write_text(token_printer *p, const char *text, size_t len) {

if (p->format_thinking) {

token_printer_process(p, text, len, false);

} else if (len) {

fwrite(text, 1, len, p->fp);

p->last_output_newline = text[len - 1] == '\n';

}

}

static void print_generated_token(void *ud, int token) {

token_printer *p = ud;

size_t len = 0;

char *text = ds4_token_text(p->engine, token, &len);

token_printer_write_text(p, text, len);

fflush(p->fp);

free(text);

}

static void build_prompt(ds4_engine *engine, const cli_generation_options *gen, ds4_tokens *out) {

if (is_rendered_chat_prompt(gen->prompt)) {

ds4_tokenize_rendered_chat(engine, gen->prompt, out);

} else {

ds4_encode_chat_prompt(engine, gen->system, gen->prompt,

cli_effective_think_mode(gen), out);

}

}

static int run_sampled_generation(ds4_engine *engine, const cli_config *cfg, const ds4_tokens *prompt) {

ds4_session *session = NULL;

if (ds4_session_create(&session, engine, cfg->gen.ctx_size) != 0) {

fprintf(stderr, "ds4: sampled CLI generation requires the Metal session backend\n");

return 1;

}

char err[160];

ds4_think_mode think_mode = cli_effective_think_mode(&cfg->gen);

token_printer printer = {

.engine = engine,

.fp = stdout,

.format_thinking = ds4_think_mode_enabled(think_mode),

.in_think = ds4_think_mode_enabled(think_mode),

.use_color = isatty(fileno(stdout)) != 0,

.last_output_newline = true,

};

cli_prefill_progress progress = {

.base_tokens = 0,

.input_tokens = prompt->len,

.use_color = ds4_log_is_tty(stderr),

};

const double t_prefill0 = cli_now_sec();

ds4_session_set_progress(session, cli_prefill_progress_cb, &progress);

if (ds4_session_sync(session, prompt, err, sizeof(err)) != 0) {

ds4_session_set_progress(session, NULL, NULL);

fprintf(stderr, "ds4: prompt processing failed: %s\n", err);

ds4_session_free(session);

return 1;

}

ds4_session_set_progress(session, NULL, NULL);

const double t_prefill1 = cli_now_sec();

int max_tokens = cfg->gen.n_predict;

int room = ds4_session_ctx(session) - ds4_session_pos(session);

if (room <= 1) max_tokens = 0;

else if (max_tokens > room - 1) max_tokens = room - 1;

uint64_t rng = cfg->gen.seed ? cfg->gen.seed :

((uint64_t)time(NULL) ^ ((uint64_t)getpid() << 32) ^ (uint64_t)clock());

int generated = 0;

const double t_decode0 = cli_now_sec();

while (generated < max_tokens && !cli_interrupt_requested()) {

int token = ds4_session_sample(session, cfg->gen.temperature, 0, cfg->gen.top_p, 0.0f, &rng);

if (token == ds4_token_eos(engine)) break;

int toks[17];

int ntok = 0;

if (cfg->gen.temperature <= 0.0f && ds4_engine_mtp_draft_tokens(engine) > 1 &&

getenv("DS4_MTP_SPEC_DISABLE") == NULL) {

ntok = ds4_session_eval_speculative_argmax(session,

token,

max_tokens - generated,

ds4_token_eos(engine),

toks,

(int)(sizeof(toks) / sizeof(toks[0])),

err,

sizeof(err));

if (ntok < 0) {

fprintf(stderr, "ds4: decode failed: %s\n", err);

ds4_session_free(session);

return 1;

}

} else {

if (ds4_session_eval(session, token, err, sizeof(err)) != 0) {

fprintf(stderr, "ds4: decode failed: %s\n", err);

ds4_session_free(session);

return 1;

}

toks[0] = token;

ntok = 1;

}

bool stop = false;

for (int j = 0; j < ntok; j++) {

if (toks[j] == ds4_token_eos(engine)) {

stop = true;

break;

}

size_t piece_len = 0;

char *piece = ds4_token_text(engine, toks[j], &piece_len);

token_printer_write_text(&printer, piece, piece_len);

fflush(stdout);

free(piece);

generated++;

if (generated >= max_tokens) break;

}

if (stop) break;

}

const double t_decode1 = cli_now_sec();

generation_done(&printer);

if (cli_interrupt_requested()) cli_interrupt_clear();

const double prefill_s = t_prefill1 - t_prefill0;

const double decode_s = t_decode1 - t_decode0;

ds4_log(stderr,

DS4_LOG_TIMING,

"ds4: prefill: %.2f t/s, generation: %.2f t/s\n",

prefill_s > 0.0 ? (double)prompt->len / prefill_s : 0.0,

decode_s > 0.0 ? (double)generated / decode_s : 0.0);

ds4_session_free(session);

return 0;

}

static bool json_utf8_valid(const char *s, size_t n) {

size_t i = 0;

while (i < n) {

unsigned char c = (unsigned char)s[i++];

if (c < 0x80) continue;

int need = 0;

if (c >= 0xc2 && c <= 0xdf) need = 1;

else if (c >= 0xe0 && c <= 0xef) need = 2;

else if (c >= 0xf0 && c <= 0xf4) need = 3;

else return false;

if (i + (size_t)need > n) return false;

unsigned char c1 = (unsigned char)s[i];

if (c == 0xe0 && c1 < 0xa0) return false;

if (c == 0xed && c1 >= 0xa0) return false;

if (c == 0xf0 && c1 < 0x90) return false;

if (c == 0xf4 && c1 >= 0x90) return false;

for (int j = 0; j < need; j++) {

unsigned char cc = (unsigned char)s[i + (size_t)j];

if ((cc & 0xc0) != 0x80) return false;

}

i += (size_t)need;

}

return true;

}

static void json_write_string(FILE *fp, const char *s, size_t n) {

bool valid_utf8 = json_utf8_valid(s, n);

fputc('"', fp);

for (size_t i = 0; i < n; i++) {

unsigned char c = (unsigned char)s[i];

if (c == '"' || c == '\\') {

fputc('\\', fp);

fputc((char)c, fp);

} else if (c == '\n') {

fputs("\\n", fp);

} else if (c == '\r') {

fputs("\\r", fp);

} else if (c == '\t') {

fputs("\\t", fp);

} else if (c < 0x20) {

fprintf(fp, "\\u%04x", (unsigned)c);

} else if (!valid_utf8 && c >= 0x80) {

/* Tokenizer pieces can be arbitrary byte fragments. The bytes

* array is authoritative; this escape keeps the JSON valid. */

fprintf(fp, "\\u%04x", (unsigned)c);

} else {

fputc((char)c, fp);

}

}

fputc('"', fp);

}

static void json_write_token(FILE *fp, ds4_engine *engine, int token) {

size_t n = 0;

char *text = ds4_token_text(engine, token, &n);

fprintf(fp, "{\"id\":%d,\"text\":", token);

json_write_string(fp, text, n);

fputs(",\"bytes\":[", fp);

for (size_t i = 0; i < n; i++) {

if (i) fputc(',', fp);

fprintf(fp, "%u", (unsigned)(unsigned char)text[i]);

}

fputc(']', fp);

fputc('}', fp);

free(text);

}

static int run_logprob_dump(ds4_engine *engine, const cli_config *cfg, const ds4_tokens *prompt) {

ds4_session *session = NULL;

if (ds4_session_create(&session, engine, cfg->gen.ctx_size) != 0) {

fprintf(stderr, "ds4: --dump-logprobs requires the Metal session backend\n");

return 1;

}

char err[160];

cli_prefill_progress progress = {

.base_tokens = 0,

.input_tokens = prompt->len,

.use_color = ds4_log_is_tty(stderr),

};

ds4_session_set_progress(session, cli_prefill_progress_cb, &progress);

if (ds4_session_sync(session, prompt, err, sizeof(err)) != 0) {

ds4_session_set_progress(session, NULL, NULL);

fprintf(stderr, "ds4: prompt processing failed: %s\n", err);

ds4_session_free(session);

return 1;

}

ds4_session_set_progress(session, NULL, NULL);

FILE *fp = fopen(cfg->gen.dump_logprobs_path, "wb");

if (!fp) {

fprintf(stderr, "ds4: failed to open --dump-logprobs file: %s\n", cfg->gen.dump_logprobs_path);

ds4_session_free(session);

return 1;

}

int k = cfg->gen.dump_logprobs_top_k > 0 ? cfg->gen.dump_logprobs_top_k : 20;

if (k > 128) k = 128;

ds4_token_score *scores = calloc((size_t)k, sizeof(scores[0]));

if (!scores) {

fclose(fp);

ds4_session_free(session);

return 1;

}

fprintf(fp, "{\n \"source\":\"ds4\",\n \"prompt_tokens\":%d,\n \"ctx\":%d,\n \"top_k\":%d,\n \"steps\":[\n",

prompt->len, cfg->gen.ctx_size, k);

int generated = 0;

int max_tokens = cfg->gen.n_predict;

int room = ds4_session_ctx(session) - ds4_session_pos(session);

if (room <= 1) max_tokens = 0;

else if (max_tokens > room - 1) max_tokens = room - 1;

for (; generated < max_tokens; generated++) {

int n = ds4_session_top_logprobs(session, scores, k);

int token = ds4_session_argmax(session);

if (generated) fputs(",\n", fp);

fprintf(fp, " {\"step\":%d,\"selected\":", generated);

json_write_token(fp, engine, token);

fputs(",\"top_logprobs\":[", fp);

for (int i = 0; i < n && scores[i].id >= 0; i++) {

if (i) fputc(',', fp);

fputs("{\"token\":", fp);

json_write_token(fp, engine, scores[i].id);

fprintf(fp, ",\"logit\":%.9g,\"logprob\":%.9g}", scores[i].logit, scores[i].logprob);

}

fputs("]}", fp);

if (token == ds4_token_eos(engine)) break;

if (ds4_session_eval(session, token, err, sizeof(err)) != 0) {

fprintf(stderr, "ds4: decode failed while dumping logprobs: %s\n", err);

free(scores);

fclose(fp);

ds4_session_free(session);

return 1;

}

}

fputs("\n ]\n}\n", fp);

if (fclose(fp) != 0) {

fprintf(stderr, "ds4: failed to close --dump-logprobs file: %s\n", cfg->gen.dump_logprobs_path);

free(scores);

ds4_session_free(session);

return 1;

}

free(scores);

ds4_session_free(session);

return 0;

}

static int run_generation(ds4_engine *engine, const cli_config *cfg) {

ds4_tokens prompt = {0};

build_prompt(engine, &cfg->gen, &prompt);

int rc = 0;

if (cfg->gen.metal_graph_test) {

rc = ds4_engine_metal_graph_test(engine, &prompt);

ds4_tokens_free(&prompt);

return rc;

}

if (cfg->gen.metal_graph_full_test) {

rc = ds4_engine_metal_graph_full_test(engine, &prompt);

ds4_tokens_free(&prompt);

return rc;

}

if (cfg->gen.metal_graph_prompt_test) {

rc = ds4_engine_metal_graph_prompt_test(engine, &prompt, cfg->gen.ctx_size);

ds4_tokens_free(&prompt);

return rc;

}

if (cfg->gen.dump_logprobs_path) {

rc = run_logprob_dump(engine, cfg, &prompt);

ds4_tokens_free(&prompt);

return rc;

}

const bool diagnostic = cfg->gen.dump_tokens ||

cfg->gen.head_test ||

cfg->gen.first_token_test;

if (cfg->gen.head_test) {

rc = ds4_engine_head_test(engine, &prompt);

}

if (rc == 0 && cfg->gen.first_token_test) {

rc = ds4_engine_first_token_test(engine, &prompt);

}

if (cfg->gen.dump_tokens) {

ds4_engine_dump_tokens(engine, &prompt);

}

if (diagnostic) {

if (rc == 0) {

fprintf(stderr, "ds4: diagnostic run completed on the native %s path.\n",

ds4_backend_name(cfg->engine.backend));

}

} else if (cfg->gen.temperature > 0.0f || ds4_engine_mtp_draft_tokens(engine) > 1) {

rc = run_sampled_generation(engine, cfg, &prompt);

} else {

token_printer printer = {

.engine = engine,

.fp = stdout,

.format_thinking = ds4_think_mode_enabled(cli_effective_think_mode(&cfg->gen)),

.in_think = ds4_think_mode_enabled(cli_effective_think_mode(&cfg->gen)),

.use_color = isatty(fileno(stdout)) != 0,

.last_output_newline = true,

};

cli_prefill_progress progress = {

.base_tokens = 0,

.input_tokens = prompt.len,

.use_color = ds4_log_is_tty(stderr),

};

rc = ds4_engine_generate_argmax(engine, &prompt, cfg->gen.n_predict,

cfg->gen.ctx_size,

print_generated_token,

generation_done,

&printer,

cli_prefill_progress_cb,

&progress);

}

ds4_tokens_free(&prompt);

return rc;

}

static char *trim_inplace(char *s) {

while (*s && isspace((unsigned char)*s)) s++;

char *end = s + strlen(s);

while (end > s && isspace((unsigned char)end[-1])) end--;

*end = '\0';

return s;

}

static void print_repl_help(void) {

puts("Commands:");

puts(" /help Show this help.");

puts(" /think Use normal thinking mode.");

puts(" /think-max Use Think Max only when context is at least 393216 tokens.");

puts(" /nothink Disable thinking mode.");

puts(" /ctx N Set context size for following prompts.");

puts(" /read FILE Read a prompt from FILE and run it.");

puts(" /quit, /exit Leave the prompt.");

puts(" Ctrl+C Stop generation and return to the prompt.");

}

static void history_file_path(char *buf, size_t len) {

const char *home = getenv("HOME");

if (!home || !home[0]) home = ".";

snprintf(buf, len, "%s/.ds4_history", home);

}

typedef struct {

ds4_session *session;

ds4_tokens transcript;

int ctx_size;

int max_prefix_tokens;

} repl_chat;

static void tokens_insert(ds4_tokens *dst, int pos, const ds4_tokens *src) {

if (!src || src->len <= 0) return;

if (pos < 0) pos = 0;

if (pos > dst->len) pos = dst->len;

while (dst->len + src->len > dst->cap) {

dst->cap = dst->cap ? dst->cap * 2 : 64;

int *next = realloc(dst->v, (size_t)dst->cap * sizeof(dst->v[0]));

if (!next) {

perror("ds4: realloc");

exit(1);

}

dst->v = next;

}

memmove(dst->v + pos + src->len, dst->v + pos,

(size_t)(dst->len - pos) * sizeof(dst->v[0]));

memcpy(dst->v + pos, src->v, (size_t)src->len * sizeof(src->v[0]));

dst->len += src->len;

}

static void tokens_remove(ds4_tokens *dst, int pos, int n) {

if (n <= 0 || pos < 0 || pos >= dst->len) return;

if (pos + n > dst->len) n = dst->len - pos;

memmove(dst->v + pos, dst->v + pos + n,

(size_t)(dst->len - pos - n) * sizeof(dst->v[0]));

dst->len -= n;

}

/* Insert/remove the Think Max prefix inside the existing transcript. The

* prefix lives after BOS, before any system/developer text, which mirrors the

* API rendering path. Changing it invalidates the session because every later

* token position would otherwise refer to the wrong prefix. */

static void repl_chat_apply_max_prefix(ds4_engine *engine, repl_chat *chat, bool enable) {

if (enable && chat->max_prefix_tokens == 0) {

ds4_tokens prefix = {0};

ds4_chat_append_max_effort_prefix(engine, &prefix);

tokens_insert(&chat->transcript, 1, &prefix);

chat->max_prefix_tokens = prefix.len;

ds4_tokens_free(&prefix);

if (chat->session) ds4_session_invalidate(chat->session);

} else if (!enable && chat->max_prefix_tokens > 0) {

tokens_remove(&chat->transcript, 1, chat->max_prefix_tokens);

chat->max_prefix_tokens = 0;

if (chat->session) ds4_session_invalidate(chat->session);

}

}

static int repl_chat_create_session(ds4_engine *engine, repl_chat *chat, int ctx_size) {

ds4_session *session = NULL;

if (ds4_session_create(&session, engine, ctx_size) != 0) {

fprintf(stderr, "ds4: interactive chat KV cache requires the Metal backend\n");

return 1;

}

if (chat->session) ds4_session_free(chat->session);

chat->session = session;

chat->ctx_size = ctx_size;

return 0;

}

static int repl_chat_init(ds4_engine *engine, repl_chat *chat, const cli_config *cfg) {

memset(chat, 0, sizeof(*chat));

ds4_chat_begin(engine, &chat->transcript);

repl_chat_apply_max_prefix(engine, chat,

cli_effective_think_mode(&cfg->gen) == DS4_THINK_MAX);

if (cfg->gen.system && cfg->gen.system[0]) {

ds4_chat_append_message(engine, &chat->transcript, "system", cfg->gen.system);

}

return repl_chat_create_session(engine, chat, cfg->gen.ctx_size);

}

static void repl_chat_free(repl_chat *chat) {

if (!chat) return;

ds4_session_free(chat->session);

ds4_tokens_free(&chat->transcript);

memset(chat, 0, sizeof(*chat));

}

static int repl_chat_set_ctx(ds4_engine *engine, repl_chat *chat, int ctx_size) {

ds4_session_free(chat->session);

chat->session = NULL;

chat->ctx_size = 0;

return repl_chat_create_session(engine, chat, ctx_size);

}

/* Run one interactive turn. The transcript is tentatively extended with user

* and assistant markers, then ds4_session_sync() decides whether this is a KV

* continuation. If prompt processing fails, the transcript rolls back before

* returning to the prompt. */

static int run_chat_turn(ds4_engine *engine, cli_config *cfg, repl_chat *chat, const char *user_text) {

if (!chat->session) {

fprintf(stderr, "ds4: no active interactive KV cache\n");

return 1;

}

ds4_think_mode think_mode = ds4_think_mode_for_context(cfg->gen.think_mode,

chat->ctx_size);

repl_chat_apply_max_prefix(engine, chat, think_mode == DS4_THINK_MAX);

const int rollback_len = chat->transcript.len;

ds4_chat_append_message(engine, &chat->transcript, "user", user_text);

ds4_chat_append_assistant_prefix(engine, &chat->transcript, think_mode);

const int old_pos = ds4_session_pos(chat->session);

const int common = ds4_session_common_prefix(chat->session, &chat->transcript);

const int cached = common == old_pos && chat->transcript.len >= old_pos ? common : 0;

const int suffix = chat->transcript.len - cached;

char err[160];

cli_prefill_progress progress = {

.base_tokens = cached,

.input_tokens = suffix,

.use_color = ds4_log_is_tty(stderr),

};

const double t_prefill0 = cli_now_sec();

ds4_session_set_progress(chat->session, cli_prefill_progress_cb, &progress);

if (ds4_session_sync(chat->session, &chat->transcript, err, sizeof(err)) != 0) {

ds4_session_set_progress(chat->session, NULL, NULL);

chat->transcript.len = rollback_len;

fprintf(stderr, "ds4: prompt processing failed: %s\n", err);

return 1;

}

ds4_session_set_progress(chat->session, NULL, NULL);

const double t_prefill1 = cli_now_sec();

token_printer printer = {

.engine = engine,

.fp = stdout,

.format_thinking = ds4_think_mode_enabled(think_mode),

.in_think = ds4_think_mode_enabled(think_mode),

.use_color = isatty(fileno(stdout)) != 0,

.last_output_newline = true,

};

int max_tokens = cfg->gen.n_predict;

int room = ds4_session_ctx(chat->session) - ds4_session_pos(chat->session);

if (room <= 1) max_tokens = 0;

else if (max_tokens > room - 1) max_tokens = room - 1;

uint64_t rng = cfg->gen.seed ? cfg->gen.seed :

((uint64_t)time(NULL) ^ ((uint64_t)getpid() << 32) ^ (uint64_t)clock());

int generated = 0;

const double t_decode0 = cli_now_sec();

while (generated < max_tokens && !cli_interrupt_requested()) {

int token = ds4_session_sample(chat->session,

cfg->gen.temperature,

0,

cfg->gen.top_p,

0.0f,

&rng);

if (token == ds4_token_eos(engine)) break;

int toks[17];

int ntok = 0;

if (cfg->gen.temperature <= 0.0f && ds4_engine_mtp_draft_tokens(engine) > 1 &&

getenv("DS4_MTP_SPEC_DISABLE") == NULL) {

ntok = ds4_session_eval_speculative_argmax(chat->session,

token,

max_tokens - generated,

ds4_token_eos(engine),

toks,

(int)(sizeof(toks) / sizeof(toks[0])),

err,

sizeof(err));

if (ntok < 0) {

fprintf(stderr, "ds4: decode failed: %s\n", err);

return 1;

}

} else {

if (ds4_session_eval(chat->session, token, err, sizeof(err)) != 0) {

fprintf(stderr, "ds4: decode failed: %s\n", err);

return 1;

}

toks[0] = token;

ntok = 1;

}

bool stop = false;

for (int j = 0; j < ntok; j++) {

if (toks[j] == ds4_token_eos(engine)) {

stop = true;

break;

}

size_t piece_len = 0;

char *piece = ds4_token_text(engine, toks[j], &piece_len);

ds4_tokens_push(&chat->transcript, toks[j]);

token_printer_write_text(&printer, piece, piece_len);

fflush(stdout);

free(piece);

generated++;

if (generated >= max_tokens) break;

}

if (stop) break;

}

const double t_decode1 = cli_now_sec();

generation_done(&printer);

const bool interrupted = cli_interrupt_requested();

if (interrupted && generated == 0) {

chat->transcript.len = rollback_len;

ds4_session_invalidate(chat->session);

} else {

ds4_tokens_push(&chat->transcript, ds4_token_eos(engine));

}

const double prefill_s = t_prefill1 - t_prefill0;

const double decode_s = t_decode1 - t_decode0;

if (interrupted) cli_interrupt_clear();

ds4_log(stderr,

DS4_LOG_TIMING,

"ds4: prefill: %.2f t/s, generation: %.2f t/s\n",

prefill_s > 0.0 ? (double)suffix / prefill_s : 0.0,

decode_s > 0.0 ? (double)generated / decode_s : 0.0);

return 0;

}