From 8c9a51e006c090e3cc7fadbff83a56b0df4c472e Mon Sep 17 00:00:00 2001 From: biondizzle Date: Sun, 17 May 2026 08:07:09 +0000 Subject: [PATCH] fix: call _ensure_stacked in warmup test --- tests/test_warmup_gs.py | 1 + 1 file changed, 1 insertion(+) 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]