[CI] Fix lint brought by #27527 (#28988)

Co-authored-by: luoyuan.luo <luoyuan.luo@antgroup.com>
This commit is contained in:
Yuan Luo
2026-06-22 20:40:06 -07:00
committed by GitHub
co-authored by luoyuan.luo
parent 6cd8d2869b
commit abb0717174
2 changed files with 68 additions and 49 deletions
+33 -14
View File
@@ -43,6 +43,7 @@ _DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
# Standalone reference implementation (original loop-based code, pre-a475156d)
# ---------------------------------------------------------------------------
def _create_custom_4d_mask_reference(
sequence_length, dtype, device, batch_size, token_type_ids
):
@@ -79,6 +80,7 @@ def _create_custom_4d_mask_reference(
# python/sglang/srt/models/deepseek_ocr.py)
# ---------------------------------------------------------------------------
def _create_custom_4d_mask_new(
sequence_length, dtype, device, batch_size, token_type_ids
):
@@ -117,6 +119,7 @@ def _create_custom_4d_mask_new(
# Helpers
# ---------------------------------------------------------------------------
def _make_token_type_ids(batch_size, seq_len, image_fraction, device):
"""First `image_fraction` tokens per sequence are image (0), rest are text (1).
@@ -172,6 +175,7 @@ def _bench(fn, run_device, n=50, **kwargs):
# Accuracy tests
# ---------------------------------------------------------------------------
class TestAccuracy(unittest.TestCase):
"""Verify new implementation produces identical masks to the reference."""
@@ -280,11 +284,17 @@ class TestAccuracy(unittest.TestCase):
for q in range(S):
for k in range(S):
if k <= q:
self.assertEqual(mask2d[q, k].item(), 0.0,
f"text[{q}] should attend to text[{k}]")
self.assertEqual(
mask2d[q, k].item(),
0.0,
f"text[{q}] should attend to text[{k}]",
)
else:
self.assertEqual(mask2d[q, k].item(), min_val,
f"text[{q}] should NOT attend to text[{k}]")
self.assertEqual(
mask2d[q, k].item(),
min_val,
f"text[{q}] should NOT attend to text[{k}]",
)
def test_image_full_attention(self):
"""Image tokens must attend to all other image tokens (bidirectional)."""
@@ -296,8 +306,9 @@ class TestAccuracy(unittest.TestCase):
mask2d = out.cpu()[0, 0]
for q in range(n_img):
for k in range(n_img):
self.assertEqual(mask2d[q, k].item(), 0.0,
f"image[{q}] should attend to image[{k}]")
self.assertEqual(
mask2d[q, k].item(), 0.0, f"image[{q}] should attend to image[{k}]"
)
def test_text_attends_to_image(self):
"""Every text token must attend to every image token."""
@@ -309,8 +320,9 @@ class TestAccuracy(unittest.TestCase):
mask2d = out.cpu()[0, 0]
for q in range(n_img, S):
for k in range(n_img):
self.assertEqual(mask2d[q, k].item(), 0.0,
f"text[{q}] should attend to image[{k}]")
self.assertEqual(
mask2d[q, k].item(), 0.0, f"text[{q}] should attend to image[{k}]"
)
# --- dtype coverage ---
@@ -352,8 +364,9 @@ class TestPerformance(unittest.TestCase):
cls.dtype = torch.float32
def _run_case(self, batch_size, seq_len, image_fraction):
ids = _make_token_type_ids(batch_size, seq_len, image_fraction,
device=self.device)
ids = _make_token_type_ids(
batch_size, seq_len, image_fraction, device=self.device
)
kwargs = dict(
sequence_length=seq_len,
dtype=self.dtype,
@@ -361,10 +374,15 @@ class TestPerformance(unittest.TestCase):
batch_size=batch_size,
token_type_ids=ids,
)
t_ref = _bench(_create_custom_4d_mask_reference, run_device=self.device,
n=BENCH_ITERS, **kwargs)
t_new = _bench(_create_custom_4d_mask_new, run_device=self.device,
n=BENCH_ITERS, **kwargs)
t_ref = _bench(
_create_custom_4d_mask_reference,
run_device=self.device,
n=BENCH_ITERS,
**kwargs,
)
t_new = _bench(
_create_custom_4d_mask_new, run_device=self.device, n=BENCH_ITERS, **kwargs
)
speedup = t_ref / t_new
dev_tag = "CUDA" if "cuda" in str(self.device) else "CPU"
print(
@@ -406,6 +424,7 @@ class TestPerformance(unittest.TestCase):
# PyTorch profiler (optional – triggered by --profile or PROFILE_TRACES=1)
# ---------------------------------------------------------------------------
def run_profiler_traces(output_dir: str = "./pt_traces", device: str = _DEVICE):
"""
Capture Chrome-trace JSON files for both implementations.