Co-authored-by: luoyuan.luo <luoyuan.luo@antgroup.com>
This commit is contained in:
@@ -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.
|
||||
|
||||
Reference in New Issue
Block a user