fix: call _ensure_stacked in warmup test

2026-05-17 08:07:09 +00:00
parent 5ba77e355f
commit 8c9a51e006
1 changed files with 1 additions and 0 deletions
--- 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]