chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

[diffusion] fix: fix torch.compile graph break caused by torch._dynamo.disable (#18336)

Browse Source
This commit is contained in:
陈一涵
2026-02-06 14:48:09 +08:00
committed by GitHub
parent d0c39bc219
commit f798ab9775
3 changed files with 38 additions and 23 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
python/sglang/jit_kernel/tests/test_fused_norm_scale_shift.py
View File

@@ -115,8 +115,6 @@ def fused_scale_residual_norm_scale_shift_ref(
def _make_tensor(index_mode: str, shape: Tuple, dtype: torch.dtype):
if index_mode == "int1":
return 1
if index_mode == "NAT":
return None
return torch.randn(*SHAPE_MAP[index_mode](*shape), device=DEVICE, dtype=dtype)
@@ -229,7 +227,7 @@ class TestFusedScaleResidualNormScaleShift:
scale_mode=index_mode, shift_mode=index_mode, norm_type=norm_type
)
@pytest.mark.parametrize("index_mode", INDEX_MODES + ["int1"])
@pytest.mark.parametrize("index_mode", INDEX_MODES)
def test_gate_index_mode(self, index_mode, norm_type):
run_scale_resi_norm_scale_shift(gate_mode=index_mode, norm_type=norm_type)