Fix RotaryEmbedding when using Triton backend for EXAONE-3.5-2.4B

[lkm2835]

Motivation

https://huggingface.co/LGAI-EXAONE/EXAONE-3.5-2.4B-Instruct
EXAONE-3.5-2.4B-Instruct used head_dim 80.

--attention-backend flashinfer does not support head dim 80.
Then, I use --attention-backend triton to avoid flashinfer.

python3 -m sglang.launch_server --model-path /model/EXAONE-3.5-2.4B-Instruct --context-length 32768 --host 0.0.0.0 --port 9000 --tp 1 --disable-radix-cache --trust-remote-code --attention-backend flashinfer

Since flashinfer is still being used, the unsupported head_dim: 80 error still occurs when doing --attention-backend triton.

Error log

  File "/sgl-workspace/sglang/python/sglang/srt/model_executor/cuda_graph_runner.py", line 245, in __init__
    self.capture()
  File "/sgl-workspace/sglang/python/sglang/srt/model_executor/cuda_graph_runner.py", line 312, in capture
    ) = self.capture_one_batch_size(bs, forward)
  File "/sgl-workspace/sglang/python/sglang/srt/model_executor/cuda_graph_runner.py", line 391, in capture_one_batch_size
    run_once()
  File "/sgl-workspace/sglang/python/sglang/srt/model_executor/cuda_graph_runner.py", line 384, in run_once
    logits_output = forward(input_ids, forward_batch.positions, forward_batch)
  File "/usr/local/lib/python3.10/dist-packages/torch/utils/_contextlib.py", line 116, in decorate_context
    return func(*args, **kwargs)
  File "/sgl-workspace/sglang/python/sglang/srt/models/exaone.py", line 312, in forward
    hidden_states = self.transformer(
  File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1736, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1747, in _call_impl
    return forward_call(*args, **kwargs)
  File "/sgl-workspace/sglang/python/sglang/srt/models/exaone.py", line 281, in forward
    hidden_states, residual = layer(
  File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1736, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1747, in _call_impl
    return forward_call(*args, **kwargs)
  File "/sgl-workspace/sglang/python/sglang/srt/models/exaone.py", line 230, in forward
    hidden_states = self.self_attn(
  File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1736, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1747, in _call_impl
    return forward_call(*args, **kwargs)
  File "/sgl-workspace/sglang/python/sglang/srt/models/exaone.py", line 167, in forward
    q, k = self.rotary_emb(positions, q, k)
  File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1736, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1747, in _call_impl
    return forward_call(*args, **kwargs)
  File "/sgl-workspace/sglang/python/sglang/srt/custom_op.py", line 14, in forward
    return self._forward_method(*args, **kwargs)
  File "/sgl-workspace/sglang/python/sglang/srt/layers/rotary_embedding.py", line 152, in forward_cuda
    apply_rope_with_cos_sin_cache_inplace(
  File "/usr/local/lib/python3.10/dist-packages/sgl_kernel/ops/__init__.py", line 55, in apply_rope_with_cos_sin_cache_inplace
    torch.ops.sgl_kernels.apply_rope_pos_ids_cos_sin_cache(
  File "/usr/local/lib/python3.10/dist-packages/torch/_ops.py", line 1116, in __call__
    return self._op(*args, **(kwargs or {}))
  File "/usr/local/lib/python3.10/dist-packages/torch/utils/_device.py", line 106, in __torch_function__
    return func(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/torch/_ops.py", line 1116, in __call__
    return self._op(*args, **(kwargs or {}))
RuntimeError: Error in function 'BatchQKApplyRotaryPosIdsCosSinCache' at /sgl-kernel/3rdparty/flashinfer/include/flashinfer/pos_enc.cuh:584: Unsupported head_dim: 80

Modifications

Checklist

• Format your code according to the Code Formatting with Pre-Commit.
• Add unit tests as outlined in the Running Unit Tests.
• Update documentation / docstrings / example tutorials as needed, according to Writing Documentation.
• Provide throughput / latency benchmark results and accuracy evaluation results as needed, according to Benchmark and Profiling and Accuracy Results.
• For reviewers: If you haven't made any contributions to this PR and are only assisting with merging the main branch, please remove yourself as a co-author when merging the PR.
• Please feel free to join our Slack channel at https://slack.sglang.ai to discuss your PR.

[Qubitium]

@lkm2835 Can you check if /sgl-kernel/3rdparty/flashinfer/include/flashinfer/pos_enc.cuh:584 has specific divisibility regarding the RotaryEmbedding shape? Some kernels requires shape/size to be wholly divisible by a fixed amount. I want to see if this is better than 80 fix value patch which may not cover all fail cases.

[lkm2835]

@Qubitium Thanks for replying.
I checked the FlashInfer code.
https://github.com/flashinfer-ai/flashinfer/blob/main/include/flashinfer/utils.cuh#L170-L197
Flashinfer only supports 64, 128, 256, 512.

[Qubitium]

@lkm2835 Looks good! All fail cases covered. This should help other models avoid the same pitfall too. Let's wait for repsonse from other reviewers.

[lkm2835]

Hi @merrymercy.
Can you review this PR for EXAONE (and others) support on the latest sglang?

[Qubitium]

@zhaochenyang20 Can you check this?

[zhaochenyang20]

cc @zhyncs

[lkm2835]

Hi @Ying1123 can you check this PR?
We want to use EXAONE-Deep-2.4B on the sglang

[Qubitium]

@lkm2835 I have pinged the SGLang Slack channel to accelerate this merge.