class RWKV7(tf.keras.layers.Layer):
    def __init__(self, config, layer_id, **kwargs):
        super(RWKV7, self).__init__(**kwargs)
        self.config = config
        self.layer_id = layer_id
        self.n_embd = config.n_embd
        self.dim_att = config.n_embd

        self.head_size = HEAD_SIZE
        self.n_head = self.dim_att // self.head_size
        assert self.dim_att % self.n_head == 0, "dim_att must be divisible by n_head"

        # Initialize parameters
        ratio_0_to_1 = layer_id / (config.n_layer - 1)
        ratio_1_to_almost0 = 1.0 - (layer_id / config.n_layer)

        # Define time_maa_* parameters with initializations similar to PyTorch

        # Initialize decay_speed

        # Define other parameters (e.g., time_faaaa, time_aaaaa, etc.)

        # Define LORA-like projections (time_maa_w1, time_maa_w2, etc.)
        D_MIX_LORA = 32
        self.time_maa_w1 = self.add_weight(
            shape=(self.n_embd, D_MIX_LORA * 4),
            initializer='zeros',
            trainable=True,
            name='time_maa_w1'
        )
        self.time_maa_w2 = self.add_weight(
            shape=(4, D_MIX_LORA, self.n_embd),
            initializer=tf.keras.initializers.Orthogonal(gain=math.sqrt(self.n_embd / (D_MIX_LORA * 4)) * 0.1),
            trainable=True,
            name='time_maa_w2'
        )
        # Initialize other LORA-like parameters (time_decay_w1, time_decay_w2, etc.) similarly...

        # Define learnable transformations
        self.receptance = layers.Dense(self.dim_att, use_bias=False, name='receptance')
        self.key = layers.Dense(self.dim_att, use_bias=False, name='key')
        self.value = layers.Dense(self.dim_att, use_bias=False, name='value')
        self.output_dense = layers.Dense(self.n_embd, use_bias=False, name='output')

        # Group Normalization
        self.ln_x = layers.GroupNormalization(groups=self.n_head, epsilon=1e-5, name='ln_x')

    def call(self, x, training=False):
        B, T, C = tf.shape(x)[0], tf.shape(x)[1], tf.shape(x)[2]
        H = self.n_head
        N = self.head_size

        # Time shift: equivalent to ZeroPad2d((0,0,1,-1)) in PyTorch
        # Replace that with slicing and concatenation to achieve the time shift
        x_shifted = tf.concat([tf.zeros_like(x[:, :1, :]), x[:, :-1, :]], axis=1)

        # Apply time_maa_* parameters
        xxx = x + x_shifted * self.time_maa_x
        xxx = tf.tanh(tf.matmul(xxx, self.time_maa_w1))
        xxx = tf.reshape(xxx, [B * T, 4, -1])
        xxx = tf.transpose(xxx, perm=[1, 0, 2])  # Shape: [4, B*T, ...]
        xxx = tf.reshape(xxx, [4, B, T, -1])
        mrg, mwa, mk, mv = tf.unstack(xxx, axis=0)  # Each has shape [B, T, ...]

        tf.print(tf.shape(mrg), tf.shape(x), tf.shape(x_shifted))
        xrg = x + x_shifted * (self.time_maa_rg + mrg)
        xwa = x + x_shifted * (self.time_maa_rg + mwa)  # Assuming similar naming
        xk = x + x_shifted * (self.time_maa_k + mk)
        xv = x + x_shifted * (self.time_maa_v + mv)

        # Compute receptive parameters
        r = self.receptance(xrg)  # Shape: [B, T, C]
        w = -tf.nn.softplus(-(self.time_decay + tf.matmul(xwa, self.time_decay_w1) @ self.time_decay_w2)) - 0.5
        k = self.key(xk)  # Shape: [B, T, C]
        v = self.value(xv)  # Shape: [B, T, C]
        g = tf.nn.tanh(tf.matmul(xrg, self.gate_w1)) @ self.gate_w2

        # Normalize keys
        kk = k + tf.matmul(tf.nn.tanh(tf.matmul(xk, self.time_kkk_w1)), self.time_kkk_w2)
        kk = tf.nn.l2_normalize(tf.reshape(kk, [B, T, H, -1]), axis=-1)
        kk = tf.reshape(kk, [B, T, C])

        a = tf.sigmoid(self.time_aaaaa + tf.matmul(tf.matmul(xwa, self.time_aaa_w1), self.time_aaa_w2))

        # Mix Keys
        ma = tf.sigmoid(self.time_misc_a + tf.matmul(tf.matmul(xwa, self.ma_w1), self.ma_w2))
        k = k * ma + k * a * (1 - ma)
        mk = tf.sigmoid(self.time_misc_k + tf.matmul(tf.matmul(xk, self.mk_w1), self.mk_w2))
        k = k * tf.clip_by_value(tf.exp(-w * mk), 0, None)  # Similar to clamp in PyTorch

        # WKV operation: To be implemented as per model's requirement
        # Placeholder for RUN_CUDA_RWKV7g equivalent
        # Here we implement a simplified version using TensorFlow operations

        # Note: The actual RUN_CUDA_RWKV7g performs efficient WKV operations using custom CUDA
        # For TensorFlow, we need to implement WKV as a custom layer or using tf.scan

        # Simplified WKV implementation using tf.scan
        def wkv_step(prev_state, inputs):
            # Inputs: [C]
            # prev_state: [C]
            # Compute state and output
            new_state = prev_state * tf.math.exp(-tf.math.exp(w)) + k * inputs + a * tf.reduce_sum(a * prev_state)
            y = tf.reduce_sum(new_state * r)
            return new_state, y

        # Initialize state
        initial_state = tf.zeros([B, C], dtype=x.dtype)
        # Iterate over T steps
        states, outputs = tf.scan(
            fn=wkv_step,
            elems=tf.transpose(v, perm=[1, 0, 2]),  # Shape: [T, B, C]
            initializer=(initial_state, tf.zeros([B, C], dtype=x.dtype))
        )
        # outputs: [T, B, C], transpose back
        outputs = tf.transpose(outputs, perm=[1, 0, 2])  # Shape: [B, T, C]

        # Apply Group Normalization
        outputs = self.ln_x(outputs)

        # Compute final output
        x = self.output_dense(outputs * g)
        return x


Update and fix the above tensorflow model to match its pytorch code:


#include <torch/extension.h>
#include "ATen/ATen.h"

typedef at::BFloat16 bf16;

void cuda_forward(int B, int T, int C, int H, bf16 *r, bf16 *w, bf16 *k, bf16 *v, bf16 *a, bf16 *b, bf16 *y, float *saa, float* sss);
void cuda_backward(int B, int T, int C, int H, bf16 *r, bf16 *w, bf16 *k, bf16 *v, bf16 *a, bf16 *b, float *saa, float* sss, float* zzz, bf16 *gy, bf16 *gr, bf16 *gw, bf16 *gk, bf16 *gv, bf16 *ga, bf16 *gb);

void forward(int64_t B, int64_t T, int64_t C, int64_t H, torch::Tensor &r, torch::Tensor &w, torch::Tensor &k, torch::Tensor &v, torch::Tensor &a, torch::Tensor &b, torch::Tensor &y, torch::Tensor &saa, torch::Tensor &sss) {
    cuda_forward(B, T, C, H, r.data_ptr<bf16>(), w.data_ptr<bf16>(), k.data_ptr<bf16>(), v.data_ptr<bf16>(), a.data_ptr<bf16>(), b.data_ptr<bf16>(), y.data_ptr<bf16>(), saa.data_ptr<float>(), sss.data_ptr<float>());
}
void backward(int64_t B, int64_t T, int64_t C, int64_t H, torch::Tensor &r, torch::Tensor &w, torch::Tensor &k, torch::Tensor &v, torch::Tensor &a, torch::Tensor &b, torch::Tensor &saa, torch::Tensor &sss, torch::Tensor &zzz, torch::Tensor &gy, torch::Tensor &gr, torch::Tensor &gw, torch::Tensor &gk, torch::Tensor &gv, torch::Tensor &ga, torch::Tensor &gb) {
    cuda_backward(B, T, C, H, r.data_ptr<bf16>(), w.data_ptr<bf16>(), k.data_ptr<bf16>(), v.data_ptr<bf16>(), a.data_ptr<bf16>(), b.data_ptr<bf16>(), saa.data_ptr<float>(), sss.data_ptr<float>(), zzz.data_ptr<float>(), gy.data_ptr<bf16>(), gr.data_ptr<bf16>(), gw.data_ptr<bf16>(), gk.data_ptr<bf16>(), gv.data_ptr<bf16>(), ga.data_ptr<bf16>(), gb.data_ptr<bf16>());
}

TORCH_LIBRARY(wkv7g, m) {
    m.def("forward", forward);
    m.def("backward", backward);
}

#include <stdio.h>
#include <assert.h>
#include "ATen/ATen.h"

typedef at::BFloat16 bf16;

template <typename F>
__global__ void kernel_forward(const int B, const int T, const int C, const int H,
                               const F *__restrict__ const _r, const F *__restrict__ const _w, const F *__restrict__ const _k, const F *__restrict__ const _v, const F *__restrict__ const _a, const F *__restrict__ const _b,
                               F *__restrict__ const _y, float *__restrict__ const _saa, float *__restrict__ const _sss)
{
    const int e = blockIdx.x / H;
    const int h = blockIdx.x % H;
    const int i = threadIdx.x;

    float state[_N_] = {0};
    __shared__ float r[_N_], k[_N_], w[_N_], a[_N_], b[_N_];

    float v[_T_];
    for (int _t = 0; _t < T; _t++)
    {
        const int t = e*T*C + h*_N_ + i + _t * C;
        v[_t] = float(_v[t]);
    }

    for (int _t = 0; _t < T; _t++)
    {
        const int t = e*T*C + h*_N_ + i + _t * C;
        __syncthreads();
        r[i] = float(_r[t]);
        w[i] = __expf(-__expf(float(_w[t])));
        k[i] = float(_k[t]);
        a[i] = float(_a[t]);
        b[i] = float(_b[t]);
        __syncthreads();

        float sa = 0;
        #pragma unroll
        for (int j = 0; j < _N_; j++)
        {
            sa += a[j] * state[j];
        }
        _saa[t] = float(sa);

        float vv = v[_t];
        float y = 0;
        #pragma unroll
        for (int j = 0; j < _N_; j++)
        {
            float& s = state[j];
            s = s * w[j] + k[j] * vv + sa * b[j];
            y += s * r[j];
        }
        _y[t] = F(y);

        if ((_t+1) % _CHUNK_LEN_ == 0)
        {
            const int a = _t / _CHUNK_LEN_;
            const int c = _T_ / _CHUNK_LEN_;
            const int p = e*C*_N_*c + h*_N_*_N_*c + a*_N_*_N_ + i;
            #pragma unroll
            for (int j = 0; j < _N_; j++)
            {
                _sss[p + j*_N_] = float(state[j]);
            }
        }
    }
}

template <typename F>
__global__ void kernel_backward_zzz(const int B, const int T, const int C, const int H,
    const F *__restrict__ const _r, const F *__restrict__ const _w, const F *__restrict__ const _k, const F *__restrict__ const _a, const F *__restrict__ const _b, const F *__restrict__ const _gy,
    float *__restrict__ const _zzz)
{
    const int e = blockIdx.x / H;
    const int h = blockIdx.x % H;
    const int i = threadIdx.x;

    __shared__ float r[_N_], w[_N_], a[_N_], b[_N_];

    const int T_1 = e*T*C + (T-1)*C + h*_N_;
    float z[_N_];
    const float gy = _gy[T_1 + i];
    __syncthreads();
    r[i] = float(_r[T_1+i]);
    __syncthreads();
    #pragma unroll
    for (int j = 0; j < _N_; j++) 
    {
        z[j] = gy * r[j];
    }

    for (int _t = T-2; _t > _CHUNK_LEN_-1; _t--)
    {
        const int t = e*T*C + h*_N_ + _t * C + i;
        const float gy = _gy[t];
        __syncthreads();
        r[i] = float(_r[t]);
        w[i] = __expf(-__expf(float(_w[t+C])));
        a[i] = float(_a[t+C]);
        b[i] = float(_b[t+C]);
        __syncthreads();

        float zz = 0;
        #pragma unroll
        for (int j = 0; j < _N_; j++)
        {
            zz += b[j] * z[j];
        }
        #pragma unroll
        for (int j = 0; j < _N_; j++)
        {
            z[j] = z[j] * w[j] + gy * r[j] + a[j] * zz;
            // printf("t %d i %d j %d z %f\n", _t, i, j, z[j]);
        }
        if (_t % _CHUNK_LEN_ == 0)
        {
            const int a = _t / _CHUNK_LEN_ - 1;
            const int c = _T_ / _CHUNK_LEN_ - 1;
            const int p = e*C*_N_*c + h*_N_*_N_*c + a*_N_*_N_ + i;
            #pragma unroll
            for (int j = 0; j < _N_; j++)
            {
                _zzz[p + j*_N_] = float(z[j]);
            }
        }
    }
}

template <typename F>
__global__ void kernel_backward_rwkv(const int B, const int T, const int C, const int H,
    const F *__restrict__ const _r, const F *__restrict__ const _w, const F *__restrict__ const _k, const F *__restrict__ const _v, const F *__restrict__ const _a, const F *__restrict__ const _b, const float *__restrict__ const _saa, const float *__restrict__ const _sss, const float *__restrict__ const _zzz,
    const F *__restrict__ const _gy, F *__restrict__ const _gr, F *__restrict__ const _gw, F *__restrict__ const _gk, F *__restrict__ const _gv, F *__restrict__ const _ga, F *__restrict__ const _gb)
{
    const int e = blockIdx.x / H;
    const int h = blockIdx.x % H;
    const int chunk = threadIdx.x;
    const int n_chunk = _T_ / _CHUNK_LEN_;
    
    float zzz[_N_*_N_] = {0}, sss[_N_*_N_] = {999}, saa[_N_] = {999};
    float r[_N_] = {999}, w[_N_] = {0}, w1[_N_] = {999}, winv[_N_] = {999}, ww[_N_] = {999};
    float k[_N_] = {0}, v[_N_] = {999}, a[_N_] = {0}, a1[_N_] = {999}, b[_N_] = {0}, b1[_N_] = {999}, gy[_N_] = {999};

    if (chunk != n_chunk - 1)
    {
        const int p = e*T*C + (chunk+1)*_CHUNK_LEN_*C + h*_N_;
        for (int i = 0; i < _N_; i++) 
        {
            k[i] = float(_k[p+i]);
            a[i] = float(_a[p+i]);
            b[i] = float(_b[p+i]);
            w[i] = __expf(-__expf(float(_w[p+i])));
            const int p = e*C*_N_*(n_chunk-1) + h*_N_*_N_*(n_chunk-1) + chunk*_N_*_N_ + i*_N_;
            #pragma unroll
            for (int j = 0; j < _N_; j++)
            {
                zzz[i*_N_+j] = float(_zzz[p+j]);
            }
        }
    }
    for (int i = 0; i < _N_; i++)
    {
        const int p = e*C*_N_*n_chunk + h*_N_*_N_*n_chunk + chunk*_N_*_N_ + i*_N_;
        #pragma unroll
        for (int j = 0; j < _N_; j++)
        {
            sss[i*_N_+j] = float(_sss[p+j]);
        }
    }

    for (int _t = _CHUNK_LEN_-1; _t > -1; _t--)
    {
        const int t = chunk * _CHUNK_LEN_ + _t;
        const int b_t_h = e*T*C + t*C + h*_N_;
        #pragma unroll
        for (int n = 0; n < _N_; n++)
        {
            w1[n] = w[n];
            a1[n] = a[n];
            b1[n] = b[n];

            r[n] = float(_r[b_t_h+n]);
            k[n] = float(_k[b_t_h+n]);
            v[n] = float(_v[b_t_h+n]);
            a[n] = float(_a[b_t_h+n]);
            b[n] = float(_b[b_t_h+n]);
            gy[n] = float(_gy[b_t_h+n]);
            saa[n] = float(_saa[b_t_h+n]);

            ww[n] = -__expf(float(_w[b_t_h+n]));
            w[n] = __expf(ww[n]);
            ww[n] = ww[n] * w[n];
            winv[n] = 1.0 / w[n];
        }

        for (int j = 0; j < _N_; j++)
        {
            float zz = 0;
            #pragma unroll
            for (int i = 0; i < _N_; i++)
            {
                zz += b1[i] * zzz[i*_N_+j];
            }
            const float gyj = gy[j];
            #pragma unroll
            for (int i = 0; i < _N_; i++)
            {
                zzz[i*_N_+j] = zzz[i*_N_+j] * w1[i] + gyj * r[i] + a1[i] * zz;
                // printf("t %d i %d j %d z %f\n",t,i,j,zzz[i*_N_+j]);
                // printf("t %d i %d j %d s %f\n",t,i,j,sss[i*_N_+j]);
            }
        }

        for (int i = 0; i < _N_; i++)
        {
            float gr = 0;
            #pragma unroll
            for (int j = 0; j < _N_; j++)
            {
                gr += gy[j] * sss[i*_N_+j];
            }
            _gr[b_t_h+i] = F(gr);
        }

        for (int i = 0; i < _N_; i++)
        {
            const float ki = k[i];
            const float bi = b[i];
            const float wi = winv[i];
            #pragma unroll
            for (int j = 0; j < _N_; j++)
            {
                sss[i*_N_+j] = (sss[i*_N_+j] - ki * v[j] - bi * saa[j]) * wi;
            }
        }

        float gv[_N_] = {0}; float as[_N_] = {0}; float bz[_N_] = {0};
        for (int i = 0; i < _N_; i++)
        {
            const float ki = k[i];
            const float ai = a[i];
            const float bi = b[i];
            float gw = 0;
            float gk = 0;
            #pragma unroll
            for (int j = 0; j < _N_; j++)
            {
                const float sij = sss[i*_N_+j];
                const float zij = zzz[i*_N_+j];
                gv[j] += ki * zij;
                as[j] += ai * sij;
                bz[j] += bi * zij;
                gw += sij * zij;
                gk += v[j] * zij;
            }
            _gw[b_t_h+i] = F(gw * ww[i]);
            _gk[b_t_h+i] = F(gk);
        }
        for (int i = 0; i < _N_; i++)
        {
            _gv[b_t_h+i] = F(gv[i]);
            float ga = 0;
            float gb = 0;
            #pragma unroll
            for (int j = 0; j < _N_; j++)
            {
                ga += bz[j] * sss[i*_N_+j];
                gb += as[j] * zzz[i*_N_+j];
            }
            _ga[b_t_h+i] = F(ga);
            _gb[b_t_h+i] = F(gb);
        }
    }
}

void cuda_forward(int B, int T, int C, int H, bf16 *r, bf16* w, bf16 *k, bf16 *v, bf16 *a, bf16 *b, bf16 *y, float *saa, float* sss)
{
    assert(H*_N_ == C);
    kernel_forward<<<dim3(B * H), dim3(_N_)>>>(B, T, C, H, r, w, k, v, a, b, y, saa, sss);
}

void cuda_backward(int B, int T, int C, int H, bf16 *r, bf16* w, bf16 *k, bf16 *v, bf16 *a, bf16 *b, float *saa, float* sss, float* zzz, bf16 *gy, bf16 *gr, bf16 *gw, bf16 *gk, bf16 *gv, bf16 *ga, bf16 *gb)
{
    assert(H*_N_ == C);
    assert(T%_CHUNK_LEN_ == 0);

    kernel_backward_zzz<<<dim3(B * H), dim3(_N_)>>>(B, T, C, H, r, w, k, a, b, gy, zzz);
    kernel_backward_rwkv<<<dim3(B * H), dim3(T/_CHUNK_LEN_)>>>(B, T, C, H, r, w, k, v, a, b, saa, sss, zzz, gy, gr, gw, gk, gv, ga, gb);
}

DTYPE = torch.bfloat16
    XTYPE = torch.float
    T = sequence_length
    CHUNK_LEN = 16

    load(name="wkv7g", sources=["rwkv_cuda/wkv7g_op.cpp", f"rwkv_cuda/wkv7g_v1.cu"], is_python_module=False,
                        verbose=True, extra_cuda_cflags=["-res-usage", "--use_fast_math", "-O3", "-Xptxas -O3", "--extra-device-vectorization", f"-D_N_={HEAD_SIZE}", f"-D_T_={T}", f"-D_CHUNK_LEN_={CHUNK_LEN}"])
    class WKV_7g(torch.autograd.Function):
        @staticmethod
        def forward(ctx, r, w, k, v, a, b):
            with torch.no_grad():
                B, T, C = r.size()
                H = C // HEAD_SIZE
                N = HEAD_SIZE
                A = T // CHUNK_LEN
                assert HEAD_SIZE == C // H
                assert T % CHUNK_LEN == 0
                assert r.dtype == DTYPE
                assert w.dtype == DTYPE
                assert k.dtype == DTYPE
                assert v.dtype == DTYPE
                assert a.dtype == DTYPE
                assert b.dtype == DTYPE
                assert r.is_contiguous()
                assert w.is_contiguous()
                assert k.is_contiguous()
                assert v.is_contiguous()
                assert a.is_contiguous()
                assert b.is_contiguous()
                ctx.B = B
                ctx.T = T
                ctx.C = C
                ctx.H = H
                y = torch.empty((B, T, C), device=k.device, dtype=DTYPE, memory_format=torch.contiguous_format)#.uniform_(-100, 100)
                saa = torch.empty((B, T, H, N), device=k.device, dtype=torch.float, memory_format=torch.contiguous_format)#.uniform_(-100, 100)
                sss = torch.empty((B, H, A, N, N), device=k.device, dtype=torch.float, memory_format=torch.contiguous_format)#.uniform_(-100, 100)
                torch.ops.wkv7g.forward(B, T, C, H, r, w, k, v, a, b, y, saa, sss)
                ctx.save_for_backward(r, w, k, v, a, b, saa, sss)
                return y
        @staticmethod
        def backward(ctx, gy):
            with torch.no_grad():
                N = HEAD_SIZE
                B = ctx.B
                T = ctx.T
                C = ctx.C
                H = ctx.H
                A = T // CHUNK_LEN
                assert gy.dtype == DTYPE
                assert gy.is_contiguous()
                r, w, k, v, a, b, saa, sss = ctx.saved_tensors
                gr = torch.empty((B, T, C), device=gy.device, requires_grad=False, dtype=DTYPE, memory_format=torch.contiguous_format)#.zero_()#.uniform_(-100, 100)
                gw = torch.empty((B, T, C), device=gy.device, requires_grad=False, dtype=DTYPE, memory_format=torch.contiguous_format)#.zero_()#.uniform_(-100, 100)
                gk = torch.empty((B, T, C), device=gy.device, requires_grad=False, dtype=DTYPE, memory_format=torch.contiguous_format)#.zero_()#.uniform_(-100, 100)
                gv = torch.empty((B, T, C), device=gy.device, requires_grad=False, dtype=DTYPE, memory_format=torch.contiguous_format)#.zero_()#.uniform_(-100, 100)
                ga = torch.empty((B, T, C), device=gy.device, requires_grad=False, dtype=DTYPE, memory_format=torch.contiguous_format)#.uniform_(-100, 100)
                gb = torch.empty((B, T, C), device=gy.device, requires_grad=False, dtype=DTYPE, memory_format=torch.contiguous_format)#.uniform_(-100, 100)
                zzz = torch.empty((B, H, A-1, N, N), device=gy.device, dtype=XTYPE, memory_format=torch.contiguous_format)#.uniform_(-100, 100)
                torch.ops.wkv7g.backward(B, T, C, H, r, w, k, v, a, b, saa, sss, zzz, gy, gr, gw, gk, gv, ga, gb)
                del saa
                del sss
                del zzz
                return (gr, gw, gk, gv, ga, gb)
    @torch.compiler.disable
    def RUN_CUDA_RWKV7g(r, w, k, v, a, b):
        return WKV_7g.apply(r, w, k, v, a, b)
class RWKV7(nn.Module):
    def __init__(self, args, layer_id):
        super().__init__()
        self.args = args
        self.layer_id = layer_id
        self.n_embd = args.n_embd
        args.dim_att = args.n_embd

        self.head_size = HEAD_SIZE
        self.n_head = args.dim_att // self.head_size
        assert args.dim_att % self.n_head == 0

        with torch.no_grad():
            ratio_0_to_1 = layer_id / (args.n_layer - 1)  # 0 to 1
            ratio_1_to_almost0 = 1.0 - (layer_id / args.n_layer)  # 1 to ~0
            ddd = torch.ones(1, 1, args.n_embd)
            for i in range(args.n_embd):
                ddd[0, 0, i] = i / args.n_embd

            # initialization comes from fitting my RWKV-6 7B runs
            # merging r&g w&a to save params
            self.time_maa_x = nn.Parameter(1.0 - torch.pow(ddd, 0.6 * ratio_1_to_almost0 ** 0.9))
            self.time_maa_rg = nn.Parameter(1.0 - torch.pow(ddd, 0.2 * ratio_1_to_almost0))
            self.time_maa_wa = nn.Parameter(1.0 - torch.pow(ddd, 0.9 * ratio_1_to_almost0))
            self.time_maa_k = nn.Parameter(1.0 - (torch.pow(ddd, 0.9 * ratio_1_to_almost0) + 0.4 * ratio_0_to_1))
            self.time_maa_v = nn.Parameter(1.0 - (torch.pow(ddd, 0.4 * ratio_1_to_almost0) + 0.6 * ratio_0_to_1))

            decay_speed = torch.ones(args.dim_att)
            for n in range(args.dim_att):
                decay_speed[n] = -7 + 5 * (n / (args.dim_att - 1)) ** (0.85 + 1.0 * ratio_0_to_1 ** 0.5)
            self.time_decay = nn.Parameter(decay_speed.reshape(1,1,args.dim_att) + 0.5) # !!! 0.5 comes from F.softplus !!!

            self.time_faaaa = nn.Parameter(torch.zeros(1,1,self.n_head,self.head_size))
            self.time_aaaaa = nn.Parameter(torch.zeros(1,1,args.dim_att))

            def ortho_init(x, scale):
                with torch.no_grad():
                    shape = x.shape
                    if len(shape) == 2:
                        gain = math.sqrt(shape[0] / shape[1]) if shape[0] > shape[1] else 1
                        nn.init.orthogonal_(x, gain=gain * scale)
                    elif len(shape) == 3:
                        gain = math.sqrt(shape[1] / shape[2]) if shape[1] > shape[2] else 1
                        for i in range(shape[0]):
                            nn.init.orthogonal_(x[i], gain=gain * scale)
                    else:
                        assert False
                    return x

            D_MIX_LORA = 32
            self.time_maa_w1 = nn.Parameter(torch.zeros(args.n_embd, D_MIX_LORA*4))
            self.time_maa_w2 = nn.Parameter(ortho_init(torch.zeros(4, D_MIX_LORA, args.n_embd), 0.1))

            D_DECAY_LORA = 64
            self.time_decay_w1 = nn.Parameter(torch.zeros(args.n_embd, D_DECAY_LORA))
            self.time_decay_w2 = nn.Parameter(ortho_init(torch.zeros(D_DECAY_LORA, args.dim_att), 0.1))

            D_AAA_LORA = 16
            self.time_aaa_w1 = nn.Parameter(torch.zeros(args.n_embd, D_AAA_LORA))
            self.time_aaa_w2 = nn.Parameter(ortho_init(torch.zeros(D_AAA_LORA, args.dim_att), 0.1))

            D_KKK_LORA = 16
            self.time_kkk_w1 = nn.Parameter(torch.zeros(args.n_embd, D_KKK_LORA))
            self.time_kkk_w2 = nn.Parameter(ortho_init(torch.zeros(D_KKK_LORA, args.dim_att), 0.1))

            D_GATE_LORA = 128
            self.gate_w1 = nn.Parameter(ortho_init(torch.zeros(args.n_embd, D_GATE_LORA), 0.1))
            self.gate_w2 = nn.Parameter(ortho_init(torch.zeros(D_GATE_LORA, args.dim_att), 0.1))

            D_MA_LORA = 16
            self.ma_w1 = nn.Parameter(torch.zeros(args.n_embd, D_MA_LORA))
            self.ma_w2 = nn.Parameter(ortho_init(torch.zeros(D_MA_LORA, args.dim_att), 0.1))
            self.time_misc_a = nn.Parameter(torch.zeros(1,1,args.n_embd))
            D_MK_LORA = 16
            self.mk_w1 = nn.Parameter(torch.zeros(args.n_embd, D_MK_LORA))
            self.mk_w2 = nn.Parameter(ortho_init(torch.zeros(D_MK_LORA, args.dim_att), 0.1))
            self.time_misc_k = nn.Parameter(torch.zeros(1,1,args.n_embd))

            self.time_shift = nn.ZeroPad2d((0, 0, 1, -1))
            self.receptance = nn.Linear(args.n_embd, args.dim_att, bias=False)
            self.key = nn.Linear(args.n_embd, args.dim_att, bias=False)
            self.value = nn.Linear(args.n_embd, args.dim_att, bias=False)
            self.output = nn.Linear(args.dim_att, args.n_embd, bias=False)
            self.ln_x = nn.GroupNorm(self.n_head, args.dim_att, eps=64e-5)

            self.receptance.weight.data.uniform_(-0.5/(self.n_embd**0.5), 0.5/(self.n_embd**0.5))
            self.key.weight.data.uniform_(-0.05/(self.n_embd**0.5), 0.05/(self.n_embd**0.5))
            self.value.weight.data.uniform_(-0.5/(self.n_embd**0.5), 0.5/(self.n_embd**0.5))
            self.output.weight.data.zero_()

    def forward(self, x):
        B, T, C = x.size()
        H = self.n_head
        xx = self.time_shift(x) - x

        xxx = x + xx * self.time_maa_x
        xxx = torch.tanh(xxx @ self.time_maa_w1).view(B*T, 4, -1).transpose(0, 1)
        xxx = torch.bmm(xxx, self.time_maa_w2).view(4, B, T, -1)
        mrg, mwa, mk, mv = xxx.unbind(dim=0)

        xrg = x + xx * (self.time_maa_rg + mrg)
        xwa = x + xx * (self.time_maa_wa + mwa)
        xk = x + xx * (self.time_maa_k + mk)
        xv = x + xx * (self.time_maa_v + mv)

        r = self.receptance(xrg)
        w = -F.softplus(-(self.time_decay + torch.tanh(xwa @ self.time_decay_w1) @ self.time_decay_w2)) - 0.5
        k = self.key(xk)
        v = self.value(xv)
        g = torch.tanh(xrg @ self.gate_w1) @ self.gate_w2

        kk = k + torch.tanh(xk @ self.time_kkk_w1) @ self.time_kkk_w2
        kk = F.normalize(kk.view(B,T,H,-1), dim=-1, p=2.0).view(B,T,C)
        a = torch.sigmoid( self.time_aaaaa + (xwa @ self.time_aaa_w1) @ self.time_aaa_w2 )

        ma = torch.sigmoid(self.time_misc_a + (xwa @ self.ma_w1) @ self.ma_w2)
        k = k * ma + k*a * (1 - ma)
        mk = torch.sigmoid(self.time_misc_k + (xk @ self.mk_w1) @ self.mk_w2)
        k = k * torch.clamp(w*mk, max=0).exp()

        x = RUN_CUDA_RWKV7g(r.bfloat16(), w.bfloat16(), k.bfloat16(), v.bfloat16(), -kk.bfloat16(), (kk*a).bfloat16())
        x = self.ln_x(x.view(B * T, C)).view(B, T, C)

        x = x + ((r.view(B,T,H,-1)*k.view(B,T,H,-1)*self.time_faaaa).sum(dim=-1, keepdim=True) * v.view(B,T,H,-1)).view(B,T,C)
        x = self.output(x * g)
        return x