[BUG]fix outcache loc (#15)

SiqiLi-Fighting · web-flow · commit 5f768eee4444 · 2025-09-28T15:19:07.000+08:00
diff --git a/python/sgl_jax/srt/speculative/eagle_worker.py b/python/sgl_jax/srt/speculative/eagle_worker.py
@@ -513,9 +513,9 @@ def draft_forward(self, schedule_batch: ScheduleBatch):
         )
         if self.hot_token_id is not None:
             topk_index = self.hot_token_id[topk_index]
-        out_cache_loc = out_cache_loc.reshape(
-            schedule_batch.batch_size(), self.topk, self.speculative_num_steps
-        )
+        out_cache_loc = out_cache_loc[
+            : (schedule_batch.batch_size() * self.topk * self.speculative_num_steps)
+        ].reshape(schedule_batch.batch_size(), self.topk, self.speculative_num_steps)
         out_cache_loc = jnp.transpose(out_cache_loc, (2, 0, 1)).reshape(
             self.speculative_num_steps, -1
         )

Original file line number	Diff line number	Diff line change
`@@ -513,9 +513,9 @@ def draft_forward(self, schedule_batch: ScheduleBatch):`
`513`	`513`	`)`
`514`	`514`	`if self.hot_token_id is not None:`
`515`	`515`	`topk_index = self.hot_token_id[topk_index]`
`516`		`- out_cache_loc = out_cache_loc.reshape(`
`517`		`- schedule_batch.batch_size(), self.topk, self.speculative_num_steps`
`518`		`- )`
	`516`	`+ out_cache_loc = out_cache_loc[`
	`517`	`+ : (schedule_batch.batch_size() * self.topk * self.speculative_num_steps)`
	`518`	`+ ].reshape(schedule_batch.batch_size(), self.topk, self.speculative_num_steps)`
`519`	`519`	`out_cache_loc = jnp.transpose(out_cache_loc, (2, 0, 1)).reshape(`
`520`	`520`	`self.speculative_num_steps, -1`
`521`	`521`	`)`