model-as-a-kernel v1 source
Browse files
mak_cuda/megakernel.cu
CHANGED
|
@@ -1,4 +1,4 @@
|
|
| 1 |
-
|
| 2 |
//
|
| 3 |
// The kernel is a phase interpreter. A "program" (int64 [n_phases, 16] array)
|
| 4 |
// lists every operation of the forward pass: embedding read, fused
|
|
|
|
| 1 |
+
// model-as-a-kernel: an entire llama-family decode step as one kernel launch.
|
| 2 |
//
|
| 3 |
// The kernel is a phase interpreter. A "program" (int64 [n_phases, 16] array)
|
| 4 |
// lists every operation of the forward pass: embedding read, fused
|
torch-ext/model_as_a_kernel/__init__.py
CHANGED
|
@@ -1,4 +1,4 @@
|
|
| 1 |
-
|
| 2 |
|
| 3 |
A phase interpreter runs the whole forward pass (embedding, every layer's
|
| 4 |
norms, QKV, rope, attention over the KV cache, SwiGLU MLP, the LM head, and
|
|
|
|
| 1 |
+
"""model-as-a-kernel: an entire llama-family decode step as one kernel launch.
|
| 2 |
|
| 3 |
A phase interpreter runs the whole forward pass (embedding, every layer's
|
| 4 |
norms, QKV, rope, attention over the KV cache, SwiGLU MLP, the LM head, and
|
torch-ext/torch_binding.cpp
CHANGED
|
@@ -1,4 +1,4 @@
|
|
| 1 |
-
|
| 2 |
|
| 3 |
#include <ATen/cuda/CUDAContext.h>
|
| 4 |
#include <c10/cuda/CUDAGuard.h>
|
|
|
|
| 1 |
+
#include <torch/library.h>
|
| 2 |
|
| 3 |
#include <ATen/cuda/CUDAContext.h>
|
| 4 |
#include <c10/cuda/CUDAGuard.h>
|