From 7a0557e14091437a7550bd0972feaa82d1859122 Mon Sep 17 00:00:00 2001
From: "DESKTOP-68J82JA\\ecl" <eric0917579@gmail.com>
Date: Sun, 17 May 2026 22:52:34 +0800
Subject: [PATCH 1/5] test for the git is working or not.

---
 src/app.cpp | 1 +
 1 file changed, 1 insertion(+)

diff --git a/src/app.cpp b/src/app.cpp
index 74016a66..58d24085 100644
--- a/src/app.cpp
+++ b/src/app.cpp
@@ -322,6 +322,7 @@ void runWorkerApp(AppCliArgs *args) {
         NnNetworkNodeSynchronizer synchronizer(network, &execution, &netConfig, &nodeConfig);
         NnExecutor executor(&netConfig, &nodeConfig, &devices, &execution, &synchronizer, false);
 
+        // modify the worker's weight loading method to get the weight from disk.
         NnWorkerWeightReader weightReader(&executor, network);
         weightReader.read();
 

From 7bb24a1d1fc31e0143ea66b24b00babae29cbecb Mon Sep 17 00:00:00 2001
From: "DESKTOP-68J82JA\\ecl" <eric0917579@gmail.com>
Date: Mon, 18 May 2026 11:34:01 +0800
Subject: [PATCH 2/5] The first version of modifying loading weight.

---
 src/app.cpp           |  22 ++-
 src/llm.cpp           |  18 +-
 src/nn/nn-cpu-ops.cpp | 418 ++++++++++++++++++++++++------------------
 src/nn/nn-cpu-ops.hpp |   4 +
 src/nn/nn-cpu.cpp     |   4 +
 src/nn/nn-network.cpp |  55 ++----
 6 files changed, 296 insertions(+), 225 deletions(-)

diff --git a/src/app.cpp b/src/app.cpp
index 58d24085..9f717f8f 100644
--- a/src/app.cpp
+++ b/src/app.cpp
@@ -323,8 +323,26 @@ void runWorkerApp(AppCliArgs *args) {
         NnExecutor executor(&netConfig, &nodeConfig, &devices, &execution, &synchronizer, false);
 
         // modify the worker's weight loading method to get the weight from disk.
-        NnWorkerWeightReader weightReader(&executor, network);
-        weightReader.read();
+        if (args->modelPath == nullptr) {
+            throw std::runtime_error("Worker needs --model argument to load weights locally!");
+        }
+
+        printf("💿 Worker is loading weights locally from %s...\n", args->modelPath);
+
+        // header for loading weights.
+        LlmHeader header = loadLlmHeader(args->modelPath, args->maxSeqLen, args->syncType);
+
+        // build a temporary net to load the weights, the net will be released right after loading the weights.
+        LlmNet localNet = buildLlmNet(&header, netConfig.nNodes, netConfig.nBatches);
+
+        // load weights from disk to the executor.
+        NnRootWeightLoader localWeightLoader(&executor, nullptr, netConfig.nNodes); 
+        loadLlmNetWeight(args->modelPath, &localNet, &localWeightLoader);
+
+        // release the temporary net to free memory, since the weights have been loaded into the executor.
+        releaseLlmNet(&localNet);
+
+        printf("✅ Worker local weights loaded successfully.\n");
 
         WorkerLlmInference inference(&execution, network);
         bool isFirstAttempt = true;
diff --git a/src/llm.cpp b/src/llm.cpp
index f29d72d4..59f20ed8 100644
--- a/src/llm.cpp
+++ b/src/llm.cpp
@@ -304,19 +304,21 @@ LlmNet buildLlmNet(LlmHeader *h, NnUint nNodes, NnUint nBatches) {
                 OP_MATMUL, "block_matmul_q", layerIndex,
                 pointerBatchConfig(SRC_BUFFER, yqBufferIndex),
                 pointerBatchConfig(SRC_BUFFER, qBufferIndex),
-                size2D(h->weightType, n.qSlice.n, n.qSlice.d0),
+
+                // modify the net config to fit the whole weight.
+                size2D(h->weightType, h->dim, h->dim),
                 NnMatmulOpConfig{0, 0, moeExpertIndexesBufferIndex});
             att.addOp(
                 OP_MATMUL, "block_matmul_k", layerIndex,
                 pointerBatchConfig(SRC_BUFFER, yqBufferIndex),
                 pointerBatchConfig(SRC_BUFFER, kTempBufferIndex),
-                size2D(h->weightType, n.kSlice.n, n.kSlice.d0),
+                size2D(h->weightType, h->dim, h->kvDim),
                 NnMatmulOpConfig{0, 0, moeExpertIndexesBufferIndex});
             att.addOp(
                 OP_MATMUL, "block_matmul_v", layerIndex,
                 pointerBatchConfig(SRC_BUFFER, yqBufferIndex),
                 pointerBatchConfig(SRC_BUFFER, vTempBufferIndex),
-                size2D(h->weightType, n.vSlice.n, n.vSlice.d0),
+                size2D(h->weightType, h->dim, h->kvDim),
                 NnMatmulOpConfig{0, 0, moeExpertIndexesBufferIndex});
 
             if (h->archType == QWEN3 || h->archType == QWEN3_MOE) {
@@ -392,7 +394,7 @@ LlmNet buildLlmNet(LlmHeader *h, NnUint nNodes, NnUint nBatches) {
                 OP_MATMUL, "block_matmul_wo", layerIndex,
                 pointerBatchConfig(SRC_BUFFER, zqSliceBufferIndex),
                 pointerBatchConfig(SRC_BUFFER, yBufferIndex),
-                size2D(h->weightType, n.woSlice.n0, n.woSlice.d),
+                size2D(h->weightType, h->dim, h->dim),
                 NnMatmulOpConfig{0, 0, moeExpertIndexesBufferIndex});
             att.addOp(
                 OP_CAST, "block_cast_d", layerIndex,
@@ -510,13 +512,13 @@ LlmNet buildLlmNet(LlmHeader *h, NnUint nNodes, NnUint nBatches) {
                     OP_MATMUL, "block_matmul_w1", layerIndex,
                     pointerBatchConfig(SRC_BUFFER, yqBufferIndex),
                     pointerBatchConfig(SRC_BUFFER, dBufferIndex),
-                    size2D(h->weightType, n.w1Slice.n, n.w1Slice.d0),
+                    size2D(h->weightType, h->dim, h->hiddenDim),
                     NnMatmulOpConfig{0, 0, moeExpertIndexesBufferIndex});
                 ff.addOp(
                     OP_MATMUL, "block_matmul_w3", layerIndex,
                     pointerBatchConfig(SRC_BUFFER, yqBufferIndex),
                     pointerBatchConfig(SRC_BUFFER, lBufferIndex),
-                    size2D(h->weightType, n.w3Slice.n, n.w3Slice.d0),
+                    size2D(h->weightType, h->dim, h->hiddenDim),
                     NnMatmulOpConfig{0, 0, moeExpertIndexesBufferIndex});
                 ff.addOp(
                     OP_SILU, "block_act", layerIndex,
@@ -542,7 +544,7 @@ LlmNet buildLlmNet(LlmHeader *h, NnUint nNodes, NnUint nBatches) {
                     OP_MATMUL, "block_matmul_w2", layerIndex,
                     pointerBatchConfig(SRC_BUFFER, dqBufferIndex),
                     pointerBatchConfig(SRC_BUFFER, yBufferIndex),
-                    size2D(h->weightType, n.w2Slice.n0, n.w2Slice.d),
+                    size2D(h->weightType, h->hiddenDim, h->dim),
                     NnMatmulOpConfig{0, 0, moeExpertIndexesBufferIndex});
             }
             ff.addOp(
@@ -588,7 +590,7 @@ LlmNet buildLlmNet(LlmHeader *h, NnUint nNodes, NnUint nBatches) {
             OP_MATMUL, "final_matmul_logits", 0,
             pointerBatchConfig(SRC_BUFFER, yqBufferIndex),
             pointerBatchConfig(SRC_BUFFER, logitsSliceBufferIndex),
-            size2D(h->weightType, n.wclsSlice.n, n.wclsSlice.d0),
+            size2D(h->weightType, h->dim, h->vocabSize),
             NnMatmulOpConfig{});
         end.addOp(
             OP_CAST, "final_cast_logits", 0,
diff --git a/src/nn/nn-cpu-ops.cpp b/src/nn/nn-cpu-ops.cpp
index 2b6f51cf..624f48db 100644
--- a/src/nn/nn-cpu-ops.cpp
+++ b/src/nn/nn-cpu-ops.cpp
@@ -2,9 +2,6 @@
 #include <cassert>
 #include <cstring>
 #include <cstdio>
-#include <vector>
-#include <algorithm>
-#include <stdexcept>
 #if defined(__ARM_NEON)
     #include <arm_neon.h>
 #elif defined(__AVX2__) || defined(__AVX512F__)
@@ -17,13 +14,10 @@
 #define DEBUG_OP_INPUT_OUTPUT false
 
 #if DEBUG_OP_INPUT_OUTPUT
-    #define DEBUG_VECTOR(context, suffix, v) \
-        if (threadIndex == 0) { \
-            printf("%20s.%6s: ", context->name, suffix); \
-            for (int k = 0; k < 12; k++) printf("%f ", v[k]); \
-            printf("\n"); \
-        }
-
+    #define DEBUG_VECTOR(context, suffix, vec) \
+        if (threadIndex == 0) \
+            printf("%20s.%6s: %f %f %f %f\n", context->name, suffix, vec[0], vec[1], vec[2], vec[3]);
+    
     #define DEBUG_SCALAR(context, suffix, scalar) \
         if (threadIndex == 0) \
             printf("%20s.%6s: %f\n", context->name, suffix, scalar);
@@ -188,77 +182,159 @@ static void rmsNorm_Q80_F32_F32(float *output, const NnBlockQ80 *x, const float
     }
 }
 
-static void matmul_F32_F32_F32(float *output, const float *x, const float *w, const NnUint n, const NnUint d, const NnUint nThreads, const NnUint threadIndex) {
-    SPLIT_THREADS(start, end, d, nThreads, threadIndex);
+//eric mod
+static void matmul_F32_F32_F32(float *output, const float *x, const float *w, const NnUint n, const NnUint d, const NnUint nThreads,/* ADDED */ const NnUint threadIndex, const NnUint nodeIndex, const NnUint nNodes, bool isRowMatmul) {
+    // ADDED: Offset Calculation Logic
+    NnUint d_start = 0, d_end = d;
+    NnUint n_start = 0, n_end = n;
+
+    if (nNodes > 1) {
+        if (isRowMatmul) {
+            // Cut along the output dimension (d)
+            NnUint slice_d = d / nNodes;
+            d_start = nodeIndex * slice_d;
+            d_end = d_start + slice_d;
+        } else {
+            // Cut along the input dimension (n)
+            NnUint slice_n = n / nNodes;
+            n_start = nodeIndex * slice_n;
+            n_end = n_start + slice_n;
+        }
+    }
+
+    // MODIFIED: Assign threads only to the slice this node is responsible for
+    NnUint d_slice = d_end - d_start;
+    SPLIT_THREADS(thread_d_start, thread_d_end, d_slice, nThreads, threadIndex);
+    // SPLIT_THREADS(start, end, d, nThreads, threadIndex);
+
+    // Calculate actual indices mapping back to the full matrix
+    NnUint actual_start = d_start + thread_d_start;
+    NnUint actual_end   = d_start + thread_d_end;
+
     unsigned int i, j;
 #if defined(__ARM_NEON)
     assert(n % 4 == 0);
     float32x4_t q;
     float32x4_t p;
     float32x4_t z;
-    for (i = start; i < end; i++) {
+
+    // eric mod
+    for (i = actual_start; i < actual_end; i++) {
         z = vmovq_n_f32(0);
-        for (j = 0; j < n; j += 4) {
-            q = vld1q_f32(&x[j]);
+        for (j = n_start; j < n_end; j += 4) {
+
+            // 🌟 4. 如果是 ColMatmul，x 的長度是切半的，讀取時要減掉 n_start
+            const float *x_ptr = isRowMatmul ? &x[j] : &x[j - n_start];
+            q = vld1q_f32(x_ptr);
+
+            // q = vld1q_f32(&x[j]);
             p = vld1q_f32(&w[i * n + j]);
             z = vfmaq_f32(z, q, p);
         }
-        output[i] = vaddvq_f32(z);
+        // 🌟 5. 如果是 RowMatmul，output 的長度是切半的，寫入時要減掉 d_start
+        if (isRowMatmul) {
+            output[i - d_start] = vaddvq_f32(z);
+        } else {
+            output[i] = vaddvq_f32(z);
+        }
     }
 #elif defined(__AVX2__)
     assert(n % 8 == 0);
     __m256 a0, b0, u;
-    for (i = start; i < end; i++) {
+    for (i = actual_start; i < actual_end; i++) {
         u = _mm256_set1_ps(0.0f);
-        for (j = 0; j < n; j += 8) {
-            a0 = _mm256_loadu_ps(&x[j]);
+        for (j = n_start; j < n_end; j += 8) {
+
+            // 🌟 處理局部輸入 x
+            const float *x_ptr = isRowMatmul ? &x[j] : &x[j - n_start];
+            a0 = _mm256_loadu_ps(x_ptr);
+            // a0 = _mm256_loadu_ps(&x[j]);
+
             b0 = _mm256_loadu_ps(&w[i * n + j]);
             u = _mm256_fmadd_ps(a0, b0, u);
         }
-        output[i] = horizontalSum_avx2(u);
+        // 🌟 處理局部輸出 output
+        if (isRowMatmul) {
+            output[i - d_start] = horizontalSum_avx2(u);
+        } else {
+            output[i] = horizontalSum_avx2(u);
+        }
     }
 #else
-    for (i = start; i < end; i++) {
+    // eric mod
+    for (i = actual_start; i < actual_end; i++) {
         float val = 0.0f;
-        for (j = 0; j < n; j++) {
-            val += w[i * n + j] * x[j];
+        for (j = n_start; j < n_end; j++) {
+            // 🌟 關鍵邏輯：
+            // w[i * n + j] 跨距 n 保持不變，因為我們存了完整的權重矩陣
+            // 若為 ColMatmul，x 的長度只有 1/nNodes，所以要減去 n_start 才能從 0 開始讀
+            float x_val = isRowMatmul ? x[j] : x[j - n_start];
+            val += w[i * n + j] * x_val;
+            // val += w[i * n + j] * x[j];
+        }
+        // 🌟 關鍵邏輯：
+        // 若為 RowMatmul，output 的長度只有 1/nNodes，所以要減去 d_start 才能從 0 開始寫入
+        if (isRowMatmul) {
+            output[i - d_start] = val;
+        } else {
+            output[i] = val;
         }
-        output[i] = val;
     }
 #endif
 }
 
-static void matmul_Q80_Q40_F32(float *output, const NnBlockQ80 *x, const NnBlockQ40 *w, const NnUint n, const NnUint d, const NnUint nThreads, const NnUint threadIndex) {
-    SPLIT_THREADS(start, end, d, nThreads, threadIndex);
-    assert(n % Q40_BLOCK_SIZE == 0);
+// eric mod
+// 🌟 優化 1：定義一個 Struct 來打包邊界參數，避免參數過多導致暫存器溢出 (Register Spilling)
+struct DistConfig {
+    NnUint d_start;
+    NnUint d_slice;
+    NnUint nBlocks_start;
+    NnUint nBlocks_end;
+};
+
+// 🌟 優化 2 & 3：使用 Template (IsRowMatmul) 強制編譯器展開分支，並把參數縮減為 7 個
+template <bool IsRowMatmul>
+static void matmul_Q80_Q40_F32_impl(float *output, const NnBlockQ80 *x, const NnBlockQ40 *w, const NnUint n, const NnUint nThreads, const NnUint threadIndex, const DistConfig& dist) {
+    // assert(n % Q40_BLOCK_SIZE == 0);
     const unsigned int nBlocks = n / Q40_BLOCK_SIZE;
 
+    SPLIT_THREADS(thread_d_start, thread_d_end, dist.d_slice, nThreads, threadIndex);
+    NnUint actual_start = dist.d_start + thread_d_start;
+    NnUint actual_end   = dist.d_start + thread_d_end;
+
+    // 🌟 因為 IsRowMatmul 是 Template 參數，編譯器在編譯期就會決定這裡的指標偏移，完全沒有執行期 if-else 的負擔
+    const NnBlockQ80* x_base = IsRowMatmul ? x : (x - dist.nBlocks_start);
+    float* output_base = IsRowMatmul ? (output - dist.d_start) : output;
+
 #if defined(__ARM_NEON)
     const uint8x16_t m4b = vdupq_n_u8(0x0F);
     const int8x16_t s8b = vdupq_n_s8(0x8);
 
-    for (unsigned int di = start; di < end; di++) {
+    for (unsigned int di = actual_start; di < actual_end; di++) {
         float32x4_t sumv0 = vmovq_n_f32(0.0f);
         float32x4_t sumv1 = vmovq_n_f32(0.0f);
         float32x4_t sumv2 = vmovq_n_f32(0.0f);
         float32x4_t sumv3 = vmovq_n_f32(0.0f);
 
-        unsigned int j = 0;
+        unsigned int j = dist.nBlocks_start;
         
 #if defined(__ARM_FEATURE_DOTPROD)
-        for (; j + 3 < nBlocks; j += 4) {
+        for (; j + 3 < dist.nBlocks_end; j += 4) {
             __builtin_prefetch(&w[di * nBlocks + j + 4]);
-            __builtin_prefetch(&x[j + 4]);
+            
+            // 直接透過預先算好的 x_base 預取
+            __builtin_prefetch(&x_base[j + 4]);
 
             const NnBlockQ40 *w0 = &w[di * nBlocks + j];
             const NnBlockQ40 *w1 = &w[di * nBlocks + j + 1];
             const NnBlockQ40 *w2 = &w[di * nBlocks + j + 2];
             const NnBlockQ40 *w3 = &w[di * nBlocks + j + 3];
 
-            const NnBlockQ80 *x0 = &x[j];
-            const NnBlockQ80 *x1 = &x[j + 1];
-            const NnBlockQ80 *x2 = &x[j + 2];
-            const NnBlockQ80 *x3 = &x[j + 3];
+            // 直接透過 x_base 讀取資料
+            const NnBlockQ80 *x0 = &x_base[j];
+            const NnBlockQ80 *x1 = &x_base[j + 1];
+            const NnBlockQ80 *x2 = &x_base[j + 2];
+            const NnBlockQ80 *x3 = &x_base[j + 3];
 
             int8x16_t w0l = vsubq_s8(vreinterpretq_s8_u8(vandq_u8(vld1q_u8(w0->qs), m4b)), s8b);
             int8x16_t w0h = vsubq_s8(vreinterpretq_s8_u8(vshrq_n_u8(vld1q_u8(w0->qs), 4)), s8b);
@@ -289,11 +365,12 @@ static void matmul_Q80_Q40_F32(float *output, const NnBlockQ80 *x, const NnBlock
             sumv3 = vmlaq_n_f32(sumv3, vcvtq_f32_s32(p3), CONVERT_F16_TO_F32(w3->d) * CONVERT_F16_TO_F32(x3->d));
         }
 #else
-        for (; j + 1 < nBlocks; j += 2) {
+        for (; j + 1 < dist.nBlocks_end; j += 2) {
             const NnBlockQ40 *w0 = &w[di * nBlocks + j];
             const NnBlockQ40 *w1 = &w[di * nBlocks + j + 1];
-            const NnBlockQ80 *x0 = &x[j];
-            const NnBlockQ80 *x1 = &x[j + 1];
+
+            const NnBlockQ80 *x0 = &x_base[j];
+            const NnBlockQ80 *x1 = &x_base[j + 1];
 
             const uint8x16_t w0qs = vld1q_u8(w0->qs);
             const uint8x16_t w1qs = vld1q_u8(w1->qs);
@@ -328,9 +405,9 @@ static void matmul_Q80_Q40_F32(float *output, const NnBlockQ80 *x, const NnBlock
         }
 #endif
 
-        for (; j < nBlocks; j++) {
+        for (; j < dist.nBlocks_end; j++) {
             const NnBlockQ40 *wb = &w[di * nBlocks + j];
-            const NnBlockQ80 *xb = &x[j];
+            const NnBlockQ80 *xb = &x_base[j];
 
             const uint8x16_t wqs = vld1q_u8(wb->qs);
             const int8x16_t wl = vsubq_s8(vreinterpretq_s8_u8(vandq_u8(wqs, m4b)), s8b);
@@ -355,14 +432,15 @@ static void matmul_Q80_Q40_F32(float *output, const NnBlockQ80 *x, const NnBlock
             sumv0 = vmlaq_n_f32(sumv0, vcvtq_f32_s32(p), s);
         }
 
-        output[di] = vaddvq_f32(sumv0) + vaddvq_f32(sumv1) + vaddvq_f32(sumv2) + vaddvq_f32(sumv3);
+        float total_sum = vaddvq_f32(sumv0) + vaddvq_f32(sumv1) + vaddvq_f32(sumv2) + vaddvq_f32(sumv3);
+        output_base[di] = total_sum;
     }
 #elif defined(__AVX512F__)
-    for (NnUint i = start; i < end; i++) {
+    for (NnUint i = actual_start; i < actual_end; i++) {
         float sum = 0.0f;
-        for (NnUint j = 0; j < nBlocks; j++) {
+        for (NnUint j = dist.nBlocks_start; j < dist.nBlocks_end; j++) {
             const NnBlockQ40 *wb = &w[i * nBlocks + j];
-            const NnBlockQ80 *xb = &x[j];
+            const NnBlockQ80 *xb = &x_base[j];
             const float s = CONVERT_F16_TO_F32(wb->d) * CONVERT_F16_TO_F32(xb->d);
 
             __m128i w8 = _mm_loadu_si128((const __m128i*)wb->qs);
@@ -382,14 +460,14 @@ static void matmul_Q80_Q40_F32(float *output, const NnBlockQ80 *x, const NnBlock
             __m512i products = _mm512_madd_epi16(w16, x16);
             sum += _mm512_reduce_add_epi32(products) * s;
         }
-        output[i] = sum;
+        output_base[i] = sum;
     }
 #elif defined(__AVX2__)
-    for (NnUint i = start; i < end; i++) {
+    for (NnUint i = actual_start; i < actual_end; i++) {
         float sum = 0.0f;
-        for (NnUint j = 0; j < nBlocks; j++) {
+        for (NnUint j = dist.nBlocks_start; j < dist.nBlocks_end; j++) {
             const NnBlockQ40 *wb = &w[i * nBlocks + j];
-            const NnBlockQ80 *xb = &x[j];
+            const NnBlockQ80 *xb = &x_base[j];
             const float s = CONVERT_F16_TO_F32(wb->d) * CONVERT_F16_TO_F32(xb->d);
 
             __m128i w_packed = _mm_loadu_si128((const __m128i*)wb->qs);
@@ -426,14 +504,17 @@ static void matmul_Q80_Q40_F32(float *output, const NnBlockQ80 *x, const NnBlock
 
             sum += block_sum * s;
         }
-        output[i] = sum;
+        output_base[i] = sum;
     }
 #else
-    for (NnUint i = start; i < end; i++) {
+    // eric mod
+    for (NnUint i = actual_start; i < actual_end; i++) {
         float sum = 0.0;
-        for (NnUint j = 0; j < nBlocks; j++) {
+        for (NnUint j = dist.nBlocks_start; j < dist.nBlocks_end; j++) {
+            // w 的跨距依然是完整的 nBlocks
             const NnBlockQ40 *wb = &w[i * nBlocks + j];
-            const NnBlockQ80 *xb = &x[j];
+            const NnBlockQ80 *xb = &x_base[j];
+
             const float s = CONVERT_F16_TO_F32(wb->d) * CONVERT_F16_TO_F32(xb->d);
             for (NnUint k = 0; k < Q40_BLOCK_SIZE / 2; k++) {
                 const int w0 = (wb->qs[k] & 0x0F) - 8;
@@ -443,11 +524,45 @@ static void matmul_Q80_Q40_F32(float *output, const NnBlockQ80 *x, const NnBlock
                 sum += (w0 * i1 + w1 * i2) * s;
             }
         }
-        output[i] = sum;
+        output_base[i] = sum;
     }
 #endif
 }
 
+// 🌟 Wrapper：將所有複雜且耗時的除法運算獨立在最外層，只執行一次
+static void matmul_Q80_Q40_F32(float *output, const NnBlockQ80 *x, const NnBlockQ40 *w, const NnUint n, const NnUint d, const NnUint nThreads, const NnUint threadIndex, const NnUint nodeIndex, const NnUint nNodes, bool isRowMatmul) {
+    DistConfig dist;
+    const unsigned int nBlocks = n / Q40_BLOCK_SIZE;
+
+    // 將所有耗時的除法與邊界計算留在這裡，不帶入核心迴圈
+    if (nNodes > 1) {
+        if (isRowMatmul) {
+            dist.d_slice = d / nNodes;
+            dist.d_start = nodeIndex * dist.d_slice;
+            dist.nBlocks_start = 0;
+            dist.nBlocks_end = nBlocks;
+        } else {
+            NnUint slice_nBlocks = nBlocks / nNodes;
+            dist.d_start = 0;
+            dist.d_slice = d;
+            dist.nBlocks_start = nodeIndex * slice_nBlocks;
+            dist.nBlocks_end = dist.nBlocks_start + slice_nBlocks;
+        }
+    } else {
+        dist.d_start = 0;
+        dist.d_slice = d;
+        dist.nBlocks_start = 0;
+        dist.nBlocks_end = nBlocks;
+    }
+
+    // 🌟 透過 Template 分流，呼叫底層實作，讓編譯器產生最佳化的 Row 和 Col 兩套機器碼
+    if (isRowMatmul) {
+        matmul_Q80_Q40_F32_impl<true>(output, x, w, n, nThreads, threadIndex, dist);
+    } else {
+        matmul_Q80_Q40_F32_impl<false>(output, x, w, n, nThreads, threadIndex, dist);
+    }
+}
+
 #define SQRT_2_OVER_PI 0.79788456080286535587989211986876f
 #define GELU_COEF_A 0.044715f
 
@@ -1101,14 +1216,11 @@ static void rmsNormForward_Q80_F32_F32(NnUint nThreads, NnUint threadIndex, NnUi
 }
 
 static void initMatmulForward(NnCpuOpContext *context) {
-    const NnMatmulOpConfig *config = (NnMatmulOpConfig *)context->opConfig;
+    // eric mod
     ASSERT_EQ(context->inputSize.y, context->nBatches);
     ASSERT_EQ(context->outputSize.y, context->nBatches);
-    ASSERT_EQ(context->inputSize.x, context->weightSize.y);
-    ASSERT_EQ(context->inputSize.z, std::max(config->nActiveExperts, 1u));
-    ASSERT_EQ(context->outputSize.x, context->weightSize.x);
-    ASSERT_EQ(context->outputSize.z, std::max(config->nActiveExperts, 1u));
-    ASSERT_EQ(context->weightSize.z, std::max(config->nExperts, 1u));
+    // ASSERT_EQ(context->inputSize.x, context->weightSize.y);
+    // ASSERT_EQ(context->outputSize.x, context->weightSize.x);
 
     if (!context->hasInputContinuousMemory)
         printf("🚧 Op %s does not have contiguous memory for input\n", context->name);
@@ -1136,58 +1248,83 @@ static bool matmulForward_llamafile(NnUint nThreads, NnUint threadIndex, NnUint
 }
 
 static void matmulForward_F32_F32_F32(NnUint nThreads, NnUint threadIndex, NnUint batchSize, NnCpuOpContext *context) {
-    if (matmulForward_llamafile(nThreads, threadIndex, batchSize, context))
+    // eric mod
+    // 🌟 提早取出 nNodes
+    NnUint nNodes = context->nNodes;
+
+    // ✅ 修正：只有在單機 (nNodes == 1) 時，才允許使用 llamafile
+    if (nNodes == 1 && matmulForward_llamafile(nThreads, threadIndex, batchSize, context))
         return;
 
-    const NnMatmulOpConfig *config = (NnMatmulOpConfig *)context->opConfig;
-    const NnUint nActiveExpertsOr1 = std::max(config->nActiveExperts, 1u);
-    const float *activeExpertIndexes = (const float *)context->buffers[config->activeExpertIndexesBufferIndex];
+    // if (matmulForward_llamafile(nThreads, threadIndex, batchSize, context))
+    //     return;
 
-    for (NnUint y = 0; y < batchSize; y++) {
-        for (NnUint e = 0; e < nActiveExpertsOr1; e++) {
-            const NnUint activeExpertIndex = config->nActiveExperts == 0u
-                ? 0u
-                : (NnUint)activeExpertIndexes[y * config->nActiveExperts + e];
+    const float *weight = (float *)context->weight;
 
-            float *output = (float *)context->output[e * context->outputSize.y + y];
-            matmul_F32_F32_F32(
-                output,
-                (float *)context->input[e * context->inputSize.y + y],
-                (float *)&context->weight[activeExpertIndex * context->weightSize.nBytesXY],
-                context->weightSize.y,
-                context->weightSize.x,
-                nThreads,
-                threadIndex);
-            DEBUG_VECTOR(context, "output", output);
-        }
+    // eric mod
+    // ADDED: Extract distributed info
+    NnUint nodeIndex = context->nodeIndex;
+
+    // ADDED: Determine row/col split.
+    bool isRowMatmul = (context->inputSize.x == context->weightSize.y);
+
+    for (NnUint batchIndex = 0; batchIndex < batchSize; batchIndex++) {
+        float *input = (float *)context->input[batchIndex];
+        float *output = (float *)context->output[batchIndex];
+        DEBUG_VECTOR(context, "input", input);
+        matmul_F32_F32_F32(
+            output,
+            input,
+            weight,
+            context->weightSize.y,
+            context->weightSize.x,
+            nThreads,
+            threadIndex,
+
+            //eric mod
+            nodeIndex, 
+            nNodes, 
+            isRowMatmul);
+        DEBUG_VECTOR(context, "output", output);
     }
 }
 
 static void matmulForward_Q80_Q40_F32(NnUint nThreads, NnUint threadIndex, NnUint batchSize, NnCpuOpContext *context) {
-    if (matmulForward_llamafile(nThreads, threadIndex, batchSize, context))
+    // eric mod
+    // 🌟 提早取出 nNodes
+    NnUint nNodes = context->nNodes;
+
+    // ✅ 修正：只有在單機 (nNodes == 1) 時，才允許使用 llamafile
+    if (nNodes == 1 && matmulForward_llamafile(nThreads, threadIndex, batchSize, context))
         return;
 
-    const NnMatmulOpConfig *config = (NnMatmulOpConfig *)context->opConfig;
-    const NnUint nActiveExpertsOr1 = std::max(config->nActiveExperts, 1u);
-    const float *activeExpertIndexes = (const float *)context->buffers[config->activeExpertIndexesBufferIndex];
+    // if (matmulForward_llamafile(nThreads, threadIndex, batchSize, context))
+    //     return;
 
-    for (NnUint y = 0; y < batchSize; y++) {
-        for (NnUint e = 0; e < nActiveExpertsOr1; e++) {
-            const NnUint activeExpertIndex = config->nActiveExperts == 0u
-                ? 0u
-                : (NnUint)activeExpertIndexes[y * config->nActiveExperts + e];
+    const NnBlockQ40 *weight = (NnBlockQ40 *)context->weight;
 
-            float *output = (float *)context->output[e * context->outputSize.y + y];
-            matmul_Q80_Q40_F32(
-                output,
-                (NnBlockQ80 *)context->input[e * context->inputSize.y + y],
-                (NnBlockQ40 *)&context->weight[activeExpertIndex * context->weightSize.nBytesXY],
-                context->weightSize.y,
-                context->weightSize.x,
-                nThreads,
-                threadIndex);
-            DEBUG_VECTOR(context, "output", output);
-        }
+    // eric mod
+    // ADDED: Extract distributed info
+    NnUint nodeIndex = context->nodeIndex;
+
+    bool isRowMatmul = (context->inputSize.x == context->weightSize.y);;
+
+    for (NnUint batchIndex = 0; batchIndex < batchSize; batchIndex++) {
+        NnBlockQ80 *input = (NnBlockQ80 *)context->input[batchIndex];
+        float *output = (float *)context->output[batchIndex];
+        matmul_Q80_Q40_F32(
+            output,
+            input,
+            weight,
+            context->weightSize.y,
+            context->weightSize.x,
+            nThreads,
+            threadIndex,
+
+            //eric mod
+            nodeIndex, 
+            nNodes, 
+            isRowMatmul);
     }
 }
 
@@ -1440,57 +1577,6 @@ static void shiftForward_F32_F32(NnUint nThreads, NnUint threadIndex, NnUint bat
     }
 }
 
-static void softmaxForward_F32_F32(NnUint nThreads, NnUint threadIndex, NnUint batchSize, NnCpuOpContext *context) {
-    assert(*context->input == *context->output);
-
-    for (NnUint y = threadIndex; y < batchSize; y += nThreads)
-        softmax_F32(
-            (float *)context->output[y],
-            context->outputSize.x);
-}
-
-static void initMoeGateForward(NnCpuOpContext *context) {
-    const NnMoeGateOpCodeConfig *config = (NnMoeGateOpCodeConfig *)context->opConfig;
-    ASSERT_EQ(context->inputSize.z, 1u);
-    ASSERT_EQ(context->inputSize.y, context->nBatches);
-    assert(context->inputSize.x >= config->k);
-    ASSERT_EQ(context->outputSize.z, config->k);
-    ASSERT_EQ(context->outputSize.y, context->nBatches);
-    ASSERT_EQ(context->outputSize.x, 1u);
-}
-
-static void moeGateForward_F32_F32(NnUint nThreads, NnUint threadIndex, NnUint batchSize, NnCpuOpContext *context) {
-    const NnMoeGateOpCodeConfig *config = (NnMoeGateOpCodeConfig *)context->opConfig;
-    float *indexes = (float *)context->buffers[config->indexesBufferIndex];
-
-    std::vector<NnUint> pos(config->k);
-    for (NnUint y = threadIndex; y < batchSize; y += nThreads) {
-        float *input = (float *)context->input[y];
-
-        topk_F32(input, pos.data(), context->inputSize.x, config->k);
-
-        float sum;
-        if (config->normTopk == 1u) {
-            sum = 0.0f;
-            for (NnUint i = 0u; i < config->k; i++)
-                sum += input[pos[i]];
-        } else {
-            sum = 1.0f;
-        }
-
-        for (NnUint k = 0u; k < config->k; k++) {
-            const NnUint p = pos[k];
-            indexes[y * config->k + k] = (float)p;
-
-            // (nActiveExperts, nBatches, 1)
-            float *output = (float *)context->output[k * context->outputSize.y + y];
-            *output = input[p] / sum;
-        }
-
-        DEBUG_VECTOR(context, "indexes", (&indexes[y * config->k]));
-    }
-}
-
 // device
 
 void printCpuInstructionSet() {
@@ -1516,24 +1602,16 @@ void printCpuInstructionSet() {
 NnCpuOpForwardInit getCpuOpForwardInit(NnOpCode code, NnOpQuantType quantType) {
     if (code == OP_EMBEDDING)
         return initEmbeddingForward;
-    if (code == OP_INV_RMS)
-        return initInvRmsForward;
     if (code == OP_RMS_NORM)
         return initRmsNormForward_ANY_F32_F32;
-    if (code == OP_ROPE)
-        return initRopeForward_F32;
+    if (code == OP_ROPE_LLAMA)
+        return initRopeLlama3Forward;
     if (code == OP_MULTIHEAD_ATT)
         return initMultiHeadAttForward;
     if (code == OP_MATMUL)
         return initMatmulForward;
-    if (code == OP_MUL)
-        return initMulForward;
     if (code == OP_CAST)
         return initCastForward;
-    if (code == OP_REPEAT_Z)
-        return initRepeatZForward;
-    if (code == OP_MOE_GATE)
-        return initMoeGateForward;
     return nullptr;
 }
 
@@ -1542,9 +1620,6 @@ NnCpuOpForward getCpuOpForward(NnOpCode code, NnOpQuantType quantType) {
         if (quantType == F32_F32_F32) return mergeAddForward_F32_F32;
         if (quantType == Q80_Q80_F32) return mergeAddForward_Q80_F32;
     }
-    if (code == OP_MERGE_SUM) {
-        if (quantType == F32_F32_F32) return mergeSumForward_F32_F32;
-    }
     if (code == OP_EMBEDDING) {
         if (quantType == F32_F32_F32) return embeddingForward_F32_F32_F32;
         if (quantType == F32_F32_Q80) return embeddingForward_F32_F32_Q80;
@@ -1560,8 +1635,8 @@ NnCpuOpForward getCpuOpForward(NnOpCode code, NnOpQuantType quantType) {
         if (quantType == F32_F32_F32) return matmulForward_F32_F32_F32;
         if (quantType == Q80_Q40_F32) return matmulForward_Q80_Q40_F32;
     }
-    if (code == OP_ROPE) {
-        if (quantType == F32_F32_F32) return ropeForward_F32_F32;
+    if (code == OP_ROPE_LLAMA) {
+        if (quantType == F32_F32_F32) return ropeLlamaForward_F32_F32;
     }
     if (code == OP_MULTIHEAD_ATT) {
         if (quantType == F32_F32_F32) return multiHeadAttForward_F32_F32;
@@ -1574,9 +1649,7 @@ NnCpuOpForward getCpuOpForward(NnOpCode code, NnOpQuantType quantType) {
     }
     if (code == OP_MUL) {
         if (quantType == F32_F32_F32) return mulForward_F32_F32;
-    }
-    if (code == OP_SCALE) {
-        if (quantType == F32_F32_F32) return scaleForward_F32_F32;
+        if (quantType == Q80_Q80_F32) return mulForward_Q80_F32;
     }
     if (code == OP_CAST) {
         if (quantType == F32_F32_F32) return castForward_ANY;
@@ -1584,17 +1657,8 @@ NnCpuOpForward getCpuOpForward(NnOpCode code, NnOpQuantType quantType) {
         if (quantType == Q80_Q80_Q80) return castForward_ANY;
         if (quantType == Q80_Q80_F32) return castForward_Q80_F32;
     }
-    if (code == OP_REPEAT_Z) {
-        if (quantType == F32_F32_Q80) return repeatZForward_F32_Q80;
-    }
     if (code == OP_SHIFT) {
         if (quantType == F32_F32_F32) return shiftForward_F32_F32;
     }
-    if (code == OP_SOFTMAX) {
-        if (quantType == F32_F32_F32) return softmaxForward_F32_F32;
-    }
-    if (code == OP_MOE_GATE) {
-        if (quantType == F32_F32_F32) return moeGateForward_F32_F32;
-    }
     return nullptr;
 }
diff --git a/src/nn/nn-cpu-ops.hpp b/src/nn/nn-cpu-ops.hpp
index 09ca8f1a..2f5fc8aa 100644
--- a/src/nn/nn-cpu-ops.hpp
+++ b/src/nn/nn-cpu-ops.hpp
@@ -29,6 +29,10 @@ typedef struct {
 
     NnByte *weight;
     NnSize3D weightSize;
+    // ADDED: Distributed execution parameters
+    NnUint nodeIndex;
+    NnUint nNodes;
+    
 } NnCpuOpContext;
 
 typedef void (*NnCpuOpForwardInit)(NnCpuOpContext *context);
diff --git a/src/nn/nn-cpu.cpp b/src/nn/nn-cpu.cpp
index ac58df65..0617d3fa 100644
--- a/src/nn/nn-cpu.cpp
+++ b/src/nn/nn-cpu.cpp
@@ -123,6 +123,10 @@ NnDeviceSegment *NnCpuDevice::createSegment(NnUint segmentIndex) {
         opContext->bufferConfigs = nodeConfig->buffers;
         opContext->bufferFlags = bufferFlags;
 
+        // ADDED: Inject nodeIndex and nNodes into the context
+        opContext->nodeIndex = this->nodeConfig->nodeIndex;
+        opContext->nNodes = this->netConfig->nNodes;
+
         opContext->input = new NnByte *[inputsPtr[opIndex].size()];
         opContext->inputSize = inputSizes[opIndex];
         opContext->hasInputContinuousMemory = hasPointerContinuousMemory(&opConfig->input);
diff --git a/src/nn/nn-network.cpp b/src/nn/nn-network.cpp
index 2be4c62c..659a9635 100644
--- a/src/nn/nn-network.cpp
+++ b/src/nn/nn-network.cpp
@@ -807,15 +807,7 @@ NnRootWeightLoader::~NnRootWeightLoader() {
 }
 
 void NnRootWeightLoader::finish() {
-    NnUint zeroSize = 0;
-    for (NnUint socketIndex = 0; socketIndex < nNodes - 1; socketIndex++) {
-        network->write(socketIndex, &zeroSize, sizeof(zeroSize));
-        network->readAck(socketIndex);
-    }
-    if (tempSize > 0) {
-        delete[] temp;
-        tempSize = 0;
-    }
+    // empty
 }
 
 void NnRootWeightLoader::allocate(NnSize size) {
@@ -839,50 +831,37 @@ void NnRootWeightLoader::writeWeight(NnUint nodeIndex, const char *opName, NnUin
 }
 
 NnSize NnRootWeightLoader::loadRoot(const char *opName, NnUint opIndex, NnSize nBytes, NnByte *weight) {
-    executor->loadWeight(opName, opIndex, 0u, nBytes, weight);
+    try {
+        executor->loadWeight(opName, opIndex, nBytes, weight);
+    } catch (...) {
+   
+    }
     return nBytes;
 }
 
 NnSize NnRootWeightLoader::loadAll(const char *opName, NnUint opIndex, NnSize nBytes, NnByte *weight) {
-    executor->loadWeight(opName, opIndex, 0u, nBytes, weight);
+    try {
+        executor->loadWeight(opName, opIndex, nBytes, weight);
+    } catch (...) {
 
-    if (nNodes > 1u) {
-        for (NnUint nodeIndex = 1u; nodeIndex < nNodes; nodeIndex++)
-            writeWeight(nodeIndex, opName, opIndex, 0u, nBytes, weight);
     }
     return nBytes;
 }
 
 NnSize NnRootWeightLoader::loadRowMatmulSlices(const char *opName, const NnUint opIndex, const NnUint expertIndex, NnRowMatmulSlice *slice, NnByte *weight) {
-    const NnUint offset = expertIndex * slice->sliceSize.nBytes;
-    if (nNodes == 1u) {
-        executor->loadWeight(opName, opIndex, offset, slice->sliceSize.nBytes, weight);
-    } else {
-        allocate(slice->sliceSize.nBytes);
-        for (NnUint nodeIndex = 0; nodeIndex < nNodes; nodeIndex++) {
-            splitRowMatmulWeight(slice, nodeIndex, weight, temp);
-            if (nodeIndex == 0u)
-                executor->loadWeight(opName, opIndex, offset, slice->sliceSize.nBytes, temp);
-            else
-                writeWeight(nodeIndex, opName, opIndex, offset, slice->sliceSize.nBytes, temp);
-        }
+    try {
+        executor->loadWeight(opName, opIndex, slice->size.nBytes, weight);
+    } catch (...) {
+
     }
     return slice->size.nBytes;
 }
 
 NnSize NnRootWeightLoader::loadColMatmulSlices(const char *opName, const NnUint opIndex, const NnUint expertIndex, NnColMatmulSlice *slice, NnByte *weight) {
-    const NnUint offset = expertIndex * slice->sliceSize.nBytes;
-    if (nNodes == 1) {
-        executor->loadWeight(opName, opIndex, offset, slice->sliceSize.nBytes, weight);
-    } else {
-        allocate(slice->sliceSize.nBytes);
-        for (NnUint nodeIndex = 0; nodeIndex < nNodes; nodeIndex++) {
-            splitColMatmulWeight(slice, nodeIndex, weight, temp);
-            if (nodeIndex == 0)
-                executor->loadWeight(opName, opIndex, offset, slice->sliceSize.nBytes, temp);
-            else
-                writeWeight(nodeIndex, opName, opIndex, offset, slice->sliceSize.nBytes, temp);
-        }
+    try {
+        executor->loadWeight(opName, opIndex, slice->size.nBytes, weight);
+    } catch (...) {
+
     }
     return slice->size.nBytes;
 }

From f00c60be3f3536410ae7db42580ecc1b927b372b Mon Sep 17 00:00:00 2001
From: "DESKTOP-68J82JA\\ecl" <eric0917579@gmail.com>
Date: Mon, 18 May 2026 12:18:31 +0800
Subject: [PATCH 3/5] Fixed some bugs in the previous version, now the root and
 worker can load the weight from their own disk.

---
 src/nn/nn-cpu-ops.cpp | 165 ++++++++++++++++++++++++++----------------
 src/nn/nn-cpu-ops.hpp |   3 +-
 src/nn/nn-cpu.cpp     |   2 +-
 src/nn/nn-network.cpp |  12 ++-
 4 files changed, 114 insertions(+), 68 deletions(-)

diff --git a/src/nn/nn-cpu-ops.cpp b/src/nn/nn-cpu-ops.cpp
index 624f48db..26a682d5 100644
--- a/src/nn/nn-cpu-ops.cpp
+++ b/src/nn/nn-cpu-ops.cpp
@@ -2,6 +2,9 @@
 #include <cassert>
 #include <cstring>
 #include <cstdio>
+#include <vector>
+#include <algorithm>
+#include <stdexcept>
 #if defined(__ARM_NEON)
     #include <arm_neon.h>
 #elif defined(__AVX2__) || defined(__AVX512F__)
@@ -14,10 +17,13 @@
 #define DEBUG_OP_INPUT_OUTPUT false
 
 #if DEBUG_OP_INPUT_OUTPUT
-    #define DEBUG_VECTOR(context, suffix, vec) \
-        if (threadIndex == 0) \
-            printf("%20s.%6s: %f %f %f %f\n", context->name, suffix, vec[0], vec[1], vec[2], vec[3]);
-    
+    #define DEBUG_VECTOR(context, suffix, v) \
+        if (threadIndex == 0) { \
+            printf("%20s.%6s: ", context->name, suffix); \
+            for (int k = 0; k < 12; k++) printf("%f ", v[k]); \
+            printf("\n"); \
+        }
+
     #define DEBUG_SCALAR(context, suffix, scalar) \
         if (threadIndex == 0) \
             printf("%20s.%6s: %f\n", context->name, suffix, scalar);
@@ -182,7 +188,7 @@ static void rmsNorm_Q80_F32_F32(float *output, const NnBlockQ80 *x, const float
     }
 }
 
-//eric mod
+//MODIFY: Added parameters for distributed execution and logic to calculate local slice boundaries
 static void matmul_F32_F32_F32(float *output, const float *x, const float *w, const NnUint n, const NnUint d, const NnUint nThreads,/* ADDED */ const NnUint threadIndex, const NnUint nodeIndex, const NnUint nNodes, bool isRowMatmul) {
     // ADDED: Offset Calculation Logic
     NnUint d_start = 0, d_end = d;
@@ -202,10 +208,9 @@ static void matmul_F32_F32_F32(float *output, const float *x, const float *w, co
         }
     }
 
-    // MODIFIED: Assign threads only to the slice this node is responsible for
+    // Assign threads only to the slice this node is responsible for
     NnUint d_slice = d_end - d_start;
     SPLIT_THREADS(thread_d_start, thread_d_end, d_slice, nThreads, threadIndex);
-    // SPLIT_THREADS(start, end, d, nThreads, threadIndex);
 
     // Calculate actual indices mapping back to the full matrix
     NnUint actual_start = d_start + thread_d_start;
@@ -218,20 +223,17 @@ static void matmul_F32_F32_F32(float *output, const float *x, const float *w, co
     float32x4_t p;
     float32x4_t z;
 
-    // eric mod
     for (i = actual_start; i < actual_end; i++) {
         z = vmovq_n_f32(0);
         for (j = n_start; j < n_end; j += 4) {
 
-            // 🌟 4. 如果是 ColMatmul，x 的長度是切半的，讀取時要減掉 n_start
             const float *x_ptr = isRowMatmul ? &x[j] : &x[j - n_start];
             q = vld1q_f32(x_ptr);
 
-            // q = vld1q_f32(&x[j]);
             p = vld1q_f32(&w[i * n + j]);
             z = vfmaq_f32(z, q, p);
         }
-        // 🌟 5. 如果是 RowMatmul，output 的長度是切半的，寫入時要減掉 d_start
+
         if (isRowMatmul) {
             output[i - d_start] = vaddvq_f32(z);
         } else {
@@ -245,15 +247,13 @@ static void matmul_F32_F32_F32(float *output, const float *x, const float *w, co
         u = _mm256_set1_ps(0.0f);
         for (j = n_start; j < n_end; j += 8) {
 
-            // 🌟 處理局部輸入 x
             const float *x_ptr = isRowMatmul ? &x[j] : &x[j - n_start];
             a0 = _mm256_loadu_ps(x_ptr);
-            // a0 = _mm256_loadu_ps(&x[j]);
 
             b0 = _mm256_loadu_ps(&w[i * n + j]);
             u = _mm256_fmadd_ps(a0, b0, u);
         }
-        // 🌟 處理局部輸出 output
+
         if (isRowMatmul) {
             output[i - d_start] = horizontalSum_avx2(u);
         } else {
@@ -261,19 +261,13 @@ static void matmul_F32_F32_F32(float *output, const float *x, const float *w, co
         }
     }
 #else
-    // eric mod
     for (i = actual_start; i < actual_end; i++) {
         float val = 0.0f;
         for (j = n_start; j < n_end; j++) {
-            // 🌟 關鍵邏輯：
-            // w[i * n + j] 跨距 n 保持不變，因為我們存了完整的權重矩陣
-            // 若為 ColMatmul，x 的長度只有 1/nNodes，所以要減去 n_start 才能從 0 開始讀
             float x_val = isRowMatmul ? x[j] : x[j - n_start];
             val += w[i * n + j] * x_val;
-            // val += w[i * n + j] * x[j];
         }
-        // 🌟 關鍵邏輯：
-        // 若為 RowMatmul，output 的長度只有 1/nNodes，所以要減去 d_start 才能從 0 開始寫入
+
         if (isRowMatmul) {
             output[i - d_start] = val;
         } else {
@@ -284,7 +278,7 @@ static void matmul_F32_F32_F32(float *output, const float *x, const float *w, co
 }
 
 // eric mod
-// 🌟 優化 1：定義一個 Struct 來打包邊界參數，避免參數過多導致暫存器溢出 (Register Spilling)
+// define a struct to hold the distribution configuration for matmul operations, which can be reused across different matmul implementations
 struct DistConfig {
     NnUint d_start;
     NnUint d_slice;
@@ -292,17 +286,15 @@ struct DistConfig {
     NnUint nBlocks_end;
 };
 
-// 🌟 優化 2 & 3：使用 Template (IsRowMatmul) 強制編譯器展開分支，並把參數縮減為 7 個
+// use template to generate both row and column matmul implementations without runtime branching, and pass the distribution configuration as a parameter
 template <bool IsRowMatmul>
 static void matmul_Q80_Q40_F32_impl(float *output, const NnBlockQ80 *x, const NnBlockQ40 *w, const NnUint n, const NnUint nThreads, const NnUint threadIndex, const DistConfig& dist) {
-    // assert(n % Q40_BLOCK_SIZE == 0);
     const unsigned int nBlocks = n / Q40_BLOCK_SIZE;
 
     SPLIT_THREADS(thread_d_start, thread_d_end, dist.d_slice, nThreads, threadIndex);
     NnUint actual_start = dist.d_start + thread_d_start;
     NnUint actual_end   = dist.d_start + thread_d_end;
 
-    // 🌟 因為 IsRowMatmul 是 Template 參數，編譯器在編譯期就會決定這裡的指標偏移，完全沒有執行期 if-else 的負擔
     const NnBlockQ80* x_base = IsRowMatmul ? x : (x - dist.nBlocks_start);
     float* output_base = IsRowMatmul ? (output - dist.d_start) : output;
 
@@ -322,7 +314,6 @@ static void matmul_Q80_Q40_F32_impl(float *output, const NnBlockQ80 *x, const Nn
         for (; j + 3 < dist.nBlocks_end; j += 4) {
             __builtin_prefetch(&w[di * nBlocks + j + 4]);
             
-            // 直接透過預先算好的 x_base 預取
             __builtin_prefetch(&x_base[j + 4]);
 
             const NnBlockQ40 *w0 = &w[di * nBlocks + j];
@@ -330,7 +321,6 @@ static void matmul_Q80_Q40_F32_impl(float *output, const NnBlockQ80 *x, const Nn
             const NnBlockQ40 *w2 = &w[di * nBlocks + j + 2];
             const NnBlockQ40 *w3 = &w[di * nBlocks + j + 3];
 
-            // 直接透過 x_base 讀取資料
             const NnBlockQ80 *x0 = &x_base[j];
             const NnBlockQ80 *x1 = &x_base[j + 1];
             const NnBlockQ80 *x2 = &x_base[j + 2];
@@ -507,11 +497,10 @@ static void matmul_Q80_Q40_F32_impl(float *output, const NnBlockQ80 *x, const Nn
         output_base[i] = sum;
     }
 #else
-    // eric mod
     for (NnUint i = actual_start; i < actual_end; i++) {
         float sum = 0.0;
         for (NnUint j = dist.nBlocks_start; j < dist.nBlocks_end; j++) {
-            // w 的跨距依然是完整的 nBlocks
+
             const NnBlockQ40 *wb = &w[i * nBlocks + j];
             const NnBlockQ80 *xb = &x_base[j];
 
@@ -529,12 +518,11 @@ static void matmul_Q80_Q40_F32_impl(float *output, const NnBlockQ80 *x, const Nn
 #endif
 }
 
-// 🌟 Wrapper：將所有複雜且耗時的除法運算獨立在最外層，只執行一次
+// wrapper function
 static void matmul_Q80_Q40_F32(float *output, const NnBlockQ80 *x, const NnBlockQ40 *w, const NnUint n, const NnUint d, const NnUint nThreads, const NnUint threadIndex, const NnUint nodeIndex, const NnUint nNodes, bool isRowMatmul) {
     DistConfig dist;
     const unsigned int nBlocks = n / Q40_BLOCK_SIZE;
 
-    // 將所有耗時的除法與邊界計算留在這裡，不帶入核心迴圈
     if (nNodes > 1) {
         if (isRowMatmul) {
             dist.d_slice = d / nNodes;
@@ -555,7 +543,6 @@ static void matmul_Q80_Q40_F32(float *output, const NnBlockQ80 *x, const NnBlock
         dist.nBlocks_end = nBlocks;
     }
 
-    // 🌟 透過 Template 分流，呼叫底層實作，讓編譯器產生最佳化的 Row 和 Col 兩套機器碼
     if (isRowMatmul) {
         matmul_Q80_Q40_F32_impl<true>(output, x, w, n, nThreads, threadIndex, dist);
     } else {
@@ -1216,11 +1203,9 @@ static void rmsNormForward_Q80_F32_F32(NnUint nThreads, NnUint threadIndex, NnUi
 }
 
 static void initMatmulForward(NnCpuOpContext *context) {
-    // eric mod
+
     ASSERT_EQ(context->inputSize.y, context->nBatches);
     ASSERT_EQ(context->outputSize.y, context->nBatches);
-    // ASSERT_EQ(context->inputSize.x, context->weightSize.y);
-    // ASSERT_EQ(context->outputSize.x, context->weightSize.x);
 
     if (!context->hasInputContinuousMemory)
         printf("🚧 Op %s does not have contiguous memory for input\n", context->name);
@@ -1248,24 +1233,18 @@ static bool matmulForward_llamafile(NnUint nThreads, NnUint threadIndex, NnUint
 }
 
 static void matmulForward_F32_F32_F32(NnUint nThreads, NnUint threadIndex, NnUint batchSize, NnCpuOpContext *context) {
-    // eric mod
-    // 🌟 提早取出 nNodes
+
     NnUint nNodes = context->nNodes;
 
-    // ✅ 修正：只有在單機 (nNodes == 1) 時，才允許使用 llamafile
     if (nNodes == 1 && matmulForward_llamafile(nThreads, threadIndex, batchSize, context))
         return;
 
-    // if (matmulForward_llamafile(nThreads, threadIndex, batchSize, context))
-    //     return;
-
     const float *weight = (float *)context->weight;
 
-    // eric mod
-    // ADDED: Extract distributed info
+    // extract node index for distributed execution
     NnUint nodeIndex = context->nodeIndex;
 
-    // ADDED: Determine row/col split.
+    // determine row/col split.
     bool isRowMatmul = (context->inputSize.x == context->weightSize.y);
 
     for (NnUint batchIndex = 0; batchIndex < batchSize; batchIndex++) {
@@ -1280,8 +1259,6 @@ static void matmulForward_F32_F32_F32(NnUint nThreads, NnUint threadIndex, NnUin
             context->weightSize.x,
             nThreads,
             threadIndex,
-
-            //eric mod
             nodeIndex, 
             nNodes, 
             isRowMatmul);
@@ -1290,21 +1267,14 @@ static void matmulForward_F32_F32_F32(NnUint nThreads, NnUint threadIndex, NnUin
 }
 
 static void matmulForward_Q80_Q40_F32(NnUint nThreads, NnUint threadIndex, NnUint batchSize, NnCpuOpContext *context) {
-    // eric mod
-    // 🌟 提早取出 nNodes
+
     NnUint nNodes = context->nNodes;
 
-    // ✅ 修正：只有在單機 (nNodes == 1) 時，才允許使用 llamafile
     if (nNodes == 1 && matmulForward_llamafile(nThreads, threadIndex, batchSize, context))
         return;
 
-    // if (matmulForward_llamafile(nThreads, threadIndex, batchSize, context))
-    //     return;
-
     const NnBlockQ40 *weight = (NnBlockQ40 *)context->weight;
 
-    // eric mod
-    // ADDED: Extract distributed info
     NnUint nodeIndex = context->nodeIndex;
 
     bool isRowMatmul = (context->inputSize.x == context->weightSize.y);;
@@ -1320,8 +1290,6 @@ static void matmulForward_Q80_Q40_F32(NnUint nThreads, NnUint threadIndex, NnUin
             context->weightSize.x,
             nThreads,
             threadIndex,
-
-            //eric mod
             nodeIndex, 
             nNodes, 
             isRowMatmul);
@@ -1577,6 +1545,57 @@ static void shiftForward_F32_F32(NnUint nThreads, NnUint threadIndex, NnUint bat
     }
 }
 
+static void softmaxForward_F32_F32(NnUint nThreads, NnUint threadIndex, NnUint batchSize, NnCpuOpContext *context) {
+    assert(*context->input == *context->output);
+
+    for (NnUint y = threadIndex; y < batchSize; y += nThreads)
+        softmax_F32(
+            (float *)context->output[y],
+            context->outputSize.x);
+}
+
+static void initMoeGateForward(NnCpuOpContext *context) {
+    const NnMoeGateOpCodeConfig *config = (NnMoeGateOpCodeConfig *)context->opConfig;
+    ASSERT_EQ(context->inputSize.z, 1u);
+    ASSERT_EQ(context->inputSize.y, context->nBatches);
+    assert(context->inputSize.x >= config->k);
+    ASSERT_EQ(context->outputSize.z, config->k);
+    ASSERT_EQ(context->outputSize.y, context->nBatches);
+    ASSERT_EQ(context->outputSize.x, 1u);
+}
+
+static void moeGateForward_F32_F32(NnUint nThreads, NnUint threadIndex, NnUint batchSize, NnCpuOpContext *context) {
+    const NnMoeGateOpCodeConfig *config = (NnMoeGateOpCodeConfig *)context->opConfig;
+    float *indexes = (float *)context->buffers[config->indexesBufferIndex];
+
+    std::vector<NnUint> pos(config->k);
+    for (NnUint y = threadIndex; y < batchSize; y += nThreads) {
+        float *input = (float *)context->input[y];
+
+        topk_F32(input, pos.data(), context->inputSize.x, config->k);
+
+        float sum;
+        if (config->normTopk == 1u) {
+            sum = 0.0f;
+            for (NnUint i = 0u; i < config->k; i++)
+                sum += input[pos[i]];
+        } else {
+            sum = 1.0f;
+        }
+
+        for (NnUint k = 0u; k < config->k; k++) {
+            const NnUint p = pos[k];
+            indexes[y * config->k + k] = (float)p;
+
+            // (nActiveExperts, nBatches, 1)
+            float *output = (float *)context->output[k * context->outputSize.y + y];
+            *output = input[p] / sum;
+        }
+
+        DEBUG_VECTOR(context, "indexes", (&indexes[y * config->k]));
+    }
+}
+
 // device
 
 void printCpuInstructionSet() {
@@ -1602,16 +1621,24 @@ void printCpuInstructionSet() {
 NnCpuOpForwardInit getCpuOpForwardInit(NnOpCode code, NnOpQuantType quantType) {
     if (code == OP_EMBEDDING)
         return initEmbeddingForward;
+    if (code == OP_INV_RMS)
+        return initInvRmsForward;
     if (code == OP_RMS_NORM)
         return initRmsNormForward_ANY_F32_F32;
-    if (code == OP_ROPE_LLAMA)
-        return initRopeLlama3Forward;
+    if (code == OP_ROPE)
+        return initRopeForward_F32;
     if (code == OP_MULTIHEAD_ATT)
         return initMultiHeadAttForward;
     if (code == OP_MATMUL)
         return initMatmulForward;
+    if (code == OP_MUL)
+        return initMulForward;
     if (code == OP_CAST)
         return initCastForward;
+    if (code == OP_REPEAT_Z)
+        return initRepeatZForward;
+    if (code == OP_MOE_GATE)
+        return initMoeGateForward;
     return nullptr;
 }
 
@@ -1620,6 +1647,9 @@ NnCpuOpForward getCpuOpForward(NnOpCode code, NnOpQuantType quantType) {
         if (quantType == F32_F32_F32) return mergeAddForward_F32_F32;
         if (quantType == Q80_Q80_F32) return mergeAddForward_Q80_F32;
     }
+    if (code == OP_MERGE_SUM) {
+        if (quantType == F32_F32_F32) return mergeSumForward_F32_F32;
+    }
     if (code == OP_EMBEDDING) {
         if (quantType == F32_F32_F32) return embeddingForward_F32_F32_F32;
         if (quantType == F32_F32_Q80) return embeddingForward_F32_F32_Q80;
@@ -1635,8 +1665,8 @@ NnCpuOpForward getCpuOpForward(NnOpCode code, NnOpQuantType quantType) {
         if (quantType == F32_F32_F32) return matmulForward_F32_F32_F32;
         if (quantType == Q80_Q40_F32) return matmulForward_Q80_Q40_F32;
     }
-    if (code == OP_ROPE_LLAMA) {
-        if (quantType == F32_F32_F32) return ropeLlamaForward_F32_F32;
+    if (code == OP_ROPE) {
+        if (quantType == F32_F32_F32) return ropeForward_F32_F32;
     }
     if (code == OP_MULTIHEAD_ATT) {
         if (quantType == F32_F32_F32) return multiHeadAttForward_F32_F32;
@@ -1649,7 +1679,9 @@ NnCpuOpForward getCpuOpForward(NnOpCode code, NnOpQuantType quantType) {
     }
     if (code == OP_MUL) {
         if (quantType == F32_F32_F32) return mulForward_F32_F32;
-        if (quantType == Q80_Q80_F32) return mulForward_Q80_F32;
+    }
+    if (code == OP_SCALE) {
+        if (quantType == F32_F32_F32) return scaleForward_F32_F32;
     }
     if (code == OP_CAST) {
         if (quantType == F32_F32_F32) return castForward_ANY;
@@ -1657,8 +1689,17 @@ NnCpuOpForward getCpuOpForward(NnOpCode code, NnOpQuantType quantType) {
         if (quantType == Q80_Q80_Q80) return castForward_ANY;
         if (quantType == Q80_Q80_F32) return castForward_Q80_F32;
     }
+    if (code == OP_REPEAT_Z) {
+        if (quantType == F32_F32_Q80) return repeatZForward_F32_Q80;
+    }
     if (code == OP_SHIFT) {
         if (quantType == F32_F32_F32) return shiftForward_F32_F32;
     }
+    if (code == OP_SOFTMAX) {
+        if (quantType == F32_F32_F32) return softmaxForward_F32_F32;
+    }
+    if (code == OP_MOE_GATE) {
+        if (quantType == F32_F32_F32) return moeGateForward_F32_F32;
+    }
     return nullptr;
-}
+}
\ No newline at end of file
diff --git a/src/nn/nn-cpu-ops.hpp b/src/nn/nn-cpu-ops.hpp
index 2f5fc8aa..c9aa865d 100644
--- a/src/nn/nn-cpu-ops.hpp
+++ b/src/nn/nn-cpu-ops.hpp
@@ -29,7 +29,8 @@ typedef struct {
 
     NnByte *weight;
     NnSize3D weightSize;
-    // ADDED: Distributed execution parameters
+    
+    // distributed execution parameters
     NnUint nodeIndex;
     NnUint nNodes;
     
diff --git a/src/nn/nn-cpu.cpp b/src/nn/nn-cpu.cpp
index 0617d3fa..8e97be9a 100644
--- a/src/nn/nn-cpu.cpp
+++ b/src/nn/nn-cpu.cpp
@@ -123,7 +123,7 @@ NnDeviceSegment *NnCpuDevice::createSegment(NnUint segmentIndex) {
         opContext->bufferConfigs = nodeConfig->buffers;
         opContext->bufferFlags = bufferFlags;
 
-        // ADDED: Inject nodeIndex and nNodes into the context
+        // Inject nodeIndex and nNodes into the context
         opContext->nodeIndex = this->nodeConfig->nodeIndex;
         opContext->nNodes = this->netConfig->nNodes;
 
diff --git a/src/nn/nn-network.cpp b/src/nn/nn-network.cpp
index 659a9635..c2089e1a 100644
--- a/src/nn/nn-network.cpp
+++ b/src/nn/nn-network.cpp
@@ -831,8 +831,9 @@ void NnRootWeightLoader::writeWeight(NnUint nodeIndex, const char *opName, NnUin
 }
 
 NnSize NnRootWeightLoader::loadRoot(const char *opName, NnUint opIndex, NnSize nBytes, NnByte *weight) {
+    //set offset to 0 for loading the whole weight.
     try {
-        executor->loadWeight(opName, opIndex, nBytes, weight);
+        executor->loadWeight(opName, opIndex, 0, nBytes, weight);
     } catch (...) {
    
     }
@@ -840,8 +841,9 @@ NnSize NnRootWeightLoader::loadRoot(const char *opName, NnUint opIndex, NnSize n
 }
 
 NnSize NnRootWeightLoader::loadAll(const char *opName, NnUint opIndex, NnSize nBytes, NnByte *weight) {
+    //set offset to 0 for loading the whole weight.
     try {
-        executor->loadWeight(opName, opIndex, nBytes, weight);
+        executor->loadWeight(opName, opIndex, 0, nBytes, weight);
     } catch (...) {
 
     }
@@ -849,8 +851,9 @@ NnSize NnRootWeightLoader::loadAll(const char *opName, NnUint opIndex, NnSize nB
 }
 
 NnSize NnRootWeightLoader::loadRowMatmulSlices(const char *opName, const NnUint opIndex, const NnUint expertIndex, NnRowMatmulSlice *slice, NnByte *weight) {
+    //set offset to 0 for loading the whole weight.
     try {
-        executor->loadWeight(opName, opIndex, slice->size.nBytes, weight);
+        executor->loadWeight(opName, opIndex, 0, slice->size.nBytes, weight);
     } catch (...) {
 
     }
@@ -858,8 +861,9 @@ NnSize NnRootWeightLoader::loadRowMatmulSlices(const char *opName, const NnUint
 }
 
 NnSize NnRootWeightLoader::loadColMatmulSlices(const char *opName, const NnUint opIndex, const NnUint expertIndex, NnColMatmulSlice *slice, NnByte *weight) {
+    //set offset to 0 for loading the whole weight.
     try {
-        executor->loadWeight(opName, opIndex, slice->size.nBytes, weight);
+        executor->loadWeight(opName, opIndex, 0, slice->size.nBytes, weight);
     } catch (...) {
 
     }

From 126316aa6f467a0fade11c2266dc1220b2b61e0c Mon Sep 17 00:00:00 2001
From: "DESKTOP-68J82JA\\ecl" <eric0917579@gmail.com>
Date: Mon, 18 May 2026 19:38:53 +0800
Subject: [PATCH 4/5] delete some comments

---
 src/nn/nn-cpu-ops.cpp | 1 -
 1 file changed, 1 deletion(-)

diff --git a/src/nn/nn-cpu-ops.cpp b/src/nn/nn-cpu-ops.cpp
index 26a682d5..d2c5d5aa 100644
--- a/src/nn/nn-cpu-ops.cpp
+++ b/src/nn/nn-cpu-ops.cpp
@@ -277,7 +277,6 @@ static void matmul_F32_F32_F32(float *output, const float *x, const float *w, co
 #endif
 }
 
-// eric mod
 // define a struct to hold the distribution configuration for matmul operations, which can be reused across different matmul implementations
 struct DistConfig {
     NnUint d_start;

From 4c1b699644d879ee695fae9b0a252cf62ce39e72 Mon Sep 17 00:00:00 2001
From: "DESKTOP-68J82JA\\ecl" <eric0917579@gmail.com>
Date: Fri, 22 May 2026 19:52:53 +0800
Subject: [PATCH 5/5] optimized the code on loading weight & matmul

---
 converter/convert-tokenizer-hf.py |  15 +-
 converter/requirements.txt        |   7 +-
 src/api-types.hpp                 |  47 +----
 src/app.cpp                       |  43 +++--
 src/llm.cpp                       |  20 +--
 src/llm.hpp                       |   2 +-
 src/nn/nn-cpu-ops.cpp             | 284 ++++++++++--------------------
 src/nn/nn-cpu-ops.hpp             |   5 -
 src/nn/nn-cpu.cpp                 |   4 -
 src/nn/nn-network.cpp             | 117 +++++++++---
 src/nn/nn-network.hpp             |  18 ++
 11 files changed, 246 insertions(+), 316 deletions(-)

diff --git a/converter/convert-tokenizer-hf.py b/converter/convert-tokenizer-hf.py
index 203c4089..e9c8e0e3 100644
--- a/converter/convert-tokenizer-hf.py
+++ b/converter/convert-tokenizer-hf.py
@@ -34,7 +34,6 @@ def __init__(self, dirPath, tokenizerConfig):
     def resolvePreTrainedTokenizerFast(self):
         utb = unicodeToBytes()
         tokenizer = PreTrainedTokenizerFast(tokenizer_file = os.path.join(self.dirPath, 'tokenizer.json'))
-        config = openJson(os.path.join(self.dirPath, 'config.json'))
         vocabLen = len(tokenizer.get_vocab())
         for i in range(vocabLen):
             tokenChars = list(tokenizer.convert_ids_to_tokens([i])[0])
@@ -47,18 +46,11 @@ def resolvePreTrainedTokenizerFast(self):
             self.tokens.append(bytes(tokenBytes))
             self.scores.append(-float(i))
 
-        # Pad tokenizer vocab to match model vocab_size if needed
-        targetVocabSize = config.get('vocab_size', vocabLen)
-        if targetVocabSize > vocabLen:
-            print(f'⚠️ Padding tokenizer vocab from {vocabLen} to {targetVocabSize}')
-            for i in range(vocabLen, targetVocabSize):
-                self.tokens.append(f'<|reserved_{i}|>'.encode('utf-8'))
-                self.scores.append(-float(i))
-
         self.bosId = tokenizer.bos_token_id
         if (tokenizer.eos_token_id):
             self.eosIds = [tokenizer.eos_token_id]
-        if (self.bosId is None or self.eosIds is None):
+        if (self.bosId is None or self.eosId is None):
+            config = openJson(os.path.join(self.dirPath, 'config.json'))
             if (self.bosId is None):
                 self.bosId = config['bos_token_id']
             if (self.eosIds is None):
@@ -91,8 +83,7 @@ def resolveLlamaTokenizer(self):
 
     def resolve(self):
         cls = self.tokenizerConfig['tokenizer_class']
-        if (cls == 'PreTrainedTokenizer' or
-            cls == 'PreTrainedTokenizerFast' or
+        if (cls == 'PreTrainedTokenizerFast' or
             cls == 'LlamaTokenizerFast' or
             cls == 'Qwen2Tokenizer'):
             return self.resolvePreTrainedTokenizerFast()
diff --git a/converter/requirements.txt b/converter/requirements.txt
index 0e866177..221c48df 100644
--- a/converter/requirements.txt
+++ b/converter/requirements.txt
@@ -1,6 +1,5 @@
-# python>=3.9
+python>=3.9
 numpy==1.23.5
-torch==2.0.1+cpu --index-url https://download.pytorch.org/whl/cpu
+pytorch==2.0.1
 safetensors==0.4.2
-sentencepiece==0.1.99
-transformers==4.57.6
+sentencepiece==0.1.99
\ No newline at end of file
diff --git a/src/api-types.hpp b/src/api-types.hpp
index 404b344e..62492082 100755
--- a/src/api-types.hpp
+++ b/src/api-types.hpp
@@ -264,8 +264,6 @@ void to_json(json& j, const ModelList& models) {
         {"data", models.data}};
 }
 
-static std::string normalizeMessageContent(const json &content);
-
 std::vector<ChatMessage> parseChatMessages(json &json){
     std::vector<ChatMessage> messages;
     messages.reserve(json.size());
@@ -274,7 +272,7 @@ std::vector<ChatMessage> parseChatMessages(json &json){
         ChatMessage msg;
         msg.role = item["role"].template get<std::string>();
         if (item.contains("content") && !item["content"].is_null())
-            msg.content = normalizeMessageContent(item["content"]);
+            msg.content = item["content"].template get<std::string>();
         if (item.contains("tool_call_id"))
             msg.tool_call_id = item["tool_call_id"].template get<std::string>();
         if (item.contains("tool_calls") && item["tool_calls"].is_array()) {
@@ -303,43 +301,6 @@ std::vector<ChatMessage> parseChatMessages(json &json){
     return messages;
 }
 
-static std::string normalizeMessageContent(const json &content) {
-    if (content.is_null())
-        return "";
-    if (content.is_string())
-        return content.template get<std::string>();
-    if (content.is_array()) {
-        std::string result;
-        for (const auto &part : content) {
-            std::string piece;
-            if (part.is_string()) {
-                piece = part.template get<std::string>();
-            } else if (part.is_object()) {
-                if (part.contains("type") && part["type"].is_string()) {
-                    const std::string type = part["type"].template get<std::string>();
-                    if ((type == "text" || type == "input_text") && part.contains("text") && part["text"].is_string()) {
-                        piece = part["text"].template get<std::string>();
-                    } else if (type == "text" && part.contains("content") && part["content"].is_string()) {
-                        piece = part["content"].template get<std::string>();
-                    }
-                } else if (part.contains("text") && part["text"].is_string()) {
-                    piece = part["text"].template get<std::string>();
-                }
-            }
-
-            if (piece.empty())
-                continue;
-            if (!result.empty() && result.back() != '\n')
-                result += ' ';
-            result += piece;
-        }
-        return result;
-    }
-    if (content.is_object())
-        return content.dump();
-    return content.dump();
-}
-
 InferenceParams parseInferenceParams(json &json, float defaultTemperature, float defaultTopp, unsigned long long defaultSeed) {
     InferenceParams params;
     params.temperature = defaultTemperature;
@@ -400,11 +361,7 @@ InferenceParams parseInferenceParams(json &json, float defaultTemperature, float
         }
     }
     if (json.contains("stop")) {
-        if (json["stop"].is_string()) {
-            params.stop = std::vector<std::string>{json["stop"].template get<std::string>()};
-        } else {
-            params.stop = json["stop"].template get<std::vector<std::string>>();
-        }
+        params.stop = json["stop"].template get<std::vector<std::string>>();
     } else {
         const std::string defaultStop = "<|eot_id|>";
         params.stop = std::vector<std::string>{defaultStop};
diff --git a/src/app.cpp b/src/app.cpp
index 9f717f8f..8e9229d3 100644
--- a/src/app.cpp
+++ b/src/app.cpp
@@ -276,8 +276,14 @@ void runInferenceApp(AppCliArgs *args, void (*handler)(AppInferenceContext *cont
     std::vector<NnExecutorDevice> devices = resolveDevices(args, &net.netConfig, rootNodeConfig, &execution);
     NnExecutor executor(&net.netConfig, rootNodeConfig, &devices, &execution, synchronizer.get(), args->benchmark);
 
-    NnRootWeightLoader weightLoader(&executor, network, nNodes);
+    NnLocalWeightLoader weightLoader(&executor, 0, nNodes);
     loadLlmNetWeight(args->modelPath, &net, &weightLoader);
+    if (network != nullptr) {
+        printf("💿 Waiting for workers to load weights...\n");
+        for (NnUint socketIndex = 0; socketIndex < nNodes - 1; socketIndex++)
+            network->readAck(socketIndex);
+        printf("💿 All workers ready\n");
+    }
 
     RootLlmInference inference(&net, &execution, &executor, network);
 
@@ -322,27 +328,26 @@ void runWorkerApp(AppCliArgs *args) {
         NnNetworkNodeSynchronizer synchronizer(network, &execution, &netConfig, &nodeConfig);
         NnExecutor executor(&netConfig, &nodeConfig, &devices, &execution, &synchronizer, false);
 
-        // modify the worker's weight loading method to get the weight from disk.
-        if (args->modelPath == nullptr) {
-            throw std::runtime_error("Worker needs --model argument to load weights locally!");
-        }
-
-        printf("💿 Worker is loading weights locally from %s...\n", args->modelPath);
-
-        // header for loading weights.
-        LlmHeader header = loadLlmHeader(args->modelPath, args->maxSeqLen, args->syncType);
+        if (args->modelPath == nullptr)
+            throw std::runtime_error("--model is required for worker mode");
 
-        // build a temporary net to load the weights, the net will be released right after loading the weights.
-        LlmNet localNet = buildLlmNet(&header, netConfig.nNodes, netConfig.nBatches);
-
-        // load weights from disk to the executor.
-        NnRootWeightLoader localWeightLoader(&executor, nullptr, netConfig.nNodes); 
-        loadLlmNetWeight(args->modelPath, &localNet, &localWeightLoader);
+        NnFloatType syncType = F_32;
+        for (NnUint i = 0; i < netConfig.nPipes; i++) {
+            if (std::strcmp(netConfig.pipes[i].name, "ZQ") == 0) {
+                syncType = netConfig.pipes[i].size.floatType;
+                break;
+            }
+        }
 
-        // release the temporary net to free memory, since the weights have been loaded into the executor.
-        releaseLlmNet(&localNet);
+        {
+            LlmHeader workerHeader = loadLlmHeader(args->modelPath, 0, syncType);
+            LlmNet workerNet = buildLlmNet(&workerHeader, netConfig.nNodes, netConfig.nBatches);
+            std::unique_ptr<LlmNet, void(*)(LlmNet *)> workerNetPtr(&workerNet, releaseLlmNet);
+            NnLocalWeightLoader weightLoader(&executor, nodeConfig.nodeIndex, netConfig.nNodes);
+            loadLlmNetWeight(args->modelPath, &workerNet, &weightLoader);
+        }
 
-        printf("✅ Worker local weights loaded successfully.\n");
+        network->writeAck(ROOT_SOCKET_INDEX);
 
         WorkerLlmInference inference(&execution, network);
         bool isFirstAttempt = true;
diff --git a/src/llm.cpp b/src/llm.cpp
index 59f20ed8..a89dad4f 100644
--- a/src/llm.cpp
+++ b/src/llm.cpp
@@ -304,21 +304,19 @@ LlmNet buildLlmNet(LlmHeader *h, NnUint nNodes, NnUint nBatches) {
                 OP_MATMUL, "block_matmul_q", layerIndex,
                 pointerBatchConfig(SRC_BUFFER, yqBufferIndex),
                 pointerBatchConfig(SRC_BUFFER, qBufferIndex),
-
-                // modify the net config to fit the whole weight.
-                size2D(h->weightType, h->dim, h->dim),
+                size2D(h->weightType, n.qSlice.n, n.qSlice.d0),
                 NnMatmulOpConfig{0, 0, moeExpertIndexesBufferIndex});
             att.addOp(
                 OP_MATMUL, "block_matmul_k", layerIndex,
                 pointerBatchConfig(SRC_BUFFER, yqBufferIndex),
                 pointerBatchConfig(SRC_BUFFER, kTempBufferIndex),
-                size2D(h->weightType, h->dim, h->kvDim),
+                size2D(h->weightType, n.kSlice.n, n.kSlice.d0),
                 NnMatmulOpConfig{0, 0, moeExpertIndexesBufferIndex});
             att.addOp(
                 OP_MATMUL, "block_matmul_v", layerIndex,
                 pointerBatchConfig(SRC_BUFFER, yqBufferIndex),
                 pointerBatchConfig(SRC_BUFFER, vTempBufferIndex),
-                size2D(h->weightType, h->dim, h->kvDim),
+                size2D(h->weightType, n.vSlice.n, n.vSlice.d0),
                 NnMatmulOpConfig{0, 0, moeExpertIndexesBufferIndex});
 
             if (h->archType == QWEN3 || h->archType == QWEN3_MOE) {
@@ -394,7 +392,7 @@ LlmNet buildLlmNet(LlmHeader *h, NnUint nNodes, NnUint nBatches) {
                 OP_MATMUL, "block_matmul_wo", layerIndex,
                 pointerBatchConfig(SRC_BUFFER, zqSliceBufferIndex),
                 pointerBatchConfig(SRC_BUFFER, yBufferIndex),
-                size2D(h->weightType, h->dim, h->dim),
+                size2D(h->weightType, n.woSlice.n0, n.woSlice.d),
                 NnMatmulOpConfig{0, 0, moeExpertIndexesBufferIndex});
             att.addOp(
                 OP_CAST, "block_cast_d", layerIndex,
@@ -512,13 +510,13 @@ LlmNet buildLlmNet(LlmHeader *h, NnUint nNodes, NnUint nBatches) {
                     OP_MATMUL, "block_matmul_w1", layerIndex,
                     pointerBatchConfig(SRC_BUFFER, yqBufferIndex),
                     pointerBatchConfig(SRC_BUFFER, dBufferIndex),
-                    size2D(h->weightType, h->dim, h->hiddenDim),
+                    size2D(h->weightType, n.w1Slice.n, n.w1Slice.d0),
                     NnMatmulOpConfig{0, 0, moeExpertIndexesBufferIndex});
                 ff.addOp(
                     OP_MATMUL, "block_matmul_w3", layerIndex,
                     pointerBatchConfig(SRC_BUFFER, yqBufferIndex),
                     pointerBatchConfig(SRC_BUFFER, lBufferIndex),
-                    size2D(h->weightType, h->dim, h->hiddenDim),
+                    size2D(h->weightType, n.w3Slice.n, n.w3Slice.d0),
                     NnMatmulOpConfig{0, 0, moeExpertIndexesBufferIndex});
                 ff.addOp(
                     OP_SILU, "block_act", layerIndex,
@@ -544,7 +542,7 @@ LlmNet buildLlmNet(LlmHeader *h, NnUint nNodes, NnUint nBatches) {
                     OP_MATMUL, "block_matmul_w2", layerIndex,
                     pointerBatchConfig(SRC_BUFFER, dqBufferIndex),
                     pointerBatchConfig(SRC_BUFFER, yBufferIndex),
-                    size2D(h->weightType, h->hiddenDim, h->dim),
+                    size2D(h->weightType, n.w2Slice.n0, n.w2Slice.d),
                     NnMatmulOpConfig{0, 0, moeExpertIndexesBufferIndex});
             }
             ff.addOp(
@@ -590,7 +588,7 @@ LlmNet buildLlmNet(LlmHeader *h, NnUint nNodes, NnUint nBatches) {
             OP_MATMUL, "final_matmul_logits", 0,
             pointerBatchConfig(SRC_BUFFER, yqBufferIndex),
             pointerBatchConfig(SRC_BUFFER, logitsSliceBufferIndex),
-            size2D(h->weightType, h->dim, h->vocabSize),
+            size2D(h->weightType, n.wclsSlice.n, n.wclsSlice.d0),
             NnMatmulOpConfig{});
         end.addOp(
             OP_CAST, "final_cast_logits", 0,
@@ -613,7 +611,7 @@ void releaseLlmNet(LlmNet *net) {
     delete[] net->nodeConfigs;
 }
 
-void loadLlmNetWeight(const char *path, LlmNet *net, NnRootWeightLoader *loader) {
+void loadLlmNetWeight(const char *path, LlmNet *net, NnLocalWeightLoader *loader) {
     MmapFile file;
     openMmapFile(&file, path, net->header->fileSize);
 #if DEBUG_USE_MMAP_FOR_WEIGHTS
diff --git a/src/llm.hpp b/src/llm.hpp
index 5ed7b983..19987848 100644
--- a/src/llm.hpp
+++ b/src/llm.hpp
@@ -99,6 +99,6 @@ LlmHeader loadLlmHeader(const char* path, const unsigned int maxSeqLen, NnFloatT
 void printLlmHeader(LlmHeader *header);
 LlmNet buildLlmNet(LlmHeader *h, NnUint nNodes, NnUint nBatches);
 void releaseLlmNet(LlmNet *net);
-void loadLlmNetWeight(const char* path, LlmNet *net, NnRootWeightLoader *loader);
+void loadLlmNetWeight(const char* path, LlmNet *net, NnLocalWeightLoader *loader);
 
 #endif
\ No newline at end of file
diff --git a/src/nn/nn-cpu-ops.cpp b/src/nn/nn-cpu-ops.cpp
index d2c5d5aa..2b6f51cf 100644
--- a/src/nn/nn-cpu-ops.cpp
+++ b/src/nn/nn-cpu-ops.cpp
@@ -188,142 +188,77 @@ static void rmsNorm_Q80_F32_F32(float *output, const NnBlockQ80 *x, const float
     }
 }
 
-//MODIFY: Added parameters for distributed execution and logic to calculate local slice boundaries
-static void matmul_F32_F32_F32(float *output, const float *x, const float *w, const NnUint n, const NnUint d, const NnUint nThreads,/* ADDED */ const NnUint threadIndex, const NnUint nodeIndex, const NnUint nNodes, bool isRowMatmul) {
-    // ADDED: Offset Calculation Logic
-    NnUint d_start = 0, d_end = d;
-    NnUint n_start = 0, n_end = n;
-
-    if (nNodes > 1) {
-        if (isRowMatmul) {
-            // Cut along the output dimension (d)
-            NnUint slice_d = d / nNodes;
-            d_start = nodeIndex * slice_d;
-            d_end = d_start + slice_d;
-        } else {
-            // Cut along the input dimension (n)
-            NnUint slice_n = n / nNodes;
-            n_start = nodeIndex * slice_n;
-            n_end = n_start + slice_n;
-        }
-    }
-
-    // Assign threads only to the slice this node is responsible for
-    NnUint d_slice = d_end - d_start;
-    SPLIT_THREADS(thread_d_start, thread_d_end, d_slice, nThreads, threadIndex);
-
-    // Calculate actual indices mapping back to the full matrix
-    NnUint actual_start = d_start + thread_d_start;
-    NnUint actual_end   = d_start + thread_d_end;
-
+static void matmul_F32_F32_F32(float *output, const float *x, const float *w, const NnUint n, const NnUint d, const NnUint nThreads, const NnUint threadIndex) {
+    SPLIT_THREADS(start, end, d, nThreads, threadIndex);
     unsigned int i, j;
 #if defined(__ARM_NEON)
     assert(n % 4 == 0);
     float32x4_t q;
     float32x4_t p;
     float32x4_t z;
-
-    for (i = actual_start; i < actual_end; i++) {
+    for (i = start; i < end; i++) {
         z = vmovq_n_f32(0);
-        for (j = n_start; j < n_end; j += 4) {
-
-            const float *x_ptr = isRowMatmul ? &x[j] : &x[j - n_start];
-            q = vld1q_f32(x_ptr);
-
+        for (j = 0; j < n; j += 4) {
+            q = vld1q_f32(&x[j]);
             p = vld1q_f32(&w[i * n + j]);
             z = vfmaq_f32(z, q, p);
         }
-
-        if (isRowMatmul) {
-            output[i - d_start] = vaddvq_f32(z);
-        } else {
-            output[i] = vaddvq_f32(z);
-        }
+        output[i] = vaddvq_f32(z);
     }
 #elif defined(__AVX2__)
     assert(n % 8 == 0);
     __m256 a0, b0, u;
-    for (i = actual_start; i < actual_end; i++) {
+    for (i = start; i < end; i++) {
         u = _mm256_set1_ps(0.0f);
-        for (j = n_start; j < n_end; j += 8) {
-
-            const float *x_ptr = isRowMatmul ? &x[j] : &x[j - n_start];
-            a0 = _mm256_loadu_ps(x_ptr);
-
+        for (j = 0; j < n; j += 8) {
+            a0 = _mm256_loadu_ps(&x[j]);
             b0 = _mm256_loadu_ps(&w[i * n + j]);
             u = _mm256_fmadd_ps(a0, b0, u);
         }
-
-        if (isRowMatmul) {
-            output[i - d_start] = horizontalSum_avx2(u);
-        } else {
-            output[i] = horizontalSum_avx2(u);
-        }
+        output[i] = horizontalSum_avx2(u);
     }
 #else
-    for (i = actual_start; i < actual_end; i++) {
+    for (i = start; i < end; i++) {
         float val = 0.0f;
-        for (j = n_start; j < n_end; j++) {
-            float x_val = isRowMatmul ? x[j] : x[j - n_start];
-            val += w[i * n + j] * x_val;
-        }
-
-        if (isRowMatmul) {
-            output[i - d_start] = val;
-        } else {
-            output[i] = val;
+        for (j = 0; j < n; j++) {
+            val += w[i * n + j] * x[j];
         }
+        output[i] = val;
     }
 #endif
 }
 
-// define a struct to hold the distribution configuration for matmul operations, which can be reused across different matmul implementations
-struct DistConfig {
-    NnUint d_start;
-    NnUint d_slice;
-    NnUint nBlocks_start;
-    NnUint nBlocks_end;
-};
-
-// use template to generate both row and column matmul implementations without runtime branching, and pass the distribution configuration as a parameter
-template <bool IsRowMatmul>
-static void matmul_Q80_Q40_F32_impl(float *output, const NnBlockQ80 *x, const NnBlockQ40 *w, const NnUint n, const NnUint nThreads, const NnUint threadIndex, const DistConfig& dist) {
+static void matmul_Q80_Q40_F32(float *output, const NnBlockQ80 *x, const NnBlockQ40 *w, const NnUint n, const NnUint d, const NnUint nThreads, const NnUint threadIndex) {
+    SPLIT_THREADS(start, end, d, nThreads, threadIndex);
+    assert(n % Q40_BLOCK_SIZE == 0);
     const unsigned int nBlocks = n / Q40_BLOCK_SIZE;
 
-    SPLIT_THREADS(thread_d_start, thread_d_end, dist.d_slice, nThreads, threadIndex);
-    NnUint actual_start = dist.d_start + thread_d_start;
-    NnUint actual_end   = dist.d_start + thread_d_end;
-
-    const NnBlockQ80* x_base = IsRowMatmul ? x : (x - dist.nBlocks_start);
-    float* output_base = IsRowMatmul ? (output - dist.d_start) : output;
-
 #if defined(__ARM_NEON)
     const uint8x16_t m4b = vdupq_n_u8(0x0F);
     const int8x16_t s8b = vdupq_n_s8(0x8);
 
-    for (unsigned int di = actual_start; di < actual_end; di++) {
+    for (unsigned int di = start; di < end; di++) {
         float32x4_t sumv0 = vmovq_n_f32(0.0f);
         float32x4_t sumv1 = vmovq_n_f32(0.0f);
         float32x4_t sumv2 = vmovq_n_f32(0.0f);
         float32x4_t sumv3 = vmovq_n_f32(0.0f);
 
-        unsigned int j = dist.nBlocks_start;
+        unsigned int j = 0;
         
 #if defined(__ARM_FEATURE_DOTPROD)
-        for (; j + 3 < dist.nBlocks_end; j += 4) {
+        for (; j + 3 < nBlocks; j += 4) {
             __builtin_prefetch(&w[di * nBlocks + j + 4]);
-            
-            __builtin_prefetch(&x_base[j + 4]);
+            __builtin_prefetch(&x[j + 4]);
 
             const NnBlockQ40 *w0 = &w[di * nBlocks + j];
             const NnBlockQ40 *w1 = &w[di * nBlocks + j + 1];
             const NnBlockQ40 *w2 = &w[di * nBlocks + j + 2];
             const NnBlockQ40 *w3 = &w[di * nBlocks + j + 3];
 
-            const NnBlockQ80 *x0 = &x_base[j];
-            const NnBlockQ80 *x1 = &x_base[j + 1];
-            const NnBlockQ80 *x2 = &x_base[j + 2];
-            const NnBlockQ80 *x3 = &x_base[j + 3];
+            const NnBlockQ80 *x0 = &x[j];
+            const NnBlockQ80 *x1 = &x[j + 1];
+            const NnBlockQ80 *x2 = &x[j + 2];
+            const NnBlockQ80 *x3 = &x[j + 3];
 
             int8x16_t w0l = vsubq_s8(vreinterpretq_s8_u8(vandq_u8(vld1q_u8(w0->qs), m4b)), s8b);
             int8x16_t w0h = vsubq_s8(vreinterpretq_s8_u8(vshrq_n_u8(vld1q_u8(w0->qs), 4)), s8b);
@@ -354,12 +289,11 @@ static void matmul_Q80_Q40_F32_impl(float *output, const NnBlockQ80 *x, const Nn
             sumv3 = vmlaq_n_f32(sumv3, vcvtq_f32_s32(p3), CONVERT_F16_TO_F32(w3->d) * CONVERT_F16_TO_F32(x3->d));
         }
 #else
-        for (; j + 1 < dist.nBlocks_end; j += 2) {
+        for (; j + 1 < nBlocks; j += 2) {
             const NnBlockQ40 *w0 = &w[di * nBlocks + j];
             const NnBlockQ40 *w1 = &w[di * nBlocks + j + 1];
-
-            const NnBlockQ80 *x0 = &x_base[j];
-            const NnBlockQ80 *x1 = &x_base[j + 1];
+            const NnBlockQ80 *x0 = &x[j];
+            const NnBlockQ80 *x1 = &x[j + 1];
 
             const uint8x16_t w0qs = vld1q_u8(w0->qs);
             const uint8x16_t w1qs = vld1q_u8(w1->qs);
@@ -394,9 +328,9 @@ static void matmul_Q80_Q40_F32_impl(float *output, const NnBlockQ80 *x, const Nn
         }
 #endif
 
-        for (; j < dist.nBlocks_end; j++) {
+        for (; j < nBlocks; j++) {
             const NnBlockQ40 *wb = &w[di * nBlocks + j];
-            const NnBlockQ80 *xb = &x_base[j];
+            const NnBlockQ80 *xb = &x[j];
 
             const uint8x16_t wqs = vld1q_u8(wb->qs);
             const int8x16_t wl = vsubq_s8(vreinterpretq_s8_u8(vandq_u8(wqs, m4b)), s8b);
@@ -421,15 +355,14 @@ static void matmul_Q80_Q40_F32_impl(float *output, const NnBlockQ80 *x, const Nn
             sumv0 = vmlaq_n_f32(sumv0, vcvtq_f32_s32(p), s);
         }
 
-        float total_sum = vaddvq_f32(sumv0) + vaddvq_f32(sumv1) + vaddvq_f32(sumv2) + vaddvq_f32(sumv3);
-        output_base[di] = total_sum;
+        output[di] = vaddvq_f32(sumv0) + vaddvq_f32(sumv1) + vaddvq_f32(sumv2) + vaddvq_f32(sumv3);
     }
 #elif defined(__AVX512F__)
-    for (NnUint i = actual_start; i < actual_end; i++) {
+    for (NnUint i = start; i < end; i++) {
         float sum = 0.0f;
-        for (NnUint j = dist.nBlocks_start; j < dist.nBlocks_end; j++) {
+        for (NnUint j = 0; j < nBlocks; j++) {
             const NnBlockQ40 *wb = &w[i * nBlocks + j];
-            const NnBlockQ80 *xb = &x_base[j];
+            const NnBlockQ80 *xb = &x[j];
             const float s = CONVERT_F16_TO_F32(wb->d) * CONVERT_F16_TO_F32(xb->d);
 
             __m128i w8 = _mm_loadu_si128((const __m128i*)wb->qs);
@@ -449,14 +382,14 @@ static void matmul_Q80_Q40_F32_impl(float *output, const NnBlockQ80 *x, const Nn
             __m512i products = _mm512_madd_epi16(w16, x16);
             sum += _mm512_reduce_add_epi32(products) * s;
         }
-        output_base[i] = sum;
+        output[i] = sum;
     }
 #elif defined(__AVX2__)
-    for (NnUint i = actual_start; i < actual_end; i++) {
+    for (NnUint i = start; i < end; i++) {
         float sum = 0.0f;
-        for (NnUint j = dist.nBlocks_start; j < dist.nBlocks_end; j++) {
+        for (NnUint j = 0; j < nBlocks; j++) {
             const NnBlockQ40 *wb = &w[i * nBlocks + j];
-            const NnBlockQ80 *xb = &x_base[j];
+            const NnBlockQ80 *xb = &x[j];
             const float s = CONVERT_F16_TO_F32(wb->d) * CONVERT_F16_TO_F32(xb->d);
 
             __m128i w_packed = _mm_loadu_si128((const __m128i*)wb->qs);
@@ -493,16 +426,14 @@ static void matmul_Q80_Q40_F32_impl(float *output, const NnBlockQ80 *x, const Nn
 
             sum += block_sum * s;
         }
-        output_base[i] = sum;
+        output[i] = sum;
     }
 #else
-    for (NnUint i = actual_start; i < actual_end; i++) {
+    for (NnUint i = start; i < end; i++) {
         float sum = 0.0;
-        for (NnUint j = dist.nBlocks_start; j < dist.nBlocks_end; j++) {
-
+        for (NnUint j = 0; j < nBlocks; j++) {
             const NnBlockQ40 *wb = &w[i * nBlocks + j];
-            const NnBlockQ80 *xb = &x_base[j];
-
+            const NnBlockQ80 *xb = &x[j];
             const float s = CONVERT_F16_TO_F32(wb->d) * CONVERT_F16_TO_F32(xb->d);
             for (NnUint k = 0; k < Q40_BLOCK_SIZE / 2; k++) {
                 const int w0 = (wb->qs[k] & 0x0F) - 8;
@@ -512,43 +443,11 @@ static void matmul_Q80_Q40_F32_impl(float *output, const NnBlockQ80 *x, const Nn
                 sum += (w0 * i1 + w1 * i2) * s;
             }
         }
-        output_base[i] = sum;
+        output[i] = sum;
     }
 #endif
 }
 
-// wrapper function
-static void matmul_Q80_Q40_F32(float *output, const NnBlockQ80 *x, const NnBlockQ40 *w, const NnUint n, const NnUint d, const NnUint nThreads, const NnUint threadIndex, const NnUint nodeIndex, const NnUint nNodes, bool isRowMatmul) {
-    DistConfig dist;
-    const unsigned int nBlocks = n / Q40_BLOCK_SIZE;
-
-    if (nNodes > 1) {
-        if (isRowMatmul) {
-            dist.d_slice = d / nNodes;
-            dist.d_start = nodeIndex * dist.d_slice;
-            dist.nBlocks_start = 0;
-            dist.nBlocks_end = nBlocks;
-        } else {
-            NnUint slice_nBlocks = nBlocks / nNodes;
-            dist.d_start = 0;
-            dist.d_slice = d;
-            dist.nBlocks_start = nodeIndex * slice_nBlocks;
-            dist.nBlocks_end = dist.nBlocks_start + slice_nBlocks;
-        }
-    } else {
-        dist.d_start = 0;
-        dist.d_slice = d;
-        dist.nBlocks_start = 0;
-        dist.nBlocks_end = nBlocks;
-    }
-
-    if (isRowMatmul) {
-        matmul_Q80_Q40_F32_impl<true>(output, x, w, n, nThreads, threadIndex, dist);
-    } else {
-        matmul_Q80_Q40_F32_impl<false>(output, x, w, n, nThreads, threadIndex, dist);
-    }
-}
-
 #define SQRT_2_OVER_PI 0.79788456080286535587989211986876f
 #define GELU_COEF_A 0.044715f
 
@@ -1202,9 +1101,14 @@ static void rmsNormForward_Q80_F32_F32(NnUint nThreads, NnUint threadIndex, NnUi
 }
 
 static void initMatmulForward(NnCpuOpContext *context) {
-
+    const NnMatmulOpConfig *config = (NnMatmulOpConfig *)context->opConfig;
     ASSERT_EQ(context->inputSize.y, context->nBatches);
     ASSERT_EQ(context->outputSize.y, context->nBatches);
+    ASSERT_EQ(context->inputSize.x, context->weightSize.y);
+    ASSERT_EQ(context->inputSize.z, std::max(config->nActiveExperts, 1u));
+    ASSERT_EQ(context->outputSize.x, context->weightSize.x);
+    ASSERT_EQ(context->outputSize.z, std::max(config->nActiveExperts, 1u));
+    ASSERT_EQ(context->weightSize.z, std::max(config->nExperts, 1u));
 
     if (!context->hasInputContinuousMemory)
         printf("🚧 Op %s does not have contiguous memory for input\n", context->name);
@@ -1232,66 +1136,58 @@ static bool matmulForward_llamafile(NnUint nThreads, NnUint threadIndex, NnUint
 }
 
 static void matmulForward_F32_F32_F32(NnUint nThreads, NnUint threadIndex, NnUint batchSize, NnCpuOpContext *context) {
-
-    NnUint nNodes = context->nNodes;
-
-    if (nNodes == 1 && matmulForward_llamafile(nThreads, threadIndex, batchSize, context))
+    if (matmulForward_llamafile(nThreads, threadIndex, batchSize, context))
         return;
 
-    const float *weight = (float *)context->weight;
-
-    // extract node index for distributed execution
-    NnUint nodeIndex = context->nodeIndex;
+    const NnMatmulOpConfig *config = (NnMatmulOpConfig *)context->opConfig;
+    const NnUint nActiveExpertsOr1 = std::max(config->nActiveExperts, 1u);
+    const float *activeExpertIndexes = (const float *)context->buffers[config->activeExpertIndexesBufferIndex];
 
-    // determine row/col split.
-    bool isRowMatmul = (context->inputSize.x == context->weightSize.y);
+    for (NnUint y = 0; y < batchSize; y++) {
+        for (NnUint e = 0; e < nActiveExpertsOr1; e++) {
+            const NnUint activeExpertIndex = config->nActiveExperts == 0u
+                ? 0u
+                : (NnUint)activeExpertIndexes[y * config->nActiveExperts + e];
 
-    for (NnUint batchIndex = 0; batchIndex < batchSize; batchIndex++) {
-        float *input = (float *)context->input[batchIndex];
-        float *output = (float *)context->output[batchIndex];
-        DEBUG_VECTOR(context, "input", input);
-        matmul_F32_F32_F32(
-            output,
-            input,
-            weight,
-            context->weightSize.y,
-            context->weightSize.x,
-            nThreads,
-            threadIndex,
-            nodeIndex, 
-            nNodes, 
-            isRowMatmul);
-        DEBUG_VECTOR(context, "output", output);
+            float *output = (float *)context->output[e * context->outputSize.y + y];
+            matmul_F32_F32_F32(
+                output,
+                (float *)context->input[e * context->inputSize.y + y],
+                (float *)&context->weight[activeExpertIndex * context->weightSize.nBytesXY],
+                context->weightSize.y,
+                context->weightSize.x,
+                nThreads,
+                threadIndex);
+            DEBUG_VECTOR(context, "output", output);
+        }
     }
 }
 
 static void matmulForward_Q80_Q40_F32(NnUint nThreads, NnUint threadIndex, NnUint batchSize, NnCpuOpContext *context) {
-
-    NnUint nNodes = context->nNodes;
-
-    if (nNodes == 1 && matmulForward_llamafile(nThreads, threadIndex, batchSize, context))
+    if (matmulForward_llamafile(nThreads, threadIndex, batchSize, context))
         return;
 
-    const NnBlockQ40 *weight = (NnBlockQ40 *)context->weight;
+    const NnMatmulOpConfig *config = (NnMatmulOpConfig *)context->opConfig;
+    const NnUint nActiveExpertsOr1 = std::max(config->nActiveExperts, 1u);
+    const float *activeExpertIndexes = (const float *)context->buffers[config->activeExpertIndexesBufferIndex];
 
-    NnUint nodeIndex = context->nodeIndex;
+    for (NnUint y = 0; y < batchSize; y++) {
+        for (NnUint e = 0; e < nActiveExpertsOr1; e++) {
+            const NnUint activeExpertIndex = config->nActiveExperts == 0u
+                ? 0u
+                : (NnUint)activeExpertIndexes[y * config->nActiveExperts + e];
 
-    bool isRowMatmul = (context->inputSize.x == context->weightSize.y);;
-
-    for (NnUint batchIndex = 0; batchIndex < batchSize; batchIndex++) {
-        NnBlockQ80 *input = (NnBlockQ80 *)context->input[batchIndex];
-        float *output = (float *)context->output[batchIndex];
-        matmul_Q80_Q40_F32(
-            output,
-            input,
-            weight,
-            context->weightSize.y,
-            context->weightSize.x,
-            nThreads,
-            threadIndex,
-            nodeIndex, 
-            nNodes, 
-            isRowMatmul);
+            float *output = (float *)context->output[e * context->outputSize.y + y];
+            matmul_Q80_Q40_F32(
+                output,
+                (NnBlockQ80 *)context->input[e * context->inputSize.y + y],
+                (NnBlockQ40 *)&context->weight[activeExpertIndex * context->weightSize.nBytesXY],
+                context->weightSize.y,
+                context->weightSize.x,
+                nThreads,
+                threadIndex);
+            DEBUG_VECTOR(context, "output", output);
+        }
     }
 }
 
@@ -1701,4 +1597,4 @@ NnCpuOpForward getCpuOpForward(NnOpCode code, NnOpQuantType quantType) {
         if (quantType == F32_F32_F32) return moeGateForward_F32_F32;
     }
     return nullptr;
-}
\ No newline at end of file
+}
diff --git a/src/nn/nn-cpu-ops.hpp b/src/nn/nn-cpu-ops.hpp
index c9aa865d..09ca8f1a 100644
--- a/src/nn/nn-cpu-ops.hpp
+++ b/src/nn/nn-cpu-ops.hpp
@@ -29,11 +29,6 @@ typedef struct {
 
     NnByte *weight;
     NnSize3D weightSize;
-    
-    // distributed execution parameters
-    NnUint nodeIndex;
-    NnUint nNodes;
-    
 } NnCpuOpContext;
 
 typedef void (*NnCpuOpForwardInit)(NnCpuOpContext *context);
diff --git a/src/nn/nn-cpu.cpp b/src/nn/nn-cpu.cpp
index 8e97be9a..ac58df65 100644
--- a/src/nn/nn-cpu.cpp
+++ b/src/nn/nn-cpu.cpp
@@ -123,10 +123,6 @@ NnDeviceSegment *NnCpuDevice::createSegment(NnUint segmentIndex) {
         opContext->bufferConfigs = nodeConfig->buffers;
         opContext->bufferFlags = bufferFlags;
 
-        // Inject nodeIndex and nNodes into the context
-        opContext->nodeIndex = this->nodeConfig->nodeIndex;
-        opContext->nNodes = this->netConfig->nNodes;
-
         opContext->input = new NnByte *[inputsPtr[opIndex].size()];
         opContext->inputSize = inputSizes[opIndex];
         opContext->hasInputContinuousMemory = hasPointerContinuousMemory(&opConfig->input);
diff --git a/src/nn/nn-network.cpp b/src/nn/nn-network.cpp
index c2089e1a..af5bc1e2 100644
--- a/src/nn/nn-network.cpp
+++ b/src/nn/nn-network.cpp
@@ -807,7 +807,15 @@ NnRootWeightLoader::~NnRootWeightLoader() {
 }
 
 void NnRootWeightLoader::finish() {
-    // empty
+    NnUint zeroSize = 0;
+    for (NnUint socketIndex = 0; socketIndex < nNodes - 1; socketIndex++) {
+        network->write(socketIndex, &zeroSize, sizeof(zeroSize));
+        network->readAck(socketIndex);
+    }
+    if (tempSize > 0) {
+        delete[] temp;
+        tempSize = 0;
+    }
 }
 
 void NnRootWeightLoader::allocate(NnSize size) {
@@ -831,41 +839,50 @@ void NnRootWeightLoader::writeWeight(NnUint nodeIndex, const char *opName, NnUin
 }
 
 NnSize NnRootWeightLoader::loadRoot(const char *opName, NnUint opIndex, NnSize nBytes, NnByte *weight) {
-    //set offset to 0 for loading the whole weight.
-    try {
-        executor->loadWeight(opName, opIndex, 0, nBytes, weight);
-    } catch (...) {
-   
-    }
+    executor->loadWeight(opName, opIndex, 0u, nBytes, weight);
     return nBytes;
 }
 
 NnSize NnRootWeightLoader::loadAll(const char *opName, NnUint opIndex, NnSize nBytes, NnByte *weight) {
-    //set offset to 0 for loading the whole weight.
-    try {
-        executor->loadWeight(opName, opIndex, 0, nBytes, weight);
-    } catch (...) {
+    executor->loadWeight(opName, opIndex, 0u, nBytes, weight);
 
+    if (nNodes > 1u) {
+        for (NnUint nodeIndex = 1u; nodeIndex < nNodes; nodeIndex++)
+            writeWeight(nodeIndex, opName, opIndex, 0u, nBytes, weight);
     }
     return nBytes;
 }
 
 NnSize NnRootWeightLoader::loadRowMatmulSlices(const char *opName, const NnUint opIndex, const NnUint expertIndex, NnRowMatmulSlice *slice, NnByte *weight) {
-    //set offset to 0 for loading the whole weight.
-    try {
-        executor->loadWeight(opName, opIndex, 0, slice->size.nBytes, weight);
-    } catch (...) {
-
+    const NnUint offset = expertIndex * slice->sliceSize.nBytes;
+    if (nNodes == 1u) {
+        executor->loadWeight(opName, opIndex, offset, slice->sliceSize.nBytes, weight);
+    } else {
+        allocate(slice->sliceSize.nBytes);
+        for (NnUint nodeIndex = 0; nodeIndex < nNodes; nodeIndex++) {
+            splitRowMatmulWeight(slice, nodeIndex, weight, temp);
+            if (nodeIndex == 0u)
+                executor->loadWeight(opName, opIndex, offset, slice->sliceSize.nBytes, temp);
+            else
+                writeWeight(nodeIndex, opName, opIndex, offset, slice->sliceSize.nBytes, temp);
+        }
     }
     return slice->size.nBytes;
 }
 
 NnSize NnRootWeightLoader::loadColMatmulSlices(const char *opName, const NnUint opIndex, const NnUint expertIndex, NnColMatmulSlice *slice, NnByte *weight) {
-    //set offset to 0 for loading the whole weight.
-    try {
-        executor->loadWeight(opName, opIndex, 0, slice->size.nBytes, weight);
-    } catch (...) {
-
+    const NnUint offset = expertIndex * slice->sliceSize.nBytes;
+    if (nNodes == 1) {
+        executor->loadWeight(opName, opIndex, offset, slice->sliceSize.nBytes, weight);
+    } else {
+        allocate(slice->sliceSize.nBytes);
+        for (NnUint nodeIndex = 0; nodeIndex < nNodes; nodeIndex++) {
+            splitColMatmulWeight(slice, nodeIndex, weight, temp);
+            if (nodeIndex == 0)
+                executor->loadWeight(opName, opIndex, offset, slice->sliceSize.nBytes, temp);
+            else
+                writeWeight(nodeIndex, opName, opIndex, offset, slice->sliceSize.nBytes, temp);
+        }
     }
     return slice->size.nBytes;
 }
@@ -918,3 +935,61 @@ void NnWorkerWeightReader::read() {
     }
     printf("💿 Weights loaded\n");
 }
+
+NnLocalWeightLoader::NnLocalWeightLoader(NnExecutor *executor, NnUint nodeIndex, NnUint nNodes) {
+    this->executor = executor;
+    this->nodeIndex = nodeIndex;
+    this->nNodes = nNodes;
+    this->tempSize = 0;
+}
+
+NnLocalWeightLoader::~NnLocalWeightLoader() {
+    if (tempSize > 0)
+        delete[] temp;
+}
+
+void NnLocalWeightLoader::finish() {}
+
+void NnLocalWeightLoader::allocate(NnSize size) {
+    if (tempSize < size) {
+        if (tempSize > 0)
+            delete[] temp;
+        tempSize = size;
+        temp = new NnByte[size];
+    }
+}
+
+NnSize NnLocalWeightLoader::loadRoot(const char *opName, NnUint opIndex, NnSize nBytes, NnByte *weight) {
+    if (nodeIndex == 0)
+        executor->loadWeight(opName, opIndex, 0u, nBytes, weight);
+    return nBytes;
+}
+
+NnSize NnLocalWeightLoader::loadAll(const char *opName, NnUint opIndex, NnSize nBytes, NnByte *weight) {
+    executor->loadWeight(opName, opIndex, 0u, nBytes, weight);
+    return nBytes;
+}
+
+NnSize NnLocalWeightLoader::loadRowMatmulSlices(const char *opName, const NnUint opIndex, const NnUint expertIndex, NnRowMatmulSlice *slice, NnByte *weight) {
+    const NnUint offset = expertIndex * slice->sliceSize.nBytes;
+    if (nNodes == 1u) {
+        executor->loadWeight(opName, opIndex, offset, slice->sliceSize.nBytes, weight);
+        return slice->size.nBytes;
+    }
+    allocate(slice->sliceSize.nBytes);
+    splitRowMatmulWeight(slice, nodeIndex, weight, temp);
+    executor->loadWeight(opName, opIndex, offset, slice->sliceSize.nBytes, temp);
+    return slice->size.nBytes;
+}
+
+NnSize NnLocalWeightLoader::loadColMatmulSlices(const char *opName, const NnUint opIndex, const NnUint expertIndex, NnColMatmulSlice *slice, NnByte *weight) {
+    const NnUint offset = expertIndex * slice->sliceSize.nBytes;
+    if (nNodes == 1u) {
+        executor->loadWeight(opName, opIndex, offset, slice->sliceSize.nBytes, weight);
+        return slice->size.nBytes;
+    }
+    allocate(slice->sliceSize.nBytes);
+    splitColMatmulWeight(slice, nodeIndex, weight, temp);
+    executor->loadWeight(opName, opIndex, offset, slice->sliceSize.nBytes, temp);
+    return slice->size.nBytes;
+}
diff --git a/src/nn/nn-network.hpp b/src/nn/nn-network.hpp
index fa2f88e5..7a173c5c 100644
--- a/src/nn/nn-network.hpp
+++ b/src/nn/nn-network.hpp
@@ -133,4 +133,22 @@ class NnWorkerWeightReader {
     void allocate(NnUint size);
 };
 
+class NnLocalWeightLoader {
+private:
+    NnExecutor *executor;
+    NnUint nodeIndex;
+    NnUint nNodes;
+    NnByte *temp;
+    NnSize tempSize;
+    void allocate(NnSize size);
+public:
+    NnLocalWeightLoader(NnExecutor *executor, NnUint nodeIndex, NnUint nNodes);
+    ~NnLocalWeightLoader();
+    NnSize loadRoot(const char *opName, NnUint opIndex, NnSize nBytes, NnByte *weight);
+    NnSize loadAll(const char *opName, NnUint opIndex, NnSize nBytes, NnByte *weight);
+    NnSize loadRowMatmulSlices(const char *opName, const NnUint opIndex, const NnUint expertIndex, NnRowMatmulSlice *slice, NnByte *weight);
+    NnSize loadColMatmulSlices(const char *opName, const NnUint opIndex, const NnUint expertIndex, NnColMatmulSlice *slice, NnByte *weight);
+    void finish();
+};
+
 #endif