Kernels:

phanerozoic
/

spec-decode-ops

Kernel card Files Files and versions

xet

Community

CharlesCNorton commited on Jul 12

Commit

335d7ff

1 Parent(s): 1e410e6

spec-decode-ops: grid-stride cast_kernel rows so M may exceed the 65535 gridDim.y cap

Browse files

Files changed (1) hide show

spec_decode_cuda/spec_decode_ops.cu +41 -8

spec_decode_cuda/spec_decode_ops.cu CHANGED Viewed

@@ -94,10 +94,19 @@ __device__ void block_argmax(float v, int64_t idx, float* out_v, int64_t* out_i)
 template <typename T>
 __global__ void cast_kernel(const T* __restrict__ logits, int64_t M, int64_t V,
                             float* __restrict__ out) {
-  const int64_t row = blockIdx.y;
-  for (int64_t v = blockIdx.x * (int64_t)blockDim.x + threadIdx.x; v < V;
-       v += (int64_t)gridDim.x * blockDim.x) {
-    out[row * V + v] = ldf(logits, row * V + v);
   }
 }
@@ -133,7 +142,27 @@ __global__ void sorted_sample_kernel(const float* __restrict__ sorted_logits,
   const float* sl = sorted_logits + row * V;
   const int* si = sorted_idx + row * V;
   const float lmax = sl[0] * inv_temp;
-  const int64_t kk = (top_k > 0 && top_k < V) ? top_k : V;
   // Z over the top_k prefix (post-temperature).
   float part = 0.f;
@@ -197,9 +226,12 @@ __global__ void sorted_sample_kernel(const float* __restrict__ sorted_logits,
   __syncthreads();
   if (write_mask) {
     float* mo = masked_out + row * V;
     for (int64_t i = threadIdx.x; i < V; i += blockDim.x) {
-      mo[si[i]] = (i < L) ? sl[i] * inv_temp : kNegInf;
     }
     return;
   }
@@ -498,7 +530,7 @@ void run_sample(const torch::Tensor& logits, const c10::optional<torch::Tensor>&
   TORCH_CHECK(M * V < INT32_MAX, "M * V must fit in int32 for the segmented sort");
   auto work = torch::empty({M, V}, opts_f);
   {
-    dim3 grid((unsigned)grid_x_for(V), (unsigned)M);
     cast_kernel<T><<<grid, kThreads, 0, stream>>>(tptr<T>(logits), M, V, work.data_ptr<float>());
     if (prev.has_value() && rep != 1.0) {
       penalty_kernel<<<M, kThreads, 0, stream>>>(work.data_ptr<float>(), V,
@@ -509,7 +541,8 @@ void run_sample(const torch::Tensor& logits, const c10::optional<torch::Tensor>&
   // segmented descending sort with indices
   auto sorted = torch::empty({M, V}, opts_f);
-  auto idx_in = torch::arange(V, logits.options().dtype(torch::kInt32)).repeat({M, 1}).contiguous();
   auto idx_out = torch::empty({M, V}, logits.options().dtype(torch::kInt32));
   auto offs = torch::arange(0, (M + 1) * V, V, logits.options().dtype(torch::kInt32)).contiguous();
   size_t temp_bytes = 0;

 template <typename T>
 __global__ void cast_kernel(const T* __restrict__ logits, int64_t M, int64_t V,
                             float* __restrict__ out) {
+  // grid-stride over rows so M may exceed the 65535 gridDim.y hardware cap.
+  for (int64_t row = blockIdx.y; row < M; row += gridDim.y) {
+    for (int64_t v = blockIdx.x * (int64_t)blockDim.x + threadIdx.x; v < V;
+         v += (int64_t)gridDim.x * blockDim.x) {
+      out[row * V + v] = ldf(logits, row * V + v);
+    }
+  }
+}
+__global__ void iota_kernel(int* __restrict__ out, int64_t M, int64_t V) {
+  for (int64_t i = blockIdx.x * (int64_t)blockDim.x + threadIdx.x; i < M * V;
+       i += (int64_t)gridDim.x * blockDim.x) {
+    out[i] = (int)(i % V);
   }
 }
   const float* sl = sorted_logits + row * V;
   const int* si = sorted_idx + row * V;
   const float lmax = sl[0] * inv_temp;
+  // top-k with tie inclusion: transformers removes only values strictly
+  // below the k-th, so equal-valued tokens past position k stay kept.
+  int64_t kk = V;
+  if (top_k > 0 && top_k < V) {
+    const float kth = sl[top_k - 1];
+    __shared__ int64_t s_kk;
+    if (threadIdx.x == 0) s_kk = V;
+    __syncthreads();
+    const int64_t start = (top_k / kThreads) * kThreads;
+    for (int64_t base = start; base < V; base += blockDim.x) {
+      const int64_t i = base + threadIdx.x;
+      if (i >= top_k && i < V && sl[i] < kth) {
+        atomicMin((unsigned long long*)&s_kk, (unsigned long long)i);
+      }
+      __syncthreads();
+      if (s_kk < base + (int64_t)kThreads) break;
+    }
+    kk = s_kk;
+    __syncthreads();
+  }
   // Z over the top_k prefix (post-temperature).
   float part = 0.f;
   __syncthreads();
   if (write_mask) {
+    // masked entries carry true -inf, matching the transformers chain's
+    // filter_value; -FLT_MAX would evaluate identically under softmax but
+    // fails isinf-based consumers.
     float* mo = masked_out + row * V;
     for (int64_t i = threadIdx.x; i < V; i += blockDim.x) {
+      mo[si[i]] = (i < L) ? sl[i] * inv_temp : -INFINITY;
     }
     return;
   }
   TORCH_CHECK(M * V < INT32_MAX, "M * V must fit in int32 for the segmented sort");
   auto work = torch::empty({M, V}, opts_f);
   {
+    dim3 grid((unsigned)grid_x_for(V), (unsigned)std::min<int64_t>(M, 65535));
     cast_kernel<T><<<grid, kThreads, 0, stream>>>(tptr<T>(logits), M, V, work.data_ptr<float>());
     if (prev.has_value() && rep != 1.0) {
       penalty_kernel<<<M, kThreads, 0, stream>>>(work.data_ptr<float>(), V,
   // segmented descending sort with indices
   auto sorted = torch::empty({M, V}, opts_f);
+  auto idx_in = torch::empty({M, V}, logits.options().dtype(torch::kInt32));
+  iota_kernel<<<grid_x_for(M * V), kThreads, 0, stream>>>(idx_in.data_ptr<int>(), M, V);
   auto idx_out = torch::empty({M, V}, logits.options().dtype(torch::kInt32));
   auto offs = torch::arange(0, (M + 1) * V, V, logits.options().dtype(torch::kInt32)).contiguous();
   size_t temp_bytes = 0;