From bc8838e48d26832d7ee664f3cd41823fafb6d12d Mon Sep 17 00:00:00 2001 From: Vaibhav Jindal Date: Fri, 7 Aug 2026 08:45:43 +0000 Subject: [PATCH] fix(rope): don't mark the unused batch-size param as tl.constexpr `_triton_rope` declares `bs: tl.constexpr`, but `bs` is never referenced in the kernel body -- it appears only in comments. Marking a dead parameter as constexpr still forces Triton/Dynamo to specialize on it, and that breaks torch.compile as soon as the batch dimension becomes dynamic: the value arrives as a SymInt, and Dynamo trips an internal assertion while tracing LigerRopeFunction through the autograd_function_apply HOP: assert subgraph_vt.is_tensor() or isinstance(subgraph_vt, SymNodeVariable) torch/_dynamo/variables/higher_order_ops.py:325 The failure only surfaced on the *second* distinct input shape, since the first call compiles statically and the second triggers automatic dynamic shapes -- which made it look like an upstream torch bug rather than a kernel signature issue. Varying seq_len alone was always fine, because `sl` is an ordinary runtime arg; only the constexpr batch size was fatal: before after vary seq_len only (auto-dynamic) OK OK vary batch only (auto-dynamic) FAIL OK vary batch, dynamic=False OK OK Dropping the constexpr annotation is sufficient. Qwen3-0.6B now compiles under a plain `torch.compile(model)` across batch sizes 1/4/8 and sequence lengths 512/1024/2048; previously that required a `dynamic=False` workaround. Explicit `torch._dynamo.mark_dynamic` on the batch and sequence dims also works now. The parameter is kept (rather than removed) so the kernel's call signature and the `q size: (bsz, ...)` documentation stay intact. Note `dynamic=True` remains unsupported: it additionally marks head_dim and the head counts dynamic, which cannot work while block sizing is derived from `triton.next_power_of_2(head_dim)` at launch time. Those dims are fixed by model architecture in practice, so this is not the case that matters. Behaviour is unchanged on the eager path, and compiled output and gradients are bit-identical to eager (max abs diff 0.0), as expected for a parameter the kernel never reads. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- src/liger_kernel/ops/rope.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/src/liger_kernel/ops/rope.py b/src/liger_kernel/ops/rope.py index bd8ded730..fb8e067dd 100644 --- a/src/liger_kernel/ops/rope.py +++ b/src/liger_kernel/ops/rope.py @@ -14,7 +14,13 @@ def _triton_rope( sin, sin_row_stride, sl, - bs: tl.constexpr, + # `bs` is unused by the kernel body and is deliberately NOT a tl.constexpr. + # Marking it constexpr forces Triton/Dynamo to specialize on the batch size, + # which fails when torch.compile makes the batch dimension dynamic: the value + # arrives as a SymInt and Dynamo trips an internal assertion while tracing + # LigerRopeFunction ("assert subgraph_vt.is_tensor() or isinstance( + # subgraph_vt, SymNodeVariable)" in _dynamo/variables/higher_order_ops.py). + bs, cos_bs: tl.constexpr, n_qh: tl.constexpr, n_kh: tl.constexpr,