From e86df587b944aa276af707c9f68c918a19dc1204 Mon Sep 17 00:00:00 2001
From: Erik Bogado <erikbogado@gmail.com>
Date: Mon, 31 Aug 2026 20:59:48 -0300
Subject: [PATCH] fix(rocm): avoid indirect host pointers in HIP graphs

---
 python/freetoken/moe/offload_cache.py |   7 +-
 tests/moe/test_fused_copy.py          | 101 +++++++++++++++++++++++++-
 2 files changed, 106 insertions(+), 2 deletions(-)

diff --git a/python/freetoken/moe/offload_cache.py b/python/freetoken/moe/offload_cache.py
index e1f20dd2f..ac9b1bed0 100644
--- a/python/freetoken/moe/offload_cache.py
+++ b/python/freetoken/moe/offload_cache.py
@@ -1001,7 +1001,12 @@ def copy_missing(self) -> None:
             for per_layer, cache in self.banks:
                 cache[: self.num_experts].copy_(per_layer[layer_id])
             return
-        if self._copy_fused_ok:
+        # HIP graphs do not reliably retain the pinned-host mappings hidden behind the
+        # fused kernel's device-side pointer table. Direct per-bank tensor arguments do.
+        use_fused = self._copy_fused_ok and not (
+            torch.version.hip and torch.cuda.is_current_stream_capturing()
+        )
+        if use_fused:
             from freetoken.kernel.fast_index_copy import fast_index_copy_multi_jit
 
             # One launch copies the missing rows for every bank (instead of one launch per
diff --git a/tests/moe/test_fused_copy.py b/tests/moe/test_fused_copy.py
index 7a0c08644..5dba20216 100644
--- a/tests/moe/test_fused_copy.py
+++ b/tests/moe/test_fused_copy.py
@@ -5,9 +5,10 @@
 
 from __future__ import annotations
 
+from types import SimpleNamespace
+
 import pytest
 import torch
-
 from freetoken.moe.offload_cache import _BANK_SCHEMAS, OffloadMoeCache
 
 CUDA = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA")
@@ -19,6 +20,104 @@
 FEATS = [8192, 512, 256, 4096, 512, 256]
 
 
+def test_rocm_graph_capture_uses_per_bank_copy(monkeypatch):
+    calls = []
+    tensor = torch.empty(1)
+    cache = SimpleNamespace(
+        banks=[([tensor], tensor)],
+        _pending_src_layer=0,
+        _pending_whole_layer=False,
+        _unpinned_layers=frozenset(),
+        _copy_fused_ok=True,
+        _copy_dst_ptrs=tensor,
+        _copy_src_ptrs=[tensor],
+        _copy_feat_bytes=tensor,
+        evict_slots=tensor,
+        src_indices=tensor,
+        num_indices=tensor,
+    )
+    monkeypatch.setattr(torch.version, "hip", "7.0", raising=False)
+    monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: True)
+    monkeypatch.setattr(
+        "freetoken.kernel.fast_index_copy.fast_index_copy_multi_jit",
+        lambda *args: calls.append("fused"),
+    )
+    monkeypatch.setattr(
+        "freetoken.kernel.fast_index_copy_jit",
+        lambda *args: calls.append("legacy"),
+    )
+
+    OffloadMoeCache.copy_missing(cache)
+
+    assert calls == ["legacy"]
+
+
+@CUDA
+@pytest.mark.slow
+def test_qwen36_sized_pinned_rows_survive_rocm_graph_replay():
+    if not torch.version.hip:
+        pytest.skip("ROCm regression")
+
+    layers, experts, cache_size, topk = 40, 256, 2117, 8
+    cache = OffloadMoeCache(
+        num_layers=layers,
+        num_experts=experts,
+        cache_size=cache_size,
+        device=torch.device("cuda"),
+        quant_format="q4_k_q5_k",
+    )
+    # Qwen3.6-35B-A3B: H=2048, I=512. Keep the real packed bytes per expert.
+    features = {"gate_up": 2 * 512 * (2048 // 256) * 144,
+                "down": 2048 * (512 // 256) * 176}
+    sources = {}
+    fingerprints = torch.arange(experts, dtype=torch.uint8)
+    for name, feature_bytes in features.items():
+        source = torch.empty((experts, feature_bytes), dtype=torch.uint8, pin_memory=True)
+        source[:, 0] = fingerprints
+        sources[name] = [source] * layers
+    cache.set_bank_sources(sources)
+    assert cache._copy_fused_ok
+
+    patterns = (
+        torch.arange(experts, dtype=torch.int32, device="cuda")[:, None]
+        + torch.arange(topk, dtype=torch.int32, device="cuda")[None, :]
+    ) % experts
+    routes = [patterns[layer * 13 % experts].clone() for layer in range(layers)]
+    slots = [torch.empty_like(route) for route in routes]
+
+    def body():
+        for layer, (route, layer_slots) in enumerate(zip(routes, slots)):
+            layer_slots.copy_(route)
+            cache.ensure_experts(layer, layer_slots)
+            cache.copy_missing()
+
+    side = torch.cuda.Stream()
+    side.wait_stream(torch.cuda.current_stream())
+    with torch.cuda.stream(side):
+        body()
+    side.synchronize()
+    cache.reset()
+    torch.cuda.synchronize()
+
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph, stream=side):
+        body()
+    cache.reset()
+    torch.cuda.synchronize()
+
+    for step in range(512):
+        for layer, route in enumerate(routes):
+            route.copy_(patterns[(step + layer * 13) % experts])
+        graph.replay()
+    torch.cuda.synchronize()
+
+    for route, layer_slots in zip(routes, slots):
+        expected = route.cpu().tolist()
+        resident = layer_slots.cpu().tolist()
+        for _, bank in cache.banks:
+            assert bank[resident, 0].cpu().tolist() == expected
+
+
 def _build_cache(num_layers, num_experts, cache_size):
     dev = torch.device("cuda")
     cache = OffloadMoeCache(