#0: fixed idle core assignment mismatch for sdpa decode in sharded lo…

…w batch case
tenstorrent · Nov 7, 2024 · e601275 · e601275
1 parent 62255a8
commit e601275
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 3 deletions.
diff --git a/...tt_eager/python_api_testing/unit_testing/misc/test_scaled_dot_product_attention_decode.py b/...tt_eager/python_api_testing/unit_testing/misc/test_scaled_dot_product_attention_decode.py
@@ -886,7 +886,10 @@ def test_sdpa_decode_paged_attention(
 )
 @pytest.mark.parametrize(
     "b, nh, nkv, s, d, grid_size",
-    ([16, 8, 1, 32768, 128, (8, 6)],),  # Llama2-70B
+    (
+        [1, 8, 1, 32768, 128, (8, 8)],
+        [16, 8, 1, 32768, 128, (8, 8)],
+    ),  # Llama2-70B
 )
 def test_sdpa_decode_sharded(device, b, nh, nkv, s, d, dtype, grid_size, q_dtype):
     ttnn.device.DisablePersistentKernelCache()

diff --git a/ttnn/cpp/ttnn/operations/transformer/sdpa_decode/device/sdpa_decode_program_factory.cpp b/ttnn/cpp/ttnn/operations/transformer/sdpa_decode/device/sdpa_decode_program_factory.cpp
@@ -148,7 +148,7 @@ operation::ProgramWithCallbacks sdpa_decode_multi_core(
 
         for (int i = 0; i < num_cores_available; ++i) {
             CoreCoord core;
-            if (i%num_cores_per_batch==0){
+            if (i%num_cores_per_batch==0 && reducer_idx<num_output_cores){
                 core = {reducer_idx % grid_size.x, reducer_idx / grid_size.x};
                 reducer_idx++;
             }
@@ -182,6 +182,7 @@ operation::ProgramWithCallbacks sdpa_decode_multi_core(
     log_debug("num_reducer_cores: {}", num_reducer_cores);
     log_debug("num_output_cores: {}", num_output_cores);
     log_debug("core_group: {}", core_group);
+    log_debug("core_group_idle: {}", core_group_idle);
 
     // These tile capacity counts for CBs need to match the number of tiles expected by the kernel (softmax.cpp)
     uint32_t q_tiles  = PNHt * DHt;
@@ -630,7 +631,7 @@ operation::ProgramWithCallbacks sdpa_decode_multi_core(
             CoreCoord core = core_group_idle[i];
             log_debug("Setting core {} to idle", core);
             // reader runtime args
-            std::vector<uint32_t> reader_rt_args = { 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0};
+            std::vector<uint32_t> reader_rt_args = { 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0};
 
             // writer runtime args
             std::vector<uint32_t> writer_rt_args = { 0, 0, 0, 0, 0, 0, 0, 0, 0, 0};