model-as-a-kernel v1 source
Browse files
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
|