fix piecewise compilation in the ubatch wrapper

Signed-off-by: Sage Moore <[email protected]>

Author: SageMoore

vllm/compilation/ubatch_wrapper.py

@@ -253,7 +253,7 @@ def __call__(self, *args, **kwargs):
 
         # If there's no ubatching, just run the runnable object
         if ubatch_slices is None:
-            if cudagraph_runtime_mode is CUDAGraphMode.NONE:
+            if cudagraph_runtime_mode in (CUDAGraphMode.NONE, CUDAGraphMode.PIECEWISE):
                 return self.runnable(*args, **kwargs)
             else:
                 assert self.cudagraph_wrapper is not None