Fix graph update for MLP with post layernorm

Author: turboderp

exllamav2/exllamav2_ext/cuda/graph.cu

@@ -133,7 +133,7 @@ void Graph::attach_label(cudaStream_t stream, int label, int sublabel)
 }
 
 template <typename T>
-void Graph::update_param(int label, int sublabel, int param, T value)
+void Graph::update_param(int label, int sublabel, int param, T value, bool debug)
 {
     for (int i = 0; i < node_labels.size(); ++i)
     {
@@ -145,19 +145,22 @@ void Graph::update_param(int label, int sublabel, int param, T value)
 
         node_needs_update[i] = true;
 
-//        printf("-----------------------------------------------------\n");
-//        printf("UPDATED:\n");
-//        DBGI(i);
-//        inspect_graph();
+        if (debug)
+        {
+            printf("-----------------------------------------------------\n");
+            printf("UPDATED: ");
+            DBGI(i);
+            inspect_graph();
+        }
     }
 }
 
-void Graph::update_param_ptr(int label, int sublabel, int param, void* value)
+void Graph::update_param_ptr(int label, int sublabel, int param, void* value, bool debug)
 {
-    update_param<void*>(label, sublabel, param, value);
+    update_param<void*>(label, sublabel, param, value, debug);
 }
 
-void Graph::update_param_int(int label, int sublabel, int param, int value)
+void Graph::update_param_int(int label, int sublabel, int param, int value, bool debug)
 {
-    update_param<int>(label, sublabel, param, value);
+    update_param<int>(label, sublabel, param, value, debug);
 }

exllamav2/exllamav2_ext/cuda/graph.cuh

@@ -46,10 +46,10 @@ public:
     void attach_label(cudaStream_t stream, int label, int sublabel);
 
     template <typename T>
-    void update_param(int label, int sublabel, int param, T value);
+    void update_param(int label, int sublabel, int param, T value, bool debug);
 
-    void update_param_ptr(int label, int sublabel, int param, void* value);
-    void update_param_int(int label, int sublabel, int param, int value);
+    void update_param_ptr(int label, int sublabel, int param, void* value, bool debug = false);
+    void update_param_int(int label, int sublabel, int param, int value, bool debug = false);
 };
 
 

exllamav2/exllamav2_ext/cuda/q_mlp.cu

@@ -109,7 +109,7 @@ void QMLP::forward_
     if (graph->count())
     {
         graph->begin_capture(stream);
-        forward_run_(stream, cublas_handle, (half*) x, rows, columns, loras, lora_temp, graph);
+        forward_run_(stream, cublas_handle, (void*) x, rows, columns, loras, lora_temp, graph);
         graph->end_capture(stream);
 //        printf("**** record ****\n");
 //        DBGI2(rows, columns);
@@ -225,7 +225,7 @@ void QMLP::forward_run_
 
     else
     {
-        gemm_half_q_half_cuda(stream, cublas_handle, temp_a, down, temp_state, rows, columns, intermediate_size, true, temp_dq, graph, 0);
+        gemm_half_q_half_cuda(stream, cublas_handle, temp_a, down, temp_state, rows, columns, intermediate_size, true, temp_dq, false, NULL, 0, false, graph, 0);
         if (layernorm_is_rms)
             rms_norm_cuda(stream, temp_state, post_layernorm, x, norm_epsilon, rows, columns, true, false, residual_fp32, graph, KernelLabels::POST_NORM);
         else