The JIT Loading Trap: Debugging Torch Dynamo's Incompatibility with FlashInfer FP4 in Piecewise CUDA Graph Capture

Introduction

In the course of optimizing inference throughput for the GLM-5-NVFP4 model on eight NVIDIA RTX PRO 6000 Blackwell GPUs, an opencode assistant encountered a stubborn and illuminating failure. The optimization being tested—SGLang's --enable-piecewise-cuda-graph flag—promised significant performance gains by using torch.compile to capture and replay CUDA graph segments for the model's dense layers. Instead, the server crashed during graph capture with a cryptic error: torch._dynamo.exc.Unsupported: Attempted to call function marked as skipped. What followed was a multi-round debugging session that peeled back layers of abstraction, revealing a fundamental tension between PyTorch's dynamo tracing engine and FlashInfer's Just-In-Time (JIT) compilation system for FP4 quantization kernels.

Message [msg 997] captures the pivotal diagnostic moment: the assistant has just realized that the FP4 operations are registered as custom ops (which should make them opaque to dynamo), but the JIT loading code that initializes those ops runs before the custom op boundary is reached. This insight—that the problem is not the custom op itself but the file I/O and subprocess calls in the module loading path—represents a critical step in the debugging journey. This article examines that message in depth, exploring the reasoning, assumptions, and technical context that make it a fascinating case study in ML systems debugging.

The Optimization Context: Why Piecewise CUDA Graphs Matter

Before diving into the message itself, it is essential to understand the stakes. The GLM-5-NVFP4 model is a Mixture-of-Experts (MoE) architecture with 4-bit floating-point (FP4) quantization, running on NVIDIA's Blackwell SM120 architecture. The assistant had already achieved impressive throughput (~3,740 tok/s) through a combination of FlashInfer CUTLASS MoE autotuning, increased max-running-requests, and careful server parameter tuning. However, the core bottleneck remained the small per-expert GEMM operations, which are memory-bandwidth-bound on Blackwell's SM120 due to its 99KB shared memory limit and lack of TMEM support.

Piecewise CUDA graphs offered a potential path forward. By using torch.compile to capture the compute graph of the dense (non-MoE) layers into reusable CUDA graphs, the server could eliminate Python interpreter overhead and kernel launch latency for those layers. The feature is enabled via --enable-piecewise-cuda-graph in SGLang and uses a custom graph runner that compiles model segments with torch.compile, captures the resulting CUDA graphs, and replays them during inference. For a model as large as GLM-5-NVFP4, this could translate to substantial throughput gains—if it worked.

The Crash Sequence: Tracing the Error to Its Source

The debugging saga began when the assistant launched the server with piecewise CUDA graphs enabled (see [msg 975]). After the model loaded and graph capture began, the server crashed. The error trace (visible in [msg 979]) pointed to flashinfer/fp4_quantization.py, specifically the gen_fp4_quantization_module function, which calls is_cuda_version_at_least("12.8"), which in turn calls get_cuda_version()—a function that uses subprocess.check_output(["nvcc", "--version"]). Torch dynamo, which powers torch.compile, cannot trace subprocess calls because they involve thread locks and system calls outside the Python interpreter's control.

The assistant's first fix was to patch get_cuda_version to use torch.version.cuda directly instead of shelling out to nvcc (see [msg 992]). This seemed clean: the function was decorated with @functools.cache, so it would only call subprocess once, and after the patch it would never call subprocess at all.

But the server crashed again. This time, the error was at os.stat(self) in pathlib.py—dynamo couldn't trace file system operations either. The FP4 quantization module's JIT loading code (build_and_load()) checks whether compiled artifacts already exist by calling path.exists(), which performs a stat system call. Dynamo traced into this code path as well.

The Key Insight in Message 997

This is where message [msg 997] enters. The assistant has just verified (in [msg 996]) that FlashInfer's FP4 quantization functions are registered as @register_custom_op—there are eight separate custom op registrations in the file. Custom ops are supposed to be opaque to dynamo; when dynamo encounters a custom op call, it should treat it as a black box and not trace into its implementation. So why was dynamo still failing?

The assistant's reasoning, captured in the message, is precise:

"The FP4 ops ARE registered as custom ops — but the JIT loading code runs before the custom op is invoked. The get_fp4_quantization_module function itself does file I/O to load the JIT module. The problem is that dynamo traces through get_fp4_quantization_module which does build_and_load() which checks is_aot which calls path.exists()."

This is the critical diagnostic leap. The assistant has traced the failure to its root cause: the custom op boundary exists after the JIT loading code. Dynamo traces the fp4_quantize function, which calls get_fp4_quantization_module(...) to load (or retrieve from cache) the compiled CUDA kernel module, and then calls the custom op on the loaded module. But dynamo never reaches the custom op because it crashes during the loading step.

The proposed fix follows logically: "The real fix: we need to ensure the FP4 module is loaded BEFORE torch.compile tracing. Let me pre-load it."

The assistant then runs a bash command to read lines 140–170 of fp4_quantization.py, examining the structure of get_fp4_quantization_module to understand how to pre-load it. The output shows the backend module mapping—a dictionary mapping architecture strings like "120", "121", "110" to generator functions—and confirms that the module is loaded via backend_modules[backend]().build_and_load().

Assumptions and Their Limits

The assistant's reasoning in this message rests on an implicit assumption: that pre-loading the FP4 module before graph capture begins will prevent dynamo from tracing through the loading code. Since get_fp4_quantization_module is decorated with @functools.cache, calling it once during server initialization would populate the cache, and subsequent calls during graph capture would return the cached result without re-executing the loading code.

However, this assumption turns out to be incorrect, as the assistant discovers in the very next message ([msg 998]). The problem is that torch dynamo does not respect @functools.cache—it explicitly warns about this ("Dynamo detected a call to a functools.lru_cache-wrapped function. Dynamo ignores the cache wrapper and directly traces the wrapped function."). Even if the module is pre-loaded, dynamo will trace through get_fp4_quantization_module as if the cache didn't exist, re-executing the file I/O and crashing again.

This is a subtle but important point about dynamo's tracing model. Dynamo captures the Python bytecode of a function and symbolically executes it to build a computational graph. It does not actually run the function—it traces it. Caching decorators like functools.cache and functools.lru_cache are Python-level optimizations that dynamo deliberately bypasses because it needs to see the actual computation to build the graph. The warning message even notes that "Silent incorrectness is only a potential risk," meaning dynamo's developers are aware that ignoring cache wrappers could theoretically produce wrong results, but they consider it acceptable for the tracing use case.

The Deeper Problem: Dynamo's Tracing Model vs. JIT Compilation

The failure in message [msg 997] reveals a fundamental architectural tension. FlashInfer's FP4 quantization uses a two-phase design:

JIT loading phase: get_fp4_quantization_module compiles (or loads pre-compiled) CUDA kernels for the specific GPU architecture. This phase involves file I/O (checking for existing compiled artifacts), subprocess calls (invoking nvcc), and dynamic library loading. It is inherently impure—it has side effects on the file system and depends on external state.
Execution phase: The loaded module exposes functions decorated with @register_custom_op, which are pure GPU operations that can be safely traced by dynamo. The boundary between these phases is the custom op registration. But dynamo traces through the entire call stack of any function it encounters, including the JIT loading code that should have been a one-time initialization step. The @functools.cache decorator is FlashInfer's attempt to ensure the loading phase runs only once, but dynamo's tracing model deliberately bypasses such caching. This is not a bug in FlashInfer or PyTorch per se—it is a design mismatch. FlashInfer's JIT system was designed for eager-mode execution where @functools.cache works correctly. The piecewise CUDA graph runner in SGLang introduces torch.compile tracing into a code path that was never designed to be traced.

Input Knowledge Required

To fully understand message [msg 997], one needs knowledge of several distinct systems:

PyTorch Dynamo: The tracing engine underlying torch.compile. Dynamo symbolically executes Python bytecode to build a computational graph, which is then sent to a backend compiler (inductor, eager, etc.). It cannot trace operations involving system calls, thread synchronization, or file I/O.
Piecewise CUDA Graphs: SGLang's optimization that uses torch.compile to capture segments of the model's forward pass into reusable CUDA graphs. The graph runner calls torch.compile on sub-functions of the model, which triggers dynamo tracing.
FlashInfer FP4 Quantization: A JIT-compiled CUDA kernel library for 4-bit floating-point quantization. It uses @functools.cache to ensure the compilation step runs only once, and registers the compiled kernels as PyTorch custom ops for use in eager mode.
Custom Ops in PyTorch: Functions registered via torch.library or @register_custom_op that are opaque to dynamo—dynamo treats them as atomic operations and does not trace into their implementations.
SM120 Architecture: NVIDIA Blackwell's streaming multiprocessor design, with specific constraints (99KB shared memory, no TMEM) that affect kernel design choices.

Output Knowledge Created

Message [msg 997] creates several pieces of knowledge:

Root cause identification: The FP4 JIT loading code, not the custom op execution, is what breaks dynamo tracing. This narrows the debugging focus from "why does the custom op fail" to "how do we prevent dynamo from tracing the loading code."
Module structure documentation: The assistant extracts and displays the backend module mapping in get_fp4_quantization_module, showing the architecture-specific generator functions (SM121, SM120, SM110, etc.) and confirming the build_and_load() call pattern.
A proposed fix direction: Pre-loading the module before graph capture. While this specific approach ultimately fails (as the next message reveals), it represents a reasonable hypothesis based on the available information.
Evidence of systematic debugging: The assistant's method—read the error, identify the problematic code path, trace it to its source, propose a fix, test it—is visible in the message's structure. The reasoning moves from observation ("FP4 ops are custom ops") to implication ("but JIT loading runs before custom op") to action ("pre-load the module").

The Thinking Process in the Reasoning

The assistant's reasoning in message [msg 997] follows a clear diagnostic pattern:

Step 1: Contradiction detection. The assistant has just verified that FP4 ops are registered as custom ops. This creates a puzzle: if the ops are custom (and thus opaque to dynamo), why is dynamo still crashing? The resolution is that the crash occurs before the custom op is reached.

Step 2: Code path tracing. The assistant traces the execution path: fp4_quantize → get_fp4_quantization_module → build_and_load() → is_aot → path.exists(). Each step in this chain is a potential dynamo failure point. The assistant has already fixed the subprocess issue (in get_cuda_version), but the file I/O in path.exists() remains.

Step 3: Hypothesis formation. Given that get_fp4_quantization_module is cached with @functools.cache, pre-loading it should make subsequent calls return the cached result without executing the loading code. This is a reasonable hypothesis that unfortunately fails due to dynamo's cache-ignoring behavior.

Step 4: Information gathering. The assistant reads the source code of get_fp4_quantization_module to understand its structure and identify the correct pre-loading mechanism. The output confirms the backend mapping and the build_and_load() call.

Conclusion

Message [msg 997] is a snapshot of a debugging process at its most productive moment: the assistant has correctly identified the root cause of a complex, multi-layered failure and formulated a targeted fix. Even though the specific fix (pre-loading) ultimately proves insufficient, the diagnostic work is sound and the reasoning is clear. The message reveals the deep technical interplay between PyTorch's tracing infrastructure and JIT-compiled kernel libraries—a tension that will likely become more common as torch.compile adoption grows and more models combine custom kernels with graph compilation.

The episode also illustrates a broader lesson about ML systems engineering: when a system composed of multiple independently-designed components (SGLang, FlashInfer, PyTorch Dynamo, CUDA) fails, the failure often occurs at the boundaries between components, not within any single one. The custom op boundary was supposed to be a clean interface between dynamo-traceable code and opaque GPU kernels, but the JIT loading code leaked across that boundary. Fixing such issues requires understanding not just each component's design, but how they interact under conditions (like dynamo tracing) that the original designers may not have anticipated.