Grouped gemm cutlass #22
Conversation
Signed-off-by: Kunshang Ji <[email protected]>
Signed-off-by: Kunshang Ji <[email protected]>
* add cutlass Signed-off-by: Kunshang Ji <[email protected]> * fix import Signed-off-by: Kunshang Ji <[email protected]> --------- Signed-off-by: Kunshang Ji <[email protected]>
Signed-off-by: Ma, Liangliang <[email protected]>
Signed-off-by: Ma, Liangliang <[email protected]>
Signed-off-by: Ma, Liangliang <[email protected]>
Signed-off-by: Ma, Liangliang <[email protected]>
Signed-off-by: Ma, Liangliang <[email protected]>
Signed-off-by: Ma, Liangliang <[email protected]>
Signed-off-by: Ma, Liangliang <[email protected]>
Signed-off-by: Ma, Liangliang <[email protected]>
Signed-off-by: Ma, Liangliang <[email protected]>
Signed-off-by: Ma, Liangliang <[email protected]>
Signed-off-by: Ma, Liangliang <[email protected]>
Signed-off-by: Ma, Liangliang <[email protected]>
Signed-off-by: Ma, Liangliang <[email protected]>
Signed-off-by: Ma, Liangliang <[email protected]>
Signed-off-by: Ma, Liangliang <[email protected]>
Signed-off-by: Ma, Liangliang <[email protected]>
Signed-off-by: Ma, Liangliang <[email protected]>
Signed-off-by: Ma, Liangliang <[email protected]>
Signed-off-by: Ma, Liangliang <[email protected]>
Signed-off-by: Ma, Liangliang <[email protected]>
Signed-off-by: Ma, Liangliang <[email protected]>
Signed-off-by: Ma, Liangliang <[email protected]>
Signed-off-by: Ma, Liangliang <[email protected]>
Fix acc and oom issue
Signed-off-by: Ma, Liangliang <[email protected]>
Signed-off-by: Ma, Liangliang <[email protected]>
Signed-off-by: Ma, Liangliang <[email protected]>
Signed-off-by: Ma, Liangliang <[email protected]>
Signed-off-by: Ma, Liangliang <[email protected]>
Signed-off-by: Ma, Liangliang <[email protected]>
CMakeLists.txt
Outdated
| FetchContent_Declare( | ||
| cutlass-sycl | ||
| GIT_REPOSITORY https://github.com/intel/cutlass-sycl | ||
| GIT_REPOSITORY https://github.com/Liangliang-Ma/cutlass-sycl |
There was a problem hiding this comment.
why using private forked cutlass-sycl?
There was a problem hiding this comment.
I will rebase to cutlass-sycl/main.
Signed-off-by: Ma, Liangliang <[email protected]>
Signed-off-by: Ma, Liangliang <[email protected]>
Signed-off-by: Ma, Liangliang <[email protected]>
Signed-off-by: Ma, Liangliang <[email protected]>
jikunshang
left a comment
jikunshang
left a comment
There was a problem hiding this comment.
some minor comments. can address in next PR.
| FUSEDMOE_AVAILABLE = True | ||
| except ImportError as e: | ||
| FUSEDMOE_UNAVAILABLE_REASON = str(e) | ||
| FUSEDMOE_AVAILABLE = False |
There was a problem hiding this comment.
maybe we need log this error even raise error directly.
| torch.ops._xpu_C.cutlass_grouped_gemm(offset=offset, N=n, K=k, **gemm_args) | ||
|
|
||
|
|
||
| def cutlass_fused_moe(hidden_states, w13, w2, topk_weights, topk_ids, |
There was a problem hiding this comment.
I think this may not be a good interface... it's ok to keep it for now.
| expert_output = input @ weight.T | ||
| ref.append(expert_output) | ||
| pre_token_sum += cur_token_num | ||
| ref = torch.cat(ref, dim=0) |
There was a problem hiding this comment.
better make this a reference function
| return result | ||
|
|
||
|
|
||
| @pytest.mark.parametrize("m,n,k", FUSED_MOE_MNK_FACTORS) |
There was a problem hiding this comment.
can you add some mini scope so we can run on simulator.
Got it. Thx |
| offset.append(0) | ||
|
|
||
| ########### gemm1 ################## | ||
| input_B = w13.transpose(-1, -2).contiguous().transpose(-1, -2) |
There was a problem hiding this comment.
You can use ColumnMajor B in the GroupGEMM kernel, so that you wouldn't have to transpose B.
To use column-major B (to avoid transposing weights), you can use a different copy atom for transposed loads . Coincidentally, because of how thread-value assignment works in the copy-atoms, the transpose copy atom for 16-bit (or 8-bits, for that matter) dtypes will load data in VNNI format (which is also true for atoms ending in _N).
However, if the latency of transpose of B + GEMM with RowMajor B is lower than GEMM with columnMajor B (highly unlikely), then you might want to retain this approach.
FWIW, once 32x32 transpose copy-atom for BF16 is added in cutlass, perf of GEMM with columnMajor B will become a bit better.
Thanks!
|
When I commented today morning, I already had this page open from yesterday, and didn't know that it had already been merged. |
8 participants
developing grouped_gemm_bf16 for Llama4-scout fused moe.
based on #11 cutlass env.