# GPU Pipelining - Phase 3: CUDA Stream Optimization

**Status:** 🔮 Future / Not Implemented

## Goal

Maximize GPU utilization by overlapping data transfer (CPU→GPU) with computation using CUDA streams. This is an advanced optimization that hides the latency of moving data to GPU memory.

## Problem

Even with Phase 1 and Phase 2 optimizations, there's still a sequential dependency:

```
1. ray.get(batch_ref)           # Fetch from object store (~1ms)
2. tensor.to("cuda")            # CPU → GPU transfer (~2-5ms)
3. model(tensor)                # GPU inference (~20ms)
4. result.to("cpu")             # GPU → CPU transfer (~1ms)
```

The GPU is idle during steps 1-2 and 4. With CUDA streams, we can overlap these operations.

## CUDA Streams Background

CUDA streams allow concurrent execution of operations on different streams:

```python
import torch

stream1 = torch.cuda.Stream()  # Data transfer stream
stream2 = torch.cuda.Stream()  # Compute stream

# These can execute concurrently:
with torch.cuda.stream(stream1):
    next_tensor = tensor.to("cuda")  # Transfer next batch

with torch.cuda.stream(stream2):
    result = model(current_tensor)   # Compute current batch
```

## Proposed Architecture: Double-Buffered Execution

```
┌─────────────────────────────────────────────────────────────────┐
│                      GPU Actor                                   │
│  ┌──────────────┐    ┌──────────────┐    ┌──────────────┐       │
│  │ Buffer A     │←──→│   Model      │←──→│ Buffer B     │       │
│  │ (loading)    │    │ (computing)  │    │ (offloading) │       │
│  └──────────────┘    └──────────────┘    └──────────────┘       │
│       Stream 1           Stream 2            Stream 3            │
└─────────────────────────────────────────────────────────────────┘
```

### Pipeline Overlap

| Time | Stream 1 (Load) | Stream 2 (Compute) | Stream 3 (Offload) |
|------|-----------------|--------------------|--------------------|
| T0   | Load batch 0    | -                  | -                  |
| T1   | Load batch 1    | Compute batch 0    | -                  |
| T2   | Load batch 2    | Compute batch 1    | Offload batch 0    |
| T3   | Load batch 3    | Compute batch 2    | Offload batch 1    |
| ...  | ...             | ...                | ...                |

## Proposed Implementation

### StreamingGPUApplier

```python
@ray.remote(num_gpus=1)
class StreamingGPUApplier:
    """GPU applier with CUDA stream optimization."""

    def __init__(self, udf: UDF, num_buffers: int = 2):
        self.udf = udf
        self.num_buffers = num_buffers

        # Create CUDA streams
        self.load_stream = torch.cuda.Stream()
        self.compute_stream = torch.cuda.Stream()
        self.offload_stream = torch.cuda.Stream()

        # Pre-allocate GPU buffers if sizes known
        self.buffers = [None] * num_buffers

    def process_batches(
        self, batch_refs: list[ObjectRef]
    ) -> Iterator[pa.RecordBatch]:
        """Process batches with stream overlap."""

        current_gpu = None
        next_gpu = None
        pending_result = None

        for i, batch_ref in enumerate(batch_refs):
            # Stream 1: Load next batch to GPU
            with torch.cuda.stream(self.load_stream):
                cpu_batch = ray.get(batch_ref)
                next_gpu = self._to_gpu_tensors(cpu_batch)

            # Stream 3: Offload previous result
            if pending_result is not None:
                with torch.cuda.stream(self.offload_stream):
                    yield self._to_arrow(pending_result)

            # Stream 2: Compute current batch
            if current_gpu is not None:
                with torch.cuda.stream(self.compute_stream):
                    pending_result = self.udf(current_gpu)

            # Synchronize streams and rotate buffers
            torch.cuda.synchronize()
            current_gpu = next_gpu
            next_gpu = None

        # Process final batch
        if current_gpu is not None:
            with torch.cuda.stream(self.compute_stream):
                pending_result = self.udf(current_gpu)
            torch.cuda.synchronize()

        if pending_result is not None:
            yield self._to_arrow(pending_result)

    def _to_gpu_tensors(self, batch: pa.RecordBatch) -> torch.Tensor:
        """Convert batch to GPU tensors (on load stream)."""
        # Implementation depends on batch format
        ...

    def _to_arrow(self, result: torch.Tensor) -> pa.RecordBatch:
        """Convert GPU result to Arrow (on offload stream)."""
        cpu_result = result.cpu()
        return pa.RecordBatch.from_pydict({"result": cpu_result.numpy()})
```

### Pinned Memory for Faster Transfers

For maximum transfer speed, use pinned (page-locked) CPU memory:

```python
class PinnedMemoryPool:
    """Pool of pinned memory buffers for fast CPU↔GPU transfer."""

    def __init__(self, buffer_size: int, num_buffers: int):
        self.buffers = [
            torch.empty(buffer_size, pin_memory=True)
            for _ in range(num_buffers)
        ]
        self.available = list(range(num_buffers))

    def get_buffer(self) -> tuple[int, torch.Tensor]:
        """Get an available pinned buffer."""
        idx = self.available.pop()
        return idx, self.buffers[idx]

    def return_buffer(self, idx: int):
        """Return buffer to pool."""
        self.available.append(idx)
```

### Async Data Loading with Prefetch

Combine with Phase 1 prefetching for full pipeline:

```python
class StreamingApplierWithPrefetch:
    """Combines prefetching with stream optimization."""

    def __init__(self, loader_pool: LoaderPool, gpu_applier: StreamingGPUApplier):
        self.loader_pool = loader_pool
        self.gpu_applier = gpu_applier

    async def process(self, read_tasks: Iterator[ReadTask]) -> Iterator[Result]:
        # Prefetch in background
        prefetch_iter = self.loader_pool.prefetch_iter(read_tasks, batch_size)

        # Collect batch refs for streaming GPU processing
        batch_refs = [task.batch_refs[0] for task in prefetch_iter]

        # Process with CUDA stream optimization
        yield from self.gpu_applier.process_batches(batch_refs)
```

## Configuration

```python
class JobConfig:
    # Phase 1 (existing)
    enable_gpu_pipelining: bool = False
    num_loader_actors: int = 4
    prefetch_depth: int = 8

    # Phase 2 (future)
    enable_cpu_preprocessing: bool = False
    num_preprocessor_actors: int = 4

    # Phase 3 (new)
    enable_cuda_streams: bool = False
    num_cuda_buffers: int = 2  # Double-buffering default
    use_pinned_memory: bool = True
```

## UDF Requirements

For stream optimization to work, UDFs must:

1. **Not synchronize internally**: Avoid `torch.cuda.synchronize()` in UDF
2. **Return tensors**: Don't convert to CPU/Arrow inside UDF
3. **Be stateless per-call**: No cross-batch state dependencies

### Stream-Aware UDF Interface

```python
@udf(num_gpus=1, stream_aware=True)
class StreamAwareUDF:
    """UDF designed for CUDA stream optimization."""

    def setup(self) -> None:
        self.model = load_model().cuda()

    def __call__(self, gpu_tensor: torch.Tensor) -> torch.Tensor:
        """Process on current CUDA stream. Do NOT synchronize."""
        # Input already on GPU, output stays on GPU
        with torch.no_grad():
            return self.model(gpu_tensor)
        # Do NOT call .cpu() or synchronize here
```

## Performance Analysis

### Theoretical Speedup

| Operation | Time (ms) | Overlapped? |
|-----------|-----------|-------------|
| Fetch from object store | 1 | Yes (prefetch) |
| CPU → GPU transfer | 3 | Yes (stream 1) |
| GPU inference | 20 | Yes (stream 2) |
| GPU → CPU transfer | 1 | Yes (stream 3) |
| **Total (sequential)** | **25** | |
| **Total (pipelined)** | **~21** | (limited by inference) |

**Speedup:** ~19% from stream optimization alone

### Combined with Phase 1 & 2

| Optimization | GPU Utilization | Bottleneck |
|--------------|-----------------|------------|
| None | ~30% | I/O, preprocessing |
| Phase 1 (Prefetch) | ~50% | Preprocessing in GPU actor |
| Phase 1 + 2 (CPU Prep) | ~80% | Data transfer latency |
| Phase 1 + 2 + 3 (Streams) | ~95% | Model inference |

## Challenges

1. **Complexity**: Requires careful synchronization management
2. **Memory**: Multiple GPU buffers increase memory pressure
3. **UDF changes**: Authors must write stream-aware code
4. **Debugging**: Async errors harder to trace
5. **Framework support**: Assumes PyTorch; other frameworks may differ

## Implementation Steps

1. Create `StreamingGPUApplier` class with CUDA stream management
2. Implement `PinnedMemoryPool` for optimized transfers
3. Add stream-aware detection to UDF interface
4. Integrate with LoaderPool from Phase 1
5. Add configuration options
6. Create benchmarks to validate performance
7. Document UDF requirements and best practices

## Open Questions

1. How to handle UDFs that internally synchronize?
2. Should we auto-detect stream compatibility?
3. How to handle variable batch sizes with pre-allocated buffers?
4. Integration with other ML frameworks (JAX, TensorFlow)?

## References

- [CUDA Streams Best Practices](https://developer.nvidia.com/blog/how-overlap-data-transfers-cuda-cc/)
- [PyTorch CUDA Semantics](https://pytorch.org/docs/stable/notes/cuda.html)
- [Pinned Memory in PyTorch](https://pytorch.org/docs/stable/notes/cuda.html#use-pinned-memory-buffers)
