docs: Add comparison with state-of-the-art weight transfer strategies

Analyze ModelExpress against Perplexity/Kimi blog posts on 1-2 second
weight transfers:

Key gaps identified:
- No pipelining (4-stage: H2D, GPU ops, RDMA, barrier)
- Individual tensor registration vs CUDACachingAllocator blocks
- Sequential transfers vs overlapped execution

Strategies that don't apply to our use case:
- FSDP/DTensor full_tensor() reconstruction
- On-the-fly projection fusion and quantization
- Training→inference transparency (WRITE vs READ)

Recommended optimizations prioritized by impact.

Author: KavinKrishnan

docs/optimization-comparison.md

@@ -0,0 +1,392 @@
+# ModelExpress vs. State-of-the-Art Weight Transfer Optimization
+
+This document compares our current ModelExpress implementation against the strategies described in recent blog posts on achieving ultra-fast (1-2 second) weight transfers for trillion-parameter RL training:
+
+**References:**
+- [Journey to 2-second Inter-node RL Weight Transfer](https://le.qun.ch/en/blog/2025/09/07/rl-weight-transfer/) (Lequn Chen, Sep 2025)
+- [Quick Follow-up on Inter-node RL Weight Transfer](https://le.qun.ch/en/blog/2025/09/17/rl-weight-transfer-2/) (Lequn Chen, Sep 2025)
+- [Weight Transfer for RL Post-Training in under 2 seconds](https://research.perplexity.ai/articles/weight-transfer-for-rl-post-training-in-under-2-seconds) (Perplexity Research, Sep 2025)
+
+---
+
+## Summary Comparison
+
+| Strategy | Blog Posts | ModelExpress | Gap |
+|----------|------------|--------------|-----|
+| **RDMA Direction** | WRITE (push) | READ (pull) | Different model |
+| **Transfer Transparency** | Invisible to inference | Target must receive | We modify loader |
+| **Routing Table** | Static, computed once | Dynamic per-transfer | Performance |
+| **Pipelining** | 4-stage overlap | Sequential | **Major gap** |
+| **Memory Registration** | CUDACachingAllocator blocks | Individual tensors | **Optimization opportunity** |
+| **DeviceMesh Groups** | Parallel groups + barriers | 1:1 rank matching | Different model |
+| **Load Balancing** | Source selection by bytes | Fixed rank matching | N/A for our use case |
+| **Global Barrier** | GLOO async (Ethernet) | Redis polling | Could improve |
+| **GPU Memory Cap** | Configurable watermark | No limit | Risk of OOM |
+| **Warmup Handling** | full_tensor() + quantize | Pre-transfer | Similar |
+
+---
+
+## Detailed Analysis
+
+### 1. RDMA Direction: WRITE vs READ
+
+**Blog Approach (WRITE - Push):**
+```python
+# Training GPU pushes weights to inference GPU
+def transfer_weights(self):
+    for entry in self.routing_table:
+        submit_rdma_write(src_mr, dst_mr, ...)  # One-sided, receiver is passive
+```
+- Receiver never knows weights changed
+- No control plane on inference side
+- Training drives all transfers
+
+**ModelExpress Approach (READ - Pull):**
+```python
+# Target worker pulls weights from source
+def receive_from_source(self, source_metadata, source_tensors):
+    handle = agent.make_prepped_xfer("READ", dst_descs, src_descs, ...)
+    agent.transfer(handle)  # Target initiates RDMA read
+```
+
+**Why We Use READ:**
+- Our use case is different: targets are new vLLM instances, not updated inference nodes
+- Targets need to actively coordinate (wait for source ready, receive, process FP8)
+- Source doesn't know about targets in advance
+
+**Gap Assessment:** This is a **design difference**, not a gap. For RL training→inference, WRITE makes sense. For our model-loading use case, READ is appropriate.
+
+---
+
+### 2. Routing Table: Static vs Dynamic
+
+**Blog Approach:**
+```python
+# Computed ONCE at initialization
+def controller_main():
+    schedule = compute_weight_transfer_schedule(trainer_params, rollout_params, ...)
+    for trainer in trainers:
+        trainer.set_routing_table(routing_table)  # Store forever
+    
+    while training:
+        train()
+        ray.get([trainer.transfer_weights.remote() for trainer in trainers])  # Just execute
+```
+
+**ModelExpress Approach:**
+```python
+# Queried dynamically each time
+def receive_from_source(...):
+    response = stub.GetMetadata(model_name)  # Query server
+    source_tensors = response.workers[rank].tensors
+    # Build transfer descriptors on the fly
+    remote_descs = [(t.addr, t.size, t.device_id) for t in source_tensors]
+```
+
+**Gap Assessment:** **Moderate gap** for repeated transfers. Our current model assumes one-time transfer per target instance. If we supported repeated weight updates (e.g., LoRA fine-tuning), we'd benefit from static routing.
+
+**Optimization Opportunity:**
+```python
+class NixlTransferManager:
+    _cached_routing: dict[str, PreparedTransfer] = {}
+    
+    def receive_from_source(self, source_metadata, source_tensors):
+        cache_key = hash((source_metadata, tuple(source_tensors)))
+        if cache_key in self._cached_routing:
+            return self._execute_cached(cache_key)
+        # ... build and cache routing table
+```
+
+---
+
+### 3. Pipelining: 4-Stage vs Sequential
+
+**Blog Approach (4-stage pipeline):**
+```
+Stage 1: H2D memcpy        ─────────────────────────────────────►
+Stage 2: GPU ops (full_tensor, fusion, quant)  ─────────────────►
+Stage 3: RDMA transfer     ─────────────────────────────────────►
+Stage 4: Global barrier    ─────────────────────────────────────►
+
+    [Task A: H2D] [Task B: H2D] [Task C: H2D]
+              [Task A: GPU]  [Task B: GPU]  [Task C: GPU]
+                       [Task A: RDMA] [Task B: RDMA] [Task C: RDMA]
+```
+
+They use CUDA events to check GPU completion without blocking Python:
+```python
+task.gpu_op_done = torch.cuda.Event()
+task.gpu_op_done.record()
+
+# Later, non-blocking check:
+if task.gpu_op_done.query():  # GPU work done
+    submit_rdma_write(...)    # Start network transfer
+```
+
+**ModelExpress Approach (Sequential):**
+```python
+def receive_from_source(self, source_tensors, ...):
+    # All blocking, no overlap
+    remote_agent_name = self._agent.add_remote_agent(source_metadata)
+    
+    for tensor in source_tensors:
+        remote_descs.append((tensor.addr, tensor.size, ...))
+    
+    src_prepped = self._agent.prep_xfer_dlist(remote_agent_name, remote_descs, ...)
+    dst_prepped = self._agent.prep_xfer_dlist("", local_descs, ...)
+    
+    handle = self._agent.make_prepped_xfer(...)
+    self._agent.transfer(handle)
+    
+    # Block until complete
+    while agent.check_xfer_state(handle) not in ("DONE", "SUCCESS"):
+        time.sleep(0.001)
+```
+
+**Gap Assessment:** **Major gap**. We have no pipelining. Each tensor transfer completes before the next starts.
+
+**Optimization Opportunity:**
+```python
+class PipelinedTransferManager:
+    def __init__(self, max_concurrent=4, max_tmp_bytes=1<<30):
+        self.max_concurrent = max_concurrent
+        self.max_tmp_bytes = max_tmp_bytes
+        
+    def receive_all(self, tensors):
+        pending = deque(tensors)
+        in_flight = []
+        
+        while pending or in_flight:
+            # Launch new transfers up to limit
+            while pending and len(in_flight) < self.max_concurrent:
+                if self._tmp_bytes + pending[0].size > self.max_tmp_bytes:
+                    break
+                tensor = pending.popleft()
+                handle = self._start_async_transfer(tensor)
+                in_flight.append((tensor, handle))
+                
+            # Poll for completion
+            completed = [t for t in in_flight if self._check_done(t[1])]
+            for t in completed:
+                in_flight.remove(t)
+                self._tmp_bytes -= t[0].size
+```
+
+---
+
+### 4. Memory Registration Strategy
+
+**Blog Approach (CUDACachingAllocator blocks):**
+```python
+# Register entire allocator blocks, not individual tensors
+blocks = torch.cuda.memory.memory_snapshot()
+for block in blocks:
+    agent.register_memory([(block['address'], block['size'], device_id, 'cuda')])
+```
+- Fewer memory registrations (hundreds vs thousands)
+- Contiguous blocks enable bulk transfers
+
+**ModelExpress Approach (Individual tensors):**
+```python
+# Register each tensor separately
+for name, tensor in tensors.items():
+    tensor_descriptors.append(TensorDescriptor(
+        name=name,
+        addr=tensor.data_ptr(),
+        size=tensor.numel() * tensor.element_size(),
+        ...
+    ))
+agent.register_memory(tensor_list, backends=["UCX"])
+```
+
+**ModelExpress Contiguous Mode (experimental, blocked):**
+```python
+# Attempt to coalesce adjacent tensors
+regions = _find_contiguous_regions(tensor_descriptors)  # ~30 regions
+agent.register_memory(region_tuples, backends=["UCX"])  # FAILS with rkey errors
+```
+
+**Gap Assessment:** **Moderate gap**. We register ~1327 tensors per GPU. Registering allocator blocks could reduce this to ~10-50.
+
+**Optimization Opportunity:**
+```python
+def register_allocator_blocks(self):
+    """Register CUDACachingAllocator blocks instead of individual tensors."""
+    snapshot = torch.cuda.memory.memory_snapshot()
+    blocks = [(b['address'], b['size'], self._device_id, 'cuda') 
+              for b in snapshot if b['state'] == 'active_allocated']
+    self._agent.register_memory(blocks, backends=["UCX"])
+    
+    # Build tensor→block mapping for transfer addressing
+    self._tensor_to_block = {}
+    for name, tensor in self._tensors.items():
+        ptr = tensor.data_ptr()
+        for block in blocks:
+            if block[0] <= ptr < block[0] + block[1]:
+                offset = ptr - block[0]
+                self._tensor_to_block[name] = (block, offset)
+                break
+```
+
+---
+
+### 5. DeviceMesh Groups and Parallelism
+
+**Blog Approach:**
+```
+Non-MoE DeviceMesh Group (FSDP shards on NVLink):
+  DeviceMesh 0: GPU 0-7, 16-23  ─── all-gather over NVLink (fast)
+  DeviceMesh 1: GPU 8-15, 24-31
+
+MoE DeviceMesh Group (FSDP shards over RDMA):
+  DeviceMesh 0: GPU 0, 16
+  DeviceMesh 1: GPU 1, 17
+  ...
+  
+Workflow:
+  [Non-MoE transfers in parallel] ──barrier──> [MoE transfers in parallel]
+```
+
+**ModelExpress Approach:**
+```
+Simple 1:1 rank matching:
+  Source Worker 0 ──RDMA──> Target Worker 0
+  Source Worker 1 ──RDMA──> Target Worker 1
+  ...
+```
+
+**Gap Assessment:** **Different model**. We don't have FSDP/DTensor sharding to deal with. Our transfers are already embarrassingly parallel (all 8 workers transfer simultaneously).
+
+**What We Could Improve:**
+- Add optional global barriers between transfer phases
+- Support for pipeline parallelism if targets span multiple nodes
+
+---
+
+### 6. Global Barrier Implementation
+
+**Blog Approach:**
+```python
+# Use GLOO over Ethernet (non-blocking, overlapped with RDMA)
+barrier = torch.distributed.barrier(async_op=True)
+
+# Kick off after last full_tensor() in mesh group
+# Completes in parallel with RDMA transfers
+```
+
+**ModelExpress Approach:**
+```python
+# Redis polling (blocking, not ideal)
+while True:
+    data = redis_client.get(f"mx:nixl_ready:{model}:worker:{id}")
+    if data and json.loads(data).get("stability_verified"):
+        break
+    time.sleep(10)  # Polling interval
+```
+
+**Gap Assessment:** **Minor gap**. Our Redis polling works but is slower than native collectives.
+
+**Optimization Opportunity:**
+```python
+# Use torch.distributed.barrier() for multi-target coordination
+if os.environ.get("MX_USE_DIST_BARRIER"):
+    torch.distributed.init_process_group(backend="gloo")
+    torch.distributed.barrier()
+else:
+    # Fallback to Redis
+    wait_for_redis_flag(...)
+```
+
+---
+
+### 7. GPU Memory Watermarking
+
+**Blog Approach:**
+```python
+def _poll_progress(self):
+    while self.tasks_not_started:
+        task = self.tasks_not_started[0]
+        if self.tmp_bytes + task.total_bytes > self.max_tmp_bytes:
+            break  # Don't start new tasks if over limit
+        ...
+```
+
+**ModelExpress Approach:**
+- No memory watermarking
+- Potential OOM if all tensors registered simultaneously
+
+**Gap Assessment:** **Minor gap** for our use case since we register tensors in the order they're loaded.
+
+**Optimization Opportunity:**
+```python
+class NixlTransferManager:
+    def __init__(self, max_tmp_bytes=None):
+        self.max_tmp_bytes = max_tmp_bytes or float('inf')
+        self._tmp_bytes = 0
+        
+    def register_tensors(self, tensors):
+        for name, tensor in tensors.items():
+            size = tensor.numel() * tensor.element_size()
+            if self._tmp_bytes + size > self.max_tmp_bytes:
+                self._flush_pending()  # Wait for previous registrations
+            self._register_single(name, tensor)
+            self._tmp_bytes += size
+```
+
+---
+
+## Strategies NOT Applicable to ModelExpress
+
+| Strategy | Reason |
+|----------|--------|
+| **full_tensor() for FSDP** | We don't use FSDP; weights are already sharded by TP |
+| **Projection fusion on transfer** | We fuse projections in vLLM, not during transfer |
+| **On-the-fly quantization** | Source has already quantized; we transfer raw FP8 |
+| **Training→Inference transparency** | Our targets are new instances, not updated engines |
+| **Multi-source load balancing** | We have 1:1 source-target matching by rank |
+
+---
+
+## Recommended Optimizations
+
+### High Priority (Large Impact)
+
+1. **Pipeline RDMA with GPU processing**
+   - Overlap FP8 processing with ongoing transfers
+   - Use CUDA events for non-blocking completion checks
+   - Expected: 20-40% reduction in total time
+
+2. **Register CUDACachingAllocator blocks**
+   - Reduce memory registrations from 1327 to ~50
+   - Enables bulk transfers without contiguous region bugs
+   - Expected: 10-20% reduction in registration overhead
+
+3. **Batch transfer preparation**
+   - Call `prep_xfer_dlist()` once for all tensors, not per-tensor
+   - Currently blocked by NIXL API understanding
+
+### Medium Priority (Moderate Impact)
+
+4. **Static routing cache**
+   - Cache prepped transfer descriptors after first transfer
+   - Useful if we add incremental weight updates
+
+5. **torch.distributed barrier**
+   - Replace Redis polling with GLOO barrier for multi-target sync
+   - Faster and more reliable
+
+### Low Priority (Minor Impact)
+
+6. **GPU memory watermarking**
+   - Prevent OOM during large transfers
+   - Currently not an issue with sequential processing
+
+---
+
+## Conclusion
+
+Our ModelExpress implementation uses solid fundamentals (NIXL/UCX, RDMA, TP-aware transfers) but lacks the **pipelining** and **batch registration** optimizations that enable 1-2 second transfers in the blog posts.
+
+The key insight from the blog posts is that weight transfer should be treated as a **4-stage pipeline** (H2D, GPU ops, RDMA, barrier) with tasks flowing through stages asynchronously. Our current sequential approach leaves significant performance on the table.
+
+However, many blog optimizations (FSDP handling, on-the-fly quantization, training→inference transparency) don't apply to our model-loading use case. Our ~40-80s transfer time for 681GB is reasonable for initial model loading, but could be reduced to ~20-30s with pipelining and batch registration.