fix(rope): don't mark the unused batch-size param as tl.constexpr

[vaibhavjindal]

Summary

_triton_rope declares bs: tl.constexpr, but bs is never referenced in the kernel body — it appears only in comments. An AST check of the kernel's parameters:

param      constexpr  used-in-body
  sl         -          yes
  bs         yes        NO   <-- dead
  cos_bs     yes        yes
  n_qh       yes        yes
  ...

Marking a dead parameter constexpr still forces Triton/Dynamo to specialize on it. 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:

File "torch/_dynamo/variables/higher_order_ops.py", line 325, in overwrite_tensor_vt_proxy
    assert subgraph_vt.is_tensor() or isinstance(subgraph_vt, SymNodeVariable)
AssertionError:

from user code:
  File "transformers/models/qwen3/modeling_qwen3.py", line 268, in forward
    query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
  File "liger_kernel/ops/rope.py", line 24, in liger_rotary_pos_emb
    return LigerRopeFunction.apply(q, k, cos, sin, position_ids, unsqueeze_dim)

The failure only surfaces on the second distinct input shape — the first call compiles statically, the second triggers automatic dynamic shapes. That timing made it look like an upstream PyTorch bug rather than a kernel signature issue.

Dropping the tl.constexpr annotation is sufficient. The parameter itself is kept so the call signature and the q size: (bsz, ...) documentation stay intact.

Details

Minimal repro (no transformers dependency), against main:

import torch, torch._dynamo as dyn
from liger_kernel.ops.rope import LigerRopeFunction

cf = torch.compile(lambda *a: LigerRopeFunction.apply(*a))   # default: auto-dynamic
def mk(bs, sl, nqh=16, nkv=8, hd=128):
    return (torch.randn(bs, nqh, sl, hd, device="cuda", dtype=torch.bfloat16, requires_grad=True),
            torch.randn(bs, nkv, sl, hd, device="cuda", dtype=torch.bfloat16, requires_grad=True),
            torch.randn(1, sl, hd, device="cuda", dtype=torch.bfloat16),
            torch.randn(1, sl, hd, device="cuda", dtype=torch.bfloat16))

cf(*mk(1, 64))   # OK
cf(*mk(2, 64))   # AssertionError on main

Varying seq_len alone was always fine, because sl is an ordinary runtime arg. Only the constexpr batch size was fatal:

scenario	before	after
vary seq_len only (auto-dynamic)	OK	OK
vary batch only (auto-dynamic)	FAIL	OK
vary batch, dynamic=False	OK	OK
vary seq_len, dynamic=False	OK	OK
dynamic=True	FAIL	FAIL (see below)

End to end, Qwen3-0.6B with apply_liger_kernel_to_qwen3(rope=True, ...) now compiles under a plain torch.compile(model) across all of 1x512, 1x2048, 4x512, 4x2048, 8x1024, 8x2048. Previously that required a dynamic=False workaround. Explicit torch._dynamo.mark_dynamic on the batch and sequence dims also works now (verified at 1x64, 2x128, 8x256).

dynamic=True remains unsupported, and I don't think it should block this. That flag 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 it isn't the case that matters — whereas batch and sequence length genuinely vary, and both work now.

Two related observations, deliberately not changed here:

1. BLOCK_SIZE is also unused in _triton_rope's body. It's harmless today (it's derived from head counts, so it never becomes a SymInt), so I left it rather than widen the diff.
2. qwen2vl_mrope.py has the same bs: tl.constexpr annotation, but there bs is genuinely used for pointer arithmetic, so it needs its own treatment rather than the same one-line change. The Ascend/NPU backends carry the same pattern and I have no way to test them here.

Testing Done

• pytest test/transformers/test_rope.py -> 36 passed

• pytest test/transformers/test_rope.py test/transformers/test_monkey_patch.py -> 98 passed

• pytest test/convergence/bf16/test_mini_models.py -k "qwen3 or llama or mistral" -> 11 passed

• make checkstyle -> passed

• Compiled vs eager output and gradients are bit-identical (max abs diff 0.0) across bs=1,2,4,8 — expected, since the kernel never reads this parameter.

• Hardware Type: H100 80GB HBM3

• run make test to ensure correctness

• run make checkstyle to ensure code style

• run make test-convergence to ensure convergence

Note: I ran the targeted suites listed above rather than the full make test /
make test-convergence sweeps, so I've left those two boxes unchecked rather than
claim more coverage than I actually have. Relying on CI for the full run.