model-as-a-kernel v1 source

Browse files

Files changed (3) hide show

mak_cuda/mak_kernels.h +4 -1
mak_cuda/megakernel.cu +12 -3
torch-ext/torch_binding.cpp +27 -3

mak_cuda/mak_kernels.h CHANGED Viewed

@@ -11,10 +11,13 @@ int mak_grid_blocks_impl(int smem_bytes);
 // two in [1, 16], chosen from the device's shared-memory budget).
 int mak_ring_stages_impl(int stage_bytes);
 cudaError_t mak_launch_mega(const long long* prog, int n_phases, int pos0,
                             int total, int slot0, int prompt_len,
                             int chunk_m, int ring_stages, int* bar, int grid,
-                            int smem_bytes, cudaStream_t stream);
 cudaError_t mak_launch_phase(const long long* prog, int phase, int pos,
                              int step_slot, int ring_stages, int grid,

 // two in [1, 16], chosen from the device's shared-memory budget).
 int mak_ring_stages_impl(int stage_bytes);
+// ts (nullable): block 0 stamps clock64() at entry (slot 0) and after each
+// phase's barrier crossing of the first step (slot p+1); size n_phases+1.
 cudaError_t mak_launch_mega(const long long* prog, int n_phases, int pos0,
                             int total, int slot0, int prompt_len,
                             int chunk_m, int ring_stages, int* bar, int grid,
+                            int smem_bytes, cudaStream_t stream,
+                            long long* ts);
 cudaError_t mak_launch_phase(const long long* prog, int phase, int pos,
                              int step_slot, int ring_stages, int grid,

mak_cuda/megakernel.cu CHANGED Viewed

@@ -940,10 +940,15 @@ __device__ void grid_barrier(int* bar, int* local_sense) {
 __global__ void __launch_bounds__(BLOCK_THREADS, 2)
     mega_kernel(const long long* __restrict__ prog, int n_phases,
                 int pos0, int total, int slot0, int prompt_len,
-                int chunk_m, int ring_stages, int* bar) {
   __shared__ int s_sense;
   if (threadIdx.x == 0) s_sense = atomicAdd(bar + BAR_STRIPES + 1, 0);
   __syncthreads();
   int pos = pos0;
   const int pos_end = pos0 + total;
   while (pos < pos_end) {
@@ -984,6 +989,9 @@ __global__ void __launch_bounds__(BLOCK_THREADS, 2)
         }
         grid_barrier(bar, &s_sense);
       }
     }
     pos += M;
   }
@@ -1089,10 +1097,11 @@ int mak_grid_blocks_impl(int smem_bytes) {
 cudaError_t mak_launch_mega(const long long* prog, int n_phases, int pos0,
                             int total, int slot0, int prompt_len,
                             int chunk_m, int ring_stages, int* bar, int grid,
-                            int smem_bytes, cudaStream_t stream) {
   mak::mega_kernel<<<grid, BLOCK_THREADS, smem_bytes, stream>>>(
       prog, n_phases, pos0, total, slot0, prompt_len, chunk_m, ring_stages,
-      bar);
   return cudaGetLastError();
 }

 __global__ void __launch_bounds__(BLOCK_THREADS, 2)
     mega_kernel(const long long* __restrict__ prog, int n_phases,
                 int pos0, int total, int slot0, int prompt_len,
+                int chunk_m, int ring_stages, int* bar,
+                long long* __restrict__ ts) {
   __shared__ int s_sense;
   if (threadIdx.x == 0) s_sense = atomicAdd(bar + BAR_STRIPES + 1, 0);
   __syncthreads();
+  // diagnostic: block 0 stamps its SM clock at entry and after each
+  // barrier crossing of the first step (pure observation, no data effect)
+  if (ts != nullptr && blockIdx.x == 0 && threadIdx.x == 0)
+    ts[0] = (long long)clock64();
   int pos = pos0;
   const int pos_end = pos0 + total;
   while (pos < pos_end) {
         }
         grid_barrier(bar, &s_sense);
       }
+      if (ts != nullptr && pos == pos0 && blockIdx.x == 0 &&
+          threadIdx.x == 0)
+        ts[p + 1] = (long long)clock64();
     }
     pos += M;
   }
 cudaError_t mak_launch_mega(const long long* prog, int n_phases, int pos0,
                             int total, int slot0, int prompt_len,
                             int chunk_m, int ring_stages, int* bar, int grid,
+                            int smem_bytes, cudaStream_t stream,
+                            long long* ts) {
   mak::mega_kernel<<<grid, BLOCK_THREADS, smem_bytes, stream>>>(
       prog, n_phases, pos0, total, slot0, prompt_len, chunk_m, ring_stages,
+      bar, ts);
   return cudaGetLastError();
 }

torch-ext/torch_binding.cpp CHANGED Viewed

@@ -71,7 +71,7 @@ void mak_run(at::Tensor prog, at::Tensor barrier, int64_t pos,
   const cudaError_t err = mak_launch_mega(
       reinterpret_cast<const long long*>(prog.data_ptr<int64_t>()),
       (int)prog.size(0), (int)pos, 1, (int)step_slot, -1, 1, sh.ring_stages,
-      bar, mak_grid_blocks_impl(sh.smem), sh.smem, stream.stream());
   TORCH_CHECK(err == cudaSuccess, "mak_run launch failed: ",
               cudaGetErrorString(err));
 }
@@ -91,7 +91,8 @@ void mak_run_steps(at::Tensor prog, at::Tensor barrier, int64_t pos0,
   for (int64_t i = 0; i < steps; ++i) {
     const cudaError_t err =
         mak_launch_mega(pp, n, (int)(pos0 + i), 1, (int)(slot0 + i), 0, 1,
-                        sh.ring_stages, bar, grid, sh.smem, stream.stream());
     TORCH_CHECK(err == cudaSuccess, "mak_run_steps launch failed: ",
                 cudaGetErrorString(err));
   }
@@ -112,11 +113,31 @@ void mak_run_seq(at::Tensor prog, at::Tensor barrier, int64_t pos0,
       reinterpret_cast<const long long*>(prog.data_ptr<int64_t>()),
       (int)prog.size(0), (int)pos0, (int)steps, (int)slot0, (int)prompt_len,
       (int)chunk_m, sh.ring_stages, bar, mak_grid_blocks_impl(sh.smem),
-      sh.smem, stream.stream());
   TORCH_CHECK(err == cudaSuccess, "mak_run_seq launch failed: ",
               cudaGetErrorString(err));
 }
 at::Tensor mak_run_phased(at::Tensor prog, int64_t pos, int64_t step_slot,
                           bool timed, int64_t max_k) {
   check_prog(prog);
@@ -169,6 +190,8 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
           "int max_k) -> Tensor");
   ops.def("mak_num_blocks(Tensor prog, int max_k) -> int");
   ops.def("mak_bw_per_sm(Tensor probe) -> float");
 #if defined(CUDA_KERNEL) || defined(ROCM_KERNEL)
   ops.impl("mak_run", torch::kCUDA, &mak_run);
@@ -177,6 +200,7 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
   ops.impl("mak_run_phased", torch::kCUDA, &mak_run_phased);
   ops.impl("mak_num_blocks", torch::kCUDA, &mak_num_blocks);
   ops.impl("mak_bw_per_sm", torch::kCUDA, &mak_bw_per_sm);
 #endif
 }

   const cudaError_t err = mak_launch_mega(
       reinterpret_cast<const long long*>(prog.data_ptr<int64_t>()),
       (int)prog.size(0), (int)pos, 1, (int)step_slot, -1, 1, sh.ring_stages,
+      bar, mak_grid_blocks_impl(sh.smem), sh.smem, stream.stream(), nullptr);
   TORCH_CHECK(err == cudaSuccess, "mak_run launch failed: ",
               cudaGetErrorString(err));
 }
   for (int64_t i = 0; i < steps; ++i) {
     const cudaError_t err =
         mak_launch_mega(pp, n, (int)(pos0 + i), 1, (int)(slot0 + i), 0, 1,
+                        sh.ring_stages, bar, grid, sh.smem, stream.stream(),
+                        nullptr);
     TORCH_CHECK(err == cudaSuccess, "mak_run_steps launch failed: ",
                 cudaGetErrorString(err));
   }
       reinterpret_cast<const long long*>(prog.data_ptr<int64_t>()),
       (int)prog.size(0), (int)pos0, (int)steps, (int)slot0, (int)prompt_len,
       (int)chunk_m, sh.ring_stages, bar, mak_grid_blocks_impl(sh.smem),
+      sh.smem, stream.stream(), nullptr);
   TORCH_CHECK(err == cudaSuccess, "mak_run_seq launch failed: ",
               cudaGetErrorString(err));
 }
+void mak_run_ts(at::Tensor prog, at::Tensor barrier, int64_t pos,
+                int64_t step_slot, int64_t max_k, at::Tensor ts) {
+  check_prog(prog);
+  check_barrier(barrier);
+  TORCH_CHECK(ts.is_cuda() && ts.scalar_type() == at::kLong &&
+                  ts.numel() >= prog.size(0) + 1,
+              "ts must be a CUDA int64 tensor with n_phases+1 elements");
+  const c10::cuda::CUDAGuard guard(prog.device());
+  auto stream = at::cuda::getCurrentCUDAStream();
+  int* bar = barrier.data_ptr<int>();
+  const LaunchShape sh = shape_for(max_k);
+  const cudaError_t err = mak_launch_mega(
+      reinterpret_cast<const long long*>(prog.data_ptr<int64_t>()),
+      (int)prog.size(0), (int)pos, 1, (int)step_slot, -1, 1, sh.ring_stages,
+      bar, mak_grid_blocks_impl(sh.smem), sh.smem, stream.stream(),
+      reinterpret_cast<long long*>(ts.data_ptr<int64_t>()));
+  TORCH_CHECK(err == cudaSuccess, "mak_run_ts launch failed: ",
+              cudaGetErrorString(err));
+}
 at::Tensor mak_run_phased(at::Tensor prog, int64_t pos, int64_t step_slot,
                           bool timed, int64_t max_k) {
   check_prog(prog);
           "int max_k) -> Tensor");
   ops.def("mak_num_blocks(Tensor prog, int max_k) -> int");
   ops.def("mak_bw_per_sm(Tensor probe) -> float");
+  ops.def("mak_run_ts(Tensor prog, Tensor barrier, int pos, int step_slot, "
+          "int max_k, Tensor ts) -> ()");
 #if defined(CUDA_KERNEL) || defined(ROCM_KERNEL)
   ops.impl("mak_run", torch::kCUDA, &mak_run);
   ops.impl("mak_run_phased", torch::kCUDA, &mak_run_phased);
   ops.impl("mak_num_blocks", torch::kCUDA, &mak_num_blocks);
   ops.impl("mak_bw_per_sm", torch::kCUDA, &mak_bw_per_sm);
+  ops.impl("mak_run_ts", torch::kCUDA, &mak_run_ts);
 #endif
 }