pytorch / pytorch/pytorch

[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)

Open
#189,940 0 comments 0 reactions 0 assignees View on GitHub
bot-triaged enhancement module: dynamic shapes module: inductor module: performance oncall: pt2 triaged
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

Open the contributing 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.