fix(dsv4): upgrade forward metadata on main stream for large PP size (#25729)

2026-05-19 20:52:00 +00:00
parent cd012ada58
commit 8322fe09a7
1 changed files with 4 additions and 0 deletions
@@ -1045,6 +1045,10 @@ class DeepseekV4Model(nn.Module):
                hidden_states = cp_split_and_rebuild_data(forward_batch, hidden_states)
            positions = cp_split_and_rebuild_position(forward_batch, positions)

+        # Upgrade lazy raw metadata on the main stream once before any layer
+        # forks alt-streams; later per-layer calls become no-ops.
+        forward_batch.attn_backend._maybe_upgrade_forward_metadata()
+
        for i in range(self.start_layer, self.end_layer):
            layer = self.layers[i]
            hidden_states = layer(