When Dynamo Meets JIT: The Architecture Clash That Blocked Piecewise CUDA Graphs

In the high-stakes world of large language model inference optimization, few things are more frustrating than a promising optimization path that collapses under the weight of an architectural incompatibility. Message [msg 996] captures exactly such a moment: the instant when an AI assistant, deep in a systematic optimization campaign for the GLM-5-NVFP4 model running on eight NVIDIA RTX PRO 6000 Blackwell GPUs, realizes that the piecewise CUDA graph optimization it has spent hours trying to enable is fundamentally incompatible with the FlashInfer FP4 quantization library upon which the entire model depends.

This message is a turning point — a moment of diagnostic clarity that terminates one line of inquiry and redirects effort toward others. It is also a masterclass in how to reason about deep software integration failures in the ML infrastructure space.

The Optimization Campaign

To understand message [msg 996], we must first understand the context. The assistant had been systematically working through a prioritized list of performance optimizations for the GLM-5-NVFP4 model, which uses 4-bit floating point (FP4) quantization via the FlashInfer library. After establishing a baseline benchmark across four concurrency levels (1, 10, 256, and 1024 concurrent requests), the assistant turned to "Tier 1.1: Piecewise CUDA Graphs" — a feature in SGLang that captures CUDA graphs for segments of the model to reduce kernel launch overhead.

The piecewise CUDA graph runner in SGLang works by using PyTorch's torch.compile infrastructure — specifically torch._dynamo — to trace through the model's forward pass, capture the sequence of operations, and compile them into reusable CUDA graphs. This is a powerful optimization technique, but it comes with a critical constraint: every operation in the traced path must be traceable by Dynamo, PyTorch's graph tracing engine.

The First Failure

The assistant's first attempt to enable piecewise CUDA graphs crashed with a torch._dynamo.exc.Unsupported error. The root cause was that FlashInfer's FP4 quantization module calls subprocess.check_output(["nvcc", "--version"]) during its JIT compilation process, and this subprocess call creates a thread lock that Dynamo cannot trace. The assistant diagnosed this problem, traced it to the get_cuda_version() function in FlashInfer's cpp_ext.py, and applied a surgical patch: rewriting the function to use torch.version.cuda directly instead of shelling out to nvcc.

This was a reasonable fix. The get_cuda_version() function was decorated with @functools.cache, so once called, the result should have been cached. But Dynamo has a documented behavior of tracing through functools.lru_cache-wrapped functions — it "ignores the cache wrapper and directly traces the wrapped function," as the warning message in the logs explicitly stated. The patch was necessary but, as we will see, insufficient.

Message 996: The Second Failure

After patching and restarting the server, the assistant waited three minutes for the model to load and the graph capture to begin, then checked the logs. The result is captured in message [msg 996]:

Still failing, but now at a different point — os.stat(self) in pathlib.py. Dynamo can't trace file system operations inside the FP4 JIT code. The fundamental issue is that the piecewise CUDA graph capture tries to torch.compile the dense layers, but the dense layers contain FP4 quantization (via modelopt_quant.py), and FlashInfer's FP4 quantization module does JIT compilation with file system and subprocess operations that dynamo can't trace.

This is the moment of generalization. The assistant has moved from diagnosing a specific error (subprocess calls) to understanding the fundamental architectural conflict. The problem is not one particular function call — it is the entire design of FlashInfer's FP4 quantization module, which performs JIT compilation at runtime, involving file system operations (os.stat, path.exists), subprocess calls (nvcc --version), and other operations that Dynamo cannot trace.

The Hypothesis and Its Refutation

The assistant then forms a new hypothesis:

The FP4 quantize function needs to be registered as a custom op so dynamo treats it as opaque. Let me check if it already is.

This is a sophisticated hypothesis. In PyTorch's Dynamo tracing system, operations that are registered as "custom ops" via torch.library or torch._dynamo.allow_in_graph are treated as opaque — Dynamo does not attempt to trace through them, instead recording them as atomic operations in the graph. If the FP4 quantization functions were already registered as custom ops, Dynamo would not try to trace through their internal JIT compilation logic.

The assistant runs a grep to check:

[bash] ssh root@10.1.230.174 'grep -n "register_custom_op\|torch.library\|custom_op\|allow_in_graph" /root/ml-env/lib/python3.12/site-packages/flashinfer/fp4_quantization.py | head -10'
41:    register_custom_op,
153:    @register_custom_op(
226:    @register_custom_op(
258:    @register_custom_op(
293:    @register_custom_op(
375:    @register_custom_op(
470:    @register_custom_op(
559:    @register_custom_op(

The result is striking. Not only is register_custom_op imported, but it is used extensively — at least seven times in the file. The FP4 quantization functions are registered as custom ops. But the problem persists.

Why? Because the JIT loading code runs before the custom op is invoked. The get_fp4_quantization_module() function calls build_and_load(), which checks is_aot, which calls path.exists() — all before the custom op registration takes effect. Dynamo traces through the forward pass, which calls the custom op, but the custom op's implementation internally triggers JIT compilation during its first invocation. Even though the op itself is opaque to Dynamo, the code path that Dynamo traces includes the setup of the op, which involves file I/O.

This is the critical insight: the custom op registration protects Dynamo from tracing the implementation of the op, but not from tracing the initialization code that loads the JIT module. And that initialization code is full of Dynamo-incompatible operations.

The Deeper Architectural Insight

Message [msg 996] reveals a fundamental tension between two different approaches to GPU kernel management. FlashInfer's FP4 quantization uses a JIT compilation model: CUDA C++ source files are compiled at runtime using nvcc, and the resulting shared objects are loaded dynamically. This requires file system access, subprocess execution, and other operations that are inherently incompatible with PyTorch's Dynamo tracing.

SGLang's piecewise CUDA graph runner, on the other hand, relies on Dynamo to trace through the model's operations and capture them as CUDA graphs. This requires that every operation in the traced path be "Dynamo-safe" — meaning it must not perform I/O, spawn threads, or call into C extensions that Dynamo cannot introspect.

These two design philosophies are fundamentally at odds. FlashInfer's JIT approach prioritizes flexibility and portability (CUDA kernels can be compiled for whatever GPU architecture is present), while SGLang's graph capture approach prioritizes traceability and determinism. When both are used in the same model, they conflict.

Assumptions and Their Limits

The assistant made several reasonable assumptions during this debugging session, each of which proved partially wrong:

  1. That patching get_cuda_version() would be sufficient. This fixed one specific Dynamo-incompatible operation, but the FP4 JIT code contains many more such operations (file system calls, path checks, etc.).
  2. That the @functools.cache decorator would prevent Dynamo from re-executing the cached function. In fact, Dynamo deliberately traces through cached functions, as its own warning message explains.
  3. That registering FP4 ops as custom ops would make Dynamo treat them as opaque. This is true for the op's implementation, but the JIT loading code runs during op initialization, which occurs within the traced path.
  4. That the problem was in FlashInfer's fp4_quantization.py specifically. The deeper issue is that the entire JIT compilation pipeline — from path resolution to nvcc invocation to shared object loading — is incompatible with Dynamo tracing.

The Knowledge Created

Message [msg 996] creates several important pieces of knowledge:

Output knowledge: The assistant now knows that piecewise CUDA graphs are blocked for this model configuration. The error signature (os.stat(self) in pathlib.py during Dynamo tracing) and the confirmed presence of custom op registrations provide a complete diagnostic picture. This knowledge directly informs the decision to abandon piecewise CUDA graphs and move to other Tier 1 optimizations (MSCCLPP, Single Batch Overlap, Expert Parallelism).

Input knowledge required: To fully understand this message, one must understand PyTorch's Dynamo tracing system, the concept of custom ops and how they interact with graph tracing, FlashInfer's JIT compilation model for FP4 quantization, and SGLang's piecewise CUDA graph runner architecture. The message assumes familiarity with torch.compile, torch._dynamo, and the register_custom_op mechanism.

The Thinking Process

The assistant's reasoning in this message follows a clear diagnostic pattern:

  1. Observe: The error has changed — now it's os.stat(self) in pathlib.py, not the subprocess issue from before.
  2. Generalize: The pattern is that Dynamo cannot trace file system operations. This is not a one-off bug but a fundamental incompatibility.
  3. Formulate root cause: The piecewise CUDA graph capture uses torch.compile on dense layers that contain FP4 quantization, and FlashInfer's FP4 module does JIT compilation with file I/O and subprocess operations.
  4. Hypothesize fix: Registering the FP4 quantize function as a custom op would make Dynamo treat it as opaque.
  5. Test hypothesis: Run a grep to check if custom ops are already registered.
  6. Interpret result: Custom ops ARE registered — the hypothesis was correct in theory but insufficient in practice, because the JIT loading code runs before the custom op is invoked. This is a textbook example of scientific debugging: form a hypothesis, test it, and refine your understanding based on the results. The assistant does not panic or try random fixes — it methodically builds a mental model of the system's architecture and uses that model to predict where the next problem will be.

The Broader Lesson

Message [msg 996] illustrates a common challenge in ML infrastructure: the tension between JIT compilation and graph tracing. JIT compilation is a powerful technique for supporting diverse hardware without pre-compiling kernels for every possible configuration. Graph tracing is a powerful technique for reducing runtime overhead by capturing and reusing execution plans. But they pull in opposite directions — JIT wants to defer decisions to runtime and perform I/O to compile kernels, while graph tracing wants to capture everything ahead of time and never deviate from the captured plan.

For the GLM-5-NVFP4 model on Blackwell GPUs, this tension means that piecewise CUDA graphs are not a viable optimization path. The assistant will need to look elsewhere — to Expert Parallelism, to MSCCLPP-based allreduce, or to more fundamental changes in how the FP4 GEMM kernels are scheduled. The debugging effort was not wasted: understanding why an optimization doesn't work is often as valuable as making it work, because it prevents wasted effort on similar approaches in the future.

This message, then, is not a failure — it is a successful diagnosis that saves the assistant from pursuing a dead end. It is the moment when a promising hypothesis meets reality and is refined into deeper understanding.