Fix reduce_scatter_minimial_async merge conflict bugs

Author: djordje-tt

models/tt_transformers/tt/generator.py

@@ -412,7 +412,6 @@ def prefill_forward_text(
                 local_kwargs["pixel_values"] = local_kwargs["pixel_values"][idx]
                 if "image_grid_thw" in local_kwargs:
                     local_kwargs["image_grid_thw"] = local_kwargs["image_grid_thw"][idx]
-            local_kwargs["user_id"] = user_id
 
             if sampling_enabled:
                 sampling_executed = True
@@ -511,11 +510,11 @@ def prefill_forward_text(
                 last_token_idx_relative = last_token_idx - num_cached_tokens
                 ttnn.synchronize_device(self.model[model_id].mesh_device)
 
-                if res["hidden_states"] is not None:
+                if "hidden_states" in res:
                     output_tensor[idx] = self.model[model_id].process_output_prefill_hidden_states(
                         res["hidden_states"], last_token_idx=(last_token_idx_relative % 32)
                     )
-                elif res["logits"] is not None:
+                elif res["sampling"]:
                     tt_tokens = res["logits"][0]
                     tt_log_probs = res["logits"][1]
                     tokens_host = ttnn.to_torch(ttnn.get_device_tensors(tt_tokens)[0]).reshape(-1)[

models/tt_transformers/tt/model_config.py

@@ -1338,6 +1338,10 @@ def _get_xattn_kv_prefill_mem_cfg(seq_len):
             self.set_tg_attention_config()
 
             self.is_multichip = self.num_devices > 1
+            self.num_reduce_scatter_links = 1
+            self.num_all_gather_links = (
+                2 if self.is_galaxy else 1
+            )  # TODO: try out 3 for short axis and 4 for long axis (TG only) <- should work but untested in model
             self.ccl_dtype = ttnn.bfloat8_b
 
             # model specific CCL configs

ttnn/cpp/ttnn/operations/experimental/ccl/reduce_scatter_minimal_async/device/reduce_scatter_minimal_async_program.cpp

@@ -959,8 +959,9 @@ ReduceScatterProgramArtifacts build_ring_reduce_scatter_minimal_async_program_ar
                         src_node_id, dst_node_id, link, program, {mux_logical_core});
                 }
                 tt::tt_metal::SetRuntimeArgs(program, mux_kernel_id, {mux_logical_core}, mux_rt_args);
+
+                termination_master_logical_core = *((termination_master_core_iter++)->begin());
             }
-            termination_master_logical_core = *((termination_master_core_iter++)->begin());
 
             for (uint32_t worker = 0; worker < num_workers_per_direction; worker++) {
                 auto core = *((worker_core_iter++)->begin());
@@ -1046,41 +1047,42 @@ ReduceScatterProgramArtifacts build_ring_reduce_scatter_minimal_async_program_ar
                         termination_master_virtual_core,
                         program,
                         writer_rt_args);
-                    if (intermediate_is_sharded) {
-                        shard_builder::extend_sharding_run_time_args(intermediate_tensor, writer_rt_args);
-                    }
-                    if (output_is_sharded) {
-                        shard_builder::extend_sharding_run_time_args(output_tensor, writer_rt_args);
-                    }
-                    if (!num_mux_cores_per_direction_per_link) {
-                        if (dir) {  // forward
-                            writer_rt_args.push_back(forward_coord.has_value());
-                            if (forward_coord.has_value()) {
-                                const auto src_fabric_node_id = mesh_device->get_fabric_node_id(sender_device_coord);
-                                const auto dst_fabric_node_id = mesh_device->get_fabric_node_id(forward_coord.value());
-                                tt::tt_fabric::append_fabric_connection_rt_args(
-                                    src_fabric_node_id, dst_fabric_node_id, link, program, {core}, writer_rt_args);
-                            }
-                            writer_rt_args.push_back(false);
-                        } else {
-                            writer_rt_args.push_back(false);
-                            writer_rt_args.push_back(backward_coord.has_value());
-                            if (backward_coord.has_value()) {
-                                const auto src_fabric_node_id = mesh_device->get_fabric_node_id(sender_device_coord);
-                                const auto dst_fabric_node_id = mesh_device->get_fabric_node_id(backward_coord.value());
-                                tt::tt_fabric::append_fabric_connection_rt_args(
-                                    src_fabric_node_id, dst_fabric_node_id, link, program, {core}, writer_rt_args);
-                            }
+                }
+                if (intermediate_is_sharded) {
+                    shard_builder::extend_sharding_run_time_args(intermediate_tensor, writer_rt_args);
+                }
+                if (output_is_sharded) {
+                    shard_builder::extend_sharding_run_time_args(output_tensor, writer_rt_args);
+                }
+                if (!num_mux_cores_per_direction_per_link) {
+                    if (dir) {  // forward
+                        writer_rt_args.push_back(forward_coord.has_value());
+                        if (forward_coord.has_value()) {
+                            const auto src_fabric_node_id = mesh_device->get_fabric_node_id(sender_device_coord);
+                            const auto dst_fabric_node_id = mesh_device->get_fabric_node_id(forward_coord.value());
+                            tt::tt_fabric::append_fabric_connection_rt_args(
+                                src_fabric_node_id, dst_fabric_node_id, link, program, {core}, writer_rt_args);
+                        }
+                        writer_rt_args.push_back(false);
+                    } else {
+                        writer_rt_args.push_back(false);
+                        writer_rt_args.push_back(backward_coord.has_value());
+                        if (backward_coord.has_value()) {
+                            const auto src_fabric_node_id = mesh_device->get_fabric_node_id(sender_device_coord);
+                            const auto dst_fabric_node_id = mesh_device->get_fabric_node_id(backward_coord.value());
+                            tt::tt_fabric::append_fabric_connection_rt_args(
+                                src_fabric_node_id, dst_fabric_node_id, link, program, {core}, writer_rt_args);
                         }
                     }
-                    tt::tt_metal::SetRuntimeArgs(program, writer_kernel_id, {core}, writer_rt_args);
-
-                    std::vector<uint32_t> reduce_rt_args = {
-                        start_tiles_read,     // start_tiles_read
-                        start_tiles_to_read,  // start_tiles_to_read
-                        dir};                 // dir
-                    tt::tt_metal::SetRuntimeArgs(program, sender_reduce_kernel_id, {core}, reduce_rt_args);
                 }
+                tt::tt_metal::SetRuntimeArgs(program, writer_kernel_id, {core}, writer_rt_args);
+
+                std::vector<uint32_t> reduce_rt_args = {
+                    start_tiles_read,     // start_tiles_read
+                    start_tiles_to_read,  // start_tiles_to_read
+                    dir};                 // dir
+                tt::tt_metal::SetRuntimeArgs(program, sender_reduce_kernel_id, {core}, reduce_rt_args);
+            }
         }
     }