[RFC][UX]: debug mode for vLLM-compile

[zou3519]

Motivation.

vLLM-compile (CompilationLevel.PIECEWISE) makes a lot of assumptions about the models that allow it to make them run really fast. There are two main assumptions that commonly lead to silent incorrectness if the models violate them. I've spent countless hours debugging user issues for it to turn out to be one of these assumptions. We should add a debug mode option for vLLM-compile that, when turned on, adds some safety checks for these assumptions at the tradeoff of some additional overhead. This will let users self-diagnose the issues without me in the loop.

This is one of the items mentioned in #20283, I'm expanding it to include some more details.

Proposed Change.

The two assumptions that bite us are:

1. the vLLM Dynamic Shapes Issue. vLLM performs one single graph capture with dynamic batch size and expects the graph to work for all batch sizes. However, the graph may not actually be valid for all batch sizes.
2. CUDAGraphs assume that the input addresses of Tensors do not change. Changing the input addresses (e.g. doing model.weight1 = new_weight1) violates the assumption.

We should add an option to have some additional "safety checks" in CompilationConfig. Here are the safety checks that we would add.

For the vLLM Dynamic Shapes Issue

torch.compile produces a symbolic expression about what batch sizes the captured graph is valid for. These might look like "size < 60000 && size % 128 == 0". At runtime, when we execute the vLLM model, we should validate that the batch size actually passes the symbolic expression.

For CUDAGraphs input addresses

We know all of the inputs being cudagraph'ed. We can record their input addresses, and, at runtime, always check that the input addresses did not change.

Feedback Period.

7/2 - 7/11

CC List.

@ProExpertProg @youkaichao @mgoin @robertgshaw2-redhat @WoosukKwon @drisspg @houseroad

Any Other Things.

thank you

Before submitting a new issue...

• Make sure you already searched for relevant issues, and asked the chatbot living at the bottom right corner of the documentation page, which can answer lots of frequently asked questions.

[ProExpertProg]

For the dynamic shape issue:

1. Padding: I think with avoid guarding on max() unnecessarily pytorch/pytorch#149635 fixed, padded allocation can always just use max. If padding requires an additional allocation, I think torch.cond can be used. Perhaps

2. Kernel dispatch: another idea is to add a way to trace the graph multiple times with different conditions (min and max args to mark_dynamic). I worry that some of the automated solutions would further add to startup time. So I think enforcing that there's a way for a single Dynamo trace and single Inductor compilation to work (even if it slightly hurts performance) would be good.

I also think solutions that wrap into a new cond op or if_statically_known are okay even if dispatching logic is inside an external library:

if statically_known(batch_size <= 64):
    external_kernel(...)
elif statically_known(batch_size <= 32):
    external_kernel(...)
else:
    general_case(...)

As long as there's a way to solve the general case, we can dispatch in vLLM code and then by the time we invoke the external library we've properly specialized on shape.