# ROCm MTP Memory Accounting and Context Optimization

Session attribution: Sol

## Current Status

Native MTP speculative decoding creates an additional llama context alongside the target context. The `sol` changes improve accounting for this context, fit fixed multi-GPU layouts per device, and make target pipeline scheduling explicit.

The current complete patch is:

```text
/home/eaman/llama/bug/sol4.patch
```

`sol4.patch` contains the original MTP fitting and pipeline changes plus the merged HIP VEC flash-attention dispatch from `vec_sol`.

It provides:

1. Complete MTP context and compute accounting.
2. MTP measurement at the target's fitted context with iterative refinement.
3. Per-device fitting when GPU placement is explicitly fixed.
4. Non-pipeline scheduling for the MTP context.
5. Explicit `--pipeline-parallel auto|on|off` control for the target context.
6. HIP VEC dispatch for supported quantized-KV attention shapes.

The best previously validated ROCm configuration used F16 MTP KV and disabled target pipeline parallelism:

```text
--pipeline-parallel off
--cache-type-k-draft f16
--cache-type-v-draft f16
```

The merged VEC update initially reached **147,712 context tokens** with Q4/Q4 MTP KV on the tested Qwen3.6 27B ROCm system, compared with 121,088 for the previous F16/off build: +26,624 tokens (+21.99%). Startup and short generation succeeded; long-run stability and throughput comparisons remain pending.

---

## Test System

| Component | Configuration |
|---|---|
| Operating system | Debian Sid |
| Target model | ThinkingCap Qwen3.6-27B Q6_K_L |
| Model layout | 64 target layers plus 1 NextN/MTP layer |
| GPU 0 | AMD Radeon RX 6800, 16 GiB |
| GPU 1 | AMD Radeon RX 6700 XT, 12 GiB |
| Backend | HIP/ROCm build with Vulkan support |
| Split mode | Layer |
| Tensor split | 0.6,0.4 |
| GPU layers | 99, explicitly fixed |
| Target KV cache | K=q8_0, V=q5_1 |
| MTP KV cache | K=F16, V=F16 |
| Maximum MTP draft length | 3 |
| Fit target | 50 MiB per device |
| Batch / microbatch | 1024 / 384 |
| Parallel slots | 1 |
| Pipeline mode | Off for target and MTP |
| Base source | build 909 (`7bd8282`) |
| Working branch | `sol` |

Model:

```text
/home/eaman/lm/models/bottlecapai/ThinkingCap-Qwen3.6-27B-Q6_K_L.ggu.gguf
```

Current launcher:

```text
/home/eaman/models/think.sh_sol
```

---

## Issues Addressed

### MTP owns a separate runtime context

Native MTP initializes another `llama_context` for the draft path:

```cpp
llama_context * ctx_dft = llama_init_from_model(model_tgt, cparams);
```

This context owns context-dependent memory in addition to the target context:

- MTP KV memory;
- a GPU compute buffer;
- a host compute buffer;
- backend-specific temporary workspace.

The complete MTP allocation must be included in the memory reservation used by the target fitter.

### MTP memory depends on the fitted context

The model is trained for 262,144 tokens, but available VRAM requires a smaller runtime context. Measuring MTP at the trained context overstates its runtime allocation. Measuring it only once at a preliminary fitted value also remains conservative after the target fit changes.

The dependency is:

```text
target context -> MTP memory -> target reservation -> fitted target context
```

The MTP allocation must be measured at the fitted target context and refined.

### Pipeline scheduling has a high memory cost on this topology

Pipeline parallelism created four scheduler copies for the target context and duplicated a large compute allocation on both GPUs. With one active server slot and these asymmetric GPUs, the extra workspace and synchronization did not provide a measured performance advantage.

The MTP graph also did not benefit from pipeline scheduling because its useful GPU work was placed on one device.

The solution therefore keeps pipeline policy explicit:

- `auto` preserves the normal llama.cpp decision;
- `on` requests pipeline scheduling and warns if unavailable;
- `off` disables it;
- MTP remains non-pipeline.

### Fixed GPU placement was fitted with a generic allowance

The tested command explicitly fixes placement:

```text
-ngl 99
--split-mode layer
--tensor-split 0.6,0.4
```

For this case, reserving memory for possible layer redistribution is unnecessary because the fitter cannot change `n_gpu_layers`. The useful constraint is the memory limit of each individual device.

The optimized fitter interpolates a context limit for every selected GPU and uses the lowest result. Automatic GPU-layer placement retains the existing conservative behavior.

---

## Solution Implemented by `sol4.patch`

### 1. Use the fitted target context for runtime MTP

File: `common/speculative.cpp`

```cpp
if (spec_mtp) {
    cparams.ctx_type = LLAMA_CONTEXT_TYPE_MTP;
    cparams.pipeline_parallel_type = LLAMA_PIPELINE_PARALLEL_TYPE_DISABLED;
    cparams.n_ctx = llama_n_ctx(ctx_tgt);
}
```

The runtime MTP context now uses the actual fitted target context instead of reserving context-dependent workspace for the full trained context.

### 2. Measure the complete MTP allocation

File: `tools/server/server-context.cpp`

```cpp
const size_t bytes =
    (measure_model_bytes ? dmd[j].model : 0) +
    dmd[j].context +
    dmd[j].compute;
```

Both MTP context storage and its independent compute workspace are added to the per-device target reservation.

### 3. Refine the MTP reservation

File: `tools/server/server-context.cpp`

The server performs two refinement passes:

```text
1. Fit the target with the current MTP reservation.
2. Measure MTP at the newly fitted context.
3. Replace the reservation with the refined per-device values.
4. Repeat once and perform the final target fit.
```

This resolves most of the fitting dependency without an open-ended loop.

### 4. Fit fixed multi-GPU layouts per device

File: `common/fit.cpp`

When `n_gpu_layers` is explicitly fixed, the fitter calculates the context supported by each GPU after its configured margin, rounds it to the required 256-token boundary, and selects the limiting device.

### 5. Keep MTP non-pipeline

Files:

```text
common/speculative.cpp
src/llama-context.cpp
tools/server/server-context.cpp
```

MTP fitting and runtime initialization both use the same non-pipeline scheduler policy. This prevents estimation and runtime allocation from selecting different compute-buffer layouts.

### 6. Add explicit target pipeline control

Files:

```text
include/llama.h
src/llama.cpp
common/common.h
common/common.cpp
common/arg.cpp
src/llama-context.cpp
```

The public context parameters and common command-line parser expose:

```text
--pipeline-parallel auto|on|off
```

The default is `auto`, so applying the patch does not silently change target scheduling. The optimized ROCm launcher explicitly selects `off`.

---

## ROCm MTP Cache Finding

For this model and HIP flash-attention path, quantized MTP KV did not minimize total GPU memory. Q4_0 K and V required a large F16 conversion workspace in the selected ROCm flash-attention implementation.

Measured MTP allocations before target pipeline optimization were:

| Draft KV | Context | ROCm1 KV | ROCm1 compute | MTP GPU total |
|---|---:|---:|---:|---:|
| Q4_0 | 99,072 | 108.84 MiB | 519.84 MiB | 628.68 MiB |
| F16 | 104,704 | 409.00 MiB | 127.70 MiB | 536.70 MiB |

F16 used more persistent KV memory but reduced the conversion workspace enough to save approximately 92 MiB of total MTP GPU memory.

F16 is therefore the best validated ROCm MTP cache for this model:

```text
--cache-type-k-draft f16
--cache-type-v-draft f16
```

This is a runtime recommendation, not hard-coded patch behavior. Vulkan and other models must be measured separately.

---

## Final Validation

### Fit sequence

With F16 MTP KV and target pipeline mode off:

```text
Target-only fit:         121,856
Initial MTP estimate:    616.26 MiB
First fitted context:    118,784
Refined MTP estimate:    602.01 MiB
Second fitted context:   119,808
Refined MTP estimate:    606.76 MiB
Final context:           119,552
```

The runtime target and MTP contexts both initialized at 119,552 tokens.

### Runtime allocation

Target context:

```text
ROCm0 KV buffer:         2116.09 MiB
ROCm1 KV buffer:         1269.66 MiB
ROCm0 compute buffer:     614.84 MiB
ROCm1 compute buffer:     614.84 MiB
Host compute buffer:      102.84 MiB
Scheduler copies:              1
```

MTP context:

```text
ROCm1 KV buffer:          467.00 MiB
ROCm1 compute buffer:     138.57 MiB
Host compute buffer:      102.58 MiB
Scheduler copies:              1
```

The validation log contained no `cudaMalloc` failure and no scheduler fallback.

### Pipeline auto versus off

| Metric | F16 / auto | F16 / off | Change |
|---|---:|---:|---:|
| Final context | 104,704 | **119,552** | **+14,848 (+14.18%)** |
| Target compute per GPU | 806.86 MiB | 614.84 MiB | -192.02 MiB |
| Target host compute | 322.87 MiB | 102.84 MiB | -220.03 MiB |
| Target scheduler copies | 4 | 1 | -3 |

### Initial performance sample

One old-prompt sample produced:

| Metric | Comparison run | `sol3`, F16/off | Change |
|---|---:|---:|---:|
| Prompt processing | 206.31 t/s | 251.09 t/s | +21.71% |
| Token generation | 32.49 t/s | 36.05 t/s | +10.96% |
| Draft acceptance | 90.640% | 88.349% | -2.291 points |

This is encouraging but not a controlled benchmark. The generated lengths, build configuration, and draft-cache configuration were not identical. Stability testing and repeated same-build measurements remain necessary.

---

## Source Files Modified

| Source file | Purpose |
|---|---|
| `common/arg.cpp` | Parse the target pipeline mode |
| `common/common.cpp` | Transfer pipeline mode to llama context parameters |
| `common/common.h` | Store the common pipeline mode parameter |
| `common/fit.cpp` | Fit fixed GPU placement against per-device limits |
| `common/speculative.cpp` | Use fitted `n_ctx` and non-pipeline mode for runtime MTP |
| `include/llama.h` | Define the public pipeline mode API |
| `src/llama-context.cpp` | Apply pipeline policy to scheduler creation |
| `src/llama.cpp` | Provide the pipeline mode name helper |
| `ggml/src/ggml-cuda/fattn.cu` | Dispatch supported HIP quantized-KV attention shapes to VEC |
| `tools/server/server-context.cpp` | Measure and refine complete MTP memory |

---

## Important Artifacts

| Artifact | Description |
|---|---|
| `/home/eaman/llama/bug/sol4.patch` | Current complete source patch |
| `/home/eaman/llama/bug/rocm_improvement.patch` | Duplicate copy of `sol4.patch` |
| `/home/eaman/llama/bug/logs/pipeline_off_f16.log` | Final F16/off validation and prompt sample, 119,552 context |
| `/home/eaman/llama/bug/logs/sol2_f16.log` | F16/auto comparison log, 104,704 context |
| `/home/eaman/llama/bug/logs/q6_current_full.log` | Q4_0/auto comparison log, 99,072 context |
| `/home/eaman/models/think.sh_sol` | Current optimized ROCm launcher |
| `/home/eaman/llama/llama.cpp/build/bin` | Source-tree build output |
| `/home/eaman/llama/bin_vulkan` | Deployed HIP and Vulkan build updated by the user |
| `/home/eaman/llama/bug/possible_improvements.md` | Detailed investigation of further optimization paths |

Patch base and checksum:

```text
Base commit: 7bd8282
Solution commit: e37aae4
Branch: sol
SHA-256: afc4216b15e833712c9c9228b41a1134d6f703614f6aec334c73e53f2ef46839
```

---

## Deployment Notes

The current build contains both HIP and Vulkan backends:

```text
GGML_HIP=ON
GGML_VULKAN=ON
```

The complete build output, including all llama and GGML shared libraries, must be deployed together. `/home/eaman/llama/bin_vulkan` has been updated by the user with the current build.

ROCm recommended settings for this model:

```text
--pipeline-parallel off
--cache-type-k-draft f16
--cache-type-v-draft f16
```

Vulkan should initially retain `--pipeline-parallel auto` and compare Q4_0 with F16 MTP KV because the ROCm conversion-workspace result does not automatically apply to Vulkan.

---

## TODO and Development Path

### Phase 0: Validate `sol3`

- [x] Build HIP and Vulkan backends successfully.
- [x] Initialize the 27B Q6 model with F16 MTP KV and pipeline mode off.
- [x] Confirm 119,552 target and MTP context without allocation fallback.
- [x] Run an initial real prompt and confirm normal MTP generation.
- [ ] Exercise the current build with the established old-prompt collection.
- [ ] Test prompts that grow close to the fitted context limit.
- [ ] Run several long generations and watch for delayed OOM, corruption, or scheduler errors.
- [ ] Record at least three same-prompt runs for prompt speed, generation speed, and acceptance.
- [ ] Perform a controlled same-build F16 `auto` versus `off` comparison at fixed context and seed.
- [ ] Smoke-test representative models on Vulkan with pipeline mode `auto`.
- [ ] Compare Q4_0 and F16 MTP KV on Vulkan before choosing a Vulkan default.

### Phase 1: Recurrent-state snapshot memory

This is the next practical context optimization after stability validation. The target context currently allocates 598.50 MiB of recurrent state for three rollback snapshots:

```text
n_rs_seq = 3
ROCm0 recurrent state = 374.06 MiB
ROCm1 recurrent state = 224.44 MiB
```

Planned work:

- [ ] Compare `--spec-draft-n-max 2` with 3 and measure context, acceptance, and tokens/second.
- [ ] Determine the exact base-state and per-snapshot allocation.
- [ ] Investigate configurable GPU rollback depth without changing MTP correctness.
- [ ] Investigate lower-precision rollback snapshots if supported by recurrent-state semantics.
- [ ] Keep the existing behavior as the default until rollback and rejection paths are validated.

Expected benefit from removing one GPU rollback snapshot is approximately 3K-5K context tokens, subject to MTP throughput and acceptance tradeoffs.

### Phase 2: MTP placement and backend overhead

- [ ] Measure free VRAM and MTP tensor placement on both GPUs after long runs.
- [ ] Investigate whether selected NextN tensors or MTP workspace can use the less constrained GPU.
- [ ] Test `--no-spec-draft-backend-sampling` to quantify the ROCm sampling fallback overhead.
- [ ] Separate context-capacity changes from performance-only backend changes.

### Phase 3: Native quantized HIP flash attention

This remains the most valuable long-term backend improvement:

- [ ] Identify the exact HIP flash-attention kernels selected for target and MTP graphs.
- [ ] Measure F16 K/V conversion allocations independently.
- [ ] Prototype direct q8_0 K support.
- [ ] Add q5_1 V and q4_0 MTP support if the first kernel is successful.
- [ ] Validate long-context numerical accuracy and performance against F16.

Native quantized kernels could remove the main reason F16 MTP KV currently uses less total ROCm memory than Q4_0.

---
## Current Patch and Validation Status

- `sol4.patch` applies cleanly to `master` (`7bd8282`).
- The full CPU/default build completed successfully in an isolated detached worktree.
- A HIP build compiled the modified `fattn.cu`, all `fattn-vec` instances, and linked `libggml-hip.so` successfully.
- The final HIP `llama-server` link was not completed because the long build was stopped manually.
- The real `master` checkout was never modified; temporary validation worktrees were removed.
- The VEC result still needs long-prompt, delayed-OOM, stability, throughput, and draft-acceptance testing.

## Scope and Limitations

- The final allocation result is validated for the stated 27B model, ROCm topology, and device split.
- One real prompt completed successfully, but broad stability testing is still in progress.
- The performance improvement is based on one non-controlled comparison and must not yet be treated as a formal benchmark.
- The patch preserves `auto` as the default target pipeline policy.
- F16 is the best validated MTP cache for this ROCm model; it is not assumed to be optimal on Vulkan or other architectures.
- The per-device fitting path applies when GPU layers are explicitly fixed. Automatic placement retains the existing conservative path.

---

## Conclusion

`sol4.patch` aligns MTP estimation with runtime allocation, fits the fixed two-GPU layout against real per-device limits, and makes target pipeline scheduling controllable. Disabling target pipeline parallelism for the tested single-slot workload reduced scheduler copies and compute workspace while increasing fitted context from 104,704 to **119,552 tokens**.

The immediate priority is stability testing with existing prompts and a controlled performance A/B. If those tests remain clean, recurrent-state snapshot memory is the next practical code path to investigate, followed by MTP placement and native quantized HIP flash-attention kernels.