remove diffs and use pad_slot_id as var in tests

Author: mzusman

csrc/mamba/causal_conv1d/causal_conv1d.cu

@@ -58,8 +58,8 @@ void set_conv_params_fwd(ConvParamsBase &params,
                          int64_t pad_slot_id,
                          const c10::optional<at::Tensor>& query_start_loc = std::nullopt,
                          const c10::optional<at::Tensor>& cache_indices = std::nullopt,
-                         const c10::optional<at::Tensor>& has_initial_state = std::nullopt
-                         ) {
+                         const c10::optional<at::Tensor>& has_initial_state = std::nullopt) {
+
     // Reset the parameters
     memset(&params, 0, sizeof(params));
 

tests/kernels/test_causal_conv1d.py

@@ -307,7 +307,8 @@ def test_causal_conv1d_update_with_batch_gather(with_padding, dim, width,
                                weight,
                                bias,
                                activation=activation,
-                               conv_state_indices=padded_state_indices)
+                               conv_state_indices=padded_state_indices,
+                               pad_slot_id=PAD_SLOT_ID)
     out_ref = causal_conv1d_update_ref(x_ref[:batch_size],
                                        conv_state_ref,
                                        weight,
@@ -397,7 +398,7 @@ def test_causal_conv1d_varlen(with_padding, dim, seqlen, width, has_bias,
 
     out = causal_conv1d_fn(x.squeeze(0), weight, bias, cumsum.cuda(),
                            padded_state_indices, has_initial_states,
-                           final_states, activation)
+                           final_states, activation, PAD_SLOT_ID)
     out_ref = []
     out_ref_b = []
 

tests/kernels/test_mamba_ssm.py

@@ -594,7 +594,8 @@ def test_selective_state_update_with_batch_indices(with_padding, dim, dstate,
                                  z=z,
                                  dt_bias=dt_bias,
                                  dt_softplus=True,
-                                 state_batch_indices=padded_state_indices)
+                                 state_batch_indices=padded_state_indices,
+                                 pad_slot_id=PAD_SLOT_ID)
     out_ref = selective_state_update_ref(state_ref,
                                          x[:batch_size],
                                          dt[:batch_size],
@@ -694,7 +695,8 @@ def test_selective_state_update_with_heads_with_batch_indices(
                                  z=z,
                                  dt_bias=dt_bias,
                                  dt_softplus=True,
-                                 state_batch_indices=state_indices)
+                                 state_batch_indices=state_indices,
+                                 pad_slot_id=PAD_SLOT_ID)
     out_ref = selective_state_update_ref(state_ref,
                                          x,
                                          dt,