From ca2ce13fa4aff68544a5fa311a8ba87d0c9f65f5 Mon Sep 17 00:00:00 2001
From: "The gemma.cpp Authors" <no-reply@google.com>
Date: Thu, 16 Jul 2026 08:53:18 -0700
Subject: [PATCH] Internal change.

PiperOrigin-RevId: 948997184
---
 gemma/weights.cc |  4 ++++
 io/io.cc         | 18 +++++++++++++++-
 ops/matmul.h     | 53 +++++++++++++++++++++++++++++++++++-------------
 3 files changed, 60 insertions(+), 15 deletions(-)

diff --git a/gemma/weights.cc b/gemma/weights.cc
index 3152cf89..87e48a4b 100644
--- a/gemma/weights.cc
+++ b/gemma/weights.cc
@@ -429,6 +429,10 @@ void LayerWeightsPtrs::Fixup(Model model,
           hwy::AllocateAligned<uint8_t>(total_bytes);
       hwy::CopyBytes(qkv_einsum_w2.RowBytes(0), tmp.get(), total_bytes);
 
+      const size_t row_bytes_old = row_bytes;
+      const size_t cols_bytes =
+          qkv_einsum_w2.Cols() * qkv_einsum_w2.ElementBytes();
+
       {
         std::lock_guard<std::mutex> lock(g_mat_owners_mutex);
         mat_owners.emplace_back();
diff --git a/io/io.cc b/io/io.cc
index 2f479b21..5816d111 100644
--- a/io/io.cc
+++ b/io/io.cc
@@ -28,6 +28,7 @@
 
 #include <stddef.h>
 
+#include <cstring>
 #include <memory>
 #include <string>
 
@@ -261,7 +262,22 @@ uint64_t IOBatch::Read(const File& file) const {
   }
 #endif  // GEMMA_IO_PREADV
 
-  // preadv disabled or no handle: use normal reads (higher kernel overhead).
+  // If no handle (remote FUSE), allocate a temporary buffer to do a single
+  // read.
+  const uint64_t total_bytes = TotalBytes();
+  std::unique_ptr<uint8_t[]> temp_buf(new (std::nothrow) uint8_t[total_bytes]);
+  if (temp_buf) {
+    if (!file.Read(offset_, total_bytes, temp_buf.get())) return 0;
+    uint8_t* src = temp_buf.get();
+    for (const IOSpan& span : spans_) {
+      std::memcpy(span.mem, src, span.bytes);
+      src += span.bytes;
+    }
+    return total_bytes;
+  }
+
+  // Fallback if allocation fails: use normal reads (higher kernel/network
+  // overhead).
   uint64_t total = 0;
   uint64_t offset = offset_;
   for (const IOSpan& span : spans_) {
diff --git a/ops/matmul.h b/ops/matmul.h
index 9e0d3f5d..9817b5f2 100644
--- a/ops/matmul.h
+++ b/ops/matmul.h
@@ -127,11 +127,17 @@ struct MMParallelWithinCluster {
     const hwy::pool::Caller caller =
         ctx.pool_callers.Get(Callers::kMMClusterForN);
 
-    ParallelPartitionWithinCluster(
-        range_n, n_multiple, inner_tasks, ctx, cluster_idx, caller,
-        [&](const IndexRange& worker_range, size_t worker) {
-          func(worker_range, worker);
-        });
+    size_t num_workers = ctx.pools.Cluster(cluster_idx).NumWorkers();
+    size_t worker_tasks = num_workers * inner_tasks;
+    if (hwy::DivCeil(range_n.Num(), worker_tasks) > kMaxNC) {
+      worker_tasks = hwy::DivCeil(range_n.Num(), kMaxNC);
+    }
+    const IndexRangePartition worker_ranges =
+        StaticPartition(range_n, worker_tasks, n_multiple);
+    ParallelForWithinCluster(worker_ranges.NumTasks(), ctx, cluster_idx, caller,
+                             [&](uint64_t w_task, size_t worker) {
+                               func(worker_ranges.Range(w_task), worker);
+                             });
   }
 
   template <class Func>
@@ -204,15 +210,34 @@ struct MMParallelHierarchical {
     (void)caller_cluster_idx;
     const hwy::pool::Caller caller = ctx.pool_callers.Get(Callers::kMMHierForN);
 
-    // Assign clusters (if any) a sub-range of `range_n` (typically hundreds).
-    ParallelPartitionAcrossClusters(
-        range_n, n_multiple, /*inner_tasks=*/1, ctx, caller,
-        [&](const IndexRange& cluster_range, size_t cluster_idx) {
-          ParallelPartitionWithinCluster(
-              cluster_range, n_multiple, inner_tasks, ctx, cluster_idx, caller,
-              [&](const IndexRange& worker_range, size_t worker) {
-                func(worker_range, worker);
-              });
+    // Calculate cluster task count to enforce kMaxNC cap per worker
+    size_t num_clusters = ctx.pools.NumClusters();
+    size_t cluster_tasks = num_clusters;
+    size_t workers_per_cluster = ctx.pools.Cluster(0).NumWorkers();
+    const size_t max_cluster_range_size =
+        workers_per_cluster * inner_tasks * kMaxNC;
+    if (hwy::DivCeil(range_n.Num(), cluster_tasks) > max_cluster_range_size) {
+      cluster_tasks = hwy::DivCeil(range_n.Num(), max_cluster_range_size);
+    }
+
+    const IndexRangePartition cluster_ranges =
+        StaticPartition(range_n, cluster_tasks, n_multiple);
+
+    ParallelForAcrossClusters(
+        cluster_ranges.NumTasks(), ctx, caller,
+        [&](uint64_t task, size_t cluster_idx) {
+          const IndexRange cluster_range = cluster_ranges.Range(task);
+          size_t num_workers = ctx.pools.Cluster(cluster_idx).NumWorkers();
+          size_t worker_tasks = num_workers * inner_tasks;
+          if (hwy::DivCeil(cluster_range.Num(), worker_tasks) > kMaxNC) {
+            worker_tasks = hwy::DivCeil(cluster_range.Num(), kMaxNC);
+          }
+          const IndexRangePartition worker_ranges =
+              StaticPartition(cluster_range, worker_tasks, n_multiple);
+          ParallelForWithinCluster(worker_ranges.NumTasks(), ctx, cluster_idx,
+                                   caller, [&](uint64_t w_task, size_t worker) {
+                                     func(worker_ranges.Range(w_task), worker);
+                                   });
         });
   }