[torch.compile] support moe models #9632
Conversation
Signed-off-by: youkaichao <[email protected]>
Signed-off-by: youkaichao <[email protected]>
Signed-off-by: youkaichao <[email protected]>
Signed-off-by: youkaichao <[email protected]>
Signed-off-by: youkaichao <[email protected]>
|
👋 Hi! Thank you for contributing to the vLLM project. Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging. To run CI, PR reviewers can do one of these:
🚀 |
| ["--quantization", "compressed-tensors" | ||
| ], 1, 1, "FLASH_ATTN", "generate", True), | ||
| ("google/gemma-2-2b-it", [], 1, 2, "FLASHINFER", "generate", True), | ||
| ("ibm/PowerMoE-3b", [], 1, 2, "FLASH_ATTN", "generate", True), |
There was a problem hiding this comment.
it does not work with FLASHINFER due to some head size config.
| gating_output=router_logits, | ||
| renormalize=renormalize) | ||
|
|
||
| forward_native = forward_cuda |
There was a problem hiding this comment.
when we use inductor, we will use forward_native . this function needs to be implemented.
There was a problem hiding this comment.
you can also write a pytorch native implementation, but I doubt if inductor can optimize it. that's why I use forward_cuda as forward_native
There was a problem hiding this comment.
basically, forward_native is the function we compile when we use inductor.
There was a problem hiding this comment.
I see. I was just surprised because I didn't know that a class method can be defined this way
There was a problem hiding this comment.
So it's essentially the same as
def forward_native(self, *args, **kwargs):
return self.forward_cuda(*args, **kwargs)?
Signed-off-by: youkaichao <[email protected]>
Signed-off-by: youkaichao <[email protected]> Signed-off-by: qishuai <[email protected]>
Signed-off-by: youkaichao <[email protected]> Signed-off-by: Randall Smith <[email protected]>
Signed-off-by: youkaichao <[email protected]> Signed-off-by: Sumit Dubey <[email protected]>
Signed-off-by: youkaichao <[email protected]>
Signed-off-by: youkaichao <[email protected]> Signed-off-by: LeiWang1999 <[email protected]>
2 participants
moe models will read config file to determine the triton config to run.
reading files during forward is a disaster for
torch.compile.this pr wraps the config reading part inside a custom op, so that it can pass
torch.compile(althoughtorch.compilewill not be able to optimize it).