diff --git a/tests/test_warmup_gs.py b/tests/test_warmup_gs.py
index ef8114c5..6c8b2c6b 100644
--- a/tests/test_warmup_gs.py
+++ b/tests/test_warmup_gs.py
@@ -83,6 +83,7 @@ def warmup_compute_gs(runner, hidden_states, topk_weights, topk_ids):
     top_k = topk_ids.shape[1]
     
     # Build slot mapping (same as runner.run())
+    runner._ensure_stacked()
     flat_ids = topk_ids.reshape(-1)
     num_slots = num_tokens * top_k
     token_indices = runner._token_indices[:num_slots]