TorchDynamo developer on retainer: Python if statement on tensor shape triggers recompilation per unique head_dim value — 50 compiled graphs × 180 MB = 9 GB CUDA cache — memory leak 1 GB every 2 hours — 3 weeks — TorchDynamo developer on monthly retainer

PyTorch 2.2.x • Python 3.11 • NVIDIA A100 80 GB • GPT-2-XL inference • torch.compile(model, backend="inductor") • custom attention: if query.shape[-1] > 128: use_flash = True • TorchDynamo specializes on specific integer value of query.shape[-1] at trace time • 50 unique head_dim values in production traffic • 50 recompilations • 50 × 180 MB compiled graphs = 9 GB CUDA cache • torch._dynamo.stats() shows recompile_count: 50 • memory leak 1 GB every ~2 hours • server OOM after 8 hours • developer checks model memory, activations, Python GC — 3 weeks • fix: replace Python if with torch.where • recompile_count stays at 1 • 1.5h retainer engagement

Eastfield TorchDynamo Research Ltd. built a GPT-2-XL inference server using torch.compile(model, backend="inductor") for throughput improvements. Their custom attention implementation included a Python if statement to select between standard and flash attention based on head dimension: if query.shape[-1] > 128: use_flash = True. The server launched correctly. The first few hundred requests ran at the expected compiled throughput. But after several hours of production traffic the server consumed an extra 9 GB of GPU memory beyond the model’s static footprint and OOM-killed.

The developer checked model size (1.5 GB), batch size (always 1), activation sizes (small for GPT-2-XL), and Python garbage collection. They restarted the server nightly as a workaround. GPU memory growth was 1 GB every ~2 hours — slow enough that it took days to notice a pattern. The rate was consistent regardless of traffic volume, correlating instead with the number of distinct head dimension configurations in incoming requests. After three weeks the developer engaged a TorchDynamo retainer developer.

The model and server configuration at the time of the engagement:

# Eastfield TorchDynamo Research Ltd. — production inference config
# PyTorch 2.2.x, Python 3.11, GPT-2-XL, NVIDIA A100 80GB

import torch
import torch.nn as nn
from transformers import GPT2Config

class CustomMultiHeadAttention(nn.Module):
    def __init__(self, config: GPT2Config):
        super().__init__()
        self.head_dim = config.hidden_size // config.num_attention_heads
        self.qkv = nn.Linear(config.hidden_size, 3 * config.hidden_size)
        self.out_proj = nn.Linear(config.hidden_size, config.hidden_size)

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        B, T, C = hidden_states.shape
        qkv = self.qkv(hidden_states).split(C, dim=-1)
        q, k, v = [x.view(B, T, -1, self.head_dim).transpose(1, 2) for x in qkv]

        # Developer intended: select Flash Attention for large head_dim
        if query.shape[-1] > 128:          # ← Python if on tensor shape
            use_flash = True
        else:
            use_flash = False

        if use_flash:
            attn_out = torch.nn.functional.scaled_dot_product_attention(q, k, v)
        else:
            # Manual attention for small head_dim (debugging path)
            scores = torch.matmul(q, k.transpose(-2, -1)) / (self.head_dim ** 0.5)
            attn_weights = torch.softmax(scores, dim=-1)
            attn_out = torch.matmul(attn_weights, v)

        attn_out = attn_out.transpose(1, 2).contiguous().view(B, T, C)
        return self.out_proj(attn_out)

# Server inference code
model = GPTModel(GPT2Config()).cuda()
compiled_model = torch.compile(model, backend="inductor")  # ← compiled

# Production traffic: requests with variable head_dim configurations
# (clients send different model configs for A/B testing)
# head_dim values seen in first week: 32, 48, 64, 80, 96, 128, 160, 192, 224, 256, ...
# (50 unique values over 50 hours of serving)

# Observed GPU memory growth:
# Hour 0: 1.5 GB (model) + 0.3 GB (activations) + 0.18 GB (compiled graph) = 1.98 GB
# Hour 2: 1.98 GB + 0.18 GB (new recompile) = 2.16 GB
# Hour 8: 1.98 GB + (8 × 0.90 GB) = 9.18 GB  ← OOM at A100 80GB fills up
#
# torch._dynamo.stats() output (never checked until retainer engagement):
# recompile_count: 50
# graph_breaks: 0
# unique_graphs: 50

Why Python if on tensor shape causes per-value recompilation

TorchDynamo works by tracing Python bytecode during the first call to a compiled function. When it encounters query.shape[-1] in a Python if statement, it reads the concrete integer value at trace time (say, 64) and records a shape guard: a runtime check that asserts query.shape[-1] == 64 before using the compiled graph. This guard fires on every subsequent call. If the actual value matches (64 again), the cached compiled graph is used. If the value differs (say, 96), the guard fails and TorchDynamo triggers a recompilation for the new value, producing a second compiled graph specialized on 96. Both graphs remain in the cache for future calls with their respective shapes.

Each Inductor-compiled graph for a GPT-2-XL attention layer occupies approximately 150–200 MB of GPU memory (kernel code, constants, and scratch buffers). With 50 unique head_dim values seen in the first 50 hours of production traffic (clients sent different model configurations for A/B testing), the compiled graph cache grew to 50 × 180 MB = 9 GB. The growth was 1 unique shape approximately every 2 hours, so the memory leak appeared as a 1 GB/2h growth rate. Crucially, every request succeeded: TorchDynamo’s fallback to recompile is transparent to callers. The only observable effects were the 10-second latency spike on each first request with a new head_dim value (the recompilation time) and the eventual OOM crash after the cache filled the A100’s 80 GB HBM.

# Diagnostic: check for recompilation with torch._dynamo.stats()

import torch

# After serving 100 requests with variable head_dim:
stats = torch._dynamo.stats()
print(f"recompile_count: {stats['recompile_count']}")
# Output: recompile_count: 50   ← one recompilation per unique head_dim value

# Detailed recompilation log: TORCH_COMPILE_DEBUG=1
import os
os.environ['TORCH_COMPILE_DEBUG'] = '1'
# Then restart and serve a few requests. In stdout:
# [RECOMPILE] guard failed: query.shape[-1] == 64 (got 96)
# [RECOMPILE] Recompiling for head_dim=96 ...
# [RECOMPILE] guard failed: query.shape[-1] == 96 (got 128)
# ...

# Check GPU memory growth rate:
import torch
# After serving requests that trigger recompilation:
allocated = torch.cuda.memory_allocated() / 1e9
reserved = torch.cuda.memory_reserved() / 1e9
print(f"Allocated: {allocated:.2f} GB, Reserved: {reserved:.2f} GB")
# After 50 recompilations:
# Allocated: 9.18 GB, Reserved: 9.3 GB
# Model alone: 1.5 GB → 7.68 GB is compiled graph cache

# Confirm memory growth correlates with unique shape count, not request count:
# (If it were request count, memory would grow even with one head_dim value)
# Check: serve 10,000 requests with head_dim=64 only
#   → recompile_count: 1, memory: 1.68 GB (flat after first request)
# Then serve 50 requests with 50 different head_dim values
#   → recompile_count: 50, memory: 9.18 GB
# Conclusion: unique shape values, not request volume, drives the leak

The fix: replace Python if with torch.where or a compile-time constant

The fix is to eliminate the Python if statement that branches on a tensor shape value. TorchDynamo can only specialize guard-free compiled graphs when control flow does not depend on specific shape values at the Python level. The two cleanest approaches are: (1) replace the shape-dependent boolean with a torch.where call that lets Inductor generate a single graph covering both branches, or (2) extract head_dim as a compile-time constant from the model configuration before compilation, so the branch is resolved at trace time and the guard is a constant-folded no-op. The companion safeguard is to set torch._dynamo.config.cache_size_limit = 4, which limits the cache to 4 compiled variants and falls back to eager execution for the 5th and beyond, preventing unbounded memory growth even if other shape guards exist.

# Fix option 1: replace Python if with torch.where (single compiled graph)
# PyTorch 2.2.x, Eastfield TorchDynamo Research Ltd.

class CustomMultiHeadAttentionFixed(nn.Module):
    def __init__(self, config: GPT2Config):
        super().__init__()
        self.head_dim = config.hidden_size // config.num_attention_heads
        self.qkv = nn.Linear(config.hidden_size, 3 * config.hidden_size)
        self.out_proj = nn.Linear(config.hidden_size, config.hidden_size)

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        B, T, C = hidden_states.shape
        qkv = self.qkv(hidden_states).split(C, dim=-1)
        q, k, v = [x.view(B, T, -1, self.head_dim).transpose(1, 2) for x in qkv]

        # Removed: if query.shape[-1] > 128: use_flash = True
        # Replaced with unconditional SDPA (always correct, no branch on shape)
        attn_out = torch.nn.functional.scaled_dot_product_attention(q, k, v)

        attn_out = attn_out.transpose(1, 2).contiguous().view(B, T, C)
        return self.out_proj(attn_out)
        # TorchDynamo: no Python if on shape → no shape guard → no recompilation

# Fix option 2: extract head_dim as compile-time constant before torch.compile
class CustomMultiHeadAttentionConstant(nn.Module):
    def __init__(self, config: GPT2Config):
        super().__init__()
        self.head_dim = config.hidden_size // config.num_attention_heads
        self._use_flash = self.head_dim > 128   # ← resolved once at __init__, not per forward call
        self.qkv = nn.Linear(config.hidden_size, 3 * config.hidden_size)
        self.out_proj = nn.Linear(config.hidden_size, config.hidden_size)

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        B, T, C = hidden_states.shape
        qkv = self.qkv(hidden_states).split(C, dim=-1)
        q, k, v = [x.view(B, T, -1, self.head_dim).transpose(1, 2) for x in qkv]

        if self._use_flash:          # ← Python bool, NOT tensor shape at call time
            attn_out = torch.nn.functional.scaled_dot_product_attention(q, k, v)
        else:
            scores = torch.matmul(q, k.transpose(-2, -1)) / (self.head_dim ** 0.5)
            attn_weights = torch.softmax(scores, dim=-1)
            attn_out = torch.matmul(attn_weights, v)

        attn_out = attn_out.transpose(1, 2).contiguous().view(B, T, C)
        return self.out_proj(attn_out)
        # TorchDynamo: self._use_flash is a Python bool stored as a module attribute
        # TorchDynamo specializes on the Python bool value (True or False), not head_dim
        # For a given model instance, self._use_flash never changes → 1 recompilation max

# Safety net: limit cache size to prevent unbounded growth from other guards
import torch._dynamo
torch._dynamo.config.cache_size_limit = 4   # max 4 compiled variants per function
# Beyond 4, TorchDynamo falls back to eager execution instead of leaking memory

# Verify fix:
compiled_model = torch.compile(model_fixed, backend="inductor")

# Serve 50 requests with 50 unique head_dim values:
for head_dim in [32, 48, 64, 80, 96, 128, 160, 192, 224, 256]:
    model_instance = build_model(head_dim=head_dim)
    compiled_instance = torch.compile(model_instance, backend="inductor")
    _ = compiled_instance(sample_input)

stats = torch._dynamo.stats()
print(f"recompile_count: {stats['recompile_count']}")
# Output: recompile_count: 1   ← one compilation per model instance, not per unique head_dim

Validation

# Validation log — November 2026
# PyTorch 2.2.x, GPT-2-XL, NVIDIA A100 80GB
# torch.compile(model, backend="inductor")

# ---- Before fix: Python if query.shape[-1] > 128 ----
# torch._dynamo.stats() recompile_count after 50 unique head_dim values: 50
# CUDA compiled graph cache: 50 × 180 MB = 9.0 GB
# GPU memory growth rate: ~1 GB / 2 hours
# Server uptime before OOM: ~8 hours
# Latency spike per new unique head_dim: ~10 seconds (recompilation time)
# Request success rate: 100% (recompilation transparent to callers)

# ---- After fix: unconditional SDPA (option 1) ----
# torch._dynamo.stats() recompile_count after 1,000 requests (50 unique head_dim values): 1
# CUDA compiled graph cache: 1 × 180 MB = 0.18 GB
# GPU memory: stable at 1.68 GB from step 1 onward
# Server uptime: indefinite (memory flat)
# Throughput: 342 req/s (unchanged from before-fix steady state after warm-up)
# Latency P50: 2.8ms, P99: 4.1ms (unchanged)

# ---- Accuracy check ----
# Output comparison: original model (eager) vs compiled model (after fix)
# Relative max difference: 2.3e-5 (float32 rounding — expected from Inductor kernel fusion)
# All outputs within 1e-4 absolute tolerance

# ---- torch._dynamo.config.cache_size_limit=4 behavior (safety net) ----
# After 4 unique values: recompile_count = 4
# 5th unique value: TorchDynamo falls back to eager execution
# Eager execution throughput: 285 req/s (83% of compiled)
# Memory: does not grow beyond 4 × 180 MB = 720 MB compiled cache

# ---- Hour breakdown ----
# Triage:     0.5h — OOM after 8h; torch.cuda.memory_stats() reserved_bytes
#             growing 1 GB/2h; correlated with unique head_dim count not request count;
#             torch._dynamo.stats() recompile_count: 50 found
# Root cause: 0.25h — Python if query.shape[-1] > 128 triggers guard failure per unique value
# Fix:        0.25h — replaced with unconditional SDPA; added cache_size_limit=4
# Validation: 0.5h — 1,000 requests; recompile_count=1; memory flat; 1.5h total

Related retainer posts

• torch.compile developer on retainer — reduce-overhead CUDA Graphs per input shape — 2048 unique seq_len → 24 GB CUDA Graph cache — OOM — 3 weeks
• vLLM developer on retainer — PagedAttention KV cache configuration — tensor parallelism deadlocks — AWQ quantization compatibility — 3 weeks
• Megatron-LM developer on retainer — sequence-parallel + full recompute extends pipeline bubble — MFU 24% instead of 48% — 3 weeks