# Task Plan: Fix GPU-only Mode Performance Issue

## Goal
Eliminate the `store_kvcache` scatter overhead in GPU-only mode by using **contiguous KV cache layout** (like offload mode), avoiding PagedAttention's blocked layout for single-sequence inference.

## Problem Summary

GPU-only mode with MInference is **slower** than CPU offload mode:

| Mode | Prefill Speed (32K tokens, Qwen3-4B) |
|------|--------------------------------------|
| GPU-only + MInference | 3383 tok/s |
| Offload + MInference | 5373 tok/s |

**Root cause**: PagedAttention's blocked layout requires expensive `index_copy_` scatter operations to convert contiguous K,V to blocked format.

## Key Insight: Why Offload is Fast

Offload mode uses **contiguous layout** for KV cache:

```python
# OffloadEngine's CPU cache layout
k_cache_cpu: [num_layers, num_blocks, block_size, kv_heads, head_dim]

# Store is simple contiguous slice assignment
self.k_cache_cpu[layer_id, block_id, :actual_size].copy_(k[start:end])
```

The K,V computed during prefill `[seq_len, kv_heads, head_dim]` matches the cache layout - no format conversion needed!

## Solution: Contiguous Layout for GPU-only Mode

For GPU-only single-sequence mode, use the **same contiguous layout as offload mode**, but on GPU:

```
Current GPU-only (PagedAttention):
  Cache: [num_blocks, block_size, kv_heads, head_dim] (blocked)
  Store: scatter via index_copy_ (SLOW)

Proposed GPU-only (Contiguous):
  Cache: [num_layers, max_seq_len, kv_heads, head_dim] (contiguous)
  Store: slice assignment k_cache[layer_id, :seq_len] = k (FAST)
```

This mirrors offload mode's architecture but keeps everything on GPU - no cross-device transfer, no layout conversion.

## Phases
- [x] Phase 1: Add contiguous GPU KV cache in GPUOnlyManager (for single-seq mode)
- [x] Phase 2: Implement `run_gpu_only_prefill()` using contiguous cache
- [x] Phase 3: Implement decode path for contiguous cache
- [x] Phase 4: Test and validate performance

## Results

| Mode | 32K Prefill Speed | Notes |
|------|-------------------|-------|
| GPU-only (before) | ~3383 tok/s | PagedAttention scatter overhead |
| GPU-only contiguous (after) | **5293 tok/s** | 56% improvement |
| Offload mode | 5391 tok/s | Baseline comparison |

**Test passed**: `test_needle.py --input-len 32768 --max-model-len 40960` - correct output retrieved.

## Detailed Design

### Phase 1: Contiguous GPU KV Cache

**File**: `nanovllm/kvcache/gpu_manager.py`

Add contiguous cache allocation for single-sequence mode:

```python
class GPUOnlyManager(KVCacheManager):
    def __init__(self, num_blocks: int, block_size: int, max_seq_len: int = 0):
        # ... existing code ...
        self.max_seq_len = max_seq_len

        # Contiguous cache for single-seq mode (allocated in allocate_cache)
        self.contiguous_k_cache = None  # [num_layers, max_seq_len, kv_heads, head_dim]
        self.contiguous_v_cache = None

    def allocate_cache(
        self,
        num_layers: int,
        num_kv_heads: int,
        head_dim: int,
        dtype: torch.dtype,
    ) -> None:
        # Existing PagedAttention cache for multi-seq/decode
        self.kv_cache = torch.empty(
            2, num_layers, self._num_blocks, self._block_size,
            num_kv_heads, head_dim,
            dtype=dtype, device="cuda"
        )

        # Contiguous cache for single-seq prefill (if max_seq_len specified)
        if self.max_seq_len > 0:
            self.contiguous_k_cache = torch.empty(
                num_layers, self.max_seq_len, num_kv_heads, head_dim,
                dtype=dtype, device="cuda"
            )
            self.contiguous_v_cache = torch.empty(
                num_layers, self.max_seq_len, num_kv_heads, head_dim,
                dtype=dtype, device="cuda"
            )
```

### Phase 2: Layer-wise GPU-only Prefill

**File**: `nanovllm/engine/model_runner.py`

Following offload pattern exactly - store K,V per-layer to contiguous cache:

```python
@torch.inference_mode()
def run_gpu_only_prefill(self, seqs: list[Sequence]) -> list[int]:
    """
    GPU-only prefill with contiguous KV cache layout.

    Mirrors run_layerwise_offload_prefill() but stores to GPU instead of CPU.
    No scatter operations - just contiguous slice assignment.
    """
    assert len(seqs) == 1, "GPU-only layer-wise prefill only supports single sequence"
    seq = seqs[0]

    num_layers = len(self.model.model.layers)
    total_tokens = len(seq)

    # Get contiguous GPU cache
    k_cache = self.kvcache_manager.contiguous_k_cache
    v_cache = self.kvcache_manager.contiguous_v_cache

    # Prepare inputs
    input_ids = torch.tensor(seq[:], dtype=torch.int64, device="cuda")
    positions = torch.arange(total_tokens, dtype=torch.int64, device="cuda")

    from flash_attn.flash_attn_interface import flash_attn_varlen_func
    cu_seqlens = torch.tensor([0, total_tokens], dtype=torch.int32, device="cuda")

    # Embedding
    hidden_states = self.model.model.embed_tokens(input_ids)
    residual = None

    # Layer-by-layer processing (same as offload prefill)
    for layer_id in range(num_layers):
        layer = self.model.model.layers[layer_id]

        # Input LayerNorm
        if residual is None:
            hidden_ln, residual = layer.input_layernorm(hidden_states), hidden_states
        else:
            hidden_ln, residual = layer.input_layernorm(hidden_states, residual)

        # QKV projection
        qkv = layer.self_attn.qkv_proj(hidden_ln)
        q, k, v = qkv.split([
            layer.self_attn.q_size,
            layer.self_attn.kv_size,
            layer.self_attn.kv_size
        ], dim=-1)

        q = q.view(total_tokens, layer.self_attn.num_heads, layer.self_attn.head_dim)
        k = k.view(total_tokens, layer.self_attn.num_kv_heads, layer.self_attn.head_dim)
        v = v.view(total_tokens, layer.self_attn.num_kv_heads, layer.self_attn.head_dim)

        # Q/K norms (Qwen3 specific)
        if not layer.self_attn.qkv_bias:
            q = layer.self_attn.q_norm(q.reshape(-1, layer.self_attn.head_dim))
            q = q.view(total_tokens, layer.self_attn.num_heads, layer.self_attn.head_dim)
            k = layer.self_attn.k_norm(k.reshape(-1, layer.self_attn.head_dim))
            k = k.view(total_tokens, layer.self_attn.num_kv_heads, layer.self_attn.head_dim)

        # RoPE
        q, k = layer.self_attn.rotary_emb(positions, q, k)

        # Store K,V to contiguous GPU cache (same layout - no conversion!)
        # This is just slice assignment, not scatter
        k_cache[layer_id, :total_tokens] = k
        v_cache[layer_id, :total_tokens] = v

        # Sparse or Full attention (uses k, v directly)
        if self.sparse_prefill_policy is not None:
            attn_output = self.sparse_prefill_policy.sparse_prefill_attention(
                q, k, v, layer_id
            )
        else:
            attn_output = flash_attn_varlen_func(
                q, k, v,
                cu_seqlens_q=cu_seqlens,
                cu_seqlens_k=cu_seqlens,
                max_seqlen_q=total_tokens,
                max_seqlen_k=total_tokens,
                softmax_scale=layer.self_attn.attn.scale,
                causal=True,
            )

        # O projection
        attn_output = attn_output.view(total_tokens, -1)
        hidden_states = layer.self_attn.o_proj(attn_output)

        # Post-attention LayerNorm + MLP
        hidden_states, residual = layer.post_attention_layernorm(hidden_states, residual)
        hidden_states = layer.mlp(hidden_states)

    # Final norm
    hidden_states, _ = self.model.model.norm(hidden_states, residual)

    # Compute logits
    logits = self.model.compute_logits(hidden_states[-1:])

    # Record prefill length for decode
    self.kvcache_manager.contiguous_seq_len = total_tokens

    # Sample
    temperatures = self.prepare_sample(seqs) if self.rank == 0 else None
    token_ids = self.sampler(logits, temperatures).tolist() if self.rank == 0 else None

    return token_ids
```

### Phase 3: Decode with Contiguous Cache

**File**: `nanovllm/engine/model_runner.py`

```python
@torch.inference_mode()
def run_gpu_only_decode(self, seqs: list[Sequence]) -> list[int]:
    """
    Decode using contiguous GPU KV cache.

    Similar to offload decode but simpler - all KV already on GPU.
    """
    assert len(seqs) == 1
    seq = seqs[0]

    num_layers = len(self.model.model.layers)
    k_cache = self.kvcache_manager.contiguous_k_cache
    v_cache = self.kvcache_manager.contiguous_v_cache
    context_len = self.kvcache_manager.contiguous_seq_len

    # Prepare inputs
    input_ids = torch.tensor([seq.last_token], dtype=torch.int64, device="cuda")
    positions = torch.tensor([len(seq) - 1], dtype=torch.int64, device="cuda")

    from flash_attn.flash_attn_interface import flash_attn_varlen_func
    cu_seqlens_q = torch.tensor([0, 1], dtype=torch.int32, device="cuda")

    # Embedding
    hidden_states = self.model.model.embed_tokens(input_ids)
    residual = None

    for layer_id in range(num_layers):
        layer = self.model.model.layers[layer_id]

        # Input LayerNorm
        if residual is None:
            hidden_ln, residual = layer.input_layernorm(hidden_states), hidden_states
        else:
            hidden_ln, residual = layer.input_layernorm(hidden_states, residual)

        # QKV projection
        qkv = layer.self_attn.qkv_proj(hidden_ln)
        q, k_new, v_new = qkv.split([
            layer.self_attn.q_size,
            layer.self_attn.kv_size,
            layer.self_attn.kv_size
        ], dim=-1)

        q = q.view(1, layer.self_attn.num_heads, layer.self_attn.head_dim)
        k_new = k_new.view(1, layer.self_attn.num_kv_heads, layer.self_attn.head_dim)
        v_new = v_new.view(1, layer.self_attn.num_kv_heads, layer.self_attn.head_dim)

        # Q/K norms
        if not layer.self_attn.qkv_bias:
            q = layer.self_attn.q_norm(q.reshape(-1, layer.self_attn.head_dim))
            q = q.view(1, layer.self_attn.num_heads, layer.self_attn.head_dim)
            k_new = layer.self_attn.k_norm(k_new.reshape(-1, layer.self_attn.head_dim))
            k_new = k_new.view(1, layer.self_attn.num_kv_heads, layer.self_attn.head_dim)

        # RoPE
        q, k_new = layer.self_attn.rotary_emb(positions, q, k_new)

        # Get cached K,V and append new token
        k_cached = k_cache[layer_id, :context_len]
        v_cached = v_cache[layer_id, :context_len]

        # Store new K,V to cache
        k_cache[layer_id, context_len] = k_new.squeeze(0)
        v_cache[layer_id, context_len] = v_new.squeeze(0)

        # Full K,V for attention
        k_full = k_cache[layer_id, :context_len + 1]
        v_full = v_cache[layer_id, :context_len + 1]

        # Attention
        cu_seqlens_k = torch.tensor([0, context_len + 1], dtype=torch.int32, device="cuda")
        attn_output = flash_attn_varlen_func(
            q, k_full, v_full,
            cu_seqlens_q=cu_seqlens_q,
            cu_seqlens_k=cu_seqlens_k,
            max_seqlen_q=1,
            max_seqlen_k=context_len + 1,
            softmax_scale=layer.self_attn.attn.scale,
            causal=False,  # Single query, no causal needed
        )

        # O projection
        attn_output = attn_output.view(1, -1)
        hidden_states = layer.self_attn.o_proj(attn_output)

        # Post-attention LayerNorm + MLP
        hidden_states, residual = layer.post_attention_layernorm(hidden_states, residual)
        hidden_states = layer.mlp(hidden_states)

    # Update context length
    self.kvcache_manager.contiguous_seq_len = context_len + 1

    # Final norm
    hidden_states, _ = self.model.model.norm(hidden_states, residual)

    # Compute logits
    logits = self.model.compute_logits(hidden_states)

    # Sample
    temperatures = self.prepare_sample(seqs) if self.rank == 0 else None
    token_ids = self.sampler(logits, temperatures).tolist() if self.rank == 0 else None

    return token_ids
```

### Phase 4: Decision Logic

```python
def _should_use_contiguous_gpu_mode(self, seqs: list[Sequence], is_prefill: bool) -> bool:
    """Check if contiguous GPU mode should be used."""
    # Must have contiguous cache allocated
    if not hasattr(self.kvcache_manager, 'contiguous_k_cache'):
        return False
    if self.kvcache_manager.contiguous_k_cache is None:
        return False

    # Must NOT be offload mode
    if hasattr(self.kvcache_manager, 'offload_engine'):
        return False

    # Single sequence only
    if len(seqs) != 1:
        return False

    # For prefill: has blocks (not warmup)
    if is_prefill and not seqs[0].block_table:
        return False

    return True


def run(self, seqs: list[Sequence], is_prefill: bool) -> list[int]:
    # Check offload mode (existing)
    if hasattr(self, 'kvcache_manager') and hasattr(self.kvcache_manager, 'offload_engine'):
        ...

    # Check contiguous GPU mode
    if self._should_use_contiguous_gpu_mode(seqs, is_prefill):
        if is_prefill:
            return self.run_gpu_only_prefill(seqs)
        else:
            return self.run_gpu_only_decode(seqs)

    # Standard PagedAttention path
    ...
```

## Architecture Comparison

| Aspect | Offload Mode | GPU-only (Proposed) | GPU-only (Current) |
|--------|--------------|---------------------|-------------------|
| Cache location | CPU (contiguous) | GPU (contiguous) | GPU (PagedAttention) |
| Cache layout | `[layers, blocks, block_size, heads, dim]` | `[layers, max_seq_len, heads, dim]` | `[blocks, block_size, heads, dim]` |
| Prefill store | Contiguous slice copy | **Slice assignment (no copy!)** | Scatter (index_copy_) |
| Decode read | H2D ring buffer | Direct GPU access | PagedAttention |

## Key Points

1. **No explicit copy_ needed**: Slice assignment `cache[layer, :len] = k` is direct memory write
2. **Same layout as computed K,V**: No format conversion required
3. **Mirrors offload architecture**: Same layer-wise processing pattern
4. **GPU advantage**: No cross-device transfer, faster than offload

## Memory Usage

Contiguous GPU cache: `2 * num_layers * max_seq_len * kv_heads * head_dim * dtype_size`

For Qwen3-4B with 32K max_seq_len:
- `2 * 28 * 32768 * 8 * 128 * 2 = 3.5GB`

Same as offload mode's CPU cache, but on GPU.

## Files to Modify

| File | Changes |
|------|---------|
| `nanovllm/kvcache/gpu_manager.py` | Add contiguous cache allocation |
| `nanovllm/engine/model_runner.py` | Add `run_gpu_only_prefill()`, `run_gpu_only_decode()`, modify `run()` |

## Expected Performance

| Metric | Before | After | Improvement |
|--------|--------|-------|-------------|
| GPU-only prefill (32K) | 3383 tok/s | ~5400+ tok/s | ~60%+ |
| Decode | Baseline | Similar | ~0% |

## Status
**Currently in Phase 1** - Ready to implement contiguous GPU cache