minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

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

Browse Source 
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
Download Patch File Download Diff File Expand all files Collapse all files
test/registered/kernels/benchmark/communication/bench_custom_all_reduce.py
+1 -1
View File
@@ -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]()
test/registered/kernels/benchmark/communication/bench_tp_qknorm.py
+1 -1
View File
@@ -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})"
)
test/registered/kernels/benchmark/diffusion/bench_fused_norm_scale_shift.py
+4 -4
View File
@@ -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)
test/registered/kernels/benchmark/diffusion/bench_residual_gate_add.py
+1 -3
View File
@@ -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:
test/registered/kernels/benchmark/kv_canary/bench_verify.py
+3 -1
View File
@@ -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,
test/registered/kernels/ops/activation/test_activation.py
+3 -3
View File
@@ -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)
test/registered/kernels/ops/attention/test_cutedsl_gdn.py
+6 -4
View File
@@ -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()
test/registered/kernels/ops/attention/test_dsv32_indexer_fusion.py
+3 -3
View File
@@ -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()}"
)
# ----------------------------------------------------------------------------
test/registered/kernels/ops/attention/test_flash_attention_4_sm120.py
+2 -6
View File
@@ -1965,9 +1965,7 @@ def test_sm120_paged_decode_transpose_is_cache_order_independent():
num_pages,
device="cuda",
dtype=torch.int64,
).to(
torch.int32
)[None]
).to(torch.int32)[None]
cache_seqlens = torch.tensor(
[max_seqlen],
device="cuda",
@@ -2097,9 +2095,7 @@ def test_sm120_paged_decode_graph_pdl_is_correct_and_eager_reusable(
num_pages,
device="cuda",
dtype=torch.int64,
).to(
torch.int32
)[None]
).to(torch.int32)[None]
cache_seqlens = torch.tensor(
[max_seqlen],
device="cuda",
test/registered/kernels/ops/attention/test_flash_mla_backends.py
+1 -3
View File
@@ -152,9 +152,7 @@ def _build_kvcache(
nope_dequant = (
nope_fp8.view(num_pages, page_size, _NUM_TILES, _TILE_SIZE)
* scale_e8m0.view(num_pages, page_size, _NUM_TILES, 1)
).view(
num_pages, page_size, _NOPE_DIM
) # float32
).view(num_pages, page_size, _NOPE_DIM) # float32
ref_per_token = torch.cat(
[nope_dequant.to(torch.bfloat16), rope_bf16_vals], dim=-1
) # (num_pages, page_size, 512) bf16
test/registered/kernels/ops/attention/test_fused_metadata_copy.py
+28 -24
View File
@@ -396,29 +396,29 @@ def test_fused_metadata_copy(bs, forward_mode, has_real_page_table, has_flashmla
)
# Compare results
assert torch.equal(
dst_ref["cache_seqlens"], dst_fused["cache_seqlens"]
), "cache_seqlens mismatch"
assert torch.equal(
dst_ref["cu_seqlens_k"], dst_fused["cu_seqlens_k"]
), "cu_seqlens_k mismatch"
assert torch.equal(
dst_ref["page_table_1"], dst_fused["page_table_1"]
), "page_table_1 mismatch"
assert torch.equal(
dst_ref["dsa_cache_seqlens"], dst_fused["dsa_cache_seqlens"]
), "dsa_cache_seqlens mismatch"
assert torch.equal(dst_ref["cache_seqlens"], dst_fused["cache_seqlens"]), (
"cache_seqlens mismatch"
)
assert torch.equal(dst_ref["cu_seqlens_k"], dst_fused["cu_seqlens_k"]), (
"cu_seqlens_k mismatch"
)
assert torch.equal(dst_ref["page_table_1"], dst_fused["page_table_1"]), (
"page_table_1 mismatch"
)
assert torch.equal(dst_ref["dsa_cache_seqlens"], dst_fused["dsa_cache_seqlens"]), (
"dsa_cache_seqlens mismatch"
)
assert torch.equal(
dst_ref["dsa_seqlens_expanded"], dst_fused["dsa_seqlens_expanded"]
), "dsa_seqlens_expanded mismatch"
assert torch.equal(
dst_ref["dsa_cu_seqlens_k"], dst_fused["dsa_cu_seqlens_k"]
), "dsa_cu_seqlens_k mismatch"
assert torch.equal(dst_ref["dsa_cu_seqlens_k"], dst_fused["dsa_cu_seqlens_k"]), (
"dsa_cu_seqlens_k mismatch"
)
if has_real_page_table:
assert torch.equal(
dst_ref["real_page_table"], dst_fused["real_page_table"]
), "real_page_table mismatch"
assert torch.equal(dst_ref["real_page_table"], dst_fused["real_page_table"]), (
"real_page_table mismatch"
)
if has_flashmla:
assert torch.equal(
@@ -643,7 +643,11 @@ def test_fused_metadata_copy_multi_dtype_validation():
# Create source tensors - one with WRONG dtype
cache_seqlens_src_wrong = torch.randint(
1, max_len, (bs,), dtype=torch.int64, device=device # Wrong dtype!
1,
max_len,
(bs,),
dtype=torch.int64,
device=device, # Wrong dtype!
)
cu_seqlens_k_src = torch.zeros(bs + 1, dtype=torch.int32, device=device)
page_indices_src = torch.randint(
@@ -829,7 +833,7 @@ def test_fused_metadata_copy_multi(bs, has_real_page_table, has_flashmla):
f"\n[VERIFY] bs={bs}, real_page_table={has_real_page_table}, flashmla={has_flashmla}"
)
print(
f"[VERIFY] Fused time: {fused_time*1000:.3f}ms, Loop time: {loop_time*1000:.3f}ms, Speedup: {speedup:.2f}x"
f"[VERIFY] Fused time: {fused_time * 1000:.3f}ms, Loop time: {loop_time * 1000:.3f}ms, Speedup: {speedup:.2f}x"
)
max_diff = 0.0
@@ -1063,7 +1067,7 @@ def test_fused_metadata_copy_multi_large_batch(bs):
speedup = loop_time / fused_time if fused_time > 0 else 0
print(
f"\n[PERF] Large batch (bs={bs}): Fused={fused_time*1000:.3f}ms, Loop={loop_time*1000:.3f}ms, Speedup={speedup:.2f}x"
f"\n[PERF] Large batch (bs={bs}): Fused={fused_time * 1000:.3f}ms, Loop={loop_time * 1000:.3f}ms, Speedup={speedup:.2f}x"
)
# Verify correctness
@@ -1076,9 +1080,9 @@ def test_fused_metadata_copy_multi_large_batch(bs):
):
for key in dst_ref:
if dst_ref[key] is not None and dst_fused[key] is not None:
assert torch.equal(
dst_ref[key], dst_fused[key]
), f"Backend {backend_idx} {key} mismatch"
assert torch.equal(dst_ref[key], dst_fused[key]), (
f"Backend {backend_idx} {key} mismatch"
)
if __name__ == "__main__":
test/registered/kernels/ops/attention/test_fused_store_index_cache.py
+3 -3
View File
@@ -392,9 +392,9 @@ def test_roundtrip_reconstruction(num_tokens: int):
per_row_energy = reconstructed.abs().sum(dim=-1)
orig_energy = original.abs().sum(dim=-1)
mask = orig_energy > 0.1
assert (
per_row_energy[mask] > 0.01
).all(), "Some tokens have zero reconstruction — kernel may not be writing output"
assert (per_row_energy[mask] > 0.01).all(), (
"Some tokens have zero reconstruction — kernel may not be writing output"
)
# TEST 4: Boundary conditions
test/registered/kernels/ops/attention/test_minimax_decode_topk_page_table.py
+6 -6
View File
@@ -159,9 +159,9 @@ def test_dp_flattened_page_table(nkv, bs, seq_len):
row = b * nkv + h
# effective KV length = sum of valid tokens over selected blocks
exp_kv = sum(min(block, seq_len - c * block) for c in blocks)
assert (
int(cache[row]) == exp_kv
), f"row {row}: {int(cache[row])} != {exp_kv}"
assert int(cache[row]) == exp_kv, (
f"row {row}: {int(cache[row])} != {exp_kv}"
)
# page table: each block -> ppb pages via req_to_token, head-minor encoded
for e in range(len(blocks) * ppb):
c = blocks[e // ppb]
@@ -169,9 +169,9 @@ def test_dp_flattened_page_table(nkv, bs, seq_len):
if tok >= max_kv:
tok = max_kv - 1
exp = int(r2t_cpu[b, tok]) // ps * nkv + h
assert (
int(pt[row, e]) == exp
), f"row {row} e {e}: {int(pt[row,e])} != {exp}"
assert int(pt[row, e]) == exp, (
f"row {row} e {e}: {int(pt[row, e])} != {exp}"
)
if __name__ == "__main__":
test/registered/kernels/ops/attention/test_paged_mqa_metadata.py
+6 -6
View File
@@ -58,12 +58,12 @@ def paged_mqa_metadata_ref(
seq_lens: torch.Tensor, num_sm: int, page_size: int
) -> torch.Tensor:
assert page_size == 64, f"page_size must be 64, got {page_size}"
assert (
seq_lens.dtype == torch.int32
), f"seq_lens dtype must be int32, got {seq_lens.dtype}"
assert (
seq_lens.dim() == 1
), f"seq_lens must be 1-D, got shape {tuple(seq_lens.shape)}"
assert seq_lens.dtype == torch.int32, (
f"seq_lens dtype must be int32, got {seq_lens.dtype}"
)
assert seq_lens.dim() == 1, (
f"seq_lens must be 1-D, got shape {tuple(seq_lens.shape)}"
)
device = seq_lens.device
batch_size = int(seq_lens.shape[0])
test/registered/kernels/ops/attention/test_pos_enc.py
+1 -1
View File
@@ -481,7 +481,7 @@ def test_performance(
print(
f"\nPerformance Test - Batch={batch_size}, SeqLen={seq_len}, Tokens={total_tokens}"
)
print(f"JIT: {jit_time*1000:.9f}ms, SGL: {sgl_time*1000:.9f}ms")
print(f"JIT: {jit_time * 1000:.9f}ms, SGL: {sgl_time * 1000:.9f}ms")
if sgl_time > 0:
speedup = sgl_time / jit_time if jit_time > 0 else float("inf")
print(f"Speedup (SGL/JIT): {speedup:.2f}x")
test/registered/kernels/ops/attention/test_q8kv8_sparse_prefill_backend.py
+4 -3
View File
@@ -442,9 +442,10 @@ def test_q8kv8_sparse_prefill_helper_builds_fp8_workspace_matching_bf16_path(
bf16_capture = _Capture()
q8_capture = _Capture()
with _patched_sparse_kernels(
bf16_capture, q8_capture
), _patched_compressed_sparse_cache_paths(compress_ratio):
with (
_patched_sparse_kernels(bf16_capture, q8_capture),
_patched_compressed_sparse_cache_paths(compress_ratio),
):
bf16_out = backend._forward_prefill_sparse(
q=q,
layer_id=0,
test/registered/kernels/ops/attention/test_qprep_bf16_fp8_sm90.py
+3 -3
View File
@@ -109,9 +109,9 @@ def test_qprep_fp64_reference_parity(h_k):
).transpose(0, 1)
err_tri = (ref8[..., :N_LORA].to(torch.float64) - ref64).abs().mean().item()
err_cuda = (out8[..., :N_LORA].to(torch.float64) - ref64).abs().mean().item()
assert (
err_cuda <= 1.05 * err_tri
), f"CUDA fp64-ref mean |err| {err_cuda:.4e} exceeds Triton's {err_tri:.4e}"
assert err_cuda <= 1.05 * err_tri, (
f"CUDA fp64-ref mean |err| {err_cuda:.4e} exceeds Triton's {err_tri:.4e}"
)
if __name__ == "__main__":
test/registered/kernels/ops/attention/test_sm120_paged_mqa_logits.py
+3 -3
View File
@@ -149,9 +149,9 @@ def _compare(
# Invalid positions in SM120 output must be -inf
positions = torch.arange(max_seq_len, device=sm120.device)
invalid = positions.unsqueeze(0) >= seq_lens.unsqueeze(1)
assert torch.all(
torch.isinf(sm120[invalid]) & (sm120[invalid] < 0)
), "SM120 output must fill invalid positions with -inf"
assert torch.all(torch.isinf(sm120[invalid]) & (sm120[invalid] < 0)), (
"SM120 output must fill invalid positions with -inf"
)
class TestSM120PagedMqaLogitsTorch(CustomTestCase):
test/registered/kernels/ops/attention/test_topk_v2.py
+3 -3
View File
@@ -96,9 +96,9 @@ def _assert_topk_close(scores_cpu, ref_raw, our_raw, bs, seq_lens, k):
print(
f"b={i} L={L} k={k}: more={list(more)[:4]} less={list(less)[:4]} mv={mv[:3]} lv={lv[:3]}"
)
assert len(our) == min(
k, L
), f"b={i} L={L} k={k}: {len(our)} valid != {min(k, L)}"
assert len(our) == min(k, L), (
f"b={i} L={L} k={k}: {len(our)} valid != {min(k, L)}"
)
assert bad <= MAX_PERMIT_ERROR, f"{bad=} > {MAX_PERMIT_ERROR}"
test/registered/kernels/ops/communication/test_amd_deterministic_custom_allreduce.py
+7 -7
View File
@@ -95,7 +95,7 @@ def worker(world_size, rank, port):
input_size_bytes = base_input.numel() * base_input.element_size()
if input_size_bytes > custom_ar.max_size and rank == 0:
print(
f"Warning: Input size ({input_size_bytes/(1024*1024):.1f} MB) exceeds buffer size ({custom_ar.max_size/(1024*1024):.1f} MB)"
f"Warning: Input size ({input_size_bytes / (1024 * 1024):.1f} MB) exceeds buffer size ({custom_ar.max_size / (1024 * 1024):.1f} MB)"
)
print(" Using unregistered mode (will copy to buffer)")
@@ -105,9 +105,9 @@ def worker(world_size, rank, port):
# TEST 1: Deterministic kernel (same batch size) - should be DETERMINISTIC
# =========================================================================
if rank == 0:
print(f"\n{'='*70}")
print(f"\n{'=' * 70}")
print("TEST 1: Deterministic kernel (same batch size)")
print(f"{'='*70}")
print(f"{'=' * 70}")
dist.barrier()
results_allreduce_only = []
@@ -125,7 +125,7 @@ def worker(world_size, rank, port):
if rank == 0:
print(
f" Trial {trial+1:2d}: sum={checksum:.6f}, first5={first_vals.tolist()}"
f" Trial {trial + 1:2d}: sum={checksum:.6f}, first5={first_vals.tolist()}"
)
# Check determinism
@@ -135,7 +135,7 @@ def worker(world_size, rank, port):
for i, (s, vals) in enumerate(results_allreduce_only[1:], 1):
if abs(ref_sum - s) > 1e-3 or not torch.allclose(ref_vals, vals, rtol=1e-3):
all_match = False
print(f" Trial {i+1} DIFFERS! ref_sum={ref_sum:.6f}, got={s:.6f}")
print(f" Trial {i + 1} DIFFERS! ref_sum={ref_sum:.6f}, got={s:.6f}")
if all_match:
print(" ✓ DETERMINISTIC KERNEL (fixed BS): DETERMINISTIC (as expected)")
@@ -151,10 +151,10 @@ def worker(world_size, rank, port):
# [a], [a, x], [a, x, x], ...
# =========================================================================
if rank == 0:
print(f"\n{'='*70}")
print(f"\n{'=' * 70}")
print("TEST 2: Deterministic kernel (different batch size)")
print("Batches: [a], [a,x], [a,x,x], ...")
print(f"{'='*70}")
print(f"{'=' * 70}")
dist.barrier()
results_allreduce_only = {trial: [] for trial in range(num_trials)}
test/registered/kernels/ops/communication/test_amd_nccl_allreduce_determinism.py
+6 -6
View File
@@ -63,9 +63,9 @@ def worker(world_size, rank, port):
# TEST 1: Default all-reduce (same batch size) - should be DETERMINISTIC
# =========================================================================
if rank == 0:
print(f"\n{'='*70}")
print(f"\n{'=' * 70}")
print("TEST 1: Default NCCL all_reduce (same batch size)")
print(f"{'='*70}")
print(f"{'=' * 70}")
dist.barrier()
results_allreduce_only = []
@@ -84,7 +84,7 @@ def worker(world_size, rank, port):
if rank == 0:
print(
f" Trial {trial+1:2d}: sum={checksum:.6f}, first5={first_vals.tolist()}"
f" Trial {trial + 1:2d}: sum={checksum:.6f}, first5={first_vals.tolist()}"
)
# Check determinism
@@ -94,7 +94,7 @@ def worker(world_size, rank, port):
for i, (s, vals) in enumerate(results_allreduce_only[1:], 1):
if abs(ref_sum - s) > 1e-3 or not torch.allclose(ref_vals, vals, rtol=1e-3):
all_match = False
print(f" Trial {i+1} DIFFERS! ref_sum={ref_sum:.6f}, got={s:.6f}")
print(f" Trial {i + 1} DIFFERS! ref_sum={ref_sum:.6f}, got={s:.6f}")
if all_match:
print(" ✓ DEFAULT ALL_REDUCE (fixed BS): DETERMINISTIC (as expected)")
@@ -108,10 +108,10 @@ def worker(world_size, rank, port):
# [a], [a, x], [a, x, x], ...
# =========================================================================
if rank == 0:
print(f"\n{'='*70}")
print(f"\n{'=' * 70}")
print("TEST 2: Default NCCL all_reduce (different batch size)")
print("Batches: [a], [a,x], [a,x,x], ...")
print(f"{'='*70}")
print(f"{'=' * 70}")
dist.barrier()
results_allreduce_only = {trial: [] for trial in range(num_trials)}
test/registered/kernels/ops/gemm/test_minimax_fused_qkv_index_gemm.py
+3 -3
View File
@@ -58,9 +58,9 @@ def test_fused_equals_separate(T, N1, N2, K):
fused = mxfp8_linear(x, w, sp)
assert fused.shape == ref.shape
assert torch.equal(
fused, ref
), f"max abs diff {(fused.float() - ref.float()).abs().max().item()}"
assert torch.equal(fused, ref), (
f"max abs diff {(fused.float() - ref.float()).abs().max().item()}"
)
if __name__ == "__main__":
test/registered/kernels/ops/kv_canary/test_kernel_config.py
+33 -33
View File
@@ -135,15 +135,15 @@ def test_verify_byte_equal_across_repeated_launches_10x() -> None:
snapshot_bufs.append(cuda_buf.clone())
for i in range(1, num_launches):
assert torch.equal(
snapshot_rings[0], snapshot_rings[i]
), f"violation_ring differs between launch 0 and {i}"
assert torch.equal(
snapshot_write_indices[0], snapshot_write_indices[i]
), f"violation_write_index differs between launch 0 and {i}"
assert torch.equal(
snapshot_bufs[0], snapshot_bufs[i]
), f"canary_buf differs between launch 0 and {i}"
assert torch.equal(snapshot_rings[0], snapshot_rings[i]), (
f"violation_ring differs between launch 0 and {i}"
)
assert torch.equal(snapshot_write_indices[0], snapshot_write_indices[i]), (
f"violation_write_index differs between launch 0 and {i}"
)
assert torch.equal(snapshot_bufs[0], snapshot_bufs[i]), (
f"canary_buf differs between launch 0 and {i}"
)
def test_write_byte_equal_across_repeated_launches_10x() -> None:
@@ -190,15 +190,15 @@ def test_write_byte_equal_across_repeated_launches_10x() -> None:
snapshot_counters.append(cuda_log.slot_run_counter.clone())
for i in range(1, num_launches):
assert torch.equal(
snapshot_bufs[0], snapshot_bufs[i]
), f"canary_buf differs between launch 0 and {i}"
assert torch.equal(
snapshot_rings[0], snapshot_rings[i]
), f"violation_ring differs between launch 0 and {i}"
assert torch.equal(
snapshot_counters[0], snapshot_counters[i]
), f"slot_run_counter differs between launch 0 and {i}"
assert torch.equal(snapshot_bufs[0], snapshot_bufs[i]), (
f"canary_buf differs between launch 0 and {i}"
)
assert torch.equal(snapshot_rings[0], snapshot_rings[i]), (
f"violation_ring differs between launch 0 and {i}"
)
assert torch.equal(snapshot_counters[0], snapshot_counters[i]), (
f"slot_run_counter differs between launch 0 and {i}"
)
def test_plan_byte_equal_across_repeated_launches_10x() -> None:
@@ -254,18 +254,18 @@ def test_plan_byte_equal_across_repeated_launches_10x() -> None:
snapshot_write_offsets.append(triton_w.write_offsets.clone())
for i in range(1, num_launches):
assert torch.equal(
snapshot_slots[0], snapshot_slots[i]
), f"verify_slot_indices differs between launch 0 and {i}"
assert torch.equal(
snapshot_positions[0], snapshot_positions[i]
), f"verify_expected_positions differs between launch 0 and {i}"
assert torch.equal(
snapshot_prevs[0], snapshot_prevs[i]
), f"verify_prev_slot_indices differs between launch 0 and {i}"
assert torch.equal(
snapshot_write_offsets[0], snapshot_write_offsets[i]
), f"write_offsets differs between launch 0 and {i}"
assert torch.equal(snapshot_slots[0], snapshot_slots[i]), (
f"verify_slot_indices differs between launch 0 and {i}"
)
assert torch.equal(snapshot_positions[0], snapshot_positions[i]), (
f"verify_expected_positions differs between launch 0 and {i}"
)
assert torch.equal(snapshot_prevs[0], snapshot_prevs[i]), (
f"verify_prev_slot_indices differs between launch 0 and {i}"
)
assert torch.equal(snapshot_write_offsets[0], snapshot_write_offsets[i]), (
f"write_offsets differs between launch 0 and {i}"
)
def test_verify_multi_launch_100x_counter_linear() -> None:
@@ -301,9 +301,9 @@ def test_verify_multi_launch_100x_counter_linear() -> None:
torch.cuda.synchronize()
assert (
int(cuda_log.kernel_run_counter[0].item()) == num_launches
), f"kernel_run_counter expected {num_launches}, got {cuda_log.kernel_run_counter[0].item()}"
assert int(cuda_log.kernel_run_counter[0].item()) == num_launches, (
f"kernel_run_counter expected {num_launches}, got {cuda_log.kernel_run_counter[0].item()}"
)
assert int(cuda_log.slot_run_counter[0].item()) == num_launches, (
f"slot_run_counter expected {num_launches} (1 active entry x 100 launches), "
f"got {cuda_log.slot_run_counter[0].item()}"
test/registered/kernels/ops/kv_canary/test_pipeline_e2e.py
+6 -6
View File
@@ -529,9 +529,9 @@ def test_pipeline_pseudo_mode_on_token_mismatch_then_verify_clean() -> None:
)
write_violations = int(log_real.write_index[0].item())
assert (
write_violations == n_tokens
), f"expected {n_tokens} write violations, got {write_violations}"
assert write_violations == n_tokens, (
f"expected {n_tokens} write violations, got {write_violations}"
)
def test_pipeline_empty_batch() -> None:
@@ -779,9 +779,9 @@ def test_pipeline_token_mismatch_detected_via_pool() -> None:
fail_bits = int(
log_real.ring[row_idx, consts.VIOLATION_FIELD_FAIL_REASON_BITS].item()
)
assert fail_bits & int(
consts.FailReason.VERIFY_TOKEN_MISMATCH
), f"row {row_idx}: VERIFY_TOKEN_MISMATCH bit missing in {fail_bits:#b}"
assert fail_bits & int(consts.FailReason.VERIFY_TOKEN_MISMATCH), (
f"row {row_idx}: VERIFY_TOKEN_MISMATCH bit missing in {fail_bits:#b}"
)
stored = int(log_real.ring[row_idx, consts.VIOLATION_FIELD_STORED_TOKEN].item())
expected = int(
log_real.ring[row_idx, consts.VIOLATION_FIELD_EXPECTED_TOKEN].item()
test/registered/kernels/ops/kv_canary/test_plan_hand.py
+27 -27
View File
@@ -282,9 +282,9 @@ class TestSeedSlot:
write_req_capacity=write_req_capacity,
)
actual_seed = int(w_plan.write_seed_slot_indices[0].item())
assert (
actual_seed == expected_seed
), f"[{label}] permuted-LUT seed expected {expected_seed} got {actual_seed}"
assert actual_seed == expected_seed, (
f"[{label}] permuted-LUT seed expected {expected_seed} got {actual_seed}"
)
def test_swa_window_head_prev_slot_is_real_predecessor(self) -> None:
"""SWA window with non-zero window_start: head entry's prev_slot != -1; it is the real predecessor."""
@@ -327,12 +327,12 @@ class TestSeedSlot:
write_req_capacity=write_req_capacity,
)
actual_prev = int(v_plan.verify_prev_slot_indices[0].item())
assert (
actual_prev != -1
), f"[{label}] SWA window head must have real predecessor, got -1"
assert (
actual_prev == expected_prev
), f"[{label}] expected prev={expected_prev} got {actual_prev}"
assert actual_prev != -1, (
f"[{label}] SWA window head must have real predecessor, got -1"
)
assert actual_prev == expected_prev, (
f"[{label}] expected prev={expected_prev} got {actual_prev}"
)
class TestPadding:
@@ -404,9 +404,9 @@ class TestPadding:
write_req_capacity=write_req_capacity,
)
actual_slots = v_plan.verify_slot_indices[:prefix].detach().cpu().tolist()
assert (
actual_slots == expected_slots
), f"[{label}] sparse-rtt slots expected {expected_slots} got {actual_slots}"
assert actual_slots == expected_slots, (
f"[{label}] sparse-rtt slots expected {expected_slots} got {actual_slots}"
)
def test_padding_row_with_garbage_prefix_does_not_oob(self) -> None:
"""rpi==0 padding row with absurd prefix_lens must not OOB-read req_to_token (row is skipped)."""
@@ -435,9 +435,9 @@ class TestPadding:
write_req_capacity=write_req_capacity,
)
assert int(v_plan.verify_num_valid[0].item()) == 8, label
assert (
int(w_plan.write_seed_slot_indices[1].item()) == -1
), f"[{label}] padding row seed must be -1"
assert int(w_plan.write_seed_slot_indices[1].item()) == -1, (
f"[{label}] padding row seed must be -1"
)
PlanInvariants.assert_all(
verify_plan=v_plan,
write_plan=w_plan,
@@ -848,9 +848,9 @@ class TestMisc:
tail_offsets = (
write_plan.write_offsets[n_active + 1 : 8].detach().cpu().tolist()
)
assert all(
v == 0 for v in tail_offsets
), f"[{label}] stale write_offsets tail not cleared: {tail_offsets}"
assert all(v == 0 for v in tail_offsets), (
f"[{label}] stale write_offsets tail not cleared: {tail_offsets}"
)
class TestVerifyContent:
@@ -983,16 +983,16 @@ class TestByteEqual:
]
for i, value in enumerate(expected_write_offsets):
assert (
int(triton_w.write_offsets[i].item()) == value
), f"write_offsets[{i}] expected {value} got {int(triton_w.write_offsets[i].item())}"
assert (
int(triton_v.verify_num_valid[0].item()) == expected_verify_num_valid
), f"verify_num_valid expected {expected_verify_num_valid}"
assert int(triton_w.write_offsets[i].item()) == value, (
f"write_offsets[{i}] expected {value} got {int(triton_w.write_offsets[i].item())}"
)
assert int(triton_v.verify_num_valid[0].item()) == expected_verify_num_valid, (
f"verify_num_valid expected {expected_verify_num_valid}"
)
for i, expected_seed in enumerate(expected_seeds):
assert (
int(triton_w.write_seed_slot_indices[i].item()) == expected_seed
), f"write_seed_slot_indices[{i}] expected {expected_seed}"
assert int(triton_w.write_seed_slot_indices[i].item()) == expected_seed, (
f"write_seed_slot_indices[{i}] expected {expected_seed}"
)
class TestBoundarySweep:
test/registered/kernels/ops/kv_canary/test_verify_hand.py
+18 -18
View File
@@ -380,9 +380,9 @@ class TestChain:
buf_pair=buf_pair, plan_pair=plan_pair, assert_equal=False
)
assert (
_n_violations(cuda_log) == 0
), f"unexpected violation at iteration token={token} position={position} slot={slot_idx}"
assert _n_violations(cuda_log) == 0, (
f"unexpected violation at iteration token={token} position={position} slot={slot_idx}"
)
def test_prev_slot_padding_skips_chain_check_arbitrary_stored_hash(self) -> None:
"""prev_slot_idx == TOKEN_TO_KV_SLOT_PADDING → chain check is skipped, regardless of stored chain hash."""
@@ -692,9 +692,9 @@ class TestViolationField:
f"(bit_to_trigger={bit_to_trigger} injection_position={injection_position})"
)
else:
assert (
_n_violations(cuda_log) > ring_capacity
), "write_index did not advance beyond ring_capacity after overflow"
assert _n_violations(cuda_log) > ring_capacity, (
"write_index did not advance beyond ring_capacity after overflow"
)
def test_position_mismatch_sets_position_bit_only(self) -> None:
"""Plan.position != stored.position with chain hash correct → only POSITION bit set."""
@@ -706,12 +706,12 @@ class TestViolationField:
cuda_log, _ = run_verify_diff(buf_pair=buf_pair, plan_pair=plan_pair)
assert _n_violations(cuda_log) == 1
bits = _fail_bits(cuda_log)
assert (
bits & consts.FailReason.VERIFY_POSITION_MISMATCH
), f"expected POSITION bit, got {bits:#b}"
assert (
bits & consts.FailReason.VERIFY_CHAIN_HASH_MISMATCH
) == 0, f"chain hash bit unexpectedly set: {bits:#b}"
assert bits & consts.FailReason.VERIFY_POSITION_MISMATCH, (
f"expected POSITION bit, got {bits:#b}"
)
assert (bits & consts.FailReason.VERIFY_CHAIN_HASH_MISMATCH) == 0, (
f"chain hash bit unexpectedly set: {bits:#b}"
)
class TestRealKvHash:
@@ -990,9 +990,9 @@ class TestRealKvHash:
assert _n_violations(cuda_log) >= 1
bits = _fail_bits(cuda_log)
assert (
bits & consts.FailReason.VERIFY_REAL_KV_HASH_MISMATCH
), f"expected REAL_KV_HASH bit, got {bits:#b}"
assert bits & consts.FailReason.VERIFY_REAL_KV_HASH_MISMATCH, (
f"expected REAL_KV_HASH bit, got {bits:#b}"
)
def test_real_kv_off_does_not_deref_real_kv_sources(self) -> None:
buf_pair = _buf_pair(num_slots=8)
@@ -1738,9 +1738,9 @@ class TestViolationRing:
plan_slot_set = set(slot_indices)
for row in range(n_violations):
kind = int(cuda_log.ring[row, consts.VIOLATION_FIELD_KERNEL_KIND].item())
assert kind == int(
launch_tag
), f"row {row} kind {kind} != {int(launch_tag)}"
assert kind == int(launch_tag), (
f"row {row} kind {kind} != {int(launch_tag)}"
)
slot = int(cuda_log.ring[row, 1].item())
assert slot in plan_slot_set, f"row {row} slot {slot} not in plan"
test/registered/kernels/ops/kv_canary/test_write_fuzz.py
+3 -3
View File
@@ -204,9 +204,9 @@ def _run_one(inputs: WriteFuzzInputs) -> None:
kernel_kind=inputs.kernel_kind,
assert_equal=False,
)
assert torch.equal(
inputs.cuda_canary_buf, inputs.ref_canary_buf
), "CUDA vs ref canary_buf diverged"
assert torch.equal(inputs.cuda_canary_buf, inputs.ref_canary_buf), (
"CUDA vs ref canary_buf diverged"
)
assert int(cuda_log.write_index[0].item()) == int(ref_log.write_index[0].item())
assert int(cuda_log.slot_run_counter[0].item()) == int(
ref_log.slot_run_counter[0].item()
test/registered/kernels/ops/kv_canary/test_write_hand.py
+21 -21
View File
@@ -362,9 +362,9 @@ class TestSeedSlot:
new_stored = read_slot_fields(canary_buf=self.buf_pair[0], slot_idx=new_slot)
assert new_stored[0] == new_token
assert new_stored[1] == new_position
assert new_stored[2] == to_signed_int64(
expected_running
), f"new slot prev_hash {new_stored[2]} != expected {to_signed_int64(expected_running)}"
assert new_stored[2] == to_signed_int64(expected_running), (
f"new slot prev_hash {new_stored[2]} != expected {to_signed_int64(expected_running)}"
)
class TestChain:
@@ -446,9 +446,9 @@ class TestChain:
stored_prev_signed, stored_real_kv_hash = read_slot_fields(
canary_buf=cuda_buf, slot_idx=slot_idx
)[2:]
assert stored_prev_signed == to_signed_int64(
running
), f"slot {slot_idx}: stored prev_hash != recomputed chain step"
assert stored_prev_signed == to_signed_int64(running), (
f"slot {slot_idx}: stored prev_hash != recomputed chain step"
)
running = splitmix64_mix3(running, token, position)
@@ -711,9 +711,9 @@ class TestSlotHandling:
for slot in range(cuda_buf.shape[0]):
if slot in (5, 7):
continue
assert torch.equal(
after[slot], cuda_buf_before_slot_view[slot]
), f"slot {slot} should not have been written"
assert torch.equal(after[slot], cuda_buf_before_slot_view[slot]), (
f"slot {slot} should not have been written"
)
def test_shrink_active_reqs_does_not_write_stale_slots(self) -> None:
"""Run write with bs=3 plan after a bs=8 run on same buffer: stale slots from bs=8 stay intact."""
@@ -749,9 +749,9 @@ class TestSlotHandling:
after = cuda_buf.view(torch.int64)
for slot in big_slots:
assert torch.equal(
after[slot], untouched_snapshot[slot]
), f"slot {slot} from earlier bs=8 run was overwritten by bs=3 run"
assert torch.equal(after[slot], untouched_snapshot[slot]), (
f"slot {slot} from earlier bs=8 run was overwritten by bs=3 run"
)
class TestRealKvHash:
@@ -927,9 +927,9 @@ class TestRealKvHash:
_, _, _, stored_real_kv_hash = read_slot_fields(
canary_buf=self.buf_pair[0], slot_idx=0
)
assert stored_real_kv_hash == to_signed_int64(
expected_hash
), f"stored_real_kv_hash={stored_real_kv_hash:#x} expected={to_signed_int64(expected_hash):#x}"
assert stored_real_kv_hash == to_signed_int64(expected_hash), (
f"stored_real_kv_hash={stored_real_kv_hash:#x} expected={to_signed_int64(expected_hash):#x}"
)
def test_paged_real_kv_hash_consistent_across_slots(self) -> None:
"""page=16: writing two slots inside same page yields independent real_kv_hash per slot."""
@@ -1008,9 +1008,9 @@ class TestRealKvHash:
fields_b = _run_with(sources_b)
assert fields_a[3] != 0
assert fields_b[3] != 0
assert (
fields_a[3] != fields_b[3]
), "reversing source order must change real_kv_hash (fold is ordered)"
assert fields_a[3] != fields_b[3], (
"reversing source order must change real_kv_hash (fold is ordered)"
)
class TestRunCounter:
@@ -1201,9 +1201,9 @@ class TestPseudoMode:
)
assert int(cuda_log.write_index[0].item()) >= 1
bits = int(cuda_log.ring[0, consts.VIOLATION_FIELD_FAIL_REASON_BITS].item())
assert (
bits & consts.FailReason.WRITE_TOKEN_MISMATCH
), f"expected WRITE_TOKEN_MISMATCH bit, got {bits:#b}"
assert bits & consts.FailReason.WRITE_TOKEN_MISMATCH, (
f"expected WRITE_TOKEN_MISMATCH bit, got {bits:#b}"
)
def test_pseudo_mode_off_skips_token_check(self) -> None:
"""enable_write_verify_inputs=False makes the caller pass no expected-input tensors."""
test/registered/kernels/ops/kvcache/test_hicache.py
+6 -6
View File
@@ -97,9 +97,9 @@ def _run_transfer_roundtrip_mha(layout: str, element_dim: int) -> None:
device_pool=device_pool,
layout=layout,
)
assert (
host_pool.can_use_jit
), f"Expected JIT HiCache kernel for MHA dim={element_dim}"
assert host_pool.can_use_jit, (
f"Expected JIT HiCache kernel for MHA dim={element_dim}"
)
for layer_id in range(NUM_LAYERS):
_copy_tensor_with_offset(device_pool.k_buffer[layer_id], layer_id)
@@ -191,9 +191,9 @@ def _run_transfer_roundtrip_mla(layout: str, element_dim: int) -> None:
device_pool=device_pool,
layout=layout,
)
assert (
host_pool.can_use_jit
), f"Expected JIT HiCache kernel for MLA dim={element_dim}"
assert host_pool.can_use_jit, (
f"Expected JIT HiCache kernel for MLA dim={element_dim}"
)
for layer_id in range(NUM_LAYERS):
_copy_tensor_with_offset(device_pool.kv_buffer[layer_id], layer_id)
test/registered/kernels/ops/layernorm/test_gemma4_fused_routing.py
+3 -3
View File
@@ -68,9 +68,9 @@ def test_matches_reference(fused_routing, dtype, T, E, K):
# Tie-break order may differ; require the same top-K set and weight sum.
ref_set = ref_i.sort(dim=-1).values
out_set = out_i.sort(dim=-1).values
assert torch.equal(
out_set, ref_set
), "fused routing picked a different top-K set than reference"
assert torch.equal(out_set, ref_set), (
"fused routing picked a different top-K set than reference"
)
torch.testing.assert_close(
out_w.sum(dim=-1).to(torch.float32),
ref_w.sum(dim=-1).to(torch.float32),
test/registered/kernels/ops/layernorm/test_rmsnorm_hf.py
+3 -3
View File
@@ -98,9 +98,9 @@ def test_rmsnorm_hf_matches_hf_not_sgl(dtype: torch.dtype) -> None:
assert (sgl_ref - hf_ref).abs().max() > 0, "inputs don't exercise the difference"
diff_hf = (out - hf_ref).abs().max().item()
diff_sgl = (out - sgl_ref).abs().max().item()
assert (
diff_hf < diff_sgl
), f"kernel closer to SGL than HF (hf={diff_hf}, sgl={diff_sgl})"
assert diff_hf < diff_sgl, (
f"kernel closer to SGL than HF (hf={diff_hf}, sgl={diff_sgl})"
)
def test_rmsnorm_hf_empty_input() -> None:
test/registered/kernels/ops/moe/test_minimax_quant_scatter.py
+3 -3
View File
@@ -94,9 +94,9 @@ def test_quant_scatter_matches_quant_plus_fill(num_tokens, topk, hidden, group):
assert torch.equal(
gi_new[e, m].view(torch.uint8), gi_ref[e, m].view(torch.uint8)
), f"fp8 mismatch token={t} slot={j} expert={e}"
assert torch.equal(
gs_new[e, :, m], gs_ref[e, :, m]
), f"scale mismatch token={t} slot={j} expert={e}"
assert torch.equal(gs_new[e, :, m], gs_ref[e, :, m]), (
f"scale mismatch token={t} slot={j} expert={e}"
)
def test_standard_deepgemm_preprocess_quantizes_with_ue8m0_scale():
test/registered/kernels/ops/moe/test_moe_lora_align_block_size.py
+3 -3
View File
@@ -159,9 +159,9 @@ def test_moe_lora_align_block_size(
# Check that all tokens in this block truly belong to 'lora_idx'
actual_owners = token_ownership[original_token_indices]
assert torch.all(
actual_owners == lora_idx
), f"Kernel put tokens from LoRA {actual_owners} into block for LoRA {lora_idx}"
assert torch.all(actual_owners == lora_idx), (
f"Kernel put tokens from LoRA {actual_owners} into block for LoRA {lora_idx}"
)
if __name__ == "__main__":
test/registered/kernels/ops/moe/test_moe_topk_sigmoid.py
+24 -22
View File
@@ -319,12 +319,14 @@ def test_topk_sigmoid_vs_ref(num_tokens, num_experts, topk, dtype, renormalize):
ref_w.sort(dim=-1)[0],
atol=1e-3,
rtol=1e-3,
), f"Weight mismatch (dtype={dtype}, n_exp={num_experts}, topk={topk}, renorm={renormalize})"
), (
f"Weight mismatch (dtype={dtype}, n_exp={num_experts}, topk={topk}, renorm={renormalize})"
)
# Exact index match is only reliable for float32 (fp16/bf16 tie-breaking may differ)
if dtype == torch.float32:
assert torch.equal(
topk_i, ref_i
), f"Index mismatch (dtype={dtype}, n_exp={num_experts}, topk={topk})"
assert torch.equal(topk_i, ref_i), (
f"Index mismatch (dtype={dtype}, n_exp={num_experts}, topk={topk})"
)
# ---------------------------------------------------------------------------
@@ -351,12 +353,12 @@ def test_topk_sigmoid_with_correction_bias(num_tokens, num_experts, topk, renorm
ref_w, ref_i = topk_sigmoid_ref(gating, topk, renormalize, correction_bias=bias)
assert torch.allclose(
topk_w, ref_w, atol=1e-3, rtol=1e-3
), f"Weight mismatch with bias (n_exp={num_experts}, topk={topk}, renorm={renormalize})"
assert torch.equal(
topk_i, ref_i
), f"Index mismatch with bias (n_exp={num_experts}, topk={topk})"
assert torch.allclose(topk_w, ref_w, atol=1e-3, rtol=1e-3), (
f"Weight mismatch with bias (n_exp={num_experts}, topk={topk}, renorm={renormalize})"
)
assert torch.equal(topk_i, ref_i), (
f"Index mismatch with bias (n_exp={num_experts}, topk={topk})"
)
# ---------------------------------------------------------------------------
@@ -394,12 +396,12 @@ def test_topk_sigmoid_with_fused_shared_experts(
gating, topk + 1, renormalize, correction_bias=bias, num_fused_shared_experts=1
)
assert torch.allclose(
topk_w, ref_w, atol=1e-3, rtol=1e-3
), f"Weight mismatch with bias (n_exp={num_experts}, topk={topk}, renorm={renormalize})"
assert torch.equal(
topk_i, ref_i
), f"Index mismatch with bias (n_exp={num_experts}, topk={topk})"
assert torch.allclose(topk_w, ref_w, atol=1e-3, rtol=1e-3), (
f"Weight mismatch with bias (n_exp={num_experts}, topk={topk}, renorm={renormalize})"
)
assert torch.equal(topk_i, ref_i), (
f"Index mismatch with bias (n_exp={num_experts}, topk={topk})"
)
# ---------------------------------------------------------------------------
@@ -484,12 +486,12 @@ def test_topk_sigmoid_vs_aot(num_tokens, num_experts, topk, dtype, renormalize):
topk_i_aot = torch.empty((num_tokens, topk), dtype=torch.int32, device="cuda")
topk_sigmoid_aot(topk_w_aot, topk_i_aot, gating, renormalize=renormalize)
assert torch.allclose(
topk_w_jit, topk_w_aot, atol=1e-3, rtol=1e-3
), f"JIT vs AOT weight mismatch (dtype={dtype}, n_exp={num_experts}, topk={topk})"
assert torch.equal(
topk_i_jit, topk_i_aot
), f"JIT vs AOT index mismatch (dtype={dtype}, n_exp={num_experts}, topk={topk})"
assert torch.allclose(topk_w_jit, topk_w_aot, atol=1e-3, rtol=1e-3), (
f"JIT vs AOT weight mismatch (dtype={dtype}, n_exp={num_experts}, topk={topk})"
)
assert torch.equal(topk_i_jit, topk_i_aot), (
f"JIT vs AOT index mismatch (dtype={dtype}, n_exp={num_experts}, topk={topk})"
)
if __name__ == "__main__":
test/registered/kernels/ops/moe/test_shuffle_rows_with_scales.py
+3 -3
View File
@@ -67,9 +67,9 @@ def _inputs(k, num_src_rows, num_dst_rows, seed):
def _assert_same_bytes(got, ref, what):
assert torch.equal(
got.view(torch.int8), ref.view(torch.int8)
), f"{what} bytes differ"
assert torch.equal(got.view(torch.int8), ref.view(torch.int8)), (
f"{what} bytes differ"
)
@pytest.mark.parametrize("k,num_src_rows,num_dst_rows", CASES)
test/registered/kernels/ops/quantization/test_per_token_group_quant_8bit_v2.py
+3 -3
View File
@@ -179,9 +179,9 @@ def test_v2_jit_masked_matches_aot(num_experts, hidden, tokens_pad):
)
torch.cuda.synchronize()
assert torch.equal(
x_q.view(torch.int8), q_ref.view(torch.int8)
), "masked fp8 differ"
assert torch.equal(x_q.view(torch.int8), q_ref.view(torch.int8)), (
"masked fp8 differ"
)
assert torch.equal(x_s, s_ref), "masked scales differ"
test/registered/kernels/test_jit_cache.py
+3 -3
View File
@@ -187,9 +187,9 @@ def test_no_unordered_container_reaches_the_key(monkeypatch):
_build_key()
def walk(value, path="parts"):
assert not isinstance(
value, (set, frozenset, dict)
), f"unordered container at {path}: {type(value).__name__}"
assert not isinstance(value, (set, frozenset, dict)), (
f"unordered container at {path}: {type(value).__name__}"
)
if isinstance(value, (list, tuple)):
for index, item in enumerate(value):
walk(item, f"{path}[{index}]")
test/registered/kernels/test_kda_replayssm_ring_ragged.py
+3 -3
View File
@@ -114,9 +114,9 @@ def _run_case(bs, gamma, HV, H, K, V, lens, acc, L, pad_last=False):
base = inter[slots[j], int(acc[j]) - 1]
fold = ckpt[slots[j]]
rel = ((fold - base).abs().max() / base.abs().max().clamp_min(1e-6)).item()
assert (
rel < 1e-3
), f"row={j} len={int(lens[j])} acc={int(acc[j])}: rel={rel:.3e}"
assert rel < 1e-3, (
f"row={j} len={int(lens[j])} acc={int(acc[j])}: rel={rel:.3e}"
)
@pytest.mark.parametrize("bs,gamma,HV,H,K,V", SHAPES, ids=SHAPE_IDS)
test/registered/kernels/test_kernel_inventory.py
+6 -4
View File
@@ -128,9 +128,9 @@ def test_registered_kernel_test_groups_are_known():
registered_root = REPO_ROOT / "test" / "registered" / "kernels"
for kind in ("ops", "benchmark"):
unknown = _directory_names(registered_root / kind) - declared_groups
assert (
not unknown
), f"Unknown {kind} kernel group directories: {sorted(unknown)}"
assert not unknown, (
f"Unknown {kind} kernel group directories: {sorted(unknown)}"
)
def test_internal_registry_target_attributes_are_declared():
@@ -172,7 +172,9 @@ def test_jit_source_declarations_exist():
function_name = (
call.func.id
if isinstance(call.func, ast.Name)
else call.func.attr if isinstance(call.func, ast.Attribute) else None
else call.func.attr
if isinstance(call.func, ast.Attribute)
else None
)
if function_name != "load_jit":
continue
test/registered/kernels/test_lean_attention.py
+3 -3
View File
@@ -209,9 +209,9 @@ def _run_pair_paged(
kv_group_num = H_Q // H_KV
sm = 1.0 / (D**0.5)
tot = B * S
assert (
tot % page_size == 0
), "test setup: total tokens must be a multiple of page_size"
assert tot % page_size == 0, (
"test setup: total tokens must be a multiple of page_size"
)
num_pages = tot // page_size
# 4-D paged KV buffers [num_pages, page_size, head, dim] (the shared-pool layout).
test/registered/kernels/test_quick_allreduce_bf16_range.py
+3 -3
View File
@@ -63,9 +63,9 @@ def _run_bf16_range_test(rank: int, world_size: int, port: int) -> None:
dist.barrier()
out = quick_all_reduce.quick_all_reduce(inp)
torch.cuda.synchronize()
assert (
torch.isfinite(out).all().item()
), f"{quant_mode=} {case_name=} produced non-finite output"
assert torch.isfinite(out).all().item(), (
f"{quant_mode=} {case_name=} produced non-finite output"
)
if quant_mode == "FP" or case_name in ("low", "ordinary"):
torch.testing.assert_close(
out,