[CI][RFC] Replace black-jupyter with ruff-format (#37210)

Co-authored-by: Alison Shao <a.shao@wustl.edu>
This commit is contained in:
Alex Nails
2026-09-02 19:46:08 -07:00
committed by GitHub
co-authored by Alison Shao
parent 2641e427be
commit 28262c20df
1411 changed files with 7766 additions and 8176 deletions
@@ -61,7 +61,7 @@ class TestDeepseekV32IndexTopkPattern(CustomTestCase):
if is_in_ci():
write_github_step_summary(
f"### test_gsm8k (deepseek-v32)\n" f'{metrics["accuracy"]=:.3f}\n'
f'### test_gsm8k (deepseek-v32)\n{metrics["accuracy"]=:.3f}\n'
)
self.assertGreater(metrics["accuracy"], 0.935)
@@ -108,7 +108,7 @@ class TestDeepseekV32IndexFreq(CustomTestCase):
if is_in_ci():
write_github_step_summary(
f"### test_gsm8k (deepseek-v32)\n" f'{metrics["accuracy"]=:.3f}\n'
f'### test_gsm8k (deepseek-v32)\n{metrics["accuracy"]=:.3f}\n'
)
self.assertGreater(metrics["accuracy"], 0.935)
@@ -246,9 +246,9 @@ class TestDeepSeekR1EvalAMD(unittest.TestCase):
for config in self.models:
display_name = config.get_display_name()
with self.subTest(model=display_name):
print(f"\n{'='*60}")
print(f"\n{'=' * 60}")
print(f"Testing: {display_name}")
print(f"{'='*60}")
print(f"{'=' * 60}")
env = os.environ.copy()
for key, value in config.env_vars.items():
@@ -98,8 +98,7 @@ class TestDeepseekV32DP(CustomTestCase):
if is_in_ci():
write_github_step_summary(
f"### test_gsm8k (deepseek-v32 DP MI325)\n"
f'{metrics["accuracy"]=:.3f}\n'
f'### test_gsm8k (deepseek-v32 DP MI325)\n{metrics["accuracy"]=:.3f}\n'
)
self.assertGreater(metrics["accuracy"], GSM8K_ACCURACY_THRESHOLD)
@@ -112,8 +111,7 @@ class TestDeepseekV32DP(CustomTestCase):
if is_in_ci():
write_github_step_summary(
f"### test_bs_1_speed (deepseek-v32 DP MI325)\n"
f"{speed=:.2f} token/s\n"
f"### test_bs_1_speed (deepseek-v32 DP MI325)\n{speed=:.2f} token/s\n"
)
self.assertGreater(speed, 10)
@@ -182,9 +182,9 @@ class TestDeepSeekV32EvalAMD(unittest.TestCase):
for config in self.models:
display_name = config.get_display_name()
with self.subTest(model=display_name):
print(f"\n{'='*60}")
print(f"\n{'=' * 60}")
print(f"Testing: {display_name}")
print(f"{'='*60}")
print(f"{'=' * 60}")
env = os.environ.copy()
for key, value in config.env_vars.items():
@@ -99,8 +99,7 @@ class TestDeepseekV32TC(CustomTestCase):
if is_in_ci():
write_github_step_summary(
f"### test_gsm8k (deepseek-v32 TC MI325)\n"
f'{metrics["accuracy"]=:.3f}\n'
f'### test_gsm8k (deepseek-v32 TC MI325)\n{metrics["accuracy"]=:.3f}\n'
)
self.assertGreater(metrics["accuracy"], GSM8K_ACCURACY_THRESHOLD)
@@ -113,8 +112,7 @@ class TestDeepseekV32TC(CustomTestCase):
if is_in_ci():
write_github_step_summary(
f"### test_bs_1_speed (deepseek-v32 TC MI325)\n"
f"{speed=:.2f} token/s\n"
f"### test_bs_1_speed (deepseek-v32 TC MI325)\n{speed=:.2f} token/s\n"
)
self.assertGreater(speed, 10)
@@ -172,9 +172,9 @@ class TestGLM51EvalAMD(unittest.TestCase):
for config in self.models:
display_name = config.get_display_name()
with self.subTest(model=display_name):
print(f"\n{'='*60}")
print(f"\n{'=' * 60}")
print(f"Testing: {display_name}")
print(f"{'='*60}")
print(f"{'=' * 60}")
env = os.environ.copy()
for key, value in config.env_vars.items():
@@ -80,8 +80,7 @@ class TestGLM51HiSparseEvalAMD(unittest.TestCase):
if is_in_ci():
write_github_step_summary(
f"### test_gsm8k (glm-5.1 hisparse mi30x)\n"
f'{metrics["score"]=:.3f}\n'
f'### test_gsm8k (glm-5.1 hisparse mi30x)\n{metrics["score"]=:.3f}\n'
)
self.assertGreater(metrics["score"], 0.93)
@@ -182,9 +182,9 @@ class TestGLM5EvalAMD(unittest.TestCase):
for config in self.models:
display_name = config.get_display_name()
with self.subTest(model=display_name):
print(f"\n{'='*60}")
print(f"\n{'=' * 60}")
print(f"Testing: {display_name}")
print(f"{'='*60}")
print(f"{'=' * 60}")
env = os.environ.copy()
for key, value in config.env_vars.items():
@@ -184,9 +184,9 @@ class TestGptOssEvalAMD(unittest.TestCase):
for config in self.models:
with self.subTest(model=config.model_path):
print(f"\n{'='*60}")
print(f"\n{'=' * 60}")
print(f"Testing: {config.model_path}")
print(f"{'='*60}")
print(f"{'=' * 60}")
env = os.environ.copy()
for key, value in config.env_vars.items():
@@ -222,9 +222,9 @@ class TestGrokEvalAMD(unittest.TestCase):
for config in self.models:
with self.subTest(model=config.model_path):
print(f"\n{'='*60}")
print(f"\n{'=' * 60}")
print(f"Testing: {config.model_path}")
print(f"{'='*60}")
print(f"{'=' * 60}")
env = os.environ.copy()
for key, value in config.env_vars.items():
@@ -181,9 +181,9 @@ def check_model_scores(results):
line = f"| {model} | {tp_size} | {score:.3f} | {threshold_str} | {startup_str} | {eval_str} | {total_str} | {status} |\n"
summary += line
print(f"\n{'='*60}")
print(f"\n{'=' * 60}")
print("SUMMARY - TP=2 Instruction Models (gsm8k)")
print(f"{'='*60}")
print(f"{'=' * 60}")
print(summary)
print(f"\n📊 Final Statistics:")
print(f" Passed: {passed_count}")
@@ -219,19 +219,19 @@ class TestNightlyGsm8KEval(unittest.TestCase):
all_results = []
total_test_start = time.time()
print(f"\n{'='*60}")
print(f"\n{'=' * 60}")
print("AMD GSM8K Evaluation Test (TP=2 Instruction Models)")
print(f"{'='*60}")
print(f"{'=' * 60}")
print(f"Benchmark: gsm8k (chat completions)")
print(f"{'='*60}\n")
print(f"{'=' * 60}\n")
for model_group, is_fp8, is_tp2 in self.model_groups:
for model in model_group:
with self.subTest(model=model):
tp_size = 2 if is_tp2 else 1
print(f"\n{'='*60}")
print(f"\n{'=' * 60}")
print(f"Testing: {model} (TP={tp_size}, FP8={is_fp8})")
print(f"{'='*60}")
print(f"{'=' * 60}")
model_start = time.time()
startup_time = None
@@ -326,7 +326,7 @@ class TestNightlyGsm8KEval(unittest.TestCase):
# Check all scores after collecting all results
check_model_scores(all_results)
print(
f"\n⏱️ Total test runtime: {total_test_time:.1f}s ({total_test_time/60:.1f} min)"
f"\n⏱️ Total test runtime: {total_test_time:.1f}s ({total_test_time / 60:.1f} min)"
)
@@ -179,9 +179,9 @@ class TestMiniMaxM25EvalAMD(unittest.TestCase):
for config in self.models:
display_name = config.get_display_name()
with self.subTest(model=display_name):
print(f"\n{'='*60}")
print(f"\n{'=' * 60}")
print(f"Testing: {display_name}")
print(f"{'='*60}")
print(f"{'=' * 60}")
env = os.environ.copy()
for key, value in config.env_vars.items():
@@ -179,9 +179,9 @@ class TestMiniMaxM27EvalAMD(unittest.TestCase):
for config in self.models:
display_name = config.get_display_name()
with self.subTest(model=display_name):
print(f"\n{'='*60}")
print(f"\n{'=' * 60}")
print(f"Testing: {display_name}")
print(f"{'='*60}")
print(f"{'=' * 60}")
env = os.environ.copy()
for key, value in config.env_vars.items():
@@ -172,14 +172,14 @@ class TestNightlyVLMMmmuEvalAMD(unittest.TestCase):
all_results = []
total_test_start = time.time()
print(f"\n{'='*60}")
print(f"\n{'=' * 60}")
print("AMD VLM MMMU Evaluation Test")
print(f"{'='*60}")
print(f"{'=' * 60}")
print(f"Benchmark: MMMU (100 samples)")
print(f"Models to test: {len(self.models)}")
for m in self.models:
print(f" - {m['model_path']} (TP={m['tp_size']})")
print(f"{'='*60}\n")
print(f"{'=' * 60}\n")
for model_config in self.models:
model_path = model_config["model_path"]
@@ -189,9 +189,9 @@ class TestNightlyVLMMmmuEvalAMD(unittest.TestCase):
error_message = None
with self.subTest(model=model_path):
print(f"\n{'='*60}")
print(f"\n{'=' * 60}")
print(f"Testing: {model_path} (TP={tp_size})")
print(f"{'='*60}")
print(f"{'=' * 60}")
model_start = time.time()
startup_time = None
@@ -358,21 +358,21 @@ class TestNightlyVLMMmmuEvalAMD(unittest.TestCase):
summary += f"| {model} | {tp_size} | {score_str} | {threshold:.2f} | {startup_str} | {eval_str} | {total_str} | {status} |\n"
print(f"\n{'='*60}")
print(f"\n{'=' * 60}")
print("SUMMARY - AMD VLM MMMU Evaluation")
print(f"{'='*60}")
print(f"{'=' * 60}")
print(summary)
print(f"\n📊 Final Statistics:")
print(f" Passed: {passed_count}")
print(f" Failed: {failed_count}")
print(
f"\n⏱️ Total test runtime: {total_test_time:.1f}s ({total_test_time/60:.1f} min)"
f"\n⏱️ Total test runtime: {total_test_time:.1f}s ({total_test_time / 60:.1f} min)"
)
if is_in_ci():
write_github_step_summary(
f"### TestNightlyVLMMmmuEvalAMD\n{summary}\n\n"
f"**Total Runtime:** {total_test_time:.1f}s ({total_test_time/60:.1f} min)"
f"**Total Runtime:** {total_test_time:.1f}s ({total_test_time / 60:.1f} min)"
)
if failed_models:
@@ -182,9 +182,9 @@ class TestDeepSeekR1EvalMI35x(unittest.TestCase):
for config in self.models:
display_name = config.get_display_name()
with self.subTest(model=display_name):
print(f"\n{'='*60}")
print(f"\n{'=' * 60}")
print(f"Testing: {display_name}")
print(f"{'='*60}")
print(f"{'=' * 60}")
env = os.environ.copy()
for key, value in config.env_vars.items():
@@ -180,9 +180,9 @@ class TestDeepSeekR1MXFP4ArFusionEvalMI35x(unittest.TestCase):
for config in self.models:
display_name = config.get_display_name()
with self.subTest(model=display_name):
print(f"\n{'='*60}")
print(f"\n{'=' * 60}")
print(f"Testing: {display_name}")
print(f"{'='*60}")
print(f"{'=' * 60}")
env = os.environ.copy()
for key, value in config.env_vars.items():
@@ -178,9 +178,9 @@ class TestDeepSeekR1MXFP4EvalMI35x(unittest.TestCase):
for config in self.models:
display_name = config.get_display_name()
with self.subTest(model=display_name):
print(f"\n{'='*60}")
print(f"\n{'=' * 60}")
print(f"Testing: {display_name}")
print(f"{'='*60}")
print(f"{'=' * 60}")
env = os.environ.copy()
for key, value in config.env_vars.items():
@@ -181,9 +181,9 @@ class TestDeepSeekR1MXFP4KvFp8EvalMI35x(unittest.TestCase):
for config in self.models:
display_name = config.get_display_name()
with self.subTest(model=display_name):
print(f"\n{'='*60}")
print(f"\n{'=' * 60}")
print(f"Testing: {display_name}")
print(f"{'='*60}")
print(f"{'=' * 60}")
env = os.environ.copy()
for key, value in config.env_vars.items():
@@ -89,8 +89,7 @@ class TestDeepseekV32DP(CustomTestCase):
if is_in_ci():
write_github_step_summary(
f"### test_gsm8k (deepseek-v32 DP MI35x)\n"
f'{metrics["accuracy"]=:.3f}\n'
f'### test_gsm8k (deepseek-v32 DP MI35x)\n{metrics["accuracy"]=:.3f}\n'
)
self.assertGreater(metrics["accuracy"], GSM8K_ACCURACY_THRESHOLD)
@@ -103,8 +102,7 @@ class TestDeepseekV32DP(CustomTestCase):
if is_in_ci():
write_github_step_summary(
f"### test_bs_1_speed (deepseek-v32 DP MI35x)\n"
f"{speed=:.2f} token/s\n"
f"### test_bs_1_speed (deepseek-v32 DP MI35x)\n{speed=:.2f} token/s\n"
)
self.assertGreater(speed, 10)
@@ -183,9 +183,9 @@ class TestDeepSeekV32EvalMI35x(unittest.TestCase):
for config in self.models:
display_name = config.get_display_name()
with self.subTest(model=display_name):
print(f"\n{'='*60}")
print(f"\n{'=' * 60}")
print(f"Testing: {display_name}")
print(f"{'='*60}")
print(f"{'=' * 60}")
env = os.environ.copy()
for key, value in config.env_vars.items():
@@ -172,9 +172,9 @@ class TestGLM51EvalMI35x(unittest.TestCase):
for config in self.models:
display_name = config.get_display_name()
with self.subTest(model=display_name):
print(f"\n{'='*60}")
print(f"\n{'=' * 60}")
print(f"Testing: {display_name}")
print(f"{'='*60}")
print(f"{'=' * 60}")
env = os.environ.copy()
for key, value in config.env_vars.items():
@@ -80,8 +80,7 @@ class TestGLM51HiSparseEvalMI35x(unittest.TestCase):
if is_in_ci():
write_github_step_summary(
f"### test_gsm8k (glm-5.1 hisparse mi35x)\n"
f'{metrics["score"]=:.3f}\n'
f'### test_gsm8k (glm-5.1 hisparse mi35x)\n{metrics["score"]=:.3f}\n'
)
self.assertGreater(metrics["score"], 0.93)
@@ -182,9 +182,9 @@ class TestGLM5EvalMI35x(unittest.TestCase):
for config in self.models:
display_name = config.get_display_name()
with self.subTest(model=display_name):
print(f"\n{'='*60}")
print(f"\n{'=' * 60}")
print(f"Testing: {display_name}")
print(f"{'='*60}")
print(f"{'=' * 60}")
env = os.environ.copy()
for key, value in config.env_vars.items():
@@ -185,9 +185,9 @@ class TestGLM5MXFP4EvalMI35x(unittest.TestCase):
for config in self.models:
display_name = config.get_display_name()
with self.subTest(model=display_name):
print(f"\n{'='*60}")
print(f"\n{'=' * 60}")
print(f"Testing: {display_name}")
print(f"{'='*60}")
print(f"{'=' * 60}")
env = os.environ.copy()
for key, value in config.env_vars.items():
@@ -196,9 +196,9 @@ class TestGptOssEvalMI35x(unittest.TestCase):
for config in self.models:
with self.subTest(model=config.model_path):
print(f"\n{'='*60}")
print(f"\n{'=' * 60}")
print(f"Testing: {config.model_path}")
print(f"{'='*60}")
print(f"{'=' * 60}")
env = os.environ.copy()
for key, value in config.env_vars.items():
@@ -185,9 +185,9 @@ class TestGptOssW4A8Mxfp4EvalMI35x(unittest.TestCase):
for config in self.models:
with self.subTest(model=config.model_path):
print(f"\n{'='*60}")
print(f"\n{'=' * 60}")
print(f"Testing: {config.model_path}")
print(f"{'='*60}")
print(f"{'=' * 60}")
env = os.environ.copy()
for key, value in config.env_vars.items():
@@ -132,9 +132,9 @@ class TestKimiK25AiterMlaEvalMI35x(unittest.TestCase):
for config in self.models:
display_name = config.get_display_name()
with self.subTest(model=display_name):
print(f"\n{'='*60}")
print(f"\n{'=' * 60}")
print(f"Testing: {display_name}")
print(f"{'='*60}")
print(f"{'=' * 60}")
env = os.environ.copy()
for key, value in config.env_vars.items():
@@ -125,9 +125,9 @@ class TestKimiK25MXFP4AiterMlaEvalMI35x(unittest.TestCase):
for config in self.models:
display_name = config.get_display_name()
with self.subTest(model=display_name):
print(f"\n{'='*60}")
print(f"\n{'=' * 60}")
print(f"Testing: {display_name}")
print(f"{'='*60}")
print(f"{'=' * 60}")
env = os.environ.copy()
for key, value in config.env_vars.items():
@@ -109,9 +109,9 @@ class TestKimiK27CodeMXFP4AiterMlaEvalMI35x(unittest.TestCase):
for config in self.models:
display_name = config.get_display_name()
with self.subTest(model=display_name):
print(f"\n{'='*60}")
print(f"\n{'=' * 60}")
print(f"Testing: {display_name}")
print(f"{'='*60}")
print(f"{'=' * 60}")
env = os.environ.copy()
for key, value in config.env_vars.items():
@@ -179,9 +179,9 @@ class TestMiniMaxM25EvalMI35x(unittest.TestCase):
for config in self.models:
display_name = config.get_display_name()
with self.subTest(model=display_name):
print(f"\n{'='*60}")
print(f"\n{'=' * 60}")
print(f"Testing: {display_name}")
print(f"{'='*60}")
print(f"{'=' * 60}")
env = os.environ.copy()
for key, value in config.env_vars.items():
@@ -188,9 +188,9 @@ class TestMiniMaxM25TP4EvalMI35x(unittest.TestCase):
for config in self.models:
display_name = config.get_display_name()
with self.subTest(model=display_name):
print(f"\n{'='*60}")
print(f"\n{'=' * 60}")
print(f"Testing: {display_name}")
print(f"{'='*60}")
print(f"{'=' * 60}")
env = os.environ.copy()
for key, value in config.env_vars.items():
@@ -179,9 +179,9 @@ class TestMiniMaxM27EvalMI35x(unittest.TestCase):
for config in self.models:
display_name = config.get_display_name()
with self.subTest(model=display_name):
print(f"\n{'='*60}")
print(f"\n{'=' * 60}")
print(f"Testing: {display_name}")
print(f"{'='*60}")
print(f"{'=' * 60}")
env = os.environ.copy()
for key, value in config.env_vars.items():
@@ -212,9 +212,9 @@ class TestMiniMaxM3TP4EvalMI35x(unittest.TestCase):
for config in self.models:
display_name = config.get_display_name()
with self.subTest(model=display_name):
print(f"\n{'='*60}")
print(f"\n{'=' * 60}")
print(f"Testing: {display_name}")
print(f"{'='*60}")
print(f"{'=' * 60}")
env = os.environ.copy()
for key, value in config.env_vars.items():
@@ -196,7 +196,7 @@ class TestQwen35Mxfp4MI35x(CustomTestCase):
if is_in_ci():
write_github_step_summary(
f"### gsm8k accuracy ({MXFP4_MODEL_PATH})\n"
f'score={metrics["score"]:.3f} '
f"score={metrics['score']:.3f} "
f"(threshold {MXFP4_ACC_THRESHOLD})\n"
)
self.assertGreater(metrics["score"], MXFP4_ACC_THRESHOLD)
@@ -233,7 +233,7 @@ class TestQwen35MoeMxfp4MI35x(CustomTestCase):
if is_in_ci():
write_github_step_summary(
f"### gsm8k run-through ({MOE_MXFP4_MODEL_PATH}, fuse_gate)\n"
f'score={metrics["score"]:.3f} (run-through, no gate)\n'
f"score={metrics['score']:.3f} (run-through, no gate)\n"
)
# Ran e2e and returned parseable answers -> fuse_gate path is healthy.
self.assertGreater(metrics["score"], 0.0)
@@ -202,9 +202,9 @@ class TestQwen3CoderNextEvalMI35x(unittest.TestCase):
for config in self.models:
display_name = config.get_display_name()
with self.subTest(model=display_name):
print(f"\n{'='*60}")
print(f"\n{'=' * 60}")
print(f"Testing: {display_name}")
print(f"{'='*60}")
print(f"{'=' * 60}")
env = os.environ.copy()
for key, value in config.env_vars.items():
@@ -164,9 +164,9 @@ class TestQwen3MoeEvalMI35x(unittest.TestCase):
for config in self.models:
with self.subTest(model=config.model_path):
print(f"\n{'='*60}")
print(f"\n{'=' * 60}")
print(f"Testing: {config.model_path}")
print(f"{'='*60}")
print(f"{'=' * 60}")
env = os.environ.copy()
for key, value in config.env_vars.items():
@@ -193,9 +193,9 @@ class TestDeepSeekR10528MXFP4EvalMI45x(unittest.TestCase):
for config in self.models:
with self.subTest(model=config.model_path):
print(f"\n{'='*60}")
print(f"\n{'=' * 60}")
print(f"Testing: {config.model_path}")
print(f"{'='*60}")
print(f"{'=' * 60}")
env = os.environ.copy()
for key, value in config.env_vars.items():
@@ -196,9 +196,9 @@ class TestDeepSeekV4FlashEvalMI45x(unittest.TestCase):
for config in self.models:
with self.subTest(model=config.model_path):
print(f"\n{'='*60}")
print(f"\n{'=' * 60}")
print(f"Testing: {config.model_path}")
print(f"{'='*60}")
print(f"{'=' * 60}")
env = os.environ.copy()
for key, value in config.env_vars.items():
@@ -182,9 +182,9 @@ class TestGptOssW4A8Mxfp4EvalMI45x(unittest.TestCase):
for config in self.models:
with self.subTest(model=config.model_path):
print(f"\n{'='*60}")
print(f"\n{'=' * 60}")
print(f"Testing: {config.model_path}")
print(f"{'='*60}")
print(f"{'=' * 60}")
env = os.environ.copy()
for key, value in config.env_vars.items():
@@ -134,12 +134,12 @@ class TestDisaggregationAccuracy(PDDisaggregationServerBase):
input_logprobs = j["meta_info"]["input_token_logprobs"]
output_logprobs = j["meta_info"]["output_token_logprobs"]
assert (
len(output_logprobs) == completion_tokens
), f"output_logprobs and completion_tokens should have the same length, but got {len(output_logprobs)} and {completion_tokens}"
assert (
len(input_logprobs) > 0
), f"input_logprobs should have at least one token, but got {len(input_logprobs)}"
assert len(output_logprobs) == completion_tokens, (
f"output_logprobs and completion_tokens should have the same length, but got {len(output_logprobs)} and {completion_tokens}"
)
assert len(input_logprobs) > 0, (
f"input_logprobs should have at least one token, but got {len(input_logprobs)}"
)
def test_structured_output(self):
json_schema = json.dumps(
@@ -76,7 +76,7 @@ class TestDeepseekR1MXFP4(CustomTestCase):
if is_in_ci():
write_github_step_summary(
f"### test_gsm8k (deepseek-r1-mxfp4)\n" f'{metrics["accuracy"]=:.3f}\n'
f'### test_gsm8k (deepseek-r1-mxfp4)\n{metrics["accuracy"]=:.3f}\n'
)
self.assertGreater(metrics["accuracy"], 0.94)
@@ -88,7 +88,7 @@ class TestDeepseekR1MXFP4(CustomTestCase):
if is_in_ci():
write_github_step_summary(
f"### test_bs_1_speed (deepseek-r1-mxfp4)\n" f"{speed=:.2f} token/s\n"
f"### test_bs_1_speed (deepseek-r1-mxfp4)\n{speed=:.2f} token/s\n"
)
self.assertGreater(speed, 75)
@@ -72,7 +72,7 @@ class TestDeepseekV32DP(CustomTestCase):
if is_in_ci():
write_github_step_summary(
f"### test_gsm8k (deepseek-v32)\n" f'{metrics["accuracy"]=:.3f}\n'
f'### test_gsm8k (deepseek-v32)\n{metrics["accuracy"]=:.3f}\n'
)
self.assertGreater(metrics["accuracy"], 0.935)
@@ -84,7 +84,7 @@ class TestDeepseekV32DP(CustomTestCase):
if is_in_ci():
write_github_step_summary(
f"### test_bs_1_speed (deepseek-v32)\n" f"{speed=:.2f} token/s\n"
f"### test_bs_1_speed (deepseek-v32)\n{speed=:.2f} token/s\n"
)
if is_in_amd_ci():
self.assertGreater(speed, 10)
@@ -140,7 +140,7 @@ class TestDeepseekV32TP(CustomTestCase):
if is_in_ci():
write_github_step_summary(
f"### test_gsm8k (deepseek-v32)\n" f'{metrics["accuracy"]=:.3f}\n'
f'### test_gsm8k (deepseek-v32)\n{metrics["accuracy"]=:.3f}\n'
)
self.assertGreater(metrics["accuracy"], 0.935)
@@ -152,7 +152,7 @@ class TestDeepseekV32TP(CustomTestCase):
if is_in_ci():
write_github_step_summary(
f"### test_bs_1_speed (deepseek-v32)\n" f"{speed=:.2f} token/s\n"
f"### test_bs_1_speed (deepseek-v32)\n{speed=:.2f} token/s\n"
)
if is_in_amd_ci():
self.assertGreater(speed, 15)
@@ -60,7 +60,7 @@ class TestDeepseekV3Basic(CustomTestCase):
if is_in_ci():
write_github_step_summary(
f"### test_gsm8k (deepseek-v3)\n" f'{metrics["accuracy"]=:.3f}\n'
f'### test_gsm8k (deepseek-v3)\n{metrics["accuracy"]=:.3f}\n'
)
self.assertGreater(metrics["accuracy"], 0.935)
@@ -72,7 +72,7 @@ class TestDeepseekV3Basic(CustomTestCase):
if is_in_ci():
write_github_step_summary(
f"### test_bs_1_speed (deepseek-v3)\n" f"{speed=:.2f} token/s\n"
f"### test_bs_1_speed (deepseek-v3)\n{speed=:.2f} token/s\n"
)
if is_in_amd_ci():
self.assertGreater(speed, 12)
@@ -64,7 +64,7 @@ class TestDeepseekV3BasicKvFp8(CustomTestCase):
if is_in_ci():
write_github_step_summary(
f"### test_gsm8k (deepseek-v3 kv-fp8)\n" f'{metrics["accuracy"]=:.3f}\n'
f'### test_gsm8k (deepseek-v3 kv-fp8)\n{metrics["accuracy"]=:.3f}\n'
)
self.assertGreater(metrics["accuracy"], 0.93)
@@ -76,7 +76,7 @@ class TestDeepseekV3BasicKvFp8(CustomTestCase):
if is_in_ci():
write_github_step_summary(
f"### test_bs_1_speed (deepseek-v3 kv-fp8)\n" f"{speed=:.2f} token/s\n"
f"### test_bs_1_speed (deepseek-v3 kv-fp8)\n{speed=:.2f} token/s\n"
)
if is_in_amd_ci():
self.assertGreater(speed, 40)
+2 -2
View File
@@ -93,7 +93,7 @@ class TestKimiK25MXFP4(CustomTestCase):
if is_in_ci():
write_github_step_summary(
f"### test_gsm8k (Kimi-K2.5-MXFP4)\n" f'{metrics["accuracy"]=:.3f}\n'
f'### test_gsm8k (Kimi-K2.5-MXFP4)\n{metrics["accuracy"]=:.3f}\n'
)
self.assertGreater(metrics["accuracy"], 0.92)
@@ -105,7 +105,7 @@ class TestKimiK25MXFP4(CustomTestCase):
if is_in_ci():
write_github_step_summary(
f"### test_bs_1_speed (Kimi-K2.5-MXFP4)\n" f"{speed=:.2f} token/s\n"
f"### test_bs_1_speed (Kimi-K2.5-MXFP4)\n{speed=:.2f} token/s\n"
)
if is_in_amd_ci():
self.assertGreater(speed, 30)
+2 -4
View File
@@ -73,8 +73,7 @@ class TestKimiK2Instruct0905(CustomTestCase):
if is_in_ci():
write_github_step_summary(
f"### test_gsm8k (Kimi-K2-Instruct-0905)\n"
f'{metrics["accuracy"]=:.3f}\n'
f'### test_gsm8k (Kimi-K2-Instruct-0905)\n{metrics["accuracy"]=:.3f}\n'
)
self.assertGreater(metrics["accuracy"], 0.94)
@@ -86,8 +85,7 @@ class TestKimiK2Instruct0905(CustomTestCase):
if is_in_ci():
write_github_step_summary(
f"### test_bs_1_speed (Kimi-K2-Instruct-0905)\n"
f"{speed=:.2f} token/s\n"
f"### test_bs_1_speed (Kimi-K2-Instruct-0905)\n{speed=:.2f} token/s\n"
)
if is_in_amd_ci():
self.assertGreater(speed, 30)
-8
View File
@@ -103,7 +103,6 @@ mtp_args = [
class TestPureDP(CustomTestCase):
@classmethod
def setUpClass(cls):
cls.model = DEFAULT_DEEPEP_MODEL_NAME_FOR_TEST
@@ -148,7 +147,6 @@ class TestPureDP(CustomTestCase):
class TestMTP(CustomTestCase):
@classmethod
def setUpClass(cls):
cls.model = DEFAULT_DEEPEP_MODEL_NAME_FOR_TEST
@@ -199,7 +197,6 @@ class TestMTP(CustomTestCase):
class TestNormal(CustomTestCase):
@classmethod
def setUpClass(cls):
cls.model = DEFAULT_DEEPEP_MODEL_NAME_FOR_TEST
@@ -247,7 +244,6 @@ class TestNormal(CustomTestCase):
class TestLowLatency(CustomTestCase):
@classmethod
def setUpClass(cls):
cls.model = DEFAULT_DEEPEP_MODEL_NAME_FOR_TEST
@@ -297,7 +293,6 @@ class TestLowLatency(CustomTestCase):
class TestTBOwithNormal(CustomTestCase):
@classmethod
def setUpClass(cls):
cls.model = DEFAULT_DEEPEP_MODEL_NAME_FOR_TEST
@@ -346,7 +341,6 @@ class TestTBOwithNormal(CustomTestCase):
class TestTBOwithLowLatency(CustomTestCase):
@classmethod
def setUpClass(cls):
cls.model = DEFAULT_DEEPEP_MODEL_NAME_FOR_TEST
@@ -397,7 +391,6 @@ class TestTBOwithLowLatency(CustomTestCase):
class TestMTPwithTBONormal(CustomTestCase):
@classmethod
def setUpClass(cls):
cls.model = DEFAULT_DEEPEP_MODEL_NAME_FOR_TEST
@@ -457,7 +450,6 @@ class TestMTPwithTBONormal(CustomTestCase):
class TestMTPwithTBOLowLatency(CustomTestCase):
@classmethod
def setUpClass(cls):
cls.model = DEFAULT_DEEPEP_MODEL_NAME_FOR_TEST
@@ -82,7 +82,7 @@ class TestQwen3CoderNext(CustomTestCase):
if is_in_ci():
write_github_step_summary(
f"### test_gsm8k (qwen3-coder-next)\n" f'{metrics["accuracy"]=:.3f}\n'
f'### test_gsm8k (qwen3-coder-next)\n{metrics["accuracy"]=:.3f}\n'
)
self.assertGreater(metrics["accuracy"], 0.90)
@@ -95,7 +95,7 @@ class TestQwen3CoderNext(CustomTestCase):
if is_in_ci():
write_github_step_summary(
f"### test_bs_1_speed (qwen3-coder-next)\n" f"{speed=:.2f} token/s\n"
f"### test_bs_1_speed (qwen3-coder-next)\n{speed=:.2f} token/s\n"
)
# self.assertGreater(speed, 50)
+2 -2
View File
@@ -68,7 +68,7 @@ class TestQwen3Instruct2507(CustomTestCase):
if is_in_ci():
write_github_step_summary(
f"### test_gsm8k ({self.model})\n" f'{metrics["accuracy"]=:.3f}\n'
f'### test_gsm8k ({self.model})\n{metrics["accuracy"]=:.3f}\n'
)
self.assertGreater(metrics["accuracy"], 0.95)
@@ -80,7 +80,7 @@ class TestQwen3Instruct2507(CustomTestCase):
if is_in_ci():
write_github_step_summary(
f"### test_bs_1_speed ({self.model})\n" f"{speed=:.2f} token/s\n"
f"### test_bs_1_speed ({self.model})\n{speed=:.2f} token/s\n"
)
if is_in_amd_ci():
self.assertGreater(speed, 50)
@@ -68,7 +68,7 @@ class TestQwen3Instruct2507FP8(CustomTestCase):
if is_in_ci():
write_github_step_summary(
f"### test_gsm8k ({self.model})\n" f'{metrics["accuracy"]=:.3f}\n'
f'### test_gsm8k ({self.model})\n{metrics["accuracy"]=:.3f}\n'
)
self.assertGreater(metrics["accuracy"], 0.95)
@@ -80,7 +80,7 @@ class TestQwen3Instruct2507FP8(CustomTestCase):
if is_in_ci():
write_github_step_summary(
f"### test_bs_1_speed ({self.model})\n" f"{speed=:.2f} token/s\n"
f"### test_bs_1_speed ({self.model})\n{speed=:.2f} token/s\n"
)
if is_in_amd_ci():
self.assertGreater(speed, 40)
@@ -72,7 +72,7 @@ class TestQwen3Instruct2507MXFP4(CustomTestCase):
if is_in_ci():
write_github_step_summary(
f"### test_gsm8k ({self.model})\n" f'{metrics["accuracy"]=:.3f}\n'
f'### test_gsm8k ({self.model})\n{metrics["accuracy"]=:.3f}\n'
)
self.assertGreater(metrics["accuracy"], 0.93)
@@ -84,7 +84,7 @@ class TestQwen3Instruct2507MXFP4(CustomTestCase):
if is_in_ci():
write_github_step_summary(
f"### test_bs_1_speed ({self.model})\n" f"{speed=:.2f} token/s\n"
f"### test_bs_1_speed ({self.model})\n{speed=:.2f} token/s\n"
)
if is_in_amd_ci():
self.assertGreater(speed, 60)
@@ -193,9 +193,9 @@ def test_kda_decode_flashinfer_matches_triton(batch_size):
idx = d["cache_indices"].long()
s_err = (st_fi[idx].float() - st_ref[idx].float()).abs()
assert s_err.max().item() < 1e-1, f"decode state max diff {s_err.max().item():.2e}"
assert (
s_err.mean().item() < 1e-2
), f"decode state mean diff {s_err.mean().item():.2e}"
assert s_err.mean().item() < 1e-2, (
f"decode state mean diff {s_err.mean().item():.2e}"
)
@pytest.mark.parametrize("batch_size,num_spec", [(1, 7), (8, 7), (32, 3)])
@@ -248,12 +248,12 @@ def test_kda_target_verify_flashinfer_checkpoint_states(
assert torch.isfinite(got).all(), "verify checkpoint states have non-finite values"
s_err = (got - ref_states).abs()
# bf16 recurrent state; same tolerance as the decode committed-state check.
assert (
s_err.max().item() < 1e-1
), f"checkpoint state max diff {s_err.max().item():.2e}"
assert (
s_err.mean().item() < 1e-2
), f"checkpoint state mean diff {s_err.mean().item():.2e}"
assert s_err.max().item() < 1e-1, (
f"checkpoint state max diff {s_err.max().item():.2e}"
)
assert s_err.mean().item() < 1e-2, (
f"checkpoint state mean diff {s_err.mean().item():.2e}"
)
def test_kda_target_verify_flashinfer_rejects_tree_spec():
@@ -127,9 +127,9 @@ def test_flashkda_matches_triton_safe_gate(seq_lens):
# bf16 cross-implementation noise (chunk=16 CUTLASS vs chunk=64 Triton);
# measured cos ~0.985 output / ~0.9999 state on H20-3e and B200.
assert _cos(ref_out, out) > 0.95, f"output cos too low: {_cos(ref_out, out):.4f}"
assert (
_cos(ref_state, st_fk[d["idx"]]) > 0.99
), f"state cos too low: {_cos(ref_state, st_fk[d['idx']]):.4f}"
assert _cos(ref_state, st_fk[d["idx"]]) > 0.99, (
f"state cos too low: {_cos(ref_state, st_fk[d['idx']]):.4f}"
)
def test_flashkda_falls_back_without_lower_bound():
@@ -191,9 +191,9 @@ def test_flashkda_spec_verify_falls_back():
assert torch.isfinite(out).all()
# Took the Triton fallback (not FlashKDA) -> matches chunk_kda closely. If
# FlashKDA had run, the cross-impl cos would be ~0.985 and this would fail.
assert (
_cos(ref_out, out) > 0.999
), f"spec-decode did not fall back: {_cos(ref_out, out):.4f}"
assert _cos(ref_out, out) > 0.999, (
f"spec-decode did not fall back: {_cos(ref_out, out):.4f}"
)
if __name__ == "__main__":
@@ -152,7 +152,6 @@ def decode_attention_fwd_torch(
class TestTritonAttention(CustomTestCase):
def _set_all_seeds(self, seed):
"""Set all random seeds for reproducibility."""
random.seed(seed)
@@ -42,9 +42,12 @@ class TestSlidingWindowAttentionTriton(CustomTestCase):
cls.short_context_prompt = "The capital of France is"
# Test prompt longer than window size
cls.long_context_prompt = """
cls.long_context_prompt = (
"""
Once upon a time, there was a mountain. In the mountain, there was a temple. In the temple, there was an old monk telling a story. The story was:
""" * 100
"""
* 100
)
cls.long_context_prompt += "\nNow, summarize the story in one sentence:"
def _test_mmlu(self):
@@ -29,7 +29,6 @@ register_amd_ci(est_time=60, suite="stage-a-test-1-gpu-small-amd")
class TestWaveAttention(unittest.TestCase):
def _set_all_seeds(self, seed):
"""Set all random seeds for reproducibility."""
random.seed(seed)
@@ -52,7 +52,6 @@ _EXTEND_CASE = DenseAttentionCase(
@unittest.skipIf(not torch.cuda.is_available(), "CUDA is required")
class TestExtendInitContract(CustomTestCase):
def _make_case(self, backend: str) -> DenseAttentionCase:
return DenseAttentionCase(
name=f"extend_no_prefix_{backend}",
@@ -179,9 +179,10 @@ class TestTRTLLMMHADenseAttentionBackendCorrectness(CustomTestCase):
def test_projected_dense_decode_cases(self):
for case_index, case in enumerate(self.DECODE_CASES):
splits = 2 if case_index == 0 else 1
with self.subTest(
case=case.name, backend=case.backend
), envs.SGLANG_TRTLLM_MHA_DECODE_SEQ_LEN_SPLITS.override(splits):
with (
self.subTest(case=case.name, backend=case.backend),
envs.SGLANG_TRTLLM_MHA_DECODE_SEQ_LEN_SPLITS.override(splits),
):
run_dense_attention_case(
self,
case,
@@ -230,9 +231,10 @@ class TestTRTLLMMHADenseAttentionBackendCorrectness(CustomTestCase):
def test_runner_mode_frozen_kv_mtp_cuda_graph_runner_cases(self):
for case in self.FROZEN_KV_MTP_RUNNER_CASES:
with self.subTest(
case=case.name, backend=case.backend
), envs.SGLANG_TRTLLM_MHA_DECODE_SEQ_LEN_SPLITS.override(2):
with (
self.subTest(case=case.name, backend=case.backend),
envs.SGLANG_TRTLLM_MHA_DECODE_SEQ_LEN_SPLITS.override(2),
):
run_dense_frozen_kv_mtp_cuda_graph_runner_case(
self,
case,
@@ -379,8 +379,9 @@ class TestFlashMLAAttentionBackendCorrectness(CustomTestCase):
)
fixture = self._build_target_verify_metadata_fixture(case)
with torch.no_grad(), forward_context(
ForwardContext(attn_backend=fixture.backend)
with (
torch.no_grad(),
forward_context(ForwardContext(attn_backend=fixture.backend)),
):
fixture.backend.init_forward_metadata(fixture.forward_batch)
@@ -62,9 +62,7 @@ def _run_rank(rank, world_size, port, scenario, result_q):
free, _total = torch.cuda.mem_get_info(rank)
target = max(free - (1 << 30), 0)
granularity_flag = (
cuda_driver.CUmemAllocationGranularity_flags.CU_MEM_ALLOC_GRANULARITY_RECOMMENDED
)
granularity_flag = cuda_driver.CUmemAllocationGranularity_flags.CU_MEM_ALLOC_GRANULARITY_RECOMMENDED
err, gran = cuda_driver.cuMemGetAllocationGranularity(
prop,
granularity_flag,
@@ -259,7 +259,6 @@ class TestSRTEndpoint(CustomTestCase):
for logprob_start_len in [0, 500, 2500, 5000, 25000]:
for return_logprob in [True, False]:
for top_logprobs_num in [0, 5]:
if logprob_start_len >= input_len:
continue
-1
View File
@@ -25,7 +25,6 @@ register_amd_ci(est_time=261, suite="stage-b-test-1-gpu-small-amd")
class TestSRTEngine(CustomTestCase):
def test_1_engine_runtime_consistency(self):
prompt = "Today is a sunny day and I like"
model_path = DEFAULT_SMALL_MODEL_NAME_FOR_TEST
+1 -2
View File
@@ -78,8 +78,7 @@ class TestDSACPV2Interleave(CustomTestCase):
if is_in_ci():
write_github_step_summary(
f"### test_a_gsm8k (dsa-cp-v2-interleave)\n"
f'{metrics["score"]=:.3f}\n'
f'### test_a_gsm8k (dsa-cp-v2-interleave)\n{metrics["score"]=:.3f}\n'
)
self.assertGreater(metrics["score"], 0.935)
-1
View File
@@ -25,7 +25,6 @@ register_cpu_ci(est_time=315, suite="base-b-tp-test-cpu")
class TestCPUGraph(CustomTestCase):
@intel_amx_benchmark(
extra_args=[
"--batch-size",
-1
View File
@@ -71,7 +71,6 @@ class TestExtendAttention(CustomTestCase):
start_q, start_kv = 0, 0
for seq_idx in range(seq_lens.shape[0]):
extend_seq_len_q = extend_seq_lens[seq_idx]
prefill_seq_len_q = extend_prefix_lens[seq_idx]
-1
View File
@@ -81,7 +81,6 @@ def flash_attn_non_varlen_ref(
class TestFlashAttn(CustomTestCase):
@parametrize(
batch=[4],
max_seqlen_q=[35, 96],
-1
View File
@@ -33,7 +33,6 @@ class Mod(nn.Module):
class TestGemm(CustomTestCase):
@parametrize(
M=[1, 101],
N=[16, 32 * 13],
@@ -25,7 +25,6 @@ register_cpu_ci(est_time=685, suite="base-b-tp-test-cpu")
class TestIntelAMXAttnBackend(CustomTestCase):
@intel_amx_benchmark(
extra_args=["--batch-size", "4", "--mem-fraction-static", "0.3"],
min_throughput=10,
@@ -18,7 +18,6 @@ register_cpu_ci(est_time=47, suite="base-b-test-cpu")
class TestIntelAMXAttnBackendQuant(CustomTestCase):
@intel_amx_benchmark(
extra_args=["--batch-size", "4", "--mem-fraction-static", "0.3"],
min_throughput=150,
@@ -18,7 +18,6 @@ register_cpu_ci(est_time=477, suite="base-b-tp-test-cpu")
class TestIntelAMXAttnBackendQuant(CustomTestCase):
@intel_amx_benchmark(
extra_args=[
"--batch-size",
-1
View File
@@ -140,7 +140,6 @@ def make_mxfp4_weights(e, out_dim, in_dim, dtype, with_bias=False):
class TestFusedExperts:
def test_unsupported_activation_is_rejected(self):
m, n, k, e, topk = 2, 32, 32, 4, 2
a = torch.randn((m, k), dtype=dtype) / 10
-5
View File
@@ -19,7 +19,6 @@ eps = 1e-6
class TestNorm:
def _forward_native(
self,
x: torch.Tensor,
@@ -201,7 +200,6 @@ class TestNorm:
class TestFusedRMSNormGated:
def _forward_native(
self,
hidden_states: torch.Tensor,
@@ -236,7 +234,6 @@ class TestFusedRMSNormGated:
class TestFusedQKRMSNorm:
@pytest.mark.parametrize("dtype", DTYPES, ids=DTYPE_IDS)
@pytest.mark.parametrize(
"batch_size,q_size,k_size,v_size",
@@ -317,7 +314,6 @@ class TestFusedQKRMSNorm:
class TestLayerNorm:
def _forward_native(
self,
x: torch.Tensor,
@@ -382,7 +378,6 @@ class TestLayerNorm:
class TestFusedQKGemmaRMSNorm:
def _gemma_rmsnorm_per_head_native(
self,
x: torch.Tensor,
@@ -226,9 +226,9 @@ class TestAssertSame(RankConsensusCheckerTestCase):
err = err_box.get()
shutdown()
assert isinstance(
err, RuntimeError
), f"Expected RuntimeError from stray-thread assert_same, got {err!r}"
assert isinstance(err, RuntimeError), (
f"Expected RuntimeError from stray-thread assert_same, got {err!r}"
)
def test_assert_same_rejects_non_scheduler_thread(self):
"""Check that assert_same() must be called in the scheduler thread. Otherwise report error."""
+13 -10
View File
@@ -56,17 +56,20 @@ class TestServerArgsIBDeviceValidation(unittest.TestCase):
real_isdir = os.path.isdir
real_listdir = os.listdir
with patch(
"sglang.srt.arg_groups.validation_hook.os.path.isdir",
side_effect=lambda path: (
True if path == "/sys/class/infiniband" else real_isdir(path)
with (
patch(
"sglang.srt.arg_groups.validation_hook.os.path.isdir",
side_effect=lambda path: (
True if path == "/sys/class/infiniband" else real_isdir(path)
),
),
), patch(
"sglang.srt.arg_groups.validation_hook.os.listdir",
side_effect=lambda path: (
available_devices
if path == "/sys/class/infiniband"
else real_listdir(path)
patch(
"sglang.srt.arg_groups.validation_hook.os.listdir",
side_effect=lambda path: (
available_devices
if path == "/sys/class/infiniband"
else real_listdir(path)
),
),
):
return validate_ib_devices(device_str)
+140 -21
View File
@@ -635,26 +635,122 @@ class TestBuildTreeKernelEfficient(CustomTestCase):
# fmt: off
torch.tensor(
[
[29889, 29974, 29945, 29900, 29974, 29922, 29930, 29958,
29889, 29974, 29930, 29945, 29974, 29922, 29930, 29958],
[22550, 4136, 16492, 8439, 29871, 2, 3001, 13,
2, 13, 29906, 29946, 2, 13, 29871, 259],
[
29889,
29974,
29945,
29900,
29974,
29922,
29930,
29958,
29889,
29974,
29930,
29945,
29974,
29922,
29930,
29958,
],
[
22550,
4136,
16492,
8439,
29871,
2,
3001,
13,
2,
13,
29906,
29946,
2,
13,
29871,
259,
],
],
),
torch.tensor(
[
[29946, 29945, 29953, 29906, 29896, 29945, 29900, 29906,
29896, 29945, 29906, 29953, 29896, 29945, 29906, 29946],
[29871, 2, 29901, 29889, 29871, 2, 395, 259,
29901, 29871, 2, 29889, 3001, 1234, 7146, 2186],
[
29946,
29945,
29953,
29906,
29896,
29945,
29900,
29906,
29896,
29945,
29906,
29953,
29896,
29945,
29906,
29946,
],
[
29871,
2,
29901,
29889,
29871,
2,
395,
259,
29901,
29871,
2,
29889,
3001,
1234,
7146,
2186,
],
],
),
torch.tensor(
[
[29946, 29974, 29945, 29930, 29889, 29922, 29974, 29930,
29974, 29946, 29930, 29922, 29889, 29974, 29945, 29922],
[29941, 29906, 2, 29946, 29871, 450, 319, 14990,
29946, 29941, 2, 29906, 29871, 2, 3001, 13],
[
29946,
29974,
29945,
29930,
29889,
29922,
29974,
29930,
29974,
29946,
29930,
29922,
29889,
29974,
29945,
29922,
],
[
29941,
29906,
2,
29946,
29871,
450,
319,
14990,
29946,
29941,
2,
29906,
29871,
2,
3001,
13,
],
],
),
# fmt: on
@@ -772,15 +868,38 @@ class TestReconstructIndicesFromTreeMask(CustomTestCase):
tree_mask = torch.tensor(
# fmt: off
[
1, 0, 0, 0,
1, 1, 0, 0,
1, 0, 1, 0,
1, 0, 1, 1,
1, 0, 0, 0,
1, 1, 0, 0,
1, 1, 1, 0,
1, 1, 1, 1,
1,
0,
0,
0,
1,
1,
0,
0,
1,
0,
1,
0,
1,
0,
1,
1,
1,
0,
0,
0,
1,
1,
0,
0,
1,
1,
1,
0,
1,
1,
1,
1,
],
# fmt: on
dtype=torch.bool,
+7 -4
View File
@@ -411,10 +411,13 @@ class TestGetDcpLens(CustomTestCase):
)
# The allocator widens from get_parallel(), not from the injected
# server_args stand-in -- drive the cause, not the effect.
with patch(
"sglang.srt.mem_cache.kv_cache_configurator.current_platform.is_out_of_tree",
return_value=False,
), rc.get_parallel().override(attn_dcp_size=dcp_size):
with (
patch(
"sglang.srt.mem_cache.kv_cache_configurator.current_platform.is_out_of_tree",
return_value=False,
),
rc.get_parallel().override(attn_dcp_size=dcp_size),
):
allocators[dcp_size] = (
KVCacheConfigurator._build_token_to_kv_pool_allocator(
configurator,
@@ -87,7 +87,7 @@ class TestQwen35TritonDCPGsm8k(CustomTestCase):
if is_in_ci():
write_github_step_summary(
f"### test_a_gsm8k (qwen3.5-triton-dcp4)\n" f'{metrics["score"]=:.3f}\n'
f'### test_a_gsm8k (qwen3.5-triton-dcp4)\n{metrics["score"]=:.3f}\n'
)
self.assertGreater(metrics["score"], GSM8K_ACCURACY_THRESHOLD)
@@ -85,8 +85,7 @@ def multiprocess_test(file: str, nproc: int, timeout: int = 120) -> None:
) from e
assert result.returncode == 0, (
f"torchrun (nproc={nproc}) failed with rc={result.returncode}\n"
f"{result.stdout}"
f"torchrun (nproc={nproc}) failed with rc={result.returncode}\n{result.stdout}"
)
@@ -473,8 +473,7 @@ class TestFormatComparisonRichMinimal:
assert_rich_tags_balanced(result)
assert result == (
"[red]❌[/] [bold red]hidden_states [/] "
"rel_diff=5.00e-01"
"[red]❌[/] [bold red]hidden_states [/] rel_diff=5.00e-01"
)
def test_shape_mismatch(self) -> None:
@@ -1046,7 +1045,7 @@ class TestFormatAbsDiffPercentilesRich:
result: str = _format_abs_diff_percentiles_rich(diff)
assert result == (
"p1=1.00e-04 p5=1.00e-04 p50=2.00e-04 " "p95=4.00e-04 p99=5.00e-04"
"p1=1.00e-04 p5=1.00e-04 p50=2.00e-04 p95=4.00e-04 p99=5.00e-04"
)
def test_high_p99_coloring(self) -> None:
@@ -1123,7 +1122,7 @@ class TestFormatReplicatedChecks:
result: str = format_replicated_checks(checks)
assert result == (
"Replicated checks:\n" " ✅ axis=tp group=0 idx=1 vs 0: n/a diff"
"Replicated checks:\n ✅ axis=tp group=0 idx=1 vs 0: n/a diff"
)
@@ -3299,9 +3299,9 @@ def _create_thd_cp_zigzag_dumps(
# Dump each rank
for cp_rank in range(cp_size):
rank_tensor: torch.Tensor = torch.cat(rank_segments[cp_rank], dim=0)
assert (
rank_tensor.shape[0] == total_per_rank
), f"rank {cp_rank}: expected {total_per_rank} tokens, got {rank_tensor.shape[0]}"
assert rank_tensor.shape[0] == total_per_rank, (
f"rank {cp_rank}: expected {total_per_rank} tokens, got {rank_tensor.shape[0]}"
)
_create_rank_dump(
directory,
@@ -4008,11 +4008,13 @@ class TestEntrypointMetaOverride:
baseline_path, target_path = self._create_single_rank_pair(tmp_path)
yaml_path: Path = tmp_path / "override.yaml"
yaml_path.write_text(textwrap.dedent("""\
yaml_path.write_text(
textwrap.dedent("""\
overrides:
- match: "hidden"
dims: "t h"
"""))
""")
)
argv = _make_argv(
baseline_path,
@@ -4148,13 +4150,15 @@ class TestEntrypointMetaOverride:
baseline_path, target_path = self._create_single_rank_pair(tmp_path)
yaml_path: Path = tmp_path / "override.yaml"
yaml_path.write_text(textwrap.dedent("""\
yaml_path.write_text(
textwrap.dedent("""\
overrides:
- match: "hidden"
dims: "t h"
- match: "hidden"
dims: "a b"
"""))
""")
)
argv = _make_argv(
baseline_path,
@@ -4169,11 +4173,13 @@ class TestEntrypointMetaOverride:
baseline_path, target_path = self._create_single_rank_pair(tmp_path)
yaml_path: Path = tmp_path / "override.yaml"
yaml_path.write_text(textwrap.dedent("""\
yaml_path.write_text(
textwrap.dedent("""\
overrides:
- match: "hidden"
dims: "a b"
"""))
""")
)
argv = _make_argv(
baseline_path,
@@ -194,11 +194,13 @@ class TestFromArgsAndConfig:
def test_cli_before_yaml(self, tmp_path: Path) -> None:
"""CLI rules are ordered before YAML rules (CLI wins on conflict)."""
yaml_path = tmp_path / "override.yaml"
yaml_path.write_text(textwrap.dedent("""\
yaml_path.write_text(
textwrap.dedent("""\
overrides:
- match: "hidden"
dims: "FROM_YAML"
"""))
""")
)
overrider = MetaOverrider.from_args_and_config(
override_dims=["hidden:FROM_CLI"],
@@ -256,14 +258,16 @@ class TestLoadYamlRules:
def test_valid_yaml(self, tmp_path: Path) -> None:
"""Valid YAML with override rules loads correctly."""
yaml_path = tmp_path / "override.yaml"
yaml_path.write_text(textwrap.dedent("""\
yaml_path.write_text(
textwrap.dedent("""\
overrides:
- match: "hidden"
dims: "b s h d"
- match: "logits"
dims: "b s v[tp]"
side: baseline
"""))
""")
)
rules = _load_yaml_rules(yaml_path)
assert len(rules) == 2
assert rules[0].dims == "b s h d"
@@ -591,9 +591,7 @@ class TestFormatAlignerPlan:
)
result: str = _format_aligner_plan(_wrap_plan(plan))
assert result == (
"Aligner Plan:\n" " baseline: (no steps)\n" " target: (no steps)"
)
assert result == ("Aligner Plan:\n baseline: (no steps)\n target: (no steps)")
def test_unsharder(self) -> None:
unsharder: UnsharderPlan = UnsharderPlan(
@@ -614,9 +612,7 @@ class TestFormatAlignerPlan:
result: str = _format_aligner_plan(_wrap_plan(plan))
assert result == (
"Aligner Plan:\n"
" baseline: (no steps)\n"
" target: [step=0: unsharder(tp)]"
"Aligner Plan:\n baseline: (no steps)\n target: [step=0: unsharder(tp)]"
)
def test_reorderer(self) -> None:
@@ -12,7 +12,7 @@ class TestApplyEdits:
"""Tests for the apply_edits() source text transformation function."""
def test_single_line_match_to_multiline_replacement(self) -> None:
source = "def foo():\n" " x = compute()\n" " return x\n"
source = "def foo():\n x = compute()\n return x\n"
edits = [
EditSpec(
match="x = compute()",
@@ -20,12 +20,10 @@ class TestApplyEdits:
)
]
result = apply_edits(source=source, edits=edits)
assert result == (
"def foo():\n" " x = compute()\n" " print(x)\n" " return x\n"
)
assert result == ("def foo():\n x = compute()\n print(x)\n return x\n")
def test_pure_insertion(self) -> None:
source = "def foo():\n" " a = 1\n" " b = 2\n"
source = "def foo():\n a = 1\n b = 2\n"
edits = [
EditSpec(
match="a = 1",
@@ -33,10 +31,10 @@ class TestApplyEdits:
)
]
result = apply_edits(source=source, edits=edits)
assert result == ("def foo():\n" " a = 1\n" " print(a)\n" " b = 2\n")
assert result == ("def foo():\n a = 1\n print(a)\n b = 2\n")
def test_pure_deletion_via_empty_replacement(self) -> None:
source = "def foo():\n" " debug_log()\n" " return 42\n"
source = "def foo():\n debug_log()\n return 42\n"
edits = [
EditSpec(
match="debug_log()",
@@ -44,10 +42,10 @@ class TestApplyEdits:
)
]
result = apply_edits(source=source, edits=edits)
assert result == ("def foo():\n" " return 42\n")
assert result == ("def foo():\n return 42\n")
def test_deletion_fewer_lines(self) -> None:
source = "def foo():\n" " a = 1\n" " b = 2\n" " c = 3\n"
source = "def foo():\n a = 1\n b = 2\n c = 3\n"
edits = [
EditSpec(
match="a = 1\nb = 2",
@@ -55,7 +53,7 @@ class TestApplyEdits:
)
]
result = apply_edits(source=source, edits=edits)
assert result == ("def foo():\n" " ab = 3\n" " c = 3\n")
assert result == ("def foo():\n ab = 3\n c = 3\n")
def test_multiline_match_to_multiline_replacement(self) -> None:
source = (
@@ -133,15 +131,7 @@ class TestApplyEdits:
def test_not_found_diagnostic_single_window_with_marker(self) -> None:
"""first line is present once but full match doesn't fit: one window with '>' on the match-region line."""
source = (
"line0\n"
"line1\n"
"line2\n"
"anchor()\n"
"wrong_next()\n"
"line5\n"
"line6\n"
)
source = "line0\nline1\nline2\nanchor()\nwrong_next()\nline5\nline6\n"
edits = [EditSpec(match="anchor()\nright_next()", replacement="x")]
with pytest.raises(PatchApplicationError) as exc_info:
apply_edits(source=source, edits=edits)
@@ -232,24 +222,22 @@ class TestApplyEdits:
assert "filler9" not in msg
def test_match_found_multiple_times_raises(self) -> None:
source = "def foo():\n" " print(1)\n" " print(1)\n"
source = "def foo():\n print(1)\n print(1)\n"
edits = [EditSpec(match="print(1)", replacement="print(2)")]
with pytest.raises(PatchApplicationError, match="multiple"):
apply_edits(source=source, edits=edits)
def test_multiple_edits_applied_sequentially(self) -> None:
source = "def foo():\n" " a = 1\n" " b = 2\n" " return a + b\n"
source = "def foo():\n a = 1\n b = 2\n return a + b\n"
edits = [
EditSpec(match="a = 1", replacement="a = 10"),
EditSpec(match="b = 2", replacement="b = 20"),
]
result = apply_edits(source=source, edits=edits)
assert result == (
"def foo():\n" " a = 10\n" " b = 20\n" " return a + b\n"
)
assert result == ("def foo():\n a = 10\n b = 20\n return a + b\n")
def test_strip_matching_ignores_leading_trailing_whitespace(self) -> None:
source = "def foo():\n" " x = compute()\n" " return x\n"
source = "def foo():\n x = compute()\n return x\n"
edits = [
EditSpec(
match=" x = compute() ",
@@ -257,11 +245,11 @@ class TestApplyEdits:
)
]
result = apply_edits(source=source, edits=edits)
assert result == ("def foo():\n" " x = replaced()\n" " return x\n")
assert result == ("def foo():\n x = replaced()\n return x\n")
def test_replacement_indented_text_realigned(self) -> None:
"""replacement text with its own indentation gets realigned to match source."""
source = "def foo():\n" " x = compute()\n" " return x\n"
source = "def foo():\n x = compute()\n return x\n"
edits = [
EditSpec(
match="x = compute()",
@@ -270,15 +258,12 @@ class TestApplyEdits:
]
result = apply_edits(source=source, edits=edits)
assert result == (
"def foo():\n"
" x = compute()\n"
" print(x)\n"
" return x\n"
"def foo():\n x = compute()\n print(x)\n return x\n"
)
def test_replacement_with_existing_indent_realigned(self) -> None:
"""replacement text already has indentation that should be rebased."""
source = "def foo():\n" " if True:\n" " x = 1\n" " return x\n"
source = "def foo():\n if True:\n x = 1\n return x\n"
edits = [
EditSpec(
match="x = 1",
@@ -296,12 +281,10 @@ class TestApplyEdits:
)
def test_append_keeps_match_and_adds_after(self) -> None:
source = "def foo():\n" " x = compute()\n" " return x\n"
source = "def foo():\n x = compute()\n return x\n"
edits = [EditSpec(match="x = compute()", append="print(x)")]
result = apply_edits(source=source, edits=edits)
assert result == (
"def foo():\n" " x = compute()\n" " print(x)\n" " return x\n"
)
assert result == ("def foo():\n x = compute()\n print(x)\n return x\n")
def test_append_multiline_match(self) -> None:
source = (
@@ -330,21 +313,18 @@ class TestApplyEdits:
)
def test_prepend_adds_before_match(self) -> None:
source = "def foo():\n" " x = compute()\n" " return x\n"
source = "def foo():\n x = compute()\n return x\n"
edits = [EditSpec(match="x = compute()", prepend="print('before')")]
result = apply_edits(source=source, edits=edits)
assert result == (
"def foo():\n"
" print('before')\n"
" x = compute()\n"
" return x\n"
"def foo():\n print('before')\n x = compute()\n return x\n"
)
def test_prepend_multiline(self) -> None:
source = "def foo():\n" " return x\n"
source = "def foo():\n return x\n"
edits = [EditSpec(match="return x", prepend="a = 1\nb = 2")]
result = apply_edits(source=source, edits=edits)
assert result == ("def foo():\n" " a = 1\n" " b = 2\n" " return x\n")
assert result == ("def foo():\n a = 1\n b = 2\n return x\n")
def test_prepend_deep_indent(self) -> None:
source = (
@@ -365,11 +345,7 @@ class TestApplyEdits:
def test_prepend_multiline_match(self) -> None:
source = (
"def foo():\n"
" result = call(\n"
" a=1,\n"
" )\n"
" return result\n"
"def foo():\n result = call(\n a=1,\n )\n return result\n"
)
edits = [
EditSpec(
@@ -401,13 +377,13 @@ class TestApplyEdits:
def test_second_edit_sees_result_of_first(self) -> None:
"""Edits are applied sequentially; second edit matches modified source."""
source = "def foo():\n" " x = 1\n" " return x\n"
source = "def foo():\n x = 1\n return x\n"
edits = [
EditSpec(match="x = 1", replacement="x = 1\ny = 2"),
EditSpec(match="y = 2", replacement="y = 20"),
]
result = apply_edits(source=source, edits=edits)
assert result == ("def foo():\n" " x = 1\n" " y = 20\n" " return x\n")
assert result == ("def foo():\n x = 1\n y = 20\n return x\n")
if __name__ == "__main__":
+52 -49
View File
@@ -734,9 +734,9 @@ def _assert_files(filenames, *, exist=(), not_exist=()):
for p in exist:
assert any(p in f for f in filenames), f"{p} not found in {filenames}"
for p in not_exist:
assert not any(
p in f for f in filenames
), f"{p} should not exist in {filenames}"
assert not any(p in f for f in filenames), (
f"{p} should not exist in {filenames}"
)
def _load_dump(path: Path) -> dict:
@@ -750,9 +750,9 @@ def _find_dump_file(tmpdir, *, rank: int = 0, name: str) -> Path:
for f in Path(tmpdir).glob("*/*.pt")
if f"rank={rank}" in f.name and name in f.name
]
assert (
len(matches) == 1
), f"Expected 1 file matching rank={rank} name={name}, got {matches}"
assert len(matches) == 1, (
f"Expected 1 file matching rank={rank} name={name}, got {matches}"
)
return matches[0]
@@ -1657,9 +1657,9 @@ class TestZmqPortIsolation:
)
resp.raise_for_status()
states = resp.json()
assert (
len(states) == 2
), f"Instance {i} (port {port}): expected 2 ranks, got {len(states)}"
assert len(states) == 2, (
f"Instance {i} (port {port}): expected 2 ranks, got {len(states)}"
)
finally:
for event in stop_events:
event.set()
@@ -1719,9 +1719,9 @@ class TestDumperHttp:
val = state
for k in keys:
val = val[k]
assert (
val == expected
), f"rank {rank}: {path}={val!r}, expected {expected!r}"
assert val == expected, (
f"rank {rank}: {path}={val!r}, expected {expected!r}"
)
def test_configure_enable_toggle(self, dumper_http_url: str):
for enable in [True, False]:
@@ -1915,9 +1915,9 @@ class TestNonIntrusiveDumper(_NonIntrusiveTestBase):
)
dumped_output = captured[f"{P}model.mutator.output"]["value"]
assert (
dumped_output == 999.0
).all(), "post-hook should capture outputs after forward"
assert (dumped_output == 999.0).all(), (
"post-hook should capture outputs after forward"
)
def test_hooks_all_module_levels(self, tmp_path):
class Attention(torch.nn.Module):
@@ -2374,9 +2374,9 @@ class TestDumperE2E:
states = requests.post(f"{base_url}/dumper/get_state", json={}).json()
assert len(states) == 2
for rank, state in enumerate(states):
assert (
state["config"]["enable"] is True
), f"rank {rank}: enable should be True after configure"
assert state["config"]["enable"] is True, (
f"rank {rank}: enable should be True after configure"
)
assert state["config"]["dir"] == dump_dir
resp = requests.post(
@@ -2403,16 +2403,16 @@ class TestDumperE2E:
)
for rank in range(2):
assert any(
f"rank={rank}" in f for f in filenames
), f"No dump files for rank {rank}"
assert any(f"rank={rank}" in f for f in filenames), (
f"No dump files for rank {rank}"
)
sample_file = dump_files[0]
loaded = torch.load(sample_file, map_location="cpu", weights_only=False)
assert isinstance(loaded, dict), f"Expected dict, got {type(loaded)}"
assert (
"value" in loaded and "meta" in loaded
), f"Missing value/meta keys: {loaded.keys()}"
assert "value" in loaded and "meta" in loaded, (
f"Missing value/meta keys: {loaded.keys()}"
)
assert "name" in loaded["meta"]
assert "rank" in loaded["meta"]
assert "step" in loaded["meta"]
@@ -2438,22 +2438,22 @@ class TestDumperE2E:
"attn_cp_size",
]
for key in expected_keys:
assert (
key in par
), f"Missing {key} in sglang_parallel_info, got: {sorted(par)}"
assert key in par, (
f"Missing {key} in sglang_parallel_info, got: {sorted(par)}"
)
rids_files = [f for f in dump_files if "name=rids" in f.name]
rids_loaded = torch.load(
rids_files[0], map_location="cpu", weights_only=False
)
rids_value = rids_loaded["value"]
assert isinstance(
rids_value, list
), f"rids should be a list, got {type(rids_value)}"
assert isinstance(rids_value, list), (
f"rids should be a list, got {type(rids_value)}"
)
assert len(rids_value) > 0, "rids should be non-empty"
assert all(
isinstance(r, str) for r in rids_value
), f"each rid should be a str, got {[type(r) for r in rids_value]}"
assert all(isinstance(r, str) for r in rids_value), (
f"each rid should be a str, got {[type(r) for r in rids_value]}"
)
finally:
kill_process_tree(proc.pid)
@@ -2912,9 +2912,9 @@ class TestRecomputeStatus:
model(torch.randn(2, 4))
for key, data in captured.items():
assert (
"recompute_status" in data["meta"]
), f"missing recompute_status in {key}"
assert "recompute_status" in data["meta"], (
f"missing recompute_status in {key}"
)
assert data["meta"]["recompute_status"] == "disabled"
def test_detect_recompute_status_default(self) -> None:
@@ -3553,8 +3553,7 @@ class TestGrafterDistributed:
# worker prepends tmp_path to sys.path so import_module sees it.
module_name = "_xform_user_basic"
(tmp_path / f"{module_name}.py").write_text(
"def transform(graft_input):\n"
" return graft_input.received_list[0] * 2\n"
"def transform(graft_input):\n return graft_input.received_list[0] * 2\n"
)
graft_port = find_available_port(29610)
_run_graft_test(
@@ -3646,7 +3645,9 @@ class TestGrafterDistributed:
7.0,
7.0,
7.0,
], f"target should be unchanged after shape-mismatch graft, got {target.tolist()}"
], (
f"target should be unchanged after shape-mismatch graft, got {target.tolist()}"
)
finally:
if grafter._pg is not None:
dist.destroy_process_group(grafter._pg)
@@ -3694,7 +3695,9 @@ class TestGrafterDistributed:
9.0,
9.0,
9.0,
], f"target must be unchanged when transform throws, got {target.tolist()}"
], (
f"target must be unchanged when transform throws, got {target.tolist()}"
)
output = captured.getvalue()
assert "transform/copy_ raised RuntimeError" in output, output
assert "intentional test error" in output, output
@@ -3783,9 +3786,9 @@ class TestGrafterDistributed:
grafter.maybe_intercept(value=target, tags={"name": "x"})
output = captured.getvalue()
if rank == 0:
assert (
"WARNING" in output
), f"expected WARNING in rank 0 output: {output}"
assert "WARNING" in output, (
f"expected WARNING in rank 0 output: {output}"
)
assert "has not completed after 2s" in output, output
finally:
if grafter._pg is not None:
@@ -3852,9 +3855,9 @@ class TestGrafterDistributed:
pg_after_first = grafter._pg
assert pg_after_first is not None
grafter.maybe_intercept(value=t2, tags={"name": "x"})
assert (
grafter._pg is pg_after_first
), "_pg must be cached across calls, not re-initialized"
assert grafter._pg is pg_after_first, (
"_pg must be cached across calls, not re-initialized"
)
else:
target1 = torch.zeros(3, device="cuda:1")
target2 = torch.zeros(3, device="cuda:1")
@@ -4190,9 +4193,9 @@ def _e2e_transform(graft_input):
the transform is just identity. Real workflows would compute a
non-trivial override (scale, reshape, decode, ...) using the extras.
"""
assert (
graft_input.received_extras_list[0]["my_extra_key"] == "my_extra_value"
), graft_input.received_extras_list
assert graft_input.received_extras_list[0]["my_extra_key"] == "my_extra_value", (
graft_input.received_extras_list
)
return graft_input.received_list[0]
@@ -290,7 +290,7 @@ def _run_e2e_scenario(
print(f"Comparator debug output: {debug_file}")
assert result.returncode == 0, (
f"Comparator failed (rc={result.returncode}). " f"Debug output: {debug_file}"
f"Comparator failed (rc={result.returncode}). Debug output: {debug_file}"
)
@@ -31,7 +31,6 @@ TEST_HIDDEN_SIZE = 32
class SimpleModel(nn.Module):
def __init__(self) -> None:
super().__init__()
self.hidden_size = TEST_HIDDEN_SIZE
@@ -76,12 +76,12 @@ class TestDisaggregationAccuracy(PauseResumeInPlaceMixin, PDDisaggregationServer
input_logprobs = j["meta_info"]["input_token_logprobs"]
output_logprobs = j["meta_info"]["output_token_logprobs"]
assert (
len(output_logprobs) == completion_tokens
), f"output_logprobs and completion_tokens should have the same length, but got {len(output_logprobs)} and {completion_tokens}"
assert (
len(input_logprobs) > 0
), f"input_logprobs should have at least one token, but got {len(input_logprobs)}"
assert len(output_logprobs) == completion_tokens, (
f"output_logprobs and completion_tokens should have the same length, but got {len(output_logprobs)} and {completion_tokens}"
)
assert len(input_logprobs) > 0, (
f"input_logprobs should have at least one token, but got {len(input_logprobs)}"
)
def test_chat_completion_top_logprobs(self):
client = openai.Client(api_key="empty", base_url=f"{self.lb_url}/v1")
@@ -461,8 +461,7 @@ class TestKimiLinearPDDCP4(GSM8KMixin, PDDisaggregationServerBase):
reference,
actual,
label=(
f"niah prompt_tokens={LONG_CONTEXT_TOKENS} "
f"depth={needle_depth}"
f"niah prompt_tokens={LONG_CONTEXT_TOKENS} depth={needle_depth}"
),
)
@@ -138,9 +138,9 @@ def test_prefill_cache_hit(base_url, input_ids, max_new_tokens, is_hicache=False
hit_label = "L1 hits" if is_hicache else "cache hits"
print(f" {hit_label}: {len(new_input_ids)}/{len(input_ids)}")
assert (
len(new_input_ids) > len(input_ids) // 2
), f"too few {hit_label}: {len(new_input_ids)}/{len(input_ids)}"
assert len(new_input_ids) > len(input_ids) // 2, (
f"too few {hit_label}: {len(new_input_ids)}/{len(input_ids)}"
)
input_logprobs = _get_input_logprobs(base_url, new_input_ids, output_logprobs)
return _compute_kl(input_logprobs, output_logprobs)
@@ -170,9 +170,9 @@ def test_decode_cache_hit(base_url, input_ids, max_new_tokens, is_hicache=False)
hit_label = "L1 decode hits" if is_hicache else "cache hits"
print(f" {hit_label}: {len(new_input_ids)}/{len(turn2_ids)}")
assert (
len(new_input_ids) > len(turn2_ids) // 2
), f"too few {hit_label}: {len(new_input_ids)}/{len(turn2_ids)}"
assert len(new_input_ids) > len(turn2_ids) // 2, (
f"too few {hit_label}: {len(new_input_ids)}/{len(turn2_ids)}"
)
input_logprobs = _get_input_logprobs(base_url, new_input_ids, output_logprobs)
return _compute_kl(input_logprobs, output_logprobs)
@@ -281,8 +281,7 @@ class _DPAttentionPrefillCudaGraphKLMixin:
return
time.sleep(0.5)
self.fail(
f"No {self.prefill_backend} prefill CUDA graph replay was logged "
f"for {case}"
f"No {self.prefill_backend} prefill CUDA graph replay was logged for {case}"
)
def test_prefill_and_decode_cache_hit_kl_is_zero(self):
@@ -173,16 +173,16 @@ def test_solve_ipm_matches_torch_reference():
max_diff = (cuda_x - torch_x).abs().max().item()
print(
f"\n[ipm-compare] converged={converged} max|cuda-torch|={max_diff:.3e} "
f"cuda={[round(v,4) for v in cuda_x.tolist()]} "
f"torch={[round(v,4) for v in torch_x.tolist()]}"
f"cuda={[round(v, 4) for v in cuda_x.tolist()]} "
f"torch={[round(v, 4) for v in torch_x.tolist()]}"
)
assert converged, (
"IPM returned the 0.5 non-convergence sentinel — the comparison would "
"be trivial. Adjust the LP instance so it converges."
)
assert torch.allclose(
cuda_x, torch_x, atol=1e-2, rtol=1e-2
), f"fused IPM diverges from torch reference: max abs diff {max_diff:.3e}"
assert torch.allclose(cuda_x, torch_x, atol=1e-2, rtol=1e-2), (
f"fused IPM diverges from torch reference: max abs diff {max_diff:.3e}"
)
@pytest.mark.skipif(
@@ -363,9 +363,9 @@ def _check_all_ranks_empty(rank: int, world_size: int, device: torch.device):
assert torch.isfinite(actual).all(), f"rank {rank}: non-finite for empty-batch"
assert (actual >= 0).all(), f"rank {rank}: negative for empty-batch"
assert torch.allclose(
actual, expected, atol=1e-4, rtol=1e-3
), f"rank {rank}: empty-batch output disagrees with all-zero oracle"
assert torch.allclose(actual, expected, atol=1e-4, rtol=1e-3), (
f"rank {rank}: empty-batch output disagrees with all-zero oracle"
)
def _check_solver_determinism(rank: int, world_size: int, device: torch.device):
@@ -309,7 +309,7 @@ class TestKimiK2DetectorSpecialTokenLeakage(unittest.TestCase):
def test_no_leak_on_error_fallback(self):
"""On parse errors, normal_text fallback has tokens stripped."""
cleaned = _strip_special_tokens(
"leaked<|tool_calls_section_begin|>" "<|tool_call_end|>content"
"leaked<|tool_calls_section_begin|><|tool_call_end|>content"
)
self.assertEqual(cleaned, "leakedcontent")
@@ -992,7 +992,7 @@ class TestKimiK2EndToEnd(unittest.TestCase):
"<|tool_call_begin|>functions.get_weather:1"
f'<|tool_call_argument_begin|>{{"city":',
' "Bad", "valid": fasle',
"<|tool_call_end|>" "<|tool_calls_section_end|>" + good_section_1,
"<|tool_call_end|><|tool_calls_section_end|>" + good_section_1,
],
}
@@ -31,7 +31,6 @@ _ROUTER_SHAPES = ((6144, 768), (3072, 384))
"requires HPC-Ops (https://github.com/Tencent/hpc-ops) and a Hopper GPU",
)
class TestLinearBf16Fp32Hpc(CustomTestCase):
@classmethod
def setUpClass(cls):
mark_hpc_bf16xfp32_gemm_enabled()
@@ -246,7 +246,7 @@ def benchmark(message_KB: int, provider: str):
)
if provider == "aot" and world_size not in AOT_SUPPORTED_WORLD_SIZES:
marker.skip(
f"AOT custom_all_reduce needs world_size in " f"{AOT_SUPPORTED_WORLD_SIZES}"
f"AOT custom_all_reduce needs world_size in {AOT_SUPPORTED_WORLD_SIZES}"
)
_init_all_backends()
backend = BACKEND_FACTORY[provider]()
@@ -88,7 +88,7 @@ def _precompile_kernels(num_gpus: List[int]) -> None:
p.join()
if p.exitcode != 0:
raise RuntimeError(
f"TP QKNorm precompile failed for {world_size=} " f"(exit {p.exitcode})"
f"TP QKNorm precompile failed for {world_size=} (exit {p.exitcode})"
)
@@ -128,12 +128,12 @@ def bench_fused_scale_residual_norm_scale_shift(
if __name__ == "__main__":
print(f"\n{'='*80}")
print(f"\n{'=' * 80}")
print("Benchmark: fused_norm_scale_shift")
print(f"{'='*80}\n")
print(f"{'=' * 80}\n")
bench_fused_norm_scale_shift.run(print_data=True)
print(f"\n{'='*80}")
print(f"\n{'=' * 80}")
print("Benchmark: fused_scale_residual_norm_scale_shift")
print(f"{'='*80}\n")
print(f"{'=' * 80}\n")
bench_fused_scale_residual_norm_scale_shift.run(print_data=True)
@@ -81,9 +81,7 @@ def benchmark() -> None:
repeats = 5 if is_in_ci() else 20
rounds = 5 if is_in_ci() else 13
print(
"| workload | gate | torch us | triton us | cuda us | reference | " "ref/cuda |"
)
print("| workload | gate | torch us | triton us | cuda us | reference | ref/cuda |")
print("|---|---|---:|---:|---:|---|---:|")
for workload in workloads:
@@ -76,7 +76,9 @@ def _verify_num_slots(case: BenchCase) -> int:
return max(2, case.bs * per_req_slots + 1)
def _build_verify_inputs(case: BenchCase, *, device: torch.device) -> Tuple[
def _build_verify_inputs(
case: BenchCase, *, device: torch.device
) -> Tuple[
torch.Tensor,
VerifyPlan,
torch.Tensor,
@@ -133,9 +133,9 @@ def test_activation_filter_expert(
if kept.any():
torch.testing.assert_close(out[kept], expected[kept], atol=atol, rtol=rtol)
if token_skip.any():
assert torch.isnan(
out[token_skip]
).all(), "filter_expert kernel touched rows whose expert_id is -1"
assert torch.isnan(out[token_skip]).all(), (
"filter_expert kernel touched rows whose expert_id is -1"
)
@pytest.mark.parametrize("op_name", OPS)
@@ -267,8 +267,9 @@ def test_cutedsl_gdn_performance(B: int):
# Benchmark
triton_times, cutedsl_times = [], []
for _ in range(bench_iters):
start, end = torch.cuda.Event(enable_timing=True), torch.cuda.Event(
enable_timing=True
start, end = (
torch.cuda.Event(enable_timing=True),
torch.cuda.Event(enable_timing=True),
)
start.record()
if graph_triton:
@@ -279,8 +280,9 @@ def test_cutedsl_gdn_performance(B: int):
torch.cuda.synchronize()
triton_times.append(start.elapsed_time(end))
start, end = torch.cuda.Event(enable_timing=True), torch.cuda.Event(
enable_timing=True
start, end = (
torch.cuda.Event(enable_timing=True),
torch.cuda.Event(enable_timing=True),
)
with torch.cuda.stream(torch_stream):
start.record()
@@ -178,9 +178,9 @@ def test_q_rope_quant_matches_reference(pos_dtype):
# scale step at the bottom of the range.
deq = q_fp8.float() * scale
err = (deq - ref).abs()
assert (
err <= 0.0625 * ref.abs() + scale
).all(), f"max fp8 dequant error {err.max().item()}"
assert (err <= 0.0625 * ref.abs() + scale).all(), (
f"max fp8 dequant error {err.max().item()}"
)
# ----------------------------------------------------------------------------

Some files were not shown because too many files have changed in this diff Show More