[DO NOT MERGE][Perf][Feat] Add generic cuteDSL LL FP32 router (GEMM) - #44343
[DO NOT MERGE][Perf][Feat] Add generic cuteDSL LL FP32 router (GEMM)#44343LopezCastroRoberto wants to merge 43 commits into
Conversation
Signed-off-by: LopezCastroRoberto <rocastro@redhat.com>
Signed-off-by: LopezCastroRoberto <rocastro@redhat.com>
Signed-off-by: LopezCastroRoberto <rocastro@redhat.com>
Signed-off-by: LopezCastroRoberto <rocastro@redhat.com>
Signed-off-by: LopezCastroRoberto <rocastro@redhat.com>
Signed-off-by: LopezCastroRoberto <rocastro@redhat.com>
Signed-off-by: LopezCastroRoberto <rocastro@redhat.com>
Signed-off-by: LopezCastroRoberto <rocastro@redhat.com>
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> Signed-off-by: LopezCastroRoberto <rocastro@redhat.com>
Signed-off-by: LopezCastroRoberto <rocastro@redhat.com>
Signed-off-by: LopezCastroRoberto <rocastro@redhat.com>
Signed-off-by: LopezCastroRoberto <rocastro@redhat.com>
Signed-off-by: LopezCastroRoberto <rocastro@redhat.com>
Signed-off-by: LopezCastroRoberto <rocastro@redhat.com>
Signed-off-by: LopezCastroRoberto <rocastro@redhat.com>
Signed-off-by: LopezCastroRoberto <rocastro@redhat.com>
Signed-off-by: LopezCastroRoberto <rocastro@redhat.com>
Signed-off-by: LopezCastroRoberto <rocastro@redhat.com>
Signed-off-by: LopezCastroRoberto <rocastro@redhat.com>
Signed-off-by: LopezCastroRoberto <rocastro@redhat.com>
Signed-off-by: LopezCastroRoberto <rocastro@redhat.com>
| - pytest -v -s tests/kernels/test_ll_bf16_gemm.py | ||
| - pytest -v -s tests/kernels/test_ll_fp32w_gemm.py |
There was a problem hiding this comment.
I think we should probably break up the kernel job, this gonna run alot unnecessarily
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
Signed-off-by: LopezCastroRoberto <rocastro@redhat.com>
Signed-off-by: LopezCastroRoberto <rocastro@redhat.com>
Signed-off-by: Roberto L. Castro <38211239+LopezCastroRoberto@users.noreply.github.com>
LopezCastroRoberto
changed the title
Signed-off-by: LopezCastroRoberto <rocastro@redhat.com>
Signed-off-by: LopezCastroRoberto <rocastro@redhat.com>
Signed-off-by: LopezCastroRoberto <rocastro@redhat.com>
|
|
||
|
|
||
| if __name__ == "__main__": | ||
| pytest.main([__file__, "-v"]) |
There was a problem hiding this comment.
can we unify tests with the other LL Gemms such that they live all in one file and parametrize over the dtype
There was a problem hiding this comment.
I think we can keep them separated for now, as there is not a clean merge: ll_bf16 has two variants (dotprod, splitk), requires K-dim to be divisible by 8, only covers M in {1-16} range... Since this is just a test file, I prefer not to over-complicate things
| cute.autovec_copy(bt0, br0) | ||
| br0_f32 = br0.load().to(cutlass.Float32) | ||
|
|
||
| if const_expr(EPB == 2): |
There was a problem hiding this comment.
i'd prefer non capital argument instead of EPB
| tidx: cutlass.Int32, | ||
| local_wid: cutlass.Int32, | ||
| M_PER_GROUP: cutlass.Constexpr, | ||
| M_OFFSET: cutlass.Int32, |
There was a problem hiding this comment.
same as above (stylistic advice)
| a_dtype: torch.dtype | ||
| token_groups: int = 1 | ||
| epb: int = 1 | ||
|
|
There was a problem hiding this comment.
do we want to inline CompileKey within class or keep it out of it?
There was a problem hiding this comment.
I migrated the kernel to the warmup contract already, so this shouldn't be a problem anymore
| _LL_BF16_WARMUP_M_RANGE = range(1, 17) | ||
|
|
||
| _LL_FP32W_WARMUP_M_RANGE = range(1, 17) | ||
|
|
There was a problem hiding this comment.
Do we need such large warmup range considering we restrict dispatch within gate_linear dispatch?
There was a problem hiding this comment.
We need to actually increase it up to 32 for this kernel :). But this shouldn't be a big deal in terms of compilation time.
|
This pull request has merge conflicts that must be resolved before it can be |
Signed-off-by: LopezCastroRoberto <rocastro@redhat.com>
Signed-off-by: Roberto L. Castro <38211239+LopezCastroRoberto@users.noreply.github.com>
Signed-off-by: LopezCastroRoberto <rocastro@redhat.com>
Signed-off-by: LopezCastroRoberto <rocastro@redhat.com>
Signed-off-by: LopezCastroRoberto <rocastro@redhat.com>
LopezCastroRoberto
changed the title
Signed-off-by: LopezCastroRoberto <rocastro@redhat.com>
3 participants
Depends on: #50174
Motivation
This kernel is intended to be a generic low-latency fp32-router kernel. We should not expect it to outperform highly specialized PTX/CUDA kernels tuned for one exact router shape. The goal is different: make the CuTeDSL implementation competitive with those specialized kernels while keeping the code flexible enough to support new router shapes without writing a new kernel each time.
That flexibility matters because router shapes keep changing across new models. Without a dedicated shape-specific CUDA path, the fallback is typically cuBLAS/
torch.mm, and the kernel-wise results show why a generic low-latency router kernel is still useful: even when it only ties or slightly trails the specialized dispatch kernel, it is dramatically faster than cuBLAS across these low-M router workloads.Kernel-level results
Speedup is computed as
baseline_us / ll_fp32w_us; values above1.0xmeanll_fp32wis faster.dispatch==fp32_router_gemm_dispatchMiniMax-M3
(K=6144, N=128)(384,1,1)(384,1,1)(384,1,1)(384,1,1)(384,1,1)(384,2,1)(256,1,1)(384,2,1)(384,1,1)(384,2,1)(384,1,1)(192,2,1)(384,1,1)(192,2,1)(256,1,1)(192,2,1)(384,1,1)(192,2,1)(256,1,1)(192,2,1)(256,1,1)(192,2,1)(384,1,1)(192,2,1)(384,1,1)(128,2,1)(384,1,1)(128,2,1)(384,1,1)(128,2,1)(384,1,1)(192,2,1)GLM5.2
(K=6144, N=256)(384,1,1)(384,1,1)(384,1,1)(128,1,1)(128,1,2)(128,1,1)(128,1,2)(128,2,2)(128,1,2)(128,2,2)(128,1,1)(128,2,2)(256,1,2)(128,2,2)(256,1,2)(128,2,2)(256,1,2)(128,2,2)(256,1,2)(128,2,2)(256,1,2)(128,2,2)(256,1,2)(128,2,2)(256,1,2)(128,2,2)(256,1,2)(128,2,2)(256,1,2)(128,2,2)(256,1,2)(128,2,2)E2E serving results (MiniMax-M3)