Linalg/XeGPU llama3 forward-pass example - #235
Conversation
charithaintc
left a comment
charithaintc
left a comment
There was a problem hiding this comment.
Overall looks good. I will wait for other reviews.
| return x / (1.0 + np.exp(-x)) | ||
|
|
||
|
|
||
| def _rope_tables(T, hs, theta=10000.0): |
There was a problem hiding this comment.
why is the meaning of _ prefix? not sure LH use it.
There was a problem hiding this comment.
not in examples, no. I don't mind the naming style, either is fine.
There was a problem hiding this comment.
worth checking if the existing test using this function runs OK e2e.
tkarna
left a comment
tkarna
left a comment
There was a problem hiding this comment.
Looks good overall, minor comments. The lowering schedule seems quite hard-coded/brittle but we can improve it later on.
| # ---- matmul: a(M,K) f16 @ b(K,N) f16 -> (M,N) f32 buffer ---- | ||
| def matmul(self, a, b, M, N, out_buf=None): | ||
| # Standard C = A @ B. `times_weights` emits linalg.matmul; we first fill the | ||
| # accumulator with 0. f16 inputs, f32 output -- matches the DPAS hardware. |
There was a problem hiding this comment.
Can you please use standard python docstrings, here and elsewhere, e.g.
def function(a, b):
"""Short description
Detailed description.
"""| 'mm' = matmul (linalg.matmul) -> DPAS systolic-array kernel | ||
| 'rms' = RMSNorm (2 generics + 1 fill) -> reduction kernel (uses shared mem) | ||
| 'fa' = flash multi-head attention -> one kernel (QK^T->softmax->@V, | ||
| online-softmax over K/V tiles; causal mask added by the schedule). | ||
| 'ew' = elementwise (cast / silu / mul / residual) -> row-parallel kernel |
There was a problem hiding this comment.
Can we make the kinds strings more human readable, e.g. "mm" -> "matmul", like you've already spelled then out in the llama3_schedule.py header.
| from llama3_payload import F32 | ||
|
|
||
|
|
||
| def _tile_one_matmul(matmul_op, anytype, mm_params): |
There was a problem hiding this comment.
Why does this function have anytype as an arg? It's not used anywhere?
| lh_transform.tile(wg_matmul, tile_sizes=[0, 0, mm_params["k_tile"]]) | ||
|
|
||
|
|
||
| def _tile_one_fused_attention_region(anytype, qkt_bmm, pv_bmm, softmax_op, fa_params): |
There was a problem hiding this comment.
anytype is an arg here as well but we can just create it where needed, or define one in the module.
| lh_transform.tile(wg_matmul, tile_sizes=[0, 0, mm_params["k_tile"]]) | ||
|
|
||
|
|
||
| def _tile_one_fused_attention_region(anytype, qkt_bmm, pv_bmm, softmax_op, fa_params): |
There was a problem hiding this comment.
nit: Would be good to mention in the docstrings what the input args are. E.g., qkt_bmm is a handle of the QK^T linalg.generic op?
In lighthouse we also tend to use type annotations, so something like
def _tile_one_fused_attention_region(
qkt_bmm: ir.Value, pv_bmm: ir.Value, softmax_op: ir.Value, fa_params: dict
) -> tuple[ir.Value, ir.Value]| gen_handles = transform.split_handle( | ||
| (anytype,) * ngen_total, match(mod, ops={"linalg.generic"}) | ||
| ) |
There was a problem hiding this comment.
this is just match_and_split, no?
| # Fusion leaves the full-size original fill DEAD at func scope (fusion only | ||
| # slices a copy inside the forall). It must be removed or the next rms finds too | ||
| # many. canonicalize (DCE) at func scope, but never apply_cse at func scope -- |
There was a problem hiding this comment.
Maybe you could just apply dead code elimination on it own, transform.apply_dce(target) ?
| ln_params["T"], | ||
| ) | ||
|
|
||
| # 2) Tile ew generics into own foralls (handles preserved across rms tiling). |
There was a problem hiding this comment.
AFAIU this creates standalone gpu kernels for the elementwise ops. Could they be fused into their producer or consumer matmuls instead?
| # NUMPY REFERENCE -- the same math in plain numpy, to CHECK the GPU result. | ||
| # `_f16` rounds through float16 to model the GPU's f16 matmul precision. | ||
| # ============================================================================= | ||
| def _rms(x, weight, eps=1e-5): |
There was a problem hiding this comment.
nit: as these helpers seem to be related to constructing the numpy reference model, maybe add _numpy in the func name to distinguish from the actual graph building methods (or wrap in a class to define a namespace).
| an=np.ones(C, np.float32), | ||
| wq=(np.random.randn(C, C) * sc).astype(np.float16), | ||
| wk=(np.random.randn(C, kv_dim) * sc).astype(np.float16), | ||
| wv=(np.random.randn(C, kv_dim) * sc).astype(np.float16), | ||
| wo=(np.random.randn(C, C) * sc).astype(np.float16), | ||
| fn=np.ones(C, np.float32), | ||
| w1=(np.random.randn(C, hidden) * sc).astype(np.float16), | ||
| w2=(np.random.randn(hidden, C) * sc).astype(np.float16), | ||
| w3=(np.random.randn(C, hidden) * sc).astype(np.float16), |
There was a problem hiding this comment.
nit: maybe these dict keys could also be made more human readable?
4 participants
This PR adds a Llama-3 model script under examples/xegpu, a full Llama-3-style transformer forward pass (6 layers, C=256, H=4 query heads / n_kv=2 KV heads, head_size=64, hidden=1024, T=256) running end-to-end on the Intel GPU via the XeGPU lowering path.
Assisted by Claude