bugfix

sidsingh-nvidia · sidsingh-nvidia · commit bc66c18bee01 · 2025-11-18T12:54:26.000-08:00
diff --git a/megatron/core/ssm/mamba_mixer.py b/megatron/core/ssm/mamba_mixer.py
@@ -486,7 +486,8 @@ def dynamic_inference(self, hidden_states: torch.Tensor, context: DynamicInferen
             zxBCdt_chunked_prefill = zxBCdt[
                 active_token_count - chunked_prefill_request_token_count : active_token_count
             ]
-            batch_index_chunked_prefill = batch_indices[context.chunked_prefill_request_id]
+            pos = torch.where(context.request_ids == context.chunked_prefill_request_id)[0][0]
+            batch_index_chunked_prefill = batch_indices[pos]
 
             y_prefill_chunked = self.ssm_prefill(
                 zxBCdt_chunked_prefill,

Original file line number	Diff line number	Diff line change
`@@ -486,7 +486,8 @@ def dynamic_inference(self, hidden_states: torch.Tensor, context: DynamicInferen`
`486`	`486`	`zxBCdt_chunked_prefill = zxBCdt[`
`487`	`487`	`active_token_count - chunked_prefill_request_token_count : active_token_count`
`488`	`488`	`]`
`489`		`- batch_index_chunked_prefill = batch_indices[context.chunked_prefill_request_id]`
	`489`	`+ pos = torch.where(context.request_ids == context.chunked_prefill_request_id)[0][0]`
	`490`	`+ batch_index_chunked_prefill = batch_indices[pos]`
`490`	`491`
`491`	`492`	`y_prefill_chunked = self.ssm_prefill(`
`492`	`493`	`zxBCdt_chunked_prefill,`