👋 Hi! Thank you for contributing to the vLLM project.

💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels.

Just a reminder: PRs would not trigger full CI run by default. Instead, it would only run fastcheck CI which starts running only a small and essential subset of CI tests to quickly catch errors. You can run other CI tests on top of those by going to your fastcheck build on Buildkite UI (linked in the PR checks section) and unblock them. If you do not have permission to unblock, ping simon-mo or khluu to add you in our Buildkite org.

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 either: Add ready label to the PR or enable auto-merge.

🚀

/gemini review

/gemini review

I'm not sure if this is the best approach. The CustomOp abstraction is meant as a simple abstraction to dispatch between torch implementations and custom/CUDA/platform kernels. For something as complex as MLA we should not use this.

Could you describe the fusion you want to do in more detail? Could this fusion in vllm-ascend instead be performed using a custom torch.compile pass?

EDIT: I think if you want to extract MLA code from the model code I think that's a great idea. But the custom op abstraction seems like the wrong tool for the job.

Thank you for your suggestion. To help clarify, I've created a diagram to demonstrate an example of one of the fused operators we intend to integrate：

Thank you for the picture. Honestly this is a common issue with attention, there can be a lot of ops hidden from torch.compile. There are two possible approaches:

1. Add custom passes that "unfuse" these ops out of attention. The passes would look similar to attention+fp8 quant fusion. This would suffer from conflict with piecewise compilation, but we're working to address that in [RFC]: Address piecewise graph splitting and attention fusion incompatibility #23261.
2. Add an extra forward layer to AttentionImpl where instead of calling the custom op directly, forward_outer is called on the impl. forward_outer just calls the custom op by default but backends can specialize it with other ops preceeding/succeeding the custom op, so that those ops are shown in .

Both approaches would require modifying the unified_attention op, and I'm not sure which is better, and if either is feasible at all.

Before we go down this rabbit hole, is Inductor supported on vllm-ascend? What about cuda graphs?

Before we go down this rabbit hole, is Inductor supported on vllm-ascend? What about cuda graphs?

Thanks for your attention. Currently, vLLM-Ascend does not support Inductor. We have enabled a graph mode whose functionality is very similar to CUDA Graph. At present, performing fusion passes in vLLM-Ascend cannot be achieved through the mechanism implemented via Inductor in vLLM. We only have a practice of performing kernel fusion through torch.fx.subgraph_rewriter, but I am not sure whether this functionality possesses the full capabilities of the pattern_matcher in Inductor.

2. Add an extra forward layer to AttentionImpl where instead of calling the custom op directly, forward_outer is called on the impl. forward_outer just calls the custom op by default but backends can specialize it with other ops preceeding/succeeding the custom op, so that those ops are shown in .

To ensure I correctly understand the second approach you mentioned, I’ll briefly describe the implementation.
First we add a forward_outer abstract to AttentionImpl:

@abstractmethod
def forward_outer(
    self,
    positions: torch.Tensor,
    hidden_states: torch.Tensor,
) -> torch.Tensor:
    raise NotImplementedError

Then we should add a static method to the CustomOp class to check whether a given op name has already been registered:

@staticmethod
def has_registered(op_name):
    """Check if the class has a registered name."""
    return op_name in {**CustomOp.op_registry, **CustomOp.op_registry_oot}

Finally, in the abstracted MLA layer, we check if the mla_outer op has been registered. If it has, we call the forward_outer method:

def forward(
    self,
    positions: torch.Tensor,
    hidden_states: torch.Tensor,
) -> torch.Tensor:
    if CustomOp.has_registered("mla_outer"):
        return self.mla_attn.impl.forward_outer(positions, hidden_states)

Does this align with what you had in mind? @ProExpertProg

Not exactly, but it would still require Inductor fusions anyway.

Does vllm-ascend support (or use) Dynamo at all? What about AotDispatcher? Because we could do passes without Inductor if AotDispatcher is used (so normalized and functional IR is produced).

How are the CUDAGraph-equivalent graphs handled?

Not exactly, but it would still require Inductor fusions anyway.

Does vllm-ascend support (or use) Dynamo at all? What about AotDispatcher? Because we could do passes without Inductor if AotDispatcher is used (so normalized and functional IR is produced).

How are the CUDAGraph-equivalent graphs handled?

vllm-ascend supports and uses Dynamo. We only implement our own piecewise backend to receive and process FX Graph. Our graph mode also works through capture and replay machanism, just like CUDA-Graph.

after chatting with @ProExpertProg I think we agree that making layers pluggable (especially MLA for the reasons you outlined) by different HW backends without the need for torch.compile is desirable; however CustomOP might not be the right mechanism for this (per @ProExpertProg comments)

I will be honest I don't know that much about how the HW pluggin infrastructure works so maybe @youkaichao can weigh in here

Yeah, to elaborate, I think a custom pass mechanism to perform fusion would be good in vllm-ascend because plugging layers and fusing manually inside will quickly become unmaintainable due to duplicated logic. It would also suffer from the same reasons vLLM uses custom passes, and fusions across layers are still going to be difficult.

However, I understand that might be too large an undertaking for this case. I still believe CustomOp is not the right abstraction because it interferes with vllm's custom op enablement mechanism in a slightly ugly way. I think we should come up with a lightweight PluggableLayer abstraction that's transparent in the vllm case but can be replaced OOT. But if @youkaichao thinks we should just use CustomOp I'm okay with that.

Thank you for your input! I completely agree that we need a more lightweight PluggableLayer to support custom model layers—this would significantly enhance the extensibility of the vLLM project through plugins. However, to implement PluggableLayer effectively, we need to conduct more detailed design work. We will subsequently propose an RFC to describe and further discuss this direction. For the current MLA-related issues, I believe we can temporarily address them using the Custom Op approach.

I will be honest I don't know that much about how the HW pluggin infrastructure works

@LucasWilkinson Let me do a simple elaboration to show how hardware plugin work:

0. Register platform: vLLM serve start with plugin register and regiter current platform.

1. Register worker: vLLM call current_platform.check_and_update_config to config worker class.

2. Init worker: vLLM init worker, this will init platform's worker:

3. Register Custom ops: vLLM hardware plugin (such as vllm-ascend) worker will call register_ascend_customop to register custom ops.

After all above steps, the ops will be replaced by hardware plugin own ops. Actually, the current custom ops mechanism is relatively flexible and supports registration at the forward, layer, and ops levels.

I think the PluggableLayer actually it is equal to custom ops with layer level + layer stable interface, This is worth describing in a detailed separate RFC (this related to #22082 and our vLLM 1.0 goal).

So, from the perspective of hardware plugin, I think custom ops is an acceptable way to go, actually, it's also the original intention of #19164

Yeah, to elaborate, I think a custom pass mechanism to perform fusion would be good in vllm-ascend because plugging layers and fusing manually inside will quickly become unmaintainable due to duplicated logic. It would also suffer from the same reasons vLLM uses custom passes, and fusions across layers are still going to be difficult.

However, I understand that might be too large an undertaking for this case. I still believe CustomOp is not the right abstraction because it interferes with vllm's custom op enablement mechanism in a slightly ugly way. I think we should come up with a lightweight PluggableLayer abstraction that's transparent in the vllm case but can be replaced OOT. But if @youkaichao thinks we should just use CustomOp I'm okay with that.

I guess both paths (the frontend pluggable layer abstraction and the backend graph fusion) can be useful. The former is direct and simpler while we have to maintain the compatibility of the layer semantics - therefore, the supported layers had better be stable enough and well-defined, e.g., MLA and fused MoE. The latter is more flexible and general but it depends on the torch.compile graph mode with extra compilation overhead as well as the complexity of graph rewrite.

From vllm-ascend, we are trying to enable some graph fusion passes by extending the vllm core. But it seems that the compilation backend is currently bound to the "inductor" backend which not every HW backend supports. In vllm-ascend, the inductor backend is not yet enabled and we are adding direct FX-based graph fusion. Currently, we have to do monkey-patches to plugin our own passes (see the rmsnorm+quant fusion PR from
@ganyi1996ppo: vllm-project/vllm-ascend#2389). We also want to reuse the existing backend-agnostic FX passes from vllm core, e.g., sequence parallel fusion. But these are all "inductor" passes.

In summary, it would be great if

1. A HW backend can extend vllm to support custom compiler backend that implements CompilerInterface.
2. A HW backend with custom compiler backend can reuse existing backend-agnostic FX passes..

Would love to get your inputs. @ProExpertProg

PS: For "2", perhaps we can manually apply the AotDispatch passes and invoke those vllm "inductor" passes from inside the custom compiler backend. Not sure if that would work well though.

@Yikun Thanks for the description! I think the main downside of using CustomOp as the plugin mechanism for layers is that it only allows you to hook into the forward call. I would suspect that for layers we would quickly run into that cases where HW plugin developers may want to do some work ahead-of-time instead of exclusively on the hot path (i.e. re-laying out weights during weight loading, or precomputing some kind of metadata). My guess would be that it would be better to have some kinda of layer abstraction that is separate from CustomOp for this reason.

@whx-sjtu is this done?

@whx-sjtu is this done?

Yes, I think we can run ci and merge this PR.

I review again on this, I also think it's ready to go.

buildkite/ci/pr/distributed-tests-2-gpus failed due to unrelated fix, @whx-sjtu Please don't re-push again, I will trigger the auto-merge check.

If there still CI issue here, we can forge merge this.