%0:2 = stablehlo.composite "odml.update_kv_cache" %arg0, %arg1, %arg2, %arg3, %arg4 {composite_attributes = {kv_cache_max = 500 : i64}, decomposition = @odml.update_kv_cache.impl_0} : (tensor<1x500x4x4xf32>, tensor<1x500x4x4xf32>, tensor<100xi64>, tensor<1x100x4x4xf32>, tensor<1x100x4x4xf32>) -> (tensor<1x500x4x4xf32>, tensor<1x500x4x4xf32>)

  // List of supported composites to represent using CustomOp.
  return llvm::is_contained(
      {"odml.update_kv_cache", "odml.scaled_dot_product_attention"},
      op.getName());
}

bool IsKVCacheCompositeOp(::mlir::stablehlo::CompositeOp op) {
  return op.getName() == "odml.update_kv_cache";
}

TFL::ConstBytesAttr CustomOption(OpBuilder* builder,
                                 const std::string& content) {