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.
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.
| 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?
| 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?
tkarna
left a comment
tkarna
left a comment
There was a problem hiding this comment.
Looks good to me.
It seems the correctness test is not passing anymore with the latest mlir version. Let's wait for a fix upstream before merging.
llvm/llvm-project#224459 fixes this |
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