fix: fix vlm cuda graph shape stability (#30868)

This commit is contained in:
Mick
2026-07-19 22:35:51 +08:00
committed by GitHub
parent a03ca46a28
commit d4801be447
11 changed files with 327 additions and 27 deletions
+6
View File
@@ -11,6 +11,7 @@ from sglang.srt.layers.rotary_embedding.rope_variant import (
DeepseekScalingRotaryEmbedding,
apply_rotary_pos_emb_native,
)
from sglang.srt.layers.rotary_embedding.utils import apply_rotary_pos_emb_native_eager
from sglang.srt.server_args import ServerArgs, set_global_server_args_for_scheduler
from sglang.test.ci.ci_register import register_cpu_ci
from sglang.test.test_utils import CustomTestCase
@@ -275,9 +276,14 @@ class TestROPE(CustomTestCase):
cos = torch.rand(num_tokens, head_size).to(sincos_dtype)
sin = torch.rand(num_tokens, head_size).to(sincos_dtype)
q_out_ref, k_out_ref = apply_rotary_pos_emb_native(query, key, cos, sin)
q_out_eager, k_out_eager = apply_rotary_pos_emb_native_eager(
query, key, cos, sin
)
q_out_sgl, k_out_sgl = torch.ops.sgl_kernel.apply_rotary_pos_emb_cpu(
query, key, cos, sin
)
torch.testing.assert_close(q_out_ref, q_out_eager)
torch.testing.assert_close(k_out_ref, k_out_eager)
torch.testing.assert_close(q_out_ref, q_out_sgl, atol=1e-2, rtol=1e-2)
torch.testing.assert_close(k_out_ref, k_out_sgl, atol=1e-2, rtol=1e-2)