remove context_parallel_strategy config param, todo: add it back later

gulsumgudukbay · gulsumgudukbay · commit c9cd3476ed4c · 2025-12-21T00:13:58.000-06:00
diff --git a/src/MaxText/layers/attention_op.py b/src/MaxText/layers/attention_op.py
@@ -1422,7 +1422,7 @@ def cudnn_flash_attention(
         window_size=sliding_window_size,
         context_parallel_causal_load_balanced=self.config.context_parallel_load_balance,
         context_parallel_axis="context",
-        context_parallel_strategy=self.config.context_parallel_strategy,
+        #context_parallel_strategy=self.config.context_parallel_strategy,
         max_segments_per_seq=max_segments_per_seq,
     )
 

Original file line number	Diff line number	Diff line change
`@@ -1422,7 +1422,7 @@ def cudnn_flash_attention(`
`1422`	`1422`	`window_size=sliding_window_size,`
`1423`	`1423`	`context_parallel_causal_load_balanced=self.config.context_parallel_load_balance,`
`1424`	`1424`	`context_parallel_axis="context",`
`1425`		`- context_parallel_strategy=self.config.context_parallel_strategy,`
	`1425`	`+ #context_parallel_strategy=self.config.context_parallel_strategy,`
`1426`	`1426`	`max_segments_per_seq=max_segments_per_seq,`
`1427`	`1427`	`)`
`1428`	`1428`