test: strengthen CG-replay coverage with prod-fill padding, metadata invariants, and pad-ratio sweep (#26658)

Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
Cheng Wan
2026-05-28 22:43:29 -07:00
committed by GitHub
co-authored by Claude Sonnet 4.6
parent 226649e3b7
commit 2dfbc3d781
8 changed files with 323 additions and 25 deletions
@@ -435,13 +435,26 @@ class TestFA3DenseAttentionBackendCorrectness(CustomTestCase):
"and produces ~82 % wrong attention values. Tracked in KNOWN_FAILURES.md §C.3."
)
for case in self.DRAFT_EXTEND_V2_CUDA_GRAPH_CASES:
with self.subTest(case=case.name, backend=case.backend):
run_dense_draft_extend_v2_cuda_graph_case(
self,
case,
head_dim=self.HEAD_DIM,
hidden_size=self.HIDDEN_SIZE,
)
for pad_style in ("small_real", "prod_fill"):
for capture_bs in (
case.batch_size,
case.batch_size * 2,
case.batch_size * 4,
):
with self.subTest(
case=case.name,
backend=case.backend,
pad_style=pad_style,
capture_bs=capture_bs,
):
run_dense_draft_extend_v2_cuda_graph_case(
self,
case,
head_dim=self.HEAD_DIM,
hidden_size=self.HIDDEN_SIZE,
cuda_graph_capture_batch_size=capture_bs,
pad_style=pad_style,
)
def test_runner_mode_eagle_draft_extend_cases(self):
for case, spec_kind in self.DRAFT_EXTEND_CASES:
@@ -424,13 +424,26 @@ class TestFA4DenseAttentionBackendCorrectness(CustomTestCase):
"and produces ~82 % wrong attention values. Tracked in KNOWN_FAILURES.md §C.3."
)
for case in self.DRAFT_EXTEND_V2_CUDA_GRAPH_CASES:
with self.subTest(case=case.name, backend=case.backend):
run_dense_draft_extend_v2_cuda_graph_case(
self,
case,
head_dim=self.HEAD_DIM,
hidden_size=self.HIDDEN_SIZE,
)
for pad_style in ("small_real", "prod_fill"):
for capture_bs in (
case.batch_size,
case.batch_size * 2,
case.batch_size * 4,
):
with self.subTest(
case=case.name,
backend=case.backend,
pad_style=pad_style,
capture_bs=capture_bs,
):
run_dense_draft_extend_v2_cuda_graph_case(
self,
case,
head_dim=self.HEAD_DIM,
hidden_size=self.HIDDEN_SIZE,
cuda_graph_capture_batch_size=capture_bs,
pad_style=pad_style,
)
def test_runner_mode_eagle_draft_extend_cases(self):
for case, spec_kind in self.DRAFT_EXTEND_CASES:
@@ -413,9 +413,27 @@ class TestTritonDenseAttentionBackendCorrectness(CustomTestCase):
)
def test_runner_mode_eagle_draft_extend_v2_cuda_graph_cases(self):
# pad_ratio is expressed as the captured batch size relative to the
# case's real batch size: 1.0x = no padding, 2.0x = 50% padded, etc.
for case in self.DRAFT_EXTEND_V2_CUDA_GRAPH_CASES:
with self.subTest(case=case.name, backend=case.backend):
run_dense_draft_extend_v2_cuda_graph_case(self, case)
for pad_style in ("small_real", "prod_fill"):
for capture_bs in (
case.batch_size,
case.batch_size * 2,
case.batch_size * 4,
):
with self.subTest(
case=case.name,
backend=case.backend,
pad_style=pad_style,
capture_bs=capture_bs,
):
run_dense_draft_extend_v2_cuda_graph_case(
self,
case,
cuda_graph_capture_batch_size=capture_bs,
pad_style=pad_style,
)
def test_runner_mode_eagle_draft_extend_v2_cuda_graph_runner_cases(self):
for case in self.EAGLE_DRAFT_EXTEND_V2_RUNNER_CASES:
@@ -336,8 +336,24 @@ class TestTritonMLAAttentionBackendCorrectness(CustomTestCase):
def test_runner_mode_eagle_draft_extend_v2_cuda_graph_cases(self):
for case in self.DRAFT_EXTEND_V2_CUDA_GRAPH_CASES:
with self.subTest(case=case.name, backend=case.backend):
run_mla_draft_extend_v2_cuda_graph_case(self, case)
for pad_style in ("small_real", "prod_fill"):
for capture_bs in (
case.batch_size,
case.batch_size * 2,
case.batch_size * 4,
):
with self.subTest(
case=case.name,
backend=case.backend,
pad_style=pad_style,
capture_bs=capture_bs,
):
run_mla_draft_extend_v2_cuda_graph_case(
self,
case,
cuda_graph_capture_batch_size=capture_bs,
pad_style=pad_style,
)
def test_runner_mode_eagle_draft_extend_v2_cuda_graph_runner_cases(self):
for case in self.EAGLE_DRAFT_EXTEND_V2_RUNNER_CASES: