vrr/koboldcpp
2
0
Fork 0
mirror of https://github.com/LostRuins/koboldcpp.git synced 2026-08-04 05:29:57 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Merge branch 'upstream' into concedo_experimental

Browse source 
# Conflicts:
#	ggml/src/ggml-opencl/ggml-opencl.cpp
#	tests/test-backend-ops.cpp
#	tests/test-quantize-fns.cpp
#	tools/server/README.md
#	tools/ui/src/lib/constants/settings-registry.ts
This commit is contained in:
Concedo 2026-07-11 09:20:40 +08:00
commit 0f0245161e
103 changed files with 2633 additions and 1170 deletions
Download patch file Download diff file Expand all files Collapse all files

8
common/arg.cpp
View file

@ -28,6 +28,7 @@
#include <cinttypes>
#include <climits>
#include <cstdarg>
#include <filesystem>
#include <fstream>
#include <list>
#include <regex>
@ -3452,9 +3453,14 @@ common_params_context common_params_parser_init(common_params & params, llama_ex
).set_env("LLAMA_ARG_LOG_FILE"));
add_opt(common_arg(
{"--log-prompts-dir"}, "PATH",
"Log prompts to directory (only used for debugging, default: disabled)",
"Log prompts to directory (auto-created if not present; only used for debugging, default: disabled)",
[](common_params & params, const std::string & value) {
params.path_prompts_log_dir = value;
std::error_code ec;
std::filesystem::create_directories(value, ec);
if (ec) {
fprintf(stderr, "warning: failed to create prompts-log-dir '%s': %s\n", value.c_str(), ec.message().c_str());
}
}
).set_examples({LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI}));
add_opt(common_arg(

1
common/common.h
View file

@ -15,6 +15,7 @@
#include <vector>
#include <map>
#include <algorithm>
#include <fstream>
#if defined(_WIN32) && !defined(_WIN32_WINNT)
#define _WIN32_WINNT 0x0A00

106
common/speculative.cpp
View file

@ -2221,6 +2221,112 @@ int32_t common_speculative_n_max(const common_params_speculative * spec) {
return n_max;
}
common_params common_base_params_to_speculative(const common_params & params) {
const bool has_draft = params.speculative.has_dft();
const auto & params_spec = params.speculative.draft;
common_params result = params;
if (has_draft) {
result.devices = params_spec.devices;
result.model = params_spec.mparams;
result.n_gpu_layers = params_spec.n_gpu_layers;
result.tensor_buft_overrides = params_spec.tensor_buft_overrides;
if (params_spec.cpuparams.n_threads > 0) {
result.cpuparams.n_threads = params_spec.cpuparams.n_threads;
result.cpuparams_batch.n_threads = params_spec.cpuparams_batch.n_threads;
}
}
result.cache_type_k = params_spec.cache_type_k;
result.cache_type_v = params_spec.cache_type_v;
result.n_outputs_max = params.n_parallel;
return result;
}
struct common_speculative_init_result::impl {
impl() = default;
~impl() = default;
// note: the order in which model, context, etc. are declared matters because their destructors will be called bottom-to-top
llama_model_ptr model;
llama_context_ptr context;
};
common_speculative_init_result::common_speculative_init_result(
common_params & params,
llama_model * model_tgt,
llama_context * ctx_tgt) :
pimpl(new impl{}) {
const bool has_draft = params.speculative.has_dft();
const bool spec_mtp = std::find(params.speculative.types.begin(),
params.speculative.types.end(),
COMMON_SPECULATIVE_TYPE_DRAFT_MTP) != params.speculative.types.end();
GGML_ASSERT(has_draft || spec_mtp);
auto mparams = common_model_params_to_llama(params);
auto cparams = common_context_params_to_llama(params);
if (spec_mtp) {
cparams.ctx_type = LLAMA_CONTEXT_TYPE_MTP;
}
// note: for small models maybe we can set this to the maximum possible draft from all speculative types
// the extra memory for small models is likely negligible?
cparams.n_rs_seq = 0;
cparams.ctx_other = ctx_tgt;
std::string model_path;
if (has_draft) {
model_path = params.speculative.draft.mparams.path;
LOG_TRC("%s: loading draft model '%s'\n", __func__, model_path.c_str());
llama_model * model_dft = llama_model_load_from_file(params.model.path.c_str(), mparams);
if (model_dft == NULL) {
LOG_ERR("%s: failed to load draft model, '%s'\n", __func__, model_path.c_str());
return;
}
pimpl->model.reset(model_dft);
llama_context * ctx_dft = llama_init_from_model(model_dft, cparams);
if (ctx_dft == nullptr) {
LOG_ERR("%s: failed to create MTP context\n", __func__);
return;
}
pimpl->context.reset(ctx_dft);
} else if (spec_mtp) {
model_path = params.model.path;
LOG_TRC("%s: creating MTP draft context against the target model '%s'\n", __func__, model_path.c_str());
llama_context * ctx_dft = llama_init_from_model(model_tgt, cparams);
if (ctx_dft == nullptr) {
LOG_ERR("%s: failed to create MTP context\n", __func__);
return;
}
pimpl->context.reset(ctx_dft);
}
}
common_speculative_init_result::~common_speculative_init_result() = default;
llama_model * common_speculative_init_result::model() {
return pimpl->model.get();
}
llama_context * common_speculative_init_result::context() {
return pimpl->context.get();
}
common_speculative_init_result_ptr common_speculative_init_from_params(common_params & params, llama_model * model_tgt, llama_context * ctx_tgt) {
return std::make_unique<common_speculative_init_result>(params, model_tgt, ctx_tgt);
}
// initialization of the speculative decoding system
//
common_speculative * common_speculative_init(common_params_speculative & params, uint32_t n_seq) {

18
common/speculative.h
View file

@ -23,6 +23,8 @@ std::string common_speculative_type_to_str(enum common_speculative_type type);
// return the max number of draft tokens based on the speculative parameters
int32_t common_speculative_n_max(const common_params_speculative * spec);
common_params common_base_params_to_speculative(const common_params & params);
common_speculative * common_speculative_init(common_params_speculative & params, uint32_t n_seq);
void common_speculative_free(common_speculative * spec);
@ -80,3 +82,19 @@ struct common_speculative_deleter {
};
typedef std::unique_ptr<common_speculative, common_speculative_deleter> common_speculative_ptr;
struct common_speculative_init_result {
common_speculative_init_result(common_params & params, llama_model * model_tgt, llama_context * ctx_tgt);
~common_speculative_init_result();
llama_model * model();
llama_context * context();
private:
struct impl;
std::unique_ptr<impl> pimpl;
};
using common_speculative_init_result_ptr = std::unique_ptr<common_speculative_init_result>;
common_speculative_init_result_ptr common_speculative_init_from_params(common_params & params, llama_model * model_tgt, llama_context * ctx_tgt);

4
examples/llama-eval/llama-eval.py
View file

@ -362,7 +362,7 @@ class EvalState:
case = cases.get(task_id, {})
status = case.get("status", "pending")
expected = case.get("expected", "")
answer = case.get("answer", "") if status == "ok" else ""
answer = case.get("answer") or "" if status == "ok" else ""
is_correct = case.get("correct", False) if status == "ok" else False
response = case.get("response", "") or ""
prompt = case.get("prompt", "") or ""
@ -647,7 +647,7 @@ class EvalState:
question, prompt, expected = self.get_case(i)
case = cases.get(task_id, {})
status = case.get("status", "pending")
answer = case.get("answer", "N/A") if status == "ok" else "N/A"
answer = case.get("answer") or "N/A" if status == "ok" else "N/A"
tokens = case.get("tokens")
tokens_str = str(tokens) if tokens is not None else "N/A"
tps_gen = case.get("tps_gen")

4
ggml/include/ggml.h
View file

@ -435,7 +435,8 @@ extern "C" {
GGML_TYPE_MXFP4 = 39, // MXFP4 (1 block)
GGML_TYPE_NVFP4 = 40, // NVFP4 (4 blocks, E4M3 scale)
GGML_TYPE_Q1_0 = 41,
GGML_TYPE_COUNT = 42,
GGML_TYPE_Q2_0 = 42,
GGML_TYPE_COUNT = 43,
};
// precision
@ -479,6 +480,7 @@ extern "C" {
GGML_FTYPE_MOSTLY_MXFP4 = 25, // except 1d tensors
GGML_FTYPE_MOSTLY_NVFP4 = 26, // except 1d tensors
GGML_FTYPE_MOSTLY_Q1_0 = 27, // except 1d tensors
GGML_FTYPE_MOSTLY_Q2_0 = 28, // except 1d tensors
};
// available tensor operations:

10
ggml/src/ggml-common.h
View file

@ -96,6 +96,9 @@ typedef sycl::half2 ggml_half2;
#define QI1_0 (QK1_0 / 32)
#define QR1_0 1
#define QI2_0 (QK2_0 / 32)
#define QR2_0 1
#define QI4_0 (QK4_0 / (4 * QR4_0))
#define QR4_0 2
@ -181,6 +184,13 @@ typedef struct {
} block_q1_0;
static_assert(sizeof(block_q1_0) == sizeof(ggml_half) + QK1_0 / 8, "wrong q1_0 block size/padding");
#define QK2_0 64
typedef struct {
ggml_half d; // delta (scale)
uint8_t qs[QK2_0 / 4]; // 2 bits per element
} block_q2_0;
static_assert(sizeof(block_q2_0) == sizeof(ggml_half) + QK2_0 / 4, "wrong q2_0 block size/padding");
#define QK4_0 32
typedef struct {
ggml_half d; // delta

7
ggml/src/ggml-cpu/arch-fallback.h
View file

@ -17,6 +17,7 @@
#define ggml_vec_dot_mxfp4_q8_0_generic ggml_vec_dot_mxfp4_q8_0
#define ggml_vec_dot_nvfp4_q8_0_generic ggml_vec_dot_nvfp4_q8_0
#define ggml_vec_dot_q1_0_q8_0_generic ggml_vec_dot_q1_0_q8_0
#define ggml_vec_dot_q2_0_q8_0_generic ggml_vec_dot_q2_0_q8_0
#define ggml_vec_dot_tq1_0_q8_K_generic ggml_vec_dot_tq1_0_q8_K
#define ggml_vec_dot_tq2_0_q8_K_generic ggml_vec_dot_tq2_0_q8_K
#define ggml_vec_dot_q2_K_q8_K_generic ggml_vec_dot_q2_K_q8_K
@ -72,6 +73,7 @@
#define ggml_gemm_q2_K_8x8_q8_K_generic ggml_gemm_q2_K_8x8_q8_K
#elif defined(__x86_64__) || defined(__i386__) || defined(_M_IX86) || defined(_M_X64)
// quants.c
#define ggml_vec_dot_q2_0_q8_0_generic ggml_vec_dot_q2_0_q8_0
// repack.cpp
#define ggml_quantize_mat_q8_0_4x4_generic ggml_quantize_mat_q8_0_4x4
#define ggml_quantize_mat_q8_K_4x4_generic ggml_quantize_mat_q8_K_4x4
@ -101,6 +103,7 @@
#define quantize_row_q8_K_generic quantize_row_q8_K
#define ggml_vec_dot_nvfp4_q8_0_generic ggml_vec_dot_nvfp4_q8_0
#define ggml_vec_dot_q1_0_q8_0_generic ggml_vec_dot_q1_0_q8_0
#define ggml_vec_dot_q2_0_q8_0_generic ggml_vec_dot_q2_0_q8_0
#define ggml_vec_dot_tq1_0_q8_K_generic ggml_vec_dot_tq1_0_q8_K
#define ggml_vec_dot_tq2_0_q8_K_generic ggml_vec_dot_tq2_0_q8_K
#define ggml_vec_dot_iq1_m_q8_K_generic ggml_vec_dot_iq1_m_q8_K
@ -144,6 +147,7 @@
#define ggml_vec_dot_mxfp4_q8_0_generic ggml_vec_dot_mxfp4_q8_0
#define ggml_vec_dot_nvfp4_q8_0_generic ggml_vec_dot_nvfp4_q8_0
#define ggml_vec_dot_q1_0_q8_0_generic ggml_vec_dot_q1_0_q8_0
#define ggml_vec_dot_q2_0_q8_0_generic ggml_vec_dot_q2_0_q8_0
// repack.cpp
#define ggml_quantize_mat_q8_0_4x4_generic ggml_quantize_mat_q8_0_4x4
#define ggml_quantize_mat_q8_0_4x8_generic ggml_quantize_mat_q8_0_4x8
@ -178,6 +182,7 @@
#elif defined(__riscv)
// quants.c
#define ggml_vec_dot_nvfp4_q8_0_generic ggml_vec_dot_nvfp4_q8_0
#define ggml_vec_dot_q2_0_q8_0_generic ggml_vec_dot_q2_0_q8_0
// repack.cpp
#define ggml_quantize_mat_q8_0_4x4_generic ggml_quantize_mat_q8_0_4x4
#define ggml_quantize_mat_q8_0_4x8_generic ggml_quantize_mat_q8_0_4x8
@ -212,6 +217,7 @@
#define quantize_row_q8_K_generic quantize_row_q8_K
#define ggml_vec_dot_nvfp4_q8_0_generic ggml_vec_dot_nvfp4_q8_0
#define ggml_vec_dot_q1_0_q8_0_generic ggml_vec_dot_q1_0_q8_0
#define ggml_vec_dot_q2_0_q8_0_generic ggml_vec_dot_q2_0_q8_0
#define ggml_vec_dot_tq1_0_q8_K_generic ggml_vec_dot_tq1_0_q8_K
#define ggml_vec_dot_tq2_0_q8_K_generic ggml_vec_dot_tq2_0_q8_K
#define ggml_vec_dot_q2_K_q8_K_generic ggml_vec_dot_q2_K_q8_K
@ -269,6 +275,7 @@
#define ggml_vec_dot_mxfp4_q8_0_generic ggml_vec_dot_mxfp4_q8_0
#define ggml_vec_dot_nvfp4_q8_0_generic ggml_vec_dot_nvfp4_q8_0
#define ggml_vec_dot_q1_0_q8_0_generic ggml_vec_dot_q1_0_q8_0
#define ggml_vec_dot_q2_0_q8_0_generic ggml_vec_dot_q2_0_q8_0
// repack.cpp
#define ggml_quantize_mat_q8_0_4x4_generic ggml_quantize_mat_q8_0_4x4
#define ggml_quantize_mat_q8_0_4x8_generic ggml_quantize_mat_q8_0_4x8

74
ggml/src/ggml-cpu/arch/arm/quants.c
View file

@ -219,6 +219,80 @@ void ggml_vec_dot_q1_0_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const voi
#endif
}
void ggml_vec_dot_q2_0_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc) {
const int qk = QK2_0;
const int nb = n / qk;
assert(n % qk == 0);
assert(nrc == 1);
UNUSED(nrc);
UNUSED(bx);
UNUSED(by);
UNUSED(bs);
const block_q2_0 * GGML_RESTRICT x = vx;
const block_q8_0 * GGML_RESTRICT y = vy;
float sumf = 0.0f;
#if defined(__ARM_NEON)
// Replicate pattern: each byte repeated 4 times
static const uint8_t tbl_idx_lo[16] = {0,0,0,0, 1,1,1,1, 2,2,2,2, 3,3,3,3};
static const uint8_t tbl_idx_hi[16] = {4,4,4,4, 5,5,5,5, 6,6,6,6, 7,7,7,7};
// Right-shift amounts: 0,2,4,6 repeated for each group of 4
static const int8_t shift_vals[16] = {0,-2,-4,-6, 0,-2,-4,-6, 0,-2,-4,-6, 0,-2,-4,-6};
const uint8x16_t idx_lo = vld1q_u8(tbl_idx_lo);
const uint8x16_t idx_hi = vld1q_u8(tbl_idx_hi);
const int8x16_t shifts = vld1q_s8(shift_vals);
const uint8x16_t mask2 = vdupq_n_u8(0x03);
const int8x16_t one = vdupq_n_s8(1);
float32x4_t sumv = vdupq_n_f32(0.0f);
for (int i = 0; i < nb; i++) {
const float d0 = GGML_CPU_FP16_TO_FP32(x[i].d);
// group 64: one Q2_0 block (64 weights) maps to two Q8_0 blocks (2 * 32 = 64)
for (int k = 0; k < 2; k++) {
const block_q8_0 * GGML_RESTRICT yb = &y[i * 2 + k];
const float d1 = GGML_CPU_FP16_TO_FP32(yb->d);
// Load 8 bytes of packed 2-bit values
const uint8x8_t raw = vld1_u8(&x[i].qs[k * 8]);
const uint8x16_t raw16 = vcombine_u8(raw, raw);
// First 16 elements: replicate bytes 0-3, shift, mask, subtract 1
uint8x16_t bytes0 = vqtbl1q_u8(raw16, idx_lo);
int8x16_t qv0 = vsubq_s8(
vreinterpretq_s8_u8(vandq_u8(vshlq_u8(bytes0, shifts), mask2)),
one);
// Second 16 elements: replicate bytes 4-7, shift, mask, subtract 1
uint8x16_t bytes1 = vqtbl1q_u8(raw16, idx_hi);
int8x16_t qv1 = vsubq_s8(
vreinterpretq_s8_u8(vandq_u8(vshlq_u8(bytes1, shifts), mask2)),
one);
// Load Q8_0 values and dot product
const int8x16_t y0 = vld1q_s8(yb->qs);
const int8x16_t y1 = vld1q_s8(yb->qs + 16);
int32x4_t p0 = ggml_vdotq_s32(vdupq_n_s32(0), qv0, y0);
int32x4_t p1 = ggml_vdotq_s32(p0, qv1, y1);
sumv = vmlaq_n_f32(sumv, vcvtq_f32_s32(p1), d0 * d1);
}
}
sumf = vaddvq_f32(sumv);
#else
ggml_vec_dot_q2_0_q8_0_generic(n, s, bs, vx, bx, vy, by, nrc);
return;
#endif
*s = sumf;
}
void ggml_vec_dot_q4_0_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc) {
const int qk = QK8_0;

6
ggml/src/ggml-cpu/ggml-cpu.c
View file

@ -231,6 +231,12 @@ static const struct ggml_type_traits_cpu type_traits_cpu[GGML_TYPE_COUNT] = {
.vec_dot_type = GGML_TYPE_Q8_0,
.nrows = 1,
},
[GGML_TYPE_Q2_0] = {
.from_float = quantize_row_q2_0,
.vec_dot = ggml_vec_dot_q2_0_q8_0,
.vec_dot_type = GGML_TYPE_Q8_0,
.nrows = 1,
},
[GGML_TYPE_Q4_0] = {
.from_float = quantize_row_q4_0,
.vec_dot = ggml_vec_dot_q4_0_q8_0,

48
ggml/src/ggml-cpu/ops.cpp
View file

@ -665,6 +665,7 @@ void ggml_compute_forward_add(
ggml_compute_forward_add_non_quantized(params, dst);
} break;
case GGML_TYPE_Q1_0:
case GGML_TYPE_Q2_0:
case GGML_TYPE_Q4_0:
case GGML_TYPE_Q4_1:
case GGML_TYPE_Q5_0:
@ -1115,6 +1116,7 @@ void ggml_compute_forward_add1(
}
} break;
case GGML_TYPE_Q1_0:
case GGML_TYPE_Q2_0:
case GGML_TYPE_Q4_0:
case GGML_TYPE_Q4_1:
case GGML_TYPE_Q5_0:
@ -1245,6 +1247,7 @@ void ggml_compute_forward_acc(
case GGML_TYPE_F16:
case GGML_TYPE_BF16:
case GGML_TYPE_Q1_0:
case GGML_TYPE_Q2_0:
case GGML_TYPE_Q4_0:
case GGML_TYPE_Q4_1:
case GGML_TYPE_Q5_0:
@ -4454,6 +4457,7 @@ void ggml_compute_forward_out_prod(
switch (src0->type) {
case GGML_TYPE_Q1_0:
case GGML_TYPE_Q2_0:
case GGML_TYPE_Q4_0:
case GGML_TYPE_Q4_1:
case GGML_TYPE_Q5_0:
@ -4730,6 +4734,7 @@ void ggml_compute_forward_set(
case GGML_TYPE_F16:
case GGML_TYPE_BF16:
case GGML_TYPE_Q1_0:
case GGML_TYPE_Q2_0:
case GGML_TYPE_Q4_0:
case GGML_TYPE_Q4_1:
case GGML_TYPE_Q5_0:
@ -4954,6 +4959,7 @@ void ggml_compute_forward_get_rows(
switch (src0->type) {
case GGML_TYPE_Q1_0:
case GGML_TYPE_Q2_0:
case GGML_TYPE_Q4_0:
case GGML_TYPE_Q4_1:
case GGML_TYPE_Q5_0:
@ -5019,8 +5025,8 @@ void ggml_compute_forward_get_rows(
//}
}
template<typename idx_t>
static void ggml_compute_forward_set_rows_f32(
template<typename src_t, typename idx_t>
static void ggml_compute_forward_set_rows_impl(
const ggml_compute_params * params,
ggml_tensor * dst) {
@ -5035,7 +5041,7 @@ static void ggml_compute_forward_set_rows_f32(
assert(ne0 == nc);
assert(ne2 == ne02);
assert(ne3 == ne03);
assert(src0->type == GGML_TYPE_F32);
GGML_ASSERT(src0->type == GGML_TYPE_F32 || (src0->type == GGML_TYPE_F16 && dst->type == GGML_TYPE_F16));
assert(ne02 % ne11 == 0);
assert(ne03 % ne12 == 0);
@ -5049,6 +5055,8 @@ static void ggml_compute_forward_set_rows_f32(
const int64_t ir0 = dr*ith;
const int64_t ir1 = std::min(ir0 + dr, nr);
const size_t rs = ggml_row_size(src0->type, nc);
ggml_from_float_t const from_float = ggml_get_type_traits_cpu(dst->type)->from_float;
for (int64_t i03 = 0; i03 < ne03; ++i03) {
@ -5062,9 +5070,18 @@ static void ggml_compute_forward_set_rows_f32(
GGML_ASSERT(i1 >= 0 && i1 < ne1);
from_float(
(const float *) ((char *) src0->data + i*nb01 + i02*nb02 + i03*nb03),
((char *) dst->data + i1*nb1 + i02*nb2 + i03*nb3), nc);
if constexpr (std::is_same_v<src_t, float>) {
from_float(
(const float *) ((char *) src0->data + i*nb01 + i02*nb02 + i03*nb03),
((char *) dst->data + i1*nb1 + i02*nb2 + i03*nb3), nc);
} else if constexpr (std::is_same_v<src_t, ggml_fp16_t>) {
memcpy(
((char *) dst->data + i1*nb1 + i02*nb2 + i03*nb3),
((char *) src0->data + i*nb01 + i02*nb02 + i03*nb03),
rs);
} else {
GGML_ABORT("src0->type = %d (%s) not supported", src0->type, ggml_type_name(src0->type));
}
}
}
}
@ -5081,13 +5098,27 @@ void ggml_compute_forward_set_rows(
case GGML_TYPE_F32:
{
if (src1->type == GGML_TYPE_I64) {
ggml_compute_forward_set_rows_f32<int64_t>(params, dst);
ggml_compute_forward_set_rows_impl<float, int64_t>(params, dst);
} else if (src1->type == GGML_TYPE_I32) {
ggml_compute_forward_set_rows_f32<int32_t>(params, dst);
ggml_compute_forward_set_rows_impl<float, int32_t>(params, dst);
} else {
GGML_ABORT("src1->type = %d (%s) not supported", src1->type, ggml_type_name(src1->type));
}
} break;
case GGML_TYPE_F16:
{
if (dst->type == GGML_TYPE_F16) {
if (src1->type == GGML_TYPE_I64) {
ggml_compute_forward_set_rows_impl<ggml_fp16_t, int64_t>(params, dst);
} else if (src1->type == GGML_TYPE_I32) {
ggml_compute_forward_set_rows_impl<ggml_fp16_t, int32_t>(params, dst);
} else {
GGML_ABORT("src1->type = %d (%s) not supported", src1->type, ggml_type_name(src1->type));
}
} else {
GGML_ABORT("dst->type = %d (%s) not supported with src0->type = %d (%s)", dst->type, ggml_type_name(dst->type), src0->type, ggml_type_name(src0->type));
}
} break;
default:
{
GGML_ABORT("src0->type = %d (%s) not supported", src0->type, ggml_type_name(src0->type));
@ -5680,6 +5711,7 @@ void ggml_compute_forward_clamp(
} break;
case GGML_TYPE_BF16:
case GGML_TYPE_Q1_0:
case GGML_TYPE_Q2_0:
case GGML_TYPE_Q4_0:
case GGML_TYPE_Q4_1:
case GGML_TYPE_Q5_0:

51
ggml/src/ggml-cpu/quants.c
View file

@ -26,6 +26,10 @@ void quantize_row_q1_0(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, in
quantize_row_q1_0_ref(x, y, k);
}
void quantize_row_q2_0(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k) {
quantize_row_q2_0_ref(x, y, k);
}
void quantize_row_q4_0(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k) {
quantize_row_q4_0_ref(x, y, k);
}
@ -170,6 +174,53 @@ void ggml_vec_dot_q1_0_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, c
*s = sumf;
}
void ggml_vec_dot_q2_0_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc) {
const int qk = QK2_0;
const int nb = n / qk;
assert(n % qk == 0);
assert(nrc == 1);
UNUSED(nrc);
UNUSED(bx);
UNUSED(by);
UNUSED(bs);
const block_q2_0 * GGML_RESTRICT x = vx;
const block_q8_0 * GGML_RESTRICT y = vy;
float sumf = 0.0f;
for (int i = 0; i < nb; i++) {
const float d0 = GGML_CPU_FP16_TO_FP32(x[i].d);
float sumi = 0.0f;
// group 64: one Q2_0 block (64 weights) maps to two Q8_0 blocks (2 * 32 = 64)
for (int k = 0; k < 2; k++) {
const block_q8_0 * GGML_RESTRICT yb = &y[i * 2 + k];
const float d1 = GGML_CPU_FP16_TO_FP32(yb->d);
int sumi_block = 0;
const uint8_t * GGML_RESTRICT qs = &x[i].qs[k * 8];
const int8_t * GGML_RESTRICT qy = yb->qs;
for (int b = 0; b < 8; ++b) {
const uint8_t byte = qs[b];
// Extract 4 two-bit values, map {0,1,2,3} -> {-1,0,1,2}
sumi_block += ((int)((byte >> 0) & 3) - 1) * qy[b*4 + 0];
sumi_block += ((int)((byte >> 2) & 3) - 1) * qy[b*4 + 1];
sumi_block += ((int)((byte >> 4) & 3) - 1) * qy[b*4 + 2];
sumi_block += ((int)((byte >> 6) & 3) - 1) * qy[b*4 + 3];
}
sumi += d1 * sumi_block;
}
sumf += d0 * sumi;
}
*s = sumf;
}
void ggml_vec_dot_q4_0_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc) {
const int qk = QK8_0;

3
ggml/src/ggml-cpu/quants.h
View file

@ -13,6 +13,7 @@ extern "C" {
// Quantization
void quantize_row_q1_0(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k);
void quantize_row_q2_0(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k);
void quantize_row_q4_0(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k);
void quantize_row_q4_1(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k);
void quantize_row_q5_0(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k);
@ -38,6 +39,7 @@ void quantize_row_iq4_xs (const float * GGML_RESTRICT x, void * GGML_RESTRICT y,
// Dot product
void ggml_vec_dot_q1_0_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc);
void ggml_vec_dot_q2_0_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc);
void ggml_vec_dot_q4_0_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc);
void ggml_vec_dot_q4_1_q8_1(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc);
void ggml_vec_dot_q5_0_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc);
@ -71,6 +73,7 @@ void quantize_row_q8_0_generic(const float * GGML_RESTRICT x, void * GGML_RESTRI
void quantize_row_q8_1_generic(const float * GGML_RESTRICT x, void * GGML_RESTRICT vy, int64_t k);
void quantize_row_q8_K_generic(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k);
void ggml_vec_dot_q1_0_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc);
void ggml_vec_dot_q2_0_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc);
void ggml_vec_dot_q4_0_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc);
void ggml_vec_dot_q4_1_q8_1_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc);
void ggml_vec_dot_q5_0_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc);

2
ggml/src/ggml-cpu/simd-gemm.h
View file

@ -78,7 +78,7 @@ static void simd_gemm(
for (int64_t i = 0; i < GEMM_RM; i++) {
float a = C[i * N + jj];
for (int64_t kk = 0; kk < K; kk++) {
a += A[i + kk] * B[kk * N + jj];
a += A[i * K + kk] * B[kk * N + jj];
}
C[i * N + jj] = a;
}

4
ggml/src/ggml-cuda/common.cuh
View file

@ -1512,12 +1512,16 @@ struct ggml_cuda_mm_fusion_args_host {
const ggml_tensor * x_bias = nullptr;
const ggml_tensor * gate = nullptr;
const ggml_tensor * gate_bias = nullptr;
const ggml_tensor * x_scale = nullptr;
const ggml_tensor * gate_scale = nullptr;
ggml_glu_op glu_op;
};
struct ggml_cuda_mm_fusion_args_device {
const void * x_bias = nullptr;
const void * gate = nullptr;
const void * gate_bias = nullptr;
const void * x_scale = nullptr;
const void * gate_scale = nullptr;
ggml_glu_op glu_op;
};

396
ggml/src/ggml-cuda/ggml-cuda.cu
View file

@ -1582,12 +1582,18 @@ static bool ggml_cuda_should_fuse_mul_mat(const ggml_tensor * ffn_up,
const ggml_tensor * ffn_gate,
const ggml_tensor * glu,
const ggml_tensor * ffn_up_bias = nullptr,
const ggml_tensor * ffn_gate_bias = nullptr) {
const ggml_tensor * ffn_gate_bias = nullptr,
const ggml_tensor * ffn_up_scale = nullptr,
const ggml_tensor * ffn_gate_scale = nullptr) {
const bool has_bias = ffn_up_bias != nullptr || ffn_gate_bias != nullptr;
const bool has_scale = ffn_up_scale != nullptr || ffn_gate_scale != nullptr;
if (has_bias && (!ffn_up_bias || !ffn_gate_bias)) {
return false;
}
if (has_scale && (!ffn_up_scale || !ffn_gate_scale)) {
return false;
}
const bool is_mul_mat = ffn_up->op == GGML_OP_MUL_MAT && ffn_gate->op == GGML_OP_MUL_MAT && glu->op == GGML_OP_GLU;
const bool is_mul_mat_id = ffn_up->op == GGML_OP_MUL_MAT_ID && ffn_gate->op == GGML_OP_MUL_MAT_ID && glu->op == GGML_OP_GLU;
@ -1599,34 +1605,45 @@ static bool ggml_cuda_should_fuse_mul_mat(const ggml_tensor * ffn_up,
}
const ggml_op expected_bias_op = is_mul_mat ? GGML_OP_ADD : GGML_OP_ADD_ID;
const ggml_tensor * ffn_up_bias_src = has_scale ? ffn_up_scale : ffn_up;
const ggml_tensor * ffn_gate_bias_src = has_scale ? ffn_gate_scale : ffn_gate;
const ggml_tensor * ffn_up_out = has_bias ? ffn_up_bias : ffn_up_bias_src;
const ggml_tensor * ffn_gate_out = has_bias ? ffn_gate_bias : ffn_gate_bias_src;
if (glu->src[0] != ffn_gate_out || glu->src[1] != ffn_up_out) {
return false;
}
if (has_scale) {
if (ffn_up_scale->op != GGML_OP_MUL || ffn_gate_scale->op != GGML_OP_MUL) {
return false;
}
const bool up_has_mm = ffn_up_scale->src[0] == ffn_up || ffn_up_scale->src[1] == ffn_up;
const bool gate_has_mm = ffn_gate_scale->src[0] == ffn_gate || ffn_gate_scale->src[1] == ffn_gate;
if (!up_has_mm || !gate_has_mm) {
return false;
}
}
if (has_bias) {
if (ffn_up_bias->op != expected_bias_op || ffn_gate_bias->op != expected_bias_op) {
return false;
}
if (glu->src[0] != ffn_gate_bias || glu->src[1] != ffn_up_bias) {
return false;
}
if (expected_bias_op == GGML_OP_ADD) {
const bool up_has_mul = ffn_up_bias->src[0] == ffn_up || ffn_up_bias->src[1] == ffn_up;
const bool gate_has_mul = ffn_gate_bias->src[0] == ffn_gate || ffn_gate_bias->src[1] == ffn_gate;
const bool up_has_mul = ffn_up_bias->src[0] == ffn_up_bias_src || ffn_up_bias->src[1] == ffn_up_bias_src;
const bool gate_has_mul = ffn_gate_bias->src[0] == ffn_gate_bias_src || ffn_gate_bias->src[1] == ffn_gate_bias_src;
if (!up_has_mul || !gate_has_mul) {
return false;
}
} else { // GGML_OP_ADD_ID
if (ffn_up_bias->src[0] != ffn_up || ffn_gate_bias->src[0] != ffn_gate) {
if (ffn_up_bias->src[0] != ffn_up_bias_src || ffn_gate_bias->src[0] != ffn_gate_bias_src) {
return false;
}
if (ffn_up_bias->src[2] != ffn_up->src[2] || ffn_gate_bias->src[2] != ffn_gate->src[2]) {
return false;
}
}
} else {
if (glu->src[0] != ffn_gate && glu->src[1] != ffn_up) {
return false;
}
}
if (ffn_up->src[0]->type != ffn_gate->src[0]->type || !ggml_are_same_shape(ffn_up->src[0], ffn_gate->src[0]) ||
@ -1638,7 +1655,7 @@ static bool ggml_cuda_should_fuse_mul_mat(const ggml_tensor * ffn_up,
return false;
}
if (ffn_up->src[2] && (ffn_up->src[2] != ffn_gate->src[2])) {
if (is_mul_mat_id && ffn_up->src[2] != ffn_gate->src[2]) {
return false;
}
@ -3212,10 +3229,240 @@ static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph
bool fused_mul_mat_vec = false;
int fused_node_count = 0;
// gate + glu + up
auto get_mul_mat_scale = [](const ggml_tensor * scale_node, const ggml_tensor * mm_node) -> const ggml_tensor * {
const bool scale_lhs_mm = scale_node->src[0] == mm_node;
const bool scale_rhs_mm = scale_node->src[1] == mm_node;
if (!scale_lhs_mm && !scale_rhs_mm) {
return nullptr;
}
const ggml_tensor * scale = scale_lhs_mm ? scale_node->src[1] : scale_node->src[0];
if (mm_node->src[0]->type != GGML_TYPE_NVFP4 || scale_node->type != GGML_TYPE_F32 ||
scale->type != GGML_TYPE_F32 || !ggml_is_contiguous(scale) || ggml_nelements(scale) != 1 ||
!ggml_are_same_shape(scale_node, mm_node)) {
return nullptr;
}
return scale;
};
auto get_mul_mat_id_scale = [](const ggml_tensor * reshape, const ggml_tensor * repeat, const ggml_tensor * getrows,
const ggml_tensor * scale_node, const ggml_tensor * mm_node) -> const ggml_tensor * {
if (repeat->src[0] != reshape || getrows->src[0] != repeat || getrows->src[1] != mm_node->src[2]) {
return nullptr;
}
if (!((scale_node->src[0] == mm_node && scale_node->src[1] == getrows) ||
(scale_node->src[0] == getrows && scale_node->src[1] == mm_node))) {
return nullptr;
}
const ggml_tensor * scale = reshape->src[0];
if (mm_node->src[0]->type != GGML_TYPE_NVFP4 || scale_node->type != GGML_TYPE_F32 ||
scale->type != GGML_TYPE_F32 || !ggml_is_contiguous(scale) || ggml_nelements(scale) != mm_node->src[0]->ne[2] ||
!ggml_are_same_shape(scale_node, mm_node)) {
return nullptr;
}
return scale;
};
auto get_bias_tensor = [](const ggml_tensor * bias_node, const ggml_tensor * mul_node, ggml_op op_bias) -> const ggml_tensor * {
if (op_bias == GGML_OP_ADD) {
if (bias_node->src[0] == mul_node) {
return bias_node->src[1];
}
if (bias_node->src[1] == mul_node) {
return bias_node->src[0];
}
return nullptr;
}
GGML_ASSERT(op_bias == GGML_OP_ADD_ID);
GGML_ASSERT(bias_node->src[0] == mul_node);
return bias_node->src[1];
};
// gate + glu + up, with optional scale/bias on both lanes.
for (ggml_op op : { GGML_OP_MUL_MAT, GGML_OP_MUL_MAT_ID }) {
const ggml_op bias_op = op == GGML_OP_MUL_MAT ? GGML_OP_ADD : GGML_OP_ADD_ID;
if (op == GGML_OP_MUL_MAT) {
for (const bool with_bias : { false, true }) {
const int gate_idx = i;
const int gate_scale_idx = i + 1;
const int gate_bias_idx = with_bias ? i + 2 : -1;
const int up_idx = with_bias ? i + 3 : i + 2;
const int up_scale_idx = up_idx + 1;
const int up_bias_idx = with_bias ? up_idx + 2 : -1;
const int glu_idx = with_bias ? up_idx + 3 : up_idx + 2;
const int out_nodes[] = { glu_idx };
ggml_op ops[7];
if (with_bias) {
ops[0] = op;
ops[1] = GGML_OP_MUL;
ops[2] = bias_op;
ops[3] = op;
ops[4] = GGML_OP_MUL;
ops[5] = bias_op;
ops[6] = GGML_OP_GLU;
} else {
ops[0] = op;
ops[1] = GGML_OP_MUL;
ops[2] = op;
ops[3] = GGML_OP_MUL;
ops[4] = GGML_OP_GLU;
}
const int n_ops = with_bias ? 7 : 5;
if (!ggml_can_fuse_subgraph(cgraph, i, n_ops, ops, out_nodes, 1) ||
!ggml_cuda_check_fusion_memory_ranges(cgraph, i, n_ops, out_nodes, 1)) {
continue;
}
ggml_tensor * gate_n = cgraph->nodes[gate_idx];
ggml_tensor * gate_scale_n = cgraph->nodes[gate_scale_idx];
ggml_tensor * gate_out_n = with_bias ? cgraph->nodes[gate_bias_idx] : gate_scale_n;
ggml_tensor * up_n = cgraph->nodes[up_idx];
ggml_tensor * up_scale_n = cgraph->nodes[up_scale_idx];
ggml_tensor * up_out_n = with_bias ? cgraph->nodes[up_bias_idx] : up_scale_n;
const ggml_tensor * glu = cgraph->nodes[glu_idx];
if (!ggml_cuda_should_fuse_mul_mat(up_n, gate_n, glu,
with_bias ? up_out_n : nullptr, with_bias ? gate_out_n : nullptr, up_scale_n, gate_scale_n)) {
continue;
}
const ggml_tensor * gate_scale = get_mul_mat_scale(gate_scale_n, gate_n);
const ggml_tensor * up_scale = get_mul_mat_scale(up_scale_n, up_n);
if (!gate_scale || !up_scale) {
continue;
}
const ggml_tensor * up_bias = with_bias ? get_bias_tensor(up_out_n, up_scale_n, bias_op) : nullptr;
const ggml_tensor * gate_bias = with_bias ? get_bias_tensor(gate_out_n, gate_scale_n, bias_op) : nullptr;
if (with_bias && (!ggml_are_same_shape(gate_out_n->src[0], gate_out_n->src[1]) ||
!ggml_are_same_shape(up_out_n->src[0], up_out_n->src[1]))) {
continue;
}
const ggml_tensor * src0 = up_n->src[0];
const ggml_tensor * src1 = up_n->src[1];
const ggml_tensor * ids = up_n->src[2];
ggml_cuda_mm_fusion_args_host fusion_data{};
fusion_data.gate = gate_n->src[0];
fusion_data.x_bias = up_bias;
fusion_data.gate_bias = gate_bias;
fusion_data.x_scale = up_scale;
fusion_data.gate_scale = gate_scale;
fusion_data.glu_op = ggml_get_glu_op(glu);
if (ggml_cuda_should_fuse_mul_mat_vec_q(up_n)) {
ggml_cuda_mul_mat_vec_q(*cuda_ctx, src0, src1, ids, cgraph->nodes[glu_idx], &fusion_data);
fused_mul_mat_vec = true;
fused_node_count = n_ops;
break;
}
}
if (fused_mul_mat_vec) {
break;
}
} else {
for (const bool with_bias : { false, true }) {
const int gate_idx = i;
const int gate_scale_idx = i + 4;
const int gate_bias_idx = with_bias ? i + 5 : -1;
const int up_idx = with_bias ? i + 6 : i + 5;
const int up_scale_idx = up_idx + 4;
const int up_bias_idx = with_bias ? up_idx + 5 : -1;
const int glu_idx = with_bias ? up_idx + 6 : up_idx + 5;
const int out_nodes[] = { glu_idx };
ggml_op ops[13];
if (with_bias) {
ops[0] = op;
ops[1] = GGML_OP_RESHAPE;
ops[2] = GGML_OP_REPEAT;
ops[3] = GGML_OP_GET_ROWS;
ops[4] = GGML_OP_MUL;
ops[5] = bias_op;
ops[6] = op;
ops[7] = GGML_OP_RESHAPE;
ops[8] = GGML_OP_REPEAT;
ops[9] = GGML_OP_GET_ROWS;
ops[10] = GGML_OP_MUL;
ops[11] = bias_op;
ops[12] = GGML_OP_GLU;
} else {
ops[0] = op;
ops[1] = GGML_OP_RESHAPE;
ops[2] = GGML_OP_REPEAT;
ops[3] = GGML_OP_GET_ROWS;
ops[4] = GGML_OP_MUL;
ops[5] = op;
ops[6] = GGML_OP_RESHAPE;
ops[7] = GGML_OP_REPEAT;
ops[8] = GGML_OP_GET_ROWS;
ops[9] = GGML_OP_MUL;
ops[10] = GGML_OP_GLU;
}
const int n_ops = with_bias ? 13 : 11;
if (!ggml_can_fuse_subgraph(cgraph, i, n_ops, ops, out_nodes, 1) ||
!ggml_cuda_check_fusion_memory_ranges(cgraph, i, n_ops, out_nodes, 1)) {
continue;
}
ggml_tensor * gate_n = cgraph->nodes[gate_idx];
ggml_tensor * gate_scale_n = cgraph->nodes[gate_scale_idx];
ggml_tensor * gate_out_n = with_bias ? cgraph->nodes[gate_bias_idx] : gate_scale_n;
ggml_tensor * up_n = cgraph->nodes[up_idx];
ggml_tensor * up_scale_n = cgraph->nodes[up_scale_idx];
ggml_tensor * up_out_n = with_bias ? cgraph->nodes[up_bias_idx] : up_scale_n;
const ggml_tensor * glu = cgraph->nodes[glu_idx];
if (!ggml_cuda_should_fuse_mul_mat(up_n, gate_n, glu,
with_bias ? up_out_n : nullptr, with_bias ? gate_out_n : nullptr, up_scale_n, gate_scale_n)) {
continue;
}
const ggml_tensor * gate_scale = get_mul_mat_id_scale(cgraph->nodes[gate_idx + 1], cgraph->nodes[gate_idx + 2],
cgraph->nodes[gate_idx + 3], gate_scale_n, gate_n);
const ggml_tensor * up_scale = get_mul_mat_id_scale(cgraph->nodes[up_idx + 1], cgraph->nodes[up_idx + 2],
cgraph->nodes[up_idx + 3], up_scale_n, up_n);
if (!gate_scale || !up_scale) {
continue;
}
const ggml_tensor * up_bias = with_bias ? get_bias_tensor(up_out_n, up_scale_n, bias_op) : nullptr;
const ggml_tensor * gate_bias = with_bias ? get_bias_tensor(gate_out_n, gate_scale_n, bias_op) : nullptr;
const ggml_tensor * src0 = up_n->src[0];
const ggml_tensor * src1 = up_n->src[1];
const ggml_tensor * ids = up_n->src[2];
ggml_cuda_mm_fusion_args_host fusion_data{};
fusion_data.gate = gate_n->src[0];
fusion_data.x_bias = up_bias;
fusion_data.gate_bias = gate_bias;
fusion_data.x_scale = up_scale;
fusion_data.gate_scale = gate_scale;
fusion_data.glu_op = ggml_get_glu_op(glu);
if (ggml_cuda_should_fuse_mul_mat_vec_q(up_n)) {
ggml_cuda_mul_mat_vec_q(*cuda_ctx, src0, src1, ids, cgraph->nodes[glu_idx], &fusion_data);
fused_mul_mat_vec = true;
fused_node_count = n_ops;
break;
}
}
if (fused_mul_mat_vec) {
break;
}
}
if (ggml_cuda_can_fuse(cgraph, i, { op, bias_op, op, bias_op, GGML_OP_GLU }, {})) {
ggml_tensor * glu = cgraph->nodes[i + 4];
ggml_tensor * gate_bias_n = glu->src[0];
@ -3235,23 +3482,8 @@ static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph
continue;
}
auto get_bias_tensor = [](const ggml_tensor * bias_node, const ggml_tensor * mul_node, ggml_op op_bias) {
if (op_bias == GGML_OP_ADD) {
if (bias_node->src[0] == mul_node) {
return bias_node->src[1];
}
if (bias_node->src[1] == mul_node) {
return bias_node->src[0];
}
return (ggml_tensor *) nullptr;
}
GGML_ASSERT(op_bias == GGML_OP_ADD_ID);
GGML_ASSERT(bias_node->src[0] == mul_node);
return bias_node->src[1];
};
ggml_tensor * up_bias_tensor = get_bias_tensor(up_bias_n, up_n, bias_op);
ggml_tensor * gate_bias_tensor = get_bias_tensor(gate_bias_n, gate_n, bias_op);
const ggml_tensor * up_bias_tensor = get_bias_tensor(up_bias_n, up_n, bias_op);
const ggml_tensor * gate_bias_tensor = get_bias_tensor(gate_bias_n, gate_n, bias_op);
if (!up_bias_tensor || !gate_bias_tensor) {
continue;
@ -3339,7 +3571,95 @@ static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph
fused_mul_mat_vec = false;
fused_node_count = 0;
// gate + add + glu + up + add
// mul_mat + scale + optional bias
for (ggml_op op : { GGML_OP_MUL_MAT, GGML_OP_MUL_MAT_ID }) {
const ggml_op bias_op = op == GGML_OP_MUL_MAT ? GGML_OP_ADD : GGML_OP_ADD_ID;
for (const bool with_bias : { false, true }) {
const int n_ops = op == GGML_OP_MUL_MAT ? (with_bias ? 3 : 2) : (with_bias ? 6 : 5);
const int out_nodes[] = { i + n_ops - 1 };
ggml_op ops[6];
if (op == GGML_OP_MUL_MAT) {
if (with_bias) {
ops[0] = op;
ops[1] = GGML_OP_MUL;
ops[2] = bias_op;
} else {
ops[0] = op;
ops[1] = GGML_OP_MUL;
}
} else {
if (with_bias) {
ops[0] = op;
ops[1] = GGML_OP_RESHAPE;
ops[2] = GGML_OP_REPEAT;
ops[3] = GGML_OP_GET_ROWS;
ops[4] = GGML_OP_MUL;
ops[5] = bias_op;
} else {
ops[0] = op;
ops[1] = GGML_OP_RESHAPE;
ops[2] = GGML_OP_REPEAT;
ops[3] = GGML_OP_GET_ROWS;
ops[4] = GGML_OP_MUL;
}
}
if (!ggml_can_fuse_subgraph(cgraph, i, n_ops, ops, out_nodes, 1) ||
!ggml_cuda_check_fusion_memory_ranges(cgraph, i, n_ops, out_nodes, 1)) {
continue;
}
ggml_tensor * mm_node = cgraph->nodes[i];
ggml_tensor * scale_node = op == GGML_OP_MUL_MAT ? cgraph->nodes[i + 1] : cgraph->nodes[i + 4];
ggml_tensor * out_node = with_bias ? cgraph->nodes[i + n_ops - 1] : scale_node;
const ggml_tensor * scale = nullptr;
if (op == GGML_OP_MUL_MAT) {
scale = get_mul_mat_scale(scale_node, mm_node);
} else {
scale = get_mul_mat_id_scale(cgraph->nodes[i + 1], cgraph->nodes[i + 2], cgraph->nodes[i + 3], scale_node, mm_node);
}
if (!scale) {
continue;
}
const ggml_tensor * bias = with_bias ? get_bias_tensor(out_node, scale_node, bias_op) : nullptr;
if (with_bias && !bias) {
continue;
}
if (with_bias && bias_op == GGML_OP_ADD && !ggml_are_same_shape(out_node->src[0], out_node->src[1])) {
continue;
}
if (with_bias && bias_op == GGML_OP_ADD_ID && out_node->src[2] != mm_node->src[2]) {
continue;
}
const ggml_tensor * src0 = mm_node->src[0];
const ggml_tensor * src1 = mm_node->src[1];
const ggml_tensor * ids = mm_node->src[2];
ggml_cuda_mm_fusion_args_host fusion_data{};
fusion_data.x_bias = bias;
fusion_data.x_scale = scale;
if (ggml_cuda_should_fuse_mul_mat_vec_q(mm_node)) {
ggml_cuda_mul_mat_vec_q(*cuda_ctx, src0, src1, ids, out_node, &fusion_data);
fused_mul_mat_vec = true;
fused_node_count = n_ops;
break;
}
}
if (fused_mul_mat_vec) {
break;
}
}
if (fused_mul_mat_vec) {
return fused_node_count - 1;
}
// mul_mat + add
for (ggml_op op : { GGML_OP_MUL_MAT, GGML_OP_MUL_MAT_ID }) {
const ggml_op bias_op = op == GGML_OP_MUL_MAT ? GGML_OP_ADD : GGML_OP_ADD_ID;
@ -3570,12 +3890,6 @@ static void ggml_cuda_graph_evaluate_and_capture(ggml_backend_cuda_context * cud
}
}
#ifdef GGML_CUDA_DEBUG
const int nodes_fused = i - prev_i - 1;
if (nodes_fused > 0) {
GGML_LOG_INFO("nodes_fused: %d\n", nodes_fused);
}
#endif
prev_i = i;
if (ggml_cuda_is_view_or_noop(node)) {
@ -3589,6 +3903,12 @@ static void ggml_cuda_graph_evaluate_and_capture(ggml_backend_cuda_context * cud
int nodes_to_skip = ggml_cuda_try_fuse(cuda_ctx, cgraph, i);
if (nodes_to_skip != 0) {
#ifdef GGML_CUDA_DEBUG
const int last_fused = i + nodes_to_skip;
GGML_LOG_INFO("nodes_fused: %d, first: %s (%s), last: %s (%s)\n",
nodes_to_skip + 1, ggml_op_name(node->op), node->name,
ggml_op_name(cgraph->nodes[last_fused]->op), cgraph->nodes[last_fused]->name);
#endif
i += nodes_to_skip;
continue;
}

75
ggml/src/ggml-cuda/mmvq.cu
View file

@ -521,9 +521,13 @@ static __global__ void mul_mat_vec_q(
bool use_gate = false;
bool use_bias = false;
bool use_gate_bias = false;
bool use_scale = false;
bool use_gate_scale = false;
[[maybe_unused]] const void * vgate = nullptr;
const float * x_bias = nullptr;
const float * gate_bias = nullptr;
const float * x_scale = nullptr;
const float * gate_scale = nullptr;
ggml_glu_op active_glu;
if constexpr (has_fusion) {
@ -534,34 +538,47 @@ static __global__ void mul_mat_vec_q(
x_bias = (const float *) fusion.x_bias;
gate_bias = (const float *) fusion.gate_bias;
active_glu = fusion.glu_op;
if constexpr (type == GGML_TYPE_NVFP4) {
use_scale = fusion.x_scale != nullptr;
use_gate_scale = fusion.gate_scale != nullptr && use_gate;
x_scale = (const float *) fusion.x_scale;
gate_scale = (const float *) fusion.gate_scale;
}
}
[[maybe_unused]] float x_biases[ncols_dst] = { 0.0f };
[[maybe_unused]] float gate_biases[ncols_dst] = { 0.0f };
[[maybe_unused]] float x_scales;
[[maybe_unused]] float gate_scales;
if constexpr (has_fusion) {
// 1. Hide latency by prefetching bias, gates and scales here
// 2. load only on threads that won't die after partial sum calculation
const uint32_t channel_bias = ids ? channel_x : channel_dst;
if (use_bias) {
x_bias = x_bias + sample_dst*stride_sample_dst + channel_bias*stride_channel_dst + row0;
// 1. Hide latency by prefetching bias and gate here
// 2. load only on threads that won't die after partial sum calculation
if (threadIdx.x < rows_per_cuda_block && threadIdx.y == 0 &&
(rows_per_cuda_block == 1 || uint32_t(row0 + threadIdx.x) < stride_col_dst)) {
if (threadIdx.x < rows_per_cuda_block && threadIdx.y == 0 &&
(rows_per_cuda_block == 1 || uint32_t(row0 + threadIdx.x) < stride_col_dst)) {
if (use_bias) {
x_bias = x_bias + sample_dst * stride_sample_dst + channel_bias * stride_channel_dst + row0;
#pragma unroll
for (int j = 0; j < ncols_dst; ++j) {
x_biases[j] = x_bias[j * stride_col_dst + threadIdx.x];
}
}
}
if (use_gate_bias) {
gate_bias = gate_bias + sample_dst*stride_sample_dst + channel_bias*stride_channel_dst + row0;
if (threadIdx.x < rows_per_cuda_block && threadIdx.y == 0 &&
(rows_per_cuda_block == 1 || uint32_t(row0 + threadIdx.x) < stride_col_dst)) {
if (use_gate_bias) {
gate_bias = gate_bias + sample_dst * stride_sample_dst + channel_bias * stride_channel_dst + row0;
#pragma unroll
for (int j = 0; j < ncols_dst; ++j) {
gate_biases[j] = gate_bias[j * stride_col_dst + threadIdx.x];
}
}
if constexpr (type == GGML_TYPE_NVFP4) {
if (use_scale) {
x_scales = x_scale[ids ? channel_x : 0];
}
if (use_gate_scale) {
gate_scales = gate_scale[ids ? channel_x : 0];
}
}
}
}
@ -643,11 +660,21 @@ static __global__ void mul_mat_vec_q(
if (threadIdx.x < rows_per_cuda_block && (rows_per_cuda_block == 1 || uint32_t(row0 + threadIdx.x) < stride_col_dst)) {
float result = tmp[j][threadIdx.x];
if constexpr (has_fusion) {
if constexpr (type == GGML_TYPE_NVFP4) {
if (use_scale) {
result *= x_scales;
}
}
if (use_bias) {
result += x_biases[j];
}
if (use_gate) {
float gate_value = tmp_gate[j][threadIdx.x];
if constexpr (type == GGML_TYPE_NVFP4) {
if (use_gate_scale) {
gate_value *= gate_scales;
}
}
if (use_gate_bias) {
gate_value += gate_biases[j];
}
@ -673,7 +700,10 @@ static __global__ void mul_mat_vec_q(
}
if constexpr (!has_fusion) {
GGML_UNUSED_VARS(use_gate, use_bias, use_gate_bias, active_glu, gate_bias, x_bias, tmp_gate);
GGML_UNUSED_VARS(use_gate, use_bias, use_gate_bias, use_scale, use_gate_scale, active_glu, gate_bias, x_bias, x_scale, gate_scale, tmp_gate);
}
if constexpr (type != GGML_TYPE_NVFP4) {
GGML_UNUSED_VARS(use_scale, use_gate_scale, x_scale, gate_scale, x_scales, gate_scales);
}
}
@ -769,7 +799,8 @@ static void mul_mat_vec_q_switch_fusion(
const dim3 & block_nums, const dim3 & block_dims, const int nbytes_shared,
const uint32_t ids_stride, cudaStream_t stream) {
const bool has_fusion = fusion.gate != nullptr || fusion.x_bias != nullptr || fusion.gate_bias != nullptr;
const bool has_fusion = fusion.gate != nullptr || fusion.x_bias != nullptr || fusion.gate_bias != nullptr ||
fusion.x_scale != nullptr || fusion.gate_scale != nullptr;
if constexpr (c_ncols_dst == 1) {
if (has_fusion) {
const ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params(block_nums, block_dims, nbytes_shared, stream);
@ -834,7 +865,6 @@ static void mul_mat_vec_q_switch_ncols_dst(
const int warp_size = ggml_cuda_info().devices[device].warp_size;
const mmvq_parameter_table_id table_id = get_device_table_id(cc);
const bool has_fusion = fusion.gate != nullptr || fusion.x_bias != nullptr || fusion.gate_bias != nullptr;
const bool has_ids = ids != nullptr;
const auto should_use_small_k = [&](int c_ncols_dst) {
@ -973,8 +1003,6 @@ static void mul_mat_vec_q_switch_ncols_dst(
GGML_ABORT("fatal error");
break;
}
GGML_UNUSED(has_fusion);
}
static void mul_mat_vec_q_switch_type(
const void * vx, const ggml_type type_x, const void * vy, const int32_t * ids, const ggml_cuda_mm_fusion_args_device fusion, float * dst,
@ -1154,6 +1182,9 @@ void ggml_cuda_mul_mat_vec_q(
if (fusion) {
GGML_ASSERT( !ids || dst->ne[2] == 1);
GGML_ASSERT( ids || dst->ne[1] == 1);
// Scale fusion is only allowed for NVFP4 currently as the cost of checking this at run-time in the prologue is
// non-negligible for some models such as gpt-oss-20b
GGML_ASSERT((fusion->x_scale == nullptr && fusion->gate_scale == nullptr) || src0->type == GGML_TYPE_NVFP4);
if (fusion->x_bias) {
GGML_ASSERT(fusion->x_bias->type == GGML_TYPE_F32);
@ -1171,6 +1202,18 @@ void ggml_cuda_mul_mat_vec_q(
GGML_ASSERT(!ids || fusion->gate_bias->ne[1] == src0->ne[2]);
fusion_local.gate_bias = fusion->gate_bias->data;
}
if (fusion->x_scale) {
GGML_ASSERT(fusion->x_scale->type == GGML_TYPE_F32);
GGML_ASSERT(ggml_is_contiguous(fusion->x_scale));
GGML_ASSERT(ggml_nelements(fusion->x_scale) == (ids ? src0->ne[2] : 1));
fusion_local.x_scale = fusion->x_scale->data;
}
if (fusion->gate_scale) {
GGML_ASSERT(fusion->gate_scale->type == GGML_TYPE_F32);
GGML_ASSERT(ggml_is_contiguous(fusion->gate_scale));
GGML_ASSERT(ggml_nelements(fusion->gate_scale) == (ids ? src0->ne[2] : 1));
fusion_local.gate_scale = fusion->gate_scale->data;
}
fusion_local.glu_op = fusion->glu_op;
}

8
ggml/src/ggml-metal/ggml-metal-device.cpp
View file

@ -160,11 +160,15 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_get_rows(ggml_me
return res;
}
ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_set_rows(ggml_metal_library_t lib, ggml_type tidx, ggml_type tdst) {
ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_set_rows(ggml_metal_library_t lib, const ggml_tensor * op) {
char base[256];
char name[256];
snprintf(base, 256, "kernel_set_rows_%s_%s", ggml_type_name(tdst), ggml_type_name(tidx));
const auto tsrc = op->src[0]->type;
const auto tidx = op->src[1]->type;
const auto tdst = op->type;
snprintf(base, 256, "kernel_set_rows_%s_%s_%s", ggml_type_name(tsrc), ggml_type_name(tidx), ggml_type_name(tdst));
snprintf(name, 256, "%s", base);
ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name);

2
ggml/src/ggml-metal/ggml-metal-device.h
View file

@ -112,7 +112,7 @@ struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_cpy
struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_pool_1d (ggml_metal_library_t lib, const struct ggml_tensor * op, enum ggml_op_pool op_pool);
struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_pool_2d (ggml_metal_library_t lib, const struct ggml_tensor * op, enum ggml_op_pool op_pool);
struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_get_rows (ggml_metal_library_t lib, enum ggml_type tsrc);
struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_set_rows (ggml_metal_library_t lib, enum ggml_type tidx, enum ggml_type tdst);
struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_set_rows (ggml_metal_library_t lib, const struct ggml_tensor * op);
struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_diag (ggml_metal_library_t lib, const struct ggml_tensor * op);
struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_repeat (ggml_metal_library_t lib, enum ggml_type tsrc);
struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_concat (ggml_metal_library_t lib, enum ggml_type tsrc);

2
ggml/src/ggml-metal/ggml-metal-device.m
View file

@ -1340,7 +1340,7 @@ bool ggml_metal_device_supports_op(ggml_metal_device_t dev, const struct ggml_te
return op->src[0]->type != GGML_TYPE_NVFP4;
case GGML_OP_SET_ROWS:
{
if (op->src[0]->type != GGML_TYPE_F32) {
if (op->src[0]->type != GGML_TYPE_F32 && op->src[0]->type != GGML_TYPE_F16) {
return false;
}

2
ggml/src/ggml-metal/ggml-metal-ops.cpp
View file

@ -1202,7 +1202,7 @@ int ggml_metal_op_set_rows(ggml_metal_op_t ctx, int idx) {
GGML_TENSOR_LOCALS( int32_t, ne, op, ne);
GGML_TENSOR_LOCALS(uint64_t, nb, op, nb);
auto pipeline = ggml_metal_library_get_pipeline_set_rows(lib, op->src[1]->type, op->type);
auto pipeline = ggml_metal_library_get_pipeline_set_rows(lib, op);
const int32_t nk0 = ne0/ggml_blck_size(op->type);

149
ggml/src/ggml-metal/ggml-metal.metal
View file

@ -42,6 +42,8 @@ typedef matrix<bfloat, 4, 4> bfloat4x4;
typedef matrix<bfloat, 2, 4> bfloat2x4;
#endif
#define QK_NL 16
constexpr constant static float kvalues_iq4nl_f[16] = {
-127.f, -104.f, -83.f, -65.f, -49.f, -35.f, -22.f, -10.f, 1.f, 13.f, 25.f, 38.f, 53.f, 69.f, 89.f, 113.f
};
@ -9386,7 +9388,40 @@ kernel void kernel_get_rows_f(
}
}
template<typename TI, typename block_q, void (*quantize_func)(device const float *, device block_q &)>
typedef decltype(kernel_get_rows_f<float, float>) get_rows_f_t;
template [[host_name("kernel_get_rows_f32")]] kernel get_rows_f_t kernel_get_rows_f<float, float>;
template [[host_name("kernel_get_rows_f16")]] kernel get_rows_f_t kernel_get_rows_f<half, float>;
template [[host_name("kernel_get_rows_i32")]] kernel get_rows_f_t kernel_get_rows_f<int32_t, int32_t>;
#if defined(GGML_METAL_HAS_BF16)
template [[host_name("kernel_get_rows_bf16")]] kernel get_rows_f_t kernel_get_rows_f<bfloat, float>;
#endif
typedef decltype(kernel_get_rows_q<block_q4_0, 2, dequantize_q4_0>) get_rows_q_t;
template [[host_name("kernel_get_rows_q1_0")]] kernel get_rows_q_t kernel_get_rows_q<block_q1_0, 8, dequantize_q1_0>;
template [[host_name("kernel_get_rows_q4_0")]] kernel get_rows_q_t kernel_get_rows_q<block_q4_0, 2, dequantize_q4_0>;
template [[host_name("kernel_get_rows_q4_1")]] kernel get_rows_q_t kernel_get_rows_q<block_q4_1, 2, dequantize_q4_1>;
template [[host_name("kernel_get_rows_q5_0")]] kernel get_rows_q_t kernel_get_rows_q<block_q5_0, 2, dequantize_q5_0>;
template [[host_name("kernel_get_rows_q5_1")]] kernel get_rows_q_t kernel_get_rows_q<block_q5_1, 2, dequantize_q5_1>;
template [[host_name("kernel_get_rows_q8_0")]] kernel get_rows_q_t kernel_get_rows_q<block_q8_0, 2, dequantize_q8_0>;
template [[host_name("kernel_get_rows_mxfp4")]] kernel get_rows_q_t kernel_get_rows_q<block_mxfp4, 2, dequantize_mxfp4>;
template [[host_name("kernel_get_rows_q2_K")]] kernel get_rows_q_t kernel_get_rows_q<block_q2_K, QK_NL, dequantize_q2_K>;
template [[host_name("kernel_get_rows_q3_K")]] kernel get_rows_q_t kernel_get_rows_q<block_q3_K, QK_NL, dequantize_q3_K>;
template [[host_name("kernel_get_rows_q4_K")]] kernel get_rows_q_t kernel_get_rows_q<block_q4_K, QK_NL, dequantize_q4_K>;
template [[host_name("kernel_get_rows_q5_K")]] kernel get_rows_q_t kernel_get_rows_q<block_q5_K, QK_NL, dequantize_q5_K>;
template [[host_name("kernel_get_rows_q6_K")]] kernel get_rows_q_t kernel_get_rows_q<block_q6_K, QK_NL, dequantize_q6_K>;
template [[host_name("kernel_get_rows_iq2_xxs")]] kernel get_rows_q_t kernel_get_rows_q<block_iq2_xxs, QK_NL, dequantize_iq2_xxs>;
template [[host_name("kernel_get_rows_iq2_xs")]] kernel get_rows_q_t kernel_get_rows_q<block_iq2_xs, QK_NL, dequantize_iq2_xs>;
template [[host_name("kernel_get_rows_iq3_xxs")]] kernel get_rows_q_t kernel_get_rows_q<block_iq3_xxs, QK_NL, dequantize_iq3_xxs>;
template [[host_name("kernel_get_rows_iq3_s")]] kernel get_rows_q_t kernel_get_rows_q<block_iq3_s, QK_NL, dequantize_iq3_s>;
template [[host_name("kernel_get_rows_iq2_s")]] kernel get_rows_q_t kernel_get_rows_q<block_iq2_s, QK_NL, dequantize_iq2_s>;
template [[host_name("kernel_get_rows_iq1_s")]] kernel get_rows_q_t kernel_get_rows_q<block_iq1_s, QK_NL, dequantize_iq1_s>;
template [[host_name("kernel_get_rows_iq1_m")]] kernel get_rows_q_t kernel_get_rows_q<block_iq1_m, QK_NL, dequantize_iq1_m>;
template [[host_name("kernel_get_rows_iq4_nl")]] kernel get_rows_q_t kernel_get_rows_q<block_iq4_nl, 2, dequantize_iq4_nl>;
template [[host_name("kernel_get_rows_iq4_xs")]] kernel get_rows_q_t kernel_get_rows_q<block_iq4_xs, QK_NL, dequantize_iq4_xs>;
template<typename TS, typename TI, typename block_q, void (*quantize_func)(device const float *, device block_q &)>
kernel void kernel_set_rows_q32(
constant ggml_metal_kargs_set_rows & args,
device const void * src0,
@ -9410,14 +9445,14 @@ kernel void kernel_set_rows_q32(
const TI i1 = ((const device TI *) ((const device char *) src1 + i10*args.nb10 + i11*args.nb11 + i12*args.nb12))[0];
device block_q * dst_row = ( device block_q *) (( device char *) dst + i1*args.nb1 + i02*args.nb2 + i03*args.nb3);
const device float * src_row = (const device float *) ((const device char *) src0 + i01*args.nb01 + i02*args.nb02 + i03*args.nb03);
const device TS * src_row = (const device TS *) ((const device char *) src0 + i01*args.nb01 + i02*args.nb02 + i03*args.nb03);
for (int ind = tiitg%tptg.x; ind < args.nk0; ind += tptg.x) {
quantize_func(src_row + 32*ind, dst_row[ind]);
}
}
template<typename T, typename TI>
template<typename TS, typename TI, typename TD>
kernel void kernel_set_rows_f(
constant ggml_metal_kargs_set_rows & args,
device const void * src0,
@ -9440,14 +9475,47 @@ kernel void kernel_set_rows_f(
const int32_t i10 = i01;
const TI i1 = ((const device TI *) ((const device char *) src1 + i10*args.nb10 + i11*args.nb11 + i12*args.nb12))[0];
device T * dst_row = ( device T *) (( device char *) dst + i1*args.nb1 + i02*args.nb2 + i03*args.nb3);
const device float * src_row = (const device float *) ((const device char *) src0 + i01*args.nb01 + i02*args.nb02 + i03*args.nb03);
device TD * dst_row = ( device TD *) (( device char *) dst + i1*args.nb1 + i02*args.nb2 + i03*args.nb3);
const device TS * src_row = (const device TS *) ((const device char *) src0 + i01*args.nb01 + i02*args.nb02 + i03*args.nb03);
for (int ind = tiitg%tptg.x; ind < args.nk0; ind += tptg.x) {
dst_row[ind] = (T) src_row[ind];
dst_row[ind] = (TD) src_row[ind];
}
}
typedef decltype(kernel_set_rows_f<float, int64_t, float>) set_rows_f_t;
template [[host_name("kernel_set_rows_f32_i64_f32")]] kernel set_rows_f_t kernel_set_rows_f<float, int64_t, float>;
template [[host_name("kernel_set_rows_f32_i32_f32")]] kernel set_rows_f_t kernel_set_rows_f<float, int32_t, float>;
template [[host_name("kernel_set_rows_f32_i64_f16")]] kernel set_rows_f_t kernel_set_rows_f<float, int64_t, half>;
template [[host_name("kernel_set_rows_f32_i32_f16")]] kernel set_rows_f_t kernel_set_rows_f<float, int32_t, half>;
#if defined(GGML_METAL_HAS_BF16)
template [[host_name("kernel_set_rows_f32_i64_bf16")]] kernel set_rows_f_t kernel_set_rows_f<float, int64_t, bfloat>;
template [[host_name("kernel_set_rows_f32_i32_bf16")]] kernel set_rows_f_t kernel_set_rows_f<float, int32_t, bfloat>;
#endif
template [[host_name("kernel_set_rows_f16_i64_f16")]] kernel set_rows_f_t kernel_set_rows_f<half, int64_t, half>;
template [[host_name("kernel_set_rows_f16_i32_f16")]] kernel set_rows_f_t kernel_set_rows_f<half, int32_t, half>;
#if defined(GGML_METAL_HAS_BF16)
template [[host_name("kernel_set_rows_bf16_i64_bf16")]] kernel set_rows_f_t kernel_set_rows_f<bfloat, int64_t, bfloat>;
template [[host_name("kernel_set_rows_bf16_i32_bf16")]] kernel set_rows_f_t kernel_set_rows_f<bfloat, int32_t, bfloat>;
#endif
typedef decltype(kernel_set_rows_q32<float, int64_t, block_q8_0, quantize_q8_0>) set_rows_q32_t;
template [[host_name("kernel_set_rows_f32_i64_q8_0")]] kernel set_rows_q32_t kernel_set_rows_q32<float, int64_t, block_q8_0, quantize_q8_0>;
template [[host_name("kernel_set_rows_f32_i32_q8_0")]] kernel set_rows_q32_t kernel_set_rows_q32<float, int32_t, block_q8_0, quantize_q8_0>;
template [[host_name("kernel_set_rows_f32_i64_q4_0")]] kernel set_rows_q32_t kernel_set_rows_q32<float, int64_t, block_q4_0, quantize_q4_0>;
template [[host_name("kernel_set_rows_f32_i32_q4_0")]] kernel set_rows_q32_t kernel_set_rows_q32<float, int32_t, block_q4_0, quantize_q4_0>;
template [[host_name("kernel_set_rows_f32_i64_q4_1")]] kernel set_rows_q32_t kernel_set_rows_q32<float, int64_t, block_q4_1, quantize_q4_1>;
template [[host_name("kernel_set_rows_f32_i32_q4_1")]] kernel set_rows_q32_t kernel_set_rows_q32<float, int32_t, block_q4_1, quantize_q4_1>;
template [[host_name("kernel_set_rows_f32_i64_q5_0")]] kernel set_rows_q32_t kernel_set_rows_q32<float, int64_t, block_q5_0, quantize_q5_0>;
template [[host_name("kernel_set_rows_f32_i32_q5_0")]] kernel set_rows_q32_t kernel_set_rows_q32<float, int32_t, block_q5_0, quantize_q5_0>;
template [[host_name("kernel_set_rows_f32_i64_q5_1")]] kernel set_rows_q32_t kernel_set_rows_q32<float, int64_t, block_q5_1, quantize_q5_1>;
template [[host_name("kernel_set_rows_f32_i32_q5_1")]] kernel set_rows_q32_t kernel_set_rows_q32<float, int32_t, block_q5_1, quantize_q5_1>;
template [[host_name("kernel_set_rows_f32_i64_iq4_nl")]] kernel set_rows_q32_t kernel_set_rows_q32<float, int64_t, block_iq4_nl, quantize_iq4_nl>;
template [[host_name("kernel_set_rows_f32_i32_iq4_nl")]] kernel set_rows_q32_t kernel_set_rows_q32<float, int32_t, block_iq4_nl, quantize_iq4_nl>;
kernel void kernel_diag_f32(
constant ggml_metal_kargs_diag & args,
device const char * src0,
@ -10190,75 +10258,6 @@ kernel void kernel_mul_mm_id(
}
}
#define QK_NL 16
//
// get rows
//
typedef decltype(kernel_get_rows_f<float, float>) get_rows_f_t;
template [[host_name("kernel_get_rows_f32")]] kernel get_rows_f_t kernel_get_rows_f<float, float>;
template [[host_name("kernel_get_rows_f16")]] kernel get_rows_f_t kernel_get_rows_f<half, float>;
template [[host_name("kernel_get_rows_i32")]] kernel get_rows_f_t kernel_get_rows_f<int32_t, int32_t>;
#if defined(GGML_METAL_HAS_BF16)
template [[host_name("kernel_get_rows_bf16")]] kernel get_rows_f_t kernel_get_rows_f<bfloat, float>;
#endif
typedef decltype(kernel_get_rows_q<block_q4_0, 2, dequantize_q4_0>) get_rows_q_t;
template [[host_name("kernel_get_rows_q1_0")]] kernel get_rows_q_t kernel_get_rows_q<block_q1_0, 8, dequantize_q1_0>;
template [[host_name("kernel_get_rows_q4_0")]] kernel get_rows_q_t kernel_get_rows_q<block_q4_0, 2, dequantize_q4_0>;
template [[host_name("kernel_get_rows_q4_1")]] kernel get_rows_q_t kernel_get_rows_q<block_q4_1, 2, dequantize_q4_1>;
template [[host_name("kernel_get_rows_q5_0")]] kernel get_rows_q_t kernel_get_rows_q<block_q5_0, 2, dequantize_q5_0>;
template [[host_name("kernel_get_rows_q5_1")]] kernel get_rows_q_t kernel_get_rows_q<block_q5_1, 2, dequantize_q5_1>;
template [[host_name("kernel_get_rows_q8_0")]] kernel get_rows_q_t kernel_get_rows_q<block_q8_0, 2, dequantize_q8_0>;
template [[host_name("kernel_get_rows_mxfp4")]] kernel get_rows_q_t kernel_get_rows_q<block_mxfp4, 2, dequantize_mxfp4>;
template [[host_name("kernel_get_rows_q2_K")]] kernel get_rows_q_t kernel_get_rows_q<block_q2_K, QK_NL, dequantize_q2_K>;
template [[host_name("kernel_get_rows_q3_K")]] kernel get_rows_q_t kernel_get_rows_q<block_q3_K, QK_NL, dequantize_q3_K>;
template [[host_name("kernel_get_rows_q4_K")]] kernel get_rows_q_t kernel_get_rows_q<block_q4_K, QK_NL, dequantize_q4_K>;
template [[host_name("kernel_get_rows_q5_K")]] kernel get_rows_q_t kernel_get_rows_q<block_q5_K, QK_NL, dequantize_q5_K>;
template [[host_name("kernel_get_rows_q6_K")]] kernel get_rows_q_t kernel_get_rows_q<block_q6_K, QK_NL, dequantize_q6_K>;
template [[host_name("kernel_get_rows_iq2_xxs")]] kernel get_rows_q_t kernel_get_rows_q<block_iq2_xxs, QK_NL, dequantize_iq2_xxs>;
template [[host_name("kernel_get_rows_iq2_xs")]] kernel get_rows_q_t kernel_get_rows_q<block_iq2_xs, QK_NL, dequantize_iq2_xs>;
template [[host_name("kernel_get_rows_iq3_xxs")]] kernel get_rows_q_t kernel_get_rows_q<block_iq3_xxs, QK_NL, dequantize_iq3_xxs>;
template [[host_name("kernel_get_rows_iq3_s")]] kernel get_rows_q_t kernel_get_rows_q<block_iq3_s, QK_NL, dequantize_iq3_s>;
template [[host_name("kernel_get_rows_iq2_s")]] kernel get_rows_q_t kernel_get_rows_q<block_iq2_s, QK_NL, dequantize_iq2_s>;
template [[host_name("kernel_get_rows_iq1_s")]] kernel get_rows_q_t kernel_get_rows_q<block_iq1_s, QK_NL, dequantize_iq1_s>;
template [[host_name("kernel_get_rows_iq1_m")]] kernel get_rows_q_t kernel_get_rows_q<block_iq1_m, QK_NL, dequantize_iq1_m>;
template [[host_name("kernel_get_rows_iq4_nl")]] kernel get_rows_q_t kernel_get_rows_q<block_iq4_nl, 2, dequantize_iq4_nl>;
template [[host_name("kernel_get_rows_iq4_xs")]] kernel get_rows_q_t kernel_get_rows_q<block_iq4_xs, QK_NL, dequantize_iq4_xs>;
//
// set rows
//
typedef decltype(kernel_set_rows_f<float, int64_t>) set_rows_f_t;
template [[host_name("kernel_set_rows_f32_i64")]] kernel set_rows_f_t kernel_set_rows_f<float, int64_t>;
template [[host_name("kernel_set_rows_f32_i32")]] kernel set_rows_f_t kernel_set_rows_f<float, int32_t>;
template [[host_name("kernel_set_rows_f16_i64")]] kernel set_rows_f_t kernel_set_rows_f<half, int64_t>;
template [[host_name("kernel_set_rows_f16_i32")]] kernel set_rows_f_t kernel_set_rows_f<half, int32_t>;
#if defined(GGML_METAL_HAS_BF16)
template [[host_name("kernel_set_rows_bf16_i64")]] kernel set_rows_f_t kernel_set_rows_f<bfloat, int64_t>;
template [[host_name("kernel_set_rows_bf16_i32")]] kernel set_rows_f_t kernel_set_rows_f<bfloat, int32_t>;
#endif
typedef decltype(kernel_set_rows_q32<int64_t, block_q8_0, quantize_q8_0>) set_rows_q32_t;
template [[host_name("kernel_set_rows_q8_0_i64")]] kernel set_rows_q32_t kernel_set_rows_q32<int64_t, block_q8_0, quantize_q8_0>;
template [[host_name("kernel_set_rows_q8_0_i32")]] kernel set_rows_q32_t kernel_set_rows_q32<int32_t, block_q8_0, quantize_q8_0>;
template [[host_name("kernel_set_rows_q4_0_i64")]] kernel set_rows_q32_t kernel_set_rows_q32<int64_t, block_q4_0, quantize_q4_0>;
template [[host_name("kernel_set_rows_q4_0_i32")]] kernel set_rows_q32_t kernel_set_rows_q32<int32_t, block_q4_0, quantize_q4_0>;
template [[host_name("kernel_set_rows_q4_1_i64")]] kernel set_rows_q32_t kernel_set_rows_q32<int64_t, block_q4_1, quantize_q4_1>;
template [[host_name("kernel_set_rows_q4_1_i32")]] kernel set_rows_q32_t kernel_set_rows_q32<int32_t, block_q4_1, quantize_q4_1>;
template [[host_name("kernel_set_rows_q5_0_i64")]] kernel set_rows_q32_t kernel_set_rows_q32<int64_t, block_q5_0, quantize_q5_0>;
template [[host_name("kernel_set_rows_q5_0_i32")]] kernel set_rows_q32_t kernel_set_rows_q32<int32_t, block_q5_0, quantize_q5_0>;
template [[host_name("kernel_set_rows_q5_1_i64")]] kernel set_rows_q32_t kernel_set_rows_q32<int64_t, block_q5_1, quantize_q5_1>;
template [[host_name("kernel_set_rows_q5_1_i32")]] kernel set_rows_q32_t kernel_set_rows_q32<int32_t, block_q5_1, quantize_q5_1>;
template [[host_name("kernel_set_rows_iq4_nl_i64")]] kernel set_rows_q32_t kernel_set_rows_q32<int64_t, block_iq4_nl, quantize_iq4_nl>;
template [[host_name("kernel_set_rows_iq4_nl_i32")]] kernel set_rows_q32_t kernel_set_rows_q32<int32_t, block_iq4_nl, quantize_iq4_nl>;
//
// matrix-matrix multiplication
//

76
ggml/src/ggml-quants.c
View file

@ -71,6 +71,44 @@ void quantize_row_q1_0_ref(const float * GGML_RESTRICT x, block_q1_0 * GGML_REST
}
}
void quantize_row_q2_0_ref(const float * GGML_RESTRICT x, block_q2_0 * GGML_RESTRICT y, int64_t k) {
static const int qk = QK2_0;
assert(k % qk == 0);
const int nb = k / qk;
for (int i = 0; i < nb; i++) {
// Compute scale as max absolute value in the block
float amax = 0.0f;
for (int j = 0; j < qk; j++) {
const float a = fabsf(x[i*qk + j]);
if (a > amax) amax = a;
}
const float d = amax;
const float id = d > 0.0f ? 1.0f / d : 0.0f;
y[i].d = GGML_FP32_TO_FP16(d);
// Clear quant bytes
for (int j = 0; j < qk / 4; ++j) {
y[i].qs[j] = 0;
}
// Encode 2-bit values: round(w/d) clamped to [-1, 2], then add 1
// 00 (-1) = -scale, 01 (0) = 0, 10 (+1) = +scale, 11 (+2) = 2*scale
for (int j = 0; j < qk; ++j) {
const float w = x[i*qk + j];
int q = (int)roundf(w * id) + 1;
if (q < 0) q = 0;
if (q > 3) q = 3;
const int byte_index = j / 4;
const int bit_offset = (j % 4) * 2;
y[i].qs[byte_index] |= ((uint8_t)q << bit_offset);
}
}
}
// reference implementation for deterministic creation of model files
void quantize_row_q4_0_ref(const float * GGML_RESTRICT x, block_q4_0 * GGML_RESTRICT y, int64_t k) {
static const int qk = QK4_0;
@ -398,6 +436,26 @@ void dequantize_row_q1_0(const block_q1_0 * GGML_RESTRICT x, float * GGML_RESTRI
}
}
void dequantize_row_q2_0(const block_q2_0 * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k) {
static const int qk = QK2_0;
assert(k % qk == 0);
const int nb = k / qk;
for (int i = 0; i < nb; i++) {
const float d = GGML_FP16_TO_FP32(x[i].d);
for (int j = 0; j < qk; ++j) {
const int byte_index = j / 4;
const int bit_offset = (j % 4) * 2;
const uint8_t q = (x[i].qs[byte_index] >> bit_offset) & 0x03;
// 00=-1, 01=0, 10=+1, 11=+2
y[i*qk + j] = ((int)q - 1) * d;
}
}
}
void dequantize_row_q4_0(const block_q4_0 * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k) {
static const int qk = QK4_0;
@ -2052,6 +2110,20 @@ size_t quantize_q1_0(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst,
return nrow * row_size;
}
size_t quantize_q2_0(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrow, int64_t n_per_row, const float * quant_weights) {
if (!quant_weights) {
quantize_row_q2_0_ref(src, dst, (int64_t)nrow*n_per_row);
return nrow * ggml_row_size(GGML_TYPE_Q2_0, n_per_row);
}
size_t row_size = ggml_row_size(GGML_TYPE_Q2_0, n_per_row);
char * qrow = (char *)dst;
for (int64_t row = 0; row < nrow; ++row) {
quantize_row_q2_0_ref(src, (block_q2_0*)qrow, n_per_row);
src += n_per_row;
qrow += row_size;
}
return nrow * row_size;
}
size_t quantize_q4_0(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrow, int64_t n_per_row, const float * quant_weights) {
if (!quant_weights) {
@ -5461,6 +5533,10 @@ bool ggml_validate_row_data(enum ggml_type type, const void * data, size_t nbyte
{
VALIDATE_ROW_DATA_D_F16_IMPL(block_q1_0, data, nb);
} break;
case GGML_TYPE_Q2_0:
{
VALIDATE_ROW_DATA_D_F16_IMPL(block_q2_0, data, nb);
} break;
case GGML_TYPE_Q4_0:
{
VALIDATE_ROW_DATA_D_F16_IMPL(block_q4_0, data, nb);

3
ggml/src/ggml-quants.h
View file

@ -15,6 +15,7 @@ extern "C" {
// Quantization
GGML_API void quantize_row_q1_0_ref(const float * GGML_RESTRICT x, block_q1_0 * GGML_RESTRICT y, int64_t k);
GGML_API void quantize_row_q2_0_ref(const float * GGML_RESTRICT x, block_q2_0 * GGML_RESTRICT y, int64_t k);
GGML_API void quantize_row_q4_0_ref(const float * GGML_RESTRICT x, block_q4_0 * GGML_RESTRICT y, int64_t k);
GGML_API void quantize_row_q4_1_ref(const float * GGML_RESTRICT x, block_q4_1 * GGML_RESTRICT y, int64_t k);
GGML_API void quantize_row_q5_0_ref(const float * GGML_RESTRICT x, block_q5_0 * GGML_RESTRICT y, int64_t k);
@ -43,6 +44,7 @@ GGML_API void quantize_row_iq2_s_ref (const float * GGML_RESTRICT x, block_iq2_
// Dequantization
GGML_API void dequantize_row_q1_0(const block_q1_0 * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k);
GGML_API void dequantize_row_q2_0(const block_q2_0 * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k);
GGML_API void dequantize_row_q4_0(const block_q4_0 * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k);
GGML_API void dequantize_row_q4_1(const block_q4_1 * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k);
GGML_API void dequantize_row_q5_0(const block_q5_0 * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k);
@ -93,6 +95,7 @@ GGML_API size_t quantize_q4_K(const float * GGML_RESTRICT src, void * GGML_RESTR
GGML_API size_t quantize_q5_K(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix);
GGML_API size_t quantize_q6_K(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix);
GGML_API size_t quantize_q1_0(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix);
GGML_API size_t quantize_q2_0(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix);
GGML_API size_t quantize_q4_0(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix);
GGML_API size_t quantize_q4_1(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix);
GGML_API size_t quantize_q5_0(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix);

21
ggml/src/ggml.c
View file

@ -685,6 +685,14 @@ static const struct ggml_type_traits type_traits[GGML_TYPE_COUNT] = {
.to_float = (ggml_to_float_t) dequantize_row_q1_0,
.from_float_ref = (ggml_from_float_t) quantize_row_q1_0_ref,
},
[GGML_TYPE_Q2_0] = {
.type_name = "q2_0",
.blck_size = QK2_0,
.type_size = sizeof(block_q2_0),
.is_quantized = true,
.to_float = (ggml_to_float_t) dequantize_row_q2_0,
.from_float_ref = (ggml_from_float_t) quantize_row_q2_0_ref,
},
[GGML_TYPE_Q4_0] = {
.type_name = "q4_0",
.blck_size = QK4_0,
@ -1433,6 +1441,7 @@ enum ggml_type ggml_ftype_to_ggml_type(enum ggml_ftype ftype) {
case GGML_FTYPE_MOSTLY_Q4_0: wtype = GGML_TYPE_Q4_0; break;
case GGML_FTYPE_MOSTLY_Q4_1: wtype = GGML_TYPE_Q4_1; break;
case GGML_FTYPE_MOSTLY_Q1_0: wtype = GGML_TYPE_Q1_0; break;
case GGML_FTYPE_MOSTLY_Q2_0: wtype = GGML_TYPE_Q2_0; break;
case GGML_FTYPE_MOSTLY_Q5_0: wtype = GGML_TYPE_Q5_0; break;
case GGML_FTYPE_MOSTLY_Q5_1: wtype = GGML_TYPE_Q5_1; break;
case GGML_FTYPE_MOSTLY_Q8_0: wtype = GGML_TYPE_Q8_0; break;
@ -3933,7 +3942,7 @@ struct ggml_tensor * ggml_set_rows(
GGML_ASSERT(b->ne[2] % c->ne[1] == 0);
GGML_ASSERT(b->ne[3] % c->ne[2] == 0);
GGML_ASSERT(c->ne[3] == 1);
GGML_ASSERT(b->type == GGML_TYPE_F32);
GGML_ASSERT(b->type == GGML_TYPE_F32 || b->type == GGML_TYPE_F16);
GGML_ASSERT(c->type == GGML_TYPE_I64 || c->type == GGML_TYPE_I32);
GGML_ASSERT(ggml_is_contiguous_rows(a));
@ -7435,6 +7444,10 @@ static int ggml_node_list_find_tensor(const struct ggml_cgraph * cgraph,
return -1;
}
static bool ggml_is_constant(const struct ggml_tensor * tensor) {
return tensor->buffer != NULL && ggml_backend_buffer_get_usage(tensor->buffer) == GGML_BACKEND_BUFFER_USAGE_WEIGHTS && (tensor->flags & GGML_TENSOR_FLAG_PARAM) == 0;
}
bool ggml_can_fuse_subgraph_ext(const struct ggml_cgraph * cgraph,
const int * node_idxs,
int count,
@ -7480,10 +7493,11 @@ bool ggml_can_fuse_subgraph_ext(const struct ggml_cgraph * cgraph,
return false;
}
// if node is a view, check if the view_src and all it's parent view_srcs are within the subgraph
// if node is a view, check if the view_src and all its parent view_srcs are within the subgraph.
// external view sources are allowed only for weight tensors, which are constant for this graph execution.
struct ggml_tensor * view_src = node->view_src;
while (view_src) {
if (ggml_node_list_find_tensor(cgraph, node_idxs, count, view_src) == -1) {
if (ggml_node_list_find_tensor(cgraph, node_idxs, count, view_src) == -1 && !ggml_is_constant(view_src)) {
return false;
}
view_src = view_src->view_src;
@ -7755,6 +7769,7 @@ size_t ggml_quantize_chunk(
switch (type) {
case GGML_TYPE_Q1_0: result = quantize_q1_0 (src + start, (char *) dst + start_row * row_size, nrows, n_per_row, imatrix); break;
case GGML_TYPE_Q2_0: result = quantize_q2_0 (src + start, (char *) dst + start_row * row_size, nrows, n_per_row, imatrix); break;
case GGML_TYPE_Q4_0: result = quantize_q4_0 (src + start, (char *) dst + start_row * row_size, nrows, n_per_row, imatrix); break;
case GGML_TYPE_Q4_1: result = quantize_q4_1 (src + start, (char *) dst + start_row * row_size, nrows, n_per_row, imatrix); break;
case GGML_TYPE_Q5_0: result = quantize_q5_0 (src + start, (char *) dst + start_row * row_size, nrows, n_per_row, imatrix); break;

3
gguf-py/gguf/constants.py
View file

@ -4533,6 +4533,7 @@ class GGMLQuantizationType(IntEnum):
MXFP4 = 39
NVFP4 = 40
Q1_0 = 41
Q2_0 = 42
class ExpertGatingFuncType(IntEnum):
@ -4588,6 +4589,7 @@ class LlamaFileType(IntEnum):
MOSTLY_MXFP4_MOE = 38 # except 1d tensors
MOSTLY_NVFP4 = 39 # except 1d tensors
MOSTLY_Q1_0 = 40 # except 1d tensors
MOSTLY_Q2_0 = 41 # except 1d tensors
GUESSED = 1024 # not specified in the model file
@ -4713,6 +4715,7 @@ GGML_QUANT_SIZES: dict[GGMLQuantizationType, tuple[int, int]] = {
GGMLQuantizationType.MXFP4: (32, 1 + 16),
GGMLQuantizationType.NVFP4: (64, 4 + 32),
GGMLQuantizationType.Q1_0: (128, 2 + 16),
GGMLQuantizationType.Q2_0: (64, 2 + 16),
}

1
include/llama.h
View file

@ -158,6 +158,7 @@ extern "C" {
LLAMA_FTYPE_MOSTLY_MXFP4_MOE = 38, // except 1d tensors
LLAMA_FTYPE_MOSTLY_NVFP4 = 39, // except 1d tensors
LLAMA_FTYPE_MOSTLY_Q1_0 = 40, // except 1d tensors
LLAMA_FTYPE_MOSTLY_Q2_0 = 41, // except 1d tensors
LLAMA_FTYPE_GUESSED = 1024, // not specified in the model file
};

74
src/llama-batch.cpp
View file

@ -505,7 +505,7 @@ llama_ubatch llama_batch_allocr::split_simple(uint32_t n_ubatch) {
return ubatch_add(idxs, idxs.size(), false);
}
llama_ubatch llama_batch_allocr::split_equal(uint32_t n_ubatch, bool sequential) {
llama_ubatch llama_batch_allocr::split_equal(uint32_t n_ubatch, bool sequential, uint32_t n_keep_tail) {
if (sequential && has_cpl) {
LLAMA_LOG_ERROR("%s: sequential split is not supported when there are coupled sequences in the input batch (you may need to use the -kvu flag)\n", __func__);
@ -548,7 +548,7 @@ llama_ubatch llama_batch_allocr::split_equal(uint32_t n_ubatch, bool sequential)
}
}
const uint32_t n_seqs = cur_seq_set.size();
uint32_t n_seqs = cur_seq_set.size();
// we are done
if (n_seqs == 0) {
@ -569,7 +569,7 @@ llama_ubatch llama_batch_allocr::split_equal(uint32_t n_ubatch, bool sequential)
std::vector<idx_vec_t> idxs_per_seq(n_seqs);
while (true) {
// we can only add new n_seq_tokens tokens if all the sequence sets have at least one more unused token and
// we can only add new n_seq_tokens tokens if all the sequence sets have at least 1 more unused tokens and
// if we haven't reached n_ubatch
bool can_expand = true;
@ -600,6 +600,72 @@ llama_ubatch llama_batch_allocr::split_equal(uint32_t n_ubatch, bool sequential)
}
}
// if n_keep_tail > 0, keep only the seqs that either finish in this ubatch or have at least
// n_keep_tail tokens remaining for a future ubatch, so that the trailing n_keep_tail tokens
// of each seq are never split across ubatches
if (n_keep_tail > 0) {
GGML_ASSERT(n_ubatch > n_keep_tail);
auto n_remaining = [&](uint32_t s) {
return (uint32_t) (seq_set_map[cur_seq_set[s]].size() - cur_idx[s]);
};
// keep the longest prefix of seqs that satisfy the constraint, to preserve sequential seq ids
uint32_t n_keep = 0;
while (n_keep < n_seqs) {
const uint32_t remaining = n_remaining(n_keep);
if (remaining != 0 && remaining < n_keep_tail) {
break;
}
n_keep++;
}
// all seqs violate the constraint - resolve the first one directly and emit it alone
if (n_keep == 0) {
auto & idxs = idxs_per_seq[0];
const auto & seq_idxs = seq_set_map[cur_seq_set[0]];
if (idxs.size() + n_remaining(0) <= n_ubatch) {
// extend the seq to completion
while (n_remaining(0) > 0) {
const int32_t idx = seq_idxs[cur_idx[0]];
idxs.push_back(idx);
used[idx] = true;
++n_used;
++cur_idx[0];
}
} else {
// truncate the seq so that at least n_keep_tail tokens remain
while (n_remaining(0) < n_keep_tail) {
used[idxs.back()] = false;
--n_used;
idxs.pop_back();
--cur_idx[0];
}
}
n_keep = 1;
}
// return the tokens of the deferred seqs back to the pool
for (uint32_t s = n_keep; s < n_seqs; ++s) {
for (const int32_t idx : idxs_per_seq[s]) {
used[idx] = false;
--n_used;
}
}
n_seqs = n_keep;
}
// concat the per-sequence-set lists
std::vector<int32_t> idxs;
@ -814,7 +880,7 @@ void llama_batch_allocr::ubatch_print(const llama_ubatch & ubatch, int debug) {
LLAMA_LOG_DEBUG("%s: output = %p\n", __func__, (void *) ubatch.output);
LLAMA_LOG_DEBUG("%s: n_outputs = %d\n", __func__, n_outputs);
if (debug > 1) {
if (debug > 0) {
int seq_id_max = 0;
for (uint32_t i = 0; i < ubatch.n_tokens; ++i) {
for (int s = 0; s < ubatch.n_seq_id[i]; ++s) {

3
src/llama-batch.h
View file

@ -104,7 +104,8 @@ public:
// make ubatches of equal-length sequences sets
// if sequential == true, the tokens in the ubatch will have increasing sequential sequence ids
llama_ubatch split_equal(uint32_t n_ubatch, bool sequential);
// n_keep_tail = minimum trailing tokens of a seq that must land in the same ubatch
llama_ubatch split_equal(uint32_t n_ubatch, bool sequential, uint32_t n_keep_tail);
// sequence-set-wise split - each ubatch contains a single sequence-set
llama_ubatch split_seq(uint32_t n_ubatch);

2
src/llama-kv-cache-dsa.cpp
View file

@ -113,7 +113,7 @@ llama_memory_context_ptr llama_kv_cache_dsa::init_batch(
std::vector<llama_ubatch> ubatches;
while (true) {
auto ubatch = n_stream == 1 ? balloc.split_simple(n_ubatch) : balloc.split_equal(n_ubatch, true);
auto ubatch = n_stream == 1 ? balloc.split_simple(n_ubatch) : balloc.split_equal(n_ubatch, true, 0);
if (ubatch.n_tokens == 0) {
break;

2
src/llama-kv-cache-dsv4.cpp
View file

@ -1110,7 +1110,7 @@ llama_memory_context_ptr llama_kv_cache_dsv4::init_batch(
if (has_coupled) {
ubatch = balloc.split_seq(n_ubatch);
} else {
ubatch = balloc.split_equal(n_ubatch, raw_per_seq || comp_per_seq);
ubatch = balloc.split_equal(n_ubatch, raw_per_seq || comp_per_seq, 0);
}
if (ubatch.n_tokens == 0) {

2
src/llama-kv-cache-iswa.cpp
View file

@ -219,7 +219,7 @@ llama_memory_context_ptr llama_kv_cache_iswa::init_batch(llama_batch_allocr & ba
std::vector<llama_ubatch> ubatches;
while (true) {
auto ubatch = balloc.split_equal(n_ubatch, !unified);
auto ubatch = balloc.split_equal(n_ubatch, !unified, 0);
if (ubatch.n_tokens == 0) {
break;

2
src/llama-kv-cache.cpp
View file

@ -706,7 +706,7 @@ llama_memory_context_ptr llama_kv_cache::init_batch(
std::vector<llama_ubatch> ubatches;
while (true) {
auto ubatch = n_stream == 1 ? balloc.split_simple(n_ubatch) : balloc.split_equal(n_ubatch, true);
auto ubatch = n_stream == 1 ? balloc.split_simple(n_ubatch) : balloc.split_equal(n_ubatch, true, 0);
if (ubatch.n_tokens == 0) {
break;

18
src/llama-memory-hybrid-iswa.cpp
View file

@ -77,15 +77,15 @@ llama_memory_context_ptr llama_memory_hybrid_iswa::init_batch(llama_batch_allocr
// if all tokens are output, split by sequence
ubatch = balloc.split_seq(n_ubatch);
} else {
if (mem_recr->n_rs_seq > 0) {
// [TAG_RECURRENT_ROLLBACK_SPLITS]
// TODO: recurrent state rollback does not support equal splits
ubatch = balloc.split_seq(n_ubatch);
} else {
// Use non-sequential split when KV cache is unified (needed for hellaswag/winogrande/multiple-choice)
const bool unified = (mem_attn->get_base()->get_n_stream() == 1);
ubatch = balloc.split_equal(n_ubatch, !unified);
}
// Use non-sequential split when KV cache is unified (needed for hellaswag/winogrande/multiple-choice)
const bool unified = (mem_attn->get_base()->get_n_stream() == 1);
// [TAG_RECURRENT_ROLLBACK_SPLITS]
// the trailing (1 + n_rs_seq) tokens of each seq must stay in the same ubatch
// so that the rollback snapshots remain valid
const uint32_t n_rs_seq = mem_recr->n_rs_seq;
ubatch = balloc.split_equal(n_ubatch, !unified, n_rs_seq > 0 ? n_rs_seq + 1 : 0);
}
if (ubatch.n_tokens == 0) {

18
src/llama-memory-hybrid.cpp
View file

@ -78,15 +78,15 @@ llama_memory_context_ptr llama_memory_hybrid::init_batch(llama_batch_allocr & ba
// if all tokens are output, split by sequence
ubatch = balloc.split_seq(n_ubatch);
} else {
if (mem_recr->n_rs_seq > 0) {
// [TAG_RECURRENT_ROLLBACK_SPLITS]
// TODO: recurrent state rollback does not support equal splits
ubatch = balloc.split_seq(n_ubatch);
} else {
// Use non-sequential split when KV cache is unified (needed for hellaswag/winogrande/multiple-choice)
const bool unified = (mem_attn->get_n_stream() == 1);
ubatch = balloc.split_equal(n_ubatch, !unified);
}
// Use non-sequential split when KV cache is unified (needed for hellaswag/winogrande/multiple-choice)
const bool unified = (mem_attn->get_n_stream() == 1);
// [TAG_RECURRENT_ROLLBACK_SPLITS]
// the trailing (1 + n_rs_seq) tokens of each seq must stay in the same ubatch
// so that the rollback snapshots remain valid
const uint32_t n_rs_seq = mem_recr->n_rs_seq;
ubatch = balloc.split_equal(n_ubatch, !unified, n_rs_seq > 0 ? n_rs_seq + 1 : 0);
}
if (ubatch.n_tokens == 0) {

15
src/llama-memory-recurrent.cpp
View file

@ -416,15 +416,12 @@ llama_memory_context_ptr llama_memory_recurrent::init_batch(llama_batch_allocr &
// if all tokens are output, split by sequence
ubatch = balloc.split_seq(n_ubatch);
} else {
if (n_rs_seq > 0) {
// [TAG_RECURRENT_ROLLBACK_SPLITS]
// TODO: recurrent state rollback does not support equal splits
ubatch = balloc.split_seq(n_ubatch);
} else {
// TODO: non-sequential equal split can be done if using unified KV cache
// for simplicity, we always use sequential equal split for now
ubatch = balloc.split_equal(n_ubatch, true);
}
// TODO: non-sequential equal split can be done if using unified KV cache
// for simplicity, we always use sequential equal split for now
// [TAG_RECURRENT_ROLLBACK_SPLITS]
// the trailing (1 + n_rs_seq) tokens of each seq must stay in the same ubatch
// so that the rollback snapshots remain valid
ubatch = balloc.split_equal(n_ubatch, true, n_rs_seq > 0 ? n_rs_seq + 1 : 0);
}
if (ubatch.n_tokens == 0) {

2
src/llama-model-loader.cpp
View file

@ -37,6 +37,7 @@ const char * llama_ftype_name(llama_ftype ftype) {
case LLAMA_FTYPE_MOSTLY_F16: name = LLAMA_FTYPE_PREFIX "F16"; break;
case LLAMA_FTYPE_MOSTLY_BF16: name = LLAMA_FTYPE_PREFIX "BF16"; break;
case LLAMA_FTYPE_MOSTLY_Q1_0: name = LLAMA_FTYPE_PREFIX "Q1_0"; break;
case LLAMA_FTYPE_MOSTLY_Q2_0: name = LLAMA_FTYPE_PREFIX "Q2_0"; break;
case LLAMA_FTYPE_MOSTLY_Q4_0: name = LLAMA_FTYPE_PREFIX "Q4_0"; break;
case LLAMA_FTYPE_MOSTLY_Q4_1: name = LLAMA_FTYPE_PREFIX "Q4_1"; break;
case LLAMA_FTYPE_MOSTLY_Q5_0: name = LLAMA_FTYPE_PREFIX "Q5_0"; break;
@ -768,6 +769,7 @@ llama_model_loader::llama_model_loader(
case GGML_TYPE_IQ3_S: ftype = LLAMA_FTYPE_MOSTLY_IQ3_S; break;
case GGML_TYPE_NVFP4: ftype = LLAMA_FTYPE_MOSTLY_NVFP4; break;
case GGML_TYPE_Q1_0: ftype = LLAMA_FTYPE_MOSTLY_Q1_0; break;
case GGML_TYPE_Q2_0: ftype = LLAMA_FTYPE_MOSTLY_Q2_0; break;
default:
{
LLAMA_LOG_WARN("%s: unknown type %s\n", __func__, ggml_type_name(type_max));

4
src/llama-quant.cpp
View file

@ -380,6 +380,7 @@ static ggml_type tensor_type_fallback(quantize_state_impl & qs, const ggml_tenso
case GGML_TYPE_IQ3_XXS:
case GGML_TYPE_IQ3_S: // types on the right: block size 32
case GGML_TYPE_IQ4_XS: return_type = GGML_TYPE_IQ4_NL; break;
case GGML_TYPE_Q2_0:
case GGML_TYPE_Q2_K:
case GGML_TYPE_Q3_K:
case GGML_TYPE_TQ1_0:
@ -482,7 +483,7 @@ static ggml_type llama_tensor_get_type_impl(quantize_state_impl & qs, ggml_type
else if (ftype == LLAMA_FTYPE_MOSTLY_IQ3_XXS) {
new_type = GGML_TYPE_IQ3_S;
}
else if (ftype == LLAMA_FTYPE_MOSTLY_TQ1_0 || ftype == LLAMA_FTYPE_MOSTLY_TQ2_0) {
else if (ftype == LLAMA_FTYPE_MOSTLY_TQ1_0 || ftype == LLAMA_FTYPE_MOSTLY_TQ2_0 || ftype == LLAMA_FTYPE_MOSTLY_Q2_0) {
new_type = GGML_TYPE_Q4_K;
}
}
@ -802,6 +803,7 @@ ggml_type llama_ftype_get_default_type(llama_ftype ftype) {
case LLAMA_FTYPE_MOSTLY_BF16: return GGML_TYPE_BF16;
case LLAMA_FTYPE_ALL_F32: return GGML_TYPE_F32;
case LLAMA_FTYPE_MOSTLY_Q1_0: return GGML_TYPE_Q1_0;
case LLAMA_FTYPE_MOSTLY_Q2_0: return GGML_TYPE_Q2_0;
case LLAMA_FTYPE_MOSTLY_MXFP4_MOE: return GGML_TYPE_MXFP4;

26
src/llama-vocab.cpp
View file

@ -1112,9 +1112,6 @@ struct llm_tokenizer_ugm : llm_tokenizer {
// blob containing XOR-compressed compact double array (XCDA) entries
uint32_t xcda_blob_size = *(const uint32_t *) &precompiled_charsmap[0];
charsmap_offset += sizeof(xcda_blob_size);
if (xcda_blob_size + charsmap_offset >= precompiled_charsmap.size()) {
throw std::runtime_error("Index out of array bounds in precompiled charsmap!");
}
// Next xcda_blob_size bytes contain entries of XOR-compressed compact
// double array (XCDA). Each entry is bit-packed into a 32-bit integer.
@ -1430,7 +1427,15 @@ private:
throw std::runtime_error("Index out of array bounds in precompiled charsmap!");
}
const char * prefix_replacement = &(tokenizer.prefix_replacements)[longest_prefix_offset];
return { prefix_replacement, strlen(prefix_replacement), longest_prefix_length };
size_t max_len = tokenizer.prefix_replacements_size - longest_prefix_offset;
size_t repl_len = 0;
while (repl_len < max_len && prefix_replacement[repl_len] != '\0') {
repl_len++;
}
if (repl_len == max_len) {
throw std::runtime_error("Unterminated string in precompiled charsmap!");
}
return { prefix_replacement, repl_len, longest_prefix_length };
}
// check if the input prefix contains a valid sequence of UTF-8 code units
@ -2254,11 +2259,18 @@ void llama_vocab::impl::load(llama_model_loader & ml, const LLM_KV & kv) {
const size_t n_precompiled_charsmap = gguf_get_arr_n(ctx, precompiled_charsmap_keyidx);
const char * pc = (const char *) gguf_get_arr_data(ctx, precompiled_charsmap_keyidx);
precompiled_charsmap.assign(pc, pc + n_precompiled_charsmap);
if (precompiled_charsmap.size() < sizeof(uint32_t)) {
throw std::runtime_error("precompiled_charsmap too small for xcda_blob_size header!");
}
uint32_t * xcda_blob_size = (uint32_t *) &precompiled_charsmap[0];
#if defined(__BYTE_ORDER__) && defined(__ORDER_BIG_ENDIAN__) && __BYTE_ORDER__ == __ORDER_BIG_ENDIAN__
*xcda_blob_size = __builtin_bswap32(*xcda_blob_size);
#endif
if (*xcda_blob_size + sizeof(uint32_t) >= precompiled_charsmap.size()) {
throw std::runtime_error("Index out of array bounds in precompiled charsmap!");
}
#if defined(__BYTE_ORDER__) && defined(__ORDER_BIG_ENDIAN__) && __BYTE_ORDER__ == __ORDER_BIG_ENDIAN__
// correct endianness of data in precompiled_charsmap binary blob
uint32_t * xcda_blob_size = (uint32_t *) &precompiled_charsmap[0];
*xcda_blob_size = __builtin_bswap32(*xcda_blob_size);
assert(*xcda_blob_size + sizeof(uint32_t) < n_precompiled_charsmap);
size_t xcda_array_size = *xcda_blob_size / sizeof(uint32_t);
uint32_t * xcda_array = (uint32_t *) &precompiled_charsmap[sizeof(uint32_t)];
for (size_t i = 0; i < xcda_array_size; ++i) {

4
src/models/delta-net-base.cpp
View file

@ -496,8 +496,8 @@ ggml_tensor * llm_build_delta_net_base::build_conv_state(
ggml_build_forward_expand(gf, ggml_cpy(ctx0, conv_state_last, conv_state_update));
} else {
// [TAG_RECURRENT_ROLLBACK_SPLITS]
// TODO: this logic incorrectly assumes that the last (n_rs_seq + 1) tokens of a sequence in a batch are
// inside the same ubatch. currently with `split_equal()` this is not correct
// this logic assumes that the last (n_rs_seq + 1) tokens of a sequence in a batch are inside
// the same ubatch, which `split_equal()` guarantees via its n_keep_tail argument
const int64_t K = (int64_t) cparams.n_rs_seq + 1;

1
tools/quantize/quantize.cpp
View file

@ -32,6 +32,7 @@ struct quant_option {
static const std::vector<quant_option> QUANT_OPTIONS = {
{ "Q1_0", LLAMA_FTYPE_MOSTLY_Q1_0, " 1.125 bpw quantization", },
{ "Q2_0", LLAMA_FTYPE_MOSTLY_Q2_0, " 2.25 bpw quantization (group 64)", },
{ "Q4_0", LLAMA_FTYPE_MOSTLY_Q4_0, " 4.34G, +0.4685 ppl @ Llama-3-8B", },
{ "Q4_1", LLAMA_FTYPE_MOSTLY_Q4_1, " 4.78G, +0.4511 ppl @ Llama-3-8B", },
{ "MXFP4_MOE",LLAMA_FTYPE_MOSTLY_MXFP4_MOE," MXFP4 MoE", },

10
tools/server/README-dev.md
View file

@ -57,7 +57,7 @@ The core architecture consists of the following components:
- `server_tokens`: Unified representation of token sequences (supports both text and multimodal tokens); used by `server_task` and `server_slot`.
- `server_prompt_checkpoint`: For recurrent (e.g., RWKV) and SWA models, stores snapshots of KV cache state. Enables reuse when subsequent requests share the same prompt prefix, saving redundant computation.
- `server_models`: Standalone component for managing multiple backend instances (used in router mode). It is completely independent of `server_context`.
- `stream_session_manager`: Process wide owner of resumable SSE stream sessions (`g_stream_sessions`), keyed by conversation id. Backs the replay buffer that lets a client reattach to a generation after an HTTP disconnect. See the "Resumable streaming" section below.
- `stream_session_manager`: process wide owner of resumable SSE stream sessions, keyed by conversation id. A file-static singleton inside `server-stream.cpp`, driven through `server_stream_session_manager_start/stop`. Backs the replay buffer that lets a client reattach to a generation after an HTTP disconnect. See the "Resumable streaming" section below.
```mermaid
graph TD
@ -127,10 +127,12 @@ It is opt in via the `X-Conversation-Id` header on `POST /v1/chat/completions`.
The feature lives entirely in `server-stream.{h,cpp}` and rests on three types:
- `stream_session`: a bounded ring buffer (4 MiB cap, oldest bytes drop first) plus a condvar. `append` pushes raw SSE bytes, `read_from` drains from any offset and blocks for live bytes or finalize, `finalize` wakes readers, `cancel` stops the producer. One conv maps to at most one live session.
- `stream_session_manager` (`g_stream_sessions`): owns all sessions keyed by conv id, enforces the one conv one session invariant via `create_or_replace`, and runs a GC thread that drops completed sessions past their TTL.
- `stream_session_manager`: a file-static singleton (`g_stream_sessions`) inside `server-stream.cpp`, owns all sessions keyed by conv id, enforces the one conv one session invariant via `create_or_replace`, and runs a GC thread that drops completed sessions past their TTL. Exposed to main only through `server_stream_session_manager_start/stop`.
- `stream_pipe_producer` / `stream_pipe_consumer`: the write and read ends. The producer owns the session lifetime and finalizes it on destruction; the consumer is read only and never finalizes, so a reader detaching cannot kill a running generation.
Producer side: `server_res_generator` attaches a producer pipe when the header is present. The HTTP content provider mirrors every chunk into the ring before writing it to the socket. While a pipe is attached, `stream_aware_should_stop` ignores peer disconnect, so a dropped socket does not stop generation: only an explicit `DELETE` does. When the peer leaves early, `on_complete` calls `close()`, which drains the rest of the generation into the ring on the http worker.
The implementation is hidden in `server-stream.cpp` (pimpl). The header exposes only the route handler factories, `server_stream_session_attach_pipe`, `server_stream_aware_should_stop`, `server_stream_conv_id_from_headers` and the GC lifecycle; the session, manager and consumer types stay in the `.cpp`.
Producer side: `server_res_generator` attaches a producer pipe when the header is present. The HTTP content provider mirrors every chunk into the ring before writing it to the socket. While a pipe is attached, `server_stream_aware_should_stop` ignores peer disconnect, so a dropped socket does not stop generation: only an explicit `DELETE` does. When the peer leaves early, `on_complete` calls `close()`, which drains the rest of the generation into the ring on the http worker.
Lifetime safety: the producer pipe holds a shared `alive` flag also captured by the session cancel hook. `~server_res_generator` calls `cleanup()` to clear that hook while the reader is still alive, so a `cancel` arriving during teardown can never call `stop()` on a freed response. This ordering is the most fragile part of the feature: finalizing or destroying the producer before `cleanup()` runs reintroduces a use after free.
@ -144,7 +146,7 @@ Routes:
Router mode binds the same paths to proxy handlers. A `conv_id -> child` map (`conv_models`), populated when a POST is routed, resolves the owning child in one lookup with no polling. The lookup groups ids per child; GET and DELETE proxy straight to the owner. This loopback REST hop is expected to move to a websocket IPC later, swapping only the transport.
Lifecycle: `g_stream_sessions.start_gc()` runs in main after common init, `stop_gc()` runs first in `clean_up()` and finalizes every live session so no reader hangs. Reader blocking and the post drop drain both run on httplib worker threads, which block on a condvar rather than spin.
Lifecycle: `server_stream_session_manager_start()` runs in main after common init, `server_stream_session_manager_stop()` runs first in `clean_up()` and finalizes every live session so no reader hangs. Reader blocking and the post drop drain both run on httplib worker threads, which block on a condvar rather than spin.
| Constant | Value | Role |
| --- | --- | --- |

167
tools/server/server-context.cpp
View file

@ -897,8 +897,10 @@ private:
server_batch batch;
llama_model_ptr model_dft;
llama_context_ptr ctx_dft;
llama_model * model_dft = nullptr;
llama_context * ctx_dft = nullptr;
common_speculative_init_result_ptr spec_init;
common_context_seq_rm_type ctx_tgt_seq_rm_type = COMMON_CONTEXT_SEQ_RM_TYPE_NO;
common_context_seq_rm_type ctx_dft_seq_rm_type = COMMON_CONTEXT_SEQ_RM_TYPE_NO;
@ -939,8 +941,10 @@ private:
void destroy() {
spec.reset();
ctx_dft.reset();
model_dft.reset();
spec_init.reset();
ctx_dft = nullptr;
model_dft = nullptr;
llama_init.reset();
@ -1084,30 +1088,15 @@ private:
// optionally reserve VRAM for the draft / MTP context before fitting the target model
if (params_base.fit_params) {
if (has_spec) {
common_params params_dft = params_base;
bool measure_model_bytes = true;
// MTP draft context lives on the target model, only context+compute are new
bool measure_model_bytes = has_draft;
if (has_draft) {
const auto & params_spec = params_base.speculative.draft;
params_dft.devices = params_spec.devices;
params_dft.model = params_spec.mparams;
params_dft.n_gpu_layers = params_spec.n_gpu_layers;
params_dft.cache_type_k = params_spec.cache_type_k;
params_dft.cache_type_v = params_spec.cache_type_v;
params_dft.tensor_buft_overrides = params_spec.tensor_buft_overrides;
} else {
// MTP draft context lives on the target model, only context+compute are new
measure_model_bytes = false;
}
params_dft.n_outputs_max = params_base.n_parallel;
common_params params_dft = common_base_params_to_speculative(params_base);
auto mparams_dft = common_model_params_to_llama(params_dft);
auto cparams_dft = common_context_params_to_llama(params_dft);
if (spec_mtp) {
cparams_dft.ctx_type = LLAMA_CONTEXT_TYPE_MTP;
cparams_dft.type_k = params_base.speculative.draft.cache_type_k;
cparams_dft.type_v = params_base.speculative.draft.cache_type_v;
}
cparams_dft.n_rs_seq = 0;
@ -1175,82 +1164,36 @@ private:
add_bos_token = llama_vocab_get_add_bos(vocab);
if (has_draft) {
// TODO speculative: move to common/speculative.cpp?
const auto & params_spec = params_base.speculative.draft;
SRV_TRC("loading draft model '%s'\n", params_spec.mparams.path.c_str());
auto params_dft = params_base;
params_dft.devices = params_spec.devices;
params_dft.model = params_spec.mparams;
params_dft.n_gpu_layers = params_spec.n_gpu_layers;
params_dft.cache_type_k = params_spec.cache_type_k;
params_dft.cache_type_v = params_spec.cache_type_v;
if (params_spec.cpuparams.n_threads > 0) {
params_dft.cpuparams.n_threads = params_spec.cpuparams.n_threads;
params_dft.cpuparams_batch.n_threads = params_spec.cpuparams_batch.n_threads;
}
params_dft.tensor_buft_overrides = params_spec.tensor_buft_overrides;
auto mparams_dft = common_model_params_to_llama(params_dft);
// progress callback
mparams_dft.progress_callback = load_progress_callback;
mparams_dft.progress_callback_user_data = &load_progress_spec;
model_dft.reset(llama_model_load_from_file(params_dft.model.path.c_str(), mparams_dft));
if (model_dft == nullptr) {
SRV_ERR("failed to load draft model, '%s'\n", params_dft.model.path.c_str());
return false;
}
auto cparams = common_context_params_to_llama(params_dft);
if (spec_mtp) {
cparams.ctx_type = LLAMA_CONTEXT_TYPE_MTP;
}
// note: for small models maybe we can set this to the maximum possible draft from all speculative types
// the extra memory for small models is likely negligible?
cparams.n_rs_seq = 0;
cparams.ctx_other = ctx_tgt;
ctx_dft.reset(llama_init_from_model(model_dft.get(), cparams));
if (ctx_dft == nullptr) {
SRV_ERR("%s", "failed to create draft context\n");
return false;
}
params_base.speculative.draft.ctx_tgt = ctx_tgt;
params_base.speculative.draft.ctx_dft = ctx_dft.get();
} else if (spec_mtp) {
// no new model load, so we simply report 0.0 and 1.0 progress
if (has_spec) {
// spec_mtp doesn't use load a model internally, so we report 0.0 and 1.0 manually
load_progress_callback(0.0f, &load_progress_spec);
load_progress_spec.t_last_load_progress_ms = 0; // reset so internal cbs aren't delayed
SRV_TRC("creating MTP draft context against the target model '%s'\n",
params_base.model.path.c_str());
{
common_params params_dft = common_base_params_to_speculative(params_base);
auto cparams_mtp = common_context_params_to_llama(params_base);
cparams_mtp.ctx_type = LLAMA_CONTEXT_TYPE_MTP;
cparams_mtp.type_k = params_base.speculative.draft.cache_type_k;
cparams_mtp.type_v = params_base.speculative.draft.cache_type_v;
cparams_mtp.n_rs_seq = 0;
cparams_mtp.n_outputs_max = params_base.n_parallel;
cparams_mtp.ctx_other = ctx_tgt;
// progress callback
params_dft.load_progress_callback = load_progress_callback;
params_dft.load_progress_callback_user_data = &load_progress_spec;
ctx_dft.reset(llama_init_from_model(model_tgt, cparams_mtp));
if (ctx_dft == nullptr) {
SRV_ERR("%s", "failed to create MTP context\n");
return false;
spec_init = common_speculative_init_from_params(params_dft, model_tgt, ctx_tgt);
model_dft = spec_init->model();
ctx_dft = spec_init->context();
if (has_draft && model_dft == nullptr) {
SRV_ERR("failed to load draft model, '%s'\n", params_dft.model.path.c_str());
return false;
}
if (ctx_dft == nullptr) {
SRV_ERR("%s", "failed to create MTP context\n");
return false;
}
params_base.speculative.draft.ctx_tgt = ctx_tgt;
params_base.speculative.draft.ctx_dft = ctx_dft;
}
params_base.speculative.draft.ctx_tgt = ctx_tgt;
params_base.speculative.draft.ctx_dft = ctx_dft.get();
load_progress_callback(1.0f, &load_progress_spec);
}
@ -1343,13 +1286,15 @@ private:
}
if (ctx_dft) {
ctx_dft_seq_rm_type = common_context_can_seq_rm(ctx_dft.get());
ctx_dft_seq_rm_type = common_context_can_seq_rm(ctx_dft);
}
if (spec) {
SRV_TRC("%s", "speculative decoding context initialized\n");
} else {
ctx_dft.reset();
spec_init.reset();
ctx_dft = nullptr;
model_dft = nullptr;
}
for (int i = 0; i < params_base.n_parallel; i++) {
@ -1357,7 +1302,7 @@ private:
slot.id = i;
slot.ctx_tgt = ctx_tgt;
slot.ctx_dft = ctx_dft.get();
slot.ctx_dft = ctx_dft;
slot.spec = spec.get();
slot.n_ctx = n_ctx_slot;
@ -2362,8 +2307,8 @@ private:
// this is not true for SWA models: https://github.com/ggml-org/llama.cpp/pull/24411#issuecomment-4677983225
cur.update_pos(slot.prompt.n_tokens() - n_tokens_cur, pos_min, pos_max);
cur.update_tgt(ctx_tgt, slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY);
cur.update_dft(ctx_dft.get(), slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY);
cur.update_tgt(ctx_tgt, slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY);
cur.update_dft(ctx_dft, slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY);
// stash the draft's speculative state with the checkpoint
common_speculative_get_state(spec.get(), slot.id, cur.data_spec);
@ -2899,8 +2844,8 @@ private:
common_context_seq_add(ctx_tgt, slot.id, n_keep + n_discard, slot.prompt.n_tokens(), -n_discard);
if (ctx_dft) {
common_context_seq_rm (ctx_dft.get(), slot.id, n_keep , n_keep + n_discard);
common_context_seq_add(ctx_dft.get(), slot.id, n_keep + n_discard, slot.prompt.tokens.pos_next(), -n_discard);
common_context_seq_rm (ctx_dft, slot.id, n_keep , n_keep + n_discard);
common_context_seq_add(ctx_dft, slot.id, n_keep + n_discard, slot.prompt.tokens.pos_next(), -n_discard);
}
// add generated tokens to cache
@ -2972,7 +2917,7 @@ private:
llama_memory_seq_pos_max(llama_get_memory(ctx_tgt), slot.id));
if (use_ckpt_dft) {
slot.spec_ckpt.update_dft(ctx_dft.get(), slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY);
slot.spec_ckpt.update_dft(ctx_dft, slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY);
}
slot.spec_prompt = slot.prompt.tokens.get_text_tokens();
@ -3009,10 +2954,10 @@ private:
if (ctx_dft) {
if (use_ckpt_dft) {
ckpt.load_dft(ctx_dft.get(), slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY);
ckpt.load_dft(ctx_dft, slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY);
}
common_context_seq_rm(ctx_dft.get(), slot.id, ckpt.pos_max + 1, -1);
common_context_seq_rm(ctx_dft, slot.id, ckpt.pos_max + 1, -1);
}
if (!draft.empty()) {
@ -3021,7 +2966,7 @@ private:
(ctx_tgt_seq_rm_type == COMMON_CONTEXT_SEQ_RM_TYPE_RS && draft.size() > llama_n_rs_seq(ctx_tgt));
const bool use_ckpt_dft =
(ctx_dft_seq_rm_type == COMMON_CONTEXT_SEQ_RM_TYPE_RS && draft.size() > llama_n_rs_seq(ctx_dft.get()));
(ctx_dft_seq_rm_type == COMMON_CONTEXT_SEQ_RM_TYPE_RS && draft.size() > llama_n_rs_seq(ctx_dft));
if (use_ckpt_tgt) {
//const int64_t t_start = ggml_time_us();
@ -3038,7 +2983,7 @@ private:
}
if (use_ckpt_dft) {
ckpt.update_dft(ctx_dft.get(), slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY);
ckpt.update_dft(ctx_dft, slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY);
}
}
});
@ -3219,8 +3164,8 @@ private:
common_context_seq_add(ctx_tgt, slot.id, head_c, head_c + n_match, kv_shift);
if (ctx_dft) {
common_context_seq_rm (ctx_dft.get(), slot.id, head_p, head_c);
common_context_seq_add(ctx_dft.get(), slot.id, head_c, head_c + n_match, kv_shift);
common_context_seq_rm (ctx_dft, slot.id, head_p, head_c);
common_context_seq_add(ctx_dft, slot.id, head_c, head_c + n_match, kv_shift);
}
for (size_t i = 0; i < n_match; i++) {
@ -3320,8 +3265,8 @@ private:
if (!do_reset) {
// restore the context checkpoint
it->load_tgt(ctx_tgt, slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY);
it->load_dft(ctx_dft.get(), slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY);
it->load_tgt(ctx_tgt, slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY);
it->load_dft(ctx_dft, slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY);
// restore the draft's speculative state
common_speculative_set_state(spec.get(), slot.id, it->data_spec);
@ -3395,7 +3340,7 @@ private:
common_context_seq_rm(ctx_tgt, slot.id, p0, -1);
if (ctx_dft) {
common_context_seq_rm(ctx_dft.get(), slot.id, p0, -1);
common_context_seq_rm(ctx_dft, slot.id, p0, -1);
}
// If using an alora, there may be uncached tokens that come
@ -4243,7 +4188,7 @@ std::unique_ptr<server_res_generator> server_routes::handle_completions_impl(
}
};
auto effective_should_stop = stream_aware_should_stop(res_this, req.should_stop);
auto effective_should_stop = server_stream_aware_should_stop(res_this, req.should_stop);
try {
if (effective_should_stop()) {
@ -4339,7 +4284,7 @@ std::unique_ptr<server_res_generator> server_routes::handle_completions_impl(
// attach a producer pipe to the response when X-Conversation-Id is present.
// the pipe mirrors SSE chunks into the ring buffer and wires up the cancel hook.
stream_session_attach_pipe(*res, req.headers);
server_stream_session_attach_pipe(*res, req.headers);
return res;
}

4
tools/server/server-models.cpp
View file

@ -1681,7 +1681,7 @@ void server_models_routes::init_routes() {
}
// remember which child serves this conversation so the stream routes can route straight
// to it without polling, keyed on the exact conv id from the header
std::string conv_id = stream_conv_id_from_headers(req.headers);
std::string conv_id = server_stream_conv_id_from_headers(req.headers);
if (!conv_id.empty()) {
models.conv_models.remember(conv_id, name);
}
@ -1896,7 +1896,7 @@ void server_models_routes::init_routes() {
if (!from.empty()) {
child_path += "?from=" + from;
}
SRV_INF("proxying stream resume to model %s on port %d, path=%s\n",
SRV_TRC("proxying stream resume to model %s on port %d, path=%s\n",
owner->name.c_str(), owner->port, child_path.c_str());
auto proxy = std::make_unique<server_http_proxy>(
"GET",

188
tools/server/server-stream.cpp
View file

@ -6,6 +6,12 @@
#include <chrono>
#include <memory>
#include <utility>
#include <shared_mutex>
enum class stream_read_status {
OK,
OFFSET_LOST,
};
namespace {
constexpr int64_t STREAM_SESSION_TTL_SECONDS = 300;
@ -13,7 +19,6 @@ constexpr size_t STREAM_SESSION_MAX_BYTES = 4 * 1024 * 1024;
constexpr int64_t STREAM_SESSION_GC_INTERVAL_SECONDS = 60;
constexpr int64_t STREAM_READ_WAKE_INTERVAL_MS = 200;
// returns unix time in seconds
int64_t now_seconds() {
return std::chrono::duration_cast<std::chrono::seconds>(
std::chrono::system_clock::now().time_since_epoch()
@ -21,6 +26,91 @@ int64_t now_seconds() {
}
}
// owns all live sessions keyed by conversation_id, one conv = at most one live session.
// a periodic GC evicts expired ones
class stream_session_manager {
public:
stream_session_manager();
~stream_session_manager();
stream_session_manager(const stream_session_manager &) = delete;
stream_session_manager & operator=(const stream_session_manager &) = delete;
// install a new session, evicting and cancelling any previous one. conversation_id must be non empty
stream_session_ptr create_or_replace(const std::string & conversation_id);
stream_session_ptr get(const std::string & conversation_id);
std::vector<stream_session_ptr> list_all() const;
void evict(const std::string & conversation_id);
void evict_and_cancel(const std::string & conversation_id);
void start_gc();
void stop_gc();
private:
void gc_loop();
mutable std::shared_mutex map_mu;
std::unordered_map<std::string, stream_session_ptr> sessions; // key: conversation_id
std::thread gc_thread;
bool running;
std::mutex gc_wake_mu;
std::condition_variable gc_wake_cv;
};
// process wide manager, lifecycle controlled by llama-server main() via start_gc/stop_gc
static stream_session_manager g_stream_sessions;
void server_stream_session_manager_start() {
g_stream_sessions.start_gc();
}
void server_stream_session_manager_stop() {
g_stream_sessions.stop_gc();
}
struct stream_session {
std::string conversation_id;
int64_t started_ts; // unix seconds at construction
stream_session(std::string conversation_id_, size_t max_bytes_);
stream_session(const stream_session &) = delete;
stream_session & operator=(const stream_session &) = delete;
bool append(const char * data, size_t len);
void finalize();
// drain from offset into sink, blocking for more bytes or finalize. OFFSET_LOST if offset
// fell below the dropped prefix
stream_read_status read_from(size_t offset,
const std::function<bool(const char *, size_t)> & sink,
const std::function<bool()> & should_stop);
bool is_done() const;
bool is_cancelled() const;
size_t total_size() const; // bytes that ever entered the session
size_t dropped_prefix() const; // bytes evicted from the front due to cap
int64_t completed_at() const; // 0 while alive, unix seconds after finalize
void set_stop_producer(std::function<void()> fn);
void cancel();
private:
mutable std::mutex mu;
std::condition_variable cv;
std::vector<char> buffer;
size_t prefix_dropped;
size_t cap_bytes;
bool done;
std::atomic<bool> cancelled; // polled lock-free by the should_stop closure, no mu
int64_t completed_ts;
std::function<void()> stop_producer;
};
stream_session::stream_session(std::string conversation_id_, size_t max_bytes_)
: conversation_id(std::move(conversation_id_))
, started_ts(now_seconds())
@ -38,7 +128,7 @@ bool stream_session::append(const char * data, size_t len) {
}
{
std::lock_guard<std::mutex> lock(mu);
if (done.load(std::memory_order_relaxed)) {
if (done) {
return false;
}
if (len >= cap_bytes) {
@ -62,11 +152,14 @@ bool stream_session::append(const char * data, size_t len) {
}
void stream_session::finalize() {
bool was_done = done.exchange(true, std::memory_order_acq_rel);
if (was_done) {
return;
{
std::lock_guard<std::mutex> lock(mu);
if (done) {
return;
}
done = true;
completed_ts = now_seconds();
}
completed_ts.store(now_seconds(), std::memory_order_release);
cv.notify_all();
}
@ -96,7 +189,7 @@ stream_read_status stream_session::read_from(size_t offset,
lock.lock();
continue;
}
if (done.load(std::memory_order_acquire)) {
if (done) {
return stream_read_status::OK;
}
// wait for new bytes, finalize, or a periodic wake to re check should_stop
@ -105,7 +198,8 @@ stream_read_status stream_session::read_from(size_t offset,
}
bool stream_session::is_done() const {
return done.load(std::memory_order_acquire);
std::lock_guard<std::mutex> lock(mu);
return done;
}
size_t stream_session::total_size() const {
@ -119,7 +213,8 @@ size_t stream_session::dropped_prefix() const {
}
int64_t stream_session::completed_at() const {
return completed_ts.load(std::memory_order_acquire);
std::lock_guard<std::mutex> lock(mu);
return completed_ts;
}
void stream_session::set_stop_producer(std::function<void()> fn) {
@ -128,7 +223,7 @@ void stream_session::set_stop_producer(std::function<void()> fn) {
}
void stream_session::cancel() {
// flip cancelled first so the producer-side stream_aware_should_stop can break out of the
// flip cancelled first so the producer-side server_stream_aware_should_stop can break out of the
// recv() wait even if remove_waiting_task_ids does not notify the condvar (the cancel task
// posted by rd.stop() will eventually notify, but we do not want to depend on that timing)
cancelled.store(true, std::memory_order_release);
@ -237,18 +332,24 @@ void stream_session_manager::evict_and_cancel(const std::string & conversation_i
}
void stream_session_manager::start_gc() {
if (running.exchange(true)) {
return;
{
std::lock_guard<std::mutex> lock(gc_wake_mu);
if (running) {
return;
}
running = true;
}
gc_thread = std::thread([this] { gc_loop(); });
}
void stream_session_manager::stop_gc() {
bool was_running = running.exchange(false);
bool was_running;
{
std::lock_guard<std::mutex> lock(gc_wake_mu);
was_running = running;
running = false;
}
if (was_running) {
{
std::lock_guard<std::mutex> lock(gc_wake_mu);
}
gc_wake_cv.notify_all();
if (gc_thread.joinable()) {
gc_thread.join();
@ -270,15 +371,15 @@ void stream_session_manager::stop_gc() {
}
void stream_session_manager::gc_loop() {
while (running.load(std::memory_order_acquire)) {
while (true) {
{
std::unique_lock<std::mutex> lock(gc_wake_mu);
gc_wake_cv.wait_for(lock,
std::chrono::seconds(STREAM_SESSION_GC_INTERVAL_SECONDS),
[this] { return !running.load(std::memory_order_acquire); });
}
if (!running.load(std::memory_order_acquire)) {
return;
[this] { return !running; });
if (!running) {
return;
}
}
int64_t cutoff = now_seconds() - STREAM_SESSION_TTL_SECONDS;
std::vector<stream_session_ptr> to_drop;
@ -301,10 +402,19 @@ void stream_session_manager::gc_loop() {
}
}
// process wide manager, lifecycle controlled by llama-server main() via start_gc/stop_gc
stream_session_manager g_stream_sessions;
// stream_pipe
// stream_pipe ---------------------------------------------------------------------------------
// consumer end: read-only replay of the ring buffer, the destructor does not finalize the session
struct stream_pipe_consumer : stream_pipe {
stream_read_status read(size_t & offset,
const std::function<bool(const char *, size_t)> & sink,
const std::function<bool()> & should_stop);
static std::shared_ptr<stream_pipe_consumer> create(stream_session_ptr session);
private:
explicit stream_pipe_consumer(stream_session_ptr session);
};
stream_pipe::stream_pipe(stream_session_ptr session)
: session_(std::move(session)) {
@ -408,12 +518,10 @@ static server_http_res_ptr make_error_response(int status, const std::string & m
return res;
}
server_http_context::handler_t make_stream_get_handler() {
server_http_context::handler_t server_stream_make_get_handler() {
return [](const server_http_req & req) -> server_http_res_ptr {
// GET /v1/stream/<conv_id>?from=N replays the SSE bytes already buffered for the
// session, blocks for more bytes when the session is still running, returns when
// the session is finalized. the body is streamed back as text/event-stream so the
// browser EventSource can attach to it like a fresh request
// GET /v1/stream/<conv_id>?from=N replays buffered SSE bytes then blocks for live
// bytes until the session finalizes, streamed as text/event-stream for EventSource
std::string conv_id = req.get_param("conv_id");
if (conv_id.empty()) {
return make_error_response(400, "Missing conversation id in path", ERROR_TYPE_INVALID_REQUEST);
@ -459,11 +567,10 @@ server_http_context::handler_t make_stream_get_handler() {
};
}
server_http_context::handler_t make_streams_lookup_handler() {
server_http_context::handler_t server_stream_make_lookup_handler() {
return [](const server_http_req & req) -> server_http_res_ptr {
// POST /v1/streams/lookup with body {"conversation_ids": ["X", "Y", ...]} returns the
// matching sessions, only for ids the caller already knows. each id matches the exact key
// and any "<id>::<model>" variant, so one lookup covers every per model session for a conv
// POST /v1/streams/lookup returns the matching sessions, only for ids the caller already
// knows. each id matches the exact key and any "<id>::<model>" per model variant
std::vector<std::string> requested;
try {
json body = json::parse(req.body);
@ -518,11 +625,10 @@ server_http_context::handler_t make_streams_lookup_handler() {
};
}
server_http_context::handler_t make_stream_delete_handler() {
server_http_context::handler_t server_stream_make_delete_handler() {
return [](const server_http_req & req) -> server_http_res_ptr {
// DELETE /v1/stream/<conv_id> is the explicit user Stop, cancels the producer hook
// wired by handle_completions_impl and evicts the buffer. idempotent, a session that
// already finalized or was never created returns 204 either way
// DELETE /v1/stream/<conv_id> is the explicit user Stop, cancels the producer and evicts
// the buffer. idempotent, returns 204 even if the session was already gone
std::string conv_id = req.get_param("conv_id");
if (conv_id.empty()) {
return make_error_response(400, "Missing conversation id in path", ERROR_TYPE_INVALID_REQUEST);
@ -536,7 +642,7 @@ server_http_context::handler_t make_stream_delete_handler() {
};
}
std::string stream_conv_id_from_headers(const std::map<std::string, std::string> & headers) {
std::string server_stream_conv_id_from_headers(const std::map<std::string, std::string> & headers) {
// case-insensitive scan for x-conversation-id
static constexpr char target[] = "x-conversation-id";
static constexpr size_t target_len = sizeof(target) - 1;
@ -555,8 +661,8 @@ std::string stream_conv_id_from_headers(const std::map<std::string, std::string>
return std::string();
}
void stream_session_attach_pipe(server_http_res & res, const std::map<std::string, std::string> & headers) {
std::string conversation_id = stream_conv_id_from_headers(headers);
void server_stream_session_attach_pipe(server_http_res & res, const std::map<std::string, std::string> & headers) {
std::string conversation_id = server_stream_conv_id_from_headers(headers);
SRV_TRC("conv_id=%s (empty=%d)\n", conversation_id.c_str(), conversation_id.empty() ? 1 : 0);
if (conversation_id.empty()) {
return;
@ -565,7 +671,7 @@ void stream_session_attach_pipe(server_http_res & res, const std::map<std::strin
res.spipe = stream_pipe_producer::create(session, res);
}
std::function<bool()> stream_aware_should_stop(server_http_res * res, std::function<bool()> fallback) {
std::function<bool()> server_stream_aware_should_stop(server_http_res * res, std::function<bool()> fallback) {
return [res, fallback = std::move(fallback)]() -> bool {
if (res->spipe) {
return res->spipe->is_cancelled();

151
tools/server/server-stream.h
View file

@ -3,81 +3,23 @@
#include "server-http.h"
#include <atomic>
#include <condition_variable>
#include <cstddef>
#include <cstdint>
#include <functional>
#include <memory>
#include <mutex>
#include <shared_mutex>
#include <string>
#include <thread>
#include <unordered_map>
#include <vector>
enum class stream_read_status {
OK,
OFFSET_LOST,
};
// streaming buffer for one generation, survives HTTP disconnect. the producer appends SSE bytes,
// readers drain from any offset via read_from. keyed by conversation_id, one conv = one live session
// streaming buffer for one generation, survives HTTP disconnect. the producer appends raw SSE
// bytes, readers drain from any offset via read_from and block until more bytes or finalize.
// keyed by conversation_id: one conv = at most one live session
struct stream_session {
std::string conversation_id;
int64_t started_ts; // unix seconds at construction, used by /v1/streams listing
stream_session(std::string conversation_id_, size_t max_bytes_);
stream_session(const stream_session &) = delete;
stream_session & operator=(const stream_session &) = delete;
// append raw bytes, drops from the front if the cap is reached.
// returns false if the session is already finalized
bool append(const char * data, size_t len);
// mark the session as complete, wakes all pending readers
void finalize();
// drain bytes from offset, calling sink for each chunk. blocks until more
// bytes arrive or finalize is called. returns OK on clean exit, OFFSET_LOST
// if offset falls below the dropped prefix
stream_read_status read_from(size_t offset,
const std::function<bool(const char *, size_t)> & sink,
const std::function<bool()> & should_stop);
bool is_done() const;
bool is_cancelled() const;
size_t total_size() const; // bytes that ever entered the session
size_t dropped_prefix() const; // bytes evicted from the front due to cap
int64_t completed_at() const; // 0 while alive, unix seconds after finalize
// attach the producer stop hook used to cancel its reader, pass an empty function to detach
void set_stop_producer(std::function<void()> fn);
// signal the producer to abort its inference asap via the stop hook, idempotent
void cancel();
private:
mutable std::mutex mu;
std::condition_variable cv;
std::vector<char> buffer;
size_t prefix_dropped;
size_t cap_bytes;
std::atomic<bool> done;
std::atomic<bool> cancelled;
std::atomic<int64_t> completed_ts;
std::function<void()> stop_producer; // protected by mu
};
struct stream_session;
using stream_session_ptr = std::shared_ptr<stream_session>;
// one end of a stream_session pipe. the base holds the session and the shared query, the
// producer and consumer ends derive from it. virtual dtor so each end runs its own teardown:
// base of the producer/consumer pipe ends. virtual dtor so each runs its own teardown:
// the producer finalizes the session, the consumer leaves it untouched
struct stream_pipe {
virtual ~stream_pipe() = default;
// true if the session was cancelled (e.g. via DELETE /v1/stream/<conv_id>)
bool is_cancelled() const;
protected:
@ -95,7 +37,6 @@ protected:
struct stream_pipe_producer : stream_pipe {
~stream_pipe_producer() override;
// append raw bytes to the session's ring buffer, returns false if already finalized
bool write(const char * data, size_t len);
// mark the natural end on the wire so a later close() is a no-op
@ -121,83 +62,21 @@ private:
server_http_res * res_ = nullptr;
};
// consumer end: read-only replay of the ring buffer, the destructor does not finalize the session
struct stream_pipe_consumer : stream_pipe {
// drain bytes from offset, calling sink for each available chunk. blocks until more data
// arrives or the session finalizes. should_stop is polled, returns OFFSET_LOST if offset
// fell below the dropped prefix
stream_read_status read(size_t & offset,
const std::function<bool(const char *, size_t)> & sink,
const std::function<bool()> & should_stop);
void server_stream_session_manager_start();
void server_stream_session_manager_stop();
static std::shared_ptr<stream_pipe_consumer> create(stream_session_ptr session);
// route handler factories wired under /v1/stream/* by server.cpp
server_http_context::handler_t server_stream_make_get_handler();
server_http_context::handler_t server_stream_make_lookup_handler();
server_http_context::handler_t server_stream_make_delete_handler();
private:
explicit stream_pipe_consumer(stream_session_ptr session);
};
// extract the X-Conversation-Id header value (case-insensitive), empty when absent
std::string server_stream_conv_id_from_headers(const std::map<std::string, std::string> & headers);
// owns all live sessions, runs a periodic GC to evict expired ones.
// the map is keyed by conversation_id, so the invariant "one conv = at most one
// live session" is enforced at the type level
class stream_session_manager {
public:
stream_session_manager();
~stream_session_manager();
stream_session_manager(const stream_session_manager &) = delete;
stream_session_manager & operator=(const stream_session_manager &) = delete;
// install a new session for this conversation, evicting and cancelling any previous one.
// the conversation_id must be non empty, the caller is responsible for that check.
// returns the new session
stream_session_ptr create_or_replace(const std::string & conversation_id);
// lookup, returns null if unknown or already evicted
stream_session_ptr get(const std::string & conversation_id);
// list every live or recently completed session, used by GET /v1/streams without filter
std::vector<stream_session_ptr> list_all() const;
// remove from the map and finalize, wakes any pending readers
void evict(const std::string & conversation_id);
// signal the producer to cancel asap then evict, used by the explicit user Stop path
void evict_and_cancel(const std::string & conversation_id);
void start_gc();
void stop_gc();
private:
void gc_loop();
mutable std::shared_mutex map_mu;
std::unordered_map<std::string, stream_session_ptr> sessions; // key: conversation_id
std::thread gc_thread;
std::atomic<bool> running;
std::mutex gc_wake_mu;
std::condition_variable gc_wake_cv;
};
// process wide manager, linked by both llama-server and llama-cli. llama-server main() drives
// start_gc/stop_gc, llama-cli leaves it idle. the dtor calls stop_gc() unconditionally so exit
// is safe whether or not the GC thread ran
extern stream_session_manager g_stream_sessions;
// route handler factories operating on g_stream_sessions, wired under /v1/stream/* by server.cpp.
// keeps the resumable stream surface confined to server-stream
server_http_context::handler_t make_stream_get_handler();
server_http_context::handler_t make_streams_lookup_handler();
server_http_context::handler_t make_stream_delete_handler();
// extract the X-Conversation-Id header value (case-insensitive), empty when absent. exposed so
// the router can track which child serves a forwarded POST
std::string stream_conv_id_from_headers(const std::map<std::string, std::string> & headers);
// on an X-Conversation-Id header, create or replace the session and attach a producer pipe to
// res. no-op when absent, called from the server_res_generator constructor
void stream_session_attach_pipe(server_http_res & res, const std::map<std::string, std::string> & headers);
// on an X-Conversation-Id header, create or replace the session and attach a producer pipe to res
void server_stream_session_attach_pipe(server_http_res & res, const std::map<std::string, std::string> & headers);
// should_stop closure that ignores peer disconnect when a pipe is attached, so only an explicit
// DELETE stops the producer and generation keeps flowing into the ring buffer. without a pipe it
// delegates to fallback, the legacy non-resumable flow
std::function<bool()> stream_aware_should_stop(server_http_res * res, std::function<bool()> fallback);
std::function<bool()> server_stream_aware_should_stop(server_http_res * res, std::function<bool()> fallback);

74
tools/server/server-task.cpp
View file

@ -730,6 +730,10 @@ json server_task_result_cmpl_final::to_json_oaicompat_resp_stream() {
}}
});
if (timings.prompt_n >= 0) {
server_sent_events.back().at("data").push_back({"timings", timings.to_json()});
}
return server_sent_events;
}
@ -1016,6 +1020,7 @@ void server_task_result_cmpl_partial::update(task_result_state & state) {
thinking_block_started = state.thinking_block_started;
text_block_started = state.text_block_started;
oai_resp_created = state.oai_resp_created;
oai_resp_id = state.oai_resp_id;
oai_resp_reasoning_id = state.oai_resp_reasoning_id;
oai_resp_message_id = state.oai_resp_message_id;
@ -1024,6 +1029,10 @@ void server_task_result_cmpl_partial::update(task_result_state & state) {
// track if the accumulated message has any reasoning content
anthropic_has_reasoning = !state.chat_msg.reasoning_content.empty();
if (res_type == TASK_RESPONSE_TYPE_OAI_RESP && !state.oai_resp_created && (is_progress || n_decoded == 1)) {
state.oai_resp_created = true;
}
// Pre-compute state updates based on diffs (for next chunk)
for (const common_chat_msg_diff & diff : oaicompat_msg_diffs) {
if (!diff.reasoning_content_delta.empty() && !state.thinking_block_started) {
@ -1181,7 +1190,7 @@ json server_task_result_cmpl_partial::to_json_oaicompat_chat() {
json server_task_result_cmpl_partial::to_json_oaicompat_resp() {
std::vector<json> events;
if (n_decoded == 1) {
if (!oai_resp_created) {
events.push_back(json {
{"event", "response.created"},
{"data", json {
@ -1204,6 +1213,18 @@ json server_task_result_cmpl_partial::to_json_oaicompat_resp() {
}},
}},
});
} else if (is_progress) {
events.push_back(json {
{"event", "response.in_progress"},
{"data", json {
{"type", "response.in_progress"},
{"response", json {
{"id", oai_resp_id},
{"object", "response"},
{"status", "in_progress"},
}},
}},
});
}
for (const common_chat_msg_diff & diff : oaicompat_msg_diffs) {
@ -1302,6 +1323,17 @@ json server_task_result_cmpl_partial::to_json_oaicompat_resp() {
});
}
}
if (!events.empty()) {
json & data = events.back().at("data");
if (timings.prompt_n >= 0) {
data.push_back({"timings", timings.to_json()});
}
if (is_progress) {
data.push_back({"prompt_progress", progress.to_json()});
}
}
return events;
}
@ -1631,7 +1663,22 @@ server_prompt * server_prompt_cache::alloc(const server_prompt & prompt, size_t
}
}
// next, remove any cached prompts that are fully contained in the current prompt
// calculate checkpoints size to see if it will fit with the prompt
size_t checkpoints_size = 0;
for (const auto & ckpt : prompt.checkpoints) {
checkpoints_size += ckpt.size();
}
const size_t state_size_new = state_size_tgt + state_size_dft + checkpoints_size;
// skip over-limit entries to avoid disturbing the cache
if (limit_size > 0 && state_size_new > limit_size) {
SRV_WRN(" - prompt state size %.3f MiB exceeds cache size limit %.3f MiB, skipping\n",
state_size_new / (1024.0 * 1024.0), limit_size / (1024.0 * 1024.0));
return nullptr;
}
// remove any cached prompts that are fully contained in the current prompt
for (auto it = states.begin(); it != states.end();) {
const int len = it->tokens.get_common_prefix(prompt.tokens);
@ -1644,6 +1691,16 @@ server_prompt * server_prompt_cache::alloc(const server_prompt & prompt, size_t
}
}
if (limit_size > 0) {
// make room before allocating the new vectors to avoid breaching the limit
while (!states.empty() && size() + state_size_new > limit_size) {
SRV_WRN(" - making room for prompt cache entry, removing oldest entry (size = %.3f MiB)\n",
states.front().size() / (1024.0 * 1024.0));
states.pop_front();
}
}
std::vector<uint8_t> state_data_tgt;
std::vector<uint8_t> state_data_dft;
@ -1752,12 +1809,7 @@ bool server_prompt_cache::load(server_prompt & prompt, const server_tokens & tok
void server_prompt_cache::update() {
if (limit_size > 0) {
// always keep at least one state, regardless of the limits
while (states.size() > 1 && size() > limit_size) {
if (states.empty()) {
break;
}
while (!states.empty() && size() > limit_size) {
SRV_WRN(" - cache size limit reached, removing oldest entry (size = %.3f MiB)\n", states.front().size() / (1024.0 * 1024.0));
states.pop_front();
@ -1771,11 +1823,7 @@ void server_prompt_cache::update() {
const size_t limit_tokens_cur = limit_size > 0 ? std::max<size_t>(limit_tokens, limit_size/size_per_token) : limit_tokens;
if (limit_tokens > 0) {
while (states.size() > 1 && n_tokens() > limit_tokens_cur) {
if (states.empty()) {
break;
}
while (!states.empty() && n_tokens() > limit_tokens_cur) {
SRV_WRN(" - cache token limit (%zu, est: %zu) reached, removing oldest entry (size = %.3f MiB)\n",
limit_tokens, limit_tokens_cur, states.front().size() / (1024.0 * 1024.0));

2
tools/server/server-task.h
View file

@ -117,6 +117,7 @@ struct task_result_state {
bool text_block_started = false;
// for OpenAI Responses streaming API
bool oai_resp_created = false;
const std::string oai_resp_id;
const std::string oai_resp_reasoning_id;
const std::string oai_resp_message_id;
@ -440,6 +441,7 @@ struct server_task_result_cmpl_partial : server_task_result {
bool text_block_started = false;
// for OpenAI Responses API
bool oai_resp_created = false;
std::string oai_resp_id;
std::string oai_resp_reasoning_id;
std::string oai_resp_message_id;

16
tools/server/server.cpp
View file

@ -85,7 +85,7 @@ int llama_server(int argc, char ** argv) {
// start the stream session manager GC right after common init, before any HTTP route can
// touch it. lifecycle is symmetric, stop_gc() runs in clean_up() before backend free
g_stream_sessions.start_gc();
server_stream_session_manager_start();
if (!common_params_parse(argc, argv, params, LLAMA_EXAMPLE_SERVER)) {
return 1;
@ -245,8 +245,8 @@ int llama_server(int argc, char ** argv) {
ctx_http.post("/slots/:id_slot", ex_wrapper(routes.post_slots));
// resumable streaming, the conversation_id is the session identity end to end. router and
// child wire different handlers under the same paths: a child binds the local g_stream_sessions
// backed factories, the router binds proxies that resolve the owning child through the
// child wire different handlers under the same paths: a child binds the local session
// factories, the router binds proxies that resolve the owning child through the
// conv_id -> model map
server_http_context::handler_t stream_get_h;
server_http_context::handler_t streams_lookup_h;
@ -256,9 +256,9 @@ int llama_server(int argc, char ** argv) {
streams_lookup_h = models_routes->router_streams_lookup;
stream_delete_h = models_routes->router_stream_delete;
} else {
stream_get_h = make_stream_get_handler();
streams_lookup_h = make_streams_lookup_handler();
stream_delete_h = make_stream_delete_handler();
stream_get_h = server_stream_make_get_handler();
streams_lookup_h = server_stream_make_lookup_handler();
stream_delete_h = server_stream_make_delete_handler();
}
ctx_http.get ("/v1/stream/:conv_id", ex_wrapper(stream_get_h));
// POST /v1/streams/lookup with body {"conversation_ids": [...]}. you can only ask for ids
@ -343,7 +343,7 @@ int llama_server(int argc, char ** argv) {
clean_up = [&models_routes]() {
SRV_INF("%s: cleaning up before exit...\n", __func__);
// stop the session GC first, it finalizes live sessions and wakes pending readers
g_stream_sessions.stop_gc();
server_stream_session_manager_stop();
if (models_routes.has_value()) {
models_routes->stopping.store(true); // maybe redundant, but just to be safe
models_routes->models.unload_all();
@ -371,7 +371,7 @@ int llama_server(int argc, char ** argv) {
clean_up = [&ctx_http, &ctx_server]() {
SRV_INF("%s: cleaning up before exit...\n", __func__);
// stop the session GC first, it finalizes live sessions and wakes pending readers
g_stream_sessions.stop_gc();
server_stream_session_manager_stop();
ctx_http.stop();
ctx_server.terminate();
llama_backend_free();

41
tools/server/tests/unit/test_compat_oai_responses.py
View file

@ -71,3 +71,44 @@ def test_responses_stream_with_openai_library():
assert r.response.output[0].id.startswith("msg_")
assert gathered_text == r.response.output_text
assert match_regex("(Suddenly)+", r.response.output_text)
def test_responses_stream_with_llama_telemetry():
global server
server.n_ctx = 256
server.n_batch = 32
server.n_slots = 1
server.start()
saw_progress = False
saw_delta_timings = False
completed = None
res = server.make_stream_request("POST", "/responses", data={
"input": "This is a test" * 10,
"max_output_tokens": 8,
"temperature": 0.8,
"stream": True,
"timings_per_token": True,
"return_progress": True,
})
for data in res:
if "prompt_progress" in data:
assert data["type"] == "response.in_progress"
assert data["prompt_progress"]["total"] > 0
assert data["prompt_progress"]["processed"] >= data["prompt_progress"]["cache"]
saw_progress = True
if "timings" in data:
assert "prompt_per_second" in data["timings"]
assert "predicted_per_second" in data["timings"]
if data["type"] == "response.output_text.delta":
saw_delta_timings = True
if data["type"] == "response.completed":
completed = data
assert saw_progress
assert saw_delta_timings
assert completed is not None
assert "usage" in completed["response"]
assert "timings" in completed

8
tools/ui/package-lock.json generated
View file

@ -11,7 +11,7 @@
"@chromatic-com/storybook": "5.0.0",
"@eslint/compat": "1.4.1",
"@eslint/js": "9.39.2",
"@internationalized/date": "3.10.1",
"@internationalized/date": "3.12.2",
"@lucide/svelte": "0.515.0",
"@modelcontextprotocol/sdk": "1.26.0",
"@playwright/test": "1.56.1",
@ -2981,9 +2981,9 @@
}
},
"node_modules/@internationalized/date": {
"version": "3.10.1",
"resolved": "https://registry.npmjs.org/@internationalized/date/-/date-3.10.1.tgz",
"integrity": "sha512-oJrXtQiAXLvT9clCf1K4kxp3eKsQhIaZqxEyowkBcsvZDdZkbWrVmnGknxs5flTD0VGsxrxKgBCZty1EzoiMzA==",
"version": "3.12.2",
"resolved": "https://registry.npmjs.org/@internationalized/date/-/date-3.12.2.tgz",
"integrity": "sha512-FY1Y+H64NDs+HAF6omlnWxm3mEpfgaCSWtL5l551ZZfImA+kGjPFgrnJrGjH6lfmLL0g8Z/mBu1R3kufeCp6Jw==",
"dev": true,
"license": "Apache-2.0",
"dependencies": {

2
tools/ui/package.json
View file

@ -30,7 +30,7 @@
"@chromatic-com/storybook": "5.0.0",
"@eslint/compat": "1.4.1",
"@eslint/js": "9.39.2",
"@internationalized/date": "3.10.1",
"@internationalized/date": "3.12.2",
"@lucide/svelte": "0.515.0",
"@modelcontextprotocol/sdk": "1.26.0",
"@playwright/test": "1.56.1",

9
tools/ui/src/lib/components/app/chat/ChatForm/ChatFormActions/ChatFormActionAdd/ChatFormActionAddDropdown.svelte
View file

@ -11,7 +11,8 @@
} from '$lib/constants';
import {
ChatFormActionAddToolsSubmenu,
ChatFormActionAddMcpServersSubmenu
ChatFormActionAddMcpServersSubmenu,
ChatFormActionAddReasoningSubmenu
} from '$lib/components/app';
import { useAttachmentMenu } from '$lib/hooks/use-attachment-menu.svelte';
@ -92,7 +93,11 @@
</Tooltip.Content>
</Tooltip.Root>
<DropdownMenu.Content align="start" class="w-48">
<DropdownMenu.Content align="start" class="w-52">
<ChatFormActionAddReasoningSubmenu />
<DropdownMenu.Separator />
<DropdownMenu.Sub>
<DropdownMenu.SubTrigger class="flex cursor-pointer items-center gap-2">
<File class="h-4 w-4" />

75
tools/ui/src/lib/components/app/chat/ChatForm/ChatFormActions/ChatFormReasoningToggle.svelte → tools/ui/src/lib/components/app/chat/ChatForm/ChatFormActions/ChatFormActionAdd/ChatFormActionAddReasoningSubmenu.svelte
View file

@ -2,7 +2,7 @@
import { Lightbulb, LightbulbOff, Check, Info } from '@lucide/svelte';
import * as DropdownMenu from '$lib/components/ui/dropdown-menu';
import * as Tooltip from '$lib/components/ui/tooltip';
import { ReasoningEffort, MessageRole } from '$lib/enums';
import { ReasoningEffort } from '$lib/enums';
import { REASONING_EFFORT_TOKENS } from '$lib/constants/reasoning-effort-tokens';
import { REASONING_EFFORT_LEVELS } from '$lib/constants/reasoning-effort';
import type { ReasoningEffortLevel } from '$lib/types';
@ -18,31 +18,23 @@
import { isRouterMode } from '$lib/stores/server.svelte';
import type { DatabaseMessage } from '$lib/types/database';
let thinkingEnabled = $derived(conversationsStore.getThinkingEnabled());
let currentEffort = $derived(conversationsStore.getReasoningEffort());
let isOff = $derived(!thinkingEnabled);
let tooltipText = $derived(thinkingEnabled ? `${currentEffort} Reasoning` : 'Disabled Reasoning');
let subOpen = $state(false);
// Get conversation model from message history
let conversationModel = $derived(
chatStore.getConversationModel(activeMessages() as DatabaseMessage[])
);
// Fallback: if model props aren't available, check if any assistant messages
// for this model in the active conversation have reasoning content.
let modelSupportsThinkingFromMessages = $derived.by(() => {
const modelId = isRouterMode() ? modelsStore.selectedModelName || conversationModel : null;
if (!modelId) return false;
const messages = conversationsStore.activeMessages;
return messages.some(
(m: DatabaseMessage) =>
m.role === MessageRole.ASSISTANT && m.model === modelId && !!m.reasoningContent
(m) => m.role === 'assistant' && m.model === modelId && !!m.reasoningContent
);
});
// Check if model supports thinking. Primary: chat template from /props.
// Fallback: message history (reasoning content in assistant messages).
let modelSupportsThinking = $derived.by(() => {
loadedModelIds();
propsCacheVersion();
@ -52,15 +44,15 @@
return checkModelSupportsThinking(modelId ?? '') || modelSupportsThinkingFromMessages;
}
// In non-router mode, use the built-in supportsThinking
return supportsThinking() || modelSupportsThinkingFromMessages;
});
// Check if current item is selected
let thinkingEnabled = $derived(conversationsStore.getThinkingEnabled());
let currentEffort = $derived(conversationsStore.getReasoningEffort());
let isOff = $derived(!thinkingEnabled);
function isSelected(item: ReasoningEffortLevel): boolean {
if (item.isOff) {
return isOff;
}
if (item.isOff) return isOff;
return thinkingEnabled && currentEffort === item.value;
}
@ -76,39 +68,30 @@
</script>
{#if modelSupportsThinking}
<DropdownMenu.Root bind:open={subOpen}>
<Tooltip.Root>
<Tooltip.Trigger>
<DropdownMenu.Trigger
class={[
'flex h-6 w-6 cursor-pointer items-center justify-center rounded-full p-0 transition-colors focus:outline-none focus-visible:ring-2 focus-visible:ring-ring focus-visible:ring-offset-2',
thinkingEnabled ? 'bg-amber-400/10 hover:bg-amber-400/20' : 'bg-muted'
]}
aria-label={`${tooltipText}. Click to configure.`}
>
{#if thinkingEnabled}
<Lightbulb class="h-3 w-3 text-amber-400" />
{:else}
<LightbulbOff class="h-3 w-3 text-muted-foreground" />
{/if}
</DropdownMenu.Trigger>
</Tooltip.Trigger>
<DropdownMenu.Sub bind:open={subOpen}>
<DropdownMenu.SubTrigger class="flex cursor-pointer items-center gap-2">
{#if thinkingEnabled}
<Lightbulb class="h-4 w-4 shrink-0 text-amber-400" />
{:else}
<LightbulbOff class="h-4 w-4 shrink-0 text-muted-foreground" />
{/if}
<Tooltip.Content>
<p class="capitalize">{tooltipText}</p>
</Tooltip.Content>
</Tooltip.Root>
<span class="text-sm inline-flex gap-2 {!thinkingEnabled ? 'text-muted-foreground' : ''}">
Reasoning
<DropdownMenu.Content
align="start"
class="w-60 rounded-xl bg-popover p-3 text-popover-foreground shadow-md outline-none"
<span class="capitalize text-muted-foreground">
{thinkingEnabled ? currentEffort : 'off'}
</span>
</span>
</DropdownMenu.SubTrigger>
<DropdownMenu.SubContent
class="w-60 bg-popover p-1.5 text-popover-foreground shadow-md outline-none"
>
<div class="mb-2 px-2.5 text-sm font-medium">Reasoning effort</div>
{#each REASONING_EFFORT_LEVELS as level (level.value)}
<button
type="button"
class="flex w-full cursor-pointer items-center gap-2 rounded-lg px-2.5 py-2 text-left text-sm transition-colors hover:bg-accent"
class="flex w-full cursor-pointer items-center gap-3 rounded-md px-2 py-1.75 text-left text-sm transition-colors hover:bg-accent"
class:bg-accent={isSelected(level)}
onclick={() => handleSelection(level)}
>
@ -140,6 +123,6 @@
{/if}
</button>
{/each}
</DropdownMenu.Content>
</DropdownMenu.Root>
</DropdownMenu.SubContent>
</DropdownMenu.Sub>
{/if}

50
tools/ui/src/lib/components/app/chat/ChatForm/ChatFormActions/ChatFormActions.svelte
View file

@ -7,14 +7,20 @@
ChatFormActionModels,
ChatFormActionRecord,
ChatFormActionSubmit,
ChatFormReasoningToggle
ChatFormContextGauge
} from '$lib/components/app';
import { FileTypeCategory } from '$lib/enums';
import { FileTypeCategory, MessageRole } from '$lib/enums';
import { mcpStore } from '$lib/stores/mcp.svelte';
import { config } from '$lib/stores/settings.svelte';
import { conversationsStore } from '$lib/stores/conversations.svelte';
import { activeMessages, conversationsStore } from '$lib/stores/conversations.svelte';
import {
activeProcessingState,
isChatStreaming,
isLoading as chatIsLoading
} from '$lib/stores/chat.svelte';
import { getFileTypeCategory } from '$lib/utils';
import { goto } from '$app/navigation';
import { page } from '$app/state';
import { ROUTES } from '$lib/constants/routes';
interface Props {
@ -93,6 +99,36 @@
let activeMessage = $derived(
conversationsStore.activeMessages[conversationsStore.activeMessages.length - 1]
);
let hasProcessedTokens = $derived.by(() => {
if (!page.params.id) return false;
const messages = activeMessages() as DatabaseMessage[];
let totalHistoricalTokens = 0;
for (const m of messages) {
if (m.role !== MessageRole.ASSISTANT) continue;
const timings = m.timings;
if (!timings) continue;
const agenticLlm = timings.agentic?.llm;
if (agenticLlm?.prompt_n != null || agenticLlm?.predicted_n != null) {
totalHistoricalTokens += (agenticLlm?.prompt_n ?? 0) + (agenticLlm?.predicted_n ?? 0);
} else {
totalHistoricalTokens += (timings.prompt_n ?? 0) + (timings.predicted_n ?? 0);
}
}
if (totalHistoricalTokens > 0) return true;
if (!chatIsLoading() && !isChatStreaming()) return false;
const processingState = activeProcessingState();
if (!processingState) return false;
const livePromptTokens = Math.max(
processingState.promptTokens ?? 0,
processingState.promptProgress?.processed ?? 0
);
const liveOutputTokens = processingState.outputTokensUsed ?? 0;
return livePromptTokens > 0 || liveOutputTokens > 0;
});
</script>
<div
@ -100,7 +136,7 @@
style="container-type: inline-size"
>
{#if showAddButton}
<div class="mr-auto flex items-center gap-3">
<div class="mr-auto flex items-center gap-2">
<ChatFormActionsAdd
{disabled}
{hasAudioModality}
@ -117,8 +153,10 @@
</div>
{/if}
<div class="flex items-center gap-2">
<ChatFormReasoningToggle />
<div class="flex items-center gap-1.5">
{#if hasProcessedTokens}
<ChatFormContextGauge />
{/if}
{#if showModelSelector}
<ChatFormActionModels

23
tools/ui/src/lib/components/app/chat/ChatForm/ChatFormActions/ChatFormReasoningEffortSubmenu.svelte
View file

@ -6,6 +6,7 @@
import { REASONING_EFFORT_TOKENS } from '$lib/constants/reasoning-effort-tokens';
import { REASONING_EFFORT_LEVELS } from '$lib/constants/reasoning-effort';
import type { ReasoningEffortLevel } from '$lib/types';
import { DIALOG_SUBMENU_CONTENT } from '$lib/constants/css-classes';
import {
modelsStore,
checkModelSupportsThinking,
@ -71,9 +72,7 @@
{#if modelSupportsThinking}
<DropdownMenu.Sub bind:open={subOpen}>
<DropdownMenu.SubTrigger
class="flex cursor-pointer items-center gap-2 rounded-md px-2.5 py-1.5 text-sm transition-colors outline-none hover:bg-accent focus:bg-accent"
>
<DropdownMenu.SubTrigger class="flex cursor-pointer items-center gap-2">
{#if thinkingEnabled}
<Lightbulb class="h-4 w-4 shrink-0 text-amber-400" />
{:else}
@ -89,23 +88,15 @@
{/if}
</DropdownMenu.SubTrigger>
<DropdownMenu.SubContent
class="w-60 rounded-xl bg-popover p-3 text-popover-foreground shadow-md outline-none data-[side=bottom]:slide-in-from-top-2 data-[side=left]:slide-in-from-right-2 data-[side=right]:slide-in-from-left-2 data-[side=top]:slide-in-from-bottom-2 data-[state=closed]:animate-out data-[state=closed]:fade-out-0 data-[state=closed]:zoom-out-95 data-[state=open]:animate-in data-[state=open]:fade-in-0 data-[state=open]:zoom-in-95"
>
<DropdownMenu.SubContent class={DIALOG_SUBMENU_CONTENT}>
{#each REASONING_EFFORT_LEVELS as level (level.value)}
<button
type="button"
class="flex w-full cursor-pointer items-center gap-2 rounded-lg px-2.5 py-2 text-left text-sm transition-colors hover:bg-accent"
class="flex w-full cursor-pointer items-center gap-2"
class:bg-accent={isSelected(level)}
onclick={() => handleSelection(level)}
>
{#if isSelected(level)}
<Check class="h-4 w-4 shrink-0 text-foreground" />
{:else}
<div class="h-4 w-4 shrink-0"></div>
{/if}
<span class="flex-1">{level.label}</span>
<span class="flex-1 text-left">{level.label}</span>
{#if !level.isOff}
<span class="text-[11px] text-muted-foreground opacity-60">
@ -125,6 +116,10 @@
</Tooltip.Content>
</Tooltip.Root>
{/if}
{#if isSelected(level)}
<Check class="h-4 w-4 shrink-0 text-foreground" />
{/if}
</button>
{/each}
</DropdownMenu.SubContent>

108
tools/ui/src/lib/components/app/chat/ChatForm/ChatFormContextGauge/ChatFormContextGauge.svelte Normal file
View file

@ -0,0 +1,108 @@
<script lang="ts">
import { untrack } from 'svelte';
import * as HoverCard from '$lib/components/ui/hover-card';
import { activeConversation, activeMessages } from '$lib/stores/conversations.svelte';
import { chatStore, isChatStreaming, isLoading } from '$lib/stores/chat.svelte';
import { formatParameters } from '$lib/utils/formatters';
import { useContextGauge } from '$lib/hooks/use-context-gauge.svelte';
import ContextGaugeDial from './ContextGaugeDial.svelte';
import ContextGaugeDetails from './ContextGaugeDetails.svelte';
import ContextGaugeLoadModel from './ContextGaugeLoadModel.svelte';
import { colorLevelBgClass, colorLevelTextClass } from './context-gauge';
const gauge = useContextGauge();
$effect(() => {
const conv = activeConversation();
untrack(() => chatStore.setActiveProcessingConversation(conv?.id ?? null));
});
$effect(() => {
const conv = activeConversation();
const messages = activeMessages() as DatabaseMessage[];
if (!conv) return;
if (isLoading() || isChatStreaming()) return;
if (messages.length === 0) {
untrack(() => chatStore.clearProcessingState(conv.id));
return;
}
untrack(() => chatStore.restoreProcessingStateFromMessages(messages, conv.id));
});
$effect(() => {
gauge.startMonitoring();
});
const showProgressBar = $derived(
gauge.contextTotal !== null &&
gauge.contextTotal > 0 &&
(gauge.activeModelId !== null || gauge.isActiveModelLoaded)
);
</script>
<HoverCard.Root>
<HoverCard.Trigger class="flex h-5 w-5 cursor-default items-center justify-center">
<ContextGaugeDial percent={gauge.contextPercent} level={gauge.colorLevel} />
</HoverCard.Trigger>
<HoverCard.Content
side="bottom"
class="z-50 w-64 rounded-lg border border-border/50 bg-popover p-3 text-popover-foreground shadow-lg"
>
<div class="flex flex-col gap-2">
<div class="flex items-center gap-2">
<span class="font-medium">Context</span>
<span class="text-muted-foreground">·</span>
<span class="font-mono text-muted-foreground">
{formatParameters(gauge.contextUsed)}
/ {gauge.contextTotal !== null ? formatParameters(gauge.contextTotal) : '-'}
</span>
</div>
{#if gauge.activeModelId !== null && !gauge.isActiveModelLoaded}
<ContextGaugeLoadModel
modelId={gauge.activeModelId}
isLoading={gauge.isActiveModelLoading}
onLoad={gauge.loadModel}
/>
{:else if showProgressBar}
<div class="h-1.5 w-full overflow-hidden rounded-full bg-muted">
<div
class="h-full rounded-full transition-all duration-300 {colorLevelBgClass(
gauge.colorLevel
)}"
style="width: {gauge.contextPercent}%"
></div>
</div>
<div class="flex justify-between text-xs text-muted-foreground">
<span>
<span class={colorLevelTextClass(gauge.colorLevel)}>{gauge.contextPercent}%</span> used
</span>
<span>
{formatParameters((gauge.contextTotal ?? 0) - gauge.contextUsed)} remaining
</span>
</div>
{:else}
<div class="text-xs text-muted-foreground">No context info available</div>
{/if}
{#if gauge.hasAnyUsage}
<ContextGaugeDetails
currentRead={gauge.currentRead}
currentFresh={gauge.currentFresh}
currentCache={gauge.currentCache}
currentOutput={gauge.currentOutput}
kvTotal={gauge.kvTotal}
cumulativeRead={gauge.cumulativeRead}
cumulativeOutput={gauge.cumulativeOutput}
cumulativeCacheTotal={gauge.cumulativeCacheTotal}
averageTokensPerSecond={gauge.averageTokensPerSecond}
transientDetails={gauge.transientDetails}
/>
{/if}
</div>
</HoverCard.Content>
</HoverCard.Root>

20
tools/ui/src/lib/components/app/chat/ChatForm/ChatFormContextGauge/ContextGaugeDetailRow.svelte Normal file
View file

@ -0,0 +1,20 @@
<script lang="ts">
interface Props {
label: string;
value: string;
subtitle?: string;
}
let { label, value, subtitle }: Props = $props();
</script>
<div class="grid gap-1.5">
<div class="flex items-baseline justify-between">
<span class="text-muted-foreground">{label}</span>
<span class="font-mono text-muted-foreground">{value}</span>
</div>
{#if subtitle}
<div class="text-[10px] leading-tight text-muted-foreground/70">{subtitle}</div>
{/if}
</div>

122
tools/ui/src/lib/components/app/chat/ChatForm/ChatFormContextGauge/ContextGaugeDetails.svelte Normal file
View file

@ -0,0 +1,122 @@
<script lang="ts">
import { ChevronDown } from '@lucide/svelte';
import * as Collapsible from '$lib/components/ui/collapsible';
import { STATS_UNITS } from '$lib/constants';
import ContextGaugeDetailRow from './ContextGaugeDetailRow.svelte';
interface Props {
currentRead: number;
currentFresh: number;
currentCache: number;
currentOutput: number;
kvTotal: number;
cumulativeRead: number;
cumulativeOutput: number;
cumulativeCacheTotal: number;
averageTokensPerSecond: number | null;
transientDetails: string[];
}
let {
currentRead,
currentFresh,
currentCache,
currentOutput,
kvTotal,
cumulativeRead,
cumulativeOutput,
cumulativeCacheTotal,
averageTokensPerSecond,
transientDetails
}: Props = $props();
let open = $state(false);
const hasCumulative = $derived(cumulativeRead > 0 || cumulativeOutput > 0);
const hasCurrent = $derived(currentRead > 0 || currentOutput > 0);
</script>
<Collapsible.Root bind:open class="mt-3 border-t border-border/50 pt-4">
<Collapsible.Trigger
class="flex w-full cursor-pointer items-center gap-1 text-xs text-muted-foreground hover:text-foreground"
>
<span>Token usage details</span>
<ChevronDown class={'ml-auto h-3 w-3 transition-transform' + (open ? ' rotate-180' : '')} />
</Collapsible.Trigger>
<Collapsible.Content class="flex flex-col gap-4 text-xs pt-4">
{#if hasCumulative}
<div>
<h3 class="text-[11px] font-medium uppercase tracking-wide text-muted-foreground/70 mb-2">
Across all turns
</h3>
<div class="flex flex-col gap-2">
{#if cumulativeRead > 0}
<ContextGaugeDetailRow
label="Prompt tokens evaluated"
value={`${cumulativeRead.toLocaleString()} tok`}
subtitle={cumulativeCacheTotal > 0
? `${cumulativeCacheTotal.toLocaleString()} reused from KV cache`
: undefined}
/>
{/if}
{#if cumulativeOutput > 0}
<ContextGaugeDetailRow
label="Tokens generated"
value={`${cumulativeOutput.toLocaleString()} tok`}
/>
{/if}
</div>
</div>
{/if}
{#if hasCurrent}
<div>
<h3 class="text-[11px] font-medium uppercase tracking-wide text-muted-foreground/70 mb-2">
This turn · KV cache
</h3>
<div class="flex flex-col gap-2">
{#if currentRead > 0}
<ContextGaugeDetailRow
label="Prompt"
value={`${currentRead.toLocaleString()} tok`}
subtitle={currentCache > 0
? `${currentFresh.toLocaleString()} fresh + ${currentCache.toLocaleString()} cached`
: undefined}
/>
{/if}
{#if currentOutput > 0}
<ContextGaugeDetailRow
label="Generated"
value={`${currentOutput.toLocaleString()} tok`}
/>
{/if}
<div class="pt-1 mt-0.5 border-t border-border/30">
<div class="flex justify-between">
<span class="text-muted-foreground">KV cache total</span>
<span class="font-mono font-medium">{kvTotal.toLocaleString()} tok</span>
</div>
</div>
</div>
</div>
{/if}
{#if averageTokensPerSecond !== null}
<div class="pt-1.5 mt-1 border-t border-border/30">
<ContextGaugeDetailRow
label="Avg speed"
value={`${averageTokensPerSecond.toFixed(1)}${STATS_UNITS.TOKENS_PER_SECOND}`}
/>
</div>
{/if}
{#each transientDetails as detail (detail)}
<div class="font-mono text-muted-foreground">{detail}</div>
{/each}
</Collapsible.Content>
</Collapsible.Root>

43
tools/ui/src/lib/components/app/chat/ChatForm/ChatFormContextGauge/ContextGaugeDial.svelte Normal file
View file

@ -0,0 +1,43 @@
<script lang="ts">
import type { ColorLevel } from './context-gauge';
import { colorLevelTextClass } from './context-gauge';
interface Props {
percent: number | null;
level: ColorLevel;
size?: 'sm' | 'md';
}
let { percent, level, size = 'sm' }: Props = $props();
const RADIUS = 11;
const CIRCUMFERENCE = 2 * Math.PI * RADIUS;
const strokeLevelClass = $derived(colorLevelTextClass(level));
const dimensions = $derived(size === 'md' ? 'h-6 w-6' : 'h-5 w-5');
const strokeWidth = $derived(size === 'md' ? 4 : 3);
</script>
<svg viewBox="0 0 32 32" fill="none" class={dimensions}>
<circle
cx="16"
cy="16"
r={RADIUS}
stroke="currentColor"
stroke-opacity="0.1"
stroke-width={strokeWidth}
/>
<circle
cx="16"
cy="16"
r={RADIUS}
class="transition-colors duration-300 {strokeLevelClass}"
stroke="currentColor"
stroke-width={strokeWidth}
stroke-linecap="round"
stroke-dasharray={CIRCUMFERENCE}
stroke-dashoffset={percent !== null ? CIRCUMFERENCE * (1 - percent / 100) : CIRCUMFERENCE}
transform="rotate(-90 16 16)"
/>
</svg>

24
tools/ui/src/lib/components/app/chat/ChatForm/ChatFormContextGauge/ContextGaugeLoadModel.svelte Normal file
View file

@ -0,0 +1,24 @@
<script lang="ts">
import { Loader2 } from '@lucide/svelte';
import { Button } from '$lib/components/ui/button';
interface Props {
modelId: string | null;
isLoading: boolean;
onLoad: () => void;
}
let { modelId, isLoading, onLoad }: Props = $props();
</script>
{#if modelId !== null && !isLoading}
<div class="flex flex-col gap-2 border-t border-border/50 pt-2 text-xs text-muted-foreground">
<span>Available context size is only visible once the model is loaded.</span>
<Button size="sm" variant="secondary" class="self-start" onclick={onLoad}>Load model</Button>
</div>
{:else if isLoading}
<div class="flex items-center gap-2 border-t border-border/50 pt-2 text-xs text-muted-foreground">
<Loader2 class="h-3.5 w-3.5 animate-spin" />
<span>Loading model...</span>
</div>
{/if}

37
tools/ui/src/lib/components/app/chat/ChatForm/ChatFormContextGauge/context-gauge.ts Normal file
View file

@ -0,0 +1,37 @@
export type ColorLevel = 'ok' | 'warning' | 'critical' | 'neutral';
const WARNING_THRESHOLD = 80;
const CRITICAL_THRESHOLD = 95;
export function colorLevelFromPercent(percent: number | null): ColorLevel {
if (percent === null) return 'neutral';
if (percent >= CRITICAL_THRESHOLD) return 'critical';
if (percent >= WARNING_THRESHOLD) return 'warning';
return 'ok';
}
export function colorLevelTextClass(level: ColorLevel): string {
switch (level) {
case 'critical':
return 'text-red-400';
case 'warning':
return 'text-amber-400';
case 'ok':
return 'text-muted-foreground';
default:
return 'text-muted-foreground';
}
}
export function colorLevelBgClass(level: ColorLevel): string {
switch (level) {
case 'critical':
return 'bg-red-500';
case 'warning':
return 'bg-amber-500';
case 'ok':
return 'bg-green-500';
default:
return 'bg-muted';
}
}

6
tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessage.svelte
View file

@ -24,6 +24,8 @@
message: DatabaseMessage;
toolMessages?: DatabaseMessage[];
isLastAssistantMessage?: boolean;
isLastUserMessage?: boolean;
nextAssistantMessage?: DatabaseMessage | null;
siblingInfo?: ChatMessageSiblingInfo | null;
}
@ -32,6 +34,8 @@
message,
toolMessages = [],
isLastAssistantMessage = false,
isLastUserMessage = false,
nextAssistantMessage = null,
siblingInfo = null
}: Props = $props();
@ -359,7 +363,9 @@
<ChatMessageUser
class={className}
{deletionInfo}
{isLastUserMessage}
{message}
{nextAssistantMessage}
onConfirmDelete={handleConfirmDelete}
onCopy={handleCopy}
onDelete={handleDelete}

80
tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageAssistant/ChatMessageAssistant.svelte
View file

@ -11,11 +11,10 @@
import { useProcessingState } from '$lib/hooks/use-processing-state.svelte';
import { isLoading, isChatStreaming } from '$lib/stores/chat.svelte';
import { copyToClipboard, deriveAgenticSections, modelLoadProgressText } from '$lib/utils';
import { AgenticSectionType } from '$lib/enums';
import { AgenticSectionType, ChatMessageStatisticsMode } from '$lib/enums';
import { REASONING_TAGS } from '$lib/constants/agentic';
import { tick } from 'svelte';
import { fade } from 'svelte/transition';
import { MessageRole, ChatMessageStatsView } from '$lib/enums';
import { MessageRole } from '$lib/enums';
import { config } from '$lib/stores/settings.svelte';
import { isRouterMode } from '$lib/stores/server.svelte';
import { modelsStore } from '$lib/stores/models.svelte';
@ -122,62 +121,6 @@
return parts.join('\n\n\n');
});
let activeStatsView = $state<ChatMessageStatsView>(ChatMessageStatsView.GENERATION);
let statsContainerEl: HTMLDivElement | undefined = $state();
function getScrollParent(el: HTMLElement): HTMLElement | null {
let parent = el.parentElement;
while (parent) {
const style = getComputedStyle(parent);
if (/(auto|scroll)/.test(style.overflowY)) {
return parent;
}
parent = parent.parentElement;
}
return null;
}
async function handleStatsViewChange(view: ChatMessageStatsView) {
const el = statsContainerEl;
if (!el) {
activeStatsView = view;
return;
}
const scrollParent = getScrollParent(el);
if (!scrollParent) {
activeStatsView = view;
return;
}
const yBefore = el.getBoundingClientRect().top;
activeStatsView = view;
await tick();
const delta = el.getBoundingClientRect().top - yBefore;
if (delta !== 0) {
scrollParent.scrollTop += delta;
}
// Correct any drift after browser paint
requestAnimationFrame(() => {
const drift = el.getBoundingClientRect().top - yBefore;
if (Math.abs(drift) > 1) {
scrollParent.scrollTop += drift;
}
});
}
let highlightAgenticTurns = $derived(
isAgentic &&
(currentConfig.alwaysShowAgenticTurns || activeStatsView === ChatMessageStatsView.SUMMARY)
);
let displayedModel = $derived(message.model ?? null);
// model being switched to while it loads, so the selector bar tracks it
@ -291,7 +234,6 @@
{toolMessages}
isStreaming={isChatStreaming()}
{isLastAssistantMessage}
highlightTurns={highlightAgenticTurns}
/>
{/if}
{:else}
@ -315,10 +257,7 @@
<div class="info my-6 grid gap-4 tabular-nums">
{#if displayedModel}
<div
bind:this={statsContainerEl}
class="inline-flex flex-wrap items-start gap-2 text-xs text-muted-foreground"
>
<div class="inline-flex flex-wrap items-start gap-2 text-xs text-muted-foreground">
{#if isRouter}
<ModelsSelectorDropdown
currentModel={pendingModel ?? displayedModel}
@ -347,28 +286,25 @@
{#if currentConfig.showMessageStats && message.timings && message.timings.predicted_n && message.timings.predicted_ms}
{@const agentic = message.timings.agentic}
<ChatMessageStatistics
mode={ChatMessageStatisticsMode.GENERATION}
promptTokens={agentic ? agentic.llm.prompt_n : message.timings.prompt_n}
promptMs={agentic ? agentic.llm.prompt_ms : message.timings.prompt_ms}
predictedTokens={agentic ? agentic.llm.predicted_n : message.timings.predicted_n}
predictedMs={agentic ? agentic.llm.predicted_ms : message.timings.predicted_ms}
agenticTimings={agentic}
onActiveViewChange={handleStatsViewChange}
/>
{:else if isLoading() && currentConfig.showMessageStats}
{@const liveStats = processingState.getLiveProcessingStats()}
{@const genStats = processingState.getLiveGenerationStats()}
{@const promptProgress = processingState.processingState?.promptProgress}
{@const isStillProcessingPrompt =
promptProgress && promptProgress.processed < promptProgress.total}
{#if liveStats || genStats}
{#if genStats}
<ChatMessageStatistics
mode={ChatMessageStatisticsMode.GENERATION}
isLive
isProcessingPrompt={!!isStillProcessingPrompt}
promptTokens={liveStats?.tokensProcessed}
promptMs={liveStats?.timeMs}
predictedTokens={genStats?.tokensGenerated}
predictedMs={genStats?.timeMs}
predictedTokens={genStats.tokensGenerated}
predictedMs={genStats.timeMs}
/>
{/if}
{/if}

72
tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageUser/ChatMessageUser.svelte
View file

@ -2,10 +2,14 @@
import {
ChatMessageActionIcons,
ChatMessageEditForm,
ChatMessageStatistics,
ChatMessageUserBubble
} from '$lib/components/app/chat';
import { getMessageEditContext } from '$lib/contexts';
import { MessageRole } from '$lib/enums';
import { useProcessingState } from '$lib/hooks/use-processing-state.svelte';
import { isLoading } from '$lib/stores/chat.svelte';
import { MessageRole, ChatMessageStatisticsMode } from '$lib/enums';
import { config } from '$lib/stores/settings.svelte';
interface Props {
class?: string;
@ -17,6 +21,8 @@
assistantMessages: number;
messageTypes: string[];
} | null;
isLastUserMessage?: boolean;
nextAssistantMessage?: DatabaseMessage | null;
showDeleteDialog: boolean;
onEdit: () => void;
onDelete: () => void;
@ -32,6 +38,8 @@
message,
siblingInfo = null,
deletionInfo,
isLastUserMessage = false,
nextAssistantMessage = null,
showDeleteDialog,
onEdit,
onDelete,
@ -44,6 +52,37 @@
// Get contexts
const editCtx = getMessageEditContext();
const processingState = useProcessingState();
const currentConfig = $derived(config());
const isActivelyProcessing = $derived(isLastUserMessage && isLoading());
// For agentic turns, prefer the cumulative agentic.llm totals over per-call timings.
let storedReadingStats = $derived.by(() => {
const timings = nextAssistantMessage?.timings;
if (!timings?.prompt_n || !timings?.prompt_ms) return null;
const agentic = timings.agentic;
return {
promptTokens: agentic ? agentic.llm.prompt_n : timings.prompt_n,
promptMs: agentic ? agentic.llm.prompt_ms : timings.prompt_ms
};
});
let showStoredReadingStats = $derived(
Boolean(currentConfig.showMessageStats) && storedReadingStats !== null
);
let showLiveReadingStats = $derived(
Boolean(currentConfig.showMessageStats) && isActivelyProcessing && storedReadingStats === null
);
$effect(() => {
if (showLiveReadingStats) {
processingState.startMonitoring();
}
});
</script>
<div
@ -60,6 +99,37 @@
renderMarkdown={true}
/>
{#if showStoredReadingStats}
<!-- Reading stats sourced from the assistant message that followed this turn -->
<div class="info my-2 grid w-full justify-items-end gap-4 tabular-nums">
<div
class="inline-flex flex-wrap items-start justify-end gap-2 text-xs text-muted-foreground"
>
<ChatMessageStatistics
mode={ChatMessageStatisticsMode.READING}
promptTokens={storedReadingStats!.promptTokens}
promptMs={storedReadingStats!.promptMs}
/>
</div>
</div>
{:else if showLiveReadingStats}
{@const liveStats = processingState.getLiveProcessingStats()}
{#if liveStats}
<div class="info my-2 grid w-full justify-items-end gap-4 tabular-nums">
<div
class="inline-flex flex-wrap items-start justify-end gap-2 text-xs text-muted-foreground"
>
<ChatMessageStatistics
mode={ChatMessageStatisticsMode.READING}
isLive
promptTokens={liveStats.tokensProcessed}
promptMs={liveStats.timeMs}
/>
</div>
</div>
{/if}
{/if}
{#if message.timestamp}
<div class="max-w-[80%]">
<ChatMessageActionIcons

8
tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageUser/ChatMessageUserPending.svelte
View file

@ -2,7 +2,6 @@
import { ActionIcon, ChatMessageEditForm, ChatMessageUserBubble } from '$lib/components/app';
import { fadeInView } from '$lib/actions/fade-in-view.svelte';
import { ArrowUp, Edit, Trash2 } from '@lucide/svelte';
import { getProcessingInfoContext } from '$lib/contexts';
import { useMessageEditContext } from '$lib/hooks/use-message-edit-context.svelte';
interface Props {
@ -23,9 +22,6 @@
onDelete
}: Props = $props();
const processingInfoCtx = getProcessingInfoContext();
let showProcessingInfo = $derived(processingInfoCtx.showProcessingInfo);
const editCtx = useMessageEditContext({
getContent: () => content,
getExtras: () => extras,
@ -36,9 +32,7 @@
<div
use:fadeInView
aria-label="Pending user message"
class="group flex flex-col items-end gap-3 transition-opacity hover:opacity-80 md:gap-2 {className} sticky {showProcessingInfo
? 'bottom-44'
: 'bottom-32'}"
class="group flex flex-col items-end gap-3 transition-opacity hover:opacity-80 md:gap-2 {className} sticky bottom-32"
role="group"
>
{#if editCtx.isEditing}

46
tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessageAgenticContent.svelte
View file

@ -41,15 +41,13 @@
toolMessages?: DatabaseMessage[];
isStreaming?: boolean;
isLastAssistantMessage?: boolean;
highlightTurns?: boolean;
}
let {
message,
toolMessages = [],
isStreaming = false,
isLastAssistantMessage = false,
highlightTurns = false
isLastAssistantMessage = false
}: Props = $props();
let expandedStates: Record<number, boolean> = $state({});
@ -57,6 +55,7 @@
const showToolCallInProgress = $derived(config().showToolCallInProgress as boolean);
const showThoughtInProgress = $derived(config().showThoughtInProgress as boolean);
const renderThinkingAsMarkdown = $derived(config().renderThinkingAsMarkdown as boolean);
const showMessageStats = $derived(config().showMessageStats as boolean);
const hasReasoningError = $derived(
isLastAssistantMessage ? !!agenticLastError(message.convId) : false
@ -354,16 +353,17 @@
{/snippet}
<div class="agentic-content">
{#if highlightTurns && turnGroups.length > 1}
{#if turnGroups.length > 1}
{#each turnGroups as turn, turnIndex (turnIndex)}
{@const turnStats = message?.timings?.agentic?.perTurn?.[turnIndex]}
<div class="agentic-turn my-2 hover:bg-muted/80 dark:hover:bg-muted/30">
<span class="agentic-turn-label">Turn {turnIndex + 1}</span>
<div class="agentic-turn group/turn grid gap-3 mb-4">
{#each turn.sections as section, sIdx (turn.flatIndices[sIdx])}
{@render renderSection(section, turn.flatIndices[sIdx])}
{/each}
{#if turnStats}
<div class="turn-stats">
{#if turnStats && showMessageStats}
<div class="turn-stats transition-opacity duration-150">
<ChatMessageStatistics
promptTokens={turnStats.llm.prompt_n}
promptMs={turnStats.llm.prompt_ms}
@ -402,39 +402,21 @@
.agentic-content {
display: flex;
flex-direction: column;
gap: 0.5rem;
width: 100%;
max-width: 48rem;
gap: 1rem;
}
.agentic-content > :global(*),
.agentic-turn > :global(*) {
min-width: 0;
}
.agentic-text {
width: 100%;
}
.agentic-turn {
position: relative;
border: 1.5px dashed var(--muted-foreground);
border-radius: 0.75rem;
padding: 1rem;
transition: background 0.1s;
}
.agentic-turn-label {
position: absolute;
top: -1rem;
left: 0.75rem;
padding: 0 0.375rem;
background: var(--background);
font-size: 0.7rem;
font-weight: 500;
color: var(--muted-foreground);
text-transform: uppercase;
letter-spacing: 0.05em;
}
.turn-stats {
margin-top: 0.75rem;
padding-top: 0.5rem;
border-top: 1px solid hsl(var(--muted) / 0.5);
}
</style>

97
tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessageStatistics/ChatMessageStatistics.svelte
View file

@ -2,7 +2,7 @@
import { Clock, Gauge, WholeWord, BookOpenText, Sparkles, Wrench, Layers } from '@lucide/svelte';
import { ChatMessageStatisticsBadge } from '$lib/components/app';
import * as Tooltip from '$lib/components/ui/tooltip';
import { ChatMessageStatsView } from '$lib/enums';
import { ChatMessageStatsView, ChatMessageStatisticsMode } from '$lib/enums';
import type { ChatMessageAgenticTimings } from '$lib/types/chat';
import { formatPerformanceTime } from '$lib/utils';
import { MS_PER_SECOND, DEFAULT_PERFORMANCE_TIME } from '$lib/constants';
@ -19,6 +19,7 @@
agenticTimings?: ChatMessageAgenticTimings;
onActiveViewChange?: (view: ChatMessageStatsView) => void;
hideSummary?: boolean;
mode?: ChatMessageStatisticsMode;
}
let {
@ -31,19 +32,30 @@
initialView = ChatMessageStatsView.GENERATION,
agenticTimings,
onActiveViewChange,
hideSummary = false
hideSummary = false,
mode = ChatMessageStatisticsMode.SWITCHABLE
}: Props = $props();
let activeView: ChatMessageStatsView = $derived(initialView);
let isSwitchable = $derived(mode === ChatMessageStatisticsMode.SWITCHABLE);
let activeView: ChatMessageStatsView = $derived(
mode === ChatMessageStatisticsMode.READING
? ChatMessageStatsView.READING
: mode === ChatMessageStatisticsMode.GENERATION
? ChatMessageStatsView.GENERATION
: initialView
);
let hasAutoSwitchedToGeneration = $state(false);
$effect(() => {
onActiveViewChange?.(activeView);
if (isSwitchable) {
onActiveViewChange?.(activeView);
}
});
// In live mode: auto-switch to GENERATION tab when prompt processing completes
$effect(() => {
if (isLive) {
if (isLive && isSwitchable) {
// Auto-switch to generation tab only when prompt processing is done (once)
if (
!hasAutoSwitchedToGeneration &&
@ -91,8 +103,7 @@
formattedPromptTime !== undefined
);
// In live mode, generation tab is disabled until we have generation stats
let isGenerationDisabled = $derived(isLive && !hasGenerationStats);
let isGenerationDisabled = $derived(isLive && isSwitchable && !hasGenerationStats);
let hasAgenticStats = $derived(agenticTimings !== undefined && agenticTimings.toolCallsCount > 0);
@ -153,44 +164,44 @@
{/snippet}
<div class="inline-flex items-center text-xs text-muted-foreground">
<div class="inline-flex items-center rounded-sm bg-muted-foreground/15 p-0.5">
{#if hasPromptStats || isLive}
{@render viewButton({
view: ChatMessageStatsView.READING,
icon: BookOpenText,
label: 'Reading',
tooltipText: 'Reading (prompt processing)'
})}
{/if}
{@render viewButton({
view: ChatMessageStatsView.GENERATION,
icon: Sparkles,
label: 'Generation',
tooltipText: isGenerationDisabled
? 'Generation (waiting for tokens...)'
: 'Generation (token output)',
disabled: isGenerationDisabled
})}
{#if hasAgenticStats}
{@render viewButton({
view: ChatMessageStatsView.TOOLS,
icon: Wrench,
label: 'Tools',
tooltipText: 'Tool calls'
})}
{#if !hideSummary}
{#if isSwitchable}
<div class="inline-flex items-center rounded-sm bg-muted-foreground/15 p-0.5">
{#if hasPromptStats || isLive}
{@render viewButton({
view: ChatMessageStatsView.SUMMARY,
icon: Layers,
label: 'Summary',
tooltipText: 'Agentic summary'
view: ChatMessageStatsView.READING,
icon: BookOpenText,
label: 'Reading',
tooltipText: 'Processing'
})}
{/if}
{/if}
</div>
{@render viewButton({
view: ChatMessageStatsView.GENERATION,
icon: Sparkles,
label: 'Generation',
tooltipText: isGenerationDisabled ? 'Waiting for tokens...' : 'Generation',
disabled: isGenerationDisabled
})}
{#if hasAgenticStats}
{@render viewButton({
view: ChatMessageStatsView.TOOLS,
icon: Wrench,
label: 'Tools',
tooltipText: 'Tool calls'
})}
{#if !hideSummary}
{@render viewButton({
view: ChatMessageStatsView.SUMMARY,
icon: Layers,
label: 'Summary',
tooltipText: 'Agentic summary'
})}
{/if}
{/if}
</div>
{/if}
<div class="flex items-center gap-1 px-2">
{#if activeView === ChatMessageStatsView.GENERATION && hasGenerationStats}
@ -256,7 +267,7 @@
value={formattedAgenticTotalTime}
tooltipLabel="Total time (LLM + tools)"
/>
{:else if hasPromptStats}
{:else if hasPromptStats && (mode === ChatMessageStatisticsMode.READING || isSwitchable)}
<ChatMessageStatisticsBadge
class="bg-transparent"
icon={WholeWord}

26
tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessages.svelte
View file

@ -186,6 +186,8 @@
message: DatabaseMessage;
toolMessages: DatabaseMessage[];
isLastAssistantMessage: boolean;
isLastUserMessage: boolean;
nextAssistantMessage: DatabaseMessage | null;
siblingInfo: ChatMessageSiblingInfo;
}> = [];
@ -236,18 +238,36 @@
message: msg,
toolMessages,
isLastAssistantMessage: false,
isLastUserMessage: false,
nextAssistantMessage: null,
siblingInfo
});
}
// Mark the last assistant message
let lastAssistantIdx = -1;
for (let i = result.length - 1; i >= 0; i--) {
if (result[i].message.role === MessageRole.ASSISTANT) {
result[i].isLastAssistantMessage = true;
lastAssistantIdx = i;
break;
}
}
if (lastAssistantIdx > 0 && result[lastAssistantIdx - 1].message.role === MessageRole.USER) {
result[lastAssistantIdx - 1].isLastUserMessage = true;
}
for (let i = 0; i < result.length; i++) {
if (result[i].message.role !== MessageRole.USER) continue;
for (let j = i + 1; j < result.length; j++) {
if (result[j].message.role === MessageRole.ASSISTANT) {
result[i].nextAssistantMessage = result[j].message;
break;
}
}
}
return result;
});
</script>
@ -257,12 +277,14 @@
{isVisible ? 'opacity-100' : 'opacity-0'}
{previousRouteId === '/(chat)/chat/[id]' ? '' : 'delay-300'}"
>
{#each displayMessages as { message, toolMessages, isLastAssistantMessage, siblingInfo } (message.id)}
{#each displayMessages as { message, toolMessages, isLastAssistantMessage, isLastUserMessage, nextAssistantMessage, siblingInfo } (message.id)}
<ChatMessage
class="mx-auto mt-12 w-full max-w-3xl"
{message}
{toolMessages}
{isLastAssistantMessage}
{isLastUserMessage}
{nextAssistantMessage}
{siblingInfo}
/>
{/each}

20
tools/ui/src/lib/components/app/chat/ChatScreen/ChatScreen.svelte
View file

@ -4,12 +4,10 @@
ChatScreenForm,
ChatMessages,
ChatScreenDragOverlay,
ChatScreenProcessingInfo,
ChatScreenStreamResumeStatus,
ServerLoadingSplash,
ChatScreenServerError
} from '$lib/components/app';
import { setProcessingInfoContext } from '$lib/contexts';
import { createAutoScrollController } from '$lib/hooks/use-auto-scroll.svelte';
import { useChatScreenActiveModel } from '$lib/hooks/use-chat-screen-active-model.svelte';
import { useChatScreenDragAndDrop } from '$lib/hooks/use-chat-screen-drag-and-drop.svelte';
@ -23,8 +21,7 @@
errorDialog,
isLoading,
isChatStreaming,
isEditing,
activeProcessingState
isEditing
} from '$lib/stores/chat.svelte';
import {
conversationsStore,
@ -42,12 +39,6 @@
let { showCenteredEmpty = false } = $props();
setProcessingInfoContext({
get showProcessingInfo() {
return showProcessingInfo;
}
});
let disableAutoScroll = $derived(Boolean(config().disableAutoScroll) || isMobile.current);
let isMobileUserScrolledUp = $state(false);
let mobileScrollDownHint = $state(false);
@ -63,11 +54,6 @@
let isServerLoading = $derived(serverLoading());
let hasPropsError = $derived(!!serverError());
let isCurrentConversationLoading = $derived(isLoading() || isChatStreaming());
let showProcessingInfo = $derived(
isCurrentConversationLoading ||
(config().keepStatsVisible && !!page.params.id) ||
activeProcessingState() !== null
);
let chatFormBottomPosition = $derived.by(() => {
if (!isMobile.current) return '1rem';
if (device.isStandalone) return '1.5rem';
@ -298,10 +284,6 @@
}}
/>
{/if}
{#if showProcessingInfo}
<ChatScreenProcessingInfo />
{/if}
</div>
<ChatScreenForm

127
tools/ui/src/lib/components/app/chat/ChatScreen/ChatScreenProcessingInfo.svelte
View file

@ -1,127 +0,0 @@
<script lang="ts">
import { untrack } from 'svelte';
import { PROCESSING_INFO_TIMEOUT } from '$lib/constants';
import { useProcessingState } from '$lib/hooks/use-processing-state.svelte';
import { chatStore, isLoading, isChatStreaming } from '$lib/stores/chat.svelte';
import { activeMessages, activeConversation } from '$lib/stores/conversations.svelte';
import { config } from '$lib/stores/settings.svelte';
const processingState = useProcessingState();
let isCurrentConversationLoading = $derived(isLoading());
let isStreaming = $derived(isChatStreaming());
let processingDetails = $derived(processingState.getTechnicalDetails());
let processingVisible = $derived(processingDetails.length > 0);
let { onVisibilityChange }: { onVisibilityChange?: (visible: boolean) => void } = $props();
$effect(() => {
onVisibilityChange?.(processingVisible);
});
$effect(() => {
const conversation = activeConversation();
untrack(() => chatStore.setActiveProcessingConversation(conversation?.id ?? null));
});
$effect(() => {
const keepStatsVisible = config().keepStatsVisible;
const shouldMonitor = keepStatsVisible || isCurrentConversationLoading || isStreaming;
if (shouldMonitor) {
processingState.startMonitoring();
}
if (!isCurrentConversationLoading && !isStreaming && !keepStatsVisible) {
const timeout = setTimeout(() => {
if (!config().keepStatsVisible && !isChatStreaming()) {
processingState.stopMonitoring();
}
}, PROCESSING_INFO_TIMEOUT);
return () => clearTimeout(timeout);
}
});
$effect(() => {
const conversation = activeConversation();
const messages = activeMessages() as DatabaseMessage[];
const keepStatsVisible = config().keepStatsVisible;
if (keepStatsVisible && conversation) {
if (messages.length === 0) {
untrack(() => chatStore.clearProcessingState(conversation.id));
return;
}
if (!isCurrentConversationLoading && !isStreaming) {
untrack(() => chatStore.restoreProcessingStateFromMessages(messages, conversation.id));
}
}
});
</script>
<div
class={[
'chat-processing-info-container pointer-events-none relative w-full hidden md:block',
processingVisible && 'visible'
]}
>
<div class="chat-processing-info-content absolute bottom-4 left-1/2 -translate-x-1/2">
{#each processingDetails as detail (detail)}
<span class="chat-processing-info-detail pointer-events-auto backdrop-blur-sm">{detail}</span>
{/each}
</div>
</div>
<style>
.chat-processing-info-container {
position: sticky;
top: 0;
z-index: 10;
padding: 0 1rem 0.75rem;
opacity: 0;
transform: translateY(50%);
transition:
opacity 300ms ease-out,
transform 300ms ease-out;
}
.chat-processing-info-container.visible {
opacity: 1;
transform: translateY(0);
}
.chat-processing-info-content {
display: flex;
flex-wrap: wrap;
align-items: center;
gap: 1rem;
justify-content: center;
max-width: 48rem;
margin: 0 auto;
}
.chat-processing-info-detail {
color: var(--muted-foreground);
font-size: 0.75rem;
padding: 0.25rem 0.75rem;
border-radius: 0.375rem;
font-family:
ui-monospace, SFMono-Regular, 'SF Mono', Consolas, 'Liberation Mono', Menlo, monospace;
white-space: nowrap;
}
@media (max-width: 768px) {
.chat-processing-info-content {
gap: 0.5rem;
}
.chat-processing-info-detail {
font-size: 0.7rem;
padding: 0.2rem 0.5rem;
}
}
</style>

21
tools/ui/src/lib/components/app/chat/index.ts
View file

@ -241,13 +241,18 @@ export { default as ChatFormActionAddToolsSubmenu } from './ChatForm/ChatFormAct
export { default as ChatFormActionAddMcpServersSubmenu } from './ChatForm/ChatFormActions/ChatFormActionAdd/ChatFormActionAddMcpServersSubmenu.svelte';
/**
* **ChatFormReasoningToggle** - Thinking toggle button with effort dropdown
* Dropdown submenu for selecting reasoning effort level.
*
* A toggle button with lightbulb icon that indicates thinking status.
* Shows the reasoning effort dropdown when clicked.
* Shows a "Reasoning" sub-menu item with a lightbulb icon indicating
* thinking status, and a nested list of effort levels.
* Only visible when the current model supports thinking.
*/
export { default as ChatFormReasoningToggle } from './ChatForm/ChatFormActions/ChatFormReasoningToggle.svelte';
export { default as ChatFormActionAddReasoningSubmenu } from './ChatForm/ChatFormActions/ChatFormActionAdd/ChatFormActionAddReasoningSubmenu.svelte';
/**
* Compact context-usage gauge with per-turn and cumulative breakdown in the tooltip.
*/
export { default as ChatFormContextGauge } from './ChatForm/ChatFormContextGauge/ChatFormContextGauge.svelte';
/**
* Hidden file input element for programmatic file selection.
@ -669,14 +674,6 @@ export { default as ChatScreenDragOverlay } from './ChatScreen/ChatScreenDragOve
*/
export { default as ChatScreenForm } from './ChatScreen/ChatScreenForm.svelte';
/**
* Processing info display during generation. Shows real-time statistics:
* tokens per second, prompt/completion token counts, and elapsed time.
* Data sourced from slotsService polling during active generation.
* Only visible when `isCurrentConversationLoading` is true.
*/
export { default as ChatScreenProcessingInfo } from './ChatScreen/ChatScreenProcessingInfo.svelte';
/**
* Server error alert displayed when the server is unreachable.
* Shows the error message with a retry button.

2
tools/ui/src/lib/components/app/content/CollapsibleContentBlock.svelte
View file

@ -76,7 +76,7 @@
open = value;
onToggle?.();
}}
class={className}
class="{className} my-0!"
>
<Card class="gap-0 border-muted bg-muted/30 py-0">
<Collapsible.Trigger class="flex w-full cursor-pointer items-start justify-between gap-2 p-3">

4
tools/ui/src/lib/components/app/content/SyntaxHighlightedCode.svelte
View file

@ -72,8 +72,8 @@
</script>
<div
class="code-preview-wrapper rounded-lg border border-border bg-muted {className}"
style="max-height: {maxHeight}; max-width: {maxWidth};"
class="code-preview-wrapper min-w-0 max-w-full overflow-x-auto rounded-lg border border-border bg-muted {className}"
style="max-height: {maxHeight}; {maxWidth ? `max-width: ${maxWidth};` : ''}"
>
<!-- Needs to be formatted as single line for proper rendering -->
<pre class="m-0"><code class="hljs text-sm leading-relaxed">{@html highlightedHtml}</code></pre>

36
tools/ui/src/lib/components/app/dialogs/DialogMcpServerRecommendations.svelte
View file

@ -4,9 +4,10 @@
import * as Dialog from '$lib/components/ui/dialog';
import { fly } from 'svelte/transition';
import { McpServerCardCompact, McpServerForm } from '$lib/components/app/mcp';
import { RECOMMENDED_MCP_SERVERS } from '$lib/constants';
import { RECOMMENDED_MCP_SERVERS, SETTINGS_KEYS } from '$lib/constants';
import { conversationsStore } from '$lib/stores/conversations.svelte';
import { mcpStore } from '$lib/stores/mcp.svelte';
import { settingsStore } from '$lib/stores/settings.svelte';
import { uuid } from '$lib/utils';
import { MCP_SERVERS_ADDED_TO_CHAT_LOCALSTORAGE_KEY, MCP_SERVER_ID_PREFIX } from '$lib/constants';
import type { MCPServerSettingsEntry } from '$lib/types';
@ -24,6 +25,22 @@
);
let addedServers = $state<MCPServerSettingsEntry[]>([]);
let didAddAny = $state(false);
let selectedRecommendedCount = $derived.by(
() => RECOMMENDED_MCP_SERVERS.filter((server) => selected[server.id]).length
);
let footerLabel = $derived.by(() => {
const recommended = selectedRecommendedCount;
const custom = addedServers.length;
const total = recommended + custom;
if (total === 0) return 'Continue';
if (recommended === 0) return custom === 1 ? 'Add server' : `Add ${custom} servers`;
if (custom === 0) return recommended === 1 ? 'Add server' : `Add ${recommended} servers`;
return `Add ${recommended} servers and ${custom} custom`;
});
let showAddForm = $state(false);
let newServerUrl = $state('');
@ -44,9 +61,14 @@
showAddForm = false;
newServerUrl = '';
newServerHeaders = '';
addedServers = [];
if (!didAddAny) {
settingsStore.updateConfig(SETTINGS_KEYS.MCP_SERVERS, []);
}
localStorage.setItem(MCP_SERVERS_ADDED_TO_CHAT_LOCALSTORAGE_KEY, 'true');
addedServers = [];
didAddAny = false;
}
open = value;
onOpenChange?.(value);
@ -59,6 +81,7 @@
}
function enableSelected() {
didAddAny = true;
localStorage.setItem(MCP_SERVERS_ADDED_TO_CHAT_LOCALSTORAGE_KEY, 'true');
for (const server of RECOMMENDED_MCP_SERVERS) {
@ -83,6 +106,8 @@
function saveNewServer() {
if (newServerUrlError) return;
didAddAny = true;
const newServerId = uuid() ?? `${MCP_SERVER_ID_PREFIX}-${Date.now()}`;
localStorage.setItem(MCP_SERVERS_ADDED_TO_CHAT_LOCALSTORAGE_KEY, 'true');
@ -174,7 +199,12 @@
<Dialog.Footer>
<Button variant="secondary" size="sm" onclick={() => handleOpenChange(false)}>Not now</Button>
<Button variant="default" size="sm" onclick={enableSelected}>Add selected</Button>
<Button
variant="default"
size="sm"
onclick={enableSelected}
disabled={footerLabel === 'Continue'}>{footerLabel}</Button
>
</Dialog.Footer>
</Dialog.Content>
</Dialog.Root>

18
tools/ui/src/lib/components/app/settings/SettingsChat/SettingsChatToolsTab.svelte
View file

@ -39,13 +39,17 @@
{@const faviconUrl = group.serverId ? mcpStore.getServerFavicon(group.serverId) : null}
<span class="inline-flex min-w-0 items-center gap-1.5 font-medium">
<McpServerIdentity
iconClass="h-4 w-4"
iconRounded="rounded-sm"
showVersion={false}
displayName={group.label}
{faviconUrl}
/>
{#if group.source === 'mcp'}
<McpServerIdentity
iconClass="h-4 w-4"
iconRounded="rounded-sm"
showVersion={false}
displayName={group.label}
{faviconUrl}
/>
{:else}
<TruncatedText text={group.label} class="font-medium" />
{/if}
</span>
<span class="ml-auto shrink-0 text-xs text-muted-foreground">

31
tools/ui/src/lib/components/ui/hover-card/hover-card-content.svelte Normal file
View file

@ -0,0 +1,31 @@
<script lang="ts">
import { LinkPreview as HoverCardPrimitive } from 'bits-ui';
import { cn, type WithoutChildrenOrChild } from '$lib/components/ui/utils.js';
import HoverCardPortal from './hover-card-portal.svelte';
import type { ComponentProps } from 'svelte';
let {
ref = $bindable(null),
class: className,
align = 'center',
sideOffset = 4,
portalProps,
...restProps
}: HoverCardPrimitive.ContentProps & {
portalProps?: WithoutChildrenOrChild<ComponentProps<typeof HoverCardPortal>>;
} = $props();
</script>
<HoverCardPortal {...portalProps}>
<HoverCardPrimitive.Content
bind:ref
data-slot="hover-card-content"
{align}
{sideOffset}
class={cn(
'data-open:animate-in data-closed:animate-out data-closed:fade-out-0 data-open:fade-in-0 data-closed:zoom-out-95 data-open:zoom-in-95 data-[side=bottom]:slide-in-from-top-2 data-[side=left]:slide-in-from-right-2 data-[side=right]:slide-in-from-left-2 data-[side=top]:slide-in-from-bottom-2 ring-foreground/10 bg-popover text-popover-foreground w-64 rounded-lg p-2.5 text-sm shadow-md ring-1 duration-100 z-50 origin-(--transform-origin) outline-hidden',
className
)}
{...restProps}
/>
</HoverCardPortal>

7
tools/ui/src/lib/components/ui/hover-card/hover-card-portal.svelte Normal file
View file

@ -0,0 +1,7 @@
<script lang="ts">
import { LinkPreview as HoverCardPrimitive } from 'bits-ui';
let { ...restProps }: HoverCardPrimitive.PortalProps = $props();
</script>
<HoverCardPrimitive.Portal {...restProps} />

7
tools/ui/src/lib/components/ui/hover-card/hover-card-trigger.svelte Normal file
View file

@ -0,0 +1,7 @@
<script lang="ts">
import { LinkPreview as HoverCardPrimitive } from 'bits-ui';
let { ref = $bindable(null), ...restProps }: HoverCardPrimitive.TriggerProps = $props();
</script>
<HoverCardPrimitive.Trigger bind:ref data-slot="hover-card-trigger" {...restProps} />

7
tools/ui/src/lib/components/ui/hover-card/hover-card.svelte Normal file
View file

@ -0,0 +1,7 @@
<script lang="ts">
import { LinkPreview as HoverCardPrimitive } from 'bits-ui';
let { open = $bindable(false), ...restProps }: HoverCardPrimitive.RootProps = $props();
</script>
<HoverCardPrimitive.Root bind:open {...restProps} />

15
tools/ui/src/lib/components/ui/hover-card/index.ts Normal file
View file

@ -0,0 +1,15 @@
import Root from './hover-card.svelte';
import Content from './hover-card-content.svelte';
import Trigger from './hover-card-trigger.svelte';
import Portal from './hover-card-portal.svelte';
export {
Root,
Content,
Trigger,
Portal,
Root as HoverCard,
Content as HoverCardContent,
Trigger as HoverCardTrigger,
Portal as HoverCardPortal
};

1
tools/ui/src/lib/constants/context-keys.ts
View file

@ -1,4 +1,3 @@
export const CONTEXT_KEY_MESSAGE_EDIT = 'chat-message-edit';
export const CONTEXT_KEY_CHAT_ACTIONS = 'chat-actions';
export const CONTEXT_KEY_CHAT_SETTINGS_CONFIG = 'chat-settings-config';
export const CONTEXT_KEY_PROCESSING_INFO = 'processing-info';

1
tools/ui/src/lib/constants/css-classes.ts
View file

@ -17,3 +17,4 @@ export const PANEL_CLASSES = `
`;
export const CHAT_FORM_POPOVER_MAX_HEIGHT = 'max-h-80';
export const DIALOG_SUBMENU_CONTENT = 'w-60';

3
tools/ui/src/lib/constants/reasoning-effort.ts
View file

@ -6,6 +6,7 @@ import type { ReasoningEffortLevel } from '$lib/types';
* Keys match the ReasoningEffort enum values for type-safe lookups.
*/
export const REASONING_EFFORT_LABELS: Record<string, string> = {
[ReasoningEffort.OFF]: 'Off',
[ReasoningEffort.LOW]: 'Low',
[ReasoningEffort.MEDIUM]: 'Medium',
[ReasoningEffort.HIGH]: 'High',
@ -13,7 +14,7 @@ export const REASONING_EFFORT_LABELS: Record<string, string> = {
};
export const REASONING_EFFORT_LEVELS: ReasoningEffortLevel[] = [
{ value: 'off', label: 'Off', isOff: true },
{ value: ReasoningEffort.OFF, label: 'Off', isOff: true },
{ value: ReasoningEffort.LOW, label: 'Low' },
{ value: ReasoningEffort.MEDIUM, label: 'Medium' },
{ value: ReasoningEffort.HIGH, label: 'High' },

2
tools/ui/src/lib/constants/settings-keys.ts
View file

@ -22,7 +22,6 @@ export const SETTINGS_KEYS = {
// Display
SHOW_MESSAGE_STATS: 'showMessageStats',
SHOW_THOUGHT_IN_PROGRESS: 'showThoughtInProgress',
KEEP_STATS_VISIBLE: 'keepStatsVisible',
AUTO_MIC_ON_EMPTY: 'autoMicOnEmpty',
RENDER_USER_CONTENT_AS_MARKDOWN: 'renderUserContentAsMarkdown',
DISABLE_AUTO_SCROLL: 'disableAutoScroll',
@ -61,7 +60,6 @@ export const SETTINGS_KEYS = {
MCP_REQUEST_TIMEOUT_SECONDS: 'mcpRequestTimeoutSeconds',
MCP_DEFAULT_SERVER_OVERRIDES: 'mcpDefaultServerOverrides',
AGENTIC_MAX_TURNS: 'agenticMaxTurns',
ALWAYS_SHOW_AGENTIC_TURNS: 'alwaysShowAgenticTurns',
AGENTIC_MAX_TOOL_PREVIEW_LINES: 'agenticMaxToolPreviewLines',
SHOW_TOOL_CALL_IN_PROGRESS: 'showToolCallInProgress',
// Performance

24
tools/ui/src/lib/constants/settings-registry.ts
View file

@ -258,18 +258,6 @@ const SETTINGS_REGISTRY: Record<string, SettingsSectionEntry> = {
paramType: SyncableParameterType.BOOLEAN
}
},
{
key: SETTINGS_KEYS.KEEP_STATS_VISIBLE,
label: 'Keep stats visible after generation',
help: 'Keep processing statistics visible after generation finishes.',
defaultValue: true,
type: SettingsFieldType.CHECKBOX,
section: SETTINGS_SECTION_SLUGS.DISPLAY,
sync: {
serverKey: SETTINGS_KEYS.KEEP_STATS_VISIBLE,
paramType: SyncableParameterType.BOOLEAN
}
},
{
key: SETTINGS_KEYS.AUTO_MIC_ON_EMPTY,
label: 'Show microphone on empty input',
@ -379,18 +367,6 @@ const SETTINGS_REGISTRY: Record<string, SettingsSectionEntry> = {
paramType: SyncableParameterType.BOOLEAN
}
},
{
key: SETTINGS_KEYS.ALWAYS_SHOW_AGENTIC_TURNS,
label: 'Always show agentic turns in conversation',
help: 'Always expand and display agentic loop turns in conversation messages.',
defaultValue: false,
type: SettingsFieldType.CHECKBOX,
section: SETTINGS_SECTION_SLUGS.DISPLAY,
sync: {
serverKey: SETTINGS_KEYS.ALWAYS_SHOW_AGENTIC_TURNS,
paramType: SyncableParameterType.BOOLEAN
}
},
{
key: SETTINGS_KEYS.SHOW_BUILD_VERSION,
label: 'Show build version information',

1
tools/ui/src/lib/constants/storage.ts
View file

@ -21,7 +21,6 @@ export const DISABLED_TOOLS_LOCALSTORAGE_KEY = `${STORAGE_APP_NAME}.disabledTool
/** Disabled tools keyed by stable selection identity, no migration from the name based key */
export const DISABLED_TOOL_KEYS_LOCALSTORAGE_KEY = `${STORAGE_APP_NAME}.disabledToolKeys`;
export const FAVORITE_MODELS_LOCALSTORAGE_KEY = `${STORAGE_APP_NAME}.favoriteModels`;
export const THINKING_ENABLED_DEFAULT_LOCALSTORAGE_KEY = `${STORAGE_APP_NAME}.thinkingEnabledDefault`;
export const REASONING_EFFORT_DEFAULT_LOCALSTORAGE_KEY = `${STORAGE_APP_NAME}.reasoningEffortDefault`;
/** Set when user has interacted with the MCP server recommendations dialog (checked servers, added custom server, or dismissed) */
export const MCP_SERVERS_ADDED_TO_CHAT_LOCALSTORAGE_KEY = `${STORAGE_APP_NAME}.mcpServersSetupDone`;

6
tools/ui/src/lib/contexts/index.ts
View file

@ -17,9 +17,3 @@ export {
setChatSettingsConfigContext,
type ChatSettingsConfigContext
} from './chat-settings-config.context';
export {
getProcessingInfoContext,
setProcessingInfoContext,
type ProcessingInfoContext
} from './processing-info.context';

16
tools/ui/src/lib/contexts/processing-info.context.ts
View file

@ -1,16 +0,0 @@
import { getContext, setContext } from 'svelte';
import { CONTEXT_KEY_PROCESSING_INFO } from '$lib/constants';
export interface ProcessingInfoContext {
readonly showProcessingInfo: boolean;
}
const PROCESSING_INFO_KEY = Symbol.for(CONTEXT_KEY_PROCESSING_INFO);
export function setProcessingInfoContext(ctx: ProcessingInfoContext): ProcessingInfoContext {
return setContext(PROCESSING_INFO_KEY, ctx);
}
export function getProcessingInfoContext(): ProcessingInfoContext {
return getContext(PROCESSING_INFO_KEY);
}

6
tools/ui/src/lib/enums/chat.enums.ts
View file

@ -5,6 +5,12 @@ export enum ChatMessageStatsView {
SUMMARY = 'summary'
}
export enum ChatMessageStatisticsMode {
SWITCHABLE = 'switchable',
READING = 'reading',
GENERATION = 'generation'
}
/**
* Connection state of a streamed completion, drives the resume status indicator.
*/

1
tools/ui/src/lib/enums/index.ts
View file

@ -10,6 +10,7 @@ export { AgenticSectionType, ContinueIntentKind, ToolCallType } from './agentic.
export {
ChatMessageStatsView,
ChatMessageStatisticsMode,
StreamConnectionState,
ContentPartType,
ConversationSelectionMode,

1
tools/ui/src/lib/enums/reasoning-effort.enums.ts
View file

@ -3,6 +3,7 @@
* These values are sent to the server and mapped to token budgets.
*/
export enum ReasoningEffort {
OFF = 'off',
LOW = 'low',
MEDIUM = 'medium',
HIGH = 'high',

295
tools/ui/src/lib/hooks/use-context-gauge.svelte.ts Normal file
View file

@ -0,0 +1,295 @@
/**
* Reactive state for the context usage gauge: resolves the active model,
* fetches its cached props, parses live server stats, and exposes per-turn
* read / fresh / cache / output and cumulative token counts.
*/
import {
modelsStore,
modelOptions,
selectedModelId,
singleModelName
} from '$lib/stores/models.svelte';
import { chatStore } from '$lib/stores/chat.svelte';
import { activeMessages } from '$lib/stores/conversations.svelte';
import { isRouterMode } from '$lib/stores/server.svelte';
import { MessageRole } from '$lib/enums';
import { STATS_UNITS } from '$lib/constants';
import type { ChatMessageTimings, DatabaseMessage } from '$lib/types';
import { useProcessingState } from './use-processing-state.svelte';
import {
colorLevelFromPercent,
type ColorLevel
} from '$lib/components/app/chat/ChatForm/ChatFormContextGauge/context-gauge';
interface LiveStats {
freshTokens: number;
promptTokens: number;
cacheTokens: number;
outputTokens: number;
}
export interface UseContextGaugeReturn {
readonly activeModelId: string | null;
readonly isActiveModelLoaded: boolean;
readonly isActiveModelLoading: boolean;
readonly contextTotal: number | null;
readonly contextUsed: number;
readonly currentRead: number;
readonly currentFresh: number;
readonly currentCache: number;
readonly currentOutput: number;
readonly kvTotal: number;
readonly cumulativeRead: number;
readonly cumulativeOutput: number;
readonly cumulativeCacheTotal: number;
readonly averageTokensPerSecond: number | null;
readonly contextPercent: number | null;
readonly colorLevel: ColorLevel;
readonly transientDetails: string[];
readonly hasAnyUsage: boolean;
loadModel(): Promise<void>;
startMonitoring(): void;
}
function lastAssistantTimings(messages: DatabaseMessage[]): ChatMessageTimings | undefined {
for (let i = messages.length - 1; i >= 0; i--) {
const m = messages[i];
if (m.role === MessageRole.ASSISTANT && m.timings) return m.timings;
}
return undefined;
}
function deriveLiveStats(
state: ReturnType<typeof useProcessingState>['processingState']
): LiveStats | null {
if (!state || (state.status !== 'preparing' && state.status !== 'generating')) {
return null;
}
const promptTokens = state.promptTokens ?? 0;
const cacheTokens = state.cacheTokens ?? 0;
return {
freshTokens: promptTokens,
promptTokens: promptTokens + cacheTokens,
cacheTokens,
outputTokens: state.outputTokensUsed ?? 0
};
}
const TRANSIENT_DETAILS_EXCLUDED_PREFIXES = ['Context:', 'Output:'];
function filterTransientDetails(raw: string[]): string[] {
return raw.filter((detail) => {
if (TRANSIENT_DETAILS_EXCLUDED_PREFIXES.some((prefix) => detail.startsWith(prefix))) {
return false;
}
return !detail.includes(STATS_UNITS.TOKENS_PER_SECOND);
});
}
export function useContextGauge(): UseContextGaugeReturn {
const processingState = useProcessingState();
// Resolve the model the gauge reports context for: explicit selection >
// last assistant model > single-model mode (mirrors useChatScreenActiveModel).
const activeModelId = $derived.by(() => {
if (!isRouterMode()) {
return singleModelName();
}
const selectedId = selectedModelId();
if (selectedId) {
const model = modelOptions().find((m) => m.id === selectedId);
if (model) return model.model;
}
return chatStore.getConversationModel(activeMessages() as DatabaseMessage[]);
});
const isActiveModelLoaded = $derived(
activeModelId !== null && modelsStore.isModelLoaded(activeModelId)
);
const isActiveModelLoading = $derived(
activeModelId !== null && modelsStore.isModelOperationInProgress(activeModelId)
);
// Pull /props on demand so n_ctx surfaces before the first chat request.
$effect(() => {
if (activeModelId && isActiveModelLoaded) {
const cached = modelsStore.getModelProps(activeModelId);
if (!cached) {
void modelsStore.fetchModelProps(activeModelId);
}
}
});
const contextTotal = $derived.by(() => {
void modelsStore.propsCacheVersion;
return activeModelId ? modelsStore.getModelContextSize(activeModelId) : null;
});
const liveStats = $derived(deriveLiveStats(processingState.processingState));
const currentRead = $derived.by(() => {
const timings = lastAssistantTimings(activeMessages() as DatabaseMessage[]);
let read = 0;
if (timings) {
read = (timings.prompt_n ?? 0) + (timings.cache_n ?? 0);
}
// live.promptTokens is already the combined reading (prompt + cache),
// so do not also add live.cacheTokens.
if (liveStats && liveStats.promptTokens > 0) {
read = Math.max(read, liveStats.promptTokens);
}
return read;
});
const currentFresh = $derived.by(() => {
const timings = lastAssistantTimings(activeMessages() as DatabaseMessage[]);
const fresh = timings?.prompt_n ?? 0;
return Math.max(fresh, liveStats?.freshTokens ?? 0);
});
const currentCache = $derived.by(() => {
const timings = lastAssistantTimings(activeMessages() as DatabaseMessage[]);
const cached = timings?.cache_n ?? 0;
if (liveStats && liveStats.promptTokens > 0) {
return Math.max(cached, liveStats.cacheTokens);
}
return cached;
});
const currentOutput = $derived.by(() => {
if (liveStats && liveStats.outputTokens > 0) return liveStats.outputTokens;
const timings = lastAssistantTimings(activeMessages() as DatabaseMessage[]);
return timings?.predicted_n ?? 0;
});
const kvTotal = $derived(currentRead + currentOutput);
const contextUsed = $derived(currentRead + currentOutput);
const cumulative = $derived.by(() => {
const messages = activeMessages() as DatabaseMessage[];
// Agentic sessions stamp the same agentic.llm totals onto every
// assistant message; cache_n is never per-turn so cache_total stays 0.
const agenticMessages = messages.filter(
(m) => m.role === MessageRole.ASSISTANT && m.timings?.agentic?.llm?.predicted_n != null
);
if (agenticMessages.length > 0) {
const llm = agenticMessages[agenticMessages.length - 1].timings!.agentic!.llm;
const output = llm.predicted_n ?? 0;
const outputMs = llm.predicted_ms ?? 0;
const averageTokensPerSecond = outputMs > 0 && output > 0 ? (output / outputMs) * 1000 : null;
return {
read: llm.prompt_n ?? 0,
output,
cacheTotal: 0,
averageTokensPerSecond
};
}
let read = 0;
let output = 0;
let outputMs = 0;
let cacheTotal = 0;
for (const m of messages) {
if (m.role !== MessageRole.ASSISTANT || !m.timings) continue;
read += m.timings.prompt_n ?? 0;
cacheTotal += m.timings.cache_n ?? 0;
output += m.timings.predicted_n ?? 0;
outputMs += m.timings.predicted_ms ?? 0;
}
const averageTokensPerSecond = outputMs > 0 && output > 0 ? (output / outputMs) * 1000 : null;
return { read, output, cacheTotal, averageTokensPerSecond };
});
const contextPercent = $derived.by(() => {
if (contextTotal === null || contextTotal <= 0) return null;
return Math.round((contextUsed / contextTotal) * 100);
});
const colorLevel = $derived(colorLevelFromPercent(contextPercent));
// Drop lines the surrounding Context / Output / speed rows already render.
const transientDetails = $derived(filterTransientDetails(processingState.getTechnicalDetails()));
const hasAnyUsage = $derived(
cumulative.read > 0 ||
cumulative.output > 0 ||
currentRead > 0 ||
currentOutput > 0 ||
cumulative.averageTokensPerSecond !== null ||
transientDetails.length > 0
);
async function loadModel() {
if (!activeModelId || isActiveModelLoading) return;
try {
await modelsStore.loadModel(activeModelId);
} catch {
// toast already surfaced by modelsStore.loadModel
}
}
return {
get activeModelId() {
return activeModelId;
},
get isActiveModelLoaded() {
return isActiveModelLoaded;
},
get isActiveModelLoading() {
return isActiveModelLoading;
},
get contextTotal() {
return contextTotal;
},
get contextUsed() {
return contextUsed;
},
get currentRead() {
return currentRead;
},
get currentFresh() {
return currentFresh;
},
get currentCache() {
return currentCache;
},
get currentOutput() {
return currentOutput;
},
get kvTotal() {
return kvTotal;
},
get cumulativeRead() {
return cumulative.read;
},
get cumulativeOutput() {
return cumulative.output;
},
get cumulativeCacheTotal() {
return cumulative.cacheTotal;
},
get averageTokensPerSecond() {
return cumulative.averageTokensPerSecond;
},
get contextPercent() {
return contextPercent;
},
get colorLevel() {
return colorLevel;
},
get transientDetails() {
return transientDetails;
},
get hasAnyUsage() {
return hasAnyUsage;
},
loadModel,
startMonitoring: () => processingState.startMonitoring()
};
}

5
tools/ui/src/lib/hooks/use-mcp-recommendations.svelte.ts
View file

@ -54,11 +54,6 @@ export function useMcpRecommendations() {
// effect, and we must not wipe the timeout that was just scheduled.
if (checked) return;
if (mcpStore.optedInRecommendationIds.size > 0) {
checked = true;
return;
}
const hasRecommendations = mcpStore
.getServers()
.some((server) => RECOMMENDED_MCP_SERVER_IDS.has(server.id));

10
tools/ui/src/lib/hooks/use-processing-state.svelte.ts
View file

@ -1,5 +1,4 @@
import { activeProcessingState } from '$lib/stores/chat.svelte';
import { config } from '$lib/stores/settings.svelte';
import { STATS_UNITS } from '$lib/constants';
import type { ApiProcessingState, LiveProcessingStats, LiveGenerationStats } from '$lib/types';
@ -46,7 +45,6 @@ export function useProcessingState(): UseProcessingStateReturn {
return activeProcessingState();
});
// Track last known state for keepStatsVisible functionality
$effect(() => {
if (processingState && isMonitoring) {
lastKnownState = processingState;
@ -88,14 +86,8 @@ export function useProcessingState(): UseProcessingStateReturn {
function stopMonitoring(): void {
if (!isMonitoring) return;
isMonitoring = false;
// Only clear last known state if keepStatsVisible is disabled
const currentConfig = config();
if (!currentConfig.keepStatsVisible) {
lastKnownState = null;
lastKnownProcessingStats = null;
}
isMonitoring = false;
}
function getProcessingMessage(): string {

3
tools/ui/src/lib/services/chat.service.ts
View file

@ -340,6 +340,7 @@ export class ChatService {
if (stream && conversationId) {
headers['X-Conversation-Id'] = streamIdentity(conversationId, options.model);
}
const response = await fetch(API_CHAT.COMPLETIONS, {
method: 'POST',
headers,
@ -1015,7 +1016,7 @@ export class ChatService {
*
* @param response - The fetch Response object containing the JSON data
* @param onComplete - Optional callback invoked when response is successfully parsed
* @param onError - Optional callback invoked if an error occurs during parsing
* @param onError - Optional callback invoked if an error occurs while parsing
* @returns {Promise<string>} Promise that resolves to the generated content string
* @throws {Error} if the response cannot be parsed or is malformed
*/

6
tools/ui/src/lib/services/migration.service.ts
View file

@ -564,9 +564,9 @@ const configTypesMigration: Migration = {
const config = JSON.parse(configRaw);
let changed = false;
// Pre-schema configs persisted booleans as the strings "true"/"false", which the
// strict server schema now rejects. Coerce those back to real booleans. No config
// string field holds exactly "true"/"false", so the match is unambiguous.
// Pre-schema configs persisted booleans as "true"/"false" strings; the strict server
// schema rejects them. No config string field holds exactly "true"/"false", so the
// match is unambiguous.
for (const key of Object.keys(config)) {
if (config[key] === 'true') {
config[key] = true;

2
tools/ui/src/lib/stores/agentic.svelte.ts
View file

@ -477,7 +477,7 @@ class AgenticStore {
conversationId: string;
messages: ApiChatMessageData[];
options: AgenticFlowOptions;
tools: ReturnType<typeof mcpStore.getToolDefinitionsForLLM>;
tools: ReturnType<typeof toolsStore.getEnabledToolsForLLM>;
agenticConfig: AgenticConfig;
callbacks: AgenticFlowCallbacks;
signal?: AbortSignal;

4
tools/ui/src/lib/stores/chat.svelte.ts
View file

@ -60,6 +60,7 @@ import {
ErrorDialogType,
MessageRole,
MessageType,
ReasoningEffort,
StreamConnectionState
} from '$lib/enums';
@ -2334,7 +2335,8 @@ class ChatStore {
if (currentConfig.excludeReasoningFromContext) apiOptions.excludeReasoningFromContext = true;
apiOptions.enableThinking = conversationsStore.getThinkingEnabled();
apiOptions.reasoningEffort = conversationsStore.getReasoningEffort();
const effort = conversationsStore.getReasoningEffort();
if (effort !== ReasoningEffort.OFF) apiOptions.reasoningEffort = effort;
if (hasValue(currentConfig.temperature))
apiOptions.temperature = Number(currentConfig.temperature);

78
tools/ui/src/lib/stores/conversations.svelte.ts
View file

@ -47,7 +47,6 @@ import {
NON_ALPHANUMERIC_REGEX,
MULTIPLE_UNDERSCORE_REGEX,
SETTINGS_KEYS,
THINKING_ENABLED_DEFAULT_LOCALSTORAGE_KEY,
REASONING_EFFORT_DEFAULT_LOCALSTORAGE_KEY
} from '$lib/constants';
@ -84,11 +83,16 @@ class ConversationsStore {
/** Pending MCP server overrides for new conversations (before first message) */
pendingMcpServerOverrides = $state<McpServerOverride[]>(ConversationsStore.loadMcpDefaults());
/** Global (non-conversation-specific) thinking toggle default */
pendingThinkingEnabled = $state(ConversationsStore.loadThinkingDefaults());
/** Global (non-conversation-specific) thinking toggle default, derived from reasoning effort */
pendingThinkingEnabled = $state(false);
/** Global (non-conversation-specific) reasoning effort default */
pendingReasoningEffort = $state<ReasoningEffort>(ConversationsStore.loadReasoningEffortDefault());
pendingReasoningEffort = $state<ReasoningEffort | ReasoningEffort.OFF>(
ConversationsStore.loadReasoningEffortDefault()
);
/** Last non-off reasoning effort, restored when re-enabling thinking globally */
private lastNonOffEffort: ReasoningEffort | null = null;
private static loadMcpDefaults(): McpServerOverride[] {
const raw = config()[SETTINGS_KEYS.MCP_DEFAULT_SERVER_OVERRIDES];
@ -112,35 +116,14 @@ class ConversationsStore {
settingsStore.updateConfig(SETTINGS_KEYS.MCP_DEFAULT_SERVER_OVERRIDES, JSON.stringify(plain));
}
/** Load thinking-enabled default from localStorage */
private static loadThinkingDefaults(): boolean {
if (typeof globalThis.localStorage === 'undefined') return true;
try {
const raw = localStorage.getItem(THINKING_ENABLED_DEFAULT_LOCALSTORAGE_KEY);
if (!raw) return true;
return raw === 'true';
} catch {
return true;
}
}
/** Persist thinking-enabled default to localStorage */
private saveThinkingDefaults(): void {
if (typeof globalThis.localStorage === 'undefined') return;
localStorage.setItem(
THINKING_ENABLED_DEFAULT_LOCALSTORAGE_KEY,
this.pendingThinkingEnabled ? 'true' : 'false'
);
}
/** Load reasoning effort default from localStorage */
private static loadReasoningEffortDefault(): ReasoningEffort {
if (typeof globalThis.localStorage === 'undefined') return ReasoningEffort.MEDIUM;
private static loadReasoningEffortDefault(): ReasoningEffort | ReasoningEffort.OFF {
if (typeof globalThis.localStorage === 'undefined') return ReasoningEffort.OFF;
try {
const raw = localStorage.getItem(REASONING_EFFORT_DEFAULT_LOCALSTORAGE_KEY);
return (raw as ReasoningEffort) || ReasoningEffort.MEDIUM;
return (raw as ReasoningEffort | ReasoningEffort.OFF) || ReasoningEffort.OFF;
} catch {
return ReasoningEffort.MEDIUM;
return ReasoningEffort.OFF;
}
}
@ -303,10 +286,17 @@ class ConversationsStore {
this.pendingMcpServerOverrides = [];
}
// Inherit global thinking default into the new conversation
conversation.thinkingEnabled = this.pendingThinkingEnabled;
// Inherit global thinking/reasoning defaults into the new conversation
const thinkingEnabled = this.getThinkingEnabled();
conversation.thinkingEnabled = thinkingEnabled;
conversation.reasoningEffort =
this.pendingReasoningEffort === ReasoningEffort.OFF ? undefined : this.pendingReasoningEffort;
await DatabaseService.updateConversation(conversation.id, {
thinkingEnabled: this.pendingThinkingEnabled
thinkingEnabled,
reasoningEffort:
this.pendingReasoningEffort === ReasoningEffort.OFF
? undefined
: this.pendingReasoningEffort
});
this.conversations = [conversation, ...this.conversations];
@ -332,7 +322,6 @@ class ConversationsStore {
}
this.pendingMcpServerOverrides = [];
this.pendingThinkingEnabled = ConversationsStore.loadThinkingDefaults();
this.activeConversation = conversation;
if (conversation.currNode) {
@ -363,7 +352,7 @@ class ConversationsStore {
this.activeMessages = [];
// reload defaults so new chats inherit persisted state
this.pendingMcpServerOverrides = ConversationsStore.loadMcpDefaults();
this.pendingThinkingEnabled = ConversationsStore.loadThinkingDefaults();
this.pendingReasoningEffort = ConversationsStore.loadReasoningEffortDefault();
}
/**
@ -794,9 +783,11 @@ class ConversationsStore {
*/
getThinkingEnabled(): boolean {
if (this.activeConversation) {
return this.activeConversation.thinkingEnabled ?? this.pendingThinkingEnabled;
if (this.activeConversation.thinkingEnabled !== undefined) {
return this.activeConversation.thinkingEnabled;
}
}
return this.pendingThinkingEnabled;
return this.getReasoningEffort() !== ReasoningEffort.OFF;
}
/**
@ -806,8 +797,17 @@ class ConversationsStore {
*/
async setThinkingEnabled(enabled: boolean): Promise<void> {
if (!this.activeConversation) {
this.pendingThinkingEnabled = enabled;
this.saveThinkingDefaults();
if (enabled) {
const effort = this.lastNonOffEffort ?? ReasoningEffort.LOW;
this.pendingReasoningEffort = effort;
this.saveReasoningEffortDefaults();
} else {
if (this.pendingReasoningEffort !== ReasoningEffort.OFF) {
this.lastNonOffEffort = this.pendingReasoningEffort;
}
this.pendingReasoningEffort = ReasoningEffort.OFF;
this.saveReasoningEffortDefaults();
}
return;
}
@ -831,7 +831,7 @@ class ConversationsStore {
* Gets the effective reasoning effort for the active conversation.
* Returns the conversation override if set, otherwise the global default.
*/
getReasoningEffort(): ReasoningEffort {
getReasoningEffort(): ReasoningEffort | ReasoningEffort.OFF {
if (this.activeConversation) {
return this.activeConversation.reasoningEffort ?? this.pendingReasoningEffort;
}

Some files were not shown because too many files have changed in this diff Show more