diff --git a/.run/TestModels.MistralRun.run.xml b/.run/TestModels.MixtralRun.run.xml
similarity index 90%
rename from .run/TestModels.MistralRun.run.xml
rename to .run/TestModels.MixtralRun.run.xml
index 2f389f8..f0fab3e 100644
--- a/.run/TestModels.MistralRun.run.xml
+++ b/.run/TestModels.MixtralRun.run.xml
@@ -1,5 +1,5 @@
 <component name="ProjectRunConfigurationManager">
-  <configuration default="false" name="TestModels.MistralRun" type="JUnit" factoryName="JUnit" nameIsGenerated="true">
+  <configuration default="false" name="TestModels.MixtralRun" type="JUnit" factoryName="JUnit" nameIsGenerated="true">
     <classpathModifications>
       <entry exclude="true" path="$PROJECT_DIR$/jlama-native/target/classes" />
       <entry path="$PROJECT_DIR$/jlama-native/target/jlama-native-0.3.1-linux-x86_64.jar" />
@@ -15,7 +15,7 @@
     <option name="ALTERNATIVE_JRE_PATH" value="graalvm-22" />
     <option name="PACKAGE_NAME" value="com.github.tjake.jlama.model" />
     <option name="MAIN_CLASS_NAME" value="com.github.tjake.jlama.model.TestModels" />
-    <option name="METHOD_NAME" value="MistralRun" />
+    <option name="METHOD_NAME" value="MixtralRun" />
     <option name="TEST_OBJECT" value="method" />
     <option name="VM_PARAMETERS" value="-ea --add-modules=jdk.incubator.vector -Djava.library.path=../jlama-native/target/native-lib-only" />
     <method v="2">
diff --git a/Dockerfile b/Dockerfile
index 1ed370a..5be82f6 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -1,7 +1,40 @@
-FROM openjdk:21-slim as builder
+FROM ubuntu:22.04 as builder
 RUN apt-get update
-RUN apt-get install -y build-essential
+RUN apt-get install -y build-essential git zip curl zlib1g-dev
 
+ENV SDKMAN_DIR /root/.sdkman
+ENV JAVA_VERSION_20 20.0.2-graalce
+ENV JAVA_VERSION_21 21.0.2-graalce
+ENV JAVA_VERSION_22 22.0.2-graalce
+
+RUN ["mkdir", "-p", "/build"]
+
+RUN rm /bin/sh && ln -s /bin/bash /bin/sh
+RUN curl -s "https://get.sdkman.io" | bash
+RUN chmod a+x "$SDKMAN_DIR/bin/sdkman-init.sh"
+
+RUN set -x \
+    && echo "sdkman_auto_answer=true" > $SDKMAN_DIR/etc/config \
+    && echo "sdkman_auto_selfupdate=false" >> $SDKMAN_DIR/etc/config \
+    && echo "sdkman_insecure_ssl=false" >> $SDKMAN_DIR/etc/config
+
+WORKDIR $SDKMAN_DIR
+RUN [[ -s "$SDKMAN_DIR/bin/sdkman-init.sh" ]] && source "$SDKMAN_DIR/bin/sdkman-init.sh" && exec "$@"
+
+RUN source /root/.bashrc
+RUN source "$SDKMAN_DIR/bin/sdkman-init.sh" && sdk install java $JAVA_VERSION_20
+RUN source "$SDKMAN_DIR/bin/sdkman-init.sh" && sdk install java $JAVA_VERSION_21
+RUN source "$SDKMAN_DIR/bin/sdkman-init.sh" && sdk install java $JAVA_VERSION_22
+
+WORKDIR /build
+RUN git clone https://github.com/tjake/sdkman-for-toolchains.git
+WORKDIR /build/sdkman-for-toolchains
+RUN source "$SDKMAN_DIR/bin/sdkman-init.sh" && ./mvnw -Pnative clean package
+RUN ["mkdir", "-p", "/root/.m2"]
+RUN printf "<?xml version=\"1.0\" encoding=\"UTF-8\"?>\n<toolchains>\n</toolchains>" > /root/.m2/toolchains.xml
+RUN source "$SDKMAN_DIR/bin/sdkman-init.sh" && target/toolchains generate
+RUN cat  /root/.m2/toolchains.xml
+WORKDIR /build
 COPY jlama-cli jlama-cli
 COPY jlama-core jlama-core
 COPY jlama-native jlama-native
@@ -11,20 +44,19 @@ COPY .mvn .mvn
 COPY pom.xml .
 COPY mvnw .
 
-RUN --mount=type=cache,target=/root/.m2 ./mvnw clean package
+RUN --mount=type=cache,target=/root/.m2/repository source "$SDKMAN_DIR/bin/sdkman-init.sh" && sdk use java $JAVA_VERSION_22 && ./mvnw clean package -DskipTests
 
 FROM openjdk:21-slim
 RUN apt-get update
-RUN apt-get install -y procps curl
+RUN apt-get install -y procps curl gzip
 
 LABEL org.opencontainers.image.source=https://github.com/tjake/Jlama
 
 COPY inlinerules.json inlinerules.json
 COPY run-cli.sh run-cli.sh
 COPY conf/logback.xml logback.xml
-COPY --from=builder jlama-cli/target/jlama-cli.jar ./jlama-cli.jar
-
-RUN curl -s -L https://github.com/async-profiler/async-profiler/releases/download/v3.0/async-profiler-3.0-linux-x64.tar.gz | tar xvz - -C /profiler
+COPY --from=builder /build/jlama-cli/target/jlama-cli.jar ./jlama-cli.jar
+RUN mkdir -p /profiler && curl -s -L https://github.com/async-profiler/async-profiler/releases/download/v3.0/async-profiler-3.0-linux-x64.tar.gz | tar zxvf - -C /profiler
 
 ENV JLAMA_PREINSTALLED_JAR=/jlama-cli.jar
 ENV JLAMA_JVM_ARGS="-Dlogback.configurationFile=./logback.xml"
diff --git a/docker-compose.yaml b/docker-compose.yaml
index e1134e1..638bab6 100644
--- a/docker-compose.yaml
+++ b/docker-compose.yaml
@@ -12,7 +12,7 @@ services:
       resources:
         limits:
           memory: 1g
-          cpus: 1
+          cpus: "1"
     command:
       - cluster-coordinator
       - --threads=2
@@ -36,7 +36,7 @@ services:
       replicas: 8
       resources:
         limits:
-          cpus: 1
+          cpus: "1"
           memory: 500m
     command:
       - cluster-worker
diff --git a/jlama-cli/src/main/java/com/github/tjake/jlama/cli/commands/ApiServiceCommand.java b/jlama-cli/src/main/java/com/github/tjake/jlama/cli/commands/ApiServiceCommand.java
index 25a8e0f..71a3205 100644
--- a/jlama-cli/src/main/java/com/github/tjake/jlama/cli/commands/ApiServiceCommand.java
+++ b/jlama-cli/src/main/java/com/github/tjake/jlama/cli/commands/ApiServiceCommand.java
@@ -19,6 +19,8 @@
 
 import com.github.tjake.jlama.model.AbstractModel;
 import java.util.Optional;
+
+import com.github.tjake.jlama.model.functions.Generator;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 import org.springframework.boot.SpringBootConfiguration;
@@ -40,10 +42,10 @@ public class ApiServiceCommand extends BaseCommand implements WebMvcConfigurer {
     @CommandLine.Option(names = { "-p", "--port" }, description = "http port (default: ${DEFAULT-VALUE})", defaultValue = "8080")
     int port = 8080;
 
-    static volatile AbstractModel m;
+    protected static volatile Generator m;
 
     @Bean
-    public AbstractModel getModelBean() {
+    public Generator getModelBean() {
         return m;
     }
 
diff --git a/jlama-cli/src/main/java/com/github/tjake/jlama/cli/commands/ClusterCoordinatorCommand.java b/jlama-cli/src/main/java/com/github/tjake/jlama/cli/commands/ClusterCoordinatorCommand.java
index e7991e2..373f27f 100644
--- a/jlama-cli/src/main/java/com/github/tjake/jlama/cli/commands/ClusterCoordinatorCommand.java
+++ b/jlama-cli/src/main/java/com/github/tjake/jlama/cli/commands/ClusterCoordinatorCommand.java
@@ -15,11 +15,23 @@
  */
 package com.github.tjake.jlama.cli.commands;
 
+import com.github.tjake.jlama.model.AbstractModel;
+import com.github.tjake.jlama.model.functions.Generator;
 import com.github.tjake.jlama.net.Coordinator;
+import org.springframework.boot.SpringBootConfiguration;
+import org.springframework.boot.autoconfigure.SpringBootApplication;
+import org.springframework.boot.builder.SpringApplicationBuilder;
+import org.springframework.context.annotation.Bean;
+import org.springframework.context.annotation.Configuration;
+import org.springframework.web.servlet.config.annotation.ResourceHandlerRegistry;
+import org.springframework.web.servlet.config.annotation.WebMvcConfigurer;
 import picocli.CommandLine;
 
 @CommandLine.Command(name = "cluster-coordinator", description = "Starts a distributed rest api for a model using cluster workers")
-public class ClusterCoordinatorCommand extends BaseCommand {
+@SpringBootApplication(scanBasePackages = { "com.github.tjake.jlama.net.openai", "com.github.tjake.jlama.cli.commands" })
+@SpringBootConfiguration
+@Configuration
+public class ClusterCoordinatorCommand extends BaseCommand implements WebMvcConfigurer {
 
     @CommandLine.Option(names = { "-w", "--worker-count" }, description = "signifies this instance is a coordinator", required = true)
     int workerCount = 1;
@@ -32,11 +44,19 @@ public class ClusterCoordinatorCommand extends BaseCommand {
         "--port" }, description = "http port to listen on (default: ${DEFAULT-VALUE})", defaultValue = "8080")
     int port = 8080;
 
+    @Override
+    public void addResourceHandlers(ResourceHandlerRegistry registry) {
+        registry.addResourceHandler("/ui/**").addResourceLocations("classpath:/static/ui/");
+    }
+
     @Override
     public void run() {
         try {
             Coordinator c = new Coordinator(model, workingDirectory, grpcPort, workerCount);
 
+            //This wires up the bean for the rest api
+            ApiServiceCommand.m = c;
+
             new Thread(() -> {
                 try {
                     c.start();
@@ -45,16 +65,13 @@ public void run() {
                 }
             }).start();
 
-            /*UndertowJaxrsServer ut = new UndertowJaxrsServer();
-            ut.deploy(new JlamaRestApi(c), APPLICATION_PATH);
-            ut.addResourcePrefixPath(
-                    "/ui",
-                    resource(new ClassPathResourceManager(ServeCommand.class.getClassLoader()))
-                            .setDirectoryListingEnabled(true)
-                            .addWelcomeFiles("index.html"));
-            
-            System.out.println("Chat UI: http://localhost:" + port + "/ui/index.html");
-            ut.start(Undertow.builder().addHttpListener(port, "0.0.0.0"));*/
+            System.out.println("Chat UI: http://localhost:" + port);
+            System.out.println("OpenAI Chat API: http://localhost:" + port + "/chat/completions");
+
+            new SpringApplicationBuilder(ClusterCoordinatorCommand.class).lazyInitialization(true)
+                    .properties("server.port", "" + port, "logging.level.org.springframework.web", "info")
+                    .build()
+                    .run();
 
         } catch (Exception e) {
             e.printStackTrace();
diff --git a/jlama-core/src/main/java/com/github/tjake/jlama/model/AbstractModel.java b/jlama-core/src/main/java/com/github/tjake/jlama/model/AbstractModel.java
index 699e1c6..38e6c1f 100644
--- a/jlama-core/src/main/java/com/github/tjake/jlama/model/AbstractModel.java
+++ b/jlama-core/src/main/java/com/github/tjake/jlama/model/AbstractModel.java
@@ -126,7 +126,7 @@ protected AbstractModel(
 
         if (workingMemoryQType != workingMemoryDType) {
             boolean supportsQType;
-            AbstractTensor tmp = makeTensor(Q8ByteBufferTensor.BLOCK_SIZE);
+            AbstractTensor tmp = makeDenseTensor(Q8ByteBufferTensor.BLOCK_SIZE);
             try (AbstractTensor tmp2 = TensorOperationsProvider.get().quantize(tmp, workingMemoryQType, 0, Q8ByteBufferTensor.BLOCK_SIZE)) {
                 supportsQType = tmp2.dType() == workingMemoryQType;
                 if (!supportsQType) {
@@ -176,14 +176,11 @@ public Optional<PromptSupport> promptSupport() {
     }
 
     public AbstractTensor makeTensor(int... shape) {
-        TensorShape s;
-        if (c.offset().isPresent() && shape[shape.length - 1] == c.embeddingLength) s = TensorShape.sparse(shape, c.offset().get());
-        else s = TensorShape.of(shape);
-
+        TensorShape s = TensorShape.of(shape);
         return c.tensorCache.get(workingDType, s);
     }
 
-    public AbstractTensor makeFullTensor(int... shape) {
+    public AbstractTensor makeDenseTensor(int... shape) {
         return c.tensorCache.get(workingDType, TensorShape.of(shape));
     }
 
@@ -219,7 +216,7 @@ public AbstractTensor forward(
         debug("EMBEDDING TOKEN", token_id);
         debug("TOKEN POSITION", pos);
 
-        for (int i = c.layerStart(); i < c.layerEnd(); i++) {
+        for (int i = c.dctx().layerStart; i < c.dctx().layerEnd; i++) {
             AbstractTensor kvlayer = kvbuf.slice(true, i);
             AbstractTensor ref = embedding; // reference so we can free
             embedding = transformerBlocks[i].forward(embedding, pos, kvlayer, normReducer, tensorReducer);
@@ -242,7 +239,7 @@ protected AbstractTensor batchForwardSlow(int[] token_ids, int startPos, Abstrac
     protected AbstractTensor batchForward(int[] token_ids, int startPos, AbstractTensor kvbuf) {
 
         AbstractTensor embedding = embedInput.batchInputsToEmbeddings(token_ids, startPos);
-        for (int i = c.layerStart(); i < c.layerEnd(); i++) {
+        for (int i = c.dctx().layerStart; i < c.dctx().layerEnd; i++) {
             AbstractTensor kvlayer = kvbuf.slice(true, i);
             AbstractTensor ref = embedding; // reference so we can free
             embedding = transformerBlocks[i].forward(embedding, startPos, kvlayer, Optional.empty(), Optional.empty());
@@ -323,7 +320,7 @@ public Response generate(
         StringBuilder responseText = new StringBuilder();
         StringBuilder responseTextWithSpecialTokens = new StringBuilder();
 
-        try (AbstractTensor logits = makeTensor(c.vocabularySize)) {
+        try (AbstractTensor logits = makeDenseTensor(c.vocabularySize)) {
             int[] promptTokens = new int[(1 + encoded.length)];
 
             promptTokens[0] = c.bosToken;
diff --git a/jlama-core/src/main/java/com/github/tjake/jlama/model/CausalSelfAttention.java b/jlama-core/src/main/java/com/github/tjake/jlama/model/CausalSelfAttention.java
index b14942f..944b715 100644
--- a/jlama-core/src/main/java/com/github/tjake/jlama/model/CausalSelfAttention.java
+++ b/jlama-core/src/main/java/com/github/tjake/jlama/model/CausalSelfAttention.java
@@ -28,6 +28,7 @@
 public class CausalSelfAttention {
     private final AbstractModel m;
     private final Config c;
+    private final DistributedContext dctx;
     private final Optional<AbstractTensor> queryAttnBias;
     private final Optional<AbstractTensor> keyAttnBias;
 
@@ -105,6 +106,7 @@ public CausalSelfAttention(
     ) {
         this.m = m;
         this.c = m.c;
+        this.dctx = m.c.dctx();
         this.queryAttnBias = queryAttnBias;
         this.keyAttnBias = keyAttnBias;
         this.valueAttnBias = valueAttnBias;
@@ -133,33 +135,33 @@ public AbstractTensor forward(
         int batchSize = input.shape().first();
 
         try (
-            AbstractTensor queryBatch = m.makeFullTensor(batchSize, attentionLength);
-            AbstractTensor tmpKeyBatch = m.makeFullTensor(batchSize, c.kvLength);
-            AbstractTensor tmpValBatch = m.makeFullTensor(batchSize, c.kvLength);
-            AbstractTensor valueBatch = m.makeFullTensor(batchSize, attentionLength)
+                AbstractTensor queryBatch = m.makeDenseTensor(batchSize, attentionLength);
+                AbstractTensor tmpKeyBatch = m.makeDenseTensor(batchSize, c.kvLength);
+                AbstractTensor tmpValBatch = m.makeDenseTensor(batchSize, c.kvLength);
+                AbstractTensor valueBatch = m.makeDenseTensor(batchSize, attentionLength)
         ) {
 
             if (c.isGQA) {
-                VectorMath.pchunk(0, attentionLength, (chunkStart, chunkLength) -> {
+                VectorMath.pchunk(dctx.attentionSegmentStart, dctx.attentionSegmentLength, (chunkStart, chunkLength) -> {
                     TensorOperationsProvider.get()
                         .dotProductChunk(
                             queryBatch,
                             input,
                             queryAttnWeights,
-                            c.embeddingSegmentStart(),
-                            c.embeddingSegmentLength(),
+                            0,
+                            c.embeddingLength,
                             chunkStart,
                             chunkLength
                         );
                 });
-                VectorMath.pchunk(0, c.kvLength, (chunkStart, chunkLength) -> {
+                VectorMath.pchunk(dctx.kvSegmentStart, dctx.kvSegmentLength, (chunkStart, chunkLength) -> {
                     TensorOperationsProvider.get()
                         .dotProductChunk(
                             tmpKeyBatch,
                             input,
                             keyAttnWeights,
-                            c.embeddingSegmentStart(),
-                            c.embeddingSegmentLength(),
+                            0,
+                            c.embeddingLength,
                             chunkStart,
                             chunkLength
                         );
@@ -168,8 +170,8 @@ public AbstractTensor forward(
                             tmpValBatch,
                             input,
                             valueAttnWeights,
-                            c.embeddingSegmentStart(),
-                            c.embeddingSegmentLength(),
+                            0,
+                            c.embeddingLength,
                             chunkStart,
                             chunkLength
                         );
@@ -180,31 +182,28 @@ public AbstractTensor forward(
                 qkvResults[2] = tmpValBatch;
 
                 // compute the query vector
-                VectorMath.pchunk(0, c.embeddingLength, (chunkStart, chunkLength) -> {
+                VectorMath.pchunk(dctx.attentionSegmentStart, dctx.attentionSegmentLength, (chunkStart, chunkLength) -> {
                     TensorOperationsProvider.get()
                         .dotProductBatchChunk(
                             qkvResults,
                             input,
                             qkvWeights,
-                            c.embeddingSegmentStart(),
-                            c.embeddingSegmentLength(),
+                            0,
+                            c.embeddingLength,
                             chunkStart,
                             chunkLength
                         );
                 });
             }
 
-            // For distributed sum of tensor
-            tensorReducer.ifPresent(func -> func.accept(List.of(queryBatch, tmpKeyBatch, tmpValBatch)));
-
             queryAttnBias.ifPresent(
-                bias -> TensorOperationsProvider.get().accumulate(queryBatch, bias, c.embeddingSegmentStart(), c.embeddingSegmentLength())
+                bias -> TensorOperationsProvider.get().accumulate(queryBatch, bias, dctx.attentionSegmentStart, dctx.attentionSegmentLength)
             );
             keyAttnBias.ifPresent(
-                bias -> TensorOperationsProvider.get().accumulate(tmpKeyBatch, bias, c.kvSegmentStart(), c.kvSegmentLength())
+                bias -> TensorOperationsProvider.get().accumulate(tmpKeyBatch, bias, dctx.kvSegmentStart, dctx.kvSegmentLength)
             );
             valueAttnBias.ifPresent(
-                bias -> TensorOperationsProvider.get().accumulate(tmpValBatch, bias, c.kvSegmentStart(), c.kvSegmentLength())
+                bias -> TensorOperationsProvider.get().accumulate(tmpValBatch, bias, dctx.kvSegmentStart, dctx.kvSegmentLength)
             );
 
             debug("query", queryBatch, 0);
@@ -233,20 +232,20 @@ public AbstractTensor forward(
                     ) {
                         key.copyFrom(
                             tmpKey2,
-                            tmpKey2.getOffset(0, c.kvSegmentStart()),
-                            key.getOffset(0, c.kvSegmentStart()),
-                            c.kvSegmentLength()
+                            tmpKey2.getOffset(0, dctx.kvSegmentStart),
+                            key.getOffset(0, dctx.kvSegmentStart),
+                            dctx.kvSegmentLength
                         );
                         val.copyFrom(
                             tmpVal2,
-                            tmpVal2.getOffset(0, c.kvSegmentStart()),
-                            val.getOffset(0, c.kvSegmentStart()),
-                            c.kvSegmentLength()
+                            tmpVal2.getOffset(0, dctx.kvSegmentStart),
+                            val.getOffset(0, dctx.kvSegmentStart),
+                            dctx.kvSegmentLength
                         );
                     }
                 } else {
-                    key.copyFrom(tmpKey, tmpKey.getOffset(0, c.kvSegmentStart()), key.getOffset(0, c.kvSegmentStart()), c.kvLength);
-                    val.copyFrom(tmpVal, tmpVal.getOffset(0, c.kvSegmentStart()), val.getOffset(0, c.kvSegmentStart()), c.kvLength);
+                    key.copyFrom(tmpKey, tmpKey.getOffset(0, dctx.kvSegmentStart), key.getOffset(0, dctx.kvSegmentStart), dctx.kvSegmentLength);
+                    val.copyFrom(tmpVal, tmpVal.getOffset(0, dctx.kvSegmentStart), val.getOffset(0, dctx.kvSegmentStart), dctx.kvSegmentLength);
                 }
 
                 // apply RoPE if present (accounting for huggingface permutation)
@@ -257,7 +256,7 @@ public AbstractTensor forward(
 
                     if (c.isGQA) {
                         // apply RoPE rotation to the q and k vectors for each head
-                        for (int h = c.headStart(); h < c.headEnd(); h++) {
+                        for (int h = dctx.headStart; h < dctx.headEnd; h++) {
                             // get the q vectors for this head
                             int offset = h * c.headSize;
 
@@ -277,7 +276,7 @@ public AbstractTensor forward(
                             }
                         }
 
-                        for (int h = c.groupHeadStart(); h < c.groupHeadEnd(); h++) {
+                        for (int h = dctx.groupHeadStart; h < dctx.groupHeadEnd; h++) {
                             // get the k vectors for this head
                             int offset = h * c.headSize;
                             if (offset >= key.shape().last()) break;
@@ -294,7 +293,7 @@ public AbstractTensor forward(
                         }
                     } else {
                         // apply RoPE rotation to the q and k vectors for each head
-                        for (int h = c.headStart(); h < c.headEnd(); h++) {
+                        for (int h = dctx.headStart; h < dctx.headEnd; h++) {
                             // get the q and k vectors for this head
                             int offset = h * c.headSize;
                             // rotate q and k by the freq theta and freq r
@@ -318,8 +317,8 @@ public AbstractTensor forward(
                 });
 
                 // Attention
-                VectorMath.pfor(c.headStart(), c.headEnd(), h -> {
-                    try (AbstractTensor attn = m.makeFullTensor(1, kvp.shape().first())) {
+                VectorMath.pfor(dctx.headStart, dctx.headEnd, h -> {
+                    try (AbstractTensor attn = m.makeDenseTensor(1, kvp.shape().first())) {
                         int xoffset = c.maybeMapToGroupHead(h) * c.headSize;
                         int yoffset = h * c.headSize;
 
@@ -342,7 +341,7 @@ public AbstractTensor forward(
 
             // matmul the projection and sum into input
             // input += c_proj_weight @ ybuf + c_proj_bias
-            AbstractTensor result = m.makeFullTensor(batchSize, c.embeddingLength);
+            AbstractTensor result = m.makeDenseTensor(batchSize, c.embeddingLength);
             try (AbstractTensor vq = m.maybeQuantize(valueBatch)) {
                 VectorMath.pchunk(0, c.embeddingLength, (chunkStart, chunkSize) -> {
                     TensorOperationsProvider.get()
@@ -350,8 +349,8 @@ public AbstractTensor forward(
                             result,
                             vq,
                             outputProjectionWeights,
-                            c.embeddingSegmentStart(),
-                            attentionQVSizeMismatch ? attentionLength : c.embeddingSegmentLength(),
+                            dctx.attentionSegmentStart,
+                            dctx.attentionSegmentLength,
                             chunkStart,
                             chunkSize
                         );
@@ -360,7 +359,7 @@ public AbstractTensor forward(
                 tensorReducer.ifPresent(func -> func.accept(Collections.singletonList(result)));
 
                 outputProjectionBias.ifPresent(
-                    bias -> TensorOperationsProvider.get().accumulate(result, bias, c.embeddingSegmentStart(), c.embeddingSegmentLength())
+                    bias -> TensorOperationsProvider.get().accumulate(result, bias, 0, c.embeddingLength)
                 );
             }
 
diff --git a/jlama-core/src/main/java/com/github/tjake/jlama/model/DistributedContext.java b/jlama-core/src/main/java/com/github/tjake/jlama/model/DistributedContext.java
new file mode 100644
index 0000000..173f3df
--- /dev/null
+++ b/jlama-core/src/main/java/com/github/tjake/jlama/model/DistributedContext.java
@@ -0,0 +1,143 @@
+package com.github.tjake.jlama.model;
+
+import com.github.tjake.jlama.safetensors.Config;
+import com.github.tjake.jlama.tensor.AbstractTensor;
+import com.github.tjake.jlama.util.Pair;
+
+import java.util.List;
+import java.util.Optional;
+import java.util.function.Consumer;
+
+public class DistributedContext {
+
+    private final Config c;
+    private final int modelShard;
+    private final int numModelShards;
+    private final int layerShard;
+    private final int numLayerShards;
+
+    private final Consumer<List<AbstractTensor>> tensorSync;
+
+    // Suppliers to store values that chance when offset is adjusted
+    public final int embeddingSegmentStart;
+    public final int embeddingSegmentLength;
+    public final int embeddingSegmentEnd;
+
+    public final int attentionSegmentStart;
+    public final int attentionSegmentLength;
+    public final int attentionSegmentEnd;
+
+    public final int hiddenSegmentStart;
+    public final int hiddenSegmentLength;
+    public final int hiddenSegmentEnd;
+
+    public final int kvSegmentStart;
+    public final int kvSegmentLength;
+    public final int kvSegmentEnd;
+
+    public final int headStart;
+    public final int headEnd;
+    public final int groupHeadStart;
+    public final int groupHeadEnd;
+
+    public final int numberOfLayers;
+    public final int layerStart;
+    public final int layerEnd;
+
+    private DistributedContext(Config c, int modelShard, int numModelShards, int layerShard, int numLayerShards, Consumer<List<AbstractTensor>> tensorSync) {
+        this.c = c;
+        this.modelShard = modelShard;
+        this.numModelShards = numModelShards;
+        this.layerShard = layerShard;
+        this.numLayerShards = numLayerShards;
+        this.tensorSync = tensorSync;
+
+        this.numberOfLayers = c.numberOfLayers / numLayerShards;
+        this.layerStart = numberOfLayers * layerShard;
+        this.layerEnd = layerStart + numberOfLayers;
+
+        this.embeddingSegmentLength = c.embeddingLength / numModelShards;
+        this.embeddingSegmentStart = embeddingSegmentLength * modelShard;
+        this.embeddingSegmentEnd = embeddingSegmentStart + embeddingSegmentLength;
+
+        this.attentionSegmentLength = c.attentionLength / numModelShards;
+        this.attentionSegmentStart = attentionSegmentLength * modelShard;
+        this.attentionSegmentEnd = attentionSegmentStart + attentionSegmentLength;
+
+        this.hiddenSegmentLength = c.hiddenLength / numModelShards;
+        this.hiddenSegmentStart = hiddenSegmentLength * modelShard;
+        this.hiddenSegmentEnd = hiddenSegmentStart + hiddenSegmentLength;
+
+        this.kvSegmentStart = embeddingSegmentStart / c.headGroupSize;
+        this.kvSegmentEnd = embeddingSegmentEnd / c.headGroupSize;
+        this.kvSegmentLength = embeddingSegmentLength / c.headGroupSize;
+
+
+        this.headStart = embeddingSegmentStart / c.headSize;
+        this.headEnd = embeddingSegmentEnd / c.headSize;
+        this.groupHeadStart = kvSegmentStart / c.headSize;
+        this.groupHeadEnd = kvSegmentEnd / c.headSize;
+    }
+
+    public boolean hasModelShard() {
+        return numModelShards > 1;
+    }
+
+    public void syncTensors(List<AbstractTensor> tensors) {
+        tensorSync.accept(tensors);
+    }
+
+    public int getShardOffsetForLength(int length) {
+        return length / numModelShards * modelShard;
+    }
+
+    public int getShardLength(int length) {
+        return length / numModelShards;
+    }
+
+    public static Builder builder(Config c) {
+        return new Builder(c);
+    }
+
+    public static class Builder {
+        private Config c;
+        private int modelShard = 0;
+        private int numModelShards = 1;
+        private int layerShard = 0;
+        private int numLayerShards = 1;
+        private Consumer<List<AbstractTensor>> tensorSync;
+
+        public Builder(Config c) {
+            this.c = c;
+        }
+
+        public Builder setModelShard(int modelShard) {
+            this.modelShard = modelShard;
+            return this;
+        }
+
+        public Builder setNumModelShards(int numModelShards) {
+            this.numModelShards = numModelShards;
+            return this;
+        }
+
+        public Builder setLayerShard(int layerShard) {
+            this.layerShard = layerShard;
+            return this;
+        }
+
+        public Builder setNumLayerShards(int numLayerShards) {
+            this.numLayerShards = numLayerShards;
+            return this;
+        }
+
+        public Builder setTensorSync(Consumer<List<AbstractTensor>> tensorSync) {
+            this.tensorSync = tensorSync;
+            return this;
+        }
+
+        public DistributedContext build() {
+            return new DistributedContext(c, modelShard, numModelShards, layerShard, numLayerShards, tensorSync);
+        }
+    }
+}
diff --git a/jlama-core/src/main/java/com/github/tjake/jlama/model/LayerNorm.java b/jlama-core/src/main/java/com/github/tjake/jlama/model/LayerNorm.java
index f89d66a..e5cee2e 100644
--- a/jlama-core/src/main/java/com/github/tjake/jlama/model/LayerNorm.java
+++ b/jlama-core/src/main/java/com/github/tjake/jlama/model/LayerNorm.java
@@ -41,7 +41,7 @@ public AbstractTensor forward(AbstractTensor input, Optional<BiFunction<Float, F
         Preconditions.checkArgument(input.shape().dims() == 2);
         int size = input.shape().last();
         Preconditions.checkArgument(size == m.c.embeddingLength);
-        return forward(input, m.c.embeddingSegmentStart(), m.c.embeddingSegmentLength(), reducer);
+        return forward(input, 0, m.c.embeddingLength, reducer);
     }
 
     public AbstractTensor forward(
@@ -64,12 +64,6 @@ public AbstractTensor forward(
                 sumSq += v * v;
             }
 
-            if (reducer.isPresent()) {
-                Pair<Float, Float> p = reducer.get().apply(sumSq, sum);
-                sumSq = p.left;
-                sum = p.right;
-            }
-
             float mean = sum / m.c.embeddingLength;
             float variance = sumSq / m.c.embeddingLength - mean * mean;
             float invStddev = 1.0f / (float) Math.sqrt(variance + m.c.layerNormEps);
diff --git a/jlama-core/src/main/java/com/github/tjake/jlama/model/MLPBlock.java b/jlama-core/src/main/java/com/github/tjake/jlama/model/MLPBlock.java
index ecf7745..579c556 100644
--- a/jlama-core/src/main/java/com/github/tjake/jlama/model/MLPBlock.java
+++ b/jlama-core/src/main/java/com/github/tjake/jlama/model/MLPBlock.java
@@ -21,6 +21,7 @@
 import com.github.tjake.jlama.tensor.AbstractTensor;
 import com.github.tjake.jlama.tensor.operations.TensorOperationsProvider;
 import java.util.ArrayList;
+import java.util.Collections;
 import java.util.List;
 import java.util.Optional;
 import java.util.function.Consumer;
@@ -30,6 +31,7 @@
  */
 public class MLPBlock implements FeedForward {
     private final AbstractModel model;
+    private final DistributedContext dctx;
     private final Optional<AbstractTensor> fullyConnectedBias;
     private final AbstractTensor fullyConnectedWeights;
 
@@ -82,6 +84,7 @@ public MLPBlock(
         AbstractTensor upProjectionWeights
     ) {
         this.model = model;
+        this.dctx = model.c.dctx();
         this.activationFunction = activationFunction;
         this.fullyConnectedBias = fullyConnectedBias;
         this.fullyConnectedWeights = fullyConnectedWeights;
@@ -105,15 +108,15 @@ public AbstractTensor forward(AbstractTensor lnemb, Optional<Consumer<List<Abstr
             batchResults[0] = buf;
             batchResults[1] = buf2;
 
-            VectorMath.pchunk(0, hiddenLength, (chunkStart, chunkSize) -> {
+            VectorMath.pchunk(dctx.hiddenSegmentStart, dctx.hiddenSegmentLength, (chunkStart, chunkSize) -> {
                 if (upProjectionWeights != null) {
                     TensorOperationsProvider.get()
                         .dotProductBatchChunk(
                             batchResults,
                             lnemb,
                             batchWeights,
-                            model.c.embeddingSegmentStart(),
-                            model.c.embeddingSegmentLength(),
+                            0,
+                            model.c.embeddingLength,
                             chunkStart,
                             chunkSize
                         );
@@ -123,25 +126,17 @@ public AbstractTensor forward(AbstractTensor lnemb, Optional<Consumer<List<Abstr
                             buf,
                             lnemb,
                             fullyConnectedWeights,
-                            model.c.embeddingSegmentStart(),
-                            model.c.embeddingSegmentLength(),
+                            0,
+                            model.c.embeddingLength,
                             chunkStart,
                             chunkSize
                         );
                 }
             });
 
-            tensorReducer.ifPresent(func -> {
-                List<AbstractTensor> ts = new ArrayList<>(2);
-                ts.add(buf);
-                if (upProjectionWeights != null) ts.add(buf2);
+            fullyConnectedBias.ifPresent(bias -> TensorOperationsProvider.get().accumulate(buf, bias, dctx.hiddenSegmentStart, dctx.hiddenSegmentLength));
 
-                func.accept(ts);
-            });
-
-            fullyConnectedBias.ifPresent(bias -> TensorOperationsProvider.get().accumulate(buf, bias, 0, hiddenLength));
-
-            VectorMath.pfor(0, hiddenLength, i -> {
+            VectorMath.pfor(dctx.hiddenSegmentStart, dctx.hiddenSegmentEnd, i -> {
                 for (int j = 0; j < batchSize; j++) {
                     float w1 = buf.get(j, i);
                     float w1a = ActivationFunction.eval(activationFunction, w1);
@@ -156,13 +151,14 @@ public AbstractTensor forward(AbstractTensor lnemb, Optional<Consumer<List<Abstr
             try (AbstractTensor bufq = model.maybeQuantize(buf)) {
                 // matmul the projection and sum into input
                 AbstractTensor result = model.makeTensor(batchSize, model.c.embeddingLength);
-                VectorMath.pchunk(model.c.embeddingSegmentStart(), model.c.embeddingSegmentLength(), (chunkStart, chunkSize) -> {
-                    TensorOperationsProvider.get().dotProductChunk(result, bufq, projectionWeights, 0, hiddenLength, chunkStart, chunkSize);
+                VectorMath.pchunk(0, model.c.embeddingLength, (chunkStart, chunkSize) -> {
+                    TensorOperationsProvider.get().dotProductChunk(result, bufq, projectionWeights, dctx.hiddenSegmentStart, dctx.hiddenSegmentLength, chunkStart, chunkSize);
                 });
 
+                tensorReducer.ifPresent(func -> func.accept(Collections.singletonList(result)));
+
                 projectionBias.ifPresent(
-                    bias -> TensorOperationsProvider.get()
-                        .accumulate(result, bias, model.c.embeddingSegmentStart(), model.c.embeddingSegmentLength())
+                    bias -> TensorOperationsProvider.get().accumulate(result, bias, 0, model.c.embeddingLength)
                 );
                 return result;
             }
diff --git a/jlama-core/src/main/java/com/github/tjake/jlama/model/MoEBlock.java b/jlama-core/src/main/java/com/github/tjake/jlama/model/MoEBlock.java
index 86e5c8a..0e0f591 100644
--- a/jlama-core/src/main/java/com/github/tjake/jlama/model/MoEBlock.java
+++ b/jlama-core/src/main/java/com/github/tjake/jlama/model/MoEBlock.java
@@ -30,6 +30,7 @@
 public class MoEBlock implements FeedForward {
 
     private final AbstractModel model;
+    private final DistributedContext dctx;
     private final AbstractTensor moeGateWeight;
     private final int numberOfExperts;
     private final int numberOfExpertsPerToken;
@@ -55,6 +56,7 @@ public MoEBlock(
         AbstractTensor[] upProjectionWeights
     ) {
         this.model = model;
+        this.dctx = model.c.dctx();
         this.numberOfExperts = numberOfExperts;
         this.numberOfExpertsPerToken = numberOfExpertsPerToken;
         this.moeGateWeight = moeGateWeight;
@@ -87,21 +89,18 @@ public AbstractTensor forward(AbstractTensor lnemb, Optional<Consumer<List<Abstr
                 // Apply each experts gate to the input
                 VectorMath.pfor(0, numberOfExperts, i -> {
                     expertResults.set(
-                        TensorOperationsProvider.get()
-                            .dotProduct(
+                            TensorOperationsProvider.get().dotProduct(
                                 lnembSlice,
                                 moeGateWeight.slice(true, i),
-                                model.c.embeddingSegmentStart(),
-                                model.c.embeddingSegmentStart(),
-                                model.c.embeddingSegmentLength()
+                                0,
+                                0,
+                                model.c.embeddingLength
                             ),
                         0,
                         i
                     );
                 });
 
-                tensorReducer.ifPresent(func -> { func.accept(Collections.singletonList(expertResults)); });
-
                 // Pick the top experts for this token
                 VectorMath.softMax(expertResults, 0, numberOfExperts);
                 topk(expertResults);
@@ -114,52 +113,41 @@ public AbstractTensor forward(AbstractTensor lnemb, Optional<Consumer<List<Abstr
                     batchResults[0] = buf;
                     batchResults[1] = buf2;
 
-                    VectorMath.pchunk(0, hiddenLength, (chunkStart, chunkSize) -> {
+                    VectorMath.pchunk(dctx.hiddenSegmentStart, dctx.hiddenSegmentLength, (chunkStart, chunkSize) -> {
                         TensorOperationsProvider.get()
                             .dotProductBatchChunk(
                                 batchResults,
                                 lnembSlice,
                                 batchWeights,
-                                model.c.embeddingSegmentStart(),
-                                model.c.embeddingSegmentLength(),
+                                0,
+                                model.c.embeddingLength,
                                 chunkStart,
                                 chunkSize
                             );
                     });
 
-                    tensorReducer.ifPresent(func -> {
-                        tmpTensors1.clear();
-                        tmpTensors1.add(buf);
-                        tmpTensors1.add(buf2);
-                        func.accept(tmpTensors1);
-                    });
-
-                    VectorMath.pfor(0, hiddenLength, iv -> {
+                    VectorMath.pfor(dctx.hiddenSegmentStart, dctx.hiddenSegmentEnd, iv -> {
                         float w1 = buf.get(0, iv);
                         float w1a = ActivationFunction.eval(activationFunction, w1);
                         buf.set(w1a, 0, iv);
                     });
 
-                    TensorOperationsProvider.get().maccumulate(buf, buf2, 0, hiddenLength);
+                    TensorOperationsProvider.get().maccumulate(buf, buf2, dctx.hiddenSegmentStart, dctx.hiddenSegmentLength);
+
+                    tensorReducer.ifPresent(func -> func.accept(Collections.singletonList(buf)));
 
                     // matmul the projection and sum into result
                     try (AbstractTensor bufq = model.maybeQuantize(buf)) {
-                        VectorMath.pchunk(model.c.embeddingSegmentStart(), model.c.embeddingSegmentLength(), (chunkStart, chunkSize) -> {
+                        VectorMath.pchunk(0, model.c.embeddingLength, (chunkStart, chunkSize) -> {
                             TensorOperationsProvider.get()
                                 .dotProductChunk(moeResult, bufq, projectionWeight, 0, hiddenLength, chunkStart, chunkSize);
                         });
                     }
 
                     if (i == 0) {
-                        result.copyFrom(
-                            moeResult,
-                            moeResult.getOffset(0, model.c.embeddingSegmentStart()),
-                            result.getOffset(b, model.c.embeddingSegmentStart()),
-                            model.c.embeddingSegmentLength()
-                        );
+                        result.copyFrom(moeResult, 0, 0, model.c.embeddingLength);
                     } else {
-                        TensorOperationsProvider.get()
-                            .accumulate(result.slice(b), moeResult, model.c.embeddingSegmentStart(), model.c.embeddingSegmentLength());
+                        TensorOperationsProvider.get().accumulate(result.slice(b), moeResult, 0, model.c.embeddingLength);
                     }
                 }
             }
diff --git a/jlama-core/src/main/java/com/github/tjake/jlama/model/ModelSupport.java b/jlama-core/src/main/java/com/github/tjake/jlama/model/ModelSupport.java
index a4a0174..f3f4949 100644
--- a/jlama-core/src/main/java/com/github/tjake/jlama/model/ModelSupport.java
+++ b/jlama-core/src/main/java/com/github/tjake/jlama/model/ModelSupport.java
@@ -44,6 +44,8 @@
 import java.lang.reflect.InvocationTargetException;
 import java.nio.file.Path;
 import java.util.*;
+import java.util.function.Function;
+
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -106,7 +108,7 @@ public static AbstractModel loadModel(
         DType workingQuantizationType,
         Optional<DType> modelQuantization,
         Optional<Integer> threadCount,
-        Optional<Pair<Integer, Integer>> offset
+        Optional<Function<Config, DistributedContext>> distributedContextLoader
     ) {
 
         if (!model.exists()) {
@@ -137,7 +139,7 @@ public static AbstractModel loadModel(
 
             ModelSupport.ModelType modelType = SafeTensorSupport.detectModel(configFile);
             Config c = om.readValue(configFile, modelType.configClass);
-            offset.ifPresent(c::setOffset);
+            distributedContextLoader.ifPresent(loader -> c.setDistributedContext(loader.apply(c)));
 
             c.setWorkingDirectory(workingDirectory);
 
diff --git a/jlama-core/src/main/java/com/github/tjake/jlama/model/RMSNorm.java b/jlama-core/src/main/java/com/github/tjake/jlama/model/RMSNorm.java
index 8c88de0..2669197 100644
--- a/jlama-core/src/main/java/com/github/tjake/jlama/model/RMSNorm.java
+++ b/jlama-core/src/main/java/com/github/tjake/jlama/model/RMSNorm.java
@@ -51,11 +51,6 @@ public AbstractTensor forward(
                 ss += v * v;
             }
 
-            if (reducer.isPresent()) {
-                Pair<Float, Float> p = reducer.get().apply((float) ss, 0f);
-                ss = p.left;
-            }
-
             ss /= m.c.embeddingLength;
             ss += m.c.layerNormEps;
             ss = (1.0 / StrictMath.sqrt(ss));
diff --git a/jlama-core/src/main/java/com/github/tjake/jlama/model/TransformerBlock.java b/jlama-core/src/main/java/com/github/tjake/jlama/model/TransformerBlock.java
index 6ff10b0..cbc4e33 100644
--- a/jlama-core/src/main/java/com/github/tjake/jlama/model/TransformerBlock.java
+++ b/jlama-core/src/main/java/com/github/tjake/jlama/model/TransformerBlock.java
@@ -105,7 +105,7 @@ public AbstractTensor forward(
 
         // residual connection
         TensorOperationsProvider.get()
-            .accumulate(postAttention, embedding, model.c.embeddingSegmentStart(), model.c.embeddingSegmentLength());
+            .accumulate(postAttention, embedding, 0, model.c.embeddingLength);
 
         debug("post_attn_res", postAttention, layerIndex);
 
@@ -120,7 +120,7 @@ public AbstractTensor forward(
         }
 
         // residual connection
-        TensorOperationsProvider.get().accumulate(postFF, postAttention, model.c.embeddingSegmentStart(), model.c.embeddingSegmentLength());
+        TensorOperationsProvider.get().accumulate(postFF, postAttention, 0, model.c.embeddingLength);
 
         debug("post_ff_res", postFF, layerIndex);
 
diff --git a/jlama-core/src/main/java/com/github/tjake/jlama/model/bert/BertModel.java b/jlama-core/src/main/java/com/github/tjake/jlama/model/bert/BertModel.java
index d5b3e8f..fe5aba0 100644
--- a/jlama-core/src/main/java/com/github/tjake/jlama/model/bert/BertModel.java
+++ b/jlama-core/src/main/java/com/github/tjake/jlama/model/bert/BertModel.java
@@ -66,7 +66,7 @@ protected EmbedInput loadInputWeights() {
         );
 
         return (inputToken, position) -> {
-            AbstractTensor embedding = makeTensor(c.embeddingLength);
+            AbstractTensor embedding = makeDenseTensor(c.embeddingLength);
 
             for (int i = 0; i < c.embeddingLength; i++) {
                 float v = we.get(inputToken, i) + wte.get(0, i) + wpe.get(position, i);
@@ -81,9 +81,9 @@ protected EmbedInput loadInputWeights() {
 
     @Override
     protected TransformerBlock[] loadTransformerBlockWeights() {
-        TransformerBlock[] transformerBlocks = new TransformerBlock[c.getNumberOfLayers()];
+        TransformerBlock[] transformerBlocks = new TransformerBlock[c.dctx().embeddingSegmentLength];
 
-        for (int i = c.layerStart(); i < c.layerEnd(); i++) {
+        for (int i = c.dctx().layerStart; i < c.dctx().layerEnd; i++) {
             String b = "encoder.layer." + i + ".";
             String prefix = b + "attention.";
 
@@ -147,7 +147,7 @@ public float[] embed(String input) {
         Preconditions.checkArgument(encoded.length < c.contextLength);
         float[] outputEmbedding = new float[c.embeddingLength];
 
-        try (AbstractTensor kvmem = makeTensor(c.getNumberOfLayers(), 2, encoded.length, c.embeddingLength)) { // 2 for key and value
+        try (AbstractTensor kvmem = makeDenseTensor(c.dctx().numberOfLayers, 2, encoded.length, c.embeddingLength)) { // 2 for key and value
 
             int promptLength = encoded.length;
             float avgp = 1.0f / promptLength;
diff --git a/jlama-core/src/main/java/com/github/tjake/jlama/model/functions/EmbedInput.java b/jlama-core/src/main/java/com/github/tjake/jlama/model/functions/EmbedInput.java
index a56b88d..cb51199 100644
--- a/jlama-core/src/main/java/com/github/tjake/jlama/model/functions/EmbedInput.java
+++ b/jlama-core/src/main/java/com/github/tjake/jlama/model/functions/EmbedInput.java
@@ -34,15 +34,14 @@ default AbstractTensor batchInputsToEmbeddings(int[] inputTokens, int startPos)
         if (inputTokens.length == 1) return t;
 
         TensorShape tbs = TensorShape.of(inputTokens.length, t.shape().last());
-        if (t.shape().isSparse()) tbs = tbs.sparsify(t.shape().sparseOffset(), t.shape().sparseLength());
 
         AbstractTensor tb = TensorCache.instance.get(t.dType(), tbs);
-        tb.copyFrom(t, 0, 0, t.shape().sparseLength());
+        tb.copyFrom(t, 0, 0, t.shape().last());
         t.close();
 
         VectorMath.pfor(1, inputTokens.length, i -> {
             AbstractTensor ti = inputTokenToEmbedding(inputTokens[i], startPos + i);
-            tb.copyFrom(ti, 0, i * ti.shape().sparseLength(), ti.shape().sparseLength());
+            tb.copyFrom(ti, 0, i * ti.shape().last(), ti.shape().last());
             ti.close();
         });
 
diff --git a/jlama-core/src/main/java/com/github/tjake/jlama/model/functions/Generator.java b/jlama-core/src/main/java/com/github/tjake/jlama/model/functions/Generator.java
index da93658..3c503b8 100644
--- a/jlama-core/src/main/java/com/github/tjake/jlama/model/functions/Generator.java
+++ b/jlama-core/src/main/java/com/github/tjake/jlama/model/functions/Generator.java
@@ -16,10 +16,12 @@
 package com.github.tjake.jlama.model.functions;
 
 import com.github.tjake.jlama.safetensors.prompt.PromptContext;
+import com.github.tjake.jlama.safetensors.prompt.PromptSupport;
 import com.github.tjake.jlama.safetensors.prompt.ToolCall;
 import com.github.tjake.jlama.safetensors.tokenizer.Tokenizer;
 import java.util.Collections;
 import java.util.List;
+import java.util.Optional;
 import java.util.UUID;
 import java.util.function.BiConsumer;
 
@@ -129,4 +131,6 @@ Response generate(
     );
 
     Tokenizer getTokenizer();
+
+    Optional<PromptSupport> promptSupport();
 }
diff --git a/jlama-core/src/main/java/com/github/tjake/jlama/model/gemma/GemmaModel.java b/jlama-core/src/main/java/com/github/tjake/jlama/model/gemma/GemmaModel.java
index 5e183c5..a9b55e1 100644
--- a/jlama-core/src/main/java/com/github/tjake/jlama/model/gemma/GemmaModel.java
+++ b/jlama-core/src/main/java/com/github/tjake/jlama/model/gemma/GemmaModel.java
@@ -78,17 +78,17 @@ protected TransformerBlock[] loadTransformerBlockWeights() {
             logger.info("Quantizing model with {} - Please hold...", qType);
         }
 
-        TransformerBlock[] transformerBlocks = new TransformerBlock[c.getNumberOfLayers()];
+        TransformerBlock[] transformerBlocks = new TransformerBlock[c.dctx().numberOfLayers];
 
-        IntStream.range(c.layerStart(), c.layerEnd()).parallel().forEach(i -> {
+        IntStream.range(c.dctx().layerStart, c.dctx().layerEnd).parallel().forEach(i -> {
             String base = "model.layers." + i + ".";
             String prefix = base + "self_attn.";
             CausalSelfAttention attention = new CausalSelfAttention(
                 this,
-                weights.load(prefix + "q_proj.weight", c.offset()).quantize(qType),
-                weights.load(prefix + "k_proj.weight", c.offset()).quantize(qType),
-                weights.load(prefix + "v_proj.weight", c.offset()).quantize(qType),
-                weights.load(prefix + "o_proj.weight", c.offset()).quantize(qType)
+                weights.load(prefix + "q_proj.weight", c.dctx(), true, false).quantize(qType),
+                weights.load(prefix + "k_proj.weight", c.dctx(), true, false).quantize(qType),
+                weights.load(prefix + "v_proj.weight", c.dctx(), true, false).quantize(qType),
+                weights.load(prefix + "o_proj.weight", c.dctx(), false, true).quantize(qType)
             );
 
             prefix = base + "mlp.";
@@ -96,17 +96,17 @@ protected TransformerBlock[] loadTransformerBlockWeights() {
             MLPBlock mlp = new MLPBlock(
                 this,
                 ActivationFunction.Type.GELU,
-                weights.load(prefix + "gate_proj.weight", c.offset()).quantize(qType), // w1
-                weights.load(prefix + "down_proj.weight").quantize(qType), // w2
-                weights.load(prefix + "up_proj.weight", c.offset()).quantize(qType)
-            ); // w3
+                weights.load(prefix + "gate_proj.weight", c.dctx(), true, false).quantize(qType), // w1
+                weights.load(prefix + "down_proj.weight", c.dctx(), false, true).quantize(qType), // w2
+                weights.load(prefix + "up_proj.weight", c.dctx(), true, false).quantize(qType) // w3
+            );
 
             transformerBlocks[i] = new TransformerBlock(
                 this,
                 i,
-                new RMSNorm(this, weights.load(base + "input_layernorm.weight", c.offset()).quantize(qType), 1.0f),
+                new RMSNorm(this, weights.load(base + "input_layernorm.weight").quantize(qType), 1.0f),
                 attention,
-                new RMSNorm(this, weights.load(base + "post_attention_layernorm.weight", c.offset()).quantize(qType), 1.0f),
+                new RMSNorm(this, weights.load(base + "post_attention_layernorm.weight").quantize(qType), 1.0f),
                 mlp
             );
         });
@@ -116,25 +116,19 @@ protected TransformerBlock[] loadTransformerBlockWeights() {
 
     @Override
     protected EmbedInput loadInputWeights() {
-
-        if (wte == null) wte = weights.load("model.embed_tokens.weight", c.offset()).quantize(workingDType); // Don't quantize this, it's
-                                                                                                             // used for the embedding layer
+        // Don't quantize this, it's used for the embedding layer
+        if (wte == null) wte = weights.load("model.embed_tokens.weight").quantize(workingDType);
 
         return (inputToken, position) -> {
-            AbstractTensor embedding = makeTensor(c.embeddingLength);
+            AbstractTensor embedding = makeDenseTensor(c.embeddingLength);
             AbstractTensor at = wte.slice(true, inputToken);
-            if (wte.dType() != embedding.dType()) at = TensorOperationsProvider.get()
-                .quantize(at, embedding.dType(), c.embeddingSegmentStart(), c.embeddingSegmentLength());
-
-            embedding.copyFrom(
-                at,
-                at.getOffset(0, c.embeddingSegmentStart()),
-                embedding.getOffset(c.embeddingSegmentStart()),
-                c.embeddingSegmentLength()
-            );
+            if (wte.dType() != embedding.dType())
+                at = TensorOperationsProvider.get().quantize(at, embedding.dType(), 0, c.embeddingLength);
+
+            embedding.copyFrom(at, 0, 0, c.embeddingLength);
 
             // This is important for Gemma, but not for Llama
-            TensorOperationsProvider.get().scale(embeddingScalingFactor, embedding, c.embeddingSegmentStart(), c.embeddingSegmentLength());
+            TensorOperationsProvider.get().scale(embeddingScalingFactor, embedding, 0, c.embeddingLength);
 
             return embedding;
         };
@@ -144,7 +138,7 @@ protected EmbedInput loadInputWeights() {
     protected SampleOutput loadOutputWeights() {
         DType qType = modelQType.orElse(this.modelDType);
 
-        if (wte == null) wte = weights.load("model.embed_tokens.weight", c.offset()).quantize(workingDType); // Don't quantize this, it's
+        if (wte == null) wte = weights.load("model.embed_tokens.weight").quantize(workingDType); // Don't quantize this, it's
                                                                                                              // used for the embedding layer
 
         final LayerNorm layerNorm = new RMSNorm(this, weights.load("model.norm.weight").quantize(qType), 1.0f);
diff --git a/jlama-core/src/main/java/com/github/tjake/jlama/model/gpt2/GPT2Model.java b/jlama-core/src/main/java/com/github/tjake/jlama/model/gpt2/GPT2Model.java
index d019f0f..8e96c26 100644
--- a/jlama-core/src/main/java/com/github/tjake/jlama/model/gpt2/GPT2Model.java
+++ b/jlama-core/src/main/java/com/github/tjake/jlama/model/gpt2/GPT2Model.java
@@ -52,9 +52,9 @@ protected EmbedInput loadInputWeights() {
         final AbstractTensor wpe = weights.load("wpe.weight");
 
         return (inputToken, position) -> {
-            AbstractTensor embedding = makeTensor(1, c.embeddingLength);
+            AbstractTensor embedding = makeDenseTensor(1, c.embeddingLength);
 
-            for (int i = c.embeddingSegmentStart(); i < c.embeddingSegmentLength(); i++) {
+            for (int i = 0; i < c.embeddingLength; i++) {
                 float v = wte.get(inputToken, i) + wpe.get(position, i);
                 embedding.set(v, 0, i);
             }
@@ -65,9 +65,9 @@ protected EmbedInput loadInputWeights() {
 
     @Override
     protected TransformerBlock[] loadTransformerBlockWeights() {
-        TransformerBlock[] transformerBlocks = new TransformerBlock[c.getNumberOfLayers()];
+        TransformerBlock[] transformerBlocks = new TransformerBlock[c.dctx().numberOfLayers];
 
-        for (int i = c.layerStart(); i < c.layerEnd(); i++) {
+        for (int i = c.dctx().layerStart; i < c.dctx().layerEnd; i++) {
             String b = "h." + i + ".";
             String prefix = b + "attn.";
 
diff --git a/jlama-core/src/main/java/com/github/tjake/jlama/model/llama/LlamaModel.java b/jlama-core/src/main/java/com/github/tjake/jlama/model/llama/LlamaModel.java
index 74f1b6a..85e8b55 100644
--- a/jlama-core/src/main/java/com/github/tjake/jlama/model/llama/LlamaModel.java
+++ b/jlama-core/src/main/java/com/github/tjake/jlama/model/llama/LlamaModel.java
@@ -65,24 +65,19 @@ public ModelSupport.ModelType getModelType() {
     @Override
     protected EmbedInput loadInputWeights() {
 
-        final AbstractTensor wte = weights.load("model.embed_tokens.weight", c.offset()).quantize(workingDType); // Don't quantize this,
-                                                                                                                 // it's used for the
-                                                                                                                 // embedding layer
+        // Don't quantize this, it's used for the embedding layer
+        final AbstractTensor wte = weights.load("model.embed_tokens.weight").quantize(workingDType);
 
         return (inputToken, position) -> {
-            AbstractTensor embedding = makeTensor(1, c.embeddingLength);
+            AbstractTensor embedding = makeDenseTensor(1, c.embeddingLength);
 
             AbstractTensor at = wte.slice(true, inputToken);
 
-            if (wte.dType() != embedding.dType()) at = TensorOperationsProvider.get()
-                .quantize(at, embedding.dType(), c.embeddingSegmentStart(), c.embeddingSegmentLength());
+            if (wte.dType() != embedding.dType())
+                at = TensorOperationsProvider.get().quantize(at, embedding.dType(), 0, c.embeddingLength);
 
-            embedding.copyFrom(
-                at,
-                at.getOffset(0, c.embeddingSegmentStart()),
-                embedding.getOffset(c.embeddingSegmentStart()),
-                c.embeddingSegmentLength()
-            );
+            //Always copy the entire embedding
+            embedding.copyFrom(at, 0, 0, c.embeddingLength);
 
             return embedding;
         };
@@ -95,17 +90,17 @@ protected TransformerBlock[] loadTransformerBlockWeights() {
             logger.info("Quantizing model with {} - Please hold...", qType);
         }
 
-        TransformerBlock[] transformerBlocks = new TransformerBlock[c.getNumberOfLayers()];
+        TransformerBlock[] transformerBlocks = new TransformerBlock[c.dctx().numberOfLayers];
 
-        IntStream.range(c.layerStart(), c.layerEnd()).parallel().forEach(i -> {
+        IntStream.range(c.dctx().layerStart, c.dctx().layerEnd).parallel().forEach(i -> {
             String base = "model.layers." + i + ".";
             String prefix = base + "self_attn.";
             CausalSelfAttention attention = new CausalSelfAttention(
                 this,
-                weights.load(prefix + "q_proj.weight", c.offset()).quantize(qType),
-                weights.load(prefix + "k_proj.weight", c.offset()).quantize(qType),
-                weights.load(prefix + "v_proj.weight", c.offset()).quantize(qType),
-                weights.load(prefix + "o_proj.weight", c.offset()).quantize(qType)
+                weights.load(prefix + "q_proj.weight", c.dctx(), true, false).quantize(qType),
+                weights.load(prefix + "k_proj.weight", c.dctx(), true, false).quantize(qType),
+                weights.load(prefix + "v_proj.weight", c.dctx(), true, false).quantize(qType),
+                weights.load(prefix + "o_proj.weight", c.dctx(), false, true).quantize(qType)
             );
 
             prefix = base + "mlp.";
@@ -113,17 +108,17 @@ protected TransformerBlock[] loadTransformerBlockWeights() {
             MLPBlock mlp = new MLPBlock(
                 this,
                 c.activationFunction,
-                weights.load(prefix + "gate_proj.weight", c.offset()).quantize(qType), // w1
-                weights.load(prefix + "down_proj.weight").quantize(qType), // w2
-                weights.load(prefix + "up_proj.weight", c.offset()).quantize(qType)
+                weights.load(prefix + "gate_proj.weight", c.dctx(), true, false).quantize(qType), // w1
+                weights.load(prefix + "down_proj.weight", c.dctx(), false, true).quantize(qType), // w2
+                weights.load(prefix + "up_proj.weight", c.dctx(), true, false).quantize(qType)
             ); // w3
 
             transformerBlocks[i] = new TransformerBlock(
                 this,
                 i,
-                new RMSNorm(this, weights.load(base + "input_layernorm.weight", c.offset()).quantize(qType)),
+                new RMSNorm(this, weights.load(base + "input_layernorm.weight").quantize(qType)),
                 attention,
-                new RMSNorm(this, weights.load(base + "post_attention_layernorm.weight", c.offset()).quantize(qType)),
+                new RMSNorm(this, weights.load(base + "post_attention_layernorm.weight").quantize(qType)),
                 mlp
             );
         });
@@ -156,8 +151,9 @@ protected AbstractTensor maybeQuantize(AbstractTensor t) {
         Preconditions.checkArgument(t.dims() == 2, "Unexpected shape");
         if (t.dType() == workingQType) return super.maybeQuantize(t);
 
-        return t.shape().last() == c.embeddingLength
-            ? TensorOperationsProvider.get().quantize(t, workingQType, c.embeddingSegmentStart(), c.embeddingSegmentLength())
-            : TensorOperationsProvider.get().quantize(t, workingQType, 0, Ints.checkedCast(t.shape().last()));
+        return //t.shape().last() == c.embeddingLength
+            //? TensorOperationsProvider.get().quantize(t, workingQType, c.embeddingSegmentStart(), c.embeddingSegmentLength())
+            //:
+        TensorOperationsProvider.get().quantize(t, workingQType, 0, Ints.checkedCast(t.shape().last()));
     }
 }
diff --git a/jlama-core/src/main/java/com/github/tjake/jlama/model/mixtral/MixtralModel.java b/jlama-core/src/main/java/com/github/tjake/jlama/model/mixtral/MixtralModel.java
index c3e3228..984febd 100644
--- a/jlama-core/src/main/java/com/github/tjake/jlama/model/mixtral/MixtralModel.java
+++ b/jlama-core/src/main/java/com/github/tjake/jlama/model/mixtral/MixtralModel.java
@@ -67,17 +67,17 @@ protected TransformerBlock[] loadTransformerBlockWeights() {
             logger.info("Quantizing model with {} - Please hold...", qType);
         }
 
-        TransformerBlock[] transformerBlocks = new TransformerBlock[c.getNumberOfLayers()];
+        TransformerBlock[] transformerBlocks = new TransformerBlock[c.dctx().numberOfLayers];
 
-        IntStream.range(c.layerStart(), c.layerEnd()).parallel().forEach(i -> {
+        IntStream.range(c.dctx().layerStart, c.dctx().layerEnd).parallel().forEach(i -> {
             String base = "model.layers." + i + ".";
             String prefix = base + "self_attn.";
             CausalSelfAttention attention = new CausalSelfAttention(
                 this,
-                weights.load(prefix + "q_proj.weight", c.offset()).quantize(qType),
-                weights.load(prefix + "k_proj.weight", c.offset()).quantize(qType),
-                weights.load(prefix + "v_proj.weight", c.offset()).quantize(qType),
-                weights.load(prefix + "o_proj.weight", c.offset()).quantize(qType)
+                weights.load(prefix + "q_proj.weight", c.dctx(), true, false).quantize(qType),
+                weights.load(prefix + "k_proj.weight", c.dctx(), true, false).quantize(qType),
+                weights.load(prefix + "v_proj.weight", c.dctx(), true, false).quantize(qType),
+                weights.load(prefix + "o_proj.weight").quantize(qType)
             );
 
             prefix = base + "block_sparse_moe.";
@@ -88,9 +88,9 @@ protected TransformerBlock[] loadTransformerBlockWeights() {
 
             for (int e = 0; e < mixtralConfig.numberOfExperts; e++) {
                 String expertPrefix = prefix + "experts." + e + ".";
-                expertGateWeights[e] = weights.load(expertPrefix + "w1.weight", c.offset()).quantize(qType);
+                expertGateWeights[e] = weights.load(expertPrefix + "w1.weight", c.dctx(), true, false).quantize(qType);
                 expertDownWeights[e] = weights.load(expertPrefix + "w2.weight").quantize(qType);
-                expertUpWeights[e] = weights.load(expertPrefix + "w3.weight", c.offset()).quantize(qType);
+                expertUpWeights[e] = weights.load(expertPrefix + "w3.weight", c.dctx(), true, false).quantize(qType);
             }
 
             MoEBlock moe = new MoEBlock(
@@ -98,7 +98,7 @@ protected TransformerBlock[] loadTransformerBlockWeights() {
                 mixtralConfig.numberOfExperts,
                 mixtralConfig.numberOfExpertsPerToken,
                 c.activationFunction,
-                weights.load(prefix + "gate.weight", c.offset()).quantize(qType),
+                weights.load(prefix + "gate.weight").quantize(qType),
                 expertGateWeights, // w1
                 expertDownWeights, // w2
                 expertUpWeights
@@ -107,9 +107,9 @@ protected TransformerBlock[] loadTransformerBlockWeights() {
             transformerBlocks[i] = new TransformerBlock(
                 this,
                 i,
-                new RMSNorm(this, weights.load(base + "input_layernorm.weight", c.offset()).quantize(qType)),
+                new RMSNorm(this, weights.load(base + "input_layernorm.weight").quantize(qType)),
                 attention,
-                new RMSNorm(this, weights.load(base + "post_attention_layernorm.weight", c.offset()).quantize(qType)),
+                new RMSNorm(this, weights.load(base + "post_attention_layernorm.weight").quantize(qType)),
                 moe
             );
         });
diff --git a/jlama-core/src/main/java/com/github/tjake/jlama/safetensors/Config.java b/jlama-core/src/main/java/com/github/tjake/jlama/safetensors/Config.java
index b2ad7f0..c632065 100644
--- a/jlama-core/src/main/java/com/github/tjake/jlama/safetensors/Config.java
+++ b/jlama-core/src/main/java/com/github/tjake/jlama/safetensors/Config.java
@@ -17,6 +17,7 @@
 
 import com.github.tjake.jlama.math.ActivationFunction;
 import com.github.tjake.jlama.math.VectorMath;
+import com.github.tjake.jlama.model.DistributedContext;
 import com.github.tjake.jlama.tensor.TensorCache;
 import com.github.tjake.jlama.util.Pair;
 import com.google.common.base.Preconditions;
@@ -28,6 +29,7 @@
 public class Config {
     public final int contextLength;
     public final int embeddingLength;
+    public final int attentionLength;
     public final int hiddenLength;
     public final int numberOfHeads;
     public final int numberOfKeyValueHeads;
@@ -36,26 +38,15 @@ public class Config {
     public final int headGroupSize;
     public final int kvLength;
     public final boolean isGQA;
-    protected final int numberOfLayers;
+    public final int numberOfLayers;
     public final float layerNormEps;
     public final int vocabularySize;
     public final int bosToken;
     public final List<Integer> eosTokens;
     public final Optional<float[][]> ropeFreqs;
-    private volatile Optional<Pair<Integer, Integer>> offset;
+    private volatile DistributedContext dctx;
     private volatile File workingDirectory;
 
-    // Suppliers to store values that chance when offset is adjusted
-    private volatile int embeddingSegmentStart;
-    private volatile int embeddingSegmentLength;
-    private volatile int embeddingSegmentEnd;
-    private volatile int kvSegmentStart;
-    private volatile int kvSegmentLength;
-    private volatile int kvSegmentEnd;
-    private volatile int headStart;
-    private volatile int headEnd;
-    private volatile int groupHeadStart;
-    private volatile int groupHeadEnd;
 
     public final TensorCache tensorCache;
 
@@ -109,6 +100,7 @@ public Config(
         Integer headSize
     ) {
         this.contextLength = contextLength;
+        this.attentionLength = numberOfHeads * headSize;
         this.embeddingLength = embeddingLength;
         this.hiddenLength = hiddenLength;
         this.numberOfHeads = numberOfHeads;
@@ -131,22 +123,11 @@ public Config(
             );
 
         // Set default values
-        setOffset(null);
+        this.dctx = DistributedContext.builder(this).build();
     }
 
-    public void setOffset(Pair<Integer, Integer> offset) {
-        this.offset = Optional.ofNullable(offset);
-
-        this.embeddingSegmentStart = this.offset.map(Pair::left).orElse(0);
-        this.embeddingSegmentLength = this.offset.map(Pair::right).orElse(embeddingLength);
-        this.embeddingSegmentEnd = embeddingSegmentStart + embeddingSegmentLength;
-        this.kvSegmentStart = embeddingSegmentStart / headGroupSize;
-        this.kvSegmentEnd = embeddingSegmentEnd / headGroupSize;
-        this.kvSegmentLength = embeddingSegmentLength / headGroupSize;
-        this.headStart = embeddingSegmentStart / headSize;
-        this.headEnd = embeddingSegmentEnd / headSize;
-        this.groupHeadStart = kvSegmentStart / headSize;
-        this.groupHeadEnd = kvSegmentEnd / headSize;
+    public void setDistributedContext(DistributedContext dctx) {
+        this.dctx = dctx;
     }
 
     public void setWorkingDirectory(File workingDirectory) {
@@ -163,56 +144,12 @@ public Optional<File> workingDirectory() {
         return Optional.ofNullable(this.workingDirectory);
     }
 
-    public Optional<Pair<Integer, Integer>> offset() {
-        return offset;
-    }
-
-    public int layerStart() {
-        return 0;
-    }
-
-    public int layerEnd() {
-        return numberOfLayers;
-    }
-
-    public int getNumberOfLayers() {
-        return numberOfLayers;
-    }
-
-    public int embeddingSegmentStart() {
-        return embeddingSegmentStart;
-    }
-
-    public int embeddingSegmentLength() {
-        return embeddingSegmentLength;
-    }
-
-    public int kvSegmentStart() {
-        return kvSegmentStart;
-    }
-
-    public int kvSegmentLength() {
-        return kvSegmentLength;
-    }
-
-    public int headStart() {
-        return headStart;
-    }
-
-    public int headEnd() {
-        return headEnd;
+    public DistributedContext dctx() {
+        return dctx;
     }
 
     public int maybeMapToGroupHead(int head) {
         if (!isGQA) return head;
         return Math.floorDiv(head, headGroupSize);
     }
-
-    public int groupHeadStart() {
-        return groupHeadStart;
-    }
-
-    public int groupHeadEnd() {
-        return groupHeadEnd;
-    }
 }
diff --git a/jlama-core/src/main/java/com/github/tjake/jlama/safetensors/SafeTensorIndex.java b/jlama-core/src/main/java/com/github/tjake/jlama/safetensors/SafeTensorIndex.java
index 55df3e7..f9c1109 100644
--- a/jlama-core/src/main/java/com/github/tjake/jlama/safetensors/SafeTensorIndex.java
+++ b/jlama-core/src/main/java/com/github/tjake/jlama/safetensors/SafeTensorIndex.java
@@ -18,6 +18,7 @@
 import com.fasterxml.jackson.annotation.JsonCreator;
 import com.fasterxml.jackson.annotation.JsonProperty;
 import com.fasterxml.jackson.databind.ObjectMapper;
+import com.github.tjake.jlama.model.DistributedContext;
 import com.github.tjake.jlama.tensor.AbstractTensor;
 import com.github.tjake.jlama.util.Pair;
 import com.google.common.collect.ImmutableMap;
@@ -179,15 +180,11 @@ public Map<String, TensorInfo> tensorInfoMap() {
     }
 
     @Override
-    public AbstractTensor load(String name, Optional<Pair<Integer, Integer>> offset) {
+    public AbstractTensor load(String name, DistributedContext dctx, boolean sparseRows, boolean sparseColumns) {
         Weights w = weightMap.get(name);
         if (w == null) throw new NoSuchElementException(name);
 
-        AbstractTensor t = w.load(name);
-        return offset.map(o -> {
-            logger.debug("Sparsifying tensor {} with shape {}", name, o);
-            return t.sparsify(o.left, o.right);
-        }).orElse(t);
+        return w.load(name, dctx, sparseRows, sparseColumns);
     }
 
     @Override
diff --git a/jlama-core/src/main/java/com/github/tjake/jlama/safetensors/SafeTensorSupport.java b/jlama-core/src/main/java/com/github/tjake/jlama/safetensors/SafeTensorSupport.java
index 4fdaef9..a847d6b 100644
--- a/jlama-core/src/main/java/com/github/tjake/jlama/safetensors/SafeTensorSupport.java
+++ b/jlama-core/src/main/java/com/github/tjake/jlama/safetensors/SafeTensorSupport.java
@@ -243,13 +243,14 @@ public static Path quantizeModel(
         String baseDirName = modelRoot.getName(modelRoot.getNameCount() - 1).toString();
         Path parentPath = modelRoot.getParent();
 
-        Path qPath = outputRoot.orElseGet(() -> Paths.get(parentPath.toString(), baseDirName + "-jlama-" + modelQuantization.name()));
+        Path qPath = outputRoot.orElseGet(() -> Paths.get(parentPath.toString(), baseDirName + "-Jlama-" + modelQuantization.name()));
         File qDir = qPath.toFile();
         qDir.mkdirs();
 
         // Copy config.json and tokenizer.json
         Files.copy(modelRoot.resolve("config.json"), qPath.resolve("config.json"));
         Files.copy(modelRoot.resolve("tokenizer.json"), qPath.resolve("tokenizer.json"));
+        Files.copy(modelRoot.resolve("README.md"), qPath.resolve("README.md"));
 
         if (Files.exists(modelRoot.resolve("tokenizer_config.json"))) Files.copy(
             modelRoot.resolve("tokenizer_config.json"),
@@ -404,7 +405,7 @@ private static Pair<InputStream, Long> getResponse(String urlString, Optional<St
 
         if (responseCode == HttpURLConnection.HTTP_OK) {
             // If the response code is 200 (HTTP_OK), return the input stream
-            return Pair.create(connection.getInputStream(), connection.getContentLengthLong());
+            return Pair.of(connection.getInputStream(), connection.getContentLengthLong());
         } else {
             // If the response code is not 200, throw an IOException
             throw new IOException("HTTP response code: " + responseCode + " for URL: " + urlString);
diff --git a/jlama-core/src/main/java/com/github/tjake/jlama/safetensors/WeightLoader.java b/jlama-core/src/main/java/com/github/tjake/jlama/safetensors/WeightLoader.java
index 8e7f51b..6aeba23 100644
--- a/jlama-core/src/main/java/com/github/tjake/jlama/safetensors/WeightLoader.java
+++ b/jlama-core/src/main/java/com/github/tjake/jlama/safetensors/WeightLoader.java
@@ -15,6 +15,7 @@
  */
 package com.github.tjake.jlama.safetensors;
 
+import com.github.tjake.jlama.model.DistributedContext;
 import com.github.tjake.jlama.tensor.AbstractTensor;
 import com.github.tjake.jlama.util.Pair;
 import java.util.Map;
@@ -27,10 +28,10 @@ public interface WeightLoader extends AutoCloseable {
     Map<String, TensorInfo> tensorInfoMap();
 
     default AbstractTensor load(String name) {
-        return load(name, Optional.empty());
+        return load(name, null, false, false);
     }
 
-    AbstractTensor load(String name, Optional<Pair<Integer, Integer>> offset);
+    AbstractTensor load(String name, DistributedContext dctx, boolean sparseRows, boolean sparseColumns);
 
     DType getModelDType();
 }
diff --git a/jlama-core/src/main/java/com/github/tjake/jlama/safetensors/Weights.java b/jlama-core/src/main/java/com/github/tjake/jlama/safetensors/Weights.java
index c14b2f6..2b20beb 100644
--- a/jlama-core/src/main/java/com/github/tjake/jlama/safetensors/Weights.java
+++ b/jlama-core/src/main/java/com/github/tjake/jlama/safetensors/Weights.java
@@ -16,6 +16,7 @@
 package com.github.tjake.jlama.safetensors;
 
 import com.github.tjake.jlama.math.FloatConversions;
+import com.github.tjake.jlama.model.DistributedContext;
 import com.github.tjake.jlama.tensor.*;
 import com.github.tjake.jlama.util.Pair;
 import com.google.common.collect.ImmutableMap;
@@ -74,16 +75,38 @@ public Map<String, TensorInfo> tensorInfoMap() {
     }
 
     @Override
-    public AbstractTensor load(String name, Optional<Pair<Integer, Integer>> offset) throws NoSuchElementException {
+    public AbstractTensor load(String name, DistributedContext dctx, boolean sparseRows, boolean sparseColumns) throws NoSuchElementException {
         TensorInfo info = tensorInfoMap.get(name);
         if (info == null) throw new NoSuchElementException(name + " not found in weights");
 
         if (info.shape.length < 1) throw new RuntimeException("Invalid shape dimensions " + info.shape.length + " encountered for " + name);
 
+        if (dctx != null && info.shape.length != 2) {
+            throw new RuntimeException("Invalid shape dimensions " + info.shape.length + " encountered for " + name + " with offset");
+        }
+
+        int positionOffset = Ints.checkedCast(info.dataOffsets[0]);
+        int positionLimit = Ints.checkedCast(info.dataOffsets[1]);
+        TensorShape shape = TensorShape.of(info.shape);
+
+        // If this is a sparse tensor, we need to fetch only the section of the tensor that is needed
+        if (dctx != null && sparseRows) {
+            int rows = info.shape[0];
+            int columnLength = info.shape[1] * info.dType.size();
+
+            //Hack for Q4
+            if (info.dType == DType.Q4)
+                columnLength /= 2;
+
+            positionOffset = Ints.checkedCast(info.dataOffsets[0]) + (dctx.getShardOffsetForLength(rows) * columnLength);
+            positionLimit = positionOffset + (dctx.getShardLength(rows) * columnLength);
+            shape = TensorShape.sparseRow(info.shape, Pair.of(dctx.getShardOffsetForLength(rows), dctx.getShardLength(rows)));
+        }
+
         ByteBuffer b = bytes.duplicate()
             .order(ByteOrder.LITTLE_ENDIAN)
-            .position(Ints.checkedCast(info.dataOffsets[0]))
-            .limit(Ints.checkedCast(info.dataOffsets[1]));
+            .position(Ints.checkedCast(positionOffset))
+            .limit(Ints.checkedCast(positionLimit));
 
         int len;
         FloatBuffer fb;
@@ -92,7 +115,7 @@ public AbstractTensor load(String name, Optional<Pair<Integer, Integer>> offset)
         switch (info.dType) {
             case F32:
                 fb = b.asFloatBuffer().slice();
-                t = new FloatBufferTensor(name, fb, TensorShape.of(info.shape), true);
+                t = new FloatBufferTensor(name, fb, shape, true);
                 break;
             case F16:
                 // If the majority of the weights are F32 then convert to F32
@@ -104,10 +127,10 @@ public AbstractTensor load(String name, Optional<Pair<Integer, Integer>> offset)
                         float v = Float.float16ToFloat(s);
                         bb.putFloat(i, v);
                     }
-                    t = new FloatBufferTensor(bb.asFloatBuffer(), TensorShape.of(info.shape), true);
+                    t = new FloatBufferTensor(bb.asFloatBuffer(), shape, true);
                 } else {
                     sb = b.asShortBuffer().slice();
-                    t = new Float16BufferTensor(name, sb, TensorShape.of(info.shape), true);
+                    t = new Float16BufferTensor(name, sb, shape, true);
                 }
                 break;
             case BF16:
@@ -119,25 +142,27 @@ public AbstractTensor load(String name, Optional<Pair<Integer, Integer>> offset)
                         float v = FloatConversions.bFloat16ToFloat32(s);
                         bb.putFloat(i, v);
                     }
-                    t = new FloatBufferTensor(bb.asFloatBuffer(), TensorShape.of(info.shape), true);
+                    t = new FloatBufferTensor(bb.asFloatBuffer(), shape, true);
                 } else {
                     sb = b.asShortBuffer().slice();
-                    t = new BFloat16BufferTensor(name, sb, TensorShape.of(info.shape), true);
+                    t = new BFloat16BufferTensor(name, sb, shape, true);
                 }
                 break;
             case Q4:
-                FloatBufferTensor qb = (FloatBufferTensor) parent.orElse(this).load(name + ".qb", offset);
-                t = new Q4ByteBufferTensor(name, b.slice(), qb, TensorShape.of(info.shape), true);
+                FloatBufferTensor qb = (FloatBufferTensor) parent.orElse(this).load(name + ".qb", dctx, sparseRows, false); //only need to sparsify once
+                t = new Q4ByteBufferTensor(name, b.slice(), qb, shape, true);
                 break;
             case I8:
-                FloatBufferTensor qb1 = (FloatBufferTensor) parent.orElse(this).load(name + ".qb", offset);
-                t = new Q8ByteBufferTensor(name, b.slice(), qb1, TensorShape.of(info.shape), true);
+                FloatBufferTensor qb1 = (FloatBufferTensor) parent.orElse(this).load(name + ".qb", dctx, sparseRows, false); //only need to sparsify once
+                t = new Q8ByteBufferTensor(name, b.slice(), qb1, shape, true);
                 break;
             default:
                 throw new IllegalArgumentException("Unsupported Tensor type: " + info.dType.name() + " for " + name);
         }
 
-        return offset.map(o -> t.sparsify(o.left, o.right)).orElse(t);
+        return dctx != null && sparseColumns && dctx.hasModelShard()
+                ? t.sparsify(dctx.getShardOffsetForLength(shape.last()), dctx.getShardLength(shape.last()))
+                : t;
     }
 
     @Override
diff --git a/jlama-core/src/main/java/com/github/tjake/jlama/tensor/AbstractTensor.java b/jlama-core/src/main/java/com/github/tjake/jlama/tensor/AbstractTensor.java
index ca5bb92..b0b9127 100644
--- a/jlama-core/src/main/java/com/github/tjake/jlama/tensor/AbstractTensor.java
+++ b/jlama-core/src/main/java/com/github/tjake/jlama/tensor/AbstractTensor.java
@@ -58,7 +58,7 @@ protected AbstractTensor(DType dType, TensorShape shape, boolean cacheSlices) {
         this.shape = shape;
         this.metadata = new HashMap<>();
         this.sliceCache = cacheSlices ? new AbstractTensor[shape.first()] : null;
-        this.stride = shape.first() > 1 && dims() == 2 ? getOffset(1, shape.sparseOffset()) : 0;
+        this.stride = shape.first() > 1 && dims() == 2 ? getOffset(shape.sparseRowOffset() + 1, shape.sparseColumnOffset()) : 0;
     }
 
     public static AbstractTensor make(DType dType, TensorShape shape) {
@@ -127,10 +127,10 @@ public AbstractTensor slice(boolean cacheInnerSlice, int... dims) {
         TensorShape slicedShape = shape.slice(dims.length);
         int totalOffset = 0;
         if (dims.length == 1 && this.shape.dims() == 2) {
-            totalOffset = shape.sparseLength() * dims[0];
+            totalOffset = shape.sparseColumnLength() * dims[0];
         } else {
             for (int d = 0; d <= dims.length - 1; d++) {
-                int offset = shape.sparseLength();
+                int offset = shape.sparseColumnLength();
                 for (int i = shape.dims() - 2; i > d; i--) { // factor scaling of each dim shape
                     offset *= shape.dim(i);
                 }
@@ -153,7 +153,7 @@ public AbstractTensor<V, T> sparsify(int offset, int length) {
 
         if (length == shape.last()) return this;
 
-        AbstractTensor<V, T> sparseT = this.make(shape.sparsify(offset, length));
+        AbstractTensor<V, T> sparseT = this.make(shape.sparsifyColumns(offset, length));
         int originalLength = shape.last();
 
         int[] cursor = new int[shape.dims()];
diff --git a/jlama-core/src/main/java/com/github/tjake/jlama/tensor/FloatBufferTensor.java b/jlama-core/src/main/java/com/github/tjake/jlama/tensor/FloatBufferTensor.java
index a32a9fd..078886a 100644
--- a/jlama-core/src/main/java/com/github/tjake/jlama/tensor/FloatBufferTensor.java
+++ b/jlama-core/src/main/java/com/github/tjake/jlama/tensor/FloatBufferTensor.java
@@ -122,7 +122,7 @@ public MemorySegment getMemorySegment() {
 
     @Override
     public void copyFrom(AbstractTensor src, int srcOffset, int destOffset, int length) {
-        // Preconditions.checkArgument(this.dType == src.dType, "Different types");
+        Preconditions.checkArgument(this.dType == src.dType, "Different types");
         // Preconditions.checkArgument(!b.isReadOnly());
         segment.asSlice(getMemorySegmentOffset(destOffset), length * dType.size())
             .copyFrom(src.getMemorySegment().asSlice(src.getMemorySegmentOffset(srcOffset), length * dType.size()));
diff --git a/jlama-core/src/main/java/com/github/tjake/jlama/tensor/KvBufferCache.java b/jlama-core/src/main/java/com/github/tjake/jlama/tensor/KvBufferCache.java
index e8dc6d5..a0b9d60 100644
--- a/jlama-core/src/main/java/com/github/tjake/jlama/tensor/KvBufferCache.java
+++ b/jlama-core/src/main/java/com/github/tjake/jlama/tensor/KvBufferCache.java
@@ -16,6 +16,8 @@
 package com.github.tjake.jlama.tensor;
 
 import com.github.tjake.jlama.model.AbstractModel;
+import com.github.tjake.jlama.model.DistributedContext;
+import com.github.tjake.jlama.safetensors.Config;
 import com.github.tjake.jlama.safetensors.DType;
 import com.github.tjake.jlama.util.Pair;
 import java.io.IOError;
@@ -52,25 +54,22 @@ public AbstractTensor getKvBuffer(UUID session) {
 
     private Pair<RandomAccessFile, AbstractTensor> makeKvBuffer(UUID session) {
         TensorShape s;
+        Config c = model.getConfig();
+        DistributedContext dctx = c.dctx();
         // FIXME: Max size should be configurable
-        int[] rawShape = new int[] { model.getConfig().getNumberOfLayers(), 2, Math.min(1024, model.getConfig().contextLength), model
-            .getConfig().kvLength };
+        int[] rawShape = new int[] { dctx.numberOfLayers, 2, Math.min(1024, c.contextLength), c.kvLength };
 
-        if (model.getConfig().offset().isPresent()) {
-            Pair<Integer, Integer> offset = model.getConfig().offset().get();
-            // Adjust the shape to be relative to the kv cache size (in case of GQA)
-            Pair<Integer, Integer> kvOffset = Pair.create(
-                offset.left / model.getConfig().headGroupSize,
-                offset.right / model.getConfig().headGroupSize
-            );
-            s = TensorShape.sparse(rawShape, kvOffset);
+        // Adjust the shape to be relative to the kv cache size (in case of GQA)
+        if (c.kvLength != dctx.kvSegmentLength) {
+            Pair<Integer, Integer> kvOffset = Pair.of(dctx.kvSegmentStart, dctx.kvSegmentEnd);
+            s = TensorShape.sparseColumn(rawShape, kvOffset);
         } else {
             s = TensorShape.of(rawShape);
         }
 
         // If we don't have a working directory, just use a FloatBufferTensor
         if (model.getConfig().workingDirectory().isEmpty()) {
-            return Pair.create(null, AbstractTensor.make(model.getWorkingDType(), s));
+            return Pair.of(null, AbstractTensor.make(model.getWorkingDType(), s));
         }
 
         // Otherwise, create a file-backed tensor
@@ -101,7 +100,7 @@ private Pair<RandomAccessFile, AbstractTensor> makeKvBuffer(UUID session) {
                 throw new UnsupportedOperationException("Only F32/BF16 is supported for now");
             }
 
-            return Pair.create(raf, t);
+            return Pair.of(raf, t);
 
         } catch (IOException e) {
             throw new IOError(e);
diff --git a/jlama-core/src/main/java/com/github/tjake/jlama/tensor/TensorShape.java b/jlama-core/src/main/java/com/github/tjake/jlama/tensor/TensorShape.java
index 66a9b75..0823219 100644
--- a/jlama-core/src/main/java/com/github/tjake/jlama/tensor/TensorShape.java
+++ b/jlama-core/src/main/java/com/github/tjake/jlama/tensor/TensorShape.java
@@ -31,37 +31,50 @@ public static TensorShape of(int... shape) {
         // Special case for vectors
         if (shape.length == 1) shape = new int[] { 1, shape[0] };
 
-        return new TensorShape(shape, Optional.empty());
+        return new TensorShape(shape, Optional.empty(), Optional.empty());
     }
 
-    public static TensorShape sparse(int[] shape, Pair<Integer, Integer> sparseOffset) {
-        return new TensorShape(shape, Optional.of(sparseOffset));
+    public static TensorShape sparseColumn(int[] shape, Pair<Integer, Integer> sparseOffset) {
+        return new TensorShape(shape, Optional.empty(), Optional.of(sparseOffset));
+    }
+
+    public static TensorShape sparseRow(int[] shape, Pair<Integer, Integer> sparseOffset) {
+        return new TensorShape(shape, Optional.of(sparseOffset), Optional.empty());
     }
 
     private final int[] tshape;
     private final long capacity;
-    private final Optional<Pair<Integer, Integer>> sparseRange;
+    private final Optional<Pair<Integer, Integer>> sparseColumnRange;
+    private final Optional<Pair<Integer, Integer>> sparseRowRange;
     private final boolean isSparse;
-    private final int sparseOffset;
-    private final int sparseLength;
+    private final int sparseColumnOffset;
+    private final int sparseColumnLength;
+    private final int sparseRowOffset;
+    private final int sparseRowLength;
 
-    private TensorShape(int[] shape, Optional<Pair<Integer, Integer>> sparseRange) {
+    private TensorShape(int[] shape, Optional<Pair<Integer,Integer>> sparseRowRange, Optional<Pair<Integer, Integer>> sparseColumnRange) {
         Preconditions.checkArgument(
             shape.length > 1,
             "Shape must have at least two dimensions, even if first is 1 (to represent a vector)"
         );
 
         this.tshape = shape;
-        this.sparseRange = sparseRange;
-        this.isSparse = sparseRange.isPresent();
-        this.sparseOffset = sparseRange.map(Pair::left).orElse(0);
-        this.sparseLength = sparseRange.map(Pair::right).orElse(shape[shape.length - 1]);
+        this.sparseColumnRange = sparseColumnRange;
+        this.sparseRowRange = sparseRowRange;
+        this.isSparse = this.sparseColumnRange.isPresent() || this.sparseRowRange.isPresent();
+
+        this.sparseColumnOffset = this.sparseColumnRange.map(Pair::left).orElse(0);
+        this.sparseColumnLength = this.sparseColumnRange.map(Pair::right).orElse(shape[shape.length - 1]);
+
+        this.sparseRowOffset = this.sparseRowRange.map(Pair::left).orElse(0);
+        this.sparseRowLength = this.sparseRowRange.map(Pair::right).orElse(shape[shape.length - 2]);
 
         long c = 1;
-        for (int i = 0; i < shape.length - 1; i++)
+        for (int i = 0; i < shape.length - 2; i++)
             c *= shape[i];
 
-        c *= sparseLength;
+        c *= sparseRowLength;
+        c *= sparseColumnLength;
         this.capacity = c;
     }
 
@@ -82,15 +95,15 @@ public final int getOffset(int... pdims) {
         // Preconditions.checkArgument(pdims.length == dims(), "Method requires all dimensions specified");
         switch (pdims.length) {
             case 1:
-                return sparseLength * pdims[0] - sparseOffset;
+                return sparseColumnLength * (pdims[0] - sparseRowOffset) - sparseColumnOffset;
             case 2:
-                return sparseLength * pdims[0] + pdims[1] - sparseOffset; // Most common case
+                return sparseColumnLength * (pdims[0] - sparseRowOffset) + pdims[1] - sparseColumnOffset; // Most common case
             case 3:
-                return (sparseLength * tshape[1] * pdims[0]) + (sparseLength * pdims[1]) + pdims[2] - sparseOffset;
+                return (sparseColumnLength * tshape[1] * (pdims[0] - sparseRowOffset)) + (sparseColumnLength * pdims[1]) + pdims[2] - sparseColumnOffset;
             default:
                 int totalOffset = 0;
                 for (int d = 0; d < pdims.length - 1; ++d) { // Stop before last dimension
-                    int offset = sparseLength;
+                    int offset = sparseColumnLength;
                     for (int i = tshape.length - 2; i > d; --i) { // factor scaling of each dim shape
                         offset *= tshape[i];
                     }
@@ -98,27 +111,31 @@ public final int getOffset(int... pdims) {
                     totalOffset += pdims[d] * offset;
                 }
 
-                return totalOffset + pdims[pdims.length - 1] - sparseOffset;
+                return totalOffset + pdims[pdims.length - 1] - sparseColumnOffset;
         }
     }
 
-    public int sparseLength() {
-        return sparseLength;
+    public int sparseColumnLength() {
+        return sparseColumnLength;
     }
 
-    public int sparseOffset() {
-        return sparseOffset;
+    public int sparseColumnOffset() {
+        return sparseColumnOffset;
     }
 
-    public int sparseAdjustment(int offset) {
-        // Preconditions.checkArgument(sparseOffset <= offset, "Offset is outside of sparse range");
-        return offset - sparseOffset;
+    public int sparseRowLength() {
+        return sparseRowLength;
     }
 
+    public int sparseRowOffset() {
+        return sparseRowOffset;
+    }
+
+
     public TensorShape scaleLastDim(float scale) {
         int[] copy = Arrays.copyOf(tshape, tshape.length);
         copy[copy.length - 1] *= scale;
-        return isSparse ? sparse(copy, Pair.create((int) (sparseOffset * scale), (int) (sparseLength * scale))) : of(copy);
+        return sparseColumnRange.isPresent() ? sparseColumn(copy, Pair.of((int) (sparseColumnOffset * scale), (int) (sparseColumnLength * scale))) : of(copy);
     }
 
     public TensorShape setDimValue(int dim, int value) {
@@ -126,7 +143,7 @@ public TensorShape setDimValue(int dim, int value) {
         int[] copy = Arrays.copyOf(tshape, tshape.length);
         copy[dim] = value;
         int newSparseLength = copy[copy.length - 1];
-        return isSparse ? sparse(copy, Pair.create(sparseOffset, newSparseLength)) : of(copy);
+        return sparseColumnRange.isPresent() ? sparseColumn(copy, Pair.of(sparseColumnOffset, newSparseLength)) : of(copy);
     }
 
     public int first() {
@@ -141,17 +158,17 @@ public long size() {
         return capacity;
     }
 
-    public TensorShape sparsify(int offset, int length) {
+    public TensorShape sparsifyColumns(int offset, int length) {
         Preconditions.checkArgument(!isSparse, "Cannot sparsify a sparse tensor");
-        return new TensorShape(tshape, Optional.of(Pair.create(offset, length)));
+        return new TensorShape(tshape, Optional.empty(), Optional.of(Pair.of(offset, length)));
     }
 
     public TensorShape slice(int numDims) {
         Preconditions.checkArgument(numDims < tshape.length, "Too many dimensions specified for tensor");
         int newLength = tshape.length - numDims;
-        if (newLength == 1) return new TensorShape(new int[] { 1, tshape[tshape.length - 1] }, sparseRange);
+        if (newLength == 1) return new TensorShape(new int[] { 1, tshape[tshape.length - 1] }, sparseRowRange, sparseColumnRange);
 
-        return new TensorShape(Arrays.copyOfRange(tshape, numDims, tshape.length), sparseRange);
+        return new TensorShape(Arrays.copyOfRange(tshape, numDims, tshape.length), sparseRowRange, sparseColumnRange);
     }
 
     @Override
@@ -159,18 +176,18 @@ public boolean equals(Object o) {
         if (this == o) return true;
         if (o == null || getClass() != o.getClass()) return false;
         TensorShape that = (TensorShape) o;
-        return Arrays.equals(tshape, that.tshape) && Objects.equals(sparseRange, that.sparseRange);
+        return Arrays.equals(tshape, that.tshape) && Objects.equals(sparseColumnRange, that.sparseColumnRange);
     }
 
     @Override
     public int hashCode() {
-        int result = Objects.hash(sparseRange);
+        int result = Objects.hash(sparseColumnRange);
         result = 31 * result + Arrays.hashCode(tshape);
         return result;
     }
 
     @Override
     public String toString() {
-        return "TensorShape{" + "tshape=" + Arrays.toString(tshape) + ", capacity=" + capacity + ", sparseRange=" + sparseRange + '}';
+        return "TensorShape{" + "tshape=" + Arrays.toString(tshape) + ", capacity=" + capacity + ", sparseRange=" + sparseColumnRange + '}';
     }
 }
diff --git a/jlama-core/src/main/java/com/github/tjake/jlama/util/Pair.java b/jlama-core/src/main/java/com/github/tjake/jlama/util/Pair.java
index f8dd871..0f54f1d 100644
--- a/jlama-core/src/main/java/com/github/tjake/jlama/util/Pair.java
+++ b/jlama-core/src/main/java/com/github/tjake/jlama/util/Pair.java
@@ -58,7 +58,7 @@ public T2 right() {
         return right;
     }
 
-    public static <X, Y> Pair<X, Y> create(X x, Y y) {
+    public static <X, Y> Pair<X, Y> of(X x, Y y) {
         return new Pair<X, Y>(x, y);
     }
 }
diff --git a/jlama-native/src/main/java/com/github/tjake/jlama/tensor/operations/NativeTensorOperations.java b/jlama-native/src/main/java/com/github/tjake/jlama/tensor/operations/NativeTensorOperations.java
index e4a394c..e8f5564 100644
--- a/jlama-native/src/main/java/com/github/tjake/jlama/tensor/operations/NativeTensorOperations.java
+++ b/jlama-native/src/main/java/com/github/tjake/jlama/tensor/operations/NativeTensorOperations.java
@@ -96,6 +96,13 @@ public void batchDotProduct(
         int N = rowChunkSize; // b.shape().dim(0);
         int K = columnLength; // a.shape().dim(1);
 
+        int aOffset = at.getOffset(0, aColumnOffset);
+        int bOffset = bt.getOffset(bt.shape().sparseRowOffset(), bColumnOffset);
+
+        int rOffset = result.shape().sparseColumnOffset() - bt.shape().sparseRowOffset();
+
+        int adjBRowOffset = bRowOffset - bt.shape().sparseRowOffset();
+
         switch (at.dType()) {
             case BF16:
                 switch (bt.dType()) {
@@ -103,14 +110,14 @@ public void batchDotProduct(
                         NativeSimd.gemm_bf16(
                             flags,
                             at.getMemorySegment(),
-                            at.getOffset(0, aColumnOffset),
+                            aOffset,
                             bt.getMemorySegment(),
-                            bt.getOffset(0, bColumnOffset),
+                            bOffset,
                             result.dType() == DType.BF16 ? result.getMemorySegment() : MemorySegment.NULL,
                             result.dType() == DType.F32 ? result.getMemorySegment() : MemorySegment.NULL,
-                            result.shape().sparseOffset(),
+                            rOffset,
                             M,
-                            bRowOffset,
+                            adjBRowOffset,
                             N,
                             K,
                             at.getStride(),
@@ -128,13 +135,13 @@ public void batchDotProduct(
                         NativeSimd.gemm_f32(
                             flags,
                             at.getMemorySegment(),
-                            at.getOffset(0, aColumnOffset),
+                            aOffset,
                             bt.getMemorySegment(),
-                            bt.getOffset(0, bColumnOffset),
+                            bOffset,
                             result.getMemorySegment(),
-                            result.shape().sparseOffset(),
+                            rOffset,
                             M,
-                            bRowOffset,
+                            adjBRowOffset,
                             N,
                             K,
                             at.getStride(),
@@ -146,14 +153,14 @@ public void batchDotProduct(
                         NativeSimd.gemm_f32_bf16(
                             flags,
                             at.getMemorySegment(),
-                            at.getOffset(0, aColumnOffset),
+                            aOffset,
                             bt.getMemorySegment(),
-                            bt.getOffset(0, bColumnOffset),
+                            bOffset,
                             result.dType() == DType.BF16 ? result.getMemorySegment() : MemorySegment.NULL,
                             result.dType() == DType.F32 ? result.getMemorySegment() : MemorySegment.NULL,
-                            result.shape().sparseOffset(),
+                            rOffset,
                             M,
-                            bRowOffset,
+                            adjBRowOffset,
                             N,
                             K,
                             at.getStride(),
@@ -170,14 +177,14 @@ public void batchDotProduct(
                                 NativeSimd.gemm_f32_q4(
                                     flags,
                                     at.getMemorySegment(),
-                                    at.getOffset(0, aColumnOffset),
+                                    aOffset,
                                     b.getBlockF().getMemorySegment(),
                                     b.getMemorySegment(),
-                                    b.getMemorySegmentOffset(b.getOffset(0, bColumnOffset)),
+                                    bOffset,
                                     result.getMemorySegment(),
-                                    result.shape().sparseOffset(),
+                                    rOffset,
                                     M,
-                                    bRowOffset,
+                                    adjBRowOffset,
                                     N,
                                     K,
                                     at.getStride(),
@@ -200,14 +207,14 @@ public void batchDotProduct(
                             flags,
                             a.getBlockF().getMemorySegment(),
                             a.getMemorySegment(),
-                            a.getOffset(0, aColumnOffset),
+                            aOffset,
                             b.getBlockF().getMemorySegment(),
                             b.getMemorySegment(),
-                            b.getMemorySegmentOffset(b.getOffset(0, bColumnOffset)),
+                            bOffset,
                             result.getMemorySegment(),
-                            result.shape().sparseOffset(),
+                            rOffset,
                             M,
-                            bRowOffset,
+                            adjBRowOffset,
                             N,
                             K,
                             a.getStride(),
@@ -251,6 +258,14 @@ public void dotProductBatchChunk(
         int N = rowChunkSize; // b.shape().dim(0);
         int K = columnLength; // a.shape().dim(1);
 
+        int aOffset = a.getOffset(0, columnOffset);
+        int bOffset = b[0].getOffset(b[0].shape().sparseRowOffset(), columnOffset);
+
+        int adjBRowOffset = bRowOffset - b[0].shape().sparseRowOffset();
+
+        //Adjusts for both sparse columns and rows
+        int rOffset = r[0].shape().sparseColumnOffset() - b[0].shape().sparseRowOffset();
+
         switch (a.dType()) {
             case BF16:
                 switch (b[0].dType()) {
@@ -259,14 +274,14 @@ public void dotProductBatchChunk(
                             flags,
                             r.length,
                             a.getMemorySegment(),
-                            a.getOffset(0, columnOffset),
+                            aOffset,
                             rb,
-                            b[0].getOffset(0, columnOffset),
+                            bOffset,
                             r[0].dType() == DType.BF16 ? ra : MemorySegment.NULL,
                             r[0].dType() == DType.F32 ? ra : MemorySegment.NULL,
-                            r[0].shape().sparseOffset(),
+                            rOffset,
                             M,
-                            bRowOffset,
+                            adjBRowOffset,
                             N,
                             K,
                             a.getStride(),
@@ -285,11 +300,11 @@ public void dotProductBatchChunk(
                             flags,
                             r.length,
                             a.getMemorySegment(),
-                            a.getOffset(0, columnOffset),
+                            aOffset,
                             rb,
-                            b[0].getOffset(0, columnOffset),
+                            bOffset,
                             ra,
-                            r[0].shape().sparseOffset(),
+                            rOffset,
                             M,
                             bRowOffset,
                             N,
@@ -304,14 +319,14 @@ public void dotProductBatchChunk(
                             flags,
                             r.length,
                             a.getMemorySegment(),
-                            a.getOffset(0, columnOffset),
+                            aOffset,
                             rb,
-                            b[0].getOffset(0, columnOffset),
+                            bOffset,
                             r[0].dType() == DType.BF16 ? ra : MemorySegment.NULL,
                             r[0].dType() == DType.F32 ? ra : MemorySegment.NULL,
-                            r[0].shape().sparseOffset(),
+                            rOffset,
                             M,
-                            bRowOffset,
+                            adjBRowOffset,
                             N,
                             K,
                             a.getStride(),
@@ -329,12 +344,12 @@ public void dotProductBatchChunk(
                                     flags,
                                     r.length,
                                     a.getMemorySegment(),
-                                    a.getOffset(0, columnOffset),
+                                    aOffset,
                                     rc,
                                     rb,
-                                    b[0].getMemorySegmentOffset(b[0].getOffset(0, columnOffset)),
+                                    bOffset,
                                     ra,
-                                    r[0].shape().sparseOffset(),
+                                    rOffset,
                                     M,
                                     bRowOffset,
                                     N,
@@ -360,14 +375,14 @@ public void dotProductBatchChunk(
                             r.length,
                             at.getBlockF().getMemorySegment(),
                             a.getMemorySegment(),
-                            a.getOffset(0, columnOffset),
+                            aOffset,
                             rc,
                             rb,
-                            bt.getMemorySegmentOffset(bt.getOffset(0, columnOffset)),
+                            bOffset,
                             ra,
-                            r[0].shape().sparseOffset(),
+                            rOffset,
                             M,
-                            bRowOffset,
+                            adjBRowOffset,
                             N,
                             K,
                             a.getStride(),
diff --git a/jlama-net/src/main/java/com/github/tjake/jlama/net/Coordinator.java b/jlama-net/src/main/java/com/github/tjake/jlama/net/Coordinator.java
index 3feba3b..17f5ad5 100644
--- a/jlama-net/src/main/java/com/github/tjake/jlama/net/Coordinator.java
+++ b/jlama-net/src/main/java/com/github/tjake/jlama/net/Coordinator.java
@@ -22,6 +22,7 @@
 import com.github.tjake.jlama.net.grpc.JlamaService;
 import com.github.tjake.jlama.safetensors.DType;
 import com.github.tjake.jlama.safetensors.prompt.PromptContext;
+import com.github.tjake.jlama.safetensors.prompt.PromptSupport;
 import com.github.tjake.jlama.safetensors.tokenizer.Tokenizer;
 import com.github.tjake.jlama.tensor.AbstractTensor;
 import com.google.common.base.Preconditions;
@@ -67,6 +68,10 @@ public Tokenizer getTokenizer() {
         return model.getTokenizer();
     }
 
+    public Optional<PromptSupport> promptSupport() {
+        return model.promptSupport();
+    }
+
     public void start() throws IOException {
         server.start();
         logger.info("Server started, listening on " + port);
@@ -107,7 +112,7 @@ public Generator.Response generate(
             long[] encoded = model.getTokenizer().encode(promptContext.getPrompt());
             Preconditions.checkArgument(encoded.length < model.getConfig().contextLength);
 
-            AbstractTensor logits = model.makeTensor(model.getConfig().vocabularySize);
+            AbstractTensor logits = model.makeDenseTensor(model.getConfig().vocabularySize);
 
             int[] promptTokens = new int[1 + encoded.length];
 
diff --git a/jlama-net/src/main/java/com/github/tjake/jlama/net/Worker.java b/jlama-net/src/main/java/com/github/tjake/jlama/net/Worker.java
index 283cf7e..733ce7d 100644
--- a/jlama-net/src/main/java/com/github/tjake/jlama/net/Worker.java
+++ b/jlama-net/src/main/java/com/github/tjake/jlama/net/Worker.java
@@ -18,6 +18,7 @@
 import static com.github.tjake.jlama.model.ModelSupport.loadModel;
 
 import com.github.tjake.jlama.model.AbstractModel;
+import com.github.tjake.jlama.model.DistributedContext;
 import com.github.tjake.jlama.safetensors.DType;
 import com.github.tjake.jlama.tensor.AbstractTensor;
 import com.github.tjake.jlama.tensor.KvBufferCache;
@@ -73,11 +74,13 @@ public Worker(
         );
         this.registerResponse = blockingClient.register(RegisterRequest.newBuilder().setWorkerid(workerIdBytes).build());
         logger.info(
-            "Registered worker {} with offset {} and length {}",
+            "Registered worker {} with shard {} of {}",
             workerId,
-            registerResponse.getOffset(),
-            registerResponse.getLength()
+            registerResponse.getModelShard(),
+            registerResponse.getNumModelShards()
         );
+
+
         this.model = loadModel(
             AbstractModel.InferenceType.FORWARD_PASS,
             modelPrefix,
@@ -86,7 +89,13 @@ public Worker(
             workingQuantizationType,
             modelQuantization,
             Optional.empty(),
-            Optional.of(Pair.create(registerResponse.getOffset(), registerResponse.getLength()))
+            Optional.of(c -> DistributedContext.builder(c)
+                        .setModelShard(registerResponse.getModelShard())
+                        .setNumModelShards(registerResponse.getNumModelShards())
+                        .setLayerShard(registerResponse.getLayerShard())
+                        .setNumLayerShards(registerResponse.getNumLayerShards())
+                        .build()
+            )
         );
     }
 
@@ -180,16 +189,7 @@ public void onNext(GenerateResponse generateResponse) {
             logger.info("Processing token {} at position {} for session {}", token, position, session);
 
             AbstractTensor output = model.forward(token, position, kvBufferCache.getKvBuffer(session), Optional.of((a, b) -> {
-                CombineRequest nr = CombineRequest.newBuilder()
-                    .setUuid(generateResponse.getSession())
-                    .setWorkerid(workerIdBytes)
-                    .setLayer(getNextRequestCount(session))
-                    .setSumSq(a)
-                    .setSum(b)
-                    .build();
-
-                CombineResponse combineResponse = getCombineResponseStream(session).request(nr).join();
-                return Pair.create(combineResponse.getSumSq(), combineResponse.getSum());
+               return null;
             }), Optional.of(t -> {
                 CombineRequest.Builder nrb = CombineRequest.newBuilder()
                     .setUuid(generateResponse.getSession())
diff --git a/jlama-net/src/main/java/com/github/tjake/jlama/net/grpc/JlamaService.java b/jlama-net/src/main/java/com/github/tjake/jlama/net/grpc/JlamaService.java
index 242a76e..cee9071 100644
--- a/jlama-net/src/main/java/com/github/tjake/jlama/net/grpc/JlamaService.java
+++ b/jlama-net/src/main/java/com/github/tjake/jlama/net/grpc/JlamaService.java
@@ -30,10 +30,8 @@
 import java.nio.ByteOrder;
 import java.nio.FloatBuffer;
 import java.util.*;
-import java.util.concurrent.ConcurrentHashMap;
-import java.util.concurrent.ConcurrentMap;
-import java.util.concurrent.CountDownLatch;
-import java.util.concurrent.TimeUnit;
+import java.util.concurrent.*;
+
 import jdk.incubator.vector.FloatVector;
 import org.jctools.queues.MpmcArrayQueue;
 import org.slf4j.Logger;
@@ -48,9 +46,7 @@ public class JlamaService extends JlamaServiceGrpc.JlamaServiceImplBase {
     private final GeneratorGroup generatorGroup;
 
     private final ConcurrentMap<String, MpmcArrayQueue<Pair<CombineRequest, StreamObserver<CombineResponse>>>> combinations;
-    private final int headSize;
     private final int headCount;
-    private final int headsPerWorker;
 
     public JlamaService(AbstractModel model, int workerCount) {
         Preconditions.checkArgument(
@@ -63,9 +59,7 @@ public JlamaService(AbstractModel model, int workerCount) {
         this.combinations = new ConcurrentHashMap<>();
         this.generatorGroup = new GeneratorGroup();
 
-        this.headCount = model.getConfig().numberOfHeads;
-        this.headSize = model.getConfig().headSize;
-        this.headsPerWorker = headCount / workerCount;
+        this.headCount = model.getConfig().numberOfKeyValueHeads;
     }
 
     public void waitForReady() {
@@ -107,12 +101,17 @@ public void register(RegisterRequest request, StreamObserver<RegisterResponse> r
                     return;
                 }
 
-                int offset = workers.size() * headsPerWorker * headSize;
-                int length = headsPerWorker * headSize;
+                int workerNum = workers.size();
+
+                RegisterResponse r = RegisterResponse.newBuilder()
+                        .setModelShard(workerNum)
+                        .setNumModelShards(workerCount)
+                        .setLayerShard(0)
+                        .setNumLayerShards(1)
+                        .build();
 
-                RegisterResponse r = RegisterResponse.newBuilder().setOffset(offset).setLength(length).build();
                 workers.put(wid, r);
-                logger.info("Registered worker {} with offset {} and length {}", wid, offset, length);
+                logger.info("Registered worker {} with workerNum {} of {}", wid, workerNum, workerCount);
 
                 responseObserver.onNext(r);
                 responseObserver.onCompleted();
@@ -143,7 +142,7 @@ public void onNext(CombineRequest request) {
                     key,
                     k -> new MpmcArrayQueue<>(workerCount + 1)
                 );
-                members.add(Pair.create(request, responseObserver));
+                members.add(Pair.of(request, responseObserver));
 
                 // If we have all the workers, then we can calculate the result and send it back
                 if (members.size() == workerCount && combinations.remove(key, members)) {
@@ -151,8 +150,6 @@ public void onNext(CombineRequest request) {
                     float sum = 0;
                     MemorySegment[] tensors = null;
                     for (Pair<CombineRequest, StreamObserver<CombineResponse>> f : members) {
-                        sumSq += f.left.getSumSq();
-                        sum += f.left.getSum();
                         if (f.left.getTensorCount() > 0) {
                             if (tensors == null) {
                                 tensors = new MemorySegment[f.left.getTensorCount()];
@@ -172,7 +169,7 @@ public void onNext(CombineRequest request) {
                         }
                     }
 
-                    CombineResponse.Builder responseBuilder = CombineResponse.newBuilder().setSumSq(sumSq).setSum(sum);
+                    CombineResponse.Builder responseBuilder = CombineResponse.newBuilder();
 
                     if (tensors != null) {
                         for (int i = 0; i < tensors.length; i++)
@@ -243,15 +240,22 @@ public AbstractTensor generateNextOutput(UUID session, int tokenId, int position
                 g.responseObserver.onNext(gr);
             }
 
-            AbstractTensor output = model.makeFullTensor(model.getConfig().embeddingLength);
+            AbstractTensor output = model.makeDenseTensor(model.getConfig().embeddingLength);
 
-            for (Generator g : generators) {
+
+            for (int j = 0; j < workerCount; j++) {
+                Generator g = generators.get(j);
                 ByteString v = g.waitForOutput(session);
                 RegisterResponse r = workers.get(g.workerId);
 
-                FloatBuffer f = v.asReadOnlyByteBuffer().order(ByteOrder.LITTLE_ENDIAN).asFloatBuffer();
-                for (int i = 0; i < r.getLength(); i++) {
-                    output.set(f.get(), 0, r.getOffset() + i);
+                if (j == 0) {
+                    FloatBuffer f = v.asReadOnlyByteBuffer()
+                            .order(ByteOrder.LITTLE_ENDIAN)
+                            .asFloatBuffer();
+                    int len = f.remaining();
+                    for (int i = 0; i < len; i++) {
+                        output.set(f.get(), 0, i);
+                    }
                 }
             }
 
diff --git a/jlama-net/src/main/java/com/github/tjake/jlama/net/openai/OpenAIChatService.java b/jlama-net/src/main/java/com/github/tjake/jlama/net/openai/OpenAIChatService.java
index 5fea9e1..da685ff 100644
--- a/jlama-net/src/main/java/com/github/tjake/jlama/net/openai/OpenAIChatService.java
+++ b/jlama-net/src/main/java/com/github/tjake/jlama/net/openai/OpenAIChatService.java
@@ -26,6 +26,9 @@
 import java.util.UUID;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.atomic.AtomicInteger;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.http.HttpStatus;
 import org.springframework.http.ResponseEntity;
@@ -37,10 +40,11 @@
 @Validated
 public class OpenAIChatService {
 
+    private static final Logger logger = LoggerFactory.getLogger(OpenAIChatService.class);
     private static final String JLAMA_SESSION_HEADER = "X-Jlama-Session";
 
     @Autowired
-    private AbstractModel model;
+    private Generator model;
 
     /**
      * POST /chat/completions : Creates a model response for the given chat conversation.
@@ -132,6 +136,8 @@ Object createChatCompletion(@RequestHeader Map<String, String> headers, @Valid @
                     );
 
                     emitter.complete();
+
+                    logger.info("Completed streaming response {} tok/sec", (r.generatedTokens/1000.0f) / r.generatedTokens);
                 } catch (IOException e) {
                     emitter.completeWithError(e);
                 }
diff --git a/jlama-net/src/main/proto/JlamaService.proto b/jlama-net/src/main/proto/JlamaService.proto
index 02039e7..b1e2284 100644
--- a/jlama-net/src/main/proto/JlamaService.proto
+++ b/jlama-net/src/main/proto/JlamaService.proto
@@ -6,14 +6,18 @@ message CombineRequest {
   bytes workerid = 1;
   bytes uuid = 2;
   int32 layer = 3;
-  float sumSq = 4;
-  float sum = 5;
   repeated bytes tensor = 6;
 }
 
+message TensorData {
+  int32 dim0 = 1; // Rows
+  int32 dim1 = 2; // Columns
+  int32 columnOffset = 3; // Offset into the column (if sparse)
+  int32 columnLength = 4; // Length of the column (if sparse)
+  bytes data = 5;
+}
+
 message CombineResponse {
-  float sumSq = 1;
-  float sum = 2;
   repeated bytes tensor = 3;
 }
 
@@ -29,13 +33,18 @@ message GenerateResponse {
   int32 position = 3;
 }
 
+/**
+ * Register a worker with the server
+ */
 message RegisterRequest {
   bytes workerid = 1;
 }
 
 message RegisterResponse {
-  int32 offset = 1;
-  int32 length = 2;
+  int32 modelShard = 1;
+  int32 numModelShards = 2;
+  int32 layerShard = 3;
+  int32 numLayerShards = 4;
 }
 
 service JlamaService {
diff --git a/jlama-net/src/test/java/com/github/tjake/jlama/net/DistributedServiceTest.java b/jlama-net/src/test/java/com/github/tjake/jlama/net/DistributedServiceTest.java
index 2a739f3..281d030 100644
--- a/jlama-net/src/test/java/com/github/tjake/jlama/net/DistributedServiceTest.java
+++ b/jlama-net/src/test/java/com/github/tjake/jlama/net/DistributedServiceTest.java
@@ -42,10 +42,6 @@ public class DistributedServiceTest {
         rootLogger.setLevel(Level.toLevel("info"));
     }
 
-    private static final ObjectMapper om = new ObjectMapper().configure(DeserializationFeature.FAIL_ON_IGNORED_PROPERTIES, false)
-        .configure(DeserializationFeature.FAIL_ON_TRAILING_TOKENS, false)
-        .configure(DeserializationFeature.FAIL_ON_UNKNOWN_PROPERTIES, false)
-        .configure(DeserializationFeature.FAIL_ON_MISSING_CREATOR_PROPERTIES, true);
 
     @Test
     void oneWorkerTestLLama() throws Exception {
@@ -79,8 +75,8 @@ void oneWorkerTestLLama() throws Exception {
 
     @Test
     void manyWorkerTestLLama() throws Exception {
-        // Path modelRoot = Paths.get("../models/Mixtral-8x7B-Instruct-v0.1-jlama-Q4");
-        Path modelRoot = Paths.get("../models/Llama-2-7b-chat-hf-jlama-Q4");
+        //Path modelRoot = Paths.get("../models/Mixtral-8x7B-Instruct-v0.1-jlama-Q4");
+        Path modelRoot = Paths.get("../models/Meta-Llama-3.1-8B-Instruct-jlama-Q4");
         Assume.assumeTrue(Files.exists(modelRoot));
 
         Coordinator coordinator = new Coordinator(modelRoot.toFile(), null, 8888, 4);
@@ -106,7 +102,7 @@ void manyWorkerTestLLama() throws Exception {
             coordinator.generate(
                 UUID.randomUUID(),
                 PromptContext.of("Simply put, the theory of relativity states that"),
-                0.7f,
+                0.3f,
                 256,
                 makeOutHandler()
             );
diff --git a/jlama-net/src/test/java/com/github/tjake/jlama/net/JlamaServiceTest.java b/jlama-net/src/test/java/com/github/tjake/jlama/net/JlamaServiceTest.java
index 39c0262..42f7389 100644
--- a/jlama-net/src/test/java/com/github/tjake/jlama/net/JlamaServiceTest.java
+++ b/jlama-net/src/test/java/com/github/tjake/jlama/net/JlamaServiceTest.java
@@ -19,6 +19,7 @@
 
 import com.github.tjake.jlama.math.ActivationFunction;
 import com.github.tjake.jlama.model.AbstractModel;
+import com.github.tjake.jlama.model.DistributedContext;
 import com.github.tjake.jlama.model.ModelSupport;
 import com.github.tjake.jlama.model.TransformerBlock;
 import com.github.tjake.jlama.model.functions.EmbedInput;
@@ -83,13 +84,13 @@ public void testRegister() {
             )
             .build();
         RegisterResponse response = blockingStub.register(request);
-        assertThat(response.getOffset()).isEqualTo(0);
-        assertThat(response.getLength()).isEqualTo(1024);
+        assertThat(response.getModelShard()).isEqualTo(0);
+        assertThat(response.getNumModelShards()).isEqualTo(4);
 
         // Should get the same response if we register again with same uuid
         response = blockingStub.register(request);
-        assertThat(response.getOffset()).isEqualTo(0);
-        assertThat(response.getLength()).isEqualTo(1024);
+        assertThat(response.getModelShard()).isEqualTo(0);
+        assertThat(response.getNumModelShards()).isEqualTo(4);
 
         // Should get a different response if we register with a different uuid
         uuid = UUID.randomUUID();
@@ -101,33 +102,8 @@ public void testRegister() {
             )
             .build();
         response = blockingStub.register(request);
-        assertThat(response.getOffset()).isEqualTo(1024);
-        assertThat(response.getLength()).isEqualTo(1024);
-    }
-
-    @Test
-    public void testNorm() {
-        UUID uuid = UUID.randomUUID();
-        CombineRequest request = CombineRequest.newBuilder()
-            .setUuid(
-                ByteString.copyFrom(
-                    ByteBuffer.allocate(128).putLong(uuid.getLeastSignificantBits()).putLong(uuid.getMostSignificantBits()).flip()
-                )
-            )
-            .setLayer(0)
-            .setSumSq(10)
-            .build();
-
-        /*ListenableFuture<NormResponse> response1 = futureStub.norm(request);
-        ListenableFuture<NormResponse> response2 = futureStub.norm(request);
-        ListenableFuture<NormResponse> response3 = futureStub.norm(request);
-        ListenableFuture<NormResponse> response4 = futureStub.norm(request);
-        
-        Futures.whenAllComplete(response1, response2, response3, response4);
-        assertThat(response1.resultNow().getSumSq()).isEqualTo(40);
-        assertThat(response2.resultNow().getSumSq()).isEqualTo(40);
-        assertThat(response3.resultNow().getSumSq()).isEqualTo(40);
-        assertThat(response4.resultNow().getSumSq()).isEqualTo(40);*/
+        assertThat(response.getModelShard()).isEqualTo(1);
+        assertThat(response.getNumModelShards()).isEqualTo(4);
     }
 
     public static class MockConfig extends Config {
@@ -169,7 +145,7 @@ public Map<String, TensorInfo> tensorInfoMap() {
         }
 
         @Override
-        public AbstractTensor load(String name, Optional<Pair<Integer, Integer>> offset) {
+        public AbstractTensor load(String name, DistributedContext dctx, boolean sparseRows, boolean sparseColumns) {
             return null;
         }
 
diff --git a/jlama-tests/src/test/java/com/github/tjake/jlama/model/Mocks.java b/jlama-tests/src/test/java/com/github/tjake/jlama/model/Mocks.java
index b1c075a..431b0c4 100644
--- a/jlama-tests/src/test/java/com/github/tjake/jlama/model/Mocks.java
+++ b/jlama-tests/src/test/java/com/github/tjake/jlama/model/Mocks.java
@@ -77,7 +77,7 @@ public Map<String, TensorInfo> tensorInfoMap() {
         }
 
         @Override
-        public AbstractTensor load(String name, Optional<Pair<Integer, Integer>> offset) {
+        public AbstractTensor load(String name, DistributedContext dctx, boolean sparseRows, boolean sparseColumns) {
             return null;
         }
 
diff --git a/jlama-tests/src/test/java/com/github/tjake/jlama/model/TestModels.java b/jlama-tests/src/test/java/com/github/tjake/jlama/model/TestModels.java
index a9c57cb..bb39f59 100644
--- a/jlama-tests/src/test/java/com/github/tjake/jlama/model/TestModels.java
+++ b/jlama-tests/src/test/java/com/github/tjake/jlama/model/TestModels.java
@@ -354,161 +354,4 @@ private BiConsumer<String, Float> makeOutHandler() {
 
         return outCallback;
     }
-
-    @Test
-    public void testDistibuted() {
-        Path model = Paths.get("../models/gpt2-medium");
-        Assume.assumeTrue(Files.exists(model));
-
-        AbstractModel mFull = ModelSupport.loadModel(model.toFile(), null, DType.F32, DType.F32, Optional.empty(), Optional.empty());
-
-        int len = mFull.getConfig().embeddingLength;
-        AbstractModel mFirstHalf = ModelSupport.loadModel(
-            AbstractModel.InferenceType.FULL_GENERATION,
-            model.toFile(),
-            null,
-            DType.F32,
-            DType.F32,
-            Optional.empty(),
-            Optional.empty(),
-            Optional.of(Pair.create(0, len / 2))
-        );
-        AbstractModel mSecondHalf = ModelSupport.loadModel(
-            AbstractModel.InferenceType.FULL_GENERATION,
-            model.toFile(),
-            null,
-            DType.F32,
-            DType.F32,
-            Optional.empty(),
-            Optional.empty(),
-            Optional.of(Pair.create(len / 2, len / 2))
-        );
-
-        int embLen = mFull.c.embeddingLength;
-
-        // Main one
-        AbstractTensor kvmem0 = new KvBufferCache(mFull).getKvBuffer(UUID.randomUUID());
-        AbstractTensor t0 = mFull.embedInput.inputTokenToEmbedding(mFull.c.bosToken, 0);
-        AbstractTensor f0 = mFull.transformerBlocks[0].preAttentionNorm.get().forward(t0);
-
-        AbstractTensor q0 = mFull.makeTensor(embLen);
-        VectorMath.pchunk(0, embLen, (chunkStart, chunkLength) -> {
-            TensorOperationsProvider.get()
-                .dotProductChunk(q0, t0, mFull.transformerBlocks[0].attention.queryAttnWeights, 0, embLen, chunkStart, chunkLength);
-        });
-
-        // Test a double dot
-        AbstractTensor qq0 = mFull.makeTensor(embLen);
-        VectorMath.pchunk(0, embLen, (chunkStart, chunkLength) -> {
-            TensorOperationsProvider.get()
-                .dotProductChunk(qq0, q0, mFull.transformerBlocks[0].attention.queryAttnWeights, 0, embLen, chunkStart, chunkLength);
-        });
-
-        AbstractTensor a0 = mFull.transformerBlocks[0].attention.forward(f0, 0, kvmem0.slice(0), Optional.empty());
-
-        // Two halves
-        AtomicDouble sum0 = new AtomicDouble(0);
-        CountDownLatch norm0Latch = new CountDownLatch(2);
-
-        BiFunction<Float, Float, Pair<Float, Float>> reducer = (a, b) -> {
-            sum0.addAndGet(a);
-            norm0Latch.countDown();
-            Uninterruptibles.awaitUninterruptibly(norm0Latch);
-
-            return Pair.create(sum0.floatValue(), 0f);
-        };
-
-        MpmcArrayQueue<AbstractTensor> sumt = new MpmcArrayQueue<>(2);
-        CountDownLatch tensor0Latch = new CountDownLatch(2);
-
-        BiFunction<AbstractTensor, AbstractTensor, Void> treducer = (a, b) -> {
-            sumt.offer(a);
-            tensor0Latch.countDown();
-            Uninterruptibles.awaitUninterruptibly(tensor0Latch);
-
-            AbstractTensor result = a.copyShape();
-            for (AbstractTensor t : sumt)
-                TensorOperationsProvider.get().accumulate(result, t, 0, Ints.checkedCast(result.size()));
-
-            b.copyFrom(result, 0, 0, Ints.checkedCast(result.size()));
-            return null;
-        };
-
-        AbstractTensor kvmem1 = mFirstHalf.makeTensor(mFull.c.getNumberOfLayers(), 10, 2, mFull.c.embeddingLength); // k and v are last 2
-                                                                                                                    // dims
-        AbstractTensor t1 = mFirstHalf.embedInput.inputTokenToEmbedding(mFull.c.bosToken, 0);
-
-        AbstractTensor kvmem2 = mSecondHalf.makeTensor(mFull.c.getNumberOfLayers(), 10, 2, mFull.c.embeddingLength); // k and v are last 2
-                                                                                                                     // dims
-        AbstractTensor t2 = mSecondHalf.embedInput.inputTokenToEmbedding(mFull.c.bosToken, 0);
-
-        CompletableFuture<AbstractTensor> f1c = CompletableFuture.supplyAsync(
-            () -> mFirstHalf.transformerBlocks[0].preAttentionNorm.get().forward(t1, Optional.of(reducer))
-        );
-        CompletableFuture<AbstractTensor> f2c = CompletableFuture.supplyAsync(
-            () -> mSecondHalf.transformerBlocks[0].preAttentionNorm.get().forward(t2, Optional.of(reducer))
-        );
-
-        CompletableFuture.allOf(f1c, f2c);
-        AbstractTensor f1 = f1c.join();
-        AbstractTensor f2 = f2c.join();
-
-        AbstractTensor q1 = mFull.makeTensor(embLen);
-        AbstractTensor qq1 = mFull.makeTensor(embLen);
-
-        int off0 = mFirstHalf.c.embeddingSegmentStart();
-        int len0 = mFirstHalf.c.embeddingSegmentLength();
-
-        VectorMath.pchunk(0, embLen, (chunkStart, chunkLength) -> {
-            TensorOperationsProvider.get()
-                .dotProductChunk(q1, t1, mFirstHalf.transformerBlocks[0].attention.queryAttnWeights, off0, len0, chunkStart, chunkLength);
-        });
-
-        VectorMath.pchunk(0, embLen, (chunkStart, chunkLength) -> {
-            TensorOperationsProvider.get()
-                .dotProductChunk(qq1, q1, mFirstHalf.transformerBlocks[0].attention.queryAttnWeights, off0, len0, chunkStart, chunkLength);
-        });
-
-        AbstractTensor q2 = mFull.makeTensor(embLen);
-        AbstractTensor qq2 = mFull.makeTensor(embLen);
-
-        int off1 = mSecondHalf.c.embeddingSegmentStart();
-        int len1 = mSecondHalf.c.embeddingSegmentLength();
-
-        VectorMath.pchunk(0, embLen, (chunkStart, chunkLength) -> {
-            TensorOperationsProvider.get()
-                .dotProductChunk(q2, t2, mSecondHalf.transformerBlocks[0].attention.queryAttnWeights, off1, len1, chunkStart, chunkLength);
-        });
-        VectorMath.pchunk(0, embLen, (chunkStart, chunkLength) -> {
-            TensorOperationsProvider.get()
-                .dotProductChunk(qq2, q2, mSecondHalf.transformerBlocks[0].attention.queryAttnWeights, off1, len1, chunkStart, chunkLength);
-        });
-
-        AbstractTensor a1 = mFirstHalf.transformerBlocks[0].attention.forward(f1, 0, kvmem1.slice(0), Optional.empty());
-        AbstractTensor a2 = mSecondHalf.transformerBlocks[0].attention.forward(f2, 0, kvmem2.slice(0), Optional.empty());
-
-        AbstractTensor tc = mFull.makeTensor(mFull.c.embeddingLength);
-        tc.copyFrom(t1, 0, 0, (int) t1.size());
-        tc.copyFrom(t2, 0, (int) t2.size(), (int) t2.size());
-        Assert.assertTrue(tensorEquals(tc, t0));
-
-        AbstractTensor fc = mFull.makeTensor(mFull.c.embeddingLength);
-        fc.copyFrom(f1, 0, 0, (int) f1.size());
-        fc.copyFrom(f2, 0, (int) f2.size(), (int) f2.size());
-        Assert.assertTrue(tensorEquals(fc, f0));
-
-        AbstractTensor qc = mFull.makeTensor(mFull.c.embeddingLength);
-        qc.copyFrom(q1, 0, 0, (int) q1.size());
-        TensorOperationsProvider.get().accumulate(qc, q2, 0, (int) qc.size());
-        Assert.assertTrue(tensorEquals(qc, q0));
-    }
-
-    static boolean tensorEquals(AbstractTensor a, AbstractTensor b) {
-        if (a.size() != b.size()) return false;
-
-        for (int i = 0; i < a.size(); i++)
-            Assert.assertEquals("Position " + i, a.get(0, i), b.get(0, i), 0.01f);
-
-        return true;
-    }
 }
diff --git a/pom.xml b/pom.xml
index 6660602..871f8d2 100644
--- a/pom.xml
+++ b/pom.xml
@@ -42,7 +42,7 @@
   <!-- Build property abstractions: versions, etc -->
   <properties>
     <project.build.sourceEncoding>UTF-8</project.build.sourceEncoding>
-    <revision>0.3.1</revision>
+    <revision>0.4.0</revision>
 
     <slf4j-api.version>2.0.7</slf4j-api.version>
     <logback.version>1.5.6</logback.version>
@@ -57,7 +57,6 @@
   </properties>
 
   <dependencies>
-
     <dependency>
       <groupId>org.slf4j</groupId>
       <artifactId>slf4j-api</artifactId>