TMP: Throwaway test commit

Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>

Author: KshitijLakhani

tests/jax/test_fused_attn.py

@@ -438,7 +438,6 @@ def _setup_inputs(self):
         self.dp_size = self.mesh.shape.get(self.mesh_resource.dp_resource, 1)
         self.cp_size = self.mesh.shape.get(self.mesh_resource.cp_resource, 1)
         self.tp_size = self.mesh.shape.get(self.mesh_resource.tpsp_resource, 1)
-        breakpoint()
 
         key = jax.random.PRNGKey(0)
         q_key, k_key, v_key, bias_key, dropout_key = jax.random.split(key, 5)
@@ -663,7 +662,6 @@ def generate_random_segment_ids(
                             self.cp_reorder_fn(self.segment_pos_kv),
                         ),
                     )
-                    breakpoint()
                 case _:
                     raise ValueError(f"Unknown {self.seq_desc_format=}")
         else:
@@ -728,7 +726,6 @@ def to_dp_shardings(x):
 
                 self.seq_desc_sharding = jax.tree.map(to_dp_shardings, self.sequence_desciptor)
 
-        #jax.debug.breakpoint()
         if self.bias_shape == BiasShape._1HSS:
             self.bias_pspec = PartitionSpec(
                 None, self.mesh_resource.tpsp_resource, self.mesh_resource.cp_resource, None

tests/jax/utils.py

@@ -1465,7 +1465,6 @@ def assert_allclose(
     mismatch_counter = 0
     has_nonzero = jnp.any(actual != 0)
     print(f"has_nonzero: {has_nonzero}")
-    breakpoint()
     with np.printoptions(threshold=sys.maxsize):
         mismatch_mask = ~np.isclose(actual, desired, **tols) # True means mismatch
         diff_indices = np.argwhere(mismatch_mask)

transformer_engine/common/fused_attn/fused_attn_f16_arbitrary_seqlen.cu

@@ -506,6 +506,7 @@ void fused_attn_arbitrary_seqlen_fwd_impl(
           else
           {
             print_tensor_elements_2<<<1, 1, 0, stream>>>(static_cast<int32_t *>(devPtrCuSeqlensQ), 1, 0, actual_b, /*does not matter for single row*/actual_b);
+            cudaDeviceSynchronize();
           }
         }
         if (devActualSeqlenQ) {
@@ -521,6 +522,7 @@ void fused_attn_arbitrary_seqlen_fwd_impl(
           }
           else {
             print_tensor_elements_2<<<1, 1, 0, stream>>>(static_cast<int32_t *>(devActualSeqlenQ), 1, 0, actual_b, /*does not matter for single row*/actual_b);
+            cudaDeviceSynchronize();
           }
         }
         if(devPtrCuSeqlensKV) {
@@ -538,6 +540,7 @@ void fused_attn_arbitrary_seqlen_fwd_impl(
           }
           else {
             print_tensor_elements_2<<<1, 1, 0, stream>>>(static_cast<int32_t *>(devPtrCuSeqlensKV), 1, 0, actual_b, /*does not matter for single row*/ actual_b);
+            cudaDeviceSynchronize();
           }
         }
         if(devActualSeqlenKV) {
@@ -556,6 +559,7 @@ void fused_attn_arbitrary_seqlen_fwd_impl(
           else
           {
             print_tensor_elements_2<<<1, 1, 0, stream>>>(static_cast<int32_t *>(devActualSeqlenKV), 1, 0, actual_b, /*does not matter for single row*/ actual_b);
+            cudaDeviceSynchronize();
           }
         }
       }
@@ -597,6 +601,7 @@ void fused_attn_arbitrary_seqlen_fwd_impl(
           layout_group, actual_b, b, h, hg, d_qk, d_v, static_cast<int32_t *>(devPtrSeqOffsetsQ),
           static_cast<int32_t *>(devPtrSeqOffsetsKV), ragged_offset_type, devOffsetsQ, devOffsetsK,
           devOffsetsV, devOffsetsO, devOffsetsS);
+      cudaDeviceSynchronize();
       NVTE_CHECK_CUDA(cudaGetLastError());
       if (print_tensors) {
         if (devPtrSeqOffsetsQ) {
@@ -614,6 +619,7 @@ void fused_attn_arbitrary_seqlen_fwd_impl(
             print_tensor_elements_2<<<1, 1, 0, stream>>>(
                 static_cast<int32_t *>(devPtrSeqOffsetsQ), 1, 0, actual_b,
                 /*does not matter for single row*/ actual_b);
+            cudaDeviceSynchronize();
           }
         }
         if (devOffsetsQ) {
@@ -631,6 +637,7 @@ void fused_attn_arbitrary_seqlen_fwd_impl(
             print_tensor_elements_2<<<1, 1, 0, stream>>>(
                 static_cast<int64_t *>(devOffsetsQ), 1, 0, actual_b,
                 /*does not matter for single row*/ actual_b);
+            cudaDeviceSynchronize();
           }
         }
         if (devPtrSeqOffsetsKV) {
@@ -648,6 +655,7 @@ void fused_attn_arbitrary_seqlen_fwd_impl(
             print_tensor_elements_2<<<1, 1, 0, stream>>>(
                 static_cast<int32_t *>(devPtrSeqOffsetsKV), 1, 0, actual_b,
                 /*does not matter for single row*/ actual_b);
+            cudaDeviceSynchronize();
           }
         }
         if (devOffsetsK) {
@@ -665,6 +673,7 @@ void fused_attn_arbitrary_seqlen_fwd_impl(
             print_tensor_elements_2<<<1, 1, 0, stream>>>(
                 static_cast<int64_t *>(devOffsetsK), 1, 0, actual_b,
                 /*does not matter for single row*/ actual_b);
+            cudaDeviceSynchronize();
           }
         }
       }