Ran into a fun issue recently testing a custom CuteDSL kernel I had written for an SSM based layer. The kernel was getting used from within vLLM and the model would generate tokens normally when run without CUDAGraph capture (“CG off”). However, when run with CUDAGraph capture (“CG on”), the very first decoded token would be fine but the following tokens would be garbage. The fix ended up being straightforward and I have no one else to blamemaybe except for Claude Code for not locating this section of the CuteDSL documentation sooner, but making sense of the bug took a little bit of profiling.
The setup is a CuteDSL kernel using TVM-FFI and running in vLLM with CUDAGraphs. And the fix is to tell CuteDSL to use the current PyTorch stream by setting the use_tvm_ffi_env_stream flag like so:
from cutlass import cute
kernel_args = ...
stream = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True)
compiled_kernel = cute.compile(
kernel, kernel_args, stream, options="--enable-tvm-ffi"
)
If we skip this explicit stream assignment because we are not good at reading documentation, here is what happens. With CG on, vLLM creates a new CUDA stream to capture the model graph (kernels/ops in the model’s forward pass) as opposed to using the pre-existing CUDA “Default Stream”. During the capture, the torch ops in the model are recorded on the “Capture Stream” but they are executed on the Default Stream during decoding. But since we didn’t tell CuteDSL/TVM-FFI to latch on to the torch stream (which during vLLM graph capture is the Capture Stream), the CuteDSL kernel runs on the Default Stream instead and is not recorded in the graph. We can verify this by comparing the NSys profiles with CG on v/s CG off during capture time:
And this leads to garbage tokens during decode because when the CG is replayed, our new fancy CuteDSL kernel is neatly skippedNote that the orange box in Figure $1$ also differs between the two profiles. This corresponds to the memory estimation pass that vLLM runs on yet another non-Default stream which is also affected by our stream issue. . This can be seen in the NSys profiles during decode time:
kernel_cutlass row is the CuteDSL kernel. The replayed graph never runs our kernel on the buggy profile while the fixed profile shows it executing.
See that empty space in the CG on profile in Figure $2$? That is where our CuteDSL kernel should have been but wasn’t because it was never recorded on the graph. The model’s forward skips our kernel entirely and uses stale values for the tensors the kernel was supposed to update/output. And that is not ideal.
The attentive reader would note that we still have at least one more outstanding question, which is, how come the very first token was generated correctly if we never recorded the kernel during graph capture? This is where we run into vLLM’s various CUDAGraphModes. The default for the v1 vLLM engine is FULL_AND_PIECEWISE which means FULL capture for decoding and PIECEWISE capture for prefill. The PIECEWISE capture for prefill ends up calling our kernel in “eager” mode (i.e. outside CG) and the first decoded token is generated by this eager prefill call (which runs on the Default Stream as desired).
In the profile in Figure $3$ below, we can see that the prefill looks similar for CG on vs CG off, but the decode has missing CuteDSL kernel calls in the buggy profile (Figure $2$ is basically zooming in to the decode portion of Figure $3$).
