[bugfix] fix TBO crashes when attn_tp_size > 1 (#13730)

Co-authored-by: Cheng Wan <54331508+ch-wan@users.noreply.github.com>
2025-12-11 18:18:40 -08:00
parent 0aa3dec5c7
commit e9e7f15eb5
20 changed files with 285 additions and 16 deletions
@@ -395,12 +395,14 @@ class GptOssDecoderLayer(nn.Module):
        self.is_layer_sparse = True
        self.is_nextn = False
        is_previous_layer_sparse = True
+        is_next_layer_sparse = True

        self.layer_scatter_modes = LayerScatterModes.init_new(
            layer_id=layer_id,
            num_layers=config.num_hidden_layers,
            is_layer_sparse=self.is_layer_sparse,
            is_previous_layer_sparse=is_previous_layer_sparse,
+            is_next_layer_sparse=is_next_layer_sparse,
        )

        if self.is_layer_sparse: