[LLM serving] Add Mamba/Jamba model kernel (#3576)

Author: blzheng

csrc/cpu/aten/Conv.cpp

@@ -9,6 +9,7 @@ namespace torch_ipex {
 namespace cpu {
 
 IPEX_DEFINE_DISPATCH(causal_conv1d_update_kernel_stub);
+IPEX_DEFINE_DISPATCH(causal_conv1d_fn_kernel_stub);
 std::vector<int64_t> calc_conv_output_size(
  at::IntArrayRef input_size,
  at::IntArrayRef kernel_size,
@@ -514,21 +515,55 @@ at::Tensor convolution_forward(
  * @param conv_weights (dim, width)
  * @param conv_bias (dim,)
  * @param silu_activation If true, apply the SiLU activation function.
+ * @param cache_seqlens (batch,) or None
  * @return (hidden_states, conv_states)
  */
 std::tuple<at::Tensor, at::Tensor> causal_conv1d_update(
  const at::Tensor& hidden_states,
  const at::Tensor& conv_states,
  const at::Tensor& conv_weights,
  const c10::optional<at::Tensor>& conv_bias,
- bool silu_activation) {
+ bool silu_activation,
+ const c10::optional<at::Tensor>& cache_seqlens) {
  RECORD_FUNCTION("causal_conv1d_update", c10::ArrayRef<c10::IValue>({}));
  return causal_conv1d_update_kernel_stub(
  kCPU,
  hidden_states,
  conv_states,
  conv_weights,
  conv_bias,
+ silu_activation,
+ cache_seqlens);
+}
+
+/**
+ * Official Python implementation: causal_conv1d_ref:
+ * https://github.com/Dao-AILab/causal-conv1d/blob/main/causal_conv1d/causal_conv1d_interface.py#L133
+ * @param x (batch, dim, seqlen)
+ * @param conv_weights (dim, width)
+ * @param conv_bias (dim,)
+ * @param initial_states (batch, dim, width - 1)
+ * @param final_states_out (batch, dim, width - 1)
+ * @param silu_activation If true, apply the SiLU activation function.
+ * @return (out, final_states_out)
+ * out: (batch, dim, seqlen)
+ * final_states_out: (batch, dim, width - 1)
+ */
+std::tuple<at::Tensor, at::Tensor> causal_conv1d_fn(
+ const at::Tensor& x,
+ const at::Tensor& conv_weights,
+ const c10::optional<at::Tensor>& conv_bias,
+ const c10::optional<at::Tensor>& initial_states,
+ const c10::optional<at::Tensor>& final_states_out,
+ bool silu_activation) {
+ RECORD_FUNCTION("causal_conv1d_fn", c10::ArrayRef<c10::IValue>({}));
+ return causal_conv1d_fn_kernel_stub(
+ kCPU,
+ x,
+ conv_weights,
+ conv_bias,
+ initial_states,
+ final_states_out,
  silu_activation);
 }
 
@@ -589,11 +624,17 @@ TORCH_LIBRARY_FRAGMENT(torch_ipex, m) {
  c10::DispatchKey::CPU,
  torch_ipex::cpu::convolution_forward_impl);
  m.def(
- "causal_conv1d_update(Tensor hidden_states, Tensor conv_states, Tensor conv_weights, Tensor? conv_bias, bool silu_activation) -> (Tensor, Tensor)");
+ "causal_conv1d_update(Tensor hidden_states, Tensor conv_states, Tensor conv_weights, Tensor? conv_bias, bool silu_activation, Tensor? cache_seqlens=None) -> (Tensor, Tensor)");
  m.impl(
  "causal_conv1d_update",
  c10::DispatchKey::CPU,
  torch_ipex::cpu::causal_conv1d_update);
+ m.def(
+ "causal_conv1d_fn(Tensor x, Tensor conv_weights, Tensor? conv_bias, Tensor? initial_states, Tensor? final_states_out, bool silu_activation) -> (Tensor, Tensor)");
+ m.impl(
+ "causal_conv1d_fn",
+ c10::DispatchKey::CPU,
+ torch_ipex::cpu::causal_conv1d_fn);
  // bw
  m.def(
  "convolution_backward(Tensor input, Tensor weight, Tensor? bias, Tensor grad_output, bool[3] out_mask, "

csrc/cpu/aten/Conv.h

@@ -57,6 +57,15 @@ std::tuple<at::Tensor, at::Tensor> causal_conv1d_update(
  const at::Tensor& conv_states,
  const at::Tensor& conv_weights,
  const c10::optional<at::Tensor>& conv_bias,
+ bool silu_activation,
+ const c10::optional<at::Tensor>& cache_seqlens);
+
+std::tuple<at::Tensor, at::Tensor> causal_conv1d_fn(
+ const at::Tensor& x,
+ const at::Tensor& conv_weights,
+ const c10::optional<at::Tensor>& conv_bias,
+ const c10::optional<at::Tensor>& initial_states,
+ const c10::optional<at::Tensor>& final_states_out,
  bool silu_activation);
 
 // IPEX customized convolution OP with n-D packed weight
@@ -108,9 +117,18 @@ using causal_conv1d_update_kernel_fn = std::tuple<at::Tensor, at::Tensor> (*)(
  const at::Tensor& conv_states,
  const at::Tensor& conv_weights,
  const c10::optional<at::Tensor>& conv_bias,
+ bool silu_activation,
+ const c10::optional<at::Tensor>& cache_seqlens);
+using causal_conv1d_fn_kernel_fn = std::tuple<at::Tensor, at::Tensor> (*)(
+ const at::Tensor& x,
+ const at::Tensor& conv_weights,
+ const c10::optional<at::Tensor>& conv_bias,
+ const c10::optional<at::Tensor>& initial_states,
+ const c10::optional<at::Tensor>& final_states_out,
  bool silu_activation);
 IPEX_DECLARE_DISPATCH(
  causal_conv1d_update_kernel_fn,
  causal_conv1d_update_kernel_stub);
+IPEX_DECLARE_DISPATCH(causal_conv1d_fn_kernel_fn, causal_conv1d_fn_kernel_stub);
 } // namespace cpu
 } // namespace torch_ipex