.

app.py

@@ -21,8 +21,11 @@ RELEASE_PREFILL_CACHE = True # saves VRAM before decode; costs one sync/cache fl
 
 ########################## text rwkv ################################################################
 
-title = "rwkv7-g1h-7.2b-20260710-ctx10240"
-model_path = hf_hub_download(repo_id="BlinkDL/rwkv7-g1", filename=f"{title}.pth")
+# title = "rwkv7-g1h-7.2b-20260710-ctx10240"
+# model_path = hf_hub_download(repo_id="BlinkDL/rwkv7-g1", filename=f"{title}.pth")
+
+title = "rwkv7-g1i_preview3260-7.2b-20260716-ctx12288"
+model_path = hf_hub_download(repo_id="BlinkDL/temp-latest-training-models", filename=f"{title}.pth")
 
 v3a.MODEL_PATH = model_path
 v3a.WKV_MODE = "fp32io16" # use "fp16" to save WKV state VRAM, with lower precision

cuda/rwkv7_fast_ops_fp16.cu

@@ -49,6 +49,7 @@ __device__ inline float sigmoid_fast(float x) {
   return 1.0f / (1.0f + __expf(-x));
 }
 
+template <bool UpdateShift>
 __global__ void tmix_mix6_kernel(
     int T,
     int C,
@@ -106,11 +107,31 @@ __global__ void tmix_mix6_kernel(
   store_h2(out_a + idx, cur.x + dx0 * xa.x, cur.y + dx1 * xa.y);
   store_h2(out_g + idx, cur.x + dx0 * xg.x, cur.y + dx1 * xg.y);
 
-  if (t == T - 1) {
-    *reinterpret_cast<__half2*>(shift_state + static_cast<int64_t>(b) * C + c) = cur2;
+  if constexpr (UpdateShift) {
+    if (t == T - 1) {
+      *reinterpret_cast<__half2*>(shift_state + static_cast<int64_t>(b) * C + c) = cur2;
+    }
   }
 }
 
+__global__ void update_shift_state_last_kernel(
+    int T,
+    int C,
+    const dtype* __restrict__ x,
+    dtype* __restrict__ shift_state,
+    int64_t total_pairs) {
+  const int64_t pair_idx = static_cast<int64_t>(blockIdx.x) * blockDim.x + threadIdx.x;
+  if (pair_idx >= total_pairs) {
+    return;
+  }
+
+  const int c_pairs = C >> 1;
+  const int b = static_cast<int>(pair_idx / c_pairs);
+  const int c = static_cast<int>(pair_idx - static_cast<int64_t>(b) * c_pairs) << 1;
+  const int64_t src_idx = (static_cast<int64_t>(b) * T + (T - 1)) * C + c;
+  *reinterpret_cast<__half2*>(shift_state + static_cast<int64_t>(b) * C + c) = load_h2(x + src_idx);
+}
+
 template <bool HalfMath, int Vec>
 __global__ void tmix_mix6_t1_c4096_kernel(
     const dtype* __restrict__ x,
@@ -439,6 +460,7 @@ __global__ void zero_vec4_kernel(dtype* __restrict__ out, int64_t n_vec4) {
   }
 }
 
+template <bool UpdateShift>
 __global__ void cmix_mix_kernel(
     int T,
     int C,
@@ -466,8 +488,10 @@ __global__ void cmix_mix_kernel(
   const float2 mix = __half22float2(load_h2(x_k + c));
   store_h2(out + idx, cur.x + (prev.x - cur.x) * mix.x, cur.y + (prev.y - cur.y) * mix.y);
 
-  if (t == T - 1) {
-    *reinterpret_cast<__half2*>(shift_state + static_cast<int64_t>(b) * C + c) = cur2;
+  if constexpr (UpdateShift) {
+    if (t == T - 1) {
+      *reinterpret_cast<__half2*>(shift_state + static_cast<int64_t>(b) * C + c) = cur2;
+    }
   }
 }
 
@@ -895,23 +919,46 @@ std::vector<at::Tensor> tmix_mix6_cuda(
   constexpr int threads = 256;
   const int64_t total_pairs = static_cast<int64_t>(B) * T * (C / 2);
   auto stream = at::cuda::getCurrentCUDAStream();
-  tmix_mix6_kernel<<<static_cast<int>(ceil_div(total_pairs, threads)), threads, 0, stream>>>(
-      T, C,
-      x.data_ptr<dtype>(),
-      shift_state.data_ptr<dtype>(),
-      x_r.data_ptr<dtype>(),
-      x_w.data_ptr<dtype>(),
-      x_k.data_ptr<dtype>(),
-      x_v.data_ptr<dtype>(),
-      x_a.data_ptr<dtype>(),
-      x_g.data_ptr<dtype>(),
-      out_r.data_ptr<dtype>(),
-      out_w.data_ptr<dtype>(),
-      out_k.data_ptr<dtype>(),
-      out_v.data_ptr<dtype>(),
-      out_a.data_ptr<dtype>(),
-      out_g.data_ptr<dtype>(),
-      total_pairs);
+  if (T == 1) {
+    tmix_mix6_kernel<true><<<static_cast<int>(ceil_div(total_pairs, threads)), threads, 0, stream>>>(
+        T, C,
+        x.data_ptr<dtype>(),
+        shift_state.data_ptr<dtype>(),
+        x_r.data_ptr<dtype>(),
+        x_w.data_ptr<dtype>(),
+        x_k.data_ptr<dtype>(),
+        x_v.data_ptr<dtype>(),
+        x_a.data_ptr<dtype>(),
+        x_g.data_ptr<dtype>(),
+        out_r.data_ptr<dtype>(),
+        out_w.data_ptr<dtype>(),
+        out_k.data_ptr<dtype>(),
+        out_v.data_ptr<dtype>(),
+        out_a.data_ptr<dtype>(),
+        out_g.data_ptr<dtype>(),
+        total_pairs);
+  } else {
+    tmix_mix6_kernel<false><<<static_cast<int>(ceil_div(total_pairs, threads)), threads, 0, stream>>>(
+        T, C,
+        x.data_ptr<dtype>(),
+        shift_state.data_ptr<dtype>(),
+        x_r.data_ptr<dtype>(),
+        x_w.data_ptr<dtype>(),
+        x_k.data_ptr<dtype>(),
+        x_v.data_ptr<dtype>(),
+        x_a.data_ptr<dtype>(),
+        x_g.data_ptr<dtype>(),
+        out_r.data_ptr<dtype>(),
+        out_w.data_ptr<dtype>(),
+        out_k.data_ptr<dtype>(),
+        out_v.data_ptr<dtype>(),
+        out_a.data_ptr<dtype>(),
+        out_g.data_ptr<dtype>(),
+        total_pairs);
+    const int64_t state_pairs = static_cast<int64_t>(B) * (C / 2);
+    update_shift_state_last_kernel<<<static_cast<int>(ceil_div(state_pairs, threads)), threads, 0, stream>>>(
+        T, C, x.data_ptr<dtype>(), shift_state.data_ptr<dtype>(), state_pairs);
+  }
   C10_CUDA_KERNEL_LAUNCH_CHECK();
   return {out_r, out_w, out_k, out_v, out_a, out_g};
 }
@@ -937,23 +984,46 @@ std::vector<at::Tensor> tmix_mix6_cfg_cuda(
   auto out_g = at::empty_like(x);
   const int64_t total_pairs = static_cast<int64_t>(B) * T * (C / 2);
   auto stream = at::cuda::getCurrentCUDAStream();
-  tmix_mix6_kernel<<<static_cast<int>(ceil_div(total_pairs, threads)), threads, 0, stream>>>(
-      T, C,
-      x.data_ptr<dtype>(),
-      shift_state.data_ptr<dtype>(),
-      x_r.data_ptr<dtype>(),
-      x_w.data_ptr<dtype>(),
-      x_k.data_ptr<dtype>(),
-      x_v.data_ptr<dtype>(),
-      x_a.data_ptr<dtype>(),
-      x_g.data_ptr<dtype>(),
-      out_r.data_ptr<dtype>(),
-      out_w.data_ptr<dtype>(),
-      out_k.data_ptr<dtype>(),
-      out_v.data_ptr<dtype>(),
-      out_a.data_ptr<dtype>(),
-      out_g.data_ptr<dtype>(),
-      total_pairs);
+  if (T == 1) {
+    tmix_mix6_kernel<true><<<static_cast<int>(ceil_div(total_pairs, threads)), threads, 0, stream>>>(
+        T, C,
+        x.data_ptr<dtype>(),
+        shift_state.data_ptr<dtype>(),
+        x_r.data_ptr<dtype>(),
+        x_w.data_ptr<dtype>(),
+        x_k.data_ptr<dtype>(),
+        x_v.data_ptr<dtype>(),
+        x_a.data_ptr<dtype>(),
+        x_g.data_ptr<dtype>(),
+        out_r.data_ptr<dtype>(),
+        out_w.data_ptr<dtype>(),
+        out_k.data_ptr<dtype>(),
+        out_v.data_ptr<dtype>(),
+        out_a.data_ptr<dtype>(),
+        out_g.data_ptr<dtype>(),
+        total_pairs);
+  } else {
+    tmix_mix6_kernel<false><<<static_cast<int>(ceil_div(total_pairs, threads)), threads, 0, stream>>>(
+        T, C,
+        x.data_ptr<dtype>(),
+        shift_state.data_ptr<dtype>(),
+        x_r.data_ptr<dtype>(),
+        x_w.data_ptr<dtype>(),
+        x_k.data_ptr<dtype>(),
+        x_v.data_ptr<dtype>(),
+        x_a.data_ptr<dtype>(),
+        x_g.data_ptr<dtype>(),
+        out_r.data_ptr<dtype>(),
+        out_w.data_ptr<dtype>(),
+        out_k.data_ptr<dtype>(),
+        out_v.data_ptr<dtype>(),
+        out_a.data_ptr<dtype>(),
+        out_g.data_ptr<dtype>(),
+        total_pairs);
+    const int64_t state_pairs = static_cast<int64_t>(B) * (C / 2);
+    update_shift_state_last_kernel<<<static_cast<int>(ceil_div(state_pairs, threads)), threads, 0, stream>>>(
+        T, C, x.data_ptr<dtype>(), shift_state.data_ptr<dtype>(), state_pairs);
+  }
   C10_CUDA_KERNEL_LAUNCH_CHECK();
   return {out_r, out_w, out_k, out_v, out_a, out_g};
 }
@@ -1298,14 +1368,28 @@ at::Tensor cmix_mix_cuda(
   constexpr int threads = 256;
   const int64_t total_pairs = static_cast<int64_t>(B) * T * (C / 2);
   auto stream = at::cuda::getCurrentCUDAStream();
-  cmix_mix_kernel<<<static_cast<int>(ceil_div(total_pairs, threads)), threads, 0, stream>>>(
-      T,
-      C,
-      x.data_ptr<dtype>(),
-      shift_state.data_ptr<dtype>(),
-      x_k.data_ptr<dtype>(),
-      out.data_ptr<dtype>(),
-      total_pairs);
+  if (T == 1) {
+    cmix_mix_kernel<true><<<static_cast<int>(ceil_div(total_pairs, threads)), threads, 0, stream>>>(
+        T,
+        C,
+        x.data_ptr<dtype>(),
+        shift_state.data_ptr<dtype>(),
+        x_k.data_ptr<dtype>(),
+        out.data_ptr<dtype>(),
+        total_pairs);
+  } else {
+    cmix_mix_kernel<false><<<static_cast<int>(ceil_div(total_pairs, threads)), threads, 0, stream>>>(
+        T,
+        C,
+        x.data_ptr<dtype>(),
+        shift_state.data_ptr<dtype>(),
+        x_k.data_ptr<dtype>(),
+        out.data_ptr<dtype>(),
+        total_pairs);
+    const int64_t state_pairs = static_cast<int64_t>(B) * (C / 2);
+    update_shift_state_last_kernel<<<static_cast<int>(ceil_div(state_pairs, threads)), threads, 0, stream>>>(
+        T, C, x.data_ptr<dtype>(), shift_state.data_ptr<dtype>(), state_pairs);
+  }
   C10_CUDA_KERNEL_LAUNCH_CHECK();
   return out;
 }
@@ -1321,14 +1405,28 @@ at::Tensor cmix_mix_cfg_cuda(
   auto out = at::empty_like(x);
   const int64_t total_pairs = static_cast<int64_t>(B) * T * (C / 2);
   auto stream = at::cuda::getCurrentCUDAStream();
-  cmix_mix_kernel<<<static_cast<int>(ceil_div(total_pairs, threads)), threads, 0, stream>>>(
-      T,
-      C,
-      x.data_ptr<dtype>(),
-      shift_state.data_ptr<dtype>(),
-      x_k.data_ptr<dtype>(),
-      out.data_ptr<dtype>(),
-      total_pairs);
+  if (T == 1) {
+    cmix_mix_kernel<true><<<static_cast<int>(ceil_div(total_pairs, threads)), threads, 0, stream>>>(
+        T,
+        C,
+        x.data_ptr<dtype>(),
+        shift_state.data_ptr<dtype>(),
+        x_k.data_ptr<dtype>(),
+        out.data_ptr<dtype>(),
+        total_pairs);
+  } else {
+    cmix_mix_kernel<false><<<static_cast<int>(ceil_div(total_pairs, threads)), threads, 0, stream>>>(
+        T,
+        C,
+        x.data_ptr<dtype>(),
+        shift_state.data_ptr<dtype>(),
+        x_k.data_ptr<dtype>(),
+        out.data_ptr<dtype>(),
+        total_pairs);
+    const int64_t state_pairs = static_cast<int64_t>(B) * (C / 2);
+    update_shift_state_last_kernel<<<static_cast<int>(ceil_div(state_pairs, threads)), threads, 0, stream>>>(
+        T, C, x.data_ptr<dtype>(), shift_state.data_ptr<dtype>(), state_pairs);
+  }
   C10_CUDA_KERNEL_LAUNCH_CHECK();
   return out;
 }

cuda/rwkv7_wkv_fp16_v2.cu

@@ -7,6 +7,7 @@
 #include <ATen/ATen.h>
 #include <ATen/cuda/CUDAContext.h>
 #include <cuda_fp16.h>
+#include <stdint.h>
 
 namespace {
 
@@ -16,12 +17,13 @@ constexpr int LDG_ELEMS = sizeof(int4) / sizeof(half);
 constexpr float TWO_NEG_41 = 4.547473508864641e-13f;
 constexpr float NEXP_HALF_LOG2_E = -0.8750387749145276f;
 constexpr float NLOG2_E = -1.4426950408889634f;
-constexpr int ROT1 = static_cast<int>(2654435769);
+constexpr uint32_t ROT1 = 2654435769u;
 using F = half;
 #define CLONE_N 64
 
 __device__ __forceinline__ float rotator1(int x) {
-  return TWO_NEG_41 * float(ROT1 * x);
+  const uint32_t bits = ROT1 * static_cast<uint32_t>(x);
+  return TWO_NEG_41 * static_cast<float>(static_cast<int32_t>(bits));
 }
 
 __device__ __forceinline__ half w_delta(float w, int phase) {
@@ -109,7 +111,7 @@ __global__ void __launch_bounds__(CLONE_N, 2) wkv_fp16_v1_clone_kernel(
     state[j] = state_smem[i][lane ^ j];
   }
 
-  __shared__ __align__(128) half2 r[CLONE_N / 2], k[CLONE_N / 2], w[CLONE_N / 2], a[CLONE_N / 2], bvec[CLONE_N / 2];
+  __shared__ __align__(128) half2 r[CLONE_N / 2], k[CLONE_N / 2], w[CLONE_N / 2], a[CLONE_N / 2], bvec[CLONE_N / 2], bvec_dummy[CLONE_N / 2];
 #pragma unroll
   for (int tt = 0; tt < T; tt++) {
     int t = b * T * C + h * CLONE_N + tt * C;
@@ -117,7 +119,8 @@ __global__ void __launch_bounds__(CLONE_N, 2) wkv_fp16_v1_clone_kernel(
     clone_cp_async<4>((half2*)(i < 32 ? w : a) + lane, (half2*)((i < 32 ? w_ptr : a_ptr) + t) + lane, true);
     clone_cp_commit();
     clone_cp_async<4>((half2*)(i < 32 ? r : k) + lane, (half2*)((i < 32 ? r_ptr : k_ptr) + t) + lane, true);
-    clone_cp_async<4>((half2*)bvec + lane, (half2*)(b_ptr + t) + lane, i < 32);
+    // src-size 0 zero-fills, so warp 1 must not race warp 0's real bvec copy.
+    clone_cp_async<4>((i < 32 ? bvec : bvec_dummy) + lane, (half2*)(b_ptr + t) + lane, i < 32);
     clone_cp_commit();
 
     half vv = v_ptr[t + i];
@@ -197,6 +200,7 @@ __device__ __forceinline__ void prefetch_token(
     half2* k,
     half2* a,
     half2* b,
+    half2* b_dummy,
     const half* r_ptr,
     const half* w_ptr,
     const half* k_ptr,
@@ -205,7 +209,8 @@ __device__ __forceinline__ void prefetch_token(
   cp_async<4>((tid < 32 ? w : a) + lane, (const half2*)(tid < 32 ? w_ptr + token : a_ptr + token) + lane, true);
   cp_commit();
   cp_async<4>((tid < 32 ? r : k) + lane, (const half2*)(tid < 32 ? r_ptr + token : k_ptr + token) + lane, true);
-  cp_async<4>(b + lane, (const half2*)(b_ptr + token) + lane, tid < 32);
+  // A predicated-off cp.async still zero-fills its shared destination.
+  cp_async<4>((tid < 32 ? b : b_dummy) + lane, (const half2*)(b_ptr + token) + lane, tid < 32);
   cp_commit();
 }
 
@@ -254,7 +259,7 @@ __global__ __launch_bounds__(N, 2) void wkv_fp16_v1_exact_kernel(
     state[j] = state_smem[i][lane ^ j];
   }
 
-  __shared__ __align__(128) half2 r[HALF2_N], k[HALF2_N], w[HALF2_N], a[HALF2_N], bvec[HALF2_N];
+  __shared__ __align__(128) half2 r[HALF2_N], k[HALF2_N], w[HALF2_N], a[HALF2_N], bvec[HALF2_N], bvec_dummy[HALF2_N];
 #pragma unroll
   for (int tt = 0; tt < T; tt++) {
     int t = b_id * T * C + h * N + tt * C;
@@ -262,7 +267,7 @@ __global__ __launch_bounds__(N, 2) void wkv_fp16_v1_exact_kernel(
     cp_async<4>((half2*)(i < 32 ? w : a) + lane, (half2*)((i < 32 ? w_ptr : a_ptr) + t) + lane, true);
     cp_commit();
     cp_async<4>((half2*)(i < 32 ? r : k) + lane, (half2*)((i < 32 ? r_ptr : k_ptr) + t) + lane, true);
-    cp_async<4>((half2*)bvec + lane, (half2*)(b_ptr + t) + lane, i < 32);
+    cp_async<4>((i < 32 ? bvec : bvec_dummy) + lane, (half2*)(b_ptr + t) + lane, i < 32);
     cp_commit();
 
     half vv = v_ptr[t + i];
@@ -350,9 +355,9 @@ __global__ __launch_bounds__(N, 2) void wkv_fp16_seq_v2_kernel(
     state[j] = state_smem[i][lane ^ j];
   }
 
-  __shared__ __align__(128) half2 r[2][HALF2_N], w[2][HALF2_N], k[2][HALF2_N], a[2][HALF2_N], bvec[2][HALF2_N];
+  __shared__ __align__(128) half2 r[2][HALF2_N], w[2][HALF2_N], k[2][HALF2_N], a[2][HALF2_N], bvec[2][HALF2_N], bvec_dummy[HALF2_N];
   int token = (b_id * T) * C + h * N;
-  prefetch_token(i, lane, token, r[0], w[0], k[0], a[0], bvec[0], r_ptr, w_ptr, k_ptr, a_ptr, b_ptr);
+  prefetch_token(i, lane, token, r[0], w[0], k[0], a[0], bvec[0], bvec_dummy, r_ptr, w_ptr, k_ptr, a_ptr, b_ptr);
 
   for (int tt = 0; tt < T; ++tt) {
     const int cur = tt & 1;
@@ -371,7 +376,7 @@ __global__ __launch_bounds__(N, 2) void wkv_fp16_seq_v2_kernel(
 
     if (tt + 1 < T) {
       int next_token = token + C;
-      prefetch_token(i, lane, next_token, r[cur ^ 1], w[cur ^ 1], k[cur ^ 1], a[cur ^ 1], bvec[cur ^ 1], r_ptr, w_ptr, k_ptr, a_ptr, b_ptr);
+      prefetch_token(i, lane, next_token, r[cur ^ 1], w[cur ^ 1], k[cur ^ 1], a[cur ^ 1], bvec[cur ^ 1], bvec_dummy, r_ptr, w_ptr, k_ptr, a_ptr, b_ptr);
     }
 
     half vv = v_ptr[token + i];
@@ -540,12 +545,13 @@ __global__ __launch_bounds__(N, 1) void wkv_fp16_one_cp_kernel(
     state[j] = state_smem[i][lane ^ j];
   }
 
-  __shared__ __align__(128) half2 r[HALF2_N], w[HALF2_N], k[HALF2_N], a[HALF2_N], bvec[HALF2_N];
+  __shared__ __align__(128) half2 r[HALF2_N], w[HALF2_N], k[HALF2_N], a[HALF2_N], bvec[HALF2_N], bvec_dummy[HALF2_N];
   const int token = b_id * C + h * N;
   cp_async<4>((half2*)(i < 32 ? w : a) + lane, (half2*)((i < 32 ? w_ptr : a_ptr) + token) + lane, true);
   cp_commit();
   cp_async<4>((half2*)(i < 32 ? r : k) + lane, (half2*)((i < 32 ? r_ptr : k_ptr) + token) + lane, true);
-  cp_async<4>((half2*)bvec + lane, (half2*)(b_ptr + token) + lane, i < 32);
+  // bvec_dummy is correctness-critical: warp 1's src-size 0 copy writes zeros.
+  cp_async<4>((i < 32 ? bvec : bvec_dummy) + lane, (half2*)(b_ptr + token) + lane, i < 32);
   cp_commit();
 
   half vv = __ldg(v_ptr + token + i);