fixing scratch buffer index

Azure · Dec 3, 2024 · 5f38d15 · 5f38d15
1 parent 02c2a46
commit 5f38d15
Showing 1 changed file with 1 addition and 5 deletions.
diff --git a/msccl/language/mscclpp/instruction_dag.py b/msccl/language/mscclpp/instruction_dag.py
@@ -487,11 +487,7 @@ def is_scratch(buffer):
             return buffer != Buffer.input and buffer != Buffer.output
 
         def get_new_index(rank, buffer, index, size, i):
-            # Scratch buffers always use batched
-            if is_scratch(buffer):
-                buf_instance_len = self.buffers[rank][buffer].instance_size()
-                return buf_instance_len * i + index
-            elif replication_policy == ReplicationPolicy.interleaved:
+            if replication_policy == ReplicationPolicy.interleaved:
                 return index * instances + i * size
             return len(self.buffers[rank][buffer]) * i + index