Avoid CUDA stream sync (#40060)

cyyever · web-flow · commit ec85d2c44f0c · 2025-08-15T12:37:15.000Z
Signed-off-by: cyy &lt;cyyever@outlook.com&gt;
diff --git a/src/transformers/modeling_flash_attention_utils.py b/src/transformers/modeling_flash_attention_utils.py
@@ -354,12 +354,12 @@ def prepare_fa_kwargs_from_position_ids(position_ids, is_packed_sequence: bool =
         max_length_q = int(q_len.max())
         max_length_k = int(last_position_ids.max()) + 1
     else:
-        position_ids = position_ids.flatten()
-        indices_q = torch.arange(position_ids.size(0), device=position_ids.device, dtype=torch.int32)
+        position_ids = position_ids.view(-1)
+        indices_q = (position_ids == 0).nonzero().view(-1)
 
         cu_seq_lens_q = torch.cat(
             (
-                indices_q[position_ids == 0],
+                indices_q,
                 torch.tensor(position_ids.size(), device=position_ids.device, dtype=torch.int32),
             )
         )

Original file line number	Diff line number	Diff line change
`@@ -354,12 +354,12 @@ def prepare_fa_kwargs_from_position_ids(position_ids, is_packed_sequence: bool =`
`354`	`354`	`max_length_q = int(q_len.max())`
`355`	`355`	`max_length_k = int(last_position_ids.max()) + 1`
`356`	`356`	`else:`
`357`		`- position_ids = position_ids.flatten()`
`358`		`- indices_q = torch.arange(position_ids.size(0), device=position_ids.device, dtype=torch.int32)`
	`357`	`+ position_ids = position_ids.view(-1)`
	`358`	`+ indices_q = (position_ids == 0).nonzero().view(-1)`
`359`	`359`
`360`	`360`	`cu_seq_lens_q = torch.cat(`
`361`	`361`	`(`
`362`		`- indices_q[position_ids == 0],`
	`362`	`+ indices_q,`
`363`	`363`	`torch.tensor(position_ids.size(), device=position_ids.device, dtype=torch.int32),`
`364`	`364`	`)`
`365`	`365`	`)`