NATTEN 0.21.7 with PyTorch 2.11.0 fails to compile na2d when the neighborhood covers the full input.
import torch
import natten
def f(q, k, v):
return natten.na2d(
q,
k,
v,
kernel_size=(8, 8),
backend="cutlass-fna",
)
shape = (1, 8, 8, 1, 64)
q, k, v = [
torch.randn(shape, device="cuda", dtype=torch.bfloat16)
for _ in range(3)
]
compiled = torch.compile(f, fullgraph=True, dynamic=True)
compiled(q, k, v)This raises:
torch._dynamo.exc.Unsupported: Failed to trace builtin operator `repr`
Developer debug context: builtin repr [SizeVariable]
The failure comes from the debug message in src/natten/functional.py:410-413. NattenLogger.debug() correctly skips logging while Dynamo is tracing, but the f-string is evaluated before debug() is called. Formatting the symbolic query.shape invokes repr(torch.Size), which Dynamo cannot trace.
Suggested fix: use lazy logging arguments, or guard message formatting before the call. For example:
logger.debug(
"query.shape=%s with kernel_size=%s, "
"has_additional_attention=%s and is_causal=%s is self attention. "
"Calling attention instead of neighborhood attention directly.",
query.shape,
kernel_size,
has_additional_attention,
is_causal,
)A regression test should compile the input_size == kernel_size self-attention fallback with fullgraph=True, dynamic=True.
NATTEN 0.21.7 with PyTorch 2.11.0 fails to compile
na2dwhen the neighborhood covers the full input.This raises:
The failure comes from the debug message in src/natten/functional.py:410-413.
NattenLogger.debug()correctly skips logging while Dynamo is tracing, but the f-string is evaluated beforedebug()is called. Formatting the symbolic query.shape invokesrepr(torch.Size), which Dynamo cannot trace.Suggested fix: use lazy logging arguments, or guard message formatting before the call. For example:
A regression test should compile the input_size == kernel_size self-attention fallback with fullgraph=True, dynamic=True.