feat: sequence-parallel Ring Attention for MLX

README.md

@@ -538,6 +538,13 @@ The `exo-bench` tool measures model prefill and token generation speed across di
 - Nodes should be running with `uv run exo` before benchmarking
 - The tool uses the `/bench/chat/completions` endpoint
 
+For Ring Attention on Metal, start every node with MLX's low-latency CPU/GPU
+synchronization enabled:
+
+```bash
+MLX_METAL_FAST_SYNCH=1 uv run exo
+```
+
 **Basic usage:**
 
 ```bash
@@ -554,7 +561,7 @@ uv run bench/exo_bench.py \
 - `--tg`: Generation lengths (comma-separated integers)
 - `--max-nodes`: Limit placements to N nodes (default: 4)
 - `--instance-meta`: Filter by `ring`, `jaccl`, or `both` (default: both)
-- `--sharding`: Filter by `pipeline`, `tensor`, or `both` (default: both)
+- `--sharding`: Filter by `pipeline`, `tensor`, `ring`, or `both` (default: both; `both` retains the pipeline/tensor comparison)
 - `--repeat`: Number of repetitions per configuration (default: 1)
 - `--warmup`: Warmup runs per placement (default: 0)
 - `--json-out`: Output file for results (default: bench/results.json)
@@ -572,6 +579,29 @@ uv run bench/exo_bench.py \
   --json-out my-results.json
 ```
 
+To compare Ring Attention against replicated pipeline prefill on the same
+Ring-compatible model and node count, run both placements with identical prompt
+lengths, generation lengths, repetitions, and warmups:
+
+```bash
+uv run bench/exo_bench.py \
+  --model Llama-3.2-1B-Instruct-4bit \
+  --pp 4096,16384,65536 --tg 128,128,128 \
+  --min-nodes 2 --max-nodes 2 --instance-meta ring \
+  --sharding ring --warmup 1 --repeat 3 \
+  --json-out bench/ring-attention.json
+
+uv run bench/exo_bench.py \
+  --model Llama-3.2-1B-Instruct-4bit \
+  --pp 4096,16384,65536 --tg 128,128,128 \
+  --min-nodes 2 --max-nodes 2 --instance-meta ring \
+  --sharding pipeline --warmup 1 --repeat 3 \
+  --json-out bench/replicated-pipeline.json
+```
+
+Use `prompt_tps` as the primary Ring prefill metric. Keep prefix caching disabled
+(the default), and record node hardware and network topology alongside results.
+
 The tool outputs performance metrics including prompt tokens per second (prompt_tps), generation tokens per second (generation_tps), and peak memory usage for each configuration.
 
 ---

dashboard/src/lib/components/ModelCard.svelte

@@ -22,7 +22,7 @@
       }>;
     } | null;
     nodes?: Record<string, NodeInfo>;
-    sharding?: "Pipeline" | "Tensor";
+    sharding?: "Pipeline" | "Tensor" | "Ring";
     runtime?: "MlxRing" | "MlxJaccl";
     onLaunch?: () => void;
     tags?: string[];

dashboard/src/lib/stores/app.svelte.ts

@@ -173,7 +173,7 @@ export interface ModelDownloadStatus {
 // Placement preview from the API
 export interface PlacementPreview {
   model_id: string;
-  sharding: "Pipeline" | "Tensor";
+  sharding: "Pipeline" | "Tensor" | "Ring";
   instance_meta: "MlxRing" | "MlxJaccl";
   instance: unknown | null;
   memory_delta_by_node: Record<string, number> | null;

dashboard/src/routes/+page.svelte

@@ -783,6 +783,7 @@
       quantization?: string;
       base_model?: string;
       capabilities?: string[];
+      supports_ring?: boolean;
     }>
   >([]);
   type ModelMemoryFitStatus =
@@ -885,14 +886,14 @@
     sendMessage(content, files, thinkingEnabled());
   }
 
-  let selectedSharding = $state<"Pipeline" | "Tensor">("Pipeline");
+  let selectedSharding = $state<"Pipeline" | "Tensor" | "Ring">("Pipeline");
   type InstanceMeta = "MlxRing" | "MlxJaccl";
 
   // Launch defaults persistence
   const LAUNCH_DEFAULTS_KEY = "exo-launch-defaults-v2";
   interface LaunchDefaults {
     modelId: string | null;
-    sharding: "Pipeline" | "Tensor";
+    sharding: "Pipeline" | "Tensor" | "Ring";
     instanceType: InstanceMeta;
     minNodes: number;
   }
@@ -932,7 +933,9 @@
     // Apply sharding and instance type unconditionally
     selectedSharding = defaults.sharding;
     selectedInstanceType =
-      defaults.instanceType === "MlxRing" ? "MlxRing" : "MlxJaccl";
+      defaults.sharding === "Ring" || defaults.instanceType === "MlxRing"
+        ? "MlxRing"
+        : "MlxJaccl";
 
     // Apply minNodes if valid (between 1 and maxNodes)
     if (
@@ -954,6 +957,20 @@
   }
 
   let selectedInstanceType = $state<InstanceMeta>("MlxRing");
+
+  const selectedModelSupportsRing = $derived(
+    models.find(
+      (model) =>
+        model.id === selectedModelId ||
+        model.hugging_face_id === selectedModelId,
+    )?.supports_ring === true,
+  );
+
+  $effect(() => {
+    if (selectedSharding === "Ring" && !selectedModelSupportsRing) {
+      selectedSharding = "Pipeline";
+    }
+  });
   let selectedMinNodes = $state<number>(1);
   let minNodesInitialized = $state(false);
   let launchingModelId = $state<string | null>(null);
@@ -2043,7 +2060,7 @@
     return inst.shardAssignments?.modelId || "Unknown Model";
   }
 
-  // Get instance details: type (MLX Ring/IBV), sharding (Pipeline/Tensor), and node names
+  // Get instance details: type (MLX Ring/IBV), sharding strategy, and node names
   function getInstanceInfo(instanceWrapped: unknown): {
     instanceType: string;
     sharding: string;
@@ -2082,6 +2099,7 @@
       const [shardTag] = getTagged(firstShardWrapped);
       if (shardTag === "PipelineShardMetadata") sharding = "Pipeline";
       else if (shardTag === "TensorShardMetadata") sharding = "Tensor";
+      else if (shardTag === "RingShardMetadata") sharding = "Ring";
       else if (shardTag === "PrefillDecodeShardMetadata")
         sharding = "Prefill/Decode";
     }
@@ -2193,7 +2211,7 @@
 
   function getOrderedRunnerNodes(
     instance: Record<string, unknown>,
-    shardType: "Pipeline" | "Tensor",
+    shardType: "Pipeline" | "Tensor" | "Ring",
   ) {
     const runnerToShard =
       (
@@ -5772,6 +5790,37 @@
                         </span>
                         Tensor
                       </button>
+                      <button
+                        disabled={!selectedModelSupportsRing}
+                        title={selectedModelSupportsRing
+                          ? "Sequence-parallel Ring Attention"
+                          : "Ring Attention is not verified for this model"}
+                        onclick={() => {
+                          if (!selectedModelSupportsRing) return;
+                          selectedSharding = "Ring";
+                          selectedInstanceType = "MlxRing";
+                          saveLaunchDefaults();
+                        }}
+                        class="flex items-center gap-2 py-1.5 px-3 text-xs font-mono border rounded transition-all duration-200 cursor-pointer {selectedSharding ===
+                        'Ring'
+                          ? 'bg-transparent text-exo-yellow border-exo-yellow'
+                          : selectedModelSupportsRing
+                            ? 'bg-transparent text-white/70 border-exo-medium-gray/50 hover:border-exo-yellow/50'
+                            : 'bg-transparent text-white/30 border-exo-medium-gray/30 cursor-not-allowed'}"
+                      >
+                        <span
+                          class="w-3 h-3 rounded-full border-2 flex items-center justify-center {selectedSharding ===
+                          'Ring'
+                            ? 'border-exo-yellow'
+                            : 'border-exo-medium-gray'}"
+                        >
+                          {#if selectedSharding === "Ring"}
+                            <span class="w-1.5 h-1.5 rounded-full bg-exo-yellow"
+                            ></span>
+                          {/if}
+                        </span>
+                        Ring Attention
+                      </button>
                     </div>
                   </div>
 
@@ -5807,6 +5856,9 @@
                       <button
                         onclick={() => {
                           selectedInstanceType = "MlxJaccl";
+                          if (selectedSharding === "Ring") {
+                            selectedSharding = "Pipeline";
+                          }
                           saveLaunchDefaults();
                         }}
                         class="flex items-center gap-2 py-1.5 px-3 text-xs font-mono border rounded transition-all duration-200 cursor-pointer {selectedInstanceType ===

resources/inference_model_cards/mlx-community--Llama-3.2-1B-Instruct-4bit.toml

@@ -3,6 +3,7 @@ n_layers = 16
 hidden_size = 2048
 num_key_value_heads = 8
 supports_tensor = true
+supports_ring = true
 tasks = ["TextGeneration"]
 family = "llama"
 quantization = "4bit"

resources/inference_model_cards/mlx-community--Llama-3.2-3B-Instruct-4bit.toml

@@ -3,6 +3,7 @@ n_layers = 28
 hidden_size = 3072
 num_key_value_heads = 8
 supports_tensor = true
+supports_ring = true
 tasks = ["TextGeneration"]
 family = "llama"
 quantization = "4bit"

resources/inference_model_cards/mlx-community--Llama-3.2-3B-Instruct-8bit.toml

@@ -3,6 +3,7 @@ n_layers = 28
 hidden_size = 3072
 num_key_value_heads = 8
 supports_tensor = true
+supports_ring = true
 tasks = ["TextGeneration"]
 family = "llama"
 quantization = "8bit"

resources/inference_model_cards/mlx-community--Llama-3.3-70B-Instruct-4bit.toml

@@ -3,6 +3,7 @@ n_layers = 80
 hidden_size = 8192
 num_key_value_heads = 8
 supports_tensor = true
+supports_ring = true
 tasks = ["TextGeneration"]
 family = "llama"
 quantization = "4bit"

resources/inference_model_cards/mlx-community--Llama-3.3-70B-Instruct-8bit.toml

@@ -3,6 +3,7 @@ n_layers = 80
 hidden_size = 8192
 num_key_value_heads = 8
 supports_tensor = true
+supports_ring = true
 tasks = ["TextGeneration"]
 family = "llama"
 quantization = "8bit"

resources/inference_model_cards/mlx-community--Meta-Llama-3.1-70B-Instruct-4bit.toml

@@ -3,6 +3,7 @@ n_layers = 80
 hidden_size = 8192
 num_key_value_heads = 8
 supports_tensor = true
+supports_ring = true
 tasks = ["TextGeneration"]
 family = "llama"
 quantization = "4bit"

resources/inference_model_cards/mlx-community--Meta-Llama-3.1-8B-Instruct-4bit.toml

@@ -3,6 +3,7 @@ n_layers = 32
 hidden_size = 4096
 num_key_value_heads = 8
 supports_tensor = true
+supports_ring = true
 tasks = ["TextGeneration"]
 family = "llama"
 quantization = "4bit"

resources/inference_model_cards/mlx-community--Meta-Llama-3.1-8B-Instruct-8bit.toml

@@ -3,6 +3,7 @@ n_layers = 32
 hidden_size = 4096
 num_key_value_heads = 8
 supports_tensor = true
+supports_ring = true
 tasks = ["TextGeneration"]
 family = "llama"
 quantization = "8bit"

resources/inference_model_cards/mlx-community--Meta-Llama-3.1-8B-Instruct-bf16.toml

@@ -3,6 +3,7 @@ n_layers = 32
 hidden_size = 4096
 num_key_value_heads = 8
 supports_tensor = true
+supports_ring = true
 tasks = ["TextGeneration"]
 family = "llama"
 quantization = "bf16"

resources/inference_model_cards/mlx-community--llama-3.3-70b-instruct-fp16.toml

@@ -3,6 +3,7 @@ n_layers = 80
 hidden_size = 8192
 num_key_value_heads = 8
 supports_tensor = true
+supports_ring = true
 tasks = ["TextGeneration"]
 family = "llama"
 quantization = "fp16"

rust/exo_rs/tests/test_python.py

@@ -13,7 +13,9 @@
 @pytest.mark.asyncio
 async def test_sleep_on_multiple_items() -> None:
     print("PYTHON: starting handle")
-    h = NetworkingHandle.new(os.urandom(16).hex().lstrip("0"), 52414, 52413)
+    h = NetworkingHandle.new(
+        os.urandom(16).hex().lstrip("0"), "exo-test", 52414, 52413
+    )
     print("PYTHON: handle started")
 
     rt = asyncio.create_task(_await_recv(h))

src/exo/api/main.py

@@ -535,6 +535,11 @@ async def get_placement_previews(
                         )
                     ]
                 )
+        if model_card.supports_ring:
+            instance_combinations.extend(
+                (Sharding.Ring, InstanceMeta.MlxRing, i)
+                for i in range(2, len(list(self.state.topology.list_nodes())) + 1)
+            )
         # TODO: PDD
         # instance_combinations.append((Sharding.PrefillDecodeDisaggregation, InstanceMeta.MlxRing, 1))
 
@@ -1805,6 +1810,7 @@ async def get_models(self, status: str | None = Query(default=None)) -> ModelLis
                     tags=[],
                     storage_size_megabytes=card.storage_size.in_mb,
                     supports_tensor=card.supports_tensor,
+                    supports_ring=card.supports_ring,
                     tasks=[task.value for task in card.tasks],
                     is_custom=card.is_custom,
                     family=card.family,
@@ -1846,6 +1852,7 @@ async def add_custom_model(self, payload: AddCustomModelParams) -> ModelListMode
             tags=[],
             storage_size_megabytes=int(card.storage_size.in_mb),
             supports_tensor=card.supports_tensor,
+            supports_ring=card.supports_ring,
             tasks=[task.value for task in card.tasks],
             is_custom=True,
         )

src/exo/api/types/api.py

@@ -42,6 +42,7 @@ class ModelListModel(BaseModel):
     tags: list[str] = Field(default=[])
     storage_size_megabytes: int = Field(default=0)
     supports_tensor: bool = Field(default=False)
+    supports_ring: bool = Field(default=False)
     tasks: list[str] = Field(default=[])
     is_custom: bool = Field(default=False)
     family: str = Field(default="")

src/exo/download/__init__.py

@@ -0,0 +1 @@
+"""Model download orchestration and storage helpers."""

src/exo/master/placement.py

@@ -4,7 +4,9 @@
 
 from exo.master.placement_utils import (
     Cycle,
+    estimate_ring_node_memory,
     filter_cycles_by_memory,
+    filter_cycles_by_replicated_memory,
     get_mlx_jaccl_coordinators,
     get_mlx_jaccl_devices_matrix,
     get_mlx_ring_hosts_by_node,
@@ -114,6 +116,18 @@ def place_instance(
     download_status: Mapping[NodeId, Sequence[DownloadProgress]] | None = None,
     node_rdma_ctl: Mapping[NodeId, NodeRdmaCtlStatus] | None = None,
 ) -> dict[InstanceId, Instance]:
+    if (
+        command.sharding is Sharding.Ring
+        and command.instance_meta is not InstanceMeta.MlxRing
+    ):
+        raise ValueError("Ring attention requires the MlxRing transport")
+    if command.sharding is Sharding.Ring and command.min_nodes < 2:
+        raise ValueError("Ring attention requires at least two nodes")
+    if command.sharding is Sharding.Ring and not command.model_card.supports_ring:
+        raise ValueError(
+            f"Model does not declare Ring attention support: {command.model_card.model_id}"
+        )
+
     cycles = topology.get_cycles()
     candidate_cycles = list(filter(lambda it: len(it) >= command.min_nodes, cycles))
 
@@ -124,9 +138,18 @@ def place_instance(
             for cycle in candidate_cycles
             if required_nodes.issubset(cycle.node_ids)
         ]
-    cycles_with_sufficient_memory = filter_cycles_by_memory(
-        candidate_cycles, node_memory, command.model_card.storage_size
-    )
+    if command.sharding is Sharding.Ring:
+        # Every ring rank replicates the weights and must also hold the
+        # long-context prefill working set, not just the model file.
+        cycles_with_sufficient_memory = filter_cycles_by_replicated_memory(
+            candidate_cycles,
+            node_memory,
+            estimate_ring_node_memory(command.model_card),
+        )
+    else:
+        cycles_with_sufficient_memory = filter_cycles_by_memory(
+            candidate_cycles, node_memory, command.model_card.storage_size
+        )
     if len(cycles_with_sufficient_memory) == 0:
         raise ValueError("No cycles found with sufficient memory")