Evaluate the need for state-stack

[JoelMathewC]

The need for state-stack was evaluated earlier in #3 but now with this documentation it turns out that we can use pytorch to handle storing backward tensors. In this discussion we'll verify if that approach works.

[JoelMathewC]

save-for-backward version

Successfully replaced the need for state-stack with pytorch save-for-backward

class KernelWrapper():
    def __init__(self):
        pass
    
    @staticmethod
    def forward(executor, kid, kernel_args, rets, bwd_tensors, *args):
        ret = executor.forward_cb(kid, kernel_args, rets, bwd_tensors, args)
        return ret

    @staticmethod
    def setup_context(ctx, inputs, output):
        # Saving executor and kid in backward cache
        ctx.save_for_backward(*inputs[4])
        ctx.backward_cache = inputs[0], inputs[1]
        ctx.set_materialize_grads(False)

    @staticmethod
    def backward(ctx, *gradout):
        saved_tensors = ctx.saved_tensors
        executor, kid = ctx.backward_cache
        return (None, None, None, None, None) + executor.backward_cb(kid, gradout, saved_tensors)

I've verified the correctness and both approaches seem to give same memory and speed benchmarks.

[JoelMathewC]

Design Decision

Due to the fact that Seastar is backend-agnostic it is best that we continue to use the state-stack and timestamp-stack. Essentially the design decision is to make the executor temporally aware and hence does not need to dependent on the backed for anything other than the input tensors during forward propagation and gradients during backward propagation.