vllm - 💡(How to fix) Fix [RFC][vLLM IR] `rms_norm` weight passing inconsistency [1 participants]

Problem

In some models, the RMSNorm layer is unweighted (has_weight=True), effectively the same as the weight tensor set to all 1s. This can also be implemented by not applying the weight tensor at all. However, some implementations (vllm_c, aiter) always require the weight tensor. While always passing the weight tensor does not seem to affect native implementation performance, it does break TPU support, because we don't pass weight as a parameter so the XLA/JAX tensor conversion doesn't port the tensor over to TPU. Passing the parameter makes torch.compile unit tests start doing backwards through the model, which we can mitigate by using @torch.inference_mode or torch.no_grad(). But this still has the issue of passing the weight where we know that weight is all ones and not needed.

We cannot include weight in supports_args because we don't want an implementation to be skipped because it doesn't support weight=None. We also don't want to allocate as it might affect eager-mode performance.

Possible solutions

• Current solution: pass weight conditionally, knowing which implementation will get selected. This does not work for fused ops. Perhaps we can do this for TPU-only?
• Alternative 1: Pass the weight tensor as a parameter and add torch.inference_mode to torch.compile unit tests
• Alternative 2: Introduce a "workspace" concept where implementations are allowed to one-time allocate workspace tensors (and possibly size them up if needed). The difficulty here is that this would happen after Dynamo tracing because it would have to happen after implementation selection during lowering. But this would be the most long-term fix as workspaces are needed elsewhere as well.
• Alternative 3: Just allocate weight in implementations that need it and eat the overhead.
• Alternative 4: Add a has_weight=bool parameter to the IR op: this would pollute the op in my opinion
• Alternative 5: Rewrite the CUDA and AITER kernels to support weight=None. This would work but isn't scalable, what if we run into something similar for other kernels?

I honestly don't know how to think through all the tradeoffs so would like some input.

CC

@lk-chen @tjtanaa @zou3519