prof: fix tensor dims formatter

common/arg.cpp

@@ -873,9 +873,7 @@ common_params_context common_params_parser_init(common_params & params, llama_ex
         sampler_type_chars += common_sampler_type_to_chr(sampler);
         sampler_type_names += common_sampler_type_to_str(sampler) + ";";
     }
-    if (!sampler_type_names.empty()) {
-        sampler_type_names.pop_back(); // remove last semicolon
-    }
+    sampler_type_names.pop_back();
 
 
     /**

convert_hf_to_gguf.py

@@ -189,10 +189,10 @@ class ModelBase:
             return tensors
 
         prefix = "model" if not self.is_mistral_format else "consolidated"
-        part_names: list[str] = ModelBase.get_model_part_names(self.dir_model, prefix, ".safetensors")
+        part_names: set[str] = set(ModelBase.get_model_part_names(self.dir_model, prefix, ".safetensors"))
         is_safetensors: bool = len(part_names) > 0
         if not is_safetensors:
-            part_names = ModelBase.get_model_part_names(self.dir_model, "pytorch_model", ".bin")
+            part_names = set(ModelBase.get_model_part_names(self.dir_model, "pytorch_model", ".bin"))
 
         tensor_names_from_index: set[str] = set()
 
@@ -209,8 +209,7 @@ class ModelBase:
                     if weight_map is None or not isinstance(weight_map, dict):
                         raise ValueError(f"Can't load 'weight_map' from {index_name!r}")
                     tensor_names_from_index.update(weight_map.keys())
-                    part_dict: dict[str, None] = dict.fromkeys(weight_map.values(), None)
-                    part_names = sorted(part_dict.keys())
+                    part_names |= set(weight_map.values())
             else:
                 weight_map = {}
         else:

ggml/CMakeLists.txt

@@ -125,6 +125,7 @@ option(GGML_CCACHE "ggml: use ccache if available"                   ON)
 option(GGML_ALL_WARNINGS           "ggml: enable all compiler warnings"                   ON)
 option(GGML_ALL_WARNINGS_3RD_PARTY "ggml: enable all compiler warnings in 3rd party libs" OFF)
 option(GGML_GPROF                  "ggml: enable gprof"                                   OFF)
+option(GGML_GRAPH_PROFILER         "ggml: enable internal Graph and Op profiler"          OFF)
 
 # build
 option(GGML_FATAL_WARNINGS    "ggml: enable -Werror flag"    OFF)

ggml/src/CMakeLists.txt

@@ -8,6 +8,10 @@ if (CMAKE_SYSTEM_NAME MATCHES "Linux")
     add_compile_definitions($<$<CONFIG:Debug>:_GLIBCXX_ASSERTIONS>)
 endif()
 
+if (GGML_GRAPH_PROFILER)
+    add_compile_definitions(GGML_GRAPH_PROFILER)
+endif()
+
 if (NOT MSVC)
     if (GGML_SANITIZE_THREAD)
         add_compile_options(-fsanitize=thread)
@@ -205,6 +209,8 @@ add_library(ggml-base
             ggml-threading.h
             ggml-quants.c
             ggml-quants.h
+            ggml-profile.h
+            ggml-profile.cpp
             gguf.cpp)
 
 set_target_properties(ggml-base PROPERTIES

ggml/src/ggml-cpu/ggml-cpu.c

@@ -13,6 +13,7 @@
 #include "binary-ops.h"
 #include "vec.h"
 #include "ops.h"
+#include "ggml-profile.h"
 #include "ggml.h"
 
 #if defined(_MSC_VER) || defined(__MINGW32__)
@@ -2943,6 +2944,8 @@ static thread_ret_t ggml_graph_compute_thread(void * data) {
             continue;
         }
 
+        ggml_graph_profile_event(cgraph, GGML_PROF_OP_START, node_n, state->ith);
+
         ggml_compute_forward(&params, node);
 
         if (state->ith == 0 && cplan->abort_callback &&
@@ -2951,9 +2954,13 @@ static thread_ret_t ggml_graph_compute_thread(void * data) {
             tp->ec    = GGML_STATUS_ABORTED;
         }
 
+        ggml_graph_profile_event(cgraph, GGML_PROF_OP_SYNC, node_n, state->ith);
+
         if (node_n + 1 < cgraph->n_nodes) {
             ggml_barrier(state->threadpool);
         }
+
+        ggml_graph_profile_event(cgraph, GGML_PROF_OP_END,  node_n, state->ith);
     }
 
     GGML_PRINT_DEBUG("thread #%d compute-done cplan %p last-graph %d \n", state->ith, cplan, state->last_graph);
@@ -3188,6 +3195,8 @@ enum ggml_status ggml_graph_compute(struct ggml_cgraph * cgraph, struct ggml_cpl
     int n_threads                               = cplan->n_threads;
     struct ggml_threadpool * threadpool = cplan->threadpool;
 
+    ggml_graph_profile_start(cgraph, n_threads);
+
     bool disposable_threadpool = false;
 
     if (threadpool == NULL) {
@@ -3246,6 +3255,8 @@ enum ggml_status ggml_graph_compute(struct ggml_cgraph * cgraph, struct ggml_cpl
     // don't leave affinity set on the main thread
     clear_numa_thread_affinity();
 
+    ggml_graph_profile_finish(cgraph, n_threads);
+
     enum ggml_status ret = threadpool->ec;
 
     if (disposable_threadpool) {

ggml/src/ggml-cuda/mma.cuh

@@ -78,25 +78,27 @@ namespace ggml_cuda_mma {
         // MIRRORED == Each data value is held exactly once per thread subgroup.
         DATA_LAYOUT_I_MAJOR           =  0, // Always used for Turing, Ampere, Ada Lovelace, consumer Blackwell, matrix A&B for RDNA4 and CDNA.
         DATA_LAYOUT_J_MAJOR           = 10, // Matrix C for CDNA and RDNA4, int and float matrix C for RDNA3.
-        DATA_LAYOUT_I_MAJOR_MIRRORED  = 20, // Volta, matrix A&B for RDNA3.
+        DATA_LAYOUT_I_MAJOR_MIRRORED  = 20,
         DATA_LAYOUT_J_MAJOR_MIRRORED  = 30,
+        DATA_LAYOUT_I_MAJOR_DUAL      = 40, // Matrix A&B for RDNA3.
     };
     // Implemented mma combinations are:
     //   - (I_MAJOR, I_MAJOR)          -> I_MAJOR
     //   - (I_MAJOR, I_MAJOR_MIRRORED) -> I_MAJOR
     //   - (I_MAJOR, J_MAJOR_MIRRORED) -> I_MAJOR
 
-    static constexpr bool is_i_major(const data_layout dl) {
+    constexpr bool is_i_major(const data_layout dl) {
         return dl == DATA_LAYOUT_I_MAJOR ||
-               dl == DATA_LAYOUT_I_MAJOR_MIRRORED;
+               dl == DATA_LAYOUT_I_MAJOR_MIRRORED ||
+               dl == DATA_LAYOUT_I_MAJOR_DUAL;
     }
 
-    static constexpr __device__ data_layout get_input_data_layout() {
-#if defined(RDNA3) || __CUDA_ARCH__ == GGML_CUDA_CC_VOLTA
-        return DATA_LAYOUT_I_MAJOR_MIRRORED;
+    constexpr data_layout get_input_data_layout() {
+#if defined(RDNA3)
+        return DATA_LAYOUT_I_MAJOR_DUAL;
 #else
         return DATA_LAYOUT_I_MAJOR;
-#endif // defined(RDNA3) || __CUDA_ARCH__ == GGML_CUDA_CC_VOLTA
+#endif // defined(RDNA3)
     }
 
     template <int I_, int J_, typename T, data_layout ds_=DATA_LAYOUT_I_MAJOR>
@@ -460,65 +462,11 @@ namespace ggml_cuda_mma {
         }
     };
 
-    template <int I_, int J_, typename T>
-    struct tile<I_, J_, T, DATA_LAYOUT_I_MAJOR_MIRRORED> {
-        static constexpr int         I  = I_;
-        static constexpr int         J  = J_;
-        static constexpr data_layout dl = DATA_LAYOUT_I_MAJOR_MIRRORED;
-
-        // RDNA3
-        static constexpr int         ne = I * J / 32 * 2;
-
-        T x[ne] = {0};
-
-        static constexpr __device__ bool supported() {
-            if (I == 16 && J == 16) return true;
-            if (I == 16 && J == 8)  return true;
-            if (I == 16 && J == 4)  return true;
-            return false;
-        }
-
-        static __device__ __forceinline__ int get_i(const int /*l*/) {
-            if constexpr (supported()) {
-                return threadIdx.x % 16;
-            } else {
-                NO_DEVICE_CODE;
-                return -1;
-            }
-        }
-
-        static __device__ __forceinline__ int get_j(const int l) {
-            if constexpr (supported()) {
-                return l;
-            } else {
-                NO_DEVICE_CODE;
-                return -1;
-            }
-        }
-    };
-
     template <int I_, int J_>
     struct tile<I_, J_, half2, DATA_LAYOUT_I_MAJOR_MIRRORED> {
         static constexpr int         I  = I_;
         static constexpr int         J  = J_;
         static constexpr data_layout dl = DATA_LAYOUT_I_MAJOR_MIRRORED;
-#if defined(RDNA3)
-        static constexpr int         ne = tile<I_, J_, float, DATA_LAYOUT_I_MAJOR_MIRRORED>::ne;
-
-        half2 x[ne] = {{0.0f, 0.0f}};
-
-        static constexpr __device__ bool supported() {
-            return tile<I_, J_, float, DATA_LAYOUT_I_MAJOR_MIRRORED>::supported();
-        }
-
-        static __device__ __forceinline__ int get_i(const int l) {
-            return tile<I_, J_, float, DATA_LAYOUT_I_MAJOR_MIRRORED>::get_i(l);
-        }
-
-        static __device__ __forceinline__ int get_j(const int l) {
-            return tile<I_, J_, float, DATA_LAYOUT_I_MAJOR_MIRRORED>::get_j(l);
-        }
-#else // Volta
         static constexpr int         ne = I * J / (WARP_SIZE/4);
 
         half2 x[ne] = {{0.0f, 0.0f}};
@@ -545,29 +493,6 @@ namespace ggml_cuda_mma {
                 return -1;
             }
         }
-#endif // defined(RDNA3)
-    };
-
-    template <int I_, int J_>
-    struct tile<I_, J_, nv_bfloat162, DATA_LAYOUT_I_MAJOR_MIRRORED> {
-        static constexpr int         I  = I_;
-        static constexpr int         J  = J_;
-        static constexpr data_layout dl = DATA_LAYOUT_I_MAJOR_MIRRORED;
-        static constexpr int         ne = tile<I_, J_, float, DATA_LAYOUT_I_MAJOR_MIRRORED>::ne;
-
-        nv_bfloat162 x[ne] = {{0.0f, 0.0f}};
-
-        static constexpr __device__ bool supported() {
-            return tile<I_, J_, float, DATA_LAYOUT_I_MAJOR_MIRRORED>::supported();
-        }
-
-        static __device__ __forceinline__ int get_i(const int l) {
-            return tile<I_, J_, float, DATA_LAYOUT_I_MAJOR_MIRRORED>::get_i(l);
-        }
-
-        static __device__ __forceinline__ int get_j(const int l) {
-            return tile<I_, J_, float, DATA_LAYOUT_I_MAJOR_MIRRORED>::get_j(l);
-        }
     };
 
     template <int I_, int J_>
@@ -603,6 +528,42 @@ namespace ggml_cuda_mma {
         }
     };
 
+    template <int I_, int J_, typename T>
+    struct tile<I_, J_, T, DATA_LAYOUT_I_MAJOR_DUAL> {
+        static constexpr int         I  = I_;
+        static constexpr int         J  = J_;
+        static constexpr data_layout dl = DATA_LAYOUT_I_MAJOR_DUAL;
+
+        static constexpr int         ne = I * J / 32 * 2;
+
+        T x[ne] = {0};
+
+        static constexpr __device__ bool supported() {
+            if (I == 16 && J == 16) return true;
+            if (I == 16 && J == 8)  return true;
+            if (I == 16 && J == 4)  return true;
+            return false;
+        }
+
+        static __device__ __forceinline__ int get_i(const int l) {
+            if constexpr (supported()) {
+                return threadIdx.x % 16;
+            } else {
+                NO_DEVICE_CODE;
+                return -1;
+            }
+        }
+
+        static __device__ __forceinline__ int get_j(const int l) {
+            if constexpr (supported()) {
+                return l;
+            } else {
+                NO_DEVICE_CODE;
+                return -1;
+            }
+        }
+    };
+
 #if defined(TURING_MMA_AVAILABLE)
     template <int I, int J>
     static __device__ __forceinline__ tile<I, J/2, half2> get_half2(const tile<I, J, float> & tile_float) {

ggml/src/ggml-impl.h

@@ -324,6 +324,8 @@ enum ggml_cgraph_eval_order {
     GGML_CGRAPH_EVAL_ORDER_COUNT
 };
 
+struct ggml_profile_data;
+
 struct ggml_cgraph {
     int size;    // maximum number of nodes/leafs/grads/grad_accs
     int n_nodes; // number of nodes currently in use
@@ -335,6 +337,8 @@ struct ggml_cgraph {
     struct ggml_tensor ** leafs;     // tensors with constant data
     int32_t             * use_counts;// number of uses of each tensor, indexed by hash table slot
 
+    struct ggml_profile_data * prof;
+
     struct ggml_hash_set visited_hash_set;
 
     enum ggml_cgraph_eval_order order;

ggml/src/ggml-profile.cpp

@@ -0,0 +1,199 @@
+#include "ggml-profile.h"
+
+#include <stdint.h>
+#include <stdlib.h>
+#include <stdio.h>
+#include <string.h>
+
+#include <string>
+#include <chrono>
+
+#ifdef GGML_GRAPH_PROFILER
+
+struct ggml_profile_output {
+    const char * prefix;
+    FILE *       stream;
+};
+
+extern "C" void ggml_graph_profile_init(struct ggml_cgraph *cg, int n_threads)
+{
+    // TODO: make this a param
+    const char *env = getenv("GGML_GRAPH_PROFILE");
+    if (!env) { return; }
+
+    // The number of threads may change between passes (pp vs tg).
+    // Allocate for max_n_threads for simplicity for now.
+    // TODO: use aligned allocator
+
+    size_t node_size = sizeof(struct ggml_profile_timing) * GGML_MAX_N_THREADS;
+    size_t pvec_size = sizeof(std::intptr_t) * cg->n_nodes;
+    size_t time_size = node_size * cg->n_nodes;
+    size_t t_size    = pvec_size + time_size + sizeof(ggml_profile_output) + sizeof(ggml_profile_data);
+
+    uint8_t * ptr = (uint8_t *) malloc(t_size);
+    if (!ptr) {
+        fprintf(stderr, "ggml-profile: failed to allocate profiling data : n_threads %d n_nodes %d\n", n_threads, cg->n_nodes);
+        return;
+    }
+    memset(ptr, 0, t_size);
+
+    // init all pointers
+    cg->prof         = (ggml_profile_data *)    ptr; ptr += sizeof(ggml_profile_data);
+    cg->prof->output = (ggml_profile_output *)  ptr; ptr += sizeof(ggml_profile_output);
+    cg->prof->timing = (ggml_profile_timing **) ptr; ptr += pvec_size;
+    for (int i=0; i < cg->n_nodes; i++) {
+        cg->prof->timing[i] = (struct ggml_profile_timing *) ptr; ptr += node_size;
+    }
+
+    // init the output
+    ggml_profile_output *out = cg->prof->output;
+    if (!strcmp("stderr", env) || !strcmp("1", env)) {
+        out->prefix = "ggml-profile:";
+        out->stream = stderr;
+    } else {
+        out->prefix = "";
+        out->stream = fopen(env, "w");
+    }
+
+}
+
+extern "C" void ggml_graph_profile_start(struct ggml_cgraph *cg, int n_threads)
+{
+    if (!cg->prof) { ggml_graph_profile_init(cg, n_threads); }
+    if (!cg->prof) { return; }
+}
+
+static inline int ggml_profile_format_tensor_dims(char *str, struct ggml_tensor *t)
+{
+    return sprintf(str, "%ld:%ld:%ld:%ld",
+        (long) t->ne[0], (long) t->ne[1], (long) t->ne[2], (long) t->ne[3]);
+}
+
+static inline void ggml_profile_format_op_dims(char *str, struct ggml_tensor *t)
+{
+    char *p = str;
+
+    // append src0 and src1 (if any)
+    if (t->src[0]) {
+       p += ggml_profile_format_tensor_dims(p, t->src[0]);
+
+       for (int i = 1; i < GGML_MAX_SRC && t->src[i]; i++) {
+           p += sprintf(p, " x ");
+           p += ggml_profile_format_tensor_dims(p, t->src[i]);
+       }
+
+       p += sprintf(p, " -> ");
+    }
+
+    // format self dims separately for better visual alignment
+    char self[64];
+    ggml_profile_format_tensor_dims(self, t);
+
+    p += sprintf(p, "%12s", self);
+}
+
+static inline void ggml_profile_format_op_types(char *str, struct ggml_tensor *t)
+{
+    char *p = str;
+
+    // append src0 and src1 (if any)
+    if (t->src[0]) {
+       p += sprintf(p, "%s", ggml_type_name(t->src[0]->type));
+
+       for (int i = 1; i < GGML_MAX_SRC && t->src[i]; i++) {
+           p += sprintf(p, " x ");
+           p += sprintf(p, "%s", ggml_type_name(t->src[i]->type));
+       }
+
+       p += sprintf(p, " -> ");
+    }
+
+    p += sprintf(p, "%3s", ggml_type_name(t->type));
+}
+
+static inline void ggml_profile_format_op_names(char *str, const struct ggml_tensor *t)
+{
+    char *p = str;
+
+    // append src0 and src1 (if any)
+    if (t->src[0]) {
+       p += sprintf(p, "%s", t->src[0]->name);
+
+       for (int i = 1; i < GGML_MAX_SRC && t->src[i]; i++) {
+           p += sprintf(p, " x ");
+           p += sprintf(p, "%s", t->src[i]->name);
+       }
+
+       p += sprintf(p, " -> ");
+    }
+
+    p += sprintf(p, "%s", t->name);
+}
+
+
+extern "C" void ggml_graph_profile_finish(struct ggml_cgraph *cg, int n_threads)
+{
+    if (!cg->prof) { return; }
+
+    ggml_profile_output *out = cg->prof->output;
+
+    fprintf(out->stream, "%s| node idx | op name | proc (nsec) | sync (nsec) | total (nsec) | op dims | op types | tensor names |\n", out->prefix);
+    fprintf(out->stream, "%s| -------: | :------ | ----------: | ----------: | -----------: | ------: | -------: | -----------: |\n", out->prefix);
+
+    char dims[64 * GGML_MAX_SRC];
+    char types[16 * GGML_MAX_SRC];
+    char names[128 * GGML_MAX_SRC];
+
+    for (int i = 0; i < cg->n_nodes; i++) {
+        uint64_t p_nsec = 0;
+        uint64_t s_nsec = 0;
+        uint64_t t_nsec = 0;
+
+        // add up per thread counters and reset them
+        for (int t=0; t < n_threads; t++) {
+            ggml_profile_timing &timing = cg->prof->timing[i][t];
+
+            p_nsec += timing.nsec[GGML_PROF_OP_SYNC] - timing.nsec[GGML_PROF_OP_START];
+            s_nsec += timing.nsec[GGML_PROF_OP_END]  - timing.nsec[GGML_PROF_OP_SYNC];
+            t_nsec += timing.nsec[GGML_PROF_OP_END]  - timing.nsec[GGML_PROF_OP_START];
+
+            timing.nsec[GGML_PROF_OP_START] = 0;
+            timing.nsec[GGML_PROF_OP_SYNC]  = 0;
+            timing.nsec[GGML_PROF_OP_END]   = 0;
+        }
+
+        ggml_profile_format_op_dims(dims, cg->nodes[i]);
+        ggml_profile_format_op_types(types, cg->nodes[i]);
+        ggml_profile_format_op_names(names, cg->nodes[i]);
+
+        fprintf(out->stream, "%s| %04d | %10s | %10lu | %10lu | %10lu | %46s | %22s | %20s |\n", out->prefix,
+            i, ggml_op_name(cg->nodes[i]->op),
+            (unsigned long) p_nsec, (unsigned long) s_nsec, (unsigned long) t_nsec,
+            dims, types, names);
+    }
+    fprintf(out->stream, "%s   \n", out->prefix); // empty line to split tables
+}
+
+extern "C" void ggml_graph_profile_free(struct ggml_cgraph *cg)
+{
+    if (!cg->prof) { return; }
+
+    ggml_profile_output *out = cg->prof->output;
+    if (out->stream != stderr) {
+        fclose(out->stream);
+    }
+
+    free(cg->prof); cg->prof = nullptr;
+}
+
+extern "C" void ggml_graph_profile_event(const struct ggml_cgraph *cg, enum ggml_profile_event e, int node_n, int ith)
+{
+    if (!cg->prof) { return; }
+
+    using clock = std::chrono::high_resolution_clock;
+
+    ggml_profile_timing &timing = cg->prof->timing[node_n][ith];
+    timing.nsec[e] = std::chrono::nanoseconds(clock::now().time_since_epoch()).count();
+}
+
+#endif // GGML_GRAPH_PROFILER

ggml/src/ggml-profile.h

@@ -0,0 +1,90 @@
+#pragma once
+
+#include "ggml-impl.h"
+
+// GGML internal header
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+// op profile events & timing (per op / per thread)
+enum ggml_profile_event {
+    GGML_PROF_OP_START,
+    GGML_PROF_OP_SYNC,
+    GGML_PROF_OP_END
+};
+
+struct ggml_profile_timing {
+    uint64_t nsec[GGML_PROF_OP_END + 1]; // event times in nsec
+};
+
+struct ggml_profile_output;
+
+struct ggml_profile_data {
+    struct ggml_profile_output *output;
+    struct ggml_profile_timing ** timing; // per op / per thread timing
+};
+
+// check if profiling is enabled for this graph
+static inline bool ggml_graph_profile_enabled(const struct ggml_cgraph *cg)
+{
+    return cg->prof != NULL;
+}
+
+// get pointer to the timing data for specific node / thread
+// can be used by the backends to populate data collected internally
+static inline struct ggml_profile_timing * ggml_graph_profile_timing(const struct ggml_cgraph *cg, int node_n, int ith)
+{
+    if (!cg->prof) { return NULL; }
+    return &cg->prof->timing[node_n][ith];
+}
+
+#ifndef GGML_GRAPH_PROFILER
+
+// Stub out all profiler functions
+
+static inline void ggml_graph_profile_init(struct ggml_cgraph *cg, int n_threads)
+{
+    GGML_UNUSED(cg);
+    GGML_UNUSED(n_threads);
+}
+
+static inline void ggml_graph_profile_start(struct ggml_cgraph *cg, int n_threads)
+{
+    GGML_UNUSED(cg);
+    GGML_UNUSED(n_threads);
+}
+
+static inline void ggml_graph_profile_finish(struct ggml_cgraph *cg, int n_threads)
+{
+    GGML_UNUSED(cg);
+    GGML_UNUSED(n_threads);
+}
+
+static inline void ggml_graph_profile_free(struct ggml_cgraph *cg)
+{
+    GGML_UNUSED(cg);
+}
+
+static inline void ggml_graph_profile_event(const struct ggml_cgraph *cg, enum ggml_profile_event e, int node_n, int ith)
+{
+    GGML_UNUSED(cg);
+    GGML_UNUSED(e);
+    GGML_UNUSED(node_n);
+    GGML_UNUSED(ith);
+}
+
+#else
+
+void ggml_graph_profile_init(struct ggml_cgraph *cg, int n_threads);
+void ggml_graph_profile_start(struct ggml_cgraph *cg, int n_threads);
+void ggml_graph_profile_finish(struct ggml_cgraph *cg, int n_threads);
+void ggml_graph_profile_free(struct ggml_cgraph *cg);
+void ggml_graph_profile_event(const struct ggml_cgraph *cg, enum ggml_profile_event e, int node_n, int ith);
+
+#endif // GGML_GRAPH_PROFILER
+
+#ifdef __cplusplus
+}
+#endif

ggml/src/ggml.c

@@ -9,6 +9,7 @@
 
 // FIXME: required here for quantization functions
 #include "ggml-quants.h"
+#include "ggml-profile.h"
 
 #ifdef GGML_USE_CPU_HBM
 #include <hbwmalloc.h>
@@ -6957,6 +6958,7 @@ struct ggml_cgraph * ggml_new_graph_custom(struct ggml_context * ctx, size_t siz
         /*.grad_accs    =*/ grad_accs_ptr,
         /*.leafs        =*/ leafs_ptr,
         /*.use_counts   =*/ use_counts_ptr,
+        /*.prof         =*/ NULL,
         /*.hash_table   =*/ { hash_size, hash_used, hash_keys_ptr },
         /*.order        =*/ GGML_CGRAPH_EVAL_ORDER_LEFT_TO_RIGHT,
     };
@@ -6984,6 +6986,7 @@ struct ggml_cgraph ggml_graph_view(struct ggml_cgraph * cgraph0, int i0, int i1)
         /*.grad_accs        =*/ NULL,
         /*.leafs            =*/ NULL,
         /*.use_counts       =*/ cgraph0->use_counts,
+        /*.prof             =*/ NULL,
         /*.visited_hash_set =*/ cgraph0->visited_hash_set,
         /*.order            =*/ cgraph0->order,
     };

gguf-py/gguf/utility.py

@@ -288,7 +288,7 @@ class LocalTensor:
     data_range: LocalTensorRange
 
     def mmap_bytes(self) -> np.ndarray:
-        return np.memmap(self.data_range.filename, mode='c', offset=self.data_range.offset, shape=self.data_range.size)
+        return np.memmap(self.data_range.filename, mode='r', offset=self.data_range.offset, shape=self.data_range.size)
 
 
 class SafetensorsLocal:

src/llama-mmap.cpp

@@ -13,10 +13,9 @@
 #ifdef __has_include
     #if __has_include(<unistd.h>)
         #include <unistd.h>
-        #include <fcntl.h>
-        #include <sys/stat.h>
         #if defined(_POSIX_MAPPED_FILES)
             #include <sys/mman.h>
+            #include <fcntl.h>
         #endif
         #if defined(_POSIX_MEMLOCK_RANGE)
             #include <sys/resource.h>
@@ -75,7 +74,7 @@ struct llama_file::impl {
         return ret;
     }
 
-    impl(const char * fname, const char * mode, [[maybe_unused]] const bool use_direct_io = false) {
+    impl(const char * fname, const char * mode) {
         fp = ggml_fopen(fname, mode);
         if (fp == NULL) {
             throw std::runtime_error(format("failed to open %s: %s", fname, strerror(errno)));
@@ -154,40 +153,13 @@ struct llama_file::impl {
         write_raw(&val, sizeof(val));
     }
 
-    void read_aligned_chunk(size_t offset, void * dest, size_t size) const {
-        throw std::runtime_error("DirectIO is not implemented on Windows.");
-    }
-
     ~impl() {
         if (fp) {
             std::fclose(fp);
         }
     }
 #else
-    impl(const char * fname, const char * mode, [[maybe_unused]] const bool use_direct_io = false) {
-#ifdef __linux__
-        // Try unbuffered I/O for read only
-        if (use_direct_io && std::strcmp(mode, "rb") == 0) {
-            fd = open(fname, O_RDONLY | O_DIRECT);
-
-            if (fd != -1) {
-                struct stat file_stats{};
-                fstat(fd, &file_stats);
-
-                size = file_stats.st_size;
-                alignment = file_stats.st_blksize;
-
-                off_t ret = lseek(fd, 0, SEEK_SET);
-                if (ret == -1) {
-                    throw std::runtime_error(format("seek error: %s", strerror(errno)));
-                }
-                return;
-            }
-
-            LLAMA_LOG_WARN("Failed to open model %s with error: %s. Falling back to buffered I/O",
-                fname, strerror(errno));
-        }
-#endif
+    impl(const char * fname, const char * mode) {
         fp = ggml_fopen(fname, mode);
         if (fp == NULL) {
             throw std::runtime_error(format("failed to open %s: %s", fname, strerror(errno)));
@@ -198,30 +170,27 @@ struct llama_file::impl {
     }
 
     size_t tell() const {
-        if (fd == -1) {
-            long ret = std::ftell(fp);
-            if (ret == -1) {
-                throw std::runtime_error(format("ftell error: %s", strerror(errno)));
-            }
-
-            return (size_t) ret;
+// TODO: this ifdef is never true?
+#ifdef _WIN32
+        __int64 ret = _ftelli64(fp);
+#else
+        long ret = std::ftell(fp);
+#endif
+        if (ret == -1) {
+            throw std::runtime_error(format("ftell error: %s", strerror(errno)));
         }
 
-        off_t pos = lseek(fd, 0, SEEK_CUR);
-        if (pos == -1) {
-            throw std::runtime_error(format("lseek error: %s", strerror(errno)));
-        }
-        return (size_t) pos;
+        return (size_t) ret;
     }
 
     void seek(size_t offset, int whence) const {
-        off_t ret = 0;
-        if (fd == -1) {
-            ret = std::fseek(fp, (long) offset, whence);
-        } else {
-            ret = lseek(fd, offset, whence);
-        }
-        if (ret == -1) {
+// TODO: this ifdef is never true?
+#ifdef _WIN32
+        int ret = _fseeki64(fp, (__int64) offset, whence);
+#else
+        int ret = std::fseek(fp, (long) offset, whence);
+#endif
+        if (ret != 0) {
             throw std::runtime_error(format("seek error: %s", strerror(errno)));
         }
     }
@@ -231,55 +200,13 @@ struct llama_file::impl {
             return;
         }
         errno = 0;
-        if (fd == -1) {
-            std::size_t ret = std::fread(ptr, len, 1, fp);
-            if (ferror(fp)) {
-                throw std::runtime_error(format("read error: %s", strerror(errno)));
-            }
-            if (ret != 1) {
-                throw std::runtime_error("unexpectedly reached end of file");
-            }
-        } else {
-            bool successful = false;
-            while (!successful) {
-                off_t ret = read(fd, ptr, len);
-
-                if (ret == -1) {
-                    if (errno == EINTR) {
-                        continue;  // Interrupted by signal, retry
-                    }
-                    throw std::runtime_error(format("read error: %s", strerror(errno)));
-                }
-                if (ret == 0) {
-                    throw std::runtime_error("unexpectedly reached end of file");
-                }
-
-                successful = true;
-            }
+        std::size_t ret = std::fread(ptr, len, 1, fp);
+        if (ferror(fp)) {
+            throw std::runtime_error(format("read error: %s", strerror(errno)));
         }
-    }
-
-    void read_aligned_chunk(size_t offset, void * dest, size_t size) const {
-        off_t aligned_offset = offset & ~(alignment - 1);
-        off_t offset_from_alignment = offset - aligned_offset;
-        size_t bytes_to_read = (offset_from_alignment + size + alignment - 1) & ~(alignment - 1);
-
-        void * raw_buffer = nullptr;
-        int ret = posix_memalign(&raw_buffer, alignment, bytes_to_read);
-        if (ret != 0) {
-            throw std::runtime_error(format("posix_memalign failed with error %d", ret));
+        if (ret != 1) {
+            throw std::runtime_error("unexpectedly reached end of file");
         }
-
-        struct aligned_buffer_deleter {
-            void operator()(void * p) const { free(p); }
-        };
-        std::unique_ptr<void, aligned_buffer_deleter> buffer(raw_buffer);
-
-        seek(aligned_offset, SEEK_SET);
-        read_raw(buffer.get(), bytes_to_read);
-
-        uintptr_t actual_data = reinterpret_cast<uintptr_t>(buffer.get()) + offset_from_alignment;
-        memcpy(dest, reinterpret_cast<void *>(actual_data), size);
     }
 
     uint32_t read_u32() const {
@@ -304,43 +231,22 @@ struct llama_file::impl {
     }
 
     ~impl() {
-        if (fd != -1) {
-            close(fd);
-        } else {
+        if (fp) {
             std::fclose(fp);
         }
     }
-    int fd = -1;
 #endif
 
-    void read_raw_at(void * ptr, size_t len, size_t offset) const {
-        if (alignment != 1) {
-            read_aligned_chunk(offset, ptr, len);
-        } else {
-            seek(offset, SEEK_SET);
-            read_raw(ptr, len);
-        }
-    }
-
-    size_t read_alignment() const {
-        return alignment;
-    }
-
-    size_t alignment = 1;
-
-    FILE * fp{};
-    size_t size{};
+    FILE * fp;
+    size_t size;
 };
 
-llama_file::llama_file(const char * fname, const char * mode, const bool use_direct_io) :
-    pimpl(std::make_unique<impl>(fname, mode, use_direct_io)) {}
+llama_file::llama_file(const char * fname, const char * mode) : pimpl(std::make_unique<impl>(fname, mode)) {}
 llama_file::~llama_file() = default;
 
 size_t llama_file::tell() const { return pimpl->tell(); }
 size_t llama_file::size() const { return pimpl->size; }
 
-size_t llama_file::read_alignment() const { return pimpl->read_alignment(); }
-
 int llama_file::file_id() const {
 #ifdef _WIN32
     return _fileno(pimpl->fp);
@@ -355,7 +261,6 @@ int llama_file::file_id() const {
 
 void llama_file::seek(size_t offset, int whence) const { pimpl->seek(offset, whence); }
 void llama_file::read_raw(void * ptr, size_t len) const { pimpl->read_raw(ptr, len); }
-void llama_file::read_raw_at(void * ptr, size_t len, size_t offset) const { pimpl->read_raw_at(ptr, len, offset); }
 
 uint32_t llama_file::read_u32() const { return pimpl->read_u32(); }
 

src/llama-mmap.h

@@ -3,7 +3,6 @@
 #include <cstdint>
 #include <memory>
 #include <vector>
-#include <cstdio>
 
 struct llama_file;
 struct llama_mmap;
@@ -14,7 +13,7 @@ using llama_mmaps  = std::vector<std::unique_ptr<llama_mmap>>;
 using llama_mlocks = std::vector<std::unique_ptr<llama_mlock>>;
 
 struct llama_file {
-    llama_file(const char * fname, const char * mode, bool use_direct_io = false);
+    llama_file(const char * fname, const char * mode);
     ~llama_file();
 
     size_t tell() const;
@@ -25,14 +24,11 @@ struct llama_file {
     void seek(size_t offset, int whence) const;
 
     void read_raw(void * ptr, size_t len) const;
-    void read_raw_at(void * ptr, size_t len, size_t offset) const;
-    void read_aligned_chunk(size_t offset, void * dest, size_t size) const;
     uint32_t read_u32() const;
 
     void write_raw(const void * ptr, size_t len) const;
     void write_u32(uint32_t val) const;
 
-    size_t read_alignment() const;
 private:
     struct impl;
     std::unique_ptr<impl> pimpl;

src/llama-model-loader.cpp

@@ -504,7 +504,7 @@ llama_model_loader::llama_model_loader(
     get_key(llm_kv(LLM_KV_GENERAL_ARCHITECTURE), arch_name, false);
     llm_kv = LLM_KV(llm_arch_from_string(arch_name));
 
-    files.emplace_back(new llama_file(fname.c_str(), "rb", !use_mmap));
+    files.emplace_back(new llama_file(fname.c_str(), "rb"));
     contexts.emplace_back(ctx);
 
     // Save tensors data offset of the main file.
@@ -572,7 +572,7 @@ llama_model_loader::llama_model_loader(
                 }
             }
 
-            files.emplace_back(new llama_file(fname_split, "rb", !use_mmap));
+            files.emplace_back(new llama_file(fname_split, "rb"));
             contexts.emplace_back(ctx);
 
             // Save tensors data offset info of the shard.
@@ -935,15 +935,7 @@ bool llama_model_loader::load_all_data(
     // 4 staging buffers for async uploads, each sized 1MB seems to be a good default for single NVMe drives.
     // NVMe raid configurations might require more / larger buffers.
     constexpr size_t n_buffers = 4;
-
-    size_t alignment = 1;
-    for (const auto & file : files) {
-        alignment = std::max(file->read_alignment(), alignment);
-    }
-
-    // Buffer size: balance between memory usage and I/O efficiency
-    // 64MB works well for NVMe drives
-    const size_t buffer_size = alignment != 1 ? 64 * 1024 * 1024 + 2 * alignment : 1 * 1024 * 1024;
+    constexpr size_t buffer_size = 1 * 1024 * 1024; // 1MB
 
     std::vector<ggml_backend_buffer_t> host_buffers;
     std::vector<ggml_backend_event_t> events;
@@ -993,7 +985,6 @@ bool llama_model_loader::load_all_data(
         // If the backend is supported, create pinned memory buffers and events for synchronisation.
         for (size_t idx = 0; idx < n_buffers; ++idx) {
             auto * buf = ggml_backend_buft_alloc_buffer(host_buft, buffer_size);
-
             if (!buf) {
                 LLAMA_LOG_DEBUG("%s: failed to allocate host buffer for async uploads for device %s\n", func,
                     ggml_backend_dev_name(dev));
@@ -1075,9 +1066,9 @@ bool llama_model_loader::load_all_data(
             }
         } else {
             const auto & file = files.at(weight->idx);
-
             if (ggml_backend_buffer_is_host(cur->buffer)) {
-                file->read_raw_at(cur->data, n_size, weight->offs);
+                file->seek(weight->offs, SEEK_SET);
+                file->read_raw(cur->data, n_size);
                 if (check_tensors) {
                     validation_result.emplace_back(std::async(std::launch::async, [cur, n_size] {
                         return std::make_pair(cur, ggml_validate_row_data(cur->type, cur->data, n_size));
@@ -1086,60 +1077,26 @@ bool llama_model_loader::load_all_data(
             } else {
                 // If upload_backend is valid load the tensor in chunks to pinned memory and upload the buffers asynchronously to the GPU.
                 if (upload_backend) {
-                    auto offset = (off_t) weight->offs;
-                    alignment = file->read_alignment();
-                    off_t aligned_offset = offset & ~(alignment - 1);
-                    off_t offset_from_alignment = offset - aligned_offset;
-                    file->seek(aligned_offset, SEEK_SET);
-
-                    // Calculate aligned read boundaries
-                    size_t read_start = aligned_offset;
-                    size_t read_end = (offset + n_size + alignment - 1) & ~(alignment - 1);
+                    file->seek(weight->offs, SEEK_SET);
 
                     size_t bytes_read = 0;
-                    size_t data_read = 0;  // Actual tensor data copied (excluding padding)
 
-                    while (bytes_read < read_end - read_start) {
-                        size_t read_size = std::min<size_t>(buffer_size, read_end - read_start - bytes_read);
+                    while (bytes_read < n_size) {
+                        size_t read_iteration = std::min<size_t>(buffer_size, n_size - bytes_read);
 
-                        // Align the destination pointer within the pinned buffer
-                        uintptr_t ptr_dest_aligned = (reinterpret_cast<uintptr_t>(host_ptrs[buffer_idx]) + alignment - 1) & ~(alignment - 1);
-
-                        // Wait for previous upload to complete before reusing buffer
                         ggml_backend_event_synchronize(events[buffer_idx]);
-
-                        // Read aligned chunk from file
-                        file->read_raw(reinterpret_cast<void *>(ptr_dest_aligned), read_size);
-
-                        // Calculate actual data portion (excluding alignment padding)
-                        uintptr_t ptr_data = ptr_dest_aligned;
-                        size_t data_to_copy = read_size;
-
-                        // Skip alignment padding at start of first chunk
-                        if (bytes_read == 0) {
-                            ptr_data += offset_from_alignment;
-                            data_to_copy -= offset_from_alignment;
-                        }
-
-                        // Trim alignment padding at end of last chunk
-                        if (aligned_offset + bytes_read + read_size > offset + n_size) {
-                            data_to_copy -= (read_end - (offset + n_size));
-                        }
-
-                        // Async upload actual data to GPU
-                        ggml_backend_tensor_set_async(upload_backend, cur,
-                                                      reinterpret_cast<void *>(ptr_data), data_read, data_to_copy);
+                        file->read_raw(host_ptrs[buffer_idx], read_iteration);
+                        ggml_backend_tensor_set_async(upload_backend, cur, host_ptrs[buffer_idx], bytes_read, read_iteration);
                         ggml_backend_event_record(events[buffer_idx], upload_backend);
 
-                        data_read += data_to_copy;
-                        bytes_read += read_size;
-
+                        bytes_read += read_iteration;
                         ++buffer_idx;
                         buffer_idx %= n_buffers;
                     }
                 } else {
                     read_buf.resize(n_size);
-                    file->read_raw_at(read_buf.data(), n_size, weight->offs);
+                    file->seek(weight->offs, SEEK_SET);
+                    file->read_raw(read_buf.data(), n_size);
                     ggml_backend_tensor_set(cur, read_buf.data(), 0, n_size);
                     if (check_tensors && !ggml_validate_row_data(cur->type, read_buf.data(), n_size)) {
                         throw std::runtime_error(format("tensor '%s' has invalid data", ggml_get_name(cur)));

src/llama-model.cpp

@@ -2378,10 +2378,10 @@ bool llama_model::load_tensors(llama_model_loader & ml) {
     if (cpu_dev == nullptr) {
         throw std::runtime_error(format("%s: no CPU backend found", __func__));
     }
-    const int i_gpu_start = std::max(int(hparams.n_layer) + 1 - n_gpu_layers, 0);
-    const int act_gpu_layers = devices.empty() ? 0 : std::min(n_gpu_layers, int(n_layer) + 1);
+    const int i_gpu_start = std::max((int) hparams.n_layer - n_gpu_layers, (int) 0);
+    const int act_gpu_layers = devices.empty() ? 0 : std::min(n_gpu_layers, (int)n_layer + 1);
     auto get_layer_buft_list = [&](int il) -> llama_model::impl::layer_dev {
-        const bool is_swa = il < int(hparams.n_layer) && hparams.is_swa(il);
+        const bool is_swa = il < (int) hparams.n_layer && hparams.is_swa(il);
         if (il < i_gpu_start || (il - i_gpu_start) >= act_gpu_layers) {
             LLAMA_LOG_DEBUG("load_tensors: layer %3d assigned to device %s, is_swa = %d\n", il, ggml_backend_dev_name(cpu_dev), is_swa);
             return {cpu_dev, &pimpl->cpu_buft_list};
@@ -6693,12 +6693,10 @@ bool llama_model::load_tensors(llama_model_loader & ml) {
     if (llama_supports_gpu_offload()) {
         const int n_gpu = std::min(n_gpu_layers, int(hparams.n_layer));
 
-        int n_repeating = n_gpu;
-        if (n_repeating > 0) {
+        LLAMA_LOG_INFO("%s: offloading %d repeating layers to GPU\n", __func__, n_gpu);
+        if (n_gpu_layers > (int) hparams.n_layer) {
             LLAMA_LOG_INFO("%s: offloading output layer to GPU\n", __func__);
-            n_repeating--;
         }
-        LLAMA_LOG_INFO("%s: offloading %d repeating layers to GPU\n", __func__, n_repeating);
 
         const int max_backend_supported_layers = hparams.n_layer + 1;
         const int max_offloadable_layers       = hparams.n_layer + 1;

src/llama.cpp

@@ -292,6 +292,10 @@ static void llama_params_fit_impl(
         if (mparams->split_mode == LLAMA_SPLIT_MODE_ROW) {
             throw std::runtime_error("changing weight allocation for LLAMA_SPLIT_MODE_ROW not implemented, abort");
         }
+        if (hp_ngl < 2*nd) {
+            throw std::runtime_error("model has only " + std::to_string(hp_ngl) + " layers but need at least "
+                + std::to_string(2*nd) + " to fit memory for " + std::to_string(nd) + " devices, abort");
+        }
     }
     if (!tensor_buft_overrides) {
         throw std::runtime_error("did not provide buffer to set tensor_buft_overrides, abort");
@@ -358,7 +362,8 @@ static void llama_params_fit_impl(
     auto set_ngl_tensor_split_tbo = [&](
             const std::vector<ngl_t> & ngl_per_device,
             const std::vector<ggml_backend_buffer_type_t> & overflow_bufts,
-            llama_model_params & mparams) {
+            llama_model_params & mparams,
+            const bool add_nonrepeating) {
         mparams.n_gpu_layers = 0;
         for (size_t id = 0; id < nd; id++) {
             mparams.n_gpu_layers += ngl_per_device[id].n_layer;
@@ -366,9 +371,13 @@ static void llama_params_fit_impl(
                 tensor_split[id] = ngl_per_device[id].n_layer;
             }
         }
-        assert(uint32_t(mparams.n_gpu_layers) <= hp_ngl + 1);
-        uint32_t il0 = hp_ngl + 1 - mparams.n_gpu_layers; // start index for tensor buft overrides
+        assert(uint32_t(mparams.n_gpu_layers) <= hp_ngl);
+        uint32_t il0 = hp_ngl - mparams.n_gpu_layers; // start index for tensor buft overrides
 
+        if (add_nonrepeating) {
+            mparams.n_gpu_layers += 1;
+            tensor_split[nd - 1] += 1;
+        }
         mparams.tensor_split = tensor_split;
 
         size_t itbo = 0;
@@ -399,9 +408,10 @@ static void llama_params_fit_impl(
     auto get_memory_for_layers = [&](
             const char * func_name,
             const std::vector<ngl_t> & ngl_per_device,
-            const std::vector<ggml_backend_buffer_type_t> & overflow_bufts) -> std::vector<int64_t> {
+            const std::vector<ggml_backend_buffer_type_t> & overflow_bufts,
+            const bool add_nonrepeating) -> std::vector<int64_t> {
         llama_model_params mparams_copy = *mparams;
-        set_ngl_tensor_split_tbo(ngl_per_device, overflow_bufts, mparams_copy);
+        set_ngl_tensor_split_tbo(ngl_per_device, overflow_bufts, mparams_copy, add_nonrepeating);
 
         const dmds_t dmd_nl = llama_get_device_memory_data(
             path_model, &mparams_copy, cparams, devs, hp_ngl, hp_nct, hp_nex, log_level);
@@ -459,6 +469,9 @@ static void llama_params_fit_impl(
         LLAMA_LOG_DEBUG("%s: id=%zu, target=%" PRId64 " MiB\n", __func__, id, targets[id]/MiB);
     }
 
+    // whether for the optimal memory use we expect to load at least some MoE tensors:
+    const bool partial_moe = hp_nex > 0 && global_surplus_cpu_moe > 0;
+
     std::vector<ggml_backend_buffer_type_t> overflow_bufts; // which bufts the partial layers of a device overflow to:
     overflow_bufts.reserve(nd);
     for (size_t id = 0; id < nd - 1; ++id) {
@@ -467,7 +480,7 @@ static void llama_params_fit_impl(
     overflow_bufts.push_back(ggml_backend_cpu_buffer_type());
 
     std::vector<ngl_t> ngl_per_device(nd);
-    std::vector<int64_t> mem = get_memory_for_layers(__func__, ngl_per_device, overflow_bufts);
+    std::vector<int64_t> mem = get_memory_for_layers(__func__, ngl_per_device, overflow_bufts, partial_moe);
     if (hp_nex > 0) {
         for (size_t id = 0; id < nd; id++) {
             ngl_per_device[id].overflow_type = LAYER_FRACTION_MOE;
@@ -480,14 +493,13 @@ static void llama_params_fit_impl(
     //   - interpolate the memory use / layer between low and high linearly to get a guess where it meets our target
     //   - check memory use of our guess, replace either the low or high bound
     //   - once we only have a difference of a single layer, stop and return the lower bound that just barely still fits
-    //   - the last device has the output layer, which cannot be a partial layer
     if (hp_nex == 0) {
         LLAMA_LOG_INFO("%s: filling dense layers back-to-front:\n", __func__);
     } else {
         LLAMA_LOG_INFO("%s: filling dense-only layers back-to-front:\n", __func__);
     }
     for (int id = nd - 1; id >= 0; id--) {
-        uint32_t n_unassigned = hp_ngl + 1;
+        uint32_t n_unassigned = hp_ngl;
         for (size_t jd = id + 1; jd < nd; ++jd) {
             assert(n_unassigned >= ngl_per_device[jd].n_layer);
             n_unassigned -= ngl_per_device[jd].n_layer;
@@ -496,10 +508,10 @@ static void llama_params_fit_impl(
         std::vector<ngl_t> ngl_per_device_high = ngl_per_device;
         ngl_per_device_high[id].n_layer = n_unassigned;
         if (hp_nex > 0) {
-            ngl_per_device_high[id].n_part = size_t(id) < nd - 1 ? ngl_per_device_high[id].n_layer : ngl_per_device_high[id].n_layer - 1;
+            ngl_per_device_high[id].n_part = ngl_per_device_high[id].n_layer;
         }
         if (ngl_per_device_high[id].n_layer > 0) {
-            std::vector<int64_t> mem_high = get_memory_for_layers(__func__, ngl_per_device_high, overflow_bufts);
+            std::vector<int64_t> mem_high = get_memory_for_layers(__func__, ngl_per_device_high, overflow_bufts, partial_moe);
             if (mem_high[id] > targets[id]) {
                 assert(ngl_per_device_high[id].n_layer > ngl_per_device[id].n_layer);
                 uint32_t delta = ngl_per_device_high[id].n_layer - ngl_per_device[id].n_layer;
@@ -514,7 +526,7 @@ static void llama_params_fit_impl(
                     if (hp_nex) {
                         ngl_per_device_test[id].n_part += step_size;
                     }
-                    const std::vector<int64_t> mem_test = get_memory_for_layers(__func__, ngl_per_device_test, overflow_bufts);
+                    const std::vector<int64_t> mem_test = get_memory_for_layers(__func__, ngl_per_device_test, overflow_bufts, partial_moe);
 
                     if (mem_test[id] <= targets[id]) {
                         ngl_per_device = ngl_per_device_test;
@@ -541,7 +553,7 @@ static void llama_params_fit_impl(
             __func__, dev_names[id].c_str(), ngl_per_device[id].n_layer, mem[id]/MiB, projected_margin/MiB);
     }
     if (hp_nex == 0 || global_surplus_cpu_moe <= 0) {
-        set_ngl_tensor_split_tbo(ngl_per_device, overflow_bufts, *mparams);
+        set_ngl_tensor_split_tbo(ngl_per_device, overflow_bufts, *mparams, partial_moe);
         return;
     }
 
@@ -564,13 +576,13 @@ static void llama_params_fit_impl(
     for (size_t id = 0; id <= id_dense_start; id++) {
         std::vector<ngl_t> ngl_per_device_high = ngl_per_device;
         for (size_t jd = id_dense_start; jd < nd; jd++) {
-            const uint32_t n_layer_move = jd < nd - 1 ? ngl_per_device_high[jd].n_layer : ngl_per_device_high[jd].n_layer - 1;
+            const uint32_t n_layer_move = ngl_per_device_high[jd].n_layer;
             ngl_per_device_high[id].n_layer += n_layer_move;
             ngl_per_device_high[jd].n_layer -= n_layer_move;
             ngl_per_device_high[jd].n_part = 0;
         }
         size_t id_dense_start_high = nd - 1;
-        std::vector<int64_t> mem_high = get_memory_for_layers(__func__, ngl_per_device_high, overflow_bufts);
+        std::vector<int64_t> mem_high = get_memory_for_layers(__func__, ngl_per_device_high, overflow_bufts, partial_moe);
 
         if (mem_high[id] > targets[id]) {
             assert(ngl_per_device_high[id].n_layer >= ngl_per_device_high[id].n_part);
@@ -598,7 +610,7 @@ static void llama_params_fit_impl(
                         break;
                     }
                 }
-                const std::vector<int64_t> mem_test = get_memory_for_layers(__func__, ngl_per_device_test, overflow_bufts);
+                const std::vector<int64_t> mem_test = get_memory_for_layers(__func__, ngl_per_device_test, overflow_bufts, partial_moe);
 
                 if (mem_test[id] <= targets[id]) {
                     ngl_per_device = ngl_per_device_test;
@@ -625,7 +637,7 @@ static void llama_params_fit_impl(
         }
 
         // try to fit at least part of one more layer
-        if (ngl_per_device[id_dense_start].n_layer > (id < nd - 1 ? 0 : 1)) {
+        if (ngl_per_device[id_dense_start].n_layer > 0) {
             std::vector<ngl_t> ngl_per_device_test = ngl_per_device;
             size_t id_dense_start_test = id_dense_start;
             ngl_per_device_test[id_dense_start_test].n_layer--;
@@ -637,7 +649,7 @@ static void llama_params_fit_impl(
             }
             ngl_per_device_test[id].overflow_type = LAYER_FRACTION_UP;
             LLAMA_LOG_DEBUG("%s: trying to fit one extra layer with overflow_type=LAYER_FRACTION_UP\n", __func__);
-            std::vector<int64_t> mem_test = get_memory_for_layers(__func__, ngl_per_device_test, overflow_bufts);
+            std::vector<int64_t> mem_test = get_memory_for_layers(__func__, ngl_per_device_test, overflow_bufts, partial_moe);
             if (mem_test[id] < targets[id]) {
                 ngl_per_device = ngl_per_device_test;
                 mem            = mem_test;
@@ -647,7 +659,7 @@ static void llama_params_fit_impl(
 
                 ngl_per_device_test[id].overflow_type = LAYER_FRACTION_GATE;
                 LLAMA_LOG_DEBUG("%s: trying to fit one extra layer with overflow_type=LAYER_FRACTION_GATE\n", __func__);
-                mem_test = get_memory_for_layers(__func__, ngl_per_device_test, overflow_bufts);
+                mem_test = get_memory_for_layers(__func__, ngl_per_device_test, overflow_bufts, partial_moe);
                 if (mem_test[id] < targets[id]) {
                     ngl_per_device = ngl_per_device_test;
                     mem            = mem_test;
@@ -658,7 +670,7 @@ static void llama_params_fit_impl(
             } else {
                 ngl_per_device_test[id].overflow_type = LAYER_FRACTION_ATTN;
                 LLAMA_LOG_DEBUG("%s: trying to fit one extra layer with overflow_type=LAYER_FRACTION_ATTN\n", __func__);
-                mem_test = get_memory_for_layers(__func__, ngl_per_device_test, overflow_bufts);
+                mem_test = get_memory_for_layers(__func__, ngl_per_device_test, overflow_bufts, partial_moe);
                 if (mem_test[id] < targets[id]) {
                     ngl_per_device = ngl_per_device_test;
                     mem            = mem_test;
@@ -675,7 +687,7 @@ static void llama_params_fit_impl(
             __func__, dev_names[id].c_str(), ngl_per_device[id].n_layer, ngl_per_device[id].n_part, mem[id]/MiB, projected_margin/MiB);
     }
 
-    set_ngl_tensor_split_tbo(ngl_per_device, overflow_bufts, *mparams);
+    set_ngl_tensor_split_tbo(ngl_per_device, overflow_bufts, *mparams, partial_moe);
 }
 
 bool llama_params_fit(

tools/server/webui/package-lock.json

@@ -2109,9 +2109,9 @@
 			}
 		},
 		"node_modules/@sveltejs/kit": {
-			"version": "2.49.2",
-			"resolved": "https://registry.npmjs.org/@sveltejs/kit/-/kit-2.49.2.tgz",
-			"integrity": "sha512-Vp3zX/qlwerQmHMP6x0Ry1oY7eKKRcOWGc2P59srOp4zcqyn+etJyQpELgOi4+ZSUgteX8Y387NuwruLgGXLUQ==",
+			"version": "2.48.5",
+			"resolved": "https://registry.npmjs.org/@sveltejs/kit/-/kit-2.48.5.tgz",
+			"integrity": "sha512-/rnwfSWS3qwUSzvHynUTORF9xSJi7PCR9yXkxUOnRrNqyKmCmh3FPHH+E9BbgqxXfTevGXBqgnlh9kMb+9T5XA==",
 			"dev": true,
 			"license": "MIT",
 			"dependencies": {
@@ -5797,9 +5797,9 @@
 			}
 		},
 		"node_modules/mdast-util-to-hast": {
-			"version": "13.2.1",
-			"resolved": "https://registry.npmjs.org/mdast-util-to-hast/-/mdast-util-to-hast-13.2.1.tgz",
-			"integrity": "sha512-cctsq2wp5vTsLIcaymblUriiTcZd0CwWtCbLvrOzYCDZoWyMNV8sZ7krj09FSnsiJi3WVsHLM4k6Dq/yaPyCXA==",
+			"version": "13.2.0",
+			"resolved": "https://registry.npmjs.org/mdast-util-to-hast/-/mdast-util-to-hast-13.2.0.tgz",
+			"integrity": "sha512-QGYKEuUsYT9ykKBCMOEDLsU5JRObWQusAolFMeko/tYPufNkRffBAQjIE+99jbA87xv6FgmjLtwjh9wBWajwAA==",
 			"license": "MIT",
 			"dependencies": {
 				"@types/hast": "^3.0.0",

tools/server/webui/src/app.d.ts

@@ -124,10 +124,3 @@ declare global {
 		SettingsConfigType
 	};
 }
-
-declare global {
-	interface Window {
-		idxThemeStyle?: number;
-		idxCodeBlock?: number;
-	}
-}

tools/server/webui/src/lib/components/app/chat/ChatScreen/ChatScreen.svelte

@@ -587,7 +587,7 @@
 
 		&::after {
 			content: '';
-			position: absolute;
+			position: fixed;
 			bottom: 0;
 			z-index: -1;
 			left: 0;

tools/server/webui/src/lib/components/app/misc/MarkdownContent.svelte

@@ -7,19 +7,15 @@
 	import remarkRehype from 'remark-rehype';
 	import rehypeKatex from 'rehype-katex';
 	import rehypeStringify from 'rehype-stringify';
-	import type { Root as HastRoot, RootContent as HastRootContent } from 'hast';
-	import type { Root as MdastRoot } from 'mdast';
-	import { browser } from '$app/environment';
-	import { onDestroy, tick } from 'svelte';
-	import { rehypeRestoreTableHtml } from '$lib/markdown/table-html-restorer';
-	import { rehypeEnhanceLinks } from '$lib/markdown/enhance-links';
-	import { rehypeEnhanceCodeBlocks } from '$lib/markdown/enhance-code-blocks';
-	import { remarkLiteralHtml } from '$lib/markdown/literal-html';
 	import { copyCodeToClipboard, preprocessLaTeX } from '$lib/utils';
+	import { rehypeRestoreTableHtml } from '$lib/markdown/table-html-restorer';
+	import { browser } from '$app/environment';
 	import '$styles/katex-custom.scss';
+
 	import githubDarkCss from 'highlight.js/styles/github-dark.css?inline';
 	import githubLightCss from 'highlight.js/styles/github.css?inline';
 	import { mode } from 'mode-watcher';
+	import { remarkLiteralHtml } from '$lib/markdown/literal-html';
 	import CodePreviewDialog from './CodePreviewDialog.svelte';
 
 	interface Props {
@@ -27,24 +23,33 @@
 		class?: string;
 	}
 
-	interface MarkdownBlock {
-		id: string;
-		html: string;
-	}
-
 	let { content, class: className = '' }: Props = $props();
 
 	let containerRef = $state<HTMLDivElement>();
-	let renderedBlocks = $state<MarkdownBlock[]>([]);
-	let unstableBlockHtml = $state('');
+	let processedHtml = $state('');
 	let previewDialogOpen = $state(false);
 	let previewCode = $state('');
 	let previewLanguage = $state('text');
 
-	let pendingMarkdown: string | null = null;
-	let isProcessing = false;
+	function loadHighlightTheme(isDark: boolean) {
+		if (!browser) return;
 
-	const themeStyleId = `highlight-theme-${(window.idxThemeStyle = (window.idxThemeStyle ?? 0) + 1)}`;
+		const existingThemes = document.querySelectorAll('style[data-highlight-theme]');
+		existingThemes.forEach((style) => style.remove());
+
+		const style = document.createElement('style');
+		style.setAttribute('data-highlight-theme', 'true');
+		style.textContent = isDark ? githubDarkCss : githubLightCss;
+
+		document.head.appendChild(style);
+	}
+
+	$effect(() => {
+		const currentMode = mode.current;
+		const isDark = currentMode === 'dark';
+
+		loadHighlightTheme(isDark);
+	});
 
 	let processor = $derived(() => {
 		return remark()
@@ -56,64 +61,139 @@
 			.use(rehypeKatex) // Render math using KaTeX
 			.use(rehypeHighlight) // Add syntax highlighting
 			.use(rehypeRestoreTableHtml) // Restore limited HTML (e.g., <br>, <ul>) inside Markdown tables
-			.use(rehypeEnhanceLinks) // Add target="_blank" to links
-			.use(rehypeEnhanceCodeBlocks) // Wrap code blocks with header and actions
-			.use(rehypeStringify, { allowDangerousHtml: true }); // Convert to HTML string
+			.use(rehypeStringify); // Convert to HTML string
 	});
 
-	/**
-	 * Removes click event listeners from copy and preview buttons.
-	 * Called on component destroy.
-	 */
-	function cleanupEventListeners() {
-		if (!containerRef) return;
-
-		const copyButtons = containerRef.querySelectorAll<HTMLButtonElement>('.copy-code-btn');
-		const previewButtons = containerRef.querySelectorAll<HTMLButtonElement>('.preview-code-btn');
-
-		for (const button of copyButtons) {
-			button.removeEventListener('click', handleCopyClick);
+	function enhanceLinks(html: string): string {
+		if (!html.includes('<a')) {
+			return html;
 		}
 
-		for (const button of previewButtons) {
-			button.removeEventListener('click', handlePreviewClick);
+		const tempDiv = document.createElement('div');
+		tempDiv.innerHTML = html;
+
+		// Make all links open in new tabs
+		const linkElements = tempDiv.querySelectorAll('a[href]');
+		let mutated = false;
+
+		for (const link of linkElements) {
+			const target = link.getAttribute('target');
+			const rel = link.getAttribute('rel');
+
+			if (target !== '_blank' || rel !== 'noopener noreferrer') {
+				mutated = true;
+			}
+
+			link.setAttribute('target', '_blank');
+			link.setAttribute('rel', 'noopener noreferrer');
+		}
+
+		return mutated ? tempDiv.innerHTML : html;
+	}
+
+	function enhanceCodeBlocks(html: string): string {
+		if (!html.includes('<pre')) {
+			return html;
+		}
+
+		const tempDiv = document.createElement('div');
+		tempDiv.innerHTML = html;
+
+		const preElements = tempDiv.querySelectorAll('pre');
+		let mutated = false;
+
+		for (const [index, pre] of Array.from(preElements).entries()) {
+			const codeElement = pre.querySelector('code');
+
+			if (!codeElement) {
+				continue;
+			}
+
+			mutated = true;
+
+			let language = 'text';
+			const classList = Array.from(codeElement.classList);
+
+			for (const className of classList) {
+				if (className.startsWith('language-')) {
+					language = className.replace('language-', '');
+					break;
+				}
+			}
+
+			const rawCode = codeElement.textContent || '';
+			const codeId = `code-${Date.now()}-${index}`;
+			codeElement.setAttribute('data-code-id', codeId);
+			codeElement.setAttribute('data-raw-code', rawCode);
+
+			const wrapper = document.createElement('div');
+			wrapper.className = 'code-block-wrapper';
+
+			const header = document.createElement('div');
+			header.className = 'code-block-header';
+
+			const languageLabel = document.createElement('span');
+			languageLabel.className = 'code-language';
+			languageLabel.textContent = language;
+
+			const copyButton = document.createElement('button');
+			copyButton.className = 'copy-code-btn';
+			copyButton.setAttribute('data-code-id', codeId);
+			copyButton.setAttribute('title', 'Copy code');
+			copyButton.setAttribute('type', 'button');
+
+			copyButton.innerHTML = `
+                                <svg xmlns="http://www.w3.org/2000/svg" width="16" height="16" viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2" stroke-linecap="round" stroke-linejoin="round" class="lucide lucide-copy-icon lucide-copy"><rect width="14" height="14" x="8" y="8" rx="2" ry="2"/><path d="M4 16c-1.1 0-2-.9-2-2V4c0-1.1.9-2 2-2h10c1.1 0 2 .9 2 2"/></svg>
+                        `;
+
+			const actions = document.createElement('div');
+			actions.className = 'code-block-actions';
+
+			actions.appendChild(copyButton);
+
+			if (language.toLowerCase() === 'html') {
+				const previewButton = document.createElement('button');
+				previewButton.className = 'preview-code-btn';
+				previewButton.setAttribute('data-code-id', codeId);
+				previewButton.setAttribute('title', 'Preview code');
+				previewButton.setAttribute('type', 'button');
+
+				previewButton.innerHTML = `
+                                        <svg xmlns="http://www.w3.org/2000/svg" width="16" height="16" viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2" stroke-linecap="round" stroke-linejoin="round" class="lucide lucide-eye lucide-eye-icon"><path d="M2.062 12.345a1 1 0 0 1 0-.69C3.5 7.73 7.36 5 12 5s8.5 2.73 9.938 6.655a1 1 0 0 1 0 .69C20.5 16.27 16.64 19 12 19s-8.5-2.73-9.938-6.655"/><circle cx="12" cy="12" r="3"/></svg>
+                                `;
+
+				actions.appendChild(previewButton);
+			}
+
+			header.appendChild(languageLabel);
+			header.appendChild(actions);
+			wrapper.appendChild(header);
+
+			const clonedPre = pre.cloneNode(true) as HTMLElement;
+			wrapper.appendChild(clonedPre);
+
+			pre.parentNode?.replaceChild(wrapper, pre);
+		}
+
+		return mutated ? tempDiv.innerHTML : html;
+	}
+
+	async function processMarkdown(text: string): Promise<string> {
+		try {
+			let normalized = preprocessLaTeX(text);
+			const result = await processor().process(normalized);
+			const html = String(result);
+			const enhancedLinks = enhanceLinks(html);
+
+			return enhanceCodeBlocks(enhancedLinks);
+		} catch (error) {
+			console.error('Markdown processing error:', error);
+
+			// Fallback to plain text with line breaks
+			return text.replace(/\n/g, '<br>');
 		}
 	}
 
-	/**
-	 * Removes this component's highlight.js theme style from the document head.
-	 * Called on component destroy to clean up injected styles.
-	 */
-	function cleanupHighlightTheme() {
-		if (!browser) return;
-
-		const existingTheme = document.getElementById(themeStyleId);
-		existingTheme?.remove();
-	}
-
-	/**
-	 * Loads the appropriate highlight.js theme based on dark/light mode.
-	 * Injects a scoped style element into the document head.
-	 * @param isDark - Whether to load the dark theme (true) or light theme (false)
-	 */
-	function loadHighlightTheme(isDark: boolean) {
-		if (!browser) return;
-
-		const existingTheme = document.getElementById(themeStyleId);
-		existingTheme?.remove();
-
-		const style = document.createElement('style');
-		style.id = themeStyleId;
-		style.textContent = isDark ? githubDarkCss : githubLightCss;
-
-		document.head.appendChild(style);
-	}
-
-	/**
-	 * Extracts code information from a button click target within a code block.
-	 * @param target - The clicked button element
-	 * @returns Object with rawCode and language, or null if extraction fails
-	 */
 	function getCodeInfoFromTarget(target: HTMLElement) {
 		const wrapper = target.closest('.code-block-wrapper');
 
@@ -129,7 +209,12 @@
 			return null;
 		}
 
-		const rawCode = codeElement.textContent ?? '';
+		const rawCode = codeElement.getAttribute('data-raw-code');
+
+		if (rawCode === null) {
+			console.error('No raw code found');
+			return null;
+		}
 
 		const languageLabel = wrapper.querySelector<HTMLElement>('.code-language');
 		const language = languageLabel?.textContent?.trim() || 'text';
@@ -137,28 +222,6 @@
 		return { rawCode, language };
 	}
 
-	/**
-	 * Generates a unique identifier for a HAST node based on its position.
-	 * Used for stable block identification during incremental rendering.
-	 * @param node - The HAST root content node
-	 * @param indexFallback - Fallback index if position is unavailable
-	 * @returns Unique string identifier for the node
-	 */
-	function getHastNodeId(node: HastRootContent, indexFallback: number): string {
-		const position = node.position;
-
-		if (position?.start?.offset != null && position?.end?.offset != null) {
-			return `hast-${position.start.offset}-${position.end.offset}`;
-		}
-
-		return `${node.type}-${indexFallback}`;
-	}
-
-	/**
-	 * Handles click events on copy buttons within code blocks.
-	 * Copies the raw code content to the clipboard.
-	 * @param event - The click event from the copy button
-	 */
 	async function handleCopyClick(event: Event) {
 		event.preventDefault();
 		event.stopPropagation();
@@ -182,25 +245,6 @@
 		}
 	}
 
-	/**
-	 * Handles preview dialog open state changes.
-	 * Clears preview content when dialog is closed.
-	 * @param open - Whether the dialog is being opened or closed
-	 */
-	function handlePreviewDialogOpenChange(open: boolean) {
-		previewDialogOpen = open;
-
-		if (!open) {
-			previewCode = '';
-			previewLanguage = 'text';
-		}
-	}
-
-	/**
-	 * Handles click events on preview buttons within HTML code blocks.
-	 * Opens a preview dialog with the rendered HTML content.
-	 * @param event - The click event from the preview button
-	 */
 	function handlePreviewClick(event: Event) {
 		event.preventDefault();
 		event.stopPropagation();
@@ -222,61 +266,6 @@
 		previewDialogOpen = true;
 	}
 
-	/**
-	 * Processes markdown content into stable and unstable HTML blocks.
-	 * Uses incremental rendering: stable blocks are cached, unstable block is re-rendered.
-	 * @param markdown - The raw markdown string to process
-	 */
-	async function processMarkdown(markdown: string) {
-		if (!markdown) {
-			renderedBlocks = [];
-			unstableBlockHtml = '';
-			return;
-		}
-
-		const normalized = preprocessLaTeX(markdown);
-		const processorInstance = processor();
-		const ast = processorInstance.parse(normalized) as MdastRoot;
-		const processedRoot = (await processorInstance.run(ast)) as HastRoot;
-		const processedChildren = processedRoot.children ?? [];
-		const stableCount = Math.max(processedChildren.length - 1, 0);
-		const nextBlocks: MarkdownBlock[] = [];
-
-		for (let index = 0; index < stableCount; index++) {
-			const hastChild = processedChildren[index];
-			const id = getHastNodeId(hastChild, index);
-			const existing = renderedBlocks[index];
-
-			if (existing && existing.id === id) {
-				nextBlocks.push(existing);
-				continue;
-			}
-
-			const html = stringifyProcessedNode(
-				processorInstance,
-				processedRoot,
-				processedChildren[index]
-			);
-
-			nextBlocks.push({ id, html });
-		}
-
-		let unstableHtml = '';
-
-		if (processedChildren.length > stableCount) {
-			const unstableChild = processedChildren[stableCount];
-			unstableHtml = stringifyProcessedNode(processorInstance, processedRoot, unstableChild);
-		}
-
-		renderedBlocks = nextBlocks;
-		await tick(); // Force DOM sync before updating unstable HTML block
-		unstableBlockHtml = unstableHtml;
-	}
-
-	/**
-	 * Attaches click event listeners to copy and preview buttons in code blocks.
-	 * Uses data-listener-bound attribute to prevent duplicate bindings.
-	 */
 	function setupCodeBlockActions() {
 		if (!containerRef) return;
 
@@ -298,97 +287,40 @@
 		}
 	}
 
-	/**
-	 * Converts a single HAST node to an enhanced HTML string.
-	 * Applies link and code block enhancements to the output.
-	 * @param processorInstance - The remark/rehype processor instance
-	 * @param processedRoot - The full processed HAST root (for context)
-	 * @param child - The specific HAST child node to stringify
-	 * @returns Enhanced HTML string representation of the node
-	 */
-	function stringifyProcessedNode(
-		processorInstance: ReturnType<typeof processor>,
-		processedRoot: HastRoot,
-		child: unknown
-	) {
-		const root: HastRoot = {
-			...(processedRoot as HastRoot),
-			children: [child as never]
-		};
+	function handlePreviewDialogOpenChange(open: boolean) {
+		previewDialogOpen = open;
 
-		return processorInstance.stringify(root);
-	}
-
-	/**
-	 * Queues markdown for processing with coalescing support.
-	 * Only processes the latest markdown when multiple updates arrive quickly.
-	 * @param markdown - The markdown content to render
-	 */
-	async function updateRenderedBlocks(markdown: string) {
-		pendingMarkdown = markdown;
-
-		if (isProcessing) {
-			return;
-		}
-
-		isProcessing = true;
-
-		try {
-			while (pendingMarkdown !== null) {
-				const nextMarkdown = pendingMarkdown;
-				pendingMarkdown = null;
-
-				await processMarkdown(nextMarkdown);
-			}
-		} catch (error) {
-			console.error('Failed to process markdown:', error);
-			renderedBlocks = [];
-			unstableBlockHtml = markdown.replace(/\n/g, '<br>');
-		} finally {
-			isProcessing = false;
+		if (!open) {
+			previewCode = '';
+			previewLanguage = 'text';
 		}
 	}
 
 	$effect(() => {
-		const currentMode = mode.current;
-		const isDark = currentMode === 'dark';
-
-		loadHighlightTheme(isDark);
+		if (content) {
+			processMarkdown(content)
+				.then((result) => {
+					processedHtml = result;
+				})
+				.catch((error) => {
+					console.error('Failed to process markdown:', error);
+					processedHtml = content.replace(/\n/g, '<br>');
+				});
+		} else {
+			processedHtml = '';
+		}
 	});
 
 	$effect(() => {
-		updateRenderedBlocks(content);
-	});
-
-	$effect(() => {
-		const hasRenderedBlocks = renderedBlocks.length > 0;
-		const hasUnstableBlock = Boolean(unstableBlockHtml);
-
-		if ((hasRenderedBlocks || hasUnstableBlock) && containerRef) {
+		if (containerRef && processedHtml) {
 			setupCodeBlockActions();
 		}
 	});
-
-	onDestroy(() => {
-		cleanupEventListeners();
-		cleanupHighlightTheme();
-	});
 </script>
 
 <div bind:this={containerRef} class={className}>
-	{#each renderedBlocks as block (block.id)}
-		<div class="markdown-block" data-block-id={block.id}>
-			<!-- eslint-disable-next-line no-at-html-tags -->
-			{@html block.html}
-		</div>
-	{/each}
-
-	{#if unstableBlockHtml}
-		<div class="markdown-block markdown-block--unstable" data-block-id="unstable">
-			<!-- eslint-disable-next-line no-at-html-tags -->
-			{@html unstableBlockHtml}
-		</div>
-	{/if}
+	<!-- eslint-disable-next-line no-at-html-tags -->
+	{@html processedHtml}
 </div>
 
 <CodePreviewDialog
@@ -399,11 +331,6 @@
 />
 
 <style>
-	.markdown-block,
-	.markdown-block--unstable {
-		display: contents;
-	}
-
 	/* Base typography styles */
 	div :global(p:not(:last-child)) {
 		margin-bottom: 1rem;

tools/server/webui/src/lib/markdown/enhance-code-blocks.ts

@@ -1,162 +0,0 @@
-/**
- * Rehype plugin to enhance code blocks with wrapper, header, and action buttons.
- *
- * Wraps <pre><code> elements with a container that includes:
- * - Language label
- * - Copy button
- * - Preview button (for HTML code blocks)
- *
- * This operates directly on the HAST tree for better performance,
- * avoiding the need to stringify and re-parse HTML.
- */
-
-import type { Plugin } from 'unified';
-import type { Root, Element, ElementContent } from 'hast';
-import { visit } from 'unist-util-visit';
-
-declare global {
-	interface Window {
-		idxCodeBlock?: number;
-	}
-}
-
-const COPY_ICON_SVG = `<svg xmlns="http://www.w3.org/2000/svg" width="16" height="16" viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2" stroke-linecap="round" stroke-linejoin="round" class="lucide lucide-copy-icon lucide-copy"><rect width="14" height="14" x="8" y="8" rx="2" ry="2"/><path d="M4 16c-1.1 0-2-.9-2-2V4c0-1.1.9-2 2-2h10c1.1 0 2 .9 2 2"/></svg>`;
-
-const PREVIEW_ICON_SVG = `<svg xmlns="http://www.w3.org/2000/svg" width="16" height="16" viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2" stroke-linecap="round" stroke-linejoin="round" class="lucide lucide-eye lucide-eye-icon"><path d="M2.062 12.345a1 1 0 0 1 0-.69C3.5 7.73 7.36 5 12 5s8.5 2.73 9.938 6.655a1 1 0 0 1 0 .69C20.5 16.27 16.64 19 12 19s-8.5-2.73-9.938-6.655"/><circle cx="12" cy="12" r="3"/></svg>`;
-
-/**
- * Creates an SVG element node from raw SVG string.
- * Since we can't parse HTML in HAST directly, we use the raw property.
- */
-function createRawHtmlElement(html: string): Element {
-	return {
-		type: 'element',
-		tagName: 'span',
-		properties: {},
-		children: [{ type: 'raw', value: html } as unknown as ElementContent]
-	};
-}
-
-function createCopyButton(codeId: string): Element {
-	return {
-		type: 'element',
-		tagName: 'button',
-		properties: {
-			className: ['copy-code-btn'],
-			'data-code-id': codeId,
-			title: 'Copy code',
-			type: 'button'
-		},
-		children: [createRawHtmlElement(COPY_ICON_SVG)]
-	};
-}
-
-function createPreviewButton(codeId: string): Element {
-	return {
-		type: 'element',
-		tagName: 'button',
-		properties: {
-			className: ['preview-code-btn'],
-			'data-code-id': codeId,
-			title: 'Preview code',
-			type: 'button'
-		},
-		children: [createRawHtmlElement(PREVIEW_ICON_SVG)]
-	};
-}
-
-function createHeader(language: string, codeId: string): Element {
-	const actions: Element[] = [createCopyButton(codeId)];
-
-	if (language.toLowerCase() === 'html') {
-		actions.push(createPreviewButton(codeId));
-	}
-
-	return {
-		type: 'element',
-		tagName: 'div',
-		properties: { className: ['code-block-header'] },
-		children: [
-			{
-				type: 'element',
-				tagName: 'span',
-				properties: { className: ['code-language'] },
-				children: [{ type: 'text', value: language }]
-			},
-			{
-				type: 'element',
-				tagName: 'div',
-				properties: { className: ['code-block-actions'] },
-				children: actions
-			}
-		]
-	};
-}
-
-function createWrapper(header: Element, preElement: Element): Element {
-	return {
-		type: 'element',
-		tagName: 'div',
-		properties: { className: ['code-block-wrapper'] },
-		children: [header, preElement]
-	};
-}
-
-function extractLanguage(codeElement: Element): string {
-	const className = codeElement.properties?.className;
-	if (!Array.isArray(className)) return 'text';
-
-	for (const cls of className) {
-		if (typeof cls === 'string' && cls.startsWith('language-')) {
-			return cls.replace('language-', '');
-		}
-	}
-
-	return 'text';
-}
-
-/**
- * Generates a unique code block ID using a global counter.
- */
-function generateCodeId(): string {
-	if (typeof window !== 'undefined') {
-		return `code-${(window.idxCodeBlock = (window.idxCodeBlock ?? 0) + 1)}`;
-	}
-	// Fallback for SSR - use timestamp + random
-	return `code-${Date.now()}-${Math.random().toString(36).slice(2, 7)}`;
-}
-
-/**
- * Rehype plugin to enhance code blocks with wrapper, header, and action buttons.
- * This plugin wraps <pre><code> elements with a container that includes:
- * - Language label
- * - Copy button
- * - Preview button (for HTML code blocks)
- */
-export const rehypeEnhanceCodeBlocks: Plugin<[], Root> = () => {
-	return (tree: Root) => {
-		visit(tree, 'element', (node: Element, index, parent) => {
-			if (node.tagName !== 'pre' || !parent || index === undefined) return;
-
-			const codeElement = node.children.find(
-				(child): child is Element => child.type === 'element' && child.tagName === 'code'
-			);
-
-			if (!codeElement) return;
-
-			const language = extractLanguage(codeElement);
-			const codeId = generateCodeId();
-
-			codeElement.properties = {
-				...codeElement.properties,
-				'data-code-id': codeId
-			};
-
-			const header = createHeader(language, codeId);
-			const wrapper = createWrapper(header, node);
-
-			// Replace pre with wrapper in parent
-			(parent.children as ElementContent[])[index] = wrapper;
-		});
-	};
-};

tools/server/webui/src/lib/markdown/enhance-links.ts

@@ -1,33 +0,0 @@
-/**
- * Rehype plugin to enhance links with security attributes.
- *
- * Adds target="_blank" and rel="noopener noreferrer" to all anchor elements,
- * ensuring external links open in new tabs safely.
- */
-
-import type { Plugin } from 'unified';
-import type { Root, Element } from 'hast';
-import { visit } from 'unist-util-visit';
-
-/**
- * Rehype plugin that adds security attributes to all links.
- * This plugin ensures external links open in new tabs safely by adding:
- * - target="_blank"
- * - rel="noopener noreferrer"
- */
-export const rehypeEnhanceLinks: Plugin<[], Root> = () => {
-	return (tree: Root) => {
-		visit(tree, 'element', (node: Element) => {
-			if (node.tagName !== 'a') return;
-
-			const props = node.properties ?? {};
-
-			// Only modify if href exists
-			if (!props.href) return;
-
-			props.target = '_blank';
-			props.rel = 'noopener noreferrer';
-			node.properties = props;
-		});
-	};
-};