Cache tensors on GPU in DAGs

[lantiga]

At the current stage, RedisAI keeps tensors "at rest" on the CPU and transfers them to GPU when needed in the context of a run (while model weights stay on the GPU).

This feature would allow DAGs to keep intermediate tensors on the GPU, if possible. The possibility and opportunity to do so (e.g. tensor produced and later directly consumed by a TORCH model located on the same GPU) can be determined by analyzing the DAG with no additions to the current syntax.

The advantage of limiting this to the DAG execution context is that we allow optimizations without incurring in issues for invalidating the GPU cache. Since GPU memory is precious, we want to make sure RedisAI is not caching all tensors on the GPU.

Another area where we could cache tensors on the GPU is optimizing access to reference data (even outside DAGs). For this we could add a flag to TENSORSET that explicitly requests copying the tensor to the GPU and keeping it there. The user would be responsible for removing it from GPU at a later stage.

/cc @gkorland

[Fawzikhattar]

Hello I'm interested in the feature of adding a flag to TENSORSET to create data directly on GPU any update on this?

[lantiga]

Hi @Fawzikhattar, this is scheduled to happen after the 1.2.0 release. Feel free to describe your use case, so we can take it into account during the design.

[Fawzikhattar]

Hello @lantiga , thank you for your reply.

My use case is very simple:

I would like to store N tensors each with dimension D=512 using redisAI, let's say N=30000 tensors. At a certain point in time I will receive a new set of tensors (around 10).

I would like to compute the distance of every tensor in the new set to a subset (or the whole set) of the N tensors already in redis.

The distance calculation speed can be significantly improved if the computation happens on GPU. However today the overhead from copying huge tensors from CPU to GPU memory is making it actually slower than a version running on CPU (which is already slow for my objective).

After this distance calculation I would like to delete some tensors in redis based on some conditions or insert new tensors directly on GPU. Currently I only tested the distance calculation using torchscript.

#def distance(y,args : List[Tensor]):
def distance(y,x):
#print(args)
#x=torch.stack(args)
r1 = torch.sum(x* x,1,keepdim=False)
r2 = torch.sum(y* y,1,keepdim=False)
r1 = torch.reshape(r1, [-1, 1])
r2 = torch.reshape(r2, [1, -1])
D = r1 - 2*torch.matmul(x, torch.transpose(y,0,1)) + torch.repeat_interleave(r2,repeats=r1.size()[0],dim=0)
D=torch.sqrt(D)
return D