[inductor] torch.compile of torch.cat with a symbolic (dynamic) concat dim is 2–2.5× slower than eager (symbolic int64 divmod in generated indexing)
- Dominant language
- Python
- Stars
- 103k
- Forks
- 29.6k
- PR merge metrics
- PR metrics pending
Description
## Describe the bug
When the concatenation dimension of `torch.cat` is a runtime symint (i.e. a non-leading dim is marked dynamic), Inductor lowers the cat to a fused `pointwise_cat` kernel whose per-element indexing contains modulo/division by a runtime symint (`xindex % ks0`, `xindex // ks0`). GPUs can't strength-reduce division by a runtime value, so this int64 divmod dominates what is otherwise a bandwidth-bound copy, making the compiled kernel 2–2.5× slower than eager torch.cat. With static shapes the same lowering folds the divmod to a compile-time constant (magic-number multiply/shift) and is instead faster than eager.
The workload below mirrors a QKV-repack in LLM backward: a per-group value grad accumulation (add) that Inductor fuses into the cat, and a nested same-dim cat.
## Minimal repro
```
import torch
def nested_cat_add(q1, k1, v1a, v1b, q2, k2, v2a, v2b):
v1 = v1a + v1b
v2 = v2a + v2b
g1 = torch.cat([q1, k1, v1], dim=-1)
g2 = torch.cat([q2, k2, v2], dim=-1)
return torch.cat([g1, g2], dim=-1)
def make_inputs(n):
widths = [2048, 256, 256, 256, 2048, 256, 256, 256] # q/k/v head widths
return [torch.randn(n, w, dtype=torch.bfloat16, device="cuda") for w in widths]
def bench_ms(fn, inputs, warmup=25, iters=100):
for _ in range(warmup):
fn(*inputs)
torch.cuda.synchronize()
s, e = (torch.cuda.Event(enable_timing=True) for _ in range(2))
s.record()
for _ in range(iters):
fn(*inputs)
e.record()
torch.cuda.synchronize()
return s.elapsed_time(e) / iters
inputs = make_inputs(4096)
eager_ms = bench_ms(nested_cat_add, inputs)
# Make the concat (last) dim a runtime symint.
for t in inputs:
torch._dynamo.mark_dynamic(t, t.dim() - 1)
compiled = torch.compile(nested_cat_add, dynamic=True)
compiled(*inputs) # triggers compilation
comp_ms = bench_ms(compiled, inputs)
print(f"eager : {eager_ms:.4f} ms/iter")
print(f"compiled : {comp_ms:.4f} ms/iter ({comp_ms / eager_ms:.2f}x eager)")
```
## Observed (NVIDIA H100, torch 2.12.0+cu126)
```
eager : 0.0815 ms/iter
compiled : 0.1653 ms/iter (2.03x eager)
```
Full matrix (eager reference; dynamic = last dim marked dynamic):
| variant | n=4096 | n=32768 |
| --- | --- | --- |
| eager `torch.cat` | 0.0815 ms | 0.5088 ms |
| compiled, dynamic (`pointwise_cat`) | 0.1653 ms (2.03× slower) | 1.2906 ms (2.54× slower) |
| compiled, static (`pointwise_cat`) | 0.0595 ms (1.37× faster) | 0.2380 ms (2.14× faster) |
## Root cause — generated Triton
Dynamic lowering carries the concat widths as runtime i64 args and reconstructs (row, col) with a symbolic divmod:
```
# dynamic: def triton_poi_fused_add_cat_0(..., ks0, ks1, ks2, ks3, ks4, ks5, ks6, xnumel, XBLOCK):
xmask = xindex < xnumel
x0 = (xindex % ks0) # <-- runtime int64 modulo by a symint
x1 = xindex // ks0 # <-- runtime int64 division by a symint
...
tmp11 = tl.load(in_ptr0 + (ks1*x1 + (x0)), tmp10 & xmask, ...) # symbolic 64-bit addr math
tmp17 = tl.load(in_ptr1 + (ks3*x1 + ((-1)*ks1 + x0)), tmp16 & xmask, ...)
```
vs. the static lowering, where the divisor is a constant and gets strength-reduced by ptxas:
```
# static:
xmask = tl.full([XBLOCK], True, tl.int1)[:]
x0 = (xindex % 5120) # constant -> magic-number multiply + shift
x1 = xindex // 5120
...
tmp11 = tl.load(in_ptr0 + (2048*x1 + (x0)), tmp10, ...) # constant strides
```
Eager avoids this because ATen’s `CatArrayBatchedCopy` precomputes fast-divmod magic constants on the host (IntDivider) and does all inputs in one launch; Inductor’s codegen emits a raw %/// by a runtime int64.
Approx. effective bandwidth (≈88 MB moved) at n=4096: static ≈1.5 TB/s (bandwidth-bound), dynamic ≈0.5 TB/s (now ALU/latency-bound on the index math).
## Alternatives considered
`ConcatKernel` (copy) path instead of `pointwise_cat` — does not help. It splits into ~8 per-input copy kernels, but each copy kernel still computes `xindex % W / xindex // W` with a symbolic `W` to write into the strided output slice: 0.1627 ms @ n=4096 (2.01× slower), essentially identical to `pointwise_cat`. This confirms the `tl.where` source-selection chain in `pointwise_cat` is not the bottleneck — the symbolic divmod (present in both) is.
## Suggested fix
Fall back to the ATen cat kernel when a non-leading dimension is symbolic (that’s exactly when the generated indexing needs a symbolic divmod; leading-dim-only dynamism never divides and should stay fused).
```
# torch/_inductor/lowering.py, in cat(...), right after `dim = _validate_dim(inputs[0], dim, 0)`
import sympy
# A symbolic extent in any NON-LEADING dim forces `idx % W` / `idx // W` with a
# runtime symint W in the generated kernel; GPUs can't strength-reduce division by
# a runtime value, so it dominates this otherwise bandwidth-bound copy (2-2.5x
# slower than eager). ATen's cat precomputes fast-divmod on the host -> fall back.
if config.fallback_dynamic_cat and any(
isinstance(s, sympy.Expr) and s.free_symbols
for inp in inputs
for s in inp.get_size()[1:]
):
return fallback_handler(aten.cat.default)(inputs, dim)
```
with a config knob (default on) in `torch/_inductor/config.py`:
```
# Fall back to the ATen cat kernel when a non-leading cat dim is symbolic, to avoid
# symbolic int64 modulo/division in the generated indexing.
fallback_dynamic_cat: bool = True
```
Notes on the condition:
* Use `sympy.Expr.free_symbols`, not `free_unbacked_symbols` — `mark_dynamic` produces backed symints.
* `get_size()[1:]` is the precise predictor: last-dim/middle-dim dynamic cat → fall back; dim=0 with static inner dims (e.g. [s0, 256]) or 1-D cat → stays fused (constant/absent divisor, already fast).
cc @chauhang @penguinwu @ezyang @bobrenjc93 @aditvenk @laithsakka @voznesenskym @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @ipiszy @kadeng @muchulee8 @amjames @aakhundov @coconutruben @jataylo
Contributor guide
Research direction
Start in torch/_inductor/lowering.py at cat(...) and review the related option in torch/_inductor/config.py. Reproduce the dynamic-versus-static benchmark from the issue, then inspect how symbolic non-leading dimensions affect cat lowering. Done means dynamic non-leading dimensions use the ATen cat path while leading-dimension-only dynamism remains fused, with the configuration behavior covered by appropriate tests.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch
- Domain
- compilers, performance
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Quiet
- Clarity
- Clearly specified
- Newbie friendliness
- 68/100